llama : add an attention split separate from the tensor split by Piggidragon · Pull Request #69 · GenerelSchwerz/llama.cpp

Piggidragon · 2026-09-03T15:28:25Z

Overview

Depends on #66. That PR's commits are included in this branch because the base repository has
no branch to target. Only the last commit, llama : add an attention split separate from the tensor split, belongs to this PR. Merge #66 first and this diff reduces to that commit.

Adds --attn-split (-as), the fraction of the attention heads each GPU gets under
--split-mode tensor.

That share decides two things: how much of a host-resident KV cache a device receives, and how much
attention work it does. Neither has to follow the memory split. On an asymmetric pair the device with
the slow host link is exactly the one that should hold less of the cache, while the memory split still
wants to fill both cards.

Unset, it follows --tensor-split and nothing changes. Only the attention tensors and the cache move;
the rest of the model still follows --tensor-split. A linear-attention layer is a different
mechanism and keeps --tensor-split. The share is rounded to whole heads by the granularity rules
that already keep GQA consistent, so a ratio the head count cannot express lands on the nearest one it
can.

Stacked on the host-cache correctness PR, which is what makes a host-resident cache work under
-sm tensor at all.

This is a new CLI and public API addition, so per CONTRIBUTING it carries a higher bar. The argument
for it over the existing mechanism: --tensor-split cannot express it. With a device-resident cache
both devices run concurrently inside a layer and the collective is a barrier, so --tensor-split
wants the compute ratio; with a host-resident cache the cache placement wants the bandwidth ratio.
One knob cannot serve both, and the measurements below show the second one is worth several times more
than the first on this pair.

Testing

Two GPUs, RTX 4070 (gen4 x16) + RTX 3060 (gen3 x4), CUDA, NCCL, stock clocks.
Qwen3.8-27B-UD-Q5_K_M, a 56935-token prompt from this repository's docs, -c 65536 -n 128 -ngl 99,
q8_0 cache, -sm tensor -ts 50,50 -nkvo --kv-cpu-pinned --recurrent-state-offload. The model has
n_head_kv = 4, so the whole-head shares available are 2,2 / 3,1 / 4,0.

`--attn-split`	rounds to	pp	tg
(follows `-ts`)	2,2	433.33	2.47
`0.75,0.25`	3,1	475.52	3.55
`3,1`	3,1	475.69	3.56
`1,0`	4,0	501.97	7.49

3.03x generation from moving the cache off the slow card, and +16% prefill. 0.75,0.25 and an
explicit 3,1 agree to within noise, which is the rounding working: a ratio the head count cannot
express takes the nearest one it can, and it lands strictly between its two neighbours.

Setting --attn-split when it is not used costs nothing. llama-bench -ngl 99 -ts 50/50 -r 3,
all 16 rows against the parent PR without this commit:

model	sm	nkvo	pp512 parent	pp512 this	tg128 parent	tg128 this
Qwen3.8-27B-UD-IQ2_M	layer	0	755.87	752.83	26.38	26.34
	layer	1	532.98	532.78	8.68	8.66
	tensor	0	610.71	609.89	30.96	30.91
	tensor	1	588.78	589.17	25.61	25.64
gemma-4-26B-A4B	layer	0	2732.69	2733.83	103.83	103.90
	layer	1	1941.45	1986.41	42.87	42.98
	tensor	0	1685.83	1688.82	85.84	86.09
	tensor	1	1339.58	1343.70	34.97	35.37

tests/test-llama-archs.cpp runs the tensor-split matrix a third time with all attention heads on the
first device (attn_split = 1,0) and the cache in host memory, checking the logits against CPU.

test-llama-archs -s 1 at 1, 2, 3 and 4 virtual CUDA devices: passes. test-arg-parser: passes.
Built with -DLLAMA_FATAL_WARNINGS=ON.

Not done here

Deriving the split automatically from the measured bandwidth.
-as for llama-bench, which only takes -sm, -ts, -ub, -ngl and -fa; the sweeps below had
to be run through llama-completion.

Requirements

I have read and agree with the contributing guidelines
AI usage disclosure: YES - implemented by an agent on my instruction, see the Assisted-by: commit trailers.

GenerelSchwerz · 2026-09-05T19:57:40Z

Automated preliminary review by Codex; the repository owner plans a separate manual review.

Verdict: FAIL.

Blocking

common/arg.cpp:2944-2955 uses std::stof without checking full-token consumption or validating the result, so negative, NaN, infinite, partially parsed, and inconsistent zero-sum splits are accepted. include/llama.h:332-334 exposes the same unchecked input to C callers, and src/llama-model.cpp:1235-1237 copies it unchanged. It then reaches cumulative division and float-to-int64_t conversion at src/llama-model.cpp:836-866. Validate complete parsing, finiteness, non-negativity, and a positive total at both the CLI/environment parser and the model API boundary.
PR ggml-meta : split a host-resident KV cache by head #66's reported Gemma4/iSWA result remains 235.0273 with host KV versus 227.2279 with device KV under tensor split, a 3.43% perplexity gap. This PR depends on ggml-meta : split a host-resident KV cache by head #66 and advertises that host-cache configuration; src/llama-model.cpp:831-866 applies attn_split to it. Resolve or explicitly exclude iSWA before merging that advertised use case.

Will slow review

The permanent CLI/environment surface at common/arg.cpp:2938-2957 and public C API field at include/llama.h:332-334 have no dedicated issue or design approval. Please get agreement on the API surface before merging it.
common/arg.cpp:2940-2943 promises the nearest whole-head share, but src/llama-model.cpp:838-866 implements rotating cumulative-floor allocation. For example, a 0.4/0.6 split over four unit-granularity heads gives the first device one head, not the nearest two. Make the behavior and documentation agree.
The newest multi-stream path at ggml/src/ggml-backend-meta.cpp:1468-1483 lacks regression coverage. The added host-KV configurations at tests/test-llama-archs.cpp:1318-1329 reach get_logits, which creates a one-sequence batch at tests/test-llama-archs.cpp:1048-1057; it does not exercise ne[3] > 1. Add a parallel host-KV/custom-split case.

Developmental progress

The broad #57 predecessor was split. This version fixes the maximum-device off-by-one at common/arg.cpp:2948, preserves scheduler-copy source names at ggml/src/ggml-backend.cpp:1037-1048, and adds baseline host-KV and custom-split coverage at tests/test-llama-archs.cpp:1318-1329.

Reviewed head: 1311c7b

A copy of an input is named "<backend>#<source>#<copy>" in a name field of fixed size. With many devices the backend label of the meta backend lists all of them and the source name is what gets cut, so a consumer can no longer tell which tensor the copy was made from. Cut the label instead. Assisted-by: Claude Opus 5

A KV cache in host memory reaches attention as a scheduler copy, which is a leaf in the compute buffer. Such a leaf never reached the device split-state callback and fell through to MIRRORED, while the queries stayed split by head: each device then attended heads whose keys live on the other device. With more than one KV head that aborts in the FlashAttention kernel, or returns wrong output where the query split happens to stay a multiple of the KV head count. Offer a copied-in leaf to the callback under the name the graph gave it, and let the callback recognise the cache there. The cache folds its heads into one flat axis but the copy arrives permuted, with the heads on an axis of their own, so its segments and granularity are rescaled to whole heads. Assisted-by: Claude Opus 5

The KV granularity was derived from the query granularity through n_gqa, which assumes the V side has the head size of the K side. Count whole KV heads and scale each side by its own head size. Assisted-by: Claude Opus 5

A fused QKV puts Kcur and Vcur in a strided view, which a host-resident cache reads back through the meta buffer. The rows split, so the chunk splice could not express the read. Assisted-by: Claude Opus 5

A linear-attention op packs the state it writes back together with its output, so that split does not line up with the one a host-resident state expects. On device the two orders agree; in host memory they disagree, and the split state of the fused op no longer resolves. The state is small next to the attention cache that -nkvo exists to move, so keep it device-resident. Assisted-by: Claude Opus 5

Run the tensor-split architecture matrix a second time with the cache in host memory, and add an 8-device CI run, where the scheduler copy name is long enough to be truncated. Assisted-by: Claude Opus 5

The strided head split path ran only for a single stream, so a host-resident cache built with --parallel N fell through to the chunk splice, which cannot express that write. Loop over ne[3] and offset each side by its own stride. Assisted-by: Claude Opus 5

A split that ends with a view of a host tensor left the subgraph bookkeeping short of the node count and aborted. That happens with more than one cache stream, which the test matrix now covers. Also make the scheduler copy name a stated contract instead of a grammar that two files reconstruct on their own, note the 2d transfer fallback at both strided cache paths, say out loud that split mode tensor overrides the recurrent state placement, and record the Gemma 4 host cache accuracy gap where it is skipped. Assisted-by: Claude Opus 5

The scheduler flags a copy as an input when it keeps more than one, which hid a host-resident cache from the split state callback. A cut source name now asserts instead of naming another tensor, and memset writes a head split like set_tensor does. Assisted-by: Claude Opus 5

The ranged copy of a host-resident window calls set_tensor_async once per stream. A meta backend only takes a whole contiguous tensor, so with two or more streams it aborts. Keep the whole-span copy for a meta destination. Assisted-by: Claude Opus 5

Piggidragon

Rebased onto llama/dev (after #94) and ran the arch-coverage suite on CPU + RTX 4070 + RTX 3060 (CUDA, Meta tensor split).

ctest workspace|generate-models|rollback|save-load|sched-copy: 6/6 pass
test-llama-archs -s 1: aborts on the first arch (llama), Meta -nkvo -np 2 row. llama/dev alone runs the full table with 0 FAIL.

ggml-backend-meta.cpp:1974: GGML_ASSERT(ggml_is_contiguous(tensor)) failed
#4 ggml_backend_meta_set_tensor_async
#5 ggml_backend_tensor_set_2d_async
#6 ggml_backend_sched_graph_compute_async
#11 get_logits

This PR carries older copies of the #66 commits that add the -np 2 row (8206130, 5073fdc), but not 7c4de91 from #66 ("sched : write a multi-stream window to a meta backend in one copy"), which fixes this exact abort. Stacking this PR on current #66 should fix it. The rows after it (including -as -np 2) were not reached, so they are untested.

The #66 review also applies here: the Parallel check from #94 fails on the -np 2 rows for gpt2/starcoder because logits_dev there is a 2-sequence decode.

Rebase note: only conflict was the <filesystem> / <map> includes in test-llama-archs.cpp.

Piggidragon · 2026-09-14T22:10:13Z

+        // with more than one stream that copy is 4d, one stride per stream
+        kv_config kvc_host_streams = kvc_host;
+        kvc_host_streams.n_seq_max = 2;
+        dev_configs.emplace_back(devices_meta, "Meta -nkvo -np 2", LLAMA_SPLIT_MODE_TENSOR, kvc_host_streams);


The full table aborts on this row (llama, first arch) with GGML_ASSERT(ggml_is_contiguous(tensor)) in ggml_backend_meta_set_tensor_async. Needs 7c4de91 from #66.

A multi-stream row decodes its reference as several sequences, so it cannot serve as a one-sequence reference - skip the check there. Give the check the row's offload_kqv, so it covers a host-resident cache with two streams. Assisted-by: Claude Opus 5

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

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

A substring of the name also caught tensors whose split does not count heads. Assisted-by: Claude Opus 5

Piggidragon · 2026-09-17T22:00:59Z

Moved to Piggidragon#7, stacked on its parent branch there.

The diff is now +104/-3 instead of +468/-49, because the parent's commits no longer show up in it. Same commits, same result.

This was referenced Sep 3, 2026

ggml-meta : split a host-resident KV cache by head #66

Closed

ggml-meta : split a host-resident KV cache by head #57

Closed

github-actions Bot added testing ggml devops labels Sep 3, 2026

Piggidragon force-pushed the mgpu/attn-split branch from f2d9059 to 1311c7b Compare September 4, 2026 22:01

Piggidragon force-pushed the mgpu/attn-split branch from 1311c7b to dd7420c Compare September 6, 2026 06:18

github-actions Bot added the documentation Improvements or additions to documentation label Sep 6, 2026

Piggidragon force-pushed the mgpu/attn-split branch 2 times, most recently from dd75518 to 7ca4fd8 Compare September 7, 2026 20:01

Piggidragon added 10 commits September 14, 2026 23:47

llama : split a V tensor on its own head size

fa56aca

The KV granularity was derived from the query granularity through n_gqa, which assumes the V side has the head size of the K side. Count whole KV heads and scale each side by its own head size. Assisted-by: Claude Opus 5

ggml : read back a strided row split from a meta buffer

5719cbf

A fused QKV puts Kcur and Vcur in a strided view, which a host-resident cache reads back through the meta buffer. The rows split, so the chunk splice could not express the read. Assisted-by: Claude Opus 5

tests : cover a host-resident KV cache split by tensor

346b38b

Run the tensor-split architecture matrix a second time with the cache in host memory, and add an 8-device CI run, where the scheduler copy name is long enough to be truncated. Assisted-by: Claude Opus 5

Piggidragon force-pushed the mgpu/attn-split branch from 7ca4fd8 to f4b6147 Compare September 14, 2026 22:09

Piggidragon commented Sep 14, 2026

View reviewed changes

Piggidragon added 4 commits September 17, 2026 21:51

llama : pick the attention split by what the tensor is measured against

6fe4274

A substring of the name also caught tensors whose split does not count heads. Assisted-by: Claude Opus 5

Piggidragon force-pushed the mgpu/attn-split branch from f4b6147 to 6fe4274 Compare September 17, 2026 20:43

Piggidragon closed this Sep 17, 2026

Navigation Menu

llama : add an attention split separate from the tensor split - #69

llama : add an attention split separate from the tensor split#69
Piggidragon wants to merge 14 commits into
GenerelSchwerz:llama/devfrom
Piggidragon:mgpu/attn-split

Piggidragon commented Sep 3, 2026

Uh oh!

GenerelSchwerz commented Sep 5, 2026

Uh oh!

Piggidragon left a comment

Uh oh!

Piggidragon Sep 14, 2026

Uh oh!

Piggidragon commented Sep 17, 2026

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

2 participants

Conversation

Piggidragon commented Sep 3, 2026

Overview

Testing

Not done here

Requirements

Uh oh!

GenerelSchwerz commented Sep 5, 2026

Blocking

Will slow review

Developmental progress

Uh oh!

Piggidragon left a comment

Choose a reason for hiding this comment

Uh oh!

Piggidragon Sep 14, 2026

Choose a reason for hiding this comment

Uh oh!

Piggidragon commented Sep 17, 2026

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

2 participants