diff --git a/TORAD.md b/TORAD.md index 65170dd3b140..76791593fe2b 100644 --- a/TORAD.md +++ b/TORAD.md @@ -135,6 +135,37 @@ which pins a commit of this branch as a submodule. | `6a565c9af` | A KDA layer's f16 or q8_0 recurrent cache takes the fused state write that `9e02027e4` and `9760e966a` gave the scalar gate only: under `-cts f16`, the served state, every glm5next KDA layer wrote its f32 state into the op's output and a CPY narrowed it into the cache. The 44-layer GLM-5.3 proxy on an RTX 5070 Ti, a token's GPU time (nsys node trace, tg32): `cpy_scalar_contiguous` 262 us against 289 for `gated_delta_net` itself. KDA stores its state as the scalar gate does, one column a warp through `gdn_store_state` (the CPY's f16 conversion and q8_0 formula); the matcher (`ggml_cuda_try_gdn_cache_fusion`) takes it and the dispatch instantiates KDA with the cache's type; the chunked prefill pipeline, f32 only, keeps its CPY. The proxy on the 5070 Ti alone, `-cts f16`, tg64 in 3 rounds alternated on the landed bin: 58.90 -> 60.23 tok/s (+2.26 %, per round +3.0, +2.4, +2.2); under `-sm tensor` on a 5080 + 5070 Ti, 4 rounds: 101.81 -> 104.49 tok/s (+2.63 %, per round +2.5, +5.7, +1.7, +2.6), the CPY's 176 and 179 us a token gone from both cards' traces. PPL at `-ub 1` the same with the CPY and fused, bit for bit: 386825.5220 on one card, 387892.2627 under `-sm tensor`. `test-backend-ops` `GATED_DELTA_NET_CACHE_FUSION` 77/77 with 19 new KDA cases and `GATED_DELTA_NET` 71/71; the f16 store or the q8_0 block scale scaled by 1.001 fails the 9 KDA f16 or the 9 KDA q8_0 cases, which pass under either mutation with the switch set | `GGML_CUDA_GDN_KDA_CACHE_LEGACY=1` | | `a6d65589c` | Off by default. With `GGML_CUDA_GDN_STATE_PERSIST=1` a Gated DeltaNet launch with the fused gather carries an L2 access policy window as a `cudaLaunchKernelEx` attribute (`ggml_cuda_kernel_launch_params.l2_window`, beside PDL's): one window over the whole buffer the recurrent states are in, the same for every layer, persisting at hitRatio set-aside / window; the set-aside (the smaller of the window and the device's most) is set once at the first window, from a thread in relaxed capture mode (a relaxed capture allows the call, a global one refuses it); the L2 issuer (`345065068`) then requests at most its share of the L2 outside the set-aside; the context's destructor returns its lines to normal. `GGML_CUDA_L2_PERSIST_CHECK=1` logs, after each capture, how many kernel nodes hold a window. A standalone graph on an RTX 5070 Ti (34 x 1 MB states behind 15 MB of weights each, 30 MB set aside): the attribute survives capture (34 of 34 nodes), the state kernel 1.67 against 2.39 us a layer (nsys); an idle process holding the lines left another's 24 MB L2-resident re-read at full speed (134.4 / 135.0 us). ncu cannot see it: its kernel profiling returns the lines to normal. The 44-layer GLM-5.3 proxy on the 5070 Ti alone, `-cts f16` (77 MB of states, hitRatio 0.40): `gated_delta_net` 289 -> 211 us a token, but tg64 -0.1 %, and -1.0 % over the fused KDA store: the set-aside costs the rest as much. Under `-sm tensor` on a 5080 + 5070 Ti over the fused store: `gated_delta_net` 177 -> 83 us a token on the 5080 (64 MB of L2) but 193 -> 188 on the 5070 Ti, whose `mul_mat_vec_q` lost 331 us (4,242 -> 4,573), and tg64 104.49 -> 103.56 tok/s (-0.9 %, 4 rounds). For the 2x RTX PRO 6000 head, 34 MB a card against 128 MB of L2, to be measured there. PPL at `-ub 1` the same with it and without, bit for bit | `GGML_CUDA_GDN_STATE_PERSIST=1` turns it on | | `3b94ce3b8` | Review of `6a565c9af` and `a6d65589c` (feature-dev code-reviewer, pr-review-toolkit silent-failure-hunter), the defects found against the tree, all in the persistence (the KDA store: none). With `GGML_CUDA_GDN_STATE_PERSIST=1`, a launch given no window says why, once a device: a device with no persisting L2, a GPU split into virtual devices, states outside a CUDA buffer, a layer's states past the largest window. A refused set-aside is a warning with the CUDA error, not an INFO "0 KiB"; the last error is cleared only when a call failed, and a limit set but not read back returns to 0. The contexts given a window are counted a device, and the last one's end resets the lines, sets the limit to 0 and the set-aside back to 0: before, any context's end reset the lines under the others still persisting, and the set-aside stayed for the process, the L2 issuer's budget short of it. `GGML_CUDA_L2_PERSIST_CHECK=1` says when no launch had a window, and counts a node whose attribute cannot be read apart. The dispatch asserts the state type is one it instantiates (f32, f16, q8_0): the assert `6a565c9af` removed, made general. `f020c0f76`: the warning names its function, not its lambda. Against `a6d65589c`'s bin on the RTX 5070 Ti (`review-check.sh`): llama-bench's two tests, two contexts in turn, print the set-aside once there and twice here; the check switch alone, nothing there and "no launch was given one" a capture here; `GGML_CUDA_DEVICES=2` on one GPU with the switch, no word there and the virtual-device warning here. `GATED_DELTA_NET_CACHE_FUSION` 77/77 with the switch on, PDL on and off (a 44,337 KiB window, 30,720 KiB set aside), and the gate's other legs at `f020c0f76`. With the switch off nothing changes | `GGML_CUDA_GDN_STATE_PERSIST=1`, as `a6d65589c` | +| `331d72782` | Qwen3.5's levers between the qkv group and the recurrence, each with its own switch: the alpha/beta pair folded into the conv-state update, the normed input reaching it by a release/acquire handoff (`ggml_cuda_ssm_conv_ab_slots`); a PQ2_0 launch's L2 prefetch of the weights of the kernels between it and the next launch; the conv's weights and cache row, and the recurrence's state, requested before their PDL waits; the recurrence and `mul_mat_vec_f` triggering the next launch at their start; the conv-state update capped at 56 registers with a 4-token instance; the PQ2_0 group writing the conv's inputs at 152 registers. Served bits on the RTX 5080 (Ternary Bonsai 2 27B PQ2_0, one slot and four, 128 greedy tokens with top-5 log-probabilities, `levers3-check.sh`): every arm bit-identical to all switches set. Open, found by the same check: the fold does not engage at this commit. The node trace's launches are the same with its switch and without (23,879), and a build whose fold sums its partials in another order is bit-identical, so the pair keeps its own launch. The caps leave 16-32 B stacks (the conv update at 56, the capped groups of 6 and 7 matrices), where the 168-register group of 7 had 40 B before. The speed sweep on the Qwen model, which sets the defaults, is still to run. GLM-5.3 proxy (no bf16 pair, so no fold): tg64 59.72 / 59.72 / 59.45 tok/s, the landed bin / a dev bin carrying these levers / the same with every switch set, alternated in one session on the RTX 5070 Ti | `GGML_CUDA_SSM_CONV_AB_LEGACY=1`, `GGML_CUDA_PQ2_MMA_GROUP_REGS_LEGACY=1`, `GGML_CUDA_PQ2_PREFETCH_BETWEEN_LEGACY=1`, `GGML_CUDA_SSM_CONV_PREWAIT_LEGACY=1`, `GGML_CUDA_SSM_CONV_STATE_PREFETCH_LEGACY=1`, `GGML_CUDA_GDN_STATE_PREFETCH_LEGACY=1`, `GGML_CUDA_GDN_TRIGGER_LEGACY=1`, `GGML_CUDA_MMVF_TRIGGER_LEGACY=1` | +| `e74ed87e1` | glm5next's KDA output gate `g_b(g_a(x))` reads the layer input, not the recurrence's output, but the graph built it after the recurrence: two small mat-vec launches on the chain between the recurrence and the output projection. It goes into the graph after the convolution, `g_a` beside `f_a` (two mat-vecs of one input, one launch by `5d58cc2`'s pair) and `g_b` before the recurrence. The 44-layer GLM-5.3 proxy on the RTX 5070 Ti alone, `-cts f16`, tg16 node trace: one small mat-vec launch fewer a KDA layer, the token's chain -0.36 % (about 60 us of 16.6 ms); PPL at `-ub 1` 386825.5220 before and after | none (graph order) | +| `7536ed775` | The recurrence's kernel wrote glm5next's KDA gated norm itself, `sigmoid(gate) * (rms_norm(o) * w)`: the head's last block by a ticket, rms_norm_f32's and unary_gated's arithmetic, the chain built in place over the attention rows (built apart, the allocator put the output on the recurrence's own inputs, and the 3-token verify graph declined every layer). Bit-identical: test-backend-ops, PPL at two shapes, served log-probabilities plain and with the MTP draft (a kernel flipping every output's last bit differs from position 0). Measured neutral against `e74ed87e1`, 44-layer proxy, rotated over 4 rounds: RTX 5070 Ti tg128 +0.24 % (95 % CI -0.95 to +1.42), pp3 at ubatch 3 -1.34 % (-3.33 to +0.65); a restructured epilogue (the ticket before the state stores) under `-sm tensor` on the 5080 + 5070 Ti, tg128 -0.71 % (-1.75 to +0.33). The rms_norm's 4 us after GDN was the wait for the recurrence's state stores to drain, which stays on the chain wherever it lands; the two launches cost about 1 us each under PDL and CUDA graphs, and the epilogue's ticket, reload and stores as much. `b850d706c` takes it out: the tree is `e74ed87e1`'s | none (taken out) | +| `bbe53707b` | The routed down projection's ring (`mmvq-moe.cu`) reads its ids, lists its experts and issues its block's first tiles before its PDL wait, then copies the tokens' vectors past it. The ids are whole before it starts: the gate/up ring before it on the stream read the same ids and triggers the next launch only past its own wait. The host marks such a launch by the stream's last ring ids, reset with the shared q8_1 copies at each evaluation's start and end. The kernel reads ids with `__ldcg`. The ring's consumers now trigger past their wait, as its producer does. The 44-layer GLM-5.3 proxy under `-sm tensor` on a 5080 + 5070 Ti, `-cts f16`, tg32 node traces over two rotated rounds (median over every MoE layer): the down 17.76 -> 16.19 us on the 5080 and 18.94 -> 17.25 on the 5070 Ti, and 17.79 / 19.01 with the switch set; the gate/up level. Each card's kernel time a token past the all-reduce: 7,779.5 -> 7,724.9 us (-0.70 %) and 8,885.7 -> 8,830.1 (-0.63 %). tg128 under `-sm tensor` over 4 rotated rounds: +0.75 % (95 % CI -1.43 to +2.94), too few rounds for a change this size. A t interval over rounds, not over samples: `ab-arms.sh` had computed it over each run's samples as if independent, about 3x too narrow. PPL bit for bit in every arm: 386825.5220 at `-ub 1` and 385507.0946 at `-ub 3` on one card, 387892.2627 under `-sm tensor`. `test-backend-ops` `MOE_FFN_CHAIN`, new: gate/up, GLU (clamped or not) and down on one ids, at n_ff 2048 and 1024, 1 and 3 tokens, 8/8. The down reading the next expert's rows fails 8/8, and dropping one tile (NMSE 5e-4, under the 5e-3 tolerance) does not. Tried after it and not landed: the down quantizing its own GLU activations in its prologue in place of the q8_1 launch between the two (MOE_FFN_CHAIN 8/8, a doubled block scale failing the 6 cases where it engages, PPL bit for bit). It lost: the down 17.47 -> 21.54 us on the 5070 Ti and 16.16 -> 19.94 on the 5080, against the q8_1 launch's 0.96 and 1.12, the MoE chain +3.1 and +2.4 us a layer. Every block quantizes every pair's vector (8 x 1,024 floats from L2) in two dependent rounds after its wait, while the tiles its producer already landed sit in the ring | `GGML_CUDA_MMVQ_MOE_IDS_EARLY_LEGACY=1` | +| `3a6886328` | The routed experts' weighted sum, `MUL(experts, weights)` and its n_used views summed by n_used - 1 ADDs, in one launch (`k_moe_weighted_sum`, `binbcast.cu`): each product and each sum rounded as the nodes round them, in their order, so bit for bit. `test-backend-ops` `MOE_WEIGHTED_SUM`, new, compares bit for bit (its error the count of differing elements, 0 allowed): n_embd 4096 and 2880, 8, 4 and 2 experts, 1, 3 and 8 tokens, 18/18; summing in reverse fails the 12 cases of 4 or 8 experts, and the 2-expert cases pass, as they must (two addends swapped are the same IEEE add). The 44-layer GLM-5.3 proxy under `-sm tensor` on a 5080 + 5070 Ti, `-cts f16`, tg32 node traces over two rotated rounds: the sum 1.83 -> 1.25 us a MoE layer on the 5080 and 1.76 -> 1.22 on the 5070 Ti, the MoE chain -0.56 / -0.58 us; each card's kernel time a token past the all-reduce 7,730.9 -> 7,705.1 us (-0.33 %) and 8,309.1 -> 8,280.7 (-0.34 %). No throughput A/B: the last one's interval was +-2.2 % at 4 rounds. PPL bit for bit: 386825.5220 at `-ub 1`, 385507.0946 at `-ub 3`, 387892.2627 under `-sm tensor`. Tried with it and not landed, each measured the same way against `bbe53707b`: (a) `mul_mat_vec_q` triggering the next launch past its PDL wait: -0.12 / +0.28 % alone, -0.46 / -0.09 % with the sum; it moves charged time between launches (the mat-vecs after the KDA q, k and v 4.5 -> 1.6 us, the recurrence +1.2 us, the rms_norm after it +1.0). (b) The shared expert's gate/up moved between the routing and the routed gate/up, with topk's ids marked whole so the routed gate/up issues its first tiles before its wait: the routed gate/up +1.0 us (`mul_mat_vec_q` triggers only at its end, so the tiles go out with the vector copy and queue ahead of it) and the token +2.3 to +3.7 %. The shared gate/up is 8.91 MB a card, a heavy launch for the L2 issuer: after it, the issue (9.9 to 15.9 MB) had prefetched the next layer's q, k and v, and moved before the routed experts its chain ends at the MUL_MAT_ID and the issue drops to 1.2 MB, the q, k and v +5.4 us a layer. (c) The issuer counting a MUL_MAT_ID as heavy, to give that issue back: +0.43 / +1.07 % on the moved order, +4.82 / +3.95 % on the landed one, where the new issue after the routed down runs while the shared gate/up does and slows it 4.16 -> 11.10 us | `GGML_CUDA_MOE_WSUM_LEGACY=1` | +| `90d1b9846` | The hyper-connection front's second kernel (`dsv4_hc_pre_gram_f32`) writes its normed mix's q8_1 copy as it writes the mix: a warp is a q8_1 block, quantized with `quantize_q8_1`'s arithmetic on the values it stores, so the same bits. The evaluation makes the copy for every such mix with a quantized reader on `mul_mat_vec_q` before any node runs, one a key, which each layer's front writes in turn (ggml-alloc gives every layer's mix the same bytes); a copy keeps the index of the node that last wrote it, so the writes of the front's own fused group leave it whole. Where the front does not fuse, or n_embd is not a multiple of 512, the first reader quantizes it as before. `test-backend-ops` `DSV4_HC_PRE_Q8_1`, new: the mix read by one or two Q8_0 or Q4_K MUL_MATs, 16/16; a doubled block scale in the copy fails the 12 cases whose readers read it, and passes the 2 where the reader quantizes and Q4_K at 8 tokens, which Blackwell runs on MMQ. The 44-layer GLM-5.3 proxy under `-sm tensor` on a 5080 + 5070 Ti, tg32 node traces: 2,640 quantize launches fewer a run of 30 evaluations, two a layer; each card's kernel time a token past the all-reduce -0.76 to -0.81 % and -0.30 to -0.34 % over three runs, -0.68 / -0.56 % with the L2 issuer off. The 5070 Ti has no all-reduce slack here: its fronts run beside an L2 issue, and there the mix_gram and pre_gram take 0.2 us more each with the copy (the front from mix_gram to the end of the q, k and v 66.37 -> 66.08 us, and 70.85 -> 70.34 without the issuer). pre_gram triggering its dependents past its mix sums, to launch the q, k and v where the quantize did, gave -0.25 % there and is not landed. PPL bit for bit: 386825.5220 at `-ub 1`, 385507.0946 at `-ub 3`, 387892.2627 under `-sm tensor`. Open, from review: no test fails if the copy stops being read (the harness allocates every tensor apart, so nothing aliases the mix); the node trace's quantize count is the check | `GGML_CUDA_MMVQ_Q8_1_PRODUCER_LEGACY=1` | +| `4bce9401c` | The hyper-connection front's second kernel (`dsv4_hc_pre_gram_f32`) ran the comb's Sinkhorn (20 row and column normalizations in one warp) before it could end, and the sublayer's first projections wait for it to end, though only the sublayer's `DSV4_HC_POST` reads the comb. It now leaves the comb's inputs in their slots of the weights, and `dsv4_hc_comb_side`, a warp a token on a stream of its own (`GGML_CUDA_HC_COMB_STREAM`, forked after the front), makes the comb from them with the same function on the same values, so the same bits, and writes it over them. The evaluation's stream waits for it before any node that reads those weights, in every `DSV4_HC_POST`, and at the evaluation's end; not while the graph runs concurrent streams. `test-backend-ops` `DSV4_HC_PRE_POST`, new: the front, then a post reading its weights at once, 5/5; without the waits the post races the comb and the 1,000-iteration case fails 5 runs of 5 (the 20-iteration ones by chance); without the side kernel `DSV4_HC_PRE_FUSED` fails 31 of 36 (the 5 that pass run unfused), `DSV4_HC_PRE_POST` 5 of 5 and `DSV4_HC_PRE_Q8_1` 16 of 16. The 44-layer GLM-5.3 proxy under `-sm tensor` on a 5080 + 5070 Ti, tg32 node traces over two rotated rounds: pre_gram ends 1.00 and 0.96 us sooner past mix_gram's end, and the next launch still starts 2.18 / 2.78 us before it ends (PDL kept across the fork); the side kernel takes 1.9 to 2.1 us beside the projections; each card's kernel time a token past the all-reduce -1.03 % and -0.69 % against the switch set, where the comb dropped outright (timing only, wrong bits) bounded it at -1.02 / -0.57 %. A graph replay runs the chain of a graph with a fork on more than one stream, so the node-trace tables count the chain as every replay stream's kernels but the side ones (`chain.py`). PPL bit for bit: 386825.5220 at `-ub 1`, 385507.0946 at `-ub 3`, 387892.2627 under `-sm tensor` | `GGML_CUDA_HC_COMB_SIDE_LEGACY=1` | +| `6588e1d6a` | `ggml_backend_compare_graph_backend` computed the backend under test, then the reference, then read both, so a kernel the backend left running past its return (a side stream's that the evaluation's end did not wait for) had finished by the read. The reference now runs first; a backend that finishes its work by the return gives the same results. `test-backend-ops` `DSV4_HC_PRE_FUSED` takes one case more, a 1,000-iteration comb with no `DSV4_HC_POST` in the graph, its front's weights an output, so only the evaluation's end waits for the comb beside the stream: without that wait the case fails 5 runs of 5 and 7 to 17 of the 20-iteration cases fail each run, where the old order passed all 37 | none (test harness) | +| `c6b4ecf1b` | `4bce9401c`'s comb beside the stream is written into the front's weights after the front's launch, and only a reader of the weights or the evaluation's end waits for it; weights read by the front's own `DSV4_HC_PRE` alone are free to the allocator once that node is placed, so a later node's output could lie over them and take the comb. The front now makes the comb beside the stream only where the weights outlive it: a use count past pre's in the whole graph (the post and comb views, which the meta backend's and the scheduler's subgraphs keep, so under `-sm tensor` the post past the all-reduce still counts), or an output; otherwise its second kernel makes the comb as before. No model builds a front without its post, so the proxy under `-sm tensor` launches 100 `dsv4_hc_comb_side` a token as before, PPL bit for bit (386825.5220, 385507.0946, 387892.2627). `test-backend-ops` cannot fail for the case this closes: it gives every tensor its own memory | `GGML_CUDA_HC_COMB_SIDE_LEGACY=1` | +| `5c03ed96f` | `ggml_cuda_check_fusion_memory_ranges` let the top-k fusion's outputs overlap its logits at one row only, and ggml-alloc places a layer's weights and ids over the logits whenever it can, so at an MTP verify's 3 tokens every MoE layer of the GLM-5.3 proxy ran the unfused routing chain, 8 kernels where a decode runs one. `topk_moe_cuda` now takes 4 rows a block, a warp each, and its warps meet at a barrier once every row's logits are read, before a weight or an id is written; `ggml_cuda_topk_moe_reads_before_writes` says so for the rows that fit one block, and the check lets those overlap. The 44-layer proxy at `-p 3 -ub 3`, nsys: the routing chain 8.32 -> 3.90 us a MoE layer on an RTX 5070 Ti, 8.32 -> 3.39 on an RTX 5080 and 8.19 -> 3.58 on the 5070 Ti under `-sm tensor`; an evaluation's kernel time past the all-reduce under `-sm tensor` -2.35 % (5080) and -1.78 % (5070 Ti) (the commit message's -1.73 / -1.55 % counted the all-reduce's second kernel as kernel time), its wall time -1.93 +- 0.49 % over 4 rotated rounds of 1,500. On the 5070 Ti alone the saving hid behind the L2 issue that ran into the routed gate/up (wall -0.24 +- 0.39 %); on `b3fc826ad`, which stops that issue, it shows: kernel time -0.65 %, wall +1.48 % (95 % CI +0.50 to +2.47, 4 rotated rounds, the switch against not). A decode is unchanged. PPL bit for bit at `-ub 1` (386825.5220) and under `-sm tensor` (387892.2627); at `-ub 3` 385634.0270 against 385507.0946, the fused kernel normalizing the weights with its own rounding as a decode's always did: KLD against the `-ub 1` logits 0.006763 +- 0.000061, the unfused chain's 0.006756 +- 0.000061. Open: `test-backend-ops` gives every tensor its own memory, so no case overlaps the outputs and the logits, and without the barrier PPL was the same | `GGML_CUDA_TOPK_MOE_ALIAS_LEGACY=1` | +| `b3fc826ad` | The paced L2 issuer (`345065068`) sizes an issue by the graph nodes of the chain beside it (`GGML_CUDA_L2_ISSUE_NODE_US` each), but where the evaluation fuses (a hyper-connection front, a top-k) many nodes are one launch, so the issue after the attention output ran on into the routed gate/up, which shared DRAM with it and ended where the issue let it. A ring launch (`mmvq_moe`) now bumps a word of the context's (`ggml_cuda_l2_issue_stop`, made before any capture as the tile counters are) as its reads start, and the issuer reads the word when it starts and every 4th piece after, and stops requesting once it has changed. The 44-layer GLM-5.3 proxy at `-p 3 -ub 3` on an RTX 5070 Ti alone, nsys over 4 rotated rounds, the switch set against not: the issue ends 9.3 us into the gate/up (p90 10.6) against 27.2, the gate/up 125.15 -> 108.26 us a KDA layer (107.94 with no issuer), an evaluation's kernel time -2.30 %; pp3 wall +2.67 % (95 % CI +0.34 to +5.01) over 4 rotated rounds of 1,500, and `5c03ed96f`'s top-k saving shows on one card. Declared and missed: the stop within 3 us of the ring's start; it takes 9.3, not explained yet (the ring still ends within 0.4 us of where it does with no issuer). Under `-sm tensor` no issue is running when a ring starts, with the switch or without, so the stop never fires: the 5080's kernel time +0.16 %, the 5070 Ti's rounds bimodal in both arms (+3.15 % pooled, +0.28 % between the lower modes). tg128 +0.31 % (-1.29 to +1.91) on one card, +0.67 % (-0.41 to +1.76) under `-sm tensor`. PPL bit for bit: 386825.5220 at `-ub 1`, 385634.0270 at `-ub 3`, 387892.2627 under `-sm tensor` | `GGML_CUDA_L2_ISSUE_STOP_LEGACY=1` | +| `3dbefa8fb` | At 3 tokens (an MTP verify's shape) the routed-expert ring was bound by its math, not by DRAM: each pair routed to an expert decoded every IQ3_XXS fragment of the expert's rows again (8 grid gathers from shared memory, 4 sign lookups and the sign ops) before its 8 dot products; the GLM-5.3 proxy's gate/up took 113-116 us at 51-53 % of DRAM with its memory pipes 72-74 % busy (ncu, RTX 5070 Ti). A launch of several tokens now runs its own instance (`mmvq_moe`, IQ3_XXS at rpw*nmat <= 2): the expert's pairs 4 at a time (`MMVQ_MOE_PB`), each fragment decoded once (`iq3_xxs_frag_decode`) and met with each pair's vector in turn, each pair's sums added in the same order, so the same bits; its own instance, as one kernel with both paths spilled. The one-token instances are unchanged: all 8,302 kernels of `b3fc826ad`'s libggml-cuda have the same SASS encodings, and the 3 added are the pairs_once ones. The 44-layer proxy at `-p 3 -ub 3`, nsys over 4 rotated rounds: on the 5070 Ti alone the gate/up 107.81 -> 75.65 us a KDA layer, the down 61.41 -> 49.66, an evaluation's kernel time -12.12 %; under `-sm tensor` -7.01 % (RTX 5080) and -6.52 % (5070 Ti). pp3 at `-ub 3`, 4 rotated rounds of 1,500: +9.68 % (95 % CI +8.58 to +10.79) on one card, +6.63 % (+4.33 to +8.93) under `-sm tensor`. Declared and missed: the down at 45 us or less; it stays 21 us over its DRAM floor (28.7 us), its tokens' vectors (55 KB at 3 tokens of 8 slots) read from global memory as they do not fit beside the ring. PPL bit for bit with the switch and without: 386825.5220, 385634.0270, 387892.2627. `test-backend-ops` gains `MUL_MAT_ID` at IQ3_XXS with 8 experts all used at 2, 3, 5 and 8 tokens, gate/up and down shapes, and `MOE_FFN_CHAIN` on them at 3 and 5 tokens; storing pair q's sums from pair q+1's fails all 16 cases of several tokens and passes the one-token ones | `GGML_CUDA_MMVQ_MOE_PAIRS_LEGACY=1` | +| `04dde5d64` | ncu of the ring at 3 tokens put the down at 48 % of DRAM with its loads hitting L1 99.5 % of the time, so not starved for its tokens' vectors as `3dbefa8fb`'s row supposed, no warp eligible 55 % of the cycles, and ~15 % of its instructions in integer divisions: each pair of each tile found its vector and its dst row from p / n_used, p % n_used and slot % nchannels_y, and each tile its expert, row tile, slot and phase from tile / ntr and i / nslots. The host now makes each pair's vector and dst offsets (`pair_y`, `pair_dst`: a constant load at an index the warp shares), and tile / ntr and i / nslots are `fast_div_modulo`; every address is the same, so the same bits. The down's several-token instance goes 3,296 -> 2,792 instructions, the gate/up's 3,920 -> 3,408, the one-token down's 2,072 -> 1,872. The 44-layer proxy at `-p 3 -ub 3`, nsys over 4 rotated rounds against `3dbefa8fb`: on the 5070 Ti alone the down 49.92 -> 40.51 us a KDA layer, the gate/up 75.01 -> 74.66 (declared: 74 or less, missed), an evaluation's kernel time -3.42 %; under `-sm tensor` both rings faster in every round on both cards (gate/up -4.2 % on the RTX 5080 and -4.8 % on the 5070 Ti, down -2.2 / -1.6 %), the 5080's kernel time -0.68 %, the 5070 Ti's +2.21 % pooled from rounds that jump in both arms; pp3 on one card +2.57 % (95 % CI +0.36 to +4.78, host load 12-22). A decode's ring launches read the same or less. On the 4-layer proxy with GLM-5.3-Flash's 288 experts (random routing, so an expert meets one pair nearly always; RTX 5080) both rings run at 89-92 % of DRAM, and the several-token instance is 0.2 % over the one-pair one under ncu, 1.2 / 2.1 % in the graph's charged time (declared: no gap, missed): it executes 11-12 % more instructions there, which DRAM mostly hides. PPL bit for bit with the pairs switch and without (386825.5220, 385634.0270, 387892.2627); adjacent pairs' dst offsets swapped fail 37 `MUL_MAT_ID` cases and `MOE_FFN_CHAIN` 10 of 10 | none (same addresses, same bits) | +| `18fa1398b` | At an MTP verify's 3 tokens glm5next's KDA gate projections ran on `mul_mat_f`, a block for each 32 rows of each dst channel, each reading a row's whole K: `ssm_f_a` and `ssm_g_a` (4096 -> 128 bf16) 8.6 us each on 4 blocks, `ssm_beta` (4096 -> 64) 6.4 us on 2 (RTX 5070 Ti). On NVIDIA from Ampere on, at 2 to 8 columns and rows x cols <= 1024, where `mul_mat_f` would launch fewer blocks than the device has SMs, `mul_mat_vec_f` takes the product, and `ggml_cuda_mul_mat_runs_mmvf` says so, so f_a and g_a run as one launch (`ggml_cuda_mul_mat_vec_f_pair`). The limit is measured: over f16/bf16 weights of 64-2048 rows, K 1024-8192 and 2-8 columns (`test-backend-ops perf`) the vector kernel took 0.21-0.68 of `mul_mat_f`'s time at rows x cols <= 1024 and 1.1-4.4x more past 2048 (a first build without the limit; `mul_mat_f`'s time there follows K, the vector kernel's rows x cols); with it all 74 such shapes take 0.18-0.68 on the 5070 Ti over 3 rounds and 0.23-0.63 on the 5080, the 118 others 0.92-1.10. The 44-layer proxy at `-p 3 -ub 3`, nsys over 4 rotated rounds: on the 5070 Ti alone a KDA layer's span from its conv to its recurrence 33.63 -> 18.18 us (declared: f_a + g_a + beta 23.65 -> 12 or less; they take ~7.3), an evaluation's kernel time -2.76 % against `2606ae528` (declared: -1.2 % or better), pp3 +4.26 % (95 % CI +1.51 to +7.02); under `-sm tensor` -5.71 % on the 5080 and -4.46 % on the 5070 Ti, the switch within 0.32 %. A decode launches the same 1,247 kernels. Not the same bits at several tokens: the vector kernel keeps the activations in f32 where `mul_mat_f` rounds them to bf16, and sums in another order, so PPL at `-ub 3` 385634.0270 -> 388534.0622 (declared: within 1e-4 relative, missed: that band did not know this random-weight proxy moves 0.31 % between `-ub 1` and `-ub 3`); against the switch's logits the mean KL divergence is 0.006795 ± 0.00006, where `-ub 1` against `-ub 3` (order alone) is 0.006763. Bit for bit at `-ub 1`, under `-sm tensor` and with the switch. `test-backend-ops` `MUL_MAT` (new cases on both sides of each limit, batched and broadcast), `MUL_MAT_PAIR` (f_a/g_a), `MUL_MAT_VEC_FUSION`, `SSM_CONV`, `SSM_CONV_STATE_UPDATE`, `GATED_DELTA_NET_CACHE_FUSION` and test-llama-archs pass on both cards; dst zeroed after the new launch fails 31 `MUL_MAT` cases, every one a shape the rule takes, and `MUL_MAT_PAIR` none (its pairs fuse) | `GGML_CUDA_MMVF_UNDERFILLED_LEGACY=1` | +| `2c5fb994c` | `3dbefa8fb`'s ring instance that decodes a fragment once for its expert's pairs pays for an expert that meets several pairs and costs its bookkeeping on one that meets one, and where 3 tokens route to 8 of GLM-5.3-Flash's 288 experts nearly every expert meets one: GLM-5.3-Flash on 2 RTX PRO 6000 under `-sm tensor` was 2.4 % and 1.4 % slower at `-p 3` and `-p 4` with it than with `GGML_CUDA_MMVQ_MOE_PAIRS_LEGACY=1` (rig-glm, `2e699cceb`, llama-bench, two rounds each), and the 4-layer 288-expert proxy's rings 1.2 / 2.1 % (`04dde5d64`'s row). The host now takes it only where the launch's pairs (ntokens x n_used) are as many as the experts or more (`ggml_cuda_mmvq_moe_args::n_experts`): not GLM-5.3-Flash at up to 8 tokens, still the 44-layer proxy (3 tokens on all 8 of 8). At `-p 3 -ub 3`, nsys over 2 rotated rounds on the RTX 5080: the 288-expert proxy launches the one-pair instance (the traces' demangled names; R3: `18fa1398b` launched the other), its gate/up 174.23 / 176.95 -> 171.67 / 175.54 us a layer (declared: -0.8 % or more, met), its down 84.00 / 85.71 -> 83.50 / 84.83 (declared: -1.5 % or more, missed; the switch's arm, the same instance, reads within 0.9 %); a 4-layer 16-expert proxy (~1.7 pairs an expert met, random routing) keeps the several-pair instance, which beats the one-pair one there: gate/up 107.17 / 109.71 against 108.58 / 113.03 us, down 51.36 / 51.71 against 55.79 / 57.22; the 44-layer proxy launches the same kernels and instances. PPL bit for bit (388534.0622 at `-ub 3`, 337649.7140 on the 288-expert proxy). `test-backend-ops` `MUL_MAT_ID` (new: 4 tokens on 32 experts, the rule's edge), `MUL_MAT_ID_FUSION`, `MOE_FFN_CHAIN` and test-llama-archs pass on both cards | `GGML_CUDA_MMVQ_MOE_PAIRS_LEGACY=1` (the one-pair instance at any counts) | +| `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) | +| `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) | — | +| `f8352b6c2` | `flash_attn_mask_to_sparse_indices` tested 2048 mask columns a round with 256 threads issuing scalar loads, 11.2 us a layer at 32K cached tokens on the proxy's decode. It now tests 32 consecutive columns a thread from four 16-byte loads issued together, 32768 columns a round, and places each thread's cells with a block-wide scan of the counts; upstream's scan stays as `flash_attn_mask_to_sparse_indices_legacy`, taken for a mask row that is not 16-byte aligned. The index lists are identical, so attention's result is bit for bit unchanged. `FLASH_ATTN_EXT` 3229/3229 on both cards each way, with two new cases at 40960 mask cells (a second scan round). PPL at `-c 8192 -b 8192 -ub 3 --chunks 1` identical by default, with the switch, and from the parent build. Three mutants fail the sparse cases and only those: a thread's offset taken inclusive of its own count; a round's count not carried into the next (fails only the 40960-cell cases); the earlier last-index drop. Scan kernel, median a launch over the last 31 decode tokens (NVTX, graphs off): 2.56 us (5080) and 3.26 us (5070 Ti) at `-d 8192`, 4.54 and 4.42 us at `-d 32768`, against 3.6 and 11.2 before; the `FLASH_ATTN_EXT` node at `-d 8192` 0.62x and 0.58x of the dense path's with all three kernels counted. Wall tg32 @ d32768 `-r 6`, reps 2-6 median, 2 interleaved runs: 89.92 and 88.36 against 84.22 and 84.92 tok/s (+6.8, +4.1 %). Declared `<= 3 us` at d8192 and `<= 4 us` at d32768: three of those four points miss by 8-13 %, recorded beside the declaration | `GGML_CUDA_FATTN_SPARSE_SCAN_LEGACY=1` (upstream's 2048-column scan) | +| `252f29c0a` | Two indexer changes, landed together because the second rewrites the kernel the first teaches to read rows. (a) `GGML_OP_LIGHTNING_INDEXER` takes an optional `src[4]`, I32 rows: key i of stream s is k's row `rows[i, s]` (`ggml_lightning_indexer_rows`). glm5next's fused path passes the f16 index cache's pooled head and `pool_reps` instead of materialising `get_rows`'s f32 copy of every pool's key — at 32K cached tokens an ~20 us gather a layer a token, 8258 blocks and 4.2 MB written, 11 times a token. CPU reads the rows too; Metal and SYCL refuse the variant; the meta backend mirrors it, every device reading the same keys through the same rows. (b) `lightning_indexer_kernel_quad` scores a key on a quad of lanes: lane s holds the dims lanes `4j + s` held, adds its 8 products in registers in `warp_reduce_sum`'s xor 16, 8, 4 order, then xor 2 and 1 across the quad — 2 shuffles a key and head where the vector kernel took 5 and served one key with them, all heads' q in shared memory at once, same grid, 64 keys a block, scores bit for bit the vector kernel's. `LIGHTNING_INDEXER` passes on both cards for every case by default, with either switch, and under the check with no trap; 26 new rows cases (f16/f32 keys, 32/64 heads, batch 1/3/64, 1/2 streams, a 77-key count off the kernel's block size). R3: the kernel ignoring the rows (row i for key i) fails the new cases; adding the in-register products in another order (xor 4's pairs first) traps under the check and passes `test-backend-ops` without it, a comparison against the CPU being blind to the summation order. PPL at `-c 8192 -b 8192`, ub 512 and ub 3 `--chunks 1`, two runs each, identical to the last digit by default and with either switch (349190.7774, 348976.2995). Quad against vector by ncu `gpu__time_duration.sum` at base clocks with all caches flushed, 20 launches a side on the 5070 Ti: 0.516x at 8258 pooled keys batch 1, 0.424x at batch 3, 0.536x at 65536 f32 keys, 0.0-0.9 % drift on a repeated arm; the same shapes at boost clocks with a warm cache 0.449x, 0.445x, 0.350x. Wall tg32 @ d32768, reps 2-6 median, 2 interleaved runs: reading pooled keys in place 89.28 and 88.86 against 86.80 and 87.77 tok/s (+2.86, +1.24 %, its `>= 1.0 %` band earned); the quad kernel 90.97 and 84.10 against 91.37 and 87.89, no gain, and not a readable result either way — two runs of the same arm span 8 % while `indexer_pool_score` is 196.8 us of a 21.7 ms token on the 5070 Ti (0.91 % of the wall), so halving it caps the gain near 0.5 %. Its declared `<= 0.4x` and `>= 0.8 %` bands both miss, the second unreachable by construction | `LLAMA_INDEXER_GATHER_LEGACY=1` (the f32 gather); `GGML_CUDA_LIGHTNING_INDEXER_VEC_LEGACY=1` (the vector kernel); `GGML_CUDA_LIGHTNING_INDEXER_CHECK=1` runs both kernels and traps on the first score whose bits differ | +| `c40b9d5ce` | Measured first (nsys at `-d 32768`, RTX 5070 Ti, main stream): a `quantize_q8_1` launch ends 3.14 us (median) after the gate/up ring it follows, 43 a token, 167 us of serial chain, and the ring's down projection then copies that q8_1 into shared memory anyway. Where a routed launch's vectors are each its own (`ne11 == n_used`, a down projection), no other MUL_MAT reads their q8_1, and the plan keeps them in shared memory, the host passes the f32 vectors and the ring's 512 consumer threads quantize them into that shared copy past the dependency wait: a thread a float4 (eight a round trip to L2), 8 lanes a q8_1 block, the max and sum in `warp_reduce_max/sum`'s own tree (xor 16, 8, 4 across the 8 lanes as xor 4, 2, 1, then xor 2 and 1 inside the lane's float4), then `d`, the rounding and `ds` as `quantize_q8_1` computes them — the launch's results bit for bit. `MUL_MAT_ID` 1076/1076 and `MOE_FFN_CHAIN` 10/10 on both cards each way and under the check with no trap. R3: a mutant adding the float4's elements as `(0 + 1) + (2 + 3)` traps under the check on the down cases and passes `test-backend-ops` without it. PPL under the check carries the legacy switch's value at `-c 1024 -ub 1 --chunks 1` and `-c 2048 -ub 3 --chunks 1` (the ring taking 1-8 tokens), and is identical to the last digit at `-c 8192 -b 8192` ub 512, 3 and 1 (349190.7774, 348976.2995, 321534.0329). Wall tg32 @ d32768, reps 2-6 median, 2 interleaved runs: 90.19 and 91.67 against 92.75 and 90.30 tok/s, the two runs disagreeing in sign (-2.76, +1.52 %); its `>= 0.6 %` band misses and was set at the instrument's own floor — 167 us is ~1 % of a token and two runs of one arm span more than that while another build shares the host. The node attribution it should be judged on (no `quantize_q8_1` after an `mmvq_moe`; the `ffn_down_exps` node at least 100 us lower on card 1) is recorded as owed | `GGML_CUDA_MMVQ_MOE_QUANTIZE_LEGACY=1` (the quantize launch); `GGML_CUDA_MMVQ_MOE_QUANTIZE_CHECK=1` runs the ring from both and traps on the first differing result | +| `6fbde0b8e` | The hyper-connection front (`dsv4_hc_front`'s Gram path) was two launches a sublayer, 88 a decode token on the 44-layer proxy: `dsv4_hc_mix_gram` (64 blocks, the dot products and Gram partials over a 64-column slice) and `dsv4_hc_pre_gram_f32` (16 blocks, each summing every partial, then the pre weights, the RMS, block 0's Sinkhorn, and its slice of the normed mix and its q8_1). On card 1 at `-d 32768` with graphs off they were 2.39 and 3.90 us a front, 210 and 343 us a token, while ncu read 0.23 and 0.05 waves with 93 % and 91 % of cycles holding no eligible warp — latency, not work. `dsv4_hc_front_one` runs mix_gram's blocks; each block's writers fence, thread 0 takes the token's ticket (`ggml_cuda_hc_front_tickets`, one unsigned int a stream and token, zeroed before a graph evaluation as the PQ2_0 tile counters are), and the block holding the last ticket sets it back to 0 and does pre_gram's work for the whole token through the same device functions, so the bits are the same; it reads the other blocks' partials streaming through L2 (`__ldcg`). `DSV4_HC_PRE_FUSED` 37/37, `DSV4_HC_PRE_POST` 5/5, `DSV4_HC_PRE_Q8_1` 16/16 and `DSV4_HC_POST` 4/4 on both cards each way and under the check with 0 check lines. R3, two mutants: the last ticket taken as `n_slices - 2` gives 320 check lines and 5 failures; the ticket never reset leaves the next launch on the stream with no last block, 544 check lines and 56 failures, collapsing the suites to 6/37, 1/5 and 1/16. PPL identical to the last digit each way at `-c 8192 -b 8192` ub 512, ub 3 and ub 1. Wall tg32 @ d32768, reps 2-6 median, 2 interleaved runs: 86.17 and 86.72 against 87.23 and 85.69 tok/s, disagreeing in sign (-1.22, +1.20 %); the `>= 0.8 %` band misses and was the wrong gate — 553 us of a token cut to 0.75x saves ~138 us, near 0.6 %, below what two runs resolve here. Its node attribution then measured the lever a REGRESSION: one launch 8.48 / 8.38 us a front against the two kernels' 2.11 + 4.22 and 2.15 + 3.01, so 754.5 / 759.4 us a token against 563.2 / 535.4, and — since a sum of two kernel times omits the gap between them, which is what fusing removes — the `DSV4_HC_POST` ranges' issue span, which does include it, 1475.4 against 1395.0 us a token on card 1. Both denominators agree. The launches a token do fall 1561.1 → 1473.3, exactly the 88 designed, and are not worth the ticket: `dsv4_hc_pre_gram_f32` summed every partial on 16 blocks in parallel, while the fused kernel's last-ticket block does all of that alone after fencing on 64 blocks — the two kernels' 93 % / 91 % of cycles with no eligible warp were a small grid, not the launch boundary. So the one launch was made **opt-in** right after it landed, as the comb beside the stream was at `98a6656ae`: the two kernels are the default and every verified property is kept (all four suites on both cards each way, 0 check lines, two R3 mutants, PPL identical at ub 512, 3 and 1). A next attempt should keep the 16-way reduction — the last block spreading pre_gram's work across its warps, or a reduction over more than one block — not one block doing all of it | `GGML_CUDA_HC_FRONT_ONE=1` (opt-in, the front in one launch); `GGML_CUDA_HC_FRONT_CHECK=1` implies it and runs the two kernels into scratch first, trapping on any differing bit of the normed mix, the weights or the q8_1 copy | +| `e11e67c29` | `4a88ac993` sent a sparse-mask filler slot to its own dump column through a liveness chain in every DSA layer (the pool bias gathered at the selected pools, an exp, a subtraction, a repeat, a copy of the top-k and their arithmetic: about ten launches a layer, 110 a token), and `build_indexer` reshaped `pool_cells` per layer, each view a split input of its own: the meta split took 30 inputs, the cap `GGML_SCHED_MAX_SPLIT_INPUTS`. The filler slots now come from the top-k itself: `build_inp_kpool` appends `select_k` dump pools to the pools the indexer scores, an F32 input of `-FLT_MAX` concatenated onto `pool_score`, so a dead pool (`-inf`) ranks below every dump pool and a dump pool below every live one; `pool_cells` gives dump pool c the cells from `n_kv + c` up, the sparse mask is widened by `kpool * select_k` columns of `-inf` (`cand_mask` stays `n_kv` wide), and `pool_cells` is reshaped once, not in every layer. `test-kpool-input` checks every dump cell is distinct and past `n_kv` and every widened column `-inf` (red with a dump cell repeated, and with the columns unwritten). On the 44-layer proxy, both RTX cards (5080 + 5070 Ti) under `-sm tensor`, against `12ef697e8`: `test-llama-archs` 0 FAIL, `test-backend-meta-sourceless` OK; one card, `-c 8192`, the logits at `-ub 3` and `-ub 1` are the parent's own floor (max KLD 2e-6, same top token 100 % and 99.951 %, the parent against itself the same), and at `-ub 512` part exactly as `4a88ac993`'s did from its parent (mean KLD 5.63e-4), the same fusion layout. `GGML_SCHED_DEBUG` at `-d 4096`: 63 split lines to 42, the widest split's inputs 30 to 15. `nsys` at `-d 32768`, graphs off, the last 31 decode tokens: 1561.1 launches a token to 1465.4 on each card, kernel time -153 / -92 us a token (card 0 / 1; `ffn_moe_gate`, which this cannot touch, -3.8 / -5.9), the window -823 us a token. Not all of it: the joint q/k `L2_NORM` that folds into KDA's `SSM_CONV` update only when the allocator put its output on none of the buffers that kernel reads (`ggml-cuda.cu`, the overlap checks after `Qwen3.5's joint q/k L2_NORM`) is refused in 22 of 33 layers a token here against 7 in the parent, +64 / +78 us, from where the changed graph lands its tensors. The launch and kernel-time cuts are the result; tg32 at `-d 32768`, graphs on, two interleaved pairs, medians of reps 2-6, is directional only: 83.27 / 87.61 and 79.18 / 81.40 t/s against the parent (+5.2 %, +2.8 %), the second pair's reps spanning 61-84 on both binaries, more than an admitted build moves a run (8.0 %). Measurement scripts: rig's `local/research/glm53-tp-rental-2026-09-28/allreduce-ll/fold-cuda.sh`, `node-attr.py` | — | +| `02f89b5c0` | Two of GLM-5.3-Flash's prefill norms, measured first on a pp4096 nsys of the 44-layer proxy under `-sm tensor -ub 1024` with each kernel billed its own time (its end less the later of its start and the end of the kernel before it on its stream: under PDL nsys's duration bills a kernel its wait on its predecessor). (a) KDA's per-head output norm (rows 128 wide, 32 a token) read 690 us a launch, 622.6 of them waiting on `gated_delta_net_cuda` and 67.8 its own. `rms_norm_f32_warp_rows` gives a row of at most 256 columns a warp and a block eight rows, each virtual warp's squares through the 256-thread block's butterfly and the eight partials through `block_reduce`'s second stage, so the outputs are the block kernel's bit for bit: own time 67.8 -> 24.5 us a launch (0.41 % of the prefill's device time); pp4096 +1.53 % in 4 interleaved pairs (95 % CI +0.04 to +3.03, t 3.182), which holds the 0.4 % the kernel time predicts and does not resolve it. (b) Each hc mix's weightless norm into `hc_fn` (BF16, on cuBLAS at prefill) wrote F32 that `convert_unary` cast to BF16 before a 21 us GEMM. Where `ggml_cuda_mul_mat_runs_cublas_bf16` (`ggml_cuda_mul_mat`'s own predicates) and `ggml_cuda_mul_mat_cublas_compute_type` say cuBLAS would cast it, the norm writes that BF16 copy itself (`rms_norm_f32`'s dst type, `ggml_cuda_cast`'s rounding) and cuBLAS reads it: norm and cast 291.6 / 325.9 -> 93.9 / 117.0 us a mix (5080 / 5070 Ti), the prefill's device time -5.5 %, pp4096 +4.40 % in 6 interleaved pairs (95 % CI +3.60 to +5.19, t 2.571). The proxy's KLD base file at `-c 2048 --chunks 8` (2,535,206,868 bytes: each scored token's scale and min log-prob as floats, its 16-bit log-probs over the vocab) is cmp-identical to the parent's with both on, with the warp rows off and with every switch set; R3: summing a lane's squares before one butterfly differs in 2,194,756,864 bytes, rounding the BF16 toward zero in 2,371,020,141. `test-backend-ops` RMS_NORM 57/57, RMS_NORM_MUL 6/6 (new: the mul-only fusion KDA's norm takes), RMS_NORM_MUL_ADD 36/36, RMS_NORM_MUL_MAT 27/27 (new: 64 columns, on cuBLAS) on both cards | `GGML_CUDA_RMS_NORM_WARP_ROWS_LEGACY=1` (the 256-thread block); `GGML_CUDA_RMS_NORM_BF16_LEGACY=1` (the F32 norm and the cast) | +| `58ef198aa` | KDA's recurrent kernel gives each state column a warp, and each warp of a head activated and exponentiated all S_v of a token's per-channel decays (`exp(raw_lb * sigmoid(-(g * raw_a[h])))`, two expf and a division an element): every value S_v times over. From 32 tokens a sequence `gdn_kda_precompute_decay` writes each once and the kernel reads it (`G_PRECOMPUTED`, with KDA leaving RAW to activate beta alone), with the kernel's own formulas, so the values are bit for bit; under 32 (decode, a verify) nothing changes, 32 being the GB10 scalar-gate precompute's threshold, not measured here. The 44-layer proxy, pp4096 under `-sm tensor -ub 1024`, own time: `gated_delta_net_cuda` 1351.2 / 1387.3 -> 1216.8 / 1274.3 us a launch plus 16.4 / 21.2 us of the precompute (5080 / 5070 Ti), net -8.7 / -6.6 %, the prefill's device time -1.0 %: the gate's instructions were not what bounds the kernel. pp4096 +1.29 % in 6 interleaved pairs (95 % CI -0.42 to +3.00, t 2.571), unresolved. The KLD base file is cmp-identical with this and `02f89b5c0` on; R3: every decay one ulp toward zero differs in 2,270,055,044 bytes (`__expf` for `expf` does not: ggml-cuda builds with `-use_fast_math`). GATED_DELTA_NET 72/72 (new: KDA at 64 tokens with snapshot slots, two sequences and GQA, at 32 with the rows-indexed state) and GATED_DELTA_NET_CACHE_FUSION 77/77 on both cards, by default and with the switch | `GGML_CUDA_KDA_DECAY_PRECOMPUTE_LEGACY=1` (the kernel computes them) | +| `7755e858a` | KDA's prefill (12.6 % of a pp4096's device time on the proxy) takes the chunked pipeline the scalar gate had to itself, in FLA's chunk_kda form, G being each key channel's log-decay summed along a 16-token chunk so every exponent is <= 0. Stage 1 (`cgdr_kda_fwdsub_intra_kernel`) activates the gate, builds G, and in one fp32 pass sharing each `exp(G[t] - G[s])` writes both the coupling L and the masked Q@K^T (stage 2 does not run), then solves for k_cumdecay and v_corr as for the scalar gate; stage 3 (`cgdr_state_wmma_kernel<..., KDA>`) scales q by `exp(G[t])` and k by `exp(G_last - G[t])` channel by channel and decays state row k by `exp(G_last[k])`; with K > 1 the tail runs on the recurrent KDA kernel. Where it runs, an f16 or q8_0 KDA state cache keeps its cpy and a gathered s0 its GET_ROWS. The 44-layer proxy, pp4096 under `-sm tensor -ub 1024`, own time: KDA 1215.6 + 16.4 (the decay precompute) -> 173.5 + 433.7 us a layer (stage 1 + stage 3) on the 5080, 1271.4 + 20.2 -> 208.3 + 441.5 on the 5070 Ti, -50.7 / -49.7 %; device time 5202.9 -> 4819.3 ms (-7.4 %), the kept GET_ROWS giving back 1.1 / 16.4 ms. pp4096 +2.93 % in 6 interleaved pairs (95 % CI +1.60 to +4.26, t 2.571); pp16384 +3.81 % in 2 (CI -2.08 to +9.69), unresolved. Not bit for bit: KL against a recurrent base 0.006462 at `-c 2048` (same top p 79.86 %) and 0.006547 at `-c 16384` (positions 8192-16383, 78.98 %), no growth along the context, the recurrent kernel against itself 0.000000; on this chaotic proxy (PPL 344,555) a `-ub 512 -> 256` change alone costs 0.0058 and a reduction order 0.0026. The trained model's KL is the first measurement for a box with the real pack. GATED_DELTA_NET 79/79 and GATED_DELTA_NET_CACHE_FUSION 80/80 on both cards, by default and with the switch, the chunked KDA cases at most NMSE 1.26e-7 under the chunked bar of 2e-7. R3: decaying each state row by its neighbour channel's G_last fails all 7 chunked KDA cases (NMSE 0.75-1.72), once KDA's test gate (`init_kda_gate`) holds a rate per channel: drawn per token, as it was, no state reached a chunk boundary and the mutant passed a 32-head 512-token case at 9e-8 | `GGML_CUDA_KDA_CHUNKED_LEGACY=1` (KDA on the recurrent kernel); `GGML_CUDA_GDN_CHUNKED=0` turns the chunked path off for both gates | +| `67275eacf` | Chunked KDA's stage 3 was 70 % of its time: each of its 128 blocks loaded three 16x128 fp32 operands a chunk with eight scalar loads a thread, converted them to fp16 and applied two exps an element, and that per-block work bounds it (BV=16 doubles the blocks and costs +61 / +78 %). Stage 1 already holds q * scale, k, G and k_cumdecay in shared memory, so it writes stage 3's operands in their final fp16 form (k_cumdecay, q * scale * exp(G), k * exp(G_last - G)) and each state row's decay exp(G_last[k]): the same formulas on the same values. `cgdr_kda_state_wmma_kernel`, KDA's own stage 3, loads each as one 16-byte load a thread, issued for the next chunk once this chunk's copy is in shared memory; the scalar gate's `cgdr_state_wmma_kernel` is back to its text before `7755e858a`. The proxy, pp4096 under `-sm tensor -ub 1024`, own time: stage 3 433.7 / 441.5 -> 212.2 / 210.8 us a layer (5080 / 5070 Ti), KDA 607.2 / 649.8 -> 379.9 / 411.8 (against the recurrent kernel's 1232.0 / 1291.6), device time 4819.3 -> 4709.4 ms (-2.3 %). pp4096 against `7755e858a`'s library (its chunk object relinked, the same executables) +1.52 % in 10 interleaved pairs (95 % CI +0.55 to +2.50, t 2.262). The KLD base file is cmp-identical to `7755e858a`'s chunked output with the final artifact; R3: q's operand one ulp toward zero differs in 2,229,794,626 bytes. GATED_DELTA_NET 79/79 and GATED_DELTA_NET_CACHE_FUSION 80/80 on both cards, by default and with the switch | `GGML_CUDA_KDA_CHUNKED_LEGACY=1` (KDA on the recurrent kernel) | Every switch in the last column is read once per process and parses as an integer: a `*_LEGACY` switch set to `0` is the same as unset (the change stays on), and `=0` turns off `GGML_CUDA_LORA_RANK1_FUSE` and diff --git a/ggml/include/ggml.h b/ggml/include/ggml.h index abefb6c1d9b2..2a93a3747e7d 100644 --- a/ggml/include/ggml.h +++ b/ggml/include/ggml.h @@ -2467,6 +2467,15 @@ extern "C" { GGML_API bool ggml_flash_attn_ext_get_mask_prefix( 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. + GGML_API void ggml_flash_attn_ext_set_n_kv_max( + struct ggml_tensor * a, + int32_t n_kv_max); + + GGML_API int32_t ggml_flash_attn_ext_get_n_kv_max( + const struct ggml_tensor * a); + // TODO: needs to be adapted to ggml_flash_attn_ext GGML_API struct ggml_tensor * ggml_flash_attn_back( struct ggml_context * ctx, @@ -2666,6 +2675,19 @@ extern "C" { struct ggml_tensor * weights, struct ggml_tensor * mask); + // the same with key i of stream s at row k_rows[i, s] of k (a view into a cache, no gathered copy), and the + // scores in f32 whatever k's type (as over the f32 rows ggml_get_rows would have copied): + // k: [n_embd_idx, 1, n_rows, ne3] + // k_rows: [n_kv, ne3] I32, each in [0, n_rows) + // mask, res: n_kv as above + GGML_API struct ggml_tensor * ggml_lightning_indexer_rows( + struct ggml_context * ctx, + struct ggml_tensor * q, + struct ggml_tensor * k, + struct ggml_tensor * k_rows, + struct ggml_tensor * weights, + struct ggml_tensor * mask); + // DeepSeek V4 hyper-connections (ref. https://arxiv.org/pdf/2512.24880) // In short these operations are replacements for the original residual connection (x = transformer(x) + x) // using a richer representation through streams. diff --git a/ggml/src/ggml-backend-meta.cpp b/ggml/src/ggml-backend-meta.cpp index 1a91f1a41d63..5876204ebd73 100644 --- a/ggml/src/ggml-backend-meta.cpp +++ b/ggml/src/ggml-backend-meta.cpp @@ -579,7 +579,8 @@ static struct ggml_backend_meta_split_state ggml_backend_meta_get_split_state(co } } if (ret.axis == GGML_BACKEND_SPLIT_AXIS_NONE) { - ret = {GGML_BACKEND_SPLIT_AXIS_UNKNOWN, {0}, {1}, 1}; + // no sources (ARANGE): every device computes the same values + ret = {GGML_BACKEND_SPLIT_AXIS_MIRRORED, {0}, {1}, 1}; } if (scalar_only && ret.axis >= 0 && ret.axis < GGML_MAX_DIMS) { ret = {GGML_BACKEND_SPLIT_AXIS_UNKNOWN, {0}, {1}, 1}; @@ -851,6 +852,8 @@ static struct ggml_backend_meta_split_state ggml_backend_meta_get_split_state(co for (size_t i = 0; i < 4; i++) { GGML_ASSERT(src_ss[i].axis == GGML_BACKEND_SPLIT_AXIS_MIRRORED); } + // ggml_lightning_indexer_rows: every device reads the same keys through the same rows + GGML_ASSERT(tensor->src[4] == nullptr || src_ss[4].axis == GGML_BACKEND_SPLIT_AXIS_MIRRORED); return {GGML_BACKEND_SPLIT_AXIS_MIRRORED, {0}, {1}, 1}; }; diff --git a/ggml/src/ggml-backend.cpp b/ggml/src/ggml-backend.cpp index f4b2251b527a..3e076085af0f 100644 --- a/ggml/src/ggml-backend.cpp +++ b/ggml/src/ggml-backend.cpp @@ -2386,9 +2386,10 @@ bool ggml_backend_compare_graph_backend(ggml_backend_t backend1, ggml_backend_t if (num_test_nodes != 0) { GGML_ASSERT(test_nodes); - // Compute the whole graph and only test the output for specific tensors - ggml_backend_graph_compute(backend1, g1); + // Compute the whole graph and only test the output for specific tensors: backend1's last, so its outputs are + // read as soon as its evaluation returns, and work it leaves running past the return reads as a wrong result ggml_backend_graph_compute(backend2, g2); + ggml_backend_graph_compute(backend1, g1); bool verified = false; for (int i = 0; i < g1->n_nodes; i++) { diff --git a/ggml/src/ggml-cpu/ops.cpp b/ggml/src/ggml-cpu/ops.cpp index 9dfcefe92409..b540bd4834bf 100644 --- a/ggml/src/ggml-cpu/ops.cpp +++ b/ggml/src/ggml-cpu/ops.cpp @@ -12163,11 +12163,13 @@ void ggml_compute_forward_lightning_indexer( const ggml_tensor * k = dst->src[1]; const ggml_tensor * w = dst->src[2]; // weights const ggml_tensor * m = dst->src[3]; // mask + const ggml_tensor * r = dst->src[4]; // ggml_lightning_indexer_rows: key i of stream s at k's row r[i, s] GGML_ASSERT(dst->type == GGML_TYPE_F32); GGML_ASSERT( q->type == GGML_TYPE_F32); GGML_ASSERT( w->type == GGML_TYPE_F32); GGML_ASSERT( m->type == GGML_TYPE_F16); + GGML_ASSERT(r == nullptr || r->type == GGML_TYPE_I32); GGML_TENSOR_LOCALS(int64_t, neq, q, ne) GGML_TENSOR_LOCALS(size_t, nbq, q, nb) @@ -12190,7 +12192,7 @@ void ggml_compute_forward_lightning_indexer( const int n_head = q->ne[1]; const int n_tokens = q->ne[2]; const int n_stream = q->ne[3]; - const int n_kv = k->ne[2]; + const int n_kv = dst->ne[0]; ggml_to_float_t const k_to_float = ggml_get_type_traits(k->type)->to_float; GGML_ASSERT((k->type == GGML_TYPE_F32 || k_to_float) && "lightning indexer: unsupported K-type"); @@ -12215,7 +12217,9 @@ void ggml_compute_forward_lightning_indexer( const ggml_fp16_t * m_row = (ggml_fp16_t *) ((char *) m->data + t*nbm1 + (s%nem3)*nbm3); float * dst_row = (float *) ((char *) dst->data + t*nb1 + s*nb3 ); for (int ik = ir0; ik < ir1; ++ik) { - char * k_row = (char *) k->data + ik*nbk2 + s*nbk3; + const int64_t i_row = r ? ((const int32_t *) ((const char *) r->data + s*r->nb[1]))[ik] : ik; + GGML_ASSERT(i_row >= 0 && i_row < nek2); + char * k_row = (char *) k->data + i_row*nbk2 + s*nbk3; if (k_to_float) { k_to_float(k_row, k_row_f32, n_embd); } else { diff --git a/ggml/src/ggml-cuda/binbcast.cu b/ggml/src/ggml-cuda/binbcast.cu index 2e38077bf67f..84555ab0207d 100644 --- a/ggml/src/ggml-cuda/binbcast.cu +++ b/ggml/src/ggml-cuda/binbcast.cu @@ -542,6 +542,45 @@ void ggml_cuda_op_fused_mul(ggml_backend_cuda_context & ctx, ggml_tensor * dst, } } +// a thread an element of a token's sum, the slots in their order: each product and each sum rounded on its own, as the +// MUL and the ADDs round them (never contracted into an FMA) +static __global__ void k_moe_weighted_sum(const float * experts, const float * weights, float * dst, const int n_embd, + const int n_used, const int64_t se1, const int64_t se2, const int64_t sw1, const int64_t sw2, const int64_t sd1) { + ggml_cuda_pdl_lc(); + const int i = blockIdx.x*blockDim.x + threadIdx.x; + const int64_t t = blockIdx.y; + if (i >= n_embd) { + return; + } + ggml_cuda_pdl_sync(); + const float * e = experts + t*se2 + i; + const float * w = weights + t*sw2; + float acc = __fmul_rn(e[0], w[0]); + for (int s = 1; s < n_used; ++s) { + acc = __fadd_rn(acc, __fmul_rn(e[s*se1], w[s*sw1])); + } + dst[t*sd1 + i] = acc; +} + +void ggml_cuda_op_moe_weighted_sum(ggml_backend_cuda_context & ctx, const ggml_tensor * experts, + const ggml_tensor * weights, ggml_tensor * dst) { + GGML_ASSERT(experts->type == GGML_TYPE_F32 && weights->type == GGML_TYPE_F32 && dst->type == GGML_TYPE_F32); + GGML_ASSERT(experts->nb[0] == sizeof(float) && dst->nb[0] == sizeof(float)); + const int n_embd = (int) experts->ne[0]; + const int n_used = (int) experts->ne[1]; + const int64_t n_tokens = experts->ne[2]; + GGML_ASSERT(weights->ne[0] == 1 && weights->ne[1] == n_used && weights->ne[2] == n_tokens); + GGML_ASSERT(dst->ne[0] == n_embd && dst->ne[1] == n_tokens); + + constexpr int block_size = 128; // k_bin_bcast's + const ggml_cuda_kernel_launch_params params = ggml_cuda_kernel_launch_params( + dim3((n_embd + block_size - 1) / block_size, (unsigned) n_tokens, 1), dim3(block_size, 1, 1), 0, ctx.stream()); + ggml_cuda_kernel_launch(k_moe_weighted_sum, params, (const float *) experts->data, (const float *) weights->data, + (float *) dst->data, n_embd, n_used, experts->nb[1] / (int64_t) sizeof(float), + experts->nb[2] / (int64_t) sizeof(float), weights->nb[1] / (int64_t) sizeof(float), + weights->nb[2] / (int64_t) sizeof(float), dst->nb[1] / (int64_t) sizeof(float)); +} + void ggml_cuda_op_repeat_back(ggml_backend_cuda_context & ctx, ggml_tensor * dst) { const ggml_tensor * src0 = dst->src[0]; diff --git a/ggml/src/ggml-cuda/binbcast.cuh b/ggml/src/ggml-cuda/binbcast.cuh index 12624785b444..8ec33c0c5fef 100644 --- a/ggml/src/ggml-cuda/binbcast.cuh +++ b/ggml/src/ggml-cuda/binbcast.cuh @@ -10,3 +10,10 @@ void ggml_cuda_op_repeat_back(ggml_backend_cuda_context & ctx, ggml_tensor * dst void ggml_cuda_op_fused_add(ggml_backend_cuda_context & ctx, ggml_tensor * dst, int n_fuse); void ggml_cuda_op_fused_mul(ggml_backend_cuda_context & ctx, ggml_tensor * dst, int n_fuse); + +// Routed experts' weighted sum as build_moe_ffn writes it, MUL(experts, weights) then the slots' views added in order +// (ADD(ADD(v0, v1), v2), ...), in one launch: dst[i, t] = ((e[i,0,t]*w[0,t] + e[i,1,t]*w[1,t]) + ...), every product and +// sum rounded as those nodes round them, so bit for bit their result. experts [n_embd, n_used, n_tokens], weights +// [1, n_used, n_tokens], dst [n_embd, n_tokens], all F32 with contiguous rows; dst lies over neither input. +void ggml_cuda_op_moe_weighted_sum(ggml_backend_cuda_context & ctx, const ggml_tensor * experts, + const ggml_tensor * weights, ggml_tensor * dst); diff --git a/ggml/src/ggml-cuda/chunk_gated_delta_net.cu b/ggml/src/ggml-cuda/chunk_gated_delta_net.cu index cc41dfcc5526..f3a8fc70afe1 100644 --- a/ggml/src/ggml-cuda/chunk_gated_delta_net.cu +++ b/ggml/src/ggml-cuda/chunk_gated_delta_net.cu @@ -11,6 +11,8 @@ // - the chunk pipeline may run on a leading token range and write its final state to any slot, // so the recurrent kernel can produce the rollback snapshots of the last tokens (see // ggml_cuda_op_gated_delta_net_impl). +// - KDA's per-channel gate (FLA's chunk_kda form): cgdr_kda_fwdsub_intra_kernel, which also writes the masked Q@K^T +// and stage 3's operands, and cgdr_kda_state_wmma_kernel; stage 2 does not run. // #include "chunk_gated_delta_net.cuh" @@ -244,6 +246,195 @@ __launch_bounds__(128, 4) __global__ void cgdr_fwdsub_intra_kernel( } } +// KDA's intra-chunk pass, the gate per channel (FLA's chunk_kda form). G[t][i] is channel i's log-decay summed along +// the chunk through token t, falling along it, so every exponent below is <= 0 and nothing overflows: +// L[t][s] = beta[t] * sum_i k[t][i] k[s][i] exp(G[t][i] - G[s][i]) (s < t) +// A[t][s] = sum_i (scale q[t][i]) k[s][i] exp(G[t][i] - G[s][i]) (s <= t), GDN's stage-2 masked Q@K^T +// in one pass, fp32, sharing each exponential; then (I + L) x = b solved for b = beta * exp(G) * k -> k_cumdecay and +// b = beta * v -> v_corr as for GDN. A goes to qk; k_cumdecay, q * exp(G) and k * exp(G_last - G) go to stage 3 as +// the fp16 operands it multiplies (cgdr_kda_state_wmma_kernel), with each state row's decay exp(G_last). RAW: beta and +// g pre-activation (ggml_gated_delta_net_set_raw_kda_gates), activated with the recurrent kernel's formulas. +// Grid (B*H, num_chunks); 128 threads, thread c channel c. +template +__launch_bounds__(128, 4) __global__ void cgdr_kda_fwdsub_intra_kernel( + const float * __restrict__ q_in, + const float * __restrict__ k_in, + const float * __restrict__ v_in, + const float * __restrict__ beta, + const float * __restrict__ g_in, + const float * __restrict__ raw_a, // [H], RAW only + const float raw_lb, // RAW only + const float scale, + float * __restrict__ v_corr, // (B, H, C, CS, V) output + __half * __restrict__ kcd_out, // (B, H, C, CS, K) output, k_cumdecay + __half * __restrict__ qg_out, // (B, H, C, CS, K) output, q * scale * exp(G) + __half * __restrict__ kg_out, // (B, H, C, CS, K) output, k * exp(G_last - G) + float * __restrict__ decay_out, // (B, H, C, K) output, exp(G_last) + float * __restrict__ qk_out, // (B, H, C, CS, CS) output + const int seq_len, // tokens this pass covers + const int H, + const int num_chunks, + const int v_dim, + const int num_k_heads, // q/k head count (H is the v-head count; GQA when smaller) + const long long sk1, const long long sk2, const long long sk3, // q/k strides: head, token, seq + const long long sv1, const long long sv2, const long long sv3, // v strides + const long long sb1, const long long sb2, const long long sb3) // beta strides; g's are BK times these +{ +#if defined(GDN_TC_AVAILABLE) + static_assert(BK == 128, "cgdr_kda_fwdsub_intra_kernel: one thread a channel, BK=128 only"); + + // SMEM: s_k, s_q (scaled), s_g (the log-decay, then G) [CS][BK+1] fp32, s_l[CS][CS] fp32, s_beta[CS]. + constexpr int sk = BK + 1; + extern __shared__ float smem[]; + float * s_k = smem; + float * s_q = s_k + CS * sk; + float * s_g = s_q + CS * sk; + float * s_l = s_g + CS * sk; + float * s_beta = s_l + CS * CS; + + const int tid = threadIdx.x; + const int pid_bh = blockIdx.x; + const int pid_chunk = blockIdx.y; + const int b = pid_bh / H; + const int h = pid_bh % H; // v-head + const int h_k = h % num_k_heads; // GQA: v-head -> shared k-head + const int t_off = pid_chunk * CS; + + const float * q_chunk = q_in + b * sk3 + t_off * sk2 + h_k * sk1; + const float * k_chunk = k_in + b * sk3 + t_off * sk2 + h_k * sk1; + const float * v_chunk = v_in + b * sv3 + t_off * sv2 + h * sv1; + const float * beta_chunk = beta + b * sb3 + t_off * sb2 + h * sb1; + const float * g_chunk = g_in + (b * sb3 + t_off * sb2 + h * sb1) * BK; + + const long long chunk_id = (long long) pid_bh * num_chunks + pid_chunk; + const int valid_cs = min(CS, seq_len - t_off); + + // 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 + 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; + if (t < valid_cs) { + qv = q_chunk[t * sk2 + c] * scale; + kv = k_chunk[t * sk2 + c]; + gv = g_chunk[t * sb2 * BK + c]; + if constexpr (RAW) { + gv = raw_lb * (1.0f / (1.0f + expf(gv * raw_a[h]))); + } + } + s_q[t * sk + c] = qv; + s_k[t * sk + c] = kv; + s_g[t * sk + c] = gv; + } + for (int i = tid; i < CS; i += 128) { + float beta_val = 0.f; + if (i < valid_cs) { + beta_val = beta_chunk[i * sb2]; + if constexpr (RAW) { + beta_val = 1.0f / (1.0f + expf(-beta_val)); + } + } + s_beta[i] = beta_val; + } + __syncthreads(); + + // Step 1: G, each channel's prefix sum along the chunk + { + float acc = 0.f; + for (int t = 0; t < CS; t++) { + acc += s_g[t * sk + tid]; + s_g[t * sk + tid] = acc; + } + } + __syncthreads(); + + // Step 2: L and A (exact FP32 scalar dot products), one exp(G[t][i] - G[s][i]) <= 1 for both + float * qk_chunk = qk_out + chunk_id * CS * CS; + for (int idx = tid; idx < CS * CS; idx += 128) { + const int t = idx / CS, s = idx % CS; + float l = 0.f, a = 0.f; + if (s <= t) { + for (int i = 0; i < BK; i++) { + const float kd = s_k[s * sk + i] * __expf(s_g[t * sk + i] - s_g[s * sk + i]); + l += s_k[t * sk + i] * kd; + a += s_q[t * sk + i] * kd; + } + } + s_l[idx] = s < t ? s_beta[t] * l : 0.f; + qk_chunk[idx] = a; + } + __syncthreads(); + + // Step 3: stage 3's q and k operands, [CS][BK] fp16 (q * scale * exp(G[t]) and k * exp(G_last - G[t]), both <= the + // unit-length vector's element), and each state row's decay over the chunk + const long long op_base = chunk_id * CS * BK; + { + const int c = tid; + const float g_last = s_g[(CS - 1) * sk + c]; + for (int t = 0; t < CS; t++) { + qg_out[op_base + t * BK + c] = cgdr_to_fp16(s_q[t * sk + c] * __expf(s_g[t * sk + c])); + kg_out[op_base + t * BK + c] = cgdr_to_fp16(s_k[t * sk + c] * __expf(g_last - s_g[t * sk + c])); + } + decay_out[chunk_id * BK + c] = __expf(g_last); + } + + // Step 4: k_cumdecay via forward substitution, b = beta * exp(G) * k + { + const int c = tid; + float xreg[CS]; + for (int t = 0; t < CS; t++) { + float xt = s_beta[t] * __expf(s_g[t * sk + c]) * s_k[t * sk + c]; + for (int s = 0; s < t; s++) { + xt -= s_l[t * CS + s] * xreg[s]; + } + xreg[t] = xt; + } + for (int t = 0; t < CS; t++) { + kcd_out[op_base + t * BK + c] = cgdr_to_fp16(xreg[t]); + } + } + __syncthreads(); + + // Step 5: v_corr via forward substitution (reuses s_k for v tile staging), as cgdr_fwdsub_intra_kernel's + const int num_vt = (v_dim + BK - 1) / BK; + for (int vt = 0; vt < num_vt; vt++) { + const int v_off = vt * BK; + const int v_cols = min(BK, v_dim - v_off); + + for (int i = tid; i < CS * BK; i += 128) { + const int t = i / BK, v = i % BK; + float val = 0.f; + if (t < valid_cs && v < v_cols) { + val = v_chunk[t * sv2 + v_off + v] * s_beta[t]; + } + s_k[t * sk + v] = val; + } + __syncthreads(); + + if (tid < v_cols) { + float xreg[CS]; + const int v_col = tid; + for (int t = 0; t < CS; t++) { + float xt = s_k[t * sk + v_col]; + for (int s = 0; s < t; s++) { + xt -= s_l[t * CS + s] * xreg[s]; + } + xreg[t] = xt; + } + for (int t = 0; t < CS; t++) { + v_corr[chunk_id * CS * v_dim + (long long) t * v_dim + (v_off + v_col)] = xreg[t]; + } + } + __syncthreads(); + } +#else + // Dispatched on Ampere+ only, with the tensor-core stage 3 its fp16 operands are for. + GGML_UNUSED_VARS(q_in, k_in, v_in, beta, g_in, raw_a, raw_lb, scale, v_corr, kcd_out, qg_out, kg_out, decay_out, + qk_out, seq_len, H, num_chunks, v_dim, num_k_heads, sk1, sk2, sk3, sv1, sv2, sv3, sb1, sb2, sb3); + NO_DEVICE_CODE; +#endif // defined(GDN_TC_AVAILABLE) +} + // Masked Q@K^T on tensor cores (fp16 WMMA, one warp per block): // qk_buf[i,j] = (Q_ch . K_ch[j]) * exp(g_cum[i] - g_cum[j]) for j <= i, else 0. // Grid (B*H, num_chunks); 32 threads. Requires CS==16, BK%16==0. @@ -561,6 +752,211 @@ __launch_bounds__(NT, OCC) __global__ void cgdr_state_wmma_kernel( #endif // defined(GDN_TC_AVAILABLE) } +// KDA's state update and fused output: cgdr_state_wmma_kernel's pass over the chunks, with the operands as +// cgdr_kda_fwdsub_intra_kernel leaves them -- k_cumdecay, q scale exp(G) and k exp(G_last - G), fp16 [CS][BK] a chunk, +// and each state row's decay exp(G_last[k]). A chunk's operand is one 16-byte load a thread, issued for the next chunk +// as soon as this chunk's copy is in shared memory: nothing but H carries from chunk to chunk, so the loads land while +// the chunk's GEMMs run. Grid (B*H, v_dim/BV); NT threads. +template +__launch_bounds__(NT, OCC) __global__ void cgdr_kda_state_wmma_kernel( + const float * __restrict__ v_corr, // (B*H, C, CS, V) + const __half * __restrict__ kcd, // (B*H, C, CS, K) + const __half * __restrict__ qg, // (B*H, C, CS, K) + const __half * __restrict__ kg, // (B*H, C, CS, K) + const float * __restrict__ decay, // (B*H, C, K) + const float * __restrict__ qk_buf, // (B*H, C, CS, CS) + float * __restrict__ output, // [B, T, H, v_dim] GGML layout -- direct output write + const float * __restrict__ init_state, + float * __restrict__ final_state, + const int num_chunks, + const int H, + const int v_dim, + const int seq_len, + const long long so2, const long long so3) { // output strides: token, seq +#if defined(GDN_TC_AVAILABLE) + static_assert(BK == 128 && CS == 16, "KDA state kernel requires BK=128, CS=16"); + static_assert(BV % 16 == 0, "BV must be a multiple of 16"); + static_assert(NT % 32 == 0, "NT must be a multiple of warp size"); + static_assert(CS * BK == 8 * NT, "an operand is one 16-byte load (8 halves) a thread"); + static_assert((CS * BV) % NT == 0 && (BK * BV) % NT == 0, "CS*BV and BK*BV must be divisible by NT"); + static_assert(CS * CS <= NT && BK <= NT, "qk and the decay are one element a thread"); + static_assert(NT / 32 >= BV / 16, "need at least BV/16 warps for WMMA n-tiles"); + + // SMEM: s_hfp16[BK*BV] fp16, s_op[CS*BK] fp16 (k_cumdecay, then q, then k, then qk), s_result[CS*BV] fp32 (also + // s_vnew fp16), s_decay[BK] fp32, s_hdelta[BK*BV] fp32. + constexpr int h_bytes = BK * BV * (int) sizeof(__half); + constexpr int op_bytes = CS * BK * (int) sizeof(__half); + constexpr int res_bytes = CS * BV * (int) sizeof(float); + constexpr int dec_bytes = BK * (int) sizeof(float); + + extern __shared__ char smem_kda_st[]; + __half * s_hfp16 = reinterpret_cast<__half *>(smem_kda_st); + __half * s_op = reinterpret_cast<__half *>(smem_kda_st + h_bytes); + float * s_result = reinterpret_cast(smem_kda_st + h_bytes + op_bytes); + float * s_decay = reinterpret_cast(smem_kda_st + h_bytes + op_bytes + res_bytes); + float * s_hdelta = reinterpret_cast(smem_kda_st + h_bytes + op_bytes + res_bytes + dec_bytes); + __half * s_vnew = reinterpret_cast<__half *>(s_result); // aliases s_result + + const int pid_bh = blockIdx.x; + const int v_off = blockIdx.y * BV; + const int tid = threadIdx.x; + const int warp_id = tid / 32; + constexpr int ept = (CS * BV) / NT; + constexpr int ept_h = (BK * BV) / NT; + constexpr int n_tiles = BV / 16; + + const long long bh_off = pid_bh; + const float * vcorr_bh = v_corr + bh_off * num_chunks * CS * v_dim; + const uint4 * kcd_bh = reinterpret_cast(kcd + bh_off * num_chunks * CS * BK); // NT a chunk + const uint4 * qg_bh = reinterpret_cast(qg + bh_off * num_chunks * CS * BK); + const uint4 * kg_bh = reinterpret_cast(kg + bh_off * num_chunks * CS * BK); + const float * decay_bh = decay + bh_off * num_chunks * BK; + const float * qk_bh = qk_buf + bh_off * num_chunks * CS * CS; + float * out_bh = output + (pid_bh / H) * so3 + (long long) (pid_bh % H) * v_dim + v_off; + + // FP32 H state in thread registers, as in cgdr_state_wmma_kernel, and its fp16 copy for the first chunk's WMMA. + float h_regs[ept_h]; + { + const long long src_base = bh_off * (long long) v_dim * BK; + for (int j = 0; j < ept_h; j++) { + const int idx = tid + j * NT; + h_regs[j] = init_state[src_base + (idx / BK + v_off) * BK + (idx % BK)]; + s_hfp16[idx] = cgdr_to_fp16(h_regs[j]); + } + } + + // Chunk 0's operands; each is reloaded for the next chunk once this chunk's copy is in shared memory. + uint4 r_kcd = kcd_bh[tid]; + uint4 r_qg = qg_bh[tid]; + uint4 r_kg = kg_bh[tid]; + float r_decay = tid < BK ? decay_bh[tid] : 0.f; + float r_qk = tid < CS * CS ? qk_bh[tid] : 0.f; + float r_vcorr[ept]; + for (int j = 0; j < ept; j++) { + const int idx = tid + j * NT; + r_vcorr[j] = vcorr_bh[(idx / BV) * v_dim + v_off + idx % BV]; + } + __syncthreads(); + + for (int ci = 0; ci < num_chunks; ci++) { + const bool next = ci + 1 < num_chunks; + const int nc = ci + 1; + const int valid_cs = min(CS, seq_len - ci * CS); // < CS on the last chunk if seq_len % CS != 0 + float vnew_regs[ept]; + float oi_regs[ept]; + + // k_cumdecay and the decay to SMEM, then the next chunk's + reinterpret_cast(s_op)[tid] = r_kcd; + if (tid < BK) { + s_decay[tid] = r_decay; + } + if (next) { + r_kcd = kcd_bh[nc * NT + tid]; + if (tid < BK) { + r_decay = decay_bh[nc * BK + tid]; + } + } + __syncthreads(); + + // s_result[CSxBV] = k_cumdecay @ H (v_new = u - w*h) + if (warp_id < n_tiles) { + cgdr_gemm_abt_16(s_op, s_hfp16 + warp_id * 16 * BK, s_result, BV, warp_id * 16); + } + __syncthreads(); + + // V_new = v_corr - s_result, kept in registers; then the next chunk's v_corr + for (int j = 0; j < ept; j++) { + const int idx = tid + j * NT; + vnew_regs[j] = r_vcorr[j] - s_result[idx]; + if (next) { + r_vcorr[j] = vcorr_bh[((long long) nc * CS + idx / BV) * v_dim + v_off + idx % BV]; + } + } + __syncthreads(); + + // q scale exp(G) to SMEM, then the next chunk's; O_inter = q @ H, kept in registers + reinterpret_cast(s_op)[tid] = r_qg; + if (next) { + r_qg = qg_bh[nc * NT + tid]; + } + __syncthreads(); + if (warp_id < n_tiles) { + cgdr_gemm_abt_16(s_op, s_hfp16 + warp_id * 16 * BK, s_result, BV, warp_id * 16); + } + __syncthreads(); + for (int j = 0; j < ept; j++) { + oi_regs[j] = s_result[tid + j * NT]; + } + __syncthreads(); + + // k exp(G_last - G) and V_new (fp16) to SMEM, then the next chunk's k; delta[v][k] = sum_t Vnew[t][v] * K[t][k] + reinterpret_cast(s_op)[tid] = r_kg; + if (next) { + r_kg = kg_bh[nc * NT + tid]; + } + for (int j = 0; j < ept; j++) { + s_vnew[tid + j * NT] = cgdr_to_fp16(vnew_regs[j]); + } + __syncthreads(); + if (warp_id < n_tiles) { + cgdr_gemm_ktv(s_vnew, s_op, s_hdelta, warp_id * 16); + } + __syncthreads(); + + // H = exp(G_last[k]) H + delta, row k by its own decay (fp32 accumulation preserved in h_regs) + for (int j = 0; j < ept_h; j++) { + h_regs[j] = s_decay[(tid + j * NT) % BK] * h_regs[j] + s_hdelta[tid + j * NT]; + } + __syncthreads(); // all reads of s_hdelta done before output WMMA overwrites it + + // Refresh s_hfp16 from fp32 h_regs for the next chunk's B-matrix. + for (int j = 0; j < ept_h; j++) { + s_hfp16[tid + j * NT] = cgdr_to_fp16(h_regs[j]); + } + + // output (fp16 WMMA): O[t][v] = O_inter + sum_t' qk[t][t'] * Vnew[t'][v] + if (tid < CS * CS) { + s_op[tid] = cgdr_to_fp16(r_qk); + if (next) { + r_qk = qk_bh[nc * CS * CS + tid]; + } + } + __syncthreads(); + if (warp_id < n_tiles) { + cgdr_gemm_qkv(s_op, s_vnew, s_hdelta, warp_id * 16); + } + __syncthreads(); + + // O = O_intra (s_hdelta[CS][BV] row-major) + O_inter. Padding tokens of the last chunk are not written. + float * out_chunk = out_bh + (long long) ci * CS * so2; +#pragma unroll + for (int j = 0; j < ept; j++) { + const int idx = tid + j * NT; + const int t_p = idx / BV; + const int v_p = idx % BV; + if (t_p < valid_cs) { + out_chunk[t_p * so2 + v_p] = s_hdelta[idx] + oi_regs[j]; + } + } + __syncthreads(); // ensure s_hfp16 refresh + output done before next chunk + } + + // Write final H state to GGML v-major [bh][v][k] (fp32 from h_regs, no conversion loss). + { + float * dst = final_state + bh_off * (long long) v_dim * BK; + for (int j = 0; j < ept_h; j++) { + const int idx = tid + j * NT; + dst[(idx / BK + v_off) * BK + (idx % BK)] = h_regs[j]; + } + } +#else + // Dispatched on Ampere+ only; body compiled out on older arches. + GGML_UNUSED_VARS(v_corr, kcd, qg, kg, decay, qk_buf, output, init_state, final_state, num_chunks, H, v_dim, + seq_len, so2, so3); + NO_DEVICE_CODE; +#endif // defined(GDN_TC_AVAILABLE) +} + // Dynamic SMEM bytes per kernel launch (<<<>>> third arg). static constexpr size_t cgdr_smem_fwdsub_intra(const int CS, const int BK) { return ((size_t) CS * (BK + 1) + (size_t) CS * CS + 2 * (size_t) CS) * sizeof(float); @@ -571,6 +967,10 @@ static constexpr size_t cgdr_smem_preqk_wmma(const int CS, const int BK) { return (size_t) 2 * CS * BK * sizeof(__half) + (size_t) (CS + CS * CS) * sizeof(float); } +static constexpr size_t cgdr_smem_kda_fwdsub_intra(const int CS, const int BK) { + return ((size_t) 3 * CS * (BK + 1) + (size_t) CS * CS + CS) * sizeof(float); +} + static constexpr size_t cgdr_smem_state_wmma(const int CS, const int BK, const int BV) { const size_t s_h = (size_t) BK * BV * sizeof(__half); const size_t s_kbuf = (size_t) CS * BK * sizeof(__half); @@ -580,6 +980,15 @@ static constexpr size_t cgdr_smem_state_wmma(const int CS, const int BK, const i return s_h + s_kbuf + s_res + s_gcum + s_hdelta; } +static constexpr size_t cgdr_smem_kda_state_wmma(const int CS, const int BK, const int BV) { + const size_t s_h = (size_t) BK * BV * sizeof(__half); + const size_t s_op = (size_t) CS * BK * sizeof(__half); + const size_t s_res = (size_t) CS * BV * sizeof(float); + const size_t s_decay = (size_t) BK * sizeof(float); + const size_t s_hdelta = (size_t) BK * BV * sizeof(float); + return s_h + s_op + s_res + s_decay + s_hdelta; +} + ggml_cuda_gdn_chunked_scratch ggml_cuda_gdn_get_chunked_scratch(const ggml_tensor * dst) { GGML_ASSERT(dst->op == GGML_OP_GATED_DELTA_NET); @@ -591,6 +1000,7 @@ ggml_cuda_gdn_chunked_scratch ggml_cuda_gdn_get_chunked_scratch(const ggml_tenso const int64_t t = src_v->ne[2]; const int64_t b = src_v->ne[3]; const int64_t k_dim = src_q->ne[0]; + const bool kda = dst->src[3]->ne[0] == v_dim; // KDA's gate, and so its G, per channel constexpr int64_t CS = 16; // ceil(T/CS); the last chunk may be partial and the kernels guard the padding tokens. Sized for @@ -604,19 +1014,28 @@ ggml_cuda_gdn_chunked_scratch ggml_cuda_gdn_get_chunked_scratch(const ggml_tenso // Every CUDA buffer is 128-aligned (ggml_backend_cuda_buffer_type_get_alignment), so each padded // offset is identical whether dst->data is the real pointer or the null it still is at // allocation time. That is what lets ggml_cuda_gdn_get_alloc_size call this before allocation. - auto carve = [&scratch](const int64_t n_floats) { - scratch.end = GGML_PAD(scratch.end, 128); - float * const p = (float *) scratch.end; - scratch.end += (uintptr_t) n_floats * sizeof(float); + auto carve_bytes = [&scratch](const int64_t n_bytes) { + scratch.end = GGML_PAD(scratch.end, 128); + const uintptr_t p = scratch.end; + scratch.end += (uintptr_t) n_bytes; return p; }; + auto carve = [&carve_bytes](const int64_t n_floats) { return (float *) carve_bytes(n_floats * sizeof(float)); }; + auto carve_half = [&carve_bytes](const int64_t n) { return (__half *) carve_bytes(n * sizeof(__half)); }; // Sized exactly: these are written only by scalar, exactly-bounded stores in the fwdsub/preqk // kernels (the WMMA store_matrix_sync writes target shared memory, not these). - scratch.v_corr = carve(bhcs * v_dim); - scratch.k_cumdecay = carve(bhcs * k_dim); - scratch.g_cum = carve(bhcs); - scratch.qk = carve(bhcs * CS); + scratch.v_corr = carve(bhcs * v_dim); + if (kda) { + scratch.kda_kcd = carve_half(bhcs * k_dim); + scratch.kda_qg = carve_half(bhcs * k_dim); + scratch.kda_kg = carve_half(bhcs * k_dim); + scratch.kda_decay = carve(bhcs / CS * k_dim); + } else { + scratch.k_cumdecay = carve(bhcs * k_dim); + scratch.g_cum = carve(bhcs); + } + scratch.qk = carve(bhcs * CS); return scratch; } @@ -648,11 +1067,21 @@ void ggml_cuda_gdn_chunked_launch(ggml_backend_cuda_context & ctx, const ggml_te // Scratch lives in the tail of dst's own allocation; see ggml_cuda_gdn_get_chunked_scratch. const ggml_cuda_gdn_chunked_scratch scratch = ggml_cuda_gdn_get_chunked_scratch(dst); - // Stage 1 -- intra pass: exact FP32 forward substitution -> v_corr, k_cumdecay, g_cum. - { + // Stage 1 -- intra pass: exact FP32 forward substitution -> v_corr, k_cumdecay, g_cum; KDA's also writes the + // masked Q@K^T, its decay per channel, and stage 3's operands in fp16. + const dim3 intra_grid(B * H, num_chunks, 1); + if (a.kda) { + constexpr size_t fs_smem = cgdr_smem_kda_fwdsub_intra(CS, BK); +#define CGDR_KDA_INTRA_LAUNCH(RAW_) \ + cgdr_kda_fwdsub_intra_kernel<<>>( \ + a.q, a.k, a.v, a.beta, a.g, RAW_ ? a.raw_a : nullptr, RAW_ ? a.raw_lb : 0.0f, a.scale, scratch.v_corr, \ + scratch.kda_kcd, scratch.kda_qg, scratch.kda_kg, scratch.kda_decay, scratch.qk, T, H, num_chunks, \ + (int) a.v_dim, (int) a.num_k_heads, a.sq1, a.sq2, a.sq3, a.sv1, a.sv2, a.sv3, a.sb1, a.sb2, a.sb3) + if (a.raw) { CGDR_KDA_INTRA_LAUNCH(true); } else { CGDR_KDA_INTRA_LAUNCH(false); } +#undef CGDR_KDA_INTRA_LAUNCH + } else { constexpr size_t fs_smem = cgdr_smem_fwdsub_intra(CS, BK); - const dim3 intra_grid(B * H, num_chunks, 1); - if (a.raw_dt_bias != nullptr) { + if (a.raw) { cgdr_fwdsub_intra_kernel<<>>( a.k, a.v, a.beta, a.g, a.raw_dt_bias, a.raw_a, scratch.v_corr, scratch.k_cumdecay, scratch.g_cum, T, H, num_chunks, (int) a.k_dim, (int) a.v_dim, (int) a.num_k_heads, @@ -666,19 +1095,28 @@ void ggml_cuda_gdn_chunked_launch(ggml_backend_cuda_context & ctx, const ggml_te } CUDA_CHECK(cudaGetLastError()); - // Stage 2 -- preqk pass: masked Q@K^T (one warp per block). - { + // Stage 2 -- preqk pass: masked Q@K^T (one warp per block); KDA's came out of stage 1. + if (!a.kda) { constexpr size_t qk_smem = cgdr_smem_preqk_wmma(CS, BK); const dim3 qk_grid(B * H, num_chunks, 1); cgdr_precompute_qk_wmma_kernel<<>>( a.q, a.k, scratch.g_cum, scratch.qk, num_chunks, a.scale, H, (int) a.num_k_heads, T, a.sq1, a.sq2, a.sq3); + CUDA_CHECK(cudaGetLastError()); } - CUDA_CHECK(cudaGetLastError()); - // Stage 3 -- state+output pass: WMMA tensor cores, fixed tile (BV=32/NT=256/OCC=4). ~30 KB - // dynamic SMEM, under the 48 KB default, so no cudaFuncAttribute opt-in needed. - { + // Stage 3 -- state+output pass: WMMA tensor cores, fixed tile (BV=32/NT=256). ~30 KB dynamic SMEM either way, + // under the 48 KB default, so no cudaFuncAttribute opt-in needed. BV=16 doubles KDA's blocks and costs it 61-78 % + // (the proxy's 5080 and 5070 Ti): each block's per-chunk operand traffic is fixed, whatever its v width. KDA's + // kernel holds its next chunk's operands in registers too, hence OCC=2 (128 blocks a card on the proxy's 84 SMs). + if (a.kda) { + constexpr int BV = 32, NT = 256, OCC = 2; + constexpr size_t st_smem = cgdr_smem_kda_state_wmma(CS, BK, BV); + const dim3 state_grid(B * H, (int) a.v_dim / BV, 1); + cgdr_kda_state_wmma_kernel<<>>( + scratch.v_corr, scratch.kda_kcd, scratch.kda_qg, scratch.kda_kg, scratch.kda_decay, scratch.qk, a.out, + a.state_in, a.state_out, num_chunks, H, (int) a.v_dim, T, a.so2, a.so3); + } else { constexpr int BV = 32, NT = 256, OCC = 4; constexpr size_t st_smem = cgdr_smem_state_wmma(CS, BK, BV); const dim3 state_grid(B * H, (int) a.v_dim / BV, 1); diff --git a/ggml/src/ggml-cuda/chunk_gated_delta_net.cuh b/ggml/src/ggml-cuda/chunk_gated_delta_net.cuh index 83b52d737680..c44c73d154b8 100644 --- a/ggml/src/ggml-cuda/chunk_gated_delta_net.cuh +++ b/ggml/src/ggml-cuda/chunk_gated_delta_net.cuh @@ -10,10 +10,13 @@ struct ggml_cuda_gdn_chunked_args { const float * q; const float * k; const float * v; - const float * g; + const float * g; // [1 or k_dim, H, tokens, seqs]: KDA's gate is per channel, at the beta strides times k_dim const float * beta; - const float * raw_dt_bias; // non-null: beta / g are pre-activation (ggml_gated_delta_net_set_raw_gates) - const float * raw_a; + bool kda; + bool raw; // beta / g pre-activation (ggml_gated_delta_net_set_raw_gates, with KDA _raw_kda_gates) + const float * raw_dt_bias; // raw, not KDA + const float * raw_a; // raw + float raw_lb; // raw KDA: g = raw_lb * sigmoid(-(g * raw_a[h])) const float * state_in; // [S, S, H, n_seqs], per (seq, head) [v][k] float * state_out; // final state after n_tokens tokens, same layout as state_in float * out; // attention output of token 0 @@ -37,9 +40,16 @@ void ggml_cuda_gdn_chunked_launch(ggml_backend_cuda_context & ctx, const ggml_te // replay without needing a separate persistent allocation to be pre-sized before capture. struct ggml_cuda_gdn_chunked_scratch { float * v_corr; + float * qk; + // the scalar gate: stage 1's fp32 k_cumdecay and g_cum, which stages 2 and 3 scale and convert float * k_cumdecay; float * g_cum; - float * qk; + // KDA: stage 3's operands as stage 1 leaves them, fp16 [CS][K] a chunk (k_cumdecay, q scale exp(G), + // k exp(G_last - G)), and each state row's decay over the chunk, exp(G_last) [K] + __half * kda_kcd; + __half * kda_qg; + __half * kda_kg; + float * kda_decay; uintptr_t end; }; diff --git a/ggml/src/ggml-cuda/common.cuh b/ggml/src/ggml-cuda/common.cuh index 342011a65d5d..f6f1cbd28a8c 100644 --- a/ggml/src/ggml-cuda/common.cuh +++ b/ggml/src/ggml-cuda/common.cuh @@ -140,6 +140,48 @@ static __device__ __forceinline__ void ggml_cuda_pdl_lc() { #endif // defined(GGML_CUDA_USE_PDL) && defined(__CUDA_ARCH__) && __CUDA_ARCH__ >= GGML_CUDA_CC_HOPPER } +// [p, p + nbytes), shrunk to the whole 16-byte units inside it, into L2 with nothing to wait on: a hint, which a kernel +// can issue before its PDL wait for data no kernel still running writes, so DRAM streams it under those kernels +// (cp.async.bulk.prefetch, sm_90 on; nothing below it). No result depends on it: L2 is where writes become coherent. +static __device__ __forceinline__ void ggml_cuda_prefetch_l2(const void * p, const int64_t nbytes) { +#if !defined(GGML_USE_HIP) && !defined(GGML_USE_MUSA) && defined(__CUDA_ARCH__) && __CUDA_ARCH__ >= GGML_CUDA_CC_HOPPER + const uintptr_t b = ((uintptr_t) p + 15) & ~(uintptr_t) 15; + const uintptr_t e = ((uintptr_t) p + nbytes) & ~(uintptr_t) 15; + if (b < e) { + asm volatile("cp.async.bulk.prefetch.L2.global [%0], %1;" :: "l"((uint64_t) b), "r"((uint32_t) (e - b)) : "memory"); + } +#else + GGML_UNUSED_VARS(p, nbytes); +#endif +} + +// A handoff between two kernels in flight together under PDL, one writing what the other reads before its dependency +// wait. In the PTX memory model a write reaches a thread of another grid only through a release that the reader's acquire +// reads from (or a barrier, or a fence.sc); griddepcontrol.wait makes a prerequisite grid's writes visible to the waiting +// grid alone, so a chain of waits proves nothing to a grid further on. The writer's threads end their stores with a block +// barrier and one of them adds with release semantics (cumulative over what the barrier ordered before it); the reader's +// thread loads with acquire semantics, then a block barrier hands the order on to its block. +static __device__ __forceinline__ void ggml_cuda_red_release_add(unsigned int * p, const unsigned int v) { +#if !defined(GGML_USE_HIP) && !defined(GGML_USE_MUSA) && defined(__CUDA_ARCH__) && __CUDA_ARCH__ >= GGML_CUDA_CC_VOLTA + asm volatile("red.release.gpu.global.add.u32 [%0], %1;" :: "l"(p), "r"(v) : "memory"); +#else + __threadfence(); + atomicAdd(p, v); +#endif +} + +static __device__ __forceinline__ unsigned int ggml_cuda_ld_acquire(const unsigned int * p) { +#if !defined(GGML_USE_HIP) && !defined(GGML_USE_MUSA) && defined(__CUDA_ARCH__) && __CUDA_ARCH__ >= GGML_CUDA_CC_VOLTA + unsigned int v; + asm volatile("ld.acquire.gpu.global.u32 %0, [%1];" : "=r"(v) : "l"(p) : "memory"); + return v; +#else + const unsigned int v = atomicAdd((unsigned int *) p, 0u); + __threadfence(); + return v; +#endif +} + #ifdef __CUDA_ARCH_LIST__ constexpr bool ggml_cuda_has_arch_impl(int) { return false; @@ -1511,15 +1553,41 @@ struct ggml_cuda_gated_delta_net_gather { #endif // defined(GGML_CUDA_USE_L2_WINDOW) }; -// The head of the next PQ2_0 launch's weights, which a launch prefetches into L2 once it has requested all of its own -// (mmvq-pq2-mma.cu): every block of a launch starts on the tiles at its matrices' heads, so while the small kernels -// between two matmuls run, DRAM streams what the next one reads first. x is the next launch's first matrix and gate the -// one it streams beside it, or null; bytes of each from its start, a multiple of 16 (0: none). A hint: no result depends +// What a PQ2_0 launch prefetches into L2 once it has requested all of its own weights (mmvq-pq2-mma.cu), as byte ranges +// its blocks share out laid end to end: first the weights the kernels between it and the next launch read (the Gated +// DeltaNet's alpha/beta matvec and conv, the norms, the rotations' signs: DRAM misses each token, each read on the chain +// between two matmuls), then the heads of the next launch's weights (its first matrix and, for a gated pair, the one it +// streams beside it: every block of a launch starts on the tiles at its matrices' heads), which DRAM streams while the +// kernels between the two run. Each range a multiple of 16 bytes from a 16-byte aligned start. A hint: no result depends // on it. The graph evaluation sets the context's before each node it dispatches (ggml_cuda_pq2_prefetch_plan). struct ggml_cuda_pq2_prefetch { - const void * x = nullptr; - const void * gate = nullptr; - int64_t bytes = 0; + static constexpr int max_ranges = 12; + + const void * ptr[max_ranges] = {}; + int64_t bytes[max_ranges] = {}; + int n = 0; + int64_t total = 0; // the ranges' bytes together + + // [p, p + nbytes) shrunk to whole 16-byte units inside it, merged into the last range when it continues it; false + // when the ranges are full + bool add(const void * p, int64_t nbytes) { + const uintptr_t b = ((uintptr_t) p + 15) & ~(uintptr_t) 15; + const uintptr_t e = ((uintptr_t) p + nbytes) & ~(uintptr_t) 15; + if (e <= b) { + return true; + } + if (n > 0 && (uintptr_t) ptr[n - 1] + bytes[n - 1] == b) { + bytes[n - 1] += e - b; + } else if (n < max_ranges) { + ptr[n] = (const void *) b; + bytes[n] = e - b; + ++n; + } else { + return false; + } + total += e - b; + return true; + } }; // One tile counter a stream for the ring launches (the PQ2_0 tensor-core ones, mmvq-pq2-mma.cu, and the routed experts', @@ -1545,12 +1613,60 @@ struct ggml_cuda_pq2_tile_counters { } }; +// The hyper-connection front's tickets in one launch (dsv4_hc_front_one, dsv4-hc.cu), one a stream and token: each +// block takes its token's after writing its partials, and the block with the last sets it back to 0 and finishes the +// token. Taken after the launch's dependency wait, like ggml_cuda_pq2_tile_counters, so one serves every launch of a +// stream, CUDA graph replays included. Made zeroed before a graph evaluation, never inside a capture. +struct ggml_cuda_hc_front_tickets { + static constexpr int n_tokens = 16; // a stream's tickets: the most tokens a front takes in one launch + + unsigned int * ptr = nullptr; // GGML_CUDA_MAX_STREAMS*n_tokens + + void ensure() { + if (ptr == nullptr) { + CUDA_CHECK(cudaMalloc(&ptr, GGML_CUDA_MAX_STREAMS * n_tokens * sizeof(unsigned int))); + CUDA_CHECK(cudaMemset(ptr, 0, GGML_CUDA_MAX_STREAMS * n_tokens * sizeof(unsigned int))); + } + } + + void release() { + if (ptr != nullptr) { + CUDA_CHECK(cudaFree(ptr)); + ptr = nullptr; + } + } +}; + +// The word a launch that keeps DRAM busy bumps as its reads start (the routed experts' ring, mmvq-moe.cu), and the paced +// L2 issuer beside the stream watches: past a change it requests no more (ggml_cuda_l2_issue), where it would take DRAM +// from that launch for bytes a later kernel reads. Only its changes mean anything. Made before a graph evaluation, never +// inside a capture, like ggml_cuda_pq2_tile_counters. +struct ggml_cuda_l2_issue_stop { + unsigned int * ptr = nullptr; + + void ensure() { + if (ptr == nullptr) { + CUDA_CHECK(cudaMalloc(&ptr, sizeof(unsigned int))); + CUDA_CHECK(cudaMemset(ptr, 0, sizeof(unsigned int))); + } + } + + void release() { + if (ptr != nullptr) { + CUDA_CHECK(cudaFree(ptr)); + ptr = nullptr; + } + } +}; + // The q8_1 copies of the F32 inputs that two or more quantized MUL_MAT / MUL_MAT_ID of one graph evaluation read on // mul_mat_vec_q (the q, k and v projections of one input; a shared expert's and the routed experts' gate/up): the first // reader quantizes into a pool allocation the evaluation holds (ggml_cuda_graph_evaluate_and_capture makes it before -// that node), and the others read the copy (ggml_cuda_mul_mat_vec_q). A copy is keyed by its input's first byte and rows, -// and its next reader quantizes it again once a node has written over any of those bytes (an in-place op, or ggml-alloc -// handing them to a later tensor). +// that node), and the others read the copy (ggml_cuda_mul_mat_vec_q). An input whose own kernel can write the copy (the +// hyper-connection front's normed mix, ggml_cuda_dsv4_hc_writes_q8_1) has it made before any node runs, for one reader +// or more, and the kernel writes it; the first reader quantizes it where it did not. A copy is keyed by its input's +// first byte and rows, and its next reader quantizes it again once a node has written over any of those bytes (an +// in-place op, or ggml-alloc handing them to a later tensor). struct ggml_cuda_mmvq_shared_q8_1 { // an input as nrows rows of ne0 floats, row r at data + r*row_stride in the order quantize_row_q8_1_cuda writes them // (r = i1 + ne1*(i2 + ne2*i3)) @@ -1569,14 +1685,20 @@ struct ggml_cuda_mmvq_shared_q8_1 { char * q8_1; size_t size; bool quantized; + // the index of the last producer to write the copy, or -1: a write by a node up to it, in its fused group or + // before it, precedes the copy on the stream and leaves it whole + int made_at = -1; }; std::vector entries; + // the nodes whose kernels write their outputs' copies (the key their copy's), by their index in the graph + std::unordered_map producers; // The MUL_MAT that leads the node group being run, or nullptr: only its input is read from a copy. The evaluation // learns of a group's writes after the group, and a MUL_MAT further in a group may read what an earlier node wrote. const ggml_tensor * head = nullptr; void reset() { entries.clear(); + producers.clear(); head = nullptr; } @@ -1611,11 +1733,10 @@ struct ggml_cuda_mmvq_shared_q8_1 { return k.nrows*GGML_PAD(k.ne0, MATRIX_ROW_PADDING)*sizeof(block_q8_1)/QK8_1; } - // the copy of t, when t is the input of head - entry * find(const ggml_tensor * t) { + // the copy of t, whoever makes it + entry * find_any(const ggml_tensor * t) { key k; - key head_k; - if (entries.empty() || head == nullptr || !key_of(t, k) || !key_of(head->src[1], head_k) || !(k == head_k)) { + if (entries.empty() || !key_of(t, k)) { return nullptr; } for (entry & e : entries) { @@ -1627,11 +1748,33 @@ struct ggml_cuda_mmvq_shared_q8_1 { return nullptr; } - // a node wrote [data, data + nbytes) - void written(const char * data, size_t nbytes) { + // the copy of t, when t is the input of head + entry * find(const ggml_tensor * t) { + key head_k; + if (head == nullptr || !key_of(head->src[1], head_k)) { + return nullptr; + } + entry * e = find_any(t); + return e != nullptr && e->k == head_k ? e : nullptr; + } + + // the copy of node's output that node's kernel writes, when node is a producer, marked made by it: a caller given + // one writes all of it on the stream + entry * produce(const ggml_tensor * node) { + const auto it = producers.find(node); + entry * e = it == producers.end() ? nullptr : find_any(node); + if (e != nullptr) { + e->quantized = true; + e->made_at = it->second; + } + return e; + } + + // the graph's node at wrote [data, data + nbytes) + void written(const char * data, size_t nbytes, const int at) { for (entry & e : entries) { const char * end = e.k.data + (e.k.nrows - 1)*e.k.row_stride + e.k.ne0*sizeof(float); - if (data < end && e.k.data < data + nbytes) { + if (data < end && e.k.data < data + nbytes && at > e.made_at) { e.quantized = false; } } @@ -1730,6 +1873,7 @@ struct ggml_cuda_fattn_kv_live_context { #define GGML_CUDA_SSM_CONV_UPDATE_D_CONV 4 // the kernel width it is built for (Qwen3-Next, Qwen3.5) #define GGML_CUDA_SSM_CONV_UPDATE_MAX_N_T 8 // new tokens per step it holds in registers #define GGML_CUDA_SSM_CONV_UPDATE_THREADS 128 // channels per block, and so the head width the L2 fold takes +#define GGML_CUDA_SSM_CONV_AB_MAX_ROWS 4 // alpha/beta rows a block of the fold computes (ggml_cuda_try_ssm_conv_ab) struct ggml_cuda_ssm_conv_state_update { const float * cache = nullptr; // conv cache rows, f32 @@ -1748,6 +1892,72 @@ struct ggml_cuda_ssm_conv_state_update { // the CONCATs): w_seg_channels rows of the SSM_CONV's src1 from each in turn; 0 channels, src1 itself const float * w_seg[3] = {}; int w_seg_channels = 0; + // Qwen3.5's alpha/beta pair, folded in (ggml_cuda_try_ssm_conv_ab): two bf16 matrices of ab_rows rows by ab_ncols + // on one activation, which the blocks compute before the dependency wait and write after it; 0 rows, not + const nv_bfloat16 * ab_w = nullptr; // the first matrix; the second is ab_s02 elements on + const float * ab_y = nullptr; // the activation, [ab_ncols, n_t] + float * ab_dst = nullptr; // the first output; the second is ab_s2 floats on + int64_t ab_s02 = 0; + int64_t ab_s2 = 0; + int64_t ab_stride_row = 0; // between a matrix's rows, in elements + int64_t ab_stride_y = 0; // between the activation's tokens, in floats + int64_t ab_stride_dst = 0; // between an output's tokens, in floats + int ab_rows = 0; + int ab_ncols = 0; + // the fold's handoff slot (ggml_cuda_ssm_conv_ab_slots) and the blocks of the activation's writer that release + // to it; set whenever that writer released to it, fold or not, since the conv's blocks are what set it back to 0 + unsigned int * ab_slot = nullptr; + int ab_writers = 0; +}; + +// Whether the conv-state update may fold in the alpha/beta pair (ggml_cuda_try_ssm_conv_ab), and so whether the PQ2_0 +// group launch writing the conv's inputs takes 152 registers, which leaves its SMs room for the conv's blocks, and +// triggers its dependents only after its own dependency wait, which keeps the fold's wait on its handoff slot from ever +// spinning beside blocks still to run (ssm_conv_ab_acquire): GGML_CUDA_SSM_CONV_AB_LEGACY=1 keeps the pair's own launch, +// and every group launch at 168 registers with the trigger at its start. +static bool ggml_cuda_ssm_conv_ab_enabled() { + static const bool legacy = ggml_env_switch("GGML_CUDA_SSM_CONV_AB_LEGACY"); + return !legacy; +} + +// The alpha/beta pair an SSM_CONV's update may fold in (ggml_cuda_try_ssm_conv_ab): its two MUL_MATs, the pair's +// activation, the conv's new inputs, the fold's handoff slot, the blocks of the kernel that wrote the activation and +// released to the slot this evaluation (0: none did), and whether the PQ2_0 group launch writing the new inputs came +// after that release, set up for the fold (ggml_cuda_ssm_conv_ab_enabled). The conv folds the pair in only with both. +struct ggml_cuda_ssm_conv_ab_plan { + const ggml_tensor * pair[2] = { nullptr, nullptr }; + const ggml_tensor * y = nullptr; + const ggml_tensor * x = nullptr; + int slot = -1; + int writers = 0; + bool fed = false; +}; + +// The alpha/beta folds' handoff slots, one unsigned int a fold of an evaluation, in the order the folds are planned. +// The kernel writing a fold's activation releases to its slot, a block at a time (the low 16 bits, rms_norm_fwht_cuda); +// every block of the conv acquires it until all the writer's blocks have, then takes a ticket (the high 16 bits), and +// the block with the grid's last ticket sets it back to 0 (ssm_conv_ab_acquire). A slot is released to and acquired once +// an evaluation, so every slot is 0 between two evaluations, and a reset reaches the slot's next use across the boundary +// between them, which PDL does not overlap (a graph launch; the copy and the host's wait for the logits). Made zeroed +// before a graph evaluation, never inside a capture, like ggml_cuda_pq2_tile_counters. +struct ggml_cuda_ssm_conv_ab_slots { + static constexpr int n = 256; + + unsigned int * ptr = nullptr; + + void ensure() { + if (ptr == nullptr) { + CUDA_CHECK(cudaMalloc(&ptr, n * sizeof(unsigned int))); + CUDA_CHECK(cudaMemset(ptr, 0, n * sizeof(unsigned int))); + } + } + + void release() { + if (ptr != nullptr) { + CUDA_CHECK(cudaFree(ptr)); + ptr = nullptr; + } + } }; // Registrations are keyed by node pointer, like ggml_cuda_gdn_gather_context, and cleared at the start of every graph @@ -1755,12 +1965,29 @@ struct ggml_cuda_ssm_conv_state_update { struct ggml_cuda_ssm_conv_update_context { std::unordered_map updates; // by SSM_CONV std::unordered_map l2_norms; // SSM_CONV -> its L2_NORM - std::unordered_set skipped; // the CONCATs, CPYs, L2_NORMs + std::unordered_map ab_plans; // by SSM_CONV + std::unordered_map ab_by_y; // a plan's activation -> its SSM_CONV + std::unordered_set skipped; // the CONCATs, CPYs, L2_NORMs, pairs + int ab_slots = 0; // handoff slots planned void reset() { updates.clear(); l2_norms.clear(); + ab_plans.clear(); + ab_by_y.clear(); skipped.clear(); + ab_slots = 0; + } + + const ggml_cuda_ssm_conv_ab_plan * ab_plan_of(const ggml_tensor * conv) const { + const auto it = ab_plans.find(conv); + return it == ab_plans.end() ? nullptr : &it->second; + } + + // the plan whose activation is y, for the kernel writing it to release to (rms_norm_fwht_cuda) + ggml_cuda_ssm_conv_ab_plan * ab_plan_of_y(const ggml_tensor * y) { + const auto it = ab_by_y.find(y); + return it == ab_by_y.end() ? nullptr : &ab_plans.at(it->second); } const ggml_cuda_ssm_conv_state_update * find(const ggml_tensor * conv) const { @@ -1946,9 +2173,16 @@ struct ggml_backend_cuda_context { ggml_cuda_pq2_prefetch pq2_next; // for the node being dispatched ggml_cuda_pq2_tile_counters pq2_tile_counters; ggml_cuda_mmvq_shared_q8_1 mmvq_shared_q8_1; // filled during a graph evaluation only + // each stream's last routed-expert ring launch's ids in this graph evaluation (ggml_cuda_mul_mat_vec_q): a later + // ring launch there reading the same ids reads them before its dependency wait (mmvq-moe.cu). A graph's tensor is + // written once, and ggml-alloc hands its bytes on only past its last reader. + const ggml_tensor * mmvq_moe_ids[GGML_CUDA_MAX_STREAMS] = {}; + ggml_cuda_ssm_conv_ab_slots ssm_conv_ab_slots; + ggml_cuda_hc_front_tickets hc_front_tickets; cudaEvent_t l2_issue_fork = nullptr; // the paced L2 issuer's fork from the evaluation's stream, and its join back cudaEvent_t l2_issue_join = nullptr; bool l2_issue_open = false; // an issuer was started and the stream has not waited for it + ggml_cuda_l2_issue_stop l2_issue_stop; bool l2_persisting = false; // a launch was given a persisting L2 window (ggml_cuda_l2_persist_window) const ggml_cgraph * l2_graph_next = nullptr; // the graph after the next computed, while it computes (ggml_backend_graph_next) // the stream is being captured with the rest of an evaluation across backends (ggml_backend_cuda_capture_begin): @@ -1960,6 +2194,11 @@ struct ggml_backend_cuda_context { // the current stream's PQ2_0 tile counter, or nullptr before the first graph evaluation made them int * pq2_tile_counter() { return pq2_tile_counters.ptr != nullptr ? pq2_tile_counters.ptr + curr_stream_no : nullptr; } + // the current stream's hyper-connection front tickets, or nullptr before the first graph evaluation made them + unsigned int * hc_front_ticket() { + return hc_front_tickets.ptr != nullptr ? hc_front_tickets.ptr + curr_stream_no*ggml_cuda_hc_front_tickets::n_tokens : nullptr; + } + cudaStream_t stream(int device, int stream) { if (streams[device][stream] == nullptr) { ggml_cuda_set_device(device); diff --git a/ggml/src/ggml-cuda/dsv4-hc.cu b/ggml/src/ggml-cuda/dsv4-hc.cu index 3593e2716687..980e8f39e983 100644 --- a/ggml/src/ggml-cuda/dsv4-hc.cu +++ b/ggml/src/ggml-cuda/dsv4-hc.cu @@ -579,10 +579,11 @@ struct dsv4_hc_post_args { int64_t sx_h; // x's stream stride (its token stride is the kernel's sx1) }; +// dsv4_hc_mix_gram's work, a block's: also the first part of dsv4_hc_front_one's template -static __global__ void __launch_bounds__(8*WARP_SIZE) dsv4_hc_mix_gram( +static __device__ __forceinline__ void dsv4_hc_mix_gram_block( float * x, const T * w, float * partials, const int64_t n_embd, - const int64_t sx1, const int64_t sw1, const bool w_prewait, const dsv4_hc_post_args pa) { + const int64_t sx1, const int64_t sw1, const bool w_prewait, const dsv4_hc_post_args & pa) { constexpr int rows = DSV4_HC_MIX/8; const int slice = blockIdx.x; @@ -675,56 +676,45 @@ static __global__ void __launch_bounds__(8*WARP_SIZE) dsv4_hc_mix_gram( } } +template +static __global__ void __launch_bounds__(8*WARP_SIZE) dsv4_hc_mix_gram( + float * x, const T * w, float * partials, const int64_t n_embd, + const int64_t sx1, const int64_t sw1, const bool w_prewait, const dsv4_hc_post_args pa) { + dsv4_hc_mix_gram_block(x, w, partials, n_embd, sx1, sw1, w_prewait, pa); +} + // a block for each DSV4_HC_PRE_GRAM_THR elements of a token's mix: every block sums the partials into the mixes and G // (from L2, a few KB), makes the pre weights, and the mix's RMS as sqrt(pre' G pre / n_embd + eps); then a thread its // element of the mix, normed, times the norm's weight, into dst. Block 0's warp 0 also makes the post and comb weights // (the comb in registers, dsv4_hc_comb_regs) into weights_out, while the other blocks write their slices. The streams' // RMS is G's trace's. base_prewait, norm_prewait: base or the norm's weight is the model's (dsv4_hc_prewait). -static __global__ void __launch_bounds__(DSV4_HC_PRE_GRAM_THR) dsv4_hc_pre_gram_f32( - const float * partials, const int n_slices, const float * x, - const float * scale, const float * base, const float * norm_w, - float * weights_out, float * dst, const int64_t n_embd, const int64_t k, - const int64_t sx1, const int64_t sx2, const int64_t ss0, const int64_t sb0, const int64_t sw0, - const int64_t sw1, const int64_t sd1, const float eps_flat, const float eps_hc, const int32_t n_iter, - const float eps_norm, const bool base_prewait, const bool norm_prewait) { - __shared__ float mix[DSV4_HC_GRAM_ROWS]; - __shared__ float base_s[DSV4_HC_MIX]; - - const int it = blockIdx.y; - const int warp = threadIdx.x / WARP_SIZE; - const int lane = threadIdx.x % WARP_SIZE; - const int64_t i = (int64_t) blockIdx.x*DSV4_HC_PRE_GRAM_THR + threadIdx.x; - - ggml_cuda_pdl_lc(); +// q8 non-null: dst's q8_1 copy as quantize_row_q8_1_cuda writes it (token it's row q8_s1 blocks on), a warp a block +// with quantize_q8_1's arithmetic on the values dst holds, so its bits: n_embd a multiple of MATRIX_ROW_PADDING. +// Its work is in the device functions below, which dsv4_hc_front_one runs too, so the two make the same bits. +struct dsv4_hc_pre_args { + const float * x; // the streams [n_embd, DSV4_HC, n_tokens], stream stride sx1, token stride sx2 + const float * scale; // stride ss0 + const float * base; // stride sb0 + const float * norm_w; + float * weights_out; // stride sw0, token stride sw1 + float * dst; // token stride sd1 + block_q8_1 * q8; // or nullptr; token stride q8_s1 blocks + int64_t n_embd, k, sx1, sx2, ss0, sb0, sw0, sw1, sd1, q8_s1; + float eps_flat, eps_hc, eps_norm; + int32_t n_iter; + bool base_prewait, norm_prewait; +}; - float nw = 0.0f; - if (base_prewait && threadIdx.x < DSV4_HC_MIX) { - base_s[threadIdx.x] = base[threadIdx.x*sb0]; - } - if (norm_prewait && i < n_embd) { - nw = norm_w[i]; - } - ggml_cuda_pdl_sync(); - if (!base_prewait && threadIdx.x < DSV4_HC_MIX) { - base_s[threadIdx.x] = base[threadIdx.x*sb0]; - } - if (!norm_prewait && i < n_embd) { - nw = norm_w[i]; - } - // the thread's streams, requested with the partials: one round trip for both - float xs[DSV4_HC] = {}; - if (i < n_embd) { - const float * xt = x + it*sx2; -#pragma unroll - for (int h = 0; h < DSV4_HC; ++h) { - xs[h] = xt[i + h*sx1]; - } - } +// a token's mixes and G from its partials pt, into mix (shared): warp w sums rows w, w + 8, ..., every row's loads issued +// before any row's shuffles, one L2 round trip, not one a row. cg: read past L1 (the launch's own blocks wrote them). The +// caller syncs the block before mix is read. +template +static __device__ __forceinline__ void dsv4_hc_gram_mix(const float * pt, const int n_slices, float * mix) { + const int warp = threadIdx.x / WARP_SIZE; + const int lane = threadIdx.x % WARP_SIZE; - // warp w sums rows w, w + 8, ...: every row's loads issued before any row's shuffles, one L2 round trip, not one a row constexpr int n_warps = DSV4_HC_PRE_GRAM_THR/WARP_SIZE; constexpr int rows_w = (DSV4_HC_GRAM_ROWS + n_warps - 1)/n_warps; - const float * pt = partials + (int64_t) it*DSV4_HC_GRAM_ROWS*n_slices; float s[rows_w]; #pragma unroll for (int q = 0; q < rows_w; ++q) { @@ -735,7 +725,7 @@ static __global__ void __launch_bounds__(DSV4_HC_PRE_GRAM_THR) dsv4_hc_pre_gram_ for (int q = 0; q < rows_w; ++q) { const int r = warp + q*n_warps; if (r < DSV4_HC_GRAM_ROWS) { - s[q] += pt[r*n_slices + sl]; + s[q] += cg ? __ldcg(pt + r*n_slices + sl) : pt[r*n_slices + sl]; } } } @@ -747,19 +737,32 @@ static __global__ void __launch_bounds__(DSV4_HC_PRE_GRAM_THR) dsv4_hc_pre_gram_ mix[r] = s[q]; } } - __syncthreads(); +} + +// element i's streams of token it (i < n_embd); cg: read past L1 (the launch's own blocks wrote them) +template +static __device__ __forceinline__ void dsv4_hc_gram_streams(const dsv4_hc_pre_args & pr, const int it, const int64_t i, + float * xs) { + const float * xt = pr.x + it*pr.sx2; +#pragma unroll + for (int h = 0; h < DSV4_HC; ++h) { + xs[h] = cg ? __ldcg(xt + i + h*pr.sx1) : xt[i + h*pr.sx1]; + } +} +// the pre weights and the mix's RMS from mix (dsv4_hc_gram_mix's) and base_s, and the streams' RMS (G's trace's) +static __device__ __forceinline__ void dsv4_hc_gram_pre(const float * mix, const float * base_s, const dsv4_hc_pre_args & pr, + float * pre, float & rms_flat, float & rms) { const float * G = mix + DSV4_HC_MIX; float trace = 0.0f; #pragma unroll for (int h = 0; h < DSV4_HC; ++h) { trace += G[dsv4_hc_gram_index(h, h)]; } - const float rms_flat = rsqrtf(trace/k + eps_flat); - float pre[DSV4_HC]; + rms_flat = rsqrtf(trace/pr.k + pr.eps_flat); #pragma unroll for (int h = 0; h < DSV4_HC; ++h) { - pre[h] = 1.0f/(1.0f + expf(-(mix[h]*rms_flat*scale[0] + base_s[h]))) + eps_hc; + pre[h] = 1.0f/(1.0f + expf(-(mix[h]*rms_flat*pr.scale[0] + base_s[h]))) + pr.eps_hc; } // pre' G pre: G symmetric, its off-diagonal entries twice float sumsq = 0.0f; @@ -772,46 +775,232 @@ static __global__ void __launch_bounds__(DSV4_HC_PRE_GRAM_THR) dsv4_hc_pre_gram_ } sumsq += pre[a]*row; } - const float rms = rsqrtf(sumsq/n_embd + eps_norm); + rms = rsqrtf(sumsq/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 +// L < 16 writing its element, lanes 16-19 a pre and a post weight +static __device__ __forceinline__ void dsv4_hc_gram_weights(const float * mix, const float * base_s, + const dsv4_hc_pre_args & pr, const int it, const float * pre, const float rms_flat) { + const int lane = threadIdx.x % WARP_SIZE; + + float m[DSV4_HC_MIX]; +#pragma unroll + for (int r = 0; r < DSV4_HC_MIX; ++r) { + m[r] = mix[r]*rms_flat; + } + float c[DSV4_HC*DSV4_HC]; + dsv4_hc_comb_regs(c, m, base_s, pr.scale[2*pr.ss0], pr.eps_hc, pr.n_iter); + float c_lane = c[0]; +#pragma unroll + for (int j = 1; j < DSV4_HC*DSV4_HC; ++j) { + c_lane = lane == j ? c[j] : c_lane; + } + const int h = lane - DSV4_HC*DSV4_HC; + float pre_h = pre[0]; + float m_h = m[DSV4_HC]; +#pragma unroll + for (int j = 1; j < DSV4_HC; ++j) { + pre_h = h == j ? pre[j] : pre_h; + m_h = h == j ? m[DSV4_HC + j] : m_h; + } + float * d = pr.weights_out + it*pr.sw1; + if (lane < DSV4_HC*DSV4_HC) { + d[(2*DSV4_HC + lane)*pr.sw0] = c_lane; + } else if (lane < DSV4_HC*DSV4_HC + DSV4_HC) { + d[h*pr.sw0] = pre_h; + d[(DSV4_HC + h)*pr.sw0] = 2.0f/(1.0f + expf(-(m_h*pr.scale[pr.ss0] + base_s[DSV4_HC + h]))); + } +} + +// token it's element i of the mix (i < n_embd) from its streams xs, normed, times the norm's weight nw, into dst, and its +// q8_1 copy when pr.q8: the warp's 32 values are a block (n_embd a multiple of 32: a warp is all in or all out) +static __device__ __forceinline__ void dsv4_hc_gram_out(const dsv4_hc_pre_args & pr, const int it, const int64_t i, + const float * xs, const float * pre, const float rms, const float nw) { + const int lane = threadIdx.x % WARP_SIZE; + + float v = xs[0]*pre[0]; +#pragma unroll + for (int h = 1; h < DSV4_HC; ++h) { + v += xs[h]*pre[h]; + } + const float xi = v*rms*nw; + pr.dst[it*pr.sd1 + i] = xi; + if (pr.q8 != nullptr) { + float amax = fabsf(xi); + float sum = xi; + amax = warp_reduce_max(amax); + sum = warp_reduce_sum(sum); + const float d = amax / 127.0f; + const int8_t q = amax == 0.0f ? 0 : roundf(xi / d); + block_q8_1 * blk = pr.q8 + it*pr.q8_s1 + i/QK8_1; + blk->qs[lane] = q; + if (lane == 0) { + blk->ds = make_half2(d, sum); + } + } +} + +static __global__ void __launch_bounds__(DSV4_HC_PRE_GRAM_THR) dsv4_hc_pre_gram_f32( + const float * partials, const int n_slices, const dsv4_hc_pre_args pr) { + __shared__ float mix[DSV4_HC_GRAM_ROWS]; + __shared__ float base_s[DSV4_HC_MIX]; + + const int it = blockIdx.y; + const int warp = threadIdx.x / WARP_SIZE; + const int64_t i = (int64_t) blockIdx.x*DSV4_HC_PRE_GRAM_THR + threadIdx.x; + + ggml_cuda_pdl_lc(); + + float nw = 0.0f; + if (pr.base_prewait && threadIdx.x < DSV4_HC_MIX) { + base_s[threadIdx.x] = pr.base[threadIdx.x*pr.sb0]; + } + if (pr.norm_prewait && i < pr.n_embd) { + nw = pr.norm_w[i]; + } + ggml_cuda_pdl_sync(); + if (!pr.base_prewait && threadIdx.x < DSV4_HC_MIX) { + base_s[threadIdx.x] = pr.base[threadIdx.x*pr.sb0]; + } + if (!pr.norm_prewait && i < pr.n_embd) { + nw = pr.norm_w[i]; + } + // the thread's streams, requested with the partials: one round trip for both + float xs[DSV4_HC] = {}; + if (i < pr.n_embd) { + dsv4_hc_gram_streams(pr, it, i, xs); + } + dsv4_hc_gram_mix(partials + (int64_t) it*DSV4_HC_GRAM_ROWS*n_slices, n_slices, mix); + __syncthreads(); + + float pre[DSV4_HC]; + float rms_flat; + float rms; + dsv4_hc_gram_pre(mix, base_s, pr, pre, rms_flat, rms); // block 0's warp 0 makes the weights first (the Sinkhorn is the longest chain), then its slice like every warp if (blockIdx.x == 0 && warp == 0) { - float m[DSV4_HC_MIX]; -#pragma unroll - for (int r = 0; r < DSV4_HC_MIX; ++r) { - m[r] = mix[r]*rms_flat; + dsv4_hc_gram_weights(mix, base_s, pr, it, pre, rms_flat); + } + if (i < pr.n_embd) { + dsv4_hc_gram_out(pr, it, i, xs, pre, rms, nw); + } +} + +// The front in one launch: dsv4_hc_mix_gram's blocks, after which the block that takes its token's last ticket (one a +// token at tickets, ggml_cuda_hc_front_tickets) sets it back to 0 and does dsv4_hc_pre_gram_f32's work for the whole +// token with its device functions, so its bits: the sums from every block's partials, the weights by warp 0, then a +// thread every DSV4_HC_PRE_GRAM_THR-th element, the streams of DSV4_HC_FRONT_ONE_CH of them requested at once. The +// partials and, under fuse_post, the streams are the launch's own blocks' writes, so they are read past L1. Every input +// dsv4_hc_pre_gram_f32 reads is read past the PDL wait mix_gram's part took. +static constexpr int DSV4_HC_FRONT_ONE_CH = 8; + +template +static __global__ void __launch_bounds__(8*WARP_SIZE) dsv4_hc_front_one( + float * x, const T * w, float * partials, const int64_t sx1, const int64_t sw1, const bool w_prewait, + const dsv4_hc_post_args pa, const dsv4_hc_pre_args pr, unsigned int * tickets) { + static_assert(8*WARP_SIZE == DSV4_HC_PRE_GRAM_THR, "the last block does a dsv4_hc_pre_gram_f32 block's work"); + __shared__ float mix[DSV4_HC_GRAM_ROWS]; + __shared__ float base_s[DSV4_HC_MIX]; + __shared__ bool last; + + dsv4_hc_mix_gram_block(x, w, partials, pr.n_embd, sx1, sw1, w_prewait, pa); + + // every writer's partials (and slice of the streams) seen device-wide before the block's ticket + const int it = blockIdx.y; + const int n_slices = gridDim.x; + __threadfence(); + __syncthreads(); + if (threadIdx.x == 0) { + const unsigned int t = atomicAdd(tickets + it, 1u); + last = t == (unsigned int) n_slices - 1; + if (last) { + tickets[it] = 0; } - float c[DSV4_HC*DSV4_HC]; - dsv4_hc_comb_regs(c, m, base_s, scale[2*ss0], eps_hc, n_iter); - float c_lane = c[0]; + } + __syncthreads(); + if (!last) { + return; + } + __threadfence(); + + const int warp = threadIdx.x / WARP_SIZE; + const int n_chunks = (int) ((pr.n_embd + DSV4_HC_PRE_GRAM_THR - 1) / DSV4_HC_PRE_GRAM_THR); + if (threadIdx.x < DSV4_HC_MIX) { + base_s[threadIdx.x] = pr.base[threadIdx.x*pr.sb0]; + } + float xs[DSV4_HC_FRONT_ONE_CH][DSV4_HC]; + float nw[DSV4_HC_FRONT_ONE_CH]; + auto request = [&](const int c0) { #pragma unroll - for (int j = 1; j < DSV4_HC*DSV4_HC; ++j) { - c_lane = lane == j ? c[j] : c_lane; + for (int c = 0; c < DSV4_HC_FRONT_ONE_CH; ++c) { + const int64_t i = (int64_t) (c0 + c)*DSV4_HC_PRE_GRAM_THR + threadIdx.x; + if (c0 + c < n_chunks && i < pr.n_embd) { + dsv4_hc_gram_streams(pr, it, i, xs[c]); + nw[c] = pr.norm_w[i]; + } } - const int h = lane - DSV4_HC*DSV4_HC; - float pre_h = pre[0]; - float m_h = m[DSV4_HC]; -#pragma unroll - for (int j = 1; j < DSV4_HC; ++j) { - pre_h = h == j ? pre[j] : pre_h; - m_h = h == j ? m[DSV4_HC + j] : m_h; + }; + // the first chunks' streams requested with the partials + request(0); + dsv4_hc_gram_mix(partials + (int64_t) it*DSV4_HC_GRAM_ROWS*n_slices, n_slices, mix); + __syncthreads(); + + float pre[DSV4_HC]; + float rms_flat; + float rms; + dsv4_hc_gram_pre(mix, base_s, pr, pre, rms_flat, rms); + if (warp == 0) { + dsv4_hc_gram_weights(mix, base_s, pr, it, pre, rms_flat); + } + for (int c0 = 0; c0 < n_chunks; c0 += DSV4_HC_FRONT_ONE_CH) { + if (c0 > 0) { + request(c0); } - float * d = weights_out + it*sw1; - if (lane < DSV4_HC*DSV4_HC) { - d[(2*DSV4_HC + lane)*sw0] = c_lane; - } else if (lane < DSV4_HC*DSV4_HC + DSV4_HC) { - d[h*sw0] = pre_h; - d[(DSV4_HC + h)*sw0] = 2.0f/(1.0f + expf(-(m_h*scale[ss0] + base_s[DSV4_HC + h]))); +#pragma unroll + for (int c = 0; c < DSV4_HC_FRONT_ONE_CH; ++c) { + const int64_t i = (int64_t) (c0 + c)*DSV4_HC_PRE_GRAM_THR + threadIdx.x; + if (c0 + c < n_chunks && i < pr.n_embd) { + dsv4_hc_gram_out(pr, it, i, xs[c], pre, rms, nw[c]); + } } } +} - if (i < n_embd) { - float v = xs[0]*pre[0]; -#pragma unroll - for (int h = 1; h < DSV4_HC; ++h) { - v += xs[h]*pre[h]; +// GGML_CUDA_HC_FRONT_CHECK=1: the one launch's normed mix, q8_1 copy and weights (one) against the two kernels' (two), +// bit for bit; a block a token +static __global__ void dsv4_hc_front_check_bits(const dsv4_hc_pre_args one, const dsv4_hc_pre_args two) { + const int it = blockIdx.y; + for (int64_t i = threadIdx.x; i < one.n_embd; i += blockDim.x) { + const uint32_t a = __float_as_uint(one.dst[it*one.sd1 + i]); + const uint32_t b = __float_as_uint(two.dst[it*two.sd1 + i]); + if (a != b) { + printf("hc front check: token %d element %lld is %08x from the one launch, %08x from the two\n", + it, (long long) i, a, b); + __trap(); + } + if (one.q8 != nullptr && i % QK8_1 == 0) { + const block_q8_1 * qa = one.q8 + it*one.q8_s1 + i/QK8_1; + const block_q8_1 * qb = two.q8 + it*two.q8_s1 + i/QK8_1; + bool same = *(const uint32_t *) &qa->ds == *(const uint32_t *) &qb->ds; + for (int j = 0; j < QK8_1; ++j) { + same = same && qa->qs[j] == qb->qs[j]; + } + if (!same) { + printf("hc front check: token %d q8_1 block %lld differs\n", it, (long long) (i/QK8_1)); + __trap(); + } + } + } + if (threadIdx.x < DSV4_HC_MIX) { + const uint32_t a = __float_as_uint(one.weights_out[it*one.sw1 + threadIdx.x*one.sw0]); + const uint32_t b = __float_as_uint(two.weights_out[it*two.sw1 + threadIdx.x*two.sw0]); + if (a != b) { + printf("hc front check: token %d weight %d is %08x from the one launch, %08x from the two\n", + it, (int) threadIdx.x, a, b); + __trap(); } - dst[it*sd1 + i] = v*rms*nw; } } @@ -828,6 +1017,33 @@ static bool dsv4_hc_pre_gram_legacy() { return legacy; } +// GGML_CUDA_HC_FRONT_ONE=1: the Gram path's front in one launch (dsv4_hc_front_one) instead of its two kernels. +// Opt-in, because it measured slower than the two on the 44-layer GLM-5.3 proxy at -d 32768 (NVTX, graphs off, the +// last 31 decode tokens, RTX 5070 Ti): one launch 8.38 us against the two kernels' 2.15 + 3.01, so 759 us a token +// against 535, and the nodes' issue span 1475 us a token against 1395. The 88 launches a token it saves do not pay +// for the serialization it adds, one block doing for the whole token what 16 blocks did in parallel. +// GGML_CUDA_HC_FRONT_CHECK=1: both, the two kernels into scratch, and a trap on any bit of the outputs that differs; +// it implies the one launch, so the check always has something to compare. +static bool dsv4_hc_front_one_enabled() { + static const bool enabled = ggml_env_switch("GGML_CUDA_HC_FRONT_ONE"); + return enabled; +} + +static bool dsv4_hc_front_check() { + static const bool check = ggml_env_switch("GGML_CUDA_HC_FRONT_CHECK"); + return check; +} + +bool ggml_cuda_dsv4_hc_writes_q8_1(const ggml_tensor * node) { + if (dsv4_hc_pre_gram_legacy() || node->op != GGML_OP_MUL || node->type != GGML_TYPE_F32 || + node->ne[0] % MATRIX_ROW_PADDING != 0) { + return false; + } + const ggml_tensor * rms = node->src[0] != nullptr && node->src[0]->op == GGML_OP_RMS_NORM ? node->src[0] : node->src[1]; + return rms != nullptr && rms->op == GGML_OP_RMS_NORM && rms->src[0] != nullptr && + rms->src[0]->op == GGML_OP_DSV4_HC_PRE; +} + bool ggml_cuda_dsv4_hc_pre_fused_supported(const ggml_tensor * rms_flat, const ggml_tensor * mm, const ggml_tensor * weights, const ggml_tensor * pre, const ggml_tensor * rms, const ggml_tensor * mul) { const ggml_tensor * flat = rms_flat->src[0]; @@ -947,38 +1163,108 @@ static void dsv4_hc_front(ggml_backend_cuda_context & ctx, const ggml_tensor * p pa.sc2 = comb_w->nb[2] / sizeof(float); pa.sx_h = x->nb[1] / sizeof(float); } - auto launch = [&](auto kernel, const auto * w) { + // the normed mix's q8_1 copy, when the evaluation holds one for its quantized readers + // (ggml_cuda_mmvq_shared_q8_1::produce): written here, not by a q8_1 launch at its first reader + ggml_cuda_mmvq_shared_q8_1::entry * q8 = n_embd % MATRIX_ROW_PADDING == 0 ? ctx.mmvq_shared_q8_1.produce(mul) : + nullptr; + + dsv4_hc_pre_args pr = {}; + pr.x = (const float *) x->data; + pr.scale = (const float *) scale->data; + pr.base = (const float *) base->data; + pr.norm_w = (const float *) norm_w->data; + pr.weights_out = (float *) weights->data; + pr.dst = (float *) mul->data; + pr.q8 = q8 != nullptr ? (block_q8_1 *) q8->q8_1 : nullptr; + pr.n_embd = n_embd; + pr.k = k; + pr.sx1 = x->nb[1] / sizeof(float); + pr.sx2 = x->nb[2] / sizeof(float); + pr.ss0 = scale->nb[0] / sizeof(float); + pr.sb0 = base->nb[0] / sizeof(float); + pr.sw0 = weights->nb[0] / sizeof(float); + pr.sw1 = weights->nb[1] / sizeof(float); + pr.sd1 = mul->nb[1] / sizeof(float); + pr.q8_s1 = n_embd / QK8_1; + pr.eps_flat = ggml_get_op_params_f32(rms_flat, 0); + pr.eps_hc = ggml_get_op_params_f32(weights, 0); + pr.eps_norm = ggml_get_op_params_f32(rms, 0); + pr.n_iter = ggml_get_op_params_i32(weights, 1); + pr.base_prewait = dsv4_hc_prewait(base); + pr.norm_prewait = dsv4_hc_prewait(norm_w); + + // the two kernels, writing p's outputs + auto launch_two = [&](const dsv4_hc_pre_args & p) { + auto launch = [&](auto kernel, const auto * w) { + ggml_cuda_kernel_launch(kernel, gram_params, + (float *) x->data, w, partials.get(), n_embd, sx1, sw1, w_prewait, pa); + }; + switch (hc_fn->type) { + case GGML_TYPE_F32: + post ? launch(dsv4_hc_mix_gram, (const float *) hc_fn->data) + : launch(dsv4_hc_mix_gram, (const float *) hc_fn->data); + break; + case GGML_TYPE_F16: + post ? launch(dsv4_hc_mix_gram, (const half *) hc_fn->data) + : launch(dsv4_hc_mix_gram, (const half *) hc_fn->data); + break; + case GGML_TYPE_BF16: + post ? launch(dsv4_hc_mix_gram, (const nv_bfloat16 *) hc_fn->data) + : launch(dsv4_hc_mix_gram, (const nv_bfloat16 *) hc_fn->data); + break; + default: + GGML_ABORT("unsupported hc_fn type %s", ggml_type_name(hc_fn->type)); + } + const int n_blocks = (int) ((n_embd + DSV4_HC_PRE_GRAM_THR - 1) / DSV4_HC_PRE_GRAM_THR); + const ggml_cuda_kernel_launch_params pre_params = + ggml_cuda_kernel_launch_params(dim3(n_blocks, n_tokens, 1), dim3(DSV4_HC_PRE_GRAM_THR, 1, 1), 0, stream); + ggml_cuda_kernel_launch(dsv4_hc_pre_gram_f32, pre_params, (const float *) partials.get(), n_slices, p); + }; + + // the tickets exist from the first graph evaluation on; a front takes at most their tokens + static_assert(DSV4_HC_PRE_FUSED_MAX_TOKENS <= ggml_cuda_hc_front_tickets::n_tokens, "a token a ticket"); + unsigned int * tickets = ctx.hc_front_ticket(); + if (!(dsv4_hc_front_one_enabled() || dsv4_hc_front_check()) || tickets == nullptr) { + launch_two(pr); + return; + } + GGML_ASSERT(n_tokens <= ggml_cuda_hc_front_tickets::n_tokens); + + ggml_cuda_pool_alloc check_dst(ctx.pool()); + ggml_cuda_pool_alloc check_weights(ctx.pool()); + ggml_cuda_pool_alloc check_q8(ctx.pool()); + dsv4_hc_pre_args two = pr; + if (dsv4_hc_front_check()) { + two.dst = check_dst.alloc(ggml_nbytes(mul) / sizeof(float)); + two.weights_out = check_weights.alloc(ggml_nbytes(weights) / sizeof(float)); + two.q8 = pr.q8 != nullptr ? check_q8.alloc(n_tokens*pr.q8_s1) : nullptr; + launch_two(two); + } + + auto launch_one = [&](auto kernel, const auto * w) { ggml_cuda_kernel_launch(kernel, gram_params, - (float *) x->data, w, partials.get(), n_embd, sx1, sw1, w_prewait, pa); + (float *) x->data, w, partials.get(), sx1, sw1, w_prewait, pa, pr, tickets); }; switch (hc_fn->type) { case GGML_TYPE_F32: - post ? launch(dsv4_hc_mix_gram, (const float *) hc_fn->data) - : launch(dsv4_hc_mix_gram, (const float *) hc_fn->data); + post ? launch_one(dsv4_hc_front_one, (const float *) hc_fn->data) + : launch_one(dsv4_hc_front_one, (const float *) hc_fn->data); break; case GGML_TYPE_F16: - post ? launch(dsv4_hc_mix_gram, (const half *) hc_fn->data) - : launch(dsv4_hc_mix_gram, (const half *) hc_fn->data); + post ? launch_one(dsv4_hc_front_one, (const half *) hc_fn->data) + : launch_one(dsv4_hc_front_one, (const half *) hc_fn->data); break; case GGML_TYPE_BF16: - post ? launch(dsv4_hc_mix_gram, (const nv_bfloat16 *) hc_fn->data) - : launch(dsv4_hc_mix_gram, (const nv_bfloat16 *) hc_fn->data); + post ? launch_one(dsv4_hc_front_one, (const nv_bfloat16 *) hc_fn->data) + : launch_one(dsv4_hc_front_one, (const nv_bfloat16 *) hc_fn->data); break; default: GGML_ABORT("unsupported hc_fn type %s", ggml_type_name(hc_fn->type)); } - const int n_blocks = (int) ((n_embd + DSV4_HC_PRE_GRAM_THR - 1) / DSV4_HC_PRE_GRAM_THR); - const ggml_cuda_kernel_launch_params pre_params = - ggml_cuda_kernel_launch_params(dim3(n_blocks, n_tokens, 1), dim3(DSV4_HC_PRE_GRAM_THR, 1, 1), 0, stream); - ggml_cuda_kernel_launch(dsv4_hc_pre_gram_f32, pre_params, - (const float *) partials.get(), n_slices, (const float *) x->data, (const float *) scale->data, - (const float *) base->data, (const float *) norm_w->data, (float *) weights->data, (float *) mul->data, - n_embd, k, x->nb[1] / sizeof(float), x->nb[2] / sizeof(float), scale->nb[0] / sizeof(float), - base->nb[0] / sizeof(float), weights->nb[0] / sizeof(float), weights->nb[1] / sizeof(float), - mul->nb[1] / sizeof(float), ggml_get_op_params_f32(rms_flat, 0), ggml_get_op_params_f32(weights, 0), - ggml_get_op_params_i32(weights, 1), ggml_get_op_params_f32(rms, 0), dsv4_hc_prewait(base), - dsv4_hc_prewait(norm_w)); + if (dsv4_hc_front_check()) { + dsv4_hc_front_check_bits<<>>(pr, two); + } return; } GGML_ASSERT(post == nullptr); diff --git a/ggml/src/ggml-cuda/dsv4-hc.cuh b/ggml/src/ggml-cuda/dsv4-hc.cuh index 31c4bddb6b1d..14f848fb5ca4 100644 --- a/ggml/src/ggml-cuda/dsv4-hc.cuh +++ b/ggml/src/ggml-cuda/dsv4-hc.cuh @@ -26,3 +26,8 @@ bool ggml_cuda_dsv4_hc_post_pre_fused_supported(const ggml_tensor * post, const void ggml_cuda_op_dsv4_hc_post_pre_fused(ggml_backend_cuda_context & ctx, ggml_tensor * post, const ggml_tensor * rms_flat, const ggml_tensor * mm, ggml_tensor * weights, const ggml_tensor * pre, const ggml_tensor * rms, ggml_tensor * mul); + +// Whether node is a front's mul, the sublayer's normed mix: the fused front writes its q8_1 copy beside it when the +// evaluation holds one for it (ggml_cuda_mmvq_shared_q8_1::produce) and n_embd is a multiple of MATRIX_ROW_PADDING, +// not under GGML_CUDA_HC_PRE_GRAM_LEGACY. Where the front does not run fused, its readers quantize it themselves. +bool ggml_cuda_dsv4_hc_writes_q8_1(const ggml_tensor * node); diff --git a/ggml/src/ggml-cuda/fattn-common.cuh b/ggml/src/ggml-cuda/fattn-common.cuh index a6a687ef184e..02e97bc9632c 100644 --- a/ggml/src/ggml-cuda/fattn-common.cuh +++ b/ggml/src/ggml-cuda/fattn-common.cuh @@ -838,6 +838,25 @@ static __global__ void flash_attn_mask_to_KV_max( KV_min[sequence*ne31 + jt] = 0; } +// The sparse mode (ggml_flash_attn_ext_set_n_kv_max): each f16 mask row's finite entries as ascending cell indices, n_kv_max +// a row (sequence, then row, of mask->ne[1] rows), -1 past the row's count. Defined in fattn.cu. +void ggml_cuda_flash_attn_ext_compact_mask( + const ggml_tensor * mask, int32_t * indices, int32_t n_kv_max, int rows_per_list, cudaStream_t stream); + +// How many consecutive query rows share one sparse index list. 1 is one list a query, which is what the ncols1 == 1 +// sparse kernel reads and the only width built today. Above 1 a list is the UNION of its rows' cells, so a tile of +// queries gathers each cell once rather than once per query -- the saving is the overlap between adjacent queries' +// selections, which is a property of the TRAINED indexer and is not measurable on a random-weight proxy. Default 1 until +// a real-weight box measures that overlap; GGML_CUDA_FATTN_SPARSE_ROWS=N asks for N (2, 4, 8 or 16). +static int ggml_cuda_fattn_sparse_rows_per_list() { + static const int rows = [] { + const char * env = getenv("GGML_CUDA_FATTN_SPARSE_ROWS"); + const int n = env ? atoi(env) : 1; + return n == 2 || n == 4 || n == 8 || n == 16 ? n : 1; + }(); + return rows; +} + // KV_live, the live KV steps of the mma kernel (nbatch_fa cells each) for n = Q tiles x sequences: // [0, n) live steps per (sequence, Q tile), at least 1 | [n] blocks done | [n+1, 2n+1) its first live step within its sequence | // [ne03] live steps per sequence | [ne03 + 1] a sequence's first unit of work (its steps times its output tiles per Q tile), then the total | @@ -1452,7 +1471,8 @@ template void launch_fattn( ggml_backend_cuda_context & ctx, ggml_tensor * dst, fattn_kernel_t fattn_kernel, const int nwarps, const size_t nbytes_shared, const int nbatch_fa, const bool need_f16_K, const bool need_f16_V, const bool stream_k, const int warp_size = WARP_SIZE, - const bool kv_live_ok = false // the kernel reads KV_live + const bool kv_live_ok = false, // the kernel reads KV_live + const bool use_sparse = false // the kernel reads K, V and the mask by the indices in KV_max, n_kv_max of them ) { constexpr int ncols = ncols1 * ncols2; @@ -1570,6 +1590,27 @@ void launch_fattn( const int ntiles_z_gqa = ((gqa_ratio + ncols2 - 1) / ncols2); const int ntiles_dst = ntiles_x * ntiles_z_gqa * K->ne[2] * Q->ne[3]; + // use_sparse: the mask rows compacted to their live cells, which the kernel reads in place of the whole cache, so none + // of the scans below + const int32_t n_kv_max = use_sparse ? ggml_flash_attn_ext_get_n_kv_max(KQV) : 0; + if (use_sparse) { + GGML_ASSERT(mask != nullptr && !mask_packed); + GGML_ASSERT(n_kv_max > 0); + + // a union list holds at most rows_per_list times one query's cells, and never more than the cache itself + const int rows_per_list = ggml_cuda_fattn_sparse_rows_per_list(); + // The builder can emit union lists; the kernel below still reads one list a query (ncols1 == 1), so a width above + // 1 would silently read the wrong list rather than fail. Abort until the ncols1 > 1 sparse kernel lands. + GGML_ASSERT(rows_per_list == 1 || + !"GGML_CUDA_FATTN_SPARSE_ROWS > 1 needs the ncols1 > 1 sparse kernel, which is not built yet"); + const int32_t n_kv_list = rows_per_list == 1 ? n_kv_max + : int32_t(std::min(int64_t(n_kv_max) * rows_per_list, mask->ne[0])); + const int64_t n_lists = (mask->ne[1] + rows_per_list - 1) / rows_per_list; + + KV_max.alloc(size_t(n_kv_list) * n_lists * mask->ne[3]); + ggml_cuda_flash_attn_ext_compact_mask(mask, KV_max.ptr, n_kv_list, rows_per_list, main_stream); + } + // Optional optimization where the mask is scanned to determine whether part of the calculation can be skipped. // Only worth the overhead if there is at lease one FATTN_KQ_STRIDE x FATTN_KQ_STRIDE square to be skipped or // multiple sequences of possibly different lengths. @@ -1583,7 +1624,7 @@ void launch_fattn( // the kernel applies the mask over the whole range. GGML_CUDA_FATTN_MASK_PREFIX_LEGACY=1: the hint is ignored. static const bool mask_prefix_legacy = ggml_env_switch("GGML_CUDA_FATTN_MASK_PREFIX_LEGACY"); const bool mask_prefix = !mask_prefix_legacy && mask && Q->ne[1] <= 16 && ggml_flash_attn_ext_get_mask_prefix(KQV); - const bool kv_range = !kv_range_legacy && !mask_prefix && mask && Q->ne[1] <= 16 && K->ne[1] >= 4096 && + const bool kv_range = !use_sparse && !kv_range_legacy && !mask_prefix && mask && Q->ne[1] <= 16 && K->ne[1] >= 4096 && (uintptr_t) mask->data % 16 == 0 && mask->nb[1] % 16 == 0 && mask->nb[3] % 16 == 0; // 16-byte mask reads // kv_live: the mma kernel reads only the KV steps a row of its Q tile sees, the masked tiles between them skipped too (with @@ -1593,7 +1634,7 @@ void launch_fattn( // The live fixup tests the blocks it combines for an empty range only when there can be one (fewer units of work than // blocks), and lets the next kernel launch as it starts. GGML_CUDA_FATTN_LIVE_FIXUP_LEGACY=1: it tests every block and does not. static const bool kv_live_fixup_legacy = ggml_env_switch("GGML_CUDA_FATTN_LIVE_FIXUP_LEGACY"); - const bool kv_live = kv_live_ok && !kv_live_legacy && !mask_prefix && stream_k && GGML_CUDA_CC_IS_NVIDIA(cc) && mask && K->ne[1] >= 4096 && + const bool kv_live = !use_sparse && kv_live_ok && !kv_live_legacy && !mask_prefix && stream_k && GGML_CUDA_CC_IS_NVIDIA(cc) && mask && K->ne[1] >= 4096 && K->ne[1] % FATTN_KQ_STRIDE == 0 && (nbatch_fa == 32 || nbatch_fa == 64) && (uintptr_t) mask->data % 16 == 0 && mask->nb[1] % 16 == 0 && mask->nb[3] % 16 == 0; // The live steps depend on the mask and the split alone, and every attention layer of a graph reads the same mask: the @@ -1652,7 +1693,7 @@ void launch_fattn( CUDA_CHECK(cudaGetLastError()); KV_live_ptr = live; } - } else if (mask && K->ne[1] % FATTN_KQ_STRIDE == 0 && (Q->ne[1] >= 1024 || Q->ne[3] > 1 || kv_range)) { + } else if (!use_sparse && mask && K->ne[1] % FATTN_KQ_STRIDE == 0 && (Q->ne[1] >= 1024 || Q->ne[3] > 1 || kv_range)) { const size_t unit = mask_packed ? sizeof(uint16_t) : sizeof(half2); const int64_t s31 = mask->nb[1] / unit; const int64_t s33 = mask->nb[3] / unit; @@ -1685,7 +1726,8 @@ void launch_fattn( GGML_ASSERT(max_blocks_per_sm > 0); int parallel_blocks = max_blocks_per_sm; - const int ntiles_KV = (K->ne[1] + nbatch_fa - 1) / nbatch_fa; // Max. number of parallel blocks limited by KV cache length. + const int64_t n_kv = use_sparse ? n_kv_max : K->ne[1]; // the cells the kernel steps over + const int ntiles_KV = (n_kv + nbatch_fa - 1) / nbatch_fa; // Max. number of parallel blocks limited by KV cache length. dim3 blocks_num; if (stream_k) { @@ -1801,7 +1843,7 @@ void launch_fattn( !stream_k && parallel_blocks > 1 ? dst_tmp.ptr : (float *) KQV->data, dst_tmp_meta.ptr, scale, max_bias, m0, m1, n_head_log2, logit_softcap, Q->ne[0], ne01, Q->ne[2], Q->ne[3], Q->nb[1], Q->nb[2], Q->nb[3], - K->ne[0], K->ne[1], K->ne[2], K->ne[3], nb11, nb12, nb13, + K->ne[0], n_kv, K->ne[2], K->ne[3], nb11, nb12, nb13, nb21, nb22, nb23, mask ? mask->ne[1] : 0, mask ? mask->ne[2] : 0, mask ? mask->ne[3] : 0, mask ? mask->nb[1] : 0, mask ? mask->nb[2] : 0, mask ? mask->nb[3] : 0, diff --git a/ggml/src/ggml-cuda/fattn-mma-f16.cuh b/ggml/src/ggml-cuda/fattn-mma-f16.cuh index f726f39ace79..ff6b57fa7f0c 100644 --- a/ggml/src/ggml-cuda/fattn-mma-f16.cuh +++ b/ggml/src/ggml-cuda/fattn-mma-f16.cuh @@ -349,20 +349,26 @@ static __host__ int ggml_cuda_fattn_mma_get_nstages(const int DKQ, const int DV, return cp_async_available(cc) && ncols2 >= 2 ? ggml_cuda_fattn_mma_get_nstages_target(DKQ, DV, ncols1*ncols2, cc) : 0; } -static constexpr __device__ int ggml_cuda_fattn_mma_get_nstages(const int DKQ, const int DV, const int ncols1, const int ncols2) { +static constexpr __device__ int ggml_cuda_fattn_mma_get_nstages( + const int DKQ, const int DV, const int ncols1, const int ncols2, const bool use_sparse) { #ifdef CP_ASYNC_AVAILABLE - return ncols2 >= 2 ? ggml_cuda_fattn_mma_get_nstages_target(DKQ, DV, ncols1*ncols2) : 0; + const int nstages_target = ncols2 >= 2 ? ggml_cuda_fattn_mma_get_nstages_target(DKQ, DV, ncols1*ncols2) : 0; + // sparse gather is not implemented for multi-stage loading + return use_sparse && nstages_target > 1 ? 1 : nstages_target; #else - GGML_UNUSED_VARS(DKQ, DV, ncols1, ncols2); + GGML_UNUSED_VARS(DKQ, DV, ncols1, ncols2, use_sparse); return 0; #endif // CP_ASYNC_AVAILABLE } // ------------------------------------------------------------------------------------------------------------------ -template +// Rows [k_VKQ_0, k_VKQ_0 + nbatch_fa) of KV, or with use_sparse the rows the indices name there (-1 or past i_sup: row 0 +// with cp_async, zeros without; the mask is -inf there). +template static __device__ __forceinline__ void flash_attn_ext_f16_load_tile( - const half2 * const __restrict__ KV, half2 * const __restrict__ tile_KV, const int D2, const int stride_KV, const int i_sup) { + const half2 * const __restrict__ KV, half2 * const __restrict__ tile_KV, const int D2, const int stride_KV, + const int k_VKQ_0, const int i_sup, const int32_t * const __restrict__ indices) { constexpr int warp_size = ggml_cuda_get_physical_warp_size(); // K/V data is loaded with decreasing granularity for D for better memory bandwidth. // The minimum granularity is 16 bytes. @@ -370,7 +376,7 @@ static __device__ __forceinline__ void flash_attn_ext_f16_load_tile( const int chunks_per_row = D2 / h2_per_chunk; if constexpr (use_cp_async) { static_assert(warp_size == 32, "bad warp_size"); - static_assert(!oob_check, "OOB check not compatible with cp_async"); + static_assert(!oob_check || use_sparse, "OOB check not compatible with cp_async"); constexpr int preload = 64; const unsigned int tile_KV_32 = ggml_cuda_cvta_generic_to_shared(tile_KV); @@ -393,11 +399,19 @@ static __device__ __forceinline__ void flash_attn_ext_f16_load_tile( break; } + int64_t i_KV; + if constexpr (use_sparse) { + const int32_t index = i < i_sup ? indices[k_VKQ_0 + i] : 0; + i_KV = index >= 0 ? index : 0; + } else { + i_KV = k_VKQ_0 + i; + } + #pragma unroll for (int k0 = k0_start; k0 < k0_stop; k0 += stride_k) { const int k = k0 + (stride_k == warp_size ? threadIdx.x : threadIdx.x % stride_k); - cp_async_cg_16(tile_KV_32 + i*(stride_tile*sizeof(half2)) + k*16, KV + i*stride_KV + k*h2_per_chunk); + cp_async_cg_16(tile_KV_32 + i*(stride_tile*sizeof(half2)) + k*16, KV + i_KV*stride_KV + k*h2_per_chunk); } } }; @@ -432,8 +446,14 @@ static __device__ __forceinline__ void flash_attn_ext_f16_load_tile( for (int k0 = k0_start; k0 < k0_stop; k0 += stride_k) { const int k = k0 + (stride_k == warp_size ? threadIdx.x : threadIdx.x % stride_k); - ggml_cuda_memcpy_1<16>(tile_KV + i*stride_tile + k*4, - !oob_check || i < i_sup ? KV + i*stride_KV + k*h2_per_chunk : zero); + const half2 * src; + if constexpr (use_sparse) { + const int32_t index = i < i_sup ? indices[k_VKQ_0 + i] : -1; + src = index >= 0 ? KV + int64_t(index)*stride_KV + k*h2_per_chunk : zero; + } else { + src = !oob_check || i < i_sup ? KV + int64_t(k_VKQ_0 + i)*stride_KV + k*h2_per_chunk : zero; + } + ggml_cuda_memcpy_1<16>(tile_KV + i*stride_tile + k*4, src); } } }; @@ -449,11 +469,28 @@ static __device__ __forceinline__ void flash_attn_ext_f16_load_tile( // Loads the mask tile for KV cells [k_VKQ_0, k_VKQ_0 + nbatch_fa) into tile_mask as halves. // mask_packed: the mask is bit-packed (16 cells per 16-bit word, bit set = attend), unpacked here so the rest of the kernel is unchanged. -template +// use_sparse: the cells the indices name there, -inf past i_sup or at -1 (one Q row a tile, an f16 mask). +template static __device__ __forceinline__ void flash_attn_ext_f16_load_mask( const half * const __restrict__ mask_base, half * const __restrict__ tile_mask, - const int stride_mask, const int k_VKQ_0, const int i_sup, const int j0, const uint3 ne01, const bool mask_packed) { + const int stride_mask, const int k_VKQ_0, const int i_sup, const int j0, const uint3 ne01, const bool mask_packed, + const int32_t * const __restrict__ indices) { constexpr int warp_size = ggml_cuda_get_physical_warp_size(); + if constexpr (use_sparse) { + static_assert(!use_cp_async, "sparse gather incompatible with cp_async"); + static_assert(ncols1 == 1, "a tile's rows share its indices"); + if (threadIdx.y == 0) { + const int j_vram = fastmodulo(j0, ne01); +#pragma unroll + for (int i0 = 0; i0 < nbatch_fa; i0 += warp_size) { + const int i = i0 + threadIdx.x; + + const int32_t index = i < i_sup ? indices[k_VKQ_0 + i] : -1; + tile_mask[i] = index >= 0 ? mask_base[int64_t(j_vram)*stride_mask + index] : half(-INFINITY); + } + } + return; + } if (mask_packed) { static_assert(nbatch_fa % 16 == 0, "bad nbatch_fa"); constexpr int nwords = nbatch_fa/16; @@ -864,13 +901,14 @@ static __device__ __forceinline__ void flash_attn_ext_raw_KQ( } template static __device__ __forceinline__ void flash_attn_ext_f16_iter( const float2 * const __restrict__ Q_f2, const half2 * const __restrict__ K_h2, const half2 * const __restrict__ V_h2, const half * const __restrict__ mask_h, + const int32_t * const __restrict__ indices, float2 * const __restrict__ dstk, float2 * const __restrict__ dstk_fixup, const float scale, @@ -906,9 +944,10 @@ static __device__ __forceinline__ void flash_attn_ext_f16_iter( constexpr int nbatch_K2 = ggml_cuda_fattn_mma_get_nbatch_K2(DKQ, DV, ncols); constexpr int nbatch_V2 = ggml_cuda_fattn_mma_get_nbatch_V2(DKQ, DV, ncols); constexpr bool Q_in_reg = ggml_cuda_fattn_mma_get_Q_in_reg (DKQ, DV, ncols); - constexpr int nstages = ggml_cuda_fattn_mma_get_nstages (DKQ, DV, ncols1, ncols2); + constexpr int nstages = ggml_cuda_fattn_mma_get_nstages (DKQ, DV, ncols1, ncols2, use_sparse); constexpr bool raw_KV = type_K != GGML_TYPE_F16; static_assert(raw_KV == (type_V != GGML_TYPE_F16), "K and V are both raw or both f16"); + static_assert(!use_sparse || !raw_KV, "the sparse gather reads f16 K/V"); static_assert(!raw_KV || cols_per_warp == 16, "raw K*Q needs 16 Q columns per warp"); // The next raw K tile loads as soon as K*Q has read this one, in flight beside this step's V through the softmax, // where 2 blocks share an SM: the 4-warp tiles of a verify, whose K then loaded only after V arrived, 1-3 % slower at @@ -931,6 +970,7 @@ static __device__ __forceinline__ void flash_attn_ext_f16_iter( if constexpr (nstages > 1) { static_assert(!oob_check, "OOB check incompatible with multi-stage pipeline"); + static_assert(!use_sparse, "sparse gather not implemented for multi-stage loading"); static_assert(!V_is_K_view, "K data reuse not implemented multi-stage loading"); static_assert(nbatch_K2 == DKQ/2, "batching not implemented for multi stage loading"); constexpr bool use_cp_async = true; @@ -944,15 +984,16 @@ static __device__ __forceinline__ void flash_attn_ext_f16_iter( cp_async_commit_group(); // the raw V tile, which the wait before the VKQ tile waits for alone } } else { - flash_attn_ext_f16_load_tile - (V_h2 + int64_t(k_VKQ_0)*stride_V, tile_V, nbatch_V2, stride_V, k_VKQ_sup); + flash_attn_ext_f16_load_tile + (V_h2, tile_V, nbatch_V2, stride_V, k_VKQ_0, k_VKQ_sup, nullptr); } } else { static_assert(!raw_KV, "raw K/V needs the multi-stage pipeline"); - constexpr bool use_cp_async = nstages == 1; + // the sparse mask values are gathered per element, always loaded synchronously + constexpr bool use_cp_async = nstages == 1 && !use_sparse; if (ncols2 > 1 || mask_h) { - flash_attn_ext_f16_load_mask - (mask_h, tile_mask, stride_mask, k_VKQ_0, k_VKQ_sup, jt*ncols1, ne01, mask_packed); + flash_attn_ext_f16_load_mask + (mask_h, tile_mask, stride_mask, k_VKQ_0, k_VKQ_sup, jt*ncols1, ne01, mask_packed, indices); } } @@ -965,8 +1006,8 @@ static __device__ __forceinline__ void flash_attn_ext_f16_iter( if constexpr (nstages <= 1) { const int k0_diff = k0_stop - k0_start; constexpr bool use_cp_async = nstages == 1; - flash_attn_ext_f16_load_tile - (K_h2 + int64_t(k_VKQ_0)*stride_K + k0_start, tile_K, k0_diff, stride_K, k_VKQ_sup); + flash_attn_ext_f16_load_tile + (K_h2 + k0_start, tile_K, k0_diff, stride_K, k_VKQ_0, k_VKQ_sup, indices); if (use_cp_async) { cp_async_wait_all(); } @@ -1310,8 +1351,8 @@ static __device__ __forceinline__ void flash_attn_ext_f16_iter( __syncthreads(); if (!last_iter) { if (ncols2 > 1 || mask_h) { - flash_attn_ext_f16_load_mask - (mask_h, tile_mask, stride_mask, kb0_next*nbatch_fa, k_VKQ_sup, jt*ncols1, ne01, mask_packed); + flash_attn_ext_f16_load_mask + (mask_h, tile_mask, stride_mask, kb0_next*nbatch_fa, k_VKQ_sup, jt*ncols1, ne01, mask_packed, nullptr); } if constexpr (raw_KV) { if constexpr (!raw_K_early) { @@ -1319,8 +1360,8 @@ static __device__ __forceinline__ void flash_attn_ext_f16_iter( ((const char *) K_h2 + int64_t(kb0_next)*nbatch_fa*stride_K, tile_raw, stride_K); } } else { - flash_attn_ext_f16_load_tile - (K_h2 + int64_t(kb0_next)*nbatch_fa*stride_K, tile_K, nbatch_K2, stride_K, k_VKQ_sup); + flash_attn_ext_f16_load_tile + (K_h2, tile_K, nbatch_K2, stride_K, kb0_next*nbatch_fa, k_VKQ_sup, nullptr); } } if constexpr (raw_KV) { @@ -1342,8 +1383,8 @@ static __device__ __forceinline__ void flash_attn_ext_f16_iter( const int i0_diff = i0_stop - i0_start; if (!V_is_K_view || i0_stop > 2*nbatch_K2) { constexpr bool use_cp_async = nstages == 1; - flash_attn_ext_f16_load_tile - (V_h2 + int64_t(k_VKQ_0)*stride_V + i0_start/2, tile_V, i0_diff/2, stride_V, k_VKQ_sup); + flash_attn_ext_f16_load_tile + (V_h2 + i0_start/2, tile_V, i0_diff/2, stride_V, k_VKQ_0, k_VKQ_sup, indices); if (use_cp_async) { cp_async_wait_all(); } @@ -1398,7 +1439,7 @@ static __device__ __forceinline__ void flash_attn_ext_f16_iter( } } #else - GGML_UNUSED_VARS(Q_f2, K_h2, V_h2, mask_h, dstk, dstk_fixup, + GGML_UNUSED_VARS(Q_f2, K_h2, V_h2, mask_h, indices, dstk, dstk_fixup, scale, slope, logit_softcap, ne01, ne02, stride_K, stride_V, stride_mask, mask_packed, tile_Q, tile_K, tile_V, tile_mask, @@ -1500,12 +1541,14 @@ template struct mma_tile_sizes { // buffer. is_fixup: the block ends inside the tile; its result and meta go to the fixup buffer. Neither: it did the whole tile. // They are arguments, not template parameters: one copy of the tile's code serves every block, so the few blocks that end a // tile run instructions the rest have already brought into the SM's instruction cache. -template +template static __device__ __forceinline__ void flash_attn_ext_f16_process_tile( const float2 * const __restrict__ Q_f2, const half2 * const __restrict__ K_h2, const half2 * const __restrict__ V_h2, const half * const __restrict__ mask_h, + const int32_t * const __restrict__ indices, // use_sparse: the tile's n_kv_max cell indices const float * const __restrict__ sinks_f, float2 * const __restrict__ dstk, float2 * const __restrict__ dstk_fixup, @@ -1550,7 +1593,7 @@ static __device__ __forceinline__ void flash_attn_ext_f16_process_tile( constexpr int nbatch_V2 = ggml_cuda_fattn_mma_get_nbatch_V2 (DKQ, DV, ncols); constexpr int nbatch_combine = ggml_cuda_fattn_mma_get_nbatch_combine(DKQ, DV, ncols); constexpr bool Q_in_reg = ggml_cuda_fattn_mma_get_Q_in_reg (DKQ, DV, ncols); - constexpr int nstages = ggml_cuda_fattn_mma_get_nstages (DKQ, DV, ncols1, ncols2); + constexpr int nstages = ggml_cuda_fattn_mma_get_nstages (DKQ, DV, ncols1, ncols2, use_sparse); if (cols_per_warp > ncols) { NO_DEVICE_CODE; @@ -1655,33 +1698,35 @@ static __device__ __forceinline__ void flash_attn_ext_f16_process_tile( // Preload mask and K data for first iteration when using cp_async with multiple stages: if constexpr (nstages > 1) { static_assert(nbatch_K2 == DKQ/2, "batching not implemented for multi-stage pipeline"); + static_assert(!use_sparse, "sparse gather not implemented for multi-stage loading"); constexpr bool use_cp_async = true; constexpr bool oob_check = false; constexpr int k_VKQ_sup = nbatch_fa; if (ncols2 > 1 || mask_h) { - flash_attn_ext_f16_load_mask - (mask_h, tile_mask, stride_mask, kb0*nbatch_fa, k_VKQ_sup, jt*ncols1, ne01, mask_packed); + flash_attn_ext_f16_load_mask + (mask_h, tile_mask, stride_mask, kb0*nbatch_fa, k_VKQ_sup, jt*ncols1, ne01, mask_packed, nullptr); } if constexpr (raw_KV) { flash_attn_ext_raw_load ((const char *) K_h2 + int64_t(kb0)*nbatch_fa*stride_K, tile_raw, stride_K); } else { - flash_attn_ext_f16_load_tile - (K_h2 + int64_t(kb0)*nbatch_fa*stride_K, tile_K, nbatch_K2, stride_K, k_VKQ_sup); + flash_attn_ext_f16_load_tile + (K_h2, tile_K, nbatch_K2, stride_K, kb0*nbatch_fa, k_VKQ_sup, nullptr); } } - // kb0_start is always < kb0_stop so the last iter can be executed unconditionally. - if constexpr (ncols2 == 1) { + // kb0_start is always < kb0_stop so the last iter can be executed unconditionally. The sparse list's last step is + // partial (n_kv_max a row), so it takes the OOB-checked iterations too. + if constexpr (ncols2 == 1 || use_sparse) { constexpr bool oob_check = true; for (int i = kb0_start; i < kb0_stop-1; ++i) { constexpr bool last_iter = false; constexpr int k_VKQ_sup = nbatch_fa; const int kb0_next = live_bits ? fattn_kv_live_next(live_bits, kb0) : kb0 + 1; flash_attn_ext_f16_iter - - (Q_f2, K_h2, V_h2, mask_h, dstk, dstk_fixup, scale, slope, logit_softcap, + (Q_f2, K_h2, V_h2, mask_h, indices, dstk, dstk_fixup, scale, slope, logit_softcap, ne01, ne02, stride_K, stride_V, stride_mask, mask_packed, tile_Q, tile_K, tile_V, tile_mask, tile_raw, Q_B, Q8, VKQ_C, KQ_max, KQ_rowsum, jt, kb0, kb0_next, k_VKQ_sup); kb0 = kb0_next; @@ -1689,9 +1734,9 @@ static __device__ __forceinline__ void flash_attn_ext_f16_process_tile( constexpr bool last_iter = true; const int k_VKQ_sup = ne11 - kb0*nbatch_fa; flash_attn_ext_f16_iter - - (Q_f2, K_h2, V_h2, mask_h, dstk, dstk_fixup, scale, slope, logit_softcap, + (Q_f2, K_h2, V_h2, mask_h, indices, dstk, dstk_fixup, scale, slope, logit_softcap, ne01, ne02, stride_K, stride_V, stride_mask, mask_packed, tile_Q, tile_K, tile_V, tile_mask, tile_raw, Q_B, Q8, VKQ_C, KQ_max, KQ_rowsum, jt, kb0, kb0 + 1, k_VKQ_sup); } else { @@ -1701,9 +1746,9 @@ static __device__ __forceinline__ void flash_attn_ext_f16_process_tile( constexpr int k_VKQ_sup = nbatch_fa; const int kb0_next = live_bits ? fattn_kv_live_next(live_bits, kb0) : kb0 + 1; flash_attn_ext_f16_iter - - (Q_f2, K_h2, V_h2, mask_h, dstk, dstk_fixup, scale, slope, logit_softcap, + (Q_f2, K_h2, V_h2, mask_h, indices, dstk, dstk_fixup, scale, slope, logit_softcap, ne01, ne02, stride_K, stride_V, stride_mask, mask_packed, tile_Q, tile_K, tile_V, tile_mask, tile_raw, Q_B, Q8, VKQ_C, KQ_max, KQ_rowsum, jt, kb0, kb0_next, k_VKQ_sup); kb0 = kb0_next; @@ -1711,9 +1756,9 @@ static __device__ __forceinline__ void flash_attn_ext_f16_process_tile( constexpr bool last_iter = true; constexpr int k_VKQ_sup = nbatch_fa; flash_attn_ext_f16_iter - - (Q_f2, K_h2, V_h2, mask_h, dstk, dstk_fixup, scale, slope, logit_softcap, + (Q_f2, K_h2, V_h2, mask_h, indices, dstk, dstk_fixup, scale, slope, logit_softcap, ne01, ne02, stride_K, stride_V, stride_mask, mask_packed, tile_Q, tile_K, tile_V, tile_mask, tile_raw, Q_B, Q8, VKQ_C, KQ_max, KQ_rowsum, jt, kb0, kb0 + 1, k_VKQ_sup); } @@ -1896,77 +1941,79 @@ static __device__ __forceinline__ void flash_attn_ext_f16_process_tile( } } - if (np > 1 && threadIdx.y % np == 0) { - // Combine the meta data for parallel warps via shared memory. - // Warps with threadIdx.y % np != 0 must NOT return early. - // All threads must return simultaneously to avoid race conditions with work on the next tile. - + if (np > 1) { + // Combine the meta data for parallel warps via shared memory. Every warp reaches the one __syncthreads() below: + // two barriers in the branches of a warp-divergent condition are a divergent barrier (compute-sanitizer + // synccheck), whose ordering the hardware does not promise. constexpr int nmeta = np*cols_per_warp >= warp_size ? np*cols_per_warp/warp_size : 1; + float KQ_cmn; + float KQ_cms[nmeta]; + float KQ_crs; + const int jc_meta = threadIdx.y*cols_per_warp + (np*cols_per_warp < warp_size ? threadIdx.x % (np*cols_per_warp) : threadIdx.x); float2 * const meta_ptr = ((float2 *) tile_Q) + jc_meta*(tile_stride/2) + nbatch_combine/2; - float2 meta[nmeta]; + + if (threadIdx.y % np == 0) { + float2 meta[nmeta]; #pragma unroll - for (int imeta = 0; imeta < nmeta; ++imeta) { - meta[imeta] = meta_ptr[imeta * warp_size * tile_stride/2]; - } + for (int imeta = 0; imeta < nmeta; ++imeta) { + meta[imeta] = meta_ptr[imeta * warp_size * tile_stride/2]; + } - float KQ_cmn = meta[0].x; // KQ combine max new, max between all parallel warps. + KQ_cmn = meta[0].x; // KQ combine max new, max between all parallel warps. #pragma unroll - for (int imeta = 1; imeta < nmeta; ++imeta) { - KQ_cmn = fmaxf(KQ_cmn, meta[imeta].x); - } + for (int imeta = 1; imeta < nmeta; ++imeta) { + KQ_cmn = fmaxf(KQ_cmn, meta[imeta].x); + } #pragma unroll - for (int offset = np*cols_per_warp/2; offset >= cols_per_warp; offset >>= 1) { - if (offset < warp_size) { - KQ_cmn = fmaxf(KQ_cmn, __shfl_xor_sync(0xFFFFFFFF, KQ_cmn, offset, warp_size)); + for (int offset = np*cols_per_warp/2; offset >= cols_per_warp; offset >>= 1) { + if (offset < warp_size) { + KQ_cmn = fmaxf(KQ_cmn, __shfl_xor_sync(0xFFFFFFFF, KQ_cmn, offset, warp_size)); + } } - } - float KQ_cms[nmeta]; // KQ combine max scale per warp. #pragma unroll - for (int imeta = 0; imeta < nmeta; ++imeta) { - KQ_cms[imeta] = expf(meta[imeta].x - KQ_cmn); - } + for (int imeta = 0; imeta < nmeta; ++imeta) { + KQ_cms[imeta] = expf(meta[imeta].x - KQ_cmn); // KQ combine max scale per warp. + } - float KQ_crs = KQ_cms[0]*meta[0].y; // KQ combine rowsum, scaled sum of all parallel warps. + KQ_crs = KQ_cms[0]*meta[0].y; // KQ combine rowsum, scaled sum of all parallel warps. #pragma unroll - for (int imeta = 1; imeta < nmeta; ++imeta) { - KQ_crs += KQ_cms[imeta]*meta[imeta].y; - } + for (int imeta = 1; imeta < nmeta; ++imeta) { + KQ_crs += KQ_cms[imeta]*meta[imeta].y; + } #pragma unroll - for (int offset = np*cols_per_warp/2; offset >= cols_per_warp; offset >>= 1) { - if (offset < warp_size) { - KQ_crs += __shfl_xor_sync(0xFFFFFFFF, KQ_crs, offset, warp_size); + for (int offset = np*cols_per_warp/2; offset >= cols_per_warp; offset >>= 1) { + if (offset < warp_size) { + KQ_crs += __shfl_xor_sync(0xFFFFFFFF, KQ_crs, offset, warp_size); + } } } __syncthreads(); - // Write back combined meta data: + if (threadIdx.y % np == 0) { + // Write back combined meta data: #pragma unroll - for (int imeta = 0; imeta < nmeta; ++imeta) { - if (np*cols_per_warp >= warp_size || threadIdx.x < np*cols_per_warp) { - // Combined KQ max scale + rowsum. - meta_ptr[imeta * warp_size * tile_stride/2] = make_float2(KQ_cms[imeta], KQ_crs); + for (int imeta = 0; imeta < nmeta; ++imeta) { + if (np*cols_per_warp >= warp_size || threadIdx.x < np*cols_per_warp) { + // Combined KQ max scale + rowsum. + meta_ptr[imeta * warp_size * tile_stride/2] = make_float2(KQ_cms[imeta], KQ_crs); + } } - } - // Combined KQ max + rowsum. - static_assert(cols_per_warp <= warp_size); - if (needs_fixup && (cols_per_warp == warp_size || threadIdx.x < cols_per_warp)) { - float2 * dstk_fixup_meta = dstk_fixup + blockIdx.x*ncols; - dstk_fixup_meta[(threadIdx.y/np)*cols_per_warp + threadIdx.x] = make_float2(KQ_cmn, KQ_crs); - } - if (is_fixup && (cols_per_warp == warp_size || threadIdx.x < cols_per_warp)) { - float2 * dstk_fixup_meta = dstk_fixup + (gridDim.x + blockIdx.x)*ncols; - dstk_fixup_meta[(threadIdx.y/np)*cols_per_warp + threadIdx.x] = make_float2(KQ_cmn, KQ_crs); + // Combined KQ max + rowsum. + static_assert(cols_per_warp <= warp_size); + if (needs_fixup && (cols_per_warp == warp_size || threadIdx.x < cols_per_warp)) { + float2 * dstk_fixup_meta = dstk_fixup + blockIdx.x*ncols; + dstk_fixup_meta[(threadIdx.y/np)*cols_per_warp + threadIdx.x] = make_float2(KQ_cmn, KQ_crs); + } + if (is_fixup && (cols_per_warp == warp_size || threadIdx.x < cols_per_warp)) { + float2 * dstk_fixup_meta = dstk_fixup + (gridDim.x + blockIdx.x)*ncols; + dstk_fixup_meta[(threadIdx.y/np)*cols_per_warp + threadIdx.x] = make_float2(KQ_cmn, KQ_crs); + } } - } else if (np > 1) { - // Warps with threadIdx.y % np == 0 execute a __syncthreads() in the if branch. - // Therefore, all other warps also need to execute a __syncthreads(). - // Otherwise the points at which warps synchronize with each other would become misaligned. - __syncthreads(); } #pragma unroll @@ -2101,7 +2148,7 @@ static __device__ __forceinline__ void flash_attn_ext_f16_process_tile( } } #else - GGML_UNUSED_VARS(Q_f2, K_h2, V_h2, mask_h, sinks_f, dstk, dstk_fixup, + GGML_UNUSED_VARS(Q_f2, K_h2, V_h2, mask_h, indices, sinks_f, dstk, dstk_fixup, scale, slope, logit_softcap, ne01, ne02, gqa_ratio, stride_Q1, stride_Q2, stride_K, stride_V, stride_mask, mask_packed, jt, kb0_start, kb0_stop, live_bits, nwords, needs_fixup, is_fixup); @@ -2109,7 +2156,17 @@ static __device__ __forceinline__ void flash_attn_ext_f16_process_tile( #endif // defined(VOLTA_MMA_AVAILABLE) || defined(TURING_MMA_AVAILABLE) || defined(AMD_WMMA_AVAILABLE) || defined(AMD_MFMA_AVAILABLE) } -template +// The instances with a sparse variant (ggml_flash_attn_ext_set_n_kv_max): one token a tile, DeepSeek's and GLM's MLA heads. +static constexpr __host__ __device__ bool ggml_cuda_flash_attn_ext_mma_f16_may_use_sparse( + const int DKQ, const int DV, const int ncols1, const int ncols2) { + return (DKQ == 512 && DV == 512 && ncols1 == 1 && ncols2 == 8) || + (DKQ == 576 && DV == 512 && ncols1 == 1 && ncols2 == 16); +} + +// use_sparse: KV_max_ptr holds each mask row's n_kv_max cell indices (ggml_cuda_flash_attn_ext_compact_mask) and ne11 is +// n_kv_max; the kernel steps over those cells in place of the cache. +template __launch_bounds__(ggml_cuda_fattn_mma_get_nthreads(DKQ, DV, ncols1*ncols2), ggml_cuda_fattn_mma_get_occupancy(DKQ, DV, ncols1*ncols2)) static __global__ void flash_attn_ext_f16( const char * Q_ptr, @@ -2142,16 +2199,25 @@ static __global__ void flash_attn_ext_f16( const char * GGML_CUDA_RESTRICT V = V_ptr; const char * GGML_CUDA_RESTRICT mask = mask_ptr; const char * GGML_CUDA_RESTRICT sinks = sinks_ptr; - const int * GGML_CUDA_RESTRICT KV_max = KV_max_ptr; + const int * GGML_CUDA_RESTRICT KV_max = use_sparse ? nullptr : KV_max_ptr; const int * GGML_CUDA_RESTRICT KV_live = KV_live_ptr; float * GGML_CUDA_RESTRICT dst = dst_ptr; float2 * GGML_CUDA_RESTRICT dst_meta = dst_meta_ptr; + const int32_t * GGML_CUDA_RESTRICT sparse_indices = use_sparse ? KV_max_ptr : nullptr; + // Skip unused kernel variants for faster compilation: if (use_logit_softcap && !(DKQ == 128 || DKQ == 256 || DKQ == 512)) { NO_DEVICE_CODE; return; } + // the sparse variant only where the host launches it, so no other instance compiles its code + constexpr bool use_sparse_arch = use_sparse && !use_logit_softcap && type_K == GGML_TYPE_F16 && + ggml_cuda_flash_attn_ext_mma_f16_may_use_sparse(DKQ, DV, ncols1, ncols2); + if (use_sparse != use_sparse_arch) { + NO_DEVICE_CODE; + return; + } if (DKQ == 192 && ncols2 != 8 && ncols2 != 16) { NO_DEVICE_CODE; return; @@ -2270,6 +2336,7 @@ static __global__ void flash_attn_ext_f16( const half2 * V_h2 = V_is_K_view ? K_h2 : (const half2 *) (V + nb23*sequence + nb22*z_KV); const float * sinks_f = sinks ? (const float *) sinks + zt_Q : nullptr; + const int32_t * indices = use_sparse_arch ? sparse_indices + (int64_t(sequence % ne33)*ne31 + jt*ncols1)*ne11 : nullptr; const float slope = ncols2 == 1 ? get_alibi_slope(max_bias, zt_Q, n_head_log2, m0, m1) : 1.0f; @@ -2284,8 +2351,9 @@ static __global__ void flash_attn_ext_f16( // a tile it did not start writes to dst and needs the fixup; one that did the whole tile writes its final result. const bool is_fixup = last_tile; const bool needs_fixup = !last_tile && kb0_start != 0; - flash_attn_ext_f16_process_tile - (Q_f2, K_h2, V_h2, mask_h, sinks_f, dstk, dst_meta, scale, slope, logit_softcap, + flash_attn_ext_f16_process_tile + (Q_f2, K_h2, V_h2, mask_h, indices, sinks_f, dstk, dst_meta, scale, slope, logit_softcap, ne01, ne02, gqa_ratio, ne11, stride_Q1, stride_Q2, stride_K, stride_V, stride_mask, mask_packed != 0, jt, zt_gqa, kb0_live, kb0_stop, live_bits, live.nwords, needs_fixup, is_fixup); } @@ -2303,6 +2371,10 @@ static __global__ void flash_attn_ext_f16( #endif // defined(FLASH_ATTN_AVAILABLE) && (defined(VOLTA_MMA_AVAILABLE) || defined(TURING_MMA_AVAILABLE) || defined(AMD_WMMA_AVAILABLE) || defined(AMD_MFMA_AVAILABLE)) } +// Whether dst runs the sparse variant: n_kv_max set, an f16 mask of one head, no ALiBi or softcap, and a cache at least +// 4096 cells and as many as the gather reads for the queries one dense pass covers. Defined in fattn.cu. +bool ggml_cuda_flash_attn_ext_mma_f16_shall_use_sparse(ggml_backend_cuda_context & ctx, const ggml_tensor * dst, int ncols2); + template void ggml_cuda_flash_attn_ext_mma_f16_case(ggml_backend_cuda_context & ctx, ggml_tensor * dst) { const ggml_tensor * KQV = dst; @@ -2350,20 +2422,38 @@ void ggml_cuda_flash_attn_ext_mma_f16_case(ggml_backend_cuda_context & ctx, ggml using fattn_kernel_ptr_t = fattn_kernel_t; #endif // defined(GGML_USE_HIP) fattn_kernel_t fattn_kernel; + bool use_sparse = false; if (logit_softcap == 0.0f) { constexpr bool use_logit_softcap = false; - fattn_kernel = flash_attn_ext_f16; +#if !defined(GGML_USE_HIP) && !defined(GGML_USE_MUSA) + if constexpr (!raw_KV && ggml_cuda_flash_attn_ext_mma_f16_may_use_sparse(DKQ, DV, ncols1, ncols2)) { + if (ggml_cuda_flash_attn_ext_mma_f16_shall_use_sparse(ctx, dst, ncols2)) { + GGML_ASSERT(nstages <= 1); // the shared memory sized here is the device's single stage + use_sparse = true; + fattn_kernel = flash_attn_ext_f16; + + static bool shared_memory_limit_raised[GGML_CUDA_MAX_DEVICES] = {false}; + if (!shared_memory_limit_raised[id]) { + CUDA_CHECK(cudaFuncSetAttribute(reinterpret_cast(fattn_kernel), cudaFuncAttributeMaxDynamicSharedMemorySize, nbytes_shared_total)); + shared_memory_limit_raised[id] = true; + } + } + } +#endif // !defined(GGML_USE_HIP) && !defined(GGML_USE_MUSA) + if (!use_sparse) { + fattn_kernel = flash_attn_ext_f16; #if !defined(GGML_USE_MUSA) - static bool shared_memory_limit_raised[GGML_CUDA_MAX_DEVICES] = {false}; - if (!shared_memory_limit_raised[id]) { - CUDA_CHECK(cudaFuncSetAttribute(reinterpret_cast(fattn_kernel), cudaFuncAttributeMaxDynamicSharedMemorySize, nbytes_shared_total)); - shared_memory_limit_raised[id] = true; - } + static bool shared_memory_limit_raised[GGML_CUDA_MAX_DEVICES] = {false}; + if (!shared_memory_limit_raised[id]) { + CUDA_CHECK(cudaFuncSetAttribute(reinterpret_cast(fattn_kernel), cudaFuncAttributeMaxDynamicSharedMemorySize, nbytes_shared_total)); + shared_memory_limit_raised[id] = true; + } #endif // !defined(GGML_USE_MUSA) + } } else { constexpr bool use_logit_softcap = true; - fattn_kernel = flash_attn_ext_f16; + fattn_kernel = flash_attn_ext_f16; #if !defined(GGML_USE_MUSA) static bool shared_memory_limit_raised[GGML_CUDA_MAX_DEVICES] = {false}; @@ -2375,7 +2465,7 @@ void ggml_cuda_flash_attn_ext_mma_f16_case(ggml_backend_cuda_context & ctx, ggml } launch_fattn - (ctx, dst, fattn_kernel, nwarps, nbytes_shared_total, nbatch_fa, !raw_KV, !raw_KV, true, warp_size_host, true); + (ctx, dst, fattn_kernel, nwarps, nbytes_shared_total, nbatch_fa, !raw_KV, !raw_KV, true, warp_size_host, true, use_sparse); } diff --git a/ggml/src/ggml-cuda/fattn.cu b/ggml/src/ggml-cuda/fattn.cu index 70233bfa11ee..a68e52caef61 100644 --- a/ggml/src/ggml-cuda/fattn.cu +++ b/ggml/src/ggml-cuda/fattn.cu @@ -5,11 +5,286 @@ #include "fattn-vec.cuh" #include "fattn.cuh" +#if !defined(GGML_USE_HIP) && !defined(GGML_USE_MUSA) +// 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). +// 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. +__launch_bounds__(256, 1) +static __global__ void flash_attn_mask_to_sparse_indices_legacy( + const half * mask_ptr, int32_t * indices_ptr, const int ne30, const int n_kv_max, + const int64_t s31, const int64_t s33) { + ggml_cuda_pdl_sync(); + + constexpr int values_per_lane = 8; + const int tid = threadIdx.x; + const int warp = tid / WARP_SIZE; + const int lane = tid % WARP_SIZE; + const int sequence = blockIdx.y; + const int query = blockIdx.x; + + const half * mask = mask_ptr + sequence*s33 + query*s31; + int32_t * indices = indices_ptr + (int64_t(sequence)*gridDim.x + query)*n_kv_max; + + __shared__ int warp_offsets[256/WARP_SIZE]; + __shared__ int row_count; + __shared__ int chunk_count; + + if (tid == 0) { + row_count = 0; + } + __syncthreads(); + + for (int i0 = 0; i0 < ne30; i0 += blockDim.x*values_per_lane) { + uint32_t selected_warp[values_per_lane]; + int warp_count = 0; +#pragma unroll + for (int item = 0; item < values_per_lane; ++item) { + const int i = i0 + (warp*values_per_lane + item)*WARP_SIZE + lane; + const bool selected = i < ne30 && isfinite(__half2float(mask[i])); + selected_warp[item] = __ballot_sync(0xFFFFFFFF, selected); + warp_count += __popc(selected_warp[item]); + } + + if (lane == 0) { + warp_offsets[warp] = warp_count; + } + __syncthreads(); + + if (tid == 0) { + int offset = 0; +#pragma unroll + for (int iw = 0; iw < 256/WARP_SIZE; ++iw) { + const int count = warp_offsets[iw]; + warp_offsets[iw] = offset; + offset += count; + } + chunk_count = offset; + } + __syncthreads(); + + const uint32_t lane_mask = lane == 0 ? 0 : (1u << lane) - 1; + int warp_item_offset = 0; +#pragma unroll + for (int item = 0; item < values_per_lane; ++item) { + const int i = i0 + (warp*values_per_lane + item)*WARP_SIZE + lane; + const int dst = row_count + warp_offsets[warp] + warp_item_offset + __popc(selected_warp[item] & lane_mask); + if ((selected_warp[item] & (uint32_t(1) << lane)) && dst < n_kv_max) { + indices[dst] = i; + } + warp_item_offset += __popc(selected_warp[item]); + } + __syncthreads(); + + if (tid == 0) { + row_count += chunk_count; + } + __syncthreads(); + } + + const int count = row_count; + for (int i = count + tid; i < n_kv_max; i += blockDim.x) { + indices[i] = -1; + } + __syncthreads(); + + // the dependent grid reads indices, signal once the row is complete + ggml_cuda_pdl_lc(); +} + +// The same lists, one block of 1024 threads a row: a thread tests 32 consecutive columns from four 16-byte loads issued +// together, so up to 32768 columns are one round trip, and a block-wide scan of the threads' counts places each +// thread's cells after the lower columns'. The row is 16-byte aligned and ne30 % 8 == 0. +static constexpr int sparse_scan_threads = 1024; +static constexpr int sparse_scan_cols = 32; + +// NROWS: how many consecutive query rows share one index list. 1 is the per-query list the ncols1 == 1 kernel reads; with +// NROWS > 1 a block emits the UNION of its rows' cells, so a tile of NROWS queries gathers each cell once instead of once +// per query. Whether that pays is entirely a question of how much adjacent queries' selections overlap, which is a +// property of the TRAINED indexer and cannot be measured on a random-weight proxy: see heads/glm-5.3-flash/evidence.md. +// Correctness does not depend on the overlap -- flash_attn_ext_f16_load_mask reads each row's own mask value at every +// gathered cell, so a cell a row does not select still reads -inf for that row. +template +__launch_bounds__(sparse_scan_threads, 1) +static __global__ void flash_attn_mask_to_sparse_indices( + const half * mask_ptr, int32_t * indices_ptr, const int ne30, const int n_kv_max, + const int64_t s31, const int64_t s33, const int ne31) { + ggml_cuda_pdl_sync(); + + constexpr int nwarps = sparse_scan_threads/WARP_SIZE; + static_assert(nwarps == WARP_SIZE, "one warp scans the warps' sums"); + const int tid = threadIdx.x; + const int warp = tid / WARP_SIZE; + const int lane = tid % WARP_SIZE; + const int sequence = blockIdx.y; + const int tile = blockIdx.x; + + const half * mask = mask_ptr + sequence*s33 + int64_t(tile)*NROWS*s31; + int32_t * indices = indices_ptr + (int64_t(sequence)*gridDim.x + tile)*n_kv_max; + + // the last tile is short when ne31 % NROWS != 0; its missing rows contribute nothing to the union + const int nrows = NROWS == 1 ? 1 : min(NROWS, ne31 - tile*NROWS); + + __shared__ int warp_sums[nwarps]; + int row_count = 0; + + for (int i0 = 0; i0 < ne30; i0 += sparse_scan_threads*sparse_scan_cols) { + const int c0 = i0 + tid*sparse_scan_cols; + + // bit c: column c0 + c is finite in ANY of the tile's rows (a half is inf or nan exactly when its exponent bits + // are all set). For NROWS == 1 this is one row's own mask and the loop folds away. + uint32_t finite = 0; + for (int r = 0; r < nrows; ++r) { + // 8 halves a load, -inf past the row (ne30 % 8 == 0: a load is all in or all out) + uint4 v[sparse_scan_cols/8]; +#pragma unroll + for (int j = 0; j < sparse_scan_cols/8; ++j) { + v[j] = c0 + 8*j < ne30 ? ((const uint4 *) (mask + r*s31 + c0))[j] : make_uint4(0xFC00FC00u, 0xFC00FC00u, 0xFC00FC00u, 0xFC00FC00u); + } + +#pragma unroll + for (int j = 0; j < sparse_scan_cols/8; ++j) { + const uint32_t w[4] = {v[j].x, v[j].y, v[j].z, v[j].w}; +#pragma unroll + for (int k = 0; k < 4; ++k) { + finite |= uint32_t((w[k] & 0x00007C00u) != 0x00007C00u) << (8*j + 2*k + 0); + finite |= uint32_t((w[k] & 0x7C000000u) != 0x7C000000u) << (8*j + 2*k + 1); + } + } + } + const int count = __popc(finite); + + // inclusive scan of the counts in the warp, then of the warps' sums + int incl = count; +#pragma unroll + for (int offset = 1; offset < WARP_SIZE; offset *= 2) { + const int t = __shfl_up_sync(0xFFFFFFFF, incl, offset); + incl += lane >= offset ? t : 0; + } + if (lane == WARP_SIZE - 1) { + warp_sums[warp] = incl; + } + __syncthreads(); + if (warp == 0) { + int s = warp_sums[lane]; +#pragma unroll + for (int offset = 1; offset < WARP_SIZE; offset *= 2) { + const int t = __shfl_up_sync(0xFFFFFFFF, s, offset); + s += lane >= offset ? t : 0; + } + warp_sums[lane] = s; + } + __syncthreads(); + + int dst = row_count + (warp == 0 ? 0 : warp_sums[warp - 1]) + incl - count; + while (finite != 0 && dst < n_kv_max) { + indices[dst++] = c0 + __ffs(finite) - 1; + finite &= finite - 1; + } + row_count += warp_sums[nwarps - 1]; + __syncthreads(); // the next chunk rewrites warp_sums + } + + for (int i = row_count + tid; i < n_kv_max; i += sparse_scan_threads) { + indices[i] = -1; + } + __syncthreads(); + + // the dependent grid reads indices, signal once the row is complete + ggml_cuda_pdl_lc(); +} +#endif // !defined(GGML_USE_HIP) && !defined(GGML_USE_MUSA) + +void ggml_cuda_flash_attn_ext_compact_mask( + const ggml_tensor * mask, int32_t * indices, int32_t n_kv_max, int rows_per_list, cudaStream_t stream) { +#if defined(GGML_USE_HIP) || defined(GGML_USE_MUSA) + GGML_UNUSED_VARS(mask, indices, n_kv_max, rows_per_list, stream); + GGML_ABORT("sparse flash attention is only supported on NVIDIA CUDA"); +#else + GGML_ASSERT(mask->type == GGML_TYPE_F16); + static const bool scan_legacy = ggml_env_switch("GGML_CUDA_FATTN_SPARSE_SCAN_LEGACY"); + const int64_t s31 = mask->nb[1] / sizeof(half); + const int64_t s33 = mask->nb[3] / sizeof(half); + const bool aligned = (uintptr_t) mask->data % 16 == 0 && mask->ne[0] % 8 == 0 && s31 % 8 == 0 && s33 % 8 == 0; + // the legacy scan has no union form, so a tiled list needs the wide one + const bool wide = (!scan_legacy || rows_per_list > 1) && aligned; + GGML_ASSERT(rows_per_list == 1 || wide); + const int64_t n_lists = (mask->ne[1] + rows_per_list - 1) / rows_per_list; + const dim3 blocks_num(n_lists, mask->ne[3], 1); + const dim3 block_dim(wide ? sparse_scan_threads : 256, 1, 1); + const ggml_cuda_kernel_launch_params launch_params(blocks_num, block_dim, 0, stream); + if (rows_per_list == 1 && !wide) { + // the legacy scan is one row a block and takes no row count + ggml_cuda_kernel_launch(flash_attn_mask_to_sparse_indices_legacy, launch_params, + (const half *) mask->data, indices, int(mask->ne[0]), n_kv_max, s31, s33); + } else if (rows_per_list == 1) { + ggml_cuda_kernel_launch(flash_attn_mask_to_sparse_indices<1>, launch_params, + (const half *) mask->data, indices, int(mask->ne[0]), n_kv_max, s31, s33, int(mask->ne[1])); + } else { + // only the widths the sparse kernel is instantiated for; ggml_cuda_fattn_sparse_rows_per_list keeps them in step + switch (rows_per_list) { + case 2: ggml_cuda_kernel_launch(flash_attn_mask_to_sparse_indices< 2>, launch_params, + (const half *) mask->data, indices, int(mask->ne[0]), n_kv_max, s31, s33, int(mask->ne[1])); break; + case 4: ggml_cuda_kernel_launch(flash_attn_mask_to_sparse_indices< 4>, launch_params, + (const half *) mask->data, indices, int(mask->ne[0]), n_kv_max, s31, s33, int(mask->ne[1])); break; + case 8: ggml_cuda_kernel_launch(flash_attn_mask_to_sparse_indices< 8>, launch_params, + (const half *) mask->data, indices, int(mask->ne[0]), n_kv_max, s31, s33, int(mask->ne[1])); break; + case 16: ggml_cuda_kernel_launch(flash_attn_mask_to_sparse_indices<16>, launch_params, + (const half *) mask->data, indices, int(mask->ne[0]), n_kv_max, s31, s33, int(mask->ne[1])); break; + default: GGML_ABORT("unsupported rows_per_list %d", rows_per_list); + } + } + CUDA_CHECK(cudaGetLastError()); +#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. +// 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 +// gave up a verify's 3 tokens at 8K cells, 0.73x, and a prefill at 16K, 0.63-0.72x). +bool ggml_cuda_flash_attn_ext_mma_f16_shall_use_sparse(ggml_backend_cuda_context & ctx, const ggml_tensor * dst, const int ncols2) { +#if defined(GGML_USE_HIP) || defined(GGML_USE_MUSA) + GGML_UNUSED_VARS(ctx, dst, ncols2); + return false; +#else + static const bool sparse_legacy = ggml_env_switch("GGML_CUDA_FATTN_SPARSE_LEGACY"); + + const ggml_tensor * Q = dst->src[0]; + const ggml_tensor * K = dst->src[1]; + const ggml_tensor * mask = dst->src[3]; + const int cc = ggml_cuda_info().devices[ctx.device].cc; + + float max_bias = 0.0f; + float logit_softcap = 0.0f; + memcpy(&max_bias, (const float *) dst->op_params + 1, sizeof(float)); + memcpy(&logit_softcap, (const float *) dst->op_params + 2, sizeof(float)); + + const int32_t n_kv_max = ggml_flash_attn_ext_get_n_kv_max(dst); + const int64_t n_gather = std::min(Q->ne[1], 64/ncols2) * n_kv_max; + return !sparse_legacy && GGML_CUDA_CC_IS_NVIDIA(cc) && turing_mma_available(cc) && + mask != nullptr && mask->type == GGML_TYPE_F16 && n_kv_max > 0 && max_bias == 0.0f && logit_softcap == 0.0f && + mask->ne[0] == K->ne[1] && mask->ne[1] >= Q->ne[1] && mask->ne[2] == 1 && + K->ne[1] >= std::max(4096, n_gather); +#endif // !defined(GGML_USE_HIP) && !defined(GGML_USE_MUSA) +} + template static void ggml_cuda_flash_attn_ext_mma_f16_switch_ncols1(ggml_backend_cuda_context & ctx, ggml_tensor * dst) { const int cc = ggml_cuda_info().devices[ggml_cuda_get_device()].cc; const ggml_tensor * Q = dst->src[0]; +#if !defined(GGML_USE_HIP) && !defined(GGML_USE_MUSA) + // the sparse variant is one token a tile, whatever the batch + if constexpr (ggml_cuda_flash_attn_ext_mma_f16_may_use_sparse(DKQ, DV, 1, ncols2)) { + if (ggml_cuda_flash_attn_ext_mma_f16_shall_use_sparse(ctx, dst, ncols2)) { + ggml_cuda_flash_attn_ext_mma_f16_case(ctx, dst); + return; + } + } +#endif // !defined(GGML_USE_HIP) && !defined(GGML_USE_MUSA) + if constexpr (ncols2 <= 8) { if (turing_mma_available(cc) && Q->ne[1] <= 8/ncols2) { ggml_cuda_flash_attn_ext_mma_f16_case(ctx, dst); diff --git a/ggml/src/ggml-cuda/fwht.cu b/ggml/src/ggml-cuda/fwht.cu index 4b4102080686..7f859bce5477 100644 --- a/ggml/src/ggml-cuda/fwht.cu +++ b/ggml/src/ggml-cuda/fwht.cu @@ -285,7 +285,7 @@ template __launch_bounds__(1024, 1) __global__ void rms_norm_fwht_cuda(const float * x, const float * w, const float * signs, float * normed, float * dst, const int ncols, const float eps, const float scale, const bool pdl_trigger, - const bool prewait) { + const bool prewait, unsigned int * release_to) { if (pdl_trigger) { ggml_cuda_pdl_lc(); } @@ -341,6 +341,15 @@ __global__ void rms_norm_fwht_cuda(const float * x, const float * w, const float fwht_block_transform(reg, s); fwht_block_store(reg, dst + e0); + + // release_to: the alpha/beta fold's handoff (ggml_cuda_ssm_conv_ab_slots), for a conv that reads normed before its + // dependency wait; the barrier puts every thread's normed stores before thread 0's release + if (release_to != nullptr) { + __syncthreads(); + if (tid == 0) { + ggml_cuda_red_release_add(release_to, 1); + } + } } static bool fwht_legacy() { @@ -532,8 +541,10 @@ bool ggml_cuda_rms_norm_fwht_supported(const ggml_tensor * rms_norm, const ggml_ (normed == nullptr || (normed->type == GGML_TYPE_F32 && ggml_is_contiguous(normed))); } -void ggml_cuda_op_rms_norm_fwht(ggml_backend_cuda_context & ctx, const ggml_tensor * rms_norm, const ggml_tensor * w, - ggml_tensor * normed, const ggml_tensor * signs, ggml_tensor * dst) { +int ggml_cuda_op_rms_norm_fwht(ggml_backend_cuda_context & ctx, const ggml_tensor * rms_norm, const ggml_tensor * w, + ggml_tensor * normed, const ggml_tensor * signs, ggml_tensor * dst, + unsigned int * release_to) { + GGML_ASSERT(release_to == nullptr || normed != nullptr); GGML_ASSERT(ggml_cuda_rms_norm_fwht_supported(rms_norm, w, normed, signs, dst)); const ggml_tensor * x = rms_norm->src[0]; const int ncols = x->ne[0]; @@ -557,13 +568,14 @@ void ggml_cuda_op_rms_norm_fwht(ggml_backend_cuda_context & ctx, const ggml_tens const ggml_cuda_kernel_launch_params lp = ggml_cuda_kernel_launch_params(grid, block, 0, ctx.stream()); switch (n) { case 1024: - ggml_cuda_kernel_launch(rms_norm_fwht_cuda<1024>, lp, x_d, w_d, signs_d, normed_d, dst_d, ncols, eps, scale, pdl_trigger, prewait); + ggml_cuda_kernel_launch(rms_norm_fwht_cuda<1024>, lp, x_d, w_d, signs_d, normed_d, dst_d, ncols, eps, scale, pdl_trigger, prewait, release_to); break; case 2048: - ggml_cuda_kernel_launch(rms_norm_fwht_cuda<2048>, lp, x_d, w_d, signs_d, normed_d, dst_d, ncols, eps, scale, pdl_trigger, prewait); + ggml_cuda_kernel_launch(rms_norm_fwht_cuda<2048>, lp, x_d, w_d, signs_d, normed_d, dst_d, ncols, eps, scale, pdl_trigger, prewait, release_to); break; default: - ggml_cuda_kernel_launch(rms_norm_fwht_cuda<4096>, lp, x_d, w_d, signs_d, normed_d, dst_d, ncols, eps, scale, pdl_trigger, prewait); + ggml_cuda_kernel_launch(rms_norm_fwht_cuda<4096>, lp, x_d, w_d, signs_d, normed_d, dst_d, ncols, eps, scale, pdl_trigger, prewait, release_to); break; } + return (int) (grid.x * grid.y); } diff --git a/ggml/src/ggml-cuda/fwht.cuh b/ggml/src/ggml-cuda/fwht.cuh index 11a2fa77eae2..1ba6aa8ad030 100644 --- a/ggml/src/ggml-cuda/fwht.cuh +++ b/ggml/src/ggml-cuda/fwht.cuh @@ -10,8 +10,10 @@ bool ggml_cuda_op_fwht_view(ggml_backend_cuda_context & ctx, const ggml_tensor * ggml_tensor * dst); // rms_norm (eps from rms_norm), its weight multiply by w, the sign flip and the transform into dst in one launch, at up to 8 tokens. -// normed: the multiply's result when other nodes read it, else nullptr. +// normed: the multiply's result when other nodes read it, else nullptr. release_to: a handoff slot each block releases to +// once its normed values are stored (ggml_cuda_ssm_conv_ab_slots), or nullptr. Returns the blocks it launched. bool ggml_cuda_rms_norm_fwht_supported(const ggml_tensor * rms_norm, const ggml_tensor * w, const ggml_tensor * normed, const ggml_tensor * signs, const ggml_tensor * dst); -void ggml_cuda_op_rms_norm_fwht(ggml_backend_cuda_context & ctx, const ggml_tensor * rms_norm, const ggml_tensor * w, - ggml_tensor * normed, const ggml_tensor * signs, ggml_tensor * dst); +int ggml_cuda_op_rms_norm_fwht(ggml_backend_cuda_context & ctx, const ggml_tensor * rms_norm, const ggml_tensor * w, + ggml_tensor * normed, const ggml_tensor * signs, ggml_tensor * dst, + unsigned int * release_to = nullptr); diff --git a/ggml/src/ggml-cuda/gated_delta_net.cu b/ggml/src/ggml-cuda/gated_delta_net.cu index b9a6c8c4b888..a0d6f3189ca4 100644 --- a/ggml/src/ggml-cuda/gated_delta_net.cu +++ b/ggml/src/ggml-cuda/gated_delta_net.cu @@ -9,6 +9,23 @@ static __global__ void gdn_precompute_exp(const float * g, float * g_exp, int64_ } } +// KDA's per-channel decay once per token, head and channel: exp(g), or with raw gates exp(raw_lb * sigmoid(-(g * +// raw_a[h]))), g contiguous [S_v, H, tokens, seqs]. In the recurrent kernel each warp owns one state column, so each +// warp of a head computed all S_v of a token's decays, every value S_v times over, two expf and a division an element. +// The formulas are the kernel's own, so the values are bit for bit. +template +static __global__ void gdn_kda_precompute_decay(const float * g, const float * raw_a, const float raw_lb, float * g_exp, + const int64_t n, const int64_t S_v, const int64_t H) { + for (int64_t i = (int64_t) blockIdx.x*blockDim.x + threadIdx.x; i < n; + i += (int64_t) blockDim.x*gridDim.x) { + float g_i = g[i]; + if constexpr (RAW) { + g_i = raw_lb * (1.0f / (1.0f + expf(g_i * raw_a[i / S_v % H]))); + } + g_exp[i] = expf(g_i); + } +} + // Writes element e of a state destination: f32, or an f16 or q8_0 cache row, with cpy's conversion (f32 -> f16 rounds // to nearest even, as ggml_cuda_cast does), so the cache holds the bytes the unfused gdn -> cpy pair would write. In // q8_0 one block is one warp-wide slice of a state column (QK8_0 == warp_size and every column starts on a block), so @@ -55,10 +72,19 @@ static __device__ __forceinline__ float gdn_load_state(const void * src, const i return ((const float *) src)[e]; } +// The byte offset of element e of a state source (gdn_load_state's types; in q8_0, e starts a block) +static __device__ __forceinline__ int64_t gdn_state_byte(const int64_t e, const ggml_type type) { + if (type == GGML_TYPE_Q8_0) { + return e / QK8_0 * (int64_t) sizeof(block_q8_0); + } + return e * (type == GGML_TYPE_F16 ? (int64_t) sizeof(half) : (int64_t) sizeof(float)); +} + // RAW: beta and g arrive pre-activation (ggml_gated_delta_net_set_raw_gates; with KDA, // ggml_gated_delta_net_set_raw_kda_gates); the kernel applies sigmoid(beta) and raw_a[h] * softplus(g + raw_dt_bias[h]), // or with KDA raw_lb * sigmoid(-(g * raw_a[h])), with the unary kernels' formulas. -// G_PRECOMPUTED: g already holds exp(g) (GB10 long-prompt path); only used with RAW == false. +// G_PRECOMPUTED: g already holds exp(g) (GB10 long-prompt path), only with RAW == false; with KDA, the decay +// gdn_kda_precompute_decay wrote, RAW then activating beta alone. template __global__ void __launch_bounds__((ggml_cuda_get_physical_warp_size() < S_v ? ggml_cuda_get_physical_warp_size() : S_v) * 4, 2) gated_delta_net_cuda(const float * q, @@ -92,7 +118,9 @@ gated_delta_net_cuda(const float * q, int64_t attn_seq_stride, const int32_t * s_ids, int64_t s_row_stride, - ggml_type s_type) { + ggml_type s_type, + bool state_prefetch, + bool lc_early) { const uint32_t h_idx = blockIdx.x; const uint32_t sequence = blockIdx.y; // Each warp owns one or more columns, using warp-level primitives to reduce across rows. @@ -123,6 +151,18 @@ gated_delta_net_cuda(const float * q, float s_shard[cols_per_warp][rows_per_lane]; // state is stored transposed: M[col][i] = S[i][col], row col is contiguous + // With the fused gather the state is a cache row, from the previous step: the block that starts a head requests the + // head's whole state into L2 before the wait, so DRAM streams it under the kernels before this one (qwen35's + // alpha/beta pair, which lets this kernel launch at its start; a hint, ggml_cuda_prefetch_l2). Then this kernel lets + // the next launch, so the gated norm lands under it. Neither changes a result. + if (state_prefetch && s_ids && blockIdx.z == 0 && threadIdx.x == 0 && threadIdx.y == 0) { + const int64_t e0 = (int64_t) s_ids[sequence] * s_row_stride + h_idx * S_v * S_v; + const int64_t b0 = gdn_state_byte(e0, s_type); + ggml_cuda_prefetch_l2((const char *) curr_state + b0, gdn_state_byte(e0 + S_v * S_v, s_type) - b0); + } + if (lc_early) { + ggml_cuda_pdl_lc(); + } ggml_cuda_pdl_sync(); const int64_t state_in_offset = (s_ids ? (int64_t) s_ids[sequence] * s_row_stride : (int64_t) sequence * H * S_v * S_v) + h_idx * S_v * S_v; @@ -199,10 +239,14 @@ gated_delta_net_cuda(const float * q, #pragma unroll for (int r = 0; r < rows_per_lane; r++) { float g_i = g_t[r * warp_size + lane]; - if constexpr (RAW) { - g_i = raw_lb * (1.0f / (1.0f + expf(g_i * raw_a[h_idx]))); + if constexpr (G_PRECOMPUTED) { + g_exp[r] = g_i; + } else { + if constexpr (RAW) { + g_i = raw_lb * (1.0f / (1.0f + expf(g_i * raw_a[h_idx]))); + } + g_exp[r] = expf(g_i); } - g_exp[r] = expf(g_i); } // kv[col] = sum_i g[i] * S[i][col] * k[i] @@ -297,13 +341,18 @@ static void launch_gated_delta_net( launch_params.l2_window = gather->l2_window; // the recurrent states' L2 persistence (ggml_cuda_l2_persist_window) } #endif // defined(GGML_CUDA_USE_L2_WINDOW) + // the state's L2 prefetch and the next launch's trigger before the wait (the kernel's comment); + // GGML_CUDA_GDN_STATE_PREFETCH_LEGACY=1 and GGML_CUDA_GDN_TRIGGER_LEGACY=1 turn them off + static const bool state_prefetch = !ggml_env_switch("GGML_CUDA_GDN_STATE_PREFETCH_LEGACY"); + static const bool lc_early = !ggml_env_switch("GGML_CUDA_GDN_TRIGGER_LEGACY"); switch (S_v) { case 16: if constexpr (STATE_T != GGML_TYPE_Q8_0) { // a q8_0 block is 32 wide: a 16-lane warp cannot own one ggml_cuda_kernel_launch(gated_delta_net_cuda<16, KDA, keep_rs_t, RAW, G_PRECOMPUTED, STATE_T>, launch_params, q_d, k_d, v_d, g_d, b_d, rb_d, ra_d, rlb, s_d, dst_d, state_d, H, n_tokens, n_seqs, sq1, sq2, sq3, sv1, sv2, sv3, - sb1, sb2, sb3, neqk1_magic, rq3_magic, scale, state_slot_stride, K, attn_seq_stride, s_ids, s_row_stride, s_type); + sb1, sb2, sb3, neqk1_magic, rq3_magic, scale, state_slot_stride, K, attn_seq_stride, s_ids, s_row_stride, s_type, + state_prefetch, lc_early); break; } GGML_ABORT("a q8_0 recurrent state needs S_v >= 32"); @@ -311,20 +360,23 @@ static void launch_gated_delta_net( ggml_cuda_kernel_launch(gated_delta_net_cuda<32, KDA, keep_rs_t, RAW, G_PRECOMPUTED, STATE_T>, launch_params, q_d, k_d, v_d, g_d, b_d, rb_d, ra_d, rlb, s_d, dst_d, state_d, H, n_tokens, n_seqs, sq1, sq2, sq3, sv1, sv2, sv3, - sb1, sb2, sb3, neqk1_magic, rq3_magic, scale, state_slot_stride, K, attn_seq_stride, s_ids, s_row_stride, s_type); + sb1, sb2, sb3, neqk1_magic, rq3_magic, scale, state_slot_stride, K, attn_seq_stride, s_ids, s_row_stride, s_type, + state_prefetch, lc_early); break; case 64: { ggml_cuda_kernel_launch(gated_delta_net_cuda<64, KDA, keep_rs_t, RAW, G_PRECOMPUTED, STATE_T>, launch_params, q_d, k_d, v_d, g_d, b_d, rb_d, ra_d, rlb, s_d, dst_d, state_d, H, n_tokens, n_seqs, sq1, sq2, sq3, sv1, sv2, sv3, - sb1, sb2, sb3, neqk1_magic, rq3_magic, scale, state_slot_stride, K, attn_seq_stride, s_ids, s_row_stride, s_type); + sb1, sb2, sb3, neqk1_magic, rq3_magic, scale, state_slot_stride, K, attn_seq_stride, s_ids, s_row_stride, s_type, + state_prefetch, lc_early); break; } case 128: { ggml_cuda_kernel_launch(gated_delta_net_cuda<128, KDA, keep_rs_t, RAW, G_PRECOMPUTED, STATE_T>, launch_params, q_d, k_d, v_d, g_d, b_d, rb_d, ra_d, rlb, s_d, dst_d, state_d, H, n_tokens, n_seqs, sq1, sq2, sq3, sv1, sv2, sv3, - sb1, sb2, sb3, neqk1_magic, rq3_magic, scale, state_slot_stride, K, attn_seq_stride, s_ids, s_row_stride, s_type); + sb1, sb2, sb3, neqk1_magic, rq3_magic, scale, state_slot_stride, K, attn_seq_stride, s_ids, s_row_stride, s_type, + state_prefetch, lc_early); break; } default: @@ -365,12 +417,18 @@ bool ggml_cuda_gdn_chunked_shape_eligible(const ggml_tensor * dst) { const int64_t K = ggml_get_op_params_i32(dst, 0); const bool kda = src_g->ne[0] == S_v; - // - scalar gate (not KDA), not the rows-indexed state read (src[6]), all f32 + // KDA's per-channel gate takes the chunked path too: on the 44-layer GLM-5.3 proxy its KL against the recurrent + // kernel (mean 0.0065 at 2k and at 16k context, no growth along it) is the size of what a -ub 512 -> 256 change + // alone costs there (0.0058), on a proxy where a reduction order alone costs 0.0026 (-sm tensor with an f32 wire + // against -sm layer). GGML_CUDA_KDA_CHUNKED_LEGACY=1 keeps it on the recurrent kernel at every size. + static const bool kda_chunked = !ggml_env_switch("GGML_CUDA_KDA_CHUNKED_LEGACY"); + + // - the scalar gate or KDA's per channel, not the rows-indexed state read (src[6]), all f32 // - 128-wide heads, q and k with one head count that divides the v-head count, no broadcast over sequences // - 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 && dst->src[6] == nullptr + return (!kda || kda_chunked) && 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 @@ -499,8 +557,11 @@ static void ggml_cuda_op_gated_delta_net_impl( args.v = v_d; args.g = g_d; args.beta = b_d; + args.kda = kda; + args.raw = raw; args.raw_dt_bias = rb_d; args.raw_a = ra_d; + args.raw_lb = rlb; args.state_in = s_d; args.state_out = chunk_state; args.out = dst_d; @@ -520,12 +581,17 @@ static void ggml_cuda_op_gated_delta_net_impl( if (n_tail > 0) { const int64_t t0 = n_chunked; -#define GDN_TAIL_LAUNCH(RAW_) \ - launch_gated_delta_net(q_d + t0 * sq2, k_d + t0 * sq2, \ - v_d + t0 * sv2, g_d + t0 * sb2, b_d + t0 * sb2, rb_d, ra_d, rlb, chunk_state, dst_d + t0 * S_v * H, state_d, \ + const int64_t g0 = t0 * sb2 * (kda ? S_v : 1); // KDA's g at the beta strides times S_v +#define GDN_TAIL_LAUNCH(KDA_, RAW_) \ + launch_gated_delta_net(q_d + t0 * sq2, k_d + t0 * sq2, \ + v_d + t0 * sv2, g_d + g0, b_d + t0 * sb2, rb_d, ra_d, rlb, chunk_state, dst_d + t0 * S_v * H, state_d, \ S_v, H, n_tail, n_seqs, sq1, sq2, sq3, sv1, sv2, sv3, \ sb1, sb2, sb3, neqk1, rq3, scale, state_slot_stride, K, S_v * H * n_tokens, nullptr, stream) - if (raw) { GDN_TAIL_LAUNCH(true); } else { GDN_TAIL_LAUNCH(false); } + if (kda) { + if (raw) { GDN_TAIL_LAUNCH(true, true); } else { GDN_TAIL_LAUNCH(true, false); } + } else { + if (raw) { GDN_TAIL_LAUNCH(false, true); } else { GDN_TAIL_LAUNCH(false, false); } + } #undef GDN_TAIL_LAUNCH } return; @@ -546,6 +612,24 @@ static void ggml_cuda_op_gated_delta_net_impl( g_precomputed = true; } + // KDA from 32 tokens a sequence: the per-channel decay once per token, head and channel (gdn_kda_precompute_decay) + // rather than in every column's warp. GGML_CUDA_KDA_DECAY_PRECOMPUTE_LEGACY=1 leaves it to the kernel. + static const bool kda_decay_precompute = !ggml_env_switch("GGML_CUDA_KDA_DECAY_PRECOMPUTE_LEGACY"); + if (kda && kda_decay_precompute && n_tokens >= 32) { + GGML_ASSERT(src_g->ne[1] == H); // the decay kernel's head of element i is i / S_v % H + const int64_t n_g = ggml_nelements(src_g); + g_exp_alloc.alloc(n_g); + const int block = 256; + const int grid = std::min((n_g + block - 1)/block, 4096); + if (raw) { + gdn_kda_precompute_decay<<>>(g_d, ra_d, rlb, g_exp_alloc.ptr, n_g, S_v, H); + } else { + gdn_kda_precompute_decay<<>>(g_d, nullptr, 0.0f, g_exp_alloc.ptr, n_g, S_v, H); + } + g_d = g_exp_alloc.ptr; + g_precomputed = true; + } + #define GDN_LAUNCH(KDA_, KEEP_, RAW_, PRE_, T_) \ launch_gated_delta_net(q_d, k_d, v_d, g_d, b_d, rb_d, ra_d, rlb, s_in, dst_d, state_d, \ S_v, H, n_tokens, n_seqs, sq1, sq2, sq3, sv1, sv2, sv3, \ @@ -555,7 +639,19 @@ static void ggml_cuda_op_gated_delta_net_impl( if (keep_rs) { GDN_LAUNCH(KDA_, true, RAW_, PRE_, T_); } else { GDN_LAUNCH(KDA_, false, RAW_, PRE_, T_); } // an f16 or q8_0 cache, fused for the scalar gate and KDA's (ggml_cuda_try_gdn_cache_fusion) - if (kda && raw) { + if (kda && raw && g_precomputed) { + switch (state_type) { + case GGML_TYPE_Q8_0: GDN_LAUNCH_KEEP(true, true, true, GGML_TYPE_Q8_0); break; + case GGML_TYPE_F16: GDN_LAUNCH_KEEP(true, true, true, GGML_TYPE_F16); break; + default: GDN_LAUNCH_KEEP(true, true, true, GGML_TYPE_F32); break; + } + } else if (kda && g_precomputed) { + switch (state_type) { + case GGML_TYPE_Q8_0: GDN_LAUNCH_KEEP(true, false, true, GGML_TYPE_Q8_0); break; + case GGML_TYPE_F16: GDN_LAUNCH_KEEP(true, false, true, GGML_TYPE_F16); break; + default: GDN_LAUNCH_KEEP(true, false, true, GGML_TYPE_F32); break; + } + } else if (kda && raw) { switch (state_type) { case GGML_TYPE_Q8_0: GDN_LAUNCH_KEEP(true, true, false, GGML_TYPE_Q8_0); break; case GGML_TYPE_F16: GDN_LAUNCH_KEEP(true, true, false, GGML_TYPE_F16); break; diff --git a/ggml/src/ggml-cuda/ggml-cuda.cu b/ggml/src/ggml-cuda/ggml-cuda.cu index 165e1826815f..c51ff8839e25 100644 --- a/ggml/src/ggml-cuda/ggml-cuda.cu +++ b/ggml/src/ggml-cuda/ggml-cuda.cu @@ -135,6 +135,35 @@ void ggml_cuda_set_device(int device) { CUDA_CHECK(cudaSetDevice(physical_device)); } +// GGML_CUDA_TIME_LAUNCH=1 reports cudaGraphLaunch's own host duration, because a profiler cannot answer the question it +// raises. nsys inflates this call about 39x: a standalone 400-node graph that its own process times at 1.3 us with +// clock_gettime reads 50.6 us by that same clock once nsys is attached, and nsys reports 51.1 us. So every host-side +// 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; + }(); + return enabled; +} + +// a running mean a device, reported every 64 launches so the print is never itself the measurement +static void ggml_cuda_time_launch_record(int device, double us) { + static std::mutex mtx; + static double sum[GGML_CUDA_MAX_DEVICES] = {}; + static int cnt[GGML_CUDA_MAX_DEVICES] = {}; + std::lock_guard lock(mtx); + sum[device] += us; + cnt[device] += 1; + if (cnt[device] % 64 == 0) { + // straight to stderr, not GGML_LOG_*: llama-bench installs a log callback that drops INFO and WARN, which cost + // two builds to discover + fprintf(stderr, "cudaGraphLaunch host cost: device %d, %d launches, mean %.1f us\n", + device, cnt[device], sum[device] / cnt[device]); + } +} + int ggml_cuda_get_device() { int id; CUDA_CHECK(cudaGetDevice(&id)); @@ -758,6 +787,18 @@ ggml_backend_cuda_context::~ggml_backend_cuda_context() { ggml_cuda_set_device(device); pq2_tile_counters.release(); } + if (l2_issue_stop.ptr != nullptr) { + ggml_cuda_set_device(device); + l2_issue_stop.release(); + } + if (ssm_conv_ab_slots.ptr != nullptr) { + ggml_cuda_set_device(device); + ssm_conv_ab_slots.release(); + } + if (hc_front_tickets.ptr != nullptr) { + ggml_cuda_set_device(device); + hc_front_tickets.release(); + } for (int i = 0; i < GGML_CUDA_MAX_DEVICES; ++i) { for (int j = 0; j < GGML_CUDA_MAX_STREAMS; ++j) { if (streams[i][j] != nullptr) { @@ -1709,7 +1750,8 @@ static void ggml_cuda_mul_mat_cublas_impl(ggml_backend_cuda_context & ctx, const } } -static void ggml_cuda_mul_mat_cublas(ggml_backend_cuda_context & ctx, const ggml_tensor * src0, const ggml_tensor * src1, ggml_tensor * dst) { +// The type cuBLAS runs a MUL_MAT in, the type its operands are cast to first +static ggml_type ggml_cuda_mul_mat_cublas_compute_type(const ggml_backend_cuda_context & ctx, const ggml_tensor * src0, const ggml_tensor * dst) { ggml_type compute_type = src0->type; if (ggml_is_quantized(compute_type)) { compute_type = fast_fp16_hardware_available(ggml_cuda_info().devices[ctx.device].cc) ? GGML_TYPE_F16 : GGML_TYPE_F32; @@ -1736,8 +1778,11 @@ static void ggml_cuda_mul_mat_cublas(ggml_backend_cuda_context & ctx, const ggml GGML_LOG_WARN("%s: unknown value for GGML_CUDA_CUBLAS_COMPUTE_TYPE: %s", __func__, env_cpp.c_str()); } } + return compute_type; +} - switch (compute_type) { +static void ggml_cuda_mul_mat_cublas(ggml_backend_cuda_context & ctx, const ggml_tensor * src0, const ggml_tensor * src1, ggml_tensor * dst) { + switch (ggml_cuda_mul_mat_cublas_compute_type(ctx, src0, dst)) { case GGML_TYPE_F32: ggml_cuda_mul_mat_cublas_impl(ctx, src0, src1, dst); break; @@ -2035,10 +2080,22 @@ static double ggml_cuda_pq2_prefetch_us() { return us; } -// For each node the evaluation may dispatch as a PQ2_0 tensor-core launch, the heads of the launch after it; after the -// graph's last launch, its first (the next evaluation of a decode graph starts there). A launch is the members that read -// one src1 (a group, or a gated pair whose GLU reads both); its heads are its first member's weights and, for a gated -// pair, the other's, which it streams beside them. Empty when nothing is prefetched. +static bool ggml_cuda_is_view_or_noop(const ggml_tensor * t); + +// GGML_CUDA_PQ2_PREFETCH_BETWEEN_LEGACY=1: a launch prefetches the next launch's heads alone, not the weights of the +// kernels before it +static bool ggml_cuda_pq2_prefetch_between() { + static const bool legacy = ggml_env_switch("GGML_CUDA_PQ2_PREFETCH_BETWEEN_LEGACY"); + return !legacy; +} + +// For each node the evaluation may dispatch as a PQ2_0 tensor-core launch, what it prefetches (ggml_cuda_pq2_prefetch): +// the weights the nodes after it and before the next launch read, in node order, up to the budget (a node's sources in a +// weights buffer; not a GET_ROWS' table, whose ids pick its rows), then the heads of the next launch; after the graph's +// last launch, the nodes to the graph's end and the first launch's heads (the next evaluation of a decode graph starts +// there). A launch is the members that read one src1 (a group, or a gated pair whose GLU reads both); its heads are its +// first member's weights and, for a gated pair, the other's, which it streams beside them, up to the budget together. +// Empty when nothing is prefetched. static std::vector ggml_cuda_pq2_prefetch_plan(ggml_backend_cuda_context & ctx, const ggml_cgraph * cgraph) { const ggml_cuda_device_info::cuda_device_info & dev = ggml_cuda_info().devices[ctx.device]; @@ -2068,9 +2125,43 @@ static std::vector ggml_cuda_pq2_prefetch_plan(ggml_back if (launches.empty()) { return {}; } - std::vector heads(launches.size()); + const bool between = ggml_cuda_pq2_prefetch_between(); + std::vector plan(cgraph->n_nodes); for (size_t k = 0; k < launches.size(); ++k) { - const launch & l = launches[k]; + ggml_cuda_pq2_prefetch p; + if (between) { + const int to = k + 1 < launches.size() ? launches[k + 1].at[0] : cgraph->n_nodes; + int64_t left = budget; + for (int i = launches[k].at[launches[k].n - 1] + 1; i < to && left > 0; ++i) { + const ggml_tensor * node = cgraph->nodes[i]; + if (ggml_cuda_is_view_or_noop(node) || (node->flags & GGML_TENSOR_FLAG_COMPUTE) == 0) { + continue; + } + // skipped, as no kernel reads them whole: a GET_ROWS table and a MUL_MAT_ID expert stack (rows by their + // ids), and a Hadamard-hint MUL_MAT's rotation table (the rotation is computed, the table never read) + const auto unread = [&](int j) { + return j == 0 && (node->op == GGML_OP_GET_ROWS || node->op == GGML_OP_MUL_MAT_ID || + (node->op == GGML_OP_MUL_MAT && ggml_get_op_params_i32(node, 1) == GGML_HINT_SRC0_IS_HADAMARD)); + }; + for (int j = 0; j < GGML_MAX_SRC && left > 0; ++j) { + const ggml_tensor * src = node->src[j]; + if (src == nullptr || src->buffer == nullptr || unread(j) || + ggml_backend_buffer_get_usage(src->buffer) != GGML_BACKEND_BUFFER_USAGE_WEIGHTS) { + continue; + } + bool seen = false; + for (int r = 0; r < p.n; ++r) { + seen |= (const char *) src->data >= (const char *) p.ptr[r] - 15 && + (const char *) src->data < (const char *) p.ptr[r] + p.bytes[r]; + } + const int64_t nbytes = std::min((int64_t) ggml_nbytes(src), left); + if (!seen && p.add(src->data, nbytes)) { + left -= nbytes; + } + } + } + } + const launch & l = launches[(k + 1) % launches.size()]; const ggml_tensor * first = l.members[0]; const ggml_tensor * gate = nullptr; if (l.n == 2) { @@ -2086,12 +2177,12 @@ static std::vector ggml_cuda_pq2_prefetch_plan(ggml_back if (gate != nullptr) { bytes = std::min(bytes, (int64_t) ggml_nbytes(gate)); } - heads[k] = { first->src[0]->data, gate != nullptr ? gate->data : nullptr, bytes & ~(int64_t) 15 }; - } - std::vector plan(cgraph->n_nodes); - for (size_t k = 0; k < launches.size(); ++k) { + p.add(first->src[0]->data, bytes); + if (gate != nullptr) { + p.add(gate->data, bytes); + } for (int m = 0; m < launches[k].n; ++m) { - plan[launches[k].at[m]] = heads[(k + 1) % launches.size()]; + plan[launches[k].at[m]] = p; } } return plan; @@ -2453,7 +2544,9 @@ static ggml_cuda_l2_issue_plan ggml_cuda_l2_issue_plan_of(ggml_backend_cuda_cont // the chain it runs beside: the nodes up to the next launch that keeps DRAM busy (the next heavy one, routed // experts, or a node moving busy_kb), each node_us of latency, and the all-reduce if it crosses one; requested // past its end the issuer takes DRAM from that launch (the 44-layer proxy on an RTX 5070 Ti, L2/2 an issue: the - // routed experts 5,970 against 5,280 us a token, qkv 4,506 against 4,252) + // routed experts 5,970 against 5,280 us a token, qkv 4,506 against 4,252). A node is not a launch where the + // evaluation fuses (a hyper-connection front's, a top-k's), so the chain may be shorter: a routed-expert ring + // stops the issue as it starts (ggml_cuda_l2_issue_fork) bool busy = false; double chain_us = 0.0; if (!is.crosses) { @@ -2496,8 +2589,13 @@ static ggml_cuda_l2_issue_plan ggml_cuda_l2_issue_plan_of(ggml_backend_cuda_cont return plan; } -// Starts the issuer for r on its own stream, after all that the evaluation's stream has been given so far +// Starts the issuer for r on its own stream, after all that the evaluation's stream has been given so far. It stops at the +// first launch that bumps the context's stop word (ggml_cuda_l2_issue_stop: a routed-expert ring, whose reads would share +// DRAM with it: GLM-5.3's proxy at 3 tokens on an RTX 5070 Ti, the issue after the attention output ran ~25 us into the +// gate/up, 125.0 us against 108.5 with no issuer, where the shared expert it brought into L2 gained back 9.2). +// GGML_CUDA_L2_ISSUE_STOP_LEGACY=1: it runs its length. static void ggml_cuda_l2_issue_fork(ggml_backend_cuda_context & ctx, const ggml_cuda_l2_ranges & r, const double rate_gbs) { + static const bool stop_legacy = ggml_env_switch("GGML_CUDA_L2_ISSUE_STOP_LEGACY"); if (ctx.l2_issue_fork == nullptr) { CUDA_CHECK(cudaEventCreateWithFlags(&ctx.l2_issue_fork, cudaEventDisableTiming)); CUDA_CHECK(cudaEventCreateWithFlags(&ctx.l2_issue_join, cudaEventDisableTiming)); @@ -2505,7 +2603,8 @@ static void ggml_cuda_l2_issue_fork(ggml_backend_cuda_context & ctx, const ggml_ cudaStream_t side = ctx.stream(ctx.device, GGML_CUDA_L2_ISSUE_STREAM); CUDA_CHECK(cudaEventRecord(ctx.l2_issue_fork, ctx.stream())); CUDA_CHECK(cudaStreamWaitEvent(side, ctx.l2_issue_fork, 0)); - ggml_cuda_l2_issue(r, rate_gbs, ggml_cuda_info().devices[ctx.device].nsm, side); + ggml_cuda_l2_issue(r, rate_gbs, ggml_cuda_info().devices[ctx.device].nsm, stop_legacy ? nullptr : ctx.l2_issue_stop.ptr, + side); ctx.l2_issue_open = true; } @@ -2599,6 +2698,24 @@ static bool ggml_cuda_should_use_mmvf_untiled(const ggml_tensor * src0, const gg && ggml_cuda_mmvf_supports(src0, src1); } +// On NVIDIA from Ampere on, ggml_cuda_should_use_mmvf leaves f16 and bf16 weights past one column (f32 past 3) to +// mul_mat_f, which launches a block for each MMF_ROWS_PER_BLOCK rows of each dst channel and reads a row's whole K in +// it: a wide projection to few rows runs on a few SMs. glm5next's ssm_f_a and ssm_g_a (4096 -> 128 bf16) took 8.6 us +// each on 4 blocks at an MTP verify's 3 columns and ssm_beta (4096 -> 64) 6.4 us on 2, RTX 5070 Ti, where at one column +// mul_mat_vec_f, a block a row, ran f_a and g_a as one launch in 5.6 us. mul_mat_f's time there follows K, the vector +// kernel's rows x cols: over f16/bf16 weights of 64-2048 rows, K 1024-8192 and 2-8 columns (test-backend-ops perf, RTX +// 5070 Ti) the vector kernel took 0.21-0.68 of mul_mat_f's time at rows x cols <= 1024, 0.76-1.13 at 2048 and 1.1-4.4 +// past it. So there, at 2 to MMVF_MAX_BATCH_SIZE columns and rows x cols <= 1024, where mul_mat_f would launch fewer +// blocks than the device has SMs, mul_mat_vec_f takes the product (and a pair of them on one src1 runs as one launch, +// ggml_cuda_mul_mat_vec_f_pair). GGML_CUDA_MMVF_UNDERFILLED_LEGACY=1: mul_mat_f. +static bool ggml_cuda_should_use_mmvf_underfilled(const ggml_tensor * src0, const ggml_tensor * src1, + const ggml_tensor * dst, const int cc, const int nsm) { + static const bool legacy = ggml_env_switch("GGML_CUDA_MMVF_UNDERFILLED_LEGACY"); + return !legacy && ampere_mma_available(cc) && src1->ne[1] >= 2 && src1->ne[1] <= MMVF_MAX_BATCH_SIZE + && src0->ne[1]*src1->ne[1] <= 1024 && ggml_cuda_mmvf_supports(src0, src1) + && src0->ne[1]/MMF_ROWS_PER_BLOCK * dst->ne[2]*dst->ne[3] < nsm; +} + static void ggml_cuda_mul_mat(ggml_backend_cuda_context & ctx, const ggml_tensor * src0, const ggml_tensor * src1, ggml_tensor * dst) { GGML_TENSOR_BINARY_OP_LOCALS @@ -2641,6 +2758,10 @@ static void ggml_cuda_mul_mat(ggml_backend_cuda_context & ctx, const ggml_tensor return; } if (ggml_cuda_should_use_mmf(src0, src1, cc, warp_size, ne11, /*mul_mat_id =*/ false)) { + if (ggml_cuda_should_use_mmvf_underfilled(src0, src1, dst, cc, ggml_cuda_info().devices[ctx.device].nsm)) { + ggml_cuda_mul_mat_vec_f(ctx, src0, src1, nullptr, dst); + return; + } ggml_cuda_mul_mat_f(ctx, src0, src1, nullptr, dst); return; } @@ -2667,18 +2788,43 @@ static void ggml_cuda_mul_mat(ggml_backend_cuda_context & ctx, const ggml_tensor // Whether ggml_cuda_mul_mat runs the MUL_MAT dst on ggml_cuda_mul_mat_vec_f with its own operands, as its choices above // go at up to MMVF_MAX_BATCH_SIZE columns (where the transposed-vector path cannot apply). -static bool ggml_cuda_mul_mat_runs_mmvf(const ggml_tensor * dst, const int cc, const int warp_size) { +static bool ggml_cuda_mul_mat_runs_mmvf(const ggml_tensor * dst, const int cc, const int warp_size, const int nsm) { const ggml_tensor * src0 = dst->src[0]; const ggml_tensor * src1 = dst->src[1]; const int64_t ne11 = src1->ne[1]; const bool bad_padding_clear = ggml_backend_buffer_get_usage(src0->buffer) == GGML_BACKEND_BUFFER_USAGE_COMPUTE && ggml_nbytes(src0) != ggml_backend_buffer_get_alloc_size(src0->buffer, src0) && src0->view_src; - return ggml_get_op_params_i32(dst, 1) != GGML_HINT_SRC0_IS_HADAMARD && !bad_padding_clear - && src1->type == GGML_TYPE_F32 && dst->type == GGML_TYPE_F32 && ne11 <= MMVF_MAX_BATCH_SIZE - && (ggml_cuda_should_use_mmvf(src0, src1, cc, ne11) - || (!ggml_cuda_should_use_mmf(src0, src1, cc, warp_size, ne11, /*mul_mat_id =*/ false) - && ggml_cuda_should_use_mmvf_untiled(src0, src1, ne11))); + if (ggml_get_op_params_i32(dst, 1) == GGML_HINT_SRC0_IS_HADAMARD || bad_padding_clear + || src1->type != GGML_TYPE_F32 || dst->type != GGML_TYPE_F32 || ne11 > MMVF_MAX_BATCH_SIZE) { + return false; + } + if (ggml_cuda_should_use_mmvf(src0, src1, cc, ne11)) { + return true; + } + return ggml_cuda_should_use_mmf(src0, src1, cc, warp_size, ne11, /*mul_mat_id =*/ false) + ? ggml_cuda_should_use_mmvf_underfilled(src0, src1, dst, cc, nsm) + : ggml_cuda_should_use_mmvf_untiled(src0, src1, ne11); +} + +// Whether ggml_cuda_mul_mat runs the MUL_MAT dst on cuBLAS in BF16, casting its F32 src1 to BF16 first: a BF16 src0 +// that none of its kernels above takes (a BF16 src0 is never quantized, nor the F32 of the transposed vector). +static bool ggml_cuda_mul_mat_runs_cublas_bf16(const ggml_backend_cuda_context & ctx, const ggml_tensor * dst) { + const ggml_tensor * src0 = dst->src[0]; + const ggml_tensor * src1 = dst->src[1]; + const int64_t ne11 = src1->ne[1]; + const int cc = ggml_cuda_info().devices[ctx.device].cc; + + if (src0->type != GGML_TYPE_BF16 || src1->type != GGML_TYPE_F32 || dst->type != GGML_TYPE_F32 || + ggml_get_op_params_i32(dst, 1) == GGML_HINT_SRC0_IS_HADAMARD) { + return false; + } + if (ggml_cuda_should_use_mmvf(src0, src1, cc, ne11) || + ggml_cuda_should_use_mmf(src0, src1, cc, ggml_cuda_info().devices[ctx.device].warp_size, ne11, /*mul_mat_id =*/ false) || + ggml_cuda_should_use_mmvf_untiled(src0, src1, ne11)) { + return false; + } + return ggml_cuda_mul_mat_cublas_compute_type(ctx, src0, dst) == GGML_TYPE_BF16; } // AMD runs a float MUL_MAT_ID of up to MMVF_MAX_BATCH_SIZE tokens on the vector kernel, which can read only operands @@ -3713,6 +3859,145 @@ static bool ggml_cuda_try_gdn_gather_skip(ggml_backend_cuda_context & ctx, const return false; } +// Qwen3.5's alpha/beta pair, folded into the conv-state update matched for the SSM_CONV at conv_idx (u, which this +// fills in): the two MUL_MATs of bf16 weights on the layer's normed input that come first after the SSM_CONV (past its +// SILU, the L2_NORM it folds and views) and that ggml_cuda_mul_mat would run on mul_mat_vec_f at 256 threads, one +// launch where there were two, and one boundary fewer on the chain from the qkv group to the recurrence. The conv's +// blocks compute the rows before their dependency wait and write them after it (ssm_conv_ab_rows), so: +// - the normed input reaches them by a release/acquire handoff on the plan's slot (ggml_cuda_ssm_conv_ab_slots): the +// kernel that writes it releases to the slot, and the conv's launch folds the pair in only if it did this evaluation +// (rms_norm_fwht_cuda, the only writer that releases; any other leaves the pair its own launch); +// - the PQ2_0 group launch writing the conv's new inputs must come after that writer (checked when the group is +// dispatched: ggml_cuda_ssm_conv_ab_plan); it then lets the kernels after it launch only after its own wait +// (ggml_cuda_ssm_conv_ab_enabled), so the writer has ended when the conv starts and the conv's blocks never wait on +// the slot beside writer blocks still to run. Nodes between the group and the conv may launch kernels (the SSM +// state's GET_ROWS, unless the recurrence gathers it); +// - the weights sit in a weights buffer, which no kernel writes; +// - the outputs, written when the conv runs instead of at their place, overlap nothing the conv reads or writes, +// nothing a node from the SSM_CONV up to the pair reads or writes, and neither the input nor each other; and none +// of those nodes writes over the input. +// Registers the plan under the SSM_CONV. CUDA only (the fold copies mul_mat_vec_f's CUDA arithmetic). +static void ggml_cuda_try_ssm_conv_ab(ggml_backend_cuda_context & ctx, const ggml_cgraph * cgraph, const int conv_idx, + const ggml_tensor * concat, const ggml_tensor * x, + const ggml_tensor * ids, const ggml_tensor * l2, ggml_cuda_ssm_conv_state_update & u) { +#if defined(GGML_USE_HIP) || defined(GGML_USE_MUSA) + GGML_UNUSED_VARS(ctx, cgraph, conv_idx, concat, x, ids, l2, u); +#else + if (!ggml_cuda_ssm_conv_ab_enabled()) { + return; + } + const int cc = ggml_cuda_info().devices[ctx.device].cc; + const int warp_size = ggml_cuda_info().devices[ctx.device].warp_size; + const int nsm = ggml_cuda_info().devices[ctx.device].nsm; + if (warp_size != WARP_SIZE) { + return; + } + + // the pair: the first nodes after the SSM_CONV but its SILU, the folded L2_NORM and views, up to 4 views apart + const ggml_tensor * silu = cgraph->nodes[conv_idx + 1]; + int ia = conv_idx + 2; + while (ia < cgraph->n_nodes && (ggml_cuda_is_view_or_noop(cgraph->nodes[ia]) || cgraph->nodes[ia] == l2)) { + ++ia; + } + int ib = ia + 1; + while (ib < cgraph->n_nodes && ib - ia <= 4 && ggml_cuda_is_view_or_noop(cgraph->nodes[ib])) { + ++ib; + } + if (ib >= cgraph->n_nodes) { + return; + } + const ggml_tensor * mm[2] = { cgraph->nodes[ia], cgraph->nodes[ib] }; + const ggml_tensor * y = mm[0]->src[1]; + for (const ggml_tensor * m : mm) { + const ggml_tensor * w = m->src[0]; + if (m->op != GGML_OP_MUL_MAT || !(m->flags & GGML_TENSOR_FLAG_COMPUTE) || m->src[1] != y || + w->type != GGML_TYPE_BF16 || w->buffer == nullptr || + ggml_backend_buffer_get_usage(w->buffer) != GGML_BACKEND_BUFFER_USAGE_WEIGHTS || + !ggml_cuda_mul_mat_runs_mmvf(m, cc, warp_size, nsm)) { + return; + } + } + if (!ggml_cuda_mmvf_pair_supports(mm[0]->src[0], mm[1]->src[0], y, mm[0], mm[1])) { + return; + } + const ggml_tensor * w0 = mm[0]->src[0]; + const int64_t n_t = concat->ne[0] - (GGML_CUDA_SSM_CONV_UPDATE_D_CONV - 1); + const int64_t n_rows = w0->ne[1]; + const int64_t n_blk = concat->ne[1] / GGML_CUDA_SSM_CONV_UPDATE_THREADS; // the conv's blocks + if (y->ne[1] != n_t || y->nb[0] != sizeof(float) || y->nb[1] % (2*sizeof(float)) != 0 || + (uintptr_t) y->data % (2*sizeof(float)) != 0 || mm[0]->ne[0] != n_rows || mm[0]->ne[1] != n_t || + mm[0]->nb[1] % sizeof(float) != 0 || 2*n_rows > GGML_CUDA_SSM_CONV_AB_MAX_ROWS*n_blk) { + return; + } + // mul_mat_vec_f's block for these columns (launch_mul_mat_vec_f_cuda) must be the 256 threads the fold copies + int64_t block = warp_size; + int64_t niter = (w0->ne[0] + 2*warp_size - 1) / (2*warp_size); + for (int64_t b = 2*warp_size; b <= 256; b += warp_size) { + const int64_t it = (w0->ne[0] + 2*b - 1) / (2*b); + if (it < niter) { + niter = it; + block = b; + } + } + if (block != 2*GGML_CUDA_SSM_CONV_UPDATE_THREADS) { + return; + } + + // memory: the outputs are written when the conv runs; every kernel before it has completed by then + const ggml_tensor * conv = cgraph->nodes[conv_idx]; + const ggml_tensor * no_touch[] = { y, x, ids, conv, silu, l2, concat }; + for (const ggml_tensor * m : mm) { + for (const ggml_tensor * t : no_touch) { + if (t != nullptr && ggml_cuda_ranges_overlap(m, t)) { + return; + } + } + for (int j = conv_idx; j < ia; ++j) { + const ggml_tensor * n = cgraph->nodes[j]; + if (ggml_cuda_is_view_or_noop(n)) { + continue; + } + if (ggml_cuda_ranges_overlap(m, n)) { + return; + } + for (int s = 0; s < GGML_MAX_SRC; ++s) { + if (n->src[s] != nullptr && ggml_cuda_ranges_overlap(m, n->src[s])) { + return; + } + } + } + } + if (ggml_cuda_ranges_overlap(mm[0], mm[1])) { + return; + } + for (int j = conv_idx; j < ia; ++j) { + const ggml_tensor * n = cgraph->nodes[j]; + if (!ggml_cuda_is_view_or_noop(n) && ggml_cuda_ranges_overlap(n, y)) { + return; + } + } + + // the handoff slot, and the ticket's 16 bits for the conv's blocks (ggml_cuda_ssm_conv_ab_slots) + ggml_cuda_ssm_conv_update_context & cu = ctx.ssm_conv_updates(); + if (cu.ab_slots >= ggml_cuda_ssm_conv_ab_slots::n || n_blk >= (1 << 16) || cu.ab_by_y.count(y) != 0) { + return; + } + + u.ab_w = (const nv_bfloat16 *) w0->data; + u.ab_y = (const float *) y->data; + u.ab_dst = (float *) mm[0]->data; + u.ab_s02 = ((const char *) mm[1]->src[0]->data - (const char *) w0->data) / (int64_t) sizeof(nv_bfloat16); + u.ab_s2 = ((const char *) mm[1]->data - (const char *) mm[0]->data) / (int64_t) sizeof(float); + u.ab_stride_row = w0->nb[1] / sizeof(nv_bfloat16); + u.ab_stride_y = y->nb[1] / sizeof(float); + u.ab_stride_dst = mm[0]->nb[1] / sizeof(float); + u.ab_rows = (int) n_rows; + u.ab_ncols = (int) w0->ne[0]; + cu.ab_plans[conv] = { { mm[0], mm[1] }, y, x, cu.ab_slots++, 0, false }; + cu.ab_by_y[y] = conv; +#endif // defined(GGML_USE_HIP) || defined(GGML_USE_MUSA) +} + // KDA builds the SSM_CONV's weights as CONCAT(CONCAT(q, k), v) along rows of three RESHAPEd model weights, every token. // When the SSM_CONV at conv_idx reads such a pair, let the fused conv kernel read the three weights itself (w_seg) and // return the two CONCATs, for the caller to skip. Needs equal weight rows, nothing else reading the CONCATs. @@ -3925,6 +4210,8 @@ static bool ggml_cuda_try_ssm_conv_state_update(ggml_backend_cuda_context & ctx, } } + ggml_cuda_try_ssm_conv_ab(ctx, cgraph, j, concat, x, ids, ctx.ssm_conv_updates().l2_norm_of(n), u); + auto & reg = ctx.ssm_conv_updates(); reg.updates[n] = u; reg.skipped.insert(concat); @@ -4253,8 +4540,11 @@ static bool ggml_cuda_check_fusion_memory_ranges(const ggml_cgraph * cgraph, }; bool is_ok = true; - // exception for topk-moe, as each row is read entirely before writing - if (ggml_nrows(cgraph->nodes[node_idx]) == 1 && is_topk_moe) { + // 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))) { return true; } @@ -4943,6 +5233,47 @@ static int ggml_cuda_try_fuse(ggml_backend_cuda_context * cuda_ctx, ggml_cgraph } } + // routed experts' weighted sum as build_moe_ffn writes it: MUL(experts, weights), each slot's view of it, and the + // views added in slot order, in one launch (ggml_cuda_op_moe_weighted_sum) where the MUL and the fused ADDs took two. + // GGML_CUDA_MOE_WSUM_LEGACY=1: the two + static const bool moe_wsum_legacy = ggml_env_switch("GGML_CUDA_MOE_WSUM_LEGACY"); + if (!moe_wsum_legacy && node->op == GGML_OP_MUL && node->ne[1] >= 2 && 2*node->ne[1] < 32 && + i + 2*node->ne[1] <= cgraph->n_nodes) { + const int n_used = (int) node->ne[1]; + const int n_ops = 2*n_used; // the MUL, n_used views, n_used - 1 ADDs + const ggml_tensor * experts = node->src[0]; + const ggml_tensor * weights = node->src[1]; + bool ok = node->type == GGML_TYPE_F32 && experts->type == GGML_TYPE_F32 && weights->type == GGML_TYPE_F32 && + ggml_are_same_shape(node, experts) && node->ne[3] == 1 && weights->ne[0] == 1 && + weights->ne[1] == n_used && weights->ne[2] == node->ne[2] && weights->ne[3] == 1 && + experts->nb[0] == sizeof(float) && ggml_is_contiguous(node); + // slot s's view: [n_embd, n_tokens] at s*nb[1], its rows nb[2] apart + for (int s = 0; ok && s < n_used; ++s) { + const ggml_tensor * v = cgraph->nodes[i + 1 + s]; + ok = v->op == GGML_OP_VIEW && v->src[0] == node && v->view_offs == s*node->nb[1] && + v->ne[0] == node->ne[0] && v->ne[1] == node->ne[2] && v->ne[2] == 1 && v->nb[1] == node->nb[2]; + } + // ADD(v0, v1), then ADD(the sum so far, v_s) + for (int s = 1; ok && s < n_used; ++s) { + const ggml_tensor * a = cgraph->nodes[i + n_used + s]; + const ggml_tensor * prev = s == 1 ? cgraph->nodes[i + 1] : cgraph->nodes[i + n_used + s - 1]; + ok = a->op == GGML_OP_ADD && a->type == GGML_TYPE_F32 && a->src[0] == prev && + a->src[1] == cgraph->nodes[i + 1 + s] && a->nb[0] == sizeof(float); + } + if (ok) { + ggml_op ops[32]; + ops[0] = GGML_OP_MUL; + std::fill(ops + 1, ops + 1 + n_used, GGML_OP_VIEW); + std::fill(ops + 1 + n_used, ops + n_ops, GGML_OP_ADD); + const int out = i + n_ops - 1; + if (ggml_can_fuse_subgraph(cgraph, i, n_ops, ops, &out, 1) && + ggml_cuda_check_fusion_memory_ranges(cgraph, i, n_ops, &out, 1)) { + ggml_cuda_op_moe_weighted_sum(*cuda_ctx, experts, weights, cgraph->nodes[out]); + return n_ops - 1; + } + } + } + // multi-(add or mul) if (node->op == GGML_OP_ADD || node->op == GGML_OP_MUL) { int n_fuse = 0; @@ -5580,14 +5911,29 @@ static int ggml_cuda_try_fuse(ggml_backend_cuda_context * cuda_ctx, ggml_cgraph GGML_ASSERT(ggml_cuda_compute_forward(*cuda_ctx, t)); } } - ggml_cuda_op_rms_norm_fwht(*cuda_ctx, node, w, normed, mul_s->src[1], mm); + // normed may be an alpha/beta fold's activation: the kernel then releases it to the fold's slot + ggml_cuda_ssm_conv_ab_plan * ab = normed != nullptr ? cuda_ctx->ssm_conv_updates().ab_plan_of_y(normed) : nullptr; + const int blocks = ggml_cuda_op_rms_norm_fwht(*cuda_ctx, node, w, normed, mul_s->src[1], mm, + ab != nullptr ? cuda_ctx->ssm_conv_ab_slots.ptr + ab->slot : nullptr); + if (ab != nullptr) { + ab->writers = blocks; + } return j + 2 - i; } } } if (const int n = ggml_cuda_pq2_mma_group_size(*cuda_ctx, cgraph, i); n > 1) { - ggml_cuda_mul_mat_vec_q_pq2_group(*cuda_ctx, cgraph->nodes + i, n); + // a group writing the new inputs of a conv whose alpha/beta fold's activation is released feeds that fold + bool feeds_fold = false; + for (auto & [conv, ab] : cuda_ctx->ssm_conv_updates().ab_plans) { + for (int g = 0; g < n && ab.writers > 0 && !ab.fed; ++g) { + if (ggml_cuda_ranges_overlap(cgraph->nodes[i + g], ab.x)) { + ab.fed = feeds_fold = true; + } + } + } + ggml_cuda_mul_mat_vec_q_pq2_group(*cuda_ctx, cgraph->nodes + i, n, feeds_fold); return n - 1; } @@ -5660,9 +6006,11 @@ static int ggml_cuda_try_fuse(ggml_backend_cuda_context * cuda_ctx, ggml_cgraph const int cc = ggml_cuda_info().devices[cuda_ctx->device].cc; const int warp_size = ggml_cuda_info().devices[cuda_ctx->device].warp_size; + const int nsm = ggml_cuda_info().devices[cuda_ctx->device].nsm; const int out_nodes[] = { i, j }; if (mm_b && mm_b->op == GGML_OP_MUL_MAT && (mm_b->flags & GGML_TENSOR_FLAG_COMPUTE) && mm_b->src[1] == node->src[1] && - !reads_first && ggml_cuda_mul_mat_runs_mmvf(node, cc, warp_size) && ggml_cuda_mul_mat_runs_mmvf(mm_b, cc, warp_size) && + !reads_first && ggml_cuda_mul_mat_runs_mmvf(node, cc, warp_size, nsm) && + ggml_cuda_mul_mat_runs_mmvf(mm_b, cc, warp_size, nsm) && ggml_cuda_mmvf_pair_supports(node->src[0], mm_b->src[0], node->src[1], node, mm_b) && ggml_cuda_check_fusion_memory_ranges(cgraph, i, j - i + 1, out_nodes, 2)) { ggml_cuda_mul_mat_vec_f_pair(*cuda_ctx, node->src[0], mm_b->src[0], node->src[1], node, mm_b); @@ -5785,6 +6133,34 @@ static int ggml_cuda_try_fuse(ggml_backend_cuda_context * cuda_ctx, ggml_cgraph } } + // The same pair at a prefill ubatch's columns, where the MUL_MAT runs on cuBLAS in BF16 and casts the norm's F32 + // output to BF16 before its GEMM: the norm writes that BF16 copy itself, the F32 one is neither written nor read, and + // cuBLAS reads the same values it would have cast (GLM-5.3's hc mixes: the norm and the cast were 94 % of norm, cast + // and GEMM, around a 21 us GEMM). GGML_CUDA_RMS_NORM_BF16_LEGACY=1 runs the norm and the cast on their own. + static const bool rms_norm_bf16_legacy = ggml_env_switch("GGML_CUDA_RMS_NORM_BF16_LEGACY"); + if (!rms_norm_bf16_legacy && node->op == GGML_OP_RMS_NORM && i + 1 < cgraph->n_nodes) { + ggml_tensor * mm = cgraph->nodes[i + 1]; + const ggml_tensor * x = node->src[0]; + const ggml_op ops[] = { GGML_OP_RMS_NORM, GGML_OP_MUL_MAT }; + const int out_nodes[] = { i + 1 }; + if (mm->op == GGML_OP_MUL_MAT && mm->src[1] == node && mm->src[0] != node && x->type == GGML_TYPE_F32 && + ggml_is_contiguous(node) && ggml_can_fuse_subgraph(cgraph, i, 2, ops, out_nodes, 1) && + ggml_cuda_mul_mat_runs_cublas_bf16(*cuda_ctx, mm) && + ggml_cuda_check_fusion_memory_ranges(cgraph, i, 2, out_nodes, 1)) { + ggml_cuda_pool_alloc y(cuda_ctx->pool(), ggml_nelements(node)); + ggml_cuda_op_rms_norm_bf16(*cuda_ctx, node, y.get()); + ggml_tensor y_bf16 = *node; + y_bf16.type = GGML_TYPE_BF16; + y_bf16.data = y.get(); + y_bf16.nb[0] = ggml_type_size(GGML_TYPE_BF16); + for (int d = 1; d < GGML_MAX_DIMS; ++d) { + y_bf16.nb[d] = y_bf16.nb[d - 1]*y_bf16.ne[d - 1]; + } + ggml_cuda_mul_mat_cublas(*cuda_ctx, mm->src[0], &y_bf16, mm); + return 1; + } + } + if (ggml_cuda_can_fuse(cgraph, i, { GGML_OP_RMS_NORM, GGML_OP_MUL, GGML_OP_ROPE, GGML_OP_VIEW, GGML_OP_SET_ROWS }, {})) { ggml_cuda_op_rms_norm_mul_rope_fused(*cuda_ctx, node, cgraph->nodes[i + 1], cgraph->nodes[i + 2], cgraph->nodes[i + 4]); return 4; @@ -5932,7 +6308,9 @@ static void ggml_cuda_graph_evaluate_and_capture(ggml_backend_cuda_context * cud // (ggml_cuda_mmvq_shared_q8_1): its copy is made before its first reader and held to the end of the evaluation, and // the readers are counted by the input's first byte at the first candidate. Off while the graph has concurrent streams. // GGML_CUDA_MMVQ_SHARED_Q8_1_LEGACY=1 quantizes it at each reader. - static const bool mmvq_shared_q8_1_legacy = ggml_env_switch("GGML_CUDA_MMVQ_SHARED_Q8_1_LEGACY"); + // GGML_CUDA_MMVQ_Q8_1_PRODUCER_LEGACY=1: no kernel writes a copy beside its output, the first reader quantizes it + static const bool mmvq_shared_q8_1_legacy = ggml_env_switch("GGML_CUDA_MMVQ_SHARED_Q8_1_LEGACY"); + static const bool mmvq_q8_1_producer_legacy = ggml_env_switch("GGML_CUDA_MMVQ_Q8_1_PRODUCER_LEGACY"); const auto mmvq_shared_q8_1_candidate = [](const ggml_tensor * node) { if ((node->op != GGML_OP_MUL_MAT && node->op != GGML_OP_MUL_MAT_ID) || node->src[0] == nullptr || node->src[1] == nullptr || !ggml_is_quantized(node->src[0]->type) || node->src[1]->type != GGML_TYPE_F32 || @@ -5958,6 +6336,7 @@ static void ggml_cuda_graph_evaluate_and_capture(ggml_backend_cuda_context * cud return it == mmvq_shared_q8_1_readers.end() ? 0 : it->second; }; cuda_ctx->mmvq_shared_q8_1.reset(); + std::fill(std::begin(cuda_ctx->mmvq_moe_ids), std::end(cuda_ctx->mmvq_moe_ids), nullptr); static const bool virtual_rms_q8_enabled = [] { const char * env = getenv("GGML_CUDA_GB10_VIRTUAL_RMS_Q8"); @@ -6088,6 +6467,27 @@ static void ggml_cuda_graph_evaluate_and_capture(ggml_backend_cuda_context * cud int mmvq_shared_q8_1_seen = 0; // the nodes before it have been checked for writes over a q8_1 copy's input + // the copies whose producers write them (a hyper-connection front's normed mix with a quantized reader), made + // before any node runs: one a key, which each of its producers writes in turn + if (!mmvq_shared_q8_1_legacy && !mmvq_q8_1_producer_legacy && !should_launch_concurrent_events) { + ggml_cuda_mmvq_shared_q8_1 & shared = cuda_ctx->mmvq_shared_q8_1; + for (int j = 0; j < cgraph->n_nodes; ++j) { + const ggml_tensor * node = cgraph->nodes[j]; + ggml_cuda_mmvq_shared_q8_1::key k; + if (!(node->flags & GGML_TENSOR_FLAG_COMPUTE) || !ggml_cuda_dsv4_hc_writes_q8_1(node) || + mmvq_shared_q8_1_reader_count(node) == 0 || !ggml_cuda_mmvq_shared_q8_1::key_of(node, k)) { + continue; + } + if (shared.find_any(node) == nullptr) { + const size_t size = ggml_cuda_mmvq_shared_q8_1::q8_1_size(k); + auto data = std::make_unique>(cuda_ctx->pool(), size); + shared.entries.push_back({ k, data->get(), size, false }); + gb10_pool_allocations.push_back(std::move(data)); + } + shared.producers[node] = j; + } + } + for (int i = 0; i < cgraph->n_nodes; i++) { ggml_tensor * node = cgraph->nodes[i]; if (is_concurrent_event_active) { @@ -6130,7 +6530,7 @@ static void ggml_cuda_graph_evaluate_and_capture(ggml_backend_cuda_context * cud for (int k = mmvq_shared_q8_1_seen; k < i; ++k) { const ggml_tensor * ran = cgraph->nodes[k]; if (!ggml_cuda_is_view_or_noop(ran) && (ran->flags & GGML_TENSOR_FLAG_COMPUTE)) { - cuda_ctx->mmvq_shared_q8_1.written((const char *) ran->data, ggml_nbytes(ran)); + cuda_ctx->mmvq_shared_q8_1.written((const char *) ran->data, ggml_nbytes(ran), k); } } } @@ -6213,9 +6613,11 @@ static void ggml_cuda_graph_evaluate_and_capture(ggml_backend_cuda_context * cud } } - // a MUL_MAT whose input other MUL_MATs read too leads its group: the first to run makes the input's copy + // a MUL_MAT whose input other MUL_MATs read too, or has a copy already, leads its group: the first to run + // makes the input's copy where no producer has if (!mmvq_shared_q8_1_legacy && !should_launch_concurrent_events && mmvq_shared_q8_1_candidate(node) && - mmvq_shared_q8_1_reader_count(node->src[1]) > 1) { + (mmvq_shared_q8_1_reader_count(node->src[1]) > 1 || (!mmvq_q8_1_producer_legacy && + cuda_ctx->mmvq_shared_q8_1.find_any(node->src[1]) != nullptr))) { ggml_cuda_mmvq_shared_q8_1 & shared = cuda_ctx->mmvq_shared_q8_1; ggml_cuda_mmvq_shared_q8_1::key k; if (ggml_cuda_mmvq_shared_q8_1::key_of(node->src[1], k)) { @@ -6354,7 +6756,16 @@ static void ggml_cuda_graph_evaluate_and_capture(ggml_backend_cuda_context * cud ggml_cuda_graph_update_executable(cuda_ctx, graph_key); } // Launch graph - CUDA_CHECK(cudaGraphLaunch(graph->instance, cuda_ctx->stream())); + if (ggml_cuda_time_launch()) { + struct timespec t0, t1; + clock_gettime(CLOCK_MONOTONIC, &t0); + CUDA_CHECK(cudaGraphLaunch(graph->instance, cuda_ctx->stream())); + clock_gettime(CLOCK_MONOTONIC, &t1); + ggml_cuda_time_launch_record(cuda_ctx->device, + (t1.tv_sec - t0.tv_sec)*1e6 + (t1.tv_nsec - t0.tv_nsec)/1e3); + } else { + CUDA_CHECK(cudaGraphLaunch(graph->instance, cuda_ctx->stream())); + } #else GGML_UNUSED(graph_key); graph_evaluated_or_captured = true; @@ -6364,6 +6775,7 @@ static void ggml_cuda_graph_evaluate_and_capture(ggml_backend_cuda_context * cud // The VMM scratch pool is stack-like, so release all persistent allocations // explicitly in reverse order across both shared-Q8 and row-scale buffers. cuda_ctx->mmvq_shared_q8_1.reset(); + std::fill(std::begin(cuda_ctx->mmvq_moe_ids), std::end(cuda_ctx->mmvq_moe_ids), nullptr); for (auto it = gb10_pool_allocations.rbegin(); it != gb10_pool_allocations.rend(); ++it) { it->reset(); } @@ -6391,6 +6803,9 @@ static enum ggml_status ggml_backend_cuda_graph_compute(ggml_backend_t backend, ggml_cuda_set_device(cuda_ctx->device); cuda_ctx->pq2_tile_counters.ensure(); // before a capture can begin + cuda_ctx->ssm_conv_ab_slots.ensure(); + cuda_ctx->l2_issue_stop.ensure(); + cuda_ctx->hc_front_tickets.ensure(); bool use_cuda_graph = false; bool cuda_graph_update_required = false; @@ -6481,6 +6896,9 @@ static void ggml_backend_cuda_capture_begin(ggml_backend_t backend) { GGML_ASSERT(!cuda_ctx->outer_capture); ggml_cuda_set_device(cuda_ctx->device); cuda_ctx->pq2_tile_counters.ensure(); // before a capture can begin + cuda_ctx->ssm_conv_ab_slots.ensure(); + cuda_ctx->l2_issue_stop.ensure(); + cuda_ctx->hc_front_tickets.ensure(); { std::lock_guard lock(ggml_cuda_lock); ggml_cuda_lock_counter.fetch_add(1, std::memory_order_relaxed); @@ -6538,7 +6956,16 @@ static void * ggml_backend_cuda_capture_end(ggml_backend_t backend) { static void ggml_backend_cuda_capture_launch(ggml_backend_t backend, void * exec) { ggml_backend_cuda_context * cuda_ctx = (ggml_backend_cuda_context *) backend->context; ggml_cuda_set_device(cuda_ctx->device); - CUDA_CHECK(cudaGraphLaunch((cudaGraphExec_t) exec, cuda_ctx->stream())); + if (ggml_cuda_time_launch()) { + struct timespec t0, t1; + clock_gettime(CLOCK_MONOTONIC, &t0); + CUDA_CHECK(cudaGraphLaunch((cudaGraphExec_t) exec, cuda_ctx->stream())); + clock_gettime(CLOCK_MONOTONIC, &t1); + ggml_cuda_time_launch_record(cuda_ctx->device, + (t1.tv_sec - t0.tv_sec)*1e6 + (t1.tv_nsec - t0.tv_nsec)/1e3); + } else { + CUDA_CHECK(cudaGraphLaunch((cudaGraphExec_t) exec, cuda_ctx->stream())); + } } static void ggml_backend_cuda_capture_free(ggml_backend_t backend, void * exec) { diff --git a/ggml/src/ggml-cuda/l2-issue.cu b/ggml/src/ggml-cuda/l2-issue.cu index 3c287c8410f5..176a488cca0c 100644 --- a/ggml/src/ggml-cuda/l2-issue.cu +++ b/ggml/src/ggml-cuda/l2-issue.cu @@ -13,16 +13,35 @@ static __device__ __forceinline__ uint64_t l2_issue_now() { return t; } -static __global__ void l2_issue_paced(const ggml_cuda_l2_ranges r, const float ns_per_byte) { +// A block reads stop every L2_ISSUE_STOP_EVERY pieces and compares what it read at the next read: the load's round trip +// to L2 overlaps the pieces between, and a block stops within 2 * L2_ISSUE_STOP_EVERY pieces of a change (~1.6 us at 40 +// GB/s a block) +#define L2_ISSUE_STOP_EVERY 4 + +static __device__ __forceinline__ unsigned int l2_issue_stop_read(const unsigned int * stop) { + unsigned int v; + asm volatile("ld.relaxed.gpu.global.u32 %0, [%1];" : "=r"(v) : "l"(stop) : "memory"); + return v; +} + +static __global__ void l2_issue_paced(const ggml_cuda_l2_ranges r, const float ns_per_byte, const unsigned int * stop) { #if !defined(GGML_USE_HIP) && !defined(GGML_USE_MUSA) && defined(__CUDA_ARCH__) && __CUDA_ARCH__ >= GGML_CUDA_CC_HOPPER - const uint64_t t0 = l2_issue_now(); - const int64_t step = (int64_t) gridDim.x * L2_ISSUE_PIECE; - int k = 0; // the range holding lo - int64_t base = 0; // its first byte's place in the ranges laid end to end - for (int64_t off = (int64_t) blockIdx.x * L2_ISSUE_PIECE; off < r.total; off += step) { + const uint64_t t0 = l2_issue_now(); + const int64_t step = (int64_t) gridDim.x * L2_ISSUE_PIECE; + const unsigned int start = stop != nullptr ? l2_issue_stop_read(stop) : 0; + unsigned int seen = start; + int k = 0; // the range holding lo + int64_t base = 0; // its first byte's place in the ranges laid end to end + for (int64_t off = (int64_t) blockIdx.x * L2_ISSUE_PIECE, n = 0; off < r.total; off += step, ++n) { const uint64_t due = (uint64_t) ((float) off * ns_per_byte); while (l2_issue_now() - t0 < due) { } + if (stop != nullptr && n % L2_ISSUE_STOP_EVERY == 0) { + if (seen != start) { + return; // a launch that keeps DRAM busy has started, and a request now would take DRAM from it + } + seen = l2_issue_stop_read(stop); + } const int64_t hi = min(off + (int64_t) L2_ISSUE_PIECE, r.total); for (int64_t lo = off; lo < hi;) { while (lo >= base + r.bytes[k]) { @@ -36,11 +55,12 @@ static __global__ void l2_issue_paced(const ggml_cuda_l2_ranges r, const float n } } #else - GGML_UNUSED_VARS(r, ns_per_byte); + GGML_UNUSED_VARS(r, ns_per_byte, stop); #endif } -void ggml_cuda_l2_issue(const ggml_cuda_l2_ranges & r, const double rate_gbs, const int nsm, cudaStream_t stream) { +void ggml_cuda_l2_issue(const ggml_cuda_l2_ranges & r, const double rate_gbs, const int nsm, const unsigned int * stop, + cudaStream_t stream) { if (r.total <= 0 || rate_gbs <= 0.0) { return; } @@ -48,6 +68,6 @@ void ggml_cuda_l2_issue(const ggml_cuda_l2_ranges & r, const double rate_gbs, co // rate, and no more blocks than pieces const int64_t pieces = (r.total + L2_ISSUE_PIECE - 1) / L2_ISSUE_PIECE; const int nb = (int) std::min({ (int64_t) nsm, pieces, (int64_t) std::ceil(rate_gbs / 40.0) }); - l2_issue_paced<<>>(r, (float) (1.0 / rate_gbs)); + l2_issue_paced<<>>(r, (float) (1.0 / rate_gbs), stop); CUDA_CHECK(cudaGetLastError()); } diff --git a/ggml/src/ggml-cuda/l2-issue.cuh b/ggml/src/ggml-cuda/l2-issue.cuh index b014bc151ce1..e594d4c3f639 100644 --- a/ggml/src/ggml-cuda/l2-issue.cuh +++ b/ggml/src/ggml-cuda/l2-issue.cuh @@ -50,6 +50,8 @@ struct ggml_cuda_l2_ranges { // Requests r into L2 at rate_gbs on stream: a burst of bulk prefetches beyond what DRAM serves at once is dropped (a // 17.8 MB matrix prefetched at once from 70 SMs of an RTX 5070 Ti lands ~4 MB of it, paced at up to 1,000 GB/s all of it), -// so the kernel spaces its requests out in time. Runs until the last range is requested, total / rate_gbs; nothing waits -// on it but the end of the graph. NVIDIA from sm_90 (cp.async.bulk.prefetch.L2); nothing elsewhere. -void ggml_cuda_l2_issue(const ggml_cuda_l2_ranges & r, double rate_gbs, int nsm, cudaStream_t stream); +// so the kernel spaces its requests out in time. Runs until the last range is requested, total / rate_gbs, or until +// stop (ggml_cuda_l2_issue_stop, or nullptr) changes from its value when the kernel started; nothing waits on it but the +// end of the graph. NVIDIA from sm_90 (cp.async.bulk.prefetch.L2); nothing elsewhere. +void ggml_cuda_l2_issue(const ggml_cuda_l2_ranges & r, double rate_gbs, int nsm, const unsigned int * stop, + cudaStream_t stream); diff --git a/ggml/src/ggml-cuda/lightning-indexer.cu b/ggml/src/ggml-cuda/lightning-indexer.cu index 38cb0fdcf22c..13130150d98d 100644 --- a/ggml/src/ggml-cuda/lightning-indexer.cu +++ b/ggml/src/ggml-cuda/lightning-indexer.cu @@ -249,7 +249,8 @@ static __global__ void lightning_indexer_kernel_vec( size_t nbk1, size_t nbk2, size_t nbk3, size_t nbw1, size_t nbw2, size_t nbw3, size_t nbm1, size_t nbm2, size_t nbm3, - int64_t nem3 + int64_t nem3, + const int32_t * rows, int64_t s_rows // ggml_lightning_indexer_rows: key i of stream s at K's row rows[s*s_rows + i] ) { constexpr int K_VECS_PER_WARP = K_VECS_PER_BLOCK / WARPS_PER_BLOCK; @@ -267,6 +268,7 @@ static __global__ void lightning_indexer_kernel_vec( const char * q_base = (const char *) Q + i_batch*nbq2 + i_stream*nbq3; const float * w_base = (const float *) ((const char *) W + i_batch*nbw1 + i_stream*nbw3); + const int32_t * rows_s = rows ? rows + i_stream*s_rows : nullptr; // phase 1 - load (and dequantize if needed) K to registers @@ -278,7 +280,8 @@ static __global__ void lightning_indexer_kernel_vec( for (int k = 0; k < K_VECS_PER_WARP; ++k) { int i_kv = start_kv + k; if (i_kv < n_kv) { - const float4 * k_base = (const float4 *) ((const char *) K + i_kv*nbk2 + i_stream*nbk3); + const int64_t i_row = rows_s ? rows_s[i_kv] : i_kv; + const float4 * k_base = (const float4 *) ((const char *) K + i_row*nbk2 + i_stream*nbk3); k_reg_f[k] = k_base[i_lane]; } else { k_reg_f[k] = make_float4(0, 0, 0, 0); @@ -291,7 +294,8 @@ static __global__ void lightning_indexer_kernel_vec( for (int k = 0; k < K_VECS_PER_WARP; ++k) { int i_kv = start_kv + k; if (i_kv < n_kv) { - const void * k_base = (const void *) ((const char *) K + i_kv*nbk2 + i_stream*nbk3); + const int64_t i_row = rows_s ? rows_s[i_kv] : i_kv; + const void * k_base = (const void *) ((const char *) K + i_row*nbk2 + i_stream*nbk3); dequantize_k(k_base, &k_reg_f[k], i_lane * 4); } else { k_reg_f[k] = make_float4(0, 0, 0, 0); @@ -382,6 +386,149 @@ static __global__ void lightning_indexer_kernel_vec( } } +// The vector kernel's scores bit for bit, a key to a quad of lanes: lane s of a quad holds the dims lanes 4j + s +// (j = 0..7) hold in the vector kernel, 16j + 4s to 16j + 4s + 3, and adds its 8 products in registers in the order +// warp_reduce_sum's xor 16, 8 and 4 add them (the bits of j), then xor 2 and 1 across the quad: 2 shuffles a key and head +// where the vector kernel takes 5 and serves one key with them. A block holds all heads' q at once. +template +static __global__ void lightning_indexer_kernel_quad( + const float * Q, const char * K, const float * W, const half * M, float * dst, + int64_t n_stream, int64_t n_batch, int64_t n_kv, + size_t nb1, size_t nb2, size_t nb3, + size_t nbq1, size_t nbq2, size_t nbq3, + size_t nbk1, size_t nbk2, size_t nbk3, + size_t nbw1, size_t nbw2, size_t nbw3, + size_t nbm1, size_t nbm2, size_t nbm3, + int64_t nem3, + const int32_t * rows, int64_t s_rows + ) { + static_assert(N_EMBD == 4*WARP_SIZE, "a lane of the vector kernel holds 4 dims"); + + constexpr int N_J = N_EMBD / 16; // float4s a lane holds of a key + constexpr int QUADS_PER_WARP = WARP_SIZE / 4; + constexpr int KEYS_PER_WARP = QUADS_PER_WARP * KEYS_PER_QUAD; + constexpr int THREADS_PER_BLOCK = WARPS_PER_BLOCK * WARP_SIZE; + static_assert(THREADS_PER_BLOCK >= N_HEAD, "a thread loads a head's weight"); + + const int i_batch = blockIdx.y; + const int i_stream = blockIdx.z; + const int i_warp = threadIdx.y; + const int i_lane = threadIdx.x; + const int tid = i_warp * WARP_SIZE + i_lane; + const int i_quad = i_lane / 4; + const int s = i_lane % 4; + + // key c of a quad: consecutive quads hold consecutive keys, so a warp writes its scores contiguously + const int start_kv = (blockIdx.x * WARPS_PER_BLOCK + i_warp) * KEYS_PER_WARP + i_quad; + + const char * q_base = (const char *) Q + i_batch*nbq2 + i_stream*nbq3; + const float * w_base = (const float *) ((const char *) W + i_batch*nbw1 + i_stream*nbw3); + const int32_t * rows_s = rows ? rows + i_stream*s_rows : nullptr; + + float4 k_reg_f[KEYS_PER_QUAD][N_J]; + +#pragma unroll + for (int c = 0; c < KEYS_PER_QUAD; ++c) { + const int i_kv = start_kv + c*QUADS_PER_WARP; + if (i_kv < n_kv) { + const int64_t i_row = rows_s ? rows_s[i_kv] : i_kv; + const char * k_base = K + i_row*nbk2 + i_stream*nbk3; + if constexpr (TYPE_K == GGML_TYPE_F32) { +#pragma unroll + for (int j = 0; j < N_J; ++j) { + k_reg_f[c][j] = ((const float4 *) k_base)[4*j + s]; + } + } else { + constexpr dequantize_V_t dequantize_k = get_dequantize_V(); +#pragma unroll + for (int j = 0; j < N_J; ++j) { + dequantize_k(k_base, &k_reg_f[c][j], 16*j + 4*s); + } + } + } else { +#pragma unroll + for (int j = 0; j < N_J; ++j) { + k_reg_f[c][j] = make_float4(0, 0, 0, 0); + } + } + } + + __shared__ float w_shared[N_HEAD]; + __shared__ float4 q_shared_f[N_HEAD][N_EMBD / 4]; + + if (tid < N_HEAD) { + w_shared[tid] = w_base[tid]; + } +#pragma unroll + for (int i_q = tid; i_q < N_HEAD * (N_EMBD / 4); i_q += THREADS_PER_BLOCK) { + const int i_head = i_q / (N_EMBD / 4); + const int i_embd = i_q % (N_EMBD / 4); + q_shared_f[i_head][i_embd] = *(const float4 *) (q_base + i_head*nbq1 + i_embd*sizeof(float4)); + } + + __syncthreads(); + + float score_k[KEYS_PER_QUAD] = { 0.0f }; + + for (int i_head = 0; i_head < N_HEAD; ++i_head) { + const float w_val = w_shared[i_head]; + float qk[KEYS_PER_QUAD][N_J]; + +#pragma unroll + for (int j = 0; j < N_J; ++j) { + const float4 q_vec = q_shared_f[i_head][4*j + s]; +#pragma unroll + for (int c = 0; c < KEYS_PER_QUAD; ++c) { + qk[c][j] = 0.0f; + ggml_cuda_mad(qk[c][j], q_vec.x, k_reg_f[c][j].x); + ggml_cuda_mad(qk[c][j], q_vec.y, k_reg_f[c][j].y); + ggml_cuda_mad(qk[c][j], q_vec.z, k_reg_f[c][j].z); + ggml_cuda_mad(qk[c][j], q_vec.w, k_reg_f[c][j].w); + } + } + +#pragma unroll + for (int c = 0; c < KEYS_PER_QUAD; ++c) { + static_assert(N_J == 8, "the in-register levels are xor 16, 8 and 4"); + const float x16_0 = qk[c][0] + qk[c][4]; + const float x16_1 = qk[c][1] + qk[c][5]; + const float x16_2 = qk[c][2] + qk[c][6]; + const float x16_3 = qk[c][3] + qk[c][7]; + const float x8_0 = x16_0 + x16_2; + const float x8_1 = x16_1 + x16_3; + float sum = x8_0 + x8_1; + sum += __shfl_xor_sync(0xffffffff, sum, 2, WARP_SIZE); + sum += __shfl_xor_sync(0xffffffff, sum, 1, WARP_SIZE); + + // ReLU, weight + sum = (sum > 0.0f) ? sum : 0.0f; + score_k[c] += sum * w_val; + } + } + + if (s == 0) { + const half * m_base = (const half *) ((const char *) M + i_batch*nbm1 + (i_stream%nem3)*nbm3); + float * dst_base = (float *) ((char *) dst + i_batch*nb1 + i_stream*nb3); +#pragma unroll + for (int c = 0; c < KEYS_PER_QUAD; ++c) { + const int i_kv = start_kv + c*QUADS_PER_WARP; + if (i_kv < n_kv) { + dst_base[i_kv] = score_k[c] + __half2float(m_base[i_kv]); + } + } + } +} + +// GGML_CUDA_LIGHTNING_INDEXER_CHECK=1: the vector kernel runs beside the quad kernel and a differing bit traps +static __global__ void lightning_indexer_check(const float * quad, const float * vec, int64_t n) { + const int64_t i = (int64_t) blockIdx.x*blockDim.x + threadIdx.x; + if (i < n && __float_as_uint(quad[i]) != __float_as_uint(vec[i])) { + printf("lightning indexer check: score %lld is %08x from the quad kernel, %08x from the vector kernel\n", + (long long) i, __float_as_uint(quad[i]), __float_as_uint(vec[i])); + __trap(); + } +} + #define LIGHTNING_INDEXER_CASE(lightning_indexer_kernel, n_embd, n_head, K, type_K) \ if (K->type == (type_K)) { \ lightning_indexer_kernel \ @@ -397,11 +544,50 @@ static __global__ void lightning_indexer_kernel_vec( ); \ } else +// the quad kernel's launch: 64 keys a block, the vector kernel's grid +static constexpr int LIGHTNING_INDEXER_QUAD_WARPS = 4; +static constexpr int LIGHTNING_INDEXER_QUAD_KEYS = 2; + +// the quad kernel into dst, the vector kernel into vec_d (dst, or the check's buffer), whichever of them run +#define LIGHTNING_INDEXER_VEC_CASE(n_embd, n_head, K, type_K) \ + if (K->type == (type_K)) { \ + if (!vec_legacy) { \ + lightning_indexer_kernel_quad \ + <<>>( \ + q_d, k_d, w_d, m_d, dst_d, \ + n_stream, n_batch, n_kv, \ + nb1, nb2, nb3, \ + nbq1, nbq2, nbq3, \ + nbk1, nbk2, nbk3, \ + nbw1, nbw2, nbw3, \ + nbm1, nbm2, nbm3, \ + nem3, \ + rows_d, s_rows \ + ); \ + } \ + if (vec_d != nullptr) { \ + lightning_indexer_kernel_vec \ + <<>>( \ + q_d, k_d, w_d, m_d, vec_d, \ + n_stream, n_batch, n_kv, \ + nb1, nb2, nb3, \ + nbq1, nbq2, nbq3, \ + nbk1, nbk2, nbk3, \ + nbw1, nbw2, nbw3, \ + nbm1, nbm2, nbm3, \ + nem3, \ + rows_d, s_rows \ + ); \ + } \ + } else + void ggml_cuda_lightning_indexer(ggml_backend_cuda_context & ctx, ggml_tensor * dst) { const ggml_tensor * q = dst->src[0]; const ggml_tensor * k = dst->src[1]; const ggml_tensor * w = dst->src[2]; // weights const ggml_tensor * m = dst->src[3]; // mask + const ggml_tensor * r = dst->src[4]; // ggml_lightning_indexer_rows, or none GGML_ASSERT(dst->type == GGML_TYPE_F32); GGML_ASSERT( q->type == GGML_TYPE_F32); @@ -438,7 +624,7 @@ void ggml_cuda_lightning_indexer(ggml_backend_cuda_context & ctx, ggml_tensor * const int n_head = q->ne[1]; const int n_batch = q->ne[2]; const int n_stream = q->ne[3]; - const int n_kv = k->ne[2]; + const int n_kv = dst->ne[0]; const float * q_d = (const float *) q->data; const char * k_d = (const char *) k->data; @@ -446,12 +632,25 @@ void ggml_cuda_lightning_indexer(ggml_backend_cuda_context & ctx, ggml_tensor * const half * m_d = (const half *) m->data; float * dst_d = ( float *) dst->data; + // rows: the vector kernel, which scores in f32 whatever k's type (the wmma kernel rounds q to f16) + const int32_t * rows_d = r ? (const int32_t *) r->data : nullptr; + const int64_t s_rows = r ? r->nb[1] / sizeof(int32_t) : 0; + + // the vector kernel's cases run the quad kernel; GGML_CUDA_LIGHTNING_INDEXER_VEC_LEGACY=1 the vector kernel, + // GGML_CUDA_LIGHTNING_INDEXER_CHECK=1 both (the vector kernel into a buffer the check compares) + static const bool vec_legacy = ggml_env_switch("GGML_CUDA_LIGHTNING_INDEXER_VEC_LEGACY"); + static const bool vec_check = !vec_legacy && ggml_env_switch("GGML_CUDA_LIGHTNING_INDEXER_CHECK"); + ggml_cuda_pool_alloc vec_check_buf(ctx.pool()); + float * vec_d = vec_legacy ? dst_d : nullptr; + const dim3 block_quad(WARP_SIZE, LIGHTNING_INDEXER_QUAD_WARPS); + const int device = ggml_cuda_get_device(); const int cc = ggml_cuda_info().devices[device].cc; if (n_embd == 128 && n_head == 64) { #if !defined(GGML_USE_HIP) && !defined(GGML_USE_MUSA) - if (GGML_CUDA_CC_IS_NVIDIA(cc) && turing_mma_available(cc) && k->type != GGML_TYPE_F32 && k->type != GGML_TYPE_BF16) { + if (GGML_CUDA_CC_IS_NVIDIA(cc) && turing_mma_available(cc) && k->type != GGML_TYPE_F32 && k->type != GGML_TYPE_BF16 && + rows_d == nullptr) { // use wmma kernel constexpr int K_VECS_PER_BLOCK = 32; constexpr int WARPS_PER_BLOCK = 8; @@ -480,19 +679,26 @@ void ggml_cuda_lightning_indexer(ggml_backend_cuda_context & ctx, ggml_tensor * int num_kv_blocks = (n_kv + (K_VECS_PER_BLOCK) - 1) / (K_VECS_PER_BLOCK); dim3 grid(num_kv_blocks, n_batch, n_stream); - LIGHTNING_INDEXER_CASE(lightning_indexer_kernel_vec, 128, 64, k, GGML_TYPE_F16) - LIGHTNING_INDEXER_CASE(lightning_indexer_kernel_vec, 128, 64, k, GGML_TYPE_Q4_0) - LIGHTNING_INDEXER_CASE(lightning_indexer_kernel_vec, 128, 64, k, GGML_TYPE_Q4_1) - LIGHTNING_INDEXER_CASE(lightning_indexer_kernel_vec, 128, 64, k, GGML_TYPE_Q5_0) - LIGHTNING_INDEXER_CASE(lightning_indexer_kernel_vec, 128, 64, k, GGML_TYPE_Q5_1) - LIGHTNING_INDEXER_CASE(lightning_indexer_kernel_vec, 128, 64, k, GGML_TYPE_Q8_0) - LIGHTNING_INDEXER_CASE(lightning_indexer_kernel_vec, 128, 64, k, GGML_TYPE_BF16) - LIGHTNING_INDEXER_CASE(lightning_indexer_kernel_vec, 128, 64, k, GGML_TYPE_F32) + static_assert(LIGHTNING_INDEXER_QUAD_WARPS * (WARP_SIZE/4) * LIGHTNING_INDEXER_QUAD_KEYS == K_VECS_PER_BLOCK, + "the quad kernel takes the vector kernel's grid"); + if (vec_check) { + vec_d = vec_check_buf.alloc(ggml_nelements(dst)); + } + + LIGHTNING_INDEXER_VEC_CASE(128, 64, k, GGML_TYPE_F16) + LIGHTNING_INDEXER_VEC_CASE(128, 64, k, GGML_TYPE_Q4_0) + LIGHTNING_INDEXER_VEC_CASE(128, 64, k, GGML_TYPE_Q4_1) + LIGHTNING_INDEXER_VEC_CASE(128, 64, k, GGML_TYPE_Q5_0) + LIGHTNING_INDEXER_VEC_CASE(128, 64, k, GGML_TYPE_Q5_1) + LIGHTNING_INDEXER_VEC_CASE(128, 64, k, GGML_TYPE_Q8_0) + LIGHTNING_INDEXER_VEC_CASE(128, 64, k, GGML_TYPE_BF16) + LIGHTNING_INDEXER_VEC_CASE(128, 64, k, GGML_TYPE_F32) GGML_ABORT("fatal error"); } } else if (n_embd == 128 && n_head == 32) { #if !defined(GGML_USE_HIP) && !defined(GGML_USE_MUSA) - if (GGML_CUDA_CC_IS_NVIDIA(cc) && turing_mma_available(cc) && k->type != GGML_TYPE_F32 && k->type != GGML_TYPE_BF16) { + if (GGML_CUDA_CC_IS_NVIDIA(cc) && turing_mma_available(cc) && k->type != GGML_TYPE_F32 && k->type != GGML_TYPE_BF16 && + rows_d == nullptr) { // use wmma kernel constexpr int K_VECS_PER_BLOCK = 32; constexpr int WARPS_PER_BLOCK = 8; @@ -521,19 +727,31 @@ void ggml_cuda_lightning_indexer(ggml_backend_cuda_context & ctx, ggml_tensor * int num_kv_blocks = (n_kv + (K_VECS_PER_BLOCK) - 1) / (K_VECS_PER_BLOCK); dim3 grid(num_kv_blocks, n_batch, n_stream); - LIGHTNING_INDEXER_CASE(lightning_indexer_kernel_vec, 128, 32, k, GGML_TYPE_F16) - LIGHTNING_INDEXER_CASE(lightning_indexer_kernel_vec, 128, 32, k, GGML_TYPE_Q4_0) - LIGHTNING_INDEXER_CASE(lightning_indexer_kernel_vec, 128, 32, k, GGML_TYPE_Q4_1) - LIGHTNING_INDEXER_CASE(lightning_indexer_kernel_vec, 128, 32, k, GGML_TYPE_Q5_0) - LIGHTNING_INDEXER_CASE(lightning_indexer_kernel_vec, 128, 32, k, GGML_TYPE_Q5_1) - LIGHTNING_INDEXER_CASE(lightning_indexer_kernel_vec, 128, 32, k, GGML_TYPE_Q8_0) - LIGHTNING_INDEXER_CASE(lightning_indexer_kernel_vec, 128, 32, k, GGML_TYPE_BF16) - LIGHTNING_INDEXER_CASE(lightning_indexer_kernel_vec, 128, 32, k, GGML_TYPE_F32) + static_assert(LIGHTNING_INDEXER_QUAD_WARPS * (WARP_SIZE/4) * LIGHTNING_INDEXER_QUAD_KEYS == K_VECS_PER_BLOCK, + "the quad kernel takes the vector kernel's grid"); + if (vec_check) { + vec_d = vec_check_buf.alloc(ggml_nelements(dst)); + } + + LIGHTNING_INDEXER_VEC_CASE(128, 32, k, GGML_TYPE_F16) + LIGHTNING_INDEXER_VEC_CASE(128, 32, k, GGML_TYPE_Q4_0) + LIGHTNING_INDEXER_VEC_CASE(128, 32, k, GGML_TYPE_Q4_1) + LIGHTNING_INDEXER_VEC_CASE(128, 32, k, GGML_TYPE_Q5_0) + LIGHTNING_INDEXER_VEC_CASE(128, 32, k, GGML_TYPE_Q5_1) + LIGHTNING_INDEXER_VEC_CASE(128, 32, k, GGML_TYPE_Q8_0) + LIGHTNING_INDEXER_VEC_CASE(128, 32, k, GGML_TYPE_BF16) + LIGHTNING_INDEXER_VEC_CASE(128, 32, k, GGML_TYPE_F32) GGML_ABORT("fatal error"); } } else { GGML_ABORT("fatal error"); } + + if (vec_check && vec_d != nullptr) { + GGML_ASSERT(ggml_is_contiguous(dst)); + const int64_t n = ggml_nelements(dst); + lightning_indexer_check<<<(n + 255) / 256, 256, 0, ctx.stream()>>>(dst_d, vec_d, n); + } } bool ggml_cuda_lightning_indexer_supported(int device, const ggml_tensor * dst) { @@ -563,6 +781,11 @@ bool ggml_cuda_lightning_indexer_supported(int device, const ggml_tensor * dst) return false; } + const ggml_tensor * r = dst->src[4]; + if (r != nullptr && (r->type != GGML_TYPE_I32 || !ggml_is_contiguous(r))) { + return false; + } + // alignment checks: the kernels read q as float4 and a float k as int2 or float4, so the data addresses must be // multiples of 16 bytes as well as the strides (a view offset by an odd number of elements would fault with a // misaligned address). This runs in supports_op, before allocation, and again at dispatch. diff --git a/ggml/src/ggml-cuda/mmvf.cu b/ggml/src/ggml-cuda/mmvf.cu index bcc05959d460..46187ffee20a 100644 --- a/ggml/src/ggml-cuda/mmvf.cu +++ b/ggml/src/ggml-cuda/mmvf.cu @@ -10,7 +10,13 @@ static __global__ void mul_mat_vec_f( const int ncols2, const uint3 nchannels_y, const int stride_row, const int stride_col_y2, const int stride_col_dst, const uint3 channel_ratio, const int stride_channel_x, const int stride_channel_y, const int stride_channel_dst, const uint3 sample_ratio, const int stride_sample_x, const int stride_sample_y, const int stride_sample_dst, - const int ids_stride) { + const int ids_stride, const bool lc_early) { + // lc_early: the next kernel may launch once every block has started, so it lands while this one reads and its + // requests before its own wait (the Gated DeltaNet's state, after qwen35's alpha/beta pair) run under this one's; + // it still waits for this grid to complete, so no result depends on where the trigger sits + if (lc_early) { + ggml_cuda_pdl_lc(); + } const T * GGML_CUDA_RESTRICT x = x_ptr; const float * GGML_CUDA_RESTRICT y = y_ptr; const int32_t * GGML_CUDA_RESTRICT ids = ids_ptr; @@ -341,7 +347,9 @@ static __global__ void mul_mat_vec_f( static_assert(std::is_same_v, "unsupported type"); } - ggml_cuda_pdl_lc(); + if (!lc_early) { + ggml_cuda_pdl_lc(); + } #pragma unroll for (int j = 0; j < ncols_dst; ++j) { sumf[j] = warp_reduce_sum(sumf[j]); @@ -651,13 +659,16 @@ static void mul_mat_vec_f_switch_fusion( const ggml_cuda_kernel_launch_params launch_params = {block_nums, block_dims, nbytes_shared, stream}; + // the kernel triggers the next launch at its start; GGML_CUDA_MMVF_TRIGGER_LEGACY=1 after its dot products + static const bool lc_legacy = ggml_env_switch("GGML_CUDA_MMVF_TRIGGER_LEGACY"); + const bool has_fusion = fusion.gate != nullptr || fusion.x_bias != nullptr || fusion.gate_bias != nullptr || fusion.rms_norm; if constexpr (ncols_dst == 1) { if (has_fusion) { ggml_cuda_kernel_launch(mul_mat_vec_f, launch_params, x, y, ids, fusion, dst, ncols, nchannels_y, stride_row, stride_col_y, stride_col_dst, channel_ratio, stride_channel_x, stride_channel_y, stride_channel_dst, - sample_ratio, stride_sample_x, stride_sample_y, stride_sample_dst, ids_stride); + sample_ratio, stride_sample_x, stride_sample_y, stride_sample_dst, ids_stride, !lc_legacy); return; } } @@ -667,7 +678,7 @@ static void mul_mat_vec_f_switch_fusion( ggml_cuda_kernel_launch(mul_mat_vec_f, launch_params, x, y, ids, fusion, dst, ncols, nchannels_y, stride_row, stride_col_y, stride_col_dst, channel_ratio, stride_channel_x, stride_channel_y, stride_channel_dst, - sample_ratio, stride_sample_x, stride_sample_y, stride_sample_dst, ids_stride); + sample_ratio, stride_sample_x, stride_sample_y, stride_sample_dst, ids_stride, !lc_legacy); } diff --git a/ggml/src/ggml-cuda/mmvq-moe.cu b/ggml/src/ggml-cuda/mmvq-moe.cu index c18124241769..f0f1c9fe39ff 100644 --- a/ggml/src/ggml-cuda/mmvq-moe.cu +++ b/ggml/src/ggml-cuda/mmvq-moe.cu @@ -36,7 +36,9 @@ // kernel spills. The math is bound by shared memory and L1, not by the dot products: the grid is gathered at random, // so past the ring the shared memory holds it a copy per lane (grid_rep[i*32 + lane]: a warp's 32 gathers in 32 banks, // whatever the indices), and a lane keeps its q8_1 fragments in registers, loaded once per token vector, not once per -// row. The other types read as mul_mat_vec_q does (vec_dot_q_cuda), from the slot, which their team releases after it. +// row. In a launch of several tokens whose pairs are as many as the experts or more, a fragment is decoded once (its grid +// gathered, its signs applied) and met with the vectors of its expert's pairs in turn (pairs_once). The other types read as mul_mat_vec_q does (vec_dot_q_cuda), from +// the slot, which their team releases after it. // // The tokens' q8_1 vectors come from a copy in shared memory past the ring when it fits there, which the producer's // first bulk copy makes, before any tile: a global load that a consumer issues behind the ring's copies waits for them. @@ -45,7 +47,17 @@ // ~3.5 us), and at 3 tokens every change of vector paid it. // // No pointer carries __restrict__: with PDL a restrict load may compile to ld.global.nc, which the compiler can move -// above the grid dependency wait (upstream #24030). Nothing is read before that wait: the experts come from ids. +// above the grid dependency wait (upstream #24030). Nothing is read before that wait but ids when they are whole before +// the launch (ids_ready), and the weights of the tiles they name; nothing is written but the stop word of the L2 issue +// beside the stream (ggml_cuda_l2_issue_stop), an atomic add that no kernel on the stream reads. +// +// A launch triggers the next only past its dependency wait, so whatever the kernels before it wrote is whole before any +// kernel after it on the stream starts. The ids a ring launch reads (a MoE layer's top-k, written by the kernels before +// it) are then whole when a later ring launch reading the same ids starts: the down projection after a gate/up. That +// launch lists its experts and issues its block's first tiles before its wait, so its first rows stream during the +// gate/up's last tiles and the kernels between (the q8_1 of the GLU's output), where they started only past both (RTX +// 5070 Ti, GLM-5.3's proxy under -sm tensor: the down 18.75 us a layer against 14.5 at the DRAM peak, the gate/up +// 34.34 against 29.0). #define MMVQ_MOE_NG 8 // a team's warps: a tile's row groups #define MMVQ_MOE_NT 2 // teams, each on its own tiles @@ -54,6 +66,8 @@ #define MMVQ_MOE_MAX_SLOTS 8 #define MMVQ_MOE_SMEM_MAX (99 * 1024) // the shared memory a block may take on sm_90 - sm_120 #define MMVQ_MOE_SLOT_TARGET (16 * 1024) // a slot's bytes past one row a warp: more slots, a shorter tail +#define MMVQ_MOE_PB 4 // an expert's pairs a decoded fragment meets at once: an MTP verify's + // 1-3 drafts and their token in one pass static_assert(MMVQ_MOE_MAX_PAIRS == 64, "the producer warp lists the pairs, two a lane"); static_assert(MMVQ_MOE_MAX_PAIRS <= 64, "an expert's pairs are a 64-bit mask"); @@ -71,25 +85,32 @@ struct mmvq_moe_dev_args { const float * gate_bias; float * dst; int64_t stride_channel_x_bytes; - int64_t stride_col_dst; - int64_t stride_channel_dst; int64_t stride_bias; int row_bytes; int box_bytes; // a matrix's part of a slot: 8*RPW rows, padded to 128 bytes int nrows; int nb; // quant blocks a row int ntr; // row tiles an expert + uint3 ntr_fd; // ntr and nslots as fast_div_modulo takes them: a tile's expert and row tile, and + uint3 nslots_fd; // a sequence's slot and phase, each a multiply and a shift in place of a division int n_used; int ntokens; int ids_stride; - int nchannels_y; - int stride_col_y; - int stride_channel_y; int glu_op; float glu_limit; int nslots; int * tile_ctr; int y_bytes; // > 0: the tokens' vectors, [y, y + y_bytes), copied into shared memory past the ring + const float * y_f32; // non-null: the vectors in f32 (slot s of token t at s*y_f32_s1 + t*y_f32_s2), which + int y_f32_s1; // the consumers quantize into that copy (mmvq_moe_quantize_y) where the producer + int y_f32_s2; // would copy y + int ids_ready; // ids whole before the launch (ggml_cuda_mmvq_moe_args): read before the dependency wait + unsigned int * l2_issue_stop; // bumped as the reads start (ggml_cuda_l2_issue_stop), or nullptr + // pair p's (token p / n_used, slot p % n_used) vector, in blocks past y_base (its slot's % nchannels_y), and its + // first dst element: made on the host, where each pair of a tile took them from 3 divisions (~15 % of the down's + // instructions at 3 tokens). Read at a pair index the warp shares, a constant load each + int pair_y[MMVQ_MOE_MAX_PAIRS]; + int64_t pair_dst[MMVQ_MOE_MAX_PAIRS]; }; // IQ3_XXS: a lane's k iterations at rpw rows a warp, 4 blocks of a row an iteration, the most its registers hold. The @@ -106,6 +127,14 @@ static constexpr bool mmvq_moe_fits(const ggml_type type, const int nmat, const return type != GGML_TYPE_IQ3_XXS || rpw*nmat*mmvq_moe_iq3_nit(rpw) <= 8; } +// Whether an instance may decode each fragment once for all the pairs of its expert (pairs_once, a launch of several +// tokens): IQ3_XXS, whose decode is most of its math, where the decoded fragments and the pairs' sums fit the registers +// beside the slot's fragments. Its own instance: beside the one-pair path's registers (a vector's fragments kept across +// tiles) the kernel spilled, and so did the one-pair path (GLM-5.3's proxy at 3 tokens: the gate/up 108 -> 139 us) +static constexpr __host__ __device__ bool mmvq_moe_pairs_fits(const ggml_type type, const int nmat, const int rpw) { + return type == GGML_TYPE_IQ3_XXS && rpw*nmat <= 2; +} + #ifdef MMVQ_MOE_AVAILABLE static __device__ __forceinline__ uint32_t mmvq_moe_smem_u32(const void * p) { return (uint32_t) __cvta_generic_to_shared(p); @@ -159,14 +188,69 @@ static __device__ __forceinline__ void mmvq_moe_bulk_load(void * dst, const void // pair p's q8_1 vector: token p / n_used's, for expert slot p % n_used, in y or in its copy at y_base static __device__ __forceinline__ const block_q8_1 * mmvq_moe_y(const mmvq_moe_dev_args & a, const block_q8_1 * y_base, const int p) { - const int t = p / a.n_used; - const int slot = p % a.n_used; - return y_base + (int64_t) (slot % a.nchannels_y)*a.stride_channel_y + (int64_t) t*a.stride_col_y; + return y_base + a.pair_y[p]; } -template +#define MMVQ_MOE_QY_CHUNK 8 // float4s a consumer thread loads before it quantizes them: a round trip to L2 for all + +// a.y_f32: the consumer threads quantize the pairs' vectors into the copy at y_smem, bit for bit as quantize_q8_1 writes +// them. A thread takes a float4 of a vector, so 8 lanes take a q8_1 block, and the block's max and sum take +// warp_reduce_max/sum's tree: its xor 16, 8 and 4 are xor 4, 2 and 1 across the 8 lanes (the block's lane +// 4m + i is lane m's element i), its xor 2 and 1 add the lane's elements 0 + 2 and 1 + 3, then those two. +static __device__ __forceinline__ void mmvq_moe_quantize_y(const mmvq_moe_dev_args & a, char * y_smem, const int ncols) { + 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) { + float4 v[MMVQ_MOE_QY_CHUNK]; +#pragma unroll + for (int c = 0; c < MMVQ_MOE_QY_CHUNK; ++c) { + const int j = j0 + c*NT; + v[c] = make_float4(0.0f, 0.0f, 0.0f, 0.0f); + if (j < total) { + const int p = j / nf4; + v[c] = *(const float4 *) (a.y_f32 + (p % a.n_used)*a.y_f32_s1 + (p / a.n_used)*a.y_f32_s2 + 4*(j - p*nf4)); + } + } +#pragma unroll + for (int c = 0; c < MMVQ_MOE_QY_CHUNK; ++c) { + float s[4] = { v[c].x, v[c].y, v[c].z, v[c].w }; + float m = fmaxf(fmaxf(fabsf(s[0]), fabsf(s[1])), fmaxf(fabsf(s[2]), fabsf(s[3]))); +#pragma unroll + for (int o = 4; o > 0; o >>= 1) { + m = fmaxf(m, __shfl_xor_sync(0xFFFFFFFF, m, o, 8)); +#pragma unroll + for (int i = 0; i < 4; ++i) { + s[i] += __shfl_xor_sync(0xFFFFFFFF, s[i], o, 8); + } + } + const float sum = (s[0] + s[2]) + (s[1] + s[3]); + + const int j = j0 + c*NT; + if (j < total) { + const float d = m / 127.0f; + const float xs[4] = { v[c].x, v[c].y, v[c].z, v[c].w }; + char4 q; + q.x = m == 0.0f ? 0 : (int8_t) roundf(xs[0] / d); + q.y = m == 0.0f ? 0 : (int8_t) roundf(xs[1] / d); + q.z = m == 0.0f ? 0 : (int8_t) roundf(xs[2] / d); + q.w = m == 0.0f ? 0 : (int8_t) roundf(xs[3] / d); + const int p = j / nf4; + const int c4 = j - p*nf4; + block_q8_1 * yb = (block_q8_1 *) y_smem + a.pair_y[p] + c4/8; + *(char4 *) &yb->qs[4*(c4 % 8)] = q; + if (c4 % 8 == 0) { + yb->ds = make_half2(d, sum); + } + } + } + } +} + +template __launch_bounds__((MMVQ_MOE_NW + 1)*32, 1) static __global__ void mmvq_moe(const mmvq_moe_dev_args a) { + static_assert(!pairs_once || mmvq_moe_pairs_fits(type, nmat, rpw), "an instance that decodes once for the pairs"); #ifdef MMVQ_MOE_AVAILABLE constexpr int qk = ggml_cuda_type_traits::qk; constexpr int qi = ggml_cuda_type_traits::qi; @@ -193,8 +277,6 @@ static __global__ void mmvq_moe(const mmvq_moe_dev_args a) { const int warp = threadIdx.x / 32; const int lane = threadIdx.x % 32; - ggml_cuda_pdl_lc(); - if (threadIdx.x == 0) { for (int s = 0; s < a.nslots; ++s) { mmvq_moe_mbar_init(&full[s], 1); @@ -210,24 +292,33 @@ static __global__ void mmvq_moe(const mmvq_moe_dev_args a) { const block_q8_1 * const y_base = a.y_bytes > 0 ? (const block_q8_1 *) y_smem : a.y; if (warp == MMVQ_MOE_NW) { - // The producer warp alone lists the distinct experts, in the order of their first pairs, as soon as the - // dependency wait lets it read ids, and starts the loads while the consumers fill their tables (they read the - // lists after a slot's barrier, which lane 0's arrival publishes): pairs lane and lane + 32, a pair past the - // launch's with an expert of its own - ggml_cuda_pdl_sync(); // ids is a previous kernel's result - if (lane == 0 && a.y_bytes > 0) { - // the tokens' vectors, previous kernels' results too, before any tile: a global load that the consumers issue - // behind the ring's bulk copies waits for them (a gate/up's first tile took 7 us to its results where the - // next took 1.8, both teams at once, while the full ring kept the SM's rows from being asked for) - uint64_t unchanged; - asm volatile("createpolicy.fractional.L2::evict_unchanged.b64 %0, 1.0;" : "=l"(unchanged)); - mmvq_moe_mbar_arrive_expect_tx(&y_full, a.y_bytes); - mmvq_moe_bulk_load(y_smem, a.y, a.y_bytes, &y_full, unchanged); + // The producer warp alone lists the distinct experts, in the order of their first pairs, as soon as it may read + // ids, and starts the loads while the consumers fill their tables (they read the lists after a slot's barrier, + // which lane 0's arrival publishes): pairs lane and lane + 32, a pair past the launch's with an expert of its own. + // ids are a previous kernel's result: read past the dependency wait, or before it when they were whole before + // the launch (ids_ready), from L2 (ld.global.cg), past any line of theirs an earlier kernel left in this SM's L1. + // The tokens' vectors are previous kernels' results too, copied past the wait before any tile: a global load that + // the consumers issue behind the ring's bulk copies waits for them (a gate/up's first tile took 7 us to its + // results where the next took 1.8, both teams at once, while the full ring kept the SM's rows from being asked + // for); with ids_ready the block's first tiles are ahead of the copy, asked for before the wait. + uint64_t unchanged; + asm volatile("createpolicy.fractional.L2::evict_unchanged.b64 %0, 1.0;" : "=l"(unchanged)); + if (!a.ids_ready) { + ggml_cuda_pdl_sync(); + ggml_cuda_pdl_lc(); // past the wait: whatever the kernels before this one wrote is whole for the next + if (lane == 0 && a.y_bytes > 0 && a.y_f32 == nullptr) { + mmvq_moe_mbar_arrive_expect_tx(&y_full, a.y_bytes); + mmvq_moe_bulk_load(y_smem, a.y, a.y_bytes, &y_full, unchanged); + } + } + // the launch's reads start here (past the wait, or before it with ids_ready): an L2 issue beside the stream stops + if (a.l2_issue_stop != nullptr && blockIdx.x == 0 && lane == 0) { + atomicAdd(a.l2_issue_stop, 1u); } const int npairs = a.ntokens * a.n_used; const int p1 = lane + 32; - const int e0 = lane < npairs ? a.ids[lane % a.n_used + (lane / a.n_used)*a.ids_stride] : -1 - lane; - const int e1 = p1 < npairs ? a.ids[p1 % a.n_used + (p1 / a.n_used)*a.ids_stride] : -1 - p1; + const int e0 = lane < npairs ? __ldcg(&a.ids[lane % a.n_used + (lane / a.n_used)*a.ids_stride]) : -1 - lane; + const int e1 = p1 < npairs ? __ldcg(&a.ids[p1 % a.n_used + (p1 / a.n_used)*a.ids_stride]) : -1 - p1; pair_e[lane] = e0; pairs_of[lane] = 0; pairs_of[p1] = 0; @@ -280,10 +371,48 @@ static __global__ void mmvq_moe(const mmvq_moe_dev_args a) { const int last = n_tiles - dyn_base + (int) gridDim.x - 1; // the launch's last ticket int n_end = 0; int ticket = -1; // the next sequence's, once asked for - for (int i = 0;; ++i) { - const int s = i % a.nslots; + + // tile into sequence i's slot + const auto issue = [&](const int i, const int tile) { + const int s = (int) fastmodulo((uint32_t) i, a.nslots_fd); + held[s] = tile; + const uint2 ur = fast_div_modulo((uint32_t) tile, a.ntr_fd); + const int u = (int) ur.x; + const int row0 = (int) ur.y * R; + const uint32_t bytes = (uint32_t) (min(R, a.nrows - row0) * a.row_bytes); + const int64_t off = expert[u]*a.stride_channel_x_bytes + (int64_t) row0*a.row_bytes; + char * slot = ring + (size_t) s*nmat*a.box_bytes; + mmvq_moe_mbar_arrive_expect_tx(&full[s], nmat*bytes); + mmvq_moe_bulk_load(slot, a.vx + off, bytes, &full[s], policy); + if constexpr (nmat == 2) { + mmvq_moe_bulk_load(slot + a.box_bytes, a.vgate + off, bytes, &full[s], policy); + } + }; + + int i = 0; + if (a.ids_ready) { + // the block's first own tiles before the wait, a ring's at most: weights, which no kernel writes. Not the + // last own tile, whose issue asks for the launch's first ticket from the stream's counter, which the + // launch before may still be taking + for (; i < a.nslots && i + 1 < n_own; ++i) { + const int t = (int) blockIdx.x + i * (int) gridDim.x; + if (t >= dyn_base) { + break; + } + issue(i, t); + } + ggml_cuda_pdl_sync(); + ggml_cuda_pdl_lc(); + if (a.y_bytes > 0 && a.y_f32 == nullptr) { + mmvq_moe_mbar_arrive_expect_tx(&y_full, a.y_bytes); + mmvq_moe_bulk_load(y_smem, a.y, a.y_bytes, &y_full, unchanged); + } + } + for (;; ++i) { + const uint2 ph = fast_div_modulo((uint32_t) i, a.nslots_fd); + const int s = (int) ph.y; if (i >= a.nslots) { - mmvq_moe_mbar_wait(&empty[s], (uint32_t) ((i / a.nslots - 1) & 1)); + mmvq_moe_mbar_wait(&empty[s], (ph.x - 1) & 1); } int tile = -1; if (n_end == 0) { @@ -307,17 +436,7 @@ static __global__ void mmvq_moe(const mmvq_moe_dev_args a) { } continue; } - held[s] = tile; - const int u = tile / a.ntr; - const int row0 = (tile % a.ntr) * R; - const uint32_t bytes = (uint32_t) (min(R, a.nrows - row0) * a.row_bytes); - const int64_t off = expert[u]*a.stride_channel_x_bytes + (int64_t) row0*a.row_bytes; - char * slot = ring + (size_t) s*nmat*a.box_bytes; - mmvq_moe_mbar_arrive_expect_tx(&full[s], nmat*bytes); - mmvq_moe_bulk_load(slot, a.vx + off, bytes, &full[s], policy); - if constexpr (nmat == 2) { - mmvq_moe_bulk_load(slot + a.box_bytes, a.vgate + off, bytes, &full[s], policy); - } + issue(i, tile); if (i + 1 >= n_own && dyn_base < n_tiles) { ticket = atomicAdd(a.tile_ctr, 1); // the next sequence's tile, asked for now this one is issued } @@ -339,7 +458,11 @@ static __global__ void mmvq_moe(const mmvq_moe_dev_args a) { asm volatile("bar.sync 1, %0;" :: "n"(MMVQ_MOE_NW*32) : "memory"); } ggml_cuda_pdl_sync(); // the tokens are the previous kernels' results, and dst may still be read - if (a.y_bytes > 0) { + ggml_cuda_pdl_lc(); + if (a.y_f32 != nullptr) { + mmvq_moe_quantize_y(a, y_smem, a.nb*qk); + asm volatile("bar.sync 1, %0;" :: "n"(MMVQ_MOE_NW*32) : "memory"); // every consumer's blocks, before any is read + } else if (a.y_bytes > 0) { mmvq_moe_mbar_wait(&y_full, 0); // the producer's copy of the tokens' vectors } @@ -348,25 +471,27 @@ static __global__ void mmvq_moe(const mmvq_moe_dev_args a) { const int kb0 = lane / lanes_per_block; // the lane's first block of a row const int kqs = vdr * (lane % lanes_per_block); // and its quant ints in each - // IQ3_XXS: the lane's q8_1 fragments, a k iteration each (its block's 8 ints and scale), of the vector y_key - // (token * nchannels_y + channel), kept across pairs and tiles until the vector changes + // IQ3_XXS: the lane's q8_1 fragments, a k iteration each (its block's 8 ints and scale), of the vector y_key (its + // pair_y), kept across pairs and tiles until the vector changes (not with pairs_once) constexpr int nit = type == GGML_TYPE_IQ3_XXS ? mmvq_moe_iq3_nit(rpw) : 1; - int yu[nit][8]; - float yd[nit]; - int y_key = -1; + [[maybe_unused]] int yu[nit][8]; + [[maybe_unused]] float yd[nit]; + [[maybe_unused]] int y_key = -1; [[maybe_unused]] const auto grid = [&](const int i) { return grid_rep[i*32 + lane]; }; [[maybe_unused]] const auto ksigns = [&](const int i) { return ksigns_s[i]; }; for (int i = team;; i += MMVQ_MOE_NT) { - const int s = i % a.nslots; - mmvq_moe_mbar_wait(&full[s], (uint32_t) ((i / a.nslots) & 1)); + const uint2 ph = fast_div_modulo((uint32_t) i, a.nslots_fd); + const int s = (int) ph.y; + mmvq_moe_mbar_wait(&full[s], ph.x & 1); const int tile = held[s]; if (tile < 0) { break; } - const int u = tile / a.ntr; + const uint2 ur = fast_div_modulo((uint32_t) tile, a.ntr_fd); + const int u = (int) ur.x; const int e = expert[u]; - const int row0 = (tile % a.ntr) * R + g*rpw; // this warp's first row + const int row0 = (int) ur.y * R + g*rpw; // this warp's first row const char * box = ring + (size_t) s*nmat*a.box_bytes; // IQ3_XXS: the lane's fragments of its rows out of the slot, and the slot released before the math, so the @@ -394,16 +519,122 @@ static __global__ void mmvq_moe(const mmvq_moe_dev_args a) { } } - // every pair that routes to the expert, each with its own tokens (a down projection's differ by slot) - for (unsigned long long pairs = pairs_of[u]; pairs != 0; pairs &= pairs - 1) { - const int p = __ffsll(pairs) - 1; - const int t = p / a.n_used; - const int slot = p % a.n_used; + // a pair's rows: each lane's sums added over the warp, and lane 0 writes them, through the GLU with a gate + const auto finish = [&](const int p, float (&acc)[rpw][nmat]) { + float * const dst = a.dst + a.pair_dst[p]; +#pragma unroll + for (int r = 0; r < rpw; ++r) { +#pragma unroll + for (int m = 0; m < nmat; ++m) { + acc[r][m] = warp_reduce_sum<32>(acc[r][m]); + } + } + +#pragma unroll + for (int r = 0; r < rpw; ++r) { + const float * v = acc[r]; + const int row = row0 + r; + if (lane == 0 && row < a.nrows) { + float result = v[0]; + if (a.x_bias != nullptr) { + result += a.x_bias[e*a.stride_bias + row]; + } + if constexpr (nmat == 2) { + float gate_value = v[1]; + if (a.gate_bias != nullptr) { + gate_value += a.gate_bias[e*a.stride_bias + row]; + } + if (a.glu_limit > 0.0f) { + gate_value = fminf(gate_value, a.glu_limit); + result = fminf(fmaxf(result, -a.glu_limit), a.glu_limit); + } + switch (a.glu_op) { + case GGML_GLU_OP_SWIGLU: + result *= ggml_cuda_op_silu_single(gate_value); + break; + case GGML_GLU_OP_GEGLU: + result *= ggml_cuda_op_gelu_single(gate_value); + break; + case GGML_GLU_OP_SWIGLU_OAI: + result = ggml_cuda_op_swiglu_oai_single(gate_value, result); + break; + default: + result = result * gate_value; + break; + } + } + dst[row] = result; + } + } + }; + + // pairs_once (IQ3_XXS, a launch of several tokens): the expert's pairs MMVQ_MOE_PB at a time, each fragment + // decoded once, its grid gathered and its signs applied, and met with the pairs' vectors in turn, where a pair at + // a time decoded it again for each (GLM-5.3's proxy at 3 tokens, every token on each of its 8 experts, RTX 5070 + // Ti under ncu: the gate/up 113-116 us at 51-53 % of DRAM, its memory pipes 72-74 % busy). A vector's fragments + // are read for each tile, not kept across tiles. Each pair's sums are added in the same order, so the same bits + if constexpr (pairs_once) { + for (unsigned long long pairs = pairs_of[u]; pairs != 0;) { + int pc[MMVQ_MOE_PB]; // the pass's pairs, -1 past them + int yo[MMVQ_MOE_PB]; // and their vectors, in blocks past y_base +#pragma unroll + for (int q = 0; q < MMVQ_MOE_PB; ++q) { + pc[q] = pairs != 0 ? __ffsll(pairs) - 1 : -1; + yo[q] = pc[q] >= 0 ? (int) (mmvq_moe_y(a, y_base, pc[q]) - y_base) : 0; + pairs &= pairs - 1; + } + float acc[MMVQ_MOE_PB][rpw][nmat] = {}; +#pragma unroll + for (int it = 0; it < nit; ++it) { + const int kb = kb0 + it*blocks_per_iter; + if (kb < a.nb) { + iq3_xxs_frag_q8 wq[rpw][nmat]; +#pragma unroll + for (int r = 0; r < rpw; ++r) { +#pragma unroll + for (int m = 0; m < nmat; ++m) { + wq[r][m] = iq3_xxs_frag_decode(wf[r][m][it], grid, ksigns); + } + } +#pragma unroll + for (int q = 0; q < MMVQ_MOE_PB; ++q) { + if (pc[q] >= 0) { + const block_q8_1 * yb = y_base + yo[q] + kb*(qk/QK8_1) + kqs/2; + int u8[8]; +#pragma unroll + for (int l = 0; l < 8; ++l) { + u8[l] = get_int_b4(yb->qs, l); + } + const float d8 = __low2float(yb->ds); +#pragma unroll + for (int r = 0; r < rpw; ++r) { +#pragma unroll + for (int m = 0; m < nmat; ++m) { + acc[q][r][m] += vec_dot_iq3_xxs_frag_q8(wq[r][m], u8, d8); + } + } + } + } + } + } +#pragma unroll + for (int q = 0; q < MMVQ_MOE_PB; ++q) { + if (pc[q] >= 0) { + finish(pc[q], acc[q]); + } + } + } + } + + // every pair that routes to the expert, each with its own tokens (a down projection's differ by slot): with + // pairs_once, none past the passes above + for (unsigned long long pairs = pairs_once ? 0 : pairs_of[u]; pairs != 0; pairs &= pairs - 1) { + const int p = __ffsll(pairs) - 1; const block_q8_1 * y = mmvq_moe_y(a, y_base, p); float acc[rpw][nmat] = {{0.0f}}; if constexpr (type == GGML_TYPE_IQ3_XXS) { - const int key = t*a.nchannels_y + slot % a.nchannels_y; + const int key = a.pair_y[p]; // the vector, whichever pairs share it if (key != y_key) { y_key = key; #pragma unroll @@ -445,50 +676,7 @@ static __global__ void mmvq_moe(const mmvq_moe_dev_args a) { } } } -#pragma unroll - for (int r = 0; r < rpw; ++r) { -#pragma unroll - for (int m = 0; m < nmat; ++m) { - acc[r][m] = warp_reduce_sum<32>(acc[r][m]); - } - } - -#pragma unroll - for (int r = 0; r < rpw; ++r) { - const float * v = acc[r]; - const int row = row0 + r; - if (lane == 0 && row < a.nrows) { - float result = v[0]; - if (a.x_bias != nullptr) { - result += a.x_bias[e*a.stride_bias + row]; - } - if constexpr (nmat == 2) { - float gate_value = v[1]; - if (a.gate_bias != nullptr) { - gate_value += a.gate_bias[e*a.stride_bias + row]; - } - if (a.glu_limit > 0.0f) { - gate_value = fminf(gate_value, a.glu_limit); - result = fminf(fmaxf(result, -a.glu_limit), a.glu_limit); - } - switch (a.glu_op) { - case GGML_GLU_OP_SWIGLU: - result *= ggml_cuda_op_silu_single(gate_value); - break; - case GGML_GLU_OP_GEGLU: - result *= ggml_cuda_op_gelu_single(gate_value); - break; - case GGML_GLU_OP_SWIGLU_OAI: - result = ggml_cuda_op_swiglu_oai_single(gate_value, result); - break; - default: - result = result * gate_value; - break; - } - } - a.dst[t*a.stride_col_dst + slot*a.stride_channel_dst + row] = result; - } - } + finish(p, acc); } if constexpr (type != GGML_TYPE_IQ3_XXS) { @@ -527,7 +715,8 @@ static int64_t mmvq_moe_row_bytes(ggml_type type, int64_t ncols_x) { // The shared memory the ring and the tokens' vectors share: MMVQ_MOE_SMEM_MAX past the instance's static shared memory // as compiled (barriers and lists, and at IQ3_XXS its tables: 36,232 bytes on sm_120), read once a device. A 36 KiB -// budget for it left an MTP verify's gate/up without the room for its 3 tokens' vectors beside its ring. +// budget for it left an MTP verify's gate/up without the room for its 3 tokens' vectors beside its ring. A pairs_once +// instance declares the same shared memory as its shape's other. template static int mmvq_moe_dyn_max() { if constexpr (mmvq_moe_fits(type, nmat, rpw)) { @@ -535,7 +724,7 @@ static int mmvq_moe_dyn_max() { const int id = ggml_cuda_get_device(); if (dyn_max[id] == 0) { cudaFuncAttributes attr; - CUDA_CHECK(cudaFuncGetAttributes(&attr, mmvq_moe)); + CUDA_CHECK(cudaFuncGetAttributes(&attr, mmvq_moe)); dyn_max[id] = MMVQ_MOE_SMEM_MAX - (int) attr.sharedSizeBytes; } return dyn_max[id]; @@ -614,28 +803,41 @@ bool ggml_cuda_mmvq_moe_usable(int cc, ggml_type type, const void * vx, const vo (type != GGML_TYPE_IQ3_XXS || (ncols_x/QK_K + 3)/4 <= mmvq_moe_iq3_nit(p.rpw)); } -template +template static void mmvq_moe_launch(const mmvq_moe_dev_args & a, const int nblocks, cudaStream_t stream) { if constexpr (mmvq_moe_fits(type, nmat, rpw)) { const size_t smem = (size_t) a.nslots * nmat * a.box_bytes + a.y_bytes; - CUDA_SET_SHARED_MEMORY_LIMIT((mmvq_moe), (mmvq_moe_dyn_max())); // every plan's + CUDA_SET_SHARED_MEMORY_LIMIT((mmvq_moe), (mmvq_moe_dyn_max())); // every plan's const ggml_cuda_kernel_launch_params params(dim3(nblocks), dim3((MMVQ_MOE_NW + 1)*32), smem, stream); - ggml_cuda_kernel_launch(mmvq_moe, params, a); + ggml_cuda_kernel_launch(mmvq_moe, params, a); } else { GGML_ABORT("%s: no plan takes %d matrices at %d rows a warp for %s", __func__, nmat, rpw, ggml_type_name(type)); } } +// the shape's pairs_once instance where it has one (mmvq_moe_pairs_fits) and pairs_once, its other otherwise +template +static void mmvq_moe_launch_pairs(const mmvq_moe_dev_args & a, const bool pairs_once, const int nblocks, + cudaStream_t stream) { + if constexpr (mmvq_moe_pairs_fits(type, nmat, rpw)) { + if (pairs_once) { + mmvq_moe_launch(a, nblocks, stream); + return; + } + } + mmvq_moe_launch(a, nblocks, stream); +} + template -static void mmvq_moe_launch_type(const mmvq_moe_dev_args & a, const int nmat, const int rpw, const int nblocks, - cudaStream_t stream) { +static void mmvq_moe_launch_type(const mmvq_moe_dev_args & a, const int nmat, const int rpw, const bool pairs_once, + const int nblocks, cudaStream_t stream) { switch (nmat*8 + rpw) { - case 1*8 + 1: mmvq_moe_launch(a, nblocks, stream); break; - case 1*8 + 2: mmvq_moe_launch(a, nblocks, stream); break; - case 1*8 + 4: mmvq_moe_launch(a, nblocks, stream); break; - case 2*8 + 1: mmvq_moe_launch(a, nblocks, stream); break; - case 2*8 + 2: mmvq_moe_launch(a, nblocks, stream); break; - case 2*8 + 4: mmvq_moe_launch(a, nblocks, stream); break; + case 1*8 + 1: mmvq_moe_launch_pairs(a, pairs_once, nblocks, stream); break; + case 1*8 + 2: mmvq_moe_launch_pairs(a, pairs_once, nblocks, stream); break; + case 1*8 + 4: mmvq_moe_launch_pairs(a, pairs_once, nblocks, stream); break; + case 2*8 + 1: mmvq_moe_launch_pairs(a, pairs_once, nblocks, stream); break; + case 2*8 + 2: mmvq_moe_launch_pairs(a, pairs_once, nblocks, stream); break; + case 2*8 + 4: mmvq_moe_launch_pairs(a, pairs_once, nblocks, stream); break; default: GGML_ABORT("%s: no instance for %d matrices at %d rows a warp", __func__, nmat, rpw); } } @@ -656,16 +858,32 @@ bool ggml_cuda_mmvq_moe_keeps_y(ggml_type type, int64_t ncols_x, bool gate, int6 return !y_global && y_bytes > 0 && p.rpw > 0 && mmvq_moe_make_plan(type, row_bytes, nmat, y_bytes).rpw == p.rpw; } +bool ggml_cuda_mmvq_moe_quantizes_y(ggml_type type, int64_t ncols_x, bool gate, int64_t n_used, int64_t ntokens, + int64_t nchannels_y, const void * y_f32, int64_t y_f32_s1, int64_t y_f32_s2) { + static const bool legacy = ggml_env_switch("GGML_CUDA_MMVQ_MOE_QUANTIZE_LEGACY"); + if (legacy || nchannels_y != n_used || ncols_x % QK8_1 != 0 || (uintptr_t) y_f32 % 16 != 0 || y_f32_s1 % 4 != 0 || + y_f32_s2 % 4 != 0 || (n_used - 1)*y_f32_s1 + (ntokens - 1)*y_f32_s2 + ncols_x > INT_MAX) { + return false; + } + // the q8_1 layout the launch reads (mmvq_src1_q8_1's: each vector's row padded to MATRIX_ROW_PADDING) + const int64_t s11 = GGML_PAD(ncols_x, MATRIX_ROW_PADDING) / QK8_1; + return ggml_cuda_mmvq_moe_keeps_y(type, ncols_x, gate, ggml_cuda_mmvq_moe_y_bytes(ncols_x, n_used, ntokens, nchannels_y, + n_used*s11, s11)); +} + void ggml_cuda_mmvq_moe(const ggml_cuda_mmvq_moe_args & args, cudaStream_t stream) { const int nmat = args.vgate != nullptr ? 2 : 1; const int64_t row_bytes = mmvq_moe_row_bytes(args.type, args.ncols_x); - // routed experts read the tokens' vectors from the producer's copy past the ring when the plan keeps them there - int64_t y_bytes = (uintptr_t) args.y % 16 != 0 ? 0 : ggml_cuda_mmvq_moe_y_bytes(args.ncols_x, + // routed experts read the tokens' vectors from the producer's copy past the ring when the plan keeps them there, or + // from the copy the consumers quantize there (y_f32, which ggml_cuda_mmvq_moe_quantizes_y admitted) + int64_t y_bytes = (uintptr_t) args.y % 16 != 0 && args.y_f32 == nullptr ? 0 : ggml_cuda_mmvq_moe_y_bytes(args.ncols_x, args.n_used, args.ntokens, args.nchannels_y, args.stride_col_y, args.stride_channel_y); if (!ggml_cuda_mmvq_moe_keeps_y(args.type, args.ncols_x, nmat == 2, y_bytes)) { + GGML_ASSERT(args.y_f32 == nullptr && "ggml_cuda_mmvq_moe_quantizes_y keeps the quantized vectors in shared memory"); y_bytes = 0; } + GGML_ASSERT(args.y_f32 == nullptr || args.nchannels_y == args.n_used); const mmvq_moe_plan p = mmvq_moe_make_plan(args.type, row_bytes, nmat, y_bytes); GGML_ASSERT(p.rpw > 0 && "ggml_cuda_mmvq_moe_usable holds a plan"); @@ -681,33 +899,54 @@ void ggml_cuda_mmvq_moe(const ggml_cuda_mmvq_moe_args & args, cudaStream_t strea a.gate_bias = args.gate_bias; a.dst = args.dst; a.stride_channel_x_bytes = args.stride_channel_x_bytes; - a.stride_col_dst = args.stride_col_dst; - a.stride_channel_dst = args.stride_channel_dst; a.stride_bias = args.stride_bias; a.row_bytes = (int) row_bytes; a.box_bytes = p.box_bytes; a.nrows = (int) args.nrows_x; a.nb = (int) (args.ncols_x / ggml_blck_size(args.type)); a.ntr = ntr; + a.ntr_fd = init_fastdiv_values(ntr); + a.nslots_fd = init_fastdiv_values(p.nslots); a.n_used = (int) args.n_used; a.ntokens = (int) args.ntokens; a.ids_stride = (int) args.ids_stride; - a.nchannels_y = (int) args.nchannels_y; - a.stride_col_y = (int) args.stride_col_y; - a.stride_channel_y = (int) args.stride_channel_y; a.glu_op = (int) args.glu_op; a.glu_limit = args.glu_limit; a.nslots = p.nslots; a.y_bytes = (int) y_bytes; + a.y_f32 = args.y_f32; + a.y_f32_s1 = (int) args.y_f32_s1; + a.y_f32_s2 = (int) args.y_f32_s2; a.tile_ctr = args.tile_ctr; + a.ids_ready = args.ids_ready ? 1 : 0; + a.l2_issue_stop = args.l2_issue_stop; + for (int pair = 0; pair < args.ntokens*args.n_used; ++pair) { + const int64_t t = pair / args.n_used; + const int64_t slot = pair % args.n_used; + const int64_t yo = (slot % args.nchannels_y)*args.stride_channel_y + t*args.stride_col_y; + GGML_ASSERT(yo <= INT_MAX); + a.pair_y[pair] = (int) yo; + a.pair_dst[pair] = t*args.stride_col_dst + slot*args.stride_channel_dst; + } + + // a launch of several tokens decodes each fragment once for an expert's pairs where the shape has that instance + // (mmvq_moe_pairs_fits) and its pairs are at least as many as the experts. Decoding once pays for an expert that + // meets several pairs, and costs its instance's bookkeeping on one that meets one: where 3 tokens route to 8 of 288 + // experts (GLM-5.3-Flash) nearly every expert a launch reads meets one pair, and that instance ran the gate/up 1.2 % + // and the down 2.1 % slower than the one-pair one (a 4-layer proxy with its 288 experts, random routing, RTX 5080), + // and GLM-5.3-Flash at 3 and 4 tokens 2.4 and 1.4 % slower end to end (2 RTX PRO 6000, -sm tensor); where each of 3 + // tokens routes to all 8 of 8 (the 44-layer proxy) its gate/up is 27 % faster. GGML_CUDA_MMVQ_MOE_PAIRS_LEGACY=1: + // each pair decodes it, whatever the counts + static const bool pairs_legacy = ggml_env_switch("GGML_CUDA_MMVQ_MOE_PAIRS_LEGACY"); + const bool pairs_once = !pairs_legacy && args.ntokens > 1 && args.ntokens*args.n_used >= args.n_experts; // one block an SM, never more than the tiles of the most distinct experts the pairs can name const int nsm = ggml_cuda_info().devices[ggml_cuda_get_device()].nsm; const int nblocks = (int) std::min(nsm, args.ntokens*args.n_used*ntr); switch (args.type) { - case GGML_TYPE_IQ3_XXS: mmvq_moe_launch_type(a, nmat, p.rpw, nblocks, stream); break; - case GGML_TYPE_Q8_0: mmvq_moe_launch_type (a, nmat, p.rpw, nblocks, stream); break; + case GGML_TYPE_IQ3_XXS: mmvq_moe_launch_type(a, nmat, p.rpw, pairs_once, nblocks, stream); break; + case GGML_TYPE_Q8_0: mmvq_moe_launch_type (a, nmat, p.rpw, pairs_once, nblocks, stream); break; default: GGML_ABORT("%s: no instance for %s", __func__, ggml_type_name(args.type)); } } diff --git a/ggml/src/ggml-cuda/mmvq-moe.cuh b/ggml/src/ggml-cuda/mmvq-moe.cuh index a4033c5a74fd..70c6a22cc99f 100644 --- a/ggml/src/ggml-cuda/mmvq-moe.cuh +++ b/ggml/src/ggml-cuda/mmvq-moe.cuh @@ -16,6 +16,12 @@ struct ggml_cuda_mmvq_moe_args { const void * vx; const void * vgate; // nullptr: dst = vx's product (+ x_bias); else the GLU of the two products const void * y; + // non-null in place of y: the tokens' vectors in f32, slot s's of token t at y_f32 + s*y_f32_s1 + t*y_f32_s2, each + // slot its own (nchannels_y == n_used), which the launch quantizes into its shared copy exactly as + // quantize_row_q8_1_cuda writes y (ggml_cuda_mmvq_moe_quantizes_y) + const float * y_f32; + int64_t y_f32_s1; + int64_t y_f32_s2; const int32_t * ids; const float * x_bias; const float * gate_bias; @@ -26,6 +32,7 @@ struct ggml_cuda_mmvq_moe_args { int64_t nrows_x; int64_t stride_channel_x_bytes; int64_t n_used; // experts a token routes to: ids' columns + int64_t n_experts; // the experts ids name: src0's channels int64_t ntokens; int64_t ids_stride; int64_t nchannels_y; @@ -36,6 +43,10 @@ struct ggml_cuda_mmvq_moe_args { int64_t stride_bias; int * tile_ctr; // the stream's tile counter (ggml_cuda_pq2_tile_counters), or nullptr: every tile its // block's own + bool ids_ready; // ids whole before the launch starts (an earlier launch here on the stream read them, and + // triggers the next only past its dependency wait): the experts listed and the block's + // first tiles issued before this one's + unsigned int * l2_issue_stop; // bumped as the launch's reads start (ggml_cuda_l2_issue_stop), or nullptr }; // Whether ggml_cuda_mmvq_moe takes these weights: an NVIDIA GPU from Hopper on (cp.async.bulk), a type and token count @@ -56,4 +67,9 @@ int64_t ggml_cuda_mmvq_moe_y_bytes(int64_t ncols_x, int64_t n_used, int64_t ntok // consumers read them from global memory. bool ggml_cuda_mmvq_moe_keeps_y(ggml_type type, int64_t ncols_x, bool gate, int64_t y_bytes); +// Whether a routed launch quantizes its tokens' f32 vectors itself (args.y_f32): each slot's vector its own, 16-byte +// aligned rows, and a plan that keeps their q8_1 in shared memory. GGML_CUDA_MMVQ_MOE_QUANTIZE_LEGACY=1: never. +bool ggml_cuda_mmvq_moe_quantizes_y(ggml_type type, int64_t ncols_x, bool gate, int64_t n_used, int64_t ntokens, + int64_t nchannels_y, const void * y_f32, int64_t y_f32_s1, int64_t y_f32_s2); + void ggml_cuda_mmvq_moe(const ggml_cuda_mmvq_moe_args & args, cudaStream_t stream); diff --git a/ggml/src/ggml-cuda/mmvq-pq2-mma.cu b/ggml/src/ggml-cuda/mmvq-pq2-mma.cu index 08d9efd8fa94..508e29606174 100644 --- a/ggml/src/ggml-cuda/mmvq-pq2-mma.cu +++ b/ggml/src/ggml-cuda/mmvq-pq2-mma.cu @@ -46,6 +46,21 @@ #define PQ2_MMA_MAX_M 7 // a box row of at most 7 x 272 bytes #define PQ2_MMA_SMEM_MAX (99 * 1024) // the shared memory a block may take on sm_120 +// The group launch writing the inputs of a conv that folds the alpha/beta pair in (the qkv and z projections before it, +// fold_beside) at most 152 registers a thread. An SM's 65,536 registers are four sub-partitions of 16,384, and a block's +// warps take them in turn, so one sub-partition holds 3 of this block's 9 warps; a block of the next kernel lands beside +// it (PDL: the kernels after a matmul start under it and wait there) only if its warp there finds its registers. At 168 +// that sub-partition kept 256 and the conv after the qkv group (40 registers) started only as the group's blocks exited; +// at 152 it keeps 1,792, a warp of up to 56 (coresid/pdl-regs.cu in rig's roofline research, RTX 5080), and the conv +// computes the pair under the group. Every other launch keeps 168 (mmvq_pq2_mma_group_168; all of them with +// GGML_CUDA_PQ2_MMA_GROUP_REGS_LEGACY=1): the norm after an output projection (1,024 threads) finds no room at any count, +// gate/up ran 1.8 us slower at 152, and alone the cap streams the qkv group 0.7 us slower and buys nothing after it. +#if defined(GGML_USE_HIP) || defined(GGML_USE_MUSA) || CUDART_VERSION < 12040 +#define PQ2_MMA_GROUP_LAUNCH_BOUNDS __launch_bounds__((PQ2_MMA_NW + 1)*32, 1) +#else +#define PQ2_MMA_GROUP_LAUNCH_BOUNDS __maxnreg__(152) // it takes no __launch_bounds__ beside it; 152 x 288 threads fit an SM +#endif + static_assert(sizeof(block_pq2_0) == 34, "PQ2_0 block layout"); static_assert(PQ2_MMA_KB_BYTES % 16 == 0, "a box row is a whole number of 16-byte TMA units"); static_assert(PQ2_MMA_MAX_M * PQ2_MMA_KB_BYTES / 8 <= 256, "a TMA box is at most 256 8-byte elements wide"); @@ -185,7 +200,7 @@ static __device__ __forceinline__ void mmvq_pq2_mma_body( const CUtensorMap * tmap, const CUtensorMap * tmap_gate, const pq2_mma_group * grp, const block_q8_1 * y, const float * x_bias, float * dst, const int nrows, const int ncols, const int nb, const int n_tiles, const int nkb, const int nslots, const int evict_first, const int stride_col_y, const int stride_col_dst, - const ggml_cuda_pq2_prefetch & next, int * tile_ctr) { + const ggml_cuda_pq2_prefetch & next, int * tile_ctr, const bool lc_after_wait) { static_assert(nmat == 1 || (nmat == 2 && !has_bias), "a gated product fuses no bias"); static_assert(!grouped || (nmat == 1 && !has_bias), "a group's matrices fuse nothing"); #ifdef PQ2_MMA_AVAILABLE @@ -203,7 +218,13 @@ static __device__ __forceinline__ void mmvq_pq2_mma_body( const int warp = threadIdx.x / 32; const int lane = threadIdx.x % 32; - ggml_cuda_pdl_lc(); + // lc_after_wait: the block lets the next kernel launch only once its consumers have passed the dependency wait, so + // that kernel's launch means every kernel before this one has ended (each waited on the one before it): the conv + // that folds the alpha/beta pair in then never waits on its handoff slot beside the writer's blocks + // (ssm_conv_ab_acquire; the handoff, not this, makes the writer's stores visible to it) + if (!lc_after_wait) { + ggml_cuda_pdl_lc(); + } if (threadIdx.x == 0) { for (int s = 0; s < nslots; ++s) { @@ -300,6 +321,9 @@ static __device__ __forceinline__ void mmvq_pq2_mma_body( const int cc1 = min(2*t + 1, ncols - 1); ggml_cuda_pdl_sync(); // the tokens and the bias are the previous kernels' results, and dst may still be read + if (lc_after_wait) { + ggml_cuda_pdl_lc(); + } const block_q8_1 * y_g = y + (int64_t) cg * stride_col_y; const block_q8_1 * y_c0 = y + (int64_t) cc0 * stride_col_y; @@ -454,17 +478,19 @@ static __device__ __forceinline__ void mmvq_pq2_mma_body( } } - // every box of this block has landed: its share of the next launch's heads, which the kernels between the two then - // stream under (requested with the producer's last box instead, they took DRAM from this launch's tail: +4 us a - // gate + up on the 5070 Ti) + // every box of this block has landed: its share of the ranges (ggml_cuda_pq2_prefetch) laid end to end, which the + // kernels up to the next launch read and then stream under (requested with the producer's last box instead, the + // next launch's heads took DRAM from this launch's tail: +4 us a gate + up on the 5070 Ti) if (threadIdx.x == 0) { - const int64_t share = (next.bytes / gridDim.x + 15) & ~(int64_t) 15; - const int64_t off = (int64_t) blockIdx.x * share; - if (off < next.bytes) { - const uint32_t n = (uint32_t) min(share, next.bytes - off); - pq2_prefetch_l2((const char *) next.x + off, n); - if (next.gate != nullptr) { - pq2_prefetch_l2((const char *) next.gate + off, n); + const int64_t share = (next.total / gridDim.x + 15) & ~(int64_t) 15; + const int64_t lo = (int64_t) blockIdx.x * share; + const int64_t hi = lo + share; + int64_t base = 0; + for (int r = 0; r < next.n && base < hi; base += next.bytes[r++]) { + const int64_t b0 = max(lo, base); + const int64_t b1 = min(hi, base + next.bytes[r]); + if (b0 < b1) { + pq2_prefetch_l2((const char *) next.ptr[r] + (b0 - base), (uint32_t) (b1 - b0)); } } } @@ -483,17 +509,27 @@ static __global__ void mmvq_pq2_mma( const int nkb, const int nslots, const int evict_first, const int stride_col_y, const int stride_col_dst, const ggml_cuda_pq2_prefetch next, int * tile_ctr) { mmvq_pq2_mma_body(&tmap, &tmap_gate, nullptr, y, x_bias, dst, nrows, ncols, nb, n_tiles, - nkb, nslots, evict_first, stride_col_y, stride_col_dst, next, tile_ctr); + nkb, nslots, evict_first, stride_col_y, stride_col_dst, next, tile_ctr, false); } template -__launch_bounds__((PQ2_MMA_NW + 1)*32, 1) +PQ2_MMA_GROUP_LAUNCH_BOUNDS static __global__ void mmvq_pq2_mma_group( const __grid_constant__ pq2_mma_group grp, const block_q8_1 * y, const int ncols, const int nb, const int n_tiles, const int nkb, const int nslots, const int evict_first, const int stride_col_y, const ggml_cuda_pq2_prefetch next, - int * tile_ctr) { + int * tile_ctr, const bool lc_after_wait) { mmvq_pq2_mma_body(nullptr, nullptr, &grp, y, nullptr, nullptr, 0, ncols, nb, n_tiles, nkb, nslots, - evict_first, stride_col_y, 0, next, tile_ctr); + evict_first, stride_col_y, 0, next, tile_ctr, lc_after_wait); +} + +template +__launch_bounds__((PQ2_MMA_NW + 1)*32, 1) +static __global__ void mmvq_pq2_mma_group_168( + const __grid_constant__ pq2_mma_group grp, const block_q8_1 * y, const int ncols, const int nb, const int n_tiles, + const int nkb, const int nslots, const int evict_first, const int stride_col_y, const ggml_cuda_pq2_prefetch next, + int * tile_ctr, const bool lc_after_wait) { + mmvq_pq2_mma_body(nullptr, nullptr, &grp, y, nullptr, nullptr, 0, ncols, nb, n_tiles, nkb, nslots, + evict_first, stride_col_y, 0, next, tile_ctr, lc_after_wait); } // --------------------------------------------------------------------------------------------------------------------- @@ -660,20 +696,29 @@ void ggml_cuda_mmvq_pq2_mma(const void * vx, const void * vgate, const void * vy template static void pq2_mma_launch_group(const pq2_mma_plan & p, const pq2_mma_group & grp, const block_q8_1 * y, int ncols, - int nb, int n_tiles, int stride_col_y, const ggml_cuda_pq2_prefetch & next, int * tile_ctr, cudaStream_t stream) { + int nb, int n_tiles, int stride_col_y, const ggml_cuda_pq2_prefetch & next, int * tile_ctr, bool fold_beside, + cudaStream_t stream) { const int nsm = ggml_cuda_info().devices[ggml_cuda_get_device()].nsm; const int nblocks = std::min(nsm, n_tiles); const ggml_cuda_kernel_launch_params params(dim3(nblocks), dim3((PQ2_MMA_NW + 1)*32), pq2_mma_smem_bytes(M, 1, p.nslots), stream); + static const bool regs_legacy = ggml_env_switch("GGML_CUDA_PQ2_MMA_GROUP_REGS_LEGACY"); + if (regs_legacy || !fold_beside) { + CUDA_SET_SHARED_MEMORY_LIMIT((mmvq_pq2_mma_group_168), PQ2_MMA_SMEM_MAX); + ggml_cuda_kernel_launch(mmvq_pq2_mma_group_168, params, grp, y, ncols, nb, n_tiles, p.nkb, p.nslots, + (int) pq2_mma_get_config().evict_first, stride_col_y, next, tile_ctr, fold_beside); + return; + } CUDA_SET_SHARED_MEMORY_LIMIT((mmvq_pq2_mma_group), PQ2_MMA_SMEM_MAX); ggml_cuda_kernel_launch(mmvq_pq2_mma_group, params, grp, y, ncols, nb, n_tiles, p.nkb, p.nslots, - (int) pq2_mma_get_config().evict_first, stride_col_y, next, tile_ctr); + (int) pq2_mma_get_config().evict_first, stride_col_y, next, tile_ctr, fold_beside); } void ggml_cuda_mmvq_pq2_mma_group(int n, const void * const * vx, float * const * dst, const int64_t * nrows_x, const int64_t * stride_row_x, const int64_t * stride_col_dst, const void * vy, int64_t ncols_x, int64_t ncols_dst, int64_t stride_col_y, - const ggml_cuda_pq2_prefetch & next, int * tile_ctr, cudaStream_t stream) { + const ggml_cuda_pq2_prefetch & next, int * tile_ctr, bool fold_beside, + cudaStream_t stream) { GGML_ASSERT(n >= 2 && n <= PQ2_MMA_MAX_GROUP); const pq2_mma_plan p = pq2_mma_make_plan(ncols_x, 1); GGML_ASSERT(p.m > 0 && "ggml_cuda_mmvq_pq2_mma_usable holds a plan"); @@ -697,7 +742,7 @@ void ggml_cuda_mmvq_pq2_mma_group(int n, const void * const * vx, float * const const block_q8_1 * y = (const block_q8_1 *) vy; switch (p.m) { #define PQ2_MMA_CASE(M) case M: pq2_mma_launch_group(p, grp, y, (int) ncols_dst, nb, (int) n_tiles, (int) stride_col_y, \ - next, pq2_mma_tile_ctr(tile_ctr), stream); break; + next, pq2_mma_tile_ctr(tile_ctr), fold_beside, stream); break; PQ2_MMA_CASE(1) PQ2_MMA_CASE(2) PQ2_MMA_CASE(3) diff --git a/ggml/src/ggml-cuda/mmvq-pq2-mma.cuh b/ggml/src/ggml-cuda/mmvq-pq2-mma.cuh index b0a29cee283e..f7b89604a775 100644 --- a/ggml/src/ggml-cuda/mmvq-pq2-mma.cuh +++ b/ggml/src/ggml-cuda/mmvq-pq2-mma.cuh @@ -15,7 +15,8 @@ bool ggml_cuda_mmvq_pq2_mma_usable(int cc, const void * vx, const void * vgate, int64_t ncols_x, int64_t nrows_x, int64_t stride_row_x, int64_t ncols_dst); -// next: the head of the launch after this one, which this one prefetches into L2 (ggml_cuda_pq2_prefetch, common.cuh). +// next: what this launch prefetches into L2 once its own weights have landed: the weights of the kernels up to the next +// launch, then that launch's heads (ggml_cuda_pq2_prefetch, common.cuh). // tile_ctr: the stream's tile counter (ggml_cuda_pq2_tile_counters, common.cuh), or nullptr: every tile to its own block. void ggml_cuda_mmvq_pq2_mma(const void * vx, const void * vgate, const void * vy, const float * x_bias, float * dst, int64_t ncols_x, int64_t nrows_x, int64_t ncols_dst, int64_t stride_row_x, @@ -26,8 +27,12 @@ void ggml_cuda_mmvq_pq2_mma(const void * vx, const void * vgate, const void * vy // n (2 to PQ2_MMA_MAX_GROUP) matrices of ncols_x columns on one activation vy, in one launch: matrix g (vx[g], nrows_x[g] // rows of stride_row_x[g] blocks) writes dst[g], column c at dst[g] + c*stride_col_dst[g]. Each matrix must be one -// ggml_cuda_mmvq_pq2_mma_usable serves unfused, and each gets that launch's result. +// ggml_cuda_mmvq_pq2_mma_usable serves unfused, and each gets that launch's result. fold_beside: a kernel after it works +// beside it before its own dependency wait (the conv folding the alpha/beta pair in): the launch takes 152 registers, +// which leaves that kernel's blocks room on its SMs, and lets the next kernel launch only after its own dependency +// wait, not at its start. void ggml_cuda_mmvq_pq2_mma_group(int n, const void * const * vx, float * const * dst, const int64_t * nrows_x, const int64_t * stride_row_x, const int64_t * stride_col_dst, const void * vy, int64_t ncols_x, int64_t ncols_dst, int64_t stride_col_y, - const ggml_cuda_pq2_prefetch & next, int * tile_ctr, cudaStream_t stream); + const ggml_cuda_pq2_prefetch & next, int * tile_ctr, bool fold_beside, + cudaStream_t stream); diff --git a/ggml/src/ggml-cuda/mmvq.cu b/ggml/src/ggml-cuda/mmvq.cu index 5e5b9a563a31..b784cee1aefe 100644 --- a/ggml/src/ggml-cuda/mmvq.cu +++ b/ggml/src/ggml-cuda/mmvq.cu @@ -1324,6 +1324,17 @@ static char * mmvq_src1_q8_1(ggml_backend_cuda_context & ctx, const ggml_tensor return q8_1; } +// GGML_CUDA_MMVQ_MOE_QUANTIZE_CHECK=1: a routed ring's results from the vectors it quantized beside those from the +// quantize launch's +static __global__ void mmvq_moe_quantize_check(const float * own, const float * launch, const int64_t n) { + const int64_t i = (int64_t) blockIdx.x*blockDim.x + threadIdx.x; + if (i < n && __float_as_uint(own[i]) != __float_as_uint(launch[i])) { + printf("mmvq_moe quantize check: result %lld is %08x from the ring's own q8_1, %08x from the quantize launch's\n", + (long long) i, __float_as_uint(own[i]), __float_as_uint(launch[i])); + __trap(); + } +} + bool ggml_cuda_mul_mat_vec_q_moe_gate_fuses(const ggml_tensor * mm, const ggml_tensor * gate, const int cc) { const ggml_tensor * src0 = mm->src[0]; const ggml_tensor * src1 = mm->src[1]; @@ -1428,8 +1439,26 @@ void ggml_cuda_mul_mat_vec_q( } const int64_t ne10_padded = GGML_PAD(ne10, MATRIX_ROW_PADDING); + + // routed experts at 1-8 tokens, each SM streaming its tiles of expert rows through a ring of bulk copies, and each + // distinct expert read once (mmvq-moe.cu) + const bool moe_ring = ids && ne03 == 1 && ne13 == 1 && + (fusion == nullptr || (fusion->x_scale == nullptr && fusion->gate_scale == nullptr)) && + ggml_cuda_mmvq_moe_usable(ggml_cuda_info().devices[ctx.device].cc, src0->type, src0->data, fusion_local.gate, ne00, + ne01, nb01 / ts_src0, nb02 / ts_src0, ne1, ne2); + // a ring whose tokens' vectors are each its own (a down projection: ne11 == ne1) and read by no other MUL_MAT + // quantizes them itself from f32, into the shared memory it would copy them to, bit for bit as the quantize launch: + // no launch between the gate/up and the down (GLM-5.3's proxy under -sm tensor at 32K cached tokens, RTX 5070 Ti: + // 3.1 us of the chain after each gate/up, 43 a token). GGML_CUDA_MMVQ_MOE_QUANTIZE_LEGACY=1: the launch; + // GGML_CUDA_MMVQ_MOE_QUANTIZE_CHECK=1: the ring runs from both, and a kernel traps on the first result whose bits differ + const bool ring_quantizes = moe_ring && ctx.mmvq_shared_q8_1.find(src1) == nullptr && + ggml_cuda_mmvq_moe_quantizes_y(src0->type, ne00, fusion_local.gate != nullptr, ne1, ne2, ne11, src1->data, + nb11 / ts_src1, nb12 / ts_src1); + static const bool ring_quantize_check = ggml_env_switch("GGML_CUDA_MMVQ_MOE_QUANTIZE_CHECK"); + ggml_cuda_pool_alloc src1_q8_1_own(ctx.pool()); - char * src1_q8_1 = mmvq_src1_q8_1(ctx, src1, src0->type, src1_q8_1_own); + char * src1_q8_1 = ring_quantizes && !ring_quantize_check ? nullptr : + mmvq_src1_q8_1(ctx, src1, src0->type, src1_q8_1_own); const int64_t s01 = src0->nb[1] / ts_src0; const int64_t s11 = ne10_padded / QK8_1; @@ -1453,11 +1482,7 @@ void ggml_cuda_mul_mat_vec_q( const int64_t ids_stride = ids ? ids->nb[1] / ggml_type_size(ids->type) : 0; - // routed experts at 1-8 tokens, each SM streaming its tiles of expert rows through a ring of bulk copies, and each - // distinct expert read once (mmvq-moe.cu) - if (ids && ne03 == 1 && ne13 == 1 && (fusion == nullptr || (fusion->x_scale == nullptr && fusion->gate_scale == nullptr)) && - ggml_cuda_mmvq_moe_usable(ggml_cuda_info().devices[ctx.device].cc, src0->type, src0->data, - fusion_local.gate, ne00, ne01, s01, s02, ne1, ne2)) { + if (moe_ring) { ggml_cuda_mmvq_moe_args args{}; args.type = src0->type; args.vx = src0->data; @@ -1473,6 +1498,7 @@ void ggml_cuda_mul_mat_vec_q( args.nrows_x = ne01; args.stride_channel_x_bytes = s02 * (int64_t) ts_src0; args.n_used = ne1; + args.n_experts = ne02; args.ntokens = ne2; args.ids_stride = ids_stride; args.nchannels_y = nchannels_y; @@ -1482,7 +1508,34 @@ void ggml_cuda_mul_mat_vec_q( args.stride_channel_dst = stride_channel_dst; args.stride_bias = stride_channel_dst; // as mul_mat_vec_q reads an expert's bias args.tile_ctr = ctx.pq2_tile_counter(); + // the ids an earlier ring launch on this stream read in this evaluation (a down projection's, the gate/up's + // before it) are whole before this launch starts, so it lists its experts and issues its first tiles before its + // dependency wait (mmvq-moe.cu). GGML_CUDA_MMVQ_MOE_IDS_EARLY_LEGACY=1: past it + static const bool ids_early_legacy = ggml_env_switch("GGML_CUDA_MMVQ_MOE_IDS_EARLY_LEGACY"); + const ggml_tensor *& ids_before = ctx.mmvq_moe_ids[ctx.curr_stream_no]; + args.ids_ready = !ids_early_legacy && ids_before == ids; + ids_before = ids; + args.l2_issue_stop = ctx.l2_issue_stop.ptr; + if (!ring_quantizes) { + ggml_cuda_mmvq_moe(args, stream); + return; + } + ggml_cuda_pool_alloc check_dst(ctx.pool()); + if (ring_quantize_check) { + GGML_ASSERT(ggml_is_contiguous(dst)); + args.dst = check_dst.alloc(ggml_nelements(dst)); + ggml_cuda_mmvq_moe(args, stream); // from the quantize launch's vectors + args.dst = dst_d; + } + args.y = nullptr; + args.y_f32 = (const float *) src1->data; + args.y_f32_s1 = nb11 / ts_src1; + args.y_f32_s2 = nb12 / ts_src1; ggml_cuda_mmvq_moe(args, stream); + if (ring_quantize_check) { + const int64_t n = ggml_nelements(dst); + mmvq_moe_quantize_check<<<(n + 255) / 256, 256, 0, stream>>>(dst_d, check_dst.get(), n); + } return; } // past one token only the ring above fuses a routed gate (ggml_cuda_mul_mat_vec_q_moe_gate_fuses) @@ -1501,7 +1554,8 @@ void ggml_cuda_mul_mat_vec_q( ne03, ne3, s03, s13, s3, ids_stride, stream); } -void ggml_cuda_mul_mat_vec_q_pq2_group(ggml_backend_cuda_context & ctx, ggml_tensor * const * dsts, const int n) { +void ggml_cuda_mul_mat_vec_q_pq2_group(ggml_backend_cuda_context & ctx, ggml_tensor * const * dsts, const int n, + const bool feeds_fold) { GGML_ASSERT(n >= 2 && n <= PQ2_MMA_MAX_GROUP); const ggml_tensor * src1 = dsts[0]->src[1]; GGML_ASSERT(src1->type == GGML_TYPE_F32 && src1->nb[0] == sizeof(float) && src1->ne[2] == 1 && src1->ne[3] == 1); @@ -1538,7 +1592,7 @@ void ggml_cuda_mul_mat_vec_q_pq2_group(ggml_backend_cuda_context & ctx, ggml_ten } ggml_cuda_mmvq_pq2_mma_group(n, vx, dst, nrows, stride_row, stride_col_dst, src1_q8_1.get(), ne10, ne11, - ne10_padded / QK8_1, ctx.pq2_next, ctx.pq2_tile_counter(), stream); + ne10_padded / QK8_1, ctx.pq2_next, ctx.pq2_tile_counter(), feeds_fold, stream); } bool ggml_cuda_mul_mat_vec_q_concat_supported(const ggml_tensor * const * w, const int n, const ggml_tensor * src1, diff --git a/ggml/src/ggml-cuda/mmvq.cuh b/ggml/src/ggml-cuda/mmvq.cuh index 025e756bb6f2..5453f17f58e1 100644 --- a/ggml/src/ggml-cuda/mmvq.cuh +++ b/ggml/src/ggml-cuda/mmvq.cuh @@ -37,8 +37,9 @@ void ggml_cuda_mul_mat_vec_q(ggml_backend_cuda_context & ctx, bool ggml_cuda_mul_mat_vec_q_moe_gate_fuses(const ggml_tensor * mm, const ggml_tensor * gate, int cc); // n MUL_MATs (2 to PQ2_MMA_MAX_GROUP) on one src1, each of which ggml_cuda_mul_mat_vec_q would send unfused to the PQ2_0 -// tensor-core kernel: src1 quantized once, one launch over all their tiles (ggml_cuda_mmvq_pq2_mma_group). -void ggml_cuda_mul_mat_vec_q_pq2_group(ggml_backend_cuda_context & ctx, ggml_tensor * const * dsts, int n); +// tensor-core kernel: src1 quantized once, one launch over all their tiles (ggml_cuda_mmvq_pq2_mma_group). feeds_fold: +// the launch writes the inputs of a conv that folds the alpha/beta pair in (ggml_cuda_try_ssm_conv_ab). +void ggml_cuda_mul_mat_vec_q_pq2_group(ggml_backend_cuda_context & ctx, ggml_tensor * const * dsts, int n, bool feeds_fold); // n MUL_MATs of src1 by the weights w[0..n) whose outputs dst concatenates in that order along dim 0 (KDA's q, k and v): // whether one mul_mat_vec_q launch runs them, the weights as its channels, which needs them one stride from each other. diff --git a/ggml/src/ggml-cuda/norm.cu b/ggml/src/ggml-cuda/norm.cu index c92a56371174..0ff4b45c362a 100644 --- a/ggml/src/ggml-cuda/norm.cu +++ b/ggml/src/ggml-cuda/norm.cu @@ -1,4 +1,5 @@ #include "norm.cuh" +#include "convert.cuh" #include "unary.cuh" #include @@ -74,9 +75,9 @@ static __global__ void group_norm_f32(const float * x, float * dst, const int gr } } -template +template static __global__ void rms_norm_f32(const float * x, - float * dst, + dst_t * dst, const int ncols, const int64_t stride_row, const int64_t stride_channel, @@ -108,6 +109,7 @@ static __global__ void rms_norm_f32(const float * x, const int tid = threadIdx.x; static_assert(!do_add || do_multiply, "fusing add is not supported without multiplying"); + static_assert(std::is_same_v || !do_multiply, "a non-F32 dst is the plain norm's only"); x += sample*stride_sample + channel*stride_channel + row*stride_row; dst += ((sample*nchannels + channel)*nrows + row)*ncols; @@ -150,11 +152,142 @@ static __global__ void rms_norm_f32(const float * x, const int mul_col = fastmodulo(col, mul_ncols_packed); dst[col] = scale * x[col] * mul[mul_col]; } else { - dst[col] = scale * x[col]; + dst[col] = ggml_cuda_cast(scale * x[col]); } } } +// rms_norm_f32<256> on rows of at most 256 columns, one warp to a row and eight rows to a block. The block kernel gives +// each such row 256 threads, half of them idle on a 128-wide row, and two barriers for its 128 values: KDA's per-head +// output norm is 32,768 of those rows a 1024-token ubatch. Here a lane plays the block's threads lane, lane+32, ..., +// lane+224: each virtual warp's squares go through the same butterfly, and the eight partials through block_reduce's +// second stage, so the sum, the scale and every output are the block kernel's bit for bit. Summing a lane's squares +// first and butterflying once, or loading a float4 to a lane, reorders the sum: test-backend-ops would still pass it. +static constexpr int rms_norm_warp_rows_per_block = 8; + +template +static __global__ void rms_norm_f32_warp_rows(const float * x, + float * dst, + const int ncols, + const int nrows, + const int64_t stride_row, + const int64_t stride_channel, + const int64_t stride_sample, + const float eps, + const float * mul, + const int64_t mul_stride_row, + const int64_t mul_stride_channel, + const int64_t mul_stride_sample, + const uint3 mul_ncols_packed, + const uint3 mul_nrows_packed, + const uint3 mul_nchannels_packed, + const uint3 mul_nsamples_packed, + const float * add, + const int64_t add_stride_row, + const int64_t add_stride_channel, + const int64_t add_stride_sample, + const uint3 add_ncols_packed, + const uint3 add_nrows_packed, + const uint3 add_nchannels_packed, + const uint3 add_nsamples_packed) { + ggml_cuda_pdl_lc(); + const int nchannels = gridDim.y; + + const int lane = threadIdx.x % WARP_SIZE; + const int row = blockIdx.x*rms_norm_warp_rows_per_block + threadIdx.x/WARP_SIZE; + const int channel = blockIdx.y; + const int sample = blockIdx.z; + if (row >= nrows) { + return; + } + + static_assert(!do_add || do_multiply, "fusing add is not supported without multiplying"); + + x += sample*stride_sample + channel*stride_channel + row*stride_row; + dst += ((sample*nchannels + channel)*nrows + row)*ncols; + + if constexpr (do_multiply) { + const uint32_t mul_row = fastmodulo(row, mul_nrows_packed); + const uint32_t mul_channel = fastmodulo(channel, mul_nchannels_packed); + const uint32_t mul_sample = fastmodulo(sample, mul_nsamples_packed); + mul += mul_sample * mul_stride_sample + mul_channel * mul_stride_channel + mul_row * mul_stride_row; + } + + if constexpr (do_add) { + const int add_row = fastmodulo(row, add_nrows_packed); + const int add_channel = fastmodulo(channel, add_nchannels_packed); + const int add_sample = fastmodulo(sample, add_nsamples_packed); + add += add_sample * add_stride_sample + add_channel * add_stride_channel + add_row * add_stride_row; + } + + ggml_cuda_pdl_sync(); + float xv[256/WARP_SIZE]; + float partial = 0.0f; // block_reduce's second stage: lane w holds virtual warp w's sum, the rest its sentinel +#pragma unroll + for (int w = 0; w < 256/WARP_SIZE; ++w) { + if (w*WARP_SIZE >= ncols) { + break; // a virtual warp past the row sums to +0.0, block_reduce's sentinel + } + const int col = w*WARP_SIZE + lane; + xv[w] = col < ncols ? x[col] : 0.0f; + const float s = warp_reduce_sum(xv[w]*xv[w]); + if (lane == w) { + partial = s; + } + } + const float tmp = warp_reduce_sum(partial); + + const float mean = tmp / ncols; + const float scale = rsqrtf(mean + eps); + +#pragma unroll + for (int w = 0; w < 256/WARP_SIZE; ++w) { + const int col = w*WARP_SIZE + lane; + if (w*WARP_SIZE >= ncols) { + break; + } + if (col >= ncols) { + continue; + } + if constexpr (do_multiply && do_add) { + const int mul_col = fastmodulo(col, mul_ncols_packed); + const int add_col = fastmodulo(col, add_ncols_packed); + dst[col] = scale * xv[w] * mul[mul_col] + add[add_col]; + } else if constexpr (do_multiply) { + const int mul_col = fastmodulo(col, mul_ncols_packed); + dst[col] = scale * xv[w] * mul[mul_col]; + } else { + dst[col] = scale * xv[w]; + } + } +} + +// GGML_CUDA_RMS_NORM_WARP_ROWS_LEGACY=1 keeps rows of at most 256 columns on the 256-thread block kernel +static bool rms_norm_warp_rows_enabled(const int ncols) { + static const bool legacy = ggml_env_switch("GGML_CUDA_RMS_NORM_WARP_ROWS_LEGACY"); + return !legacy && ncols <= 256; +} + +template +static void rms_norm_f32_warp_rows_launch( + const float * x, const float * mul, const float * add, float * dst, const int ncols, const int nrows, + const int nchannels, const int nsamples, const int64_t stride_row, const int64_t stride_channel, + const int64_t stride_sample, const int64_t mul_stride_row, const int64_t mul_stride_channel, + const int64_t mul_stride_sample, const uint3 mul_ncols_packed, const uint3 mul_nrows_packed, + const uint3 mul_nchannels_packed, const uint3 mul_nsamples_packed, const int64_t add_stride_row, + const int64_t add_stride_channel, const int64_t add_stride_sample, const uint3 add_ncols_packed, + const uint3 add_nrows_packed, const uint3 add_nchannels_packed, const uint3 add_nsamples_packed, const float eps, + cudaStream_t stream) { + const dim3 blocks_num((nrows + rms_norm_warp_rows_per_block - 1) / rms_norm_warp_rows_per_block, nchannels, nsamples); + const dim3 block_dims(rms_norm_warp_rows_per_block*WARP_SIZE, 1, 1); + const ggml_cuda_kernel_launch_params launch_params = {blocks_num, block_dims, 0, stream}; + ggml_cuda_kernel_launch(rms_norm_f32_warp_rows, launch_params, + x, dst, ncols, nrows, stride_row, stride_channel, stride_sample, eps, mul, mul_stride_row, mul_stride_channel, + mul_stride_sample, mul_ncols_packed, mul_nrows_packed, mul_nchannels_packed, mul_nsamples_packed, add, + add_stride_row, add_stride_channel, add_stride_sample, add_ncols_packed, add_nrows_packed, add_nchannels_packed, + add_nsamples_packed); +} + // Qwen3.5's gated norm, silu(gate) * (rms_norm(x) * w) with a one-row weight w, in one kernel: rms_norm_f32's sum of // squares and scale at the same block size, then the weight multiply and the GLU's product in the unfused order, so the // values match the RMS_NORM -> MUL -> GLU(swiglu split) chain bit for bit. @@ -422,6 +555,10 @@ static void rms_norm_f32_cuda( // underlying cudaLaunchKernelEx does not support default params nullptr, 0, 0, 0, make_uint3(0, 0, 0), make_uint3(0, 0, 0), make_uint3(0, 0, 0), make_uint3(0, 0, 0), nullptr, 0, 0, 0, make_uint3(0, 0, 0), make_uint3(0, 0, 0), make_uint3(0, 0, 0), make_uint3(0, 0, 0)); + } else if (rms_norm_warp_rows_enabled(ncols)) { + const uint3 z = make_uint3(0, 0, 0); + rms_norm_f32_warp_rows_launch(x, nullptr, nullptr, dst, ncols, nrows, nchannels, nsamples, + stride_row, stride_channel, stride_sample, 0, 0, 0, z, z, z, z, 0, 0, 0, z, z, z, z, eps, stream); } else if (ncols < 1024) { const dim3 block_dims(256, 1, 1); const ggml_cuda_kernel_launch_params launch_params = {blocks_num, block_dims, block_dims.x > WARP_SIZE ? 32 * sizeof(float): 0, stream}; @@ -491,6 +628,12 @@ static void rms_norm_mul_f32_cuda(const float * x, mul_stride_sample, mul_ncols_packed, mul_nrows_packed, mul_nchannels_packed, mul_nsamples_packed, // underlying cudaLaunchKernelEx does not support default params nullptr, 0, 0, 0, make_uint3(0, 0, 0), make_uint3(0, 0, 0), make_uint3(0, 0, 0), make_uint3(0, 0, 0)); + } else if (rms_norm_warp_rows_enabled(ncols)) { + const uint3 z = make_uint3(0, 0, 0); + rms_norm_f32_warp_rows_launch(x, mul, nullptr, dst, ncols, nrows, nchannels, nsamples, + stride_row, stride_channel, stride_sample, mul_stride_row, mul_stride_channel, mul_stride_sample, + mul_ncols_packed, mul_nrows_packed, mul_nchannels_packed, mul_nsamples_packed, 0, 0, 0, z, z, z, z, + eps, stream); } else if (ncols < 1024) { const dim3 block_dims(256, 1, 1); const ggml_cuda_kernel_launch_params launch_params = ggml_cuda_kernel_launch_params{blocks_num, block_dims, block_dims.x > WARP_SIZE ? 32 * sizeof(float): 0, stream}; @@ -518,7 +661,13 @@ static void rms_norm_mul_f32_cuda(const float * x, const uint3 add_nrows_packed = init_fastdiv_values(add_nrows); const uint3 add_nchannels_packed = init_fastdiv_values(add_nchannels); const uint3 add_nsamples_packed = init_fastdiv_values(add_nsamples); - if (ncols < 1024) { + if (rms_norm_warp_rows_enabled(ncols)) { + rms_norm_f32_warp_rows_launch(x, mul, add, dst, ncols, nrows, nchannels, nsamples, stride_row, + stride_channel, stride_sample, mul_stride_row, mul_stride_channel, mul_stride_sample, mul_ncols_packed, + mul_nrows_packed, mul_nchannels_packed, mul_nsamples_packed, add_stride_row, add_stride_channel, + add_stride_sample, add_ncols_packed, add_nrows_packed, add_nchannels_packed, add_nsamples_packed, eps, + stream); + } else if (ncols < 1024) { const dim3 block_dims(256, 1, 1); const ggml_cuda_kernel_launch_params launch_params = ggml_cuda_kernel_launch_params{blocks_num, block_dims,block_dims.x > WARP_SIZE ? 32 * sizeof(float): 0, stream}; ggml_cuda_kernel_launch(rms_norm_f32<256, true, true>, launch_params, @@ -630,6 +779,37 @@ void ggml_cuda_op_rms_norm(ggml_backend_cuda_context & ctx, ggml_tensor * dst) { rms_norm_f32_cuda(src0_d, dst_d, ne00, ne01, ne02, ne03, s01, s02, s03, eps, stream); } +void ggml_cuda_op_rms_norm_bf16(ggml_backend_cuda_context & ctx, const ggml_tensor * dst, nv_bfloat16 * y) { + const ggml_tensor * src0 = dst->src[0]; + const float * src0_d = (const float *) src0->data; + cudaStream_t stream = ctx.stream(); + + GGML_ASSERT(src0->type == GGML_TYPE_F32); + GGML_ASSERT( dst->type == GGML_TYPE_F32); + + GGML_TENSOR_UNARY_OP_LOCALS; + + float eps; + memcpy(&eps, dst->op_params, sizeof(float)); + GGML_ASSERT(eps >= 0.0f); + + const size_t ts0 = ggml_type_size(src0->type); + GGML_ASSERT(nb00 == ts0); + const int64_t s01 = nb01 / ts0; + const int64_t s02 = nb02 / ts0; + const int64_t s03 = nb03 / ts0; + + // rms_norm_f32_cuda's block sizes: its 128-thread and warp-per-row kernels sum in the 256-thread block's order + const dim3 blocks_num(ne01, ne02, ne03); + const dim3 block_dims(ne00 < 1024 ? 256 : 1024, 1, 1); + const ggml_cuda_kernel_launch_params launch_params = {blocks_num, block_dims, 32 * sizeof(float), stream}; + const uint3 z = make_uint3(0, 0, 0); + ggml_cuda_kernel_launch(ne00 < 1024 ? rms_norm_f32<256, false, false, nv_bfloat16> : rms_norm_f32<1024, false, false, nv_bfloat16>, + launch_params, src0_d, y, (int) ne00, s01, s02, s03, eps, + // underlying cudaLaunchKernelEx does not support default params + nullptr, 0, 0, 0, z, z, z, z, nullptr, 0, 0, 0, z, z, z, z); +} + void ggml_cuda_op_add_rms_norm_fused( ggml_backend_cuda_context & ctx, ggml_tensor * add, ggml_tensor * rms_norm, ggml_tensor * mul) { const ggml_tensor * a = add->src[0]; diff --git a/ggml/src/ggml-cuda/norm.cuh b/ggml/src/ggml-cuda/norm.cuh index a4c455d31713..7f64ee769085 100644 --- a/ggml/src/ggml-cuda/norm.cuh +++ b/ggml/src/ggml-cuda/norm.cuh @@ -6,6 +6,10 @@ void ggml_cuda_op_group_norm(ggml_backend_cuda_context & ctx, ggml_tensor * dst) void ggml_cuda_op_rms_norm(ggml_backend_cuda_context & ctx, ggml_tensor * dst); +// The weightless RMS_NORM dst written to y as BF16 rather than to dst, rounded as convert_unary rounds F32 to BF16: the +// cast cuBLAS makes of an F32 src1 for a BF16 GEMM, of exactly the values ggml_cuda_op_rms_norm writes +void ggml_cuda_op_rms_norm_bf16(ggml_backend_cuda_context & ctx, const ggml_tensor * dst, nv_bfloat16 * y); + void ggml_cuda_op_add_rms_norm_fused( ggml_backend_cuda_context & ctx, ggml_tensor * add, ggml_tensor * rms_norm, ggml_tensor * mul); diff --git a/ggml/src/ggml-cuda/ssm-conv.cu b/ggml/src/ggml-cuda/ssm-conv.cu index 7bb4f24310f1..e406b3b5b2fa 100644 --- a/ggml/src/ggml-cuda/ssm-conv.cu +++ b/ggml/src/ggml-cuda/ssm-conv.cu @@ -123,6 +123,121 @@ static __global__ void ssm_conv_long_token_f32(const float * __restrict__ src0, } } +// Qwen3.5's alpha/beta pair folded into the conv-state update (u.ab_rows > 0, ggml_cuda_try_ssm_conv_ab): the block takes +// rows blockIdx.x, blockIdx.x + gridDim.x, ... of the two matrices laid end to end, at most +// GGML_CUDA_SSM_CONV_AB_MAX_ROWS, and computes each as mul_mat_vec_f does at its 256 threads, a thread here being two of +// its threads (tid and tid + 128): each one's column pairs strided by 256 in order, one warp reduction for each, then +// warp 0 over the eight partials and 24 zeros, so each value is the pair launch's bit for bit. The results wait in +// ab_res for the writes after the dependency wait. Everything here is read before that wait: the weights, which no +// kernel writes, and the activation, once its writer's release reaches the block (ssm_conv_ab_acquire), through L2 +// (ld.global.cg), where that kernel's writes are, not through this SM's L1. +// The fold's handoff (ggml_cuda_ssm_conv_ab_slots): thread 0 acquires the slot until every one of the activation's +// writer's blocks has released to it, the barrier orders the block's reads after that, and thread 0 takes a ticket; the +// block with the grid's last sets the slot back to 0, after every block's acquire. The wait costs one load: the group +// launch writing this conv's inputs, which comes after the writer, lets the kernels after it launch only after its own +// wait (ggml_cuda_ssm_conv_ab_enabled), when the writer has ended. A slot that never fills would be a planning fault: +// the block traps. +static __device__ __forceinline__ void ssm_conv_ab_acquire(const ggml_cuda_ssm_conv_state_update & u) { + if (threadIdx.x == 0) { + for (int i = 0; (ggml_cuda_ld_acquire(u.ab_slot) & 0xffffu) < (unsigned int) u.ab_writers; ++i) { + if (i == 1 << 22) { + __trap(); + } + } + if (atomicAdd(u.ab_slot, 1u << 16) >> 16 == gridDim.x - 1) { + atomicExch(u.ab_slot, 0u); + } + } + __syncthreads(); +} + +template +static __device__ __forceinline__ int ssm_conv_ab_rows(const ggml_cuda_ssm_conv_state_update & u, const int n_t, + float (*ab_res)[max_n_t]) { + constexpr int block = 256; // mul_mat_vec_f's block for the pair's rows (launch_mul_mat_vec_f_cuda: 5,120 columns) + constexpr int half = GGML_CUDA_SSM_CONV_UPDATE_THREADS; + static_assert(2*half == block, "a thread here is two of mul_mat_vec_f's"); + __shared__ float part[block/WARP_SIZE][max_n_t]; + + const int warp = threadIdx.x / WARP_SIZE; + const int lane = threadIdx.x % WARP_SIZE; + const int ncols2 = u.ab_ncols / 2; + const int rows = 2*u.ab_rows; + + // the block's rows into L2 first: their loads below then wait on L2, not DRAM + if (threadIdx.x == 0) { + for (int r = blockIdx.x; r < rows; r += gridDim.x) { + const nv_bfloat16 * x = u.ab_w + (r / u.ab_rows)*u.ab_s02 + (int64_t) (r % u.ab_rows)*u.ab_stride_row; + ggml_cuda_prefetch_l2(x, (int64_t) u.ab_ncols*sizeof(nv_bfloat16)); + } + } + ssm_conv_ab_acquire(u); + + int n = 0; + for (int r = blockIdx.x; r < rows; r += gridDim.x, ++n) { + const nv_bfloat162 * x2 = (const nv_bfloat162 *) (u.ab_w + (r / u.ab_rows)*u.ab_s02 + + (int64_t) (r % u.ab_rows)*u.ab_stride_row); + const float2 * y2 = (const float2 *) u.ab_y; + + float lo[max_n_t]; // mul_mat_vec_f's thread tid + float hi[max_n_t]; // and its thread tid + 128 +#pragma unroll + for (int j = 0; j < max_n_t; ++j) { + lo[j] = 0.0f; + hi[j] = 0.0f; + } +#pragma unroll 5 + for (int col2 = threadIdx.x; col2 < ncols2; col2 += block) { + const nv_bfloat162 tmpx = x2[col2]; +#pragma unroll + for (int j = 0; j < max_n_t; ++j) { + if (j < n_t) { + const float2 tmpy = __ldcg(y2 + j*(u.ab_stride_y/2) + col2); + ggml_cuda_mad(lo[j], tmpx.x, tmpy.x); + ggml_cuda_mad(lo[j], tmpx.y, tmpy.y); + } + } + if (col2 + half < ncols2) { + const nv_bfloat162 tmpx_hi = x2[col2 + half]; +#pragma unroll + for (int j = 0; j < max_n_t; ++j) { + if (j < n_t) { + const float2 tmpy = __ldcg(y2 + j*(u.ab_stride_y/2) + col2 + half); + ggml_cuda_mad(hi[j], tmpx_hi.x, tmpy.x); + ggml_cuda_mad(hi[j], tmpx_hi.y, tmpy.y); + } + } + } + } +#pragma unroll + for (int j = 0; j < max_n_t; ++j) { + if (j < n_t) { + lo[j] = warp_reduce_sum(lo[j]); + hi[j] = warp_reduce_sum(hi[j]); + if (lane == 0) { + part[warp][j] = lo[j]; + part[warp + half/WARP_SIZE][j] = hi[j]; + } + } + } + __syncthreads(); + if (warp == 0) { +#pragma unroll + for (int j = 0; j < max_n_t; ++j) { + if (j < n_t) { + float v = lane < block/WARP_SIZE ? part[lane][j] : 0.0f; + v = warp_reduce_sum(v); + if (lane == 0) { + ab_res[n][j] = v; + } + } + } + } + __syncthreads(); + } + return n; +} + // The weights of channel c: its row of the SSM_CONV's src1, or of the one of the three tensors src1 is a concatenation of // that holds it (ggml_cuda_ssm_conv_state_update::w_seg) static __device__ __forceinline__ const float * ssm_conv_w_row(const ggml_cuda_ssm_conv_state_update & u, const float * w, @@ -142,17 +257,58 @@ static __device__ __forceinline__ const float * ssm_conv_w_row(const ggml_cuda_s // head of GGML_CUDA_SSM_CONV_UPDATE_THREADS channels, so the blocks of the leading u.l2_heads heads also write the L2_NORM // of their outputs (u.l2_dst) with l2_norm_f32's arithmetic: one warp sums the squares lane by lane in its column order, // then the same warp reduction, rsqrtf and product, so the values match the L2_NORM's bit for bit. +// At most 56 registers a thread: a block of it lands beside a PQ2_0 group launch's block at 152 registers (a sub-partition +// keeps 1,792 of its 16,384, one warp of up to 56, mmvq-pq2-mma.cu), so it starts under the qkv matmul and its work +// before the dependency wait (the weights, the alpha/beta fold) runs there, off the chain +#if defined(GGML_USE_HIP) || defined(GGML_USE_MUSA) || CUDART_VERSION < 12040 +#define SSM_CONV_UPDATE_LAUNCH_BOUNDS __launch_bounds__(GGML_CUDA_SSM_CONV_UPDATE_THREADS) +#else +#define SSM_CONV_UPDATE_LAUNCH_BOUNDS __maxnreg__(56) +#endif + template +SSM_CONV_UPDATE_LAUNCH_BOUNDS static __global__ void ssm_conv_state_update_f32(const ggml_cuda_ssm_conv_state_update u, const float * w, - const int w_stride, float * dst, const int dst_stride, const int n_t) { + const int w_stride, float * dst, const int dst_stride, const int n_t, + const bool prewait, const bool state_prefetch) { ggml_cuda_pdl_lc(); const int c = blockIdx.x * blockDim.x + threadIdx.x; + // state_prefetch: the block's columns of the cache row into L2 before the wait, so their DRAM latency is paid under + // the matmul before this kernel (a hint: no result depends on it, ggml_cuda_prefetch_l2) + if (state_prefetch && threadIdx.x == 0) { + ggml_cuda_prefetch_l2(u.cache + (int64_t) u.ids[0] * u.row_stride + (int64_t) blockIdx.x * blockDim.x * (d_conv - 1), + (int64_t) blockDim.x * (d_conv - 1) * sizeof(float)); + } + + __shared__ float ab_res[GGML_CUDA_SSM_CONV_AB_MAX_ROWS][max_n_t]; + int n_ab = 0; + if (u.ab_rows > 0) { + n_ab = ssm_conv_ab_rows(u, n_t, ab_res); + } else if (u.ab_slot != nullptr) { + ssm_conv_ab_acquire(u); // released to, but not folded: the slot still goes back to 0 + } + float x[d_conv - 1 + max_n_t]; float wc[d_conv]; const float * w_row = ssm_conv_w_row(u, w, w_stride, c); + // prewait: w is a model weight, which no kernel writes, so it is requested before the dependency wait and arrives + // while the qkv matmul before this kernel streams (its blocks leave this one room beside them: PQ2_MMA_LAUNCH_BOUNDS) + if (prewait) { +#pragma unroll + for (int j = 0; j < d_conv; ++j) { + wc[j] = w_row[j]; + } + } ggml_cuda_pdl_sync(); + // the pair's results: their memory is free now that every kernel before this one has completed + if ((int) threadIdx.x < n_t) { + for (int k = 0; k < n_ab; ++k) { + const int r = blockIdx.x + k*gridDim.x; + u.ab_dst[(r / u.ab_rows)*u.ab_s2 + threadIdx.x*u.ab_stride_dst + r % u.ab_rows] = ab_res[k][threadIdx.x]; + } + } const float * state = u.cache + (int64_t) u.ids[0] * u.row_stride + (int64_t) c * (d_conv - 1); #pragma unroll for (int j = 0; j < d_conv - 1; ++j) { @@ -164,9 +320,11 @@ static __global__ void ssm_conv_state_update_f32(const ggml_cuda_ssm_conv_state_ x[d_conv - 1 + t] = u.x[t * u.x_stride + c]; } } + if (!prewait) { #pragma unroll - for (int j = 0; j < d_conv; ++j) { - wc[j] = w_row[j]; + for (int j = 0; j < d_conv; ++j) { + wc[j] = w_row[j]; + } } // the same sum as ssm_conv_f32, including its zero bias (it turns a -0.0f sum into +0.0f) @@ -310,18 +468,54 @@ void ggml_cuda_op_ssm_conv(ggml_backend_cuda_context & ctx, ggml_tensor * dst, g const ggml_cuda_kernel_launch_params launch_params(dim3(nr / threads), dim3(threads), 0, stream); const int w_stride = src1->nb[1] / sizeof(float); const int dst_stride = out->nb[1] / sizeof(float); + // GGML_CUDA_SSM_CONV_PREWAIT_LEGACY=1: the weights after the dependency wait, as everything else + static const bool prewait_legacy = ggml_env_switch("GGML_CUDA_SSM_CONV_PREWAIT_LEGACY"); + const bool prewait = !prewait_legacy && (u->w_seg_channels != 0 || (src1->buffer != nullptr && + ggml_backend_buffer_get_usage(src1->buffer) == GGML_BACKEND_BUFFER_USAGE_WEIGHTS)); + // GGML_CUDA_SSM_CONV_STATE_PREFETCH_LEGACY=1: no L2 prefetch of the cache row before the wait + static const bool state_prefetch = !ggml_env_switch("GGML_CUDA_SSM_CONV_STATE_PREFETCH_LEGACY"); + // the alpha/beta pair folds in only if its activation's writer released to the plan's slot and the PQ2_0 group + // launch writing this conv's inputs came after it, set up for the fold (ggml_cuda_try_ssm_conv_ab); then the + // pair's own nodes, which come later, are skipped. A slot released to is acquired either way, which sets it + // back to 0 + ggml_cuda_ssm_conv_state_update uab = *u; + const ggml_cuda_ssm_conv_ab_plan * ab = ctx.ssm_conv_updates().ab_plan_of(dst); + if (ab != nullptr && ab->writers > 0) { + uab.ab_slot = ctx.ssm_conv_ab_slots.ptr + ab->slot; + uab.ab_writers = ab->writers; + } + if (uab.ab_slot != nullptr && ab->fed) { + ctx.ssm_conv_updates().skipped.insert(ab->pair, ab->pair + 2); + } else { + uab.ab_rows = 0; + } if (fuse_silu) { // the L2_NORM normalizes the SILU's output: run it here only with the SILU fused, and only then skip it if (const ggml_tensor * l2 = ctx.ssm_conv_updates().l2_norm_of(dst)) { ctx.ssm_conv_updates().skipped.insert(l2); } - ggml_cuda_kernel_launch(ssm_conv_state_update_f32, - launch_params, *u, src1_d, w_stride, dst_d, dst_stride, (int) n_t); + // up to 4 tokens (decode, an MTP verify of up to 3 drafts) with half the registers of 8 + if (n_t <= 4) { + ggml_cuda_kernel_launch(ssm_conv_state_update_f32, + launch_params, uab, src1_d, w_stride, dst_d, dst_stride, (int) n_t, prewait, + state_prefetch); + } else { + ggml_cuda_kernel_launch(ssm_conv_state_update_f32, + launch_params, uab, src1_d, w_stride, dst_d, dst_stride, (int) n_t, prewait, + state_prefetch); + } } else { - ggml_cuda_ssm_conv_state_update raw = *u; + ggml_cuda_ssm_conv_state_update raw = uab; raw.l2_heads = 0; - ggml_cuda_kernel_launch(ssm_conv_state_update_f32, - launch_params, raw, src1_d, w_stride, dst_d, dst_stride, (int) n_t); + if (n_t <= 4) { + ggml_cuda_kernel_launch(ssm_conv_state_update_f32, + launch_params, raw, src1_d, w_stride, dst_d, dst_stride, (int) n_t, prewait, + state_prefetch); + } else { + ggml_cuda_kernel_launch(ssm_conv_state_update_f32, + launch_params, raw, src1_d, w_stride, dst_d, dst_stride, (int) n_t, prewait, + state_prefetch); + } } return; } diff --git a/ggml/src/ggml-cuda/topk-moe.cu b/ggml/src/ggml-cuda/topk-moe.cu index c8cec70bb320..ec9f1c32ec86 100644 --- a/ggml/src/ggml-cuda/topk-moe.cu +++ b/ggml/src/ggml-cuda/topk-moe.cu @@ -87,8 +87,11 @@ __device__ void sqrt_softplus_warp_inplace(float (&vals)[experts_per_thread], co It is intended as fusion of softmax->top-k->get_rows pipeline for MoE models */ +// the rows a block of topk_moe_cuda takes, a warp each +static constexpr int topk_moe_rows_per_block = 4; + template -__launch_bounds__(4 * WARP_SIZE, 1) __global__ void topk_moe_cuda(const float * logits, +__launch_bounds__(topk_moe_rows_per_block * WARP_SIZE, 1) __global__ void topk_moe_cuda(const float * logits, float * weights, int32_t * ids, float * bias, @@ -97,10 +100,8 @@ __launch_bounds__(4 * WARP_SIZE, 1) __global__ void topk_moe_cuda(const float * const float clamp_val, const float scale_val, const topk_moe_config config) { - const int row = blockIdx.x * blockDim.y + threadIdx.y; - if (row >= n_rows) { - return; - } + const int row = blockIdx.x * blockDim.y + threadIdx.y; + const bool has_row = row < n_rows; // a warp past the rows only meets the block's barrier past the reads logits += n_experts * row; weights += n_expert_used * row; @@ -117,10 +118,18 @@ __launch_bounds__(4 * WARP_SIZE, 1) __global__ void topk_moe_cuda(const float * } ggml_cuda_pdl_sync(); + if (has_row) { #pragma unroll - for (int i = 0; i < n_experts; i += WARP_SIZE) { - const int expert = i + threadIdx.x; - wt[i / WARP_SIZE] = (n_experts % WARP_SIZE == 0 || expert < n_experts) ? logits[expert] : -INFINITY; + for (int i = 0; i < n_experts; i += WARP_SIZE) { + const int expert = i + threadIdx.x; + wt[i / WARP_SIZE] = (n_experts % WARP_SIZE == 0 || expert < n_experts) ? logits[expert] : -INFINITY; + } + } + // the block's rows' logits are all read before a warp writes: its weights and ids may lie over another row's logits + // (ggml_cuda_topk_moe_reads_before_writes) + __syncthreads(); + if (!has_row) { + return; } if (!config.delayed_softmax) { @@ -282,9 +291,8 @@ static void launch_topk_moe_cuda(ggml_backend_cuda_context & ctx, const topk_moe_config config) { GGML_ASSERT(!(config.with_norm && config.delayed_softmax) && "delayed softmax is not supported with weight normalization"); - const int rows_per_block = 4; - dim3 grid_dims((n_rows + rows_per_block - 1) / rows_per_block, 1, 1); - dim3 block_dims(WARP_SIZE, rows_per_block, 1); + dim3 grid_dims((n_rows + topk_moe_rows_per_block - 1) / topk_moe_rows_per_block, 1, 1); + dim3 block_dims(WARP_SIZE, topk_moe_rows_per_block, 1); cudaStream_t stream = ctx.stream(); const ggml_cuda_kernel_launch_params launch_params = ggml_cuda_kernel_launch_params(grid_dims, block_dims, 0, stream); @@ -391,6 +399,10 @@ void ggml_cuda_op_topk_moe(ggml_backend_cuda_context & ctx, } } +bool ggml_cuda_topk_moe_reads_before_writes(const int64_t n_rows) { + return n_rows <= topk_moe_rows_per_block; +} + bool ggml_cuda_should_use_topk_moe(const ggml_tensor * gating_op, const ggml_tensor * weights, const ggml_tensor * logits, diff --git a/ggml/src/ggml-cuda/topk-moe.cuh b/ggml/src/ggml-cuda/topk-moe.cuh index 091ef02a415a..28af518b5044 100644 --- a/ggml/src/ggml-cuda/topk-moe.cuh +++ b/ggml/src/ggml-cuda/topk-moe.cuh @@ -22,6 +22,11 @@ void ggml_cuda_op_topk_moe(ggml_backend_cuda_context & ctx, const ggml_tensor * bias, const ggml_cuda_topk_moe_args & args); +// Whether ggml_cuda_op_topk_moe over n_rows reads every row's logits before it writes a weight or an id: its rows fit one +// block, whose warps meet at a barrier past their reads, so its weights and ids may lie over its logits (a decode's, an +// MTP verify's). +bool ggml_cuda_topk_moe_reads_before_writes(int64_t n_rows); + bool ggml_cuda_should_use_topk_moe(const ggml_tensor * gating_op, const ggml_tensor * weights, const ggml_tensor * logits, diff --git a/ggml/src/ggml-cuda/vecdotq.cuh b/ggml/src/ggml-cuda/vecdotq.cuh index 710f290ea5d8..c4dfda2209be 100644 --- a/ggml/src/ggml-cuda/vecdotq.cuh +++ b/ggml/src/ggml-cuda/vecdotq.cuh @@ -1393,43 +1393,84 @@ static __device__ __forceinline__ iq3_xxs_frag iq3_xxs_frag_load(const block_iq3 __half2float(bq3->d) }; } -// A fragment against the 8 ints u of the q8_1 block it meets and its scale d8. grid(i) and ksigns(i) read iq3xxs_grid -// and ksigns64 wherever the caller keeps them (mmvq-moe.cu: the grid a copy per lane in shared memory, so a warp's 32 +// A fragment's 32 weights as signed bytes, 4 an int in the order of the q8_1 ints they meet, and its scale: what +// vec_dot_iq3_xxs_frag makes of it before the dot products, for a caller that meets one fragment with several q8_1 +// blocks (mmvq-moe.cu, the tokens that route to one expert) +struct iq3_xxs_frag_q8 { + int q[8]; + int ls; // the 4-bit scale + float d; +}; + +// A fragment's weights 4*l0 to 4*l0 + 7 (l0 even) as signed bytes, 4 an int. grid(i) and ksigns(i) read iq3xxs_grid and +// ksigns64 wherever the caller keeps them (mmvq-moe.cu: the grid a copy per lane in shared memory, so a warp's 32 // gathers meet no bank conflict). template -static __device__ __forceinline__ float vec_dot_iq3_xxs_frag( - const iq3_xxs_frag & w, const int * u, const float d8, grid_t grid, ksigns_t ksigns) { - +static __device__ __forceinline__ int2 iq3_xxs_frag_pair(const iq3_xxs_frag & w, const int l0, grid_t grid, ksigns_t ksigns) { const uint8_t * q3 = (const uint8_t *) &w.q3; const uint32_t aux32 = w.aux32; - int sumi = 0; -#pragma unroll - for (int l0 = 0; l0 < 8; l0 += 2) { - const int2 grid_pos = make_int2(grid(q3[l0 + 0]), grid(q3[l0 + 1])); - // the 8 weights' signs as byte masks, 0xFF where negative (the 8th: the parity of the 7 stored), and a negative - // weight as ~g + 1, with no carry into the next byte as every grid byte is 4-62: 3 integer ops for 4 weights, - // where __vcmpne4 and __vsub4 are emulated in several each (built with -DGGML_CUDA_IQ3_XXS_SIGNS_LEGACY) + const int2 grid_pos = make_int2(grid(q3[l0 + 0]), grid(q3[l0 + 1])); + // the 8 weights' signs as byte masks, 0xFF where negative (the 8th: the parity of the 7 stored), and a negative + // weight as ~g + 1, with no carry into the next byte as every grid byte is 4-62: 3 integer ops for 4 weights, + // where __vcmpne4 and __vsub4 are emulated in several each (built with -DGGML_CUDA_IQ3_XXS_SIGNS_LEGACY) #ifdef GGML_CUDA_IQ3_XXS_SIGNS_LEGACY - GGML_UNUSED(ksigns); - const uint32_t signs = unpack_ksigns(aux32 >> (7*l0/2)); - const int signs0 = __vcmpne4(signs & 0x08040201, 0); - const int grid_l = __vsub4(grid_pos.x ^ signs0, signs0); - const int signs1 = __vcmpne4(signs & 0x80402010, 0); - const int grid_h = __vsub4(grid_pos.y ^ signs1, signs1); + GGML_UNUSED(ksigns); + const uint32_t signs = unpack_ksigns(aux32 >> (7*l0/2)); + const int signs0 = __vcmpne4(signs & 0x08040201, 0); + const int grid_l = __vsub4(grid_pos.x ^ signs0, signs0); + const int signs1 = __vcmpne4(signs & 0x80402010, 0); + const int grid_h = __vsub4(grid_pos.y ^ signs1, signs1); #else - const uint64_t signs = ksigns((aux32 >> (7*l0/2)) & 0x7F); - const int signs0 = (int) (uint32_t) signs; - const int signs1 = (int) (uint32_t) (signs >> 32); - const int grid_l = (grid_pos.x ^ signs0) + (signs0 & 0x01010101); - const int grid_h = (grid_pos.y ^ signs1) + (signs1 & 0x01010101); + const uint64_t signs = ksigns((aux32 >> (7*l0/2)) & 0x7F); + const int signs0 = (int) (uint32_t) signs; + const int signs1 = (int) (uint32_t) (signs >> 32); + const int grid_l = (grid_pos.x ^ signs0) + (signs0 & 0x01010101); + const int grid_h = (grid_pos.y ^ signs1) + (signs1 & 0x01010101); #endif // GGML_CUDA_IQ3_XXS_SIGNS_LEGACY + return make_int2(grid_l, grid_h); +} + +// The whole fragment decoded (iq3_xxs_frag_pair's grid and ksigns) +template +static __device__ __forceinline__ iq3_xxs_frag_q8 iq3_xxs_frag_decode(const iq3_xxs_frag & w, grid_t grid, ksigns_t ksigns) { + iq3_xxs_frag_q8 wq; +#pragma unroll + for (int l0 = 0; l0 < 8; l0 += 2) { + const int2 q = iq3_xxs_frag_pair(w, l0, grid, ksigns); + wq.q[l0 + 0] = q.x; + wq.q[l0 + 1] = q.y; + } + wq.ls = w.aux32 >> 28; + wq.d = w.d; + return wq; +} - sumi = ggml_cuda_dp4a(grid_l, u[l0 + 0], sumi); - sumi = ggml_cuda_dp4a(grid_h, u[l0 + 1], sumi); +// A decoded fragment against the 8 ints u of the q8_1 block it meets and its scale d8 +static __device__ __forceinline__ float vec_dot_iq3_xxs_frag_q8(const iq3_xxs_frag_q8 & wq, const int * u, const float d8) { + int sumi = 0; +#pragma unroll + for (int l = 0; l < 8; ++l) { + sumi = ggml_cuda_dp4a(wq.q[l], u[l], sumi); } + sumi = (wq.ls*sumi + sumi/2)/2; + return wq.d * d8 * sumi; +} - const int ls = aux32 >> 28; +// A fragment against the 8 ints u of the q8_1 block it meets and its scale d8 (iq3_xxs_frag_pair's grid and ksigns), +// each pair of ints met as it is decoded: the whole fragment decoded first, the one-pair ring's down ran 61.4 -> 66.9 us +// a layer (GLM-5.3's proxy at 3 tokens, RTX 5070 Ti), its q8_1 reloads branched around where they were predicated +template +static __device__ __forceinline__ float vec_dot_iq3_xxs_frag( + const iq3_xxs_frag & w, const int * u, const float d8, grid_t grid, ksigns_t ksigns) { + int sumi = 0; +#pragma unroll + for (int l0 = 0; l0 < 8; l0 += 2) { + const int2 q = iq3_xxs_frag_pair(w, l0, grid, ksigns); + sumi = ggml_cuda_dp4a(q.x, u[l0 + 0], sumi); + sumi = ggml_cuda_dp4a(q.y, u[l0 + 1], sumi); + } + const int ls = w.aux32 >> 28; sumi = (ls*sumi + sumi/2)/2; return w.d * d8 * sumi; } diff --git a/ggml/src/ggml-metal/ggml-metal-device.m b/ggml/src/ggml-metal/ggml-metal-device.m index da9343414046..0fc89c44f0fc 100644 --- a/ggml/src/ggml-metal/ggml-metal-device.m +++ b/ggml/src/ggml-metal/ggml-metal-device.m @@ -1634,7 +1634,8 @@ bool ggml_metal_device_supports_op(ggml_metal_device_t dev, const struct ggml_te return has_simdgroup_mm; // TODO: over-restricted for vec-kernels case GGML_OP_LIGHTNING_INDEXER: if (op->src[0]->ne[0] != OP_LIGHTNING_INDEXER_DK || - op->src[0]->ne[1] != OP_LIGHTNING_INDEXER_NH) { + op->src[0]->ne[1] != OP_LIGHTNING_INDEXER_NH || + op->src[4] != NULL) { // ggml_lightning_indexer_rows: no Metal kernel reads k by row index return false; } if (!has_simdgroup_mm || diff --git a/ggml/src/ggml-sycl/ggml-sycl.cpp b/ggml/src/ggml-sycl/ggml-sycl.cpp index 23bce9a01794..a29c386b0298 100644 --- a/ggml/src/ggml-sycl/ggml-sycl.cpp +++ b/ggml/src/ggml-sycl/ggml-sycl.cpp @@ -6085,7 +6085,8 @@ static bool do_ggml_backend_sycl_device_supports_op(ggml_backend_dev_t dev, cons op->src[2]->type == GGML_TYPE_F32 && op->src[3]->type == GGML_TYPE_F32 && op->type == GGML_TYPE_F32; case GGML_OP_LIGHTNING_INDEXER: - return op->src[0]->type == GGML_TYPE_F32 && + // ggml_lightning_indexer_rows (src[4]): no SYCL kernel reads k by row index + return op->src[4] == nullptr && op->src[0]->type == GGML_TYPE_F32 && (op->src[1]->type == GGML_TYPE_F16 || op->src[1]->type == GGML_TYPE_F32 || op->src[1]->type == GGML_TYPE_BF16 || op->src[1]->type == GGML_TYPE_Q8_0 || op->src[1]->type == GGML_TYPE_Q5_1 || op->src[1]->type == GGML_TYPE_Q5_0 || diff --git a/ggml/src/ggml.c b/ggml/src/ggml.c index 2de1a2965d71..a5d42ecc2248 100644 --- a/ggml/src/ggml.c +++ b/ggml/src/ggml.c @@ -5549,6 +5549,22 @@ bool ggml_flash_attn_ext_get_mask_prefix( return ggml_get_op_params_i32(a, 4) != 0; } +void ggml_flash_attn_ext_set_n_kv_max( + struct ggml_tensor * a, + int32_t n_kv_max) { + GGML_ASSERT(a->op == GGML_OP_FLASH_ATTN_EXT); + GGML_ASSERT(n_kv_max >= 0); + + ggml_set_op_params_i32(a, 5, n_kv_max); // after the mask_prefix hint +} + +int32_t ggml_flash_attn_ext_get_n_kv_max( + const struct ggml_tensor * a) { + GGML_ASSERT(a->op == GGML_OP_FLASH_ATTN_EXT); + + return ggml_get_op_params_i32(a, 5); +} + // ggml_flash_attn_back struct ggml_tensor * ggml_flash_attn_back( @@ -6514,6 +6530,44 @@ struct ggml_tensor * ggml_lightning_indexer( return result; } +struct ggml_tensor * ggml_lightning_indexer_rows( + struct ggml_context * ctx, + struct ggml_tensor * q, + struct ggml_tensor * k, + struct ggml_tensor * k_rows, + struct ggml_tensor * weights, + struct ggml_tensor * mask) { + GGML_ASSERT( q->type == GGML_TYPE_F32); + GGML_ASSERT( weights->type == GGML_TYPE_F32); + GGML_ASSERT( mask->type == GGML_TYPE_F16); + GGML_ASSERT( k_rows->type == GGML_TYPE_I32); + GGML_ASSERT(ggml_is_contiguous(k_rows)); + GGML_ASSERT( q->ne[0] == k->ne[0]); + GGML_ASSERT( k->ne[1] == 1); + GGML_ASSERT( mask->ne[0] == k_rows->ne[0]); + GGML_ASSERT( k_rows->ne[1] == k->ne[3] && k_rows->ne[2] == 1 && k_rows->ne[3] == 1); + GGML_ASSERT( q->ne[1] == weights->ne[0]); + GGML_ASSERT( mask->ne[1] == q->ne[2]); + GGML_ASSERT( q->ne[2] == weights->ne[1]); + GGML_ASSERT(weights->ne[2] == 1); + GGML_ASSERT( mask->ne[2] == 1); + GGML_ASSERT( q->ne[3] == k->ne[3]); + GGML_ASSERT( k->ne[3] == weights->ne[3]); + GGML_ASSERT(weights->ne[3] % mask->ne[3] == 0); + + int64_t ne[4] = { k_rows->ne[0], q->ne[2], 1, q->ne[3] }; + struct ggml_tensor * result = ggml_new_tensor(ctx, GGML_TYPE_F32, 4, ne); + + result->op = GGML_OP_LIGHTNING_INDEXER; + result->src[0] = q; + result->src[1] = k; + result->src[2] = weights; + result->src[3] = mask; + result->src[4] = k_rows; + + return result; +} + // ggml_dsv4_hc_comb struct ggml_tensor * ggml_dsv4_hc_comb( diff --git a/src/llama-context.cpp b/src/llama-context.cpp index 4ffd92403fbf..6f7e11a5be52 100644 --- a/src/llama-context.cpp +++ b/src/llama-context.cpp @@ -28,6 +28,89 @@ // llama_context // +// LLAMA_TIME_DECODE=1 reports where a decode token's HOST time goes, measured with clock_gettime in-process. +// It exists because a profiler cannot answer this on this machine: nsys inflates CUPTI RUNTIME durations roughly 30-40x +// (a cudaGraphLaunch the engine times at 5.9-13.6 us reads 246-294 us in a trace), which put 66 % of a decode token's +// ~430 us device-idle gap on the graph launch when the launch is really about 3 % of it. Three fixes aimed at that +// phantom -- fusing nodes, launching the two devices' graphs from separate threads, and the scheduler's extra input +// copies -- each measured null, correctly. So the ~415 us is elsewhere in this function, and these phases say where. +// Device timestamps from a trace remain trustworthy; host attribution has to come from here. +// Off (the default) each timer is one test of a static flag. +enum llama_decode_phase { + LLAMA_DECODE_PHASE_BUILD_GRAPH, + LLAMA_DECODE_PHASE_ALLOC_GRAPH, + LLAMA_DECODE_PHASE_SET_INPUTS, + LLAMA_DECODE_PHASE_GRAPH_COMPUTE, + LLAMA_DECODE_PHASE_DECODE_TOTAL, + LLAMA_DECODE_PHASE_COUNT, +}; + +static const char * llama_decode_phase_name(llama_decode_phase p) { + switch (p) { + case LLAMA_DECODE_PHASE_BUILD_GRAPH: return "build_graph"; + case LLAMA_DECODE_PHASE_ALLOC_GRAPH: return "alloc_graph"; + case LLAMA_DECODE_PHASE_SET_INPUTS: return "set_inputs"; + case LLAMA_DECODE_PHASE_GRAPH_COMPUTE: return "graph_compute"; + case LLAMA_DECODE_PHASE_DECODE_TOTAL: return "DECODE TOTAL"; + default: return "?"; + } +} + +static bool llama_time_decode() { + static const bool enabled = [] { + const char * env = getenv("LLAMA_TIME_DECODE"); + return env != nullptr && atoi(env) == 1; + }(); + return enabled; +} + +// An RAII accumulator. Nested phases are fine: each one charges only its own bracket, and DECODE_TOTAL is the +// denominator the others are read against, so whatever they do not account for is named rather than assumed. +struct llama_decode_phase_timer { + llama_decode_phase phase; + struct timespec t0; + bool on; + + llama_decode_phase_timer(llama_decode_phase p) : phase(p), on(llama_time_decode()) { + if (on) { + clock_gettime(CLOCK_MONOTONIC, &t0); + } + } + + ~llama_decode_phase_timer() { + if (!on) { + return; + } + struct timespec t1; + clock_gettime(CLOCK_MONOTONIC, &t1); + const double us = (t1.tv_sec - t0.tv_sec)*1e6 + (t1.tv_nsec - t0.tv_nsec)/1e3; + static double sum[LLAMA_DECODE_PHASE_COUNT] = {}; + static int cnt[LLAMA_DECODE_PHASE_COUNT] = {}; + sum[phase] += us; + cnt[phase] += 1; + if (phase == LLAMA_DECODE_PHASE_DECODE_TOTAL && cnt[phase] % 64 == 0) { + // straight to stderr: llama-bench installs a log callback that drops INFO and WARN + const double tot = sum[LLAMA_DECODE_PHASE_DECODE_TOTAL] / cnt[LLAMA_DECODE_PHASE_DECODE_TOTAL]; + fprintf(stderr, "decode host phases over %d calls, mean us a call:\n", cnt[phase]); + double named = 0; + for (int i = 0; i < LLAMA_DECODE_PHASE_COUNT; i++) { + if (cnt[i] == 0) { + continue; + } + const double mean = sum[i] / cnt[i]; + if (i != LLAMA_DECODE_PHASE_DECODE_TOTAL) { + named += mean * cnt[i] / cnt[LLAMA_DECODE_PHASE_DECODE_TOTAL]; + } + fprintf(stderr, " %-14s %9.1f us %5.1f %% of decode (%d calls)\n", + llama_decode_phase_name((llama_decode_phase) i), mean, + tot > 0 ? 100.0*mean*cnt[i]/cnt[LLAMA_DECODE_PHASE_DECODE_TOTAL]/tot : 0.0, cnt[i]); + } + fprintf(stderr, " %-14s %9.1f us %5.1f %% of decode <-- in decode but in none of the phases above\n", + "unattributed", tot - named, tot > 0 ? 100.0*(tot - named)/tot : 0.0); + } + } +}; + // Verify that every Hadamard-folded weight consumed by the graph receives its // activation-side transform, and every latent lookup table gets the inverse. // An architecture whose matmul path bypasses the transform helpers would @@ -1634,9 +1717,10 @@ llm_graph_result * llama_context::process_ubatch(const llama_ubatch & ubatch, ll //const auto t_start_us = ggml_time_us(); - gf = model.build_graph(gparams); - - //LLAMA_LOG_INFO("graph build time: %.3f ms\n", (ggml_time_us() - t_start_us)/1000.0); + { + llama_decode_phase_timer t(LLAMA_DECODE_PHASE_BUILD_GRAPH); + gf = model.build_graph(gparams); + } if (!gf) { LLAMA_LOG_ERROR("%s: failed to initialize graph\n", __func__); @@ -1644,6 +1728,7 @@ llm_graph_result * llama_context::process_ubatch(const llama_ubatch & ubatch, ll return nullptr; } + llama_decode_phase_timer t_alloc(LLAMA_DECODE_PHASE_ALLOC_GRAPH); if (!ggml_backend_sched_alloc_graph(sched_gf, gf)) { LLAMA_LOG_ERROR("%s: failed to allocate graph\n", __func__); ret = GGML_STATUS_ALLOC_FAILED; @@ -1658,15 +1743,17 @@ llm_graph_result * llama_context::process_ubatch(const llama_ubatch & ubatch, ll // set the input data for the input tensors { - //const auto t_start_us = ggml_time_us(); + llama_decode_phase_timer t(LLAMA_DECODE_PHASE_SET_INPUTS); // FIXME this call causes a crash if any model inputs were not used in the graph and were therefore not allocated res->set_inputs(&ubatch); - - //LLAMA_LOG_INFO("graph set inputs time: %.3f ms\n", (ggml_time_us() - t_start_us)/1000.0); } - const auto status = graph_compute(res->get_gf(), ubatch.n_tokens > 1, sched_gf); + ggml_status status; + { + llama_decode_phase_timer t(LLAMA_DECODE_PHASE_GRAPH_COMPUTE); + status = graph_compute(res->get_gf(), ubatch.n_tokens > 1, sched_gf); + } if (status != GGML_STATUS_SUCCESS) { LLAMA_LOG_ERROR("%s: failed to compute graph, compute status: %d\n", __func__, status); ret = status; @@ -1917,6 +2004,8 @@ static bool needs_raw_logits(const llama_ubatch & ubatch, const std::map hparams.indexer_top_k + kpool - 1; + dims.n_dump = dims.scoring ? llama_kpool_select_k(dims.n_pools, hparams.indexer_top_k, kpool) : 0; return dims; } @@ -622,13 +623,17 @@ bool llm_graph_input_kpool::shapes_match(const llm_graph_input_kpool_dims & dims return res; } - res &= inp.pool_cells->ne[0] == (int64_t) inp.kpool*dims.n_pools; + res &= inp.pool_cells->ne[0] == (int64_t) inp.kpool*(dims.n_pools + dims.n_dump); res &= inp.pool_cells->ne[1] == dims.n_stream; res &= inp.pool_bias->ne[0] == dims.n_pools; res &= inp.pool_bias->ne[1] == dims.n_tps; res &= inp.pool_bias->ne[2] == dims.n_stream; + res &= inp.pool_dump->ne[0] == dims.n_dump; + res &= inp.pool_dump->ne[1] == dims.n_tps; + res &= inp.pool_dump->ne[2] == dims.n_stream; + if (inp.pool_bias_f16) { res &= inp.pool_bias_f16->ne[0] == dims.n_pools; res &= inp.pool_bias_f16->ne[1] == dims.n_tps; @@ -636,7 +641,7 @@ bool llm_graph_input_kpool::shapes_match(const llm_graph_input_kpool_dims & dims res &= inp.pool_bias_f16->ne[3] == dims.n_stream; } - res &= inp.sel_mask->ne[0] == dims.n_kv; + res &= inp.sel_mask->ne[0] == dims.n_kv + (int64_t) inp.kpool*dims.n_dump; res &= inp.sel_mask->ne[1] == dims.n_tps; res &= inp.sel_mask->ne[2] == 1; res &= inp.sel_mask->ne[3] == dims.n_stream; @@ -2883,7 +2888,8 @@ ggml_tensor * llm_graph_context::build_attn_mha( ggml_tensor * v_mla, float kq_scale, int il, - bool mask_is_prefix) const { + bool mask_is_prefix, + int64_t n_kv_max) const { const bool v_trans = v->nb[1] > v->nb[2]; // split the batch into streams if needed @@ -2928,6 +2934,10 @@ ggml_tensor * llm_graph_context::build_attn_mha( ggml_flash_attn_ext_set_mask_prefix(cur, (mask_is_prefix || sparse_prefix_legacy) && cparams.n_seq_max == 1 && cparams.causal_attn && il >= 0 && !hparams.is_swa(il)); + // n_kv_max: a bound on each mask row's live cells, which a backend may then read alone (the sparse mode) + GGML_ASSERT(n_kv_max >= 0 && n_kv_max <= INT32_MAX); + ggml_flash_attn_ext_set_n_kv_max(cur, static_cast(n_kv_max)); + if (v_mla) { #if 0 // v_mla can be applied as a matrix-vector multiplication with broadcasting across dimension 3 == n_tokens. @@ -3355,7 +3365,9 @@ ggml_tensor * llm_graph_context::build_attn( ggml_tensor * k = mctx_cur->get_k(ctx0, il); ggml_tensor * v = ggml_view_4d(ctx0, k, v_cur->ne[0], k->ne[1], k->ne[2], k->ne[3], k->nb[1], k->nb[2], k->nb[3], 0); - ggml_tensor * cur = build_attn_mha(q, k, v, kq_b, kq_mask_top_k, sinks, v_mla, kq_scale, il, /*mask_is_prefix =*/ false); + // a row's live cells are among its top_k: the kernel may read those alone + ggml_tensor * cur = build_attn_mha(q, k, v, kq_b, kq_mask_top_k, sinks, v_mla, kq_scale, il, /*mask_is_prefix =*/ false, + top_k->ne[0]); cb(cur, "kqv_out", il); if (wo) { @@ -4000,19 +4012,26 @@ llm_graph_input_kpool * llm_graph_context::build_inp_kpool( GGML_ASSERT(n_ps >= 1 && (int64_t) ubatch.n_seqs_unq == n_ps*n_stream); const int64_t n_pools = llama_kpool_n_pools(n_kv, kpool, n_ps); + const int64_t n_dump = llama_kpool_select_k(n_pools, hparams.indexer_top_k, kpool); GGML_ASSERT(kq_mask->ne[0] == n_kv && kq_mask->ne[3] == n_stream); GGML_ASSERT(kq_mask->ne[1] == n_tps && "the pooled indexer needs an unpadded KQ mask"); - inp->pool_cells = ggml_new_tensor_2d(ctx0, GGML_TYPE_I32, kpool*n_pools, n_stream); + inp->pool_cells = ggml_new_tensor_2d(ctx0, GGML_TYPE_I32, kpool*(n_pools + n_dump), n_stream); ggml_set_input(inp->pool_cells); ggml_set_name(inp->pool_cells, "kpool_pool_cells"); + inp->pool_cells_3d = ggml_reshape_3d(ctx0, inp->pool_cells, kpool, n_pools + n_dump, n_stream); + inp->pool_bias = ggml_new_tensor_3d(ctx0, GGML_TYPE_F32, n_pools, n_tps, n_stream); ggml_set_input(inp->pool_bias); ggml_set_name(inp->pool_bias, "kpool_pool_bias"); + inp->pool_dump = ggml_new_tensor_3d(ctx0, GGML_TYPE_F32, n_dump, n_tps, n_stream); + ggml_set_input(inp->pool_dump); + ggml_set_name(inp->pool_dump, "kpool_pool_dump"); + // the fused indexer wants f16; built once, shared by every indexer layer if (cparams.fused_lid) { inp->pool_bias_f16 = ggml_cast(ctx0, @@ -4022,7 +4041,7 @@ llm_graph_input_kpool * llm_graph_context::build_inp_kpool( } // lossless in f16 (only 0.0f and -INFINITY), and f16 + f32 -> f16 adds the KQ mask uncast - inp->sel_mask = ggml_new_tensor_4d(ctx0, GGML_TYPE_F16, n_kv, n_tps, 1, n_stream); + inp->sel_mask = ggml_new_tensor_4d(ctx0, GGML_TYPE_F16, n_kv + kpool*n_dump, n_tps, 1, n_stream); ggml_set_input(inp->sel_mask); ggml_set_name(inp->sel_mask, "kpool_sel_mask"); @@ -4087,9 +4106,16 @@ ggml_tensor * llm_graph_context::build_attn_sparse( GGML_ASSERT(sel_mask->type == GGML_TYPE_F16 || sel_mask->type == GGML_TYPE_F32); GGML_ASSERT(sel_mask->type == cand_mask->type); - GGML_ASSERT(ggml_are_same_shape(sel_mask, cand_mask)); - GGML_ASSERT(sel_mask->ne[0] == kq_mask->ne[0] && sel_mask->ne[1] == kq_mask->ne[1] && - sel_mask->ne[3] == kq_mask->ne[3]); + GGML_ASSERT(cand_mask->ne[0] == kq_mask->ne[0] && cand_mask->ne[1] == kq_mask->ne[1] && + cand_mask->ne[3] == kq_mask->ne[3]); + GGML_ASSERT(sel_mask->ne[1] == cand_mask->ne[1] && sel_mask->ne[2] == cand_mask->ne[2] && + sel_mask->ne[3] == cand_mask->ne[3]); + + // sel_mask's columns past n_kv are the dump pools' (llm_graph_input_kpool::pool_dump): a row the top-k fills with + // dump pools scatters into them, never twice into one cell, and they are dropped below + const int64_t n_kv = cand_mask->ne[0]; + + GGML_ASSERT(sel_mask->ne[0] >= n_kv); // ggml_set_rows writes THROUGH, and sel_mask is shared per ubatch: scatter into a copy ggml_tensor * mask_all = ggml_dup(ctx0, sel_mask); @@ -4107,7 +4133,8 @@ ggml_tensor * llm_graph_context::build_attn_sparse( ggml_tensor * mask_top_k = ggml_set_rows(ctx0, mask_all, zeros, top_k_3d); - mask_top_k = ggml_view_4d(ctx0, mask_top_k, mask_top_k->ne[1], mask_top_k->ne[2], 1, mask_top_k->ne[3], + // the first n_kv columns are the mask; the dump columns are dropped here + mask_top_k = ggml_view_4d(ctx0, mask_top_k, n_kv, mask_top_k->ne[2], 1, mask_top_k->ne[3], mask_top_k->nb[2], mask_top_k->nb[3], mask_top_k->nb[3], 0); mask_top_k = ggml_add(ctx0, mask_top_k, cand_mask); @@ -4125,7 +4152,13 @@ ggml_tensor * llm_graph_context::build_attn_sparse( ggml_tensor * k = mctx_cur->get_k(ctx0, il); ggml_tensor * v = ggml_view_4d(ctx0, k, v_cur->ne[0], k->ne[1], k->ne[2], k->ne[3], k->nb[1], k->nb[2], k->nb[3], 0); - ggml_tensor * cur = build_attn_mha(q, k, v, kq_b, mask_top_k, sinks, v_mla, kq_scale, il, /*mask_is_prefix =*/ false); + // 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); + 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); if (wo) { diff --git a/src/llama-graph.h b/src/llama-graph.h index 17630424a714..b138c71d846b 100644 --- a/src/llama-graph.h +++ b/src/llama-graph.h @@ -1282,7 +1282,8 @@ struct llm_graph_context { ggml_tensor * v_mla, // [n_embd_head_v_mla, n_embd_head_v, n_head_v] float kq_scale, int il, - bool mask_is_prefix = true) const; // false: kq_mask selects cells (sparse attention's top-k) + bool mask_is_prefix = true, // false: kq_mask selects cells (sparse attention's top-k) + int64_t n_kv_max = 0) const; // > 0: no mask row has more live cells (ggml_flash_attn_ext_set_n_kv_max) llm_graph_input_attn_no_cache * build_attn_inp_no_cache() const; @@ -1493,8 +1494,8 @@ struct llm_graph_context { ggml_tensor * kq_b, ggml_tensor * sinks, // [n_head_q] ggml_tensor * v_mla, // [n_embd_head_v_mla, n_embd_head_v, n_head_v] - ggml_tensor * top_k, // I32 [n_select, n_tokens/n_stream, n_stream] - ggml_tensor * sel_mask, // F16/F32 [n_kv, n_batch, 1, n_stream] + 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] float kq_scale, int il) const; diff --git a/src/llama-kv-cache-kpool.cpp b/src/llama-kv-cache-kpool.cpp index eab9b9d4081a..a41b2764cd15 100644 --- a/src/llama-kv-cache-kpool.cpp +++ b/src/llama-kv-cache-kpool.cpp @@ -1,11 +1,15 @@ #include "llama-kv-cache-kpool.h" #include "llama-batch.h" +#include "llama-impl.h" #include "llama-kv-cache.h" #include "llama-kv-cells.h" #include +#include +#include #include +#include #include uint32_t llama_kpool_n_pools(uint32_t n_kv, uint32_t kpool, uint32_t n_seqs) { @@ -39,12 +43,22 @@ static void kpool_mask_fill(T * dst, int64_t n) { std::fill(dst, dst + n, kpool_mask_of::from(-INFINITY)); } +static void kpool_mask_fill(char * dst, int64_t n, bool f16) { + if (f16) { + kpool_mask_fill((ggml_fp16_t *) dst, n); + } else { + kpool_mask_fill((float *) dst, n); + } +} + +// pool_of[j] - b_base is the pool of cell j among the run, or negative when its pool is not complete template static void kpool_mask_row( T * cur_sel, T * cur_cand, const llama_pos * pos_at, const int32_t * pool_of, + int32_t b_base, int64_t n_kv, llama_pos q, llama_pos tail_start, @@ -54,7 +68,7 @@ static void kpool_mask_row( 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] < (uint32_t) bo_vis; + const bool pooled = (uint32_t) (pool_of[j] - b_base) < (uint32_t) bo_vis; const bool tail = pos_at[j] >= tail_start; cur_sel [j] = vis && tail ? v_sel : v_mask; @@ -62,8 +76,214 @@ static void kpool_mask_row( } } -void llama_kv_cache_set_input_kpool( - const llama_kv_cache * kv, +// 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 +// is read at the last cell of a block only. One cell per slot keeps the rows of the write different (set_rows writes them from +// several threads). -1 when none is left +struct kpool_spare_cells { + const llama_kv_cells & cells; + llama_pos r; + const int64_t * rows; // the rows of the real slots + int64_t n_rows; + int64_t row0; // the row of cell 0 + std::vector written; + bool sorted = false; + uint32_t next = 0; + + int32_t take() { + if (!sorted) { + sorted = true; + for (int64_t i = 0; i < n_rows; ++i) { + written.push_back(rows[i] - row0); + } + std::sort(written.begin(), written.end()); + } + + for (; next < cells.size(); ++next) { + if ((cells.is_empty(next) || cells.pos_get(next) % r != r - 1) && !std::binary_search(written.begin(), written.end(), (int64_t) next)) { + return (int32_t) next++; + } + } + return -1; + } +}; + +// where one stream's maps go: the tensors of llama_kpool_set_input, at this stream's first element +struct kpool_stream_out { + int32_t * pool_cells = nullptr; + char * sel_mask = nullptr; + char * cand_mask = nullptr; + float * pool_bias = nullptr; + int32_t * pool_reps = nullptr; // nullptr with the last three when the key cache is off + int32_t * new_cells = nullptr; + int64_t * new_reps = nullptr; + + int64_t strm_row = 0; // the stream's first row of the key cache + int64_t n_kv = 0; + int64_t n_sel = 0; // a sel_mask row: n_kv and the dump columns + int64_t n_tps = 0; + int64_t n_pools = 0; + int64_t n_new_max = 0; + int64_t r = 0; + + bool mask_f16 = false; + size_t mask_ts = 0; +}; + +// 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( + const kpool_stream_out & o, + const llama_kpool_views::view & v, + const llama_kv_cells & cells, + llama_seq_id seq, + const llama_ubatch * ubatch, + int64_t s, + bool rebuild) { + const int64_t r = o.r; + + // the run is every pool from the lowest position to the highest, and starts at the lowest + int64_t b_base = 0; + int64_t n_run = 0; + + if (v.n > 0) { + b_base = cells.seq_pos_min(seq)/r; + n_run = cells.seq_pos_max(seq)/r - b_base + 1; + } + + const int64_t off = v.n > 0 ? b_base - v.b_lo : 0; + + GGML_ASSERT(n_run <= o.n_pools && off >= 0 && off + n_run <= (int64_t) v.cap); + + const uint16_t * fill = v.fill.data() + off; + + std::copy_n(v.slot.data() + off*r, n_run*r, o.pool_cells); + + int64_t n_new = 0; + + // 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; + + if (o.pool_reps) { + std::copy_n(v.reps.data() + off, n_run, o.pool_reps); + + for (int64_t p = 0; p < n_run; ++p) { + if (fill[p] == r) { + any_rep_src = o.pool_cells + p*r; + break; + } + } + + auto emit = [&](int64_t p) { + // bounded while a sequence's ubatch tokens are a contiguous run (llama-batch.cpp enforces); + // fail loudly, clamping would serve a stale key + GGML_ASSERT(n_new < o.n_new_max && "k-pool: more pools completed than the fixed bound"); + + const int32_t * member = o.pool_cells + p*r; + + std::copy(member, member + r, o.new_cells + n_new*r); + + o.new_reps[n_new] = o.strm_row + member[r - 1]; + + n_new++; + }; + + if (rebuild) { + for (int64_t p = 0; p < n_run; ++p) { + if (fill[p] == r) { + emit(p); + } + } + } else { + // the pools this ubatch's tokens are in, in pool order + std::vector touched; + + for (int64_t ii = 0; ii < o.n_tps; ++ii) { + const int64_t i = s*o.n_tps + ii; + + if (ubatch->seq_id[i][0] != seq) { + continue; + } + + const int64_t bo = ubatch->pos[i]/r - b_base; + + if (bo >= 0 && bo < n_run) { + touched.push_back(bo); + } + } + + std::sort(touched.begin(), touched.end()); + touched.erase(std::unique(touched.begin(), touched.end()), touched.end()); + + for (const int64_t p : touched) { + if (fill[p] == r) { + emit(p); + } + } + } + } + + int64_t n_done = 0; + + for (int64_t ii = 0; ii < o.n_tps; ++ii) { + const int64_t i = s*o.n_tps + ii; + + if (ubatch->seq_id[i][0] != seq) { + continue; + } + + const llama_pos q = ubatch->pos[i]; + + GGML_ASSERT(q >= 0); + + n_done++; + + const llama_pos tail_start = (q + 1)/r*r; + + // the reference tests visibility at a pool's LAST member, so a straddled pool drops whole + const int64_t bo_vis = std::max(0, tail_start/r - b_base); + + 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) { + 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 { + 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); + } + + float * q_pool_bias = o.pool_bias + ii*o.n_pools; + + for (int64_t p = 0, n_vis = std::min(n_run, bo_vis); p < n_vis; ++p) { + q_pool_bias[p] = fill[p] == r ? 0.0f : -INFINITY; + } + } + + // exactly one partition per row, or a query reads another sequence's pools + GGML_ASSERT(n_done == o.n_tps && "every query must belong to a sequence of the ubatch"); + + if (o.pool_reps) { + // the fixed row count means unused slots must name a safe destination: a spare cell each (their members stay cell 0) + kpool_spare_cells spare = { cells, (llama_pos) r, o.new_reps, n_new, o.strm_row, {} }; + + for (int64_t p = n_new; p < o.n_new_max; ++p) { + if (const int32_t cell = spare.take(); cell >= 0) { + o.new_reps[p] = o.strm_row + cell; + } else if (any_rep_src) { + std::copy(any_rep_src, any_rep_src + r, o.new_cells + p*r); + o.new_reps[p] = o.strm_row + any_rep_src[r - 1]; + } else { + o.new_reps[p] = o.strm_row; + } + } + } +} + +void llama_kpool_set_input( + const llama_kpool_cells_fn & cells_of, + llama_kpool_views * views, ggml_tensor * cell_pool, ggml_tensor * pool_cells, ggml_tensor * bias, @@ -78,7 +298,7 @@ void llama_kv_cache_set_input_kpool( bool rebuild, const llama_ubatch * ubatch, uint32_t kpool) { - GGML_ASSERT(kv != nullptr); + GGML_ASSERT(cells_of != nullptr); GGML_ASSERT(kpool > 0); GGML_ASSERT(ggml_backend_buffer_is_host(pool_cells->buffer)); @@ -97,7 +317,8 @@ void llama_kv_cache_set_input_kpool( GGML_ASSERT(ggml_is_contiguous(sel_mask)); GGML_ASSERT(ggml_is_contiguous(cand_mask)); - const int64_t n_kv = sel_mask->ne[0]; + const int64_t n_kv = cand_mask->ne[0]; + const int64_t n_sel = sel_mask->ne[0]; const int64_t n_ns = sel_mask->ne[3]; const int64_t r = kpool; const int64_t n_tokens = ubatch->n_tokens; @@ -107,15 +328,17 @@ void llama_kv_cache_set_input_kpool( GGML_ASSERT(n_ns == 1 || (int64_t) ubatch->n_seqs_unq == n_ns); const int64_t n_ps = (int64_t) ubatch->n_seqs_unq/n_ns; - const int64_t n_pools = pool_cells->ne[0]/r; + const int64_t n_pools = pool_bias->ne[0]; + const int64_t n_dump = pool_cells->ne[0]/r - n_pools; GGML_ASSERT(n_ps > 0 && (int64_t) ubatch->n_seqs_unq == n_ns*n_ps); GGML_ASSERT(pool_cells->ne[0] % r == 0); GGML_ASSERT(n_pools >= 2*n_ps); + GGML_ASSERT(n_dump >= 0 && n_sel == n_kv + r*n_dump && "sel_mask carries kpool dump columns for each dump pool"); GGML_ASSERT(pool_cells->ne[1] == n_ns); GGML_ASSERT(sel_mask->ne[2] == 1); - GGML_ASSERT(ggml_are_same_shape(cand_mask, sel_mask)); - GGML_ASSERT(pool_bias->ne[0] == n_pools && pool_bias->ne[2] == n_ns); + GGML_ASSERT(cand_mask->ne[1] == sel_mask->ne[1] && cand_mask->ne[2] == 1 && cand_mask->ne[3] == n_ns); + GGML_ASSERT(pool_bias->ne[2] == n_ns); GGML_ASSERT(n_tokens % n_ns == 0); const int64_t n_tps = n_tokens/n_ns; @@ -185,8 +408,8 @@ void llama_kv_cache_set_input_kpool( const size_t mask_ts = ggml_type_size(sel_mask->type); // -1 marks a cell with no usable pool; host side only, never copied into cell_pool - std::vector pool_of(n_kv); - std::vector filled(n_pools); + std::vector pool_of; + std::vector filled; std::vector pos_at; std::vector run_off(n_ps); @@ -196,36 +419,80 @@ void llama_kv_cache_set_input_kpool( return n_ps == 1 ? ubatch->seq_id[s*n_tps][0] : ubatch->seq_id_unq[ps]; }; + if (views) { + views->begin(); + } + for (int64_t s = 0; s < n_ns; ++s) { - int32_t * cur_pool_cells = dst_pool_cells + s*(r*n_pools); - char * cur_sel_mask = dst_sel_mask + s*(n_padq*n_kv)*mask_ts; + int32_t * cur_pool_cells = dst_pool_cells + s*(r*(n_pools + n_dump)); + char * cur_sel_mask = dst_sel_mask + s*(n_padq*n_sel)*mask_ts; char * cur_cand_mask = dst_cand_mask + s*(n_padq*n_kv)*mask_ts; float * cur_pool_bias = dst_pool_bias + s*(n_tps*n_pools); std::fill(cur_pool_cells, cur_pool_cells + r*n_pools, 0); std::fill(cur_pool_bias, cur_pool_bias + n_tps*n_pools, -INFINITY); + // the dump pools name the dump columns, one cell each, so no two slots of a row name one cell + for (int64_t c = 0; c < r*n_dump; ++c) { + cur_pool_cells[r*n_pools + c] = (int32_t) (n_kv + c); + } + int32_t * cur_pool_reps = kcache ? dst_pool_reps + s*n_pools : nullptr; int32_t * cur_new_cells = kcache ? dst_new_cells + s*(r*n_new_max) : nullptr; int64_t * cur_new_reps = kcache ? dst_new_reps + s*n_new_max : nullptr; int64_t n_new = 0; - // pads the fixed-size write; recomputing a complete pool is idempotent, so a repeat is safe + // 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); std::fill(cur_new_cells, cur_new_cells + r*n_new_max, 0); } - if (mask_f16) { - kpool_mask_fill((ggml_fp16_t *) (cur_sel_mask + n_tps*n_kv*mask_ts), (n_padq - n_tps)*n_kv); - kpool_mask_fill((ggml_fp16_t *) (cur_cand_mask + n_tps*n_kv*mask_ts), (n_padq - n_tps)*n_kv); - } else { - kpool_mask_fill((float *) (cur_sel_mask + n_tps*n_kv*mask_ts), (n_padq - n_tps)*n_kv); - kpool_mask_fill((float *) (cur_cand_mask + n_tps*n_kv*mask_ts), (n_padq - n_tps)*n_kv); + // the padded rows whole and a row's dump columns: the rows' first n_kv are kpool_mask_row's + for (int64_t ii = 0; ii < n_padq; ++ii) { + const int64_t j0 = ii < n_tps ? n_kv : 0; + + kpool_mask_fill(cur_sel_mask + (ii*n_sel + j0)*mask_ts, n_sel - j0, mask_f16); + } + kpool_mask_fill(cur_cand_mask + n_tps*n_kv*mask_ts, (n_padq - n_tps)*n_kv, mask_f16); + + // a sequence alone in its stream is served from its view; else (or when the view cannot) from the cells below + if (views && n_ps == 1 && !cell_pool && !bias) { + const llama_seq_id seq = seq_of(s, 0); + + if (const auto * v = views->serve(cells_of(seq), seq, r, n_kv, n_pools)) { + kpool_stream_out o; + + o.pool_cells = cur_pool_cells; + o.sel_mask = cur_sel_mask; + o.cand_mask = cur_cand_mask; + o.pool_bias = cur_pool_bias; + o.pool_reps = cur_pool_reps; + o.new_cells = cur_new_cells; + o.new_reps = cur_new_reps; + o.strm_row = kcache ? (int64_t) strm_of[s]*kv_size : 0; + o.n_kv = n_kv; + o.n_sel = n_sel; + o.n_tps = n_tps; + o.n_pools = n_pools; + o.n_new_max = n_new_max; + o.r = r; + o.mask_f16 = mask_f16; + o.mask_ts = mask_ts; + + kpool_stream_from_view(o, *v, cells_of(seq), seq, ubatch, s, rebuild); + + continue; + } + + views->get_stats().n_direct++; } // [TAG_KPOOL_PACK] one packed run per sequence, NOT one full-width table: the indexer @@ -235,7 +502,7 @@ void llama_kv_cache_set_input_kpool( for (int64_t ps = 0; ps < n_ps; ++ps) { const llama_seq_id seq = seq_of(s, ps); - const auto & cells = kv->get_cells(seq); + const auto & cells = cells_of(seq); int64_t b_min = 0; int64_t b_max = 0; @@ -277,15 +544,15 @@ void llama_kv_cache_set_input_kpool( for (int64_t ps = 0; ps < n_ps; ++ps) { const llama_seq_id seq_of_pool = seq_of(s, ps); - const auto & cells = kv->get_cells(seq_of_pool); + const auto & cells = cells_of(seq_of_pool); const int64_t n_run = run_len[ps]; int32_t * cur_cell_pool = dst_cell_pool ? dst_cell_pool + s*n_kv : nullptr; int32_t * part_pool_cells = cur_pool_cells + run_off[ps]*r; - std::fill(pool_of.begin(), pool_of.end(), -1); - std::fill(filled.begin(), filled.end(), 0); + pool_of.assign(n_kv, -1); + filled .assign(n_pools, 0); pos_at.resize(n_kv); for (int64_t j = 0; j < n_kv; ++j) { @@ -376,6 +643,15 @@ void llama_kv_cache_set_input_kpool( continue; } + // a rep is written once: sequences that share cells, or cells that share a position, give the same rep twice + const int32_t rep = part_pool_cells[p*r + (r - 1)]; + + if (emitted[rep]) { + continue; + } + + emitted[rep] = 1; + // bounded while a sequence's ubatch tokens are a contiguous run (llama-batch.cpp enforces); // fail loudly, clamping would serve a stale key GGML_ASSERT(n_new < n_new_max && "k-pool: more pools completed than the fixed bound"); @@ -383,7 +659,7 @@ void llama_kv_cache_set_input_kpool( std::copy(part_pool_cells + p*r, part_pool_cells + (p + 1)*r, cur_new_cells + n_new*r); - cur_new_reps[n_new] = (int64_t) strm_of[s]*kv_size + part_pool_cells[p*r + (r - 1)]; + cur_new_reps[n_new] = (int64_t) strm_of[s]*kv_size + rep; n_new++; } @@ -408,15 +684,15 @@ void llama_kv_cache_set_input_kpool( const int64_t bo_vis = std::max(0, tail_start/r - b_base); float * cur_bias = dst_bias ? dst_bias + i*n_kv : nullptr; - char * cur_sel = cur_sel_mask + ii*n_kv*mask_ts; - char * cur_cand = cur_cand_mask + ii*n_kv*mask_ts; + 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) { kpool_mask_row((ggml_fp16_t *) cur_sel, (ggml_fp16_t *) cur_cand, - pos_at.data(), pool_of.data(), n_kv, q, tail_start, bo_vis); + pos_at.data(), pool_of.data(), 0, n_kv, q, tail_start, bo_vis); } else { kpool_mask_row((float *) cur_sel, (float *) cur_cand, - pos_at.data(), pool_of.data(), n_kv, q, tail_start, bo_vis); + pos_at.data(), pool_of.data(), 0, n_kv, q, tail_start, bo_vis); } if (cur_bias) { @@ -443,10 +719,13 @@ void llama_kv_cache_set_input_kpool( GGML_ASSERT(n_done == n_tps && "every query must belong to a sequence of the ubatch"); if (kcache) { - // the fixed row count means unused slots must name a safe destination: repeat a complete - // pool (recompute is a no-op), or cell 0 when none exists (nothing reads its pooled third) + // the fixed row count means unused slots must name a safe destination: a spare cell each (their members stay cell 0) + kpool_spare_cells spare = { cells_of(seq_of(s, 0)), (llama_pos) r, cur_new_reps, n_new, (int64_t) strm_of[s]*kv_size, {} }; + for (int64_t p = n_new; p < n_new_max; ++p) { - if (any_rep_src) { + if (const int32_t cell = spare.take(); cell >= 0) { + cur_new_reps[p] = (int64_t) strm_of[s]*kv_size + cell; + } else if (any_rep_src) { std::copy(any_rep_src, any_rep_src + r, cur_new_cells + p*r); cur_new_reps[p] = (int64_t) strm_of[s]*kv_size + any_rep_src[r - 1]; } else { @@ -457,6 +736,305 @@ void llama_kv_cache_set_input_kpool( } } +void llama_kpool_views::view::rebuild(const llama_kv_cells & c, llama_seq_id seq) { + built = true; + ok = false; + + const uint32_t size = c.size(); + + cap = size/r + 2; + n = c.seq_cell_count(seq); + + // two cells at one position: a slot holds one cell, so the maps come from the cells + if (n != c.seq_pos_distinct(seq)) { + return; + } + + int64_t b_min = 0; + + if (n > 0) { + b_min = c.seq_pos_min(seq)/r; + + if (c.seq_pos_max(seq)/r - b_min + 1 > (int64_t) cap) { + return; + } + } + + b_lo = b_min; + + pos_at.assign(size, -1); + pblk .assign(size, -1); + slot .assign((size_t) cap*r, 0); + fill .assign(cap, 0); + reps .assign(cap, 0); + + if (n > 0) { + uint32_t found = 0; + + for (uint32_t j = 0; j < size; ++j) { + if (c.is_empty(j) || !c.seq_has(j, seq)) { + continue; + } + + const llama_pos p = c.pos_get(j); + + pos_at[j] = p; + slot[p - b_lo*r] = (int32_t) j; + fill[p/r - b_lo]++; + + found++; + } + + if (found != n) { + return; + } + + for (int64_t bi = 0; bi < (int64_t) cap; ++bi) { + if (fill[bi] == r) { + for (uint32_t k = 0; k < r; ++k) { + pblk[slot[bi*r + k]] = (int32_t) (b_lo + bi); + } + + reps[bi] = slot[bi*r + r - 1]; + } + } + } + + ok = true; +} + +bool llama_kpool_views::view::insert(uint32_t j, llama_pos p) { + const int64_t b = p/r; + const int64_t bi = b - b_lo; + + if (bi < 0 || bi >= (int64_t) cap) { + return false; + } + + int32_t & at = slot[p - b_lo*r]; + + // the slot's cell still holds this position: two cells claim it + if (pos_at[at] == p) { + return false; + } + + at = (int32_t) j; + pos_at[j] = p; + + n++; + + if (++fill[bi] == r) { + for (uint32_t k = 0; k < r; ++k) { + pblk[slot[bi*r + k]] = (int32_t) b; + } + + reps[bi] = slot[bi*r + r - 1]; + } + + return true; +} + +void llama_kpool_views::view::remove(uint32_t j) { + const llama_pos p = pos_at[j]; + const int64_t bi = p/r - b_lo; + + // a complete pool has every slot filled, so each of its cells drops the pool's mark + if (fill[bi] == r) { + for (uint32_t k = 0; k < r; ++k) { + pblk[slot[bi*r + k]] = -1; + } + + reps[bi] = 0; + } + + slot[p - b_lo*r] = 0; + fill[bi]--; + pos_at[j] = -1; + + n--; +} + +bool llama_kpool_views::view::apply(const llama_kv_cells & c, llama_seq_id seq, const std::vector & chg) { + upd.clear(); + + for (const uint32_t j : chg) { + GGML_ASSERT(j < pos_at.size()); + + const llama_pos p = !c.is_empty(j) && c.seq_has(j, seq) ? c.pos_get(j) : -1; + + if (p != pos_at[j]) { + upd.emplace_back(j, p); + } + } + + // every removal first: a cell leaves a position before another takes it, in whatever order they changed + for (const auto & u : upd) { + if (pos_at[u.first] >= 0) { + remove(u.first); + } + } + + for (const auto & u : upd) { + if (u.second >= 0 && !insert(u.first, u.second)) { + return false; + } + } + + return true; +} + +void llama_kpool_views::sync(const llama_kv_cells & cells, uint32_t r) { + uint64_t & at = synced[&cells]; + + if (at == call) { + return; + } + + at = call; + + // a list past 1/16 of the cells costs more to apply than the views cost to rebuild + cells.changes().track(std::max(256, cells.size()/16)); + + const bool listed = cells.changes().take(chg); + + std::sort(chg.begin(), chg.end()); + chg.erase(std::unique(chg.begin(), chg.end()), chg.end()); + + for (auto & [seq, v] : views) { + if (v.cells != &cells) { + continue; + } + + // no list: the view is stale, and is rebuilt if a stream asks for it + if (!listed || !v.built || v.r != r) { + v.built = false; + continue; + } + + if (!v.ok || !v.apply(cells, seq, chg)) { + v.rebuild(cells, seq); + st.n_rebuilt++; + } + } +} + +const llama_kpool_views::view * llama_kpool_views::serve( + const llama_kv_cells & cells, llama_seq_id seq, uint32_t r, int64_t n_kv, int64_t n_pools) { + sync(cells, r); + + auto it = views.find(seq); + + if (it == views.end()) { + if (views.size() >= n_views_max) { + return nullptr; + } + + it = views.emplace(seq, view()).first; + } + + view & v = it->second; + + // a view that lost the cells (the list is gone), or disagrees with their count of the sequence + if (!v.built || v.cells != &cells || v.r != r || (v.ok && v.n != cells.seq_cell_count(seq))) { + v.cells = &cells; + v.r = r; + + v.rebuild(cells, seq); + st.n_rebuilt++; + } + + if (!v.ok || (int64_t) cells.used_max_p1() > n_kv) { + return nullptr; + } + + if (v.n > 0 && cells.seq_pos_max(seq)/r - cells.seq_pos_min(seq)/r + 1 > n_pools) { + return nullptr; + } + + st.n_served++; + + return &v; +} + +void llama_kv_cache_set_input_kpool( + const llama_kv_cache * kv, + ggml_tensor * cell_pool, + ggml_tensor * pool_cells, + ggml_tensor * bias, + ggml_tensor * pool_bias, + ggml_tensor * sel_mask, + ggml_tensor * cand_mask, + ggml_tensor * pool_reps, + ggml_tensor * new_pool_cells, + ggml_tensor * new_pool_reps, + const uint32_t * strm_of, + int64_t kv_size, + bool rebuild, + const llama_ubatch * ubatch, + uint32_t kpool) { + GGML_ASSERT(kv != nullptr); + + static const bool legacy = ggml_env_switch("LLAMA_KPOOL_INPUT_LEGACY"); + static const bool check = ggml_env_switch("LLAMA_KPOOL_INPUT_CHECK"); + + const llama_kpool_cells_fn cells_of = [kv](llama_seq_id seq) -> const llama_kv_cells & { + return kv->get_cells(seq); + }; + + llama_kpool_views * views = legacy ? nullptr : &kv->get_kpool_views(); + + 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; + } + + // the same maps from the cells; the views must have written the same bytes + ggml_tensor * outs[] = { + cell_pool, pool_cells, bias, pool_bias, sel_mask, cand_mask, pool_reps, new_pool_cells, new_pool_reps, + }; + + std::vector> kept; + + for (const ggml_tensor * t : outs) { + kept.emplace_back(t ? ggml_nbytes(t) : 0); + + if (t) { + memcpy(kept.back().data(), t->data, ggml_nbytes(t)); + } + } + + llama_kpool_set_input(cells_of, nullptr, + 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); + + for (size_t i = 0; i < kept.size(); ++i) { + if (outs[i] == nullptr) { + continue; + } + + const uint8_t * now = (const uint8_t *) outs[i]->data; + + for (size_t k = 0; k < kept[i].size(); ++k) { + if (kept[i][k] != now[k]) { + GGML_ABORT("k-pool input: the view maps and the maps from the cells differ in %s at byte %zu", + ggml_get_name(outs[i]), k); + } + } + } + + static uint64_t n_checked = 0; + + if (++n_checked % 512 == 0) { + const auto & st = views->get_stats(); + + 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); + } +} + void llm_graph_input_kpool::set_input(const llama_ubatch * ubatch) { // unconditional: the key/gate STORE runs on the dense path too, or cells below n_select // would have no indexer state when the first ubatch crosses it @@ -476,6 +1054,9 @@ 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( mctx_attn->get_kv(), /* cell_pool */ nullptr, pool_cells, /* bias */ nullptr, pool_bias, diff --git a/src/llama-kv-cache-kpool.h b/src/llama-kv-cache-kpool.h index f13a33014b84..0e563d504d69 100644 --- a/src/llama-kv-cache-kpool.h +++ b/src/llama-kv-cache-kpool.h @@ -5,10 +5,14 @@ #include "llama-graph.h" #include +#include +#include +#include struct llama_ubatch; class llama_kv_cache; class llama_kv_cache_context; +class llama_kv_cells; // GLM-5-Next indexer pooling. no input may hold a negative index (ggml_set_rows asserts // i1 >= 0, ggml_get_rows has no sentinel), so unusable entries are clamped and masked. @@ -19,9 +23,101 @@ uint32_t llama_kpool_n_pools(uint32_t n_kv, uint32_t kpool, uint32_t n_seqs = 1) // select_k of Glm5NextTextIndexer.forward: must run over POOLS, a cell cut takes partial pools uint32_t llama_kpool_select_k(uint32_t n_pools, uint32_t indexer_top_k, uint32_t kpool); +// The pool maps of a decode step are a function of the cells alone, and between two steps a few cells change. A view +// is one sequence's cells in pool terms (per-cell positions, pool table, completeness), kept up to date from the cells +// llama_kv_cells_log names, so a step reads only the cells that changed. A view that cannot follow (the log no longer +// reaches back, two cells share a position, the positions span more pools than a map holds) is rebuilt from the cells, +// or the maps are built from the cells (llama_kpool_set_input with views == nullptr, the reference). +class llama_kpool_views { +public: + struct view { + const llama_kv_cells * cells = nullptr; + + uint32_t r = 0; + uint32_t cap = 0; // pools the tables hold: llama_kpool_n_pools of the whole cache + bool built = false; // the tables were built, and follow the cells since + bool ok = false; // the tables can serve; built && !ok: build from the cells + + int64_t b_lo = 0; // first pool of the tables: position p is at slot p - b_lo*r, pool b at b - b_lo + uint32_t n = 0; // cells that hold the sequence + + std::vector pos_at; // per cell: the position if the cell holds the sequence, else -1 + std::vector pblk; // per cell: its pool (p/r) if that pool is complete, else -1 + std::vector slot; // [cap*r] position -> cell, 0 if none + std::vector fill; // [cap] cells in the pool + std::vector reps; // [cap] cell of the last position of a complete pool, else 0 + + void rebuild(const llama_kv_cells & c, llama_seq_id seq); + + // false: the cells disagree with the tables (two cells at one position, a position off the tables); rebuild + bool apply(const llama_kv_cells & c, llama_seq_id seq, const std::vector & chg); + + private: + bool insert(uint32_t j, llama_pos p); + void remove(uint32_t j); + + std::vector> upd; + }; + + struct stats { + uint64_t n_served = 0; // streams built from a view + uint64_t n_rebuilt = 0; // views rebuilt from the cells + uint64_t n_direct = 0; // streams built from the cells + }; + + // one call of llama_kpool_set_input: every cells object takes its changes again + void begin() { ++call; } + + // the view of seq after the changes of its cells, nullptr when it cannot serve a map of n_pools over n_kv cells + const view * serve(const llama_kv_cells & cells, llama_seq_id seq, uint32_t r, int64_t n_kv, int64_t n_pools); + + stats & get_stats() { return st; } + +private: + void sync(const llama_kv_cells & cells, uint32_t r); + + // a sequence past this many takes the maps built from the cells + static constexpr size_t n_views_max = 8; + + std::map views; + + // the call in which a cells object last took its changes + std::map synced; + + std::vector chg; + + uint64_t call = 1; + stats st; +}; + +using llama_kpool_cells_fn = std::function; + +// 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( + const llama_kpool_cells_fn & cells_of, + llama_kpool_views * views, + ggml_tensor * cell_pool, + ggml_tensor * pool_cells, + ggml_tensor * bias, + ggml_tensor * pool_bias, + ggml_tensor * sel_mask, + ggml_tensor * cand_mask, + ggml_tensor * pool_reps, + ggml_tensor * new_pool_cells, + ggml_tensor * new_pool_reps, + const uint32_t * strm_of, + int64_t kv_size, + bool rebuild, + const llama_ubatch * ubatch, + uint32_t kpool); + // `kv` must be the ATTENTION (MLA) cache; the indexer cache shares its slot layout. -// pool_cells pool member -> cell, 0 if not resident +// LLAMA_KPOOL_INPUT_LEGACY=1 builds the maps from the cells every call; LLAMA_KPOOL_INPUT_CHECK=1 builds them both +// ways and aborts on a difference. +// pool_cells pool member -> cell, 0 if not resident; past pool_bias's n_pools, n_dump dump pools of cells n_kv + d*kpool + t // pool_bias computed, NOT gathered at the last member, which an incomplete pool lacks +// 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 @@ -53,6 +149,7 @@ struct llm_graph_input_kpool_dims { int64_t n_tps = 0; int64_t n_ps = 0; int64_t n_pools = 0; + int64_t n_dump = 0; int64_t n_new_max = 0; bool rebuild = false; bool scoring = false; @@ -84,9 +181,16 @@ class llm_graph_input_kpool : public llm_graph_input_i { static bool shapes_match(const llm_graph_input_kpool_dims & dims, const llm_graph_input_kpool & inp); ggml_tensor * k_idxs = nullptr; // I32 [n_tokens] - ggml_tensor * pool_cells = nullptr; // I32 [kpool*n_pools, n_stream] + ggml_tensor * pool_cells = nullptr; // I32 [kpool*(n_pools + n_dump), n_stream] ggml_tensor * pool_bias = nullptr; // F32 [n_pools, n_tps, n_stream] + // the top-k's dump pools, n_dump = select_k of them after the n_pools real ones: each scores -FLT_MAX, above a dead pool's + // -inf and under every live score, so a row takes one only where fewer than select_k pools are live, and its cells are its + // own, past n_kv (sel_mask's dump columns), so the rows of the mask's scatter stay unique + ggml_tensor * pool_dump = nullptr; // F32 [n_dump, n_tps, n_stream], every element -FLT_MAX + ggml_tensor * pool_cells_3d = nullptr; // pool_cells as [kpool, n_pools + n_dump, n_stream]; built once: a view of an input + // made in each layer is a split input, and a host copy, of its own + // pooled-key cache: a pool's value lives in the row of its LAST member, a pure function of // cell content (seq_cp shares it, rebase leaves it alone) ggml_tensor * pool_reps = nullptr; // I32 [n_pools, n_stream] stream-local rep cell @@ -103,7 +207,7 @@ class llm_graph_input_kpool : public llm_graph_input_i { // 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] - ggml_tensor * sel_mask = nullptr; // F16 [n_kv, n_batch, 1, n_stream] + ggml_tensor * sel_mask = nullptr; // F16 [n_kv + kpool*n_dump, n_batch, 1, n_stream]; the dump columns -inf ggml_tensor * cand_mask = nullptr; // F16 [n_kv, n_batch, 1, n_stream] const llama_kv_cache_context * mctx_attn; diff --git a/src/llama-kv-cache.cpp b/src/llama-kv-cache.cpp index 3c45220f6fec..7897ea6fadda 100644 --- a/src/llama-kv-cache.cpp +++ b/src/llama-kv-cache.cpp @@ -3,6 +3,7 @@ #include "gguf.h" #include "llama-impl.h" #include "llama-io.h" +#include "llama-kv-cache-kpool.h" #include "llama-model.h" #include "llama-context.h" @@ -1296,6 +1297,14 @@ void llama_kv_cache::clear_kpool_dirty() const { kpool_dirty = false; } +llama_kpool_views & llama_kv_cache::get_kpool_views() const { + if (!kpool_views) { + kpool_views = std::make_shared(); + } + + return *kpool_views; +} + ggml_type llama_kv_cache::type_k() const { return layers[0].k->type; } diff --git a/src/llama-kv-cache.h b/src/llama-kv-cache.h index 58d70bb87cc9..c89858d3ac01 100644 --- a/src/llama-kv-cache.h +++ b/src/llama-kv-cache.h @@ -5,6 +5,7 @@ #include "llama-kv-cells.h" #include "llama-memory.h" +#include #include #include @@ -13,6 +14,8 @@ struct llama_hparams; struct llama_model; struct llama_context; +class llama_kpool_views; + // // llama_kv_cache // @@ -164,6 +167,9 @@ class llama_kv_cache : public llama_memory_i { bool get_kpool_dirty() const; void clear_kpool_dirty() 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; + ggml_type type_k() const; ggml_type type_v() const; @@ -268,6 +274,8 @@ class llama_kv_cache : public llama_memory_i { // see set_kpool_dirty. mutable: its only consumer runs from set_input, holding a const cache mutable bool kpool_dirty = false; + mutable std::shared_ptr kpool_views; + const uint32_t n_seq_max = 1; const uint32_t n_stream = 1; diff --git a/src/llama-kv-cells.h b/src/llama-kv-cells.h index 14ed247b7015..51c9e50eda4b 100644 --- a/src/llama-kv-cells.h +++ b/src/llama-kv-cells.h @@ -28,6 +28,66 @@ struct llama_kv_cell_ext { } }; +// the cells whose position or sequences changed, for ONE consumer that keeps a view of the cells +// (llama_kpool_views): the consumer takes the list, and is told to rebuild the view when it cannot be given one +class llama_kv_cells_log { +public: + llama_kv_cells_log() = default; + + // a copy starts untracked and empty; an assignment replaces the cells, so any of them may have changed + llama_kv_cells_log(const llama_kv_cells_log &) {} + llama_kv_cells_log & operator=(const llama_kv_cells_log &) { all(); return *this; } + + // consumer side: const, because the consumer holds the cells const + void track(uint32_t max_cells) const { + if (!on) { + on = true; + everyone = true; + } + cap = max_cells; + } + + // the cells changed since the last take() go to out; false = any cell may have changed (the first take, a reset, a + // resize, an assignment, a list that passed its cap): out is empty then + bool take(std::vector & out) const { + const bool res = !everyone; + + out.clear(); + if (res) { + out.swap(list); + } + list.clear(); + everyone = false; + + return res; + } + + void all() { + if (on) { + everyone = true; + list.clear(); + } + } + + void add(uint32_t i) { + if (!on || everyone || (!list.empty() && list.back() == i)) { + return; + } + if (list.size() >= cap) { + all(); + return; + } + list.push_back(i); + } + +private: + mutable bool on = false; + mutable bool everyone = false; + mutable uint32_t cap = 0; + + mutable std::vector list; +}; + // meta information about KV cells that can be part of multiple sequences at the same time // TODO: add unit tests class llama_kv_cells { @@ -46,10 +106,13 @@ class llama_kv_cells { for (uint32_t s = 0; s < LLAMA_MAX_SEQ; ++s) { seq_pos[s].clear(); + seq_n[s] = 0; std::fill(seq_bits[s].begin(), seq_bits[s].end(), 0); } std::fill(pos_max_g.begin(), pos_max_g.end(), -1); + + chg.all(); } void reset_shift() { @@ -180,6 +243,8 @@ class llama_kv_cells { for (uint32_t j = 0; j < other.pos.size(); ++j) { const auto idx = i + j; + chg.add(idx); + if (pos[idx] == -1 && other.pos[j] != -1) { used.insert(i + j); } @@ -216,6 +281,8 @@ class llama_kv_cells { for (uint32_t j = 0; j < other.pos.size(); ++j) { const auto idx = idxs[j]; + chg.add(idx); + if (pos[idx] == -1 && other.pos[j] != -1) { used.insert(idx); } @@ -250,6 +317,8 @@ class llama_kv_cells { assert(i < pos.size()); assert(pos[i] != -1); + chg.add(i); + seq_pos_rm(i); seq_bits_clr(i); seq[i].reset(); @@ -270,6 +339,8 @@ class llama_kv_cells { assert(pos[i] != -1); assert(seq_id >= 0); + chg.add(i); + seq[i].reset(seq_id); seq_bit_clr(i, seq_id); seq_pos_dec(seq_id, pos[i]); @@ -292,6 +363,8 @@ class llama_kv_cells { bool seq_keep(uint32_t i, llama_seq_id seq_id) { assert(i < pos.size()); + chg.add(i); + if (seq[i].test(seq_id)) { seq_pos_rm(i); seq_bits_clr(i); @@ -346,6 +419,8 @@ class llama_kv_cells { assert(pos[i] != -1); assert(!seq[i].test(seq_id)); + chg.add(i); + seq[i].set(seq_id); seq_bit_set(i, seq_id); seq_pos_inc(seq_id, pos[i]); @@ -395,6 +470,25 @@ class llama_kv_cells { return seq_pos[seq_id].rbegin()->first; } + // the number of cells that hold sequence seq_id, and the number of different positions among them: + // equal when no two of them share a position + uint32_t seq_cell_count(llama_seq_id seq_id) const { + assert(seq_id >= 0 && seq_id < LLAMA_MAX_SEQ); + + return seq_n[seq_id]; + } + + uint32_t seq_pos_distinct(llama_seq_id seq_id) const { + assert(seq_id >= 0 && seq_id < LLAMA_MAX_SEQ); + + return seq_pos[seq_id].size(); + } + + // the cells changed since the consumer last took them + const llama_kv_cells_log & changes() const { + return chg; + } + // the cells in groups of 64, for a KQ mask row to take a group whole where no position in it can be masked // (llama_kv_cache_set_input_kq_mask): bit k set = cell 64*g + k holds seq_id uint64_t seq_cells64(llama_seq_id seq_id, uint32_t g) const { @@ -464,6 +558,8 @@ class llama_kv_cells { assert(pos[i] == -1); assert(seq[i].none()); + chg.add(i); + pos[i] = p; used.insert(i); @@ -482,6 +578,8 @@ class llama_kv_cells { assert(i < pos.size()); assert(pos[i] != -1); + chg.add(i); + seq_pos_rm(i); pos[i] += d; @@ -516,6 +614,8 @@ class llama_kv_cells { const llama_pos p_old = pos[i]; + chg.add(i); + seq_pos_rm(i); pos[i] /= d; @@ -569,6 +669,11 @@ class llama_kv_cells { // std::map seq_pos[LLAMA_MAX_SEQ]; + // seq_n[s]: the cells that hold sequence s, the sum of the counts in seq_pos[s] + uint32_t seq_n[LLAMA_MAX_SEQ] = {}; + + llama_kv_cells_log chg; + // seq_bits[s]: the bits of seq_cells64() for sequence s, sized at its first cell (a sequence never placed has none); // pos_max_g: the bounds of pos_max64() std::vector seq_bits[LLAMA_MAX_SEQ]; @@ -622,12 +727,15 @@ class llama_kv_cells { auto it = seq_pos[s].find(p); assert(it != seq_pos[s].end()); + seq_n[s]--; + if (--it->second == 0) { seq_pos[s].erase(it); } } void seq_pos_inc(llama_seq_id s, llama_pos p) { + seq_n[s]++; seq_pos[s][p]++; } diff --git a/src/llama-quant.cpp b/src/llama-quant.cpp index 39af78137f4a..092d7862321b 100644 --- a/src/llama-quant.cpp +++ b/src/llama-quant.cpp @@ -1082,7 +1082,13 @@ static void llama_model_quantize_impl(const std::string & fname_inp, const std:: metadata[i].target_type = tensor->type; } - metadata[i].requires_imatrix = tensor_requires_imatrix(tensor->name, metadata[i].target_type, ftype); + // A tensor already in its target type is copied verbatim further down (the cur_type != new_type check), so it + // needs no imatrix however demanding that type is. Asking the target type alone made requantizing a model that + // already holds IQ-family tensors impossible without an imatrix even when those tensors were left untouched -- + // and because --dry-run only records the requirement instead of raising it, the dry run passed and the real run + // failed on the same arguments. + metadata[i].requires_imatrix = tensor->type != metadata[i].target_type && + tensor_requires_imatrix(tensor->name, metadata[i].target_type, ftype); if (params->imatrix) { metadata[i].remapped_imatrix_name = remap_imatrix(tensor->name, mapped); diff --git a/src/models/glm5next.cpp b/src/models/glm5next.cpp index 1ade59ed6c09..d083b95c0f4a 100644 --- a/src/models/glm5next.cpp +++ b/src/models/glm5next.cpp @@ -314,8 +314,17 @@ ggml_tensor * llama_model_glm5next::graph::build_kda_layer( nb_head, nb_qkv, nb_qkv*n_seq_tokens, ggml_row_size(conv_out->type, 2*d_inner)); + // the output gate reads the layer input, not the recurrence's output: its two low-rank products go into the graph + // here, after the convolution, g_a beside f_a (two mat-vecs of one input, which ggml-cuda runs as one launch) and + // g_b before the recurrence, so that nothing after the recurrence waits on them + ggml_build_forward_expand(gf, qk); + ggml_tensor * f_a = ggml_mul_mat(ctx0, layer.ssm_f_a, inp); + ggml_tensor * gate = ggml_mul_mat(ctx0, layer.ssm_g_b, ggml_mul_mat(ctx0, layer.ssm_g_a, inp)); + ggml_build_forward_expand(gf, f_a); + ggml_build_forward_expand(gf, gate); + // g = lower_bound * sigmoid(exp(A_log)*(f_b(f_a(x)) + dt_bias)); it scales, not clamps - ggml_tensor * g_raw = ggml_mul_mat(ctx0, layer.ssm_f_b, ggml_mul_mat(ctx0, layer.ssm_f_a, inp)); + ggml_tensor * g_raw = ggml_mul_mat(ctx0, layer.ssm_f_b, f_a); g_raw = ggml_add(ctx0, g_raw, layer.ssm_dt_b); ggml_tensor * g = ggml_reshape_3d(ctx0, g_raw, head_dim, n_head, n_tokens); g = ggml_mul(ctx0, g, ggml_reshape_3d(ctx0, layer.ssm_a, 1, n_head, 1)); @@ -348,7 +357,6 @@ ggml_tensor * llama_model_glm5next::graph::build_kda_layer( ggml_tensor * o = ggml_reshape_3d(ctx0, build_cont(out), head_dim, n_head, n_tokens); cb(o, "kda_scan_out", il); - ggml_tensor * gate = ggml_mul_mat(ctx0, layer.ssm_g_b, ggml_mul_mat(ctx0, layer.ssm_g_a, inp)); gate = ggml_reshape_3d(ctx0, gate, head_dim, n_head, n_tokens); // plain sigmoid gate, not the SiLU that FusedRMSNormGated defaults to @@ -404,7 +412,7 @@ ggml_tensor * llama_model_glm5next::graph::build_indexer( const int64_t n_kv = kbuf->ne[2]; const int64_t n_stream = kbuf->ne[3]; const int64_t n_tps = n_tokens/n_stream; - const int64_t n_pools = inp_kp->pool_cells->ne[0]/r; + const int64_t n_pools = inp_kp->pool_bias->ne[0]; GGML_ASSERT(kbuf->ne[0] == d_idx && kbuf->ne[1] == 3 && "the pooled indexer cache needs a key head, a gate head and a pooled head"); @@ -447,12 +455,16 @@ ggml_tensor * llama_model_glm5next::graph::build_indexer( ggml_build_forward_expand(gf, mctx_idx->cpy_k_part(ctx0, pool_new, inp_kp->new_pool_reps, il, d_idx, 2*d_idx)); - ggml_tensor * pooled_rd = ggml_view_3d(ctx0, kbuf, d_idx, n_kv, n_stream, - kbuf->nb[2], kbuf->nb[3], 2*d_idx*kbuf->nb[0]); + // 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 = [&]() { + ggml_tensor * pooled_rd = ggml_view_3d(ctx0, kbuf, d_idx, n_kv, n_stream, + kbuf->nb[2], kbuf->nb[3], 2*d_idx*kbuf->nb[0]); - ggml_tensor * pool_k = ggml_get_rows(ctx0, pooled_rd, inp_kp->pool_reps); - pool_k = ggml_reshape_4d(ctx0, pool_k, d_idx, n_pools, 1, n_stream); - cb(pool_k, "indexer_pool_k", il); + ggml_tensor * pool_k = ggml_get_rows(ctx0, pooled_rd, inp_kp->pool_reps); + pool_k = ggml_reshape_4d(ctx0, pool_k, d_idx, n_pools, 1, n_stream); + cb(pool_k, "indexer_pool_k", il); + return pool_k; + }; // no rope: n_rot() is 0 for the whole text tower ggml_tensor * iq = ggml_mul_mat(ctx0, layer.indexer_attn_q_b, qr); @@ -469,17 +481,28 @@ ggml_tensor * llama_model_glm5next::graph::build_indexer( ggml_tensor * pool_score = nullptr; if (cparams.fused_lid) { - // pool_k stays f32 so the kernel takes its f32 path; f16 wmma would undo the prec - ggml_tensor * pool_kf = ggml_reshape_4d(ctx0, pool_k, d_idx, 1, n_pools, n_stream); + static const bool gather_legacy = ggml_env_switch("LLAMA_INDEXER_GATHER_LEGACY"); + + if (!gather_legacy) { + // the pooled keys read where the cache holds them, pool p at cell pool_reps[p], scored in f32 like the + // gathered rows (ggml_lightning_indexer_rows): no copy of every pool's key each token + ggml_tensor * pooled = ggml_view_4d(ctx0, kbuf, d_idx, 1, n_kv, n_stream, + kbuf->nb[1], kbuf->nb[2], kbuf->nb[3], 2*d_idx*kbuf->nb[0]); - pool_score = ggml_lightning_indexer(ctx0, iq, pool_kf, w, inp_kp->pool_bias_f16); + pool_score = ggml_lightning_indexer_rows(ctx0, iq, pooled, inp_kp->pool_reps, w, inp_kp->pool_bias_f16); + } else { + // pool_k stays f32 so the kernel takes its f32 path; f16 wmma would undo the prec + ggml_tensor * pool_kf = ggml_reshape_4d(ctx0, gather_pool_k(), d_idx, 1, n_pools, n_stream); + + pool_score = ggml_lightning_indexer(ctx0, iq, pool_kf, w, inp_kp->pool_bias_f16); + } cb(pool_score, "indexer_pool_score", il); res->add_fused_node({LLM_FUSED_OP_LIGHTNING_INDEXER, pool_score, il}); pool_score = ggml_reshape_3d(ctx0, pool_score, n_pools, n_tps, n_stream); } else { - ggml_tensor * kq = ggml_mul_mat(ctx0, pool_k, ggml_permute(ctx0, iq, 0, 2, 1, 3)); + ggml_tensor * kq = ggml_mul_mat(ctx0, gather_pool_k(), ggml_permute(ctx0, iq, 0, 2, 1, 3)); // the ReLU sits BETWEEN the per-head dot and the head weighting; either side differs kq = ggml_cont(ctx0, ggml_permute(ctx0, kq, 2, 1, 0, 3)); @@ -497,14 +520,18 @@ ggml_tensor * llama_model_glm5next::graph::build_indexer( // top-k over POOLS then expand: a cell-level top-k is wrong, relu ties span pool bounds const int64_t select_k = llama_kpool_select_k(n_pools, hparams.indexer_top_k, r); GGML_ASSERT(select_k > 0 && select_k <= n_pools); + GGML_ASSERT(inp_kp->pool_dump->ne[0] == select_k && "one dump pool for each slot of the top-k"); + + // fewer than select_k live pools (finite scores) fill the top-k with dump pools (-FLT_MAX), never with a dead pool + // (-inf), whose cells may be another slot's: each dump pool's cells are its own, past n_kv (llm_graph_input_kpool) + pool_score = ggml_concat(ctx0, pool_score, inp_kp->pool_dump, 0); ggml_tensor * sel = ggml_cont(ctx0, ggml_top_k(ctx0, pool_score, (int) select_k)); cb(sel, "indexer_top_k_pools", il); - ggml_tensor * pc3 = ggml_reshape_3d(ctx0, inp_kp->pool_cells, r, n_pools, n_stream); ggml_tensor * sel_flat = ggml_reshape_2d(ctx0, sel, select_k*n_tps, n_stream); - ggml_tensor * top_k = ggml_get_rows(ctx0, pc3, sel_flat); + ggml_tensor * top_k = ggml_get_rows(ctx0, inp_kp->pool_cells_3d, sel_flat); GGML_ASSERT(top_k->type == GGML_TYPE_I32 && "pool_cells is I32, so the gather stays I32"); top_k = ggml_reshape_3d(ctx0, top_k, r*select_k, n_tps, n_stream); cb(top_k, "indexer_top_k", il); diff --git a/src/models/models.h b/src/models/models.h index 314bfdf93f9d..576493fefb6b 100644 --- a/src/models/models.h +++ b/src/models/models.h @@ -1399,7 +1399,8 @@ struct llama_model_glm5next : public llama_model_base { ggml_tensor * cur, int il) const; - // always stores the key and gate; when `scoring`, returns the selected CELL indices + // always stores the key and gate; when `scoring`, returns the selected CELL indices, a dump pool's cells where + // fewer than select_k pools are live ggml_tensor * build_indexer( const llama_layer & layer, llm_graph_input_kpool * inp_kp, diff --git a/tests/CMakeLists.txt b/tests/CMakeLists.txt index 81fe71c8008b..fbe9cba0278d 100644 --- a/tests/CMakeLists.txt +++ b/tests/CMakeLists.txt @@ -162,6 +162,7 @@ if (NOT WIN32 OR NOT BUILD_SHARED_LIBS) llama_build_and_test(test-batch-alloc.cpp) llama_build_and_test(test-kv-mask.cpp) llama_build_and_test(test-kpool-can-reuse.cpp) + llama_build_and_test(test-kpool-input.cpp) 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) @@ -338,6 +339,7 @@ if (NOT GGML_BACKEND_DL) llama_build_and_test(test-backend-meta-views.cpp) llama_build_and_test(test-backend-meta-split.cpp) llama_build_and_test(test-backend-meta-capture.cpp) + llama_build_and_test(test-backend-meta-sourceless.cpp) llama_build_and_test(test-pq2-mma-device-memory.cpp) llama_build_and_test(test-ptq1_0-element-map.cpp) llama_build_and_test(test-ptq1_0-cuda-dot.cpp) diff --git a/tests/test-backend-meta-sourceless.cpp b/tests/test-backend-meta-sourceless.cpp new file mode 100644 index 000000000000..faa16707ce57 --- /dev/null +++ b/tests/test-backend-meta-sourceless.cpp @@ -0,0 +1,81 @@ +// The meta backend (-sm tensor) derives each node's split from its sources' splits. A node with no sources, as +// ggml_arange, had none to derive it from: its split was unknown, and allocating the graph aborted +// (GGML_ASSERT(ret.axis != GGML_BACKEND_SPLIT_AXIS_UNKNOWN)). GLM-5.3's sparse attention makes one for the dump columns of +// its filler slots (build_attn_sparse), so under -sm tensor every decode aborted. Every device computes the same values +// for such a node: it is mirrored. +// +// y = x + arange(n) is allocated and computed through the meta backend, over the first GPU twice (it splits over two +// devices as two GPUs do), and must hold x[i] + i; with no GPU the test is skipped. + +#include "ggml.h" +#include "ggml-alloc.h" +#include "ggml-backend.h" + +#include +#include + +static constexpr int64_t n = 64; + +static ggml_backend_meta_split_state mirrored(const ggml_tensor * /*tensor*/, void * /*userdata*/) { + return {GGML_BACKEND_SPLIT_AXIS_MIRRORED, {0}, {1}, 1}; +} + +int main() { + ggml_backend_dev_t gpu = nullptr; + for (size_t i = 0; i < ggml_backend_dev_count() && gpu == nullptr; i++) { + if (ggml_backend_dev_type(ggml_backend_dev_get(i)) == GGML_BACKEND_DEVICE_TYPE_GPU) { + gpu = ggml_backend_dev_get(i); + } + } + if (gpu == nullptr) { + printf("no GPU for the meta backend, skipped\n"); + return 0; + } + ggml_backend_dev_t devs[2] = {gpu, gpu}; + ggml_backend_dev_t meta = ggml_backend_meta_device(devs, 2, mirrored, nullptr); + + ggml_backend_t cpu = ggml_backend_init_by_type(GGML_BACKEND_DEVICE_TYPE_CPU, nullptr); + ggml_backend_t backend = ggml_backend_dev_init(meta, nullptr); + + const ggml_init_params params = { + /*.mem_size =*/ 4*ggml_tensor_overhead() + ggml_graph_overhead(), + /*.mem_buffer =*/ nullptr, + /*.no_alloc =*/ true, + }; + ggml_context * ctx = ggml_init(params); + ggml_tensor * x = ggml_new_tensor_1d(ctx, GGML_TYPE_F32, n); + ggml_set_input(x); + ggml_tensor * y = ggml_add(ctx, x, ggml_arange(ctx, 0.0f, (float) n, 1.0f)); + ggml_set_output(y); + ggml_cgraph * gf = ggml_new_graph(ctx); + ggml_build_forward_expand(gf, y); + + ggml_backend_t backends[2] = {backend, cpu}; + ggml_backend_sched_t sched = ggml_backend_sched_new(backends, nullptr, 2, GGML_DEFAULT_GRAPH_SIZE, false, true); + GGML_ASSERT(ggml_backend_sched_alloc_graph(sched, gf)); + + std::vector x_data(n); + for (int64_t i = 0; i < n; i++) { + x_data[i] = 0.5f + 100.0f*i; + } + ggml_backend_tensor_set(x, x_data.data(), 0, ggml_nbytes(x)); + GGML_ASSERT(ggml_backend_sched_graph_compute(sched, gf) == GGML_STATUS_SUCCESS); + std::vector y_data(n); + ggml_backend_tensor_get(y, y_data.data(), 0, ggml_nbytes(y)); + + bool ok = true; + for (int64_t i = 0; i < n && ok; i++) { + if (y_data[i] != x_data[i] + (float) i) { + printf("FAIL y[%lld] = %g, expected %g\n", (long long) i, y_data[i], x_data[i] + (float) i); + ok = false; + } + } + + ggml_backend_sched_free(sched); + ggml_free(ctx); + ggml_backend_free(backend); + ggml_backend_free(cpu); + + printf("%s\n", ok ? "OK" : "FAIL"); + return ok ? 0 : 1; +} diff --git a/tests/test-backend-ops.cpp b/tests/test-backend-ops.cpp index e1f9d328fcdc..cc88ed199a34 100644 --- a/tests/test-backend-ops.cpp +++ b/tests/test-backend-ops.cpp @@ -299,6 +299,33 @@ static void init_tensor_kq_mask_band(ggml_tensor * tensor, int band) { set_tensor_kq_mask(tensor, data_f32); } +// n_kv_max (ggml_flash_attn_ext_set_n_kv_max): each row finite at n_kv_max - row % 17 random cells, as a sparse attention's +// top-k selects them, the other cells -inf; rows differ in their count, so a row's list ends at different points +static void init_tensor_kq_mask_sparse(ggml_tensor * tensor, int64_t n_kv_max) { + GGML_ASSERT(tensor->type == GGML_TYPE_F16); + GGML_ASSERT(n_kv_max > 1 && n_kv_max <= tensor->ne[0]); + + const int64_t ne0 = tensor->ne[0]; + const int64_t nrows = ggml_nrows(tensor); + std::vector data_f32(ggml_nelements(tensor), -INFINITY); + std::vector order(ne0); + for (int64_t i = 0; i < ne0; ++i) { + order[i] = i; + } + + std::mt19937 gen(0x5A17); + for (int64_t row = 0; row < nrows; ++row) { + std::shuffle(order.begin(), order.end(), gen); + const int64_t count = n_kv_max - row % std::min(n_kv_max, 17); + std::sort(order.begin(), order.begin() + count); + for (int64_t i = 0; i < count; ++i) { + data_f32[row*ne0 + order[i]] = -0.03125f * (1 + (i + row) % 7); + } + } + + set_tensor_kq_mask(tensor, data_f32); +} + // generate a lower triangular matrix static void init_tensor_tril(ggml_tensor * tensor, float min = -1.0f, float max = 1.0f) { GGML_ASSERT(tensor->type == GGML_TYPE_F32); @@ -3734,6 +3761,47 @@ struct test_rms_norm_mul_add : public test_case { } }; +// GGML_OP_RMS_NORM + GGML_OP_MUL with nothing after it, which CUDA fuses on its own; the weight one row broadcast over +// the rest, as a model's norm weight is, or full size +struct test_rms_norm_mul : public test_case { + const std::array ne; + const float eps; + const bool broadcast; + + std::string op_desc(ggml_tensor * t) override { + GGML_UNUSED(t); + return "RMS_NORM_MUL"; + } + + bool run_whole_graph() override { return true; } + + std::string vars() override { + return VARS_TO_STR3(ne, eps, broadcast); + } + + test_rms_norm_mul(std::array ne = {64, 5, 4, 3}, float eps = 1e-6f, bool broadcast = false) + : ne(ne), eps(eps), broadcast(broadcast) {} + + ggml_tensor * build_graph(ggml_context * ctx) override { + ggml_tensor * a = ggml_new_tensor(ctx, GGML_TYPE_F32, 4, ne.data()); + ggml_tensor * b = ggml_new_tensor(ctx, GGML_TYPE_F32, 4, ne.data()); + ggml_tensor * w = broadcast ? ggml_new_tensor_1d(ctx, GGML_TYPE_F32, ne[0]) : b; + + // the norm's input computed in the graph, as in test_rms_norm_mul_add + a = ggml_add(ctx, a, b); + ggml_tensor * out = ggml_mul(ctx, ggml_rms_norm(ctx, a, eps), w); + ggml_set_name(out, "out"); + + return out; + } + + void initialize_tensors(ggml_context * ctx) override { + for (ggml_tensor * t = ggml_get_first_tensor(ctx); t != NULL; t = ggml_get_next_tensor(ctx, t)) { + init_tensor_uniform(t, -10.f, 10.f); + } + } +}; + // GGML_OP_RMS_NORM without a weight straight into GGML_OP_MUL_MAT (the hyper-connection mixes of DeepSeek V4 and // GLM-5.3); CUDA runs both as one matvec at one column struct test_rms_norm_mul_mat : public test_case { @@ -4416,6 +4484,7 @@ struct test_dsv4_hc_pre_fused : public test_dsv4_hc { ggml_tensor * mixes = ggml_mul_mat(ctx, hc_fn, ggml_rms_norm(ctx, flat, 1e-5f)); weights = ggml_dsv4_hc_weights(ctx, mixes, scale, base, 1e-6f, n_iter); ggml_set_name(weights, "hc_weights"); + ggml_set_output(weights); // read back after the evaluation, as the tests compare it ggml_tensor * pre = ggml_view_2d(ctx, weights, hc, n_tokens, weights->nb[1], 0); post = ggml_view_2d(ctx, weights, hc, n_tokens, weights->nb[1], hc*weights->nb[0]); @@ -4445,6 +4514,71 @@ struct test_dsv4_hc_pre_fused : public test_dsv4_hc { } }; +// 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. +struct test_dsv4_hc_pre_q8_1 : public test_dsv4_hc_pre_fused { + const ggml_type type_q; + const int n_readers; + + ggml_tensor * sum = nullptr; + + std::string op_desc(ggml_tensor * t) override { + GGML_UNUSED(t); + return "DSV4_HC_PRE_Q8_1"; + } + + std::vector fusion_test_nodes() override { return { weights, out, sum }; } + + std::string vars() override { + return VARS_TO_STR5(type_q, n_readers, n_embd, n_tokens, with_post); + } + + // the readers' activations quantized, to q8_1 on CUDA and to the weights' vec_dot_type on the CPU + double max_nmse_err() override { return 5e-4; } + + test_dsv4_hc_pre_q8_1(ggml_type type_q, int n_readers, int64_t n_embd, int64_t n_tokens, bool with_post) + : test_dsv4_hc_pre_fused(GGML_TYPE_BF16, n_embd, n_tokens, 20, false, true, with_post), type_q(type_q), + n_readers(n_readers) {} + + ggml_tensor * build_graph(ggml_context * ctx, ggml_context * ctx_weights) override { + ggml_tensor * mix = test_dsv4_hc_pre_fused::build_graph(ctx, ctx_weights); + ggml_context * ctx_w = ctx_weights != nullptr ? ctx_weights : ctx; + for (int r = 0; r < n_readers; ++r) { + ggml_tensor * w = ggml_new_tensor_2d(ctx_w, type_q, n_embd, 256); + ggml_set_name(w, ("proj" + std::to_string(r)).c_str()); + ggml_tensor * y = ggml_mul_mat(ctx, w, mix); + sum = r == 0 ? y : ggml_add(ctx, sum, y); + } + ggml_set_name(sum, "sum"); + return sum; + } +}; + +// The front, then a DSV4_HC_POST that reads its post and comb weights at once (its output the normed mix, its residual +// the front's streams): no model's order, but the tightest one, the post's reads of the weights right after the +// front's writes of them. +struct test_dsv4_hc_pre_post : public test_dsv4_hc_pre_fused { + ggml_tensor * post_out = nullptr; + + std::string op_desc(ggml_tensor * t) override { + GGML_UNUSED(t); + return "DSV4_HC_PRE_POST"; + } + + std::vector fusion_test_nodes() override { return { weights, out, post_out }; } + + test_dsv4_hc_pre_post(ggml_type type_w, int64_t n_embd, int64_t n_tokens, int32_t n_iter) + : test_dsv4_hc_pre_fused(type_w, n_embd, n_tokens, n_iter, false, true, false) {} + + ggml_tensor * build_graph(ggml_context * ctx, ggml_context * ctx_weights) override { + ggml_tensor * mix = test_dsv4_hc_pre_fused::build_graph(ctx, ctx_weights); + post_out = ggml_dsv4_hc_post(ctx, mix, hc_pre->src[0], post, comb); + ggml_set_name(post_out, "post_out"); + return post_out; + } +}; + // GLM-5.3's KDA q, k and v: three MUL_MATs of one input by weights of m rows, created one after another as // create_tensor_qkv creates them, joined along dim 0 by two CONCATs; CUDA runs the five as one launch at up to 8 tokens // when the weights lie one stride from each other. spacer: a tensor made between k's weight and v's, so they do not. @@ -4915,6 +5049,35 @@ struct test_rwkv_wkv6 : public test_case { }; // GGML_OP_GATED_DELTA_NET +// Both gates: a rate per (ne0, ne1) slot and a little jitter per token, so a slow row stays slow. For KDA's gate +// ([head_size, heads, tokens, seqs]) that is a rate per channel and head; for the scalar gate ([1, heads, tokens, seqs]) +// `i % (ne0*ne1)` reduces to the head index, so it is a rate per head. Activated, the rate is log-uniform over the old +// (-20, -1e-4) a token, and about two thirds of the rows (those under 1/16) carry their state across a 16-token chunk; +// raw, the pre-activation is uniform over (-3, 3), and the rows whose g * a[h] is well above 0 carry it. +// +// Drawn independently per token, as BOTH gates were, a row is fast on some token of every chunk, little state reaches a +// chunk boundary, and nothing tests what carries it there. The ranges are identical either way; only the correlation +// across tokens changes, and that correlation is the whole coverage. For the activated gate a row carries when |g| < +// -ln(1e-3)/16 = 0.43: drawn per token that needs all 16 draws to land in the slowest 2.2 % of the range, which happens +// with probability 2e-27, while a per-row rate makes it 0.69. With the i.i.d. draw, a mutant decaying each chunk's state +// row by its neighbour's decay passed a 32-head 512-token case at NMSE 9e-8 (float64 says its true effect is 1e-13 to +// 2e-11, because the data never stressed it); with this one, all seven chunked KDA cases fail it at NMSE 0.75-1.72. +static void init_decay_gate(ggml_tensor * t, bool raw) { + std::mt19937 gen(std::random_device{}()); + std::uniform_real_distribution rate(raw ? -3.0f : logf(1e-4f), raw ? 3.0f : logf(20.0f)); + std::uniform_real_distribution jitter(-0.25f, 0.25f); + std::vector rates(t->ne[0] * t->ne[1]); + for (float & r : rates) { + r = rate(gen); + } + std::vector g(ggml_nelements(t)); + for (size_t i = 0; i < g.size(); i++) { + const float x = rates[i % rates.size()] + jitter(gen); + g[i] = raw ? x : -expf(x); + } + ggml_backend_tensor_set(t, g.data(), 0, ggml_nbytes(t)); +} + struct test_gated_delta_net : public test_case { const ggml_type type; @@ -4944,7 +5107,7 @@ struct test_gated_delta_net : public test_case { double max_nmse_err() override { // the CUDA chunked prefill path (fp16 tensor-core GEMMs, fp32 accumulation) needs a slightly // higher threshold than the default; the shapes it cannot take keep the default - const bool chunk_shape = head_size == 128 && n_seq_tokens >= 128 && !kda && !rows_mode; + const bool chunk_shape = head_size == 128 && n_seq_tokens >= 128 && !rows_mode; return chunk_shape ? 2e-7 : 1e-7; } @@ -5037,8 +5200,9 @@ struct test_gated_delta_net : public test_case { for (ggml_tensor * t = ggml_get_first_tensor(ctx); t != nullptr; t = ggml_get_next_tensor(ctx, t)) { if (ggml_is_view_op(t->op)) { continue; } if (strcmp(t->name, "g") == 0) { + // the scalar gate too: its rate is per head, and a head must be able to stay slow across a whole chunk. // raw: pre-softplus alpha; a[h] * softplus(g + dt_bias) lands in the usual decay range - init_tensor_uniform(t, raw_gates ? -3.0f : -20.0f, raw_gates ? 3.0f : -1e-4f); + init_decay_gate(t, raw_gates); } else if (strcmp(t->name, "beta") == 0) { init_tensor_uniform(t, raw_gates ? -4.0f : 0.0f, raw_gates ? 4.0f : 1.0f); } else if (strcmp(t->name, "dt_bias") == 0) { @@ -5234,8 +5398,8 @@ struct test_gated_delta_net_cache_fusion : public test_case { for (ggml_tensor * t = ggml_get_first_tensor(ctx); t != nullptr; t = ggml_get_next_tensor(ctx, t)) { if (ggml_is_view_op(t->op)) { continue; } if (strcmp(t->name, "g") == 0) { - // raw: see test_gated_delta_net::initialize_tensors - init_tensor_uniform(t, raw_gates ? -3.0f : -20.0f, raw_gates ? 3.0f : -1e-4f); + // see test_gated_delta_net::initialize_tensors + init_decay_gate(t, raw_gates); } else if (strcmp(t->name, "beta") == 0) { init_tensor_uniform(t, raw_gates ? -4.0f : 0.0f, raw_gates ? 4.0f : 1.0f); } else if (strcmp(t->name, "dt_bias") == 0) { @@ -7772,6 +7936,130 @@ struct test_mul_mat_vec_fusion : public test_case { } }; +// A routed FFN as llama-graph builds GLM-5.3-Flash's: gate and up of the tokens' input (the SwiGLU limit's clamps when +// glu_limit > 0), the GLU, then the down projection of its output, the three on one ids. The CUDA backend runs the +// gate/up and then the down on mmvq-moe.cu's ring, and the down, its ids read by the launch before it, lists its experts +// and issues its first tiles before its dependency wait (GGML_CUDA_MMVQ_MOE_IDS_EARLY_LEGACY=1: past it). +struct test_moe_ffn_chain : public test_case { + const ggml_type type; + const int n_mats; + const int n_used; + const int64_t n_embd; + const int64_t n_ff; + const int64_t n_tokens; + const float glu_limit; + + std::string vars() override { + return VARS_TO_STR7(type, n_mats, n_used, n_embd, n_ff, n_tokens, glu_limit); + } + + std::string op_desc(ggml_tensor * t) override { + GGML_UNUSED(t); + return "MOE_FFN_CHAIN"; + } + + bool run_whole_graph() override { return true; } + + double max_nmse_err() override { + return 5e-3; + } + + test_moe_ffn_chain(ggml_type type, int n_mats, int n_used, int64_t n_embd, int64_t n_ff, int64_t n_tokens, + float glu_limit) + : type(type), n_mats(n_mats), n_used(n_used), n_embd(n_embd), n_ff(n_ff), n_tokens(n_tokens), + glu_limit(glu_limit) { + GGML_ASSERT(n_used <= n_mats); + } + + ggml_tensor * build_graph(ggml_context * ctx) override { + ggml_tensor * gates = ggml_new_tensor_3d(ctx, type, n_embd, n_ff, n_mats); + ggml_tensor * ups = ggml_new_tensor_3d(ctx, type, n_embd, n_ff, n_mats); + ggml_tensor * downs = ggml_new_tensor_3d(ctx, type, n_ff, n_embd, n_mats); + ggml_tensor * ids = ggml_new_tensor_2d(ctx, GGML_TYPE_I32, n_mats, n_tokens); + if (n_used != n_mats) { + ids = ggml_view_2d(ctx, ids, n_used, n_tokens, ids->nb[1], 0); + } + ggml_tensor * cur = ggml_new_tensor_3d(ctx, GGML_TYPE_F32, n_embd, 1, n_tokens); + ggml_set_name(cur, "cur"); + + ggml_tensor * gate = ggml_mul_mat_id(ctx, gates, cur, ids); + ggml_tensor * up = ggml_mul_mat_id(ctx, ups, cur, ids); + if (glu_limit > 0.0f) { + gate = ggml_clamp(ctx, gate, -INFINITY, glu_limit); + up = ggml_clamp(ctx, up, -glu_limit, glu_limit); + } + ggml_tensor * act = ggml_swiglu_split(ctx, gate, up); + ggml_tensor * out = ggml_mul_mat_id(ctx, downs, act, ids); + ggml_set_name(out, "out"); + return out; + } + + void initialize_tensors(ggml_context * ctx) override { + init_mul_mat_id_tensors(ctx, n_mats); + } +}; + +// Routed experts' weighted sum as build_moe_ffn writes it: MUL(experts, weights), each slot's view, and the views added in +// slot order. The CUDA backend runs the three as one launch (ggml_cuda_op_moe_weighted_sum, GGML_CUDA_MOE_WSUM_LEGACY=1: +// the MUL and the fused ADDs), which rounds each product and each sum as the nodes do: the check is bit for bit. +struct test_moe_weighted_sum : public test_case { + const int64_t n_embd; + const int n_used; + const int64_t n_tokens; + + std::vector order; + + std::string vars() override { + return VARS_TO_STR3(n_embd, n_used, n_tokens); + } + + std::string op_desc(ggml_tensor * t) override { + GGML_UNUSED(t); + return "MOE_WEIGHTED_SUM"; + } + + bool run_whole_graph() override { return true; } + + std::vector forward_first() override { return order; } + + // the elements whose bits differ + double err(const float * a, const float * b, size_t n) override { + size_t differ = 0; + for (size_t i = 0; i < n; ++i) { + differ += memcmp(&a[i], &b[i], sizeof(float)) != 0; + } + return (double) differ; + } + + double max_nmse_err() override { + return 0.0; + } + + test_moe_weighted_sum(int64_t n_embd, int n_used, int64_t n_tokens) + : n_embd(n_embd), n_used(n_used), n_tokens(n_tokens) {} + + ggml_tensor * build_graph(ggml_context * ctx) override { + ggml_tensor * experts = ggml_new_tensor_3d(ctx, GGML_TYPE_F32, n_embd, n_used, n_tokens); + ggml_tensor * weights = ggml_new_tensor_3d(ctx, GGML_TYPE_F32, 1, n_used, n_tokens); + ggml_set_name(experts, "experts"); + ggml_set_name(weights, "weights"); + + ggml_tensor * weighted = ggml_mul(ctx, experts, weights); + order = { weighted }; + std::vector views; + for (int s = 0; s < n_used; ++s) { + views.push_back(ggml_view_2d(ctx, weighted, n_embd, n_tokens, weighted->nb[2], s*weighted->nb[1])); + order.push_back(views.back()); + } + ggml_tensor * out = views[0]; + for (int s = 1; s < n_used; ++s) { + out = ggml_add(ctx, out, views[s]); + } + ggml_set_name(out, "out"); + return out; + } +}; + // PQ2_0 matmuls on one activation, next to each other in the graph, as qwen35's qkv and z (or q, k and v) sit once // ggml_backend_cuda_graph_optimize has moved them: the CUDA backend runs them as one launch over all their tiles // (GGML_CUDA_PQ2_MMA_GROUP_LEGACY=1: one launch each). Every matrix's output is checked, and the row counts put partial @@ -8495,10 +8783,12 @@ struct test_flash_attn_ext : public test_case { const bool mask_bits; // bit-packed mask (GGML_TYPE_I16, 16 cells per element) instead of f16 const int mask_band; // 0: random mask; else every row sees only a band of the cells, see kq_mask_band const bool mask_prefix; // ggml_flash_attn_ext_set_mask_prefix: a hint only, the result must not depend on it + const int64_t n_kv_max; // ggml_flash_attn_ext_set_n_kv_max: each mask row has at most this many finite cells std::string vars() override { return VARS_TO_STR16(hsk, hsv, nh, nr23, kv, nb, mask, sinks, max_bias, logit_softcap, prec, type_K, type_V, permute, kv_view, v_is_view_of_k) - + "," + VAR_TO_STR(mask_bits) + (mask_band ? "," + VAR_TO_STR(mask_band) : "") + (mask_prefix ? "," + VAR_TO_STR(mask_prefix) : ""); + + "," + VAR_TO_STR(mask_bits) + (mask_band ? "," + VAR_TO_STR(mask_band) : "") + (mask_prefix ? "," + VAR_TO_STR(mask_prefix) : "") + + (n_kv_max ? "," + VAR_TO_STR(n_kv_max) : ""); } double max_nmse_err() override { @@ -8515,10 +8805,11 @@ struct test_flash_attn_ext : public test_case { test_flash_attn_ext(int64_t hsk = 128, int64_t hsv = 128, int64_t nh = 32, std::array nr23 = {1, 1}, int64_t kv = 96, int64_t nb = 8, bool mask = true, bool sinks = false, float max_bias = 0.0f, float logit_softcap = 0.0f, ggml_prec prec = GGML_PREC_F32, ggml_type type_K = GGML_TYPE_F16, ggml_type type_V = GGML_TYPE_F16, std::array permute = {0, 1, 2, 3}, - bool kv_view = true, bool v_is_view_of_k = false, bool mask_bits = false, int mask_band = 0, bool mask_prefix = false) + bool kv_view = true, bool v_is_view_of_k = false, bool mask_bits = false, int mask_band = 0, bool mask_prefix = false, + int64_t n_kv_max = 0) : hsk(hsk), hsv(hsv), nh(nh), nr23(nr23), kv(kv), nb(nb), mask(mask), sinks(sinks), max_bias(max_bias), logit_softcap(logit_softcap), prec(prec), type_K(type_K), type_V(type_V), permute(permute), kv_view(kv_view), v_is_view_of_k(v_is_view_of_k), mask_bits(mask_bits), mask_band(mask_band), - mask_prefix(mask_prefix) {} + mask_prefix(mask_prefix), n_kv_max(n_kv_max) {} ggml_tensor * build_graph(ggml_context * ctx) override { const int64_t hsk_padded = GGML_PAD(hsk, ggml_blck_size(type_K)); @@ -8581,6 +8872,7 @@ struct test_flash_attn_ext : public test_case { ggml_flash_attn_ext_add_sinks(out, s); ggml_flash_attn_ext_set_prec (out, prec); ggml_flash_attn_ext_set_mask_prefix(out, mask_prefix); + ggml_flash_attn_ext_set_n_kv_max(out, n_kv_max); ggml_set_name(out, "out"); return out; @@ -8592,7 +8884,9 @@ struct test_flash_attn_ext : public test_case { // make the sink values more noticeable in order to trigger a test failure when the implementation is wrong init_tensor_uniform(t, -10.0f, 10.0f); } else if (strcmp(t->name, "m") == 0) { - if (mask_band) { + if (n_kv_max > 0) { + init_tensor_kq_mask_sparse(t, n_kv_max); + } else if (mask_band) { init_tensor_kq_mask_band(t, mask_band); } else { init_tensor_kq_mask(t); @@ -9130,6 +9424,75 @@ struct test_lightning_indexer_view : public test_case { } }; +// ggml_lightning_indexer_rows: kv keys read through random rows (repeats included) of a cache laid out as GLM-5.3's pooled +// indexer cache, 3 heads a cell and the key the third, so k is a strided view at the third head's offset. +struct test_lightning_indexer_rows : public test_case { + const int64_t nh; // num indexer heads + const int64_t kv; // keys scored (pools) + const int64_t n_rows; // cache cells + const int64_t nb; // batch size + const int64_t ns; // num streams + const ggml_type type_K; + + std::string vars() override { + return VARS_TO_STR6(nh, kv, n_rows, nb, ns, type_K); + } + + double max_nmse_err() override { + return 1e-6; + } + + test_lightning_indexer_rows(int64_t nh = 32, int64_t kv = 256, int64_t n_rows = 1027, int64_t nb = 1, int64_t ns = 1, + ggml_type type_K = GGML_TYPE_F16) + : nh(nh), kv(kv), n_rows(n_rows), nb(nb), ns(ns), type_K(type_K) {} + + ggml_tensor * build_graph(ggml_context * ctx) override { + const int64_t hsk = 128; + + ggml_tensor * q = ggml_new_tensor_4d(ctx, GGML_TYPE_F32, hsk, nh, nb, ns); + ggml_set_name(q, "q"); + + ggml_tensor * cache = ggml_new_tensor_4d(ctx, type_K, hsk, 3, n_rows, ns); + ggml_set_name(cache, "cache"); + ggml_tensor * k = ggml_view_4d(ctx, cache, hsk, 1, n_rows, ns, cache->nb[1], cache->nb[2], cache->nb[3], 2*cache->nb[1]); + ggml_set_name(k, "k"); + + ggml_tensor * rows = ggml_new_tensor_2d(ctx, GGML_TYPE_I32, kv, ns); + ggml_set_name(rows, "rows"); + + ggml_tensor * w = ggml_new_tensor_4d(ctx, GGML_TYPE_F32, nh, nb, 1, ns); + ggml_set_name(w, "w"); + + ggml_tensor * m = ggml_new_tensor_4d(ctx, GGML_TYPE_F16, kv, nb, 1, ns); + ggml_set_name(m, "m"); + + ggml_tensor * out = ggml_lightning_indexer_rows(ctx, q, k, rows, w, m); + ggml_set_name(out, "out"); + return out; + } + + void initialize_tensors(ggml_context * ctx) override { + std::mt19937 rng(0x1D3); + for (ggml_tensor * t = ggml_get_first_tensor(ctx); t != NULL; t = ggml_get_next_tensor(ctx, t)) { + if (t->view_src != nullptr) { + continue; // k: the cache's values + } + if (strcmp(t->name, "m") == 0) { + init_tensor_kq_mask(t); + } else if (t->type == GGML_TYPE_I32) { + std::uniform_int_distribution row(0, (int32_t) n_rows - 1); + std::vector data(ggml_nelements(t)); + for (auto & x : data) { + x = row(rng); + } + ggml_backend_tensor_set(t, data.data(), 0, ggml_nbytes(t)); + } else { + init_tensor_uniform(t); + } + } + } +}; + // Deserializable generic test case struct input_tensor { ggml_type type; @@ -9916,6 +10279,9 @@ static std::vector> make_test_cases_eval() { 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)); + // a comb of 1,000 iterations with no reader in the graph: any of it left running past the evaluation's return would + // show in the weights read back + test_cases.emplace_back(new test_dsv4_hc_pre_fused(GGML_TYPE_BF16, 4096, 3, 1000)); // the previous sublayer's DSV4_HC_POST before the front, which CUDA's Gram path fuses: its streams are an output for (int64_t n_tokens : { 1, 2, 3, 16, 17 }) { test_cases.emplace_back(new test_dsv4_hc_pre_fused(GGML_TYPE_BF16, 4096, n_tokens, 20, false, true, true)); @@ -9925,6 +10291,27 @@ static std::vector> make_test_cases_eval() { test_cases.emplace_back(new test_dsv4_hc_pre_fused(GGML_TYPE_BF16, 4096, 2, 20, false, false, true)); test_cases.emplace_back(new test_dsv4_hc_pre_fused(GGML_TYPE_BF16, 7168, 2, 20, false, true, true)); test_cases.emplace_back(new test_dsv4_hc_pre_fused(GGML_TYPE_BF16, 64, 5, 4, false, true, true)); + // the front's mix read by quantized MUL_MATs: its q8_1 copy written by the front on CUDA and read where the readers + // run on mul_mat_vec_q (not Q4_K's at 8 tokens on Blackwell, which run on MMQ); at 256 (not a multiple of 512) and at + // 12 tokens (past mul_mat_vec_q's columns) the readers quantize it + for (ggml_type type_q : { GGML_TYPE_Q8_0, GGML_TYPE_Q4_K }) { + for (int n_readers : { 1, 2 }) { + for (int64_t n_tokens : { 1, 3, 8 }) { + test_cases.emplace_back(new test_dsv4_hc_pre_q8_1(type_q, n_readers, 4096, n_tokens, false)); + } + } + } + test_cases.emplace_back(new test_dsv4_hc_pre_q8_1(GGML_TYPE_Q8_0, 2, 4096, 2, true)); + test_cases.emplace_back(new test_dsv4_hc_pre_q8_1(GGML_TYPE_Q4_K, 1, 7168, 3, true)); + test_cases.emplace_back(new test_dsv4_hc_pre_q8_1(GGML_TYPE_Q8_0, 2, 256, 3, false)); + test_cases.emplace_back(new test_dsv4_hc_pre_q8_1(GGML_TYPE_Q8_0, 2, 4096, 12, false)); + // a DSV4_HC_POST reading the front's weights at once: at 1,000 iterations (tens of us) a read of the comb before the + // front has written it would show + for (int64_t n_tokens : { 1, 3, 8 }) { + test_cases.emplace_back(new test_dsv4_hc_pre_post(GGML_TYPE_BF16, 4096, n_tokens, 20)); + } + test_cases.emplace_back(new test_dsv4_hc_pre_post(GGML_TYPE_BF16, 4096, 3, 1000)); + test_cases.emplace_back(new test_dsv4_hc_pre_post(GGML_TYPE_F32, 4096, 2, 2)); // KDA's q, k and v: one launch on CUDA at 1 to 8 tokens; unfused at 9, with the weights' stride broken, and at 101 // rows, where a block of 2 rows would cross from one weight's rows into the next @@ -10637,6 +11024,17 @@ static std::vector> make_test_cases_eval() { // in-place tests test_cases.emplace_back(new test_rms_norm(GGML_TYPE_F32, {64, 5, 4, 3}, false, 1e-6f, true)); + // rows of at most 256 columns over more than one 8-row block of the CUDA warp-per-row kernel, the last one partial + // (17 rows), and KDA's per-head output norm: 128 wide, 32 rows a token + for (uint32_t n : { 33, 128, 256 }) { + test_cases.emplace_back(new test_rms_norm(GGML_TYPE_F32, { n, 17, 3, 2 }, false, 1e-6f)); + test_cases.emplace_back(new test_rms_norm(GGML_TYPE_F32, { n, 17, 3, 2 }, true, 1e-6f)); + test_cases.emplace_back(new test_rms_norm_mul( { n, 32, 6, 1 }, 1e-6f, false)); + test_cases.emplace_back(new test_rms_norm_mul( { n, 32, 6, 1 }, 1e-6f, true)); + test_cases.emplace_back(new test_rms_norm_mul_add(GGML_TYPE_F32, { n, 32, 6, 1 }, 1e-6f, false)); + test_cases.emplace_back(new test_rms_norm_mul_add(GGML_TYPE_F32, { n, 32, 6, 1 }, 1e-6f, true)); + } + for (float eps : { 0.0f, 1e-6f, 1e-4f, 1e-1f, 1.0f }) { for (uint32_t n : { 64, 1025 }) { test_cases.emplace_back(new test_rms_norm_mul_add(GGML_TYPE_F32, { n, 5, 4, 3 }, eps, false)); @@ -10668,7 +11066,7 @@ static std::vector> make_test_cases_eval() { for (ggml_type type_w : {GGML_TYPE_BF16, GGML_TYPE_F16, GGML_TYPE_F32}) { for (int64_t k : {256, 4096, 16384}) { - for (int64_t n : {1, 3}) { + for (int64_t n : {1, 3, 64}) { // 64: past the vector and mmf kernels, on cuBLAS test_cases.emplace_back(new test_rms_norm_mul_mat(type_w, k, 24, n)); } } @@ -10938,6 +11336,21 @@ static std::vector> make_test_cases_eval() { test_cases.emplace_back(new test_mul_mat(type_a, GGML_TYPE_F32, 48, 12, 5120, {2, 1}, {1, 1})); test_cases.emplace_back(new test_mul_mat(type_a, GGML_TYPE_F32, 40, 3, 5122, {1, 1}, {1, 1})); } + // a float src0 mmf tiles into fewer blocks than an SM count (glm5next's ssm_beta and ssm_f_a: 64 and 128 x 4096), + // which CUDA runs on the vector kernel at 2-8 columns and rows x cols <= 1024: both sides of each limit, plain, + // batched and broadcast + for (ggml_type type_a : {GGML_TYPE_F32, GGML_TYPE_F16, GGML_TYPE_BF16}) { + for (int64_t m : {64, 128}) { + for (int64_t n : {2, 3, 5, 8, 9}) { + test_cases.emplace_back(new test_mul_mat(type_a, GGML_TYPE_F32, m, n, 4096, {1, 1}, {1, 1})); + } + } + for (int64_t n : {2, 3}) { + test_cases.emplace_back(new test_mul_mat(type_a, GGML_TYPE_F32, 512, n, 1024, {1, 1}, {1, 1})); + } + test_cases.emplace_back(new test_mul_mat(type_a, GGML_TYPE_F32, 64, 3, 256, {2, 3}, {1, 1})); + test_cases.emplace_back(new test_mul_mat(type_a, GGML_TYPE_F32, 64, 3, 256, {2, 3}, {2, 1})); + } // operands the float kernels cannot read, which must fall back (test_mul_mat_view). Strides: at 48 rows, an src1 // stride odd in floats at every batch the vector kernel takes, and under the fused gate/up vector kernel at one // token; at 64 rows, which mmf tiles, src1 strides of 1 and 2 mod 4 floats and an src0 row stride of 2 mod 4 @@ -11092,11 +11505,40 @@ static std::vector> make_test_cases_eval() { // GLM-5.3-Flash's routed experts (IQ3_XXS, 8 used; 32 experts here, of its 288) at decode and an MTP verify, where // the CUDA backend streams them through mmvq-moe.cu's ring: gate/up rows of 4,096 weights with a token's vector - // shared by its experts, down rows of 2,048 with a vector each, and three tokens that share experts - for (int n : { 1, 3 }) { + // shared by its experts, down rows of 2,048 with a vector each, and three tokens that share experts (fewer pairs than + // experts: each pair decodes its fragments) or four (as many: an expert's pairs decode them once) + for (int n : { 1, 3, 4 }) { test_cases.emplace_back(new test_mul_mat_id(GGML_TYPE_IQ3_XXS, GGML_TYPE_F32, 32, 8, true, 2048, n, 4096)); test_cases.emplace_back(new test_mul_mat_id(GGML_TYPE_IQ3_XXS, GGML_TYPE_F32, 32, 8, false, 4096, n, 2048)); } + // and every token routed to each of 8 experts (the 44-layer proxy's routing), so each expert meets every token: the + // ring meets a fragment with 2 to 8 vectors, past its 4 at once at 5 and 8 tokens + for (int n : { 2, 3, 5, 8 }) { + test_cases.emplace_back(new test_mul_mat_id(GGML_TYPE_IQ3_XXS, GGML_TYPE_F32, 8, 8, true, 2048, n, 4096)); + test_cases.emplace_back(new test_mul_mat_id(GGML_TYPE_IQ3_XXS, GGML_TYPE_F32, 8, 8, false, 4096, n, 2048)); + } + // and its routed FFN, gate/up then down on one ids: its FFN on one card (2,048) and on each of two under -sm tensor + // (1,024), with and without the SwiGLU limit + for (int64_t n_ff : { 2048, 1024 }) { + for (int n : { 1, 3 }) { + for (float glu_limit : { 0.0f, 7.0f }) { + test_cases.emplace_back(new test_moe_ffn_chain(GGML_TYPE_IQ3_XXS, 16, 8, 4096, n_ff, n, glu_limit)); + } + } + } + // and the proxy's, every token on each of 8 experts, its gated gate/up and down each meeting 3 or 5 vectors a row + 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 experts' weighted sum after the down: GLM's (4,096 and 8 slots) at a decode, an MTP verify and a batch, and + // other widths and slot counts (gpt-oss's 2,880 and 4, a row past a block) + for (int64_t n_embd : { 4096, 2880 }) { + for (int n_used : { 8, 4, 2 }) { + for (int64_t n : { 1, 3, 8 }) { + test_cases.emplace_back(new test_moe_weighted_sum(n_embd, n_used, n)); + } + } + } for (ggml_type type_a : all_types) { test_cases.emplace_back(new test_mul_mat_id(type_a, GGML_TYPE_F32, 4, 2, false, 64, 16, 3*ggml_blck_size(type_a))); @@ -11785,6 +12227,24 @@ static std::vector> make_test_cases_eval() { test_cases.emplace_back(new test_flash_attn_ext(576, 512, 1, {20, 1}, 1024, 1, true, false, 0, 0, GGML_PREC_F32, GGML_TYPE_Q8_0, GGML_TYPE_Q8_0, {0, 1, 2, 3}, true, true)); test_cases.emplace_back(new test_flash_attn_ext(576, 512, 1, {20, 1}, 1024, 64, true, false, 0, 0, GGML_PREC_F32, GGML_TYPE_Q8_0, GGML_TYPE_Q8_0, {0, 1, 2, 3}, true, true)); + // Sparse mask hint (n_kv_max): supported decode/prefill layouts and dense fallbacks (a cache under the cells the + // queries' gathers read, a head size with no sparse variant); then GLM-5.3's DSA layers, 64 heads on one latent at its + // n_select (2048 + 4 - 1): a decode, an MTP verify of 3 tokens, 2 streams + test_cases.emplace_back(new test_flash_attn_ext(512, 512, 1, { 8, 1}, 4096, 1, true, false, 0, 0, GGML_PREC_F32, GGML_TYPE_F16, GGML_TYPE_F16, {0, 1, 2, 3}, true, false, false, 0, false, 512)); + test_cases.emplace_back(new test_flash_attn_ext(512, 512, 1, { 8, 2}, 4096, 3, true, false, 0, 0, GGML_PREC_F32, GGML_TYPE_F16, GGML_TYPE_F16, {0, 1, 2, 3}, true, false, false, 0, false, 768)); + test_cases.emplace_back(new test_flash_attn_ext(576, 512, 1, {16, 1}, 4096, 1, true, false, 0, 0, GGML_PREC_F32, GGML_TYPE_F16, GGML_TYPE_F16, {0, 1, 2, 3}, true, true, false, 0, false, 512)); + test_cases.emplace_back(new test_flash_attn_ext(576, 512, 1, {16, 2}, 4096, 2, true, false, 0, 0, GGML_PREC_F32, GGML_TYPE_F16, GGML_TYPE_F16, {0, 1, 2, 3}, true, true, false, 0, false, 768)); + test_cases.emplace_back(new test_flash_attn_ext(512, 512, 1, { 8, 1}, 4096, 2, true, false, 0, 0, GGML_PREC_F32, GGML_TYPE_F16, GGML_TYPE_F16, {0, 1, 2, 3}, true, false, false, 0, false, 2304)); + test_cases.emplace_back(new test_flash_attn_ext(256, 256, 1, { 8, 1}, 4096, 1, true, false, 0, 0, GGML_PREC_F32, GGML_TYPE_F16, GGML_TYPE_F16, {0, 1, 2, 3}, true, false, false, 0, false, 512)); + for (int64_t nb : { 1, 3 }) { + test_cases.emplace_back(new test_flash_attn_ext(512, 512, 1, {64, 1}, 8192, nb, true, false, 0, 0, GGML_PREC_F32, GGML_TYPE_F16, GGML_TYPE_F16, {0, 1, 2, 3}, true, true, false, 0, false, 2051)); + } + test_cases.emplace_back(new test_flash_attn_ext(512, 512, 1, {64, 2}, 8192, 1, true, false, 0, 0, GGML_PREC_F32, GGML_TYPE_F16, GGML_TYPE_F16, {0, 1, 2, 3}, true, true, false, 0, false, 2051)); + // rows past one round of the mask scan (32768 columns): a decode and a verify on one card of two under -sm tensor + for (int64_t nb : { 1, 3 }) { + test_cases.emplace_back(new test_flash_attn_ext(512, 512, 1, {32, 1}, 40960, nb, true, false, 0, 0, GGML_PREC_F32, GGML_TYPE_F16, GGML_TYPE_F16, {0, 1, 2, 3}, true, true, false, 0, false, 2080)); + } + // more V-is-sub-view-of-K cases: other head shapes, and full views with equal head sizes test_cases.emplace_back(new test_flash_attn_ext(320, 256, 1, {32, 1}, 512, 1, true, false, 0, 0, GGML_PREC_F32, GGML_TYPE_F16, GGML_TYPE_F16, {0, 1, 2, 3}, true, true)); test_cases.emplace_back(new test_flash_attn_ext(192, 128, 4, {8, 1}, 512, 1, true, false, 0, 0, GGML_PREC_F32, GGML_TYPE_F16, GGML_TYPE_F16, {0, 1, 2, 3}, true, true)); @@ -11991,6 +12451,12 @@ static std::vector> make_test_cases_eval() { test_cases.emplace_back(new test_mul_mat_pair(GGML_TYPE_BF16, 48, n, 5120, views)); } } + // glm5next's ssm_f_a and ssm_g_a (128 x 4096 bf16), which mmf tiles into fewer blocks than an SM count + for (ggml_type type : {GGML_TYPE_F16, GGML_TYPE_BF16}) { + for (int64_t n : {2, 3, 8}) { + test_cases.emplace_back(new test_mul_mat_pair(type, 128, n, 4096, 0)); + } + } for (auto gate : {GATING_FUNC_SOFTMAX, GATING_FUNC_SIGMOID, GATING_FUNC_SOFTMAX_WEIGHT, GATING_FUNC_SQRT_SOFTPLUS}) { for (bool with_norm : {false, true}) { @@ -12045,6 +12511,10 @@ static std::vector> make_test_cases_eval() { test_cases.emplace_back(new test_gated_delta_net(GGML_TYPE_F32, 4, 64, 4, 2, 1, false, true, 1, false, -1, true)); test_cases.emplace_back(new test_gated_delta_net(GGML_TYPE_F32, 4, 16, 4, 2, 1, false, true, 1, false, -1, true)); test_cases.emplace_back(new test_gated_delta_net(GGML_TYPE_F32, 4, 64, 4, 2, 1, true, true, 1, false, -1, true)); + // from 32 tokens a sequence CUDA computes KDA's decay once per token, head and channel ahead of the recurrence: + // with snapshot slots, two sequences and GQA, and with the rows-indexed state + test_cases.emplace_back(new test_gated_delta_net(GGML_TYPE_F32, 4, 128, 64, 2, 2, false, true, 4, false, -1, true)); + test_cases.emplace_back(new test_gated_delta_net(GGML_TYPE_F32, 4, 128, 32, 2, 1, false, true, 1, true, -1, true)); // chunked path: multi-chunk and non-multiple-of-chunk-size (chunk_size=64 GDN, 16 KDA) test_cases.emplace_back(new test_gated_delta_net(GGML_TYPE_F32, 4, 64, 64, 1)); test_cases.emplace_back(new test_gated_delta_net(GGML_TYPE_F32, 4, 64, 127, 1)); @@ -12088,6 +12558,15 @@ static std::vector> make_test_cases_eval() { // the served Bonsai 2 27B prefill ubatch: 16 k-heads, 48 v-heads, raw gates, qwen35 views, MTP n_max 2 (K = 3) test_cases.emplace_back(new test_gated_delta_net(GGML_TYPE_F32, 16, 128, 512, 1, 3, false, false, 3, false, -1, true, true)); test_cases.emplace_back(new test_gated_delta_net(GGML_TYPE_F32, 16, 128, 300, 2, 3, false, false, 3, false, -1, true, true)); + // KDA on the chunked path (its gate per channel): activated and raw gates, the partial last chunk over two + // sequences, GQA, permuted q/k/v, and the snapshot tail from the chunked state at K 4 and 9 + test_cases.emplace_back(new test_gated_delta_net(GGML_TYPE_F32, 8, 128, 128, 1, 1, false, true)); + test_cases.emplace_back(new test_gated_delta_net(GGML_TYPE_F32, 4, 128, 143, 2, 1, false, true)); + test_cases.emplace_back(new test_gated_delta_net(GGML_TYPE_F32, 16, 128, 256, 1, 2, false, true)); + test_cases.emplace_back(new test_gated_delta_net(GGML_TYPE_F32, 4, 128, 256, 2, 1, true, true)); + 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)); // 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)); @@ -12156,6 +12635,10 @@ static std::vector> make_test_cases_eval() { } } test_cases.emplace_back(new test_gated_delta_net_cache_fusion(GGML_TYPE_F32, 16, 128, 1, 1, 1, 1, true, GGML_TYPE_F32, 0, 0, -1, true)); + // KDA prefill on the chunked path: the f32 cache fused, an f16 one behind its kept cpy, and the gather it keeps + test_cases.emplace_back(new test_gated_delta_net_cache_fusion(GGML_TYPE_F32, 16, 128, 256, 1, 3, 1, true, GGML_TYPE_F32, 0, 0, -1, true)); + test_cases.emplace_back(new test_gated_delta_net_cache_fusion(GGML_TYPE_F32, 16, 128, 200, 2, 3, 1, true, GGML_TYPE_F16, 0, 0, -1, true)); + test_cases.emplace_back(new test_gated_delta_net_cache_fusion(GGML_TYPE_F32, 16, 128, 256, 1, 3, 1, true, GGML_TYPE_F32, 3, 1, 4, true)); // K > 1: output keeps the last min(n_tokens, K) per-token snapshots, ordered most-recent-first // (slot 0 = final state, slot s = state s tokens back). @@ -12214,6 +12697,19 @@ static std::vector> make_test_cases_eval() { test_cases.emplace_back(new test_lightning_indexer_view(type_K, 0, 1)); } + // keys read through rows of a pooled cache: a decode, an MTP verify and a prefill ubatch, 1 and 2 streams, a count + // off the kernel's blocks + for (ggml_type type_K : {GGML_TYPE_F16, GGML_TYPE_F32}) { + for (int64_t nh : { 32, 64 }) { + for (int64_t nb : { 1, 3, 64 }) { + for (int64_t ns : { 1, 2 }) { + test_cases.emplace_back(new test_lightning_indexer_rows(nh, 256, 1027, nb, ns, type_K)); + } + } + } + test_cases.emplace_back(new test_lightning_indexer_rows(32, 77, 311, 3, 1, type_K)); + } + return test_cases; } #ifdef _MSC_VER @@ -12407,6 +12903,23 @@ static std::vector> make_test_cases_perf() { test_cases.emplace_back(new test_mul_mat(GGML_TYPE_BF16, GGML_TYPE_F32, 4096, n, 1536, {1, 1}, {1, 1})); test_cases.emplace_back(new test_mul_mat(GGML_TYPE_F32, GGML_TYPE_F32, 32, n, 4096, {1, 1}, {1, 1})); } + // float weights mmf tiles into fewer blocks than an SM count, which CUDA runs on the vector kernel at 2-8 columns and + // rows x cols <= 1024 (GGML_CUDA_MMVF_UNDERFILLED_LEGACY=1: mmf): rows on both sides of that, up to the SM counts + // of an RTX 5070 Ti and 5080 times 32 + for (ggml_type type : {GGML_TYPE_F16, GGML_TYPE_BF16}) { + for (int64_t m : {64, 128, 256, 512, 1024, 2048}) { + for (int64_t k : {1024, 4096, 8192}) { + for (int64_t n : {2, 3, 4, 8}) { + test_cases.emplace_back(new test_mul_mat(type, GGML_TYPE_F32, m, n, k, {1, 1}, {1, 1})); + } + } + } + } + for (int64_t m : {64, 1024}) { + for (int64_t n : {4, 8}) { + test_cases.emplace_back(new test_mul_mat(GGML_TYPE_F32, GGML_TYPE_F32, m, n, 4096, {1, 1}, {1, 1})); + } + } // FWHT tests test_cases.emplace_back(new test_mul_mat_hadamard(GGML_TYPE_F32, GGML_TYPE_F32, 128, 1, 128)); @@ -12482,10 +12995,18 @@ static std::vector> make_test_cases_perf() { test_cases.emplace_back(new test_mul_mat_id_fusion(type_a, GGML_TYPE_F32, 32, 8, false, 2048, bs, 4096, 1)); } } + // glm5next's real non-expert shapes, at one row and at an MTP verify's width. The candidate 4-bit formats run + // beside Q8_0 because the 8.9 GB of non-expert weight is 70 % of what a token reads, and a format that cuts bytes + // 47 % while losing 30 % of achieved bandwidth is not a 47 % win. 4096 x 154880 is the output head, read every + // token and what an MTP draft pass is bandwidth-dominated by. At bs=1 these take the vector kernel (mmvq), where + // NVFP4's and MXFP4's block-scaled MMA cannot apply; bs=3 is where it can, so the ranking may differ between them. for (int bs : {1, 3}) { - test_cases.emplace_back(new test_mul_mat(GGML_TYPE_Q8_0, GGML_TYPE_F32, 8192, bs, 4096, {1, 1}, {1, 1})); - test_cases.emplace_back(new test_mul_mat(GGML_TYPE_Q8_0, GGML_TYPE_F32, 4096, bs, 8192, {1, 1}, {1, 1})); - test_cases.emplace_back(new test_mul_mat(GGML_TYPE_Q8_0, GGML_TYPE_F32, 4096, bs, 16384, {1, 1}, {1, 1})); + for (ggml_type type_a : {GGML_TYPE_Q8_0, GGML_TYPE_Q4_K, GGML_TYPE_IQ4_XS, GGML_TYPE_MXFP4, GGML_TYPE_NVFP4}) { + test_cases.emplace_back(new test_mul_mat(type_a, GGML_TYPE_F32, 8192, bs, 4096, {1, 1}, {1, 1})); + test_cases.emplace_back(new test_mul_mat(type_a, GGML_TYPE_F32, 4096, bs, 8192, {1, 1}, {1, 1})); + test_cases.emplace_back(new test_mul_mat(type_a, GGML_TYPE_F32, 4096, bs, 16384, {1, 1}, {1, 1})); + test_cases.emplace_back(new test_mul_mat(type_a, GGML_TYPE_F32, 154880, bs, 4096, {1, 1}, {1, 1})); + } } for (int K : {3, 5}) { @@ -12510,6 +13031,20 @@ static std::vector> make_test_cases_perf() { test_cases.emplace_back(new test_flash_attn_ext(64, 64, 8, {8, 1}, 7680, 1, true, false, 0, 0, GGML_PREC_F32, GGML_TYPE_Q8_0, GGML_TYPE_Q8_0)); test_cases.emplace_back(new test_flash_attn_ext(64, 64, 8, {8, 1}, 7680, 512, true, false, 0, 0, GGML_PREC_F32, GGML_TYPE_Q8_0, GGML_TYPE_Q8_0)); + // sparse decode at long context: the whole cache under the mask (n_kv_max 0) and its 2048 live cells alone; then + // GLM-5.3's DSA shape on one card of two under -sm tensor (32 heads on the latent) at its bound (2048 + 4 - 1, padded + // to 2080), where the cache at 8K, 16K and 32K cached tokens and a batch from a decode to a prefill ubatch set when + // the gather beats the dense pass (GGML_CUDA_FATTN_SPARSE_LEGACY=1 runs the same cases dense) + test_cases.emplace_back(new test_flash_attn_ext(512, 512, 1, { 8, 1}, 49152, 1, true, false, 0, 0, GGML_PREC_F32, GGML_TYPE_F16, GGML_TYPE_F16, {0, 1, 2, 3}, true, false, false, 0, false, 0)); + test_cases.emplace_back(new test_flash_attn_ext(512, 512, 1, { 8, 1}, 49152, 1, true, false, 0, 0, GGML_PREC_F32, GGML_TYPE_F16, GGML_TYPE_F16, {0, 1, 2, 3}, true, false, false, 0, false, 2048)); + test_cases.emplace_back(new test_flash_attn_ext(576, 512, 1, {16, 1}, 49152, 1, true, false, 0, 0, GGML_PREC_F32, GGML_TYPE_F16, GGML_TYPE_F16, {0, 1, 2, 3}, true, true, false, 0, false, 0)); + test_cases.emplace_back(new test_flash_attn_ext(576, 512, 1, {16, 1}, 49152, 1, true, false, 0, 0, GGML_PREC_F32, GGML_TYPE_F16, GGML_TYPE_F16, {0, 1, 2, 3}, true, true, false, 0, false, 2048)); + for (int64_t kv : { 8448, 16640, 33280 }) { + for (int64_t nb : { 1, 2, 3, 4, 8, 16, 64, 512 }) { + test_cases.emplace_back(new test_flash_attn_ext(512, 512, 1, {32, 1}, kv, nb, true, false, 0, 0, GGML_PREC_F32, GGML_TYPE_F16, GGML_TYPE_F16, {0, 1, 2, 3}, true, true, false, 0, false, 2080)); + } + } + // q8_0 KV cases with long context (decode and prompt) test_cases.emplace_back(new test_flash_attn_ext(256, 256, 2, {16, 1}, 128, 1, true, false, 0, 0, GGML_PREC_F32, GGML_TYPE_Q8_0, GGML_TYPE_Q8_0)); test_cases.emplace_back(new test_flash_attn_ext(256, 256, 2, {16, 1}, 512, 1, true, false, 0, 0, GGML_PREC_F32, GGML_TYPE_Q8_0, GGML_TYPE_Q8_0)); @@ -12709,6 +13244,10 @@ static std::vector> make_test_cases_perf() { } } } + // GLM-5.3-Flash's pooled indexer at 32K cached tokens: 8258 pools read by rows from the f16 cache, decode and a 3-token verify + for (int bs : { 1, 3 }) { + test_cases.emplace_back(new test_lightning_indexer_rows(32, 8258, 33040, bs, 1, GGML_TYPE_F16)); + } return test_cases; } diff --git a/tests/test-kpool-can-reuse.cpp b/tests/test-kpool-can-reuse.cpp index 4aef882fc263..052c21e6a4b7 100644 --- a/tests/test-kpool-can-reuse.cpp +++ b/tests/test-kpool-can-reuse.cpp @@ -35,11 +35,13 @@ static constexpr int64_t N_TPS = 1; static constexpr int64_t N_PS = 1; static constexpr int64_t N_POOLS = N_KV/KPOOL + 2*N_PS; // 66, llama_kpool_n_pools static constexpr int64_t N_NEW_MAX = N_TPS/KPOOL + N_PS; // 1, decode, not rebuilding +static constexpr int64_t N_DUMP = 64/KPOOL; // 16, llama_kpool_select_k: the top-k's dump pools struct case_tensors { ggml_tensor * k_idxs = nullptr; ggml_tensor * pool_cells = nullptr; ggml_tensor * pool_bias = nullptr; + ggml_tensor * pool_dump = nullptr; ggml_tensor * pool_bias_f16 = nullptr; ggml_tensor * sel_mask = nullptr; ggml_tensor * cand_mask = nullptr; @@ -54,10 +56,11 @@ static case_tensors make_tensors(ggml_context * ctx, bool scoring, int64_t n_new if (!scoring) { return t; } - t.pool_cells = ggml_new_tensor_2d(ctx, GGML_TYPE_I32, KPOOL*N_POOLS, N_STREAM); + t.pool_cells = ggml_new_tensor_2d(ctx, GGML_TYPE_I32, KPOOL*(N_POOLS + N_DUMP), N_STREAM); t.pool_bias = ggml_new_tensor_3d(ctx, GGML_TYPE_F32, N_POOLS, N_TPS, N_STREAM); + t.pool_dump = ggml_new_tensor_3d(ctx, GGML_TYPE_F32, N_DUMP, N_TPS, N_STREAM); t.pool_bias_f16 = ggml_new_tensor_4d(ctx, GGML_TYPE_F16, N_POOLS, N_TPS, 1, N_STREAM); - t.sel_mask = ggml_new_tensor_4d(ctx, GGML_TYPE_F16, N_KV, N_TPS, 1, N_STREAM); + t.sel_mask = ggml_new_tensor_4d(ctx, GGML_TYPE_F16, N_KV + KPOOL*N_DUMP, N_TPS, 1, N_STREAM); t.cand_mask = ggml_new_tensor_4d(ctx, GGML_TYPE_F16, N_KV, N_TPS, 1, N_STREAM); t.pool_reps = ggml_new_tensor_2d(ctx, GGML_TYPE_I32, N_POOLS, N_STREAM); t.new_pool_cells = ggml_new_tensor_2d(ctx, GGML_TYPE_I32, KPOOL*n_new_max, N_STREAM); @@ -70,6 +73,7 @@ static llm_graph_input_kpool make_inp(const case_tensors & t, bool scoring, int6 inp.k_idxs = t.k_idxs; inp.pool_cells = t.pool_cells; inp.pool_bias = t.pool_bias; + inp.pool_dump = t.pool_dump; inp.pool_bias_f16 = t.pool_bias_f16; inp.sel_mask = t.sel_mask; inp.cand_mask = t.cand_mask; @@ -90,6 +94,7 @@ static llm_graph_input_kpool_dims base_dims() { d.n_tps = N_TPS; d.n_ps = N_PS; d.n_pools = N_POOLS; + d.n_dump = N_DUMP; d.n_new_max = N_NEW_MAX; d.rebuild = false; d.scoring = true; @@ -161,6 +166,10 @@ int main() { d.scoring = false; CHECK(!llm_graph_input_kpool::shapes_match(d, inp), "scoring flip must refuse"); + d = base_dims(); + d.n_dump = N_DUMP - 1; + CHECK(!llm_graph_input_kpool::shapes_match(d, inp), "a dump pool count change 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 new file mode 100644 index 000000000000..f6c33feb0a79 --- /dev/null +++ b/tests/test-kpool-input.cpp @@ -0,0 +1,793 @@ +// The pooled-indexer inputs built from a sequence's view (llama_kpool_views) against the maps built from the cells. +// +// llama_kpool_set_input with views keeps, for each sequence, its positions, pool table and completeness, and updates them +// from the cells llama_kv_cells_log names; with views == nullptr it builds every map from the cells each call. The two +// must write the same bytes into every input tensor. +// +// Random cell states the way a server makes them (as test-kv-mask.cpp): sequences in one unified pool or one stream each, +// tokens appended, drafts rolled back, conversations ended, prefixes shared, cells removed and kept, saved and restored +// (llama_kv_cache::prepare), positions shifted and divided with a rebuild after, two cells at one position, positions +// sparse enough to span more pools than a map holds, prompts large enough to pass the log's cap, the cells copied over. +// One llama_kpool_views lives through a whole round, so the view follows many small changes; after every few operations +// the inputs of a decode or verify ubatch are built both ways and compared byte for byte. +// +// Two controls: a view whose changes were taken by someone else must be caught by the comparison, and the log of the cells +// is checked against a snapshot of the cells on its own. + +#include "../src/llama-batch.h" +#include "../src/llama-kv-cache-kpool.h" +#include "../src/llama-kv-cells.h" + +#include "ggml.h" +#include "ggml-backend.h" + +#include +#include +#include +#include +#include +#include +#include + +static int n_fail = 0; + +#define CHECK(cond, ...) do { if (!(cond)) { ++n_fail; if (n_fail <= 20) { fprintf(stderr, "FAIL %s:%d: ", __FILE__, __LINE__); fprintf(stderr, __VA_ARGS__); fprintf(stderr, "\n"); } } } while (0) + +struct pool { + std::vector v_cells; // one per stream + std::vector seq_to_stream; + std::vector head; // per stream + uint32_t n_seq; + std::mt19937 & rng; + + // prepare() between its apply and its restore: the cells saved and the ubatch placed, the restore still to come + struct pending_t { + bool on = false; + llama_seq_id s = 0; + std::vector idxs; + llama_kv_cells saved; + } pending; + + pool(uint32_t size, uint32_t n_seq, bool unified, std::mt19937 & rng) : n_seq(n_seq), rng(rng) { + const uint32_t n_stream = unified ? 1 : n_seq; + v_cells.resize(n_stream); + for (auto & cells : v_cells) { + cells.resize(size); + } + head.assign(n_stream, 0); + for (uint32_t s = 0; s < n_seq; ++s) { + seq_to_stream.push_back(unified ? 0 : s); + } + } + + llama_kv_cells & cells_of(llama_seq_id s) { return v_cells[seq_to_stream[s]]; } + + llama_pos p_next(llama_seq_id s) { return cells_of(s).seq_pos_max(s) + 1; } + + // the next empty cell from the stream's head on, wrapping (llama_kv_cache::find_slot, not contiguous) + int64_t find_empty(llama_seq_id s) { + auto & cells = cells_of(s); + uint32_t & h = head[seq_to_stream[s]]; + for (uint32_t n = 0; n < cells.size(); ++n) { + const uint32_t i = (h + n) % cells.size(); + if (cells.is_empty(i)) { + h = (i + 1) % cells.size(); + return i; + } + } + return -1; + } + + // n tokens of the sequence at the positions after its last, as a ubatch is applied; the positions placed + std::vector append(llama_seq_id s, uint32_t n) { + std::vector placed; + auto & cells = cells_of(s); + for (uint32_t k = 0; k < n; ++k) { + const int64_t i = find_empty(s); + if (i < 0) { + break; + } + const llama_pos p = p_next(s); + cells.pos_set(i, p); + cells.seq_add(i, s); + placed.push_back(p); + } + return placed; + } + + // drop the sequence's positions from p0 on (a rejected draft: llama_kv_cache::seq_rm) + void rollback(llama_seq_id s, llama_pos p0, llama_pos p1 = INT32_MAX) { + auto & cells = cells_of(s); + uint32_t new_head = cells.size(); + for (uint32_t i = 0; i < cells.size(); ++i) { + if (cells.pos_in(i, p0, p1) && cells.seq_has(i, s) && cells.seq_rm(i, s)) { + new_head = std::min(new_head, i); + } + } + uint32_t & h = head[seq_to_stream[s]]; + if (new_head < h) { + h = new_head; + } + } + + // a position of the sequence taken a second time, by another cell + void duplicate(llama_seq_id s) { + auto & cells = cells_of(s); + const llama_pos p = cells.seq_pos_max(s); + const int64_t i = find_empty(s); + if (p >= 0 && i >= 0) { + cells.pos_set(i, rng() % 2 ? p : (llama_pos) (rng() % (p + 1))); + cells.seq_add(i, s); + } + } + + // true when a position moved: the pooled keys are stale, the next inputs re-emit every pool + bool random_op() { + const llama_seq_id s = rng() % n_seq; + auto & cells = cells_of(s); + switch (rng() % 16) { + case 0: case 1: case 2: case 3: // a prompt chunk or a few decoded tokens + append(s, 1 + rng() % (rng() % 8 == 0 ? 300 : 6)); + break; + case 4: case 5: case 6: // a verify's rejected draft + if (p_next(s) > 0) { + rollback(s, std::max(0, p_next(s) - 1 - (llama_pos) (rng() % 4))); + } + break; + case 7: // the conversation ends; the next one starts at 0 in whatever cells are free + if (rng() % 3 == 0) { + rollback(s, 0); + } + break; + case 8: // a range out of the middle + if (p_next(s) > 4) { + const llama_pos p0 = rng() % p_next(s); + rollback(s, p0, p0 + 1 + rng() % 12); + } + break; + case 9: { // a second sequence of the same stream takes this one's prefix (a cached prompt: seq_cp) + const llama_seq_id t = rng() % n_seq; + if (t != s && seq_to_stream[t] == seq_to_stream[s] && cells.seq_pos_max(t) < 0 && p_next(s) > 1) { + const llama_pos p1 = 1 + rng() % p_next(s); + for (uint32_t i = 0; i < cells.size(); ++i) { + if (cells.pos_in(i, 0, p1) && cells.seq_has(i, s) && !cells.seq_has(i, t)) { + cells.seq_add(i, t); + } + } + } + } break; + case 10: { // one cell removed whatever holds it + const uint32_t i = rng() % cells.size(); + if (!cells.is_empty(i)) { + cells.rm(i); + } + } break; + case 11: { // prepare(): cells saved, a ubatch placed; the cells restored by a later operation + if (pending.on) { + cells_of(pending.s).set(pending.idxs, pending.saved); + pending.on = false; + break; + } + std::vector idxs; + const uint32_t n = 1 + rng() % 8; + for (uint32_t k = 0; k < n; ++k) { + idxs.push_back(rng() % cells.size()); + } + std::sort(idxs.begin(), idxs.end()); + idxs.erase(std::unique(idxs.begin(), idxs.end()), idxs.end()); + pending.saved = cells.cp(idxs); + for (const uint32_t i : idxs) { + if (!cells.is_empty(i)) { + cells.rm(i); + } + cells.pos_set(i, p_next(s) + (llama_pos) (rng() % 4)); + cells.seq_add(i, s); + } + pending.on = true; + pending.s = s; + pending.idxs = idxs; + } break; + case 12: { // positions shifted (a context shift) or divided, applied + const bool div = rng() % 4 == 0; + const llama_pos d = (rng() % 2 ? 1 : -1) * (llama_pos) (1 + rng() % 64); + const llama_pos p0 = p_next(s) > 0 ? rng() % p_next(s) : 0; + for (uint32_t i = 0; i < cells.size(); ++i) { + if (cells.pos_in(i, p0, INT32_MAX) && cells.seq_has(i, s) && cells.seq_count(i) == 1) { + if (div) { + cells.pos_div(i, 2); + } else { + cells.pos_add(i, d); + } + } + } + cells.reset_shift(); + return true; + } + case 13: // seq_keep of one sequence + if (rng() % 8 == 0) { + for (uint32_t i = 0; i < cells.size(); ++i) { + cells.seq_keep(i, s); + } + } + break; + case 14: // a position taken twice, or cells sparse enough to span more pools than a map holds + if (rng() % 2) { + duplicate(s); + } else if (rng() % 4 == 0) { + const int64_t i = find_empty(s); + if (i >= 0) { + cells.pos_set(i, p_next(s) + 1000 + rng() % 100000); + cells.seq_add(i, s); + } + } else if (rng() % 2) { // a position with no sequence yet + const int64_t i = find_empty(s); + if (i >= 0) { + cells.pos_set(i, rng() % 512); + } + } + break; + case 15: // the cells all cleared, or copied over + if (rng() % 16 == 0) { + if (rng() % 2) { + cells.reset(); + } else { + llama_kv_cells copy; + copy.resize(cells.size()); + cells = copy; + } + } + break; + } + return false; + } +}; + +// the cells of one round, as a snapshot: the log must name every cell that differs from it +struct snap { + std::vector pos; + std::vector seqs; + + void take(const llama_kv_cells & cells, uint32_t n_seq) { + pos.assign(cells.size(), -1); + seqs.assign(cells.size(), 0); + for (uint32_t i = 0; i < cells.size(); ++i) { + if (cells.is_empty(i)) { + continue; + } + pos[i] = cells.pos_get(i); + for (uint32_t s = 0; s < n_seq; ++s) { + if (cells.seq_has(i, s)) { + seqs[i] |= (uint64_t) 1 << s; + } + } + } + } +}; + +static void check_log(llama_kv_cells & cells, snap & was, uint32_t n_seq) { + std::vector list; + const bool listed = cells.changes().take(list); + + snap now; + now.take(cells, n_seq); + + if (listed) { + const std::set in_list(list.begin(), list.end()); + for (uint32_t i = 0; i < cells.size(); ++i) { + if (was.pos[i] != now.pos[i] || was.seqs[i] != now.seqs[i]) { + CHECK(in_list.count(i), "cell %u changed (position %d -> %d) and is not in the log", i, was.pos[i], now.pos[i]); + } + } + } + + was = std::move(now); +} + +// the counts against their definition +static void check_counts(const llama_kv_cells & cells, uint32_t n_seq) { + for (uint32_t s = 0; s < n_seq; ++s) { + std::set distinct; + uint32_t n = 0; + for (uint32_t i = 0; i < cells.size(); ++i) { + if (!cells.is_empty(i) && cells.seq_has(i, s)) { + distinct.insert(cells.pos_get(i)); + n++; + } + } + CHECK(cells.seq_cell_count(s) == n, "seq %u: %u cells, counted %u", s, cells.seq_cell_count(s), n); + CHECK(cells.seq_pos_distinct(s) == distinct.size(), "seq %u: %u positions, counted %zu", s, cells.seq_pos_distinct(s), distinct.size()); + } +} + +// the input tensors of one build, in host memory +struct maps { + ggml_context * ctx = nullptr; + ggml_backend_buffer_t buf = nullptr; + + ggml_tensor * pool_cells = nullptr; + ggml_tensor * pool_bias = nullptr; + ggml_tensor * sel_mask = nullptr; + ggml_tensor * cand_mask = nullptr; + ggml_tensor * pool_reps = nullptr; + ggml_tensor * new_pool_cells = nullptr; + ggml_tensor * new_pool_reps = nullptr; + + // n_dump: the top-k's dump pools past the n_pools real ones, with their columns in sel_mask (llm_graph_input_kpool) + maps(int64_t n_kv, int64_t n_ns, int64_t n_tps, int64_t n_ps, uint32_t r, bool kcache, bool f16, bool rebuild, int64_t n_dump, int fill) { + ggml_init_params ip = { 32*ggml_tensor_overhead(), nullptr, true }; + ctx = ggml_init(ip); + + const int64_t n_pools = llama_kpool_n_pools(n_kv, r, n_ps); + const int64_t n_new_max = rebuild ? n_pools : n_tps/r + n_ps; + const ggml_type tm = f16 ? GGML_TYPE_F16 : GGML_TYPE_F32; + + pool_cells = ggml_new_tensor_2d(ctx, GGML_TYPE_I32, r*(n_pools + n_dump), n_ns); + pool_bias = ggml_new_tensor_3d(ctx, GGML_TYPE_F32, n_pools, n_tps, n_ns); + sel_mask = ggml_new_tensor_4d(ctx, tm, n_kv + r*n_dump, n_tps, 1, n_ns); + cand_mask = ggml_new_tensor_4d(ctx, tm, n_kv, n_tps, 1, n_ns); + + if (kcache) { + pool_reps = ggml_new_tensor_2d(ctx, GGML_TYPE_I32, n_pools, n_ns); + new_pool_cells = ggml_new_tensor_2d(ctx, GGML_TYPE_I32, r*n_new_max, n_ns); + new_pool_reps = ggml_new_tensor_1d(ctx, GGML_TYPE_I64, n_new_max*n_ns); + } + + buf = ggml_backend_alloc_ctx_tensors_from_buft(ctx, ggml_backend_cpu_buffer_type()); + // a byte the build does not write stays different between the two builds + ggml_backend_buffer_clear(buf, fill); + } + + ~maps() { + ggml_backend_buffer_free(buf); + ggml_free(ctx); + } + + maps(const maps &) = delete; + + std::vector all() const { + return { pool_cells, pool_bias, sel_mask, cand_mask, pool_reps, new_pool_cells, new_pool_reps }; + } +}; + +// the ubatch: seq_id[i][0] and pos[i] are read, and seq_id_unq and n_seqs_unq +struct ubatch_data { + std::vector pos; + std::vector seq; + std::vector seq_ptr; + std::vector unq; + llama_ubatch ub = {}; + + void add(llama_seq_id s, llama_pos p) { + pos.push_back(p); + seq.push_back(s); + } + + const llama_ubatch * get() { + seq_ptr.clear(); + unq.clear(); + for (auto & s : seq) { + seq_ptr.push_back(&s); + if (std::find(unq.begin(), unq.end(), s) == unq.end()) { + unq.push_back(s); + } + } + ub.n_tokens = pos.size(); + ub.n_seq_tokens = pos.size(); + ub.n_seqs = unq.size(); + ub.n_seqs_unq = unq.size(); + ub.n_pos = 1; + ub.pos = pos.data(); + ub.seq_id = seq_ptr.data(); + ub.seq_id_unq = unq.data(); + return &ub; + } +}; + +static const char * const names[] = { "pool_cells", "pool_bias", "sel_mask", "cand_mask", "pool_reps", "new_pool_cells", "new_pool_reps" }; + +// every byte of every tensor; the first one that differs is named +static bool same(const maps & a, const maps & b, char * what, size_t n_what) { + const auto ta = a.all(); + const auto tb = b.all(); + for (size_t t = 0; t < ta.size(); ++t) { + if (ta[t] == nullptr) { + continue; + } + const size_t n = ggml_nbytes(ta[t]); + const uint8_t * pa = (const uint8_t *) ta[t]->data; + const uint8_t * pb = (const uint8_t *) tb[t]->data; + for (size_t k = 0; k < n; ++k) { + if (pa[k] != pb[k]) { + snprintf(what, n_what, "%s byte %zu of %zu: %02x, reference %02x", names[t], k, n, pa[k], pb[k]); + return false; + } + } + } + return true; +} + +struct config { + uint32_t r; + bool unified; + bool kcache; + bool f16; + bool small_kv = true; // now and then an n_kv below the used cells, which the views must decline +}; + +enum control { + CONTROL_NONE, + CONTROL_STEAL, // someone else takes the changes the view is waiting for + CONTROL_CORRUPT, // a byte of the views' output is flipped +}; + +// the rows the pooled keys are written to: all different, because set_rows writes them from several threads, and a slot +// that names the last cell of a block must be the one that computes that cell's pool, because that cell's pooled key is read +static bool write_rows_safe(const maps & m, const pool & w, int64_t n_ns, uint32_t r, char * what, size_t n_what) { + if (m.new_pool_reps == nullptr) { + return true; + } + + const int64_t n_slots = m.new_pool_reps->ne[0]; + const int64_t n_new_max = n_slots/n_ns; + const int64_t kv_size = w.v_cells[0].size(); + const int64_t * rows = (const int64_t *) m.new_pool_reps->data; + const int32_t * cells = (const int32_t *) m.new_pool_cells->data; + + // r 1 has no spare cell (every cell is a pool), so an unused slot repeats a pool and writes the same key again + for (int64_t i = 0; r > 1 && i < n_slots; ++i) { + for (int64_t j = i + 1; j < n_slots; ++j) { + if (rows[i] == rows[j]) { + const llama_kv_cells & c = w.v_cells[0]; + snprintf(what, n_what, "new_pool_reps: row %lld written by slots %lld and %lld of %lld (cell %s, pos %u, n_ns %lld) rows %lld %lld %lld %lld %lld %lld; used %u/%u", (long long) rows[i], + (long long) i, (long long) j, (long long) n_slots, c.is_empty(rows[i] % kv_size) ? "empty" : "used", (unsigned) (c.is_empty(rows[i] % kv_size) ? 9999 : c.pos_get(rows[i] % kv_size)), (long long) n_ns, + (long long) rows[0], (long long) rows[1], (long long) rows[std::min(2, n_slots - 1)], (long long) rows[std::min(3, n_slots - 1)], + (long long) rows[std::min(4, n_slots - 1)], (long long) rows[std::min(5, n_slots - 1)], c.get_used(), c.size()); + return false; + } + } + } + + for (int64_t i = 0; i < n_slots; ++i) { + const int64_t s = i/n_new_max; + const int64_t cell = rows[i] - s*kv_size; + + if (cell < 0 || cell >= kv_size) { + snprintf(what, n_what, "new_pool_reps[%lld] = %lld: outside stream %lld", (long long) i, (long long) rows[i], (long long) s); + return false; + } + + const llama_kv_cells & c = w.v_cells[s]; + + if (!c.is_empty(cell) && c.pos_get(cell) % (llama_pos) r == (llama_pos) r - 1 && + cells[s*r*n_new_max + (i % n_new_max)*r + r - 1] != (int32_t) cell) { + snprintf(what, n_what, "new_pool_reps[%lld]: cell %lld is the last of its block but the slot computes another pool", (long long) i, (long long) cell); + return false; + } + } + + return true; +} + +// the dump pools name cells n_kv + c, one each, and the real pools' cells are under n_kv, so the slots of a row of the top-k, +// real pools and dump pools together, never name one cell twice (the mask's set_rows writes them from several threads); a +// dump column is masked in every row, and cand_mask has none +static bool dump_safe(const maps & m, int64_t n_kv, uint32_t r, char * what, size_t n_what) { + const int64_t n_ns = m.pool_cells->ne[1]; + const int64_t n_pools = m.pool_bias->ne[0]; + const int64_t n_all = m.pool_cells->ne[0]/r; + const int64_t n_sel = m.sel_mask->ne[0]; + const int64_t n_rows = m.sel_mask->ne[1]; + + if (m.cand_mask->ne[0] != n_kv || n_sel != n_kv + r*(n_all - n_pools)) { + snprintf(what, n_what, "dump: sel_mask %lld columns, cand_mask %lld, for n_kv %lld and %lld dump pools", (long long) n_sel, + (long long) m.cand_mask->ne[0], (long long) n_kv, (long long) (n_all - n_pools)); + return false; + } + + for (int64_t s = 0; s < n_ns; ++s) { + const int32_t * pc = (const int32_t *) m.pool_cells->data + s*r*n_all; + + for (int64_t c = 0; c < r*n_all; ++c) { + const bool ok = c < r*n_pools ? pc[c] >= 0 && pc[c] < n_kv : pc[c] == n_kv + (c - r*n_pools); + if (!ok) { + snprintf(what, n_what, "dump: stream %lld, pool_cells[%lld] = %d (a %s pool; %lld real pools, n_kv %lld)", (long long) s, (long long) c, + pc[c], c < r*n_pools ? "real" : "dump", (long long) n_pools, (long long) n_kv); + return false; + } + } + + for (int64_t ii = 0; ii < n_rows; ++ii) { + for (int64_t j = n_kv; j < n_sel; ++j) { + const int64_t k = (s*n_rows + ii)*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]; + if (!(std::isinf(x) && x < 0)) { + snprintf(what, n_what, "dump: stream %lld, sel_mask row %lld dump column %lld = %g", (long long) s, (long long) ii, (long long) j, x); + return false; + } + } + } + } + + return true; +} + +// 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) { + const uint32_t n_ns = cfg.unified ? 1 : w.n_seq; + const uint32_t n_tps = rng() % 2 ? 1 : 2 + rng() % 4; + const bool two_ps = cfg.unified && w.n_seq > 1 && rng() % 6 == 0; + + ubatch_data ud; + + if (cfg.unified) { + // one sequence, or two sharing the ubatch (the maps are then built from the cells) + for (uint32_t k = 0; k < (two_ps ? 2u : 1u); ++k) { + const llama_seq_id s = two_ps ? k : (llama_seq_id) (rng() % w.n_seq); + for (const llama_pos p : w.append(s, two_ps ? 1 + n_tps/2 : n_tps)) { + ud.add(s, p); + } + } + } else { + // one stream each, the same count of tokens from each + std::vector> placed; + size_t n_min = n_tps; + for (uint32_t s = 0; s < n_ns; ++s) { + placed.push_back(w.append(s, n_tps)); + n_min = std::min(n_min, placed.back().size()); + } + for (uint32_t s = 0; s < n_ns; ++s) { + // a stream that placed more than another keeps only as many, so the streams hold the same count; the ubatch is + // the last of its sequence, nothing lies ahead of it + for (size_t k = 0; k < n_min; ++k) { + ud.add(s, placed[s][k]); + } + if (placed[s].size() > n_min) { + w.rollback(s, placed[s][n_min]); + } + } + } + + if (ud.pos.empty() || ud.pos.size() % n_ns != 0) { + return true; + } + + const llama_ubatch * ub = ud.get(); + + if (ub->n_seqs_unq % n_ns != 0 || (n_ns > 1 && ub->n_seqs_unq != n_ns)) { + return true; + } + + const int64_t n_tokens = ub->n_tokens; + const int64_t n_ps = ub->n_seqs_unq/n_ns; + + // the tokens of a stream must be the stream's own, and one run of positions each (llama-batch.cpp) + if (n_ns > 1) { + for (int64_t i = 0; i < n_tokens; ++i) { + if (ub->seq_id[i][0] != (llama_seq_id) (i/(n_tokens/n_ns))) { + return true; + } + } + } + + uint32_t used_max = 0; + for (const auto & cells : w.v_cells) { + used_max = std::max(used_max, cells.used_max_p1()); + } + + // the cache's n_kv, padded; now and then a bound below the used cells, which the views must decline + int64_t n_kv = std::min(w.v_cells[0].size(), std::max(32, (used_max + 31)/32*32)); + if (cfg.small_kv && rng() % 24 == 0 && used_max > 8) { + n_kv = 1 + rng() % (used_max - 1); + } + + std::vector strm_of(n_ns); + for (uint32_t s = 0; s < n_ns; ++s) { + strm_of[s] = s; + } + + const llama_kpool_cells_fn cells_of = [&](llama_seq_id s) -> const llama_kv_cells & { return w.cells_of(s); }; + + if (ctl == CONTROL_STEAL) { + std::vector tmp; + for (auto & cells : w.v_cells) { + cells.changes().take(tmp); + } + } + + // none, a few, or as many as a top-k over every pool takes + const int64_t n_pools = llama_kpool_n_pools(n_kv, cfg.r, n_ps); + const int64_t n_dump = rng() % 3 == 0 ? 0 : rng() % 2 ? 1 + rng() % 4 : n_pools; + + maps ref(n_kv, n_ns, n_tokens/n_ns, n_ps, cfg.r, cfg.kcache, cfg.f16, rebuild, n_dump, 0xA5); + 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, + 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, + ref.pool_reps, ref.new_pool_cells, ref.new_pool_reps, strm_of.data(), w.v_cells[0].size(), rebuild, ub, cfg.r); + + if (ctl == CONTROL_CORRUPT) { + ((uint8_t *) inc.sel_mask->data)[0] ^= 0x80; + } + + return write_rows_safe(ref, w, n_ns, cfg.r, what, n_what) && write_rows_safe(inc, w, n_ns, cfg.r, what, n_what) && + dump_safe(ref, n_kv, cfg.r, what, n_what) && same(inc, ref, what, n_what); +} + +static void run_round(std::mt19937 & rng, const config & cfg, uint32_t size, uint32_t n_seq, int n_ops, llama_kpool_views::stats & total) { + pool w(size, n_seq, cfg.unified, rng); + llama_kpool_views views; + + bool rebuild = false; + + for (int k = 0; k < n_ops; ++k) { + rebuild |= w.random_op(); + + if (rng() % 3 != 0) { + continue; + } + + char what[256]; + if (!build_both(w, views, cfg, rebuild, rng, CONTROL_NONE, what, sizeof(what))) { + CHECK(false, "r %u, %s, %s: %s", cfg.r, cfg.unified ? "unified" : "streams", cfg.f16 ? "f16" : "f32", what); + } + rebuild = rng() % 8 == 0; + + for (size_t i = 0; i < w.v_cells.size(); ++i) { + check_counts(w.v_cells[i], n_seq); + } + } + + const auto & st = views.get_stats(); + total.n_served += st.n_served; + total.n_rebuilt += st.n_rebuilt; + total.n_direct += st.n_direct; +} + +// the log on its own: a consumer that takes it after every few operations finds every changed cell in it +static void run_log_round(std::mt19937 & rng, uint32_t size, uint32_t n_seq, bool unified, int n_ops) { + pool w(size, n_seq, unified, rng); + + std::vector was(w.v_cells.size()); + for (size_t i = 0; i < w.v_cells.size(); ++i) { + w.v_cells[i].changes().track(size/16); + was[i].take(w.v_cells[i], n_seq); + } + + for (int k = 0; k < n_ops; ++k) { + w.random_op(); + + if (rng() % 3 == 0) { + for (size_t i = 0; i < w.v_cells.size(); ++i) { + check_log(w.v_cells[i], was[i], n_seq); + check_counts(w.v_cells[i], n_seq); + } + } + } +} + +// a decode loop with rollbacks, nothing else: the view must follow it without a rebuild. With moves, a hole in the cells +// lies below the newest token, and every few steps the token is moved into it (a defrag): a cell takes the position another +// cell, of a higher index, leaves in the same changes +static void run_steady(std::mt19937 & rng, bool moves) { + const config cfg = { 4, true, true, true, false }; + + pool w(4096, 1, true, rng); + llama_kpool_views views; + char what[256]; + + for (int k = 0; k < 1200; ++k) { + if (k % 5 == 4) { + w.rollback(0, w.p_next(0) - 1 - rng() % 3); + } + + auto & cells = w.cells_of(0); + + if (moves && k == 200) { + for (uint32_t i = 0; i < cells.size(); ++i) { + if (!cells.is_empty(i) && cells.pos_get(i) == 100) { + cells.rm(i); + } + } + } + + if (moves && k > 200 && k % 25 == 0) { + uint32_t hi = 0; + uint32_t lo = 0; + while (lo < cells.size() && !cells.is_empty(lo)) { + lo++; + } + for (uint32_t i = 0; i < cells.size(); ++i) { + if (!cells.is_empty(i) && cells.pos_get(i) == cells.seq_pos_max(0)) { + hi = i; + } + } + if (lo < hi) { + const llama_pos p = cells.pos_get(hi); + cells.rm(hi); + cells.pos_set(lo, p); + cells.seq_add(lo, 0); + } + } + + if (!build_both(w, views, cfg, false, rng, CONTROL_NONE, what, sizeof(what))) { + CHECK(false, "steady, step %d: %s", k, what); + break; + } + } + + const auto & st = views.get_stats(); + printf("steady%s: %llu served, %llu rebuilt, %llu from the cells\n", moves ? " with moves" : "", (unsigned long long) st.n_served, + (unsigned long long) st.n_rebuilt, (unsigned long long) st.n_direct); + CHECK(st.n_rebuilt == 1, "steady decode rebuilt the view %llu times, once is the first build", (unsigned long long) st.n_rebuilt); + CHECK(st.n_direct == 0, "steady decode built %llu streams from the cells", (unsigned long long) st.n_direct); +} + +// the comparison fails on a difference; and a view that lost changes which left its count of cells off is rebuilt, not served +static void run_controls(std::mt19937 & rng) { + const config cfg = { 4, true, true, true, false }; + + pool w(512, 1, true, rng); + llama_kpool_views views; + char what[256]; + + for (int k = 0; k < 40; ++k) { + CHECK(build_both(w, views, cfg, false, rng, CONTROL_NONE, what, sizeof(what)), "controls, warm-up: %s", what); + } + + CHECK(!build_both(w, views, cfg, false, rng, CONTROL_CORRUPT, what, sizeof(what)), "a flipped byte passed the comparison"); + + const uint64_t n_rebuilt = views.get_stats().n_rebuilt; + + for (int k = 0; k < 20; ++k) { + CHECK(build_both(w, views, cfg, false, rng, CONTROL_STEAL, what, sizeof(what)), "a view that lost its changes was served: %s", what); + } + + CHECK(views.get_stats().n_rebuilt > n_rebuilt, "a view whose changes were taken was not rebuilt"); +} + +int main() { + std::mt19937 rng(20260930); + + llama_kpool_views::stats total; + + const config configs[] = { + { 4, true, true, true }, + { 4, true, true, false }, + { 2, true, false, true }, + { 8, true, true, true }, + { 1, true, true, true }, + { 4, false, true, true }, + { 4, false, false, false }, + }; + + for (int round = 0; round < 420; ++round) { + const config & cfg = configs[round % (sizeof(configs)/sizeof(configs[0]))]; + const uint32_t size = std::vector{ 64, 256, 1024, 2048 }[rng() % 4]; + const uint32_t n_seq = cfg.unified ? 1 + rng() % 3 : 2 + rng() % 2; + run_round(rng, cfg, size, n_seq, 200 + rng() % 300, total); + } + + for (int round = 0; round < 200; ++round) { + run_log_round(rng, std::vector{ 64, 256, 1024, 4096 }[rng() % 4], 1 + rng() % 4, rng() % 3 != 0, rng() % 400); + } + + run_steady(rng, false); + run_steady(rng, true); + run_controls(rng); + + printf("random rounds: %llu streams from views, %llu views rebuilt, %llu streams from the cells\n", + (unsigned long long) total.n_served, (unsigned long long) total.n_rebuilt, (unsigned long long) total.n_direct); + + // the rounds must reach the view path, not only the fallback + 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); + + if (n_fail == 0) { + printf("ok\n"); + } + return n_fail != 0; +} diff --git a/tools/llama-bench/llama-bench.cpp b/tools/llama-bench/llama-bench.cpp index b7b4abe3f853..fa06d5336d2d 100644 --- a/tools/llama-bench/llama-bench.cpp +++ b/tools/llama-bench/llama-bench.cpp @@ -50,7 +50,8 @@ static uint64_t get_time_ns() { // An NVTX range "gen" in the domain "llama-bench" around each repetition's generation, so a profiler can capture the // tokens and none of the depth prefill before them: nsys profile -c nvtx -p gen@llama-bench --capture-range-end=stop // records the first repetition's generation (Ternary Bonsai 2 27B at a depth of 245,760: 16 tokens launch 18,816 kernels, -// the prefill before them 1.19M). The message is a registered string, which nsys matches without NSYS_NVTX_PROFILER_REGISTER_ONLY=0. Without +// the prefill before them 1.19M). "prompt" is the same around the measured prompt pass (-p), not the depth's: a pp4 at a +// depth, an MTP verify's shape, captured alone. The message is a registered string, which nsys matches without NSYS_NVTX_PROFILER_REGISTER_ONLY=0. Without // the NVTX headers (a build without the CUDA toolkit) the range compiles out. struct bench_phase_range { #ifdef LLAMA_BENCH_NVTX @@ -2527,6 +2528,7 @@ int llama_bench(int argc, char ** argv) { fprintf(stderr, "llama-bench: benchmark %d/%zu: prompt run %d/%d\n", params_idx, params_count, i + 1, params.reps); } + const bench_phase_range prompt_range("prompt"); bool res = test_prompt(ctx, t.n_prompt, t.n_batch, t.n_threads); if (!res) { fprintf(stderr, "%s: error: failed to run prompt\n", __func__); diff --git a/tools/perplexity/perplexity.cpp b/tools/perplexity/perplexity.cpp index ba41287d8e3d..6e454036e527 100644 --- a/tools/perplexity/perplexity.cpp +++ b/tools/perplexity/perplexity.cpp @@ -9,6 +9,7 @@ #include #include #include +#include #include #include #include @@ -172,6 +173,21 @@ static void process_logits(std::ostream& out, int n_vocab, const float * logits, out.write((const char *)log_probs.data(), size_t(n_token)*nv*sizeof(uint16_t)); } +// The first position in each chunk that gets scored. Both the base-logits writer and the KL-divergence reader must +// agree on it, because it sets how many values a chunk occupies in the base file; they therefore read it from here. +// +// The default of n_ctx/2 discards the first half of every chunk, which for a recurrent architecture discards exactly +// what matters: a KDA layer carries its state forward, so a quantization error written into that state accumulates +// with position, and a window starting at n_ctx/2 cannot show the accumulation beginning. Scoring from 0 is what makes +// a position-binned dKLD possible. +static int kld_first_position(int n_ctx) { + const char * env = getenv("LLAMA_KLD_FIRST"); + if (env == nullptr) { + return n_ctx/2; + } + return std::min(std::max(atoi(env), 0), n_ctx - 2); +} + struct kl_divergence_result { double sum_nll = 0.0; double sum_nll2 = 0.0; @@ -253,11 +269,11 @@ static std::pair log_softmax(int n_vocab, const float * logits, c static void process_logits(int n_vocab, const float * logits, const int * tokens, int n_token, std::vector & workers, const std::vector & base_log_probs, kl_divergence_result & kld, - float * kld_values, float * p_diff_values) { + float * kld_values, float * p_diff_values, uint8_t * same_top_values) { std::mutex mutex; const int nv = 2*((n_vocab + 1)/2) + 4; int counter = 0; - auto compute = [&mutex, &counter, &base_log_probs, &kld, n_vocab, logits, tokens, n_token, nv, kld_values, p_diff_values] () { + auto compute = [&mutex, &counter, &base_log_probs, &kld, n_vocab, logits, tokens, n_token, nv, kld_values, p_diff_values, same_top_values] () { kl_divergence_result local_kld; while (true) { std::unique_lock lock(mutex); @@ -279,9 +295,11 @@ static void process_logits(int n_vocab, const float * logits, const int * tokens break; } lock.unlock(); + const size_t same_before = local_kld.n_same_top; std::pair v = log_softmax(n_vocab, logits + size_t(i)*n_vocab, base_log_probs.data() + size_t(i)*nv, tokens[i+1], local_kld); - kld_values[i] = (float)v.first; - p_diff_values[i] = v.second; + kld_values[i] = (float)v.first; + p_diff_values[i] = v.second; + same_top_values[i] = local_kld.n_same_top != same_before; } }; for (auto & w : workers) { @@ -539,7 +557,7 @@ static results_perplexity perplexity(llama_context * ctx, const common_params & // Example, we have a context window of 512, we will compute perplexity for each of the // last 256 tokens. Then, we split the input up into context window size chunks to // process the entire prompt. - const int first = n_ctx/2; + const int first = kld_first_position(n_ctx); for (int i = 0; i < n_chunk; i += n_seq) { const int start = i * n_ctx; @@ -1739,6 +1757,26 @@ static void kl_divergence(llama_context * ctx, const common_params & params) { return; } + // The header carries n_ctx, n_vocab and n_chunk but not the first scored position, and that position sets how many + // values each chunk occupies. A base file written at the default and read under LLAMA_KLD_FIRST=0 would misalign + // every chunk and report a plausible, meaningless dKLD, so check the payload length against what first implies. + { + const size_t nv_bytes = size_t(2*((n_vocab + 1)/2) + 4) * sizeof(uint16_t); + const std::streampos here = in.tellg(); + in.seekg(0, std::ios::end); + const size_t payload = size_t(in.tellg() - here); + in.seekg(here); + const size_t want = size_t(n_ctx - 1 - kld_first_position(n_ctx)) * nv_bytes * size_t(n_chunk); + if (payload != want) { + LOG_ERR("%s: %s holds %zu payload bytes but this run expects %zu (n_ctx %u, n_chunk %d, first %d).\n" + "%s: the base file was written with a different first scored position -- regenerate it with the " + "same LLAMA_KLD_FIRST.\n", + __func__, params.logits_file.c_str(), payload, want, n_ctx, n_chunk, kld_first_position(n_ctx), + __func__); + return; + } + } + const int n_batch = params.n_batch; const int num_batches = (static_cast(n_ctx) + n_batch - 1) / n_batch; // Calculate n_seq based on the logits file's n_ctx, but cap it at what the context supports @@ -1755,9 +1793,12 @@ static void kl_divergence(llama_context * ctx, const common_params & params) { llama_batch batch = llama_batch_init(std::min(n_batch, static_cast(n_ctx)*n_seq), 0, 1); - std::vector log_probs_uint16(size_t(n_ctx - 1 - n_ctx/2) * nv); - std::vector kld_values(size_t(n_ctx - 1 - n_ctx/2)*n_chunk); - std::vector p_diff_values(size_t(n_ctx - 1 - n_ctx/2)*n_chunk); + // sized from the SAME first the writer used, or a lowered LLAMA_KLD_FIRST overruns all three + const int first = kld_first_position(n_ctx); + std::vector log_probs_uint16(size_t(n_ctx - 1 - first) * nv); + std::vector kld_values(size_t(n_ctx - 1 - first)*n_chunk); + std::vector p_diff_values(size_t(n_ctx - 1 - first)*n_chunk); + std::vector same_top_values(size_t(n_ctx - 1 - first)*n_chunk); std::vector logits; if (num_batches > 1) { logits.reserve(size_t(n_ctx) * n_vocab); @@ -1788,8 +1829,12 @@ static void kl_divergence(llama_context * ctx, const common_params & params) { kl_divergence_result kld; auto kld_ptr = kld_values.data(); auto p_diff_ptr = p_diff_values.data(); + auto same_ptr = same_top_values.data(); - const int first = n_ctx/2; + if (first != n_ctx/2) { + LOG_INF("%s: scoring from position %d of %d (LLAMA_KLD_FIRST); the base file must have been written with the " + "same setting, which the size check below enforces\n", __func__, first, n_ctx); + } for (int i = 0; i < n_chunk; i += n_seq) { const int start = i * n_ctx; @@ -1869,9 +1914,10 @@ static void kl_divergence(llama_context * ctx, const common_params & params) { const float * all_logits = num_batches > 1 ? logits.data() : llama_get_logits_ith(ctx, seq*n_ctx + first); process_logits(n_vocab, all_logits, tokens.data() + start + seq*n_ctx + first, n_ctx - 1 - first, - workers, log_probs_uint16, kld, kld_ptr, p_diff_ptr); + workers, log_probs_uint16, kld, kld_ptr, p_diff_ptr, same_ptr); p_diff_ptr += n_ctx - 1 - first; kld_ptr += n_ctx - 1 - first; + same_ptr += n_ctx - 1 - first; LOG("%4d", i + seq + 1); @@ -1909,6 +1955,60 @@ static void kl_divergence(llama_context * ctx, const common_params & params) { if (kld.count < 100) return; // we do not wish to do statistics on so few values + // dKLD BY POSITION, before the sort below discards position for good. + // + // A single mean over the whole corpus cannot distinguish a format that is uniformly slightly worse from one whose + // error compounds: in a recurrent layer the quantized tensors written INTO the state (KDA attn_k, attn_v) carry + // their error forward, so their damage grows with distance from the start of the sequence while a tensor feeding + // only the readout (attn_q) stays flat. Those two look identical in the aggregate and are not remotely the same + // risk. Each chunk restarts from cleared state, so position WITHIN a chunk is the axis, and a chunk shorter than + // the bin edges simply leaves the later bins empty -- which is itself the warning that the run cannot see + // accumulation at all. + { + const size_t per_chunk = size_t(n_ctx - 1 - first); + const size_t filled = size_t(kld_ptr - kld_values.data()); + static const int edges[] = {0, 512, 4096, INT_MAX}; + const char * labels[] = {"0-512", "512-4k", "4k+"}; + std::vector bin[3]; + double p_abs[3] = {0.0, 0.0, 0.0}; + // same top by position: where the head picks a different token than the base, which a divergence spread over + // the tail of the distribution does not show + size_t same[3] = {0, 0, 0}; + for (size_t idx = 0; idx < filled; ++idx) { + const int pos = first + int(idx % per_chunk); + for (int b = 0; b < 3; ++b) { + if (pos >= edges[b] && pos < edges[b+1]) { + bin[b].push_back(kld_values[idx]); + p_abs[b] += std::fabs(p_diff_values[idx]); + same[b] += same_top_values[idx]; + break; + } + } + } + LOG("====== dKLD by position in sequence (scored from %d; %zu values over %zu per chunk) ======\n", + first, filled, per_chunk); + LOG("%-8s %9s %12s %12s %12s %12s %16s\n", "bin", "count", "mean", "median", "p99", "mean|dp|", "same top"); + for (int b = 0; b < 3; ++b) { + if (bin[b].empty()) { + LOG("%-8s %9d %12s (no scored position lands here at n_ctx %u, first %d)\n", + labels[b], 0, "-", n_ctx, first); + continue; + } + std::sort(bin[b].begin(), bin[b].end()); + double sum = 0.0; + for (const float v : bin[b]) { + sum += v; + } + const size_t n = bin[b].size(); + const double top = double(same[b])/n; + LOG("%-8s %9zu %12.6f %12.6f %12.6f %12.6f %7.3f ± %5.3f %%\n", labels[b], n, sum/n, + bin[b][n/2], bin[b][std::min(n - 1, size_t(0.99*double(n)))], p_abs[b]/n, + 100.0*top, n > 1 ? 100.0*sqrt(top*(1.0 - top)/(n - 1)) : 0.0); + } + LOG("A bin's mean RISING with position is error accumulating through the recurrent state: that class stays at\n" + "the higher precision. Flat across bins means the error enters the readout only and does not compound.\n\n"); + } + std::sort(kld_values.begin(), kld_values.end()); std::sort(p_diff_values.begin(), p_diff_values.end());