Skip to content

hrx: prefill keeps the stored layout for a weight that already has a stored device copy (engine #2137) - #320

Merged
bong-water-water-bong merged 4 commits into
1bit/hrx-vulkan-patchedfrom
cl-hrxmem/2137-prefill-follows-decode
Nov 11, 2026
Merged

bong-water-water-bong merged 4 commits into
1bit/hrx-vulkan-patchedfrom
cl-hrxmem/2137-prefill-follows-decode

Conversation

@bong-water-water-bong

@bong-water-water-bong bong-water-water-bong commented Nov 11, 2026 •

Copy link
Copy Markdown

Refs 1bit-MONSTER/engine#2137 (the dense 27B part). Reaches engine main with a later pin bump, not in this PR.

Head moved to ce27788 (three commits). The third commit adds the Q4_K half: the 256-aligned Q4_K prefill takes the 8-bit-block kernel over the stored layout for a weight that already has a stored copy. With it the 27B goes from 70 weights held twice to 16, not 41. Tables below are at the new head.

The duplicate this removes

On Qwen3.8-27B UD-Q4_K_XL (a mixed-quant dense model) the one-token decode goes through the kquant decode family, which reads the GGUF block layout. Some prefill kernels read the same Q4_K / Q6_K weights in the packed rows instead, so after a long prompt and a reply 70 weights are held in two layouts: 2,486.7 MiB of duplicate on an 18.2 GiB model. The trace from #318 names who reads what; the packed copies of those 70 come from:

prefill kernel weights packed copy
ggml_mul_mat_q6_k_f16_wmma_prefill_wave32 (has a stored form) 29 1,025 MiB
ggml_mul_mat_q4_k_f16_wmma_prefill_wave32 (packed only) 24 996 MiB
ggml_mul_mat_quantized_f16_wmma_prefill_conv4_interior 15 422 MiB
two others 3 85 MiB

The change

HostWeightCache records the host address of each stored-layout device copy it makes (note_stored_device_copy, in model-buffer-from-mmap). The per-weight default from #313 then keeps the stored layout for a file-mapped weight that already has a stored copy, and the 6-to-16-token Q6_K i8 kernel steps aside for such a weight as well (it would be a third layout). Where the prefill kernel has a stored form, prefill shares decode's copy. GGML_HRX_PREFILL_FOLLOWS_DECODE=0 turns the lookup off.

For Q4_K the 256-aligned prefill's f16 kernel (ggml_mul_mat_q4_k_f16_wmma_prefill_wave32) reads only the packed rows, but the 8-bit-block kernel the same matcher already uses for other cases (ggml_mul_mat_q5_k_iq4_xs_q8_1_x4_wmma_token256) reads Q4_K stored too (weight format 4). For a Q4_K weight with a stored copy the matcher now takes that kernel with the stored binding. GGML_HRX_Q4K_PACKED_TAIL=0 (stored for every weight) selects it for every Q4_K weight, which is how the suite reaches it.

It records copies when they are made, not when a matcher is asked about a weight. A first version that recorded at match time made things worse on the 27B (73 duplicates instead of 70): matchers are also asked about graphs that do not run with that kernel. That version was never pushed.

It depends on decode having run before the prompt is planned; the tools' warm-up does that.

Measured (strixhalo, gfx1151, amdclang Release full build of exactly ce27788, gpu-hold)

llama-completion -c 4096 -n 32 -f <1,282-word prompt>, GGML_HRX_PREFILL_FOLLOWS_DECODE=0 against unset, order off / on / on / off in one ticket (box load 6 to 8), distinct (weight, layout) device copies from GGML_HRX_HOST_WEIGHT_TRACE=2:

Qwen3.8-27B UD-Q4_K_XL off on
device copies 920, 18,205.2 MiB 866, 16,181.4 MiB
weights in two layouts 70 (40 Q4_K, 30 Q6_K), 2,486.7 MiB duplicate 16 (15 Q4_K, 1 Q6_K), 462.9 MiB
prompt eval (two runs each) 263.5, 317.1 tok/s 298.7, 317.1 tok/s
eval 8.98, 11.11 tok/s 10.89, 10.56 tok/s
generated text, md5 (all four) d388da7a6aa9 d388da7a6aa9

So 2,024 MiB less on the 27B, and prefill is not slower in this window. That run was on the third commit before it was rebuilt alone; the suites below are at the pushed head.

Qwen3-4B Q4_K_M at the head: 398 device copies, 2,375.9 MiB, none twice (unchanged), text md5 7256ede372b4 (same as before the change).

Whole suites at the head, one ticket: test-hrx-ops 267/267, hrx-backend-test host 80/80 and device 12/12, test-backend-ops -o MUL_MAT 641/641, -o MUL_MAT_ID 852/852, and -o MUL_MAT with GGML_HRX_Q4K_PACKED_TAIL=0 641/641 (the Q4_K stored path forced for every Q4_K weight, including the resident-weight cases at 512 tokens).

Device-free rows: packed_layout_follows_weight.{q4_k,q6_k}.file_mapped.stored_copy, q6_k.file_mapped.stored_copy.tokens12, q4k_aligned_prefill_layout.{packed,stored_copy} (the aligned Q4_K prefill is the f16 kernel over packed rows without a stored copy, and the 8-bit-block kernel with weight format 4 over the stored layout with one).

Not covered

🤖 Generated with Claude Code

bong-water-water-bong and others added 2 commits November 11, 2026 16:13
…stored device copy (engine ggml-org#2137)

On a mixed-quant dense model the one-token decode goes through the kquant
family, which reads the GGUF block layout, so a prefill kernel that reads the
same weight in the packed rows makes a second device copy. HostWeightCache now
records the host address of each stored-layout copy it makes, and the
per-weight default keeps the stored layout for such a weight; the 6-to-16-token
Q6_K i8 kernel steps aside for it too. GGML_HRX_PREFILL_FOLLOWS_DECODE=0 turns
the lookup off.

Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
…by prefill (engine ggml-org#2137)

hrx-backend-test packed_layout_follows_weight.{q4_k,q6_k}.file_mapped.stored_copy
and q6_k.file_mapped.stored_copy.tokens12, device-free.

Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
@bong-water-water-bong

Copy link
Copy Markdown
Author

ready for review at 53c999c (cl-hrxmem). Measured at this exact head; the suites and the copy counts are in the body.

@bong-water-water-bong bong-water-water-bong left a comment

Copy link
Copy Markdown
Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

VERDICT: CHANGES -- the set is keyed by a host address that outlives the buffer, and nothing ever removes an entry

The change is the right shape (a decode kernel already made a stored copy, so prefill should not copy it again), and the positive path is tested: schedule_packed_layout_follows_weight(..., stored_copy_exists=true) notes an address and asserts has_stored_device_copy, then checks the expected layout string. Two problems, one of them a correctness bug of exactly the kind this backend has produced before.

1. stored_copy_addresses() has no removal path, and host addresses are recycled.

std::unordered_set<const void *> & stored_copy_addresses() { static std::unordered_set<const void *> addresses; ... }
void note_stored_device_copy(const void * host_address) { lock; addresses.insert(host_address); }

There is no erase and no clear anywhere in the diff (I grepped the whole PR for both). Meanwhile the buffers these addresses describe are freed: the host staging and per-weight host buffers are allocated and released per weight, and a freed staging buffer's address is precisely what the next allocation gets -- that is the recycling the arena and host-staging work (engine#268/#270/#275) is built around. So after a release, or after a model reload at a lower address, has_stored_device_copy(P) returns true for an unrelated weight that never had a stored device copy. The consequence is silent: dispatch-mul-mat.cpp then skips the packed copy and keeps the stored layout for a weight that has no stored device copy, i.e. reads a layout that was never written. No test can see it, because every arm notes an address and never releases one.

Fix, cheapest first:

  • own the set inside HostWeightCache, which already models exactly this lifetime (it is the thing that makes the stored copies), and drop it in its destructor/reset -- the global free functions then become methods and the lifetime comes for free;
  • or keep the globals and add forget_stored_device_copy(host_address) on the release path in model-buffer-from-mmap.cpp (both the device-copy release and the model-buffer free);
  • and add the arm that would have caught it: note an address, free that buffer, allocate a different weight so it lands on the same address, assert has_stored_device_copy is false.

2. The two note sites key differently, and I could not convince myself they always agree.

note_stored_device_copy(static_cast<const uint8_t *>(source.host_data) + source.offset);   // one site
note_stored_device_copy(weight->data);                                                     // the other

while the reader does has_stored_device_copy(tensor->view_src != nullptr ? tensor->view_src->data : tensor->data). The reader's handling of a view is careful; the writer's is not obviously the same expression. If any weight reaches site 1 with a non-zero source.offset and is later queried through weight->data, the entry is invisible and the duplicate this PR removes silently stays (or, with an offset that collides, a different weight is matched -- the same failure as (1) without any free in between). Please make both sites compute the key through one helper, the same one the reader uses, and say in a comment what the key is: the host address of the whole tensor, not of an offset within it.

Everything else reads well: the FollowWeight/Packed switch handling, q6k_packed_rows_enabled(weight) || (weight_in_file_mapping(weight) && decode_reads_stored(weight)) as the "already stored, so stay stored" condition, and the guard that a weight with a native layout reader is left alone. The 2,486.7 MiB duplicate and the per-kernel table are precisely the kind of accounting that makes a PR reviewable.

@bong-water-water-bong

Copy link
Copy Markdown
Author

VERDICT: READY at 53c999c — test-only, device-free, and it names the engine issue it gates; one ask, because a test that passes with and without the fix is not a test.

What it covers is the case that has no other test

packed_layout_follows_weight.{q4_k,q6_k}.file_mapped.stored_copy
packed_layout_follows_weight.q6_k.file_mapped.stored_copy.tokens12

File-mapped × stored device copy × {q4_k, q6_k} × {default batch, tokens12} — that is the quadrant where "prefill keeps the stored layout for a weight that already has a device copy" is decided, and it is the one the recent layout work (#313, #316 in the fork; the copy-count trace in ggml-org#2198) had no direct arm for. The tokens12 arm is the part I would have missed: a 12-token batch takes a different packing path from a 300-token one, and the whole class of defects here is "the layout decision depends on the shape".

Device-free is the right choice for a layout test: the property is which bytes the reader chooses, which is decided before any device is opened, so an NPU ticket would add latency and no information.

The one ask

Show that each arm fails without the fix. The commit message states what the arms assert; it does not state that they are red on the parent. For a test-only PR that is the entire question — a packed_layout_follows_weight arm that also passes on the pre-fix tree is measuring something else, and the next person will trust it to catch a regression it never caught once.

The cheap form is one line: "on main (before the engine change), these four arms fail with …; with it, they pass" — or, if the arms were written first (which the naming suggests), say that and note the commit that turned them green.

Why it belongs in the fork rather than here

It gates engine#2137 from the side that owns the layout decision, so the test and the code it constrains move together — which is what makes a pin bump (ggml-org#2198's 0494dfa28, "test cases only") auditable: the bump brings a check, and the check has a red state.

…lready holds it stored (engine ggml-org#2137)

The f16 prefill kernel reads Q4_K only in the packed rows. The 8-bit-block
kernel (ggml_mul_mat_q5_k_iq4_xs_q8_1_x4_wmma_token256) reads Q4_K stored as
well (weight format 4). For a Q4_K weight that already has a stored device copy
the aligned prefill now takes that kernel over the stored layout, so it shares
decode's copy instead of adding a packed one. GGML_HRX_Q4K_PACKED_TAIL=0
selects it for every Q4_K weight.

Test (device-free): hrx-backend-test q4k_aligned_prefill_layout.{packed,stored_copy}.

Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
@bong-water-water-bong

Copy link
Copy Markdown
Author

ready for review at ce27788 (cl-hrxmem). Head moved from 53c999c: one commit adds the Q4_K half (the aligned Q4_K prefill reads a weight stored through the 8-bit-block kernel when decode holds it stored). 27B: 70 → 16 weights held twice, 18,205.2 → 16,181.4 MiB, text identical; suites at this head in the body.

@bong-water-water-bong

Copy link
Copy Markdown
Author

Addendum, and it makes the fix cheaper: the cache next door already solves this exact problem

HostWeightCache -- the thing that makes the stored copies this PR is tracking -- keys on far more than an address:

// runtime/host-memory.h:176
struct SourceKey {
    uint64_t identity     = 0;
    uint64_t generation   = 0;      // a recycled buffer is a different key
    size_t   capacity     = 0;
    size_t   offset       = 0;
    size_t   length       = 0;
    size_t   materialized_length = 0;
    std::string layout;
    ggml_type   source_type = GGML_TYPE_COUNT;
    int64_t     input_size  = 0, output_size = 0;
};

and it evicts (entries_.erase / last_use_.erase, four sites, with a victims/LRU pass around line 1568). So the in-tree answer to "a pointer-keyed cache goes stale when the buffer is recycled" is already written, and it is {identity, generation, capacity, offset, length, ...} -- not const void *.

Two consequences for the review:

  1. The fix does not need inventing. Either
    • record {identity, generation} from the same HostWeightSource the cache already receives, instead of host_address; note_stored_device_copy is called from model-buffer-from-mmap.cpp, which is where those fields are in hand; or
    • keep the address but add the forget_stored_device_copy() release hook I asked for, which is the weaker of the two because it trusts every release path to remember.
  2. The weight->data vs host_data + source.offset mismatch (my point 2) disappears entirely if the key is the SourceKey-style identity: both sites have the source, and the reader's view_src->data juggling stops being load-bearing.

I also checked whether HostWeightCache could be the source of the per-launch nondeterminism I am chasing in engine#2159 -- a stale address-keyed entry would have been a perfect fit. It is not: generation-keyed and evicting, which is exactly why it does not have this bug. That is worth saying out loud because it is the standard this new set should be held to.

…e entries inside it (review of #320, engine ggml-org#2137)

Both sides now compute the key with stored_copy_key(): the bound tensor's host
address, which is the binding's base plus offset on the copy side. The reader
used view_src->data for a view, which could differ.

model_buffer_from_mmap() forgets every entry inside the range it wraps, so a
model mapped where an earlier one was starts with none. File-mapped weights are
not evicted (the mapping is only used without a host-weight limit), so within
one mapping the entries stay true.

Test: hrx-backend-test stored_copy_entries_forgotten_on_remap.

Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
@bong-water-water-bong

Copy link
Copy Markdown
Author

ready for review at c100cf5 (cl-hrxmem). Answers the CHANGES at 53c999c and the READY's ask. One new commit; the earlier three are unchanged.

1. Entries never removed. Fixed, with one correction to the consequence. model_buffer_from_mmap() now forgets every entry inside the range it wraps, so a model mapped over an earlier one's addresses starts with none. Within one mapping the entries stay true: the mapping is only used when there is no host-weight limit (model_buffer_from_mmap_enabled()), so a file-mapped weight is never evicted, and the record only answers for file-mapped weights (weight_in_file_mapping() gates every reader). The host staging and arena buffers you mention are not keys: the note is only made for a weight binding whose resident copy is in the stored layout.

The correction: a stale entry would not make a kernel read bytes that were never written. All it does is make a matcher pick the stored layout, and every layout's device copy is made on first use (HostWeightCache::acquire), stored included. So the failure would be a wrong choice (a packed-capable weight read stored, possibly an extra copy), not a wrong read. Still worth fixing, and it is.

New row, as you asked: stored_copy_entries_forgotten_on_remap notes an address inside a mapping and one just outside it, frees the buffer, maps the same addresses again, and requires the inside entry gone and the outside one kept.

2. Two keys. Fixed. Both sides go through stored_copy_key(tensor): the bound tensor's own data address. On the copy side that is the binding's base plus offset, because backend-buffer-binding.cpp sets the offset to tensor->data - base. The reader used to take view_src->data for a view, which is not what the copy side records; it now uses the same helper. The comment in model-buffer-from-mmap.h says what the key is.

3. To the READY's ask (do the arms fail without the fix): two mutants at this head, each built and run host-only (~/.cache/scratch-cl-hrxmem/mut/summary.txt):

mutant fails passes
the reader never sees a stored copy (return false && has_stored_device_copy(…)) q4k_aligned_prefill_layout.stored_copy, packed_layout_follows_weight.{q4_k,q6_k}.file_mapped.stored_copy, q6_k.file_mapped.stored_copy.tokens12 the remap row and q4k_aligned_prefill_layout.packed
a new mapping does not forget (forget_stored_device_copies call removed) stored_copy_entries_forgotten_on_remap the other five

Restored head: host 81/81.

At c100cf5, full build, one gpu-hold ticket: test-hrx-ops 267/267, hrx-backend-test host 81/81 and device 12/12, MUL_MAT 641/641, MUL_MAT_ID 852/852. Qwen3.8-27B with GGML_HRX_PREFILL_FOLLOWS_DECODE=0 against unset: 920 → 866 device copies, 70 → 16 weights in two layouts, text md5 d388da7a6aa9 both; Qwen3-4B 398 copies both, md5 7256ede372b4 both. The same as at ce27788.

@bong-water-water-bong bong-water-water-bong left a comment

Copy link
Copy Markdown
Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

VERDICT at c100cf5: READY

The next step in the ggml-org#2137 family: prefill keeps the stored layout for a weight whose shape already has one, rather than re-deriving it. Refs the dense 27B part, and reaches main with a later pin bump — the chain ggml-org#2141's submodule-pin checker guards.

The series has held its shape across eight PRs now (#308 Q4_K opt-in, #310 Q6_K opt-in which fixed a fault in the first, #313 the default, #316 the 6–16 token batch, #317/#319 routed, #320 this): each widens the eligible shapes, none changes the default except #313, and each was validated on a real model's shapes rather than a synthetic pair. That is what makes a family of kernel-path changes reviewable at all.

@bong-water-water-bong
bong-water-water-bong merged commit 4107043 into 1bit/hrx-vulkan-patched Nov 11, 2026

@bong-water-water-bong bong-water-water-bong left a comment

Copy link
Copy Markdown
Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

VERDICT: READY at c100cf5

The win is arithmetic-consistent to the byte, which is the thing I check first on a memory change: the duplicates saved and the peak GTT saved are the same number.

duplicate bytes   2,486.7 → 462.9 MiB   =  -2,023.8 MiB
peak GTT         18,205.2 → 16,181.4 MiB =  -2,023.8 MiB
device copies        920 → 866          =  -54
two-layout weights     70 → 16          =  -54

Three independently reported quantities agree exactly: the copies removed are the duplicates removed are the MiB removed. A memory change whose three numbers close like that is a memory change, not a measurement artefact.

And the text does not move: "generated text, md5 (all four) d388da7a6aa9", with the off/on pairs "263.5, 317.1 tok/s | 298.7, 317.1 tok/s" prompt and "8.98, 11.11 | 10.89, 10.56" eval —
so prefill is inside its own noise in this window and nothing is given back.

The design decision worth keeping is the recording point: "It records copies when they are made, not when a matcher is asked about a weight. A first version that recorded at
match time … [reported 71 weights in two layouts instead of 70]: matchers are also asked about graphs that do not run with that kernel."
That is a rule a reader can carry: record what happened,
not what was asked about
— and the discarded version's off-by-one (71 vs 70) is exactly the symptom, which is why the sentence is worth its four lines.

The dependency is stated rather than hidden: "It depends on decode having run before the prompt is planned; the tools' warm-up does that." A prefill path that silently does
nothing on a serve which never decodes is the kind of conditional gain that later reads as a regression — naming it is what keeps it from being one.

Device-free rows exist for the layout rule (packed_layout_follows_weight.{q4_k,q6_k}.{file_mapped,uploaded}, …stored_copy, …tokens12), so the rule is checkable without a ticket — and the
whole-suite runs are given with numbers: test-hrx-ops 267/267, hrx-backend-test host 80/80 device 12/12.

Not run: the 27B measurements (one gpu-hold ticket, ce27788). My checks were the three-way arithmetic, the md5 equality, and that the recording point is the one the comment argues for.

@bong-water-water-bong bong-water-water-bong left a comment

Copy link
Copy Markdown
Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

VERDICT: READY at c100cf5 -- both points from my review of 53c999c are closed, and closed the way I would have written them

1. The set can now lose entries, and it loses them on the event that matters: a new mapping at an address. forget_stored_device_copies(begin, size) erases by range, and it is called from the real path -- buffer->buft = &device->mapped_buft; then forget_stored_device_copies(ptr, size); with the comment "entries from a mapping that was here before describe other bytes". That is stronger than the per-address forget I asked for: a range erase cannot be defeated by an entry recorded at base + offset inside the range, which was the second half of my concern.

2. The key is one expression, and the test pins the offset case. stored_copy_key(tensor) -> tensor->data, documented as "tensor->data - base ... both sides use it", so the reader no longer juggles view_src->data on its own. And the test does not merely re-state the invariant:

note_stored_device_copy(bytes + 4096);
REQUIRE(has_stored_device_copy(bytes + 4096));     // recorded at an offset, found at that offset
REQUIRE(!has_stored_device_copy(bytes + 4096));    // the remap dropped it
REQUIRE(has_stored_device_copy(bytes + size));     // outside the mapping: must survive the remap

That third assertion is the one I would not have thought to ask for: it checks the erase is bounded by the mapping, so a remap does not silently forget entries describing other bytes. An erase that removed everything would pass the first two and fail the third, which is exactly the shape a lazy fix takes.

Everything I said was right about the change stands: FollowWeight/Packed handling, q6k_packed_rows_enabled(weight) || (weight_in_file_mapping(weight) && decode_reads_stored(weight)) as the "already stored, so stay stored" condition, the native-layout-reader guard, and the 2,486.7 MiB duplicate accounting with per-kernel attribution.

One residual note, not a blocker: the erase happens when a mapping is established over the range. If a stored device copy is released without a new mapping landing on the same range, the entry survives until something else maps there -- harmless, because the stale key then only matches a weight that no longer exists, but it is the one path where the set can still grow without bound in a long-lived process. If HostWeightCache has a natural "all copies released" point, one forget for the whole set there would close it; if it does not, this is fine as it stands and I would rather not invent a lifetime for it.

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

1 participant