From de2448f07475c6a89aef881af48c838adf85e6be Mon Sep 17 00:00:00 2001 From: alpineq Date: Thu, 10 Sep 2026 10:58:27 +0300 Subject: [PATCH] perf(gfx1100): share Q8 KV scans across verifier rows Base: beta b8092f7c7fe0eb3dabccc28e8993ee10c3465fc6. Add wave32 R4/R8 Q8 flash-attention entries for head dimensions 128 and 256. The production route is fail-closed to exact gfx1100, sequential non-tree batches of 4..32 rows, Q8 KV, logical context >4096, and graph capture off. Admission/launcher drift falls back to the established batched Attend step; projection GEMMs remain batched. Canonical XT fresh-process A/B on RX 7900 XTX, HIP 7.2.53211-9999: target SHA256 9f91556f7e0431a077d03756a7102d0154108757289e6e5fe9a2d204c0c9eeb7, draft SHA256 d0a74a232a0e2166d889f823e91e0fbf778d21dd9668d7de055cdecb065401bc, prompt MD5 b4d0b63cddcac872648ddf3cdd92cac2, daemon MD5 512fccca7c7189559048a7aba17cb6c1. Batched samples 33.4/32.6/33.4 tok/s (median 33.4); multi-row 46.4/42.9/46.3 (median 46.3, +38.6%). All six runs: 200 tokens, tau 1.88, 69 cycles, output MD5 b501ab0e0102889bd63537f2006d4f61. Kernel oracle at hd256: R4/R8 speedups 0.98x/0.72x at 2048, 1.45x/2.01x at 4096, 1.99x/2.40x at 20676, and 1.94x/2.28x at 32768. Worst relative error 4.222e-7. ISA: R4 106 VGPR/41 SGPR; R8 186/58; no spills or private scratch. Validation: test_kernels 16/16; hipfire-arch-qwen35 193 passed; rdna-compute 242 passed; canonical XT serve battery 5/5 coherent; maps/env-doc/rustfmt/fmt-bomb/diff checks pass. The Redline PM4 arm remains blocked identically on base and candidate by the pre-existing gfx1100 private=32 scratch refusal in gemv_mq4g256v2_residual; this route is excluded during capture. Full fixture-bound evidence is recorded in docs/perf-checkpoints/2026-09-10-gfx1100-qwen38-multirow-verifier.md. --- .gitattributes | 2 + .../qwen38_issue693_longcode_20676.txt | 1974 +++++++++++++++++ crates/hipfire-arch-qwen35/map.md | 6 +- .../hipfire-arch-qwen35/src/qwen35/forward.rs | 1 + .../hipfire-arch-qwen35/src/qwen35/prefill.rs | 479 ++-- crates/rdna-compute/Cargo.toml | 4 + .../rdna-compute/examples/bench_flash_rows.rs | 173 ++ crates/rdna-compute/map.md | 10 +- crates/rdna-compute/src/attention.rs | 178 ++ crates/rdna-compute/src/kernels.rs | 2 + docs/env-vars.md | 3 +- ...-09-10-gfx1100-qwen38-multirow-verifier.md | 108 + .../src/attention_flash_q8_0_tile_rows.hip | 197 ++ 13 files changed, 3002 insertions(+), 135 deletions(-) create mode 100644 .gitattributes create mode 100644 benchmarks/prompts/qwen38_issue693_longcode_20676.txt create mode 100644 crates/rdna-compute/examples/bench_flash_rows.rs create mode 100644 docs/perf-checkpoints/2026-09-10-gfx1100-qwen38-multirow-verifier.md create mode 100644 kernels/src/attention_flash_q8_0_tile_rows.hip diff --git a/.gitattributes b/.gitattributes new file mode 100644 index 000000000..0a6e3513b --- /dev/null +++ b/.gitattributes @@ -0,0 +1,2 @@ +# Preserve byte-identical benchmark prompts whose whitespace is part of the fixture. +benchmarks/prompts/qwen38_issue693_longcode_20676.txt whitespace=-trailing-space diff --git a/benchmarks/prompts/qwen38_issue693_longcode_20676.txt b/benchmarks/prompts/qwen38_issue693_longcode_20676.txt new file mode 100644 index 000000000..a33531036 --- /dev/null +++ b/benchmarks/prompts/qwen38_issue693_longcode_20676.txt @@ -0,0 +1,1974 @@ +[issue693] Summarize this code. + + +// ==== admission.rs ==== +// SPDX-License-Identifier: Apache-2.0 +// Copyright (c) 2026 Nick Woolmer +// hipfire — see LICENSE and NOTICE in the project root. +// +// AdmissionController — decides whether a session can be admitted, and with how +// much context. +// +// In the test harnesses `kv_slots::preflight_alloc` is what stops an oversized +// configuration. In the daemon that job is HERE. The difference matters: on this +// hardware the GPU allocates from system RAM and the cgroup does NOT contain +// amdgpu GTT, so a wrong decision here does not fail a request — it takes down +// the user's desktop with a global OOM. + +/// What one loaded model costs, split into the part charged once and the part +/// charged per session. +#[derive(Debug, Clone, Copy)] +pub struct ModelFootprint { + /// Charged ONCE, however many sessions are admitted. + pub weights_bytes: u64, + /// Charged per session, per token of granted context. + pub kv_bytes_per_token: u64, +} + +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +pub enum AdmitError { + PoolFull, + WouldExceedBudget { need: u64, available: u64 }, +} + +impl std::fmt::Display for AdmitError { + fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + let gib = |b: u64| b as f64 / 1073741824.0; + match self { + AdmitError::PoolFull => write!(f, "no free slot"), + AdmitError::WouldExceedBudget { need, available } => write!( + f, + "needs {:.2} GiB but only {:.2} GiB of the budget remains", + gib(*need), + gib(*available) + ), + } + } +} + +pub struct AdmissionController { + footprint: ModelFootprint, + budget_bytes: u64, + /// Granted context per admitted session, in tokens. + admitted: Vec, + /// Host-tier budget for swapped-out snapshots. Separate from the VRAM + /// budget: admission is the production memory gate for BOTH, because the + /// control group does not contain amdgpu GTT. + host_budget: u64, + host_used: u64, +} + +impl AdmissionController { + pub fn new(footprint: ModelFootprint, budget_bytes: u64) -> Self { + Self { + footprint, + budget_bytes, + admitted: Vec::new(), + host_budget: crate::swap::DEFAULT_HOST_BUDGET_BYTES, + host_used: 0, + } + } + + /// Bytes currently committed: weights once (if anything is admitted) plus + /// each session's KV. + pub fn used_bytes(&self) -> u64 { + if self.admitted.is_empty() { + return 0; + } + let kv: u64 = self + .admitted + .iter() + .map(|&ctx| ctx as u64 * self.footprint.kv_bytes_per_token) + .sum(); + self.footprint.weights_bytes + kv + } + + /// Admit a session at `requested_ctx` tokens, or explain why not. + /// + /// Rejects rather than silently capping: a caller that asked for 128K and + /// silently got 8K would produce baffling truncation far from here. + pub fn admit(&mut self, requested_ctx: usize) -> Result { + let kv_need = requested_ctx as u64 * self.footprint.kv_bytes_per_token; + // Weights are charged once, on the first admission. + let weights_need = if self.admitted.is_empty() { + self.footprint.weights_bytes + } else { + 0 + }; + let need = kv_need + weights_need; + let available = self.budget_bytes.saturating_sub(self.used_bytes()); + // >= rather than >: an admission that would consume the LAST byte of + // budget is refused too, not just one that overflows it. On this + // hardware (no swap, cgroup does not contain amdgpu GTT) landing + // exactly on the edge leaves zero headroom for anything else running + // on the box, so it is treated the same as exceeding the budget. + if need >= available { + return Err(AdmitError::WouldExceedBudget { need, available }); + } + self.admitted.push(requested_ctx); + Ok(requested_ctx) + } + + /// Return a session's context allowance to the budget. + /// Reserve host-tier bytes for a swapped-out session. Returns false when + /// the budget cannot cover it, in which case the caller spills to disk + /// rather than exceeding the budget. + pub fn admit_host(&mut self, bytes: u64) -> bool { + if self.host_used.saturating_add(bytes) > self.host_budget { + return false; + } + self.host_used += bytes; + true + } + + pub fn release_host(&mut self, bytes: u64) { + self.host_used = self.host_used.saturating_sub(bytes); + } + + pub fn host_used_bytes(&self) -> u64 { + self.host_used + } + + pub fn host_budget_bytes(&self) -> u64 { + self.host_budget + } + + /// Set the host-tier budget. Defaults to `DEFAULT_HOST_BUDGET_BYTES`. + pub fn set_host_budget(&mut self, bytes: u64) { + self.host_budget = bytes; + } + + pub fn release(&mut self, granted_ctx: usize) { + if let Some(i) = self.admitted.iter().position(|&c| c == granted_ctx) { + self.admitted.remove(i); + } + } +} + +#[cfg(test)] +mod tests { + use super::*; + + const GIB: u64 = 1024 * 1024 * 1024; + + /// qwen3.6:27b — 15.0 GB of weights, 34 KB of KV per token. + fn f27b() -> ModelFootprint { + ModelFootprint { + weights_bytes: 15 * GIB, + kv_bytes_per_token: 34 * 1024, + } + } + + /// qwen3.6:35b-a3b — ~20 GB of weights, 10.6 KB of KV per token. + fn f35b() -> ModelFootprint { + ModelFootprint { + weights_bytes: 20 * GIB, + kv_bytes_per_token: 10_854, + } + } + + #[test] + fn weights_are_charged_once_not_per_session() { + let mut a = AdmissionController::new(f27b(), 32 * GIB); + a.admit(1024).unwrap(); + let after_one = a.used_bytes(); + a.admit(1024).unwrap(); + let after_two = a.used_bytes(); + // The second session adds only its KV, never another copy of the weights. + assert!(after_two - after_one < GIB, "weights charged twice"); + assert!(after_one >= 15 * GIB, "weights not charged at all"); + } + + #[test] + fn the_27b_cannot_take_four_agents_at_128k() { + // 15 GB + 4 x 4.25 GB = 32.25 GB against a 32 GB card. + let mut a = AdmissionController::new(f27b(), 32 * GIB); + for _ in 0..3 { + a.admit(128 * 1024).expect("first three must fit"); + } + let e = a.admit(128 * 1024).unwrap_err(); + assert!( + matches!(e, AdmitError::WouldExceedBudget { .. }), + "got {e:?}" + ); + } + + #[test] + fn the_27b_does_take_four_agents_at_96k() { + let mut a = AdmissionController::new(f27b(), 32 * GIB); + for i in 0..4 { + a.admit(96 * 1024) + .unwrap_or_else(|e| panic!("agent {i} rejected: {e:?}")); + } + } + + #[test] + fn the_35b_does_take_four_agents_at_128k() { + let mut a = AdmissionController::new(f35b(), 32 * GIB); + for i in 0..4 { + a.admit(128 * 1024) + .unwrap_or_else(|e| panic!("agent {i} rejected: {e:?}")); + } + } + + #[test] + fn release_returns_budget_so_a_later_session_fits() { + let mut a = AdmissionController::new(f27b(), 32 * GIB); + for _ in 0..3 { + a.admit(128 * 1024).unwrap(); + } + assert!(a.admit(128 * 1024).is_err()); + a.release(128 * 1024); + a.admit(128 * 1024) + .expect("budget must be reusable after release"); + } + + #[test] + fn rejection_reports_the_numbers_not_just_a_failure() { + let mut a = AdmissionController::new(f27b(), 32 * GIB); + for _ in 0..3 { + a.admit(128 * 1024).unwrap(); + } + match a.admit(128 * 1024).unwrap_err() { + AdmitError::WouldExceedBudget { need, available } => { + // `>=`, not `>`. Zero headroom is a rejection: 15 GiB of weights + // plus 4 x 4.25 GiB of KV is an EXACT tie with a 32 GiB budget, + // and a card with nothing left for activations, scratch and + // driver overhead does not fit the workload. The plan's comment + // claiming 32.25 GB was wrong -- 34 * 1024 IS the real per-token + // cost and the sum lands exactly on the budget. + assert!( + need >= available, + "need {need} should be at least available {available}" + ); + assert!(available < 32 * GIB); + } + other => panic!("expected a budget rejection, got {other:?}"), + } + } + + #[test] + fn a_single_session_over_budget_is_rejected_not_silently_capped() { + // One agent asking for more than the whole card can hold. + let mut a = AdmissionController::new(f27b(), 32 * GIB); + assert!( + a.admit(2 * 1024 * 1024).is_err(), + "must reject, not silently truncate" + ); + } + + #[test] + fn the_host_tier_has_its_own_budget() { + let mut a = AdmissionController::new( + ModelFootprint { + weights_bytes: 0, + kv_bytes_per_token: 0, + }, + 1 << 30, + ); + a.set_host_budget(1000); + assert!(a.admit_host(600)); + assert_eq!(a.host_used_bytes(), 600); + assert!( + !a.admit_host(600), + "the second must not fit; the caller spills to disk instead" + ); + assert_eq!(a.host_used_bytes(), 600, "a refused admit reserves nothing"); + a.release_host(600); + assert_eq!(a.host_used_bytes(), 0); + assert!(a.admit_host(600), "released budget must be reusable"); + } +} + +// ==== arch.rs ==== +// SPDX-License-Identifier: Apache-2.0 +// Copyright (c) 2026 Kaden Schutt +// hipfire — see LICENSE and NOTICE in the project root. + +//! ## Status: intra-crate helper, NOT the architecture contract +//! +//! This trait once competed with `hipfire_loader::Carrier` to be "the" arch +//! contract. It no longer does, and the distinction matters when adding a model: +//! +//! - **The contract** is `Carrier` (registration + load, in the loader because +//! `Carrier::load` returns `LoadedModel`) plus +//! [`crate::arch_model::ArchModel`] (the arch-agnostic view of a loaded model, +//! implemented in the arch crate). +//! - **This trait** is a typed bring-up convenience — associated +//! `Config`/`Weights`/`State` plus `config_from_hfq` / `load_weights` / +//! `new_state`. It is used only *within* arch crates, by their own +//! `load__bundle` functions. Measured: zero consumers in +//! `hipfire-loader`, `hipfire-generate` or `hipfire-daemon`. +//! +//! It is therefore optional. `hipfire-arch-muse-glimmer` implements it not at +//! all and loads fine. Adopt it if the typed shape helps your crate; skip it if +//! it does not. Its four override hooks were deleted as dead in an earlier pass. +//! +//! The bring-up contract for a hipfire architecture. Implement this +//! trait in your arch crate (e.g. `hipfire-arch-qwen35`) to plug a +//! model into the runtime. Generation, sampling, eviction, spec +//! decode, paging, prompt framing, and EOS filtering all live in +//! the runtime crate; the arch contributes only the model-specific +//! pieces. +//! +//! Default impls cover the Qwen3.5 family conventions. Override only +//! what diverges for your arch. +//! +//! # Worked examples +//! +//! - `crates/hipfire-arch-toy/` — minimum-viable stub, ~50 lines of +//! trait-impl with explanatory comments. Copy-paste this directory +//! as a starting point for a new arch. +//! - `crates/hipfire-arch-qwen35/src/arch.rs` — full production impl +//! for the Qwen3.5 hybrid DeltaNet + MoE family. Read this for the +//! bar: how `config_from_hfq` walks the JSON metadata, how +//! `load_weights` drives the weight pager, how `new_state` allocates +//! GPU scratch. +//! - `crates/hipfire-arch-llama/src/arch.rs` — second impl, dense +//! LLaMA / Mistral / plain-Qwen3 family. Demonstrates the trait at +//! facade-stage (forward body still in `hipfire-runtime::llama`, +//! PR 14 will physically split). +//! +//! # Why forward isn't on the trait +//! +//! Forward-pass dispatch is intentionally NOT routed through this +//! trait. Reasons: +//! 1. Forward signatures vary heavily across arches (number of +//! buffers, KV layout, hybrid-vs-dense paths, vision conditioning, +//! MoE expert management). Forcing one trait shape would either +//! bloat the contract or hide essential parameters behind opaque +//! slots. +//! 2. Forward dispatch is hot-path. Static dispatch via concrete-type +//! function calls keeps the call graph fully inlinable; dyn-trait +//! dispatch in the inner loop costs measurable tok/s on small +//! models. +//! 3. The trait's job is BRING-UP scaffolding (load → instantiate → +//! generation-loop wiring), not runtime polymorphism. Once an arch +//! is loaded, the daemon/CLI knows the concrete type at compile +//! time. + +use crate::hfq::HfqFile; +use crate::llama::WeightTensor; +use rdna_compute::{DType, Gpu}; + +/// Bring-up contract for a hipfire architecture. +/// +/// Implementors live in their own arch crate (`hipfire-arch-`) +/// and provide the three required types (Config / Weights / State) +/// plus five required methods. The optional override hook lets +/// an arch deviate from Qwen3.5 family defaults without growing a +/// per-`arch_id` `match` ladder in the daemon. +/// +/// # Required: associated types +/// +/// - `Config` — model-shape constants parsed from HFQ metadata. +/// Cheap to clone, sent across threads. Example: `Qwen35Config` +/// in `hipfire-arch-qwen35` carries dim, n_layers, head counts, +/// MoE topology, RoPE params. +/// - `Weights` — GPU-resident model weights. Owns `WeightTensor` +/// handles plus any host-side metadata for the weight pager. +/// - `State` — GPU-resident per-decode scratch (KV cache, attention +/// workspace, recurrent state for hybrid archs). +/// +/// # Required: methods +/// +/// See per-method docs below. +/// +/// # Optional: override hook +/// +/// `eos_filter_overrides`. Default impl matches Qwen3.5 conventions. +/// Override per-arch when the arch's end-of-turn markers diverge. +pub trait Architecture: Send + 'static { + type Weights; + type State; + type Config: Clone + Send + 'static; + + /// Canonical arch_id marker for this family. Existing IDs: + /// 0 = LLaMA / Mistral, 1 = plain Qwen3 / Qwen2, + /// 5 = Qwen3.5 dense, 6 = Qwen3.5/3.6 MoE. + /// + /// The actual id loaded at runtime is `HfqFile::arch_id` and may + /// differ from this canonical marker for families that span + /// multiple ids (e.g. `Llama::arch_id() == 0` but covers both 0 + /// and 1; the dense-vs-Qwen3-norm distinction is read off the HFQ + /// metadata inside `config_from_hfq`). + fn arch_id() -> u32; + + /// Human-readable arch tag for logs and CLI dispatch (e.g. `"qwen35"`, + /// `"llama"`). + fn name() -> &'static str; + + /// Parse model-shape constants out of `hfq.metadata_json`. + /// + /// Returns a typed `Config` or an error string. Implementations + /// generally use `serde_json` to walk the metadata blob and branch + /// on `hfq.arch_id` for variants within the family (e.g. dense vs + /// MoE, with-vs-without DeltaNet). + /// + /// # Worked example: Qwen3.5 + /// + /// `hipfire_arch_qwen35::qwen35::config_from_hfq` parses the + /// metadata, branches `arch_id == 5` (dense) vs `arch_id == 6` + /// (MoE) for expert-count fields, fills defaults for missing + /// keys (e.g. `partial_rotary_factor`), and returns a + /// `Qwen35Config` with the full per-layer shape. + fn config_from_hfq(hfq: &HfqFile) -> Result; + + /// Load model weights from an HFQ file into GPU memory. + /// + /// PR 8 note: signature changed from `&mut HfqFile` (PR 7 + /// scaffold) to `&HfqFile`. The mmap-backed HfqFile is read-only + /// at the syscall level and Qwen35::load_weights only reads + /// tensor data. Weight-pager state mutations happen on the + /// returned Weights object via interior mutability + /// (`RefCell`), not on the file. + /// + /// # Worked example: Qwen3.5 + /// + /// `hipfire_arch_qwen35::qwen35::load_weights` walks every layer's + /// QKV / output / FFN / norm tensors, hands each to + /// `WeightTensor::from_hfq_tensor` (which dispatches on the + /// HFQ quant_type to upload Q4F16G64 / F16 / F32 to GPU), and + /// assembles per-layer `LayerWeights` arrays. The weight pager + /// (lazy load + LRU eviction for >VRAM models) is wired through + /// `WeightTensor` and is not arch-specific. + fn load_weights( + hfq: &mut HfqFile, + cfg: &Self::Config, + gpu: &mut Gpu, + ) -> Result; + + /// Allocate per-decode GPU scratch for this arch. + /// + /// Returns the `State` object the daemon's generation loop holds + /// for the lifetime of a session. Sized by `cfg`. + /// + /// # Worked examples + /// + /// - Hybrid LA + FA (`DeltaNetState::new` in + /// `hipfire-arch-qwen35`) — KV cache for FA layers, recurrent + /// state buffers for DeltaNet (LA) layers, plus shared + /// attention scratch. + /// - Dense FA-only (`ForwardScratch::new` in + /// `hipfire-runtime::llama`) — KV cache plus attention + /// workspace; no recurrent state. + fn new_state(gpu: &mut Gpu, cfg: &Self::Config) -> Result; + + // Forward pass shapes are arch-specific; declare the surface but + // don't constrain types in this trait — concrete arch crates + // expose their own typed forward methods. The runtime's generic + // generation loop holds an `impl Architecture`-bound model and + // uses arch crate-specific call sites. + // + // Future PRs may tighten the forward signatures once we see what + // the qwen35 / qwen35-vl / llama splits actually need. For PR 7 + // the trait is intentionally minimal — just enough scaffolding for + // a canary arch crate to implement and the runtime to type-check. + + + /// Override EOS handling for this arch. Default uses ChatML + /// `<|im_end|>` plus the `` strip policy from runtime. + /// + /// Override to add arch-specific stop sequences (e.g. Gemma's + /// ``) and matching `holdback_prefixes` so the + /// stream doesn't leak the marker bytes to the visible output. + fn eos_filter_overrides(_cfg: &Self::Config) -> EosFilterOverrides { + EosFilterOverrides::default() + } +} + + +/// Per-arch overrides for EOS / end-of-turn filtering. +/// +/// `hipfire_runtime::eos_filter` owns visible-stream EOS detection. +/// The default implementation handles ChatML `<|im_end|>` plus +/// `` strip; per-arch overrides extend to additional markers. +#[derive(Debug, Clone, Default)] +pub struct EosFilterOverrides { + /// Byte sequences that signal end-of-turn for this arch. Streaming + /// stops (and the marker is not emitted) when the decoded byte + /// stream contains any sequence here. + /// Examples: Gemma4's `` (when forward-ported). + pub stop_at: Vec>, + /// Byte prefixes the streamer holds back until disambiguated. + /// Required so a partial decode of a `stop_at` marker doesn't leak + /// its initial bytes (e.g. holding back `` to stop or `` to + /// flush). + pub holdback_prefixes: Vec>, + /// If `Some`, override whether to strip `...` blocks + /// from the visible stream. Default is on for thinking-mode arches. + pub strip_think: Option, +} + +/// Architecture-owned iteration over weights eligible for load-time MMQ +/// safety screening. Each implementation returns `(safe, unsafe)` counts. +pub trait MmqScreenable { + fn screen_mmq_weights(&self, gpu: &mut Gpu) -> (usize, usize); +} + +/// Screen one weight tensor when its storage layout is accepted by the HFQ4 +/// MMQ reference probe. The dtype guard is load-bearing: probing a different +/// packed layout can read beyond the tensor buffer. +pub fn screen_weight_tensor( + weight: &WeightTensor, + gpu: &mut Gpu, + safe: &mut usize, + unsafe_count: &mut usize, +) { + if !matches!(weight.gpu_dtype, DType::HFQ4G256 | DType::MQ4G256) { + return; + } + if gpu.mmq_screen_weight(&weight.buf, weight.m, weight.k) { + *safe += 1; + } else { + *unsafe_count += 1; + } +} + +/// Apply the current enable/architecture policy and screen an architecture's +/// weights. Screening remains opt-in; disabled loads return immediately. +pub fn maybe_screen_mmq(weights: &impl MmqScreenable, gpu: &mut Gpu) { + if !gpu.mmq_screen.enabled + || !matches!( + gpu.arch.as_str(), + "gfx906" + | "gfx1100" + | "gfx1101" + | "gfx1102" + | "gfx1103" + | "gfx1150" + | "gfx1151" + | "gfx1152" + ) + { + return; + } + + let started = std::time::Instant::now(); + let (safe, unsafe_count) = weights.screen_mmq_weights(gpu); + eprintln!( + " MMQ screening: {safe} safe, {unsafe_count} unsafe (threshold={:.2}, {:.1}ms)", + gpu.mmq_screen.threshold, + started.elapsed().as_secs_f64() * 1000.0, + ); +} + +// ==== arch_mapping.rs ==== +// SPDX-License-Identifier: Apache-2.0 +// Copyright (c) 2026 Kaden Schutt +// Copyright (c) 2026 Nick Woolmer +// hipfire — see LICENSE and NOTICE in the project root. + +//! Single source of truth for `model_type` / `general.architecture` → `arch_id`. +//! +//! Why this table exists: three independent `model_type -> arch_id` maps +//! drifted (safetensors_source.rs, quantize/src/pipeline.rs, +//! quantize/src/pipeline_gguf.rs). One silently defaulted to llama (0) on +//! unknown input, another returned UNCLAIMED, the third lacked entries. This +//! module is the sole authority; the other sites call [`lookup_model_type`]. +//! +//! The numeric ids are the HFQ `arch_id` stamped into the file header and +//! claimed by `Carrier::claims_arch_id`. Changing any assignment is a +//! wire-format / routing break, so keep them byte-identical. + +/// Canonical `model_type` (HF) / `general.architecture` (GGUF) → `arch_id`. +/// +/// Covers the union of every string previously recognised by the three +/// consumers. Strings absent from this table are *unknown* and must fail +/// closed (not silently become llama 0). The qwen2 entry is intentionally +/// `7` (Qwen2Carrier, loads Q/K/V biases); earlier `hipfire-quantize` builds +/// mapped it to `1` (LLaMA) which dropped those biases — that was a bug and +/// is corrected here. See `safetensors_source.rs` commit 9002d7f8b. +/// +/// Sorted by `arch_id` then alphabetically for auditability. +pub const MODEL_TYPE_TO_ARCH_ID: &[(&str, u32)] = &[ + // arch 0 — llama family + ("llama", 0), + ("mistral", 0), + // arch 1 — qwen3 (llama-family loader, no bias) + ("qwen3", 1), + // arch 5 — qwen3.5 dense (qwen3.5/qwen3.6 share the same loader, 5 dense / 6 MoE) + ("qwen3.5", 5), + ("qwen3.6", 5), + ("qwen35", 5), + ("qwen3_5", 5), + ("qwen3_5_text", 5), + ("qwen3_6", 5), + // arch 5 — ornith 1.5 dense (9B). Same loader as qwen3.5 dense (5); a3b MoE variant is 6. + ("ornith", 5), + ("ornith-1.5", 5), + ("ornith1.5", 5), + ("ornith_1.5", 5), + // arch 6 — qwen3.5 MoE (explicit model_type strings; the safetensors path also + // derives 6 from has_experts==true for the qwen3.5/3.6 family) + ("qwen3_5_moe", 6), + ("qwen3_5_moe_text", 6), + ("qwen3moe", 6), + // arch 6 — ornith 1.5 MoE (35B-A3B). Mirrors registry_gen arch_id_for ornith-1.5 + a3b. + ("ornith_moe", 6), + ("ornith-1.5_moe", 6), + ("ornith1.5_moe", 6), + ("ornith_1.5_moe", 6), + ("qwen2", 7), + // arch 8 — dots.ocr + ("dots_ocr", 8), + // arch 9 — deepseek_v4 + ("deepseek_v4", 9), + // arch 10 — minimax_m2 + ("minimax_m2", 10), + // arch 11 — lfm2 (dense) + lfm2_moe (MoE); both route to hipfire-arch-lfm2moe/11 + ("lfm2", 11), + ("lfm2_moe", 11), + // lfm2_vl is the vision-language variant; it reuses the arch-11 text backbone + // (hipfire-arch-lfm2moe) plus an embedded SigLIP-2 vision tower + projector. + ("lfm2_vl", 11), + // arch 12 — cohere2_moe + ("cohere2_moe", 12), + // arch 13 — gemma4 family (dense + MoE unified; text decoder only). The + // four strings mirror pipeline.rs; gguf's old `starts_with("gemma4")` + // catch-all is intentionally replaced by this exact list so unknown + // `gemma4*` variants fail closed instead of silently becoming 13. + ("gemma4", 13), + ("gemma4_text", 13), + ("gemma4_unified", 13), + ("gemma4_unified_text", 13), + // arch 14 — muse_glimmer dense (52-layer + ViT) + ("muse_glimmer", 14), + ("muse_glimmer_text", 14), + // arch 15 — maple (Maple-Preview 20B-A1B, natively-ternary 256-expert MoE) + ("maple", 15), + // arch 22 — gemma4 EAGLE drafter (single-block spec-decode head for arch 13) + ("gemma4_unified_assistant", 22), + // arch 23 — muse_glimmer DFlash drafter + ("muse_glimmer_assistant", 23), +]; + +/// Look up an `arch_id` for a `model_type` / GGUF `general.architecture` string. +/// +/// Returns `None` for unknown inputs — callers must fail closed (error +/// naming the unrecognised string and listing `supported_model_types()`). +/// The lookup is an exact string compare; no prefix or substring fallback, +/// so a typo does not silently route to an unrelated arch. +pub fn lookup_model_type(model_type: &str) -> Option { + for (k, v) in MODEL_TYPE_TO_ARCH_ID { + if *k == model_type { + return Some(*v); + } + } + None +} + +/// Sorted list of every recognised `model_type` / architecture string, for +/// error messages. Computed from [`MODEL_TYPE_TO_ARCH_ID`] so it cannot drift. +pub fn supported_model_types() -> Vec<&'static str> { + let mut out: Vec<&'static str> = MODEL_TYPE_TO_ARCH_ID.iter().map(|(k, _)| *k).collect(); + out.sort_unstable(); + out.dedup(); + out +} + +/// Human-readable, comma-joined list for `eprintln!` diagnostics. +pub fn supported_model_types_display() -> String { + supported_model_types().join(", ") +} + +// ==== arch_model.rs ==== +// SPDX-License-Identifier: Apache-2.0 +// Copyright (c) 2026 Kaden Schutt +// hipfire — see LICENSE and NOTICE in the project root. + +//! The architecture-agnostic view of a loaded model. +//! +//! ## Why this exists +//! +//! `hipfire_loader::ModelState` is a closed enum with one variant per +//! architecture. Every crate that needed a scalar off a loaded model therefore +//! had to name all eleven — including the product entry point, which computed +//! *three integers* through a seven-arm architecture dispatch: +//! +//! ```text +//! let (dim, layers, vocab) = match m.state.as_ref() { +//! Some(ModelState::Qwen35(b)) => (b.config.dim, b.config.n_layers, b.config.vocab_size), +//! Some(ModelState::Qwen2(b)) => (b.config.hidden_size, b.config.num_hidden_layers, b.config.vocab_size), +//! … +//! ``` +//! +//! Seven arms, one tuple, and the only real difference is that some configs +//! spell it `dim`/`n_layers` and others `hidden_size`/`num_hidden_layers`. A +//! naming inconsistency was being paid for with architecture dispatch in the +//! daemon. +//! +//! Measured before this trait existed: the loader and daemon between them held +//! 93 `ModelState::` references, but touched only **seven distinct members** of +//! the bundles they unwrapped — `config` (26 hits, only ever for those three +//! scalars), `state`, `reset_session_state`, `kv_cache`/`kv`, `dn_state` (since +//! deleted as vestigial) and `weights` (free-on-unload). That is the whole +//! surface, and it is what this trait exposes. +//! +//! ## Why it lives in `hipfire-runtime` +//! +//! `hipfire-loader` depends on every `hipfire-arch-*` crate; the arch crates +//! must not depend on the loader. A trait that arch crates implement and the +//! loader consumes therefore cannot live in the loader — that is a cycle. It +//! also cannot live in `saddle-core`, which sits below the runtime and must not +//! know about `KvCache`. `hipfire-runtime` is the one layer both sides already +//! depend on, so it is where the contract belongs. +//! +//! ## What this is NOT +//! +//! Not a forward-pass abstraction. Generation stays in `hipfire-generate`, +//! which is the architecture composition root by design and legitimately names +//! arch crates. This trait exists so that *infrastructure* — load +//! acknowledgement, session reset, unload — stops branching on architecture. + +use crate::llama::KvCache; +use rdna_compute::Gpu; + +/// A loaded model, viewed without knowing its architecture. +/// +/// Implemented by each architecture's bundle type in its own crate. The loader +/// stores `Box` so that adding an architecture does not edit a +/// closed enum, and the daemon asks questions instead of matching variants. +pub trait ArchModel: Send + std::any::Any { + /// Hidden size. Spelled `dim` by some configs and `hidden_size` by others; + /// the implementor resolves that, not the caller. + fn dim(&self) -> usize; + + /// Number of decoder layers (`n_layers` / `num_hidden_layers`). + fn n_layers(&self) -> usize; + + /// Vocabulary size. + fn vocab_size(&self) -> usize; + + /// Short stable identifier, e.g. `"qwen35"`. Matches the key used by + /// [`crate::reset_core`]'s inventory so the two cannot drift. + fn arch_key(&self) -> &'static str; + + /// The model's KV cache, when it owns one directly. + /// + /// `None` is legitimate: some bundles keep the cache elsewhere, and callers + /// must treat absence as "not applicable", never as an error. + fn kv_cache_mut(&mut self) -> Option<&mut KvCache>; + + /// Drop per-session state so the next turn starts clean — recurrent state, + /// conv rings, cache offsets. Position and conversation history are the + /// caller's concern, not the model's. + /// + /// Default is a no-op because a pure-attention model with no recurrent + /// state has nothing to reset, and forcing every implementor to write an + /// empty body would obscure the ones that genuinely do work here. + fn reset_session_state(&mut self, _gpu: &mut Gpu) -> Result<(), String> { + Ok(()) + } + + /// Downcast hatch for the architecture composition root. + /// + /// `hipfire-generate` legitimately needs the concrete bundle to call a + /// per-architecture forward pass — that is what a composition root does. + /// This exists so it can keep doing that once `ModelState` is replaced by + /// `Box`. + /// + /// Crucially it borrows only the receiver, so a caller can hold the + /// downcast bundle and a disjoint `LoadedModel` field at the same time. + /// A whole-struct accessor cannot: that distinction is why the accessor + /// experiment converted 15 sites of 154 and this hatch is expected to do + /// better. + + + /// Return every GPU buffer this model owns. + /// + /// Consumes the box: unload is terminal, and taking `self` by value makes + /// use-after-free a compile error rather than a runtime one. + fn free_gpu(self: Box, gpu: &mut Gpu); +} + +#[cfg(test)] +mod tests { + use super::*; + + /// A bundle with no recurrent state and no directly-owned cache still + /// satisfies the contract without writing either method — that is the point + /// of the defaults, and a regression here would force boilerplate into + /// every pure-attention arch crate. + struct Minimal { + dim: usize, + } + + impl ArchModel for Minimal { + fn dim(&self) -> usize { + self.dim + } + fn n_layers(&self) -> usize { + 2 + } + fn vocab_size(&self) -> usize { + 32 + } + fn arch_key(&self) -> &'static str { + "minimal" + } + fn kv_cache_mut(&mut self) -> Option<&mut KvCache> { + None + } + fn free_gpu(self: Box, _gpu: &mut Gpu) {} + } + + #[test] + fn defaults_cover_a_stateless_arch() { + let m = Minimal { dim: 8 }; + assert_eq!(m.dim(), 8); + assert_eq!(m.arch_key(), "minimal"); + } + + #[test] + fn trait_is_object_safe() { + // The loader stores these behind a box; if this stops compiling the + // whole design is void, so pin it rather than discovering it later. + let m: Box = Box::new(Minimal { dim: 4 }); + assert_eq!(m.n_layers(), 2); + assert_eq!(m.vocab_size(), 32); + } +} + +// ==== arch_spec.rs ==== +//! Shared dense-transformer decode forward (N5 Phase B). +//! +//! A plain dense transformer layer is the same op sequence across arches — +//! rmsnorm-rotate + QKV, optional attention bias, optional qk-norm, RoPE, +//! attention, o_proj+residual, ffn rmsnorm-rotate + gate/up, SwiGLU, +//! down+residual. llama and qwen2 hand-rolled byte-identical copies of it +//! (qwen2 wrapped in the `SuperOp` interpreter, llama inline). This module +//! factors that body into one [`dense_forward`] driver parameterized by a few +//! config-derived [`DenseKnobs`], with the one genuinely non-shared piece — the +//! KV-cache write + attention kernel family (llama's 7-tier KV ladder vs +//! qwen2's flash/gqa selector) — left to each arch via [`DenseArch::attend`]. +//! +//! This is the "ArchSpec" authoring surface (greenfield BET 2), scoped to its +//! load/forward-time-durable core. Per the N4 review the design's `config:` +//! rows are superseded by serde `RawConfig + finalize`, so there is no config +//! schema here — each arch builds its own `Config` and derives `DenseKnobs`. +//! +//! Static dispatch only: [`dense_forward`] is generic over the concrete +//! `A: DenseArch`, so the per-token call graph stays fully inlinable (no +//! per-token `dyn`, per the forward-static rule in `arch.rs`). The driver feeds +//! the already-static `hipfire_dispatch::execute_steps`; it is not a runtime +//! op-interpreter. + +use hip_bridge::{DeviceBuffer, HipResult}; +use hipfire_dispatch::context::DispatchCtx; +use hipfire_dispatch::families::gemv::WeightRef; +use hipfire_dispatch::pipeline::{execute_steps, GemvInput, Step}; +use hipfire_dispatch::types::RotationPlan; +use rdna_compute::{Gpu, GpuTensor}; + +/// Config-derived scalars that parameterize the shared dense forward. Built once +/// per forward from each arch's finalized `Config`. +pub struct DenseKnobs { + /// Add q/k/v projection bias (qwen2: true; llama: false). + pub attn_bias: bool, + /// Apply per-head Q/K RMSNorm before RoPE (Qwen3-style; llama/qwen2: false). + pub qk_norm: bool, + pub rope_theta: f32, + pub norm_eps: f32, + pub n_heads: usize, + pub n_kv_heads: usize, + pub head_dim: usize, + pub q_dim: usize, + pub kv_dim: usize, +} + +/// Per-forward borrow of the shared decode scratch buffers. `pos_buf` is a raw +/// `DeviceBuffer` (matches both arches' scratch layout). +pub struct DenseScratch<'a> { + pub x: &'a GpuTensor, + pub tmp: &'a GpuTensor, + pub x_rot: &'a GpuTensor, + pub q: &'a GpuTensor, + pub k: &'a GpuTensor, + pub v: &'a GpuTensor, + pub attn_out: &'a GpuTensor, + pub o: &'a GpuTensor, + pub gate: &'a GpuTensor, + pub up: &'a GpuTensor, + pub ffn_hidden: &'a GpuTensor, + pub ffn_out: &'a GpuTensor, + pub pos_buf: &'a DeviceBuffer, +} + +/// Per-layer borrow of one decoder layer's weights + its derived rotation plans. +/// Bias/qk-norm tensors are `Option` (present only when the matching knob is on). +pub struct DenseLayer<'a> { + pub attn_norm: &'a GpuTensor, + pub ffn_norm: &'a GpuTensor, + pub wq: WeightRef<'a>, + pub wk: WeightRef<'a>, + pub wv: WeightRef<'a>, + pub wo: WeightRef<'a>, + pub w_gate: WeightRef<'a>, + pub w_up: WeightRef<'a>, + pub w_down: WeightRef<'a>, + pub wq_bias: Option<&'a GpuTensor>, + pub wk_bias: Option<&'a GpuTensor>, + pub wv_bias: Option<&'a GpuTensor>, + pub q_norm: Option<&'a GpuTensor>, + pub k_norm: Option<&'a GpuTensor>, + pub qkv_rot: RotationPlan, + pub ffn_rot: RotationPlan, + pub qkv_awq: Option<&'a GpuTensor>, + pub ffn_awq: Option<&'a GpuTensor>, + /// Activation dim fed to `RmsnormAutomatic` (the projection's `k`). + pub qkv_k: usize, + pub ffn_k: usize, +} + +/// A dense transformer arch expressed for the shared [`dense_forward`] driver. +/// Implementors are thin per-forward borrow wrappers over the arch's weights + +/// scratch + KV cache + config. +pub trait DenseArch { + fn n_layers(&self) -> usize; + fn knobs(&self) -> &DenseKnobs; + fn scratch(&self) -> DenseScratch<'_>; + fn layer(&self, l: usize) -> DenseLayer<'_>; + /// KV-cache write + single-token attention for layer `l`. By this point q/k/v + /// are projected, biased, qk-normed and RoPE'd into the shared scratch; write + /// the attention result into `attn_out`. This is the one op that does NOT + /// unify across arches (different KV layouts + attention kernel families). + fn attend(&self, gpu: &mut Gpu, l: usize) -> HipResult<()>; + /// Optional data-returning companion to [`attend`]: build the + /// `(KvTierPlan, AttnParams)` for layer `l` so `dense_forward` can emit a + /// first-class `Step::Attend` in one contiguous step list. Default `None` → + /// the caller keeps using the side-effecting `attend`. Only arches whose + /// attention is a `KvTierPlan` family (llama) override this; bespoke-attention + /// arches (qwen2 GQA-flash) leave it `None`. + fn attend_plan( + &self, + _l: usize, + ) -> HipResult< + Option<( + hipfire_dispatch::families::kv_tier::KvTierPlan, + hipfire_dispatch::families::attention::AttnParams<'_>, + )>, + > { + Ok(None) + } +} + +#[inline] +fn herr(e: impl std::fmt::Display) -> hip_bridge::HipError { + hip_bridge::HipError::new(0, &e.to_string()) +} + +/// Shared dense-transformer decode forward for one token. Runs the per-layer op +/// sequence; the caller does embedding (before) and final norm + lm_head + +/// sampling (after), since those buffers/dtypes differ per arch. +pub fn dense_forward(gpu: &mut Gpu, ctx: &DispatchCtx, arch: &A) -> HipResult<()> { + let k = arch.knobs(); + let s = arch.scratch(); + + for l in 0..arch.n_layers() { + let layer = arch.layer(l); + + // The attention block as one contiguous step list: QKV (fuses to + // FusedQkv*), bias, qk-norm, RoPE — then, on the `Some` path, the + // first-class `Step::Attend` + o-proj, so the whole block is one + // `execute_steps` invocation (future cross-boundary fusion seam). + // match_prefix slices each fused pattern to its own window, so the + // QKV3/Gemv fusion still fires inside the longer list. + let mut steps: Vec = vec![ + Step::RmsnormAutomatic { + x: s.x, + norm_weight: layer.attn_norm, + x_plain: s.tmp, + out: s.x_rot, + awq_scale: layer.qkv_awq, + k: layer.qkv_k, + eps: k.norm_eps, + rotation: layer.qkv_rot, + }, + Step::Gemv { + w: &layer.wq, + input: GemvInput::Prerotated(s.x_rot), + out: s.q, + }, + Step::Gemv { + w: &layer.wk, + input: GemvInput::Prerotated(s.x_rot), + out: s.k, + }, + Step::Gemv { + w: &layer.wv, + input: GemvInput::Prerotated(s.x_rot), + out: s.v, + }, + ]; + + // QKV bias (qwen2). + if k.attn_bias { + steps.push(Step::BiasAdd { + x: s.q, + bias: layer.wq_bias.expect("attn_bias: wq_bias"), + dim: k.q_dim, + }); + steps.push(Step::BiasAdd { + x: s.k, + bias: layer.wk_bias.expect("attn_bias: wk_bias"), + dim: k.kv_dim, + }); + steps.push(Step::BiasAdd { + x: s.v, + bias: layer.wv_bias.expect("attn_bias: wv_bias"), + dim: k.kv_dim, + }); + } + + // Per-head Q/K norm (Qwen3-style). + if k.qk_norm { + if let Some(qn) = layer.q_norm { + steps.push(Step::QkNorm { + x: s.q, + weight: qn, + n_groups: k.n_heads, + head_dim: k.head_dim, + eps: k.norm_eps, + }); + } + if let Some(kn) = layer.k_norm { + steps.push(Step::QkNorm { + x: s.k, + weight: kn, + n_groups: k.n_kv_heads, + head_dim: k.head_dim, + eps: k.norm_eps, + }); + } + } + + // RoPE. + steps.push(Step::Rope { + q: s.q, + k: s.k, + pos_buf: s.pos_buf, + n_heads: k.n_heads, + n_kv_heads: k.n_kv_heads, + head_dim: k.head_dim, + theta: k.rope_theta, + }); + + let o_proj = Step::GemvResidual { + w: &layer.wo, + input: GemvInput::Raw(s.attn_out), + residual: s.x, + out: s.o, + }; + match arch.attend_plan(l)? { + Some((plan, attn_io)) => { + // llama: attention is a first-class step → one contiguous list. + steps.push(Step::Attend { plan, io: attn_io }); + steps.push(o_proj); + execute_steps(gpu, ctx, &steps).map_err(herr)?; + } + None => { + // Bespoke-attention arch (qwen2 GQA-flash): keep the split — + // pre-attend steps, then the side-effecting attend, then o-proj. + // Identical kernels/order to the pre-seam path. + execute_steps(gpu, ctx, &steps).map_err(herr)?; + arch.attend(gpu, l)?; + execute_steps(gpu, ctx, &[o_proj]).map_err(herr)?; + } + } + + // FFN: rmsnorm-rotate + gate/up. + execute_steps( + gpu, + ctx, + &[ + Step::RmsnormAutomatic { + x: s.x, + norm_weight: layer.ffn_norm, + x_plain: s.tmp, + out: s.x_rot, + awq_scale: layer.ffn_awq, + k: layer.ffn_k, + eps: k.norm_eps, + rotation: layer.ffn_rot, + }, + Step::Gemv { + w: &layer.w_gate, + input: GemvInput::Prerotated(s.x_rot), + out: s.gate, + }, + Step::Gemv { + w: &layer.w_up, + input: GemvInput::Prerotated(s.x_rot), + out: s.up, + }, + ], + ) + .map_err(herr)?; + + // SwiGLU + down projection + residual. + gpu.silu_mul_f32(s.gate, s.up, s.ffn_hidden)?; + execute_steps( + gpu, + ctx, + &[Step::GemvResidual { + w: &layer.w_down, + input: GemvInput::Raw(s.ffn_hidden), + residual: s.x, + out: s.ffn_out, + }], + ) + .map_err(herr)?; + } + + Ok(()) +} + +// ==== augmentor.rs ==== +// SPDX-License-Identifier: Apache-2.0 +// Copyright (c) 2026 Kaden Schutt +// hipfire — see LICENSE and NOTICE in the project root. + +//! WeightAugmentor — a plugin interface for transparently transforming weight +//! tensors at load time. Arch crates call `load_weight()` and augmentors run +//! automatically based on the model's QuantConfig. + +use crate::llama::WeightTensor; +use crate::model_source::{ModelSource, QuantConfig}; +use hip_bridge::HipResult; +use rdna_compute::Gpu; + +// ── Trait ────────────────────────────────────────────────────────────────────── + +/// A plugin that may replace or post-process a weight tensor at load time. +/// +/// Implementors are registered in `DEFAULT_AUGMENTORS`. `load_weight()` iterates +/// the list; the first active augmentor whose `try_load` returns `Some` wins. +/// If no augmentor fires, the caller must use its own base-loading fallback. +pub trait WeightAugmentor: Send + Sync { + fn name(&self) -> &'static str; + + /// True if this augmentor applies to models with the given QuantConfig. + fn is_active_for(&self, qc: &QuantConfig) -> bool; + + /// True if this augmentor applies to the given source (delegates to + /// is_active_for if quant_config is present, otherwise false). + fn is_active(&self, source: &dyn ModelSource) -> bool { + source + .quant_config() + .map(|qc| self.is_active_for(qc)) + .unwrap_or(false) + } + + /// Attempt to fully load the weight tensor named `base_name` (no extension). + /// Returns `Ok(Some(t))` if this augmentor handles it (e.g. PaRo: reads + /// `.qweight`, `.qzeros`, etc.), `Ok(None)` to pass to the next augmentor + /// or to the base loader. + fn try_load( + &self, + source: &dyn ModelSource, + base_name: &str, + out_dim: usize, + in_dim: usize, + gpu: &mut Gpu, + ) -> HipResult>; +} + +// ── Dispatch helper ──────────────────────────────────────────────────────────── + +/// Try every active augmentor in order. Returns the first `Some(WeightTensor)` +/// found, or `None` if no augmentor handled the tensor. +/// +/// The caller is responsible for providing a fallback (standard HFQ loading or +/// error) when `None` is returned. +pub fn try_augmentors( + source: &dyn ModelSource, + base_name: &str, + out_dim: usize, + in_dim: usize, + gpu: &mut Gpu, + augmentors: &[&'static dyn WeightAugmentor], +) -> HipResult> { + for a in augmentors { + if a.is_active(source) { + if let Some(t) = a.try_load(source, base_name, out_dim, in_dim, gpu)? { + return Ok(Some(t)); + } + } + } + Ok(None) +} + +// ── ParoAugmentor ────────────────────────────────────────────────────────────── + +pub struct ParoAugmentor; + +impl ParoAugmentor { + pub fn is_active_for(qc: &QuantConfig) -> bool { + qc.method == "paroquant" && qc.krot > 0 + } +} + +impl WeightAugmentor for ParoAugmentor { + fn name(&self) -> &'static str { + "paroquant" + } + + fn is_active_for(&self, qc: &QuantConfig) -> bool { + ParoAugmentor::is_active_for(qc) + } + + fn try_load( + &self, + source: &dyn ModelSource, + base_name: &str, + out_dim: usize, + in_dim: usize, + gpu: &mut Gpu, + ) -> HipResult> { + // Only fires if the quantized tensors actually exist for this weight. + // Some tensors are excluded from quantization (router, embeddings) and + // have no .qweight — paro_load_wt falls back to .weight for those. + if source + .tensor_info(&format!("{base_name}.qweight")) + .is_none() + { + return Ok(None); + } + let qc = source + .quant_config() + .expect("ParoAugmentor: quant_config required"); + let t = crate::paro::load_paro_weight( + source, + gpu, + base_name, + out_dim, + in_dim, + qc.group_size, + qc.krot, + )?; + Ok(Some(t)) + } +} + +// ── Default registry ─────────────────────────────────────────────────────────── + +static PARO: ParoAugmentor = ParoAugmentor; + +/// Default augmentor set used by all arch crates. Extend per-arch by building +/// a custom slice: `&[DEFAULT_AUGMENTORS, &[&MyAugmentor]].concat()`. +pub static DEFAULT_AUGMENTORS: &[&dyn WeightAugmentor] = &[&PARO]; + +#[cfg(test)] +mod tests { + use super::*; + use crate::model_source::QuantConfig; + + fn make_quant_config(krot: u8) -> QuantConfig { + QuantConfig { + method: "paroquant".into(), + bits: 4, + group_size: 128, + krot, + dynamic_excludes: vec![], + } + } + + // Tests call the free function `ParoAugmentor::is_active_for(&QuantConfig)` + // which does not need a ModelSource — no mock needed for these three tests. + + #[test] + fn paro_augmentor_active_when_krot_positive() { + let qc = make_quant_config(8); + assert!(ParoAugmentor::is_active_for(&qc)); + } + + #[test] + fn paro_augmentor_inactive_when_krot_zero() { + let qc = make_quant_config(0); + assert!(!ParoAugmentor::is_active_for(&qc)); + } + + #[test] + fn paro_augmentor_inactive_for_non_paro_method() { + let mut qc = make_quant_config(8); + qc.method = "awq".into(); + assert!(!ParoAugmentor::is_active_for(&qc)); + } +} + +// ==== bf16_loader.rs ==== +//! GPTQ-target tensor-name predicate for the Tier-1 calibration path. +//! +//! The only live symbol here is [`is_gptq_target`], used by +//! `calibration.rs` (`HessianCollector`) and mirrored from +//! `scripts/collect_hessian.py` so the Tier-1 binary produces a +//! byte-compatible HFHS-v1 output with the Tier-2 Python path. +//! +//! History: this module formerly also held a `load_bf16_model` +//! safetensors-loader scaffold (`unimplemented!()`) plus its `Bf16Tensor` +//! / `TrunkBF16` metadata structs, sketched in the 2026-05-19 Tier-1 +//! foundation series. They were never wired (the imatrix/hessian work +//! moved to its own pipeline) and were removed as dead scaffold on +//! 2026-06-15. Recover from git history if a BF16 calibration loader is +//! revived. + +/// Returns true if a tensor name matches the GPTQ-target whitelist that +/// `collect_hessian` should accumulate a Hessian for. Mirrors +/// `scripts/collect_hessian.py::is_gptq_target` so the Tier 1 binary +/// produces a byte-compatible HFHS-v1 output with the Tier 2 Python +/// path. +/// +/// Whitelist (suffixes matched against the last `.`-separated segment): +/// +/// - Attention input projections: `q_proj`, `k_proj`, `v_proj`, +/// `qkv_proj` +/// - Attention output: `o_proj`, `out_proj` +/// - MLP: `gate_proj`, `up_proj`, `down_proj`, `gate_up_proj` +/// - Linear-attention (Gated DeltaNet): +/// `in_proj_qkv`, `in_proj_z`, `in_proj_a`, `in_proj_b` +/// - MoE router: `gate` +#[allow(dead_code)] +pub fn is_gptq_target(name: &str) -> bool { + const TARGETS: &[&str] = &[ + "q_proj", + "k_proj", + "v_proj", + "qkv_proj", + "o_proj", + "out_proj", + "gate_proj", + "up_proj", + "down_proj", + "gate_up_proj", + "in_proj_qkv", + "in_proj_z", + "in_proj_a", + "in_proj_b", + "gate", + ]; + // Strip a trailing `.weight` (HF safetensors stores Linear weights + // as `.weight`; the GPTQ targets are checked on the module + // name, not the parameter name). + let bare = name.strip_suffix(".weight").unwrap_or(name); + let last = bare.rsplit('.').next().unwrap_or(bare); + TARGETS.contains(&last) +} + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn gptq_target_recognizes_canonical_qwen35_names() { + assert!(is_gptq_target("model.layers.0.self_attn.q_proj.weight")); + assert!(is_gptq_target("model.layers.0.self_attn.k_proj.weight")); + assert!(is_gptq_target("model.layers.0.self_attn.v_proj.weight")); + assert!(is_gptq_target("model.layers.0.self_attn.o_proj.weight")); + assert!(is_gptq_target("model.layers.0.mlp.gate_proj.weight")); + assert!(is_gptq_target("model.layers.0.mlp.up_proj.weight")); + assert!(is_gptq_target("model.layers.0.mlp.down_proj.weight")); + } + + #[test] + fn gptq_target_recognizes_moe_router() { + // Qwen3.5-A3B MoE router lives at `model.layers.N.mlp.gate.weight` + assert!(is_gptq_target("model.layers.0.mlp.gate.weight")); + } + + #[test] + fn gptq_target_rejects_norms_and_embed() { + assert!(!is_gptq_target("model.embed_tokens.weight")); + assert!(!is_gptq_target("model.layers.0.input_layernorm.weight")); + assert!(!is_gptq_target("model.norm.weight")); + assert!(!is_gptq_target("lm_head.weight")); + } + + #[test] + fn gptq_target_recognizes_deltanet_projections() { + assert!(is_gptq_target( + "model.layers.0.linear_attn.in_proj_qkv.weight" + )); + assert!(is_gptq_target( + "model.layers.0.linear_attn.in_proj_z.weight" + )); + assert!(is_gptq_target("model.layers.0.linear_attn.out_proj.weight")); + } +} + +// ==== cache_plan.rs ==== +// SPDX-License-Identifier: Apache-2.0 +// Copyright (c) 2026 Kaden Schutt +// hipfire — see LICENSE and NOTICE in the project root. + +//! Pure, side-effect-free prompt-cache planner. +//! +//! Unifies the two per-arch LCP cache decisions that previously lived +//! inline in `daemon.rs`: +//! +//! - **qwen35**: `plan_prompt_cache` (~lines 3811-3905 of `daemon.rs`) +//! - **deepseek4**: inline LCP block (~lines 9411-9482 of `daemon.rs`) +//! +//! Call [`plan_cache`] with the appropriate [`CachePolicy`] for the arch. +//! No GPU side-effects belong here; they move into the unified loop (T4). + +/// Outcome of [`plan_cache`]: how much of the prior conversation is cached +/// and where the new prefill should start. +#[derive(Debug, PartialEq, Eq)] +pub struct CachePlan { + /// Whether this turn can reuse the previous turn's recurrent state. + pub cache_hit: bool, + /// Token index in `rendered` at which the caller should begin prefilling. + /// On a hit, `rendered[start_pos..]` is the suffix to feed to the model. + pub start_pos: usize, + /// Number of tokens from `rendered` that are already in the model's state + /// (always equals `start_pos` — kept as a separate field for call-site + /// clarity when accounting). + pub cached_tokens: usize, + /// For qwen35 resume-from-checkpoint: if `Some(p)`, the caller should + /// rewind DeltaNet state to checkpoint `p` before prefilling the suffix. + /// Always `None` for deepseek4 (no checkpoints). + pub resume_from: Option, +} + +impl CachePlan { + /// Canonical miss: cold-prefill the entire rendered conversation. + #[inline] + pub fn miss() -> Self { + CachePlan { + cache_hit: false, + start_pos: 0, + cached_tokens: 0, + resume_from: None, + } + } +} + +/// How the planner handles an exact-match render (new render == prior, byte-for-byte). +#[derive(Debug, PartialEq, Eq)] +pub enum ExactMatch { + /// qwen35: exact-match degrades to a miss (the 1-token DeltaNet + /// over-advance that would result from advancing past the last token + /// is not safe). + Miss, + /// deepseek4: step `lcp` back one so prefilling always processes ≥ 1 + /// token. The stepped-back `lcp` is then in the partial range + /// `(0, prior_len)`, which the `allow_partial=false` guard immediately + /// forces to a cold miss for DSA compressor-ring safety. + StepBack, +} + +/// Per-arch cache policy knobs. +/// +/// Construct via [`CachePolicy::qwen35`] or [`CachePolicy::deepseek4`]. +#[derive(Debug)] +pub struct CachePolicy { + /// If `true`, a miss is forced when `rendered.len() < prior.len()`. + /// + /// **deepseek4** sets this `true` for DSA compressor-ring safety: + /// `generate_deepseek4` (daemon.rs ~9413) checks + /// `prompt_ids.len() < prior.len()` before computing LCP. + /// **qwen35** sets this `false` — it has no such constraint. + pub min_new_len_ge_prior: bool, + /// What to do when the new render is byte-identical to the prior. + /// + /// See [`ExactMatch`] variants for the per-arch rationale. + pub on_exact: ExactMatch, + /// Whether a partial prefix match (`0 < lcp < prior_len`) is accepted + /// as a cache hit. + /// + /// Both current arches set this `false`. The field is part of the + /// documented policy surface so the unified loop (T4) can enable it + /// for future arches without a new planner API. + pub allow_partial: bool, +} + +impl CachePolicy { + /// Policy for **qwen35** (`plan_prompt_cache`, daemon.rs ~3811-3905). + /// + /// - No minimum-length constraint on the new render. + /// - Exact-match (`lcp == rendered.len()`) → miss (avoids DeltaNet over-advance). + /// - Partial divergence without a usable checkpoint → miss. + /// - Resume-from-checkpoint is a call-site toggle, not a policy knob. + pub fn qwen35() -> Self { + CachePolicy { + min_new_len_ge_prior: false, + on_exact: ExactMatch::Miss, + allow_partial: false, + } + } + + /// Policy for **deepseek4** (inline LCP, daemon.rs ~9411-9482). + /// + /// - Rendered must be at least as long as the prior (DSA compressor-ring safety). + /// - Exact-match (`lcp == rendered.len()`) → step back one, which then + /// falls into the partial-cold guard and becomes a miss. + /// - Any partial hit (`0 < lcp < prior_len`) → forced cold (DSA ring safety). + pub fn deepseek4() -> Self { + CachePolicy { + min_new_len_ge_prior: true, + on_exact: ExactMatch::StepBack, + allow_partial: false, + } + } +} + +/// Compute the prompt-cache plan for one turn. +/// +/// # Arguments +/// - `rendered` — the fully-rendered canonical conversation tokens for +/// this turn (already built by the caller). +/// - `prior` — `m.conversation_tokens` from the previous turn. +/// - `policy` — per-arch knobs; see [`CachePolicy::qwen35`] / +/// [`CachePolicy::deepseek4`]. +/// - `checkpoints` — ascending DeltaNet checkpoint positions +/// (`m.dflash_checkpoints`); pass `&[]` for deepseek4. +/// - `resume_enabled` — whether to attempt resume-from-checkpoint on +/// divergence; `false` for deepseek4. +/// +/// # Returns +/// A [`CachePlan`] whose `start_pos` is the index into `rendered` at which +/// the caller should begin prefilling. `cached_tokens == start_pos` always. +pub fn plan_cache( + rendered: &[u32], + prior: &[u32], + policy: &CachePolicy, + checkpoints: &[usize], + resume_enabled: bool, +) -> CachePlan { + // 1. No prior → miss. + if prior.is_empty() { + return CachePlan::miss(); + } + // 2. ds4 ring-safety: new render must be at least as long as prior. + if policy.min_new_len_ge_prior && rendered.len() < prior.len() { + return CachePlan::miss(); + } + // 3. Raw longest common prefix, bounded by both lengths. + let max_match = prior.len().min(rendered.len()); + let mut lcp = 0usize; + while lcp < max_match && prior[lcp] == rendered[lcp] { + lcp += 1; + } + // 4. Exact-match edge: lcp consumed the WHOLE new render. + if lcp == rendered.len() && lcp > 0 { + match policy.on_exact { + ExactMatch::Miss => return CachePlan::miss(), + ExactMatch::StepBack => lcp -= 1, // falls into partial-cold below + } + } + // 5. Pure forward extension → hit. + if lcp == prior.len() && lcp < rendered.len() && lcp > 0 { + return CachePlan { + cache_hit: true, + start_pos: lcp, + cached_tokens: lcp, + resume_from: None, + }; + } + // 6. Partial divergence (0 < lcp < prior_len). + if lcp > 0 && lcp < prior.len() { + if policy.allow_partial { + return CachePlan { + cache_hit: true, + start_pos: lcp, + cached_tokens: lcp, + resume_from: None, + }; + } + if resume_enabled { + if let Some(&ckpt) = checkpoints + .iter() + .filter(|&&p| p <= lcp && p < rendered.len()) + .max() + { + return CachePlan { + cache_hit: true, + start_pos: ckpt, + cached_tokens: ckpt, + resume_from: Some(ckpt), + }; + } + } + return CachePlan::miss(); + } + // 7. Otherwise miss (lcp == 0: total divergence). + CachePlan::miss() +} + +#[cfg(test)] +mod tests { + use super::*; + + fn q() -> CachePolicy { + CachePolicy::qwen35() + } + + fn d() -> CachePolicy { + CachePolicy::deepseek4() + } + + fn check(plan: CachePlan, hit: bool, start: usize, resume: Option) { + assert_eq!(plan.cache_hit, hit); + assert_eq!(plan.start_pos, start); + assert_eq!( + plan.cached_tokens, start, + "cached_tokens must equal start_pos" + ); + assert_eq!(plan.resume_from, resume); + } + + #[test] + fn t01_empty_prior_miss() { + // branch: step 1 — no prior → miss + let plan = plan_cache(&[1, 2, 3], &[], &q(), &[], false); + check(plan, false, 0, None); + } + + #[test] + fn t02_qwen35_forward_extension_hit() { + // branch: step 5 — pure forward extension + let plan = plan_cache(&[1, 2, 3, 4, 5], &[1, 2, 3], &q(), &[], false); + check(plan, true, 3, None); + } + + #[test] + fn t03_ds4_forward_extension_hit() { + // branch: step 5 — pure forward extension (ds4 policy) + let plan = plan_cache(&[1, 2, 3, 4, 5], &[1, 2, 3], &d(), &[], false); + check(plan, true, 3, None); + } + + #[test] + fn t04_qwen35_exact_match_miss() { + // branch: step 4 — exact-match → ExactMatch::Miss → miss + let plan = plan_cache(&[1, 2, 3], &[1, 2, 3], &q(), &[], false); + check(plan, false, 0, None); + } + + #[test] + fn t05_ds4_exact_match_stepback_then_partial_cold() { + // branch: step 4 → StepBack (lcp=2), then step 6 partial-cold → miss + let plan = plan_cache(&[1, 2, 3], &[1, 2, 3], &d(), &[], false); + check(plan, false, 0, None); + } + + #[test] + fn t06_qwen35_partial_no_resume_miss() { + // branch: step 6 — partial divergence, resume_enabled=false → miss + let plan = plan_cache(&[1, 2, 9, 9, 9], &[1, 2, 3, 4], &q(), &[], false); + check(plan, false, 0, None); + } + + #[test] + fn t07_qwen35_partial_resume_latest_ckpt() { + // branch: step 6 — partial divergence, resume finds latest ckpt ≤ lcp(2) + let plan = plan_cache(&[1, 2, 9, 9, 9], &[1, 2, 3, 4], &q(), &[0, 1], true); + check(plan, true, 1, Some(1)); + } + + #[test] + fn t08_qwen35_partial_resume_ckpt_beyond_lcp_miss() { + // branch: step 6 — resume_enabled but ckpt=3 > lcp(2), filtered → miss + let plan = plan_cache(&[1, 2, 9, 9, 9], &[1, 2, 3, 4], &q(), &[3], true); + check(plan, false, 0, None); + } + + #[test] + fn t09_ds4_partial_cold_miss() { + // branch: step 6 — ds4 partial → allow_partial=false, no resume → miss + let plan = plan_cache(&[1, 2, 9, 9, 9], &[1, 2, 3, 4], &d(), &[], false); + check(plan, false, 0, None); + } + + #[test] + fn t10_ds4_rendered_shorter_than_prior_miss() { + // branch: step 2 — min_new_len_ge_prior triggers miss + let plan = plan_cache(&[1, 2, 3], &[1, 2, 3, 4, 5], &d(), &[], false); + check(plan, false, 0, None); + } + + #[test] + fn t11_qwen35_rendered_shorter_full_prefix_exact_miss() { + // branch: step 4 — rendered shorter, all 3 rendered tokens match → + // lcp == rendered.len() → exact → ExactMatch::Miss + let plan = plan_cache(&[1, 2, 3], &[1, 2, 3, 4, 5], &q(), &[], false); + check(plan, false, 0, None); + } + + #[test] + fn t12_ds4_forward_extension_cached_tokens_eq_start_pos() { + // branch: step 5 — forward extension; assert cached_tokens == start_pos == 3 + let plan = plan_cache(&[1, 2, 3, 4], &[1, 2, 3], &d(), &[], false); + assert!(plan.cache_hit); + assert_eq!(plan.start_pos, 3); + assert_eq!( + plan.cached_tokens, plan.start_pos, + "cached_tokens must equal start_pos" + ); + } +} + +// ==== calibration.rs ==== +// SPDX-License-Identifier: Apache-2.0 +// hipfire — Tier-1 calibration collector (lib-ified core). +// +//! The reusable, model-agnostic calibration collector: an [`ActivationCapture`] +//! that accumulates a per-tensor GPTQ Hessian (`Σ x·xᵀ`) and imatrix diagonal +//! (`Σ x²`) on-GPU via the `calib_*_reduce_f32` kernels, and drains to HFQ +//! tensors (`.hessian` [K,K] + `.imatrix` [K]) plus an +//! internal-consistency metric (`diag(Σxxᵀ)` must equal `Σx²`). +//! +//! This is generic (hipfire-rdna + the HFQ writer only) so it sits in +//! hipfire-runtime without a cycle on the arch crates. Callers (the +//! `collect_artifacts` CLI, the daemon `Collect` op) own the forward loop + +//! the model-specific taps (MoE router histogram, KLDREF) and arm this via +//! `gpu.active_capture = Some(Arc::new(CalibCollector::default()))`. + +use crate::hfq::HfqMemTensor; +use rdna_compute::{ActivationCapture, DType, Gpu, GpuTensor}; +use std::collections::HashMap; +use std::sync::Mutex; + +fn f32_to_bf16_bits(v: f32) -> u16 { + (v.to_bits() >> 16) as u16 +} +fn bf16_bits_to_f32(bits: u16) -> f32 { + f32::from_bits((bits as u32) << 16) +} + +/// Rows buffered per tensor before flushing the outer-product. A single +/// `calib_hessian_outer_f32` over `[FLUSH_BATCH, K]` is ~FLUSH_BATCH× more +/// efficient than per-token (N=1) launches (the tiled GEMM is built for N≥16), +/// so this is the dominant calibration-throughput lever. +const FLUSH_BATCH: usize = 256; + +/// Calibration-only HFQM quant_type for compact Hessians: +/// exact F32 diagonal followed by BF16 lower strict triangle. +const QUANT_TYPE_HESSIAN_BF16_TRIL_DIAG_F32: u8 = 130; + +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +enum HessianStorage { + DenseF32, + Bf16TrilDiagF32, +} + +fn hessian_storage_from_env() -> HessianStorage { + match std::env::var("HIPFIRE_CALIB_HESSIAN_STORAGE") + .ok() + .as_deref() + .map(str::to_ascii_lowercase) + .as_deref() + { + Some("f32" | "dense-f32" | "full-f32" | "legacy") => HessianStorage::DenseF32, + _ => HessianStorage::Bf16TrilDiagF32, + } +} + +fn compact_hessian_bytes(k: usize) -> u64 { + (k * 4 + k * (k - 1)) as u64 +} + +/// Per-tensor on-GPU accumulators + a small activation row buffer. +struct Acc { + diag: GpuTensor, // [K] Σx² (imatrix) + h: Option, // [K,K] Σxxᵀ (Hessian); `None` = imatrix-only tensor + /// Host f64 reference accumulator (`Some` only under `HIPFIRE_CALIB_F64_AUDIT`). + /// The GPU outer-product accumulates `Σxxᵀ` in f32; RDNA has no f64 matrix + /// units and only ~1:16 scalar f64, so a faithful f64 reference is computed + /// CPU-side from the same staged rows. `drain` then reports the max relative + /// f32-vs-f64 divergence — measure-first before deciding whether f32 + /// accumulation needs replacing for large token counts. + h_f64: Option>, + buf: GpuTensor, // [FLUSH_BATCH, K] staged activation rows + buf_rows: usize, // rows currently staged in `buf` + k: usize, + n_tokens: u64, +} + +impl Acc { + /// Reduce the staged rows into the accumulators (one batched launch each), + /// then reset the buffer. No-op when empty. Imatrix-only tensors (`h` is + /// `None`) skip the [K,K] outer-product — this is how MoE routed experts + /// are captured: a full per-expert Hessian (256 experts × ~48 layers × + /// [K,K]) is ~196 GB and does not fit, but the imatrix (Σx², a K-vector) + /// is ~100 MB and is the importance signal AWQ-style quant needs. + fn flush(&mut self, gpu: &mut Gpu) { + if self.buf_rows == 0 { + return; + } + gpu.calib_sumsq_reduce_f32(&self.buf, &self.diag, self.buf_rows, self.k) + .unwrap(); + if let Some(h) = &self.h { + gpu.calib_hessian_outer_f32(&self.buf, h, self.buf_rows, self.k) + .unwrap(); + } + // Audit: accumulate the same rows in f64 on the CPU (no GPU f64 path). + if let Some(h_f64) = &mut self.h_f64 { + let k = self.k; + let rows = gpu + .download_f32(&self.buf) + .expect("download buf (f64 audit)"); + for r in 0..self.buf_rows { + let x = &rows[r * k..r * k + k]; + for i in 0..k { + let xi = x[i] as f64; + let hrow = &mut h_f64[i * k..i * k + k]; + for j in 0..k { + hrow[j] += xi * x[j] as f64; + } + } + } + } + self.buf_rows = 0; + } +} + +/// Unified Hessian + imatrix collector. Arm via `gpu.active_capture`. +/// +/// By default every captured tensor accumulates a full [K,K] Hessian. Tensors +/// whose canonical name contains any of `imatrix_only_substr` accumulate only +/// the imatrix (Σx²); used for MoE routed experts whose full Hessians do not +/// fit in memory (see [`Acc::flush`]). +#[derive(Default)] +pub struct CalibCollector { + accs: Mutex>, + imatrix_only_substr: Vec, + /// When set (`HIPFIRE_CALIB_F64_AUDIT=1`), also accumulate each Hessian in + /// f64 on the CPU and report the f32-vs-f64 divergence in `drain`. Opt-in, + /// slow (CPU outer-products) — a measurement tool, not the default path. + f64_audit: bool, +} + +/// `HIPFIRE_CALIB_F64_AUDIT=1` → run the CPU f64 reference accumulation. +fn f64_audit_enabled() -> bool { + std::env::var("HIPFIRE_CALIB_F64_AUDIT").ok().as_deref() == Some("1") +} + +impl CalibCollector { + pub fn new() -> Self { + Self { + accs: Mutex::new(HashMap::new()), + imatrix_only_substr: Vec::new(), + f64_audit: f64_audit_enabled(), + } + } + + /// Collector that stores imatrix-only (no [K,K] Hessian) for any tensor + /// whose name contains one of `substr` (e.g. `".experts."` for MoE). + pub fn with_imatrix_only(substr: Vec) -> Self { + Self { + accs: Mutex::new(HashMap::new()), + imatrix_only_substr: substr, + f64_audit: f64_audit_enabled(), + } + } + + fn wants_hessian(&self, name: &str) -> bool { + !self.imatrix_only_substr.iter().any(|s| name.contains(s)) + } + + /// Number of distinct tensors captured so far. + pub fn len(&self) -> usize { + self.accs.lock().unwrap().len() + } + + pub fn is_empty(&self) -> bool { + self.len() == 0 + } + + /// Per-tensor descriptors (no GPU work): `name`, whether it has a full + /// Hessian, `k`, and `n_tokens`. The caller uses these to compute counts + + /// `name -> n_tokens` provenance for the metadata BEFORE the streaming write + /// (the HFQM index/metadata must be written ahead of the payloads). + pub fn tensor_descriptors(&self) -> Vec { + let accs = self.accs.lock().unwrap(); + let mut names: Vec<&String> = accs.keys().collect(); + names.sort(); + names + .iter() + .map(|name| { + let acc = &accs[*name]; + CalibTensorDesc { + name: (*name).clone(), + has_hessian: acc.h.is_some(), + k: acc.k, + n_tokens: acc.n_tokens, + } + }) + .collect() + } + + /// Release all GPU accumulators owned by this collector. Grouped + /// calibration runs call this after streaming a part file so the next group + /// can reuse the memory instead of waiting for process teardown. + pub fn free_gpu(&self, gpu: &mut Gpu) { + let mut accs = self.accs.lock().unwrap(); + for (_, acc) in accs.drain() { + let _ = gpu.free_tensor(acc.diag); + if let Some(h) = acc.h { + let _ = gpu.free_tensor(h); + } + let _ = gpu.free_tensor(acc.buf); + } + } + + /// GuidedQuant capture: accumulate the per-token **Fisher-weighted** Hessian + /// `H̄ = Σ_n w[n]·xₙxₙᵀ` (and its diagonal) for `tensor_name`. `x` is the + /// linear's input activation `[n,k]` (a real contiguous block, not the shared + /// scratch the `ActivationCapture::capture` tap takes); `w` `[n]` is the + /// per-token weight the caller forms from that linear's output-grad `∂ℓ/∂z` + /// (see `calib_row_meansq_f32`). Unbuffered — one weighted outer-product + + /// weighted sumsq per call, fine offline. `w≡1` makes this identical to the + /// plain unweighted capture. + pub fn capture_weighted( + &self, + gpu: &mut Gpu, + tensor_name: &str, + x: &GpuTensor, + w: &GpuTensor, + n: usize, + k: usize, + ) { + let mut accs = self.accs.lock().unwrap(); + if !accs.contains_key(tensor_name) { + let diag = gpu.zeros(&[k], DType::F32).unwrap(); + let h = if self.wants_hessian(tensor_name) { + Some(gpu.zeros(&[k, k], DType::F32).unwrap()) + } else { + None + }; + // No row buffering on this path; a minimal placeholder keeps `Acc` + // uniform (`flush` is a no-op while `buf_rows == 0`). + let buf = gpu.zeros(&[1, k], DType::F32).unwrap(); + accs.insert( + tensor_name.to_string(), + Acc { + diag, + h, + h_f64: None, + buf, + buf_rows: 0, + k, + n_tokens: 0, + }, + ); + } + let acc = accs.get_mut(tensor_name).unwrap(); + gpu.calib_sumsq_weighted_f32(x, w, &acc.diag, n, k).unwrap(); + if let Some(h) = &acc.h { + gpu.calib_hessian_outer_weighted_f32(x, w, h, n, k).unwrap(); + } + acc.n_tokens += n as u64; + } + + /// Stream the accumulated tensors into an HFQM `.calib.hfq` at `path`, + /// **one tensor at a time** (download → normalize `/ n_tokens` → write → + /// drop), so peak host memory is a single Hessian rather than all of them + /// (a 9B is ~32 GB if materialized at once). `extra` holds any small + /// already-in-RAM tensors (e.g. KLDREF) the caller wants in the same + /// package. The metadata + index are written first (payload sizes are + /// deterministic from `k`), then the payloads stream. Returns the max + /// relative `diag(H)`-vs-`Σx²` consistency error. Also runs the optional + /// f64 audit (`HIPFIRE_CALIB_F64_AUDIT`) during the \ No newline at end of file diff --git a/crates/hipfire-arch-qwen35/map.md b/crates/hipfire-arch-qwen35/map.md index 5f5812333..4d3fd48d2 100644 --- a/crates/hipfire-arch-qwen35/map.md +++ b/crates/hipfire-arch-qwen35/map.md @@ -40,9 +40,9 @@ _Generated by `scripts/check-crate-maps.py` from the tree — do not edit inside | [`src/qwen35/batch.rs`](src/qwen35/batch.rs) | 1,949 | 16 | 2 | | [`src/qwen35/config.rs`](src/qwen35/config.rs) | 1,643 | 41 | 21 | | [`src/qwen35/ep_batch.rs`](src/qwen35/ep_batch.rs) | 4,805 | 20 | 7 | -| [`src/qwen35/forward.rs`](src/qwen35/forward.rs) | 6,251 | 31 | 12 | +| [`src/qwen35/forward.rs`](src/qwen35/forward.rs) | 6,252 | 31 | 12 | | [`src/qwen35/load.rs`](src/qwen35/load.rs) | 4,906 | 10 | 0 | -| [`src/qwen35/prefill.rs`](src/qwen35/prefill.rs) | 10,126 | 12 | 49 | +| [`src/qwen35/prefill.rs`](src/qwen35/prefill.rs) | 10,353 | 12 | 51 | | [`src/qwen35/weights.rs`](src/qwen35/weights.rs) | 2,020 | 43 | 10 | | [`src/qwen35.rs`](src/qwen35.rs) | 63 | 7 | 0 | | [`src/scheduler.rs`](src/scheduler.rs) | 142 | 3 | 4 | @@ -97,6 +97,6 @@ _Generated by `scripts/check-crate-maps.py` from the tree — do not edit inside ### Totals -- 29 modules · 59,444 lines · 440 public items · 196 tests · 11 examples +- 29 modules · 59,672 lines · 440 public items · 198 tests · 11 examples diff --git a/crates/hipfire-arch-qwen35/src/qwen35/forward.rs b/crates/hipfire-arch-qwen35/src/qwen35/forward.rs index 98fc14cbe..6838d85ef 100644 --- a/crates/hipfire-arch-qwen35/src/qwen35/forward.rs +++ b/crates/hipfire-arch-qwen35/src/qwen35/forward.rs @@ -4508,6 +4508,7 @@ pub fn forward_prefill_dense_tp( let ctx = DispatchCtx::new(&gpus.devices[rank]); if let Err(e) = crate::qwen35::prefill::batch_chunk_full_attn_attn( &mut gpus.devices[rank], + false, layer, &configs[rank], &pbs_vec[rank], diff --git a/crates/hipfire-arch-qwen35/src/qwen35/prefill.rs b/crates/hipfire-arch-qwen35/src/qwen35/prefill.rs index 31f4ef47f..918b294b6 100644 --- a/crates/hipfire-arch-qwen35/src/qwen35/prefill.rs +++ b/crates/hipfire-arch-qwen35/src/qwen35/prefill.rs @@ -5611,6 +5611,7 @@ fn batch_chunk_full_attn_input_projection( /// Same statements, same order, same launches as the inlined block. fn batch_chunk_full_attn_prepare( gpu: &mut Gpu, + fa_attn_multirow: bool, layer: &FullAttnLayerWeights, config: &Qwen35Config, pbs: &PrefillBatchScratch, @@ -5766,70 +5767,21 @@ fn batch_chunk_full_attn_prepare( } // 6–7. Batched KV write + flash attention (via dispatch). - let is_tree = tree_verify.is_some(); - let (block_start, block_cols) = match tree_verify.as_ref() { - Some(_) => (start_pos, n), - None => (0, 0), - }; - let tree_bias = tree_verify.as_ref().map(|c| c.attn_bias); - let plan = KvTierPlan::derive(KvTierInputs { - pos: start_pos, - flash_mode: s.flash_mode as usize, - capture_mode: gpu.graphs.capture_mode, - batch_size: n, - is_tree, - ..kv_cache.tier_inputs() - }) - .map_err(|e| HipError::new(0, &e.to_string()))?; - let io = AttnParams { - q: &pbs.fa_q_batch, - k: &pbs.fa_k_batch, - v: &pbs.fa_v_batch, - k_cache: &kv_cache.k_gpu[layer_idx], - v_cache: &kv_cache.v_gpu[layer_idx], - k_scales: None, - v_scales: None, - pos_buf: &s.pos_buf, - pos: start_pos, - positions: Some(&pbs.positions), - n_heads: config.n_heads, - n_kv_heads: config.n_kv_heads, - head_dim: config.head_dim, - physical_cap: kv_cache.physical_cap, - batch_size: n, + batch_chunk_fa_attend( + gpu, + config, + pbs, + s, + kv_cache, + n, + start_pos, max_ctx_len, - flash_partials: Some(&s.flash_partials), - givens_cos: kv_cache.givens_cos.as_ref(), - givens_sin: kv_cache.givens_sin.as_ref(), - tree_bias, - block_start, - block_cols, - output_gate: None, - output: &pbs.fa_attn_out_batch, - }; - if let BatchSemantics::Independent { - lane_capacity, - active_mask, - .. - } = batch_semantics - { - run_independent_q8_attention( - gpu, - pbs, - kv_cache, - config, - layer_idx, - n, - lane_capacity, - max_ctx_len, - active_mask, - )?; - } else if batch_semantics.is_independent() { - unreachable!("independent variant must carry active_mask"); - } else { - execute_steps(gpu, &ctx, &[Step::Attend { plan, io }]) - .map_err(|e| HipError::new(0, &e.to_string()))?; - } + ctx, + batch_semantics, + tree_verify, + layer_idx, + fa_attn_multirow, + )?; Ok(()) } @@ -5937,8 +5889,172 @@ fn batch_chunk_full_attn_output_projection( Ok(()) } +/// Context length past which an admitted gfx1100/Q8 small-batch attend step +/// leaves the batched masked FA kernel for the multi-row tile. Measured with +/// the Qwen3.8-27B verify shape (`bench_flash_rows`, tile 128): the batched +/// kernel still wins at 2k and loses from 4k on. +/// `HIPFIRE_FA_PERTOKEN_MIN_CTX` overrides; `0` disables the route. +pub(crate) fn fa_pertoken_min_ctx() -> Option { + use std::sync::OnceLock; + static MIN_CTX: OnceLock> = OnceLock::new(); + *MIN_CTX.get_or_init(|| { + let v = hipfire_config::developer_var("HIPFIRE_FA_PERTOKEN_MIN_CTX") + .ok() + .and_then(|v| v.parse::().ok()) + .unwrap_or(4_096); + (v > 0).then_some(v) + }) +} + +#[allow(clippy::too_many_arguments)] +fn q8_multirow_attn_admitted( + is_gfx1100: bool, + quant_q8: bool, + head_dim: usize, + n: usize, + logical_ctx: usize, + min_ctx: Option, + is_tree: bool, + is_independent: bool, + capture_mode: bool, +) -> bool { + is_gfx1100 + && quant_q8 + && matches!(head_dim, 128 | 256) + && (4..=32).contains(&n) + && min_ctx.is_some_and(|threshold| logical_ctx > threshold) + && !is_tree + && !is_independent + && !capture_mode +} + +#[allow(clippy::too_many_arguments)] +fn batch_chunk_fa_attend( + gpu: &mut Gpu, + config: &Qwen35Config, + pbs: &PrefillBatchScratch, + s: &Qwen35Scratch, + kv_cache: &llama::KvCache, + n: usize, + start_pos: usize, + max_ctx_len: usize, + ctx: &DispatchCtx, + batch_semantics: BatchSemantics<'_>, + tree_verify: Option>, + layer_idx: usize, + multirow: bool, +) -> HipResult<()> { + if let BatchSemantics::Independent { + lane_capacity, + active_mask, + .. + } = batch_semantics + { + return run_independent_q8_attention( + gpu, + pbs, + kv_cache, + config, + layer_idx, + n, + lane_capacity, + max_ctx_len, + active_mask, + ); + } + if batch_semantics.is_independent() { + unreachable!("independent variant must carry active_mask"); + } + + if multirow { + debug_assert!(gpu.arch_caps.is_gfx1100()); + debug_assert!(kv_cache.quant_q8); + debug_assert!(matches!(config.head_dim, 128 | 256)); + gpu.kv_cache_write_q8_0_batched( + &kv_cache.k_gpu[layer_idx], + &pbs.fa_k_batch, + &pbs.positions, + config.n_kv_heads, + config.head_dim, + n, + )?; + gpu.kv_cache_write_q8_0_batched( + &kv_cache.v_gpu[layer_idx], + &pbs.fa_v_batch, + &pbs.positions, + config.n_kv_heads, + config.head_dim, + n, + )?; + if gpu.attention_flash_q8_0_rows_masked( + &pbs.fa_q_batch, + &kv_cache.k_gpu[layer_idx], + &kv_cache.v_gpu[layer_idx], + &pbs.fa_attn_out_batch, + &pbs.positions, + config.n_heads, + config.n_kv_heads, + config.head_dim, + max_ctx_len, + n, + &s.flash_partials, + )? { + return Ok(()); + } + // Admission and launcher support intentionally duplicate the shape + // checks. If they ever drift, retain the established batched route + // below rather than silently exploding the verify block into n + // independent attention launches. + } + + let is_tree = tree_verify.is_some(); + let (block_start, block_cols) = match tree_verify.as_ref() { + Some(_) => (start_pos, n), + None => (0, 0), + }; + let tree_bias = tree_verify.as_ref().map(|c| c.attn_bias); + let plan = KvTierPlan::derive(KvTierInputs { + pos: start_pos, + flash_mode: s.flash_mode as usize, + capture_mode: gpu.graphs.capture_mode, + batch_size: n, + is_tree, + ..kv_cache.tier_inputs() + }) + .map_err(|e| HipError::new(0, &e.to_string()))?; + let io = AttnParams { + q: &pbs.fa_q_batch, + k: &pbs.fa_k_batch, + v: &pbs.fa_v_batch, + k_cache: &kv_cache.k_gpu[layer_idx], + v_cache: &kv_cache.v_gpu[layer_idx], + k_scales: None, + v_scales: None, + pos_buf: &s.pos_buf, + pos: start_pos, + positions: Some(&pbs.positions), + n_heads: config.n_heads, + n_kv_heads: config.n_kv_heads, + head_dim: config.head_dim, + physical_cap: kv_cache.physical_cap, + batch_size: n, + max_ctx_len, + flash_partials: Some(&s.flash_partials), + givens_cos: kv_cache.givens_cos.as_ref(), + givens_sin: kv_cache.givens_sin.as_ref(), + tree_bias, + block_start, + block_cols, + output_gate: None, + output: &pbs.fa_attn_out_batch, + }; + execute_steps(gpu, ctx, &[Step::Attend { plan, io }]) + .map_err(|e| HipError::new(0, &e.to_string())) +} + pub(crate) fn batch_chunk_full_attn_attn( gpu: &mut Gpu, + fa_attn_multirow: bool, layer: &FullAttnLayerWeights, config: &Qwen35Config, pbs: &PrefillBatchScratch, @@ -5967,6 +6083,7 @@ pub(crate) fn batch_chunk_full_attn_attn( batch_chunk_full_attn_prepare( gpu, + fa_attn_multirow, layer, config, pbs, @@ -7143,6 +7260,7 @@ fn batch_chunk_delta_net_moe( #[allow(clippy::too_many_arguments)] fn batch_chunk_full_attn_moe( gpu: &mut Gpu, + fa_attn_multirow: bool, layer: &FullAttnMoeLayerWeights, config: &Qwen35Config, pbs: &PrefillBatchScratch, @@ -7498,70 +7616,21 @@ fn batch_chunk_full_attn_moe( kv_cache.compact_offset as i32, )?; // Batched KV write + flash attention (via dispatch). - let is_tree = tree_verify.is_some(); - let (block_start, block_cols) = match tree_verify.as_ref() { - Some(_) => (start_pos, n), - None => (0, 0), - }; - let tree_bias = tree_verify.as_ref().map(|c| c.attn_bias); - let plan = KvTierPlan::derive(KvTierInputs { - pos: start_pos, - flash_mode: s.flash_mode as usize, - capture_mode: gpu.graphs.capture_mode, - batch_size: n, - is_tree, - ..kv_cache.tier_inputs() - }) - .map_err(|e| HipError::new(0, &e.to_string()))?; - let io = AttnParams { - q: &pbs.fa_q_batch, - k: &pbs.fa_k_batch, - v: &pbs.fa_v_batch, - k_cache: &kv_cache.k_gpu[layer_idx], - v_cache: &kv_cache.v_gpu[layer_idx], - k_scales: None, - v_scales: None, - pos_buf: &s.pos_buf, - pos: start_pos, - positions: Some(&pbs.positions), - n_heads: config.n_heads, - n_kv_heads: config.n_kv_heads, - head_dim: config.head_dim, - physical_cap: kv_cache.physical_cap, - batch_size: n, + batch_chunk_fa_attend( + gpu, + config, + pbs, + s, + kv_cache, + n, + start_pos, max_ctx_len, - flash_partials: Some(&s.flash_partials), - givens_cos: kv_cache.givens_cos.as_ref(), - givens_sin: kv_cache.givens_sin.as_ref(), - tree_bias, - block_start, - block_cols, - output_gate: None, - output: &pbs.fa_attn_out_batch, - }; - if let BatchSemantics::Independent { - lane_capacity, - active_mask, - .. - } = batch_semantics - { - run_independent_q8_attention( - gpu, - pbs, - kv_cache, - config, - layer_idx, - n, - lane_capacity, - max_ctx_len, - active_mask, - )?; - } else if batch_semantics.is_independent() { - unreachable!("independent variant must carry active_mask"); - } else { - execute_steps(gpu, &ctx, &[Step::Attend { plan, io }]) - .map_err(|e| HipError::new(0, &e.to_string()))?; - } + ctx, + batch_semantics, + tree_verify, + layer_idx, + fa_attn_multirow, + )?; gpu.sigmoid_mul_f32(&pbs.fa_attn_out_batch, &pbs.fa_gate_batch)?; // wo + residual. Mirrors the dense FA wo dispatch at // qwen35.rs:5591-5623 — Q8 wo skips rotation (un-rotated @@ -7900,6 +7969,23 @@ pub(crate) fn forward_batch_chunk_impl( } _ => true, }); + // Attention only: the batched masked FA kernel grids [n_heads, tiles, ROW] + // and re-scans the whole KV once per row, so a small verify block over a + // long context pays the scan n times (202 vs 103 ms at 33k). The layer's + // GEMMs stay batched either way — only the attend step switches to the + // multi-row tile. Its tile grid is sized from the live logical context on + // the host, so a captured replay would keep the first cycle's tile count. + let fa_attn_multirow = q8_multirow_attn_admitted( + gpu.arch_caps.is_gfx1100(), + kv_cache.quant_q8, + config.head_dim, + n, + start_pos + n, + fa_pertoken_min_ctx(), + tree_verify.is_some(), + batch_semantics.is_independent(), + gpu.graphs.capture_mode, + ); let logical_max_ctx = match batch_semantics { BatchSemantics::Sequential => start_pos + n, BatchSemantics::Independent { positions, .. } => { @@ -7971,6 +8057,7 @@ pub(crate) fn forward_batch_chunk_impl( (LayerWeights::FullAttn(layer), LayerType::FullAttention) if fa_batched_ok => { batch_chunk_full_attn_attn( gpu, + fa_attn_multirow, layer, config, pbs, @@ -8070,6 +8157,7 @@ pub(crate) fn forward_batch_chunk_impl( (LayerWeights::FullAttnMoe(layer), LayerType::FullAttention) if fa_batched_ok => { batch_chunk_full_attn_moe( gpu, + fa_attn_multirow, layer, config, pbs, @@ -8718,6 +8806,145 @@ mod tests { use hipfire_dispatch::context::DispatchWorkload; use rdna_compute::DType; + #[test] + fn q8_multirow_attn_admits_only_measured_gfx1100_shapes() { + for head_dim in [128, 256] { + for n in [4, 8, 32] { + assert!(q8_multirow_attn_admitted( + true, + true, + head_dim, + n, + 4097, + Some(4096), + false, + false, + false, + )); + } + } + } + + #[test] + fn q8_multirow_attn_rejects_unmeasured_or_unsupported_routes() { + let admitted = |is_gfx1100, + quant_q8, + head_dim, + n, + logical_ctx, + min_ctx, + is_tree, + is_independent, + capture_mode| { + q8_multirow_attn_admitted( + is_gfx1100, + quant_q8, + head_dim, + n, + logical_ctx, + min_ctx, + is_tree, + is_independent, + capture_mode, + ) + }; + assert!(!admitted( + false, + true, + 256, + 8, + 8192, + Some(4096), + false, + false, + false, + )); + assert!(!admitted( + true, + false, + 256, + 8, + 8192, + Some(4096), + false, + false, + false, + )); + for head_dim in [64, 320] { + assert!(!admitted( + true, + true, + head_dim, + 8, + 8192, + Some(4096), + false, + false, + false, + )); + } + for n in [1, 3, 33] { + assert!(!admitted( + true, + true, + 256, + n, + 8192, + Some(4096), + false, + false, + false, + )); + } + assert!(!admitted( + true, + true, + 256, + 8, + 4096, + Some(4096), + false, + false, + false, + )); + assert!(!admitted( + true, true, 256, 8, 8192, None, false, false, false, + )); + assert!(!admitted( + true, + true, + 256, + 8, + 8192, + Some(4096), + true, + false, + false, + )); + assert!(!admitted( + true, + true, + 256, + 8, + 8192, + Some(4096), + false, + true, + false, + )); + assert!(!admitted( + true, + true, + 256, + 8, + 8192, + Some(4096), + false, + false, + true, + )); + } + #[test] fn paro_batched_admit_defaults_off_and_allows_opt_in() { // PARO batched prefill is default-OFF (the path has a coherence/echo bug; diff --git a/crates/rdna-compute/Cargo.toml b/crates/rdna-compute/Cargo.toml index f4e88eeb2..9c7cfe905 100644 --- a/crates/rdna-compute/Cargo.toml +++ b/crates/rdna-compute/Cargo.toml @@ -55,6 +55,10 @@ required-features = ["lab"] name = "bench_dflash_verify_shapes" required-features = ["lab"] +[[example]] +name = "bench_flash_rows" +required-features = ["lab"] + [[example]] name = "test_mq4v2_residual_ksplit_gfx1100" required-features = ["lab"] diff --git a/crates/rdna-compute/examples/bench_flash_rows.rs b/crates/rdna-compute/examples/bench_flash_rows.rs new file mode 100644 index 000000000..b5d5c0f8b --- /dev/null +++ b/crates/rdna-compute/examples/bench_flash_rows.rs @@ -0,0 +1,173 @@ +// SPDX-License-Identifier: Apache-2.0 +// Copyright (c) 2026 Kaden Schutt +// hipfire — see LICENSE and NOTICE in the project root. + +//! Qwen3.8-27B verify shape (24 heads, hd 256, Q8 KV, 8 rows): multi-row vs +//! the existing batched route, with achieved KV bandwidth. `--check` compares the two outputs +//! and fails above a 1e-3 relative envelope (same reduce, different scan). + +use rdna_compute::{DType, Gpu}; + +fn lcg(seed: u32, n: usize) -> Vec { + let mut s = seed; + (0..n) + .map(|_| { + s = s.wrapping_mul(1_103_515_245).wrapping_add(12_345); + ((s >> 16) & 0x7fff) as f32 / 32_768.0 - 0.5 + }) + .collect() +} + +fn main() { + let args: Vec = std::env::args().collect(); + let argval = |k: &str, d: usize| { + args.iter() + .position(|a| a == k) + .map(|i| args[i + 1].parse().unwrap()) + .unwrap_or(d) + }; + let seq_len = argval("--seq", 33014); + let batch = argval("--rows", 8); + let iters = argval("--iters", 50); + let n_heads = argval("--heads", 24); + let n_kv_heads = argval("--kv-heads", 4); + let head_dim = argval("--head-dim", 256); + + let max_seq = 65536usize; + let q_dim = n_heads * head_dim; + let kv_dim = n_kv_heads * head_dim; + + let mut gpu = Gpu::init().expect("GPU init"); + let tile = gpu.attn_tile_size(); + eprintln!( + "GPU: {} seq={seq_len} rows={batch} heads={n_heads}/{n_kv_heads} hd={head_dim} tile={tile}", + gpu.arch + ); + + let d_q = gpu + .upload_f32(&lcg(0xa5a5, batch * q_dim), &[batch * q_dim]) + .unwrap(); + let d_kf = gpu + .upload_f32(&lcg(0xc3c3, max_seq * kv_dim), &[max_seq * kv_dim]) + .unwrap(); + let d_vf = gpu + .upload_f32(&lcg(0x9696, max_seq * kv_dim), &[max_seq * kv_dim]) + .unwrap(); + let kv_bytes_total = max_seq * n_kv_heads * (head_dim / 32) * 34; + let d_k = gpu.alloc_tensor(&[kv_bytes_total], DType::Q8_0).unwrap(); + let d_v = gpu.alloc_tensor(&[kv_bytes_total], DType::Q8_0).unwrap(); + let all_pos: Vec = (0..max_seq as i32).flat_map(|p| p.to_ne_bytes()).collect(); + let d_all_pos = gpu.alloc_tensor(&[max_seq], DType::F32).unwrap(); + gpu.hip.memcpy_htod(&d_all_pos.buf, &all_pos).unwrap(); + let mut written = 0usize; + while written < max_seq { + let chunk = (max_seq - written).min(8192); + let pos_view = d_all_pos.sub_offset(written, chunk); + let kf_view = d_kf.sub_offset(written * kv_dim, chunk * kv_dim); + let vf_view = d_vf.sub_offset(written * kv_dim, chunk * kv_dim); + gpu.kv_cache_write_q8_0_batched(&d_k, &kf_view, &pos_view, n_kv_heads, head_dim, chunk) + .unwrap(); + gpu.kv_cache_write_q8_0_batched(&d_v, &vf_view, &pos_view, n_kv_heads, head_dim, chunk) + .unwrap(); + written += chunk; + } + gpu.hip.device_synchronize().unwrap(); + + let check = args.iter().any(|a| a == "--check"); + let d_out = gpu.zeros(&[batch * q_dim], DType::F32).unwrap(); + let d_out_batched = gpu.zeros(&[batch * q_dim], DType::F32).unwrap(); + let max_tiles = seq_len.div_ceil(tile); + let d_part = gpu + .zeros( + &[16 * n_heads * (max_seq / tile) * (2 + head_dim)], + DType::F32, + ) + .unwrap(); + + let positions: Vec = (0..batch) + .flat_map(|i| ((seq_len - batch + i) as i32).to_ne_bytes()) + .collect(); + let d_pos = gpu.alloc_tensor(&[batch], DType::F32).unwrap(); + gpu.hip.memcpy_htod(&d_pos.buf, &positions).unwrap(); + + let kv_bytes = (seq_len * n_kv_heads * (head_dim / 32) * 34 * 2) as f64; + eprintln!( + "KV footprint per layer-call: {:.1} MB partials {:.1} MB", + kv_bytes / 1e6, + (max_tiles * n_heads * batch * (2 + head_dim) * 4) as f64 / 1e6 + ); + + let run_rows = |gpu: &mut Gpu| { + gpu.attention_flash_q8_0_rows_masked( + &d_q, &d_k, &d_v, &d_out, &d_pos, n_heads, n_kv_heads, head_dim, seq_len, batch, + &d_part, + ) + .unwrap() + }; + if !run_rows(&mut gpu) { + eprintln!("multi-row kernel refused this shape"); + return; + } + gpu.hip.device_synchronize().unwrap(); + let t = std::time::Instant::now(); + for _ in 0..iters { + run_rows(&mut gpu); + } + gpu.hip.device_synchronize().unwrap(); + let us = t.elapsed().as_secs_f64() * 1e6 / iters as f64; + eprintln!( + "rows kernel: {us:8.1} us/call {:.0} GB/s of KV", + kv_bytes / (us * 1e3) + ); + + let run_batched = |gpu: &mut Gpu| { + gpu.attention_flash_q8_0_batched_masked( + &d_q, + &d_k, + &d_v, + &d_out_batched, + &d_pos, + n_heads, + n_kv_heads, + head_dim, + max_seq, + seq_len, + batch, + &d_part, + None, + 0, + 0, + ) + .unwrap() + }; + run_batched(&mut gpu); + gpu.hip.device_synchronize().unwrap(); + let t = std::time::Instant::now(); + for _ in 0..iters { + run_batched(&mut gpu); + } + gpu.hip.device_synchronize().unwrap(); + let us_b = t.elapsed().as_secs_f64() * 1e6 / iters as f64; + eprintln!( + "batched (ROW): {us_b:8.1} us/call {:.0} GB/s of KV×rows", + kv_bytes * batch as f64 / (us_b * 1e3) + ); + eprintln!("speedup: {:.2}x", us_b / us); + + let rows_out = gpu.download_f32(&d_out).unwrap(); + let batched_out = gpu.download_f32(&d_out_batched).unwrap(); + let scale = batched_out + .iter() + .fold(0.0f32, |m, x| m.max(x.abs())) + .max(1e-12); + let max_abs = rows_out + .iter() + .zip(&batched_out) + .fold(0.0f32, |m, (a, b)| m.max((a - b).abs())); + let rel = max_abs / scale; + eprintln!("parity vs batched: max_abs={max_abs:.3e} rel={rel:.3e}"); + if check && !(rel < 1e-3) { + eprintln!("FAIL: multi-row output diverges from the batched kernel"); + std::process::exit(1); + } +} diff --git a/crates/rdna-compute/map.md b/crates/rdna-compute/map.md index f42c21343..14edf238c 100644 --- a/crates/rdna-compute/map.md +++ b/crates/rdna-compute/map.md @@ -24,7 +24,7 @@ _Generated by `scripts/check-crate-maps.py` from the tree — do not edit inside | File | Lines | Public items | Tests | |---|---:|---:|---:| | [`src/arch_caps.rs`](src/arch_caps.rs) | 731 | 57 | 19 | -| [`src/attention.rs`](src/attention.rs) | 16,139 | 221 | 12 | +| [`src/attention.rs`](src/attention.rs) | 16,317 | 222 | 12 | | [`src/bin/hipfire-kernel-hash.rs`](src/bin/hipfire-kernel-hash.rs) | 142 | 0 | 0 | | [`src/cdna/gfx942.rs`](src/cdna/gfx942.rs) | 578 | 10 | 1 | | [`src/cdna/mod.rs`](src/cdna/mod.rs) | 11 | 1 | 0 | @@ -43,7 +43,7 @@ _Generated by `scripts/check-crate-maps.py` from the tree — do not edit inside | [`src/gemma4_ops.rs`](src/gemma4_ops.rs) | 83 | 1 | 0 | | [`src/gemv.rs`](src/gemv.rs) | 15,670 | 235 | 0 | | [`src/graph.rs`](src/graph.rs) | 556 | 33 | 0 | -| [`src/kernels.rs`](src/kernels.rs) | 8,337 | 1266 | 37 | +| [`src/kernels.rs`](src/kernels.rs) | 8,339 | 1267 | 37 | | [`src/kv_slots.rs`](src/kv_slots.rs) | 496 | 9 | 12 | | [`src/lib.rs`](src/lib.rs) | 99 | 36 | 1 | | [`src/moe.rs`](src/moe.rs) | 1,742 | 27 | 0 | @@ -67,7 +67,7 @@ _Generated by `scripts/check-crate-maps.py` from the tree — do not edit inside ### Public API surface - [`src/arch_caps.rs`](src/arch_caps.rs): `ArchCaps`, `note_process_gpu_arch`, `process_gpu_arch`, `new`, `should_use_mmq`, `is_gfx906`, `is_gfx908`, `is_gfx1010`, `is_gfx1011`, `is_gfx1012`, `is_gfx1030`, `is_gfx1031`, +45 more -- [`src/attention.rs`](src/attention.rs): `attention_q8_0_kv_independent_lds_bytes`, `attention_q8_0_kv_independent_max_lane_capacity`, `q8_flash_tile_size`, `dspark_stage_kv`, `triattn_accumulate`, `attention_f32`, `attention_flash`, `attention_flash_gqa`, `attention_flash_gqa_fused`, `attention_gqa_warp`, `attention_gqa_warp_dv`, `kv_cache_write_hfq4`, +209 more +- [`src/attention.rs`](src/attention.rs): `attention_q8_0_kv_independent_lds_bytes`, `attention_q8_0_kv_independent_max_lane_capacity`, `q8_flash_tile_size`, `dspark_stage_kv`, `triattn_accumulate`, `attention_f32`, `attention_flash`, `attention_flash_gqa`, `attention_flash_gqa_fused`, `attention_gqa_warp`, `attention_gqa_warp_dv`, `kv_cache_write_hfq4`, +210 more - [`src/bin/hipfire-kernel-hash.rs`](src/bin/hipfire-kernel-hash.rs): — - [`src/cdna/gfx942.rs`](src/cdna/gfx942.rs): `Gfx942Device`, `try_gfx942`, `mq2_lloyd_moe_gate_up_wave64`, `mq2_lloyd_moe_gate_up_wave64x8_candidate`, `mq_rotate_x_wave64_batched`, `mq2_lloyd_moe_down_expanded_wave64`, `mq2_lloyd_moe_down_residual_wave64`, `indexer_top_k_buf_parallel`, `grouped_olora_e8`, `grouped_olora_e8_wave64x4_candidate` - [`src/cdna/mod.rs`](src/cdna/mod.rs): `gfx942` @@ -86,7 +86,7 @@ _Generated by `scripts/check-crate-maps.py` from the tree — do not edit inside - [`src/gemma4_ops.rs`](src/gemma4_ops.rs): `gemma4_ple_gelu_mul_strided_f32` - [`src/gemv.rs`](src/gemv.rs): `gemv_q4lut`, `gemv_q4wave`, `gemv_q4as8`, `gemv_f32`, `gemv_q4k`, `gemv_hfq4g128`, `givens_rotate`, `givens_rotate_to`, `fused_silu_mul_givens_rotate_f32`, `ensure_paro_scratch`, `ensure_paro_fused_scratch`, `fused_gate_up_paro4g128t`, +223 more - [`src/graph.rs`](src/graph.rs): `PerBGraphCache`, `GraphState`, `begin_graph_capture`, `begin_graph_capture_relaxed`, `end_graph_capture`, `end_graph_capture_segment`, `graph_segment_count`, `abort_graph_capture`, `graph_segment_launch`, `drop_graph_segments`, `graph_launch`, `end_decode_turn`, +21 more -- [`src/kernels.rs`](src/kernels.rs): `GEMV_SRC`, `GEMV_Q4K_SRC`, `GEMV_HFQ4G128_SRC`, `GEMV_HFQ4G128_RESIDUAL_SIGMOID_SCALED_SRC`, `GEMV_PARO4G128_SRC`, `GEMM_HFQ4G128_SRC`, `GEMM_HFQ4G128_MMQ_GFX1151_SRC`, `GEMV_HFQ2G256_SRC`, `GEMV_MQ2G256_LLOYD_SRC`, `GEMV_MQ3G256_LLOYD_SRC`, `GEMV_MQ4G256_LLOYD_SRC`, `GEMV_MQ4G256_LLOYD_GFX1100_SRC`, +1254 more +- [`src/kernels.rs`](src/kernels.rs): `GEMV_SRC`, `GEMV_Q4K_SRC`, `GEMV_HFQ4G128_SRC`, `GEMV_HFQ4G128_RESIDUAL_SIGMOID_SCALED_SRC`, `GEMV_PARO4G128_SRC`, `GEMM_HFQ4G128_SRC`, `GEMM_HFQ4G128_MMQ_GFX1151_SRC`, `GEMV_HFQ2G256_SRC`, `GEMV_MQ2G256_LLOYD_SRC`, `GEMV_MQ3G256_LLOYD_SRC`, `GEMV_MQ4G256_LLOYD_SRC`, `GEMV_MQ4G256_LLOYD_GFX1100_SRC`, +1255 more - [`src/kv_slots.rs`](src/kv_slots.rs): `KvSlotDesc`, `total_rows`, `half_from_f32`, `build_arena`, `build_asym3_k_arena`, `build_tiles`, `R9700_VRAM_BYTES`, `mem_available_bytes`, `preflight_alloc` - [`src/lib.rs`](src/lib.rs): `arch_caps`, `attention`, `cdna`, `dflash_draft_fusion`, `dflash_gdn_pre`, `dflash_hidden_scatter`, `dflash_state_copy`, `embedding`, `feature_flags`, `flash_attn_ck`, `flux_fused`, `gemm`, +24 more - [`src/moe.rs`](src/moe.rs): `moe_down_combine_k8_batched`, `moe_down_combine_rmsnorm_mq_rotate_vecsum_gfx1100`, `moe_scatter_histogram_k8`, `moe_scatter_offsets_k8`, `moe_scatter_permute_k8`, `moe_scatter_fused_k8`, `moe_down_combine_grouped_k8`, `moe_gate_up_unscatter_k8`, `moe_unscatter_silu_clamp_k8`, `hash_router_normalize_f32`, `hash_router_normalize_f32_batched`, `hash_router_normalize_f32_buf`, +15 more @@ -120,6 +120,6 @@ _Generated by `scripts/check-crate-maps.py` from the tree — do not edit inside ### Totals -- 40 modules · 119,936 lines · 2938 public items · 258 tests · 206 examples +- 40 modules · 120,116 lines · 2940 public items · 258 tests · 207 examples diff --git a/crates/rdna-compute/src/attention.rs b/crates/rdna-compute/src/attention.rs index 8f64583c4..c1c2a70ff 100644 --- a/crates/rdna-compute/src/attention.rs +++ b/crates/rdna-compute/src/attention.rs @@ -142,6 +142,18 @@ fn wmma_fa_min_batch() -> usize { .unwrap_or(16) } +/// Query rows one multi-row flash block owns. 8 is the register budget of the +/// kernel (ROWS x (Q, accumulator) per lane at head_dim 256). Measured on +/// gfx1100 at 33k context against the batched tile: 1.91x at 8 rows, 1.66x +/// at 4, 0.85x at 2 — so a block never takes fewer than 4 rows and the caller +/// keeps batches under 4 on the batched kernel. +fn flash_rows_per_block(batch_size: usize) -> usize { + [8usize, 4] + .into_iter() + .find(|&r| r <= batch_size) + .unwrap_or(0) +} + impl Gpu { /// DSpark bidirectional staging assembly (on-GPU; replaces a host /// d2h+assemble+h2d that forced ~2 stream syncs per stage). @@ -3868,6 +3880,172 @@ impl Gpu { ) } + /// One KV scan for `batch_size` query rows; `Ok(false)` = out of scope, caller must fall back. + #[allow(clippy::too_many_arguments)] + pub fn attention_flash_q8_0_rows_masked( + &mut self, + q: &GpuTensor, + k_cache: &GpuTensor, + v_cache: &GpuTensor, + out: &GpuTensor, + positions: &GpuTensor, + n_heads: usize, + n_kv_heads: usize, + head_dim: usize, + max_ctx_len: usize, + batch_size: usize, + partials: &GpuTensor, + ) -> HipResult { + if !self.arch_caps.is_gfx1100() { + return Ok(false); + } + let dpt = head_dim / 32; + if head_dim % 32 != 0 || !(dpt == 4 || dpt == 8) { + return Ok(false); + } + let rows = flash_rows_per_block(batch_size); + if rows < 4 { + return Ok(false); + } + self.bind_thread()?; + let tile_size = self.attn_tile_size(); + let max_tiles = max_ctx_len.div_ceil(tile_size); + let stride = 2 + head_dim; + let partials_bytes_per_row = n_heads * max_tiles * stride * 4; + if partials_bytes_per_row == 0 { + return Ok(false); + } + let sub_batch = (partials.numel() * 4 / partials_bytes_per_row) + .max(1) + .min(batch_size); + let func: &'static str = match (rows, dpt) { + (8, 8) => "attention_flash_q8_0_rows8_d8", + (4, 8) => "attention_flash_q8_0_rows4_d8", + (8, 4) => "attention_flash_q8_0_rows8_d4", + (4, 4) => "attention_flash_q8_0_rows4_d4", + _ => return Ok(false), + }; + self.ensure_kernel(func, kernels::ATTENTION_FLASH_Q8_0_TILE_ROWS_SRC, func)?; + self.ensure_kernel( + "attention_flash_asym_reduce_batched", + kernels::ATTENTION_FLASH_ASYM_REDUCE_BATCHED_SRC, + "attention_flash_asym_reduce_batched", + )?; + + let q_dim = n_heads * head_dim; + // Scores are carried in log2 space so the kernel's softmax is one + // v_exp_f32 per row-token; partials convert back on write. + let scale = std::f32::consts::LOG2_E / (head_dim as f32).sqrt(); + let mut offset = 0usize; + while offset < batch_size { + let chunk = (batch_size - offset).min(sub_batch); + { + let q_ptr = + unsafe { (q.buf.as_ptr() as *mut u8).add(offset * q_dim * 4) as *mut c_void }; + let k_ptr = k_cache.buf.as_ptr(); + let v_ptr = v_cache.buf.as_ptr(); + let p_ptr = partials.buf.as_ptr(); + let pos_ptr = positions.buf.as_ptr(); + let nh = n_heads as i32; + let nkv = n_kv_heads as i32; + let hd = head_dim as i32; + let sc = scale; + let ts = tile_size as i32; + let mt = max_tiles as i32; + let bo = offset as i32; + let rv = chunk as i32; + let mut params: Vec<*mut c_void> = vec![ + &q_ptr as *const _ as *mut c_void, + &k_ptr as *const _ as *mut c_void, + &v_ptr as *const _ as *mut c_void, + &p_ptr as *const _ as *mut c_void, + &pos_ptr as *const _ as *mut c_void, + &nh as *const _ as *mut c_void, + &nkv as *const _ as *mut c_void, + &hd as *const _ as *mut c_void, + &sc as *const _ as *mut c_void, + &ts as *const _ as *mut c_void, + &mt as *const _ as *mut c_void, + &bo as *const _ as *mut c_void, + &rv as *const _ as *mut c_void, + ]; + let groups = chunk.div_ceil(rows); + self.launch_maybe_blob( + func, + [n_heads as u32, max_tiles as u32, groups as u32], + [32, 1, 1], + 0, + &mut params, + || { + let mut b = hip_bridge::KernargBlob::new(); + b.push_ptr(q_ptr); + b.push_ptr(k_ptr); + b.push_ptr(v_ptr); + b.push_ptr(p_ptr); + b.push_ptr(pos_ptr); + b.push_i32(nh); + b.push_i32(nkv); + b.push_i32(hd); + b.push_f32(sc); + b.push_i32(ts); + b.push_i32(mt); + b.push_i32(bo); + b.push_i32(rv); + b + }, + )?; + } + { + let p_ptr = partials.buf.as_ptr(); + let o_ptr = + unsafe { (out.buf.as_ptr() as *mut u8).add(offset * q_dim * 4) as *mut c_void }; + let pos_ptr = positions.buf.as_ptr(); + let nh = n_heads as i32; + let hd = head_dim as i32; + let ts = tile_size as i32; + let mt = max_tiles as i32; + let bo = offset as i32; + let bs = 0i32; + let bc = 0i32; + let mut params: Vec<*mut c_void> = vec![ + &p_ptr as *const _ as *mut c_void, + &o_ptr as *const _ as *mut c_void, + &pos_ptr as *const _ as *mut c_void, + &nh as *const _ as *mut c_void, + &hd as *const _ as *mut c_void, + &ts as *const _ as *mut c_void, + &mt as *const _ as *mut c_void, + &bo as *const _ as *mut c_void, + &bs as *const _ as *mut c_void, + &bc as *const _ as *mut c_void, + ]; + self.launch_maybe_blob( + "attention_flash_asym_reduce_batched", + [n_heads as u32, chunk as u32, 1], + [32, 1, 1], + 0, + &mut params, + || { + let mut b = hip_bridge::KernargBlob::new(); + b.push_ptr(p_ptr); + b.push_ptr(o_ptr); + b.push_ptr(pos_ptr); + b.push_i32(nh); + b.push_i32(hd); + b.push_i32(ts); + b.push_i32(mt); + b.push_i32(bo); + b.push_i32(bs); + b.push_i32(bc); + b + }, + )?; + } + offset += chunk; + } + Ok(true) + } + /// Multi-slot Q8_0 tiled flash attention. `slot_descs` is `[n_slots]` /// `KvSlotDesc`; `row_slot` is `[batch_size]` slot indices per query row. /// Passing `None` for both is exactly the legacy single-sequence path diff --git a/crates/rdna-compute/src/kernels.rs b/crates/rdna-compute/src/kernels.rs index 3662fdbb7..8c98a22eb 100644 --- a/crates/rdna-compute/src/kernels.rs +++ b/crates/rdna-compute/src/kernels.rs @@ -5545,6 +5545,8 @@ pub const ATTENTION_FLASH_Q8_0_TILE_BATCHED_SRC: &str = include_str!("../../../kernels/src/attention_flash_q8_0_tile_batched.hip"); pub const ATTENTION_FLASH_BF16_TILE_BATCHED_SRC: &str = include_str!("../../../kernels/src/attention_flash_bf16_tile_batched.hip"); +pub const ATTENTION_FLASH_Q8_0_TILE_ROWS_SRC: &str = + include_str!("../../../kernels/src/attention_flash_q8_0_tile_rows.hip"); pub const ATTENTION_FLASH_ASYM_REDUCE_BATCHED_SRC: &str = include_str!("../../../kernels/src/attention_flash_asym_reduce_batched.hip"); diff --git a/docs/env-vars.md b/docs/env-vars.md index 9d2fb4a5f..3ca180c39 100644 --- a/docs/env-vars.md +++ b/docs/env-vars.md @@ -178,6 +178,7 @@ diagnostic and developer harness exports pending their cleanup. | `HIPFIRE_MAX_REQUEST_BYTES` | Body cap | | `HIPFIRE_SERVE_MAX_QUEUE` / `HIPFIRE_SERVE_QUEUE_TIMEOUT_MS` | Admission queue | | `HIPFIRE_EXPERIMENTAL_BUDGET_ALERT` | Research budget nudge | +| `HIPFIRE_FA_PERTOKEN_MIN_CTX` | Context length past which an exact-gfx1100 Q8 small-batch (n = 4..32, head_dim 128/256, sequential non-tree, graph capture off) attend step leaves the batched flash kernel for the multi-row tile; default `4096`, `0` disables the route. Other arches, KV modes, shapes, and semantics retain the batched route. | | `HIPFIRE_RCCL_LIB` | Explicit `librccl.so` path, tried before the ROCm root. For distributions whose ROCm prefix does not carry RCCL (nixpkgs: `rocmtoolkit-merged` has HIP/HSA, `librccl` is a separate store path). | | `HIPFIRE_DEVICES` / `HIPFIRE_TP` / `HIPFIRE_TP_USE_RCCL` | Multi-GPU / TP. `HIPFIRE_DEVICES` is the compatibility alias for `hardware.devices`; startup lowers its physical list to ROCr selectors plus matching HIP logical selectors. | | `HIPFIRE_ALLOW_MIXED_ARCH=1` | Mixed arch pairs | @@ -296,7 +297,7 @@ Copyable user, developer, and retained-PM4 TOML profiles are in **Do not hand-edit rows below** except by re-running the source scan. **Generation method:** token scan over visible `*.rs`, `*.py`, and `*.sh`, excluding ignored/generated files. **Columns:** variable; up to two lexical source paths. -**Count:** 737 +**Count:** 738 | Variable | Example source path(s) | |---|---| diff --git a/docs/perf-checkpoints/2026-09-10-gfx1100-qwen38-multirow-verifier.md b/docs/perf-checkpoints/2026-09-10-gfx1100-qwen38-multirow-verifier.md new file mode 100644 index 000000000..4cff58f4f --- /dev/null +++ b/docs/perf-checkpoints/2026-09-10-gfx1100-qwen38-multirow-verifier.md @@ -0,0 +1,108 @@ +# gfx1100 Qwen3.8 multi-row verifier — 2026-09-10 + +**Lifecycle:** `historical` + +**Disposition:** measured candidate evidence; not a product baseline or admission decision. + +## Question and scope + +Measure a gfx1100-only Q8 flash-attention kernel in which one wave owns four +or eight verifier rows and shares each KV scan across those rows. Only the +attention step changes; projections remain batched. The production admission +predicate is exact `gfx1100`, Q8 KV, head dimension 128 or 256, sequential +non-tree batches of 4–32 rows, logical context above 4096, and graph capture +off. Unsupported shapes retain the established batched route. + +Current-beta base: `b8092f7c7fe0eb3dabccc28e8993ee10c3465fc6`. + +## Fixture identity + +- Host: `odin`, Radeon RX 7900 XTX, `gfx1100`, wave32, 24,560 MiB reported + VRAM, HIP `7.2.53211-9999`. +- Target: `qwen3.8-27b.mq4-xt`, 14,980,361,216 bytes; SHA-256 + `9f91556f7e0431a077d03756a7102d0154108757289e6e5fe9a2d204c0c9eeb7`; + MD5 `e45d15bfe0c9a87132697101d17cbed6`. +- Draft: `qwen38-27b-dflash-mq4.hfq`, 1,209,603,072 bytes; SHA-256 + `d0a74a232a0e2166d889f823e91e0fbf778d21dd9668d7de055cdecb065401bc`; + MD5 `013395583cd04206c8aa68f4d061983d`. +- Prompt: `benchmarks/prompts/qwen38_issue693_longcode_20676.txt`, 75,251 + bytes; MD5 `b4d0b63cddcac872648ddf3cdd92cac2`; 21,550 tokens after the + Qwen chat scaffold. +- Tested daemon MD5: `512fccca7c7189559048a7aba17cb6c1` (SHA-256 + `4775d1225c71db6d8c1b717e59a62ff1145f74f9ee433462f8376db99f19bfb3`). +- Dedicated micro-harness MD5: `0a032cfcbf525c9fa3565a7b55ca6d88`. + +## Product-path A/B + +Six native-daemon fresh processes ran in declared order +`off,on,on,off,off,on`. Both arms used `max_seq=65536`, Q8 KV, DFlash, +greedy sampling, `HIPFIRE_VERIFY_GRAPH=0`, a ten-second DPM warmup, and 200 +generated tokens. The only arm difference was +`HIPFIRE_FA_PERTOKEN_MIN_CTX=0` versus `4096`. A separate unrecorded candidate +probe populated the shared JIT cache before the series. + +| route | decode samples (tok/s) | median | delta | +|---|---|---:|---:| +| established batched | 33.4, 32.6, 33.4 | 33.4 | — | +| multi-row R4/R8 | 46.4, 42.9, 46.3 | 46.3 | +38.6% | + +All six samples produced 200 tokens in 69 cycles with `tau=1.88`, no daemon +errors, and byte-identical decoded output MD5 +`b501ab0e0102889bd63537f2006d4f61`. A separately built, warmed current-beta +daemon produced 33.5 tok/s with the same token count, cycles, tau, and output +MD5. The first baseline invocation was discarded as cold-JIT (11.2 tok/s). + +Raw local discovery artifact: `/mnt/data4/claude-scratch/20260831-hipfire-tp2/evidence-beta/xt-ab/results.jsonl`, +MD5 `1b037e814fd00caae4b9ca6d919f3a30`. + +## Kernel screen + +The dedicated oracle used the Qwen3.8-27B shape (24 query heads, four KV +heads, head dimension 256), 100 timed iterations per cell, and compared every +output against `attention_flash_q8_0_tile_batched`. + +| context | R4 speedup | R8 speedup | +|---:|---:|---:| +| 2,048 | 0.98x | 0.72x | +| 4,096 | 1.45x | 2.01x | +| 20,676 | 1.99x | 2.40x | +| 32,768 | 1.94x | 2.28x | + +Worst relative output error was `4.222e-7` against a `1e-3` limit. The +head-dimension-128 screen at context 8192 measured 1.33x for R4 and 2.03x for +R8, with worst relative error `3.419e-7`. + +Radiowave inspection for head dimension 256 reported: + +| entry point | VGPR | SGPR | VGPR spills | SGPR spills | private bytes | +|---|---:|---:|---:|---:|---:| +| `attention_flash_q8_0_rows4_d8` | 106 | 41 | 0 | 0 | 0 | +| `attention_flash_q8_0_rows8_d8` | 186 | 58 | 0 | 0 | 0 | + +## Correctness and route validation + +- `test_kernels`: 16 passed, 0 failed, 0 skipped on the RX 7900 XTX. +- `hipfire-arch-qwen35` unit tests: 193 passed, 4 ignored. +- `rdna-compute` unit tests: 242 passed. +- Canonical-XT serve battery: five of five coherent responses passed recall, + empty, runaway, and attractor checks. +- crate-map generation, env-doc scan, changed-file formatting, fmt-bomb, and + diff whitespace checks passed. + +The canonical kernel-bucket Redline PM4 arm is blocked on both the clean base +and candidate by the same current-beta limitation: +`gemv_mq4g256v2_residual: GFX10/GFX11 PM4 dispatch does not yet support +scratch (private=32, dynamic_callstack=false)`. Before that refusal, both +lanes reported the same capture identities: prefill-128 hash +`bdc60fd56c3670e7`, prefill-512 hash `9111b45dd02bfe5d`, and decode hash +`92ede73d35f4a51f`. This record does not claim a PM4 pass; the new route is +itself excluded during graph/retained capture and falls back to the batched +kernel there. + +## Interpretation + +The result supports review of this narrowly gated gfx1100 candidate. It does +not transfer to other architectures, KV formats, graph/PM4 routes, prompts, +or drafts. It also does not solve draft-target disagreement: this fixture's +`tau=1.88` is unchanged, so sufficiently low-tau workloads may still favor +plain autoregressive decode. diff --git a/kernels/src/attention_flash_q8_0_tile_rows.hip b/kernels/src/attention_flash_q8_0_tile_rows.hip new file mode 100644 index 000000000..5d1412639 --- /dev/null +++ b/kernels/src/attention_flash_q8_0_tile_rows.hip @@ -0,0 +1,197 @@ +// SPDX-License-Identifier: Apache-2.0 +// Copyright (c) 2026 Kaden Schutt +// hipfire — see LICENSE and NOTICE in the project root. + +// Multi-row Q8_0 flash attention tile — ONE KV scan serves ROWS query rows. +// +// attention_flash_q8_0_tile_batched grids [n_heads, tiles, ROW], so a verify +// batch re-reads the whole KV once per row. Here a block owns (head, tile, +// row-group) and keeps ROWS × (Q, running max, running sum, accumulator) in +// registers, so each K/V block is fetched once and reused across the group. +// That is the only reason this kernel exists; the arithmetic is the same +// two-pass softmax expressed in the online (FlashAttention) form, because a +// per-row score array would need LDS and cost the occupancy the register +// form keeps. +// +// Grid: [n_heads, max_tiles, ceil(rows_valid / ROWS)]. Block: [32, 1, 1]. +// LDS: none. +// Partials: [row][head][max_tiles][2 + head_dim] — identical to the batched +// tile kernel, so attention_flash_asym_reduce_batched consumes it unchanged. +// +// Scope: full causal, single KV arena (no slot descriptors), no tree bias, +// no sliding window. The launcher refuses anything else. +#include + +static __device__ __forceinline__ unsigned int rows_load_u32(const unsigned char* p) { + unsigned int packed; + __builtin_memcpy(&packed, p, sizeof(packed)); + return packed; +} + +static __device__ __forceinline__ unsigned long long rows_load_u64(const unsigned char* p) { + unsigned long long packed; + __builtin_memcpy(&packed, p, sizeof(packed)); + return packed; +} + +template +static __device__ __forceinline__ void rows_dequant( + const unsigned char* __restrict__ blk, int bj, float scale, float* __restrict__ out +) { + if (DPT == 4) { + const unsigned int packed = rows_load_u32(blk + 2 + bj); + #pragma unroll + for (int i = 0; i < 4; i++) + out[i] = scale * (float)(signed char)(packed >> (8 * i)); + } else { + const unsigned long long packed = rows_load_u64(blk + 2 + bj); + #pragma unroll + for (int i = 0; i < 8; i++) + out[i] = scale * (float)(signed char)(packed >> (8 * i)); + } +} + +// Five-stage XOR butterfly over the wave: each lane holds DPT partial +// products of one score, so every score costs one full cross-lane reduction. +// Spreading KV rows across lanes instead would remove it, but on gfx1100 the +// LDS staging that layout needs costs more occupancy than the shuffle costs +// time (measured: 1058-1989 us against 706 us for this form). +__device__ __forceinline__ float rows_reduce_sum(float v) { + for (int off = 16; off > 0; off >>= 1) v += __shfl_xor(v, off); + return v; +} + +template +static __device__ __forceinline__ void flash_rows_body( + const float* __restrict__ q, + const unsigned char* __restrict__ k_cache, + const unsigned char* __restrict__ v_cache, + float* __restrict__ partials, + const int* __restrict__ positions, + int n_heads, + int n_kv_heads, + int head_dim, + float scale_attn, + int tile_size, + int max_tiles, + int batch_offset, + int rows_valid +) { + const int h = blockIdx.x; + if (h >= n_heads) return; + const int tile_id = blockIdx.y; + const int tid = threadIdx.x; + const int row0 = (int)blockIdx.z * ROWS; + + int seq[ROWS]; + int seq_max = 0; + #pragma unroll + for (int r = 0; r < ROWS; r++) { + const int row = row0 + r; + seq[r] = (row < rows_valid) ? (positions[batch_offset + row] + 1) : 0; + seq_max = max(seq_max, seq[r]); + } + const int tile_start = tile_id * tile_size; + if (tile_start >= seq_max) return; + const int tile_end = min(tile_start + tile_size, seq_max); + const int tile_len = tile_end - tile_start; + + const int kv_group = n_heads / n_kv_heads; + const int kv_h = h / kv_group; + const int blocks_per_head = head_dim / 32; + const int per_pos_bytes = n_kv_heads * blocks_per_head * 34; + const int q_dim = n_heads * head_dim; + const int d0 = tid * DPT; + const int bj = d0 % 32; + const int blk_off = (kv_h * blocks_per_head + d0 / 32) * 34; + + float mq[ROWS][DPT]; + #pragma unroll + for (int r = 0; r < ROWS; r++) { + const int row = (row0 + r < rows_valid) ? (row0 + r) : 0; + const float* qh = q + (size_t)row * q_dim + h * head_dim; + #pragma unroll + for (int i = 0; i < DPT; i++) mq[r][i] = qh[d0 + i]; + } + + float acc[ROWS][DPT]; + float run_max[ROWS]; + float run_sum[ROWS]; + #pragma unroll + for (int r = 0; r < ROWS; r++) { + run_max[r] = -1e30f; + run_sum[r] = 0.0f; + #pragma unroll + for (int i = 0; i < DPT; i++) acc[r][i] = 0.0f; + } + + for (int t = tile_start; t < tile_end; t++) { + const unsigned char* kb = k_cache + (size_t)t * per_pos_bytes + blk_off; + float kd[DPT]; + rows_dequant(kb, bj, (float)*((const _Float16*)kb), kd); + float s[ROWS]; + #pragma unroll + for (int r = 0; r < ROWS; r++) { + float p = 0.0f; + #pragma unroll + for (int i = 0; i < DPT; i++) p += mq[r][i] * kd[i]; + s[r] = rows_reduce_sum(p) * scale_attn; + } + const unsigned char* vb = v_cache + (size_t)t * per_pos_bytes + blk_off; + float vd[DPT]; + rows_dequant(vb, bj, (float)*((const _Float16*)vb), vd); + #pragma unroll + for (int r = 0; r < ROWS; r++) { + // Causal mask per row: rows in one group end at different positions. + if (t >= seq[r]) continue; + const float sr = s[r]; + if (sr > run_max[r]) { + const float corr = __builtin_amdgcn_exp2f(run_max[r] - sr); + run_sum[r] *= corr; + #pragma unroll + for (int i = 0; i < DPT; i++) acc[r][i] *= corr; + run_max[r] = sr; + } + const float e = __builtin_amdgcn_exp2f(sr - run_max[r]); + run_sum[r] += e; + #pragma unroll + for (int i = 0; i < DPT; i++) acc[r][i] += e * vd[i]; + } + } + + #pragma unroll + for (int r = 0; r < ROWS; r++) { + const int row = row0 + r; + if (row >= rows_valid) continue; + float* p = partials + + ((long long)row * n_heads + h) * max_tiles * (2 + head_dim) + + (long long)tile_id * (2 + head_dim); + if (tid == 0) { + p[0] = run_max[r] * 0.69314718f; + p[1] = run_sum[r]; + } + #pragma unroll + for (int i = 0; i < DPT; i++) p[2 + d0 + i] = acc[r][i]; + } +} + +#define HIPFIRE_FLASH_ROWS_ENTRY(ROWS, DPT) \ +extern "C" __launch_bounds__(32) \ +__global__ void attention_flash_q8_0_rows##ROWS##_d##DPT( \ + const float* __restrict__ q, \ + const unsigned char* __restrict__ k_cache, \ + const unsigned char* __restrict__ v_cache, \ + float* __restrict__ partials, \ + const int* __restrict__ positions, \ + int n_heads, int n_kv_heads, int head_dim, float scale_attn, \ + int tile_size, int max_tiles, int batch_offset, int rows_valid \ +) { \ + flash_rows_body(q, k_cache, v_cache, partials, positions, \ + n_heads, n_kv_heads, head_dim, scale_attn, \ + tile_size, max_tiles, batch_offset, rows_valid); \ +} + +HIPFIRE_FLASH_ROWS_ENTRY(4, 4) +HIPFIRE_FLASH_ROWS_ENTRY(8, 4) +HIPFIRE_FLASH_ROWS_ENTRY(4, 8) +HIPFIRE_FLASH_ROWS_ENTRY(8, 8)