Repository navigation
llama : attention split on main, for the KV cache in VRAM - #15
Merged
Merged
Conversation
The share of the attention heads a device gets decides both how much of a host- resident cache it receives and how much attention work it does, and neither has to follow the memory split. Add --attn-split (-as) to set it; unset it follows --tensor-split and nothing changes. The share is rounded to whole heads by the granularity rules that already keep GQA consistent. Assisted-by: Claude Opus 5 (cherry picked from commit 88c3a20)
std::stof took a partial parse, a negative, a NaN and an all-zero set without a word, and those reach a cumulative division and a conversion to whole heads. Reject them at the command line and fall back to the tensor split at the API boundary. Also say what the rounding does instead of promising the nearest share. Assisted-by: Claude Opus 5 (cherry picked from commit aef345a)
A substring of the name also caught tensors whose split does not count heads. Assisted-by: Claude Opus 5 (cherry picked from commit 6fe4274)
… log the result The attention split now applies to a device-resident KV cache on main, without the host KV changes it was stacked on. - each layer gets whole KV heads and the rounding goes to a different device from layer to layer, so the share over all layers follows the ratio; the help said the shares are rounded down - the model load logs the KV heads each device got over all layers - llama-bench takes -as like -ts - test-llama-archs runs the tensor split with all heads on the first device and with heads split 1/3 Qwen3.8-27B UD-Q4_K_S on an RTX 4070 + RTX 3060, 262144 tokens of q5_1 cache in 2 slots and MTP: -ts 0.6,0.4 alone does not fit, -ts 0.63,0.37 -as 1,3 does (11880 / 11362 MiB) and is 11-15% faster than 50/50 at short context and 4-11% at 32k (pp512 652 -> 727, pp3 74.8 -> 86.0, tg32 29.8 -> 33.1). Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
- the split boundaries are computed in double and with a margin below whole units: a share that is inexact in binary rounded the tensors of a layer to different head counts (-as 0,0.9 or -ts 0.45,0.45 aborted at the first decode on Qwen3.8-27B) - the attention split spreads the rounding over the layers of a kind, so the heads over all layers follow the ratio within one head: -as 0.3,0.7 gives 20 of 68 KV heads (29.4%) instead of 25 (36.8%) - the shares are kept as fractions of their sum, computed in double, so large shares cannot overflow into NaN - the load log skips layers in host memory and layers without their own KV, and warns when the shares of the devices in use add up to 0 or no layer of the model splits its heads by attn_output.weight Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
1 participant
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Rebases
--attn-split(#7) ontomain, without the host-KV stack it depended on, so it works with the KV cache in VRAM. #7 stays unchanged.With
-sm tensoron GPUs of different speed, the faster GPU can now get more of the weights. The slower one takes more KV heads in exchange, so both still fit in memory.What changed against #7
-as 0.3,0.7gives 20 of 68 KV heads (29.4%). Before, the rounding alternated between the GPUs and gave 25 heads (36.8%).attention split, KV heads of all layers per device: device 0: 17 (25.0%), device 1: 51 (75.0%).llama-bench -as: works like-tsand is shown in the results table. llama : add an attention split separate from the tensor split GenerelSchwerz/llama.cpp#69 listed it as missing.test-llama-archsaddsMeta -as 1/0(all heads on the first GPU) andMeta -as 1/3, with the KV cache in VRAM. 214 OK on 4070 + 3060. The only failures are maion-coder MoE in the "Mixed batch" column, which also fails on a single GPU and onmain.Review
A review (reviewer-xhigh) found 5 issues; all are fixed in ff1a44e:
-as 0,0.9, gave Q/O and K/V/KV cache of one layer different head counts, and the server aborted at the first decode.-ts 0.45,0.45did the same; that bug is older than this PR. The boundaries are now computed in double, with a small margin below whole units.-as 3e38,3e38) gave NaN. The shares are now stored as fractions of their sum, computed in double.-ngl) and layers without their own KV cache are no longer counted.Checked after the fixes on the 27B:
-as 0,0.9,-ts 0.45,0.45,-as 0.3,0.7,-as 3e38,3e38,-as 0,0,1and the production config all load and answer. test-llama-archs gives 214 OK for the-asconfigs.Memory
Qwen3.8-27B UD-Q4_K_S. What follows which split:
-ts-ts-ts-ts-as-asConfigurations that load with 2 slots, 262144 tokens and MTP on RTX 4070 + RTX 3060 (VRAM in use, MiB):
-ts-asThese are the sizes at load. While running, the 4070 needs about 150 MiB more (CUDA graphs, pools).
-ts 0.63,0.37 -as 1,3loads, but it runs out of memory at the first long request (cudaGraphInstantiate).-ts 0.625,0.375 -as 1,3runs at full context with 10k and 95k-token prompts and keeps about 550 MiB free on both GPUs at peak. The linear-attention layers split in 16 head groups, so-tsmoves their share in coarse steps (0.625 is exactly 10 of 16). That is also why 0.61 already overfills the 3060.Speed
llama-bench, q5_1 cache, native FA,-ub 512 -b 1024, NCCL, t/s:-ts 0.63 -as 1/3-ts 0.76 -as 0/1-ts 0.56 -as 0.375Decode is bound by memory bandwidth (4070: 504 GB/s, 3060: 360 GB/s), prefill by compute (about 2.2:1). 0.63 sits between the two ratios and is the best point overall. All KV heads on the 3060 helps only prefill, because it moves the long-context attention to the slower card.
Together with #14 (small all-reduces through the internal pipeline),
-ts 0.63 -as 1/3against production 50/50:🤖 Generated with Claude Code