diff --git a/Cargo.lock b/Cargo.lock index 4ed89d3d8d..ad89c18ee8 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -1259,6 +1259,7 @@ dependencies = [ "hipfire-arch-qwen35-vl", "hipfire-config", "hipfire-dispatch", + "hipfire-quantize", "hipfire-reap", "hipfire-runtime", "rdna-compute", @@ -1380,6 +1381,7 @@ dependencies = [ name = "hipfire-dispatch" version = "0.3.0" dependencies = [ + "half", "hip-bridge", "hipfire-config", "rdna-compute", @@ -1515,6 +1517,7 @@ dependencies = [ "safetensors", "serde", "serde_json", + "sha2", "tempfile", ] @@ -2704,8 +2707,10 @@ dependencies = [ name = "rdna-compute" version = "0.3.0" dependencies = [ + "half", "hip-bridge", "hipfire-config", + "hipfire-quantize", "libloading", "radiowave", "rayon", diff --git a/crates/hipfire-arch-cohere2moe/src/forward.rs b/crates/hipfire-arch-cohere2moe/src/forward.rs index 49eae1b47d..1b8cdbbf35 100644 --- a/crates/hipfire-arch-cohere2moe/src/forward.rs +++ b/crates/hipfire-arch-cohere2moe/src/forward.rs @@ -34,10 +34,11 @@ use crate::cohere2moe::{Cohere2MoeState, Cohere2MoeWeights, Ffn}; use crate::config::{AttnKind, Cohere2MoeConfig}; use hipfire_dispatch::context::DispatchCtx; use hipfire_dispatch::families::moe::{MoeDtypes, MoePrefillParams}; +use hipfire_runtime::llama::KvCacheExt; use hipfire_runtime::llama::{ fused_silu_mul_rotate_mq_batched_for, moe_family, rotate_x_mq_batched_for, rotate_x_mq_for, - weight_gemv, weight_gemv_residual}; -use hipfire_runtime::llama::KvCacheExt; + weight_gemv, weight_gemv_residual, +}; use rdna_compute::{DType, Gpu, GpuTensor}; /// Grouped-MoE prefill tiling constant — must match `run_moe_prefill`'s @@ -854,6 +855,7 @@ pub fn forward_batch( routed_has_mixed_experts: false, per_expert_gate_up: None, per_expert_down: None, + routed_escha_transforms: false, has_paro_shared: false, }, batch_size: b, @@ -889,6 +891,13 @@ pub fn forward_batch( paro_down: None, down_awq_scale: None, routed_out: None, + // Not an escha model: the escha branch in + // `run_moe_prefill` is skipped and Path 1 / Path 2 run + // exactly as before, and `check_moe_prefill_supported` is + // a no-op for `layer_is_escha == false`. + escha: None, + layer_is_escha: false, + hidden, }; moe_family() .run_prefill(&ctx, gpu, ¶ms) diff --git a/crates/hipfire-arch-qwen35/Cargo.toml b/crates/hipfire-arch-qwen35/Cargo.toml index 6b61184baf..96f6a09d98 100644 --- a/crates/hipfire-arch-qwen35/Cargo.toml +++ b/crates/hipfire-arch-qwen35/Cargo.toml @@ -23,6 +23,32 @@ hipfire-reap = { path = "../hipfire-reap" } saddle-core = { path = "../saddle-core" } serde = { version = "1", features = ["derive"] } serde_json = "1" + +[dev-dependencies] +# G4b (Escha-W2 Task 9): `examples/escha_router_contract.rs` reads the escha +# fixture's f16 golden buffers via the same decoder the frozen CPU oracle +# uses (`hipfire_quantize::float16::f16_to_f32`). +hipfire-quantize = { path = "../hipfire-quantize" } + +[[example]] +name = "escha_model_smoke" +path = "examples/escha_model_smoke.rs" +required-features = ["deltanet"] + +[[example]] +name = "escha_moe_block_gate" +path = "examples/escha_moe_block_gate.rs" +required-features = ["deltanet"] + +[[example]] +name = "escha_router_contract" +path = "examples/escha_router_contract.rs" + +[[example]] +name = "escha_prefill_batch_gate" +path = "examples/escha_prefill_batch_gate.rs" +required-features = ["deltanet"] + [[example]] name = "test_qwen35_load_multi" path = "examples/test_qwen35_load_multi.rs" diff --git a/crates/hipfire-arch-qwen35/examples/escha_config_smoke.rs b/crates/hipfire-arch-qwen35/examples/escha_config_smoke.rs new file mode 100644 index 0000000000..0c2a5a850e --- /dev/null +++ b/crates/hipfire-arch-qwen35/examples/escha_config_smoke.rs @@ -0,0 +1,16 @@ +//! Prove the converted .hfq's metadata satisfies arch-6's config loader. +fn main() { + let path = std::env::args().nth(1).expect("usage: "); + let hfq = hipfire_runtime::hfq::HfqFile::open(std::path::Path::new(&path)) + .expect("open hfq"); + println!("arch_id = {}", hfq.arch_id); + println!("tensors = {}", hfq.tensors().len()); + match hipfire_arch_qwen35::qwen35::config::config_from_hfq(&hfq) { + Ok(c) => println!( + "config OK: dim={} layers={} experts={} top_k={} moe_inter={} vocab={} is_vl_text={}", + c.dim, c.n_layers, c.num_experts, c.num_experts_per_tok, + c.moe_intermediate_size, c.vocab_size, c.is_vl_text + ), + Err(e) => { eprintln!("config FAILED: {e}"); std::process::exit(1); } + } +} diff --git a/crates/hipfire-arch-qwen35/examples/escha_model_smoke.rs b/crates/hipfire-arch-qwen35/examples/escha_model_smoke.rs new file mode 100644 index 0000000000..07d0ea6c95 --- /dev/null +++ b/crates/hipfire-arch-qwen35/examples/escha_model_smoke.rs @@ -0,0 +1,265 @@ +// SPDX-License-Identifier: Apache-2.0 +// Copyright (c) 2026 Kaden Schutt +// hipfire — see LICENSE and NOTICE in the project root. +//! Escha-W2 end-to-end smoke (Task 10): load the WHOLE `.hfq` single-GPU and +//! run PREFILL and then decode through the production forward paths. +//! +//! The G4 block gate (`escha_moe_block_gate`) calls the routed executor +//! directly, with routing injected. That proves the maths; it does NOT prove +//! that a real `qwen35::forward` ever reaches it. This does: it asserts layer +//! 0 came through the escha loader, then decodes and reads the H128 launch +//! counter, which must be exactly `4 * n_layers` per token — the batched +//! budget. A regression to a per-expert wiring shows up here as `4 * k * +//! n_layers` (1280 at A3B) rather than 160, with no numerical change at all. +//! +//! # The prefill phase, and why the launch counter is the load-bearing assert +//! +//! Escha layers route to an escha executor in prefill too, but through +//! `escha_routed_prefill_indexed` rather than the decode one. The H128 launch +//! budget is what identifies which of THREE things happened, and the three are +//! indistinguishable by looking at the logits: +//! +//! | launches (8-token prompt, 40 layers) | what ran | +//! |---|---| +//! | **160** = `4 * n_layers` per CHUNK | batched escha prefill — correct | +//! | 1 280 = `n * 4 * n_layers` | silently fell back to the per-token loop | +//! | 0 | a batched MoE body with NO escha awareness | +//! +//! Zero is the dangerous one: the generic batched routed body would run the +//! Q8_0 experts without the H128 pair and emit finite, fluent, ~1e-1-wrong +//! hidden state, which finiteness alone would never catch. 1 280 is not wrong, +//! just 3.6x slower — but a silent fallback is exactly how a performance fix +//! rots, so it fails here too rather than being tolerated. +//! +//! The count is per CHUNK because that is the whole point of the batched body: +//! one launch of `escha_h128_in_batched` covers `n_tokens * k` slots. A prompt +//! longer than the prefill chunk ceiling would legitimately show one budget +//! per chunk; this gate keeps the prompt inside one chunk so the expected +//! number is exact. +//! +//! Token ids are arbitrary here on purpose: this gate is about the launch +//! budget and the structural invariants (finite logits, a non-degenerate +//! argmax), not about semantics. Semantic checking is Task 11's — the +//! converter now embeds the tokenizer, chat template and generation_config, so +//! the daemon DOES drive this checkpoint (`scripts/_coherence_runner.py`, and +//! §10.4 of the design doc). +//! +//! COST: **37.6 GB resident** (37 587 996 672 B), measured as an amdgpu GTT +//! delta on gfx1151 (`scripts/escha-gtt-probe.sh`: 40.94 GB peak over a +//! 3.36 GB idle baseline). 34.2 GB of that is the Q8_0 routed experts and +//! ~3.3 GB is everything else. It was 67.9 GB until the experts were packed +//! one device buffer per (layer, projection): while each of the 20,480 +//! per-expert buffers was its own allocation, the HIP allocator's 2 MiB +//! granule rounded the 2.125 MiB gate_up up to 4 MiB and the 1.0625 MiB down +//! up to 2 MiB, spending 30 GB on rounding. Still not free on a 128 GB +//! workstation with other applications running — check headroom first. +//! See design doc §10.3, which now records this figure rather than the 67.9 GB +//! it predated. +//! +//! Run: +//! cargo run --release -p hipfire-arch-qwen35 \ +//! --example escha_model_smoke -- /data/hipfire-models/escha-35b.hfq +use hipfire_arch_qwen35::qwen35; +use hipfire_runtime::hfq::HfqFile; +use hipfire_runtime::loader_api::{CaskConfig, LoadCtx, ModelSource, SpecLoadCfg}; +use rdna_compute::Gpu; +use std::path::Path; + +fn main() -> Result<(), String> { + let path = std::env::args() + .nth(1) + .unwrap_or_else(|| "/data/hipfire-models/escha-35b.hfq".to_string()); + let hfq = HfqFile::open(Path::new(&path)).map_err(|e| format!("open: {e:?}"))?; + let mut gpu = Gpu::init().map_err(|e| format!("gpu: {e:?}"))?; + let cask = CaskConfig::default(); + let src = ModelSource::Hfq(hfq); + let mut ctx = LoadCtx { + path: &path, + max_seq: 512, + deepseek4_compute_placement: Default::default(), + deepseek4_experts_per_token: None, + draft_path: None, + kv_mode_override: None, + kv_backend: hipfire_runtime::kv_backend::KvBackend::Contiguous, + kv_adaptive_override: None, + state_quant_override: None, + cask: &cask, + pp: 1, + spec: SpecLoadCfg::default(), + gpu: &mut gpu, + gemma4_drafter_path: None, + gemma4_draft_len: 3, + }; + let t0 = std::time::Instant::now(); + let mut b = hipfire_arch_qwen35::load_qwen35_bundle(src, &mut ctx)?; + eprintln!("loaded in {:?}", t0.elapsed()); + + // Layer 0 must have come through the escha loader, and its experts must + // hold one of the containers that loader produces — not whatever the + // generic per-expert path would have found. + // + // The exact container depends on `HIPFIRE_ESCHA_EXPERT_STORE` and is not + // what this gate is about, so it is asserted as a SET rather than pinned + // to one value. It is asserted at all because the failure it catches is + // "the escha loader did not run and some other path filled these slots", + // which is a different bug from a wrong store. + match &b.weights.layers[0] { + qwen35::LayerWeights::DeltaNetMoe(l) => { + assert!(l.ffn.escha.is_some(), "layer 0 carries no escha tables"); + assert!( + matches!( + l.ffn.experts[0].gate_up.gpu_dtype, + rdna_compute::DType::Escha2T16 + | rdna_compute::DType::Escha3T16 + | rdna_compute::DType::Q8_0 + ), + "layer 0 routed experts are {:?}, which no escha store produces", + l.ffn.experts[0].gate_up.gpu_dtype + ); + eprintln!( + "layer0: escha=Some experts={} gate_up dtype={:?} m={} k={}", + l.ffn.experts.len(), + l.ffn.experts[0].gate_up.gpu_dtype, + l.ffn.experts[0].gate_up.m, + l.ffn.experts[0].gate_up.k + ); + } + _ => panic!("layer 0 is not a DeltaNet+MoE layer"), + } + + let want_launches = + hipfire_dispatch::pipeline::escha::escha_launches_per_token(b.config.n_layers); + + // ── Phase 1: PREFILL ───────────────────────────────────────────────── + // 8 tokens, matching the G4 fixture width, through the real batched + // prefill entry point (which is expected to fall through to its per-token + // loop — see the module docs). + const PROMPT: [u32; 8] = [1000, 2000, 3000, 4000, 5000, 6000, 7000, 8000]; + let before_prefill = rdna_compute::escha_h128_launches(); + let t = std::time::Instant::now(); + qwen35::forward_prefill_batch( + ctx.gpu, + &b.weights, + &b.config, + &PROMPT, + 0, + &mut b.kv_cache, + &mut b.dn_state, + &b.scratch, + None, // hidden ring + None, // per-token hidden out — keep last-token logits enabled + None, // gdn tape + None, // tree verify + ) + .map_err(|e| format!("prefill: {e:?}"))?; + ctx.gpu + .hip + .device_synchronize() + .map_err(|e| format!("sync: {e:?}"))?; + let prefill_launches = rdna_compute::escha_h128_launches() - before_prefill; + let prefill_logits = ctx + .gpu + .download_f32(&b.scratch.logits) + .map_err(|e| format!("download prefill logits: {e:?}"))?; + let prefill_bad = prefill_logits + .iter() + .take(b.config.vocab_size) + .filter(|v| !v.is_finite()) + .count(); + eprintln!( + "prefill n={}: H128 launches={prefill_launches} (want {want_launches} for one \ + batched chunk; {} would be the per-token fallback), non-finite logits={prefill_bad}/{}, \ + {:?}", + PROMPT.len(), + PROMPT.len() * want_launches, + b.config.vocab_size, + t.elapsed() + ); + assert_eq!( + prefill_bad, + 0, + "non-finite logits after an {}-token prefill", + PROMPT.len() + ); + // BOTH correct routes are accepted, and everything else fails. + // + // `HIPFIRE_PREFILL_BATCHED=0` is a supported escape hatch, so a gate that + // demanded the batched count would fail the model under a configuration it + // is meant to survive — and a gate that has to be run with one specific + // env is a gate people stop running. What must never be accepted is ZERO: + // that is the generic batched MoE body running escha weights without the + // H128 pair, the finite-fluent-wrong case no logit check would catch. + let per_token_total = PROMPT.len() * want_launches; + let route = match prefill_launches as usize { + n if n == want_launches => "batched escha prefill body", + n if n == per_token_total => "per-token fallback (HIPFIRE_PREFILL_BATCHED=0?)", + _ => "UNKNOWN", + }; + eprintln!("prefill route: {route}"); + assert!( + prefill_launches as usize == want_launches || prefill_launches as usize == per_token_total, + "PREFILL issued {prefill_launches} H128 launches, which is neither the batched \ + budget ({want_launches} = 4 x {} layers, once for the whole chunk) nor the \ + per-token one ({per_token_total}). ZERO in particular means the model reached a \ + BATCHED MoE body with NO escha awareness: it omits both Hadamard transforms and \ + emits finite, fluent, ~1e-1-wrong hidden state that no finiteness or argmax check \ + would catch. Check that the escha branch at the top of `run_moe_prefill` still \ + fires before Path 1 / Path 2.", + b.config.n_layers + ); + // Under the DEFAULT configuration the batched body is the expected route; + // a silent fall back to per-token is correct but 3.6x slower, and a + // performance fix that quietly stops applying is how this regresses. + if hipfire_runtime::config::get().prefill_batched { + assert_eq!( + prefill_launches as usize, want_launches, + "default config, but prefill took the per-token route ({prefill_launches} \ + launches). Run with HIPFIRE_DEBUG_BATCH=1 to see which layer refused." + ); + } + + // ── Phase 2: DECODE, continuing from the prefilled context ─────────── + let mut prev = rdna_compute::escha_h128_launches(); + for (i, &tok) in [9000u32, 10000, 11000, 12000].iter().enumerate() { + let pos = PROMPT.len() + i; + let t = std::time::Instant::now(); + let logits = qwen35::forward( + ctx.gpu, + &b.weights, + &b.config, + tok, + pos, + &mut b.kv_cache, + &mut b.dn_state, + ) + .map_err(|e| format!("forward: {e:?}"))?; + + let n_bad = logits.iter().filter(|v| !v.is_finite()).count(); + let mut best = f32::NEG_INFINITY; + let mut argmax = 0usize; + for (j, &v) in logits.iter().enumerate() { + if v > best { + best = v; + argmax = j; + } + } + let mean = logits.iter().sum::() / logits.len() as f32; + let now = rdna_compute::escha_h128_launches(); + let launches = now - prev; + prev = now; + eprintln!( + "pos {pos} tok {tok}: {} logits, non-finite={n_bad}, argmax={argmax} ({best:.4}), \ + mean={mean:.4}, H128 launches={launches}, {:?}", + logits.len(), + t.elapsed() + ); + assert_eq!(n_bad, 0, "non-finite logits at pos {pos}"); + assert!(best > mean, "degenerate logit distribution at pos {pos}"); + assert_eq!( + launches as usize, want_launches, + "H128 launches per token drifted from the batched budget" + ); + } + eprintln!("escha_model_smoke: PASS"); + Ok(()) +} diff --git a/crates/hipfire-arch-qwen35/examples/escha_moe_block_gate.rs b/crates/hipfire-arch-qwen35/examples/escha_moe_block_gate.rs new file mode 100644 index 0000000000..0126e9e5be --- /dev/null +++ b/crates/hipfire-arch-qwen35/examples/escha_moe_block_gate.rs @@ -0,0 +1,830 @@ +// SPDX-License-Identifier: Apache-2.0 +// Copyright (c) 2026 Kaden Schutt +// hipfire — see LICENSE and NOTICE in the project root. +//! G4 (Escha-W2 port, Task 10): arch-6 must reproduce EschaLabs' layer-0 MoE +//! block. +//! +//! Runs the shipped `moeblk_x.f16` through hipfire's layer-0 MoE block with +//! `moeblk_ids.i64` / `moeblk_scores.f32` **injected** — the fixture ships the +//! routing precisely because it does not gate the router; that was Task 9's +//! job (`examples/escha_router_contract.rs`). What this gates is the part +//! Task 10 built: expert loading (trellis decode -> transpose -> Q8_0) and the +//! H128-wrapped, batched-across-experts routed executor. +//! +//! The golden is `routed + shared expert`, with no residual add — verified by +//! decomposition, not assumed (the routed sum alone lands at cos 0.266 against +//! the golden and 22% of its magnitude; adding the shared expert takes it to +//! cos 1.00000). +//! +//! # This is a TOLERANCE gate, and it has two arms +//! +//! The golden came from EschaLabs' Metal path, not from `ref.py`, so exact +//! agreement is not available at any weight precision. The codec goldens +//! (G2 `test_escha_decode_gpu_vs_cpu`, G3 `test_escha_h128_gpu_vs_cpu`) ARE +//! bit-exact; do NOT generalise the bounds below to them. +//! +//! Two arms run, because the two error sources are independent and must not +//! be allowed to hide each other: +//! +//! * **F32 arm** — experts stored as the exactly-decoded fp16 widened to f32, +//! no re-quantisation. This isolates the WIRING: transpose orientation, +//! H128 placement, SwiGLU half order, the f16(score) combine. If the H128 +//! pair were missing, this arm lands near 1e-1, not 1e-4. +//! * **Q8_0 arm** — production storage. The delta between the arms IS the cost +//! of the 8-bit re-quantisation, reported explicitly rather than buried in +//! a single pass/fail number. +//! +//! Run: +//! cargo run --release -p hipfire-arch-qwen35 \ +//! --example escha_moe_block_gate -- /data/hipfire-models/escha-35b.hfq + +use hipfire_arch_qwen35::qwen35::escha::{load_escha_moe_experts, EschaWeightStore}; +use hipfire_dispatch::context::DispatchCtx; +use hipfire_dispatch::pipeline::escha::{ + escha_launches_per_token, escha_routed_decode, escha_routed_decode_indexed, + escha_routed_prefill_indexed, EschaIndexedRouting, +}; +use hipfire_quantize::float16::f16_to_f32; +use hipfire_runtime::hfq::{load_weight_tensor_pread, HfqFile}; +use hipfire_runtime::llama::{weight_gemv, WeightTensor}; +use rdna_compute::{DType, Gpu, GpuTensor}; +use std::path::PathBuf; + +fn fixture(name: &str) -> PathBuf { + PathBuf::from(env!("CARGO_MANIFEST_DIR")) + .join("../hipfire-quantize/tests/data/escha") + .join(name) +} + +fn read_f16(name: &str) -> Vec { + std::fs::read(fixture(name)) + .expect("run crates/hipfire-quantize/tests/data/escha/fetch-goldens.sh first") + .chunks_exact(2) + .map(|c| f16_to_f32(u16::from_le_bytes([c[0], c[1]]))) + .collect() +} + +/// The candidate-name expander the qwen35 loader uses. The escha `.hfq` +/// already carries fully-qualified `model.language_model.*` names, so this is +/// the identity for every name below; passing the real expander keeps the gate +/// on the same lookup path production takes. +fn exact_or_prefixed(name: &str) -> Vec { + if name.starts_with("model.") { + vec![name.to_string()] + } else { + vec![ + format!("model.language_model.{name}"), + format!("model.{name}"), + name.to_string(), + ] + } +} + +fn upload_f32(gpu: &Gpu, v: &[f32]) -> GpuTensor { + let bytes: &[u8] = + unsafe { std::slice::from_raw_parts(v.as_ptr() as *const u8, std::mem::size_of_val(v)) }; + gpu.upload_raw(bytes, &[v.len()]).expect("upload") +} + +struct SharedExpert { + gate: WeightTensor, + up: WeightTensor, + down: WeightTensor, + scalar_gate: WeightTensor, +} + +/// The shared expert, run exactly as `run_moe_decode_cpu_fallback`'s generic +/// (non-MQ4) shared-down arm runs it — sigmoid(gate·x) scaling a SwiGLU MLP, +/// accumulated into the output. Unchanged arch-6 code; Task 10 does not touch +/// it, and it is here only because the golden includes it. +fn run_shared_expert( + gpu: &mut Gpu, + w: &SharedExpert, + x: &GpuTensor, + out: &GpuTensor, + smi: usize, + hidden: usize, +) { + let scalar = gpu.alloc_tensor(&[1], DType::F32).unwrap(); + let g = gpu.alloc_tensor(&[smi], DType::F32).unwrap(); + let u = gpu.alloc_tensor(&[smi], DType::F32).unwrap(); + let h = gpu.alloc_tensor(&[smi], DType::F32).unwrap(); + let y = gpu.alloc_tensor(&[hidden], DType::F32).unwrap(); + + weight_gemv(gpu, &w.scalar_gate, x, &scalar).unwrap(); + gpu.sigmoid_f32(&scalar).unwrap(); + weight_gemv(gpu, &w.gate, x, &g).unwrap(); + weight_gemv(gpu, &w.up, x, &u).unwrap(); + gpu.silu_mul_f32(&g, &u, &h).unwrap(); + weight_gemv(gpu, &w.down, &h, &y).unwrap(); + gpu.scaled_add_inplace_gpu_scalar_f32(out, &y, &scalar) + .unwrap(); + + for t in [scalar, g, u, h, y] { + let _ = gpu.free_tensor(t); + } +} + +struct ArmResult { + max_abs: f32, + mean_abs: f32, + got: Vec, + /// `Some(n)` when the indexed (GPU-top-K) route was cross-checked on this + /// arm: the number of ROUTED-ONLY output floats that did NOT match the + /// CPU-top-K route bit-for-bit, across every token. Must be 0. + indexed_mismatches: Option, + /// `Some(n)` when the BATCHED-PREFILL routed executor was cross-checked on + /// this arm: the number of routed-only output floats where running all + /// `n_tok` tokens in ONE call disagreed, bit-for-bit, with running them one + /// at a time through `escha_routed_decode_indexed`. + /// + /// Must be 0, and equality — not a tolerance — is the right standard here + /// for the same reason it is for indexed-vs-host: every kernel in the + /// escha routed pipeline is purely slot-parallel, so slot `s` performs the + /// identical FLOPs in the identical order regardless of how many slots the + /// launch carried. A difference means a wrong stride, a wrong `x_group` + /// row, or a scratch aliasing bug — never rounding. + batched_mismatches: Option, +} + +/// Device pointer table over the loaded experts, in the `[n_exp]` packed-u64 +/// layout every indexed MoE GEMV consumes. Built here the same way +/// `qwen35::load` builds it for production, because the indexed executor +/// reaches the weights ONLY through this table — a gate that passed its own +/// hand-rolled table would not be gating the production addressing. +fn expert_ptr_table(gpu: &Gpu, ptrs: &[u64]) -> GpuTensor { + let bytes: Vec = ptrs.iter().flat_map(|p| p.to_ne_bytes()).collect(); + gpu.upload_raw(&bytes, &[2 * ptrs.len()]) + .expect("ptr table") +} + +#[allow(clippy::too_many_arguments)] +fn run_arm( + gpu: &mut Gpu, + hfq: &HfqFile, + layer_prefix: &str, + shared: &SharedExpert, + store: EschaWeightStore, + x: &[f32], + want: &[f32], + ids: &[i64], + scores: &[f32], + n_tok: usize, + top_k: usize, + n_exp: usize, + hidden: usize, + mi: usize, + smi: usize, +) -> ArmResult { + let all: Vec = (0..n_exp).collect(); + let (experts, tables, owners) = load_escha_moe_experts( + hfq, + gpu, + layer_prefix, + &all, + n_exp, + hidden, + mi, + top_k, + store, + exact_or_prefixed, + ) + .expect("escha expert load"); + + let refs = tables.refs(); + let routed: Vec<_> = experts + .iter() + .map(|e| (e.gate_up.dispatch_ref(), e.down.dispatch_ref())) + .collect(); + let ctx = DispatchCtx::new(gpu); + + let out = gpu.alloc_tensor(&[hidden], DType::F32).unwrap(); + let zeros = vec![0.0f32; hidden]; + let mut got = vec![0.0f32; n_tok * hidden]; + + // ── Indexed (GPU-top-K) route cross-check ──────────────────────────── + // + // Production decode and prefill take `escha_routed_decode_indexed`, not + // the host-routed executor this gate was originally written against. The + // two must agree BIT-FOR-BIT: they are the same eight phases over the + // same weights, differing only in where the routing lives, so any + // difference at all is a defect (a changed GEMV accumulate order, a + // device-vs-host disagreement in the f16 score rounding, a wrong slot + // stride) — not a tolerance question. Asserting equality rather than a + // bound is what lets the golden tolerances below keep meaning ONE thing + // for both routes. + // + // Q8_0 and Native only: the indexed GEMVs decode one specific container + // each (a 34 B/32-element Q8_0 block, or a 16x16 trellis tile), so the F32 + // weight-exact control arm has no indexed counterpart — and needs none, it + // exists to isolate wiring, which every route shares. + // + // Native is the reverse case: it has no HOST counterpart. There is no + // per-expert native GEMV (`GemvFamily::run_auto` refuses an escha dtype + // outright, which is the fail-closed behaviour that keeps escha off a + // Plain GEMV), so for that store the indexed route is not a cross-check — + // it IS the route, and its output is what the golden comparison is made + // against. `host_route_supported` is the one flag that distinguishes them. + let host_route_supported = !matches!(store, EschaWeightStore::Native); + let indexed = if matches!(store, EschaWeightStore::Q8_0 | EschaWeightStore::Native) { + let gu_ptrs: Vec = experts + .iter() + .map(|e| e.gate_up.buf.buf.as_ptr() as u64) + .collect(); + let dn_ptrs: Vec = experts + .iter() + .map(|e| e.down.buf.buf.as_ptr() as u64) + .collect(); + Some(( + expert_ptr_table(gpu, &gu_ptrs), + expert_ptr_table(gpu, &dn_ptrs), + gpu.alloc_tensor(&[top_k], DType::F32).unwrap(), // ids (i32 bits) + gpu.alloc_tensor(&[top_k], DType::F32).unwrap(), // raw scores + gpu.alloc_tensor(&[hidden], DType::F32).unwrap(), // routed-only out + )) + } else { + None + }; + // Only meaningful when there are two routes to compare. + let mut indexed_mismatches = indexed + .as_ref() + .filter(|_| host_route_supported) + .map(|_| 0usize); + let (gu_dtype, dn_dtype) = (experts[0].gate_up.gpu_dtype, experts[0].down.gpu_dtype); + let mut per_token_indexed: Vec = Vec::new(); + + for t in 0..n_tok { + let x_gpu = upload_f32(gpu, &x[t * hidden..(t + 1) * hidden]); + gpu.hip + .memcpy_htod(&out.buf, unsafe { + std::slice::from_raw_parts(zeros.as_ptr() as *const u8, hidden * 4) + }) + .unwrap(); + + let slot_ids: Vec = ids[t * top_k..(t + 1) * top_k] + .iter() + .map(|&v| v as usize) + .collect(); + let slot_w = &scores[t * top_k..(t + 1) * top_k]; + if host_route_supported { + escha_routed_decode( + &ctx, gpu, &refs, &routed, &slot_ids, slot_w, &x_gpu, &out, hidden, mi, + ) + .expect("escha routed decode"); + } + + if let Some((gu_tbl, dn_tbl, ids_dev, wts_dev, out_idx)) = indexed.as_ref() { + // Routing goes up ONCE, as the device buffers the GPU top-K + // kernel would have written: ids as i32 bits in an F32 tensor, + // scores UNROUNDED (the executor's own kernel does the f16 + // round-trip — that is part of what is being gated). + let id_bytes: Vec = slot_ids + .iter() + .flat_map(|&i| (i as i32).to_le_bytes()) + .collect(); + gpu.hip.memcpy_htod(&ids_dev.buf, &id_bytes).unwrap(); + let w_bytes: Vec = slot_w.iter().flat_map(|w| w.to_le_bytes()).collect(); + gpu.hip.memcpy_htod(&wts_dev.buf, &w_bytes).unwrap(); + gpu.hip + .memcpy_htod(&out_idx.buf, unsafe { + std::slice::from_raw_parts(zeros.as_ptr() as *const u8, hidden * 4) + }) + .unwrap(); + + escha_routed_decode_indexed( + gpu, + &refs, + &EschaIndexedRouting { + expert_gate_up_ptrs: gu_tbl, + expert_down_ptrs: dn_tbl, + topk_indices: ids_dev, + topk_weights: wts_dev, + n_experts: n_exp, + gate_up_dtype: gu_dtype, + down_dtype: dn_dtype, + gate_up_m: 2 * mi, + gate_up_k: hidden, + down_m: hidden, + down_k: mi, + }, + out_idx, + &x_gpu, + hidden, + mi, + top_k, + ) + .expect("escha routed decode (indexed)"); + + gpu.hip.device_synchronize().unwrap(); + let idx_route = gpu.download_f32(out_idx).unwrap(); + if let Some(bad) = indexed_mismatches.as_mut() { + let host_route = gpu.download_f32(&out).unwrap(); + *bad += host_route[..hidden] + .iter() + .zip(idx_route[..hidden].iter()) + .filter(|(a, b)| a.to_bits() != b.to_bits()) + .count(); + } else { + // Native: the indexed route is the only route, so its routed + // output is what the shared expert accumulates onto and what + // the golden comparison sees. + gpu.hip + .memcpy_dtod_at(&out.buf, 0, &out_idx.buf, 0, hidden * 4) + .unwrap(); + } + // Keep the per-token indexed result as the oracle for the batched + // executor below. + per_token_indexed.extend_from_slice(&idx_route[..hidden]); + } + + run_shared_expert(gpu, shared, &x_gpu, &out, smi, hidden); + + gpu.hip.device_synchronize().unwrap(); + let row = gpu.download_f32(&out).unwrap(); + got[t * hidden..(t + 1) * hidden].copy_from_slice(&row[..hidden]); + let _ = gpu.free_tensor(x_gpu); + } + + // ── Batched-prefill routed executor cross-check ────────────────────── + // + // The whole point of `escha_routed_prefill_indexed` is that `slots` grows + // from `k` to `n_tok * k` and NOTHING else changes. This runs all `n_tok` + // tokens in one call and requires the result to be bit-identical to the + // per-token indexed route captured above — the routed half of the §5.4 + // batched-prefill gate, at the layer where a failure is diagnosable. + // + // The tokens deliberately have DIFFERENT expert sets and different + // activations (they come from EschaLabs' shipped fixture), so an executor + // that broadcast token 0's activation to every slot — the `x_group` + // mistake this is here to catch — fails on tokens 1.. rather than passing + // by luck. + let batched_mismatches = indexed.as_ref().map(|(gu_tbl, dn_tbl, _, _, _)| { + let slots = n_tok * top_k; + let x_all = upload_f32(gpu, x); + let id_bytes: Vec = ids.iter().flat_map(|&i| (i as i32).to_le_bytes()).collect(); + let ids_dev = gpu.upload_raw(&id_bytes, &[slots]).unwrap(); + let w_bytes: Vec = scores.iter().flat_map(|w| w.to_le_bytes()).collect(); + let wts_dev = gpu.upload_raw(&w_bytes, &[slots]).unwrap(); + let out_b = gpu.alloc_tensor(&[n_tok * hidden], DType::F32).unwrap(); + let zb = vec![0.0f32; n_tok * hidden]; + gpu.hip + .memcpy_htod(&out_b.buf, unsafe { + std::slice::from_raw_parts(zb.as_ptr() as *const u8, n_tok * hidden * 4) + }) + .unwrap(); + let scratch = gpu + .ensure_escha_prefill_scratch(slots, hidden, mi) + .expect("escha prefill scratch"); + escha_routed_prefill_indexed( + gpu, + &refs, + &scratch, + &EschaIndexedRouting { + expert_gate_up_ptrs: gu_tbl, + expert_down_ptrs: dn_tbl, + topk_indices: &ids_dev, + topk_weights: &wts_dev, + n_experts: n_exp, + gate_up_dtype: gu_dtype, + down_dtype: dn_dtype, + gate_up_m: 2 * mi, + gate_up_k: hidden, + down_m: hidden, + down_k: mi, + }, + &out_b, + &x_all, + hidden, + mi, + top_k, + n_tok, + ) + .expect("escha routed prefill (batched)"); + gpu.hip.device_synchronize().unwrap(); + let batched = gpu.download_f32(&out_b).unwrap(); + let bad = batched[..n_tok * hidden] + .iter() + .zip(per_token_indexed.iter()) + .filter(|(a, b)| a.to_bits() != b.to_bits()) + .count(); + for t in [x_all, ids_dev, wts_dev, out_b] { + let _ = gpu.free_tensor(t); + } + bad + }); + + let diffs: Vec = got + .iter() + .zip(want.iter()) + .map(|(a, b)| (a - b).abs()) + .collect(); + let max_abs = diffs.iter().cloned().fold(0.0f32, f32::max); + let mean_abs = diffs.iter().sum::() / diffs.len() as f32; + + let _ = gpu.free_tensor(out); + if let Some((gu_tbl, dn_tbl, ids_dev, wts_dev, out_idx)) = indexed { + for t in [gu_tbl, dn_tbl, ids_dev, wts_dev, out_idx] { + let _ = gpu.free_tensor(t); + } + } + // Expert slots are non-owning views into `owners` (one blob per + // projection), so return the two blobs and NOT the 512 views. `free_all` + // on a view is refused by `free_tensor` and would leak the blob — this + // gate runs the F32 arm at 256 experts, i.e. 2 GiB per projection, and + // both arms run in one process. + drop(experts); + let _ = gpu.free_tensor(owners.gate_up); + let _ = gpu.free_tensor(owners.down); + tables.free_gpu(gpu); + + ArmResult { + max_abs, + mean_abs, + got, + indexed_mismatches, + batched_mismatches, + } +} + +fn main() { + let path = std::env::args() + .nth(1) + .unwrap_or_else(|| "/data/hipfire-models/escha-35b.hfq".to_string()); + let hfq = HfqFile::open(std::path::Path::new(&path)).expect("open hfq"); + let config = hipfire_arch_qwen35::qwen35::config::config_from_hfq(&hfq).expect("config"); + let hidden = config.dim; + let mi = config.moe_intermediate_size; + let smi = config.shared_expert_intermediate_size; + let n_exp = config.num_experts; + let top_k = config.num_experts_per_tok; + let layer_prefix = "model.language_model.layers.0"; + + let x = read_f16("moeblk_x.f16"); + let want = read_f16("moeblk_out.f16"); + let n_tok = x.len() / hidden; + assert_eq!(want.len(), n_tok * hidden, "fixture shape mismatch"); + let ids: Vec = std::fs::read(fixture("moeblk_ids.i64")) + .unwrap() + .chunks_exact(8) + .map(|c| i64::from_le_bytes(c.try_into().unwrap())) + .collect(); + let scores: Vec = std::fs::read(fixture("moeblk_scores.f32")) + .unwrap() + .chunks_exact(4) + .map(|c| f32::from_le_bytes(c.try_into().unwrap())) + .collect(); + assert_eq!(ids.len(), n_tok * top_k); + assert_eq!(scores.len(), n_tok * top_k); + + let mut gpu = Gpu::init().expect("gpu"); + assert!( + hipfire_arch_qwen35::qwen35::escha::layer_is_escha(&hfq, layer_prefix, exact_or_prefixed), + "layer 0 of {path} does not carry Escha-W2 routed experts" + ); + + let shared = SharedExpert { + gate: load_weight_tensor_pread( + &hfq, + &gpu, + &format!("{layer_prefix}.mlp.shared_expert.gate_proj.weight"), + smi, + hidden, + exact_or_prefixed, + ) + .unwrap(), + up: load_weight_tensor_pread( + &hfq, + &gpu, + &format!("{layer_prefix}.mlp.shared_expert.up_proj.weight"), + smi, + hidden, + exact_or_prefixed, + ) + .unwrap(), + down: load_weight_tensor_pread( + &hfq, + &gpu, + &format!("{layer_prefix}.mlp.shared_expert.down_proj.weight"), + hidden, + smi, + exact_or_prefixed, + ) + .unwrap(), + scalar_gate: load_weight_tensor_pread( + &hfq, + &gpu, + &format!("{layer_prefix}.mlp.shared_expert_gate.weight"), + 1, + hidden, + exact_or_prefixed, + ) + .unwrap(), + }; + + let mag = want.iter().map(|v| v.abs()).sum::() / want.len() as f32; + println!("tokens={n_tok} hidden={hidden} top_k={top_k} experts={n_exp}"); + println!("golden mean magnitude: {mag:.4e}"); + + // ── Arm 1: weight-exact (F32) — isolates the wiring ────────────────── + let launches_before = rdna_compute::escha_h128_launches(); + let f32_arm = run_arm( + &mut gpu, + &hfq, + layer_prefix, + &shared, + EschaWeightStore::F32, + &x, + &want, + &ids, + &scores, + n_tok, + top_k, + n_exp, + hidden, + mi, + smi, + ); + let launches_one_layer_all_tokens = rdna_compute::escha_h128_launches() - launches_before; + println!( + "MoE block [F32 experts, weight-exact]: max|diff|={:.3e} mean|diff|={:.3e}", + f32_arm.max_abs, f32_arm.mean_abs + ); + + // ── Arm 2: production (Q8_0) ───────────────────────────────────────── + let q8_arm = run_arm( + &mut gpu, + &hfq, + layer_prefix, + &shared, + EschaWeightStore::Q8_0, + &x, + &want, + &ids, + &scores, + n_tok, + top_k, + n_exp, + hidden, + mi, + smi, + ); + println!( + "MoE block [Q8_0 experts, production]: max|diff|={:.3e} mean|diff|={:.3e}", + q8_arm.max_abs, q8_arm.mean_abs + ); + + // ── Arm 3: production (Native — the trellis code, fused GEMV) ──────── + // + // Phase 2's production store. It is WEIGHT-EXACT: the fused GEMV consumes + // the same fp16 values `escha_decode_tiles` produces, so this arm carries + // no re-quantisation error at all and its bounds are the F32 arm's, not + // the Q8_0 arm's. That is the assertion below, and it is the point of + // running it here rather than trusting the kernel-level gate: if the store + // or the executor wiring lost the exactness that G7 proved at the GEMV, + // this arm lands on the Q8_0 arm's numbers instead of the F32 arm's. + let native_arm = run_arm( + &mut gpu, + &hfq, + layer_prefix, + &shared, + EschaWeightStore::Native, + &x, + &want, + &ids, + &scores, + n_tok, + top_k, + n_exp, + hidden, + mi, + smi, + ); + println!( + "MoE block [Native code, fused GEMV]: max|diff|={:.3e} mean|diff|={:.3e}", + native_arm.max_abs, native_arm.mean_abs + ); + + // ── Indexed route == host route, bit-for-bit ───────────────────────── + // Production decode/prefill run `escha_routed_decode_indexed`. The + // tolerances asserted below were measured on the host-routed executor, so + // they only describe production if the two routes are the SAME numbers — + // which they are by construction (identical phases, identical GEMV + // accumulate order, the f16 score rounding moved from host to kernel) and + // must therefore be provable exactly, not to a tolerance. + let n_indexed = q8_arm + .indexed_mismatches + .expect("the Q8_0 arm must cross-check the indexed route"); + println!("indexed route vs host route: {n_indexed} differing floats (want 0)"); + assert_eq!( + n_indexed, 0, + "the indexed (GPU-top-K) escha route disagreed with the CPU-top-K route on \ + {n_indexed} routed-output floats. These must be BIT-identical: same phases, same \ + weights, same H128 pair, same GEMV accumulate order. Look at (1) the wide-vs-narrow \ + Q8_0 kernel choice in `escha_gemv_q8_0_moe_k8_indexed_batched` — it must reuse \ + `gemv_q8_0`'s `k <= 1536` rule, because the wide kernel folds four interleaved \ + accumulators and the narrow one a single sum; (2) `escha_round_weights_f16_rne` \ + against the host's `half::f16::from_f32`; (3) the per-slot x/y strides. Until this is \ + 0 the tolerances below describe a route production does not take." + ); + assert!( + f32_arm.indexed_mismatches.is_none(), + "the F32 control arm has no indexed counterpart (each indexed GEMV decodes one specific \ + container) — a Some here means the cross-check silently ran against the wrong container" + ); + assert!( + native_arm.indexed_mismatches.is_none(), + "the Native arm has no HOST counterpart (there is no per-expert native GEMV) — a Some \ + here means the host route ran on trellis code, which `GemvFamily::run_auto` is supposed \ + to refuse outright" + ); + + // ── Batched-prefill route == per-token indexed route, bit-for-bit ──── + // The routed half of the §5.4 batched-prefill gate. The dense half of a + // batched prefill is NOT bit-identical (a batched WMMA GEMM does not + // accumulate like a batch-1 GEMV) — but the ROUTED half must be, because + // every kernel in the escha pipeline is purely slot-parallel and slot `s` + // does the same FLOPs in the same order at 8 slots as at 2 048. Asserting + // equality here is what makes the whole-model logit delta attributable to + // the dense half alone. + let n_batched = q8_arm + .batched_mismatches + .expect("the Q8_0 arm must cross-check the batched-prefill route"); + println!( + "batched prefill route vs per-token indexed route: {n_batched} differing floats (want 0)" + ); + assert_eq!( + n_batched, 0, + "the batched-prefill escha route disagreed with the per-token indexed route on \ + {n_batched} routed-output floats. These must be BIT-identical — only `slots` \ + changes. Look at (1) `escha_h128_in_batched`'s `x_group`: the gate_up input side \ + must be `Grouped(k)` (slot s reads token s/k), not `Broadcast` (every slot reads \ + token 0) and not `PerSlot`; (2) the token-major slot layout, which must match what \ + `moe_topk_renorm_k8_batched` writes and `moe_down_combine_k8_batched` reads; \ + (3) the scratch views, which must be cut to exactly `n_tok * k` slots." + ); + assert!( + f32_arm.batched_mismatches.is_none(), + "the F32 control arm has no batched counterpart (each indexed GEMV decodes one specific \ + container)" + ); + // The same slot-parallel invariance has to hold for the fused kernels. It + // is not free: they are the only escha GEMVs whose BLOCK spans 16 output + // rows, so a `blockIdx`/slot mix-up would show up here and nowhere else. + let n_batched_native = native_arm + .batched_mismatches + .expect("the Native arm must cross-check the batched-prefill route"); + println!( + "batched prefill route vs per-token indexed route [Native]: {n_batched_native} differing \ + floats (want 0)" + ); + assert_eq!( + n_batched_native, 0, + "the batched-prefill escha route disagreed with the per-token indexed route on \ + {n_batched_native} routed-output floats with the fused native GEMV. Same three \ + suspects as the Q8_0 case, plus one specific to these kernels: their grid is \ + (m/16, slots) with a 512-thread block, so check that `blockIdx.y` is still the slot \ + and that `m % 16 == 0` held." + ); + + let dq: Vec = q8_arm + .got + .iter() + .zip(f32_arm.got.iter()) + .map(|(a, b)| (a - b).abs()) + .collect(); + let dq_max = dq.iter().cloned().fold(0.0f32, f32::max); + let dq_mean = dq.iter().sum::() / dq.len() as f32; + println!("Q8_0 re-quantisation cost (arm2 - arm1): max={dq_max:.3e} mean={dq_mean:.3e}"); + + // ── Launch budget ──────────────────────────────────────────────────── + // Measured, not asserted from a comment: the counter in + // `Gpu::escha_h128_batched` ticks once per batched transform launch. + let per_layer_per_token = launches_one_layer_all_tokens as f64 / n_tok as f64; + let per_token = per_layer_per_token * config.n_layers as f64; + println!( + "H128 launches: {per_layer_per_token} per (layer, token) -> {per_token} per token at \ + {} layers (a per-expert wiring would be {})", + config.n_layers, + 4 * top_k * config.n_layers + ); + assert_eq!( + per_token as usize, + escha_launches_per_token(config.n_layers), + "H128 launch budget drifted from the batched contract" + ); + + // ── Bounds ─────────────────────────────────────────────────────────── + // Arm 1 (weight-exact) carries the brief's measured tolerance: this is + // the arm that says "the wiring is right". A missing H128 pair lands at + // ~1e-1 here, three orders of magnitude outside it. + assert!( + f32_arm.max_abs <= 2e-4, + "F32 arm max|diff| {:.3e} exceeds 2e-4 — with weight-exact experts this can only be a \ + wiring defect (transpose orientation, H128 placement/side, SwiGLU half order, or the \ + f16(score) combine). If it is ~1e-1 the H128 pair is not being applied at all: check \ + that the escha dtypes did not reach a Plain GEMV.", + f32_arm.max_abs + ); + // The mean bound is 1.2e-5, NOT the brief's 1e-5. Derivation, because this + // number is otherwise a knife-edge that would misdiagnose: + // + // measured (deterministic, no sampling) 9.673e-6 + // sensitivity to the SHARED expert's rounding contract 7.32e-7 + // (report 5.2 CPU sweep: 8.942e-6 -> 9.674e-6 mean when only the + // shared expert's rounding flips) + // bound = 9.673e-6 + 3 x 7.32e-7 = 1.186e-5, rounded to 1.2e-5 + // + // The brief's 1e-5 left 3.3% headroom = 0.45x that single sensitivity. The + // shared expert is arch-6 code Task 10 does not own: any benign change to + // `silu_mul_f32`, `sigmoid_f32`, the shared-down GEMV selection, or the HIP + // compiler could move the metric further than the entire remaining margin + // and trip an assert whose message blames escha WIRING. That is a false + // diagnosis sending the next person after a bug that does not exist. + // + // Option (a) from the review — exclude the shared expert from the + // comparison — was considered and rejected as not practical: the shipped + // golden `moeblk_out.f16` is `routed + shared` and no routed-only golden + // exists. Subtracting hipfire's OWN shared output from both sides is + // algebraically a no-op on the diff: + // (r_h + s_h) - (r_e + s_e) == (r_h + s_h - s_h) - (r_e + s_e - s_h) + // so the shared expert's Metal-vs-hipfire divergence stays in the measured + // quantity either way. Widening with the derivation written down is the + // honest option; 1.2e-5 is still ~4 orders below the ~1e-1 a missing H128 + // pair produces, so the gate keeps all of its diagnostic power. + assert!( + f32_arm.mean_abs <= 1.2e-5, + "F32 arm mean|diff| {:.3e} exceeds 1.2e-5. NOTE before you go hunting escha wiring: the \ + SHARED expert is inside this measured quantity (the golden is routed + shared, and no \ + routed-only golden ships), and it is arch-6 code Task 10 does not own. A move of order \ + 1e-6 is consistent with a change to silu_mul_f32 / sigmoid_f32 / the shared-down GEMV \ + selection / the HIP compiler, NOT with an escha defect. Only a jump of 1e-4 or more — \ + and especially ~1e-1 — indicates the escha wiring (transpose orientation, H128 \ + placement/side, SwiGLU half order, the f16(score) combine, or the H128 pair not being \ + applied at all). Check the max|diff| assert above and the two bit-exact codec gates \ + (test_escha_decode_gpu_vs_cpu, test_escha_h128_gpu_vs_cpu) first.", + f32_arm.mean_abs + ); + + // Arm 2 adds the 8-bit re-quantisation on top. That cost is real, + // irreducible at this storage format, and MEASURED — see the report for + // the CPU-side derivation that agrees with it. The bound below is set + // from that measurement with headroom, and its job is to catch a + // REGRESSION in the quantiser (a wrong block axis, a dropped clamp, a + // truncating instead of RNE scale), not to re-prove the wiring. + assert!( + q8_arm.max_abs <= 4e-4, + "Q8_0 arm max|diff| {:.3e} exceeds 4e-4", + q8_arm.max_abs + ); + assert!( + q8_arm.mean_abs <= 6e-5, + "Q8_0 arm mean|diff| {:.3e} exceeds 6e-5", + q8_arm.mean_abs + ); + assert!( + dq_mean <= 5e-5, + "Q8_0 re-quantisation cost {dq_mean:.3e} exceeds 5e-5 — the quantiser regressed" + ); + + // ── Arm 3 bounds: weight-exact, so the F32 arm's bounds ────────────── + // + // Deliberately NOT the Q8_0 arm's looser 4e-4 / 6e-5. The fused GEMV + // decodes to the same fp16 the F32 store holds, so the only thing between + // this arm and the F32 arm is f32 summation ORDER (the indexed kernels' + // lane-strided partials against `gemv_f32`'s), which moves the last bit — + // not the fourth digit. Holding it to the F32 bounds is what makes "the + // Q8_0 re-quantisation error is gone" a checked claim rather than a hope. + assert!( + native_arm.max_abs <= 2e-4, + "Native arm max|diff| {:.3e} exceeds 2e-4 (the WEIGHT-EXACT bound). If this sits near \ + the Q8_0 arm's 2.6e-4 instead, the layer is not running the fused native GEMV at all \ + — check that the store resolved to Native, that \ + `MoeResolution::routed_indexable_escha_native` admitted the layer, and that \ + `escha_routed_gemv` dispatched on the escha dtype rather than falling through to the \ + Q8_0 arm.", + native_arm.max_abs + ); + assert!( + native_arm.mean_abs <= 1.2e-5, + "Native arm mean|diff| {:.3e} exceeds 1.2e-5 (the WEIGHT-EXACT bound; see the F32 arm's \ + derivation of that figure)", + native_arm.mean_abs + ); + // And state the relationship directly rather than leaving it to two + // separate bounds: the fused arm must be no worse than the weight-exact + // control, and strictly better than the re-quantised one. + println!( + "weight-exactness: native mean {:.3e} vs F32 {:.3e} vs Q8_0 {:.3e}", + native_arm.mean_abs, f32_arm.mean_abs, q8_arm.mean_abs + ); + assert!( + native_arm.mean_abs < q8_arm.mean_abs, + "the Native arm's mean|diff| {:.3e} is not better than the Q8_0 arm's {:.3e}. The fused \ + GEMV uses exactly-decoded weights and Q8_0 re-quantises them, so this can only mean \ + the Native arm did not actually run the fused path.", + native_arm.mean_abs, + q8_arm.mean_abs + ); + + println!("G4 PASS"); +} diff --git a/crates/hipfire-arch-qwen35/examples/escha_prefill_batch_gate.rs b/crates/hipfire-arch-qwen35/examples/escha_prefill_batch_gate.rs new file mode 100644 index 0000000000..85c688aa03 --- /dev/null +++ b/crates/hipfire-arch-qwen35/examples/escha_prefill_batch_gate.rs @@ -0,0 +1,508 @@ +// SPDX-License-Identifier: Apache-2.0 +// hipfire — see LICENSE and NOTICE in the project root. +//! G6: batched Escha-W2 prefill against the per-token route, whole model. +//! +//! This is the gate specified in task-perf-2 §5.4. It runs ONE prompt through +//! the model twice — once through `forward_prefill_batch` (batched) and once +//! through the per-token `forward` loop the non-batched route is byte-identical +//! to — and reports what changed. +//! +//! # What must be equal, what must not be, and why +//! +//! The routed (escha) half is asserted EQUAL, bit-for-bit, by +//! `escha_moe_block_gate` (G4): `escha_routed_prefill_indexed` differs from +//! `escha_routed_decode_indexed` only in `slots`, and every kernel in that +//! pipeline is purely slot-parallel — slot `s` performs the same FLOPs in the +//! same order whether the launch carried 8 slots or 2 048. That is a real +//! equality claim and it is asserted as equality there, not as a tolerance. +//! +//! The DENSE half is NOT bit-identical and cannot be. Batched prefill runs the +//! attention projections, the router and the shared expert as batched WMMA +//! GEMMs (`gemm_q8_0_residual_wmma`, `gemm_f16_wmma_mb8`, ...) where the +//! per-token route runs batch-1 GEMVs. A 16x16x16 WMMA tile accumulates a +//! different K-order than a warp-reduced GEMV, so the two disagree at the +//! last bits of every dot product and the difference compounds across 40 +//! layers. That is a summation-order difference, not a defect. +//! +//! So the whole-model claim this gate can make is: the final-token logits +//! agree to a MEASURED bound, and the ARGMAX does not move. +//! +//! The measured values at n=64 on escha-35b are `max|delta| = 4.393e-1` and +//! `mean|delta| = 7.160e-2`, and both are now asserted rather than printed. +//! Note that these are two orders of magnitude ABOVE what pure accumulation +//! reordering would give: the dominant term is not reordering at all but the +//! expert-selection divergence documented below, which is 24.1% of (token, +//! layer) decisions over the whole stack. A token whose expert set differs +//! computes a different hidden state, and that compounds with depth. An +//! earlier version of this file quoted "~1e-3" here; that figure describes +//! two arms with identical routing, which these are not. +//! +//! **The argmax assertion is the load-bearing one.** Do not relax this gate to +//! "max delta < tol" and stop there. A dropped H128 transform, a stale +//! activation cache, or a wrong expert-slot stride all produce finite, fluent +//! output that is wrong by ~1e-1 per element — and a tolerance chosen loosely +//! enough to absorb the real divergence can absorb those too. The argmax +//! moving is the signal. (When this port's stale-FP16-activation bug was live, +//! this gate's max delta was ~1e+1 and the argmax moved from 25760 to 220.) +//! +//! # The expert-selection divergence +//! +//! Escha's router logits pass through `router_logits_round_f16_rne` on BOTH +//! routes (that is what makes the comparison meaningful at all), but the +//! logits fed to it differ, so some decisions land on the other side of an f16 +//! rounding boundary and the two routes legitimately select different experts. +//! That is a property of the format, not a defect — but it has to be MEASURED, +//! which is what this gate does: with `HIPFIRE_ESCHA_ROUTE_TRACE` set, both +//! arms record the `topk_indices` every MoE layer actually indexed with. +//! +//! The measured rate is **2.96-3.13% per expert slot** (24.1% of (token, +//! layer) SETS over the whole stack, 0.00% at layer 0). Task 9's design-time +//! estimate was ~0.42%, i.e. 8x low, and it was low because it modelled the +//! wrong term: the dominant perturbation is not f32 accumulation reordering +//! but the **f16 downcast** the batched dense half applies to its activations +//! for the WMMA GEMMs, which the per-token F32 GEMV route does not. See the +//! design doc §10.5(b). Do not re-quote 0.42% anywhere. +//! +//! # State reset between the arms +//! +//! The DeltaNet recurrent state is the only thing that carries between the two +//! arms, and it must not: arm B would otherwise start from the state arm A +//! left behind and the comparison would be meaningless. A fresh +//! `DeltaNetState` is allocated for each arm rather than zeroing in place, so +//! "reset" means exactly what a fresh process means. The KV cache needs no +//! reset — arm B prefills from position 0 and rewrites the same slots, and +//! attention only reads positions below its own `start_pos + n`. +//! +//! COST: loads the whole model, ~37.6 GB resident. +//! +//! Run: +//! cargo run --release -p hipfire-arch-qwen35 \ +//! --example escha_prefill_batch_gate -- /data/hipfire-models/escha-35b.hfq [n] + +use hipfire_arch_qwen35::qwen35; +use hipfire_arch_qwen35::qwen35::DeltaNetState; +use hipfire_runtime::hfq::HfqFile; +use hipfire_runtime::loader_api::{CaskConfig, LoadCtx, ModelSource, SpecLoadCfg}; +use rdna_compute::Gpu; +use std::path::Path; + +/// Same deterministic filler as `escha_prefill_bench`, for the same reason: a +/// constant token would route every position to the same experts and would +/// understate both the routed work and the selection divergence. +fn prompt_of(n: usize, vocab: usize) -> Vec { + let mut s: u64 = 0x2545_F491_4F6C_DD1D; + (0..n) + .map(|_| { + s ^= s << 13; + s ^= s >> 7; + s ^= s << 17; + (s % vocab as u64) as u32 + }) + .collect() +} + +/// One trace record: `n_tokens` rows of `k` expert ids, token-major. +struct TraceRec { + n_tokens: usize, + k: usize, + ids: Vec, +} + +fn read_trace(path: &str) -> Result, String> { + let data = std::fs::read(path).map_err(|e| format!("read {path}: {e}"))?; + let mut out = Vec::new(); + let mut off = 0usize; + while off + 8 <= data.len() { + let n = u32::from_le_bytes(data[off..off + 4].try_into().unwrap()) as usize; + let k = u32::from_le_bytes(data[off + 4..off + 8].try_into().unwrap()) as usize; + off += 8; + let want = n * k * 4; + if off + want > data.len() { + return Err(format!("{path}: truncated record")); + } + let ids = data[off..off + want] + .chunks_exact(4) + .map(|c| i32::from_le_bytes(c.try_into().unwrap())) + .collect(); + off += want; + out.push(TraceRec { + n_tokens: n, + k, + ids, + }); + } + if off != data.len() { + return Err(format!("{path}: trailing bytes")); + } + Ok(out) +} + +/// Flatten a trace into `(token, layer) -> sorted expert set`, given the number +/// of MoE layers. Every MoE layer emits exactly one record per forward, so +/// record order determines the grid; the caller checks the totals. +fn flatten( + recs: &[TraceRec], + n_layers_moe: usize, + n_tokens: usize, +) -> Result>, String> { + let mut out: Vec> = vec![Vec::new(); n_tokens * n_layers_moe]; + let mut layer = 0usize; + let mut token_base = 0usize; + for r in recs { + for t in 0..r.n_tokens { + let tok = token_base + t; + if tok >= n_tokens { + return Err("trace covers more tokens than the prompt".into()); + } + let mut ids: Vec = r.ids[t * r.k..(t + 1) * r.k].to_vec(); + ids.sort_unstable(); + out[tok * n_layers_moe + layer] = ids; + } + layer += 1; + if layer == n_layers_moe { + layer = 0; + token_base += r.n_tokens; + } + } + if token_base != n_tokens { + return Err(format!( + "trace covered {token_base} tokens, expected {n_tokens} — record order does not \ + match the assumed (chunk, layer) grid" + )); + } + if out.iter().any(|v| v.is_empty()) { + return Err("trace left a (token, layer) cell unfilled".into()); + } + Ok(out) +} + +fn main() -> Result<(), String> { + let path = std::env::args() + .nth(1) + .unwrap_or_else(|| "/data/hipfire-models/escha-35b.hfq".to_string()); + let n: usize = std::env::args() + .nth(2) + .unwrap_or_else(|| "64".to_string()) + .trim() + .parse() + .expect("prefill length"); + + let trace_dir = std::env::var("TMPDIR").unwrap_or_else(|_| "/tmp".to_string()); + let trace_a = format!("{trace_dir}/escha-route-batched.bin"); + let trace_b = format!("{trace_dir}/escha-route-pertoken.bin"); + + let hfq = HfqFile::open(Path::new(&path)).map_err(|e| format!("open: {e:?}"))?; + let mut gpu = Gpu::init().map_err(|e| format!("gpu: {e:?}"))?; + let cask = CaskConfig::default(); + let src = ModelSource::Hfq(hfq); + let mut ctx = LoadCtx { + path: &path, + max_seq: n + 64, + deepseek4_compute_placement: Default::default(), + deepseek4_experts_per_token: None, + draft_path: None, + kv_mode_override: None, + kv_backend: hipfire_runtime::kv_backend::KvBackend::Contiguous, + kv_adaptive_override: None, + // `ESCHA_GATE_STATE_QUANT=f32` runs both arms with an FP32 DeltaNet + // state. That is an ATTRIBUTION lever, not a mode: the Q8 recurrent + // state is requantised after every token on both routes, but the + // batched GDN kernel's requant frame is documented as + // "distributionally equivalent to decode, not byte-identical", so some + // of the prefill-vs-decode divergence measured below predates any of + // this and belongs to that kernel rather than to the batched dense + // half. Running FP32 removes the requant and shows how much. + state_quant_override: std::env::var("ESCHA_GATE_STATE_QUANT") + .ok() + .map(|_| "f32") + .filter(|_| std::env::var("ESCHA_GATE_STATE_QUANT").as_deref() == Ok("f32")), + cask: &cask, + pp: 1, + spec: SpecLoadCfg::default(), + gpu: &mut gpu, + gemma4_drafter_path: None, + gemma4_draft_len: 3, + }; + let mut b = hipfire_arch_qwen35::load_qwen35_bundle(src, &mut ctx)?; + let vocab = b.config.vocab_size; + let prompt = prompt_of(n, vocab); + let n_moe_layers = b + .weights + .layers + .iter() + .filter(|l| { + matches!( + l, + qwen35::LayerWeights::DeltaNetMoe(_) | qwen35::LayerWeights::FullAttnMoe(_) + ) + }) + .count(); + println!("tokens={n} vocab={vocab} moe_layers={n_moe_layers}"); + + // ── Arm A: batched ─────────────────────────────────────────────────── + // + // The H128 launch count identifies the route with no ambiguity, which + // matters because `forward_prefill_batch` falls back to the per-token loop + // SILENTLY when a layer is inadmissible. Without this check a refused + // model would compare the per-token route against itself and pass a gate + // that proves nothing. + // Open the trace explicitly rather than via HIPFIRE_ESCHA_ROUTE_TRACE: + // `developer_var` reads a config snapshot resolved at process start, so a + // `set_var` here would never be seen and the gate would silently compare + // two empty traces and report 0% divergence. + hipfire_dispatch::pipeline::route_trace::reopen(&trace_a); + let before = rdna_compute::escha_h128_launches(); + qwen35::forward_prefill_batch( + ctx.gpu, + &b.weights, + &b.config, + &prompt, + 0, + &mut b.kv_cache, + &mut b.dn_state, + &b.scratch, + None, + None, + None, + None, + ) + .map_err(|e| format!("batched prefill: {e:?}"))?; + ctx.gpu + .hip + .device_synchronize() + .map_err(|e| format!("sync: {e:?}"))?; + let batched_launches = rdna_compute::escha_h128_launches() - before; + let logits_a = ctx + .gpu + .download_f32(&b.scratch.logits) + .map_err(|e| format!("download A: {e:?}"))?; + + let per_token_budget = + hipfire_dispatch::pipeline::escha::escha_launches_per_token(b.config.n_layers); + println!( + "arm A (batched): H128 launches={batched_launches} \ + (per-token route would be {})", + per_token_budget * n + ); + if batched_launches >= (per_token_budget * n) as u64 { + return Err(format!( + "arm A issued {batched_launches} H128 launches, which is the PER-TOKEN budget \ + ({per_token_budget} x {n} tokens) — `forward_prefill_batch` fell back to the \ + per-token loop, so this gate would be comparing that route against itself. Run \ + with HIPFIRE_DEBUG_BATCH=1 to see which layer refused." + )); + } + + // ── Arm B: per-token ───────────────────────────────────────────────── + // + // Fresh recurrent state; see the module docs for why the KV cache does not + // need one. + let fresh = DeltaNetState::new_with_quant(ctx.gpu, &b.config, b.dn_state.quant) + .map_err(|e| format!("fresh dn_state: {e:?}"))?; + let stale = std::mem::replace(&mut b.dn_state, fresh); + stale.free_gpu(ctx.gpu); + + // The trace sink is initialised on first use — inside arm A's first MoE + // layer — so setting the env var again here would be ignored. Redirect it + // explicitly. + hipfire_dispatch::pipeline::route_trace::reopen(&trace_b); + let mut logits_b = Vec::new(); + for (i, &tok) in prompt.iter().enumerate() { + logits_b = qwen35::forward( + ctx.gpu, + &b.weights, + &b.config, + tok, + i, + &mut b.kv_cache, + &mut b.dn_state, + ) + .map_err(|e| format!("per-token forward at {i}: {e:?}"))?; + } + + // ── Final-token logits ─────────────────────────────────────────────── + let (mut max_d, mut sum_d) = (0.0f32, 0.0f64); + for i in 0..vocab { + let d = (logits_a[i] - logits_b[i]).abs(); + if d > max_d { + max_d = d; + } + sum_d += d as f64; + } + let mean_d = sum_d / vocab as f64; + let am = |v: &[f32]| { + v.iter() + .take(vocab) + .enumerate() + .fold((0usize, f32::NEG_INFINITY), |(bi, bv), (i, &x)| { + if x > bv { + (i, x) + } else { + (bi, bv) + } + }) + }; + let (arg_a, best_a) = am(&logits_a); + let (arg_b, best_b) = am(&logits_b); + println!("final-token logits: max|delta|={max_d:.3e} mean|delta|={mean_d:.3e}"); + println!("argmax: batched={arg_a} ({best_a:.4}) per-token={arg_b} ({best_b:.4})"); + let nonfinite = logits_a + .iter() + .take(vocab) + .filter(|v| !v.is_finite()) + .count(); + println!("non-finite logits (batched): {nonfinite}"); + assert_eq!(nonfinite, 0, "batched prefill produced non-finite logits"); + + // ── Expert-selection divergence ────────────────────────────────────── + // + // Reported THREE ways, because one number would be misleading. + // + // * LAYER 0 is where the two routes have accumulated the LEAST difference: + // the embedding they start from is identical. It is NOT a matched-input + // measurement — by the time layer 0's router runs, the hidden state has + // already been through the batched attention projections AND the batched + // GDN recurrence, whose Q8 state requant is documented as + // "distributionally equivalent to decode, not byte-identical". So layer + // 0's rate is the floor of the whole effect, not the f16-boundary rate + // in isolation. Run with `ESCHA_GATE_STATE_QUANT=f32` to remove the + // requant term and see what is left. This is still the number to assert + // on, because it is the one that cannot have compounded. + // * Layers 1.. see inputs that already differ, because a flip at layer L + // changes that token's hidden state for every later layer. The rate + // therefore COMPOUNDS with depth and plateaus. That is arithmetic, not + // a defect, and it would happen on any model whose batched prefill is + // not bit-identical to its decode. + // * The per-EXPERT-SLOT rate says how much a differing set differs. A + // boundary straddle swaps the 8th expert for the 9th, so it should be + // ~1/k of the set rate; a much larger ratio would mean the routes are + // choosing genuinely different experts, not neighbouring ones. + let ra = read_trace(&trace_a)?; + let rb = read_trace(&trace_b)?; + let fa = flatten(&ra, n_moe_layers, n)?; + let fb = flatten(&rb, n_moe_layers, n)?; + let total = fa.len(); + let differing = fa.iter().zip(fb.iter()).filter(|(a, c)| a != c).count(); + let pct = 100.0 * differing as f64 / total as f64; + + let mut per_layer = vec![0usize; n_moe_layers]; + let (mut slots_total, mut slots_diff) = (0usize, 0usize); + for tok in 0..n { + for l in 0..n_moe_layers { + let a = &fa[tok * n_moe_layers + l]; + let c = &fb[tok * n_moe_layers + l]; + if a != c { + per_layer[l] += 1; + } + slots_total += a.len(); + slots_diff += a.iter().filter(|e| !c.contains(e)).count(); + } + } + let l0_pct = 100.0 * per_layer[0] as f64 / n as f64; + let slot_pct = 100.0 * slots_diff as f64 / slots_total as f64; + println!( + "expert-selection divergence, layer 0 (least-accumulated): {}/{n} = {l0_pct:.4}%", + per_layer[0] + ); + println!( + "expert-selection divergence, whole stack: {differing} of {total} \ + (token, layer) decisions = {pct:.4}%" + ); + println!("expert-slot divergence, whole stack: {slots_diff} of {slots_total} = {slot_pct:.4}%"); + print!("per-layer set-divergence %:"); + for l in 0..n_moe_layers { + print!(" {:.1}", 100.0 * per_layer[l] as f64 / n as f64); + } + println!(); + + // Layer 0 is the assertion. Both routes feed it the SAME embedding, so a + // flip there can only come from this layer's own accumulation reordering + // landing on an f16 boundary. Measured on this model at n=64: 0.00% — + // layer 0 does not flip at all. The 10% bound leaves room for sampling + // noise at small `n` while still failing loudly if the dense half is wrong + // (when the stale-FP16-activation bug was live this was well above it). + // The whole-stack rate is a different quantity and is much higher (24.1% + // of sets, 2.96-3.13% of slots) because it compounds; see §10.5(b). + assert!( + l0_pct < 10.0, + "layer-0 expert-selection divergence {l0_pct:.4}% is far above the few percent \ + expected from one layer of batched-vs-GEMV dense arithmetic plus the batched GDN \ + requant frame. Layer 0 cannot have compounded — it is the first layer — so a large \ + rate here means the batched dense half is computing materially different router \ + logits, not that a few decisions straddled an f16 boundary." + ); + // A differing set should differ by about one expert in k (the boundary + // swap). Much more than that is not a boundary effect. + assert!( + slot_pct * (2.0 * ra[0].k as f64) < pct.max(1e-9) * 3.0 + 5.0, + "sets that differ are differing by {slot_pct:.3}% of slots against a {pct:.3}% set \ + rate — a boundary straddle swaps ONE expert, so the slot rate should be near the \ + set rate divided by k. The routes are picking genuinely different experts." + ); + + // ── The logit-delta bound ──────────────────────────────────────────── + // + // Until now `max_d` and `mean_d` were computed and only PRINTED, so the + // bound the module docs call load-bearing was enforced by a human reading + // stdout. That is not enforcement, and the defect this gate exists to + // catch produced ~1e+1 in exactly this quantity. + // + // THE BOUNDS BELOW ARE MEASURED, not derived from the "~1e-3" figure this + // file used to quote. At n=64 on this model the actual values are + // + // max|delta| = 4.393e-1 mean|delta| = 7.160e-2 + // + // reproduced identically across builds. The old ~1e-3 estimate described + // pure accumulation reordering with IDENTICAL routing on both arms, and + // that is not what these two arms do: 24.1% of (token, layer) routing + // decisions differ between them (see the divergence report above). A + // token routed to a different expert set computes a genuinely different + // hidden state, and 40 layers of that lands two orders of magnitude above + // the reordering-only figure. The estimate was wrong, not the run. + // + // So the headroom here is real but narrow: 4.4e-1 measured against the + // ~1e+1 the stale-FP16-activation bug produced is a factor of ~23, and the + // bound has to sit inside it. 2.0 is ~4.6x above the measurement and ~5x + // below the known-bad value — the honest split. The mean is the steadier + // statistic (it is an average over 248 320 logits rather than one extreme + // order statistic), so it is bounded more tightly at ~7x headroom. + // + // This bound is NOT what makes the gate work; the argmax assertion below + // is. What it adds is a trip-wire for a structural error that happens not + // to move the argmax on this one prompt. + const MAX_ABS_LOGIT_DELTA: f32 = 2.0; + const MAX_MEAN_LOGIT_DELTA: f64 = 5e-1; + assert!( + max_d < MAX_ABS_LOGIT_DELTA, + "final-token logit max|delta| {max_d:.3e} exceeds {MAX_ABS_LOGIT_DELTA:.1}. Measured \ + on this model at n=64: 4.393e-1, from f32 accumulation reordering compounded by the \ + 24% prefill-vs-decode expert-selection divergence reported above. A value near 1e+1 \ + is the signature of a structural error — a dropped H128 transform, a stale \ + activation-conversion cache, or a wrong expert-slot stride — every one of which \ + stays finite and fluent. Check the divergence percentages above first: if THEY are \ + unchanged and this moved, the dense half changed." + ); + assert!( + mean_d < MAX_MEAN_LOGIT_DELTA, + "final-token logit mean|delta| {mean_d:.3e} exceeds {MAX_MEAN_LOGIT_DELTA:.1e}. \ + Measured on this model at n=64: 7.160e-2. The mean is averaged over the whole \ + vocabulary, so unlike max|delta| it does not move on a single outlier logit — a \ + mean this far out means the two routes disagree broadly, not at one index." + ); + + // The load-bearing assertion. See the module docs: a tolerance alone would + // absorb the failure class this port keeps catching. + assert_eq!( + arg_a, arg_b, + "batched prefill chose a different next token than the per-token route \ + (batched {arg_a} @ {best_a}, per-token {arg_b} @ {best_b}). The dense half is not \ + bit-identical by design (batched WMMA vs batch-1 GEMV accumulation order) and the \ + routing genuinely diverges on ~24% of decisions, which together move these logits by \ + ~4e-1 — but that must never move the argmax. A moved argmax means a structural \ + error — a dropped H128 transform, a stale activation-conversion cache, or a wrong \ + expert-slot stride — all of which stay finite and fluent." + ); + + println!("G6 PASS"); + Ok(()) +} diff --git a/crates/hipfire-arch-qwen35/examples/escha_prefill_bench.rs b/crates/hipfire-arch-qwen35/examples/escha_prefill_bench.rs new file mode 100644 index 0000000000..802362a0fc --- /dev/null +++ b/crates/hipfire-arch-qwen35/examples/escha_prefill_bench.rs @@ -0,0 +1,202 @@ +// SPDX-License-Identifier: Apache-2.0 +// Copyright (c) 2026 Kaden Schutt +// hipfire — see LICENSE and NOTICE in the project root. +//! Escha-W2 prefill/decode timing harness — the measurement behind the +//! Phase-2 prefill work. +//! +//! `escha_model_smoke` prefills 8 tokens because it is a *gate*: it exists to +//! prove the escha executor ran for every (layer, token) and the logits are +//! finite. Eight tokens is far too short to say anything about prefill THROUGHPUT +//! — the load and the first-launch JIT dominate. This harness prefills a +//! realistic prompt (512 or 2048), reports tok/s, and reports the +//! H128 launch count so the route taken is never in doubt: +//! +//! * `4 * n_layers` launches **per token** => the per-token `forward_scratch` +//! fallback (prefill doing decode's work once per token). +//! * `4 * n_layers` launches **per chunk** => a genuinely batched escha +//! prefill body. +//! +//! Prompt token ids are arbitrary but FIXED (a deterministic LCG), because the +//! point is a stable timing comparison, not semantics. Routing depends on the +//! activations, so a fixed prompt also keeps the expert-selection work +//! comparable between runs. +//! +//! COST: loads the whole model, ~37.6 GB resident. See `escha_model_smoke`. +//! +//! Run: +//! cargo run --release -p hipfire-arch-qwen35 \ +//! --example escha_prefill_bench -- /data/hipfire-models/escha-35b.hfq 512 [decode_tokens] +use hipfire_arch_qwen35::qwen35; +use hipfire_runtime::hfq::HfqFile; +use hipfire_runtime::loader_api::{CaskConfig, LoadCtx, ModelSource, SpecLoadCfg}; +use rdna_compute::Gpu; +use std::path::Path; + +/// Deterministic filler ids. A fixed sequence, not a constant one: a constant +/// token would route every position to the same experts and understate the +/// routed work. +fn prompt_of(n: usize, vocab: usize) -> Vec { + let mut s: u64 = 0x2545_F491_4F6C_DD1D; + (0..n) + .map(|_| { + s ^= s << 13; + s ^= s >> 7; + s ^= s << 17; + (s % vocab as u64) as u32 + }) + .collect() +} + +fn main() -> Result<(), String> { + let path = std::env::args() + .nth(1) + .unwrap_or_else(|| "/data/hipfire-models/escha-35b.hfq".to_string()); + // ONE length per process, on purpose. Prefill mutates the KV cache and the + // DeltaNet recurrent state, and a second prefill in the same process would + // either continue that context (so its attention cost is not comparable) or + // need a state reset whose completeness is one more thing to be wrong + // about. A fresh process is the cheap, unambiguous control. + let n: usize = std::env::args() + .nth(2) + .unwrap_or_else(|| "512".to_string()) + .trim() + .parse() + .expect("prefill length"); + let decode_tokens: usize = std::env::args() + .nth(3) + .map(|s| s.parse().expect("decode count")) + .unwrap_or(6); + + let hfq = HfqFile::open(Path::new(&path)).map_err(|e| format!("open: {e:?}"))?; + let mut gpu = Gpu::init().map_err(|e| format!("gpu: {e:?}"))?; + let cask = CaskConfig::default(); + let src = ModelSource::Hfq(hfq); + let mut ctx = LoadCtx { + path: &path, + // Room for the longest prefill plus the decode tail. + max_seq: n + decode_tokens + 64, + deepseek4_compute_placement: Default::default(), + deepseek4_experts_per_token: None, + draft_path: None, + kv_mode_override: None, + kv_backend: hipfire_runtime::kv_backend::KvBackend::Contiguous, + kv_adaptive_override: None, + state_quant_override: None, + cask: &cask, + pp: 1, + spec: SpecLoadCfg::default(), + gpu: &mut gpu, + gemma4_drafter_path: None, + gemma4_draft_len: 3, + }; + let t0 = std::time::Instant::now(); + let mut b = hipfire_arch_qwen35::load_qwen35_bundle(src, &mut ctx)?; + eprintln!("loaded in {:?}", t0.elapsed()); + + let per_token_budget = + hipfire_dispatch::pipeline::escha::escha_launches_per_token(b.config.n_layers); + eprintln!( + "n_layers={} vocab={} per-token H128 budget={per_token_budget}", + b.config.n_layers, b.config.vocab_size + ); + + { + let prompt = prompt_of(n, b.config.vocab_size); + let before = rdna_compute::escha_h128_launches(); + let t = std::time::Instant::now(); + qwen35::forward_prefill_batch( + ctx.gpu, + &b.weights, + &b.config, + &prompt, + 0, + &mut b.kv_cache, + &mut b.dn_state, + &b.scratch, + None, + None, + None, + None, + ) + .map_err(|e| format!("prefill {n}: {e:?}"))?; + ctx.gpu + .hip + .device_synchronize() + .map_err(|e| format!("sync: {e:?}"))?; + let wall = t.elapsed(); + let launches = rdna_compute::escha_h128_launches() - before; + let logits = ctx + .gpu + .download_f32(&b.scratch.logits) + .map_err(|e| format!("download logits: {e:?}"))?; + let bad = logits + .iter() + .take(b.config.vocab_size) + .filter(|v| !v.is_finite()) + .count(); + let (argmax, best) = logits.iter().take(b.config.vocab_size).enumerate().fold( + (0usize, f32::NEG_INFINITY), + |(bi, bv), (i, &v)| if v > bv { (i, v) } else { (bi, bv) }, + ); + eprintln!( + "PREFILL n={n}: {:.1} ms, {:.1} tok/s, {:.3} ms/token | H128 launches={launches} \ + ({:.1} per token) | non-finite={bad} argmax={argmax} ({best:.4})", + wall.as_secs_f64() * 1e3, + n as f64 / wall.as_secs_f64(), + wall.as_secs_f64() * 1e3 / n as f64, + launches as f64 / n as f64, + ); + if bad != 0 { + return Err(format!("non-finite logits after an {n}-token prefill")); + } + // Full final-token logits, for byte-level A/B between two runs of the + // SAME binary (e.g. `HIPFIRE_ESCHA_INDEXED=0` vs not, to check the two + // routed routes are bit-identical). An argmax and a 4-decimal top + // logit cannot establish bit-identity; this can. Nothing else in the + // harness depends on it. + if let Ok(dump) = std::env::var("HIPFIRE_BENCH_LOGITS_OUT") { + let mut bytes = Vec::with_capacity(b.config.vocab_size * 4); + for v in logits.iter().take(b.config.vocab_size) { + bytes.extend_from_slice(&v.to_le_bytes()); + } + std::fs::write(&dump, &bytes).map_err(|e| format!("write {dump}: {e}"))?; + eprintln!("wrote {} logits to {dump}", b.config.vocab_size); + } + + // Decode continuation, from the state this prefill just built. The + // first token is discarded: it pays any first-shape JIT. + let mut times = Vec::new(); + for i in 0..decode_tokens { + let tok = prompt[i % prompt.len()]; + let t = std::time::Instant::now(); + let logits = qwen35::forward( + ctx.gpu, + &b.weights, + &b.config, + tok, + n + i, + &mut b.kv_cache, + &mut b.dn_state, + ) + .map_err(|e| format!("decode: {e:?}"))?; + times.push(t.elapsed().as_secs_f64()); + if logits.iter().any(|v| !v.is_finite()) { + return Err(format!("non-finite decode logits at pos {}", n + i)); + } + } + let warm = ×[1..]; + let mean = warm.iter().sum::() / warm.len() as f64; + let lo = warm.iter().cloned().fold(f64::INFINITY, f64::min); + let hi = warm.iter().cloned().fold(0.0f64, f64::max); + eprintln!( + "DECODE after n={n}: {:.2} ms/token mean ({:.2}-{:.2}), {:.1} tok/s, n={}", + mean * 1e3, + lo * 1e3, + hi * 1e3, + 1.0 / mean, + warm.len() + ); + } + eprintln!("escha_prefill_bench: done"); + Ok(()) +} diff --git a/crates/hipfire-arch-qwen35/examples/escha_router_contract.rs b/crates/hipfire-arch-qwen35/examples/escha_router_contract.rs new file mode 100644 index 0000000000..9b64f38d71 --- /dev/null +++ b/crates/hipfire-arch-qwen35/examples/escha_router_contract.rs @@ -0,0 +1,66 @@ +// SPDX-License-Identifier: Apache-2.0 +// Copyright (c) 2026 Kaden Schutt +// hipfire — see LICENSE and NOTICE in the project root. + +//! G4b: hipfire's arch-6 router must select the same experts Escha does. +//! +//! Escha rounds router logits to f16 BEFORE top-k (`ref.py`: the logits are +//! computed as f16 then widened to f32 to select). Selecting on unrounded f32 +//! logits is a different function, and the rounding manufactures exact ties +//! that f32 never produces. +//! +//! Asserts the SET, not the order: the combine is a sum over slots, so intra-k +//! order cannot change the output. On the fixture, token 3 has two experts on +//! identical f16 logits (1.80078), and which one lands in which slot is +//! implementation-defined. +//! +//! Run: +//! cargo run --release -p hipfire-arch-qwen35 \ +//! --example escha_router_contract -- /data/hipfire-models/escha-35b.hfq +use std::collections::HashSet; +use std::path::PathBuf; + +fn fixture(name: &str) -> PathBuf { + PathBuf::from(env!("CARGO_MANIFEST_DIR")) + .join("../hipfire-quantize/tests/data/escha") + .join(name) +} + +fn read_f16_as_f32(name: &str) -> Vec { + let raw = std::fs::read(fixture(name)).expect("run fetch-goldens.sh first"); + raw.chunks_exact(2) + .map(|c| hipfire_quantize::float16::f16_to_f32(u16::from_le_bytes([c[0], c[1]]))) + .collect() +} + +fn main() { + let hfq = std::env::args().nth(1).expect("usage: "); + let x = read_f16_as_f32("moeblk_x.f16"); // [8, 2048] + let raw_ids = std::fs::read(fixture("moeblk_ids.i64")).unwrap(); + let want_ids: Vec = raw_ids + .chunks_exact(8) + .map(|c| i64::from_le_bytes(c.try_into().unwrap())) + .collect(); // [8, 8] + + // Call hipfire's real router for layer 0 on each token. + let got = hipfire_arch_qwen35::escha_router_topk_for_test(&hfq, 0, &x, 8, 2048, 8) + .expect("router"); + + let mut bad = 0usize; + for t in 0..8 { + let want: HashSet = want_ids[t * 8..(t + 1) * 8].iter().copied().collect(); + let mine: HashSet = got[t * 8..(t + 1) * 8].iter().map(|&v| v as i64).collect(); + if want != mine { + bad += 1; + println!("token {t}: escha={:?}", &want_ids[t * 8..(t + 1) * 8]); + println!(" hipfire={:?}", &got[t * 8..(t + 1) * 8]); + } + } + println!("tokens with a differing top-8 SET: {bad}/8"); + assert_eq!( + bad, 0, + "arch-6 router selects different experts than escha. Most likely cause: \ + it is not rounding logits to f16 before top-k." + ); + println!("G4b PASS"); +} diff --git a/crates/hipfire-arch-qwen35/examples/test_escha_dense_linear_gpu_vs_cpu.rs b/crates/hipfire-arch-qwen35/examples/test_escha_dense_linear_gpu_vs_cpu.rs new file mode 100644 index 0000000000..054b0347d1 --- /dev/null +++ b/crates/hipfire-arch-qwen35/examples/test_escha_dense_linear_gpu_vs_cpu.rs @@ -0,0 +1,325 @@ +//! Gate: one escha-coded DENSE linear on GPU against the `escha_ref` oracle. +//! +//! This exists to be run BEFORE the dense path is wired into ten call sites in +//! `forward.rs`/`prefill.rs`. Every failure mode of that wiring — a missing +//! H128, the two H128s swapped, a bias applied before the output transform +//! instead of after, rin/rout transposed — produces a full-rank, finite, +//! plausible activation rather than a crash. Debugging that through a whole +//! 64-layer model is far more expensive than pinning the single linear first. +//! +//! The oracle is `escha_ref::expert_linear`, which is the same +//! input_transform -> matmul -> output_transform that `ref.py::dense_linear` +//! specifies, and which G2/G3 already gate bit-exact at the decode and H128 +//! level. Bias is added on top here because the reference helper covers the +//! coded linear only. +//! +//! Usage: +//! cargo run --release -p hipfire-arch-qwen35 \ +//! --example test_escha_dense_linear_gpu_vs_cpu -- [proj] +//! +//! `proj` defaults to `mlp.gate_proj` on layer 0; pass e.g. +//! `linear_attn.in_proj_qkv` to check another. + +use hipfire_arch_qwen35::qwen35::escha::{EschaProj, + escha_dense_leaf, escha_dense_linear_forward, load_escha_dense_linear, EschaWeightStore, +}; +use hipfire_runtime::hfq::HfqFile; +use rdna_compute::{DType, Gpu}; + +/// Same candidate expansion the loader uses, so a bare `layers.0.…` resolves +/// against the `model.language_model.…` names actually in the file. +fn candidates(name: &str) -> Vec { + hipfire_arch_qwen35::qwen35::load::qwen35_tensor_name_candidates(name) +} + +fn find(hfq: &HfqFile, name: &str) -> Option<(hipfire_runtime::hfq::HfqTensorInfo, Vec)> { + for c in candidates(name) { + if let Some((info, data)) = hfq.tensor_data(&c) { + return Some((info.clone(), data.to_vec())); + } + if let Some((info, buf)) = hfq.tensor_data_pread(&c) { + return Some((info.clone(), buf.to_vec())); + } + } + None +} + +fn main() -> Result<(), Box> { + let mut args = std::env::args().skip(1); + let path = args.next().expect("usage: [proj]"); + let proj = args.next().unwrap_or_else(|| "mlp.gate_proj".to_string()); + let p = "layers.0"; + + let hfq = HfqFile::open(std::path::Path::new(&path))?; + let mut gpu = Gpu::init()?; + + // Shape from the code tensor's own dims: escha stores [in/16, out/16, 16K]. + let (code_info, code_bytes) = + find(&hfq, &escha_dense_leaf(p, &proj, "code")).expect("escha_code not found"); + let k = match code_info.quant_type { + 42 => 2usize, + 43 => 3usize, + other => panic!("{proj}: quant_type {other} is not escha (42/43)"), + }; + let dims: Vec = code_info.shape.iter().map(|&d| d as usize).collect(); + assert_eq!(dims.len(), 3, "{proj}: escha_code should be 3-D, got {dims:?}"); + let (ic, oc) = (dims[0] * 16, dims[1] * 16); + println!("{proj}: ic={ic} oc={oc} K={k}"); + + // ── CPU reference ──────────────────────────────────────────────────── + let code_i16: Vec = code_bytes + .chunks_exact(2) + .map(|c| i16::from_le_bytes([c[0], c[1]])) + .collect(); + let w_bits = hipfire_quantize::escha_ref::reconstruct(&code_i16, ic, oc, k); + + let read_f32 = |leaf: &str, want: usize| -> Vec { + let (_, d) = find(&hfq, &escha_dense_leaf(p, &proj, leaf)) + .unwrap_or_else(|| panic!("{leaf} not found")); + let v: Vec = d + .chunks_exact(4) + .map(|c| f32::from_le_bytes([c[0], c[1], c[2], c[3]])) + .collect(); + assert_eq!(v.len(), want, "{leaf}: {} elements, want {want}", v.len()); + v + }; + let rin = read_f32("rin_eff", ic); + let rout = read_f32("rout_eff", oc); + + // Deterministic input; no rand dependency and reproducible across runs. + let x: Vec = (0..ic) + .map(|i| (((i * 2654435761usize) % 1000) as f32 / 500.0) - 1.0) + .collect(); + + let y_ref_bits = hipfire_quantize::escha_ref::expert_linear(&x, &w_bits, &rin, &rout); + let mut y_ref: Vec = y_ref_bits + .iter() + .map(|&b| hipfire_runtime::llama::f16_to_f32(b)) + .collect(); + + let mut bias_ref: Vec = vec![0.0; oc]; + // Bias AFTER the output transform, matching `ref.py::dense_linear`. + let bias_name = format!("{p}.{proj}.bias"); + if let Some((info, d)) = find(&hfq, &bias_name) { + let b: Vec = match info.quant_type { + 1 => d + .chunks_exact(2) + .map(|c| hipfire_runtime::llama::f16_to_f32(u16::from_le_bytes([c[0], c[1]]))) + .collect(), + _ => d + .chunks_exact(4) + .map(|c| f32::from_le_bytes([c[0], c[1], c[2], c[3]])) + .collect(), + }; + assert_eq!(b.len(), oc, "bias length"); + for (yi, bi) in y_ref.iter_mut().zip(b.iter()) { + *yi += bi; + } + bias_ref = b; + println!("bias: present ({oc} elements), applied after the output transform"); + } else { + println!("bias: absent (optional leaf)"); + } + + // ── FOLD CHECK ─────────────────────────────────────────────────────── + // Can the two H128s and both diagonals be folded into the weight, so an + // escha dense linear becomes an ORDINARY weight that every existing fused + // path can consume untouched? The algebra says yes: + // + // mid = W^T xh, xh = RS*H*diag(rin)*x, y = RS*diag(rout)*H*mid + // => W_eff[i][o] = RS^2 * rin_i * (H W H)[i][o] * rout_o + // + // The only deviation from the reference is that folding SKIPS the fp16 + // rounding of xh, so it should land at or slightly better than the + // runtime path — not worse. If this holds, the 27B needs no forward-path + // changes at all. + { + const RS: f32 = 0.088_388_347_648; + let mut wf: Vec = w_bits + .iter() + .map(|&b| hipfire_runtime::llama::f16_to_f32(b)) + .collect(); + // H along the contiguous o-axis (each of the ic rows, length oc). + for row in wf.chunks_exact_mut(oc) { + hipfire_quantize::escha_ref::h128_inplace(row); + } + // H along the strided i-axis (each of the oc columns, stride oc). + let mut col = vec![0.0f32; ic]; + for o in 0..oc { + for i in 0..ic { + col[i] = wf[i * oc + o]; + } + hipfire_quantize::escha_ref::h128_inplace(&mut col); + for i in 0..ic { + wf[i * oc + o] = col[i]; + } + } + // Both diagonals and RS^2. + for i in 0..ic { + let s = RS * RS * rin[i]; + for o in 0..oc { + wf[i * oc + o] *= s * rout[o]; + } + } + // Plain matmul, no transforms at all. + let mut y_fold = vec![0.0f32; oc]; + for i in 0..ic { + let a = x[i]; + let row = &wf[i * oc..(i + 1) * oc]; + for (m, w) in y_fold.iter_mut().zip(row) { + *m += a * w; + } + } + if let Some((info, d)) = find(&hfq, &bias_name) { + let b: Vec = match info.quant_type { + 1 => d + .chunks_exact(2) + .map(|c| hipfire_runtime::llama::f16_to_f32(u16::from_le_bytes([c[0], c[1]]))) + .collect(), + _ => d + .chunks_exact(4) + .map(|c| f32::from_le_bytes([c[0], c[1], c[2], c[3]])) + .collect(), + }; + for (yi, bi) in y_fold.iter_mut().zip(b.iter()) { + *yi += bi; + } + } + let mut num = 0.0f64; + let mut den = 0.0f64; + for (a, b) in y_fold.iter().zip(y_ref.iter()) { + let d = (*a - *b) as f64; + num += d * d; + den += (*b as f64) * (*b as f64); + } + println!( + " FOLDED (rotations baked into the weight, no runtime transform): rel_rms {:.3e}", + (num / den.max(1e-30)).sqrt() + ); + } + + // ── GPU ────────────────────────────────────────────────────────────── + // F16 store: the decode to fp16 is exact, so any difference is the + // forward path rather than a re-quantisation. Q8_0 is reported after. + for store in [EschaWeightStore::F16, EschaWeightStore::Native, EschaWeightStore::Q8_0] { + let lin = load_escha_dense_linear(&hfq, &mut gpu, p, &proj, ic, oc, store, candidates)?; + let xg = gpu.upload_f32(&x, &[ic])?; + let xh = gpu.alloc_tensor(&[ic], DType::F32)?; + let mid = gpu.alloc_tensor(&[oc], DType::F32)?; + let yg = gpu.alloc_tensor(&[oc], DType::F32)?; + escha_dense_linear_forward(&mut gpu, &lin, &xg, &xh, &mid, &yg)?; + let y = gpu.download_f32(&yg)?; + + // Stage-by-stage norms: a zero at a known stage names the broken step. + let l2 = |v: &[f32]| (v.iter().map(|a| (*a as f64) * (*a as f64)).sum::()).sqrt(); + let xh_h = gpu.download_f32(&xh)?; + let mid_h = gpu.download_f32(&mid)?; + println!( + " |x|={:.4} |xh|={:.4} |mid|={:.4} |y|={:.4} |y_ref|={:.4} w.dtype={:?} w.m={} w.k={}", + l2(&x), l2(&xh_h), l2(&mid_h), l2(&y), l2(&y_ref), lin.w.gpu_dtype, lin.w.m, lin.w.k + ); + + let mut worst = 0.0f32; + let mut num = 0.0f64; + let mut den = 0.0f64; + let mut nonfinite = 0usize; + for (a, b) in y.iter().zip(y_ref.iter()) { + let (a, b): (&f32, &f32) = (a, b); + if !a.is_finite() { + nonfinite += 1; + } + let d = (a - b).abs(); + if d > worst { + worst = d; + } + num += (d as f64) * (d as f64); + den += (*b as f64) * (*b as f64); + } + let rel_rms = (num / den.max(1e-30)).sqrt(); + // Timed: is the native trellis path fast enough to justify wiring it + // into the layer forward? Weight bytes differ per store, so report + // achieved bandwidth as well as raw time — the native store moves ~3x + // fewer bytes than Q8_0 for the same maths. + let iters = 200usize; + for _ in 0..10 { + escha_dense_linear_forward(&mut gpu, &lin, &xg, &xh, &mid, &yg)?; + } + gpu.hip.device_synchronize()?; + let t0 = std::time::Instant::now(); + for _ in 0..iters { + escha_dense_linear_forward(&mut gpu, &lin, &xg, &xh, &mid, &yg)?; + } + gpu.hip.device_synchronize()?; + let us = t0.elapsed().as_secs_f64() * 1e6 / iters as f64; + let wbytes = match store { + EschaWeightStore::Native => (ic * oc * k) as f64 / 8.0, + EschaWeightStore::F16 => (ic * oc * 2) as f64, + EschaWeightStore::Q8_0 => (ic * oc) as f64 * 34.0 / 32.0, + EschaWeightStore::F32 => (ic * oc * 4) as f64, + }; + println!( + " store={store:?}: rel_rms {rel_rms:.3e} worst_abs {worst:.3e} non-finite {nonfinite} {us:.1} us/call {:.1} GB/s weights {:.1} MB", + wbytes / us / 1e3, + wbytes / 1e6 + ); + // F16 must be tight; Q8_0 carries its own re-quantisation error and is + // reported rather than gated, so the two are not held to one bar. + // Native decodes the trellis INSIDE the GEMV, so it should track F16 + // (both deliver exactly-decoded weights); only the accumulation order + // differs. Q8_0 carries its own re-quantisation and is reported, not + // gated. + if matches!(store, EschaWeightStore::F16 | EschaWeightStore::Native) { + assert_eq!(nonfinite, 0, "non-finite output"); + assert!( + rel_rms < 2e-3, + "F16 store: rel_rms {rel_rms:.3e} exceeds 2e-3 — the dense forward disagrees \ + with escha_ref by more than fp16 accumulation explains" + ); + } + } + // BATCHED: the indexed GEMV must serve a dense linear as `slots` copies + // of expert 0. Verified before wiring it into the batched prefill path, + // because a wrong slot stride there is per-token garbage that only shows + // up as a bad PPL much later. + { + let lin = load_escha_dense_linear( + &hfq, &mut gpu, p, &proj, ic, oc, EschaWeightStore::Native, candidates)?; + let ep = EschaProj { + rin: lin.rin, rout: lin.rout, ptr0: lin.ptr0.expect("native has ptr0"), + }; + for slots in [1usize, 4] { + let mut xb = Vec::with_capacity(slots * ic); + for _ in 0..slots { xb.extend_from_slice(&x); } + let xg = gpu.upload_f32(&xb, &[slots * ic])?; + let ids = gpu.upload_f32(&vec![f32::from_bits(0); slots], &[slots])?; + let xh = gpu.alloc_tensor(&[slots * ic], DType::F32)?; + let mid = gpu.alloc_tensor(&[slots * oc], DType::F32)?; + let yg = gpu.alloc_tensor(&[slots * oc], DType::F32)?; + // Exercise BOTH shapes: the per-slot GEMV and, above 1 slot, the + // grouped WMMA GEMM that batched prefill uses. + let off_bytes: Vec = [0i32, slots as i32] + .iter().flat_map(|v| v.to_le_bytes()).collect(); + let offsets = gpu.upload_raw(&off_bytes, &[2])?; + let iota: Vec = (0..slots).map(|i| f32::from_bits(i as u32)).collect(); + let iota = gpu.upload_f32(&iota, &[slots])?; + let grouped = if slots > 1 { Some((&offsets, &iota)) } else { None }; + ep.forward(&mut gpu, &lin.w, &ids, &xg, &xh, &mid, &yg, slots, grouped)?; + let y = gpu.download_f32(&yg)?; + // Every slot fed the same x, so every slot must equal y_ref + // (minus the bias, which EschaProj deliberately does not add). + let mut worst = 0.0f64; + for sl in 0..slots { + for o in 0..oc { + let got = y[sl * oc + o] as f64; + let want = y_ref[o] as f64 - bias_ref[o] as f64; + let d = (got - want).abs() / (want.abs().max(1e-3)); + if d > worst { worst = d; } + } + } + println!(" BATCHED slots={slots}: worst_rel {worst:.3e}"); + assert!(worst < 5e-2, "batched slots={slots} worst_rel {worst:.3e}"); + } + } + println!("PASS"); + Ok(()) +} diff --git a/crates/hipfire-arch-qwen35/src/layer_driver.rs b/crates/hipfire-arch-qwen35/src/layer_driver.rs index 41cd836193..91969efda7 100644 --- a/crates/hipfire-arch-qwen35/src/layer_driver.rs +++ b/crates/hipfire-arch-qwen35/src/layer_driver.rs @@ -6,6 +6,7 @@ //! `WeightBackend`. `load_weights` (HFQ), `load_weights_paroquant` (PaRo), and //! `load_layer_into` (multi-GPU HFQ) all funnel through `load_layer`. +use crate::qwen35::weights::{DeltaNetBiases, DeltaNetEscha, FullAttnBiases, FullAttnEscha}; use crate::qwen35::{ DeltaNetLayerWeights, DeltaNetMoeLayerWeights, FullAttnLayerWeights, FullAttnMoeLayerWeights, LayerType, LayerWeights, MoeFfnWeights, Qwen35Config, @@ -16,6 +17,61 @@ use hipfire_runtime::weight_backend::WeightBackend; /// Load one layer's weights. `load_moe` builds the MoE FFN block for MoE layers /// (format-specific: HFQ `load_moe_ffn` vs PaRo `paro_load_moe_ffn`), supplied by /// the caller so MoE layout stays arch-owned. +/// Largest token batch an escha layer's `ids` table serves. Decode reads a +/// 1-element prefix; batched prefill reads `n`. 4096 covers every prefill +/// chunk hipfire issues and costs 16 KB per layer. +const ESCHA_MAX_SLOTS: usize = 4096; + +/// Build an `EschaProj` from the backend's sidecars, or `None` when the +/// weight is not a trellis code. +fn eproj( + b: &mut B, + rel: &str, + w: &hipfire_runtime::llama::WeightTensor, +) -> hip_bridge::HipResult> { + Ok(b.escha_sidecars(rel, w)? + .map(|(rin, rout, ptr0)| crate::qwen35::escha::EschaProj { rin, rout, ptr0 })) +} + +/// Every coded projection in a layer is escha or none is — the export does not +/// mix. So probe one and require the rest, exactly as the biases do: a +/// half-escha layer is a corrupt checkpoint, and falling back per projection +/// would silently run some through the fused MQ paths on trellis bytes. +fn need_eproj( + b: &mut B, + rel: &str, + w: &hipfire_runtime::llama::WeightTensor, +) -> hip_bridge::HipResult { + eproj(b, rel, w)?.ok_or_else(|| { + hip_bridge::HipError::new( + 0, + &format!("{rel}: layer has escha projections but this one is not coded"), + ) + }) +} + +/// A bias that MUST be there. +/// +/// Escha ships all of a layer's biases together or none, so once the first +/// probe finds one the rest are mandatory. A half-biased layer is a corrupt +/// checkpoint, and silently substituting zeros would degrade output without +/// failing — the same trap as the dropped MTP head. +fn need_bias( + b: &mut B, + rel: &str, + n: usize, +) -> hip_bridge::HipResult { + b.bias_opt(rel, n)?.ok_or_else(|| { + hip_bridge::HipError::new( + 0, + &format!( + "{rel}: layer has some escha biases but not this one — a partially biased \ + layer is a corrupt checkpoint, not a model to run with zeros" + ), + ) + }) +} + pub(crate) fn load_layer( b: &mut B, config: &Qwen35Config, @@ -32,10 +88,33 @@ pub(crate) fn load_layer( let o_in = config.n_heads * config.head_dim; Ok(match (config.layer_types[layer_idx], is_moe) { - (LayerType::LinearAttention, false) => LayerWeights::DeltaNet(DeltaNetLayerWeights { + (LayerType::LinearAttention, false) => { + // Coded projections bound to locals FIRST: the escha sidecars are + // keyed off each weight's dtype and device pointer, so they have + // to be built from the loaded tensor rather than alongside it. + let wqkv = b.proj("linear_attn.in_proj_qkv", qkv_dim, config.dim)?; + let wz = b.proj("linear_attn.in_proj_z", d_inner, config.dim)?; + let wo = b.proj("linear_attn.out_proj", config.dim, d_inner)?; + let w_gate = b.proj("mlp.gate_proj", config.hidden_dim, config.dim)?; + let w_up = b.proj("mlp.up_proj", config.hidden_dim, config.dim)?; + let w_down = b.proj("mlp.down_proj", config.dim, config.hidden_dim)?; + let escha = match eproj(b, "linear_attn.in_proj_qkv", &wqkv)? { + None => None, + Some(qkv) => Some(DeltaNetEscha { + qkv, + z: need_eproj(b, "linear_attn.in_proj_z", &wz)?, + o: need_eproj(b, "linear_attn.out_proj", &wo)?, + gate: need_eproj(b, "mlp.gate_proj", &w_gate)?, + up: need_eproj(b, "mlp.up_proj", &w_up)?, + down: need_eproj(b, "mlp.down_proj", &w_down)?, + ids: b.zeros_i32(ESCHA_MAX_SLOTS)?, + iota: b.iota_i32(ESCHA_MAX_SLOTS)?, + }), + }; + LayerWeights::DeltaNet(DeltaNetLayerWeights { attn_norm: b.norm("input_layernorm.weight", &[config.dim])?, - wqkv: b.proj("linear_attn.in_proj_qkv", qkv_dim, config.dim)?, - wz: b.proj("linear_attn.in_proj_z", d_inner, config.dim)?, + wqkv, + wz, w_alpha: b.proj( "linear_attn.in_proj_a", config.linear_num_value_heads, @@ -53,25 +132,78 @@ pub(crate) fn load_layer( qkv_dim * config.conv_kernel_dim, )?, norm_weight: b.raw_f32("linear_attn.norm.weight", config.linear_value_head_dim)?, - wo: b.proj("linear_attn.out_proj", config.dim, d_inner)?, + wo, ffn_norm: b.norm("post_attention_layernorm.weight", &[config.dim])?, - w_gate: b.proj("mlp.gate_proj", config.hidden_dim, config.dim)?, - w_up: b.proj("mlp.up_proj", config.hidden_dim, config.dim)?, - w_down: b.proj("mlp.down_proj", config.dim, config.hidden_dim)?, - }), - (LayerType::FullAttention, false) => LayerWeights::FullAttn(FullAttnLayerWeights { + w_gate, + w_up, + w_down, + // Escha dense exports (Qwen3.8-27B) carry an additive output bias + // on every coded projection. All six are present together or not + // at all, so one probe decides. `in_proj_a`/`in_proj_b` have none + // — escha's `ignore` list keeps them as plain weights. + biases: match b.bias_opt("linear_attn.in_proj_qkv.bias", qkv_dim)? { + None => None, + Some(qkv) => Some(DeltaNetBiases { + qkv, + z: need_bias(b, "linear_attn.in_proj_z.bias", d_inner)?, + o: need_bias(b, "linear_attn.out_proj.bias", config.dim)?, + gate: need_bias(b, "mlp.gate_proj.bias", config.hidden_dim)?, + up: need_bias(b, "mlp.up_proj.bias", config.hidden_dim)?, + down: need_bias(b, "mlp.down_proj.bias", config.dim)?, + }), + }, + escha, + }) + } + (LayerType::FullAttention, false) => { + let wq = b.proj("self_attn.q_proj", q_out_dim, config.dim)?; + let wk = b.proj("self_attn.k_proj", kv_dim, config.dim)?; + let wv = b.proj("self_attn.v_proj", kv_dim, config.dim)?; + let wo = b.proj("self_attn.o_proj", config.dim, o_in)?; + let w_gate = b.proj("mlp.gate_proj", config.hidden_dim, config.dim)?; + let w_up = b.proj("mlp.up_proj", config.hidden_dim, config.dim)?; + let w_down = b.proj("mlp.down_proj", config.dim, config.hidden_dim)?; + let escha = match eproj(b, "self_attn.q_proj", &wq)? { + None => None, + Some(q) => Some(FullAttnEscha { + q, + k: need_eproj(b, "self_attn.k_proj", &wk)?, + v: need_eproj(b, "self_attn.v_proj", &wv)?, + o: need_eproj(b, "self_attn.o_proj", &wo)?, + gate: need_eproj(b, "mlp.gate_proj", &w_gate)?, + up: need_eproj(b, "mlp.up_proj", &w_up)?, + down: need_eproj(b, "mlp.down_proj", &w_down)?, + ids: b.zeros_i32(ESCHA_MAX_SLOTS)?, + iota: b.iota_i32(ESCHA_MAX_SLOTS)?, + }), + }; + LayerWeights::FullAttn(FullAttnLayerWeights { attn_norm: b.norm("input_layernorm.weight", &[config.dim])?, - wq: b.proj("self_attn.q_proj", q_out_dim, config.dim)?, - wk: b.proj("self_attn.k_proj", kv_dim, config.dim)?, - wv: b.proj("self_attn.v_proj", kv_dim, config.dim)?, - wo: b.proj("self_attn.o_proj", config.dim, o_in)?, + wq, + wk, + wv, + wo, q_norm: b.norm("self_attn.q_norm.weight", &[config.head_dim])?, k_norm: b.norm("self_attn.k_norm.weight", &[config.head_dim])?, ffn_norm: b.norm("post_attention_layernorm.weight", &[config.dim])?, - w_gate: b.proj("mlp.gate_proj", config.hidden_dim, config.dim)?, - w_up: b.proj("mlp.up_proj", config.hidden_dim, config.dim)?, - w_down: b.proj("mlp.down_proj", config.dim, config.hidden_dim)?, - }), + w_gate, + w_up, + w_down, + biases: match b.bias_opt("self_attn.q_proj.bias", q_out_dim)? { + None => None, + Some(q) => Some(FullAttnBiases { + q, + k: need_bias(b, "self_attn.k_proj.bias", kv_dim)?, + v: need_bias(b, "self_attn.v_proj.bias", kv_dim)?, + o: need_bias(b, "self_attn.o_proj.bias", config.dim)?, + gate: need_bias(b, "mlp.gate_proj.bias", config.hidden_dim)?, + up: need_bias(b, "mlp.up_proj.bias", config.hidden_dim)?, + down: need_bias(b, "mlp.down_proj.bias", config.dim)?, + }), + }, + escha, + }) + } (LayerType::LinearAttention, true) => LayerWeights::DeltaNetMoe(DeltaNetMoeLayerWeights { attn_norm: b.norm("input_layernorm.weight", &[config.dim])?, wqkv: b.proj("linear_attn.in_proj_qkv", qkv_dim, config.dim)?, diff --git a/crates/hipfire-arch-qwen35/src/lib.rs b/crates/hipfire-arch-qwen35/src/lib.rs index fa4d01eb64..5e04edbc74 100644 --- a/crates/hipfire-arch-qwen35/src/lib.rs +++ b/crates/hipfire-arch-qwen35/src/lib.rs @@ -119,3 +119,129 @@ pub use carrier::{free_qwen35_bundle, load_bundle as load_qwen35_bundle, Qwen35B pub use mtp_compose::{spec_step_dflash_mtp_tree, MtpComposeTreeResult, MtpComposeTreeState}; #[cfg(feature = "deltanet")] pub use mtp_speculator::{build_qwen35_mtp_speculator, Qwen35MtpDrafter}; + +/// G4b (Escha-W2 port, Task 9): expose the arch-6 MoE router selection step +/// so `examples/escha_router_contract.rs` can call the *actual* production +/// router — not a reimplementation of it — on an arbitrary `.hfq` layer and +/// compare its top-K expert SET against EschaLabs' shipped fixture. +/// +/// This calls exactly the same primitives the production decode/prefill path +/// calls (`hipfire_runtime::llama::weight_gemv` for the router GEMV, then +/// whichever top-K kernel `hipfire_dispatch::pipeline::run_moe_decode` would +/// pick for this GPU's arch — the fused `moe_router_softmax_topk_k8_wave64_exact` +/// on gfx1100/gfx1151, or the reference two-launch `softmax_f32` + +/// `moe_topk_renorm_k8` everywhere else). No selection math is duplicated +/// here; the kernel-choice arch check itself is +/// `hipfire_dispatch::pipeline::exact_wave64_router_predicate` — the exact +/// function `run_moe_decode` calls, not a copy of its logic — so this +/// helper cannot silently drift from production's actual gate. +/// +/// This test model's `.hfq` (`gate.weight` quant_type=1/F16) does not carry +/// escha routed-expert dtypes, so it never exercises the escha-only f16 +/// router-logits round-trip (review Fix 1); this helper intentionally omits +/// that step to stay a pure probe of the pre-existing selection kernels. +#[cfg(feature = "deltanet")] +pub fn escha_router_topk_for_test( + hfq_path: &str, + layer: usize, + x: &[f32], + n_tokens: usize, + hidden: usize, + top_k: usize, +) -> Result, String> { + use hipfire_dispatch::context::DispatchCtx; + use hipfire_runtime::hfq::{load_weight_tensor_pread, HfqFile}; + use hipfire_runtime::llama::weight_gemv; + use rdna_compute::{DType, Gpu}; + + if x.len() != n_tokens * hidden { + return Err(format!( + "escha_router_topk_for_test: x.len()={} != n_tokens*hidden={}", + x.len(), + n_tokens * hidden + )); + } + + let hfq = HfqFile::open(std::path::Path::new(hfq_path)).map_err(|e| e.to_string())?; + let config = qwen35::config_from_hfq(&hfq)?; + let n_exp = config.num_experts; + let norm_topk = config.norm_topk_prob; + + let mut gpu = Gpu::init().map_err(|e| e.to_string())?; + let ctx = DispatchCtx::new(&gpu); + // The *actual* predicate `run_moe_decode` (hipfire-dispatch/src/pipeline/mod.rs) + // uses to pick the fused exact-wave64 router kernel over the reference + // two-launch path — extracted to `exact_wave64_router_predicate` so this + // helper cannot silently drift from production's real gate (review Fix 2: + // the previous `is_gfx1151() || is_gfx1100()` copy here happened to agree + // with production for this model/GPU, but wasn't actually the same check — + // production also requires `n_exp == 256` and honors the + // `HIPFIRE_GFX1100_ROUTER_W64` override on gfx1100). + // HIPFIRE_MOE_ROUTER_SHARED_FUSE-gated shared-expert fusion is a pure perf + // variant of the same math and is left out here — it doesn't change which + // experts get selected. + let gfx1100_router_mode = hipfire_config::developer_var("HIPFIRE_GFX1100_ROUTER_W64").ok(); + let use_exact_wave64 = hipfire_dispatch::pipeline::exact_wave64_router_predicate( + n_exp, + &ctx.arch, + gfx1100_router_mode.as_deref(), + ); + + fn exact_name(s: &str) -> Vec { + vec![s.to_string()] + } + let weight_name = format!("model.language_model.layers.{layer}.mlp.gate.weight"); + let router = load_weight_tensor_pread(&hfq, &gpu, &weight_name, n_exp, hidden, exact_name) + .map_err(|e| e.to_string())?; + + let logits = gpu + .alloc_tensor(&[n_exp], DType::F32) + .map_err(|e| e.to_string())?; + // topk_idx carries raw i32 selections in an f32-tagged buffer — the same + // "i32-in-F32 alias" convention `moe_ffn_decode_impl` uses for its scratch + // (see qwen35/forward.rs `capture_expert_stats`'s `ti[krank].to_bits()`). + let topk_idx = gpu + .alloc_tensor(&[top_k], DType::F32) + .map_err(|e| e.to_string())?; + let topk_w = gpu + .alloc_tensor(&[top_k], DType::F32) + .map_err(|e| e.to_string())?; + + let mut out = Vec::with_capacity(n_tokens * top_k); + for t in 0..n_tokens { + let row = &x[t * hidden..(t + 1) * hidden]; + let bytes: &[u8] = + unsafe { std::slice::from_raw_parts(row.as_ptr() as *const u8, row.len() * 4) }; + let x_gpu = gpu + .upload_raw(bytes, &[hidden]) + .map_err(|e| e.to_string())?; + + weight_gemv(&mut gpu, &router, &x_gpu, &logits).map_err(|e| e.to_string())?; + + if use_exact_wave64 { + gpu.moe_router_softmax_topk_k8_wave64_exact( + &logits, &topk_idx, &topk_w, n_exp, norm_topk, + ) + .map_err(|e| e.to_string())?; + } else { + gpu.softmax_f32(&logits).map_err(|e| e.to_string())?; + gpu.moe_topk_renorm_k8(&logits, &topk_idx, &topk_w, n_exp, norm_topk) + .map_err(|e| e.to_string())?; + } + + gpu.hip.device_synchronize().map_err(|e| e.to_string())?; + let idx_f32 = gpu.download_f32(&topk_idx).map_err(|e| e.to_string())?; + for v in idx_f32.iter().take(top_k) { + out.push((v.to_bits() as i32) as u32); + } + + let _ = gpu.free_tensor(x_gpu); + } + + let _ = gpu.free_tensor(logits); + let _ = gpu.free_tensor(topk_idx); + let _ = gpu.free_tensor(topk_w); + router.free_all(&mut gpu); + + Ok(out) +} diff --git a/crates/hipfire-arch-qwen35/src/mtp_head.rs b/crates/hipfire-arch-qwen35/src/mtp_head.rs index 68e4eb798a..493b269289 100644 --- a/crates/hipfire-arch-qwen35/src/mtp_head.rs +++ b/crates/hipfire-arch-qwen35/src/mtp_head.rs @@ -822,11 +822,192 @@ pub fn load_mtp_head(path: &Path, gpu: &mut Gpu, max_seq: usize) -> HipResult = &'a dyn Fn(&str) -> String; + +/// Identity: the bare names `mtp_extract` writes. +fn bare_namer(n: &str) -> String { + n.to_string() +} + +/// HF-style names as the escha converter passes them through from upstream. +/// Every logical name the loader asks for has exactly one counterpart; an +/// unmapped name is a programming error rather than a missing tensor, so it +/// panics instead of silently looking up a name no container has. +fn trunk_namer(n: &str) -> String { + let m = match n { + "eh_proj" => "mtp.fc.weight", + "enorm" => "mtp.pre_fc_norm_embedding.weight", + "hnorm" => "mtp.pre_fc_norm_hidden.weight", + "shared_head_norm" => "mtp.norm.weight", + "attn_norm" => "mtp.layers.0.input_layernorm.weight", + "attn_post_norm" => "mtp.layers.0.post_attention_layernorm.weight", + "attn_q_norm" => "mtp.layers.0.self_attn.q_norm.weight", + "attn_k_norm" => "mtp.layers.0.self_attn.k_norm.weight", + "wq" => "mtp.layers.0.self_attn.q_proj.weight", + "wk" => "mtp.layers.0.self_attn.k_proj.weight", + "wv" => "mtp.layers.0.self_attn.v_proj.weight", + "wo" => "mtp.layers.0.self_attn.o_proj.weight", + "ffn_gate" => "mtp.layers.0.mlp.gate_proj.weight", + "ffn_up" => "mtp.layers.0.mlp.up_proj.weight", + "ffn_down" => "mtp.layers.0.mlp.down_proj.weight", + "moe_router" => "mtp.layers.0.mlp.gate.weight", + "moe_shared_gate" => "mtp.layers.0.mlp.shared_expert.gate_proj.weight", + "moe_shared_up" => "mtp.layers.0.mlp.shared_expert.up_proj.weight", + "moe_shared_down" => "mtp.layers.0.mlp.shared_expert.down_proj.weight", + "moe_shared_expert_gate" => "mtp.layers.0.mlp.shared_expert_gate.weight", + other => panic!("trunk_namer: no HF spelling for MTP tensor '{other}'"), + }; + m.to_string() +} + +impl Qwen35MtpHeadConfig { + /// Derive the head's config from a TRUNK model's `text_config`. + /// + /// A standalone `.mtp` carries flat metadata describing the head; a trunk + /// carries HF config describing the trunk. The escha MTP head is one + /// transformer layer over the trunk's own geometry (`mtp_num_hidden_layers` + /// is 1 and `mtp_use_dedicated_embeddings` is false), so every dimension + /// comes from the trunk — verified against the shipped 27B, whose head + /// tensors match these exactly: q_proj `[2*head_dim*n_head, n_embd]` = + /// [12288, 5120] under attn_output_gate, k/v `[head_dim*n_head_kv, n_embd]` + /// = [1024, 5120], mlp `[intermediate_size, n_embd]` = [17408, 5120]. + pub fn from_trunk_text_config(tc: &serde_json::Value, max_seq: usize) -> Option { + let gu = |k: &str| tc.get(k).and_then(|v| v.as_u64()).map(|v| v as usize); + let n_embd = gu("hidden_size")?; + let n_head = gu("num_attention_heads")?; + let n_head_kv = gu("num_key_value_heads")?; + let head_dim = gu("head_dim")?; + let moe_int = gu("moe_intermediate_size").unwrap_or(0); + let n_ff = gu("intermediate_size").unwrap_or(moe_int); + let num_experts = gu("num_experts").unwrap_or(0); + let ffn_kind = if num_experts > 0 { + Qwen35MtpFfnKind::Moe + } else { + Qwen35MtpFfnKind::Dense + }; + let prf = tc + .get("partial_rotary_factor") + .and_then(|v| v.as_f64()) + .unwrap_or(1.0); + let rope_theta = tc + .get("rope_parameters") + .and_then(|r| r.get("rope_theta")) + .or_else(|| tc.get("rope_theta")) + .and_then(|v| v.as_f64()) + .unwrap_or(10_000_000.0) as f32; + Some(Self { + n_embd, + n_head, + n_head_kv, + head_dim, + n_ff, + ffn_kind, + num_experts, + num_experts_per_tok: gu("num_experts_per_tok").unwrap_or(0), + moe_intermediate_size: moe_int, + shared_expert_intermediate_size: gu("shared_expert_intermediate_size").unwrap_or(0), + norm_topk_prob: tc + .get("norm_topk_prob") + .and_then(|v| v.as_bool()) + .unwrap_or(false), + vocab_size: gu("vocab_size")?, + rope_theta, + n_rot: (head_dim as f64 * prf) as usize, + rms_norm_eps: tc + .get("rms_norm_eps") + .and_then(|v| v.as_f64()) + .unwrap_or(1e-6) as f32, + max_seq, + tie_word_embeddings: tc + .get("tie_word_embeddings") + .and_then(|v| v.as_bool()) + .unwrap_or(false), + }) + } +} + +/// Load an MTP head carried as ordinary `mtp.*` tensors inside a TRUNK +/// container, which is how the escha converter passes upstream's head +/// through. Returns `Ok(None)` when the trunk has no such tensors. +/// +/// This is the third MTP packaging hipfire sees, after the bundled +/// `HFBNDMTP` trailer and the sibling `.mtp` sidecar. The escha builds ship +/// 849 MB (27B) of head that neither of those resolvers could find. +pub fn load_mtp_head_from_trunk( + path: &Path, + gpu: &mut Gpu, + max_seq: usize, +) -> HipResult> { + let hfq = match HfqFile::open(path) { + Ok(h) => h, + Err(_) => return Ok(None), + }; + if hfq.find_tensor_info("mtp.fc.weight").is_none() { + return Ok(None); + } + let meta: serde_json::Value = match serde_json::from_str(&hfq.metadata_json) { + Ok(m) => m, + Err(_) => return Ok(None), + }; + let tc = meta + .get("config") + .and_then(|c| c.get("text_config")) + .or_else(|| meta.get("config")); + let Some(cfg) = tc.and_then(|t| Qwen35MtpHeadConfig::from_trunk_text_config(t, max_seq)) else { + return Ok(None); + }; + // An MoE head needs its routed experts. Upstream's 35B ships only the + // router and shared expert inside the trunk (17 mtp.* tensors, none of + // them `experts.*`), so `load_mtp_moe_ffn` would panic partway through + // loading 256 experts that are not there. Decline with a reason instead: + // the head is present but not self-sufficient. + if matches!(cfg.ffn_kind, Qwen35MtpFfnKind::Moe) + && hfq + .find_tensor_info("mtp.layers.0.mlp.experts.0.gate_up_proj.weight") + .is_none() + && hfq + .find_tensor_info("mtp.layers.0.mlp.experts.0.up_proj.weight") + .is_none() + { + eprintln!( + " MTP head present in trunk but not loadable: MoE head with \ + num_experts={} ships no routed experts (router and shared expert \ + only). Speculative decode stays off.", + cfg.num_experts + ); + return Ok(None); + } + drop(hfq); + let head = load_mtp_head_at_offset_named(path, gpu, max_seq, 0, &trunk_namer, Some(cfg))?; + Ok(Some(head)) +} + pub fn load_mtp_head_at_offset( path: &Path, gpu: &mut Gpu, max_seq: usize, base_offset: u64, +) -> HipResult { + load_mtp_head_at_offset_named(path, gpu, max_seq, base_offset, &bare_namer, None) +} + +/// As [`load_mtp_head_at_offset`], but the caller supplies the tensor-name +/// spelling and, when the container is a trunk rather than a `.mtp`, the +/// already-derived head config (a trunk's metadata describes the trunk, not +/// the head). +pub fn load_mtp_head_at_offset_named( + path: &Path, + gpu: &mut Gpu, + max_seq: usize, + base_offset: u64, + namer: MtpNamer<'_>, + config_override: Option, ) -> HipResult { let hfq = HfqFile::open_at_offset(path, base_offset).unwrap_or_else(|e| { panic!( @@ -834,17 +1015,26 @@ pub fn load_mtp_head_at_offset( path.display() ) }); - assert_eq!( - hfq.arch_id, - 21, - ".mtp file at {} has arch_id={} (expected 21 = QWEN35_MTP_HEAD); \ - is this actually an MTP head extracted by mtp_extract?", - path.display(), - hfq.arch_id - ); + // A standalone `.mtp` is its own container and must be arch 21. A trunk + // carrying the head as in-container `mtp.*` tensors is arch 5/6 and the + // check does not apply — `config_override` being set is exactly the + // signal that the caller resolved the geometry from the trunk instead. + if config_override.is_none() { + assert_eq!( + hfq.arch_id, + 21, + ".mtp file at {} has arch_id={} (expected 21 = QWEN35_MTP_HEAD); \ + is this actually an MTP head extracted by mtp_extract?", + path.display(), + hfq.arch_id + ); + } let meta: serde_json::Value = serde_json::from_str(&hfq.metadata_json).expect(".mtp metadata JSON parse failed"); - let config = Qwen35MtpHeadConfig::from_metadata(&meta, max_seq); + let config = match config_override { + Some(c) => c, + None => Qwen35MtpHeadConfig::from_metadata(&meta, max_seq), + }; // ── Norms (F32, 1D) ───────────────────────────────────────────────── // @@ -859,29 +1049,29 @@ pub fn load_mtp_head_at_offset( // verified 2026-05-15 A/B: removing +1.0 regressed K=3 from τ=3.08 to // τ=2.00 on 27B-3.5 LRU bench. The MTP head trains its `mtp.norm` with // the trunk per-layer convention, NOT the trunk final-norm convention. - let shared_head_norm = load_norm_raw(&hfq, gpu, "shared_head_norm", n_embd)?; - let enorm = load_norm_raw(&hfq, gpu, "enorm", n_embd)?; - let hnorm = load_norm_raw(&hfq, gpu, "hnorm", n_embd)?; - let attn_norm = load_norm_raw(&hfq, gpu, "attn_norm", n_embd)?; - let attn_post_norm = load_norm_raw(&hfq, gpu, "attn_post_norm", n_embd)?; - let attn_q_norm = load_norm_raw(&hfq, gpu, "attn_q_norm", head_dim)?; - let attn_k_norm = load_norm_raw(&hfq, gpu, "attn_k_norm", head_dim)?; + let shared_head_norm = load_norm_raw(&hfq, gpu, &namer("shared_head_norm"), n_embd)?; + let enorm = load_norm_raw(&hfq, gpu, &namer("enorm"), n_embd)?; + let hnorm = load_norm_raw(&hfq, gpu, &namer("hnorm"), n_embd)?; + let attn_norm = load_norm_raw(&hfq, gpu, &namer("attn_norm"), n_embd)?; + let attn_post_norm = load_norm_raw(&hfq, gpu, &namer("attn_post_norm"), n_embd)?; + let attn_q_norm = load_norm_raw(&hfq, gpu, &namer("attn_q_norm"), head_dim)?; + let attn_k_norm = load_norm_raw(&hfq, gpu, &namer("attn_k_norm"), head_dim)?; // ── 2D weights ────────────────────────────────────────────────────── let q_full_dim = 2 * head_dim * config.n_head; let kv_dim = head_dim * config.n_head_kv; let q_dim = head_dim * config.n_head; - let eh_proj = load_weight_raw(&hfq, gpu, "eh_proj", n_embd, 2 * n_embd)?; - let wq = load_weight_raw(&hfq, gpu, "wq", q_full_dim, n_embd)?; - let wk = load_weight_raw(&hfq, gpu, "wk", kv_dim, n_embd)?; - let wv = load_weight_raw(&hfq, gpu, "wv", kv_dim, n_embd)?; - let wo = load_weight_raw(&hfq, gpu, "wo", n_embd, q_dim)?; + let eh_proj = load_weight_raw(&hfq, gpu, &namer("eh_proj"), n_embd, 2 * n_embd)?; + let wq = load_weight_raw(&hfq, gpu, &namer("wq"), q_full_dim, n_embd)?; + let wk = load_weight_raw(&hfq, gpu, &namer("wk"), kv_dim, n_embd)?; + let wv = load_weight_raw(&hfq, gpu, &namer("wv"), kv_dim, n_embd)?; + let wo = load_weight_raw(&hfq, gpu, &namer("wo"), n_embd, q_dim)?; let ffn = match config.ffn_kind { Qwen35MtpFfnKind::Dense => Qwen35MtpFfnWeights::Dense(Qwen35MtpDenseFfnWeights { - gate: load_weight_raw(&hfq, gpu, "ffn_gate", config.n_ff, n_embd)?, - up: load_weight_raw(&hfq, gpu, "ffn_up", config.n_ff, n_embd)?, - down: load_weight_raw(&hfq, gpu, "ffn_down", n_embd, config.n_ff)?, + gate: load_weight_raw(&hfq, gpu, &namer("ffn_gate"), config.n_ff, n_embd)?, + up: load_weight_raw(&hfq, gpu, &namer("ffn_up"), config.n_ff, n_embd)?, + down: load_weight_raw(&hfq, gpu, &namer("ffn_down"), n_embd, config.n_ff)?, }), Qwen35MtpFfnKind::Moe => { assert_eq!( @@ -910,7 +1100,7 @@ pub fn load_mtp_head_at_offset( .expect("metadata claims has_compressed_lm_head_draft but lacks compressed_vocab_size") as usize; assert!(cvs > 0, "compressed_vocab_size must be positive"); - let lm_d = load_weight_raw(&hfq, gpu, "lm_head_draft.weight", cvs, n_embd)?; + let lm_d = load_weight_raw(&hfq, gpu, &namer("lm_head_draft.weight"), cvs, n_embd)?; let (vmap_info, vmap_bytes) = hfq .tensor_data_vec("lm_head_draft.vocab_map") .expect("compressed sidecar missing vocab_map tensor"); @@ -1083,14 +1273,19 @@ fn load_mtp_moe_ffn( gpu.hip.memcpy_htod(&expert_gate_up_ptrs.buf, &gu_bytes)?; gpu.hip.memcpy_htod(&expert_down_ptrs.buf, &dn_bytes)?; - Ok(Qwen35MtpMoeFfnWeights { + let ffn = Qwen35MtpMoeFfnWeights { router, shared_expert, shared_expert_gate, experts, expert_gate_up_ptrs, expert_down_ptrs, - }) + }; + // Fail at LOAD, not at the first speculative step: a head that this + // forward cannot run correctly must not be reported as loaded. + // See `mtp_moe_refuse_unsupported_rotation`. + mtp_moe_refuse_unsupported_rotation(&ffn)?; + Ok(ffn) } /// Cross-check the on-disk shape against the caller's expected (m, k). @@ -1716,6 +1911,70 @@ pub fn mtp_head_forward_block_only_with_pos_buf( Ok(()) } +/// Refuse any MTP MoE weight whose dtype needs a rotation this function does +/// not apply. +/// +/// # Why this exists +/// +/// [`mtp_moe_ffn_decode`] below is a SECOND, independent MoE forward: it does +/// not go through `hipfire_dispatch::pipeline::run_moe_decode`, so none of the +/// escha guards there (`check_moe_decode_supported` arm (c), the escha branch +/// that routes to `pipeline::escha`) protect it. It applies exactly one +/// rotation, `rotate_x_mq_for`'s FWHT, and combines the raw expert outputs. +/// +/// An Escha-W2 layer here would therefore skip BOTH Hadamard transforms and +/// multiply rotated-domain weights by an unrotated activation — finite, +/// fluent, ~1e-1-wrong draft tokens. Speculative decoding would then verify +/// them against a correct trunk and reject nearly all of them: the failure +/// would present as "the speculator is useless", not as "the speculator is +/// wrong", which is the shape of bug that survives a release. +/// +/// It is unreachable today by two accidents, neither of which is a decision: +/// the MTP loader reads per-expert `moe_experts.{i}.{gate_up,down}` tensor +/// names that an escha checkpoint does not contain, and the escha SKU ships no +/// MTP sidecar at all. Wiring a speculator to escha must fail loudly, here. +/// +/// The `match` is deliberately EXHAUSTIVE — no `_` arm. A new `RotationPlan` +/// variant will not compile until someone decides whether this forward +/// implements it or refuses it. +fn mtp_moe_refuse_unsupported_rotation(ffn: &Qwen35MtpMoeFfnWeights) -> HipResult<()> { + use hipfire_dispatch::types::{dtype_rotation_plan, RotationPlan}; + + let mut check = |label: &str, dt: DType| -> HipResult<()> { + match dtype_rotation_plan(dt) { + // Applied by `rotate_x_mq_for` / folded into the kernels below. + RotationPlan::None + | RotationPlan::FwhtG256 + | RotationPlan::FwhtG128 + | RotationPlan::Mq8Internal + | RotationPlan::Givens => Ok(()), + RotationPlan::EschaH128 => Err(HipError::new( + 0, + &format!( + "MTP head: {label} is {dt:?}, whose rotation plan is EschaH128. The MTP \ + MoE forward is a separate implementation from run_moe_decode and applies \ + NO H128 pair, so it would multiply rotated-domain weights by an \ + unrotated activation and emit finite, fluent, ~1e-1-wrong draft tokens. \ + Refusing. To wire a speculator to Escha-W2, teach this forward the H128 \ + transforms (see hipfire_dispatch::pipeline::escha) — do not delete this \ + check." + ), + )), + } + }; + + check("moe_router", ffn.router.gpu_dtype)?; + check("moe_shared_expert_gate", ffn.shared_expert_gate.gpu_dtype)?; + check("moe_shared_gate", ffn.shared_expert.gate.gpu_dtype)?; + check("moe_shared_up", ffn.shared_expert.up.gpu_dtype)?; + check("moe_shared_down", ffn.shared_expert.down.gpu_dtype)?; + for (i, e) in ffn.experts.iter().enumerate() { + check(&format!("moe_experts.{i}.gate_up"), e.gate_up.gpu_dtype)?; + check(&format!("moe_experts.{i}.down"), e.down.gpu_dtype)?; + } + Ok(()) +} + fn mtp_moe_ffn_decode( gpu: &mut Gpu, ffn: &Qwen35MtpMoeFfnWeights, @@ -1724,6 +1983,12 @@ fn mtp_moe_ffn_decode( cfg: &Qwen35MtpHeadConfig, scratch: &Qwen35MtpHeadScratch, ) -> HipResult<()> { + // Before any GPU work. `load_mtp_moe_ffn` makes the same call so an escha + // sidecar is refused at LOAD rather than at the first speculative step; + // this one is the guard at the point of danger, and it is what protects a + // `Qwen35MtpMoeFfnWeights` built by any future path that skips the loader. + mtp_moe_refuse_unsupported_rotation(ffn)?; + let dim = cfg.n_embd; let mi = cfg.moe_intermediate_size; let smi = cfg.shared_expert_intermediate_size; diff --git a/crates/hipfire-arch-qwen35/src/paro_moe.rs b/crates/hipfire-arch-qwen35/src/paro_moe.rs index ff6661030e..bd9181109a 100644 --- a/crates/hipfire-arch-qwen35/src/paro_moe.rs +++ b/crates/hipfire-arch-qwen35/src/paro_moe.rs @@ -218,5 +218,7 @@ pub(crate) fn paro_load_moe_ffn( paro_shared: Some(shared), global_expert_dtypes: None, ep_dummy_buffers: Vec::new(), + // ParoQuant, not Escha-W2. + escha: None, }) } diff --git a/crates/hipfire-arch-qwen35/src/qwen35.rs b/crates/hipfire-arch-qwen35/src/qwen35.rs index 31fda38b1f..431bd9a786 100644 --- a/crates/hipfire-arch-qwen35/src/qwen35.rs +++ b/crates/hipfire-arch-qwen35/src/qwen35.rs @@ -11,6 +11,9 @@ pub mod batch; pub mod config; pub mod ep_batch; +/// Escha-W2 routed-expert loading (Task 10): trellis decode -> transpose -> +/// Q8_0, plus the per-layer H128 transform tables. +pub mod escha; pub mod forward; pub mod load; pub mod prefill; @@ -31,6 +34,7 @@ pub use ep_batch::{ forward_ep, forward_prefill_batch_ep, forward_prefill_batch_multi, forward_scratch_multi, validate_ep_batch_compatibility, Qwen35DecodeBatchEpState, }; +pub use escha::{EschaMoeTables, EschaWeightStore}; pub use forward::{ dump_expert_stats, forward, forward_gpu, forward_prefill_dense_tp, forward_scratch, forward_scratch_dense_tp, forward_scratch_embed, forward_scratch_embed_mrope, diff --git a/crates/hipfire-arch-qwen35/src/qwen35/batch.rs b/crates/hipfire-arch-qwen35/src/qwen35/batch.rs index c7c7f802f7..d7729e0a74 100644 --- a/crates/hipfire-arch-qwen35/src/qwen35/batch.rs +++ b/crates/hipfire-arch-qwen35/src/qwen35/batch.rs @@ -63,6 +63,18 @@ pub struct PrefillBatchScratch { pub up_batch: GpuTensor, // SwiGLU output (FWHT-rotated for MQ4) feeding w_down. pub ffn_hidden_batch: GpuTensor, + /// Escha trellis scratch: the H128-rotated activations feeding one + /// projection's batched GEMV, `[max_batch, max(hidden_dim, dim)]`. + /// + /// Cannot share `x_rot_batch` or `ffn_hidden_batch`: each escha projection + /// rotates the SAME input with its OWN `rin`, and for `down_proj` the + /// ffn hidden state IS the input, so writing xh there would destroy it. + pub escha_xh_batch: GpuTensor, + /// Escha projection output before it is accumulated into the residual, + /// `[max_batch, dim]`. The fused epilogue writes straight into the + /// residual; the trellis GEMV has no residual variant, so out_proj and + /// down_proj land here first. + pub escha_y_batch: GpuTensor, // FWHT-rotated dn_normed [N × v_dim] feeding wo for MQ4 weights. // Decode path handles this via an internal mq_x_rot scratch inside @@ -262,6 +274,8 @@ impl PrefillBatchScratch { gate_ffn_batch: alloc!(&[max_batch * hidden_dim], DType::F32), up_batch: alloc!(&[max_batch * hidden_dim], DType::F32), ffn_hidden_batch: alloc!(&[max_batch * hidden_dim], DType::F32), + escha_xh_batch: alloc!(&[max_batch * hidden_dim.max(dim)], DType::F32), + escha_y_batch: alloc!(&[max_batch * dim], DType::F32), dn_normed_rot_batch: alloc!(&[max_batch * v_dim], DType::F32), // F32 dtype = 4 bytes/element, same layout as i32. The rope / // attention / kv_write kernels cast the pointer to `const int*`, diff --git a/crates/hipfire-arch-qwen35/src/qwen35/escha.rs b/crates/hipfire-arch-qwen35/src/qwen35/escha.rs new file mode 100644 index 0000000000..79cc779e7b --- /dev/null +++ b/crates/hipfire-arch-qwen35/src/qwen35/escha.rs @@ -0,0 +1,1160 @@ +// SPDX-License-Identifier: Apache-2.0 +// Copyright (c) 2026 Kaden Schutt +// hipfire — see LICENSE and NOTICE in the project root. +//! Escha-W2 routed-expert loading for arch-6 (Task 10). +//! +//! An Escha-W2 `.hfq` stores each MoE projection as ONE tensor per layer +//! covering all `num_experts` experts: +//! +//! | tensor | qt | shape | meaning | +//! |---|---|---|---| +//! | `…experts.gate_up_proj.escha_code` | 42 (`ESCHA2T16`) | `[E, ic/16, oc/16, 16*2]` i16 | K=2 trellis code | +//! | `…experts.gate_up_proj.escha_rin_eff` | 2 (F32) | `[E, ic]` | folded input scales | +//! | `…experts.gate_up_proj.escha_rout_eff` | 2 (F32) | `[E, oc]` | folded output scales (carries the prune mask) | +//! | `…experts.down_proj.escha_code` | 43 (`ESCHA3T16`) | `[E, ic/16, oc/16, 16*3]` i16 | K=3 trellis code | +//! | `…experts.down_proj.escha_rin_eff` / `…rout_eff` | 2 (F32) | `[E, ic]` / `[E, oc]` | | +//! +//! ## Orientation — the one thing that must not be got wrong +//! +//! Escha's code tile grid is **IN-MAJOR** (`[in/16, out/16]`) and +//! `Gpu::escha_decode_tiles` writes bare fp16 **row-major `[in_features, +//! out_features]`**. hipfire's expert slots are **OUT-MAJOR** — +//! `experts[X].gate_up` is `[2*moe_intermediate, hidden]` (see +//! `weights.rs:69`) and every hipfire GEMV walks K contiguously along a row. +//! So on the three DECODING stores a transpose happens, and it happens exactly +//! once, folded into the store pass (`Gpu::escha_bare_to_q8_0` and friends). A +//! wrong orientation still yields a full-rank, plausible weight matrix, so it +//! is gated by the G4 block gate, never by "the output looks sane". +//! +//! [`EschaWeightStore::Native`] keeps escha's in-major grid instead, because +//! it stores the code and never materialises a matrix at all; the fused GEMV +//! addresses the in-major tile grid directly. Same hazard, same gate. +//! +//! hipfire's `gate_up` slot is already FUSED (gate ‖ up), matching escha's +//! single fused `gate_up_proj`, so there is no concat step. +//! +//! ## Why production stores the CODE (Phase 2) +//! +//! Phase 1 decoded to `Q8_0` at load: 1.0625 B/weight, 34.2 GB of routed +//! experts, 37.55 GB resident, and 1.07 GB of routed-expert traffic on every +//! decode token. That put a hard 69 tok/s roofline on a box measured at +//! 209 GB/s — under the 71.8 tok/s the comparable `qwen3.6:35b-a3b-mq4r` SKU +//! already reaches, so Q8_0 could not have won at any efficiency. +//! +//! [`EschaWeightStore::Native`] stores the trellis code verbatim (2.00 bpw for +//! the K=2 gate_up, 3.00 for the K=3 down) and decodes it inside the routed +//! GEMV. It is both smaller AND weight-exact — the `Q8_0` re-quantisation that +//! dominates the G4 block gate's error simply does not happen. The remaining +//! decoding stores are measurement arms: [`EschaWeightStore::Q8_0`] is Phase 1, +//! kept because every published Phase 1 number was measured on it, +//! [`EschaWeightStore::F16`] is the G5 KLD reference, and +//! [`EschaWeightStore::F32`] is a small-layer weight-exact control. + +use hip_bridge::HipError; +use hip_bridge::HipResult; +use hipfire_runtime::hfq::HfqFile; +use hipfire_runtime::llama::WeightTensor; +use rdna_compute::DType; +use rdna_compute::Gpu; +use rdna_compute::GpuTensor; + +use super::weights::ExpertWeights; +use super::weights::PackedExpertOwners; + +/// How the decoded fp16 expert weight is stored in the expert slot. +#[derive(Clone, Copy, PartialEq, Eq, Debug)] +pub enum EschaWeightStore { + /// PRODUCTION (Phase 2): the trellis code, verbatim, not decoded at all. + /// 0.25 B/weight at K=2 and 0.375 at K=3 — a quarter of Q8_0 — and the + /// routed GEMV decodes it in-register + /// (`Gpu::escha_gemv_native_moe_k8_indexed_batched`). + /// + /// This is the arm that makes the port fast, and it is fast for exactly + /// one reason: the expanded copy stops existing, so it stops crossing the + /// bus. At A3B shapes the routed half of a decode token moves 0.294 GB + /// instead of 1.07 GB, and the whole model 2.23 GB/token instead of 3.01 + /// — a 94 tok/s roofline instead of 69 on this box's measured 209 GB/s. + /// Q8_0's 69 was BELOW the 71.8 the comparable `qwen3.6:35b-a3b-mq4r` SKU + /// already achieves, so no amount of efficiency could have closed it. + /// + /// It is also weight-EXACT: the fused GEMV consumes the same fp16 values + /// `escha_decode_tiles` would have produced, so unlike [`Self::Q8_0`] it + /// carries no re-quantisation error at all. Gated bit-for-bit against the + /// F16 store and against `escha_ref` by + /// `rdna-compute/examples/test_escha_native_gemv_gpu_vs_cpu.rs`. + /// + /// Requires the indexed (GPU-top-K) route: there is no per-expert + /// native GEMV, so `HIPFIRE_ESCHA_INDEXED=0` with this store fails loudly + /// in `GemvFamily::run_auto` (`RotationPlan::EschaH128` has no plain GEMV) + /// rather than running unrotated. The A/B lever for this store is + /// `HIPFIRE_ESCHA_EXPERT_STORE=q8_0`, which restores Phase 1 entirely. + Native, + /// Phase 1 production, now the A/B arm: transpose + Q8_0 re-quantise + /// (1.0625 B/weight). Kept because it is what every published Phase 1 + /// number was measured on — including the G4 block gate's Q8_0 arm and the + /// G5 KLD headline — and because it is the only routed store that also + /// works through the per-expert host route. + Q8_0, + /// Diagnostic control arm: transpose only, F32 store (4 B/weight), so a + /// caller can separate "the H128 wiring is wrong" from "Q8_0 costs this + /// much". Do not use for a whole model. + F32, + /// Weight-exact arm that DOES fit a whole model: transpose only, F16 + /// store (2 B/weight). The decode already produced fp16, so this holds + /// bit-identically the same values as [`EschaWeightStore::F32`] in half + /// the bytes. It is the G5 KLD reference arm. + /// + /// It costs **2x production's expert bytes**, and that is now the whole + /// difference. It did not used to be: while every per-expert buffer was + /// its own allocation, the HIP allocator's 2 MiB granule rounded Q8_0's + /// 2.125 MiB gate_up / 1.0625 MiB down up to exactly the 4 MiB / 2 MiB + /// F16 needed outright, so both arms sat at 60 GiB of experts (measured + /// 67.9 GB of GTT for the whole Q8_0 model on gfx1151, against a 34.2 GB + /// logical expert size) and F16 was free. Since the projections are packed + /// one buffer per (layer, projection) — see [`PackedExpertOwners`] — the + /// granule is charged 80 times instead of 20,480 and Q8_0 is measured at + /// 37.6 GB. F16 would be ~32 GB more. It remains the G5 KLD reference arm + /// and still fits; it is no longer a free upgrade. + /// + /// Like F32 this loses the indexed GPU-top-K fast path (admission is + /// `routed_gate_up == Q8_0 && routed_down == Q8_0`, see + /// hipfire-dispatch `families/moe.rs`) and runs host-routed instead. That + /// is slower and numerically identical. + F16, +} + +/// One layer's Escha-W2 transform tables plus the per-layer decode scratch the +/// batched routed executor needs. +/// +/// The `[E, ·]` tables stay resident in full — they are 5.5 MB/layer at A3B +/// shapes (2+1+0.5+2), i.e. 220 MB for the whole model, and keeping them whole +/// is precisely what lets one H128 launch serve all `top_k` experts: slot `s` +/// indexes row `ids[s]`, no gather. +/// +/// The scratch is per-layer rather than model-global purely so ownership is +/// simple (it is freed with the layer). At `k=8` / A3B shapes it is ~272 KB +/// per layer, 11 MB for the model. +pub struct EschaMoeTables { + pub gate_up_rin: GpuTensor, + pub gate_up_rout: GpuTensor, + pub down_rin: GpuTensor, + pub down_rout: GpuTensor, + pub ids: GpuTensor, + pub weights: GpuTensor, + pub xh_gu: GpuTensor, + pub mid_gu: GpuTensor, + pub y_gu: GpuTensor, + pub h: GpuTensor, + pub xh_dn: GpuTensor, + pub mid_dn: GpuTensor, + pub y_dn: GpuTensor, + pub hidden: usize, + pub mi: usize, + pub k: usize, +} + +impl EschaMoeTables { + /// Borrow as the dispatch-crate view. Logic-free adapter. + pub fn refs(&self) -> hipfire_dispatch::pipeline::escha::EschaRoutedRefs<'_> { + hipfire_dispatch::pipeline::escha::EschaRoutedRefs { + gate_up_rin: &self.gate_up_rin, + gate_up_rout: &self.gate_up_rout, + down_rin: &self.down_rin, + down_rout: &self.down_rout, + ids: &self.ids, + weights: &self.weights, + xh_gu: &self.xh_gu, + mid_gu: &self.mid_gu, + y_gu: &self.y_gu, + h: &self.h, + xh_dn: &self.xh_dn, + mid_dn: &self.mid_dn, + y_dn: &self.y_dn, + } + } + + pub fn free_gpu(self, gpu: &mut Gpu) { + for t in [ + self.gate_up_rin, + self.gate_up_rout, + self.down_rin, + self.down_rout, + self.ids, + self.weights, + self.xh_gu, + self.mid_gu, + self.y_gu, + self.h, + self.xh_dn, + self.mid_dn, + self.y_dn, + ] { + let _ = gpu.free_tensor(t); + } + } +} + +/// Kill switch for the escha INDEXED (GPU-resident top-K) routed route. +/// +/// `HIPFIRE_ESCHA_INDEXED=0` withholds `MoeDtypes::routed_escha_transforms`, +/// which drops `routed_indexable_escha_q8`, which drops `use_gpu_topk`, which +/// sends the layer back down the CPU-top-K route and its host-routed escha +/// executor. Everything stays consistent on the way — including +/// `check_moe_decode_supported`, which sees a non-indexed escha layer and +/// admits it — so this is a genuine A/B of the two routes in ONE build, not a +/// half-disabled state. +/// +/// It exists because the two routes are BIT-IDENTICAL (gated by +/// `examples/escha_moe_block_gate.rs`) and differ only in cost, so the +/// performance claim for the indexed route is checkable at any time without a +/// rebuild or a revert. It is also the escape hatch if the indexed route ever +/// needs to be taken out of service in the field. +/// +/// Default ON. Read once. +pub fn escha_indexed_route_enabled() -> bool { + static ENABLED: std::sync::OnceLock = std::sync::OnceLock::new(); + *ENABLED.get_or_init(|| { + hipfire_config::developer_var("HIPFIRE_ESCHA_INDEXED").as_deref() != Ok("0") + }) +} + +/// `.hfq` tensor name for one of the six escha MoE leaves of a layer, BEFORE +/// candidate expansion. `p` is the bare layer prefix `load.rs` uses +/// (`layers.N`); the caller's `resolve` is what turns that into the +/// checkpoint's actual `model.language_model.layers.N…` name, exactly as +/// every other tensor in this loader is resolved. +pub fn escha_leaf(p: &str, proj: &str, leaf: &str) -> String { + format!("{p}.mlp.experts.{proj}_proj.escha_{leaf}") +} + +/// Candidate-expanding lookup. Mirrors `hfq::load_weight_tensor`'s contract so +/// an escha layer resolves through the same name aliasing as everything else +/// in the checkpoint (`layers.0.…` -> `model.language_model.layers.0.…`). +pub type NameResolver = fn(&str) -> Vec; + +/// `tensor_data_vec`, NOT `tensor_data`: on a unified-memory APU the qwen35 +/// loader drops the mmap in `prepare()` (mapped pages cannot be evicted while +/// the mapping exists, and they starve `hipMalloc`), after which +/// `tensor_data` returns `None` for every tensor while `find_tensor_info` +/// keeps working. Reading through the mmap here therefore fails only on the +/// full-model path and not in a single-layer probe — exactly the shape of bug +/// that ships. `tensor_data_vec` takes the pread + `FADV_DONTNEED` route the +/// rest of the loader uses. +fn find<'a>( + hfq: &'a HfqFile, + name: &str, + resolve: NameResolver, +) -> Option<(&'a hipfire_runtime::hfq::HfqTensorInfo, Vec)> { + resolve(name) + .into_iter() + .find_map(|c| hfq.tensor_data_vec(&c)) +} + +/// True iff this layer's routed experts are Escha-W2 coded. Keyed on the +/// `gate_up` code tensor's presence AND its quant type, so a checkpoint that +/// happened to carry a same-named tensor of another format is rejected by the +/// loader rather than mis-decoded. +/// Leaf name for a DENSE escha linear: `{p}.{proj}.escha_{leaf}`. +/// +/// Distinct from `escha_leaf`, which is MoE-shaped +/// (`{p}.mlp.experts.{proj}_proj.escha_{leaf}`). Here `proj` is the full path +/// below the layer — `linear_attn.in_proj_qkv`, `self_attn.q_proj`, +/// `mlp.gate_proj` — because the dense export codes every projection in place +/// rather than gathering experts under one name. +pub fn escha_dense_leaf(p: &str, proj: &str, leaf: &str) -> String { + format!("{p}.{proj}.escha_{leaf}") +} + +/// Is this individual dense projection escha-coded? +/// +/// Keyed on the CODE tensor's quant type, exactly as `layer_is_escha` is: a +/// projection is escha if and only if its code is qt=42/43. Presence of the +/// name alone is not enough — `escha_config` and friends are optional leaves +/// (§1.4 of the design doc), and the required trio is code/rin/rout. +/// +/// Per-projection rather than per-layer because the dense 27B mixes K within +/// a layer: `mlp.gate_proj` is K=2 while `mlp.up_proj` is K=3, so nothing at +/// layer granularity can describe it. +pub fn dense_proj_is_escha(hfq: &HfqFile, p: &str, proj: &str, resolve: NameResolver) -> bool { + resolve(&escha_dense_leaf(p, proj, "code")) + .into_iter() + .find_map(|c| hfq.find_tensor_info(&c)) + .is_some_and(|i| i.quant_type == 42 || i.quant_type == 43) +} + +pub fn layer_is_escha(hfq: &HfqFile, p: &str, resolve: NameResolver) -> bool { + resolve(&escha_leaf(p, "gate_up", "code")) + .into_iter() + .find_map(|c| hfq.find_tensor_info(&c)) + .is_some_and(|i| i.quant_type == 42 || i.quant_type == 43) +} + +fn read_f32_tensor( + hfq: &HfqFile, + gpu: &Gpu, + name: &str, + want: usize, + resolve: NameResolver, +) -> HipResult { + let (info, data) = find(hfq, name, resolve) + .ok_or_else(|| HipError::new(0, &format!("escha: tensor not found: {name}")))?; + if info.quant_type != 2 { + return Err(HipError::new( + 0, + &format!( + "escha: {name} has quant_type {} (expected 2 = F32)", + info.quant_type + ), + )); + } + if data.len() != want * 4 { + return Err(HipError::new( + 0, + &format!( + "escha: {name} is {} bytes, expected {} ({want} f32)", + data.len(), + want * 4 + ), + )); + } + gpu.upload_raw(&data, &[want]) +} + +/// K (trellis order) implied by the on-disk quant type. +fn k_from_quant_type(qt: u8, name: &str) -> HipResult { + match qt { + 42 => Ok(2), + 43 => Ok(3), + other => Err(HipError::new( + 0, + &format!("escha: {name} has quant_type {other}, expected 42 (K=2) or 43 (K=3)"), + )), + } +} + +/// Decode one layer's escha experts into hipfire's expert slots, and build the +/// layer's transform tables. +/// +/// `expert_ids` selects which experts to materialise, in slot order — the +/// caller's REAP/EP mapping, or simply `0..n_exp`. Passing a short list is how +/// the G4 gate keeps a single-layer probe cheap. +/// +/// ## Ownership +/// +/// The returned [`ExpertWeights`] are **non-owning views** into the returned +/// [`PackedExpertOwners`] pair — one device buffer per projection covering +/// every requested expert. The caller must keep the owners alive for as long +/// as the views are used and free the owners (not the views) exactly once. In +/// the model loader that is `MoeFfnWeights::packed_expert_owners`, whose +/// existing free path (`free_moe_ffn`) already frees per-expert metadata only +/// and returns the two blobs; a direct caller such as the G4 gate must do the +/// same. `Gpu::free_tensor` refuses a borrowed view, so a caller that gets +/// this wrong gets an error rather than a double free — but it also leaks the +/// blob, so it is not a substitute for freeing the owners. +#[allow(clippy::too_many_arguments)] +pub fn load_escha_moe_experts( + hfq: &HfqFile, + gpu: &mut Gpu, + p: &str, + expert_ids: &[usize], + n_exp: usize, + hidden: usize, + mi: usize, + k: usize, + store: EschaWeightStore, + resolve: NameResolver, +) -> HipResult<(Vec, EschaMoeTables, PackedExpertOwners)> { + // gate_up: [ic = hidden, oc = 2*mi]; down: [ic = mi, oc = hidden]. + let gu = (hidden, 2 * mi); + let dn = (mi, hidden); + + let tables = EschaMoeTables { + gate_up_rin: read_f32_tensor( + hfq, + gpu, + &escha_leaf(p, "gate_up", "rin_eff"), + n_exp * gu.0, + resolve, + )?, + gate_up_rout: read_f32_tensor( + hfq, + gpu, + &escha_leaf(p, "gate_up", "rout_eff"), + n_exp * gu.1, + resolve, + )?, + down_rin: read_f32_tensor( + hfq, + gpu, + &escha_leaf(p, "down", "rin_eff"), + n_exp * dn.0, + resolve, + )?, + down_rout: read_f32_tensor( + hfq, + gpu, + &escha_leaf(p, "down", "rout_eff"), + n_exp * dn.1, + resolve, + )?, + // DELIBERATE DTYPE REINTERPRETATION: `ids` holds `k` 32-bit signed + // INTEGERS — the H128 batched kernels bind it as `const int*`. It is + // declared `DType::F32` only because `rdna_compute::DType` has no + // integer variant; F32 is the 4-byte-per-element stand-in, and the + // allocation size is therefore correct. This mirrors + // `qwen35::forward`'s `topk_indices`, which does the same thing for + // the same reason. + // + // Consequence: `gpu.download_f32(ids)` returns GARBAGE (int bit + // patterns reinterpreted as floats), and so would any f32 kernel + // pointed at it. Read it back with a raw byte download and + // `i32::from_le_bytes`. Fixing this properly means adding an integer + // DType to rdna-compute, which is out of scope here. + ids: gpu.alloc_tensor(&[k], DType::F32)?, + // `weights` genuinely IS f32 (the f16-rounded combine scores). + weights: gpu.alloc_tensor(&[k], DType::F32)?, + xh_gu: gpu.alloc_tensor(&[k * gu.0], DType::F32)?, + mid_gu: gpu.alloc_tensor(&[k * gu.1], DType::F32)?, + y_gu: gpu.alloc_tensor(&[k * gu.1], DType::F32)?, + h: gpu.alloc_tensor(&[k * mi], DType::F32)?, + xh_dn: gpu.alloc_tensor(&[k * dn.0], DType::F32)?, + mid_dn: gpu.alloc_tensor(&[k * dn.1], DType::F32)?, + y_dn: gpu.alloc_tensor(&[k * dn.1], DType::F32)?, + hidden, + mi, + k, + }; + + let (mut gate_ups, gate_up_owner) = decode_projection( + hfq, gpu, p, "gate_up", expert_ids, n_exp, gu, store, resolve, escha_leaf, + )?; + let (mut downs, down_owner) = match decode_projection( + hfq, gpu, p, "down", expert_ids, n_exp, dn, store, resolve, escha_leaf, + ) { + Ok(ok) => ok, + Err(error) => { + // The gate_up blob is already on the device and its per-expert + // views are about to be dropped without ever reaching a + // caller, so nothing else can free it. Return it here or the + // whole projection (544 MiB at A3B shapes) leaks on every + // failed layer load. + let _ = gpu.free_tensor(gate_up_owner); + return Err(error); + } + }; + + let experts = gate_ups + .drain(..) + .zip(downs.drain(..)) + .map(|(gate_up, down)| ExpertWeights { gate_up, down }) + .collect(); + Ok(( + experts, + tables, + PackedExpertOwners { + gate_up: gate_up_owner, + down: down_owner, + }, + )) +} + +/// Bytes and elements one expert slot of this projection occupies, for a given +/// store. `(elems_per_slot, dtype)` — `sub_offset` counts in `dtype.size()` +/// units, and `DType::Q8_0::size()` is 1, so the Q8_0 arm's "elements" are +/// bytes. Pure, so the packing arithmetic is checkable without a GPU. +fn slot_extent(store: EschaWeightStore, ic: usize, oc: usize, trellis_k: usize) -> (usize, DType) { + match store { + // Native holds the code stream itself: `(ic/16) * (oc/16)` tiles of + // `16 * trellis_k` int16. `Escha{2,3}T16::size()` is 1 (a byte dtype, + // like Q8_0), so these "elements" are BYTES. + EschaWeightStore::Native => ( + (ic / 16) * (oc / 16) * 16 * trellis_k * 2, + escha_dtype(trellis_k), + ), + // Q8_0 rows are `ic/32` blocks of 34 B (32 int8 + one f16 scale). + EschaWeightStore::Q8_0 => (oc * (ic / 32) * 34, DType::Q8_0), + EschaWeightStore::F32 => (ic * oc, DType::F32), + EschaWeightStore::F16 => (ic * oc, DType::F16), + } +} + +/// The `DType` that names a trellis order. This is the dtype the routed expert +/// slot CARRIES under [`EschaWeightStore::Native`], and it is what +/// `MoeResolution::routed_indexable_escha_native` and the batched-prefill +/// admission arm key on — so the layer's route is decided by the same fact the +/// GEMV's bit geometry is decided by, not by two independently-maintained +/// flags. +fn escha_dtype(trellis_k: usize) -> DType { + if trellis_k == 2 { + DType::Escha2T16 + } else { + DType::Escha3T16 + } +} + +/// Decode every requested expert of ONE projection into ONE device buffer. +/// +/// Staging is reused across experts: one device code buffer, one device bare +/// buffer. At A3B gate_up shapes that is 512 KB + 4 MB held for the whole +/// layer instead of 256 allocations, and the decode never round-trips through +/// the host (`escha_decode_tiles` is the device-resident entry; the `_host` +/// wrapper exists only for the G2 parity gate). +/// +/// The returned `WeightTensor`s are non-owning `sub_offset` views into the +/// returned owner buffer — see [`load_escha_moe_experts`] for why, and +/// [`PackedExpertOwners`] for how much it is worth. Each slot's byte offset is +/// `slot * slot_extent(...)`; at A3B shapes that stride is a multiple of 1024, +/// so every view is at least as aligned as an independent allocation would be +/// and no kernel's vector loads are disturbed. The values written are +/// byte-identical to the per-allocation version: `escha_bare_to_*` takes a +/// base pointer and a size, and both are unchanged. +#[allow(clippy::too_many_arguments)] +/// How a projection's leaves are named. `escha_leaf` for the MoE export, +/// `escha_dense_leaf` for the dense one — the two namespaces are disjoint and +/// nothing else about the decode differs, so the namer is a parameter rather +/// than a second copy of this function. +pub type LeafNamer = fn(&str, &str, &str) -> String; + +/// One escha-coded DENSE linear, loaded and ready for the forward pass. +/// +/// The dense export (Qwen3.8-27B) codes every projection in place rather than +/// gathering experts, so there is no expert table and no routing — just a +/// weight, the two rotation vectors, and the additive bias the end-to-end +/// fine-tune leaves behind. +pub struct EschaDenseLinear { + /// Decoded weight when `store` is Q8_0/F16, or the verbatim trellis code + /// when it is Native. + pub w: WeightTensor, + /// `escha_rin_eff`, `[ic]` f32 — pre-multiplied into x before the input + /// H128. + pub rin: GpuTensor, + /// `escha_rout_eff`, `[oc]` f32 — applied after the output H128. + pub rout: GpuTensor, + /// `bias`, `[oc]`. Present on the 27B, absent on the 35B. Base + /// Qwen3.8-27B has `attention_bias: false` and no MLP bias, so this is + /// purely Escha's additive output correction and is applied AFTER the + /// output transform, per `ref.py::dense_linear`. Applying it before the + /// H128 would be silently wrong rather than a crash. + pub bias: Option, + /// Buffer owning `w`'s bytes; freed with the layer. + pub owner: GpuTensor, + /// One-element `[0]` slot table, so the BATCHED H128 kernels can serve a + /// dense linear as the degenerate single-slot case. See + /// `escha_dense_linear_forward` for why the batched form and not the + /// single one. + pub ids0: GpuTensor, + /// One-element expert-pointer table holding `w.buf`'s device address. + /// + /// Present only for `EschaWeightStore::Native`, where `w` IS the trellis + /// code and there is no decoded weight for a normal GEMV to read. Every + /// escha GEMV kernel is expert-INDEXED, so rather than write a second + /// near-identical kernel, a dense linear is served as the degenerate + /// one-expert case: `expert_ptrs = [&code]`, `ids = [0]`, `slots = 1`. + /// That reuses the kernel G2 already gates bit-exact against the oracle + /// instead of forking the trellis inner loop. + pub ptr0: Option, +} + +/// Per-projection escha runtime data for ONE dense linear, held alongside the +/// `WeightTensor` rather than replacing it. +/// +/// The weight itself stays a `WeightTensor` (dtype `Escha2T16`/`Escha3T16`, +/// buffer = verbatim trellis code) so every existing `layer.wqkv.gpu_dtype` +/// check keeps working. What a trellis weight needs BEYOND that — the two +/// rotation vectors and the one-element pointer table the indexed GEMV wants — +/// lives here. +pub struct EschaProj { + pub rin: GpuTensor, + pub rout: GpuTensor, + pub ptr0: GpuTensor, +} + +impl EschaProj { + /// Run this projection: H128 in -> trellis GEMV -> H128 out. + /// + /// Bias is NOT applied here — it is added by the existing per-op bias + /// path, so there is exactly one place that knows bias ordering. + /// + /// `slots` is the token count: 1 for decode, n for batched prefill. The + /// indexed GEMV serves a dense linear as `slots` copies of expert 0, so + /// `ids` must be a slots-long run of zeros. + #[allow(clippy::too_many_arguments)] + pub fn forward( + &self, + gpu: &mut Gpu, + w: &WeightTensor, + ids: &GpuTensor, + x: &GpuTensor, + xh: &GpuTensor, + mid: &GpuTensor, + y: &GpuTensor, + slots: usize, + // `(expert_offsets, sorted_slot_index)` for the grouped GEMM, or + // `None` to force the per-slot GEMV. Decode passes `None`. + grouped: Option<(&GpuTensor, &GpuTensor)>, + ) -> HipResult<()> { + let (ic, oc) = (w.k, w.m); + let tk = match w.gpu_dtype { + DType::Escha2T16 => 2u32, + DType::Escha3T16 => 3u32, + other => { + return Err(HipError::new( + 0, + &format!("EschaProj::forward: dtype {other:?} is not a trellis code"), + )) + } + }; + let xg = if slots == 1 { + rdna_compute::EschaXGroup::Broadcast + } else { + rdna_compute::EschaXGroup::PerSlot + }; + gpu.escha_h128_batched( + "escha_h128_in_batched", + x, + &self.rin, + ids, + xh, + ic, + slots, + xg, + )?; + match grouped { + // BATCHED: one group holding every slot. The indexed GEMV re-reads + // the weight ONCE PER SLOT — correct for MoE, where each slot is a + // different expert, and 512x the weight traffic for a dense linear + // where every token shares one weight. The grouped WMMA GEMM reads + // it once per (layer, batch) instead, which is the same fix that + // took the 35B's expert path from 4.525 to 2.657 ms/token. + Some((offsets, iota)) if slots > 1 => { + // nt_major = true: dense escha codes are transposed at load + // by `escha_tiles_to_nt_major`. MoE experts stay kt-major. + gpu.escha_gemm_native_moe_grouped_wmma( + &self.ptr0, offsets, iota, xh, mid, oc, ic, slots, 1, tk, true, + )?; + } + _ => { + gpu.escha_gemv_native_moe_k8_indexed_batched( + &self.ptr0, ids, xh, mid, oc, ic, slots, tk, true, + )?; + } + } + gpu.escha_h128_batched( + "escha_h128_out_batched", + mid, + &self.rout, + ids, + y, + oc, + slots, + rdna_compute::EschaXGroup::PerSlot, + )?; + Ok(()) + } +} + +/// Load the escha runtime data for a projection whose weight is a trellis +/// code. `None` when the weight is any other dtype — that is the signal a +/// layer is not escha and should take its ordinary path. +pub fn load_escha_proj( + hfq: &HfqFile, + gpu: &mut Gpu, + p: &str, + proj: &str, + w: &WeightTensor, + resolve: NameResolver, +) -> HipResult> { + if !matches!(w.gpu_dtype, DType::Escha2T16 | DType::Escha3T16) { + return Ok(None); + } + let (ic, oc) = (w.k, w.m); + let rin = read_f32_tensor(hfq, gpu, &escha_dense_leaf(p, proj, "rin_eff"), ic, resolve)?; + let rout = read_f32_tensor( + hfq, + gpu, + &escha_dense_leaf(p, proj, "rout_eff"), + oc, + resolve, + )?; + let addr = w.buf.buf.as_ptr() as u64; + let ptr0 = gpu.upload_raw(&addr.to_le_bytes(), &[1])?; + Ok(Some(EschaProj { rin, rout, ptr0 })) +} + +/// Load one dense escha linear: `{p}.{proj}` with `[ic, oc]`. +/// +/// `proj` is the full path below the layer (`linear_attn.in_proj_qkv`, +/// `mlp.gate_proj`, `self_attn.q_proj`). Reuses `decode_projection` at +/// `n_exp = 1` — a dense linear is exactly the degenerate case of one expert, +/// and duplicating that decode would mean two places to keep bit-exact +/// against the oracle. +#[allow(clippy::too_many_arguments)] +pub fn load_escha_dense_linear( + hfq: &HfqFile, + gpu: &mut Gpu, + p: &str, + proj: &str, + ic: usize, + oc: usize, + store: EschaWeightStore, + resolve: NameResolver, +) -> HipResult { + let (mut ws, owner) = decode_projection( + hfq, + gpu, + p, + proj, + &[0], + 1, + (ic, oc), + store, + resolve, + escha_dense_leaf, + )?; + if ws.len() != 1 { + return Err(HipError::new( + 0, + &format!("escha: {p}.{proj} decoded {} slots, expected 1", ws.len()), + )); + } + let w = ws.remove(0); + let rin = read_f32_tensor(hfq, gpu, &escha_dense_leaf(p, proj, "rin_eff"), ic, resolve)?; + let rout = read_f32_tensor( + hfq, + gpu, + &escha_dense_leaf(p, proj, "rout_eff"), + oc, + resolve, + )?; + + // Bias is OPTIONAL by the leaf contract (§1.4): an export without the + // end-to-end stage ships none and must still load. So absence is not an + // error — but a bias of the wrong length is, because it would broadcast + // or truncate into plausible-looking output. + let bias_name = format!("{p}.{proj}.bias"); + let bias = match find(hfq, &bias_name, resolve) { + None => None, + Some((info, data)) => { + let elems = match info.quant_type { + 1 => data.len() / 2, // F16 + 2 => data.len() / 4, // F32 + other => { + return Err(HipError::new( + 0, + &format!("escha: {bias_name} has quant_type {other} (expected F16 or F32)"), + )) + } + }; + if elems != oc { + return Err(HipError::new( + 0, + &format!("escha: {bias_name} has {elems} elements, expected oc = {oc}"), + )); + } + Some(read_bias_f32(gpu, info.quant_type, &data, oc)?) + } + }; + // `ids` holds a 32-bit signed INTEGER, declared F32 only because + // rdna_compute::DType has no integer variant — the same deliberate + // reinterpretation `EschaMoeTables::ids` documents. A dense linear is + // slot 0 of a one-entry table, so the bytes are four zeros. + let ids0 = gpu.upload_f32(&[f32::from_bits(0)], &[1])?; + // Native store keeps the code verbatim, so the GEMV needs its address in + // a device-side table. Decoded stores (Q8_0/F16/F32) go through the + // ordinary GEMV and need none. + let ptr0 = if matches!(store, EschaWeightStore::Native) { + let addr = w.buf.buf.as_ptr() as u64; + Some(gpu.upload_raw(&addr.to_le_bytes(), &[1])?) + } else { + None + }; + Ok(EschaDenseLinear { + w, + rin, + rout, + bias, + owner, + ids0, + ptr0, + }) +} + +/// Run one escha-coded dense linear. +/// +/// The whole point of the format lives in these three steps, in this order: +/// +/// ```text +/// xh = f16( H128(x * rin) * RS ) escha_h128_in +/// mid = xh @ W plain GEMV on the decoded weight +/// y = H128(mid) * RS * rout + bias escha_h128_out, then the bias +/// ``` +/// +/// Skipping either H128 does NOT crash. It produces a full-rank, finite, +/// entirely plausible activation that is simply wrong — which is why this is +/// one function rather than three calls open-coded at ten call sites. +/// +/// The bias goes on AFTER the output transform, per `ref.py::dense_linear`. +/// Folding it in before the H128 would be rotated along with the signal and +/// is the single easiest way to get this silently wrong. +/// +/// `xh` and `mid` are caller-owned scratch so a layer can reuse one pair +/// across its projections instead of allocating per call. +pub fn escha_dense_linear_forward( + gpu: &mut Gpu, + lin: &EschaDenseLinear, + x: &GpuTensor, + xh: &GpuTensor, + mid: &GpuTensor, + y: &GpuTensor, +) -> HipResult<()> { + // BATCHED H128, even for a single vector. The two variants differ in + // OUTPUT TYPE, not just shape: `escha_h128_in` writes `__half*` (it is the + // G3 parity form, matching `ref.py`'s f16 return) while + // `escha_h128_in_batched` writes `float*`. Using the single form with an + // f32 scratch buffer silently produced a zero activation — the f16 pairs + // reinterpret as denormal-scale f32 — which then flowed through the GEMV + // as zeros and left only the bias in the output. Caught by + // `test_escha_dense_linear_gpu_vs_cpu` before any of this was wired. + let ic = lin.w.k; + let oc = lin.w.m; + gpu.escha_h128_batched( + "escha_h128_in_batched", + x, + &lin.rin, + &lin.ids0, + xh, + ic, + 1, + rdna_compute::EschaXGroup::Broadcast, + )?; + match lin.ptr0.as_ref() { + // NATIVE: the trellis code decoded inside the GEMV, served as the + // degenerate one-expert case of the indexed kernel. + Some(ptr0) => { + let tk = match lin.w.gpu_dtype { + DType::Escha2T16 => 2u32, + DType::Escha3T16 => 3u32, + other => { + return Err(HipError::new( + 0, + &format!("escha dense native: unexpected dtype {other:?}"), + )) + } + }; + gpu.escha_gemv_native_moe_k8_indexed_batched( + ptr0, &lin.ids0, xh, mid, oc, ic, 1, tk, true, + )?; + } + // Plain `weight_gemv`, NOT `weight_gemv_prerotated`. The decoded + // stores are Q8_0/F16, neither of which wants an FWHT rotation, and + // `xh` is already H128-rotated — the prerotated path would rotate a + // rotated activation. + None => hipfire_runtime::llama::weight_gemv(gpu, &lin.w, xh, mid)?, + } + gpu.escha_h128_batched( + "escha_h128_out_batched", + mid, + &lin.rout, + &lin.ids0, + y, + oc, + 1, + rdna_compute::EschaXGroup::Broadcast, + )?; + if let Some(b) = lin.bias.as_ref() { + gpu.add_inplace_f32(y, b)?; + } + Ok(()) +} + +/// Upload a bias as f32 regardless of whether it was stored F16 or F32. +fn read_bias_f32(gpu: &mut Gpu, qt: u8, data: &[u8], oc: usize) -> HipResult { + let mut v = Vec::with_capacity(oc); + match qt { + 1 => { + for c in data.chunks_exact(2) { + v.push(hipfire_runtime::llama::f16_to_f32(u16::from_le_bytes([ + c[0], c[1], + ]))); + } + } + _ => { + for c in data.chunks_exact(4) { + v.push(f32::from_le_bytes([c[0], c[1], c[2], c[3]])); + } + } + } + v.truncate(oc); + gpu.upload_f32(&v, &[oc]) +} + +#[allow(clippy::too_many_arguments)] +fn decode_projection( + hfq: &HfqFile, + gpu: &mut Gpu, + p: &str, + proj: &str, + expert_ids: &[usize], + n_exp: usize, + shape: (usize, usize), + store: EschaWeightStore, + resolve: NameResolver, + namer: LeafNamer, +) -> HipResult<(Vec, GpuTensor)> { + let (ic, oc) = shape; + let name = namer(p, proj, "code"); + let (info, data) = find(hfq, &name, resolve) + .ok_or_else(|| HipError::new(0, &format!("escha: tensor not found: {name}")))?; + let k = k_from_quant_type(info.quant_type, &name)?; + + let words_per_expert = (ic / 16) * (oc / 16) * 16 * k as usize; + let bytes_per_expert = words_per_expert * 2; + if data.len() != n_exp * bytes_per_expert { + return Err(HipError::new( + 0, + &format!( + "escha: {name} is {} bytes, expected {} for {n_exp} experts of {ic}x{oc} K={k}", + data.len(), + n_exp * bytes_per_expert + ), + )); + } + + // Reject out-of-range ids BEFORE allocating anything, so the error path + // has nothing to unwind. (Previously this check sat inside the decode loop + // and had to free the staging buffers by hand.) + if let Some(&bad) = expert_ids.iter().find(|&&x| x >= n_exp) { + return Err(HipError::new( + 0, + &format!("escha: expert id {bad} out of range for {n_exp} experts ({name})"), + )); + } + + // ONE buffer for the whole projection. See `PackedExpertOwners`: the 2 MiB + // allocation granule is charged once here instead of once per expert. + let (slot_elems, slot_dtype) = slot_extent(store, ic, oc, k as usize); + let total_elems = slot_elems + .checked_mul(expert_ids.len()) + .ok_or_else(|| HipError::new(0, &format!("escha: {name} packed size overflow")))?; + let owner = gpu.alloc_tensor(&[total_elems], slot_dtype)?; + + // ── Native: there is nothing to decode ──────────────────────────────── + // The slot IS the code. No `escha_decode_tiles`, no transpose, no + // requantise, no staging buffers — the bytes go from the file to their + // final resting place and the GEMV decodes them per token. This is also + // why an escha model loads faster on this store than on Q8_0: the decode + // that used to run 20 480 times at load does not run at all. + // + // The code keeps escha's own IN-major `[ic/16, oc/16, 16*k]` tile grid. + // The out-major transpose the other three stores fold in is not skipped + // here so much as absorbed: the fused GEMV addresses tiles in the in-major + // grid directly (see `escha_moe_gemv_native.hip`). + if store == EschaWeightStore::Native { + debug_assert_eq!(slot_elems, bytes_per_expert); + let mut out = Vec::with_capacity(expert_ids.len()); + for (slot, &x) in expert_ids.iter().enumerate() { + let src = &data[x * bytes_per_expert..(x + 1) * bytes_per_expert]; + let buf = owner.sub_offset(slot * slot_elems, slot_elems); + if let Err(error) = gpu.hip.memcpy_htod(&buf.buf, src) { + let _ = gpu.free_tensor(owner); + return Err(error); + } + out.push(WeightTensor { + buf, + gpu_dtype: slot_dtype, + // `m` / `k` stay the LOGICAL matrix shape, exactly as on every + // other store, because that is what the executor passes to the + // GEMV. The dtype is what says the bytes are trellis code. + m: oc, + k: ic, + row_stride: 0, + paro: None, + awq_scale: None, + }); + } + return Ok((out, owner)); + } + + // `escha_decode_tiles` validates `code.numel()` in SHORTS, so the staging + // tensor's logical length must be the i16 count (F16 gives the right + // 2-bytes-per-element sizing; the payload is trellis code, not floats). + let code_stage = match gpu.alloc_tensor(&[words_per_expert], DType::F16) { + Ok(t) => t, + Err(error) => { + let _ = gpu.free_tensor(owner); + return Err(error); + } + }; + let bare = match gpu.alloc_tensor(&[ic * oc], DType::F16) { + Ok(t) => t, + Err(error) => { + let _ = gpu.free_tensor(code_stage); + let _ = gpu.free_tensor(owner); + return Err(error); + } + }; + + let mut out = Vec::with_capacity(expert_ids.len()); + let mut decode = |gpu: &mut Gpu| -> HipResult<()> { + for (slot, &x) in expert_ids.iter().enumerate() { + let src = &data[x * bytes_per_expert..(x + 1) * bytes_per_expert]; + gpu.hip.memcpy_htod(&code_stage.buf, src)?; + gpu.escha_decode_tiles(&code_stage, &bare, ic as u32, oc as u32, k)?; + + // Non-owning window onto this expert's slice of the layer blob. + // The device pointer this yields is what lands in + // `expert_{gate_up,down}_ptrs`, so the indexed GEMV addresses the + // expert exactly as it did when each slot was its own allocation. + let buf = owner.sub_offset(slot * slot_elems, slot_elems); + + // The transpose to hipfire's OUT-major slot lives here, folded + // into the store. See the module docs. + match store { + EschaWeightStore::Q8_0 => gpu.escha_bare_to_q8_0(&bare, &buf, ic, oc)?, + EschaWeightStore::F32 => gpu.escha_bare_to_f32(&bare, &buf, ic, oc)?, + EschaWeightStore::F16 => gpu.escha_bare_to_f16(&bare, &buf, ic, oc)?, + // Returned above, before any staging buffer was allocated. + EschaWeightStore::Native => unreachable!("native store returns before decoding"), + } + out.push(WeightTensor { + buf, + gpu_dtype: slot_dtype, + m: oc, + k: ic, + row_stride: 0, + paro: None, + awq_scale: None, + }); + } + Ok(()) + }; + let result = decode(gpu); + let _ = gpu.free_tensor(code_stage); + let _ = gpu.free_tensor(bare); + if let Err(error) = result { + let _ = gpu.free_tensor(owner); + return Err(error); + } + Ok((out, owner)) +} + +#[cfg(test)] +mod tests { + use super::escha_dense_leaf; + use super::escha_leaf; + use super::slot_extent; + use super::EschaWeightStore; + use rdna_compute::DType; + + /// The dense and MoE leaf namers are NOT interchangeable, and using the + /// wrong one yields a name that simply is not in the file — which reads + /// as "this projection is not escha" and silently takes the plain-weight + /// path. Pin both against names copied out of the real checkpoints. + #[test] + fn dense_and_moe_leaf_names_do_not_collide() { + let p = "model.language_model.layers.0"; + // 27B dense export, verbatim from model.safetensors.index.json. + assert_eq!( + escha_dense_leaf(p, "linear_attn.in_proj_qkv", "code"), + "model.language_model.layers.0.linear_attn.in_proj_qkv.escha_code" + ); + assert_eq!( + escha_dense_leaf(p, "mlp.gate_proj", "rin"), + "model.language_model.layers.0.mlp.gate_proj.escha_rin" + ); + assert_eq!( + escha_dense_leaf(p, "self_attn.q_proj", "rout"), + "model.language_model.layers.0.self_attn.q_proj.escha_rout" + ); + // 35B MoE export gathers experts under one name instead. + assert_eq!( + escha_leaf(p, "gate_up", "code"), + "model.language_model.layers.0.mlp.experts.gate_up_proj.escha_code" + ); + assert_ne!( + escha_dense_leaf(p, "mlp.gate_proj", "code"), + escha_leaf(p, "gate", "code") + ); + } + + /// The packing arithmetic, at the real A3B shapes, against the sizes the + /// allocator-granularity diagnosis is built on. A slot stride that is not + /// a multiple of the Q8_0 block (34 B) or that disagrees with + /// `escha_bare_to_q8_0`'s own `oc * (ic/32) * 34` would put every expert + /// after slot 0 at a wrong offset — plausible, finite, wrong weights. + #[test] + fn q8_0_slot_extent_matches_the_a3b_projection_sizes() { + // gate_up: ic = hidden = 2048, oc = 2*mi = 1024. + let (gu, gu_dtype) = slot_extent(EschaWeightStore::Q8_0, 2048, 1024, 2); + assert_eq!(gu_dtype, DType::Q8_0); + assert_eq!(gu, 2_228_224, "gate_up slot is 2.125 MiB"); + // down: ic = mi = 512, oc = hidden = 2048. + let (dn, _) = slot_extent(EschaWeightStore::Q8_0, 512, 2048, 3); + assert_eq!(dn, 1_114_112, "down slot is 1.0625 MiB"); + // 256 experts x 40 layers x both projections = the 34.2 GB of real + // weight bytes the 67.9 GB of granules was hiding. + assert_eq!((gu + dn) * 256 * 40, 34_225_520_640); + } + + /// The Native (Phase 2) store's slot arithmetic, at the real A3B shapes, + /// against the code sizes on disk. These are the numbers the whole Phase-2 + /// memory claim rests on, so they are asserted rather than asserted-about: + /// a wrong stride would place every expert after slot 0 at a wrong offset + /// and — because trellis code decodes to *something* from any bit pattern + /// — produce finite, plausible, wrong weights rather than a fault. + #[test] + fn native_slot_extent_matches_the_a3b_code_sizes() { + // gate_up: ic = hidden = 2048, oc = 2*mi = 1024, K=2. + let (gu, gu_dtype) = slot_extent(EschaWeightStore::Native, 2048, 1024, 2); + assert_eq!(gu_dtype, DType::Escha2T16); + assert_eq!(gu_dtype.size(), 1, "escha code offsets are byte offsets"); + assert_eq!(gu, 524_288, "gate_up code is 512 KiB (2.00 bpw)"); + // down: ic = mi = 512, oc = hidden = 2048, K=3. + let (dn, dn_dtype) = slot_extent(EschaWeightStore::Native, 512, 2048, 3); + assert_eq!(dn_dtype, DType::Escha3T16); + assert_eq!(dn, 393_216, "down code is 384 KiB (3.00 bpw)"); + + // Exactly 2.00 / 3.00 bits per weight — the format's own figures, so + // this also catches a tile-count or word-count slip. + assert_eq!(gu * 8, 2048 * 1024 * 2); + assert_eq!(dn * 8, 512 * 2048 * 3); + + // 256 experts x 40 layers x both projections. The Q8_0 store's 34.2 GB + // of the same weights is 3.73x this. + assert_eq!((gu + dn) * 256 * 40, 9_395_240_960); + let (q8_gu, _) = slot_extent(EschaWeightStore::Q8_0, 2048, 1024, 2); + let (q8_dn, _) = slot_extent(EschaWeightStore::Q8_0, 512, 2048, 3); + assert!((q8_gu + q8_dn) > 3 * (gu + dn)); + } + + /// `sub_offset` counts in `dtype.size()` units. Q8_0 is a byte dtype, so + /// the Q8_0 stride is a byte stride while F16/F32 strides are element + /// counts. Getting that wrong scales every offset by 2 or 4. + #[test] + fn slot_extent_is_in_dtype_units_not_bytes() { + let (f32_elems, f32_dtype) = slot_extent(EschaWeightStore::F32, 2048, 1024, 2); + assert_eq!(f32_dtype, DType::F32); + assert_eq!(f32_elems, 2048 * 1024); + assert_eq!(f32_elems * DType::F32.size(), 8 * 1024 * 1024); + + let (f16_elems, f16_dtype) = slot_extent(EschaWeightStore::F16, 2048, 1024, 2); + assert_eq!(f16_dtype, DType::F16); + assert_eq!(f16_elems, 2048 * 1024); + assert_eq!(f16_elems * DType::F16.size(), 4 * 1024 * 1024); + + let (q8_elems, q8_dtype) = slot_extent(EschaWeightStore::Q8_0, 2048, 1024, 2); + assert_eq!(q8_dtype.size(), 1, "Q8_0 offsets are byte offsets"); + assert_eq!(q8_elems * q8_dtype.size(), 2_228_224); + } + + /// Every A3B slot stride is a multiple of 1024 B, so no expert view is + /// less aligned than the 2 MiB-granule allocation it replaces and the + /// kernels' vector loads are undisturbed. + #[test] + fn a3b_slot_strides_are_widely_aligned() { + for (ic, oc) in [(2048usize, 1024usize), (512, 2048)] { + let (elems, dtype) = slot_extent(EschaWeightStore::Q8_0, ic, oc, 2); + assert_eq!(elems * dtype.size() % 1024, 0, "{ic}x{oc} stride alignment"); + } + } +} diff --git a/crates/hipfire-arch-qwen35/src/qwen35/forward.rs b/crates/hipfire-arch-qwen35/src/qwen35/forward.rs index 00c86bd151..0c7e116709 100644 --- a/crates/hipfire-arch-qwen35/src/qwen35/forward.rs +++ b/crates/hipfire-arch-qwen35/src/qwen35/forward.rs @@ -621,6 +621,12 @@ fn moe_ffn_decode_impl( has_paro_shared: ffn.paro_shared.is_some(), per_expert_gate_up, per_expert_down, + // Escha-W2: same single source of truth as `MoeParams::escha` below. + // The loader has already turned the trellis experts into Q8_0, so the + // transform tables are the ONLY remaining evidence that this layer is + // escha — and they are also what makes the escha indexed executor + // callable, which is exactly what the resolver's Q8_0 arm gates on. + routed_escha_transforms: ffn.escha.is_some() && super::escha::escha_indexed_route_enabled(), }; // Resolution is owned by the MoeFamily (Ship 4.1). The model passes only // the dtype snapshot + k; the executor computes MoeResolution from MoeDtypes. @@ -712,6 +718,11 @@ fn moe_ffn_decode_impl( topk_indices: s.topk_indices, topk_weights: s.topk_weights, down_expanded: s.down_expanded, + // Escha-W2 (Task 10). `Some` only for Escha-W2 layers; it is both the + // transform tables the H128-wrapped routed executor needs AND the + // layer's escha marker (the loader has already turned the trellis + // experts into Q8_0, so no routed dtype says "escha" any more). + escha: ffn.escha.as_ref().map(|e| e.refs()), }; // Build one DispatchCtx per token (the family threads it through every // inner GEMV — no internal DispatchCtx::new reconstructions). @@ -833,7 +844,7 @@ fn forward_from_x_gpu( // Run the production pipeline forward_scratch_layers( - gpu, weights, config, pos, kv_cache, dn_state, &scratch, None, None, + gpu, weights, config, pos, kv_cache, dn_state, &scratch, None, None, true, )?; // DEBUG_LAYERS: dump per-layer residual norms @@ -896,6 +907,19 @@ pub struct Qwen35Scratch { pub gate_ffn: GpuTensor, // [hidden_dim] pub up: GpuTensor, // [hidden_dim] pub ffn_hidden: GpuTensor, // [hidden_dim] + /// Escha trellis scratch: the H128-rotated activation `xh` feeding one + /// projection's GEMV. Sized to the LARGEST `ic` any projection uses + /// (hidden_dim, for down_proj) so one buffer serves them all. + /// + /// Separate from `x`/`tmp` because each escha projection rotates the SAME + /// input with its OWN `rin` — the whole reason the fused MQ paths cannot + /// serve a trellis layer. Reusing the layer input here would corrupt the + /// next projection's source. + /// + /// `mid` needs no buffer: `escha_h128_out_batched` stages its 128-lane + /// block into LDS and syncs before writing, so it is safe in place and + /// the projection's own output tensor serves as both. + pub escha_xh: GpuTensor, // [max(hidden_dim, dim)] pub ffn_out: GpuTensor, // [dim] // Sampling @@ -1068,6 +1092,9 @@ impl Qwen35Scratch { gate_ffn: tracked_tensor!(gpu.alloc_tensor(&[config.hidden_dim], DType::F32)), up: tracked_tensor!(gpu.alloc_tensor(&[config.hidden_dim], DType::F32)), ffn_hidden: tracked_tensor!(gpu.alloc_tensor(&[config.hidden_dim], DType::F32)), + escha_xh: tracked_tensor!( + gpu.alloc_tensor(&[config.hidden_dim.max(config.dim)], DType::F32) + ), ffn_out: tracked_tensor!(gpu.alloc_tensor(&[dim], DType::F32)), logits: tracked_tensor!(gpu.alloc_tensor(&[config.vocab_size], DType::F32)), @@ -1385,6 +1412,19 @@ fn ar_graph_trace_enabled() -> bool { }) } +/// Whether this forward may capture or replay the plain-AR hipGraph. +/// +/// `emit_logits == false` produces a DIFFERENT kernel sequence (no lm_head), +/// so a logits-suppressed forward must neither capture the graph nor replay +/// one captured from a full forward: replaying a full graph would re-run the +/// lm_head the caller meant to skip, and capturing a suppressed one would +/// leave a later plain decode replaying a graph that never writes +/// `scratch.logits` — stale logits, no error, no NaN. +#[inline] +fn ar_graph_eligible_for(requested: bool, compact_offset: usize, emit_logits: bool) -> bool { + emit_logits && ar_graph_eligible_for_kv(requested, compact_offset) +} + #[inline] fn ar_graph_eligible_for_kv(requested: bool, compact_offset: usize) -> bool { // The captured single-token route is built while compact_offset is zero. @@ -1406,6 +1446,48 @@ pub fn forward_scratch( kv_cache: &mut llama::KvCache, dn_state: &mut DeltaNetState, scratch: &Qwen35Scratch, +) -> HipResult<()> { + forward_scratch_opts( + gpu, weights, config, token, pos, kv_cache, dn_state, scratch, true, + ) +} + +/// [`forward_scratch`] with the final lm_head projection made optional. +/// +/// `emit_logits == false` runs every layer, the KV/DeltaNet state updates and +/// the final output norm exactly as before, and stops short of the vocabulary +/// GEMV. `scratch.tmp` (the post-output-norm hidden state) is still written, so +/// the per-token-hidden extraction in `forward_prefill_batch`'s fallback is +/// unaffected; only `scratch.logits` is left holding the previous call's value. +/// +/// # Why this exists +/// +/// A model that fails batched-prefill admission prefills through a per-token +/// `forward_scratch` loop, and every one of those tokens computed a full +/// `[vocab, hidden]` projection whose result the next token immediately +/// overwrote. Only the LAST token's logits are ever read. On escha-35b +/// (`vocab = 248 320`, Q8_0 lm_head) `rocprofv3 --kernel-trace` prices that at +/// **2.32 ms of a 24.55 ms prefill token — 508 MB of weight traffic, 9.5 % of +/// the token** — spent producing a value that is discarded. +/// +/// It is a real saving rather than a bookkeeping one because this model is +/// bandwidth-bound: the lm_head GEMV moves the whole 508 MB weight matrix per +/// call and achieves ~219 GB/s doing it. +/// +/// CONTRACT: pass `false` only when the caller will not read `scratch.logits` +/// for that position. The prefill fallback passes `true` for the final token, +/// which is the only one whose logits survive the loop. +#[allow(clippy::too_many_arguments)] +pub fn forward_scratch_opts( + gpu: &mut Gpu, + weights: &Qwen35Weights, + config: &Qwen35Config, + token: u32, + pos: usize, + kv_cache: &mut llama::KvCache, + dn_state: &mut DeltaNetState, + scratch: &Qwen35Scratch, + emit_logits: bool, ) -> HipResult<()> { let required_tokens = checked_kv_end(pos, 1, "forward_scratch")?; // Grow before any possible AR graph capture/replay. Stable virtual @@ -1522,8 +1604,19 @@ pub fn forward_scratch( // capture or replay in a non-sequential context. An ineligible call also // INVALIDATES any captured graph (forces re-capture on the next plain call). let requested_graph_eligible = std::mem::replace(&mut gpu.graphs.ar_graph_eligible, true); - let graph_eligible = - ar_graph_eligible_for_kv(requested_graph_eligible, kv_cache.compact_offset); + // A logits-suppressed forward emits a DIFFERENT kernel sequence (no + // lm_head), so it must never capture the plain-AR graph nor replay one + // captured from a full forward — a replay would either re-run the lm_head + // this call meant to skip or, captured the other way round, leave a later + // plain decode reading stale logits. Today's only `emit_logits == false` + // caller (the prefill fallback) already sets `ar_graph_eligible = false`; + // this makes the invariant a property of the function rather than of its + // callers, so a new caller cannot reintroduce the hazard. + let graph_eligible = ar_graph_eligible_for( + requested_graph_eligible, + kv_cache.compact_offset, + emit_logits, + ); // Redline's plain-AR capture/replay has the same eligibility contract as // the AR HipGraph. MTP/spec re-seed and verify calls must not contaminate // or consume the immutable single-token replay sequence. @@ -1609,7 +1702,16 @@ pub fn forward_scratch( gpu.hip .memcpy_htod(&scratch.pos_buf, &pos_i32.to_ne_bytes())?; forward_scratch_layers( - gpu, weights, config, pos, kv_cache, dn_state, scratch, None, None, + gpu, + weights, + config, + pos, + kv_cache, + dn_state, + scratch, + None, + None, + emit_logits, )?; gpu.graphs.ar_forward_kernel_dirty = false; } else if use_graph { @@ -1630,7 +1732,16 @@ pub fn forward_scratch( gpu.active_stream.as_ref().unwrap(), )?; forward_scratch_layers( - gpu, weights, config, pos, kv_cache, dn_state, scratch, None, None, + gpu, + weights, + config, + pos, + kv_cache, + dn_state, + scratch, + None, + None, + emit_logits, )?; gpu.graphs.end_graph_capture( &gpu.hip, @@ -1653,7 +1764,16 @@ pub fn forward_scratch( gpu.hip .memcpy_htod(&scratch.pos_buf, &pos_i32.to_ne_bytes())?; forward_scratch_layers( - gpu, weights, config, pos, kv_cache, dn_state, scratch, None, None, + gpu, + weights, + config, + pos, + kv_cache, + dn_state, + scratch, + None, + None, + emit_logits, )?; } if gpu.replay.should_auto_finalize_capture() { @@ -1764,6 +1884,7 @@ pub fn forward_scratch_with_hidden( scratch, Some(hidden_rb), None, + true, )?; hidden_rb.advance_head(); Ok(()) @@ -1794,7 +1915,7 @@ pub fn forward_scratch_embed( }; gpu.hip.memcpy_htod(&scratch.x.buf, bytes)?; forward_scratch_layers( - gpu, weights, config, pos, kv_cache, dn_state, scratch, None, None, + gpu, weights, config, pos, kv_cache, dn_state, scratch, None, None, true, ) } @@ -1885,6 +2006,7 @@ pub fn forward_scratch_mrope( scratch, None, Some(mc), + true, ) } @@ -1936,11 +2058,21 @@ pub fn forward_scratch_embed_mrope( scratch, None, Some(mc), + true, ) } // ── Forward scratch layers (dispatch family version) ──────────────────── +/// `emit_logits == false` runs the whole layer stack and the final output +/// norm but skips the lm_head GEMV. See [`forward_scratch_opts`] for why. +/// Honoured on BOTH arms: the hand-written one below and the lowered super-op +/// executor (`forward_scratch_layers_lowered`), which is the DEFAULT +/// (`HIPFIRE_FORWARD_LOWERED` opts out with `0`). Threading it through only +/// the hand arm makes the saving invisible on the default path — that mistake +/// was made once here and caught by counting lm_head dispatches in the kernel +/// trace, not by the wall clock, which moved 0.3 %. +#[allow(clippy::too_many_arguments)] fn forward_scratch_layers( gpu: &mut Gpu, weights: &Qwen35Weights, @@ -1951,6 +2083,7 @@ fn forward_scratch_layers( s: &Qwen35Scratch, hidden_rb: Option<&mut HiddenStateRingBuffer>, mrope: Option<&MropeCtx>, + emit_logits: bool, ) -> HipResult<()> { // #397 Ship 6 — forward-as-pipeline. When HIPFIRE_FORWARD_LOWERED=1, route // single-GPU decode through the lowered super-op executor. Skipped when a @@ -1964,7 +2097,16 @@ fn forward_scratch_layers( // silently reinstate sequential positions. VL therefore always takes the // hand arms below, which DO branch on `mrope`. if forward_lowered_enabled() && hidden_rb.is_none() && mrope.is_none() { - return forward_scratch_layers_lowered(gpu, weights, config, pos, kv_cache, dn_state, s); + return forward_scratch_layers_lowered( + gpu, + weights, + config, + pos, + kv_cache, + dn_state, + s, + emit_logits, + ); } let k_dim = config.linear_num_key_heads * config.linear_key_head_dim; @@ -2611,9 +2753,12 @@ fn forward_scratch_layers( dump_hidden_localize(gpu, &s.x, 1, pos, config.dim, layer_idx, "pertoken"); } - // Final norm + logits into scratch.logits + // Final norm — ALWAYS. `s.tmp` is the post-norm hidden state and is read + // by callers that never look at the logits (per-token hidden extraction in + // the prefill fallback, hidden-ring staging), so it is not part of what + // `emit_logits == false` skips. gpu.rmsnorm_f32(&s.x, &weights.output_norm, &s.tmp, config.norm_eps)?; - { + if emit_logits { let ctx = DispatchCtx::new(gpu); let wr = weights.output.dispatch_ref(); let step = Step::Gemv { @@ -4754,6 +4899,39 @@ fn qkv_from_prerotated_mq( } #[allow(clippy::too_many_arguments)] +/// True when all four QKVZA weights sit in the container +/// `fused_qkvza_hfq4g256` actually reads. +/// +/// qt=6 (HFQ4G256) and qt=13 (MQ4G256) share one 136 B group layout, so MQ4 +/// has always borrowed HFQ4's kernel and that is correct — only the +/// activations differ, which is what `precomputed_attn_x_rot` handles. NO +/// other MQ container shares it. +/// +/// [`qkvza_from_prerotated_mq`] hardcodes that kernel with no dtype check, so +/// without this predicate any other MQ container on `wqkv` — or an F16 +/// sibling, which is how escha-35b stores `w_alpha`/`w_beta` — is read at +/// HFQ4 stride. Measured: down-quantising ONLY `in_proj_qkv` on escha-35b +/// scored KLD 12.63 / PPL 2,375,141 against a 7.68 baseline, identically +/// under MQ6G256, MQ6G256V2 and MQ4G256V2, while `out_proj` (not in this +/// launch) was unaffected at KLD 0.0076. Finite, fluent, wrong. +/// +/// Same failure family as `prefill::all_q8_0`, which was added for the Q8_0 +/// arms after escha-35b made mixed layers reachable; the per-token MQ path +/// never got the equivalent. +fn qkvza_hfq4_container( + wqkv: &WeightTensor, + wz: &WeightTensor, + w_beta: &WeightTensor, + w_alpha: &WeightTensor, +) -> bool { + [wqkv, wz, w_beta, w_alpha].iter().all(|w| { + matches!( + w.gpu_dtype, + rdna_compute::DType::HFQ4G256 | rdna_compute::DType::MQ4G256 + ) + }) +} + fn qkvza_from_prerotated_mq( gpu: &mut Gpu, wqkv: &WeightTensor, @@ -4810,7 +4988,213 @@ fn op_code(op: &OpBinding) -> u32 { op.weights.first().map(|w| w.0).unwrap_or(u32::MAX) } +/// Run one projection op for a layer whose weights are escha trellis codes. +/// +/// Returns `Ok(false)` when the layer is not escha, so the caller falls +/// through to its ordinary dispatch untouched. +/// +/// WHY THIS BYPASSES THE FUSED PATHS ENTIRELY: every escha projection rotates +/// the SAME normed input with its OWN `rin` before its GEMV. FusedQkv / +/// FusedQkvza / gate_up exist precisely to share one rotated activation across +/// several weights, so there is nothing for them to share here — and they +/// cannot read a trellis code in any case. A layer is all-escha or none +/// (`need_eproj` enforces that at load), so the bypass is wholesale. +/// +/// `in_proj_a` / `in_proj_b` are NOT coded — escha's `ignore` list keeps them +/// plain — so PROJ_QKVZA runs those two through the normal GEMV. +fn escha_run_proj( + gpu: &mut Gpu, + op: &OpBinding, + layer: &LayerWeights, + s: &Qwen35Scratch, + config: &Qwen35Config, +) -> Result { + let hip = |e: hip_bridge::HipError| DispatchError::Hip(e.to_string()); + match (op_code(op), layer) { + (q35_op::PROJ_QKVZA, LayerWeights::DeltaNet(l)) => { + let Some(e) = l.escha.as_ref() else { return Ok(false) }; + // Plain RMSNorm, NOT the fused rmsnorm+rotate: escha applies its + // own H128 per projection and a pre-rotated input would be + // rotated twice. + gpu.rmsnorm_f32(&s.x, &l.attn_norm, &s.tmp, config.norm_eps) + .map_err(hip)?; + e.qkv.forward(gpu, &l.wqkv, &e.ids, &s.tmp, &s.escha_xh, &s.dn_qkv, &s.dn_qkv, 1, None) + .map_err(hip)?; + e.z.forward(gpu, &l.wz, &e.ids, &s.tmp, &s.escha_xh, &s.dn_z, &s.dn_z, 1, None) + .map_err(hip)?; + hipfire_runtime::llama::weight_gemv(gpu, &l.w_beta, &s.tmp, &s.dn_beta).map_err(hip)?; + hipfire_runtime::llama::weight_gemv(gpu, &l.w_alpha, &s.tmp, &s.dn_alpha) + .map_err(hip)?; + Ok(true) + } + (q35_op::PROJ_QKV, LayerWeights::FullAttn(l)) => { + let Some(e) = l.escha.as_ref() else { return Ok(false) }; + gpu.rmsnorm_f32(&s.x, &l.attn_norm, &s.tmp, config.norm_eps) + .map_err(hip)?; + e.q.forward(gpu, &l.wq, &e.ids, &s.tmp, &s.escha_xh, &s.fa_q_full, &s.fa_q_full, 1, None) + .map_err(hip)?; + e.k.forward(gpu, &l.wk, &e.ids, &s.tmp, &s.escha_xh, &s.fa_k, &s.fa_k, 1, None) + .map_err(hip)?; + e.v.forward(gpu, &l.wv, &e.ids, &s.tmp, &s.escha_xh, &s.fa_v, &s.fa_v, 1, None) + .map_err(hip)?; + Ok(true) + } + (q35_op::PROJ_GATE_UP, LayerWeights::DeltaNet(l)) => { + let Some(e) = l.escha.as_ref() else { return Ok(false) }; + gpu.rmsnorm_f32(&s.x, &l.ffn_norm, &s.tmp, config.norm_eps) + .map_err(hip)?; + e.gate.forward(gpu, &l.w_gate, &e.ids, &s.tmp, &s.escha_xh, &s.gate_ffn, &s.gate_ffn, 1, None) + .map_err(hip)?; + e.up.forward(gpu, &l.w_up, &e.ids, &s.tmp, &s.escha_xh, &s.up, &s.up, 1, None) + .map_err(hip)?; + Ok(true) + } + (q35_op::PROJ_GATE_UP, LayerWeights::FullAttn(l)) => { + let Some(e) = l.escha.as_ref() else { return Ok(false) }; + gpu.rmsnorm_f32(&s.x, &l.ffn_norm, &s.tmp, config.norm_eps) + .map_err(hip)?; + e.gate.forward(gpu, &l.w_gate, &e.ids, &s.tmp, &s.escha_xh, &s.gate_ffn, &s.gate_ffn, 1, None) + .map_err(hip)?; + e.up.forward(gpu, &l.w_up, &e.ids, &s.tmp, &s.escha_xh, &s.up, &s.up, 1, None) + .map_err(hip)?; + Ok(true) + } + _ => Ok(false), + } +} + +/// Escha counterpart of the residual ops. `out_proj` and `down_proj` write +/// into the residual stream, which the fused epilogue normally does in one +/// launch; here the projection and the accumulate are separate because the +/// trellis GEMV has no residual variant. +/// +/// Adding into `s.x` afterwards is exact, not an approximation — both the +/// residual and the projection output are plain f32 adds. +fn escha_run_resid( + gpu: &mut Gpu, + op: &OpBinding, + layer: &LayerWeights, + s: &Qwen35Scratch, +) -> Result { + let hip = |e: hip_bridge::HipError| DispatchError::Hip(e.to_string()); + let (esch, wo, w_down, dn_in) = match layer { + LayerWeights::DeltaNet(l) => match l.escha.as_ref() { + None => return Ok(false), + Some(e) => ( + (&e.o, &e.down), + &l.wo, + &l.w_down, + &s.dn_normed, + ), + }, + LayerWeights::FullAttn(l) => match l.escha.as_ref() { + None => return Ok(false), + Some(e) => ((&e.o, &e.down), &l.wo, &l.w_down, &s.fa_attn_out), + }, + _ => return Ok(false), + }; + match op_code(op) { + q35_op::RESID_WO => { + esch.0 + .forward(gpu, wo, escha_ids(layer), dn_in, &s.escha_xh, &s.o, &s.o, 1, None) + .map_err(hip)?; + gpu.add_inplace_f32(&s.x, &s.o).map_err(hip)?; + Ok(true) + } + q35_op::RESID_DOWN_SWIGLU => { + // SwiGLU first — the fused `weight_gemv_swiglu_residual` folds it + // in, but that kernel cannot read a trellis code. + gpu.silu_mul_f32(&s.gate_ffn, &s.up, &s.ffn_hidden) + .map_err(hip)?; + esch.1 + .forward( + gpu, + w_down, + escha_ids(layer), + &s.ffn_hidden, + &s.escha_xh, + &s.ffn_out, + &s.ffn_out, + 1, + None, + ) + .map_err(hip)?; + gpu.add_inplace_f32(&s.x, &s.ffn_out).map_err(hip)?; + Ok(true) + } + _ => Ok(false), + } +} + +/// The layer's shared zero `ids` table. Only called where `escha` is `Some`. +fn escha_ids(layer: &LayerWeights) -> &GpuTensor { + match layer { + LayerWeights::DeltaNet(l) => &l.escha.as_ref().expect("escha layer").ids, + LayerWeights::FullAttn(l) => &l.escha.as_ref().expect("escha layer").ids, + _ => unreachable!("escha_ids on a non-escha layer kind"), + } +} + +/// Add the escha dense export's additive output biases. +/// +/// Applied here, at the ONE exit of `run_proj`, rather than inside each +/// arm: every op has several branches that fill the same output buffers +/// (prerotated / scalar-prep / execute-steps), and a bias added in some +/// branches but not others is a silent wrong answer, not a crash. +/// +/// Correct AFTER the projection for the same reason it is correct after a +/// residual add — both are additive and commute. +/// +/// No-op for every model without escha biases, which is all of them but +/// the 27B. +fn apply_proj_biases( +gpu: &mut Gpu, +op: &OpBinding, +layer: &LayerWeights, +s: &Qwen35Scratch, +) -> Result<(), DispatchError> { + // `bias_add_f32` and not `add_inplace_f32`: the same helper serves decode + // (batch = 1) and the batched prefill path, so both use one primitive and + // cannot drift. + let add = |gpu: &mut Gpu, y: &GpuTensor, b: &GpuTensor| -> Result<(), DispatchError> { + let n = b.numel(); + gpu.bias_add_f32(y, b, 1, n) + .map_err(|e| DispatchError::Hip(e.to_string())) + }; + match (op_code(op), layer) { + (q35_op::PROJ_QKVZA, LayerWeights::DeltaNet(l)) => { + if let Some(b) = l.biases.as_ref() { + add(gpu, &s.dn_qkv, &b.qkv)?; + add(gpu, &s.dn_z, &b.z)?; + } + } + (q35_op::PROJ_QKV, LayerWeights::FullAttn(l)) => { + if let Some(b) = l.biases.as_ref() { + add(gpu, &s.fa_q_full, &b.q)?; + add(gpu, &s.fa_k, &b.k)?; + add(gpu, &s.fa_v, &b.v)?; + } + } + (q35_op::PROJ_GATE_UP, LayerWeights::DeltaNet(l)) => { + if let Some(b) = l.biases.as_ref() { + add(gpu, &s.gate_ffn, &b.gate)?; + add(gpu, &s.up, &b.up)?; + } + } + (q35_op::PROJ_GATE_UP, LayerWeights::FullAttn(l)) => { + if let Some(b) = l.biases.as_ref() { + add(gpu, &s.gate_ffn, &b.gate)?; + add(gpu, &s.up, &b.up)?; + } + } + _ => {} + } + Ok(()) +} + impl<'a> ForwardBindings for Qwen35Bindings<'a> { + + fn run_proj( &mut self, gpu: &mut Gpu, @@ -4819,6 +5203,9 @@ impl<'a> ForwardBindings for Qwen35Bindings<'a> { ) -> Result<(), DispatchError> { let s = self.s; let config = self.config; + if escha_run_proj(gpu, op, self.layer, s, config)? { + return apply_proj_biases(gpu, op, self.layer, s); + } let res: HipResult<()> = match op_code(op) { q35_op::PROJ_QKV => match self.layer { LayerWeights::FullAttn(l) => { @@ -4909,7 +5296,9 @@ impl<'a> ForwardBindings for Qwen35Bindings<'a> { )); } }; - if self.precomputed_attn_x_rot { + if self.precomputed_attn_x_rot + && qkvza_hfq4_container(wqkv, wz, w_beta, w_alpha) + { qkvza_from_prerotated_mq( gpu, wqkv, @@ -5016,9 +5405,12 @@ impl<'a> ForwardBindings for Qwen35Bindings<'a> { }, other => return Err(DispatchError::Hip(format!("unknown PROJ opcode {other}"))), }; - res.map_err(|e| DispatchError::Hip(e.to_string())) + res.map_err(|e| DispatchError::Hip(e.to_string()))?; + apply_proj_biases(gpu, op, self.layer, self.s) } + + fn run_residual_gemv( &mut self, gpu: &mut Gpu, @@ -5026,6 +5418,26 @@ impl<'a> ForwardBindings for Qwen35Bindings<'a> { op: &OpBinding, ) -> Result<(), DispatchError> { let s = self.s; + if escha_run_resid(gpu, op, self.layer, s)? { + let bias = match self.layer { + LayerWeights::DeltaNet(l) => l.biases.as_ref().map(|b| (&b.o, &b.down)), + LayerWeights::FullAttn(l) => l.biases.as_ref().map(|b| (&b.o, &b.down)), + _ => None, + }; + if let Some((bo, bdown)) = bias { + let which = match op_code(op) { + q35_op::RESID_WO => Some(bo), + q35_op::RESID_DOWN_SWIGLU => Some(bdown), + _ => None, + }; + if let Some(b) = which { + let n = b.numel(); + gpu.bias_add_f32(&s.x, b, 1, n) + .map_err(|e| DispatchError::Hip(e.to_string()))?; + } + } + return Ok(()); + } let res: HipResult<()> = (|| match op_code(op) { q35_op::RESID_WO => { let (wo, input) = match self.layer { @@ -5102,7 +5514,29 @@ impl<'a> ForwardBindings for Qwen35Bindings<'a> { } other => Err(HipError::new(0, &format!("unknown RESID opcode {other}"))), })(); - res.map_err(|e| DispatchError::Hip(e.to_string())) + res.map_err(|e| DispatchError::Hip(e.to_string()))?; + // `out_proj`/`o_proj` and `down_proj` write into the residual stream, + // so their bias lands on `s.x`. Adding it after the residual add is + // the same value as adding it before — both additive. + let s = self.s; + let bias = match self.layer { + LayerWeights::DeltaNet(l) => l.biases.as_ref().map(|b| (&b.o, &b.down)), + LayerWeights::FullAttn(l) => l.biases.as_ref().map(|b| (&b.o, &b.down)), + _ => None, + }; + if let Some((bo, bdown)) = bias { + let which = match op_code(op) { + q35_op::RESID_WO => Some(bo), + q35_op::RESID_DOWN_SWIGLU => Some(bdown), + _ => None, + }; + if let Some(b) = which { + let n = b.numel(); + gpu.bias_add_f32(&s.x, b, 1, n) + .map_err(|e| DispatchError::Hip(e.to_string()))?; + } + } + Ok(()) } fn run_norm( @@ -5904,6 +6338,7 @@ fn moe_combine_next_rms_enabled(gpu: &Gpu, weights: &Qwen35Weights, config: &Qwe /// to `forward_scratch_layers`'s hand arms (validated byte-identical via the /// external committed-token md5 gate). Builds a coarse-super-op `LayerProgram` /// per layer and runs it through the dispatch substrate's executor. +#[allow(clippy::too_many_arguments)] fn forward_scratch_layers_lowered( gpu: &mut Gpu, weights: &Qwen35Weights, @@ -5912,6 +6347,7 @@ fn forward_scratch_layers_lowered( kv_cache: &mut llama::KvCache, dn_state: &DeltaNetState, s: &Qwen35Scratch, + emit_logits: bool, ) -> HipResult<()> { let k_dim = config.linear_num_key_heads * config.linear_key_head_dim; let v_dim = config.linear_num_value_heads * config.linear_value_head_dim; @@ -5982,7 +6418,7 @@ fn forward_scratch_layers_lowered( // Final norm + logits into scratch.logits (mirrors forward_scratch_layers). gpu.rmsnorm_f32(&s.x, &weights.output_norm, &s.tmp, config.norm_eps)?; - { + if emit_logits { let ctx = DispatchCtx::new(gpu); let wr = weights.output.dispatch_ref(); let step = Step::Gemv { @@ -6119,6 +6555,24 @@ mod tests { assert_eq!(lower_variant(Q35Variant::FullAttnMoe).len(), 4); } + /// A logits-suppressed forward must never touch the plain-AR graph. + /// + /// Without this, the prefill fallback's `emit_logits = false` tokens could + /// replay a graph captured from a full forward (re-running the lm_head the + /// skip exists to remove — the optimisation silently does nothing), or + /// capture a logits-free graph that a later plain decode replays, leaving + /// `scratch.logits` holding a previous token's values with no error and no + /// NaN to catch it. + #[test] + fn logits_suppressed_forward_is_never_ar_graph_eligible() { + // Otherwise-perfect conditions: requested, no KV compaction. + assert!(ar_graph_eligible_for(true, 0, true)); + assert!(!ar_graph_eligible_for(true, 0, false)); + // emit_logits cannot RE-enable a forward the other conditions refuse. + assert!(!ar_graph_eligible_for(false, 0, true)); + assert!(!ar_graph_eligible_for(true, 128, true)); + } + #[test] fn ar_graph_is_ineligible_after_kv_compaction() { assert!(ar_graph_eligible_for_kv(true, 0)); diff --git a/crates/hipfire-arch-qwen35/src/qwen35/load.rs b/crates/hipfire-arch-qwen35/src/qwen35/load.rs index 8d8d317626..d41a15efaa 100644 --- a/crates/hipfire-arch-qwen35/src/qwen35/load.rs +++ b/crates/hipfire-arch-qwen35/src/qwen35/load.rs @@ -8,6 +8,8 @@ use super::config::f16_lm_head_mode_from_config; use super::config::F16LmHeadMode; use super::config::Qwen35Config; +use super::escha; +use super::escha::EschaWeightStore; use super::forward::layers_have_mq6_moe; use super::weights::dtype_from_quant_type; use super::weights::mixed_expert_tag; @@ -63,7 +65,11 @@ const _: () = assert!(QWEN35_NORM_BIAS == 1.0); // ─── Weight loading ───────────────────────────────────────────────────── -fn qwen35_tensor_name_candidates(name: &str) -> Vec { +/// Public so gates outside the loader (e.g. +/// `examples/test_escha_dense_linear_gpu_vs_cpu.rs`) resolve names through the +/// SAME aliasing the production load uses. A gate with its own copy of this +/// would stop testing what actually runs the moment the two drifted. +pub fn qwen35_tensor_name_candidates(name: &str) -> Vec { let mut out = Vec::with_capacity(4); let mut push = |s: String| { if !out.iter().any(|x| x == &s) { @@ -85,6 +91,26 @@ fn qwen35_tensor_name_candidates(name: &str) -> Vec { push(format!("model.{name}")); push(name.to_string()); } + + // Escha alias. A trellis-coded projection has NO `.weight` — it ships + // `escha_code` + `escha_rin` + `escha_rout`. Offering the code under the + // same lookup lets the ordinary `b.proj(...)` path find it and build a + // `WeightTensor` with dtype Escha2T16/3T16; `escha_sidecars` then attaches + // the rotation vectors. + // + // Ordered LAST on purpose: a checkpoint carrying both a real `.weight` and + // a stale `escha_code` must resolve to the weight, not silently prefer a + // code the runtime would then decode against the wrong activation. + if let Some(stem) = name.strip_suffix(".weight") { + let alias = format!("{stem}.escha_code"); + if alias.starts_with("model.") { + push(alias); + } else { + push(format!("model.language_model.{alias}")); + push(format!("model.{alias}")); + push(alias); + } + } out } @@ -306,6 +332,56 @@ fn load_weight_tensor_raw( awq_scale: None, }) } + 42 | 43 => { + // Escha-W2 trellis code, kept VERBATIM — the 2-bit/3-bit stream is + // decoded inside the GEMV, never at load. That is the whole point + // of the format: an 11.16 GB resident 27B instead of 22.63 GB + // folded, at better quality (PPL 11.8654 vs 13.6957). + // + // Opaque raw buffer like the MQ arms above, but the resemblance + // ends there: an escha weight is NOT self-contained. It needs its + // `escha_rin_eff`/`escha_rout_eff` vectors and an H128 on both + // sides of the GEMV, which is why `EschaDenseLinear` exists and + // why the fused MQ paths (FusedQkv/FusedQkvza/gate_up) CANNOT + // consume one — each projection needs its own rin-rotated + // activation. A layer holding these must route through + // `escha::escha_dense_linear_forward`. + // + // `m`/`k` are the logical output/input dims; the buffer length is + // the tile-packed code, not m*k of anything. + // + // TILE GRID TRANSPOSED HERE, kt-major -> nt-major. Every escha + // kernel holds one output tile column `nt` fixed and walks `kt`; + // in the checkpoint's order consecutive `kt` are a full tile-row + // apart (139 KB on the 27B's gate_proj), so each step is a fresh + // 64-bit address against a cold line. Adjacent instead: measured + // 243.9 -> 186.3 us on the decode GEMV, 24%. + // + // ONLY THIS (DENSE) LOADER PERMUTES. MoE experts come through + // `escha::load_escha_moe_experts` and stay kt-major, which is why + // the kernels take an `nt_major` flag rather than assuming a + // layout — the dense call sites pass `true`, MoE passes `false`. + // + // At LOAD, not in the converter, so the payload stays verbatim + // from upstream: no re-convert, no re-upload, no format version. + // Whole tiles move and their contents are untouched, so every + // decoded weight is identical. Gated on mean KLD = 0.000000. + let permuted = escha_tiles_to_nt_major(data, m, k, quant_type)?; + let buf = gpu.upload_raw(&permuted, &[permuted.len()])?; + Ok(WeightTensor { + buf, + gpu_dtype: if quant_type == 42 { + DType::Escha2T16 + } else { + DType::Escha3T16 + }, + m, + k, + row_stride: 0, + paro: None, + awq_scale: None, + }) + } 31 => { // MQ5-G256 — MagnumQuant FWHT-rotated 5-bit (168 bytes/group, 5.25 bpw). // Opaque raw buffer, same pattern as MQ4(13)/MQ6(15); the GEMV @@ -1435,6 +1511,8 @@ fn paro_load_moe_ffn( paro_shared: Some(shared), global_expert_dtypes: None, ep_dummy_buffers: Vec::new(), + // ParoQuant/paged, not Escha-W2. + escha: None, }) } @@ -3072,6 +3150,32 @@ pub fn preflight_weights_dense_tp( let conv = config.conv_kernel_dim; let validate_proj = |bare: &str, m: usize, k: usize| -> Result<(), String> { + // An escha-coded DENSE projection (Qwen3.8-27B) has no `.weight` at + // all — it ships `escha_code` + `escha_rin` + `escha_rout`. Validate + // the trio that the leaf contract makes REQUIRED (§1.4) and return; + // the shape lives in the code tensor's own dims and is checked when + // it is decoded, not here. + if let Some(stem) = bare.strip_suffix(".weight") { + let code = format!("{stem}.escha_code"); + if let Some((info, _)) = find_qwen35_tensor(hfq, &code) { + if info.quant_type == 42 || info.quant_type == 43 { + for leaf in ["escha_rin_eff", "escha_rout_eff"] { + let n = format!("{stem}.{leaf}"); + if find_qwen35_tensor(hfq, &n).is_none() { + return Err(format!( + "preflight: {code} is escha-coded but {n} is missing — an \ + incomplete escha linear must fail loudly at load, not \ + decode into noise" + )); + } + } + if k % 256 != 0 { + return Err(format!("preflight: {bare} K={k} not G256 aligned")); + } + return Ok(()); + } + } + } let (info, cand) = find_qwen35_tensor(hfq, bare) .ok_or_else(|| format!("preflight: missing tensor {bare}"))?; validate_mq4_proj_info(info, m, k, bare)?; @@ -3625,6 +3729,8 @@ pub fn load_weights_dense_tp_rank( w_gate: w_gate_opt.take().unwrap(), w_up: w_up_opt.take().unwrap(), w_down: w_down_opt.take().unwrap(), + biases: None, + escha: None, })) })(); match layer_res { @@ -3821,6 +3927,8 @@ pub fn load_weights_dense_tp_rank( w_gate: w_gate_opt.take().unwrap(), w_up: w_up_opt.take().unwrap(), w_down: w_down_opt.take().unwrap(), + biases: None, + escha: None, })) })(); match layer_res { @@ -4357,6 +4465,40 @@ fn e8_aos_to_soa(aos: &[u8], m: usize, k: usize) -> Vec { /// `[n_exp]` with dummy pointers for non-owned slots (which contribute 0 to the /// all-reduce because their gate_up is a zeroed buffer). Uniform files only — /// graded/AWQ EP would need the full per-expert dtype map and is rejected here. +/// Escha-W2 expert storage. Production (Phase 2) is `Native` — the trellis +/// code itself, 0.25/0.375 B/weight, decoded inside the routed GEMV. The three +/// other values select DECODING stores, each of which exists to make a +/// specific measurement possible rather than to be run: +/// +/// * `HIPFIRE_ESCHA_EXPERT_STORE=q8_0` (also `q8`) — Phase 1: transpose + +/// Q8_0 re-quantise, 1.0625 B/weight, 37.55 GB resident. This is the A/B arm +/// for every Phase-2 performance claim, and the arm every published Phase-1 +/// number (G4's Q8_0 arm, the G5 KLD headline) was measured on. It is also +/// the only routed store that works on the per-expert HOST route, so it is +/// what `HIPFIRE_ESCHA_INDEXED=0` needs. +/// * `HIPFIRE_ESCHA_EXPERT_STORE=f16` — 2 B/weight, weight-exact, ~64 GB of +/// experts. The arm the G5 KLD reference is built with. +/// * `HIPFIRE_ESCHA_EXPERT_STORE=f32` — 4 B/weight, ~129 GB of experts on the +/// 35B. Equally exact and does NOT fit; small-layer diagnostic only (the G4 +/// block gate uses it). +/// +/// `f16` and `f32` lose the indexed GPU-top-K path and run host-routed; +/// `native` REQUIRES it — there is no per-expert native GEMV, so an escha +/// layer that reaches the host route with this store fails loudly in +/// `GemvFamily::run_auto` (no plain GEMV exists for `RotationPlan::EschaH128`) +/// instead of running unrotated. See `qwen35/escha.rs`. +/// +/// An unrecognised value falls through to production rather than erroring, +/// matching every other developer var in this loader. +fn escha_weight_store() -> EschaWeightStore { + match hipfire_config::developer_var("HIPFIRE_ESCHA_EXPERT_STORE").as_deref() { + Ok("f32") | Ok("F32") => EschaWeightStore::F32, + Ok("f16") | Ok("F16") => EschaWeightStore::F16, + Ok("q8_0") | Ok("Q8_0") | Ok("q8") | Ok("Q8") => EschaWeightStore::Q8_0, + _ => EschaWeightStore::Native, + } +} + pub(crate) fn load_moe_ffn( hfq: &HfqFile, gpu: &mut Gpu, @@ -4371,7 +4513,30 @@ pub(crate) fn load_moe_ffn( .reap_keep .as_ref() .map(|r| r.expert_plan(layer_idx as usize)); + // Detect Escha-W2 BEFORE the EP-shard block. An escha layer carries one + // trellis code tensor per projection for all experts, not the per-expert + // `experts.{x}.gate_up_proj.weight` tensors the EP path fishes out by + // index — so if the EP block runs first it panics on a tensor that does + // not exist, and the escha refusal further down is never reached. The + // daemon sets an EP shard even on a single GPU, which is exactly how that + // happened: `hipfire bench` panicked with + // "tensor not found: layers.0.mlp.experts.0.gate_up_proj.weight". + // + // A single-rank shard splits nothing, so escha simply ignores it; only a + // genuine multi-rank split is refused (below, and again after the router + // load for the REAP keep-map case). + let escha_layer = escha::layer_is_escha(hfq, p, qwen35_tensor_name_candidates); let ep_shard = current_ep_expert_shard(); + let ep_shard = match (escha_layer, ep_shard) { + (true, Some((ref sc, _))) if sc.tp_size > 1 => { + return Err(HipError::new( + 0, + "qwen35: Escha-W2 routed experts do not support EP sharding across >1 rank (it re-maps experts across the per-expert tensors escha does not have)", + )) + } + (true, _) => None, + (_, other) => other, + }; if ep.is_some() && ep_shard.is_some() { return Err(HipError::new( 0, @@ -4750,42 +4915,93 @@ pub(crate) fn load_moe_ffn( .map(|slot| ep.as_ref().map(|e| e.src(slot)).unwrap_or(slot)) .filter(|&x| owns_orig(x)) .collect(); - let packed = if ep_shard.is_none() && packed_mq4_experts_supported(gpu) { - try_load_packed_mq4_experts(hfq, gpu, p, &expert_ids, mi, config.dim)? - } else { - None - }; - let (mut experts, packed_expert_owners) = if let Some((experts, owners)) = packed { + // ── Escha-W2 routed experts (Task 10) ──────────────────────────────── + // An Escha-W2 layer carries ONE trellis code tensor per projection for + // all experts, not the per-expert `experts.{x}.gate_up_proj.weight` + // tensors every other path fishes out by index, so it bypasses both the + // packed-MQ4 fast path and the generic per-expert loop below. + if escha_layer && (ep_shard.is_some() || ep.is_some()) { + return Err(HipError::new( + 0, + "qwen35: Escha-W2 routed experts do not support EP sharding or a REAP keep-map \ + (both re-map experts across the per-expert tensors escha does not have)", + )); + } + let escha_tables = if escha_layer { + let store = escha_weight_store(); + let (experts, tables, owners) = escha::load_escha_moe_experts( + hfq, + gpu, + p, + &expert_ids, + n_exp, + config.dim, + mi, + config.num_experts_per_tok, + store, + qwen35_tensor_name_candidates, + )?; if layer_idx == 0 { eprintln!( - " routed MQ4 expert packing: {} per-expert weight buffers -> 2 layer blobs", + " Escha-W2 routed experts: {} experts, store {store:?} ({}), {} per-expert \ + weight buffers -> 2 layer blobs", + experts.len(), + match store { + EschaWeightStore::Native => "trellis code kept verbatim, decoded in the GEMV", + _ => "decoded from the trellis at load", + }, 2 * experts.len() ); } - (experts, Some(owners)) + Some((experts, tables, owners)) } else { - let mut experts = Vec::with_capacity(expert_ids.len()); - for x in expert_ids { - let gate_up = load_weight_tensor( - hfq, - gpu, - &format!("{p}.mlp.experts.{x}.gate_up_proj.weight"), - 2 * mi, - config.dim, - qwen35_tensor_name_candidates, - )?; - let down = load_weight_tensor( - hfq, - gpu, - &format!("{p}.mlp.experts.{x}.down_proj.weight"), - config.dim, - mi, - qwen35_tensor_name_candidates, - )?; - experts.push(ExpertWeights { gate_up, down }); - } - (experts, None) + None }; + + let packed = if !escha_layer && ep_shard.is_none() && packed_mq4_experts_supported(gpu) { + try_load_packed_mq4_experts(hfq, gpu, p, &expert_ids, mi, config.dim)? + } else { + None + }; + let (mut experts, packed_expert_owners, escha_tables) = + if let Some((e, t, owners)) = escha_tables { + // Escha expert slots are views into `owners`, exactly like the packed + // MQ4 path's — so they ride the SAME `packed_expert_owners` free path + // (`free_moe_ffn` frees per-expert metadata only, then the two blobs). + // Publishing them here rather than in a bespoke field is what keeps + // teardown from double-freeing or leaking 32 GB. + (e, Some(owners), Some(t)) + } else if let Some((experts, owners)) = packed { + if layer_idx == 0 { + eprintln!( + " routed MQ4 expert packing: {} per-expert weight buffers -> 2 layer blobs", + 2 * experts.len() + ); + } + (experts, Some(owners), None) + } else { + let mut experts = Vec::with_capacity(expert_ids.len()); + for x in expert_ids { + let gate_up = load_weight_tensor( + hfq, + gpu, + &format!("{p}.mlp.experts.{x}.gate_up_proj.weight"), + 2 * mi, + config.dim, + qwen35_tensor_name_candidates, + )?; + let down = load_weight_tensor( + hfq, + gpu, + &format!("{p}.mlp.experts.{x}.down_proj.weight"), + config.dim, + mi, + qwen35_tensor_name_candidates, + )?; + experts.push(ExpertWeights { gate_up, down }); + } + (experts, None, None) + }; if e8_soa_experts() && gpu.arch_caps.is_rdna3_dgpu() && ep_shard.is_none() { let mut converted = 0usize; for ew in experts.iter_mut() { @@ -4902,5 +5118,44 @@ pub(crate) fn load_moe_ffn( paro_shared: None, global_expert_dtypes: None, ep_dummy_buffers, + escha: escha_tables, }) } + +/// Transpose an escha code blob's TILE GRID from the checkpoint's +/// `[ic/16][oc/16]` (kt-major) to `[oc/16][ic/16]` (nt-major). See the call +/// site. Moves whole tiles only, so it is bit-exact by construction. +fn escha_tiles_to_nt_major(data: &[u8], m: usize, k: usize, quant_type: u8) -> HipResult> { + let tk = if quant_type == 42 { 2usize } else { 3usize }; + if m % 16 != 0 || k % 16 != 0 { + return Err(hip_bridge::HipError::new( + 0, + &format!("escha code: m={m} k={k}; both must be multiples of 16"), + )); + } + let (ktiles, ntiles) = (k / 16, m / 16); + let tile_bytes = 16 * tk * 2; + let grid = ktiles * ntiles * tile_bytes; + if grid == 0 || data.len() % grid != 0 { + return Err(hip_bridge::HipError::new( + 0, + &format!( + "escha code: {} bytes is not a whole number of {grid}-byte grids \ + (ktiles={ktiles} ntiles={ntiles})", + data.len() + ), + )); + } + let mut out = vec![0u8; data.len()]; + for e in 0..(data.len() / grid) { + let base = e * grid; + for kt in 0..ktiles { + for nt in 0..ntiles { + let src = base + (kt * ntiles + nt) * tile_bytes; + let dst = base + (nt * ktiles + kt) * tile_bytes; + out[dst..dst + tile_bytes].copy_from_slice(&data[src..src + tile_bytes]); + } + } + } + Ok(out) +} diff --git a/crates/hipfire-arch-qwen35/src/qwen35/prefill.rs b/crates/hipfire-arch-qwen35/src/qwen35/prefill.rs index 418e23b449..248528b2ef 100644 --- a/crates/hipfire-arch-qwen35/src/qwen35/prefill.rs +++ b/crates/hipfire-arch-qwen35/src/qwen35/prefill.rs @@ -14,6 +14,7 @@ use super::config::Qwen35Config; use super::config::TreeVerifyCtx; use super::forward::checked_kv_end; use super::forward::forward_scratch; +use super::forward::forward_scratch_opts; use super::forward::forward_scratch_with_hidden; use super::forward::kv_cache_attention_dispatch; use super::forward::moe_ffn_has_mq3_experts_uniform; @@ -100,7 +101,6 @@ fn dispatch_batched_gemm_epilogue( let is_mq3 = matches!(w.gpu_dtype, DType::MQ3G256); let is_fp4 = matches!(w.gpu_dtype, DType::HFP4G32 | DType::MFP4G32); let is_q8 = matches!(w.gpu_dtype, DType::Q8_0); - let is_lowbit = matches!(w.gpu_dtype, DType::TQ2G128 | DType::BQ1G128); match epilogue { BatchEpilogue::Residual => { if is_6bit { @@ -128,11 +128,11 @@ fn dispatch_batched_gemm_epilogue( k, n, ); - } else if is_q8 || is_lowbit { + } else if is_unfused_plain_gemm_dtype(w.gpu_dtype) { let scratch = pbs.x_rot_batch.sub_offset(0, n * m); run_plain_gemm_key( gpu, - plain_gemm_key_for(w.gpu_dtype), + plain_gemm_key_for(w.gpu_dtype)?, &w.buf, w.gpu_dtype, input, @@ -221,10 +221,10 @@ fn dispatch_batched_gemm_epilogue( k, n, ); - } else if is_q8 || is_lowbit { + } else if is_unfused_plain_gemm_dtype(w.gpu_dtype) { return run_plain_gemm_key( gpu, - plain_gemm_key_for(w.gpu_dtype), + plain_gemm_key_for(w.gpu_dtype)?, &w.buf, w.gpu_dtype, input, @@ -349,6 +349,13 @@ const PREFILL_DEFAULT_BATCH_GFX1100: usize = 512; /// Exact `gfx1201` only — not gfx1200 or other gfx12 variants. const PREFILL_DEFAULT_BATCH_GFX1201: usize = 384; +/// gfx1151-measured default prefill chunk size (Strix Halo). +/// +/// Measured on the native escha 27B at an 8k prompt: 256 -> 52 tok/s, +/// **512 -> 73 tok/s**, 1024 -> 69 tok/s. gfx1151 had been falling through to +/// the generic 256 and leaving 40% of prefill on the table. +const PREFILL_DEFAULT_BATCH_GFX1151: usize = 512; + /// Architecture default for prefill chunk size when /// `HIPFIRE_PREFILL_MAX_BATCH` is unset or invalid. #[inline] @@ -357,6 +364,8 @@ fn prefill_max_batch_for_arch(arch: &str) -> usize { PREFILL_DEFAULT_BATCH_GFX1100 } else if arch == "gfx1201" { PREFILL_DEFAULT_BATCH_GFX1201 + } else if arch == "gfx1151" { + PREFILL_DEFAULT_BATCH_GFX1151 } else { PREFILL_MAX_BATCH } @@ -1188,7 +1197,16 @@ fn forward_prefill_batch_with_pbs_opts_inner( // One-shot: mark this forward AR-graph-eligible iff it's plain // single-token decode (consumed inside forward_scratch). gpu.graphs.ar_graph_eligible = plain_ar_graph_eligible; - forward_scratch( + // Only the LAST token's logits survive this loop — every + // earlier token's are overwritten by the next iteration and + // never read. Skipping the lm_head for them is not a + // shortcut, it deletes work that has no consumer: measured + // 2.32 ms of a 24.55 ms escha-35b prefill token (508 MB of + // weight traffic at vocab 248 320), 9.5 % of prefill, on + // every token but one. `scratch.tmp` (post-output-norm + // hidden) is still written, so the `per_token_hidden_out` + // copy below is unaffected. + forward_scratch_opts( gpu, weights, config, @@ -1197,6 +1215,7 @@ fn forward_prefill_batch_with_pbs_opts_inner( kv_cache, dn_state, scratch, + i + 1 == tokens.len(), )?; } if let Some(dst) = per_token_hidden_out { @@ -1343,14 +1362,99 @@ fn forward_prefill_batch_with_pbs_opts_inner( /// /// Q8 keeps the chunked kernel it already used, so this is behaviour-preserving /// for every existing model; the low-bit formats route to their tiled prefill -/// GEMMs. They share the Q8 call sites deliberately: none of the three has a -/// fused qkvza/gate_up/qkv kernel, so all three want the same unfused strategy. -fn plain_gemm_key_for(dt: DType) -> hipfire_dispatch::types::KernelKey { +/// GEMMs. They share the Q8 call sites deliberately: none of the four has a +/// fused qkvza/gate_up/qkv kernel, so all four want the same unfused strategy. +/// +/// # Why this is an EXHAUSTIVE match that errors, not a `_ =>` guess +/// +/// This used to end in `_ => GemmQ8_0BatchedChunked`. That default is a +/// silent-corruption hazard, and it became reachable the moment +/// [`is_batchable_la`] admitted `F16`: a Q8_0 GEMM walks 34-byte blocks of +/// 32 int8 plus an f16 scale, so pointing it at a plain f16 weight matrix +/// reads the wrong stride and produces finite, fluent, wrong output with +/// nothing to catch it — the exact failure class this port keeps finding. +/// An unrecognised dtype is now a hard error at the dispatch boundary, which +/// surfaces as a refused batched prefill (the caller falls back per-token) +/// rather than as bad tokens. +/// +/// `F16 => GemmF16WmmaMb8` is the escha arm: `linear_attn.in_proj_a` / +/// `in_proj_b` (`w_alpha` / `w_beta`) are stored at qt=1 = F16 and load +/// native. `gemm_f16_wmma_mb8` takes an F16 weight against an F32 activation +/// and was already the E8 shared-expert batched GEMM, so the kernel is +/// shipped and exercised on gfx11. +fn plain_gemm_key_for(dt: DType) -> HipResult { use hipfire_dispatch::types::KernelKey as K; match dt { - DType::TQ2G128 => K::GemmTQ2G128Prefill, - DType::BQ1G128 => K::GemmBQ1G128Prefill, - _ => K::GemmQ8_0BatchedChunked, + DType::TQ2G128 => Ok(K::GemmTQ2G128Prefill), + DType::BQ1G128 => Ok(K::GemmBQ1G128Prefill), + DType::F16 => Ok(K::GemmF16WmmaMb8), + DType::Q8_0 => Ok(K::GemmQ8_0BatchedChunked), + other => Err(HipError::new( + 0, + &format!( + "plain_gemm_key_for: no batched plain-GEMM kernel for {other:?}. \ + This dispatcher used to fall through to GemmQ8_0BatchedChunked for \ + any unknown dtype, which reads non-Q8_0 bytes at a Q8_0 stride — \ + add an explicit arm (and admit the dtype in is_batchable_la in the \ + SAME change) rather than restoring the default." + ), + )), + } +} + +/// Weight dtypes the batched-prefill matchers serve with the UNFUSED, +/// per-weight [`run_plain_gemm_key`] strategy — exactly the set +/// [`plain_gemm_key_for`] resolves. +/// +/// Keep the two in lockstep. A dtype admitted here with no arm there errors at +/// the dispatch boundary (safe: the layer refuses and prefill falls back +/// per-token). A dtype with an arm there but missing here falls through to a +/// FUSED matcher instead, which assumes one shared stride across all of a +/// layer's projections — that direction is silent corruption, so this predicate +/// is the one that must be widened first. +#[inline] +pub(crate) fn is_unfused_plain_gemm_dtype(dt: DType) -> bool { + matches!( + dt, + DType::Q8_0 | DType::TQ2G128 | DType::BQ1G128 | DType::F16 + ) +} + +/// Runtime precondition of the fused Q8_0 QKVZA / QKV / gate+up kernels: every +/// weight the ONE launch reads must actually be Q8_0. +/// +/// These matchers key their arm selection on a single representative weight +/// (`wqkv`, `wq`, `w_gate`) and used to guard the rest with a `debug_assert!`, +/// which is compiled out of a release build. That was unreachable only because +/// `is_batchable_la` refused F16 up front; escha-35b stores `w_alpha` / `w_beta` +/// as F16 beside a Q8_0 `wqkv`, so admitting F16 makes a Q8_0-strided read of +/// F16 bytes reachable in release — finite, fluent, wrong. A false here sends +/// the layer down the per-weight `run_plain_gemm_key` path instead, which reads +/// each weight at its own stride. +#[inline] +pub(crate) fn all_q8_0(dtypes: &[DType]) -> bool { + dtypes.iter().all(|dt| matches!(dt, DType::Q8_0)) +} + +/// The same precondition as [`all_q8_0`], for the arms whose representative +/// weight is an MQ-family container rather than Q8_0: every weight the ONE +/// fused launch reads must share the container the arm was selected from. +/// +/// `all_q8_0` was added because escha-35b stores `w_alpha` / `w_beta` as F16 +/// beside a Q8_0 `wqkv`. It fixed the Q8_0 arms and left the MQ arms — which +/// select on `layer.wqkv.gpu_dtype` alone in exactly the same way — still able +/// to read those F16 siblings at MQ stride. +/// +/// Measured before this guard existed, down-quantising ONLY `in_proj_qkv` on +/// escha-35b and scoring against the untouched build: KLD 12.63, PPL +/// 2,375,141 against a 7.68 baseline. Identical under MQ6G256, MQ6G256V2 and +/// MQ4G256V2, while `out_proj` — the one GDN projection not in the fused +/// launch — was unaffected at KLD 0.0076. Finite, fluent and wrong, which is +/// the signature this guard family exists to prevent. +pub(crate) fn all_same_dtype(dtypes: &[DType]) -> bool { + match dtypes.split_first() { + Some((head, rest)) => rest.iter().all(|dt| dt == head), + None => true, } } @@ -1395,10 +1499,58 @@ pub(crate) fn is_batchable_la(dt: DType, arch: &str) -> bool { // models unaffected because no production checkpoint sets // wqkv.gpu_dtype = ParoQ4G128 outside the shisa-PARO codepath. | DType::ParoQ4G128 | DType::F32 + // Escha-W2 trellis codes. Admissible NOT because any batched GEMM can + // read them — none can — but because the escha arms added to the LA + // and FFN chunk functions intercept BEFORE those matchers and run + // `EschaProj::forward` for the whole batch instead. Refusing here + // would drop the layer to the per-token path, which is what made a + // native 27B prefill at 11.7 tok/s against a decode of 11.0: prefill + // was doing decode's work once per token. + | DType::Escha2T16 | DType::Escha3T16 ); if always_ok { return true; } + + // F16 dense projections (Escha-W2 blocker 1), WMMA arches only. + // + // escha-35b stores `linear_attn.in_proj_a` / `in_proj_b` (w_alpha / + // w_beta) at qt=1 = F16 and the default `f16_lm_head_mode` is Native, so + // they load as DType::F16 beside a Q8_0 wqkv/wz/wo. Refusing F16 here + // refused 30 of escha's 40 layers BEFORE the MoE predicate was ever + // consulted, which is what pinned prefill to the per-token fallback. + // + // Admitting it is only safe because the SAME change (a) gave + // `plain_gemm_key_for` an explicit `F16 => GemmF16WmmaMb8` arm and turned + // its `_ =>` default from GemmQ8_0BatchedChunked into an error, and (b) + // converted the fused-QKVZA / QKV / gate+up `debug_assert!` stride guards + // into runtime `all_q8_0` predicates, so a layer that mixes Q8_0 and F16 + // takes the per-weight unfused path instead of one fused launch at the + // wrong stride. See the all-together rule in + // docs/plans/mq-lloyd-batched-prefill-followup.md. + // + // Arch-gated rather than `always_ok` on purpose: the only batched F16 GEMM + // is `gemm_f16_wmma_mb8`, which hard-errors without wave32 WMMA. Admitting + // a layer whose dispatch then errors is NOT a graceful fallback — the + // chunk propagates the error out of `forward_prefill_batch` — so a + // non-WMMA arch must be refused HERE, where refusal means "keep taking the + // per-token path you take today". Same arch set as the MQ3/FP4 WMMA arms. + let f16_with_wmma = matches!(dt, DType::F16) + && matches!( + arch, + "gfx1100" + | "gfx1101" + | "gfx1102" + | "gfx1103" + | "gfx1150" + | "gfx1151" + | "gfx1152" + | "gfx1200" + | "gfx1201" + ); + if f16_with_wmma { + return true; + } // MQ3 (uniform / HFQ3 family) is batchable on archs with a WMMA // family ported. As of this commit: // - gfx11 (gfx1100/1101/1102/1150/1151): wave32 WMMA via the @@ -1837,7 +1989,24 @@ pub fn qwen35_layer_batch_admissible( arch, ); if !moe_ffn_batched_admissible(&l.ffn, admit_mq6, arch) { - return Err(HipError::new(0, "DeltaNetMoe moe_ffn not batch-admissible")); + // Name the dtypes. "not batch-admissible" alone sends the + // reader to a 180-line predicate; the six dtypes it keys on + // identify the missing arm directly. + return Err(HipError::new( + 0, + &format!( + "DeltaNetMoe moe_ffn not batch-admissible: router={:?} \ + shared_gate={:?} shared=({:?},{:?},{:?}) routed=({:?},{:?}) escha={}", + l.ffn.router.gpu_dtype, + l.ffn.shared_expert_gate.gpu_dtype, + l.ffn.shared_expert.gate.gpu_dtype, + l.ffn.shared_expert.up.gpu_dtype, + l.ffn.shared_expert.down.gpu_dtype, + l.ffn.experts.first().map(|e| e.gate_up.gpu_dtype), + l.ffn.experts.first().map(|e| e.down.gpu_dtype), + l.ffn.escha.is_some(), + ), + )); } Ok(()) } @@ -1907,7 +2076,24 @@ pub fn qwen35_layer_batch_admissible( arch, ); if !moe_ffn_batched_admissible(&l.ffn, admit_mq6, arch) { - return Err(HipError::new(0, "FullAttnMoe moe_ffn not batch-admissible")); + // Name the dtypes. "not batch-admissible" alone sends the + // reader to a 180-line predicate; the six dtypes it keys on + // identify the missing arm directly. + return Err(HipError::new( + 0, + &format!( + "FullAttnMoe moe_ffn not batch-admissible: router={:?} \ + shared_gate={:?} shared=({:?},{:?},{:?}) routed=({:?},{:?}) escha={}", + l.ffn.router.gpu_dtype, + l.ffn.shared_expert_gate.gpu_dtype, + l.ffn.shared_expert.gate.gpu_dtype, + l.ffn.shared_expert.up.gpu_dtype, + l.ffn.shared_expert.down.gpu_dtype, + l.ffn.experts.first().map(|e| e.gate_up.gpu_dtype), + l.ffn.experts.first().map(|e| e.down.gpu_dtype), + l.ffn.escha.is_some(), + ), + )); } Ok(()) } @@ -2021,6 +2207,17 @@ struct MoePrefillDtypes { /// fails admission and silently drops to the per-token prefill fallback (the /// merged kernel never fires — observed as ~decode-speed prefill). routed_mixed_merged: bool, + /// This layer's routed experts are Escha-W2 coded: the H128 transform + /// tables are resident (`MoeFfnWeights::escha`) AND the indexed + /// (device-resident top-K) route is enabled. + /// + /// Both halves matter. The tables are what make the H128 pair callable at + /// all; the indexed route is what the batched prefill executor mirrors — + /// there is no batched analogue of the CPU-top-K host route, so under + /// `HIPFIRE_ESCHA_INDEXED=0` this stays false and escha prefills per-token + /// exactly as it did before, keeping that env var a genuine A/B of two + /// working routes rather than a half-disabled state. + escha: bool, } impl MoePrefillDtypes { @@ -2037,6 +2234,7 @@ impl MoePrefillDtypes { expert_gate_up_uniform: true, expert_down_uniform: true, routed_mixed_merged: false, + escha: false, } } @@ -2058,6 +2256,7 @@ impl MoePrefillDtypes { expert_gate_up_uniform: global.iter().all(|(g, _)| *g == first.0), expert_down_uniform: global.iter().all(|(_, d)| *d == first.1), routed_mixed_merged: ffn.expert_dtype_tags.is_some(), + escha: ffn.escha.is_some() && super::escha::escha_indexed_route_enabled(), }); } let first = ffn.experts.first()?; @@ -2078,6 +2277,7 @@ impl MoePrefillDtypes { .iter() .all(|e| e.down.gpu_dtype == first.down.gpu_dtype), routed_mixed_merged: ffn.expert_dtype_tags.is_some(), + escha: ffn.escha.is_some() && super::escha::escha_indexed_route_enabled(), }) } } @@ -2197,13 +2397,25 @@ fn moe_ffn_batched_admissible_for_dtypes( admit_e8: bool, admit_codebook: bool, ) -> bool { + // F16 (Escha-W2 blocker 2). `mlp.gate.weight` and + // `mlp.shared_expert_gate.weight` are F16 on escha-35b, and these are the + // FIRST two checks in this predicate — so the routed arms below were never + // even reached for an escha layer. Both are small ([n_exp, dim] and + // [1, dim]) and dispatch through `GemmF16WmmaMb8` against the UN-rotated + // `x_norm_batch`, matching the Q8_0 convention (F16 weights are dense and + // unrotated, so they are quantised against the unrotated activation). + // + // Every arm below is unchanged, so a model whose router / scalar gate is + // already MQ4 / MQ4V2 / Q8_0 / F32 keeps exactly the admission decision it + // has today. The only models this widens are ones REFUSED today, which by + // definition prefill per-token and cannot regress. let router_ok = matches!( dtypes.router, - DType::MQ4G256 | DType::MQ4G256V2 | DType::Q8_0 | DType::F32 + DType::MQ4G256 | DType::MQ4G256V2 | DType::Q8_0 | DType::F32 | DType::F16 ); let shared_gate_ok = matches!( dtypes.shared_expert_scalar_gate, - DType::MQ4G256 | DType::MQ4G256V2 | DType::Q8_0 | DType::F32 + DType::MQ4G256 | DType::MQ4G256V2 | DType::Q8_0 | DType::F32 | DType::F16 ); // Graded (mixed-dtype) routed experts are served by the merged grouped-WMMA // prefill kernel, so the per-expert *uniform* requirement is waived for the @@ -2236,6 +2448,51 @@ fn moe_ffn_batched_admissible_for_dtypes( return shared_gu_ok && shared_dn_ok; } + // ── Escha-W2 routed experts (blocker 3) ────────────────────────────── + // + // The routed experts are the Q8_0 the trellis decoded into, on BOTH + // projections, and the layer carries the H128 transform tables. This is + // the prefill twin of `MoeResolution::routed_indexable_escha_q8`, and it + // admits the layer to ONE executor only: `escha_routed_prefill_indexed`, + // reached through the escha branch at the top of `run_moe_prefill`. It + // must never reach the generic Path 1 / Path 2 routed bodies, which know + // nothing about the Hadamard domain escha weights live in and would emit + // finite, fluent, ~1e-1-wrong output. + // + // The shared expert is required Q8_0 on all three projections because that + // is what the escha checkpoint ships and what the Q8 arm of the batched + // body above serves; a different shared dtype would be a different model + // and should get its own arm rather than silently borrow this one. + // + // The ROUTED side admits exactly the two containers the escha indexed + // executor has a GEMV for, and the pair must be uniform-in-kind: + // + // * `Escha2T16` / `Escha3T16` — Phase 2 production. The routed experts + // are the trellis CODE and `escha_gemv_native_*` decodes it inside the + // GEMV. Either order on either projection (the shipped file is K=2 + // gate_up / K=3 down; the reverse allocation is equally valid), which + // is exactly what `MoeResolution::routed_indexable_escha_native` + // admits on the decode side — this is its prefill twin and the two + // must agree or a layer batches in prefill and does not in decode. + // * `Q8_0` on both — Phase 1, the A/B arm. + // + // The container is load-bearing for the same reason as every other uniform + // arm: each escha GEMV hard-codes one bit geometry (a 34 B/32-element Q8_0 + // block, or a 16x16 trellis tile) and handing it the other reads different + // weights out of the same bytes — silent corruption, not a fault. + let escha_routed_ok = (dtypes.expert_gate_up == DType::Q8_0 + && dtypes.expert_down == DType::Q8_0) + || (matches!(dtypes.expert_gate_up, DType::Escha2T16 | DType::Escha3T16) + && matches!(dtypes.expert_down, DType::Escha2T16 | DType::Escha3T16)); + if dtypes.escha + && escha_routed_ok + && dtypes.shared_expert_gate == DType::Q8_0 + && dtypes.shared_expert_up == DType::Q8_0 + && dtypes.shared_expert_down == DType::Q8_0 + { + return true; + } + // mfp4-E8 routed experts with Q8 shared expert (original arm): // gfx1151-native A3B checkpoint. Shared expert is Q8 (gate/up/down); // router/scalar-gate are Q8 (validated by router_ok/shared_gate_ok above). @@ -2483,6 +2740,44 @@ pub fn prefill_batch_pbs_eligible( .as_deref() == Some("1") { + // Per-layer attention-projection dtypes. The refusal lines below name + // only the FIRST weight that failed, which is enough to know why a + // layer refused but not enough to know what the layer IS — and an + // ADMITTED layer printed nothing at all. Both questions come up every + // time a mixed-dtype checkpoint (escha: Q8_0 wqkv/wz/wo beside F16 + // w_alpha/w_beta) picks a fused-vs-unfused arm. + for (i, lw) in weights.layers.iter().enumerate() { + let dts = match lw { + LayerWeights::DeltaNetMoe(l) => Some(( + "DeltaNetMoe", + [ + l.wqkv.gpu_dtype, + l.wz.gpu_dtype, + l.w_beta.gpu_dtype, + l.w_alpha.gpu_dtype, + l.wo.gpu_dtype, + ], + )), + LayerWeights::DeltaNet(l) => Some(( + "DeltaNet", + [ + l.wqkv.gpu_dtype, + l.wz.gpu_dtype, + l.w_beta.gpu_dtype, + l.w_alpha.gpu_dtype, + l.wo.gpu_dtype, + ], + )), + _ => None, + }; + if let Some((kind, d)) = dts { + eprintln!( + "[hipfire::batch_eligible] L{i} {kind} wqkv={:?} wz={:?} \ + w_beta={:?} w_alpha={:?} wo={:?}", + d[0], d[1], d[2], d[3], d[4] + ); + } + } eprintln!( "[hipfire::batch_eligible] result={result} \ arch={arch} n={n} n>={MIN_BATCH}={} \ @@ -2492,6 +2787,40 @@ pub fn prefill_batch_pbs_eligible( all_layers_ok={all_layers_ok}", n >= MIN_BATCH, ); + // `all_layers_ok=false` on its own says a model prefills at decode + // speed but not WHY, and the reason is a string inside an `is_ok()` + // that nothing prints. Report each DISTINCT refusal once, with how + // many layers it covers, so a per-token prefill is diagnosable from + // one run instead of a source read. (Escha-W2 was diagnosed this way.) + if !all_layers_ok { + let mut seen: Vec<(String, Vec)> = Vec::new(); + for (i, lw) in weights.layers.iter().enumerate() { + let reason = if matches!( + lw, + LayerWeights::DeltaNetMoe(_) | LayerWeights::FullAttnMoe(_) + ) && !moe_router_logits_present + { + Some("MoE layer with no batched router-logits buffer".to_string()) + } else { + qwen35_layer_batch_admissible(lw, config, arch) + .err() + .map(|e| e.message.clone()) + }; + if let Some(reason) = reason { + match seen.iter_mut().find(|(r, _)| *r == reason) { + Some((_, layers)) => layers.push(i), + None => seen.push((reason, vec![i])), + } + } + } + for (reason, layers) in &seen { + eprintln!( + "[hipfire::batch_eligible] REFUSED {} layer(s) (first is {}): {reason}", + layers.len(), + layers[0], + ); + } + } } result } @@ -3010,9 +3339,17 @@ pub(crate) fn prefill_moe_ffn_body_batched( hipfire_dispatch::types::KernelKey::GemmF32Batched, &pbs.x_norm_batch, ), + // F16 (escha `mlp.gate.weight`). Dense and unrotated, so it + // reads `x_norm_batch` like Q8_0/F32 — feeding it the + // FWHT-rotated `x_rot_batch` would be the classic silent + // "rotated activation into an unrotated weight" failure. + DType::F16 => ( + hipfire_dispatch::types::KernelKey::GemmF16WmmaMb8, + &pbs.x_norm_batch, + ), other => panic!( "prefill_moe_ffn_body_batched: unexpected router dtype {other:?} \ - — moe_ffn_batched_admitted admits MQ4G256, Q8_0, F32" + — moe_ffn_batched_admissible admits MQ4G256, MQ4G256V2, Q8_0, F32, F16" ), }; let w = WeightRef { @@ -3047,9 +3384,12 @@ pub(crate) fn prefill_moe_ffn_body_batched( DType::MQ4G256 => (KernelKey::GemmHfq4G256, &pbs.x_rot_batch), DType::MQ4G256V2 => (KernelKey::GemmMq4G256V2, &pbs.x_rot_batch), DType::F32 => (KernelKey::GemmF32Batched, &pbs.x_norm_batch), + // F16 (escha `mlp.shared_expert_gate.weight`), un-rotated — see + // the router arm above. + DType::F16 => (KernelKey::GemmF16WmmaMb8, &pbs.x_norm_batch), other => panic!( "prefill_moe_ffn_body_batched: unexpected shared_expert_gate dtype {other:?} \ - — moe_ffn_batched_admissible admits MQ4G256, Q8_0, F32" + — moe_ffn_batched_admissible admits MQ4G256, MQ4G256V2, Q8_0, F32, F16" ), }; run_plain_gemm_key( @@ -3260,6 +3600,31 @@ pub(crate) fn prefill_moe_ffn_body_batched( // so prefill activations match the CPU-reference softmax math // exactly. router_logits is allocated 1D as [n × n_exp]; alias it // into a 2D view so gpu.softmax_f32 takes rows = n. + // Escha-only router-logits f16 round-trip, the batched twin of the call in + // `run_moe_decode`. EschaLabs' runtime computes router logits as + // f16(x @ gate_w.T) and only widens to F32 to select top-k; hipfire keeps + // them F32 end to end, and escha's recovery fine-tune was trained against + // the rounding runtime. + // + // This is NOT optional polish for the batched path. Decode applies the + // rounding unconditionally, so omitting it here would make batched prefill + // select from UNROUNDED logits while decode selects from rounded ones — + // a systematic prefill-vs-decode route divergence on every token, on top + // of the residual divergence that genuinely straddles an f16 boundary + // (measured 2.96-3.13% per expert slot; design doc §10.5(b)). + // Applied to the shared `router_logits` buffer BEFORE the softmax, exactly + // where decode applies it, so both routes see identical inputs to top-k. + // + // Keyed on the layer's own transform tables (a load-time model-state + // property), never an env var, so every non-escha arch-6 model skips the + // launch entirely and keeps its selection bit-for-bit. + if ffn.escha.is_some() { + // `router_logits` is the [max_batch x n_exp] scratch; only the first + // n rows are live. Round exactly those — the kernel is numel-driven, + // so handing it the whole buffer would launch over stale tail rows. + let live = router_logits.sub_offset(0, n * n_exp); + gpu.router_logits_round_f16_rne(&live)?; + } let router_logits_2d = GpuTensor { buf: unsafe { router_logits.buf.alias() }, shape: vec![n, n_exp], @@ -3275,6 +3640,13 @@ pub(crate) fn prefill_moe_ffn_body_batched( n, )?; + // Selection capture, the batched twin of the call in `run_moe_decode` + // (off unless HIPFIRE_ESCHA_ROUTE_TRACE is set). Same point in the + // pipeline: after top-k, before any routed body. + if hipfire_dispatch::pipeline::route_trace::enabled() { + hipfire_dispatch::pipeline::route_trace::record(gpu, topk_indices, n, k_top); + } + // ── 4. Shared-expert SwiGLU + FWHT, batched over N tokens ── // // fused_silu_mul_rotate_mq_batched expects [batch × k] gate/up with @@ -3516,6 +3888,11 @@ pub(crate) fn prefill_moe_ffn_body_batched( has_paro_shared: ffn.paro_shared.is_some(), per_expert_gate_up, per_expert_down, + // Escha-W2. Same marker as the decode snapshot, and it must stay the + // same: the two routes resolving one layer differently is exactly the + // failure this field exists to prevent. It now also gates the escha + // branch at the top of `run_moe_prefill` (via `escha` below). + routed_escha_transforms: ffn.escha.is_some() && super::escha::escha_indexed_route_enabled(), }; let paro_gate_up = @@ -3580,6 +3957,27 @@ pub(crate) fn prefill_moe_ffn_body_batched( paro_down, down_awq_scale, routed_out, + // Escha-W2 batched routed executor. `Some` publishes the layer's four + // `[E, ·]` transform tables to `run_moe_prefill`'s escha branch; the + // `[k]`-sized decode scratch inside `refs()` is unused there (batched + // prefill uses the model-global `[n × k]` scratch instead). + // + // Gated identically to `routed_escha_transforms` above and to the + // `escha` field of `MoePrefillDtypes` that admitted this layer, so + // "admitted" and "has the tables" cannot come apart. + escha: if ffn.escha.is_some() && super::escha::escha_indexed_route_enabled() { + ffn.escha.as_ref().map(|e| e.refs()) + } else { + None + }, + // UNGATED, deliberately — the one marker in this struct that is not + // ANDed with `escha_indexed_route_enabled()`. Keyed on the layer's own + // transform tables, exactly as the router f16 rounding above is. With + // the indexed route off, `escha` is `None` while this stays true, and + // `check_moe_prefill_supported` refuses the layer rather than letting + // it fall into a transform-free Path 1 / Path 2. + layer_is_escha: ffn.escha.is_some(), + hidden: dim, }; hipfire_runtime::llama::moe_family() .run_prefill(ctx, gpu, &moe_prefill_params) @@ -3625,6 +4023,46 @@ pub(crate) struct PrefillBandCtx<'a> { /// divergence by diffing `.batched` vs `.pertoken` per layer. Requires /// `HIPFIRE_GRAPH=0` (does a synchronous D2H readback, which is illegal under /// graph capture). +/// Layer the `HIPFIRE_DUMP_HIDDEN` stage dumps (`q_b`/`k_b`/`v_b`/`alpha_b`/ +/// `beta_b`) fire on. `HIPFIRE_DUMP_HIDDEN_LAYER`, default 0. +pub(crate) fn dump_diag_layer() -> usize { + static LAYER: std::sync::OnceLock = std::sync::OnceLock::new(); + *LAYER.get_or_init(|| { + hipfire_config::developer_var("HIPFIRE_DUMP_HIDDEN_LAYER") + .ok() + .and_then(|v| v.parse().ok()) + .unwrap_or(0) + }) +} + +/// `HIPFIRE_DUMP_HIDDEN`'s file prefix, resolved ONCE. +/// +/// The call sites below are unconditional — every qwen35 model pays them, once +/// per layer per chunk, not just the escha localisation runs they were written +/// for. Read once behind a `OnceLock` (the pattern +/// `escha::escha_indexed_route_enabled` and `route_trace` already use) so the +/// steady-state cost is one relaxed atomic load rather than an environment +/// lookup and a `String` allocation. Mid-run env mutation is not honoured, in +/// common with every other developer var in this crate. +fn dump_hidden_prefix() -> Option<&'static str> { + static PREFIX: std::sync::OnceLock> = std::sync::OnceLock::new(); + PREFIX + .get_or_init(|| hipfire_config::developer_var("HIPFIRE_DUMP_HIDDEN").ok()) + .as_deref() +} + +/// `HIPFIRE_DUMP_HIDDEN_POS` (default 0), resolved once. See +/// [`dump_hidden_prefix`]. +fn dump_hidden_pos() -> usize { + static POS: std::sync::OnceLock = std::sync::OnceLock::new(); + *POS.get_or_init(|| { + hipfire_config::developer_var("HIPFIRE_DUMP_HIDDEN_POS") + .ok() + .and_then(|v| v.parse().ok()) + .unwrap_or(0) + }) +} + pub(crate) fn dump_hidden_localize( gpu: &Gpu, x: &GpuTensor, @@ -3634,14 +4072,12 @@ pub(crate) fn dump_hidden_localize( layer_idx: usize, tag: &str, ) { - let prefix = match hipfire_config::developer_var("HIPFIRE_DUMP_HIDDEN") { - Ok(p) => p, - Err(_) => return, + // One relaxed atomic load in the (overwhelmingly common) off case — this + // is called unconditionally, per layer per chunk, for EVERY qwen35 model. + let Some(prefix) = dump_hidden_prefix() else { + return; }; - let target: usize = hipfire_config::developer_var("HIPFIRE_DUMP_HIDDEN_POS") - .ok() - .and_then(|v| v.parse().ok()) - .unwrap_or(0); + let target: usize = dump_hidden_pos(); if target < abs_pos_of_row0 { return; } @@ -4105,11 +4541,6 @@ pub(crate) fn batch_chunk_delta_net_attn( let is_mq3_lloyd = matches!(layer.wqkv.gpu_dtype, DType::MQ3G256Lloyd); let is_fp4 = matches!(layer.wqkv.gpu_dtype, DType::HFP4G32 | DType::MFP4G32); let is_q8 = matches!(layer.wqkv.gpu_dtype, DType::Q8_0); - // TQ2G128/BQ1G128 have no fused qkvza/gate_up/qkv kernel, so they take - // the same UNFUSED plain-GEMM strategy as Q8 rather than falling through - // to the HFQ4 arm, which would read these packed blocks at the wrong - // stride and produce fluent-but-wrong tokens. - let is_lowbit = matches!(layer.wqkv.gpu_dtype, DType::TQ2G128 | DType::BQ1G128); // Batched rmsnorm (+ FWHT for MQ) for the LA preamble. // x_batch / x_rot_batch are [N × dim] contiguous. For HFQ @@ -4139,7 +4570,65 @@ pub(crate) fn batch_chunk_delta_net_attn( } // Batched 4-way LA projection (wqkv + wz + w_beta + w_alpha). - if is_6bit { + // + // ESCHA FIRST: a trellis layer bypasses every fused arm below. Each + // projection rotates the same normed input with its OWN rin, so the fused + // kernels have nothing to share and cannot read a trellis code. `is_mq` is + // false for Escha2T16/3T16, so the rmsnorm above already left the normed + // (unrotated) activation in `x_rot_batch` — exactly what escha wants, + // since it applies its own H128 per projection. + // + // `in_proj_a`/`in_proj_b` stay on the ordinary batched GEMM: escha's + // `ignore` list leaves them uncoded. + if let Some(e) = layer.escha.as_ref() { + // One group holding every slot: `expert_offsets = [0, n]` and the + // identity permutation. Built here rather than per projection so the + // 2-int upload happens once per layer, not six times. + let off_bytes: Vec = [0i32, n as i32] + .iter() + .flat_map(|v| v.to_le_bytes()) + .collect(); + let offsets = gpu.upload_raw(&off_bytes, &[2])?; + let grouped = Some((&offsets, &e.iota)); + + e.qkv.forward( + gpu, + &layer.wqkv, + &e.ids, + &pbs.x_rot_batch, + &pbs.escha_xh_batch, + &pbs.dn_qkv_batch, + &pbs.dn_qkv_batch, + n, + grouped, + )?; + e.z.forward( + gpu, + &layer.wz, + &e.ids, + &pbs.x_rot_batch, + &pbs.escha_xh_batch, + &pbs.dn_z_batch, + &pbs.dn_z_batch, + n, + grouped, + )?; + batched_gemm_single_weight(gpu, &layer.w_beta, &pbs.x_rot_batch, &pbs.dn_beta_batch, n)?; + batched_gemm_single_weight( + gpu, + &layer.w_alpha, + &pbs.x_rot_batch, + &pbs.dn_alpha_batch, + n, + )?; + } else if is_6bit + && all_same_dtype(&[ + layer.wqkv.gpu_dtype, + layer.wz.gpu_dtype, + layer.w_beta.gpu_dtype, + layer.w_alpha.gpu_dtype, + ]) + { run_fused_qkvza_key( gpu, hipfire_dispatch::types::KernelKey::FusedQkvzaHfq6G256, @@ -4159,17 +4648,22 @@ pub(crate) fn batch_chunk_delta_net_attn( layer.wqkv.k, n, )?; - } else if is_q8 && q8_wmma_arch { - // `is_q8` only inspects `wqkv` (the routing anchor). The fused - // kernel assumes ALL four weights share the Q8_0 stride; a - // mixed-dtype layer would silently re-introduce the Tier-1 - // kernel-vs-stride corruption mode. - debug_assert!( - matches!(layer.wz.gpu_dtype, DType::Q8_0) - && matches!(layer.w_beta.gpu_dtype, DType::Q8_0) - && matches!(layer.w_alpha.gpu_dtype, DType::Q8_0), - "LA qkvza Q8 WMMA dispatch requires all of wqkv/wz/w_beta/w_alpha to be Q8_0", - ); + } else if is_q8 + && q8_wmma_arch + && all_q8_0(&[ + layer.wqkv.gpu_dtype, + layer.wz.gpu_dtype, + layer.w_beta.gpu_dtype, + layer.w_alpha.gpu_dtype, + ]) + { + // `is_q8` only inspects `wqkv` (the routing anchor). The fused kernel + // reads all FOUR weights in one launch at the Q8_0 stride, so the + // other three are checked HERE, at runtime — this used to be a + // `debug_assert!` (compiled out of release) and became reachable in + // release the moment `is_batchable_la` admitted F16. A mixed layer now + // falls through to the per-weight unfused arm below, which reads each + // weight at its own stride. run_fused_qkvza_key( gpu, hipfire_dispatch::types::KernelKey::FusedQkvzaQ8_0, @@ -4189,14 +4683,14 @@ pub(crate) fn batch_chunk_delta_net_attn( layer.wqkv.k, n, )?; - } else if is_q8 || is_lowbit { + } else if is_unfused_plain_gemm_dtype(layer.wqkv.gpu_dtype) { // #397 Ship 5.2 slice1: four plain Q8 batched GEMMs // (wqkv/wz/w_beta/w_alpha) → GemmFamily::run_key with the // GemmQ8_0BatchedChunked dispatcher-entry key → identical // gpu.gemm_q8_0_batched_chunked method, byte-for-byte. run_plain_gemm_key( gpu, - plain_gemm_key_for(layer.wqkv.gpu_dtype), + plain_gemm_key_for(layer.wqkv.gpu_dtype)?, &layer.wqkv.buf, layer.wqkv.gpu_dtype, &pbs.x_rot_batch, @@ -4207,7 +4701,7 @@ pub(crate) fn batch_chunk_delta_net_attn( )?; run_plain_gemm_key( gpu, - plain_gemm_key_for(layer.wz.gpu_dtype), + plain_gemm_key_for(layer.wz.gpu_dtype)?, &layer.wz.buf, layer.wz.gpu_dtype, &pbs.x_rot_batch, @@ -4218,7 +4712,7 @@ pub(crate) fn batch_chunk_delta_net_attn( )?; run_plain_gemm_key( gpu, - plain_gemm_key_for(layer.w_beta.gpu_dtype), + plain_gemm_key_for(layer.w_beta.gpu_dtype)?, &layer.w_beta.buf, layer.w_beta.gpu_dtype, &pbs.x_rot_batch, @@ -4229,7 +4723,7 @@ pub(crate) fn batch_chunk_delta_net_attn( )?; run_plain_gemm_key( gpu, - plain_gemm_key_for(layer.w_alpha.gpu_dtype), + plain_gemm_key_for(layer.w_alpha.gpu_dtype)?, &layer.w_alpha.buf, layer.w_alpha.gpu_dtype, &pbs.x_rot_batch, @@ -4238,7 +4732,14 @@ pub(crate) fn batch_chunk_delta_net_attn( layer.w_alpha.k, n, )?; - } else if is_mq3_lloyd { + } else if is_mq3_lloyd + && all_same_dtype(&[ + layer.wqkv.gpu_dtype, + layer.wz.gpu_dtype, + layer.w_beta.gpu_dtype, + layer.w_alpha.gpu_dtype, + ]) + { // 112 B/group Lloyd-MQ3 stride; X is already FWHT-rotated. run_fused_qkvza_key( gpu, @@ -4259,7 +4760,14 @@ pub(crate) fn batch_chunk_delta_net_attn( layer.wqkv.k, n, )?; - } else if is_mq3 { + } else if is_mq3 + && all_same_dtype(&[ + layer.wqkv.gpu_dtype, + layer.wz.gpu_dtype, + layer.w_beta.gpu_dtype, + layer.w_alpha.gpu_dtype, + ]) + { // 104 B/group HFQ3-stride; X is already FWHT-rotated by // fused_rmsnorm_rotate_mq_batched above. The FusedQkvzaHfq3G256 // run-arm replicates the call-site WMMA-vs-base arch split @@ -4284,7 +4792,14 @@ pub(crate) fn batch_chunk_delta_net_attn( layer.wqkv.k, n, )?; - } else if is_fp4 { + } else if is_fp4 + && all_same_dtype(&[ + layer.wqkv.gpu_dtype, + layer.wz.gpu_dtype, + layer.w_beta.gpu_dtype, + layer.w_alpha.gpu_dtype, + ]) + { // HFP4G32: 17-B blocks (vs HFQ4's 136-B groups), per-row 16-B header. // MFP4G32: same storage as HFP4 + offline-FWHT weights; X is already // rotated above when is_mq, so this branch handles both unrotated @@ -4330,6 +4845,16 @@ pub(crate) fn batch_chunk_delta_net_attn( )?; } + // Escha dense biases, applied at the ONE point the whole if/else chain + // above converges — five branches fill `dn_qkv_batch`/`dn_z_batch`, and a + // bias added in some but not others is a silent wrong answer. Must land + // before sigmoid/conv1d consume these. `in_proj_a`/`in_proj_b` (beta/ + // alpha) have no bias: escha's `ignore` list keeps them plain weights. + if let Some(b) = layer.biases.as_ref() { + gpu.bias_add_f32(&pbs.dn_qkv_batch, &b.qkv, n, b.qkv.numel())?; + gpu.bias_add_f32(&pbs.dn_z_batch, &b.z, n, b.z.numel())?; + } + // Fused sigmoid(beta) + alpha_gate(alpha) — [N × n_v_heads] each. gpu.fused_sigmoid_alpha_gate_f32_batched( &pbs.dn_beta_batch, @@ -4688,16 +5213,52 @@ pub(crate) fn batch_chunk_delta_net_attn( } else { &pbs.dn_normed_batch }; - dispatch_batched_gemm_epilogue( - gpu, - pbs, - &layer.wo, - wo_input, - &epilogue, - n, - q8_wmma_arch, - arch_has_wmma, - )?; + if let Some(e) = layer.escha.as_ref() { + // One group holding every slot: `expert_offsets = [0, n]` and the + // identity permutation. Built here rather than per projection so the + // 2-int upload happens once per layer, not six times. + let off_bytes: Vec = [0i32, n as i32] + .iter() + .flat_map(|v| v.to_le_bytes()) + .collect(); + let offsets = gpu.upload_raw(&off_bytes, &[2])?; + let grouped = Some((&offsets, &e.iota)); + + // Trellis wo: project into the ffn scratch, then accumulate into the + // residual. The fused epilogue folds those together but cannot read a + // trellis code; the split add is exact, both terms being plain f32. + // Input is the UNROTATED gated-norm output: escha applies its own + // H128, so `dn_normed_batch` and not the MQ-rotated variant. + e.o.forward( + gpu, + &layer.wo, + &e.ids, + &pbs.dn_normed_batch, + &pbs.escha_xh_batch, + &pbs.escha_y_batch, + &pbs.escha_y_batch, + n, + grouped, + )?; + gpu.add_inplace_f32(&pbs.x_batch, &pbs.escha_y_batch)?; + } else { + dispatch_batched_gemm_epilogue( + gpu, + pbs, + &layer.wo, + wo_input, + &epilogue, + n, + q8_wmma_arch, + arch_has_wmma, + )?; + } + + // out_proj's bias lands on the residual stream. After the residual add is + // the same value as before it — both additive. + if let Some(b) = layer.biases.as_ref() { + gpu.bias_add_f32(&pbs.x_batch, &b.o, n, b.o.numel())?; + } Ok(()) } @@ -4715,6 +5276,76 @@ pub(crate) fn batch_chunk_delta_net_ffn( arch_has_wmma: bool, epilogue: BatchEpilogue<'_>, ) -> HipResult<()> { + // ── ESCHA TRELLIS FFN ─────────────────────────────────────────────── + // Whole FFN in one branch: plain rmsnorm (escha applies its own H128, so a + // pre-rotated input would be rotated twice), gate and up as separate + // trellis GEMVs, SwiGLU, then down accumulated into the residual. The + // fused gate_up kernel cannot serve this — the two projections have + // different rin and it could not read a trellis code regardless. + if let Some(e) = layer.escha.as_ref() { + // One group holding every slot: `expert_offsets = [0, n]` and the + // identity permutation. Built here rather than per projection so the + // 2-int upload happens once per layer, not six times. + let off_bytes: Vec = [0i32, n as i32] + .iter() + .flat_map(|v| v.to_le_bytes()) + .collect(); + let offsets = gpu.upload_raw(&off_bytes, &[2])?; + let grouped = Some((&offsets, &e.iota)); + + gpu.rmsnorm_batched( + &pbs.x_batch, + &layer.ffn_norm, + &pbs.x_norm_batch, + n, + layer.w_gate.k, + config.norm_eps, + )?; + e.gate.forward( + gpu, + &layer.w_gate, + &e.ids, + &pbs.x_norm_batch, + &pbs.escha_xh_batch, + &pbs.gate_ffn_batch, + &pbs.gate_ffn_batch, + n, + grouped, + )?; + e.up.forward( + gpu, + &layer.w_up, + &e.ids, + &pbs.x_norm_batch, + &pbs.escha_xh_batch, + &pbs.up_batch, + &pbs.up_batch, + n, + grouped, + )?; + if let Some(b) = layer.biases.as_ref() { + gpu.bias_add_f32(&pbs.gate_ffn_batch, &b.gate, n, b.gate.numel())?; + gpu.bias_add_f32(&pbs.up_batch, &b.up, n, b.up.numel())?; + } + gpu.silu_mul_f32(&pbs.gate_ffn_batch, &pbs.up_batch, &pbs.ffn_hidden_batch)?; + e.down.forward( + gpu, + &layer.w_down, + &e.ids, + &pbs.ffn_hidden_batch, + &pbs.escha_xh_batch, + &pbs.escha_y_batch, + &pbs.escha_y_batch, + n, + grouped, + )?; + gpu.add_inplace_f32(&pbs.x_batch, &pbs.escha_y_batch)?; + if let Some(b) = layer.biases.as_ref() { + gpu.bias_add_f32(&pbs.x_batch, &b.down, n, b.down.numel())?; + } + return Ok(()); + } + // FFN: rmsnorm (+ rotate for MQ). let ffn_is_mq = matches!( layer.w_gate.gpu_dtype, @@ -4735,11 +5366,6 @@ pub(crate) fn batch_chunk_delta_net_ffn( let ffn_is_mq3_lloyd = matches!(layer.w_gate.gpu_dtype, DType::MQ3G256Lloyd); let ffn_is_fp4 = matches!(layer.w_gate.gpu_dtype, DType::HFP4G32 | DType::MFP4G32); let ffn_is_q8 = matches!(layer.w_gate.gpu_dtype, DType::Q8_0); - // TQ2G128/BQ1G128 have no fused qkvza/gate_up/qkv kernel, so they take - // the same UNFUSED plain-GEMM strategy as Q8 rather than falling through - // to the HFQ4 arm, which would read these packed blocks at the wrong - // stride and produce fluent-but-wrong tokens. - let ffn_is_lowbit = matches!(layer.w_gate.gpu_dtype, DType::TQ2G128 | DType::BQ1G128); if ffn_is_mq { // AWQ-aware: next linear is w_gate (gate/up share input → same AWQ scale). fused_rmsnorm_rotate_mq_batched_for( @@ -4784,11 +5410,11 @@ pub(crate) fn batch_chunk_delta_net_ffn( layer.w_gate.k, n, )?; - } else if ffn_is_q8 && q8_wmma_arch { - debug_assert!( - matches!(layer.w_up.gpu_dtype, DType::Q8_0), - "LA FFN Q8 WMMA dispatch requires both w_gate and w_up to be Q8_0", - ); + } else if ffn_is_q8 && q8_wmma_arch && all_q8_0(&[layer.w_gate.gpu_dtype, layer.w_up.gpu_dtype]) + { + // `ffn_is_q8` inspects w_gate only; the fused kernel reads BOTH + // weights in one launch at the Q8_0 stride. Runtime, not + // `debug_assert!` — see `all_q8_0`. run_fused_gate_up_key( gpu, hipfire_dispatch::types::KernelKey::FusedGateUpQ8_0, @@ -4802,10 +5428,10 @@ pub(crate) fn batch_chunk_delta_net_ffn( layer.w_gate.k, n, )?; - } else if ffn_is_q8 || ffn_is_lowbit { + } else if is_unfused_plain_gemm_dtype(layer.w_gate.gpu_dtype) { run_plain_gemm_key( gpu, - plain_gemm_key_for(layer.w_gate.gpu_dtype), + plain_gemm_key_for(layer.w_gate.gpu_dtype)?, &layer.w_gate.buf, layer.w_gate.gpu_dtype, &pbs.x_rot_batch, @@ -4816,7 +5442,7 @@ pub(crate) fn batch_chunk_delta_net_ffn( )?; run_plain_gemm_key( gpu, - plain_gemm_key_for(layer.w_up.gpu_dtype), + plain_gemm_key_for(layer.w_up.gpu_dtype)?, &layer.w_up.buf, layer.w_up.gpu_dtype, &pbs.x_rot_batch, @@ -4883,6 +5509,14 @@ pub(crate) fn batch_chunk_delta_net_ffn( )?; } + // Escha dense biases on gate/up, at the point the branches converge and + // before SwiGLU consumes them — a bias applied after the activation would + // be a different function, not a rounding difference. + if let Some(b) = layer.biases.as_ref() { + gpu.bias_add_f32(&pbs.gate_ffn_batch, &b.gate, n, b.gate.numel())?; + gpu.bias_add_f32(&pbs.up_batch, &b.up, n, b.up.numel())?; + } + // SwiGLU activation feeding w_down. For MQ, we need the // output FWHT-rotated so it matches the pre-rotated w_down // weights. For HFQ, plain silu_mul is enough. silu_mul_f32 @@ -4929,6 +5563,11 @@ pub(crate) fn batch_chunk_delta_net_ffn( arch_has_wmma, )?; + // down_proj's bias, likewise onto the residual stream. + if let Some(b) = layer.biases.as_ref() { + gpu.bias_add_f32(&pbs.x_batch, &b.down, n, b.down.numel())?; + } + Ok(()) } @@ -4977,11 +5616,6 @@ pub(crate) fn batch_chunk_full_attn_attn( let qkv_is_mq3_lloyd = matches!(layer.wq.gpu_dtype, DType::MQ3G256Lloyd); let qkv_is_fp4 = matches!(layer.wq.gpu_dtype, DType::HFP4G32 | DType::MFP4G32); let qkv_is_q8 = matches!(layer.wq.gpu_dtype, DType::Q8_0); - // TQ2G128/BQ1G128 have no fused qkvza/gate_up/qkv kernel, so they take - // the same UNFUSED plain-GEMM strategy as Q8 rather than falling through - // to the HFQ4 arm, which would read these packed blocks at the wrong - // stride and produce fluent-but-wrong tokens. - let qkv_is_lowbit = matches!(layer.wq.gpu_dtype, DType::TQ2G128 | DType::BQ1G128); // Fused QKV kernels require all three weights to share a // dtype — they treat wq/wk/wv as same-stride byte arrays. // When kmap mode 2 promotes only `v_proj` (issue #249), the @@ -5019,7 +5653,57 @@ pub(crate) fn batch_chunk_full_attn_attn( } // 2. Batched 3-way QKV projection (wq+wk+wv). - if qkv_is_6bit && qkv_same_dtype { + // ESCHA FIRST: a trellis layer takes none of the arms below — they cannot + // read a trellis code, and each projection needs its own rin-rotated + // activation so there is nothing for a fused q/k/v kernel to share. + // `qkv_is_mq` is false for Escha2T16/3T16, so the rmsnorm above left the + // normed (unrotated) activation in `x_rot_batch`, which is what escha + // wants. + if let Some(e) = layer.escha.as_ref() { + // One group holding every slot: `expert_offsets = [0, n]` and the + // identity permutation. Built here rather than per projection so the + // 2-int upload happens once per layer, not six times. + let off_bytes: Vec = [0i32, n as i32] + .iter() + .flat_map(|v| v.to_le_bytes()) + .collect(); + let offsets = gpu.upload_raw(&off_bytes, &[2])?; + let grouped = Some((&offsets, &e.iota)); + + e.q.forward( + gpu, + &layer.wq, + &e.ids, + &pbs.x_rot_batch, + &pbs.escha_xh_batch, + &pbs.fa_q_full_batch, + &pbs.fa_q_full_batch, + n, + grouped, + )?; + e.k.forward( + gpu, + &layer.wk, + &e.ids, + &pbs.x_rot_batch, + &pbs.escha_xh_batch, + &pbs.fa_k_batch, + &pbs.fa_k_batch, + n, + grouped, + )?; + e.v.forward( + gpu, + &layer.wv, + &e.ids, + &pbs.x_rot_batch, + &pbs.escha_xh_batch, + &pbs.fa_v_batch, + &pbs.fa_v_batch, + n, + grouped, + )?; + } else if qkv_is_6bit && qkv_same_dtype { run_fused_qkv_key( gpu, hipfire_dispatch::types::KernelKey::FusedQkvHfq6G256, @@ -5095,11 +5779,14 @@ pub(crate) fn batch_chunk_full_attn_attn( layer.wq.k, n, )?; - } else if qkv_is_q8 && q8_wmma_arch && qkv_same_dtype { - debug_assert!( - matches!(layer.wk.gpu_dtype, DType::Q8_0) && matches!(layer.wv.gpu_dtype, DType::Q8_0), - "FA qkv Q8 WMMA dispatch requires all of wq/wk/wv to be Q8_0", - ); + } else if qkv_is_q8 + && q8_wmma_arch + && all_q8_0(&[layer.wq.gpu_dtype, layer.wk.gpu_dtype, layer.wv.gpu_dtype]) + { + // All three checked at runtime rather than by `qkv_same_dtype` plus a + // `debug_assert!`: the equality predicate and the Q8_0 anchor are two + // facts, and only their conjunction licenses one fused Q8_0-stride + // launch over three buffers. run_fused_qkv_key( gpu, hipfire_dispatch::types::KernelKey::FusedQkvQ8_0, @@ -5116,10 +5803,10 @@ pub(crate) fn batch_chunk_full_attn_attn( layer.wq.k, n, )?; - } else if (qkv_is_q8 || qkv_is_lowbit) && qkv_same_dtype { + } else if is_unfused_plain_gemm_dtype(layer.wq.gpu_dtype) && qkv_same_dtype { run_plain_gemm_key( gpu, - plain_gemm_key_for(layer.wq.gpu_dtype), + plain_gemm_key_for(layer.wq.gpu_dtype)?, &layer.wq.buf, layer.wq.gpu_dtype, &pbs.x_rot_batch, @@ -5130,7 +5817,7 @@ pub(crate) fn batch_chunk_full_attn_attn( )?; run_plain_gemm_key( gpu, - plain_gemm_key_for(layer.wk.gpu_dtype), + plain_gemm_key_for(layer.wk.gpu_dtype)?, &layer.wk.buf, layer.wk.gpu_dtype, &pbs.x_rot_batch, @@ -5141,7 +5828,7 @@ pub(crate) fn batch_chunk_full_attn_attn( )?; run_plain_gemm_key( gpu, - plain_gemm_key_for(layer.wv.gpu_dtype), + plain_gemm_key_for(layer.wv.gpu_dtype)?, &layer.wv.buf, layer.wv.gpu_dtype, &pbs.x_rot_batch, @@ -5177,6 +5864,15 @@ pub(crate) fn batch_chunk_full_attn_attn( batched_gemm_single_weight(gpu, &layer.wv, &pbs.x_rot_batch, &pbs.fa_v_batch, n)?; } + // Escha dense biases on q/k/v, where the branches converge and BEFORE the + // Q/gate deinterleave and q_norm consume them — after either would be a + // different function, not a rounding difference. + if let Some(b) = layer.biases.as_ref() { + gpu.bias_add_f32(&pbs.fa_q_full_batch, &b.q, n, b.q.numel())?; + gpu.bias_add_f32(&pbs.fa_k_batch, &b.k, n, b.k.numel())?; + gpu.bias_add_f32(&pbs.fa_v_batch, &b.v, n, b.v.numel())?; + } + // 3. Batched deinterleave Q + gate: one kernel launch for all N tokens. gpu.deinterleave_f32_batched( &pbs.fa_q_full_batch, @@ -5369,16 +6065,50 @@ pub(crate) fn batch_chunk_full_attn_attn( } else { &pbs.fa_attn_out_batch }; - dispatch_batched_gemm_epilogue( - gpu, - pbs, - &layer.wo, - fa_wo_input, - &epilogue, - n, - q8_wmma_arch, - arch_has_wmma, - )?; + if let Some(e) = layer.escha.as_ref() { + // One group holding every slot: `expert_offsets = [0, n]` and the + // identity permutation. Built here rather than per projection so the + // 2-int upload happens once per layer, not six times. + let off_bytes: Vec = [0i32, n as i32] + .iter() + .flat_map(|v| v.to_le_bytes()) + .collect(); + let offsets = gpu.upload_raw(&off_bytes, &[2])?; + let grouped = Some((&offsets, &e.iota)); + + // Trellis o_proj from the UNROTATED attention output, then accumulate. + // The fused epilogue is SKIPPED, not supplemented: it writes into the + // residual itself, so running both would count this projection twice. + e.o.forward( + gpu, + &layer.wo, + &e.ids, + &pbs.fa_attn_out_batch, + &pbs.escha_xh_batch, + &pbs.escha_y_batch, + &pbs.escha_y_batch, + n, + grouped, + )?; + gpu.add_inplace_f32(&pbs.x_batch, &pbs.escha_y_batch)?; + } else { + dispatch_batched_gemm_epilogue( + gpu, + pbs, + &layer.wo, + fa_wo_input, + &epilogue, + n, + q8_wmma_arch, + arch_has_wmma, + )?; + } + + // o_proj's bias, onto the residual stream — additive, so after the + // residual add is the same value as before it. + if let Some(b) = layer.biases.as_ref() { + gpu.bias_add_f32(&pbs.x_batch, &b.o, n, b.o.numel())?; + } Ok(()) } @@ -5396,6 +6126,76 @@ pub(crate) fn batch_chunk_full_attn_ffn( arch_has_wmma: bool, epilogue: BatchEpilogue<'_>, ) -> HipResult<()> { + // ── ESCHA TRELLIS FFN ─────────────────────────────────────────────── + // Whole FFN in one branch: plain rmsnorm (escha applies its own H128, so a + // pre-rotated input would be rotated twice), gate and up as separate + // trellis GEMVs, SwiGLU, then down accumulated into the residual. The + // fused gate_up kernel cannot serve this — the two projections have + // different rin and it could not read a trellis code regardless. + if let Some(e) = layer.escha.as_ref() { + // One group holding every slot: `expert_offsets = [0, n]` and the + // identity permutation. Built here rather than per projection so the + // 2-int upload happens once per layer, not six times. + let off_bytes: Vec = [0i32, n as i32] + .iter() + .flat_map(|v| v.to_le_bytes()) + .collect(); + let offsets = gpu.upload_raw(&off_bytes, &[2])?; + let grouped = Some((&offsets, &e.iota)); + + gpu.rmsnorm_batched( + &pbs.x_batch, + &layer.ffn_norm, + &pbs.x_norm_batch, + n, + layer.w_gate.k, + config.norm_eps, + )?; + e.gate.forward( + gpu, + &layer.w_gate, + &e.ids, + &pbs.x_norm_batch, + &pbs.escha_xh_batch, + &pbs.gate_ffn_batch, + &pbs.gate_ffn_batch, + n, + grouped, + )?; + e.up.forward( + gpu, + &layer.w_up, + &e.ids, + &pbs.x_norm_batch, + &pbs.escha_xh_batch, + &pbs.up_batch, + &pbs.up_batch, + n, + grouped, + )?; + if let Some(b) = layer.biases.as_ref() { + gpu.bias_add_f32(&pbs.gate_ffn_batch, &b.gate, n, b.gate.numel())?; + gpu.bias_add_f32(&pbs.up_batch, &b.up, n, b.up.numel())?; + } + gpu.silu_mul_f32(&pbs.gate_ffn_batch, &pbs.up_batch, &pbs.ffn_hidden_batch)?; + e.down.forward( + gpu, + &layer.w_down, + &e.ids, + &pbs.ffn_hidden_batch, + &pbs.escha_xh_batch, + &pbs.escha_y_batch, + &pbs.escha_y_batch, + n, + grouped, + )?; + gpu.add_inplace_f32(&pbs.x_batch, &pbs.escha_y_batch)?; + if let Some(b) = layer.biases.as_ref() { + gpu.bias_add_f32(&pbs.x_batch, &b.down, n, b.down.numel())?; + } + return Ok(()); + } + // 10. FFN: rmsnorm (+ rotate for MQ), gate+up, silu_mul // (+ rotate for MQ), w_down residual. let fa_ffn_is_mq = matches!( @@ -5417,11 +6217,6 @@ pub(crate) fn batch_chunk_full_attn_ffn( let fa_ffn_is_mq3_lloyd = matches!(layer.w_gate.gpu_dtype, DType::MQ3G256Lloyd); let fa_ffn_is_fp4 = matches!(layer.w_gate.gpu_dtype, DType::HFP4G32 | DType::MFP4G32); let fa_ffn_is_q8 = matches!(layer.w_gate.gpu_dtype, DType::Q8_0); - // TQ2G128/BQ1G128 have no fused qkvza/gate_up/qkv kernel, so they take - // the same UNFUSED plain-GEMM strategy as Q8 rather than falling through - // to the HFQ4 arm, which would read these packed blocks at the wrong - // stride and produce fluent-but-wrong tokens. - let fa_ffn_is_lowbit = matches!(layer.w_gate.gpu_dtype, DType::TQ2G128 | DType::BQ1G128); if fa_ffn_is_mq { // AWQ-aware: next linear is w_gate (FA-FFN, gate/up share input). fused_rmsnorm_rotate_mq_batched_for( @@ -5462,11 +6257,13 @@ pub(crate) fn batch_chunk_full_attn_ffn( layer.w_gate.k, n, )?; - } else if fa_ffn_is_q8 && q8_wmma_arch { - debug_assert!( - matches!(layer.w_up.gpu_dtype, DType::Q8_0), - "FA FFN Q8 WMMA dispatch requires both w_gate and w_up to be Q8_0", - ); + } else if fa_ffn_is_q8 + && q8_wmma_arch + && all_q8_0(&[layer.w_gate.gpu_dtype, layer.w_up.gpu_dtype]) + { + // `fa_ffn_is_q8` inspects w_gate only; the fused kernel reads BOTH + // weights in one launch at the Q8_0 stride. Runtime, not + // `debug_assert!` — see `all_q8_0`. run_fused_gate_up_key( gpu, hipfire_dispatch::types::KernelKey::FusedGateUpQ8_0, @@ -5480,10 +6277,10 @@ pub(crate) fn batch_chunk_full_attn_ffn( layer.w_gate.k, n, )?; - } else if fa_ffn_is_q8 || fa_ffn_is_lowbit { + } else if is_unfused_plain_gemm_dtype(layer.w_gate.gpu_dtype) { run_plain_gemm_key( gpu, - plain_gemm_key_for(layer.w_gate.gpu_dtype), + plain_gemm_key_for(layer.w_gate.gpu_dtype)?, &layer.w_gate.buf, layer.w_gate.gpu_dtype, &pbs.x_rot_batch, @@ -5494,7 +6291,7 @@ pub(crate) fn batch_chunk_full_attn_ffn( )?; run_plain_gemm_key( gpu, - plain_gemm_key_for(layer.w_up.gpu_dtype), + plain_gemm_key_for(layer.w_up.gpu_dtype)?, &layer.w_up.buf, layer.w_up.gpu_dtype, &pbs.x_rot_batch, @@ -5560,6 +6357,11 @@ pub(crate) fn batch_chunk_full_attn_ffn( n, )?; } + // Escha gate/up biases, before SwiGLU consumes them. + if let Some(b) = layer.biases.as_ref() { + gpu.bias_add_f32(&pbs.gate_ffn_batch, &b.gate, n, b.gate.numel())?; + gpu.bias_add_f32(&pbs.up_batch, &b.up, n, b.up.numel())?; + } let fa_w_down_is_mq = matches!( layer.w_down.gpu_dtype, DType::MQ4G256 @@ -5598,6 +6400,11 @@ pub(crate) fn batch_chunk_full_attn_ffn( arch_has_wmma, )?; + // down_proj's bias, onto the residual stream. + if let Some(b) = layer.biases.as_ref() { + gpu.bias_add_f32(&pbs.x_batch, &b.down, n, b.down.numel())?; + } + Ok(()) } @@ -5705,11 +6512,6 @@ fn batch_chunk_delta_net_moe( ); let is_6bit = matches!(layer.wqkv.gpu_dtype, DType::MQ6G256 | DType::HFQ6G256); let is_q8 = matches!(layer.wqkv.gpu_dtype, DType::Q8_0); - // TQ2G128/BQ1G128 have no fused qkvza/gate_up/qkv kernel, so they take - // the same UNFUSED plain-GEMM strategy as Q8 rather than falling through - // to the HFQ4 arm, which would read these packed blocks at the wrong - // stride and produce fluent-but-wrong tokens. - let is_lowbit = matches!(layer.wqkv.gpu_dtype, DType::TQ2G128 | DType::BQ1G128); // Phase 1.5: PARO mode for DeltaNetMoe — wqkv/wz are // ParoQ4G128 (each with its own Givens rotation tables); // w_alpha/w_beta are F32 (no rotation, no quantization). @@ -5855,18 +6657,21 @@ fn batch_chunk_delta_net_moe( layer.wqkv.k, n, )?; - } else if is_q8 && q8_wmma_arch { - // Fused Q8 QKVZA WMMA — assumes all 4 weights share Q8_0 - // stride; mixed Q8/other layers within DNMoe are rejected - // upstream by `moe_ffn_batched_admissible` (router/gate Q8 OK, but - // shared_expert + experts must be MQ4) and would otherwise - // re-introduce Tier-1 stride corruption. - debug_assert!( - matches!(layer.wz.gpu_dtype, DType::Q8_0) - && matches!(layer.w_beta.gpu_dtype, DType::Q8_0) - && matches!(layer.w_alpha.gpu_dtype, DType::Q8_0), - "DNMoe LA qkvza Q8 WMMA dispatch requires all of wqkv/wz/w_beta/w_alpha to be Q8_0", - ); + } else if is_q8 + && q8_wmma_arch + && all_q8_0(&[ + layer.wqkv.gpu_dtype, + layer.wz.gpu_dtype, + layer.w_beta.gpu_dtype, + layer.w_alpha.gpu_dtype, + ]) + { + // Fused Q8 QKVZA WMMA — reads all four weights in ONE launch at the + // Q8_0 stride, so all four are checked at runtime here. This is the + // arm escha-35b would have hit: its wqkv/wz are Q8_0 but w_alpha / + // w_beta are F16, and the guard used to be a `debug_assert!` that a + // release build compiled out. It now falls through to the per-weight + // unfused arm below, which dispatches each weight by its own dtype. run_fused_qkvza_key( gpu, hipfire_dispatch::types::KernelKey::FusedQkvzaQ8_0, @@ -5886,12 +6691,12 @@ fn batch_chunk_delta_net_moe( layer.wqkv.k, n, )?; - } else if is_q8 || is_lowbit { + } else if is_unfused_plain_gemm_dtype(layer.wqkv.gpu_dtype) { // #397 Ship 5.2 slice1: four plain Q8 batched GEMMs // (wqkv/wz/w_beta/w_alpha), sibling DeltaNet QKVZA path. run_plain_gemm_key( gpu, - plain_gemm_key_for(layer.wqkv.gpu_dtype), + plain_gemm_key_for(layer.wqkv.gpu_dtype)?, &layer.wqkv.buf, layer.wqkv.gpu_dtype, &pbs.x_rot_batch, @@ -5902,7 +6707,7 @@ fn batch_chunk_delta_net_moe( )?; run_plain_gemm_key( gpu, - plain_gemm_key_for(layer.wz.gpu_dtype), + plain_gemm_key_for(layer.wz.gpu_dtype)?, &layer.wz.buf, layer.wz.gpu_dtype, &pbs.x_rot_batch, @@ -5913,7 +6718,7 @@ fn batch_chunk_delta_net_moe( )?; run_plain_gemm_key( gpu, - plain_gemm_key_for(layer.w_beta.gpu_dtype), + plain_gemm_key_for(layer.w_beta.gpu_dtype)?, &layer.w_beta.buf, layer.w_beta.gpu_dtype, &pbs.x_rot_batch, @@ -5924,7 +6729,7 @@ fn batch_chunk_delta_net_moe( )?; run_plain_gemm_key( gpu, - plain_gemm_key_for(layer.w_alpha.gpu_dtype), + plain_gemm_key_for(layer.w_alpha.gpu_dtype)?, &layer.w_alpha.buf, layer.w_alpha.gpu_dtype, &pbs.x_rot_batch, @@ -6072,8 +6877,13 @@ fn batch_chunk_delta_net_moe( gpu.memcpy_dtod_auto(&pbs.dn_q_batch.buf, &pbs.dn_q_raw_batch.buf, n * k_dim * 4)?; gpu.memcpy_dtod_auto(&pbs.dn_k_batch.buf, &pbs.dn_k_raw_batch.buf, n * k_dim * 4)?; } - // DIAG: dump GDN inputs (batched, MoE branch) - if layer_idx == 0 { + // DIAG: dump GDN inputs (batched, MoE branch). + // + // The layer is selectable (`HIPFIRE_DUMP_HIDDEN_LAYER`, default 0) because + // a per-layer x_batch diff localises a divergence to ONE layer and the + // next question is always "what did that layer's GDN see?" — which was + // unanswerable while this was hard-wired to layer 0. + if layer_idx == dump_diag_layer() { let qk_dim = n_v_heads * hd; dump_hidden_localize(gpu, &pbs.dn_q_batch, n, start_pos, qk_dim, 0, "q_b"); dump_hidden_localize(gpu, &pbs.dn_k_batch, n, start_pos, qk_dim, 0, "k_b"); @@ -6253,8 +7063,8 @@ fn batch_chunk_delta_net_moe( config.linear_value_head_dim, )?, } - // DIAG: dump GDN attention output at layer 0 - if layer_idx == 0 { + // DIAG: dump GDN attention output at the selected diag layer. + if layer_idx == dump_diag_layer() { dump_hidden_localize( gpu, &pbs.dn_attn_out_batch, @@ -6276,6 +7086,26 @@ fn batch_chunk_delta_net_moe( config.norm_eps, n, )?; + if layer_idx == dump_diag_layer() { + let vd = n_v_heads * config.linear_value_head_dim; + dump_hidden_localize(gpu, &pbs.dn_normed_batch, n, start_pos, vd, 0, "dnnorm_b"); + dump_hidden_localize(gpu, &pbs.dn_z_batch, n, start_pos, vd, 0, "dnz_b"); + } + // DIAG: the two points that split a per-layer divergence inside the LA + // block. `prewo_b` is the residual stream just BEFORE the wo add, so + // `attn_b - prewo_b` isolates wo alone; `larot_b` is the layer's rmsnorm + // output, whose magnitude against `dnnorm_b` is what identifies wo reading + // the wrong activation buffer (a ~400x ratio, in the case this port hit). + dump_hidden_localize(gpu, &pbs.x_batch, n, start_pos, dim, layer_idx, "prewo_b"); + dump_hidden_localize( + gpu, + &pbs.x_rot_batch, + n, + start_pos, + dim, + layer_idx, + "larot_b", + ); // wo + residual. Q8 wo lands un-rotated (Q8 weights were // quantized against un-rotated activations); MQ4/MQ6 wo // require FWHT(awq_scale-adjusted) rotation. Mirrors the @@ -6285,14 +7115,9 @@ fn batch_chunk_delta_net_moe( // stream when dispatched through the HFQ4 kernel against // 200 B/group MQ6-layout bytes. let dn_wo_is_q8 = matches!(layer.wo.gpu_dtype, DType::Q8_0); - // TQ2G128/BQ1G128 have no fused qkvza/gate_up/qkv kernel, so they take - // the same UNFUSED plain-GEMM strategy as Q8 rather than falling through - // to the HFQ4 arm, which would read these packed blocks at the wrong - // stride and produce fluent-but-wrong tokens. - let dn_wo_is_lowbit = matches!(layer.wo.gpu_dtype, DType::TQ2G128 | DType::BQ1G128); let dn_wo_is_6bit = matches!(layer.wo.gpu_dtype, DType::MQ6G256 | DType::HFQ6G256); let dn_wo_is_paro = matches!(layer.wo.gpu_dtype, DType::ParoQ4G128); - let dn_wo_input = if dn_wo_is_q8 { + let dn_wo_input = if dn_wo_is_q8 || matches!(layer.wo.gpu_dtype, DType::F16) { &pbs.dn_normed_batch } else if dn_wo_is_paro { // PARO wo: rotate dn_normed by wo's own Givens tables @@ -6349,14 +7174,14 @@ fn batch_chunk_delta_net_moe( layer.wo.k, n, )?; - } else if dn_wo_is_q8 || dn_wo_is_lowbit { + } else if is_unfused_plain_gemm_dtype(layer.wo.gpu_dtype) { // Non-WMMA Q8: gemm into a scratch then add into x_batch. // Reuse `dn_normed_rot_batch` (free since the MQ4 rotate // path didn't run here) as the GEMM scratch. let scratch = pbs.dn_normed_rot_batch.sub_offset(0, n * layer.wo.m); run_plain_gemm_key( gpu, - plain_gemm_key_for(layer.wo.gpu_dtype), + plain_gemm_key_for(layer.wo.gpu_dtype)?, &layer.wo.buf, layer.wo.gpu_dtype, dn_wo_input, @@ -6404,6 +7229,10 @@ fn batch_chunk_delta_net_moe( // silu_mul + w_down) block. Takes pbs.x_batch as input AND // accumulates the FFN output residual back into it via the // batched indexed down kernel's atomicAdd path. + // DIAG: x_batch after the LA/attention residual, BEFORE the MoE FFN. + // Splits a per-layer divergence into "attention half" vs "FFN half" + // without another build. + dump_hidden_localize(gpu, &pbs.x_batch, n, start_pos, dim, layer_idx, "attn_b"); prefill_moe_ffn_body_batched( gpu, &layer.ffn, @@ -6472,11 +7301,6 @@ fn batch_chunk_full_attn_moe( ); let qkv_is_6bit = matches!(layer.wq.gpu_dtype, DType::MQ6G256 | DType::HFQ6G256); let qkv_is_q8 = matches!(layer.wq.gpu_dtype, DType::Q8_0); - // TQ2G128/BQ1G128 have no fused qkvza/gate_up/qkv kernel, so they take - // the same UNFUSED plain-GEMM strategy as Q8 rather than falling through - // to the HFQ4 arm, which would read these packed blocks at the wrong - // stride and produce fluent-but-wrong tokens. - let qkv_is_lowbit = matches!(layer.wq.gpu_dtype, DType::TQ2G128 | DType::BQ1G128); // Phase 1.6 (PARO FullAttnMoe): wq/wk/wv are ParoQ4G128 // (each with its own Givens rotation tables). The fused-QKV // kernels can't handle this — they assume one shared @@ -6618,11 +7442,14 @@ fn batch_chunk_full_attn_moe( layer.wq.k, n, )?; - } else if qkv_is_q8 && q8_wmma_arch && qkv_same_dtype { - debug_assert!( - matches!(layer.wk.gpu_dtype, DType::Q8_0) && matches!(layer.wv.gpu_dtype, DType::Q8_0), - "FAMoe qkv Q8 WMMA dispatch requires all of wq/wk/wv to be Q8_0", - ); + } else if qkv_is_q8 + && q8_wmma_arch + && all_q8_0(&[layer.wq.gpu_dtype, layer.wk.gpu_dtype, layer.wv.gpu_dtype]) + { + // All three checked at runtime rather than by `qkv_same_dtype` plus a + // `debug_assert!`: the equality predicate and the Q8_0 anchor are two + // facts, and only their conjunction licenses one fused Q8_0-stride + // launch over three buffers. run_fused_qkv_key( gpu, hipfire_dispatch::types::KernelKey::FusedQkvQ8_0, @@ -6639,10 +7466,10 @@ fn batch_chunk_full_attn_moe( layer.wq.k, n, )?; - } else if (qkv_is_q8 || qkv_is_lowbit) && qkv_same_dtype { + } else if is_unfused_plain_gemm_dtype(layer.wq.gpu_dtype) && qkv_same_dtype { run_plain_gemm_key( gpu, - plain_gemm_key_for(layer.wq.gpu_dtype), + plain_gemm_key_for(layer.wq.gpu_dtype)?, &layer.wq.buf, layer.wq.gpu_dtype, &pbs.x_rot_batch, @@ -6653,7 +7480,7 @@ fn batch_chunk_full_attn_moe( )?; run_plain_gemm_key( gpu, - plain_gemm_key_for(layer.wk.gpu_dtype), + plain_gemm_key_for(layer.wk.gpu_dtype)?, &layer.wk.buf, layer.wk.gpu_dtype, &pbs.x_rot_batch, @@ -6664,7 +7491,7 @@ fn batch_chunk_full_attn_moe( )?; run_plain_gemm_key( gpu, - plain_gemm_key_for(layer.wv.gpu_dtype), + plain_gemm_key_for(layer.wv.gpu_dtype)?, &layer.wv.buf, layer.wv.gpu_dtype, &pbs.x_rot_batch, @@ -6851,18 +7678,13 @@ fn batch_chunk_full_attn_moe( // — catastrophic stride mismatch produces a single-token // attractor on AWQ A3B's 4/40 FA layers with MQ6 wo). let fa_wo_is_q8 = matches!(layer.wo.gpu_dtype, DType::Q8_0); - // TQ2G128/BQ1G128 have no fused qkvza/gate_up/qkv kernel, so they take - // the same UNFUSED plain-GEMM strategy as Q8 rather than falling through - // to the HFQ4 arm, which would read these packed blocks at the wrong - // stride and produce fluent-but-wrong tokens. - let fa_wo_is_lowbit = matches!(layer.wo.gpu_dtype, DType::TQ2G128 | DType::BQ1G128); let fa_wo_is_6bit = matches!(layer.wo.gpu_dtype, DType::MQ6G256 | DType::HFQ6G256); // Phase 1.6 (PARO FullAttnMoe wo): own Givens rotation table, // 72 B/group HFQ4G128 layout. Rotate fa_attn_out_batch by wo's // paro into fa_attn_out_rot_batch, then HFQ4G128 GEMM into a // scratch, then add into x_batch. let fa_wo_is_paro = matches!(layer.wo.gpu_dtype, DType::ParoQ4G128); - let fa_wo_input = if fa_wo_is_q8 { + let fa_wo_input = if fa_wo_is_q8 || matches!(layer.wo.gpu_dtype, DType::F16) { &pbs.fa_attn_out_batch } else if fa_wo_is_paro { let paro_wo = layer.wo.paro.as_ref().unwrap_or_else(|| { @@ -6916,14 +7738,14 @@ fn batch_chunk_full_attn_moe( layer.wo.k, n, )?; - } else if fa_wo_is_q8 || fa_wo_is_lowbit { + } else if is_unfused_plain_gemm_dtype(layer.wo.gpu_dtype) { // Non-WMMA Q8: GEMM into a scratch then add into x_batch. // Reuse `fa_attn_out_rot_batch` (free since MQ4 rotate // didn't run here) as scratch. let scratch = pbs.fa_attn_out_rot_batch.sub_offset(0, n * layer.wo.m); run_plain_gemm_key( gpu, - plain_gemm_key_for(layer.wo.gpu_dtype), + plain_gemm_key_for(layer.wo.gpu_dtype)?, &layer.wo.buf, layer.wo.gpu_dtype, fa_wo_input, @@ -6968,6 +7790,10 @@ fn batch_chunk_full_attn_moe( } // Batched MoE FFN. + // DIAG: x_batch after the LA/attention residual, BEFORE the MoE FFN. + // Splits a per-layer divergence into "attention half" vs "FFN half" + // without another build. + dump_hidden_localize(gpu, &pbs.x_batch, n, start_pos, dim, layer_idx, "attn_b"); prefill_moe_ffn_body_batched( gpu, &layer.ffn, @@ -7532,6 +8358,18 @@ fn run_fa_layer_body( weight_gemv_prerotated(gpu, &layer.wv, &s.tmp, x_rot, &s.fa_v)?; } + // Escha q/k/v biases. THIS is the path the 27B's full-attention layers + // actually take: `fa_batched_ok` refuses MQ6 weights, so the batched arm + // never runs and `batch_chunk_full_attn_fallback` walks tokens through + // here instead. Adding them to `batch_chunk_full_attn_attn` alone moved + // PPL by exactly zero, which is how the wrong route was caught. + // Must precede the Q/gate deinterleave and q_norm. + if let Some(b) = layer.biases.as_ref() { + gpu.bias_add_f32(&s.fa_q_full, &b.q, 1, b.q.numel())?; + gpu.bias_add_f32(&s.fa_k, &b.k, 1, b.k.numel())?; + gpu.bias_add_f32(&s.fa_v, &b.v, 1, b.v.numel())?; + } + gpu.deinterleave_f32( &s.fa_q_full, &s.fa_q, @@ -7638,6 +8476,12 @@ fn run_fa_layer_body( .map_err(|e| hip_bridge::HipError::new(0, &e.to_string()))?; } + // o_proj bias onto the residual stream (additive, so order with the + // residual add does not matter). + if let Some(b) = layer.biases.as_ref() { + gpu.bias_add_f32(&s.x, &b.o, 1, b.o.numel())?; + } + // FFN: fused rmsnorm + rotate for w_gate/w_up. let x_rot = fused_rmsnorm_rotate_for_mq( gpu, @@ -7743,7 +8587,17 @@ fn run_fa_layer_body( weight_gemv_prerotated(gpu, &layer.w_gate, &s.tmp, x_rot, &s.gate_ffn)?; weight_gemv_prerotated(gpu, &layer.w_up, &s.tmp, x_rot, &s.up)?; } + + // gate/up biases BEFORE the SwiGLU that `weight_gemv_swiglu_residual` + // applies; after it would be a different function. + if let Some(b) = layer.biases.as_ref() { + gpu.bias_add_f32(&s.gate_ffn, &b.gate, 1, b.gate.numel())?; + gpu.bias_add_f32(&s.up, &b.up, 1, b.up.numel())?; + } weight_gemv_swiglu_residual(gpu, &layer.w_down, &s.gate_ffn, &s.up, &s.ffn_hidden, &s.x)?; + if let Some(b) = layer.biases.as_ref() { + gpu.bias_add_f32(&s.x, &b.down, 1, b.down.numel())?; + } Ok(()) } @@ -7774,6 +8628,23 @@ fn batched_gemm_single_weight( n: usize, ) -> HipResult<()> { match w.gpu_dtype { + // F16: `linear_attn.in_proj_a`/`in_proj_b` on the escha exports, which + // escha's `ignore` list leaves uncoded while every sibling projection + // is a trellis code. `plain_gemm_key_for` already sanctions this exact + // mapping; this match simply had no arm for it, which surfaced as + // "weight dtype F16 has no single-weight batched dispatch yet" the + // moment escha layers were admitted to batched prefill. + DType::F16 => run_plain_gemm_key( + gpu, + hipfire_dispatch::types::KernelKey::GemmF16WmmaMb8, + &w.buf, + w.gpu_dtype, + x, + y, + w.m, + w.k, + n, + ), DType::MQ4G256 | DType::HFQ4G256 => run_plain_gemm_key( gpu, hipfire_dispatch::types::KernelKey::GemmHfq4G256, @@ -8006,8 +8877,8 @@ mod tests { #[test] fn prefill_max_batch_arch_defaults() { - // Exact gfx1100 / gfx1201 alone get the measured defaults; every other - // string keeps the conservative PREFILL_MAX_BATCH=256 ceiling. + // Exact gfx1100 / gfx1201 / gfx1151 get the measured defaults; every + // other string keeps the conservative PREFILL_MAX_BATCH=256 ceiling. // Pure helper — no process env mutation. assert_eq!(prefill_max_batch_for_arch("gfx1100"), 512); assert_eq!( @@ -8019,7 +8890,14 @@ mod tests { prefill_max_batch_for_arch("gfx1201"), PREFILL_DEFAULT_BATCH_GFX1201 ); - for arch in ["gfx1200", "gfx1151", "gfx942", "unknown"] { + // gfx1151 measured on the native escha 27B at an 8k prompt: + // 256 -> 52 tok/s, 512 -> 73, 1024 -> 69. + assert_eq!(prefill_max_batch_for_arch("gfx1151"), 512); + assert_eq!( + prefill_max_batch_for_arch("gfx1151"), + PREFILL_DEFAULT_BATCH_GFX1151 + ); + for arch in ["gfx1200", "gfx1152", "gfx942", "unknown"] { assert_eq!( prefill_max_batch_for_arch(arch), PREFILL_MAX_BATCH, @@ -9309,4 +10187,165 @@ mod tests { let full_chunk = moe_grouped_m_total_bound(2048, 256); assert_eq!(full_chunk, 5888); } + + // ── Escha-W2 batched prefill (task perf-3) ─────────────────────────── + + /// `is_batchable_la` admits F16 on WMMA arches only. + /// + /// The arch gate is not decoration. The only batched F16 GEMM is + /// `gemm_f16_wmma_mb8`, which hard-ERRORS without wave32 WMMA, and an + /// admitted-then-erroring layer is not a graceful fallback: the error + /// propagates out of `forward_prefill_batch`. Refusal here is what keeps a + /// non-WMMA arch on the per-token path it already takes. + #[test] + fn is_batchable_la_f16_is_wmma_arch_only() { + for arch in [ + "gfx1100", "gfx1101", "gfx1102", "gfx1103", "gfx1150", "gfx1151", "gfx1152", "gfx1200", + "gfx1201", + ] { + assert!( + is_batchable_la(DType::F16, arch), + "F16 must be batchable on {arch}" + ); + } + for arch in [ + "gfx906", "gfx908", "gfx942", "gfx1010", "gfx1030", "gfx1031", + ] { + assert!( + !is_batchable_la(DType::F16, arch), + "F16 must NOT be batchable on {arch} — gemm_f16_wmma_mb8 errors there" + ); + } + } + + /// `plain_gemm_key_for` resolves exactly the dtypes + /// `is_unfused_plain_gemm_dtype` admits, and ERRORS on everything else. + /// + /// The two must stay in lockstep. The dangerous direction is a dtype the + /// matchers route to the unfused path with no arm here: before this change + /// the `_ =>` default was `GemmQ8_0BatchedChunked`, so any such dtype was + /// read at a Q8_0 stride — 34-byte blocks of 32 int8 plus an f16 scale — + /// over bytes with a different layout. That is finite, fluent, wrong. + #[test] + fn plain_gemm_key_matches_the_unfused_dtype_set() { + use hipfire_dispatch::types::KernelKey as K; + // Behaviour-preserving for every dtype that reached this before. + assert_eq!( + plain_gemm_key_for(DType::Q8_0).unwrap(), + K::GemmQ8_0BatchedChunked + ); + assert_eq!( + plain_gemm_key_for(DType::TQ2G128).unwrap(), + K::GemmTQ2G128Prefill + ); + assert_eq!( + plain_gemm_key_for(DType::BQ1G128).unwrap(), + K::GemmBQ1G128Prefill + ); + // New arm. NOT GemmF16 / GemmF16Tiled: those write Y as [M, N] while + // every batched-prefill consumer reads [N, M], so picking one of them + // would transpose the output silently. + assert_eq!(plain_gemm_key_for(DType::F16).unwrap(), K::GemmF16WmmaMb8); + for dt in [DType::Q8_0, DType::TQ2G128, DType::BQ1G128, DType::F16] { + assert!( + is_unfused_plain_gemm_dtype(dt), + "{dt:?} resolves a key but is not in the unfused set" + ); + } + // Everything else must ERROR, not guess. + for dt in [ + DType::MQ4G256, + DType::MQ6G256, + DType::MQ3G256, + DType::F32, + DType::ParoQ4G128, + DType::MFP4G32E8, + ] { + assert!( + !is_unfused_plain_gemm_dtype(dt), + "{dt:?} must not be in the unfused set" + ); + assert!( + plain_gemm_key_for(dt).is_err(), + "{dt:?} must error rather than fall through to a Q8_0-stride GEMM" + ); + } + } + + /// The fused Q8_0 QKVZA/QKV/gate+up kernels read several weights in ONE + /// launch at the Q8_0 stride, so the arm must check ALL of them. + /// + /// This was a `debug_assert!` — compiled out of release — and it was + /// unreachable only because `is_batchable_la` refused F16. escha-35b has + /// Q8_0 `wqkv`/`wz` beside F16 `w_alpha`/`w_beta`, so admitting F16 made a + /// Q8_0-strided read of F16 bytes reachable in a release build. + #[test] + fn all_q8_0_rejects_a_mixed_layer() { + assert!(all_q8_0(&[DType::Q8_0; 4])); + assert!(all_q8_0(&[])); + // The exact escha shape: the anchor weight is Q8_0, two others are not. + assert!(!all_q8_0(&[ + DType::Q8_0, + DType::Q8_0, + DType::F16, + DType::F16 + ])); + assert!(!all_q8_0(&[DType::Q8_0, DType::MQ4G256])); + } + + /// The escha admission arm, and the proof that it is ADDITIVE. + /// + /// The same six dtypes without the escha marker must still be refused — + /// otherwise this arm would be admitting some other model's layer to an + /// executor that assumes escha's rotated weight domain and its H128 pair. + #[test] + fn moe_prefill_escha_arm_requires_the_escha_marker() { + let mut d = MoePrefillDtypes::uniform(DType::Q8_0); + d.router = DType::F16; + d.shared_expert_scalar_gate = DType::F16; + // shared expert + routed experts are Q8_0 from `uniform`. + assert!( + !moe_ffn_batched_admissible_for_dtypes(&d, false, false, false, false), + "Q8_0 routed + Q8_0 shared must stay refused WITHOUT the escha marker — no other \ + arm serves that combination and the escha executor must not see a non-escha layer" + ); + d.escha = true; + assert!( + moe_ffn_batched_admissible_for_dtypes(&d, false, false, false, false), + "the escha arm must admit F16 router + F16 scalar gate + Q8_0 shared + Q8_0 routed" + ); + // The marker alone is not enough: the executor hard-codes the Q8_0 + // block decode, so a non-Q8_0 routed projection must still refuse. + let mut mixed = d; + mixed.expert_down = DType::MQ4G256; + assert!(!moe_ffn_batched_admissible_for_dtypes( + &mixed, false, false, false, false + )); + } + + /// Widening `router_ok` / `shared_gate_ok` to F16 must not change any + /// decision for a model whose router and scalar gate were already + /// admissible — the additivity claim, checked rather than asserted. + #[test] + fn moe_prefill_f16_router_widening_is_additive() { + for base in [DType::MQ4G256, DType::MQ4G256V2] { + let d = MoePrefillDtypes::uniform(base); + // Admitted before and after; the F16 arms are simply not reached. + assert!(moe_ffn_batched_admissible_for_dtypes( + &d, false, false, false, false + )); + } + // A router dtype that is in NEITHER the old nor the new allowlist must + // still refuse — the widening added F16 and nothing else. + let mut d = MoePrefillDtypes::uniform(DType::MQ4G256); + d.router = DType::MQ6G256; + assert!(!moe_ffn_batched_admissible_for_dtypes( + &d, true, false, false, false + )); + d.router = DType::MQ4G256; + d.shared_expert_scalar_gate = DType::MQ3G256; + assert!(!moe_ffn_batched_admissible_for_dtypes( + &d, true, false, false, false + )); + } } diff --git a/crates/hipfire-arch-qwen35/src/qwen35/weights.rs b/crates/hipfire-arch-qwen35/src/qwen35/weights.rs index cf25aa7ead..8b0569b450 100644 --- a/crates/hipfire-arch-qwen35/src/qwen35/weights.rs +++ b/crates/hipfire-arch-qwen35/src/qwen35/weights.rs @@ -37,6 +37,47 @@ pub struct DeltaNetLayerWeights { pub w_gate: WeightTensor, // mlp.gate_proj pub w_up: WeightTensor, // mlp.up_proj pub w_down: WeightTensor, // mlp.down_proj + /// Additive output biases, present only on escha dense exports + /// (Qwen3.8-27B). Base Qwen3.8-27B has `attention_bias: false` and no MLP + /// bias — these are Escha's end-to-end output correction, and they cannot + /// be folded into a weight because they are additive. `None` for every + /// other model, which is why they are Option rather than a zero vector: + /// a zero add per projection per layer is real work for no effect. + pub biases: Option, + /// Escha trellis metadata, one per coded projection. `Some` only when the + /// weights are `Escha2T16`/`Escha3T16`, which is the signal this layer + /// must bypass the fused MQ paths entirely — each projection needs its + /// own rin-rotated activation, so FusedQkvza/gate_up have nothing to + /// share. + pub escha: Option, +} + +/// See `DeltaNetLayerWeights::escha`. +pub struct DeltaNetEscha { + pub qkv: crate::qwen35::escha::EschaProj, + pub z: crate::qwen35::escha::EschaProj, + pub o: crate::qwen35::escha::EschaProj, + pub gate: crate::qwen35::escha::EschaProj, + pub up: crate::qwen35::escha::EschaProj, + pub down: crate::qwen35::escha::EschaProj, + /// `slots`-long run of zeros for the indexed GEMV; sized for the largest + /// batch the model was built for so decode (slots=1) reads a prefix. + pub ids: GpuTensor, + /// `0..MAX` — the grouped GEMM's slot permutation, which for a dense + /// linear is the identity. + pub iota: GpuTensor, +} + +/// See `DeltaNetLayerWeights::biases`. Each is `[oc]` f32, applied to the +/// projection's output. `in_proj_a`/`in_proj_b` have none — they are on +/// escha's `ignore` list and ship as plain weights. +pub struct DeltaNetBiases { + pub qkv: GpuTensor, + pub z: GpuTensor, + pub o: GpuTensor, + pub gate: GpuTensor, + pub up: GpuTensor, + pub down: GpuTensor, } /// Weights for a full attention (gated) layer — similar to Qwen3 but with q+gate split. @@ -52,6 +93,35 @@ pub struct FullAttnLayerWeights { pub w_gate: WeightTensor, pub w_up: WeightTensor, pub w_down: WeightTensor, + /// See `DeltaNetLayerWeights::biases`. + pub biases: Option, + /// See `DeltaNetLayerWeights::escha`. + pub escha: Option, +} + +/// See `DeltaNetLayerWeights::escha`. +pub struct FullAttnEscha { + pub q: crate::qwen35::escha::EschaProj, + pub k: crate::qwen35::escha::EschaProj, + pub v: crate::qwen35::escha::EschaProj, + pub o: crate::qwen35::escha::EschaProj, + pub gate: crate::qwen35::escha::EschaProj, + pub up: crate::qwen35::escha::EschaProj, + pub down: crate::qwen35::escha::EschaProj, + pub ids: GpuTensor, + /// See `DeltaNetEscha::iota`. + pub iota: GpuTensor, +} + +/// See `DeltaNetLayerWeights::biases`. +pub struct FullAttnBiases { + pub q: GpuTensor, + pub k: GpuTensor, + pub v: GpuTensor, + pub o: GpuTensor, + pub gate: GpuTensor, + pub up: GpuTensor, + pub down: GpuTensor, } // ─── MoE FFN weights (Qwen3.5-MoE / A3B) ──────────────────────────────── @@ -83,15 +153,33 @@ pub struct ExpertWeights { pub down: WeightTensor, // [hidden, moe_intermediate] } -/// Owning storage for a layer's packed uniform-MQ4 routed experts. +/// Owning storage for a layer's packed routed experts — one device buffer per +/// (layer, projection) covering ALL experts. /// /// `experts` still carries one [`WeightTensor`] view per routed expert so the /// CPU fallback and every existing indexed dispatch keep their exact metadata /// and pointer-table ABI. Those views are non-owning subranges of these two /// buffers; only this owner pair may be returned to the GPU pool. -pub(crate) struct PackedExpertOwners { - pub(crate) gate_up: GpuTensor, - pub(crate) down: GpuTensor, +/// +/// Two producers build this: `try_load_packed_mq4_experts` (uniform MQ4) and +/// `escha::load_escha_moe_experts` (Escha-W2). It is `pub` because the latter +/// hands the owners back across the crate boundary to +/// `examples/escha_moe_block_gate`, which loads a layer's experts directly and +/// must free them exactly once. +/// +/// ## Why this is not merely tidier +/// +/// The HIP allocator rounds every allocation up to a 2 MiB granule. At A3B +/// shapes a Q8_0 gate_up is 2.125 MiB and a Q8_0 down is 1.0625 MiB, so 20,480 +/// independent per-expert buffers (40 layers x 256 experts x 2 projections) +/// occupy 4 MiB and 2 MiB each — 64.4 GB of granules for 34.2 GB of weights. +/// Packing each (layer, projection) into ONE buffer pays the rounding once per +/// buffer instead of once per expert and recovers ~30 GB. Measured: 67.9 GB -> +/// 37.6 GB of GTT for the whole escha-35b model on gfx1151 (37 587 996 672 B +/// delta, `scripts/escha-gtt-probe.sh`). +pub struct PackedExpertOwners { + pub gate_up: GpuTensor, + pub down: GpuTensor, } /// SP2: build the per-expert (gate_up, down) quant-tier tables that @@ -300,6 +388,17 @@ pub struct MoeFfnWeights { /// non-owned storage layout. Non-owned global slots alias into the /// matching entry. Owned so `free_moe_ffn` can reclaim them. pub(crate) ep_dummy_buffers: Vec, + + /// Escha-W2 (Task 10): per-layer H128 transform tables + decode scratch. + /// `Some` only for layers loaded from an Escha-W2 checkpoint. + /// + /// This is also the layer's escha MARKER. The loader decodes the trellis + /// and stores the experts as `Q8_0`, so `experts[i].gate_up.gpu_dtype` no + /// longer says "escha" by the time dispatch resolves the layer; only this + /// field does. `moe_ffn_decode_impl` threads it into + /// `MoeParams::escha`, whose `has_escha()` drives both the f16 + /// router-logit round-trip and the H128-wrapped routed executor. + pub escha: Option, } /// Owning storage for the per-layer shared ParoQuant rotation sidecars. @@ -1048,6 +1147,10 @@ impl PendingEpMoeFfn { paro_shared: None, global_expert_dtypes: self.global_dtypes, ep_dummy_buffers: self.dummy_buffers, + // The EP / pending-commit path does not carry escha layers (an + // Escha-W2 checkpoint has one code tensor per layer, not the + // per-expert tensors EP sharding streams). + escha: None, } } } @@ -1340,6 +1443,11 @@ impl MmqScreenable for Qwen35Weights { } fn free_moe_ffn(gpu: &mut Gpu, ffn: MoeFfnWeights) { + // Escha-W2 transform tables + decode scratch. Owned outright (nothing + // aliases them), so free before the experts they describe. + if let Some(e) = ffn.escha { + e.free_gpu(gpu); + } ffn.router.free_all(gpu); ffn.shared_expert_gate.free_all(gpu); ffn.shared_expert.gate.free_all(gpu); diff --git a/crates/hipfire-dispatch-tests/src/dtype.rs b/crates/hipfire-dispatch-tests/src/dtype.rs index 24541330a5..a5c45ac2f5 100644 --- a/crates/hipfire-dispatch-tests/src/dtype.rs +++ b/crates/hipfire-dispatch-tests/src/dtype.rs @@ -2,33 +2,63 @@ use rdna_compute::DType; /// Every DType variant that represents a quantized format (byte-level). const QUANTIZED_DTYPES: &[DType] = &[ - DType::Q4K, DType::Q6K, DType::Q8_0, - DType::Q4F16G64, DType::Q4F16G32, DType::Q8HFQ, - DType::HFQ4G256, DType::HFQ4G128, - DType::HFQ3G256, DType::HFQ3G128, - DType::MQ4G256, DType::MQ4G128, - DType::MQ8G256, DType::MQ6G256, - DType::MQ3G256, DType::MQ2G256, - DType::MQ2G256Lloyd, DType::MQ3G256Lloyd, DType::MQ4G256Lloyd, - DType::HFP4G32, DType::MFP4G32, - DType::HFQ2G256, DType::HFQ2G128, DType::HFQ6G256, - DType::ParoQ4G128, DType::Raw, + DType::Q4K, + DType::Q6K, + DType::Q8_0, + DType::Q4F16G64, + DType::Q4F16G32, + DType::Q8HFQ, + DType::HFQ4G256, + DType::HFQ4G128, + DType::HFQ3G256, + DType::HFQ3G128, + DType::MQ4G256, + DType::MQ4G128, + DType::MQ8G256, + DType::MQ6G256, + DType::MQ3G256, + DType::MQ2G256, + DType::MQ2G256Lloyd, + DType::MQ3G256Lloyd, + DType::MQ4G256Lloyd, + DType::HFP4G32, + DType::MFP4G32, + DType::HFQ2G256, + DType::HFQ2G128, + DType::HFQ6G256, + DType::ParoQ4G128, + // Escha-W2. Listed so `rotation_plan_matches_legacy_needs_fwht` actually + // covers the two dtypes whose rotation plan is `EschaH128`. This branch + // reflowed that list without adding them, which let + // `dtype_needs_rotation` return false for an `EschaH128` dtype with no + // test to catch it. + DType::Escha2T16, + DType::Escha3T16, + DType::Raw, ]; /// DTypes that are MQ-family (FWHT-rotated MagnumQuant). const MAGNUMQUANT_DTYPES: &[DType] = &[ - DType::MQ4G256, DType::MQ4G128, - DType::MQ8G256, DType::MQ6G256, - DType::MQ3G256, DType::MQ2G256, - DType::MQ2G256Lloyd, DType::MQ3G256Lloyd, DType::MQ4G256Lloyd, + DType::MQ4G256, + DType::MQ4G128, + DType::MQ8G256, + DType::MQ6G256, + DType::MQ3G256, + DType::MQ2G256, + DType::MQ2G256Lloyd, + DType::MQ3G256Lloyd, + DType::MQ4G256Lloyd, DType::MFP4G32, ]; /// DTypes that are HFQ-family (flat quant with inline f32 scale+zero). const HFQ_DTYPES: &[DType] = &[ - DType::HFQ4G256, DType::HFQ4G128, - DType::HFQ3G256, DType::HFQ3G128, - DType::HFQ2G256, DType::HFQ2G128, + DType::HFQ4G256, + DType::HFQ4G128, + DType::HFQ3G256, + DType::HFQ3G128, + DType::HFQ2G256, + DType::HFQ2G128, DType::HFQ6G256, ]; @@ -69,12 +99,24 @@ fn awq_sidecar_not_on_non_awq_dtypes() { for dt in MAGNUMQUANT_DTYPES { if matches!( *dt, - DType::MQ4G256 | DType::MQ3G256 | DType::MQ2G256 | DType::MQ3G256Lloyd | DType::MQ2G256Lloyd - ) { continue; } - assert!(!dt.supports_awq_sidecar(), "DType::{dt:?} should NOT support AWQ"); + DType::MQ4G256 + | DType::MQ3G256 + | DType::MQ2G256 + | DType::MQ3G256Lloyd + | DType::MQ2G256Lloyd + ) { + continue; + } + assert!( + !dt.supports_awq_sidecar(), + "DType::{dt:?} should NOT support AWQ" + ); } for dt in HFQ_DTYPES { - assert!(!dt.supports_awq_sidecar(), "DType::{dt:?} should NOT support AWQ"); + assert!( + !dt.supports_awq_sidecar(), + "DType::{dt:?} should NOT support AWQ" + ); } assert!(!DType::F32.supports_awq_sidecar()); assert!(!DType::F16.supports_awq_sidecar()); @@ -103,25 +145,38 @@ fn rotation_plan_covers_every_dtype() { assert_eq!(dtype_rotation_plan(DType::MQ6G256), RotationPlan::FwhtG256); assert_eq!(dtype_rotation_plan(DType::MFP4G32), RotationPlan::FwhtG256); assert_eq!(dtype_rotation_plan(DType::MQ4G128), RotationPlan::FwhtG128); - assert_eq!(dtype_rotation_plan(DType::MQ8G256), RotationPlan::Mq8Internal); + assert_eq!( + dtype_rotation_plan(DType::MQ8G256), + RotationPlan::Mq8Internal + ); assert_eq!(dtype_rotation_plan(DType::ParoQ4G128), RotationPlan::Givens); + assert_eq!( + dtype_rotation_plan(DType::Escha2T16), + RotationPlan::EschaH128 + ); + assert_eq!( + dtype_rotation_plan(DType::Escha3T16), + RotationPlan::EschaH128 + ); } #[test] fn rotation_plan_matches_legacy_needs_fwht() { - use hipfire_dispatch::types::{dtype_rotation_plan, dtype_needs_rotation, RotationPlan}; + use hipfire_dispatch::types::{dtype_needs_rotation, dtype_rotation_plan, RotationPlan}; for d in QUANTIZED_DTYPES { assert_eq!( dtype_rotation_plan(*d) != RotationPlan::None, dtype_needs_rotation(*d), - "rotation_plan/needs_fwht disagree for {:?}", d + "rotation_plan/needs_fwht disagree for {:?}", + d ); } for d in [DType::F32, DType::F16, DType::Q8_0] { assert_eq!( dtype_rotation_plan(d) != RotationPlan::None, dtype_needs_rotation(d), - "rotation_plan/needs_fwht disagree for {:?}", d + "rotation_plan/needs_fwht disagree for {:?}", + d ); } } @@ -130,16 +185,31 @@ fn rotation_plan_matches_legacy_needs_fwht() { fn post_rotation_variant_paro_is_plain_mq_is_prerotated() { use hipfire_dispatch::types::{dtype_post_rotation_variant, GemvVariant}; use rdna_compute::DType; - assert_eq!(dtype_post_rotation_variant(DType::ParoQ4G128), GemvVariant::Plain); - assert_eq!(dtype_post_rotation_variant(DType::MQ4G256), GemvVariant::Prerotated); - assert_eq!(dtype_post_rotation_variant(DType::MQ8G256), GemvVariant::Prerotated); - assert_eq!(dtype_post_rotation_variant(DType::MQ4G128), GemvVariant::Prerotated); - assert_eq!(dtype_post_rotation_variant(DType::HFQ4G256), GemvVariant::Plain); + assert_eq!( + dtype_post_rotation_variant(DType::ParoQ4G128), + GemvVariant::Plain + ); + assert_eq!( + dtype_post_rotation_variant(DType::MQ4G256), + GemvVariant::Prerotated + ); + assert_eq!( + dtype_post_rotation_variant(DType::MQ8G256), + GemvVariant::Prerotated + ); + assert_eq!( + dtype_post_rotation_variant(DType::MQ4G128), + GemvVariant::Prerotated + ); + assert_eq!( + dtype_post_rotation_variant(DType::HFQ4G256), + GemvVariant::Plain + ); } #[test] fn q8hfq_resolves_to_plain_gemv_key() { - use hipfire_dispatch::types::{KernelKey, GemvVariant}; + use hipfire_dispatch::types::{GemvVariant, KernelKey}; use rdna_compute::DType; let key = KernelKey::for_gemv(DType::Q8HFQ, GemvVariant::Plain, false) .expect("Q8HFQ Plain must resolve"); @@ -150,32 +220,114 @@ fn q8hfq_resolves_to_plain_gemv_key() { fn rotation_tag_distinguishes_awq_and_batched() { use hipfire_dispatch::families::gemv::RotationTag; use hipfire_dispatch::types::RotationPlan; - let base = RotationTag { plan: RotationPlan::FwhtG256, awq: false, batched: false }; - let awq = RotationTag { plan: RotationPlan::FwhtG256, awq: true, batched: false }; - let bat = RotationTag { plan: RotationPlan::FwhtG256, awq: false, batched: true }; + let base = RotationTag { + plan: RotationPlan::FwhtG256, + awq: false, + batched: false, + }; + let awq = RotationTag { + plan: RotationPlan::FwhtG256, + awq: true, + batched: false, + }; + let bat = RotationTag { + plan: RotationPlan::FwhtG256, + awq: false, + batched: true, + }; assert_ne!(base, awq, "AWQ vs non-AWQ must not compare equal"); assert_ne!(base, bat, "batched vs non-batched must not compare equal"); - assert_eq!(base, RotationTag { plan: RotationPlan::FwhtG256, awq: false, batched: false }); + assert_eq!( + base, + RotationTag { + plan: RotationPlan::FwhtG256, + awq: false, + batched: false + } + ); } #[test] fn run_rejects_tag_plan_mismatch() { use hipfire_dispatch::families::gemv::{check_rotation_tag, RotationTag}; use hipfire_dispatch::types::RotationPlan; - let want = RotationTag { plan: RotationPlan::FwhtG256, awq: false, batched: false }; - let givens = RotationTag { plan: RotationPlan::Givens, awq: false, batched: false }; - let awq = RotationTag { plan: RotationPlan::FwhtG256, awq: true, batched: false }; + let want = RotationTag { + plan: RotationPlan::FwhtG256, + awq: false, + batched: false, + }; + let givens = RotationTag { + plan: RotationPlan::Givens, + awq: false, + batched: false, + }; + let awq = RotationTag { + plan: RotationPlan::FwhtG256, + awq: true, + batched: false, + }; assert!(check_rotation_tag(want, want).is_ok()); - assert!(check_rotation_tag(want, givens).is_err(), "plan mismatch must reject"); - assert!(check_rotation_tag(want, awq).is_err(), "awq mismatch must reject"); + assert!( + check_rotation_tag(want, givens).is_err(), + "plan mismatch must reject" + ); + assert!( + check_rotation_tag(want, awq).is_err(), + "awq mismatch must reject" + ); } #[test] fn rotate_variant_selection() { use hipfire_dispatch::families::gemv::select_rotation_variant; use hipfire_dispatch::types::{RotationPlan, RotationVariant}; - assert_eq!(select_rotation_variant(RotationPlan::FwhtG256, false, false), RotationVariant::Plain); - assert_eq!(select_rotation_variant(RotationPlan::FwhtG256, true, false), RotationVariant::WithRmsnorm); - assert_eq!(select_rotation_variant(RotationPlan::FwhtG256, false, true), RotationVariant::WithSwiGLU); - assert_eq!(select_rotation_variant(RotationPlan::FwhtG128, false, false), RotationVariant::PlainG128); - assert_eq!(select_rotation_variant(RotationPlan::Givens, false, false), RotationVariant::Givens); + assert_eq!( + select_rotation_variant(RotationPlan::FwhtG256, false, false), + RotationVariant::Plain + ); + assert_eq!( + select_rotation_variant(RotationPlan::FwhtG256, true, false), + RotationVariant::WithRmsnorm + ); + assert_eq!( + select_rotation_variant(RotationPlan::FwhtG256, false, true), + RotationVariant::WithSwiGLU + ); + assert_eq!( + select_rotation_variant(RotationPlan::FwhtG128, false, false), + RotationVariant::PlainG128 + ); + assert_eq!( + select_rotation_variant(RotationPlan::Givens, false, false), + RotationVariant::Givens + ); +} + +// ─── Escha-W2 registration (Task 4) ───────────────────────────── + +#[test] +fn escha_types_use_the_escha_rotation_plan() { + use hipfire_dispatch::types::{dtype_rotation_plan, RotationPlan}; + assert_eq!( + dtype_rotation_plan(DType::Escha2T16), + RotationPlan::EschaH128 + ); + assert_eq!( + dtype_rotation_plan(DType::Escha3T16), + RotationPlan::EschaH128 + ); +} + +/// Escha weights are stored in the rotated domain. Reaching a Plain GEMV +/// without the H128 pair does not crash — it produces coherent-looking +/// garbage. Both types must therefore refuse to resolve to Plain, exactly as +/// MQ4G128 does (see coverage_tests.rs). +#[test] +fn escha_types_never_resolve_to_plain() { + use hipfire_dispatch::types::{GemvVariant, KernelKey}; + for dt in [DType::Escha2T16, DType::Escha3T16] { + assert!( + KernelKey::for_gemv(dt, GemvVariant::Plain, false).is_err(), + "{dt:?} must not have a Plain GEMV arm — that would skip the H128 pair" + ); + } } diff --git a/crates/hipfire-dispatch-tests/src/qwen35.rs b/crates/hipfire-dispatch-tests/src/qwen35.rs index a94bad38fc..786d305f89 100644 --- a/crates/hipfire-dispatch-tests/src/qwen35.rs +++ b/crates/hipfire-dispatch-tests/src/qwen35.rs @@ -144,6 +144,7 @@ fn mq4_dtypes() -> MoeDtypes { has_paro_shared: false, per_expert_gate_up: None, per_expert_down: None, + routed_escha_transforms: false, } } diff --git a/crates/hipfire-dispatch/Cargo.toml b/crates/hipfire-dispatch/Cargo.toml index cc27bdc504..4ee4197e44 100644 --- a/crates/hipfire-dispatch/Cargo.toml +++ b/crates/hipfire-dispatch/Cargo.toml @@ -8,6 +8,10 @@ license.workspace = true hip-bridge = { path = "../hip-bridge" } hipfire-config = { path = "../hipfire-config" } rdna-compute = { path = "../rdna-compute" } +# Task 10 (Escha-W2): the routed combine must multiply by f16(score); `half` +# is the workspace's RNE f32->f16 encoder and is already a transitive dep via +# rdna-compute, so this adds no build cost. +half.workspace = true [features] default = [] diff --git a/crates/hipfire-dispatch/src/coverage_tests.rs b/crates/hipfire-dispatch/src/coverage_tests.rs index 4ff6b83e5f..e97ae04c78 100644 --- a/crates/hipfire-dispatch/src/coverage_tests.rs +++ b/crates/hipfire-dispatch/src/coverage_tests.rs @@ -439,6 +439,7 @@ fn non_k8_and_q8_routed_moe_has_a_dispatch_plan() { has_paro_shared: false, per_expert_gate_up: None, per_expert_down: None, + routed_escha_transforms: false, }; let res = MoeResolution::resolve(&d, u.k); // (a) These layers MUST take the fallback, not the k8 indexed path. @@ -499,6 +500,7 @@ fn moe_decode_pre_guard_admits_fallback_and_rejects_invalid() { has_paro_shared: false, per_expert_gate_up: None, per_expert_down: None, + routed_escha_transforms: false, }; let res_k4 = MoeResolution::resolve(&mq4_k4, 4); assert!( @@ -511,7 +513,9 @@ fn moe_decode_pre_guard_admits_fallback_and_rejects_invalid() { res_k4.use_gpu_topk, 4, /*n_exp=*/ 64, - /*resident=*/ true + /*resident=*/ true, + /*has_escha=*/ false, + /*escha_indexed_supported=*/ false ) .is_ok(), "MQ4G256 k=4 with resident experts is a VALID fallback case — guard must not reject it" @@ -529,27 +533,356 @@ fn moe_decode_pre_guard_admits_fallback_and_rejects_invalid() { "MQ4G256 k=8 must be GPU-top-K-indexable" ); assert!( - check_moe_decode_supported(res_k8.use_gpu_topk, 8, 64, /*resident=*/ false).is_ok(), + check_moe_decode_supported( + res_k8.use_gpu_topk, + 8, + 64, + /*resident=*/ false, + /*has_escha=*/ false, + /*escha_indexed_supported=*/ false + ) + .is_ok(), "GPU-top-K path is valid even under paged (non-resident) residency" ); // (a) out-of-range k errors gracefully (no panic): k == 0 and k > n_exp. assert!( - check_moe_decode_supported(false, 0, 64, true).is_err(), + check_moe_decode_supported(false, 0, 64, true, false, false).is_err(), "k == 0 must be rejected (would panic select_nth_unstable_by(k-1))" ); assert!( - check_moe_decode_supported(false, 65, 64, true).is_err(), + check_moe_decode_supported(false, 65, 64, true, false, false).is_err(), "k > n_exp must be rejected (would panic select_nth_unstable_by(k-1))" ); // (b) routed dtype on NEITHER path: not GPU-top-K AND no resident experts. assert!( - check_moe_decode_supported(/*use_gpu_topk=*/ false, 4, 64, /*resident=*/ false).is_err(), + check_moe_decode_supported( + /*use_gpu_topk=*/ false, 4, 64, /*resident=*/ false, + /*has_escha=*/ false, /*escha_indexed_supported=*/ false + ) + .is_err(), "non-fast-path dtype with no resident experts has no runnable path — reject gracefully" ); } +/// LAYER 1f — Escha-W2 on the indexed GPU-top-K path: SUPPORTED only in the +/// one shape that keeps the H128 pair, and failing closed in every other. +/// +/// Escha-W2 weights live in a rotated domain. Only the executors in +/// `pipeline::escha` wrap the routed GEMVs in the H128 pair; an escha layer +/// that reached the GENERIC indexed routed body would skip both Hadamard +/// transforms and emit finite, fluent output wrong by ~1e-1. Nothing would +/// fire — not a NaN check, not a shape check, not a dtype check. +/// +/// This test used to assert that escha could never be on the indexed path at +/// all, because the only thing keeping it off was a negative accident (the +/// loader materialises the experts as `Q8_0`, and no `routed_indexable_*` arm +/// admitted `Q8_0`). That accident has been replaced by real support: +/// `routed_indexable_escha_q8` admits Q8_0-on-both-projections experts that +/// carry the transform tables, and `run_moe_decode` routes them to +/// `escha::escha_routed_decode_indexed`, which applies the same H128 pair as +/// the CPU-top-K executor. +/// +/// So the contract the guard now encodes is narrower and stronger — the +/// SPECIFIC combination is supported and everything adjacent to it is not: +/// 1. escha + indexed + `escha_indexed_supported` → ADMITTED (production). +/// 2. escha + indexed + tables MISSING → REJECTED (the executor could not +/// have been called; the generic body would run and drop the transforms). +/// 3. escha + indexed via some OTHER arm (non-Q8_0 routed dtype) → REJECTED. +/// The escha indexed GEMV hard-codes the Q8_0 block layout, so this is +/// the mirror-image silent corruption of case 2. +/// 4. escha + `use_gpu_topk == false` → admitted (the CPU-top-K route). +/// 5. non-escha + indexed → admitted (every other MoE model). +/// +/// Cases 2 and 3 are what go red if arm (c) is weakened to `true`; case 1 is +/// what goes red if it is left as the old blanket refusal. +#[test] +fn escha_layer_must_not_take_the_indexed_gpu_topk_path() { + use crate::pipeline::check_moe_decode_supported; + + let base = MoeDtypes { + router: Q8_0, + shared_gate: Q8_0, + shared_expert_gate: Q8_0, + shared_expert_up: Q8_0, + shared_expert_down: Q8_0, + experts_all_gate_up_mq4: false, + routed_gate_up: Q8_0, + routed_down: Q8_0, + routed_has_mixed_experts: false, + has_paro_shared: false, + per_expert_gate_up: None, + per_expert_down: None, + routed_escha_transforms: false, + }; + + // ── 1. Production shape: Q8_0 experts + resident transform tables ────── + // `MoeResolution` is asked for the real answer rather than hand-set, so + // this stays honest if the arms are reworked. + let escha_q8 = MoeDtypes { + routed_escha_transforms: true, + per_expert_gate_up: None, + per_expert_down: None, + ..base + }; + let res_escha = MoeResolution::resolve(&escha_q8, 8); + assert!( + res_escha.routed_indexable_escha_q8, + "Q8_0-on-both-projections experts with resident H128 tables ARE the escha indexed arm" + ); + assert!( + res_escha.use_gpu_topk, + "the escha arm must admit the layer to GPU-resident top-K — that is the whole point of \ + it (no per-layer topk D2H, no per-expert GEMV launch storm, hipGraph-capturable)" + ); + assert!( + check_moe_decode_supported( + res_escha.use_gpu_topk, + 8, + /*n_exp=*/ 64, + /*resident=*/ true, + /*has_escha=*/ true, + /*escha_indexed_supported=*/ true, + ) + .is_ok(), + "escha on the indexed path WITH the escha executor behind it is the supported \ + production shape — the guard must admit it. If this is red the guard is still the old \ + blanket refusal and escha is stuck on the CPU-top-K route." + ); + + // ── 2. Tables missing: the executor could not run, so refuse ─────────── + let err = check_moe_decode_supported( + res_escha.use_gpu_topk, + 8, + 64, + /*resident=*/ true, + /*has_escha=*/ true, + /*escha_indexed_supported=*/ false, + ) + .expect_err( + "an escha layer on the indexed path WITHOUT its transform tables MUST be refused: the \ + escha executor cannot be called, so the generic indexed body would run the experts \ + with no H128 pair and emit finite, fluent, ~1e-1-wrong output with nothing to catch \ + it.", + ); + match err { + DispatchError::UnsupportedVariant { + family, variant, .. + } => { + assert_eq!(family, "moe"); + assert_eq!( + variant, "escha-routed-experts-on-indexed-gpu-topk-path", + "the refusal must name escha and the unsupported path — a generic error here \ + means the escha arm did not fire and something else rejected the case" + ); + } + other => panic!("expected UnsupportedVariant, got {other:?}"), + } + + // ── 3. Indexable through a DIFFERENT arm while still escha ───────────── + // A future graded/mixed escha file whose representative routed dtype is + // not Q8_0 would resolve indexable via the MQ4 arm. The escha indexed + // GEMV hard-codes the Q8_0 34 B/32-element block layout, so dispatching + // it there is silent corruption; so is letting the generic body have it. + let escha_via_mq4 = MoeDtypes { + experts_all_gate_up_mq4: true, + routed_gate_up: MQ4G256, + routed_down: MQ4G256, + routed_escha_transforms: true, + per_expert_gate_up: None, + per_expert_down: None, + ..base + }; + let res_mq4 = MoeResolution::resolve(&escha_via_mq4, 8); + assert!( + res_mq4.use_gpu_topk && !res_mq4.routed_indexable_escha_q8, + "fixture precondition: indexable, but NOT through the escha arm" + ); + assert!( + check_moe_decode_supported( + res_mq4.use_gpu_topk, + 8, + 64, + true, + /*has_escha=*/ true, + // What `run_moe_decode` computes: the escha arm did not fire, so + // the escha executor is not the one that would run. + /*escha_indexed_supported=*/ + res_mq4.routed_indexable_escha_q8, + ) + .is_err(), + "an escha layer that reached the indexed path through a NON-escha arm must be refused — \ + the escha executor only decodes Q8_0, and the generic body applies no transforms" + ); + + // ── 4. CPU-top-K route stays admitted ────────────────────────────────── + let escha_k4 = MoeDtypes { + routed_escha_transforms: true, + per_expert_gate_up: None, + per_expert_down: None, + ..base + }; + let res_k4 = MoeResolution::resolve(&escha_k4, 4); + assert!( + !res_k4.use_gpu_topk, + "k != 8 has no indexed kernel on any arm, escha included" + ); + assert!( + check_moe_decode_supported(res_k4.use_gpu_topk, 4, 64, true, true, false).is_ok(), + "escha on the CPU-top-K fallback is still a supported shape — the guard must not \ + reject it" + ); + + // ── 5. Non-escha models on the indexed path are untouched ────────────── + assert!( + check_moe_decode_supported( + res_mq4.use_gpu_topk, + 8, + 64, + true, + /*has_escha=*/ false, + /*escha_indexed_supported=*/ false + ) + .is_ok(), + "the escha guard must not affect any non-escha MoE model" + ); +} + +/// LAYER 1g — the BATCHED-PREFILL twin of 1f: an escha layer must never fall +/// into a prefill path that applies no Hadamard transform. +/// +/// `run_moe_prefill`'s escha branch keys on `MoePrefillParams::escha`, which is +/// gated on `escha_indexed_route_enabled()`. `None` there falls through into +/// Path 1 / Path 2, which feed the activation straight into the expert GEMMs +/// and combine the raw result — no H128 pair, no error, ~1e-1-wrong output. +/// +/// That is unreachable TODAY only because no admission arm outside escha's own +/// admits Q8_0 routed experts to batched prefill: a property of a dtype table +/// in `families::moe`, not of the branch that depends on it. The next planned +/// change (a Q8_0 grouped GEMM over sorted expert groups) adds exactly such an +/// arm. `check_moe_prefill_supported` therefore states the requirement at the +/// point of danger, keyed on the UNGATED `layer_is_escha` marker. +/// +/// The cases, and what each one catches: +/// 1. escha layer + tables present → ADMITTED (production; the escha +/// branch runs and returns before Path 1 / Path 2 are reached). +/// 2. escha layer + tables ABSENT → REJECTED. This is the whole point: +/// `HIPFIRE_ESCHA_INDEXED=0` produces exactly this state, and so would a +/// future generic Q8_0 prefill arm. Goes GREEN-to-RED if the guard is +/// removed. +/// 3. non-escha + tables absent → ADMITTED (every other MoE model — +/// the guard must be a no-op for them). +/// 4. the marker must be UNGATED → asserted structurally below. +#[test] +fn escha_layer_must_not_take_a_non_escha_prefill_path() { + use crate::pipeline::check_moe_prefill_supported; + + // ── 1. Production shape ──────────────────────────────────────────────── + assert!( + check_moe_prefill_supported(/*layer_is_escha=*/ true, /*escha_tables_present=*/ true) + .is_ok(), + "an escha layer WITH its transform tables is the supported batched-prefill shape — the \ + escha branch in run_moe_prefill runs it and returns. If this is red the guard is \ + refusing production." + ); + + // ── 2. The silent-wrong-output case ──────────────────────────────────── + let err = check_moe_prefill_supported( + /*layer_is_escha=*/ true, /*escha_tables_present=*/ false, + ) + .expect_err( + "an escha layer that reaches run_moe_prefill WITHOUT its transform tables MUST be \ + refused. Path 1 and Path 2 apply NO Hadamard transform and raise no error, so this \ + is a rotated-domain weight multiplied by an unrotated activation: finite, fluent, \ + ~1e-1-wrong output with nothing to catch it. This is reachable today by setting \ + HIPFIRE_ESCHA_INDEXED=0, and will be reachable by default the moment a generic Q8_0 \ + routed prefill arm exists.", + ); + match err { + DispatchError::UnsupportedVariant { + family, variant, .. + } => { + assert_eq!(family, "moe"); + assert_eq!( + variant, "escha-routed-experts-on-non-escha-prefill-path", + "the refusal must name escha and the unsupported path" + ); + } + other => panic!("expected UnsupportedVariant, got {other:?}"), + } + + // ── 3. Every non-escha MoE model is untouched ────────────────────────── + assert!( + check_moe_prefill_supported(false, false).is_ok(), + "the guard must be a no-op for non-escha models — they legitimately have no tables" + ); + assert!( + check_moe_prefill_supported(false, true).is_ok(), + "tables without the escha marker is not a state this guard has an opinion about" + ); + + // ── 4. The marker must be UNGATED ────────────────────────────────────── + // + // `escha_tables_present` is `MoePrefillParams::escha.is_some()`, which the + // model ANDs with `escha_indexed_route_enabled()`. If `layer_is_escha` + // were built the same way, case 2 could never arise — the guard would be + // structurally dead, always seeing `(false, false)`. Stated as an + // assertion over the pair rather than left as prose: the ONLY pair the + // guard rejects is the one a gated marker cannot produce. + let rejected: Vec<(bool, bool)> = [(false, false), (false, true), (true, false), (true, true)] + .into_iter() + .filter(|&(l, t)| check_moe_prefill_supported(l, t).is_err()) + .collect(); + assert_eq!( + rejected, + vec![(true, false)], + "exactly one input pair may be refused: escha layer, no tables. A gated `layer_is_escha` \ + would never produce it, which is why the field is documented as UNGATED." + ); +} + +/// LAYER 1f (companion) — the Q8_0 indexed arm must stay ESCHA-SCOPED. +/// +/// `routed_indexable_escha_q8` is the first arm in `MoeResolution` that +/// admits `Q8_0` routed experts, and `Q8_0` is the most common dtype in this +/// codebase. If the escha gate on it were ever dropped, every plain-Q8_0 MoE +/// model in the tree would silently move from the CPU-top-K fallback onto an +/// indexed path — either the generic body (whose kernels expect a different +/// container for the projections that arm implies) or, worse, the escha +/// executor, which would apply a Hadamard pair those weights were never +/// packed in. Both are fluent wrong output, not a fault. +#[test] +fn plain_q8_0_routed_experts_stay_off_the_indexed_path() { + let plain_q8 = MoeDtypes { + router: Q8_0, + shared_gate: Q8_0, + shared_expert_gate: Q8_0, + shared_expert_up: Q8_0, + shared_expert_down: Q8_0, + experts_all_gate_up_mq4: false, + routed_gate_up: Q8_0, + routed_down: Q8_0, + routed_has_mixed_experts: false, + has_paro_shared: false, + per_expert_gate_up: None, + per_expert_down: None, + // The ONLY difference from the escha production fixture above. + routed_escha_transforms: false, + }; + let res = MoeResolution::resolve(&plain_q8, 8); + assert!( + !res.routed_indexable_escha_q8, + "the Q8_0 indexed arm must require the escha transform tables" + ); + assert!( + !res.use_gpu_topk, + "a plain Q8_0 MoE model must keep taking the CPU-top-K fallback. If this is red, the \ + escha gate came off the Q8_0 arm and every Q8_0 MoE model in the tree just changed \ + execution path (and, on the escha executor, answer)." + ); +} + /// LAYER 1c — Q8/Paro were gapped in MULTIPLE GEMV variants: o_proj used Residual, /// then the FFN/qkv used Prerotated (the second panic domino). Lock every variant /// these dtypes are actually dispatched through. @@ -1376,6 +1709,7 @@ fn moe_dtypes_uniform(gate_up: DType, down: DType) -> MoeDtypes { has_paro_shared: false, per_expert_gate_up: None, per_expert_down: None, + routed_escha_transforms: false, } } diff --git a/crates/hipfire-dispatch/src/families/gemv.rs b/crates/hipfire-dispatch/src/families/gemv.rs index a6fc86ea0e..9f086e50e1 100644 --- a/crates/hipfire-dispatch/src/families/gemv.rs +++ b/crates/hipfire-dispatch/src/families/gemv.rs @@ -144,6 +144,14 @@ pub fn select_rotation_variant( RotationVariant::Plain } } + // Escha-W2: stated explicitly rather than left to the FwhtG256 `_` + // fallthrough below, even though no RotationVariant here is ever + // actually launched for this plan — `prepare_rotation_scratch` + // errors on `RotationPlan::EschaH128` before `rotate()` calls + // `self.rotation.run()`, so whatever is returned here is inert. + // Do NOT read this arm as "Escha shares the FwhtG256 fusion axis"; + // it does not — it has its own H128 transform with no kernel yet. + RotationPlan::EschaH128 => RotationVariant::Plain, // FwhtG256 shares the fusion axis. _ => { if has_swiglu { @@ -452,6 +460,20 @@ fn prepare_rotation_scratch( arch: "", quant: "", }), + // Escha-W2's 128-point Hadamard has no scratch buffer / rotate kernel + // yet (Task 4 registers the dtype + rotation plan only). This is the + // guard that actually matters: it is what stops `rotate()` from ever + // reaching `self.rotation.run()` for these types, regardless of which + // RotationVariant `select_rotation_variant`'s catch-all picked for + // `EschaH128` above. Must become a real scratch-buffer allocation + // when the H128 rotate kernel lands — do not widen this to reuse the + // FwhtG256/FwhtG128 scratch, which is the wrong transform size. + RotationPlan::EschaH128 => Err(DispatchError::UnsupportedVariant { + family: "gemv", + variant: "escha-h128-rotate-unimplemented", + arch: "", + quant: "", + }), } } diff --git a/crates/hipfire-dispatch/src/families/moe.rs b/crates/hipfire-dispatch/src/families/moe.rs index b2882225e9..e060e6ce1d 100644 --- a/crates/hipfire-dispatch/src/families/moe.rs +++ b/crates/hipfire-dispatch/src/families/moe.rs @@ -88,6 +88,25 @@ pub struct MoeDtypes { pub per_expert_gate_up: Option>, /// Per-expert down tiers (parallel to `per_expert_gate_up`). Same semantics. pub per_expert_down: Option>, + /// This layer's routed experts came from an Escha-W2 checkpoint AND the + /// per-expert H128 transform tables are resident, so the escha routed + /// executor (`crate::pipeline::escha`) can run them. + /// + /// It is a separate flag rather than a dtype because the escha loader + /// decodes the trellis at load time and materialises the experts as + /// `Q8_0`: by the time dispatch sees the layer, NO routed dtype says + /// "escha" any more (that is exactly what `has_escha_experts` can no + /// longer see). The model sets it from `MoeFfnWeights.escha.is_some()` — + /// the same single source of truth that populates `MoeParams::escha`. + /// + /// It gates `routed_indexable_escha_q8`. Scoping the Q8_0 indexed arm to + /// escha, rather than admitting `Q8_0` routed experts in general, is + /// deliberate in BOTH directions: escha's indexed route is a different + /// executor (it wraps every GEMV in the H128 pair) that a plain-Q8_0 MoE + /// model must not be pulled onto, and a plain-Q8_0 MoE model must equally + /// not be pulled onto the generic indexed body on the strength of an arm + /// added for escha's benefit. + pub routed_escha_transforms: bool, } impl MoeDtypes { @@ -104,6 +123,30 @@ impl MoeDtypes { // the gfx1151 MQ4-i8 grouped fence via `force_mq4_grouped_fp16`. .any(|dt| matches!(*dt, DType::MQ6G256 | DType::MQ6G256V2)) } + + /// True iff this layer's routed experts carry an Escha-W2 dtype + /// (`Escha2T16` gate_up / `Escha3T16` down, today's only pairing — + /// checked on both `routed_*` and, defensively, any per-expert tier + /// table, in case a future graded-tier layer mixes escha in). Drives + /// the escha-only router-logits f16 round-trip in `run_moe_decode` + /// (see `kernels/src/router_logits_round_f16_rne.hip`): EschaLabs' + /// runtime rounds router logits to f16 before top-k, hipfire keeps them + /// F32 end-to-end everywhere else, and this flag scopes the rounding to + /// exactly the models that need to match Escha's f16 selection — + /// `qwen3.6:35b-a3b-*` and every other arch-6 SKU must stay bit-exact. + pub fn has_escha_experts(&self) -> bool { + let is_escha = |dt: DType| matches!(dt, DType::Escha2T16 | DType::Escha3T16); + is_escha(self.routed_gate_up) + || is_escha(self.routed_down) + || self + .per_expert_gate_up + .as_ref() + .is_some_and(|v| v.iter().copied().any(is_escha)) + || self + .per_expert_down + .as_ref() + .is_some_and(|v| v.iter().copied().any(is_escha)) + } } /// Resolved fused-vs-fallback eligibility for one MoE decode layer. This IS the @@ -180,6 +223,27 @@ pub struct MoeResolution { /// gate_up and down. Binds the same indexed MQ2-Lloyd kernels as qt19 but /// consumes x in the natural basis (`needs_x_rot_local == false`). pub routed_indexable_mq2lloyd_u: bool, + /// Escha-W2 routed experts (Q8_0 on both projections, H128 transform + /// tables resident). Admits the layer to GPU-resident top-K ONLY — the + /// routed body it reaches is `pipeline::escha::escha_routed_decode_indexed`, + /// never the generic indexed body, which has no escha awareness. See the + /// arm in `resolve_arch` and `pipeline::check_moe_decode_supported`. + pub routed_indexable_escha_q8: bool, + /// Escha-W2 routed experts stored as the TRELLIS CODE (`Escha2T16` / + /// `Escha3T16` on either projection, H128 transform tables resident) — the + /// Phase-2 production shape. Admits the layer to exactly the same place + /// [`Self::routed_indexable_escha_q8`] does: GPU-resident top-K, reaching + /// `pipeline::escha::escha_routed_decode_indexed` and never the generic + /// indexed body. + /// + /// It is a SEPARATE flag rather than a widened `routed_indexable_escha_q8` + /// because the two select different GEMVs — the code arm dispatches + /// `escha_gemv_native_*` with a trellis order taken from each projection's + /// own dtype, the Q8_0 arm dispatches the block-decode kernels. A single + /// flag would leave the executor guessing from the dtype anyway, and the + /// fail-closed guard would no longer be able to say WHICH shape it + /// admitted. + pub routed_indexable_escha_native: bool, pub use_gpu_topk: bool, pub needs_x_rot_local: bool, /// True when a per-expert tier table is `Some` AND contains >1 distinct @@ -291,6 +355,44 @@ impl MoeResolution { && routed_gate_up_e8 && matches!(d.routed_down, MFP4G32E8 | MFP3G32E8 | MFP2G32E8); + // Escha-W2. The routed experts are the Q8_0 the trellis decoded into, + // BOTH projections, and the layer carries the H128 transform tables. + // + // This arm does NOT admit the layer to the generic indexed routed body + // below — escha weights are in a rotated domain and that body would + // omit both Hadamard transforms, producing finite, fluent, ~1e-1-wrong + // output. What it admits is GPU-resident top-K: `run_moe_decode` + // branches to `pipeline::escha::escha_routed_decode_indexed` (which + // keeps the H128 pair) before the generic body, and + // `check_moe_decode_supported` refuses any escha layer that arrives on + // the indexed path WITHOUT those tables. + // + // Q8_0 is required on both sides for the same reason every other + // uniform arm requires it: the indexed GEMV decodes a 34 B/32-element + // block layout, and handing it a different container is silent + // corruption, not a fault. + let routed_indexable_escha_q8 = + d.routed_escha_transforms && d.routed_gate_up == Q8_0 && d.routed_down == Q8_0; + + // Escha-W2, Phase 2: the routed experts are the TRELLIS CODE itself and + // the fused GEMV decodes it in-register. Everything the Q8_0 arm above + // says applies unchanged — this admits GPU-resident top-K only, and the + // executor it reaches is still `escha_routed_decode_indexed` with its + // H128 pair, never the generic indexed body. + // + // Either escha dtype is accepted on either projection rather than + // hard-coding today's (K=2 gate_up, K=3 down) pairing: the trellis + // order is a per-projection property that the executor reads back off + // the SAME dtype to pick the kernel, so a file that allocated the bits + // the other way round is served correctly instead of being silently + // refused. What is NOT accepted is a mix with any other container — + // the fused kernel's bit geometry is the format's, and handing it + // anything else is silent corruption rather than a fault. + let is_escha_code = |dt: DType| matches!(dt, Escha2T16 | Escha3T16); + let routed_indexable_escha_native = d.routed_escha_transforms + && is_escha_code(d.routed_gate_up) + && is_escha_code(d.routed_down); + let routed_dtype_indexable = routed_indexable_mq4 || routed_indexable_mq4v2 || routed_indexable_mq5 @@ -303,7 +405,9 @@ impl MoeResolution { || routed_indexable_mq3lloyd || routed_indexable_mixed_lloyd || routed_indexable_paro - || routed_indexable_e8; + || routed_indexable_e8 + || routed_indexable_escha_q8 + || routed_indexable_escha_native; let use_gpu_topk = k == 8 && routed_dtype_indexable; let needs_x_rot_local = gate_side_mq4 @@ -361,6 +465,8 @@ impl MoeResolution { routed_indexable_mixed_lloyd, routed_indexable_mixed_per_expert, routed_indexable_paro, + routed_indexable_escha_q8, + routed_indexable_escha_native, use_gpu_topk, needs_x_rot_local, mixed, @@ -484,6 +590,32 @@ pub struct MoeParams<'a> { pub topk_indices: &'a GpuTensor, pub topk_weights: &'a GpuTensor, pub down_expanded: &'a GpuTensor, + + /// Escha-W2 (Task 10): per-layer H128 transform tables + the phase + /// scratch the batched routed executor needs. `Some` only for layers + /// loaded from an Escha-W2 checkpoint; `None` leaves every other model + /// byte-identical. + /// + /// This is ALSO the escha marker `MoeDtypes::has_escha_experts` can no + /// longer be: the loader decodes the trellis at load time and stores the + /// experts as `Q8_0`, so by the time dispatch sees the layer no routed + /// dtype says "escha" any more. See [`MoeParams::has_escha`]. + pub escha: Option>, +} + +impl MoeParams<'_> { + /// True iff this layer must take escha semantics — the f16 router-logit + /// round-trip and the H128-wrapped routed executor. + /// + /// Two sources, deliberately OR-ed: `dtypes.has_escha_experts()` still + /// catches a layer whose routed dtype is literally `Escha2T16`/`Escha3T16` + /// (what a future on-the-fly-decode GEMV would present), and + /// `escha.is_some()` catches today's shape, where the loader has already + /// materialised the experts as `Q8_0` and only the transform tables + /// remain as evidence. + pub fn has_escha(&self) -> bool { + self.dtypes.has_escha_experts() || self.escha.is_some() + } } // ── DeepSeek-V4 bias-aware decode parameters ─────────── @@ -809,6 +941,37 @@ pub struct MoePrefillParams<'a> { /// no all-reduce). `None` (the default) accumulates routed into `x_batch`, /// byte-identical to pre-EP behavior. pub routed_out: Option<&'a GpuTensor>, + /// Escha-W2 transform tables for this layer. `Some` iff the layer's routed + /// experts are escha-coded AND the indexed route is enabled — the same + /// marker `MoePrefillDtypes::escha` gates admission on, so a layer cannot + /// be admitted to batched prefill and then arrive here without the tables + /// the executor needs. `None` for every other model: the escha branch in + /// [`crate::pipeline::run_moe_prefill`] is skipped entirely and Path 1 / + /// Path 2 run exactly as they do today. + /// + /// Only the FOUR `[E, ·]` transform tables are read — the `[k]`-sized + /// decode scratch fields of `EschaRoutedRefs` are ignored here, because + /// batched prefill uses the model-global `[n_tokens × k]` scratch from + /// `Gpu::ensure_escha_prefill_scratch` instead. + pub escha: Option>, + /// UNGATED marker: this layer's routed experts are escha-coded, full stop. + /// + /// Set from `ffn.escha.is_some()` alone — never ANDed with + /// `escha_indexed_route_enabled()` or any other env lever. That is the + /// whole point of the field: [`escha`](Self::escha) above IS gated, so + /// `escha.is_none()` cannot distinguish "not an escha layer" from "an + /// escha layer with the indexed route switched off", and the second of + /// those must never be allowed to run Path 1 / Path 2. + /// + /// Consumed by `pipeline::check_moe_prefill_supported`, which refuses + /// `layer_is_escha && escha.is_none()` before any GPU work. `false` for + /// every non-escha model, where the check is a no-op. + pub layer_is_escha: bool, + /// Model hidden size. Decode's `MoeParams` already carries this; prefill + /// did not need it until the escha branch, whose H128 transforms are sized + /// by it (`down_m` happens to equal it, but relying on that coincidence is + /// how a shape bug gets written). + pub hidden: usize, } /// Resolved dispatch plan for the qwen35 batched MoE prefill routed block. @@ -1053,6 +1216,7 @@ mod tests { has_paro_shared: false, per_expert_gate_up: None, per_expert_down: None, + routed_escha_transforms: false, } } diff --git a/crates/hipfire-dispatch/src/pipeline/escha.rs b/crates/hipfire-dispatch/src/pipeline/escha.rs new file mode 100644 index 0000000000..2814c367ce --- /dev/null +++ b/crates/hipfire-dispatch/src/pipeline/escha.rs @@ -0,0 +1,923 @@ +// SPDX-License-Identifier: MIT OR Apache-2.0 +// hipfire — see LICENSE and NOTICE in the project root. +//! Escha-W2 routed-expert decode executors (Tasks 10 and 11). +//! +//! Two entry points, one for each of the two routes into a routed MoE layer. +//! Both apply the SAME phase structure and the SAME H128 pair; they differ +//! only in where the routing lives. +//! +//! * [`escha_routed_decode`] — CPU-top-K route. `topk_ids` / `topk_weights` +//! arrive on the HOST, already downloaded by +//! [`super::run_moe_decode_cpu_fallback`]. Used when the layer is not +//! admitted to the indexed path (k != 8, F32 control experts in the G4 +//! gate, …). +//! * [`escha_routed_decode_indexed`] — GPU-top-K route, the production +//! decode/prefill path. Routing stays on the device end to end: the ids +//! feed the transforms and the GEMVs as a device buffer, the combine +//! weights are f16-rounded by a device kernel, and there is no D2H +//! anywhere in the layer. +//! +//! Everything around them — the router GEMV, the f16 logit rounding, top-k +//! selection, and the shared expert — is unchanged arch-6 code. +//! +//! # Why the indexed route exists (Task 11, measured) +//! +//! `rocprofv3 --kernel-trace` over a warm 40-layer decode token on gfx1151: +//! 22.0 ms of GPU-busy against a 29.7 ms wall (profiled; 26.0 ms unprofiled), +//! of which the routed experts were 6.5 ms of kernel time, ~2.1 ms of +//! launch-gap across their 640 per-expert GEMV launches, and ~2.1 ms of +//! copy/stall around the per-layer `topk_indices` + `topk_weights` D2H and the +//! ids/weights H2D that followed it. The host round trip alone did NOT +//! dominate — but the round trip and the launch storm together were the whole +//! addressable overhead, and the round trip is also what made the layer +//! non-capturable under hipGraph. This route removes both: 640 routed GEMV +//! launches per token become 80, and the per-layer sync disappears. +//! +//! # Why the routed loop needs replacing at all +//! +//! Escha's weights live in a ROTATED domain: the stored matrix is +//! `H·W·H` (up to the folded per-channel scales), so a matmul against it is +//! only the intended linear if the activation is Hadamard-transformed going +//! in and the result is Hadamard-transformed coming out. `escha_h128_in` +//! before the GEMV and `escha_h128_out` after are not a normalisation detail; +//! omitting them yields plausible-looking output that is wrong by ~1e-1 +//! rather than ~1e-4. +//! +//! # Why it is phase-structured rather than a per-expert loop +//! +//! Task 8 measured the H128 kernels LAUNCH-bound: an empty kernel at the same +//! grid/block costs 1.74–1.78 us against a real launch's 2.4 us, and +//! overhead-subtracted time is nearly flat (0.59 → 0.69 us) from 16 to 136 +//! blocks. A per-expert wiring costs `40 layers × 8 experts × 4 transforms = +//! 1280` launches/token = 3.07 ms = a **326 tok/s ceiling from the transforms +//! alone**, before any GEMV work. Running the token's `k` experts in phases — +//! all inputs transformed, then all GEMVs, then all outputs transformed — is +//! **4 H128-family launches per layer, 160 per token, 0.38 ms**. +//! +//! Phase order (per layer, decode, one token). The `host` column is +//! [`escha_routed_decode`], `idx` is [`escha_routed_decode_indexed`]: +//! +//! | # | host | idx | what | +//! |---|------|-----|------| +//! | 0 | 0 (host-side rounding) | 1 | `escha_round_weights_f16_rne` | +//! | 1 | 1 | 1 | `escha_h128_in_batched` — gate_up input side, x broadcast | +//! | 2 | k | 1 | Q8_0 gate_up GEMV (per expert / all experts indexed) | +//! | 3 | 1 | 1 | `escha_h128_out_batched` — gate_up output side | +//! | 4 | 1 | 1 | `escha_swiglu_batched` | +//! | 5 | 1 | 1 | `escha_h128_in_batched` — down input side, per-slot x | +//! | 6 | k | 1 | Q8_0 down GEMV (per expert / all experts indexed) | +//! | 7 | 1 | 1 | `escha_h128_out_batched` — down output side | +//! | 8 | 1 | 1 | `moe_down_combine_k8_batched` with f16-rounded scores | +//! +//! At k=8 that is 22 launches + 2 H2D + (2 D2H upstream) per layer on the +//! host route against 9 launches and no host transfer at all on the indexed +//! one. The H128 budget — 4 per layer — is identical on both, which is why +//! `escha_launches_per_token` and the gates that read it are route-agnostic. +//! +//! The per-expert `rin_eff` / `rout_eff` rows are an extra index into the +//! already-resident `[E, IC]` / `[E, OC]` tensors — the batching is a kernel +//! indexing change plus a grid change, not new maths, and it is gated +//! bit-exactly against `escha_ref` by +//! `rdna-compute/examples/test_escha_h128_gpu_vs_cpu.rs`. + +use rdna_compute::{EschaXGroup, Gpu, GpuTensor}; + +use crate::context::DispatchCtx; +use crate::families::gemv::{GemvFamily, WeightRef}; +use crate::types::DispatchError; + +/// Borrowed view of one layer's Escha-W2 transform tables plus the decode +/// scratch the phase structure needs. Built by the model +/// (`hipfire_arch_qwen35::qwen35::escha::EschaMoeTables::refs`); this crate +/// never owns or allocates any of it. +pub struct EschaRoutedRefs<'a> { + /// `[n_exp, hidden]` f32 — folded `rin` for the gate_up projection. + pub gate_up_rin: &'a GpuTensor, + /// `[n_exp, 2*mi]` f32 — folded `rout` for gate_up. Carries the per-expert + /// prune mask (zeros); see the zero contract in `escha_h128.hip`. + pub gate_up_rout: &'a GpuTensor, + /// `[n_exp, mi]` f32 — folded `rin` for the down projection. + pub down_rin: &'a GpuTensor, + /// `[n_exp, hidden]` f32 — folded `rout` for down. + pub down_rout: &'a GpuTensor, + /// `[k]` i32 — this token's selected expert ids (device). + pub ids: &'a GpuTensor, + /// `[k]` f32 — this token's combine weights, ALREADY f16-rounded. + pub weights: &'a GpuTensor, + /// `[k, hidden]` f32 scratch. + pub xh_gu: &'a GpuTensor, + /// `[k, 2*mi]` f32 scratch. + pub mid_gu: &'a GpuTensor, + /// `[k, 2*mi]` f32 scratch. + pub y_gu: &'a GpuTensor, + /// `[k, mi]` f32 scratch. + pub h: &'a GpuTensor, + /// `[k, mi]` f32 scratch. + pub xh_dn: &'a GpuTensor, + /// `[k, hidden]` f32 scratch. + pub mid_dn: &'a GpuTensor, + /// `[k, hidden]` f32 scratch — the per-slot expert outputs the combine + /// reduces. + pub y_dn: &'a GpuTensor, +} + +/// Number of H128-family launches this executor issues per call. Pinned as a +/// constant so the launch budget in the module docs is a checked claim rather +/// than a comment. +/// +/// It is checked by the GATES, not by this function: `escha_routed_decode` +/// itself contains no assert (an earlier version of this doc claimed a +/// debug-build assert that never existed). The claim is enforced by reading +/// `rdna_compute::escha_h128_launches()` across a real forward pass — +/// `hipfire-arch-qwen35/examples/escha_moe_block_gate.rs` (one layer, per +/// (layer, token)) and `examples/escha_model_smoke.rs` (whole model, decode +/// AND prefill, against `escha_launches_per_token`). +/// +/// This is an H128-TRANSFORM budget, not the layer's whole launch cost: a +/// decode layer at k=8 issues 22 launches in total (4 H128 + 1 SwiGLU + +/// 1 combine + 16 GEMV). See the phase table in the module docs. +pub const ESCHA_H128_LAUNCHES_PER_LAYER: usize = 4; + +/// H128 launches for a whole decode step. Independent of `k` — that is the +/// entire point of batching across experts. +pub fn escha_launches_per_token(n_layers: usize) -> usize { + n_layers * ESCHA_H128_LAUNCHES_PER_LAYER +} + +/// Run one routed GEMV phase against whichever container this layer's expert +/// slots hold. +/// +/// Both arms compute the SAME sum in the SAME order — the fused native kernels +/// are transcriptions of the Q8_0 ones with only the weight's provenance +/// changed (see `kernels/src/escha_moe_gemv_native.hip`). They do not produce +/// the same NUMBERS, because Q8_0 is a lossy re-quantisation of the weights +/// the code decodes to; the native arm is the weight-exact one. +/// +/// The `_ =>` arm is load-bearing and must stay an error. `expert_ptrs` are +/// raw device addresses with no length or type attached, so dispatching the +/// wrong kernel at one of them reads a different byte geometry out of the same +/// bytes: finite, plausible, wrong. `MoeResolution` should already have +/// refused anything that lands here, which is exactly why this is the place to +/// notice that it did not. +fn escha_routed_gemv( + gpu: &mut Gpu, + dtype: rdna_compute::DType, + expert_ptrs: &GpuTensor, + ids: &GpuTensor, + x: &GpuTensor, + y: &GpuTensor, + m: usize, + k: usize, + slots: usize, +) -> Result<(), DispatchError> { + use rdna_compute::DType; + let r = match dtype { + DType::Escha2T16 => gpu.escha_gemv_native_moe_k8_indexed_batched( + expert_ptrs, + ids, + x, + y, + m, + k, + slots, + 2, + false, + ), + DType::Escha3T16 => gpu.escha_gemv_native_moe_k8_indexed_batched( + expert_ptrs, + ids, + x, + y, + m, + k, + slots, + 3, + false, + ), + DType::Q8_0 => { + gpu.escha_gemv_q8_0_moe_k8_indexed_batched(expert_ptrs, ids, x, y, m, k, slots) + } + _ => { + return Err(DispatchError::UnsupportedVariant { + family: "moe", + variant: "escha-routed-expert-container-not-code-or-q8_0", + arch: "", + quant: "", + }) + } + }; + r.map_err(|err| DispatchError::Hip(err.to_string())) +} + +/// Expert-grouped routed GEMM — the prefill form of [`escha_routed_gemv`]. +/// +/// The slot-parallel GEMV reads an expert's weights once per (token, expert) +/// SLOT, so prefill does not amortise across the batch at all. Measured on the +/// shipped 35B before this existed: 4.832 / 4.525 / 4.440 ms/token at +/// n = 128 / 512 / 2048 — flat across a 16x batch increase. At n=512 that is +/// 512 tokens x 8 slots x ~0.92 MB x 40 layers = ~150 GB of routed weight +/// traffic per batch. Loading each expert's code once per (layer, batch) +/// instead is ~9.4 GB. +/// +/// The grouping is the pre-existing SGLang-style scatter pipeline (histogram -> +/// padded exclusive scan -> permute); only the GEMM is escha-specific, and it +/// decodes the trellis code once per expert rather than once per slot. +/// +/// NOT bit-identical to the slot-parallel route: a group sums over a different +/// partition of the contraction, so accumulation order differs. The +/// microbenchmark measures K=2 as bit-identical anyway and K=3 within +/// 1.5e-5 max / 1.5e-6 mean — far inside the Q8_0 arm's 2.633e-4 / 3.027e-5. +/// +/// Falls back to the slot-parallel GEMV for any non-trellis container (there is +/// no grouped Q8_0 escha kernel) and for slot counts too small to pay for the +/// grouping launches — which is also what keeps decode (slots == k) off it. +#[allow(clippy::too_many_arguments)] +fn escha_routed_gemm_grouped( + gpu: &mut Gpu, + dtype: rdna_compute::DType, + expert_ptrs: &GpuTensor, + ids: &GpuTensor, + x: &GpuTensor, + y: &GpuTensor, + m: usize, + k: usize, + slots: usize, + n_exp: usize, +) -> Result<(), DispatchError> { + use rdna_compute::DType; + let trellis_k = match dtype { + DType::Escha2T16 => 2u32, + DType::Escha3T16 => 3u32, + _ => return escha_routed_gemv(gpu, dtype, expert_ptrs, ids, x, y, m, k, slots), + }; + + let (rows, _ctiles) = rdna_compute::escha_grouped_tile(m); + // Grouping pays as soon as an expert has more than about one slot to + // amortise its code read over — that is independent of the tile HEIGHT. + // Keying on `n_exp * rows` was a bug: it silently disabled grouping for + // any tile with rows >= 16 at the 256-token prefill chunk (2048 slots), + // which measured as a clean fall back to the un-grouped 223 tok/s. + // + // Decode has `slots == k` (8), far below `2 * n_exp`, so this is also what + // keeps the decode route on the slot-parallel GEMV. + let _ = rows; + // HIPFIRE_ESCHA_GROUP_ALWAYS forces the grouped/WMMA route at decode too. + // Measured negative result, kept because the reasoning is not obvious: + // decode 47.7 -> 47.6 tok/s, i.e. nothing. In PREFILL each decoded tile + // feeds ~8 slots, so FMAs dominate and WMMA wins 2.2x. At decode there is + // ONE slot per expert, so the same decoded tile feeds a single MAC — the + // trellis decode dominates and there are no FMAs worth moving to the + // matrix cores. Same kernel, opposite bottleneck, purely from + // slots-per-expert. + let force = std::env::var("HIPFIRE_ESCHA_GROUP_ALWAYS").is_ok(); + if !force && slots < 2 * n_exp { + return escha_routed_gemv(gpu, dtype, expert_ptrs, ids, x, y, m, k, slots); + } + + macro_rules! hip { + ($ex:expr) => { + $ex.map_err(|err| DispatchError::Hip(err.to_string())) + }; + } + + // Padding rounds every expert's bucket up to `rows`, so the permuted array + // is at most `slots + n_exp * (rows - 1)`; padded entries carry the -1 + // sentinel the GEMM skips. + let m_total_max = slots + n_exp * rows; + let counts = hip!(gpu.alloc_tensor(&[n_exp], DType::F32))?; + let offsets = hip!(gpu.alloc_tensor(&[n_exp + 1], DType::F32))?; + let sorted = hip!(gpu.alloc_tensor(&[m_total_max], DType::F32))?; + let tile_ids = hip!(gpu.alloc_tensor(&[m_total_max / rows + 1], DType::F32))?; + let inverse = hip!(gpu.alloc_tensor(&[slots], DType::F32))?; + + hip!(gpu.moe_scatter_histogram_k8(ids, &counts, slots, n_exp))?; + hip!(gpu.moe_scatter_offsets_k8(&counts, &offsets, n_exp, rows))?; + hip!(gpu.moe_scatter_permute_k8( + ids, + &offsets, + &sorted, + &tile_ids, + &inverse, + slots, + n_exp, + m_total_max, + rows, + ))?; + + // WMMA by default. Measured against the scalar grouped kernel on the + // shipped projection shapes: 2.18x (K=2) and 2.49x (K=3), i.e. 5.6-6.2x + // over the slot-parallel GEMV. + // + // The first WMMA attempt measured only 1.02x, and hardware counters said + // why: it staged the activation fragment through LDS, which cost 0.52 G + // LDS instructions against the scalar kernel's 0.05 G and cancelled the + // 2.2x VALU saving cycle-for-cycle (busy 1.48 G vs 1.47 G). Reading B + // straight from global — each lane wants a different slot's contiguous + // 16-float run, so there is nothing to share — is what unlocked it. + if std::env::var("HIPFIRE_ESCHA_GROUPED_SCALAR").is_ok() { + return hip!(gpu.escha_gemm_native_moe_grouped( + expert_ptrs, + &offsets, + &sorted, + x, + y, + m, + k, + slots, + n_exp, + trellis_k, + false + )); + } + hip!(gpu.escha_gemm_native_moe_grouped_wmma( + expert_ptrs, + &offsets, + &sorted, + x, + y, + m, + k, + slots, + n_exp, + trellis_k, + false + )) +} + +/// SAFETY: `src` is a device buffer of at least `offset_elems + len_elems` +/// f32; the returned view is non-owning and must not outlive `src`. +unsafe fn view(src: &GpuTensor, offset_elems: usize, len_elems: usize) -> GpuTensor { + let ptr = (src.buf.as_ptr() as *mut u8).add(offset_elems * 4); + GpuTensor { + buf: hip_bridge::DeviceBuffer::from_raw(ptr as *mut _, len_elems * 4), + shape: vec![len_elems], + dtype: rdna_compute::DType::F32, + } +} + +/// Run the routed half of one Escha-W2 MoE layer for one token. +/// +/// `topk_ids` / `topk_weights` are host-side and already selected+renormalised +/// by the caller (production: `run_moe_decode_cpu_fallback`; the G4 gate +/// injects EschaLabs' shipped fixture instead, which is why this boundary is +/// public). `out` is accumulated into, never overwritten. +/// +/// The combine multiplies by **`f16(score)`** — one of the three load-bearing +/// rounding points of the format. It is applied here, on the host copy, so the +/// caller's `topk_weights` device buffer is left untouched for any other +/// consumer (e.g. `capture_expert_stats`). +#[allow(clippy::too_many_arguments)] +pub fn escha_routed_decode( + ctx: &DispatchCtx, + gpu: &mut Gpu, + e: &EschaRoutedRefs<'_>, + routed_experts: &[(WeightRef<'_>, WeightRef<'_>)], + topk_ids: &[usize], + topk_weights: &[f32], + x_norm: &GpuTensor, + out: &GpuTensor, + hidden: usize, + mi: usize, +) -> Result<(), DispatchError> { + macro_rules! hip { + ($ex:expr) => { + $ex.map_err(|err| DispatchError::Hip(err.to_string())) + }; + } + let k = topk_ids.len(); + // `moe_down_combine_k8_batched` unrolls to a hard 8 slots (`k < K_TOP` + // guard inside a `for k in 0..8`), so it silently drops slots 8.. rather + // than failing. Every escha SKU is k=8; reject anything else loudly here + // instead of returning a quietly truncated sum. + if k > 8 { + return Err(DispatchError::UnsupportedVariant { + family: "moe", + variant: "escha-routed-decode-supports-k<=8", + arch: "", + quant: "", + }); + } + if k == 0 || k != topk_weights.len() { + return Err(DispatchError::UnsupportedVariant { + family: "moe", + variant: "escha-topk-id-weight-length-mismatch", + arch: "", + quant: "", + }); + } + for &id in topk_ids { + if id >= routed_experts.len() { + return Err(DispatchError::UnsupportedVariant { + family: "moe", + variant: "escha-topk-id-out-of-range", + arch: "", + quant: "", + }); + } + } + + // ids + f16-rounded combine weights -> device. + let ids_i32: Vec = topk_ids.iter().map(|&i| i as i32).collect(); + let id_bytes: Vec = ids_i32.iter().flat_map(|v| v.to_le_bytes()).collect(); + hip!(gpu.hip.memcpy_htod(&e.ids.buf, &id_bytes))?; + let w_bytes: Vec = topk_weights + .iter() + .map(|&w| f32::from(half::f16::from_f32(w))) + .flat_map(|w| w.to_le_bytes()) + .collect(); + hip!(gpu.hip.memcpy_htod(&e.weights.buf, &w_bytes))?; + + // ── 1. gate_up input transform, all k slots, ONE launch ─────────────── + hip!(gpu.escha_h128_batched( + "escha_h128_in_batched", + x_norm, + e.gate_up_rin, + e.ids, + e.xh_gu, + hidden, + k, + EschaXGroup::Broadcast, + ))?; + + // ── 2. gate_up GEMV per selected expert ─────────────────────────────── + static GEMV: std::sync::OnceLock = std::sync::OnceLock::new(); + let gemv = GEMV.get_or_init(GemvFamily::new); + for (s, &id) in topk_ids.iter().enumerate() { + let x = unsafe { view(e.xh_gu, s * hidden, hidden) }; + let y = unsafe { view(e.mid_gu, s * 2 * mi, 2 * mi) }; + gemv.run_auto(ctx, gpu, &routed_experts[id].0, &x, &y)?; + } + + // ── 3. gate_up output transform, ONE launch ─────────────────────────── + hip!(gpu.escha_h128_batched( + "escha_h128_out_batched", + e.mid_gu, + e.gate_up_rout, + e.ids, + e.y_gu, + 2 * mi, + k, + EschaXGroup::PerSlot, + ))?; + + // ── 4. SwiGLU on the f16-rounded merged output, gate = FIRST half ───── + hip!(gpu.escha_swiglu_batched(e.y_gu, e.h, mi, k))?; + + // ── 5. down input transform, ONE launch (per-slot activation) ───────── + hip!(gpu.escha_h128_batched( + "escha_h128_in_batched", + e.h, + e.down_rin, + e.ids, + e.xh_dn, + mi, + k, + EschaXGroup::PerSlot, + ))?; + + // ── 6. down GEMV per selected expert ────────────────────────────────── + for (s, &id) in topk_ids.iter().enumerate() { + let x = unsafe { view(e.xh_dn, s * mi, mi) }; + let y = unsafe { view(e.mid_dn, s * hidden, hidden) }; + gemv.run_auto(ctx, gpu, &routed_experts[id].1, &x, &y)?; + } + + // ── 7. down output transform, ONE launch ────────────────────────────── + hip!(gpu.escha_h128_batched( + "escha_h128_out_batched", + e.mid_dn, + e.down_rout, + e.ids, + e.y_dn, + hidden, + k, + EschaXGroup::PerSlot, + ))?; + + // ── 8. weighted combine into the residual, ONE launch ───────────────── + hip!(gpu.moe_down_combine_k8_batched(e.y_dn, e.weights, out, hidden, k, 1))?; + Ok(()) +} + +/// Device-resident routing for [`escha_routed_decode_indexed`]. +/// +/// These are the buffers the GPU top-K kernel wrote. Nothing here is ever +/// read by the host: that is the point of the route. +pub struct EschaIndexedRouting<'a> { + /// `[n_exp]` u64 weight-base pointers for the Q8_0 gate_up slots, packed + /// into an F32 tensor (2 f32 per pointer) — the same table the other + /// indexed MoE GEMVs consume. + pub expert_gate_up_ptrs: &'a GpuTensor, + /// `[n_exp]` u64 weight-base pointers for the Q8_0 down slots. + pub expert_down_ptrs: &'a GpuTensor, + /// `[k]` selected expert ids as i32 (stored in an F32 tensor — same + /// 4 B/elem — exactly as `moe_topk_renorm_k8` and friends write them). + pub topk_indices: &'a GpuTensor, + /// `[k]` combine weights, NOT yet f16-rounded. Left untouched; the + /// rounded copy goes into the layer's own `weights` scratch. + pub topk_weights: &'a GpuTensor, + /// Container of the gate_up expert slots the pointers above address: + /// `Escha2T16` / `Escha3T16` (Phase 2 — the trellis code, decoded inside + /// the GEMV) or `Q8_0` (Phase 1 — decoded at load). This is read off the + /// layer's own routed dtype, the SAME fact + /// `MoeResolution::routed_indexable_escha_{native,q8}` admitted the layer + /// on, so the kernel choice and the admission decision cannot disagree. + /// Anything else is refused by [`escha_routed_gemv`] rather than + /// mis-decoded. + /// Number of routed experts in this layer. Needed by the grouped prefill + /// GEMM to size the expert-offset table and to decide whether grouping + /// pays for itself at this slot count. + pub n_experts: usize, + pub gate_up_dtype: rdna_compute::DType, + /// Container of the down expert slots. Independent of `gate_up_dtype`: the + /// shipped A3B file allocates K=2 to gate_up and K=3 to down, and a file + /// that allocated them the other way round is equally valid. + pub down_dtype: rdna_compute::DType, + /// Rows of the gate_up weight matrix (`2 * mi`). + pub gate_up_m: usize, + /// Columns of the gate_up weight matrix (`hidden`). + pub gate_up_k: usize, + /// Rows of the down weight matrix (`hidden`). + pub down_m: usize, + /// Columns of the down weight matrix (`mi`). + pub down_k: usize, +} + +/// Run the routed half of one Escha-W2 MoE layer for one token, with the +/// routing left on the device. +/// +/// Same eight phases as [`escha_routed_decode`] and the same H128 pair; the +/// only differences are that phases 2 and 6 are ONE indexed GEMV launch each +/// instead of `k`, and that the f16 rounding of the combine weights happens +/// in a kernel rather than on the host copy. +/// +/// `out` is accumulated into, never overwritten. +/// +/// # Why this is not just "the indexed path with escha bolted on" +/// +/// The generic indexed routed body in [`super::run_moe_decode`] is not +/// reachable for escha and must never become reachable: it feeds the raw +/// activation straight into the expert GEMVs and combines the raw result, so +/// it would omit both Hadamard transforms and emit finite, fluent output +/// wrong by ~1e-1. `run_moe_decode` therefore branches to THIS function +/// before that body, and `check_moe_decode_supported` refuses any escha layer +/// that reaches the indexed path without the transform tables that make this +/// function callable at all. +#[allow(clippy::too_many_arguments)] +pub fn escha_routed_decode_indexed( + gpu: &mut Gpu, + e: &EschaRoutedRefs<'_>, + r: &EschaIndexedRouting<'_>, + out: &GpuTensor, + x_norm: &GpuTensor, + hidden: usize, + mi: usize, + k: usize, +) -> Result<(), DispatchError> { + macro_rules! hip { + ($ex:expr) => { + $ex.map_err(|err| DispatchError::Hip(err.to_string())) + }; + } + // Same hard k<=8 bound as the host route: `moe_down_combine_k8_batched` + // unrolls to 8 slots and silently DROPS the rest rather than failing. + if k == 0 || k > 8 { + return Err(DispatchError::UnsupportedVariant { + family: "moe", + variant: "escha-routed-decode-supports-k<=8", + arch: "", + quant: "", + }); + } + // The GEMVs read `expert_ptrs[topk_indices[slot]]` on device, so a + // mis-sized weight table is an out-of-bounds READ (undefined behaviour), + // not a wrong answer. Both projections' shapes must also agree with the + // scratch the transforms were sized for — a mismatch there would have the + // GEMV write past the end of `mid_gu` / `mid_dn`. + if r.gate_up_m != 2 * mi || r.gate_up_k != hidden || r.down_m != hidden || r.down_k != mi { + return Err(DispatchError::UnsupportedVariant { + family: "moe", + variant: "escha-indexed-routed-shape-mismatch", + arch: "", + quant: "", + }); + } + + // ── 0. f16-round the combine weights, ONE launch, out-of-place ──────── + // `f16(score)` is one of the three load-bearing rounding points of the + // format. The host route does this on the downloaded copy; here it is a + // kernel, writing the layer's own scratch so `topk_weights` stays intact + // for any other consumer. + hip!(gpu.escha_round_weights_f16_rne(r.topk_weights, e.weights, k))?; + + // ── 1. gate_up input transform, all k slots, ONE launch ─────────────── + hip!(gpu.escha_h128_batched( + "escha_h128_in_batched", + x_norm, + e.gate_up_rin, + r.topk_indices, + e.xh_gu, + hidden, + k, + EschaXGroup::Broadcast, + ))?; + + // ── 2. gate_up GEMV for ALL k experts, ONE launch ───────────────────── + escha_routed_gemv( + gpu, + r.gate_up_dtype, + r.expert_gate_up_ptrs, + r.topk_indices, + e.xh_gu, + e.mid_gu, + r.gate_up_m, + r.gate_up_k, + k, + )?; + + // ── 3. gate_up output transform, ONE launch ─────────────────────────── + hip!(gpu.escha_h128_batched( + "escha_h128_out_batched", + e.mid_gu, + e.gate_up_rout, + r.topk_indices, + e.y_gu, + 2 * mi, + k, + EschaXGroup::PerSlot, + ))?; + + // ── 4. SwiGLU on the f16-rounded merged output, gate = FIRST half ───── + hip!(gpu.escha_swiglu_batched(e.y_gu, e.h, mi, k))?; + + // ── 5. down input transform, ONE launch (per-slot activation) ───────── + hip!(gpu.escha_h128_batched( + "escha_h128_in_batched", + e.h, + e.down_rin, + r.topk_indices, + e.xh_dn, + mi, + k, + EschaXGroup::PerSlot, + ))?; + + // ── 6. down GEMV for ALL k experts, ONE launch ──────────────────────── + escha_routed_gemv( + gpu, + r.down_dtype, + r.expert_down_ptrs, + r.topk_indices, + e.xh_dn, + e.mid_dn, + r.down_m, + r.down_k, + k, + )?; + + // ── 7. down output transform, ONE launch ────────────────────────────── + hip!(gpu.escha_h128_batched( + "escha_h128_out_batched", + e.mid_dn, + e.down_rout, + r.topk_indices, + e.y_dn, + hidden, + k, + EschaXGroup::PerSlot, + ))?; + + // ── 8. weighted combine into the residual, ONE launch ───────────────── + hip!(gpu.moe_down_combine_k8_batched(e.y_dn, e.weights, out, hidden, k, 1))?; + Ok(()) +} + +/// Run the routed half of one Escha-W2 MoE layer for a BATCH of `n_tokens`, +/// with the routing left on the device. +/// +/// The batched-prefill twin of [`escha_routed_decode_indexed`]. Same eight +/// phases, same H128 pair, same kernels — the ONLY change is that `slots` is +/// `n_tokens * k` instead of `k`. +/// +/// # Why that is enough, and why it is bit-identical per slot +/// +/// Every kernel in the pipeline is already purely slot-parallel: +/// +/// * `escha_gemv_q8_0_moe_k8_indexed_batched` (and its wide sibling) take +/// `krank = blockIdx.y`, read `x_batch + krank*K`, write `y_batch + krank*M` +/// and address the expert through `expert_ptrs[topk_indices[krank]]`. No +/// loop bound, accumulator, unroll or reduction depends on how many slots +/// there are, so slot `s` computes exactly the same sum in exactly the same +/// order whether the launch carried 8 slots or 2 048. +/// * `escha_h128_out_batched`, `escha_swiglu_batched` and +/// `moe_down_combine_k8_batched` are already `[slots, ...]` / +/// `[N, K_TOP, M]`-shaped and need nothing. +/// * `escha_h128_in_batched` needed the one change: its input-side activation +/// was either broadcast to every slot or one row per slot, and batched +/// prefill needs one row per TOKEN shared by that token's `k` slots. Hence +/// [`EschaXGroup::Grouped`]. +/// +/// That per-slot invariance is the reason the routed half of batched prefill +/// is asserted EQUAL to the per-token route rather than close to it — see the +/// gate in `hipfire-arch-qwen35/examples/escha_prefill_batch_gate.rs`. (The +/// DENSE half is not bit-identical: a batched WMMA GEMM does not accumulate +/// like a batch-1 GEMV.) +/// +/// # Slot layout +/// +/// Token-major: slot `s` is token `s / k`, rank `s % k`. This is the layout +/// `moe_topk_renorm_k8_batched` already writes `topk_indices` / `topk_weights` +/// in, and the layout `moe_down_combine_k8_batched` already reads +/// `expert_outputs` in, so nothing is permuted anywhere in this function. +/// +/// `out` (`[n_tokens, hidden]`) is accumulated into, never overwritten. +#[allow(clippy::too_many_arguments)] +pub fn escha_routed_prefill_indexed( + gpu: &mut Gpu, + tables: &EschaRoutedRefs<'_>, + scratch: &rdna_compute::scratch::EschaPrefillViews, + r: &EschaIndexedRouting<'_>, + out: &GpuTensor, + x_norm_batch: &GpuTensor, + hidden: usize, + mi: usize, + k: usize, + n_tokens: usize, +) -> Result<(), DispatchError> { + macro_rules! hip { + ($ex:expr) => { + $ex.map_err(|err| DispatchError::Hip(err.to_string())) + }; + } + // `moe_down_combine_k8_batched` unrolls to a hard 8 slots per token + // (`k < K_TOP` guard inside a `for k in 0..8`), so it silently DROPS ranks + // 8.. rather than failing. Same bound as both decode routes. + if k == 0 || k > 8 { + return Err(DispatchError::UnsupportedVariant { + family: "moe", + variant: "escha-routed-prefill-supports-k<=8", + arch: "", + quant: "", + }); + } + if n_tokens == 0 { + return Err(DispatchError::UnsupportedVariant { + family: "moe", + variant: "escha-routed-prefill-empty-batch", + arch: "", + quant: "", + }); + } + let slots = n_tokens + .checked_mul(k) + .ok_or(DispatchError::UnsupportedVariant { + family: "moe", + variant: "escha-routed-prefill-slot-overflow", + arch: "", + quant: "", + })?; + // The indexed GEMVs put `slots` on grid.y, which HIP caps at 65 535. Above + // that the launch would be truncated, not rejected — tokens past the cap + // would silently contribute nothing to the residual, which reads as a + // mildly worse model rather than as a failure. + if slots > 65_535 { + return Err(DispatchError::UnsupportedVariant { + family: "moe", + variant: "escha-routed-prefill-slots-exceed-grid-y", + arch: "", + quant: "", + }); + } + // The scratch is model-global and grows on demand, and its views are cut + // to a specific slot count; a caller that cut them for a different chunk + // would have every kernel here read or write the wrong extent. The + // wrappers also length-check exactly, so this is belt and braces — but it + // names the mistake instead of reporting a numel mismatch six frames down. + if scratch.slots != slots { + return Err(DispatchError::UnsupportedVariant { + family: "moe", + variant: "escha-prefill-scratch-slot-mismatch", + arch: "", + quant: "", + }); + } + // Same shape contract as the decode route: the GEMVs read + // `expert_ptrs[topk_indices[slot]]` on device, so a mis-sized weight table + // is an out-of-bounds READ, not a wrong answer, and a projection whose + // shape disagrees with the scratch would have the GEMV write past the end + // of `mid_gu` / `mid_dn`. + if r.gate_up_m != 2 * mi || r.gate_up_k != hidden || r.down_m != hidden || r.down_k != mi { + return Err(DispatchError::UnsupportedVariant { + family: "moe", + variant: "escha-indexed-routed-shape-mismatch", + arch: "", + quant: "", + }); + } + + let weights = &scratch.weights; + let xh_gu = &scratch.xh_gu; + let mid_gu = &scratch.mid_gu; + let y_gu = &scratch.y_gu; + let h = &scratch.h; + let xh_dn = &scratch.xh_dn; + let mid_dn = &scratch.mid_dn; + let y_dn = &scratch.y_dn; + // `topk_indices` / `topk_weights` are the [max_batch x k] prefill scratch; + // only the first `slots` entries are this chunk's. + let ids = r.topk_indices.sub_offset(0, slots); + let raw_weights = r.topk_weights.sub_offset(0, slots); + + // ── 0. f16-round the combine weights, ONE launch, out-of-place ──────── + hip!(gpu.escha_round_weights_f16_rne(&raw_weights, weights, slots))?; + + // ── 1. gate_up input transform, all slots, ONE launch ───────────────── + // Grouped(k): slot s reads token s/k's row of `x_norm_batch`. This is the + // one place batched prefill differs from decode, which broadcasts a single + // row to all k slots. + hip!(gpu.escha_h128_batched( + "escha_h128_in_batched", + x_norm_batch, + tables.gate_up_rin, + &ids, + xh_gu, + hidden, + slots, + EschaXGroup::Grouped(k), + ))?; + + // ── 2. gate_up GEMM for ALL slots, ONE launch ───────────────────────── + // Grouped by expert: at prefill batch sizes the slot-parallel form re-reads + // each expert's code once per slot and does not amortise at all. + escha_routed_gemm_grouped( + gpu, + r.gate_up_dtype, + r.expert_gate_up_ptrs, + &ids, + xh_gu, + mid_gu, + r.gate_up_m, + r.gate_up_k, + slots, + r.n_experts, + )?; + + // ── 3. gate_up output transform, ONE launch ─────────────────────────── + hip!(gpu.escha_h128_batched( + "escha_h128_out_batched", + mid_gu, + tables.gate_up_rout, + &ids, + y_gu, + 2 * mi, + slots, + EschaXGroup::PerSlot, + ))?; + + // ── 4. SwiGLU on the f16-rounded merged output, gate = FIRST half ───── + hip!(gpu.escha_swiglu_batched(y_gu, h, mi, slots))?; + + // ── 5. down input transform, ONE launch (per-slot activation) ───────── + hip!(gpu.escha_h128_batched( + "escha_h128_in_batched", + h, + tables.down_rin, + &ids, + xh_dn, + mi, + slots, + EschaXGroup::PerSlot, + ))?; + + // ── 6. down GEMM for ALL slots, ONE launch ──────────────────────────── + escha_routed_gemm_grouped( + gpu, + r.down_dtype, + r.expert_down_ptrs, + &ids, + xh_dn, + mid_dn, + r.down_m, + r.down_k, + slots, + r.n_experts, + )?; + + // ── 7. down output transform, ONE launch ────────────────────────────── + hip!(gpu.escha_h128_batched( + "escha_h128_out_batched", + mid_dn, + tables.down_rout, + &ids, + y_dn, + hidden, + slots, + EschaXGroup::PerSlot, + ))?; + + // ── 8. weighted combine into the residual, ONE launch ───────────────── + // `[n_tokens, k, hidden]` folded into `[n_tokens, hidden]`; blockIdx.y is + // already the token, so this is the same kernel decode uses at n = 1. + hip!(gpu.moe_down_combine_k8_batched(y_dn, weights, out, hidden, k, n_tokens))?; + Ok(()) +} diff --git a/crates/hipfire-dispatch/src/pipeline/mod.rs b/crates/hipfire-dispatch/src/pipeline/mod.rs index b4e80d7f48..017146483b 100644 --- a/crates/hipfire-dispatch/src/pipeline/mod.rs +++ b/crates/hipfire-dispatch/src/pipeline/mod.rs @@ -17,6 +17,11 @@ pub use steps::{execute_steps, FusedPattern, GemvInput, Step}; // only at this step; not on any live path until wired behind HIPFIRE_FORWARD_LOWERED). pub mod superop; +/// Escha-W2 routed-expert decode executor (Task 10). Replaces step 4 of the +/// CPU-top-K fallback for escha layers; everything else stays arch-6 code. +pub mod escha; +pub mod route_trace; + pub struct Pipeline { pub ops: &'static [PipelineOp], } @@ -174,7 +179,7 @@ pub fn check_moe_decode_batch_size(batch_size: usize) -> Result<(), DispatchErro Ok(()) } -/// GPU-free pre-guard for MoE decode (#397 Ship 4c). Rejects the two +/// GPU-free pre-guard for MoE decode (#397 Ship 4c). Rejects the /// truly-unsupported cases up front — *before* any GPU work — so the caller /// gets a clean [`DispatchError`] instead of a deep panic in the CPU-top-K /// fallback (`select_nth_unstable_by(k-1)` panics when `k == 0 || k > n_exp`) @@ -192,14 +197,61 @@ pub fn check_moe_decode_batch_size(batch_size: usize) -> Result<(), DispatchErro /// issue but experts are resident, the fallback runs it and its inner /// `gemv.run_auto` surfaces any genuinely-unsupported dtype as its own clean /// `DispatchError` — so we must NOT reject that case here.) +/// - **(c)** an Escha-W2 layer on the indexed GPU-top-K path WITHOUT the escha +/// indexed executor behind it. See below — this one is a silent-wrong-output +/// guard, not a panic guard. /// /// `routed_experts_resident` mirrors `!MoeParams::routed_experts.is_empty()` /// (false under paged residency, where only the GPU-top-K path is available). +/// +/// # (c) — why escha must fail closed on an UNSUPPORTED indexed path +/// +/// Escha-W2 weights live in a ROTATED domain: only the escha routed executors +/// in [`crate::pipeline::escha`] wrap the GEMVs in the H128 pair. The GENERIC +/// indexed routed body in [`run_moe_decode`] knows nothing about escha; +/// running an escha layer through it omits the transforms and emits finite, +/// fluent output that is wrong by ~1e-1 — no crash, no NaN, no test fires. +/// +/// Escha now has TWO supported indexed variants, and arm (c) is the assertion +/// that the layer is one of them: +/// +/// * `routed_indexable_escha_native` (Phase 2, production) — the routed +/// experts are the trellis CODE (`Escha2T16` / `Escha3T16`) and the fused +/// `escha_gemv_native_*` decodes it inside the GEMV; +/// * `routed_indexable_escha_q8` (Phase 1, the A/B arm) — the experts are the +/// Q8_0 the trellis decoded into at load. +/// +/// Both additionally require the H128 transform tables to be resident, and +/// both reach `escha::escha_routed_decode_indexed`, which `run_moe_decode` +/// branches to before the generic body ever runs. `escha_indexed_supported` is +/// the caller's assertion that a supported container AND the tables hold for +/// this layer. +/// +/// The arm was WIDENED for Phase 2 rather than removed, and it is still the +/// specific-combination test it always was. What it catches is every way escha +/// could arrive on the indexed path other than through those two arms: +/// +/// * transform tables missing (`MoeParams::escha == None`) while the layer +/// is still marked escha by dtype — the executor could not be called; +/// * an escha layer that resolved indexable through some OTHER arm, e.g. a +/// future graded/mixed escha file whose representative routed dtype is +/// neither escha-coded nor Q8_0. Each escha GEMV hard-codes its container's +/// bit geometry, so dispatching one on a different container is silent +/// corruption too — this hazard is the mirror image of the original one, +/// and both are refused here by requiring the *specific* supported +/// combinations rather than merely "escha, somehow, on the indexed path". +/// +/// It deliberately ERRORS rather than forcing `use_gpu_topk = false`. Forcing +/// would keep escha correct while hiding the fact that a new indexed arm needs +/// to be taught about escha (or explicitly excluded from it) — the whole point +/// is that the next person has to make that decision consciously. pub fn check_moe_decode_supported( use_gpu_topk: bool, k: usize, n_exp: usize, routed_experts_resident: bool, + has_escha: bool, + escha_indexed_supported: bool, ) -> Result<(), DispatchError> { // (a) k-range — required by BOTH the GPU-top-K path and the CPU fallback's // `select_nth_unstable_by(k-1)`. Universal precondition, not a k==8 check. @@ -222,6 +274,65 @@ pub fn check_moe_decode_supported( quant: "", }); } + // (c) escha on the indexed GPU-top-K path WITHOUT the escha indexed + // executor behind it: fail closed. The generic indexed body never applies + // the H128 pair, and each escha GEMV hard-codes one container's bit + // geometry — either mismatch is silently-wrong output rather than an + // error. See the module-level rationale on this function. + if has_escha && use_gpu_topk && !escha_indexed_supported { + return Err(DispatchError::UnsupportedVariant { + family: "moe", + variant: "escha-routed-experts-on-indexed-gpu-topk-path", + arch: "", + quant: "", + }); + } + Ok(()) +} + +/// The BATCHED-PREFILL mirror of [`check_moe_decode_supported`]'s arm (c). +/// +/// [`run_moe_prefill`] branches to the escha routed executor on +/// `MoePrefillParams::escha.is_some()`. If that is `None`, control falls +/// through into Path 1 / Path 2, which apply NO Hadamard transform and raise +/// no error — the same finite, fluent, ~1e-1-wrong output arm (c) exists to +/// prevent on the decode side. +/// +/// # Why the layer marker must be UNGATED +/// +/// `MoePrefillParams::escha` is `Some` only when the layer is escha AND +/// `escha_indexed_route_enabled()`. So `escha.is_none()` on its own cannot +/// tell "this is a plain Q8_0 MoE layer" apart from "this is an escha layer +/// with `HIPFIRE_ESCHA_INDEXED=0`". `layer_is_escha` is therefore taken +/// straight from the layer's own transform tables (`ffn.escha.is_some()`, a +/// load-time model-state property), never from an env var — exactly as the +/// router f16 rounding in `qwen35::prefill` already does. +/// +/// # Why this is not "safe today, therefore unnecessary" +/// +/// It is safe today only because no admission arm outside escha's own admits +/// Q8_0 routed experts to batched prefill — a property of a dtype table in +/// another crate, 200 lines from the branch that depends on it. The next +/// planned work is a Q8_0 grouped GEMM over sorted expert groups, which is +/// precisely a generic Q8_0 routed arm; when it lands, an escha layer with the +/// indexed route disabled becomes indistinguishable from a plain Q8_0 MoE +/// layer inside `run_moe_prefill` and takes the transform-free path. This +/// makes that a loud refusal at the point of danger instead. +/// +/// Like arm (c), it ERRORS rather than silently rerouting: a new prefill arm +/// has to be taught about escha, or explicitly excluded from it, consciously. +pub fn check_moe_prefill_supported( + layer_is_escha: bool, + escha_tables_present: bool, +) -> Result<(), DispatchError> { + if layer_is_escha && !escha_tables_present { + return Err(DispatchError::UnsupportedVariant { + family: "moe", + variant: "escha-routed-experts-on-non-escha-prefill-path", + arch: "", + quant: "", + }); + } Ok(()) } @@ -576,6 +687,32 @@ pub fn run_uniform_moe_down_expanded( } } +/// The exact arch predicate `run_moe_decode` uses to pick the fused +/// exact-wave64 router kernel (`moe_router_softmax_topk_k8_wave64_exact`) +/// over the reference two-launch (`softmax_f32` + `moe_topk_renorm_k8`) +/// path. Extracted to a standalone function — rather than left inline — +/// so any caller that needs to reproduce production's kernel choice (today: +/// the escha router-contract test helper in hipfire-arch-qwen35) calls the +/// real predicate instead of a hand-copied approximation that can silently +/// drift from it. +/// +/// `gfx1100_router_mode` is `HIPFIRE_GFX1100_ROUTER_W64`'s value (pass +/// `hipfire_config::developer_var("HIPFIRE_GFX1100_ROUTER_W64").ok().as_deref()` +/// to match production exactly): `"0"`/`"approx"` opt the gfx1100 production +/// path back out to the reference two-launch kernel or the old +/// non-bit-exact research kernel, respectively; anything else (including +/// unset) keeps the fused exact kernel on gfx1100. gfx1151 always takes the +/// fused exact kernel unconditionally. +pub fn exact_wave64_router_predicate( + n_exp: usize, + arch: &rdna_compute::arch_caps::ArchCaps, + gfx1100_router_mode: Option<&str>, +) -> bool { + n_exp == 256 + && ((arch.is_gfx1100() && !matches!(gfx1100_router_mode, Some("0" | "approx"))) + || arch.is_gfx1151()) +} + /// MoE decode executor. Ports the body of `moe_ffn_decode_impl` verbatim, /// substituting `ffn.*`/`config.*`/`s.*` references with `MoeParams` fields. /// Resolution is owned here (computed from `MoeDtypes` + k), and `ctx` is @@ -615,7 +752,31 @@ pub fn run_moe_decode( // deep `select_nth_unstable_by` panic in the fallback into a clean error. // NOTE: k != 8 is intentionally NOT rejected — the fallback handles k ∈ // [1, n_exp] (MQ4 k=4, F32 k=2, …). - check_moe_decode_supported(res.use_gpu_topk, p.k, p.n_exp, !p.routed_experts.is_empty())?; + // + // `p.has_escha()` feeds arm (c): an Escha-W2 layer on the indexed + // GPU-top-K path is refused HERE, before any GPU work, UNLESS it is one of + // the two supported shapes — resolved through + // `routed_indexable_escha_native` (Phase 2: the trellis code, fused GEMV) + // or `routed_indexable_escha_q8` (Phase 1: the decoded Q8_0) AND carrying + // the transform tables the escha indexed executor needs. Both arguments + // must stay wired to the real values; a constant re-opens exactly the hole + // they close. + // + // `escha_indexed_supported` is deliberately the AND of the resolver's + // escha arm and the tables' presence, and it is computed once here so the + // guard and the dispatch below cannot drift apart — the branch to + // `escha_routed_decode_indexed` re-reads THIS binding rather than + // recomputing the predicate. + let escha_indexed_supported = + (res.routed_indexable_escha_q8 || res.routed_indexable_escha_native) && p.escha.is_some(); + check_moe_decode_supported( + res.use_gpu_topk, + p.k, + p.n_exp, + !p.routed_experts.is_empty(), + p.has_escha(), + escha_indexed_supported, + )?; // EP (Ship 6 substrate-EP): when `routed_out` is set, the shared-down and // routed-combine accumulate into that zeroed partial (all-reduced by the EP @@ -859,6 +1020,20 @@ pub fn run_moe_decode( // shared-expert down → generic per-expert routed loop, then returns. It // does NOT fall through to the indexed GPU-top-K path below (which assumes // k=8 + an indexable routed dtype). + // Escha-only router-logits f16 round-trip. This is HOISTED ABOVE the + // CPU-fallback return on purpose (Task 10 fix): Task 9 placed it further + // down, on the GPU-top-K path only, but escha's routed experts are stored + // Q8_0 and Q8_0 is not an indexable routed dtype — so every real escha + // layer takes the `!use_gpu_topk` branch below and the rounding was + // unreachable on the one model family that needs it. The rationale for + // the rounding itself is unchanged; see the comment at the (now + // no-op-for-escha) second call site below and + // `MoeDtypes::has_escha_experts`. `p.has_escha()` also admits layers whose + // routed dtype has been rewritten to Q8_0 by the escha loader, which + // `has_escha_experts()` alone can no longer see. + if p.has_escha() { + hip!(gpu.router_logits_round_f16_rne(p.router_logits))?; + } if !res.use_gpu_topk { return run_moe_decode_cpu_fallback(ctx, gpu, p, &shared_gate, &shared_up); } @@ -882,15 +1057,31 @@ pub fn run_moe_decode( } } let gfx1100_router_mode = hipfire_config::developer_var("HIPFIRE_GFX1100_ROUTER_W64").ok(); - let gfx1151_radiowave_fusions = ctx.arch.is_gfx1151(); - let exact_wave64_router = p.n_exp == 256 - && ((ctx.arch.is_gfx1100() - // The exact fused router is the production gfx1100 path. `0` retains - // the two-launch reference path for A/B diagnosis; `approx` retains - // the old non-bit-exact research kernel without exposing it by - // accident through the former `1` opt-in. - && !matches!(gfx1100_router_mode.as_deref(), Some("0" | "approx"))) - || gfx1151_radiowave_fusions); + let exact_wave64_router = + exact_wave64_router_predicate(p.n_exp, &ctx.arch, gfx1100_router_mode.as_deref()); + // Escha-only router-logits f16 round-trip (review Fix 1). EschaLabs' + // runtime computes router logits as f16(x @ gate_w.T) and only then + // widens to F32 to select top-k; hipfire keeps logits F32 end-to-end. + // The two selections differ whenever two experts' F32 logits round to + // the same f16 value AND straddle the top-k boundary (measured ~0.42% + // of router decisions across 11 layers). Escha's recovery fine-tune was + // trained against the f16-rounding runtime, so this model family must + // reproduce that rounding to avoid a silent, unexplained divergence + // from Escha's own expert choice. + // + // Gated on `MoeDtypes::has_escha_experts` — a model-state property fixed + // at load time by the routed-expert dtype actually on disk, never a + // global env var — so every `qwen3.6:35b-a3b-*` SKU and every other + // arch-6 model takes the branch below as a no-op and keeps its current + // selection bit-for-bit. Applied to the shared `router_logits` buffer + // BEFORE the kernel-selection `if`/`else` below, so both routes (the + // fused exact-wave64 kernel and the softmax_f32 + moe_topk_renorm_k8 + // fallback pair) see identically-rounded logits. + // (The rounding itself was applied above, hoisted so the CPU-top-K + // fallback gets it too. Re-rounding here would be numerically a no-op — + // f16(f16(x)) == f16(x) — but the launch would not be free, so it is not + // repeated. This site keeps the rationale next to the selection kernels + // it protects.) static ROUTER_SHARED_FUSE: OnceLock = OnceLock::new(); let router_shared_fuse = exact_wave64_router && p.batch_size == 1 @@ -1099,6 +1290,87 @@ pub fn run_moe_decode( } } + // ── Escha-W2 routed experts, indexed (device-resident) routing ─────────── + // Escha weights are in a rotated domain; the generic indexed body below + // would run them without the H128 pair and emit ~1e-1-wrong output with + // nothing to catch it. This branch runs the escha executor instead — same + // eight phases and the same transforms as the CPU-top-K route, with the + // routing never leaving the device. + // + // It returns rather than falling through: it has already accumulated the + // routed contribution into `out_target`, so the generic body must not run. + // Selection capture (off unless HIPFIRE_ESCHA_ROUTE_TRACE is set). Placed + // after top-k and before any routed body, so it records what the expert + // GEMVs are about to index with on EITHER decode route. + if crate::pipeline::route_trace::enabled() { + crate::pipeline::route_trace::record(gpu, p.topk_indices, p.batch_size, p.k); + } + if escha_indexed_supported { + let escha = p + .escha + .as_ref() + .expect("escha_indexed_supported implies escha tables"); + // Same refusals the CPU-top-K escha branch makes, for the same + // reasons: the executor has no AWQ / graded-tier arm, and Hessian + // capture keyed on `x_norm` would record the H128 outputs instead of + // the raw pre-rotation activations and silently poison the Hessians. + if p.expert_down_awq_ptrs.is_some() || p.expert_dtype_tags.is_some() { + return Err(DispatchError::UnsupportedVariant { + family: "moe", + variant: "escha-routed-with-awq-or-graded-tiers", + arch: "", + quant: "", + }); + } + if gpu.hessian_capture.is_some() { + return Err(DispatchError::UnsupportedVariant { + family: "moe", + variant: "escha-routed-hessian-capture-unsupported", + arch: "", + quant: "", + }); + } + // The escha executor always folds the weighted combine into the layer + // (phase 8). `defer_routed_combine` promises the caller an EXPANDED, + // uncombined `down_expanded` it will fold itself — honouring the flag + // is not possible here, and ignoring it would double-count. + if p.defer_routed_combine { + return Err(DispatchError::UnsupportedVariant { + family: "moe", + variant: "escha-routed-cannot-defer-combine", + arch: "", + quant: "", + }); + } + return crate::pipeline::escha::escha_routed_decode_indexed( + gpu, + escha, + &crate::pipeline::escha::EschaIndexedRouting { + expert_gate_up_ptrs: p.expert_gate_up_ptrs, + expert_down_ptrs: p.expert_down_ptrs, + topk_indices: p.topk_indices, + topk_weights: p.topk_weights, + n_experts: p.n_exp, + // The container the expert slots hold, straight off the layer's + // routed dtype — the same fact `escha_indexed_supported` was + // resolved from, so the GEMV the executor picks and the arm the + // guard admitted are the same decision read twice, not two + // decisions that could drift. + gate_up_dtype: p.dtypes.routed_gate_up, + down_dtype: p.dtypes.routed_down, + gate_up_m: 2 * p.mi, + gate_up_k: p.routed_gate_up_k, + down_m: p.routed_down_m, + down_k: p.routed_down_k, + }, + out_target, + p.x_norm, + p.hidden, + p.mi, + p.k, + ); + } + // ── Indexed routed experts ──────────────────────────────────────────────── // Signs back the FWHT used by every MQ4/MQ6 gate_up rotation + silu-rotate // (idempotent/cached). Only the paro path is sign-free. @@ -2037,6 +2309,47 @@ fn run_moe_decode_cpu_fallback( }); } + // ── 4a. Escha-W2 routed experts: the H128-wrapped, batched executor ────── + // Escha weights are in a rotated domain; a plain per-expert `run_auto` + // here would silently produce ~1e-1-wrong output. The executor also + // batches the transforms across the token's k experts, which is a + // measured hard requirement (see pipeline::escha module docs), so it + // replaces the loop below wholesale rather than wrapping each iteration. + if let Some(escha) = p.escha.as_ref() { + if p.expert_down_awq_ptrs.is_some() || p.expert_dtype_tags.is_some() { + return Err(DispatchError::UnsupportedVariant { + family: "moe", + variant: "escha-routed-with-awq-or-graded-tiers", + arch: "", + quant: "", + }); + } + if gpu.hessian_capture.is_some() { + // The capture keys off the RAW pre-rotation activations; on the + // escha path those are the H128 outputs, not `x_norm`/`silu(g)*u`, + // so silently reusing the loop below's keys would poison the + // Hessians. Refuse rather than record the wrong thing. + return Err(DispatchError::UnsupportedVariant { + family: "moe", + variant: "escha-routed-hessian-capture-unsupported", + arch: "", + quant: "", + }); + } + return crate::pipeline::escha::escha_routed_decode( + ctx, + gpu, + escha, + p.routed_experts, + &topk_indices, + &topk_weights, + p.x_norm, + p.x_residual, + p.hidden, + mi, + ); + } + // ── 4. Per-expert routed loop (master's generic `weight_gemv` arm) ──────── static GEMV_FB: OnceLock = OnceLock::new(); let gemv = GEMV_FB.get_or_init(GemvFamily::new); @@ -3109,6 +3422,12 @@ pub fn run_moe_prefill( }; } + // BEFORE any GPU work, and before the escha branch below: an escha layer + // that reaches here without its transform tables would fall into Path 1 / + // Path 2, which apply no Hadamard transform and raise no error. See + // `check_moe_prefill_supported`. + check_moe_prefill_supported(p.layer_is_escha, p.escha.is_some())?; + let res = MoePrefillResolution::resolve(&p.dtypes, &ctx.arch, &ctx.flags); let force_mq4_grouped_fp16 = res.force_mq4_grouped_fp16 || p.force_mq4_grouped_fp16; if hipfire_config::developer_var("HIPFIRE_MOE_PREFILL_TRACE") @@ -3145,6 +3464,79 @@ pub fn run_moe_prefill( // partials equals the full single-GPU routed combine. let out_target: &GpuTensor = p.routed_out.unwrap_or(p.x_batch); + // ── Escha-W2 routed experts, batched + indexed ─────────────────────────── + // + // Mirrors the escha branch in `run_moe_decode`, and for the same reason: + // escha weights live in a ROTATED domain, so Path 1 and Path 2 below — + // which feed the activation straight into the expert GEMVs and combine the + // raw result — would omit both Hadamard transforms and emit finite, + // fluent, ~1e-1-wrong output with nothing to catch it. This branch RETURNS; + // it has already accumulated the routed contribution into `out_target`. + // + // It runs BEFORE the Path 2 scatter so escha never pays for a scatter it + // does not use. + if let Some(escha) = p.escha.as_ref() { + // The same four refusals the two decode escha branches make, for the + // same reasons. Kept verbatim rather than factored out: each one is a + // claim about THIS executor, and a shared helper would let a future + // divergence between the routes go unnoticed. + // + // AWQ / graded tiers: the escha executor has no arm for either. + if p.expert_down_awq_ptrs.is_some() || p.expert_dtype_tags.is_some() { + return Err(DispatchError::UnsupportedVariant { + family: "moe", + variant: "escha-routed-with-awq-or-graded-tiers", + arch: "", + quant: "", + }); + } + // Hessian capture keyed on the activation would record the H128 + // outputs instead of the raw pre-rotation activations, silently + // poisoning the Hessians. + if gpu.hessian_capture.is_some() { + return Err(DispatchError::UnsupportedVariant { + family: "moe", + variant: "escha-routed-hessian-capture-unsupported", + arch: "", + quant: "", + }); + } + // The escha executor always folds the weighted combine into the layer + // (phase 8); there is no expanded, uncombined `down_expanded` to hand + // back. (Prefill has no `defer_routed_combine` flag today — the decode + // branch refuses one — so this is recorded as a comment rather than a + // dead check. If prefill gains the flag, it must refuse here.) + let slots = total_slots; + let scratch = hip!(gpu.ensure_escha_prefill_scratch(slots, p.hidden, mi))?; + return crate::pipeline::escha::escha_routed_prefill_indexed( + gpu, + escha, + &scratch, + &crate::pipeline::escha::EschaIndexedRouting { + expert_gate_up_ptrs: p.expert_gate_up_ptrs, + expert_down_ptrs: p.expert_down_ptrs, + topk_indices: p.topk_indices, + topk_weights: p.topk_weights, + n_experts: p.n_exp, + // See the decode branch: the container comes off the layer's + // own routed dtype, which is what the batched-prefill + // admission arm keyed on too. + gate_up_dtype: p.dtypes.routed_gate_up, + down_dtype: p.dtypes.routed_down, + gate_up_m: 2 * mi, + gate_up_k, + down_m, + down_k, + }, + out_target, + p.x_norm_batch, + p.hidden, + mi, + k_top, + n, + ); + } + // ── Path 2 scatter pipeline ─────────────────────────────────────── let mut path2_m_total: usize = 0; if res.use_path2 { diff --git a/crates/hipfire-dispatch/src/pipeline/route_trace.rs b/crates/hipfire-dispatch/src/pipeline/route_trace.rs new file mode 100644 index 0000000000..36524b91d1 --- /dev/null +++ b/crates/hipfire-dispatch/src/pipeline/route_trace.rs @@ -0,0 +1,144 @@ +// SPDX-License-Identifier: MIT OR Apache-2.0 +// hipfire — see LICENSE and NOTICE in the project root. +//! Opt-in capture of the routed-expert SELECTION, so "prefill and decode +//! sometimes choose different experts" can be a measured rate instead of an +//! estimate. +//! +//! # Why the selection and not the router logits +//! +//! It is tempting to capture the logits and re-run top-k on the host. That +//! measures a different thing: the device kernels +//! (`moe_router_softmax_topk_k8_*` / `moe_topk_renorm_k8_batched`) have their +//! own ordering and tie-breaking, and a host re-implementation would attribute +//! its own disagreements with them to the routes. What matters is what the +//! model actually computed with, so this records the `topk_indices` buffer the +//! expert GEMVs are about to index with. +//! +//! # Why it is a file and not a return value +//! +//! The two routes being compared are a batched prefill and a per-token loop; +//! they are called from different places, produce a different number of +//! records per call, and neither has a natural channel back to the caller that +//! wants the comparison. A trace file keyed by call order is the cheapest +//! thing that lets one gate read both. +//! +//! Off unless `HIPFIRE_ESCHA_ROUTE_TRACE` names a path. When off, the only +//! cost is one relaxed atomic load per MoE layer. +//! +//! # Record format +//! +//! Little-endian, appended in CALL ORDER: +//! +//! ```text +//! u32 n_tokens // rows in this record (1 for decode, chunk size for prefill) +//! u32 k // experts per token +//! i32 ids[n_tokens*k] // token-major: row t is ids[t*k .. (t+1)*k] +//! ``` +//! +//! A reader reconstructs `(token, layer)` from call order: every MoE layer +//! emits exactly one record per forward, so with `L` MoE layers the j-th +//! record of a per-token run is `(token = j / L, layer = j % L)` and the j-th +//! record of a batched run covering one chunk is `(layer = j % L)` with its +//! `n_tokens` rows being that chunk's tokens in order. Readers MUST check the +//! totals agree rather than assuming — see the gate. + +use std::fs::File; +use std::io::Write; +use std::sync::atomic::{AtomicBool, Ordering}; +use std::sync::Mutex; +use std::sync::OnceLock; + +use rdna_compute::{Gpu, GpuTensor}; + +/// Mirror of "`sink()` currently holds a `Some`", so [`enabled`] can answer +/// with a relaxed load instead of taking the sink mutex. Written only by +/// `sink()`'s initialiser and by [`reopen`], both of which happen O(1) times +/// per process; read once per MoE layer. +static ACTIVE: AtomicBool = AtomicBool::new(false); + +fn sink() -> &'static Mutex> { + static SINK: OnceLock>> = OnceLock::new(); + SINK.get_or_init(|| { + let f = hipfire_config::developer_var("HIPFIRE_ESCHA_ROUTE_TRACE") + .ok() + .and_then(|path| match File::create(&path) { + Ok(f) => Some(f), + Err(e) => { + eprintln!("[escha-route-trace] cannot create {path}: {e}"); + None + } + }); + ACTIVE.store(f.is_some(), Ordering::Relaxed); + Mutex::new(f) + }) +} + +/// Redirect the trace to a new file, closing the previous one. +/// +/// Exists because the comparison this module serves needs TWO traces from ONE +/// process (a batched arm and a per-token arm over the same prompt, with the +/// model loaded once — it is 37.6 GB). Re-reading the env var would not work: +/// the sink is initialised on first use, which is inside the first MoE layer, +/// long before the second arm starts. +pub fn reopen(path: &str) { + let f = match File::create(path) { + Ok(f) => Some(f), + Err(e) => { + eprintln!("[escha-route-trace] cannot create {path}: {e}"); + None + } + }; + if let Ok(mut slot) = sink().lock() { + ACTIVE.store(f.is_some(), Ordering::Relaxed); + *slot = f; + } +} + +/// True when tracing is on. Callers should check this before doing anything +/// expensive (the capture below forces a device sync). +/// +/// This is the "one relaxed atomic load per MoE layer" the module docs +/// promise, and it has to be: it is called from the MoE routed path on every +/// forward of every model, tracing or not. It used to be +/// `sink().lock().map(...)` — a mutex acquisition per layer, contending with +/// nothing, to answer a question that changes O(1) times per process. +/// +/// The `OnceLock` below exists only to force the lazy env-derived +/// initialisation of `sink()` exactly once; after that the answer comes from +/// [`ACTIVE`]. It cannot be replaced by reading the env var directly: +/// [`reopen`] installs a sink with no env var set at all (that is how the G6 +/// gate takes two traces from one process), so "tracing was requested via the +/// environment" is not the same predicate as "there is somewhere to write". +pub fn enabled() -> bool { + static INIT: OnceLock<()> = OnceLock::new(); + INIT.get_or_init(|| { + let _ = sink(); + }); + ACTIVE.load(Ordering::Relaxed) +} + +/// Append one layer's selection. +/// +/// `topk_indices` is the device buffer the routed GEMVs index with: `i32` +/// expert ids stored in an F32 tensor (the same 4-bytes-per-element +/// reinterpretation `moe_topk_renorm_k8_batched` writes). Only the first +/// `n_tokens * k` entries are read — the buffer itself is usually the larger +/// `[max_batch x k]` scratch. +/// +/// Synchronises the device: this is a diagnostic, and a capture that raced the +/// kernel that produced the values would measure nothing. +pub fn record(gpu: &Gpu, topk_indices: &GpuTensor, n_tokens: usize, k: usize) { + let want = n_tokens * k; + if gpu.hip.device_synchronize().is_err() { + return; + } + let mut bytes = vec![0u8; want * 4]; + if gpu.hip.memcpy_dtoh(&mut bytes, &topk_indices.buf).is_err() { + return; + } + let Ok(mut slot) = sink().lock() else { return }; + let Some(f) = slot.as_mut() else { return }; + let _ = f.write_all(&(n_tokens as u32).to_le_bytes()); + let _ = f.write_all(&(k as u32).to_le_bytes()); + let _ = f.write_all(&bytes); +} diff --git a/crates/hipfire-dispatch/src/tables/gemv_table.rs b/crates/hipfire-dispatch/src/tables/gemv_table.rs index 7a638c2031..5a76cc9243 100644 --- a/crates/hipfire-dispatch/src/tables/gemv_table.rs +++ b/crates/hipfire-dispatch/src/tables/gemv_table.rs @@ -62,11 +62,14 @@ fn register_plain(registry: &mut KernelRegistry) { let Ok(key) = KernelKey::for_gemv(dtype, GemvVariant::Plain, false) else { continue; }; + let Ok(steps) = KernelKey::gemv_steps(dtype, GemvVariant::Plain) else { + continue; + }; registry.register(KernelVariant { key, arch_required: KernelKey::dtype_arch_predicate(dtype), shape_gate: None, - steps: KernelKey::gemv_steps(dtype, GemvVariant::Plain), + steps, has_awq: dtype == DType::MQ4G256, tile: TileImpl::None, }); @@ -100,11 +103,14 @@ fn register_prerotated(registry: &mut KernelRegistry) { let Ok(key) = KernelKey::for_gemv_prerotated(dtype) else { continue; }; + let Ok(steps) = KernelKey::gemv_steps(dtype, GemvVariant::Prerotated) else { + continue; + }; registry.register(KernelVariant { key, arch_required: KernelKey::dtype_arch_predicate(dtype), shape_gate: None, - steps: KernelKey::gemv_steps(dtype, GemvVariant::Prerotated), + steps, has_awq: dtype == DType::MQ4G256, tile: TileImpl::None, }); @@ -133,11 +139,14 @@ fn register_residual(registry: &mut KernelRegistry) { let Ok(key) = KernelKey::for_gemv_residual(dtype) else { continue; }; + let Ok(steps) = KernelKey::gemv_steps(dtype, GemvVariant::WithResidual) else { + continue; + }; registry.register(KernelVariant { key, arch_required: KernelKey::dtype_arch_predicate(dtype), shape_gate: None, - steps: KernelKey::gemv_steps(dtype, GemvVariant::WithResidual), + steps, has_awq: dtype == DType::MQ4G256, tile: TileImpl::None, }); @@ -177,11 +186,14 @@ fn register_swiglu_residual(registry: &mut KernelRegistry) { let Ok(key) = KernelKey::for_gemv_swiglu_residual(dtype) else { continue; }; + let Ok(steps) = KernelKey::gemv_steps(dtype, GemvVariant::WithSwiGLUResidual) else { + continue; + }; registry.register(KernelVariant { key, arch_required: KernelKey::dtype_arch_predicate(dtype), shape_gate: None, - steps: KernelKey::gemv_steps(dtype, GemvVariant::WithSwiGLUResidual), + steps, has_awq: dtype == DType::MQ4G256, tile: TileImpl::None, }); diff --git a/crates/hipfire-dispatch/src/tables/mod.rs b/crates/hipfire-dispatch/src/tables/mod.rs index 1b863e1524..0abf92fc53 100644 --- a/crates/hipfire-dispatch/src/tables/mod.rs +++ b/crates/hipfire-dispatch/src/tables/mod.rs @@ -114,6 +114,11 @@ impl ArchPredicate { Self::HasCdna3LdsGemv => ctx.arch.has_cdna3_lds_gemv(), Self::HasDp4a => ctx.arch.gemv_dp4a_enabled(), Self::IsGfx942 => ctx.arch.is_gfx942(), + // No kernel exists yet for whatever dtype carries this predicate. + // Fail closed on every arch rather than advertise availability + // nothing backs. See the variant's doc comment on why this must + // never collapse to `Always`. + Self::Unimplemented => false, } } } diff --git a/crates/hipfire-dispatch/src/tests.rs b/crates/hipfire-dispatch/src/tests.rs index 51aebb0226..b18aef9c67 100644 --- a/crates/hipfire-dispatch/src/tests.rs +++ b/crates/hipfire-dispatch/src/tests.rs @@ -80,6 +80,17 @@ fn dp4a_variant(key: KernelKey) -> KernelVariant { } } +fn unimplemented_variant(key: KernelKey) -> KernelVariant { + KernelVariant { + key, + arch_required: ArchPredicate::Unimplemented, + shape_gate: None, + steps: &[], + has_awq: false, + tile: TileImpl::None, + } +} + // ── ShapePredicate::eval ────────────────────────────────────────────────────── #[test] @@ -200,6 +211,21 @@ fn arch_has_mmq_on_rdna3_or_rdna4() { assert!(ArchPredicate::HasMmq.eval_arch(&ctx_rdna4())); // RDNA4 MQ6/HFQ6 } +/// `Unimplemented` must evaluate to `false` on every architecture this test +/// can construct — it exists to fail closed for a dtype with no kernel yet +/// (currently Escha2T16/Escha3T16). Unlike the other predicates above, there +/// is no arch that should ever flip this to `true`; that is the entire point +/// of the variant (see its doc comment in `types.rs`). +#[test] +fn arch_unimplemented_is_false_on_every_arch() { + assert!(!ArchPredicate::Unimplemented.eval_arch(&ctx_rdna1())); + assert!(!ArchPredicate::Unimplemented.eval_arch(&ctx_rdna2())); + assert!(!ArchPredicate::Unimplemented.eval_arch(&ctx_rdna3())); + assert!(!ArchPredicate::Unimplemented.eval_arch(&ctx_rdna4())); + assert!(!ArchPredicate::Unimplemented.eval_arch(&ctx_gfx906())); + assert!(!ArchPredicate::Unimplemented.eval_arch(&DispatchCtx::for_test("gfx942"))); +} + #[test] fn arch_gemv_dp4a_gfx906_only() { // HasDp4a (=v_dot4_i32_i8, gfx906-only) @@ -342,6 +368,31 @@ fn registry_resolve_falls_through_to_second_variant() { ); } +/// A kernel gated on `ArchPredicate::Unimplemented` must never resolve as +/// available, on any architecture — that is the fail-closed contract for a +/// dtype with no kernel yet (e.g. Escha2T16/Escha3T16 in +/// `dtype_arch_predicate`). This is the scenario Finding 1 describes: a +/// future contributor adds a real registration for such a dtype to +/// `gemv_table.rs` without touching `dtype_arch_predicate` — the predicate +/// alone must be what stops `resolve()` from advertising it everywhere. +#[test] +fn registry_resolve_unimplemented_arch_never_resolves_on_any_arch() { + let mut reg = KernelRegistry::new(); + reg.register(unimplemented_variant(KernelKey::GemvF32)); + for arch in [ + "gfx1010", "gfx1030", "gfx1100", "gfx1200", "gfx906", "gfx942", + ] { + let ctx = DispatchCtx::for_test(arch); + let err = reg + .resolve(KernelKey::GemvF32, &ctx, None) + .expect_err("Unimplemented must never resolve to an available kernel"); + assert!( + matches!(err, DispatchError::MissingImpl { .. }), + "expected MissingImpl on {arch}, got {err:?}" + ); + } +} + #[test] fn gemm_q8_0_batched_wide_exact_resolves_on_wmma_only() { use crate::families::gemm::GemmFamily; @@ -747,7 +798,7 @@ fn gemv_steps_rotation_matches_plan() { DType::ParoQ4G128, DType::HFQ4G256, ] { - let steps = KernelKey::gemv_steps(dtype, GemvVariant::Plain); + let steps = KernelKey::gemv_steps(dtype, GemvVariant::Plain).unwrap(); let plan = dtype_rotation_plan(dtype); let has_fwht = steps.contains(&PipelineOp::RotateFwht); let has_givens = steps.contains(&PipelineOp::GivensRotate); @@ -768,10 +819,48 @@ fn gemv_steps_rotation_matches_plan() { assert!(!has_fwht && !has_givens, "{dtype:?}: no rotation"); } RotationPlan::Mq8Internal => {} + // Not exercised by the fixed dtype list above (no Escha GEMV kernel + // exists yet to produce steps for) — present only so this exhaustive + // match keeps compiling once RotationPlan gained EschaH128. + RotationPlan::EschaH128 => {} } } } +/// `gemv_steps(Plain, _)` must reject Escha-W2 dtypes with an explicit `Err` +/// (no H128 rotate/GEMV kernel exists yet — see the `RotationPlan::EschaH128` +/// arm), while every existing dtype keeps resolving to `Ok` exactly as +/// before this function was converted from a panicking `&'static [_]` return +/// to a `Result`. +#[test] +fn gemv_steps_rejects_escha_and_keeps_existing_dtypes_ok() { + for dtype in [DType::Escha2T16, DType::Escha3T16] { + assert!( + KernelKey::gemv_steps(dtype, GemvVariant::Plain).is_err(), + "{dtype:?}: gemv_steps(Plain) should be Err — no Escha GEMV kernel exists yet" + ); + } + + // MQ4G256: rotated dtype (RotationPlan::FwhtG256), exercises the `_` + // catch-all arm — must still succeed with the RotateFwht+Gemv steps. + let rotated = KernelKey::gemv_steps(DType::MQ4G256, GemvVariant::Plain) + .expect("MQ4G256 (rotated) must still resolve to Ok"); + assert_eq!( + rotated, + &[PipelineOp::RotateFwht, PipelineOp::Gemv], + "MQ4G256 (rotated) must keep the same step list" + ); + + // F32: unrotated dtype (RotationPlan::None) — must still succeed. + let unrotated = KernelKey::gemv_steps(DType::F32, GemvVariant::Plain) + .expect("F32 (unrotated) must still resolve to Ok"); + assert_eq!( + unrotated, + &[PipelineOp::Gemv], + "F32 (unrotated) must keep the same step list" + ); +} + // ── GemvFamily::resolve via populated table ─────────────────────────────────── #[test] @@ -969,6 +1058,7 @@ fn dtypes_all_mq4() -> MoeDtypes { has_paro_shared: false, per_expert_gate_up: None, per_expert_down: None, + routed_escha_transforms: false, } } @@ -2194,6 +2284,7 @@ fn moe_dtypes_mq4() -> MoeDtypes { has_paro_shared: false, per_expert_gate_up: None, per_expert_down: None, + routed_escha_transforms: false, } } diff --git a/crates/hipfire-dispatch/src/types.rs b/crates/hipfire-dispatch/src/types.rs index a191089a0e..8eaca42ecd 100644 --- a/crates/hipfire-dispatch/src/types.rs +++ b/crates/hipfire-dispatch/src/types.rs @@ -111,6 +111,9 @@ pub enum RotationPlan { FwhtG128, Mq8Internal, Givens, + /// Escha-W2: unnormalised 128-point Walsh-Hadamard on BOTH sides, + /// RS = 1/sqrt(128), signs folded into rin/rout rather than seeded. + EschaH128, } /// Sign-domain plan for a dtype. `None` <=> no activation rotation required. @@ -137,6 +140,12 @@ pub fn dtype_rotation_plan(dtype: DType) -> RotationPlan { // whoever adds the next Lloyd variant should be forced to decide which // side of this line it belongs on. MQ2G256LloydU => RotationPlan::None, + // Escha-W2: rotated-domain weights, 128-point Hadamard on both sides. + // Stated explicitly rather than left to the `_` fallthrough below — + // reaching RotationPlan::None here would skip the H128 pair and feed + // an unrotated activation through a rotated-domain weight, which is + // silent, fluent-looking garbage rather than a crash. + Escha2T16 | Escha3T16 => RotationPlan::EschaH128, _ => RotationPlan::None, } } @@ -155,6 +164,15 @@ pub fn dtype_post_rotation_variant(dtype: DType) -> GemvVariant { | MQ4G256Lloyd | MFP4G32 | MFP4G32Lloyd | MFP4G32P | MFP4G32E8 | MFP4G32E8SOA | MQ4G128 => { GemvVariant::Prerotated } + // Escha-W2: stated explicitly rather than left to the `_` fallthrough. + // No GEMV kernel exists yet, so this value is never actually launched — + // `for_gemv(_, Plain)` has no arm for these types (always Err), and + // `prepare_rotation_scratch` errors on `RotationPlan::EschaH128` before + // any rotate kernel runs, so the end-to-end path is Err regardless of + // what is returned here. Kept as `Plain` (the same value the `_` arm + // would produce) purely so a future reader sees this dtype was + // considered, not missed. + Escha2T16 | Escha3T16 => GemvVariant::Plain, _ => GemvVariant::Plain, } } @@ -579,6 +597,16 @@ pub enum ArchPredicate { /// BF16 MFMA GEMM. Narrower than `is_cdna3()` on purpose: the wrapper /// refuses non-gfx942 outright, so the predicate must match the wrapper. IsGfx942, + /// No kernel exists for this dtype yet; fail closed. Evaluates to `false` + /// on every architecture (see `eval_arch`), so any table registration + /// gated on this predicate is dead-on-arrival rather than silently + /// advertised as available everywhere. This is the correct default for a + /// dtype whose GEMV/GEMM kernel is future work — `Always` would be an + /// affirmatively wrong claim that compiles cleanly and only breaks at + /// runtime, on real hardware, the day someone registers the kernel. + /// MUST be replaced (not removed) with the kernel's real arch gate when + /// that kernel lands. + Unimplemented, } #[derive(Clone, Debug)] @@ -881,21 +909,58 @@ impl KernelKey { // inherits the identical arch gating. MQ2G256Lloyd | MQ2G256LloydU | MQ3G256Lloyd | MQ4G256Lloyd => ArchPredicate::HasWave32, MQ2G256GL | MQ3G256GL | MQ4G256V2 | MQ2G256V2 | MQ3G256V2 | MQ5G256V2 | MQ6G256V2 | MQ4CG256 => ArchPredicate::HasWave32, + // Escha-W2: no GEMV kernel exists yet (Task 4 registers the ids/rotation + // plan only; the HIP kernels are future work), so there is no real + // hardware requirement to encode yet. `Unimplemented` fails closed on + // every arch (see its doc comment) rather than claiming — as `Always` + // would — that a kernel runs everywhere. It is inert today the same way + // `Always` was inert: `for_gemv` has no arm for these types (always Err) + // and `dtype_post_rotation_variant` never maps them to `Prerotated`, so + // neither call site in `pipeline/mod.rs` that guards on + // `dtype_arch_predicate(..).eval_arch()` can reach this value for Escha. + // The difference is what happens the day someone registers an Escha + // entry in `gemv_table.rs` without touching this function: `Always` + // would silently advertise it on every arch; `Unimplemented` fails + // closed with an explicit dispatch error instead. Replace with the + // real arch gate when the Escha GEMV kernel lands. + Escha2T16 | Escha3T16 => ArchPredicate::Unimplemented, Q8HFQ | Raw => ArchPredicate::Always, } } /// Pipeline steps required for a given (DType, GemvVariant) pair. - pub fn gemv_steps(dtype: DType, variant: GemvVariant) -> &'static [PipelineOp] { + pub fn gemv_steps( + dtype: DType, + variant: GemvVariant, + ) -> Result<&'static [PipelineOp], DispatchError> { use DType::*; use GemvVariant::*; match variant { Plain => match dtype_rotation_plan(dtype) { - RotationPlan::Givens => &[PipelineOp::GivensRotate, PipelineOp::Gemv], - RotationPlan::None => &[PipelineOp::Gemv], - _ => &[PipelineOp::RotateFwht, PipelineOp::Gemv], + RotationPlan::Givens => Ok(&[PipelineOp::GivensRotate, PipelineOp::Gemv]), + RotationPlan::None => Ok(&[PipelineOp::Gemv]), + // Escha-W2: stated explicitly rather than left to the `_` + // catch-all below, which would mislabel the 128-point + // Hadamard as a `RotateFwht` step — a plausible-looking step + // list for a rotate kernel that does not exist. No caller + // reaches this today: `for_gemv(_, Plain)` already returns + // Err for Escha2T16/Escha3T16 (see + // `escha_types_never_resolve_to_plain`), so this arm is dead + // code, not a load-bearing runtime path. Return an explicit + // error instead of guessing a step list (or panicking) — + // callers already handle `for_gemv`'s Err the same way via + // `let Ok(..) else { continue }`. Replace with the real + // `[EschaRotate, Gemv]`-shaped step list when the H128 + // rotate/GEMV kernels land. + RotationPlan::EschaH128 => Err(DispatchError::UnsupportedVariant { + family: "gemv", + variant: "Plain", + arch: "", + quant: "escha-w2 (no GEMV kernel exists yet)", + }), + _ => Ok(&[PipelineOp::RotateFwht, PipelineOp::Gemv]), }, - Prerotated => &[PipelineOp::Gemv], + Prerotated => Ok(&[PipelineOp::Gemv]), WithResidual => { let steps: &[PipelineOp] = match dtype { MQ4G256 | MQ4G256V2 | MQ2G256V2 | MQ3G256V2 | MQ5G256V2 | MQ6G256V2 @@ -906,7 +971,7 @@ impl KernelKey { ], _ => &[PipelineOp::Gemv, PipelineOp::ResidualAdd], }; - steps + Ok(steps) } WithSwiGLUResidual => { let steps: &[PipelineOp] = match dtype { @@ -920,7 +985,7 @@ impl KernelKey { PipelineOp::ResidualAdd, ], }; - steps + Ok(steps) } } } @@ -958,6 +1023,17 @@ pub fn dtype_needs_rotation(dtype: DType) -> bool { | MFP3G32E8 | MFP2G32E8 | ParoQ4G128 + // Escha-W2. Its plan is `RotationPlan::EschaH128` (see + // `dtype_rotation_plan` above), so it MUST answer true here. + // The two functions disagreeing is precisely the state in which + // one caller applies the H128 pair and another skips it — an + // unrotated activation through a rotated-domain weight, which is + // finite, fluent, wrong output rather than a crash. + // `rotation_plan_matches_legacy_needs_fwht` in dispatch-tests + // enforces the agreement; both types are in its + // `QUANTIZED_DTYPES` list so the enforcement is not vacuous. + | Escha2T16 + | Escha3T16 ) } diff --git a/crates/hipfire-generate/src/ar.rs b/crates/hipfire-generate/src/ar.rs index f3f9082ee7..bd616952f2 100644 --- a/crates/hipfire-generate/src/ar.rs +++ b/crates/hipfire-generate/src/ar.rs @@ -37,6 +37,29 @@ use std::any::Any; use std::io::Write; use std::time::Instant; +/// Stage accounting for the AR decode loop, enabled by +/// `HIPFIRE_DECODE_PROFILE=1`. Statics rather than locals because the +/// detokenise calls sit inside closures handed to the semantic producer. +pub(crate) mod decode_prof { + use std::sync::atomic::{AtomicU64, Ordering}; + pub static NS_FORWARD: AtomicU64 = AtomicU64::new(0); + pub static NS_SAMPLE: AtomicU64 = AtomicU64::new(0); + pub static NS_DETOK: AtomicU64 = AtomicU64::new(0); + pub static DETOK_CALLS: AtomicU64 = AtomicU64::new(0); + pub static DETOK_TOKENS: AtomicU64 = AtomicU64::new(0); + + pub fn on() -> bool { + static ON: std::sync::OnceLock = std::sync::OnceLock::new(); + *ON.get_or_init(|| std::env::var("HIPFIRE_DECODE_PROFILE").as_deref() == Ok("1")) + } + pub fn add(c: &AtomicU64, ns: u128) { + c.fetch_add(ns as u64, Ordering::Relaxed); + } + pub fn get(c: &AtomicU64) -> u64 { + c.load(Ordering::Relaxed) + } +} + #[cfg(feature = "serve-fault-inject")] thread_local! { static FAULT_AFTER_PREFILL_ARMED: std::cell::Cell = @@ -3514,7 +3537,32 @@ pub fn generate( // `while` instead of `for 0..max_tokens` so budget-alert injection // (which increments `generated` beyond the iteration count) can't // push generated past max_tokens: each loop start rechecks the cap. + let dp_t0 = Instant::now(); while generated < max_tokens { + if decode_prof::on() && generated > 0 && generated % 40 == 0 { + let tot = dp_t0.elapsed().as_nanos().max(1) as f64; + let f = decode_prof::get(&decode_prof::NS_FORWARD) as f64; + let sm = decode_prof::get(&decode_prof::NS_SAMPLE) as f64; + let dt = decode_prof::get(&decode_prof::NS_DETOK) as f64; + let calls = decode_prof::get(&decode_prof::DETOK_CALLS); + let dtoks = decode_prof::get(&decode_prof::DETOK_TOKENS); + eprintln!( + "[ar-profile] n={generated} total={:.0}ms ({:.2} tok/s) | \ + forward {:.1}ms/tok ({:.0}%) | sample {:.1}ms/tok ({:.0}%) | \ + detok {:.1}ms/tok ({:.0}%) [{} calls, {} toks re-decoded] | other {:.0}%", + tot / 1e6, + generated as f64 / (tot / 1e9), + f / 1e6 / generated as f64, + 100.0 * f / tot, + sm / 1e6 / generated as f64, + 100.0 * sm / tot, + dt / 1e6 / generated as f64, + 100.0 * dt / tot, + calls, + dtoks, + 100.0 * (tot - f - sm - dt).max(0.0) / tot, + ); + } // Decode-side abort check. Client cancel (Pi 4-min idle // timeout firing while the CLI buffers tokens for tool-call // detection — wire shows zero output until `done`) sends @@ -3577,7 +3625,19 @@ pub fn generate( next_token, QwenArRawCommitDisposition::ClassifiedVisible, ); - let all_bytes = tokenizer.decode_bytes(&streamed_tokens); + let all_bytes = { + let _t = std::time::Instant::now(); + let r = tokenizer.decode_bytes(&streamed_tokens); + if decode_prof::on() { + decode_prof::add(&decode_prof::NS_DETOK, _t.elapsed().as_nanos()); + decode_prof::add(&decode_prof::DETOK_CALLS, 1); + decode_prof::add( + &decode_prof::DETOK_TOKENS, + streamed_tokens.len() as u128, + ); + } + r + }; let new_bytes = all_bytes[prev_fed..].to_vec(); bytes_fed_to_filter = all_bytes.len(); (pos, new_bytes) @@ -3692,7 +3752,19 @@ pub fn generate( || force_answer_latched || max_total_think > 0 { - let raw_so_far = tokenizer.decode_bytes(&streamed_tokens); + let raw_so_far = { + let _t = std::time::Instant::now(); + let r = tokenizer.decode_bytes(&streamed_tokens); + if decode_prof::on() { + decode_prof::add(&decode_prof::NS_DETOK, _t.elapsed().as_nanos()); + decode_prof::add(&decode_prof::DETOK_CALLS, 1); + decode_prof::add( + &decode_prof::DETOK_TOKENS, + streamed_tokens.len() as u128, + ); + } + r + }; let raw_str = std::str::from_utf8(&raw_so_far).unwrap_or(""); let in_think = currently_in_think(raw_str, started_in_think); // Total-think bound (re-arm-proof). Count every think token; at the @@ -3806,7 +3878,19 @@ pub fn generate( m.seq_pos = new_phys; } } - let all_bytes = tokenizer.decode_bytes(&streamed_tokens); + let all_bytes = { + let _t = std::time::Instant::now(); + let r = tokenizer.decode_bytes(&streamed_tokens); + if decode_prof::on() { + decode_prof::add(&decode_prof::NS_DETOK, _t.elapsed().as_nanos()); + decode_prof::add(&decode_prof::DETOK_CALLS, 1); + decode_prof::add( + &decode_prof::DETOK_TOKENS, + streamed_tokens.len() as u128, + ); + } + r + }; let new_bytes = all_bytes[prev_fed..].to_vec(); bytes_fed_to_filter = all_bytes.len(); (pos, new_bytes) @@ -3874,7 +3958,19 @@ pub fn generate( // answer with a system-alert string. Check the raw decoded // text rather than token IDs since tokenizes as a // multi-token sequence in Qwen3.5's vocab. - let raw_so_far = tokenizer.decode_bytes(&streamed_tokens); + let raw_so_far = { + let _t = std::time::Instant::now(); + let r = tokenizer.decode_bytes(&streamed_tokens); + if decode_prof::on() { + decode_prof::add(&decode_prof::NS_DETOK, _t.elapsed().as_nanos()); + decode_prof::add(&decode_prof::DETOK_CALLS, 1); + decode_prof::add( + &decode_prof::DETOK_TOKENS, + streamed_tokens.len() as u128, + ); + } + r + }; let raw_str = std::str::from_utf8(&raw_so_far).unwrap_or(""); let in_think = currently_in_think(raw_str, started_in_think); if !in_think { @@ -3913,7 +4009,7 @@ pub fn generate( &grammar_mask, &mut logits, ); - sampler::sample_cpu(&mut logits, ngram_scope, &cfg) + { let _t = std::time::Instant::now(); let r = sampler::sample_cpu(&mut logits, ngram_scope, &cfg); if decode_prof::on() { decode_prof::add(&decode_prof::NS_SAMPLE, _t.elapsed().as_nanos()); } r } } else { sampler::sample( gpu, @@ -4083,18 +4179,25 @@ pub fn generate( &grammar_mask, &mut logits, ); - sampler::sample_cpu(&mut logits, ngram_scope, &cfg) + { let _t = std::time::Instant::now(); let r = sampler::sample_cpu(&mut logits, ngram_scope, &cfg); if decode_prof::on() { decode_prof::add(&decode_prof::NS_SAMPLE, _t.elapsed().as_nanos()); } r } } else { - sampler::sample( - gpu, - &scratch.logits, - &scratch.sample_buf, - &scratch.repeat_buf, - vocab_size, - ngram_scope, - &cfg, - &mut rng_state, - ) + { + let _t = std::time::Instant::now(); + let r = sampler::sample( + gpu, + &scratch.logits, + &scratch.sample_buf, + &scratch.repeat_buf, + vocab_size, + ngram_scope, + &cfg, + &mut rng_state, + ); + if decode_prof::on() { + decode_prof::add(&decode_prof::NS_SAMPLE, _t.elapsed().as_nanos()); + } + r + } }; if grammar_active { let text = tokenizer.decode(&[next_token]); diff --git a/crates/hipfire-loader/src/lib.rs b/crates/hipfire-loader/src/lib.rs index c099ef2818..b9ba7617c0 100644 --- a/crates/hipfire-loader/src/lib.rs +++ b/crates/hipfire-loader/src/lib.rs @@ -2054,6 +2054,32 @@ fn finish_qwen35_load( } } } + // Third packaging: the head carried as ordinary `mtp.*` tensors inside + // the trunk's own container, which is how the escha converter passes + // upstream's head through. Neither the HFBNDMTP trailer nor a sibling + // `.mtp` exists for those builds, so without this the 849 MB head in + // qwen3.8-27b.escha is downloaded and never used. + if head_opt.is_none() { + match hipfire_arch_qwen35::mtp_head::load_mtp_head_from_trunk( + trunk_path, + ctx.gpu, + physical_cap, + ) { + Ok(Some(h)) => { + eprintln!( + " MTP head loaded (in-trunk mtp.* tensors): n_embd={} vocab={}", + h.config.n_embd, h.config.vocab_size + ); + head_opt = Some(h); + load_err = None; + } + Ok(None) => {} + Err(e) => { + load_err = Some(format!("in-trunk mtp.* load failed: {e}")); + } + } + } + if head_opt.is_none() { if ctx.spec.mtp == Some(true) { return Err(rollback_unfinished_qwen35( diff --git a/crates/hipfire-quantize/Cargo.toml b/crates/hipfire-quantize/Cargo.toml index 2b0b6d756e..e6ecd47208 100644 --- a/crates/hipfire-quantize/Cargo.toml +++ b/crates/hipfire-quantize/Cargo.toml @@ -26,6 +26,7 @@ hipfire-runtime = { path = "../hipfire-runtime", default-features = false } [dev-dependencies] tempfile = "3" +sha2 = "0.10" # CPU-only reader for the overlay round-trip test: HfqFile parses the .hfq # container header/index by name without touching the GPU (builds on CPU). # hipfire-quantize is a binary crate (no lib target), so the overlay diff --git a/crates/hipfire-quantize/src/escha_fold.rs b/crates/hipfire-quantize/src/escha_fold.rs new file mode 100644 index 0000000000..8c11ba1239 --- /dev/null +++ b/crates/hipfire-quantize/src/escha_fold.rs @@ -0,0 +1,237 @@ +//! Fold an escha linear's H128 rotations and diagonals into its weight. +//! +//! An escha linear evaluates +//! +//! ```text +//! xh = RS * H * diag(rin) * x +//! mid = W^T * xh +//! y = RS * diag(rout) * H * mid +//! ``` +//! +//! Every one of those operators except `W` is linear and input-independent, +//! so they collapse into the weight: +//! +//! ```text +//! W_eff[i][o] = RS^2 * rin_i * (H W H)[i][o] * rout_o +//! ``` +//! +//! WHY THIS MATTERS: a folded escha linear is an ORDINARY dense weight. Every +//! fused path in the runtime — FusedQkv, FusedQkvza, gate_up, the batched +//! prefill arms — consumes it untouched, with no escha awareness anywhere in +//! the forward pass. That is the difference between a contained converter +//! change and threading a new weight kind through two multi-thousand-line +//! files. +//! +//! WHAT IT COSTS: the folded matrix is dense, so the 2-bit residency of the +//! trellis code is gone; what survives is escha's quantisation QUALITY, baked +//! into whatever container it is re-quantised to. This is a deliberate +//! trade, not an oversight — see the module docs on the converter flag. +//! +//! NUMERICS: folding SKIPS the reference's fp16 rounding of `xh`, so it is not +//! bit-identical to `escha_ref::expert_linear`. Measured against it on the +//! real 27B (`linear_attn.in_proj_z`, ic=5120 oc=6144): rel_rms 2.868e-4, +//! against 9.192e-5 for the runtime-transform path — same order, and the +//! difference is f32 accumulation across two Hadamard passes rather than a +//! systematic bias. + +/// `1/sqrt(128)`, applied once per H128 — hence `RS*RS` in the fold. +pub const ESCHA_RS: f32 = 0.088_388_347_648; + +/// Unnormalised 128-point Walsh-Hadamard, Sylvester order, in place, over +/// every contiguous 128-lane block of `x`. Byte-for-byte the same butterfly +/// as `escha_ref::h128_inplace`; duplicated rather than shared so the +/// reference stays a reference (G3 gates the GPU H128 against it, and +/// generating one from the other would make that circular). +#[inline] +fn h128_blocks(x: &mut [f32]) { + debug_assert_eq!(x.len() % 128, 0); + for block in x.chunks_exact_mut(128) { + let mut h = 1; + while h < 128 { + let mut i = 0; + while i < 128 { + for j in i..i + h { + let (a, b) = (block[j], block[j + h]); + block[j] = a + b; + block[j + h] = a - b; + } + i += 2 * h; + } + h *= 2; + } + } +} + +/// Blocked transpose `[rows, cols] -> [cols, rows]`. +/// +/// 64x64 tiles: the naive version strides one side by a full row and thrashes +/// for anything the size of a real projection (gate_proj is 5120x17408 = +/// 89 M elements, 356 MB as f32). +fn transpose_blocked(src: &[f32], rows: usize, cols: usize) -> Vec { + const T: usize = 64; + let mut dst = vec![0.0f32; rows * cols]; + for r0 in (0..rows).step_by(T) { + let r1 = (r0 + T).min(rows); + for c0 in (0..cols).step_by(T) { + let c1 = (c0 + T).min(cols); + for r in r0..r1 { + let srow = &src[r * cols..r * cols + cols]; + for c in c0..c1 { + dst[c * rows + r] = srow[c]; + } + } + } + } + dst +} + +/// Fold one escha linear into an ordinary dense weight. +/// +/// `w_bits` is the decoded weight as `escha_ref::reconstruct` returns it: +/// fp16 bits, IN-major `[ic, oc]` (`w[i * oc + o]`). +/// +/// Returns f32 in OUT-major `[oc, ic]` — hipfire's dense convention, and +/// exactly what `quantize_mq6g256v2(data, oc, ic, ..)` and friends expect, so +/// no further shuffling is needed before quantisation. +/// +/// Both Hadamard passes run over CONTIGUOUS lanes: pass A along `o` while the +/// matrix is still in-major, then one blocked transpose, then pass B along `i` +/// which is now contiguous. Doing pass B strided instead (gathering each of +/// `oc` columns at stride `oc`) is correct but cache-hostile enough to be +/// unusable across 400 projections. +pub fn fold_escha_linear( + w_bits: &[u16], + ic: usize, + oc: usize, + rin: &[f32], + rout: &[f32], +) -> Result, String> { + if w_bits.len() != ic * oc { + return Err(format!( + "fold: weight has {} elements, expected ic*oc = {}", + w_bits.len(), + ic * oc + )); + } + if rin.len() != ic || rout.len() != oc { + return Err(format!( + "fold: rin {} (want {ic}), rout {} (want {oc})", + rin.len(), + rout.len() + )); + } + if ic % 128 != 0 || oc % 128 != 0 { + return Err(format!( + "fold: H128 needs both dims a multiple of 128, got ic={ic} oc={oc}" + )); + } + + // in-major [ic, oc], f32. + let mut w: Vec = w_bits + .iter() + .map(|&b| crate::float16::f16_to_f32(b)) + .collect(); + + // Pass A — H along `o`, contiguous within each in-major row. + for row in w.chunks_exact_mut(oc) { + h128_blocks(row); + } + + // -> out-major [oc, ic]. + let mut wt = transpose_blocked(&w, ic, oc); + drop(w); + + // Pass B — H along `i`, now contiguous within each out-major row. + for row in wt.chunks_exact_mut(ic) { + h128_blocks(row); + } + + // Diagonals and RS^2. Row `o`, column `i`. + let rs2 = ESCHA_RS * ESCHA_RS; + for (o, row) in wt.chunks_exact_mut(ic).enumerate() { + let s = rs2 * rout[o]; + for (i, v) in row.iter_mut().enumerate() { + *v *= s * rin[i]; + } + } + Ok(wt) +} + +#[cfg(test)] +mod tests { + use super::*; + + /// H128 is its own inverse up to a factor of 128, which is the cheapest + /// end-to-end check that the butterfly is the right transform and not + /// merely *a* transform. + #[test] + fn h128_is_an_involution_scaled_by_128() { + let mut x: Vec = (0..256).map(|i| ((i * 37) % 91) as f32 - 45.0).collect(); + let orig = x.clone(); + h128_blocks(&mut x); + h128_blocks(&mut x); + for (a, b) in x.iter().zip(orig.iter()) { + assert!((a - b * 128.0).abs() < 1e-3, "{a} vs {}", b * 128.0); + } + } + + /// The folded weight must reproduce the un-folded evaluation. Built on a + /// tiny random-ish linear so the whole path is checked without a + /// checkpoint: fold, then a plain matmul, against transform-matmul- + /// transform done longhand. + #[test] + fn folded_weight_matches_the_transform_pipeline() { + let (ic, oc) = (128usize, 256usize); + let w_bits: Vec = (0..ic * oc) + .map(|n| crate::float16::f32_to_f16((((n * 31) % 17) as f32 - 8.0) / 8.0)) + .collect(); + let rin: Vec = (0..ic).map(|i| 1.0 + (i % 5) as f32 * 0.1).collect(); + let rout: Vec = (0..oc).map(|o| 0.5 + (o % 7) as f32 * 0.05).collect(); + let x: Vec = (0..ic).map(|i| ((i % 13) as f32 - 6.0) / 6.0).collect(); + + // Longhand: xh = RS*H*diag(rin)*x ; mid = W^T xh ; y = RS*diag(rout)*H*mid + let mut xh: Vec = x.iter().zip(&rin).map(|(a, b)| a * b).collect(); + h128_blocks(&mut xh); + for v in xh.iter_mut() { + *v *= ESCHA_RS; + } + let mut mid = vec![0.0f32; oc]; + for i in 0..ic { + let a = xh[i]; + for o in 0..oc { + mid[o] += a * crate::float16::f16_to_f32(w_bits[i * oc + o]); + } + } + h128_blocks(&mut mid); + let want: Vec = mid + .iter() + .zip(&rout) + .map(|(m, r)| m * ESCHA_RS * r) + .collect(); + + // Folded: one plain matmul on the out-major weight. + let wf = fold_escha_linear(&w_bits, ic, oc, &rin, &rout).unwrap(); + let got: Vec = (0..oc) + .map(|o| { + let row = &wf[o * ic..(o + 1) * ic]; + row.iter().zip(&x).map(|(w, xi)| w * xi).sum() + }) + .collect(); + + let num: f64 = got + .iter() + .zip(&want) + .map(|(a, b)| ((a - b) as f64).powi(2)) + .sum(); + let den: f64 = want.iter().map(|b| (*b as f64).powi(2)).sum(); + let rel = (num / den.max(1e-30)).sqrt(); + assert!(rel < 1e-5, "folded vs longhand rel_rms {rel:.3e}"); + } + + #[test] + fn rejects_dims_that_are_not_h128_shaped() { + let e = fold_escha_linear(&[0u16; 100 * 128], 100, 128, &[1.0; 100], &[1.0; 128]) + .unwrap_err(); + assert!(e.contains("multiple of 128"), "{e}"); + } +} diff --git a/crates/hipfire-quantize/src/escha_ref.rs b/crates/hipfire-quantize/src/escha_ref.rs new file mode 100644 index 0000000000..598e967102 --- /dev/null +++ b/crates/hipfire-quantize/src/escha_ref.rs @@ -0,0 +1,729 @@ +//! Portable CPU reference for the escha codec — the numerical oracle for +//! every GPU kernel in this port. +//! +//! Ported from `escha_mlx/ref.py` (EschaLabs/escha-mlx, Apache-2.0), which +//! declares itself "the semantic contract for every Metal kernel in this +//! package". Rounding points are deliberate; do not "simplify" them. + +use crate::float16::f16_to_f32; + +/// 1/sqrt(128) — the exact f32 constant the format pins. +pub const RS: f32 = 0.088388347648; + +/// Round `v` to fp16 bits using round-to-nearest-even. +/// +/// This is the one RNE encode site in the module. `crate::float16::f32_to_f16` +/// is documented to truncate ("Hipfire's historical truncating conversion") +/// to keep existing HFQ encoder output byte-stable, whereas every `f16(...)` +/// in the escha contract is RNE — truncating flips the low bit on roughly +/// half of all 65536 codebook states (state 3 alone: truncation gives +/// 0x3ab7, the published/measured value is 0x3ab8). `half` is already a +/// workspace dependency used inside `float16.rs` itself, so this is not a +/// new dependency, just reaching past the crate's non-RNE convenience +/// wrapper for the one place where RNE is the spec. +#[inline] +pub fn f16_rne(v: f32) -> u16 { + half::f16::from_f32(v).to_bits() +} + +/// Decode one 16-bit trellis state to fp16 **bits** via the cbA codebook. +/// +/// `decode(x) = f16_lo(r) + f16_hi(r)` with fp16 RNE addition, where +/// `r = ((x * 0xCBAC1FED) & 0x8FFF8FFF) ^ 0x3B603B60` in 32-bit arithmetic. +/// +/// Adding in f32 and rounding once is exactly an fp16 RNE add: the exact sum +/// of two fp16 values is always representable in f32, so the single rounding +/// here is the correctly-rounded fp16 result. +/// +/// There are 65536 reachable values, so a lookup table would be 128 KB and +/// will not fit gfx1151's 64 KB LDS. This is five integer/FP ops and no +/// memory traffic — keep it that way in the kernels. +#[inline] +pub fn cba_decode(state: u16) -> u16 { + let r = ((state as u32).wrapping_mul(0xCBAC_1FED) & 0x8FFF_8FFF) ^ 0x3B60_3B60; + let lo = f16_to_f32((r & 0xFFFF) as u16); + let hi = f16_to_f32((r >> 16) as u16); + f16_rne(lo + hi) +} + +/// The 8 states lane `lane` owns, K=2. `words` is the tile's 16 u32. +/// +/// DELIBERATE DUPLICATION: `kernels/src/escha_decode_tiles.hip` implements +/// this same lane maths independently. That is the G2 gate — the GPU decode +/// is asserted bit-exact against this one. Generating either from the other, +/// or sharing a source, would make G2 circular: both paths could be wrong in +/// exactly the same way and still agree. Two independent implementations of +/// a published spec is the point. Do not deduplicate. +pub fn decode8_k2(words: &[u32; 16], lane: usize) -> [u16; 8] { + let t_off = lane * 8; + let i1 = t_off >> 4; + let i0 = (i1 + 15) & 15; + let merged = ((words[i0] as u64) << 32) | words[i1] as u64; + let shift = ((!t_off) & 8) << 1; // 16 for even lanes, 0 for odd + let w = ((merged >> shift) & 0xFFFF_FFFF) as u32; + let mut out = [0u16; 8]; + for (j, o) in out.iter_mut().enumerate() { + *o = (w >> (2 * (7 - j))) as u16; + } + out +} + +/// The 8 states lane `lane` owns, K=3. `words` is the tile's 24 u32. +/// +/// Structurally different from K=2 — 24 words, a computed bit offset, and a +/// modular wrap. Do not attempt to unify the two. +pub fn decode8_k3(words: &[u32; 24], lane: usize) -> [u16; 8] { + const BITS: usize = 3; + let t_off = lane * 8; + let b1 = (t_off + 257) * BITS; + let b0 = b1 - 16; + let b2 = b1 + BITS * 7; + let i0 = b0 >> 5; + let i2 = (b2 - 1) >> 5; + let s2 = ((i2 + 1) << 5) - b2; + let merged = ((words[i0 % 24] as u64) << 32) | words[i2 % 24] as u64; + let w7 = (merged >> s2) & 0xFFFF_FFFF; + let w3 = (merged >> (s2 + BITS * 4)) & 0xFFFF_FFFF; + [ + (w3 >> 9) as u16, + (w3 >> 6) as u16, + (w3 >> 3) as u16, + w3 as u16, + (w7 >> 9) as u16, + (w7 >> 6) as u16, + (w7 >> 3) as u16, + w7 as u16, + ] +} + +/// `(row, col)` inside the 16x16 tile for each of the lane's 8 values. +/// +/// This permutation is the single easiest thing to get subtly wrong: a wrong +/// shuffle still yields a full-rank, plausible weight matrix. It is gated +/// directly on golden vectors, never on end-to-end coherence. +pub fn lane_positions(lane: usize) -> [(usize, usize); 8] { + let l0 = lane & !4; + let c_off = (lane >> 2) & 1; + let mut out = [(0usize, 0usize); 8]; + for (j, o) in out.iter_mut().enumerate() { + let fi = j >> 1; + let row = (lane & 3) * 2 + (j & 1) + (fi & 1) * 8; + let col = 2 * ((l0 >> 3) + if j >= 4 { 4 } else { 0 }) + c_off; + *o = (row, col); + } + out +} + +/// Decode one packed tile (`16*K` i16) to a 16x16 fp16-bit tile, row-major. +pub fn decode_tile(tile: &[i16], k: usize) -> [u16; 256] { + debug_assert_eq!(tile.len(), 16 * k); + let mut words = [0u32; 24]; + for (i, w) in words.iter_mut().enumerate().take(8 * k) { + *w = (tile[2 * i] as u16 as u32) | ((tile[2 * i + 1] as u16 as u32) << 16); + } + let mut out = [0u16; 256]; + for lane in 0..32 { + let states = match k { + 2 => { + let mut w16 = [0u32; 16]; + w16.copy_from_slice(&words[..16]); + decode8_k2(&w16, lane) + } + 3 => decode8_k3(&words, lane), + _ => panic!("unsupported escha K={k}"), + }; + for (j, (r, c)) in lane_positions(lane).into_iter().enumerate() { + out[r * 16 + c] = cba_decode(states[j]); + } + } + out +} + +/// Packed `(in/16, out/16, 16K)` i16 -> `(in, out)` fp16 bits, row-major. +pub fn reconstruct(code: &[i16], in_features: usize, out_features: usize, k: usize) -> Vec { + let (tk, tn) = (in_features / 16, out_features / 16); + assert_eq!(code.len(), tk * tn * 16 * k, "escha code length mismatch"); + let mut out = vec![0u16; in_features * out_features]; + for kt in 0..tk { + for nt in 0..tn { + let base = (kt * tn + nt) * 16 * k; + let tile = decode_tile(&code[base..base + 16 * k], k); + for r in 0..16 { + let dst = (kt * 16 + r) * out_features + nt * 16; + out[dst..dst + 16].copy_from_slice(&tile[r * 16..r * 16 + 16]); + } + } + } + out +} + +/// Unnormalised 128-point Walsh-Hadamard (Sylvester / natural order), applied +/// independently to each contiguous 128-element block of `x`. +/// +/// `x.len()` must be a multiple of 128. Every dimension in both checkpoints +/// satisfies this (512, 1024, 2048, 5120, 6144, 10240, 17408 are all +/// multiples of 128), including the gate|up split point at 512 — so a block +/// never straddles the gate/up boundary. +pub fn h128_inplace(x: &mut [f32]) { + assert_eq!( + x.len() % 128, + 0, + "H128 needs a multiple of 128, got {}", + x.len() + ); + for block in x.chunks_exact_mut(128) { + let mut h = 1; + while h < 128 { + let mut i = 0; + while i < 128 { + for j in i..i + h { + let (a, b) = (block[j], block[j + h]); + block[j] = a + b; + block[j + h] = a - b; + } + i += 2 * h; + } + h *= 2; + } + } +} + +/// `xh = f16( H128(x * rin) * RS )`. Returns fp16 bits. +pub fn input_transform(x: &[f32], rin: &[f32]) -> Vec { + let ic = rin.len(); + assert_eq!(x.len() % ic, 0); + let mut buf: Vec = x + .iter() + .zip(rin.iter().cycle()) + .map(|(a, b)| a * b) + .collect(); + for row in buf.chunks_exact_mut(ic) { + h128_inplace(row); + } + buf.iter().map(|v| f16_rne(v * RS)).collect() +} + +/// `y = f16( H128(mid) * RS * rout )`. Returns fp16 bits. +pub fn output_transform(mid: &[f32], rout: &[f32]) -> Vec { + let oc = rout.len(); + assert_eq!(mid.len() % oc, 0); + let mut buf = mid.to_vec(); + for row in buf.chunks_exact_mut(oc) { + h128_inplace(row); + } + buf.iter() + .zip(rout.iter().cycle()) + .map(|(v, s)| f16_rne(v * RS * s)) + .collect() +} + +/// Fold the optional end-to-end scales into the transform vectors. +/// +/// `s_in` multiplies the activation at exactly the point `rin` does, and +/// `s_out` at exactly the point `rout` does, so the pair collapses with no new +/// kernel and no new tensor. Folding keeps both products in f32 and rounds +/// once — one rounding point FEWER than applying the scales separately. +/// `None` returns that vector unchanged (as f32), which is the path MoE +/// exports and end-to-end-free exports both take. +pub fn fold_scales( + rin: &[u16], + rout: &[u16], + s_in: Option<&[f32]>, + s_out: Option<&[f32]>, +) -> (Vec, Vec) { + let mut ri: Vec = rin.iter().map(|&b| f16_to_f32(b)).collect(); + let mut ro: Vec = rout.iter().map(|&b| f16_to_f32(b)).collect(); + if let Some(s) = s_in { + assert_eq!(s.len(), ri.len()); + for (a, b) in ri.iter_mut().zip(s) { + *a *= b; + } + } + if let Some(s) = s_out { + assert_eq!(s.len(), ro.len()); + for (a, b) in ro.iter_mut().zip(s) { + *a *= b; + } + } + (ri, ro) +} + +/// Full single-expert linear for one token: `x [ic] -> [oc]` fp16 bits. +/// +/// `w_bits` is the decoded bare weight, row-major `[ic, oc]` — decode it once +/// with `reconstruct` and reuse it. `ref.moe_block` in the Python original +/// re-decodes per (token, slot), which is 128 full tile decodes for an +/// 8-token fixture; do not reproduce that. +pub fn expert_linear(x: &[f32], w_bits: &[u16], rin: &[f32], rout: &[f32]) -> Vec { + let (ic, oc) = (rin.len(), rout.len()); + assert_eq!(x.len(), ic); + assert_eq!(w_bits.len(), ic * oc); + let xh = input_transform(x, rin); + let mut mid = vec![0.0f32; oc]; + for i in 0..ic { + let a = f16_to_f32(xh[i]); + let row = &w_bits[i * oc..(i + 1) * oc]; + for (m, &wb) in mid.iter_mut().zip(row) { + *m += a * f16_to_f32(wb); + } + } + output_transform(&mid, rout) +} + +/// `silu(g) * u` on the fp16-rounded merged output; gate is the first half. +pub fn swiglu(gate_up_bits: &[u16], inter: usize) -> Vec { + assert_eq!(gate_up_bits.len(), 2 * inter); + let mut out = Vec::with_capacity(inter); + for i in 0..inter { + let g = f16_to_f32(gate_up_bits[i]); + let s = f16_to_f32(f16_rne(g / (1.0 + (-g).exp()))); + out.push(f16_rne(s * f16_to_f32(gate_up_bits[inter + i]))); + } + out +} + +/// `y = f16( x @ f16(w8 * scale)^T )`. `w8` is `[oc, ic]`, `scale` is `[oc]` +/// fp16 bits — Escha's int8 is per-output-row, not per-block. +pub fn w8a16(x: &[f32], w8: &[i8], scale: &[u16], oc: usize, ic: usize) -> Vec { + assert_eq!(w8.len(), oc * ic); + assert_eq!(scale.len(), oc); + assert_eq!(x.len(), ic); + let mut out = Vec::with_capacity(oc); + for o in 0..oc { + let s = f16_to_f32(scale[o]); + let mut acc = 0.0f32; + for i in 0..ic { + acc += x[i] * f16_to_f32(f16_rne(w8[o * ic + i] as f32 * s)); + } + out.push(f16_rne(acc)); + } + out +} + +#[cfg(test)] +mod tests { + use super::*; + use std::path::PathBuf; + + fn data(name: &str) -> PathBuf { + PathBuf::from(env!("CARGO_MANIFEST_DIR")) + .join("tests/data/escha") + .join(name) + } + + fn sha256_hex(bytes: &[u8]) -> String { + use sha2::{Digest, Sha256}; + let mut h = Sha256::new(); + h.update(bytes); + h.finalize().iter().map(|b| format!("{b:02x}")).collect() + } + + /// The codebook is a pure function of the 16-bit state. These eight values + /// were computed from the published constants and pin the hash, the + /// masking, and the fp16 RNE add all at once. + #[test] + fn cba_decode_matches_published_constants() { + let want: [u16; 8] = [ + 0x3f60, 0x304e, 0xba13, 0x3ab8, 0x3952, 0xb75f, 0xbea4, 0xbc71, + ]; + for (state, &bits) in want.iter().enumerate() { + assert_eq!(cba_decode(state as u16), bits, "state {state}"); + } + } + + #[test] + fn reconstruct_k2_matches_golden() { + let raw = std::fs::read(data("packed_gu_e0_k2.i16")).unwrap(); + let code: Vec = raw + .chunks_exact(2) + .map(|c| i16::from_le_bytes([c[0], c[1]])) + .collect(); + assert_eq!(code.len(), 262144); + let out = reconstruct(&code, 2048, 1024, 2); + assert_eq!(out.len(), 2048 * 1024); + let bytes: Vec = out.iter().flat_map(|v| v.to_le_bytes()).collect(); + assert_eq!( + sha256_hex(&bytes), + "51ddde9a07613aafcc9f5db79702349d19e18357d9fb910f48b38eeea028dcab", + "decoded K=2 tensor does not match expected_gu_e0_k2.f16" + ); + } + + #[test] + fn reconstruct_k3_matches_golden() { + let raw = std::fs::read(data("packed_down_e0_k3.i16")).unwrap(); + let code: Vec = raw + .chunks_exact(2) + .map(|c| i16::from_le_bytes([c[0], c[1]])) + .collect(); + assert_eq!(code.len(), 196608); + let out = reconstruct(&code, 512, 2048, 3); + assert_eq!(out.len(), 512 * 2048); + let bytes: Vec = out.iter().flat_map(|v| v.to_le_bytes()).collect(); + assert_eq!( + sha256_hex(&bytes), + "51c99817d00282f4aa9d618140eb4503083e1238cb59dd71a670aa4c320f7438", + "decoded K=3 tensor does not match expected_down_e0_k3.f16" + ); + } + + /// Every one of the 256 tile slots must be written exactly once. A + /// permutation bug that drops and duplicates slots still produces a + /// full-rank, plausible-looking matrix, so check the permutation directly. + #[test] + fn lane_positions_is_a_permutation_of_the_tile() { + let mut seen = [0u8; 256]; + for lane in 0..32 { + for (r, c) in lane_positions(lane) { + assert!(r < 16 && c < 16, "lane {lane} -> ({r},{c})"); + seen[r * 16 + c] += 1; + } + } + assert!( + seen.iter().all(|&n| n == 1), + "lane_positions is not a bijection" + ); + } + + /// H128 is its own inverse up to a factor of 128. That is necessary but + /// NOT sufficient: a wrong butterfly order is also self-inverse and would + /// pass this alone. The Hadamard-of-a-basis-vector check below pins the + /// actual transform. + #[test] + fn h128_roundtrip_scales_by_128() { + let mut x: Vec = (0..256).map(|i| (i as f32 * 0.37).sin()).collect(); + let orig = x.clone(); + h128_inplace(&mut x); + h128_inplace(&mut x); + for (a, b) in x.iter().zip(orig.iter()) { + assert!((a - b * 128.0).abs() < 1e-2, "{a} vs {}", b * 128.0); + } + } + + /// H128 applied to e_0 must give all ones (Sylvester, unnormalised). + /// Applied to e_1 it must give the alternating +1/-1 pattern of row 1. + #[test] + fn h128_matches_sylvester_order() { + let mut e0 = vec![0.0f32; 128]; + e0[0] = 1.0; + h128_inplace(&mut e0); + assert!(e0.iter().all(|&v| v == 1.0), "row 0 must be all ones"); + + let mut e1 = vec![0.0f32; 128]; + e1[1] = 1.0; + h128_inplace(&mut e1); + for (i, &v) in e1.iter().enumerate() { + let want = if i % 2 == 0 { 1.0 } else { -1.0 }; + assert_eq!(v, want, "index {i}"); + } + } + + /// Blocks are independent: H128 must never mix across a 128 boundary. + #[test] + fn h128_does_not_mix_across_blocks() { + let mut x = vec![0.0f32; 256]; + x[0] = 1.0; + h128_inplace(&mut x); + assert!(x[..128].iter().all(|&v| v == 1.0)); + assert!( + x[128..].iter().all(|&v| v == 0.0), + "second block was contaminated" + ); + } + + /// MoE exports ship all-ones s_in/s_out; dense exports ship real values; + /// and an export without the end-to-end stage ships neither. All three + /// must go through one code path. + #[test] + fn fold_scales_handles_absent_scales() { + let rin = [f16_rne(2.0), f16_rne(-3.0)]; + let rout = [f16_rne(0.5)]; + let (a, b) = fold_scales(&rin, &rout, None, None); + assert_eq!(a, vec![2.0, -3.0]); + assert_eq!(b, vec![0.5]); + let (c, d) = fold_scales(&rin, &rout, Some(&[3.0, 2.0]), Some(&[4.0])); + assert_eq!(c, vec![6.0, -6.0]); + assert_eq!(d, vec![2.0]); + } + + /// Mixed case: `s_in` present, `s_out` absent. `fold_scales` has four + /// (s_in, s_out) combinations; the all-Some and all-None corners are + /// covered above, but a branch that only handles the symmetric cases + /// would still pass those. This and the next test pin the two mixed + /// corners. + #[test] + fn fold_scales_handles_s_in_only() { + let rin = [f16_rne(2.0), f16_rne(-3.0)]; + let rout = [f16_rne(0.5)]; + let (a, b) = fold_scales(&rin, &rout, Some(&[3.0, 2.0]), None); + assert_eq!(a, vec![6.0, -6.0]); + assert_eq!( + b, + vec![0.5], + "rout must pass through unscaled when s_out is None" + ); + } + + /// Mixed case: `s_in` absent, `s_out` present. + #[test] + fn fold_scales_handles_s_out_only() { + let rin = [f16_rne(2.0), f16_rne(-3.0)]; + let rout = [f16_rne(0.5)]; + let (a, b) = fold_scales(&rin, &rout, None, Some(&[4.0])); + assert_eq!( + a, + vec![2.0, -3.0], + "rin must pass through unscaled when s_in is None" + ); + assert_eq!(b, vec![2.0]); + } + + /// `input_transform` is specified as `f16( H128(x * rin) * RS )` — scale + /// applied BEFORE the Hadamard transform. H128 mixes elements within a + /// block, so for a non-constant `rin` that is a materially different + /// operation from scaling after the transform. This is the exact + /// ordering contract the GPU kernels get gated against, so pin it + /// directly against an independently-built expected value, and prove + /// the chosen inputs actually discriminate the two orderings. + #[test] + fn input_transform_scales_before_not_after() { + let ic = 128; + let rin: Vec = (0..ic).map(|i| 1.0 + (i as f32) * 0.05).collect(); + let x: Vec = (0..ic).map(|i| ((i as f32) * 0.13).cos()).collect(); + + // Correct: scale THEN transform THEN RS. + let mut correct: Vec = x.iter().zip(rin.iter()).map(|(a, b)| a * b).collect(); + h128_inplace(&mut correct); + let correct_bits: Vec = correct.iter().map(|v| f16_rne(v * RS)).collect(); + + // Wrong: transform THEN scale THEN RS — the order a later edit + // could plausibly swap to. + let mut wrong = x.clone(); + h128_inplace(&mut wrong); + let wrong_bits: Vec = wrong + .iter() + .zip(rin.iter()) + .map(|(v, s)| f16_rne(v * s * RS)) + .collect(); + + assert_ne!( + correct_bits, wrong_bits, + "chosen rin/x must make the two orderings diverge, or this test proves nothing" + ); + + assert_eq!(input_transform(&x, &rin), correct_bits); + } + + /// `output_transform` is specified as `f16( H128(mid) * RS * rout )` — + /// scale applied AFTER the Hadamard transform, the mirror image of + /// `input_transform`'s contract. Same reasoning as above: a + /// non-constant `rout` makes pre- and post-transform scaling diverge. + #[test] + fn output_transform_scales_after_not_before() { + let oc = 128; + let rout: Vec = (0..oc).map(|i| 1.0 + (i as f32) * 0.05).collect(); + let mid: Vec = (0..oc).map(|i| ((i as f32) * 0.19).sin()).collect(); + + // Correct: transform THEN RS*scale. + let mut correct = mid.clone(); + h128_inplace(&mut correct); + let correct_bits: Vec = correct + .iter() + .zip(rout.iter()) + .map(|(v, s)| f16_rne(v * RS * s)) + .collect(); + + // Wrong: scale THEN transform THEN RS. + let mut wrong: Vec = mid.iter().zip(rout.iter()).map(|(a, b)| a * b).collect(); + h128_inplace(&mut wrong); + let wrong_bits: Vec = wrong.iter().map(|v| f16_rne(v * RS)).collect(); + + assert_ne!( + correct_bits, wrong_bits, + "chosen rout/mid must make the two orderings diverge, or this test proves nothing" + ); + + assert_eq!(output_transform(&mid, &rout), correct_bits); + } + + /// `input_transform` broadcasts `rin` across rows via `.iter().cycle()`. + /// A multi-row input must apply the exact same per-channel scale to row + /// 2 as row 1, and the two rows must transform independently (H128 + /// never mixes across the 128 boundary — see + /// `h128_does_not_mix_across_blocks` — so row 2 must not see row 1's + /// data either). The expected vector is built by hand, per row, rather + /// than by calling `input_transform`. + #[test] + fn input_transform_broadcasts_scale_across_rows() { + let ic = 128; + let rin: Vec = (0..ic).map(|i| 1.0 + (i as f32) * 0.03).collect(); + let row0: Vec = (0..ic).map(|i| ((i as f32) * 0.11).sin()).collect(); + let row1: Vec = row0.iter().map(|v| v + 10.0).collect(); + let mut x = row0.clone(); + x.extend_from_slice(&row1); + + let mut expected = Vec::with_capacity(2 * ic); + for row in [&row0, &row1] { + let mut buf: Vec = row.iter().zip(rin.iter()).map(|(a, b)| a * b).collect(); + h128_inplace(&mut buf); + expected.extend(buf.iter().map(|v| f16_rne(v * RS))); + } + + let actual = input_transform(&x, &rin); + assert_eq!( + actual, expected, + "row 2 must see the same per-channel rin as row 1" + ); + assert_ne!( + actual[..ic], + actual[ic..], + "rows must transform independently, not collapse to the same output" + ); + } + + /// Mirror of the above for `output_transform`'s `rout` broadcast. + #[test] + fn output_transform_broadcasts_scale_across_rows() { + let oc = 128; + let rout: Vec = (0..oc).map(|i| 1.0 + (i as f32) * 0.03).collect(); + let row0: Vec = (0..oc).map(|i| ((i as f32) * 0.17).cos()).collect(); + let row1: Vec = row0.iter().map(|v| v + 5.0).collect(); + let mut mid = row0.clone(); + mid.extend_from_slice(&row1); + + let mut expected = Vec::with_capacity(2 * oc); + for row in [&row0, &row1] { + let mut buf = row.clone(); + h128_inplace(&mut buf); + expected.extend( + buf.iter() + .zip(rout.iter()) + .map(|(v, s)| f16_rne(v * RS * s)), + ); + } + + let actual = output_transform(&mid, &rout); + assert_eq!( + actual, expected, + "row 2 must see the same per-channel rout as row 1" + ); + assert_ne!( + actual[..oc], + actual[oc..], + "rows must transform independently, not collapse to the same output" + ); + } + + /// A pruned output channel must be EXACTLY zero, not approximately. + /// gate_up.rout carries a per-expert prune mask (design §1.2): on the + /// shipped layer-0 expert 0, 560 of 1024 channels are hard zeros. A kernel + /// that "optimises away" the zero multiply must preserve exact zero. + #[test] + fn zero_rout_gives_exactly_zero_output() { + let ic = 128; + let oc = 128; + let w: Vec = (0..ic * oc) + .map(|i| f16_rne((i % 7) as f32 - 3.0)) + .collect(); + let rin = vec![1.0f32; ic]; + let mut rout = vec![1.0f32; oc]; + rout[3] = 0.0; + rout[57] = 0.0; + let x: Vec = (0..ic).map(|i| (i as f32 * 0.11).cos()).collect(); + let y = expert_linear(&x, &w, &rin, &rout); + assert_eq!(y.len(), oc); + assert_eq!( + f16_to_f32(y[3]), + 0.0, + "pruned channel 3 must be exactly zero" + ); + assert_eq!( + f16_to_f32(y[57]), + 0.0, + "pruned channel 57 must be exactly zero" + ); + assert!(y + .iter() + .enumerate() + .any(|(i, &v)| i != 3 && i != 57 && f16_to_f32(v) != 0.0)); + } + + /// `expert_linear` is a plain matmul (`mid = xh_f32 @ W_f32`) and this + /// module is the numerical ORACLE that GPU kernels are gated against. + /// Skipping a zero activation as a "the term is zero anyway" shortcut is + /// only valid if every weight it would multiply is finite: IEEE-754 says + /// `0.0 * NaN = NaN`, so a real matmul must let a non-finite weight + /// contaminate the output even where the paired activation is exactly + /// zero. Silently substituting 0 there would mask a corrupted decode + /// instead of surfacing it. + /// + /// An all-zero `x` makes every activation exactly zero + /// (`input_transform`: `f16(H128(0 * rin) * RS) == 0`), so this poisons + /// one weight with NaN and asserts the oracle still reports NaN rather + /// than swallowing it to 0.0. + #[test] + fn expert_linear_propagates_nan_weight_through_zero_activation() { + let ic = 128; + let oc = 128; + let rin = vec![1.0f32; ic]; + let rout = vec![1.0f32; oc]; + let x = vec![0.0f32; ic]; // -> every activation channel is exactly 0.0 + let mut w: Vec = (0..ic * oc) + .map(|i| f16_rne((i % 7) as f32 - 3.0)) + .collect(); + w[5 * oc + 10] = f16_rne(f32::NAN); // one poisoned weight, row 5 col 10 + let y = expert_linear(&x, &w, &rin, &rout); + assert!( + y.iter().any(|&b| f16_to_f32(b).is_nan()), + "0.0 * NaN must propagate as NaN somewhere in the output, not be skipped to \ + all-zero: a faithful matmul cannot discard a NaN weight just because the paired \ + activation happens to be zero" + ); + } + + /// SwiGLU splits the f16-ROUNDED merged output at the halfway point, gate + /// first. Rounding before the split is part of the contract. + /// + /// The first case (gate=[0,0]) does NOT by itself discriminate a + /// gate/up swap: silu(0) == 0 zeroes the output whichever half is + /// treated as gate, so it only pins `silu(0) == 0`. The second case + /// (gate=[10,10], up=[2,3]) is the one that actually catches a swap — + /// silu(large) ~= large makes the output track gate*up, so swapping + /// gate and up would swap which operand tracks toward ~1 and change + /// the result. + #[test] + fn swiglu_uses_gate_first_half() { + let inter = 2; + // gate = [0, 0], up = [5, 7]; silu(0) == 0 so both outputs are zero. + let gu: Vec = [0.0, 0.0, 5.0, 7.0].iter().map(|&v| f16_rne(v)).collect(); + let h = swiglu(&gu, inter); + assert_eq!(h.len(), inter); + assert_eq!(f16_to_f32(h[0]), 0.0); + assert_eq!(f16_to_f32(h[1]), 0.0); + // gate = [large, large] -> silu(x) ~ x, so out ~ gate*up. + let gu2: Vec = [10.0, 10.0, 2.0, 3.0].iter().map(|&v| f16_rne(v)).collect(); + let h2 = swiglu(&gu2, inter); + assert!( + (f16_to_f32(h2[0]) - 20.0).abs() < 0.1, + "{}", + f16_to_f32(h2[0]) + ); + assert!( + (f16_to_f32(h2[1]) - 30.0).abs() < 0.2, + "{}", + f16_to_f32(h2[1]) + ); + } + + /// Escha's int8 is per-output-ROW: y = f16(x @ f16(w8*scale)^T). + #[test] + fn w8a16_applies_per_row_scale() { + let (ic, oc) = (4, 2); + let w8: Vec = vec![1, 2, 3, 4, 5, 6, 7, 8]; + let scale: Vec = vec![f16_rne(0.5), f16_rne(2.0)]; + let x = vec![1.0f32, 1.0, 1.0, 1.0]; + let y = w8a16(&x, &w8, &scale, oc, ic); + assert_eq!(f16_to_f32(y[0]), 5.0); // (1+2+3+4)*0.5 + assert_eq!(f16_to_f32(y[1]), 52.0); // (5+6+7+8)*2.0 + } +} diff --git a/crates/hipfire-quantize/src/hfq.rs b/crates/hipfire-quantize/src/hfq.rs index 860e765107..342cd8ee84 100644 --- a/crates/hipfire-quantize/src/hfq.rs +++ b/crates/hipfire-quantize/src/hfq.rs @@ -87,6 +87,8 @@ impl QuantType { 49 => Some(Self::MQ3G256V2), 50 => Some(Self::MQ2G256V2), 51 => Some(Self::MQ2G256LloydU), + 42 => Some(Self::ESCHA2T16), + 43 => Some(Self::ESCHA3T16), _ => None, } } @@ -235,6 +237,11 @@ pub(crate) enum QuantType { /// never indexed. 2.25 bpw. `K % 256 == 0`. /// See `docs/design/2026-08-22-maple-preview-20b-a1b.md`. MQ2G256LloydU = 51, + /// Escha-W2 trellis, K=2, 16x16 tile, cbA hash codebook (2.00 bpw). + /// Codes are stored verbatim from the source safetensors. + ESCHA2T16 = 42, + /// Escha-W2 trellis, K=3, 16x16 tile, cbA hash codebook (3.00 bpw). + ESCHA3T16 = 43, } /// Per-tensor precision level assigned by the K-map pre-pass. @@ -825,4 +832,15 @@ mod maple_dtype_tests { } } } + + /// from_u8 and the enum discriminants must agree — the doc comment on + /// from_u8 makes this a contract, and a drifted pair silently mislabels + /// every tensor written after it. + #[test] + fn escha_quant_types_round_trip() { + assert_eq!(QuantType::from_u8(42), Some(QuantType::ESCHA2T16)); + assert_eq!(QuantType::from_u8(43), Some(QuantType::ESCHA3T16)); + assert_eq!(QuantType::ESCHA2T16 as u8, 42); + assert_eq!(QuantType::ESCHA3T16 as u8, 43); + } } diff --git a/crates/hipfire-quantize/src/lib.rs b/crates/hipfire-quantize/src/lib.rs index 5b4ad34eda..3be78312c8 100644 --- a/crates/hipfire-quantize/src/lib.rs +++ b/crates/hipfire-quantize/src/lib.rs @@ -1,5 +1,7 @@ //! Shared model-format helpers used by the quantizer binaries. +pub mod escha_fold; +pub mod escha_ref; pub mod float16; pub mod gptq; pub mod hessian_io; diff --git a/crates/hipfire-quantize/src/main.rs b/crates/hipfire-quantize/src/main.rs index f68497c09c..853b222bee 100644 --- a/crates/hipfire-quantize/src/main.rs +++ b/crates/hipfire-quantize/src/main.rs @@ -24,6 +24,7 @@ mod maple; mod model_filter; mod pipeline; mod pipeline_deepseek; +mod pipeline_escha; mod pipeline_gguf; mod pipeline_maple; mod quant_e8; diff --git a/crates/hipfire-quantize/src/pipeline.rs b/crates/hipfire-quantize/src/pipeline.rs index 950ec402c2..80e60bc1bf 100644 --- a/crates/hipfire-quantize/src/pipeline.rs +++ b/crates/hipfire-quantize/src/pipeline.rs @@ -3038,6 +3038,20 @@ fn handle_early_special_formats(args: &QuantizeArgs) -> bool { run_qwen3_dspark(args); return true; } + // ── escha: EschaLabs Escha-W2 trellis checkpoint -> .hfq ──────────────── + // Input is the safetensors DIRECTORY (config.json + shards). Dispatches + // on quant_method inside config.json's quantization_config (escha = + // dense, arch 5; eschamoe = MoE, arch 6) — see pipeline_escha.rs. + // hipfire-quantize --format escha --input --output + if format == "escha" || format == "escha-w2" || format == "eschamoe" { + if let Err(e) = + crate::pipeline_escha::convert_escha(Path::new(input_dir), Path::new(output_path)) + { + eprintln!("error: {e}"); + std::process::exit(2); + } + return true; + } false } diff --git a/crates/hipfire-quantize/src/pipeline_escha.rs b/crates/hipfire-quantize/src/pipeline_escha.rs new file mode 100644 index 0000000000..7bd1ccd36a --- /dev/null +++ b/crates/hipfire-quantize/src/pipeline_escha.rs @@ -0,0 +1,1113 @@ +//! Converter for EschaLabs Escha-W2 checkpoints (`quant_method` = `escha` / +//! `eschamoe`) into `.hfq`. +//! +//! Code streams are copied byte-for-byte; `memcmp` on the round-trip is a +//! post-condition. See docs/plans/escha-w2-port-design.md. + +use hipfire_quantize::float16::f16_to_f32; +use crate::hfq::QuantType; +use std::path::Path; + +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +pub(crate) enum Leaf { + Code, + Rin, + Rout, + SIn, + SOut, + Config, + Bias, + Int8, + Int8Scale, + Passthrough, + UnknownEscha, +} + +/// The complete escha leaf namespace. Anything else beginning `escha_` is a +/// format mismatch from a newer exporter and must stop conversion. +/// +/// `ignore` in `quantization_config` means "not escha-coded", NOT "not +/// quantized" — both Escha releases list `embed_tokens` and `lm_head` there +/// and still ship them as `weight_int8`. Classification therefore keys off +/// the tensor suffix actually present, never off the config's `ignore` list. +pub(crate) fn classify_leaf(name: &str) -> Leaf { + let suffix = name.rsplit('.').next().unwrap_or(""); + match suffix { + "escha_code" => Leaf::Code, + "escha_rin" => Leaf::Rin, + "escha_rout" => Leaf::Rout, + "escha_s_in" => Leaf::SIn, + "escha_s_out" => Leaf::SOut, + "escha_config" => Leaf::Config, + "bias" => Leaf::Bias, + "weight_int8" => Leaf::Int8, + "weight_scale" => Leaf::Int8Scale, + s if s.starts_with("escha_") => Leaf::UnknownEscha, + _ => Leaf::Passthrough, + } +} + +/// `K` from the code tensor's own shape: the last dimension is `16 * K`. +/// +/// This is the ONLY source of truth. `escha_config` is optional — an export +/// made without the end-to-end fine-tune ships none — and `layer_meta.bits` +/// is self-inconsistent across releases (the same projection is recorded as +/// bits 3.0 in one release and bits 2.0 in the other, both with K=3). +pub(crate) fn k_from_code_shape(shape: &[usize]) -> Result { + let last = *shape.last().ok_or("escha_code has no dimensions")?; + if last % 16 != 0 { + return Err(format!( + "escha_code last dim {last} is not a multiple of 16" + )); + } + let k = last / 16; + if k != 2 && k != 3 { + return Err(format!( + "unsupported escha code rate K={k} (expected 2 or 3)" + )); + } + Ok(k) +} + +pub(crate) fn quant_type_for_k(k: usize) -> Result { + match k { + 2 => Ok(QuantType::ESCHA2T16), + 3 => Ok(QuantType::ESCHA3T16), + _ => Err(format!("unsupported escha code rate K={k}")), + } +} + +/// Required: code, rin, rout. Optional: s_in, s_out, config, bias. Missing +/// any of the required three is a hard error — Escha's own test for this is +/// `rejects_incomplete_linear`, whose docstring requires failing loudly at +/// load rather than decoding into noise. +pub(crate) fn check_linear_complete(proj: &str, present: &[Leaf]) -> Result<(), String> { + for req in [Leaf::Code, Leaf::Rin, Leaf::Rout] { + if !present.contains(&req) { + return Err(format!( + "incomplete escha linear '{proj}': missing {req:?}; \ + refusing to decode into noise" + )); + } + } + Ok(()) +} + +use crate::hfq::{write_hfq, HfqTensor}; +use hipfire_quantize::escha_ref::fold_scales; +use hipfire_quantize::safetensors_file::SafetensorsFile; +use std::collections::BTreeMap; + +/// Convert an Escha-W2 checkpoint directory into a single `.hfq`. +/// +/// `arch` is 6 for `eschamoe` (MoE) and 5 for `escha` (dense). +pub(crate) fn convert_escha(src_dir: &Path, out: &Path) -> Result<(), String> { + let cfg: serde_json::Value = serde_json::from_slice( + &std::fs::read(src_dir.join("config.json")).map_err(|e| e.to_string())?, + ) + .map_err(|e| e.to_string())?; + let qc = &cfg["quantization_config"]; + let method = qc["quant_method"].as_str().unwrap_or_default(); + let version = qc["format_version"].as_str().unwrap_or_default(); + if version != "2.0" { + return Err(format!( + "unsupported escha format_version {version:?}; expected \"2.0\"" + )); + } + let arch: u32 = match method { + "eschamoe" => 6, + "escha" => 5, + other => return Err(format!("not an escha checkpoint: quant_method {other:?}")), + }; + + // Tensors can straddle shards (the 27B's mlp.up_proj has its escha_code in + // shard 2 while its metadata sits in shard 1), so resolve through every + // shard rather than per-file. + let mut shards = Vec::new(); + let mut paths: Vec<_> = std::fs::read_dir(src_dir) + .map_err(|e| e.to_string())? + .filter_map(|e| e.ok()) + .map(|e| e.path()) + .filter(|p| p.extension().is_some_and(|x| x == "safetensors")) + .collect(); + + // MTP lives in a SUBDIRECTORY on the dense export, inline in the main + // shards on the MoE one. `read_dir` is not recursive, so the 35B's + // `mtp.*` tensors rode along as passthrough while the 27B's `mtp/` was + // invisible — a converted 27B silently came out with no MTP head at all + // (verified: 0 mtp tensors in the output). Silent is the wrong failure: + // nobody notices a missing speculative-decode head until they wonder why + // it is slow. + // + // Names inside the subdirectory are already `mtp.`-prefixed in the + // checkpoint, so no rewriting is needed — they land in the same flat + // namespace the 35B produces. + let mtp_dir = src_dir.join("mtp"); + if mtp_dir.is_dir() { + let mut mtp_paths: Vec<_> = std::fs::read_dir(&mtp_dir) + .map_err(|e| e.to_string())? + .filter_map(|e| e.ok()) + .map(|e| e.path()) + .filter(|p| p.extension().is_some_and(|x| x == "safetensors")) + .collect(); + mtp_paths.sort(); + if mtp_paths.is_empty() { + return Err(format!( + "{}: mtp/ exists but holds no .safetensors — refusing to convert a \ + model whose MTP head would be silently dropped", + mtp_dir.display() + )); + } + paths.extend(mtp_paths); + } + paths.sort(); + for p in &paths { + shards.push(SafetensorsFile::open(p).map_err(|e| e.to_string())?); + } + let find = |name: &str| shards.iter().find_map(|s| s.tensor_data(name)); + + // Group leaves by projection prefix so completeness can be checked. + let mut by_proj: BTreeMap> = BTreeMap::new(); + let mut passthrough: Vec = Vec::new(); + for s in &shards { + for name in s.tensor_names() { + let leaf = classify_leaf(name); + match leaf { + Leaf::UnknownEscha => { + return Err(format!( + "unknown escha tensor '{name}': this build implements \ + escha_code/rin/rout/s_in/s_out/config only. A newer \ + exporter shipped a leaf we do not decode; refusing." + )) + } + Leaf::Passthrough | Leaf::Int8 | Leaf::Int8Scale => { + passthrough.push(name.to_string()) + } + _ => { + let prefix = name + .rsplit_once('.') + .ok_or_else(|| format!("{name}: escha leaf name has no '.' separator"))? + .0 + .to_string(); + by_proj + .entry(prefix) + .or_default() + .push((name.to_string(), leaf)); + } + } + } + } + + let mut tensors: Vec = Vec::new(); + for (proj, leaves) in &by_proj { + let kinds: Vec = leaves.iter().map(|(_, l)| *l).collect(); + check_linear_complete(proj, &kinds)?; + + let (meta, data) = find(&format!("{proj}.escha_code")) + .ok_or_else(|| format!("{proj}: escha_code vanished between passes"))?; + let k = k_from_code_shape(&meta.shape)?; + let qt = quant_type_for_k(k)?; + + // Fold the optional end-to-end scales into rin/rout — one f32 pair per + // projection, per row when the tensor is E-stacked. + let (rin_m, rin_d) = find(&format!("{proj}.escha_rin")) + .ok_or_else(|| format!("{proj}: escha_rin vanished between passes"))?; + let (rout_m, rout_d) = find(&format!("{proj}.escha_rout")) + .ok_or_else(|| format!("{proj}: escha_rout vanished between passes"))?; + let s_in = find(&format!("{proj}.escha_s_in")).map(|(_, d)| as_f32(d)); + let s_out = find(&format!("{proj}.escha_s_out")).map(|(_, d)| as_f32(d)); + let (ri, ro) = fold_scales( + &as_u16(rin_d), + &as_u16(rout_d), + s_in.as_deref(), + s_out.as_deref(), + ); + + // ── FOLD MODE ──────────────────────────────────────────────────── + // `HIPFIRE_ESCHA_FOLD=mq6|mq4v2` bakes the H128 rotations and both + // diagonals into the weight and emits an ORDINARY `{proj}.weight`, + // so the runtime needs no escha awareness in the forward pass at all + // (see `escha_fold`). The trellis code is NOT emitted in this mode — + // shipping both would double the file for no reader. + // + // Trade, stated plainly: 2-bit residency is exchanged for escha's + // quantisation QUALITY in a container hipfire already runs at full + // speed. Default is off; unfolded output is unchanged byte-for-byte. + if let Some(fold_fmt) = fold_format() { + let ri_f: Vec = ri.clone(); + let ro_f: Vec = ro.clone(); + let (ic_f, oc_f) = (ri_f.len(), ro_f.len()); + let code_i16: Vec = data + .chunks_exact(2) + .map(|c| i16::from_le_bytes([c[0], c[1]])) + .collect(); + let w_bits = hipfire_quantize::escha_ref::reconstruct(&code_i16, ic_f, oc_f, k as usize); + let folded = hipfire_quantize::escha_fold::fold_escha_linear(&w_bits, ic_f, oc_f, &ri_f, &ro_f)?; + let s1 = crate::quant_fwht::gen_fwht_signs(42, 256); + let s2 = crate::quant_fwht::gen_fwht_signs(1042, 256); + let (bytes, qt_f) = match fold_fmt { + DenseFormat::F16Fold => { + let mut b = Vec::with_capacity(folded.len() * 2); + for v in &folded { + b.extend_from_slice(&hipfire_quantize::float16::f32_to_f16(*v).to_le_bytes()); + } + (b, QuantType::F16) + } + DenseFormat::Mq4V2 => ( + crate::quant_fwht::quantize_mq4g256v2(&folded, oc_f, ic_f, &s1, &s2), + QuantType::MQ4G256V2, + ), + _ => ( + crate::quant_fwht::quantize_mq6g256v2(&folded, oc_f, ic_f, &s1, &s2), + QuantType::MQ6G256V2, + ), + }; + tensors.push(HfqTensor { + name: format!("{proj}.weight"), + quant_type: qt_f, + shape: vec![oc_f as u32, ic_f as u32], + group_size: 256, + data: bytes, + spilled_len: 0, + }); + // A bias CANNOT be folded — it is additive, and no weight matrix + // absorbs it. `WeightTensor` has no bias slot today, so a folded + // model with biases in the file has them silently IGNORED at load. + // That is the same failure class as the dropped MTP head: it does + // not crash, it quietly degrades. Measured on the 27B, the bias is + // ~1.3% of a projection's output magnitude — small per layer, + // compounding over 64. + // + // So refuse, unless the caller says explicitly that they want an + // artifact with the biases dropped (which is legitimate for + // measuring the fold in isolation, and is how the first folded 27B + // was benchmarked at PPL 13.91). + if let Some((bm, bd)) = find(&format!("{proj}.bias")) { + if std::env::var("HIPFIRE_ESCHA_FOLD_DROP_BIAS").as_deref() != Ok("1") { + return Err(format!( + "{proj}: fold mode cannot represent the additive bias, and the runtime \ + has no bias slot to apply it — the folded model would silently ignore \ + it. Add bias support, or set HIPFIRE_ESCHA_FOLD_DROP_BIAS=1 to \ + acknowledge shipping an artifact without it." + )); + } + tensors.push(HfqTensor { + name: format!("{proj}.bias"), + quant_type: QuantType::F16, + shape: bm.shape.iter().map(|&d| d as u32).collect(), + group_size: 0, + data: bd.to_vec(), + spilled_len: 0, + }); + } + continue; + } + + // Verbatim: the code stream is copied byte-for-byte. memcmp on the + // round-trip is the post-condition (G1). + tensors.push(HfqTensor { + name: format!("{proj}.escha_code"), + quant_type: qt, + shape: meta.shape.iter().map(|&d| d as u32).collect(), + group_size: 16, + data: data.to_vec(), + spilled_len: 0, + }); + + tensors.push(f32_tensor( + &format!("{proj}.escha_rin_eff"), + &rin_m.shape, + ri, + )); + tensors.push(f32_tensor( + &format!("{proj}.escha_rout_eff"), + &rout_m.shape, + ro, + )); + + + if let Some((bm, bd)) = find(&format!("{proj}.bias")) { + tensors.push(HfqTensor { + name: format!("{proj}.bias"), + quant_type: QuantType::F16, + shape: bm.shape.iter().map(|&d| d as u32).collect(), + group_size: 0, + data: bd.to_vec(), + spilled_len: 0, + }); + } + } + + for name in &passthrough { + match classify_leaf(name) { + // Consumed alongside its weight_int8 sibling. + Leaf::Int8Scale => continue, + Leaf::Int8 => { + let prefix = name + .rsplit_once('.') + .ok_or_else(|| format!("{name}: escha leaf name has no '.' separator"))? + .0; + let (m, d) = find(name) + .ok_or_else(|| format!("{name}: weight_int8 vanished between passes"))?; + let (_, sd) = find(&format!("{prefix}.weight_scale")).ok_or_else(|| { + format!("{name}: weight_int8 without a matching weight_scale") + })?; + let oc = m.shape[0]; + let ic = m.shape[1]; + let w8: Vec = d.iter().map(|&b| b as i8).collect(); + // The embedding TABLE stays Q8_0 regardless: only one row is + // read per token, so down-quantising it saves nothing per + // token, and the embedding loader supports a fixed format set + // (qt 15 is not in it — it errors with "unsupported embedding + // quant_type 15"). lm_head is a different matter: it IS read + // in full every token (0.54 GB, 22% of the budget). + let is_embed = prefix.contains("embed_tokens") || prefix.contains("token_embd"); + let fmt = if is_embed { DenseFormat::Q8_0 } else { dense_format_for(prefix) }; + let (data, qt, gs) = match fmt { + // F16Fold only ever reaches the FOLD branch above; the + // non-escha dense tensors it never touches stay Q8_0. + DenseFormat::F16Fold | DenseFormat::Q8_0 => ( + int8_rows_to_q8_0(&w8, &as_u16(sd), oc, ic)?, + QuantType::Q8F16, + 32u32, + ), + // Down-quantise the tensors Escha did NOT trellis-code. + // They are int8 in the checkpoint and are 88% of the bytes + // touched per decode token (the 2-bit experts are only + // 12%), so this is the only large lever left on decode — + // at the cost of a SECOND quantisation on top of escha's + // own. Measured relative RMS weight error on + // layers.0.linear_attn.in_proj_qkv: MQ6 2.8%, MQ5 5.9%, + // MQ4V2 9.9%. + DenseFormat::Mq4V2 | DenseFormat::Mq6 => { + let scales = as_u16(sd); + let mut f32_data = Vec::with_capacity(oc * ic); + for o in 0..oc { + let s = f16_to_f32(scales[o]); + for i in 0..ic { + f32_data.push(w8[o * ic + i] as f32 * s); + } + } + let s1 = crate::quant_fwht::gen_fwht_signs(42, 256); + let s2 = crate::quant_fwht::gen_fwht_signs(1042, 256); + // MQ6G256**V2**, not MQ6G256. `fused_qkvza_key_for` + // (forward_slots.rs) matches only the V2 MQ variants + // and has a `_ => FusedQkvzaHfq4G256` catch-all, so a + // plain MQ6G256 QKVZA weight dispatches a 4-bit HFQ4 + // kernel over 6-bit MQ6 bytes: finite output, garbage + // values. Measured with MQ6G256 on the GDN input + // projections: KLD 12.6, PPL 2,375,141 — while + // out_proj, which is not on the fused path, was fine + // at KLD 0.0076. + if fmt == DenseFormat::Mq4V2 { + ( + crate::quant_fwht::quantize_mq4g256v2( + &f32_data, oc, ic, &s1, &s2, + ), + QuantType::MQ4G256V2, + 256u32, + ) + } else { + ( + crate::quant_fwht::quantize_mq6g256v2( + &f32_data, oc, ic, &s1, &s2, + ), + QuantType::MQ6G256V2, + 256u32, + ) + } + } + }; + tensors.push(HfqTensor { + name: format!("{prefix}.weight"), + quant_type: qt, + shape: vec![oc as u32, ic as u32], + group_size: gs, + data, + spilled_len: 0, + }); + } + _ => { + let (m, d) = find(name) + .ok_or_else(|| format!("{name}: passthrough tensor vanished between passes"))?; + + // The GDN input projections all consume the SAME normed x. + // An MQ container needs that activation FWHT-rotated; F16 and + // Q8_0 need it un-rotated. So quantising only some of them + // leaves the rest reading a rotated input they were never + // quantised against — which is what made a single MQ + // `in_proj_qkv` score PPL 2.4M while `out_proj` (which + // consumes the GDN OUTPUT, not the shared input) was fine. + // + // in_proj_a / in_proj_b are F16 in the checkpoint, so they + // must be brought along or the layer stays mixed. + let is_gdn_in = name.contains("in_proj_a") || name.contains("in_proj_b"); + if is_gdn_in + && m.dtype == "F16" + && dense_format_for(name) != DenseFormat::Q8_0 + && m.shape.len() == 2 + && m.shape[1] % 256 == 0 + { + let (oc, ic) = (m.shape[0], m.shape[1]); + let f32_data: Vec = d + .chunks_exact(2) + .map(|c| f16_to_f32(u16::from_le_bytes([c[0], c[1]]))) + .collect(); + let s1 = crate::quant_fwht::gen_fwht_signs(42, 256); + let s2 = crate::quant_fwht::gen_fwht_signs(1042, 256); + let (data, qt) = if dense_format_for(name) == DenseFormat::Mq4V2 { + ( + crate::quant_fwht::quantize_mq4g256v2(&f32_data, oc, ic, &s1, &s2), + QuantType::MQ4G256V2, + ) + } else { + ( + crate::quant_fwht::quantize_mq6g256v2(&f32_data, oc, ic, &s1, &s2), + QuantType::MQ6G256V2, + ) + }; + tensors.push(HfqTensor { + name: name.clone(), + quant_type: qt, + shape: vec![oc as u32, ic as u32], + group_size: 256, + data, + spilled_len: 0, + }); + continue; + } + + tensors.push(HfqTensor { + name: name.clone(), + quant_type: match m.dtype.as_str() { + "F16" => QuantType::F16, + "F32" => QuantType::F32, + "BF16" => QuantType::BF16, + other => return Err(format!("{name}: unhandled dtype {other}")), + }, + shape: m.shape.iter().map(|&d| d as u32).collect(), + group_size: 0, + data: d.to_vec(), + spilled_len: 0, + }); + } + } + } + + let metadata = build_metadata(src_dir, &cfg, version, method)?; + write_hfq(out, arch, &metadata, &tensors, None).map_err(|e| e.to_string()) +} + +/// Build the HFQ `metadata_json` envelope. +/// +/// The envelope shape is NOT free-form. It mirrors `pipeline.rs`'s envelope +/// key-for-key and only ADDS the `escha` provenance key: +/// +/// * `config` — the parsed config.json verbatim. `config_from_metadata_json` +/// (hipfire-arch-qwen35) requires it to reconstruct the arch config at load +/// time, and it self-detects the nested `text_config`/`vision_config` these +/// VL-shaped checkpoints carry — do not flatten or pre-process it here. +/// * `tokenizer` — tokenizer.json verbatim as a STRING, not a nested object. +/// That is what `Tokenizer::from_hfq_metadata` expects; anything else and it +/// returns `MetadataMissing { field: "tokenizer | gguf_meta" }`. vocab.json +/// and merges.txt are NOT carried separately — tokenizer.json already holds +/// the BPE vocab and merge table, and no reader looks for the sidecars. +/// * `tokenizer_config` — carries `chat_template`, the ONLY key +/// `HfqFile::chat_template()` reads and hence the only source +/// `resolve_chat_template` has for arch 5/6. +/// * `generation_config` — authoritative bos/eos ids. Escha's is an ARRAY eos +/// `[248046, 248044]`, so `from_hfq_metadata` keeps its heuristic eos; the +/// bos scalar 248044 still overrides. +fn build_metadata( + src_dir: &Path, + cfg: &serde_json::Value, + version: &str, + method: &str, +) -> Result { + let read_json = |name: &str| -> Option { + std::fs::read_to_string(src_dir.join(name)) + .ok() + .and_then(|s| serde_json::from_str(&s).ok()) + }; + + let tokenizer_str = std::fs::read_to_string(src_dir.join("tokenizer.json")).ok(); + if tokenizer_str.is_none() { + return Err(format!( + "escha: no tokenizer.json in {} — the .hfq would convert cleanly and \ + then be unservable (Tokenizer::from_hfq_metadata would fail). \ + Refusing to write a model that cannot be driven.", + src_dir.display() + )); + } + + // Some checkpoints ship the Jinja template in a `chat_template.jinja` + // sidecar rather than inside tokenizer_config.json. Fold it in only when + // tokenizer_config lacks one — an existing template wins, same rule as + // `pipeline.rs`. + let tokenizer_config = { + let mut tc = read_json("tokenizer_config.json"); + let jinja_path = src_dir.join("chat_template.jinja"); + if jinja_path.exists() { + let has_template = tc + .as_ref() + .and_then(|v| v.get("chat_template")) + .map(|v| !v.is_null()) + .unwrap_or(false); + if !has_template { + if let Ok(jinja) = std::fs::read_to_string(&jinja_path) { + let n = jinja.len(); + let obj = tc.get_or_insert_with(|| serde_json::json!({})); + if let Some(map) = obj.as_object_mut() { + map.insert( + "chat_template".to_string(), + serde_json::Value::String(jinja), + ); + eprintln!( + " embedded chat_template.jinja into tokenizer_config ({n} bytes)" + ); + } + } + } + } + tc + }; + if tokenizer_config + .as_ref() + .and_then(|v| v.get("chat_template")) + .and_then(|v| if v.is_null() { None } else { Some(v) }) + .is_none() + { + eprintln!( + "escha: warning: no chat_template in tokenizer_config.json and no \ + chat_template.jinja sidecar — the daemon will fall back to a \ + hand-rolled frame for this instruct model" + ); + } + + let metadata = serde_json::json!({ + "config": cfg, + "tokenizer": tokenizer_str.as_deref().unwrap_or("{}"), + "tokenizer_config": tokenizer_config, + "generation_config": read_json("generation_config.json"), + "escha": { "format_version": version, "quant_method": method }, + }); + serde_json::to_string(&metadata).map_err(|e| format!("serialize metadata: {e}")) +} + +/// Escha's int8 is per-output-ROW; hipfire's `Q8_0` is per-32-element block +/// (34 bytes: f16 scale then 32 int8). Replicating the row scale into every +/// block of that row passes the int8 bytes through unchanged, so the +/// reconstruction is bit-identical to Escha's `w8a16`. Cost is 2 bytes per 32 +/// elements — 6.25% — for scales that are all equal within a row. +/// +/// Do NOT recompute per-block scales from the dequantised values. That is a +/// second quantisation and adds avoidable error. + +/// Container for the tensors Escha did NOT trellis-code (`weight_int8` in the +/// checkpoint: linear-attention projections, attention q/k/v/o, shared expert, +/// router, embeddings, lm_head). +/// +/// These are 88% of the bytes touched per decode token — the 2-bit routed +/// experts are only 12% — so they, not the codec, set decode speed. Measured +/// per-token totals on the shipped 35B: 2.448 GB at Q8_0 against a comparable +/// MQ4 SKU's 1.886 GB, i.e. escha reads 1.30x more DESPITE being 12.33 GB on +/// disk against 19.00 GB. +/// +/// Default is `Q8_0`, which preserves escha's shipped int8 exactly (the repack +/// is bit-exact — see `int8_rows_to_q8_0`). Anything else is a SECOND +/// quantisation and changes the model; it must be justified by a KLD number, +/// not by the byte saving alone. +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +pub(crate) enum DenseFormat { + Q8_0, + Mq6, + /// MQ4G256V2 — the container the comparable mq4 SKU uses for exactly + /// these tensors, and the one with the widest kernel coverage. + Mq4V2, + F16Fold, +} + +/// `HIPFIRE_ESCHA_DENSE=mq6` down-quantises every non-embedding int8 tensor. +/// `HIPFIRE_ESCHA_DENSE=mq6:` limits it to tensors whose name contains +/// `` — used to isolate WHICH path a quality regression comes from, +/// since a whole-model number cannot distinguish "this format costs accuracy" +/// from "this one path is not applying the activation rotation MQ requires". +/// Pick the container for one dense (non-escha-coded) tensor. +/// +/// `HIPFIRE_ESCHA_DENSE` selects it: +/// unset / `q8` every dense tensor stays `Q8_0` — the default +/// `mq6` / `mq4v2` that format everywhere EXCEPT the router and the +/// shared expert (see `DENSE_DEFAULT_KEEP_Q8`) +/// `:a,b,c` ONLY tensors matching a listed substring +/// `!a,b,c` every tensor EXCEPT those matching one +/// +/// `Q8_0` is the default because it is the most faithful of the three shipped +/// recipes and the one that reproduces the upstream checkpoint's precision: +/// Escha's own int8 is per-ROW, and this repacks it into hipfire's per-32-block +/// `Q8_0` by replicating the row scale into every block, which is bit-exact and +/// costs 6.25% on those tensors only (12.34 GB against upstream's 12.30 GB). +/// +/// The two down-quants are shipped alternatives, measured on the same corpus: +/// +/// dense size PPL decode note +/// q8 12.34 GB 7.6864 ~46 tok/s default, most faithful +/// mq6 11.84 GB 7.6940 ~56 tok/s +21% decode, +0.10% PPL +/// mq4v2 11.39 GB 8.0643 ~62 tok/s +35% decode, +4.9% PPL +/// +/// `mq6` and `mq4v2` hold the router and shared expert at `Q8_0` regardless, +/// because moving either costs batched prefill outright — see +/// `DENSE_DEFAULT_KEEP_Q8`. Anything other than these three spellings is a +/// research knob, not a shipped recipe. +/// +/// Down-quanting dense is nearly free in quality and worth ~17% of decode: +/// the dense weights are the half of escha's per-token byte traffic that its +/// 2-bit experts already win handily. KLD against the int8-dense build on the +/// same corpus (self-control 0.000000, baseline PPL 7.6769): +/// +/// lm_head 0.000838 PPL 7.6984 +/// self_attn 0.005879 PPL 7.6782 +/// linear_attn (GDN, all 4) 0.004637 PPL 7.6633 +/// everything 0.009785 PPL 7.6965 +/// +/// An earlier reading of this table had GDN at 0.448186 (PPL 13.41) and a +/// lone MQ6 `in_proj_qkv` at 12.639980 (PPL 2.4M), and concluded GDN's +/// recurrent state compounds quantisation error along the sequence where +/// attention's does not. That conclusion was WRONG. Those numbers were a +/// mixed-rotation bug in this converter: `in_proj_a` / `in_proj_b` are F16 in +/// the checkpoint and fell through to the passthrough branch, so quantising +/// their siblings FWHT-rotated the shared normed x out from under them. Once +/// all four projections share a container, GDN quantises better than baseline. +/// The tell was that `out_proj` was unharmed (0.0076) — it consumes the GDN +/// output, not the shared input. +/// +/// The `!` form exists because `moe_ffn_batched_admissible_for_dtypes` pins +/// the router (`mlp.gate`) and the shared expert to `Q8_0` on the escha arm. +/// Moving either off it costs batched prefill — measured 160 H128 launches +/// per token and 109.5 tok/s, against 0.3 and 708.0 with both excluded — for +/// 0.0002 KLD, since neither is where down-quanting pays. +/// Tensors the default recipe deliberately leaves at `Q8_0`. +/// +/// `moe_ffn_batched_admissible_for_dtypes` pins the router (`mlp.gate`) and +/// the shared expert to `Q8_0` on the escha arm. Moving either off it costs +/// batched prefill — measured 160 H128 launches per token and 109.5 tok/s, +/// against 0.3 and 708.0 with both excluded — and buys 0.0002 KLD, because +/// neither is where down-quanting pays. +const DENSE_DEFAULT_KEEP_Q8: [&str; 2] = ["mlp.gate", "shared_expert"]; + +/// `HIPFIRE_ESCHA_FOLD` — when set, escha-coded linears are FOLDED into +/// ordinary `{proj}.weight` tensors instead of shipped as trellis code. +/// +/// `mq6` (recommended) or `mq4v2`. Unset means the normal escha output, +/// byte-for-byte unchanged. See `escha_fold` for the algebra and the trade. +fn fold_format() -> Option { + match std::env::var("HIPFIRE_ESCHA_FOLD").ok().as_deref() { + Some("mq6") => Some(DenseFormat::Mq6), + Some("mq4v2") => Some(DenseFormat::Mq4V2), + // `f16` folds WITHOUT re-quantising. It is an ATTRIBUTION tool, not a + // shipping format (~50 GB): it isolates escha's own 2-bit quality by + // removing the fold's second quantisation step, so a gap against a + // native MQ6 baseline can be split between "escha's codec" and "my + // re-quantisation". + Some("f16") => Some(DenseFormat::F16Fold), + _ => None, + } +} + +fn dense_format_for(name: &str) -> DenseFormat { + let Ok(spec) = std::env::var("HIPFIRE_ESCHA_DENSE") else { + return DenseFormat::Q8_0; + }; + let (fmt, rule) = match spec.find(['!', ':']) { + Some(i) => (&spec[..i], Some((spec.as_bytes()[i], &spec[i + 1..]))), + None => (spec.as_str(), None), + }; + let fmt = match fmt { + "mq6" => DenseFormat::Mq6, + "mq4v2" => DenseFormat::Mq4V2, + // Everything else, `q8` included, is the faithful default. Spelling + // `q8` explicitly is supported so a build script can state the recipe + // rather than rely on the variable being absent. + _ => return DenseFormat::Q8_0, + }; + match rule { + // A bare `mq6` / `mq4v2` still holds the router and shared expert at + // Q8_0. Losing batched prefill is never what the caller meant, and the + // exclusion costs 0.0002 KLD; `mq6!` with an empty list opts out. + None => { + if DENSE_DEFAULT_KEEP_Q8.iter().any(|sub| name.contains(sub)) { + return DenseFormat::Q8_0; + } + fmt + } + Some((sep, list)) => { + let hit = list + .split(',') + .any(|sub| !sub.is_empty() && name.contains(sub)); + // `:` = allow-list (hit takes fmt), `!` = deny-list (hit stays Q8). + if hit == (sep == b':') { + fmt + } else { + DenseFormat::Q8_0 + } + } + } +} + +pub(crate) fn int8_rows_to_q8_0( + w8: &[i8], + scale_f16: &[u16], + oc: usize, + ic: usize, +) -> Result, String> { + if w8.len() != oc * ic { + return Err(format!( + "int8 tensor is {} bytes, expected {oc}x{ic}", + w8.len() + )); + } + if scale_f16.len() != oc { + return Err(format!("expected {oc} row scales, got {}", scale_f16.len())); + } + if ic % 32 != 0 { + return Err(format!("Q8_0 needs a multiple of 32 per row, got ic={ic}")); + } + let mut out = Vec::with_capacity(oc * (ic / 32) * 34); + for o in 0..oc { + let s = scale_f16[o].to_le_bytes(); + for blk in 0..ic / 32 { + out.extend_from_slice(&s); + let base = o * ic + blk * 32; + out.extend(w8[base..base + 32].iter().map(|&v| v as u8)); + } + } + Ok(out) +} + +fn as_u16(d: &[u8]) -> Vec { + d.chunks_exact(2) + .map(|c| u16::from_le_bytes([c[0], c[1]])) + .collect() +} + +fn as_f32(d: &[u8]) -> Vec { + d.chunks_exact(4) + .map(|c| f32::from_le_bytes([c[0], c[1], c[2], c[3]])) + .collect() +} + +fn f32_tensor(name: &str, shape: &[usize], v: Vec) -> HfqTensor { + HfqTensor { + name: name.to_string(), + quant_type: QuantType::F32, + shape: shape.iter().map(|&d| d as u32).collect(), + group_size: 0, + data: v.iter().flat_map(|x| x.to_le_bytes()).collect(), + spilled_len: 0, + } +} + +#[cfg(test)] +mod tests { + use super::*; + use hipfire_quantize::escha_ref::f16_rne; + + #[test] + fn k_comes_from_the_code_shape_not_metadata() { + // gate_up: [E, in/16, out/16, 16K] with K=2 -> last dim 32 + assert_eq!(k_from_code_shape(&[256, 128, 64, 32]), Ok(2)); + // down: K=3 -> last dim 48 + assert_eq!(k_from_code_shape(&[256, 32, 128, 48]), Ok(3)); + // dense exports have no E axis + assert_eq!(k_from_code_shape(&[320, 1088, 32]), Ok(2)); + assert!(k_from_code_shape(&[256, 128, 64, 33]).is_err()); + } + + #[test] + fn quant_type_follows_k() { + assert_eq!(quant_type_for_k(2), Ok(QuantType::ESCHA2T16)); + assert_eq!(quant_type_for_k(3), Ok(QuantType::ESCHA3T16)); + assert!(quant_type_for_k(4).is_err()); + } + + /// `ignore` means "not escha-coded", NOT "not quantized" — both models + /// list embed_tokens and lm_head there and ship them as weight_int8. + /// Classification must key off the tensor suffix actually present. + #[test] + fn classify_leaf_keys_off_the_suffix() { + assert_eq!( + classify_leaf("l.0.mlp.experts.gate_up_proj.escha_code"), + Leaf::Code + ); + assert_eq!( + classify_leaf("l.0.mlp.experts.gate_up_proj.escha_rin"), + Leaf::Rin + ); + assert_eq!( + classify_leaf("l.0.mlp.experts.gate_up_proj.escha_s_out"), + Leaf::SOut + ); + assert_eq!(classify_leaf("lm_head.weight_int8"), Leaf::Int8); + assert_eq!(classify_leaf("lm_head.weight_scale"), Leaf::Int8Scale); + assert_eq!( + classify_leaf("l.0.input_layernorm.weight"), + Leaf::Passthrough + ); + } + + /// A future export carrying a rotation variant this version does not + /// implement must stop conversion, not decode under the wrong rotation. + #[test] + fn unknown_escha_leaf_is_rejected() { + assert_eq!( + classify_leaf("l.0.mlp.gate_proj.escha_rotation_theta"), + Leaf::UnknownEscha + ); + } + + /// Required: code, rin, rout. Missing any is "incomplete escha linear" — + /// fail loudly at load, never a partial decode. + #[test] + fn incomplete_linear_is_rejected() { + let mut present = vec![Leaf::Code, Leaf::Rin, Leaf::Rout]; + assert!(check_linear_complete("proj", &present).is_ok()); + present.pop(); + let err = check_linear_complete("proj", &present).unwrap_err(); + assert!(err.contains("incomplete escha linear"), "{err}"); + } + + /// Optional: s_in, s_out, config, bias. An export without the end-to-end + /// stage ships none of them and must still convert. + #[test] + fn optional_leaves_may_all_be_absent() { + assert!(check_linear_complete("proj", &[Leaf::Code, Leaf::Rin, Leaf::Rout]).is_ok()); + } + + /// The row scale must be replicated into every block with the int8 bytes + /// untouched — that is what makes the repack bit-exact. Recomputing block + /// scales would be a second quantisation. + #[test] + fn int8_repack_replicates_the_row_scale() { + let oc = 2; + let ic = 64; // two Q8_0 blocks per row + let w8: Vec = (0..(oc * ic)).map(|i| (i % 127) as i8).collect(); + let scale = vec![f16_rne(0.5), f16_rne(2.0)]; + let q8 = int8_rows_to_q8_0(&w8, &scale, oc, ic).unwrap(); + assert_eq!(q8.len(), oc * (ic / 32) * 34); + // Both blocks of row 0 carry row 0's scale, unchanged. + assert_eq!(&q8[0..2], &scale[0].to_le_bytes()); + assert_eq!(&q8[34..36], &scale[0].to_le_bytes()); + // Row 1's blocks carry row 1's scale. + assert_eq!(&q8[68..70], &scale[1].to_le_bytes()); + // Payload bytes are passed through verbatim. + assert_eq!(q8[2] as i8, w8[0]); + assert_eq!(q8[36] as i8, w8[32]); + } + + #[test] + fn int8_repack_rejects_a_ragged_row() { + assert!(int8_rows_to_q8_0(&[0i8; 20], &[0u16], 1, 20).is_err()); + } +} + +/// `convert_escha` integration test: a minimal synthetic checkpoint directory +/// (config.json + one safetensors shard) through the real converter, then a +/// real `HfqFile` read-back. This lives in-module (not under `tests/`) for +/// the same reason as `reap_overlay::integ`: `hipfire-quantize` is a binary +/// crate with no library target, so a `tests/` integration target can't reach +/// `convert_escha`, which is crate-private. `hipfire-runtime` (CPU-only HFQ +/// container reader) is a dev-dependency. +/// +/// This is the test that Finding 1 (missing top-level `config` key) proves +/// would have caught the regression: it asserts the round-tripped metadata +/// carries `config` with the fields written into config.json. +#[cfg(test)] +mod convert_escha_tests { + use super::*; + use hipfire_quantize::escha_ref::f16_rne; + use hipfire_runtime::hfq::HfqFile; + use safetensors::tensor::TensorView; + use safetensors::Dtype; + use std::collections::HashMap; + + /// escha_code payload: deterministic, recognizable bytes so the + /// verbatim-repack assertion can compare against a value computed + /// independently of what `build_fixture` wrote. + fn code_bytes() -> Vec { + (0..(4 * 32)).map(|i| (i % 251) as u8).collect() + } + + /// Write a minimal Escha-W2 checkpoint directory: config.json with a + /// `quantization_config` block plus a couple of recognisable top-level + /// config fields, and one safetensors shard with one complete escha + /// linear (escha_code last dim 16*K=32 => K=2, plus escha_rin/rout) and + /// one int8 passthrough pair (weight_int8 + weight_scale). + fn build_fixture(dir: &Path) { + let cfg = serde_json::json!({ + "hidden_size": 64, + "num_hidden_layers": 2, + "num_attention_heads": 4, + "vocab_size": 100, + "quantization_config": { + "quant_method": "eschamoe", + "format_version": "2.0", + }, + }); + std::fs::write(dir.join("config.json"), cfg.to_string()).unwrap(); + + // A real escha checkpoint always ships these; without them the .hfq + // converts cleanly and is then unservable, which is the regression + // `tokenizer_and_chat_template_land_in_metadata` pins. + std::fs::write( + dir.join("tokenizer.json"), + r#"{"model":{"type":"BPE","vocab":{"a":0,"b":1},"merges":[]},"added_tokens":[]}"#, + ) + .unwrap(); + std::fs::write( + dir.join("tokenizer_config.json"), + r#"{"add_bos_token":false,"chat_template":"ESCHA_TEMPLATE"}"#, + ) + .unwrap(); + std::fs::write( + dir.join("generation_config.json"), + r#"{"bos_token_id":1,"eos_token_id":[1,0]}"#, + ) + .unwrap(); + + let code = code_bytes(); + + let rin: Vec = (0..4).map(|i| f16_rne(1.0 + i as f32 * 0.1)).collect(); + let rout: Vec = (0..4).map(|i| f16_rne(2.0 + i as f32 * 0.1)).collect(); + let rin_bytes: Vec = rin.iter().flat_map(|v| v.to_le_bytes()).collect(); + let rout_bytes: Vec = rout.iter().flat_map(|v| v.to_le_bytes()).collect(); + + // int8 pair: oc=2, ic=32 (one Q8_0 block per row). + let w8: Vec = (0..(2 * 32)).map(|i| (i as i8) as u8).collect(); + let scale: Vec = vec![f16_rne(0.5), f16_rne(1.5)]; + let scale_bytes: Vec = scale.iter().flat_map(|v| v.to_le_bytes()).collect(); + + let mut tensors: HashMap = HashMap::new(); + tensors.insert( + "layers.0.mlp.up_proj.escha_code".to_string(), + TensorView::new(Dtype::U8, vec![4, 32], &code).unwrap(), + ); + tensors.insert( + "layers.0.mlp.up_proj.escha_rin".to_string(), + TensorView::new(Dtype::F16, vec![4], &rin_bytes).unwrap(), + ); + tensors.insert( + "layers.0.mlp.up_proj.escha_rout".to_string(), + TensorView::new(Dtype::F16, vec![4], &rout_bytes).unwrap(), + ); + tensors.insert( + "lm_head.weight_int8".to_string(), + TensorView::new(Dtype::I8, vec![2, 32], &w8).unwrap(), + ); + tensors.insert( + "lm_head.weight_scale".to_string(), + TensorView::new(Dtype::F16, vec![2], &scale_bytes).unwrap(), + ); + + let bytes = safetensors::serialize(&tensors, None).unwrap(); + std::fs::write(dir.join("model.safetensors"), bytes).unwrap(); + } + + /// Unique scratch dir under the system temp root, cleaned up on drop. + struct TempCheckpointDir(std::path::PathBuf); + impl TempCheckpointDir { + fn new(tag: &str) -> Self { + let nanos = std::time::SystemTime::now() + .duration_since(std::time::UNIX_EPOCH) + .unwrap() + .as_nanos(); + let dir = std::env::temp_dir().join(format!( + "hipfire_escha_convert_test_{tag}_{}_{nanos}", + std::process::id() + )); + std::fs::create_dir_all(&dir).unwrap(); + Self(dir) + } + } + impl Drop for TempCheckpointDir { + fn drop(&mut self) { + std::fs::remove_dir_all(&self.0).ok(); + } + } + + #[test] + fn convert_escha_embeds_config_and_repacks_code_verbatim() { + let src = TempCheckpointDir::new("ok"); + build_fixture(&src.0); + let out = src.0.join("out.hfq"); + + convert_escha(&src.0, &out).expect("conversion of a well-formed fixture must succeed"); + + let hf = HfqFile::open(&out).expect("convert_escha's output must parse as a valid .hfq"); + + // The regression Finding 1 caught: metadata must carry a top-level + // `config` key, or every arch loader's `config_from_metadata_json` + // fails before a tensor is read. + let meta: serde_json::Value = serde_json::from_str(&hf.metadata_json) + .expect("metadata_json must itself be valid JSON"); + let config = meta + .get("config") + .expect("metadata must carry a top-level `config` key"); + assert_eq!(config["hidden_size"], 64); + assert_eq!(config["num_hidden_layers"], 2); + assert_eq!(config["num_attention_heads"], 4); + assert_eq!(config["vocab_size"], 100); + + // The verbatim-repack contract (G1): escha_code bytes in the output + // are byte-identical to the input. + let (_, out_code) = hf + .tensor_data("layers.0.mlp.up_proj.escha_code") + .expect("escha_code tensor must survive conversion"); + assert_eq!(out_code, code_bytes().as_slice()); + } + + /// The converter originally emitted only `config` + `escha`, so the .hfq + /// carried no tokenizer, no chat template and no generation_config: it + /// converted cleanly, passed G1, and could not be served at all. Pin the + /// exact keys the readers use — `Tokenizer::from_hfq_metadata` wants + /// `tokenizer` as a STRING, and `HfqFile::chat_template()` reads ONLY + /// `tokenizer_config.chat_template`. + #[test] + fn tokenizer_and_chat_template_land_in_metadata() { + let src = TempCheckpointDir::new("tok"); + build_fixture(&src.0); + let out = src.0.join("out.hfq"); + convert_escha(&src.0, &out).expect("conversion must succeed"); + let hf = HfqFile::open(&out).expect("output must parse"); + + let meta: serde_json::Value = serde_json::from_str(&hf.metadata_json).unwrap(); + let tok = meta["tokenizer"] + .as_str() + .expect("`tokenizer` must be a STRING holding tokenizer.json verbatim"); + assert!(tok.contains("\"vocab\""), "tokenizer.json must be verbatim"); + assert_eq!(meta["tokenizer_config"]["add_bos_token"], false); + assert_eq!(meta["generation_config"]["bos_token_id"], 1); + // The provenance key must survive alongside the new siblings. + assert_eq!(meta["escha"]["quant_method"], "eschamoe"); + + // The reader paths themselves, not just the raw keys. + assert_eq!( + hf.chat_template().as_deref(), + Some("ESCHA_TEMPLATE"), + "resolve_chat_template (arch 5|6) reads this and nothing else" + ); + hipfire_runtime::tokenizer::Tokenizer::from_hfq_metadata(&hf.metadata_json) + .expect("the embedded metadata must build a working Tokenizer"); + } + + /// A checkpoint with no tokenizer.json cannot produce a servable model. + /// Fail at convert time rather than shipping a 12 GB file that only fails + /// when someone tries to run it — same fail-closed rule this converter + /// applies to unknown escha leaves and incomplete linears. + #[test] + fn missing_tokenizer_is_rejected() { + let src = TempCheckpointDir::new("notok"); + build_fixture(&src.0); + std::fs::remove_file(src.0.join("tokenizer.json")).unwrap(); + let err = convert_escha(&src.0, &src.0.join("out.hfq")).unwrap_err(); + assert!(err.contains("no tokenizer.json"), "{err}"); + } +} diff --git a/crates/hipfire-quantize/tests/data/escha/expected_down_e0_k3.f16 b/crates/hipfire-quantize/tests/data/escha/expected_down_e0_k3.f16 new file mode 100644 index 0000000000..242941acdf Binary files /dev/null and b/crates/hipfire-quantize/tests/data/escha/expected_down_e0_k3.f16 differ diff --git a/crates/hipfire-quantize/tests/data/escha/expected_gu_e0_k2.f16 b/crates/hipfire-quantize/tests/data/escha/expected_gu_e0_k2.f16 new file mode 100644 index 0000000000..e084a54894 Binary files /dev/null and b/crates/hipfire-quantize/tests/data/escha/expected_gu_e0_k2.f16 differ diff --git a/crates/hipfire-quantize/tests/data/escha/fetch-goldens.sh b/crates/hipfire-quantize/tests/data/escha/fetch-goldens.sh new file mode 100755 index 0000000000..4440ca2c80 --- /dev/null +++ b/crates/hipfire-quantize/tests/data/escha/fetch-goldens.sh @@ -0,0 +1,40 @@ +#!/usr/bin/env bash +# Fetch the escha-mlx golden vectors that are NOT committed to this repo. +# +# WHAT IS AND IS NOT COMMITTED +# +# Committed (a clean checkout can run G2, G3, G4, G4b with no network): +# packed_gu_e0_k2.i16 packed_down_e0_k3.i16 — codec inputs (G2) +# moeblk_x.f16 moeblk_out.f16 — MoE block in/out (G4) +# moeblk_ids.i64 moeblk_scores.f32 — injected routing (G4/G4b) +# +# NOT committed, fetched by this script (6.0 MB): +# expected_gu_e0_k2.f16 expected_down_e0_k3.f16 +# +# Only that LAST pair is "needed only to regenerate the digests" — the +# `escha_ref` unit tests compare against sha256 digests recorded in +# escha_ref.rs, so the expected tensors themselves are not needed to run +# them. The moeblk_* fixtures are a different case entirely: G4 and G4b +# assert against `moeblk_out.f16` directly and cannot run without it, which +# is why they are now in the repo. (They previously were not, and those two +# gates were unrunnable from a clean checkout while this comment claimed +# every uncommitted file was digest-regeneration only.) +# +# THE REF IS PINNED, deliberately. It used to be `HEAD`, which means the +# goldens this gate scores against could change under the repo without a +# single line of it changing — the classic silently-moving-oracle. This SHA's +# eight files were verified byte-for-byte (sha256) against the copies now in +# the tree. To move to a newer upstream: bump REF, re-run, diff the printed +# digests against the committed files, and say in the commit message what +# changed and why. +set -euo pipefail +cd "$(dirname "$0")" +REF=22f7f4c3bb128cacee0d7ae19af9812b212bdc2c +B=https://raw.githubusercontent.com/EschaLabs/escha-mlx/$REF/tests/data +for f in codec/packed_gu_e0_k2.i16 codec/expected_gu_e0_k2.f16 \ + codec/packed_down_e0_k3.i16 codec/expected_down_e0_k3.f16 \ + qwen3_5_moe/moeblk_x.f16 qwen3_5_moe/moeblk_out.f16 \ + qwen3_5_moe/moeblk_ids.i64 qwen3_5_moe/moeblk_scores.f32; do + curl -sL --fail "$B/$f" -o "$(basename "$f")" +done +sha256sum ./*.f16 ./*.i16 ./*.i64 ./*.f32 diff --git a/crates/hipfire-quantize/tests/data/escha/moeblk_ids.i64 b/crates/hipfire-quantize/tests/data/escha/moeblk_ids.i64 new file mode 100644 index 0000000000..e1aad03b12 Binary files /dev/null and b/crates/hipfire-quantize/tests/data/escha/moeblk_ids.i64 differ diff --git a/crates/hipfire-quantize/tests/data/escha/moeblk_out.f16 b/crates/hipfire-quantize/tests/data/escha/moeblk_out.f16 new file mode 100644 index 0000000000..ab8b7897c9 Binary files /dev/null and b/crates/hipfire-quantize/tests/data/escha/moeblk_out.f16 differ diff --git a/crates/hipfire-quantize/tests/data/escha/moeblk_scores.f32 b/crates/hipfire-quantize/tests/data/escha/moeblk_scores.f32 new file mode 100644 index 0000000000..78baa39d91 --- /dev/null +++ b/crates/hipfire-quantize/tests/data/escha/moeblk_scores.f32 @@ -0,0 +1,3 @@ +JL>~H'>z >===º==li>8>{='= +Y=}=y==$>L>===I===k)>1># >>ʮ=W=9=9=r3>R>} +>6==U=h=ɣ=T>2>N>==+E==,=<>6,>݉=}=5=#=J=S=m>i >SD>h=R== =[= \ No newline at end of file diff --git a/crates/hipfire-quantize/tests/data/escha/moeblk_x.f16 b/crates/hipfire-quantize/tests/data/escha/moeblk_x.f16 new file mode 100644 index 0000000000..6ebb8183d2 Binary files /dev/null and b/crates/hipfire-quantize/tests/data/escha/moeblk_x.f16 differ diff --git a/crates/hipfire-quantize/tests/data/escha/packed_down_e0_k3.i16 b/crates/hipfire-quantize/tests/data/escha/packed_down_e0_k3.i16 new file mode 100644 index 0000000000..ee1f10b14f Binary files /dev/null and b/crates/hipfire-quantize/tests/data/escha/packed_down_e0_k3.i16 differ diff --git a/crates/hipfire-quantize/tests/data/escha/packed_gu_e0_k2.i16 b/crates/hipfire-quantize/tests/data/escha/packed_gu_e0_k2.i16 new file mode 100644 index 0000000000..87e046335a Binary files /dev/null and b/crates/hipfire-quantize/tests/data/escha/packed_gu_e0_k2.i16 differ diff --git a/crates/hipfire-runtime/src/spec.rs b/crates/hipfire-runtime/src/spec.rs index 1499f231a2..6fc2025f23 100644 --- a/crates/hipfire-runtime/src/spec.rs +++ b/crates/hipfire-runtime/src/spec.rs @@ -1116,6 +1116,17 @@ impl MtpSpeculator { } } +/// Whether `HIPFIRE_MTP_ACCEPT_STATS=1` asked for per-window accept logging. +fn mtp_accept_stats_enabled() -> bool { + use std::sync::OnceLock; + static ON: OnceLock = OnceLock::new(); + *ON.get_or_init(|| { + std::env::var("HIPFIRE_MTP_ACCEPT_STATS") + .map(|v| v == "1" || v.eq_ignore_ascii_case("true")) + .unwrap_or(false) + }) +} + /// Lower an [`MtpWindow`] to the generic [`SpecStep`]. `committed` already /// excludes the seed and includes the bonus, so it maps 1:1 to `emit`; the next /// window's seed is the last committed token (the daemon's `position += @@ -1127,6 +1138,28 @@ fn lower_mtp_window(w: MtpWindow) -> Result { .committed .last() .ok_or("MtpSpeculator: drafter committed 0 tokens (would stall the decode loop)")?; + // Every MTP window funnels through here, so this is the one place a + // cumulative accept rate can be observed without threading stats out to + // each caller. Speed alone cannot distinguish "drafts rejected" from + // "drafts accepted but verification is expensive"; HIPFIRE_MTP_ACCEPT_STATS=1 + // prints the rate a speculative-decode change has to move. + if mtp_accept_stats_enabled() { + use std::sync::atomic::{AtomicUsize, Ordering}; + static WINDOWS: AtomicUsize = AtomicUsize::new(0); + static DRAFTED: AtomicUsize = AtomicUsize::new(0); + static ACCEPTED: AtomicUsize = AtomicUsize::new(0); + static EMITTED: AtomicUsize = AtomicUsize::new(0); + let n = WINDOWS.fetch_add(1, Ordering::Relaxed) + 1; + let d = DRAFTED.fetch_add(w.drafts_generated, Ordering::Relaxed) + w.drafts_generated; + let a = ACCEPTED.fetch_add(w.accepted, Ordering::Relaxed) + w.accepted; + let e = EMITTED.fetch_add(w.committed.len(), Ordering::Relaxed) + w.committed.len(); + eprintln!( + "[mtp-accept] windows={n} drafted={d} accepted={a} emitted={e} \ + accept_rate={:.3} tokens_per_window={:.3}", + if d > 0 { a as f64 / d as f64 } else { 0.0 }, + e as f64 / n as f64, + ); + } Ok(SpecStep::new( w.committed.iter().copied(), next_seed, diff --git a/crates/hipfire-runtime/src/weight_backend.rs b/crates/hipfire-runtime/src/weight_backend.rs index 43a7d09dc9..27e6145394 100644 --- a/crates/hipfire-runtime/src/weight_backend.rs +++ b/crates/hipfire-runtime/src/weight_backend.rs @@ -1265,6 +1265,33 @@ pub trait WeightBackend { fn raw_f32(&mut self, rel: &str, n: usize) -> HipResult; /// Load a bias vector (f32). Only qwen2 attention biases use this today. fn bias(&mut self, rel: &str, n: usize) -> HipResult; + /// Load a bias vector if the checkpoint has one, else `None`. + /// + /// `bias` PANICS on a missing tensor, which is right for qwen2 where the + /// bias is mandatory. Escha's dense export makes it optional per the leaf + /// contract (§1.4) — an export without the end-to-end stage ships none and + /// must still load — so that path needs absence to be a value, not a + /// crash. + fn bias_opt(&mut self, rel: &str, n: usize) -> HipResult>; + /// Escha trellis sidecars for a projection: `escha_rin_eff`, + /// `escha_rout_eff`, and a one-element device pointer table for `w`. + /// + /// `(rin, rout, ptr0)` rather than a typed struct because + /// `hipfire-runtime` must not depend on the arch crate that owns + /// `EschaProj`. `None` for any weight that is not a trellis code. + fn escha_sidecars( + &mut self, + rel: &str, + w: &WeightTensor, + ) -> HipResult>; + /// `n` zeroed 32-bit slots. Used for the escha indexed GEMV's `ids`, which + /// holds INTEGERS declared F32 because `DType` has no integer variant — + /// the same deliberate reinterpretation `EschaMoeTables::ids` documents. + fn zeros_i32(&mut self, n: usize) -> HipResult; + /// `0..n` as 32-bit ints. The grouped escha GEMM's `sorted_slot_index`: + /// a dense linear is one group in slot order, so the permutation is the + /// identity. + fn iota_i32(&mut self, n: usize) -> HipResult; } /// HFQ backend. `norm_bias`: `1.0` (qwen3.5/gemma) or `0.0` (qwen2/llama). @@ -1307,6 +1334,60 @@ impl<'a> WeightBackend for HfqBackend<'a> { .unwrap_or_else(|| panic!("tensor not found: {name}")); dequant_f32(self.gpu, info.quant_type, &data, n) } + fn zeros_i32(&mut self, n: usize) -> HipResult { + self.gpu.upload_f32(&vec![f32::from_bits(0); n], &[n]) + } + + fn iota_i32(&mut self, n: usize) -> HipResult { + let v: Vec = (0..n).map(|i| f32::from_bits(i as u32)).collect(); + self.gpu.upload_f32(&v, &[n]) + } + + + fn escha_sidecars( + &mut self, + rel: &str, + w: &WeightTensor, + ) -> HipResult> { + use rdna_compute::DType; + if !matches!(w.gpu_dtype, DType::Escha2T16 | DType::Escha3T16) { + return Ok(None); + } + let read = |gpu: &mut Gpu, name: &str, want: usize| -> HipResult { + let (info, data) = read_first(self.hfq, name, self.candidates) + .ok_or_else(|| hip_bridge::HipError::new(0, &format!("escha: {name} missing")))?; + let t = dequant_f32(gpu, info.quant_type, &data, want)?; + if t.numel() != want { + return Err(hip_bridge::HipError::new( + 0, + &format!("escha: {name} has {} elements, want {want}", t.numel()), + )); + } + Ok(t) + }; + let base = hfq_plain_name(self.layer, rel); + let rin = read(&mut *self.gpu, &format!("{base}.escha_rin_eff"), w.k)?; + let rout = read(&mut *self.gpu, &format!("{base}.escha_rout_eff"), w.m)?; + let addr = w.buf.buf.as_ptr() as u64; + let ptr0 = self.gpu.upload_raw(&addr.to_le_bytes(), &[1])?; + Ok(Some((rin, rout, ptr0))) + } + + fn bias_opt(&mut self, rel: &str, n: usize) -> HipResult> { + let name = hfq_plain_name(self.layer, rel); + let Some((info, data)) = read_first(self.hfq, &name, self.candidates) else { + return Ok(None); + }; + let t = dequant_f32(self.gpu, info.quant_type, &data, n)?; + if t.numel() != n { + return Err(hip_bridge::HipError::new( + 0, + &format!("bias {name} has {} elements, expected {n}", t.numel()), + )); + } + Ok(Some(t)) + } + fn bias(&mut self, rel: &str, n: usize) -> HipResult { let name = hfq_plain_name(self.layer, rel); let (info, data) = read_first(self.hfq, &name, self.candidates) @@ -1379,6 +1460,31 @@ impl<'a> WeightBackend for ParoBackend<'a> { "ParoBackend: attention biases unsupported", )) } + + /// `None`, not an error: ParoQuant checkpoints simply have no biases, and + /// the optional loader's contract is that absence is a value. + fn bias_opt(&mut self, _rel: &str, _n: usize) -> HipResult> { + Ok(None) + } + + fn zeros_i32(&mut self, n: usize) -> HipResult { + self.gpu.upload_f32(&vec![f32::from_bits(0); n], &[n]) + } + + fn iota_i32(&mut self, n: usize) -> HipResult { + let v: Vec = (0..n).map(|i| f32::from_bits(i as u32)).collect(); + self.gpu.upload_f32(&v, &[n]) + } + + + /// ParoQuant is not escha. + fn escha_sidecars( + &mut self, + _rel: &str, + _w: &WeightTensor, + ) -> HipResult> { + Ok(None) + } } #[cfg(test)] diff --git a/crates/rdna-compute/Cargo.toml b/crates/rdna-compute/Cargo.toml index 61f95849dd..dd47bb8eb0 100644 --- a/crates/rdna-compute/Cargo.toml +++ b/crates/rdna-compute/Cargo.toml @@ -21,6 +21,16 @@ rayon = "1" [dev-dependencies] redline-rocr = { path = "../redline-rocr" } +# G2 parity gate only (test_escha_decode_gpu_vs_cpu): escha_ref::reconstruct is +# the CPU oracle the GPU tile decode is asserted bit-exact against. Note +# hipfire-quantize already carries a dev-dependency back on rdna-compute (for +# its own GL-codebook parity test); Cargo permits a dev-dependency cycle since +# dev-deps are excluded from the normal build graph. +hipfire-quantize = { path = "../hipfire-quantize" } +# Task 9 review Fix 1 divergence proof (test_escha_router_f16_boundary): the +# host-side oracle for "top-k of f16-widened router logits" needs an f32<->f16 +# round trip. Not a production dependency of rdna-compute itself. +half.workspace = true # --------------------------------------------------------------------------- # Archived research probes. diff --git a/crates/rdna-compute/examples/bench_escha_decode_tiles.rs b/crates/rdna-compute/examples/bench_escha_decode_tiles.rs new file mode 100644 index 0000000000..56072b1b8b --- /dev/null +++ b/crates/rdna-compute/examples/bench_escha_decode_tiles.rs @@ -0,0 +1,64 @@ +// SPDX-License-Identifier: Apache-2.0 +// Copyright (c) 2026 Kaden Schutt +// hipfire — see LICENSE and NOTICE in the project root. + +//! Throughput benchmark for `escha_decode_tiles`, kernel-only (no host +//! round trip): upload the K=2 golden code once, then launch the +//! device-resident decode in a loop, syncing once at the end. +//! +//! Temporary tool for the Finding-3 before/after measurement in the Task 7 +//! review (dynamic `words[]` array indexing vs. two direct global loads). +//! Not wired into any gate. +use rdna_compute::{DType, Gpu}; +use std::time::Instant; + +fn main() { + let (ic, oc, k) = (2048usize, 1024usize, 2usize); + let path = format!( + "{}/../hipfire-quantize/tests/data/escha/packed_gu_e0_k2.i16", + env!("CARGO_MANIFEST_DIR") + ); + let code_bytes = std::fs::read(&path).expect("run fetch-goldens.sh first"); + assert_eq!(code_bytes.len(), (ic / 16) * (oc / 16) * 16 * k * 2); + + let mut gpu = Gpu::init().expect("gpu"); + let d_code = gpu + .upload_raw(&code_bytes, &[code_bytes.len() / 2]) + .expect("upload code"); + let d_bare = gpu + .alloc_tensor(&[ic * oc], DType::F16) + .expect("alloc bare"); + + // Warm up: first call JIT-compiles the kernel (or loads it from the + // on-disk cache), which must not be counted. + for _ in 0..8 { + gpu.escha_decode_tiles(&d_code, &d_bare, ic as u32, oc as u32, k as u32) + .expect("warmup decode"); + } + gpu.hip.device_synchronize().expect("sync after warmup"); + + const REPS: u32 = 2000; + let start = Instant::now(); + for _ in 0..REPS { + gpu.escha_decode_tiles(&d_code, &d_bare, ic as u32, oc as u32, k as u32) + .expect("decode"); + } + gpu.hip.device_synchronize().expect("sync after loop"); + let elapsed = start.elapsed(); + + let elems = (ic * oc) as f64; + let per_launch = elapsed.as_secs_f64() / REPS as f64; + // Bytes moved per launch: code read once (ic*oc/16 elements packed at + // k*2 bytes per 16-element tile -> code bytes) + bare fp16 written once. + let code_read_bytes = code_bytes.len() as f64; + let bare_write_bytes = elems * 2.0; + let bytes_per_launch = code_read_bytes + bare_write_bytes; + + println!("escha_decode_tiles K=2 {ic}x{oc}: {REPS} reps in {elapsed:?}"); + println!( + " {:.3} us/launch, {:.3} Gelem/s, {:.3} GB/s (code-read + bare-write)", + per_launch * 1e6, + elems / per_launch / 1e9, + bytes_per_launch / per_launch / 1e9 + ); +} diff --git a/crates/rdna-compute/examples/bench_escha_dense_decode.rs b/crates/rdna-compute/examples/bench_escha_dense_decode.rs new file mode 100644 index 0000000000..e053c41f98 --- /dev/null +++ b/crates/rdna-compute/examples/bench_escha_dense_decode.rs @@ -0,0 +1,169 @@ +// SPDX-License-Identifier: Apache-2.0 +// hipfire — see LICENSE and NOTICE in the project root. +// +//! Escha native GEMV at the shapes the 27B DENSE model decodes at. +//! +//! WHY THIS EXISTS SEPARATELY FROM `bench_escha_grouped_gemm`: that bench runs +//! one real A3B prefill chunk — 2048 slots, ic=2048, oc=1024 — and at that +//! shape the slot-parallel kernel already moves 2.147 GB at ~175 GB/s against +//! this box's ~209-220 GB/s ceiling. It is 80%+ BANDWIDTH-BOUND, so no +//! instruction-level change can move it, and using it to evaluate one is a +//! category error. Decoding the dense 27B is the opposite regime: ONE slot, +//! weights streamed once, measured at 86 GB/s. That is where decode cost shows. +//! +//! Shapes are the real 27B projections (`hidden 5120`, 64 layers), with the +//! documented per-projection K split: `gate_proj` is K=2, `up_proj` K=3. + +use rdna_compute::{DType, Gpu, GpuTensor}; + +struct Rng(u64); +impl Rng { + fn next_u32(&mut self) -> u32 { + self.0 ^= self.0 >> 12; + self.0 ^= self.0 << 25; + self.0 ^= self.0 >> 27; + (self.0.wrapping_mul(0x2545_F491_4F6C_DD1D) >> 32) as u32 + } + fn next_f32(&mut self) -> f32 { + (self.next_u32() as f32 / u32::MAX as f32) - 0.5 + } +} + +fn ptr_table(gpu: &Gpu, owner: &GpuTensor, n_exp: usize, stride_bytes: usize) -> GpuTensor { + let bytes: Vec = (0..n_exp) + .map(|e| owner.buf.as_ptr() as u64 + (e * stride_bytes) as u64) + .flat_map(|p| p.to_ne_bytes()) + .collect(); + gpu.upload_raw(&bytes, &[2 * n_exp]).expect("ptr table") +} + +struct Case { + name: &'static str, + ic: usize, + oc: usize, + trellis_k: u32, +} + +fn main() { + // The 27B dense projections. `gate_proj` K=2 / `up_proj` K=3 is the + // documented trap: they cannot share one kernel call. + let cases = [ + Case { + name: "gate_proj", + ic: 5120, + oc: 17408, + trellis_k: 2, + }, + Case { + name: "up_proj", + ic: 5120, + oc: 17408, + trellis_k: 3, + }, + Case { + name: "down_proj", + ic: 17408, + oc: 5120, + trellis_k: 3, + }, + Case { + name: "in_proj_z", + ic: 5120, + oc: 6144, + trellis_k: 2, + }, + ]; + let slots = 1usize; // decode: one token, one "expert" + let n_exp = 1usize; + let iters: usize = std::env::var("ITERS") + .ok() + .and_then(|v| v.parse().ok()) + .unwrap_or(200); + + let mut gpu = Gpu::init().expect("gpu"); + let mut rng = Rng(0x1234_5678_9ABC_DEF0); + + let ids: Vec = (0..slots).flat_map(|_| 0i32.to_le_bytes()).collect(); + let d_ids = gpu.upload_raw(&ids, &[slots]).expect("ids"); + + println!("escha native GEMV, DENSE DECODE regime (slots=1)"); + println!("box ceiling ~209-220 GB/s; the A3B prefill bench sits at ~175 GB/s\n"); + + for case in &cases { + let words_per_expert = (case.ic / 16) * (case.oc / 16) * 16 * case.trellis_k as usize; + let expert_bytes = words_per_expert * 2; + + let mut crng = Rng(0x5EED_0000_0000_0001 ^ case.ic as u64); + let mut code_bytes = vec![0u8; n_exp * expert_bytes]; + for chunk in code_bytes.chunks_exact_mut(4) { + chunk.copy_from_slice(&crng.next_u32().to_le_bytes()); + } + let d_code = gpu + .upload_raw(&code_bytes, &[code_bytes.len()]) + .expect("upload code"); + drop(code_bytes); + let code_ptrs = ptr_table(&gpu, &d_code, n_exp, expert_bytes); + + let x_bytes: Vec = (0..slots * case.ic) + .flat_map(|_| rng.next_f32().to_le_bytes()) + .collect(); + let d_x = gpu.upload_raw(&x_bytes, &[slots * case.ic]).expect("x"); + drop(x_bytes); + let d_y = gpu.alloc_tensor(&[slots * case.oc], DType::F32).expect("y"); + + // Warm up, then take the MINIMUM over `iters` — the fastest observed + // launch is the one least polluted by other work on the box. + for _ in 0..20 { + gpu.escha_gemv_native_moe_k8_indexed_batched( + &code_ptrs, + &d_ids, + &d_x, + &d_y, + case.oc, + case.ic, + slots, + case.trellis_k, + false, + ) + .expect("gemv warmup"); + } + gpu.hip.device_synchronize().expect("sync"); + + let mut best = f64::INFINITY; + for _ in 0..iters { + gpu.hip.device_synchronize().expect("sync"); + let t = std::time::Instant::now(); + gpu.escha_gemv_native_moe_k8_indexed_batched( + &code_ptrs, + &d_ids, + &d_x, + &d_y, + case.oc, + case.ic, + slots, + case.trellis_k, + false, + ) + .expect("gemv"); + gpu.hip.device_synchronize().expect("sync"); + best = best.min(t.elapsed().as_secs_f64()); + } + + let weights = case.ic * case.oc; + let code_gb = expert_bytes as f64 / 1e9; + let wide = case.ic <= 1536; + println!( + " {:<10} ic={:<6} oc={:<6} K={} {:<6} {:8.1} us \ + {:.3} GB code -> {:6.1} GB/s ({:.2} G weights/s)", + case.name, + case.ic, + case.oc, + case.trellis_k, + if wide { "wide" } else { "narrow" }, + best * 1e6, + code_gb, + code_gb / best, + weights as f64 / best / 1e9, + ); + } +} diff --git a/crates/rdna-compute/examples/bench_escha_grouped_gemm.rs b/crates/rdna-compute/examples/bench_escha_grouped_gemm.rs new file mode 100644 index 0000000000..82723d2d1f --- /dev/null +++ b/crates/rdna-compute/examples/bench_escha_grouped_gemm.rs @@ -0,0 +1,403 @@ +// SPDX-License-Identifier: Apache-2.0 +// hipfire — see LICENSE and NOTICE in the project root. +//! Escha-W2 routed GROUPED GEMM: correctness against the slot-parallel kernel, +//! and the tile-shape sweep that chose `escha_grouped_tile`'s default. +//! +//! Both shipped A3B projections at a real prefill chunk's shape +//! (`n_exp = 256`, `slots = 256 tokens * k=8 = 2048`): +//! +//! | projection | K | M | trellis K | slot-parallel variant | +//! |------------|------|------|-----------|-----------------------| +//! | `gate_up` | 2048 | 1024 | 2 | narrow (`K > 1536`) | +//! | `down` | 512 | 2048 | 3 | wide (`K <= 1536`) | +//! +//! # What the correctness arm asserts, and why it differs per projection +//! +//! The grouped kernel keeps the slot-parallel NARROW form per (token, output +//! row). So: +//! +//! * `gate_up` — the slot-parallel path is narrow too, so the two must agree +//! **BIT FOR BIT**. Asserted as an exact `to_bits()` equality. +//! * `down` — the slot-parallel path is WIDE (four interleaved accumulators +//! folded `(a0+a1)+(a2+a3)`); the grouped one sums the identical 512 products +//! in one sequential chain. They cannot be bit-equal, so this arm reports the +//! max/mean delta and asserts only that it stays inside f32 summation noise +//! for K=512. +//! +//! The codes here are RANDOM bits, not a golden fixture: any 16-bit window +//! decodes to a valid fp16 through `escha_cba`, so random codes exercise the +//! decode arithmetic just as well while letting the fixture be the full 256 +//! experts a real chunk touches. Bit-exactness of the DECODE itself against +//! the frozen oracle is G2/G7's job, not this one's — this example is about +//! the grouping. +//! +//! Run: +//! cargo run --release -p rdna-compute --example bench_escha_grouped_gemm +use rdna_compute::{DType, Gpu, GpuTensor}; + +struct Rng(u64); +impl Rng { + fn next_u32(&mut self) -> u32 { + self.0 ^= self.0 >> 12; + self.0 ^= self.0 << 25; + self.0 ^= self.0 >> 27; + (self.0.wrapping_mul(0x2545_F491_4F6C_DD1D) >> 32) as u32 + } + fn next_f32(&mut self) -> f32 { + (self.next_u32() as f32 / u32::MAX as f32) * 2.0 - 1.0 + } +} + +fn ptr_table(gpu: &Gpu, owner: &GpuTensor, n_exp: usize, stride_bytes: usize) -> GpuTensor { + let bytes: Vec = (0..n_exp) + .map(|e| owner.buf.as_ptr() as u64 + (e * stride_bytes) as u64) + .flat_map(|p| p.to_ne_bytes()) + .collect(); + gpu.upload_raw(&bytes, &[2 * n_exp]).expect("ptr table") +} + +struct Case { + name: &'static str, + ic: usize, + oc: usize, + trellis_k: u32, + /// True when the slot-parallel path picks the WIDE accumulator form + /// (`ic <= 1536`) and the grouped kernel therefore cannot be bit-equal. + slot_parallel_is_wide: bool, +} + +fn main() { + let cases = [ + Case { + name: "gate_up", + ic: 2048, + oc: 1024, + trellis_k: 2, + slot_parallel_is_wide: false, + }, + Case { + name: "down", + ic: 512, + oc: 2048, + trellis_k: 3, + slot_parallel_is_wide: true, + }, + ]; + // One real prefill chunk of the shipped A3B file. + let n_exp = 256usize; + let n_tokens = 256usize; + let k_top = 8usize; + let slots = n_tokens * k_top; + // Sweep set = every instantiation in `escha_moe_gemm_grouped.hip`. + let tiles = [(4, 2), (8, 2), (8, 4), (8, 8), (16, 2), (16, 4)]; + let iters = 5usize; + + let mut gpu = Gpu::init().expect("gpu"); + let mut failures = 0usize; + + // Routing: each token picks k_top DISTINCT experts, as the router does. A + // token that could pick the same expert twice would put two slots of the + // same token in one group, which is legal for the kernel but is not the + // distribution the tile shape is being tuned for. + let mut rng = Rng(0xE5C8_A000_0000_0001); + let mut ids = Vec::with_capacity(slots); + for _ in 0..n_tokens { + let mut chosen: Vec = Vec::with_capacity(k_top); + while chosen.len() < k_top { + let e = (rng.next_u32() as usize % n_exp) as i32; + if !chosen.contains(&e) { + chosen.push(e); + } + } + ids.extend_from_slice(&chosen); + } + let mut hist = vec![0usize; n_exp]; + for &e in &ids { + hist[e as usize] += 1; + } + let live = hist.iter().filter(|c| **c > 0).count(); + println!( + "fixture: n_exp={n_exp} slots={slots} live experts={live} group size min/mean/max = \ + {}/{:.1}/{}", + hist.iter().filter(|c| **c > 0).min().unwrap(), + slots as f64 / live as f64, + hist.iter().max().unwrap() + ); + + let id_bytes: Vec = ids.iter().flat_map(|v| v.to_le_bytes()).collect(); + let d_ids = gpu.upload_raw(&id_bytes, &[slots]).expect("ids"); + + // Sort scratch — `block_m = 1`, so no padding: `expert_offsets` is the + // exact exclusive scan and `sorted_slot_index` is a pure permutation. + let d_counts = gpu.alloc_tensor(&[n_exp], DType::F32).expect("counts"); + let d_offsets = gpu.alloc_tensor(&[n_exp + 1], DType::F32).expect("offsets"); + let d_sorted = gpu.alloc_tensor(&[slots], DType::F32).expect("sorted"); + let d_tile_ids = gpu.alloc_tensor(&[slots], DType::F32).expect("tile ids"); + let d_inv = gpu.alloc_tensor(&[slots], DType::F32).expect("inv"); + gpu.moe_scatter_fused_k8( + &d_ids, + &d_counts, + &d_offsets, + &d_sorted, + &d_tile_ids, + &d_inv, + slots, + n_exp, + slots, + 1, + ) + .expect("scatter"); + gpu.hip.device_synchronize().expect("sync"); + + for case in &cases { + let words_per_expert = (case.ic / 16) * (case.oc / 16) * 16 * case.trellis_k as usize; + let expert_bytes = words_per_expert * 2; + let mut crng = Rng(0x5EED_0000_0000_0001 ^ case.ic as u64); + let mut code_bytes = vec![0u8; n_exp * expert_bytes]; + for chunk in code_bytes.chunks_exact_mut(4) { + chunk.copy_from_slice(&crng.next_u32().to_le_bytes()); + } + let d_code = gpu + .upload_raw(&code_bytes, &[code_bytes.len()]) + .expect("upload code"); + drop(code_bytes); + let code_ptrs = ptr_table(&gpu, &d_code, n_exp, expert_bytes); + + let x: Vec = (0..slots * case.ic).map(|_| rng.next_f32()).collect(); + let x_bytes: Vec = x.iter().flat_map(|v| v.to_le_bytes()).collect(); + let d_x = gpu.upload_raw(&x_bytes, &[slots * case.ic]).expect("x"); + drop(x_bytes); + let d_y_slot = gpu + .alloc_tensor(&[slots * case.oc], DType::F32) + .expect("y slot"); + let d_y_grp = gpu + .alloc_tensor(&[slots * case.oc], DType::F32) + .expect("y grouped"); + let d_y_wmma = gpu + .alloc_tensor(&[slots * case.oc], DType::F32) + .expect("y wmma"); + + // ── the slot-parallel baseline ─────────────────────────────────────── + // + // Re-measured immediately before EVERY tile below, not once. This box + // drifts: the untouched slot-parallel kernel has been seen to move 25% + // (12.95 -> 16.27 ms) between two consecutive processes, which is + // larger than several of the differences the sweep is trying to rank. + // Every speedup printed here is therefore a ratio against a control + // taken in the same second, and the absolute millisecond figures are + // only good to ~15%. + let mut measure_slot = |gpu: &mut Gpu| -> f64 { + let mut best = f64::INFINITY; + for _ in 0..iters { + gpu.hip.device_synchronize().expect("sync"); + let t = std::time::Instant::now(); + gpu.escha_gemv_native_moe_k8_indexed_batched( + &code_ptrs, + &d_ids, + &d_x, + &d_y_slot, + case.oc, + case.ic, + slots, + case.trellis_k, + false, + ) + .expect("slot-parallel gemv"); + gpu.hip.device_synchronize().expect("sync"); + best = best.min(t.elapsed().as_secs_f64() * 1e3); + } + best + }; + let slot_ms = measure_slot(&mut gpu); + let y_slot = gpu.download_f32(&d_y_slot).expect("dl slot"); + // A comparison of two all-zero buffers would report zero difference and + // prove nothing. + let nonzero = y_slot.iter().filter(|v| **v != 0.0).count(); + assert!( + nonzero > y_slot.len() / 2, + "{}: only {nonzero} of {} slot-parallel outputs are non-zero — the fixture is \ + degenerate and every comparison below would be vacuous", + case.name, + y_slot.len() + ); + // Logical expert bytes the slot-parallel kernel moves: every slot reads + // its whole expert. + let slot_w_gb = (slots * expert_bytes) as f64 / 1e9; + let slot_x_gb = ((case.oc / 16) * slots * case.ic * 4) as f64 / 1e9; + println!( + "\n=== {} (K={} M={} tk={} slot-parallel {}) ===", + case.name, + case.ic, + case.oc, + case.trellis_k, + if case.slot_parallel_is_wide { + "wide" + } else { + "narrow" + } + ); + println!( + " slot-parallel : {slot_ms:8.3} ms weights {slot_w_gb:.3} GB + x {slot_x_gb:.3} GB \ + = {:.3} GB -> {:.1} GB/s", + slot_w_gb + slot_x_gb, + (slot_w_gb + slot_x_gb) / (slot_ms / 1e3) + ); + + for &(rows, ctiles) in &tiles { + if case.oc % (16 * ctiles) != 0 { + continue; + } + let ctl_ms = measure_slot(&mut gpu); + // The `_tiled` entry point, not the production one: the latter + // memoises its shape in a `OnceLock`, so a sweep through it would + // measure the first shape six times over. + let mut grp_ms = f64::INFINITY; + let mut launched = true; + for _ in 0..iters { + gpu.hip.device_synchronize().expect("sync"); + let t = std::time::Instant::now(); + let r = gpu.escha_gemm_native_moe_grouped_tiled( + &code_ptrs, + &d_offsets, + &d_sorted, + &d_x, + &d_y_grp, + case.oc, + case.ic, + slots, + n_exp, + case.trellis_k, + false, + rows, + ctiles, + ); + if let Err(e) = r { + println!(" r{rows}c{ctiles} : launch refused: {e}"); + launched = false; + break; + } + gpu.hip.device_synchronize().expect("sync"); + grp_ms = grp_ms.min(t.elapsed().as_secs_f64() * 1e3); + } + if !launched { + continue; + } + // Logical expert bytes: each expert's code is read once per pass of + // ROWS rows through its group. + let passes: usize = hist.iter().map(|g| g.div_ceil(rows)).sum(); + let grp_w_gb = (passes * expert_bytes) as f64 / 1e9; + let grp_x_gb = ((case.oc / (16 * ctiles)) * slots * case.ic * 4) as f64 / 1e9; + println!( + " r{rows}c{ctiles} : {grp_ms:8.3} ms \ + {:.2}x (control {ctl_ms:.3} ms) weights {grp_w_gb:.3} GB + x {grp_x_gb:.3} GB \ + = {:.3} GB -> {:.1} GB/s", + ctl_ms / grp_ms, + grp_w_gb + grp_x_gb, + (grp_w_gb + grp_x_gb) / (grp_ms / 1e3) + ); + + // ── WMMA arm ────────────────────────────────────────────── + // Same grouping inputs, matrix cores instead of scalar FMAs. + // Run once per (rows, ctiles) sweep entry is wasteful — it does + // not take a register tile — so only do it on the first entry. + if rows == 8 && ctiles == 4 { + let mut w_ms = f64::INFINITY; + let mut ok = true; + for _ in 0..iters { + gpu.hip.device_synchronize().expect("sync"); + let t = std::time::Instant::now(); + let r = gpu.escha_gemm_native_moe_grouped_wmma( + &code_ptrs, + &d_offsets, + &d_sorted, + &d_x, + &d_y_wmma, + case.oc, + case.ic, + slots, + n_exp, + case.trellis_k, + false, + ); + if let Err(e) = r { + println!(" wmma : launch refused: {e}"); + ok = false; + break; + } + gpu.hip.device_synchronize().expect("sync"); + w_ms = w_ms.min(t.elapsed().as_secs_f64() * 1e3); + } + if ok { + let yw = gpu.download_f32(&d_y_wmma).expect("dl wmma"); + let mut worst = 0.0f32; + let mut sum = 0.0f64; + let mut nf = 0usize; + for (a, b) in y_slot.iter().zip(&yw) { + let d = (a - b).abs(); + worst = worst.max(d); + sum += d as f64; + if !b.is_finite() { + nf += 1; + } + } + println!( + " wmma : {w_ms:8.3} ms {:.2}x vs scalar-grouped, {:.2}x vs slot-parallel", + grp_ms / w_ms, + ctl_ms / w_ms + ); + println!( + " vs slot-parallel: max {worst:e}, mean {:e}, non-finite {nf}", + sum / yw.len() as f64 + ); + if nf != 0 { + failures += 1; + } + } + } + + let y_grp = gpu.download_f32(&d_y_grp).expect("dl grouped"); + let mut diff = 0usize; + let mut worst = 0.0f32; + let mut sum_abs = 0.0f64; + for (a, b) in y_slot.iter().zip(&y_grp) { + if a.to_bits() != b.to_bits() { + diff += 1; + } + let d = (a - b).abs(); + worst = worst.max(d); + sum_abs += d as f64; + } + let mean = sum_abs / y_grp.len() as f64; + let nonfinite = y_grp.iter().filter(|v| !v.is_finite()).count(); + println!( + " vs slot-parallel: {diff} differing floats of {} (max {worst:e}, \ + mean {mean:e}), non-finite {nonfinite}", + y_grp.len() + ); + if nonfinite != 0 { + failures += 1; + } + if case.slot_parallel_is_wide { + // Cannot be bit-equal by construction (see the module docs); + // hold it to f32 summation noise instead. + if worst > 1e-3 { + println!(" FAIL: delta {worst:e} is beyond f32 summation noise"); + failures += 1; + } + } else if diff != 0 { + println!(" FAIL: the narrow arm must be bit-identical"); + failures += 1; + } + } + + for t in [d_code, code_ptrs, d_x, d_y_slot, d_y_grp] { + let _ = gpu.free_tensor(t); + } + } + + for t in [d_ids, d_counts, d_offsets, d_sorted, d_tile_ids, d_inv] { + let _ = gpu.free_tensor(t); + } + assert_eq!(failures, 0, "grouped GEMM sweep had {failures} failures"); + println!("\nbench_escha_grouped_gemm: OK"); +} diff --git a/crates/rdna-compute/examples/bench_escha_h128.rs b/crates/rdna-compute/examples/bench_escha_h128.rs new file mode 100644 index 0000000000..e783b78981 --- /dev/null +++ b/crates/rdna-compute/examples/bench_escha_h128.rs @@ -0,0 +1,273 @@ +// SPDX-License-Identifier: Apache-2.0 +// Copyright (c) 2026 Kaden Schutt +// hipfire — see LICENSE and NOTICE in the project root. + +//! Throughput benchmark for `escha_h128_in` / `escha_h128_out`, kernel-only +//! (no host round trip): upload fixed input buffers once, then launch the +//! device-resident `Gpu::escha_h128` in a loop, syncing once at the end. +//! +//! These kernels run on every token, on both sides of every escha matmul — +//! unlike `escha_decode_tiles` (once per expert at load time), this is the +//! hot decode path. Used for the Task 8 naive-vs-parallel-butterfly +//! before/after measurement. Not wired into any gate. +//! +//! Task 8 review fix (finding 3): a single n=2048 shape does not establish +//! whether the parallel-butterfly speedup survives at production widths, and +//! a raw us/launch number conflates real kernel work with fixed launch +//! overhead. This version: +//! - launches a genuinely empty kernel (`escha_h128_noop`, same grid/block +//! config as the real launch, zero kernarg bytes) at each shape to +//! measure the launch-overhead floor, and reports both the raw and the +//! overhead-subtracted per-launch time; +//! - sweeps n = 2048, 6144, 17408 (16, 48, 136 blocks of 128) — production +//! escha hidden/intermediate widths, not just the smallest one. +use hip_bridge::KernargBlob; +use rdna_compute::{DType, Gpu}; +use std::ffi::c_void; +use std::time::Instant; + +/// A completely empty kernel launched with the same grid/block config as +/// `escha_h128_in`/`escha_h128_out`, used only to measure the fixed +/// per-launch overhead (queue submission, dispatch packet, sync) that has +/// nothing to do with the H128 butterfly itself. Takes no arguments, so its +/// kernarg blob is zero bytes — it does not touch any memory. +const NOOP_SRC: &str = r#" +extern "C" __global__ void escha_h128_noop() {} +"#; + +/// The pre-parallelisation kernel (commit `2a73edd46`, before +/// `61e3ab8bb`'s parallel butterfly), reproduced verbatim here under +/// renamed entry points (`_naive` suffix) so this benchmark can measure the +/// naive-vs-parallel speedup at production widths without touching the +/// shipped `escha_h128.hip` (which now only carries the parallel version — +/// there is no live "before" to compare against otherwise). This is a +/// benchmark-only artifact, not part of the shipped kernel; it must never be +/// registered in `kernels.rs`/`dispatch.rs`. +const NAIVE_SRC: &str = r#" +#include +#include + +#define ESCHA_RS 0.0883883476f + +__device__ __forceinline__ __half f2h_rne_naive(float v) { + if (v == 0.0f) { + unsigned int bits = __float_as_uint(v); + return __ushort_as_half((unsigned short)(bits >> 16)); + } + return __float2half(v); +} + +__device__ __forceinline__ void h128_block_naive(float* v) { + for (int h = 1; h < 128; h <<= 1) { + for (int i = 0; i < 128; i += (h << 1)) { + for (int j = i; j < i + h; ++j) { + float a = v[j], b = v[j + h]; + v[j] = a + b; + v[j + h] = a - b; + } + } + } +} + +extern "C" __global__ void escha_h128_in_naive( + const float* __restrict__ x, const float* __restrict__ rin, + __half* __restrict__ xh, int n) { + __shared__ float s[128]; + int g = blockIdx.x, t = threadIdx.x; + int idx = g * 128 + t; + if (idx >= n) return; + s[t] = x[idx] * rin[idx]; + __syncthreads(); + if (t == 0) h128_block_naive(s); + __syncthreads(); + xh[idx] = f2h_rne_naive(s[t] * ESCHA_RS); +} + +extern "C" __global__ void escha_h128_out_naive( + const float* __restrict__ mid, const float* __restrict__ rout, + __half* __restrict__ y, int n) { + __shared__ float s[128]; + int g = blockIdx.x, t = threadIdx.x; + int idx = g * 128 + t; + if (idx >= n) return; + s[t] = mid[idx]; + __syncthreads(); + if (t == 0) h128_block_naive(s); + __syncthreads(); + y[idx] = f2h_rne_naive(s[t] * ESCHA_RS * rout[idx]); +} +"#; + +const REPS: u32 = 20_000; + +/// Time `REPS` back-to-back launches of `f`, returning us/launch. `f` must +/// not itself synchronize; the caller syncs once after the loop. +fn time_launches(gpu: &mut Gpu, mut f: impl FnMut(&mut Gpu)) -> f64 { + // Warm up: first call JIT-compiles (or loads from the on-disk cache), + // which must not be counted. + for _ in 0..8 { + f(gpu); + } + gpu.hip.device_synchronize().expect("sync after warmup"); + + let start = Instant::now(); + for _ in 0..REPS { + f(gpu); + } + gpu.hip.device_synchronize().expect("sync after loop"); + start.elapsed().as_secs_f64() / REPS as f64 +} + +fn main() { + let mut gpu = Gpu::init().expect("gpu"); + gpu.ensure_kernel_public("escha_h128_noop", NOOP_SRC, "escha_h128_noop") + .expect("jit noop"); + gpu.ensure_kernel_public("escha_h128_naive", NAIVE_SRC, "escha_h128_in_naive") + .expect("jit naive in"); + gpu.ensure_kernel_public("escha_h128_naive", NAIVE_SRC, "escha_h128_out_naive") + .expect("jit naive out"); + + println!( + "{:<18} {:>10} {:>16} {:>22} {:>10}", + "shape", "kernel", "us/launch", "overhead-sub us", "GB/s" + ); + + // Production escha hidden/intermediate widths (all multiples of 128): + // 2048 matches the packed_gu_e0_k2 golden fixture's `ic`; 6144 and 17408 + // are the wider matmul dimensions in the real model, included so the + // occupancy trend (16 -> 48 -> 136 blocks) is visible, not just the + // smallest case. + for n in [2048usize, 6144, 17408] { + let blocks = (n / 128) as u32; + + let x: Vec = (0..n).map(|i| ((i * 37) as f32 * 0.017).sin()).collect(); + let rin: Vec = (0..n) + .map(|i| if i % 3 == 0 { -0.0023 } else { 0.0023 }) + .collect(); + let mut rout: Vec = (0..n).map(|i| 1.0 + (i % 5) as f32 * 0.1).collect(); + // Pruned channels, same convention as the G3 gate fixture. Both + // indices are < 2048, the smallest shape swept, so they exist at + // every n benchmarked here. + rout[7] = 0.0; + rout[1000] = 0.0; + + let x_bytes: Vec = x.iter().flat_map(|v| v.to_le_bytes()).collect(); + let rin_bytes: Vec = rin.iter().flat_map(|v| v.to_le_bytes()).collect(); + let rout_bytes: Vec = rout.iter().flat_map(|v| v.to_le_bytes()).collect(); + let d_x = gpu.upload_raw(&x_bytes, &[n]).expect("upload x"); + let d_rin = gpu.upload_raw(&rin_bytes, &[n]).expect("upload rin"); + let d_rout = gpu.upload_raw(&rout_bytes, &[n]).expect("upload rout"); + let d_out_in = gpu.alloc_tensor(&[n], DType::F16).expect("alloc out_in"); + let d_out_out = gpu.alloc_tensor(&[n], DType::F16).expect("alloc out_out"); + + // Launch-overhead floor for this grid/block config: an empty kernel, + // same [blocks, 1, 1] x [128, 1, 1] launch shape, zero kernarg bytes. + let overhead_us = time_launches(&mut gpu, |gpu| { + let mut kb = KernargBlob::new(); + gpu.launch_kernel_blob( + "escha_h128_noop", + [blocks, 1, 1], + [128, 1, 1], + 0, + kb.as_mut_slice(), + ) + .expect("noop launch"); + }) * 1e6; + + let per_launch_in = time_launches(&mut gpu, |gpu| { + gpu.escha_h128("escha_h128_in", &d_x, &d_rin, &d_out_in) + .expect("h128 in"); + }) * 1e6; + + let per_launch_out = time_launches(&mut gpu, |gpu| { + gpu.escha_h128("escha_h128_out", &d_x, &d_rout, &d_out_out) + .expect("h128 out"); + }) * 1e6; + + // Naive (pre-parallelisation) kernel at the same shape, launched + // directly via the raw kernarg-blob path since `Gpu::escha_h128` + // always compiles against the shipped `ESCHA_H128_SRC`. + let per_launch_in_naive = time_launches(&mut gpu, |gpu| { + let mut kb = KernargBlob::new(); + kb.push_ptr(d_x.buf.as_ptr() as *const c_void); + kb.push_ptr(d_rin.buf.as_ptr() as *const c_void); + kb.push_ptr(d_out_in.buf.as_ptr() as *const c_void); + kb.push_i32(n as i32); + gpu.launch_kernel_blob( + "escha_h128_in_naive", + [blocks, 1, 1], + [128, 1, 1], + 0, + kb.as_mut_slice(), + ) + .expect("naive in launch"); + }) * 1e6; + + let per_launch_out_naive = time_launches(&mut gpu, |gpu| { + let mut kb = KernargBlob::new(); + kb.push_ptr(d_x.buf.as_ptr() as *const c_void); + kb.push_ptr(d_rout.buf.as_ptr() as *const c_void); + kb.push_ptr(d_out_out.buf.as_ptr() as *const c_void); + kb.push_i32(n as i32); + gpu.launch_kernel_blob( + "escha_h128_out_naive", + [blocks, 1, 1], + [128, 1, 1], + 0, + kb.as_mut_slice(), + ) + .expect("naive out launch"); + }) * 1e6; + + // Bytes moved per launch: two f32 reads (a, vec_in) + one f16 write. + let bytes_per_launch = n as f64 * (4.0 + 4.0 + 2.0); + let shape = format!("n={n} ({blocks} blk)"); + + for (name, raw_us) in [ + ("escha_h128_in", per_launch_in), + ("escha_h128_out", per_launch_out), + ] { + let sub_us = (raw_us - overhead_us).max(0.0); + // GB/s computed against the overhead-subtracted time: the launch + // overhead moves no bytes, so folding it into the denominator + // would understate the kernel's own bandwidth at small n, where + // overhead is a proportionally larger share of the raw time. + let gbs = if sub_us > 0.0 { + bytes_per_launch / (sub_us * 1e-6) / 1e9 + } else { + f64::NAN + }; + println!("{shape:<18} {name:>10} {raw_us:>16.3} {sub_us:>22.3} {gbs:>10.3}"); + } + println!( + "{shape:<18} {:>10} {overhead_us:>16.3} {:>22} {:>10}", + "noop", "-", "-" + ); + + for (name, raw_us) in [ + ("in_naive", per_launch_in_naive), + ("out_naive", per_launch_out_naive), + ] { + let sub_us = (raw_us - overhead_us).max(0.0); + let gbs = if sub_us > 0.0 { + bytes_per_launch / (sub_us * 1e-6) / 1e9 + } else { + f64::NAN + }; + println!("{shape:<18} {name:>10} {raw_us:>16.3} {sub_us:>22.3} {gbs:>10.3}"); + } + + let speedup_in_raw = per_launch_in_naive / per_launch_in; + let speedup_out_raw = per_launch_out_naive / per_launch_out; + let speedup_in_sub = + (per_launch_in_naive - overhead_us).max(0.0) / (per_launch_in - overhead_us).max(1e-9); + let speedup_out_sub = (per_launch_out_naive - overhead_us).max(0.0) + / (per_launch_out - overhead_us).max(1e-9); + println!( + "{shape:<18} speedup(in) raw={speedup_in_raw:.3}x overhead-sub={speedup_in_sub:.3}x" + ); + println!( + "{shape:<18} speedup(out) raw={speedup_out_raw:.3}x overhead-sub={speedup_out_sub:.3}x" + ); + } +} diff --git a/crates/rdna-compute/examples/test_escha_decode_gpu_vs_cpu.rs b/crates/rdna-compute/examples/test_escha_decode_gpu_vs_cpu.rs new file mode 100644 index 0000000000..a5e8e8d787 --- /dev/null +++ b/crates/rdna-compute/examples/test_escha_decode_gpu_vs_cpu.rs @@ -0,0 +1,93 @@ +//! G2: GPU tile decode must match escha_ref::reconstruct EXACTLY in fp16, +//! for both K, at the golden shapes AND at production shapes up to 89M +//! elements. Run: +//! cargo run --release -p rdna-compute --example test_escha_decode_gpu_vs_cpu +use hipfire_quantize::escha_ref; + +/// xorshift64* — same tiny inline PRNG convention used by the other GPU +/// parity examples in this crate (no `rand` dependency). +struct Rng(u64); +impl Rng { + fn next_u32(&mut self) -> u32 { + self.0 ^= self.0 >> 12; + self.0 ^= self.0 << 25; + self.0 ^= self.0 >> 27; + (self.0.wrapping_mul(0x2545F4914F6CDD1D) >> 32) as u32 + } +} + +fn main() { + for (name, ic, oc, k) in [ + ("packed_gu_e0_k2.i16", 2048usize, 1024usize, 2usize), + ("packed_down_e0_k3.i16", 512, 2048, 3), + ] { + let path = format!( + "{}/../hipfire-quantize/tests/data/escha/{name}", + env!("CARGO_MANIFEST_DIR") + ); + let raw = std::fs::read(&path).expect("run fetch-goldens.sh first"); + let code: Vec = raw + .chunks_exact(2) + .map(|c| i16::from_le_bytes([c[0], c[1]])) + .collect(); + let want = escha_ref::reconstruct(&code, ic, oc, k); + + let mut gpu = rdna_compute::Gpu::init().expect("gpu"); + let got = gpu + .escha_decode_tiles_host(&code, ic as u32, oc as u32, k as u32) + .expect("decode"); + + let bad = want.iter().zip(&got).filter(|(a, b)| a != b).count(); + println!("{name}: {bad} mismatched of {} elements", want.len()); + assert_eq!(bad, 0, "{name}: GPU decode diverges from the CPU reference"); + } + println!("G2 PASS"); + + // Widen the gate: the two golden fixtures above are 2048x1024 and + // 512x2048 — both comfortably small. Cover at least one production + // shape from the dense 27B checkpoint at each K, on pseudo-random code, + // to prove the tile/lane indexing generalises past the two golden + // shapes rather than happening to work only for them. + let mut rng = Rng(0xE5CA_5EED_1234_5678u64); + for (ic, oc, k) in [(5120usize, 17408usize, 2usize), (17408, 5120, 3usize)] { + let n_tiles = (ic / 16) * (oc / 16); + let code_len = n_tiles * 16 * k; + let code: Vec = (0..code_len).map(|_| rng.next_u32() as i16).collect(); + let want = escha_ref::reconstruct(&code, ic, oc, k); + + let mut gpu = rdna_compute::Gpu::init().expect("gpu"); + let got = gpu + .escha_decode_tiles_host(&code, ic as u32, oc as u32, k as u32) + .expect("wide-shape decode"); + + let bad = want.iter().zip(&got).filter(|(a, b)| a != b).count(); + println!( + "wide shape {ic}x{oc} K={k}: {bad} mismatched of {} elements", + want.len() + ); + assert_eq!( + bad, 0, + "{ic}x{oc} K={k}: GPU decode diverges from the CPU reference at production scale" + ); + } + println!("wide-shape gate PASS"); + + // The "device-resident vs host-roundtrip equivalence" check that used to + // sit here has been DELETED, not moved. + // + // It compared `escha_decode_tiles_host(...)` against a hand-rolled + // upload -> `escha_decode_tiles` -> download of the same input. But + // `escha_decode_tiles_host` IS that sequence — it calls + // `escha_decode_tiles` internally — so both sides ran the same kernel on + // the same bytes. The only thing it could ever have caught is kernel + // nondeterminism, while its message claimed to prove the two entry points + // equivalent. A check that cannot fail for the reason it names is worse + // than no check: it reads as coverage. + // + // The real gate for the device-resident path is the G2 arm at the top of + // this file. `escha_decode_tiles_host` is a thin wrapper over + // `escha_decode_tiles`, so asserting bit-exactness against + // `escha_ref::reconstruct` — the frozen oracle — already asserts it for + // the device kernel, at both K and at production shapes up to 89M + // elements. Nothing was lost. +} diff --git a/crates/rdna-compute/examples/test_escha_h128_gpu_vs_cpu.rs b/crates/rdna-compute/examples/test_escha_h128_gpu_vs_cpu.rs new file mode 100644 index 0000000000..68ab806ad1 --- /dev/null +++ b/crates/rdna-compute/examples/test_escha_h128_gpu_vs_cpu.rs @@ -0,0 +1,351 @@ +//! G3: the H128 kernels must match escha_ref directly. +//! +//! A round-trip check (H128 . H128 == 128 I) is NOT sufficient — a wrong +//! butterfly order is also self-inverse and would pass it while being wrong. +use hipfire_quantize::escha_ref; +use hipfire_quantize::float16::f16_to_f32; +use rdna_compute::EschaXGroup; + +fn main() { + let n = 2048usize; + let x: Vec = (0..n).map(|i| ((i * 37) as f32 * 0.017).sin()).collect(); + let rin: Vec = (0..n) + .map(|i| if i % 3 == 0 { -0.0023 } else { 0.0023 }) + .collect(); + let mut rout: Vec = (0..n).map(|i| 1.0 + (i % 5) as f32 * 0.1).collect(); + rout[7] = 0.0; + rout[1000] = 0.0; // pruned channels must stay exactly zero + + let want_in = escha_ref::input_transform(&x, &rin); + let want_out = escha_ref::output_transform(&x, &rout); + + let mut gpu = rdna_compute::Gpu::init().expect("gpu"); + let got_in = gpu.escha_h128_in_host(&x, &rin).expect("h128 in"); + let got_out = gpu.escha_h128_out_host(&x, &rout).expect("h128 out"); + + let bad_in = want_in.iter().zip(&got_in).filter(|(a, b)| a != b).count(); + let bad_out = want_out + .iter() + .zip(&got_out) + .filter(|(a, b)| a != b) + .count(); + println!("h128_in : {bad_in} mismatched of {n}"); + println!("h128_out: {bad_out} mismatched of {n}"); + assert_eq!( + f16_to_f32(got_out[7]), + 0.0, + "pruned channel 7 must be exactly zero" + ); + assert_eq!( + f16_to_f32(got_out[1000]), + 0.0, + "pruned channel 1000 must be exactly zero" + ); + assert_eq!(bad_in, 0); + assert_eq!(bad_out, 0); + + batched_vs_ref(&mut gpu, n); + println!("G3 PASS"); +} + +/// G3b (Task 10): the BATCHED forms must agree with `escha_ref` element for +/// element, exactly as the per-expert forms do. Batching is an indexing +/// change — slot `s` reads row `ids[s]` of the resident `[E, n]` table — so a +/// wrong index would still produce plausible, full-magnitude output. It is +/// gated against the oracle, never against the per-expert kernel alone. +/// +/// The slot list deliberately repeats an expert (two slots on row 1) and is +/// not sorted: nothing in the kernel may assume distinct or ordered ids. +fn batched_vs_ref(gpu: &mut rdna_compute::Gpu, n: usize) { + use rdna_compute::DType; + let n_exp = 5usize; + let ids: Vec = vec![3, 1, 0, 1]; + let slots = ids.len(); + + // Per-expert transform vectors. Expert 1 (used twice) carries the pruned + // channels so the zero contract is exercised through the batched index. + let mut r_in = vec![0.0f32; n_exp * n]; + let mut r_out = vec![0.0f32; n_exp * n]; + for e in 0..n_exp { + for i in 0..n { + r_in[e * n + i] = if (i + e) % 3 == 0 { -0.0023 } else { 0.0019 } * (1.0 + e as f32); + r_out[e * n + i] = 1.0 + ((i + 2 * e) % 5) as f32 * 0.1; + } + } + r_out[1 * n + 7] = 0.0; + r_out[1 * n + 1000] = 0.0; + + // Broadcast activation (gate_up input side) and a per-slot one (down side). + let x1: Vec = (0..n).map(|i| ((i * 37) as f32 * 0.017).sin()).collect(); + let xk: Vec = (0..slots * n) + .map(|i| ((i * 11) as f32 * 0.013).cos() * 0.5) + .collect(); + + let up = |g: &rdna_compute::Gpu, v: &[f32]| { + let b: Vec = v.iter().flat_map(|x| x.to_le_bytes()).collect(); + g.upload_raw(&b, &[v.len()]).expect("upload") + }; + let d_rin = up(gpu, &r_in); + let d_rout = up(gpu, &r_out); + let d_x1 = up(gpu, &x1); + let d_xk = up(gpu, &xk); + let id_bytes: Vec = ids.iter().flat_map(|v| v.to_le_bytes()).collect(); + let d_ids = gpu.upload_raw(&id_bytes, &[slots]).expect("ids"); + let d_out = gpu.alloc_tensor(&[slots * n], DType::F32).expect("out"); + + // ── in, broadcast x ────────────────────────────────────────────────── + gpu.escha_h128_batched( + "escha_h128_in_batched", + &d_x1, + &d_rin, + &d_ids, + &d_out, + n, + slots, + EschaXGroup::Broadcast, + ) + .expect("in batched (broadcast)"); + let got = gpu.download_f32(&d_out).expect("dl"); + let mut bad = 0usize; + for (s, &e) in ids.iter().enumerate() { + let want = escha_ref::input_transform(&x1, &r_in[e as usize * n..(e as usize + 1) * n]); + for i in 0..n { + if got[s * n + i] != f16_to_f32(want[i]) { + bad += 1; + } + } + } + println!( + "h128_in_batched (broadcast x): {bad} mismatched of {}", + slots * n + ); + assert_eq!(bad, 0); + + // ── in, per-slot x ─────────────────────────────────────────────────── + gpu.escha_h128_batched( + "escha_h128_in_batched", + &d_xk, + &d_rin, + &d_ids, + &d_out, + n, + slots, + EschaXGroup::PerSlot, + ) + .expect("in batched (per-slot)"); + let got = gpu.download_f32(&d_out).expect("dl"); + let mut bad = 0usize; + for (s, &e) in ids.iter().enumerate() { + let want = escha_ref::input_transform( + &xk[s * n..(s + 1) * n], + &r_in[e as usize * n..(e as usize + 1) * n], + ); + for i in 0..n { + if got[s * n + i] != f16_to_f32(want[i]) { + bad += 1; + } + } + } + println!( + "h128_in_batched (per-slot x): {bad} mismatched of {}", + slots * n + ); + assert_eq!(bad, 0); + + // ── in, GROUPED x (Task perf-3: batched prefill) ───────────────────── + // + // The third `x_group` case, gated exactly like the other two: bit-exact + // against `escha_ref`, never against the other kernel arms. + // + // This is the case batched prefill needs and neither existing case covers: + // `slots = n_tokens * k` laid out token-major, with all `k` slots of a + // token reading THAT TOKEN's activation row. `slots = 4`, `g = 2` means + // two tokens of two experts each: slots 0,1 read x row 0 and slots 2,3 + // read x row 1. + // + // Getting the group arithmetic wrong is silent, not loud. `slot * n` + // (the PerSlot formula) against a `[slots/g, n]` buffer reads off the end + // for the later slots; `slot / g` with the wrong `g` reads a real, + // full-magnitude activation belonging to a DIFFERENT token — plausible + // output, wrong answer. Only an oracle comparison distinguishes them, + // which is why this is here rather than a round-trip or a self-check. + // + // The ids list still repeats an expert and is still unsorted, and the two + // tokens' rows differ, so an implementation that broadcast row 0 to + // everything (the `x_group <= 0` arm firing by accident) fails on slots + // 2 and 3 rather than passing by luck. + let group = 2usize; + assert_eq!(slots % group, 0, "grouped case needs g | slots"); + let n_tokens = slots / group; + let xg: Vec = (0..n_tokens * n) + .map(|i| ((i * 23) as f32 * 0.0091).sin() * 0.75) + .collect(); + let d_xg = up(gpu, &xg); + gpu.escha_h128_batched( + "escha_h128_in_batched", + &d_xg, + &d_rin, + &d_ids, + &d_out, + n, + slots, + EschaXGroup::Grouped(group), + ) + .expect("in batched (grouped)"); + let got = gpu.download_f32(&d_out).expect("dl"); + let mut bad = 0usize; + for (s, &e) in ids.iter().enumerate() { + let tok = s / group; + let want = escha_ref::input_transform( + &xg[tok * n..(tok + 1) * n], + &r_in[e as usize * n..(e as usize + 1) * n], + ); + for i in 0..n { + if got[s * n + i] != f16_to_f32(want[i]) { + bad += 1; + } + } + } + println!( + "h128_in_batched (grouped x, g={group}): {bad} mismatched of {}", + slots * n + ); + assert_eq!(bad, 0); + + // Grouped(1) must be byte-identical to PerSlot — the backward-compatibility + // claim the kernel change rests on, checked rather than asserted in a + // comment. (Grouped(0) is rejected by the wrapper, not silently treated as + // broadcast, so there is nothing to check for it here.) + gpu.escha_h128_batched( + "escha_h128_in_batched", + &d_xk, + &d_rin, + &d_ids, + &d_out, + n, + slots, + EschaXGroup::Grouped(1), + ) + .expect("in batched (grouped g=1)"); + let got_g1 = gpu.download_f32(&d_out).expect("dl"); + gpu.escha_h128_batched( + "escha_h128_in_batched", + &d_xk, + &d_rin, + &d_ids, + &d_out, + n, + slots, + EschaXGroup::PerSlot, + ) + .expect("in batched (per-slot recheck)"); + let got_ps = gpu.download_f32(&d_out).expect("dl"); + assert_eq!( + got_g1, got_ps, + "Grouped(1) must be byte-identical to PerSlot" + ); + println!("h128_in_batched: Grouped(1) == PerSlot, byte-identical"); + + // A group that does not divide `slots` must be REFUSED, not truncated: + // a ragged tail would read a wrong row for the last slots, silently. + assert!( + gpu.escha_h128_batched( + "escha_h128_in_batched", + &d_xg, + &d_rin, + &d_ids, + &d_out, + n, + slots, + EschaXGroup::Grouped(3), + ) + .is_err(), + "a group size that does not divide slots must be rejected" + ); + + // ── out ────────────────────────────────────────────────────────────── + gpu.escha_h128_batched( + "escha_h128_out_batched", + &d_xk, + &d_rout, + &d_ids, + &d_out, + n, + slots, + EschaXGroup::PerSlot, + ) + .expect("out batched"); + let got = gpu.download_f32(&d_out).expect("dl"); + let mut bad = 0usize; + for (s, &e) in ids.iter().enumerate() { + let want = escha_ref::output_transform( + &xk[s * n..(s + 1) * n], + &r_out[e as usize * n..(e as usize + 1) * n], + ); + for i in 0..n { + if got[s * n + i] != f16_to_f32(want[i]) { + bad += 1; + } + } + } + println!("h128_out_batched: {bad} mismatched of {}", slots * n); + assert_eq!(bad, 0); + // Pruned channels of expert 1 land in slots 1 and 3. + for s in [1usize, 3] { + for ch in [7usize, 1000] { + assert_eq!( + got[s * n + ch], + 0.0, + "slot {s} channel {ch} must be exactly zero" + ); + } + } + + // ── batched SwiGLU ─────────────────────────────────────────────────── + // Input must be f16-representable (it is the H128 output in production); + // reuse the transform result above so the gate feeds the real shape. + let inter = n / 2; + let d_h = gpu + .alloc_tensor(&[slots * inter], DType::F32) + .expect("h buf"); + gpu.escha_swiglu_batched(&d_out, &d_h, inter, slots) + .expect("swiglu"); + let got_h = gpu.download_f32(&d_h).expect("dl"); + let mut ulp1 = 0usize; + let mut worse = 0usize; + for s in 0..slots { + let bits: Vec = got[s * n..(s + 1) * n] + .iter() + .map(|&v| escha_ref::f16_rne(v)) + .collect(); + let want = escha_ref::swiglu(&bits, inter); + for i in 0..inter { + let gb = escha_ref::f16_rne(got_h[s * inter + i]); + if gb != want[i] { + let d = (gb as i32 - want[i] as i32).abs(); + if d == 1 { + ulp1 += 1; + } else { + worse += 1; + } + } + } + } + println!( + "swiglu_batched vs escha_ref::swiglu: {ulp1} at 1 f16 ulp, {worse} worse, of {}", + slots * inter + ); + // Device `expf` and Rust `f32::exp` are both <1 ulp but not the same + // function, so a handful of values straddle an f16 rounding boundary. + // Anything beyond 1 ulp is a real defect (wrong half, wrong slot stride, + // missing rounding), not a libm difference. + assert_eq!( + worse, 0, + "swiglu differs from the oracle by more than 1 ulp" + ); + assert!( + ulp1 * 1000 <= slots * inter, + "swiglu 1-ulp mismatches {ulp1} exceed 0.1% of {} — not a libm difference", + slots * inter + ); +} diff --git a/crates/rdna-compute/examples/test_escha_native_gemv_gpu_vs_cpu.rs b/crates/rdna-compute/examples/test_escha_native_gemv_gpu_vs_cpu.rs new file mode 100644 index 0000000000..9bb56035b2 --- /dev/null +++ b/crates/rdna-compute/examples/test_escha_native_gemv_gpu_vs_cpu.rs @@ -0,0 +1,415 @@ +//! G7: the FUSED routed GEMV — the one that decodes escha's trellis code +//! inside the matmul instead of reading an expanded copy — must produce +//! **bit-identical** f32 output to +//! +//! 1. the same GEMV reading the **F16 expert store** +//! (`escha_decode_tiles` -> `escha_bare_to_f16`, i.e. the exactly-decoded +//! weights with nothing re-quantised), and +//! 2. a CPU replica built from **`escha_ref`**, the frozen oracle, summed in +//! the kernel's own order. +//! +//! Run: +//! cargo run --release -p rdna-compute --example test_escha_native_gemv_gpu_vs_cpu +//! +//! # Why the comparison is against F16 and NOT against the Q8_0 store +//! +//! The Q8_0 expert store is a LOSSY re-quantisation of the decoded weight +//! (that loss is the dominant term in the G4 block gate: 2.633e-4 max against +//! the F32 arm's 1.828e-4). The fused kernel uses the decoded fp16 value +//! itself. Asserting equality against Q8_0 would therefore be asserting +//! something false; the arm that holds the same values the fused kernel +//! decodes is the F16 store, and that is what arm 1 compares against. +//! +//! # Why arm 2 exists as well +//! +//! Arm 1 alone would be satisfied by two GPU paths that are wrong in the same +//! way — they share `escha_decode_tiles`' idea of where a weight lives. Arm 2 +//! closes that: it takes the weights from `escha_ref::reconstruct`, the CPU +//! oracle that has been frozen since commit 11 and that every bit-exact claim +//! in this port rests on, and it re-does the summation in the kernel's exact +//! order (lane-strided partial sums, the accumulator count the variant uses, +//! the `__shfl_down` ladder). So arm 2 checks the DECODE against the oracle +//! and the ACCUMULATION against the transcription contract at once. +//! +//! Arm 2 also reports WHICH f32 contraction the compiler chose (fused +//! multiply-add or separate multiply and add) rather than assuming one: the +//! two differ in the last bit and the gate is an exact-equality gate, so +//! guessing would make it flaky rather than wrong. +//! +//! # Coverage +//! +//! Both shipped A3B projections, i.e. both trellis orders AND both accumulator +//! variants, on the SHIPPED golden code (not synthetic): +//! +//! | projection | ic (K) | oc (M) | trellis K | variant | +//! |---|---|---|---|---| +//! | `gate_up` | 2048 | 1024 | 2 | narrow (K > 1536) | +//! | `down` | 512 | 2048 | 3 | wide (K <= 1536) | +//! +//! `n_exp = 4` experts, each a distinct rotation of the golden code, and eight +//! slots whose ids repeat and are out of order — so a kernel that ignored +//! `expert_ptrs[topk_indices[krank]]`, or that mixed slots up, fails here +//! rather than passing on an accidentally-uniform fixture. +use hipfire_quantize::escha_ref; +use rdna_compute::{DType, Gpu, GpuTensor}; + +/// xorshift64* — the same inline PRNG the other GPU parity examples in this +/// crate use (no `rand` dependency). +struct Rng(u64); +impl Rng { + fn next_u32(&mut self) -> u32 { + self.0 ^= self.0 >> 12; + self.0 ^= self.0 << 25; + self.0 ^= self.0 >> 27; + (self.0.wrapping_mul(0x2545_F491_4F6C_DD1D) >> 32) as u32 + } + /// Activations in roughly the range the H128 input transform emits. + fn next_f32(&mut self) -> f32 { + (self.next_u32() as f32 / u32::MAX as f32) * 2.0 - 1.0 + } +} + +fn f16_to_f32(bits: u16) -> f32 { + half::f16::from_bits(bits).to_f32() +} + +/// How the kernel folds one product into an accumulator. HIP compiles with +/// `-ffp-contract=fast` by default, so `acc += w * x` becomes a single-rounding +/// `v_fmac_f32`; but that is a toolchain default, not a guarantee, and the two +/// forms differ in the last bit. The gate determines which one the build +/// actually produced instead of assuming — see `main`. +#[derive(Clone, Copy, PartialEq, Eq, Debug)] +enum Fold { + Fma, + MulAdd, +} +impl Fold { + #[inline(always)] + fn apply(self, acc: f32, w: f32, x: f32) -> f32 { + match self { + Fold::Fma => w.mul_add(x, acc), + Fold::MulAdd => acc + w * x, + } + } +} + +/// The `__shfl_down(sum, 16/8/4/2/1)` ladder both kernel variants end with, +/// evaluated for lane 0 — the only lane whose result is stored. +/// +/// Every value lane 0 consumes comes from an in-range lane, so the ladder is +/// exactly a balanced binary tree over the 32 partial sums and the +/// out-of-range-returns-self behaviour of `__shfl_down` never enters. +fn warp_reduce(mut lanes: [f32; 32]) -> f32 { + let mut width = 16; + while width > 0 { + for t in 0..width { + lanes[t] += lanes[t + width]; + } + width >>= 1; + } + lanes[0] +} + +/// CPU replica of `escha_gemv_native_*_moe_k8_indexed_batched` for one output +/// row: the NARROW form (one accumulator, `K > 1536`). +/// +/// `w_row[i]` is the weight for contraction index `i`, already fp16-decoded. +fn cpu_row_narrow(w_row: &dyn Fn(usize) -> f32, x: &[f32], k: usize, fold: Fold) -> f32 { + let blocks = k / 32; + let mut lanes = [0.0f32; 32]; + for (t, lane) in lanes.iter_mut().enumerate() { + let mut sum = 0.0f32; + for bi in 0..blocks { + let i = bi * 32 + t; + sum = fold.apply(sum, w_row(i), x[i]); + } + *lane = sum; + } + warp_reduce(lanes) +} + +/// CPU replica of the WIDE form (`K <= 1536`): four interleaved accumulators +/// folded `(acc0 + acc1) + (acc2 + acc3)`, tail blocks landing in `acc[t]`. +fn cpu_row_wide(w_row: &dyn Fn(usize) -> f32, x: &[f32], k: usize, fold: Fold) -> f32 { + let blocks = k / 32; + let quads = blocks >> 2; + let tail = blocks & 3; + let mut lanes = [0.0f32; 32]; + for (t, lane) in lanes.iter_mut().enumerate() { + let mut acc = [0.0f32; 4]; + for q in 0..quads { + let bi = q << 2; + for (s, a) in acc.iter_mut().enumerate() { + let i = (bi + s) * 32 + t; + *a = fold.apply(*a, w_row(i), x[i]); + } + } + for s in 0..tail { + let bi = (quads << 2) + s; + let i = bi * 32 + t; + let contrib = w_row(i) * x[i]; + // Verbatim from the kernel, `t == 3` included in its unreachability. + if s < 3 { + acc[s] += contrib; + } + } + *lane = (acc[0] + acc[1]) + (acc[2] + acc[3]); + } + warp_reduce(lanes) +} + +/// The `[n_exp]` u64 weight-base table the indexed GEMVs index with +/// `expert_ptrs[topk_indices[krank]]`, packed into an F32 tensor (2 f32 per +/// pointer) exactly as the model loader packs it. +fn ptr_table(gpu: &Gpu, owner: &GpuTensor, n_exp: usize, stride_bytes: usize) -> GpuTensor { + let bytes: Vec = (0..n_exp) + .map(|e| owner.buf.as_ptr() as u64 + (e * stride_bytes) as u64) + .flat_map(|p| p.to_ne_bytes()) + .collect(); + gpu.upload_raw(&bytes, &[2 * n_exp]).expect("ptr table") +} + +struct Case { + name: &'static str, + fixture: &'static str, + ic: usize, + oc: usize, + trellis_k: usize, +} + +fn main() { + let cases = [ + Case { + name: "gate_up", + fixture: "packed_gu_e0_k2.i16", + ic: 2048, + oc: 1024, + trellis_k: 2, + }, + Case { + name: "down", + fixture: "packed_down_e0_k3.i16", + ic: 512, + oc: 2048, + trellis_k: 3, + }, + ]; + + let n_exp = 4usize; + let slots = 8usize; + // Repeating, out-of-order ids: an implementation that dropped the indirection + // (or that used `krank` where it meant `ids[krank]`) cannot pass. + let ids: [i32; 8] = [3, 0, 2, 1, 0, 3, 1, 2]; + + let mut gpu = Gpu::init().expect("gpu"); + let mut failures = 0usize; + + for case in &cases { + let path = format!( + "{}/../hipfire-quantize/tests/data/escha/{}", + env!("CARGO_MANIFEST_DIR"), + case.fixture + ); + let raw = std::fs::read(&path).expect("run fetch-goldens.sh first"); + let base: Vec = raw + .chunks_exact(2) + .map(|c| i16::from_le_bytes([c[0], c[1]])) + .collect(); + let words_per_expert = (case.ic / 16) * (case.oc / 16) * 16 * case.trellis_k; + assert_eq!( + base.len(), + words_per_expert, + "{}: fixture is {} i16, expected {words_per_expert}", + case.name, + base.len() + ); + + // Four DIFFERENT experts out of one fixture: rotating the code stream + // keeps it a valid trellis code (every tile is still 16*K words) while + // making no two experts share a single tile. + let experts: Vec> = (0..n_exp) + .map(|e| { + let shift = (e * 16 * case.trellis_k * 37) % words_per_expert; + let mut v = base.clone(); + v.rotate_left(shift); + v + }) + .collect(); + + // ── the two device-side weight stores ──────────────────────────────── + // NATIVE: the code bytes, verbatim, one contiguous buffer with the + // per-expert slot stride the packed loader uses. + let code_bytes: Vec = experts + .iter() + .flat_map(|e| e.iter().flat_map(|v| v.to_le_bytes())) + .collect(); + let d_code = gpu + .upload_raw(&code_bytes, &[code_bytes.len()]) + .expect("upload code"); + + // F16: the production store path — `escha_decode_tiles` then + // `escha_bare_to_f16` — into the same kind of packed buffer. + let d_f16 = gpu + .alloc_tensor(&[n_exp * case.ic * case.oc], DType::F16) + .expect("alloc f16 store"); + let bare = gpu + .alloc_tensor(&[case.ic * case.oc], DType::F16) + .expect("alloc bare"); + let stage = gpu + .alloc_tensor(&[words_per_expert], DType::F16) + .expect("alloc stage"); + for (e, code) in experts.iter().enumerate() { + let bytes: Vec = code.iter().flat_map(|v| v.to_le_bytes()).collect(); + gpu.hip.memcpy_htod(&stage.buf, &bytes).expect("stage code"); + gpu.escha_decode_tiles( + &stage, + &bare, + case.ic as u32, + case.oc as u32, + case.trellis_k as u32, + ) + .expect("decode"); + let slot = d_f16.sub_offset(e * case.ic * case.oc, case.ic * case.oc); + gpu.escha_bare_to_f16(&bare, &slot, case.ic, case.oc) + .expect("bare->f16"); + } + + // ── pointer tables (the same [n_exp] u64 packing the loader builds) ── + let code_ptrs = ptr_table(&gpu, &d_code, n_exp, words_per_expert * 2); + let f16_ptrs = ptr_table(&gpu, &d_f16, n_exp, case.ic * case.oc * 2); + + let id_bytes: Vec = ids.iter().flat_map(|v| v.to_le_bytes()).collect(); + let d_ids = gpu.upload_raw(&id_bytes, &[slots]).expect("ids"); + + // ── activations: one distinct vector per slot ──────────────────────── + let mut rng = Rng(0x5EED_0000_0000_0001 ^ (case.ic as u64)); + let x: Vec = (0..slots * case.ic).map(|_| rng.next_f32()).collect(); + let x_bytes: Vec = x.iter().flat_map(|v| v.to_le_bytes()).collect(); + let d_x = gpu + .upload_raw(&x_bytes, &[slots * case.ic]) + .expect("upload x"); + let d_y_native = gpu + .alloc_tensor(&[slots * case.oc], DType::F32) + .expect("y native"); + let d_y_f16 = gpu + .alloc_tensor(&[slots * case.oc], DType::F32) + .expect("y f16"); + + gpu.escha_gemv_native_moe_k8_indexed_batched( + &code_ptrs, + &d_ids, + &d_x, + &d_y_native, + case.oc, + case.ic, + slots, + case.trellis_k as u32, + false, + ) + .expect("native gemv"); + gpu.escha_gemv_f16_moe_k8_indexed_batched( + &f16_ptrs, &d_ids, &d_x, &d_y_f16, case.oc, case.ic, slots, + ) + .expect("f16 gemv"); + + let y_native = gpu.download_f32(&d_y_native).expect("dl native"); + let y_f16 = gpu.download_f32(&d_y_f16).expect("dl f16"); + + // ── arm 1: fused native vs the F16 expert store, bit for bit ───────── + let diff_f16 = y_native + .iter() + .zip(&y_f16) + .filter(|(a, b)| a.to_bits() != b.to_bits()) + .count(); + let worst = y_native + .iter() + .zip(&y_f16) + .map(|(a, b)| (a - b).abs()) + .fold(0.0f32, f32::max); + println!( + "{}: fused-native vs F16 store — {diff_f16} differing floats of {} (max |delta| {worst:e})", + case.name, + y_native.len() + ); + if diff_f16 != 0 { + failures += 1; + } + + // ── arm 2: fused native vs escha_ref, summed in the kernel's order ─── + // `reconstruct` is IN-major `[ic, oc]`; the GEMV's output row `o` + // contracts over `i`, so the weight is `bare[i * oc + o]`. + let decoded: Vec> = experts + .iter() + .map(|c| escha_ref::reconstruct(c, case.ic, case.oc, case.trellis_k)) + .collect(); + let wide = case.ic <= 1536; + + // Which f32 contraction did this build produce? Decide it from ONE + // element rather than assuming, then hold the whole gate to it. + let probe_row = |fold: Fold, slot: usize, o: usize| -> f32 { + let w = &decoded[ids[slot] as usize]; + let w_row = |i: usize| f16_to_f32(w[i * case.oc + o]); + let xs = &x[slot * case.ic..(slot + 1) * case.ic]; + if wide { + cpu_row_wide(&w_row, xs, case.ic, fold) + } else { + cpu_row_narrow(&w_row, xs, case.ic, fold) + } + }; + let fold = if probe_row(Fold::Fma, 0, 0).to_bits() == y_native[0].to_bits() { + Fold::Fma + } else { + Fold::MulAdd + }; + println!("{}: CPU replica folding as {fold:?}", case.name); + + let mut diff_ref = 0usize; + let mut worst_ref = 0.0f32; + for slot in 0..slots { + let w = &decoded[ids[slot] as usize]; + let xs = &x[slot * case.ic..(slot + 1) * case.ic]; + for o in 0..case.oc { + let w_row = |i: usize| f16_to_f32(w[i * case.oc + o]); + let want = if wide { + cpu_row_wide(&w_row, xs, case.ic, fold) + } else { + cpu_row_narrow(&w_row, xs, case.ic, fold) + }; + let got = y_native[slot * case.oc + o]; + if want.to_bits() != got.to_bits() { + diff_ref += 1; + worst_ref = worst_ref.max((want - got).abs()); + } + } + } + println!( + "{}: fused-native vs escha_ref (kernel order) — {diff_ref} differing floats of {} \ + (max |delta| {worst_ref:e})", + case.name, + y_native.len() + ); + if diff_ref != 0 { + failures += 1; + } + + // A gate that compared two all-zero buffers would report 0 differences + // and prove nothing. Assert the output is actually a GEMV result. + let nonzero = y_native.iter().filter(|v| **v != 0.0).count(); + assert!( + nonzero > y_native.len() / 2, + "{}: only {nonzero} of {} outputs are non-zero — the fixture or the launch is \ + degenerate and the equality above is vacuous", + case.name, + y_native.len() + ); + + for t in [ + d_code, d_f16, bare, stage, code_ptrs, f16_ptrs, d_ids, d_x, d_y_native, d_y_f16, + ] { + let _ = gpu.free_tensor(t); + } + } + + assert_eq!(failures, 0, "G7: the fused native GEMV is not bit-exact"); + println!("G7 PASS"); +} diff --git a/crates/rdna-compute/examples/test_escha_router_f16_boundary.rs b/crates/rdna-compute/examples/test_escha_router_f16_boundary.rs new file mode 100644 index 0000000000..87cae0d860 --- /dev/null +++ b/crates/rdna-compute/examples/test_escha_router_f16_boundary.rs @@ -0,0 +1,212 @@ +// SPDX-License-Identifier: Apache-2.0 +// Copyright (c) 2026 Kevin Read +// hipfire — see LICENSE and NOTICE in the project root. + +//! Task 9 review Fix 1 — proof that `gpu.router_logits_round_f16_rne` +//! actually changes MoE top-k SELECTION on a constructed boundary case, +//! on BOTH selection routes (`moe_router_softmax_topk_k8_wave64_exact`, +//! the production gfx1151 kernel, and the reference two-launch +//! `softmax_f32` + `moe_topk_renorm_k8` pair). +//! +//! Construction: 256 router logits. +//! - experts 0..6: distinct large values (20.0 down to 14.0) — always +//! in the top-8, on either side of rounding. +//! - expert A_IDX (50): value `v_a`. +//! - expert B_IDX (200): value `v_b`, the LARGEST f32 strictly greater +//! than `v_a` that still rounds to `v_a`'s f16 bit pattern (found by +//! walking `v_a`'s bit pattern up one ULP at a time and checking +//! `half::f16::from_f32` — no hand-derived f16 arithmetic). +//! - every other expert: -50.0 (never competitive). +//! +//! Raw f32 order: v_b > v_a, so unrounded top-8 = {0..6, B_IDX} and A_IDX +//! is the (correctly) excluded 9th-place expert. +//! +//! Every kernel on this path (`moe_router_softmax_topk_k8_wave64_exact`'s +//! `ROUTER_EXACT_CONSIDER`/`router_exact_wave32_chunk_argmax`, and +//! `moe_topk_renorm_k8`'s `r4_topk_pair_reduce`) picks a new candidate only +//! on a STRICT `>` — so on an exact tie the lower index always wins. Once +//! `v_a` and `v_b` are rounded to the identical f16-widened f32 value, that +//! tie-break flips the winner from B_IDX (higher index) to A_IDX (lower +//! index): rounded top-8 = {0..6, A_IDX}, and B_IDX is now excluded. +//! +//! This is exactly the class of divergence the background describes: +//! EschaLabs selects top-k from f16(logits) widened back to f32; hipfire's +//! default path selects top-k from full-precision F32 logits. The escha-only +//! `router_logits_round_f16_rne` step makes hipfire's escha decode path +//! reproduce Escha's selection; every other model's `run_moe_decode` call +//! never invokes it and keeps the current (raw F32) selection untouched. +//! +//! Run: `cargo run --release -p rdna-compute --example test_escha_router_f16_boundary` + +use rdna_compute::{DType, Gpu}; + +const N_EXP: usize = 256; +const TOP_K: usize = 8; +const A_IDX: usize = 50; +const B_IDX: usize = 200; + +fn main() { + // ── Construct the boundary pair on the host, byte-exact, no hand f16 math ── + let v_a: f32 = 3.0; + let h_a = half::f16::from_f32(v_a); + let mut bits = v_a.to_bits(); + loop { + let candidate = f32::from_bits(bits + 1); + if half::f16::from_f32(candidate) != h_a { + break; + } + bits += 1; + } + let v_b = f32::from_bits(bits); + assert!(v_b > v_a, "expected v_b > v_a, got v_b={v_b} v_a={v_a}"); + assert_eq!( + half::f16::from_f32(v_a), + half::f16::from_f32(v_b), + "v_a={v_a} and v_b={v_b} must round to the identical f16 value" + ); + println!( + "boundary pair: v_a(idx {A_IDX})={v_a:.9} v_b(idx {B_IDX})={v_b:.9} \ + both round to f16 {:#06x} ({:.9})", + h_a.to_bits(), + h_a.to_f32() + ); + + let mut logits = vec![-50.0f32; N_EXP]; + for i in 0..7 { + logits[i] = 20.0 - i as f32; + } + logits[A_IDX] = v_a; + logits[B_IDX] = v_b; + + // Host oracle: "escha selection" = top-8 by f32(f16(logit)), stable on + // ties in ascending-index order (matches every kernel's strict `>` + // tie-break: lower index wins). Independent of any GPU kernel. + let mut rounded: Vec<(usize, f32)> = logits + .iter() + .enumerate() + .map(|(i, &v)| (i, half::f16::from_f32(v).to_f32())) + .collect(); + rounded.sort_by(|a, b| b.1.partial_cmp(&a.1).unwrap()); + let mut oracle_escha_set: Vec = rounded[..TOP_K].iter().map(|(i, _)| *i).collect(); + oracle_escha_set.sort_unstable(); + + let mut raw: Vec<(usize, f32)> = logits.iter().enumerate().map(|(i, &v)| (i, v)).collect(); + raw.sort_by(|a, b| b.1.partial_cmp(&a.1).unwrap()); + let mut oracle_raw_set: Vec = raw[..TOP_K].iter().map(|(i, _)| *i).collect(); + oracle_raw_set.sort_unstable(); + + assert!( + oracle_raw_set.contains(&B_IDX) && !oracle_raw_set.contains(&A_IDX), + "sanity: raw f32 oracle should include B_IDX={B_IDX} and exclude A_IDX={A_IDX}, got {oracle_raw_set:?}" + ); + assert!( + oracle_escha_set.contains(&A_IDX) && !oracle_escha_set.contains(&B_IDX), + "sanity: f16-rounded oracle should include A_IDX={A_IDX} and exclude B_IDX={B_IDX}, got {oracle_escha_set:?}" + ); + assert_ne!( + oracle_raw_set, oracle_escha_set, + "constructed case failed to produce a differing SET at the host-oracle level" + ); + println!("host oracle: raw-F32 top-8 set = {oracle_raw_set:?}"); + println!("host oracle: f16-rounded top-8 set = {oracle_escha_set:?}"); + + // ── GPU: exact-wave64 kernel (production gfx1151/gfx1100 route) ── + let mut gpu = Gpu::init().expect("GPU init"); + let unrounded_exact = run_wave64_exact(&mut gpu, &logits, false); + let rounded_exact = run_wave64_exact(&mut gpu, &logits, true); + println!("GPU exact-wave64 kernel: unrounded set = {unrounded_exact:?}"); + println!("GPU exact-wave64 kernel: rounded set = {rounded_exact:?}"); + assert_eq!( + unrounded_exact, oracle_raw_set, + "exact-wave64 kernel's unrounded selection must match the raw-F32 oracle \ + (this is what every non-escha model gets today, unchanged)" + ); + assert_eq!( + rounded_exact, oracle_escha_set, + "exact-wave64 kernel's rounded selection must match the f16-widened oracle \ + (this is what the escha decode path now gets)" + ); + assert_ne!( + unrounded_exact, rounded_exact, + "router_logits_round_f16_rne had no effect on the exact-wave64 route" + ); + + // ── GPU: reference two-launch fallback (softmax_f32 + moe_topk_renorm_k8) ── + let unrounded_fallback = run_softmax_fallback(&mut gpu, &logits, false); + let rounded_fallback = run_softmax_fallback(&mut gpu, &logits, true); + println!("GPU softmax+renorm fallback: unrounded set = {unrounded_fallback:?}"); + println!("GPU softmax+renorm fallback: rounded set = {rounded_fallback:?}"); + assert_eq!( + unrounded_fallback, oracle_raw_set, + "fallback pair's unrounded selection must match the raw-F32 oracle" + ); + assert_eq!( + rounded_fallback, oracle_escha_set, + "fallback pair's rounded selection must match the f16-widened oracle" + ); + assert_ne!( + unrounded_fallback, rounded_fallback, + "router_logits_round_f16_rne had no effect on the softmax+renorm fallback route" + ); + + println!( + "PASS: router_logits_round_f16_rne flips the top-8 SET at the constructed \ + boundary (A_IDX={A_IDX} in, B_IDX={B_IDX} out) on both selection routes, \ + matching the f16-widened oracle; the unrounded routes are byte-identical \ + to the pre-existing raw-F32 selection." + ); +} + +/// Runs `moe_router_softmax_topk_k8_wave64_exact`, optionally rounding the +/// logits through `router_logits_round_f16_rne` first (mirrors exactly what +/// `run_moe_decode` now does when `MoeDtypes::has_escha_experts()` is true). +/// Returns the selected top-8 expert indices, sorted ascending. +fn run_wave64_exact(gpu: &mut Gpu, logits: &[f32], round_f16: bool) -> Vec { + let logits_gpu = gpu.upload_f32(logits, &[N_EXP]).expect("upload logits"); + if round_f16 { + gpu.router_logits_round_f16_rne(&logits_gpu) + .expect("round logits to f16"); + } + let idx = gpu.zeros(&[TOP_K], DType::F32).expect("idx tensor"); + let w = gpu.zeros(&[TOP_K], DType::F32).expect("weight tensor"); + gpu.moe_router_softmax_topk_k8_wave64_exact(&logits_gpu, &idx, &w, N_EXP, true) + .expect("exact wave64 router"); + let set = download_index_set(gpu, &idx); + gpu.free_tensor(logits_gpu).expect("free logits"); + gpu.free_tensor(idx).expect("free idx"); + gpu.free_tensor(w).expect("free w"); + set +} + +/// Runs the reference two-launch fallback (`softmax_f32` + `moe_topk_renorm_k8`), +/// optionally rounding first. Same contract as `run_wave64_exact`. +fn run_softmax_fallback(gpu: &mut Gpu, logits: &[f32], round_f16: bool) -> Vec { + let logits_gpu = gpu.upload_f32(logits, &[N_EXP]).expect("upload logits"); + if round_f16 { + gpu.router_logits_round_f16_rne(&logits_gpu) + .expect("round logits to f16"); + } + let idx = gpu.zeros(&[TOP_K], DType::F32).expect("idx tensor"); + let w = gpu.zeros(&[TOP_K], DType::F32).expect("weight tensor"); + gpu.softmax_f32(&logits_gpu).expect("softmax"); + gpu.moe_topk_renorm_k8(&logits_gpu, &idx, &w, N_EXP, true) + .expect("topk renorm"); + let set = download_index_set(gpu, &idx); + gpu.free_tensor(logits_gpu).expect("free logits"); + gpu.free_tensor(idx).expect("free idx"); + gpu.free_tensor(w).expect("free w"); + set +} + +/// Downloads a `[TOP_K]` "i32-in-F32 alias" index buffer (the same +/// bit-reinterpretation convention `moe_ffn_decode_impl::capture_expert_stats` +/// and `escha_router_topk_for_test` use) and returns it as a sorted `Vec`. +fn download_index_set(gpu: &Gpu, idx: &rdna_compute::GpuTensor) -> Vec { + let idx_f32 = gpu.download_f32(idx).expect("download idx"); + let mut set: Vec = idx_f32 + .iter() + .map(|v| (v.to_bits() as i32) as usize) + .collect(); + set.sort_unstable(); + set +} diff --git a/crates/rdna-compute/src/dispatch.rs b/crates/rdna-compute/src/dispatch.rs index f269ce54db..6bc5143525 100644 --- a/crates/rdna-compute/src/dispatch.rs +++ b/crates/rdna-compute/src/dispatch.rs @@ -326,6 +326,16 @@ pub enum DType { /// `[6..8)` fp16 z1, `[8..104)` 96 B 3-bit payload (8/3 B). Same half /// semantics as MQ6G256V2. `K % 256 == 0`, 3.25 bpw. MQ3G256V2, + /// Escha-W2 trellis, K=2, 16x16 tile, cbA hash codebook (hfq qt=42, 2.00 bpw). + /// Weights are stored in the ROTATED domain — a 128-point unnormalised + /// Walsh-Hadamard is applied to activations on BOTH sides of the matmul + /// (see `RotationPlan::EschaH128`). Reaching an unrotated Plain GEMV with + /// this dtype produces fluent-looking but silently wrong output, not a + /// crash — see `dtype_rotation_plan` / `KernelKey::for_gemv`. + Escha2T16, + /// Escha-W2 trellis, K=3, 16x16 tile, cbA hash codebook (hfq qt=43, 3.00 bpw). + /// Same rotated-domain contract as `Escha2T16`. + Escha3T16, /// MQ2-G256 v2 (qt=50): FWHT-rotated, 72 B/group, neutral Magnum V2. /// Per-group 72 B: `[0..2)` fp16 s0, `[2..4)` fp16 z0, `[4..6)` fp16 s1, /// `[6..8)` fp16 z1, `[8..72)` 64 B 2-bit payload (4/B). Same half @@ -432,6 +442,8 @@ impl DType { | DType::MQ4G256Lloyd | DType::MQ2G256GL | DType::MQ3G256GL + | DType::Escha2T16 + | DType::Escha3T16 | DType::HFP4G32 | DType::MFP4G32 | DType::MFP4G32Lloyd @@ -1223,6 +1235,7 @@ impl Gpu { mq_x_scales: None, mq_rmsnorm_wavegrid_scratch: None, gemv_residual_tmp: None, + escha_prefill: None, paro_x_scratch: None, paro_fused_scratch: None, fp16_x_scratch: None, @@ -2915,6 +2928,25 @@ impl Gpu { .ensure_gemv_residual_tmp(&self.hip, self.device_id, min_elems) } + /// Model-global Escha-W2 batched-prefill routed scratch, allocated on + /// first use and grown on demand. See + /// [`crate::scratch::EschaPrefillScratch`]. + /// Returns VIEWS by value, not a borrow: the caller launches kernels + /// through `&mut Gpu` immediately afterwards, so a borrow of + /// `self.scratch` could not survive. + pub fn ensure_escha_prefill_scratch( + &mut self, + slots: usize, + hidden: usize, + mi: usize, + ) -> HipResult { + // bind_thread: skip — delegated to scratch.rs (takes device_id explicitly). + Ok(self + .scratch + .ensure_escha_prefill(&self.hip, self.device_id, slots, hidden, mi)? + .views(slots)) + } + pub fn alloc_tensor(&mut self, shape: &[usize], dtype: DType) -> HipResult { self.bind_thread()?; let numel: usize = shape.iter().product(); @@ -4253,6 +4285,27 @@ impl Gpu { _ => {} } + // Escha-W2 tile decode: standalone utility, not gated on weight_quant + // (it runs ahead of the normal GEMV dispatch to materialize bare fp16 + // weights from the packed trellis code). + specs.push(( + "escha_decode_tiles", + kernels::ESCHA_DECODE_TILES_SRC.to_string(), + )); + specs.push(("escha_h128", kernels::ESCHA_H128_SRC.to_string())); + specs.push(( + "escha_bare_to_outmajor", + kernels::ESCHA_BARE_TO_OUTMAJOR_SRC.to_string(), + )); + specs.push(( + "escha_moe_gemv_k8_indexed", + kernels::ESCHA_MOE_GEMV_K8_INDEXED_SRC.to_string(), + )); + specs.push(( + "escha_moe_gemv_native", + kernels::ESCHA_MOE_GEMV_NATIVE_SRC.to_string(), + )); + // Embedding kernels — Q8_0 is most common, also cover HFQ4G256/G128 variants specs.push(("embedding_q8", kernels::EMBEDDING_Q8_SRC.to_string())); specs.push(( @@ -4582,6 +4635,36 @@ impl Gpu { "sample_topk_partial_fast65", "sample_topk_finalize_fast65", ], + // Escha-W2 (Tasks 8/10): two multi-entry modules whose module + // name is NOT a symbol. Without these arms the `other => + // vec![other]` default asks hipModuleGetFunction for a symbol + // named "escha_h128" / "escha_bare_to_outmajor", which does not + // exist, and the whole precompile batch fails. + "escha_h128" => vec![ + "escha_h128_in", + "escha_h128_out", + "escha_h128_in_batched", + "escha_h128_out_batched", + "escha_swiglu_batched", + ], + "escha_bare_to_outmajor" => vec![ + "escha_bare_to_q8_0", + "escha_bare_to_f32", + "escha_bare_to_f16", + ], + "escha_moe_gemv_k8_indexed" => vec![ + "escha_gemv_q8_0_moe_k8_indexed_batched", + "escha_gemv_q8_0_wide_moe_k8_indexed_batched", + "escha_round_weights_f16_rne", + ], + "escha_moe_gemv_native" => vec![ + "escha_gemv_native_k2_moe_k8_indexed_batched", + "escha_gemv_native_k3_moe_k8_indexed_batched", + "escha_gemv_native_k2_wide_moe_k8_indexed_batched", + "escha_gemv_native_k3_wide_moe_k8_indexed_batched", + "escha_gemv_f16_moe_k8_indexed_batched", + "escha_gemv_f16_wide_moe_k8_indexed_batched", + ], other => vec![other], }; // Compile and ensure the module is loaded once. diff --git a/crates/rdna-compute/src/gemm.rs b/crates/rdna-compute/src/gemm.rs index 1ea8aa32ce..cb6d337af4 100644 --- a/crates/rdna-compute/src/gemm.rs +++ b/crates/rdna-compute/src/gemm.rs @@ -22549,7 +22549,29 @@ impl Gpu { kernels::GEMM_Q8_0_RESIDUAL_WMMA_SRC, "gemm_q8_0_residual_wmma", )?; - let x_f16_ptr = self.ensure_fp16_x(x, batch_size * k)?; + // Unconditional F32->F16 conversion, NOT the pointer-keyed + // `ensure_fp16_x` — the same call the `gemm_q8_0_wmma` sibling already + // makes, for the same reason, now applied to the residual variant too. + // + // `x` here is a batched-prefill activation: a FIXED scratch allocation + // (`dn_normed_batch`, `fa_attn_out_batch`, `ffn_hidden_batch`, ...) + // whose contents are rewritten EVERY layer. Pointer-keyed caching + // therefore hits on layer 1 and hands layer 0's activation to layer + // 1's weights. It went unnoticed because no shipped model reached this + // kernel with a per-layer buffer and no intervening conversion until + // Escha-W2's Q8_0 `wo` did; there the residual stream picked up a + // ~400x-too-large term, stayed finite and fluent, and only moved the + // argmax. One extra elementwise kernel per call against a full + // [M, K] GEMM is not a measurable cost. + // + // There is deliberately no env lever back to the pointer-keyed path: + // the measurement it existed for is recorded in commit d4dff4a26, and + // the only thing the lever could do now is reinstate the defect. + let x_f16_ptr = if matches!(x.dtype, DType::F16) { + x.buf.as_ptr() + } else { + self.convert_fp16_x_uncached(x, batch_size * k)? + }; let mut a_p = a.buf.as_ptr(); let mut xp = x_f16_ptr; @@ -23032,7 +23054,29 @@ impl Gpu { kernels::GEMM_Q8_0_RESIDUAL_WMMA_GFX12_SRC, "gemm_q8_0_residual_wmma_gfx12", )?; - let x_f16_ptr = self.ensure_fp16_x(x, batch_size * k)?; + // Unconditional F32->F16 conversion, NOT the pointer-keyed + // `ensure_fp16_x` — the same call the `gemm_q8_0_wmma` sibling already + // makes, for the same reason, now applied to the residual variant too. + // + // `x` here is a batched-prefill activation: a FIXED scratch allocation + // (`dn_normed_batch`, `fa_attn_out_batch`, `ffn_hidden_batch`, ...) + // whose contents are rewritten EVERY layer. Pointer-keyed caching + // therefore hits on layer 1 and hands layer 0's activation to layer + // 1's weights. It went unnoticed because no shipped model reached this + // kernel with a per-layer buffer and no intervening conversion until + // Escha-W2's Q8_0 `wo` did; there the residual stream picked up a + // ~400x-too-large term, stayed finite and fluent, and only moved the + // argmax. One extra elementwise kernel per call against a full + // [M, K] GEMM is not a measurable cost. + // + // There is deliberately no env lever back to the pointer-keyed path: + // the measurement it existed for is recorded in commit d4dff4a26, and + // the only thing the lever could do now is reinstate the defect. + let x_f16_ptr = if matches!(x.dtype, DType::F16) { + x.buf.as_ptr() + } else { + self.convert_fp16_x_uncached(x, batch_size * k)? + }; let mut a_p = a.buf.as_ptr(); let mut xp = x_f16_ptr; diff --git a/crates/rdna-compute/src/gemv.rs b/crates/rdna-compute/src/gemv.rs index e1f5f7070b..70fd71a588 100644 --- a/crates/rdna-compute/src/gemv.rs +++ b/crates/rdna-compute/src/gemv.rs @@ -168,6 +168,78 @@ pub(crate) fn e8_soa_experts_enabled() -> bool { }) } +/// How the input activation of `escha_h128_in_batched` maps onto its slots. +/// +/// Replaces a bare `x_batched: bool`. The bool covered the two decode cases +/// (one shared `x`, or one `x` per slot); batched prefill needs a third — one +/// `x` per TOKEN, shared by that token's `k` expert slots — and a bool cannot +/// express it. Making it an enum rather than a raw `i32` means a caller cannot +/// pass `2` when it meant "true": the grouped case has to name its group size, +/// which is exactly the value that would otherwise be silently wrong. +/// +/// The kernarg encoding (`<= 0` broadcast, else `slot / g`) keeps 0 and 1 +/// meaning what the bool's `false` and `true` meant, so the kernel change is +/// backward compatible for every existing call site. +#[derive(Clone, Copy, PartialEq, Eq, Debug)] +pub enum EschaXGroup { + /// One `[n]` activation shared by every slot. + Broadcast, + /// `[slots, n]` — one activation per slot. + PerSlot, + /// `[slots / g, n]` — slot `s` reads row `s / g`. `g` must divide `slots`. + Grouped(usize), +} + +impl EschaXGroup { + /// The `x_group` kernarg. `PerSlot` is `Grouped(1)` by construction. + #[inline] + pub fn as_kernarg(self) -> i32 { + match self { + EschaXGroup::Broadcast => 0, + EschaXGroup::PerSlot => 1, + EschaXGroup::Grouped(g) => g as i32, + } + } +} + +/// Register-tile shape of the escha grouped GEMM: `(ROWS, CTILES)` — how many +/// of an expert's token rows one pass holds, and how many adjacent 16-wide tile +/// columns one block owns. +/// +/// A lane holds `ROWS * 2 * CTILES` f32 accumulators, so this is a register +/// budget, and it trades the two traffic terms against each other: bigger ROWS +/// re-reads the expert code fewer times (`ceil(G_e / ROWS)` passes per expert), +/// bigger CTILES re-reads the ACTIVATION fewer times (`m / (16*CTILES)` blocks +/// per slot instead of `m / 16`). +/// +/// `(8, 4)` is the swept default — see +/// `rdna-compute/examples/bench_escha_grouped_gemm.rs`, which measures the +/// whole instantiated set at both shipped projection shapes. `CTILES` falls +/// back when it does not divide `m / 16`; every shipped escha projection has +/// `m ∈ {1024, 2048}`, so the fallback is unreachable today and exists so a +/// future shape gets a smaller tile rather than a rejected launch. +/// +/// `HIPFIRE_ESCHA_GROUPED_TILE=RxC` overrides it. That is a TUNING knob, not a +/// route switch: every instantiation computes the same sums in the same order, +/// so moving it changes speed and nothing else. +pub fn escha_grouped_tile(m: usize) -> (usize, usize) { + static TILE: std::sync::OnceLock<(usize, usize)> = std::sync::OnceLock::new(); + let (rows, ctiles) = *TILE.get_or_init(|| { + std::env::var("HIPFIRE_ESCHA_GROUPED_TILE") + .ok() + .and_then(|s| { + let (r, c) = s.split_once('x')?; + Some((r.trim().parse().ok()?, c.trim().parse().ok()?)) + }) + .unwrap_or((8, 4)) + }); + let mut c = ctiles; + while c > 1 && m % (16 * c) != 0 { + c /= 2; + } + (rows, c) +} + impl Gpu { /// Q4_LUT GEMV: 4-bit with LDS codebook lookup. 48 bytes per 32 elements. pub fn gemv_q4lut( @@ -13637,180 +13709,1548 @@ impl Gpu { result } - /// y = A_q8hfq * x (split-metadata Q8 GEMV, row_stride = padded row bytes) - pub fn gemv_q8hfq( + /// Shared validation for both `escha_decode_tiles` entry points: catches + /// a bad shape/K/code-length combination before it becomes an + /// out-of-bounds device read. The kernel launches + /// `(in_features/16)*(out_features/16)` blocks, each reading `16*K` + /// shorts starting at `tile*16*K` — a short `code` slice, an `in_features` + /// or `out_features` that is not a multiple of 16, or an unsupported `K` + /// all lead to a device-side OOB read (undefined behaviour, not just a + /// wrong answer) rather than a clean failure. Mirrors the assertion the + /// CPU oracle already makes (`escha_ref::reconstruct`). + fn escha_validate_tile_shape( + in_features: u32, + out_features: u32, + k: u32, + code_len: usize, + ) -> HipResult<()> { + if in_features % 16 != 0 { + return Err(hip_bridge::HipError::new( + 0, + &format!("escha_decode_tiles: in_features {in_features} is not a multiple of 16"), + )); + } + if out_features % 16 != 0 { + return Err(hip_bridge::HipError::new( + 0, + &format!("escha_decode_tiles: out_features {out_features} is not a multiple of 16"), + )); + } + if k != 2 && k != 3 { + return Err(hip_bridge::HipError::new( + 0, + &format!("escha_decode_tiles: unsupported K={k} (only 2 and 3 are defined)"), + )); + } + let want_len = (in_features as usize / 16) * (out_features as usize / 16) * 16 * k as usize; + if code_len != want_len { + return Err(hip_bridge::HipError::new( + 0, + &format!( + "escha_decode_tiles: code length mismatch: got {code_len} shorts, expected \ + {want_len} for in_features={in_features} out_features={out_features} K={k}" + ), + )); + } + Ok(()) + } + + /// Decode an escha code stream already resident on the GPU into a bare + /// fp16 weight matrix `[in_features, out_features]`, also GPU-resident. + /// This is the load-path form: no host round trip. `escha_decode_tiles_host` + /// is the host-roundtrip convenience wrapper used by the G2 parity gate + /// and by callers that do not already have the code on-device; it calls + /// this function rather than duplicating the launch. + pub fn escha_decode_tiles( &mut self, - a_raw: &GpuTensor, - x: &GpuTensor, - y: &GpuTensor, - m: usize, - k: usize, - row_stride: usize, + code: &GpuTensor, + bare_out: &GpuTensor, + in_features: u32, + out_features: u32, + k: u32, ) -> HipResult<()> { self.bind_thread()?; - let mut a_ptr = a_raw.buf.as_ptr(); - let mut x_ptr = x.buf.as_ptr(); - let mut y_ptr = y.buf.as_ptr(); - let mut m_val = m as i32; - let mut k_val = k as i32; - let mut rs_val = row_stride as i32; + // Validate against the tensors' LOGICAL shapes, not `buf.size()`: pooled + // allocations (`alloc_tensor`) can hand back a physically larger buffer + // than requested (see `GpuPool::alloc`), so the physical capacity is not + // proof of how much real code/output data is present. + Self::escha_validate_tile_shape(in_features, out_features, k, code.numel())?; + let n_elems = (in_features as usize) * (out_features as usize); + if bare_out.numel() != n_elems { + return Err(hip_bridge::HipError::new( + 0, + &format!( + "escha_decode_tiles: bare_out has {} elements, need exactly {} for \ + {in_features}x{out_features} fp16", + bare_out.numel(), + n_elems + ), + )); + } + let n_tiles = (in_features / 16) * (out_features / 16); + self.ensure_kernel( + "escha_decode_tiles", + kernels::ESCHA_DECODE_TILES_SRC, + "escha_decode_tiles", + )?; + let mut code_ptr = code.buf.as_ptr(); + let mut bare_ptr = bare_out.buf.as_ptr(); + let mut ic = in_features as i32; + let mut oc = out_features as i32; + let mut kk = k as i32; let mut params: Vec<*mut c_void> = vec![ - &mut a_ptr as *mut _ as *mut c_void, - &mut x_ptr as *mut _ as *mut c_void, - &mut y_ptr as *mut _ as *mut c_void, - &mut m_val as *mut _ as *mut c_void, - &mut k_val as *mut _ as *mut c_void, - &mut rs_val as *mut _ as *mut c_void, + &mut code_ptr as *mut _ as *mut c_void, + &mut bare_ptr as *mut _ as *mut c_void, + &mut ic as *mut _ as *mut c_void, + &mut oc as *mut _ as *mut c_void, + &mut kk as *mut _ as *mut c_void, ]; - - if k <= 1536 { - self.ensure_kernel( - "gemv_q8hfq_wide", - kernels::GEMV_Q8HFQ_WIDE_SRC, - "gemv_q8hfq_wide", - )?; - let func = &self.functions["gemv_q8hfq_wide"]; - let block_size = 64u32; - let grid = ((m + 1) / 2) as u32; - return unsafe { - self.hip - .launch_kernel(func, [grid, 1, 1], [block_size, 1, 1], 0, None, &mut params) - }; - } - - self.ensure_kernel("gemv_q8hfq", kernels::GEMV_Q8HFQ_SRC, "gemv_q8hfq")?; - let func = &self.functions["gemv_q8hfq"]; + let func = &self.functions["escha_decode_tiles"]; unsafe { self.hip - .launch_kernel(func, [m as u32, 1, 1], [32, 1, 1], 0, None, &mut params) + .launch_kernel(func, [n_tiles, 1, 1], [32, 1, 1], 0, None, &mut params) } } - /// y = A_q6k * x (quantized GEMV for Q6_K) - pub fn gemv_q6k( + /// Decode an escha code stream to a bare fp16 weight matrix `[ic, oc]`. + /// Host-side helper used by the G2 parity gate; the load path uses the + /// device-resident `escha_decode_tiles` above, which this calls. + pub fn escha_decode_tiles_host( &mut self, - a_raw: &GpuTensor, - x: &GpuTensor, - y: &GpuTensor, - m: usize, - k: usize, - ) -> HipResult<()> { + code: &[i16], + in_features: u32, + out_features: u32, + k: u32, + ) -> HipResult> { self.bind_thread()?; - self.ensure_kernel("gemv_q6k", kernels::GEMV_Q6K_SRC, "gemv_q6k")?; - let func = &self.functions["gemv_q6k"]; + Self::escha_validate_tile_shape(in_features, out_features, k, code.len())?; + let n_elems = (in_features as usize) * (out_features as usize); + let code_bytes: Vec = code.iter().flat_map(|v| v.to_le_bytes()).collect(); + let d_code = self.upload_raw(&code_bytes, &[code.len()])?; + let d_bare = self.alloc_tensor(&[n_elems], DType::F16)?; + + self.escha_decode_tiles(&d_code, &d_bare, in_features, out_features, k)?; + + let mut out = vec![0u8; n_elems * 2]; + self.hip.memcpy_dtoh(&mut out, &d_bare.buf)?; + Ok(out + .chunks_exact(2) + .map(|c| u16::from_le_bytes([c[0], c[1]])) + .collect()) + } - let mut a_ptr = a_raw.buf.as_ptr(); - let mut x_ptr = x.buf.as_ptr(); - let mut y_ptr = y.buf.as_ptr(); - let mut m_val = m as i32; - let mut k_val = k as i32; + /// Launch one of the two H128 activation-transform entry points on + /// already GPU-resident buffers. `entry` selects `escha_h128_in` (scale + /// THEN transform) or `escha_h128_out` (transform THEN scale) — see + /// `kernels/src/escha_h128.hip`. This is the load-bearing form for the + /// forward path: `escha_h128_in_host`/`escha_h128_out_host` below are + /// the host-roundtrip convenience wrappers used by the G3 parity gate + /// and the benchmark; they call this rather than duplicating the launch. + pub fn escha_h128( + &mut self, + entry: &str, + a: &GpuTensor, + vec_in: &GpuTensor, + out: &GpuTensor, + ) -> HipResult<()> { + self.bind_thread()?; + let n = a.numel(); + assert_eq!(n, vec_in.numel(), "escha_h128: a/vec_in length mismatch"); + assert_eq!(n, out.numel(), "escha_h128: a/out length mismatch"); + assert_eq!(n % 128, 0, "H128 needs a multiple of 128"); + self.ensure_kernel("escha_h128", kernels::ESCHA_H128_SRC, entry)?; + let mut a_ptr = a.buf.as_ptr(); + let mut v_ptr = vec_in.buf.as_ptr(); + let mut o_ptr = out.buf.as_ptr(); + let mut n_val = n as i32; let mut params: Vec<*mut c_void> = vec![ &mut a_ptr as *mut _ as *mut c_void, - &mut x_ptr as *mut _ as *mut c_void, - &mut y_ptr as *mut _ as *mut c_void, - &mut m_val as *mut _ as *mut c_void, - &mut k_val as *mut _ as *mut c_void, + &mut v_ptr as *mut _ as *mut c_void, + &mut o_ptr as *mut _ as *mut c_void, + &mut n_val as *mut _ as *mut c_void, ]; - - let block_size = 256u32; - let shared_mem = block_size * 4; + let func = &self.functions[entry]; unsafe { self.hip.launch_kernel( func, - [m as u32, 1, 1], - [block_size, 1, 1], - shared_mem, - self.stream_ref(), + [(n / 128) as u32, 1, 1], + [128, 1, 1], + 0, + None, &mut params, ) } } - /// y = A_q4f16 * x (RDNA-native Q4_F16 GEMV, group size 64) - /// a_raw: raw Q4_F16_G64 bytes on GPU, x: F32 input, y: F32 output - /// Block: 36 bytes per 64 elements. K must be multiple of 64. - /// Uses 128 threads (4 warps) with shared memory reduction for increased MLP. - pub fn gemv_q4f16_g64( - &mut self, - a_raw: &GpuTensor, - x: &GpuTensor, - y: &GpuTensor, - m: usize, - k: usize, - ) -> HipResult<()> { - self.bind_thread()?; - self.ensure_kernel( - "gemv_q4f16_g64", - kernels::GEMV_Q4F16_G64_SRC, - "gemv_q4f16_g64", - )?; - let func = &self.functions["gemv_q4f16_g64"]; + /// `xh = f16( H128(x * rin) * RS )` on device. Host-side helper for the + /// G3 parity gate; the forward path uses the device-resident form above. + pub fn escha_h128_in_host(&mut self, x: &[f32], rin: &[f32]) -> HipResult> { + self.escha_h128_host_impl("escha_h128_in", x, rin) + } - let mut a_ptr = a_raw.buf.as_ptr(); - let mut x_ptr = x.buf.as_ptr(); - let mut y_ptr = y.buf.as_ptr(); - let mut m_val = m as i32; - let mut k_val = k as i32; + /// `y = f16( H128(mid) * RS * rout )` on device. + pub fn escha_h128_out_host(&mut self, mid: &[f32], rout: &[f32]) -> HipResult> { + self.escha_h128_host_impl("escha_h128_out", mid, rout) + } - let mut params: Vec<*mut c_void> = vec![ - &mut a_ptr as *mut _ as *mut c_void, - &mut x_ptr as *mut _ as *mut c_void, - &mut y_ptr as *mut _ as *mut c_void, - &mut m_val as *mut _ as *mut c_void, - &mut k_val as *mut _ as *mut c_void, - ]; + fn escha_h128_host_impl( + &mut self, + entry: &str, + a: &[f32], + vec_in: &[f32], + ) -> HipResult> { + assert_eq!(a.len(), vec_in.len()); + assert_eq!(a.len() % 128, 0, "H128 needs a multiple of 128"); + self.bind_thread()?; + let n = a.len(); + let a_bytes: Vec = a.iter().flat_map(|v| v.to_le_bytes()).collect(); + let v_bytes: Vec = vec_in.iter().flat_map(|v| v.to_le_bytes()).collect(); + let d_a = self.upload_raw(&a_bytes, &[n])?; + let d_v = self.upload_raw(&v_bytes, &[n])?; + let d_out = self.alloc_tensor(&[n], DType::F16)?; + + self.escha_h128(entry, &d_a, &d_v, &d_out)?; + + let mut raw = vec![0u8; n * 2]; + self.hip.memcpy_dtoh(&mut raw, &d_out.buf)?; + Ok(raw + .chunks_exact(2) + .map(|c| u16::from_le_bytes([c[0], c[1]])) + .collect()) + } - let block_size = 32u32; // single warp — no shared memory - unsafe { - self.hip.launch_kernel( - func, - [m as u32, 1, 1], - [block_size, 1, 1], + /// One H128 launch covering ALL `slots` top-k experts of a token + /// (Task 10). `entry` is `escha_h128_in_batched` or + /// `escha_h128_out_batched`. + /// + /// This is a HARD REQUIREMENT of the escha forward path, not an + /// optimisation: Task 8 measured these kernels launch-bound (an empty + /// kernel at the same grid/block is 70-75% of a real launch's cost), so + /// the per-expert form costs 1280 launches/token = 3.07 ms = a 326 tok/s + /// ceiling before any GEMV work. The batched form is 160 launches. + /// + /// - `x_group` (INPUT side only; ignored by `escha_h128_out_batched`) says + /// how many consecutive slots share one row of `a`: + /// * `EschaXGroup::Broadcast` — `a` is `[n]`, every slot reads it. The + /// decode gate_up input side: a token's top-k experts see the same + /// post-rmsnorm activation and differ only in `rin`. + /// * `EschaXGroup::PerSlot` — `a` is `[slots, n]`. The down input side. + /// * `EschaXGroup::Grouped(g)` — `a` is `[slots / g, n]`; slot `s` reads + /// row `s / g`. The batched-prefill gate_up input side with `g = k`: + /// slots are token-major (`token * k + krank`), so all k of a token's + /// experts read that token's activation. `slots % g != 0` is rejected + /// rather than truncated — a ragged tail would silently read a wrong + /// row for the last few slots. + /// - `r_table`: the whole resident `[E, n]` `escha_rin_eff` / + /// `escha_rout_eff` tensor. Slot `s` reads row `ids[s]` — that indexing + /// IS the batching; no per-expert vector is gathered or copied. + /// - `ids`: `[slots]` i32 expert ids, device-resident. + /// - `out`: `[slots, n]` F32 holding f16-ROUNDED values (see the kernel). + pub fn escha_h128_batched( + &mut self, + entry: &str, + a: &GpuTensor, + r_table: &GpuTensor, + ids: &GpuTensor, + out: &GpuTensor, + n: usize, + slots: usize, + x_group: EschaXGroup, + ) -> HipResult<()> { + self.bind_thread()?; + if n % 128 != 0 { + return Err(hip_bridge::HipError::new( 0, - self.stream_ref(), - &mut params, - ) + &format!("escha_h128_batched: n={n} is not a multiple of 128"), + )); + } + let want_a = match x_group { + EschaXGroup::Broadcast => n, + EschaXGroup::PerSlot => slots * n, + EschaXGroup::Grouped(g) => { + if g == 0 || slots % g != 0 { + return Err(hip_bridge::HipError::new( + 0, + &format!( + "escha_h128_batched: x_group={g} does not divide slots={slots}; \ + a ragged group would make the last slots read the wrong row of a" + ), + )); + } + (slots / g) * n + } + }; + // `>=`, not `==`, for the same reason as `out` below: the kernel reads + // exactly `want_a` elements, so an oversized source is safe. Equality + // here broke batched prefill for the DENSE escha path outright. Prefill + // scratch is sized for the maximum chunk, so any shorter chunk was + // rejected for being too BIG — a 2009-token prompt yields a 217-slot + // chunk against 256-slot scratch and failed with "a has 1310720 + // elements, need 1111040". The dense 27B could then only prefill + // token-by-token through the decode path, at ~10 tok/s. + if a.numel() < want_a { + return Err(hip_bridge::HipError::new( + 0, + &format!( + "escha_h128_batched: a has {} elements, need at least {want_a}", + a.numel() + ), + )); + } + // `>=`, not `==`: the kernel writes exactly `slots * n` elements and + // never reads `out`, so an oversized destination is safe. The dense + // path relies on this — its `xh` scratch is sized to the LARGEST `ic` + // any projection uses so one buffer serves them all, and an exact + // check would reject every projection but the widest. + // + // Undersized is still fatal: that is a real out-of-bounds write. + if out.numel() < slots * n { + return Err(hip_bridge::HipError::new( + 0, + &format!( + "escha_h128_batched: out has {} elements, need at least {}", + out.numel(), + slots * n + ), + )); + } + if r_table.numel() % n != 0 || r_table.numel() < n { + return Err(hip_bridge::HipError::new( + 0, + &format!( + "escha_h128_batched: r table has {} elements, not a whole number of {n}-wide rows", + r_table.numel() + ), + )); + } + if ids.numel() < slots { + return Err(hip_bridge::HipError::new( + 0, + &format!( + "escha_h128_batched: ids has {} elements, need {slots}", + ids.numel() + ), + )); } + self.ensure_kernel("escha_h128", kernels::ESCHA_H128_SRC, entry)?; + let mut a_ptr = a.buf.as_ptr(); + let mut r_ptr = r_table.buf.as_ptr(); + let mut i_ptr = ids.buf.as_ptr(); + let mut o_ptr = out.buf.as_ptr(); + let mut n_val = n as i32; + let mut xb = x_group.as_kernarg(); + let mut params: Vec<*mut c_void> = if entry == "escha_h128_in_batched" { + vec![ + &mut a_ptr as *mut _ as *mut c_void, + &mut r_ptr as *mut _ as *mut c_void, + &mut i_ptr as *mut _ as *mut c_void, + &mut o_ptr as *mut _ as *mut c_void, + &mut n_val as *mut _ as *mut c_void, + &mut xb as *mut _ as *mut c_void, + ] + } else { + vec![ + &mut a_ptr as *mut _ as *mut c_void, + &mut r_ptr as *mut _ as *mut c_void, + &mut i_ptr as *mut _ as *mut c_void, + &mut o_ptr as *mut _ as *mut c_void, + &mut n_val as *mut _ as *mut c_void, + ] + }; + let grid = (slots * (n / 128)) as u32; + // Counted so the "160 H128 launches per token" budget is a MEASURED + // number in the G4 gate, not a claim in a comment. Relaxed ordering: + // this is a diagnostic tally, nothing synchronises on it. + crate::ESCHA_H128_LAUNCHES.fetch_add(1, std::sync::atomic::Ordering::Relaxed); + // `launch_maybe_blob`, NOT a raw `launch_kernel(.., None, ..)`. + // + // This runs INSIDE the forward pass, so under graph capture a + // null-stream launch is not ordered against the captured stream that + // the surrounding GEMVs use — those go through `launch_maybe_blob` + // already. The mixture was a genuine data race: the G5 gate's + // reference arm scored 0.000361 and 0.000152 on two runs of the SAME + // binary against the SAME reference, which failed its own negative + // control and made the whole gate unusable. Setting either + // HIP_LAUNCH_BLOCKING=1 or HIPFIRE_GRAPH=0 restored determinism — + // that is what identified capture as the trigger. + let entry_name = entry.to_string(); + let (a_p, r_p, i_p, o_p, n_v, xb_v) = (a_ptr, r_ptr, i_ptr, o_ptr, n_val, xb); + let is_in = entry == "escha_h128_in_batched"; + self.launch_maybe_blob( + &entry_name, + [grid, 1, 1], + [128, 1, 1], + 0, + &mut params, + || { + let mut b = hip_bridge::KernargBlob::new(); + b.push_ptr(a_p); + b.push_ptr(r_p); + b.push_ptr(i_p); + b.push_ptr(o_p); + b.push_i32(n_v); + if is_in { + b.push_i32(xb_v); + } + b + }, + ) } - /// y = A_q4f16 * x (256-thread wide variant for occupancy testing) - /// Element-strided access pattern matching F32 GEMV. Shared memory reduction. - pub fn gemv_q4f16_g64_wide( + /// SwiGLU over the f16-rounded merged `gate_up` output, batched across + /// the token's top-k slots. `y` is `[slots, 2*inter]` (gate = FIRST + /// half), `h` is `[slots, inter]`. One launch for the whole token. + pub fn escha_swiglu_batched( &mut self, - a_raw: &GpuTensor, - x: &GpuTensor, y: &GpuTensor, - m: usize, - k: usize, + h: &GpuTensor, + inter: usize, + slots: usize, ) -> HipResult<()> { self.bind_thread()?; + if y.numel() != slots * 2 * inter || h.numel() != slots * inter { + return Err(hip_bridge::HipError::new( + 0, + &format!( + "escha_swiglu_batched: y={} h={} for slots={slots} inter={inter}", + y.numel(), + h.numel() + ), + )); + } self.ensure_kernel( - "gemv_q4f16_g64_wide", - kernels::GEMV_Q4F16_G64_WIDE_SRC, - "gemv_q4f16_g64_wide", + "escha_h128", + kernels::ESCHA_H128_SRC, + "escha_swiglu_batched", )?; - let func = &self.functions["gemv_q4f16_g64_wide"]; - - let mut a_ptr = a_raw.buf.as_ptr(); - let mut x_ptr = x.buf.as_ptr(); let mut y_ptr = y.buf.as_ptr(); - let mut m_val = m as i32; - let mut k_val = k as i32; - + let mut h_ptr = h.buf.as_ptr(); + let mut inter_i = inter as i32; let mut params: Vec<*mut c_void> = vec![ - &mut a_ptr as *mut _ as *mut c_void, - &mut x_ptr as *mut _ as *mut c_void, &mut y_ptr as *mut _ as *mut c_void, - &mut m_val as *mut _ as *mut c_void, - &mut k_val as *mut _ as *mut c_void, + &mut h_ptr as *mut _ as *mut c_void, + &mut inter_i as *mut _ as *mut c_void, ]; + let bx = 256u32; + let gx = (inter as u32).div_ceil(bx); + // Capture-aware, for the same reason as `escha_h128_batched` above: + // a null-stream launch inside the forward pass is not ordered against + // the captured stream the surrounding GEMVs run on. + let (y_p, h_p, inter_v) = (y_ptr, h_ptr, inter_i); + self.launch_maybe_blob( + "escha_swiglu_batched", + [gx, slots as u32, 1], + [bx, 1, 1], + 0, + &mut params, + || { + let mut b = hip_bridge::KernargBlob::new(); + b.push_ptr(y_p); + b.push_ptr(h_p); + b.push_i32(inter_v); + b + }, + ) + } - let block_size = 256u32; - let shared_mem = block_size * 4; // one float per thread - unsafe { - self.hip.launch_kernel( + /// Escha-W2 routed GEMV for the indexed (GPU-top-K) decode path: one + /// launch computes all `slots` experts' `y[s] = W[ids[s]] · x[s]`. + /// + /// `x_batch` is `[slots, k]`, `y_batch` is `[slots, m]`; every slot has + /// its own input because the escha input transform folds a per-expert + /// `rin_eff` row into it, and every slot keeps its own output because the + /// escha output transform has to run before anything is combined. + /// + /// # Numerics + /// + /// The wide/narrow choice re-uses [`Self::gemv_q8_0`]'s own `k <= 1536` + /// threshold, and each entry point is a verbatim transcription of the + /// corresponding non-indexed kernel. Both facts are load-bearing: the + /// escha routed path previously ran these projections through + /// `GemvFamily::run_auto` -> `gemv_q8_0`, and the G4 block gate's + /// tolerances are calibrated against those exact sums. `gemv_q8_0_wide` + /// folds four interleaved accumulators where `gemv_q8_0` uses one, so + /// choosing the other variant here would silently change the answer. + pub fn escha_gemv_q8_0_moe_k8_indexed_batched( + &mut self, + expert_ptrs: &GpuTensor, + topk_indices: &GpuTensor, + x_batch: &GpuTensor, + y_batch: &GpuTensor, + m: usize, + k: usize, + slots: usize, + ) -> HipResult<()> { + self.bind_thread()?; + if k % 32 != 0 { + return Err(hip_bridge::HipError::new( + 0, + &format!("escha_gemv_q8_0_moe_k8_indexed_batched: k={k} is not a multiple of 32"), + )); + } + if slots == 0 || m == 0 { + return Err(hip_bridge::HipError::new( + 0, + &format!("escha_gemv_q8_0_moe_k8_indexed_batched: slots={slots} m={m}"), + )); + } + if x_batch.numel() < slots * k || y_batch.numel() < slots * m { + return Err(hip_bridge::HipError::new( + 0, + &format!( + "escha_gemv_q8_0_moe_k8_indexed_batched: x has {} elements (need {}), y has \ + {} (need {})", + x_batch.numel(), + slots * k, + y_batch.numel(), + slots * m + ), + )); + } + if topk_indices.numel() < slots { + return Err(hip_bridge::HipError::new( + 0, + &format!( + "escha_gemv_q8_0_moe_k8_indexed_batched: topk_indices has {} elements, need \ + {slots}", + topk_indices.numel() + ), + )); + } + // Same rule as `gemv_q8_0`: wide kernel for small K. See the doc above + // — this is a NUMERICAL selection, not only a performance one. + let wide = k <= 1536; + let entry = if wide { + "escha_gemv_q8_0_wide_moe_k8_indexed_batched" + } else { + "escha_gemv_q8_0_moe_k8_indexed_batched" + }; + self.ensure_kernel( + "escha_moe_gemv_k8_indexed", + kernels::ESCHA_MOE_GEMV_K8_INDEXED_SRC, + entry, + )?; + let pp = expert_ptrs.buf.as_ptr(); + let ip = topk_indices.buf.as_ptr(); + let xp = x_batch.buf.as_ptr(); + let yp = y_batch.buf.as_ptr(); + let m_val = m as i32; + let k_val = k as i32; + let mut params: Vec<*mut c_void> = vec![ + &pp as *const _ as *mut c_void, + &ip as *const _ as *mut c_void, + &xp as *const _ as *mut c_void, + &yp as *const _ as *mut c_void, + &m_val as *const _ as *mut c_void, + &k_val as *const _ as *mut c_void, + ]; + let (grid_x, block_x) = if wide { + (m.div_ceil(2) as u32, 64u32) + } else { + (m as u32, 32u32) + }; + self.launch_maybe_blob( + entry, + [grid_x, slots as u32, 1], + [block_x, 1, 1], + 0, + &mut params, + || { + let mut b = hip_bridge::KernargBlob::new(); + b.push_ptr(pp); + b.push_ptr(ip); + b.push_ptr(xp); + b.push_ptr(yp); + b.push_i32(m_val); + b.push_i32(k_val); + b + }, + ) + } + + /// Escha-W2 routed GEMV that reads the TRELLIS CODE DIRECTLY — the Phase-2 + /// fused kernel. Same signature and same launch shape contract as + /// [`Self::escha_gemv_q8_0_moe_k8_indexed_batched`], so the executor swaps + /// one for the other and nothing else moves. + /// + /// `expert_ptrs` here point at the raw `[in/16, out/16, 16*trellis_k]` + /// int16 code stream of each expert slot — the bytes that came off disk, + /// never an expanded copy. `trellis_k` is 2 (`Escha2T16`, hfq qt=42) or 3 + /// (`Escha3T16`, qt=43); the two have structurally different bit geometry + /// and get separate kernels rather than a runtime branch. + /// + /// # Numerics + /// + /// The wide/narrow choice re-uses the SAME `k <= 1536` threshold as the + /// Q8_0 sibling, and each entry point reproduces that sibling's lane + /// mapping, accumulator count, loop order and final reduction exactly. The + /// result is therefore bit-identical to running the Q8_0 kernel's + /// arithmetic on exactly-decoded fp16 weights — asserted against both + /// [`Self::escha_gemv_f16_moe_k8_indexed_batched`] and `escha_ref` by + /// `rdna-compute/examples/test_escha_native_gemv_gpu_vs_cpu.rs`. + /// + /// # Shape + /// + /// Grid `(m/16, slots)`, block 512: a block owns a whole 16-wide tile + /// column so that every tile it reads is used in full (see the .hip + /// header). Hence the extra `m % 16 == 0` requirement the Q8_0 sibling + /// does not have — satisfied by every escha projection (`2*mi` and + /// `hidden` are both multiples of 16), and rejected loudly rather than + /// silently dropping the tail rows if it ever is not. + #[allow(clippy::too_many_arguments)] + pub fn escha_gemv_native_moe_k8_indexed_batched( + &mut self, + expert_ptrs: &GpuTensor, + topk_indices: &GpuTensor, + x_batch: &GpuTensor, + y_batch: &GpuTensor, + m: usize, + k: usize, + slots: usize, + trellis_k: u32, + nt_major: bool, + ) -> HipResult<()> { + let wide = self.escha_indexed_gemv_preflight( + "escha_gemv_native_moe_k8_indexed_batched", + topk_indices, + x_batch, + y_batch, + m, + k, + slots, + )?; + let entry = match (trellis_k, wide) { + (2, false) => "escha_gemv_native_k2_moe_k8_indexed_batched", + (2, true) => "escha_gemv_native_k2_wide_moe_k8_indexed_batched", + (3, false) => "escha_gemv_native_k3_moe_k8_indexed_batched", + (3, true) => "escha_gemv_native_k3_wide_moe_k8_indexed_batched", + (other, _) => { + return Err(hip_bridge::HipError::new( + 0, + &format!( + "escha_gemv_native_moe_k8_indexed_batched: trellis_k={other}, expected \ + 2 (Escha2T16) or 3 (Escha3T16)" + ), + )) + } + }; + self.escha_launch_native_family( + entry, + expert_ptrs, + topk_indices, + x_batch, + y_batch, + m, + k, + slots, + Some(nt_major), + ) + } + + /// The F16 reference arm of [`Self::escha_gemv_native_moe_k8_indexed_batched`]. + /// + /// Reads an OUT-major `[m, k]` fp16 expert slot — what `escha_bare_to_f16` + /// writes, i.e. the exactly-decoded weights with nothing re-quantised — + /// through the identical grid, lane mapping, accumulator structure and + /// reduction. It exists so "decoding inside the GEMV changes nothing" is a + /// checkable claim about the DECODE rather than about floating point. + /// + /// It is a GATE arm, not a production route: `EschaWeightStore::F16` still + /// runs host-routed, because that arm is what the published G5 KLD + /// reference is built from and changing its GEMV would move a published + /// number. + #[allow(clippy::too_many_arguments)] + pub fn escha_gemv_f16_moe_k8_indexed_batched( + &mut self, + expert_ptrs: &GpuTensor, + topk_indices: &GpuTensor, + x_batch: &GpuTensor, + y_batch: &GpuTensor, + m: usize, + k: usize, + slots: usize, + ) -> HipResult<()> { + let wide = self.escha_indexed_gemv_preflight( + "escha_gemv_f16_moe_k8_indexed_batched", + topk_indices, + x_batch, + y_batch, + m, + k, + slots, + )?; + let entry = if wide { + "escha_gemv_f16_wide_moe_k8_indexed_batched" + } else { + "escha_gemv_f16_moe_k8_indexed_batched" + }; + self.escha_launch_native_family( + entry, + expert_ptrs, + topk_indices, + x_batch, + y_batch, + m, + k, + slots, + None, + ) + } + + /// Shared argument validation for the fused-native GEMV family. Returns + /// `wide` — the SAME `k <= 1536` variant choice the Q8_0 sibling makes, + /// restated here rather than duplicated as a literal so the two cannot + /// drift apart silently (they must agree: it is a numerical selection). + #[allow(clippy::too_many_arguments)] + fn escha_indexed_gemv_preflight( + &mut self, + what: &str, + topk_indices: &GpuTensor, + x_batch: &GpuTensor, + y_batch: &GpuTensor, + m: usize, + k: usize, + slots: usize, + ) -> HipResult { + self.bind_thread()?; + if k % 32 != 0 { + return Err(hip_bridge::HipError::new( + 0, + &format!("{what}: k={k} is not a multiple of 32"), + )); + } + // A block owns a 16-wide tile column, so a non-multiple-of-16 `m` + // would leave the tail rows uncomputed — reading whatever `y` held. + if m % 16 != 0 { + return Err(hip_bridge::HipError::new( + 0, + &format!("{what}: m={m} is not a multiple of 16"), + )); + } + if slots == 0 || m == 0 { + return Err(hip_bridge::HipError::new( + 0, + &format!("{what}: slots={slots} m={m}"), + )); + } + if x_batch.numel() < slots * k || y_batch.numel() < slots * m { + return Err(hip_bridge::HipError::new( + 0, + &format!( + "{what}: x has {} elements (need {}), y has {} (need {})", + x_batch.numel(), + slots * k, + y_batch.numel(), + slots * m + ), + )); + } + if topk_indices.numel() < slots { + return Err(hip_bridge::HipError::new( + 0, + &format!( + "{what}: topk_indices has {} elements, need {slots}", + topk_indices.numel() + ), + )); + } + // Wide (four independent accumulator chains) vs narrow (one chain with + // `#pragma unroll 4`). The quad structure is what hides the weight-load + // latency, and it needs four blocks to fill — hence K >= 128, since a + // block is 32 contraction elements. + // + // This used to be `k <= 1536`, inherited from the Q8_0 twin with no + // stated reason, which sent the one shipped escha projection above that + // line — gate_up at K=2048 — down the narrow path. Measured on + // escha-35b, rocprof kernel trace, 40 decode tokens: + // + // gate_up narrow 83.20 us/call + // gate_up wide 75.98 us/call -8.7% + // + // down_proj (K=512) was already wide and is unchanged at ~46 us/call. + // The kernel source records the same effect as 33 GB/s narrow vs + // 105 GB/s wide, so this is the structure working as documented, not a + // shape-specific fluke. + Ok(k >= 128) + } + + /// Launch one of the six `escha_moe_gemv_native` entry points. They all + /// share a kernarg list and a grid, so the launch is written once. + #[allow(clippy::too_many_arguments)] + fn escha_launch_native_family( + &mut self, + entry: &str, + expert_ptrs: &GpuTensor, + topk_indices: &GpuTensor, + x_batch: &GpuTensor, + y_batch: &GpuTensor, + m: usize, + k: usize, + slots: usize, + nt_major: Option, + ) -> HipResult<()> { + self.ensure_kernel( + "escha_moe_gemv_native", + kernels::ESCHA_MOE_GEMV_NATIVE_SRC, + entry, + )?; + let pp = expert_ptrs.buf.as_ptr(); + let ip = topk_indices.buf.as_ptr(); + let xp = x_batch.buf.as_ptr(); + let yp = y_batch.buf.as_ptr(); + let m_val = m as i32; + let k_val = k as i32; + let mut params: Vec<*mut c_void> = vec![ + &pp as *const _ as *mut c_void, + &ip as *const _ as *mut c_void, + &xp as *const _ as *mut c_void, + &yp as *const _ as *mut c_void, + &m_val as *const _ as *mut c_void, + &k_val as *const _ as *mut c_void, + ]; + // The escha NATIVE kernels take a trailing `int nt_major` selecting the + // tile-grid order; the F16 reference arm does not, so it passes `None` + // and its kernarg list is unchanged. + let nt_major_val = nt_major.map(|v| i32::from(v)); + if let Some(v) = nt_major_val.as_ref() { + params.push(v as *const i32 as *mut c_void); + } + // Grid (m/16, slots), block 256: a block owns a 16-wide tile column, + // and each of its EIGHT warps owns two of that column's output rows + // (`w` and `w + 8`) because those two share a decode window. See the + // .hip header. + self.launch_maybe_blob( + entry, + [(m / 16) as u32, slots as u32, 1], + [256, 1, 1], + 0, + &mut params, + || { + let mut b = hip_bridge::KernargBlob::new(); + b.push_ptr(pp); + b.push_ptr(ip); + b.push_ptr(xp); + b.push_ptr(yp); + b.push_i32(m_val); + b.push_i32(k_val); + // MUST mirror `params` above — this closure is the SECOND, + // independent kernarg description used on the capture/blob + // path. Updating only `params` silently sends the old arg list + // whenever a graph is being recorded. + if let Some(v) = nt_major_val { + b.push_i32(v); + } + b + }, + ) + } + + /// Escha-W2 routed-expert GROUPED GEMM — decode each expert's trellis code + /// once per (layer, batch) and spend it across every token that routed to + /// it, instead of once per (token, expert) slot. + /// + /// The batched-prefill replacement for + /// [`Self::escha_gemv_native_moe_k8_indexed_batched`]. `expert_offsets` + /// (`[n_exp + 1]`, exclusive scan) and `sorted_slot_index` (`[slots]`) are + /// what `moe_scatter_fused_k8` writes when it is run with `block_m = 1`; + /// `x_batch` / `y_batch` stay in the caller's ORIGINAL token-major slot + /// order — the sort is an index permutation the kernel follows, nothing is + /// physically gathered, so every other phase of the escha layer is + /// untouched. + /// + /// # Numerics + /// + /// Per (token, output row) this reproduces the slot-parallel NARROW form + /// exactly — same lane -> contraction map, same `bi` order, one sequential + /// accumulator, same `__shfl_down` ladder — so for a projection the + /// slot-parallel path also runs narrow (`k > 1536`) the two agree bit for + /// bit. For a projection it runs WIDE (`k <= 1536`) they do not: that form + /// folds four interleaved accumulators, which would cost 4x the + /// accumulator registers here and force the tile back to one column. See + /// the .hip header, and the grouped arm of the G4 block gate for the + /// measured cost. + /// + /// # Shape + /// + /// Grid `(m / (16*ctiles), n_exp)`, block 256. `m % (16*ctiles) == 0` and + /// `k % 32 == 0` are enforced rather than truncated: a block owns whole + /// tile columns, so a short tail would leave rows holding whatever `y` + /// happened to contain. + #[allow(clippy::too_many_arguments)] + /// Expert-grouped routed GEMM on the RDNA3 matrix cores. + /// + /// Same grouping contract as [`Self::escha_gemm_native_moe_grouped`] — + /// `expert_offsets` is the padded exclusive scan, `sorted_slot_index` maps + /// sorted position to flat slot, `-1` is the padding sentinel — but the + /// inner product runs on WMMA instead of scalar FMAs. + /// + /// Prefill is compute bound once grouping has fixed the weight traffic + /// (weights are ~4% of prefill time), so this is where the remaining time + /// is: the scalar path measured 1.75 TFLOP/s against a WMMA comparator's + /// 4.59. + /// + /// NOT bit-identical to the scalar grouped path: WMMA accumulates over a + /// different partition of the contraction. The decoded weight VALUES are + /// identical; only summation order moves. + #[allow(clippy::too_many_arguments)] + pub fn escha_gemm_native_moe_grouped_wmma( + &mut self, + expert_ptrs: &GpuTensor, + expert_offsets: &GpuTensor, + sorted_slot_index: &GpuTensor, + x_batch: &GpuTensor, + y_batch: &GpuTensor, + m: usize, + k: usize, + slots: usize, + n_exp: usize, + trellis_k: u32, + nt_major: bool, + ) -> HipResult<()> { + self.bind_thread()?; + let what = "escha_gemm_native_moe_grouped_wmma"; + if k % 16 != 0 || m % 16 != 0 { + return Err(hip_bridge::HipError::new( + 0, + &format!("{what}: m={m} k={k} must both be multiples of 16"), + )); + } + if slots == 0 || m == 0 || n_exp == 0 { + return Err(hip_bridge::HipError::new( + 0, + &format!("{what}: slots={slots} m={m} n_exp={n_exp}"), + )); + } + if x_batch.numel() < slots * k || y_batch.numel() < slots * m { + return Err(hip_bridge::HipError::new( + 0, + &format!( + "{what}: x has {} elements (need {}), y has {} (need {})", + x_batch.numel(), + slots * k, + y_batch.numel(), + slots * m + ), + )); + } + let entry = match trellis_k { + 2 => "escha_gemm_grouped_wmma_k2", + 3 => "escha_gemm_grouped_wmma_k3", + other => { + return Err(hip_bridge::HipError::new( + 0, + &format!("{what}: unsupported trellis K={other}"), + )) + } + }; + self.ensure_kernel(entry, kernels::ESCHA_MOE_GEMM_GROUPED_WMMA_SRC, entry)?; + + let mut ep = expert_ptrs.buf.as_ptr(); + let mut off = expert_offsets.buf.as_ptr(); + let mut idx = sorted_slot_index.buf.as_ptr(); + let mut xp = x_batch.buf.as_ptr(); + let mut yp = y_batch.buf.as_ptr(); + let mut m_val = m as i32; + let mut k_val = k as i32; + let mut params: Vec<*mut c_void> = vec![ + &mut ep as *mut _ as *mut c_void, + &mut off as *mut _ as *mut c_void, + &mut idx as *mut _ as *mut c_void, + &mut xp as *mut _ as *mut c_void, + &mut yp as *mut _ as *mut c_void, + &mut m_val as *mut _ as *mut c_void, + &mut k_val as *mut _ as *mut c_void, + ]; + let mut ntm = i32::from(nt_major); + params.push(&mut ntm as *mut _ as *mut c_void); + let func = &self.functions[entry]; + unsafe { + self.hip.launch_kernel( + func, + [(m / 16) as u32, n_exp as u32, 1], + [32, 1, 1], + 0, + None, + &mut params, + ) + } + } + + pub fn escha_gemm_native_moe_grouped( + &mut self, + expert_ptrs: &GpuTensor, + expert_offsets: &GpuTensor, + sorted_slot_index: &GpuTensor, + x_batch: &GpuTensor, + y_batch: &GpuTensor, + m: usize, + k: usize, + slots: usize, + n_exp: usize, + trellis_k: u32, + nt_major: bool, + ) -> HipResult<()> { + let (rows, ctiles) = escha_grouped_tile(m); + self.escha_gemm_native_moe_grouped_tiled( + expert_ptrs, + expert_offsets, + sorted_slot_index, + x_batch, + y_batch, + m, + k, + slots, + n_exp, + trellis_k, + nt_major, + rows, + ctiles, + ) + } + + /// [`Self::escha_gemm_native_moe_grouped`] with the register tile named + /// explicitly instead of taken from [`escha_grouped_tile`]. + /// + /// Exists for the sweep in `bench_escha_grouped_gemm`, which has to drive + /// every instantiation inside ONE process: `escha_grouped_tile` memoises + /// its env read in a `OnceLock`, so a sweep that went through it would + /// silently measure the first shape six times. Production goes through the + /// non-`_tiled` entry point. + #[allow(clippy::too_many_arguments)] + pub fn escha_gemm_native_moe_grouped_tiled( + &mut self, + expert_ptrs: &GpuTensor, + expert_offsets: &GpuTensor, + sorted_slot_index: &GpuTensor, + x_batch: &GpuTensor, + y_batch: &GpuTensor, + m: usize, + k: usize, + slots: usize, + n_exp: usize, + trellis_k: u32, + nt_major: bool, + rows: usize, + ctiles: usize, + ) -> HipResult<()> { + self.bind_thread()?; + let what = "escha_gemm_native_moe_grouped"; + if k % 32 != 0 { + return Err(hip_bridge::HipError::new( + 0, + &format!("{what}: k={k} is not a multiple of 32"), + )); + } + if slots == 0 || m == 0 || n_exp == 0 { + return Err(hip_bridge::HipError::new( + 0, + &format!("{what}: slots={slots} m={m} n_exp={n_exp}"), + )); + } + if x_batch.numel() < slots * k || y_batch.numel() < slots * m { + return Err(hip_bridge::HipError::new( + 0, + &format!( + "{what}: x has {} elements (need {}), y has {} (need {})", + x_batch.numel(), + slots * k, + y_batch.numel(), + slots * m + ), + )); + } + // The kernel reads `expert_offsets[e]` and `[e+1]` for every `e` on + // grid.y, and dereferences `sorted_slot_index` across that range. A + // short table is an out-of-bounds READ — undefined behaviour, not a + // wrong answer — so check both exactly. + if expert_offsets.numel() < n_exp + 1 { + return Err(hip_bridge::HipError::new( + 0, + &format!( + "{what}: expert_offsets has {} elements, need {}", + expert_offsets.numel(), + n_exp + 1 + ), + )); + } + if sorted_slot_index.numel() < slots { + return Err(hip_bridge::HipError::new( + 0, + &format!( + "{what}: sorted_slot_index has {} elements, need {slots}", + sorted_slot_index.numel() + ), + )); + } + if ctiles == 0 || m % (16 * ctiles) != 0 { + return Err(hip_bridge::HipError::new( + 0, + &format!("{what}: m={m} is not a multiple of {}", 16 * ctiles), + )); + } + let entry = match (trellis_k, rows, ctiles) { + (2, 4, 2) => "escha_gemm_grouped_k2_r4_c2", + (2, 8, 2) => "escha_gemm_grouped_k2_r8_c2", + (2, 8, 4) => "escha_gemm_grouped_k2_r8_c4", + (2, 8, 8) => "escha_gemm_grouped_k2_r8_c8", + (2, 16, 2) => "escha_gemm_grouped_k2_r16_c2", + (2, 16, 4) => "escha_gemm_grouped_k2_r16_c4", + (3, 4, 2) => "escha_gemm_grouped_k3_r4_c2", + (3, 8, 2) => "escha_gemm_grouped_k3_r8_c2", + (3, 8, 4) => "escha_gemm_grouped_k3_r8_c4", + (3, 8, 8) => "escha_gemm_grouped_k3_r8_c8", + (3, 16, 2) => "escha_gemm_grouped_k3_r16_c2", + (3, 16, 4) => "escha_gemm_grouped_k3_r16_c4", + (tk, r, c) => { + return Err(hip_bridge::HipError::new( + 0, + &format!("{what}: no entry point for trellis_k={tk} rows={r} ctiles={c}"), + )) + } + }; + self.ensure_kernel( + "escha_moe_gemm_grouped", + kernels::ESCHA_MOE_GEMM_GROUPED_SRC, + entry, + )?; + let pp = expert_ptrs.buf.as_ptr(); + let op = expert_offsets.buf.as_ptr(); + let sp = sorted_slot_index.buf.as_ptr(); + let xp = x_batch.buf.as_ptr(); + let yp = y_batch.buf.as_ptr(); + let m_val = m as i32; + let k_val = k as i32; + let mut params: Vec<*mut c_void> = vec![ + &pp as *const _ as *mut c_void, + &op as *const _ as *mut c_void, + &sp as *const _ as *mut c_void, + &xp as *const _ as *mut c_void, + &yp as *const _ as *mut c_void, + &m_val as *const _ as *mut c_void, + &k_val as *const _ as *mut c_void, + ]; + self.launch_maybe_blob( + entry, + [(m / (16 * ctiles)) as u32, n_exp as u32, 1], + [256, 1, 1], + 0, + &mut params, + || { + let mut b = hip_bridge::KernargBlob::new(); + b.push_ptr(pp); + b.push_ptr(op); + b.push_ptr(sp); + b.push_ptr(xp); + b.push_ptr(yp); + b.push_i32(m_val); + b.push_i32(k_val); + b + }, + ) + } + + /// Device-side out-of-place `f32 -> f16 -> f32` round-trip of the escha + /// combine weights. + /// + /// The escha combine scales each expert by `f16(score)`. The CPU-top-K + /// route does that on the host copy it already downloaded; the indexed + /// route never downloads, so it happens here. Out-of-place because + /// `src` is the shared `topk_weights` buffer other consumers still read + /// unrounded. + pub fn escha_round_weights_f16_rne( + &mut self, + src: &GpuTensor, + dst: &GpuTensor, + n: usize, + ) -> HipResult<()> { + self.bind_thread()?; + if src.numel() < n || dst.numel() < n { + return Err(hip_bridge::HipError::new( + 0, + &format!( + "escha_round_weights_f16_rne: src has {} / dst has {} elements, need {n}", + src.numel(), + dst.numel() + ), + )); + } + self.ensure_kernel( + "escha_moe_gemv_k8_indexed", + kernels::ESCHA_MOE_GEMV_K8_INDEXED_SRC, + "escha_round_weights_f16_rne", + )?; + let sp = src.buf.as_ptr(); + let dp = dst.buf.as_ptr(); + let n_val = n as i32; + let mut params: Vec<*mut c_void> = vec![ + &sp as *const _ as *mut c_void, + &dp as *const _ as *mut c_void, + &n_val as *const _ as *mut c_void, + ]; + let bx = 256u32; + self.launch_maybe_blob( + "escha_round_weights_f16_rne", + [(n as u32).div_ceil(bx), 1, 1], + [bx, 1, 1], + 0, + &mut params, + || { + let mut b = hip_bridge::KernargBlob::new(); + b.push_ptr(sp); + b.push_ptr(dp); + b.push_i32(n_val); + b + }, + ) + } + + /// Load path: transpose the bare in-major `[ic, oc]` fp16 that + /// `escha_decode_tiles` produced into hipfire's OUT-major expert slot, + /// re-quantising to Q8_0 in the same pass. `out` must be + /// `oc * (ic/32) * 34` bytes. + pub fn escha_bare_to_q8_0( + &mut self, + bare: &GpuTensor, + out: &GpuTensor, + ic: usize, + oc: usize, + ) -> HipResult<()> { + self.bind_thread()?; + if ic % 32 != 0 { + return Err(hip_bridge::HipError::new( + 0, + &format!("escha_bare_to_q8_0: ic={ic} is not a multiple of 32"), + )); + } + if bare.numel() != ic * oc { + return Err(hip_bridge::HipError::new( + 0, + &format!( + "escha_bare_to_q8_0: bare has {} elements, need {}", + bare.numel(), + ic * oc + ), + )); + } + let want = oc * (ic / 32) * 34; + if out.byte_size() != want { + return Err(hip_bridge::HipError::new( + 0, + &format!( + "escha_bare_to_q8_0: out is {} bytes, need {want}", + out.byte_size() + ), + )); + } + self.ensure_kernel( + "escha_bare_to_outmajor", + kernels::ESCHA_BARE_TO_OUTMAJOR_SRC, + "escha_bare_to_q8_0", + )?; + let mut b_ptr = bare.buf.as_ptr(); + let mut o_ptr = out.buf.as_ptr(); + let mut ic_i = ic as i32; + let mut oc_i = oc as i32; + let mut params: Vec<*mut c_void> = vec![ + &mut b_ptr as *mut _ as *mut c_void, + &mut o_ptr as *mut _ as *mut c_void, + &mut ic_i as *mut _ as *mut c_void, + &mut oc_i as *mut _ as *mut c_void, + ]; + let func = &self.functions["escha_bare_to_q8_0"]; + let grid = (oc * (ic / 32)) as u32; + unsafe { + self.hip + .launch_kernel(func, [grid, 1, 1], [32, 1, 1], 0, None, &mut params) + } + } + + /// Weight-exact control arm of [`Self::escha_bare_to_q8_0`]: same + /// transpose, F32 store, no re-quantisation. 4 B/weight — diagnostic + /// only (the G4 gate uses it to separate wiring error from Q8_0 error). + pub fn escha_bare_to_f32( + &mut self, + bare: &GpuTensor, + out: &GpuTensor, + ic: usize, + oc: usize, + ) -> HipResult<()> { + self.bind_thread()?; + if bare.numel() != ic * oc || out.numel() != ic * oc { + return Err(hip_bridge::HipError::new( + 0, + &format!( + "escha_bare_to_f32: bare={} out={} need {} each", + bare.numel(), + out.numel(), + ic * oc + ), + )); + } + self.ensure_kernel( + "escha_bare_to_outmajor", + kernels::ESCHA_BARE_TO_OUTMAJOR_SRC, + "escha_bare_to_f32", + )?; + let mut b_ptr = bare.buf.as_ptr(); + let mut o_ptr = out.buf.as_ptr(); + let mut ic_i = ic as i32; + let mut oc_i = oc as i32; + let mut params: Vec<*mut c_void> = vec![ + &mut b_ptr as *mut _ as *mut c_void, + &mut o_ptr as *mut _ as *mut c_void, + &mut ic_i as *mut _ as *mut c_void, + &mut oc_i as *mut _ as *mut c_void, + ]; + let func = &self.functions["escha_bare_to_f32"]; + let bx = 256u32; + unsafe { + self.hip.launch_kernel( + func, + [(ic as u32).div_ceil(bx), oc as u32, 1], + [bx, 1, 1], + 0, + None, + &mut params, + ) + } + } + + /// Weight-exact arm of [`Self::escha_bare_to_q8_0`] that fits a whole + /// model: same transpose, F16 store, no re-quantisation. + /// + /// The decode already produced fp16, so this is pure data movement and + /// the stored weight is bit-identical to `escha_ref::reconstruct`. + /// [`Self::escha_bare_to_f32`] is equally exact but 4 B/weight, which is + /// 129 GB of experts on the 35B; this is 2 B/weight, and because every + /// per-expert buffer is separately allocated and rounded to a 2 MiB + /// granule it occupies the SAME 60 GiB the Q8_0 arm already occupies. + /// That is what makes a model-scale weight-exact KLD reference runnable + /// (G5, docs/plans/escha-w2-port-design.md). + pub fn escha_bare_to_f16( + &mut self, + bare: &GpuTensor, + out: &GpuTensor, + ic: usize, + oc: usize, + ) -> HipResult<()> { + self.bind_thread()?; + if bare.numel() != ic * oc || out.numel() != ic * oc { + return Err(hip_bridge::HipError::new( + 0, + &format!( + "escha_bare_to_f16: bare={} out={} need {} each", + bare.numel(), + out.numel(), + ic * oc + ), + )); + } + self.ensure_kernel( + "escha_bare_to_outmajor", + kernels::ESCHA_BARE_TO_OUTMAJOR_SRC, + "escha_bare_to_f16", + )?; + let mut b_ptr = bare.buf.as_ptr(); + let mut o_ptr = out.buf.as_ptr(); + let mut ic_i = ic as i32; + let mut oc_i = oc as i32; + let mut params: Vec<*mut c_void> = vec![ + &mut b_ptr as *mut _ as *mut c_void, + &mut o_ptr as *mut _ as *mut c_void, + &mut ic_i as *mut _ as *mut c_void, + &mut oc_i as *mut _ as *mut c_void, + ]; + let func = &self.functions["escha_bare_to_f16"]; + let bx = 256u32; + unsafe { + self.hip.launch_kernel( + func, + [(ic as u32).div_ceil(bx), oc as u32, 1], + [bx, 1, 1], + 0, + None, + &mut params, + ) + } + } + + /// y = A_q8hfq * x (split-metadata Q8 GEMV, row_stride = padded row bytes) + pub fn gemv_q8hfq( + &mut self, + a_raw: &GpuTensor, + x: &GpuTensor, + y: &GpuTensor, + m: usize, + k: usize, + row_stride: usize, + ) -> HipResult<()> { + self.bind_thread()?; + let mut a_ptr = a_raw.buf.as_ptr(); + let mut x_ptr = x.buf.as_ptr(); + let mut y_ptr = y.buf.as_ptr(); + let mut m_val = m as i32; + let mut k_val = k as i32; + let mut rs_val = row_stride as i32; + + let mut params: Vec<*mut c_void> = vec![ + &mut a_ptr as *mut _ as *mut c_void, + &mut x_ptr as *mut _ as *mut c_void, + &mut y_ptr as *mut _ as *mut c_void, + &mut m_val as *mut _ as *mut c_void, + &mut k_val as *mut _ as *mut c_void, + &mut rs_val as *mut _ as *mut c_void, + ]; + + if k <= 1536 { + self.ensure_kernel( + "gemv_q8hfq_wide", + kernels::GEMV_Q8HFQ_WIDE_SRC, + "gemv_q8hfq_wide", + )?; + let func = &self.functions["gemv_q8hfq_wide"]; + let block_size = 64u32; + let grid = ((m + 1) / 2) as u32; + return unsafe { + self.hip + .launch_kernel(func, [grid, 1, 1], [block_size, 1, 1], 0, None, &mut params) + }; + } + + self.ensure_kernel("gemv_q8hfq", kernels::GEMV_Q8HFQ_SRC, "gemv_q8hfq")?; + let func = &self.functions["gemv_q8hfq"]; + unsafe { + self.hip + .launch_kernel(func, [m as u32, 1, 1], [32, 1, 1], 0, None, &mut params) + } + } + + /// y = A_q6k * x (quantized GEMV for Q6_K) + pub fn gemv_q6k( + &mut self, + a_raw: &GpuTensor, + x: &GpuTensor, + y: &GpuTensor, + m: usize, + k: usize, + ) -> HipResult<()> { + self.bind_thread()?; + self.ensure_kernel("gemv_q6k", kernels::GEMV_Q6K_SRC, "gemv_q6k")?; + let func = &self.functions["gemv_q6k"]; + + let mut a_ptr = a_raw.buf.as_ptr(); + let mut x_ptr = x.buf.as_ptr(); + let mut y_ptr = y.buf.as_ptr(); + let mut m_val = m as i32; + let mut k_val = k as i32; + + let mut params: Vec<*mut c_void> = vec![ + &mut a_ptr as *mut _ as *mut c_void, + &mut x_ptr as *mut _ as *mut c_void, + &mut y_ptr as *mut _ as *mut c_void, + &mut m_val as *mut _ as *mut c_void, + &mut k_val as *mut _ as *mut c_void, + ]; + + let block_size = 256u32; + let shared_mem = block_size * 4; + unsafe { + self.hip.launch_kernel( + func, + [m as u32, 1, 1], + [block_size, 1, 1], + shared_mem, + self.stream_ref(), + &mut params, + ) + } + } + + /// y = A_q4f16 * x (RDNA-native Q4_F16 GEMV, group size 64) + /// a_raw: raw Q4_F16_G64 bytes on GPU, x: F32 input, y: F32 output + /// Block: 36 bytes per 64 elements. K must be multiple of 64. + /// Uses 128 threads (4 warps) with shared memory reduction for increased MLP. + pub fn gemv_q4f16_g64( + &mut self, + a_raw: &GpuTensor, + x: &GpuTensor, + y: &GpuTensor, + m: usize, + k: usize, + ) -> HipResult<()> { + self.bind_thread()?; + self.ensure_kernel( + "gemv_q4f16_g64", + kernels::GEMV_Q4F16_G64_SRC, + "gemv_q4f16_g64", + )?; + let func = &self.functions["gemv_q4f16_g64"]; + + let mut a_ptr = a_raw.buf.as_ptr(); + let mut x_ptr = x.buf.as_ptr(); + let mut y_ptr = y.buf.as_ptr(); + let mut m_val = m as i32; + let mut k_val = k as i32; + + let mut params: Vec<*mut c_void> = vec![ + &mut a_ptr as *mut _ as *mut c_void, + &mut x_ptr as *mut _ as *mut c_void, + &mut y_ptr as *mut _ as *mut c_void, + &mut m_val as *mut _ as *mut c_void, + &mut k_val as *mut _ as *mut c_void, + ]; + + let block_size = 32u32; // single warp — no shared memory + unsafe { + self.hip.launch_kernel( + func, + [m as u32, 1, 1], + [block_size, 1, 1], + 0, + self.stream_ref(), + &mut params, + ) + } + } + + /// y = A_q4f16 * x (256-thread wide variant for occupancy testing) + /// Element-strided access pattern matching F32 GEMV. Shared memory reduction. + pub fn gemv_q4f16_g64_wide( + &mut self, + a_raw: &GpuTensor, + x: &GpuTensor, + y: &GpuTensor, + m: usize, + k: usize, + ) -> HipResult<()> { + self.bind_thread()?; + self.ensure_kernel( + "gemv_q4f16_g64_wide", + kernels::GEMV_Q4F16_G64_WIDE_SRC, + "gemv_q4f16_g64_wide", + )?; + let func = &self.functions["gemv_q4f16_g64_wide"]; + + let mut a_ptr = a_raw.buf.as_ptr(); + let mut x_ptr = x.buf.as_ptr(); + let mut y_ptr = y.buf.as_ptr(); + let mut m_val = m as i32; + let mut k_val = k as i32; + + let mut params: Vec<*mut c_void> = vec![ + &mut a_ptr as *mut _ as *mut c_void, + &mut x_ptr as *mut _ as *mut c_void, + &mut y_ptr as *mut _ as *mut c_void, + &mut m_val as *mut _ as *mut c_void, + &mut k_val as *mut _ as *mut c_void, + ]; + + let block_size = 256u32; + let shared_mem = block_size * 4; // one float per thread + unsafe { + self.hip.launch_kernel( func, [m as u32, 1, 1], [block_size, 1, 1], @@ -15989,3 +17429,26 @@ impl Gpu { result } } + +#[cfg(test)] +mod escha_x_group_tests { + use super::EschaXGroup; + + /// The kernarg encoding must keep the two meanings the old `x_batched: + /// bool` had, or the two pre-existing decode call sites (and G3's two + /// pre-existing cases) change behaviour under a change that is supposed to + /// be purely additive. + #[test] + fn x_group_kernarg_is_backward_compatible() { + assert_eq!(EschaXGroup::Broadcast.as_kernarg(), 0, "was `false`"); + assert_eq!(EschaXGroup::PerSlot.as_kernarg(), 1, "was `true`"); + // Grouped(1) is PerSlot by construction — the kernel computes + // `slot / x_group`, and `slot / 1 == slot`. + assert_eq!( + EschaXGroup::Grouped(1).as_kernarg(), + EschaXGroup::PerSlot.as_kernarg() + ); + // The batched-prefill case: k slots per token. + assert_eq!(EschaXGroup::Grouped(8).as_kernarg(), 8); + } +} diff --git a/crates/rdna-compute/src/kernels.rs b/crates/rdna-compute/src/kernels.rs index 9dc44fe085..523b610bc5 100644 --- a/crates/rdna-compute/src/kernels.rs +++ b/crates/rdna-compute/src/kernels.rs @@ -4809,6 +4809,64 @@ pub const GEMM_Q8_0_BATCHED_WIDE_EXACT_SRC: &str = pub const GEMV_Q8_0_SRC: &str = include_str!("../../../kernels/src/gemv_q8_0.hip"); +/// Escha-W2 one-shot tile decode: packed trellis code -> bare fp16 weights. +/// No codebook LUT (65536 fp16 entries would be 128 KB, gfx1151 LDS is 64 KB); +/// the codebook is computed inline per element. See the file for the G2 gate +/// rationale (deliberately duplicates escha_ref.rs's lane maths in Rust). +pub const ESCHA_DECODE_TILES_SRC: &str = + include_str!("../../../kernels/src/escha_decode_tiles.hip"); + +/// Escha-W2 activation transforms: the 128-point Walsh-Hadamard applied to +/// both sides of every escha matmul (`escha_h128_in`, `escha_h128_out`). Two +/// entry points sharing one source file — see `ensure_kernel` call sites. +pub const ESCHA_H128_SRC: &str = include_str!("../../../kernels/src/escha_h128.hip"); + +/// Escha-W2 load-path transpose: `escha_decode_tiles` writes bare fp16 +/// IN-major `[ic, oc]` (escha's tile grid is in-major); hipfire's expert +/// slots are OUT-major `[oc, ic]`. Two entry points — `escha_bare_to_q8_0` +/// (production: transpose + Q8_0 re-quantise in one pass) and +/// `escha_bare_to_f32` (the weight-exact control arm the G4 gate uses to +/// separate wiring error from re-quantisation error). +pub const ESCHA_BARE_TO_OUTMAJOR_SRC: &str = + include_str!("../../../kernels/src/escha_bare_to_outmajor.hip"); + +/// Escha-W2 routed-expert GEMVs for the GPU-top-K (indexed) decode path, plus +/// the device-side f16 round-trip of the combine weights. Three entry points: +/// `escha_gemv_q8_0_moe_k8_indexed_batched`, +/// `escha_gemv_q8_0_wide_moe_k8_indexed_batched` and +/// `escha_round_weights_f16_rne`. The two GEMVs are per-slot-in/per-slot-out +/// (both escha phases need that shape) verbatim transcriptions of `gemv_q8_0` +/// / `gemv_q8_0_wide` — see the header of the .hip for why the pre-existing +/// indexed Q8_0 MoE kernels cannot serve escha and why the accumulate order +/// is copied rather than improved. +pub const ESCHA_MOE_GEMV_K8_INDEXED_SRC: &str = + include_str!("../../../kernels/src/escha_moe_gemv_k8_indexed.hip"); + +/// Escha-W2 routed-expert GEMVs that read the TRELLIS CODE DIRECTLY — the +/// Phase-2 fused kernels. Six entry points: the four +/// `escha_gemv_native_k{2,3}[_wide]_moe_k8_indexed_batched` production kernels +/// and the two `escha_gemv_f16[_wide]_moe_k8_indexed_batched` reference arms +/// they are gated bit-exactly against. Same lane mapping, accumulator count +/// and reduction as [`ESCHA_MOE_GEMV_K8_INDEXED_SRC`]; the only difference is +/// that the weight is decoded in-register from the code instead of read from +/// an expanded Q8_0 copy. See the .hip header for why a block has to own a +/// whole 16-wide tile column to reach the format's 0.25 B/weight floor. +pub const ESCHA_MOE_GEMV_NATIVE_SRC: &str = + include_str!("../../../kernels/src/escha_moe_gemv_native.hip"); + +/// Escha-W2 routed-expert GROUPED GEMM — the Phase-3 batched-prefill kernel. +/// Same trellis decode as [`ESCHA_MOE_GEMV_NATIVE_SRC`] and the same lane -> +/// contraction map, but `blockIdx.y` is an EXPERT rather than a (token, +/// expert) slot: the expert's code is decoded once and spent across every +/// token that routed to it, and a block owns `CTILES` adjacent tile columns so +/// the activation is read once for all of them. See the .hip header for the +/// traffic arithmetic and for which projections stay bit-identical to the +/// slot-parallel kernel (the narrow ones) and which do not (the wide ones). +pub const ESCHA_MOE_GEMM_GROUPED_WMMA_SRC: &str = + include_str!("../../../kernels/src/escha_moe_gemm_grouped_wmma.hip"); +pub const ESCHA_MOE_GEMM_GROUPED_SRC: &str = + include_str!("../../../kernels/src/escha_moe_gemm_grouped.hip"); + /// Batched Q8_0 GEMM. Same per-row math as gemv_q8_0 but holds MAX_BATCH /// per-row accumulators in registers, broadcasting each weight load across /// all batch elements. Saves the (batch_size - 1)× weight re-reads of the @@ -5800,6 +5858,16 @@ pub const CAST_F32_TO_F16_SRC: &str = include_str!("../../../kernels/src/cast_f3 /// points. See `kernels/src/bf16_round_trip.hip`. pub const BF16_ROUND_TRIP_SRC: &str = include_str!("../../../kernels/src/bf16_round_trip.hip"); +/// In-place F32 router-logits round-trip through f16 (round-to-nearest-even). +/// Escha-only: EschaLabs' runtime selects MoE top-k from f16-rounded router +/// logits; hipfire keeps logits F32 end-to-end everywhere else. Applied to +/// `router_logits` before top-k ONLY when the layer's routed experts are +/// Escha2T16/Escha3T16 (see `MoeDtypes::has_escha_experts` in +/// hipfire-dispatch and `run_moe_decode`'s call site). See +/// `kernels/src/router_logits_round_f16_rne.hip`. +pub const ROUTER_LOGITS_ROUND_F16_RNE_SRC: &str = + include_str!("../../../kernels/src/router_logits_round_f16_rne.hip"); + /// Batched partial-interleaved RoPE — per-row positions read from a /// positions[] array. Used by the batched prefill FA path. #[cfg(feature = "deltanet")] diff --git a/crates/rdna-compute/src/lib.rs b/crates/rdna-compute/src/lib.rs index 3822d3797c..2c68cdc43f 100644 --- a/crates/rdna-compute/src/lib.rs +++ b/crates/rdna-compute/src/lib.rs @@ -40,6 +40,8 @@ pub use dispatch::{ MQ6G256V2_GROUP_BYTES, }; pub use feature_flags::FeatureFlags; +/// Slot-to-activation mapping for `Gpu::escha_h128_batched`'s input side. +pub use gemv::{escha_grouped_tile, EschaXGroup}; pub use hip_bridge::{HipError, HipResult}; use std::sync::OnceLock; @@ -86,3 +88,23 @@ mod tests { ); } } + +/// Process-wide count of `escha_h128_in_batched` / `escha_h128_out_batched` +/// launches. The Escha-W2 forward path is launch-bound (Task 8), so the +/// per-token launch budget is a correctness-adjacent property: the G4 gate +/// reads this to report the achieved launches/token instead of asserting a +/// comment. Diagnostic only — nothing synchronises on it. +pub static ESCHA_H128_LAUNCHES: std::sync::atomic::AtomicU64 = std::sync::atomic::AtomicU64::new(0); + +/// Read the H128 batched-launch tally. +/// +/// Monotonic for the life of the process — there is NO reset, by design (an +/// earlier version of this doc said "and optionally reset"; no such affordance +/// exists). Callers wanting a delta snapshot the counter before and after the +/// region of interest, which is what every escha gate does. +/// +/// The count reflects launches ISSUED, not completed — nothing synchronises on +/// it. Sample it after a `device_synchronize` if that distinction matters. +pub fn escha_h128_launches() -> u64 { + ESCHA_H128_LAUNCHES.load(std::sync::atomic::Ordering::Relaxed) +} diff --git a/crates/rdna-compute/src/norm.rs b/crates/rdna-compute/src/norm.rs index e77e6b70f2..26fbd360cd 100644 --- a/crates/rdna-compute/src/norm.rs +++ b/crates/rdna-compute/src/norm.rs @@ -1805,6 +1805,62 @@ impl Gpu { result } + /// In-place F32 → f16 → F32 round-trip on MoE router logits + /// (round-to-nearest-even). Escha-only precision-matching step: see + /// `kernels/src/router_logits_round_f16_rne.hip` for why and + /// `hipfire_dispatch::families::moe::MoeDtypes::has_escha_experts` for + /// the gate. Callers other than the escha routed paths must not call + /// this — every other model's router logits stay F32 end-to-end. + /// + /// Numel-driven and layout-agnostic: `x` may be one decode token's + /// `[n_exp]` logits or a batched prefill chunk's `[n x n_exp]` block. Both + /// call sites exist and MUST both round, or batched prefill would select + /// experts from unrounded logits while decode selects from rounded ones — + /// a systematic route divergence, not the ~0.42% f16-boundary straddle + /// that is inherent to the format. A batched caller should pass a view of + /// exactly the live rows; the whole scratch would also round stale tail + /// rows (harmless, but a larger launch for nothing). + pub fn router_logits_round_f16_rne(&mut self, x: &GpuTensor) -> HipResult<()> { + self.bind_thread()?; + self.ensure_kernel( + "router_logits_round_f16_rne", + kernels::ROUTER_LOGITS_ROUND_F16_RNE_SRC, + "router_logits_round_f16_rne", + )?; + let xp = x.buf.as_ptr(); + let n = x.numel() as i32; + let mut params: Vec<*mut c_void> = vec![ + &xp as *const _ as *mut c_void, + &n as *const _ as *mut c_void, + ]; + let block_size = 256u32; + let grid = (((n as u32) + block_size - 1) / block_size).max(1); + let bytes = crate::profile::elementwise_bytes(n as usize); + let timer = crate::profile::begin_timer( + &self.hip, + "router_logits_round_f16_rne", + "router_logits_round_f16_rne", + bytes, + ); + let result = self.launch_maybe_blob( + "router_logits_round_f16_rne", + [grid, 1, 1], + [block_size, 1, 1], + 0, + &mut params, + || { + let mut b = hip_bridge::KernargBlob::new(); + b.push_ptr(xp); + b.push_i32(n); + b + }, + ); + if let Some(t) = timer { + t.finish(&self.hip); + } + result + } + /// Sigmoid activation, in-place. #[cfg(feature = "deltanet")] /// Repeat-interleave Q and K key heads up to value heads count. diff --git a/crates/rdna-compute/src/scratch.rs b/crates/rdna-compute/src/scratch.rs index 7f27a8c36d..a54fd44958 100644 --- a/crates/rdna-compute/src/scratch.rs +++ b/crates/rdna-compute/src/scratch.rs @@ -60,6 +60,116 @@ pub struct ScratchState { /// in one allocation; grows-never-shrinks. pub sample_partials: Option, pub sample_partials_bytes: usize, + /// Escha-W2 BATCHED-PREFILL routed scratch. See [`EschaPrefillScratch`]. + /// + /// Lives here, on the per-GPU scratch state, rather than on the model's + /// `PrefillBatchScratch`, for two reasons. It is MODEL-GLOBAL — one copy + /// serves all 40 layers, because the routed half of a layer is fully + /// consumed before the next layer's begins — and the per-layer escha + /// scratch it mirrors is only `[k]` slots (~272 KB); a `[max_batch x k]` + /// version PER LAYER would be ~3 GB. It is also grows-never-shrinks and + /// lazily allocated, exactly like `gemv_residual_tmp` and + /// `paro_fused_scratch` beside it, so a non-escha model never pays a byte. + pub escha_prefill: Option, +} + +/// Scratch for `escha_routed_prefill_indexed`: the per-slot buffers of the +/// eight-phase escha routed pipeline, sized for `slots = n_tokens * k` rather +/// than decode's `k`. +/// +/// ONE device allocation carved into seven views. The packing matters: the HIP +/// allocator rounds every allocation to a 2 MiB granule, and seven separate +/// buffers would charge that seven times for buffers that are always allocated +/// and freed together (the same lesson `PackedExpertOwners` records at a much +/// larger scale — 20,480 allocations became 80 and 30.4 GB came back). +/// +/// At A3B shapes with `max_batch = 256`, `k = 8` (2 048 slots, hidden 2 048, +/// mi 512) this is `2048 * (3*2048 + 6*512 + 1) * 4 B` = ~75 MB total, for the +/// whole model. +pub struct EschaPrefillScratch { + /// Slot capacity this was sized for (`n_tokens * k`). + pub slots: usize, + /// Model hidden size the views were carved against. + pub hidden: usize, + /// Routed-expert intermediate size the views were carved against. + pub mi: usize, + /// The single owning allocation. Every field below is a non-owning view. + owner: DeviceBuffer, + /// `[slots]` f32 — f16-rounded combine weights. + pub weights: GpuTensor, + /// `[slots, hidden]` f32. + pub xh_gu: GpuTensor, + /// `[slots, 2*mi]` f32. + pub mid_gu: GpuTensor, + /// `[slots, 2*mi]` f32. + pub y_gu: GpuTensor, + /// `[slots, mi]` f32. + pub h: GpuTensor, + /// `[slots, mi]` f32. + pub xh_dn: GpuTensor, + /// `[slots, hidden]` f32. + pub mid_dn: GpuTensor, + /// `[slots, hidden]` f32 — the per-slot expert outputs the combine reduces. + pub y_dn: GpuTensor, +} + +/// Non-owning views of exactly the LIVE prefix of an [`EschaPrefillScratch`]. +/// +/// Returned by value so the caller does not hold a borrow of `Gpu` across the +/// kernel launches that consume it — `ensure_escha_prefill_scratch` needs +/// `&mut Gpu`, and so does every launch, so a `&EschaPrefillScratch` tied to +/// `gpu.scratch` could not survive the first launch. +/// +/// Every field is a `sub_offset` alias: dropping one is a no-op, and none of +/// them may be passed to `free_tensor`. The owning allocation stays in +/// `ScratchState::escha_prefill`. +pub struct EschaPrefillViews { + /// Live slot count these views were cut to (`n_tokens * k`). + pub slots: usize, + pub weights: GpuTensor, + pub xh_gu: GpuTensor, + pub mid_gu: GpuTensor, + pub y_gu: GpuTensor, + pub h: GpuTensor, + pub xh_dn: GpuTensor, + pub mid_dn: GpuTensor, + pub y_dn: GpuTensor, +} + +impl EschaPrefillScratch { + /// Cut views of the first `slots` slots. + /// + /// The `escha_h128_batched` / GEMV wrappers validate buffer lengths + /// EXACTLY, so they must be handed the live prefix rather than the whole + /// (larger) capacity — and that exactness is deliberate: it is what turns + /// a slot-count mistake into a rejected launch instead of a silent + /// out-of-range write. + pub fn views(&self, slots: usize) -> EschaPrefillViews { + let (hidden, mi) = (self.hidden, self.mi); + EschaPrefillViews { + slots, + weights: self.weights.sub_offset(0, slots), + xh_gu: self.xh_gu.sub_offset(0, slots * hidden), + mid_gu: self.mid_gu.sub_offset(0, slots * 2 * mi), + y_gu: self.y_gu.sub_offset(0, slots * 2 * mi), + h: self.h.sub_offset(0, slots * mi), + xh_dn: self.xh_dn.sub_offset(0, slots * mi), + mid_dn: self.mid_dn.sub_offset(0, slots * hidden), + y_dn: self.y_dn.sub_offset(0, slots * hidden), + } + } + + /// f32 elements one slot occupies across all seven buffers. + /// + /// Pure, so the packing arithmetic is checkable without a GPU — and it + /// must be, because a wrong stride here puts every slot after the first at + /// a wrong offset, which is finite, plausible, wrong output rather than a + /// fault. + pub fn elems_per_slot(hidden: usize, mi: usize) -> usize { + // weights(1) + xh_gu(hidden) + mid_gu(2mi) + y_gu(2mi) + h(mi) + // + xh_dn(mi) + mid_dn(hidden) + y_dn(hidden) + 1 + 3 * hidden + 6 * mi + } } // ── Shared kernel dispatch helpers ────────────────────────────────────── @@ -158,50 +268,40 @@ pub(crate) fn launch_maybe_blob( if record { // Single decision point for how a launch is recorded: same // artifact lookup shape as `Gpu::launch_maybe_blob_bound`. - let artifact = compiler - .as_ref() - .and_then(|c| { - c - .compiled_kernels() - .get(func_name) - .or_else(|| match func_name { - "mq_rotate_x" => c.compiled_kernels().get("gemv_mq4g256"), - "deinterleave_f32_batched" => { - c.compiled_kernels().get("deinterleave_batched") - } - name if name.starts_with("gemv_hfq4g256_residual_sigmoid_scaled_gpu") => { - c + let artifact = compiler.as_ref().and_then(|c| { + c.compiled_kernels() + .get(func_name) + .or_else(|| match func_name { + "mq_rotate_x" => c.compiled_kernels().get("gemv_mq4g256"), + "deinterleave_f32_batched" => { + c.compiled_kernels().get("deinterleave_batched") + } + name if name.starts_with("gemv_hfq4g256_residual_sigmoid_scaled_gpu") => { + c.compiled_kernels().get("gemv_hfq4g256_residual_scaled") + } + "gemv_hfq4g256_moe_gate_up_k8_indexed" => c .compiled_kernels() - .get("gemv_hfq4g256_residual_scaled") - } - "gemv_hfq4g256_moe_gate_up_k8_indexed" => c - .compiled_kernels() - .get("gemv_hfq4g256_moe_gate_up_indexed"), - name if name.starts_with("gemv_hfq4g256_multirow_r") => c - .compiled_kernels() - .get("gemv_hfq4g256_multirow_default") - .or_else(|| { - c - .compiled_kernels() - .get("gemv_hfq4g256_multirow_rdna3") - }), - name if name.starts_with("gemv_hfq4g256_residual_multirow_r") => c - .compiled_kernels() - .get("gemv_hfq4g256_residual_multirow_default") - .or_else(|| { - c - .compiled_kernels() - .get("gemv_hfq4g256_residual_multirow_rdna3") - }), - _ => None, - }) - .or_else(|| { - func_name - .strip_suffix("_f32") - .and_then(|name| c.compiled_kernels().get(name)) - }) - .cloned() - }); + .get("gemv_hfq4g256_moe_gate_up_indexed"), + name if name.starts_with("gemv_hfq4g256_multirow_r") => c + .compiled_kernels() + .get("gemv_hfq4g256_multirow_default") + .or_else(|| c.compiled_kernels().get("gemv_hfq4g256_multirow_rdna3")), + name if name.starts_with("gemv_hfq4g256_residual_multirow_r") => c + .compiled_kernels() + .get("gemv_hfq4g256_residual_multirow_default") + .or_else(|| { + c.compiled_kernels() + .get("gemv_hfq4g256_residual_multirow_rdna3") + }), + _ => None, + }) + .or_else(|| { + func_name + .strip_suffix("_f32") + .and_then(|name| c.compiled_kernels().get(name)) + }) + .cloned() + }); replay.as_mut().unwrap().record_hip_launch_typed_bound( hip, func_name, @@ -217,11 +317,15 @@ pub(crate) fn launch_maybe_blob( capture_blobs.push(blob.into_vec()); let buf = capture_blobs.last_mut().unwrap(); let func = &functions[func_name]; - unsafe { hip.launch_kernel_blob(func, grid, block, shared_mem, stream, buf.as_mut_slice()) } + unsafe { + hip.launch_kernel_blob(func, grid, block, shared_mem, stream, buf.as_mut_slice()) + } } else { let mut bytes = blob.into_vec(); let func = &functions[func_name]; - unsafe { hip.launch_kernel_blob(func, grid, block, shared_mem, stream, bytes.as_mut_slice()) } + unsafe { + hip.launch_kernel_blob(func, grid, block, shared_mem, stream, bytes.as_mut_slice()) + } } } else { let func = &functions[func_name]; @@ -394,6 +498,80 @@ impl ScratchState { Ok(self.gemv_residual_tmp.as_ref().unwrap()) } + /// Ensure the Escha-W2 batched-prefill routed scratch can serve `slots` + /// slots at this model's `hidden` / `mi`, growing on demand. + /// + /// Reallocates whenever the request exceeds the current capacity OR the + /// shapes differ — `hidden` / `mi` are fixed per model, so a shape change + /// means a different model on the same `Gpu`, and reusing views carved for + /// the old shapes would silently read the wrong strides. + pub fn ensure_escha_prefill( + &mut self, + hip: &HipRuntime, + device_id: i32, + slots: usize, + hidden: usize, + mi: usize, + ) -> HipResult<&EschaPrefillScratch> { + crate::graph::bind_thread(hip, device_id)?; + let fits = self + .escha_prefill + .as_ref() + .is_some_and(|e| e.slots >= slots && e.hidden == hidden && e.mi == mi); + if !fits { + let per_slot = EschaPrefillScratch::elems_per_slot(hidden, mi); + let total = slots.checked_mul(per_slot).ok_or_else(|| { + hip_bridge::HipError::new(0, "escha prefill scratch size overflow") + })?; + let owner = hip.malloc(total * 4)?; + // Carve seven views out of one allocation, in declaration order. + // `off` counts f32 ELEMENTS; each view is a non-owning alias, so + // dropping them is a no-op and only `owner` holds the memory. + let base = owner.as_ptr() as *mut u8; + let mut off = 0usize; + let mut carve = |len: usize| -> GpuTensor { + // SAFETY: `off + len <= total` by construction (the sum of the + // seven lengths is exactly `slots * per_slot`), and the view + // never outlives `owner`, which is moved into the struct below + // and only replaced by this same function. + let t = GpuTensor { + buf: unsafe { DeviceBuffer::from_raw(base.add(off * 4) as *mut _, len * 4) }, + shape: vec![len], + dtype: DType::F32, + }; + off += len; + t + }; + let weights = carve(slots); + let xh_gu = carve(slots * hidden); + let mid_gu = carve(slots * 2 * mi); + let y_gu = carve(slots * 2 * mi); + let h = carve(slots * mi); + let xh_dn = carve(slots * mi); + let mid_dn = carve(slots * hidden); + let y_dn = carve(slots * hidden); + debug_assert_eq!(off, total, "escha prefill scratch carve must be exact"); + if let Some(prev) = self.escha_prefill.take() { + hip.free(prev.owner)?; + } + self.escha_prefill = Some(EschaPrefillScratch { + slots, + hidden, + mi, + owner, + weights, + xh_gu, + mid_gu, + y_gu, + h, + xh_dn, + mid_dn, + y_dn, + }); + } + Ok(self.escha_prefill.as_ref().unwrap()) + } + /// Lazily initialize MagnumQuant FWHT sign tables (256 floats each, seeds 42 and 1042). pub fn ensure_mq_signs( &mut self, @@ -544,6 +722,23 @@ impl ScratchState { /// When either recorder is active (`capture_mode` or `replay.is_recording()`) /// the kernel always runs so the tape stays complete; the skip only applies /// to the live non-recording path. Returns the FP16 device pointer. + /// + /// # The pointer key is only sound for back-to-back same-`x` dispatches + /// + /// It identifies the SOURCE BUFFER, not its CONTENTS. A caller that hands + /// this a stable scratch allocation whose contents are rewritten between + /// calls — a per-layer activation buffer, which is most of them — gets the + /// first layer's conversion for every subsequent layer, silently. Two + /// call sites have already been bitten (the MTP lm_head, τ 1.85 → 1.01; + /// Escha-W2's `wo`) and both were fixed by switching to + /// [`Self::convert_fp16_x_uncached`]. + /// + /// So: use this ONLY when the same `x` is consumed by several dispatches + /// with nothing writing to it in between (the Q/K/V case it was built + /// for). If `x` is a per-layer or per-step buffer, use + /// `convert_fp16_x_uncached`; one extra elementwise kernel is far cheaper + /// than the GEMM it feeds, and is always correct. Writers of a cached + /// buffer may instead call [`Self::invalidate_x_caches_for`]. pub fn ensure_fp16_x( &mut self, hip: &HipRuntime, @@ -582,7 +777,12 @@ impl ScratchState { self.fp16_x_source_ptr = std::ptr::null_mut(); // force reconversion after realloc } - let must_convert = scratch_must_convert(capture_mode, replay.is_recording(), self.fp16_x_source_ptr, src_ptr); + let must_convert = scratch_must_convert( + capture_mode, + replay.is_recording(), + self.fp16_x_source_ptr, + src_ptr, + ); if must_convert { let in_ptr = src_ptr; let out_ptr = self.fp16_x_scratch.as_ref().unwrap().as_ptr(); @@ -630,6 +830,27 @@ impl ScratchState { /// every layer), where pointer-keyed caching would read stale FP16. /// Always launches; both recorders observe the same launch via /// `launch_maybe_blob`'s unified `record || capture_mode || force_blob` gate. + /// + /// # It also INVALIDATES the cache, and must + /// + /// This writes the SHARED `fp16_x_scratch` — the same buffer + /// [`Self::ensure_fp16_x`] hands out — so after it runs, any cached + /// `fp16_x_source_ptr` marker describes a buffer that no longer holds that + /// pointer's data. Leaving the marker set makes the very next + /// `ensure_fp16_x` for that pointer a CACHE HIT onto this call's + /// conversion of a completely different tensor. + /// + /// That is not hypothetical. On Escha-W2 batched prefill it fired every + /// layer: `wo` (`gemm_q8_0_residual_wmma`, cached) converted + /// `dn_normed_batch` in layer 0; `wqkv` (`gemm_q8_0_wmma`, uncached) + /// overwrote the scratch with `x_rot_batch` in layer 1; layer 1's `wo` + /// then hit the stale marker and multiplied its weights by the LA + /// rmsnorm output (magnitude ~7.5e-1) instead of the gated-norm output + /// (~1.9e-3) — a ~400x amplification, finite and fluent, that showed up + /// only as a moved argmax. The DeepSeek-V4 gfx942 call site had already + /// discovered this and nulls the marker by hand after calling; doing it + /// here makes that hand-repair unnecessary and closes the same hole for + /// every other caller. pub fn convert_fp16_x_uncached( &mut self, hip: &HipRuntime, @@ -678,14 +899,14 @@ impl ScratchState { ]; let grid = ((n_elems + 255) / 256) as u32; launch_maybe_blob( - hip, - Some(&*compiler), - functions, - stream, - capture_blobs, - capture_mode, - force_blob_path, - Some(replay), + hip, + Some(&*compiler), + functions, + stream, + capture_blobs, + capture_mode, + force_blob_path, + Some(replay), "convert_f32_to_f16", [grid, 1, 1], [256, 1, 1], @@ -699,6 +920,14 @@ impl ScratchState { b }, )?; + // The shared scratch no longer holds whatever `fp16_x_source_ptr` + // says it holds. See this function's doc comment — dropping the + // marker is part of the contract, not an optimisation, so it is + // unconditional. There is deliberately no lever that restores the + // pre-fix behaviour: the only thing it could do is reinstate a + // known silent-wrong-output defect, for any model with a Q8_0 + // `wo`/`w_down` in batched prefill. + self.fp16_x_source_ptr = std::ptr::null_mut(); Ok(self.fp16_x_scratch.as_ref().unwrap().as_ptr()) } @@ -745,7 +974,12 @@ impl ScratchState { self.fp8_x_source_ptr = std::ptr::null_mut(); } - let must_convert = scratch_must_convert(capture_mode, replay.is_recording(), self.fp8_x_source_ptr, src_ptr); + let must_convert = scratch_must_convert( + capture_mode, + replay.is_recording(), + self.fp8_x_source_ptr, + src_ptr, + ); if must_convert { let in_ptr = src_ptr; let out_ptr = self.fp8_x_scratch.as_ref().unwrap().as_ptr(); @@ -929,14 +1163,14 @@ impl ScratchState { let bytes = crate::profile::mq_rotate_bytes(k); let timer = crate::profile::begin_timer(hip, "fwht", "mq_rotate_x", bytes); let result = launch_maybe_blob( - hip, - Some(compiler), - functions, - stream, - capture_blobs, - capture_mode, - force_blob_path, - Some(replay), + hip, + Some(compiler), + functions, + stream, + capture_blobs, + capture_mode, + force_blob_path, + Some(replay), "mq_rotate_x", [n_groups, 1, 1], [32, 1, 1], @@ -997,14 +1231,14 @@ impl ScratchState { let bytes = crate::profile::mq_rotate_bytes(k) * batch_size; let timer = crate::profile::begin_timer(hip, "fwht", "mq_rotate_x_batched", bytes); let result = launch_maybe_blob( - hip, - Some(compiler), - functions, - stream, - capture_blobs, - capture_mode, - force_blob_path, - Some(replay), + hip, + Some(compiler), + functions, + stream, + capture_blobs, + capture_mode, + force_blob_path, + Some(replay), "mq_rotate_x", [n_groups * batch_size as u32, 1, 1], [32, 1, 1], @@ -1062,14 +1296,14 @@ impl ScratchState { let bytes = crate::profile::mq_rotate_bytes(k); let timer = crate::profile::begin_timer(hip, "fwht", "mq_rotate_x_128", bytes); let result = launch_maybe_blob( - hip, - Some(compiler), - functions, - stream, - capture_blobs, - capture_mode, - force_blob_path, - Some(replay), + hip, + Some(compiler), + functions, + stream, + capture_blobs, + capture_mode, + force_blob_path, + Some(replay), "mq_rotate_x_128", [n_groups, 1, 1], [32, 1, 1], @@ -1129,14 +1363,14 @@ impl ScratchState { let bytes = k * 4 * 3 + 2 * 256 * 4; let timer = crate::profile::begin_timer(hip, "fwht", "rotate_x_mq_awq", bytes); let result = launch_maybe_blob( - hip, - Some(compiler), - functions, - stream, - capture_blobs, - capture_mode, - force_blob_path, - Some(replay), + hip, + Some(compiler), + functions, + stream, + capture_blobs, + capture_mode, + force_blob_path, + Some(replay), "rotate_x_mq_awq", [n_groups, 1, 1], [32, 1, 1], @@ -1202,14 +1436,14 @@ impl ScratchState { let bytes = (k * 4 * 3 + 2 * 256 * 4) * batch_size; let timer = crate::profile::begin_timer(hip, "fwht", "rotate_x_mq_awq_batched", bytes); let result = launch_maybe_blob( - hip, - Some(compiler), - functions, - stream, - capture_blobs, - capture_mode, - force_blob_path, - Some(replay), + hip, + Some(compiler), + functions, + stream, + capture_blobs, + capture_mode, + force_blob_path, + Some(replay), "rotate_x_mq_awq", [n_groups, batch_size as u32, 1], [32, 1, 1], @@ -1285,14 +1519,14 @@ impl ScratchState { let bytes = crate::profile::mq_rotate_bytes(k) + k; let timer = crate::profile::begin_timer(hip, "fwht", "mq_rotate_x_dual_fp8", bytes); let result = launch_maybe_blob( - hip, - Some(compiler), - functions, - stream, - capture_blobs, - capture_mode, - force_blob_path, - Some(replay), + hip, + Some(compiler), + functions, + stream, + capture_blobs, + capture_mode, + force_blob_path, + Some(replay), "mq_rotate_x_dual_fp8_gfx12", [n_groups, 1, 1], [32, 1, 1], @@ -1413,3 +1647,27 @@ fn alloc_tensor_on( dtype, }) } + +#[cfg(test)] +mod escha_prefill_scratch_tests { + use super::EschaPrefillScratch; + + /// The carve arithmetic at the real A3B shapes. A wrong `elems_per_slot` + /// puts every buffer after the first at a wrong offset — plausible, + /// finite, wrong values rather than a fault — and the `debug_assert_eq!` + /// in `ensure_escha_prefill` that would catch it is compiled out of the + /// release build this actually runs in. + #[test] + fn elems_per_slot_matches_the_a3b_shapes() { + // hidden = 2048, mi = 512. + // weights 1 + xh_gu 2048 + mid_gu 1024 + y_gu 1024 + h 512 + // + xh_dn 512 + mid_dn 2048 + y_dn 2048 + assert_eq!(EschaPrefillScratch::elems_per_slot(2048, 512), 9217); + // The whole-model footprint the design is sized against: max_batch 256 + // x k 8 = 2048 slots. ~75 MB, ONE copy for all 40 layers (a per-layer + // one would be ~3 GB). + let bytes = 2048 * EschaPrefillScratch::elems_per_slot(2048, 512) * 4; + assert_eq!(bytes, 75_505_664); + assert!(bytes < 80 << 20); + } +} diff --git a/crates/saddle-lab/examples/infer.rs b/crates/saddle-lab/examples/infer.rs index 15162f5a7b..51c35989e3 100644 --- a/crates/saddle-lab/examples/infer.rs +++ b/crates/saddle-lab/examples/infer.rs @@ -37,7 +37,7 @@ fn main() { } let args: Vec = std::env::args().collect(); if args.len() < 2 { - eprintln!("Usage: infer [--image ] [--no-think] [prompt...]"); + eprintln!("Usage: infer [--image ] [--no-think] [--kv-seq N] [--max-tokens N] [prompt...]"); std::process::exit(1); } @@ -49,10 +49,23 @@ fn main() { .position(|a| a == "--max-tokens") .and_then(|i| args.get(i + 1).and_then(|v| v.parse().ok())) .unwrap_or(2048); + // KV capacity. Was hardwired to 4096; a longer prompt then ran off the end + // of the cache and surfaced as `hipMemcpy H2D: illegal memory access` from + // deep inside the prefill loop, with nothing pointing at the real cause. + let kv_seq: usize = args + .iter() + .position(|a| a == "--kv-seq") + .and_then(|i| args.get(i + 1).and_then(|v| v.parse().ok())) + .unwrap_or(4096); + // `--kv-full` keeps K/V unquantised. The 10 full-attention layers are what + // carry long-range retrieval, so KV precision is the thing to vary when + // deep-context recall fails. let kv_mode: &str = if args.iter().any(|a| a == "--givens4") { "givens4" } else if args.iter().any(|a| a == "--givens2") { "givens2" + } else if args.iter().any(|a| a == "--kv-full") { + "full" } else { "q8" }; @@ -69,17 +82,29 @@ fn main() { skip_next = false; continue; } - if a == "--no-think" || a == "--debug-compare" || a == "--givens4" || a == "--givens2" { + if a == "--no-think" + || a == "--debug-compare" + || a == "--givens4" + || a == "--givens2" + || a == "--kv-full" + || a == "--no-ngram-block" + { continue; } - if a == "--image" || a == "--max-tokens" { + if a == "--image" + || a == "--max-tokens" + || a == "--kv-seq" + || a == "--repeat-penalty" + || a == "--dn-state" + || a == "--temp" + { skip_next = true; continue; } positional.push(a.as_str()); } let model_path = positional.first().unwrap_or_else(|| { - eprintln!("Usage: infer [--image ] [--no-think] [prompt...]"); + eprintln!("Usage: infer [--image ] [--no-think] [--kv-seq N] [--max-tokens N] [prompt...]"); std::process::exit(1); }); let prompt_text = if positional.len() > 1 { @@ -222,8 +247,7 @@ fn main() { } .expect("failed to load text weights"); - let kv_seq = 4096usize; - eprintln!("KV cache: {kv_mode}"); + eprintln!("KV cache: {kv_mode} (capacity {kv_seq} tokens)"); let mut kv_cache = match kv_mode { "givens4" => llama::KvCache::new_gpu_asym3( &mut gpu, @@ -241,6 +265,14 @@ fn main() { kv_seq, ) .unwrap(), + "full" => llama::KvCache::new_gpu( + &mut gpu, + text_config.n_layers, + text_config.n_kv_heads, + text_config.head_dim, + kv_seq, + ) + .unwrap(), _ => llama::KvCache::new_gpu_q8( &mut gpu, text_config.n_layers, @@ -250,7 +282,22 @@ fn main() { ) .unwrap(), }; - let mut dn_state = DeltaNetState::new(&mut gpu, &text_config).unwrap(); + // DeltaNet's recurrent S matrix defaults to Q8. 30 of this model's 40 layers + // are recurrent, so at long context that state is requantised thousands of + // times in sequence and the noise compounds — the suspected cause of the + // ~8k degeneracy. `--dn-state fp32` is the ground-truth comparison. + let dn_quant = match args + .iter() + .position(|a| a == "--dn-state") + .and_then(|i| args.get(i + 1)) + .map(|s| s.as_str()) + { + Some("fp32") => qwen35::StateQuant::FP32, + Some("q4") => qwen35::StateQuant::Q4, + _ => qwen35::StateQuant::Q8, + }; + eprintln!("DeltaNet state: {dn_quant:?}"); + let mut dn_state = DeltaNetState::new_with_quant(&mut gpu, &text_config, dn_quant).unwrap(); if debug_cmp { let mut kv2 = llama::KvCache::new_gpu( @@ -366,14 +413,87 @@ fn main() { } ); - let sc = llama::SamplingConfig::text_thinking(); + // Refuse to start rather than fault the GPU partway through prefill. + let needed = prompt_tokens.len() + max_tokens; + if needed > kv_seq { + eprintln!( + "error: prompt ({} tokens) + --max-tokens ({}) = {} exceeds KV \ + capacity ({}). Re-run with --kv-seq {}.", + prompt_tokens.len(), + max_tokens, + needed, + kv_seq, + needed.next_power_of_two(), + ); + std::process::exit(2); + } + + // Verbatim-recall tasks are hostile to a repetition penalty: re-emitting a + // proper noun you just wrote is exactly what it suppresses. Overridable so + // that effect can be measured rather than assumed. + let mut sc = llama::SamplingConfig::text_thinking(); + if let Some(v) = args + .iter() + .position(|a| a == "--repeat-penalty") + .and_then(|i| args.get(i + 1).and_then(|v| v.parse::().ok())) + { + sc.repeat_penalty = v; + } + // Default temp is 0.3, i.e. NOT greedy. Comparing one sample per config + // across that is comparing noise; `--temp 0` makes an A/B mean something. + if let Some(v) = args + .iter() + .position(|a| a == "--temp") + .and_then(|i| args.get(i + 1).and_then(|v| v.parse::().ok())) + { + sc.think_temp = v; + sc.answer_temp = v; + } + eprintln!( + "Sampling: think_temp {} answer_temp {} top_p {} repeat_penalty {} window {}", + sc.think_temp, sc.answer_temp, sc.top_p, sc.repeat_penalty, sc.repeat_window + ); let scratch = qwen35::Qwen35Scratch::new(&mut gpu, &text_config, sc.repeat_window) .expect("failed to create scratch"); - // Prefill (zero-alloc scratch path) + let no_ngram_block = args.iter().any(|a| a == "--no-ngram-block"); + let logit_probe: Option = std::env::var("HIPFIRE_LOGIT_PROBE") + .ok() + .and_then(|v| v.parse().ok()) + .filter(|&n: &usize| n > 0); + + // Prefill. Text-only takes the BATCHED path; the per-token loop below is + // only for VL (interleaved image embeddings) and for the logit probe, which + // needs a distribution per position. + // + // The loop runs the DECODE kernel once per prompt token, so it prefills at + // decode speed — ~13 minutes for an 8k prompt on the dense 27B. Batched + // prefill does the same work through the MMQ/WMMA path in seconds. Nothing + // required the slow path for text; it was just what this example did. let t_pf = Instant::now(); + let use_batched = !vl_mode && logit_probe.is_none(); + if use_batched { + qwen35::forward_prefill_batch( + &mut gpu, + &weights, + &text_config, + &prompt_tokens, + 0, + &mut kv_cache, + &mut dn_state, + &scratch, + None, + None, + None, + None, + ) + .expect("forward_prefill_batch failed"); + } let mut visual_idx = 0usize; for (pos, &token) in prompt_tokens.iter().enumerate() { + if use_batched { + break; + } if vl_mode && token == IMAGE_PAD_ID && visual_idx < n_visual_tokens { let vt = visual_tokens.as_ref().unwrap(); let emb = &vt[visual_idx * text_config.dim..(visual_idx + 1) * text_config.dim]; @@ -401,6 +521,38 @@ fn main() { &scratch, ) .expect("forward_scratch failed"); + + // HIPFIRE_LOGIT_PROBE=N: every N prefill positions, report the + // shape of the next-token distribution. The question this answers + // is whether long-context collapse is a SAMPLING problem or a + // LOGIT problem: if greedy stays coherent while temp-1.0 sampling + // derails, the argmax is fine and the tail has grown — which shows + // up here as top-1 mass falling and top-20 mass falling with it. + if let Some(every) = logit_probe { + if pos % every == 0 { + let lg = gpu.download_f32(&scratch.logits).unwrap(); + let mx = lg.iter().cloned().fold(f32::NEG_INFINITY, f32::max); + let mut ex: Vec = lg.iter().map(|&v| (v - mx).exp()).collect(); + let sum: f32 = ex.iter().sum(); + for v in ex.iter_mut() { + *v /= sum; + } + let ent: f32 = -ex + .iter() + .filter(|&&p| p > 0.0) + .map(|&p| p * p.ln()) + .sum::(); + let mut srt = ex.clone(); + srt.sort_by(|a, b| b.partial_cmp(a).unwrap()); + let top1 = srt[0]; + let top20: f32 = srt.iter().take(20).sum(); + let top100: f32 = srt.iter().take(100).sum(); + println!( + "PROBE pos={pos} top1={top1:.4} top20={top20:.4} \ + top100={top100:.4} entropy={ent:.3} max_logit={mx:.2}" + ); + } + } } } let prefill_len = prompt_tokens.len(); @@ -434,7 +586,18 @@ fn main() { let mut next_token = llama::sample_top_p(&logits, temp, sc.top_p); let t_gen = Instant::now(); - let mut token_history: Vec = prompt_tokens.clone(); + // Anti-repeat state covers ONLY the model's own output, never the prompt. + // + // This used to be seeded with `prompt_tokens.clone()`, which fed the whole + // prompt to `apply_ngram_block`. That routine hard-blocks (-INF) whatever + // token followed any repeated 3/4/5/6-gram, so the instant the model emitted + // a 3-gram occurring in the prompt — unavoidable when quoting the prompt — + // the next token of the quote was banned. Verbatim quotation was impossible + // by construction: `VIOLET-ANVIL-62` came back as `VIOLETANVIL62`, and the + // model then noticed the mangling and retried forever, which read as + // long-context "degeneracy". `test_long_ctx.rs` already documents this + // hazard and slices both corrections to the current turn; this file did not. + let mut token_history: Vec = Vec::new(); let mut generated = Vec::new(); let mut think_count = 0usize; @@ -495,7 +658,41 @@ fn main() { ) .expect("forward_scratch failed"); logits = gpu.download_f32(&scratch.logits).unwrap(); - if !in_thinking { + // Generation-time distribution shape. The prefill probe showed the + // representation is healthy out to 8k, so if collapse is real it must + // appear HERE — either bad from the first generated token (a state + // problem) or degrading over the first few (a feedback spiral). + // Sampled BEFORE ngram-block / repeat-penalty so it reflects the raw + // model distribution, not the harness's corrections. + if let Some(_) = logit_probe { + if generated.len() < 40 || generated.len() % 50 == 0 { + let mx = logits.iter().cloned().fold(f32::NEG_INFINITY, f32::max); + let mut ex: Vec = logits.iter().map(|&v| (v - mx).exp()).collect(); + let sum: f32 = ex.iter().sum(); + for v in ex.iter_mut() { + *v /= sum; + } + let ent: f32 = -ex + .iter() + .filter(|&&p| p > 0.0) + .map(|&p| p * p.ln()) + .sum::(); + let mut srt = ex.clone(); + srt.sort_by(|a, b| b.partial_cmp(a).unwrap()); + let top20: f32 = srt.iter().take(20).sum(); + println!( + "GPROBE step={} top1={:.4} top20={:.4} entropy={:.3} max_logit={:.2}", + generated.len(), + srt[0], + top20, + ent, + mx + ); + } + } + // `token_history` holds generated tokens only, so this can only block + // the model repeating ITSELF — never block it quoting the prompt. + if !in_thinking && !no_ngram_block { llama::apply_ngram_block(&mut logits, &token_history); } llama::apply_repeat_penalty( diff --git a/crates/saddle-lab/examples/run.rs b/crates/saddle-lab/examples/run.rs index 67f414f82a..ce73a0c20c 100644 --- a/crates/saddle-lab/examples/run.rs +++ b/crates/saddle-lab/examples/run.rs @@ -273,7 +273,7 @@ fn main() { } let input_norm = hipfire_runtime::tokenizer::maybe_normalize_prompt(input); let input: &str = &input_norm; - if hipfire_runtime::config::get().prompt_token_heat { + if hipfire_runtime::config::get().prompt_token_heat { tokenizer.dump_prompt_heat(input); } @@ -482,10 +482,18 @@ fn main() { target_slot.forward(&mut gpu, next_token, pos).unwrap(); logits = gpu.download_f32(&target_slot.scratch.logits).unwrap(); if !no_penalty { - llama::apply_ngram_block(&mut logits, &conversation_tokens); + // Scope both corrections to THIS turn's own output. Passing + // the whole conversation lets `apply_ngram_block` hard-block + // (-INF) any token that followed a repeated 3/4/5/6-gram in + // the user's text, so the model cannot quote back what it + // was just given — it emits a mangled version, notices, and + // retries. Same hazard `test_long_ctx.rs` documents. + let turn_start = conversation_tokens.len() - generated; + let turn = &conversation_tokens[turn_start..]; + llama::apply_ngram_block(&mut logits, turn); llama::apply_repeat_penalty( &mut logits, - &conversation_tokens, + turn, sc.repeat_window, sc.repeat_penalty, ); diff --git a/crates/saddle-lab/examples/tokenizer_roundtrip.rs b/crates/saddle-lab/examples/tokenizer_roundtrip.rs new file mode 100644 index 0000000000..4ea7c30160 --- /dev/null +++ b/crates/saddle-lab/examples/tokenizer_roundtrip.rs @@ -0,0 +1,69 @@ +// SPDX-License-Identifier: Apache-2.0 +// hipfire — see LICENSE and NOTICE in the project root. +// +//! Encode → decode round-trip for a model's embedded tokenizer. +//! +//! Written to chase a specific symptom: the 35B reports the planted checksum +//! `VIOLET-ANVIL-62` as `VIOLETANVIL62` — and volunteers "no hyphens" — +//! even when the source sentence is 500 tokens away, where retrieval is +//! plainly working. Characters going missing at that distance is not memory +//! decay; it points at the text pipeline. If a string does not survive +//! encode→decode, the model never saw what we think we put in the prompt and +//! every downstream "recall" measurement built on it is measuring the wrong +//! thing. + +use hipfire_runtime::hfq::HfqFile; +use std::path::Path; + +fn main() { + let args: Vec = std::env::args().collect(); + let model = args.get(1).map(String::as_str).unwrap_or_else(|| { + eprintln!("usage: tokenizer_roundtrip [extra strings...]"); + std::process::exit(1); + }); + + let hfq = HfqFile::open(Path::new(model)).expect("open hfq"); + let tk = hipfire_runtime::tokenizer::Tokenizer::from_hfq_metadata(&hfq.metadata_json) + .expect("tokenizer not in hfq metadata"); + + let mut cases: Vec = vec![ + "VIOLET-ANVIL-62".into(), + "the checksum phrase VIOLET-ANVIL-62 before".into(), + "Kestrel".into(), + "Reykjavik".into(), + "Reykjavík".into(), + "14 March 2019".into(), + "ratified in Reykjavik on 14 March 2019 by exactly eleven".into(), + "3.7-second".into(), + "the Halvorsen Gap".into(), + "escha_types_never_resolve".into(), + "a-b-c".into(), + "well-known".into(), + "2019".into(), + "62".into(), + ]; + cases.extend(args.iter().skip(2).cloned()); + + let mut bad = 0; + for s in &cases { + let ids = tk.encode(s); + let back = tk.decode(&ids); + let ok = &back == s; + if !ok { + bad += 1; + } + println!( + "{} {:?}\n -> {} ids {:?}\n -> {:?}", + if ok { "ok " } else { "FAIL" }, + s, + ids.len(), + ids, + back + ); + } + println!( + "\n{bad} of {} strings did not survive encode->decode", + cases.len() + ); + std::process::exit(if bad > 0 { 1 } else { 0 }); +} diff --git a/docs/MODELS.md b/docs/MODELS.md index b8efd4f1e8..fbc8dcd17c 100644 --- a/docs/MODELS.md +++ b/docs/MODELS.md @@ -63,6 +63,12 @@ Sizes below are **registry declarations**, not a substitute for runtime MoE layo | `qwen3.5:35b-a3b` | `qwen3.5-35b-a3b.mq4` | 19.7 | 22 | q8 | 35B / 3B-active | | `qwen3.6:35b-a3b` | `qwen3.6-35b-a3b.mq4p` | 19.8 | 22 | q8 | Default graded mq4p SKU | | `qwen3.6:35b-a3b-mq2` | `qwen3.6-35b-a3b.mq2` | 11.6 | 14 | | Floor SKU | +| `qwen3.6:35b-a3b-escha-xt` | `qwen3.6-35b-a3b.escha-xt` | 11.4 | 15 | q8 | Escha-W2 2-bit trellis, MQ4V2 dense. Fastest/smallest. PPL 8.0643 | +| `qwen3.6:35b-a3b-escha` | `qwen3.6-35b-a3b.escha` | 11.8 | 15 | q8 | **Default.** Escha-W2 2-bit trellis experts stored verbatim and decoded in the GEMV — no decode-at-load. MQ6 dense. PPL 7.6940, 725 tok/s prefill, 55 decode | +| `qwen3.6:35b-a3b-escha-pro` | `qwen3.6-35b-a3b.escha-pro` | 12.3 | 16 | q8 | Q8_0 dense, most faithful. PPL 7.6864. KLD reference for the other two | +| `qwen3.8:27b-escha-xt` | `qwen3.8-27b.escha-xt` | 10.5 | 14 | q8 | Escha-W2 2-bit trellis, MQ4V2 dense. PPL 9.7242, 122 tok/s prefill, 12.3 decode | +| `qwen3.8:27b-escha` | `qwen3.8-27b.escha` | 10.8 | 15 | q8 | **Default.** Dense arch-5 sibling of the 35B. MQ6 dense. PPL 9.6753, 119 tok/s prefill, 12.1 decode. Beats `qwen3.8-27b.mq3` on quality while smaller | +| `qwen3.8:27b-escha-pro` | `qwen3.8-27b.escha-pro` | 11.2 | 15 | q8 | Q8_0 dense, most faithful. PPL 9.6486 | | `qwen3.6:35b-a3b-mq3p` | `qwen3.6-35b-a3b.mq3p` | 17.2 | 20 | | MQ3+P graded | | `qwen3.6:35b-a3b-mq4p` | `qwen3.6-35b-a3b.mq4p` | 19.8 | 22 | | MQ4+P graded | | `qwen3.6:35b-a3b-mfp4` | `qwen3.6-35b-a3b.mfp4` | 20.2 | 22 | | MFP4-E8 | diff --git a/docs/VALIDATION.md b/docs/VALIDATION.md index 394e5ad1bc..91de64e480 100644 --- a/docs/VALIDATION.md +++ b/docs/VALIDATION.md @@ -204,7 +204,7 @@ Narrow roles. Do not widen a harness into a universal gate. | Harness | Path | Role | Not this harness | |---|---|---|---| -| **gates.sh** | [`scripts/gates.sh`](../scripts/gates.sh) | Maintained **manual** wrapper: optional Redline capture, generic serve battery, optional fresh-process perf compare (`probe_commits.sh`). Requires `--model`. | Not CI-default. Not universal. Does not call retired coherence-gate scripts. | +| **gates.sh** | [`scripts/gates.sh`](../scripts/gates.sh) | Maintained **manual** wrapper: optional Redline capture, generic serve battery, optional fresh-process perf compare (`probe_commits.sh`), and (opt-in, `--escha`) the Escha-W2 G1-G6 correctness battery. Requires `--model`. | Not CI-default. Not universal. Does not call retired coherence-gate scripts. The `--escha` arm is checkpoint-specific and is **off by default** — it is not a gate for any other model. | | **serve_harness.py** | [`scripts/serve_harness.py`](../scripts/serve_harness.py) | **Model-agnostic** user-facing serve behavior (battery / chain / session): finish reasons, runaway/empty, prefix cache, prefill/decode timing, recall hooks. | Not LFM thinking-frame specifics. Not Redline route proof. | | **serve_harness.py (LFM tag)** | [`scripts/serve_harness.py`](../scripts/serve_harness.py) | LFM2.5 serve smoke with the exact registry tag; use registry sampling or `recipe:nothink` for non-thinking framing. | Not a substitute for numerical parity oracles. | | **redline_daemon_harness.py** | [`scripts/redline_daemon_harness.py`](../scripts/redline_daemon_harness.py) | Resident-daemon **Redline** capture, phase fingerprint, shadow/parity, and timing evidence under manual-capture env. | Discovery/correctness evidence ≠ product timed-arm route proof by itself. Does not enable AQL routing. | @@ -225,6 +225,38 @@ Use only when the claim class below names them. They are not universal. | Redline certification ladder | [`docs/REDLINE.md`](REDLINE.md) | What evidence is required before Redline-attributed promotion. | | Path-specific parity / state oracle | Arch-owned example or test named by the change (when one exists) | Hidden-state, logit, KV/conv, or graph parity. If none exists for the surface → **blocked**. | +### Escha-W2 correctness gates (G1-G6) + +Checkpoint-specific, **manual**, GPU + model. Run them with +`scripts/gates.sh --escha-only --model /path/to/escha-35b.hfq`, or +individually. These are the required route for any change to the escha codec, +the H128 transforms, the escha routed executors, or the escha loader — a green +serve battery is **not** evidence for any of them. + +Measured values live in +[`escha-w2-port-design.md`](plans/escha-w2-port-design.md) §10.6. A gate that +passes with a number materially different from the one recorded there has not +passed; re-derive before re-recording. + +| Gate | Command | Asserts | Recorded result | +|---|---|---|---| +| **G1** | `python3 scripts/escha-verify-roundtrip.py ` | Verbatim repack: every `escha_code` tensor byte-identical to the source safetensors. Count asserted against `model.safetensors.index.json`; fails on zero shards, zero code tensors, or a count mismatch. | 80/80 byte-identical | +| **G2** | `cargo run --release -p rdna-compute --example test_escha_decode_gpu_vs_cpu` | GPU tile decode == `escha_ref::reconstruct`, bit-exact in fp16, at golden and 89M-element shapes. | 0 mismatched | +| **G3** | `cargo run --release -p rdna-compute --example test_escha_h128_gpu_vs_cpu` | The H128 pair == `escha_ref`, bit-exact, every launch form (single, batched broadcast / per-slot / grouped, out_batched, swiglu). | 0 mismatched | +| **G4b** | `cargo run --release -p hipfire-arch-qwen35 --example escha_router_contract -- ` | arch-6 router selects the same experts as escha's reference routing. | 0/8 differing sets | +| **G4** | `cargo run --release -p hipfire-arch-qwen35 --example escha_moe_block_gate -- ` | Whole MoE block vs escha's `moeblk_out.f16`; plus indexed-vs-host and batched-vs-per-token **equality**. | F32 max 1.828e-4 / mean 9.673e-6; Q8_0 max 2.633e-4 / mean 3.027e-5; 0 differing floats on both route comparisons | +| **G5** | `scripts/escha-kld.sh ` | KLD vs the weight-exact escha arm on a fixed teacher-forced corpus, with an asserted negative control and an asserted upper bound. | 0.0027576 nats (CI 0.0019491-0.0038610), PPL 7.6585, control prints 0.000000 | +| **G6** | `cargo run --release -p hipfire-arch-qwen35 --example escha_prefill_batch_gate -- ` | Batched prefill vs the per-token route, whole model: argmax stable, logit deltas within measured bounds, no non-finite logits. | argmax stable; max\|delta\| 4.393e-1, mean\|delta\| 7.160e-2 | + +`escha_ref` (`crates/hipfire-quantize/src/escha_ref.rs`) is the **frozen +oracle** every bit-exactness claim above rests on. It is a transcription of +EschaLabs' `ref.py` and must not be edited to make a gate pass; a gate that +disagrees with it is reporting a defect in hipfire. + +G1-G4b need only the checkpoint and the fixtures committed under +`crates/hipfire-quantize/tests/data/escha/`. G5 and G6 load the whole model +(37.6 GB resident) and take minutes each. + ## Claim → route map | Claim / change class | Minimum route(s) | Evidence kind | @@ -243,6 +275,7 @@ Use only when the claim class below names them. They are not universal. | Arch port | `methodology/arch-port-validation.md` (channel + speed; no retired coherence battery as acceptance) | Manual | | Model/route **admission** (registry evidence) | Row in [`admissions.yml`](admissions.yml) | Schema v2; exactly one evidence-bound record (LFM2.5-350M MQ4 gfx1201 retained-PM4). No inferred/wildcard rows. Registry admission/evidence is distinct from runtime wiring: the sealed LFM row does **not** select a runtime default and current automatic selection does not use it. | | MQ4R **runtime** automatic Redline default | Source predicate `mq4r_redline_default` in `crates/hipfire-runtime/src/config.rs`; policy in [`REDLINE.md`](REDLINE.md) | **Only** current automatic runtime predicate. Runtime-only: exact GPU arch `gfx1100`, `gfx1151`, or `gfx1201`; PP=1; TP=1; case-insensitive `.mq4r` → retained PM4/Auto unless disabled with the config wizard's built-in `hip` profile, another explicit backend selection, or `HIPFIRE_REPLAY_BACKEND=hip`. Model-family agnostic (no `arch_id` gate). `gfx1200` and all other arches remain opt-in. Existing LFM `.mq4` registry evidence is not auto-selected because it is not `.mq4r`, not because LFM is categorically exempt; any usable non-default retained route must still prove route support and fail closed when unsupported. **Not** registry admission, **not** Section 7 certification, and **not** a sealed-fixture claim for every default-eligible `.mq4r` model. | +| Escha-W2 codec / H128 transforms / escha routed executors / escha loader | `scripts/gates.sh --escha-only --model ` (G1-G6; see the Escha-W2 table above), against the recorded values in [`escha-w2-port-design.md`](plans/escha-w2-port-design.md) §10.6 | Manual GPU + model. Bit-exactness claims are against the frozen `escha_ref` oracle. **Not** `serve_harness.py`, and **not** a subset — G4's two `0 differing floats` rows are equality claims and a change that turns either into a tolerance needs its own argument | | Unknown surface | **Blocked** until an owner adds a row here | Fail closed | ## Retired coherence-gate scripts diff --git a/docs/investigations/2026-09-04-infer-ngram-block-prompt-history.md b/docs/investigations/2026-09-04-infer-ngram-block-prompt-history.md new file mode 100644 index 0000000000..3afa4298fb --- /dev/null +++ b/docs/investigations/2026-09-04-infer-ngram-block-prompt-history.md @@ -0,0 +1,170 @@ +# `infer.rs` banned the tokens needed to quote its own prompt + +**Status:** fixed. Lab harness only — no shipped path affected. + +A long-context coherence probe on `escha-35b` appeared to show the model +degenerating at ~8.5k tokens: mangled proper nouns, then an unbounded +self-correction loop, then word salad. None of that was the model. The +sampling harness was hard-banning the tokens required to write the answer. + +--- + +## The bug + +`crates/saddle-lab/examples/infer.rs` seeded its anti-repeat state with the +entire prompt: + +```rust +let mut token_history: Vec = prompt_tokens.clone(); // <-- the bug +``` + +and passed that to `llama::apply_ngram_block`, which scans **all** history for +repeated 3/4/5/6-grams and sets whatever token followed each earlier +occurrence to **`-INF`** — a hard ban, not a penalty. + +Quoting the prompt necessarily emits 3-grams that occur in the prompt. So the +instant the model began reproducing a planted sentence, the *next token of +that sentence* was banned. **Verbatim quotation of the prompt was impossible +by construction.** + +`crates/saddle-lab/examples/run.rs` had the same defect, passing +`&conversation_tokens` — which includes the user's own messages. + +### The hazard was already known in this repo + +`saddle-lab/examples/test_long_ctx.rs:317` documents it exactly: + +> apply_ngram_block is DISABLED across the full conversation history because +> it will aggressively block legitimate tokens that happen to follow n-grams +> from the user's earlier turns […] Only apply it over the current turn's own +> tokens. + +That file slices **both** corrections to `&history[turn_start..]`. `infer.rs` +and `run.rs` did neither. + +--- + +## Symptoms, and why each one followed + +Planted needle: `Every calibration run must terminate with the checksum phrase +VIOLET-ANVIL-62 before its results are considered admissible.` + +| symptom | cause | +|---|---| +| `VIOLET-ANVIL-62` → `VIOLETANVIL62` | hyphen (token 12) banned as the continuation of a prompt 3-gram | +| `Kestral` / `Krestrel` / `Kvestrel` in one output | each retry forced to deviate from the previously-banned path | +| `14 March 2019` → `1 March 2` → `4 March` | digits banned mid-sequence | +| unbounded `Wait, I'll copy exactly:` retry loop | model quotes → output mangled → it notices → retries → banned again | +| paraphrase-only task ran clean | it never quotes, so nothing is banned | + +The model even volunteered *"(note: no hyphens or spaces)"* — confidently +describing text that was not in the prompt — at 1k context, where retrieval +plainly worked. That was the tell, and it was visible early. + +--- + +## Ablation (needle at 1k, greedy, `--temp 0`) + +| run | output | | +|---|---|---| +| A default (pen 1.15 + ngram + prompt-history) | `VIOLETANVIL62` | wrong | +| B `--repeat-penalty 1.0` | `VIOLETANVIL62` | wrong | +| C `--no-ngram-block` | `VIOLET-ANVIL-62` | ok | +| D pen 1.0 + ngram off | `VIOLET-ANVIL-62` | ok | +| **E prompt-free history, both corrections ON** | `VIOLET-ANVIL-62` | **ok** | + +**E is the load-bearing row.** Keep both corrections enabled and merely scope +the history to generated tokens: correct. The anti-repeat machinery was never +the problem — feeding it the prompt was. B rules out the repetition penalty. + +--- + +## The model was never the problem + +Context-length sweep, needle pinned ~607 tokens in, only the distance between +it and the question varying (524 → 7,487 tokens): + +| ctx | 1000 | 2000 | 3000 | 4000 | 4500 | 5000 | 6000 | 8000 | +|---|---|---|---|---|---|---|---|---| +| result | EXACT | EXACT | budget | EXACT | EXACT | miss | budget | EXACT | + +No cliff, no decay. (`3000`/`6000` were still inside `` at the +120-token cap; `5000` is the single genuine miss, hallucinating +`VIOLET-CHIME`.) + +At `ctx=8000` **before** the fix, the model said the phrase +`"should appear as "VIOLET- ANVIL- 62" (with hyphens)"` — it knew the hyphens +were there and could not emit them contiguously. + +--- + +## The fix + +```rust +// infer.rs — anti-repeat state covers ONLY the model's own output +let mut token_history: Vec = Vec::new(); + +// run.rs — scope both corrections to this turn +let turn_start = conversation_tokens.len() - generated; +let turn = &conversation_tokens[turn_start..]; +``` + +`--no-ngram-block` is retained as a diagnostic flag. + +### Verified + +| | before | after | +|---|---|---| +| needle at 1k, greedy | `VIOLETANVIL62` | **`VIOLET-ANVIL-62`** | +| needle at 8k, greedy | `VIOLETANVIL62, which should appear as "VIOLET- ANVIL- 62"` | **`VIOLET-ANVIL-62`** | + +Full coherence probe at 8.5k, temp 1.0, **no** repetition penalty, 1600 +tokens — the configuration that previously collapsed — now yields 903 words +of coherent prose and reproduces a planted sentence word for word: + +> `Kestrel Protocol was ratified in Reykjavik on 14 March 2019 by exactly +> eleven signatory laboratories` + +--- + +## Severity + +**Lab only.** `apply_ngram_block` has no caller in `hipfire-engine`; the +callers are three `saddle-lab` examples plus a speculative-decode path behind +`HIPFIRE_DFLASH_NGRAM_BLOCK=1`. No shipped model or user-facing path is +affected, and PR #694 is not blocked by this. + +--- + +## Retracted + +An earlier draft of this document reported a long-context degeneracy in +escha-35b and attributed it, in turn, to trellis-codec quality, logit-tail +corruption at long context, DeltaNet recurrent-state precision, and the +capacity limits of a 30-of-40-layer linear-attention architecture. **All of +those are withdrawn.** Every measurement behind them ran through this harness +with the blocker active. Specifically retracted: + +- all needle-recall tables +- the Q8 → FP32 DeltaNet state result (1/3 → 2/3 needles) +- the recommendation to run `--repeat-penalty 1.05 --dn-state fp32` +- the claim that escha's long-range verbatim retrieval is imperfect + +Measurements that remain valid, because they do not depend on quoting the +prompt: the prefill logit sweep (no distributional degradation across 8k; +top-20 mass 0.80 → 0.9998, entropy falling, logit scale flat, and an mq3 +control tracking it closely), and the tokenizer encode→decode round-trip +(14/14 strings exact, including `VIOLET-ANVIL-62`). + +## Method notes worth keeping + +- **Read the output, not its statistics.** A counting-loop collapse scored + *85% unique words* — better than the coherent run it was being compared + against. Every proxy used here (unique-word ratio, top-20 mass, entropy) + ranked the failures wrong at least once. The bug was found by reading text. +- **Sample the whole generation.** Probing only the first 40 tokens showed + long and short context as identical (tail mass 0.00313 vs 0.00303); those + tokens are the structural preamble. Across the full 1200, long context + carried ~3.5× the tail mass. +- **Default sampling is not greedy** (temp 0.3). Single-sample A/Bs across it + are noise; `--temp 0` was added so comparisons mean something. diff --git a/docs/plans/escha-w2-phase1.md b/docs/plans/escha-w2-phase1.md new file mode 100644 index 0000000000..6c6cf98610 --- /dev/null +++ b/docs/plans/escha-w2-phase1.md @@ -0,0 +1,2369 @@ +# Escha-W2 Phase 1 (35B-A3B) Implementation Plan + +> **For agentic workers:** REQUIRED SUB-SKILL: Use superpowers:subagent-driven-development (recommended) or superpowers:executing-plans to implement this plan task-by-task. Steps use checkbox (`- [ ]`) syntax for tracking. + +**Goal:** Serve `EschaLabs/Qwen3.6-35B-A3B-Escha-W2` in hipfire on gfx1151 by repacking its trellis codes verbatim into `.hfq` and decoding them to `Q8_0` resident weights at load, with escha's two H128 activation transforms running at runtime. + +**Architecture:** Three layers, each gated against the layer above it. A pure-Rust CPU reference (`escha_ref`) ported from Escha's Apache-2.0 `ref.py` is the numerical oracle. A converter turns safetensors into `.hfq` with two new quant types whose code streams are byte-identical to the source. At load, one GPU kernel expands those codes to `Q8_0`; two more apply the input and output Hadamards at runtime, and every GEMV downstream is existing hipfire code. + +**Tech Stack:** Rust (workspace crates `hipfire-quantize`, `hipfire-runtime`, `hipfire-dispatch`, `rdna-compute`, `hipfire-arch-qwen35`), HIP/ROCm 7.2.2 targeting gfx1151, `cargo test`. + +**Design spec:** `docs/plans/escha-w2-port-design.md`. Read §1.1–§1.4 before Task 1. + +## Global Constraints + +- Branch `nw_escha_w2`, worktree `~/repos/hipfire-escha`, off `origin/master` @ `8cd15a62b`. Upstream is deliberately unset; push with `git push origin nw_escha_w2`. +- Quant type ids are **`ESCHA2T16 = 42`** and **`ESCHA3T16 = 43`**. The authoritative registry is `crates/hipfire-quantize/src/hfq.rs` — the `#[repr(u8)] enum QuantType` **and** its `from_u8`, which its own doc comment requires be kept in sync. Do not consult the stale partial enum in the `loop/gfx1151` checkout. +- `RS = 0.088388347648` exactly (`1/sqrt(128)`). `kernels/src/gemv_mq4g128.hip:116` already pins `0.0883883476f`. +- Codebook hash constants, exact: multiplier `0xCBAC1FED`, mask `0x8FFF8FFF`, xor `0x3B603B60`. +- **Every `f16(...)` in the escha contract is round-to-nearest-even.** Do NOT use `crate::float16::f32_to_f16` for it — that helper **truncates**, deliberately, to keep existing HFQ bytes stable (see its module doc). Truncating breaks the codebook: it misses published constants at states 3, 6 and 7. Use `escha_ref::f16_rne` (Task 2), which routes through `half::f16::from_f32`. Decoding with `crate::float16::f16_to_f32` is fine — only the encode direction differs. Found by TDD in Task 1; upstream `ref.py` states the contract outright: "numpy f16+f16 rounds RNE, like the GPU". +- **No codebook LUT in any kernel.** 65536 × f16 = 128 KB; gfx1151 has 64 KB LDS (`crates/rdna-compute/src/profiler.rs`, `lds_per_cu: 65536`). Decode inline. +- Required leaves are `escha_code`, `escha_rin`, `escha_rout`. Optional: `escha_s_in`, `escha_s_out`, `escha_config`, `bias`. Unknown `escha_*` leaves are a hard error. See spec §1.4. +- `K` comes from `code.shape[-1] / 16`. Never from `escha_config` (optional) and never from `layer_meta.bits` (self-inconsistent across releases). +- Build: `cargo build --release --workspace --all-targets --locked`. Never run bare `cargo fmt` — it rewrites the workspace and buries the change (`CLAUDE.md:109`). +- hipcc must be on PATH and the kernel cache must be single-toolchain, or you will chase attractor garbage that mimics a codec bug. See `hipfire_kernel_rebuild_gfx1151`. + +## Scope + +**This plan is Phase 1, 35B-A3B only.** Out of scope, each needing its own plan: Phase 2 fused decode+GEMV; the Qwen3.8-27B dense model (larger kernel surface, needs arch-5 bias slots, and cannot ship on the decode-at-load tier). + +## File Structure + +| File | Responsibility | +|---|---| +| `crates/hipfire-quantize/src/escha_ref.rs` (new) | CPU oracle: codebook, tile decode, H128, transforms, MoE block. No GPU, no hipfire deps. | +| `crates/hipfire-quantize/src/lib.rs` (modify) | Add `pub mod escha_ref;` | +| `crates/hipfire-quantize/tests/data/escha/` (new) | Vendored packed golden inputs + digest constants | +| `crates/hipfire-quantize/src/hfq.rs` (modify) | `ESCHA2T16 = 42`, `ESCHA3T16 = 43` in enum + `from_u8` | +| `crates/hipfire-quantize/src/pipeline_escha.rs` (new) | safetensors → `.hfq` converter | +| `crates/hipfire-quantize/src/main.rs` (modify) | `mod pipeline_escha;` + CLI arm | +| `crates/rdna-compute/src/dispatch.rs` (modify) | `DType::Escha2T16`, `DType::Escha3T16` | +| `crates/hipfire-dispatch/src/types.rs` (modify) | `RotationPlan::EschaH128` + `dtype_rotation_plan` arms | +| `kernels/src/escha_decode_tiles.hip` (new) | One-shot tile → `Q8_0` expansion | +| `kernels/src/escha_h128.hip` (new) | Input and output H128 transforms | +| `crates/rdna-compute/src/kernels.rs` (modify) | `include_str!` consts for both kernels | +| `crates/hipfire-arch-qwen35/src/` (modify) | Load escha experts, call the transforms | +| `registry/v1.json` (modify) | `qwen3.6:35b-a3b-escha` | + +--- + +### Task 1: `escha_ref` codebook and tile decode + +The trellis decode is the heart of the port. Everything else is gated against it, so it is built first and proven against Escha's own golden vectors. + +**Files:** +- Create: `crates/hipfire-quantize/src/escha_ref.rs` +- Create: `crates/hipfire-quantize/tests/data/escha/fetch-goldens.sh` +- Modify: `crates/hipfire-quantize/src/lib.rs:3` (add module) +- Test: inline `#[cfg(test)]` in `escha_ref.rs` (this crate has no `tests/` dir; all tests are inline) + +**Interfaces:** +- Consumes: `crate::float16::f16_to_f32`, `half::f16::from_f32` (RNE encode — see Global Constraints) +- Produces: `pub const RS: f32`, `pub fn cba_decode(state: u16) -> u16`, `pub fn decode8_k2(words: &[u32; 16], lane: usize) -> [u16; 8]`, `pub fn decode8_k3(words: &[u32; 24], lane: usize) -> [u16; 8]`, `pub fn lane_positions(lane: usize) -> [(usize, usize); 8]`, `pub fn reconstruct(code: &[i16], in_features: usize, out_features: usize, k: usize) -> Vec` (returns f16 **bits**, row-major `[in_features, out_features]`) + +- [ ] **Step 1: Vendor the packed golden inputs** + +Only the packed inputs are committed (0.9 MB). The expected outputs are 6.3 MB and are asserted by SHA-256 digest instead, so the repo stays light and the gate stays exact. + +```bash +mkdir -p crates/hipfire-quantize/tests/data/escha +cd crates/hipfire-quantize/tests/data/escha +B=https://raw.githubusercontent.com/EschaLabs/escha-mlx/HEAD/tests/data/codec +curl -sL "$B/packed_gu_e0_k2.i16" -o packed_gu_e0_k2.i16 +curl -sL "$B/packed_down_e0_k3.i16" -o packed_down_e0_k3.i16 +sha256sum packed_gu_e0_k2.i16 packed_down_e0_k3.i16 +``` + +Expected output, exactly: + +``` +c164583731eed50d52ae3bfcc6a58a72b50329f8d4c40e5d62f940d67991ec1b packed_gu_e0_k2.i16 +aa2f2dd1f165d03ae868ce14913c05313876fa4d6c0e52d2ee16c60ff22eb062 packed_down_e0_k3.i16 +``` + +If either digest differs, stop — upstream changed the fixtures and the constants in Step 3 are stale. + +- [ ] **Step 2: Write the fetch script for the full goldens** + +Create `crates/hipfire-quantize/tests/data/escha/fetch-goldens.sh`: + +```bash +#!/usr/bin/env bash +# Fetch the full escha-mlx golden vectors (6.3 MB of expected outputs, not +# committed). Only needed to regenerate the digests in escha_ref.rs; the +# committed packed inputs plus those digests are a complete gate. +set -euo pipefail +cd "$(dirname "$0")" +B=https://raw.githubusercontent.com/EschaLabs/escha-mlx/HEAD/tests/data +for f in codec/packed_gu_e0_k2.i16 codec/expected_gu_e0_k2.f16 \ + codec/packed_down_e0_k3.i16 codec/expected_down_e0_k3.f16 \ + qwen3_5_moe/moeblk_x.f16 qwen3_5_moe/moeblk_out.f16 \ + qwen3_5_moe/moeblk_ids.i64 qwen3_5_moe/moeblk_scores.f32; do + curl -sL --fail "$B/$f" -o "$(basename "$f")" +done +sha256sum ./*.f16 ./*.i16 ./*.i64 ./*.f32 +``` + +```bash +chmod +x crates/hipfire-quantize/tests/data/escha/fetch-goldens.sh +``` + +- [ ] **Step 3: Write the failing tests** + +Create `crates/hipfire-quantize/src/escha_ref.rs` containing **only** this test module for now: + +```rust +//! Portable CPU reference for the escha codec — the numerical oracle for +//! every GPU kernel in this port. +//! +//! Ported from `escha_mlx/ref.py` (EschaLabs/escha-mlx, Apache-2.0), which +//! declares itself "the semantic contract for every Metal kernel in this +//! package". Rounding points are deliberate; do not "simplify" them. + +#[cfg(test)] +mod tests { + use super::*; + use std::path::PathBuf; + + fn data(name: &str) -> PathBuf { + PathBuf::from(env!("CARGO_MANIFEST_DIR")).join("tests/data/escha").join(name) + } + + fn sha256_hex(bytes: &[u8]) -> String { + use sha2::{Digest, Sha256}; + let mut h = Sha256::new(); + h.update(bytes); + h.finalize().iter().map(|b| format!("{b:02x}")).collect() + } + + /// The codebook is a pure function of the 16-bit state. These eight values + /// were computed from the published constants and pin the hash, the + /// masking, and the fp16 RNE add all at once. + #[test] + fn cba_decode_matches_published_constants() { + let want: [u16; 8] = + [0x3f60, 0x304e, 0xba13, 0x3ab8, 0x3952, 0xb75f, 0xbea4, 0xbc71]; + for (state, &bits) in want.iter().enumerate() { + assert_eq!(cba_decode(state as u16), bits, "state {state}"); + } + } + + #[test] + fn reconstruct_k2_matches_golden() { + let raw = std::fs::read(data("packed_gu_e0_k2.i16")).unwrap(); + let code: Vec = + raw.chunks_exact(2).map(|c| i16::from_le_bytes([c[0], c[1]])).collect(); + assert_eq!(code.len(), 262144); + let out = reconstruct(&code, 2048, 1024, 2); + assert_eq!(out.len(), 2048 * 1024); + let bytes: Vec = out.iter().flat_map(|v| v.to_le_bytes()).collect(); + assert_eq!( + sha256_hex(&bytes), + "51ddde9a07613aafcc9f5db79702349d19e18357d9fb910f48b38eeea028dcab", + "decoded K=2 tensor does not match expected_gu_e0_k2.f16" + ); + } + + #[test] + fn reconstruct_k3_matches_golden() { + let raw = std::fs::read(data("packed_down_e0_k3.i16")).unwrap(); + let code: Vec = + raw.chunks_exact(2).map(|c| i16::from_le_bytes([c[0], c[1]])).collect(); + assert_eq!(code.len(), 196608); + let out = reconstruct(&code, 512, 2048, 3); + assert_eq!(out.len(), 512 * 2048); + let bytes: Vec = out.iter().flat_map(|v| v.to_le_bytes()).collect(); + assert_eq!( + sha256_hex(&bytes), + "51c99817d00282f4aa9d618140eb4503083e1238cb59dd71a670aa4c320f7438", + "decoded K=3 tensor does not match expected_down_e0_k3.f16" + ); + } + + /// Every one of the 256 tile slots must be written exactly once. A + /// permutation bug that drops and duplicates slots still produces a + /// full-rank, plausible-looking matrix, so check the permutation directly. + #[test] + fn lane_positions_is_a_permutation_of_the_tile() { + let mut seen = [0u8; 256]; + for lane in 0..32 { + for (r, c) in lane_positions(lane) { + assert!(r < 16 && c < 16, "lane {lane} -> ({r},{c})"); + seen[r * 16 + c] += 1; + } + } + assert!(seen.iter().all(|&n| n == 1), "lane_positions is not a bijection"); + } +} +``` + +Add to `crates/hipfire-quantize/src/lib.rs` after line 3 (`pub mod float16;`): + +```rust +pub mod escha_ref; +``` + +Add the test-only digest dependency to `crates/hipfire-quantize/Cargo.toml` under `[dev-dependencies]` (create the section if absent): + +```toml +[dev-dependencies] +sha2 = "0.10" +``` + +- [ ] **Step 4: Run the tests to verify they fail** + +```bash +cargo test -p hipfire-quantize --lib escha_ref 2>&1 | tail -20 +``` + +Expected: compile error, `cannot find function 'cba_decode' in this scope` (and the same for `reconstruct`, `lane_positions`). A compile failure is the correct "red" here — the functions do not exist yet. + +- [ ] **Step 5: Implement the codec** + +Insert above the `#[cfg(test)] mod tests` block in `escha_ref.rs`: + +```rust +use crate::float16::f16_to_f32; + +/// 1/sqrt(128) — the exact f32 constant the format pins. +pub const RS: f32 = 0.088388347648; + +/// Round f32 to fp16 bits, round-to-nearest-even. +/// +/// Every `f16(...)` in the escha contract is RNE. `crate::float16::f32_to_f16` +/// TRUNCATES — deliberately, to keep existing HFQ bytes stable — and using it +/// here silently corrupts the codec: it misses the published cbA constants at +/// states 3, 6 and 7. Do not "simplify" this back to the crate helper. +#[inline] +pub fn f16_rne(v: f32) -> u16 { + half::f16::from_f32(v).to_bits() +} + +/// Decode one 16-bit trellis state to fp16 **bits** via the cbA codebook. +/// +/// `decode(x) = f16_lo(r) + f16_hi(r)` with fp16 RNE addition, where +/// `r = ((x * 0xCBAC1FED) & 0x8FFF8FFF) ^ 0x3B603B60` in 32-bit arithmetic. +/// +/// Adding in f32 and rounding once is exactly an fp16 RNE add: the exact sum +/// of two fp16 values is always representable in f32, so the single rounding +/// here is the correctly-rounded fp16 result. +/// +/// The round MUST be `half::f16::from_f32`, not `crate::float16::f32_to_f16` +/// — the latter truncates by design and misses states 3, 6 and 7. +/// +/// There are 65536 reachable values, so a lookup table would be 128 KB and +/// will not fit gfx1151's 64 KB LDS. This is five integer/FP ops and no +/// memory traffic — keep it that way in the kernels. +#[inline] +pub fn cba_decode(state: u16) -> u16 { + let r = ((state as u32).wrapping_mul(0xCBAC_1FED) & 0x8FFF_8FFF) ^ 0x3B60_3B60; + let lo = f16_to_f32((r & 0xFFFF) as u16); + let hi = f16_to_f32((r >> 16) as u16); + half::f16::from_f32(lo + hi).to_bits() +} + +/// The 8 states lane `lane` owns, K=2. `words` is the tile's 16 u32. +/// +/// DELIBERATE DUPLICATION: `kernels/src/escha_decode_tiles.hip` implements +/// this same lane maths independently. That is the G2 gate — the GPU decode +/// is asserted bit-exact against this one. Generating either from the other, +/// or sharing a source, would make G2 circular: both paths could be wrong in +/// exactly the same way and still agree. Two independent implementations of +/// a published spec is the point. Do not deduplicate. +pub fn decode8_k2(words: &[u32; 16], lane: usize) -> [u16; 8] { + let t_off = lane * 8; + let i1 = t_off >> 4; + let i0 = (i1 + 15) & 15; + let merged = ((words[i0] as u64) << 32) | words[i1] as u64; + let shift = ((!t_off) & 8) << 1; // 16 for even lanes, 0 for odd + let w = ((merged >> shift) & 0xFFFF_FFFF) as u32; + let mut out = [0u16; 8]; + for (j, o) in out.iter_mut().enumerate() { + *o = (w >> (2 * (7 - j))) as u16; + } + out +} + +/// The 8 states lane `lane` owns, K=3. `words` is the tile's 24 u32. +/// +/// Structurally different from K=2 — 24 words, a computed bit offset, and a +/// modular wrap. Do not attempt to unify the two. +pub fn decode8_k3(words: &[u32; 24], lane: usize) -> [u16; 8] { + const BITS: usize = 3; + let t_off = lane * 8; + let b1 = (t_off + 257) * BITS; + let b0 = b1 - 16; + let b2 = b1 + BITS * 7; + let i0 = b0 >> 5; + let i2 = (b2 - 1) >> 5; + let s2 = ((i2 + 1) << 5) - b2; + let merged = ((words[i0 % 24] as u64) << 32) | words[i2 % 24] as u64; + let w7 = (merged >> s2) & 0xFFFF_FFFF; + let w3 = (merged >> (s2 + BITS * 4)) & 0xFFFF_FFFF; + [ + (w3 >> 9) as u16, + (w3 >> 6) as u16, + (w3 >> 3) as u16, + w3 as u16, + (w7 >> 9) as u16, + (w7 >> 6) as u16, + (w7 >> 3) as u16, + w7 as u16, + ] +} + +/// `(row, col)` inside the 16x16 tile for each of the lane's 8 values. +/// +/// This permutation is the single easiest thing to get subtly wrong: a wrong +/// shuffle still yields a full-rank, plausible weight matrix. It is gated +/// directly on golden vectors, never on end-to-end coherence. +pub fn lane_positions(lane: usize) -> [(usize, usize); 8] { + let l0 = lane & !4; + let c_off = (lane >> 2) & 1; + let mut out = [(0usize, 0usize); 8]; + for (j, o) in out.iter_mut().enumerate() { + let fi = j >> 1; + let row = (lane & 3) * 2 + (j & 1) + (fi & 1) * 8; + let col = 2 * ((l0 >> 3) + if j >= 4 { 4 } else { 0 }) + c_off; + *o = (row, col); + } + out +} + +/// Decode one packed tile (`16*K` i16) to a 16x16 fp16-bit tile, row-major. +pub fn decode_tile(tile: &[i16], k: usize) -> [u16; 256] { + debug_assert_eq!(tile.len(), 16 * k); + let mut words = [0u32; 24]; + for (i, w) in words.iter_mut().enumerate().take(8 * k) { + *w = (tile[2 * i] as u16 as u32) | ((tile[2 * i + 1] as u16 as u32) << 16); + } + let mut out = [0u16; 256]; + for lane in 0..32 { + let states = match k { + 2 => { + let mut w16 = [0u32; 16]; + w16.copy_from_slice(&words[..16]); + decode8_k2(&w16, lane) + } + 3 => decode8_k3(&words, lane), + _ => panic!("unsupported escha K={k}"), + }; + for (j, (r, c)) in lane_positions(lane).into_iter().enumerate() { + out[r * 16 + c] = cba_decode(states[j]); + } + } + out +} + +/// Packed `(in/16, out/16, 16K)` i16 -> `(in, out)` fp16 bits, row-major. +pub fn reconstruct(code: &[i16], in_features: usize, out_features: usize, k: usize) -> Vec { + let (tk, tn) = (in_features / 16, out_features / 16); + assert_eq!(code.len(), tk * tn * 16 * k, "escha code length mismatch"); + let mut out = vec![0u16; in_features * out_features]; + for kt in 0..tk { + for nt in 0..tn { + let base = (kt * tn + nt) * 16 * k; + let tile = decode_tile(&code[base..base + 16 * k], k); + for r in 0..16 { + let dst = (kt * 16 + r) * out_features + nt * 16; + out[dst..dst + 16].copy_from_slice(&tile[r * 16..r * 16 + 16]); + } + } + } + out +} +``` + +- [ ] **Step 6: Run the tests to verify they pass** + +```bash +cargo test -p hipfire-quantize --lib escha_ref 2>&1 | tail -15 +``` + +Expected: `test result: ok. 4 passed; 0 failed`. + +If `reconstruct_k2_matches_golden` fails while `cba_decode_matches_published_constants` passes, the bug is in `lane_positions` or the word packing, not the codebook — check `lane_positions_is_a_permutation_of_the_tile` first. + +- [ ] **Step 7: Commit** + +```bash +git add crates/hipfire-quantize/src/escha_ref.rs \ + crates/hipfire-quantize/src/lib.rs \ + crates/hipfire-quantize/Cargo.toml \ + crates/hipfire-quantize/tests/data/escha/ +git commit -m "feat(escha): trellis codec CPU reference, gated on golden vectors" +``` + +--- + +### Task 2: `escha_ref` H128 and the transform pair + +**Files:** +- Modify: `crates/hipfire-quantize/src/escha_ref.rs` + +**Interfaces:** +- Consumes: `RS` from Task 1 +- Produces: `pub fn f16_rne(v: f32) -> u16`, `pub fn h128_inplace(x: &mut [f32])`, `pub fn input_transform(x: &[f32], rin: &[f32]) -> Vec`, `pub fn output_transform(mid: &[f32], rout: &[f32]) -> Vec`, `pub fn fold_scales(rin: &[u16], rout: &[u16], s_in: Option<&[f32]>, s_out: Option<&[f32]>) -> (Vec, Vec)` +- Note: Task 1 inlined `half::f16::from_f32(..).to_bits()` in `cba_decode`. Extract that into `f16_rne` here and have `cba_decode` call it, so there is exactly one RNE encode site in the module. + +- [ ] **Step 1: Write the failing tests** + +Add inside the existing `mod tests` block in `escha_ref.rs`: + +```rust + /// H128 is its own inverse up to a factor of 128. That is necessary but + /// NOT sufficient: a wrong butterfly order is also self-inverse and would + /// pass this alone. The Hadamard-of-a-basis-vector check below pins the + /// actual transform. + #[test] + fn h128_roundtrip_scales_by_128() { + let mut x: Vec = (0..256).map(|i| (i as f32 * 0.37).sin()).collect(); + let orig = x.clone(); + h128_inplace(&mut x); + h128_inplace(&mut x); + for (a, b) in x.iter().zip(orig.iter()) { + assert!((a - b * 128.0).abs() < 1e-2, "{a} vs {}", b * 128.0); + } + } + + /// H128 applied to e_0 must give all ones (Sylvester, unnormalised). + /// Applied to e_1 it must give the alternating +1/-1 pattern of row 1. + #[test] + fn h128_matches_sylvester_order() { + let mut e0 = vec![0.0f32; 128]; + e0[0] = 1.0; + h128_inplace(&mut e0); + assert!(e0.iter().all(|&v| v == 1.0), "row 0 must be all ones"); + + let mut e1 = vec![0.0f32; 128]; + e1[1] = 1.0; + h128_inplace(&mut e1); + for (i, &v) in e1.iter().enumerate() { + let want = if i % 2 == 0 { 1.0 } else { -1.0 }; + assert_eq!(v, want, "index {i}"); + } + } + + /// Blocks are independent: H128 must never mix across a 128 boundary. + #[test] + fn h128_does_not_mix_across_blocks() { + let mut x = vec![0.0f32; 256]; + x[0] = 1.0; + h128_inplace(&mut x); + assert!(x[..128].iter().all(|&v| v == 1.0)); + assert!(x[128..].iter().all(|&v| v == 0.0), "second block was contaminated"); + } + + /// MoE exports ship all-ones s_in/s_out; dense exports ship real values; + /// and an export without the end-to-end stage ships neither. All three + /// must go through one code path. + #[test] + fn fold_scales_handles_absent_scales() { + let rin = [f16_rne(2.0), f16_rne(-3.0)]; + let rout = [f16_rne(0.5)]; + let (a, b) = fold_scales(&rin, &rout, None, None); + assert_eq!(a, vec![2.0, -3.0]); + assert_eq!(b, vec![0.5]); + let (c, d) = fold_scales(&rin, &rout, Some(&[3.0, 2.0]), Some(&[4.0])); + assert_eq!(c, vec![6.0, -6.0]); + assert_eq!(d, vec![2.0]); + } +``` + +- [ ] **Step 2: Run the tests to verify they fail** + +```bash +cargo test -p hipfire-quantize --lib escha_ref 2>&1 | tail -15 +``` + +Expected: compile error, `cannot find function 'h128_inplace' in this scope`. + +- [ ] **Step 3: Implement the transforms** + +Append to the non-test portion of `escha_ref.rs`: + +```rust +/// Unnormalised 128-point Walsh-Hadamard (Sylvester / natural order), applied +/// independently to each contiguous 128-element block of `x`. +/// +/// `x.len()` must be a multiple of 128. Every dimension in both checkpoints +/// satisfies this (512, 1024, 2048, 5120, 6144, 10240, 17408 are all +/// multiples of 128), including the gate|up split point at 512 — so a block +/// never straddles the gate/up boundary. +pub fn h128_inplace(x: &mut [f32]) { + assert_eq!(x.len() % 128, 0, "H128 needs a multiple of 128, got {}", x.len()); + for block in x.chunks_exact_mut(128) { + let mut h = 1; + while h < 128 { + let mut i = 0; + while i < 128 { + for j in i..i + h { + let (a, b) = (block[j], block[j + h]); + block[j] = a + b; + block[j + h] = a - b; + } + i += 2 * h; + } + h *= 2; + } + } +} + +/// `xh = f16( H128(x * rin) * RS )`. Returns fp16 bits. +pub fn input_transform(x: &[f32], rin: &[f32]) -> Vec { + let ic = rin.len(); + assert_eq!(x.len() % ic, 0); + let mut buf: Vec = x.iter().zip(rin.iter().cycle()).map(|(a, b)| a * b).collect(); + for row in buf.chunks_exact_mut(ic) { + h128_inplace(row); + } + buf.iter().map(|v| f16_rne(v * RS)).collect() +} + +/// `y = f16( H128(mid) * RS * rout )`. Returns fp16 bits. +pub fn output_transform(mid: &[f32], rout: &[f32]) -> Vec { + let oc = rout.len(); + assert_eq!(mid.len() % oc, 0); + let mut buf = mid.to_vec(); + for row in buf.chunks_exact_mut(oc) { + h128_inplace(row); + } + buf.iter() + .zip(rout.iter().cycle()) + .map(|(v, s)| f16_rne(v * RS * s)) + .collect() +} + +/// Fold the optional end-to-end scales into the transform vectors. +/// +/// `s_in` multiplies the activation at exactly the point `rin` does, and +/// `s_out` at exactly the point `rout` does, so the pair collapses with no new +/// kernel and no new tensor. Folding keeps both products in f32 and rounds +/// once — one rounding point FEWER than applying the scales separately. +/// `None` returns that vector unchanged (as f32), which is the path MoE +/// exports and end-to-end-free exports both take. +pub fn fold_scales( + rin: &[u16], + rout: &[u16], + s_in: Option<&[f32]>, + s_out: Option<&[f32]>, +) -> (Vec, Vec) { + let mut ri: Vec = rin.iter().map(|&b| f16_to_f32(b)).collect(); + let mut ro: Vec = rout.iter().map(|&b| f16_to_f32(b)).collect(); + if let Some(s) = s_in { + assert_eq!(s.len(), ri.len()); + for (a, b) in ri.iter_mut().zip(s) { + *a *= b; + } + } + if let Some(s) = s_out { + assert_eq!(s.len(), ro.len()); + for (a, b) in ro.iter_mut().zip(s) { + *a *= b; + } + } + (ri, ro) +} +``` + +- [ ] **Step 4: Run the tests to verify they pass** + +```bash +cargo test -p hipfire-quantize --lib escha_ref 2>&1 | tail -15 +``` + +Expected: `test result: ok. 8 passed; 0 failed`. + +- [ ] **Step 5: Commit** + +```bash +git add crates/hipfire-quantize/src/escha_ref.rs +git commit -m "feat(escha): H128 transforms and scale folding in the CPU reference" +``` + +--- + +### Task 3: `escha_ref` expert linear, SwiGLU and w8a16 + +Completes the oracle. `expert_linear` is what Task 7's GPU decode is gated against, and `swiglu`/`w8a16` pin the two rounding points the Task 10 wiring must reproduce. + +No `moe_block` port is needed: Task 10's G4 gate compares hipfire directly against Escha's shipped `moeblk_out.f16` golden, so a Rust reimplementation of the block would add a second thing to keep in sync without gating anything. + +**Files:** +- Modify: `crates/hipfire-quantize/src/escha_ref.rs` + +**Interfaces:** +- Consumes: `reconstruct`, `input_transform`, `output_transform` from Tasks 1–2 +- Produces: `pub fn expert_linear(x: &[f32], w_bits: &[u16], rin: &[f32], rout: &[f32]) -> Vec`, `pub fn swiglu(gate_up_bits: &[u16], inter: usize) -> Vec`, `pub fn w8a16(x: &[f32], w8: &[i8], scale: &[u16], oc: usize, ic: usize) -> Vec` + +- [ ] **Step 1: Write the failing test** + +Add inside `mod tests`: + +```rust + /// A pruned output channel must be EXACTLY zero, not approximately. + /// gate_up.rout carries a per-expert prune mask (design §1.2): on the + /// shipped layer-0 expert 0, 560 of 1024 channels are hard zeros. A kernel + /// that "optimises away" the zero multiply must preserve exact zero. + #[test] + fn zero_rout_gives_exactly_zero_output() { + let ic = 128; + let oc = 128; + let w: Vec = (0..ic * oc).map(|i| f16_rne((i % 7) as f32 - 3.0)).collect(); + let rin = vec![1.0f32; ic]; + let mut rout = vec![1.0f32; oc]; + rout[3] = 0.0; + rout[57] = 0.0; + let x: Vec = (0..ic).map(|i| (i as f32 * 0.11).cos()).collect(); + let y = expert_linear(&x, &w, &rin, &rout); + assert_eq!(y.len(), oc); + assert_eq!(f16_to_f32(y[3]), 0.0, "pruned channel 3 must be exactly zero"); + assert_eq!(f16_to_f32(y[57]), 0.0, "pruned channel 57 must be exactly zero"); + assert!(y.iter().enumerate().any(|(i, &v)| i != 3 && i != 57 && f16_to_f32(v) != 0.0)); + } + + /// SwiGLU splits the f16-ROUNDED merged output at the halfway point, gate + /// first. Rounding before the split is part of the contract. + #[test] + fn swiglu_uses_gate_first_half() { + let inter = 2; + // gate = [0, 0], up = [5, 7]; silu(0) == 0 so both outputs are zero. + let gu: Vec = [0.0, 0.0, 5.0, 7.0].iter().map(|&v| f16_rne(v)).collect(); + let h = swiglu(&gu, inter); + assert_eq!(h.len(), inter); + assert_eq!(f16_to_f32(h[0]), 0.0); + assert_eq!(f16_to_f32(h[1]), 0.0); + // gate = [large, large] -> silu(x) ~ x, so out ~ gate*up. + let gu2: Vec = [10.0, 10.0, 2.0, 3.0].iter().map(|&v| f16_rne(v)).collect(); + let h2 = swiglu(&gu2, inter); + assert!((f16_to_f32(h2[0]) - 20.0).abs() < 0.1, "{}", f16_to_f32(h2[0])); + assert!((f16_to_f32(h2[1]) - 30.0).abs() < 0.2, "{}", f16_to_f32(h2[1])); + } + + /// Escha's int8 is per-output-ROW: y = f16(x @ f16(w8*scale)^T). + #[test] + fn w8a16_applies_per_row_scale() { + let (ic, oc) = (4, 2); + let w8: Vec = vec![1, 2, 3, 4, 5, 6, 7, 8]; + let scale: Vec = vec![f16_rne(0.5), f16_rne(2.0)]; + let x = vec![1.0f32, 1.0, 1.0, 1.0]; + let y = w8a16(&x, &w8, &scale, oc, ic); + assert_eq!(f16_to_f32(y[0]), 5.0); // (1+2+3+4)*0.5 + assert_eq!(f16_to_f32(y[1]), 52.0); // (5+6+7+8)*2.0 + } +``` + +- [ ] **Step 2: Run the tests to verify they fail** + +```bash +cargo test -p hipfire-quantize --lib escha_ref 2>&1 | tail -15 +``` + +Expected: compile error, `cannot find function 'expert_linear' in this scope`. + +- [ ] **Step 3: Implement** + +Append to the non-test portion of `escha_ref.rs`: + +```rust +/// Full single-expert linear for one token: `x [ic] -> [oc]` fp16 bits. +/// +/// `w_bits` is the decoded bare weight, row-major `[ic, oc]` — decode it once +/// with `reconstruct` and reuse it. `ref.moe_block` in the Python original +/// re-decodes per (token, slot), which is 128 full tile decodes for an +/// 8-token fixture; do not reproduce that. +pub fn expert_linear(x: &[f32], w_bits: &[u16], rin: &[f32], rout: &[f32]) -> Vec { + let (ic, oc) = (rin.len(), rout.len()); + assert_eq!(x.len(), ic); + assert_eq!(w_bits.len(), ic * oc); + let xh = input_transform(x, rin); + let mut mid = vec![0.0f32; oc]; + for i in 0..ic { + // Unconditional MAC — no zero-activation skip. This is the oracle, so + // it must be a faithful matmul: 0.0 * NaN is NaN, not 0, and skipping + // would MASK a corrupted decode instead of surfacing it. + let a = f16_to_f32(xh[i]); + let row = &w_bits[i * oc..(i + 1) * oc]; + for (m, &wb) in mid.iter_mut().zip(row) { + *m += a * f16_to_f32(wb); + } + } + output_transform(&mid, rout) +} + +/// `silu(g) * u` on the fp16-rounded merged output; gate is the first half. +pub fn swiglu(gate_up_bits: &[u16], inter: usize) -> Vec { + assert_eq!(gate_up_bits.len(), 2 * inter); + let mut out = Vec::with_capacity(inter); + for i in 0..inter { + let g = f16_to_f32(gate_up_bits[i]); + let s = f16_to_f32(f16_rne(g / (1.0 + (-g).exp()))); + out.push(f16_rne(s * f16_to_f32(gate_up_bits[inter + i]))); + } + out +} + +/// `y = f16( x @ f16(w8 * scale)^T )`. `w8` is `[oc, ic]`, `scale` is `[oc]` +/// fp16 bits — Escha's int8 is per-output-row, not per-block. +pub fn w8a16(x: &[f32], w8: &[i8], scale: &[u16], oc: usize, ic: usize) -> Vec { + assert_eq!(w8.len(), oc * ic); + assert_eq!(scale.len(), oc); + assert_eq!(x.len(), ic); + let mut out = Vec::with_capacity(oc); + for o in 0..oc { + let s = f16_to_f32(scale[o]); + let mut acc = 0.0f32; + for i in 0..ic { + acc += x[i] * f16_to_f32(f16_rne(w8[o * ic + i] as f32 * s)); + } + out.push(f16_rne(acc)); + } + out +} +``` + +- [ ] **Step 4: Run the tests to verify they pass** + +```bash +cargo test -p hipfire-quantize --lib escha_ref 2>&1 | tail -15 +``` + +Expected: `test result: ok. 11 passed; 0 failed`. + +- [ ] **Step 5: Commit** + +```bash +git add crates/hipfire-quantize/src/escha_ref.rs +git commit -m "feat(escha): expert linear, SwiGLU and w8a16 in the CPU reference" +``` + +--- + +### Task 4: Register the quant types and the rotation plan + +Two ids, one rotation plan, and the guard that stops them silently falling through to an unrotated kernel. That fallthrough is the highest-severity failure mode in the port: it produces coherent-looking text rather than a crash. + +**Files:** +- Modify: `crates/hipfire-quantize/src/hfq.rs` (enum + `from_u8`) +- Modify: `crates/rdna-compute/src/dispatch.rs:285-310` (`DType`) +- Modify: `crates/hipfire-dispatch/src/types.rs:107-130` (`RotationPlan`, `dtype_rotation_plan`) +- Test: `crates/hipfire-dispatch-tests/src/dtype.rs` + +**Interfaces:** +- Produces: `hfq::QuantType::{ESCHA2T16, ESCHA3T16}` (bytes 42, 43), `DType::{Escha2T16, Escha3T16}`, `RotationPlan::EschaH128` + +- [ ] **Step 1: Write the failing tests** + +Add to `crates/hipfire-dispatch-tests/src/dtype.rs`: + +```rust +#[test] +fn escha_types_use_the_escha_rotation_plan() { + assert_eq!(dtype_rotation_plan(DType::Escha2T16), RotationPlan::EschaH128); + assert_eq!(dtype_rotation_plan(DType::Escha3T16), RotationPlan::EschaH128); +} + +/// Escha weights are stored in the rotated domain. Reaching a Plain GEMV +/// without the H128 pair does not crash — it produces coherent-looking +/// garbage. Both types must therefore refuse to resolve to Plain, exactly as +/// MQ4G128 does (see coverage_tests.rs). +#[test] +fn escha_types_never_resolve_to_plain() { + for dt in [DType::Escha2T16, DType::Escha3T16] { + assert!( + KernelKey::for_gemv(dt, GemvVariant::Plain, false).is_err(), + "{dt:?} must not have a Plain GEMV arm — that would skip the H128 pair" + ); + } +} +``` + +Add to `crates/hipfire-quantize/src/hfq.rs`, inside its existing `#[cfg(test)] mod tests`: + +```rust + /// from_u8 and the enum discriminants must agree — the doc comment on + /// from_u8 makes this a contract, and a drifted pair silently mislabels + /// every tensor written after it. + #[test] + fn escha_quant_types_round_trip() { + assert_eq!(QuantType::from_u8(42), Some(QuantType::ESCHA2T16)); + assert_eq!(QuantType::from_u8(43), Some(QuantType::ESCHA3T16)); + assert_eq!(QuantType::ESCHA2T16 as u8, 42); + assert_eq!(QuantType::ESCHA3T16 as u8, 43); + } +``` + +- [ ] **Step 2: Run the tests to verify they fail** + +```bash +cargo test -p hipfire-quantize --lib hfq::tests::escha 2>&1 | tail -10 +cargo test -p hipfire-dispatch-tests escha 2>&1 | tail -10 +``` + +Expected: compile errors — `no variant named 'ESCHA2T16'`, `no variant or associated item named 'Escha2T16'`. + +- [ ] **Step 3: Add the quant types** + +In `crates/hipfire-quantize/src/hfq.rs`, add to the `#[repr(u8)] enum QuantType` after `MQ2G256LloydU`: + +```rust + /// Escha-W2 trellis, K=2, 16x16 tile, cbA hash codebook (2.00 bpw). + /// Codes are stored verbatim from the source safetensors. + ESCHA2T16 = 42, + /// Escha-W2 trellis, K=3, 16x16 tile, cbA hash codebook (3.00 bpw). + ESCHA3T16 = 43, +``` + +And in `from_u8`, before the `_ => None` arm: + +```rust + 42 => Some(Self::ESCHA2T16), + 43 => Some(Self::ESCHA3T16), +``` + +In `crates/rdna-compute/src/dispatch.rs`, add to `enum DType` after `MQ4G256V2`: + +```rust + Escha2T16, + Escha3T16, +``` + +In `crates/hipfire-dispatch/src/types.rs`, add to `enum RotationPlan` after `Givens`: + +```rust + /// Escha-W2: unnormalised 128-point Walsh-Hadamard on BOTH sides, + /// RS = 1/sqrt(128), signs folded into rin/rout rather than seeded. + EschaH128, +``` + +And add an arm to `dtype_rotation_plan`, before the `_ => RotationPlan::None` catch-all: + +```rust + Escha2T16 | Escha3T16 => RotationPlan::EschaH128, +``` + +- [ ] **Step 4: Run the tests to verify they pass** + +```bash +cargo test -p hipfire-quantize --lib hfq::tests::escha 2>&1 | tail -10 +cargo test -p hipfire-dispatch-tests escha 2>&1 | tail -10 +``` + +Expected: both `test result: ok`. + +If `escha_types_never_resolve_to_plain` fails, a `_ =>` catch-all in `KernelKey::for_gemv` is swallowing the new types. Find it and make the escha arms explicit errors — do not leave the catch-all to handle them. + +- [ ] **Step 5: Verify the whole workspace still builds** + +```bash +cargo build --release --workspace --all-targets --locked 2>&1 | tail -20 +``` + +Expected: no errors. Non-exhaustive-match errors elsewhere are the point — every site that matches on `DType` or `RotationPlan` now has to say what it does with escha. Add explicit `Escha2T16 | Escha3T16 => Err(...)` arms rather than folding them into existing catch-alls. + +- [ ] **Step 6: Commit** + +```bash +git add crates/hipfire-quantize/src/hfq.rs crates/rdna-compute/src/dispatch.rs \ + crates/hipfire-dispatch/src/types.rs crates/hipfire-dispatch-tests/src/dtype.rs +git commit -m "feat(escha): register ESCHA2T16/ESCHA3T16 and the EschaH128 rotation plan" +``` + +--- + +### Task 5: Converter — safetensors to `.hfq` + +**Files:** +- Create: `crates/hipfire-quantize/src/pipeline_escha.rs` +- Modify: `crates/hipfire-quantize/src/main.rs` (add `mod pipeline_escha;` beside the other `pipeline_*` modules, and a CLI arm) + +**Interfaces:** +- Consumes: `crate::hfq::{HfqTensor, QuantType, write_hfq}`, `crate::safetensors_file::SafetensorsFile`, `crate::escha_ref::fold_scales` +- Produces: `pub(crate) fn convert_escha(src_dir: &Path, out: &Path) -> Result<(), String>`, `pub(crate) fn classify_leaf(name: &str) -> Leaf`, `pub(crate) enum Leaf { Code, Rin, Rout, SIn, SOut, Config, Bias, Int8, Int8Scale, Passthrough, UnknownEscha }` + +- [ ] **Step 1: Write the failing tests** + +Create `crates/hipfire-quantize/src/pipeline_escha.rs` with only this test module: + +```rust +#[cfg(test)] +mod tests { + use super::*; + use crate::escha_ref::f16_rne; + + #[test] + fn k_comes_from_the_code_shape_not_metadata() { + // gate_up: [E, in/16, out/16, 16K] with K=2 -> last dim 32 + assert_eq!(k_from_code_shape(&[256, 128, 64, 32]), Ok(2)); + // down: K=3 -> last dim 48 + assert_eq!(k_from_code_shape(&[256, 32, 128, 48]), Ok(3)); + // dense exports have no E axis + assert_eq!(k_from_code_shape(&[320, 1088, 32]), Ok(2)); + assert!(k_from_code_shape(&[256, 128, 64, 33]).is_err()); + } + + #[test] + fn quant_type_follows_k() { + assert_eq!(quant_type_for_k(2), Ok(QuantType::ESCHA2T16)); + assert_eq!(quant_type_for_k(3), Ok(QuantType::ESCHA3T16)); + assert!(quant_type_for_k(4).is_err()); + } + + /// `ignore` means "not escha-coded", NOT "not quantized" — both models + /// list embed_tokens and lm_head there and ship them as weight_int8. + /// Classification must key off the tensor suffix actually present. + #[test] + fn classify_leaf_keys_off_the_suffix() { + assert_eq!(classify_leaf("l.0.mlp.experts.gate_up_proj.escha_code"), Leaf::Code); + assert_eq!(classify_leaf("l.0.mlp.experts.gate_up_proj.escha_rin"), Leaf::Rin); + assert_eq!(classify_leaf("l.0.mlp.experts.gate_up_proj.escha_s_out"), Leaf::SOut); + assert_eq!(classify_leaf("lm_head.weight_int8"), Leaf::Int8); + assert_eq!(classify_leaf("lm_head.weight_scale"), Leaf::Int8Scale); + assert_eq!(classify_leaf("l.0.input_layernorm.weight"), Leaf::Passthrough); + } + + /// A future export carrying a rotation variant this version does not + /// implement must stop conversion, not decode under the wrong rotation. + #[test] + fn unknown_escha_leaf_is_rejected() { + assert_eq!( + classify_leaf("l.0.mlp.gate_proj.escha_rotation_theta"), + Leaf::UnknownEscha + ); + } + + /// Required: code, rin, rout. Missing any is "incomplete escha linear" — + /// fail loudly at load, never a partial decode. + #[test] + fn incomplete_linear_is_rejected() { + let mut present = vec![Leaf::Code, Leaf::Rin, Leaf::Rout]; + assert!(check_linear_complete("proj", &present).is_ok()); + present.pop(); + let err = check_linear_complete("proj", &present).unwrap_err(); + assert!(err.contains("incomplete escha linear"), "{err}"); + } + + /// Optional: s_in, s_out, config, bias. An export without the end-to-end + /// stage ships none of them and must still convert. + #[test] + fn optional_leaves_may_all_be_absent() { + assert!(check_linear_complete("proj", &[Leaf::Code, Leaf::Rin, Leaf::Rout]).is_ok()); + } + + /// The row scale must be replicated into every block with the int8 bytes + /// untouched — that is what makes the repack bit-exact. Recomputing block + /// scales would be a second quantisation. + #[test] + fn int8_repack_replicates_the_row_scale() { + let oc = 2; + let ic = 64; // two Q8_0 blocks per row + let w8: Vec = (0..(oc * ic)).map(|i| (i % 127) as i8).collect(); + let scale = vec![f16_rne(0.5), f16_rne(2.0)]; + let q8 = int8_rows_to_q8_0(&w8, &scale, oc, ic).unwrap(); + assert_eq!(q8.len(), oc * (ic / 32) * 34); + // Both blocks of row 0 carry row 0's scale, unchanged. + assert_eq!(&q8[0..2], &scale[0].to_le_bytes()); + assert_eq!(&q8[34..36], &scale[0].to_le_bytes()); + // Row 1's blocks carry row 1's scale. + assert_eq!(&q8[68..70], &scale[1].to_le_bytes()); + // Payload bytes are passed through verbatim. + assert_eq!(q8[2] as i8, w8[0]); + assert_eq!(q8[36] as i8, w8[32]); + } + + #[test] + fn int8_repack_rejects_a_ragged_row() { + assert!(int8_rows_to_q8_0(&[0i8; 20], &[0u16], 1, 20).is_err()); + } +} +``` + +- [ ] **Step 2: Run the tests to verify they fail** + +```bash +cargo test -p hipfire-quantize --bin hipfire-quantize pipeline_escha 2>&1 | tail -15 +``` + +Expected: compile error, `cannot find function 'k_from_code_shape'`. + +- [ ] **Step 3: Implement the classification core** + +Insert above the test module in `pipeline_escha.rs`: + +```rust +//! Converter for EschaLabs Escha-W2 checkpoints (`quant_method` = `escha` / +//! `eschamoe`) into `.hfq`. +//! +//! Code streams are copied byte-for-byte; `memcmp` on the round-trip is a +//! post-condition. See docs/plans/escha-w2-port-design.md. + +use crate::hfq::QuantType; +use std::path::Path; + +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +pub(crate) enum Leaf { + Code, + Rin, + Rout, + SIn, + SOut, + Config, + Bias, + Int8, + Int8Scale, + Passthrough, + UnknownEscha, +} + +/// The complete escha leaf namespace. Anything else beginning `escha_` is a +/// format mismatch from a newer exporter and must stop conversion. +pub(crate) fn classify_leaf(name: &str) -> Leaf { + let suffix = name.rsplit('.').next().unwrap_or(""); + match suffix { + "escha_code" => Leaf::Code, + "escha_rin" => Leaf::Rin, + "escha_rout" => Leaf::Rout, + "escha_s_in" => Leaf::SIn, + "escha_s_out" => Leaf::SOut, + "escha_config" => Leaf::Config, + "bias" => Leaf::Bias, + "weight_int8" => Leaf::Int8, + "weight_scale" => Leaf::Int8Scale, + s if s.starts_with("escha_") => Leaf::UnknownEscha, + _ => Leaf::Passthrough, + } +} + +/// `K` from the code tensor's own shape: the last dimension is `16 * K`. +/// +/// This is the ONLY source of truth. `escha_config` is optional (spec §1.4) +/// and `layer_meta.bits` is self-inconsistent across releases — the 35B +/// records down_proj as bits 3.0 / K 3, the 27B as bits 2.0 / K 3 (spec §1.3). +pub(crate) fn k_from_code_shape(shape: &[u64]) -> Result { + let last = *shape.last().ok_or("escha_code has no dimensions")? as usize; + if last % 16 != 0 { + return Err(format!("escha_code last dim {last} is not a multiple of 16")); + } + let k = last / 16; + if k != 2 && k != 3 { + return Err(format!("unsupported escha code rate K={k} (expected 2 or 3)")); + } + Ok(k) +} + +pub(crate) fn quant_type_for_k(k: usize) -> Result { + match k { + 2 => Ok(QuantType::ESCHA2T16), + 3 => Ok(QuantType::ESCHA3T16), + _ => Err(format!("unsupported escha code rate K={k}")), + } +} + +/// Required: code, rin, rout. Optional: s_in, s_out, config, bias. +pub(crate) fn check_linear_complete(proj: &str, present: &[Leaf]) -> Result<(), String> { + for req in [Leaf::Code, Leaf::Rin, Leaf::Rout] { + if !present.contains(&req) { + return Err(format!( + "incomplete escha linear '{proj}': missing {req:?}; \ + refusing to decode into noise" + )); + } + } + Ok(()) +} +``` + +- [ ] **Step 4: Run the tests to verify they pass** + +```bash +cargo test -p hipfire-quantize --bin hipfire-quantize pipeline_escha 2>&1 | tail -15 +``` + +Expected: `test result: ok. 8 passed; 0 failed`. + +- [ ] **Step 5: Implement `convert_escha`** + +Append to `pipeline_escha.rs`: + +```rust +use crate::escha_ref::fold_scales; +use crate::hfq::{write_hfq, HfqTensor}; +use crate::safetensors_file::SafetensorsFile; +use std::collections::BTreeMap; + +/// Convert an Escha-W2 checkpoint directory into a single `.hfq`. +/// +/// `arch` is 6 for `eschamoe` (MoE) and 5 for `escha` (dense). +pub(crate) fn convert_escha(src_dir: &Path, out: &Path) -> Result<(), String> { + let cfg: serde_json::Value = serde_json::from_slice( + &std::fs::read(src_dir.join("config.json")).map_err(|e| e.to_string())?, + ) + .map_err(|e| e.to_string())?; + let qc = &cfg["quantization_config"]; + let method = qc["quant_method"].as_str().unwrap_or_default(); + let version = qc["format_version"].as_str().unwrap_or_default(); + if version != "2.0" { + return Err(format!("unsupported escha format_version {version:?}; expected \"2.0\"")); + } + let arch: u32 = match method { + "eschamoe" => 6, + "escha" => 5, + other => return Err(format!("not an escha checkpoint: quant_method {other:?}")), + }; + + // Tensors can straddle shards (the 27B's mlp.up_proj has its escha_code in + // shard 2 while its metadata sits in shard 1), so resolve through every + // shard rather than per-file. + let mut shards = Vec::new(); + let mut paths: Vec<_> = std::fs::read_dir(src_dir) + .map_err(|e| e.to_string())? + .filter_map(|e| e.ok()) + .map(|e| e.path()) + .filter(|p| p.extension().is_some_and(|x| x == "safetensors")) + .collect(); + paths.sort(); + for p in &paths { + shards.push(SafetensorsFile::open(p).map_err(|e| e.to_string())?); + } + let find = |name: &str| shards.iter().find_map(|s| s.tensor_data(name)); + + // Group leaves by projection prefix so completeness can be checked. + let mut by_proj: BTreeMap> = BTreeMap::new(); + let mut passthrough: Vec = Vec::new(); + for s in &shards { + for name in s.tensor_names() { + let leaf = classify_leaf(name); + match leaf { + Leaf::UnknownEscha => { + return Err(format!( + "unknown escha tensor '{name}': this build implements \ + escha_code/rin/rout/s_in/s_out/config only. A newer \ + exporter shipped a leaf we do not decode; refusing." + )) + } + Leaf::Passthrough | Leaf::Int8 | Leaf::Int8Scale => { + passthrough.push(name.to_string()) + } + _ => { + let prefix = name.rsplit_once('.').unwrap().0.to_string(); + by_proj.entry(prefix).or_default().push((name.to_string(), leaf)); + } + } + } + } + + let mut tensors: Vec = Vec::new(); + for (proj, leaves) in &by_proj { + let kinds: Vec = leaves.iter().map(|(_, l)| *l).collect(); + check_linear_complete(proj, &kinds)?; + + let (meta, data) = find(&format!("{proj}.escha_code")) + .ok_or_else(|| format!("{proj}: escha_code vanished between passes"))?; + let k = k_from_code_shape(&meta.shape)?; + let qt = quant_type_for_k(k)?; + + // Verbatim: the code stream is copied byte-for-byte. memcmp on the + // round-trip is the post-condition (G1). + tensors.push(HfqTensor { + name: format!("{proj}.escha_code"), + quant_type: qt, + shape: meta.shape.iter().map(|&d| d as u32).collect(), + group_size: 16, + data: data.to_vec(), + spilled_len: 0, + }); + + // Fold the optional end-to-end scales into rin/rout — one f32 pair per + // projection, per row when the tensor is E-stacked. + let (rin_m, rin_d) = find(&format!("{proj}.escha_rin")).unwrap(); + let (rout_m, rout_d) = find(&format!("{proj}.escha_rout")).unwrap(); + let s_in = find(&format!("{proj}.escha_s_in")).map(|(_, d)| as_f32(d)); + let s_out = find(&format!("{proj}.escha_s_out")).map(|(_, d)| as_f32(d)); + let (ri, ro) = fold_scales( + &as_u16(rin_d), + &as_u16(rout_d), + s_in.as_deref(), + s_out.as_deref(), + ); + tensors.push(f32_tensor(&format!("{proj}.escha_rin_eff"), &rin_m.shape, ri)); + tensors.push(f32_tensor(&format!("{proj}.escha_rout_eff"), &rout_m.shape, ro)); + + if let Some((bm, bd)) = find(&format!("{proj}.bias")) { + tensors.push(HfqTensor { + name: format!("{proj}.bias"), + quant_type: QuantType::F16, + shape: bm.shape.iter().map(|&d| d as u32).collect(), + group_size: 0, + data: bd.to_vec(), + spilled_len: 0, + }); + } + } + + for name in &passthrough { + match classify_leaf(name) { + // Consumed alongside its weight_int8 sibling. + Leaf::Int8Scale => continue, + Leaf::Int8 => { + let prefix = name.rsplit_once('.').unwrap().0; + let (m, d) = find(name).unwrap(); + let (_, sd) = find(&format!("{prefix}.weight_scale")).ok_or_else(|| { + format!("{name}: weight_int8 without a matching weight_scale") + })?; + let oc = m.shape[0] as usize; + let ic = m.shape[1] as usize; + let w8: Vec = d.iter().map(|&b| b as i8).collect(); + let q8 = int8_rows_to_q8_0(&w8, &as_u16(sd), oc, ic)?; + tensors.push(HfqTensor { + name: format!("{prefix}.weight"), + quant_type: QuantType::Q8F16, + shape: vec![oc as u32, ic as u32], + group_size: 32, + data: q8, + spilled_len: 0, + }); + } + _ => { + let (m, d) = find(name).unwrap(); + tensors.push(HfqTensor { + name: name.clone(), + quant_type: match m.dtype.as_str() { + "F16" => QuantType::F16, + "F32" => QuantType::F32, + "BF16" => QuantType::BF16, + other => return Err(format!("{name}: unhandled dtype {other}")), + }, + shape: m.shape.iter().map(|&d| d as u32).collect(), + group_size: 0, + data: d.to_vec(), + spilled_len: 0, + }); + } + } + } + + // The top-level "config" key is REQUIRED: `hipfire-arch-qwen35`'s + // `config_from_hfq` -> `config_from_metadata_json` errors "qwen35: missing + // config" before touching a single tensor, so an .hfq without it is + // unloadable no matter how correct the codec is. Every sibling converter + // (pipeline_gguf.rs, pipeline_deepseek.rs, pipeline_maple.rs) embeds the + // parsed config.json the same way. `from_config_value` self-detects the + // nested text_config/vision_config these VL-shaped checkpoints carry. + let metadata = serde_json::json!({ + "config": cfg, + "escha": { "format_version": version, "quant_method": method }, + }) + .to_string(); + write_hfq(out, arch, &metadata, &tensors, None).map_err(|e| e.to_string()) +} + +/// Escha's int8 is per-output-ROW; hipfire's `Q8_0` is per-32-element block +/// (34 bytes: f16 scale then 32 int8, per `llama.rs:148`). Replicating the row +/// scale into every block of that row passes the int8 bytes through unchanged, +/// so the reconstruction is bit-identical to Escha's `w8a16`. Cost is 2 bytes +/// per 32 elements — 6.25% — for scales that are all equal within a row. +/// +/// Do NOT recompute per-block scales from the dequantised values. That is a +/// second quantisation and adds avoidable error (design §4.2.1). +pub(crate) fn int8_rows_to_q8_0( + w8: &[i8], + scale_f16: &[u16], + oc: usize, + ic: usize, +) -> Result, String> { + if w8.len() != oc * ic { + return Err(format!("int8 tensor is {} bytes, expected {oc}x{ic}", w8.len())); + } + if scale_f16.len() != oc { + return Err(format!("expected {oc} row scales, got {}", scale_f16.len())); + } + if ic % 32 != 0 { + return Err(format!("Q8_0 needs a multiple of 32 per row, got ic={ic}")); + } + let mut out = Vec::with_capacity(oc * (ic / 32) * 34); + for o in 0..oc { + let s = scale_f16[o].to_le_bytes(); + for blk in 0..ic / 32 { + out.extend_from_slice(&s); + let base = o * ic + blk * 32; + out.extend(w8[base..base + 32].iter().map(|&v| v as u8)); + } + } + Ok(out) +} + +fn as_u16(d: &[u8]) -> Vec { + d.chunks_exact(2).map(|c| u16::from_le_bytes([c[0], c[1]])).collect() +} + +fn as_f32(d: &[u8]) -> Vec { + d.chunks_exact(4).map(|c| f32::from_le_bytes([c[0], c[1], c[2], c[3]])).collect() +} + +fn f32_tensor(name: &str, shape: &[u64], v: Vec) -> HfqTensor { + HfqTensor { + name: name.to_string(), + quant_type: QuantType::F32, + shape: shape.iter().map(|&d| d as u32).collect(), + group_size: 0, + data: v.iter().flat_map(|x| x.to_le_bytes()).collect(), + spilled_len: 0, + } +} +``` + +Add to `crates/hipfire-quantize/src/main.rs`, beside the existing `mod pipeline_deepseek;` / `mod pipeline_gguf;` declarations: + +```rust +mod pipeline_escha; +``` + +- [ ] **Step 6: Build and run the full crate test suite** + +```bash +cargo build --release -p hipfire-quantize 2>&1 | tail -20 +cargo test -p hipfire-quantize 2>&1 | tail -15 +``` + +Expected: build succeeds; all tests pass. + +- [ ] **Step 7: Commit** + +```bash +git add crates/hipfire-quantize/src/pipeline_escha.rs crates/hipfire-quantize/src/main.rs +git commit -m "feat(escha): safetensors -> hfq converter with the leaf contract enforced" +``` + +--- + +### Task 6: Convert the real 35B and prove the round-trip (G1) + +The first gate that touches the actual model. It catches shape and classification errors that no synthetic test reaches. + +**Files:** +- Modify: `crates/hipfire-quantize/src/main.rs` (CLI arm) +- Create: `scripts/escha-verify-roundtrip.py` + +**Interfaces:** +- Consumes: `pipeline_escha::convert_escha` +- Produces: CLI `hipfire-quantize escha --src --out ` + +- [ ] **Step 1: Wire the CLI arm** + +In `crates/hipfire-quantize/src/main.rs`, add to the subcommand dispatch (follow the existing `deepseek4` / `gguf` arms verbatim for style): + +```rust + Some("escha") => { + let src = arg_value(&args, "--src").expect("--src required"); + let out = arg_value(&args, "--out").expect("--out required"); + pipeline_escha::convert_escha(Path::new(&src), Path::new(&out)) + .unwrap_or_else(|e| panic!("escha conversion failed: {e}")); + println!("wrote {out}"); + } +``` + +- [ ] **Step 2: Write the round-trip verifier** + +Create `scripts/escha-verify-roundtrip.py`: + +```python +#!/usr/bin/env python3 +"""G1: every escha_code tensor in the .hfq must be byte-identical to source. + +Verbatim repack is the whole basis for claiming no codec loss, so this is a +memcmp against the tensor at its indexed offset — not a substring search, +which would be quadratic over a 12 GB file. + +HFQ layout (see hipfire-quantize/src/hfq.rs::write_hfq): + header 32B : magic[4] "HFQM", version u32, arch u32, n_tensors u32, + metadata_offset u64, data_offset u64 + metadata : JSON at metadata_offset + index : n_tensors u32, then per tensor + name_len u16, name, quant_type u8, ndim u8, + dims u32*ndim, group_size u32, data_len u64 + data : at data_offset (4096-aligned), tensors concatenated in order +""" +import json, mmap, struct, sys +from pathlib import Path + +ESCHA_QT = {42: "ESCHA2T16", 43: "ESCHA3T16"} + + +def hfq_tensors(mm): + assert mm[:4] == b"HFQM", "not an HFQ file" + version, arch, n_tensors = struct.unpack_from(" HipResult<()>` + +- [ ] **Step 1: Write the kernel** + +Create `kernels/src/escha_decode_tiles.hip`: + +```cpp +// Escha-W2 one-shot tile decode -> Q8_0 resident weights (Phase 1). +// +// Reads the verbatim int16 code stream, decodes 16x16 trellis tiles, and +// writes hipfire Q8_0 (34 bytes per 32 elements: f16 scale + 32 int8). +// +// The codebook is computed inline. A 65536-entry fp16 LUT would be 128 KB and +// gfx1151 has 64 KB LDS, so there is no table anywhere in this file. +// +// Escha's tile grid is in-major [in/16, out/16]; hipfire stores weights +// out-major [out, in]. This kernel transposes on the way out. +// +// DELIBERATE DUPLICATION: hipfire-quantize/src/escha_ref.rs implements this +// same lane maths in Rust. That is the G2 gate — this kernel is asserted +// bit-exact against it. Generating either from the other would make G2 +// circular. Do not deduplicate. + +#include +#include + +__device__ __forceinline__ __half escha_cba(unsigned short state) { + unsigned int r = ((unsigned int)state * 0xCBAC1FEDu) & 0x8FFF8FFFu; + r ^= 0x3B603B60u; + __half lo = __ushort_as_half((unsigned short)(r & 0xFFFFu)); + __half hi = __ushort_as_half((unsigned short)(r >> 16)); + return __hadd(lo, hi); // fp16 RNE add — matches the reference exactly +} + +__device__ __forceinline__ void escha_decode8_k2( + const unsigned int* w, int lane, unsigned short* out) { + int t_off = lane * 8; + int i1 = t_off >> 4; + int i0 = (i1 + 15) & 15; + unsigned long long merged = ((unsigned long long)w[i0] << 32) | w[i1]; + int shift = ((~t_off) & 8) << 1; // 16 for even lanes, 0 for odd + unsigned int v = (unsigned int)((merged >> shift) & 0xFFFFFFFFull); + #pragma unroll + for (int j = 0; j < 8; ++j) out[j] = (unsigned short)(v >> (2 * (7 - j))); +} + +__device__ __forceinline__ void escha_decode8_k3( + const unsigned int* w, int lane, unsigned short* out) { + const int BITS = 3; + int t_off = lane * 8; + int b1 = (t_off + 257) * BITS; + int b0 = b1 - 16; + int b2 = b1 + BITS * 7; + int i0 = b0 >> 5; + int i2 = (b2 - 1) >> 5; + int s2 = ((i2 + 1) << 5) - b2; + unsigned long long merged = + ((unsigned long long)w[i0 % 24] << 32) | w[i2 % 24]; + unsigned int w7 = (unsigned int)((merged >> s2) & 0xFFFFFFFFull); + unsigned int w3 = (unsigned int)((merged >> (s2 + BITS * 4)) & 0xFFFFFFFFull); + out[0] = (unsigned short)(w3 >> 9); out[1] = (unsigned short)(w3 >> 6); + out[2] = (unsigned short)(w3 >> 3); out[3] = (unsigned short)(w3); + out[4] = (unsigned short)(w7 >> 9); out[5] = (unsigned short)(w7 >> 6); + out[6] = (unsigned short)(w7 >> 3); out[7] = (unsigned short)(w7); +} + +// One block per tile. 32 lanes, 8 values each = the 256 tile slots. +extern "C" __global__ void escha_decode_tiles( + const short* __restrict__ code, // [in/16, out/16, 16K] + __half* __restrict__ bare, // [in, out] fp16 scratch + int in_features, int out_features, int K) { + int tile = blockIdx.x; + int tn = out_features / 16; + int kt = tile / tn, nt = tile % tn; + int lane = threadIdx.x; + if (lane >= 32) return; + + unsigned int words[24]; + const short* src = code + (size_t)tile * 16 * K; + for (int i = 0; i < 8 * K; ++i) + words[i] = ((unsigned int)(unsigned short)src[2 * i]) | + (((unsigned int)(unsigned short)src[2 * i + 1]) << 16); + + unsigned short st[8]; + if (K == 2) escha_decode8_k2(words, lane, st); + else escha_decode8_k3(words, lane, st); + + int l0 = lane & ~4; + int c_off = (lane >> 2) & 1; + #pragma unroll + for (int j = 0; j < 8; ++j) { + int fi = j >> 1; + int row = (lane & 3) * 2 + (j & 1) + (fi & 1) * 8; + int col = 2 * ((l0 >> 3) + (j >= 4 ? 4 : 0)) + c_off; + bare[(size_t)(kt * 16 + row) * out_features + (nt * 16 + col)] = escha_cba(st[j]); + } +} +``` + +- [ ] **Step 2: Register the kernel source** + +In `crates/rdna-compute/src/kernels.rs`, beside line 4810: + +```rust +pub const ESCHA_DECODE_TILES_SRC: &str = + include_str!("../../../kernels/src/escha_decode_tiles.hip"); +``` + +In `crates/rdna-compute/src/dispatch.rs` near line 4251: + +```rust + specs.push(("escha_decode_tiles", kernels::ESCHA_DECODE_TILES_SRC.to_string())); +``` + +- [ ] **Step 3: Write the GPU-vs-CPU parity test** + +Create `crates/rdna-compute/examples/test_escha_decode_gpu_vs_cpu.rs`: + +```rust +//! G2: GPU tile decode must match escha_ref::reconstruct EXACTLY in fp16, +//! for both K. Run: +//! cargo run --release -p rdna-compute --example test_escha_decode_gpu_vs_cpu +use hipfire_quantize::escha_ref; + +fn main() { + for (name, ic, oc, k) in [ + ("packed_gu_e0_k2.i16", 2048usize, 1024usize, 2usize), + ("packed_down_e0_k3.i16", 512, 2048, 3), + ] { + let path = format!( + "{}/../hipfire-quantize/tests/data/escha/{name}", + env!("CARGO_MANIFEST_DIR") + ); + let raw = std::fs::read(&path).expect("run fetch-goldens.sh first"); + let code: Vec = + raw.chunks_exact(2).map(|c| i16::from_le_bytes([c[0], c[1]])).collect(); + let want = escha_ref::reconstruct(&code, ic, oc, k); + + let mut gpu = rdna_compute::Gpu::new().expect("gpu"); + let got = gpu.escha_decode_tiles_host(&code, ic as u32, oc as u32, k as u32) + .expect("decode"); + + let bad = want.iter().zip(&got).filter(|(a, b)| a != b).count(); + println!("{name}: {bad} mismatched of {} elements", want.len()); + assert_eq!(bad, 0, "{name}: GPU decode diverges from the CPU reference"); + } + println!("G2 PASS"); +} +``` + +Add `hipfire-quantize` to `crates/rdna-compute/Cargo.toml` under `[dev-dependencies]`: + +```toml +hipfire-quantize = { path = "../hipfire-quantize" } +``` + +- [ ] **Step 4: Run it and confirm it fails** + +```bash +cargo run --release -p rdna-compute --example test_escha_decode_gpu_vs_cpu 2>&1 | tail -10 +``` + +Expected: compile error, `no method named 'escha_decode_tiles_host'`. + +- [ ] **Step 5: Implement the host wrapper** + +Add to `crates/rdna-compute/src/gemv.rs` beside `gemv_q8_0` (line 13563), following its `ensure_kernel` + launch style: + +```rust + /// Decode an escha code stream to a bare fp16 weight matrix `[ic, oc]`. + /// Host-side helper used by the G2 parity gate; the load path uses the + /// device-resident form. + pub fn escha_decode_tiles_host( + &mut self, + code: &[i16], + in_features: u32, + out_features: u32, + k: u32, + ) -> HipResult> { + self.bind_thread()?; + let n_elems = (in_features as usize) * (out_features as usize); + let n_tiles = (in_features / 16) * (out_features / 16); + let code_bytes: Vec = code.iter().flat_map(|v| v.to_le_bytes()).collect(); + let d_code = self.upload_raw(&code_bytes, &[code.len()])?; + let d_bare = self.alloc_tensor(&[n_elems], DType::F16)?; + self.ensure_kernel( + "escha_decode_tiles", + kernels::ESCHA_DECODE_TILES_SRC, + "escha_decode_tiles", + )?; + + let mut code_ptr = d_code.buf.as_ptr(); + let mut bare_ptr = d_bare.buf.as_ptr(); + let mut ic = in_features as i32; + let mut oc = out_features as i32; + let mut kk = k as i32; + let mut params: Vec<*mut c_void> = vec![ + &mut code_ptr as *mut _ as *mut c_void, + &mut bare_ptr as *mut _ as *mut c_void, + &mut ic as *mut _ as *mut c_void, + &mut oc as *mut _ as *mut c_void, + &mut kk as *mut _ as *mut c_void, + ]; + let func = &self.functions["escha_decode_tiles"]; + unsafe { + self.hip + .launch_kernel(func, [n_tiles, 1, 1], [32, 1, 1], 0, None, &mut params)?; + } + + let mut out = vec![0u8; n_elems * 2]; + self.hip.memcpy_dtoh(&mut out, &d_bare.buf)?; + Ok(out + .chunks_exact(2) + .map(|c| u16::from_le_bytes([c[0], c[1]])) + .collect()) + } +``` + +This follows the launch convention used throughout `gemv.rs` (see `gemv_q8_0` +at line 13563): `bind_thread` → `upload_raw`/`alloc_tensor` → `ensure_kernel` → +build a `Vec<*mut c_void>` of `&mut` locals → look the function up in +`self.functions` → `unsafe { self.hip.launch_kernel(...) }`. There is no +`Gpu::launch_kernel(name, ...)` helper taking buffers directly; do not invent +one. `c_void` is already imported in `gemv.rs`. + +- [ ] **Step 6: Run the parity gate (G2)** + +```bash +which hipcc || echo "FIX: hipcc must be on PATH or the daemon cannot JIT" +cargo run --release -p rdna-compute --example test_escha_decode_gpu_vs_cpu 2>&1 | tail -10 +``` + +Expected: + +``` +packed_gu_e0_k2.i16: 0 mismatched of 2097152 elements +packed_down_e0_k3.i16: 0 mismatched of 1048576 elements +G2 PASS +``` + +Any nonzero mismatch count is a decode bug, not a rounding artifact — `__hadd` is RNE and the CPU reference rounds identically. If K=2 passes and K=3 fails, the bug is in `escha_decode8_k3`'s modular wrap. + +- [ ] **Step 7: Commit** + +```bash +git add kernels/src/escha_decode_tiles.hip crates/rdna-compute/src/kernels.rs \ + crates/rdna-compute/src/dispatch.rs crates/rdna-compute/src/gemv.rs \ + crates/rdna-compute/examples/test_escha_decode_gpu_vs_cpu.rs \ + crates/rdna-compute/Cargo.toml +git commit -m "feat(escha): GPU tile decode kernel, bit-exact against the CPU reference" +``` + +--- + +### Task 8: H128 input and output transforms (G3) + +**Files:** +- Create: `kernels/src/escha_h128.hip` +- Modify: `crates/rdna-compute/src/kernels.rs`, `crates/rdna-compute/src/dispatch.rs` +- Test: `crates/rdna-compute/examples/test_escha_h128_gpu_vs_cpu.rs` + +**Interfaces:** +- Consumes: `escha_ref::{h128_inplace, input_transform, output_transform, RS}` +- Produces: kernels `escha_h128_in`, `escha_h128_out`; `Gpu::escha_h128_in(...)`, `Gpu::escha_h128_out(...)` + +- [ ] **Step 1: Write the kernel** + +Create `kernels/src/escha_h128.hip`: + +```cpp +// Escha-W2 activation transforms. +// in : xh = f16( H128(x * rin) * RS ) +// out: y = f16( H128(mid) * RS * rout ) +// +// H128 is the UNNORMALISED 128-point Walsh-Hadamard in Sylvester (natural) +// order. Escha folds its sign flips into rin/rout, so unlike hipfire's +// gemv_mq4g128 there is no sign-seed stage here. +// +// A pruned output channel (rout == 0) must come out EXACTLY zero — do not +// reorder the final multiply in a way that could produce -0.0 or a denormal. + +#include +#include + +#define ESCHA_RS 0.0883883476f // 1/sqrt(128) + +__device__ __forceinline__ void h128_block(float* v) { + for (int h = 1; h < 128; h <<= 1) { + for (int i = 0; i < 128; i += (h << 1)) { + for (int j = i; j < i + h; ++j) { + float a = v[j], b = v[j + h]; + v[j] = a + b; + v[j + h] = a - b; + } + } + } +} + +// One block per 128-channel group. 128 threads cooperate via LDS. +extern "C" __global__ void escha_h128_in( + const float* __restrict__ x, const float* __restrict__ rin, + __half* __restrict__ xh, int n) { + __shared__ float s[128]; + int g = blockIdx.x, t = threadIdx.x; + int idx = g * 128 + t; + if (idx >= n) return; + s[t] = x[idx] * rin[idx]; + __syncthreads(); + if (t == 0) h128_block(s); + __syncthreads(); + xh[idx] = __float2half(s[t] * ESCHA_RS); +} + +extern "C" __global__ void escha_h128_out( + const float* __restrict__ mid, const float* __restrict__ rout, + __half* __restrict__ y, int n) { + __shared__ float s[128]; + int g = blockIdx.x, t = threadIdx.x; + int idx = g * 128 + t; + if (idx >= n) return; + s[t] = mid[idx]; + __syncthreads(); + if (t == 0) h128_block(s); + __syncthreads(); + y[idx] = __float2half(s[t] * ESCHA_RS * rout[idx]); +} +``` + +- [ ] **Step 2: Register both kernels** + +In `crates/rdna-compute/src/kernels.rs`: + +```rust +pub const ESCHA_H128_SRC: &str = include_str!("../../../kernels/src/escha_h128.hip"); +``` + +In `crates/rdna-compute/src/dispatch.rs` near line 4251: + +```rust + specs.push(("escha_h128", kernels::ESCHA_H128_SRC.to_string())); +``` + +- [ ] **Step 3: Write the parity test** + +Create `crates/rdna-compute/examples/test_escha_h128_gpu_vs_cpu.rs`: + +```rust +//! G3: the H128 kernels must match escha_ref directly. +//! +//! A round-trip check (H128 . H128 == 128 I) is NOT sufficient — a wrong +//! butterfly order is also self-inverse and would pass it while being wrong. +use hipfire_quantize::escha_ref; +use hipfire_quantize::float16::f16_to_f32; + +fn main() { + let n = 2048usize; + let x: Vec = (0..n).map(|i| ((i * 37) as f32 * 0.017).sin()).collect(); + let rin: Vec = (0..n).map(|i| if i % 3 == 0 { -0.0023 } else { 0.0023 }).collect(); + let mut rout: Vec = (0..n).map(|i| 1.0 + (i % 5) as f32 * 0.1).collect(); + rout[7] = 0.0; + rout[1000] = 0.0; // pruned channels must stay exactly zero + + let want_in = escha_ref::input_transform(&x, &rin); + let want_out = escha_ref::output_transform(&x, &rout); + + let mut gpu = rdna_compute::Gpu::new().expect("gpu"); + let got_in = gpu.escha_h128_in_host(&x, &rin).expect("h128 in"); + let got_out = gpu.escha_h128_out_host(&x, &rout).expect("h128 out"); + + let bad_in = want_in.iter().zip(&got_in).filter(|(a, b)| a != b).count(); + let bad_out = want_out.iter().zip(&got_out).filter(|(a, b)| a != b).count(); + println!("h128_in : {bad_in} mismatched of {n}"); + println!("h128_out: {bad_out} mismatched of {n}"); + assert_eq!(f16_to_f32(got_out[7]), 0.0, "pruned channel 7 must be exactly zero"); + assert_eq!(f16_to_f32(got_out[1000]), 0.0, "pruned channel 1000 must be exactly zero"); + assert_eq!(bad_in, 0); + assert_eq!(bad_out, 0); + println!("G3 PASS"); +} +``` + +- [ ] **Step 4: Run it and confirm it fails** + +```bash +cargo run --release -p rdna-compute --example test_escha_h128_gpu_vs_cpu 2>&1 | tail -10 +``` + +Expected: compile error, `no method named 'escha_h128_in_host'`. + +- [ ] **Step 5: Implement the host wrappers** + +Add to `crates/rdna-compute/src/gemv.rs`: + +```rust + /// `xh = f16( H128(x * rin) * RS )` on device. Host-side helper for the + /// G3 parity gate; the forward path uses the device-resident form. + pub fn escha_h128_in_host(&mut self, x: &[f32], rin: &[f32]) -> HipResult> { + self.escha_h128_host_impl("escha_h128_in", x, rin) + } + + /// `y = f16( H128(mid) * RS * rout )` on device. + pub fn escha_h128_out_host(&mut self, mid: &[f32], rout: &[f32]) -> HipResult> { + self.escha_h128_host_impl("escha_h128_out", mid, rout) + } + + fn escha_h128_host_impl( + &mut self, + entry: &str, + a: &[f32], + vec_in: &[f32], + ) -> HipResult> { + assert_eq!(a.len(), vec_in.len()); + assert_eq!(a.len() % 128, 0, "H128 needs a multiple of 128"); + self.bind_thread()?; + let n = a.len(); + let a_bytes: Vec = a.iter().flat_map(|v| v.to_le_bytes()).collect(); + let v_bytes: Vec = vec_in.iter().flat_map(|v| v.to_le_bytes()).collect(); + let d_a = self.upload_raw(&a_bytes, &[n])?; + let d_v = self.upload_raw(&v_bytes, &[n])?; + let d_out = self.alloc_tensor(&[n], DType::F16)?; + self.ensure_kernel("escha_h128", kernels::ESCHA_H128_SRC, entry)?; + + let mut a_ptr = d_a.buf.as_ptr(); + let mut v_ptr = d_v.buf.as_ptr(); + let mut o_ptr = d_out.buf.as_ptr(); + let mut n_val = n as i32; + let mut params: Vec<*mut c_void> = vec![ + &mut a_ptr as *mut _ as *mut c_void, + &mut v_ptr as *mut _ as *mut c_void, + &mut o_ptr as *mut _ as *mut c_void, + &mut n_val as *mut _ as *mut c_void, + ]; + let func = &self.functions[entry]; + unsafe { + self.hip + .launch_kernel(func, [(n / 128) as u32, 1, 1], [128, 1, 1], 0, None, &mut params)?; + } + + let mut raw = vec![0u8; n * 2]; + self.hip.memcpy_dtoh(&mut raw, &d_out.buf)?; + Ok(raw + .chunks_exact(2) + .map(|c| u16::from_le_bytes([c[0], c[1]])) + .collect()) + } +``` + +Note `ensure_kernel(module, SRC, entry)` is called once per entry point — the +two entries live in the same `.hip` source but are separate device functions, +so each needs its own `ensure_kernel` call and its own `self.functions` lookup +key. + +- [ ] **Step 6: Run the parity gate (G3)** + +```bash +cargo run --release -p rdna-compute --example test_escha_h128_gpu_vs_cpu 2>&1 | tail -10 +``` + +Expected: + +``` +h128_in : 0 mismatched of 2048 +h128_out: 0 mismatched of 2048 +G3 PASS +``` + +- [ ] **Step 7: Commit** + +```bash +git add kernels/src/escha_h128.hip crates/rdna-compute/src/kernels.rs \ + crates/rdna-compute/src/dispatch.rs crates/rdna-compute/src/gemv.rs \ + crates/rdna-compute/examples/test_escha_h128_gpu_vs_cpu.rs +git commit -m "feat(escha): H128 input/output transform kernels, gated against the reference" +``` + +--- + +### Task 9: Verify the arch-6 router contract (G4b) + +Done before wiring, because if the router disagrees the wiring is built on +sand. The spec assumed arch-6's router is reusable; this exercises the real +router and proves or disproves it. + +**Files:** +- Create: `crates/hipfire-arch-qwen35/examples/escha_router_contract.rs` +- Possibly modify: the arch-6 router path (only if Step 4 shows a mismatch) + +**Interfaces:** +- Consumes: the `.hfq` from Task 6 (for `layers.0.mlp.gate.weight`), the golden + fixture from `crates/hipfire-quantize/tests/data/escha/` +- Produces: a passing G4b assertion, or a fix to the router + +- [ ] **Step 1: Fetch the fixture** + +```bash +crates/hipfire-quantize/tests/data/escha/fetch-goldens.sh +``` + +Expected: eight files. `moeblk_ids.i64` must have digest +`0781eecdacd5fbfe30887f6d1f6af5d5ca001d32253da8bb3e1294230c2ed649`. + +- [ ] **Step 2: Write the failing test** + +This calls hipfire's **actual** arch-6 router — not a reimplementation of it — +and asserts against Escha's shipped selection. + +Create `crates/hipfire-arch-qwen35/examples/escha_router_contract.rs`: + +```rust +//! G4b: hipfire's arch-6 router must select the same experts Escha does. +//! +//! Escha rounds router logits to f16 BEFORE top-k (`ref.py`: the logits are +//! computed as f16 then widened to f32 to select). Selecting on unrounded f32 +//! logits is a different function, and the rounding manufactures exact ties +//! that f32 never produces. +//! +//! Asserts the SET, not the order: the combine is a sum over slots, so intra-k +//! order cannot change the output. On the fixture, token 3 has two experts on +//! identical f16 logits (1.80078), and which one lands in which slot is +//! implementation-defined. +//! +//! Run: +//! cargo run --release -p hipfire-arch-qwen35 \ +//! --example escha_router_contract -- /data/hipfire-models/escha-35b.hfq +use std::collections::HashSet; +use std::path::PathBuf; + +fn fixture(name: &str) -> PathBuf { + PathBuf::from(env!("CARGO_MANIFEST_DIR")) + .join("../hipfire-quantize/tests/data/escha") + .join(name) +} + +fn read_f16_as_f32(name: &str) -> Vec { + let raw = std::fs::read(fixture(name)).expect("run fetch-goldens.sh first"); + raw.chunks_exact(2) + .map(|c| hipfire_quantize::float16::f16_to_f32(u16::from_le_bytes([c[0], c[1]]))) + .collect() +} + +fn main() { + let hfq = std::env::args().nth(1).expect("usage: "); + let x = read_f16_as_f32("moeblk_x.f16"); // [8, 2048] + let raw_ids = std::fs::read(fixture("moeblk_ids.i64")).unwrap(); + let want_ids: Vec = raw_ids + .chunks_exact(8) + .map(|c| i64::from_le_bytes(c.try_into().unwrap())) + .collect(); // [8, 8] + + // Call hipfire's real router for layer 0 on each token. + let got = hipfire_arch_qwen35::escha_router_topk_for_test(&hfq, 0, &x, 8, 2048, 8) + .expect("router"); + + let mut bad = 0usize; + for t in 0..8 { + let want: HashSet = want_ids[t * 8..(t + 1) * 8].iter().copied().collect(); + let mine: HashSet = got[t * 8..(t + 1) * 8].iter().map(|&v| v as i64).collect(); + if want != mine { + bad += 1; + println!("token {t}: escha={:?}", &want_ids[t * 8..(t + 1) * 8]); + println!(" hipfire={:?}", &got[t * 8..(t + 1) * 8]); + } + } + println!("tokens with a differing top-8 SET: {bad}/8"); + assert_eq!( + bad, 0, + "arch-6 router selects different experts than escha. Most likely cause: \ + it is not rounding logits to f16 before top-k." + ); + println!("G4b PASS"); +} +``` + +- [ ] **Step 3: Run it and confirm it fails** + +```bash +cargo run --release -p hipfire-arch-qwen35 --example escha_router_contract \ + -- /data/hipfire-models/escha-35b.hfq 2>&1 | tail -15 +``` + +Expected: compile error — `escha_router_topk_for_test` does not exist yet. + +- [ ] **Step 4: Expose the router and run the gate** + +Read the arch-6 router path first: + +```bash +grep -rn "moe_topk\|router\|norm_topk_prob" crates/hipfire-arch-qwen35/src/ | head -20 +grep -rn "moe_topk" kernels/src/*.hip | head +``` + +Add a thin `pub fn escha_router_topk_for_test(hfq_path: &str, layer: usize, x: &[f32], n_tokens: usize, hidden: usize, top_k: usize) -> Result, String>` that loads `layers.{layer}.mlp.gate.weight` from the `.hfq` and runs **the production router path** — not a reimplementation. If that path is not callable in isolation, extract the selection step into a function both it and this helper call, and leave the production behaviour unchanged. + +Then run the gate: + +```bash +cargo run --release -p hipfire-arch-qwen35 --example escha_router_contract \ + -- /data/hipfire-models/escha-35b.hfq 2>&1 | tail -15 +``` + +- **If it passes:** the router already matches. Record that in the commit message and change nothing. +- **If it fails:** check whether the router rounds logits to f16 before top-k. If it does not, add the rounding **on the escha path only**, keyed on the model carrying `ESCHA2T16`/`ESCHA3T16` experts, so existing `qwen3.6:35b-a3b-*` SKUs keep their current selection bit-for-bit. Re-run until `G4b PASS`. + +- [ ] **Step 5: Confirm no existing SKU changed** + +Only required if Step 4 modified the router. + +```bash +cargo test -p hipfire-arch-qwen35 2>&1 | tail -15 +cargo test -p hipfire-dispatch-tests qwen35 2>&1 | tail -10 +``` + +Expected: all pass, unchanged. + +- [ ] **Step 6: Commit** + +```bash +git add crates/hipfire-arch-qwen35 +git commit -m "test(escha): G4b — arch-6 router selects the same experts as escha" +``` + +--- + +### Task 10: Wire arch-6 to load and run escha experts (G4) + +**Files:** +- Modify: `crates/hipfire-arch-qwen35/src/qwen35/weights.rs:69-80` (expert loading) +- Modify: `crates/hipfire-arch-qwen35/src/qwen35.rs` (MoE forward: call the transforms) +- Modify: `crates/hipfire-loader/src/lib.rs` (accept the new quant types) +- Test: `crates/hipfire-arch-qwen35/examples/escha_moe_block_gate.rs` (new) + +**Interfaces:** +- Consumes: `Gpu::escha_decode_tiles_host` (Task 7), `Gpu::escha_h128_in_host` / `Gpu::escha_h128_out_host` (Task 8), `escha_ref::{expert_linear, swiglu, w8a16}`. The load path needs device-resident variants of all three that take an existing `GpuTensor` instead of host slices; add them beside the `_host` helpers rather than round-tripping through the CPU per expert. +- Produces: arch-6 loads `ESCHA2T16`/`ESCHA3T16` experts as `Q8_0` resident + +- [ ] **Step 1: Fetch the MoE golden fixture** + +```bash +crates/hipfire-quantize/tests/data/escha/fetch-goldens.sh +``` + +Expected: eight files, with `moeblk_out.f16` at digest `484abdf7257631105d90e5e8f7794974d02fab6c7dd587e30955f660a688c4d4`. + +- [ ] **Step 2: Write the G4 gate** + +Create `crates/hipfire-arch-qwen35/examples/escha_moe_block_gate.rs`. It loads `/data/hipfire-models/escha-35b.hfq`, runs the layer-0 MoE block on `moeblk_x.f16` with `moeblk_ids.i64` / `moeblk_scores.f32` injected, and compares against `moeblk_out.f16` with these **measured** tolerances: + +```rust + let max_abs = diffs.iter().cloned().fold(0.0f32, f32::max); + let mean_abs = diffs.iter().sum::() / diffs.len() as f32; + println!("MoE block: max|diff|={max_abs:.3e} mean|diff|={mean_abs:.3e}"); + // Measured against the shipped layer-0 weights: max 1.22e-4, mean 2.1e-6, + // on outputs of mean magnitude 0.0185. The golden came from the Metal + // path, not ref.py, so this is a TOLERANCE gate. The codec goldens in G0 + // are bit-exact; do not generalise these bounds to them. + assert!(max_abs <= 2e-4, "max|diff| {max_abs:.3e} exceeds 2e-4"); + assert!(mean_abs <= 1e-5, "mean|diff| {mean_abs:.3e} exceeds 1e-5"); + println!("G4 PASS"); +``` + +Inject the routing rather than computing it — the fixture ships ids/scores precisely because it does not gate the router (that is Task 9's job). + +- [ ] **Step 3: Run it and confirm it fails** + +```bash +cargo run --release -p hipfire-arch-qwen35 --example escha_moe_block_gate 2>&1 | tail -15 +``` + +Expected: failure — the loader does not yet accept `ESCHA2T16`. + +- [ ] **Step 4: Implement expert loading** + +In `crates/hipfire-arch-qwen35/src/qwen35/weights.rs`, extend the expert loader documented at line 69 (`experts[X].gate_up: [2*moe_intermediate, hidden]`). For each expert: + +1. Read the `ESCHA2T16` `gate_up` and `ESCHA3T16` `down` code streams. +2. Call the device-resident `escha_decode_tiles` to produce bare fp16 `[ic, oc]`. +3. Quantise to `Q8_0` per output row and store in the existing `experts[X].gate_up` / `.down` slots — hipfire is out-major, escha's grid is in-major, so this is where the transpose lands. +4. Keep `escha_rin_eff` / `escha_rout_eff` as f32 device tensors alongside. + +In the MoE forward path, wrap each expert projection: the device-resident +`escha_h128_in` before the `Q8_0` GEMV, `escha_h128_out` after. + +**BATCH THE TRANSFORMS ACROSS EXPERTS — this is a hard requirement, not an +optimisation.** Task 8 measured the H128 kernels to be *launch-bound*, not +bandwidth-bound: an empty kernel at the same grid/block costs 1.74–1.78 us, +which is 70–75% of the 2.4 us a real launch takes, and the overhead-subtracted +time stays nearly flat (0.59 → 0.69 us) from 16 to 136 blocks. + +A naive one-launch-per-expert-per-projection wiring costs +`40 layers x 8 experts x 4 transforms = 1280` launches per token: + +| wiring | launches/token | H128 cost | ceiling from H128 alone | +|---|---|---|---| +| per-expert (naive) | 1280 | 3.07 ms | **326 tok/s** | +| batched across experts | 160 | 0.38 ms | ~2600 tok/s | + +326 tok/s is a hard ceiling *before any GEMV work*, which would make Phase 1 +pointless even as a correctness baseline. Issue ONE launch per (layer, +projection, side) covering all `top_k` experts — the per-expert `rin_eff` / +`rout_eff` rows are just an extra index into the already-resident +`[E, IC]` / `[E, OC]` tensors, so this is an indexing change in the kernel and +a grid change on the host, not new maths. Verify the batched form against +`escha_ref` exactly as the per-expert form was. SwiGLU consumes the **f16-rounded** merged `gate_up` output; the combine multiplies by `f16(score)`. + +- [ ] **Step 5: Run the G4 gate** + +```bash +cargo run --release -p hipfire-arch-qwen35 --example escha_moe_block_gate 2>&1 | tail -15 +``` + +Expected: `max|diff|` around `1.2e-4`, `mean|diff|` around `2e-6`, then `G4 PASS`. + +If `max|diff|` is around 1e-1 rather than 1e-4, the H128 pair is not being applied — check that the escha types did not fall through to a Plain GEMV (Task 4's guard should have made that impossible; if it did happen, the guard has a hole). + +- [ ] **Step 6: Commit** + +```bash +git add crates/hipfire-arch-qwen35/src crates/hipfire-loader/src/lib.rs \ + crates/hipfire-arch-qwen35/examples/escha_moe_block_gate.rs +git commit -m "feat(escha): load escha experts as Q8_0 and run the H128 pair in arch-6" +``` + +--- + +### Task 11: Registry entry, coherence and KLD (G5) + +**Files:** +- Modify: `registry/v1.json` +- Create: `scripts/escha-kld.sh` + +- [ ] **Step 1: Add the registry entry** + +In `registry/v1.json`, beside the existing `qwen3.6:35b-a3b` entries: + +```json +{ "name": "qwen3.6:35b-a3b-escha", "arch_id": 6, "quant": "escha" } +``` + +Match the surrounding entries' exact field set — copy a neighbouring `qwen3.6:35b-a3b-mq2` object and change only `name` and `quant`. + +- [ ] **Step 2: Serve and check coherence** + +```bash +which hipcc || echo "FIX: hipcc must be on PATH" +HIPFIRE_MODEL=/data/hipfire-models/escha-35b.hfq ~/.hipfire/serve-stable.sh & +sleep 60 +curl -s localhost:8080/v1/chat/completions -H 'Content-Type: application/json' \ + -d '{"model":"qwen3.6:35b-a3b-escha","messages":[{"role":"user","content":"What is the capital of France? Answer in one word."}],"max_tokens":16}' \ + | python3 -c "import sys,json;print(json.load(sys.stdin)['choices'][0]['message']['content'])" +``` + +Expected: `Paris`. Garbled output on a first run is usually the kernel cache, not the codec — confirm a single-toolchain cache before investigating the decode. + +- [ ] **Step 3: Measure VRAM against the prediction** + +```bash +rocm-smi --showmeminfo vram | head -20 +``` + +Expected: about **36.7 GB** resident. A number far below that means experts are not all resident; far above means something is being held at fp16 rather than `Q8_0`. + +- [ ] **Step 4: Run the KLD gate (G5)** + +Create `scripts/escha-kld.sh`: + +```bash +#!/usr/bin/env bash +# G5: two KLD comparisons on a FIXED corpus slice, teacher-forced. +# +# The reference is escha_ref on CPU, NOT any Escha runtime: escha-mlx is +# Metal, the escha wheel is CUDA, and ZML needs an NVIDIA driver, so none of +# them execute on gfx1151. ref.py declares itself the semantic contract for +# their kernels and is gated on the goldens, so this is exact rather than +# cross-machine. +# +# Score on the fixed corpus, never on the model's own greedy output: for ds4 +# that scored 8x better on the median and was optimistic. +set -euo pipefail +HFQ=${1:-/data/hipfire-models/escha-35b.hfq} +SLICE=benchmarks/quality-baselines/slice/wikitext2-1024s-2048ctx.txt +POS=${POS:-192} + +echo "== 1/2 hipfire vs escha_ref (CPU oracle) ==" +cargo run --release -p hipfire-runtime --example eval_hipfire -- \ + --model "$HFQ" --corpus "$SLICE" --positions "$POS" \ + --reference escha-ref --teacher-forced --report kld + +echo "== 2/2 hipfire vs bf16 parent Qwen/Qwen3.6-35B-A3B ==" +echo " (compare this number against the existing qwen3.6:35b-a3b-mq2 result)" +cargo run --release -p hipfire-runtime --example eval_hipfire -- \ + --model "$HFQ" --corpus "$SLICE" --positions "$POS" \ + --reference /data/hipfire-models/qwen3.6-35b-a3b-bf16 \ + --teacher-forced --report kld +``` + +```bash +chmod +x scripts/escha-kld.sh +./scripts/escha-kld.sh 2>&1 | tail -30 +``` + +`eval_hipfire` currently has no `--reference escha-ref` arm — add one that calls `escha_ref` per position. Teacher forcing is load-bearing and easy to get subtly wrong: two quants diverge at roughly token 0, so an unforced position-wise KL compares unrelated continuations. Force at **both** sites (the pre-loop sampled token and the decode loop) and key on a call counter, not `pos` — the two share a `pos`. Assert the committed token streams are identical; without that assertion a wrong number reads as fact. + +- [ ] **Step 5: Record the results** + +Expected: comparison 1 near zero, dominated by the `Q8_0` intermediate rather than the codec. A handful of divergent positions traceable to boundary ties in the router is expected and is not a codec bug (§5) — attribute before investigating. + +Write both numbers into `docs/plans/escha-w2-port-design.md` under a new "Phase 1 results" section, alongside the measured tok/s. State plainly that Phase 1 loses to `qwen3.6:35b-a3b-mq2` on speed; that is the expected outcome and Phase 2's job. + +- [ ] **Step 6: Commit** + +```bash +git add registry/v1.json scripts/escha-kld.sh docs/plans/escha-w2-port-design.md +git commit -m "feat(escha): registry entry, coherence and KLD results for Phase 1" +``` + +--- + +## Definition of done + +- G0–G5 all pass, with G2 and G3 bit-exact and G4 inside the stated tolerances. +- `cargo build --release --workspace --all-targets --locked` is clean. +- `qwen3.6:35b-a3b-escha` serves coherent text at roughly 36.7 GB resident. +- Both KLD numbers are recorded in the design doc. + +## Deferred to later plans + +- **Phase 2**: fused decode+GEMV; the prune-mask optimisation (skip `down_proj` input rows for the channels `gate_up.rout` zeroes — spec §4.5). +- **Qwen3.8-27B**: 400 escha tensors over 10 projections, the 3-way `FusedQkvQ8_0` path, and arch-5 bias slots. diff --git a/docs/plans/escha-w2-port-design.md b/docs/plans/escha-w2-port-design.md new file mode 100644 index 0000000000..32db42e97e --- /dev/null +++ b/docs/plans/escha-w2-port-design.md @@ -0,0 +1,880 @@ +# Escha-W2 port — design + +**Status:** Phase 1 implemented and gated; see §10 for measured results and §10.5 for the limitations that remain open. +**Date:** 2026-09-02 +**Branch:** `nw_escha_w2` (worktree `~/repos/hipfire-escha`, off `origin/master` @ `8cd15a62b`) +**Targets:** `EschaLabs/Qwen3.6-35B-A3B-Escha-W2` (first), `EschaLabs/Qwen3.8-27B-Escha-W2` (second) + +## 1. What Escha-W2 is + +Escha Labs publishes a 2-bit *quantization format*, not a model family. Two +checkpoints are in scope: + +| release | HF repo | on disk | `quant_method` | hipfire `arch_id` | +|---|---|---|---|---| +| 35B-A3B | `EschaLabs/Qwen3.6-35B-A3B-Escha-W2` | 12.3 GB | `eschamoe` | 6 | +| 27B | `EschaLabs/Qwen3.8-27B-Escha-W2` | 10.15 GB | `escha` | 5 | + +**Essentially no architecture work.** `Qwen3.8-27B`'s `config.json` differs from +`Qwen3.6-27B`'s only in the `quantization_config` block — every architectural +field is identical. Both models are hybrid GatedDeltaNet + full-attention +(`full_attention_interval = 4`), already served by `hipfire-arch-qwen35` as +`arch_id` 5 (dense) and 6 (MoE). The registry already ships `qwen3.6:27b` and +`qwen3.6:35b-a3b`. + +Base identity was checked both ways: the 35B's config differs from +`Qwen/Qwen3.6-35B-A3B` only in `transformers_version`, and the 27B's differs +from `Qwen/Qwen3.6-27B` only in the `quantization_config` block. Both bases are +VL-shaped composites (nested `text_config` plus a `vision_config` carrying no +weights), which is exactly the case `Qwen35Config::is_vl_text` already covers — +the converter takes that branch rather than treating them as plain text configs. + +The one exception: the 27B's escha linears carry fp16 biases the base +architecture does not have, so arch-5 gains bias slots (§1.3, §9). Otherwise +this is a codec + loader project. + +### 1.1 Format + +The codec is open and bit-exactly specified: `EschaLabs/escha-mlx` is Apache-2.0 +and ships `escha_mlx/ref.py`, a NumPy reference that is the format contract, +plus golden vectors under `tests/data/`. Nothing needs reverse-engineering. +Their `THIRD_PARTY_LICENSES` vendors exllamav3, and the codec is QTIP lineage. + +**Weight stream.** 16x16 tiles, K bits/weight, packed `int16[in/16, out/16, 16K]` +(MoE exports carry a leading `[E]` axis; dense exports do not). Verified against +the shipped safetensors headers: + +- `gate_up_proj.escha_code` `[256, 128, 64, 32]`, `in=2048 out=1024 K=2` -> 2.00 bpw +- `down_proj.escha_code` `[256, 32, 128, 48]`, `in=512 out=2048 K=3` -> 3.00 bpw + +**Codebook.** Not a table — a 3-op integer hash. For a 16-bit state `s`: + +``` +r = ((s * 0xCBAC1FED) & 0x8FFF8FFF) ^ 0x3B603B60 # 32-bit +value = f16_lo(r) + f16_hi(r) # fp16 RNE add +``` + +**This is a trellis, not a per-weight codebook.** Each weight's value is a full +16-bit state indexing 65536 fp16 values; consecutive states are overlapping +windows of the bitstream sliding by K bits. The K bpw is amortized across the +overlap. This is why no lossless repack into MQ2 exists (see §2). + +**Rotation.** Unnormalized 128-point Walsh-Hadamard (Sylvester / natural order) +on contiguous 128-channel blocks, applied on *both* sides, with +`RS = 1/sqrt(128) = 0.088388347648`: + +``` +xh = f16( H128(x_f32 * rin_f32) * RS ) +mid = xh_f32 @ W_f32 +y = f16( H128(mid) * RS * rout_f32 ) +``` + +Escha folds its sign flips into `rin`/`rout` rather than carrying a separate +sign stage. `s_in`/`s_out` are the end-to-end fine-tune scales; MoE exports ship +them all-ones, dense exports ship real values. Both collapse into `rin`/`rout` +by `fold_scales`, which keeps the product in f32 and rounds once. + +**Verified 2026-09-02.** The above was checked, not inferred: `ref.py`'s +`reconstruct_fast` run against the committed goldens reproduces +`expected_gu_e0_k2.f16` and `expected_down_e0_k3.f16` **bit-exactly** at the +stated tile grid and packing, for both K. A single expert projection uses +**10,746 distinct fp16 values** — direct evidence for the trellis claim and for +§2. The goldens are also not synthetic: `packed_gu_e0_k2.i16` is byte-identical +to the shipped layer-0 / expert-0 `gate_up_proj.escha_code`, so G0 and G2 gate +against real model data. + +**Coverage differs between the two models**, and the dense one is the larger +surface: + +- **35B (`eschamoe`)** — escha covers routed experts only, as a *single fused* + `gate_up_proj` (K=2, `[256,128,64,32]`) plus `down_proj` (K=3, + `[256,32,128,48]`). Everything else is int8 W8A16: attention projections, + linear-attn `in_proj_qkv`/`in_proj_z`/`out_proj`, shared expert, embeddings, + `lm_head`. No bias tensors anywhere. +- **27B (`escha`)** — escha covers *every* projection: both attention families + and all three MLP legs, with `gate_proj`/`up_proj` as **separate tensors + carrying different K**: + + | tensor | layers | shape | in x out | K | + |---|---|---|---|---| + | `linear_attn.in_proj_qkv` | 48 | `[320, 640, 32]` | 5120 x 10240 | 2 | + | `linear_attn.in_proj_z` | 48 | `[320, 384, 32]` | 5120 x 6144 | 2 | + | `linear_attn.out_proj` | 48 | `[384, 320, 32]` | 6144 x 5120 | 2 | + | `self_attn.q_proj` | 16 | — | 5120 x 6144 | 2 | + | `self_attn.k_proj` | 16 | — | 5120 x 1024 | 2 | + | `self_attn.v_proj` | 16 | — | 5120 x 1024 | 2 | + | `self_attn.o_proj` | 16 | — | 6144 x 5120 | 2 | + | `mlp.gate_proj` | 64 | `[320, 1088, 32]` | 5120 x 17408 | 2 | + | `mlp.up_proj` | 64 | `[320, 1088, 48]` | 5120 x 17408 | 3 | + | `mlp.down_proj` | 64 | `[1088, 320, 48]` | 17408 x 5120 | 3 | + + That is 10 distinct projections over 64 layers (48 linear-attention, 16 full + attention) — 400 escha tensors. **The full-attention layers are escha-coded + too**, which the linear-attention-only layer 0 does not reveal; sample a + layer 3 as well as a layer 0 when validating the converter. Consequence for + dispatch: the 27B needs the 3-way `FusedQkvQ8_0` path as well as the 4-way + `FusedQkvzaQ8_0` one. + + Only `in_proj_a`/`in_proj_b` and the norms sit outside — but see §1.3 on what + `ignore` does and does not mean. + +### 1.2 `rout` carries a per-expert channel prune mask + +`rin` is a clean sign x scale vector (no zeros, tight magnitude spread). **`rout` +on `gate_up_proj` is not.** Measured on layer 0: + +| expert | `gate_up.rout` zeros | non-zero magnitude range | +|---|---|---| +| 0 | 560 / 1024 (54.7%) | 1.81 – 3.26 | +| 1 | 594 / 1024 (58.0%) | 1.99 – 3.21 | +| 2 | 142 / 1024 (13.9%) | 0.991 – 1.62 | +| 7 | 378 / 1024 (36.9%) | 0.979 – 1.02 | + +The zeros are exact, the rate varies per expert, and the masks are **not shared** +between experts (pairwise agreement ~0.47, i.e. chance). `down_proj.rout` by +contrast has **zero** zeros on every expert sampled. + +The structure is exact, not incidental. `rout` is applied last +(`y = f16(H128(mid) * RS * rout)`), so a zero hard-zeroes that output channel +for every input. For expert 0, the 560 zeros split as **280 in the gate half and +the same 280 channels in the up half** — `gate-only = 0`, `up-only = 0`. +Confirmed end-to-end: running `expert_linear` on random inputs yields 560/1024 +output channels identically zero, and after SwiGLU **280 of the 512 intermediate +channels are dead for all inputs**. + +So Escha's fine-tune leaves behind **structured, per-expert width pruning**, +carried in `rout` rather than in a mask tensor. Consequences: + +- The "signs folded into `rin`/`rout`" description is correct for `rin` and for + `down_proj.rout`, and incomplete for `gate_up.rout`, which is + sign x scale x prune-mask. +- It is exploitable — see §4.5. It is also a correctness trap: a kernel that + "optimizes away" the zero multiply without preserving exact-zero output would + change results. +- The pruned `gate_up` columns are still stored in the code stream at 2 bits + each and are never used. Physically dropping them would shrink the model but + would break the verbatim/`memcmp` property, so it is not done in this port. + +### 1.3 Four metadata traps + +**`escha_config` has two lengths — and is optional (§1.4).** When present, MoE +exports ship `[9]` += `[16, K, 2, 1, E, in, out, in_p, out_p]`; dense exports ship `[6]` += `[16, K, 2, 1, in, out]`. Fields 0 (tile = 16), 1 (K) and the trailing dims +are identified. **Fields 2 and 3 (values `2` and `1`) are not identified** — most +likely the QTIP vector dim V and a version/flag. They are asserted equal to +their observed values, never interpreted; if a future release changes them, +conversion fails loudly. + +**Trust `K`, not `bits`.** `quantization_config.layer_meta` disagrees with +itself across the two releases: the 35B records `down_proj` as +`{"bits": 3.0, "K": 3}` while the 27B records it as `{"bits": 2.0, "K": 3}`. +`bits` tracks the marketing rate on one and the true rate on the other. `K` is +consistent in both, and matches `escha_config[1]` and the code-tensor shapes. +The converter keys off `K` and uses `bits` only as a cross-check it is allowed +to fail. + +**`ignore` means "not escha-coded", not "not quantized".** Both models list +`embed_tokens` and `lm_head` in `quantization_config.ignore`, and both ship them +as `weight_int8` + `weight_scale` anyway (`int8_embedding: true`). A converter +that reads `ignore` as "keep at source precision" will go looking for f16 +tensors that do not exist. Classify on the tensor suffix actually present, not +on the ignore list. The 27B's list is shorter still (`in_proj_a`, `in_proj_b`, +`lm_head`) because its norms simply never match. + +**The 27B carries biases the base model does not have.** Every escha linear in +the dense export ships an F16 `bias` (`in_proj_qkv.bias [10240]`, +`in_proj_z.bias`, `out_proj.bias`, `mlp.{gate,up,down}_proj.bias`). Base +Qwen3.8-27B has `attention_bias: false` and no MLP bias — these are the additive +fp16 output correction Escha's end-to-end fine-tune leaves behind, applied after +the output transform per `ref.py::dense_linear`. **`hipfire-arch-qwen35`'s +arch-5 path has no bias on these projections today and must gain one.** The 35B +MoE path needs no such change. + +MTP packaging also differs: the 35B ships inline `mtp.*` tensors in the main +shards; the 27B ships a separate `mtp/` subdirectory with its own `config.json` +and `model.safetensors`. + +### 1.4 Leaf contract: three required, four optional + +Escha's own tests state the loader contract, and it is stricter and looser than +the spec first assumed. The namespace is +`CODED_LEAVES = (escha_code, escha_rin, escha_rout, escha_s_in, escha_s_out, +escha_config)`, plus `bias`. + +**Required — `escha_code`, `escha_rin`, `escha_rout`.** A coded linear missing +any transform vector must fail loudly. Their test is named +`rejects_incomplete_linear` and its docstring is explicit that this must "fail +loudly at load, not decode into noise". Our converter and loader adopt the same +rule. + +**Optional — `escha_s_in`, `escha_s_out`, `escha_config`, `bias`.** An export +produced without the end-to-end fine-tune stage ships none of them and must +still load and run (`test_dense_checkpoint_without_optional_leaves`). Note their +assertion is `bias is None`, not a zero bias — absence is a distinct state from +zero, even though the two are numerically identical once applied. + +**Unknown `escha_*` leaves are a format mismatch and must be rejected**, not +silently ignored and not allowed to fail deep inside a parameter-name error. +Their named example is `escha_rotation_theta` — evidence the format anticipates +a theta-parameterized (Givens-style) rotation variant that today's checkpoints +do not use. hipfire already has `RotationPlan::Givens` for ParoQuant, so such a +variant would not be alien, but it is **out of scope here**: if a future release +ships it, conversion must stop rather than decode the codes under the wrong +rotation. + +**Consequence for the converter.** Because `escha_config` is optional, it cannot +be the source of truth for `K`. `K` is always derivable from the code tensor's +own shape — the last dimension is `16K` — so that is the primary source, with +`escha_config[1]` and `layer_meta` used as cross-checks *when present*. This is +strictly more robust than §1.3's `bits`/`K` disagreement work-around: the shape +cannot disagree with itself. + +## 2. Why there is no lossless repack into an existing MQ codec + +MQ2 assigns each weight one of 4 levels within a linearly grouped run sharing +one scale. Escha assigns each weight any of 65536 fp16 values, at 2.0 amortized +bits, with no block scale. The alphabets are not nested and no MQ group scale +recovers the mapping — a lossless transcode into MQ2 is not merely lossy, it is +impossible. MQ8 would be near but still not lossless (256 levels vs arbitrary +fp16). Only fp16 storage is exactly lossless, at 8x the bytes. + +This is measured, not argued from the spec: decoding the shipped layer-0 / +expert-0 `gate_up_proj` yields **10,746 distinct fp16 values** across the +2048 x 1024 matrix, spanning [-3.949, 3.949]. MQ2 can express 4 per group. + +Separately, folding the Hadamards into an effective weight +`W_eff = diag(rin) . H128 . W . H128 . diag(rout) . RS^2` is exactly computable +(the Hadamards are block-diagonal), but it skips the intermediate f16 rounding +of `xh`, so it is a different numerical contract from the deployed one. Escha +flag the same deviation for their own Q8 repack. + +**Therefore:** the repack is lossless into hipfire's *container* (codes stored +byte-for-byte, `memcmp` as post-condition), and the kernels are mandatory. +"Lossless repack" and "no new kernels" cannot both hold. + +## 3. Decisions + +| decision | value | +|---|---| +| quant types | `ESCHA2T16 = 42`, `ESCHA3T16 = 43` | +| registry quant label | `escha` | +| registry model ids | `qwen3.6:35b-a3b-escha`, `qwen3.8:27b-escha` | +| rotation plan | new `RotationPlan::EschaH128` | +| Phase 1 GPU path | decode tiles to `Q8_0` resident at load; H128 kept at runtime | +| Phase 2 GPU path | fused decode+GEMV; rotations unchanged | +| model order | 35B-A3B first, 27B second | + +**Why ids 42/43.** The authoritative registry is +`crates/hipfire-quantize/src/hfq.rs` — both the `#[repr(u8)] enum QuantType` +and its `from_u8`, whose doc comment requires the two be kept in sync when a +variant is added. Everything in the crate imports `crate::hfq::QuantType`. + +On `origin/master` the taken ids are **0–22, 24, 28–41, 44, 45, 47–51**, with +23/25/26/27 documented do-not-reuse reservations. **42 and 43 are the lowest +free pair**, and 46 and 52+ are also free. + +Note what is *already merged* here, because it contradicts what the branch +survey suggests: 38/39 are `MQ2G256GL`/`MQ3G256GL` and **40/41 are +`TQ2G128`/`BQ1G128`** — the Bonsai ternary/binary types. Anyone re-deriving +this must read `hfq.rs` on master, not a `QuantType` enum in another checkout: +the `loop/gfx1151` branch carries a stale partial copy that stops at +`MFP2G32E8 = 37` and would suggest 38–41 are free. It also means the Neutrino +plan's reservation of 40/41 for `FV5G256`/`FV5B256` is stale and needs +renumbering independently of this port. + +**Why two types, not one.** `decode8_k2` reads a 16-word tile at a fixed 16-bit +stride; `decode8_k3` walks 24 words at a computed bit offset with a modular +wrap. They are structurally different, and hipfire dispatches kernels off +`(QuantType, RotationPlan)` — one type with K in `group_size` would force a +runtime K branch through every existing match arm. + +**Why `T16` and not `G256`.** The block *is* 256 weights, so `ESCHA2G256` would +look consistent with `MQ2G256`. It would also mislead: everywhere else in +hipfire `G256` means a linear run of 256 contiguous weights along a row sharing +one scale, whereas escha's 256 is a 16x16 two-dimensional tile with no block +scale. `T16` reads as "16x16 tile" and cannot be mistaken for linear grouping. + +## 4. Components + +### 4.1 `escha-ref` — CPU reference + +A Rust port of `ref.py`: `cba_decode`, `decode_tile`, `reconstruct`, `h128`, +`input_transform`, `output_transform`, `expert_linear`, `fold_scales`, `swiglu`, +`w8a16`. Pure functions, no GPU and no hipfire dependencies. + +One practical note from running it: `ref.moe_block` calls `expert_linear` +without a pre-decoded weight, so a literal port re-decodes the same expert once +per (token, slot) — 128 full tile decodes for an 8-token fixture. Memoize +`reconstruct` per (expert, projection) or the reference is unusably slow at +G4/G5 scale. + +**This is the numerical oracle for every other component.** The Phase 1 +Q8-resident build is explicitly *not* the oracle — it carries its own +quantization error. Two different artifacts, two different roles. + +### 4.2 Converter — `hipfire-quantize` + +Arch-detect on `quant_method in {escha, eschamoe}`, emitting `.hfq`. Every +tensor in the safetensors index is classified; there is no default-skip branch. + +| source | destination | +|---|---| +| `*.escha_code` | `ESCHA2T16`/`ESCHA3T16` verbatim, K from `shape[-1] / 16` (§1.4) | +| `*.escha_{rin,rout,s_in,s_out}` | folded to one f32 pair per projection | +| `*.weight_int8` + `*.weight_scale` | `Q8_0`, row scale replicated (§4.2.1) | +| `*.bias` (27B only) | F16, new arch-5 bias slots | +| norms, `A_log`, `conv1d`, `dt_bias`, `in_proj_a/b`, `mlp.gate`, `shared_expert_gate` | F16/F32, as today | +| `mtp.*` (35B inline) / `mtp/` dir (27B) | existing `.mq4-mtp` trailer | + +`escha_config`, when present, is read at both lengths (`[9]` MoE, `[6]` dense) +and cross-checked against `quantization_config.layer_meta` rather than trusted. +`K` comes from the code tensor's shape; `escha_config[1]`, `layer_meta.K` and +`layer_meta.bits` are cross-checks, and `bits` is one that is allowed to fail +(§1.3, §1.4). + +#### 4.2.1 The int8 repack is lossless only if done one specific way + +Escha's int8 is **per-output-row**: `w8 [O, K] int8` with `scale [O]` f16, and +`ref.py::w8a16` dequantizes as `f16(w8 * scale)`. hipfire's `Q8_0` is +**per-32-element block** (34 bytes per 32: 32 int8 plus one f16 scale). + +Do not recompute per-block scales from the dequantized values — that is a second +quantization and adds avoidable error. Instead **replicate the row scale into +every block of that row** and pass the int8 bytes through unchanged. The +reconstruction is then bit-identical to Escha's, at a cost of 2 bytes per 32 +elements (6.25% overhead) for scales that are all equal within a row. + +Note also that a single logical tensor can straddle safetensors shards (the 27B's +`mlp.up_proj` has its `escha_code` in shard 2 while its metadata sits in shard 1), +so the converter resolves tensors through the index, never per-file. + +### 4.3 Dispatch + +`DType::Escha2T16` / `Escha3T16` and `RotationPlan::EschaH128`. Both types need +the guard that they can never fall through to `GemvVariant::Plain` — the rule +`coverage_tests.rs:521` already enforces for `MQ4G128`, where falling through +would double-rotate. Here it would *un*-rotate, which is worse: the output is +coherent-looking text rather than a crash. + +### 4.4 Phase 1 kernels + +- `escha_decode_tiles.hip` — one-shot expansion to `Q8_0` resident at load. +- `escha_h128_in.hip` / `escha_h128_out.hip` — the two activation transforms. + The butterfly is reused from `gemv_mq4g128.hip`, which already pins the exact + `0.0883883476f`; the sign-seed stage (43, 1043) is dropped, since escha folds + its signs into `rin`/`rout`. + +Everything downstream is untouched: `GemvQ8_0`, `FusedGateUpQ8_0` and +`FusedQkvzaQ8_0` already exist, and the last is exactly the shape the 27B's +`in_proj_qkv` / `in_proj_z` need. + +**Why Q8 is a good intermediate.** The decoded weights live in the rotated +domain, where incoherence processing has already made them near-Gaussian and +outlier-free. Per-row Q8 is close to the best case for that distribution — +which is also why Escha themselves ship a Q8 repack path. + +### 4.5 Phase 2 kernels + +Fused decode+GEMV: hash decode inline, lane bit-extraction, both H128 +transforms. MoE indexed variants for K=2 `gate_up` and K=3 `down`; dense +variants for the 27B including the QKVZA shapes. Rotations unchanged from +Phase 1, so the only new thing under test is the fusion. + +The dense gate+up is the one place fusion does not come free: `gate_proj` is +K=2 and `up_proj` is K=3 (§1.1), so a fused dense gate+up kernel must either +run mixed-K or stay split. Decide by measurement, not up front. + +**Exploiting the prune mask (MoE only).** Per §1.2 a large, per-expert fraction +of `gate_up`'s output channels is identically zero — 55% on layer-0 expert 0, +14–58% across the experts sampled. What that does and does not buy: + +- **Not skippable:** the `gate_up` GEMV itself. `rout` is applied *after* the + output H128, and the Hadamard mixes all 128 channels of a block, so producing + the surviving channels still requires the full `mid` vector. Dropping GEMV + columns for pruned outputs would corrupt their block-mates. +- **Skippable:** the final `rout` multiply and SwiGLU for pruned channels, and — + the real prize — the corresponding **input rows of `down_proj`**. Those + intermediate activations are exactly zero, so for layer-0 expert 0 the K=3 + `down_proj` GEMV can skip 280 of its 512 input rows. `down_proj` is the more + expensive of the two (K=3, 3.0 bpw), so this lands on the dominant half of the + expert. + +The mask is static per expert, so it can be precomputed once at load into a +compacted row index rather than tested per token. Treat this as a Phase 2 +optimization gated on measurement, not a Phase 1 requirement — and note the +correctness trap in §1.2: pruned outputs must stay *exactly* zero, not +approximately. + +### 4.6 Loader + registry + +`qwen3.6:35b-a3b-escha` (arch 6), `qwen3.8:27b-escha` (arch 5), quant label +`escha`. + +## 5. Data flow — decode step + +Per escha linear: + +``` +x -> *rin_eff -> H128 blockwise -> *RS -> round f16 + -> Q8_0 GEMV -> mid f32 + -> H128 -> *RS -> *rout_eff -> f16 -> (+bias, dense only) +``` + +For the MoE, the router, top-k and shared expert are *intended* to be untouched +existing arch-6 code; only the two expert projections change. Three rounding +points are load-bearing and easy to lose: + +- SwiGLU runs on the **f16-rounded merged** `gate_up` output, gate first half. +- The expert combine multiplies by `f16(score)`, not the f32 score. +- **Router logits are rounded to f16 before top-k**: `ref.py` computes + `f16(x @ gate_w.T)` and only then widens to f32 to select. Selecting on + unrounded f32 logits is a different function. + +**f16 logit ties are real, and they are the one place "untouched arch-6 router" +needs checking rather than assuming.** Rounding to f16 before top-k manufactures +exact ties that f32 would not produce, and they are not rare: in an 8-token +fixture, one token has two experts on identical logits (§7 G4b). Consequences: + +- A tie *inside* the selected k is harmless — the combine is a sum over slots, + so slot order does not change the output. +- A tie *at the k / k+1 boundary* changes the selected set and therefore the + output. It is implementation-defined which expert wins, so hipfire and + `escha-ref` may legitimately disagree on such a token. +- Therefore: if hipfire's arch-6 router does not round logits to f16 before + top-k, it will diverge from Escha more often than tie-breaking alone explains. + Verify this before assuming the router is reusable as-is. +- And when reading G5: a handful of divergent positions traceable to boundary + ties is expected, not a codec bug. Attribute before investigating. + +**Fusion and orientation.** `qwen35/weights.rs:69` documents +`experts[X].gate_up: [2*moe_intermediate, hidden]`, i.e. hipfire already stores +the MoE `gate_up` **fused and out-major** — matching Escha's single fused +`gate_up_proj`. But escha's tile grid is **in-major** (`[in/16, out/16, 16K]`), +so the decode kernel transposes on the way out. The dense 27B is the mirror +problem: `gate_proj` and `up_proj` are separate tensors with **different K** +(2 and 3), so they can only reach a fused gate+up kernel after Phase 1 has +normalized both to `Q8_0`. A Phase 2 fused gate+up for the dense model would +have to be mixed-K, or stay split. + +## 6. Constraints + +**No codebook LUT.** 65536 x f16 = 128 KB; gfx1151 has 64 KB LDS +(`profiler.rs` records `lds_per_cu: 65536`). Decode inline +instead — multiply, and, xor, two f16 unpacks, one f16 add. Five ops and no +memory traffic, which is what makes Phase 2 plausible at all. + +**`lane_positions` is the trap.** It is a non-obvious permutation, and a wrong +one still yields a full-rank, plausible-looking weight matrix. Gate it directly +on golden vectors, never on end-to-end coherence. + +**Every `ESCHA2T16` site needs an `ESCHA3T16` twin.** Same shape as the Neutrino +`FV5G256`/`FV5B256` rule. `weight_backend.rs::dequant_f32` is specifically where +Bonsai lost hours to a missing arm that surfaced only at e2e as a `token_embd` +panic. + +## 7. Gates + +- **G0** — `escha-ref` bit-exact against committed goldens: `packed_gu_e0_k2` -> + `expected_gu_e0_k2` and `packed_down_e0_k3` -> `expected_down_e0_k3` exact; + `w8a16` fixture. The MoE-block fixture is **not** a bit-exact gate — see + below. + **Already demonstrated in NumPy** (§1.1), so G0 is a port-fidelity gate on the + Rust translation, not an open question about the format. The goldens are real + shipped tensors, so passing G0 means decoding the actual model correctly. +- **G1** — converter `memcmp` round-trip on code streams; zero unclassified + tensors in the index; and the §1.4 leaf contract enforced in both directions — + a checkpoint with a transform vector removed must be rejected, and one with + `s_in`/`s_out`/`config`/`bias` absent must convert cleanly. +- **G2** — GPU decode vs `escha-ref::reconstruct`, exact fp16 on every tile of a + sampled expert set, **both K**. +- **G3** — H128 kernels vs `escha-ref::h128` directly. A round-trip check + (`H128 . H128 = 128 I`) is **not** sufficient: a wrong butterfly order is also + self-inverse, so it passes while being wrong. +- **G4** — single-expert `expert_linear` on GPU vs reference, then the MoE block + against `moeblk_out`. Two things about this fixture, both measured: + + **It is a tolerance gate, not a bit-exact one.** Running `ref.moe_block` + against it with the real shipped layer-0 weights gives `max|diff| = 1.22e-4`, + `mean|diff| = 2.1e-6`, with 4752 of 16384 values differing at ULP level + against outputs whose mean magnitude is 0.0185. The golden was evidently + produced by the Metal path, not by `ref.py`, so "within fp16 rounding" is too + vague to gate on. Use **`max|diff| <= 2e-4` and `mean|diff| <= 1e-5`**, and do + not assert equality. The codec goldens in G0 *are* bit-exact; do not + generalize this tolerance to them. + + **It does not gate the router.** The fixture ships `moeblk_ids.i64` and + `moeblk_scores.f32` and injects them, bypassing selection entirely. Router + correctness needs its own check — G4b below. + +- **G4b** — router, gated separately since G4 cannot see it. Reproducing + selection from `mlp.gate.weight` on the fixture's `x` gives the **identical + top-8 set for all 8 tokens** and scores agreeing to 3e-8, confirming the + contract. Assert the set and the scores, **not the order**: token 3 orders two + experts differently because experts 65 and 43 have *exactly equal* f16 logits + (both 1.80078). See §5 on why that is benign here and when it is not. +- **G5** — e2e coherence, then KLD on a fixed wikitext slice. Not on the model's + own output: for ds4 that scored 8x better on the median and was optimistic. + + **The reference is `escha-ref` on CPU, not any Escha runtime.** None of + Escha's three runtimes can run on this box: `escha-mlx` is Metal / Apple + Silicon, the `escha` wheel is CUDA (sm_80–sm_120), and ZML requires an NVIDIA + driver. There is no cross-checking against their engine on gfx1151, and a gate + that cannot be executed proves nothing. This is not a downgrade: `ref.py` + declares itself "the semantic contract for every Metal kernel in this package" + and is itself gated on the goldens, so agreeing with `escha-ref` *is* + agreeing with their runtime — and it is exact rather than cross-machine. + Cost is CPU time; budget a few hundred positions, not thousands. + + Run a second KLD against the **bf16 parent** (`Qwen/Qwen3.6-35B-A3B`) as well. + That one answers a different and independently useful question — whether + Escha's 2-bit delivers what they claim — and it is the number to compare + against the existing `qwen3.6:35b-a3b-mq2` measurement. +- **G6** (Phase 2) — fused GEMV compared against the Phase 1 path at Q8 + precision, then KLD against `escha-ref`. + +## 8. Error handling + +Refuse rather than guess: + +- unknown `quant_method`, or `format_version != "2.0"` — hard error +- `escha_config` disagreeing with `layer_meta` — hard error. This is what + catches a shape assumption silently breaking on a future Escha release. +- any tensor in the safetensors index left unclassified — hard error, not skip +- a coded linear missing `escha_code`, `escha_rin` or `escha_rout` — hard error + ("incomplete escha linear"), never a partial decode (§1.4) +- an `escha_*` leaf outside the known six — hard error naming the leaf. This is + the guard that stops a future `escha_rotation_theta` export from being decoded + under the wrong rotation (§1.4). +- a missing `ESCHA3T16` dispatch arm must never resolve to `Plain` + +Conversely, absence of `escha_s_in`, `escha_s_out`, `escha_config` or `bias` is +**not** an error — those are optional by contract and the fold/decode paths must +run without them. + +## 9. Risks and expectations + +**Phase 1 perf will be bad, and that is the expected outcome.** Worked through +at `Q8_0`'s 34 bytes per 32 elements: + +| | 35B experts | 35B non-expert | 35B total | +|---|---|---|---| +| escha native | 9.4 GB | 2.3 GB | ~11.7 GB (published 12.3) | +| Phase 1 `Q8_0` | 34.2 GB | 2.5 GB | **~36.7 GB** | + +MoE reads 8 of 256 experts per token, so per-token expert traffic goes from +**0.29 GB to 1.07 GB — 3.6x**. It will lose to `qwen3.6:35b-a3b-mq2` on tok/s. +The Phase 1 deliverable is correctness plus a servable artifact; speed is +Phase 2's job. 36.7 GB is comfortable on 128 GB but is not free — check it +against the standing heap baseline before assuming headroom. + +The same arithmetic applied to the 27B gives ~10 GB native against a published +10.15 GB, and ~28 GB at `Q8_0`. That both reconstructions land within a few +percent of the published sizes is independent confirmation that the format model +in §1.1 is right — a wrong bpw or a missed tensor class would not close. + +**The 27B is not benchmarked before Phase 2.** At ~28 GB dense, re-read every +token, a decode-at-load build is not worth timing. It still gets built and +gated for correctness through G4 — that is how the dense code path, the +different-K gate/up split and the new bias slots get exercised — but no +tok/s number is reported for it until fused kernels exist. + +**The 27B needs a change outside the codec.** Its escha linears carry fp16 +biases the base architecture does not have (§1.3), so `hipfire-arch-qwen35`'s +arch-5 path gains bias slots on `in_proj_qkv`, `in_proj_z`, `out_proj` and the +three MLP projections. This is the only work in the port that touches an +existing architecture, and it must not perturb the existing `qwen3.6:27b` SKUs +that share that path — those have no bias tensors and must keep taking the +no-bias branch. + +**Kernel cache hygiene.** Mixed-toolchain blobs on gfx1151 produce attractor +garbage that will look exactly like a codec bug. Single-toolchain cache, hipcc +on PATH. See `hipfire_kernel_rebuild_gfx1151`. + +**Golden coverage is thin** — one expert of one layer per K, plus one small MoE +block. Broader expert sampling is checked against `escha-ref`, not against +goldens. + +**Format drift.** `format_version` is pinned and asserted; a future Escha +release that changes tile size or codebook constants must fail loudly at +conversion, not decode into garbage. + +## 10. Phase 1 results (measured) + +Artifact: `/data/hipfire-models/escha-35b.hfq`, sha256 +`bd186b37037ee6c1cb58ce6a5c053b785e9ba30b7f1de04fda7ef4f4f06e3010`, 12.34 GB, +arch_id 6, registered as `qwen3.6:35b-a3b-escha`. Host gfx1151, ROCm 7.2.2. + +### 10.1 Quality (G5) + +Reproduce with `scripts/escha-kld.sh`. + +| | value | +|---|---| +| corpus | `benchmarks/quality-baselines/slice/wikitext2-1024s-2048ctx.txt`, n_ctx 384, 6 chunks | +| scored positions | 1146, teacher-forced | +| **mean KLD vs weight-exact escha** | **0.00276 nats** (95% CI 0.0019–0.0039) | +| p99 KLD | 0.054 nats | +| reference PPL / NLL | 7.6651 / 2.036678 | +| candidate PPL / NLL | 7.6585 / 2.035821 | +| negative control (exact vs itself) | **0.000000** nats as printed (actual 2.1341e-10, p99 6.03e-09) | +| KLD vs the bf16 parent | **not run** — see §10.2 | + +Near-zero, as §7 predicted, and it is the `Q8_0` intermediate rather than the +codec: the reference arm is the same forward with the experts stored +weight-exactly (`HIPFIRE_ESCHA_EXPERT_STORE=f16`, bit-identical to +`escha_ref::reconstruct`), so the only thing that differs between the arms is +the 8-bit re-quantisation. The p99 sits ~20x the mean, consistent with the +boundary-tie behaviour in §5 concentrating the error on a handful of positions +rather than spreading it; that shape is expected and is not a codec bug. + +Three things about the method are load-bearing and were each got wrong once +before being fixed: + +- **Teacher forcing is structural, not a flag.** `build_kld_ref_native` writes + the token stream into the HFKLDR file and `eval_hipfire` reads the tokens + from that file, so both arms are scored on one identical committed stream by + construction. Nothing is scored on a model's own greedy output. +- **`--kv-mode f32` matters more than it looks.** The reference builder uses an + unquantised F32 KV cache; `eval_hipfire` defaults to `asym3`. Leaving that + default in place gives **0.018357** nats on the identical reference — 6.5x + the real figure, and almost all of it is KV quantisation, not the codec. +- **The negative control is the thing that makes the number attributable.** + Scoring the exact arm against its own reference prints 0.000000. The + underlying float is not bit-zero — it is 2.1341004702939135e-10 (p99 + 6.029504362788427e-09) — and re-running it against the same reference + reproduces both figures BIT-FOR-BIT. That reproducibility is the point: a + nondeterminism floor would not repeat. The residue is the fixed difference + between two programs computing the same forward (`build_kld_ref_native` + writes the reference, `eval_hipfire` scores it), and it is ~1.3e7 times + below the 0.0027576 the production arm reports. `scripts/escha-kld.sh` now + asserts both halves of that — the control rounds to 0.000000 at the printed + precision, and it is at least 1e4x below the headline number — rather than + leaving "must print exactly 0.000000" to a human reading stdout. + +The reference could not be `escha_ref` driving a CPU forward: `escha_ref` is a +block-level oracle (codec, H128, `expert_linear`, `swiglu`) and this repo has +no CPU transformer, so a CPU reference would have meant hand-writing a +40-layer hybrid DeltaNet MoE forward and then trusting it more than the thing +under test. Storing escha's own decoded fp16 weights and reusing the gated GPU +forward keeps `escha_ref` as the authority for every step it actually defines +(G2 and G3 gate the decode and the H128 pair bit-exact against it). + +### 10.2 The bf16 parent comparison was not run + +`Qwen/Qwen3.6-35B-A3B` is not on this box in any form. `/data/hipfire-models` +has no safetensors copy, and the only cached artifact of the parent is +`unsloth/Qwen3.6-35B-A3B-GGUF: Qwen3.6-35B-A3B-UD-Q4_K_XL.gguf` — a 4-bit +quant, which is a peer rather than a reference. Fetching the parent is ~70 GB. +Skipped deliberately and left open rather than substituted with the Q4 GGUF, +which would have produced a number that reads like the answer and is not. + +Consequence: **the "does Escha's 2-bit deliver" question is still open**, and +there is no comparable figure against `qwen3.6:35b-a3b-mq2` on quality. §10.1 +prices the port, not the codec. + +### 10.3 Speed and memory + +All figures re-measured on the branch head (gfx1151, ROCm 7.2.2) with +`escha_prefill_bench` and `scripts/escha-gtt-probe.sh`. An earlier revision of +this section reported the state before `6547c78b6` (contiguous expert buffers) +and `06ab4db8e` (reachable batched prefill) and understated prefill by 4x; the +numbers below supersede it. + +| | escha (Phase 1) | `qwen3.6:35b-a3b-mq4r` | +|---|---|---| +| decode | **40.0–40.2 tok/s** | 71.8 tok/s (MTP engaged, tau 1.52) | +| prefill, n=512 | **161.9 tok/s** | — | +| prefill, n=2048 | **165.0 tok/s** | — | +| prefill, short prompt | — | 289.6 tok/s | +| file on disk | 12.34 GB | 18.7 GB | +| resident (GTT) | **37.58 GB** | — | + +Batched prefill IS admissible and IS taken: `escha_routed_prefill_indexed` +runs the routed half over `n_tokens * k` slots in one launch per (layer, +projection, side), which the bench confirms as 160 H128 launches for a +512-token prompt (a per-token fallback would be 81 920). Prefill is therefore +~4.1x decode, not equal to it. + +The remaining prefill gap to `mq4r` is ~1.8x, and §10.5 says what closing it +would cost. The decode row is still not like-for-like — `mq4r` has its MTP +sidecar attached and escha has no speculator wired (and, per the MTP refusal +added on this branch, cannot have one until the H128 pair is taught to that +forward). + +`qwen3.6:35b-a3b-mq2` remains the honest speed comparison and is still not on +this box, so the table uses `mq4r` from the same family; mq2 is 11.6 GB +against mq4r's 18.7 GB, reads strictly fewer expert bytes per token and takes +the same indexed GPU-top-K path, so it is at least as fast as the column +shown. + +**Resident memory is 37.58 GB**, measured as an amdgpu GTT delta over a +3.36 GB idle baseline (40.94 GB peak). 34.2 GB of that is the Q8_0 routed +experts and ~3.3 GB is everything else. + +It was **67.9 GB** until `6547c78b6`, and the cause was allocation +granularity, not the codec or the carrier. While each of the 20 480 per-expert +projections was its own device allocation, the HIP allocator's 2 MiB granule +rounded up every one of them: + +| | logical | rounded to 2 MiB granules | +|---|---|---| +| gate_up `[1024, 2048]` Q8_0 | 2.125 MiB | 4 MiB | +| down `[2048, 512]` Q8_0 | 1.0625 MiB | 2 MiB | +| x 10240 experts | 31.9 GiB (34.2 GB) | 60 GiB (64.4 GB) | + +64.4 GB of granules plus ~3.4 GB of everything else is 67.8 GB against 67.9 GB +observed — the arithmetic closed to 0.1%, which is why it was attributed to +granularity rather than to a leak, and why packing the experts into one buffer +per (layer, projection) recovered the predicted ~30 GB exactly. The logical +34.2 GB now IS the resident cost. + +One consequence survives the fix: the weight-exact **F16** expert store is no +longer free. It was free at 67.9 GB, because F16's 4 MiB / 2 MiB projections +were exactly what the rounded Q8_0 already occupied. Packed, F16 is a genuine +2x on the expert half (~68 GB resident). That still fits on this box and is +what the §10.1 reference arm uses, but it is now a real cost rather than an +accounting artifact, and F32 (129 GB) still does not fit. + +### 10.4 Coherence + +Verbatim, greedy (`temperature: 0.0`), via `scripts/_coherence_runner.py`: + + prompt: /no_think What is the capital of France? Answer in one word. + reasoning channel: (a short numbered plan, ending) 4. Final Output + Generation: - Paris - Matches all constraints. + answer channel: Paris + + prompt: /no_think Write three sentences describing a harbour at dawn. + Pale gold light spills across the still water, painting the masts of + sleeping boats in soft amber hues. A gentle mist clings to the wooden + docks while the distant cry of a gull breaks the quiet morning air. + Slowly, the first fishing trawler rumbles to life, its engine echoing + through the quiet cove as the sun finally crests the horizon. + +Attractor detector clean (unique-token ratio 0.617, max token frequency 0.055 +over the first 128 tokens; no 3-gram flag). First run was already coherent — +no kernel-cache garbage, single-toolchain cache with hipcc on PATH. + +One thing to know rather than to fix: the model ignores `/no_think` as a mode +switch and reasons anyway, so a short `max_tokens` truncates mid-think and the +daemon reports `open think span at end of generation (validation)`. That is a +budget artifact, not a decode fault — the same prompt with headroom closes the +span and emits a normal answer. + +### 10.5 Limitations that remain open + +Two, both measured, both required reading before this is treated as finished. + +**(a) The routed half does not amortise in prefill.** Batched prefill amortises +the *launch* cost of the H128 transforms and the *activation* traffic, and that +is what took prefill from decode speed to ~163 tok/s. It does not amortise the +expert **weight** traffic: each token's top-8 experts are read for that token, +so a batch of `n` tokens still moves ~`n` x 34.2 GB / (n_layers x tokens) +worth of expert bytes — 34.2 GB of expert traffic per full pass over the batch, +the same total as the per-token route. That is why escha stays ~1.8x behind +`mq4r`'s 289.6 tok/s at prefill despite the batched body. + +Closing it needs a Q8_0 grouped GEMM over **sorted expert groups**: sort the +`n * k` slots by expert, then run one GEMM per expert over all the tokens that +chose it, so each expert's weights are read once per batch instead of once per +token. That is the same Path-2 scatter/grouped structure the MQ dtypes already +use, and it is a real change rather than a port: grouping changes the order in +which a token's `k` expert contributions accumulate, so the result would no +longer be bit-identical to the per-token route. G4's +`batched prefill route vs per-token indexed route: 0 differing floats` — an +equality claim, not a tolerance — would have to be restated as a bound, and +the argument for the new bound would have to be made from scratch. Do not land +that work while leaving G4 asserting equality; the gate would go red and the +temptation would be to loosen it. + +That grouped arm is also what would spring the prefill hole this branch closed +with `check_moe_prefill_supported`: it is the first generic Q8_0 routed prefill +arm, and without the guard an escha layer with `HIPFIRE_ESCHA_INDEXED=0` would +silently take it and drop both Hadamard transforms. + +**(b) Prefill and decode select different experts, at 2.96% per expert slot.** +The design-time estimate in §5 was ~0.42% of routing decisions straddling an +f16 boundary. The measured figure is **8x that**, and — more importantly — the +dominant term is not the one the estimate was about. + +`escha_prefill_batch_gate` records the `topk_indices` both routes actually +indexed with. Measured at n=64: + +| | rate | +|---|---| +| expert-SLOT divergence, whole stack | **2.96–3.13%** | +| (token, layer) SET divergence, whole stack | 24.1% | +| (token, layer) SET divergence, layer 0 | 0.00% | + +Layer 0 is 0% because both routes feed it the same embedding, and it is the +only layer that cannot have compounded. The whole-stack rate is higher because +a flip at layer L changes that token's hidden state for every later layer; the +per-layer series rises from 0% to a ~25-35% plateau, which is compounding, not +a defect. + +The dominant contribution is an **f16 activation downcast**, not accumulation +reordering. The batched dense half converts its F32 activations to F16 for the +WMMA GEMMs where the per-token route runs F32 GEMVs; that downcast is a much +coarser perturbation of the router input than the last-bit differences of a +reordered f32 sum, and it is what pushes decisions across the f16 rounding +boundary that `router_logits_round_f16_rne` then quantises them to. Attributing +this to "accumulation reordering" — as the 0.42% estimate did — is why the +estimate was 8x low. + +Consequences: escha's prefill and decode are NOT interchangeable for anything +that assumes identical routing (KV-cache reuse across a prefill/decode boundary +is fine; a cache of routed-expert outputs would not be). And the final-token +logit delta between the routes is `max 4.393e-1 / mean 7.160e-2` — two orders +of magnitude above what reordering alone gives, dominated by this divergence. +The argmax is stable, which is what G6 asserts, but the logits are not close. + +Neither limitation is a blocker for Phase 1, whose claim is correctness plus a +servable artifact. Both are stated here because §10.3 reads like a success and +a reader who stops there would carry away two wrong beliefs. + +### 10.6 Correctness gates (measured) + +The six gates, their commands, and the results on the branch head. All are +re-runnable; G1-G4b need no model beyond the checkpoint and the committed +fixtures. Registered in `docs/VALIDATION.md` and `scripts/gates.sh`. + +| Gate | What it asserts | Result | +|---|---|---| +| **G1** | every `escha_code` tensor in the `.hfq` is byte-identical to the source safetensors | **80/80 byte-identical** (40 layers x 2 projections; count asserted against `model.safetensors.index.json`, not hardcoded) | +| **G2** | GPU tile decode == `escha_ref::reconstruct` exactly, in fp16 | **bit-exact**, 0 mismatched, at both golden shapes and at production shapes of 89 128 960 elements (K=2 and K=3) | +| **G3** | the H128 pair == `escha_ref`, all launch forms | **bit-exact**, 0 mismatched: `h128_in`, `h128_out`, batched broadcast / per-slot / grouped input mappings, `out_batched`, and swiglu | +| **G4b** | arch-6 router selects the same experts as escha's reference routing | **0/8 tokens with a differing top-8 set** | +| **G4** | the whole MoE block against escha's `moeblk_out.f16` golden | F32 (weight-exact) arm **max 1.828e-4 / mean 9.673e-6**; Q8_0 (production) arm **max 2.633e-4 / mean 3.027e-5**, against a golden mean magnitude of 1.85e-2. Indexed route vs host route: **0 differing floats**. Batched prefill route vs per-token indexed route: **0 differing floats**. Q8_0 re-quantisation cost, arm2 - arm1: max 1.687e-4 / mean 2.864e-5 | +| **G5** | KLD against the weight-exact escha arm on a fixed teacher-forced corpus | **0.0027576 nats** (95% CI 0.0019491-0.0038610), p99 0.054, **PPL 7.6585**; negative control **2.13e-10 (prints 0.000000)** and bit-reproducible across runs — see §10.1 | +| **G6** | batched prefill vs the per-token route, whole model | **argmax stable** (135744 on both arms); max\|delta\| 4.393e-1, mean\|delta\| 7.160e-2, 0 non-finite logits. Both deltas are now asserted, not printed — see §10.5(b) for why they are this large | + +The two `0 differing floats` rows in G4 are EQUALITY claims, not tolerances, +and are only true because every kernel in the escha routed pipeline is purely +slot-parallel: slot `s` performs the same FLOPs in the same order whether the +launch carried 8 slots or 2048. §10.5(a) explains what would break them. + +G4's F32 arm is the codec+transform accuracy; its Q8_0 arm is that plus the +8-bit expert store. The gap between them (max 1.687e-4) is the price of the +Q8_0 carrier at block level, and G5's 0.0027576 nats is the same price +end-to-end over 40 layers. + +## 11. Out of scope + +- Reproducing Escha's quantization or recovery fine-tune. We consume their + weights; we do not build a quantizer. +- The vision tower. Declared in both configs, no weights shipped, in the + quantization `ignore` list. Text-only. +- Transcoding escha into MQ formats. Established impossible-to-do-losslessly in + §2; a lossy MQ transcode would discard exactly the quality that motivates the + port. +- MTP-driven speculative decode. The head is carried through conversion so it is + available, but wiring it is separate work. + +## 12. References + +- `EschaLabs/escha-mlx` (Apache-2.0) — `escha_mlx/ref.py` is the format + contract; `tests/data/codec/` holds the golden vectors. +- `EschaLabs/Qwen3.6-35B-A3B-Escha-W2`, `EschaLabs/Qwen3.8-27B-Escha-W2`. +- `docs/architecture-ids.md` — arch 5 and 6 are `hipfire-arch-qwen35`. +- `kernels/src/gemv_mq4g128.hip` — existing 128-point FWHT butterfly. diff --git a/docs/quant-formats/qt-register.txt b/docs/quant-formats/qt-register.txt index 67bdf4fd0c..11304a0852 100644 --- a/docs/quant-formats/qt-register.txt +++ b/docs/quant-formats/qt-register.txt @@ -76,6 +76,8 @@ 39 MQ3G256GL passthrough 40 TQ2G128 passthrough PR #597 Bonsai ternary 41 BQ1G128 passthrough PR #597 Bonsai 1-bit; won the qt=41 adjudication 2026-08-18 +42 ESCHA2T16 arch-loaded Escha-W2 K=2 trellis; qwen35 escha.rs decodes it on device to Q8_0/F16 at load, so it never reaches a GEMV as itself +43 ESCHA3T16 arch-loaded Escha-W2 K=3 trellis; same load path as qt=42 44 MQ4G256V2 passthrough 45 MQ4CG256 passthrough 47 MQ6G256V2 passthrough diff --git a/kernels/src/escha_bare_to_outmajor.hip b/kernels/src/escha_bare_to_outmajor.hip new file mode 100644 index 0000000000..6ff41fc22b --- /dev/null +++ b/kernels/src/escha_bare_to_outmajor.hip @@ -0,0 +1,117 @@ +// SPDX-License-Identifier: Apache-2.0 +// Copyright (c) 2026 Kaden Schutt +// hipfire — see LICENSE and NOTICE in the project root. +// +// Escha-W2 load path, step 3 of 3 (Task 10). `escha_decode_tiles` writes the +// decoded weight BARE and IN-MAJOR — row-major `[in_features, out_features]` +// fp16 — because escha's code tile grid is in-major (`[in/16, out/16]`). +// hipfire's expert slots are OUT-major: `experts[X].gate_up` is +// `[2*moe_intermediate, hidden]` (see hipfire-arch-qwen35 weights.rs:69), and +// every hipfire GEMV walks K contiguously along a row. So the transpose +// happens HERE, folded into the store, and it is the single place in the +// port where a wrong orientation would still produce a full-rank, +// plausible-looking weight matrix. Do not "simplify" the indexing. +// +// Both entry points read `bare[i * oc + o]` (in-major) and write element +// `(o, i)` of an out-major `[oc, ic]` matrix. + +#include +#include + +// Q8_0 target layout (hipfire's, see kernels/src/gemv_q8_0.hip): per output +// row `o`, `ic/32` consecutive blocks of 34 bytes — fp16 scale at +0, then 32 +// signed int8. Row stride is therefore `(ic/32) * 34` bytes and there is no +// padding. +// +// One block per (o, blk) pair, 32 threads; thread `t` owns element +// `i = blk*32 + t`. The 32 reads are strided by `oc` fp16 in `bare` — that is +// inherent to transposing, and this runs ONCE PER EXPERT AT LOAD TIME, not +// per token, so it is not on any hot path. +// +// Scale rule matches llama.cpp/hipfire Q8_0: `d = amax/127` rounded to fp16, +// values `roundf(v/d)` clamped to [-127, 127]. The clamp is load-bearing: +// when `f16(amax/127) < amax/127` the extreme element rounds to +/-128, which +// would wrap in int8. +extern "C" __global__ void escha_bare_to_q8_0( + const __half* __restrict__ bare, // [ic, oc] in-major + unsigned char* __restrict__ out, // [oc][ic/32][34] + int ic, int oc) { + const int blocks_per_row = ic >> 5; + const int o = blockIdx.x / blocks_per_row; + const int blk = blockIdx.x - o * blocks_per_row; + const int t = threadIdx.x; + const int i = (blk << 5) + t; + + const float v = __half2float(bare[(size_t)i * oc + o]); + + float amax = fabsf(v); + #pragma unroll + for (int off = 16; off > 0; off >>= 1) { + amax = fmaxf(amax, __shfl_xor(amax, off, 32)); + } + + const __half dh = __float2half(amax / 127.0f); + const float d = __half2float(dh); + int q = 0; + if (d != 0.0f) { + float r = roundf(v / d); + q = (int)fminf(fmaxf(r, -127.0f), 127.0f); + } + + unsigned char* blkp = out + ((size_t)o * blocks_per_row + blk) * 34; + if (t == 0) { + unsigned short sb = __half_as_ushort(dh); + blkp[0] = (unsigned char)(sb & 0xFF); + blkp[1] = (unsigned char)(sb >> 8); + } + blkp[2 + t] = (unsigned char)(signed char)q; +} + +// F32 out-major transpose — the bit-exact control arm. `escha_bare_to_q8_0` +// re-quantises the decoded weight to 8 bits, which costs real accuracy (see +// the G4 gate); this entry stores the decoded fp16 value widened to f32 with +// no loss at all, so a caller can measure the re-quantisation cost against a +// weight-exact baseline instead of guessing at it. 4 bytes/weight, so it is a +// diagnostic/small-layer path, not a way to run the 35B model. +extern "C" __global__ void escha_bare_to_f32( + const __half* __restrict__ bare, // [ic, oc] in-major + float* __restrict__ out, // [oc, ic] out-major + int ic, int oc) { + const int i = blockIdx.x * blockDim.x + threadIdx.x; + const int o = blockIdx.y; + if (i >= ic) return; + out[(size_t)o * ic + i] = __half2float(bare[(size_t)i * oc + o]); +} + +// F16 out-major transpose — the weight-exact arm that fits a whole model. +// +// This is PURE DATA MOVEMENT: `escha_decode_tiles` already produced fp16, so +// the stored value is bit-identical to `escha_ref::reconstruct`'s output. F32 +// (above) is exactly as exact and twice the bytes, which is why it is only a +// small-layer diagnostic; this one is what a model-scale weight-exact +// reference arm has to use. +// +// The 35B's A3B shapes make the difference stark, and it is an allocator +// effect rather than an arithmetic one. Per expert the two projections are +// gate_up `[1024, 2048]` and down `[2048, 512]`: +// +// store gate_up down per-expert x10240 experts +// Q8_0 2.125 MiB 1.0625 MiB 3.1875 MiB 31.9 GiB (34.2 GB) +// F16 4 MiB 2 MiB 6 MiB 60 GiB (64.4 GB) +// F32 8 MiB 4 MiB 12 MiB 120 GiB (129 GB) +// +// but each buffer is a separate allocation that the HIP allocator rounds up +// to a 2 MiB granule, so the Q8_0 arm's 2.125/1.0625 MiB buffers occupy 4 and +// 2 MiB anyway — 60 GiB resident, the SAME as F16. Measured on gfx1151: +// 67.9 GB of GTT for the Q8_0 model against 64.4 GB of expert granules plus +// ~3.4 GB of everything else. So F16 is weight-exact at no resident cost over +// production, whereas F32 would be 2x and would not fit. +extern "C" __global__ void escha_bare_to_f16( + const __half* __restrict__ bare, // [ic, oc] in-major + __half* __restrict__ out, // [oc, ic] out-major + int ic, int oc) { + const int i = blockIdx.x * blockDim.x + threadIdx.x; + const int o = blockIdx.y; + if (i >= ic) return; + out[(size_t)o * ic + i] = bare[(size_t)i * oc + o]; +} diff --git a/kernels/src/escha_decode_tiles.hip b/kernels/src/escha_decode_tiles.hip new file mode 100644 index 0000000000..0b811356d5 --- /dev/null +++ b/kernels/src/escha_decode_tiles.hip @@ -0,0 +1,105 @@ +// Escha-W2 one-shot tile decode -> bare fp16 (Phase 1/2). +// +// Reads the verbatim int16 code stream and decodes 16x16 trellis tiles, +// writing raw __half values into `bare`, a row-major [in_features, +// out_features] buffer — the SAME in-major orientation Escha's packed code +// already uses. There is no quantisation and no transpose: this kernel is a +// pure decode, not a requantiser. (A later stage owns any Q8_0/HFQ +// requantisation and any transpose to hipfire's preferred layout.) +// +// The codebook is computed inline. A 65536-entry fp16 LUT would be 128 KB and +// gfx1151 has 64 KB LDS, so there is no table anywhere in this file. +// +// DELIBERATE DUPLICATION: hipfire-quantize/src/escha_ref.rs implements this +// same lane maths in Rust. That is the G2 gate — this kernel is asserted +// bit-exact against it. Generating either from the other would make G2 +// circular. Do not deduplicate. + +#include +#include + +__device__ __forceinline__ __half escha_cba(unsigned short state) { + unsigned int r = ((unsigned int)state * 0xCBAC1FEDu) & 0x8FFF8FFFu; + r ^= 0x3B603B60u; + __half lo = __ushort_as_half((unsigned short)(r & 0xFFFFu)); + __half hi = __ushort_as_half((unsigned short)(r >> 16)); + return __hadd(lo, hi); // fp16 RNE add — matches the reference exactly +} + +// Load tile word `i` (a little-endian pair of int16s) directly from the code +// stream. Two of these per lane replace what used to be a 16- or 24-entry +// `unsigned int words[]` staged in registers and indexed at runtime by `i`. +// That runtime-indexed array measurably cost ~90 v_cndmask/94 v_cmp select +// instructions per invocation (compiler-generated select cascades to keep a +// dynamically-indexed array out of scratch memory) even though it never +// actually spilled to scratch (private_segment_fixed_size was 0). Loading +// the two needed words directly turns that ALU-bound cascade into two +// small global loads that land in the same one or two cache lines every +// other lane in the block is already touching. +__device__ __forceinline__ unsigned int escha_load_word( + const short* __restrict__ src, int i) { + return ((unsigned int)(unsigned short)src[2 * i]) | + (((unsigned int)(unsigned short)src[2 * i + 1]) << 16); +} + +__device__ __forceinline__ void escha_decode8_k2( + const short* __restrict__ src, int lane, unsigned short* out) { + int t_off = lane * 8; + int i1 = t_off >> 4; + int i0 = (i1 + 15) & 15; + unsigned int w0 = escha_load_word(src, i0); + unsigned int w1 = escha_load_word(src, i1); + unsigned long long merged = ((unsigned long long)w0 << 32) | w1; + int shift = ((~t_off) & 8) << 1; // 16 for even lanes, 0 for odd + unsigned int v = (unsigned int)((merged >> shift) & 0xFFFFFFFFull); + #pragma unroll + for (int j = 0; j < 8; ++j) out[j] = (unsigned short)(v >> (2 * (7 - j))); +} + +__device__ __forceinline__ void escha_decode8_k3( + const short* __restrict__ src, int lane, unsigned short* out) { + const int BITS = 3; + int t_off = lane * 8; + int b1 = (t_off + 257) * BITS; + int b0 = b1 - 16; + int b2 = b1 + BITS * 7; + int i0 = b0 >> 5; + int i2 = (b2 - 1) >> 5; + int s2 = ((i2 + 1) << 5) - b2; + unsigned int w0 = escha_load_word(src, i0 % 24); + unsigned int w2 = escha_load_word(src, i2 % 24); + unsigned long long merged = ((unsigned long long)w0 << 32) | w2; + unsigned int w7 = (unsigned int)((merged >> s2) & 0xFFFFFFFFull); + unsigned int w3 = (unsigned int)((merged >> (s2 + BITS * 4)) & 0xFFFFFFFFull); + out[0] = (unsigned short)(w3 >> 9); out[1] = (unsigned short)(w3 >> 6); + out[2] = (unsigned short)(w3 >> 3); out[3] = (unsigned short)(w3); + out[4] = (unsigned short)(w7 >> 9); out[5] = (unsigned short)(w7 >> 6); + out[6] = (unsigned short)(w7 >> 3); out[7] = (unsigned short)(w7); +} + +// One block per tile. 32 lanes, 8 values each = the 256 tile slots. +extern "C" __global__ void escha_decode_tiles( + const short* __restrict__ code, // [in/16, out/16, 16K] + __half* __restrict__ bare, // [in, out] fp16 scratch + int in_features, int out_features, int K) { + int tile = blockIdx.x; + int tn = out_features / 16; + int kt = tile / tn, nt = tile % tn; + int lane = threadIdx.x; + + const short* src = code + (size_t)tile * 16 * K; + + unsigned short st[8]; + if (K == 2) escha_decode8_k2(src, lane, st); + else escha_decode8_k3(src, lane, st); + + int l0 = lane & ~4; + int c_off = (lane >> 2) & 1; + #pragma unroll + for (int j = 0; j < 8; ++j) { + int fi = j >> 1; + int row = (lane & 3) * 2 + (j & 1) + (fi & 1) * 8; + int col = 2 * ((l0 >> 3) + (j >= 4 ? 4 : 0)) + c_off; + bare[(size_t)(kt * 16 + row) * out_features + (nt * 16 + col)] = escha_cba(st[j]); + } +} diff --git a/kernels/src/escha_h128.hip b/kernels/src/escha_h128.hip new file mode 100644 index 0000000000..2cf8240a53 --- /dev/null +++ b/kernels/src/escha_h128.hip @@ -0,0 +1,272 @@ +// Escha-W2 activation transforms. +// in : xh = f16( H128(x * rin) * RS ) +// out: y = f16( H128(mid) * RS * rout ) +// +// H128 is the UNNORMALISED 128-point Walsh-Hadamard in Sylvester (natural) +// order. Escha folds its sign flips into rin/rout, so unlike hipfire's +// gemv_mq4g128 there is no sign-seed stage here. +// +// A pruned output channel (rout == 0) must come out EXACTLY zero — do not +// reorder the final multiply in a way that could produce -0.0 or a denormal. +// +// These kernels run on every token, on both sides of every escha matmul — +// the hot decode path, unlike escha_decode_tiles (once per expert at load +// time) — so the butterfly is parallel across all 128 threads of the block, +// not a single thread doing all 7 stages serially. See h128_parallel below. + +#include +#include + +#define ESCHA_RS 0.0883883476f // 1/sqrt(128) + +// NOT a general "__float2half(-0.0f) is broken on this toolchain" bug — a +// bare literal -0.0f converts correctly (confirmed both as an immediate and +// loaded from device memory: 0x8000, correct). The trigger is narrower: on +// this ROCm 7.2.2/gfx1151 build, __float2half loses the sign bit only when +// its argument is a RUNTIME MULTIPLY-PRODUCED exact zero — the shape this +// kernel's own `v * RS` / `v * RS * rout` produces when a pruned (`rout == +// 0`) channel meets a negative `v`. A reviewer independently reproduced this +// standalone against `/opt/rocm-7.2.2/bin/hipcc --offload-arch=gfx1151`: +// +// __float2half(-0.0f) = 0x8000 (literal, CORRECT) +// __float2half(-0.0f) = 0x8000 (loaded from memory, CORRECT) +// __float2half(-1e-9f) = 0x8000 (CORRECT) +// __float2half(-1e-30f) = 0x8000 (CORRECT) +// __float2half(a[i]*c[i]) where product == -0.0 -> 0x0000 (WRONG, want 0x8000) +// __float2half(a[i]*CONST*c[i]) (this kernel's exact shape) -> 0x0000 (WRONG) +// float pre = a*b*c; __float2half(pre) -> 0x0000 (WRONG) +// +// Critically, a tiny NONZERO underflowing product computed the identical way +// DOES preserve sign (-1e-10 * 1e-10 -> 0x8000, correct) — so the defect is +// scoped to exact-zero products, not underflow in general. That's why the +// `v == 0.0f` guard below is correctly targeted at exact zero and must NOT +// be broadened to "small magnitude" or any other range. +// +// Auditor's test, if you're checking another __float2half call site for the +// same bug: converting a bare -0.0f literal will NOT reproduce it and will +// wrongly look safe. You must reproduce a -0.0 as the runtime result of a +// multiply (or chain of multiplies) feeding directly into __float2half. +// +// This is the ONLY divergence from escha_ref across the whole G3 sweep — +// every non-zero value, including values that underflow fp16 (e.g. 1e-10, +// 1e-30, either sign), round correctly. escha_ref's `f16_rne` uses +// `half::f16::from_f32`, which does preserve the sign of an exact zero (a +// pruned `rout` channel multiplying a negative pre-scale value yields -0.0 +// there), so the G3 gate requires bit-exact agreement including that sign. +// Route only the exact -0.0f or 0.0f case around __float2half by hand: the +// top 16 bits of an IEEE-754 float zero (sign + 8 exponent + top 8 mantissa, +// all zero except sign) collapse exactly onto the fp16 zero's +// sign+exponent+mantissa layout, so `bits >> 16` is the correct fp16 +// encoding for either signed zero. +// +// Is the workaround needed for CORRECTNESS, or only for bit-exactness with +// the CPU oracle? Only the latter. IEEE-754 defines -0.0 == 0.0, and nothing +// downstream of this kernel in the escha pipeline divides by an activation +// (which is the only operation where the sign of a zero would change a +// result), so the pruned-channel contract — "a pruned channel reads back as +// exactly zero" — holds numerically with EITHER sign of zero. Model output +// would be identical without this workaround. It exists purely to keep this +// kernel bit-for-bit against `escha_ref` for the G3 gate; a port of this +// workaround to another kernel should be justified by its own bit-exactness +// requirement, not assumed to be needed for correctness. +__device__ __forceinline__ __half f2h_rne(float v) { + if (v == 0.0f) { + unsigned int bits = __float_as_uint(v); + return __ushort_as_half((unsigned short)(bits >> 16)); + } + return __float2half(v); +} + +// Parallel 128-point Hadamard butterfly, one thread per element, 7 stages +// with a `__syncthreads()` between each. +// +// Ping-pongs between two LDS buffers instead of updating `bufA` in place: +// the naive single-thread version reads `v[j]`/`v[j+h]` and writes both +// before any other pair is touched, so within one stage a pair's inputs are +// always the PREVIOUS stage's outputs. Doing this in place with 128 threads +// racing would let a "high" thread (writing v[j+h]) read v[j] after some +// other "low" thread has already overwritten it, or vice versa. Ping-ponging +// reads every stage entirely from the buffer the previous stage finished +// (and `__syncthreads()`-published), so every thread's inputs are the exact +// pre-stage values regardless of scheduling — same operand pairing as the +// reference's sequential stage-by-stage sweep, just computed out of order, +// and each pair's `a+b`/`a-b` depends only on that pair's own two operands. +// Order-independence across pairs is what makes this bit-exact with the +// naive kernel and with escha_ref: no reassociation, no changed rounding, +// only a different (but result-identical) schedule of the same 64 adds and +// 64 subtracts per stage. +__device__ __forceinline__ float* h128_parallel(float* bufA, float* bufB, int t) { + float* cur = bufA; + float* nxt = bufB; +#pragma unroll + for (int h = 1; h < 128; h <<= 1) { + int period = h << 1; + bool low = (t & (period - 1)) < h; + int j = low ? t : (t - h); + float a = cur[j]; + float b = cur[j + h]; + nxt[t] = low ? (a + b) : (a - b); + __syncthreads(); + float* tmp = cur; + cur = nxt; + nxt = tmp; + } + return cur; +} + +// One block per 128-channel group. 128 threads cooperate via LDS. +// +// NO `if (idx >= n) return;` BOUNDS GUARD, deliberately. `Gpu::escha_h128` +// asserts `n % 128 == 0` and launches exactly `n / 128` blocks of 128 +// threads, so `idx` is in range for every thread of every block — the guard +// could never fire. If it ever DID, an early `return` before the +// `__syncthreads()` below is a divergent barrier: undefined behaviour, and on +// RDNA a hang or silently wrong LDS contents rather than the clean truncation +// it looks like. The batched forms further down correctly have no such guard. +// `n` is still taken as a parameter so the signature matches the batched +// kernels and the host launcher. +extern "C" __global__ void escha_h128_in( + const float* __restrict__ x, const float* __restrict__ rin, + __half* __restrict__ xh, int n) { + (void)n; + __shared__ float bufA[128]; + __shared__ float bufB[128]; + int g = blockIdx.x, t = threadIdx.x; + int idx = g * 128 + t; + bufA[t] = x[idx] * rin[idx]; + __syncthreads(); + float* out = h128_parallel(bufA, bufB, t); + xh[idx] = f2h_rne(out[t] * ESCHA_RS); +} + +// No bounds guard, for the reason spelled out on `escha_h128_in` above. +extern "C" __global__ void escha_h128_out( + const float* __restrict__ mid, const float* __restrict__ rout, + __half* __restrict__ y, int n) { + (void)n; + __shared__ float bufA[128]; + __shared__ float bufB[128]; + int g = blockIdx.x, t = threadIdx.x; + int idx = g * 128 + t; + bufA[t] = mid[idx]; + __syncthreads(); + float* out = h128_parallel(bufA, bufB, t); + y[idx] = f2h_rne(out[t] * ESCHA_RS * rout[idx]); +} + +// ── Batched-across-experts forms (Task 10) ─────────────────────────────── +// +// WHY THESE EXIST: Task 8 measured the H128 kernels as LAUNCH-bound, not +// bandwidth-bound — an empty kernel at the same grid/block costs 1.74-1.78 us +// against a real launch's 2.4 us, and the overhead-subtracted time is nearly +// flat (0.59 -> 0.69 us) from 16 to 136 blocks. Issuing one launch per +// (expert, projection, side) costs 40 layers x 8 experts x 4 transforms = +// 1280 launches/token = 3.07 ms = a 326 tok/s ceiling from the transforms +// ALONE, before any GEMV work. One launch per (layer, projection, side) +// covering all top_k experts is 160 launches = 0.38 ms. +// +// The maths is identical to the per-expert form; the only change is that +// `r` is now the whole resident `[E, n]` table and each slot picks its row +// via `ids[slot]`. Grid is `slots * (n/128)` blocks of 128 threads; block +// `g` serves slot `g / (n/128)`, group `g % (n/128)`. +// +// OUTPUT DTYPE: these write f32 holding the f16-ROUNDED value (widened), +// not __half. The value set is identical to the per-expert kernels' — the +// rounding happens at exactly the same point — but the consumer is +// hipfire's f32 GEMV, so widening here saves a separate convert launch. +// +// `x_group` says how many CONSECUTIVE slots share one row of `x`: +// +// x_group <= 0 broadcast — every slot reads the single `[n]` row `x`. +// (Decode gate_up input side: a token's top_k experts all see +// the same post-rmsnorm x, only `rin` differs.) +// x_group == 1 one row per slot, `x` is `[slots, n]`. +// (Down input side: each slot's activation is its own SwiGLU +// output.) +// x_group == g row `slot / g`, `x` is `[slots/g, n]`. +// (Batched-prefill gate_up input side, g = k: slots are laid +// out token-major as `token * k + krank`, so `slot / k` is the +// token and every one of its k experts reads that token's x.) +// +// 0 and 1 keep the exact meanings the old binary `x_batched` flag had, so the +// two pre-existing call sites and the two pre-existing G3 cases are unchanged +// by construction; the grouped case is the only new behaviour and G3 gates it +// bit-exactly against escha_ref like the other two. +// +// Deliberately `slot / x_group` rather than a second `tokens` kernarg: the +// division is once per block (128 threads share it) against a value already +// in a register, and expressing it as "how many slots share a row" keeps ALL +// THREE cases one expression instead of a three-way branch per element. +extern "C" __global__ void escha_h128_in_batched( + const float* __restrict__ x, const float* __restrict__ r, + const int* __restrict__ ids, float* __restrict__ out, + int n, int x_group) { + __shared__ float bufA[128]; + __shared__ float bufB[128]; + int groups = n >> 7; + int slot = blockIdx.x / groups; + int grp = blockIdx.x - slot * groups; + int t = threadIdx.x; + int lane = grp * 128 + t; + int e = ids[slot]; + float xv = (x_group <= 0) ? x[lane] + : x[(size_t)(slot / x_group) * n + lane]; + bufA[t] = xv * r[(size_t)e * n + lane]; + __syncthreads(); + float* h = h128_parallel(bufA, bufB, t); + out[(size_t)slot * n + lane] = __half2float(f2h_rne(h[t] * ESCHA_RS)); +} + +extern "C" __global__ void escha_h128_out_batched( + const float* __restrict__ mid, const float* __restrict__ r, + const int* __restrict__ ids, float* __restrict__ out, int n) { + __shared__ float bufA[128]; + __shared__ float bufB[128]; + int groups = n >> 7; + int slot = blockIdx.x / groups; + int grp = blockIdx.x - slot * groups; + int t = threadIdx.x; + int lane = grp * 128 + t; + int e = ids[slot]; + bufA[t] = mid[(size_t)slot * n + lane]; + __syncthreads(); + float* h = h128_parallel(bufA, bufB, t); + // Pruned channel (r == 0) must read back EXACTLY zero — f2h_rne routes + // the exact-zero product around __float2half by hand (see the comment on + // f2h_rne), so a negative pre-scale value against a pruned channel gives + // -0.0, not a denormal and not a dropped sign. Do NOT reorder or + // "optimise away" this multiply. + out[(size_t)slot * n + lane] = + __half2float(f2h_rne(h[t] * ESCHA_RS * r[(size_t)e * n + lane])); +} + +// SwiGLU over the f16-ROUNDED merged gate_up output, batched across slots. +// `y` is `[slots, 2*inter]`, gate is the FIRST half. `y` is already the +// output of escha_h128_out_batched, i.e. every element is an exactly +// representable f16 value — that is the escha contract ("SwiGLU runs on the +// f16-rounded merged gate_up output"), and it is why this kernel must be fed +// the transform's output and never a pre-transform accumulator. +// +// The two interior f16 roundings are `escha_ref::swiglu`'s, not an accident: +// the frozen oracle rounds silu(g) and then rounds the product. Output is +// stored as f32 holding that f16 value (the down-projection input transform +// consumes f32). +// +// `expf`, not `__expf`: the fast intrinsic is a different function; hipfire's +// own `silu_mul_f32` (kernels/src/silu_mul.hip) uses `expf` and so does the +// oracle. Note this kernel is NOT bit-exact against `escha_ref::swiglu` for +// every input — device `expf` and Rust's `f32::exp` can differ by 1 ulp in +// f32, and about 1 in 2^13 of those land on an f16 rounding boundary. The G3 +// gate therefore bounds the swiglu mismatch at <=1 f16 ulp rather than +// demanding equality; the H128 transforms themselves ARE bit-exact and are +// gated as such. +extern "C" __global__ void escha_swiglu_batched( + const float* __restrict__ y, float* __restrict__ h, int inter) { + int i = blockIdx.x * blockDim.x + threadIdx.x; + int slot = blockIdx.y; + if (i >= inter) return; + const float* row = y + (size_t)slot * 2 * inter; + float g = row[i]; + float s = __half2float(f2h_rne(g / (1.0f + expf(-g)))); + h[(size_t)slot * inter + i] = __half2float(f2h_rne(s * row[inter + i])); +} diff --git a/kernels/src/escha_moe_gemm_grouped.hip b/kernels/src/escha_moe_gemm_grouped.hip new file mode 100644 index 0000000000..27e760bfe4 --- /dev/null +++ b/kernels/src/escha_moe_gemm_grouped.hip @@ -0,0 +1,358 @@ +// SPDX-License-Identifier: Apache-2.0 +// hipfire — see LICENSE and NOTICE in the project root. + +#include +#include + +// Escha-W2 routed-expert GROUPED GEMM (Phase 3) — decode each expert's trellis +// code ONCE per (layer, batch) and multiply it against every token that routed +// to it, instead of once per (token, expert) slot. +// +// # The measurement this exists for +// +// `escha_moe_gemv_native.hip` is SLOT-parallel: `krank = blockIdx.y` is one +// (token, expert) pair, so a batch of `n` tokens at `k = 8` re-reads and +// re-decodes each selected expert's whole code `n*k / n_exp` times over. On a +// warm 256-token prefill chunk of the shipped A3B file (`n_exp = 256`, +// `slots = 2048`), `rocprofv3 --kernel-trace` put the two routed GEMV launches +// at 1 779 ms of a 2 370 ms prefill — 75.1% of all GPU time — and prefill +// ms/token was FLAT from n=128 to n=2048 (4.83 / 4.53 / 4.44), i.e. the batch +// bought nothing at all. +// +// Two things are amortised here, and both mattered: +// +// * WEIGHT TRAFFIC. Logical expert bytes per (layer, chunk) go from +// `slots * expert_bytes` to `sum_e ceil(G_e / ROWS) * expert_bytes`, where +// `G_e` is the number of this chunk's slots that chose expert `e`. +// * DECODE ALU. `escha_native_cba` is five ops and the window extraction is +// two loads plus a shift; the slot-parallel kernel pays that per (weight, +// slot), this one pays it per (weight, ROWS slots). +// +// # Why a block owns CTILES tile columns, not one +// +// The slot-parallel kernel's block owns ONE 16-wide tile column, so every slot +// re-reads its whole `x` row once per tile column — `(M/16) * slots * K * 4` +// bytes, which at the gate_up shape is 1.07 GB per (layer, chunk), the SAME +// order as the expert traffic the grouping removes. Grouping alone would have +// left that term untouched and capped the win at ~1.8x. +// +// A block that owns CTILES ADJACENT tile columns reads `x` once for all of +// them: the `x` term divides by CTILES while the expert term is unchanged +// (there are CTILES-times fewer blocks, each reading CTILES-times more code). +// The decode slot descriptor (`word_a`, `word_b`, `shift`) depends only on +// (row-in-tile, column-in-tile), so the extra tile columns cost only their own +// code loads — no extra slot maths, no extra registers beyond the accumulators. +// +// # Registers +// +// A lane holds `ROWS * 2 * CTILES` f32 accumulators (two output columns per +// warp — `w` and `w + 8` share a decode window). ROWS=8, CTILES=4 is 64, which +// fits with room for the `x` values in flight; the host wrapper picks the +// instantiation and the tuning sweep behind +// `HIPFIRE_ESCHA_GROUPED_TILE` is what chose it. +// +// # No codebook lookup table, still +// +// `escha_native_cba` is five inline ops. 65 536 fp16 entries would be 128 KB +// against gfx1151's 64 KB of LDS, so there is no table here either — same +// constraint, same answer, as `escha_decode_tiles.hip` and +// `escha_moe_gemv_native.hip`. +// +// # K=2 and K=3 stay separate +// +// The two trellis orders extract their 16-bit window from structurally +// different stream geometries. Everything that decides WHICH two words and +// WHICH shift is duplicated per K and resolved at compile time, exactly as in +// `escha_moe_gemv_native.hip`; neither pays a runtime branch. +// +// # DELIBERATE DUPLICATION +// +// `escha_native_cba`, `escha_native_word`, the lane/j inversion and the two +// slot descriptors below are copied VERBATIM from +// `kernels/src/escha_moe_gemv_native.hip`, which in turn reproduces +// `escha_decode_tiles.hip` and `hipfire_quantize::escha_ref`. That chain is +// what G2 and G7 assert bit-exact against the frozen CPU oracle. Sharing the +// source between the kernel under test and the kernel it is tested against +// would make those gates circular. Do not deduplicate. +// +// # Numerics — this kernel is NOT bit-identical to the slot-parallel one for +// # every projection, and that is stated rather than hidden +// +// Per (token, output row) this kernel keeps the slot-parallel NARROW form +// exactly: the same lane -> contraction-index map (`lane` strided by 32), the +// same `bi` order, ONE sequential accumulator, and the same `__shfl_down` +// ladder. So for a projection the slot-parallel path also runs narrow +// (`K > 1536` — the shipped gate_up, K = 2048) the two agree BIT FOR BIT. +// +// For a projection the slot-parallel path runs WIDE (`K <= 1536` — the shipped +// down, K = 512) it does not: that form folds FOUR interleaved accumulators as +// `(acc0+acc1)+(acc2+acc3)`, and reproducing it here would cost 4x the +// accumulator registers (ROWS*2*CTILES*4 = 256 at the shipped tile) and force +// CTILES back to 1, which is the term that had to shrink. The grouped kernel +// therefore sums that projection as one sequential chain. The difference is a +// different PARTITION of the same contraction, not a different contraction: +// both are exact-order f32 sums of the identical 512 products. The measured +// cost is reported by the G4 gate's grouped arm. + +// Contraction-loop unroll factor. 1 — i.e. NO unroll — and that is a MEASURED +// choice, not an oversight. +// +// The slot-parallel narrow kernel needs `#pragma unroll 4` because it has one +// FMA chain and one memory round trip of latency to hide per iteration, with +// nothing to hide it with. This kernel already carries ROWS*2*CTILES = 64 +// independent chains and CTILES*2 independent code loads per iteration, so the +// ILP is there without unrolling. Unrolling on top of it only widens the live +// range of the in-flight loads: measured on the shipped gate_up shape (drift +// corrected against the untouched slot-parallel control in the same process), +// unroll 2 took r8c4 from 2.38x to 2.00x and left r8c8 unchanged. +// +// Overridable at build time so a future shape can re-derive the choice rather +// than inherit it. +#ifndef ESCHA_GROUPED_UNROLL +#define ESCHA_GROUPED_UNROLL 1 +#endif + +// ── the codebook, inline ───────────────────────────────────────────────────── +__device__ __forceinline__ __half escha_native_cba(unsigned short state) { + unsigned int r = ((unsigned int)state * 0xCBAC1FEDu) & 0x8FFF8FFFu; + r ^= 0x3B603B60u; + __half lo = __ushort_as_half((unsigned short)(r & 0xFFFFu)); + __half hi = __ushort_as_half((unsigned short)(r >> 16)); + return __hadd(lo, hi); // fp16 RNE add — matches the reference exactly +} + +// ADDRESS SPACE 1, explicitly — `expert_ptrs[...]` is a runtime 64-bit integer, +// so without the cast the compiler emits `flat_load_b32` and pays an aperture +// check the whole inner loop then waits on. See the same typedef in +// `escha_moe_gemv_native.hip`, where it was measured worth 1.6x on its own. +typedef __attribute__((address_space(1))) const unsigned int* EschaCodePtr; + +__device__ __forceinline__ unsigned int escha_native_word(EschaCodePtr tile, int i) { + return tile[i]; +} + +__device__ __forceinline__ int escha_native_lane_of(int row, int col) { + return ((col >> 1) & 3) * 8 + (col & 1) * 4 + ((row & 7) >> 1); +} +__device__ __forceinline__ int escha_native_j_of(int row, int col) { + return ((col >> 3) & 1) * 4 + ((row >> 3) & 1) * 2 + (row & 1); +} + +__device__ __forceinline__ void escha_native_slot_k2( + int row, int col, int* word_a, int* word_b, int* shift) { + const int lane = escha_native_lane_of(row, col); + const int j = escha_native_j_of(row, col); + const int t_off = lane * 8; + const int i1 = t_off >> 4; + const int i0 = (i1 + 15) & 15; + *word_a = i0; + *word_b = i1; + *shift = (((~t_off) & 8) << 1) + 2 * (7 - j); +} + +__device__ __forceinline__ void escha_native_slot_k3( + int row, int col, int* word_a, int* word_b, int* shift) { + const int BITS = 3; + const int lane = escha_native_lane_of(row, col); + const int j = escha_native_j_of(row, col); + const int t_off = lane * 8; + const int b1 = (t_off + 257) * BITS; + const int b0 = b1 - 16; + const int b2 = b1 + BITS * 7; + const int i0 = b0 >> 5; + const int i2 = (b2 - 1) >> 5; + const int s2 = ((i2 + 1) << 5) - b2; + *word_a = i0 % 24; + *word_b = i2 % 24; + *shift = s2 + (j < 4 ? BITS * 4 : 0) + BITS * (3 - (j & 3)); +} + +template +__device__ __forceinline__ void escha_native_slot( + int row, int col, int* word_a, int* word_b, int* shift) { + if (TRELLIS_K == 2) escha_native_slot_k2(row, col, word_a, word_b, shift); + else escha_native_slot_k3(row, col, word_a, word_b, shift); +} + +// ── the grouped body ───────────────────────────────────────────────────────── +// +// Grid: (M / (16*CTILES), n_exp, 1). Block: 256 = eight warps; warp `w` owns +// output columns `w` and `w + 8` of each of the block's CTILES tile columns. +// +// `expert_offsets[e] .. expert_offsets[e+1]` is expert `e`'s half-open range in +// the SORTED slot order, and `sorted_slot_index[p]` is the flat slot at sorted +// position `p`. Both come from `moe_scatter_fused_k8` run with `block_m = 1`, +// i.e. WITHOUT the grouped-WMMA path's BLOCK_M padding: padding would cost real +// `x` reads and real FMAs for rows that are then discarded, and at the shipped +// shape (256 experts, 2048 slots, mean group 8) a BLOCK_M of 8 would have +// inflated the `x` term — the dominant one — by ~40%. Ranging over the exact +// group instead costs one extra `expert_offsets` load and a runtime loop bound. +// +// An expert with no slots this chunk exits before touching `expert_ptrs`, so a +// batch that lights up 30 of 256 experts does 30 experts' work. +template +__device__ __forceinline__ void escha_gemm_grouped_body( + const unsigned long long* __restrict__ expert_ptrs, + const int* __restrict__ expert_offsets, // [n_exp + 1] exclusive scan + const int* __restrict__ sorted_slot_index, // [slots] sorted pos -> flat slot + const float* __restrict__ x_batch, // [slots, K] + float* __restrict__ y_batch, // [slots, M] + int M, int K, int nt_major +) { + const int e = blockIdx.y; + const int g_start = expert_offsets[e]; + const int g_end = expert_offsets[e + 1]; + if (g_start >= g_end) return; + + const int warp = threadIdx.x >> 5; + const int lane = threadIdx.x & 31; + + // Loop-invariant decode geometry. `in_row` is this lane's row inside a + // tile, `tile_hi` says which of the iteration's two stacked tiles it falls + // in — identical to the slot-parallel kernel, because the lane -> weight + // map is what the bit-exactness contract is about. + const int tn = M >> 4; + const int in_row = lane & 15; + const int tile_hi = lane >> 4; + int word_a, word_b, shift_lo, shift_hi; + int word_a2, word_b2; + escha_native_slot(in_row, warp, &word_a, &word_b, &shift_lo); + escha_native_slot(in_row, warp + 8, &word_a2, &word_b2, &shift_hi); + // Columns `c` and `c + 8` provably share a decode window (see + // `escha_moe_gemv_native.hip`). If that ever stops holding the two columns + // would silently read each other's weights, so poison loudly instead — + // every escha gate asserts finiteness and nothing else manufactures a NaN. + if (word_a2 != word_a || word_b2 != word_b) { + if (lane == 0) { + for (int p = g_start; p < g_end; p++) { + const int s = sorted_slot_index[p]; + if (s < 0) continue; + #pragma unroll + for (int c = 0; c < CTILES; c++) { + const int row = (blockIdx.x * CTILES + c) * 16 + warp; + y_batch[(size_t)s * M + row] = __builtin_nanf(""); + y_batch[(size_t)s * M + row + 8] = __builtin_nanf(""); + } + } + } + return; + } + + EschaCodePtr code = (EschaCodePtr)(expert_ptrs[e]); + // Tile `kt` of tile column `nt` starts at `(kt*tn + nt) * 8*TRELLIS_K` + // DWORDS. This lane's tile in iteration `bi` is `kt = bi*2 + tile_hi`. + const size_t ktiles = (size_t)K >> 4; + const size_t strip = nt_major ? (size_t)(8 * TRELLIS_K) + : (size_t)tn * (8 * TRELLIS_K); + const size_t colstride = nt_major ? ktiles * (size_t)(8 * TRELLIS_K) + : (size_t)(8 * TRELLIS_K); + const int nt0 = blockIdx.x * CTILES; + EschaCodePtr strip0 = code + (size_t)nt0 * colstride + (size_t)tile_hi * strip; + + const int blocks_per_row = K / 32; + + for (int g = g_start; g < g_end; g += ROWS) { + // The chunk's rows. Out-of-range lanes carry -1 and are skipped by a + // BLOCK-UNIFORM branch (every lane sees the same `slot[r]`), so this is + // a scalar branch, not divergence. + int slot[ROWS]; + #pragma unroll + for (int r = 0; r < ROWS; r++) { + const int p = g + r; + slot[r] = (p < g_end) ? sorted_slot_index[p] : -1; + } + + float acc[ROWS][2 * CTILES]; + #pragma unroll + for (int r = 0; r < ROWS; r++) { + #pragma unroll + for (int c = 0; c < 2 * CTILES; c++) acc[r][c] = 0.0f; + } + + #pragma unroll ESCHA_GROUPED_UNROLL + for (int bi = 0; bi < blocks_per_row; bi++) { + // Decode this iteration's CTILES weights pairs ONCE, then spend + // them across every row of the chunk. This is the whole point of + // the kernel: `ROWS * CTILES * 2` FMAs per `CTILES * 2` decodes. + float w_lo[CTILES], w_hi[CTILES]; + #pragma unroll + for (int c = 0; c < CTILES; c++) { + EschaCodePtr tile = + strip0 + (size_t)c * colstride + (size_t)bi * 2 * strip; + const unsigned long long merged = + ((unsigned long long)escha_native_word(tile, word_a) << 32) | + escha_native_word(tile, word_b); + w_lo[c] = __half2float(escha_native_cba((unsigned short)(merged >> shift_lo))); + w_hi[c] = __half2float(escha_native_cba((unsigned short)(merged >> shift_hi))); + } + const int xo = bi * 32 + lane; + #pragma unroll + for (int r = 0; r < ROWS; r++) { + if (slot[r] < 0) continue; + const float xv = x_batch[(size_t)slot[r] * K + xo]; + #pragma unroll + for (int c = 0; c < CTILES; c++) { + acc[r][2 * c] += w_lo[c] * xv; + acc[r][2 * c + 1] += w_hi[c] * xv; + } + } + } + + // The slot-parallel kernel's reduction, verbatim, once per (row, + // column). Every value lane 0 consumes comes from an in-range lane, so + // the ladder is a balanced binary tree over the 32 partial sums. + #pragma unroll + for (int r = 0; r < ROWS; r++) { + if (slot[r] < 0) continue; + #pragma unroll + for (int c = 0; c < CTILES; c++) { + float sum_lo = acc[r][2 * c]; + float sum_hi = acc[r][2 * c + 1]; + for (int offset = 16; offset > 0; offset >>= 1) { + sum_lo += __shfl_down(sum_lo, offset); + sum_hi += __shfl_down(sum_hi, offset); + } + if (lane == 0) { + const int row = (nt0 + c) * 16 + warp; + y_batch[(size_t)slot[r] * M + row] = sum_lo; + y_batch[(size_t)slot[r] * M + row + 8] = sum_hi; + } + } + } + } +} + +// ── entry points ───────────────────────────────────────────────────────────── +// +// One per (trellis order, tile shape). The host wrapper enforces +// `M % (16*CTILES) == 0` and `K % 32 == 0`, and picks the shape; the set below +// is what the sweep in `bench_escha_grouped_gemm` covers. + +#define ESCHA_GROUPED_GEMM(NAME, TK, ROWS, CTILES) \ + __launch_bounds__(256) \ + extern "C" __global__ void NAME( \ + const unsigned long long* __restrict__ expert_ptrs, \ + const int* __restrict__ expert_offsets, \ + const int* __restrict__ sorted_slot_index, \ + const float* __restrict__ x_batch, \ + float* __restrict__ y_batch, \ + int M, int K, int nt_major \ + ) { \ + escha_gemm_grouped_body( \ + expert_ptrs, expert_offsets, sorted_slot_index, x_batch, y_batch, \ + M, K, nt_major); \ + } + +ESCHA_GROUPED_GEMM(escha_gemm_grouped_k2_r4_c2, 2, 4, 2) +ESCHA_GROUPED_GEMM(escha_gemm_grouped_k2_r8_c2, 2, 8, 2) +ESCHA_GROUPED_GEMM(escha_gemm_grouped_k2_r8_c4, 2, 8, 4) +ESCHA_GROUPED_GEMM(escha_gemm_grouped_k2_r16_c2, 2, 16, 2) +ESCHA_GROUPED_GEMM(escha_gemm_grouped_k2_r16_c4, 2, 16, 4) +ESCHA_GROUPED_GEMM(escha_gemm_grouped_k2_r8_c8, 2, 8, 8) + +ESCHA_GROUPED_GEMM(escha_gemm_grouped_k3_r4_c2, 3, 4, 2) +ESCHA_GROUPED_GEMM(escha_gemm_grouped_k3_r8_c2, 3, 8, 2) +ESCHA_GROUPED_GEMM(escha_gemm_grouped_k3_r8_c4, 3, 8, 4) +ESCHA_GROUPED_GEMM(escha_gemm_grouped_k3_r16_c2, 3, 16, 2) +ESCHA_GROUPED_GEMM(escha_gemm_grouped_k3_r16_c4, 3, 16, 4) +ESCHA_GROUPED_GEMM(escha_gemm_grouped_k3_r8_c8, 3, 8, 8) diff --git a/kernels/src/escha_moe_gemm_grouped_wmma.hip b/kernels/src/escha_moe_gemm_grouped_wmma.hip new file mode 100644 index 0000000000..70d1e94647 --- /dev/null +++ b/kernels/src/escha_moe_gemm_grouped_wmma.hip @@ -0,0 +1,232 @@ +// SPDX-License-Identifier: MIT OR Apache-2.0 +// hipfire — see LICENSE and NOTICE in the project root. +// +// Escha-W2 expert-grouped routed GEMM on the RDNA3 matrix cores. +// +// WHY THIS EXISTS +// --------------- +// The scalar grouped GEMM (`escha_moe_gemm_grouped.hip`) fixed prefill's +// WEIGHT traffic — each expert's trellis code is read once per (layer, batch) +// instead of once per (token, expert) slot, which took n=512 prefill from +// 4.525 to 2.657 ms/token. But that made weight traffic 4% of prefill time, +// and left the other 96% running scalar `acc += w * x` FMAs on the VALU: +// +// escha (scalar grouped) 1360 ms / 512 tok -> 1.75 TFLOP/s +// ornith (WMMA grouped) 520 ms / 512 tok -> 4.59 TFLOP/s +// +// Same 2.6x as the throughput gap. Prefill is COMPUTE bound, so the fix is +// the matrix cores, not more bandwidth work. +// +// WHY THE TILE SHAPE IS A GIFT +// ---------------------------- +// Escha codes a 16x16 weight tile as its unit. WMMA on gfx11 is 16x16x16. +// One decoded escha tile is therefore exactly one A fragment — no repacking, +// no padding, no partial-tile bookkeeping in the K direction. +// +// ORIENTATION +// ----------- +// The GEMM is `y[slot, m] = sum_k x[slot, k] * W[k, m]`, and escha stores W +// IN-major: tile (kt, nt) holds rows `kt*16 + r` (contraction) by columns +// `nt*16 + c` (output). WMMA wants A as [m x k], so lane `l` takes COLUMN `l` +// of the decoded tile — i.e. `W[kt*16 + 0..15, nt*16 + l]`. Decoding into LDS +// first makes that a plain strided read; decoding straight to registers would +// need a cross-lane transpose for no benefit. +// +// NUMERICS +// -------- +// WMMA accumulates in f32 over a different partition of the contraction than +// the scalar path, so this is NOT bit-identical to it, exactly as the scalar +// grouped path is not bit-identical to the slot-parallel one. The weight +// VALUES are the same decoded fp16 either way — only the summation order +// moves. Gated by measured bound, not by equality; see the gate. +// +// The codebook stays inline (no LUT): 65536 fp16 entries is 128 KB and +// gfx1151 has 64 KB of LDS. + +#include +#include + +typedef _Float16 __attribute__((ext_vector_type(16))) half16_t; +typedef float __attribute__((ext_vector_type(8))) float8_t; + +// ── codec (transcribed from escha_decode_tiles.hip; see the deliberate +// duplication note there — the CPU reference is the gate for both) ───────── + +__device__ __forceinline__ __half escha_w_cba(unsigned short state) { + unsigned int r = ((unsigned int)state * 0xCBAC1FEDu) & 0x8FFF8FFFu; + r ^= 0x3B603B60u; + __half lo = __ushort_as_half((unsigned short)(r & 0xFFFFu)); + __half hi = __ushort_as_half((unsigned short)(r >> 16)); + return __hadd(lo, hi); +} + +__device__ __forceinline__ unsigned int escha_w_word( + const short* __restrict__ src, int i) { + return ((unsigned int)(unsigned short)src[2 * i]) | + (((unsigned int)(unsigned short)src[2 * i + 1]) << 16); +} + +__device__ __forceinline__ void escha_w_dec8_k2( + const short* __restrict__ src, int lane, unsigned short* out) { + int t_off = lane * 8; + int i1 = t_off >> 4; + int i0 = (i1 + 15) & 15; + unsigned long long merged = + ((unsigned long long)escha_w_word(src, i0) << 32) | escha_w_word(src, i1); + int shift = ((~t_off) & 8) << 1; + unsigned int v = (unsigned int)((merged >> shift) & 0xFFFFFFFFull); +#pragma unroll + for (int j = 0; j < 8; ++j) out[j] = (unsigned short)(v >> (2 * (7 - j))); +} + +__device__ __forceinline__ void escha_w_dec8_k3( + const short* __restrict__ src, int lane, unsigned short* out) { + const int BITS = 3; + int t_off = lane * 8; + int b1 = (t_off + 257) * BITS; + int b0 = b1 - 16; + int b2 = b1 + BITS * 7; + int i0 = b0 >> 5; + int i2 = (b2 - 1) >> 5; + int s2 = ((i2 + 1) << 5) - b2; + unsigned long long merged = ((unsigned long long)escha_w_word(src, i0 % 24) << 32) | + escha_w_word(src, i2 % 24); + unsigned int w7 = (unsigned int)((merged >> s2) & 0xFFFFFFFFull); + unsigned int w3 = (unsigned int)((merged >> (s2 + BITS * 4)) & 0xFFFFFFFFull); + out[0] = (unsigned short)(w3 >> 9); + out[1] = (unsigned short)(w3 >> 6); + out[2] = (unsigned short)(w3 >> 3); + out[3] = (unsigned short)(w3); + out[4] = (unsigned short)(w7 >> 9); + out[5] = (unsigned short)(w7 >> 6); + out[6] = (unsigned short)(w7 >> 3); + out[7] = (unsigned short)(w7); +} + +/// Decode one 16x16 escha tile into `dst` (row-major, 16 rows of `stride`). +/// All 32 lanes of the wave participate: 32 lanes x 8 values = the 256 slots. +template +__device__ __forceinline__ void escha_decode_tile_lds( + const short* __restrict__ code_tile, _Float16* dst, int stride, int lane) { + unsigned short st[8]; + if (TK == 2) escha_w_dec8_k2(code_tile, lane, st); + else escha_w_dec8_k3(code_tile, lane, st); + + const int l0 = lane & ~4; + const int c_off = (lane >> 2) & 1; +#pragma unroll + for (int j = 0; j < 8; ++j) { + const int fi = j >> 1; + const int row = (lane & 3) * 2 + (j & 1) + (fi & 1) * 8; + const int col = 2 * ((l0 >> 3) + (j >= 4 ? 4 : 0)) + c_off; + dst[row * stride + col] = (_Float16)escha_w_cba(st[j]); + } +} + +// ── the GEMM ────────────────────────────────────────────────────────────── +// +// grid.x = output tile index (M / 16) +// grid.y = expert id +// block = 32 lanes (one wave32) +// +// Each block owns one 16-wide output column tile for one expert, and walks +// that expert's slot group 16 slots at a time. `+1` on the LDS strides keeps +// the column reads off a single bank. + +#define ESCHA_WMMA_GEMM(NAME, TK, NT) \ + __launch_bounds__(32) extern "C" __global__ void NAME( \ + const unsigned long long* __restrict__ expert_ptrs, \ + const int* __restrict__ expert_offsets, \ + const int* __restrict__ sorted_slot_index, \ + const float* __restrict__ x_batch, \ + float* __restrict__ y_batch, \ + int M, int K, int nt_major) { \ + const int nt = blockIdx.x; \ + const int e = blockIdx.y; \ + const int lane = threadIdx.x; \ + const int g_start = expert_offsets[e]; \ + const int g_end = expert_offsets[e + 1]; \ + if (g_start >= g_end) return; \ + \ + const short* code = (const short*)expert_ptrs[e]; \ + const int ktiles = K / 16; \ + const int ntiles = M / 16; \ + \ + __shared__ _Float16 w_lds[16 * 17]; \ + \ + for (int base = g_start; base < g_end; base += 16 * NT) { \ + float8_t acc[NT]; \ + _Pragma("unroll") for (int t = 0; t < NT; ++t) \ + acc[t] = (float8_t){0.f, 0.f, 0.f, 0.f, 0.f, 0.f, 0.f, 0.f}; \ + \ + for (int kt = 0; kt < ktiles; ++kt) { \ + /* A: decode escha tile (kt, nt); lane takes column `lane`. */ \ + /* kt-major `[ic/16][oc/16]` or nt-major `[oc/16][ic/16]`; \ + see escha_moe_gemv_native.hip. Two loop-invariant strides. */\ + const size_t tile_off = nt_major \ + ? (size_t)(nt * ktiles + kt) : (size_t)(kt * ntiles + nt); \ + const short* tile = code + tile_off * 16 * TK; \ + escha_decode_tile_lds(tile, w_lds, 17, lane); \ + \ + /* Stage the weight tile AND all NT activation tiles, then a \ + SINGLE barrier. An earlier version put a __syncthreads() \ + inside the NT loop against one shared x_lds; that serialised\ + the staging and measured 2.4x SLOWER than NT=1 (13.06 vs \ + 5.40 ms) even though the numerics were identical. */ \ + __syncthreads(); \ + const int m_lane = lane & 15; \ + half16_t a_frag; \ + _Pragma("unroll") for (int i = 0; i < 16; ++i) \ + a_frag[i] = w_lds[i * 17 + m_lane]; /* column = out col */ \ + \ + /* B comes STRAIGHT FROM GLOBAL, not via LDS. Lane `l` needs \ + x[slot_l, kt*16 .. +16] — already contiguous, and each lane \ + wants a different slot, so there is nothing to share and \ + nothing to transpose. Staging it through LDS was measured \ + to cost 0.52 G LDS instructions against the scalar kernel's \ + 0.05 G, which exactly cancelled the 2.2x VALU saving WMMA \ + bought (busy cycles 1.48 G vs 1.47 G). */ \ + _Pragma("unroll") for (int t = 0; t < NT; ++t) { \ + const int p = base + t * 16 + m_lane; \ + const int s = (p < g_end) ? sorted_slot_index[p] : -1; \ + half16_t b_frag; \ + if (s >= 0) { \ + const float* xr = x_batch + (size_t)s * K + kt * 16; \ + _Pragma("unroll") for (int i = 0; i < 16; ++i) \ + b_frag[i] = (_Float16)xr[i]; \ + } else { \ + _Pragma("unroll") for (int i = 0; i < 16; ++i) \ + b_frag[i] = (_Float16)0.f; \ + } \ + acc[t] = __builtin_amdgcn_wmma_f32_16x16x16_f16_w32( \ + a_frag, b_frag, acc[t]); \ + } \ + __syncthreads(); \ + } \ + \ + /* gfx11 wave32 16x16 f32 accumulator layout (see the writeback \ + in gemm_bq1g128_wmma.hip, which is the authoritative local \ + example): the A-side index is `2*i + (lane>>4)` and the B-side \ + index is `lane & 15`. Here A is the OUTPUT COLUMN and B is the \ + SLOT — getting these the wrong way round still runs and still \ + produces finite numbers, it just produces the wrong ones \ + (measured max |diff| 2.2e2 against the scalar path). */ \ + const int slot_in_tile = lane & 15; \ + _Pragma("unroll") for (int t = 0; t < NT; ++t) { \ + const int p = base + t * 16 + slot_in_tile; \ + if (p < g_end) { \ + const int s = sorted_slot_index[p]; \ + if (s >= 0) { \ + _Pragma("unroll") for (int i = 0; i < 8; ++i) { \ + const int oc = nt * 16 + 2 * i + (lane >> 4); \ + y_batch[(size_t)s * M + oc] = acc[t][i]; \ + } \ + } \ + } \ + } \ + __syncthreads(); \ + } \ + } + +ESCHA_WMMA_GEMM(escha_gemm_grouped_wmma_k2, 2, 8) +ESCHA_WMMA_GEMM(escha_gemm_grouped_wmma_k3, 3, 8) diff --git a/kernels/src/escha_moe_gemv_k8_indexed.hip b/kernels/src/escha_moe_gemv_k8_indexed.hip new file mode 100644 index 0000000000..4590b15584 --- /dev/null +++ b/kernels/src/escha_moe_gemv_k8_indexed.hip @@ -0,0 +1,212 @@ +// SPDX-License-Identifier: Apache-2.0 +// hipfire — see LICENSE and NOTICE in the project root. + +#include +#include + +// Escha-W2 routed-expert GEMVs for the GPU-top-K (indexed) decode path. +// +// # Why these exist when `gemv_q8_0_moe_{gate_up,down}_*_k8_indexed` already do +// +// The two pre-existing indexed Q8_0 MoE kernels are shaped for the ORDINARY +// routed path and neither fits Escha-W2: +// +// * `gemv_q8_0_moe_gate_up_k8_indexed` broadcasts ONE `x` vector to all k +// ranks and splits its result into separate `y_gate` / `y_up` buffers. +// Escha needs the opposite on both counts: every rank has its OWN input +// (the gate_up input transform folds a PER-EXPERT `rin_eff` row into x, so +// the k slots are k different vectors), and the output must stay a +// contiguous `[k, 2*mi]` block because the gate_up OUTPUT transform reads +// a per-expert `rout_eff` row of width `2*mi` spanning both halves. +// * `gemv_q8_0_moe_down_residual_scaled_k8_indexed` folds the weighted +// combine into the GEMV via atomicAdd. Escha cannot: the down result is +// still in the ROTATED domain at that point and must pass through the +// down output transform BEFORE it is combined. +// +// So both escha phases want the same thing — per-slot x in, per-slot y out, +// no combine — and that is exactly one kernel used twice. +// +// # Bit-exactness contract (load-bearing) +// +// These are transcriptions of `gemv_q8_0.hip` / `gemv_q8_0_wide.hip` with two +// changes and no others: the weight base comes from +// `expert_ptrs[topk_indices[krank]]` instead of a kernarg, and `x` / `y` are +// offset by `krank`. The accumulate order, the unroll structure, the +// accumulator count and the final reduction are copied verbatim, because the +// escha routed path must produce the SAME numbers on the indexed route as it +// does on the per-expert `GemvFamily::run_auto` route it replaces — the G4 +// block gate's tolerances are calibrated against those numbers and a changed +// summation order would move them. +// +// The narrow/wide split mirrors `Gpu::gemv_q8_0`'s own `k <= 1536` rule for +// the same reason: `gemv_q8_0_wide` sums into FOUR interleaved accumulators +// and `gemv_q8_0` into one, so picking the other variant for a given +// projection is a numerical change, not just a performance one. The host +// wrapper re-uses that threshold rather than restating it. +// +// Q8_0 layout: 34 bytes per block of 32 elements (2 B f16 scale + 32 B int8). +// Expert ptrs: [n_exp] weight-base pointers packed as `unsigned long long`. +// No kernarg depends on the routing result, so both replay under hipGraph. + +// Narrow variant. Grid: (M, K_TOP, 1). Block: 32. +// x: [K_TOP, K], y: [K_TOP, M]. +__launch_bounds__(32, 20) +extern "C" __global__ void escha_gemv_q8_0_moe_k8_indexed_batched( + const unsigned long long* __restrict__ expert_ptrs, // [n_exp] + const int* __restrict__ topk_indices, // [K_TOP] + const float* __restrict__ x_batch, // [K_TOP, K] + float* __restrict__ y_batch, // [K_TOP, M] + int M, int K +) { + const int row = blockIdx.x; + if (row >= M) return; + const int krank = blockIdx.y; + const int tid = threadIdx.x; + + const unsigned char* __restrict__ A_q8 = + reinterpret_cast(expert_ptrs[topk_indices[krank]]); + const float* __restrict__ x = x_batch + (size_t)krank * K; + float* __restrict__ y = y_batch + (size_t)krank * M; + + const int blocks_per_row = K / 32; + const unsigned char* row_data = A_q8 + (size_t)row * blocks_per_row * 34; + + float sum = 0.0f; + + // Process 8 Q8_0 blocks (256 elements) per outer iteration + const int outer_iters = blocks_per_row / 8; + for (int oi = 0; oi < outer_iters; oi++) { + const unsigned char* base = row_data + oi * 8 * 34; + const float* xb = x + oi * 256; + + // Unrolled: 8 blocks, each: load scale (f16→f32), load byte, FMA + #pragma unroll + for (int sub = 0; sub < 8; sub++) { + const unsigned char* block = base + sub * 34; + float d = (float)*((const _Float16*)block); + signed char qval = (signed char)block[2 + tid]; + sum += d * (float)qval * xb[sub * 32 + tid]; + } + } + + // Handle remaining blocks (if K is not multiple of 256) + for (int bi = outer_iters * 8; bi < blocks_per_row; bi++) { + const unsigned char* block = row_data + bi * 34; + float d = (float)*((const _Float16*)block); + signed char qval = (signed char)block[2 + tid]; + sum += d * (float)qval * x[bi * 32 + tid]; + } + + for (int offset = 16; offset > 0; offset >>= 1) + sum += __shfl_down(sum, offset); + if (tid == 0) y[row] = sum; +} + +// Wide variant: 2 rows per block, one row per warp. Grid: (ceil(M/2), K_TOP, 1). +// Block: 64. x: [K_TOP, K], y: [K_TOP, M]. +// +// Four interleaved accumulators with the `(acc0 + acc1) + (acc2 + acc3)` fold, +// copied from `gemv_q8_0_wide.hip` — see that file for why the ordering is +// fixed rather than incidental. +extern "C" __global__ void escha_gemv_q8_0_wide_moe_k8_indexed_batched( + const unsigned long long* __restrict__ expert_ptrs, // [n_exp] + const int* __restrict__ topk_indices, // [K_TOP] + const float* __restrict__ x_batch, // [K_TOP, K] + float* __restrict__ y_batch, // [K_TOP, M] + int M, int K +) { + const int tid = threadIdx.x; + const int warp_id = tid / 32; + const int lane = tid & 31; + const int row = blockIdx.x * 2 + warp_id; + if (row >= M) return; + const int krank = blockIdx.y; + + const unsigned char* __restrict__ A_q8 = + reinterpret_cast(expert_ptrs[topk_indices[krank]]); + const float* __restrict__ x = x_batch + (size_t)krank * K; + float* __restrict__ y = y_batch + (size_t)krank * M; + + const int blocks_per_row = K / 32; + const unsigned char* row_data = A_q8 + (size_t)row * blocks_per_row * 34; + + float acc0 = 0.0f, acc1 = 0.0f, acc2 = 0.0f, acc3 = 0.0f; + const int quads = blocks_per_row >> 2; + const int tail = blocks_per_row & 3; + + for (int q = 0; q < quads; q++) { + const int bi = q << 2; + const unsigned char* bp0 = row_data + (bi + 0) * 34; + const unsigned char* bp1 = row_data + (bi + 1) * 34; + const unsigned char* bp2 = row_data + (bi + 2) * 34; + const unsigned char* bp3 = row_data + (bi + 3) * 34; + + float d0 = (float)*((const _Float16*)bp0); + float d1 = (float)*((const _Float16*)bp1); + float d2 = (float)*((const _Float16*)bp2); + float d3 = (float)*((const _Float16*)bp3); + + signed char q0 = (signed char)bp0[2 + lane]; + signed char q1 = (signed char)bp1[2 + lane]; + signed char q2 = (signed char)bp2[2 + lane]; + signed char q3 = (signed char)bp3[2 + lane]; + + acc0 += d0 * (float)q0 * x[(bi + 0) * 32 + lane]; + acc1 += d1 * (float)q1 * x[(bi + 1) * 32 + lane]; + acc2 += d2 * (float)q2 * x[(bi + 2) * 32 + lane]; + acc3 += d3 * (float)q3 * x[(bi + 3) * 32 + lane]; + } + + // Tail blocks accumulate into acc[bi % 4]. + for (int t = 0; t < tail; t++) { + const int bi = (quads << 2) + t; + const unsigned char* bp = row_data + bi * 34; + float d = (float)*((const _Float16*)bp); + signed char qv = (signed char)bp[2 + lane]; + float contrib = d * (float)qv * x[bi * 32 + lane]; + if (t == 0) acc0 += contrib; + else if (t == 1) acc1 += contrib; + else if (t == 2) acc2 += contrib; + } + + float sum = (acc0 + acc1) + (acc2 + acc3); + for (int offset = 16; offset > 0; offset >>= 1) + sum += __shfl_down(sum, offset); + if (lane == 0) y[row] = sum; +} + +// Out-of-place F32 -> f16 -> F32 round-trip of the k combine weights. +// +// The escha combine multiplies each expert's contribution by `f16(score)` — +// one of the three load-bearing rounding points of the format (EschaLabs' +// runtime stores the routing weights as f16). On the CPU-top-K route that +// rounding is applied host-side, on the already-downloaded weights. The +// indexed route never downloads them, so it must be done on device. +// +// It is OUT-OF-PLACE on purpose: `topk_weights` is the shared routing buffer +// and other consumers (expert-stats capture, any future reader) must keep +// seeing the unrounded F32 values. `dst` is the escha layer's own [k] scratch. +// +// The exact-zero guard is the same `f2h_rne` workaround the other escha +// kernels carry: on this ROCm 7.2.2/gfx1151 build `__float2half` loses the +// sign of a runtime-produced exact zero. Softmax weights are non-negative so +// the sign cannot matter here, but the guard is free and keeps one rounding +// idiom across the escha kernels rather than two. +extern "C" __global__ void escha_round_weights_f16_rne( + const float* __restrict__ src, + float* __restrict__ dst, + int n +) { + int i = blockIdx.x * blockDim.x + threadIdx.x; + if (i >= n) return; + float v = src[i]; + if (!isfinite(v)) { dst[i] = v; return; } + __half h; + if (v == 0.0f) { + unsigned int bits = __float_as_uint(v); + h = __ushort_as_half((unsigned short)(bits >> 16)); + } else { + h = __float2half(v); + } + dst[i] = __half2float(h); +} diff --git a/kernels/src/escha_moe_gemv_native.hip b/kernels/src/escha_moe_gemv_native.hip new file mode 100644 index 0000000000..b199f0653a --- /dev/null +++ b/kernels/src/escha_moe_gemv_native.hip @@ -0,0 +1,545 @@ +// SPDX-License-Identifier: Apache-2.0 +// hipfire — see LICENSE and NOTICE in the project root. + +#include +#include + +// Escha-W2 routed-expert GEMVs that read the TRELLIS CODE DIRECTLY (Phase 2). +// +// # Why +// +// Phase 1 decoded escha's codes to Q8_0 at load: 1.0625 B/weight resident and +// 1.0625 B/weight across the bus on every token. The codes themselves are +// 0.25 B/weight (K=2) and 0.375 B/weight (K=3). At A3B shapes the routed half +// of a decode token moves 1.07 GB as Q8_0 against 0.294 GB native, and the +// whole model is 3.01 GB/token against 2.23 — a 69 tok/s roofline against 94 +// on this box's measured 209 GB/s. Q8_0 could not reach the comparable mq4r +// SKU at ANY efficiency. So the expanded copy has to stop existing: these +// kernels decode inside the GEMV and the code is what crosses the bus. +// +// # The one structural constraint: a block owns a whole 16-wide tile column +// +// Escha's code is a 16x16 TRELLIS tile: 256 weights encoded as one 512-bit +// (K=2) or 768-bit (K=3) sliding-window bit stream. Any single weight needs a +// 16-bit window out of that stream, so a kernel that wants ONE weight from a +// tile must still touch 4 bytes of it. Decoding a tile for one output column +// therefore costs ~12 B of the 64 B tile for 16 weights (0.75 B/weight) — +// three times the format's own 0.25, and no better than Q8_0. +// +// A block that owns all SIXTEEN columns of a tile column-strip reads each tile +// exactly once and uses all 256 of its weights: 0.25 B/weight, the format's +// floor. That is why the block is 16 warps and `blockIdx.x` is a TILE column +// (`out/16` of them) rather than an output row. +// +// # Bit-exactness contract (load-bearing) +// +// Warp `w` of the block computes output row `blockIdx.x*16 + w` with the SAME +// lane->contraction-index mapping, the SAME accumulator count, the SAME loop +// order and the SAME final reduction as `escha_moe_gemv_k8_indexed.hip`. Only +// the weight's provenance changes: `__half2float(escha_cba(state))` instead of +// `d * (float)qval`. So this is bit-identical to running the Q8_0 kernel's +// arithmetic on exactly-decoded fp16 weights — which is exactly what +// `escha_gemv_f16_*` at the bottom of this file does, and what the gate +// `rdna-compute/examples/test_escha_native_gemv_gpu_vs_cpu.rs` asserts to the +// bit against both that kernel and `escha_ref`, the frozen CPU oracle. +// +// The narrow/wide split mirrors the Q8_0 kernels' `k <= 1536` rule for the +// same reason it exists there: the wide form folds FOUR interleaved +// accumulators and the narrow one, so picking the other variant for a given +// projection changes the answer. The host wrapper re-uses that threshold. +// +// # No codebook table, still +// +// `escha_cba` is five inline ops. A 65536-entry fp16 LUT would be 128 KB and +// gfx1151 has 64 KB of LDS, so there is no table here either — see +// `escha_decode_tiles.hip`, whose per-lane maths this reproduces. +// +// # K=2 and K=3 are separate, deliberately +// +// The two orders extract their 16-bit window from structurally different +// stream geometries (16 32-bit words with a 2-bit stride vs 24 words with a +// 3-bit stride and a modular wrap). They share the final "merge two words, +// shift, truncate" step because that step IS the same; everything that decides +// WHICH two words and WHICH shift is duplicated per K on purpose, and each K +// gets its own kernel entry points so neither pays a runtime branch. + +// ── the codebook, inline ───────────────────────────────────────────────────── +// Byte-for-byte the same five ops as `escha_decode_tiles.hip` and +// `escha_ref::cba`. Duplicated rather than shared: G2 asserts the GPU decode +// bit-exact against the Rust oracle, and generating either from the other +// would make that gate circular. Same reasoning here. +__device__ __forceinline__ __half escha_native_cba(unsigned short state) { + unsigned int r = ((unsigned int)state * 0xCBAC1FEDu) & 0x8FFF8FFFu; + r ^= 0x3B603B60u; + __half lo = __ushort_as_half((unsigned short)(r & 0xFFFFu)); + __half hi = __ushort_as_half((unsigned short)(r >> 16)); + return __hadd(lo, hi); // fp16 RNE add — matches the reference exactly +} + +// Load tile word `i` — the little-endian pair of int16s +// `escha_decode_tiles::escha_load_word` assembles from two `short` reads. +// +// ONE dword load, not two ushort loads. It is the same value on any +// little-endian target, and it is 2x fewer memory INSTRUCTIONS per weight, +// which is what this kernel is actually short of: at the shapes that matter it +// moves 4 MiB per launch and is nowhere near bandwidth-bound, so the vector +// memory issue rate is the limit. The compiler cannot make this merge itself — +// it would have to prove the short offset is even. +// +// Alignment: a tile begins at `(kt*tn + nt) * 32*TRELLIS_K` bytes from the +// expert's base, which is a multiple of 64 (K=2) or 96 (K=3), and every expert +// slot is at a multiple of the whole projection's byte size. So the dword view +// is always 4-byte aligned. +// +// ADDRESS SPACE 1, explicitly. `expert_ptrs[...]` is a 64-bit integer loaded at +// runtime, so the compiler cannot prove the pointer it becomes is global and +// emits `flat_load_b32` — which on gfx1151 pays an aperture check the whole +// inner loop then waits on. The cast makes it `global_load_b32`. Measured on +// the shipped gate_up projection this was worth 1.6x on its own. +typedef __attribute__((address_space(1))) const unsigned int* EschaCodePtr; + +__device__ __forceinline__ unsigned int escha_native_word(EschaCodePtr tile, int i) { + return tile[i]; +} + +// Where the state for tile element (row, col) lives. +// +// `escha_decode_tiles` runs 32 lanes x 8 values FORWARD: lane `l` decodes 8 +// states out of one 32-bit window and scatters them to 8 (row, col) slots. A +// GEMV needs the INVERSE — "which window, which shift, for THIS (row, col)" — +// because the lane that must own a weight is fixed by the contraction index, +// not by the tile geometry. +// +// Inverting the forward map (`escha_decode_tiles.hip:96-103`): +// row = 2*(l & 3) + (j & 1) + 8*((j >> 1) & 1) +// col = 2*((l >> 3) + 4*(j >= 4)) + ((l >> 2) & 1) +// gives, uniquely, +// l = ((col >> 1) & 3)*8 + (col & 1)*4 + ((row & 7) >> 1) +// j = ((col >> 3) & 1)*4 + ((row >> 3) & 1)*2 + (row & 1) +// which is what both helpers below start from. A wrong inversion yields a +// full-rank, finite, plausible weight matrix — never a fault — so it is gated +// against `escha_ref` rather than eyeballed. +__device__ __forceinline__ int escha_native_lane_of(int row, int col) { + return ((col >> 1) & 3) * 8 + (col & 1) * 4 + ((row & 7) >> 1); +} +__device__ __forceinline__ int escha_native_j_of(int row, int col) { + return ((col >> 3) & 1) * 4 + ((row >> 3) & 1) * 2 + (row & 1); +} + +// K=2: 16 words per tile, `escha_decode8_k2`'s geometry. +// +// That routine merges words (i0, i1) into a 64-bit value and takes a 32-bit +// window at `shift` (16 for even lanes, 0 for odd), then value `j` is the +// 16 bits at `2*(7-j)` inside the window. Composing the two shifts is exact: +// the widest total is 30, and 30 + 16 <= 64, so the truncation to +// `unsigned short` selects the identical bits. +__device__ __forceinline__ void escha_native_slot_k2( + int row, int col, int* word_a, int* word_b, int* shift) { + const int lane = escha_native_lane_of(row, col); + const int j = escha_native_j_of(row, col); + const int t_off = lane * 8; + const int i1 = t_off >> 4; + const int i0 = (i1 + 15) & 15; + *word_a = i0; // the high half of `merged` + *word_b = i1; // the low half + *shift = (((~t_off) & 8) << 1) + 2 * (7 - j); +} + +// K=3: 24 words per tile, `escha_decode8_k3`'s geometry — a 3-bit stride with +// a modular wrap, not a scaled copy of K=2. +// +// There the merge is of words (i0, i2), the two 32-bit windows are at `s2` +// (values 4..7) and `s2 + 12` (values 0..3), and value `j` is the 16 bits at +// `3*(3 - (j & 3))` inside its window. `s2` is in 1..32, so the largest total +// shift is 32 + 12 + 9 = 53 — under 64, so composing is defined, and the bits +// past 63 that the original's 32-bit truncation dropped are zeros either way. +__device__ __forceinline__ void escha_native_slot_k3( + int row, int col, int* word_a, int* word_b, int* shift) { + const int BITS = 3; + const int lane = escha_native_lane_of(row, col); + const int j = escha_native_j_of(row, col); + const int t_off = lane * 8; + const int b1 = (t_off + 257) * BITS; + const int b0 = b1 - 16; + const int b2 = b1 + BITS * 7; + const int i0 = b0 >> 5; + const int i2 = (b2 - 1) >> 5; + const int s2 = ((i2 + 1) << 5) - b2; + *word_a = i0 % 24; // the high half of `merged` + *word_b = i2 % 24; // the low half + *shift = s2 + (j < 4 ? BITS * 4 : 0) + BITS * (3 - (j & 3)); +} + +// ── the two accumulation forms ─────────────────────────────────────────────── +// +// Both are transcriptions of `escha_moe_gemv_k8_indexed.hip`. The ONLY edit is +// the weight expression. Everything that decides the value of the sum — which +// lane owns which contraction index, the iteration order, the accumulator +// count, the `(acc0+acc1)+(acc2+acc3)` fold, the `__shfl_down` ladder — is +// copied. `TRELLIS_K` is a template parameter so the slot maths is resolved at +// compile time and the inner loop carries no branch on it. +// +// `code` addresses tiles as `[in/16][out/16][16*TRELLIS_K]` int16 — escha's +// own IN-major grid, unchanged from the file. `nt` is this block's tile +// column, `tn = M/16` the grid's width. + +template +__device__ __forceinline__ void escha_native_slot( + int row, int col, int* word_a, int* word_b, int* shift) { + if (TRELLIS_K == 2) escha_native_slot_k2(row, col, word_a, word_b, shift); + else escha_native_slot_k3(row, col, word_a, word_b, shift); +} + +// Narrow form (Q8_0 twin: `escha_gemv_q8_0_moe_k8_indexed_batched`, K > 1536). +// One accumulator per output row; lane `t` owns contraction indices +// t, t+32, t+64, ... +template +__device__ __forceinline__ void escha_gemv_native_narrow_body( + const unsigned long long* __restrict__ expert_ptrs, + const int* __restrict__ topk_indices, + const float* __restrict__ x_batch, + float* __restrict__ y_batch, + int M, int K, int nt_major +) { + const int warp = threadIdx.x >> 5; + const int lane = threadIdx.x & 31; + // Two output rows per warp: tile columns `warp` and `warp + 8`. + const int row_lo = blockIdx.x * 16 + warp; + const int row_hi = row_lo + 8; + if (row_hi >= M) return; + const int krank = blockIdx.y; + + EschaCodePtr code = (EschaCodePtr)(expert_ptrs[topk_indices[krank]]); + const float* __restrict__ x = x_batch + (size_t)krank * K; + float* __restrict__ y = y_batch + (size_t)krank * M; + + // Loop-invariant: the block's tile column, the lane's row inside a tile, + // and which of the iteration's two tiles the lane falls in. The slot + // descriptor depends only on (row-in-tile, column-in-tile), so it is + // hoisted out of the contraction loop entirely. + const int tn = M >> 4; + const int nt = blockIdx.x; + const int in_row = lane & 15; + const int tile_hi = lane >> 4; + int word_a, word_b, shift_lo, shift_hi; + int word_a2, word_b2; + escha_native_slot(in_row, warp, &word_a, &word_b, &shift_lo); + escha_native_slot(in_row, warp + 8, &word_a2, &word_b2, &shift_hi); + // Columns `c` and `c + 8` share a window — see the comment on this + // function's grid, and the algebra: the lane index the tile decoder uses is + // `((col>>1)&3)*8 + (col&1)*4 + ((row&7)>>1)`, and `(col>>1)&3` is + // unchanged by +8 while `col&1` is unchanged outright. Only `j` moves, by + // 4, which is a shift. If that ever stops holding, the two columns would + // silently read each other's weights, so poison the output loudly instead: + // every escha gate asserts finiteness, and nothing else in the pipeline + // manufactures a NaN. + if (word_a2 != word_a || word_b2 != word_b) { + if (lane == 0) { y[row_lo] = __builtin_nanf(""); y[row_hi] = __builtin_nanf(""); } + return; + } + + // The tile grid comes in one of two orders and `nt_major` says which. Both + // reduce to a pair of loop-invariant strides, so there is no hot-loop cost + // and no second kernel: + // + // kt-major `[ic/16][oc/16]` (MoE, and any un-permuted code): a `kt` step + // crosses a whole tile ROW, `tn * 8*TRELLIS_K`. + // nt-major `[oc/16][ic/16]` (dense, permuted at load by + // `escha_tiles_to_nt_major`): a `kt` step is ONE TILE, so the four + // tiles a quad reads sit 32 dwords apart at K=2 instead of two 139 KB + // strides. Worth a measured 24% on this kernel. + // + // This lane's tile in iteration `bi` is `kt = bi*2 + tile_hi`. + const size_t ktiles = (size_t)K >> 4; + const size_t strip = nt_major ? (size_t)(8 * TRELLIS_K) + : (size_t)tn * (8 * TRELLIS_K); + const size_t colstride = nt_major ? ktiles * (size_t)(8 * TRELLIS_K) + : (size_t)(8 * TRELLIS_K); + EschaCodePtr strip0 = code + (size_t)nt * colstride + (size_t)tile_hi * strip; + + const int blocks_per_row = K / 32; + float sum_lo = 0.0f, sum_hi = 0.0f; + // UNROLL 4 is not cosmetic. Un-unrolled, the compiler emits the two code + // loads and then `s_waitcnt vmcnt(1)` in the same iteration: the loop has + // exactly one memory round trip of latency to hide per 42 instructions and + // nothing to hide it with. Unrolling puts eight loads in flight before the + // first wait. It does NOT touch the arithmetic: `sum_lo`/`sum_hi` stay a + // single sequential FMA chain each, in `bi` order, which is what the + // bit-exactness contract is about. + // + // The wide form needs no such pragma — its quad structure already carries + // four independent iterations, which is exactly why it was measured at + // 105 GB/s while this one sat at 33. + #pragma unroll 4 + for (int bi = 0; bi < blocks_per_row; bi++) { + EschaCodePtr tile = strip0 + (size_t)bi * 2 * strip; + const unsigned long long merged = + ((unsigned long long)escha_native_word(tile, word_a) << 32) | + escha_native_word(tile, word_b); + const __half w_lo = escha_native_cba((unsigned short)(merged >> shift_lo)); + const __half w_hi = escha_native_cba((unsigned short)(merged >> shift_hi)); + const float xv = x[bi * 32 + lane]; + sum_lo += __half2float(w_lo) * xv; + sum_hi += __half2float(w_hi) * xv; + } + + for (int offset = 16; offset > 0; offset >>= 1) { + sum_lo += __shfl_down(sum_lo, offset); + sum_hi += __shfl_down(sum_hi, offset); + } + if (lane == 0) { + y[row_lo] = sum_lo; + y[row_hi] = sum_hi; + } +} + +// Wide form (Q8_0 twin: `escha_gemv_q8_0_wide_moe_k8_indexed_batched`, +// K <= 1536). Four interleaved accumulators per output row with the +// `(acc0+acc1)+(acc2+acc3)` fold. The tail's `acc[t]` assignment is copied +// verbatim, including the fact that `t == 3` is unreachable. +template +__device__ __forceinline__ void escha_gemv_native_wide_body( + const unsigned long long* __restrict__ expert_ptrs, + const int* __restrict__ topk_indices, + const float* __restrict__ x_batch, + float* __restrict__ y_batch, + int M, int K, int nt_major +) { + const int warp = threadIdx.x >> 5; + const int lane = threadIdx.x & 31; + const int row_lo = blockIdx.x * 16 + warp; + const int row_hi = row_lo + 8; + if (row_hi >= M) return; + const int krank = blockIdx.y; + + EschaCodePtr code = (EschaCodePtr)(expert_ptrs[topk_indices[krank]]); + const float* __restrict__ x = x_batch + (size_t)krank * K; + float* __restrict__ y = y_batch + (size_t)krank * M; + + const int tn = M >> 4; + const int nt = blockIdx.x; + const int in_row = lane & 15; + const int tile_hi = lane >> 4; + int word_a, word_b, shift_lo, shift_hi; + int word_a2, word_b2; + escha_native_slot(in_row, warp, &word_a, &word_b, &shift_lo); + escha_native_slot(in_row, warp + 8, &word_a2, &word_b2, &shift_hi); + // See the narrow form: columns `c` and `c + 8` provably share a window. + if (word_a2 != word_a || word_b2 != word_b) { + if (lane == 0) { y[row_lo] = __builtin_nanf(""); y[row_hi] = __builtin_nanf(""); } + return; + } + + // The tile grid comes in one of two orders and `nt_major` says which. Both + // reduce to a pair of loop-invariant strides, so there is no hot-loop cost + // and no second kernel: + // + // kt-major `[ic/16][oc/16]` (MoE, and any un-permuted code): a `kt` step + // crosses a whole tile ROW, `tn * 8*TRELLIS_K`. + // nt-major `[oc/16][ic/16]` (dense, permuted at load by + // `escha_tiles_to_nt_major`): a `kt` step is ONE TILE, so the four + // tiles a quad reads sit 32 dwords apart at K=2 instead of two 139 KB + // strides. Worth a measured 24% on this kernel. + // + // This lane's tile in iteration `bi` is `kt = bi*2 + tile_hi`. + const size_t ktiles = (size_t)K >> 4; + const size_t strip = nt_major ? (size_t)(8 * TRELLIS_K) + : (size_t)tn * (8 * TRELLIS_K); + const size_t colstride = nt_major ? ktiles * (size_t)(8 * TRELLIS_K) + : (size_t)(8 * TRELLIS_K); + EschaCodePtr strip0 = code + (size_t)nt * colstride + (size_t)tile_hi * strip; + + const int blocks_per_row = K / 32; + float lo0 = 0.0f, lo1 = 0.0f, lo2 = 0.0f, lo3 = 0.0f; + float hi0 = 0.0f, hi1 = 0.0f, hi2 = 0.0f, hi3 = 0.0f; + const int quads = blocks_per_row >> 2; + const int tail = blocks_per_row & 3; + + for (int q = 0; q < quads; q++) { + const int bi = q << 2; + EschaCodePtr t0 = strip0 + (size_t)(bi + 0) * 2 * strip; + EschaCodePtr t1 = strip0 + (size_t)(bi + 1) * 2 * strip; + EschaCodePtr t2 = strip0 + (size_t)(bi + 2) * 2 * strip; + EschaCodePtr t3 = strip0 + (size_t)(bi + 3) * 2 * strip; + + const unsigned long long m0 = + ((unsigned long long)escha_native_word(t0, word_a) << 32) | + escha_native_word(t0, word_b); + const unsigned long long m1 = + ((unsigned long long)escha_native_word(t1, word_a) << 32) | + escha_native_word(t1, word_b); + const unsigned long long m2 = + ((unsigned long long)escha_native_word(t2, word_a) << 32) | + escha_native_word(t2, word_b); + const unsigned long long m3 = + ((unsigned long long)escha_native_word(t3, word_a) << 32) | + escha_native_word(t3, word_b); + + const float x0 = x[(bi + 0) * 32 + lane]; + const float x1 = x[(bi + 1) * 32 + lane]; + const float x2 = x[(bi + 2) * 32 + lane]; + const float x3 = x[(bi + 3) * 32 + lane]; + + lo0 += __half2float(escha_native_cba((unsigned short)(m0 >> shift_lo))) * x0; + lo1 += __half2float(escha_native_cba((unsigned short)(m1 >> shift_lo))) * x1; + lo2 += __half2float(escha_native_cba((unsigned short)(m2 >> shift_lo))) * x2; + lo3 += __half2float(escha_native_cba((unsigned short)(m3 >> shift_lo))) * x3; + + hi0 += __half2float(escha_native_cba((unsigned short)(m0 >> shift_hi))) * x0; + hi1 += __half2float(escha_native_cba((unsigned short)(m1 >> shift_hi))) * x1; + hi2 += __half2float(escha_native_cba((unsigned short)(m2 >> shift_hi))) * x2; + hi3 += __half2float(escha_native_cba((unsigned short)(m3 >> shift_hi))) * x3; + } + + for (int t = 0; t < tail; t++) { + const int bi = (quads << 2) + t; + EschaCodePtr tp = strip0 + (size_t)bi * 2 * strip; + const unsigned long long m = + ((unsigned long long)escha_native_word(tp, word_a) << 32) | + escha_native_word(tp, word_b); + const float xv = x[bi * 32 + lane]; + const float c_lo = + __half2float(escha_native_cba((unsigned short)(m >> shift_lo))) * xv; + const float c_hi = + __half2float(escha_native_cba((unsigned short)(m >> shift_hi))) * xv; + if (t == 0) { lo0 += c_lo; hi0 += c_hi; } + else if (t == 1) { lo1 += c_lo; hi1 += c_hi; } + else if (t == 2) { lo2 += c_lo; hi2 += c_hi; } + } + + float sum_lo = (lo0 + lo1) + (lo2 + lo3); + float sum_hi = (hi0 + hi1) + (hi2 + hi3); + for (int offset = 16; offset > 0; offset >>= 1) { + sum_lo += __shfl_down(sum_lo, offset); + sum_hi += __shfl_down(sum_hi, offset); + } + if (lane == 0) { + y[row_lo] = sum_lo; + y[row_hi] = sum_hi; + } +} + +// ── entry points ───────────────────────────────────────────────────────────── +// Grid: (M/16, slots, 1). Block: 256 — EIGHT warps, each covering two columns +// of the 16-wide strip (`w` and `w + 8`), which share a decode window. The +// host wrapper enforces `M % 16 == 0` and `K % 32 == 0`. + +#define ESCHA_NATIVE_GEMV(NAME, FORM, TK) \ + __launch_bounds__(256) \ + extern "C" __global__ void NAME( \ + const unsigned long long* __restrict__ expert_ptrs, \ + const int* __restrict__ topk_indices, \ + const float* __restrict__ x_batch, \ + float* __restrict__ y_batch, \ + int M, int K, int nt_major \ + ) { \ + FORM(expert_ptrs, topk_indices, x_batch, y_batch, M, K, nt_major); \ + } + +ESCHA_NATIVE_GEMV(escha_gemv_native_k2_moe_k8_indexed_batched, + escha_gemv_native_narrow_body, 2) +ESCHA_NATIVE_GEMV(escha_gemv_native_k3_moe_k8_indexed_batched, + escha_gemv_native_narrow_body, 3) +ESCHA_NATIVE_GEMV(escha_gemv_native_k2_wide_moe_k8_indexed_batched, + escha_gemv_native_wide_body, 2) +ESCHA_NATIVE_GEMV(escha_gemv_native_k3_wide_moe_k8_indexed_batched, + escha_gemv_native_wide_body, 3) + +// ── the F16 reference arm ──────────────────────────────────────────────────── +// +// These read an OUT-major `[M, K]` fp16 expert slot — exactly what +// `escha_bare_to_f16` writes, i.e. the exactly-decoded weights with nothing +// re-quantised. Their whole purpose is to be the thing the fused kernels above +// are compared against: same grid, same block, same lane mapping, same +// accumulator structure, same reduction, and the weight expression written +// character-for-character the same (`__half2float(w) * x[...]`), so the ONLY +// difference between the two is where `w` came from. Equality is therefore a +// statement about the DECODE, not about floating-point luck. +// +// They are a gate arm, not a production route: `EschaWeightStore::F16` still +// runs host-routed. Admitting F16 to the indexed path would change the arm the +// G5 KLD reference is built from, and that reference has published numbers. + +template +__device__ __forceinline__ void escha_gemv_f16_body( + const unsigned long long* __restrict__ expert_ptrs, + const int* __restrict__ topk_indices, + const float* __restrict__ x_batch, + float* __restrict__ y_batch, + int M, int K +) { + const int warp = threadIdx.x >> 5; + const int lane = threadIdx.x & 31; + const int krank = blockIdx.y; + + const __half* __restrict__ A = + reinterpret_cast(expert_ptrs[topk_indices[krank]]); + const float* __restrict__ x = x_batch + (size_t)krank * K; + float* __restrict__ y = y_batch + (size_t)krank * M; + + // Same two-rows-per-warp geometry as the fused kernels, so the two are + // launch-compatible and differ ONLY in where the weight comes from. + #pragma unroll + for (int half = 0; half < 2; half++) { + const int row = blockIdx.x * 16 + warp + half * 8; + if (row >= M) continue; + const __half* __restrict__ row_data = A + (size_t)row * K; + const int blocks_per_row = K / 32; + float sum = 0.0f; + if (!WIDE) { + for (int bi = 0; bi < blocks_per_row; bi++) { + const __half w = row_data[bi * 32 + lane]; + sum += __half2float(w) * x[bi * 32 + lane]; + } + } else { + float acc0 = 0.0f, acc1 = 0.0f, acc2 = 0.0f, acc3 = 0.0f; + const int quads = blocks_per_row >> 2; + const int tail = blocks_per_row & 3; + for (int q = 0; q < quads; q++) { + const int bi = q << 2; + const __half w0 = row_data[(bi + 0) * 32 + lane]; + const __half w1 = row_data[(bi + 1) * 32 + lane]; + const __half w2 = row_data[(bi + 2) * 32 + lane]; + const __half w3 = row_data[(bi + 3) * 32 + lane]; + acc0 += __half2float(w0) * x[(bi + 0) * 32 + lane]; + acc1 += __half2float(w1) * x[(bi + 1) * 32 + lane]; + acc2 += __half2float(w2) * x[(bi + 2) * 32 + lane]; + acc3 += __half2float(w3) * x[(bi + 3) * 32 + lane]; + } + for (int t = 0; t < tail; t++) { + const int bi = (quads << 2) + t; + const __half w = row_data[bi * 32 + lane]; + const float contrib = __half2float(w) * x[bi * 32 + lane]; + if (t == 0) acc0 += contrib; + else if (t == 1) acc1 += contrib; + else if (t == 2) acc2 += contrib; + } + sum = (acc0 + acc1) + (acc2 + acc3); + } + for (int offset = 16; offset > 0; offset >>= 1) + sum += __shfl_down(sum, offset); + if (lane == 0) y[row] = sum; + } +} + +__launch_bounds__(256) +extern "C" __global__ void escha_gemv_f16_moe_k8_indexed_batched( + const unsigned long long* __restrict__ expert_ptrs, + const int* __restrict__ topk_indices, + const float* __restrict__ x_batch, + float* __restrict__ y_batch, + int M, int K, int nt_major +) { + escha_gemv_f16_body(expert_ptrs, topk_indices, x_batch, y_batch, M, K); +} + +__launch_bounds__(256) +extern "C" __global__ void escha_gemv_f16_wide_moe_k8_indexed_batched( + const unsigned long long* __restrict__ expert_ptrs, + const int* __restrict__ topk_indices, + const float* __restrict__ x_batch, + float* __restrict__ y_batch, + int M, int K, int nt_major +) { + escha_gemv_f16_body(expert_ptrs, topk_indices, x_batch, y_batch, M, K); +} diff --git a/kernels/src/gemm_gate_up_hfq6g256_wmma.hip b/kernels/src/gemm_gate_up_hfq6g256_wmma.hip index b75b01646c..72dbeebbd2 100644 --- a/kernels/src/gemm_gate_up_hfq6g256_wmma.hip +++ b/kernels/src/gemm_gate_up_hfq6g256_wmma.hip @@ -87,7 +87,14 @@ extern "C" __global__ void gemm_gate_up_hfq6g256_wmma( unsigned int d0a = *(const unsigned int*)(dp_a); unsigned int d1a = *(const unsigned int*)(dp_a + 3); unsigned int d2a = *(const unsigned int*)(dp_a + 6); - unsigned int d3a = *(const unsigned int*)(dp_a + 9); + // Bytes 9,10,11 of the tile, NOT a dword at +9 — see the same fix + // in gemm_qkv_mq6g256v2_wmma.hip. The four 6-bit values this feeds + // sit at shifts 0/6/12/18, so only bits 0..23 are ever used and + // byte 12 is read then discarded. At kt=15 a +9 dword reads one + // byte past the 200-byte group, and on the last group of the last + // row that is one byte past the weight tensor. Load at +8 and + // shift down: identical bits 0..23, never leaves the tile. + unsigned int d3a = (*(const unsigned int*)(dp_a + 8)) >> 8; // === Tile B (kt+1) — start loading early === const int byte_off_b = (kt + 1) * 12; @@ -95,7 +102,7 @@ extern "C" __global__ void gemm_gate_up_hfq6g256_wmma( unsigned int d0b = *(const unsigned int*)(dp_b); unsigned int d1b = *(const unsigned int*)(dp_b + 3); unsigned int d2b = *(const unsigned int*)(dp_b + 6); - unsigned int d3b = *(const unsigned int*)(dp_b + 9); + unsigned int d3b = (*(const unsigned int*)(dp_b + 8)) >> 8; // see d3a // Load both X tiles early (consecutive addresses — prefetcher-friendly) half16_t b_a = *(const half16_t*)(x_base + kt * 16); diff --git a/kernels/src/gemm_gate_up_mq6g256v2_wmma.hip b/kernels/src/gemm_gate_up_mq6g256v2_wmma.hip index dac4e9f297..5847b03a52 100644 --- a/kernels/src/gemm_gate_up_mq6g256v2_wmma.hip +++ b/kernels/src/gemm_gate_up_mq6g256v2_wmma.hip @@ -94,14 +94,21 @@ extern "C" __global__ void gemm_gate_up_mq6g256v2_wmma( unsigned int d0a = *(const unsigned int*)(dp_a); unsigned int d1a = *(const unsigned int*)(dp_a + 3); unsigned int d2a = *(const unsigned int*)(dp_a + 6); - unsigned int d3a = *(const unsigned int*)(dp_a + 9); + // Bytes 9,10,11 of the tile, NOT a dword at +9 — see the same fix + // in gemm_qkv_mq6g256v2_wmma.hip. The four 6-bit values this feeds + // sit at shifts 0/6/12/18, so only bits 0..23 are ever used and + // byte 12 is read then discarded. At kt=15 a +9 dword reads one + // byte past the 200-byte group, and on the last group of the last + // row that is one byte past the weight tensor. Load at +8 and + // shift down: identical bits 0..23, never leaves the tile. + unsigned int d3a = (*(const unsigned int*)(dp_a + 8)) >> 8; const int byte_off_b = (kt + 1) * 12; const unsigned char* dp_b = (const unsigned char*)(gp + 8 + byte_off_b); unsigned int d0b = *(const unsigned int*)(dp_b); unsigned int d1b = *(const unsigned int*)(dp_b + 3); unsigned int d2b = *(const unsigned int*)(dp_b + 6); - unsigned int d3b = *(const unsigned int*)(dp_b + 9); + unsigned int d3b = (*(const unsigned int*)(dp_b + 8)) >> 8; // see d3a half16_t b_a = *(const half16_t*)(x_base + kt * 16); half16_t b_b = *(const half16_t*)(x_base + (kt + 1) * 16); diff --git a/kernels/src/gemm_hfq6g256_moe_grouped_wmma.gfx1151.hip b/kernels/src/gemm_hfq6g256_moe_grouped_wmma.gfx1151.hip index cacc0baa79..542b4cf870 100644 --- a/kernels/src/gemm_hfq6g256_moe_grouped_wmma.gfx1151.hip +++ b/kernels/src/gemm_hfq6g256_moe_grouped_wmma.gfx1151.hip @@ -78,11 +78,18 @@ extern "C" __global__ void gemm_hfq6g256_moe_grouped_wmma_gfx1151( unsigned int d0a = *(const unsigned int*)(dp_a); unsigned int d1a = *(const unsigned int*)(dp_a + 3); unsigned int d2a = *(const unsigned int*)(dp_a + 6); - unsigned int d3a = *(const unsigned int*)(dp_a + 9); + // Bytes 9,10,11 of the tile, NOT a dword at +9 — see the same fix + // in gemm_qkv_mq6g256v2_wmma.hip. The four 6-bit values this feeds + // sit at shifts 0/6/12/18, so only bits 0..23 are ever used and + // byte 12 is read then discarded. At kt=15 a +9 dword reads one + // byte past the 200-byte group, and on the last group of the last + // row that is one byte past the weight tensor. Load at +8 and + // shift down: identical bits 0..23, never leaves the tile. + unsigned int d3a = (*(const unsigned int*)(dp_a + 8)) >> 8; unsigned int d0b = *(const unsigned int*)(dp_b); unsigned int d1b = *(const unsigned int*)(dp_b + 3); unsigned int d2b = *(const unsigned int*)(dp_b + 6); - unsigned int d3b = *(const unsigned int*)(dp_b + 9); + unsigned int d3b = (*(const unsigned int*)(dp_b + 8)) >> 8; // see d3a half16_t b_a, b_b; if (xg != nullptr) { diff --git a/kernels/src/gemm_hfq6g256_residual_wmma_k2.hip b/kernels/src/gemm_hfq6g256_residual_wmma_k2.hip index da4e64d486..637aa40c08 100644 --- a/kernels/src/gemm_hfq6g256_residual_wmma_k2.hip +++ b/kernels/src/gemm_hfq6g256_residual_wmma_k2.hip @@ -78,7 +78,14 @@ extern "C" __global__ void gemm_hfq6g256_residual_wmma_k2( unsigned int d0a = *(const unsigned int*)(dp_a); unsigned int d1a = *(const unsigned int*)(dp_a + 3); unsigned int d2a = *(const unsigned int*)(dp_a + 6); - unsigned int d3a = *(const unsigned int*)(dp_a + 9); + // Bytes 9,10,11 of the tile, NOT a dword at +9 — see the same fix + // in gemm_qkv_mq6g256v2_wmma.hip. The four 6-bit values this feeds + // sit at shifts 0/6/12/18, so only bits 0..23 are ever used and + // byte 12 is read then discarded. At kt=15 a +9 dword reads one + // byte past the 200-byte group, and on the last group of the last + // row that is one byte past the weight tensor. Load at +8 and + // shift down: identical bits 0..23, never leaves the tile. + unsigned int d3a = (*(const unsigned int*)(dp_a + 8)) >> 8; // === Tile B (kt+1) -- start loading early === const int byte_off_b = (kt + 1) * 12; @@ -86,7 +93,7 @@ extern "C" __global__ void gemm_hfq6g256_residual_wmma_k2( unsigned int d0b = *(const unsigned int*)(dp_b); unsigned int d1b = *(const unsigned int*)(dp_b + 3); unsigned int d2b = *(const unsigned int*)(dp_b + 6); - unsigned int d3b = *(const unsigned int*)(dp_b + 9); + unsigned int d3b = (*(const unsigned int*)(dp_b + 8)) >> 8; // see d3a // Load both X tiles early (consecutive addresses -- prefetcher-friendly) half16_t b_a = *(const half16_t*)(xg + kt * 16); diff --git a/kernels/src/gemm_mq6g256v2_residual_wmma.hip b/kernels/src/gemm_mq6g256v2_residual_wmma.hip index b11cc3c9da..89e1cfcaeb 100644 --- a/kernels/src/gemm_mq6g256v2_residual_wmma.hip +++ b/kernels/src/gemm_mq6g256v2_residual_wmma.hip @@ -90,7 +90,14 @@ extern "C" __global__ void gemm_mq6g256v2_residual_wmma( unsigned int d0a = *(const unsigned int*)(dp_a); unsigned int d1a = *(const unsigned int*)(dp_a + 3); unsigned int d2a = *(const unsigned int*)(dp_a + 6); - unsigned int d3a = *(const unsigned int*)(dp_a + 9); + // Bytes 9,10,11 of the tile, NOT a dword at +9 — see the same fix + // in gemm_qkv_mq6g256v2_wmma.hip. The four 6-bit values this feeds + // sit at shifts 0/6/12/18, so only bits 0..23 are ever used and + // byte 12 is read then discarded. At kt=15 a +9 dword reads one + // byte past the 200-byte group, and on the last group of the last + // row that is one byte past the weight tensor. Load at +8 and + // shift down: identical bits 0..23, never leaves the tile. + unsigned int d3a = (*(const unsigned int*)(dp_a + 8)) >> 8; // === Tile B (kt+1) — start loading early === const int byte_off_b = (kt + 1) * 12; @@ -98,7 +105,7 @@ extern "C" __global__ void gemm_mq6g256v2_residual_wmma( unsigned int d0b = *(const unsigned int*)(dp_b); unsigned int d1b = *(const unsigned int*)(dp_b + 3); unsigned int d2b = *(const unsigned int*)(dp_b + 6); - unsigned int d3b = *(const unsigned int*)(dp_b + 9); + unsigned int d3b = (*(const unsigned int*)(dp_b + 8)) >> 8; // see d3a half16_t b_a = *(const half16_t*)(xg + kt * 16); half16_t b_b = *(const half16_t*)(xg + (kt + 1) * 16); diff --git a/kernels/src/gemm_qkv_hfq6g256_wmma.hip b/kernels/src/gemm_qkv_hfq6g256_wmma.hip index 15d1912940..1068e5c8b1 100644 --- a/kernels/src/gemm_qkv_hfq6g256_wmma.hip +++ b/kernels/src/gemm_qkv_hfq6g256_wmma.hip @@ -89,7 +89,14 @@ extern "C" __global__ void gemm_qkv_hfq6g256_wmma( unsigned int d0a = *(const unsigned int*)(dp_a); unsigned int d1a = *(const unsigned int*)(dp_a + 3); unsigned int d2a = *(const unsigned int*)(dp_a + 6); - unsigned int d3a = *(const unsigned int*)(dp_a + 9); + // Bytes 9,10,11 of the tile, NOT a dword at +9 — see the same fix + // in gemm_qkv_mq6g256v2_wmma.hip. The four 6-bit values this feeds + // sit at shifts 0/6/12/18, so only bits 0..23 are ever used and + // byte 12 is read then discarded. At kt=15 a +9 dword reads one + // byte past the 200-byte group, and on the last group of the last + // row that is one byte past the weight tensor. Load at +8 and + // shift down: identical bits 0..23, never leaves the tile. + unsigned int d3a = (*(const unsigned int*)(dp_a + 8)) >> 8; // === Tile B (kt+1) — start loading early === const int byte_off_b = (kt + 1) * 12; @@ -97,7 +104,7 @@ extern "C" __global__ void gemm_qkv_hfq6g256_wmma( unsigned int d0b = *(const unsigned int*)(dp_b); unsigned int d1b = *(const unsigned int*)(dp_b + 3); unsigned int d2b = *(const unsigned int*)(dp_b + 6); - unsigned int d3b = *(const unsigned int*)(dp_b + 9); + unsigned int d3b = (*(const unsigned int*)(dp_b + 8)) >> 8; // see d3a // Load both X tiles early (consecutive addresses — prefetcher-friendly) half16_t b_a = *(const half16_t*)(x_base + kt * 16); diff --git a/kernels/src/gemm_qkv_mq6g256v2_wmma.hip b/kernels/src/gemm_qkv_mq6g256v2_wmma.hip index 0212065dfb..347c26fb66 100644 --- a/kernels/src/gemm_qkv_mq6g256v2_wmma.hip +++ b/kernels/src/gemm_qkv_mq6g256v2_wmma.hip @@ -97,14 +97,25 @@ extern "C" __global__ void gemm_qkv_mq6g256v2_wmma( unsigned int d0a = *(const unsigned int*)(dp_a); unsigned int d1a = *(const unsigned int*)(dp_a + 3); unsigned int d2a = *(const unsigned int*)(dp_a + 6); - unsigned int d3a = *(const unsigned int*)(dp_a + 9); + // Bytes 9,10,11 of the tile, NOT a dword at +9. The four 6-bit + // values this dword feeds live at shifts 0/6/12/18, so only bits + // 0..23 — bytes 9,10,11 — are ever used; byte 12 is read and + // discarded. Reading at +9 therefore touches one byte PAST the + // 12-byte K-tile, and for kt=15 that is one byte past the whole + // 200-byte group. On the last group of the last row that runs off + // the end of the weight tensor: measured as an illegal memory + // access on escha-35b (q_m=8192, k=2048) for any N < 96, which is + // exactly where this base kernel is reached — N >= 96 routes to + // the batch-tile path and hid the bug. Load the dword at +8 and + // shift down instead: same bits 0..23, never leaves the tile. + unsigned int d3a = (*(const unsigned int*)(dp_a + 8)) >> 8; const int byte_off_b = (kt + 1) * 12; const unsigned char* dp_b = (const unsigned char*)(gp + 8 + byte_off_b); unsigned int d0b = *(const unsigned int*)(dp_b); unsigned int d1b = *(const unsigned int*)(dp_b + 3); unsigned int d2b = *(const unsigned int*)(dp_b + 6); - unsigned int d3b = *(const unsigned int*)(dp_b + 9); + unsigned int d3b = (*(const unsigned int*)(dp_b + 8)) >> 8; // see d3a half16_t b_a = *(const half16_t*)(x_base + kt * 16); half16_t b_b = *(const half16_t*)(x_base + (kt + 1) * 16); diff --git a/kernels/src/gemm_qkvza_hfq6g256_wmma.hip b/kernels/src/gemm_qkvza_hfq6g256_wmma.hip index c9ededa8c2..2ebd815c69 100644 --- a/kernels/src/gemm_qkvza_hfq6g256_wmma.hip +++ b/kernels/src/gemm_qkvza_hfq6g256_wmma.hip @@ -93,7 +93,14 @@ extern "C" __global__ void gemm_qkvza_hfq6g256_wmma( unsigned int d0a = *(const unsigned int*)(dp_a); unsigned int d1a = *(const unsigned int*)(dp_a + 3); unsigned int d2a = *(const unsigned int*)(dp_a + 6); - unsigned int d3a = *(const unsigned int*)(dp_a + 9); + // Bytes 9,10,11 of the tile, NOT a dword at +9 — see the same fix + // in gemm_qkv_mq6g256v2_wmma.hip. The four 6-bit values this feeds + // sit at shifts 0/6/12/18, so only bits 0..23 are ever used and + // byte 12 is read then discarded. At kt=15 a +9 dword reads one + // byte past the 200-byte group, and on the last group of the last + // row that is one byte past the weight tensor. Load at +8 and + // shift down: identical bits 0..23, never leaves the tile. + unsigned int d3a = (*(const unsigned int*)(dp_a + 8)) >> 8; // === Tile B (kt+1) — start loading early === const int byte_off_b = (kt + 1) * 12; @@ -101,7 +108,7 @@ extern "C" __global__ void gemm_qkvza_hfq6g256_wmma( unsigned int d0b = *(const unsigned int*)(dp_b); unsigned int d1b = *(const unsigned int*)(dp_b + 3); unsigned int d2b = *(const unsigned int*)(dp_b + 6); - unsigned int d3b = *(const unsigned int*)(dp_b + 9); + unsigned int d3b = (*(const unsigned int*)(dp_b + 8)) >> 8; // see d3a // Load both X tiles early (consecutive addresses — prefetcher-friendly) half16_t b_a = *(const half16_t*)(x_base + kt * 16); diff --git a/kernels/src/gemm_qkvza_mq6g256v2_wmma.hip b/kernels/src/gemm_qkvza_mq6g256v2_wmma.hip index 09a8d24173..73647d8989 100644 --- a/kernels/src/gemm_qkvza_mq6g256v2_wmma.hip +++ b/kernels/src/gemm_qkvza_mq6g256v2_wmma.hip @@ -101,14 +101,21 @@ extern "C" __global__ void gemm_qkvza_mq6g256v2_wmma( unsigned int d0a = *(const unsigned int*)(dp_a); unsigned int d1a = *(const unsigned int*)(dp_a + 3); unsigned int d2a = *(const unsigned int*)(dp_a + 6); - unsigned int d3a = *(const unsigned int*)(dp_a + 9); + // Bytes 9,10,11 of the tile, NOT a dword at +9 — see the same fix + // in gemm_qkv_mq6g256v2_wmma.hip. The four 6-bit values this feeds + // sit at shifts 0/6/12/18, so only bits 0..23 are ever used and + // byte 12 is read then discarded. At kt=15 a +9 dword reads one + // byte past the 200-byte group, and on the last group of the last + // row that is one byte past the weight tensor. Load at +8 and + // shift down: identical bits 0..23, never leaves the tile. + unsigned int d3a = (*(const unsigned int*)(dp_a + 8)) >> 8; const int byte_off_b = (kt + 1) * 12; const unsigned char* dp_b = (const unsigned char*)(gp + 8 + byte_off_b); unsigned int d0b = *(const unsigned int*)(dp_b); unsigned int d1b = *(const unsigned int*)(dp_b + 3); unsigned int d2b = *(const unsigned int*)(dp_b + 6); - unsigned int d3b = *(const unsigned int*)(dp_b + 9); + unsigned int d3b = (*(const unsigned int*)(dp_b + 8)) >> 8; // see d3a half16_t b_a = *(const half16_t*)(x_base + kt * 16); half16_t b_b = *(const half16_t*)(x_base + (kt + 1) * 16); diff --git a/kernels/src/gemv_mq6g256v2_multirow.hip b/kernels/src/gemv_mq6g256v2_multirow.hip index 2d0b69d316..a000685316 100644 --- a/kernels/src/gemv_mq6g256v2_multirow.hip +++ b/kernels/src/gemv_mq6g256v2_multirow.hip @@ -258,6 +258,20 @@ __device__ __forceinline__ void gemv_mq6g256v2_multirow_body( // same math (4-way acc[g%4] + pairwise combine, 5302926) and same // branchless dual-row clamp, but removes x-load → DOG serialization so // dual-issue can cover more of the weight-load latency under L2-hit. +// R=2 goes through the SHARED TEMPLATE, like R=4 and R=8. +// +// It used to have its own hand-written dual-row body — an lm-head experiment +// carrying raw-buffer VMEM loads, an X-buffer SRD, and a compile-time K=2048 +// specialization, all selected by macros that only fire on gfx1100/gfx1201. +// On any other target (gfx1151 among them) every one of those #ifs took its +// `#else`, leaving a pointer-load path that was never validated: it produced +// KLD 10.64 / PPL 248,303 against 0.0071 for the R=4 and R=8 template bodies +// on the same weights. The entry point was unreferenced by dispatch, so +// nothing caught it. +// +// If the specialization is revived for gfx1201/gfx1100, gate it on those +// targets explicitly and keep this template as the fallback — do not let a +// tuned variant become the only implementation of a width again. __launch_bounds__(32, 18) extern "C" __global__ void gemv_mq6g256v2_multirow_r2( const char* __restrict__ A, @@ -265,169 +279,7 @@ extern "C" __global__ void gemv_mq6g256v2_multirow_r2( float* __restrict__ y, int M, int K ) { - const int row0 = blockIdx.x << 1; - if (row0 >= M) return; - const int row1 = row0 + 1; - const bool have_row1 = row1 < M; - const int tid = threadIdx.x; - -#if defined(HIPFIRE_GFX1151_LM_HEAD_K2048) - const int groups_per_row = 8; -#else - const int groups_per_row = K / 256; -#endif - const long long row_stride = (long long)groups_per_row * 200; - // Clamp OOB row1 to row0 so weight reads stay in-bounds; store is masked. - const char* row_ptr0 = A + (long long)row0 * row_stride; - const char* row_ptr1 = A + (long long)(have_row1 ? row1 : row0) * row_stride; - const int boff = tid * 6; - const unsigned int row_weight_offset0 = (unsigned int)((long long)row0 * row_stride); - const unsigned int row_weight_offset1 = - (unsigned int)((long long)(have_row1 ? row1 : row0) * row_stride); -#if HIPFIRE_GFX12_WEIGHT_BUFFER_LOADS - // The full lm-head matrix fits the raw-buffer 32-bit offset range; one SRD - // avoids extending the live descriptor set in this 94-VGPR kernel. - const auto weight_rsrc = HIPFIRE_WEIGHT_RSRC(A); -#endif -#if defined(HIPFIRE_GFX1151_LM_HEAD_X_BUFFER) - const auto x_rsrc = __builtin_amdgcn_make_buffer_rsrc( - const_cast(x), 0, 0xffffffffu, 0x31004000u); -#endif - - float a0 = 0.0f, a1 = 0.0f, a2 = 0.0f, a3 = 0.0f; - float b0 = 0.0f, b1 = 0.0f, b2 = 0.0f, b3 = 0.0f; - const int quads = groups_per_row >> 2; - const int tail = groups_per_row & 3; - -#if defined(HIPFIRE_GFX1151_LM_HEAD_X_BUFFER) - #define LOAD_X8_INTO(b, pfx) \ - const hipfire_lm_f32x4 pfx##v0 = __builtin_bit_cast( \ - hipfire_lm_f32x4, __builtin_amdgcn_raw_buffer_load_b128( \ - x_rsrc, (unsigned int)(b) * 4u, 0, 0)); \ - const hipfire_lm_f32x4 pfx##v1 = __builtin_bit_cast( \ - hipfire_lm_f32x4, __builtin_amdgcn_raw_buffer_load_b128( \ - x_rsrc, ((unsigned int)(b) + 4u) * 4u, 0, 0)); \ - const float pfx##0 = pfx##v0.x, pfx##1 = pfx##v0.y, pfx##2 = pfx##v0.z, pfx##3 = pfx##v0.w; \ - const float pfx##4 = pfx##v1.x, pfx##5 = pfx##v1.y, pfx##6 = pfx##v1.z, pfx##7 = pfx##v1.w -#else - #define LOAD_X8_INTO(b, pfx) \ - const float4 pfx##v0 = *reinterpret_cast(x + (b)); \ - const float4 pfx##v1 = *reinterpret_cast(x + (b) + 4); \ - const float pfx##0 = pfx##v0.x, pfx##1 = pfx##v0.y, pfx##2 = pfx##v0.z, pfx##3 = pfx##v0.w; \ - const float pfx##4 = pfx##v1.x, pfx##5 = pfx##v1.y, pfx##6 = pfx##v1.z, pfx##7 = pfx##v1.w -#endif - -#if defined(HIPFIRE_GFX1151_LM_HEAD_HYBRID_BUFFER) - // Uniform per-group headers are cheaper on the scalar/global path. Only - // the lane-divergent packed dword uses temporal buffer VMEM. - #define LOAD_LM_HEADER(gp, weight_offset) (*(const unsigned int*)(gp)) -#else - #define LOAD_LM_HEADER(gp, weight_offset) \ - HIPFIRE_WEIGHT_LOAD_PTR_U32(weight_rsrc, gp, weight_offset) -#endif - #define LOAD_LM_PACKED(gp, weight_offset) \ - HIPFIRE_WEIGHT_LOAD_PTR_U32(weight_rsrc, gp, weight_offset) - - #define DOG_X8(gp, weight_offset, a, pfx) do { \ - const unsigned int hA = LOAD_LM_HEADER(gp, weight_offset); \ - const unsigned int hB = LOAD_LM_HEADER((gp) + 4, (weight_offset) + 4); \ - const unsigned int hs = (tid < 16) ? hA : hB; \ - float sc = __half2float(__ushort_as_half((unsigned short)(hs & 0xFFFFu))); \ - float zp = __half2float(__ushort_as_half((unsigned short)(hs >> 16))); \ - const unsigned char* d = (const unsigned char*)((gp) + 8 + boff); \ - unsigned char b0=d[0], b1=d[1], b2=d[2], b3=d[3], b4=d[4], b5=d[5]; \ - int q0=b0 &63; int q1=(b0>>6)|((b1&0xF)<<2); int q2=(b1>>4)|((b2&3)<<4); int q3=b2>>2; \ - int q4=b3 &63; int q5=(b3>>6)|((b4&0xF)<<2); int q6=(b4>>4)|((b5&3)<<4); int q7=b5>>2; \ - (a) += (sc*(float)q0+zp)*pfx##0 + (sc*(float)q1+zp)*pfx##1 + (sc*(float)q2+zp)*pfx##2 + (sc*(float)q3+zp)*pfx##3 \ - + (sc*(float)q4+zp)*pfx##4 + (sc*(float)q5+zp)*pfx##5 + (sc*(float)q6+zp)*pfx##6 + (sc*(float)q7+zp)*pfx##7; \ - } while (0) - - for (int q = 0; q < quads; q++) { - const int g = q << 2; - const int base0 = g * 256 + tid * 8; - - // Hoist all 32 x values for this group-quad before any DOG. - LOAD_X8_INTO(base0, xa); - LOAD_X8_INTO(base0 + 256, xb); - LOAD_X8_INTO(base0 + 512, xc); - LOAD_X8_INTO(base0 + 768, xd); -#if defined(HIPFIRE_GFX1151_LM_HEAD_X_BUFFER) - asm volatile("" : : - "v"(xav0), "v"(xav1), "v"(xbv0), "v"(xbv1), - "v"(xcv0), "v"(xcv1), "v"(xdv0), "v"(xdv1) : "memory"); -#endif - - // Branchless dual-row DOG with all x resident (no mid-quad x reload). - DOG_X8(row_ptr0 + g * 200, row_weight_offset0 + g * 200, a0, xa); - DOG_X8(row_ptr1 + g * 200, row_weight_offset1 + g * 200, b0, xa); - DOG_X8(row_ptr0 + (g + 1) * 200, row_weight_offset0 + (g + 1) * 200, a1, xb); - DOG_X8(row_ptr1 + (g + 1) * 200, row_weight_offset1 + (g + 1) * 200, b1, xb); - DOG_X8(row_ptr0 + (g + 2) * 200, row_weight_offset0 + (g + 2) * 200, a2, xc); - DOG_X8(row_ptr1 + (g + 2) * 200, row_weight_offset1 + (g + 2) * 200, b2, xc); - DOG_X8(row_ptr0 + (g + 3) * 200, row_weight_offset0 + (g + 3) * 200, a3, xd); - DOG_X8(row_ptr1 + (g + 3) * 200, row_weight_offset1 + (g + 3) * 200, b3, xd); - } - - // Tail uses sequential x (at most 3 groups; hoist buys little). - #define LOAD_X8(b) \ - const float4 xv0 = *reinterpret_cast(x + (b)); \ - const float4 xv1 = *reinterpret_cast(x + (b) + 4); \ - const float x0 = xv0.x, x1 = xv0.y, x2 = xv0.z, x3 = xv0.w; \ - const float x4 = xv1.x, x5 = xv1.y, x6 = xv1.z, x7 = xv1.w - - #define DOG_X8_TAIL(gp, weight_offset, a) do { \ - const unsigned int hA = LOAD_LM_HEADER(gp, weight_offset); \ - const unsigned int hB = LOAD_LM_HEADER((gp) + 4, (weight_offset) + 4); \ - const unsigned int hs = (tid < 16) ? hA : hB; \ - float sc = __half2float(__ushort_as_half((unsigned short)(hs & 0xFFFFu))); \ - float zp = __half2float(__ushort_as_half((unsigned short)(hs >> 16))); \ - const unsigned char* d = (const unsigned char*)((gp) + 8 + boff); \ - unsigned char b0=d[0], b1=d[1], b2=d[2], b3=d[3], b4=d[4], b5=d[5]; \ - int q0=b0 &63; int q1=(b0>>6)|((b1&0xF)<<2); int q2=(b1>>4)|((b2&3)<<4); int q3=b2>>2; \ - int q4=b3 &63; int q5=(b3>>6)|((b4&0xF)<<2); int q6=(b4>>4)|((b5&3)<<4); int q7=b5>>2; \ - (a) += (sc*(float)q0+zp)*x0 + (sc*(float)q1+zp)*x1 + (sc*(float)q2+zp)*x2 + (sc*(float)q3+zp)*x3 \ - + (sc*(float)q4+zp)*x4 + (sc*(float)q5+zp)*x5 + (sc*(float)q6+zp)*x6 + (sc*(float)q7+zp)*x7; \ - } while (0) - - if (tail >= 1) { - const int g = quads << 2; - const int base = g * 256 + tid * 8; - LOAD_X8(base); - DOG_X8_TAIL(row_ptr0 + g * 200, row_weight_offset0 + g * 200, a0); - DOG_X8_TAIL(row_ptr1 + g * 200, row_weight_offset1 + g * 200, b0); - } - if (tail >= 2) { - const int g = (quads << 2) + 1; - const int base = g * 256 + tid * 8; - LOAD_X8(base); - DOG_X8_TAIL(row_ptr0 + g * 200, row_weight_offset0 + g * 200, a1); - DOG_X8_TAIL(row_ptr1 + g * 200, row_weight_offset1 + g * 200, b1); - } - if (tail >= 3) { - const int g = (quads << 2) + 2; - const int base = g * 256 + tid * 8; - LOAD_X8(base); - DOG_X8_TAIL(row_ptr0 + g * 200, row_weight_offset0 + g * 200, a2); - DOG_X8_TAIL(row_ptr1 + g * 200, row_weight_offset1 + g * 200, b2); - } - #undef DOG_X8_TAIL - #undef LOAD_X8 - #undef DOG_X8 - #undef LOAD_X8_INTO - #undef LOAD_LM_PACKED - #undef LOAD_LM_HEADER - - float s0 = (a0 + a1) + (a2 + a3); - float s1 = (b0 + b1) + (b2 + b3); - for (int offset = 16; offset > 0; offset >>= 1) { - s0 += __shfl_down(s0, offset); - s1 += __shfl_down(s1, offset); - } - - if (tid == 0) { - y[row0] = s0; - if (have_row1) y[row1] = s1; - } + gemv_mq6g256v2_multirow_body<2>(A, x, y, M, K); } __launch_bounds__(32, 18) diff --git a/kernels/src/router_logits_round_f16_rne.hip b/kernels/src/router_logits_round_f16_rne.hip new file mode 100644 index 0000000000..2e99ef2e26 --- /dev/null +++ b/kernels/src/router_logits_round_f16_rne.hip @@ -0,0 +1,59 @@ +// SPDX-License-Identifier: Apache-2.0 +// Copyright (c) 2026 Kevin Read +// hipfire — see LICENSE and NOTICE in the project root. + +#include +#include + +// Round-trip an F32 router-logits array through f16 in-place +// (F32 -> f16 -> F32, round-to-nearest-even). +// +// EschaLabs' runtime computes MoE router logits as `f16(x @ gate_w.T)` and +// only THEN widens back to F32 to run top-k. hipfire's router keeps logits +// in F32 end to end. The two selection functions differ whenever two +// experts' F32 logits round to the same f16 value AND straddle the top-k +// boundary (measured ~0.42% of router decisions, concentrated in a handful +// of layers). Escha's recovery fine-tune was trained and validated against +// the f16-rounding runtime, so hipfire's escha decode path must reproduce +// that rounding before top-k to avoid a silent, unexplained expert-choice +// divergence. +// +// Caller applies this ONLY on the escha decode path (gated on the model's +// routed-expert dtype being Escha2T16/Escha3T16 — see +// `MoeDtypes::has_escha_experts` in hipfire-dispatch), applied to the F32 +// `router_logits` buffer BEFORE any top-k selection kernel runs (fused +// exact-wave64 kernel or the softmax_f32 + moe_topk_renorm_k8 fallback +// pair) so both routes see identically-rounded logits. Every other model / +// dtype takes the unmodified F32 path, byte-for-byte unchanged. +// +// Reuses the `f2h_rne` workaround pattern from escha_h128.hip: on this +// ROCm 7.2.2/gfx1151 build, `__float2half` loses the sign of a RUNTIME +// MULTIPLY-PRODUCED exact zero (see that file's comment for the full +// characterisation). Router logits here come from a GEMV reduction, not a +// bare multiply, so they are not known to hit that exact shape — but the +// exact-zero guard below is correct and free for any provenance, so it is +// applied unconditionally rather than relying on that distinction holding. +// `__float2half` itself is round-to-nearest-even for every other input, so +// no other special-casing is needed to satisfy the RNE requirement. +// +// Layout: in-place, single 1-D buffer of `n` F32 elements (one router's +// logits for one decode token). Grid: ceil(n / 256), Block: 256. +__device__ __forceinline__ __half f2h_rne_logit(float v) { + if (v == 0.0f) { + unsigned int bits = __float_as_uint(v); + return __ushort_as_half((unsigned short)(bits >> 16)); + } + return __float2half(v); +} + +extern "C" __global__ void router_logits_round_f16_rne( + float* __restrict__ logits, + int n +) { + int i = blockIdx.x * blockDim.x + threadIdx.x; + if (i >= n) return; + float v = logits[i]; + if (!isfinite(v)) return; + __half h = f2h_rne_logit(v); + logits[i] = __half2float(h); +} diff --git a/registry/models.json b/registry/models.json index d18e5332e9..3c8d9e595b 100644 --- a/registry/models.json +++ b/registry/models.json @@ -1,5 +1,5 @@ { - "_comment": "Curated hipfire model registry overlay. Keep entry shape: { repo, file, size_gb, min_vram_gb, desc }. Empty repo = local-only. Aliases are simple string→string redirects. Run scripts/registry_gen.py after editing.", + "_comment": "Curated hipfire model registry overlay. Keep entry shape: { repo, file, size_gb, min_vram_gb, desc }. Empty repo = local-only. Aliases are simple string\u2192string redirects. Run scripts/registry_gen.py after editing.", "models": { "qwen3.5:0.8b": { "repo": "hipfire-models/qwen3.5-0.8b", @@ -416,7 +416,253 @@ "file": "qwen3.6-35b-a3b.mq2", "size_gb": 11.6, "min_vram_gb": 14, - "desc": "Floor SKU — smallest, coherent but degraded." + "desc": "Floor SKU \u2014 smallest, coherent but degraded." + }, + "qwen3.6:35b-a3b-escha": { + "recommended_settings": { + "temperature": 1.0, + "top_p": 0.95, + "top_k": 20, + "min_p": 0.0, + "presence_penalty": 1.5, + "repeat_penalty": 1.0 + }, + "sampling_profiles": { + "general": { + "temperature": 1.0, + "top_p": 0.95, + "top_k": 20, + "min_p": 0.0, + "presence_penalty": 1.5, + "repeat_penalty": 1.0 + }, + "coding": { + "temperature": 0.6, + "top_p": 0.95, + "top_k": 20, + "min_p": 0.0, + "presence_penalty": 0.0, + "repeat_penalty": 1.0 + }, + "instruct": { + "temperature": 0.7, + "top_p": 0.8, + "top_k": 20, + "min_p": 0.0, + "presence_penalty": 1.5, + "repeat_penalty": 1.0 + } + }, + "repo": "hipfire-models/qwen3.6-35b-a3b-escha", + "file": "qwen3.6-35b-a3b.escha", + "size_gb": 11.84, + "min_vram_gb": 15, + "desc": "EschaLabs Escha-W2 2-bit trellis experts, stored verbatim and decoded inside the GEMV (no decode-at-load). Base model Qwen3.6-35B-A3B. Measured on gfx1151: Dense down-quantised to MQ6. +17% decode for +0.10% PPL \u2014 the best speed/quality trade of the three. PPL 7.6940, KLD 0.0079 vs the q8 arm, 725 tok/s prefill, 55 tok/s decode, 12.45 GB resident. Experts are identical across all three; only the dense tensors differ." + }, + "qwen3.6:35b-a3b-escha-pro": { + "recommended_settings": { + "temperature": 1.0, + "top_p": 0.95, + "top_k": 20, + "min_p": 0.0, + "presence_penalty": 1.5, + "repeat_penalty": 1.0 + }, + "sampling_profiles": { + "general": { + "temperature": 1.0, + "top_p": 0.95, + "top_k": 20, + "min_p": 0.0, + "presence_penalty": 1.5, + "repeat_penalty": 1.0 + }, + "coding": { + "temperature": 0.6, + "top_p": 0.95, + "top_k": 20, + "min_p": 0.0, + "presence_penalty": 0.0, + "repeat_penalty": 1.0 + }, + "instruct": { + "temperature": 0.7, + "top_p": 0.8, + "top_k": 20, + "min_p": 0.0, + "presence_penalty": 1.5, + "repeat_penalty": 1.0 + } + }, + "repo": "hipfire-models/qwen3.6-35b-a3b-escha", + "file": "qwen3.6-35b-a3b.escha-pro", + "size_gb": 12.34, + "min_vram_gb": 16, + "desc": "EschaLabs Escha-W2 2-bit trellis experts, stored verbatim and decoded inside the GEMV (no decode-at-load). Base model Qwen3.6-35B-A3B. Measured on gfx1151: Default and most faithful. Dense weights are a bit-exact repack of Escha's per-row int8 into per-32-block Q8_0. PPL 7.6864, 684 tok/s prefill, 47 tok/s decode, 12.94 GB resident. Experts are identical across all three; only the dense tensors differ." + }, + "qwen3.6:35b-a3b-escha-xt": { + "recommended_settings": { + "temperature": 1.0, + "top_p": 0.95, + "top_k": 20, + "min_p": 0.0, + "presence_penalty": 1.5, + "repeat_penalty": 1.0 + }, + "sampling_profiles": { + "general": { + "temperature": 1.0, + "top_p": 0.95, + "top_k": 20, + "min_p": 0.0, + "presence_penalty": 1.5, + "repeat_penalty": 1.0 + }, + "coding": { + "temperature": 0.6, + "top_p": 0.95, + "top_k": 20, + "min_p": 0.0, + "presence_penalty": 0.0, + "repeat_penalty": 1.0 + }, + "instruct": { + "temperature": 0.7, + "top_p": 0.8, + "top_k": 20, + "min_p": 0.0, + "presence_penalty": 1.5, + "repeat_penalty": 1.0 + } + }, + "repo": "hipfire-models/qwen3.6-35b-a3b-escha", + "file": "qwen3.6-35b-a3b.escha-xt", + "size_gb": 11.39, + "min_vram_gb": 15, + "desc": "EschaLabs Escha-W2 2-bit trellis experts, stored verbatim and decoded inside the GEMV (no decode-at-load). Base model Qwen3.6-35B-A3B. Measured on gfx1151: Dense down-quantised to MQ4. Fastest, and the only one with a visible quality cost: PPL 8.0643 (+4.9%), KLD 0.0590 vs the q8 arm. 886 tok/s prefill, 63 tok/s decode, 12.04 GB resident. Experts are identical across all three; only the dense tensors differ." + }, + "qwen3.8:27b-escha": { + "recommended_settings": { + "temperature": 1.0, + "top_p": 0.95, + "top_k": 20, + "min_p": 0.0, + "presence_penalty": 1.5, + "repeat_penalty": 1.0 + }, + "sampling_profiles": { + "general": { + "temperature": 1.0, + "top_p": 0.95, + "top_k": 20, + "min_p": 0.0, + "presence_penalty": 1.5, + "repeat_penalty": 1.0 + }, + "coding": { + "temperature": 0.6, + "top_p": 0.95, + "top_k": 20, + "min_p": 0.0, + "presence_penalty": 0.0, + "repeat_penalty": 1.0 + }, + "instruct": { + "temperature": 0.7, + "top_p": 0.8, + "top_k": 20, + "min_p": 0.0, + "presence_penalty": 1.5, + "repeat_penalty": 1.0 + } + }, + "repo": "hipfire-models/qwen3.8-27b-escha", + "file": "qwen3.8-27b.escha", + "size_gb": 10.77, + "min_vram_gb": 15, + "desc": "EschaLabs Escha-W2 2-bit trellis, stored verbatim and decoded inside the GEMV (no decode-at-load). Base model Qwen3.8-27B (dense, arch 5). Default: MQ6 dense tensors, PPL 9.6753, 119 tok/s prefill, 12.1 decode. Measured on gfx1151: PPL on a fixed wikitext-2 slice, 1020 scored tokens. Beats qwen3.8-27b.mq3 on quality while being smaller; slower to decode than the plain MQ quants because it also decodes the trellis (~127 GB/s against their ~198)." + }, + "qwen3.8:27b-escha-pro": { + "recommended_settings": { + "temperature": 1.0, + "top_p": 0.95, + "top_k": 20, + "min_p": 0.0, + "presence_penalty": 1.5, + "repeat_penalty": 1.0 + }, + "sampling_profiles": { + "general": { + "temperature": 1.0, + "top_p": 0.95, + "top_k": 20, + "min_p": 0.0, + "presence_penalty": 1.5, + "repeat_penalty": 1.0 + }, + "coding": { + "temperature": 0.6, + "top_p": 0.95, + "top_k": 20, + "min_p": 0.0, + "presence_penalty": 0.0, + "repeat_penalty": 1.0 + }, + "instruct": { + "temperature": 0.7, + "top_p": 0.8, + "top_k": 20, + "min_p": 0.0, + "presence_penalty": 1.5, + "repeat_penalty": 1.0 + } + }, + "repo": "hipfire-models/qwen3.8-27b-escha", + "file": "qwen3.8-27b.escha-pro", + "size_gb": 11.16, + "min_vram_gb": 15, + "desc": "EschaLabs Escha-W2 2-bit trellis, stored verbatim and decoded inside the GEMV (no decode-at-load). Base model Qwen3.8-27B (dense, arch 5). Q8_0 dense tensors, the most faithful build: PPL 9.6486, 113 tok/s prefill, 10.8 decode. Measured on gfx1151: PPL on a fixed wikitext-2 slice, 1020 scored tokens. Beats qwen3.8-27b.mq3 on quality while being smaller; slower to decode than the plain MQ quants because it also decodes the trellis (~127 GB/s against their ~198)." + }, + "qwen3.8:27b-escha-xt": { + "recommended_settings": { + "temperature": 1.0, + "top_p": 0.95, + "top_k": 20, + "min_p": 0.0, + "presence_penalty": 1.5, + "repeat_penalty": 1.0 + }, + "sampling_profiles": { + "general": { + "temperature": 1.0, + "top_p": 0.95, + "top_k": 20, + "min_p": 0.0, + "presence_penalty": 1.5, + "repeat_penalty": 1.0 + }, + "coding": { + "temperature": 0.6, + "top_p": 0.95, + "top_k": 20, + "min_p": 0.0, + "presence_penalty": 0.0, + "repeat_penalty": 1.0 + }, + "instruct": { + "temperature": 0.7, + "top_p": 0.8, + "top_k": 20, + "min_p": 0.0, + "presence_penalty": 1.5, + "repeat_penalty": 1.0 + } + }, + "repo": "hipfire-models/qwen3.8-27b-escha", + "file": "qwen3.8-27b.escha-xt", + "size_gb": 10.45, + "min_vram_gb": 14, + "desc": "EschaLabs Escha-W2 2-bit trellis, stored verbatim and decoded inside the GEMV (no decode-at-load). Base model Qwen3.8-27B (dense, arch 5). MQ4V2 dense tensors, smallest and fastest: PPL 9.7242, 122 tok/s prefill, 12.3 decode. Measured on gfx1151: PPL on a fixed wikitext-2 slice, 1020 scored tokens. Beats qwen3.8-27b.mq3 on quality while being smaller; slower to decode than the plain MQ quants because it also decodes the trellis (~127 GB/s against their ~198)." }, "qwen3.6:35b-a3b-mq3p": { "recommended_settings": { @@ -460,7 +706,7 @@ "file": "qwen3.6-35b-a3b.mq3p", "size_gb": 17.2, "min_vram_gb": 20, - "desc": "MQ3+P graded — MQ4 quality at MQ3 size." + "desc": "MQ3+P graded \u2014 MQ4 quality at MQ3 size." }, "qwen3.6:35b-a3b-mq4p": { "recommended_settings": { @@ -504,7 +750,7 @@ "file": "qwen3.6-35b-a3b.mq4p", "size_gb": 19.8, "min_vram_gb": 22, - "desc": "MQ4+P graded — robust 4-bit default, beats uniform MQ4." + "desc": "MQ4+P graded \u2014 robust 4-bit default, beats uniform MQ4." }, "qwen3.6:35b-a3b-mfp4": { "recommended_settings": { @@ -548,7 +794,7 @@ "file": "qwen3.6-35b-a3b.mfp4", "size_gb": 20.2, "min_vram_gb": 22, - "desc": "MFP4-E8 vector quant — the 4-bit quality leader." + "desc": "MFP4-E8 vector quant \u2014 the 4-bit quality leader." }, "qwen3.6:35b-a3b-mq5": { "recommended_settings": { @@ -592,7 +838,7 @@ "file": "qwen3.6-35b-a3b.mq5", "size_gb": 23.7, "min_vram_gb": 26, - "desc": "Quality SKU — ~80% MQ4->f32." + "desc": "Quality SKU \u2014 ~80% MQ4->f32." }, "qwen3.6:35b-a3b-mq6": { "recommended_settings": { @@ -1074,7 +1320,7 @@ "file": "qwen3.5-9b.mq3", "size_gb": 4.57, "min_vram_gb": 6.1, - "desc": "MQ3 alpha (3.25 bpw, gfx11/gfx12). Smaller than MQ4, comparable decode. Quality eval pending — see issue #113. Sub-9B MQ3 not shipped — see #114.", + "desc": "MQ3 alpha (3.25 bpw, gfx11/gfx12). Smaller than MQ4, comparable decode. Quality eval pending \u2014 see issue #113. Sub-9B MQ3 not shipped \u2014 see #114.", "default_kv_mode": "q8", "quant_recipe": "v3-awq-f1" }, @@ -1083,7 +1329,7 @@ "file": "qwen3.5-27b.mq3", "size_gb": 10.7, "min_vram_gb": 12, - "desc": "MQ3 alpha — fits 128K asym3 ctx on 24 GB (MQ4 OOMs at ~98K). gfx11/gfx12 only." + "desc": "MQ3 alpha \u2014 fits 128K asym3 ctx on 24 GB (MQ4 OOMs at ~98K). gfx11/gfx12 only." }, "qwen3.6:27b-mq3": { "recommended_settings": { @@ -1096,14 +1342,14 @@ "file": "qwen3.6-27b.mq3", "size_gb": 10.7, "min_vram_gb": 12, - "desc": "MQ3 alpha — fits 128K asym3 ctx on 24 GB. Pairs well with mq4 DFlash draft (126 tok/s τ=7.0)." + "desc": "MQ3 alpha \u2014 fits 128K asym3 ctx on 24 GB. Pairs well with mq4 DFlash draft (126 tok/s \u03c4=7.0)." }, "qwen3.5:27b-draft-mq3": { "repo": "hipfire-models/qwen3.5-27b", "file": "qwen35-27b-dflash-mq3.hfq", "size_gb": 0.67, "min_vram_gb": 12, - "desc": "MQ3 DFlash draft for qwen3.5:27b — 671 MB vs 920 MB MQ4 draft. ~30% lower τ on code prompts; trades draft VRAM for ctx fit." + "desc": "MQ3 DFlash draft for qwen3.5:27b \u2014 671 MB vs 920 MB MQ4 draft. ~30% lower \u03c4 on code prompts; trades draft VRAM for ctx fit." }, "qwen3.6:27b-draft-mq3": { "recommended_settings": { @@ -1159,7 +1405,7 @@ "file": "qwen38-27b-dflash-mq4.hfq", "size_gb": 1.21, "min_vram_gb": 16, - "desc": "MQ4 DFlash draft for qwen3.8:27b (arch 20). Recommended controller draft across the MQ3–MQ6 V2 ladder; same-bit mq3/mq5/mq6 alternatives available." + "desc": "MQ4 DFlash draft for qwen3.8:27b (arch 20). Recommended controller draft across the MQ3\u2013MQ6 V2 ladder; same-bit mq3/mq5/mq6 alternatives available." }, "qwen3.8:27b-draft-mq5": { "repo": "hipfire-models/qwen3.8-27b", @@ -1180,14 +1426,14 @@ "file": "qwen35-9b-dflash-mq4.hfq", "size_gb": 0.55, "min_vram_gb": 6, - "desc": "DFlash draft for qwen3.5:9b — pairs with target for 2-3× decode on code/instruct" + "desc": "DFlash draft for qwen3.5:9b \u2014 pairs with target for 2-3\u00d7 decode on code/instruct" }, "qwen3.5:27b-draft": { "repo": "hipfire-models/qwen3.5-27b", "file": "qwen35-27b-dflash-mq4.hfq", "size_gb": 0.92, "min_vram_gb": 16, - "desc": "DFlash draft for qwen3.5:27b — pairs with target for 4× decode on code (212 tok/s peak)" + "desc": "DFlash draft for qwen3.5:27b \u2014 pairs with target for 4\u00d7 decode on code (212 tok/s peak)" }, "qwen3.6:27b-draft": { "recommended_settings": { @@ -1200,7 +1446,7 @@ "file": "qwen36-27b-dflash-mq4.hfq", "size_gb": 0.92, "min_vram_gb": 16, - "desc": "DFlash draft for qwen3.6:27b — pairs with target for ~4× decode on code (refreshed 2026-04-27 from z-lab@0919688)" + "desc": "DFlash draft for qwen3.6:27b \u2014 pairs with target for ~4\u00d7 decode on code (refreshed 2026-04-27 from z-lab@0919688)" }, "carnice:9b": { "repo": "hipfire-models/carnice-9b", @@ -1414,7 +1660,7 @@ "sha256": "47ccfccddbef5b8e14040bae567c87129669fa9451f0c6b15d8f35964fce7503", "min_vram_gb": 24, "default_kv_mode": "q8", - "desc": "Speed SKU of muse-glimmer: MQ4 body AND MQ4 attention with a Q8 lm_head, 2.35 GB smaller and ~14% faster AR decode than the quality trunk (32.3 vs 28.4 tok/s, gfx1201, 2026-08-14 — registry text, not a live baseline). The .mq4r suffix marks it an MQ4R Redline SKU, but it is not lowered to Redline PM4 yet, so automatic Redline admission is withheld.", + "desc": "Speed SKU of muse-glimmer: MQ4 body AND MQ4 attention with a Q8 lm_head, 2.35 GB smaller and ~14% faster AR decode than the quality trunk (32.3 vs 28.4 tok/s, gfx1201, 2026-08-14 \u2014 registry text, not a live baseline). The .mq4r suffix marks it an MQ4R Redline SKU, but it is not lowered to Redline PM4 yet, so automatic Redline admission is withheld.", "recommended_settings": { "temperature": 1.0, "top_p": 0.95, @@ -1443,7 +1689,7 @@ "size_bytes": 1357990400, "sha256": "6fc0988e51689abf24c2e43d5b1f44794be3409484f99a115d36484ea0e23d2d", "min_vram_gb": 26, - "desc": "DFlash draft for muse-glimmer (muse_glimmer_assistant, arch 23) — 5-layer block-diffusion head that reuses the target's embed and lm_head. Pairs with either SKU. Attach with HIPFIRE_DFLASH_DRAFT; there is no filename auto-pairing." + "desc": "DFlash draft for muse-glimmer (muse_glimmer_assistant, arch 23) \u2014 5-layer block-diffusion head that reuses the target's embed and lm_head. Pairs with either SKU. Attach with HIPFIRE_DFLASH_DRAFT; there is no filename auto-pairing." }, "bonsai:27b-bq1": { "repo": "hipfire-models/bonsai-27b", diff --git a/registry/v1.json b/registry/v1.json index 65aa9982a6..c90ef9adcf 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-04T00:37:30Z", "_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": { @@ -474,6 +474,51 @@ "arch_id": 6, "quant": "mq2" }, + "qwen3.6:35b-a3b-escha": { + "recommended_settings": { + "temperature": 1.0, + "top_p": 0.95, + "top_k": 20, + "min_p": 0.0, + "presence_penalty": 1.5, + "repeat_penalty": 1.0 + }, + "sampling_profiles": { + "general": { + "temperature": 1.0, + "top_p": 0.95, + "top_k": 20, + "min_p": 0.0, + "presence_penalty": 1.5, + "repeat_penalty": 1.0 + }, + "coding": { + "temperature": 0.6, + "top_p": 0.95, + "top_k": 20, + "min_p": 0.0, + "presence_penalty": 0.0, + "repeat_penalty": 1.0 + }, + "instruct": { + "temperature": 0.7, + "top_p": 0.8, + "top_k": 20, + "min_p": 0.0, + "presence_penalty": 1.5, + "repeat_penalty": 1.0 + } + }, + "repo": "hipfire-models/qwen3.6-35b-a3b-escha", + "file": "qwen3.6-35b-a3b-escha-q8.escha", + "size_gb": 12.34, + "min_vram_gb": 16, + "desc": "EschaLabs Escha-W2 2-bit trellis experts, stored verbatim and decoded inside the GEMV (no decode-at-load). Base model Qwen3.6-35B-A3B. Measured on gfx1151: Default and most faithful. Dense weights are a bit-exact repack of Escha's per-row int8 into per-32-block Q8_0. PPL 7.6864, 684 tok/s prefill, 47 tok/s decode, 12.94 GB resident. Experts are identical across all three; only the dense tensors differ.", + "sha256": "bd186b37037ee6c1cb58ce6a5c053b785e9ba30b7f1de04fda7ef4f4f06e3010", + "size_bytes": 12344899840, + "arch_id": 6, + "quant": "escha" + }, "qwen3.6:35b-a3b-mq3p": { "recommended_settings": { "temperature": 1.0, @@ -1898,6 +1943,96 @@ "desc": "Official Ornith 1.5 35B-A3B uniform MQ4G256V2 Redline SKU, quantized from ornith-ai/Ornith-1.5-35B-A3B@10fbf86f with --format mq4 --no-q8-router --uniform. Census: 20,871 MQ4G256V2 (qt44), 31 Q8F16, 191 F16, zero qt13 and zero qt15. Parent-card sampling is temperature 0.6, top_p 0.95, top_k 20, no presence penalty. The optional .mtp is the separately published replacement head documented by the model card. Low/medium/xhigh effort steering defaults to xhigh and has no implicit think-token cap.", "arch_id": 6, "quant": "mq4r" + }, + "qwen3.6:35b-a3b-escha-mq6": { + "recommended_settings": { + "temperature": 1.0, + "top_p": 0.95, + "top_k": 20, + "min_p": 0.0, + "presence_penalty": 1.5, + "repeat_penalty": 1.0 + }, + "sampling_profiles": { + "general": { + "temperature": 1.0, + "top_p": 0.95, + "top_k": 20, + "min_p": 0.0, + "presence_penalty": 1.5, + "repeat_penalty": 1.0 + }, + "coding": { + "temperature": 0.6, + "top_p": 0.95, + "top_k": 20, + "min_p": 0.0, + "presence_penalty": 0.0, + "repeat_penalty": 1.0 + }, + "instruct": { + "temperature": 0.7, + "top_p": 0.8, + "top_k": 20, + "min_p": 0.0, + "presence_penalty": 1.5, + "repeat_penalty": 1.0 + } + }, + "repo": "hipfire-models/qwen3.6-35b-a3b-escha", + "file": "qwen3.6-35b-a3b-escha-mq6.escha", + "size_gb": 11.84, + "min_vram_gb": 15, + "desc": "EschaLabs Escha-W2 2-bit trellis experts, stored verbatim and decoded inside the GEMV (no decode-at-load). Base model Qwen3.6-35B-A3B. Measured on gfx1151: Dense down-quantised to MQ6. +17% decode for +0.10% PPL \u2014 the best speed/quality trade of the three. PPL 7.6940, KLD 0.0079 vs the q8 arm, 725 tok/s prefill, 55 tok/s decode, 12.45 GB resident. Experts are identical across all three; only the dense tensors differ.", + "sha256": "1c8e81f1ca3d07922f0a6b79abae0b119a77bb584121efbb0a040275d4c203ce", + "size_bytes": 11837282560, + "arch_id": 6, + "quant": "escha" + }, + "qwen3.6:35b-a3b-escha-mq4": { + "recommended_settings": { + "temperature": 1.0, + "top_p": 0.95, + "top_k": 20, + "min_p": 0.0, + "presence_penalty": 1.5, + "repeat_penalty": 1.0 + }, + "sampling_profiles": { + "general": { + "temperature": 1.0, + "top_p": 0.95, + "top_k": 20, + "min_p": 0.0, + "presence_penalty": 1.5, + "repeat_penalty": 1.0 + }, + "coding": { + "temperature": 0.6, + "top_p": 0.95, + "top_k": 20, + "min_p": 0.0, + "presence_penalty": 0.0, + "repeat_penalty": 1.0 + }, + "instruct": { + "temperature": 0.7, + "top_p": 0.8, + "top_k": 20, + "min_p": 0.0, + "presence_penalty": 1.5, + "repeat_penalty": 1.0 + } + }, + "repo": "hipfire-models/qwen3.6-35b-a3b-escha", + "file": "qwen3.6-35b-a3b-escha-mq4.escha", + "size_gb": 11.39, + "min_vram_gb": 14, + "desc": "EschaLabs Escha-W2 2-bit trellis experts, stored verbatim and decoded inside the GEMV (no decode-at-load). Base model Qwen3.6-35B-A3B. Measured on gfx1151: Dense down-quantised to MQ4. Fastest, and the only one with a visible quality cost: PPL 8.0643 (+4.9%), KLD 0.0590 vs the q8 arm. 886 tok/s prefill, 63 tok/s decode, 12.04 GB resident. Experts are identical across all three; only the dense tensors differ.", + "sha256": "0cad1caaf21783cc6046345739d39b79be6b26bf9fe67562342761dcc1e0576d", + "size_bytes": 11389344000, + "arch_id": 6, + "quant": "escha" } }, "aliases": { diff --git a/scripts/escha-gtt-probe.sh b/scripts/escha-gtt-probe.sh new file mode 100755 index 0000000000..0705fb2087 --- /dev/null +++ b/scripts/escha-gtt-probe.sh @@ -0,0 +1,35 @@ +#!/usr/bin/env bash +# Measure the escha model's resident footprint as an amdgpu GTT delta. +# +# gfx1151 is a unified-memory APU: "VRAM" is system RAM handed to the GPU +# through GTT, so `mem_info_gtt_used` is the model's true resident cost and +# process RSS is not (the weights are device allocations, not process pages). +# Reports the idle baseline, the peak, and the delta, sampling at 200 ms. +# +# Usage: scripts/escha-gtt-probe.sh [args...] +set -u + +CARD="${ESCHA_GTT_CARD:-/sys/class/drm/card1/device/mem_info_gtt_used}" +[ -r "$CARD" ] || { echo "no readable GTT node at $CARD" >&2; exit 2; } + +base=$(cat "$CARD") +echo "[gtt] baseline: $base bytes ($(echo "scale=2; $base/1000000000" | bc) GB)" + +"$@" & +pid=$! + +peak=$base +while kill -0 "$pid" 2>/dev/null; do + now=$(cat "$CARD" 2>/dev/null || echo 0) + [ "$now" -gt "$peak" ] && peak=$now + sleep 0.2 +done +wait "$pid" +status=$? + +after=$(cat "$CARD") +echo "[gtt] peak: $peak bytes ($(echo "scale=2; $peak/1000000000" | bc) GB)" +echo "[gtt] delta: $((peak - base)) bytes ($(echo "scale=2; ($peak-$base)/1000000000" | bc) GB)" +echo "[gtt] after: $after bytes ($(echo "scale=2; $after/1000000000" | bc) GB)" +echo "[gtt] exit: $status" +exit $status diff --git a/scripts/escha-kld.sh b/scripts/escha-kld.sh new file mode 100755 index 0000000000..c3b1ca8285 --- /dev/null +++ b/scripts/escha-kld.sh @@ -0,0 +1,259 @@ +#!/usr/bin/env bash +# G5 — Escha-W2 quality gate: KLD on a FIXED corpus slice, teacher-forced. +# +# COMPARISON 1 (this script): hipfire-escha production vs `escha_ref` semantics. +# +# The reference is escha_ref, NOT any Escha runtime: escha-mlx is Metal, the +# escha wheel is CUDA (sm_80-sm_120) and ZML needs an NVIDIA driver, so none +# of the three execute on gfx1151. `ref.py` declares itself "the semantic +# contract for every Metal kernel in this package" and is gated on the +# goldens, so agreeing with escha_ref IS agreeing with their runtime, and it +# is exact rather than cross-machine. +# +# escha_ref is a BLOCK-level oracle (codec, H128, expert_linear, swiglu) — +# there is no CPU transformer in this repo and writing one for a 40-layer +# hybrid DeltaNet MoE would make the reference itself the least-trusted +# component. So the reference arm is the SAME hipfire forward with the escha +# experts stored weight-exactly, `HIPFIRE_ESCHA_EXPERT_STORE=f16`. That is +# bit-identical to `escha_ref::reconstruct`'s output (the decode already +# produces fp16; G2 gates it bit-exact against escha_ref, G3 gates the H128 +# pair bit-exact), so the ONLY thing that differs between the two arms is the +# Q8_0 re-quantisation of the expert weights — which is precisely what the +# design doc predicts will dominate this number. +# +# The f16 arm costs no more resident memory than production: per-expert +# buffers are rounded to 2 MiB granules and Q8_0's 2.125/1.0625 MiB +# projections already occupy the 4/2 MiB that f16 needs outright. +# +# COMPARISON 2 (bf16 parent Qwen/Qwen3.6-35B-A3B): NOT RUN. The bf16 parent is +# not on this box in any form — /data/hipfire-models has no safetensors copy, +# and the only cached artifact of the parent is +# `unsloth/Qwen3.6-35B-A3B-GGUF: Qwen3.6-35B-A3B-UD-Q4_K_XL.gguf`, a 4-bit +# quant, which is not a reference. Fetching the parent is ~70 GB. Skipped +# deliberately rather than substituted; see the design doc's Phase 1 results. +# +# TEACHER FORCING is structural here, not a flag. `build_kld_ref_native` writes +# the token stream into the HFKLDR file and `eval_hipfire` reads the tokens +# FROM that file rather than from any generation, so both arms are scored on +# one identical committed token stream by construction. Nothing is ever scored +# on a model's own greedy output — on ds4 that scored 8x better on the median +# and was optimistic. +# +# Both arms use --scoring-mode per-token so the candidate walks the same +# `forward_scratch` path the reference builder walks; the prefill-batch body is +# not admissible for escha anyway. +# +# --kv-mode f32 IS LOAD-BEARING. `build_kld_ref_native` builds its reference +# with an unquantised F32 KV cache. eval_hipfire defaults to `asym3`, and +# leaving that default in place folds the KV-quantisation error into the +# number: measured 0.018357 nats with asym3 against 0.002829 on the identical +# reference with f32, i.e. 6.5x, almost all of it KV rather than codec. +# +# Stage 3 is a NEGATIVE CONTROL, not decoration. It scores the f16 arm against +# its own reference and must be exactly 0.000000; anything else means the +# harness is measuring the run-to-run noise floor rather than the codec, and +# the stage-2 number cannot be attributed. +# +# BOTH the control and the headline number are now ASSERTED by this script +# (stage 4). They used to be enforced only by a human reading stdout — the +# control's result directory was written and never read at all — which is not +# enforcement. A gate whose pass condition lives in a person's attention is a +# gate that passes. +set -euo pipefail +cd "$(dirname "$0")/.." + +HFQ=${1:-/data/hipfire-models/escha-35b.hfq} +SLICE=benchmarks/quality-baselines/slice/wikitext2-1024s-2048ctx.txt +OUT=${ESCHA_KLD_OUT:-/tmp/escha-kld} +# n_ctx 384 => scored_per_chunk = 384 - 1 - 192 = 191 positions per chunk, so +# CHUNKS=1 is the design doc's "~192 positions". 6 chunks (1146 positions) is +# the default because one chunk is one sequence and gives no CI at all. +NCTX=${NCTX:-384} +CHUNKS=${CHUNKS:-6} +TOPK=${TOPK:-256} +ARCH=${ARCH:-gfx1151} + +# Upper bound on the headline KLD. Measured: 0.0027576 nats (95% CI +# 0.0019491-0.0038610) with PPL 7.6585. +# +# 0.010 is ~3.6x the measured mean and ~2.6x the upper CI bound. Chosen loose +# on purpose: the quantity being bounded is a Q8_0 re-quantisation error over +# a 6-chunk sample, so it has real sampling spread, and a bound tight enough +# to flake on a reseed would get raised rather than investigated the first +# time it fired. It is still far below anything a genuine defect produces — +# the failure modes this port has actually hit (a dropped H128 transform, a +# stale activation cache, an unrounded router) move the logits by ~1e-1 and +# the KLD by orders of magnitude, not by a factor of three. +KLD_MAX=${KLD_MAX:-0.010} + +# The reference cache key includes everything the reference DEPENDS on, not +# just its shape. It used to be keyed on "${NCTX}x${CHUNKS}" alone and skipped +# whenever a file of that name was non-empty, so pointing the script at a +# different $HFQ — or rebuilding build_kld_ref_native — silently scored the +# new candidate against the OLD model's reference. That is a stale-oracle +# false negative: the number it prints is meaningless and looks fine. +# +# Keyed on: the model's identity and bytes-in-place (realpath, size, mtime), +# the slice contents, the reference builder binary, and the sampling +# parameters. Cheap: no 12 GB hash, and any of these changing changes the key. +ref_key() { + local builder=./target/release/examples/build_kld_ref_native + { + printf '%s\n' "$(readlink -f "$HFQ")" + stat -c '%s %Y' "$HFQ" + sha256sum "$SLICE" | cut -d' ' -f1 + stat -c '%s %Y' "$builder" + printf '%s %s %s\n' "$NCTX" "$CHUNKS" "$TOPK" + } | sha256sum | cut -c1-16 +} +KEY=$(ref_key) +REF="$OUT/escha-35b-f16-exact-${NCTX}x${CHUNKS}-${KEY}.kldref.bin" + +# Each run gets its OWN result directories, named by the same key. +# +# `kld_reduce.py` reduces a whole DIRECTORY, so a single shared `per-seq/` +# means any `.kldseq` left behind by an earlier experiment silently joins this +# run's reduction — and the stage-4 assertions would then be scoring a mixture +# of runs, or asserting against a row this run did not produce. Per-key +# directories make "the reduction describes exactly this run" structural +# rather than something the operator has to remember to clean up. +RUNDIR="$OUT/run-${NCTX}x${CHUNKS}-${KEY}" +PER_SEQ="$RUNDIR/per-seq" +CONTROL="$RUNDIR/control" +mkdir -p "$PER_SEQ" "$CONTROL" + +echo "== 1/4 weight-exact reference (escha_ref semantics, f16 expert store) ==" +echo " cache key: $(basename "$REF")" +if [ ! -s "$REF" ]; then + HIPFIRE_ESCHA_EXPERT_STORE=f16 \ + ./target/release/examples/build_kld_ref_native \ + --model "$HFQ" --slice "$SLICE" --top-k "$TOPK" \ + --n-ctx "$NCTX" --max-chunks "$CHUNKS" --output "$REF" +else + echo " reusing $REF" +fi + +echo "== 2/4 score the production Q8_0 arm on the SAME token stream ==" +./target/release/examples/eval_hipfire \ + --model "$HFQ" --ref "$REF" \ + --scoring-mode per-token --kv-mode f32 \ + --output "$PER_SEQ/escha-35b-q8_0__${ARCH}__per-token.kldseq" + +echo "== 3/4 negative control: the reference arm against its own reference ==" +echo " (must be exactly 0.000000, or stage 2 is unattributable)" +HIPFIRE_ESCHA_EXPERT_STORE=f16 \ +./target/release/examples/eval_hipfire \ + --model "$HFQ" --ref "$REF" \ + --scoring-mode per-token --kv-mode f32 \ + --output "$CONTROL/escha-35b-f16-selfcontrol__${ARCH}__per-token.kldseq" + +echo "== 4/4 reduce and ASSERT ==" +python3 benchmarks/quality-baselines/harness/kld_reduce.py \ + --result-dir "$PER_SEQ" \ + --out-md "$RUNDIR/result-table.md" \ + --out-json "$RUNDIR/result-data.json" +# The control arm gets reduced too. Writing a result directory and never +# reading it is what let "must print exactly 0.000000" be a comment. +python3 benchmarks/quality-baselines/harness/kld_reduce.py \ + --result-dir "$CONTROL" \ + --out-md "$RUNDIR/control-table.md" \ + --out-json "$RUNDIR/control-data.json" +cat "$RUNDIR/result-table.md" +echo +cat "$RUNDIR/control-table.md" +echo + +python3 - "$RUNDIR/result-data.json" "$RUNDIR/control-data.json" "$KLD_MAX" <<'PYEOF' +import json, sys + +result, control, kld_max = sys.argv[1], sys.argv[2], float(sys.argv[3]) +fail = [] + +def one(path, what): + rows = json.load(open(path)) + if len(rows) != 1: + fail.append( + f"{what}: expected exactly one row in {path}, got {len(rows)} " + f"({[r['variant'] for r in rows]}). A stale .kldseq from an earlier " + "run is in the result directory; the reduction is not describing " + "this run." + ) + return None + return rows[0] + +# What "exactly 0.000000" means, precisely. +# +# The contract in the header is about the value eval_hipfire PRINTS, which is +# 6 decimal places. The underlying float is NOT bit-zero: measured +# +# mean_kld = 2.1341004702939135e-10 p99_kld = 6.029504362788427e-09 +# +# and — this is the part that matters — those two figures reproduce BIT-FOR-BIT +# across repeated runs against the same reference. So the residue is not a +# nondeterminism floor. It is the fixed difference between two programs +# computing the same forward: the reference is written by +# `build_kld_ref_native` and scored by `eval_hipfire`, which are separate +# binaries with their own scratch reuse and launch order. A constant ~1e-10 +# between them is f32 last-bit noise, and it is ~1.3e7 times smaller than the +# 2.7576e-3 the production arm reports. +# +# The assertion therefore encodes what actually has to hold for stage 2 to be +# attributable, in two parts: +# (1) the control rounds to 0.000000 at the printed precision — the literal +# documented contract; and +# (2) the control is at least 10 000x below the production number, so no +# part of the headline figure can be the floor. +# A control that had drifted to genuine run-to-run noise would break (2) long +# before it broke (1), which is why (2) is here at all. +CONTROL_ABS_MAX = 5e-7 # rounds to 0.000000 at 6 dp +CONTROL_RATIO = 1e4 # measured margin is ~1.3e7 + +ctl = one(control, "negative control") +if ctl is not None: + if not (abs(ctl["mean_kld"]) < CONTROL_ABS_MAX and abs(ctl["p99_kld"]) < CONTROL_ABS_MAX): + fail.append( + f"negative control does not round to 0.000000: " + f"mean_kld={ctl['mean_kld']!r} p99_kld={ctl['p99_kld']!r}. The " + "reference arm scored against its own reference must agree with " + "it to the printed precision; anything visible at 6 dp means the " + "harness has a floor of its own and the stage-2 KLD cannot be " + "attributed to the Q8_0 expert re-quantisation." + ) + else: + print( + f"negative control: mean_kld={ctl['mean_kld']:.3e} " + f"p99_kld={ctl['p99_kld']:.3e} (prints as 0.000000), OK" + ) + +res = one(result, "production arm") +if res is not None: + print( + f"production arm: mean_kld={res['mean_kld']:.7f} nats " + f"(95% CI {res['mean_kld_ci_lo']:.7f}-{res['mean_kld_ci_hi']:.7f}) " + f"ppl={res['ppl']:.4f}" + ) + if not (res["mean_kld"] <= kld_max): + fail.append( + f"KLD {res['mean_kld']:.7f} nats exceeds the bound {kld_max}. " + "See the KLD_MAX rationale at the top of escha-kld.sh: this bound " + "is ~3.6x the recorded 0.0027576, so exceeding it is not sampling " + "spread." + ) + if ctl is not None and not ( + abs(ctl["mean_kld"]) * CONTROL_RATIO < res["mean_kld"] + ): + fail.append( + f"the negative control ({ctl['mean_kld']:.3e}) is within " + f"{CONTROL_RATIO:.0e}x of the production KLD " + f"({res['mean_kld']:.7f}). Whatever stage 2 is measuring, a " + "material fraction of it is the harness floor rather than the " + "codec, and the number must not be reported as a codec result." + ) + +if fail: + print() + for f in fail: + print("G5 FAIL:", f) + sys.exit(1) +print("G5 PASS") +PYEOF diff --git a/scripts/escha-verify-roundtrip.py b/scripts/escha-verify-roundtrip.py new file mode 100755 index 0000000000..d057942eeb --- /dev/null +++ b/scripts/escha-verify-roundtrip.py @@ -0,0 +1,220 @@ +#!/usr/bin/env python3 +"""G1: every escha_code tensor in the .hfq must be byte-identical to source. + +Verbatim repack is the whole basis for claiming no codec loss, so this is a +memcmp against the tensor at its indexed offset — not a substring search, +which would be quadratic over a 12 GB file. + +Because it is the SOLE evidence for that contract, the count is asserted, not +merely printed. An earlier version compared each discovered code tensor and +printed PASS when it discovered none — which a non-recursive `*.safetensors` +glob against an HF cache root (shards live under `snapshots//`) produced +silently. The gate now fails on: no shards, no index, no code tensors, or a +count that disagrees with `model.safetensors.index.json`. + +HFQ layout (see hipfire-quantize/src/hfq.rs::write_hfq): + header 32B : magic[4] "HFQM", version u32, arch u32, n_tensors u32, + metadata_offset u64, data_offset u64 + metadata : JSON at metadata_offset + index : n_tensors u32, then per tensor + name_len u16, name, quant_type u8, ndim u8, + dims u32*ndim, group_size u32, data_len u64 + data : at data_offset (4096-aligned), tensors concatenated in order +""" +import json, mmap, struct, sys +from pathlib import Path + +ESCHA_QT = {42: "ESCHA2T16", 43: "ESCHA3T16"} + + +def hfq_tensors(mm): + assert mm[:4] == b"HFQM", "not an HFQ file" + version, arch, n_tensors = struct.unpack_from("/`, and a non-recursive search there finds nothing. + """ + root = Path(d) + if not root.is_dir(): + raise GateFailure(f"source {d!r} is not a directory") + idxs = sorted(root.rglob("model.safetensors.index.json")) + if not idxs: + raise GateFailure( + f"no model.safetensors.index.json under {d!r} — cannot establish the " + "expected escha_code count, so a zero-comparison PASS could not be " + "distinguished from a real one" + ) + if len(idxs) > 1: + raise GateFailure( + "more than one model.safetensors.index.json under " + f"{d!r} ({', '.join(str(p) for p in idxs)}) — ambiguous source" + ) + weight_map = json.loads(idxs[0].read_text())["weight_map"] + names = {k for k in weight_map if k.endswith(".escha_code")} + if not names: + raise GateFailure( + f"{idxs[0]} lists no .escha_code tensors — this is not an escha " + "checkpoint, and G1 would otherwise compare nothing and pass" + ) + return names + + +def safetensors_tensors(d): + out = {} + shards = sorted(Path(d).rglob("*.safetensors")) + if not shards: + raise GateFailure( + f"no *.safetensors shards under {d!r} — nothing to compare. " + "(This is the failure mode the count assertions exist for: the " + "loop below would not execute and every counter would stay 0.)" + ) + for shard in shards: + raw = shard.read_bytes() + (n,) = struct.unpack_from(" ") + sys.exit(2) + sys.exit(main(sys.argv[1], sys.argv[2])) \ No newline at end of file diff --git a/scripts/gates.sh b/scripts/gates.sh index 75e4b0fe12..9735c828a0 100755 --- a/scripts/gates.sh +++ b/scripts/gates.sh @@ -21,6 +21,11 @@ WORK_DIR="${HIPFIRE_GATE_WORK_DIR:-$ROOT/.redline-work/gates}" RUN_REDLINE=1 RUN_SERVE=1 RUN_PERF=1 +# Escha-W2 G1-G6. OFF by default: these are checkpoint-specific (they need the +# escha .hfq and, for G1, the source safetensors tree), so running them against +# whatever `--model` happens to be would either fail or, worse, pass vacuously. +RUN_ESCHA=0 +ESCHA_SRC="${ESCHA_SRC:-/data/hipfire-models/escha-35b}" PM4=1 PERF_BASE="HEAD~1" @@ -36,6 +41,21 @@ Options: --perf REF compare performance against REF (default HEAD~1) --aql shadow the retained AQL path instead of one PM4 IB --work-dir PATH artifact directory + --escha ALSO run the Escha-W2 correctness battery G1-G6 + --escha-only run ONLY G1-G6 (--model must be the escha .hfq) + --escha-src PATH source safetensors tree for G1 (or ESCHA_SRC) + default /data/hipfire-models/escha-35b + +Escha-W2 gates (see docs/plans/escha-w2-port-design.md §10.6): + G1 verbatim repack: every escha_code tensor byte-identical to source + G2 GPU tile decode == escha_ref::reconstruct, bit-exact + G3 the H128 pair == escha_ref, bit-exact, every launch form + G4 the whole MoE block against escha's moeblk_out.f16 golden + G4b arch-6 router selects the same experts as escha + G5 KLD on a fixed teacher-forced corpus, with a negative control + G6 batched prefill vs the per-token route, whole model +G5 and G6 load the model (37.6 GB resident) and take minutes; G1-G4b do not +need a GPU-resident model beyond the checkpoint and the committed fixtures. EOF } @@ -47,6 +67,9 @@ while [ $# -gt 0 ]; do --no-perf) RUN_PERF=0 ;; --perf) PERF_BASE="${2:?--perf requires a git ref}"; shift ;; --aql) PM4=0 ;; + --escha) RUN_ESCHA=1 ;; + --escha-only) RUN_ESCHA=1; RUN_REDLINE=0; RUN_SERVE=0; RUN_PERF=0 ;; + --escha-src) ESCHA_SRC="${2:?--escha-src requires a path}"; shift ;; --work-dir) WORK_DIR="${2:?--work-dir requires a path}"; shift ;; -h|--help) usage; exit 0 ;; *) echo "gates.sh: unknown argument: $1" >&2; usage >&2; exit 2 ;; @@ -97,4 +120,40 @@ if [ "$RUN_PERF" -eq 1 ]; then BENCH_MODEL="$MODEL" scripts/probe_commits.sh "$BASE_SHA" "$HEAD_SHA" fi +if [ "$RUN_ESCHA" -eq 1 ]; then + echo "== Escha-W2 correctness battery (G1-G6) ==" + echo " model: $MODEL" + echo " source: $ESCHA_SRC" + echo + # Built once; every gate below is a release example. + cargo build --release --workspace --all-targets --locked + + echo "-- G1: verbatim repack (expect 80/80 byte-identical) --" + python3 scripts/escha-verify-roundtrip.py "$ESCHA_SRC" "$MODEL" + + echo "-- G2: tile decode vs escha_ref (expect 0 mismatched) --" + cargo run --release -p rdna-compute --example test_escha_decode_gpu_vs_cpu + + echo "-- G3: H128 pair vs escha_ref (expect 0 mismatched) --" + cargo run --release -p rdna-compute --example test_escha_h128_gpu_vs_cpu + + echo "-- G4b: router contract (expect 0/8 differing sets) --" + cargo run --release -p hipfire-arch-qwen35 \ + --example escha_router_contract -- "$MODEL" + + echo "-- G4: MoE block vs golden (expect F32 1.828e-4/9.673e-6," \ + "Q8_0 2.633e-4/3.027e-5, 0 differing floats on both routes) --" + cargo run --release -p hipfire-arch-qwen35 \ + --example escha_moe_block_gate -- "$MODEL" + + echo "-- G6: batched prefill vs per-token (expect a stable argmax) --" + cargo run --release -p hipfire-arch-qwen35 \ + --example escha_prefill_batch_gate -- "$MODEL" + + echo "-- G5: KLD (expect 0.0027576 nats, PPL 7.6585, control 0.000000) --" + scripts/escha-kld.sh "$MODEL" + + echo "Escha-W2 G1-G6: all green." +fi + echo "runtime validation artifacts: $WORK_DIR" diff --git a/scripts/registry_gen.py b/scripts/registry_gen.py index eef423544f..a40bdc7a86 100644 --- a/scripts/registry_gen.py +++ b/scripts/registry_gen.py @@ -80,6 +80,12 @@ "hf4", "hf6", "q8", + # EschaLabs Escha-W2 trellis codec (qt=42 ESCHA2T16 / qt=43 ESCHA3T16), + # converted verbatim by `hipfire-quantize --format escha`. It is not an MQ + # format and cannot be transcoded into one without discarding exactly the + # quality that motivates it (docs/plans/escha-w2-port-design.md §2), so it + # gets its own suffix rather than being filed under an mqN. + "escha", "hfq", } # Allowlist for the optional per-entry `default_kv_mode` field (the registry is