Repository navigation
hrx: prefill keeps the stored layout for a weight that already has a stored device copy (engine #2137) - #320
Conversation
…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>
|
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
left a comment
There was a problem hiding this comment.
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 inmodel-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_copyis 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 otherwhile 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.
|
VERDICT: READY at What it covers is the case that has no other testFile-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 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 askShow 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 The cheap form is one line: "on Why it belongs in the fork rather than hereIt gates |
…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>
|
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. |
Addendum, and it makes the fix cheaper: the cache next door already solves this exact problem
// 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 ( Two consequences for the review:
I also checked whether |
…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>
|
ready for review at c100cf5 (cl-hrxmem). Answers the CHANGES at 1. Entries never removed. Fixed, with one correction to the consequence. 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 ( New row, as you asked: 2. Two keys. Fixed. Both sides go through 3. To the READY's ask (do the arms fail without the fix): two mutants at this head, each built and run host-only (
Restored head: host 81/81. At |
bong-water-water-bong
left a comment
There was a problem hiding this comment.
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
left a comment
There was a problem hiding this comment.
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
left a comment
There was a problem hiding this comment.
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.
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:
ggml_mul_mat_q6_k_f16_wmma_prefill_wave32(has a stored form)ggml_mul_mat_q4_k_f16_wmma_prefill_wave32(packed only)ggml_mul_mat_quantized_f16_wmma_prefill_conv4_interiorThe change
HostWeightCacherecords the host address of each stored-layout device copy it makes (note_stored_device_copy, inmodel-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=0turns 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=0against unset, order off / on / on / off in one ticket (box load 6 to 8), distinct (weight, layout) device copies fromGGML_HRX_HOST_WEIGHT_TRACE=2:d388da7a6aa9d388da7a6aa9So 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-ops267/267,hrx-backend-testhost 80/80 and device 12/12,test-backend-ops -o MUL_MAT641/641,-o MUL_MAT_ID852/852, and-o MUL_MATwithGGML_HRX_Q4K_PACKED_TAIL=0641/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
ggml_mul_mat_quantized_f16_wmma_prefill_conv4_interior, which has no stored form, and 1 Q6_K. That is on engine Support for xgen and codegen models from Salesforce ggml-org/llama.cpp#2137.🤖 Generated with Claude Code