Skip to content

ggml-hrx: accelerate row64 prefill and share weight layouts - #171

Closed
AaronStGeorge wants to merge 9 commits into
hrx-graph-develop-v2from
users/astgeorg/row64-prefill-combined
Closed

AaronStGeorge wants to merge 9 commits into
hrx-graph-develop-v2from
users/astgeorg/row64-prefill-combined

Conversation

@AaronStGeorge

@AaronStGeorge AaronStGeorge commented Nov 2, 2026 •

Copy link
Copy Markdown
Collaborator

Overview

Builds on merged #160. This draft now targets hrx-graph-develop-v2 at 930e43fe920a9c519eefad7d6fccc8498b757050; rebased head is 9c52414d8be8c5ed480965d3f15e742a47432c9e. The diff contains the row64 consumer and prefill changes on top of persistent demand-based materialization.

In the original six-pair study below, on Llama 3.1 8B Instruct Q4_K_M / RX 7900 XTX, the combined change delivers about 18% faster 66-token prefill, 27% faster 128-token prefill, and 15% faster actual Lemonade prefill, while reducing managed resident weight storage by 4,071,161,856 bytes (4.07 GB) compared with #160. The gain is confirmed separately with default embedding placement and with the embedding explicitly placed on HRX. Decode measurements remain near baseline. Numerical limitations below remain open for review; this is not an unconditional correctness pass.

Implementation

  • Make eligible Q4_K/Q6_K consumers use one row64 representation across prefill, decode, QKV, paired gate/up, and embedding/get_rows where eligible. This lets ggml-hrx: materialize weights on demand and keep them resident #160 materialize only that requested layout instead of retaining different layouts for different consumers. CPU-origin and other native-only source boundaries remain intact.
  • Reuse Q4 block factors and scale words across eight K32 groups in the generic single and paired tiled readers. The decoder receives the logical row directly. Payload decoding and arithmetic order are preserved; dedicated large-prefill and decode readers are unchanged by this optimization.
  • Select the existing masked wave32 M128/N128 Q6 prefill body for eligible 64-127-token inputs. It consumes F32 activations and converts inside the kernel, avoiding a separate short-prompt conversion dispatch. Existing aligned 128-2048-token routes remain unchanged.
  • Add scalar-dequantization, tail-boundary, independent cache-publication, and metadata-reuse reference cases. Actual pp66 selects 31 specialized Q6 commands; it also loses 15 residual fusions, taking the program from 323 to 338 commands. The reported measurements include this tradeoff.

Scheduling and metadata-reuse ideas were informed by the Loom kernels in hrx-solutions. HRX stays pinned; there is no runtime upgrade, cache-policy change, eviction, or quantization-format change.

Post-rebase quick validation 2026-11-03

The rebased head 9c52414d8be8c5ed480965d3f15e742a47432c9e was rebuilt and measured against current develop 930e43fe920a9c519eefad7d6fccc8498b757050. Both use the same pinned HRX 82214dfb029975205d1c3a59164b678e2e21fc41, matched build options, identical runtime payloads, frozen executable/dependency snapshots, and exclusive RX 7900 XTX access. The Llama 3.1 8B Instruct Q4_K_M model is unchanged.

Two fresh AB/BA pairs, default embedding placement, five repetitions per invocation, primary mean of repetitions 2–5:

Workload Throughput change Nominal 95% interval
pp66 +18.71% +16.54% to +20.92%
pp128 +26.86% +26.44% to +27.28%
pp256 -0.19% -4.92% to +4.76%
pp512 +0.01% -3.03% to +3.15%
pp1024 -0.08% -3.78% to +3.76%
tg128 +0.01% -0.11% to +0.13%

Actual Lemonade 2026.43.3 chat uses two fresh AB/BA pairs, one warmup and three measured retained-model requests per invocation, temperature zero, and no prompt cache or reload. All 12 measured requests complete with 66 input tokens, 256 output tokens, and zero cached prompt tokens. Prefill improves 15.09% (nominal interval +7.40% to +23.33%); backend prompt duration falls from 62.281 to 54.125 ms. Generation changes +1.21% and client duration -1.53%, with both intervals spanning zero. All 12 request timing records reconcile with backend logs; TTFT here is backend prompt duration. These two-pair intervals use one degree of freedom and are a bounded regression screen.

Focused buffer/backend/row64/source-boundary/persistence tests in replay and direct modes report 130 passes and four inherited Q4 reference failures out of 134 checks, with no unexpected failures. The same two paired Q4 reference cases fail per mode as in the original campaign. Four-chunk PPL checks at context/batch512 pass the unchanged HRX/Vulkan ratio <=1.10 gate:

Microbatch Baseline HRX PPL Candidate HRX PPL Vulkan PPL on both builds Candidate HRX/Vulkan
66 6.3398 6.3404 6.3466 0.999023
512 6.3447 6.3447 6.3423 1.000378

Exact chat equality remains incomplete: 2/6 paired responses match, all at request position two. Each build's sequence repeats across fresh invocations, but text varies across its request positions. This differs from the original campaign's agreement pattern; its cause is not isolated. All results and text differences are retained, and PPL does not establish equal logits.

The rebase conflicts were confined to test adaptation for the merged base. All eleven kernel/dispatch files are byte-identical to the original candidate; the buffer-binding delta retains upstream cleanup. A final test-only commit removes an assertion on the deleted weight_peak_bytes API while retaining exact live bytes, one-materialization, and zero-duplicate checks. Full suites, memory tracing, and broader model coverage were not rerun in this quick pass. Detailed report and indexed raw evidence are in the workspace record docs/hrx-row64-rebase-validation.md and investigation/row64-rebase-20261103.

Original pre-rebase validation

Historical evidence: the following six-pair, full-suite, memory, and longer-PPL results apply to the exact original revisions below. They are not measurements of the rebased head.

Candidate: 0505f403573adf990d4e8956044bcc701bfca80c. Baseline: #160 at 463fd0b285835d36cf15544b7bc121fc37400ca7. Both use HRX 82214dfb029975205d1c3a59164b678e2e21fc41 on the same RX 7900 XTX (gfx1100). Builds use matched options and frozen executable/dependency snapshots. Performance runs are uninstrumented and use exclusive GPU access. These results cover one model quantization and one GPU, not the 30-model coverage of the parent PR.

Whole-model prefill and decode

Six fresh alternating AB/BA pairs per placement policy, independently confirmed after exploratory selection. llama-bench uses five repetitions per invocation; the predeclared primary is the mean of repetitions 2-5. Both sides use eight threads, flash attention, F16 KV, mmap, and batch/microbatch 512. Percentages and nominal paired 95% intervals use log ratios; positive throughput changes are faster. First timed and all-five results remain preserved; the first timed repetition follows built-in warmup and is not a cold-load measurement.

Workload Default embedding: change [95% interval] HRX embedding: change [95% interval]
pp66 +17.927% [+17.752, +18.102] +18.293% [+17.792, +18.795]
pp128 +26.812% [+26.403, +27.223] +26.899% [+26.301, +27.500]
pp256 +0.017% [-0.296, +0.331] -0.028% [-0.494, +0.440]
pp512 -0.145% [-0.451, +0.162] -0.178% [-0.435, +0.079]
pp1024 -0.135% [-0.236, -0.035] -0.181% [-0.540, +0.178]
tg128 +0.029% [-0.132, +0.191] -0.033% [-0.144, +0.077]

The small default-placement pp1024 decrease is retained. Intervals spanning zero do not establish formal equivalence. The combined candidate is confirmed against #160; its incremental benefit over Q4 metadata reuse alone remains uncertain in the smaller selection experiment.

Actual Lemonade chat-long-output

Lemonade 2026.43.3, six new AB/BA pairs per placement policy, one warmup and three measured requests per invocation. The primary includes all three measured requests. The protocol uses --no-reload, --no-cache-prompt, context 4096, parallel 1, temperature zero, natural EOS, a 256-token output limit, and speculation disabled. All requests have 66 input tokens, 256 output tokens, and zero cached prompt tokens. Default placement and token_embd.weight=HRX0 are separate matched comparisons.

Metric Default embedding: change [95% interval] HRX embedding: change [95% interval]
Prefill tokens/s +15.172% [+14.639, +15.707] +15.492% [+14.840, +16.147]
Backend prompt duration -13.136% [-13.542, -12.727] -13.342% [-13.776, -12.906]
Generation tokens/s +0.762% [-0.824, +2.374] +0.723% [-1.756, +3.264]
Backend total request time -1.121% [-2.624, +0.405] -1.077% [-3.452, +1.357]
Client request time -1.109% [-2.597, +0.402] -1.036% [-3.335, +1.317]

Default-placement prompt time falls from 62.190 to 54.021 ms; HRX-placement prompt time falls from 62.008 to 53.735 ms. Prefill is only about 3% of the approximately 2.13-second backend request, so saving about 8 ms predicts roughly 0.4% total improvement alone. A resolved total-time gain is not a prefill acceptance requirement. No confirmed generation or total-request gain is claimed.

Lemonade's reported TTFT here is backend prompt duration, not streaming first-byte latency. Generation timing includes first autoregressive preparation. All 36 measured requests per policy reconcile with backend logs. Recorded CLI/daemon hashes match the pinned release archive and scenario; original HTTP bodies were not captured, so retrospective wire-payload equality is not claimed.

Resident weights and repeated serving

Managed weight storage falls from 8,984,060,160 to 4,912,898,304 bytes under either placement policy. With HRX embedding, each of all 226 eligible Q4/Q6 weights has one row64 representation; 66 F32 weights remain separate. Default placement has 225 eligible row64 weights plus the CPU-origin embedding's native representation. Materialization misses fall from 513 to 292. Across 20 subsequent prefill/decode requests, both candidate policies perform no additional weight allocations or uploads. All 20 native diagnostic outputs match their same-policy baselines at the tested 16-token output length.

These byte counts describe managed weight representations, not total process/device memory. The source GGUF remains available on the host.

Numerical checks and known limits

The complete buffer, backend, operations, and Loom JIT suites in replay and direct modes report 1,050 passes and four inherited Q4 reference failures out of 1,054 checks, with no unexpected failures. The four failures are two new input cases in each mode, reproduced on unchanged row64 control kernels:

  • The paired constant-input scalar check has 1.666% relative L2 error against a 1% gate, with maximum absolute error 7.47e-8.
  • The paired patterned-input check passes scalar relative L2 at 0.635%, but fails the existing CPU absolute/relative assertion.

Complete single-reader and both paired-case output arrays are byte-identical between the unchanged row64 control and the optimized Q4 kernels in both execution modes. This supports evaluating the optimization's delta; it does not waive the failed absolute-reference gates or establish correctness of every row64 input.

Q6 independent scalar checks cover 63, 64, 65, 66, 95, 127, 128, and 129 tokens, with K=768 and N=192 to exercise block and tail boundaries. The short route measures about 0.075% relative L2. The independent embedding/RMS/projection/cache-publication reference measures 0.0088% relative L2 over written rows, with all 481 untouched rows retaining their exact sentinel. Native fused-path coverage is preserved.

PPL uses context/batch 512, matched explicit embedding placement, the same pinned corpus, and Vulkan on the same physical GPU. There is one run per build/backend/regime. All three regimes pass the unchanged maximum HRX/Vulkan PPL ratio of 1.10:

Microbatch / chunks Baseline HRX PPL Candidate HRX PPL Vulkan PPL Candidate HRX / Vulkan
512 / 32 7.0713 7.0713 7.0760 0.999336
1 / 2 5.4104 5.4104 5.4195 0.998321
66 / 32 7.0996 7.0846 7.0755 1.001286

The short-route PPL difference is not a demonstrated quality improvement, and these checks do not establish equal logits.

Long-response equality is incomplete. For the same fixed Lemonade prompt repeated three times, every pair under both policies produces baseline A/B/B and candidate B/B/B. Thus 12 of 18 paired responses match per policy; the first measured response differs despite equal lengths. Candidate output is stable across all three positions and matches the baseline's later responses. The differing texts share 870 characters. The cause is not isolated; the changed Q6 accumulation path is a plausible contributor. These differences limit attribution of generation timing and remain part of numerical review. The original all-three-request performance primary includes them; no samples were discarded to obtain agreement.

Raw runs, failed attempts, exact source exports, and the detailed experiment report are preserved in the workspace evidence record for lemonade-ws-p6dm. Parent review verified all 169 indexed review-file sizes/hashes and reconstructed the exact candidate source tree from the saved patches. Broader model/GPU coverage and resolution of the numerical limits remain outside this completed experiment.

Requirements

  • I have read and agree with the contributing guidelines - human review/sign-off remains pending for this draft.
  • AI usage disclosure: YES. Codex assisted implementation, tests, experiment orchestration, review, and this draft description under the contributor's direction. The contributor remains responsible for reviewing and maintaining the submitted changes.

Base automatically changed from users/astgeorg/hrx-weight-residency to hrx-graph-develop-v2 November 3, 2026 14:49
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