context : preserve quantized block sizes in device state I/O

[arc-uri-el]

Overview

Device checkpoint I/O converts byte counts into 1D tensor dimensions with bytes / ggml_element_size(tensor). ggml_element_size returns ggml_type_size, which is the byte size of a whole quantization block. For Q8_0 and Q4_0, the resulting view has one thirty-second of the intended logical elements.

Multiply by ggml_blck_size in the writer, reader and fragmented-copy subview. No public API, serialization metadata, allocation policy, kernels or sampling code changes. The existing save/load test gains a fragmented-to-compact round trip that compares the resulting host state byte-for-byte.

The current-base Q8_0 and Q4_0 tests segfault during the on-device scatter restore on both CPU and Radeon Vulkan. The isolated patch makes both pass. F32 passes before and after. This was investigated on a Strix Halo machine while qualifying quantized prompt checkpoints, not inferred from an unrelated device.

Measurements

Device:     Ryzen AI MAX+ 395 / ASUS ROG Flow Z13 GZ302EAC
Memory:     128 GB LPDDR5X
Power:      AC; performance platform profile and CPU governor; GPU clocks auto
Firmware:   BIOS 304; PPT settings 80/92/93 W; VRAM 512 MiB
Kernel:     Linux 7.3.0-rc1; ttm.pages_limit=30408704; GTT 116 GiB
Backend:    CPU and Radeon 8060S Vulkan RADV, Mesa 26.2.0
Build:      GCC 15.2, Release, shared libraries, native CPU, Vulkan, OpenMP; -j8
Baseline:   7449a0fe9710ab584c5f9a6d25e7a31eea2708b8
Change:     fc349bd785d8bd878235039d8e35a186d872d0cb
Model:      test-llama-archs generated llama-dense; separate official Qwen3.8-27B Q4_K_M control

The baseline and candidate were built separately in this session with the same settings and no compiler warnings. The contribution's attribution was subsequently shortened; its complete source tree is byte-identical to the tested tree. Runtime libraries were selected from their corresponding build directories. The later master merge 99a40a3e6 changes speculative replay, not these I/O or test lines.

Baseline / after:

Backend	KV	Baseline	Candidate
CPU	F32	pass	pass, all 9 tests
CPU	Q8_0	SIGSEGV at device scatter restore	pass, all 9 tests
CPU	Q4_0	SIGSEGV at device scatter restore	pass, all 9 tests
Vulkan0	F32	pass	pass, all 9 tests
Vulkan0	Q8_0	SIGSEGV at device scatter restore	pass, all 9 tests
Vulkan0	Q4_0	SIGSEGV at device scatter restore	pass, all 9 tests

The existing host save/load and generated-token checks pass before the failing baseline scatter case. The candidate passes those checks, the host/device scatter byte comparison and the added compact-destination comparison. A second, independently generated synthetic llama fixture also reproduces the quantized failure and passes with the fix.

Reproduce with the existing tools (run in a disposable working directory because the save/load test writes dump_state.bin):

cmake --build build --target test-llama-archs test-save-load-state
mkdir -p fixtures
build/bin/test-llama-archs -a llama -s 1234 -o fixtures
build/bin/test-save-load-state -m fixtures/llama-dense.gguf \
  -c 256 -b 64 -ub 64 -n 8 -t 1 -ngl all -dev Vulkan0 \
  -ctk q8_0 -ctv q8_0 -fa on

Repeat with q4_0. CPU control: -ngl 0 -dev none. F32 control: -ctk f32 -ctv f32 -fa off.

Correctness:

On the unchanged official ggml-org/Qwen3.8-27B-GGUF Q4_K_M, separate baseline/candidate Vulkan runs also produced byte-identical prompt IDs, nine complete logit rows (prefill plus eight greedy decode steps) and selected IDs for each of this repository's prose, code, structured and numeric corpora. Each corpus was repeated and truncated to 512 input tokens; both used Q8_0 KV, FA on and the same model file. This is a bounded fresh-inference control, not a claim that every state/backend path is covered.

No throughput improvement is claimed. llama-bench, perplexity and the full backend-op suite were not run for this I/O-only patch; the acceptance evidence is the failing and passing state round trips and exact output controls above.

Requirements

• I have read and agree with the contributing guidelines.
• This change is justified by the CPU and actual Strix Halo Vulkan state measurements above. No GPU-specific behavior is inferred for untested backends.
• AI usage disclosure: AGENT-AUTHORED. GPT 6 Astra investigated the failure, isolated the patch, extended and executed tests, and prepared this PR at the explicit direction of the submitting account's owner. The submitting account owns review and follow-up. Original broader local investigations used LLM assistance as well. No human Reviewed-by or Tested-by attestation is claimed.
• What was NOT verified: other GPU backends, other operating systems, all model architectures, full occupied long contexts, perplexity and the full backend-op/performance matrix. The test proves the stated sizing and restoration boundary; it does not claim global correctness.

[LaurentZuijdwijk]

Thank you for your PR. Reviewing it now

[LaurentZuijdwijk]

Approving. The fix is correct and I reproduced both the failure and the fix locally.

ggml_element_size() returns ggml_type_size() — the byte size of a whole quantization block — so dividing a byte count by it yielded a block count that was then handed to ggml_view_1d as an element count. Multiplying by ggml_blck_size is the right correction, and it is algebraically a no-op for F32/F16/BF16 where blck_size == 1, which matches the reported "F32 passes before and after".

The SIGSEGV also has a coherent mechanism, which is worth stating since an undersized view on its own would under-copy silently rather than crash: each ggml_nbytes(src_t) becomes 1/32 of the real chunk, so their sum never reaches the total computed from the true rinfo.size values. The while (src_pos < total) cursor therefore walks src_j/dst_i off the end of mbuf_cur.cpy/mbuf.org and indexes out of bounds.

I also checked that the byte offsets cannot land mid-block: ggml_nbytes of a quantized tensor is a multiple of ggml_type_size, n_copy is a min of differences of such multiples, and offsets only advance by n_copy — so every n_copy / ggml_type_size divides exactly.

Verification

Built the PR tree, then rebuilt with only src/llama-context.cpp reverted to the base commit (keeping the new Test 9) to isolate the fix.

Release, shared libs, GCC, CPU backend only (-DGGML_VULKAN=OFF), llama-dense.gguf from test-llama-archs:

KV	baseline (fix reverted)	PR 35
f32	passes all 9	passes all 9
q8_0	SIGSEGV at Test 7	passes all 9
q4_0	SIGSEGV at Test 7	passes all 9

The crash lands on Test 7, which already existed — so this is a reachable bug on current master for anyone using -ctk q8_0/q4_0 with device checkpoint I/O, not something the new test invented.

Not verified

• The Vulkan path was not exercised — no GPU was available in this session. Only the CPU backend was run.
• I did not independently re-derive the measurement table in the PR description. The reproduction above was run against this PR's own tree, so it stands on its own rather than relying on those numbers.

Non-blocking nits

• In the fragmented-copy hunk, dst_v is sized from src_t->type rather than dst_t->type. Same type in practice, but correct-by-assumption rather than self-evidently correct.
• (size / ggml_type_size(t)) * ggml_blck_size(t) now appears three times verbatim; a small helper would name the intent.

Disclosure per AGENTS.md: this review was produced with Claude Code acting as an agent in my session. The builds, the revert-and-rebuild, and the test runs above were performed by the agent on this machine; I have not separately re-run them by hand.