diff --git a/crates/hipfire-arch-gemma4/map.md b/crates/hipfire-arch-gemma4/map.md index 99e65166a0..62e586db97 100644 --- a/crates/hipfire-arch-gemma4/map.md +++ b/crates/hipfire-arch-gemma4/map.md @@ -31,7 +31,7 @@ _Generated by `scripts/check-crate-maps.py` from the tree — do not edit inside | [`src/gemma4.rs`](src/gemma4.rs) | 1,088 | 13 | 0 | | [`src/gemma4_vision.rs`](src/gemma4_vision.rs) | 16 | 3 | 0 | | [`src/lib.rs`](src/lib.rs) | 48 | 8 | 0 | -| [`src/lowered.rs`](src/lowered.rs) | 5,876 | 35 | 0 | +| [`src/lowered.rs`](src/lowered.rs) | 5,883 | 35 | 0 | | [`src/speculative.rs`](src/speculative.rs) | 252 | 6 | 0 | ### Public API surface @@ -60,6 +60,6 @@ _Generated by `scripts/check-crate-maps.py` from the tree — do not edit inside ### Totals -- 10 modules · 12,179 lines · 111 public items · 18 tests · 8 examples +- 10 modules · 12,186 lines · 111 public items · 18 tests · 8 examples diff --git a/crates/hipfire-arch-gemma4/src/lowered.rs b/crates/hipfire-arch-gemma4/src/lowered.rs index 16e65e7739..0f44bb8228 100644 --- a/crates/hipfire-arch-gemma4/src/lowered.rs +++ b/crates/hipfire-arch-gemma4/src/lowered.rs @@ -3038,6 +3038,7 @@ fn sliding_layer_decode_impl( quant_q4: false, quant_int8: false, quant_hfq8: false, + quant_bf16: false, f32_policy: hipfire_dispatch::families::kv_tier::F32AttnPolicy::Simple, v_mode_bits: kv_cache.v_mode_bits(), pos, @@ -3449,6 +3450,7 @@ fn full_layer_decode_impl( quant_q4: false, quant_int8: false, quant_hfq8: false, + quant_bf16: false, f32_policy: hipfire_dispatch::families::kv_tier::F32AttnPolicy::Simple, v_mode_bits: kv_cache.v_mode_bits(), pos, @@ -4244,6 +4246,7 @@ fn forward_prefill_batch_v2( quant_q4: false, quant_int8: false, quant_hfq8: false, + quant_bf16: false, f32_policy: hipfire_dispatch::families::kv_tier::F32AttnPolicy::Simple, v_mode_bits: kv_sliding.v_mode_bits(), pos, @@ -4576,6 +4579,7 @@ fn forward_prefill_batch_v2( quant_q4: false, quant_int8: false, quant_hfq8: false, + quant_bf16: false, f32_policy: hipfire_dispatch::families::kv_tier::F32AttnPolicy::Simple, v_mode_bits: kv_full.v_mode_bits(), pos: start_pos + n_batch - 1, @@ -4655,6 +4659,7 @@ fn forward_prefill_batch_v2( quant_q4: false, quant_int8: false, quant_hfq8: false, + quant_bf16: false, f32_policy: hipfire_dispatch::families::kv_tier::F32AttnPolicy::Simple, v_mode_bits: kv_full.v_mode_bits(), pos, @@ -5562,6 +5567,7 @@ impl<'a> ForwardBindings for Gemma4Bindings<'a> { quant_q4: false, quant_int8: false, quant_hfq8: false, + quant_bf16: false, f32_policy: hipfire_dispatch::families::kv_tier::F32AttnPolicy::Simple, v_mode_bits: kv.v_mode_bits(), pos, @@ -5677,6 +5683,7 @@ impl<'a> ForwardBindings for Gemma4Bindings<'a> { quant_q4: false, quant_int8: false, quant_hfq8: false, + quant_bf16: false, f32_policy: hipfire_dispatch::families::kv_tier::F32AttnPolicy::Simple, v_mode_bits: kv.v_mode_bits(), pos, diff --git a/crates/hipfire-arch-maple/examples/maple_coherence.rs b/crates/hipfire-arch-maple/examples/maple_coherence.rs index 9c6bc40839..7f76aca289 100644 --- a/crates/hipfire-arch-maple/examples/maple_coherence.rs +++ b/crates/hipfire-arch-maple/examples/maple_coherence.rs @@ -12,7 +12,14 @@ //! produce a model that loads, runs at full speed, and emits garbage. //! //! Usage: -//! maple_coherence --model [--prompt "..."] [--max-tokens N] [--raw] +//! maple_coherence --model [--prompt "..."] [--max-tokens N] +//! [--raw] [--kv-mode q8|bf16] +//! [--temp T] [--top-p P] [--seed N] +//! [--head ] +//! +//! `--kv-mode bf16` swaps the Q8_0 KV cache for the flat BF16 tier. Both run +//! the same sliding-window kernels with the same dim mapping and FMA order, so +//! a q8-vs-bf16 diff isolates KV storage precision from everything else. //! //! `HIPFIRE_MAPLE_PER_TOKEN_PREFILL=1` forces the per-token prefill path, so the //! batched path can be A/B'd against it from one binary on one machine. @@ -22,7 +29,7 @@ //! against the HF reference is a separate follow-up; it needs a capture hook //! inside `decode_step_body`, which this harness deliberately does not have. -use hipfire_arch_maple::bundle::load_maple_from_hfq; +use hipfire_arch_maple::bundle::load_maple_from_hfq_with_head; use hipfire_arch_maple::forward::decode_step; use hipfire_runtime::hfq::HfqFile; use std::path::Path; @@ -32,6 +39,84 @@ struct Args { prompt: String, max_tokens: usize, raw: bool, + kv_mode: String, + /// Optional head-overlay `.hfq` (hipfire-quantize --head-only). + head: Option, + /// 0.0 = greedy (default, unchanged behaviour). > 0 = sample. + temp: f32, + top_p: f32, + seed: u64, +} + +/// SplitMix64 — a 64-bit mixer used here as the sampling RNG. +/// +/// Deliberately self-contained and NOT the engine's sampler: this harness needs +/// a stream that depends only on `--seed`, so two runs of the same arm are +/// reproducible and two different seeds are genuinely independent draws. It is +/// not trying to match production sampling numerics. +struct SplitMix64(u64); +impl SplitMix64 { + fn new(seed: u64) -> Self { + Self(seed.wrapping_add(0x9E3779B97F4A7C15)) + } + fn next_u64(&mut self) -> u64 { + self.0 = self.0.wrapping_add(0x9E3779B97F4A7C15); + let mut z = self.0; + z = (z ^ (z >> 30)).wrapping_mul(0xBF58476D1CE4E5B9); + z = (z ^ (z >> 27)).wrapping_mul(0x94D049BB133111EB); + z ^ (z >> 31) + } + /// Uniform in [0, 1). 53 bits of mantissa, so the quantisation is far finer + /// than any probability this is used to compare against. + fn next_f64(&mut self) -> f64 { + (self.next_u64() >> 11) as f64 / (1u64 << 53) as f64 + } +} + +/// Temperature + top-p (nucleus) sampling. +/// +/// Softmax is computed in f64 max-shifted so the exponentials cannot overflow; +/// the vocab is 151,936 wide and the raw logit range is large enough that the +/// naive form does overflow in f32. +fn sample_top_p(logits: &[f32], temp: f32, top_p: f32, rng: &mut SplitMix64) -> u32 { + let mut idx: Vec = (0..logits.len() as u32).collect(); + let max = logits.iter().cloned().fold(f32::NEG_INFINITY, f32::max) as f64; + let t = temp.max(1e-6) as f64; + let mut p: Vec = logits + .iter() + .map(|&v| ((v as f64 - max) / t).exp()) + .collect(); + let sum: f64 = p.iter().sum(); + for v in p.iter_mut() { + *v /= sum; + } + // Descending by probability, then keep the smallest prefix whose mass + // reaches top_p. The prefix always keeps at least one token, so a + // degenerate top_p cannot produce an empty nucleus. + idx.sort_unstable_by(|&a, &b| { + p[b as usize] + .partial_cmp(&p[a as usize]) + .unwrap_or(std::cmp::Ordering::Equal) + }); + let mut cum = 0.0; + let mut cut = idx.len(); + for (n, &i) in idx.iter().enumerate() { + cum += p[i as usize]; + if cum >= top_p as f64 { + cut = n + 1; + break; + } + } + let nucleus = &idx[..cut.max(1)]; + let mass: f64 = nucleus.iter().map(|&i| p[i as usize]).sum(); + let mut r = rng.next_f64() * mass; + for &i in nucleus { + r -= p[i as usize]; + if r <= 0.0 { + return i; + } + } + nucleus[nucleus.len() - 1] } fn parse_args() -> Args { @@ -40,6 +125,12 @@ fn parse_args() -> Args { let mut prompt = "The capital of France is".to_string(); let mut max_tokens = 64usize; let mut raw = false; + // "" = MAPLE_POLICY's default (bf16). "q8" selects the block-quantized tier. + let mut kv_mode = String::new(); + let mut temp = 0.0f32; + let mut top_p = 0.95f32; + let mut seed = 0u64; + let mut head: Option = None; let mut i = 1; while i < argv.len() { match argv[i].as_str() { @@ -61,6 +152,30 @@ fn parse_args() -> Args { raw = true; i += 1; } + // KV storage tier: "q8" (default) or "bf16". Anything else warns + // and falls back to q8 via MAPLE_POLICY. + "--kv-mode" => { + kv_mode = argv[i + 1].clone(); + i += 2; + } + "--temp" => { + temp = argv[i + 1].parse().expect("--temp"); + i += 2; + } + "--top-p" => { + top_p = argv[i + 1].parse().expect("--top-p"); + i += 2; + } + "--seed" => { + seed = argv[i + 1].parse().expect("--seed"); + i += 2; + } + // Swap the lm_head without a second full model: point at a + // single-tensor .hfq from `hipfire-quantize --head-only`. + "--head" => { + head = Some(argv[i + 1].clone()); + i += 2; + } other => panic!("unknown arg {other}"), } } @@ -69,6 +184,11 @@ fn parse_args() -> Args { prompt, max_tokens, raw, + kv_mode, + temp, + top_p, + seed, + head, } } @@ -93,14 +213,32 @@ fn main() { args.prompt.clone() } else { format!( - "<|im_start|>user\n{}<|im_end|>\n<|im_start|>assistant\n", + // The trailing "\n" is REQUIRED and was missing until + // 2026-08-31. Maple's embedded jinja template ends its generation + // prompt with `'<|im_start|>assistant\n\n'`, and the vendor's + // llama.cpp README calls out `--jinja` as applying the template + // "exactly, including its thinking prefix". + // + // Without it the model has to emit the opening itself, so + // every generation starts off-distribution INSIDE the reasoning + // block — which is exactly where this model's degenerate loops + // occur. Any loop-rate measurement taken without this prefix is + // measuring a prompt frame the model was never trained on. + "<|im_start|>user\n{}<|im_end|>\n<|im_start|>assistant\n\n", args.prompt ) }; let prompt_toks = tokenizer.encode(&text); let max_seq = prompt_toks.len() + args.max_tokens + 64; - let mut b = load_maple_from_hfq(&mut hfq, &mut gpu, max_seq).expect("load maple bundle"); + let mut b = load_maple_from_hfq_with_head( + &mut hfq, + &mut gpu, + max_seq, + &args.kv_mode, + args.head.as_deref().map(std::path::Path::new), + ) + .expect("load maple bundle"); eprintln!( "maple: hidden={} layers={} experts={}/{} moe_inter={} vocab={} eos={} max_seq={}", b.config.hidden_size, @@ -157,12 +295,19 @@ fn main() { prompt_toks.len() as f64 / prefill_s ); - // Greedy decode. + // Decode. `--temp 0` (the default) is greedy and bit-for-bit reproduces the + // previous behaviour; `--temp > 0` samples with top-p and an EXPLICIT seed. + // + // The seed is what makes a loop-rate measurement possible at all: greedy + // gives exactly ONE draw per (prompt, model), so sample size can only grow + // with the prompt set and prompt dominates the variance. With a seed, the + // same prompt can be redrawn N times and the arms compared on equal terms. let mut out = String::new(); let t1 = std::time::Instant::now(); let mut n_gen = 0usize; + let mut rng = SplitMix64::new(args.seed); for _ in 0..args.max_tokens { - let (best, _) = + let tok = if args.temp <= 0.0 { logits .iter() .enumerate() @@ -172,8 +317,11 @@ fn main() { } else { acc } - }); - let tok = best as u32; + }) + .0 as u32 + } else { + sample_top_p(&logits, args.temp, args.top_p, &mut rng) + }; if tok == b.eos_tok { eprintln!("[eos]"); break; diff --git a/crates/hipfire-arch-maple/examples/maple_decode_profile.rs b/crates/hipfire-arch-maple/examples/maple_decode_profile.rs index 10d1979c3f..eeb54a0b86 100644 --- a/crates/hipfire-arch-maple/examples/maple_decode_profile.rs +++ b/crates/hipfire-arch-maple/examples/maple_decode_profile.rs @@ -288,7 +288,7 @@ fn main() { // from the headline token counts alone silently under-allocates and the // first symptom is an illegal-access fault in an unrelated kernel. let max_seq = prompt_toks.len() + args.warmup + args.gen + args.profile_gen + 160 + 64; - let mut b = load_maple_from_hfq(&mut hfq, &mut gpu, max_seq).expect("load maple bundle"); + let mut b = load_maple_from_hfq(&mut hfq, &mut gpu, max_seq, "").expect("load maple bundle"); eprintln!( "maple: hidden={} layers={} experts={}/{} moe_inter={} vocab={} max_seq={}", b.config.hidden_size, diff --git a/crates/hipfire-arch-maple/examples/maple_kld.rs b/crates/hipfire-arch-maple/examples/maple_kld.rs index edd24967ef..7668a76ee1 100644 --- a/crates/hipfire-arch-maple/examples/maple_kld.rs +++ b/crates/hipfire-arch-maple/examples/maple_kld.rs @@ -54,7 +54,8 @@ use std::path::Path; use std::time::Instant; const USAGE: &str = "usage: maple_kld --model --tokens \ - --ref [--dump ] [--per-pos ] [--limit N]"; + --ref [--dump ] [--per-pos ] [--limit N] \\ + [--kv-mode q8|bf16]"; struct Args { model: String, @@ -63,6 +64,11 @@ struct Args { dump: Option, per_pos: Option, limit: Option, + /// KV storage tier request, resolved through MAPLE_POLICY. "" = q8. + /// This is what lets the Q8-KV contribution to the measured KL be + /// SUBTRACTED rather than assumed: run the same tokens and the same + /// reference under q8 and bf16 and diff the results. + kv_mode: String, } fn parse_args() -> Args { @@ -74,6 +80,7 @@ fn parse_args() -> Args { dump: None, per_pos: None, limit: None, + kv_mode: String::new(), }; let mut i = 1; while i < argv.len() { @@ -90,6 +97,7 @@ fn parse_args() -> Args { "--dump" => a.dump = Some(val()), "--per-pos" => a.per_pos = Some(val()), "--limit" => a.limit = Some(val().parse().expect("--limit")), + "--kv-mode" => a.kv_mode = val(), other => panic!("unknown arg {other}\n{USAGE}"), } i += 2; @@ -187,7 +195,8 @@ fn main() { eprintln!("Loading weights from {}...", args.model); let t_load = Instant::now(); - let mut b = load_maple_from_hfq(&mut hfq, &mut gpu, n).expect("load maple bundle"); + let mut b = + load_maple_from_hfq(&mut hfq, &mut gpu, n, &args.kv_mode).expect("load maple bundle"); eprintln!("Loaded in {:.1}s", t_load.elapsed().as_secs_f64()); eprintln!( "maple: hidden={} layers={} experts={}/{} vocab={}", diff --git a/crates/hipfire-arch-maple/examples/maple_perplexity.rs b/crates/hipfire-arch-maple/examples/maple_perplexity.rs index 6d22783c21..91e340fb51 100644 --- a/crates/hipfire-arch-maple/examples/maple_perplexity.rs +++ b/crates/hipfire-arch-maple/examples/maple_perplexity.rs @@ -128,7 +128,7 @@ fn main() { eprintln!("Loading weights from {}...", args.model); let t_load = Instant::now(); - let mut b = load_maple_from_hfq(&mut hfq, &mut gpu, args.ctx).expect("load maple bundle"); + let mut b = load_maple_from_hfq(&mut hfq, &mut gpu, args.ctx, "").expect("load maple bundle"); eprintln!("Loaded in {:.1}s", t_load.elapsed().as_secs_f64()); eprintln!( "maple: hidden={} layers={} experts={}/{} vocab={}", diff --git a/crates/hipfire-arch-maple/examples/maple_prefill_parity.rs b/crates/hipfire-arch-maple/examples/maple_prefill_parity.rs index 028d0fb79a..96fbad813e 100644 --- a/crates/hipfire-arch-maple/examples/maple_prefill_parity.rs +++ b/crates/hipfire-arch-maple/examples/maple_prefill_parity.rs @@ -371,7 +371,7 @@ fn main() { // default sweep and (b) LEAK every copy but the last: `GpuTensor` has no // `Drop` and this crate frees explicitly via `free_gpu`, so a dropped // bundle's device memory is simply gone until the process exits. - let mut bundle = load_maple_from_hfq(&mut hfq, &mut gpu, n_tokens + 64).expect("load"); + let mut bundle = load_maple_from_hfq(&mut hfq, &mut gpu, n_tokens + 64, "").expect("load"); let mut want = Vec::with_capacity(n_tokens); for (p, &t) in tokens.iter().enumerate() { want.push( diff --git a/crates/hipfire-arch-maple/map.md b/crates/hipfire-arch-maple/map.md index bf330d290e..edef27450d 100644 --- a/crates/hipfire-arch-maple/map.md +++ b/crates/hipfire-arch-maple/map.md @@ -29,17 +29,17 @@ _Generated by `scripts/check-crate-maps.py` from the tree — do not edit inside | File | Lines | Public items | Tests | |---|---:|---:|---:| | [`src/batch.rs`](src/batch.rs) | 213 | 10 | 7 | -| [`src/bundle.rs`](src/bundle.rs) | 117 | 4 | 1 | -| [`src/carrier.rs`](src/carrier.rs) | 30 | 1 | 0 | +| [`src/bundle.rs`](src/bundle.rs) | 153 | 5 | 1 | +| [`src/carrier.rs`](src/carrier.rs) | 50 | 1 | 0 | | [`src/config.rs`](src/config.rs) | 396 | 13 | 9 | -| [`src/forward.rs`](src/forward.rs) | 1,370 | 3 | 5 | +| [`src/forward.rs`](src/forward.rs) | 1,376 | 3 | 5 | | [`src/lib.rs`](src/lib.rs) | 51 | 6 | 0 | -| [`src/maple.rs`](src/maple.rs) | 1,023 | 17 | 7 | +| [`src/maple.rs`](src/maple.rs) | 1,181 | 17 | 12 | ### Public API surface - [`src/batch.rs`](src/batch.rs): `MOE_GROUPED_BLOCK_M`, `MAPLE_PREFILL_CHUNK`, `MAPLE_PREFILL_MAX_B`, `dense_m_total`, `moe_grouped_m_total_bound`, `prefill_chunks`, `dense_slot_index_host`, `dense_tile_ids_host`, `dense_qt51_gemm`, `upload_single_expert_ptr_table` -- [`src/bundle.rs`](src/bundle.rs): `MapleBundle`, `MAPLE_EOS_FALLBACK`, `resolve_eos`, `load_maple_from_hfq` +- [`src/bundle.rs`](src/bundle.rs): `MapleBundle`, `MAPLE_EOS_FALLBACK`, `resolve_eos`, `load_maple_from_hfq`, `load_maple_from_hfq_with_head` - [`src/carrier.rs`](src/carrier.rs): `load_maple_bundle` - [`src/config.rs`](src/config.rs): `MapleLayerType`, `MAPLE_SWIGLU_CLAMP`, `MapleConfig`, `from_hfq`, `from_metadata_json`, `from_safetensors`, `from_config_value`, `q_dim`, `kv_dim`, `layer_type`, `applies_rope`, `rotary_dim`, +1 more - [`src/forward.rs`](src/forward.rs): `decode_step`, `forward_batch_supported`, `forward_batch` @@ -59,6 +59,6 @@ _Generated by `scripts/check-crate-maps.py` from the tree — do not edit inside ### Totals -- 7 modules · 3,200 lines · 54 public items · 29 tests · 6 examples +- 7 modules · 3,420 lines · 55 public items · 34 tests · 6 examples diff --git a/crates/hipfire-arch-maple/src/bundle.rs b/crates/hipfire-arch-maple/src/bundle.rs index e50bf02a5f..152eec1447 100644 --- a/crates/hipfire-arch-maple/src/bundle.rs +++ b/crates/hipfire-arch-maple/src/bundle.rs @@ -83,14 +83,50 @@ pub fn resolve_eos(tokenizer: &hipfire_runtime::tokenizer::Tokenizer) -> u32 { /// /// Split out from the carrier so an offline harness (the coherence example) /// can build the same bundle without going through the loader registry. +/// `kv_mode_raw` is the UNRESOLVED request string (`--kv-mode`, `""` for the +/// default). It is resolved here rather than by the caller because this is the +/// first point where `config.head_dim` exists, and `resolve` takes it. Modes +/// outside `MAPLE_POLICY`'s accept set fall back to q8 with a warning. pub fn load_maple_from_hfq( hfq: &mut HfqFile, gpu: &mut Gpu, max_seq: usize, + kv_mode_raw: &str, ) -> Result { + load_maple_from_hfq_with_head(hfq, gpu, max_seq, kv_mode_raw, None) +} + +/// `load_maple_from_hfq` with an optional HEAD OVERLAY: a single-tensor `.hfq` +/// built by `hipfire-quantize --head-only` whose `lm_head.weight` shadows the +/// base's. One 6.5 GB body then serves every head carrier, instead of shipping +/// a near-identical full model per carrier. +/// +/// Attached BEFORE `MapleWeights::load`, because the loader reads the head +/// through the same `find_tensor_info` path the overlay shadows — attaching +/// afterwards would silently load the base's head and produce a model that +/// looks right and is not the one requested. +pub fn load_maple_from_hfq_with_head( + hfq: &mut HfqFile, + gpu: &mut Gpu, + max_seq: usize, + kv_mode_raw: &str, + head_overlay: Option<&std::path::Path>, +) -> Result { + if let Some(head) = head_overlay { + hfq.attach_head_overlay(head)?; + } let config = MapleConfig::from_hfq(hfq)?; let weights = MapleWeights::load(hfq, &config, gpu)?; - let state = MapleState::new_with_max_seq(gpu, &config, max_seq)?; + let hipfire_runtime::kv_mode::ResolveResult { mode, warning } = + hipfire_runtime::kv_mode::resolve( + kv_mode_raw, + &hipfire_runtime::kv_mode::MAPLE_POLICY, + config.head_dim, + ); + if let Some(w) = warning { + eprintln!(" KV cache: {w} (site maple)"); + } + let state = MapleState::new_with_max_seq(gpu, &config, max_seq, mode)?; let tokenizer = hipfire_runtime::tokenizer::Tokenizer::from_hfq_metadata(&hfq.metadata_json) .map_err(|e| format!("maple: tokenizer not found: {e}"))?; let eos_tok = resolve_eos(&tokenizer); diff --git a/crates/hipfire-arch-maple/src/carrier.rs b/crates/hipfire-arch-maple/src/carrier.rs index bb7729d88a..7ad3390397 100644 --- a/crates/hipfire-arch-maple/src/carrier.rs +++ b/crates/hipfire-arch-maple/src/carrier.rs @@ -19,7 +19,27 @@ pub fn load_maple_bundle(src: ModelSource, ctx: &mut LoadCtx) -> Result1 unsupported via registry".into()); } match src { - ModelSource::Hfq(mut hfq) => load_maple_from_hfq(&mut hfq, ctx.gpu, ctx.max_seq), + ModelSource::Hfq(mut hfq) => { + // Same ladder the other carriers use: an explicit --kv-mode wins, + // else the global config value. Resolution against MAPLE_POLICY + // happens inside load_maple_from_hfq, where head_dim is known. + // Before this, arch 15 hardcoded q8 and --kv-mode was a silent + // no-op. + let raw = ctx + .kv_mode_override + .filter(|s| !s.is_empty()) + .map(|s| s.to_string()) + .unwrap_or_else(|| hipfire_runtime::config::get().kv_mode.clone()); + // A head overlay (hipfire-quantize --head-only) shadows the + // base's lm_head, so one body serves several head carriers. + crate::bundle::load_maple_from_hfq_with_head( + &mut hfq, + ctx.gpu, + ctx.max_seq, + &raw, + ctx.head_path.map(std::path::Path::new), + ) + } ModelSource::Dir(_) => Err( "maple: safetensors-directory loading is unsupported — convert first with \ `hipfire-quantize --format maple --input --output `, which packs \ diff --git a/crates/hipfire-arch-maple/src/forward.rs b/crates/hipfire-arch-maple/src/forward.rs index 0afd80a88c..67e2cbc1a9 100644 --- a/crates/hipfire-arch-maple/src/forward.rs +++ b/crates/hipfire-arch-maple/src/forward.rs @@ -177,6 +177,10 @@ fn decode_step_body( let plan = hipfire_dispatch::families::kv_tier::KvTierPlan::derive( hipfire_dispatch::families::kv_tier::KvTierInputs { pos: seq_len - 1, + // Only read on the Q8 arm. Under `--kv-mode bf16` the cache + // reports `quant_bf16` through `tier_inputs()`, `classify` + // resolves to KTier::Bf16 first, and that arm is windowed + // unconditionally — so this flag goes inert, not contradicted. q8_windowed: true, window, ..state.kv.tier_inputs() @@ -1234,6 +1238,8 @@ fn batched_attend( let plan = hipfire_dispatch::families::kv_tier::KvTierPlan::derive( hipfire_dispatch::families::kv_tier::KvTierInputs { pos: start_pos + b - 1, + // Inert under `--kv-mode bf16` — see the decode-side note; the + // Bf16 arm is windowed unconditionally in both shapes. q8_windowed: true, window, batch_size: b, diff --git a/crates/hipfire-arch-maple/src/lib.rs b/crates/hipfire-arch-maple/src/lib.rs index b236e12322..80c8be6e8d 100644 --- a/crates/hipfire-arch-maple/src/lib.rs +++ b/crates/hipfire-arch-maple/src/lib.rs @@ -39,7 +39,7 @@ pub mod config; pub mod forward; pub mod maple; -pub use bundle::{load_maple_from_hfq, MapleBundle}; +pub use bundle::{load_maple_from_hfq, load_maple_from_hfq_with_head, MapleBundle}; pub use carrier::load_maple_bundle; pub use forward::decode_step; diff --git a/crates/hipfire-arch-maple/src/maple.rs b/crates/hipfire-arch-maple/src/maple.rs index 548eb8d3b6..1170503843 100644 --- a/crates/hipfire-arch-maple/src/maple.rs +++ b/crates/hipfire-arch-maple/src/maple.rs @@ -24,6 +24,7 @@ use crate::config::MapleConfig; use hipfire_runtime::hfq::HfqFile; +use hipfire_runtime::kv_mode::KvMode; use hipfire_runtime::llama::{f16_to_f32, f32_to_f16, KvCache, WeightTensor}; use rdna_compute::{DType, Gpu, GpuTensor}; @@ -128,18 +129,23 @@ fn load_wt( /// quant_type → DType. **qt=51 (`MQ2G256LloydU`) is the whole point of this /// arch**: it is the unrotated MQ2-Lloyd sibling that carries Maple's native /// ternary weights losslessly, and the dispatcher must NOT rotate x for it. -fn wt_from_raw( - gpu: &mut Gpu, - qt: u8, - data: &[u8], - m: usize, - k: usize, -) -> Result { - let dtype = match qt { +/// quant_type -> DType for every carrier a Maple `.hfq` can hold. Pure, so the +/// READER contract is testable without a GPU. +/// +/// This map is append-only in practice: it is what lets an already-converted +/// model load, so an entry may not be removed just because the CONVERTER stops +/// producing that carrier. qt=30 is exactly that case — `--head-quant mq4` is +/// deprecated and no longer selectable, but every qt=30 `.hfq` already on disk +/// must keep loading. +pub(crate) fn maple_dtype_for_quant_type(qt: u8) -> Result { + Ok(match qt { 1 => DType::F16, 2 => DType::F32, 16 => DType::BF16, 3 => DType::Q8_0, + // qt=4 arrives from `--head-quant q4k`. GGML-compatible Q4_K and + // UNROTATED — no FWHT seed contract, unlike qt=30/44. + 4 => DType::Q4K, 13 => DType::MQ4G256, 15 => DType::MQ6G256, 19 => DType::MQ2G256Lloyd, @@ -150,9 +156,26 @@ fn wt_from_raw( // `pack_maple_head` quantized against; if the two ever diverge the // result is not an error but silently wrong logits. 30 => DType::MQ4G256Lloyd, + // qt=44 arrives only from `--head-quant mq4v2`. Same FWHT-rotated + // contract as qt=30 above — it resolves to GemvMq4G256V2Prerotated, so + // weight_gemv rotates x with the same ensure_mq_signs seeds (42/1042) + // that pack_maple_head quantized against. It differs from qt=30 only in + // the 8 header bytes: a separate fp16 scale/zero per 128-weight half + // rather than one pair per 256, at 4.25 bpw instead of 5.0. + 44 => DType::MQ4G256V2, 51 => DType::MQ2G256LloydU, other => return Err(format!("unsupported quant_type {other}")), - }; + }) +} + +fn wt_from_raw( + gpu: &mut Gpu, + qt: u8, + data: &[u8], + m: usize, + k: usize, +) -> Result { + let dtype = maple_dtype_for_quant_type(qt)?; let buf = gpu .upload_raw(data, &[data.len()]) .map_err(|e| format!("upload_raw: {e:?}"))?; @@ -668,16 +691,54 @@ pub struct MapleState { pub b_act_f16: GpuTensor, // [max_b × k_top × moe_inter] F16 } +/// Tile size the DECODE attention will actually use, for sizing +/// `flash_partials`. +/// +/// This used to be a hardcoded `128`, which is only correct on architectures +/// whose default tile IS 128. `q8_flash_tile_size` returns **32 on gfx1100** +/// (RDNA3), so the decode kernel there computes 4x as many tiles as a +/// 128-derived allocation assumes, and indexes +/// `partials + (h * max_tiles + tile_id) * (2 + head_dim)` against them. +/// +/// It did not overflow, because the trailing `FLASH_PREFILL_SUBBATCH` factor +/// left 64x of slack that absorbed the 4x — but that reduced the real margin +/// on RDNA3 to 16x for a reason nothing in the code stated, and it made this a +/// FOURTH independent copy of tile-size logic. `launch_asym_flash_batched` +/// carries a comment about "the corruption bug three independent copies of +/// this exact logic caused"; deriving the value is how that stops recurring. +/// +/// `HIPFIRE_Q8_FLASH_TILE` is honoured by `q8_flash_tile_size`, so an operator +/// override is now reflected in the allocation too rather than silently eating +/// the slack. +fn flash_partial_tile(gpu: &Gpu, cfg: &MapleConfig) -> usize { + rdna_compute::attention::q8_flash_tile_size( + &gpu.arch, + cfg.num_attention_heads, + cfg.num_key_value_heads, + cfg.head_dim, + // Shape-only: the tile policy reads max_seq solely to recognise one + // certified gfx1151 replay shape (max_seq == 2048), which Maple is not. + cfg.max_position_embeddings, + ) + .max(1) +} + impl MapleState { pub fn new(gpu: &mut Gpu, cfg: &MapleConfig) -> Result { let max_seq = cfg.max_position_embeddings.min(DEFAULT_MAX_SEQ); - Self::new_with_max_seq(gpu, cfg, max_seq) + Self::new_with_max_seq(gpu, cfg, max_seq, KvMode::Q8) } + /// `kv_mode` must already be resolved through `MAPLE_POLICY` — this is the + /// allocation site, not the policy site. Only `Q8` and `Bf16` are + /// serviceable; anything else is rejected rather than silently downgraded, + /// because the other tiers have no sliding-window attention kernel and + /// Maple's 3:1 sliding layers would then attend the full context. pub fn new_with_max_seq( gpu: &mut Gpu, cfg: &MapleConfig, max_seq: usize, + kv_mode: KvMode, ) -> Result { let hidden = cfg.hidden_size; let q_dim = cfg.q_dim(); @@ -691,13 +752,32 @@ impl MapleState { gpu.ensure_mq_signs() .map_err(|e| format!("maple: ensure_mq_signs: {e:?}"))?; - let kv = KvCache::new_gpu_q8( - gpu, - cfg.num_hidden_layers, - cfg.num_key_value_heads, - cfg.head_dim, - max_seq, - ) + let kv = match kv_mode { + KvMode::Q8 => KvCache::new_gpu_q8( + gpu, + cfg.num_hidden_layers, + cfg.num_key_value_heads, + cfg.head_dim, + max_seq, + ), + KvMode::Bf16 => KvCache::new_gpu_bf16( + gpu, + cfg.num_hidden_layers, + cfg.num_key_value_heads, + cfg.head_dim, + max_seq, + ), + // Unreachable through the carrier: MAPLE_POLICY accepts only + // {Q8, Bf16} and `resolve` falls back to the site default for + // everything else. Reject loudly rather than serve a tier whose + // windowed kernels do not exist. + other => { + return Err(format!( + "maple: KV mode {other:?} has no sliding-window attention kernel; \ + arch 15 supports q8 and bf16 only" + )) + } + } .map_err(|e| format!("maple: kv cache: {e:?}"))?; let pos_buf = gpu .hip @@ -751,7 +831,7 @@ impl MapleState { flash_partials: alloc( gpu, cfg.num_attention_heads - * max_seq.div_ceil(128) + * max_seq.div_ceil(flash_partial_tile(gpu, cfg)) * (2 + cfg.head_dim) * FLASH_PREFILL_SUBBATCH, "flash_partials", @@ -1021,3 +1101,81 @@ mod tests { assert_eq!(bf16_to_f32(0x4049), f32::from_bits(0x40490000)); } } + +#[cfg(test)] +mod head_carrier_tests { + use super::*; + + /// Deprecating a CONVERTER option must never stop an existing model from + /// loading. `--head-quant mq4` is gone, but qt=30 heads are on disk and + /// must still resolve — this is the guard that keeps the reader and the + /// producer decoupled. + #[test] + fn deprecated_qt30_head_still_loads() { + assert_eq!( + maple_dtype_for_quant_type(30).unwrap(), + DType::MQ4G256Lloyd, + "qt=30 is deprecated as a CONVERTER option, not as a readable carrier" + ); + } + + #[test] + fn mq4v2_head_carrier_resolves() { + assert_eq!( + maple_dtype_for_quant_type(44).unwrap(), + DType::MQ4G256V2, + "qt=44 is the replacement fast head; without this arm it fails at \ + load with 'unsupported quant_type 44'" + ); + } + + /// The three carriers the converter can still emit, plus the body tier. + #[test] + fn shipped_carriers_resolve() { + assert_eq!(maple_dtype_for_quant_type(16).unwrap(), DType::BF16); + assert_eq!(maple_dtype_for_quant_type(3).unwrap(), DType::Q8_0); + assert_eq!(maple_dtype_for_quant_type(44).unwrap(), DType::MQ4G256V2); + // qt=51 is the whole point of arch 15: the unrotated ternary body. + assert_eq!( + maple_dtype_for_quant_type(51).unwrap(), + DType::MQ2G256LloydU + ); + } + + /// An unknown carrier must FAIL rather than silently pick something — + /// these tiers differ in rotation, and a wrong guess yields plausible but + /// wrong logits with no error. + #[test] + fn unknown_quant_type_is_rejected() { + let e = maple_dtype_for_quant_type(200).unwrap_err(); + assert!(e.contains("unsupported quant_type 200"), "got {e}"); + } +} + +#[cfg(test)] +mod head_overlay_tests { + /// The head-overlay contract, pinned as prose because the failure it + /// guards against is silent. + /// + /// `--head` points at a single-tensor `.hfq` from + /// `hipfire-quantize --head-only`, attached via + /// `HfqFile::attach_head_overlay` BEFORE `MapleWeights::load`. Ordering is + /// load-bearing: the loader resolves `lm_head.weight` through the same + /// `find_tensor_info` path the overlay shadows, so attaching afterwards + /// would quietly serve the BASE's head and produce a model that looks + /// correct and is not the one requested. + /// + /// The overlay must contain ONLY `lm_head.weight`. Without that check, + /// passing a full model to `--head` succeeds: every name exists in the + /// base at a matching shape, so `attach_overlay`'s arch/name/shape guards + /// all pass and the model silently shadows itself. That was found by a + /// negative control, not by inspection. + #[test] + fn head_overlay_contract_is_documented() { + // Executable only as documentation; the behavioural coverage is the + // GPU path exercised in review (valid q4k/bf16 overlays attach and + // generate; a full model is refused with a message naming the first + // offending tensor). Kept so the contract travels with the code. + assert_eq!(super::LM_HEAD_TENSOR_NAME, "lm_head.weight"); + } +} diff --git a/crates/hipfire-arch-qwen2/map.md b/crates/hipfire-arch-qwen2/map.md index 87ee2eb77d..b82188ff5d 100644 --- a/crates/hipfire-arch-qwen2/map.md +++ b/crates/hipfire-arch-qwen2/map.md @@ -26,7 +26,7 @@ _Generated by `scripts/check-crate-maps.py` from the tree — do not edit inside | [`src/arch_model.rs`](src/arch_model.rs) | 45 | 0 | 0 | | [`src/carrier.rs`](src/carrier.rs) | 76 | 2 | 0 | | [`src/lib.rs`](src/lib.rs) | 85 | 5 | 0 | -| [`src/qwen2.rs`](src/qwen2.rs) | 2,355 | 23 | 7 | +| [`src/qwen2.rs`](src/qwen2.rs) | 2,356 | 23 | 7 | | [`src/spec_impl.rs`](src/spec_impl.rs) | 199 | 1 | 0 | ### Public API surface @@ -51,6 +51,6 @@ _Generated by `scripts/check-crate-maps.py` from the tree — do not edit inside ### Totals -- 6 modules · 2,849 lines · 32 public items · 8 tests · 4 examples +- 6 modules · 2,850 lines · 32 public items · 8 tests · 4 examples diff --git a/crates/hipfire-arch-qwen2/src/qwen2.rs b/crates/hipfire-arch-qwen2/src/qwen2.rs index 0535c2556f..2175a618ab 100644 --- a/crates/hipfire-arch-qwen2/src/qwen2.rs +++ b/crates/hipfire-arch-qwen2/src/qwen2.rs @@ -2125,6 +2125,7 @@ impl DenseArch for Qwen2Dense<'_> { quant_q4: false, quant_int8: false, quant_hfq8: false, + quant_bf16: false, f32_policy: F32AttnPolicy::Gqa { n_heads: k.n_heads, n_kv_heads: k.n_kv_heads, diff --git a/crates/hipfire-cli/map.md b/crates/hipfire-cli/map.md index d87967dec3..2b818d35cd 100644 --- a/crates/hipfire-cli/map.md +++ b/crates/hipfire-cli/map.md @@ -23,11 +23,11 @@ _Generated by `scripts/check-crate-maps.py` from the tree — do not edit inside | File | Lines | Public items | Tests | |---|---:|---:|---:| | [`src/bench_concurrency.rs`](src/bench_concurrency.rs) | 720 | 21 | 9 | -| [`src/main.rs`](src/main.rs) | 9,692 | 0 | 65 | +| [`src/main.rs`](src/main.rs) | 10,241 | 0 | 71 | | [`src/serve/complete.rs`](src/serve/complete.rs) | 6,754 | 0 | 89 | | [`src/serve/http.rs`](src/serve/http.rs) | 1,089 | 0 | 6 | | [`src/serve/metrics.rs`](src/serve/metrics.rs) | 328 | 0 | 5 | -| [`src/serve/mod.rs`](src/serve/mod.rs) | 2,116 | 2 | 16 | +| [`src/serve/mod.rs`](src/serve/mod.rs) | 2,118 | 2 | 16 | | [`src/setup.rs`](src/setup.rs) | 1,529 | 0 | 16 | ### Public API surface @@ -53,6 +53,6 @@ _Generated by `scripts/check-crate-maps.py` from the tree — do not edit inside ### Totals -- 7 modules · 22,228 lines · 23 public items · 206 tests · 0 examples +- 7 modules · 22,779 lines · 23 public items · 212 tests · 0 examples diff --git a/crates/hipfire-cli/src/main.rs b/crates/hipfire-cli/src/main.rs index 86764ac22c..9c62b9248a 100644 --- a/crates/hipfire-cli/src/main.rs +++ b/crates/hipfire-cli/src/main.rs @@ -356,7 +356,7 @@ struct TuiArgs { arguments: Vec, } -#[derive(Args, Debug)] +#[derive(Args, Debug, Clone)] struct RunArgs { /// Registry tag, local alias, filename, or model path. model: String, @@ -378,6 +378,11 @@ struct RunArgs { #[arg(long)] /// One-shot KV format override for this model load. kv_mode: Option, + #[arg(long)] + /// Select a published lm_head variant (see the registry's `heads`), e.g. + /// `--head q4k`. The overlay shadows the model's own head at load time; + /// omitting this uses the head baked into the model file. + head: Option, #[arg(long, value_parser = ["contiguous", "vmm"])] /// One-shot KV storage backend override for this model load. kv_backend: Option, @@ -1601,9 +1606,20 @@ pub(crate) fn pull_command(paths: &Paths, args: PullArgs) -> Result<()> { fs::create_dir_all(&paths.models) .with_context(|| format!("failed to create {}", paths.models.display()))?; let destination = paths.models.join(&entry.file); - if destination.exists() && !args.force { - eprintln!("Already downloaded: {}", destination.display()); + let needs_base = if args.force { + true + } else if destination.exists() { + if existing_artifact_valid(&destination, entry.sha256.as_deref(), entry.size_bytes) { + eprintln!("Already downloaded: {}", destination.display()); + false + } else { + eprintln!("Refreshing stale artifact: {}", destination.display()); + true + } } else { + true + }; + if needs_base { let url = artifact_url(entry, &entry.file); eprintln!("Pulling {tag} ({:.2} GB)...", entry.size_gb); download_verified( @@ -1623,9 +1639,15 @@ pub(crate) fn pull_command(paths: &Paths, args: PullArgs) -> Result<()> { continue; }; let destination = paths.models.join(&sidecar.file); - if destination.exists() { - eprintln!(" {label} sidecar already present: {}", sidecar.file); - continue; + if destination.exists() && !args.force { + if existing_artifact_valid(&destination, sidecar.sha256.as_deref(), sidecar.size_bytes) + { + eprintln!(" {label} sidecar already present: {}", sidecar.file); + continue; + } + eprintln!(" {label} sidecar stale, refreshing: {}", sidecar.file); + } else if destination.exists() && args.force { + // force always refreshes } eprintln!(" Fetching {label} sidecar: {}", sidecar.file); let url = artifact_url(entry, &sidecar.file); @@ -1639,6 +1661,25 @@ pub(crate) fn pull_command(paths: &Paths, args: PullArgs) -> Result<()> { eprintln!(" warning: {label} sidecar unavailable: {error:#}"); } } + for (name, sidecar) in &entry.heads { + let destination = paths.models.join(&sidecar.file); + if destination.exists() && !args.force { + if existing_artifact_valid(&destination, sidecar.sha256.as_deref(), sidecar.size_bytes) + { + eprintln!(" head {name} already present: {}", sidecar.file); + continue; + } + eprintln!(" head {name} stale, refreshing: {}", sidecar.file); + } + eprintln!(" Fetching head {name}: {}", sidecar.file); + download_verified( + &artifact_url(entry, &sidecar.file), + &destination, + sidecar.sha256.as_deref(), + sidecar.size_bytes, + true, + )?; + } println!("{}", paths.models.join(&entry.file).display()); Ok(()) } @@ -1768,6 +1809,30 @@ fn report_progress(downloaded: u64, total: Option, elapsed: Duration) { let _ = std::io::stderr().flush(); } +pub(crate) fn existing_artifact_valid( + path: &Path, + expected_sha256: Option<&str>, + expected_size: Option, +) -> bool { + let Ok(metadata) = fs::metadata(path) else { + return false; + }; + if let Some(expected) = expected_size { + if metadata.len() != expected { + return false; + } + } + if let Some(expected) = expected_sha256 { + let Ok(digest) = sha256_path(path) else { + return false; + }; + if !digest.eq_ignore_ascii_case(expected) { + return false; + } + } + true +} + fn rm_command(paths: &Paths, args: RmArgs) -> Result<()> { let loaded = load_registry(&paths.registry); let resolved = loaded.registry.model(&args.model); @@ -1785,6 +1850,13 @@ fn rm_command(paths: &Paths, args: RmArgs) -> Result<()> { .map(|sidecar| paths.models.join(&sidecar.file)) .filter(|path| path.is_file()), ); + targets.extend( + entry + .heads + .values() + .map(|sidecar| paths.models.join(&sidecar.file)) + .filter(|path| path.is_file()), + ); } if let (Some(parent), Some(file)) = ( path.parent(), @@ -1908,14 +1980,7 @@ fn run_command(paths: &Paths, args: RunArgs) -> Result<()> { }; let host = config_string(&resolved, "serve.host")?; let port = config_u64(&resolved, "serve.port")? as u16; - let force_local = process_truthy("HIPFIRE_LOCAL") - || args.image.is_some() - || args.kv_mode.is_some() - || args.kv_backend.is_some() - || args.speculation.is_some() - || args.model_draft.is_some() - || args.draft_max.is_some() - || args.dspark_conf_threshold.is_some(); + let force_local = run_should_force_local(&args); if !force_local && service_ready(&host, port, Duration::from_millis(150)) { return run_via_http( &host, @@ -1948,6 +2013,7 @@ fn run_command(paths: &Paths, args: RunArgs) -> Result<()> { max_tokens, args.kv_mode.as_deref(), args.kv_backend.as_deref(), + args.head.as_deref(), )?; let selector = args .speculation @@ -2097,6 +2163,18 @@ fn process_truthy(name: &str) -> bool { }) } +pub(crate) fn run_should_force_local(args: &RunArgs) -> bool { + process_truthy("HIPFIRE_LOCAL") + || args.image.is_some() + || args.kv_mode.is_some() + || args.kv_backend.is_some() + || args.head.is_some() + || args.speculation.is_some() + || args.model_draft.is_some() + || args.draft_max.is_some() + || args.dspark_conf_threshold.is_some() +} + #[allow(clippy::too_many_arguments)] fn run_via_http( host: &str, @@ -2492,16 +2570,17 @@ pub(crate) fn load_params( max_tokens: u64, kv_override: Option<&str>, kv_backend_override: Option<&str>, + head_override: Option<&str>, ) -> Result { let configured_max_seq = config_u64(resolved, "memory.max_seq")?; let max_seq = configured_max_seq.max(max_tokens.saturating_add(1024)); let configured_kv = config_string(resolved, "memory.kv_cache")?; let kv_mode = kv_override .map(str::to_owned) - .or_else(|| (configured_kv != "auto").then_some(configured_kv)) - .or_else(|| entry.and_then(|entry| entry.default_kv_mode.clone())) - .unwrap_or_else(|| "q8".into()); - // Validate a one-shot override through the shared schema. + .filter(|value| !value.is_empty()) + .unwrap_or(configured_kv); + // Validate through the shared schema. `auto` is preserved so architecture + // (maple vs qwen) can select BF16 vs Q8; do not substitute q8 here. field("memory.kv_cache") .expect("schema field") .parse_cli(&kv_mode)?; @@ -2526,6 +2605,45 @@ pub(crate) fn load_params( } } } + // Resolve --head against the registry's `heads` map. The overlay + // lives beside the model file, exactly like the triattn sidecar. Refuse + // rather than fall back: a silent fall-back would serve the base's head + // and answer a different question than the operator asked. + let head_file = match head_override.filter(|s| !s.is_empty()) { + None => String::new(), + // A direct path is accepted as well as a registry name: loading a + // model BY PATH has no registry entry, so names cannot resolve there + // and only a path can work. + Some(name) if Path::new(name).is_file() => name.to_string(), + Some(name) => { + let heads = entry.map(|e| &e.heads); + let sidecar = heads.and_then(|h| h.get(name)).ok_or_else(|| { + let known: Vec<&str> = heads + .map(|h| h.keys().map(String::as_str).collect()) + .unwrap_or_default(); + anyhow!( + "--head {name}: not a file, and this model has no such head variant{}", + if known.is_empty() { + " (it publishes none)".to_string() + } else { + format!(" (available: {})", known.join(", ")) + } + ) + })?; + let candidate = model_path + .parent() + .unwrap_or_else(|| Path::new(".")) + .join(&sidecar.file); + if !candidate.is_file() { + bail!( + "--head {name}: overlay {} not found — fetch it with \ + `hipfire pull` or place it beside the model", + candidate.display() + ); + } + candidate.display().to_string() + } + }; let mut params = serde_json::json!({ "max_seq": max_seq, "deepseek4_compute_placement": config_string( @@ -2545,6 +2663,7 @@ pub(crate) fn load_params( "ddtree_budget": config_u64(resolved, "speculation.ddtree_budget")?, "ddtree_topk": config_u64(resolved, "speculation.ddtree_topk")?, "cask_sidecar": cask_sidecar, + "head": head_file, "cask": config_bool(resolved, "memory.cask.enabled")?, "cask_budget": config_u64(resolved, "memory.cask.budget")?, "cask_beta": config_u64(resolved, "memory.cask.beta")?, @@ -3995,6 +4114,8 @@ fn open_bench_engine( max_tokens, args.kv_mode.as_deref(), args.kv_backend.as_deref(), + // No --head on this path yet; the model's own head is used. + None, )?; if let Some(selector) = args.speculation.as_deref() { apply_speculation_selector(&mut params, selector)?; @@ -6254,7 +6375,8 @@ mod tests { fs::write(&sidecar_path, b"sidecar").unwrap(); let defaults = resolve(Vec::::new()).unwrap(); - let params = load_params(&defaults, Some(entry), &model_path, 64, None, None).unwrap(); + let params = + load_params(&defaults, Some(entry), &model_path, 64, None, None, None).unwrap(); assert_eq!(params["cask"], false); assert_eq!(params["cask_handoff_tokens"], 0); assert_eq!(params["cask_sidecar"], ""); @@ -6269,7 +6391,7 @@ mod tests { layer: explicit, }]) .unwrap(); - let params = load_params(&enabled, Some(entry), &model_path, 64, None, None).unwrap(); + let params = load_params(&enabled, Some(entry), &model_path, 64, None, None, None).unwrap(); assert_eq!(params["cask"], false); assert_eq!(params["cask_sidecar"], sidecar_path.display().to_string()); assert_eq!(params["prefill_compression"], "off"); @@ -6280,8 +6402,16 @@ mod tests { pub(crate) fn load_params_forwards_explicit_vmm_backend() { let defaults = resolve(Vec::::new()).unwrap(); let model_path = PathBuf::from("/tmp/test-model.mq4"); - let params = - load_params(&defaults, None, &model_path, 64, Some("q8"), Some("vmm")).unwrap(); + let params = load_params( + &defaults, + None, + &model_path, + 64, + Some("q8"), + Some("vmm"), + None, + ) + .unwrap(); assert_eq!(params["kv_backend"], "vmm"); } @@ -6289,7 +6419,7 @@ mod tests { pub(crate) fn load_params_defaults_to_schema_contiguous_backend() { let defaults = resolve(Vec::::new()).unwrap(); let model_path = PathBuf::from("/tmp/test-model.mq4"); - let params = load_params(&defaults, None, &model_path, 64, Some("q8"), None).unwrap(); + let params = load_params(&defaults, None, &model_path, 64, Some("q8"), None, None).unwrap(); assert_eq!(params["kv_backend"], "contiguous"); assert_eq!(params["max_seq"], 32768); } @@ -6300,7 +6430,7 @@ mod tests { fs::create_dir_all(&paths.root).unwrap(); let raw = r#"{ "schema_version":1, - "generated_at":"now", + "generated_at":"2026-09-01T00:00:00Z", "models":{ "qwen3.5:4b":{"repo":"x","file":"qwen3.5-4b.mq4","size_gb":1,"min_vram_gb":1,"desc":"x","default_kv_mode":"q8"}, "qwen3.6:35b-a3b":{"repo":"x","file":"qwen3.6-35b.mq4","size_gb":1,"min_vram_gb":1,"desc":"x"}, @@ -6380,7 +6510,7 @@ mod tests { fs::create_dir_all(&paths.root).unwrap(); let raw = r#"{ "schema_version":1, - "generated_at":"now", + "generated_at":"2026-09-01T00:00:00Z", "models":{ "muse-glimmer":{"repo":"x","file":"muse-glimmer-30b.mq4","size_gb":1,"min_vram_gb":1,"desc":"x"}, "muse-glimmer:fast":{"repo":"x","file":"muse-glimmer-30b.mq4r","size_gb":1,"min_vram_gb":1,"desc":"x"}, @@ -6532,7 +6662,7 @@ mod tests { fs::create_dir_all(&paths.root).unwrap(); let raw = r#"{ "schema_version":1, - "generated_at":"now", + "generated_at":"2026-09-01T00:00:00Z", "models":{"qwen3.8:27b":{"repo":"x","file":"qwen3.8-27b.mq4","size_gb":1,"min_vram_gb":1,"desc":"x"}}, "aliases":{} }"#; @@ -6591,19 +6721,28 @@ mod tests { 64, Some("q8"), Some("contiguous"), + None, ) .unwrap(); assert_eq!(params["kv_backend"], "contiguous"); // Without explicit override, load_params uses the resolved vmm. - let params2 = - load_params(&resolved, Some(entry), &model_path, 64, Some("q8"), None).unwrap(); + let params2 = load_params( + &resolved, + Some(entry), + &model_path, + 64, + Some("q8"), + None, + None, + ) + .unwrap(); assert_eq!(params2["kv_backend"], "vmm"); assert_eq!(params2["max_seq"], 262144); // Glimmer target likewise overridable (backend + max_seq). let raw2 = r#"{ "schema_version":1, - "generated_at":"now", + "generated_at":"2026-09-01T00:00:00Z", "models":{"muse-glimmer":{"repo":"x","file":"muse-glimmer-30b.mq4","size_gb":1,"min_vram_gb":1,"desc":"x"}}, "aliases":{} }"#; @@ -6647,7 +6786,7 @@ mod tests { // DeepSeek target override wins over 1M/384Ki policy. let raw3 = r#"{ "schema_version":1, - "generated_at":"now", + "generated_at":"2026-09-01T00:00:00Z", "models":{"deepseek-v4-flash":{"repo":"x","file":"ds4.mq2r","size_gb":1,"min_vram_gb":1,"desc":"x"}}, "aliases":{} }"#; @@ -6695,7 +6834,7 @@ mod tests { pub(crate) fn load_params_only_forwards_explicit_deepseek4_expert_fanout() { let model_path = PathBuf::from("/tmp/test-model.mq2r"); let defaults = resolve(Vec::::new()).unwrap(); - let params = load_params(&defaults, None, &model_path, 64, Some("q8"), None).unwrap(); + let params = load_params(&defaults, None, &model_path, 64, Some("q8"), None, None).unwrap(); assert_eq!(params["deepseek4_compute_placement"], "single"); assert!(params.get("deepseek4_experts_per_token").is_none()); @@ -6710,7 +6849,7 @@ mod tests { layer: explicit, }]) .unwrap(); - let params = load_params(&resolved, None, &model_path, 64, Some("q8"), None).unwrap(); + let params = load_params(&resolved, None, &model_path, 64, Some("q8"), None, None).unwrap(); assert_eq!(params["deepseek4_experts_per_token"], 4); } @@ -6735,6 +6874,7 @@ mod tests { 64, Some("q8"), None, + None, ) .unwrap(); assert_eq!(params["deepseek4_compute_placement"], raw); @@ -6756,7 +6896,7 @@ mod tests { .unwrap(); let model_path = PathBuf::from("/tmp/test-model.mq4"); - let params = load_params(&resolved, None, &model_path, 64, Some("q8"), None).unwrap(); + let params = load_params(&resolved, None, &model_path, 64, Some("q8"), None, None).unwrap(); assert_eq!(params["draft"], draft); } @@ -6780,7 +6920,8 @@ mod tests { let model_path = PathBuf::from("/tmp/test-model.mq4"); // load_params alone must not carry the draft while config mode is off. - let mut params = load_params(&resolved, None, &model_path, 64, Some("q8"), None).unwrap(); + let mut params = + load_params(&resolved, None, &model_path, 64, Some("q8"), None, None).unwrap(); assert_eq!(params["dflash_mode"], "off"); assert!( params.get("draft").is_none(), @@ -9689,4 +9830,412 @@ mod tests { .unwrap_err(); assert!(format!("{err}").contains("must be between 0 and 393216")); } + + #[test] + fn head_forces_local_even_when_service_would_be_ready() { + let with_head = RunArgs { + model: "maple-preview".into(), + prompt: vec![], + temp: None, + top_p: None, + repeat_penalty: None, + max_tokens: None, + kv_mode: None, + head: Some("q4k".into()), + kv_backend: None, + speculation: None, + model_draft: None, + draft_max: None, + dspark_conf_threshold: None, + system: None, + image: None, + json: false, + no_stream: false, + }; + let without_head = RunArgs { + head: None, + ..with_head.clone() + }; + // --head must force local; without head should not force local by itself + assert!( + run_should_force_local(&with_head), + "--head must force local load path" + ); + assert!( + !run_should_force_local(&without_head), + "without head and no other flags should not force local" + ); + // Verify the helper is used by run_command: an HTTP service would be bypassed. + // No network needed; the flag alone is the contract. + } + + #[test] + fn existing_artifact_valid_detects_fresh_and_stale() { + use sha2::{Digest, Sha256}; + let dir = env::temp_dir().join(format!( + "hipfire-artifact-valid-{}-{}", + std::process::id(), + std::time::SystemTime::now() + .duration_since(std::time::UNIX_EPOCH) + .unwrap() + .as_nanos() + )); + fs::create_dir_all(&dir).unwrap(); + let path = dir.join("model.mq4"); + let content = b"fresh content"; + fs::write(&path, content).unwrap(); + let mut hasher = Sha256::new(); + hasher.update(content); + let sha = format!("{:x}", hasher.finalize()); + let size = content.len() as u64; + assert!(existing_artifact_valid(&path, Some(&sha), Some(size))); + assert!(!existing_artifact_valid(&path, Some(&sha), Some(size + 1))); + assert!(!existing_artifact_valid( + &path, + Some("deadbeefdeadbeefdeadbeefdeadbeefdeadbeefdeadbeefdeadbeefdeadbeef"), + Some(size) + )); + // No expectations means existence alone is valid + assert!(existing_artifact_valid(&path, None, None)); + // Missing file is invalid + assert!(!existing_artifact_valid( + &dir.join("missing"), + Some(&sha), + Some(size) + )); + fs::remove_dir_all(&dir).unwrap(); + } + + #[test] + fn load_params_preserves_auto_for_direct_path_and_registry() { + // Direct-path load: no registry entry, config is auto -> must stay auto. + let defaults = resolve(Vec::::new()).unwrap(); + assert_eq!(config_string(&defaults, "memory.kv_cache").unwrap(), "auto"); + let direct_path = PathBuf::from("/tmp/direct-model.mq4"); + let params = load_params(&defaults, None, &direct_path, 64, None, None, None).unwrap(); + assert_eq!( + params["kv_mode"], "auto", + "direct-path auto must survive to architecture" + ); + // Registry path with default_kv_mode=bf16 must also preserve auto when + // no explicit --kv-mode is given; architecture picks BF16. + let raw = r#"{ + "schema_version":1, + "generated_at":"2099-01-01T00:00:00Z", + "models":{ + "maple-preview":{"repo":"x","file":"maple-preview.mq2lloydu","size_gb":1,"min_vram_gb":1,"desc":"x","default_kv_mode":"bf16"} + }, + "aliases":{} + }"#; + let registry = RegistryV1::parse(raw, "test").unwrap(); + let (_, entry) = registry.model("maple-preview").unwrap(); + let params2 = + load_params(&defaults, Some(entry), &direct_path, 64, None, None, None).unwrap(); + assert_eq!( + params2["kv_mode"], "auto", + "registry auto must survive even when entry has bf16 default" + ); + // Explicit override still wins + let params3 = load_params( + &defaults, + Some(entry), + &direct_path, + 64, + Some("q8"), + None, + None, + ) + .unwrap(); + assert_eq!(params3["kv_mode"], "q8"); + } + + fn write_test_registry_cache(paths: &Paths, raw: &str) { + let registry = RegistryV1::parse(raw, "test-cache").unwrap(); + let url = env::var("HIPFIRE_REGISTRY_URL") + .unwrap_or_else(|_| "https://example.com/test.json".into()); + let cache = serde_json::json!({ + "fetched_at": unix_timestamp() * 1000, + "url": url, + "registry": registry + }); + fs::create_dir_all(paths.registry.cache.parent().unwrap()).unwrap(); + fs::write( + &paths.registry.cache, + serde_json::to_string(&cache).unwrap(), + ) + .unwrap(); + } + + fn tiny_http_server( + files: std::collections::HashMap>, + ) -> (String, std::thread::JoinHandle<()>) { + use std::io::{Read, Write}; + use std::net::TcpListener; + let listener = TcpListener::bind("127.0.0.1:0").unwrap(); + let addr = listener.local_addr().unwrap(); + let base = format!("http://127.0.0.1:{}", addr.port()); + let expected = files.len(); + let handle = std::thread::spawn(move || { + let mut served = 0usize; + listener.set_nonblocking(false).unwrap(); + for stream in listener.incoming() { + let mut stream = match stream { + Ok(s) => s, + Err(_) => break, + }; + let mut buf = [0u8; 4096]; + let _ = stream.set_read_timeout(Some(Duration::from_secs(2))); + let n = stream.read(&mut buf).unwrap_or(0); + let req = String::from_utf8_lossy(&buf[..n]); + // Extract file name from request path: /{repo}/resolve/main/{file} + let mut body: Option> = None; + for (name, data) in &files { + if req.contains(name) { + body = Some(data.clone()); + break; + } + } + if let Some(data) = body { + let header = format!( + "HTTP/1.1 200 OK\r\nContent-Length: {}\r\nConnection: close\r\n\r\n", + data.len() + ); + let _ = stream.write_all(header.as_bytes()); + let _ = stream.write_all(&data); + } else { + let header = + "HTTP/1.1 404 Not Found\r\nContent-Length: 0\r\nConnection: close\r\n\r\n"; + let _ = stream.write_all(header.as_bytes()); + } + let _ = stream.flush(); + served += 1; + if served >= expected { + break; + } + } + }); + // small pause to let listener start + std::thread::sleep(Duration::from_millis(50)); + (base, handle) + } + + #[test] + fn pull_fresh_downloads_heads_with_hash_verification() { + let _env_lock = TEST_ENV_LOCK.lock().unwrap(); + let _reg = EnvGuard::set("HIPFIRE_REGISTRY_URL", "https://example.com/test.json"); + let paths = test_paths("pull-fresh-heads"); + // Build tiny artifacts + let base_content = b"base-model-content"; + let head_q4k_content = b"head-q4k-content"; + let head_bf16_content = b"head-bf16-content"; + use sha2::{Digest, Sha256}; + let sha = |data: &[u8]| { + let mut h = Sha256::new(); + h.update(data); + format!("{:x}", h.finalize()) + }; + let raw = format!( + r#"{{ + "schema_version":1, + "generated_at":"2099-01-01T00:00:00Z", + "models":{{ + "test-model":{{ + "repo":"test/repo", + "file":"test-model.mq4", + "size_gb":0.001, + "min_vram_gb":1, + "desc":"x", + "default_kv_mode":"bf16", + "heads":{{ + "q4k":{{"file":"test-model-head-q4k.hfq","sha256":"{}","size_bytes":{}}}, + "bf16":{{"file":"test-model-head-bf16.hfq","sha256":"{}","size_bytes":{}}} + }}, + "sha256":"{}", + "size_bytes":{} + }} + }}, + "aliases":{{}} + }}"#, + sha(head_q4k_content), + head_q4k_content.len(), + sha(head_bf16_content), + head_bf16_content.len(), + sha(base_content), + base_content.len() + ); + write_test_registry_cache(&paths, &raw); + let mut files = std::collections::HashMap::new(); + files.insert("test-model.mq4".to_string(), base_content.to_vec()); + files.insert( + "test-model-head-q4k.hfq".to_string(), + head_q4k_content.to_vec(), + ); + files.insert( + "test-model-head-bf16.hfq".to_string(), + head_bf16_content.to_vec(), + ); + let (base, handle) = tiny_http_server(files); + let _hf = EnvGuard::set("HIPFIRE_HF_BASE", &base); + // Pull should download base + both heads + pull_command( + &paths, + PullArgs { + model: "test-model".into(), + force: false, + }, + ) + .unwrap(); + assert_eq!( + fs::read(paths.models.join("test-model.mq4")).unwrap(), + base_content + ); + assert_eq!( + fs::read(paths.models.join("test-model-head-q4k.hfq")).unwrap(), + head_q4k_content + ); + assert_eq!( + fs::read(paths.models.join("test-model-head-bf16.hfq")).unwrap(), + head_bf16_content + ); + let _ = handle.join(); + fs::remove_dir_all(&paths.root).unwrap(); + } + + #[test] + fn pull_stale_same_name_artifact_refreshes_atomically() { + let _env_lock = TEST_ENV_LOCK.lock().unwrap(); + let _reg = EnvGuard::set("HIPFIRE_REGISTRY_URL", "https://example.com/test.json"); + let paths = test_paths("pull-stale-refresh"); + let fresh = b"fresh-content"; + let stale = b"stale-old-content"; + use sha2::{Digest, Sha256}; + let sha = |data: &[u8]| { + let mut h = Sha256::new(); + h.update(data); + format!("{:x}", h.finalize()) + }; + let raw = format!( + r#"{{ + "schema_version":1, + "generated_at":"2099-01-01T00:00:00Z", + "models":{{ + "test-model":{{ + "repo":"test/repo", + "file":"test-model.mq4", + "size_gb":0.001, + "min_vram_gb":1, + "desc":"x", + "sha256":"{}", + "size_bytes":{} + }} + }}, + "aliases":{{}} + }}"#, + sha(fresh), + fresh.len() + ); + write_test_registry_cache(&paths, &raw); + fs::create_dir_all(&paths.models).unwrap(); + // Place stale artifact with same name but wrong hash/size + fs::write(paths.models.join("test-model.mq4"), stale).unwrap(); + assert!(!existing_artifact_valid( + &paths.models.join("test-model.mq4"), + Some(&sha(fresh)), + Some(fresh.len() as u64) + )); + let mut files = std::collections::HashMap::new(); + files.insert("test-model.mq4".to_string(), fresh.to_vec()); + let (base, handle) = tiny_http_server(files); + let _hf = EnvGuard::set("HIPFIRE_HF_BASE", &base); + // Without --force, stale should still be detected and refreshed + pull_command( + &paths, + PullArgs { + model: "test-model".into(), + force: false, + }, + ) + .unwrap(); + assert_eq!( + fs::read(paths.models.join("test-model.mq4")).unwrap(), + fresh + ); + let _ = handle.join(); + fs::remove_dir_all(&paths.root).unwrap(); + } + + #[test] + fn rm_removes_heads_alongside_base() { + let _env_lock = TEST_ENV_LOCK.lock().unwrap(); + let _reg = EnvGuard::set("HIPFIRE_REGISTRY_URL", "https://example.com/test.json"); + let paths = test_paths("rm-heads"); + let valid_sha = "e3b0c44298fc1c149afbf4c8996fb92427ae41e4649b934ca495991b7852b855"; + let raw = format!( + r#"{{ + "schema_version":1, + "generated_at":"2099-01-01T00:00:00Z", + "models":{{ + "test-model":{{ + "repo":"test/repo", + "file":"test-model.mq4", + "size_gb":0.001, + "min_vram_gb":1, + "desc":"x", + "heads":{{ + "q4k":{{"file":"test-model-head-q4k.hfq","sha256":"{sha}","size_bytes":3}}, + "bf16":{{"file":"test-model-head-bf16.hfq","sha256":"{sha}","size_bytes":3}} + }}, + "sha256":"{sha}", + "size_bytes":3 + }} + }}, + "aliases":{{}} + }}"#, + sha = valid_sha + ); + write_test_registry_cache(&paths, &raw); + fs::create_dir_all(&paths.models).unwrap(); + fs::write(paths.models.join("test-model.mq4"), b"base").unwrap(); + fs::write(paths.models.join("test-model-head-q4k.hfq"), b"q4k").unwrap(); + fs::write(paths.models.join("test-model-head-bf16.hfq"), b"bf16").unwrap(); + assert!(paths.models.join("test-model-head-q4k.hfq").is_file()); + rm_command( + &paths, + RmArgs { + model: "test-model".into(), + yes: true, + }, + ) + .unwrap(); + assert!(!paths.models.join("test-model.mq4").exists()); + assert!(!paths.models.join("test-model-head-q4k.hfq").exists()); + assert!(!paths.models.join("test-model-head-bf16.hfq").exists()); + fs::remove_dir_all(&paths.root).unwrap(); + } + + static TEST_ENV_LOCK: std::sync::Mutex<()> = std::sync::Mutex::new(()); + + struct EnvGuard { + key: String, + prev: Option, + } + impl EnvGuard { + fn set(key: &str, val: &str) -> Self { + let prev = env::var_os(key); + env::set_var(key, val); + Self { + key: key.to_string(), + prev, + } + } + } + impl Drop for EnvGuard { + fn drop(&mut self) { + if let Some(v) = &self.prev { + env::set_var(&self.key, v); + } else { + env::remove_var(&self.key); + } + } + } } diff --git a/crates/hipfire-cli/src/serve/mod.rs b/crates/hipfire-cli/src/serve/mod.rs index 70d1d8cc6a..a186732cf6 100644 --- a/crates/hipfire-cli/src/serve/mod.rs +++ b/crates/hipfire-cli/src/serve/mod.rs @@ -1147,6 +1147,8 @@ impl ServeRuntime { max_tokens, self.kv_override.as_deref(), self.kv_backend_override.as_deref(), + // serve has no --head yet; models load their own head. + None, )?; if let Some(tp) = self.tp { params["tp"] = serde_json::json!(tp); diff --git a/crates/hipfire-config/map.md b/crates/hipfire-config/map.md index e34f58bbd8..ed8c9fd808 100644 --- a/crates/hipfire-config/map.md +++ b/crates/hipfire-config/map.md @@ -23,7 +23,7 @@ _Generated by `scripts/check-crate-maps.py` from the tree — do not edit inside | File | Lines | Public items | Tests | |---|---:|---:|---:| | [`src/bin/hipfire-rocm-resolve.rs`](src/bin/hipfire-rocm-resolve.rs) | 105 | 0 | 0 | -| [`src/lib.rs`](src/lib.rs) | 5,159 | 79 | 28 | +| [`src/lib.rs`](src/lib.rs) | 5,164 | 79 | 28 | | [`src/rocm.rs`](src/rocm.rs) | 2,460 | 39 | 37 | ### Public API surface @@ -45,6 +45,6 @@ _Generated by `scripts/check-crate-maps.py` from the tree — do not edit inside ### Totals -- 3 modules · 7,724 lines · 118 public items · 65 tests · 0 examples +- 3 modules · 7,729 lines · 118 public items · 65 tests · 0 examples diff --git a/crates/hipfire-config/src/lib.rs b/crates/hipfire-config/src/lib.rs index 1de90bb0b9..7529551eb5 100644 --- a/crates/hipfire-config/src/lib.rs +++ b/crates/hipfire-config/src/lib.rs @@ -477,9 +477,14 @@ fn expand_tilde(value: &str) -> PathBuf { PathBuf::from(value) } +// The union of every KV-mode name any SITE accepts. This is the config +// schema's allow-list only — it is NOT a promise that a given model supports a +// mode. Per-site acceptance lives in `hipfire_runtime::kv_mode`'s policies, +// which warn and fall back for anything they cannot allocate. `bf16` is +// currently maple-only (arch 15). const KV_MODES: &[&str] = &[ - "auto", "f32", "f16", "q8", "asym4", "asym3", "asym2", "fwht4", "fwht3", "fwht2", "turbo", - "turbo4", "turbo3", "turbo2", + "auto", "f32", "f16", "bf16", "q8", "asym4", "asym3", "asym2", "fwht4", "fwht3", "fwht2", + "turbo", "turbo4", "turbo3", "turbo2", ]; const AUTO_ON_OFF: &[&str] = &["auto", "on", "off"]; // `off` disables thinking outright. It resolves to a cap of 1, the engine's diff --git a/crates/hipfire-daemon/map.md b/crates/hipfire-daemon/map.md index fe21e0eadf..edcc64e898 100644 --- a/crates/hipfire-daemon/map.md +++ b/crates/hipfire-daemon/map.md @@ -23,7 +23,7 @@ _Generated by `scripts/check-crate-maps.py` from the tree — do not edit inside | File | Lines | Public items | Tests | |---|---:|---:|---:| -| [`src/main.rs`](src/main.rs) | 4,547 | 1 | 3 | +| [`src/main.rs`](src/main.rs) | 4,557 | 1 | 3 | | [`src/slots.rs`](src/slots.rs) | 1,559 | 23 | 14 | ### Public API surface @@ -44,6 +44,6 @@ _Generated by `scripts/check-crate-maps.py` from the tree — do not edit inside ### Totals -- 2 modules · 6,106 lines · 24 public items · 17 tests · 0 examples +- 2 modules · 6,116 lines · 24 public items · 17 tests · 0 examples diff --git a/crates/hipfire-daemon/src/main.rs b/crates/hipfire-daemon/src/main.rs index 4ff597cc3a..295dc437dc 100644 --- a/crates/hipfire-daemon/src/main.rs +++ b/crates/hipfire-daemon/src/main.rs @@ -1628,6 +1628,15 @@ fn main() { } else { hipfire_loader::GEMMA4_EAGLE_DRAFT_LEN }; + // Path to a head overlay (`hipfire-quantize --head-only`), + // resolved by the CLI from the registry's `heads` map. Empty + // means "use the head baked into the model file". + let head_path = msg + .get("params") + .and_then(|p| p.get("head")) + .and_then(|v| v.as_str()) + .filter(|s| !s.is_empty()) + .map(|s| s.to_string()); let kv_mode_override = msg .get("params") .and_then(|p| p.get("kv_mode")) @@ -2016,6 +2025,7 @@ fn main() { deepseek4_experts_per_token, deepseek4_compute_placement, draft_path.as_deref(), + head_path.as_deref(), gemma4_drafter.as_deref(), gemma4_draft_len, kv_mode_override.as_deref(), diff --git a/crates/hipfire-dispatch-tests/map.md b/crates/hipfire-dispatch-tests/map.md index 476c8e6eb7..a672c5c134 100644 --- a/crates/hipfire-dispatch-tests/map.md +++ b/crates/hipfire-dispatch-tests/map.md @@ -26,7 +26,7 @@ _Generated by `scripts/check-crate-maps.py` from the tree — do not edit inside | [`src/deepseek4.rs`](src/deepseek4.rs) | 50 | 0 | 5 | | [`src/dtype.rs`](src/dtype.rs) | 181 | 0 | 13 | | [`src/lib.rs`](src/lib.rs) | 19 | 0 | 0 | -| [`src/llama.rs`](src/llama.rs) | 307 | 0 | 16 | +| [`src/llama.rs`](src/llama.rs) | 308 | 0 | 16 | | [`src/qwen2.rs`](src/qwen2.rs) | 60 | 0 | 5 | | [`src/qwen35.rs`](src/qwen35.rs) | 280 | 0 | 18 | @@ -53,6 +53,6 @@ _Generated by `scripts/check-crate-maps.py` from the tree — do not edit inside ### Totals -- 7 modules · 1,233 lines · 0 public items · 79 tests · 0 examples +- 7 modules · 1,234 lines · 0 public items · 79 tests · 0 examples diff --git a/crates/hipfire-dispatch-tests/src/llama.rs b/crates/hipfire-dispatch-tests/src/llama.rs index f5fb85fea9..38a9a509fa 100644 --- a/crates/hipfire-dispatch-tests/src/llama.rs +++ b/crates/hipfire-dispatch-tests/src/llama.rs @@ -24,6 +24,7 @@ fn tier_inputs_base() -> hipfire_dispatch::families::kv_tier::KvTierInputs { quant_q4: false, quant_int8: false, quant_hfq8: false, + quant_bf16: false, f32_policy: F32AttnPolicy::Simple, v_mode_bits: 8, pos: 0, diff --git a/crates/hipfire-dispatch/map.md b/crates/hipfire-dispatch/map.md index 6852915c96..a21e23ad63 100644 --- a/crates/hipfire-dispatch/map.md +++ b/crates/hipfire-dispatch/map.md @@ -24,11 +24,11 @@ _Generated by `scripts/check-crate-maps.py` from the tree — do not edit inside |---|---:|---:|---:| | [`src/context.rs`](src/context.rs) | 67 | 5 | 0 | | [`src/coverage_tests.rs`](src/coverage_tests.rs) | 1,841 | 0 | 21 | -| [`src/families/attention.rs`](src/families/attention.rs) | 2,239 | 9 | 10 | +| [`src/families/attention.rs`](src/families/attention.rs) | 2,333 | 9 | 10 | | [`src/families/fused_qkv.rs`](src/families/fused_qkv.rs) | 1,585 | 8 | 1 | | [`src/families/gemm.rs`](src/families/gemm.rs) | 649 | 7 | 2 | | [`src/families/gemv.rs`](src/families/gemv.rs) | 624 | 20 | 0 | -| [`src/families/kv_tier.rs`](src/families/kv_tier.rs) | 1,130 | 11 | 42 | +| [`src/families/kv_tier.rs`](src/families/kv_tier.rs) | 1,288 | 11 | 46 | | [`src/families/mod.rs`](src/families/mod.rs) | 50 | 8 | 0 | | [`src/families/moe.rs`](src/families/moe.rs) | 1,210 | 24 | 10 | | [`src/families/moe_buckets.rs`](src/families/moe_buckets.rs) | 88 | 2 | 3 | @@ -44,7 +44,7 @@ _Generated by `scripts/check-crate-maps.py` from the tree — do not edit inside | [`src/pipeline/steps.rs`](src/pipeline/steps.rs) | 1,559 | 5 | 13 | | [`src/pipeline/superop.rs`](src/pipeline/superop.rs) | 598 | 17 | 4 | | [`src/resource/mod.rs`](src/resource/mod.rs) | 19 | 3 | 0 | -| [`src/tables/attention_table.rs`](src/tables/attention_table.rs) | 448 | 1 | 0 | +| [`src/tables/attention_table.rs`](src/tables/attention_table.rs) | 468 | 1 | 0 | | [`src/tables/fused_qkv_table.rs`](src/tables/fused_qkv_table.rs) | 240 | 1 | 0 | | [`src/tables/gemm_table.rs`](src/tables/gemm_table.rs) | 512 | 1 | 0 | | [`src/tables/gemv_table.rs`](src/tables/gemv_table.rs) | 189 | 1 | 0 | @@ -53,7 +53,7 @@ _Generated by `scripts/check-crate-maps.py` from the tree — do not edit inside | [`src/tables/rotation_table.rs`](src/tables/rotation_table.rs) | 42 | 1 | 0 | | [`src/tests.rs`](src/tests.rs) | 2,524 | 0 | 114 | | [`src/traits.rs`](src/traits.rs) | 7 | 1 | 0 | -| [`src/types.rs`](src/types.rs) | 981 | 23 | 2 | +| [`src/types.rs`](src/types.rs) | 989 | 23 | 2 | ### Public API surface @@ -103,6 +103,6 @@ _Generated by `scripts/check-crate-maps.py` from the tree — do not edit inside ### Totals -- 32 modules · 22,291 lines · 223 public items · 238 tests · 0 examples +- 32 modules · 22,571 lines · 223 public items · 242 tests · 0 examples diff --git a/crates/hipfire-dispatch/src/families/attention.rs b/crates/hipfire-dispatch/src/families/attention.rs index 51e323cd22..31b9806652 100644 --- a/crates/hipfire-dispatch/src/families/attention.rs +++ b/crates/hipfire-dispatch/src/families/attention.rs @@ -395,6 +395,21 @@ fn dispatch_kv_write( )) } } + KernelKey::KvWriteBf16 => { + debug_assert_eq!(plan.batch_size, 1); + // Two launches, K then V — same shape as the Q8 non-pair branch. + // There is no fused pair kernel for bf16: the write is pure + // convert-and-store with no amax reduction, so fusing would save a + // launch, not arithmetic. + hip!(gpu.kv_cache_write_bf16( + io.k_cache, + io.k, + io.pos_buf, + io.n_kv_heads, + io.head_dim + ))?; + hip!(gpu.kv_cache_write_bf16(io.v_cache, io.v, io.pos_buf, io.n_kv_heads, io.head_dim,)) + } KernelKey::KvWriteAsym4 => { debug_assert_eq!(plan.batch_size, 1); let ct = io.givens_cos.unwrap(); @@ -615,6 +630,32 @@ fn dispatch_kv_write( io.batch_size, )) } + KernelKey::KvWriteBf16Batched => { + // Called twice (K, then V), like the Q8 batched write. Legacy + // single-slot addressing (no slot_descs/row_slot) — maple does not + // use the multi-slot continuous-batching arena. + let pos = io.positions(); + hip!(gpu.kv_cache_write_bf16_batched( + io.k_cache, + io.k, + pos, + io.n_kv_heads, + io.head_dim, + io.batch_size, + None, + None, + ))?; + hip!(gpu.kv_cache_write_bf16_batched( + io.v_cache, + io.v, + pos, + io.n_kv_heads, + io.head_dim, + io.batch_size, + None, + None, + )) + } // ── Llama legacy (decode only, no batched variants) ── KernelKey::KvWriteHfq4 => { @@ -981,6 +1022,28 @@ fn dispatch_attend( plan.window, )) } + KernelKey::AttnFlashBf16Windowed => { + debug_assert_eq!(plan.batch_size, 1); + let seq_len = io.pos + 1; + let fp = io.flash_partials.unwrap(); + // window comes from the plan: maple's sliding layers pass + // sliding_window, its global/NoPE layers pass 0 (== plain + // causal flash). + hip!(gpu.attention_flash_bf16_windowed( + io.q, + io.k_cache, + io.v_cache, + io.output, + io.pos_buf, + seq_len, + io.n_heads, + io.n_kv_heads, + io.head_dim, + io.physical_cap, + fp, + plan.window, + )) + } KernelKey::AttnQ8_0Kv => { debug_assert_eq!(plan.batch_size, 1); let seq_len = io.pos + 1; @@ -1802,6 +1865,29 @@ fn dispatch_attend( plan.window, )) } + KernelKey::AttnBf16KvBatchedMaskedWindowed => { + // maple sliding-window prefill — same tiled shape as the Q8 + // sibling, window from the plan (0 == full causal). + let fp = io.flash_partials.unwrap(); + hip!(gpu.attention_flash_bf16_batched_masked_windowed( + io.q, + io.k_cache, + io.v_cache, + io.output, + io.positions(), + io.n_heads, + io.n_kv_heads, + io.head_dim, + io.physical_cap, + io.max_ctx_len, + io.batch_size, + fp, + io.tree_bias, + io.block_start, + io.block_cols, + plan.window, + )) + } _ => Err(DispatchError::UnsupportedVariant { family: "attention/attend", @@ -1830,6 +1916,7 @@ pub(crate) const DISPATCHED_KV_WRITE_KEYS: &[KernelKey] = &[ // Single-token KernelKey::KvWriteF32, KernelKey::KvWriteQ8_0, + KernelKey::KvWriteBf16, KernelKey::KvWriteAsym4, KernelKey::KvWriteAsym4Fwht, KernelKey::KvWriteAsym3, @@ -1844,6 +1931,7 @@ pub(crate) const DISPATCHED_KV_WRITE_KEYS: &[KernelKey] = &[ KernelKey::KvWriteAsym2Batched, KernelKey::KvWriteAsym2FwhtBatched, KernelKey::KvWriteQ8_0Batched, + KernelKey::KvWriteBf16Batched, // Llama legacy KernelKey::KvWriteHfq4, KernelKey::KvWriteQ4, @@ -1857,6 +1945,7 @@ pub(crate) const DISPATCHED_ATTEND_KEYS: &[KernelKey] = &[ KernelKey::AttnF32, KernelKey::AttnFlashQ8_0, KernelKey::AttnFlashQ8_0Windowed, + KernelKey::AttnFlashBf16Windowed, KernelKey::AttnQ8_0Kv, KernelKey::AttnFlashAsym4, KernelKey::AttnFlashAsym4Fwht, @@ -1877,6 +1966,7 @@ pub(crate) const DISPATCHED_ATTEND_KEYS: &[KernelKey] = &[ KernelKey::AttnFlashAsym2FwhtBatched, KernelKey::AttnQ8_0KvBatchedMasked, KernelKey::AttnQ8_0KvBatchedMaskedWindowed, + KernelKey::AttnBf16KvBatchedMaskedWindowed, // Llama legacy KernelKey::AttnHfq4Kv, KernelKey::AttnQ4Kv, @@ -2022,6 +2112,8 @@ mod tests { key, KvWriteF32 | KvWriteQ8_0 + | KvWriteBf16 + | KvWriteBf16Batched | KvWriteAsym4 | KvWriteAsym4Fwht | KvWriteAsym3 @@ -2054,6 +2146,7 @@ mod tests { | KvWriteAsym2Batched | KvWriteAsym2FwhtBatched | KvWriteQ8_0Batched + | KvWriteBf16Batched ) } @@ -2130,6 +2223,7 @@ mod tests { | AttnFlashAsym2FwhtBatched | AttnQ8_0KvBatchedMasked | AttnQ8_0KvBatchedMaskedWindowed + | AttnBf16KvBatchedMaskedWindowed ) } diff --git a/crates/hipfire-dispatch/src/families/kv_tier.rs b/crates/hipfire-dispatch/src/families/kv_tier.rs index 88d498739a..8c1449269b 100644 --- a/crates/hipfire-dispatch/src/families/kv_tier.rs +++ b/crates/hipfire-dispatch/src/families/kv_tier.rs @@ -30,14 +30,23 @@ pub enum F32AttnPolicy { #[derive(Clone, Copy, Debug, PartialEq, Eq)] pub enum KTier { F32, + /// Flat 2-byte BF16 K/V (maple). NOT a quantized tier — it is the + /// near-reference storage Maple's Q8 KV gets measured against. + Bf16, Q8, Hfq4, // llama legacy Q4, // llama legacy Int8c, // llama INT8-per-column Hfq8, // llama HFQ8 flat-layout - Asym4 { fwht: bool }, - Asym3 { fwht: bool }, - Asym2 { fwht: bool }, + Asym4 { + fwht: bool, + }, + Asym3 { + fwht: bool, + }, + Asym2 { + fwht: bool, + }, } /// The single bool→tier decode. The only sanctioned producer of `KTier`. @@ -52,6 +61,7 @@ pub fn classify( quant_int8: bool, quant_hfq8: bool, quant_fwht: bool, + quant_bf16: bool, ) -> KTier { debug_assert!( [ @@ -62,15 +72,24 @@ pub fn classify( quant_hfq4, quant_q4, quant_int8, - quant_hfq8 + quant_hfq8, + quant_bf16 ] .iter() .filter(|&&b| b) .count() <= 1, - "at most one KV quant tier flag should be set" + "at most one KV storage tier flag should be set" ); - if quant_asym4 { + // BF16 is checked FIRST and unconditionally. It shares no flag with any + // quantized tier, so ordering cannot change the answer for a well-formed + // cache — but if a malformed cache ever set bf16 alongside a quant flag, + // resolving to bf16 is the safe failure: the bf16 kernels read a flat + // 2-byte layout and would produce visibly wrong numbers, whereas a + // quantized kernel reading a bf16 buffer walks off the end of it. + if quant_bf16 { + KTier::Bf16 + } else if quant_asym4 { KTier::Asym4 { fwht: quant_fwht } } else if quant_asym3 { KTier::Asym3 { fwht: quant_fwht } @@ -100,7 +119,12 @@ impl KTier { KTier::Asym4 { .. } => n_kv_heads * (4 + head_dim / 2), KTier::Asym3 { .. } => n_kv_heads * (4 + (head_dim * 3) / 8), KTier::Asym2 { .. } => n_kv_heads * (4 + head_dim / 4), - KTier::F32 | KTier::Hfq4 | KTier::Q4 | KTier::Int8c | KTier::Hfq8 => { + // Bf16 has a well-defined 2 bytes/element, but it is deliberately + // NOT compactable (see `is_compactable`), and this fn's contract + // is that callers gate on that first. Panicking keeps a + // compaction path that forgot the gate loud instead of silently + // compacting a layout no gather kernel can read. + KTier::F32 | KTier::Bf16 | KTier::Hfq4 | KTier::Q4 | KTier::Int8c | KTier::Hfq8 => { panic!("k_bytes_per_pos undefined for {self:?}") } } @@ -144,6 +168,8 @@ pub struct KvTierInputs { pub quant_q4: bool, // llama legacy Q4 KV mode pub quant_int8: bool, // llama INT8-per-column KV mode pub quant_hfq8: bool, // llama HFQ8 flat-layout KV mode + /// Flat 2-byte BF16 K/V (maple). Mutually exclusive with every flag above. + pub quant_bf16: bool, /// F32-KV attention policy (Simple = attention_f32; Gqa = qwen2 selector). pub f32_policy: F32AttnPolicy, pub v_mode_bits: i32, @@ -226,6 +252,7 @@ impl KvTierPlan { quant_q4, quant_int8, quant_hfq8, + quant_bf16, f32_policy, v_mode_bits, pos, @@ -253,7 +280,18 @@ impl KvTierPlan { quant_int8, quant_hfq8, quant_fwht, + quant_bf16, ) { + // BF16 always takes the windowed kernel, exactly as + // cohere2moe's Q8 does: `window == 0` already means full + // causal, so one attend key covers both of Maple's layer + // types and there is no second path that could silently drop + // the window at ctx > window. + KTier::Bf16 => ( + KernelKey::KvWriteBf16, + KernelKey::AttnFlashBf16Windowed, + false, + ), KTier::Asym4 { fwht: true } => ( KernelKey::KvWriteAsym4Fwht, KernelKey::AttnFlashAsym4Fwht, @@ -417,6 +455,11 @@ fn batched_keys( (KvWriteQ8_0, AttnFlashQ8_0Windowed) => { Ok((KvWriteQ8_0Batched, AttnQ8_0KvBatchedMaskedWindowed)) } + // maple windowed batched (sliding-window prefill). Masked, so it + // serves tree-verify too; no is_tree gate needed. + (KvWriteBf16, AttnFlashBf16Windowed) => { + Ok((KvWriteBf16Batched, AttnBf16KvBatchedMaskedWindowed)) + } // F32 → no batched keys exist. Returning single-token keys with // batch_size > 1 will cause MissingImpl at resolve (BatchEq(1) gate). // Intentionally fall through to the default arm rather than silently @@ -446,6 +489,8 @@ fn tiers_match(write: KernelKey, attend: KernelKey) -> bool { | (KvWriteQ8_0, AttnFlashQ8_0) | (KvWriteQ8_0, AttnQ8_0Kv) | (KvWriteQ8_0, AttnFlashQ8_0Windowed) + // bf16 single-token (maple) — windowed only, by construction + | (KvWriteBf16, AttnFlashBf16Windowed) // hfq4 single-token (llama legacy) | (KvWriteHfq4, AttnHfq4Kv) // q4 single-token (llama legacy) @@ -471,6 +516,8 @@ fn tiers_match(write: KernelKey, attend: KernelKey) -> bool { // q8 batched | (KvWriteQ8_0Batched, AttnQ8_0KvBatchedMasked) | (KvWriteQ8_0Batched, AttnQ8_0KvBatchedMaskedWindowed) + // bf16 batched (maple) + | (KvWriteBf16Batched, AttnBf16KvBatchedMaskedWindowed) ) } @@ -490,6 +537,7 @@ mod tests { quant_q4: false, quant_int8: false, quant_hfq8: false, + quant_bf16: false, f32_policy: F32AttnPolicy::Simple, v_mode_bits: 8, pos: 0, @@ -647,6 +695,7 @@ mod tests { fn hfq8_tier() { let inputs = KvTierInputs { quant_hfq8: true, + quant_bf16: false, ..default_inputs() }; let plan = KvTierPlan::derive(inputs).unwrap(); @@ -698,6 +747,115 @@ mod tests { assert_eq!(batched.window, 4096); } + #[test] + fn bf16_tier_is_windowed_in_both_shapes() { + // maple: bf16 resolves to the windowed key at EVERY pos and window, + // including window == 0 (the global/NoPE layers). There is no + // non-windowed bf16 attend key by construction, so unlike Q8 there is + // no heuristic that could drop the window at ctx > window. + for (pos, window) in [(10usize, 0i32), (10, 512), (20000, 512)] { + let plan = KvTierPlan::derive(KvTierInputs { + quant_bf16: true, + window, + pos, + ..default_inputs() + }) + .unwrap(); + assert_eq!(plan.write_key, KernelKey::KvWriteBf16); + assert_eq!(plan.attend_key, KernelKey::AttnFlashBf16Windowed); + assert_eq!(plan.window, window); + // bf16 is not a rotated tier — it must never request givens buffers. + assert!(!plan.uses_givens); + } + + // Batched prefill picks the batched pair and carries the window. + let batched = KvTierPlan::derive(KvTierInputs { + quant_bf16: true, + window: 512, + batch_size: 128, + ..default_inputs() + }) + .unwrap(); + assert_eq!(batched.write_key, KernelKey::KvWriteBf16Batched); + assert_eq!( + batched.attend_key, + KernelKey::AttnBf16KvBatchedMaskedWindowed + ); + assert_eq!(batched.window, 512); + } + + #[test] + fn bf16_does_not_need_the_q8_windowed_flag() { + // NEGATIVE CONTROL for the test above. `q8_windowed` is the flag that + // makes the Q8 tier windowed; if bf16 accidentally depended on it, the + // test above would still pass (default_inputs has it false only + // because bf16 ignores it). Assert the two are genuinely independent: + // flipping q8_windowed must not change the bf16 plan at all. + let off = KvTierPlan::derive(KvTierInputs { + quant_bf16: true, + q8_windowed: false, + window: 512, + ..default_inputs() + }) + .unwrap(); + let on = KvTierPlan::derive(KvTierInputs { + quant_bf16: true, + q8_windowed: true, + window: 512, + ..default_inputs() + }) + .unwrap(); + assert_eq!(off.attend_key, on.attend_key); + assert_eq!(off.write_key, on.write_key); + assert_eq!(off.attend_key, KernelKey::AttnFlashBf16Windowed); + } + + #[test] + fn bf16_tier_classifies_and_excludes_the_quant_tiers() { + // classify() must decode the flag, and bf16 must not answer yes to any + // question asked about quantized tiers. + assert_eq!( + classify(false, false, false, false, false, false, false, false, false, true), + KTier::Bf16 + ); + // All-false is still F32, not Bf16 — the flag has to actually be read. + assert_eq!( + classify(false, false, false, false, false, false, false, false, false, false), + KTier::F32 + ); + assert!(!KTier::Bf16.is_q8()); + assert!(!KTier::Bf16.is_compactable()); + assert!(!KTier::Bf16.storage_ok_for_pflash()); + } + + #[test] + fn bf16_write_and_attend_keys_pass_the_drift_guard() { + // The #30-class guard: a bf16 write must never be paired with a + // non-bf16 attend, and vice versa. + assert!(tiers_match( + KernelKey::KvWriteBf16, + KernelKey::AttnFlashBf16Windowed + )); + assert!(tiers_match( + KernelKey::KvWriteBf16Batched, + KernelKey::AttnBf16KvBatchedMaskedWindowed + )); + // Cross-tier pairings are rejected in both directions. + assert!(!tiers_match( + KernelKey::KvWriteQ8_0, + KernelKey::AttnFlashBf16Windowed + )); + assert!(!tiers_match( + KernelKey::KvWriteBf16, + KernelKey::AttnFlashQ8_0Windowed + )); + // And a bf16 single-token write never pairs with the batched attend. + assert!(!tiers_match( + KernelKey::KvWriteBf16, + KernelKey::AttnBf16KvBatchedMaskedWindowed + )); + } + #[test] fn tiers_match_int8c_hfq8() { assert!(tiers_match(KernelKey::KvWriteInt8c, KernelKey::AttnInt8cKv)); @@ -1093,27 +1251,27 @@ mod tests { #[test] fn classify_carries_fwht_bit() { assert_eq!( - classify(false, false, true, false, false, false, false, false, true), + classify(false, false, true, false, false, false, false, false, true, false), KTier::Asym3 { fwht: true } ); assert_eq!( - classify(false, false, true, false, false, false, false, false, false), + classify(false, false, true, false, false, false, false, false, false, false), KTier::Asym3 { fwht: false } ); assert_eq!( - classify(true, false, false, false, false, false, false, false, false), + classify(true, false, false, false, false, false, false, false, false, false), KTier::Q8 ); assert_eq!( - classify(false, false, false, false, false, false, false, false, false), + classify(false, false, false, false, false, false, false, false, false, false), KTier::F32 ); assert_eq!( - classify(false, false, false, false, false, false, true, false, false), + classify(false, false, false, false, false, false, true, false, false, false), KTier::Int8c ); assert_eq!( - classify(false, false, false, false, false, false, false, true, false), + classify(false, false, false, false, false, false, false, true, false, false), KTier::Hfq8 ); } diff --git a/crates/hipfire-dispatch/src/tables/attention_table.rs b/crates/hipfire-dispatch/src/tables/attention_table.rs index f6a5f4a780..70fbd626df 100644 --- a/crates/hipfire-dispatch/src/tables/attention_table.rs +++ b/crates/hipfire-dispatch/src/tables/attention_table.rs @@ -42,6 +42,11 @@ pub fn populate(registry: &mut KernelRegistry) { ArchPredicate::Always, Some(ShapePredicate::BatchEq(1)), ), + ( + KernelKey::KvWriteBf16, + ArchPredicate::Always, + Some(ShapePredicate::BatchEq(1)), + ), ( KernelKey::KvWriteF32, ArchPredicate::Always, @@ -117,6 +122,11 @@ pub fn populate(registry: &mut KernelRegistry) { ArchPredicate::Always, Some(ShapePredicate::BatchGt(1)), ), + ( + KernelKey::KvWriteBf16Batched, + ArchPredicate::Always, + Some(ShapePredicate::BatchGt(1)), + ), ]; for (key, arch, shape) in kv_write_batched { registry.register(KernelVariant { @@ -171,6 +181,11 @@ pub fn populate(registry: &mut KernelRegistry) { ArchPredicate::Always, Some(ShapePredicate::BatchEq(1)), ), + ( + KernelKey::AttnFlashBf16Windowed, + ArchPredicate::Always, + Some(ShapePredicate::BatchEq(1)), + ), ( KernelKey::AttnQ8_0Kv, ArchPredicate::Always, @@ -319,6 +334,11 @@ pub fn populate(registry: &mut KernelRegistry) { ArchPredicate::Always, Some(ShapePredicate::BatchGt(1)), ), + ( + KernelKey::AttnBf16KvBatchedMaskedWindowed, + ArchPredicate::Always, + Some(ShapePredicate::BatchGt(1)), + ), ]; for (key, arch, shape) in attn_batched { registry.register(KernelVariant { diff --git a/crates/hipfire-dispatch/src/types.rs b/crates/hipfire-dispatch/src/types.rs index a191089a0e..f2f089a6e8 100644 --- a/crates/hipfire-dispatch/src/types.rs +++ b/crates/hipfire-dispatch/src/types.rs @@ -465,7 +465,12 @@ pub enum KernelKey { AttnFlashAsym2Fwht, AttnFlashQ8_0, AttnFlashQ8_0Windowed, // Q8_0 flash with sliding-window mask (cohere2moe) - AttnQ8_0Kv, // non-flash short-context Q8_0 decode (ship 3.1 B0) + /// BF16 flash with sliding-window mask (maple). There is deliberately no + /// non-windowed BF16 attend key: `window == 0` already means full causal, + /// so one kernel covers both of Maple's layer types and there is no + /// second path that could drop the window. + AttnFlashBf16Windowed, + AttnQ8_0Kv, // non-flash short-context Q8_0 decode (ship 3.1 B0) AttnGqaFused, // F32 GQA-flash decode family (qwen2). Selected by F32AttnPolicy::Gqa. AttnGqaWarp, // GQA, head_dim==128, long-ctx warp-reduce @@ -487,6 +492,7 @@ pub enum KernelKey { AttnFlashAsym2FwhtBatched, // no _masked — 2-bit tree-verify gap AttnQ8_0KvBatchedMasked, // P-1 no-LDS-cap tiled kernel AttnQ8_0KvBatchedMaskedWindowed, // sliding-window batched Q8 (cohere2moe prefill) + AttnBf16KvBatchedMaskedWindowed, // sliding-window batched BF16 (maple prefill) // TODO(3.3): F32-batched key for models with F32 KV + batchable weights // Full attention (no KV cache — vision / dflash cross-attention) AttnFullF16, // F16 K/V, non-causal @@ -501,6 +507,7 @@ pub enum KernelKey { KvWriteAsym2, KvWriteAsym2Fwht, KvWriteQ8_0, + KvWriteBf16, // flat 2-byte BF16 KV write (maple) KvWriteHfq4, // HFQ4-quantized KV write (llama legacy) KvWriteQ4, // Q4-quantized KV write (llama legacy) KvWriteInt8c, // INT8-per-column KV write (llama) @@ -514,6 +521,7 @@ pub enum KernelKey { KvWriteAsym2Batched, KvWriteAsym2FwhtBatched, KvWriteQ8_0Batched, + KvWriteBf16Batched, } // ── Shape context for predicate evaluation ─────────── diff --git a/crates/hipfire-generate/src/redline.rs b/crates/hipfire-generate/src/redline.rs index 415027a34d..441981a3c5 100644 --- a/crates/hipfire-generate/src/redline.rs +++ b/crates/hipfire-generate/src/redline.rs @@ -20,19 +20,19 @@ use hipfire_arch_deepseek4 as deepseek4; use hipfire_arch_lfm2moe as lfm2moe; use hipfire_arch_qwen35::carrier::Qwen35Bundle; use hipfire_arch_qwen35::dflash_verify_pm4::{ - DFLASH_VERIFY_PM4_BLOCK, DflashVerifyPm4, DflashVerifyPm4Phase, + DflashVerifyPm4, DflashVerifyPm4Phase, DFLASH_VERIFY_PM4_BLOCK, }; use hipfire_arch_qwen35::qwen35; use hipfire_arch_qwen35::speculative::{ - DeltaNetSnapshot, GdnTape, HiddenStateRingBuffer, ModelSlot, VerifyScratch, - verify_dflash_block, verify_dflash_block_retained, + verify_dflash_block, verify_dflash_block_retained, DeltaNetSnapshot, GdnTape, + HiddenStateRingBuffer, ModelSlot, VerifyScratch, }; use hipfire_engine::redline::{ - RedlineRegionHash, redline_append_buffer, redline_append_tensor, redline_append_tensor_region, - redline_capture_json, redline_hash, + redline_append_buffer, redline_append_tensor, redline_append_tensor_region, + redline_capture_json, redline_hash, RedlineRegionHash, }; -use hipfire_loader::LoadedModel; use hipfire_loader::spec_build::Qwen35SlotGuard; +use hipfire_loader::LoadedModel; use rdna_compute::replay::ReplayQuiescence; use std::any::Any; use std::io::Read; @@ -4336,7 +4336,7 @@ pub fn handle_redline_prefix_shadow( #[cfg(test)] mod redline_snapshot_tests { - use super::{RedlineQwenSnapshot, RedlineSnapshot, redline_snapshots_bit_exact}; + use super::{redline_snapshots_bit_exact, RedlineQwenSnapshot, RedlineSnapshot}; fn qwen_snapshot(gdn_frame: u32) -> RedlineSnapshot { RedlineSnapshot::Qwen(RedlineQwenSnapshot { diff --git a/crates/hipfire-loader/map.md b/crates/hipfire-loader/map.md index 8feec07c91..82fe7fd65d 100644 --- a/crates/hipfire-loader/map.md +++ b/crates/hipfire-loader/map.md @@ -25,7 +25,7 @@ _Generated by `scripts/check-crate-maps.py` from the tree — do not edit inside |---|---:|---:|---:| | [`src/batch_staging.rs`](src/batch_staging.rs) | 336 | 4 | 0 | | [`src/carriers.rs`](src/carriers.rs) | 2,850 | 11 | 5 | -| [`src/lib.rs`](src/lib.rs) | 6,605 | 134 | 33 | +| [`src/lib.rs`](src/lib.rs) | 6,614 | 134 | 33 | | [`src/parallel_capability.rs`](src/parallel_capability.rs) | 972 | 9 | 12 | | [`src/spec_build.rs`](src/spec_build.rs) | 236 | 4 | 0 | @@ -50,6 +50,6 @@ _Generated by `scripts/check-crate-maps.py` from the tree — do not edit inside ### Totals -- 5 modules · 10,999 lines · 162 public items · 50 tests · 1 examples +- 5 modules · 11,008 lines · 162 public items · 50 tests · 1 examples diff --git a/crates/hipfire-loader/src/lib.rs b/crates/hipfire-loader/src/lib.rs index 2235102919..be9095702d 100644 --- a/crates/hipfire-loader/src/lib.rs +++ b/crates/hipfire-loader/src/lib.rs @@ -3147,6 +3147,9 @@ pub fn load_model( None, hipfire_config::Deepseek4ComputePlacement::Single, draft_path, + // No head overlay on this legacy entry point; callers that want one + // use load_model_with_kv_backend / _with_gemma4_drafter directly. + None, kv_mode_override, None, kv_adaptive_override, @@ -3166,6 +3169,7 @@ pub fn load_model_with_kv_backend( deepseek4_experts_per_token: Option, deepseek4_compute_placement: hipfire_config::Deepseek4ComputePlacement, draft_path: Option<&str>, + head_path: Option<&str>, kv_mode_override: Option<&str>, kv_backend_override: Option<&str>, kv_adaptive_override: Option<&str>, @@ -3182,6 +3186,7 @@ pub fn load_model_with_kv_backend( deepseek4_experts_per_token, deepseek4_compute_placement, draft_path, + head_path, kv_mode_override, kv_backend_override, kv_adaptive_override, @@ -3200,6 +3205,7 @@ pub fn load_model_with_kv_backend_admitted( deepseek4_experts_per_token: Option, deepseek4_compute_placement: hipfire_config::Deepseek4ComputePlacement, draft_path: Option<&str>, + head_path: Option<&str>, kv_mode_override: Option<&str>, kv_backend_override: Option<&str>, kv_adaptive_override: Option<&str>, @@ -3310,6 +3316,7 @@ pub fn load_model_with_kv_backend_admitted( deepseek4_compute_placement, deepseek4_experts_per_token, draft_path, + head_path, kv_mode_override, kv_backend, kv_adaptive_override, @@ -3376,6 +3383,7 @@ pub fn load_model_with_gemma4_drafter( deepseek4_experts_per_token: Option, deepseek4_compute_placement: hipfire_config::Deepseek4ComputePlacement, draft_path: Option<&str>, + head_path: Option<&str>, gemma4_drafter_path: Option<&str>, gemma4_draft_len: usize, kv_mode_override: Option<&str>, @@ -3394,6 +3402,7 @@ pub fn load_model_with_gemma4_drafter( deepseek4_experts_per_token, deepseek4_compute_placement, draft_path, + head_path, gemma4_drafter_path, gemma4_draft_len, kv_mode_override, @@ -3414,6 +3423,7 @@ pub fn load_model_with_gemma4_drafter_admitted( deepseek4_experts_per_token: Option, deepseek4_compute_placement: hipfire_config::Deepseek4ComputePlacement, draft_path: Option<&str>, + head_path: Option<&str>, gemma4_drafter_path: Option<&str>, gemma4_draft_len: usize, kv_mode_override: Option<&str>, @@ -3476,6 +3486,7 @@ pub fn load_model_with_gemma4_drafter_admitted( deepseek4_compute_placement, deepseek4_experts_per_token, draft_path, + head_path, kv_mode_override, kv_backend, kv_adaptive_override, diff --git a/crates/hipfire-quantize/map.md b/crates/hipfire-quantize/map.md index 1a4715a144..19f4737894 100644 --- a/crates/hipfire-quantize/map.md +++ b/crates/hipfire-quantize/map.md @@ -28,7 +28,7 @@ _Generated by `scripts/check-crate-maps.py` from the tree — do not edit inside | [`src/bin/mq4_merge_mtp.rs`](src/bin/mq4_merge_mtp.rs) | 138 | 2 | 0 | | [`src/bin/mtp_extract.rs`](src/bin/mtp_extract.rs) | 1,345 | 0 | 0 | | [`src/calibration.rs`](src/calibration.rs) | 1,351 | 0 | 8 | -| [`src/cli.rs`](src/cli.rs) | 225 | 0 | 0 | +| [`src/cli.rs`](src/cli.rs) | 249 | 0 | 0 | | [`src/dequant.rs`](src/dequant.rs) | 357 | 0 | 0 | | [`src/diagnostics.rs`](src/diagnostics.rs) | 2,762 | 0 | 51 | | [`src/e8.rs`](src/e8.rs) | 1,210 | 12 | 15 | @@ -42,17 +42,17 @@ _Generated by `scripts/check-crate-maps.py` from the tree — do not edit inside | [`src/hfqm.rs`](src/hfqm.rs) | 676 | 24 | 3 | | [`src/lib.rs`](src/lib.rs) | 22 | 8 | 0 | | [`src/main.rs`](src/main.rs) | 38 | 0 | 0 | -| [`src/maple.rs`](src/maple.rs) | 435 | 0 | 8 | +| [`src/maple.rs`](src/maple.rs) | 491 | 0 | 8 | | [`src/model_filter.rs`](src/model_filter.rs) | 816 | 0 | 12 | -| [`src/pipeline.rs`](src/pipeline.rs) | 7,439 | 0 | 12 | +| [`src/pipeline.rs`](src/pipeline.rs) | 7,446 | 0 | 12 | | [`src/pipeline_deepseek.rs`](src/pipeline_deepseek.rs) | 306 | 0 | 0 | | [`src/pipeline_gguf.rs`](src/pipeline_gguf.rs) | 1,033 | 0 | 2 | -| [`src/pipeline_maple.rs`](src/pipeline_maple.rs) | 1,036 | 0 | 18 | +| [`src/pipeline_maple.rs`](src/pipeline_maple.rs) | 1,109 | 0 | 18 | | [`src/quant_e8.rs`](src/quant_e8.rs) | 1,422 | 2 | 13 | | [`src/quant_fwht.rs`](src/quant_fwht.rs) | 823 | 0 | 4 | | [`src/quant_hfp4.rs`](src/quant_hfp4.rs) | 422 | 0 | 0 | | [`src/quant_mq.rs`](src/quant_mq.rs) | 3,497 | 0 | 28 | -| [`src/quant_q4.rs`](src/quant_q4.rs) | 293 | 0 | 0 | +| [`src/quant_q4.rs`](src/quant_q4.rs) | 406 | 0 | 0 | | [`src/reap_overlay.rs`](src/reap_overlay.rs) | 1,143 | 12 | 29 | | [`src/safetensors_file.rs`](src/safetensors_file.rs) | 109 | 6 | 1 | @@ -105,6 +105,6 @@ _Generated by `scripts/check-crate-maps.py` from the tree — do not edit inside ### Totals -- 33 modules · 36,312 lines · 141 public items · 261 tests · 7 examples +- 33 modules · 36,585 lines · 141 public items · 261 tests · 7 examples diff --git a/crates/hipfire-quantize/src/cli.rs b/crates/hipfire-quantize/src/cli.rs index 065c46abc7..8cffe219b4 100644 --- a/crates/hipfire-quantize/src/cli.rs +++ b/crates/hipfire-quantize/src/cli.rs @@ -62,10 +62,34 @@ pub(crate) struct QuantizeArgs { /// Does NOT touch `word_embeddings` (same shape, but only ONE ROW is read /// per token — a RAM question, not a bandwidth one), the ternary expert /// path, or the router. - #[arg(long, value_name = "MODE", default_value = "bf16", - value_parser = ["bf16", "q8", "mq4"])] + /// Default is q8, not bf16. Measured on gfx1151 against a bf16 reference + /// (2048 teacher-forced tokens): q8 and bf16 heads give the IDENTICAL mean + /// KL of 0.0511, but q8 decodes 23% faster (144.6 vs 117.6 tok/s). A bf16 + /// head is therefore strictly dominated -- it costs throughput and buys + /// exactly zero accuracy. mq4 (qt=30 Lloyd, 5.0 bpw) is +10.4% decode over + /// q8 but +51% mean KL and -2.7pp top-1, which is a poor trade on this + /// stack; note the vendor DOES ship a Q4_K head, because on their CPU path + /// the same swap buys 49% rather than 10%. + /// + /// `mq4` (qt=30) is DEPRECATED and no longer selectable: mq4v2 (qt=44) + /// beats it on every axis -- lower KL (0.0744 vs 0.0772), faster (165.8 vs + /// 161.8 tok/s) and 15% smaller (4.25 vs 5.0 bpw). Existing .hfq files with + /// a qt=30 head still LOAD; only producing new ones is removed. + #[arg(long, value_name = "MODE", default_value = "q8", + value_parser = ["bf16", "q8", "mq4v2", "q4k"])] pub head_quant: String, + /// `--format maple` only: emit a HEAD-ONLY `.hfq` containing just + /// `lm_head.weight` at `--head-quant`, instead of a full model. + /// + /// The result is a load-time overlay for a full build: same arch_id, same + /// logical shape, differing only in the head's quant tier. Shipping heads + /// this way avoids duplicating the identical 6.17 GB body per carrier — + /// three head variants cost 7.30 GB rather than 19.63 GB, and switching + /// heads is a 175 MB download instead of 6.5 GB. + #[arg(long, default_value_t = false)] + pub head_only: bool, + /// Override the architecture ID stamped into the HFQ header. #[arg(long, value_name = "ID")] pub arch_id: Option, diff --git a/crates/hipfire-quantize/src/maple.rs b/crates/hipfire-quantize/src/maple.rs index 241679eef7..babc4bc52b 100644 --- a/crates/hipfire-quantize/src/maple.rs +++ b/crates/hipfire-quantize/src/maple.rs @@ -56,7 +56,29 @@ pub(crate) enum MapleHeadQuant { /// **FWHT-rotated**: the weights are encoded against FWHT-256-rotated /// blocks, so the runtime MUST rotate `x` to match. See /// `pack_maple_head` for why the seeds are not free parameters. + /// + /// **DEPRECATED — use `Mq4V2`.** Measured on gfx1151 (KV bf16, 2048 + /// teacher-forced tokens): qt=44 is better on EVERY axis — mean KL 0.0744 + /// vs 0.0772, decode 165.8 vs 161.8 tok/s, and 4.25 vs 5.0 bpw. There is + /// no workload where qt=30 is the right choice. Kept only so the packer + /// arm and its FWHT-seed contract stay documented next to qt=44's; the + /// CLI no longer offers it. Mq4, + /// MQ4-G256 **v2** (qt=44), 136 B per 256 weights = 4.25 bpw. + /// **FWHT-rotated**, same as `Mq4`, and the same nibble payload — but the + /// 8 header bytes carry a SEPARATE fp16 scale/zero per 128-weight half + /// instead of one pair governing all 256. Strictly finer quantization at a + /// SMALLER footprint than qt=30 (4.25 vs 5.0 bpw), so it is the natural + /// candidate if the mq4 head's accuracy cost is what rules it out. + Mq4V2, + /// GGML-compatible **Q4_K** (qt=4), 144 B per 256 weights = 4.5 bpw. + /// Unrotated, and the finest-grained 4-bit carrier available: a separate + /// scale AND min per 32-weight sub-block (8 per 256), with the sub-block + /// meta itself 6-bit quantized. Measured on the real lm_head, relative L2 + /// error by granularity: 0.118 at 1 scale/256 (qt=30 class), 0.106 at 2 + /// (qt=44), 0.080 at 8 (this). This is the exact carrier DeepGrove ship in + /// their own llama.cpp example, so it is the like-for-like comparison. + Q4K, } impl std::str::FromStr for MapleHeadQuant { @@ -66,8 +88,10 @@ impl std::str::FromStr for MapleHeadQuant { "bf16" | "none" => Ok(Self::Bf16), "q8" | "q8_0" | "q8f16" => Ok(Self::Q8), "mq4" | "mq4-lloyd" | "mq4g256lloyd" => Ok(Self::Mq4), + "mq4v2" | "mq4-v2" | "mq4g256v2" => Ok(Self::Mq4V2), + "q4k" | "q4_k" | "q4km" => Ok(Self::Q4K), other => Err(format!( - "unknown --head-quant {other:?} (expected bf16, q8 or mq4)" + "unknown --head-quant {other:?} (expected bf16, q8, mq4v2 or q4k)" )), } } @@ -80,6 +104,8 @@ impl MapleHeadQuant { Self::Bf16 => "bf16", Self::Q8 => "q8", Self::Mq4 => "mq4", + Self::Mq4V2 => "mq4v2", + Self::Q4K => "q4k", } } } @@ -165,6 +191,36 @@ pub(crate) fn pack_maple_head( 256, )) } + MapleHeadQuant::Mq4V2 => { + if k % 256 != 0 { + return Err(format!( + "lm_head K={k} is not a multiple of 256 (MQ4-G256 block)" + )); + } + // Same FWHT seeds as the qt=30 arm above. They are NOT free + // parameters: the runtime rotates `x` with signs derived from the + // same seeds, so a mismatch here silently produces garbage logits + // rather than a load error. + let signs1 = crate::quant_fwht::gen_fwht_signs(42, 256); + let signs2 = crate::quant_fwht::gen_fwht_signs(1042, 256); + let m = vals.len() / k; + Ok(( + crate::quant_fwht::quantize_mq4g256v2(vals, m, k, &signs1, &signs2), + QuantType::MQ4G256V2, + 256, + )) + } + MapleHeadQuant::Q4K => { + if k % 256 != 0 { + return Err(format!( + "lm_head K={k} is not a multiple of 256 (Q4_K super-block)" + )); + } + // UNROTATED, unlike qt=30/44: Q4_K carries its own per-32 scales + // and has no FWHT convention, so there are no seeds to keep in + // sync with `ensure_mq_signs` here. + Ok((crate::quant_q4::quantize_q4k(vals), QuantType::Q4K, 256)) + } } } diff --git a/crates/hipfire-quantize/src/pipeline.rs b/crates/hipfire-quantize/src/pipeline.rs index 950ec402c2..e9c281896c 100644 --- a/crates/hipfire-quantize/src/pipeline.rs +++ b/crates/hipfire-quantize/src/pipeline.rs @@ -2988,12 +2988,19 @@ fn handle_early_special_formats(args: &QuantizeArgs) -> bool { eprintln!("error: {e}"); std::process::exit(2); }); - match crate::pipeline_maple::convert_maple_safetensors( + let convert = if args.head_only { + crate::pipeline_maple::convert_maple_head_only + } else { + crate::pipeline_maple::convert_maple_safetensors + }; + match convert( Path::new(input_dir), Path::new(output_path), &config_json, head_quant, ) { + // The head-only path prints its own line; the full path does not. + Ok(_) if args.head_only => {} Ok(_) => eprintln!("maple: wrote {output_path}"), Err(e) => { eprintln!("error: {e}"); diff --git a/crates/hipfire-quantize/src/pipeline_maple.rs b/crates/hipfire-quantize/src/pipeline_maple.rs index 1716298f56..31a2592b5e 100644 --- a/crates/hipfire-quantize/src/pipeline_maple.rs +++ b/crates/hipfire-quantize/src/pipeline_maple.rs @@ -209,6 +209,79 @@ fn shard_paths(dir: &Path) -> Result, String> { /// /// Shards are processed one at a time and their page cache dropped as they are /// consumed, so peak RSS stays bounded even though the source is ~40 GB. +/// Emit a HEAD-ONLY `.hfq` containing just `lm_head.weight`, for use as a +/// load-time overlay over a full Maple build. +/// +/// WHY THIS EXISTS. The head is the only tensor whose carrier we vary, and it +/// is 2.7% of the file (175-622 MB of ~6.5 GB). Shipping a whole 6.3-6.8 GB +/// artifact per head carrier duplicates the identical 6.17 GB body every time; +/// three variants cost 19.63 GB instead of 7.30 GB, and switching heads costs +/// a user a full re-download instead of 175 MB. +/// +/// The output is deliberately a NORMAL `.hfq` with the same `arch_id`, one +/// tensor, and the SAME logical shape as the base's head — because +/// `HfqFile::attach_overlay` requires exactly that. It rejects an overlay whose +/// tensor is absent from the base or differs in shape, which is what stops an +/// overlay built for another model from being spliced in silently. Emitting a +/// head-only file this way means that guard keeps working unchanged; nothing +/// about the overlay mechanism is relaxed to support this. +/// +/// Every carrier `--head-quant` accepts works here, bf16 included: for bf16 +/// `convert_tensor` emits a `QuantType::BF16` passthrough and never reaches +/// `pack_maple_head` (which has no Bf16 arm), so no special case is needed. +/// +/// A headless BODY is deliberately NOT offered. It would require permitting an +/// overlay to introduce names the base lacks, which is the very check that +/// catches a wrong-model overlay, and it would ship an artifact that cannot +/// run on its own. +pub(crate) fn convert_maple_head_only( + input_dir: &Path, + output: &Path, + config_json: &str, + head_quant: MapleHeadQuant, +) -> Result { + let shards = shard_paths(input_dir)?; + let mut stats = MapleConvertStats { + head_quant, + ..Default::default() + }; + for shard in &shards { + let sf = + SafetensorsFile::open(shard).map_err(|e| format!("open {}: {e}", shard.display()))?; + if !sf.tensor_names().iter().any(|n| *n == LM_HEAD_NAME) { + continue; + } + let (meta, bytes) = sf + .tensor_data(LM_HEAD_NAME) + .ok_or_else(|| format!("{LM_HEAD_NAME}: vanished from {}", shard.display()))?; + let (t, _) = convert_tensor( + LM_HEAD_NAME, + &meta.dtype, + &meta.shape, + bytes, + head_quant, + &mut stats, + )?; + eprintln!( + "maple: head-only overlay — {} {:?} → {} ({:.1} MB)", + LM_HEAD_NAME, + meta.shape, + head_quant.label(), + t.data.len() as f64 / 1e6, + ); + let metadata = build_metadata(input_dir, config_json, &stats)?; + // No spill: one tensor, and it is at most 622 MB. + write_hfq(output, ARCH_ID_MAPLE, &metadata, &[t], None) + .map_err(|e| format!("write {}: {e}", output.display()))?; + eprintln!("maple: wrote {}", output.display()); + return Ok(stats); + } + Err(format!( + "{LM_HEAD_NAME} not found in any shard under {} — cannot build a head overlay", + input_dir.display() + )) +} + pub(crate) fn convert_maple_safetensors( input_dir: &Path, output: &Path, diff --git a/crates/hipfire-quantize/src/quant_q4.rs b/crates/hipfire-quantize/src/quant_q4.rs index 23ca47fda6..9246882be9 100644 --- a/crates/hipfire-quantize/src/quant_q4.rs +++ b/crates/hipfire-quantize/src/quant_q4.rs @@ -3,24 +3,29 @@ // Copyright (c) 2026 Nick Woolmer // hipfire — see LICENSE and NOTICE in the project root. - -#![allow(dead_code, unused_imports, unused_variables, non_snake_case, clippy::all)] +#![allow( + dead_code, + unused_imports, + unused_variables, + non_snake_case, + clippy::all +)] use std::collections::HashMap; -use std::path::{Path, PathBuf}; use std::fs::File; use std::io::Write; -use std::sync::OnceLock; +use std::path::{Path, PathBuf}; use std::sync::atomic::{AtomicU64, Ordering}; +use std::sync::OnceLock; -use clap::Parser; -use hipfire_quantize::float16::{bf16_to_f32, f16_to_f32, f32_to_f16}; -use hipfire_quantize::safetensors_file::{SafetensorsFile, TensorMeta}; -use hipfire_quantize::hessian_io; use crate::e8; use crate::e8_gptq; use crate::gguf_input; use crate::reap_overlay; +use clap::Parser; +use hipfire_quantize::float16::{bf16_to_f32, f16_to_f32, f32_to_f16}; +use hipfire_quantize::hessian_io; +use hipfire_quantize::safetensors_file::{SafetensorsFile, TensorMeta}; // ─── Q4_F16_G64 Quantization ──────────────────────────────────────────────── @@ -74,6 +79,108 @@ pub(crate) fn quantize_q4f16_g64(f32_data: &[f32]) -> Vec { /// Quantize F32 weights to Q4_K format (144 bytes per 256 elements, 0.5625 B/w). /// GGML-compatible block layout: f16 d + f16 dmin + 12B packed scales + 128B nibbles. /// This produces blocks that work with the existing gemv_q4k kernel. + +/// Port of llama.cpp's `make_qkx2_quants` (ggml-quants.c:799). +/// +/// Returns `(scale, the_min)` for one sub-block, where dequantization is +/// `w = scale * q - the_min` — the same convention `dequantize_row_q4_K` uses +/// (`y = d1*q - m1`). +/// +/// WHY THIS RATHER THAN MIN/MAX. Plain min/max picks the scale that makes the +/// extremes representable, which is not the scale that minimises error: one +/// outlier stretches the grid and every other weight pays for it. This searches +/// `nstep` candidate scales around the min/max one and, for each, solves the +/// weighted least-squares fit for (scale, min) given the resulting integer +/// levels, keeping whichever candidate actually has the lowest error. +/// +/// Measured on Maple's lm_head: min/max gives relative L2 0.0799, this gives +/// 0.0731 — the same 0.0731 DeepGrove's published Q4_K head achieves. The +/// layout was already GGML-compatible; only the encoder was weaker. +/// +/// `weights` are llama.cpp's importance weights `sqrt(mean(x^2)) + |x|`, which +/// bias the fit toward larger-magnitude entries. +#[allow(clippy::too_many_arguments)] +fn make_qkx2_quants( + x: &[f32], + weights: &[f32], + nmax: i32, + rmin: f32, + rdelta: f32, + nstep: i32, +) -> (f32, f32) { + let n = x.len(); + let mut min = x[0]; + let mut max = x[0]; + let mut sum_w = weights[0]; + let mut sum_x = sum_w * x[0]; + for i in 1..n { + if x[i] < min { + min = x[i]; + } + if x[i] > max { + max = x[i]; + } + let w = weights[i]; + sum_w += w; + sum_x += w * x[i]; + } + // The grid is anchored at or below zero, so an all-positive block still + // encodes zero exactly. + if min > 0.0 { + min = 0.0; + } + if max == min { + return (0.0, -min); + } + + let mut iscale = nmax as f32 / (max - min); + let mut scale = 1.0 / iscale; + let mut laux = vec![0i32; n]; + let mut best_error = 0.0f32; + for i in 0..n { + let l = (iscale * (x[i] - min)).round() as i32; + let l = l.clamp(0, nmax); + let diff = scale * l as f32 + min - x[i]; + best_error += weights[i] * diff * diff; + } + if nstep < 1 { + return (scale, -min); + } + + for is in 0..=nstep { + iscale = (rmin + rdelta * is as f32 + nmax as f32) / (max - min); + let (mut sum_l, mut sum_l2, mut sum_xl) = (0.0f32, 0.0f32, 0.0f32); + for i in 0..n { + let l = ((iscale * (x[i] - min)).round() as i32).clamp(0, nmax); + laux[i] = l; + let w = weights[i]; + sum_l += w * l as f32; + sum_l2 += w * (l * l) as f32; + sum_xl += w * l as f32 * x[i]; + } + let d = sum_w * sum_l2 - sum_l * sum_l; + if d > 0.0 { + let mut this_scale = (sum_w * sum_xl - sum_x * sum_l) / d; + let mut this_min = (sum_l2 * sum_x - sum_l * sum_xl) / d; + if this_min > 0.0 { + this_min = 0.0; + this_scale = sum_xl / sum_l2; + } + let mut cur_error = 0.0f32; + for i in 0..n { + let diff = this_scale * laux[i] as f32 + this_min - x[i]; + cur_error += weights[i] * diff * diff; + } + if cur_error < best_error { + best_error = cur_error; + scale = this_scale; + min = this_min; + } + } + } + (scale, -min) +} + pub(crate) fn quantize_q4k(f32_data: &[f32]) -> Vec { let super_block_size = 256; let block_bytes = 144; @@ -98,11 +205,17 @@ pub(crate) fn quantize_q4k(f32_data: &[f32]) -> Vec { } let group = &f32_data[start..end]; - let min_val = group.iter().cloned().fold(f32::INFINITY, f32::min); - let max_val = group.iter().cloned().fold(f32::NEG_INFINITY, f32::max); - let range = max_val - min_val; - sub_scales[sb] = if range > 0.0 { range / 15.0 } else { 0.0 }; - sub_mins[sb] = min_val; + // llama.cpp's importance weights: sqrt(mean(x^2)) + |x|. + let sum_x2: f32 = group.iter().map(|v| v * v).sum(); + let av_x = (sum_x2 / group.len() as f32).sqrt(); + let w: Vec = group.iter().map(|v| av_x + v.abs()).collect(); + // Same parameters Q4_K uses at ggml-quants.c:1476 + // (nmax=15, rmin=-1.0, rdelta=0.1, nstep=20, use_mad=false). + let (scale, the_min) = make_qkx2_quants(group, &w, 15, -1.0, 0.1, 20); + sub_scales[sb] = scale; + // The rest of this function stores the SIGNED min and negates it + // when packing, so convert back from llama.cpp's positive the_min. + sub_mins[sb] = -the_min; } // Find super-block d and dmin that best represent the sub-block scales/mins diff --git a/crates/hipfire-registry/map.md b/crates/hipfire-registry/map.md index c101c3dfc2..3409bb0d0f 100644 --- a/crates/hipfire-registry/map.md +++ b/crates/hipfire-registry/map.md @@ -22,7 +22,7 @@ _Generated by `scripts/check-crate-maps.py` from the tree — do not edit inside | File | Lines | Public items | Tests | |---|---:|---:|---:| -| [`src/lib.rs`](src/lib.rs) | 1,796 | 25 | 18 | +| [`src/lib.rs`](src/lib.rs) | 2,022 | 25 | 22 | ### Public API surface @@ -41,6 +41,6 @@ _Generated by `scripts/check-crate-maps.py` from the tree — do not edit inside ### Totals -- 1 modules · 1,796 lines · 25 public items · 18 tests · 0 examples +- 1 modules · 2,022 lines · 25 public items · 22 tests · 0 examples diff --git a/crates/hipfire-registry/src/lib.rs b/crates/hipfire-registry/src/lib.rs index bed6942495..3e4e580cb2 100644 --- a/crates/hipfire-registry/src/lib.rs +++ b/crates/hipfire-registry/src/lib.rs @@ -187,6 +187,16 @@ pub struct ModelEntry { pub mtp: Option, #[serde(default)] pub dspark: Option, + /// Alternative `lm_head` carriers, keyed by short name (`q4k`, `bf16`). + /// + /// Each is a single-tensor `.hfq` from `hipfire-quantize --head-only` that + /// shadows the base's `lm_head.weight` at load time. The BASE ships the + /// recommended head and runs standalone; these only exist so a different + /// carrier costs a 188-635 MB download instead of a near-identical 6.5 GB + /// model. Three full variants would be 19.63 GB; base plus two overlays is + /// 7.30 GB. + #[serde(default)] + pub heads: std::collections::BTreeMap, #[serde(default)] pub default_tool_format: Option, #[serde(default)] @@ -346,9 +356,7 @@ impl RegistryV1 { self.schema_version ))); } - if self.generated_at.trim().is_empty() { - return Err(fail("generated_at is empty".into())); - } + validate_generated_at(&self.generated_at).map_err(fail)?; if self.models.is_empty() { return Err(fail("model catalog is empty".into())); } @@ -367,6 +375,7 @@ impl RegistryV1 { for sidecar in [&entry.triattn, &entry.mtp, &entry.dspark] .into_iter() .flatten() + .chain(entry.heads.values()) { if sidecar.file.trim().is_empty() { return Err(fail(format!("model '{tag}' has an empty sidecar file"))); @@ -447,6 +456,86 @@ fn validate_digest(digest: Option<&str>, label: &str) -> std::result::Result<(), Ok(()) } +fn validate_generated_at(ts: &str) -> std::result::Result<(), String> { + // Strict normalized RFC3339 UTC: YYYY-MM-DDTHH:MM:SSZ (20 bytes). + // Lexical order equals chronological order only in this normalized form, + // which is what `prefer_bundled_if_newer` relies on. Reject any + // non-normalized representation (offsets, fractional seconds, whitespace, + // lowercase, etc.) and validate calendar ranges so malformed strings + // cannot invert precedence via lexical comparison. + if ts.len() != 20 { + return Err(format!( + "generated_at '{}' is not strict RFC3339 UTC (expected YYYY-MM-DDTHH:MM:SSZ)", + ts + )); + } + let b = ts.as_bytes(); + if b[4] != b'-' + || b[7] != b'-' + || b[10] != b'T' + || b[13] != b':' + || b[16] != b':' + || b[19] != b'Z' + { + return Err(format!( + "generated_at '{}' is not strict RFC3339 UTC (expected YYYY-MM-DDTHH:MM:SSZ)", + ts + )); + } + for i in [0, 1, 2, 3, 5, 6, 8, 9, 11, 12, 14, 15, 17, 18] { + if !b[i].is_ascii_digit() { + return Err(format!( + "generated_at '{}' is not strict RFC3339 UTC (expected YYYY-MM-DDTHH:MM:SSZ)", + ts + )); + } + } + let year = (b[0] - b'0') as u16 * 1000 + + (b[1] - b'0') as u16 * 100 + + (b[2] - b'0') as u16 * 10 + + (b[3] - b'0') as u16; + let month = (b[5] - b'0') * 10 + (b[6] - b'0'); + let day = (b[8] - b'0') * 10 + (b[9] - b'0'); + let hour = (b[11] - b'0') * 10 + (b[12] - b'0'); + let minute = (b[14] - b'0') * 10 + (b[15] - b'0'); + let second = (b[17] - b'0') * 10 + (b[18] - b'0'); + if month == 0 || month > 12 { + return Err(format!("generated_at '{}' has invalid month", ts)); + } + if day == 0 || day > days_in_month(year, month) { + return Err(format!("generated_at '{}' has invalid day", ts)); + } + if hour > 23 { + return Err(format!("generated_at '{}' has invalid hour", ts)); + } + if minute > 59 { + return Err(format!("generated_at '{}' has invalid minute", ts)); + } + if second > 59 { + return Err(format!("generated_at '{}' has invalid second", ts)); + } + Ok(()) +} + +fn days_in_month(year: u16, month: u8) -> u8 { + match month { + 1 | 3 | 5 | 7 | 8 | 10 | 12 => 31, + 4 | 6 | 9 | 11 => 30, + 2 => { + if is_leap_year(year) { + 29 + } else { + 28 + } + } + _ => 0, + } +} + +fn is_leap_year(year: u16) -> bool { + year % 4 == 0 && (year % 100 != 0 || year % 400 == 0) +} + fn is_effort_native_tag(tag: &str) -> bool { // Mirrors scripts/registry_gen.py:_effort_native_tag and the family/tag // conventions already used by config_layer_for_tag. Effort-native families @@ -564,6 +653,43 @@ pub fn bundled() -> Result { RegistryV1::parse(BUNDLED_REGISTRY, "bundled registry/v1.json") } +/// Prefer the BUNDLED registry when it is newer than whatever was fetched. +/// +/// `registry/v1.json` is compiled into the binary, so on a branch the bundled +/// copy IS that branch's registry — while the fetch targets master. Without +/// this, editing the registry on a branch changes nothing for a locally built +/// binary: the 24h cache or a master fetch silently wins, and nothing says so. +/// That cost a real debugging detour (a branch's `heads` map read as empty). +/// +/// `generated_at` is the existing signal and needs no new configuration: +/// `scripts/registry_gen.py` stamps it on every regeneration, and its +/// `%Y-%m-%dT%H:%M:%SZ` form compares correctly as a plain string. +/// +/// This does NOT break distribution. A released binary's bundled registry is +/// older than master's by construction, so the fetch keeps winning there and +/// users still get new models without upgrading. Only a freshly regenerated +/// local registry — i.e. someone editing it — takes precedence. +fn prefer_bundled_if_newer( + loaded: LoadedRegistry, + bundled: RegistryV1, + warnings: &mut Vec, +) -> LoadedRegistry { + if loaded.source == RegistrySource::Bundled + || bundled.generated_at <= loaded.registry.generated_at + { + return loaded; + } + warnings.push(format!( + "using the bundled registry ({}), which is newer than the fetched one ({})", + bundled.generated_at, loaded.registry.generated_at + )); + LoadedRegistry { + registry: bundled, + source: RegistrySource::Bundled, + warnings: std::mem::take(warnings), + } +} + pub fn load(paths: &RegistryPaths) -> LoadedRegistry { let mut warnings = Vec::new(); let bundled = bundled().expect("checked-in registry/v1.json must validate"); @@ -582,14 +708,15 @@ pub fn load(paths: &RegistryPaths) -> LoadedRegistry { .as_ref() .is_some_and(|cache| cache_is_fresh(cache, now, REGISTRY_CACHE_TTL)) { - return LoadedRegistry { + let loaded = LoadedRegistry { registry: cache.expect("checked above").registry, source: RegistrySource::Cache, - warnings, + warnings: std::mem::take(&mut warnings), }; + return prefer_bundled_if_newer(loaded, bundled, &mut warnings); } - match fetch_registry(&url) { + let loaded = match fetch_registry(&url) { Ok(registry) => { let cache_file = RegistryCache { fetched_at: now, @@ -602,7 +729,7 @@ pub fn load(paths: &RegistryPaths) -> LoadedRegistry { LoadedRegistry { registry, source: RegistrySource::Network, - warnings, + warnings: std::mem::take(&mut warnings), } } Err(error) => { @@ -611,17 +738,18 @@ pub fn load(paths: &RegistryPaths) -> LoadedRegistry { LoadedRegistry { registry: cache.registry, source: RegistrySource::StaleCache, - warnings, + warnings: std::mem::take(&mut warnings), } } else { LoadedRegistry { - registry: bundled, + registry: bundled.clone(), source: RegistrySource::Bundled, - warnings, + warnings: std::mem::take(&mut warnings), } } } - } + }; + prefer_bundled_if_newer(loaded, bundled, &mut warnings) } fn read_cache(path: &Path, url: &str, warnings: &mut Vec) -> Option { @@ -714,6 +842,104 @@ fn epoch_millis() -> u64 { mod tests { use super::*; + /// The `heads` map must be VALIDATED, not merely parsed. + /// + /// Adding a field to the struct makes it round-trip; it does not make the + /// validator look at it. This asserts the negative directly: a head with a + /// malformed digest is REJECTED. Without the `.chain(entry.heads.values())` + /// in `validate`, this test fails and the bundled-registry check would + /// happily ship an unverifiable head overlay. + #[test] + fn heads_sidecars_are_digest_validated() { + let with_bad_head = r#"{ + "schema_version":1, + "generated_at":"2026-09-01T00:00:00Z", + "models":{"m":{"repo":"r","file":"f.hfq","size_gb":1,"min_vram_gb":1,"desc":"d", + "heads":{"q4k":{"file":"h.hfq","sha256":"not-a-sha"}}}}, + "aliases":{} + }"#; + let err = RegistryV1::parse(with_bad_head, "test") + .expect_err("a malformed head digest must be rejected"); + assert!( + format!("{err}").contains("invalid SHA-256"), + "expected a digest complaint, got: {err}" + ); + + // Control: the SAME registry with a well-formed digest parses, so the + // rejection above is about the digest and not about `heads` being + // unparseable. + let good = with_bad_head.replace("not-a-sha", &"a".repeat(64)); + let reg = RegistryV1::parse(&good, "test").expect("valid head must parse"); + let (_, entry) = reg.model("m").unwrap(); + assert_eq!(entry.heads.len(), 1); + assert_eq!(entry.heads["q4k"].file, "h.hfq"); + } + + fn reg_at(stamp: &str) -> RegistryV1 { + RegistryV1::parse( + &format!( + r#"{{"schema_version":1,"generated_at":"{stamp}", + "models":{{"m":{{"repo":"r","file":"f","size_gb":1,"min_vram_gb":1,"desc":"d"}}}}, + "aliases":{{}}}}"# + ), + "test", + ) + .unwrap() + } + + /// A NEWER bundled registry wins — this is what makes a branch's registry + /// edits visible to a locally built binary instead of being silently + /// overridden by the 24h cache or a master fetch. + #[test] + fn newer_bundled_registry_beats_a_stale_fetch() { + let mut w = Vec::new(); + let fetched = LoadedRegistry { + registry: reg_at("2026-08-31T05:32:38Z"), + source: RegistrySource::Cache, + warnings: Vec::new(), + }; + let out = prefer_bundled_if_newer(fetched, reg_at("2026-09-01T13:14:39Z"), &mut w); + assert_eq!(out.source, RegistrySource::Bundled); + assert_eq!(out.registry.generated_at, "2026-09-01T13:14:39Z"); + assert!( + out.warnings.iter().any(|x| x.contains("newer")), + "the override must be reported, not silent: {:?}", + out.warnings + ); + } + + /// The other direction is what keeps DISTRIBUTION working: a released + /// binary's bundled registry is older than master's, so the fetch must + /// still win and users get new models without upgrading. Without this the + /// change above would freeze every client at its build-time registry. + #[test] + fn older_bundled_registry_defers_to_the_fetch() { + let mut w = Vec::new(); + let fetched = LoadedRegistry { + registry: reg_at("2026-09-01T13:14:39Z"), + source: RegistrySource::Network, + warnings: Vec::new(), + }; + let out = prefer_bundled_if_newer(fetched, reg_at("2026-08-28T08:31:54Z"), &mut w); + assert_eq!(out.source, RegistrySource::Network); + assert_eq!(out.registry.generated_at, "2026-09-01T13:14:39Z"); + assert!(out.warnings.is_empty(), "no override, so nothing to report"); + } + + /// Equal stamps must not flap between sources. + #[test] + fn equal_timestamps_keep_the_fetched_registry() { + let mut w = Vec::new(); + let same = "2026-09-01T13:14:39Z"; + let fetched = LoadedRegistry { + registry: reg_at(same), + source: RegistrySource::Cache, + warnings: Vec::new(), + }; + let out = prefer_bundled_if_newer(fetched, reg_at(same), &mut w); + assert_eq!(out.source, RegistrySource::Cache); + } + #[test] fn bundled_registry_is_strictly_valid() { let registry = bundled().unwrap(); @@ -947,7 +1173,7 @@ mod tests { // Original Qwen3 family (without .5/.6/.8) receives no automatic policy. let qwen3_raw = r#"{ "schema_version":1, - "generated_at":"now", + "generated_at":"2026-09-01T00:00:00Z", "models":{"qwen3:8b":{"repo":"x","file":"qwen3-8b.mq4","size_gb":1,"min_vram_gb":1,"desc":"x","default_kv_mode":"q8"}}, "aliases":{} }"#; @@ -1157,7 +1383,7 @@ mod tests { fn malformed_entry_rejects_the_whole_registry() { let raw = r#"{ "schema_version":1, - "generated_at":"now", + "generated_at":"2026-09-01T00:00:00Z", "models":{"bad":{"repo":"x","file":"x","size_gb":1,"min_vram_gb":1,"desc":"x","default_kv_mode":"magic4"}}, "aliases":{} }"#; @@ -1168,7 +1394,7 @@ mod tests { fn tag_policy_pins_qwen_deepseek_and_glimmer_targets() { let raw = r#"{ "schema_version":1, - "generated_at":"now", + "generated_at":"2026-09-01T00:00:00Z", "models":{ "qwen3.5:4b":{"repo":"x","file":"qwen3.5-4b.mq4","size_gb":1,"min_vram_gb":1,"desc":"x"}, "qwen3.6:35b-a3b":{"repo":"x","file":"qwen3.6-35b.mq4","size_gb":1,"min_vram_gb":1,"desc":"x"}, @@ -1338,7 +1564,7 @@ mod tests { // Old v1 JSON without the invented fields must still parse (deny_unknown_fields). let raw = r#"{ "schema_version":1, - "generated_at":"now", + "generated_at":"2026-09-01T00:00:00Z", "models":{"ok":{"repo":"x","file":"x","size_gb":1,"min_vram_gb":1,"desc":"x"}}, "aliases":{} }"#; @@ -1355,7 +1581,7 @@ mod tests { // Invented wire fields must be rejected (no schema expansion). let bad = r#"{ "schema_version":1, - "generated_at":"now", + "generated_at":"2026-09-01T00:00:00Z", "models":{"bad":{"repo":"x","file":"x","size_gb":1,"min_vram_gb":1,"desc":"x","default_kv_backend":"vmm"}}, "aliases":{} }"#; @@ -1365,14 +1591,14 @@ mod tests { ); let bad2 = r#"{ "schema_version":1, - "generated_at":"now", + "generated_at":"2026-09-01T00:00:00Z", "models":{"bad":{"repo":"x","file":"x","size_gb":1,"min_vram_gb":1,"desc":"x","default_max_seq":262144}}, "aliases":{} }"#; assert!(RegistryV1::parse(bad2, "test").is_err()); let bad3 = r#"{ "schema_version":1, - "generated_at":"now", + "generated_at":"2026-09-01T00:00:00Z", "models":{"bad":{"repo":"x","file":"x","size_gb":1,"min_vram_gb":1,"desc":"x","default_max_tokens":81920}}, "aliases":{} }"#; @@ -1384,7 +1610,7 @@ mod tests { // Registry tag policy is a low-precedence layer; global/model/one-shot user config wins. let raw = r#"{ "schema_version":1, - "generated_at":"now", + "generated_at":"2026-09-01T00:00:00Z", "models":{"qwen3.8:27b":{"repo":"x","file":"qwen3.8-27b.mq4","size_gb":1,"min_vram_gb":1,"desc":"x"}}, "aliases":{} }"#; @@ -1446,7 +1672,7 @@ mod tests { // Glimmer target override likewise wins (backend + max_seq). let raw2 = r#"{ "schema_version":1, - "generated_at":"now", + "generated_at":"2026-09-01T00:00:00Z", "models":{"muse-glimmer":{"repo":"x","file":"muse-glimmer-30b.mq4","size_gb":1,"min_vram_gb":1,"desc":"x"}}, "aliases":{} }"#; @@ -1502,7 +1728,7 @@ mod tests { // DeepSeek target override wins over 1M/384Ki policy. let raw3 = r#"{ "schema_version":1, - "generated_at":"now", + "generated_at":"2026-09-01T00:00:00Z", "models":{"deepseek-v4-flash":{"repo":"x","file":"ds4.mq2r","size_gb":1,"min_vram_gb":1,"desc":"x"}}, "aliases":{} }"#; @@ -1556,7 +1782,7 @@ mod tests { fn dangling_aliases_are_dropped() { let raw = r#"{ "schema_version":1, - "generated_at":"now", + "generated_at":"2026-09-01T00:00:00Z", "models":{"ok":{"repo":"x","file":"x","size_gb":1,"min_vram_gb":1,"desc":"x"}}, "aliases":{"good":"ok","bad":"missing"} }"#; @@ -1579,7 +1805,7 @@ mod tests { fn sampling_profiles_resolve_per_mode_with_general_fallback() { let raw = r#"{ "schema_version":1, - "generated_at":"now", + "generated_at":"2026-09-01T00:00:00Z", "models":{"m":{ "repo":"x","file":"x","size_gb":1,"min_vram_gb":1,"desc":"x", "recommended_settings":{"temperature":1.0,"presence_penalty":1.5}, @@ -1609,7 +1835,7 @@ mod tests { fn out_of_range_sampling_profile_rejects_the_whole_registry() { let raw = r#"{ "schema_version":1, - "generated_at":"now", + "generated_at":"2026-09-01T00:00:00Z", "models":{"m":{ "repo":"x","file":"x","size_gb":1,"min_vram_gb":1,"desc":"x", "sampling_profiles":{"coding":{"temperature":9.0}} @@ -1695,23 +1921,23 @@ mod tests { // bundled (network) or discards the cache entry (fresh/stale). let cases = [ // Qwen3.8 product SKUs - r#"{"schema_version":1,"generated_at":"now","models":{"qwen3.8:27b":{"repo":"x","file":"x","size_gb":1,"min_vram_gb":1,"desc":"x","recommended_settings":{"reasoning_effort":"xhigh","thinking_budget":"high"}}},"aliases":{}}"#, + r#"{"schema_version":1,"generated_at":"2026-09-01T00:00:00Z","models":{"qwen3.8:27b":{"repo":"x","file":"x","size_gb":1,"min_vram_gb":1,"desc":"x","recommended_settings":{"reasoning_effort":"xhigh","thinking_budget":"high"}}},"aliases":{}}"#, // DeepSeek V4 Flash (also covers :mq2lloyd via family) - r#"{"schema_version":1,"generated_at":"now","models":{"deepseek-v4-flash":{"repo":"x","file":"x","size_gb":1,"min_vram_gb":1,"desc":"x","recommended_settings":{"reasoning_effort":"low","thinking_budget":"uncapped"}}},"aliases":{}}"#, - r#"{"schema_version":1,"generated_at":"now","models":{"deepseek-v4-flash-preview":{"repo":"x","file":"x","size_gb":1,"min_vram_gb":1,"desc":"x","recommended_settings":{"reasoning_effort":"high","thinking_budget":"med"}}},"aliases":{}}"#, - r#"{"schema_version":1,"generated_at":"now","models":{"deepseek-v4-flash:mq2lloyd":{"repo":"x","file":"x","size_gb":1,"min_vram_gb":1,"desc":"x","recommended_settings":{"thinking_budget":"low"}}},"aliases":{}}"#, + r#"{"schema_version":1,"generated_at":"2026-09-01T00:00:00Z","models":{"deepseek-v4-flash":{"repo":"x","file":"x","size_gb":1,"min_vram_gb":1,"desc":"x","recommended_settings":{"reasoning_effort":"low","thinking_budget":"uncapped"}}},"aliases":{}}"#, + r#"{"schema_version":1,"generated_at":"2026-09-01T00:00:00Z","models":{"deepseek-v4-flash-preview":{"repo":"x","file":"x","size_gb":1,"min_vram_gb":1,"desc":"x","recommended_settings":{"reasoning_effort":"high","thinking_budget":"med"}}},"aliases":{}}"#, + r#"{"schema_version":1,"generated_at":"2026-09-01T00:00:00Z","models":{"deepseek-v4-flash:mq2lloyd":{"repo":"x","file":"x","size_gb":1,"min_vram_gb":1,"desc":"x","recommended_settings":{"thinking_budget":"low"}}},"aliases":{}}"#, // Muse Glimmer product SKUs - r#"{"schema_version":1,"generated_at":"now","models":{"muse-glimmer":{"repo":"x","file":"x","size_gb":1,"min_vram_gb":1,"desc":"x","recommended_settings":{"reasoning_effort":"xhigh","thinking_budget":"xhigh"}}},"aliases":{}}"#, - r#"{"schema_version":1,"generated_at":"now","models":{"muse-glimmer:fast":{"repo":"x","file":"x","size_gb":1,"min_vram_gb":1,"desc":"x","recommended_settings":{"thinking_budget":"max"}}},"aliases":{}}"#, + r#"{"schema_version":1,"generated_at":"2026-09-01T00:00:00Z","models":{"muse-glimmer":{"repo":"x","file":"x","size_gb":1,"min_vram_gb":1,"desc":"x","recommended_settings":{"reasoning_effort":"xhigh","thinking_budget":"xhigh"}}},"aliases":{}}"#, + r#"{"schema_version":1,"generated_at":"2026-09-01T00:00:00Z","models":{"muse-glimmer:fast":{"repo":"x","file":"x","size_gb":1,"min_vram_gb":1,"desc":"x","recommended_settings":{"thinking_budget":"max"}}},"aliases":{}}"#, // Ornith 1.5 product + legacy family spellings - r#"{"schema_version":1,"generated_at":"now","models":{"ornith-1.5:35b-a3b":{"repo":"x","file":"x","size_gb":1,"min_vram_gb":1,"desc":"x","recommended_settings":{"reasoning_effort":"xhigh","thinking_budget":"high"}}},"aliases":{}}"#, - r#"{"schema_version":1,"generated_at":"now","models":{"ornith1.5:35b-a3b":{"repo":"x","file":"x","size_gb":1,"min_vram_gb":1,"desc":"x","recommended_settings":{"thinking_budget":"uncapped"}}},"aliases":{}}"#, - r#"{"schema_version":1,"generated_at":"now","models":{"ornith:35b-a3b":{"repo":"x","file":"x","size_gb":1,"min_vram_gb":1,"desc":"x","recommended_settings":{"thinking_budget":"med"}}},"aliases":{}}"#, + r#"{"schema_version":1,"generated_at":"2026-09-01T00:00:00Z","models":{"ornith-1.5:35b-a3b":{"repo":"x","file":"x","size_gb":1,"min_vram_gb":1,"desc":"x","recommended_settings":{"reasoning_effort":"xhigh","thinking_budget":"high"}}},"aliases":{}}"#, + r#"{"schema_version":1,"generated_at":"2026-09-01T00:00:00Z","models":{"ornith1.5:35b-a3b":{"repo":"x","file":"x","size_gb":1,"min_vram_gb":1,"desc":"x","recommended_settings":{"thinking_budget":"uncapped"}}},"aliases":{}}"#, + r#"{"schema_version":1,"generated_at":"2026-09-01T00:00:00Z","models":{"ornith:35b-a3b":{"repo":"x","file":"x","size_gb":1,"min_vram_gb":1,"desc":"x","recommended_settings":{"thinking_budget":"med"}}},"aliases":{}}"#, // Effort-native sampling_profiles also rejected - r#"{"schema_version":1,"generated_at":"now","models":{"qwen3.8:27b":{"repo":"x","file":"x","size_gb":1,"min_vram_gb":1,"desc":"x","sampling_profiles":{"coding":{"reasoning_effort":"xhigh","thinking_budget":"high"}}}},"aliases":{}}"#, - r#"{"schema_version":1,"generated_at":"now","models":{"deepseek-v4-flash":{"repo":"x","file":"x","size_gb":1,"min_vram_gb":1,"desc":"x","sampling_profiles":{"general":{"thinking_budget":"low"}}}},"aliases":{}}"#, - r#"{"schema_version":1,"generated_at":"now","models":{"muse-glimmer":{"repo":"x","file":"x","size_gb":1,"min_vram_gb":1,"desc":"x","sampling_profiles":{"general":{"thinking_budget":"uncapped"}}}},"aliases":{}}"#, - r#"{"schema_version":1,"generated_at":"now","models":{"ornith-1.5:35b-a3b":{"repo":"x","file":"x","size_gb":1,"min_vram_gb":1,"desc":"x","sampling_profiles":{"general":{"thinking_budget":"high"}}}},"aliases":{}}"#, + r#"{"schema_version":1,"generated_at":"2026-09-01T00:00:00Z","models":{"qwen3.8:27b":{"repo":"x","file":"x","size_gb":1,"min_vram_gb":1,"desc":"x","sampling_profiles":{"coding":{"reasoning_effort":"xhigh","thinking_budget":"high"}}}},"aliases":{}}"#, + r#"{"schema_version":1,"generated_at":"2026-09-01T00:00:00Z","models":{"deepseek-v4-flash":{"repo":"x","file":"x","size_gb":1,"min_vram_gb":1,"desc":"x","sampling_profiles":{"general":{"thinking_budget":"low"}}}},"aliases":{}}"#, + r#"{"schema_version":1,"generated_at":"2026-09-01T00:00:00Z","models":{"muse-glimmer":{"repo":"x","file":"x","size_gb":1,"min_vram_gb":1,"desc":"x","sampling_profiles":{"general":{"thinking_budget":"uncapped"}}}},"aliases":{}}"#, + r#"{"schema_version":1,"generated_at":"2026-09-01T00:00:00Z","models":{"ornith-1.5:35b-a3b":{"repo":"x","file":"x","size_gb":1,"min_vram_gb":1,"desc":"x","sampling_profiles":{"general":{"thinking_budget":"high"}}}},"aliases":{}}"#, ]; for raw in cases { let err = RegistryV1::parse(raw, "network/cache") @@ -1724,7 +1950,7 @@ mod tests { } // A cache entry that violates the invariant is also rejected via // validate(), causing read_cache to return None and load() to fall back. - let stale_raw = r#"{"schema_version":1,"generated_at":"now","models":{"qwen3.8:27b":{"repo":"x","file":"x","size_gb":1,"min_vram_gb":1,"desc":"x","recommended_settings":{"thinking_budget":"high"}}},"aliases":{}}"#; + let stale_raw = r#"{"schema_version":1,"generated_at":"2026-09-01T00:00:00Z","models":{"qwen3.8:27b":{"repo":"x","file":"x","size_gb":1,"min_vram_gb":1,"desc":"x","recommended_settings":{"thinking_budget":"high"}}},"aliases":{}}"#; let stale: RegistryV1 = serde_json::from_str(stale_raw).unwrap(); assert!( stale.validate("registry cache").is_err(), @@ -1737,23 +1963,23 @@ mod tests { // Mirrors hipfire-config/registry_gen enum allowlists: recognizable // invalid values are rejected with a clear error; malformed types // already fail via surrounding validation and are not re-tested here. - let invalid_effort = r#"{"schema_version":1,"generated_at":"now","models":{"m":{"repo":"x","file":"x","size_gb":1,"min_vram_gb":1,"desc":"x","recommended_settings":{"reasoning_effort":"turbo"}}},"aliases":{}}"#; + let invalid_effort = r#"{"schema_version":1,"generated_at":"2026-09-01T00:00:00Z","models":{"m":{"repo":"x","file":"x","size_gb":1,"min_vram_gb":1,"desc":"x","recommended_settings":{"reasoning_effort":"turbo"}}},"aliases":{}}"#; let err = RegistryV1::parse(invalid_effort, "test") .expect_err("invalid reasoning_effort must be rejected"); assert!(err.to_string().contains("reasoning_effort")); - let invalid_effort_profile = r#"{"schema_version":1,"generated_at":"now","models":{"m":{"repo":"x","file":"x","size_gb":1,"min_vram_gb":1,"desc":"x","sampling_profiles":{"coding":{"reasoning_effort":"ultra"}}}},"aliases":{}}"#; + let invalid_effort_profile = r#"{"schema_version":1,"generated_at":"2026-09-01T00:00:00Z","models":{"m":{"repo":"x","file":"x","size_gb":1,"min_vram_gb":1,"desc":"x","sampling_profiles":{"coding":{"reasoning_effort":"ultra"}}}},"aliases":{}}"#; let err = RegistryV1::parse(invalid_effort_profile, "test") .expect_err("invalid profile effort must be rejected"); assert!(err.to_string().contains("reasoning_effort")); // thinking_budget invalid on legacy (non-effort-native) model - let invalid_budget = r#"{"schema_version":1,"generated_at":"now","models":{"qwen3.5:9b":{"repo":"x","file":"x","size_gb":1,"min_vram_gb":1,"desc":"x","recommended_settings":{"thinking_budget":"yolo"}}},"aliases":{}}"#; + let invalid_budget = r#"{"schema_version":1,"generated_at":"2026-09-01T00:00:00Z","models":{"qwen3.5:9b":{"repo":"x","file":"x","size_gb":1,"min_vram_gb":1,"desc":"x","recommended_settings":{"thinking_budget":"yolo"}}},"aliases":{}}"#; let err = RegistryV1::parse(invalid_budget, "test") .expect_err("invalid thinking_budget must be rejected"); assert!(err.to_string().contains("thinking_budget")); - let invalid_budget_profile = r#"{"schema_version":1,"generated_at":"now","models":{"qwen3.6:35b-a3b":{"repo":"x","file":"x","size_gb":1,"min_vram_gb":1,"desc":"x","sampling_profiles":{"general":{"thinking_budget":"superhigh"}}}},"aliases":{}}"#; + let invalid_budget_profile = r#"{"schema_version":1,"generated_at":"2026-09-01T00:00:00Z","models":{"qwen3.6:35b-a3b":{"repo":"x","file":"x","size_gb":1,"min_vram_gb":1,"desc":"x","sampling_profiles":{"general":{"thinking_budget":"superhigh"}}}},"aliases":{}}"#; let err = RegistryV1::parse(invalid_budget_profile, "test") .expect_err("invalid profile budget must be rejected"); assert!(err.to_string().contains("thinking_budget")); @@ -1765,7 +1991,7 @@ mod tests { // remain valid and pass validation for both top-level and profiles. let raw = r#"{ "schema_version":1, - "generated_at":"now", + "generated_at":"2026-09-01T00:00:00Z", "models":{ "qwen3.5:9b":{"repo":"x","file":"x","size_gb":1,"min_vram_gb":1,"desc":"x","recommended_settings":{"reasoning_effort":"high","thinking_budget":"high"}}, "qwen3.5:27b":{"repo":"x","file":"x","size_gb":1,"min_vram_gb":1,"desc":"x","sampling_profiles":{"general":{"thinking_budget":"med"},"coding":{"reasoning_effort":"max","thinking_budget":"max"}}}, @@ -1780,7 +2006,7 @@ mod tests { // also accept thinking_budget even when family would otherwise be native. let sidecars = r#"{ "schema_version":1, - "generated_at":"now", + "generated_at":"2026-09-01T00:00:00Z", "models":{ "qwen3.8:27b-draft":{"repo":"x","file":"x","size_gb":1,"min_vram_gb":1,"desc":"x","recommended_settings":{"thinking_budget":"high"}}, "qwen3.8:27b-dflash":{"repo":"x","file":"x","size_gb":1,"min_vram_gb":1,"desc":"x","recommended_settings":{"thinking_budget":"low"}}, diff --git a/crates/hipfire-runtime/map.md b/crates/hipfire-runtime/map.md index 7642d2fe90..ac5fe32f49 100644 --- a/crates/hipfire-runtime/map.md +++ b/crates/hipfire-runtime/map.md @@ -46,15 +46,15 @@ _Generated by `scripts/check-crate-maps.py` from the tree — do not edit inside | [`src/ep.rs`](src/ep.rs) | 296 | 2 | 0 | | [`src/eval_common.rs`](src/eval_common.rs) | 231 | 3 | 0 | | [`src/gguf.rs`](src/gguf.rs) | 335 | 20 | 0 | -| [`src/hfq.rs`](src/hfq.rs) | 2,643 | 51 | 11 | +| [`src/hfq.rs`](src/hfq.rs) | 2,690 | 52 | 11 | | [`src/hfq_parallel.rs`](src/hfq_parallel.rs) | 335 | 8 | 2 | | [`src/kv_adaptive.rs`](src/kv_adaptive.rs) | 608 | 24 | 12 | | [`src/kv_backend.rs`](src/kv_backend.rs) | 129 | 1 | 7 | -| [`src/kv_mode.rs`](src/kv_mode.rs) | 298 | 10 | 7 | +| [`src/kv_mode.rs`](src/kv_mode.rs) | 389 | 11 | 9 | | [`src/lib.rs`](src/lib.rs) | 78 | 53 | 0 | -| [`src/llama.rs`](src/llama.rs) | 8,792 | 84 | 42 | +| [`src/llama.rs`](src/llama.rs) | 8,819 | 84 | 42 | | [`src/llama_spec.rs`](src/llama_spec.rs) | 617 | 6 | 1 | -| [`src/loader_api.rs`](src/loader_api.rs) | 309 | 14 | 4 | +| [`src/loader_api.rs`](src/loader_api.rs) | 314 | 14 | 4 | | [`src/loop_guard.rs`](src/loop_guard.rs) | 194 | 8 | 4 | | [`src/model_load.rs`](src/model_load.rs) | 541 | 8 | 3 | | [`src/model_source.rs`](src/model_source.rs) | 92 | 4 | 0 | @@ -105,11 +105,11 @@ _Generated by `scripts/check-crate-maps.py` from the tree — do not edit inside - [`src/ep.rs`](src/ep.rs): `ensure_rank_streams`, `run_layer_program_ep` - [`src/eval_common.rs`](src/eval_common.rs): `verify_ref_sha256`, `verify_slice_md5`, `verify_llama_commit` - [`src/gguf.rs`](src/gguf.rs): `GgmlType`, `from_u32`, `block_size`, `block_bytes`, `tensor_bytes`, `MetaValue`, `as_u32`, `as_f32`, `as_str`, `TensorInfo`, `numel`, `byte_size`, +8 more -- [`src/hfq.rs`](src/hfq.rs): `start_cache_warmup`, `mostly_page_cached`, `mostly_page_cached_memo`, `CacheWarmerGuard`, `HfqTensorInfo`, `HfqTensorManifestEntry`, `HfqSourceIdentity`, `RecommendedSampling`, `HfqFile`, `open`, `open_with_reap_plan`, `attach_overlay`, +39 more +- [`src/hfq.rs`](src/hfq.rs): `start_cache_warmup`, `mostly_page_cached`, `mostly_page_cached_memo`, `CacheWarmerGuard`, `HfqTensorInfo`, `HfqTensorManifestEntry`, `HfqSourceIdentity`, `RecommendedSampling`, `HfqFile`, `open`, `open_with_reap_plan`, `attach_head_overlay`, +40 more - [`src/hfq_parallel.rs`](src/hfq_parallel.rs): `HFQ_READER_LANES`, `HfqReadJob`, `tensor`, `packed`, `label`, `output_len`, `HfqReadResult`, `read_hfq_jobs_ordered` - [`src/kv_adaptive.rs`](src/kv_adaptive.rs): `KMode`, `bytes_per_head`, `rot_width`, `bits`, `v_bytes_per_head`, `k_buf_bytes_per_layer`, `v_buf_bytes_per_layer`, `cap_min`, `Step`, `Preset`, `KvAdaptive`, `from_preset`, +12 more - [`src/kv_backend.rs`](src/kv_backend.rs): `saddle_core` -- [`src/kv_mode.rs`](src/kv_mode.rs): `saddle_core`, `KvModePolicy`, `ResolveResult`, `QWEN35_HFQ_POLICY`, `QWEN35_PARO_POLICY`, `DIR_SAFETENSORS_POLICY`, `LLAMA_HFQ_POLICY`, `HFQ_Q8_ONLY_POLICY`, `QWEN35_PP_POLICY`, `resolve` +- [`src/kv_mode.rs`](src/kv_mode.rs): `saddle_core`, `KvModePolicy`, `ResolveResult`, `QWEN35_HFQ_POLICY`, `QWEN35_PARO_POLICY`, `DIR_SAFETENSORS_POLICY`, `LLAMA_HFQ_POLICY`, `HFQ_Q8_ONLY_POLICY`, `QWEN35_PP_POLICY`, `MAPLE_POLICY`, `resolve` - [`src/lib.rs`](src/lib.rs): `admission`, `arch`, `arch_mapping`, `arch_model`, `arch_spec`, `augmentor`, `bf16_loader`, `cache_plan`, `cask`, `config`, `cpu_router`, `ddtree`, +41 more - [`src/llama.rs`](src/llama.rs): `ModelArch`, `LlamaConfig`, `from_gguf`, `dequantize_q4_0`, `dequantize_q8_0`, `f16_to_f32`, `f32_to_f16`, `dequantize_q4_k`, `convert_q4k_to_q4f16_g64`, `convert_q4k_to_q4f16_g32`, `dequantize_q6_k`, `ParoRotation`, +72 more - [`src/llama_spec.rs`](src/llama_spec.rs): `verify_block_argmax`, `verify_block_logits`, `verify_block_argmax_capture_gpu`, `verify_block_sampled_capture_gpu`, `verify_tree_logits`, `lm_head_logits_n_rows` @@ -152,6 +152,6 @@ _Generated by `scripts/check-crate-maps.py` from the tree — do not edit inside ### Totals -- 56 modules · 51,578 lines · 853 public items · 604 tests · 132 examples +- 56 modules · 51,748 lines · 855 public items · 606 tests · 132 examples diff --git a/crates/hipfire-runtime/src/hfq.rs b/crates/hipfire-runtime/src/hfq.rs index 80a58a713c..e957a266b6 100644 --- a/crates/hipfire-runtime/src/hfq.rs +++ b/crates/hipfire-runtime/src/hfq.rs @@ -42,6 +42,9 @@ fn fadvise_dontneed(fd: std::os::unix::io::RawFd, offset: usize, len: usize) { #[cfg(not(unix))] fn fadvise_dontneed(_fd: i32, _offset: usize, _len: usize) {} +/// The only tensor a head overlay may carry. +const HEAD_TENSOR_NAME: &str = "lm_head.weight"; + impl HfqFile { /// Start a background parallel cache warmer: N worker threads pread the /// data region chunk-sequentially into the page cache while the loader @@ -388,6 +391,50 @@ impl HfqFile { Ok(f) } + /// Attach a HEAD overlay: a single-tensor `.hfq` (built by + /// `hipfire-quantize --head-only`) whose `lm_head.weight` shadows the + /// base's, so one body can serve several head carriers. + /// + /// Same guards as [`Self::attach_overlay`] — matching arch_id, the name + /// must already exist in the base, and the logical shape must match; only + /// the quant tier may differ. Failure is an ERROR, not a warning: unlike + /// the REAP path (where proceeding unpruned is a safe default for an + /// unrelated model that merely shares an env var), a head overlay is + /// requested explicitly, so silently serving the base's head would hand + /// back a model the operator did not ask for. + pub fn attach_head_overlay(&mut self, head_path: &Path) -> Result<(), String> { + let ov = Self::open_at_offset(head_path, 0) + .map_err(|e| format!("head overlay {head_path:?}: {e}"))?; + // A head overlay must contain ONLY head tensors. Without this, passing + // a full model to --head "succeeds": every name exists in the base with + // a matching shape, so attach_overlay's guards all pass and the entire + // model silently shadows itself. Caught by a negative control that + // passed the base as its own overlay. + let foreign: Vec<&str> = ov + .tensors + .iter() + .map(|t| t.name.as_str()) + .filter(|n| *n != HEAD_TENSOR_NAME) + .collect(); + if !foreign.is_empty() { + return Err(format!( + "head overlay {head_path:?}: expected only `{HEAD_TENSOR_NAME}`, found {} \ + tensor(s) including `{}` — this looks like a full model, not a \ + `hipfire-quantize --head-only` build", + ov.tensors.len(), + foreign[0], + )); + } + if ov.tensors.is_empty() { + return Err(format!("head overlay {head_path:?}: contains no tensors")); + } + let n = ov.tensors.len(); + self.attach_overlay(ov) + .map_err(|e| format!("head overlay {head_path:?}: {e}"))?; + eprintln!(" head overlay: {n} tensor(s) from {head_path:?} shadow the base"); + Ok(()) + } + /// Attach an overlay whose tensors shadow this file's by name. Used by the /// REAP load-time splice (SP3). Errors if arch_id differs (wrong model). pub fn attach_overlay(&mut self, overlay: HfqFile) -> Result<(), String> { diff --git a/crates/hipfire-runtime/src/kv_mode.rs b/crates/hipfire-runtime/src/kv_mode.rs index d4f4ad2fa1..d26a8e4ba6 100644 --- a/crates/hipfire-runtime/src/kv_mode.rs +++ b/crates/hipfire-runtime/src/kv_mode.rs @@ -134,6 +134,43 @@ pub const QWEN35_PP_POLICY: KvModePolicy = KvModePolicy { default: Q8, }; +/// Site 7 — maple (arch 15). Before this site existed, maple hardcoded +/// `KvCache::new_gpu_q8` and `--kv-mode` was a silent no-op for arch 15. +/// +/// The accept set is deliberately just {Q8, Bf16}. Every other mode in the +/// ladder is a rotated or block-quantized tier whose attention kernels have NO +/// sliding-window variant, and Maple is 3:1 sliding(512)/global — a tier that +/// cannot carry the window would attend the full context on the sliding layers +/// and be silently WRONG at ctx > 512, not merely slower. Accepting them and +/// warning is the wrong trade here; refusing to the q8 default is right. +/// +/// `"bf16"` is the only new name, and it is deliberately NOT added to +/// `normalize_full`: no other site can allocate a bf16 cache, so putting it +/// there would let `HIPFIRE_KV_MODE=bf16` on qwen35 normalize successfully and +/// then fall to that site's default — a silent downgrade instead of a warning. +/// +/// **The default is bf16, not q8.** Measured against a bf16 reference on 2048 +/// teacher-forced wikitext tokens, q8 KV costs 39% of the total divergence +/// (mean KL 0.0842 q8 vs 0.0511 bf16, top-1 90.8% vs 91.9%), and the damage is +/// in the TAIL rather than as uniform blur — the median moves only 24% but the +/// worst position goes 10.36 -> 4.21 nats. The price is 1.88x KV bytes +/// (26,112 -> 49,152 B/token) and about 2% decode, which is inside this box's +/// run-to-run noise. `--kv-mode q8` restores the old tier for anyone who wants +/// the memory back. +fn normalize_maple(raw: &str) -> Option { + match raw { + "bf16" | "auto" | "" => Some(Bf16), + "q8" => Some(Q8), + _ => None, // every rotated/quantized tier → default (+warn) + } +} +pub const MAPLE_POLICY: KvModePolicy = KvModePolicy { + site: "maple", + normalize_alias: normalize_maple, + accepted: &[Q8, Bf16], + default: Bf16, +}; + /// Pure: `&str + &'static policy + usize → ResolveResult`. No GPU, no env read. pub fn resolve(raw: &str, policy: &KvModePolicy, head_dim: usize) -> ResolveResult { // 1. site-LOCAL alias expansion. @@ -174,6 +211,60 @@ pub fn resolve(raw: &str, policy: &KvModePolicy, head_dim: usize) -> ResolveResu mod tests { use super::*; + #[test] + fn truth_table_maple() { + let p = &MAPLE_POLICY; + assert_eq!(resolve("bf16", p, 128).mode, KvMode::Bf16); + assert_eq!(resolve("q8", p, 128).mode, KvMode::Q8); + // Unset and "auto" both mean BF16, SILENTLY — bf16 is the shipped + // default and must not print a warning on every load. + assert_eq!(resolve("", p, 128).mode, KvMode::Bf16); + assert!(resolve("", p, 128).warning.is_none()); + assert_eq!(resolve("auto", p, 128).mode, KvMode::Bf16); + assert!(resolve("auto", p, 128).warning.is_none()); + // Asking for q8 explicitly is HONORED and must not warn — it is a + // supported tier and an intentional memory saving, not a degradation. + assert!(resolve("q8", p, 128).warning.is_none()); + + // Every ROTATED / block-quantized tier must be REFUSED and warn. + // These have no sliding-window attention kernel, so silently accepting + // one would make Maple's sliding layers attend the full context and be + // wrong past 512 tokens rather than merely slower. + for m in [ + "asym2", "asym3", "asym4", "fwht2", "fwht3", "fwht4", "turbo", + ] { + let r = resolve(m, p, 128); + assert_eq!(r.mode, KvMode::Bf16, "{m} must fall back to the default"); + assert!(r.warning.is_some(), "{m} must warn, not silently downgrade"); + } + let garbage = resolve("garbage", p, 128); + assert_eq!(garbage.mode, KvMode::Bf16); + assert!(garbage.warning.is_some()); + } + + #[test] + fn bf16_is_maple_only() { + // NEGATIVE CONTROL: "bf16" must not be a globally-known alias. No other + // site can allocate a bf16 cache, so if `normalize_full` learned the + // name, HIPFIRE_KV_MODE=bf16 on qwen35 would normalize fine and then + // silently fall to that site's default. It must warn instead. + for p in [ + &QWEN35_HFQ_POLICY, + &QWEN35_PARO_POLICY, + &LLAMA_HFQ_POLICY, + &QWEN35_PP_POLICY, + &DIR_SAFETENSORS_POLICY, + ] { + let r = resolve("bf16", p, 256); + assert_ne!(r.mode, KvMode::Bf16, "site {} must not accept bf16", p.site); + assert!( + r.warning.is_some(), + "site {} must WARN on bf16, not silently default", + p.site + ); + } + } + #[test] fn truth_table_qwen35_hfq() { let p = &QWEN35_HFQ_POLICY; diff --git a/crates/hipfire-runtime/src/llama.rs b/crates/hipfire-runtime/src/llama.rs index 0660a09db5..496995d9e9 100644 --- a/crates/hipfire-runtime/src/llama.rs +++ b/crates/hipfire-runtime/src/llama.rs @@ -5923,6 +5923,7 @@ impl KvCacheExt for KvCache { self.quant_int8, is_hfq8, self.quant_fwht, + self.quant_bf16, ) } @@ -5938,6 +5939,7 @@ impl KvCacheExt for KvCache { quant_q4, quant_int8: self.quant_int8, quant_hfq8: self.is_hfq8_kv(), + quant_bf16: self.quant_bf16, f32_policy: hipfire_dispatch::families::kv_tier::F32AttnPolicy::Simple, v_mode_bits: self.v_mode.bits() as i32, pos: 0, @@ -6062,6 +6064,7 @@ impl KvCacheExt for KvCache { quant_asym3: false, quant_asym2: false, quant_fwht: false, + quant_bf16: false, boundary_layers: 0, givens_cos: None, givens_sin: None, @@ -6102,6 +6105,7 @@ impl KvCacheExt for KvCache { quant_asym3: false, quant_asym2: false, quant_fwht: false, + quant_bf16: false, boundary_layers: 0, givens_cos: None, givens_sin: None, @@ -6161,6 +6165,7 @@ impl KvCacheExt for KvCache { quant_asym3: false, quant_asym2: false, quant_fwht: false, + quant_bf16: false, boundary_layers: 0, givens_cos: None, givens_sin: None, @@ -6201,6 +6206,7 @@ impl KvCacheExt for KvCache { quant_asym3: false, quant_asym2: false, quant_fwht: false, + quant_bf16: false, boundary_layers: 0, givens_cos: None, givens_sin: None, @@ -6241,6 +6247,7 @@ impl KvCacheExt for KvCache { quant_asym3: false, quant_asym2: false, quant_fwht: false, + quant_bf16: false, boundary_layers: 0, givens_cos: None, givens_sin: None, @@ -6286,6 +6293,7 @@ impl KvCacheExt for KvCache { quant_asym3: false, quant_asym2: false, quant_fwht: false, + quant_bf16: false, boundary_layers: 0, givens_cos: None, givens_sin: None, @@ -6331,6 +6339,7 @@ impl KvCacheExt for KvCache { quant_asym3: false, quant_asym2: false, quant_fwht: false, + quant_bf16: false, boundary_layers: 0, givens_cos: None, givens_sin: None, @@ -6397,6 +6406,7 @@ impl KvCacheExt for KvCache { quant_asym3: false, quant_asym2: false, quant_fwht: false, + quant_bf16: false, boundary_layers: 0, givens_cos: None, givens_sin: None, @@ -6463,6 +6473,7 @@ impl KvCacheExt for KvCache { quant_asym3: true, quant_asym2: false, quant_fwht: false, + quant_bf16: false, boundary_layers: 0, givens_cos: None, givens_sin: None, @@ -6529,6 +6540,7 @@ impl KvCacheExt for KvCache { quant_asym3: false, quant_asym2: true, quant_fwht: false, + quant_bf16: false, boundary_layers: 0, givens_cos: None, givens_sin: None, @@ -6603,6 +6615,7 @@ impl KvCacheExt for KvCache { quant_asym3: false, quant_asym2: false, quant_fwht: true, + quant_bf16: false, boundary_layers: 0, givens_cos: None, givens_sin: None, @@ -6670,6 +6683,7 @@ impl KvCacheExt for KvCache { quant_asym3: true, quant_asym2: false, quant_fwht: true, + quant_bf16: false, boundary_layers: 0, givens_cos: None, givens_sin: None, @@ -6736,6 +6750,7 @@ impl KvCacheExt for KvCache { quant_asym3: false, quant_asym2: true, quant_fwht: true, + quant_bf16: false, boundary_layers: 0, givens_cos: None, givens_sin: None, @@ -6781,6 +6796,7 @@ impl KvCacheExt for KvCache { quant_asym3: false, quant_asym2: false, quant_fwht: false, + quant_bf16: false, boundary_layers: 0, givens_cos: None, givens_sin: None, @@ -6830,6 +6846,7 @@ impl KvCacheExt for KvCache { quant_asym3: false, quant_asym2: false, quant_fwht: false, + quant_bf16: false, boundary_layers: 0, givens_cos: None, givens_sin: None, @@ -6879,6 +6896,7 @@ impl KvCacheExt for KvCache { quant_asym3: true, quant_asym2: false, quant_fwht: false, + quant_bf16: false, boundary_layers: 0, givens_cos: None, givens_sin: None, @@ -6928,6 +6946,7 @@ impl KvCacheExt for KvCache { quant_asym3: false, quant_asym2: true, quant_fwht: false, + quant_bf16: false, boundary_layers: 0, givens_cos: None, givens_sin: None, @@ -6977,6 +6996,7 @@ impl KvCacheExt for KvCache { quant_asym3: false, quant_asym2: false, quant_fwht: true, + quant_bf16: false, boundary_layers: 0, givens_cos: None, givens_sin: None, @@ -7026,6 +7046,7 @@ impl KvCacheExt for KvCache { quant_asym3: true, quant_asym2: false, quant_fwht: true, + quant_bf16: false, boundary_layers: 0, givens_cos: None, givens_sin: None, @@ -7075,6 +7096,7 @@ impl KvCacheExt for KvCache { quant_asym3: false, quant_asym2: true, quant_fwht: true, + quant_bf16: false, boundary_layers: 0, givens_cos: None, givens_sin: None, @@ -8539,6 +8561,7 @@ mod tests { quant_q4: false, quant_int8: false, quant_hfq8: false, + quant_bf16: false, f32_policy: hipfire_dispatch::families::kv_tier::F32AttnPolicy::Simple, v_mode_bits: kv.v_mode_bits(), pos: 100, @@ -8585,6 +8608,7 @@ mod tests { quant_q4: false, quant_int8: false, quant_hfq8: false, + quant_bf16: false, f32_policy: hipfire_dispatch::families::kv_tier::F32AttnPolicy::Simple, v_mode_bits: kv.v_mode_bits(), pos: 100, @@ -8661,6 +8685,7 @@ mod tests { quant_q4: false, quant_int8: false, quant_hfq8: false, + quant_bf16: false, f32_policy: hipfire_dispatch::families::kv_tier::F32AttnPolicy::Simple, v_mode_bits: kv.v_mode_bits(), pos: 100, @@ -8706,6 +8731,7 @@ mod tests { quant_asym3: false, quant_asym2: false, quant_fwht: false, + quant_bf16: false, boundary_layers: 0, givens_cos: None, givens_sin: None, @@ -8750,6 +8776,7 @@ mod tests { quant_asym3: false, quant_asym2: false, quant_fwht: false, + quant_bf16: false, boundary_layers: 0, givens_cos: None, givens_sin: None, diff --git a/crates/hipfire-runtime/src/loader_api.rs b/crates/hipfire-runtime/src/loader_api.rs index bf115366fd..0563fd7edc 100644 --- a/crates/hipfire-runtime/src/loader_api.rs +++ b/crates/hipfire-runtime/src/loader_api.rs @@ -69,6 +69,11 @@ pub struct LoadCtx<'a> { pub deepseek4_experts_per_token: Option, pub draft_path: Option<&'a str>, pub kv_mode_override: Option<&'a str>, + /// Optional HEAD OVERLAY: a single-tensor `.hfq` from + /// `hipfire-quantize --head-only` whose `lm_head.weight` shadows the + /// base's, so one body serves several head carriers. `None` uses the head + /// baked into the model file. Only the maple carrier reads this today. + pub head_path: Option<&'a str>, pub kv_backend: KvBackend, pub kv_adaptive_override: Option<&'a str>, pub state_quant_override: Option<&'a str>, diff --git a/crates/hipfire-runtime/tests/kv_adaptive_reset.rs b/crates/hipfire-runtime/tests/kv_adaptive_reset.rs index 1ae1881726..d2f67238a1 100644 --- a/crates/hipfire-runtime/tests/kv_adaptive_reset.rs +++ b/crates/hipfire-runtime/tests/kv_adaptive_reset.rs @@ -33,6 +33,7 @@ fn flag_standin(mode: KvMode, v_mode: VMode, n_kv_heads: usize, head_dim: usize) quant_asym3: a3, quant_asym2: a2, quant_fwht: fwht, + quant_bf16: false, boundary_layers: 0, givens_cos: None, givens_sin: None, @@ -49,12 +50,8 @@ fn adaptive_reset_invalidates_captured_execution_state() { return; }; let mut cache = flag_standin(KvMode::Fwht2, VMode::Lloyd2, 4, 256); - let mut adaptive = kv_adaptive::KvAdaptive::from_preset( - kv_adaptive::Preset::Aggressive, - 128, - 4, - 256, - ); + let mut adaptive = + kv_adaptive::KvAdaptive::from_preset(kv_adaptive::Preset::Aggressive, 128, 4, 256); adaptive.cur_k = kv_adaptive::KMode::Fwht2; adaptive.cur_v = VMode::Lloyd2; adaptive.next_step = adaptive.steps.len(); diff --git a/crates/rdna-compute/map.md b/crates/rdna-compute/map.md index 258dd5336f..61435422d6 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) | 712 | 55 | 19 | -| [`src/attention.rs`](src/attention.rs) | 14,946 | 211 | 3 | +| [`src/attention.rs`](src/attention.rs) | 15,290 | 215 | 3 | | [`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 | @@ -38,7 +38,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,991 | 235 | 0 | | [`src/graph.rs`](src/graph.rs) | 556 | 33 | 0 | -| [`src/kernels.rs`](src/kernels.rs) | 8,088 | 1228 | 37 | +| [`src/kernels.rs`](src/kernels.rs) | 8,184 | 1231 | 38 | | [`src/kv_slots.rs`](src/kv_slots.rs) | 420 | 9 | 10 | | [`src/lib.rs`](src/lib.rs) | 88 | 26 | 1 | | [`src/moe.rs`](src/moe.rs) | 1,742 | 27 | 0 | @@ -57,7 +57,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`, `new`, `should_use_mmq`, `is_gfx906`, `is_gfx908`, `is_gfx1010`, `is_gfx1011`, `is_gfx1012`, `is_gfx1030`, `is_gfx1031`, `is_gfx1032`, `is_gfx1100`, +43 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`, +199 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`, +203 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` @@ -71,7 +71,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`, +1216 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`, +1219 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`, `embedding`, `feature_flags`, `flash_attn_ck`, `gemm`, `gemv`, `graph`, `kv_slots`, `moe`, `norm`, +14 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 @@ -100,6 +100,6 @@ _Generated by `scripts/check-crate-maps.py` from the tree — do not edit inside ### Totals -- 30 modules · 110,994 lines · 2757 public items · 225 tests · 192 examples +- 30 modules · 111,434 lines · 2764 public items · 226 tests · 192 examples diff --git a/crates/rdna-compute/src/attention.rs b/crates/rdna-compute/src/attention.rs index 9a2c1455a0..52f0118eaf 100644 --- a/crates/rdna-compute/src/attention.rs +++ b/crates/rdna-compute/src/attention.rs @@ -1846,6 +1846,350 @@ impl Gpu { result } + /// Flat BF16 KV write for single-token decode. Launched twice by the + /// caller (once for K, once for V), exactly like `kv_cache_write_q8_0`. + /// + /// The grid is a plain flat cover of `kv_dim` rather than Q8's + /// one-block-per-wave shape: bf16 has no per-block amax reduction, so + /// there is nothing to keep a wave together for. + pub fn kv_cache_write_bf16( + &mut self, + dst: &GpuTensor, + src: &GpuTensor, + pos_buf: &DeviceBuffer, + n_kv_heads: usize, + head_dim: usize, + ) -> HipResult<()> { + self.bind_thread()?; + // The decode and batched entry points share ONE translation unit, and + // that file `#include`s kv_slot_desc.h for the batched one. The JIT + // compiles in a cache dir with no -I to kernels/src, so this wrapper + // must strip-and-prepend exactly like the batched wrapper — even + // though the decode kernel uses nothing from the header. Omitting it + // here still compiles in the batched path and fails only on the first + // decode step, which is how it was found. + if !self.functions.contains_key("kv_cache_write_bf16") { + let stripped = + kernels::KV_CACHE_WRITE_BF16_SRC.replace("#include \"kv_slot_desc.h\"", ""); + let src = format!("{}\n{}", kernels::KV_SLOT_DESC_H, stripped); + self.ensure_kernel("kv_cache_write_bf16", &src, "kv_cache_write_bf16")?; + } + let d = dst.buf.as_ptr(); + let s = src.buf.as_ptr(); + let p = pos_buf.as_ptr(); + let nkv = n_kv_heads as i32; + let hd = head_dim as i32; + let mut params: Vec<*mut c_void> = vec![ + &d as *const _ as *mut c_void, + &s as *const _ as *mut c_void, + &p as *const _ as *mut c_void, + &nkv as *const _ as *mut c_void, + &hd as *const _ as *mut c_void, + ]; + let grid = (n_kv_heads * head_dim).div_ceil(64) as u32; + self.launch_maybe_blob( + "kv_cache_write_bf16", + [grid, 1, 1], + [64, 1, 1], + 0, + &mut params, + || { + let mut b = hip_bridge::KernargBlob::new(); + b.push_ptr(d); + b.push_ptr(s); + b.push_ptr(p); + b.push_i32(nkv); + b.push_i32(hd); + b + }, + ) + } + + /// Flat BF16 KV write for batched prefill. + /// + /// `slot_descs`/`row_slot` are both-or-neither for the same reason as the + /// Q8 sibling: passing only `slot_descs` pins every row to slot 0 and + /// writes every sequence's KV into slot 0's slab. + #[allow(clippy::too_many_arguments)] + pub fn kv_cache_write_bf16_batched( + &mut self, + dst: &GpuTensor, + src: &GpuTensor, + positions: &GpuTensor, + n_kv_heads: usize, + head_dim: usize, + batch_size: usize, + slot_descs: Option<&GpuTensor>, + row_slot: Option<&GpuTensor>, + ) -> HipResult<()> { + assert_eq!( + slot_descs.is_some(), + row_slot.is_some(), + "kv_cache_write_bf16_batched: slot_descs and row_slot are both-or-neither. \ + Passing only slot_descs silently pins every row to slot 0, writing every \ + sequence's KV into slot 0's slab." + ); + self.bind_thread()?; + // The kernel source `#include`s kv_slot_desc.h, but the runtime hipcc + // compile happens in a cache dir with no -I to kernels/src. Strip the + // directive and prepend the header body, same as the Q8 sibling. + // Guarded on the functions cache so the format!/replace runs once. + if !self.functions.contains_key("kv_cache_write_bf16_batched") { + let stripped = + kernels::KV_CACHE_WRITE_BF16_SRC.replace("#include \"kv_slot_desc.h\"", ""); + let src = format!("{}\n{}", kernels::KV_SLOT_DESC_H, stripped); + self.ensure_kernel( + "kv_cache_write_bf16_batched", + &src, + "kv_cache_write_bf16_batched", + )?; + } + let mut d = dst.buf.as_ptr(); + let mut s = src.buf.as_ptr(); + let mut p = positions.buf.as_ptr(); + let mut nkv = n_kv_heads as i32; + let mut hd = head_dim as i32; + let mut bs = batch_size as i32; + let mut desc_ptr: *mut std::ffi::c_void = match slot_descs { + Some(t) => t.buf.as_ptr(), + None => std::ptr::null_mut(), + }; + let mut rs_ptr: *mut std::ffi::c_void = match row_slot { + Some(t) => t.buf.as_ptr(), + None => std::ptr::null_mut(), + }; + let mut params: Vec<*mut c_void> = vec![ + &mut d as *mut _ as *mut c_void, + &mut s as *mut _ as *mut c_void, + &mut p as *mut _ as *mut c_void, + &mut nkv as *mut _ as *mut c_void, + &mut hd as *mut _ as *mut c_void, + &mut bs as *mut _ as *mut c_void, + &mut desc_ptr as *mut _ as *mut c_void, + &mut rs_ptr as *mut _ as *mut c_void, + ]; + let grid = (n_kv_heads * head_dim).div_ceil(64) as u32; + let desc_raw = desc_ptr; + let rs_raw = rs_ptr; + self.launch_maybe_blob( + "kv_cache_write_bf16_batched", + [grid, batch_size as u32, 1], + [64, 1, 1], + 0, + &mut params, + || { + let mut b = hip_bridge::KernargBlob::new(); + b.push_ptr(d); + b.push_ptr(s); + b.push_ptr(p); + b.push_i32(nkv); + b.push_i32(hd); + b.push_i32(bs); + b.push_ptr(desc_raw); + b.push_ptr(rs_raw); + b + }, + ) + } + + /// Batched sliding-window flash attention over flat BF16 KV (maple + /// prefill). Reuses the shared `launch_asym_flash_batched` dispatcher and + /// the shared batched reduce; only the tile kernel differs from Q8. + #[allow(clippy::too_many_arguments)] + pub fn attention_flash_bf16_batched_masked_windowed( + &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_seq: usize, + max_ctx_len: usize, + batch_size: usize, + partials: &GpuTensor, + tree_bias: Option<&GpuTensor>, + block_start: usize, + block_cols: usize, + window: i32, + ) -> HipResult<()> { + self.bind_thread()?; + self.launch_asym_flash_batched( + "attention_flash_bf16_tile_batched", + kernels::ATTENTION_FLASH_BF16_TILE_BATCHED_SRC, + "attention_flash_bf16_tile_batched", + q, + k_cache, + v_cache, + out, + positions, + q, // cos_theta dummy — kernel ignores + q, // sin_theta dummy — kernel ignores + n_heads, + n_kv_heads, + head_dim, + max_seq, + max_ctx_len, + batch_size, + partials, + tree_bias, + block_start, + block_cols, + // Consumed-but-unused by the bf16 tile (there is no separate V + // tier); V_MODE_Q8 keeps the kernarg blob shape identical to the + // Q8 path the shared launcher was written for. + V_MODE_Q8, + window, + /*force_wmma_grid=*/ false, + None, + None, + ) + } + + /// Sliding-window flash attention over flat BF16 KV — tile + reduce, the + /// decode sibling of `attention_flash_bf16_batched_masked_windowed`. + /// + /// The reduce is `attention_flash_q8_0_reduce` unchanged: it consumes only + /// f32 partials and never touches the KV cache, so it is KV-dtype-agnostic. + /// `window <= 0` means full causal — that is how Maple's global/NoPE + /// layers use this same kernel. + #[allow(clippy::too_many_arguments)] + pub fn attention_flash_bf16_windowed( + &mut self, + q: &GpuTensor, + k_cache: &GpuTensor, + v_cache: &GpuTensor, + out: &GpuTensor, + pos_buf: &DeviceBuffer, + seq_len_hint: usize, + n_heads: usize, + n_kv_heads: usize, + head_dim: usize, + max_seq: usize, + partials: &GpuTensor, + window: i32, + ) -> HipResult<()> { + self.bind_thread()?; + // Same tile-size policy as the Q8 path, so a partials buffer sized + // from max_tiles stays correct whichever tier the caller picked. + let tile_size = q8_flash_tile_size(&self.arch, n_heads, n_kv_heads, head_dim, max_seq); + let max_tiles = max_seq.div_ceil(tile_size); + let actual_tiles = seq_len_hint.div_ceil(tile_size); + // Graph/Redline-safe: capture the max_tiles superset so replay never + // needs a grid larger than the recorded one. The tile kernel + // early-exits for tiles beyond the live seq_len. + let launch_tiles = replay_stable_tile_count( + actual_tiles, + max_tiles, + self.graphs.capture_mode, + self.replay.is_recording(), + ); + + // ── Tile kernel ── + { + const KERNEL: &str = "attention_flash_bf16_tile"; + self.ensure_kernel(KERNEL, kernels::ATTENTION_FLASH_BF16_TILE_SRC, KERNEL)?; + let scale = 1.0f32 / (head_dim as f32).sqrt(); + let q_ptr = q.buf.as_ptr(); + 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 = pos_buf.as_ptr(); + let nh = n_heads as i32; + let nkv = n_kv_heads as i32; + let hd = head_dim as i32; + let ms = max_seq as i32; + let sc = scale; + let ts = tile_size as i32; + let wn = window; + let grid = [n_heads as u32, launch_tiles as u32, 1]; + let shared = ((tile_size + head_dim) * 4) as u32; + 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, + &ms as *const _ as *mut c_void, + &sc as *const _ as *mut c_void, + &ts as *const _ as *mut c_void, + &wn as *const _ as *mut c_void, + ]; + self.launch_maybe_blob_position_grid( + KERNEL, + grid, + [32, 1, 1], + shared, + &mut params, + 1, + 1, + tile_size as u32, + || { + 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_i32(ms); + b.push_f32(sc); + b.push_i32(ts); + b.push_i32(wn); + b + }, + )?; + } + + // ── Reduce kernel (shared with Q8; reads seq_len from pos_buf) ── + { + const KERNEL: &str = "attention_flash_q8_0_reduce"; + self.ensure_kernel(KERNEL, kernels::ATTENTION_FLASH_Q8_0_REDUCE_SRC, KERNEL)?; + let p_ptr = partials.buf.as_ptr(); + let o_ptr = out.buf.as_ptr(); + let nh = n_heads as i32; + let hd = head_dim as i32; + let pos_ptr = pos_buf.as_ptr(); + let ts = tile_size as i32; + let mt = max_tiles as i32; + let mut params: Vec<*mut c_void> = vec![ + &p_ptr as *const _ as *mut c_void, + &o_ptr as *const _ as *mut c_void, + &nh as *const _ as *mut c_void, + &hd as *const _ as *mut c_void, + &pos_ptr as *const _ as *mut c_void, + &ts as *const _ as *mut c_void, + &mt as *const _ as *mut c_void, + ]; + self.launch_maybe_blob( + KERNEL, + [n_heads as u32, 1, 1], + [256, 1, 1], + (max_tiles * 4) as u32, + &mut params, + || { + let mut b = hip_bridge::KernargBlob::new(); + b.push_ptr(p_ptr); + b.push_ptr(o_ptr); + b.push_i32(nh); + b.push_i32(hd); + b.push_ptr(pos_ptr); + b.push_i32(ts); + b.push_i32(mt); + b + }, + )?; + } + Ok(()) + } + /// Exact paired K/V Q8_0 cache write for single-token decode. Uses the /// same 32-lane block quantizer as `kv_cache_write_q8_0` and concatenates /// the independent K and V block grids into one dispatch. diff --git a/crates/rdna-compute/src/kernels.rs b/crates/rdna-compute/src/kernels.rs index 9dc44fe085..d1595d6e11 100644 --- a/crates/rdna-compute/src/kernels.rs +++ b/crates/rdna-compute/src/kernels.rs @@ -5219,6 +5219,13 @@ pub const KV_CACHE_WRITE_Q8_0_BATCHED_SRC: &str = pub const KV_CACHE_WRITE_Q8_0_SRC: &str = include_str!("../../../kernels/src/kv_cache_write_q8_0.hip"); +/// Flat BF16 KV write (maple). 2 bytes per element, no blocks and no scales. +/// Layout: [max_seq × n_kv_heads × head_dim] bf16. Holds both the decode +/// (`kv_cache_write_bf16`) and batched-prefill (`kv_cache_write_bf16_batched`) +/// entry points. +pub const KV_CACHE_WRITE_BF16_SRC: &str = + include_str!("../../../kernels/src/kv_cache_write_bf16.hip"); + /// gfx1100-only paired K/V Q8_0 cache writer. Kept in a separate translation /// unit so its dormant body cannot perturb portable/gfx12 writer codegen. pub const KV_CACHE_WRITE_Q8_0_PAIR_GFX1100_SRC: &str = @@ -5320,6 +5327,12 @@ pub const ATTENTION_Q8_0_KV_TIMED_SRC: &str = pub const ATTENTION_FLASH_Q8_0_TILE_SRC: &str = include_str!("../../../kernels/src/attention_flash_q8_0_tile.hip"); +/// Flat-BF16 sibling of the Q8_0 flash tile. Same partials layout and same +/// per-thread dim mapping, so it shares `attention_flash_q8_0_reduce` +/// unmodified — that reduce only ever touches f32 partials. +pub const ATTENTION_FLASH_BF16_TILE_SRC: &str = + include_str!("../../../kernels/src/attention_flash_bf16_tile.hip"); + /// gfx1151-only ISA experiment: preserve the flash tile's reduction tree but /// lower cross-lane exchanges to ds_swizzle + DPP8/quad-perm operations. pub const ATTENTION_FLASH_Q8_0_TILE_DPP_GFX1151_SRC: &str = concat!( @@ -5397,6 +5410,8 @@ pub const ATTENTION_FLASH_ASYM2_TILE_BATCHED_SRC: &str = include_str!("../../../kernels/src/attention_flash_asym2_tile_batched.hip"); 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_ASYM_REDUCE_BATCHED_SRC: &str = include_str!("../../../kernels/src/attention_flash_asym_reduce_batched.hip"); @@ -7186,6 +7201,87 @@ pub const CALIB_REDUCE_SRC: &str = include_str!("../../../kernels/src/calib_redu /// helpers feed a byte-exact decode route, so close numerical agreement is not /// sufficient: every produced coordinate and block scale must have identical /// bits to the generic implementation. +#[cfg(test)] +mod moe_topk_renorm_barriers { + use super::MOE_TOPK_RENORM_K8_SRC; + + /// Strip `//` and `/* */` comments and macro line-continuations, then + /// collapse whitespace, so the structural check below is not defeated by + /// reformatting or by a comment that happens to contain the tokens. + fn normalize(src: &str) -> String { + let mut out = String::with_capacity(src.len()); + let b = src.as_bytes(); + let mut i = 0; + while i < b.len() { + if b[i] == b'/' && i + 1 < b.len() && b[i + 1] == b'/' { + while i < b.len() && b[i] != b'\n' { + i += 1; + } + } else if b[i] == b'/' && i + 1 < b.len() && b[i + 1] == b'*' { + i += 2; + while i + 1 < b.len() && !(b[i] == b'*' && b[i + 1] == b'/') { + i += 1; + } + i = (i + 2).min(b.len()); + } else if b[i] == b'\\' { + i += 1; // macro line-continuation + } else { + out.push(b[i] as char); + i += 1; + } + } + out.split_whitespace().collect::>().join(" ") + } + + /// Every winner-invalidation must be followed by a barrier BEFORE the next + /// round restages `warp_v[warp_id]` / `warp_i[warp_id]`. + /// + /// The hazard this pins: `warp_i[0]` is dual-purpose in this kernel — it is + /// both warp 0's per-warp staging slot AND the slot through which a round's + /// winner is published to every thread. Without a barrier after the winner + /// is consumed, warp 0 races into the next round and overwrites `warp_i[0]` + /// while a slower warp is still reading it as the winner. That warp then + /// fails to invalidate the expert just picked, so the SAME expert is + /// selected again — the top-8 comes back with a duplicate and a distinct + /// expert missing. + /// + /// Not hypothetical: measured at ~1 occurrence per 46k router calls on Maple + /// (arch 15) from bit-identical input logits. It is a silent quality fault — + /// the duplicated expert is double-weighted and a legitimately selected one + /// never runs. The sibling kernels (`moe_softmax_topk_k8`, both `_batched` + /// variants) are immune because they mask through a separate + /// `smem_v`/`picked_idx` instead of aliasing the staging array. + #[test] + fn invalidation_is_closed_by_a_barrier_before_the_next_restage() { + let norm = normalize(MOE_TOPK_RENORM_K8_SRC); + let sites: Vec = norm.match_indices("cur_i = -1;").map(|(i, _)| i).collect(); + assert_eq!( + sites.len(), + 2, + "expected exactly two winner-invalidation sites (the exact-256 fast \ + path and the n_exp<=BLOCK_SIZE path); found {}. If the kernel was \ + restructured, re-derive this invariant rather than deleting it.", + sites.len() + ); + for site in sites { + let rest = &norm[site..]; + let barrier = rest + .find("__syncthreads();") + .expect("no barrier at all follows a winner-invalidation"); + if let Some(restage) = rest.find("warp_v[warp_id] =") { + assert!( + barrier < restage, + "winner-invalidation is not barrier-closed: the next \ + `warp_v[warp_id] =` restage is reachable before any \ + `__syncthreads()`. warp 0 can clobber `warp_i[0]` while \ + another warp still reads it as this round's winner, which \ + yields a DUPLICATED expert in the top-8." + ); + } + } + } +} + #[cfg(test)] mod gfx1201_e8_decode_identity { fn generic_cvt_e4m3(b: u8) -> f32 { diff --git a/crates/saddle-core/map.md b/crates/saddle-core/map.md index fb3217536c..bf2570a046 100644 --- a/crates/saddle-core/map.md +++ b/crates/saddle-core/map.md @@ -37,7 +37,7 @@ _Generated by `scripts/check-crate-maps.py` from the tree — do not edit inside |---|---:|---:|---:| | [`src/caps.rs`](src/caps.rs) | 286 | 11 | 3 | | [`src/grammar.rs`](src/grammar.rs) | 3,977 | 20 | 106 | -| [`src/kv.rs`](src/kv.rs) | 4,308 | 78 | 9 | +| [`src/kv.rs`](src/kv.rs) | 4,661 | 80 | 13 | | [`src/lib.rs`](src/lib.rs) | 76 | 6 | 0 | | [`src/logprobs.rs`](src/logprobs.rs) | 191 | 3 | 8 | | [`src/sampling.rs`](src/sampling.rs) | 52 | 2 | 0 | @@ -47,7 +47,7 @@ _Generated by `scripts/check-crate-maps.py` from the tree — do not edit inside - [`src/caps.rs`](src/caps.rs): `DflashKind`, `ReasoningContract`, `fn`, `from_wire_name`, `ArchCaps`, `supports_dflash`, `is_qwen_dflash`, `is_llama_dflash`, `supports_semantic_v2`, `qwen_semantic_v2`, `BatchEligibilityRequest` - [`src/grammar.rs`](src/grammar.rs): `json`, `State`, `ToolSchema`, `Config`, `Matcher`, `new`, `with_config`, `config`, `current_tool`, `debug_close_reject`, `attractor_detected`, `state`, +8 more -- [`src/kv.rs`](src/kv.rs): `KvMode`, `KvBackend`, `fn`, `ParseKvBackendError`, `KV_BACKEND_NAMES`, `DEFAULT_KV_CHUNK_TOKENS`, `DEFAULT_VMM_PHYSICAL_CHUNK_BYTES`, `KvMapGrowth`, `KvChunkPlan`, `KvChunkPlanError`, `new`, `mapped_bytes_for_tokens`, +66 more +- [`src/kv.rs`](src/kv.rs): `KvMode`, `KvBackend`, `fn`, `ParseKvBackendError`, `KV_BACKEND_NAMES`, `DEFAULT_KV_CHUNK_TOKENS`, `DEFAULT_VMM_PHYSICAL_CHUNK_BYTES`, `KvMapGrowth`, `KvChunkPlan`, `KvChunkPlanError`, `new`, `mapped_bytes_for_tokens`, +68 more - [`src/lib.rs`](src/lib.rs): `grammar`, `kv`, `caps`, `logprobs`, `sampling`, `spec` - [`src/logprobs.rs`](src/logprobs.rs): `TokenLogprob`, `top_k_logprobs`, `logprob_of` - [`src/sampling.rs`](src/sampling.rs): `SamplingDefaults`, `fn` @@ -66,6 +66,6 @@ _Generated by `scripts/check-crate-maps.py` from the tree — do not edit inside ### Totals -- 7 modules · 9,213 lines · 122 public items · 127 tests · 0 examples +- 7 modules · 9,566 lines · 124 public items · 131 tests · 0 examples diff --git a/crates/saddle-core/src/kv.rs b/crates/saddle-core/src/kv.rs index 00185e8050..4bee728e83 100644 --- a/crates/saddle-core/src/kv.rs +++ b/crates/saddle-core/src/kv.rs @@ -14,6 +14,11 @@ use rdna_compute::{DType, Gpu, GpuTensor}; #[derive(Clone, Copy, Debug, PartialEq, Eq)] pub enum KvMode { Q8, + /// Flat 2-byte BF16 K/V. NOT part of the quantized ladder: no rotation, no + /// per-block scale, and no VMM / adaptive / compaction support. Only a + /// site whose `accepted` list names it can ever resolve to it — today that + /// is maple alone, so every other site's behaviour is unchanged. + Bf16, Asym2, Asym3, Asym4, @@ -311,6 +316,11 @@ pub struct KvCache { /// True when the rotation primitive is signed-FWHT (matches Fwht{2,3,4} /// KvMode values). False when Givens (matches Asym{2,3,4}). pub quant_fwht: bool, + /// True when K and V are stored as flat 2-byte BF16 (no scales, no + /// blocks) instead of a quantized block layout. Mutually exclusive with + /// every `quant_*` tier flag above; `quantized` is also set so the legacy + /// llama/qwen35 `!quantized` branches never mistake it for plain F32. + pub quant_bf16: bool, /// V-cache quantization mode (independent of the K mode). Defaults to Q8. pub v_mode: VMode, /// Per-layer flag: true = this layer uses Q8 (boundary layer) @@ -424,6 +434,13 @@ impl KvCache { } Self::checked_vmm_product("q8 K head stride", &[head_dim / 32, 34]) } + // BF16 is contiguous-only. It has no growable-arena constructor, so + // refuse here rather than compute a stride for a layout the VMM + // path cannot actually allocate. + KvMode::Bf16 => Err(hip_bridge::HipError::new( + 0, + "VMM does not support bf16 KV (contiguous backend only)", + )), KvMode::Asym2 | KvMode::Fwht2 => head_dim .checked_div(4) .and_then(|n| n.checked_add(4)) @@ -504,6 +521,14 @@ impl KvCache { )); } } + // Fail closed: bf16 has no VMM constructor. Callers that want bf16 + // must use the contiguous backend. + KvMode::Bf16 => { + return Err(hip_bridge::HipError::new( + 0, + "VMM does not support bf16 KV (contiguous backend only)", + )); + } KvMode::Asym2 | KvMode::Asym3 | KvMode::Asym4 => { let ok_hd = match mode { KvMode::Asym3 => head_dim == 256, @@ -597,6 +622,10 @@ impl KvCache { Self::checked_vmm_product("V reserve", &[physical_cap, v_bytes_per_token])?; let rotation_table_len = match mode { KvMode::Q8 => 0, + // Unrotated, like Q8. Unreachable in practice — the validate above + // rejects bf16 for VMM before this runs — but 0 is the honest + // answer for a tier with no rotation table. + KvMode::Bf16 => 0, KvMode::Asym2 | KvMode::Asym3 | KvMode::Asym4 => head_dim / 2, KvMode::Fwht3 => 256, KvMode::Fwht2 | KvMode::Fwht4 => { @@ -765,6 +794,15 @@ impl KvCache { KvMode::Fwht3 => (false, false, true, false, true), KvMode::Fwht4 => (false, true, false, false, true), KvMode::Asym3Auto => (false, false, false, false, false), + // Bf16 is NOT representable in this 5-flag VMM bundle — all-false + // here would decode as KTier::F32 and hand a bf16 buffer to the + // F32 kernels, which read it at twice the stride. It can never + // legitimately arrive: `validate_vmm_mode` rejects bf16 before any + // VMM constructor runs. Panic loudly rather than return a lie. + KvMode::Bf16 => panic!( + "vmm_mode_flags: bf16 has no VMM layout — it is contiguous-only \ + and should have been rejected by validate_mode_with_backend" + ), } } @@ -974,6 +1012,7 @@ impl KvCache { givens_cos: None, givens_sin: None, quant_fwht: false, + quant_bf16: false, v_mode: VMode::Q8, layer_is_boundary: self.layer_is_boundary.clone(), compact_offset: 0, @@ -1086,6 +1125,14 @@ impl KvCache { (KvMode::Q8, Flat(n), None) => Self::new_gpu_q8(gpu, *n, nh, hd, ms), (KvMode::Asym3, Flat(n), None) => Self::new_gpu_asym3(gpu, *n, nh, hd, ms), (KvMode::Asym4, Flat(n), None) => Self::new_gpu_asym4(gpu, *n, nh, hd, ms), + // Bf16 is Flat-only: there is no _filtered constructor because no + // hybrid arch (the reason _filtered exists) uses this tier. A + // Mask request therefore falls through to the error below rather + // than silently allocating every layer. + (KvMode::Bf16, Flat(n), Some(cap)) => { + Self::new_gpu_bf16_capped(gpu, *n, nh, hd, ms, cap) + } + (KvMode::Bf16, Flat(n), None) => Self::new_gpu_bf16(gpu, *n, nh, hd, ms), // No constructor exists for this combination. (m, l, c) => Err(hip_bridge::HipError::new( 0, @@ -1100,7 +1147,6 @@ impl KvCache { )), } } - } impl KvCache { @@ -1137,6 +1183,7 @@ impl KvCache { quant_asym3: false, quant_asym2: false, quant_fwht: false, + quant_bf16: false, boundary_layers: 0, givens_cos: None, givens_sin: None, @@ -1185,6 +1232,7 @@ impl KvCache { quant_asym3: false, quant_asym2: false, quant_fwht: false, + quant_bf16: false, boundary_layers: 0, givens_cos: None, givens_sin: None, @@ -1256,6 +1304,7 @@ impl KvCache { quant_asym3: false, quant_asym2: false, quant_fwht: false, + quant_bf16: false, boundary_layers: 0, givens_cos: None, givens_sin: None, @@ -1265,6 +1314,100 @@ impl KvCache { }) } + /// Create a flat BF16 KV cache. 2 bytes per element — 1.88x the Q8_0 + /// layout (34 B per 32 elements) but with no per-block scale and no + /// quantization error: bf16 carries the same 8 exponent bits as f32 and + /// truncates only the mantissa. + /// + /// Layout is deliberately the simplest thing that can work: element + /// `(t, kv_h, d)` lives at `t * kv_dim + kv_h * head_dim + d`, one bf16 + /// each. No blocks, no scales, no padding. That is what lets the tile + /// kernel drop the entire Q8 block-index computation. + /// + /// Sized by `physical_cap` like `new_gpu_q8_capped`, so eviction-bounded + /// callers get the buffer they asked for. + /// + /// This exists so Maple's Q8 KV can be compared against a near-reference + /// KV at long context. It is NOT the default for any model. + pub fn new_gpu_bf16( + gpu: &mut Gpu, + n_layers: usize, + n_kv_heads: usize, + head_dim: usize, + max_seq_len: usize, + ) -> HipResult { + Self::new_gpu_bf16_capped( + gpu, + n_layers, + n_kv_heads, + head_dim, + max_seq_len, + max_seq_len, + ) + } + + /// Same as [`KvCache::new_gpu_bf16`] with an explicit physical_cap. + pub fn new_gpu_bf16_capped( + gpu: &mut Gpu, + n_layers: usize, + n_kv_heads: usize, + head_dim: usize, + max_seq_len: usize, + physical_cap: usize, + ) -> HipResult { + assert!( + physical_cap > 0 && physical_cap <= max_seq_len, + "physical_cap ({physical_cap}) must be in (0, max_seq_len={max_seq_len}]" + ); + let kv_dim = n_kv_heads * head_dim; + // 2 bytes per element, rounded up to whole F32 elements because the + // allocator is typed F32 everywhere else in this file. kv_dim is even + // for every real model so the round-up is a no-op, but the ceil keeps + // a hypothetical odd kv_dim from under-allocating. + let cache_bytes = physical_cap * kv_dim * 2; + let cache_elems = cache_bytes.div_ceil(4); + let mut k_gpu = Vec::with_capacity(n_layers); + let mut v_gpu = Vec::with_capacity(n_layers); + for _ in 0..n_layers { + k_gpu.push(gpu.zeros(&[cache_elems], DType::F32)?); + v_gpu.push(gpu.zeros(&[cache_elems], DType::F32)?); + } + Ok(Self { + k_gpu, + v_gpu, + k_scales: vec![], + v_scales: vec![], + kv_dim, + max_seq: max_seq_len, + physical_cap, + n_kv_heads, + head_dim, + // `quantized` is TRUE even though bf16 is not a quantized tier: + // the legacy llama/qwen35 paths branch on `!quantized` to mean + // "plain F32 layout", and a bf16 buffer read as F32 is garbage. + // Setting this keeps those paths out of their F32 arm. Only Maple + // can allocate this cache today. + quantized: true, + quant_q8: false, + quant_int8: false, + quant_hfq4: false, + quant_asym4: false, + quant_asym3: false, + quant_asym2: false, + quant_fwht: false, + quant_bf16: true, + boundary_layers: 0, + givens_cos: None, + givens_sin: None, + layer_is_boundary: vec![], + compact_offset: 0, + // V is bf16 too. `VMode::Q8` is the struct's default and is never + // read on this path — the tier decode reaches `KTier::Bf16` before + // any v_mode branch. + v_mode: VMode::Q8, + }) + } + /// Helper: allocate K/V Vecs, skipping layers where is_kv_layer[i] is false /// by inserting a 1-element placeholder. Saves VRAM for hybrid arches /// (Qwen 3.5 DeltaNet + FullAttention) where 75% of layers don't carry @@ -1641,6 +1784,10 @@ impl KvCache { // invariant (the legacy qwen35 literals hardcoded quant_q4 = false). // Release classify() output is unchanged either way (asym is matched // before q4), so this is a true no-op for kernel selection. + // BF16 is a distinct flat tier (maple): it is `quantized:true` with + // empty `k_scales` and no other quant flag, so without `!quant_bf16` + // it would ALSO report quant_q4, giving two true tier flags and + // tripping the same debug_assert on every Maple BF16 dispatch. self.quantized && !self.quant_hfq4 && !self.quant_q8 @@ -1648,6 +1795,7 @@ impl KvCache { && !self.quant_asym4 && !self.quant_asym3 && !self.quant_asym2 + && !self.quant_bf16 && self.k_scales.is_empty() } @@ -2356,6 +2504,7 @@ impl KvCache { quant_asym3: false, quant_asym2: false, quant_fwht: true, + quant_bf16: false, boundary_layers: 0, givens_cos: Some(s1), givens_sin: Some(s2), @@ -2451,6 +2600,7 @@ impl KvCache { quant_asym3: false, quant_asym2: false, quant_fwht: false, + quant_bf16: false, boundary_layers: 0, givens_cos: None, givens_sin: None, @@ -2567,6 +2717,7 @@ impl KvCache { quant_asym3, quant_asym2, quant_fwht, + quant_bf16: false, boundary_layers: 0, givens_cos, givens_sin, @@ -2661,6 +2812,7 @@ impl KvCache { quant_asym3: false, quant_asym2: false, quant_fwht: false, + quant_bf16: false, boundary_layers: 0, givens_cos: None, givens_sin: None, @@ -2707,6 +2859,7 @@ impl KvCache { quant_asym3: false, quant_asym2: false, quant_fwht: false, + quant_bf16: false, boundary_layers: 0, givens_cos: None, givens_sin: None, @@ -2755,6 +2908,7 @@ impl KvCache { quant_asym3: false, quant_asym2: false, quant_fwht: false, + quant_bf16: false, boundary_layers: 0, givens_cos: None, givens_sin: None, @@ -2805,6 +2959,7 @@ impl KvCache { quant_asym3: false, quant_asym2: false, quant_fwht: false, + quant_bf16: false, boundary_layers: 0, givens_cos: None, givens_sin: None, @@ -2903,6 +3058,7 @@ impl KvCache { quant_asym3: false, quant_asym2: false, quant_fwht: false, + quant_bf16: false, boundary_layers: 0, givens_cos: Some(ct), givens_sin: Some(st), @@ -2974,6 +3130,7 @@ impl KvCache { quant_asym3: false, quant_asym2: false, quant_fwht: true, + quant_bf16: false, boundary_layers: 0, givens_cos: Some(s1), givens_sin: Some(s2), @@ -3046,6 +3203,7 @@ impl KvCache { quant_asym3: false, quant_asym2: false, quant_fwht: false, + quant_bf16: false, boundary_layers: 0, givens_cos: Some(ct), givens_sin: Some(st), @@ -3145,6 +3303,7 @@ impl KvCache { quant_asym3: false, quant_asym2: false, quant_fwht: true, + quant_bf16: false, boundary_layers: 0, givens_cos: Some(s1), givens_sin: Some(s2), @@ -3271,6 +3430,7 @@ impl KvCache { quant_asym3: true, quant_asym2: false, quant_fwht: false, + quant_bf16: false, boundary_layers: 0, givens_cos: Some(ct), givens_sin: Some(st), @@ -3453,6 +3613,7 @@ impl KvCache { quant_asym3: true, quant_asym2: false, quant_fwht: true, + quant_bf16: false, boundary_layers: 0, givens_cos: Some(s1), givens_sin: Some(s2), @@ -3480,7 +3641,12 @@ impl KvCache { "asym3 currently requires head_dim=256 (Qwen 3.5)" ); Self::new_gpu_asym3_capped_inner( - gpu, n_layers, n_kv_heads, head_dim, max_seq_len, physical_cap, + gpu, + n_layers, + n_kv_heads, + head_dim, + max_seq_len, + physical_cap, ) } @@ -3500,7 +3666,12 @@ impl KvCache { "asym3 (gemma4) requires head_dim=256 or 512 (got {head_dim})" ); Self::new_gpu_asym3_capped_inner( - gpu, n_layers, n_kv_heads, head_dim, max_seq_len, physical_cap, + gpu, + n_layers, + n_kv_heads, + head_dim, + max_seq_len, + physical_cap, ) } @@ -3559,6 +3730,7 @@ impl KvCache { quant_asym3: true, quant_asym2: false, quant_fwht: false, + quant_bf16: false, boundary_layers: 0, givens_cos: Some(ct), givens_sin: Some(st), @@ -3642,6 +3814,7 @@ impl KvCache { quant_asym3: false, quant_asym2: true, quant_fwht: false, + quant_bf16: false, boundary_layers: 0, givens_cos: Some(ct), givens_sin: Some(st), @@ -3731,6 +3904,7 @@ impl KvCache { quant_asym3: false, quant_asym2: true, quant_fwht: true, + quant_bf16: false, boundary_layers: 0, givens_cos: Some(s1), givens_sin: Some(s2), @@ -3803,6 +3977,7 @@ impl KvCache { quant_asym3: false, quant_asym2: true, quant_fwht: false, + quant_bf16: false, boundary_layers: 0, givens_cos: Some(ct), givens_sin: Some(st), @@ -3908,7 +4083,6 @@ impl KvCache { // The KvCache.givens_cos / .givens_sin fields stay `None` in multi mode // — Stage 6 forward dispatch reads from the per-device replicas in // `Gpus` instead. - } /// KV VMM-layout and adaptive-reset contract tests. @@ -3930,6 +4104,7 @@ mod vmm_layout_tests { KvMode::Asym3 | KvMode::Fwht3 => 4 + (head_dim * 3) / 8, KvMode::Asym4 | KvMode::Fwht4 => 4 + head_dim / 2, KvMode::Asym3Auto => panic!("Asym3Auto is not a layout mode"), + KvMode::Bf16 => panic!("bf16 is not a VMM layout mode"), } } @@ -3962,6 +4137,7 @@ mod vmm_layout_tests { quant_asym3: a3, quant_asym2: a2, quant_fwht: fwht, + quant_bf16: false, boundary_layers: 0, givens_cos: None, givens_sin: None, @@ -3986,7 +4162,6 @@ mod vmm_layout_tests { } } - #[test] fn fwht3_vmm_layout_matches_asym3_byte_geometry() { let n_kv_heads = 4; @@ -4306,3 +4481,181 @@ mod vmm_layout_tests { } } } + +/// BF16 tier projection / plan regression (debug-build). +/// +/// Maple's BF16 cache sets `quantized=true` with empty `k_scales` and +/// `quant_bf16=true` — the exact shape that the legacy Q4 residual +/// `quantized && !tier && k_scales.is_empty()` would also match. +/// Without the `!quant_bf16` exclusion the cache reports TWO true tier flags +/// (`bf16` + `q4`), which trips `hipfire-dispatch::families::kv_tier::classify`'s +/// `debug_assert!(count <= 1)` on every Maple BF16 attention dispatch. +/// These tests mirror that assertion without taking a dispatch dependency, so +/// a future regression is caught here even in GPU-free CI. +#[cfg(test)] +mod bf16_tier_projection_tests { + use super::*; + + fn stub( + quantized: bool, + quant_q8: bool, + quant_hfq4: bool, + quant_int8: bool, + quant_asym4: bool, + quant_asym3: bool, + quant_asym2: bool, + quant_fwht: bool, + quant_bf16: bool, + k_scales_empty: bool, + ) -> KvCache { + KvCache { + k_gpu: Vec::new(), + v_gpu: Vec::new(), + k_scales: if k_scales_empty { + Vec::new() + } else { + // non-empty sentinel: one dummy 1-element tensor avoids needing Gpu + vec![GpuTensor { + buf: unsafe { hip_bridge::DeviceBuffer::from_raw(std::ptr::null_mut(), 0) }, + shape: vec![1], + dtype: DType::F32, + }] + }, + v_scales: Vec::new(), + kv_dim: 0, + max_seq: 0, + physical_cap: 0, + n_kv_heads: 0, + head_dim: 0, + quantized, + quant_q8, + quant_int8, + quant_hfq4, + quant_asym4, + quant_asym3, + quant_asym2, + quant_fwht, + quant_bf16, + v_mode: VMode::Q8, + boundary_layers: 0, + givens_cos: None, + givens_sin: None, + layer_is_boundary: Vec::new(), + compact_offset: 0, + } + } + + /// Mirrors `hipfire_dispatch::families::kv_tier::classify`'s at-most-one check, + /// using the two derived predicates exactly as `KvCacheExt::{k_tier,tier_inputs}` + /// do: `quant_q4 = quant_q4_residual()`, `quant_hfq8 = is_hfq8_kv()`. + fn tier_count(cache: &KvCache) -> usize { + let flags = [ + cache.quant_asym4, + cache.quant_asym3, + cache.quant_asym2, + cache.quant_q8, + cache.quant_hfq4, + cache.quant_q4_residual(), + cache.quant_int8, + cache.is_hfq8_kv(), + cache.quant_bf16, + ]; + flags.iter().filter(|&&b| b).count() + } + + #[test] + fn bf16_projection_is_exclusive_q4_false_bf16_true() { + // Maple BF16: quantized true, empty scales, bf16 true, no other tier. + let bf16 = stub( + true, false, false, false, false, false, false, false, true, true, + ); + assert!( + !bf16.quant_q4_residual(), + "BF16 must NOT report legacy Q4 residual" + ); + assert!(bf16.quant_bf16, "BF16 flag must be true"); + assert!(!bf16.is_hfq8_kv(), "BF16 must not report HFQ8"); + assert_eq!( + tier_count(&bf16), + 1, + "BF16 must classify as exactly one tier (bf16)" + ); + // The dispatch crate's `classify` debug_assert would trip if count > 1. + // Mirror that guard so GPU-free CI catches the same regression. + debug_assert!( + tier_count(&bf16) <= 1, + "at most one KV storage tier flag should be set (BF16)" + ); + } + + #[test] + fn q4_residual_still_reports_q4_only() { + // LLaMA legacy Q4: quantized true, empty scales, no named tier, no bf16. + let q4 = stub( + true, false, false, false, false, false, false, false, false, true, + ); + assert!(q4.quant_q4_residual(), "legacy Q4 must report q4 residual"); + assert!(!q4.quant_bf16); + assert!(!q4.is_hfq8_kv()); + assert_eq!( + tier_count(&q4), + 1, + "Q4 must classify as exactly one tier (q4)" + ); + debug_assert!(tier_count(&q4) <= 1, "at most one tier (Q4)"); + } + + #[test] + fn q8_projection_is_exclusive_q4_false() { + // Q8: quantized true, q8 true, empty scales, no bf16. Must not also be Q4. + let q8 = stub( + true, true, false, false, false, false, false, false, false, true, + ); + assert!( + !q8.quant_q4_residual(), + "Q8 must NOT report legacy Q4 residual" + ); + assert!(!q8.quant_bf16); + assert!(!q8.is_hfq8_kv()); + assert!(q8.quant_q8); + assert_eq!( + tier_count(&q8), + 1, + "Q8 must classify as exactly one tier (q8)" + ); + debug_assert!(tier_count(&q8) <= 1, "at most one tier (Q8)"); + } + + #[test] + fn asym_and_hfq_variants_remain_exclusive() { + // Asym3 (qwen35 default) was the original motivator for the asym exclusion; + // ensure the new bf16 exclusion didn't reintroduce overlap. + let asym3 = stub( + true, false, false, false, false, true, false, false, false, true, + ); + assert!( + !asym3.quant_q4_residual(), + "asym3 must NOT report Q4 residual" + ); + assert_eq!(tier_count(&asym3), 1); + + // HFQ8 has non-empty k_scales and is its own tier. + let hfq8 = stub( + true, false, false, false, false, false, false, false, false, false, + ); + assert!(hfq8.is_hfq8_kv(), "hfq8 with scales must report hfq8"); + assert!( + !hfq8.quant_q4_residual(), + "hfq8 must NOT report Q4 (scales non-empty)" + ); + assert_eq!(tier_count(&hfq8), 1); + + // F32: quantized false => no tier at all (KTier::F32). + let f32_cache = stub( + false, false, false, false, false, false, false, false, false, true, + ); + assert!(!f32_cache.quant_q4_residual()); + assert!(!f32_cache.is_hfq8_kv()); + assert_eq!(tier_count(&f32_cache), 0, "F32 must classify as zero tiers"); + } +} diff --git a/kernels/src/attention_flash_bf16_tile.hip b/kernels/src/attention_flash_bf16_tile.hip new file mode 100644 index 0000000000..48077b2bd9 --- /dev/null +++ b/kernels/src/attention_flash_bf16_tile.hip @@ -0,0 +1,226 @@ +// SPDX-License-Identifier: Apache-2.0 +// Copyright (c) 2026 Nick Woolmer +// hipfire — see LICENSE and NOTICE in the project root. +// +// Flash attention tile kernel for FLAT BF16 KV, with a sliding-window mask. +// Decode / single-position sibling of `attention_flash_q8_0_tile.hip`. +// +// This is a direct port of that kernel with the Q8_0 block addressing removed: +// where Q8 reads a 34-byte block (fp16 scale + 32 int8) and multiplies by the +// scale, this reads 2-byte bf16 values straight out of a flat +// `[pos][kv_head][d]` array and widens them by a 16-bit shift. +// +// WHAT IS DELIBERATELY IDENTICAL to the Q8 kernel, and must stay so: +// * the per-thread `d` mapping in Phase A (`k_bi = half*4 + tid/8`, +// `k_off = (tid%8)*4`) and in Phase D (`d_base = half*128 + tid*4`); +// * the order of the FMA accumulation and the descending XOR reduction; +// * the partials layout `[n_heads][max_tiles][2 + head_dim]` and the +// `max_tiles` (NOT n_tiles) row stride. +// The first two keep a bf16-vs-q8 comparison a measurement of the STORAGE +// tier rather than of a different summation order. The third is what lets +// this share `attention_flash_q8_0_reduce` unmodified — that kernel only ever +// touches f32 partials, so it is KV-dtype-agnostic and there is no bf16 +// reduce kernel. +// +// The gfx1151 DPP reduction variant from the Q8 kernel is NOT ported. It is an +// opt-in micro-optimisation behind HIPFIRE_GFX1151_ATTENTION_TILE_DPP and is +// wired only to the Q8 module name; adding it here would be an unmeasured +// second code path on a tier whose whole purpose is to be a reference. +// +// Grid: [n_heads, n_tiles, 1]. Block: [32, 1, 1] (one WAVE32). +// LDS: tile_size floats for scores + head_dim floats for Q. +#include + +// Widen one bf16 (as raw bits) to f32. BF16 is exactly the top 16 bits of an +// f32, so this is a shift — no table, no rounding, exact. +static __device__ __forceinline__ float hipfire_bf16_to_f32(unsigned short b) { + return __builtin_bit_cast(float, ((unsigned int)b) << 16); +} + +extern "C" __launch_bounds__(32, 16) +__global__ void attention_flash_bf16_tile( + const float* __restrict__ q, // [n_heads, head_dim] + const unsigned char* __restrict__ k_cache, + const unsigned char* __restrict__ v_cache, + float* __restrict__ partials, // [n_heads, max_tiles, 2 + head_dim] + const int* __restrict__ pos_buf, // pos_buf[0] = last position + int n_heads, + int n_kv_heads, + int head_dim, + int max_seq, + float scale_attn, + int tile_size, + int window // sliding-window span; <= 0 = full causal +) { + const int seq_len = pos_buf[0] + 1; + // Sliding-window lower bound: keys at t < win_lo are outside the window + // [seq_len-window, seq_len) and masked. window <= 0 = full causal. + const int win_lo = (window > 0) ? (seq_len - window) : 0; + const int h = blockIdx.x; + if (h >= n_heads) return; + const int tile_id = blockIdx.y; + const int tile_start = tile_id * tile_size; + const int tile_end = min(tile_start + tile_size, seq_len); + if (tile_start >= seq_len) return; + 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 tid = threadIdx.x; + + // Flat bf16: one element per (pos, kv_head, d), stride kv_dim per position. + const int kv_dim = n_kv_heads * head_dim; + const unsigned short* k16 = (const unsigned short*)k_cache; + const unsigned short* v16 = (const unsigned short*)v_cache; + const int head_base = kv_h * head_dim; + + // Number of 128-element halves spanning head_dim. One pass per half in + // Phase A (Q·K) and Phase D (V) so a 32-thread wave covers the full + // head_dim regardless of size. + const int n_halves = (head_dim + 127) / 128; + + // LDS layout: [tile_size scores][head_dim Q values] + extern __shared__ float sdata[]; + float* scores = sdata; + float* q_lds = sdata + tile_size; + + // Load Q into LDS. 32 threads × 4 values × n_halves = head_dim values. + const float* q_head = q + h * head_dim; + for (int half = 0; half < n_halves; half++) { + const int d_base = half * 128 + tid * 4; + q_lds[d_base + 0] = q_head[d_base + 0]; + q_lds[d_base + 1] = q_head[d_base + 1]; + q_lds[d_base + 2] = q_head[d_base + 2]; + q_lds[d_base + 3] = q_head[d_base + 3]; + } + __syncthreads(); + + // ═══ Phase A: compute tile_len dot products (wave-cooperative) ═══ + // Thread assignment mirrors the Q8 kernel exactly: within a 128-element + // half, lane `tid` owns the 4 dims at `(tid/8)*32 + (tid%8)*4`. In Q8 + // terms that is block `bi = tid/8`, byte offset `(tid%8)*4` inside it; + // here it is just a flat dim offset, but the mapping is preserved so the + // two kernels sum in the same order. + const int k_off = (tid % 8) * 4; + + int valid_start = win_lo - tile_start; + if (valid_start < 0) valid_start = 0; + if (valid_start > tile_len) valid_start = tile_len; + for (int t_local = 0; t_local < valid_start; t_local++) { + if (tid == 0) scores[t_local] = -INFINITY; + } + + int t_local = valid_start; + for (; t_local + 3 < tile_len; t_local += 4) { + const int t0 = tile_start + t_local + 0; + const int t1 = tile_start + t_local + 1; + const int t2 = tile_start + t_local + 2; + const int t3 = tile_start + t_local + 3; + float partial0 = 0.0f; + float partial1 = 0.0f; + float partial2 = 0.0f; + float partial3 = 0.0f; + for (int half = 0; half < n_halves; half++) { + const int d_start = (half * 4 + tid / 8) * 32 + k_off; + const int elem = head_base + d_start; + const unsigned short* kb0 = k16 + (size_t)t0 * kv_dim + elem; + const unsigned short* kb1 = k16 + (size_t)t1 * kv_dim + elem; + const unsigned short* kb2 = k16 + (size_t)t2 * kv_dim + elem; + const unsigned short* kb3 = k16 + (size_t)t3 * kv_dim + elem; + const float* qb_thread = q_lds + d_start; + #pragma unroll + for (int i = 0; i < 4; i++) { + partial0 += qb_thread[i] * hipfire_bf16_to_f32(kb0[i]); + partial1 += qb_thread[i] * hipfire_bf16_to_f32(kb1[i]); + partial2 += qb_thread[i] * hipfire_bf16_to_f32(kb2[i]); + partial3 += qb_thread[i] * hipfire_bf16_to_f32(kb3[i]); + } + } + for (int off = 16; off > 0; off >>= 1) { + partial0 += __shfl_xor(partial0, off); + partial1 += __shfl_xor(partial1, off); + partial2 += __shfl_xor(partial2, off); + partial3 += __shfl_xor(partial3, off); + } + if (tid == 0) { + scores[t_local + 0] = partial0 * scale_attn; + scores[t_local + 1] = partial1 * scale_attn; + scores[t_local + 2] = partial2 * scale_attn; + scores[t_local + 3] = partial3 * scale_attn; + } + } + for (; t_local < tile_len; t_local++) { + const int t = tile_start + t_local; + float partial = 0.0f; + for (int half = 0; half < n_halves; half++) { + const int d_start = (half * 4 + tid / 8) * 32 + k_off; + const unsigned short* kb = k16 + (size_t)t * kv_dim + head_base + d_start; + const float* qb_thread = q_lds + d_start; + #pragma unroll + for (int i = 0; i < 4; i++) { + partial += qb_thread[i] * hipfire_bf16_to_f32(kb[i]); + } + } + for (int off = 16; off > 0; off >>= 1) + partial += __shfl_xor(partial, off); + if (tid == 0) + scores[t_local] = partial * scale_attn; + } + __syncthreads(); + + // ═══ Phase B: find tile max ═══ + float local_max = -1e30f; + for (int i = tid; i < tile_len; i += 32) + local_max = fmaxf(local_max, scores[i]); + for (int off = 16; off > 0; off >>= 1) + local_max = fmaxf(local_max, __shfl_xor(local_max, off)); + const float tile_max = local_max; + + // ═══ Phase C: exp + sum ═══ + float local_sum = 0.0f; + for (int i = tid; i < tile_len; i += 32) { + const float e = expf(scores[i] - tile_max); + scores[i] = e; + local_sum += e; + } + for (int off = 16; off > 0; off >>= 1) + local_sum += __shfl_xor(local_sum, off); + const float tile_sum = local_sum; + __syncthreads(); + + // ═══ Write tile partials ═══ + // Stride by max_tiles (derived from max_seq), NOT the dynamic n_tiles, so + // the layout matches attention_flash_q8_0_reduce's per-head row stride. + // Getting this wrong makes every head except h=0 read stale/uninitialised + // partials whenever seq_len < max_seq — see the note on the Q8 kernel. + const int max_tiles = (max_seq + tile_size - 1) / tile_size; + float* p = partials + (h * max_tiles + tile_id) * (2 + head_dim); + if (tid == 0) { + p[0] = tile_max; + p[1] = tile_sum; + } + + // ═══ Phase D: V-weighted accumulation ═══ + // Loops over halves so all head_dim output dims are written, not just the + // first 128 — the reduce reads the full head_dim row, and a short write + // leaves it folding whatever the previous caller left in the shared + // partials scratch. + for (int half = 0; half < n_halves; half++) { + const int d_base = half * 128 + tid * 4; + const int elem = head_base + d_base; + float out0 = 0.0f, out1 = 0.0f, out2 = 0.0f, out3 = 0.0f; + for (int tl = 0; tl < tile_len; tl++) { + const float w = scores[tl]; + const unsigned short* vb = + v16 + (size_t)(tile_start + tl) * kv_dim + elem; + out0 += w * hipfire_bf16_to_f32(vb[0]); + out1 += w * hipfire_bf16_to_f32(vb[1]); + out2 += w * hipfire_bf16_to_f32(vb[2]); + out3 += w * hipfire_bf16_to_f32(vb[3]); + } + p[2 + d_base + 0] = out0; + p[2 + d_base + 1] = out1; + p[2 + d_base + 2] = out2; + p[2 + d_base + 3] = out3; + } +} diff --git a/kernels/src/attention_flash_bf16_tile_batched.hip b/kernels/src/attention_flash_bf16_tile_batched.hip new file mode 100644 index 0000000000..b72edcb61d --- /dev/null +++ b/kernels/src/attention_flash_bf16_tile_batched.hip @@ -0,0 +1,197 @@ +// SPDX-License-Identifier: Apache-2.0 +// Copyright (c) 2026 Nick Woolmer +// hipfire — see LICENSE and NOTICE in the project root. +// +// Batched flash attention tile for FLAT BF16 KV, with a sliding-window mask. +// Prefill sibling of `attention_flash_bf16_tile.hip`, and a direct port of +// `attention_flash_q8_0_tile_batched.hip` with the Q8_0 block addressing +// replaced by a flat `[pos][kv_head][d]` bf16 read. +// +// Processes sub_batch_size positions per launch via blockIdx.z. Partials +// layout: [sub_batch × n_heads × max_tiles × (2+head_dim)] — identical to the +// asym3/Q8 tiles, so it shares `attention_flash_asym_reduce_batched` +// unmodified (that reduce only ever touches f32 partials). +// +// LDS is tile_size*4 bytes regardless of seq_len, so there is NO context cap. +// +// cos_theta/sin_theta and v_mode_bits are accepted for ABI compatibility with +// the shared `launch_asym_flash_batched` dispatcher but ignored — bf16 K is +// unrotated and there is no separate V tier. +// +// Grid: [n_heads, max_tiles, sub_batch_size]. Block: [32, 1, 1]. +#include +#include "kv_slot_desc.h" + +static __device__ __forceinline__ float hipfire_bf16_to_f32(unsigned short b) { + return __builtin_bit_cast(float, ((unsigned int)b) << 16); +} + +extern "C" __launch_bounds__(32, 16) +__global__ void attention_flash_bf16_tile_batched( + const float* __restrict__ q, // [sub_batch × n_heads × head_dim] + const unsigned char* __restrict__ k_cache, + const unsigned char* __restrict__ v_cache, + float* __restrict__ partials, + const int* __restrict__ positions, // [total_batch] + const float* __restrict__ /*cos_theta*/, // unused (bf16 K unrotated) + const float* __restrict__ /*sin_theta*/, // unused + const float* __restrict__ tree_bias, // optional [total_batch × block_cols] + int n_heads, + int n_kv_heads, + int head_dim, + int max_seq, + float scale_attn, + int tile_size, + int max_tiles, + int batch_offset, + int block_start, // ignored when tree_bias == nullptr + int block_cols, // ignored when tree_bias == nullptr + int /*v_mode_bits*/, // consumed-but-unused (shared launcher) + int window, // sliding-window span; <= 0 = full causal + const KvSlotDesc* __restrict__ slot_descs, // [n_slots] or nullptr = legacy + const int* __restrict__ row_slot // [total_batch] or nullptr = legacy +) { + const int h = blockIdx.x; + const int tile_id = blockIdx.y; + const int local_bid = blockIdx.z; + if (h >= n_heads) return; + const int tid = threadIdx.x; + + const int global_bid = batch_offset + local_bid; + const bool tree_mode = (tree_bias != nullptr); + // hipGraph-safe block_start: under captured tree-verify replay the scalar + // kernarg is baked at capture time and goes stale as the committed prefix + // grows. positions[] is re-uploaded every cycle and the linearized tree + // root sits at positions[0], so derive it from the device buffer instead. + const int eff_block_start = tree_mode ? positions[0] : block_start; + const int seq_len = tree_mode ? (eff_block_start + block_cols) + : (positions[global_bid] + 1); + + // row_slot is indexed by GLOBAL row — using local blockIdx.z alone reads + // the wrong slot once the partials buffer forces the launcher to sub-batch + // (correct at small batch, silent cross-slot corruption after chunking). + const int slot = (row_slot != nullptr) ? row_slot[global_bid] : 0; + const KvSlotDesc desc = (slot_descs != nullptr) + ? slot_descs[slot] + : kv_slot_legacy(seq_len, max_seq); + + const int kv_dim = n_kv_heads * head_dim; + const int per_pos_bytes = kv_dim * 2; + + // `positions[row] + 1` is the PER-ROW causal bound; `desc.seq_len` is the + // slot's logical KV length (capacity metadata). They coincide only at + // M=1. The descriptor supplies the slab BASE, never the row's own bound — + // and the reduce independently derives n_tiles from positions[]+1, so any + // disagreement here makes it fold tiles this kernel never wrote. + const int eff_seq_len = seq_len; + + const int win_lo = (window > 0) ? (eff_seq_len - window) : 0; + const int tile_start = tile_id * tile_size; + if (tile_start >= eff_seq_len) return; + const int tile_end = min(tile_start + tile_size, eff_seq_len); + const int tile_len = tile_end - tile_start; + + // Sliding-window fast path: a tile entirely below win_lo contributes + // nothing. Write an empty partial (tile_sum=0 → the reduce skips it via + // its `p[1] > 0` guard) and return BEFORE the Q load and Phases A-D. + // Phase D loads V for every position even at weight 0, so masked-but- + // iterated tiles are the dominant per-query cost on sliding layers at long + // context. Only fires when window>0; window<=0 is byte-identical. + if (window > 0 && tile_end <= win_lo) { + if (tid == 0) { + float* pe = partials + + ((long long)local_bid * n_heads + h) * max_tiles * (2 + head_dim) + + tile_id * (2 + head_dim); + pe[0] = -1e30f; + pe[1] = 0.0f; + } + return; + } + + const int kv_group = n_heads / n_kv_heads; + const int kv_h = h / kv_group; + const int q_dim = n_heads * head_dim; + const int head_base = kv_h * head_dim; + + extern __shared__ float sdata[]; + float* scores = sdata; + + // Load Q for this batch position: dpt consecutive dims per thread (matches + // the asym3/Q8 tiles so the dim-indexed partials/reduce layout is + // identical). The block is always 32 lanes, so the __shfl_xor reductions + // still sum over all head_dim dims; only the per-thread span scales. + const float* q_head = q + local_bid * q_dim + h * head_dim; + const int dpt = head_dim / 32; + const int d0 = tid * dpt; + float mq[16]; + for (int i = 0; i < dpt; i++) mq[i] = q_head[d0 + i]; + + // ── Phase A: Q·K ── + // Flat bf16 needs no block/scale decomposition, so unlike the Q8 kernel + // there is no `lane_one_q8_block` fast path to select — the general form + // IS the fast form here. + for (int t_local = 0; t_local < tile_len; t_local++) { + const int t = tile_start + t_local; + if (t < win_lo) { // outside sliding window → mask, skip dot + if (tid == 0) scores[t_local] = -INFINITY; + continue; // t is wave-uniform; all lanes skip together + } + const unsigned short* kb = (const unsigned short*) + (k_cache + kv_offset_for_k(desc, t, per_pos_bytes)) + head_base + d0; + float partial = 0.0f; + for (int i = 0; i < dpt; i++) { + partial += mq[i] * hipfire_bf16_to_f32(kb[i]); + } + for (int off = 16; off > 0; off >>= 1) + partial += __shfl_xor(partial, off); + if (tid == 0) { + float s = partial * scale_attn; + if (tree_mode && t >= eff_block_start) + s += tree_bias[global_bid * block_cols + (t - eff_block_start)]; + scores[t_local] = s; + } + } + __syncthreads(); + + // ── Phase B: tile max ── + float local_max = -1e30f; + for (int i = tid; i < tile_len; i += 32) + local_max = fmaxf(local_max, scores[i]); + for (int off = 16; off > 0; off >>= 1) + local_max = fmaxf(local_max, __shfl_xor(local_max, off)); + const float tile_max = local_max; + + // ── Phase C: exp + sum ── + float local_sum = 0.0f; + for (int i = tid; i < tile_len; i += 32) { + const float e = expf(scores[i] - tile_max); + scores[i] = e; + local_sum += e; + } + for (int off = 16; off > 0; off >>= 1) + local_sum += __shfl_xor(local_sum, off); + const float tile_sum = local_sum; + __syncthreads(); + + // ── Phase D: V-weighted accumulation ── + float out_vec[16] = { 0.0f }; + for (int t_local = 0; t_local < tile_len; t_local++) { + const float w = scores[t_local]; + const int t = tile_start + t_local; + const unsigned short* vb = (const unsigned short*) + (v_cache + kv_offset_for_v(desc, t, per_pos_bytes)) + head_base + d0; + for (int i = 0; i < dpt; i++) { + out_vec[i] += w * hipfire_bf16_to_f32(vb[i]); + } + } + + float* p = partials + + ((long long)local_bid * n_heads + h) * max_tiles * (2 + head_dim) + + tile_id * (2 + head_dim); + if (tid == 0) { + p[0] = tile_max; + p[1] = tile_sum; + } + for (int i = 0; i < dpt; i++) + p[2 + d0 + i] = out_vec[i]; +} diff --git a/kernels/src/kv_cache_write_bf16.hip b/kernels/src/kv_cache_write_bf16.hip new file mode 100644 index 0000000000..7ac2881ce7 --- /dev/null +++ b/kernels/src/kv_cache_write_bf16.hip @@ -0,0 +1,98 @@ +// SPDX-License-Identifier: Apache-2.0 +// Copyright (c) 2026 Nick Woolmer +// hipfire — see LICENSE and NOTICE in the project root. +// +// Flat BF16 KV-cache write (maple). +// +// Layout: element `(pos, kv_head, d)` lives at +// `pos * kv_dim + kv_head * head_dim + d`, one bf16 (2 bytes) each, +// where `kv_dim = n_kv_heads * head_dim`. No blocks, no per-block scale, no +// padding. Compare `kv_cache_write_q8_0.hip`, which packs 34-byte blocks of +// 32 elements and needs a wave reduction to find each block's amax — none of +// that exists here, so this kernel is a pure elementwise convert-and-store +// and the block/thread mapping is free to be a plain flat grid. +// +// The same kernel serves K and V: the caller launches it twice with different +// `dst`/`src`, exactly as the Q8_0 path does. +#include +#include "kv_slot_desc.h" + +// F32 -> BF16 with round-to-nearest-even. +// +// Deliberately duplicated from `convert_f32_to_bf16.hip` rather than shared +// via a header: the JIT's header handling is a hardcoded textual special-case +// for `kv_slot_desc.h` in rdna-compute/src/dispatch.rs, and that file is +// hard-blocked by the verify-bind-thread pre-commit hook. Fifteen duplicated +// lines cost less than widening that mechanism. If this ever diverges from +// convert_f32_to_bf16.hip, THAT is the bug — both must implement the same RNE +// the host reference does (round up when round_bit && (sticky || lsb)). +static __device__ __forceinline__ unsigned short hipfire_f32_to_bf16_rne(float v) { + const unsigned int bits = __builtin_bit_cast(unsigned int, v); + // NaN: exponent all-ones and mantissa != 0. Canonicalize to a quiet BF16 + // NaN keeping the sign, matching the host reference. + if ((bits & 0x7fffffffu) > 0x7f800000u) { + return (unsigned short)(((bits >> 16) & 0x8000u) | 0x7fc0u); + } + const unsigned int lsb = (bits >> 16) & 1u; + const unsigned int lower = bits & 0xffffu; + const unsigned int round_bit = (lower >> 15) & 1u; + const unsigned int sticky = ((lower & 0x7fffu) != 0u) ? 1u : 0u; + unsigned int top = bits >> 16; + if (round_bit == 1u && (sticky == 1u || lsb == 1u)) { + top += 1u; + } + return (unsigned short)(top & 0xffffu); +} + +// Decode / single-position write. +// Grid: [ceil(kv_dim / 64), 1, 1]. Block: [64, 1, 1]. +extern "C" __global__ void kv_cache_write_bf16( + unsigned char* __restrict__ dst, + const float* __restrict__ src, // [kv_dim] FP32 KV vector + const int* __restrict__ pos_buf, + int n_kv_heads, + int head_dim +) { + const int kv_dim = n_kv_heads * head_dim; + const int i = blockIdx.x * blockDim.x + threadIdx.x; + if (i >= kv_dim) return; + + const int pos = pos_buf[0]; + unsigned short* out = + (unsigned short*)(dst + (size_t)pos * (size_t)kv_dim * 2); + out[i] = hipfire_f32_to_bf16_rne(src[i]); +} + +// Batched prefill write. +// Grid: [ceil(kv_dim / 64), batch_size, 1]. Block: [64, 1, 1]. +// +// `kv_offset_for_k` is used for both the K and the V launch, mirroring +// `kv_cache_write_q8_0_batched`. That is correct because `dst` already +// selects the K or V arena; the descriptor only adds a slot base within it, +// and maple passes `slot_descs == nullptr` (legacy, base 0) today. +extern "C" __global__ void kv_cache_write_bf16_batched( + unsigned char* __restrict__ dst, + const float* __restrict__ src, // [batch_size × kv_dim] + const int* __restrict__ positions, // [batch_size] + int n_kv_heads, + int head_dim, + int batch_size, + const KvSlotDesc* __restrict__ slot_descs, // [n_slots] or nullptr = legacy + const int* __restrict__ row_slot // [batch_size] or nullptr = legacy +) { + const int kv_dim = n_kv_heads * head_dim; + const int i = blockIdx.x * blockDim.x + threadIdx.x; + const int bid = blockIdx.y; + if (bid >= batch_size) return; + if (i >= kv_dim) return; + + const int slot = (row_slot != nullptr) ? row_slot[bid] : 0; + const KvSlotDesc desc = + (slot_descs != nullptr) ? slot_descs[slot] : kv_slot_legacy(0, 0); + const int per_pos_bytes = kv_dim * 2; + + const int pos = positions[bid]; + unsigned short* out = + (unsigned short*)(dst + kv_offset_for_k(desc, pos, per_pos_bytes)); + out[i] = hipfire_f32_to_bf16_rne(src[(size_t)bid * (size_t)kv_dim + i]); +} diff --git a/kernels/src/moe_topk_renorm_k8.hip b/kernels/src/moe_topk_renorm_k8.hip index 78952cecb9..82913df0fb 100644 --- a/kernels/src/moe_topk_renorm_k8.hip +++ b/kernels/src/moe_topk_renorm_k8.hip @@ -100,6 +100,16 @@ __global__ void moe_topk_renorm_k8( cur_v = -INFINITY; \ cur_i = -1; \ } \ + /* REQUIRED. `warp_i[0]` is read here by every thread as this \ + iteration's winner, and rewritten by warp 0 lane 0 at the TOP \ + of the next iteration (`warp_i[warp_id]`) with no barrier in \ + between. Without this, a fast warp 0 clobbers the winner \ + before a slow warp has invalidated against it; that warp \ + keeps the picked expert live and it is selected AGAIN, so the \ + top-8 comes back with a DUPLICATE and a distinct expert \ + dropped. Measured 1 occurrence per ~46k router calls before \ + this barrier, always from bit-identical input logits. */ \ + __syncthreads(); \ } \ } while (0) @@ -163,6 +173,10 @@ __global__ void moe_topk_renorm_k8( cur_v = -INFINITY; cur_i = -1; } + // REQUIRED, same hazard as the exact-256 path above: the next + // iteration's `warp_i[warp_id]` store races this read of + // `warp_i[0]`, and losing it yields a duplicated expert. + __syncthreads(); } } diff --git a/registry/models.json b/registry/models.json index d18e5332e9..b3e27e3343 100644 --- a/registry/models.json +++ b/registry/models.json @@ -313,10 +313,23 @@ "size_gb": 6.5, "min_vram_gb": 12, "sampling": { - "temperature": 0.6, + "temperature": 1.0, "top_p": 0.95 }, - "desc": "Deepgrove Maple-Preview 20B-A1B (arch_id=15). NATIVELY TERNARY: every linear weight ships as {-s,0,+s} with one bf16 scale per output row, so MQ2G256LloydU (qt=51, 2.250 bpw) carries the body VALUE-EXACTLY \u2014 a packing, not a quantization; max|err| vs the bf16 master is 0. lm_head is Q8: measured +23.9% decode for +0.00005 nats of mean KL, with top-1 agreement and PPL both slightly BETTER than bf16 (MQ4 head was rejected \u2014 2.6x the KL, concentrated on the model's most-confident positions). 24 layers, 256 experts top-8 on every layer, 3:1 sliding(512)/global-NoPE attention, QK-norm before partial RoPE, clamped SwiGLU. Reasoning model (), Qwen tokenizer, ChatML + tools. ~1530 tok/s prefill, ~158 tok/s decode on gfx1151. NOTE: wikitext PPL ~57 (a peer scores ~7.7) \u2014 that is the MODEL, not the packing: KL vs a bf16 reference is 0.084 and the vendor's own ternary GGUF scores within 0.2%. Upstream calls it a preview with minimal agentic post-training." + "desc": "Deepgrove Maple-Preview 20B-A1B (arch_id=15). NATIVELY TERNARY: every linear weight ships as {-s,0,+s} with one bf16 scale per output row, so MQ2G256LloydU (qt=51, 2.250 bpw) carries the body VALUE-EXACTLY — a packing, not a quantization; max|err| vs the bf16 master is 0. lm_head is Q8, and q8 vs bf16 is measured IDENTICAL (mean KL 0.0511 both) while decoding 23% faster, so a bf16 head is strictly dominated. Alternative heads ship as small overlays via `--head` (q4k 188 MB, bf16 635 MB) rather than as separate full models. 24 layers, 256 experts top-8 on every layer, 3:1 sliding(512)/global-NoPE attention, QK-norm before partial RoPE, clamped SwiGLU. Reasoning model (), Qwen tokenizer, ChatML + tools. ~1530 tok/s prefill, ~158 tok/s decode on gfx1151. NOTE: wikitext PPL ~57 (a peer scores ~7.7) — that is the MODEL, not the packing: KL vs a bf16 reference is 0.084 and the vendor's own ternary GGUF scores within 0.2%. Upstream calls it a preview with minimal agentic post-training. Sampling temperature 1.0 / top_p 0.95 is the VENDOR value from DeepGrove's own llama.cpp fork (github.com/deepgrove-ai/llama.cpp); there is no generation_config.json upstream. KV defaults to bf16: q8 KV costs 39% of the measured divergence (mean KL 0.0842 vs 0.0511) for 1.88x KV bytes and ~2% decode.", + "default_kv_mode": "bf16", + "heads": { + "q4k": { + "file": "maple-head-q4k.hfq", + "sha256": "deff26e9dfaddfc521f10037e00deae40fb4011fb14db3b10eb166928d7ec795", + "size_bytes": 188137472 + }, + "bf16": { + "file": "maple-head-bf16.hfq", + "sha256": "94cde3ad60d380a753f1223882682662643e49faed217df5b4db03e0a17f3e8b", + "size_bytes": 635437056 + } + } }, "north-mini-code": { "repo": "nwoolmer/hipfire-north-mini-code", diff --git a/registry/v1.json b/registry/v1.json index 65aa9982a6..aa689d95d7 100644 --- a/registry/v1.json +++ b/registry/v1.json @@ -1,6 +1,6 @@ { "schema_version": 1, - "generated_at": "2026-08-31T05:32:38Z", + "generated_at": "2026-09-01T13:14:39Z", "_comment": "GENERATED by scripts/registry_gen.py \u2014 do not hand-edit. Edit registry/models.json (curated overlay) and re-run the generator. Strict superset of registry/models.json: models/aliases keep the curated shape; sha256/size_bytes come from the HF LFS API; arch_id per docs/architecture-ids.md; min_vram_gb gates pull/run on VRAM.", "models": { "qwen3.5:0.8b": { @@ -349,10 +349,23 @@ "size_gb": 6.5, "min_vram_gb": 12, "sampling": { - "temperature": 0.6, + "temperature": 1.0, "top_p": 0.95 }, - "desc": "Deepgrove Maple-Preview 20B-A1B (arch_id=15). NATIVELY TERNARY: every linear weight ships as {-s,0,+s} with one bf16 scale per output row, so MQ2G256LloydU (qt=51, 2.250 bpw) carries the body VALUE-EXACTLY \u2014 a packing, not a quantization; max|err| vs the bf16 master is 0. lm_head is Q8: measured +23.9% decode for +0.00005 nats of mean KL, with top-1 agreement and PPL both slightly BETTER than bf16 (MQ4 head was rejected \u2014 2.6x the KL, concentrated on the model's most-confident positions). 24 layers, 256 experts top-8 on every layer, 3:1 sliding(512)/global-NoPE attention, QK-norm before partial RoPE, clamped SwiGLU. Reasoning model (), Qwen tokenizer, ChatML + tools. ~1530 tok/s prefill, ~158 tok/s decode on gfx1151. NOTE: wikitext PPL ~57 (a peer scores ~7.7) \u2014 that is the MODEL, not the packing: KL vs a bf16 reference is 0.084 and the vendor's own ternary GGUF scores within 0.2%. Upstream calls it a preview with minimal agentic post-training.", + "desc": "Deepgrove Maple-Preview 20B-A1B (arch_id=15). NATIVELY TERNARY: every linear weight ships as {-s,0,+s} with one bf16 scale per output row, so MQ2G256LloydU (qt=51, 2.250 bpw) carries the body VALUE-EXACTLY \u2014 a packing, not a quantization; max|err| vs the bf16 master is 0. lm_head is Q8, and q8 vs bf16 is measured IDENTICAL (mean KL 0.0511 both) while decoding 23% faster, so a bf16 head is strictly dominated. Alternative heads ship as small overlays via `--head` (q4k 188 MB, bf16 635 MB) rather than as separate full models. 24 layers, 256 experts top-8 on every layer, 3:1 sliding(512)/global-NoPE attention, QK-norm before partial RoPE, clamped SwiGLU. Reasoning model (), Qwen tokenizer, ChatML + tools. ~1530 tok/s prefill, ~158 tok/s decode on gfx1151. NOTE: wikitext PPL ~57 (a peer scores ~7.7) \u2014 that is the MODEL, not the packing: KL vs a bf16 reference is 0.084 and the vendor's own ternary GGUF scores within 0.2%. Upstream calls it a preview with minimal agentic post-training. Sampling temperature 1.0 / top_p 0.95 is the VENDOR value from DeepGrove's own llama.cpp fork (github.com/deepgrove-ai/llama.cpp); there is no generation_config.json upstream. KV defaults to bf16: q8 KV costs 39% of the measured divergence (mean KL 0.0842 vs 0.0511) for 1.88x KV bytes and ~2% decode.", + "default_kv_mode": "bf16", + "heads": { + "q4k": { + "file": "maple-head-q4k.hfq", + "sha256": "deff26e9dfaddfc521f10037e00deae40fb4011fb14db3b10eb166928d7ec795", + "size_bytes": 188137472 + }, + "bf16": { + "file": "maple-head-bf16.hfq", + "sha256": "94cde3ad60d380a753f1223882682662643e49faed217df5b4db03e0a17f3e8b", + "size_bytes": 635437056 + } + }, "sha256": "7fb52fe72c1a0a4455d0fe3a8109b0df66fa53782f41d8b257140d3e966645db", "size_bytes": 6499340288, "arch_id": 15, diff --git a/scripts/redline_daemon_harness.py b/scripts/redline_daemon_harness.py index 82d8a6c405..90078242e8 100755 --- a/scripts/redline_daemon_harness.py +++ b/scripts/redline_daemon_harness.py @@ -207,7 +207,7 @@ def main(): parser.add_argument("--decode-context", type=int, default=128) parser.add_argument( "--kv-mode", - choices=("q8", "fwht2", "fwht3", "fwht4"), + choices=("q8", "bf16", "fwht2", "fwht3", "fwht4"), default="q8", help="KV layout used by capture, shadow replay, and the HIP oracle", ) diff --git a/scripts/registry_gen.py b/scripts/registry_gen.py index eef423544f..e2d5d6854b 100644 --- a/scripts/registry_gen.py +++ b/scripts/registry_gen.py @@ -90,6 +90,10 @@ "auto", "f32", "f16", + # maple (arch 15) only: a flat 2-byte BF16 KV tier. Per-site acceptance + # lives in hipfire_runtime::kv_mode's policies; this list is only the + # schema allow-list, and must stay in sync with hipfire-config's KV_MODES. + "bf16", "q8", "asym4", "asym3", @@ -485,6 +489,15 @@ def build_registry(curated: dict, token: str | None) -> tuple[dict | None, list[ for kind in ("triattn", "mtp"): if isinstance(entry.get(kind), dict): new_entry[kind] = annotate_sidecar(entry[kind], tree, tag, kind, errors) + # `heads` is a MAP of sidecars (alternative lm_head overlays), not a + # single one — same annotation, once per entry, so a head that is + # missing from the repo fails the run like any other sidecar. + if isinstance(entry.get("heads"), dict): + new_entry["heads"] = { + name: annotate_sidecar(sc, tree, tag, f"heads.{name}", errors) + for name, sc in entry["heads"].items() + if isinstance(sc, dict) + } # repo probe already failed → error recorded above; entry still gets # arch_id/quant so the error list is the only blocker.