Skip to content

fix(arch-llama): unify Site A attention dispatch + expand llama HFQ KV-mode policy - #465

Closed
fivetide wants to merge 3 commits into
warpfront:masterfrom
fivetide:fix/llama-site-a-kv-dispatch
Closed

fivetide wants to merge 3 commits into
warpfront:masterfrom
fivetide:fix/llama-site-a-kv-dispatch

Conversation

@fivetide

Copy link
Copy Markdown
Collaborator

Summary

Deferred items from PR #463 (transparent loading). Three commits:

1. fix(arch-llama): unify Site A attention dispatch (arch.rs:242)

Llama::forward_scratch_layers in hipfire-arch-llama had a hardcoded Q8-shaped KvTierInputs literal (quant_asym*: false, v_mode_bits: 8, givens_cos/sin: None). For asym or HFQ8 KV caches this would either panic or silently select the wrong kernel.

Changes in crates/hipfire-arch-llama/src/arch.rs:

  • Replace the literal with ..kv_cache.tier_inputs() struct update (all quantization flags now come from cache state)
  • Thread givens_cos/sin from the cache into AttnParams (asym attend kernels unwrap these)
  • Pass flash_partials: Some(&scratch.attn_partials) (required for flash attend tiles)
  • Add conditional (k_scales, v_scales) for HFQ8 (attend kernel panics on None for that tier)
  • Use gpu.graphs.capture_mode instead of hardcoded false

Byte-identical no-op for Q8 caches.

2. test(dispatch-tests): KV-tier dispatch sweep

9 GPU-free tier-routing tests in hipfire-dispatch-tests/src/llama.rs covering Q8, asym2/3/4, FWHT+asym4, HFQ4, HFQ8, Q4, F32. Each verifies uses_givens and the paired write_key/attend_key.

3. feat(kv-mode): expand LLAMA_HFQ_POLICY + wire kv_mode_override

  • kv_mode.rs: LLAMA_HFQ_POLICY now accepts [Q8, Asym3, Asym4] (was [Q8] only). Asym2/Fwht* have no Flat KV constructor and silently fall to Q8.
  • llama.rs::KvCache::from_mode: new clean-error gate for Asym3 alongside the existing Asym4 gate — at head_dim≠256 returns a user-readable error instead of panicking in new_gpu_asym3_capped.
  • carrier.rs::load_bundle: was hardcoding "" for kv_mode_override, silently ignoring explicit --kv-mode requests for HFQ llama loads since the carrier was introduced. Now passes ctx.kv_mode_override.

Validated sweep on qwen3-0.6b-llama.mq4 (head_dim=128, gfx1151):

--kv-mode Result
q8 LOADED, coherent output
asym3 Clean error: "asym3 KV cache requires head_dim=256"
asym4 Clean error: "asym4 KV cache requires head_dim=256"
asym2 LOADED → Q8 (silently downgraded, no Flat constructor)

Test plan

  • cargo test -p hipfire-dispatch-tests — 9 new tier tests pass
  • cargo test -p hipfire-runtime --lib -- kv_mode — truth-table tests pass
  • Q8 coherence: qwen3-0.6b-llama.mq4 with --kv-mode q8 produces coherent output
  • Asym gate: --kv-mode asym3 on head_dim=128 model returns clean load error

🤖 Generated with Claude Code

fivetide and others added 3 commits June 28, 2026 08:44
Replace the hardcoded Q8-shaped KvTierInputs literal (quant_asym*:
false, v_mode_bits: 8) with kv_cache.tier_inputs() so asym, FWHT,
HFQ8, and all other KV tiers are dispatched correctly.

Thread givens_cos/sin from the cache into AttnParams (asym attend
kernels unwrap these), flash_partials from ForwardScratch (required
for flash attend tiles), and k_scales/v_scales conditionally for
HFQ8 (attend.rs panics on None for that tier). Also use
gpu.graphs.capture_mode instead of the hardcoded false.

Byte-identical no-op for Q8 caches; prevents panics/garbage when the
arch-llama modular forward is wired to asym or HFQ8 KV caches.

Deferred item from PR warpfront#463.

Co-Authored-By: Claude Sonnet 4.6 <noreply@anthropic.com>
9 GPU-free tier-routing tests covering Q8, asym2/3/4, FWHT-asym4,
HFQ4, HFQ8, Q4, and F32. Each constructs the KvTierInputs that
tier_inputs() now produces for that cache mode and verifies:
  - uses_givens is true iff rotation buffers are needed (asym/fwht)
  - write_key and attend_key select the correct paired kernel tier

Validates the dispatch path that arch.rs Site A (hipfire-arch-llama)
now exercises via kv_cache.tier_inputs() instead of the hardcoded
Q8-shaped literal.

Co-Authored-By: Claude Sonnet 4.6 <noreply@anthropic.com>
…d_dim gates

Three changes:

1. llama.rs::KvCache::from_mode: add a clean-error gate for Asym3 alongside
   the existing Asym4 gate. Both require head_dim=256 (validated on Qwen 3.5);
   at head_dim=128 the constructors assert and the flash-decode paths are
   untested. Previously asym3 would panic in new_gpu_asym3_capped; now it
   returns a user-readable HipError.

2. kv_mode.rs::LLAMA_HFQ_POLICY: expand accepted from [Q8] to [Q8, Asym3,
   Asym4]. Asym2/Fwht* have no Flat KV constructor and silently fall to Q8
   (unimplemented, not broken). Truth-table test updated.

3. carrier.rs::load_bundle: pass ctx.kv_mode_override to resolve() instead
   of hardcoding "". The override was being silently ignored since the carrier
   was created — explicit --kv-mode requests had no effect on HFQ llama loads.

Result: --kv-mode asym3/asym4 on a head_dim=128 HFQ llama model now gives
a clean "requires head_dim=256" error instead of either silent Q8 or a
constructor panic. At head_dim=256, asym3/asym4 load and run correctly.

Co-Authored-By: Claude Sonnet 4.6 <noreply@anthropic.com>
@fivetide
fivetide force-pushed the fix/llama-site-a-kv-dispatch branch from 32a93b0 to 00271a9 Compare June 28, 2026 06:47
@fivetide

Copy link
Copy Markdown
Collaborator Author

Rebased onto current master (d44f89e7), conflict resolved.

Conflict was entirely stale-base noise. The branch forked at 2026-06-21, before the large KV/dispatch churn from the merged speculator-abstraction PR (#477), so it carried duplicate loader/tui changes that now collide. The PR's real work is just 3 commits / 5 files, none of which touch the conflicting files — so I cherry-picked the 3 real commits onto fresh master (all applied cleanly) and dropped the 4 trailing style: rustfmt/clippy commits and the stale base. Net diff: 214 insertions / 45 deletions across arch.rs, carrier.rs, kv_mode.rs, llama.rs, and the new dispatch-tests sweep.

Still needed on master: I confirmed master's Site A is unfixed — crates/hipfire-arch-llama/src/arch.rs still has givens_cos/sin: None and the hardcoded Q8-shaped KvTierInputs literal. This PR replaces that with kv_cache.tier_inputs() + threaded givens/scales/flash_partials, which is byte-identical for Q8 caches and only changes the asym/HFQ8 paths (which currently panic).

Validation (CPU):

  • New Site-A KV-tier dispatch sweep (added by this PR): 16/16 pass — incl. llama_kv_tier_asym{2,3,4}_uses_givens, asym4_fwht_uses_givens (the fix), and q8_uses_q8_kernels_no_givens (byte-identical common path).
  • kv_mode truth-table tests: 7/7 pass, incl. truth_table_llama_hfq_expanded (the asym3/asym4 LLAMA_HFQ_POLICY change).
  • Full workspace build (lib + bins + examples): clean.

Scope caveat: validation is the targeted dispatch/policy unit sweep + byte-identical-Q8 design — I did not run an end-to-end GPU coherence pass with asym KV on a llama-Dir model this round (the path the fix actually enables). Recommend that as the pre-merge gate, since it's the surface my notes previously flagged as panic/garbage-prone. The Q8 default path is byte-identical and unit-confirmed, so the common case carries no regression risk.

Now MERGEABLE.

@fivetide

fivetide commented Jul 7, 2026

Copy link
Copy Markdown
Collaborator Author

✅ Validated in a local integration stack (#497→#473→#495→#468→#466→#465) on gfx1151 — cargo build --release --workspace --locked green; scripts/coherence-gate.sh correctness PASS (no hard errors) and scripts/qwen2-reset-gate.sh PASS (0 ceiling resets, 0 panics). Merges clean onto current master with no conflicts. (The pflash perf stage false-fails on this box: gfx1100 baseline vs gfx1151 — untouched path.)

Copy link
Copy Markdown
Collaborator

Folded into #534 as a cleanly rebased/cherry-picked delta on beta for v0.3.0, preserving the newer Redline and modular-runtime line. Closing this source PR in favor of the consolidated release PR. Thank you for the contribution.

fivetide pushed a commit to fivetide/hipfire that referenced this pull request Aug 29, 2026
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.

2 participants