From eb2490b4b50a56f671a802ac09366334398270f3 Mon Sep 17 00:00:00 2001 From: Bjoern Agent Date: Fri, 2 Oct 2026 18:53:15 +0200 Subject: [PATCH 01/13] refactor(gemma4): run the eager Gemma 4 layers as declarative Steps EAGLE verify target) is one typed step list `[SandwichAttention, SandwichMlp, PerLayerInput?, Scale?]` built by `hipfire-arch-gemma4/src/program.rs`; the head is `[RmsnormAutomatic, Gemv, Softcap]`. Executors live in the new model-neutral `hipfire-dispatch/src/pipeline/sandwich.rs` with `rows == 1` (decode) and `rows > 1` (batched prefill / EAGLE verify) bodies moved verbatim from forward.rs, so every route issues the same launches. RoPE is an explicit typed `Rope { kind: RotateHalf | PartialHalved { rot_pairs }, theta }`. KV-shared layers and the K=V full layers are op fields, not hand branches. `AttentionFamily::run_attend_only` serves read-only (shared-KV) attention. The fused-route developer knobs (HIPFIRE_GEMMA4_FUSED_*, EAGLE strict) are read once in dispatch with unchanged names and defaults. Deleted from forward.rs: sliding/full layer arms, attn_q8_swa, finish_attn_and_ffn, both PLE branches, batch_attn_block, batch_ffn_block, proj_gemm_batched and the fused-key helpers (-1.6k lines). Dispatch-bypass ledger: gemma4 17 -> 6. Gate (gfx1151, pre-port daemon 8a6c4710 vs this build, greedy, 5 prompts x 128 tokens, HIPFIRE_GEMMA4_LOGIT_TRACE_DIR top-16 logits compared with cmp): - gemma4-12b.mq4 decode (graph): tokens + logit trace byte-identical - gemma4-12b.mq4 HIPFIRE_GEMMA4_EAGLE=1 (unfused route): byte-identical - gemma4-12b.mq4 HIPFIRE_GEMMA4_PREFILL_BATCH=64 (rows>1 prefill): identical - gemma4-e2b / e4b q8 decode and PREFILL_BATCH=64 (PLE, KV sharing): identical - gemma4-12b.mq4 + gemma-4-12B-it-assistant drafter, EAGLE spec 3: identical text, tau per prompt unchanged (3.368/3.459/2.667/3.447/3.514) gfx1100/gfx1201-only batched fusions (Q8 fused prefill, wide PLE GEMM, fused PLE activation) moved verbatim; not exercised on this host. (cherry picked from commit 52308e1cc565a29cd61a629e0b6f7f2f5d9771f4) --- crates/hipfire-arch-gemma4/src/forward.rs | 1793 +---------------- crates/hipfire-arch-gemma4/src/lib.rs | 1 + crates/hipfire-arch-gemma4/src/program.rs | 349 ++++ .../src/families/attention.rs | 25 + crates/hipfire-dispatch/src/pipeline/mod.rs | 1 + .../hipfire-dispatch/src/pipeline/sandwich.rs | 953 +++++++++ crates/hipfire-dispatch/src/pipeline/steps.rs | 24 +- docs/design/gemma4-declarative-program.md | 85 + docs/governance/debt-dispatch-bypass.txt | 2 +- 9 files changed, 1541 insertions(+), 1692 deletions(-) create mode 100644 crates/hipfire-arch-gemma4/src/program.rs create mode 100644 crates/hipfire-dispatch/src/pipeline/sandwich.rs create mode 100644 docs/design/gemma4-declarative-program.md diff --git a/crates/hipfire-arch-gemma4/src/forward.rs b/crates/hipfire-arch-gemma4/src/forward.rs index 6efea6d731..30b82d0d88 100644 --- a/crates/hipfire-arch-gemma4/src/forward.rs +++ b/crates/hipfire-arch-gemma4/src/forward.rs @@ -2,10 +2,10 @@ // Copyright (c) 2026 Kaden Schutt // hipfire — see LICENSE and NOTICE in the project root. -//! Gemma 4 dense forward pass (free functions — hot-path static dispatch). -//! -//! Ported from the old branch's `forward_scratch` / `sliding_layer_decode` / -//! `full_layer_decode`, dropping MoE / E-series / vision. Per-token pipeline: +//! Gemma 4 dense forward driver: embedding, per-layer-input staging, hipGraph +//! capture and the batched verify entry. Every decoder layer runs as the +//! declarative step program in `program.rs` (`rows == 1` decode, `rows > 1` +//! batched prefill / verify); the layer math is, per token: //! //! x = embed(token) * sqrt(dim) //! for each layer (sandwich RMSNorm around BOTH attn and FFN): @@ -54,202 +54,22 @@ //! (`gemm_qkv_q8_0_wmma` is batched-prefill WMMA only), so q+k is the most that //! fuses on the Q8 attention path. -use crate::config::{Gemma4Config, LayerType, RopeType}; -use crate::gemma4::{ - FullLayerWeights, Gemma4State, Gemma4Weights, LayerWeights, PerLayerBranchWeights, - SlidingLayerWeights, GEMMA4_FORWARD_BATCH_MAX, -}; -use hipfire_runtime::llama::{ - rotate_x_mq_batched_for, weight_gemv, weight_gemv_prerotated, KvCache, WeightTensor, -}; +use crate::config::Gemma4Config; +use crate::gemma4::{Gemma4State, Gemma4Weights, LayerWeights, GEMMA4_FORWARD_BATCH_MAX}; +use crate::program::{head, LayerScratch, PleScratch, ProgramBinding}; +use hipfire_dispatch::context::DispatchCtx; +use hipfire_dispatch::pipeline::execute_steps; +use hipfire_runtime::llama::{weight_gemv, WeightTensor}; use rdna_compute::{DType, Gpu, GpuTensor}; /// Greedy EAGLE promises byte-identical target decisions. Its batched verify /// must therefore use the same arithmetic family as eager decode rather than /// numerically-close fused/WMMA variants whose small drift accumulates in KV. fn eagle_strict_enabled() -> bool { - hipfire_config::developer_var("HIPFIRE_GEMMA4_EAGLE").ok().as_deref() == Some("1") -} - -/// Master switch for the qwen35-mirror fused-projection FFN path -/// (`fused_rmsnorm_rotate_mq` + `fused_gate_up_hfq4g256`). Default ON; opt out -/// with `HIPFIRE_GEMMA4_FUSED_FFN=0`. Only fires when the FFN gate/up weights -/// are MQ4G256 (byte-compatible with the HFQ4G256 fused kernel given a -/// pre-FWHT-rotated input). -fn fused_ffn_enabled() -> bool { - !matches!( - hipfire_config::developer_var("HIPFIRE_GEMMA4_FUSED_FFN").ok().as_deref(), - Some("0") | Some("off") | Some("false") - ) -} - -/// Master switch for the fused Q8 q+k projection path -/// (`fused_gate_up_q8_0`, 2 Q8 GEMVs → 1 launch, shared rmsnorm input). -/// Default ON; opt out with `HIPFIRE_GEMMA4_FUSED_QK=0`. gemma4 attention is -/// Q8 (no Q8 fused-QKV decode kernel exists), so we fuse q+k via the 2-way -/// Q8 gate_up fuser and leave v / o separate. Coherence is preserved (the -/// fused kernel is byte-equivalent to two separate Q8 GEMVs). -fn fused_qk_enabled() -> bool { - !matches!( - hipfire_config::developer_var("HIPFIRE_GEMMA4_FUSED_QK").ok().as_deref(), - Some("0") | Some("off") | Some("false") - ) -} - -/// Master switch for the fused sandwich-postnorm + residual-add path -/// (`rmsnorm_residual_add_f32`, gemma4 L4). Fuses post_attention_layernorm and -/// post_feedforward_layernorm each from 3 launches (rmsnorm + memcpy(x<-res) + -/// add_inplace) down to 1, removing ~96 tiny launches/token. Default ON; opt -/// out with `HIPFIRE_GEMMA4_FUSED_POSTNORM=0`. Not byte-identical (the residual -/// add rounds inside the norm kernel) -- coherence-validated. -fn fused_postnorm_enabled() -> bool { - !eagle_strict_enabled() - && !matches!( - hipfire_config::developer_var("HIPFIRE_GEMMA4_FUSED_POSTNORM") - .ok() - .as_deref(), - Some("0") | Some("off") | Some("false") - ) -} - -/// Master switch for the fused per-head weighted q/k RMSNorm + q prescale + -/// dual RoPE path (`fused_gemma4_qk_norm_rope_f32`, gemma4 L3). Collapses -/// (q_norm + k_norm + scale_f32(q) + rope) from 4 launches to 1 on the AR -/// decode path; V's weight-less RMSNorm and the k_eq_v k->v capture stay -/// separate. Default ON; opt out with `HIPFIRE_GEMMA4_FUSED_QK_ROPE=0`. Not -/// byte-identical (fused rsqrt/rope rounds differently) -- coherence-validated. -fn fused_qk_rope_enabled() -> bool { - !eagle_strict_enabled() - && !matches!( - hipfire_config::developer_var("HIPFIRE_GEMMA4_FUSED_QK_ROPE") - .ok() - .as_deref(), - Some("0") | Some("off") | Some("false") - ) -} - -/// q = q_proj(x); k = k_proj(x). Fused into one launch via `fused_gate_up_q8_0` -/// when both are Q8_0 (same input `x`); else two `weight_gemv` calls. -fn qk_proj( - gpu: &mut Gpu, - q_proj: &hipfire_runtime::llama::WeightTensor, - k_proj: &hipfire_runtime::llama::WeightTensor, - x: &rdna_compute::GpuTensor, - q_out: &rdna_compute::GpuTensor, - k_out: &rdna_compute::GpuTensor, -) -> Result<(), String> { - let both_q8 = q_proj.gpu_dtype == DType::Q8_0 && k_proj.gpu_dtype == DType::Q8_0; - if fused_qk_enabled() && both_q8 { - gpu.fused_gate_up_q8_0( - &q_proj.buf, - &k_proj.buf, - x, - q_out, - k_out, - q_proj.m, - k_proj.m, - q_proj.k, - ) - .map_err(|e| format!("gemma4: fused q+k: {e:?}")) - } else { - weight_gemv(gpu, q_proj, x, q_out).map_err(|e| format!("gemma4: q_proj: {e}"))?; - weight_gemv(gpu, k_proj, x, k_out).map_err(|e| format!("gemma4: k_proj: {e}")) - } -} - -/// Master switch for the fused attention-input norm path. Default ON; opt out -/// with `HIPFIRE_GEMMA4_FUSED_ATTN_NORM=0`. Only fires when q/k/v are MQ4G256 -/// (this model's attention projections), folding the pre-attention -/// `input_layernorm` rmsnorm + the shared FWHT rotate into ONE launch -/// (`fused_rmsnorm_rotate_mq` -> tmp_rot), then feeding the prerotated input to -/// each projection's prerotated GEMV (rotate ONCE, reuse across q/k/v). This -/// removes the per-layer input `rmsnorm_f32` launch and collapses the 3 -/// redundant per-projection rotates (q/k/v each re-rotated the identical input) -/// into 1. Byte-equivalent: `fused_rmsnorm_rotate_mq` + `gemv_*_prerotated` is -/// the exact same math as `rmsnorm_f32` + rotation-doing `weight_gemv`, fused. -fn fused_attn_norm_enabled() -> bool { - !matches!( - hipfire_config::developer_var("HIPFIRE_GEMMA4_FUSED_ATTN_NORM") - .ok() - .as_deref(), - Some("0") | Some("off") | Some("false") - ) -} - -/// input_layernorm(x) -> q/k(/v) projections. Fused norm+rotate when MQ4G256 -/// (see `fused_attn_norm_enabled`); else plain rmsnorm -> rotation-doing GEMVs. -/// -/// `v_proj = None` => k_eq_v (the caller copies v from the PRE-k_norm k output -/// AFTER this returns); we only project q+k in that case. When `Some`, v is -/// projected here from the same (pre)normed input. -#[allow(clippy::too_many_arguments)] -fn attn_input_qkv( - gpu: &mut Gpu, - input_ln: &rdna_compute::GpuTensor, - tmp: &rdna_compute::GpuTensor, - tmp_rot: &rdna_compute::GpuTensor, - q_proj: &WeightTensor, - k_proj: &WeightTensor, - v_proj: Option<&WeightTensor>, - x: &rdna_compute::GpuTensor, - q_out: &rdna_compute::GpuTensor, - k_out: &rdna_compute::GpuTensor, - v_out: &rdna_compute::GpuTensor, - dim: usize, - eps: f32, - label: &str, -) -> Result<(), String> { - let all_mq4 = q_proj.gpu_dtype == DType::MQ4G256 - && k_proj.gpu_dtype == DType::MQ4G256 - && v_proj.map_or(true, |w| w.gpu_dtype == DType::MQ4G256); - if fused_attn_norm_enabled() && all_mq4 { - // One launch: rmsnorm(x, input_ln) then FWHT-rotate -> tmp_rot. - gpu.fused_rmsnorm_rotate_mq(x, input_ln, tmp_rot, dim, eps) - .map_err(|e| format!("gemma4 {label}: fused input rmsnorm+rotate: {e:?}"))?; - // Prerotated GEMVs share tmp_rot (NO re-rotation -- byte-identical math). - weight_gemv_prerotated(gpu, q_proj, x, Some(tmp_rot), q_out) - .map_err(|e| format!("gemma4 {label}: q_proj (prerot): {e}"))?; - weight_gemv_prerotated(gpu, k_proj, x, Some(tmp_rot), k_out) - .map_err(|e| format!("gemma4 {label}: k_proj (prerot): {e}"))?; - if let Some(vw) = v_proj { - weight_gemv_prerotated(gpu, vw, x, Some(tmp_rot), v_out) - .map_err(|e| format!("gemma4 {label}: v_proj (prerot): {e}"))?; - } - Ok(()) - } else { - // Eager fallback: plain rmsnorm -> tmp, then rotation-doing GEMVs. - gpu.rmsnorm_f32(x, input_ln, tmp, eps) - .map_err(|e| format!("gemma4 {label}: input rmsnorm: {e:?}"))?; - qk_proj(gpu, q_proj, k_proj, tmp, q_out, k_out)?; - if let Some(vw) = v_proj { - weight_gemv(gpu, vw, tmp, v_out).map_err(|e| format!("gemma4 {label}: v_proj: {e}"))?; - } - Ok(()) - } -} - -fn attn_input_q_only( - gpu: &mut Gpu, - input_ln: &GpuTensor, - tmp: &GpuTensor, - tmp_rot: &GpuTensor, - q_proj: &WeightTensor, - x: &GpuTensor, - q_out: &GpuTensor, - dim: usize, - eps: f32, - label: &str, -) -> Result<(), String> { - if fused_attn_norm_enabled() && q_proj.gpu_dtype == DType::MQ4G256 { - gpu.fused_rmsnorm_rotate_mq(x, input_ln, tmp_rot, dim, eps) - .map_err(|e| format!("gemma4 {label}: fused input rmsnorm+rotate: {e:?}"))?; - weight_gemv_prerotated(gpu, q_proj, x, Some(tmp_rot), q_out) - .map_err(|e| format!("gemma4 {label}: q_proj (prerot): {e}")) - } else { - gpu.rmsnorm_f32(x, input_ln, tmp, eps) - .map_err(|e| format!("gemma4 {label}: input rmsnorm: {e:?}"))?; - weight_gemv(gpu, q_proj, tmp, q_out).map_err(|e| format!("gemma4 {label}: q_proj: {e}")) - } + hipfire_config::developer_var("HIPFIRE_GEMMA4_EAGLE") + .ok() + .as_deref() + == Some("1") } /// Decode one token (eager); returns the full logits vector. Used for prefill, @@ -309,14 +129,16 @@ pub fn decode_step_with_graph( ) -> Result, String> { use std::sync::OnceLock; static GRAPH_ENV: OnceLock> = OnceLock::new(); - let env_override = - *GRAPH_ENV.get_or_init( - || match hipfire_config::developer_var("HIPFIRE_GEMMA4_GRAPH").ok().as_deref() { - Some("1") => Some(true), - Some("0") => Some(false), - _ => None, - }, - ); + let env_override = *GRAPH_ENV.get_or_init(|| { + match hipfire_config::developer_var("HIPFIRE_GEMMA4_GRAPH") + .ok() + .as_deref() + { + Some("1") => Some(true), + Some("0") => Some(false), + _ => None, + } + }); // The captured path is the default. Set HIPFIRE_GEMMA4_GRAPH=0 to retain // the eager fallback for diagnostics. let graph_on = env_override.unwrap_or(true); @@ -616,15 +438,15 @@ fn prepare_per_layer_inputs_batched( && b > 1 && ple.model_projection.gpu_dtype == DType::Q8_0 { - proj_gemm_batched( + hipfire_dispatch::pipeline::sandwich::gemm_rows( gpu, - &ple.model_projection, + &ple.model_projection.dispatch_ref(), x, projection_all, x_rot, b, - "ple model_projection", - )?; + ) + .map_err(|e| format!("gemma4 forward_batch ple model_projection: {e}"))?; } else { for row in 0..b { let x_row = x.sub_offset(row * cfg.dim, cfg.dim); @@ -667,7 +489,6 @@ fn decode_step_body( state.max_seq )); } - let eps = cfg.norm_eps; // Device position scalar (i32). Staged from the heap-stable `state.pos_host` // so the captured memcpy re-reads it on replay (memcpy_htod_auto → async on @@ -680,45 +501,21 @@ fn decode_step_body( .map_err(|e| format!("gemma4: htod pos: {e:?}"))?; } - // Per-layer forward. + // Per-layer program: [SandwichAttention, SandwichMlp, PerLayerInput?, Scale?]. + let ctx = DispatchCtx::new(gpu); + let mut steps = Vec::with_capacity(4); for layer_idx in 0..cfg.n_layers { - let slot = state.kv_slot_for_layer[layer_idx]; - let shared_source_slot = match cfg.kv_shared_source_layer_idx(layer_idx) { - Some(source_layer) => Some(state.kv_slot_for_layer[source_layer]), - None if cfg.is_kv_shared_layer(layer_idx) => { - return Err(format!( - "gemma4 layer {layer_idx}: missing same-type KV sharing source" - )); - } - None => None, + steps.clear(); + let binding = ProgramBinding { + cfg, + state, + rows: 1, + position: position as usize, + positions: None, + scratch: LayerScratch::decode(state), }; - match (cfg.layer_types[layer_idx], &weights.layers[layer_idx]) { - (LayerType::Sliding, LayerWeights::Sliding(lw)) => { - sliding_layer_decode( - gpu, - cfg, - lw, - layer_idx, - position, - slot, - shared_source_slot, - state, - )?; - } - (LayerType::Full, LayerWeights::Full(lw)) => { - full_layer_decode( - gpu, - cfg, - lw, - layer_idx, - position, - slot, - shared_source_slot, - state, - )?; - } - _ => return Err(format!("gemma4 layer {layer_idx} type/weights mismatch")), - } + binding.layer(layer_idx, &weights.layers[layer_idx], &mut steps)?; + execute_steps(gpu, &ctx, &steps).map_err(|e| format!("gemma4 L{layer_idx}: {e}"))?; if let Some(cap) = capture.as_deref_mut() { let h = gpu .download_f32(&state.x) @@ -726,607 +523,22 @@ fn decode_step_body( cap[layer_idx].extend_from_slice(&h); } } + drop(steps); state.n_tokens = position as usize + 1; - // 4) Final RMSNorm → tmp. - gpu.rmsnorm_f32(&state.x, &weights.final_norm, &state.tmp, eps) - .map_err(|e| format!("gemma4: final rmsnorm: {e:?}"))?; - - // 5) LM head → logits (tied embed bytes via lm_head.buf alias). - weight_gemv(gpu, &weights.lm_head, &state.tmp, &state.logits) - .map_err(|e| format!("gemma4: lm_head: {e}"))?; - - // 6) Final logit softcap: logits = tanh(logits / cap) * cap. - if cfg.final_logit_softcapping > 0.0 { - gpu.logit_softcap_f32(&state.logits, cfg.vocab_size, cfg.final_logit_softcapping) - .map_err(|e| format!("gemma4: logit softcap: {e:?}"))?; - } - Ok(()) -} - -/// One sliding-window attention layer (head_dim 256, own v_proj, full RoPE). -fn sliding_layer_decode( - gpu: &mut Gpu, - cfg: &Gemma4Config, - lw: &SlidingLayerWeights, - layer_idx: usize, - pos: u32, - kv_slot: usize, - shared_source_slot: Option, - state: &mut Gemma4State, -) -> Result<(), String> { - let dim = cfg.dim; - let head_dim = cfg.sliding_head_dim; - let n_heads = cfg.n_heads; - let n_kv = cfg.sliding_n_kv_heads; - let eps = cfg.norm_eps; - let dim_bytes = dim * 4; - - // residual = x - gpu.memcpy_dtod_auto(&state.residual.buf, &state.x.buf, dim_bytes) - .map_err(|e| format!("gemma4 sliding: save residual: {e:?}"))?; - - if shared_source_slot.is_some() { - attn_input_q_only( - gpu, - &lw.input_layernorm, - &state.tmp, - &state.tmp_rot, - &lw.q_proj, - &state.x, - &state.q, - dim, - eps, - "sliding shared", - )?; - } else { - attn_input_qkv( - gpu, - &lw.input_layernorm, - &state.tmp, - &state.tmp_rot, - &lw.q_proj, - &lw.k_proj, - Some(&lw.v_proj), - &state.x, - &state.q, - &state.k, - &state.v, - dim, - eps, - "sliding", - )?; - gpu.rmsnorm_batched(&state.v, &state.v_norm_ones, &state.v, n_kv, head_dim, eps) - .map_err(|e| format!("gemma4 sliding: v_norm: {e:?}"))?; - } - - // Per-head q_norm/k_norm (weighted RMS) + q prescale (sqrt(head_dim) so the - // attn kernel's 1/sqrt(head_dim) cancels → effective Gemma4 scale 1.0) + - // full rotate-half RoPE (n_rot_pairs = head_dim/2, theta = 10000). Fused - // (L3) into one launch; eager fallback = 4 separate launches. - if shared_source_slot.is_some() { - gpu.rmsnorm_batched(&state.q, &lw.q_norm, &state.q, n_heads, head_dim, eps) - .map_err(|e| format!("gemma4 sliding shared: q_norm: {e:?}"))?; - gpu.scale_f32(&state.q, (head_dim as f32).sqrt()) - .map_err(|e| format!("gemma4 sliding shared: q scale: {e:?}"))?; - gpu.rope_f32( - &state.q, - &state.k, - &state.pos_buf, - n_heads, - 0, - head_dim, - cfg.sliding_rope_theta, - ) - .map_err(|e| format!("gemma4 sliding shared: rope q: {e:?}"))?; - } else if fused_qk_rope_enabled() { - gpu.fused_gemma4_qk_norm_rope_f32( - &state.q, - &state.k, - &lw.q_norm, - &lw.k_norm, - &state.pos_buf, - n_heads, - n_kv, - head_dim, - head_dim / 2, // full rotate-half - (head_dim as f32).sqrt(), - cfg.sliding_rope_theta, - eps, - ) - .map_err(|e| format!("gemma4 sliding: fused qk norm+rope: {e:?}"))?; - } else { - gpu.rmsnorm_batched(&state.q, &lw.q_norm, &state.q, n_heads, head_dim, eps) - .map_err(|e| format!("gemma4 sliding: q_norm: {e:?}"))?; - gpu.rmsnorm_batched(&state.k, &lw.k_norm, &state.k, n_kv, head_dim, eps) - .map_err(|e| format!("gemma4 sliding: k_norm: {e:?}"))?; - gpu.scale_f32(&state.q, (head_dim as f32).sqrt()) - .map_err(|e| format!("gemma4 sliding: q scale: {e:?}"))?; - gpu.rope_f32( - &state.q, - &state.k, - &state.pos_buf, - n_heads, - n_kv, - head_dim, - cfg.sliding_rope_theta, - ) - .map_err(|e| format!("gemma4 sliding: rope: {e:?}"))?; - } - - // KV write (Q8) + windowed attention (window = sliding_window). - attn_q8_swa( - gpu, - &mut state.kv_sliding, - if shared_source_slot.is_some() { - None - } else { - Some(kv_slot) - }, - shared_source_slot.unwrap_or(kv_slot), - &state.k, - &state.v, - &state.q, - &state.attn_out, - &state.pos_buf, - n_heads, - n_kv, - head_dim, - pos as usize + 1, - &state.q8_flash_partials, - cfg.sliding_window, - )?; - - // o_proj → tmp, post_attention_layernorm(tmp), x = residual + tmp. - finish_attn_and_ffn(gpu, cfg, state, &lw_common_sliding(lw, layer_idx))?; - Ok(()) -} - -/// One full (global) attention layer (head_dim 512, K=V sharing, partial RoPE). -fn full_layer_decode( - gpu: &mut Gpu, - cfg: &Gemma4Config, - lw: &FullLayerWeights, - layer_idx: usize, - pos: u32, - kv_slot: usize, - shared_source_slot: Option, - state: &mut Gemma4State, -) -> Result<(), String> { - let dim = cfg.dim; - let head_dim = cfg.full_head_dim; - let n_heads = cfg.n_heads; - let n_kv = cfg.full_n_kv_heads; - let eps = cfg.norm_eps; - let dim_bytes = dim * 4; - let kv_bytes = n_kv * head_dim * 4; - - // residual = x - gpu.memcpy_dtod_auto(&state.residual.buf, &state.x.buf, dim_bytes) - .map_err(|e| format!("gemma4 full: save residual: {e:?}"))?; - - // n1 = input_layernorm(x) -> q/k(/v). Fused (norm+rotate+prerotated GEMVs, - // shared rotate) when MQ4G256; else plain rmsnorm -> rotation-doing GEMVs. - // When v_proj is Some, v is projected inside the helper from the same - // (pre)normed input; when None (k_eq_v) the helper projects only q+k and we - // capture V below from the PRE-k_norm k output. - if shared_source_slot.is_some() { - attn_input_q_only( - gpu, - &lw.input_layernorm, - &state.tmp, - &state.tmp_rot, - &lw.q_proj, - &state.x, - &state.q, - dim, - eps, - "full shared", - )?; - } else { - attn_input_qkv( - gpu, - &lw.input_layernorm, - &state.tmp, - &state.tmp_rot, - &lw.q_proj, - &lw.k_proj, - lw.v_proj.as_ref(), - &state.x, - &state.q, - &state.k, - &state.v, - dim, - eps, - "full", - )?; - if lw.v_proj.is_none() { - gpu.memcpy_dtod_auto(&state.v.buf, &state.k.buf, kv_bytes) - .map_err(|e| format!("gemma4 full: k->v copy: {e:?}"))?; - } - gpu.rmsnorm_batched(&state.v, &state.v_norm_ones, &state.v, n_kv, head_dim, eps) - .map_err(|e| format!("gemma4 full: v_norm: {e:?}"))?; - } - - // Proportional / partial RoPE: rotate the first `partial_rotary_factor × - // head_dim` dims; theta = full_rope_theta. n_rot_pairs = factor*head_dim/2. - let n_rot_pairs = match cfg.full_rope_type { - RopeType::Proportional => { - ((head_dim as f32) * cfg.full_partial_rotary_factor * 0.5) as usize - } - // Default → all pairs rotate (head_dim/2). rope_partial_halved with - // n_rot_pairs = head_dim/2 == full rotate-half. - RopeType::Default => head_dim / 2, - }; - - // Per-head q_norm/k_norm (weighted RMS) + q prescale (sqrt(head_dim)) + - // partial rotate-half RoPE (n_rot_pairs, theta = full_rope_theta). Fused - // (L3) into one launch; eager fallback = 4 separate launches. - if shared_source_slot.is_some() { - gpu.rmsnorm_batched(&state.q, &lw.q_norm, &state.q, n_heads, head_dim, eps) - .map_err(|e| format!("gemma4 full shared: q_norm: {e:?}"))?; - gpu.scale_f32(&state.q, (head_dim as f32).sqrt()) - .map_err(|e| format!("gemma4 full shared: q scale: {e:?}"))?; - gpu.rope_partial_halved_f32( - &state.q, - &state.k, - &state.pos_buf, - n_heads, - 0, - head_dim, - n_rot_pairs, - cfg.full_rope_theta, - ) - .map_err(|e| format!("gemma4 full shared: rope q: {e:?}"))?; - } else if fused_qk_rope_enabled() { - gpu.fused_gemma4_qk_norm_rope_f32( - &state.q, - &state.k, - &lw.q_norm, - &lw.k_norm, - &state.pos_buf, - n_heads, - n_kv, - head_dim, - n_rot_pairs, - (head_dim as f32).sqrt(), - cfg.full_rope_theta, - eps, - ) - .map_err(|e| format!("gemma4 full: fused qk norm+rope: {e:?}"))?; - } else { - gpu.rmsnorm_batched(&state.q, &lw.q_norm, &state.q, n_heads, head_dim, eps) - .map_err(|e| format!("gemma4 full: q_norm: {e:?}"))?; - gpu.rmsnorm_batched(&state.k, &lw.k_norm, &state.k, n_kv, head_dim, eps) - .map_err(|e| format!("gemma4 full: k_norm: {e:?}"))?; - gpu.scale_f32(&state.q, (head_dim as f32).sqrt()) - .map_err(|e| format!("gemma4 full: q scale: {e:?}"))?; - gpu.rope_partial_halved_f32( - &state.q, - &state.k, - &state.pos_buf, - n_heads, - n_kv, - head_dim, - n_rot_pairs, - cfg.full_rope_theta, - ) - .map_err(|e| format!("gemma4 full: rope: {e:?}"))?; - } - - // KV write (Q8) + full causal attention (window = 0). - attn_q8_swa( - gpu, - &mut state.kv_full, - if shared_source_slot.is_some() { - None - } else { - Some(kv_slot) - }, - shared_source_slot.unwrap_or(kv_slot), - &state.k, - &state.v, - &state.q, - &state.attn_out, - &state.pos_buf, - n_heads, - n_kv, - head_dim, - pos as usize + 1, - &state.q8_flash_partials, - 0, - )?; - - finish_attn_and_ffn(gpu, cfg, state, &lw_common_full(lw, layer_idx))?; - Ok(()) -} - -/// KV write (Q8) + windowed/full tile-reduce attention. -/// `window = 0` ⇒ full causal; `window > 0` ⇒ sliding window. -#[allow(clippy::too_many_arguments)] -fn attn_q8_swa( - gpu: &mut Gpu, - kv: &mut KvCache, - write_slot: Option, - read_slot: usize, - k: &rdna_compute::GpuTensor, - v: &rdna_compute::GpuTensor, - q: &rdna_compute::GpuTensor, - attn_out: &rdna_compute::GpuTensor, - pos_buf: &hip_bridge::DeviceBuffer, - n_heads: usize, - n_kv: usize, - head_dim: usize, - seq_len: usize, - flash_partials: &rdna_compute::GpuTensor, - window: usize, -) -> Result<(), String> { - if let Some(slot) = write_slot { - gpu.kv_cache_write_q8_0(&kv.k_gpu[slot], k, pos_buf, n_kv, head_dim) - .map_err(|e| format!("gemma4: kv write k: {e:?}"))?; - gpu.kv_cache_write_q8_0(&kv.v_gpu[slot], v, pos_buf, n_kv, head_dim) - .map_err(|e| format!("gemma4: kv write v: {e:?}"))?; - } - // DIAG: HIPFIRE_GEMMA4_BASELINE_ATTN routes through the proven baseline - // attention_q8_0_kv (no window) to isolate the new _swa kernel. - if hipfire_config::developer_var_os("HIPFIRE_GEMMA4_BASELINE_ATTN").is_some() { - return gpu - .attention_q8_0_kv( - q, - &kv.k_gpu[read_slot], - &kv.v_gpu[read_slot], - attn_out, - pos_buf, - seq_len, - n_heads, - n_kv, - head_dim, - kv.physical_cap, - ) - .map_err(|e| format!("gemma4: attention baseline: {e:?}")); - } - gpu.attention_flash_q8_0_windowed( - q, - &kv.k_gpu[read_slot], - &kv.v_gpu[read_slot], - attn_out, - pos_buf, - seq_len, - n_heads, - n_kv, - head_dim, - kv.physical_cap, - flash_partials, - window as i32, - ) - .map_err(|e| format!("gemma4: flash attention: {e:?}")) -} - -/// Common per-layer tail shared by sliding + full layers: o_proj, post-attn -/// norm, attn residual add, pre-FFN norm, SwiGLU(gelu_tanh), post-FFN norm, -/// FFN residual add, learned layer_scalar. Operates on `state` scratch. -struct LayerTail<'a> { - layer_idx: usize, - o_proj: &'a hipfire_runtime::llama::WeightTensor, - post_attention_layernorm: &'a rdna_compute::GpuTensor, - pre_feedforward_layernorm: &'a rdna_compute::GpuTensor, - post_feedforward_layernorm: &'a rdna_compute::GpuTensor, - gate_proj: &'a hipfire_runtime::llama::WeightTensor, - up_proj: &'a hipfire_runtime::llama::WeightTensor, - down_proj: &'a hipfire_runtime::llama::WeightTensor, - ffn_hidden_dim: usize, - per_layer: Option<&'a PerLayerBranchWeights>, - layer_scalar_host: f32, -} - -fn lw_common_sliding<'a>(lw: &'a SlidingLayerWeights, layer_idx: usize) -> LayerTail<'a> { - LayerTail { - layer_idx, - o_proj: &lw.o_proj, - post_attention_layernorm: &lw.post_attention_layernorm, - pre_feedforward_layernorm: &lw.pre_feedforward_layernorm, - post_feedforward_layernorm: &lw.post_feedforward_layernorm, - gate_proj: &lw.gate_proj, - up_proj: &lw.up_proj, - down_proj: &lw.down_proj, - ffn_hidden_dim: lw.ffn_hidden_dim, - per_layer: lw.per_layer.as_ref(), - layer_scalar_host: lw.layer_scalar_host, - } -} - -fn lw_common_full<'a>(lw: &'a FullLayerWeights, layer_idx: usize) -> LayerTail<'a> { - LayerTail { - layer_idx, - o_proj: &lw.o_proj, - post_attention_layernorm: &lw.post_attention_layernorm, - pre_feedforward_layernorm: &lw.pre_feedforward_layernorm, - post_feedforward_layernorm: &lw.post_feedforward_layernorm, - gate_proj: &lw.gate_proj, - up_proj: &lw.up_proj, - down_proj: &lw.down_proj, - ffn_hidden_dim: lw.ffn_hidden_dim, - per_layer: lw.per_layer.as_ref(), - layer_scalar_host: lw.layer_scalar_host, - } -} - -fn finish_attn_and_ffn( - gpu: &mut Gpu, - cfg: &Gemma4Config, - state: &mut Gemma4State, - tail: &LayerTail, -) -> Result<(), String> { - let dim = cfg.dim; - let eps = cfg.norm_eps; - let dim_bytes = dim * 4; - let ffn_hd = tail.ffn_hidden_dim; - - // o_proj(attn_out) → tmp. - weight_gemv(gpu, tail.o_proj, &state.attn_out, &state.tmp) - .map_err(|e| format!("gemma4: o_proj: {e}"))?; - - // Sandwich post-attn norm + residual add: x = residual + post_attn_norm(tmp). - // Fused (L4) into one launch; eager fallback = rmsnorm + memcpy + add. - if fused_postnorm_enabled() { - gpu.rmsnorm_residual_add_f32( - &state.tmp, - tail.post_attention_layernorm, - &state.residual, - &state.x, - eps, - ) - .map_err(|e| format!("gemma4: fused post_attn norm+residual: {e:?}"))?; - } else { - // Sandwich post-attn norm (in-place on tmp). - gpu.rmsnorm_f32(&state.tmp, tail.post_attention_layernorm, &state.tmp, eps) - .map_err(|e| format!("gemma4: post_attn rmsnorm: {e:?}"))?; - - // x = residual + tmp. - gpu.memcpy_dtod_auto(&state.x.buf, &state.residual.buf, dim_bytes) - .map_err(|e| format!("gemma4: reset x: {e:?}"))?; - gpu.add_inplace_f32(&state.x, &state.tmp) - .map_err(|e| format!("gemma4: attn residual add: {e:?}"))?; - } - - // residual = x (FFN residual stream). - gpu.memcpy_dtod_auto(&state.residual.buf, &state.x.buf, dim_bytes) - .map_err(|e| format!("gemma4: save ffn residual: {e:?}"))?; - - // SwiGLU FFN with gelu_pytorch_tanh: - // gate = gate_proj(pre_ffn_norm(x)); up = up_proj(pre_ffn_norm(x)); - // hidden = gelu_tanh(gate) * up; ffn_out = down_proj(hidden). - // - // Fused path (MQ4G256 gate/up): fuse the pre-FFN rmsnorm + FWHT rotation - // into one launch (fused_rmsnorm_rotate_mq → tmp_rot), then gate+up in one - // launch via fused_gate_up_hfq4g256 (MQ4G256 bytes are HFQ4G256-compatible - // given a pre-rotated input; the kernel does NOT re-rotate). Mirrors the old - // branch's HIPFIRE_GEMMA4_FUSED_PROJ path + qwen35. - let fuse_gate_up = fused_ffn_enabled() - && tail.gate_proj.gpu_dtype == DType::MQ4G256 - && tail.up_proj.gpu_dtype == DType::MQ4G256; - if fuse_gate_up { - gpu.fused_rmsnorm_rotate_mq( - &state.x, - tail.pre_feedforward_layernorm, - &state.tmp_rot, - dim, - eps, - ) - .map_err(|e| format!("gemma4: fused pre_ffn rmsnorm+rotate: {e:?}"))?; - gpu.fused_gate_up_hfq4g256( - &tail.gate_proj.buf, - &tail.up_proj.buf, - &state.tmp_rot, - &state.gate_ffn, - &state.up_ffn, - tail.gate_proj.m, - tail.up_proj.m, - tail.gate_proj.k, - ) - .map_err(|e| format!("gemma4: fused gate_up: {e:?}"))?; - } else { - // Eager fallback: plain rmsnorm → two rotation-doing GEMVs. - gpu.rmsnorm_f32(&state.x, tail.pre_feedforward_layernorm, &state.tmp, eps) - .map_err(|e| format!("gemma4: pre_ffn rmsnorm: {e:?}"))?; - weight_gemv(gpu, tail.gate_proj, &state.tmp, &state.gate_ffn) - .map_err(|e| format!("gemma4: gate_proj: {e}"))?; - weight_gemv(gpu, tail.up_proj, &state.tmp, &state.up_ffn) - .map_err(|e| format!("gemma4: up_proj: {e}"))?; - } - gpu.gelu_tanh_f32(&state.gate_ffn, &state.ffn_hidden, ffn_hd) - .map_err(|e| format!("gemma4: gelu_tanh: {e:?}"))?; - gpu.mul_f32(&state.ffn_hidden, &state.up_ffn, &state.ffn_hidden) - .map_err(|e| format!("gemma4: silu mul: {e:?}"))?; - weight_gemv(gpu, tail.down_proj, &state.ffn_hidden, &state.ffn_out) - .map_err(|e| format!("gemma4: down_proj: {e}"))?; - - // Sandwich post-FFN norm + residual add: x = residual + post_ffn_norm(ffn_out). - // Fused (L4) into one launch; eager fallback = rmsnorm + memcpy + add. - if fused_postnorm_enabled() { - gpu.rmsnorm_residual_add_f32( - &state.ffn_out, - tail.post_feedforward_layernorm, - &state.residual, - &state.x, - eps, - ) - .map_err(|e| format!("gemma4: fused post_ffn norm+residual: {e:?}"))?; - } else { - // Sandwich post-FFN norm (ffn_out → tmp). - gpu.rmsnorm_f32( - &state.ffn_out, - tail.post_feedforward_layernorm, - &state.tmp, - eps, - ) - .map_err(|e| format!("gemma4: post_ffn rmsnorm: {e:?}"))?; - - // x = residual + tmp. - gpu.memcpy_dtod_auto(&state.x.buf, &state.residual.buf, dim_bytes) - .map_err(|e| format!("gemma4: reset x (ffn): {e:?}"))?; - gpu.add_inplace_f32(&state.x, &state.tmp) - .map_err(|e| format!("gemma4: ffn residual add: {e:?}"))?; - } - - apply_per_layer_input_branch(gpu, cfg, state, tail)?; - - // Learned per-layer scalar multiplier (no-op = 1.0 when tensor absent). - if tail.layer_scalar_host != 1.0 { - gpu.scale_f32(&state.x, tail.layer_scalar_host) - .map_err(|e| format!("gemma4: layer_scalar: {e:?}"))?; - } - Ok(()) -} - -fn apply_per_layer_input_branch( - gpu: &mut Gpu, - cfg: &Gemma4Config, - state: &mut Gemma4State, - tail: &LayerTail, -) -> Result<(), String> { - let Some(ple) = tail.per_layer else { - return Ok(()); - }; - let ple_dim = cfg.hidden_size_per_layer_input; - if ple_dim == 0 { - return Ok(()); - } - let per_layer_inputs = state - .ple_projection_all - .as_ref() - .ok_or_else(|| "gemma4: missing ple_projection_all scratch".to_string())?; - let ple_gate = state - .ple_gate - .as_ref() - .ok_or_else(|| "gemma4: missing ple_gate scratch".to_string())?; - let ple_hidden = state - .ple_hidden - .as_ref() - .ok_or_else(|| "gemma4: missing ple_hidden scratch".to_string())?; - let ple_out = state - .ple_out - .as_ref() - .ok_or_else(|| "gemma4: missing ple_out scratch".to_string())?; - let layer_input = per_layer_inputs.sub_offset(tail.layer_idx * ple_dim, ple_dim); - let dim_bytes = cfg.dim * 4; - - gpu.memcpy_dtod_auto(&state.residual.buf, &state.x.buf, dim_bytes) - .map_err(|e| format!("gemma4: save ple residual: {e:?}"))?; - weight_gemv(gpu, &ple.input_gate, &state.x, ple_gate) - .map_err(|e| format!("gemma4: ple input_gate: {e}"))?; - gpu.gelu_tanh_f32(ple_gate, ple_hidden, ple_dim) - .map_err(|e| format!("gemma4: ple gelu_tanh: {e:?}"))?; - gpu.mul_f32(ple_hidden, &layer_input, ple_hidden) - .map_err(|e| format!("gemma4: ple mul: {e:?}"))?; - weight_gemv(gpu, &ple.projection, ple_hidden, ple_out) - .map_err(|e| format!("gemma4: ple projection: {e}"))?; - gpu.rmsnorm_f32(ple_out, &ple.post_input_norm, &state.tmp, cfg.norm_eps) - .map_err(|e| format!("gemma4: ple post norm: {e:?}"))?; - gpu.memcpy_dtod_auto(&state.x.buf, &state.residual.buf, dim_bytes) - .map_err(|e| format!("gemma4: reset x (ple): {e:?}"))?; - gpu.add_inplace_f32(&state.x, &state.tmp) - .map_err(|e| format!("gemma4: ple residual add: {e:?}")) + // Final norm → tied LM head → logit softcap. + let lm_head = weights.lm_head.dispatch_ref(); + let mut head_steps = Vec::with_capacity(3); + head( + cfg, + &weights.final_norm, + &lm_head, + &state.x, + &state.tmp, + &state.logits, + &mut head_steps, + ); + execute_steps(gpu, &ctx, &head_steps).map_err(|e| format!("gemma4 head: {e}")) } // ════════════════════════════════════════════════════════════════════════ @@ -1352,169 +564,12 @@ fn apply_per_layer_input_branch( // ADDITIVE: does not touch `decode_step` / `decode_step_with_graph`. The eager // path is unchanged. -/// One projection GEMM, batched, dispatched by weight dtype. Mirrors the -/// dtypes `weight_gemv` handles on the gemma4 weight set: -/// - Q8_0 → `gemm_q8_0_batched_chunked` (no rotation; auto-routes to -/// the WMMA Q8 GEMM on gfx12/RDNA4 — the scalar -/// `gemm_q8_0_batched` was 334 us/call at b=4 proj widths = -/// ~61 ms/round = the gemma4 EAGLE slowdown on gfx1201; -/// elsewhere identical scalar kernel, b ≤ 64 = one sub-batch) -/// - F32 → `gemm_f32_batched`; F16/BF16 projection tensors are widened -/// to F32 by the loader. -/// - MQ4G256 → FWHT-rotate x (batched) then `gemm_hfq4g256_batched_lmhead` -/// (MQ4G256 bytes are HFQ4G256-compatible given a pre-rotated -/// input — same equivalence the eager fused FFN path relies on). -/// -/// `x` is [B, k] row-major, `y` is [B, m] row-major. `x_rot` is a [B, k] -/// scratch used only on the MQ4 path. Both buffers must be ≥ the needed size. -#[allow(clippy::too_many_arguments)] -fn proj_gemm_batched( - gpu: &mut Gpu, - w: &WeightTensor, - x: &GpuTensor, - y: &GpuTensor, - x_rot: &GpuTensor, - b: usize, - label: &str, -) -> Result<(), String> { - match w.gpu_dtype { - DType::F32 => gpu - .gemm_f32_batched(&w.buf, x, y, w.m, w.k, b) - .map_err(|e| format!("gemma4 batch {label} (f32): {e:?}")), - DType::Q8_0 => { - let result = if eagle_strict_enabled() { - gpu.gemm_q8_0_batched(&w.buf, x, y, w.m, w.k, b) - } else { - gpu.gemm_q8_0_batched_chunked(&w.buf, x, y, w.m, w.k, b) - }; - result.map_err(|e| format!("gemma4 batch {label} (q8): {e:?}")) - } - DType::MQ4G256 | DType::HFQ4G256 => { - // FWHT-rotate the shared input once, then run the prerotated GEMM. - // rotate_x_mq_batched_for handles the AWQ-aware branch (no-op AWQ - // on these HFQs → plain rotate_x_mq_batched). - rotate_x_mq_batched_for(gpu, w, x, x_rot, w.k, b) - .map_err(|e| format!("gemma4 batch {label} rotate: {e:?}"))?; - gpu.gemm_hfq4g256_batched_lmhead(&w.buf, x_rot, y, w.m, w.k, b) - .map_err(|e| format!("gemma4 batch {label} (mq4): {e:?}")) - } - DType::MQ6G256 => { - // Same two-step as MQ4G256: rotate the shared input once, then - // run the prerotated scalar batched GEMM. MQ6 = HFQ6 packing - // (200 B/group, 6-bit values) with the weights pre-rotated at - // quant time; the inner loop is rotation-free. - rotate_x_mq_batched_for(gpu, w, x, x_rot, w.k, b) - .map_err(|e| format!("gemma4 batch {label} rotate: {e:?}"))?; - gpu.gemm_mq6g256_batched_lmhead(&w.buf, x_rot, y, w.m, w.k, b) - .map_err(|e| format!("gemma4 batch {label} (mq6): {e:?}")) - } - other => Err(format!( - "gemma4 batch {label}: dtype {other:?} has no batched proj kernel" - )), - } -} - -/// Route a plain GEMM through `GemmFamily::run_key` (explicit catalog key). -#[inline] -#[allow(clippy::too_many_arguments)] -fn run_plain_gemm_key( - gpu: &mut Gpu, - key: hipfire_dispatch::types::KernelKey, - w: &WeightTensor, - x: &GpuTensor, - y: &GpuTensor, - batch_size: usize, -) -> Result<(), String> { - use hipfire_dispatch::context::DispatchCtx; - use hipfire_dispatch::families::gemm::GemmParams; - let ctx = DispatchCtx::new(gpu); - let wr = w.dispatch_ref(); - let params = GemmParams { - w: &wr, - x, - y, - batch_size, - }; - hipfire_runtime::llama::gemm_family() - .run_key(key, &ctx, gpu, ¶ms) - .map_err(|e| format!("{e:?}")) -} - -/// Route a 2-way fused gate+up (or Q/K) GEMM through `FusedQkvFamily`. -#[inline] -#[allow(clippy::too_many_arguments)] -fn run_fused_gate_up_key( - gpu: &mut Gpu, - key: hipfire_dispatch::types::KernelKey, - w_gate: &WeightTensor, - w_up: &WeightTensor, - x: &GpuTensor, - y_gate: &GpuTensor, - y_up: &GpuTensor, - batch_size: usize, -) -> Result<(), String> { - use hipfire_dispatch::context::DispatchCtx; - use hipfire_dispatch::families::fused_qkv::FusedQkvParams; - let ctx = DispatchCtx::new(gpu); - let params = FusedQkvParams { - kind: key, - weights: &[&w_gate.buf, &w_up.buf], - x, - outputs: &[y_gate, y_up], - m: &[w_gate.m, w_up.m], - k: w_gate.k, - rot_scratch: &[], - batch_size: Some(batch_size), - }; - hipfire_runtime::llama::fused_qkv_family() - .run(&ctx, gpu, ¶ms) - .map_err(|e| format!("{e:?}")) -} - -/// Route a 3-way fused QKV GEMM through `FusedQkvFamily`. -#[inline] -#[allow(clippy::too_many_arguments)] -fn run_fused_qkv_key( - gpu: &mut Gpu, - key: hipfire_dispatch::types::KernelKey, - wq: &WeightTensor, - wk: &WeightTensor, - wv: &WeightTensor, - x: &GpuTensor, - y_q: &GpuTensor, - y_k: &GpuTensor, - y_v: &GpuTensor, - batch_size: usize, -) -> Result<(), String> { - use hipfire_dispatch::context::DispatchCtx; - use hipfire_dispatch::families::fused_qkv::FusedQkvParams; - let ctx = DispatchCtx::new(gpu); - let params = FusedQkvParams { - kind: key, - weights: &[&wq.buf, &wk.buf, &wv.buf], - x, - outputs: &[y_q, y_k, y_v], - m: &[wq.m, wk.m, wv.m], - k: wq.k, - rot_scratch: &[], - batch_size: Some(batch_size), - }; - hipfire_runtime::llama::fused_qkv_family() - .run(&ctx, gpu, ¶ms) - .map_err(|e| format!("{e:?}")) -} - -fn supports_batched_projection_dtype(dtype: DType) -> bool { - matches!( - dtype, - DType::F32 | DType::Q8_0 | DType::MQ4G256 | DType::HFQ4G256 | DType::MQ6G256 - ) -} - /// Whether every projection used by batched prefill has a matching kernel. /// Other loadable quantizations must remain on eager prefill. pub fn supports_batched_prefill(weights: &Gemma4Weights) -> bool { - let supports = |weight: &WeightTensor| supports_batched_projection_dtype(weight.gpu_dtype); + let supports = |weight: &WeightTensor| { + hipfire_dispatch::pipeline::sandwich::supports_batched_projection(weight.gpu_dtype) + }; weights.layers.iter().all(|layer| match layer { LayerWeights::Sliding(layer) => { supports(&layer.q_proj) @@ -1668,7 +723,7 @@ pub fn forward_batch_spec( let ple_dim = cfg.hidden_size_per_layer_input; let ple_packed = cfg.n_layers * ple_dim; // seq_len after this batch = absolute positions [start_pos, start_pos+B). - let seq_len = checked_batch_seq_len(start_pos, b, state.max_seq)?; + checked_batch_seq_len(start_pos, b, state.max_seq)?; let mut scratch = BatchScratchLedger::new(gpu); let mut alloc = |_g: &mut Gpu, n: usize, label: &str| scratch.alloc(n, label); @@ -1686,7 +741,6 @@ pub fn forward_batch_spec( let k = alloc(gpu, b * max_kv, "k")?; let v = alloc(gpu, b * max_kv, "v")?; let attn_out = alloc(gpu, b * max_q, "attn_out")?; - let o = alloc(gpu, b * dim, "o")?; let gate_ffn = alloc(gpu, b * ffn_hd, "gate_ffn")?; let up_ffn = alloc(gpu, b * ffn_hd, "up_ffn")?; let ffn_hidden = alloc(gpu, b * ffn_hd, "ffn_hidden")?; @@ -1790,149 +844,50 @@ pub fn forward_batch_spec( )?; } - for layer_idx in 0..cfg.n_layers { - let slot = state.kv_slot_for_layer[layer_idx]; - let shared_source_slot = match cfg.kv_shared_source_layer_idx(layer_idx) { - Some(source_layer) => Some(state.kv_slot_for_layer[source_layer]), - None if cfg.is_kv_shared_layer(layer_idx) => { - return Err(format!( - "gemma4 forward_batch L{layer_idx}: shared KV layer has no source" - )); - } - None => None, - }; - match (cfg.layer_types[layer_idx], &weights.layers[layer_idx]) { - (LayerType::Sliding, LayerWeights::Sliding(lw)) => { - let hd = cfg.sliding_head_dim; - let n_kv = cfg.sliding_n_kv_heads; - batch_attn_block( - gpu, - cfg, - state, - BatchAttn { - b, - hd, - n_kv, - write_slot: shared_source_slot.is_none().then_some(slot), - read_slot: shared_source_slot.unwrap_or(slot), - seq_len, - is_full: false, - rope: RopeKind::Full(cfg.sliding_rope_theta), - q_proj: &lw.q_proj, - k_proj: &lw.k_proj, - v_proj: Some(&lw.v_proj), - o_proj: &lw.o_proj, - q_norm: &lw.q_norm, - k_norm: &lw.k_norm, - input_layernorm: &lw.input_layernorm, - post_attention_layernorm: &lw.post_attention_layernorm, - }, - &x, - &residual, - &nrm, - &x_rot, - &q, - &k, - &v, - &attn_out, - &o, - &pos_array, - )?; - batch_ffn_block( - gpu, - cfg, - state, - &lw_common_sliding(lw, layer_idx), - b, - &x, - &residual, - &nrm, - &gate_ffn, - &up_ffn, - &ffn_hidden, - &ffn_out, - &ffn_rot, - BatchPle::from_parts( - cfg, - ple_projection_all.as_ref(), - ple_gate.as_ref(), - ple_hidden.as_ref(), - ple_out.as_ref(), - ), - )?; - } - (LayerType::Full, LayerWeights::Full(lw)) => { - let hd = cfg.full_head_dim; - let n_kv = cfg.full_n_kv_heads; - let n_rot_pairs = match cfg.full_rope_type { - RopeType::Proportional => { - ((hd as f32) * cfg.full_partial_rotary_factor * 0.5) as usize - } - RopeType::Default => hd / 2, - }; - batch_attn_block( - gpu, - cfg, - state, - BatchAttn { - b, - hd, - n_kv, - write_slot: shared_source_slot.is_none().then_some(slot), - read_slot: shared_source_slot.unwrap_or(slot), - seq_len, - is_full: true, - rope: RopeKind::PartialHalved(cfg.full_rope_theta, n_rot_pairs), - q_proj: &lw.q_proj, - k_proj: &lw.k_proj, - v_proj: lw.v_proj.as_ref(), - o_proj: &lw.o_proj, - q_norm: &lw.q_norm, - k_norm: &lw.k_norm, - input_layernorm: &lw.input_layernorm, - post_attention_layernorm: &lw.post_attention_layernorm, - }, - &x, - &residual, - &nrm, - &x_rot, - &q, - &k, - &v, - &attn_out, - &o, - &pos_array, - )?; - batch_ffn_block( - gpu, - cfg, - state, - &lw_common_full(lw, layer_idx), - b, - &x, - &residual, - &nrm, - &gate_ffn, - &up_ffn, - &ffn_hidden, - &ffn_out, - &ffn_rot, - BatchPle::from_parts( - cfg, - ple_projection_all.as_ref(), - ple_gate.as_ref(), - ple_hidden.as_ref(), - ple_out.as_ref(), - ), - )?; - } - _ => { - return Err(format!( - "gemma4 forward_batch L{layer_idx} type/weights mismatch" - )) - } - } + let ctx = DispatchCtx::new(gpu); + let ple = match ( + ple_projection_all.as_ref(), + ple_gate.as_ref(), + ple_hidden.as_ref(), + ple_out.as_ref(), + ) { + (Some(inputs), Some(gate), Some(act), Some(out)) => Some(PleScratch { + inputs, + gate, + act, + out, + }), + _ => None, + }; + let binding = ProgramBinding { + cfg, + state, + rows: b, + position: start_pos, + positions: Some(&pos_array), + scratch: LayerScratch { + x: &x, + residual: &residual, + normed: &nrm, + attn_rot: &x_rot, + mlp_rot: &ffn_rot, + q: &q, + k: &k, + v: &v, + attn_out: &attn_out, + gate: &gate_ffn, + up: &up_ffn, + act: &ffn_hidden, + mlp_out: &ffn_out, + ple, + }, + }; + let mut steps = Vec::with_capacity(4 * cfg.n_layers); + for (layer_idx, layer) in weights.layers.iter().enumerate() { + binding.layer(layer_idx, layer, &mut steps)?; } + execute_steps(gpu, &ctx, &steps).map_err(|e| format!("gemma4 forward_batch: {e}"))?; + drop(steps); state.n_tokens = start_pos + b; // ── Spec-decode out-params (ADDITIVE; both default-off). ── @@ -2097,525 +1052,14 @@ fn embed_lookup_row( } } -/// RoPE flavour for a batched layer. -enum RopeKind { - /// Full rotate-half over the whole head_dim (sliding layers). theta. - Full(f32), - /// Partial proportional rotate-half (full layers). (theta, n_rot_pairs). - PartialHalved(f32, usize), -} - -/// Bundled batched-attention inputs for one layer (keeps arg count sane). -struct BatchAttn<'a> { - b: usize, - hd: usize, - n_kv: usize, - write_slot: Option, - read_slot: usize, - seq_len: usize, - is_full: bool, - rope: RopeKind, - q_proj: &'a WeightTensor, - k_proj: &'a WeightTensor, - v_proj: Option<&'a WeightTensor>, - o_proj: &'a WeightTensor, - q_norm: &'a GpuTensor, - k_norm: &'a GpuTensor, - input_layernorm: &'a GpuTensor, - post_attention_layernorm: &'a GpuTensor, -} - -#[derive(Clone, Copy)] -struct BatchPle<'a> { - projection_all: &'a GpuTensor, - gate: &'a GpuTensor, - hidden: &'a GpuTensor, - out: &'a GpuTensor, -} - -impl<'a> BatchPle<'a> { - fn from_parts( - cfg: &Gemma4Config, - projection_all: Option<&'a GpuTensor>, - gate: Option<&'a GpuTensor>, - hidden: Option<&'a GpuTensor>, - out: Option<&'a GpuTensor>, - ) -> Option { - if cfg.hidden_size_per_layer_input == 0 { - return None; - } - Some(Self { - projection_all: projection_all.expect("PLE projection_all allocated"), - gate: gate.expect("PLE gate allocated"), - hidden: hidden.expect("PLE hidden allocated"), - out: out.expect("PLE out allocated"), - }) - } -} - -/// Batched attention sub-block: residual save, input rmsnorm, q/k/v proj, -/// q/k/v norms, q-scale, RoPE, KV write, masked attention, o_proj, residual -/// add. Mirrors `sliding_layer_decode` / `full_layer_decode` (attention half) -/// batched across B. On exit `x += o_proj(attn)` per row. -#[allow(clippy::too_many_arguments)] -fn batch_attn_block( - gpu: &mut Gpu, - cfg: &Gemma4Config, - state: &mut Gemma4State, - a: BatchAttn, - x: &GpuTensor, - residual: &GpuTensor, - nrm: &GpuTensor, - x_rot: &GpuTensor, - q: &GpuTensor, - k: &GpuTensor, - v: &GpuTensor, - attn_out: &GpuTensor, - o: &GpuTensor, - pos_array: &GpuTensor, -) -> Result<(), String> { - let b = a.b; - let dim = cfg.dim; - let hd = a.hd; - let n_heads = cfg.n_heads; - let n_kv = a.n_kv; - let eps = cfg.norm_eps; - let kv_dim = n_kv * hd; - - // residual = x - gpu.hip - .memcpy_dtod_at(&residual.buf, 0, &x.buf, 0, b * dim * 4) - .map_err(|e| format!("gemma4 batch attn save residual: {e:?}"))?; - - // n1 = input_layernorm(x) → nrm. (per-row, batch over B) - gpu.rmsnorm_batched(x, a.input_layernorm, nrm, b, dim, eps) - .map_err(|e| format!("gemma4 batch attn input rmsnorm: {e:?}"))?; - - // Shared-KV layers only project Q. K/V are read from the latest preceding - // layer with the same attention type. On exact gfx1100, opt in to the - // existing fused Q8 WMMA projections so the common activation is staged - // once instead of once per matrix. - let fuse_q8 = gpu.arch == "gfx1100" - && gpu.flags.gemma4_q8_fused_prefill - && a.write_slot.is_some() - && a.q_proj.gpu_dtype == DType::Q8_0 - && a.k_proj.gpu_dtype == DType::Q8_0 - && a.q_proj.k == a.k_proj.k - && a.q_proj.k % 32 == 0; - match (fuse_q8, a.v_proj) { - (true, Some(vw)) if vw.gpu_dtype == DType::Q8_0 && vw.k == a.q_proj.k => { - run_fused_qkv_key( - gpu, - hipfire_dispatch::types::KernelKey::FusedQkvQ8_0, - a.q_proj, - a.k_proj, - vw, - nrm, - q, - k, - v, - b, - ) - .map_err(|e| format!("gemma4 batch fused qkv (q8): {e}"))?; - } - (true, None) => { - run_fused_gate_up_key( - gpu, - hipfire_dispatch::types::KernelKey::FusedGateUpQ8_0, - a.q_proj, - a.k_proj, - nrm, - q, - k, - b, - ) - .map_err(|e| format!("gemma4 batch fused qk (q8): {e}"))?; - gpu.hip - .memcpy_dtod_at(&v.buf, 0, &k.buf, 0, b * kv_dim * 4) - .map_err(|e| format!("gemma4 batch attn k→v copy: {e:?}"))?; - } - _ => { - proj_gemm_batched(gpu, a.q_proj, nrm, q, x_rot, b, "q_proj")?; - if a.write_slot.is_some() { - proj_gemm_batched(gpu, a.k_proj, nrm, k, x_rot, b, "k_proj")?; - match a.v_proj { - Some(vw) => { - proj_gemm_batched(gpu, vw, nrm, v, x_rot, b, "v_proj")?; - } - None => { - gpu.hip - .memcpy_dtod_at(&v.buf, 0, &k.buf, 0, b * kv_dim * 4) - .map_err(|e| format!("gemma4 batch attn k→v copy: {e:?}"))?; - } - } - } - } - } - - // Per-head q_norm / k_norm over head_dim (batch over B*heads), weight-less - // V norm (ones). The flat [B, heads, head_dim] layout makes B*heads rows. - gpu.rmsnorm_batched(q, a.q_norm, q, b * n_heads, hd, eps) - .map_err(|e| format!("gemma4 batch attn q_norm: {e:?}"))?; - if a.write_slot.is_some() { - gpu.rmsnorm_batched(k, a.k_norm, k, b * n_kv, hd, eps) - .map_err(|e| format!("gemma4 batch attn k_norm: {e:?}"))?; - gpu.rmsnorm_batched(v, &state.v_norm_ones, v, b * n_kv, hd, eps) - .map_err(|e| format!("gemma4 batch attn v_norm: {e:?}"))?; - } - - // Pre-scale Q by √head_dim (cancels the kernel's 1/√head_dim). - gpu.scale_f32(q, (hd as f32).sqrt()) - .map_err(|e| format!("gemma4 batch attn q scale: {e:?}"))?; - - // RoPE (per-row positions). - let rope_n_kv = if a.write_slot.is_some() { n_kv } else { 0 }; - match a.rope { - RopeKind::Full(theta) => { - gpu.rope_batched_f32(q, k, pos_array, n_heads, rope_n_kv, hd, theta, b) - .map_err(|e| format!("gemma4 batch attn rope (full-rot): {e:?}"))?; - } - RopeKind::PartialHalved(theta, n_rot_pairs) => { - gpu.rope_partial_halved_f32_batched( - q, - k, - pos_array, - n_heads, - rope_n_kv, - hd, - n_rot_pairs, - theta, - b, - ) - .map_err(|e| format!("gemma4 batch attn rope (partial): {e:?}"))?; - } - } - - // KV write (Q8) for all B positions. - let kv = if a.is_full { - &state.kv_full - } else { - &state.kv_sliding - }; - let physical_cap = kv.physical_cap; - if let Some(slot) = a.write_slot { - gpu.kv_cache_write_q8_0_batched(&kv.k_gpu[slot], k, pos_array, n_kv, hd, b) - .map_err(|e| format!("gemma4 batch attn kv write k: {e:?}"))?; - gpu.kv_cache_write_q8_0_batched(&kv.v_gpu[slot], v, pos_array, n_kv, hd, b) - .map_err(|e| format!("gemma4 batch attn kv write v: {e:?}"))?; - } - - // Ordinary prefill is not tree verification: passing a dense causal mask - // as `tree_bias` makes the graph-safe tree kernel reinterpret positions[0] - // as block_start. From the second chunk onward that expands seq_len and - // indexes beyond the bias rows. Use the native causal path instead, and - // the existing windowed tile path for sliding attention. - if !a.is_full { - gpu.attention_flash_q8_0_batched_masked_windowed( - q, - &kv.k_gpu[a.read_slot], - &kv.v_gpu[a.read_slot], - attn_out, - pos_array, - n_heads, - n_kv, - hd, - physical_cap, - a.seq_len, - b, - &state.q8_flash_partials, - None, - 0, - 0, - cfg.sliding_window as i32, - ) - .map_err(|e| format!("gemma4 batch sliding flash attn: {e:?}"))?; - } else if a.seq_len > 8_192 { - gpu.attention_flash_q8_0_batched_masked( - q, - &kv.k_gpu[a.read_slot], - &kv.v_gpu[a.read_slot], - attn_out, - pos_array, - n_heads, - n_kv, - hd, - physical_cap, - a.seq_len, - b, - &state.q8_flash_partials, - None, - 0, - 0, - ) - .map_err(|e| format!("gemma4 batch flash attn masked: {e:?}"))?; - } else { - gpu.attention_q8_0_kv_batched_masked( - q, - &kv.k_gpu[a.read_slot], - &kv.v_gpu[a.read_slot], - attn_out, - pos_array, - n_heads, - n_kv, - hd, - physical_cap, - a.seq_len, - b, - None, - 0, - 0, - ) - .map_err(|e| format!("gemma4 batch attn masked: {e:?}"))?; - } - - // o_proj(attn_out) → o. (Mirrors finish_attn_and_ffn's o_proj.) - proj_gemm_batched(gpu, a.o_proj, attn_out, o, x_rot, b, "o_proj")?; - - // Sandwich post-attention norm (in-place on o), then x = residual + o. - gpu.rmsnorm_batched(o, a.post_attention_layernorm, o, b, dim, eps) - .map_err(|e| format!("gemma4 batch attn post_attn rmsnorm: {e:?}"))?; - gpu.hip - .memcpy_dtod_at(&x.buf, 0, &residual.buf, 0, b * dim * 4) - .map_err(|e| format!("gemma4 batch attn reset x: {e:?}"))?; - gpu.add_inplace_f32(x, o) - .map_err(|e| format!("gemma4 batch attn residual add: {e:?}"))?; - Ok(()) -} - -/// Batched FFN sub-block. On entry `x` holds the post-attention residual -/// stream (written by `batch_attn_block`). Mirrors the FFN half of -/// `finish_attn_and_ffn` batched across B: pre-FFN norm, gate/up proj, -/// gelu_tanh·mul, down proj, post-FFN norm, residual add, layer_scalar. -#[allow(clippy::too_many_arguments)] -fn batch_ffn_block( - gpu: &mut Gpu, - cfg: &Gemma4Config, - _state: &mut Gemma4State, - tail: &LayerTail, - b: usize, - x: &GpuTensor, - residual: &GpuTensor, - nrm: &GpuTensor, - gate_ffn: &GpuTensor, - up_ffn: &GpuTensor, - ffn_hidden: &GpuTensor, - ffn_out: &GpuTensor, - ffn_rot: &GpuTensor, - ple: Option, -) -> Result<(), String> { - let dim = cfg.dim; - let eps = cfg.norm_eps; - let ffn_hd = tail.ffn_hidden_dim; - - // 1) pre-FFN norm over the post-attention residual stream `x`. - gpu.rmsnorm_batched(x, tail.pre_feedforward_layernorm, nrm, b, dim, eps) - .map_err(|e| format!("gemma4 batch ffn pre rmsnorm: {e:?}"))?; - - // residual = x (FFN residual stream). - gpu.hip - .memcpy_dtod_at(&residual.buf, 0, &x.buf, 0, b * dim * 4) - .map_err(|e| format!("gemma4 batch ffn save residual: {e:?}"))?; - - // 2) gate / up projections (shared nrm input). - if gpu.arch == "gfx1100" - && gpu.flags.gemma4_q8_fused_prefill - && tail.gate_proj.gpu_dtype == DType::Q8_0 - && tail.up_proj.gpu_dtype == DType::Q8_0 - && tail.gate_proj.k == tail.up_proj.k - && tail.gate_proj.k % 32 == 0 - { - run_fused_gate_up_key( - gpu, - hipfire_dispatch::types::KernelKey::FusedGateUpQ8_0, - tail.gate_proj, - tail.up_proj, - nrm, - gate_ffn, - up_ffn, - b, - ) - .map_err(|e| format!("gemma4 batch fused gate+up (q8): {e}"))?; - } else { - proj_gemm_batched(gpu, tail.gate_proj, nrm, gate_ffn, ffn_rot, b, "gate_proj")?; - proj_gemm_batched(gpu, tail.up_proj, nrm, up_ffn, ffn_rot, b, "up_proj")?; - } - - // 3) hidden = gelu_tanh(gate) * up (elementwise over B*ffn_hd). - gpu.gelu_tanh_f32(gate_ffn, ffn_hidden, b * ffn_hd) - .map_err(|e| format!("gemma4 batch ffn gelu_tanh: {e:?}"))?; - gpu.mul_f32(ffn_hidden, up_ffn, ffn_hidden) - .map_err(|e| format!("gemma4 batch ffn silu mul: {e:?}"))?; - - // 4) down_proj(hidden) → ffn_out. - proj_gemm_batched( - gpu, - tail.down_proj, - ffn_hidden, - ffn_out, - ffn_rot, - b, - "down_proj", - )?; - - // 5) post-FFN norm (ffn_out → nrm). - gpu.rmsnorm_batched(ffn_out, tail.post_feedforward_layernorm, nrm, b, dim, eps) - .map_err(|e| format!("gemma4 batch ffn post rmsnorm: {e:?}"))?; - - // 6) x = residual + nrm. - gpu.hip - .memcpy_dtod_at(&x.buf, 0, &residual.buf, 0, b * dim * 4) - .map_err(|e| format!("gemma4 batch ffn reset x: {e:?}"))?; - gpu.add_inplace_f32(x, nrm) - .map_err(|e| format!("gemma4 batch ffn residual add: {e:?}"))?; - - apply_per_layer_input_branch_batched(gpu, cfg, tail, b, x, residual, nrm, ple)?; - - // 7) learned per-layer scalar. - if tail.layer_scalar_host != 1.0 { - gpu.scale_f32(x, tail.layer_scalar_host) - .map_err(|e| format!("gemma4 batch ffn layer_scalar: {e:?}"))?; - } - Ok(()) -} - -#[allow(clippy::too_many_arguments)] -fn apply_per_layer_input_branch_batched( - gpu: &mut Gpu, - cfg: &Gemma4Config, - tail: &LayerTail, - b: usize, - x: &GpuTensor, - residual: &GpuTensor, - nrm: &GpuTensor, - ple: Option, -) -> Result<(), String> { - let Some(ple_weights) = tail.per_layer else { - return Ok(()); - }; - let Some(ple) = ple else { - return Err(format!( - "gemma4 forward_batch L{}: missing PLE scratch", - tail.layer_idx - )); - }; - let ple_dim = cfg.hidden_size_per_layer_input; - if ple_dim == 0 { - return Ok(()); - } - let dim = cfg.dim; - let packed_dim = cfg.n_layers * ple_dim; - - gpu.hip - .memcpy_dtod_at(&residual.buf, 0, &x.buf, 0, b * dim * 4) - .map_err(|e| format!("gemma4 batch ple save residual: {e:?}"))?; - let exact_batching = supports_gemma4_batched_prefill_arch(&gpu.arch) - && gpu.flags.gemma4_ple_branch_batched_prefill - && b > 1; - let batch_input_gate = exact_batching - && supports_exact_batched_q8_ple_projection( - ple_weights.input_gate.gpu_dtype, - ple_weights.input_gate.k, - ); - let batch_projection = exact_batching - && supports_exact_batched_q8_ple_projection( - ple_weights.projection.gpu_dtype, - ple_weights.projection.k, - ); - if batch_input_gate { - run_plain_gemm_key( - gpu, - hipfire_dispatch::types::KernelKey::GemmQ8_0BatchedWideExact, - &ple_weights.input_gate, - x, - ple.gate, - b, - ) - .map_err(|e| format!("gemma4 batch ple input_gate (q8 wide-exact): {e}"))?; - } else { - for row in 0..b { - let x_row = x.sub_offset(row * dim, dim); - let gate_row = ple.gate.sub_offset(row * ple_dim, ple_dim); - weight_gemv(gpu, &ple_weights.input_gate, &x_row, &gate_row) - .map_err(|e| format!("gemma4 batch ple input_gate row {row}: {e}"))?; - } - } - if supports_gemma4_batched_prefill_arch(&gpu.arch) - && gpu.flags.gemma4_ple_activation_fused_prefill - && b > 1 - { - gpu.gemma4_ple_gelu_mul_strided_f32( - ple.gate, - ple.projection_all, - ple.hidden, - b, - ple_dim, - packed_dim, - tail.layer_idx, - ) - .map_err(|e| format!("gemma4 batch fused PLE activation: {e:?}"))?; - } else { - gpu.gelu_tanh_f32(ple.gate, ple.hidden, b * ple_dim) - .map_err(|e| format!("gemma4 batch ple gelu_tanh: {e:?}"))?; - for row in 0..b { - let hidden_row = ple.hidden.sub_offset(row * ple_dim, ple_dim); - let layer_input = ple - .projection_all - .sub_offset(row * packed_dim + tail.layer_idx * ple_dim, ple_dim); - gpu.mul_f32(&hidden_row, &layer_input, &hidden_row) - .map_err(|e| format!("gemma4 batch ple mul row {row}: {e:?}"))?; - } - } - if batch_projection { - run_plain_gemm_key( - gpu, - hipfire_dispatch::types::KernelKey::GemmQ8_0BatchedWideExact, - &ple_weights.projection, - ple.hidden, - ple.out, - b, - ) - .map_err(|e| format!("gemma4 batch ple projection (q8 wide-exact): {e}"))?; - } else { - for row in 0..b { - let hidden_row = ple.hidden.sub_offset(row * ple_dim, ple_dim); - let out_row = ple.out.sub_offset(row * dim, dim); - weight_gemv(gpu, &ple_weights.projection, &hidden_row, &out_row) - .map_err(|e| format!("gemma4 batch ple projection row {row}: {e}"))?; - } - } - gpu.rmsnorm_batched( - ple.out, - &ple_weights.post_input_norm, - nrm, - b, - dim, - cfg.norm_eps, - ) - .map_err(|e| format!("gemma4 batch ple post norm: {e:?}"))?; - gpu.hip - .memcpy_dtod_at(&x.buf, 0, &residual.buf, 0, b * dim * 4) - .map_err(|e| format!("gemma4 batch ple reset x: {e:?}"))?; - gpu.add_inplace_f32(x, nrm) - .map_err(|e| format!("gemma4 batch ple residual add: {e:?}")) -} - #[inline] fn supports_gemma4_batched_prefill_arch(arch: &str) -> bool { matches!(arch, "gfx1100" | "gfx1201") } -#[inline] -fn supports_exact_batched_q8_ple_projection(dtype: DType, k: usize) -> bool { - dtype == DType::Q8_0 && k > 0 && k <= 1536 && k % 32 == 0 -} - #[cfg(test)] mod tests { - use super::{ - checked_batch_seq_len, supports_batched_projection_dtype, - supports_exact_batched_q8_ple_projection, supports_gemma4_batched_prefill_arch, - }; - use rdna_compute::DType; + use super::{checked_batch_seq_len, supports_gemma4_batched_prefill_arch}; #[test] fn batch_window_must_fit_allocated_kv_capacity() { @@ -2624,37 +1068,6 @@ mod tests { assert!(checked_batch_seq_len(usize::MAX, 1, usize::MAX).is_err()); } - #[test] - fn batched_prefill_projection_formats_are_explicit() { - for dtype in [ - DType::F32, - DType::Q8_0, - DType::MQ4G256, - DType::HFQ4G256, - DType::MQ6G256, - ] { - assert!(supports_batched_projection_dtype(dtype)); - } - for dtype in [ - DType::HFQ4G128, - DType::HFQ6G256, - DType::HFQ2G256, - DType::HFQ3G256, - DType::MQ3G256, - ] { - assert!(!supports_batched_projection_dtype(dtype)); - } - } - - #[test] - fn exact_ple_batching_matches_the_wide_gemv_dispatch_boundary() { - assert!(supports_exact_batched_q8_ple_projection(DType::Q8_0, 1536)); - assert!(!supports_exact_batched_q8_ple_projection(DType::Q8_0, 1535)); - assert!(!supports_exact_batched_q8_ple_projection(DType::Q8_0, 1537)); - assert!(!supports_exact_batched_q8_ple_projection(DType::Q8_0, 0)); - assert!(!supports_exact_batched_q8_ple_projection(DType::F32, 256)); - } - #[test] fn batched_prefill_arch_scope_is_explicit() { assert!(supports_gemma4_batched_prefill_arch("gfx1100")); diff --git a/crates/hipfire-arch-gemma4/src/lib.rs b/crates/hipfire-arch-gemma4/src/lib.rs index 318c003602..060b414679 100644 --- a/crates/hipfire-arch-gemma4/src/lib.rs +++ b/crates/hipfire-arch-gemma4/src/lib.rs @@ -34,6 +34,7 @@ pub mod drafter; pub mod forward; pub mod gemma4; pub mod lowered; +mod program; pub mod speculative; pub use carrier::{ gemma4_context_admission, gemma4_source_uses_lowered, gemma4_use_lowered, load_gemma4_bundle, diff --git a/crates/hipfire-arch-gemma4/src/program.rs b/crates/hipfire-arch-gemma4/src/program.rs new file mode 100644 index 0000000000..85b5f629eb --- /dev/null +++ b/crates/hipfire-arch-gemma4/src/program.rs @@ -0,0 +1,349 @@ +// SPDX-License-Identifier: MIT OR Apache-2.0 +// Copyright (c) 2026 Björn Bösel +// hipfire — see LICENSE and NOTICE in the project root. + +//! Declarative Gemma 4 decoder-layer program. +//! +//! Each layer is `[SandwichAttention, SandwichMlp, PerLayerInput?, Scale?]`: +//! the per-layer-input branch exists on E-series checkpoints, the scale step +//! when the learned layer scalar is not 1. This module only binds resident +//! weights, KV storage and scratch to the shared +//! `hipfire_dispatch::pipeline::sandwich` operations; route selection and +//! executor bodies live in dispatch. The same program serves single-token +//! decode (`rows == 1`) and batched prefill / verify (`rows > 1`). + +use crate::config::{Gemma4Config, LayerType, RopeType}; +use crate::gemma4::{Gemma4State, LayerWeights, PerLayerBranchWeights}; +use hipfire_dispatch::families::gemv::WeightRef; +use hipfire_dispatch::pipeline::sandwich::{ + Activation, PerLayerInputOp, Rope, RopeKind, SandwichAttentionOp, SandwichKv, SandwichMlpOp, + SandwichStream, ScaleOp, SoftcapOp, +}; +use hipfire_dispatch::pipeline::{GemvInput, Step}; +use hipfire_dispatch::types::RotationPlan; +use hipfire_runtime::llama::{KvCache, KvCacheExt, WeightTensor}; +use rdna_compute::GpuTensor; + +/// The typed RoPE of one layer type. +pub(crate) fn layer_rope(cfg: &Gemma4Config, layer_type: LayerType) -> Rope { + match layer_type { + LayerType::Sliding => Rope { + kind: RopeKind::RotateHalf, + theta: cfg.sliding_rope_theta, + }, + LayerType::Full => { + let head_dim = cfg.full_head_dim; + let rot_pairs = match cfg.full_rope_type { + RopeType::Proportional => { + ((head_dim as f32) * cfg.full_partial_rotary_factor * 0.5) as usize + } + RopeType::Default => head_dim / 2, + }; + Rope { + kind: RopeKind::PartialHalved { rot_pairs }, + theta: cfg.full_rope_theta, + } + } + } +} + +/// Per-layer-input scratch; present on E-series checkpoints. +pub(crate) struct PleScratch<'a> { + /// `[rows, n_layers * width]` per-layer inputs of every row. + pub inputs: &'a GpuTensor, + pub gate: &'a GpuTensor, + pub act: &'a GpuTensor, + pub out: &'a GpuTensor, +} + +/// Activation scratch of one forward, `rows` rows each. +pub(crate) struct LayerScratch<'a> { + pub x: &'a GpuTensor, + pub residual: &'a GpuTensor, + pub normed: &'a GpuTensor, + pub attn_rot: &'a GpuTensor, + pub mlp_rot: &'a GpuTensor, + pub q: &'a GpuTensor, + pub k: &'a GpuTensor, + pub v: &'a GpuTensor, + pub attn_out: &'a GpuTensor, + pub gate: &'a GpuTensor, + pub up: &'a GpuTensor, + pub act: &'a GpuTensor, + pub mlp_out: &'a GpuTensor, + pub ple: Option>, +} + +impl<'a> LayerScratch<'a> { + /// The resident single-row scratch of `state`. + pub fn decode(state: &'a Gemma4State) -> Self { + let ple = match ( + &state.ple_projection_all, + &state.ple_gate, + &state.ple_hidden, + &state.ple_out, + ) { + (Some(inputs), Some(gate), Some(act), Some(out)) => Some(PleScratch { + inputs, + gate, + act, + out, + }), + _ => None, + }; + Self { + x: &state.x, + residual: &state.residual, + normed: &state.tmp, + attn_rot: &state.tmp_rot, + mlp_rot: &state.tmp_rot, + q: &state.q, + k: &state.k, + v: &state.v, + attn_out: &state.attn_out, + gate: &state.gate_ffn, + up: &state.up_ffn, + act: &state.ffn_hidden, + mlp_out: &state.ffn_out, + ple, + } + } +} + +/// Binding of one forward over `rows` consecutive positions from `position`. +pub(crate) struct ProgramBinding<'a> { + pub cfg: &'a Gemma4Config, + pub state: &'a Gemma4State, + pub rows: usize, + pub position: usize, + /// Per-row i32 positions (`rows > 1`). + pub positions: Option<&'a GpuTensor>, + pub scratch: LayerScratch<'a>, +} + +struct LayerRefs<'a> { + input_norm: &'a GpuTensor, + post_attn_norm: &'a GpuTensor, + pre_ffn_norm: &'a GpuTensor, + post_ffn_norm: &'a GpuTensor, + layer_scalar: f32, + q: &'a WeightTensor, + k: &'a WeightTensor, + v: Option<&'a WeightTensor>, + o: &'a WeightTensor, + q_norm: &'a GpuTensor, + k_norm: &'a GpuTensor, + gate: &'a WeightTensor, + up: &'a WeightTensor, + down: &'a WeightTensor, + ffn_hidden_dim: usize, + per_layer: Option<&'a PerLayerBranchWeights>, +} + +fn layer_refs(layer: &LayerWeights) -> LayerRefs<'_> { + match layer { + LayerWeights::Sliding(l) => LayerRefs { + input_norm: &l.input_layernorm, + post_attn_norm: &l.post_attention_layernorm, + pre_ffn_norm: &l.pre_feedforward_layernorm, + post_ffn_norm: &l.post_feedforward_layernorm, + layer_scalar: l.layer_scalar_host, + q: &l.q_proj, + k: &l.k_proj, + v: Some(&l.v_proj), + o: &l.o_proj, + q_norm: &l.q_norm, + k_norm: &l.k_norm, + gate: &l.gate_proj, + up: &l.up_proj, + down: &l.down_proj, + ffn_hidden_dim: l.ffn_hidden_dim, + per_layer: l.per_layer.as_ref(), + }, + LayerWeights::Full(l) => LayerRefs { + input_norm: &l.input_layernorm, + post_attn_norm: &l.post_attention_layernorm, + pre_ffn_norm: &l.pre_feedforward_layernorm, + post_ffn_norm: &l.post_feedforward_layernorm, + layer_scalar: l.layer_scalar_host, + q: &l.q_proj, + k: &l.k_proj, + v: l.v_proj.as_ref(), + o: &l.o_proj, + q_norm: &l.q_norm, + k_norm: &l.k_norm, + gate: &l.gate_proj, + up: &l.up_proj, + down: &l.down_proj, + ffn_hidden_dim: l.ffn_hidden_dim, + per_layer: l.per_layer.as_ref(), + }, + } +} + +impl<'a> ProgramBinding<'a> { + fn stream(&self) -> SandwichStream<'a> { + SandwichStream { + rows: self.rows, + hidden: self.cfg.dim, + eps: self.cfg.norm_eps, + x: self.scratch.x, + residual: self.scratch.residual, + normed: self.scratch.normed, + } + } + + /// Append the steps of `layer_idx` to `steps`. + pub fn layer( + &self, + layer_idx: usize, + layer: &'a LayerWeights, + steps: &mut Vec>, + ) -> Result<(), String> { + let cfg = self.cfg; + let st = self.state; + let sc = &self.scratch; + let layer_type = cfg.layer_types[layer_idx]; + if !matches!( + (layer_type, layer), + (LayerType::Sliding, LayerWeights::Sliding(_)) + | (LayerType::Full, LayerWeights::Full(_)) + ) { + return Err(format!("gemma4 layer {layer_idx} type/weights mismatch")); + } + let shared = match cfg.kv_shared_source_layer_idx(layer_idx) { + Some(source) => Some(st.kv_slot_for_layer[source]), + None if cfg.is_kv_shared_layer(layer_idx) => { + return Err(format!( + "gemma4 layer {layer_idx}: missing same-type KV sharing source" + )); + } + None => None, + }; + let (kv, head_dim, n_kv_heads, window): (&KvCache, _, _, _) = match layer_type { + LayerType::Sliding => ( + &st.kv_sliding, + cfg.sliding_head_dim, + cfg.sliding_n_kv_heads, + cfg.sliding_window, + ), + LayerType::Full => (&st.kv_full, cfg.full_head_dim, cfg.full_n_kv_heads, 0), + }; + let w = layer_refs(layer); + let slot = shared.unwrap_or(st.kv_slot_for_layer[layer_idx]); + let writes = shared.is_none(); + steps.push(Step::SandwichAttention(SandwichAttentionOp { + stream: self.stream(), + position: self.position, + n_heads: cfg.n_heads, + n_kv_heads, + head_dim, + input_norm: w.input_norm, + wq: w.q.dispatch_ref(), + wk: writes.then(|| w.k.dispatch_ref()), + wv: w.v.filter(|_| writes).map(WeightTensor::dispatch_ref), + q_norm: w.q_norm, + k_norm: w.k_norm, + v_norm: Some(&st.v_norm_ones), + q_scale: (head_dim as f32).sqrt(), + rope: layer_rope(cfg, layer_type), + kv: SandwichKv { + tier: kv.tier_inputs(), + k_cache: &kv.k_gpu[slot], + v_cache: &kv.v_gpu[slot], + physical_cap: kv.physical_cap, + givens_cos: kv.givens_cos.as_ref(), + givens_sin: kv.givens_sin.as_ref(), + window, + write: writes, + }, + pos_buf: &st.pos_buf, + positions: self.positions, + wo: w.o.dispatch_ref(), + post_norm: w.post_attn_norm, + x_rot: sc.attn_rot, + q: sc.q, + k: sc.k, + v: sc.v, + attn_out: sc.attn_out, + flash_partials: &st.q8_flash_partials, + })); + steps.push(Step::SandwichMlp(SandwichMlpOp { + stream: self.stream(), + pre_norm: w.pre_ffn_norm, + w_gate: w.gate.dispatch_ref(), + w_up: w.up.dispatch_ref(), + w_down: w.down.dispatch_ref(), + activation: Activation::GeluTanh, + hidden_dim: w.ffn_hidden_dim, + post_norm: w.post_ffn_norm, + x_rot: sc.mlp_rot, + gate: sc.gate, + up: sc.up, + act: sc.act, + out: sc.mlp_out, + })); + if let Some(ple) = w.per_layer.filter(|_| cfg.hidden_size_per_layer_input != 0) { + let s = sc + .ple + .as_ref() + .ok_or_else(|| format!("gemma4 layer {layer_idx}: missing PLE scratch"))?; + steps.push(Step::PerLayerInput(PerLayerInputOp { + stream: self.stream(), + layer: layer_idx, + layer_width: cfg.hidden_size_per_layer_input, + n_layers: cfg.n_layers, + inputs: s.inputs, + w_gate: ple.input_gate.dispatch_ref(), + w_proj: ple.projection.dispatch_ref(), + post_norm: &ple.post_input_norm, + gate: s.gate, + act: s.act, + out: s.out, + })); + } + if w.layer_scalar != 1.0 { + steps.push(Step::Scale(ScaleOp { + x: sc.x, + factor: w.layer_scalar, + })); + } + Ok(()) + } +} + +/// Final norm, tied LM head and logit soft cap of one row `x` into `logits`. +/// `lm_head` is the caller-held `weights.lm_head.dispatch_ref()`. +#[allow(clippy::too_many_arguments)] +pub(crate) fn head<'a>( + cfg: &Gemma4Config, + final_norm: &'a GpuTensor, + lm_head: &'a WeightRef<'a>, + x: &'a GpuTensor, + normed: &'a GpuTensor, + logits: &'a GpuTensor, + steps: &mut Vec>, +) { + steps.push(Step::RmsnormAutomatic { + x, + norm_weight: final_norm, + x_plain: normed, + out: normed, + awq_scale: None, + k: cfg.dim, + eps: cfg.norm_eps, + rotation: RotationPlan::None, + }); + steps.push(Step::Gemv { + w: lm_head, + input: GemvInput::Raw(normed), + out: logits, + }); + if cfg.final_logit_softcapping > 0.0 { + steps.push(Step::Softcap(SoftcapOp { + logits, + n: cfg.vocab_size, + cap: cfg.final_logit_softcapping, + })); + } +} diff --git a/crates/hipfire-dispatch/src/families/attention.rs b/crates/hipfire-dispatch/src/families/attention.rs index 77bfff87cd..55df5581b8 100644 --- a/crates/hipfire-dispatch/src/families/attention.rs +++ b/crates/hipfire-dispatch/src/families/attention.rs @@ -193,6 +193,31 @@ impl AttentionFamily { }) } + /// Attend against an already-populated cache without writing the current + /// row: KV-shared layers and draft heads read another layer's cache. + pub fn run_attend_only( + &self, + ctx: &DispatchCtx, + gpu: &mut Gpu, + plan: &crate::families::kv_tier::KvTierPlan, + io: &AttnParams, + ) -> Result<(), DispatchError> { + let shape = ShapeInfo { + batch_size: plan.batch_size, + head_dim: io.head_dim, + m: if plan.batch_size > 1 { + io.max_ctx_len + } else { + io.pos + 1 + }, + is_tree: io.tree_bias.is_some(), + }; + let attend_var = self.resolve(plan.attend_key, ctx, Some(&shape))?; + dispatch_attend(ctx, gpu, plan.attend_key, attend_var.tile, plan, io).map_err(|error| { + DispatchError::Hip(format!("attention {:?}: {error}", plan.attend_key)) + }) + } + /// Full-attention entry point (no KV cache — vision / DFlash cross-attention). /// Resolves under the given key (AttnFullF16 / AttnFullF32 / causal variants) /// and dispatches on the resolved variant's `tile`. The caller is responsible diff --git a/crates/hipfire-dispatch/src/pipeline/mod.rs b/crates/hipfire-dispatch/src/pipeline/mod.rs index 4c4d5f944a..b5c0a507bd 100644 --- a/crates/hipfire-dispatch/src/pipeline/mod.rs +++ b/crates/hipfire-dispatch/src/pipeline/mod.rs @@ -30,6 +30,7 @@ pub mod moe_dump; pub(crate) mod draft_head; pub use draft_head::{DraftHead, DraftHeadLayout, DraftHeadPolicy, DraftHeadRequestState}; pub(crate) mod layer_ops; +pub mod sandwich; pub(crate) mod steps; pub use layer_ops::{ execute_argmax, execute_broadcast_add, execute_clear, execute_embedding, execute_final_hyper, diff --git a/crates/hipfire-dispatch/src/pipeline/sandwich.rs b/crates/hipfire-dispatch/src/pipeline/sandwich.rs new file mode 100644 index 0000000000..c1a77dab37 --- /dev/null +++ b/crates/hipfire-dispatch/src/pipeline/sandwich.rs @@ -0,0 +1,953 @@ +// SPDX-License-Identifier: MIT OR Apache-2.0 +// Copyright (c) 2026 Björn Bösel +// hipfire — see LICENSE and NOTICE in the project root. + +//! Sandwich-norm decoder layer operations. +//! +//! A sandwich-norm decoder normalizes each sublayer's input AND its output +//! before the residual add (`x = residual + post_norm(f(pre_norm(x)))`). +//! Architecture crates declare one attention and one MLP operation per layer +//! (plus an optional per-layer-input branch and layer scale) and bind their +//! resident weights, KV storage and fixed scratch; every dtype-, tier-, arch- +//! and fusion-gated route choice lives here. +//! +//! Executors take `rows`: `rows == 1` is single-token decode, `rows > 1` a +//! batch of consecutive positions (prefill or speculative verify) laid out +//! row-major. The bodies are the former Gemma 4 eager decode and batched +//! verify arms moved verbatim, so every route issues the same launches. + +use crate::context::DispatchCtx; +use crate::families::attention::{AttentionFamily, AttnParams}; +use crate::families::fused_qkv::{FusedQkvFamily, FusedQkvParams}; +use crate::families::gemm::{GemmFamily, GemmParams}; +use crate::families::gemv::{GemvFamily, GemvParams, WeightRef}; +use crate::families::kv_tier::{KvTierInputs, KvTierPlan}; +use crate::types::{DispatchError, GemvVariant, KernelKey}; +use hip_bridge::DeviceBuffer; +use rdna_compute::{DType, Gpu, GpuTensor}; +use std::sync::LazyLock; + +fn hip(e: impl std::fmt::Display) -> DispatchError { + DispatchError::Hip(e.to_string()) +} + +fn hip_dbg(e: impl std::fmt::Debug) -> DispatchError { + DispatchError::Hip(format!("{e:?}")) +} + +/// Developer route switches. Each fused route is on unless its variable is +/// `0`/`off`/`false`. `HIPFIRE_GEMMA4_EAGLE=1` (greedy EAGLE) additionally +/// turns off the fusions that are not byte-identical to their unfused form, +/// so batched verify and single-row decode share one arithmetic. +struct Routes { + attn_norm: bool, + q8_qk: bool, + qk_norm_rope: bool, + post_norm: bool, + gate_up: bool, + /// Batched Q8 GEMMs use the unchunked scalar kernel (decode's arithmetic). + strict_q8_gemm: bool, +} + +static ROUTES: LazyLock = LazyLock::new(|| { + let on = |name: &str| { + !matches!( + hipfire_config::developer_var(name).ok().as_deref(), + Some("0") | Some("off") | Some("false") + ) + }; + let strict = hipfire_config::developer_var("HIPFIRE_GEMMA4_EAGLE") + .ok() + .as_deref() + == Some("1"); + Routes { + attn_norm: on("HIPFIRE_GEMMA4_FUSED_ATTN_NORM"), + q8_qk: on("HIPFIRE_GEMMA4_FUSED_QK"), + qk_norm_rope: !strict && on("HIPFIRE_GEMMA4_FUSED_QK_ROPE"), + post_norm: !strict && on("HIPFIRE_GEMMA4_FUSED_POSTNORM"), + gate_up: on("HIPFIRE_GEMMA4_FUSED_FFN"), + strict_q8_gemm: strict, + } +}); + +static GEMV: LazyLock = LazyLock::new(GemvFamily::new); +static GEMM: LazyLock = LazyLock::new(GemmFamily::new); +static FUSED_QKV: LazyLock = LazyLock::new(FusedQkvFamily::new); +static ATTENTION: LazyLock = LazyLock::new(AttentionFamily::new); + +/// Arches whose batched sandwich prefill admits the exact fused and wide +/// routes below. Elsewhere batched rows take the per-row fallbacks. +fn batched_fusion_arch(arch: &str) -> bool { + matches!(arch, "gfx1100" | "gfx1201") +} + +fn gemv( + gpu: &mut Gpu, + ctx: &DispatchCtx, + w: &WeightRef, + x: &GpuTensor, + y: &GpuTensor, +) -> Result<(), DispatchError> { + GEMV.run_auto(ctx, gpu, w, x, y) +} + +fn gemv_prerotated( + gpu: &mut Gpu, + ctx: &DispatchCtx, + w: &WeightRef, + x_rot: &GpuTensor, + y: &GpuTensor, +) -> Result<(), DispatchError> { + GEMV.run( + ctx, + gpu, + &GemvParams { + w, + x: x_rot, + y, + variant: GemvVariant::Prerotated, + residual: None, + gate: None, + up: None, + }, + ) +} + +/// FWHT rotation of `rows` rows of `x` for the MQ weight `w` (AWQ-aware). +fn rotate_x_mq_rows( + gpu: &mut Gpu, + w: &WeightRef, + x: &GpuTensor, + x_rot: &GpuTensor, + rows: usize, +) -> Result<(), DispatchError> { + match w.awq_scale { + Some(awq) => gpu.rotate_x_mq_awq_batched(x, awq, x_rot, w.k, rows), + None => gpu.rotate_x_mq_batched(x, x_rot, w.k, rows), + } + .map_err(hip) +} + +/// `y[rows, m] = x[rows, k] · Wᵀ`. MagnumQuant weights rotate `x` into +/// `x_rot` (`[rows, k]`) first. +pub fn gemm_rows( + gpu: &mut Gpu, + w: &WeightRef, + x: &GpuTensor, + y: &GpuTensor, + x_rot: &GpuTensor, + rows: usize, +) -> Result<(), DispatchError> { + match w.dtype { + DType::F32 => gpu + .gemm_f32_batched(w.buf, x, y, w.m, w.k, rows) + .map_err(hip), + DType::Q8_0 if ROUTES.strict_q8_gemm => gpu + .gemm_q8_0_batched(w.buf, x, y, w.m, w.k, rows) + .map_err(hip), + DType::Q8_0 => gpu + .gemm_q8_0_batched_chunked(w.buf, x, y, w.m, w.k, rows) + .map_err(hip), + DType::MQ4G256 | DType::HFQ4G256 => { + rotate_x_mq_rows(gpu, w, x, x_rot, rows)?; + gpu.gemm_hfq4g256_batched_lmhead(w.buf, x_rot, y, w.m, w.k, rows) + .map_err(hip) + } + DType::MQ6G256 => { + rotate_x_mq_rows(gpu, w, x, x_rot, rows)?; + gpu.gemm_mq6g256_batched_lmhead(w.buf, x_rot, y, w.m, w.k, rows) + .map_err(hip) + } + other => Err(DispatchError::Hip(format!( + "sandwich: dtype {other:?} has no batched projection kernel" + ))), + } +} + +/// Weight formats [`gemm_rows`] can project with `rows > 1`. +pub fn supports_batched_projection(dtype: DType) -> bool { + matches!( + dtype, + DType::F32 | DType::Q8_0 | DType::MQ4G256 | DType::HFQ4G256 | DType::MQ6G256 + ) +} + +fn fused_projection( + gpu: &mut Gpu, + ctx: &DispatchCtx, + key: KernelKey, + weights: &[&WeightRef], + x: &GpuTensor, + outputs: &[&GpuTensor], + rows: usize, +) -> Result<(), DispatchError> { + let bufs: Vec<&GpuTensor> = weights.iter().map(|w| w.buf).collect(); + let m: Vec = weights.iter().map(|w| w.m).collect(); + FUSED_QKV.run( + ctx, + gpu, + &FusedQkvParams { + kind: key, + weights: &bufs, + x, + outputs, + m: &m, + k: weights[0].k, + rot_scratch: &[], + batch_size: Some(rows), + }, + ) +} + +/// gfx1100 opt-in: Q8 projections sharing one input run as one fused WMMA GEMM. +fn q8_fused_prefill(gpu: &Gpu, weights: &[&WeightRef]) -> bool { + gpu.arch == "gfx1100" + && gpu.flags.gemma4_q8_fused_prefill + && weights + .iter() + .all(|w| w.dtype == DType::Q8_0 && w.k == weights[0].k) + && weights[0].k.is_multiple_of(32) +} + +fn scale(gpu: &mut Gpu, x: &GpuTensor, factor: f32) -> Result<(), DispatchError> { + rdna_compute::tensor_ops::scale_f32( + gpu, + &rdna_compute::tensor_ops::ScaleF32 { + values: x, + scale: factor, + }, + ) + .map_err(hip) +} + +/// Residual stream and the fixed scratch every sandwich sublayer shares. +/// With `rows > 1` each tensor holds `rows` row-major rows. +#[derive(Clone, Copy)] +pub struct SandwichStream<'a> { + pub rows: usize, + pub hidden: usize, + pub eps: f32, + /// Residual stream; updated in place. + pub x: &'a GpuTensor, + /// Holds `x` at sublayer entry. + pub residual: &'a GpuTensor, + /// Normed activation / sublayer output scratch. + pub normed: &'a GpuTensor, +} + +impl SandwichStream<'_> { + fn bytes(&self) -> usize { + self.rows * self.hidden * 4 + } + + fn save_residual(&self, gpu: &mut Gpu) -> Result<(), DispatchError> { + if self.rows == 1 { + gpu.memcpy_dtod_auto(&self.residual.buf, &self.x.buf, self.bytes()) + } else { + gpu.hip + .memcpy_dtod_at(&self.residual.buf, 0, &self.x.buf, 0, self.bytes()) + } + .map_err(hip) + } + + fn restore_residual(&self, gpu: &mut Gpu) -> Result<(), DispatchError> { + if self.rows == 1 { + gpu.memcpy_dtod_auto(&self.x.buf, &self.residual.buf, self.bytes()) + } else { + gpu.hip + .memcpy_dtod_at(&self.x.buf, 0, &self.residual.buf, 0, self.bytes()) + } + .map_err(hip) + } + + fn norm( + &self, + gpu: &mut Gpu, + x: &GpuTensor, + w: &GpuTensor, + out: &GpuTensor, + ) -> Result<(), DispatchError> { + if self.rows == 1 { + gpu.rmsnorm_f32(x, w, out, self.eps) + } else { + gpu.rmsnorm_batched(x, w, out, self.rows, self.hidden, self.eps) + } + .map_err(hip) + } + + /// `x = residual + post_norm(out)`; `out` may alias `normed`. + fn post_norm_residual( + &self, + gpu: &mut Gpu, + out: &GpuTensor, + post_norm: &GpuTensor, + ) -> Result<(), DispatchError> { + if self.rows == 1 && ROUTES.post_norm { + return gpu + .rmsnorm_residual_add_f32(out, post_norm, self.residual, self.x, self.eps) + .map_err(hip); + } + self.norm(gpu, out, post_norm, self.normed)?; + self.restore_residual(gpu)?; + gpu.add_inplace_f32(self.x, self.normed).map_err(hip) + } +} + +/// Rotary position embedding applied to Q and K. +#[derive(Clone, Copy, Debug, PartialEq)] +pub enum RopeKind { + /// Every channel pair of the head rotates. + RotateHalf, + /// The first `rot_pairs` pairs of each half rotate; the rest pass through. + PartialHalved { rot_pairs: usize }, +} + +#[derive(Clone, Copy, Debug, PartialEq)] +pub struct Rope { + pub kind: RopeKind, + pub theta: f32, +} + +impl Rope { + fn rot_pairs(&self, head_dim: usize) -> usize { + match self.kind { + RopeKind::RotateHalf => head_dim / 2, + RopeKind::PartialHalved { rot_pairs } => rot_pairs, + } + } +} + +/// The K/V cache one attention sublayer reads, and whether it writes the +/// current rows first. `write == false` reads another layer's populated cache +/// (KV sharing, draft heads). +pub struct SandwichKv<'a> { + pub tier: KvTierInputs, + pub k_cache: &'a GpuTensor, + pub v_cache: &'a GpuTensor, + pub physical_cap: usize, + pub givens_cos: Option<&'a GpuTensor>, + pub givens_sin: Option<&'a GpuTensor>, + /// Sliding window; `0` = full causal. + pub window: usize, + pub write: bool, +} + +/// Attention sublayer: +/// `x = residual + post_norm(o(attend(rope(q_norm(q)), rope(k_norm(k)), v_norm(v))))` +/// over `input_norm(x)`. +/// +/// `wk == None` projects only Q and attends a cache it does not write. +/// `wv == None` with `wk` present takes V from the pre-norm K (K=V). +/// Row `r` sits at absolute position `position + r`. +pub struct SandwichAttentionOp<'a> { + pub stream: SandwichStream<'a>, + pub position: usize, + pub n_heads: usize, + pub n_kv_heads: usize, + pub head_dim: usize, + pub input_norm: &'a GpuTensor, + pub wq: WeightRef<'a>, + pub wk: Option>, + pub wv: Option>, + pub q_norm: &'a GpuTensor, + pub k_norm: &'a GpuTensor, + /// Weight-less V RMSNorm, bound as a ones vector of `head_dim`. + pub v_norm: Option<&'a GpuTensor>, + /// Q multiplier applied after q_norm (sets the softmax scale against the + /// kernel's `1/sqrt(head_dim)`). + pub q_scale: f32, + pub rope: Rope, + pub kv: SandwichKv<'a>, + /// Device position scalar (`rows == 1`). + pub pos_buf: &'a DeviceBuffer, + /// Per-row i32 positions (`rows > 1`). + pub positions: Option<&'a GpuTensor>, + pub wo: WeightRef<'a>, + pub post_norm: &'a GpuTensor, + /// FWHT rotation scratch, `[rows, max(hidden, n_heads*head_dim)]`. + pub x_rot: &'a GpuTensor, + pub q: &'a GpuTensor, + pub k: &'a GpuTensor, + pub v: &'a GpuTensor, + pub attn_out: &'a GpuTensor, + pub flash_partials: &'a GpuTensor, +} + +impl SandwichAttentionOp<'_> { + fn projections(&self) -> [Option<(&WeightRef<'_>, &GpuTensor)>; 3] { + [ + Some((&self.wq, self.q)), + self.wk.as_ref().map(|w| (w, self.k)), + self.wv.as_ref().map(|w| (w, self.v)), + ] + } + + fn kv_bytes(&self) -> usize { + self.stream.rows * self.n_kv_heads * self.head_dim * 4 + } + + /// `input_norm(x)` → Q, K, V projections (V from K when K=V). + fn project(&self, gpu: &mut Gpu, ctx: &DispatchCtx) -> Result<(), DispatchError> { + let s = &self.stream; + if s.rows > 1 { + return self.project_rows(gpu, ctx); + } + let all_mq4 = self + .projections() + .iter() + .flatten() + .all(|(w, _)| w.dtype == DType::MQ4G256); + if ROUTES.attn_norm && all_mq4 { + gpu.fused_rmsnorm_rotate_mq(s.x, self.input_norm, self.x_rot, s.hidden, s.eps) + .map_err(hip)?; + for (w, out) in self.projections().iter().flatten() { + gemv_prerotated(gpu, ctx, w, self.x_rot, out)?; + } + } else { + s.norm(gpu, s.x, self.input_norm, s.normed)?; + match &self.wk { + Some(wk) + if ROUTES.q8_qk && self.wq.dtype == DType::Q8_0 && wk.dtype == DType::Q8_0 => + { + gpu.fused_gate_up_q8_0( + self.wq.buf, + wk.buf, + s.normed, + self.q, + self.k, + self.wq.m, + wk.m, + self.wq.k, + ) + .map_err(hip)?; + } + Some(wk) => { + gemv(gpu, ctx, &self.wq, s.normed, self.q)?; + gemv(gpu, ctx, wk, s.normed, self.k)?; + } + None => gemv(gpu, ctx, &self.wq, s.normed, self.q)?, + } + if let Some(wv) = &self.wv { + gemv(gpu, ctx, wv, s.normed, self.v)?; + } + } + if self.wk.is_some() && self.wv.is_none() { + gpu.memcpy_dtod_auto(&self.v.buf, &self.k.buf, self.kv_bytes()) + .map_err(hip)?; + } + Ok(()) + } + + fn project_rows(&self, gpu: &mut Gpu, ctx: &DispatchCtx) -> Result<(), DispatchError> { + let s = &self.stream; + let rows = s.rows; + s.norm(gpu, s.x, self.input_norm, s.normed)?; + let mut v_projected = false; + match (&self.wk, &self.wv) { + (Some(wk), Some(wv)) if q8_fused_prefill(gpu, &[&self.wq, wk, wv]) => { + fused_projection( + gpu, + ctx, + KernelKey::FusedQkvQ8_0, + &[&self.wq, wk, wv], + s.normed, + &[self.q, self.k, self.v], + rows, + )?; + v_projected = true; + } + (Some(wk), None) if q8_fused_prefill(gpu, &[&self.wq, wk]) => { + fused_projection( + gpu, + ctx, + KernelKey::FusedGateUpQ8_0, + &[&self.wq, wk], + s.normed, + &[self.q, self.k], + rows, + )?; + } + (wk, wv) => { + gemm_rows(gpu, &self.wq, s.normed, self.q, self.x_rot, rows)?; + if let Some(wk) = wk { + gemm_rows(gpu, wk, s.normed, self.k, self.x_rot, rows)?; + } + if let (Some(_), Some(wv)) = (wk, wv) { + gemm_rows(gpu, wv, s.normed, self.v, self.x_rot, rows)?; + v_projected = true; + } + } + } + if self.wk.is_some() && !v_projected { + gpu.hip + .memcpy_dtod_at(&self.v.buf, 0, &self.k.buf, 0, self.kv_bytes()) + .map_err(hip)?; + } + Ok(()) + } + + /// Per-head q_norm/k_norm, weight-less v_norm, Q prescale, RoPE. + fn position_encode(&self, gpu: &mut Gpu) -> Result<(), DispatchError> { + let eps = self.stream.eps; + let rows = self.stream.rows; + let (n_heads, n_kv, hd) = (self.n_heads, self.n_kv_heads, self.head_dim); + let writes_k = self.wk.is_some(); + if writes_k { + if let Some(ones) = self.v_norm { + gpu.rmsnorm_batched(self.v, ones, self.v, rows * n_kv, hd, eps) + .map_err(hip)?; + } + } + if rows == 1 && writes_k && ROUTES.qk_norm_rope { + return gpu + .fused_gemma4_qk_norm_rope_f32( + self.q, + self.k, + self.q_norm, + self.k_norm, + self.pos_buf, + n_heads, + n_kv, + hd, + self.rope.rot_pairs(hd), + self.q_scale, + self.rope.theta, + eps, + ) + .map_err(hip); + } + gpu.rmsnorm_batched(self.q, self.q_norm, self.q, rows * n_heads, hd, eps) + .map_err(hip)?; + if writes_k { + gpu.rmsnorm_batched(self.k, self.k_norm, self.k, rows * n_kv, hd, eps) + .map_err(hip)?; + } + scale(gpu, self.q, self.q_scale)?; + let rope_kv = if writes_k { n_kv } else { 0 }; + let theta = self.rope.theta; + match (self.rope.kind, rows) { + (RopeKind::RotateHalf, 1) => { + gpu.rope_f32(self.q, self.k, self.pos_buf, n_heads, rope_kv, hd, theta) + } + (RopeKind::PartialHalved { rot_pairs }, 1) => gpu.rope_partial_halved_f32( + self.q, + self.k, + self.pos_buf, + n_heads, + rope_kv, + hd, + rot_pairs, + theta, + ), + (RopeKind::RotateHalf, _) => gpu.rope_batched_f32( + self.q, + self.k, + self.positions()?, + n_heads, + rope_kv, + hd, + theta, + rows, + ), + (RopeKind::PartialHalved { rot_pairs }, _) => gpu.rope_partial_halved_f32_batched( + self.q, + self.k, + self.positions()?, + n_heads, + rope_kv, + hd, + rot_pairs, + theta, + rows, + ), + } + .map_err(hip) + } + + fn positions(&self) -> Result<&GpuTensor, DispatchError> { + self.positions.ok_or_else(|| { + DispatchError::Hip("sandwich attention: rows > 1 needs per-row positions".into()) + }) + } + + fn attend(&self, gpu: &mut Gpu, ctx: &DispatchCtx) -> Result<(), DispatchError> { + if self.stream.rows > 1 { + return self.attend_rows(gpu); + } + let kv = &self.kv; + let plan = KvTierPlan::derive(KvTierInputs { + pos: self.position, + batch_size: 1, + q8_windowed: true, + window: kv.window as i32, + ..kv.tier + }) + .map_err(hip)?; + let io = AttnParams { + q: self.q, + k: self.k, + v: self.v, + k_cache: kv.k_cache, + v_cache: kv.v_cache, + k_scales: None, + v_scales: None, + pos_buf: self.pos_buf, + pos: self.position, + positions: None, + n_heads: self.n_heads, + n_kv_heads: self.n_kv_heads, + head_dim: self.head_dim, + physical_cap: kv.physical_cap, + batch_size: 1, + max_ctx_len: 0, + flash_partials: Some(self.flash_partials), + givens_cos: kv.givens_cos, + givens_sin: kv.givens_sin, + tree_bias: None, + block_start: 0, + block_cols: 0, + output_gate: None, + output_awq_scale: None, + output: self.attn_out, + }; + if kv.write { + ATTENTION.run_attention(ctx, gpu, &plan, &io) + } else { + ATTENTION.run_attend_only(ctx, gpu, &plan, &io) + } + } + + /// Causal attention for `rows` consecutive positions over a Q8 cache. + fn attend_rows(&self, gpu: &mut Gpu) -> Result<(), DispatchError> { + let kv = &self.kv; + if !kv.tier.quant_q8 { + return Err(DispatchError::Hip( + "sandwich attention: rows > 1 needs a Q8 KV cache".into(), + )); + } + let rows = self.stream.rows; + let positions = self.positions()?; + let (n_heads, n_kv, hd) = (self.n_heads, self.n_kv_heads, self.head_dim); + let seq_len = self.position + rows; + if kv.write { + gpu.kv_cache_write_q8_0_batched(kv.k_cache, self.k, positions, n_kv, hd, rows) + .map_err(hip)?; + gpu.kv_cache_write_q8_0_batched(kv.v_cache, self.v, positions, n_kv, hd, rows) + .map_err(hip)?; + } + // Prefill is not tree verification: the native causal kernels, and the + // windowed tile kernel for sliding layers. + if kv.window > 0 { + gpu.attention_flash_q8_0_batched_masked_windowed( + self.q, + kv.k_cache, + kv.v_cache, + self.attn_out, + positions, + n_heads, + n_kv, + hd, + kv.physical_cap, + seq_len, + rows, + self.flash_partials, + None, + 0, + 0, + kv.window as i32, + ) + } else if seq_len > 8_192 { + gpu.attention_flash_q8_0_batched_masked( + self.q, + kv.k_cache, + kv.v_cache, + self.attn_out, + positions, + n_heads, + n_kv, + hd, + kv.physical_cap, + seq_len, + rows, + self.flash_partials, + None, + 0, + 0, + ) + } else { + gpu.attention_q8_0_kv_batched_masked( + self.q, + kv.k_cache, + kv.v_cache, + self.attn_out, + positions, + n_heads, + n_kv, + hd, + kv.physical_cap, + seq_len, + rows, + None, + 0, + 0, + ) + } + .map_err(hip) + } +} + +pub fn execute_sandwich_attention( + gpu: &mut Gpu, + ctx: &DispatchCtx, + op: &SandwichAttentionOp<'_>, +) -> Result<(), DispatchError> { + if op.wk.is_none() && op.kv.write { + return Err(DispatchError::Hip( + "sandwich attention: a Q-only layer cannot write its KV row".into(), + )); + } + let s = op.stream; + s.save_residual(gpu)?; + op.project(gpu, ctx)?; + op.position_encode(gpu)?; + op.attend(gpu, ctx)?; + if s.rows == 1 { + gemv(gpu, ctx, &op.wo, op.attn_out, s.normed)?; + } else { + gemm_rows(gpu, &op.wo, op.attn_out, s.normed, op.x_rot, s.rows)?; + } + s.post_norm_residual(gpu, s.normed, op.post_norm) +} + +/// Gated MLP activation. +#[derive(Clone, Copy, Debug, PartialEq)] +pub enum Activation { + /// `gelu_pytorch_tanh(gate) * up`. + GeluTanh, +} + +/// MLP sublayer: `x = residual + post_norm(down(act(gate(n)) * up(n)))`, +/// `n = pre_norm(x)`. +pub struct SandwichMlpOp<'a> { + pub stream: SandwichStream<'a>, + pub pre_norm: &'a GpuTensor, + pub w_gate: WeightRef<'a>, + pub w_up: WeightRef<'a>, + pub w_down: WeightRef<'a>, + pub activation: Activation, + pub hidden_dim: usize, + pub post_norm: &'a GpuTensor, + /// FWHT rotation scratch, `[rows, max(hidden, hidden_dim)]`. + pub x_rot: &'a GpuTensor, + pub gate: &'a GpuTensor, + pub up: &'a GpuTensor, + pub act: &'a GpuTensor, + pub out: &'a GpuTensor, +} + +pub fn execute_sandwich_mlp( + gpu: &mut Gpu, + ctx: &DispatchCtx, + op: &SandwichMlpOp<'_>, +) -> Result<(), DispatchError> { + let s = op.stream; + let rows = s.rows; + s.save_residual(gpu)?; + if rows > 1 { + s.norm(gpu, s.x, op.pre_norm, s.normed)?; + if q8_fused_prefill(gpu, &[&op.w_gate, &op.w_up]) { + fused_projection( + gpu, + ctx, + KernelKey::FusedGateUpQ8_0, + &[&op.w_gate, &op.w_up], + s.normed, + &[op.gate, op.up], + rows, + )?; + } else { + gemm_rows(gpu, &op.w_gate, s.normed, op.gate, op.x_rot, rows)?; + gemm_rows(gpu, &op.w_up, s.normed, op.up, op.x_rot, rows)?; + } + } else if ROUTES.gate_up && op.w_gate.dtype == DType::MQ4G256 && op.w_up.dtype == DType::MQ4G256 + { + gpu.fused_rmsnorm_rotate_mq(s.x, op.pre_norm, op.x_rot, s.hidden, s.eps) + .map_err(hip)?; + gpu.fused_gate_up_hfq4g256( + op.w_gate.buf, + op.w_up.buf, + op.x_rot, + op.gate, + op.up, + op.w_gate.m, + op.w_up.m, + op.w_gate.k, + ) + .map_err(hip)?; + } else { + s.norm(gpu, s.x, op.pre_norm, s.normed)?; + gemv(gpu, ctx, &op.w_gate, s.normed, op.gate)?; + gemv(gpu, ctx, &op.w_up, s.normed, op.up)?; + } + match op.activation { + Activation::GeluTanh => { + gpu.gelu_tanh_f32(op.gate, op.act, rows * op.hidden_dim) + .map_err(hip)?; + gpu.mul_f32(op.act, op.up, op.act).map_err(hip)?; + } + } + if rows == 1 { + gemv(gpu, ctx, &op.w_down, op.act, op.out)?; + } else { + gemm_rows(gpu, &op.w_down, op.act, op.out, op.x_rot, rows)?; + } + s.post_norm_residual(gpu, op.out, op.post_norm) +} + +/// Per-layer-input branch: +/// `x = residual + post_norm(proj(gelu(gate(x)) * inputs[row, layer]))`. +/// `inputs` holds `layer_width`-wide slices for `n_layers` layers per row. +pub struct PerLayerInputOp<'a> { + pub stream: SandwichStream<'a>, + pub layer: usize, + pub layer_width: usize, + pub n_layers: usize, + pub inputs: &'a GpuTensor, + pub w_gate: WeightRef<'a>, + pub w_proj: WeightRef<'a>, + pub post_norm: &'a GpuTensor, + pub gate: &'a GpuTensor, + pub act: &'a GpuTensor, + pub out: &'a GpuTensor, +} + +/// Q8 projections whose batched GEMM reproduces the single-row GEMV exactly. +fn exact_wide_q8_gemm(dtype: DType, k: usize) -> bool { + dtype == DType::Q8_0 && k > 0 && k <= 1536 && k.is_multiple_of(32) +} + +pub fn execute_per_layer_input( + gpu: &mut Gpu, + ctx: &DispatchCtx, + op: &PerLayerInputOp<'_>, +) -> Result<(), DispatchError> { + let s = op.stream; + let (rows, width, hidden) = (s.rows, op.layer_width, s.hidden); + let packed = op.n_layers * width; + s.save_residual(gpu)?; + if rows == 1 { + let layer_input = op.inputs.sub_offset(op.layer * width, width); + gemv(gpu, ctx, &op.w_gate, s.x, op.gate)?; + gpu.gelu_tanh_f32(op.gate, op.act, width).map_err(hip)?; + gpu.mul_f32(op.act, &layer_input, op.act).map_err(hip)?; + gemv(gpu, ctx, &op.w_proj, op.act, op.out)?; + } else { + let exact = batched_fusion_arch(&gpu.arch) && gpu.flags.gemma4_ple_branch_batched_prefill; + let wide = |gpu: &mut Gpu, w: &WeightRef, x: &GpuTensor, y: &GpuTensor| { + GEMM.run_key( + KernelKey::GemmQ8_0BatchedWideExact, + ctx, + gpu, + &GemmParams { + w, + x, + y, + batch_size: rows, + }, + ) + }; + if exact && exact_wide_q8_gemm(op.w_gate.dtype, op.w_gate.k) { + wide(gpu, &op.w_gate, s.x, op.gate)?; + } else { + for row in 0..rows { + let x_row = s.x.sub_offset(row * hidden, hidden); + let gate_row = op.gate.sub_offset(row * width, width); + gemv(gpu, ctx, &op.w_gate, &x_row, &gate_row)?; + } + } + if batched_fusion_arch(&gpu.arch) && gpu.flags.gemma4_ple_activation_fused_prefill { + gpu.gemma4_ple_gelu_mul_strided_f32( + op.gate, op.inputs, op.act, rows, width, packed, op.layer, + ) + .map_err(hip_dbg)?; + } else { + gpu.gelu_tanh_f32(op.gate, op.act, rows * width) + .map_err(hip)?; + for row in 0..rows { + let act_row = op.act.sub_offset(row * width, width); + let layer_input = op.inputs.sub_offset(row * packed + op.layer * width, width); + gpu.mul_f32(&act_row, &layer_input, &act_row).map_err(hip)?; + } + } + if exact && exact_wide_q8_gemm(op.w_proj.dtype, op.w_proj.k) { + wide(gpu, &op.w_proj, op.act, op.out)?; + } else { + for row in 0..rows { + let act_row = op.act.sub_offset(row * width, width); + let out_row = op.out.sub_offset(row * hidden, hidden); + gemv(gpu, ctx, &op.w_proj, &act_row, &out_row)?; + } + } + } + s.norm(gpu, op.out, op.post_norm, s.normed)?; + s.restore_residual(gpu)?; + gpu.add_inplace_f32(s.x, s.normed).map_err(hip) +} + +/// In-place `x *= factor`. +pub struct ScaleOp<'a> { + pub x: &'a GpuTensor, + pub factor: f32, +} + +pub fn execute_scale(gpu: &mut Gpu, op: &ScaleOp<'_>) -> Result<(), DispatchError> { + scale(gpu, op.x, op.factor) +} + +/// In-place final-logit soft cap: `x = tanh(x / cap) * cap`. +pub struct SoftcapOp<'a> { + pub logits: &'a GpuTensor, + pub n: usize, + pub cap: f32, +} + +pub fn execute_softcap(gpu: &mut Gpu, op: &SoftcapOp<'_>) -> Result<(), DispatchError> { + gpu.logit_softcap_f32(op.logits, op.n, op.cap).map_err(hip) +} + +#[cfg(test)] +mod tests { + use super::{exact_wide_q8_gemm, supports_batched_projection}; + use rdna_compute::DType; + + #[test] + fn batched_projection_formats_are_explicit() { + for dtype in [ + DType::F32, + DType::Q8_0, + DType::MQ4G256, + DType::HFQ4G256, + DType::MQ6G256, + ] { + assert!(supports_batched_projection(dtype)); + } + for dtype in [ + DType::HFQ4G128, + DType::HFQ6G256, + DType::HFQ2G256, + DType::HFQ3G256, + DType::MQ3G256, + ] { + assert!(!supports_batched_projection(dtype)); + } + } + + #[test] + fn exact_wide_q8_gemm_matches_the_wide_gemv_boundary() { + assert!(exact_wide_q8_gemm(DType::Q8_0, 1536)); + assert!(!exact_wide_q8_gemm(DType::Q8_0, 1535)); + assert!(!exact_wide_q8_gemm(DType::Q8_0, 1537)); + assert!(!exact_wide_q8_gemm(DType::Q8_0, 0)); + assert!(!exact_wide_q8_gemm(DType::F32, 256)); + } +} diff --git a/crates/hipfire-dispatch/src/pipeline/steps.rs b/crates/hipfire-dispatch/src/pipeline/steps.rs index 568ec42da2..ccb27769d9 100644 --- a/crates/hipfire-dispatch/src/pipeline/steps.rs +++ b/crates/hipfire-dispatch/src/pipeline/steps.rs @@ -105,6 +105,16 @@ pub enum Step<'a> { Project(crate::pipeline::layer_ops::ProjectOp<'a>), /// Per-row broadcast add of one activation row. BroadcastAdd(crate::pipeline::layer_ops::BroadcastAddOp<'a>), + /// Sandwich-norm attention sublayer. + SandwichAttention(crate::pipeline::sandwich::SandwichAttentionOp<'a>), + /// Sandwich-norm gated MLP sublayer. + SandwichMlp(crate::pipeline::sandwich::SandwichMlpOp<'a>), + /// Sandwich-norm per-layer-input branch. + PerLayerInput(crate::pipeline::sandwich::PerLayerInputOp<'a>), + /// In-place scalar multiply. + Scale(crate::pipeline::sandwich::ScaleOp<'a>), + /// In-place final-logit soft cap. + Softcap(crate::pipeline::sandwich::SoftcapOp<'a>), /// Complete validated MoE program; granular operands come from its shared lowerer. Moe(sealed_moe::SealedMoeCall<'a>), /// Validated granular MoE stage; the sealed call is still the public authority. @@ -136,7 +146,12 @@ fn op_kind(step: &Step) -> Option { | Step::Project(_) | Step::BroadcastAdd(_) | Step::Moe(_) - | Step::MoeStage(..) => None, + | Step::MoeStage(..) + | Step::SandwichAttention(_) + | Step::SandwichMlp(_) + | Step::PerLayerInput(_) + | Step::Scale(_) + | Step::Softcap(_) => None, } } @@ -1486,6 +1501,13 @@ fn launch_op(gpu: &mut Gpu, ctx: &DispatchCtx, step: &Step) -> Result<(), Dispat Step::BroadcastAdd(op) => crate::pipeline::layer_ops::execute_broadcast_add(gpu, op), Step::Moe(call) => sealed_moe::execute_sealed(gpu, call), Step::MoeStage(op, stage) => op.execute_stage(gpu, *stage), + Step::SandwichAttention(op) => { + crate::pipeline::sandwich::execute_sandwich_attention(gpu, ctx, op) + } + Step::SandwichMlp(op) => crate::pipeline::sandwich::execute_sandwich_mlp(gpu, ctx, op), + Step::PerLayerInput(op) => crate::pipeline::sandwich::execute_per_layer_input(gpu, ctx, op), + Step::Scale(op) => crate::pipeline::sandwich::execute_scale(gpu, op), + Step::Softcap(op) => crate::pipeline::sandwich::execute_softcap(gpu, op), } } fn rmsnorm_out<'a>(step: &Step<'a>) -> &'a rdna_compute::GpuTensor { diff --git a/docs/design/gemma4-declarative-program.md b/docs/design/gemma4-declarative-program.md new file mode 100644 index 0000000000..eaca2a1a00 --- /dev/null +++ b/docs/design/gemma4-declarative-program.md @@ -0,0 +1,85 @@ +# Gemma 4 declarative layer program (#666 G6) + +Status: in progress on `feat/gemma-declarative` (cut from #795 +`feat/qwen35-declarative` at `d101a15d0`). + +## Goal + +Gemma 4 declares each decoder layer as one typed `Step` list, as Qwen3.5 +(#795) and Qwen4 do. `hipfire-dispatch` owns every executor body and every +arch-, dtype- or tier-gated route choice. One program serves decode +(`rows == 1`), batched prefill and EAGLE verify (`rows > 1`), and the arch-22 +EAGLE/MTP drafter. + +The architecture crate keeps: + +- weight and state ownership (one config, one weight set, KV allocation); +- the per-layer program shape (`src/program.rs`: which ops, in which order, + bound to which tensors); +- lifecycle: hipGraph capture, the spec-decode driver, Redline shadow, the + chunk loop. + +The #397 super-op facade (`lowered.rs` `lower_variant` / `Gemma4Bindings`) +and the hand-written decode arms in `forward.rs` and `lowered.rs` are deleted. +Gemma 4 has no EP/TP route, so nothing keeps the super-op program alive for it. + +## Today + +Two parallel stacks, each with its own config, weights and loader: + +| Stack | Fixtures | KV | Notes | +|---|---|---|---| +| eager (`config.rs`, `gemma4.rs`, `forward.rs`) | 12B dense (default), E2B/E4B, EAGLE target | Q8 both tiers, `physical_cap = max_seq` | fused qk-norm+RoPE and fused post-norm+residual; hipGraph decode; `forward_batch{,_spec}` B≤64 | +| lowered (`lowered.rs`) | 26B-A4B MoE (always); 12B with `HIPFIRE_{BATCHED,WMMA}_PREFILL=1` | sliding Q8 ring (`window`), full asym3 hd512 | super-op facade; per-token serve prefill; `forward_prefill_batch` only for calibration tools | + +## Vocabulary + +New model-neutral ops live in `hipfire-dispatch/src/pipeline/sandwich.rs` +(not `deltanet`-gated). Each executor takes `rows`. + +| Step | Semantics | Route choice owned by dispatch | +|---|---|---| +| `SandwichAttention` | `x = residual + post_norm(o_proj(attend(rope(norm_h(q)), rope(norm_h(k)), norm(v))))` on `input_norm(x)`. Optional V projection (`None` = K=V from pre-norm K); optional weightless V norm; Q prescale; **explicit `Rope { kind: RotateHalf \| PartialHalved { rot_pairs }, theta }`**; KV `write: Option` (`None` = read-only shared slot, used by E-series KV sharing and the drafter); `window` | fused norm+FWHT-rotate + prerotated GEMVs (MQ4), fused Q8 q+k, fused qk-norm+RoPE, fused post-norm+residual, KV tier via `KvTierPlan` | +| `SandwichMlp` | `x = residual + post_norm(down(act(gate(n)) * up(n)))`, `n = pre_norm(x)`, `act ∈ {GeluTanh}` | fused norm+rotate + fused gate/up (MQ4), fused post-norm+residual | +| `ParallelMoeMlp` | Gemma 4 MoE block: dense GeGLU MLP ∥ routed experts (router input `rms(x)·router_scale/√dim`, softmax top-k renorm, per-expert scale, GeGLU), `post_ffn_norm(norm1(mlp) + norm2(moe))`, residual add | expert dtype routes (MQ4 / HFQ4 / Q8 indexed) | +| `PerLayerInput` | E-series PLE branch: `x += post_norm(proj(gelu(gate(x)) * ple[layer]))` | batched strided PLE kernel | +| `Scale` | `x *= s` (layer scalar), replay-safe recorded scale | — | +| `Softcap` | `logits = tanh(logits / c) · c` | — | + +Projections, the final norm and the LM head reuse `Gemv` and +`RmsnormAutomatic`. + +## Slices + +1. **Eager stack** (dense, E-series, EAGLE target). `SandwichAttention`, + `SandwichMlp`, `PerLayerInput`, `Scale`, `Softcap` with `rows == 1` and + `rows > 1` executors; `decode_step_body` and `forward_batch_spec` run the + same program. The token-level PLE staging and embedding stay in the arch + crate, outside the captured body. +2. **MoE.** `ParallelMoeMlp` (`rows == 1` and `rows > 1`); 26B-A4B runs the + shared program on its ring/asym3 KV. +3. **EAGLE/MTP drafter.** The drafter block is `[SandwichAttention(q-only, + write: None, target KV), SandwichMlp, Scale]`; pre/post projection and head + are `Gemv` steps. +4. **Deletion.** `lowered.rs` super-op facade, hand arms, duplicate + config/weights/loader; calibration tools move to the shared batched program. + +## Gates per slice + +- 12B dense, E2B, E4B AR: greedy token stream and the per-step top-16 logit + trace (`HIPFIRE_GEMMA4_LOGIT_TRACE_DIR`) byte-identical to the pre-port + daemon (`8a6c4710`). +- 26B-A4B MoE: the shared ops take the eager route's certified fusions + (qk-norm+RoPE, post-norm+residual) that the lowered path never used, so its + logits move in the last bits. Gate: coherent decoded text and recorded + divergence point against the pre-port daemon, same relaxation #795 applied + to its MTP and decode-arm slices. +- EAGLE drafter: tau parity and coherent text; greedy EAGLE output equals AR + output where it did before. +- Redline: the 26B shadow oracle (`redline_shadow_gemma4`) stays exact. + +## Progress + +| Piece | State | Commit | +|---|---|---| +| Eager decode + batched prefill/verify + E-series PLE as `[SandwichAttention, SandwichMlp, PerLayerInput?, Scale?]` (`program.rs`, dispatch `sandwich.rs`); hand arms deleted | landed; bitwise | (this commit) | diff --git a/docs/governance/debt-dispatch-bypass.txt b/docs/governance/debt-dispatch-bypass.txt index 4ab303a18b..570d09c9be 100644 --- a/docs/governance/debt-dispatch-bypass.txt +++ b/docs/governance/debt-dispatch-bypass.txt @@ -34,7 +34,7 @@ hipfire-arch-qwen35 partial 184 191 migrate last; MTP routed MoE now dispatch-layer owned; +10 = DFlash fp16-X GEMM entries (S3/S4 launch fusion); -20 = MQ v2 lm_head match deduped in speculative.rs; +7 = MQ4-Lloyd (qt=52) FP8-LUT prefill arms (gemm_*_fp8_lloyd) pending registry keys with LUT plumbing (gfx11 MMQ-LUT PR); +15 = MQ4-Lloyd gfx11 MMQ-LUT prefill arms (gemm_*_mmq_lloyd), same registry-key debt as the FP8-LUT arms; +40 = speculation-decode additions recorded by the SCS hygiene pass (was carried unrecorded on the beta integration base); +1 from the SCS port PR series (landed with the SCS fold); +2 landed on beta since this ledger was frozen (recorded on rebase, not new debt); -5 = sealed MoE prefill moved into dispatch (PR #755 via #774); +4 = gfx1151 A4-fusion gate/up + down-sidecar arms (gemm_gate_up_silu_a4_*, gemm_*_iu4_down_prepared) beside the Iu4H direct calls they mirror, same registry-key debt hipfire-arch-deepseek4 debt 19 0 cheapest full conversion; no registry use to reconcile -hipfire-arch-gemma4 partial 17 4 +hipfire-arch-gemma4 partial 6 4 hipfire-arch-lfm2moe debt 17 0 cheapest full conversion; no registry use to reconcile hipfire-arch-minimax debt 12 0 hipfire-arch-muse-glimmer partial 12 17 From 1e58135c3c8f9fb3b8f09b0fdad6ca2ab7b4a6dd Mon Sep 17 00:00:00 2001 From: Bjoern Agent Date: Fri, 2 Oct 2026 19:26:43 +0200 Subject: [PATCH 02/13] refactor(gemma4): MoE and EAGLE draft head on the declarative program; drop super-op path arch-22 EAGLE/MTP draft head now run the same typed step program as the eager stack: - `ParallelMoeMlp` (dispatch `sandwich.rs`): dense GeGLU MLP in parallel with the routed experts (router norm/scale, softmax top-8 renorm, indexed gate/up + scaled indexed down, per-expert scale, triple post-norms), moved verbatim from lowered.rs `apply_moe_branch`. Expert formats without an indexed kernel pair refuse at load instead of taking a host D2H loop. - `program.rs` binds either weight stack (eager, lowered, drafter) through one `Geometry` / `LayerRefs` / `ProgramBinding`; K/V projection is a typed `Option` (query-only layers read a cache they do not write). - Draft step = `[Gemv(pre_projection), query-only [SandwichAttention, SandwichMlp, Scale?] x n, RmsnormAutomatic, Gemv(lm_head), Gemv(post_projection)]` over the target's last slot of each type. - Deleted: lowered.rs super-op facade (`Gemma4Variant`, `lower_variant`, `Gemma4Bindings`, `HIPFIRE_FORWARD_LOWERED` gate), the lowered decode hand arms, dead v1 prefill, `drafter_layer`, and the Gemma-only dead super-op vocabulary in superop.rs (`AttnFlavor`, `ActFlavor`, `RopeFlavor`, `EscapeKind::GemmaLogitSoftcap`). lowered.rs -2.5k lines. - Residual / K=V copies are `copy_f32_buffer` launches (visible to hipGraph capture and the Redline recorder); a D2D memcpy made PM4 replay inexact. Gate (gfx1151, pre-port daemon 8a6c4710, greedy 5 prompts x 128 tokens): - 12B, 12B EAGLE-strict, 12B PREFILL_BATCH=64, E2B/E4B decode and PREFILL_BATCH=64: tokens + top-16 logit traces byte-identical. - 12B + gemma-4-12B-it-assistant (q8f16) EAGLE spec 3: text identical, tau per prompt unchanged 3.368/3.459/2.667/3.447/3.514. - gemma4-26b-a4b.q8-experts-reference: with HIPFIRE_GEMMA4_FUSED_QK_ROPE=0 HIPFIRE_GEMMA4_FUSED_POSTNORM=0 byte-identical (faithful move); default takes the eager fused routes: 2/5 identical, 3/5 diverge late, coherent. - 12B HIPFIRE_BATCHED_PREFILL=1 (lowered opt-in): byte-identical with every HIPFIRE_GEMMA4_FUSED_* at 0. - 26B redline_daemon_harness --pm4 --skip-prefill: shadow exact, 1083 launches (was 1263), tape hash 249f24e5f39d425d (was 0909f3792962c37d). - gemma4-26b-a4b.mq4 (MQ6G256 gate_up / HFQ4G128 down) produced garbage pre-port via the host expert loop; it now refuses at load. Not ported: lowered::forward_prefill_batch (tool-only batched prefill with calibration taps); see docs/design/gemma4-declarative-program.md. Rebased onto beta: beta's lowered hipGraph capture/replay (lowered_graph_binding, gfx1201 default) now captures the program body; the fingerprint drops the deleted HIPFIRE_FORWARD_LOWERED switch. (cherry picked from commit 02e070a0ebfa3755683c6dccae7196868db938e8) --- crates/hipfire-arch-gemma4/src/carrier.rs | 25 + crates/hipfire-arch-gemma4/src/drafter.rs | 321 +- crates/hipfire-arch-gemma4/src/forward.rs | 41 +- crates/hipfire-arch-gemma4/src/lowered.rs | 2671 ++--------------- crates/hipfire-arch-gemma4/src/program.rs | 622 +++- .../hipfire-dispatch/src/pipeline/sandwich.rs | 309 +- crates/hipfire-dispatch/src/pipeline/steps.rs | 6 + .../hipfire-dispatch/src/pipeline/superop.rs | 56 +- docs/design/gemma4-declarative-program.md | 27 +- 9 files changed, 1115 insertions(+), 2963 deletions(-) diff --git a/crates/hipfire-arch-gemma4/src/carrier.rs b/crates/hipfire-arch-gemma4/src/carrier.rs index 18d23c7276..8dd61ce61a 100644 --- a/crates/hipfire-arch-gemma4/src/carrier.rs +++ b/crates/hipfire-arch-gemma4/src/carrier.rs @@ -118,6 +118,27 @@ fn append_cleanup_context(op_err: String, cleanup: Result<(), String>) -> String } } +/// Refuse MoE checkpoints whose expert formats have no indexed device kernel +/// pair (the routed-expert step is capture-safe and never falls back to a +/// host loop). +fn unsupported_expert_formats(weights: &lowered::Gemma4Weights) -> Option { + weights.layers.iter().find_map(|layer| { + let moe = match layer { + lowered::LayerWeights::Sliding(l) => l.moe.as_ref(), + lowered::LayerWeights::Full(l) => l.moe.as_ref(), + }?; + let expert = moe.experts.first()?; + let (gate_up, down) = (expert.gate_up_proj.gpu_dtype, expert.down_proj.gpu_dtype); + (!hipfire_dispatch::pipeline::sandwich::RoutedExperts::supports(gate_up, down)).then(|| { + format!( + "gemma4 MoE: expert formats gate_up={gate_up:?} down={down:?} have no indexed \ + kernel; requantize experts to MQ4G256/HFQ4G256/HFQ6G256/Q8_0 gate_up with \ + Q8_0/HFQ4G128 down" + ) + }) + }) +} + fn free_lowered_weights(weights: lowered::Gemma4Weights, gpu: &mut Gpu) { weights.free_gpu(gpu); } @@ -253,6 +274,10 @@ pub fn load_gemma4_bundle(src: ModelSource, ctx: &mut LoadCtx) -> Result v, Err(e) => { diff --git a/crates/hipfire-arch-gemma4/src/drafter.rs b/crates/hipfire-arch-gemma4/src/drafter.rs index 99ea7cb260..99d8f460bf 100644 --- a/crates/hipfire-arch-gemma4/src/drafter.rs +++ b/crates/hipfire-arch-gemma4/src/drafter.rs @@ -39,17 +39,16 @@ //! forward at a FIXED position. `use_ordered_embeddings`/centroid masking is //! disabled (the shipped head ties lm_head directly). //! -//! REUSE: this introduces NO new kernels. The pre/post projections and lm_head -//! are plain `weight_gemv`; attention is the existing `attention_q8_0_kv_swa` -//! (sliding) / `attention_q8_0_kv` (full) called WITHOUT the kv-write step -//! (drafter does not own the cache); norms/rope/gelu/rmsnorm are the same -//! helpers `crate::forward` uses for the target. The target's `decode_step`, -//! `forward_batch`, weights and state are READ-ONLY here and byte-unchanged. +//! A draft step is one declarative step list: `Gemv(pre_projection)`, one +//! query-only `[SandwichAttention, SandwichMlp, Scale?]` block per layer over +//! the target's last cache slot of the matching type (`crate::program`, the +//! same ops the target runs), the final norm, the tied `lm_head` and +//! `post_projection`. The target's weights and state are READ-ONLY here. use crate::config::{Gemma4Config, LayerType, RopeType}; use crate::gemma4::{Gemma4State, Gemma4Weights}; use hipfire_runtime::hfq::{load_awq_scale, HfqFile}; -use hipfire_runtime::llama::{f16_to_f32, weight_gemv, WeightTensor}; +use hipfire_runtime::llama::{f16_to_f32, WeightTensor}; use rdna_compute::{DType, Gpu, GpuTensor}; pub const DRAFTER_ARCH_ID: u32 = 22; @@ -648,6 +647,8 @@ pub struct Gemma4DrafterScratch { pub x: GpuTensor, // [hidden] = 1024 residual stream pub residual: GpuTensor, // [hidden] pub tmp: GpuTensor, // [hidden] norm / o_proj scratch + /// FWHT rotation scratch for MagnumQuant projections, `[max_q_dim]`. + pub x_rot: GpuTensor, pub q: GpuTensor, // [max_q_dim] pub attn_out: GpuTensor, // [max_q_dim] pub gate_ffn: GpuTensor, // [hidden_dim] @@ -659,6 +660,8 @@ pub struct Gemma4DrafterScratch { pub logits: GpuTensor, // [vocab] /// device i32 query position scalar (constant across a round). pub pos_buf: hip_bridge::DeviceBuffer, + /// Host copy of the query position staged in `pos_buf`. + pub query_pos: usize, } impl Gemma4DrafterScratch { @@ -677,6 +680,7 @@ impl Gemma4DrafterScratch { x: alloc(gpu, cfg.hidden, "x")?, residual: alloc(gpu, cfg.hidden, "residual")?, tmp: alloc(gpu, cfg.hidden, "tmp")?, + x_rot: alloc(gpu, cfg.max_q_dim().max(cfg.hidden), "x_rot")?, q: alloc(gpu, cfg.max_q_dim(), "q")?, attn_out: alloc(gpu, cfg.max_q_dim(), "attn_out")?, gate_ffn: alloc(gpu, cfg.hidden_dim, "gate_ffn")?, @@ -687,6 +691,7 @@ impl Gemma4DrafterScratch { post_proj: alloc(gpu, cfg.backbone_hidden, "post_proj")?, logits: alloc(gpu, cfg.vocab_size, "logits")?, pos_buf, + query_pos: 0, }) } @@ -699,6 +704,7 @@ impl Gemma4DrafterScratch { let _ = gpu.free_tensor(self.x); let _ = gpu.free_tensor(self.residual); let _ = gpu.free_tensor(self.tmp); + let _ = gpu.free_tensor(self.x_rot); let _ = gpu.free_tensor(self.q); let _ = gpu.free_tensor(self.attn_out); let _ = gpu.free_tensor(self.gate_ffn); @@ -759,6 +765,7 @@ pub fn drafter_step( let pos_bytes = unsafe { std::slice::from_raw_parts(pos_host.as_ptr() as *const u8, 4) }; gpu.memcpy_htod_auto(&ds.pos_buf, pos_bytes) .map_err(|e| format!("gemma4-drafter: htod pos: {e:?}"))?; + ds.query_pos = query_pos; } // ── (a) concat = [ target_embed(prev)·√bb ‖ hidden_backbone ] (7680) ── @@ -817,6 +824,7 @@ pub fn drafter_step_from_concat( let pos_bytes = unsafe { std::slice::from_raw_parts(pos_host.as_ptr() as *const u8, 4) }; gpu.memcpy_htod_auto(&ds.pos_buf, pos_bytes) .map_err(|e| format!("gemma4-drafter: htod pos: {e:?}"))?; + ds.query_pos = query_pos; } if concat.len() != dcfg.pre_proj_in() { return Err(format!( @@ -841,48 +849,109 @@ fn drafter_step_from_concat_inner( target_state: &Gemma4State, target_cfg: &Gemma4Config, ) -> Result { - let eps = dcfg.norm_eps; + use crate::program::{Geometry, LayerKv, LayerRefs, LayerScratch, ProgramBinding, Resident}; + use hipfire_dispatch::pipeline::{execute_steps, GemvInput, Step}; - // ── (b) x = pre_projection · concat (7680 → 1024, no bias) ── - weight_gemv(gpu, &dw.pre_projection, &ds.concat, &ds.x) - .map_err(|e| format!("gemma4-drafter: pre_projection: {e}"))?; + // Each block attends the target's LAST cache slot of its type. + let last_slot = |n: usize, kind: &str| { + n.checked_sub(1) + .ok_or_else(|| format!("gemma4-drafter: target has no {kind} layers to share KV from")) + }; + let sliding_slot = last_slot(target_cfg.n_sliding_layers(), "sliding")?; + let full_slot = last_slot(target_cfg.n_full_layers(), "full")?; + let pos = ds.query_pos; + let binding = ProgramBinding { + geo: Geometry::drafter(dcfg), + resident: Resident { + kv_sliding: &target_state.kv_sliding, + kv_full: &target_state.kv_full, + pos_buf: &ds.pos_buf, + v_norm_ones: &target_state.v_norm_ones, + flash_partials: &target_state.q8_flash_partials, + }, + rows: 1, + position: pos, + positions: None, + scratch: LayerScratch { + x: &ds.x, + residual: &ds.residual, + normed: &ds.tmp, + attn_rot: &ds.x_rot, + mlp_rot: &ds.x_rot, + q: &ds.q, + k: &ds.q, + v: &ds.q, + attn_out: &ds.attn_out, + gate: &ds.gate_ffn, + up: &ds.up_ffn, + act: &ds.ffn_hidden, + mlp_out: &ds.ffn_out, + ple: None, + moe: None, + }, + }; - // ── (c) drafter layers ── - for layer_idx in 0..dcfg.n_layers { - drafter_layer( - gpu, - dcfg, - target_cfg, - &dw.layers[layer_idx], - target_state, - ds, - dcfg.layer_types[layer_idx], + // pre_projection → blocks → final norm → tied lm_head (no softcap) and + // post_projection, one step list. + let pre = dw.pre_projection.dispatch_ref(); + let lm_head = dw.lm_head.dispatch_ref(); + let post = dw.post_projection.dispatch_ref(); + let mut steps = Vec::with_capacity(3 * dcfg.n_layers + 4); + steps.push(Step::Gemv { + w: &pre, + input: GemvInput::Raw(&ds.concat), + out: &ds.x, + }); + for (layer_idx, layer) in dw.layers.iter().enumerate() { + let layer_type = dcfg.layer_types[layer_idx]; + let slot = match layer_type { + LayerType::Sliding => sliding_slot, + LayerType::Full => full_slot, + }; + binding.layer( + layer_idx, + LayerRefs::drafter(layer, layer_type, dcfg.hidden_dim), + LayerKv { + slot, + writes: false, + }, + &mut steps, )?; } + steps.push(Step::RmsnormAutomatic { + x: &ds.x, + norm_weight: &dw.final_norm, + x_plain: &ds.normed, + out: &ds.normed, + awq_scale: None, + k: dcfg.hidden, + eps: dcfg.norm_eps, + rotation: hipfire_dispatch::types::RotationPlan::None, + }); + steps.push(Step::Gemv { + w: &lm_head, + input: GemvInput::Raw(&ds.normed), + out: &ds.logits, + }); + steps.push(Step::Gemv { + w: &post, + input: GemvInput::Raw(&ds.normed), + out: &ds.post_proj, + }); + let ctx = hipfire_dispatch::context::DispatchCtx::new(gpu); + execute_steps(gpu, &ctx, &steps).map_err(|e| format!("gemma4-drafter: {e}"))?; + drop(steps); - // ── (d) n = model.norm(x) (1024) ── - gpu.rmsnorm_f32(&ds.x, &dw.final_norm, &ds.normed, eps) - .map_err(|e| format!("gemma4-drafter: final rmsnorm: {e:?}"))?; - - // ── (e) logits = n · embed_tokens.T (tied, full vocab, NO softcap) → argmax ── - weight_gemv(gpu, &dw.lm_head, &ds.normed, &ds.logits) - .map_err(|e| format!("gemma4-drafter: lm_head: {e}"))?; - // (final_logit_softcapping is null on the drafter — do NOT apply.) let logits = gpu .download_f32(&ds.logits) .map_err(|e| format!("gemma4-drafter: download logits: {e:?}"))?; let argmax = argmax_f32(&logits); - - // ── (f) post_proj = post_projection · n (1024 → 3840) ── - weight_gemv(gpu, &dw.post_projection, &ds.normed, &ds.post_proj) - .map_err(|e| format!("gemma4-drafter: post_projection: {e}"))?; let post_proj_hidden = gpu .download_f32(&ds.post_proj) .map_err(|e| format!("gemma4-drafter: download post_proj: {e:?}"))?; let normed_hidden = gpu .download_f32(&ds.normed) .map_err(|e| format!("gemma4-drafter: download normed: {e:?}"))?; - let _ = dw.embd_q8; Ok(DrafterStepOut { argmax, @@ -892,190 +961,6 @@ fn drafter_step_from_concat_inner( }) } -/// One drafter layer: sandwich-norm block over `ds.x`, KV-shared attention -/// against the target's last-of-type slot. Mirrors the target's -/// sliding/full `*_layer_decode` attention half + the shared FFN tail, but -/// with NO k/v projection and NO cache write. -#[allow(clippy::too_many_arguments)] -fn drafter_layer( - gpu: &mut Gpu, - dcfg: &Gemma4DrafterConfig, - target_cfg: &Gemma4Config, - lw: &DrafterLayerWeights, - target_state: &Gemma4State, - ds: &mut Gemma4DrafterScratch, - lt: LayerType, -) -> Result<(), String> { - let hidden = dcfg.hidden; - let eps = dcfg.norm_eps; - let hidden_bytes = hidden * 4; - let n_heads = dcfg.n_heads; - - let (head_dim, n_kv, window, rope) = match lt { - LayerType::Sliding => ( - dcfg.sliding_head_dim, - dcfg.sliding_n_kv_heads, - dcfg.sliding_window, - RopeFlavour::Full(dcfg.sliding_rope_theta), - ), - LayerType::Full => { - let n_rot_pairs = match dcfg.full_rope_type { - RopeType::Proportional => { - ((dcfg.full_head_dim as f32) * dcfg.full_partial_rotary_factor * 0.5) as usize - } - RopeType::Default => dcfg.full_head_dim / 2, - }; - ( - dcfg.full_head_dim, - dcfg.full_n_kv_heads, - 0usize, - RopeFlavour::PartialHalved(dcfg.full_rope_theta, n_rot_pairs), - ) - } - }; - - // The shared KV slot in the TARGET state: the LAST slot of the matching - // per-type cache (the target's last sliding / last full layer). - let (kv_cache, kv_slot) = match lt { - LayerType::Sliding => { - let slot = target_cfg - .n_sliding_layers() - .checked_sub(1) - .ok_or("gemma4-drafter: target has no sliding layers to share KV from")?; - (&target_state.kv_sliding, slot) - } - LayerType::Full => { - let slot = target_cfg - .n_full_layers() - .checked_sub(1) - .ok_or("gemma4-drafter: target has no full layers to share KV from")?; - (&target_state.kv_full, slot) - } - }; - - // residual = x - gpu.memcpy_dtod_auto(&ds.residual.buf, &ds.x.buf, hidden_bytes) - .map_err(|e| format!("gemma4-drafter: save residual: {e:?}"))?; - - // n1 = input_layernorm(x) → tmp - gpu.rmsnorm_f32(&ds.x, &lw.input_layernorm, &ds.tmp, eps) - .map_err(|e| format!("gemma4-drafter: input rmsnorm: {e:?}"))?; - - // q = q_proj(n1) - weight_gemv(gpu, &lw.q_proj, &ds.tmp, &ds.q) - .map_err(|e| format!("gemma4-drafter: q_proj: {e}"))?; - // per-head q_norm over head_dim - gpu.rmsnorm_batched(&ds.q, &lw.q_norm, &ds.q, n_heads, head_dim, eps) - .map_err(|e| format!("gemma4-drafter: q_norm: {e:?}"))?; - // q ·= √head_dim (cancels the kernel's 1/√head_dim → effective scale 1.0). - gpu.scale_f32(&ds.q, (head_dim as f32).sqrt()) - .map_err(|e| format!("gemma4-drafter: q scale: {e:?}"))?; - - // RoPE the drafter Q ONLY (n_heads_k = 0 ⇒ the kernel skips K untouched). - match rope { - RopeFlavour::Full(theta) => { - gpu.rope_f32(&ds.q, &ds.q, &ds.pos_buf, n_heads, 0, head_dim, theta) - .map_err(|e| format!("gemma4-drafter: rope: {e:?}"))?; - } - RopeFlavour::PartialHalved(theta, n_rot_pairs) => { - gpu.rope_partial_halved_f32( - &ds.q, - &ds.q, - &ds.pos_buf, - n_heads, - 0, - head_dim, - n_rot_pairs, - theta, - ) - .map_err(|e| format!("gemma4-drafter: rope partial: {e:?}"))?; - } - } - - // Attention over the target's shared K/V (NO write). pos_buf = query_pos so - // the kernel reads seq_len = query_pos+1 from pos_buf[0]+1. - let max_seq = target_state.max_seq; - let phys_cap = kv_cache.physical_cap; - let k_cache = &kv_cache.k_gpu[kv_slot]; - let v_cache = &kv_cache.v_gpu[kv_slot]; - if window > 0 { - gpu.attention_q8_0_kv_swa( - &ds.q, - k_cache, - v_cache, - &ds.attn_out, - &ds.pos_buf, - max_seq, - n_heads, - n_kv, - head_dim, - phys_cap, - window, - ) - .map_err(|e| format!("gemma4-drafter: attention swa: {e:?}"))?; - } else { - gpu.attention_q8_0_kv( - &ds.q, - k_cache, - v_cache, - &ds.attn_out, - &ds.pos_buf, - max_seq, - n_heads, - n_kv, - head_dim, - phys_cap, - ) - .map_err(|e| format!("gemma4-drafter: attention full: {e:?}"))?; - } - - // o_proj(attn_out) → tmp; post_attention_layernorm; x = residual + tmp. - weight_gemv(gpu, &lw.o_proj, &ds.attn_out, &ds.tmp) - .map_err(|e| format!("gemma4-drafter: o_proj: {e}"))?; - gpu.rmsnorm_f32(&ds.tmp, &lw.post_attention_layernorm, &ds.tmp, eps) - .map_err(|e| format!("gemma4-drafter: post_attn rmsnorm: {e:?}"))?; - gpu.memcpy_dtod_auto(&ds.x.buf, &ds.residual.buf, hidden_bytes) - .map_err(|e| format!("gemma4-drafter: reset x: {e:?}"))?; - gpu.add_inplace_f32(&ds.x, &ds.tmp) - .map_err(|e| format!("gemma4-drafter: attn residual add: {e:?}"))?; - - // residual = x; FFN: GeGLU(gelu_tanh) → down_proj; post_ffn norm; x += ffn. - gpu.memcpy_dtod_auto(&ds.residual.buf, &ds.x.buf, hidden_bytes) - .map_err(|e| format!("gemma4-drafter: save ffn residual: {e:?}"))?; - gpu.rmsnorm_f32(&ds.x, &lw.pre_feedforward_layernorm, &ds.tmp, eps) - .map_err(|e| format!("gemma4-drafter: pre_ffn rmsnorm: {e:?}"))?; - weight_gemv(gpu, &lw.gate_proj, &ds.tmp, &ds.gate_ffn) - .map_err(|e| format!("gemma4-drafter: gate_proj: {e}"))?; - weight_gemv(gpu, &lw.up_proj, &ds.tmp, &ds.up_ffn) - .map_err(|e| format!("gemma4-drafter: up_proj: {e}"))?; - gpu.gelu_tanh_f32(&ds.gate_ffn, &ds.ffn_hidden, dcfg.hidden_dim) - .map_err(|e| format!("gemma4-drafter: gelu_tanh: {e:?}"))?; - gpu.mul_f32(&ds.ffn_hidden, &ds.up_ffn, &ds.ffn_hidden) - .map_err(|e| format!("gemma4-drafter: geglu mul: {e:?}"))?; - weight_gemv(gpu, &lw.down_proj, &ds.ffn_hidden, &ds.ffn_out) - .map_err(|e| format!("gemma4-drafter: down_proj: {e}"))?; - gpu.rmsnorm_f32(&ds.ffn_out, &lw.post_feedforward_layernorm, &ds.tmp, eps) - .map_err(|e| format!("gemma4-drafter: post_ffn rmsnorm: {e:?}"))?; - gpu.memcpy_dtod_auto(&ds.x.buf, &ds.residual.buf, hidden_bytes) - .map_err(|e| format!("gemma4-drafter: reset x (ffn): {e:?}"))?; - gpu.add_inplace_f32(&ds.x, &ds.tmp) - .map_err(|e| format!("gemma4-drafter: ffn residual add: {e:?}"))?; - - // learned per-layer scalar. - if lw.layer_scalar_host != 1.0 { - gpu.scale_f32(&ds.x, lw.layer_scalar_host) - .map_err(|e| format!("gemma4-drafter: layer_scalar: {e:?}"))?; - } - Ok(()) -} - -enum RopeFlavour { - /// Full rotate-half over the whole head_dim (sliding). θ. - Full(f32), - /// Partial proportional rotate-half (full layer). (θ, n_rot_pairs). - PartialHalved(f32, usize), -} - fn argmax_f32(v: &[f32]) -> u32 { let mut bi = 0u32; let mut bv = f32::NEG_INFINITY; diff --git a/crates/hipfire-arch-gemma4/src/forward.rs b/crates/hipfire-arch-gemma4/src/forward.rs index 30b82d0d88..cafe70aa8e 100644 --- a/crates/hipfire-arch-gemma4/src/forward.rs +++ b/crates/hipfire-arch-gemma4/src/forward.rs @@ -56,7 +56,9 @@ use crate::config::Gemma4Config; use crate::gemma4::{Gemma4State, Gemma4Weights, LayerWeights, GEMMA4_FORWARD_BATCH_MAX}; -use crate::program::{head, LayerScratch, PleScratch, ProgramBinding}; +use crate::program::{ + eager_layer_kv, head, Geometry, LayerRefs, LayerScratch, PleScratch, ProgramBinding, Resident, +}; use hipfire_dispatch::context::DispatchCtx; use hipfire_dispatch::pipeline::execute_steps; use hipfire_runtime::llama::{weight_gemv, WeightTensor}; @@ -475,6 +477,16 @@ fn prepare_per_layer_inputs_batched( .map_err(|e| format!("gemma4 forward_batch ple combine scale: {e:?}")) } +fn eager_resident(state: &Gemma4State) -> Resident<'_> { + Resident { + kv_sliding: &state.kv_sliding, + kv_full: &state.kv_full, + pos_buf: &state.pos_buf, + v_norm_ones: &state.v_norm_ones, + flash_partials: &state.q8_flash_partials, + } +} + fn decode_step_body( cfg: &Gemma4Config, weights: &Gemma4Weights, @@ -507,14 +519,19 @@ fn decode_step_body( for layer_idx in 0..cfg.n_layers { steps.clear(); let binding = ProgramBinding { - cfg, - state, + geo: Geometry::eager(cfg), + resident: eager_resident(state), rows: 1, position: position as usize, positions: None, - scratch: LayerScratch::decode(state), + scratch: LayerScratch::eager(state), }; - binding.layer(layer_idx, &weights.layers[layer_idx], &mut steps)?; + binding.layer( + layer_idx, + LayerRefs::eager(&weights.layers[layer_idx]), + eager_layer_kv(cfg, state, layer_idx)?, + &mut steps, + )?; execute_steps(gpu, &ctx, &steps).map_err(|e| format!("gemma4 L{layer_idx}: {e}"))?; if let Some(cap) = capture.as_deref_mut() { let h = gpu @@ -530,7 +547,7 @@ fn decode_step_body( let lm_head = weights.lm_head.dispatch_ref(); let mut head_steps = Vec::with_capacity(3); head( - cfg, + &Geometry::eager(cfg), &weights.final_norm, &lm_head, &state.x, @@ -860,8 +877,8 @@ pub fn forward_batch_spec( _ => None, }; let binding = ProgramBinding { - cfg, - state, + geo: Geometry::eager(cfg), + resident: eager_resident(state), rows: b, position: start_pos, positions: Some(&pos_array), @@ -880,11 +897,17 @@ pub fn forward_batch_spec( act: &ffn_hidden, mlp_out: &ffn_out, ple, + moe: None, }, }; let mut steps = Vec::with_capacity(4 * cfg.n_layers); for (layer_idx, layer) in weights.layers.iter().enumerate() { - binding.layer(layer_idx, layer, &mut steps)?; + binding.layer( + layer_idx, + LayerRefs::eager(layer), + eager_layer_kv(cfg, state, layer_idx)?, + &mut steps, + )?; } execute_steps(gpu, &ctx, &steps).map_err(|e| format!("gemma4 forward_batch: {e}"))?; drop(steps); diff --git a/crates/hipfire-arch-gemma4/src/lowered.rs b/crates/hipfire-arch-gemma4/src/lowered.rs index bff3cad2ed..02dc6d7821 100644 --- a/crates/hipfire-arch-gemma4/src/lowered.rs +++ b/crates/hipfire-arch-gemma4/src/lowered.rs @@ -2270,6 +2270,8 @@ pub struct Gemma4Scratch { pub x: GpuTensor, // [dim] — hidden state pub residual: GpuTensor, // [dim] — saved for sandwich residual pub tmp: GpuTensor, // [dim] — norm output scratch + /// `[dim]` FWHT-rotated normed activation shared by MagnumQuant projections. + pub x_rot: GpuTensor, /// Position buffer (single i32 on device, updated per decode step). pub pos_buf: DeviceBuffer, @@ -2477,6 +2479,7 @@ impl Gemma4Scratch { let i_x = alloc!(&[dim], DType::F32); let i_residual = alloc!(&[dim], DType::F32); let i_tmp = alloc!(&[dim], DType::F32); + let i_x_rot = alloc!(&[dim], DType::F32); match alloc_pos(gpu) { Ok(pos) => { @@ -2594,6 +2597,7 @@ impl Gemma4Scratch { x: take!(i_x), residual: take!(i_residual), tmp: take!(i_tmp), + x_rot: take!(i_x_rot), pos_buf: pos_slot.take().expect("gemma4 scratch pos_buf staged once"), q: take!(i_q), k: take!(i_k), @@ -2665,6 +2669,7 @@ impl Gemma4Scratch { let _ = gpu.free_tensor(self.x); let _ = gpu.free_tensor(self.residual); let _ = gpu.free_tensor(self.tmp); + let _ = gpu.free_tensor(self.x_rot); // pos_buf is a raw DeviceBuffer (not a GpuTensor): free it explicitly — // DeviceBuffer has no Drop-side free. Eager precedent: gemma4.rs free_gpu. let _ = gpu.hip.free(self.pos_buf); @@ -3476,300 +3481,6 @@ fn moe_repoint_f32_row( view.buf = unsafe { hip_bridge::DeviceBuffer::from_raw(ptr, byte_len) }; } -/// Apply the Gemma 4 MoE parallel branch (26B-A4B). Called from each layer -/// AFTER `down_proj` produces `scratch.ffn_out`, REPLACING the standalone -/// `post_feedforward_layernorm` call. On exit, `scratch.tmp` holds the -/// combined `post_norm(cur_mlp + cur_moe)`, ready for `x = residual + tmp`. -/// -/// Shared indexed MoE semantics via `moe_token_indexed` (fused indexed -/// gate_up + scaled indexed down, top-K device-resident). The legacy CPU -/// per-expert loop below is retained for unsupported quant mixes only. -/// -/// HF reference (modeling_gemma4.py Gemma4MoeBlock + Gemma4MoeMLP): -/// cur_mlp = post_feedforward_layernorm_1(ffn_out) # standard SwiGLU out, normed -/// pre2 = pre_feedforward_layernorm_2(attn_out) # MoE branch input -/// router_in = rmsnorm(attn_out, router_scale) / sqrt(dim) -/// router_logits = router_proj @ router_in -/// topk_idx, topk_w = softmax_topk_renorm(router_logits, k=8) -/// cur_moe = sum_k [ topk_w[k] * per_expert_scale[i_k] * -/// down_proj_{i_k}( gelu_tanh(gate) * up -/// where (gate, up) = split(gate_up_proj_{i_k} @ pre2) ) ] -/// cur_moe = post_feedforward_layernorm_2(cur_moe) -/// tmp = post_feedforward_layernorm(cur_mlp + cur_moe) -/// -/// `attn_out` parameter is the layer's post-attention residual stream -/// (= `scratch.residual` at the call site, since the caller stored -/// `residual = x` after the attention sandwich). -fn apply_moe_branch( - gpu: &mut Gpu, - config: &Gemma4Config, - scratch: &Gemma4Scratch, - moe: &MoeLayerExtras, - post_ffn_norm: &GpuTensor, - attn_out: &GpuTensor, -) -> HipResult<()> { - let dim = config.dim; - let mi = config.moe_intermediate_size; - let n_exp = config.num_experts; - let k_top = config.top_k_experts; - if k_top != 8 { - return Err(hip_bridge::HipError::new( - 0, - &format!("MoE top_k_experts={k_top} unsupported (kernel hardcoded to 8)"), - )); - } - - // 1) cur_mlp = post_feedforward_layernorm_1(ffn_out) - gpu.rmsnorm_f32( - &scratch.ffn_out, - &moe.post_feedforward_layernorm_1, - &scratch.moe_cur_mlp, - config.norm_eps, - )?; - - // 2) pre2 = pre_feedforward_layernorm_2(attn_out) - gpu.rmsnorm_f32( - attn_out, - &moe.pre_feedforward_layernorm_2, - &scratch.moe_pre2, - config.norm_eps, - )?; - - // Dump moe_pre2 on first call - { - use std::sync::atomic::{AtomicUsize, Ordering}; - static PRE2_DUMP: AtomicUsize = AtomicUsize::new(0); - let pc = PRE2_DUMP.fetch_add(1, Ordering::Relaxed); - if pc == 0 { - if let Ok(data) = gpu.download_f32(&scratch.moe_pre2) { - let sum: f64 = data.iter().map(|&v| v as f64).sum(); - eprintln!( - "[moe pre2] first4={:?} sum={sum:.4}", - &data[..4.min(data.len())] - ); - } - } - } - // 3) Router input: rmsnorm(attn_out, router_scale) / sqrt(dim). - // Equivalent to ref `rms_norm(x) * router_scale / sqrt(dim)` since - // rmsnorm_f32(x, w) = w * x / sqrt(mean(x²) + eps) — elementwise commutative. - gpu.rmsnorm_f32( - attn_out, - &moe.router_scale, - &scratch.moe_router_in, - config.norm_eps, - )?; - gpu.scale_f32_recorded(&scratch.moe_router_in, 1.0 / (dim as f32).sqrt())?; - - // 4) Router GEMV → logits [n_exp] - weight_gemv( - gpu, - &moe.router_proj, - &scratch.moe_router_in, - &scratch.moe_router_logits, - )?; - - // 5) Top-K softmax + renorm on device. Kernel hardcoded to k_top=8. - gpu.moe_softmax_topk_renorm_k8( - &scratch.moe_router_logits, - &scratch.moe_topk_indices, - &scratch.moe_topk_weights, - n_exp, - true, - )?; - - // 6) Dispatch top-K experts. Two paths: - // - // (a) Indexed (fast) path: gate_up is MQ4G256/MG4G256 and down is - // Q8_0. Two kernel launches replace the 8-iteration CPU loop + - // 2 D2H syncs/layer (=60 syncs/token on 30-layer Gemma 4 26B). - // Whole MoE branch becomes hipGraph-capturable. - // - // (b) Legacy path: any other quant mix. CPU per-expert loop with - // the 2 D2H downloads per layer. - // - // 26B-A4B-it (Gemma 4): gate_up=MQ4G256 (dim=2816 / 256-aligned), - // down=Q8_0 (mi=704 not 256-aligned → quantizer falls back to Q8_0). - // Hits the fast path. Other Gemma 4 variants might land in legacy. - let first = &moe.experts[0]; - let gate_mq4 = first.gate_up_proj.gpu_dtype == rdna_compute::DType::MQ4G256; - let gate_hfq4 = first.gate_up_proj.gpu_dtype == rdna_compute::DType::HFQ4G256; - let gate_hfq6 = first.gate_up_proj.gpu_dtype == rdna_compute::DType::HFQ6G256; - let gate_q8 = first.gate_up_proj.gpu_dtype == rdna_compute::DType::Q8_0; - let down_q8 = first.down_proj.gpu_dtype == rdna_compute::DType::Q8_0; - let down_hfq4g128 = first.down_proj.gpu_dtype == rdna_compute::DType::HFQ4G128; - let fast = (gate_mq4 || gate_hfq4 || gate_hfq6 || gate_q8) && (down_q8 || down_hfq4g128); - { - use std::sync::OnceLock; - static LOGGED: OnceLock<()> = OnceLock::new(); - LOGGED.get_or_init(|| { - eprintln!( - "[gemma4 MoE] dispatch path: {} (gate_up={:?} down={:?})", - if fast { - "indexed-fast" - } else { - "legacy-cpu-loop" - }, - first.gate_up_proj.gpu_dtype, - first.down_proj.gpu_dtype - ); - }); - } - - if fast { - // Zero accumulator (memset is sync but tiny — 11 KB for dim=2816). - // Required for Q8 down (atomicAdd); harmless for HFQ4G128 (assign). - gpu.zero_f32(&scratch.moe_cur_moe)?; - moe_token_indexed( - gpu, - moe, - first.gate_up_proj.gpu_dtype, - down_q8, - &scratch.moe_pre2, - &scratch.moe_pre2_rot, - &scratch.moe_topk_indices, - &scratch.moe_topk_weights, - &scratch.moe_expert_gate_batch, - &scratch.moe_expert_up_batch, - &scratch.moe_expert_hidden_batch, - &scratch.moe_cur_moe, - dim, - mi, - k_top, - )?; - } else { - // ── Legacy CPU per-expert path (quant mix doesn't match the - // fast kernels). 60 D2H syncs/token, no graph capture. ── - let idx_bytes = gpu.download_f32(&scratch.moe_topk_indices)?; - let topk_indices: Vec = - unsafe { std::slice::from_raw_parts(idx_bytes.as_ptr() as *const i32, k_top) } - .iter() - .map(|&i| i as usize) - .collect(); - let topk_weights = gpu.download_f32(&scratch.moe_topk_weights)?; - for &e in topk_indices.iter().take(k_top) { - if e >= n_exp { - return Err(hip_bridge::HipError::new( - 0, - &format!("MoE topk index {e} out of range (n_exp={n_exp})"), - )); - } - } - gpu.zero_f32(&scratch.moe_cur_moe)?; - // Dump router info for first MoE layer - { - use std::sync::atomic::{AtomicUsize, Ordering}; - static CALL: AtomicUsize = AtomicUsize::new(0); - let c = CALL.fetch_add(1, Ordering::Relaxed); - if c == 0 { - eprintln!( - "[moe diag] first call: topk_indices={:?} topk_weights={:?}\n per_expert_scale[0..8]={:?}", - &topk_indices[..k_top], - &topk_weights[..k_top], - &moe.per_expert_scale_host[..8.min(n_exp)] - ); - } - } - for ki in 0..k_top { - let e = topk_indices[ki]; - let weight = topk_weights[ki] * moe.per_expert_scale_host[e]; - let expert = &moe.experts[e]; - weight_gemv( - gpu, - &expert.gate_up_proj, - &scratch.moe_pre2, - &scratch.moe_expert_gate_up, - )?; - let gate = scratch.moe_expert_gate_up.sub_offset(0, mi); - let up = scratch.moe_expert_gate_up.sub_offset(mi, mi); - gpu.gelu_tanh_f32(&gate, &scratch.moe_expert_hidden, mi)?; - gpu.mul_f32(&scratch.moe_expert_hidden, &up, &scratch.moe_expert_hidden)?; - weight_gemv( - gpu, - &expert.down_proj, - &scratch.moe_expert_hidden, - &scratch.moe_expert_out, - )?; - gpu.scaled_add_inplace_cpu_scalar_f32( - &scratch.moe_cur_moe, - &scratch.moe_expert_out, - weight, - )?; - // Dump first expert's contribution - if ki == 0 { - use std::sync::atomic::{AtomicUsize, Ordering}; - static EXPERT_CALL: AtomicUsize = AtomicUsize::new(0); - let ec = EXPERT_CALL.fetch_add(1, Ordering::Relaxed); - if ec == 0 { - // gate_up output - if let Ok(gu_data) = gpu.download_f32(&scratch.moe_expert_gate_up) { - let sum: f64 = gu_data.iter().map(|&v| v as f64).sum(); - let gate_sum: f64 = gu_data[..mi].iter().map(|&v| v as f64).sum(); - let up_sum: f64 = gu_data[mi..].iter().map(|&v| v as f64).sum(); - eprintln!( - "[moe expert] expert={e} GATE_UP: gate_first4={:?} gate_sum={gate_sum:.4} up_first4={:?} up_sum={up_sum:.4}", - &gu_data[..4.min(mi)], - &gu_data[mi..mi + 4.min(mi)], - ); - } - // gate-up shape: gate_up_proj m=2*mi=1408, k=dim=2816 - eprintln!( - "[moe expert] expert={e} gate_up_proj m={} k={} dtype={:?}", - expert.gate_up_proj.m, expert.gate_up_proj.k, expert.gate_up_proj.gpu_dtype - ); - // down_proj output - if let Ok(data) = gpu.download_f32(&scratch.moe_expert_out) { - let sum: f64 = data.iter().map(|&v| v as f64).sum(); - eprintln!( - "[moe expert] expert={e} weight={weight:.6} down_dtype={:?} expert_out_sum={sum:.4} first4={:?}", - expert.down_proj.gpu_dtype, - &data[..4.min(data.len())] - ); - } - } - } - } - } - - // 9) cur_moe = post_feedforward_layernorm_2(cur_moe) — in-place - // Dump MoE branch intermediates on first call - { - use std::sync::atomic::{AtomicUsize, Ordering}; - static MOE_END: AtomicUsize = AtomicUsize::new(0); - let mc = MOE_END.fetch_add(1, Ordering::Relaxed); - if mc == 0 { - if let Ok(mlp_data) = gpu.download_f32(&scratch.moe_cur_mlp) { - let sum: f64 = mlp_data.iter().map(|&v| v as f64).sum(); - eprintln!( - "[moe branch] cur_mlp first4={:?} sum={sum:.4}", - &mlp_data[..4.min(mlp_data.len())] - ); - } - if let Ok(moe_data) = gpu.download_f32(&scratch.moe_cur_moe) { - let sum: f64 = moe_data.iter().map(|&v| v as f64).sum(); - eprintln!( - "[moe branch] cur_moe (before norm2) first4={:?} sum={sum:.4}", - &moe_data[..4.min(moe_data.len())] - ); - } - } - } - gpu.rmsnorm_f32( - &scratch.moe_cur_moe, - &moe.post_feedforward_layernorm_2, - &scratch.moe_cur_moe, - config.norm_eps, - )?; - - // 10) combined = cur_mlp + cur_moe → scratch.tmp - gpu.add_f32(&scratch.moe_cur_mlp, &scratch.moe_cur_moe, &scratch.tmp)?; - - // 11) tmp = post_feedforward_layernorm(combined) - gpu.rmsnorm_f32(&scratch.tmp, post_ffn_norm, &scratch.tmp, config.norm_eps)?; - - Ok(()) -} - /// Batched MoE branch — processes N tokens at once. Mirrors `apply_moe_branch` /// but every per-token tensor becomes a `[N × *]` row-major batch and every /// kernel call uses the `_batched` variant. @@ -4160,7 +3871,6 @@ fn lowered_graph_binding( config.full_head_dim, config.full_n_kv_heads, config.hidden_dim, config.enable_moe_block as usize, config.moe_intermediate_size, config.num_experts, config.top_k_experts, config.attention_k_eq_v as usize, - forward_lowered_enabled() as usize, ] { mix(v as u64); } @@ -4411,131 +4121,56 @@ fn forward_scratch_inner( kv_full: &mut hipfire_runtime::llama::KvCache, scratch: &Gemma4Scratch, ) -> HipResult<()> { - // #397 Ship 6 — forward-as-pipeline. When HIPFIRE_FORWARD_LOWERED=1, - // route decode through the lowered super-op executor. - if forward_lowered_enabled() { - return forward_scratch_inner_lowered( - gpu, weights, config, pos, kv_sliding, kv_full, scratch, - ); - } - - // 3) Per-layer forward (legacy hand path). - let mut sliding_kv_idx = 0usize; - let mut full_kv_idx = 0usize; - for (layer_idx, layer_type) in config.layer_types.iter().copied().enumerate() { - // Diagnostic: dump residual before layer - if hipfire_config::developer_var("HIPFIRE_GEMMA4_DUMP") - .ok() - .as_deref() - == Some("1") - && layer_idx < 2 - { - let data = gpu.download_f32(&scratch.x).unwrap_or_default(); - let sum: f64 = data.iter().map(|&v| v as f64).sum(); - let min = data.iter().fold(f32::INFINITY, |a, &b| a.min(b)); - let max = data.iter().fold(f32::NEG_INFINITY, |a, &b| a.max(b)); - eprintln!( - "[gemma4 diag] pos={pos} L{layer_idx} before: first4={:?} sum={sum:.4e} min={min:.4} max={max:.4}", - &data[..4.min(data.len())] - ); - } - match (layer_type, &weights.layers[layer_idx]) { - (LayerType::Sliding, LayerWeights::Sliding(lw)) => { - sliding_layer_decode(gpu, config, lw, pos, kv_sliding, sliding_kv_idx, scratch)?; - sliding_kv_idx += 1; - } - (LayerType::Full, LayerWeights::Full(lw)) => { - full_layer_decode(gpu, config, lw, pos, kv_full, full_kv_idx, scratch)?; - full_kv_idx += 1; - } - _ => { - return Err(hip_bridge::HipError::new( - 0, - &format!("Gemma 4 layer {} type/weights mismatch", layer_idx), - )); - } - } - // Diagnostic: dump residual after layer - if hipfire_config::developer_var("HIPFIRE_GEMMA4_DUMP") - .ok() - .as_deref() - == Some("1") - { - let data = gpu.download_f32(&scratch.x).unwrap_or_default(); - let sum: f64 = data.iter().map(|&v| v as f64).sum(); - let min = data.iter().fold(f32::INFINITY, |a, &b| a.min(b)); - let max = data.iter().fold(f32::NEG_INFINITY, |a, &b| a.max(b)); - eprintln!( - "[gemma4 diag] pos={pos} L{layer_idx} hidden: first4={:?} sum={sum:.4e} min={min:.4} max={max:.4}", - &data[..4.min(data.len())] - ); - } + use crate::program::{ + head, Geometry, LayerKv, LayerRefs, LayerScratch, ProgramBinding, Resident, + }; + let hip = |e: String| hip_bridge::HipError::new(0, &e); + let geo = Geometry::lowered(config); + let binding = ProgramBinding { + geo, + resident: Resident { + kv_sliding, + kv_full, + pos_buf: &scratch.pos_buf, + v_norm_ones: &scratch.v_norm_ones_full, + flash_partials: &scratch.flash_partials, + }, + rows: 1, + position: pos, + positions: None, + scratch: LayerScratch::lowered(scratch, config.enable_moe_block), + }; + // One cache slot per layer of each type, in layer order. + let mut slots = [0usize; 2]; + let mut steps = Vec::with_capacity(3 * config.n_layers + 3); + for (layer_idx, layer) in weights.layers.iter().enumerate() { + let class = matches!(layer, LayerWeights::Full(_)) as usize; + let kv = LayerKv { + slot: slots[class], + writes: true, + }; + slots[class] += 1; + binding + .layer( + layer_idx, + LayerRefs::lowered(layer, config.hidden_dim), + kv, + &mut steps, + ) + .map_err(hip)?; } - - // 4) Final RMSNorm. - gpu.rmsnorm_f32( - &scratch.x, + let lm_head = weights.lm_head.dispatch_ref(); + head( + &geo, &weights.final_norm, + &lm_head, + &scratch.x, &scratch.tmp, - config.norm_eps, - )?; - - // Diagnostic: dump hidden state before lm_head - if hipfire_config::developer_var("HIPFIRE_GEMMA4_DUMP") - .ok() - .as_deref() - == Some("1") - { - let data = gpu.download_f32(&scratch.tmp).unwrap_or_default(); - let sum: f64 = data.iter().map(|&v| v as f64).sum(); - let min = data.iter().fold(f32::INFINITY, |a, &b| a.min(b)); - let max = data.iter().fold(f32::NEG_INFINITY, |a, &b| a.max(b)); - eprintln!( - "[gemma4 diag] pos={pos} hidden pre-lm_head: first4={:?} sum={sum:.4e} min={min:.4} max={max:.4}", - &data[..4.min(data.len())] - ); - } - - // 5) LM head → logits (reads tied embed bytes via lm_head.buf alias). + &scratch.logits, + &mut steps, + ); let ctx = DispatchCtx::new(gpu); - let wr_lm = weights.lm_head.dispatch_ref(); - execute_steps( - gpu, - &ctx, - &[Step::Gemv { - w: &wr_lm, - input: GemvInput::Raw(&scratch.tmp), - out: &scratch.logits, - }], - ) - .map_err(|e| hip_bridge::HipError::new(0, &e.to_string()))?; - - // 6) Final logit softcap (Gemma 4): logits = tanh(logits / cap) * cap. - if config.final_logit_softcapping > 0.0 { - gpu.logit_softcap_f32( - &scratch.logits, - config.vocab_size, - config.final_logit_softcapping, - )?; - } - - // Diagnostic: dump logits - if hipfire_config::developer_var("HIPFIRE_GEMMA4_DUMP") - .ok() - .as_deref() - == Some("1") - { - let data = gpu.download_f32(&scratch.logits).unwrap_or_default(); - let top5: Vec<(usize, f32)> = { - let mut indexed: Vec<(usize, f32)> = - data.iter().enumerate().map(|(i, &v)| (i, v)).collect(); - indexed.sort_by(|a, b| b.1.partial_cmp(&a.1).unwrap()); - indexed.into_iter().take(5).collect() - }; - eprintln!("[gemma4 diag] pos={pos} logits top5: {:?}", top5); - } - - Ok(()) + execute_steps(gpu, &ctx, &steps).map_err(|e| hip(format!("gemma4: {e}"))) } /// Single sliding-window attention layer. @@ -4573,1090 +4208,56 @@ fn forward_scratch_inner( // all 32 lanes participate in __shfl_xor; window_size=0 = byte-identical // for Qwen 3.5/3.6. Runtime NRMSE validation: verify_against_torch.rs Phase 2 // (D2.5) — per-layer attention NRMSE against the bf16 PyTorch reference dump. -fn sliding_layer_decode( - gpu: &mut Gpu, - config: &Gemma4Config, - lw: &SlidingLayerWeights, - pos: usize, - kv_cache: &mut hipfire_runtime::llama::KvCache, - kv_layer_idx: usize, - scratch: &Gemma4Scratch, -) -> HipResult<()> { - sliding_layer_decode_impl(gpu, config, lw, pos, kv_cache, kv_layer_idx, scratch, false) -} -/// As `sliding_layer_decode` but with `stop_before_moe=true`, the function -/// returns immediately after the dense FFN computes `scratch.ffn_out`, -/// leaving `scratch.residual` holding the post-attention x. Used by -/// `forward_prefill_chunk` to interleave per-token attention with a -/// batched MoE call across all N tokens. -fn sliding_layer_attn_ffn_only( +/// Token-batched prefill. Processes up to `scratch.max_prefill_batch` tokens +/// at once, amortizing MoE-branch launch overhead across the batch via the +/// batched-indexed kernels (router GEMM, top-K, gate_up, gelu*mul, down). +/// +/// Per-token operations (attention, dense projections, RoPE, KV writes) +/// stay per-token in V1 — no batched flash-prefill kernel for asym3 +/// sliding window. The win comes from collapsing N per-token MoE calls +/// (~10 launches each) into one batched MoE call per layer. +/// +/// On exit: `kv_sliding` / `kv_full` are updated for every token. The KV +/// cache is the source of truth for downstream sampling; this function +/// does NOT compute logits (the caller should follow with a per-token +/// `forward_scratch` for the LAST token if it needs logits). +pub fn forward_prefill_batch( gpu: &mut Gpu, + weights: &Gemma4Weights, config: &Gemma4Config, - lw: &SlidingLayerWeights, - pos: usize, - kv_cache: &mut hipfire_runtime::llama::KvCache, - kv_layer_idx: usize, + tokens: &[u32], + start_pos: usize, + kv_sliding: &mut hipfire_runtime::llama::KvCache, + kv_full: &mut hipfire_runtime::llama::KvCache, scratch: &Gemma4Scratch, ) -> HipResult<()> { - sliding_layer_decode_impl(gpu, config, lw, pos, kv_cache, kv_layer_idx, scratch, true) + // v2 — batched dense projections + batched MoE. The +55% prefill win + // from 521161f8 is back: the regressing bug was in `gemm_hfq4g128`'s + // partial-trailing-group handling (used floor instead of ceil for + // groups_per_row, silently dropped 64 input dims at K=2112 dense + // down_proj → garbage). Fixed in kernels/src/gemm_hfq4g128.hip. + forward_prefill_batch_v2( + gpu, weights, config, tokens, start_pos, kv_sliding, kv_full, scratch, + ) } -fn sliding_layer_decode_impl( +/// V2 batched prefill — batches both the dense projections AND the MoE +/// branch across N tokens. Per-token RoPE and attention stay sequential +/// because the flash kernels haven't been ported to prefill-batched layout. +/// +/// Fixed 2026-05-19 by the gemm_hfq4g128 partial-trailing-group fix +/// (kernels/src/gemm_hfq4g128.hip). The v2 dispatch path is byte-identical +/// to v1 through attention; the divergence was the missing 64-element +/// partial group in the dense down_proj GEMM at K=2112. +fn forward_prefill_batch_v2( gpu: &mut Gpu, + weights: &Gemma4Weights, config: &Gemma4Config, - lw: &SlidingLayerWeights, - pos: usize, - kv_cache: &mut hipfire_runtime::llama::KvCache, - kv_layer_idx: usize, - scratch: &Gemma4Scratch, - stop_before_moe: bool, -) -> HipResult<()> { - let dim = config.dim; - let head_dim = config.sliding_head_dim; - let n_heads = config.n_heads; - let n_kv = config.sliding_n_kv_heads; - let dim_bytes = dim * 4; - - // residual = x - if let Some(s) = gpu.active_stream.as_ref() { - gpu.hip - .memcpy_dtod_async_at(&scratch.residual.buf, 0, &scratch.x.buf, 0, dim_bytes, s)?; - } else { - gpu.hip - .memcpy_dtod(&scratch.residual.buf, &scratch.x.buf, dim_bytes)?; - } - let _dump_on = hipfire_config::developer_var("HIPFIRE_GEMMA4_DUMP") - .ok() - .as_deref() - == Some("1") - && (pos == 0 || pos == 1) - && kv_layer_idx == 0; - if _dump_on { - dbg_dump(gpu, "[v1] L0 input scratch.x", &scratch.x, dim); - } - let ctx = DispatchCtx::new(gpu); - - // tmp = input_layernorm(x) - gpu.rmsnorm_f32( - &scratch.x, - &lw.input_layernorm, - &scratch.tmp, - config.norm_eps, - )?; - if _dump_on { - dbg_dump(gpu, "[v1] L0 after input_norm", &scratch.tmp, dim); - } - - // Q/K/V projections: q[n_heads*head_dim], k/v[n_kv*head_dim]. - let wr_q = lw.q_proj.dispatch_ref(); - let wr_k = lw.k_proj.dispatch_ref(); - execute_steps( - gpu, - &ctx, - &[ - Step::Gemv { - w: &wr_q, - input: GemvInput::Raw(&scratch.tmp), - out: &scratch.q, - }, - Step::Gemv { - w: &wr_k, - input: GemvInput::Raw(&scratch.tmp), - out: &scratch.k, - }, - ], - ) - .map_err(|e| hip_bridge::HipError::new(0, &e.to_string()))?; - weight_gemv(gpu, &lw.v_proj, &scratch.tmp, &scratch.v)?; - if _dump_on { - dbg_dump(gpu, "[v1] L0 after q_proj", &scratch.q, n_heads * head_dim); - dbg_dump(gpu, "[v1] L0 after k_proj", &scratch.k, n_kv * head_dim); - dbg_dump(gpu, "[v1] L0 after v_proj", &scratch.v, n_kv * head_dim); - } - - // q_norm + k_norm + no-scale v_norm across head_dim (in-place). - // v_norm matches HF Gemma 4 `value_states = v_norm(v_proj(x))` (no_scale=True - // RMSNorm — divide-only, ones buffer as weight). Same pattern as full_layer_decode. - // Omitting v_norm compounds attention-output bias across 48 sliding layers and - // produces single-token garbage end-to-end while still passing Phase 2 kernel - // NRMSE (which tests q_norm + k_norm but not v_norm — see d2.5-results.md). - gpu.rmsnorm_batched( - &scratch.q, - &lw.q_norm, - &scratch.q, - n_heads, - head_dim, - config.norm_eps, - )?; - gpu.rmsnorm_batched( - &scratch.k, - &lw.k_norm, - &scratch.k, - n_kv, - head_dim, - config.norm_eps, - )?; - gpu.rmsnorm_batched( - &scratch.v, - &scratch.v_norm_ones_full, - &scratch.v, - n_kv, - head_dim, - config.norm_eps, - )?; - if _dump_on { - dbg_dump(gpu, "[v1] L0 after q_norm", &scratch.q, n_heads * head_dim); - dbg_dump(gpu, "[v1] L0 after k_norm", &scratch.k, n_kv * head_dim); - dbg_dump(gpu, "[v1] L0 after v_norm", &scratch.v, n_kv * head_dim); - } - - // Pre-scale Q by sqrt(head_dim) so the flash-attn kernel's internal - // 1/sqrt(head_dim) cancels, leaving the effective Gemma 4 scale of 1.0. - // Only the first n_heads*head_dim elements of scratch.q are live. - gpu.scale_f32(&scratch.q, (head_dim as f32).sqrt())?; - if _dump_on { - dbg_dump(gpu, "[v1] L0 after scale_q", &scratch.q, n_heads * head_dim); - } - - // Full rotate_half RoPE, theta=10000, head_dim=256 (all dims rotate). - gpu.rope_f32( - &scratch.q, - &scratch.k, - &scratch.pos_buf, - n_heads, - n_kv, - head_dim, - config.sliding_rope_theta, - )?; - if _dump_on { - dbg_dump(gpu, "[v1] L0 after rope_q", &scratch.q, n_heads * head_dim); - dbg_dump(gpu, "[v1] L0 after rope_k", &scratch.k, n_kv * head_dim); - // Dump raw Q/K/V after RoPE for score analysis - let q_data = gpu.download_f32(&scratch.q).unwrap_or_default(); - let k_data = gpu.download_f32(&scratch.k).unwrap_or_default(); - let v_data = gpu.download_f32(&scratch.v).unwrap_or_default(); - let _ = std::fs::write("/tmp/gemma4_q_rope.bin", unsafe { - std::slice::from_raw_parts(q_data.as_ptr() as *const u8, q_data.len() * 4) - }); - let _ = std::fs::write("/tmp/gemma4_k_rope.bin", unsafe { - std::slice::from_raw_parts(k_data.as_ptr() as *const u8, k_data.len() * 4) - }); - let _ = std::fs::write("/tmp/gemma4_v.bin", unsafe { - std::slice::from_raw_parts(v_data.as_ptr() as *const u8, v_data.len() * 4) - }); - eprintln!( - "[v1] Dumped post-RoPE Q({}), K({}), V({})", - q_data.len(), - k_data.len(), - v_data.len() - ); - } - - // KV cache write + flash attention via dispatch framework (Step::Attend). - // flash_mode=2 (forced) because sliding layers always need flash for window masking. - { - let tier_inputs = KvTierInputs { - quant_asym4: kv_cache.quant_asym4, - quant_asym3: kv_cache.quant_asym3, - quant_asym2: kv_cache.quant_asym2, - quant_q8: kv_cache.quant_q8, - quant_fwht: kv_cache.quant_fwht, - quant_hfq4: false, - quant_q4: false, - quant_int8: false, - quant_hfq8: false, - quant_bf16: false, - quant_fp8: false, - f32_policy: hipfire_dispatch::families::kv_tier::F32AttnPolicy::Simple, - v_mode_bits: kv_cache.v_mode_bits(), - pos, - flash_mode: 2, // forced flash — sliding layers need window masking - capture_mode: gpu.graphs.capture_mode, - batch_size: 1, - is_tree: false, - is_boundary: false, - q8_windowed: true, - window: config.sliding_window as i32, - }; - let plan = KvTierPlan::derive(tier_inputs) - .map_err(|e| hip_bridge::HipError::new(0, &format!("{:?}", e)))?; - let io = AttnParams { - q: &scratch.q, - k: &scratch.k, - v: &scratch.v, - k_cache: &kv_cache.k_gpu[kv_layer_idx], - v_cache: &kv_cache.v_gpu[kv_layer_idx], - k_scales: None, - v_scales: None, - pos_buf: &scratch.pos_buf, - pos, - positions: None, - n_heads, - n_kv_heads: n_kv, - head_dim, - physical_cap: kv_cache.physical_cap, - batch_size: 1, - max_ctx_len: 0, - flash_partials: Some(&scratch.flash_partials), - givens_cos: kv_cache.givens_cos.as_ref(), - givens_sin: kv_cache.givens_sin.as_ref(), - tree_bias: None, - block_start: 0, - block_cols: 0, - output_gate: None, - output_awq_scale: None, - output: &scratch.attn_out, - }; - execute_steps(gpu, &ctx, &[Step::Attend { plan, io }]) - .map_err(|e| hip_bridge::HipError::new(0, &e.to_string()))?; - } - // Dump attention output regardless of KV cache branch. - if _dump_on { - dbg_dump( - gpu, - "[v1] L0 after attention", - &scratch.attn_out, - n_heads * head_dim, - ); - let data = gpu.download_f32(&scratch.attn_out).unwrap_or_default(); - let gqa_ratio = n_heads / n_kv; - for h in 0..n_heads { - let start = h * head_dim; - let end = start + head_dim; - let h_sum: f64 = data[start..end].iter().map(|&v| v as f64).sum(); - eprintln!( - "[v1] head {:2} (kv={}): sum={:+10.4} first2=[{:.4}, {:.4}]", - h, - h / gqa_ratio, - h_sum, - data[start], - data[start + 1] - ); - } - } - - // o_proj → tmp (reuse tmp, overwriting input_layernorm output). - let wr_o = lw.o_proj.dispatch_ref(); - execute_steps( - gpu, - &ctx, - &[Step::Gemv { - w: &wr_o, - input: GemvInput::Raw(&scratch.attn_out), - out: &scratch.tmp, - }], - ) - .map_err(|e| hip_bridge::HipError::new(0, &e.to_string()))?; - if _dump_on { - dbg_dump(gpu, "[v1] L0 after o_proj", &scratch.tmp, dim); - } - - // Sandwich post-attn norm (in-place on tmp). - gpu.rmsnorm_f32( - &scratch.tmp, - &lw.post_attention_layernorm, - &scratch.tmp, - config.norm_eps, - )?; - if _dump_on { - dbg_dump(gpu, "[v1] L0 after post_attn_norm", &scratch.tmp, dim); - } - - // x = residual + tmp. (Reset x first since earlier ops mutated it.) - if let Some(s) = gpu.active_stream.as_ref() { - gpu.hip - .memcpy_dtod_async_at(&scratch.x.buf, 0, &scratch.residual.buf, 0, dim_bytes, s)?; - } else { - gpu.hip - .memcpy_dtod(&scratch.x.buf, &scratch.residual.buf, dim_bytes)?; - } - gpu.add_inplace_f32(&scratch.x, &scratch.tmp)?; - if _dump_on { - dbg_dump(gpu, "[v1] L0 after attn_residual", &scratch.x, dim); - } - - // residual = x (for the FFN residual stream). - if let Some(s) = gpu.active_stream.as_ref() { - gpu.hip - .memcpy_dtod_async_at(&scratch.residual.buf, 0, &scratch.x.buf, 0, dim_bytes, s)?; - } else { - gpu.hip - .memcpy_dtod(&scratch.residual.buf, &scratch.x.buf, dim_bytes)?; - } - - // Pre-FFN norm. - gpu.rmsnorm_f32( - &scratch.x, - &lw.pre_feedforward_layernorm, - &scratch.tmp, - config.norm_eps, - )?; - if _dump_on { - dbg_dump(gpu, "[v1] L0 after pre_ffn_norm", &scratch.tmp, dim); - } - - // SwiGLU(gelu_pytorch_tanh): gate_proj, up_proj, gelu_tanh(gate) * up → down_proj. - let wr_gate = lw.gate_proj.dispatch_ref(); - let wr_up = lw.up_proj.dispatch_ref(); - execute_steps( - gpu, - &ctx, - &[ - Step::Gemv { - w: &wr_gate, - input: GemvInput::Raw(&scratch.tmp), - out: &scratch.gate_ffn, - }, - Step::Gemv { - w: &wr_up, - input: GemvInput::Raw(&scratch.tmp), - out: &scratch.up_ffn, - }, - ], - ) - .map_err(|e| hip_bridge::HipError::new(0, &e.to_string()))?; - if _dump_on { - dbg_dump( - gpu, - "[v1] L0 after gate_proj", - &scratch.gate_ffn, - config.hidden_dim, - ); - dbg_dump( - gpu, - "[v1] L0 after up_proj", - &scratch.up_ffn, - config.hidden_dim, - ); - } - gpu.gelu_tanh_f32(&scratch.gate_ffn, &scratch.ffn_hidden, config.hidden_dim)?; - gpu.mul_f32(&scratch.ffn_hidden, &scratch.up_ffn, &scratch.ffn_hidden)?; - if _dump_on { - dbg_dump( - gpu, - "[v1] L0 after gelu*up", - &scratch.ffn_hidden, - config.hidden_dim, - ); - } - let wr_down = lw.down_proj.dispatch_ref(); - execute_steps( - gpu, - &ctx, - &[Step::Gemv { - w: &wr_down, - input: GemvInput::Raw(&scratch.ffn_hidden), - out: &scratch.ffn_out, - }], - ) - .map_err(|e| hip_bridge::HipError::new(0, &e.to_string()))?; - if _dump_on { - dbg_dump(gpu, "[v1] L0 after down_proj", &scratch.ffn_out, dim); - } - - // Batched prefill hand-off point: caller wants to batch the MoE branch - // across N tokens. Return now with scratch.ffn_out + scratch.residual - // populated; caller assembles the batched MoE call externally. - if stop_before_moe { - return Ok(()); - } - - // Sandwich post-FFN norm. On MoE layers (26B-A4B) this is folded into - // apply_moe_branch (which adds the parallel MoE branch + sandwich norms - // 1 and 2 before this outer norm); on dense layers we just call the - // standalone post_feedforward_layernorm. - let moe_bypass = hipfire_config::developer_var("HIPFIRE_MOE_BYPASS") - .ok() - .as_deref() - == Some("1"); - match (lw.moe.as_ref(), moe_bypass) { - (Some(moe), false) => apply_moe_branch( - gpu, - config, - scratch, - moe, - &lw.post_feedforward_layernorm, - &scratch.residual, - )?, - _ => gpu.rmsnorm_f32( - &scratch.ffn_out, - &lw.post_feedforward_layernorm, - &scratch.tmp, - config.norm_eps, - )?, - } - - // x = residual + tmp (again, reset x from saved residual). - if let Some(s) = gpu.active_stream.as_ref() { - gpu.hip - .memcpy_dtod_async_at(&scratch.x.buf, 0, &scratch.residual.buf, 0, dim_bytes, s)?; - } else { - gpu.hip - .memcpy_dtod(&scratch.x.buf, &scratch.residual.buf, dim_bytes)?; - } - gpu.add_inplace_f32(&scratch.x, &scratch.tmp)?; - - // Learned per-layer scalar multiplier. - gpu.scale_f32(&scratch.x, lw.layer_scalar_host)?; - - Ok(()) -} - -/// Single full (global) attention layer with K=V weight sharing. -/// -/// Key differences from sliding: -/// • head_dim = 512 (global_head_dim), 4 KV heads (vs sliding's 256 / 16). -/// • V is the *pre*-k_norm output of k_proj — CRITICAL ordering (line 1214 -/// of modeling_gemma4.py). In Python: -/// key_states = k_proj(x) -/// value_states = v_proj(x) if v_proj else key_states # bound BEFORE norm -/// key_states = k_norm(key_states) # rebind, value_states holds pre-norm -/// value_states = v_norm(value_states) -/// Our translation: write k_proj output into `scratch.k`, memcpy into -/// `scratch.v`, then apply k_norm in-place on scratch.k. -/// • v_norm is `no_scale=true` RMSNorm — divide only, no learned gain. -/// We call the existing `rmsnorm_batched` with the ones-filled -/// `scratch.v_norm_ones_full` as the weight vector. -/// • RoPE is partial_rotary_factor=0.25 proportional: -/// pairs (i, i+head_dim/2) for i in [0, 64) rotate with theta=1e6; -/// pairs [64, 256) are NoPE (identity). See `rope_partial_halved_f32`. -/// • No sliding window (window_size=0 = full causal). -/// • Attention scale = 1.0 (same as sliding — Gemma 4 sets -/// `self.scaling = 1.0`; we compensate by pre-scaling Q by sqrt(head_dim)). -fn full_layer_decode( - gpu: &mut Gpu, - config: &Gemma4Config, - lw: &FullLayerWeights, - pos: usize, - kv_cache: &mut hipfire_runtime::llama::KvCache, - kv_layer_idx: usize, - scratch: &Gemma4Scratch, -) -> HipResult<()> { - full_layer_decode_impl(gpu, config, lw, pos, kv_cache, kv_layer_idx, scratch, false) -} - -fn full_layer_attn_ffn_only( - gpu: &mut Gpu, - config: &Gemma4Config, - lw: &FullLayerWeights, - pos: usize, - kv_cache: &mut hipfire_runtime::llama::KvCache, - kv_layer_idx: usize, - scratch: &Gemma4Scratch, -) -> HipResult<()> { - full_layer_decode_impl(gpu, config, lw, pos, kv_cache, kv_layer_idx, scratch, true) -} - -fn full_layer_decode_impl( - gpu: &mut Gpu, - config: &Gemma4Config, - lw: &FullLayerWeights, - pos: usize, - kv_cache: &mut hipfire_runtime::llama::KvCache, - kv_layer_idx: usize, - scratch: &Gemma4Scratch, - stop_before_moe: bool, -) -> HipResult<()> { - let dim = config.dim; - let head_dim = config.full_head_dim; - let n_heads = config.n_heads; - let n_kv = config.full_n_kv_heads; - let dim_bytes = dim * 4; - let ctx = DispatchCtx::new(gpu); - let kv_bytes = n_kv * head_dim * 4; - - // residual = x - if let Some(s) = gpu.active_stream.as_ref() { - gpu.hip - .memcpy_dtod_async_at(&scratch.residual.buf, 0, &scratch.x.buf, 0, dim_bytes, s)?; - } else { - gpu.hip - .memcpy_dtod(&scratch.residual.buf, &scratch.x.buf, dim_bytes)?; - } - - // tmp = input_layernorm(x) - gpu.rmsnorm_f32( - &scratch.x, - &lw.input_layernorm, - &scratch.tmp, - config.norm_eps, - )?; - - let _fdump = hipfire_config::developer_var("HIPFIRE_GEMMA4_DUMP") - .ok() - .as_deref() - == Some("1") - && pos == 1 - && kv_layer_idx == 0; - if _fdump { - dbg_dump(gpu, "[FL] L5 input scratch.x", &scratch.x, dim); - } - if _fdump { - dbg_dump(gpu, "[FL] L5 post_input_norm", &scratch.tmp, dim); - } - - // Q + K projections. V is derived from K's pre-norm output below. - weight_gemv(gpu, &lw.q_proj, &scratch.tmp, &scratch.q)?; - weight_gemv(gpu, &lw.k_proj, &scratch.tmp, &scratch.k)?; - if _fdump { - dbg_dump(gpu, "[FL] L5 post_q_proj", &scratch.q, n_heads * head_dim); - } - if _fdump { - dbg_dump(gpu, "[FL] L5 post_k_proj", &scratch.k, n_kv * head_dim); - } - - // CRITICAL: capture pre-k_norm bytes as V before applying k_norm. - if let Some(s) = gpu.active_stream.as_ref() { - gpu.hip - .memcpy_dtod_async_at(&scratch.v.buf, 0, &scratch.k.buf, 0, kv_bytes, s)?; - } else { - gpu.hip - .memcpy_dtod(&scratch.v.buf, &scratch.k.buf, kv_bytes)?; - } - - // q_norm, k_norm, and no-scale v_norm (all head_dim = 512). - gpu.rmsnorm_batched( - &scratch.q, - &lw.q_norm, - &scratch.q, - n_heads, - head_dim, - config.norm_eps, - )?; - gpu.rmsnorm_batched( - &scratch.k, - &lw.k_norm, - &scratch.k, - n_kv, - head_dim, - config.norm_eps, - )?; - gpu.rmsnorm_batched( - &scratch.v, - &scratch.v_norm_ones_full, - &scratch.v, - n_kv, - head_dim, - config.norm_eps, - )?; - if _fdump { - dbg_dump(gpu, "[FL] L5 post_q_norm", &scratch.q, n_heads * head_dim); - } - if _fdump { - dbg_dump(gpu, "[FL] L5 post_k_norm", &scratch.k, n_kv * head_dim); - } - if _fdump { - dbg_dump(gpu, "[FL] L5 post_v_norm", &scratch.v, n_kv * head_dim); - } - - // Pre-scale Q by sqrt(head_dim=512) so the flash kernel's 1/sqrt(head_dim) - // cancels (Gemma 4 attention scaling is 1.0). - gpu.scale_f32(&scratch.q, (head_dim as f32).sqrt())?; - - // Proportional RoPE: rotate_half of the first 64 pairs of every 512-dim head. - let n_rot_pairs = ((head_dim as f32) * config.full_partial_rotary_factor * 0.5) as usize; - gpu.rope_partial_halved_f32( - &scratch.q, - &scratch.k, - &scratch.pos_buf, - n_heads, - n_kv, - head_dim, - n_rot_pairs, - config.full_rope_theta, - )?; - if _fdump { - dbg_dump(gpu, "[FL] L5 post_rope_q", &scratch.q, n_heads * head_dim); - } - if _fdump { - dbg_dump(gpu, "[FL] L5 post_rope_k", &scratch.k, n_kv * head_dim); - } - - // KV cache write + attention via dispatch framework (Step::Attend). - // Full-attention layers: window_size=0 (full causal), cache_capacity=0. - { - let tier_inputs = KvTierInputs { - quant_asym4: kv_cache.quant_asym4, - quant_asym3: kv_cache.quant_asym3, - quant_asym2: kv_cache.quant_asym2, - quant_q8: kv_cache.quant_q8, - quant_fwht: kv_cache.quant_fwht, - quant_hfq4: false, - quant_q4: false, - quant_int8: false, - quant_hfq8: false, - quant_bf16: false, - quant_fp8: false, - f32_policy: hipfire_dispatch::families::kv_tier::F32AttnPolicy::Simple, - v_mode_bits: kv_cache.v_mode_bits(), - pos, - flash_mode: 2, - capture_mode: gpu.graphs.capture_mode, - batch_size: 1, - is_tree: false, - is_boundary: false, - q8_windowed: false, - window: 0, - }; - let plan = KvTierPlan::derive(tier_inputs) - .map_err(|e| hip_bridge::HipError::new(0, &format!("{:?}", e)))?; - let io = AttnParams { - q: &scratch.q, - k: &scratch.k, - v: &scratch.v, - k_cache: &kv_cache.k_gpu[kv_layer_idx], - v_cache: &kv_cache.v_gpu[kv_layer_idx], - k_scales: None, - v_scales: None, - pos_buf: &scratch.pos_buf, - pos, - positions: None, - n_heads, - n_kv_heads: n_kv, - head_dim, - physical_cap: kv_cache.max_seq, - batch_size: 1, - max_ctx_len: 0, - flash_partials: Some(&scratch.flash_partials), - givens_cos: kv_cache.givens_cos.as_ref(), - givens_sin: kv_cache.givens_sin.as_ref(), - tree_bias: None, - block_start: 0, - block_cols: 0, - output_gate: None, - output_awq_scale: None, - output: &scratch.attn_out, - }; - execute_steps(gpu, &ctx, &[Step::Attend { plan, io }]) - .map_err(|e| hip_bridge::HipError::new(0, &e.to_string()))?; - } - - // o_proj → tmp. - weight_gemv(gpu, &lw.o_proj, &scratch.attn_out, &scratch.tmp)?; - if _fdump { - dbg_dump( - gpu, - "[FL] L5 post_attn_out", - &scratch.attn_out, - n_heads * head_dim, - ); - } - if _fdump { - dbg_dump(gpu, "[FL] L5 post_o_proj", &scratch.tmp, dim); - } - - // Sandwich post-attn norm. - gpu.rmsnorm_f32( - &scratch.tmp, - &lw.post_attention_layernorm, - &scratch.tmp, - config.norm_eps, - )?; - - // x = residual + tmp. - if let Some(s) = gpu.active_stream.as_ref() { - gpu.hip - .memcpy_dtod_async_at(&scratch.x.buf, 0, &scratch.residual.buf, 0, dim_bytes, s)?; - } else { - gpu.hip - .memcpy_dtod(&scratch.x.buf, &scratch.residual.buf, dim_bytes)?; - } - gpu.add_inplace_f32(&scratch.x, &scratch.tmp)?; - - // Save new residual. - if let Some(s) = gpu.active_stream.as_ref() { - gpu.hip - .memcpy_dtod_async_at(&scratch.residual.buf, 0, &scratch.x.buf, 0, dim_bytes, s)?; - } else { - gpu.hip - .memcpy_dtod(&scratch.residual.buf, &scratch.x.buf, dim_bytes)?; - } - - // Pre-FFN norm. - gpu.rmsnorm_f32( - &scratch.x, - &lw.pre_feedforward_layernorm, - &scratch.tmp, - config.norm_eps, - )?; - - // SwiGLU with gelu_pytorch_tanh activation. - weight_gemv(gpu, &lw.gate_proj, &scratch.tmp, &scratch.gate_ffn)?; - weight_gemv(gpu, &lw.up_proj, &scratch.tmp, &scratch.up_ffn)?; - gpu.gelu_tanh_f32(&scratch.gate_ffn, &scratch.ffn_hidden, config.hidden_dim)?; - gpu.mul_f32(&scratch.ffn_hidden, &scratch.up_ffn, &scratch.ffn_hidden)?; - weight_gemv(gpu, &lw.down_proj, &scratch.ffn_hidden, &scratch.ffn_out)?; - - // Batched prefill hand-off (see sliding_layer_decode_impl for rationale). - if stop_before_moe { - return Ok(()); - } - - // Sandwich post-FFN norm. Same MoE dispatch as sliding_layer_decode. - let moe_bypass = hipfire_config::developer_var("HIPFIRE_MOE_BYPASS") - .ok() - .as_deref() - == Some("1"); - match (lw.moe.as_ref(), moe_bypass) { - (Some(moe), false) => apply_moe_branch( - gpu, - config, - scratch, - moe, - &lw.post_feedforward_layernorm, - &scratch.residual, - )?, - _ => gpu.rmsnorm_f32( - &scratch.ffn_out, - &lw.post_feedforward_layernorm, - &scratch.tmp, - config.norm_eps, - )?, - } - - // x = residual + tmp. - if let Some(s) = gpu.active_stream.as_ref() { - gpu.hip - .memcpy_dtod_async_at(&scratch.x.buf, 0, &scratch.residual.buf, 0, dim_bytes, s)?; - } else { - gpu.hip - .memcpy_dtod(&scratch.x.buf, &scratch.residual.buf, dim_bytes)?; - } - gpu.add_inplace_f32(&scratch.x, &scratch.tmp)?; - - // Learned per-layer scalar multiplier. - gpu.scale_f32(&scratch.x, lw.layer_scalar_host)?; - - Ok(()) -} - -/// Token-batched prefill. Processes up to `scratch.max_prefill_batch` tokens -/// at once, amortizing MoE-branch launch overhead across the batch via the -/// batched-indexed kernels (router GEMM, top-K, gate_up, gelu*mul, down). -/// -/// Per-token operations (attention, dense projections, RoPE, KV writes) -/// stay per-token in V1 — no batched flash-prefill kernel for asym3 -/// sliding window. The win comes from collapsing N per-token MoE calls -/// (~10 launches each) into one batched MoE call per layer. -/// -/// On exit: `kv_sliding` / `kv_full` are updated for every token. The KV -/// cache is the source of truth for downstream sampling; this function -/// does NOT compute logits (the caller should follow with a per-token -/// `forward_scratch` for the LAST token if it needs logits). -pub fn forward_prefill_batch( - gpu: &mut Gpu, - weights: &Gemma4Weights, - config: &Gemma4Config, - tokens: &[u32], - start_pos: usize, - kv_sliding: &mut hipfire_runtime::llama::KvCache, - kv_full: &mut hipfire_runtime::llama::KvCache, - scratch: &Gemma4Scratch, -) -> HipResult<()> { - // v2 — batched dense projections + batched MoE. The +55% prefill win - // from 521161f8 is back: the regressing bug was in `gemm_hfq4g128`'s - // partial-trailing-group handling (used floor instead of ceil for - // groups_per_row, silently dropped 64 input dims at K=2112 dense - // down_proj → garbage). Fixed in kernels/src/gemm_hfq4g128.hip. - forward_prefill_batch_v2( - gpu, weights, config, tokens, start_pos, kv_sliding, kv_full, scratch, - ) -} - -fn forward_prefill_batch_v1( - gpu: &mut Gpu, - weights: &Gemma4Weights, - config: &Gemma4Config, - tokens: &[u32], - start_pos: usize, - kv_sliding: &mut hipfire_runtime::llama::KvCache, - kv_full: &mut hipfire_runtime::llama::KvCache, - scratch: &Gemma4Scratch, -) -> HipResult<()> { - let n_batch = tokens.len(); - if n_batch == 0 { - return Ok(()); - } - if n_batch > scratch.max_prefill_batch { - return Err(hip_bridge::HipError::new( - 0, - &format!( - "forward_prefill_batch: n_batch={n_batch} > max_prefill_batch={}; \ - caller must chunk", - scratch.max_prefill_batch - ), - )); - } - let dim = config.dim; - let dim_bytes = dim * 4; - - // ── Step 1: per-token embed + scale into pb_residual[i]. ───────────── - for (i, &tok) in tokens.iter().enumerate() { - match weights.embd_format { - EmbeddingFormat::HFQ4G256 => { - gpu.embedding_lookup_hfq4g256(&weights.embed_tokens, &scratch.x, tok, dim)? - } - EmbeddingFormat::HFQ4G128 => { - gpu.embedding_lookup_hfq4g128(&weights.embed_tokens, &scratch.x, tok, dim)? - } - EmbeddingFormat::Q8_0 => { - gpu.embedding_lookup_q8(&weights.embed_tokens, &scratch.x, tok, dim)? - } - EmbeddingFormat::F32 => { - gpu.embedding_lookup(&weights.embed_tokens, &scratch.x, tok, dim)? - } - _ => { - return Err(hip_bridge::HipError::new( - 0, - "unsupported Gemma 4 embed format", - )); - } - } - gpu.scale_f32(&scratch.x, config.embed_scale)?; - // pb_residual[i] = scratch.x - let stream = gpu.active_stream.as_ref(); - if let Some(s) = stream { - gpu.hip.memcpy_dtod_async_at( - &scratch.pb_residual.buf, - i * dim_bytes, - &scratch.x.buf, - 0, - dim_bytes, - s, - )?; - } else { - gpu.hip.memcpy_dtod_at( - &scratch.pb_residual.buf, - i * dim_bytes, - &scratch.x.buf, - 0, - dim_bytes, - )?; - } - } - - // ── Step 2: layer loop. Each layer: per-token attn+FFN, then batched MoE. ── - let mut sliding_kv_idx = 0usize; - let mut full_kv_idx = 0usize; - for (layer_idx, layer_type) in config.layer_types.iter().copied().enumerate() { - // Stage A — per-token attn + FFN. Fills pb_attn_out[i] = post-attn x, - // pb_ffn_out[i] = dense FFN out. Also writes the per-token KV at - // position start_pos+i. - let stream_present = gpu.active_stream.is_some(); - for i in 0..n_batch { - let pos = start_pos + i; - let pos_i32 = pos as i32; - gpu.hip - .memcpy_htod(&scratch.pos_buf, &pos_i32.to_ne_bytes())?; - // Load this token's residual into scratch.x. - if stream_present { - let s = gpu.active_stream.as_ref().unwrap(); - gpu.hip.memcpy_dtod_async_at( - &scratch.x.buf, - 0, - &scratch.pb_residual.buf, - i * dim_bytes, - dim_bytes, - s, - )?; - } else { - gpu.hip.memcpy_dtod_at( - &scratch.x.buf, - 0, - &scratch.pb_residual.buf, - i * dim_bytes, - dim_bytes, - )?; - } - match (layer_type, &weights.layers[layer_idx]) { - (LayerType::Sliding, LayerWeights::Sliding(lw)) => { - sliding_layer_attn_ffn_only( - gpu, - config, - lw, - pos, - kv_sliding, - sliding_kv_idx, - scratch, - )?; - } - (LayerType::Full, LayerWeights::Full(lw)) => { - full_layer_attn_ffn_only(gpu, config, lw, pos, kv_full, full_kv_idx, scratch)?; - } - _ => { - return Err(hip_bridge::HipError::new( - 0, - &format!("Gemma 4 layer {} type/weights mismatch", layer_idx), - )); - } - } - // Copy outputs into batch slots: - // pb_attn_out[i] = scratch.residual (post-attn x; input to MoE pre2/router_in) - // pb_ffn_out [i] = scratch.ffn_out (dense FFN output) - // pb_residual[i] = scratch.residual (running residual) - if stream_present { - let s = gpu.active_stream.as_ref().unwrap(); - gpu.hip.memcpy_dtod_async_at( - &scratch.pb_attn_out.buf, - i * dim_bytes, - &scratch.residual.buf, - 0, - dim_bytes, - s, - )?; - gpu.hip.memcpy_dtod_async_at( - &scratch.pb_ffn_out.buf, - i * dim_bytes, - &scratch.ffn_out.buf, - 0, - dim_bytes, - s, - )?; - gpu.hip.memcpy_dtod_async_at( - &scratch.pb_residual.buf, - i * dim_bytes, - &scratch.residual.buf, - 0, - dim_bytes, - s, - )?; - } else { - gpu.hip.memcpy_dtod_at( - &scratch.pb_attn_out.buf, - i * dim_bytes, - &scratch.residual.buf, - 0, - dim_bytes, - )?; - gpu.hip.memcpy_dtod_at( - &scratch.pb_ffn_out.buf, - i * dim_bytes, - &scratch.ffn_out.buf, - 0, - dim_bytes, - )?; - gpu.hip.memcpy_dtod_at( - &scratch.pb_residual.buf, - i * dim_bytes, - &scratch.residual.buf, - 0, - dim_bytes, - )?; - } - } - match layer_type { - LayerType::Sliding => sliding_kv_idx += 1, - LayerType::Full => full_kv_idx += 1, - } - - // Stage B — batched MoE branch (or dense post-FFN on layers without MoE). - let layer_scalar = match (layer_type, &weights.layers[layer_idx]) { - (LayerType::Sliding, LayerWeights::Sliding(lw)) => lw.layer_scalar_host, - (LayerType::Full, LayerWeights::Full(lw)) => lw.layer_scalar_host, - _ => unreachable!(), - }; - let (moe_opt, post_ffn_norm) = match (layer_type, &weights.layers[layer_idx]) { - (LayerType::Sliding, LayerWeights::Sliding(lw)) => { - (lw.moe.as_ref(), &lw.post_feedforward_layernorm) - } - (LayerType::Full, LayerWeights::Full(lw)) => { - (lw.moe.as_ref(), &lw.post_feedforward_layernorm) - } - _ => unreachable!(), - }; - match moe_opt { - Some(moe) => { - apply_moe_branch_batched(gpu, config, scratch, moe, post_ffn_norm, n_batch)?; - // Stage C — batched residual add + layer scale. - // pb_residual[0..N×dim] += pb_moe_pre2[0..N×dim] - gpu.add_inplace_f32(&scratch.pb_residual, &scratch.pb_moe_pre2)?; - gpu.scale_f32(&scratch.pb_residual, layer_scalar)?; - } - None => { - // Dense path — fall back to per-token finalization to keep - // semantics identical to single-token forward. - for i in 0..n_batch { - if stream_present { - let s = gpu.active_stream.as_ref().unwrap(); - gpu.hip.memcpy_dtod_async_at( - &scratch.x.buf, - 0, - &scratch.pb_residual.buf, - i * dim_bytes, - dim_bytes, - s, - )?; - gpu.hip.memcpy_dtod_async_at( - &scratch.residual.buf, - 0, - &scratch.pb_residual.buf, - i * dim_bytes, - dim_bytes, - s, - )?; - gpu.hip.memcpy_dtod_async_at( - &scratch.ffn_out.buf, - 0, - &scratch.pb_ffn_out.buf, - i * dim_bytes, - dim_bytes, - s, - )?; - } - gpu.rmsnorm_f32( - &scratch.ffn_out, - post_ffn_norm, - &scratch.tmp, - config.norm_eps, - )?; - if stream_present { - let s = gpu.active_stream.as_ref().unwrap(); - gpu.hip.memcpy_dtod_async_at( - &scratch.x.buf, - 0, - &scratch.residual.buf, - 0, - dim_bytes, - s, - )?; - } - gpu.add_inplace_f32(&scratch.x, &scratch.tmp)?; - gpu.scale_f32(&scratch.x, layer_scalar)?; - if stream_present { - let s = gpu.active_stream.as_ref().unwrap(); - gpu.hip.memcpy_dtod_async_at( - &scratch.pb_residual.buf, - i * dim_bytes, - &scratch.x.buf, - 0, - dim_bytes, - s, - )?; - } - } - } - } - } - Ok(()) -} - -/// V2 batched prefill — batches both the dense projections AND the MoE -/// branch across N tokens. Per-token RoPE and attention stay sequential -/// because the flash kernels haven't been ported to prefill-batched layout. -/// -/// Fixed 2026-05-19 by the gemm_hfq4g128 partial-trailing-group fix -/// (kernels/src/gemm_hfq4g128.hip). The v2 dispatch path is byte-identical -/// to v1 through attention; the divergence was the missing 64-element -/// partial group in the dense down_proj GEMM at K=2112. -fn forward_prefill_batch_v2( - gpu: &mut Gpu, - weights: &Gemma4Weights, - config: &Gemma4Config, - tokens: &[u32], - start_pos: usize, - kv_sliding: &mut hipfire_runtime::llama::KvCache, - kv_full: &mut hipfire_runtime::llama::KvCache, + tokens: &[u32], + start_pos: usize, + kv_sliding: &mut hipfire_runtime::llama::KvCache, + kv_full: &mut hipfire_runtime::llama::KvCache, scratch: &Gemma4Scratch, ) -> HipResult<()> { let n_batch = tokens.len(); @@ -6667,1004 +5268,131 @@ fn forward_prefill_batch_v2( n_batch, Some(&scratch.pb_bf16), )?; - } else { - run_prefill_gemm( - gpu, - &lw.gate_proj, - &scratch.pb_tmp, - &scratch.pb_gate, - n_batch, - Some(&scratch.pb_bf16), - )?; - run_prefill_gemm( - gpu, - &lw.up_proj, - &scratch.pb_tmp, - &scratch.pb_up, - n_batch, - Some(&scratch.pb_bf16), - )?; - } - gpu.gelu_tanh_f32( - &scratch.pb_gate, - &scratch.pb_ffn_hidden, - n_batch * config.hidden_dim, - )?; - gpu.mul_f32( - &scratch.pb_ffn_hidden, - &scratch.pb_up, - &scratch.pb_ffn_hidden, - )?; - run_prefill_gemm( - gpu, - &lw.down_proj, - &scratch.pb_ffn_hidden, - &scratch.pb_ffn_out, - n_batch, - Some(&scratch.pb_bf16), - )?; - } - _ => unreachable!(), - } - - // MoE branch (or dense fallback). HIPFIRE_MOE_BYPASS=1 forces dense - // path even on MoE layers (parity with v1 — used to isolate whether - // a regression lives in apply_moe_branch_batched vs the dense path). - let moe_bypass = hipfire_config::developer_var("HIPFIRE_MOE_BYPASS") - .ok() - .as_deref() - == Some("1"); - match (moe_opt, moe_bypass) { - (Some(moe), false) => { - apply_moe_branch_batched(gpu, config, scratch, moe, post_ffn_norm_ref, n_batch)?; - gpu.add_inplace_f32(&scratch.pb_residual, &scratch.pb_moe_pre2)?; - gpu.scale_f32(&scratch.pb_residual, layer_scalar)?; - } - _ => { - gpu.rmsnorm_batched( - &scratch.pb_ffn_out, - post_ffn_norm_ref, - &scratch.pb_ffn_out, - n_batch, - dim, - config.norm_eps, - )?; - gpu.add_inplace_f32(&scratch.pb_residual, &scratch.pb_ffn_out)?; - gpu.scale_f32(&scratch.pb_residual, layer_scalar)?; - } - } - if hipfire_config::developer_var("HIPFIRE_GEMMA4_DUMP") - .ok() - .as_deref() - == Some("1") - { - let data = gpu.download_f32(&scratch.pb_residual).unwrap_or_default(); - let last = &data[(n_batch - 1) * dim..n_batch * dim]; - let sum: f64 = last.iter().map(|&v| v as f64).sum(); - let min = last.iter().fold(f32::INFINITY, |a, &b| a.min(b)); - let max = last.iter().fold(f32::NEG_INFINITY, |a, &b| a.max(b)); - eprintln!( - "[v2 diag] L{layer_idx} {:?} last-tok hidden: first4={:?} sum={sum:.4e} min={min:.4} max={max:.4}", - layer_type, - &last[..4.min(last.len())] - ); - } - } - - // ── Final logits: extract last token, run final_norm + lm_head + softcap ── - // The decode loop samples from scratch.logits, which must be filled - // with the last token's logits. Copy last position from pb_residual - // to scratch.tmp, run final_norm, lm_head, and softcap. - let last_offset = (n_batch - 1) * dim * 4; - if let Some(_s) = gpu.active_stream.as_ref() { - gpu.hip.memcpy_dtod_async_at( - &scratch.tmp.buf, - 0, - &scratch.pb_residual.buf, - last_offset, - dim * 4, - _s, - )?; - } else { - gpu.hip.memcpy_dtod_at( - &scratch.tmp.buf, - 0, - &scratch.pb_residual.buf, - last_offset, - dim * 4, - )?; - } - gpu.rmsnorm_f32( - &scratch.tmp, - &weights.final_norm, - &scratch.tmp, - config.norm_eps, - )?; - let ctx = DispatchCtx::new(gpu); - let wr_lm = weights.lm_head.dispatch_ref(); - execute_steps( - gpu, - &ctx, - &[Step::Gemv { - w: &wr_lm, - input: GemvInput::Raw(&scratch.tmp), - out: &scratch.logits, - }], - ) - .map_err(|e| hip_bridge::HipError::new(0, &e.to_string()))?; - if config.final_logit_softcapping > 0.0 { - gpu.logit_softcap_f32( - &scratch.logits, - config.vocab_size, - config.final_logit_softcapping, - )?; - } - - Ok(()) -} - -// ── Forward-as-pipeline (#397 Ship 6) — lowered super-op substrate ────────── -// -// Migrates Gemma 4's decode forward from per-token execute_steps resolution -// to pre-resolved LayerPrograms executed via run_layer_program + ForwardBindings. -// Behind HIPFIRE_FORWARD_LOWERED gate (default ON since byte-parity validated -// 2026-06-08; set HIPFIRE_FORWARD_LOWERED=0 to force the legacy hand path). -// See docs/plans/gemma4_forward_as_pipeline.md for the full plan. - -use hipfire_dispatch::pipeline::superop::{ - self, EscapeKind, ForwardBindings, LayerProgram, OpBinding, SuperOp, SuperOpKind, WeightSlot, -}; - -// ── Variant enum ────────────────────────────────────────────────────────── - -/// The four Gemma 4 decoder-layer shapes. Derived from the layer type -/// (Sliding vs Full) and the presence of MoE extras. Pure → unit-testable. -#[derive(Clone, Copy, Debug, PartialEq, Eq)] -enum Gemma4Variant { - SlidingDense, - SlidingMoe, - FullDense, - FullMoe, -} - -fn gemma4_variant_of(layer_type: LayerType, layer: &LayerWeights) -> Gemma4Variant { - let has_moe = match layer { - LayerWeights::Sliding(lw) => lw.moe.is_some(), - LayerWeights::Full(lw) => lw.moe.is_some(), - }; - match (layer_type, has_moe) { - (LayerType::Sliding, false) => Gemma4Variant::SlidingDense, - (LayerType::Sliding, true) => Gemma4Variant::SlidingMoe, - (LayerType::Full, false) => Gemma4Variant::FullDense, - (LayerType::Full, true) => Gemma4Variant::FullMoe, - } -} - -// ── Opcodes ─────────────────────────────────────────────────────────────── - -/// Arch-local opcodes encoded into `OpBinding.weights[0]` as `WeightSlot(code)`. -/// `SuperOpKind` routes to the `ForwardBindings` method; the opcode disambiguates -/// *which* op of that kind within the layer. -mod g4_op { - // Proj (projection clusters) - pub const PROJ_QK_SLIDING: u32 = 0; // sliding: q_proj + k_proj fused (execute_steps) - pub const PROJ_V_SLIDING: u32 = 1; // sliding: v_proj standalone (weight_gemv) - pub const PROJ_Q_FULL: u32 = 2; // full: q_proj standalone (weight_gemv) - pub const PROJ_K_FULL: u32 = 3; // full: k_proj standalone; V←K copied in Attend(FULL) - pub const PROJ_O: u32 = 4; // o_proj - pub const PROJ_GATE_UP: u32 = 5; // gate_proj + up_proj fused - pub const PROJ_DOWN: u32 = 6; // down_proj (folds gelu_tanh + mul activation) - - // Attend - pub const ATTEND_SLIDING: u32 = 0; // window=1024, head_dim=256, full rope - pub const ATTEND_FULL: u32 = 1; // window=0, head_dim=512, partial rope, k_eq_v - - // Norm - pub const NORM_INPUT: u32 = 0; // input_layernorm - pub const NORM_POST_ATTN: u32 = 1; // post_attention_layernorm - pub const NORM_PRE_FFN: u32 = 2; // pre_feedforward_layernorm - pub const NORM_POST_FFN: u32 = 3; // post_feedforward_layernorm (dense only) - - // ResidualGemv - pub const RESID_POST_ATTN: u32 = 0; // residual plumbing only (no norm) - pub const RESID_POST_FFN: u32 = 1; // residual + scale(layer_scalar) - - // Moe - pub const MOE_BRANCH: u32 = 0; // encapsulated MoE branch (replaces NORM_POST_FFN) -} - -// ── Lowering ────────────────────────────────────────────────────────────── - -fn g4_superop(kind: SuperOpKind, code: u32) -> SuperOp { - SuperOp { - kind, - binding: OpBinding { - key: None, - weights: vec![WeightSlot(code)], - scratch: Vec::new(), - flavor: superop::OpFlavor::None, - }, - } -} - -fn lower_variant(v: Gemma4Variant) -> LayerProgram { - use g4_op::*; - use SuperOpKind::*; - match v { - Gemma4Variant::SlidingDense => vec![ - g4_superop(Norm, NORM_INPUT), - g4_superop(Proj, PROJ_QK_SLIDING), - g4_superop(Proj, PROJ_V_SLIDING), - g4_superop(Attend, ATTEND_SLIDING), - g4_superop(Proj, PROJ_O), - g4_superop(Norm, NORM_POST_ATTN), - g4_superop(ResidualGemv, RESID_POST_ATTN), - g4_superop(Norm, NORM_PRE_FFN), - g4_superop(Proj, PROJ_GATE_UP), - g4_superop(Proj, PROJ_DOWN), // folds gelu_tanh + mul - g4_superop(Norm, NORM_POST_FFN), // dense only - g4_superop(ResidualGemv, RESID_POST_FFN), - ], - Gemma4Variant::SlidingMoe => vec![ - g4_superop(Norm, NORM_INPUT), - g4_superop(Proj, PROJ_QK_SLIDING), - g4_superop(Proj, PROJ_V_SLIDING), - g4_superop(Attend, ATTEND_SLIDING), - g4_superop(Proj, PROJ_O), - g4_superop(Norm, NORM_POST_ATTN), - g4_superop(ResidualGemv, RESID_POST_ATTN), - g4_superop(Norm, NORM_PRE_FFN), - g4_superop(Proj, PROJ_GATE_UP), - g4_superop(Proj, PROJ_DOWN), - // MoE replaces standalone Norm(POST_FFN) — do NOT emit both - g4_superop(Moe, MOE_BRANCH), - g4_superop(ResidualGemv, RESID_POST_FFN), // identical to dense path - ], - Gemma4Variant::FullDense => vec![ - g4_superop(Norm, NORM_INPUT), - g4_superop(Proj, PROJ_Q_FULL), - g4_superop(Proj, PROJ_K_FULL), // no v_proj; V←K copied in Attend(FULL) - g4_superop(Attend, ATTEND_FULL), - g4_superop(Proj, PROJ_O), - g4_superop(Norm, NORM_POST_ATTN), - g4_superop(ResidualGemv, RESID_POST_ATTN), - g4_superop(Norm, NORM_PRE_FFN), - g4_superop(Proj, PROJ_GATE_UP), - g4_superop(Proj, PROJ_DOWN), - g4_superop(Norm, NORM_POST_FFN), - g4_superop(ResidualGemv, RESID_POST_FFN), - ], - Gemma4Variant::FullMoe => vec![ - g4_superop(Norm, NORM_INPUT), - g4_superop(Proj, PROJ_Q_FULL), - g4_superop(Proj, PROJ_K_FULL), - g4_superop(Attend, ATTEND_FULL), - g4_superop(Proj, PROJ_O), - g4_superop(Norm, NORM_POST_ATTN), - g4_superop(ResidualGemv, RESID_POST_ATTN), - g4_superop(Norm, NORM_PRE_FFN), - g4_superop(Proj, PROJ_GATE_UP), - g4_superop(Proj, PROJ_DOWN), - g4_superop(Moe, MOE_BRANCH), - g4_superop(ResidualGemv, RESID_POST_FFN), - ], - } -} - -// ── Bindings struct ─────────────────────────────────────────────────────── - -fn op_code(op: &OpBinding) -> u32 { - op.weights.first().map(|w| w.0).unwrap_or(u32::MAX) -} - -/// Per-layer execution context for the lowered decode path. Holds the current -/// layer's weights + shared scratch by reference; rebuilt each layer iteration -/// so the borrows stay scoped. -struct Gemma4Bindings<'a> { - layer: &'a LayerWeights, - config: &'a Gemma4Config, - scratch: &'a Gemma4Scratch, - kv_sliding: &'a mut llama::KvCache, - kv_full: &'a mut llama::KvCache, - pos: usize, - sliding_kv_idx: usize, - full_kv_idx: usize, -} - -impl<'a> ForwardBindings for Gemma4Bindings<'a> { - fn run_proj( - &mut self, - gpu: &mut Gpu, - ctx: &DispatchCtx, - op: &OpBinding, - ) -> Result<(), hipfire_dispatch::types::DispatchError> { - let s = self.scratch; - let config = self.config; - let res: HipResult<()> = match op_code(op) { - g4_op::PROJ_QK_SLIDING => match self.layer { - LayerWeights::Sliding(lw) => { - let wr_q = lw.q_proj.dispatch_ref(); - let wr_k = lw.k_proj.dispatch_ref(); - execute_steps( - gpu, - ctx, - &[ - Step::Gemv { - w: &wr_q, - input: GemvInput::Raw(&s.tmp), - out: &s.q, - }, - Step::Gemv { - w: &wr_k, - input: GemvInput::Raw(&s.tmp), - out: &s.k, - }, - ], - ) - .map_err(|e| hip_bridge::HipError::new(0, &e.to_string())) - } - _ => Err(hip_bridge::HipError::new( - 0, - "PROJ_QK_SLIDING on non-Sliding layer", - )), - }, - g4_op::PROJ_V_SLIDING => match self.layer { - LayerWeights::Sliding(lw) => { - let wr = lw.v_proj.dispatch_ref(); - execute_steps( - gpu, - ctx, - &[Step::Gemv { - w: &wr, - input: GemvInput::Raw(&s.tmp), - out: &s.v, - }], - ) - .map_err(|e| hip_bridge::HipError::new(0, &e.to_string())) - } - _ => Err(hip_bridge::HipError::new( - 0, - "PROJ_V_SLIDING on non-Sliding layer", - )), - }, - g4_op::PROJ_Q_FULL => match self.layer { - LayerWeights::Full(lw) => { - let wr = lw.q_proj.dispatch_ref(); - execute_steps( - gpu, - ctx, - &[Step::Gemv { - w: &wr, - input: GemvInput::Raw(&s.tmp), - out: &s.q, - }], - ) - .map_err(|e| hip_bridge::HipError::new(0, &e.to_string())) - } - _ => Err(hip_bridge::HipError::new( - 0, - "PROJ_Q_FULL on non-Full layer", - )), - }, - g4_op::PROJ_K_FULL => match self.layer { - LayerWeights::Full(lw) => { - let wr = lw.k_proj.dispatch_ref(); - execute_steps( - gpu, - ctx, - &[Step::Gemv { - w: &wr, - input: GemvInput::Raw(&s.tmp), - out: &s.k, - }], - ) - .map_err(|e| hip_bridge::HipError::new(0, &e.to_string())) - } - _ => Err(hip_bridge::HipError::new( - 0, - "PROJ_K_FULL on non-Full layer", - )), - }, - g4_op::PROJ_O => match self.layer { - LayerWeights::Sliding(lw) => { - let wr = lw.o_proj.dispatch_ref(); - execute_steps( - gpu, - ctx, - &[Step::Gemv { - w: &wr, - input: GemvInput::Raw(&s.attn_out), - out: &s.tmp, - }], - ) - .map_err(|e| hip_bridge::HipError::new(0, &e.to_string())) - } - LayerWeights::Full(lw) => { - let wr = lw.o_proj.dispatch_ref(); - execute_steps( - gpu, - ctx, - &[Step::Gemv { - w: &wr, - input: GemvInput::Raw(&s.attn_out), - out: &s.tmp, - }], - ) - .map_err(|e| hip_bridge::HipError::new(0, &e.to_string())) - } - }, - g4_op::PROJ_GATE_UP => match self.layer { - LayerWeights::Sliding(lw) => { - let wr_gate = lw.gate_proj.dispatch_ref(); - let wr_up = lw.up_proj.dispatch_ref(); - execute_steps( - gpu, - ctx, - &[ - Step::Gemv { - w: &wr_gate, - input: GemvInput::Raw(&s.tmp), - out: &s.gate_ffn, - }, - Step::Gemv { - w: &wr_up, - input: GemvInput::Raw(&s.tmp), - out: &s.up_ffn, - }, - ], - ) - .map_err(|e| hip_bridge::HipError::new(0, &e.to_string())) - } - LayerWeights::Full(lw) => { - let wr_gate = lw.gate_proj.dispatch_ref(); - let wr_up = lw.up_proj.dispatch_ref(); - execute_steps( - gpu, - ctx, - &[ - Step::Gemv { - w: &wr_gate, - input: GemvInput::Raw(&s.tmp), - out: &s.gate_ffn, - }, - Step::Gemv { - w: &wr_up, - input: GemvInput::Raw(&s.tmp), - out: &s.up_ffn, - }, - ], - ) - .map_err(|e| hip_bridge::HipError::new(0, &e.to_string())) - } - }, - g4_op::PROJ_DOWN => { - // Folded activation: gelu_tanh(gate) * up, then down_proj. - let hidden_dim = config.hidden_dim; - gpu.gelu_tanh_f32(&s.gate_ffn, &s.ffn_hidden, hidden_dim) - .map_err(|e| hipfire_dispatch::types::DispatchError::Hip(e.to_string()))?; - gpu.mul_f32(&s.ffn_hidden, &s.up_ffn, &s.ffn_hidden) - .map_err(|e| hipfire_dispatch::types::DispatchError::Hip(e.to_string()))?; - match self.layer { - LayerWeights::Sliding(lw) => { - let wr = lw.down_proj.dispatch_ref(); - execute_steps( - gpu, - ctx, - &[Step::Gemv { - w: &wr, - input: GemvInput::Raw(&s.ffn_hidden), - out: &s.ffn_out, - }], - ) - .map_err(|e| hip_bridge::HipError::new(0, &e.to_string())) - } - LayerWeights::Full(lw) => { - let wr = lw.down_proj.dispatch_ref(); - execute_steps( - gpu, - ctx, - &[Step::Gemv { - w: &wr, - input: GemvInput::Raw(&s.ffn_hidden), - out: &s.ffn_out, - }], - ) - .map_err(|e| hip_bridge::HipError::new(0, &e.to_string())) - } - } - } - other => Err(hip_bridge::HipError::new( - 0, - &format!("unknown PROJ opcode {other}"), - )), - }; - res.map_err(|e| hipfire_dispatch::types::DispatchError::Hip(e.to_string())) - } - fn run_residual_gemv( - &mut self, - gpu: &mut Gpu, - _ctx: &DispatchCtx, - op: &OpBinding, - ) -> Result<(), hipfire_dispatch::types::DispatchError> { - let s = self.scratch; - let dim = self.config.dim; - let hip_to_dispatch = - |e: hip_bridge::HipError| hipfire_dispatch::types::DispatchError::Hip(e.to_string()); - - let res: HipResult<()> = match op_code(op) { - g4_op::RESID_POST_ATTN => { - // First: save x → residual (this is the first time residual is set - // in this layer — Norm(INPUT) wrote to tmp, x is still intact). - gpu.copy_f32_buffer(&s.residual, &s.x, dim) - .map_err(hip_to_dispatch)?; - // x = residual + tmp (tmp holds post_attn_norm output). - gpu.copy_f32_buffer(&s.x, &s.residual, dim) - .map_err(hip_to_dispatch)?; - gpu.add_inplace_f32(&s.x, &s.tmp).map_err(hip_to_dispatch)?; - // Save x → residual for the FFN residual stream. - gpu.copy_f32_buffer(&s.residual, &s.x, dim) - .map_err(hip_to_dispatch)?; - Ok(()) - } - g4_op::RESID_POST_FFN => { - // x = residual + tmp; x *= layer_scalar. - // Identical for dense and MoE — tmp already holds normalized output. - gpu.copy_f32_buffer(&s.x, &s.residual, dim) - .map_err(hip_to_dispatch)?; - gpu.add_inplace_f32(&s.x, &s.tmp).map_err(hip_to_dispatch)?; - let layer_scalar = match self.layer { - LayerWeights::Sliding(lw) => lw.layer_scalar_host, - LayerWeights::Full(lw) => lw.layer_scalar_host, - }; - gpu.scale_f32_recorded(&s.x, layer_scalar) - .map_err(hip_to_dispatch)?; - Ok(()) - } - other => Err(hip_bridge::HipError::new( - 0, - &format!("unknown RESID_GEMV opcode {other}"), - )), - }; - res.map_err(|e| hipfire_dispatch::types::DispatchError::Hip(e.to_string())) - } - fn run_norm( - &mut self, - gpu: &mut Gpu, - _ctx: &DispatchCtx, - op: &OpBinding, - ) -> Result<(), hipfire_dispatch::types::DispatchError> { - let s = self.scratch; - let config = self.config; - let res: HipResult<()> = match op_code(op) { - g4_op::NORM_INPUT => match self.layer { - LayerWeights::Sliding(lw) => { - gpu.rmsnorm_f32(&s.x, &lw.input_layernorm, &s.tmp, config.norm_eps) - } - LayerWeights::Full(lw) => { - gpu.rmsnorm_f32(&s.x, &lw.input_layernorm, &s.tmp, config.norm_eps) - } - }, - g4_op::NORM_POST_ATTN => match self.layer { - LayerWeights::Sliding(lw) => gpu.rmsnorm_f32( - &s.tmp, - &lw.post_attention_layernorm, - &s.tmp, - config.norm_eps, - ), - LayerWeights::Full(lw) => gpu.rmsnorm_f32( - &s.tmp, - &lw.post_attention_layernorm, - &s.tmp, - config.norm_eps, - ), - }, - g4_op::NORM_PRE_FFN => match self.layer { - LayerWeights::Sliding(lw) => { - gpu.rmsnorm_f32(&s.x, &lw.pre_feedforward_layernorm, &s.tmp, config.norm_eps) - } - LayerWeights::Full(lw) => { - gpu.rmsnorm_f32(&s.x, &lw.pre_feedforward_layernorm, &s.tmp, config.norm_eps) + } else { + run_prefill_gemm( + gpu, + &lw.gate_proj, + &scratch.pb_tmp, + &scratch.pb_gate, + n_batch, + Some(&scratch.pb_bf16), + )?; + run_prefill_gemm( + gpu, + &lw.up_proj, + &scratch.pb_tmp, + &scratch.pb_up, + n_batch, + Some(&scratch.pb_bf16), + )?; } - }, - g4_op::NORM_POST_FFN => match self.layer { - // Dense-only norm. MoE layers use Moe(MOE_BRANCH) instead. - LayerWeights::Sliding(lw) => gpu.rmsnorm_f32( - &s.ffn_out, - &lw.post_feedforward_layernorm, - &s.tmp, - config.norm_eps, - ), - LayerWeights::Full(lw) => gpu.rmsnorm_f32( - &s.ffn_out, - &lw.post_feedforward_layernorm, - &s.tmp, - config.norm_eps, - ), - }, - other => { - return Err(hipfire_dispatch::types::DispatchError::Hip(format!( - "unknown NORM opcode {other}" - ))); + gpu.gelu_tanh_f32( + &scratch.pb_gate, + &scratch.pb_ffn_hidden, + n_batch * config.hidden_dim, + )?; + gpu.mul_f32( + &scratch.pb_ffn_hidden, + &scratch.pb_up, + &scratch.pb_ffn_hidden, + )?; + run_prefill_gemm( + gpu, + &lw.down_proj, + &scratch.pb_ffn_hidden, + &scratch.pb_ffn_out, + n_batch, + Some(&scratch.pb_bf16), + )?; } - }; - res.map_err(|e| hipfire_dispatch::types::DispatchError::Hip(e.to_string())) - } - fn run_attend( - &mut self, - gpu: &mut Gpu, - _ctx: &DispatchCtx, - op: &OpBinding, - ) -> Result<(), hipfire_dispatch::types::DispatchError> { - let s = self.scratch; - let config = self.config; - let pos = self.pos; - let hip_to_dispatch = - |e: hip_bridge::HipError| hipfire_dispatch::types::DispatchError::Hip(e.to_string()); - - let res: HipResult<()> = match op_code(op) { - g4_op::ATTEND_SLIDING => { - let lw = match self.layer { - LayerWeights::Sliding(lw) => lw, - _ => { - return Err(hipfire_dispatch::types::DispatchError::Hip( - "ATTEND_SLIDING on non-Sliding layer".into(), - )); - } - }; - let head_dim = config.sliding_head_dim; - let n_heads = config.n_heads; - let n_kv = config.sliding_n_kv_heads; - - // q/k/v norms (v_norm is no-scale — ones buffer). - gpu.rmsnorm_batched(&s.q, &lw.q_norm, &s.q, n_heads, head_dim, config.norm_eps) - .map_err(hip_to_dispatch)?; - gpu.rmsnorm_batched(&s.k, &lw.k_norm, &s.k, n_kv, head_dim, config.norm_eps) - .map_err(hip_to_dispatch)?; - gpu.rmsnorm_batched( - &s.v, - &s.v_norm_ones_full, - &s.v, - n_kv, - head_dim, - config.norm_eps, - ) - .map_err(hip_to_dispatch)?; - - // Pre-scale Q by sqrt(head_dim) — Gemma 4 attention scale is 1.0. - gpu.scale_f32_recorded(&s.q, (head_dim as f32).sqrt()) - .map_err(hip_to_dispatch)?; - - // Full rotate_half RoPE (all dims rotate). - gpu.rope_f32( - &s.q, - &s.k, - &s.pos_buf, - n_heads, - n_kv, - head_dim, - config.sliding_rope_theta, - ) - .map_err(hip_to_dispatch)?; + _ => unreachable!(), + } - // KV write + flash attention via dispatch. - let kv = &mut *self.kv_sliding; - let kv_layer_idx = self.sliding_kv_idx; - let ctx = DispatchCtx::new(gpu); - let tier_inputs = KvTierInputs { - quant_asym4: kv.quant_asym4, - quant_asym3: kv.quant_asym3, - quant_asym2: kv.quant_asym2, - quant_q8: kv.quant_q8, - quant_fwht: kv.quant_fwht, - quant_hfq4: false, - quant_q4: false, - quant_int8: false, - quant_hfq8: false, - quant_bf16: false, - quant_fp8: false, - f32_policy: hipfire_dispatch::families::kv_tier::F32AttnPolicy::Simple, - v_mode_bits: kv.v_mode_bits(), - pos, - flash_mode: 2, - capture_mode: gpu.graphs.capture_mode, - batch_size: 1, - is_tree: false, - is_boundary: false, - q8_windowed: true, - window: config.sliding_window as i32, - }; - let plan = KvTierPlan::derive(tier_inputs) - .map_err(|e| hipfire_dispatch::types::DispatchError::Hip(format!("{:?}", e)))?; - let io = AttnParams { - q: &s.q, - k: &s.k, - v: &s.v, - k_cache: &kv.k_gpu[kv_layer_idx], - v_cache: &kv.v_gpu[kv_layer_idx], - k_scales: None, - v_scales: None, - pos_buf: &s.pos_buf, - pos, - positions: None, - n_heads, - n_kv_heads: n_kv, - head_dim, - physical_cap: kv.physical_cap, - batch_size: 1, - max_ctx_len: 0, - flash_partials: Some(&s.flash_partials), - givens_cos: kv.givens_cos.as_ref(), - givens_sin: kv.givens_sin.as_ref(), - tree_bias: None, - block_start: 0, - block_cols: 0, - output_gate: None, - output_awq_scale: None, - output: &s.attn_out, - }; - execute_steps(gpu, &ctx, &[Step::Attend { plan, io }]) - .map_err(|e| hip_bridge::HipError::new(0, &e.to_string())) + // MoE branch (or dense fallback). HIPFIRE_MOE_BYPASS=1 forces dense + // path even on MoE layers (parity with v1 — used to isolate whether + // a regression lives in apply_moe_branch_batched vs the dense path). + let moe_bypass = hipfire_config::developer_var("HIPFIRE_MOE_BYPASS") + .ok() + .as_deref() + == Some("1"); + match (moe_opt, moe_bypass) { + (Some(moe), false) => { + apply_moe_branch_batched(gpu, config, scratch, moe, post_ffn_norm_ref, n_batch)?; + gpu.add_inplace_f32(&scratch.pb_residual, &scratch.pb_moe_pre2)?; + gpu.scale_f32(&scratch.pb_residual, layer_scalar)?; } - g4_op::ATTEND_FULL => { - let lw = match self.layer { - LayerWeights::Full(lw) => lw, - _ => { - return Err(hipfire_dispatch::types::DispatchError::Hip( - "ATTEND_FULL on non-Full layer".into(), - )); - } - }; - let head_dim = config.full_head_dim; - let n_heads = config.n_heads; - let n_kv = config.full_n_kv_heads; - let kv_bytes = n_kv * head_dim * 4; - - // CRITICAL: capture pre-k_norm K as V before applying k_norm. - gpu.copy_f32_buffer(&s.v, &s.k, kv_bytes / 4) - .map_err(hip_to_dispatch)?; - - // q/k/v norms (v_norm is no-scale — ones buffer). - gpu.rmsnorm_batched(&s.q, &lw.q_norm, &s.q, n_heads, head_dim, config.norm_eps) - .map_err(hip_to_dispatch)?; - gpu.rmsnorm_batched(&s.k, &lw.k_norm, &s.k, n_kv, head_dim, config.norm_eps) - .map_err(hip_to_dispatch)?; + _ => { gpu.rmsnorm_batched( - &s.v, - &s.v_norm_ones_full, - &s.v, - n_kv, - head_dim, + &scratch.pb_ffn_out, + post_ffn_norm_ref, + &scratch.pb_ffn_out, + n_batch, + dim, config.norm_eps, - ) - .map_err(hip_to_dispatch)?; - - // Pre-scale Q by sqrt(head_dim). - gpu.scale_f32_recorded(&s.q, (head_dim as f32).sqrt()) - .map_err(hip_to_dispatch)?; - - // Proportional partial RoPE (only first n_rot_pairs pairs rotate). - let n_rot_pairs = - ((head_dim as f32) * config.full_partial_rotary_factor * 0.5) as usize; - gpu.rope_partial_halved_f32( - &s.q, - &s.k, - &s.pos_buf, - n_heads, - n_kv, - head_dim, - n_rot_pairs, - config.full_rope_theta, - ) - .map_err(hip_to_dispatch)?; - - // KV write + flash attention via dispatch. - let kv = &mut *self.kv_full; - let kv_layer_idx = self.full_kv_idx; - let ctx = DispatchCtx::new(gpu); - let tier_inputs = KvTierInputs { - quant_asym4: kv.quant_asym4, - quant_asym3: kv.quant_asym3, - quant_asym2: kv.quant_asym2, - quant_q8: kv.quant_q8, - quant_fwht: kv.quant_fwht, - quant_hfq4: false, - quant_q4: false, - quant_int8: false, - quant_hfq8: false, - quant_bf16: false, - quant_fp8: false, - f32_policy: hipfire_dispatch::families::kv_tier::F32AttnPolicy::Simple, - v_mode_bits: kv.v_mode_bits(), - pos, - flash_mode: 2, - capture_mode: gpu.graphs.capture_mode, - batch_size: 1, - is_tree: false, - is_boundary: false, - q8_windowed: false, - window: 0, - }; - let plan = KvTierPlan::derive(tier_inputs) - .map_err(|e| hipfire_dispatch::types::DispatchError::Hip(format!("{:?}", e)))?; - let io = AttnParams { - q: &s.q, - k: &s.k, - v: &s.v, - k_cache: &kv.k_gpu[kv_layer_idx], - v_cache: &kv.v_gpu[kv_layer_idx], - k_scales: None, - v_scales: None, - pos_buf: &s.pos_buf, - pos, - positions: None, - n_heads, - n_kv_heads: n_kv, - head_dim, - physical_cap: kv.max_seq, - batch_size: 1, - max_ctx_len: 0, - flash_partials: Some(&s.flash_partials), - givens_cos: kv.givens_cos.as_ref(), - givens_sin: kv.givens_sin.as_ref(), - tree_bias: None, - block_start: 0, - block_cols: 0, - output_gate: None, - output_awq_scale: None, - output: &s.attn_out, - }; - execute_steps(gpu, &ctx, &[Step::Attend { plan, io }]) - .map_err(|e| hip_bridge::HipError::new(0, &e.to_string())) - } - other => Err(hip_bridge::HipError::new( - 0, - &format!("unknown ATTEND opcode {other}"), - )), - }; - res.map_err(|e| hipfire_dispatch::types::DispatchError::Hip(e.to_string())) - } - fn run_moe( - &mut self, - gpu: &mut Gpu, - _ctx: &DispatchCtx, - _op: &OpBinding, - ) -> Result<(), hipfire_dispatch::types::DispatchError> { - let moe = match self.layer { - LayerWeights::Sliding(lw) => lw.moe.as_ref(), - LayerWeights::Full(lw) => lw.moe.as_ref(), - }; - let moe = match moe { - Some(m) => m, - None => { - return Err(hipfire_dispatch::types::DispatchError::Hip( - "MOE_BRANCH on layer without MoE extras".into(), - )); + )?; + gpu.add_inplace_f32(&scratch.pb_residual, &scratch.pb_ffn_out)?; + gpu.scale_f32(&scratch.pb_residual, layer_scalar)?; } - }; - let post_ffn_norm = match self.layer { - LayerWeights::Sliding(lw) => &lw.post_feedforward_layernorm, - LayerWeights::Full(lw) => &lw.post_feedforward_layernorm, - }; - apply_moe_branch( - gpu, - self.config, - self.scratch, - moe, - post_ffn_norm, - &self.scratch.residual, - ) - .map_err(|e| hipfire_dispatch::types::DispatchError::Hip(e.to_string())) - } - fn run_recurrent( - &mut self, - _gpu: &mut Gpu, - _ctx: &DispatchCtx, - _op: &OpBinding, - ) -> Result<(), hipfire_dispatch::types::DispatchError> { - Err(hipfire_dispatch::types::DispatchError::Hip( - "gemma4 has no recurrent ops".into(), - )) - } - fn run_conv( - &mut self, - _gpu: &mut Gpu, - _ctx: &DispatchCtx, - _op: &OpBinding, - ) -> Result<(), hipfire_dispatch::types::DispatchError> { - Err(hipfire_dispatch::types::DispatchError::Hip( - "gemma4 has no conv ops".into(), - )) - } - fn run_escape( - &mut self, - _gpu: &mut Gpu, - _ctx: &DispatchCtx, - _op: &OpBinding, - _kind: EscapeKind, - ) -> Result<(), hipfire_dispatch::types::DispatchError> { - Err(hipfire_dispatch::types::DispatchError::Hip( - "gemma4 escape not yet implemented (Step 6)".into(), - )) - } -} - -// ── Gate + lowered forward ─────────────────────────────────────────────── - -/// Cached `HIPFIRE_FORWARD_LOWERED` toggle. Default ON (byte-parity validated -/// 2026-06-08); set `HIPFIRE_FORWARD_LOWERED=0` to force the legacy hand path. -fn forward_lowered_enabled() -> bool { - static F: std::sync::OnceLock = std::sync::OnceLock::new(); - *F.get_or_init(|| { - // Default ON (byte-parity validated 2026-06-08). - // Set HIPFIRE_FORWARD_LOWERED=0 to force legacy hand path. - hipfire_config::developer_var("HIPFIRE_FORWARD_LOWERED") + } + if hipfire_config::developer_var("HIPFIRE_GEMMA4_DUMP") .ok() .as_deref() - != Some("0") - }) -} - -/// Lowered (#397 Ship 6) decode layer loop. Behaviorally equivalent to -/// `forward_scratch_inner`'s hand arms. Builds a coarse-super-op `LayerProgram` -/// per layer and runs it through the dispatch substrate's executor. -fn forward_scratch_inner_lowered( - gpu: &mut Gpu, - weights: &Gemma4Weights, - config: &Gemma4Config, - pos: usize, - kv_sliding: &mut llama::KvCache, - kv_full: &mut llama::KvCache, - scratch: &Gemma4Scratch, -) -> HipResult<()> { - let ctx = DispatchCtx::new(gpu); - let mut sliding_kv_idx = 0usize; - let mut full_kv_idx = 0usize; - - for (layer_idx, layer_type) in config.layer_types.iter().copied().enumerate() { - let layer = &weights.layers[layer_idx]; - let variant = gemma4_variant_of(layer_type, layer); - let program = lower_variant(variant); + == Some("1") { - let mut bind = Gemma4Bindings { - layer, - config, - scratch, - kv_sliding, - kv_full, - pos, - sliding_kv_idx, - full_kv_idx, - }; - superop::run_layer_program(gpu, &ctx, &program, &mut bind) - .map_err(|e| hip_bridge::HipError::new(0, &e.to_string()))?; - } - match layer_type { - LayerType::Sliding => sliding_kv_idx += 1, - LayerType::Full => full_kv_idx += 1, + let data = gpu.download_f32(&scratch.pb_residual).unwrap_or_default(); + let last = &data[(n_batch - 1) * dim..n_batch * dim]; + let sum: f64 = last.iter().map(|&v| v as f64).sum(); + let min = last.iter().fold(f32::INFINITY, |a, &b| a.min(b)); + let max = last.iter().fold(f32::NEG_INFINITY, |a, &b| a.max(b)); + eprintln!( + "[v2 diag] L{layer_idx} {:?} last-tok hidden: first4={:?} sum={sum:.4e} min={min:.4} max={max:.4}", + layer_type, + &last[..4.min(last.len())] + ); } } - // Output stage: final norm + lm_head + softcap (same as hand path). + // ── Final logits: extract last token, run final_norm + lm_head + softcap ── + // The decode loop samples from scratch.logits, which must be filled + // with the last token's logits. Copy last position from pb_residual + // to scratch.tmp, run final_norm, lm_head, and softcap. + let last_offset = (n_batch - 1) * dim * 4; + if let Some(_s) = gpu.active_stream.as_ref() { + gpu.hip.memcpy_dtod_async_at( + &scratch.tmp.buf, + 0, + &scratch.pb_residual.buf, + last_offset, + dim * 4, + _s, + )?; + } else { + gpu.hip.memcpy_dtod_at( + &scratch.tmp.buf, + 0, + &scratch.pb_residual.buf, + last_offset, + dim * 4, + )?; + } gpu.rmsnorm_f32( - &scratch.x, + &scratch.tmp, &weights.final_norm, &scratch.tmp, config.norm_eps, )?; - { - let ctx = DispatchCtx::new(gpu); - let wr = weights.lm_head.dispatch_ref(); - execute_steps( - gpu, - &ctx, - &[Step::Gemv { - w: &wr, - input: GemvInput::Raw(&scratch.tmp), - out: &scratch.logits, - }], - ) - .map_err(|e| hip_bridge::HipError::new(0, &e.to_string()))?; - } + let ctx = DispatchCtx::new(gpu); + let wr_lm = weights.lm_head.dispatch_ref(); + execute_steps( + gpu, + &ctx, + &[Step::Gemv { + w: &wr_lm, + input: GemvInput::Raw(&scratch.tmp), + out: &scratch.logits, + }], + ) + .map_err(|e| hip_bridge::HipError::new(0, &e.to_string()))?; if config.final_logit_softcapping > 0.0 { gpu.logit_softcap_f32( &scratch.logits, @@ -7672,5 +5400,6 @@ fn forward_scratch_inner_lowered( config.final_logit_softcapping, )?; } + Ok(()) } diff --git a/crates/hipfire-arch-gemma4/src/program.rs b/crates/hipfire-arch-gemma4/src/program.rs index 85b5f629eb..5db64ece14 100644 --- a/crates/hipfire-arch-gemma4/src/program.rs +++ b/crates/hipfire-arch-gemma4/src/program.rs @@ -4,49 +4,356 @@ //! Declarative Gemma 4 decoder-layer program. //! -//! Each layer is `[SandwichAttention, SandwichMlp, PerLayerInput?, Scale?]`: -//! the per-layer-input branch exists on E-series checkpoints, the scale step -//! when the learned layer scalar is not 1. This module only binds resident -//! weights, KV storage and scratch to the shared -//! `hipfire_dispatch::pipeline::sandwich` operations; route selection and -//! executor bodies live in dispatch. The same program serves single-token -//! decode (`rows == 1`) and batched prefill / verify (`rows > 1`). +//! Each layer is +//! `[SandwichAttention, SandwichMlp | ParallelMoeMlp, PerLayerInput?, Scale?]`: +//! MoE checkpoints run the parallel dense+expert block, E-series checkpoints +//! add the per-layer-input branch, and the scale step exists when the learned +//! layer scalar is not 1. This module only binds resident weights, KV storage +//! and scratch to the shared `hipfire_dispatch::pipeline::sandwich` +//! operations; route selection and executor bodies live in dispatch. The same +//! program serves single-token decode (`rows == 1`) and batched prefill / +//! verify (`rows > 1`), for both weight stacks (`gemma4` and `lowered`). use crate::config::{Gemma4Config, LayerType, RopeType}; -use crate::gemma4::{Gemma4State, LayerWeights, PerLayerBranchWeights}; +use crate::gemma4::{self, PerLayerBranchWeights}; +use crate::lowered; +use hip_bridge::DeviceBuffer; use hipfire_dispatch::families::gemv::WeightRef; use hipfire_dispatch::pipeline::sandwich::{ - Activation, PerLayerInputOp, Rope, RopeKind, SandwichAttentionOp, SandwichKv, SandwichMlpOp, - SandwichStream, ScaleOp, SoftcapOp, + Activation, KvProjection, ParallelMoeMlpOp, PerLayerInputOp, Rope, RopeKind, RoutedExperts, + RoutedScratch, SandwichAttentionOp, SandwichKv, SandwichMlpOp, SandwichStream, ScaleOp, + SoftcapOp, }; use hipfire_dispatch::pipeline::{GemvInput, Step}; use hipfire_dispatch::types::RotationPlan; use hipfire_runtime::llama::{KvCache, KvCacheExt, WeightTensor}; use rdna_compute::GpuTensor; -/// The typed RoPE of one layer type. -pub(crate) fn layer_rope(cfg: &Gemma4Config, layer_type: LayerType) -> Rope { - match layer_type { - LayerType::Sliding => Rope { - kind: RopeKind::RotateHalf, - theta: cfg.sliding_rope_theta, - }, - LayerType::Full => { - let head_dim = cfg.full_head_dim; - let rot_pairs = match cfg.full_rope_type { - RopeType::Proportional => { - ((head_dim as f32) * cfg.full_partial_rotary_factor * 0.5) as usize - } - RopeType::Default => head_dim / 2, - }; - Rope { - kind: RopeKind::PartialHalved { rot_pairs }, - theta: cfg.full_rope_theta, - } +/// Attention geometry of one layer type. +#[derive(Clone, Copy)] +pub(crate) struct AttnGeometry { + pub head_dim: usize, + pub n_kv_heads: usize, + /// Sliding window; `0` = full causal. + pub window: usize, + pub rope: Rope, +} + +/// Model geometry the program reads; built from either config type. +#[derive(Clone, Copy)] +pub(crate) struct Geometry { + pub dim: usize, + pub eps: f32, + pub n_heads: usize, + pub n_layers: usize, + pub vocab: usize, + pub softcap: f32, + /// Per-layer-input width (E-series); `0` without PLE. + pub ple_width: usize, + /// Routed experts per token and per-expert FFN width (MoE checkpoints). + pub moe_top_k: usize, + pub moe_hidden: usize, + pub sliding: AttnGeometry, + pub full: AttnGeometry, +} + +fn full_rope(head_dim: usize, rope_type: RopeType, factor: f32, theta: f32) -> Rope { + let rot_pairs = match rope_type { + RopeType::Proportional => ((head_dim as f32) * factor * 0.5) as usize, + RopeType::Default => head_dim / 2, + }; + Rope { + kind: RopeKind::PartialHalved { rot_pairs }, + theta, + } +} + +impl Geometry { + pub fn eager(cfg: &Gemma4Config) -> Self { + Self { + dim: cfg.dim, + eps: cfg.norm_eps, + n_heads: cfg.n_heads, + n_layers: cfg.n_layers, + vocab: cfg.vocab_size, + softcap: cfg.final_logit_softcapping, + ple_width: cfg.hidden_size_per_layer_input, + moe_top_k: 0, + moe_hidden: 0, + sliding: AttnGeometry { + head_dim: cfg.sliding_head_dim, + n_kv_heads: cfg.sliding_n_kv_heads, + window: cfg.sliding_window, + rope: Rope { + kind: RopeKind::RotateHalf, + theta: cfg.sliding_rope_theta, + }, + }, + full: AttnGeometry { + head_dim: cfg.full_head_dim, + n_kv_heads: cfg.full_n_kv_heads, + window: 0, + rope: full_rope( + cfg.full_head_dim, + cfg.full_rope_type, + cfg.full_partial_rotary_factor, + cfg.full_rope_theta, + ), + }, + } + } + + /// The arch-22 draft head's blocks (no softcap; the head is bound + /// separately). + pub fn drafter(cfg: &crate::drafter::Gemma4DrafterConfig) -> Self { + Self { + dim: cfg.hidden, + eps: cfg.norm_eps, + n_heads: cfg.n_heads, + n_layers: cfg.n_layers, + vocab: cfg.vocab_size, + softcap: 0.0, + ple_width: 0, + moe_top_k: 0, + moe_hidden: 0, + sliding: AttnGeometry { + head_dim: cfg.sliding_head_dim, + n_kv_heads: cfg.sliding_n_kv_heads, + window: cfg.sliding_window, + rope: Rope { + kind: RopeKind::RotateHalf, + theta: cfg.sliding_rope_theta, + }, + }, + full: AttnGeometry { + head_dim: cfg.full_head_dim, + n_kv_heads: cfg.full_n_kv_heads, + window: 0, + rope: full_rope( + cfg.full_head_dim, + cfg.full_rope_type, + cfg.full_partial_rotary_factor, + cfg.full_rope_theta, + ), + }, + } + } + + pub fn lowered(cfg: &lowered::Gemma4Config) -> Self { + let rope_type = match cfg.full_rope_type { + lowered::RopeType::Default => RopeType::Default, + lowered::RopeType::Proportional => RopeType::Proportional, + }; + Self { + dim: cfg.dim, + eps: cfg.norm_eps, + n_heads: cfg.n_heads, + n_layers: cfg.n_layers, + vocab: cfg.vocab_size, + softcap: cfg.final_logit_softcapping, + ple_width: 0, + moe_top_k: cfg.top_k_experts, + moe_hidden: cfg.moe_intermediate_size, + sliding: AttnGeometry { + head_dim: cfg.sliding_head_dim, + n_kv_heads: cfg.sliding_n_kv_heads, + window: cfg.sliding_window, + rope: Rope { + kind: RopeKind::RotateHalf, + theta: cfg.sliding_rope_theta, + }, + }, + full: AttnGeometry { + head_dim: cfg.full_head_dim, + n_kv_heads: cfg.full_n_kv_heads, + window: 0, + rope: full_rope( + cfg.full_head_dim, + rope_type, + cfg.full_partial_rotary_factor, + cfg.full_rope_theta, + ), + }, + } + } +} + +#[derive(Clone, Copy)] +struct LayerKvWeights<'a> { + k: &'a WeightTensor, + v: Option<&'a WeightTensor>, + k_norm: &'a GpuTensor, +} + +/// Per-layer weights the program binds, from any Gemma 4 weight stack. +pub(crate) struct LayerRefs<'a> { + layer_type: LayerType, + input_norm: &'a GpuTensor, + post_attn_norm: &'a GpuTensor, + pre_ffn_norm: &'a GpuTensor, + post_ffn_norm: &'a GpuTensor, + layer_scalar: f32, + q: &'a WeightTensor, + q_norm: &'a GpuTensor, + /// Own K projection, V projection (`None` = K=V) and K norm; `None` on + /// query-only layers that read another layer's cache. + kv: Option>, + o: &'a WeightTensor, + gate: &'a WeightTensor, + up: &'a WeightTensor, + down: &'a WeightTensor, + ffn_hidden_dim: usize, + per_layer: Option<&'a PerLayerBranchWeights>, + moe: Option<&'a lowered::MoeLayerExtras>, +} + +impl<'a> LayerRefs<'a> { + pub fn eager(layer: &'a gemma4::LayerWeights) -> Self { + match layer { + gemma4::LayerWeights::Sliding(l) => Self { + layer_type: LayerType::Sliding, + input_norm: &l.input_layernorm, + post_attn_norm: &l.post_attention_layernorm, + pre_ffn_norm: &l.pre_feedforward_layernorm, + post_ffn_norm: &l.post_feedforward_layernorm, + layer_scalar: l.layer_scalar_host, + q: &l.q_proj, + q_norm: &l.q_norm, + kv: Some(LayerKvWeights { + k: &l.k_proj, + v: Some(&l.v_proj), + k_norm: &l.k_norm, + }), + o: &l.o_proj, + gate: &l.gate_proj, + up: &l.up_proj, + down: &l.down_proj, + ffn_hidden_dim: l.ffn_hidden_dim, + per_layer: l.per_layer.as_ref(), + moe: None, + }, + gemma4::LayerWeights::Full(l) => Self { + layer_type: LayerType::Full, + input_norm: &l.input_layernorm, + post_attn_norm: &l.post_attention_layernorm, + pre_ffn_norm: &l.pre_feedforward_layernorm, + post_ffn_norm: &l.post_feedforward_layernorm, + layer_scalar: l.layer_scalar_host, + q: &l.q_proj, + q_norm: &l.q_norm, + kv: Some(LayerKvWeights { + k: &l.k_proj, + v: l.v_proj.as_ref(), + k_norm: &l.k_norm, + }), + o: &l.o_proj, + gate: &l.gate_proj, + up: &l.up_proj, + down: &l.down_proj, + ffn_hidden_dim: l.ffn_hidden_dim, + per_layer: l.per_layer.as_ref(), + moe: None, + }, + } + } + + /// `hidden_dim` is the lowered config's dense FFN width. + pub fn lowered(layer: &'a lowered::LayerWeights, hidden_dim: usize) -> Self { + match layer { + lowered::LayerWeights::Sliding(l) => Self { + layer_type: LayerType::Sliding, + input_norm: &l.input_layernorm, + post_attn_norm: &l.post_attention_layernorm, + pre_ffn_norm: &l.pre_feedforward_layernorm, + post_ffn_norm: &l.post_feedforward_layernorm, + layer_scalar: l.layer_scalar_host, + q: &l.q_proj, + q_norm: &l.q_norm, + kv: Some(LayerKvWeights { + k: &l.k_proj, + v: Some(&l.v_proj), + k_norm: &l.k_norm, + }), + o: &l.o_proj, + gate: &l.gate_proj, + up: &l.up_proj, + down: &l.down_proj, + ffn_hidden_dim: hidden_dim, + per_layer: None, + moe: l.moe.as_ref(), + }, + lowered::LayerWeights::Full(l) => Self { + layer_type: LayerType::Full, + input_norm: &l.input_layernorm, + post_attn_norm: &l.post_attention_layernorm, + pre_ffn_norm: &l.pre_feedforward_layernorm, + post_ffn_norm: &l.post_feedforward_layernorm, + layer_scalar: l.layer_scalar_host, + q: &l.q_proj, + q_norm: &l.q_norm, + kv: Some(LayerKvWeights { + k: &l.k_proj, + v: None, + k_norm: &l.k_norm, + }), + o: &l.o_proj, + gate: &l.gate_proj, + up: &l.up_proj, + down: &l.down_proj, + ffn_hidden_dim: hidden_dim, + per_layer: None, + moe: l.moe.as_ref(), + }, + } + } +} + +impl<'a> LayerRefs<'a> { + /// A draft-head block: query-only attention over the target's cache. + pub fn drafter( + layer: &'a crate::drafter::DrafterLayerWeights, + layer_type: LayerType, + hidden_dim: usize, + ) -> Self { + Self { + layer_type, + input_norm: &layer.input_layernorm, + post_attn_norm: &layer.post_attention_layernorm, + pre_ffn_norm: &layer.pre_feedforward_layernorm, + post_ffn_norm: &layer.post_feedforward_layernorm, + layer_scalar: layer.layer_scalar_host, + q: &layer.q_proj, + q_norm: &layer.q_norm, + kv: None, + o: &layer.o_proj, + gate: &layer.gate_proj, + up: &layer.up_proj, + down: &layer.down_proj, + ffn_hidden_dim: hidden_dim, + per_layer: None, + moe: None, } } } +/// Resident buffers every layer of a forward shares. +pub(crate) struct Resident<'a> { + pub kv_sliding: &'a KvCache, + pub kv_full: &'a KvCache, + /// Device position scalar (`rows == 1`). + pub pos_buf: &'a DeviceBuffer, + /// Ones vector of the larger head dim (weight-less V norm). + pub v_norm_ones: &'a GpuTensor, + pub flash_partials: &'a GpuTensor, +} + +/// Which cache slot a layer attends, and whether it writes its own row. +#[derive(Clone, Copy)] +pub(crate) struct LayerKv { + pub slot: usize, + pub writes: bool, +} + /// Per-layer-input scratch; present on E-series checkpoints. pub(crate) struct PleScratch<'a> { /// `[rows, n_layers * width]` per-layer inputs of every row. @@ -72,11 +379,12 @@ pub(crate) struct LayerScratch<'a> { pub act: &'a GpuTensor, pub mlp_out: &'a GpuTensor, pub ple: Option>, + pub moe: Option>, } impl<'a> LayerScratch<'a> { - /// The resident single-row scratch of `state`. - pub fn decode(state: &'a Gemma4State) -> Self { + /// The resident single-row scratch of the eager state. + pub fn eager(state: &'a gemma4::Gemma4State) -> Self { let ple = match ( &state.ple_projection_all, &state.ple_gate, @@ -106,14 +414,48 @@ impl<'a> LayerScratch<'a> { act: &state.ffn_hidden, mlp_out: &state.ffn_out, ple, + moe: None, + } + } + + /// The resident single-row scratch of the lowered state. + pub fn lowered(s: &'a lowered::Gemma4Scratch, moe: bool) -> Self { + Self { + x: &s.x, + residual: &s.residual, + normed: &s.tmp, + attn_rot: &s.x_rot, + mlp_rot: &s.x_rot, + q: &s.q, + k: &s.k, + v: &s.v, + attn_out: &s.attn_out, + gate: &s.gate_ffn, + up: &s.up_ffn, + act: &s.ffn_hidden, + mlp_out: &s.ffn_out, + ple: None, + moe: moe.then_some(RoutedScratch { + input: &s.moe_pre2, + input_rot: &s.moe_pre2_rot, + router_in: &s.moe_router_in, + router_logits: &s.moe_router_logits, + topk_indices: &s.moe_topk_indices, + topk_weights: &s.moe_topk_weights, + gate: &s.moe_expert_gate_batch, + up: &s.moe_expert_up_batch, + act: &s.moe_expert_hidden_batch, + out: &s.moe_cur_moe, + dense_normed: &s.moe_cur_mlp, + }), } } } /// Binding of one forward over `rows` consecutive positions from `position`. pub(crate) struct ProgramBinding<'a> { - pub cfg: &'a Gemma4Config, - pub state: &'a Gemma4State, + pub geo: Geometry, + pub resident: Resident<'a>, pub rows: usize, pub position: usize, /// Per-row i32 positions (`rows > 1`). @@ -121,143 +463,62 @@ pub(crate) struct ProgramBinding<'a> { pub scratch: LayerScratch<'a>, } -struct LayerRefs<'a> { - input_norm: &'a GpuTensor, - post_attn_norm: &'a GpuTensor, - pre_ffn_norm: &'a GpuTensor, - post_ffn_norm: &'a GpuTensor, - layer_scalar: f32, - q: &'a WeightTensor, - k: &'a WeightTensor, - v: Option<&'a WeightTensor>, - o: &'a WeightTensor, - q_norm: &'a GpuTensor, - k_norm: &'a GpuTensor, - gate: &'a WeightTensor, - up: &'a WeightTensor, - down: &'a WeightTensor, - ffn_hidden_dim: usize, - per_layer: Option<&'a PerLayerBranchWeights>, -} - -fn layer_refs(layer: &LayerWeights) -> LayerRefs<'_> { - match layer { - LayerWeights::Sliding(l) => LayerRefs { - input_norm: &l.input_layernorm, - post_attn_norm: &l.post_attention_layernorm, - pre_ffn_norm: &l.pre_feedforward_layernorm, - post_ffn_norm: &l.post_feedforward_layernorm, - layer_scalar: l.layer_scalar_host, - q: &l.q_proj, - k: &l.k_proj, - v: Some(&l.v_proj), - o: &l.o_proj, - q_norm: &l.q_norm, - k_norm: &l.k_norm, - gate: &l.gate_proj, - up: &l.up_proj, - down: &l.down_proj, - ffn_hidden_dim: l.ffn_hidden_dim, - per_layer: l.per_layer.as_ref(), - }, - LayerWeights::Full(l) => LayerRefs { - input_norm: &l.input_layernorm, - post_attn_norm: &l.post_attention_layernorm, - pre_ffn_norm: &l.pre_feedforward_layernorm, - post_ffn_norm: &l.post_feedforward_layernorm, - layer_scalar: l.layer_scalar_host, - q: &l.q_proj, - k: &l.k_proj, - v: l.v_proj.as_ref(), - o: &l.o_proj, - q_norm: &l.q_norm, - k_norm: &l.k_norm, - gate: &l.gate_proj, - up: &l.up_proj, - down: &l.down_proj, - ffn_hidden_dim: l.ffn_hidden_dim, - per_layer: l.per_layer.as_ref(), - }, - } -} - impl<'a> ProgramBinding<'a> { fn stream(&self) -> SandwichStream<'a> { SandwichStream { rows: self.rows, - hidden: self.cfg.dim, - eps: self.cfg.norm_eps, + hidden: self.geo.dim, + eps: self.geo.eps, x: self.scratch.x, residual: self.scratch.residual, normed: self.scratch.normed, } } - /// Append the steps of `layer_idx` to `steps`. + /// Append the steps of `layer_idx` to `steps`. A layer writes its own + /// cache rows iff it has K/V weights and `kv_slot.writes`. pub fn layer( &self, layer_idx: usize, - layer: &'a LayerWeights, + w: LayerRefs<'a>, + kv_slot: LayerKv, steps: &mut Vec>, ) -> Result<(), String> { - let cfg = self.cfg; - let st = self.state; + let geo = &self.geo; + let r = &self.resident; let sc = &self.scratch; - let layer_type = cfg.layer_types[layer_idx]; - if !matches!( - (layer_type, layer), - (LayerType::Sliding, LayerWeights::Sliding(_)) - | (LayerType::Full, LayerWeights::Full(_)) - ) { - return Err(format!("gemma4 layer {layer_idx} type/weights mismatch")); - } - let shared = match cfg.kv_shared_source_layer_idx(layer_idx) { - Some(source) => Some(st.kv_slot_for_layer[source]), - None if cfg.is_kv_shared_layer(layer_idx) => { - return Err(format!( - "gemma4 layer {layer_idx}: missing same-type KV sharing source" - )); - } - None => None, - }; - let (kv, head_dim, n_kv_heads, window): (&KvCache, _, _, _) = match layer_type { - LayerType::Sliding => ( - &st.kv_sliding, - cfg.sliding_head_dim, - cfg.sliding_n_kv_heads, - cfg.sliding_window, - ), - LayerType::Full => (&st.kv_full, cfg.full_head_dim, cfg.full_n_kv_heads, 0), + let (kv, attn): (&KvCache, _) = match w.layer_type { + LayerType::Sliding => (r.kv_sliding, geo.sliding), + LayerType::Full => (r.kv_full, geo.full), }; - let w = layer_refs(layer); - let slot = shared.unwrap_or(st.kv_slot_for_layer[layer_idx]); - let writes = shared.is_none(); + let kv_weights = w.kv.filter(|_| kv_slot.writes); steps.push(Step::SandwichAttention(SandwichAttentionOp { stream: self.stream(), position: self.position, - n_heads: cfg.n_heads, - n_kv_heads, - head_dim, + n_heads: geo.n_heads, + n_kv_heads: attn.n_kv_heads, + head_dim: attn.head_dim, input_norm: w.input_norm, wq: w.q.dispatch_ref(), - wk: writes.then(|| w.k.dispatch_ref()), - wv: w.v.filter(|_| writes).map(WeightTensor::dispatch_ref), q_norm: w.q_norm, - k_norm: w.k_norm, - v_norm: Some(&st.v_norm_ones), - q_scale: (head_dim as f32).sqrt(), - rope: layer_rope(cfg, layer_type), + kv_proj: kv_weights.map(|kv| KvProjection { + wk: kv.k.dispatch_ref(), + wv: kv.v.map(WeightTensor::dispatch_ref), + k_norm: kv.k_norm, + v_norm: Some(r.v_norm_ones), + }), + q_scale: (attn.head_dim as f32).sqrt(), + rope: attn.rope, kv: SandwichKv { tier: kv.tier_inputs(), - k_cache: &kv.k_gpu[slot], - v_cache: &kv.v_gpu[slot], + k_cache: &kv.k_gpu[kv_slot.slot], + v_cache: &kv.v_gpu[kv_slot.slot], physical_cap: kv.physical_cap, givens_cos: kv.givens_cos.as_ref(), givens_sin: kv.givens_sin.as_ref(), - window, - write: writes, + window: attn.window, }, - pos_buf: &st.pos_buf, + pos_buf: r.pos_buf, positions: self.positions, wo: w.o.dispatch_ref(), post_norm: w.post_attn_norm, @@ -266,9 +527,9 @@ impl<'a> ProgramBinding<'a> { k: sc.k, v: sc.v, attn_out: sc.attn_out, - flash_partials: &st.q8_flash_partials, + flash_partials: r.flash_partials, })); - steps.push(Step::SandwichMlp(SandwichMlpOp { + let mlp = SandwichMlpOp { stream: self.stream(), pre_norm: w.pre_ffn_norm, w_gate: w.gate.dispatch_ref(), @@ -282,8 +543,41 @@ impl<'a> ProgramBinding<'a> { up: sc.up, act: sc.act, out: sc.mlp_out, - })); - if let Some(ple) = w.per_layer.filter(|_| cfg.hidden_size_per_layer_input != 0) { + }; + match w.moe { + None => steps.push(Step::SandwichMlp(mlp)), + Some(moe) => { + let scratch = sc + .moe + .as_ref() + .ok_or_else(|| format!("gemma4 layer {layer_idx}: missing MoE scratch"))?; + let expert = moe + .experts + .first() + .ok_or_else(|| format!("gemma4 layer {layer_idx}: MoE layer has no experts"))?; + steps.push(Step::ParallelMoeMlp(ParallelMoeMlpOp { + mlp, + dense_post_norm: &moe.post_feedforward_layernorm_1, + experts: RoutedExperts { + pre_norm: &moe.pre_feedforward_layernorm_2, + router_norm: &moe.router_scale, + router_input_scale: 1.0 / (geo.dim as f32).sqrt(), + router: moe.router_proj.dispatch_ref(), + n_experts: moe.experts.len(), + top_k: geo.moe_top_k, + hidden_dim: geo.moe_hidden, + gate_up_dtype: expert.gate_up_proj.gpu_dtype, + down_dtype: expert.down_proj.gpu_dtype, + gate_up_ptrs: &moe.experts_gate_up_ptrs, + down_ptrs: &moe.experts_down_ptrs, + per_expert_scale: &moe.per_expert_scale, + post_norm: &moe.post_feedforward_layernorm_2, + }, + scratch: RoutedScratch { ..*scratch }, + })); + } + } + if let Some(ple) = w.per_layer.filter(|_| geo.ple_width != 0) { let s = sc .ple .as_ref() @@ -291,8 +585,8 @@ impl<'a> ProgramBinding<'a> { steps.push(Step::PerLayerInput(PerLayerInputOp { stream: self.stream(), layer: layer_idx, - layer_width: cfg.hidden_size_per_layer_input, - n_layers: cfg.n_layers, + layer_width: geo.ple_width, + n_layers: geo.n_layers, inputs: s.inputs, w_gate: ple.input_gate.dispatch_ref(), w_proj: ple.projection.dispatch_ref(), @@ -312,11 +606,31 @@ impl<'a> ProgramBinding<'a> { } } +/// KV slot and write ownership of every eager layer (E-series KV sharing). +pub(crate) fn eager_layer_kv( + cfg: &Gemma4Config, + state: &gemma4::Gemma4State, + layer_idx: usize, +) -> Result { + match cfg.kv_shared_source_layer_idx(layer_idx) { + Some(source) => Ok(LayerKv { + slot: state.kv_slot_for_layer[source], + writes: false, + }), + None if !cfg.is_kv_shared_layer(layer_idx) => Ok(LayerKv { + slot: state.kv_slot_for_layer[layer_idx], + writes: true, + }), + None => Err(format!( + "gemma4 layer {layer_idx}: missing same-type KV sharing source" + )), + } +} + /// Final norm, tied LM head and logit soft cap of one row `x` into `logits`. /// `lm_head` is the caller-held `weights.lm_head.dispatch_ref()`. -#[allow(clippy::too_many_arguments)] pub(crate) fn head<'a>( - cfg: &Gemma4Config, + geo: &Geometry, final_norm: &'a GpuTensor, lm_head: &'a WeightRef<'a>, x: &'a GpuTensor, @@ -330,8 +644,8 @@ pub(crate) fn head<'a>( x_plain: normed, out: normed, awq_scale: None, - k: cfg.dim, - eps: cfg.norm_eps, + k: geo.dim, + eps: geo.eps, rotation: RotationPlan::None, }); steps.push(Step::Gemv { @@ -339,11 +653,11 @@ pub(crate) fn head<'a>( input: GemvInput::Raw(normed), out: logits, }); - if cfg.final_logit_softcapping > 0.0 { + if geo.softcap > 0.0 { steps.push(Step::Softcap(SoftcapOp { logits, - n: cfg.vocab_size, - cap: cfg.final_logit_softcapping, + n: geo.vocab, + cap: geo.softcap, })); } } diff --git a/crates/hipfire-dispatch/src/pipeline/sandwich.rs b/crates/hipfire-dispatch/src/pipeline/sandwich.rs index c1a77dab37..c51520fce9 100644 --- a/crates/hipfire-dispatch/src/pipeline/sandwich.rs +++ b/crates/hipfire-dispatch/src/pipeline/sandwich.rs @@ -236,28 +236,20 @@ pub struct SandwichStream<'a> { } impl SandwichStream<'_> { - fn bytes(&self) -> usize { - self.rows * self.hidden * 4 + fn len(&self) -> usize { + self.rows * self.hidden } + /// Copies are kernel launches so hipGraph capture and the retained-replay + /// recorder both see them (a D2D memcpy is invisible to Redline PM4). fn save_residual(&self, gpu: &mut Gpu) -> Result<(), DispatchError> { - if self.rows == 1 { - gpu.memcpy_dtod_auto(&self.residual.buf, &self.x.buf, self.bytes()) - } else { - gpu.hip - .memcpy_dtod_at(&self.residual.buf, 0, &self.x.buf, 0, self.bytes()) - } - .map_err(hip) + gpu.copy_f32_buffer(self.residual, self.x, self.len()) + .map_err(hip) } fn restore_residual(&self, gpu: &mut Gpu) -> Result<(), DispatchError> { - if self.rows == 1 { - gpu.memcpy_dtod_auto(&self.x.buf, &self.residual.buf, self.bytes()) - } else { - gpu.hip - .memcpy_dtod_at(&self.x.buf, 0, &self.residual.buf, 0, self.bytes()) - } - .map_err(hip) + gpu.copy_f32_buffer(self.x, self.residual, self.len()) + .map_err(hip) } fn norm( @@ -317,9 +309,9 @@ impl Rope { } } -/// The K/V cache one attention sublayer reads, and whether it writes the -/// current rows first. `write == false` reads another layer's populated cache -/// (KV sharing, draft heads). +/// The K/V cache one attention sublayer reads. A sublayer with its own +/// [`KvProjection`] writes the current rows first; a query-only sublayer reads +/// another layer's populated cache (KV sharing, draft heads). pub struct SandwichKv<'a> { pub tier: KvTierInputs, pub k_cache: &'a GpuTensor, @@ -329,15 +321,23 @@ pub struct SandwichKv<'a> { pub givens_sin: Option<&'a GpuTensor>, /// Sliding window; `0` = full causal. pub window: usize, - pub write: bool, +} + +/// K/V projections of an attention sublayer that writes its own cache rows. +pub struct KvProjection<'a> { + pub wk: WeightRef<'a>, + /// `None` takes V from the pre-norm K (K=V). + pub wv: Option>, + pub k_norm: &'a GpuTensor, + /// Weight-less V RMSNorm, bound as a ones vector of `head_dim`. + pub v_norm: Option<&'a GpuTensor>, } /// Attention sublayer: /// `x = residual + post_norm(o(attend(rope(q_norm(q)), rope(k_norm(k)), v_norm(v))))` /// over `input_norm(x)`. /// -/// `wk == None` projects only Q and attends a cache it does not write. -/// `wv == None` with `wk` present takes V from the pre-norm K (K=V). +/// `kv_proj == None` projects only Q and attends a cache it does not write. /// Row `r` sits at absolute position `position + r`. pub struct SandwichAttentionOp<'a> { pub stream: SandwichStream<'a>, @@ -347,12 +347,8 @@ pub struct SandwichAttentionOp<'a> { pub head_dim: usize, pub input_norm: &'a GpuTensor, pub wq: WeightRef<'a>, - pub wk: Option>, - pub wv: Option>, pub q_norm: &'a GpuTensor, - pub k_norm: &'a GpuTensor, - /// Weight-less V RMSNorm, bound as a ones vector of `head_dim`. - pub v_norm: Option<&'a GpuTensor>, + pub kv_proj: Option>, /// Q multiplier applied after q_norm (sets the softmax scale against the /// kernel's `1/sqrt(head_dim)`). pub q_scale: f32, @@ -374,16 +370,24 @@ pub struct SandwichAttentionOp<'a> { } impl SandwichAttentionOp<'_> { + fn wk(&self) -> Option<&WeightRef<'_>> { + self.kv_proj.as_ref().map(|p| &p.wk) + } + + fn wv(&self) -> Option<&WeightRef<'_>> { + self.kv_proj.as_ref().and_then(|p| p.wv.as_ref()) + } + fn projections(&self) -> [Option<(&WeightRef<'_>, &GpuTensor)>; 3] { [ Some((&self.wq, self.q)), - self.wk.as_ref().map(|w| (w, self.k)), - self.wv.as_ref().map(|w| (w, self.v)), + self.wk().map(|w| (w, self.k)), + self.wv().map(|w| (w, self.v)), ] } - fn kv_bytes(&self) -> usize { - self.stream.rows * self.n_kv_heads * self.head_dim * 4 + fn kv_len(&self) -> usize { + self.stream.rows * self.n_kv_heads * self.head_dim } /// `input_norm(x)` → Q, K, V projections (V from K when K=V). @@ -396,7 +400,7 @@ impl SandwichAttentionOp<'_> { .projections() .iter() .flatten() - .all(|(w, _)| w.dtype == DType::MQ4G256); + .all(|(w, _)| w.dtype == DType::MQ4G256 && w.awq_scale.is_none()); if ROUTES.attn_norm && all_mq4 { gpu.fused_rmsnorm_rotate_mq(s.x, self.input_norm, self.x_rot, s.hidden, s.eps) .map_err(hip)?; @@ -405,7 +409,7 @@ impl SandwichAttentionOp<'_> { } } else { s.norm(gpu, s.x, self.input_norm, s.normed)?; - match &self.wk { + match self.wk() { Some(wk) if ROUTES.q8_qk && self.wq.dtype == DType::Q8_0 && wk.dtype == DType::Q8_0 => { @@ -427,12 +431,12 @@ impl SandwichAttentionOp<'_> { } None => gemv(gpu, ctx, &self.wq, s.normed, self.q)?, } - if let Some(wv) = &self.wv { + if let Some(wv) = self.wv() { gemv(gpu, ctx, wv, s.normed, self.v)?; } } - if self.wk.is_some() && self.wv.is_none() { - gpu.memcpy_dtod_auto(&self.v.buf, &self.k.buf, self.kv_bytes()) + if self.wk().is_some() && self.wv().is_none() { + gpu.copy_f32_buffer(self.v, self.k, self.kv_len()) .map_err(hip)?; } Ok(()) @@ -443,7 +447,7 @@ impl SandwichAttentionOp<'_> { let rows = s.rows; s.norm(gpu, s.x, self.input_norm, s.normed)?; let mut v_projected = false; - match (&self.wk, &self.wv) { + match (self.wk(), self.wv()) { (Some(wk), Some(wv)) if q8_fused_prefill(gpu, &[&self.wq, wk, wv]) => { fused_projection( gpu, @@ -478,9 +482,8 @@ impl SandwichAttentionOp<'_> { } } } - if self.wk.is_some() && !v_projected { - gpu.hip - .memcpy_dtod_at(&self.v.buf, 0, &self.k.buf, 0, self.kv_bytes()) + if self.wk().is_some() && !v_projected { + gpu.copy_f32_buffer(self.v, self.k, self.kv_len()) .map_err(hip)?; } Ok(()) @@ -491,20 +494,17 @@ impl SandwichAttentionOp<'_> { let eps = self.stream.eps; let rows = self.stream.rows; let (n_heads, n_kv, hd) = (self.n_heads, self.n_kv_heads, self.head_dim); - let writes_k = self.wk.is_some(); - if writes_k { - if let Some(ones) = self.v_norm { - gpu.rmsnorm_batched(self.v, ones, self.v, rows * n_kv, hd, eps) - .map_err(hip)?; - } + if let Some(ones) = self.kv_proj.as_ref().and_then(|p| p.v_norm) { + gpu.rmsnorm_batched(self.v, ones, self.v, rows * n_kv, hd, eps) + .map_err(hip)?; } - if rows == 1 && writes_k && ROUTES.qk_norm_rope { + if let (1, Some(p), true) = (rows, &self.kv_proj, ROUTES.qk_norm_rope) { return gpu .fused_gemma4_qk_norm_rope_f32( self.q, self.k, self.q_norm, - self.k_norm, + p.k_norm, self.pos_buf, n_heads, n_kv, @@ -518,12 +518,12 @@ impl SandwichAttentionOp<'_> { } gpu.rmsnorm_batched(self.q, self.q_norm, self.q, rows * n_heads, hd, eps) .map_err(hip)?; - if writes_k { - gpu.rmsnorm_batched(self.k, self.k_norm, self.k, rows * n_kv, hd, eps) + if let Some(p) = &self.kv_proj { + gpu.rmsnorm_batched(self.k, p.k_norm, self.k, rows * n_kv, hd, eps) .map_err(hip)?; } scale(gpu, self.q, self.q_scale)?; - let rope_kv = if writes_k { n_kv } else { 0 }; + let rope_kv = if self.kv_proj.is_some() { n_kv } else { 0 }; let theta = self.rope.theta; match (self.rope.kind, rows) { (RopeKind::RotateHalf, 1) => { @@ -610,7 +610,7 @@ impl SandwichAttentionOp<'_> { output_awq_scale: None, output: self.attn_out, }; - if kv.write { + if self.kv_proj.is_some() { ATTENTION.run_attention(ctx, gpu, &plan, &io) } else { ATTENTION.run_attend_only(ctx, gpu, &plan, &io) @@ -629,7 +629,7 @@ impl SandwichAttentionOp<'_> { let positions = self.positions()?; let (n_heads, n_kv, hd) = (self.n_heads, self.n_kv_heads, self.head_dim); let seq_len = self.position + rows; - if kv.write { + if self.kv_proj.is_some() { gpu.kv_cache_write_q8_0_batched(kv.k_cache, self.k, positions, n_kv, hd, rows) .map_err(hip)?; gpu.kv_cache_write_q8_0_batched(kv.v_cache, self.v, positions, n_kv, hd, rows) @@ -701,11 +701,6 @@ pub fn execute_sandwich_attention( ctx: &DispatchCtx, op: &SandwichAttentionOp<'_>, ) -> Result<(), DispatchError> { - if op.wk.is_none() && op.kv.write { - return Err(DispatchError::Hip( - "sandwich attention: a Q-only layer cannot write its KV row".into(), - )); - } let s = op.stream; s.save_residual(gpu)?; op.project(gpu, ctx)?; @@ -749,6 +744,16 @@ pub fn execute_sandwich_mlp( gpu: &mut Gpu, ctx: &DispatchCtx, op: &SandwichMlpOp<'_>, +) -> Result<(), DispatchError> { + dense_mlp(gpu, ctx, op)?; + op.stream.post_norm_residual(gpu, op.out, op.post_norm) +} + +/// Residual save, then `out = down(act(gate(n)) * up(n))`, `n = pre_norm(x)`. +fn dense_mlp( + gpu: &mut Gpu, + ctx: &DispatchCtx, + op: &SandwichMlpOp<'_>, ) -> Result<(), DispatchError> { let s = op.stream; let rows = s.rows; @@ -801,7 +806,193 @@ pub fn execute_sandwich_mlp( } else { gemm_rows(gpu, &op.w_down, op.act, op.out, op.x_rot, rows)?; } - s.post_norm_residual(gpu, op.out, op.post_norm) + Ok(()) +} + +/// Routed-expert half of a parallel dense+MoE block. +pub struct RoutedExperts<'a> { + /// Norm of the residual stream feeding the experts. + pub pre_norm: &'a GpuTensor, + /// Router input norm weight; the normed input is scaled by + /// `router_input_scale` before the router projection. + pub router_norm: &'a GpuTensor, + pub router_input_scale: f32, + pub router: WeightRef<'a>, + pub n_experts: usize, + pub top_k: usize, + /// Per-expert FFN width. + pub hidden_dim: usize, + pub gate_up_dtype: DType, + pub down_dtype: DType, + /// `[n_experts]` device pointers to each expert's fused gate/up weight. + pub gate_up_ptrs: &'a GpuTensor, + /// `[n_experts]` device pointers to each expert's down weight. + pub down_ptrs: &'a GpuTensor, + /// `[n_experts]` learned scale applied to each expert's output. + pub per_expert_scale: &'a GpuTensor, + /// Norm of the combined expert output. + pub post_norm: &'a GpuTensor, +} + +impl RoutedExperts<'_> { + /// Expert formats with an indexed device-resident kernel pair. + pub fn supports(gate_up: DType, down: DType) -> bool { + matches!( + gate_up, + DType::MQ4G256 | DType::HFQ4G256 | DType::HFQ6G256 | DType::Q8_0 + ) && matches!(down, DType::Q8_0 | DType::HFQ4G128) + } +} + +/// Scratch of one routed-expert pass. +#[derive(Clone, Copy)] +pub struct RoutedScratch<'a> { + /// `pre_norm(residual)`, and its FWHT rotation for MagnumQuant experts. + pub input: &'a GpuTensor, + pub input_rot: &'a GpuTensor, + pub router_in: &'a GpuTensor, + pub router_logits: &'a GpuTensor, + pub topk_indices: &'a GpuTensor, + pub topk_weights: &'a GpuTensor, + /// `[top_k, hidden_dim]` each. + pub gate: &'a GpuTensor, + pub up: &'a GpuTensor, + pub act: &'a GpuTensor, + /// Weighted sum of the selected experts. + pub out: &'a GpuTensor, + /// `dense_post_norm(dense MLP output)`. + pub dense_normed: &'a GpuTensor, +} + +/// Parallel dense+MoE MLP sublayer: +/// `x = residual + post_norm(dense_post_norm(mlp(n)) + experts.post_norm(moe(r)))` +/// with `n = mlp.pre_norm(x)` and `r = experts.pre_norm(x)`; the router reads +/// `router_norm(x) * router_input_scale`. `mlp.post_norm` is the block's final +/// post-norm. +pub struct ParallelMoeMlpOp<'a> { + pub mlp: SandwichMlpOp<'a>, + pub dense_post_norm: &'a GpuTensor, + pub experts: RoutedExperts<'a>, + pub scratch: RoutedScratch<'a>, +} + +pub fn execute_parallel_moe_mlp( + gpu: &mut Gpu, + ctx: &DispatchCtx, + op: &ParallelMoeMlpOp<'_>, +) -> Result<(), DispatchError> { + let mlp = &op.mlp; + let s = mlp.stream; + let e = &op.experts; + let r = &op.scratch; + if s.rows != 1 { + return Err(DispatchError::Hip(format!( + "parallel MoE MLP rows={} has no batched executor", + s.rows + ))); + } + if e.top_k != 8 { + return Err(DispatchError::Hip(format!( + "parallel MoE MLP: top_k={} unsupported (indexed kernels are k=8)", + e.top_k + ))); + } + if !RoutedExperts::supports(e.gate_up_dtype, e.down_dtype) { + return Err(DispatchError::Hip(format!( + "parallel MoE MLP: expert formats gate_up={:?} down={:?} have no indexed kernel", + e.gate_up_dtype, e.down_dtype + ))); + } + dense_mlp(gpu, ctx, mlp)?; + s.norm(gpu, mlp.out, op.dense_post_norm, r.dense_normed)?; + + // Router over the post-attention residual stream. + s.norm(gpu, s.residual, e.pre_norm, r.input)?; + s.norm(gpu, s.residual, e.router_norm, r.router_in)?; + scale(gpu, r.router_in, e.router_input_scale)?; + gemv(gpu, ctx, &e.router, r.router_in, r.router_logits)?; + gpu.moe_softmax_topk_renorm_k8( + r.router_logits, + r.topk_indices, + r.topk_weights, + e.n_experts, + true, + ) + .map_err(hip)?; + + // Q8 down accumulates atomically; HFQ4G128 down assigns. + gpu.zero_f32(r.out).map_err(hip)?; + let (hidden, mi, k) = (s.hidden, e.hidden_dim, e.top_k); + match e.gate_up_dtype { + DType::MQ4G256 => { + gpu.rotate_x_mq(r.input, r.input_rot, hidden).map_err(hip)?; + gpu.gemv_mq4g256_moe_gate_up_k8_indexed( + e.gate_up_ptrs, + r.topk_indices, + r.input_rot, + r.gate, + r.up, + 2 * mi, + hidden, + ) + .map_err(hip)?; + } + DType::HFQ4G256 | DType::HFQ6G256 => crate::pipeline::run_uniform_moe_gate_up( + gpu, + e.gate_up_dtype, + e.gate_up_ptrs, + r.topk_indices, + r.input, + r.gate, + r.up, + 2 * mi, + hidden, + k, + )?, + _ => gpu + .gemv_q8_0_moe_gate_up_k8_indexed( + e.gate_up_ptrs, + r.topk_indices, + r.input, + r.gate, + r.up, + 2 * mi, + hidden, + ) + .map_err(hip)?, + } + gpu.gelu_tanh_f32(r.gate, r.act, k * mi).map_err(hip)?; + gpu.mul_f32(r.act, r.up, r.act).map_err(hip)?; + if e.down_dtype == DType::Q8_0 { + gpu.gemv_q8_0_moe_down_residual_scaled_k8_indexed( + e.down_ptrs, + r.topk_indices, + r.topk_weights, + e.per_expert_scale, + r.act, + r.out, + hidden, + mi, + ) + } else { + gpu.gemv_hfq4g128_moe_down_residual_scaled_k8_indexed( + e.down_ptrs, + r.topk_indices, + r.topk_weights, + e.per_expert_scale, + r.act, + r.out, + hidden, + mi, + ) + } + .map_err(hip)?; + s.norm(gpu, r.out, e.post_norm, r.out)?; + + gpu.add_f32(r.dense_normed, r.out, s.normed).map_err(hip)?; + s.norm(gpu, s.normed, mlp.post_norm, s.normed)?; + s.restore_residual(gpu)?; + gpu.add_inplace_f32(s.x, s.normed).map_err(hip) } /// Per-layer-input branch: diff --git a/crates/hipfire-dispatch/src/pipeline/steps.rs b/crates/hipfire-dispatch/src/pipeline/steps.rs index ccb27769d9..5724352d59 100644 --- a/crates/hipfire-dispatch/src/pipeline/steps.rs +++ b/crates/hipfire-dispatch/src/pipeline/steps.rs @@ -109,6 +109,8 @@ pub enum Step<'a> { SandwichAttention(crate::pipeline::sandwich::SandwichAttentionOp<'a>), /// Sandwich-norm gated MLP sublayer. SandwichMlp(crate::pipeline::sandwich::SandwichMlpOp<'a>), + /// Parallel dense + routed-expert MLP sublayer. + ParallelMoeMlp(crate::pipeline::sandwich::ParallelMoeMlpOp<'a>), /// Sandwich-norm per-layer-input branch. PerLayerInput(crate::pipeline::sandwich::PerLayerInputOp<'a>), /// In-place scalar multiply. @@ -149,6 +151,7 @@ fn op_kind(step: &Step) -> Option { | Step::MoeStage(..) | Step::SandwichAttention(_) | Step::SandwichMlp(_) + | Step::ParallelMoeMlp(_) | Step::PerLayerInput(_) | Step::Scale(_) | Step::Softcap(_) => None, @@ -1505,6 +1508,9 @@ fn launch_op(gpu: &mut Gpu, ctx: &DispatchCtx, step: &Step) -> Result<(), Dispat crate::pipeline::sandwich::execute_sandwich_attention(gpu, ctx, op) } Step::SandwichMlp(op) => crate::pipeline::sandwich::execute_sandwich_mlp(gpu, ctx, op), + Step::ParallelMoeMlp(op) => { + crate::pipeline::sandwich::execute_parallel_moe_mlp(gpu, ctx, op) + } Step::PerLayerInput(op) => crate::pipeline::sandwich::execute_per_layer_input(gpu, ctx, op), Step::Scale(op) => crate::pipeline::sandwich::execute_scale(gpu, op), Step::Softcap(op) => crate::pipeline::sandwich::execute_softcap(gpu, op), diff --git a/crates/hipfire-dispatch/src/pipeline/superop.rs b/crates/hipfire-dispatch/src/pipeline/superop.rs index 4124ffa060..6e9b6acece 100644 --- a/crates/hipfire-dispatch/src/pipeline/superop.rs +++ b/crates/hipfire-dispatch/src/pipeline/superop.rs @@ -26,7 +26,7 @@ //! //! ## Coverage (the whole served fleet, one substrate) //! - `Proj` / `ResidualGemv` / `Moe` → qwen35, MiniMax(reuse), cohere2moe(reuse) -//! - `Attend` (flavor-carrying) → all; Gemma SWA/qk-norm/softcap/k_eq_v live here as flavors +//! - `Attend` → all //! - `Recurrent` → qwen35 DeltaNet linear-attention state //! - `Conv` → LFM2 depthwise causal short-conv mixer (+ conv state) //! - `Escape(EscapeKind)` → irregular/stateful: deepseek4 compressor/indexer/SWA, etc. @@ -62,54 +62,10 @@ pub enum ScratchSlot { Cache(u32), } -/// FFN/gate-up activation flavor. SiLU for qwen-family; GeLU-tanh (GeGLU) for -/// Gemma (`gelu_tanh(gate)·up`). -#[derive(Clone, Copy, Debug, PartialEq, Eq)] -pub enum ActFlavor { - SiluMul, - GeluTanhMul, -} - -/// RoPE flavor carried by an `Attend` super-op (resolved from config at load). -#[derive(Clone, Copy, Debug, PartialEq)] -pub enum RopeFlavor { - None, - /// Standard rotate-half (most archs). `theta` = rope base. - HalfRotate { - theta: f32, - }, - /// Interleaved full-dim RoPE (e.g. cohere2moe). - Interleaved { - theta: f32, - }, -} - -/// Attention-block flavor — everything that distinguishes one arch's attention -/// from another, resolved at load so the per-token `Attend` is branch-free. -/// Gemma exercises the full surface (SWA window, per-head qk-norm, q·√hd scaling, -/// the k_eq_v weightless-V-RMSNorm prelude, logit softcap). -#[derive(Clone, Copy, Debug, PartialEq)] -pub struct AttnFlavor { - /// Sliding-window size; 0 = full (global) attention. Gemma alternates - /// Sliding(1024)/Full per layer. - pub window: u32, - /// Per-head q_norm/k_norm over head_dim (Gemma, qwen35). - pub qk_norm: bool, - /// q *= sqrt(head_dim) (Gemma query scaling). - pub q_scale_sqrt_hd: bool, - /// V = copy of K before k_norm + weightless RMSNorm on V (Gemma full layers). - pub k_eq_v: bool, - /// Attention-logit softcap value; 0.0 = none (Gemma-2 style). - pub logit_softcap: f32, - pub rope: RopeFlavor, -} - /// Per-super-op flavor payload (None for ops with no flavor axis, e.g. Proj). #[derive(Clone, Copy, Debug, PartialEq)] pub enum OpFlavor { None, - Attn(AttnFlavor), - Act(ActFlavor), } /// Irregular/stateful ops that don't map onto a single fused kernel. Each is a @@ -123,8 +79,6 @@ pub enum EscapeKind { Deepseek4IndexerTopK, /// deepseek4 sparse SWA over the gathered top-K KV. Deepseek4SwaTopK, - /// Gemma final logit softcap (output stage). - GemmaLogitSoftcap, } /// One coarse super-op. For `Proj`/`ResidualGemv`/`Moe` the `key` is the @@ -146,7 +100,7 @@ pub enum SuperOpKind { /// projection cluster did NOT fuse — mirrors execute_steps' unfused /// per-step launch_op path (rmsnorm, then individual Proj gemvs). Norm, - /// Attention block (flavor in `OpFlavor::Attn`). + /// Attention block. Attend, /// MoE FFN block (routes to MoeFamily / run_moe_decode). Moe, @@ -249,9 +203,9 @@ fn lower_walk( /// the frozen `KernelKey`s match `execute_steps` byte-for-byte — no fusion drift. /// /// TODO(arch-migration, Step 3+): populate `WeightSlot`/`ScratchSlot` bindings -/// and `AttnFlavor`/`ActFlavor` from the arch's weight/scratch tables + config. -/// This step freezes the FUSION STRUCTURE + kernel keys; the per-arch migration -/// supplies the operand slots + flavor data (and the executor binds them). +/// from the arch's weight/scratch tables. This step freezes the FUSION +/// STRUCTURE + kernel keys; the per-arch migration supplies the operand slots +/// (and the executor binds them). pub fn lower_layer(steps: &[Step], ctx: &DispatchCtx) -> LayerProgram { lower_walk( steps.len(), diff --git a/docs/design/gemma4-declarative-program.md b/docs/design/gemma4-declarative-program.md index eaca2a1a00..5bd32258a7 100644 --- a/docs/design/gemma4-declarative-program.md +++ b/docs/design/gemma4-declarative-program.md @@ -77,9 +77,34 @@ Projections, the final norm and the LM head reuse `Gemv` and - EAGLE drafter: tau parity and coherent text; greedy EAGLE output equals AR output where it did before. - Redline: the 26B shadow oracle (`redline_shadow_gemma4`) stays exact. + Residual and K=V copies are `copy_f32_buffer` launches, not D2D memcpys, + so the retained recorder sees them; a memcpy there made PM4 replay inexact. +- Measured (gfx1151, pre-port daemon `8a6c4710`): every 12B/E2B/E4B row + bitwise; 26B with `HIPFIRE_GEMMA4_FUSED_QK_ROPE=0 + HIPFIRE_GEMMA4_FUSED_POSTNORM=0` byte-identical, default 2 of 5 prompts + identical and 3 diverging late, coherent; EAGLE per-prompt tau + 3.368/3.459/2.667/3.447/3.514 before and after; the 12B lowered opt-in + route (`HIPFIRE_BATCHED_PREFILL=1`) byte-identical with every + `HIPFIRE_GEMMA4_FUSED_*` knob at 0, since it too now takes the eager fused + routes; 26B + `redline_daemon_harness.py --pm4 --skip-prefill` shadow exact (1083 + launches, was 1263; tape hash `249f24e5f39d425d`, was `0909f3792962c37d`). ## Progress | Piece | State | Commit | |---|---|---| -| Eager decode + batched prefill/verify + E-series PLE as `[SandwichAttention, SandwichMlp, PerLayerInput?, Scale?]` (`program.rs`, dispatch `sandwich.rs`); hand arms deleted | landed; bitwise | (this commit) | +| Eager decode + batched prefill/verify + E-series PLE as `[SandwichAttention, SandwichMlp, PerLayerInput?, Scale?]` (`program.rs`, dispatch `sandwich.rs`); hand arms deleted | landed; bitwise | `30b55f753` | +| MoE decode (`ParallelMoeMlp`), lowered decode on the shared program; super-op facade (`lower_variant`, `Gemma4Bindings`) and lowered hand arms deleted; unsupported expert formats refuse at load | landed; 26B byte-identical with the two non-bitwise fusions off, coherent with them on | (this commit) | +| EAGLE draft head as one step list (`Gemv` pre-projection, query-only `[SandwichAttention, SandwichMlp, Scale?]` blocks over the target's last slot, norm, `lm_head`, post-projection) | landed; EAGLE text and per-prompt tau identical | (this commit) | + +## Remaining + +- `lowered::forward_prefill_batch` (the tool-only batched prefill used by + `calib_sweep`, `eval_hipfire`, the KLD reference builders and the lowered + examples) still runs hand-written layer bodies. It carries the calibration + activation taps (`run_prefill_gemm`), BF16/WMMA prefill GEMM routes and + batched asym3 / per-token ring attention that the serve program does not. + Porting it needs a decision on where calibration taps live in dispatch. +- Two weight stacks remain (`gemma4.rs` eager, `lowered.rs` MoE / batched + opt-in) with duplicate config and loaders; both bind to the same program. From bf51655f5ed47e4a2de9c57ec019e4ba6c2e78e7 Mon Sep 17 00:00:00 2001 From: Bjoern Agent Date: Fri, 2 Oct 2026 19:35:31 +0200 Subject: [PATCH 03/13] docs(gemma4): record serve_harness A/B for the declarative port (cherry picked from commit da0acc5b9bfcfc7b074774b34d2a64ba68df290f) --- docs/design/gemma4-declarative-program.md | 8 ++++++-- 1 file changed, 6 insertions(+), 2 deletions(-) diff --git a/docs/design/gemma4-declarative-program.md b/docs/design/gemma4-declarative-program.md index 5bd32258a7..35273b8ec4 100644 --- a/docs/design/gemma4-declarative-program.md +++ b/docs/design/gemma4-declarative-program.md @@ -89,14 +89,18 @@ Projections, the final norm and the LM head reuse `Gemv` and routes; 26B `redline_daemon_harness.py --pm4 --skip-prefill` shadow exact (1083 launches, was 1263; tape hash `249f24e5f39d425d`, was `0909f3792962c37d`). +- `serve_harness.py --thinking off --sampling greedy --compare-transcript` + against the pre-port daemon: 12B battery and chain, E4B battery + `transcript_byte_identical=true`; 26B battery coherent (no runaway, empty or + attractor turns), turn 2 diverges. ## Progress | Piece | State | Commit | |---|---|---| | Eager decode + batched prefill/verify + E-series PLE as `[SandwichAttention, SandwichMlp, PerLayerInput?, Scale?]` (`program.rs`, dispatch `sandwich.rs`); hand arms deleted | landed; bitwise | `30b55f753` | -| MoE decode (`ParallelMoeMlp`), lowered decode on the shared program; super-op facade (`lower_variant`, `Gemma4Bindings`) and lowered hand arms deleted; unsupported expert formats refuse at load | landed; 26B byte-identical with the two non-bitwise fusions off, coherent with them on | (this commit) | -| EAGLE draft head as one step list (`Gemv` pre-projection, query-only `[SandwichAttention, SandwichMlp, Scale?]` blocks over the target's last slot, norm, `lm_head`, post-projection) | landed; EAGLE text and per-prompt tau identical | (this commit) | +| MoE decode (`ParallelMoeMlp`), lowered decode on the shared program; super-op facade (`lower_variant`, `Gemma4Bindings`) and lowered hand arms deleted; unsupported expert formats refuse at load | landed; 26B byte-identical with the two non-bitwise fusions off, coherent with them on | `287e67d5b` | +| EAGLE draft head as one step list (`Gemv` pre-projection, query-only `[SandwichAttention, SandwichMlp, Scale?]` blocks over the target's last slot, norm, `lm_head`, post-projection) | landed; EAGLE text and per-prompt tau identical | `287e67d5b` | ## Remaining From 110a7425447613e28e43fa50019101412430f4a0 Mon Sep 17 00:00:00 2001 From: Bjoern Agent Date: Sat, 3 Oct 2026 07:28:58 +0200 Subject: [PATCH 04/13] refactor(gemma4): drop the HIPFIRE_GEMMA4_FUSED_* developer switches The fused Gemma 4 routes (attn-norm+FWHT, Q8 q+k, qk-norm+RoPE, post-norm+residual, gate+up) are always on; only greedy EAGLE (HIPFIRE_GEMMA4_EAGLE=1) still turns off the two non-bitwise fusions. Default behaviour unchanged: 12B decode tokens and top-16 logits bitwise vs the pre-port daemon, EAGLE text and tau (3.368/3.459/2.667/3.447/3.514) identical. Daemon md5 d02826078f1f343392fed7517defbfea. (cherry picked from commit 9f53a8680540859b7bbdfb5204d859842485a12c) --- crates/hipfire-arch-gemma4/src/forward.rs | 11 ++-- .../hipfire-dispatch/src/pipeline/sandwich.rs | 53 +++++-------------- docs/design/gemma4-declarative-program.md | 8 +-- 3 files changed, 22 insertions(+), 50 deletions(-) diff --git a/crates/hipfire-arch-gemma4/src/forward.rs b/crates/hipfire-arch-gemma4/src/forward.rs index cafe70aa8e..b5d4ddf6b5 100644 --- a/crates/hipfire-arch-gemma4/src/forward.rs +++ b/crates/hipfire-arch-gemma4/src/forward.rs @@ -40,12 +40,11 @@ //! //! * `HIPFIRE_GEMMA4_GRAPH` (default ON; set =0 to disable) — hipGraph //! 48-layer body + lm_head (`decode_step_with_graph`). +2.6%. -//! * `HIPFIRE_GEMMA4_FUSED_FFN` (default ON) — fold pre-FFN rmsnorm+FWHT into -//! one launch then gate+up into one (`fused_rmsnorm_rotate_mq` + -//! `fused_gate_up_hfq4g256`; MQ4G256 bytes are HFQ4G256-compatible given a -//! pre-rotated input). +1.0–1.2%. -//! * `HIPFIRE_GEMMA4_FUSED_QK` (default ON) — fuse the Q8 q+k projections into -//! one launch (`fused_gate_up_q8_0`, shared rmsnorm input). +1.1%. +//! * fused FFN — fold pre-FFN rmsnorm+FWHT into one launch then gate+up into +//! one (`fused_rmsnorm_rotate_mq` + `fused_gate_up_hfq4g256`; MQ4G256 bytes +//! are HFQ4G256-compatible given a pre-rotated input). +1.0–1.2%. +//! * fused Q8 q+k projections in one launch (`fused_gate_up_q8_0`, shared +//! rmsnorm input). +1.1%. //! //! Full stack: 46.8 → 50.7 tok/s (+8.3%). The 70–75 tok/s target is NOT //! reachable via fusion/graph alone — it would require reading ~40% fewer diff --git a/crates/hipfire-dispatch/src/pipeline/sandwich.rs b/crates/hipfire-dispatch/src/pipeline/sandwich.rs index c51520fce9..f1851b8f0f 100644 --- a/crates/hipfire-dispatch/src/pipeline/sandwich.rs +++ b/crates/hipfire-dispatch/src/pipeline/sandwich.rs @@ -35,39 +35,15 @@ fn hip_dbg(e: impl std::fmt::Debug) -> DispatchError { DispatchError::Hip(format!("{e:?}")) } -/// Developer route switches. Each fused route is on unless its variable is -/// `0`/`off`/`false`. `HIPFIRE_GEMMA4_EAGLE=1` (greedy EAGLE) additionally -/// turns off the fusions that are not byte-identical to their unfused form, -/// so batched verify and single-row decode share one arithmetic. -struct Routes { - attn_norm: bool, - q8_qk: bool, - qk_norm_rope: bool, - post_norm: bool, - gate_up: bool, - /// Batched Q8 GEMMs use the unchunked scalar kernel (decode's arithmetic). - strict_q8_gemm: bool, -} - -static ROUTES: LazyLock = LazyLock::new(|| { - let on = |name: &str| { - !matches!( - hipfire_config::developer_var(name).ok().as_deref(), - Some("0") | Some("off") | Some("false") - ) - }; - let strict = hipfire_config::developer_var("HIPFIRE_GEMMA4_EAGLE") +/// Greedy EAGLE (`HIPFIRE_GEMMA4_EAGLE=1`) turns off the fusions that are not +/// byte-identical to their unfused form and keeps batched Q8 GEMMs on the +/// unchunked scalar kernel, so batched verify and single-row decode share one +/// arithmetic. +static EAGLE_STRICT: LazyLock = LazyLock::new(|| { + hipfire_config::developer_var("HIPFIRE_GEMMA4_EAGLE") .ok() .as_deref() - == Some("1"); - Routes { - attn_norm: on("HIPFIRE_GEMMA4_FUSED_ATTN_NORM"), - q8_qk: on("HIPFIRE_GEMMA4_FUSED_QK"), - qk_norm_rope: !strict && on("HIPFIRE_GEMMA4_FUSED_QK_ROPE"), - post_norm: !strict && on("HIPFIRE_GEMMA4_FUSED_POSTNORM"), - gate_up: on("HIPFIRE_GEMMA4_FUSED_FFN"), - strict_q8_gemm: strict, - } + == Some("1") }); static GEMV: LazyLock = LazyLock::new(GemvFamily::new); @@ -142,7 +118,7 @@ pub fn gemm_rows( DType::F32 => gpu .gemm_f32_batched(w.buf, x, y, w.m, w.k, rows) .map_err(hip), - DType::Q8_0 if ROUTES.strict_q8_gemm => gpu + DType::Q8_0 if *EAGLE_STRICT => gpu .gemm_q8_0_batched(w.buf, x, y, w.m, w.k, rows) .map_err(hip), DType::Q8_0 => gpu @@ -274,7 +250,7 @@ impl SandwichStream<'_> { out: &GpuTensor, post_norm: &GpuTensor, ) -> Result<(), DispatchError> { - if self.rows == 1 && ROUTES.post_norm { + if self.rows == 1 && !*EAGLE_STRICT { return gpu .rmsnorm_residual_add_f32(out, post_norm, self.residual, self.x, self.eps) .map_err(hip); @@ -401,7 +377,7 @@ impl SandwichAttentionOp<'_> { .iter() .flatten() .all(|(w, _)| w.dtype == DType::MQ4G256 && w.awq_scale.is_none()); - if ROUTES.attn_norm && all_mq4 { + if all_mq4 { gpu.fused_rmsnorm_rotate_mq(s.x, self.input_norm, self.x_rot, s.hidden, s.eps) .map_err(hip)?; for (w, out) in self.projections().iter().flatten() { @@ -410,9 +386,7 @@ impl SandwichAttentionOp<'_> { } else { s.norm(gpu, s.x, self.input_norm, s.normed)?; match self.wk() { - Some(wk) - if ROUTES.q8_qk && self.wq.dtype == DType::Q8_0 && wk.dtype == DType::Q8_0 => - { + Some(wk) if self.wq.dtype == DType::Q8_0 && wk.dtype == DType::Q8_0 => { gpu.fused_gate_up_q8_0( self.wq.buf, wk.buf, @@ -498,7 +472,7 @@ impl SandwichAttentionOp<'_> { gpu.rmsnorm_batched(self.v, ones, self.v, rows * n_kv, hd, eps) .map_err(hip)?; } - if let (1, Some(p), true) = (rows, &self.kv_proj, ROUTES.qk_norm_rope) { + if let (1, Some(p), true) = (rows, &self.kv_proj, !*EAGLE_STRICT) { return gpu .fused_gemma4_qk_norm_rope_f32( self.q, @@ -774,8 +748,7 @@ fn dense_mlp( gemm_rows(gpu, &op.w_gate, s.normed, op.gate, op.x_rot, rows)?; gemm_rows(gpu, &op.w_up, s.normed, op.up, op.x_rot, rows)?; } - } else if ROUTES.gate_up && op.w_gate.dtype == DType::MQ4G256 && op.w_up.dtype == DType::MQ4G256 - { + } else if op.w_gate.dtype == DType::MQ4G256 && op.w_up.dtype == DType::MQ4G256 { gpu.fused_rmsnorm_rotate_mq(s.x, op.pre_norm, op.x_rot, s.hidden, s.eps) .map_err(hip)?; gpu.fused_gate_up_hfq4g256( diff --git a/docs/design/gemma4-declarative-program.md b/docs/design/gemma4-declarative-program.md index 35273b8ec4..b5faf8912b 100644 --- a/docs/design/gemma4-declarative-program.md +++ b/docs/design/gemma4-declarative-program.md @@ -80,12 +80,12 @@ Projections, the final norm and the LM head reuse `Gemv` and Residual and K=V copies are `copy_f32_buffer` launches, not D2D memcpys, so the retained recorder sees them; a memcpy there made PM4 replay inexact. - Measured (gfx1151, pre-port daemon `8a6c4710`): every 12B/E2B/E4B row - bitwise; 26B with `HIPFIRE_GEMMA4_FUSED_QK_ROPE=0 - HIPFIRE_GEMMA4_FUSED_POSTNORM=0` byte-identical, default 2 of 5 prompts - identical and 3 diverging late, coherent; EAGLE per-prompt tau + bitwise; 26B with the fused qk-norm+RoPE and post-norm routes off + byte-identical (measured before those switches were removed), with them on + 2 of 5 prompts identical and 3 diverging late, coherent; EAGLE per-prompt tau 3.368/3.459/2.667/3.447/3.514 before and after; the 12B lowered opt-in route (`HIPFIRE_BATCHED_PREFILL=1`) byte-identical with every - `HIPFIRE_GEMMA4_FUSED_*` knob at 0, since it too now takes the eager fused + fusion off, since it too now takes the eager fused routes; 26B `redline_daemon_harness.py --pm4 --skip-prefill` shadow exact (1083 launches, was 1263; tape hash `249f24e5f39d425d`, was `0909f3792962c37d`). From 105b3151237219399001ba360b24ba820d457ff4 Mon Sep 17 00:00:00 2001 From: Bjoern Agent Date: Sat, 3 Oct 2026 09:24:02 +0200 Subject: [PATCH 05/13] refactor(gemma4): run the calibration tools on the declarative program through lowered::forward_prefill_batch, now a thin driver over the shared layer program (rows <= 64 per chunk, Q8 KV). The hand-written batched prefill (v1/v2 layer bodies, run_prefill_gemm, batched MoE, pb_* scratch) is deleted. - Calibration taps live in dispatch: the sandwich projection helpers (gemv, gemm_rows) and the fused-QKV family call maybe_capture_activation with the unrotated input, so any architecture on these steps calibrates without per-architecture taps. While a collector is armed, rows == 1 fusions that never materialize the unrotated input step aside. - gemm_rows stages BF16 teacher inputs for the MFMA GEMM and runs formats without a batched kernel as one GEMV per row. - ParallelMoeMlp gains a rows > 1 executor (batched dense half, routed experts row by row); rows == 1 issues the same launches as before. - HIPFIRE_BATCHED_PREFILL / HIPFIRE_WMMA_PREFILL are removed: dense Gemma 4 always loads eager, MoE on the lowered stack. - eval_hipfire's Gemma arm sat inside the llama branch, so Gemma models scored 0 tokens; it is its own branch again. Evidence (gfx1151, daemon 21384cc30cb4df9580cb6a057fd91649): - prefill_parity_gemma4, 725-token prompt: batched vs per-token same argmax and identical 24-token continuation; 12B mq4 35.2s -> 4.4s, 26B-A4B q8-experts 19.1s -> 7.9s. - calib_sweep 12B mq4 (2048 tokens, seq 512): coverage 328/328, consistency 0, q/k/(v) and gate/up identity PASS. - eval_hipfire 12B mq4, synthetic ref: prefill KLD 21.170 vs per-token 21.122. - Serve gates vs pre-port daemon 8a6c4710: 12B, 12B pb64, E2B pb64, E4B tokens and top-16 logits bitwise; EAGLE text and tau 3.368/3.459/2.667/3.447/3.514 unchanged; 26B text identical to the first port; 26B Redline PM4 shadow exact, 1083 launches, tape 249f24e5f39d425d. Rebased onto beta: lowered.rs keeps beta's gemma4_decode_flash_partials_len (and its test), the loader files keep beta's formatting with only the gemma4_source_uses_lowered / opt-in removal applied, and the carrier log line reports beta's full-tier KV mode. (cherry picked from commit 9774fefb5a6644039d7bb1c55c26bed7b86a4bdb) --- .../examples/prefill_parity_gemma4.rs | 64 +- crates/hipfire-arch-gemma4/src/carrier.rs | 98 +- crates/hipfire-arch-gemma4/src/forward.rs | 156 +- crates/hipfire-arch-gemma4/src/gemma4.rs | 24 +- crates/hipfire-arch-gemma4/src/lib.rs | 4 +- crates/hipfire-arch-gemma4/src/lowered.rs | 2458 ++--------------- crates/hipfire-arch-gemma4/src/program.rs | 191 +- .../hipfire-dispatch/src/pipeline/sandwich.rs | 153 +- crates/hipfire-loader/src/admission.rs | 2 +- crates/hipfire-loader/src/carriers.rs | 23 +- crates/hipfire-loader/src/lib.rs | 6 - .../hipfire-runtime/examples/calib_sweep.rs | 65 +- .../hipfire-runtime/examples/eval_hipfire.rs | 55 +- docs/design/gemma4-declarative-program.md | 41 +- docs/governance/debt-dispatch-bypass.txt | 2 +- scripts/mi300x_calib_rest.sh | 4 - 16 files changed, 616 insertions(+), 2730 deletions(-) diff --git a/crates/hipfire-arch-gemma4/examples/prefill_parity_gemma4.rs b/crates/hipfire-arch-gemma4/examples/prefill_parity_gemma4.rs index 366ca7080c..b9040366dd 100644 --- a/crates/hipfire-arch-gemma4/examples/prefill_parity_gemma4.rs +++ b/crates/hipfire-arch-gemma4/examples/prefill_parity_gemma4.rs @@ -5,17 +5,14 @@ //! Batched-vs-per-token prefill parity harness for the gemma4 lowered path. //! //! Runs the SAME prompt through (A) per-token `forward_scratch` and -//! (B) `forward_prefill_batch` with fresh KV each, compares prefill logits +//! (B) `forward_prefill_batch` over all but the last token plus one +//! `forward_scratch`, with fresh KV each; compares the prompt's last logits //! (max-abs diff, argmax) and an N-token greedy continuation (per-token decode -//! from both prefill states). Uses Q8 KV for both attention tiers, matching -//! the daemon's explicit `--kv-mode q8` route. +//! from both prefill states). Uses unwindowed Q8 KV for both attention tiers, +//! the layout batched prefill writes. //! //! Usage: //! prefill_parity_gemma4 --model [--prompt ] [--decode N] -//! -//! Env: HIPFIRE_GEMMA4_DUMP=1 + HIPFIRE_BATCHED_PREFILL / HIPFIRE_WMMA_PREFILL -//! affect only the lowered internals (run_prefill_gemm WMMA arm); this harness -//! calls forward_prefill_batch unconditionally for run B. #[cfg(not(feature = "deltanet"))] fn main() { @@ -76,11 +73,7 @@ fn main() { let scratch = lowered::Gemma4Scratch::new(&mut gpu, &cfg, max_seq).expect("scratch"); lowered::init_scratch_constants(&mut gpu, &scratch, cfg.full_head_dim).expect("scratch consts"); - eprintln!( - "prompt tokens = {} (chunk={})", - ids.len(), - scratch.max_prefill_batch - ); + eprintln!("prompt tokens = {}", ids.len()); let fnv = |bytes: &[u8]| -> u64 { let mut h: u64 = 0xcbf29ce484222325; @@ -103,13 +96,12 @@ fn main() { }; let mut run = |label: &str, batched: bool| -> (Vec, Vec) { - let mut kv_sliding = KvCache::new_gpu_q8_capped( + let mut kv_sliding = KvCache::new_gpu_q8( &mut gpu, cfg.n_layers, cfg.sliding_n_kv_heads, cfg.sliding_head_dim, max_seq, - cfg.sliding_window, ) .expect("kv sliding"); let mut kv_full = KvCache::new_gpu_q8( @@ -122,27 +114,29 @@ fn main() { .expect("kv full"); let t0 = std::time::Instant::now(); if batched { - let chunk = std::env::var("HIPFIRE_PREFILL_CHUNK") - .ok() - .and_then(|v| v.parse().ok()) - .unwrap_or(scratch.max_prefill_batch) - .max(1); - let mut off = 0usize; - while off < ids.len() { - let end = (off + chunk).min(ids.len()); - lowered::forward_prefill_batch( - &mut gpu, - &weights, - &cfg, - &ids[off..end], - off, - &mut kv_sliding, - &mut kv_full, - &scratch, - ) - .expect("batched prefill"); - off = end; - } + let last = ids.len() - 1; + lowered::forward_prefill_batch( + &mut gpu, + &weights, + &cfg, + &ids[..last], + 0, + &mut kv_sliding, + &mut kv_full, + &scratch, + ) + .expect("batched prefill"); + lowered::forward_scratch( + &mut gpu, + &weights, + &cfg, + ids[last], + last, + &mut kv_sliding, + &mut kv_full, + &scratch, + ) + .expect("last prompt token"); } else { for (p, &t) in ids.iter().enumerate() { lowered::forward_scratch( diff --git a/crates/hipfire-arch-gemma4/src/carrier.rs b/crates/hipfire-arch-gemma4/src/carrier.rs index 8dd61ce61a..97a0363f24 100644 --- a/crates/hipfire-arch-gemma4/src/carrier.rs +++ b/crates/hipfire-arch-gemma4/src/carrier.rs @@ -21,21 +21,11 @@ use rdna_compute::Gpu; // ─── Helpers moved verbatim from carriers.rs ───────────────────────────── -/// Route selection shared by the carrier load path and source admission. -pub fn gemma4_use_lowered( - enable_moe_block: bool, - want_batched: bool, - has_drafter: bool, - is_e_series: bool, -) -> bool { - enable_moe_block || (want_batched && !has_drafter && !is_e_series) -} - /// Pure context admission for Gemma 4. Refuses the lowered route when /// `max_seq < 128`. Eager loads are unrestricted (no eager assert exists). /// -/// Callers compute `use_lowered` via [`gemma4_use_lowered`] / source probes so -/// the refusal string is byte-identical at admission and carrier layers. +/// Callers compute `use_lowered` via [`gemma4_source_uses_lowered`] so the +/// refusal string is byte-identical at admission and carrier layers. pub fn gemma4_context_admission(max_seq: usize, use_lowered: bool) -> Result<(), String> { if use_lowered && max_seq < 128 { return Err(format!( @@ -45,29 +35,10 @@ pub fn gemma4_context_admission(max_seq: usize, use_lowered: bool) -> Result<(), Ok(()) } -/// Mirror carrier route selection from an already-open HFQ source + env gates. -/// `has_drafter` is the EAGLE `params.drafter` presence (not DFlash draft). -pub fn gemma4_source_uses_lowered(hfq: &hipfire_runtime::hfq::HfqFile, has_drafter: bool) -> bool { - let lowered_cfg = lowered::config_from_hfq(hfq); - let want_batched = lowered::batched_prefill_enabled() || lowered::wmma_prefill_enabled(); - let Some(lcfg) = &lowered_cfg else { - return false; - }; - let lowered_is_moe = lcfg.enable_moe_block; - let is_e_series = if lowered_is_moe { - false - } else { - match Gemma4Config::from_hfq(hfq) { - Ok(cfg) => cfg.hidden_size_per_layer_input != 0 || cfg.num_kv_shared_layers != 0, - Err(_) => false, - } - }; - gemma4_use_lowered( - lcfg.enable_moe_block, - want_batched, - has_drafter, - is_e_series, - ) +/// MoE checkpoints run on the lowered stack; every other Gemma 4 load is +/// eager. +pub fn gemma4_source_uses_lowered(hfq: &hipfire_runtime::hfq::HfqFile) -> bool { + lowered::config_from_hfq(hfq).is_some_and(|cfg| cfg.enable_moe_block) } fn gemma4_validate_drafter_route(is_e_series: bool, has_drafter: bool) -> Result<(), String> { @@ -193,9 +164,9 @@ fn resolve_lowered_full_kv_mode( /// Build the Gemma 4 GPU bundle from an HFQ source. /// /// `ModelSource::Dir` returns the same error string the carrier previously -/// emitted inline. HFQ path is verbatim: lowered/eager selection, -/// `want_batched` env gate, E-series validation, weight/state/KV allocation, -/// and the preserved `eprintln!` diagnostics for the chosen path. +/// emitted inline. HFQ path: lowered (MoE) / eager selection, E-series +/// validation, weight/state/KV allocation, and the `eprintln!` diagnostics +/// for the chosen path. pub fn load_gemma4_bundle(src: ModelSource, ctx: &mut LoadCtx) -> Result { if ctx.kv_backend != hipfire_runtime::kv_backend::KvBackend::Legacy { return Err("gemma4: sliding/full KV owners require legacy backend".into()); @@ -207,17 +178,11 @@ pub fn load_gemma4_bundle(src: ModelSource, ctx: &mut LoadCtx) -> Result Result Result 1 && ple.model_projection.gpu_dtype == DType::Q8_0 { + let ctx = DispatchCtx::new(gpu); hipfire_dispatch::pipeline::sandwich::gemm_rows( gpu, + &ctx, &ple.model_projection.dispatch_ref(), x, projection_all, @@ -608,48 +611,6 @@ pub fn supports_batched_prefill(weights: &Gemma4Weights) -> bool { }) } -struct BatchScratchLedger { - gpu: *mut Gpu, - tensors: Vec, -} - -impl BatchScratchLedger { - fn new(gpu: &mut Gpu) -> Self { - Self { - gpu, - tensors: Vec::with_capacity(24), - } - } - - fn alloc(&mut self, n: usize, label: &str) -> Result { - // SAFETY: the ledger is scoped inside forward_batch_spec and the Gpu - // reference outlives it. No access occurs after the function returns. - let gpu = unsafe { &mut *self.gpu }; - let tensor = gpu - .alloc_tensor(&[n], DType::F32) - .map_err(|e| format!("gemma4 forward_batch alloc {label}: {e:?}"))?; - let view = GpuTensor { - // SAFETY: the owning tensor remains in this ledger until all views - // have dropped and the ledger frees it exactly once. - buf: unsafe { tensor.buf.alias() }, - shape: tensor.shape.clone(), - dtype: tensor.dtype, - }; - self.tensors.push(tensor); - Ok(view) - } -} - -impl Drop for BatchScratchLedger { - fn drop(&mut self) { - // SAFETY: see new(); the ledger cannot outlive forward_batch_spec's Gpu. - let gpu = unsafe { &mut *self.gpu }; - for tensor in self.tensors.drain(..) { - gpu.free_tensor(tensor).ok(); - } - } -} - /// argmax over a logits row (spec-decode greedy per-position prediction). fn argmax_f32_row(v: &[f32]) -> u32 { let mut bi = 0u32; @@ -741,61 +702,22 @@ pub fn forward_batch_spec( // seq_len after this batch = absolute positions [start_pos, start_pos+B). checked_batch_seq_len(start_pos, b, state.max_seq)?; - let mut scratch = BatchScratchLedger::new(gpu); - let mut alloc = |_g: &mut Gpu, n: usize, label: &str| scratch.alloc(n, label); - - // ── Batched scratch (per call; verify is not the hot per-kernel loop). ── - let x = alloc(gpu, b * dim, "x")?; - let residual = alloc(gpu, b * dim, "residual")?; - let nrm = alloc(gpu, b * dim, "nrm")?; // rmsnorm output / shared proj input - // x_rot is the SHARED FWHT-rotate scratch for every projection in the layer - // (proj_gemm_batched rotates b*w.k floats into it). The largest w.k is the - // o_proj on FULL attention layers: k = n_heads*full_head_dim (= max_q here), - // which exceeds dim. Size for the max so the o_proj rotation can't OOB-write. - let x_rot = alloc(gpu, b * max_q.max(dim), "x_rot")?; // FWHT scratch (MQ4 proj path) - let q = alloc(gpu, b * max_q, "q")?; - let k = alloc(gpu, b * max_kv, "k")?; - let v = alloc(gpu, b * max_kv, "v")?; - let attn_out = alloc(gpu, b * max_q, "attn_out")?; - let gate_ffn = alloc(gpu, b * ffn_hd, "gate_ffn")?; - let up_ffn = alloc(gpu, b * ffn_hd, "up_ffn")?; - let ffn_hidden = alloc(gpu, b * ffn_hd, "ffn_hidden")?; - let ffn_out = alloc(gpu, b * dim, "ffn_out")?; - // FFN rotation scratch must hold B*ffn_hd for the down_proj input. - let ffn_rot = alloc(gpu, b * ffn_hd, "ffn_rot")?; - let ple_token_inputs = if ple_dim != 0 { - Some(alloc(gpu, b * ple_packed, "ple_token_inputs")?) - } else { - None + let mut bufs = RowBuffers::new(gpu); + let widths = RowWidths { + dim, + max_q, + max_kv, + ffn: ffn_hd, }; - let ple_projection_all = if ple_dim != 0 { - Some(alloc(gpu, b * ple_packed, "ple_projection_all")?) - } else { - None - }; - let ple_gate = if ple_dim != 0 { - Some(alloc(gpu, b * ple_dim, "ple_gate")?) - } else { - None - }; - let ple_hidden = if ple_dim != 0 { - Some(alloc(gpu, b * ple_dim, "ple_hidden")?) - } else { - None - }; - let ple_out = if ple_dim != 0 { - Some(alloc(gpu, b * dim, "ple_out")?) - } else { - None - }; - - // positions [B] i32 (kernels read this buffer as i32). - let pos_data: Vec = (0..b).map(|i| (start_pos + i) as i32).collect(); - let pos_bytes: Vec = pos_data.iter().flat_map(|p| p.to_ne_bytes()).collect(); - let pos_array = alloc(gpu, b, "pos_array")?; - gpu.hip - .memcpy_htod(&pos_array.buf, &pos_bytes) - .map_err(|e| format!("gemma4 forward_batch htod pos: {e:?}"))?; + let acts = RowActivations::new(&mut bufs, gpu, b, start_pos, widths)?; + let (x, x_rot) = (&acts.x, &acts.attn_rot); + let mut ple_buf = + |n: usize, label: &str| (ple_dim != 0).then(|| bufs.alloc(n, label)).transpose(); + let ple_token_inputs = ple_buf(b * ple_packed, "ple_token_inputs")?; + let ple_projection_all = ple_buf(b * ple_packed, "ple_projection_all")?; + let ple_gate = ple_buf(b * ple_dim, "ple_gate")?; + let ple_hidden = ple_buf(b * ple_dim, "ple_hidden")?; + let ple_out = ple_buf(b * dim, "ple_out")?; let batched_embedding_requested = supports_gemma4_batched_prefill_arch(&gpu.arch) && gpu.flags.gemma4_batched_embedding_prefill @@ -810,7 +732,7 @@ pub fn forward_batch_spec( let token_bytes: &[u8] = unsafe { std::slice::from_raw_parts(token_data.as_ptr() as *const u8, token_data.len() * 4) }; - let token_ids = alloc(gpu, b, "token_ids")?; + let token_ids = bufs.alloc(b, "token_ids")?; gpu.hip .memcpy_htod(&token_ids.buf, token_bytes) .map_err(|e| format!("gemma4 forward_batch htod token ids: {e:?}"))?; @@ -825,7 +747,7 @@ pub fn forward_batch_spec( gpu, weights.embd_format, &weights.embed_tokens, - &x, + x, token_ids, b, dim, @@ -835,7 +757,7 @@ pub fn forward_batch_spec( false }; if !batched_embedding { - let x_single = alloc(gpu, dim, "x_single")?; + let x_single = bufs.alloc(dim, "x_single")?; for (i, &tok) in tokens.iter().enumerate() { embed_lookup_row(cfg, weights, gpu, &x_single, tok)?; gpu.hip @@ -844,7 +766,7 @@ pub fn forward_batch_spec( } } // √dim scale on the whole [B*dim] buffer (uniform — matches eager scale_f32). - gpu.scale_f32(&x, cfg.embed_scale) + gpu.scale_f32(x, cfg.embed_scale) .map_err(|e| format!("gemma4 forward_batch embed scale: {e:?}"))?; if ple_dim != 0 { prepare_per_layer_inputs_batched( @@ -853,8 +775,8 @@ pub fn forward_batch_spec( gpu, tokens, token_ids.as_ref(), - &x, - &x_rot, + x, + x_rot, ple_token_inputs.as_ref().unwrap(), ple_projection_all.as_ref().unwrap(), )?; @@ -880,24 +802,8 @@ pub fn forward_batch_spec( resident: eager_resident(state), rows: b, position: start_pos, - positions: Some(&pos_array), - scratch: LayerScratch { - x: &x, - residual: &residual, - normed: &nrm, - attn_rot: &x_rot, - mlp_rot: &ffn_rot, - q: &q, - k: &k, - v: &v, - attn_out: &attn_out, - gate: &gate_ffn, - up: &up_ffn, - act: &ffn_hidden, - mlp_out: &ffn_out, - ple, - moe: None, - }, + positions: Some(&acts.positions), + scratch: acts.scratch(ple, None), }; let mut steps = Vec::with_capacity(4 * cfg.n_layers); for (layer_idx, layer) in weights.layers.iter().enumerate() { @@ -935,14 +841,14 @@ pub fn forward_batch_spec( let local_hidden = if want_hidden { None } else { - Some(alloc(gpu, b * dim, "spec_normed")?) + Some(bufs.alloc(b * dim, "spec_normed")?) }; let normed_hidden: &GpuTensor = match per_token_hidden_out { Some(h) => h, None => local_hidden.as_ref().unwrap(), }; // Batched final RMSNorm over all B rows. - gpu.rmsnorm_batched(&x, &weights.final_norm, normed_hidden, b, dim, eps) + gpu.rmsnorm_batched(x, &weights.final_norm, normed_hidden, b, dim, eps) .map_err(|e| format!("gemma4 forward_batch_spec final rmsnorm: {e:?}"))?; if let Some(out) = per_pos_argmax_out { @@ -961,7 +867,7 @@ pub fn forward_batch_spec( // Softcap is SKIPPED: it's a strictly-monotonic per-element map // (tanh-scaled), so argmax(softcap(z)) == argmax(z). The accept // decision is an argmax, so this is bit-exact in the decision. - let logits_b = alloc(gpu, b * vocab, "spec_logits_b")?; + let logits_b = bufs.alloc(b * vocab, "spec_logits_b")?; let lm_result = if eagle_strict_enabled() { gpu.gemm_q8_0_batched( &weights.lm_head.buf, @@ -984,7 +890,7 @@ pub fn forward_batch_spec( lm_result .map_err(|e| format!("gemma4 forward_batch_spec batched lm_head: {e:?}"))?; // GPU per-row argmax over [B, vocab]; only B indices land on PCIe. - let idx_buf = alloc(gpu, b, "spec_argmax_idx")?; + let idx_buf = bufs.alloc(b, "spec_argmax_idx")?; gpu.argmax_f32_batched(&logits_b, &idx_buf, vocab, b) .map_err(|e| format!("gemma4 forward_batch_spec batched argmax: {e:?}"))?; let mut idx_i32 = vec![0i32; b]; @@ -1026,7 +932,7 @@ pub fn forward_batch_spec( // ── Final RMSNorm + tied lm_head on the LAST row only (verify needs the // last position's logits). ── - let x_last = alloc(gpu, dim, "x_last")?; + let x_last = bufs.alloc(dim, "x_last")?; gpu.hip .memcpy_dtod_at(&x_last.buf, 0, &x.buf, (b - 1) * dim * 4, dim * 4) .map_err(|e| format!("gemma4 forward_batch last copy: {e:?}"))?; diff --git a/crates/hipfire-arch-gemma4/src/gemma4.rs b/crates/hipfire-arch-gemma4/src/gemma4.rs index 224c40f460..a289235fa8 100644 --- a/crates/hipfire-arch-gemma4/src/gemma4.rs +++ b/crates/hipfire-arch-gemma4/src/gemma4.rs @@ -800,24 +800,12 @@ fn build_kv_slot_map(cfg: &Gemma4Config) -> Result, String> { } fn q8_flash_partials_len(gpu: &Gpu, cfg: &Gemma4Config, max_seq: usize) -> usize { - [ - (cfg.sliding_n_kv_heads, cfg.sliding_head_dim), - (cfg.full_n_kv_heads, cfg.full_head_dim), - ] - .into_iter() - .map(|(n_kv_heads, head_dim)| { - let tile = rdna_compute::attention::q8_flash_tile_size( - &gpu.arch, - cfg.n_heads, - n_kv_heads, - head_dim, - max_seq, - ); - let max_tiles = max_seq.div_ceil(tile); - GEMMA4_FORWARD_BATCH_MAX * cfg.n_heads * max_tiles * (2 + head_dim) - }) - .max() - .unwrap_or(0) + crate::program::q8_flash_partials_len( + gpu, + &crate::program::Geometry::eager(cfg), + max_seq, + GEMMA4_FORWARD_BATCH_MAX, + ) } impl Gemma4State { diff --git a/crates/hipfire-arch-gemma4/src/lib.rs b/crates/hipfire-arch-gemma4/src/lib.rs index 060b414679..0dccd62756 100644 --- a/crates/hipfire-arch-gemma4/src/lib.rs +++ b/crates/hipfire-arch-gemma4/src/lib.rs @@ -37,8 +37,8 @@ pub mod lowered; mod program; pub mod speculative; pub use carrier::{ - gemma4_context_admission, gemma4_source_uses_lowered, gemma4_use_lowered, load_gemma4_bundle, - Gemma4Bundle, Gemma4EagerBundle, Gemma4LoweredBundle, + gemma4_context_admission, gemma4_source_uses_lowered, load_gemma4_bundle, Gemma4Bundle, + Gemma4EagerBundle, Gemma4LoweredBundle, }; pub use arch::{Gemma4, ARCH_ID}; diff --git a/crates/hipfire-arch-gemma4/src/lowered.rs b/crates/hipfire-arch-gemma4/src/lowered.rs index 02dc6d7821..9bdc7113c1 100644 --- a/crates/hipfire-arch-gemma4/src/lowered.rs +++ b/crates/hipfire-arch-gemma4/src/lowered.rs @@ -20,354 +20,12 @@ use hip_bridge::HipResult; use hipfire_dispatch::context::DispatchCtx; -use hipfire_dispatch::families::attention::AttnParams; -use hipfire_dispatch::families::gemm::GemmParams; -use hipfire_dispatch::families::gemv::WeightRef; -use hipfire_dispatch::families::kv_tier::{KvTierInputs, KvTierPlan}; -use hipfire_dispatch::pipeline::{execute_steps, run_uniform_moe_gate_up, GemvInput, Step}; +use hipfire_dispatch::pipeline::execute_steps; use hipfire_runtime::hfq::HfqFile; -use hipfire_runtime::llama::{self, f16_to_f32, weight_gemv, EmbeddingFormat, WeightTensor}; +use hipfire_runtime::llama::{f16_to_f32, EmbeddingFormat, WeightTensor}; use hipfire_runtime::weight_store::upload_pooled_bytes; use rdna_compute::{DType, Gpu, GpuTensor}; -/// #397 Ship 5.2: route a single PLAIN-batched prefill GEMM through -/// [`GemmFamily::run_key`] against an explicit dispatcher-entry key. -/// -/// Byte-identical to the old `weight_gemm()` call for every (dtype × arch). -/// Only HFQ4G256, HFQ4G128 have batched kernels; other dtypes fall back to -/// the repeated-GEMV loop (same as `weight_gemm`'s fallback). -#[inline] -/// Env gate for WMMA prefill. Default OFF because WMMA F16 input quantization -/// loses ~3 mantissa bits vs F32 scalar, so WMMA results are NOT byte-identical -/// to the scalar GEMV path. Set HIPFIRE_WMMA_PREFILL=1 to opt in. -pub fn wmma_prefill_enabled() -> bool { - static GATE: std::sync::OnceLock = std::sync::OnceLock::new(); - *GATE.get_or_init(|| { - hipfire_config::developer_var("HIPFIRE_WMMA_PREFILL").map_or(false, |v| v == "1") - }) -} - -/// Env gate for batched prefill (v2). Independent from WMMA. -/// Set HIPFIRE_BATCHED_PREFILL=1 to use batched projections + per-token attention. -pub fn batched_prefill_enabled() -> bool { - static GATE: std::sync::OnceLock = std::sync::OnceLock::new(); - *GATE.get_or_init(|| { - hipfire_config::developer_var("HIPFIRE_BATCHED_PREFILL").map_or(false, |v| v == "1") - }) -} - -/// Batched GEMM for prefill projections. -/// -/// Scalar path: routes through `GemmFamily::run_key` with the appropriate -/// kernel key for the weight dtype. -/// -/// WMMA path (opt-in via `HIPFIRE_WMMA_PREFILL=1`): routes through -/// `GemmFamily::run` which resolves WMMA-eligible dtypes to their WMMA -/// kernel variant. The WMMA path converts F32 input to F16 automatically -/// (Bug 1 fix: `gemm_hfq4g256_wmma` now calls `ensure_fp16_x`). -/// -/// MQ4G256 weights (Bug 2 fix): `GemmFamily::resolve` now maps -/// `DType::MQ4G256` → `GemmHfq4G256Wmma` / `GemmHfq4G256`, sharing the -/// same kernel binary since the layout is identical (136 B/group). -/// -/// For dtypes that have no GEMM kernel, falls back to per-token GEMV. -fn run_prefill_gemm( - gpu: &mut Gpu, - w: &WeightTensor, - x: &GpuTensor, - y: &GpuTensor, - batch_size: usize, - bf16_scratch: Option<&GpuTensor>, -) -> HipResult<()> { - // MQ4G256 weights are FWHT-rotated at quantize time. The per-token GEMV - // path (weight_gemv) FWHT-rotates x before running the prerotated kernel; - // the batched GEMM kernels (scalar gemm_hfq4g256 AND its WMMA siblings) - // decode the same 136 B/group layout but do NOT rotate internally, so the - // input must be pre-rotated here -- exactly like the eager - // forward.rs::proj_gemm_batched (EAGLE verify) does. Skipping this fed raw - // activations to the GEMM and produced garbage logits from token 0 on - // every arch (root-caused on gfx1201 2026-06-10; gfx11 batched-prefill - // validation predated the *.union.mq4 files, which is why it never fired - // there). - let x_rot_holder = if matches!(w.gpu_dtype, DType::MQ4G256) { - let xr = gpu.alloc_tensor(&[batch_size, w.k], DType::F32)?; - llama::rotate_x_mq_batched_for(gpu, w, x, &xr, w.k, batch_size)?; - Some(xr) - } else { - None - }; - let x_gemm: &GpuTensor = x_rot_holder.as_ref().unwrap_or(x); - - let result = run_prefill_gemm_inner(gpu, w, x, x_gemm, y, batch_size, bf16_scratch); - if let Some(t) = x_rot_holder { - gpu.free_tensor(t)?; - } - result -} - -/// Inner body of `run_prefill_gemm`: `x_gemm` is the (possibly FWHT-rotated) -/// GEMM input; `x_raw` is the original activation, used only by the per-token -/// GEMV fallback + verify hook (both rotate internally via ). -fn run_prefill_gemm_inner( - gpu: &mut Gpu, - w: &WeightTensor, - x_raw: &GpuTensor, - x_gemm: &GpuTensor, - y: &GpuTensor, - batch_size: usize, - bf16_scratch: Option<&GpuTensor>, -) -> HipResult<()> { - // BF16 MFMA path — calibration override, DEFAULT path (no WMMA gate). - // Stages F32 activation to BF16 via persistent scratch, then dispatches - // GemmBf16Mfma. The wrapper refuses F32 x, so this must be BF16. - if w.gpu_dtype == DType::BF16 { - let nelems = batch_size * w.k; - if x_gemm.dtype == DType::BF16 { - // Hoisted: caller already staged to BF16 (q/k/v or gate+up reuse). - // F32 source was captured at the hoist site before staging, so the - // run_key tap (which skips BF16 x) is bypassed correctly there. - let w_ref = WeightRef { - buf: &w.buf, - dtype: w.gpu_dtype, - m: w.m, - k: w.k, - row_stride: w.k, - rotation: None, - awq_scale: None, - lloyd_lut_e4m3: None, - lloyd_lut_f16: None, - lloyd_lut_c16: None, - }; - let ctx = DispatchCtx::new(gpu); - let params = GemmParams { - w: &w_ref, - x: x_gemm, - y, - batch_size, - }; - let family = llama::gemm_family(); - family - .run_key( - hipfire_dispatch::types::KernelKey::GemmBf16Mfma, - &ctx, - gpu, - ¶ms, - ) - .map_err(|e| hip_bridge::HipError::new(0, &e.to_string()))?; - return Ok(()); - } - // Non-hoisted BF16 path: capture the true F32 activation before - // staging to BF16. The GemmFamily::run_key tap deliberately skips - // BF16 x (would overrun and record rounded values), so this explicit - // F32 capture restores coverage. Zero-cost when unarmed. - gpu.maybe_capture_activation(&w.buf, x_gemm, batch_size, w.k); - if let Some(scratch) = bf16_scratch { - let bf16_view = scratch.sub_offset(0, nelems); - gpu.convert_f32_to_bf16(x_gemm, &bf16_view, nelems)?; - let w_ref = WeightRef { - buf: &w.buf, - dtype: w.gpu_dtype, - m: w.m, - k: w.k, - row_stride: w.k, - rotation: None, - awq_scale: None, - lloyd_lut_e4m3: None, - lloyd_lut_f16: None, - lloyd_lut_c16: None, - }; - let ctx = DispatchCtx::new(gpu); - let params = GemmParams { - w: &w_ref, - x: &bf16_view, - y, - batch_size, - }; - let family = llama::gemm_family(); - family - .run_key( - hipfire_dispatch::types::KernelKey::GemmBf16Mfma, - &ctx, - gpu, - ¶ms, - ) - .map_err(|e| hip_bridge::HipError::new(0, &e.to_string()))?; - return Ok(()); - } - // No persistent scratch — fallback per-call alloc (not hoisted, but functional). - let tmp = gpu.alloc_tensor(&[nelems], DType::BF16)?; - gpu.convert_f32_to_bf16(x_gemm, &tmp, nelems)?; - let w_ref = WeightRef { - buf: &w.buf, - dtype: w.gpu_dtype, - m: w.m, - k: w.k, - row_stride: w.k, - rotation: None, - awq_scale: None, - lloyd_lut_e4m3: None, - lloyd_lut_f16: None, - lloyd_lut_c16: None, - }; - let ctx = DispatchCtx::new(gpu); - let params = GemmParams { - w: &w_ref, - x: &tmp, - y, - batch_size, - }; - let family = llama::gemm_family(); - let res = family - .run_key( - hipfire_dispatch::types::KernelKey::GemmBf16Mfma, - &ctx, - gpu, - ¶ms, - ) - .map_err(|e| hip_bridge::HipError::new(0, &e.to_string())); - gpu.free_tensor(tmp)?; - return res; - } - let ctx = DispatchCtx::new(gpu); - let w_ref = WeightRef { - buf: &w.buf, - dtype: w.gpu_dtype, - m: w.m, - k: w.k, - row_stride: w.k, - rotation: None, - awq_scale: None, - lloyd_lut_e4m3: None, - lloyd_lut_f16: None, - lloyd_lut_c16: None, - }; - let params = GemmParams { - w: &w_ref, - x: x_gemm, - y, - batch_size, - }; - let family = llama::gemm_family(); - - if wmma_prefill_enabled() { - // WMMA path: use GemmFamily::run which resolves the best kernel - // for each dtype (WMMA where available, scalar fallback otherwise). - // The F32->F16 conversion is handled inside the kernel methods. - // - // HARD ERROR on failure -- no silent fallthrough. The old code fell - // through to the scalar key path when errored (e.g. missing - // WMMA kernel for this dtype x arch), which silently produced - // special-token garbage from token 0 on gfx1201 (RDNA4) where the - // gfx11 WMMA kernels do not exist. Refuse-dont-degrade: the - // operator opted into WMMA explicitly, so a missing kernel is an - // error, not a degrade. Unset HIPFIRE_WMMA_PREFILL (or use the - // scalar HIPFIRE_BATCHED_PREFILL=1 path) on archs without coverage. - return family.run(&ctx, gpu, ¶ms).map_err(|e| { - hip_bridge::HipError::new(0, &format!( - "gemma4 WMMA prefill: GemmFamily::run failed for dtype {:?} on arch {} (HIPFIRE_WMMA_PREFILL=1 requires a WMMA kernel for every projection dtype -- unset it to use the scalar batched path): {e}", - w.gpu_dtype, gpu.arch, - )) - }); - } - - // Scalar path (default) or fallback for unsupported WMMA dtypes. - let key = match w.gpu_dtype { - DType::HFQ4G256 => hipfire_dispatch::types::KernelKey::GemmHfq4G256, - DType::HFQ4G128 => hipfire_dispatch::types::KernelKey::GemmHfq4G128, - // Same kernel as HFQ4G256 (layout-identical); input pre-rotated above. - DType::MQ4G256 => hipfire_dispatch::types::KernelKey::GemmHfq4G256, - // Explicit F32 batched path: default Q8 WMMA rounds F32 inputs and - // dequant weights to F16, which breaks Gemma prefill/decode continuation. - DType::Q8_0 => hipfire_dispatch::types::KernelKey::GemmQ8_0BatchedF32Chunked, - // No batched GEMM kernel for this dtype -- fall back to repeated GEMV - // on the RAW input (weight_gemv applies any needed rotation itself). - // This matches the old fallback path. - _ => { - let x_tok = gpu.alloc_tensor(&[w.k], DType::F32)?; - let y_tok = gpu.alloc_tensor(&[w.m], DType::F32)?; - for b in 0..batch_size { - gpu.hip - .memcpy_dtod_at(&x_tok.buf, 0, &x_raw.buf, b * w.k * 4, w.k * 4)?; - weight_gemv(gpu, w, &x_tok, &y_tok)?; - gpu.hip - .memcpy_dtod_at(&y.buf, b * w.m * 4, &y_tok.buf, 0, w.m * 4)?; - } - gpu.free_tensor(x_tok)?; - gpu.free_tensor(y_tok)?; - return Ok(()); - } - }; - family - .run_key(key, &ctx, gpu, ¶ms) - .map_err(|e| hip_bridge::HipError::new(0, &e.to_string()))?; - // Debug parity hook: re-run token 0 through the per-token GEMV path - // (raw input; weight_gemv rotates internally) and diff against the GEMM. - if hipfire_config::developer_var("HIPFIRE_GEMMA4_GEMM_VERIFY") - .ok() - .as_deref() - == Some("1") - { - let x_tok = gpu.alloc_tensor(&[w.k], DType::F32)?; - let y_tok = gpu.alloc_tensor(&[w.m], DType::F32)?; - gpu.hip - .memcpy_dtod_at(&x_tok.buf, 0, &x_raw.buf, 0, w.k * 4)?; - weight_gemv(gpu, w, &x_tok, &y_tok)?; - let yv = gpu.download_f32(&y_tok)?; - let yg = gpu.download_f32(y)?; - let mut worst = 0f32; - let mut wi = 0usize; - for i in 0..w.m { - let d = (yv[i] - yg[i]).abs(); - if d > worst { - worst = d; - wi = i; - } - } - eprintln!( - "[gemm-verify] dtype={:?} m={} k={} b={} key={:?} worst={:.5} at {} gemv={:.4} gemm={:.4} head_gemv={:?} head_gemm={:?}", - w.gpu_dtype, - w.m, - w.k, - batch_size, - key, - worst, - wi, - yv[wi], - yg[wi], - &yv[..2], - &yg[..2] - ); - gpu.free_tensor(x_tok)?; - gpu.free_tensor(y_tok)?; - } - Ok(()) -} - -/// Env-gated dump helper for v1-vs-v2 root-cause work. -/// Set HIPFIRE_GEMMA4_DUMP=1 to enable. Prints first 4 floats + sum + nan/inf count. -#[allow(dead_code)] -fn dbg_dump(gpu: &mut Gpu, label: &str, t: &GpuTensor, take: usize) { - if hipfire_config::developer_var("HIPFIRE_GEMMA4_DUMP") - .ok() - .as_deref() - != Some("1") - { - return; - } - let data = match gpu.download_f32(t) { - Ok(d) => d, - Err(_) => return, - }; - let take = take.min(data.len()); - let head: Vec = data[..take.min(4)].iter().copied().collect(); - let sum: f64 = data[..take].iter().map(|&v| v as f64).sum(); - let nans = data[..take].iter().filter(|&&v| v.is_nan()).count(); - let infs = data[..take].iter().filter(|&&v| v.is_infinite()).count(); - eprintln!( - "[dump] {label:42} sum={sum:>+14.4e} n={take:>6} head={head:?} nan={nans} inf={infs}" - ); -} - // ─── Config ───────────────────────────────────────────────────────────── #[derive(Debug, Clone, Copy, PartialEq)] @@ -2211,9 +1869,6 @@ use hip_bridge::DeviceBuffer; /// Flash tile size for gemma4 lowered path. Matches the HIP partition kernel's /// `TILE_SIZE = 128`. pub const GEMMA4_FLASH_TILE: usize = 128; -/// Max prefill batch size for lowered gemma4. Sized once; batch flash partials -/// scale linearly with this. -pub const GEMMA4_MAX_PREFILL_BATCH: usize = 128; /// Pure geometry: single-query flash partial length for `max_seq`. /// `n_heads * ceil(max_seq / TILE) * (2 + head_dim)` floats. @@ -2223,23 +1878,12 @@ pub fn gemma4_flash_partials_len(max_seq: usize, n_heads: usize, full_head_dim: n_heads * tiles * (2 + full_head_dim) } -/// Pure geometry: batched flash partial length for `max_seq`. -#[inline] -pub fn gemma4_pb_flash_partials_len(max_seq: usize, n_heads: usize, full_head_dim: usize) -> usize { - GEMMA4_MAX_PREFILL_BATCH * gemma4_flash_partials_len(max_seq, n_heads, full_head_dim) -} - /// Convenience for `Gemma4Config`. #[inline] pub fn gemma4_flash_partials_len_for_config(max_seq: usize, config: &Gemma4Config) -> usize { gemma4_flash_partials_len(max_seq, config.n_heads, config.full_head_dim) } -#[inline] -pub fn gemma4_pb_flash_partials_len_for_config(max_seq: usize, config: &Gemma4Config) -> usize { - gemma4_pb_flash_partials_len(max_seq, config.n_heads, config.full_head_dim) -} - /// Single-query flash partial length that covers every decode attend tier the /// lowered path can run on `arch`: the asym3 hd512 tile ([`GEMMA4_FLASH_TILE`]) /// and the Q8 decode tile (`q8_flash_tile_size`) for the full tier at `max_seq` @@ -2324,73 +1968,6 @@ pub struct Gemma4Scratch { pub moe_expert_gate_batch: GpuTensor, // [k_top × mi] pub moe_expert_up_batch: GpuTensor, // [k_top × mi] - // ── Prefill-batch (N tokens at a time) scratch ───────────────────── - // Sized for max_prefill_batch tokens. Used by forward_prefill_chunk to - // amortize MoE launch overhead across N tokens via the batched-indexed - // kernels. None on models with no MoE (n_experts==0). - pub max_prefill_batch: usize, - /// `[N × dim]` — per-token attention output after o_proj (before residual). - pub pb_attn_out: GpuTensor, - /// `[N × dim]` — per-token dense FFN output (gemma4 26B-A4B-it computes - /// this in parallel to MoE; both feed the sandwich norm). - pub pb_ffn_out: GpuTensor, - /// `[N × dim]` — pre-norm-2(attn_out) for each token. - pub pb_moe_pre2: GpuTensor, - /// `[N × dim]` — FWHT-rotated pre2 (MQ4 gate_up expects rotated x). - pub pb_moe_pre2_rot: GpuTensor, - /// `[N × dim]` — router input (rmsnorm(attn_out, router_scale) / sqrt(dim)). - pub pb_moe_router_in: GpuTensor, - /// `[N × n_experts]` — router logits batched. - pub pb_moe_router_logits: GpuTensor, - /// `[N × k_top]` i32 — top-K expert indices per token. - pub pb_moe_topk_indices: GpuTensor, - /// `[N × k_top]` — renormalized top-K weights per token. - pub pb_moe_topk_weights: GpuTensor, - /// `[n_exp + 1]` i32 — prefix-sum bucket offsets for routing-bucketed MoE GEMM. - pub pb_moe_expert_offsets: GpuTensor, - /// `[N × k_top]` i32 — packed (token_idx × k_top + krank), sorted by expert. - pub pb_moe_expert_token_list: GpuTensor, - /// `[N × k_top × mi]` — gate output across tokens × experts. - pub pb_moe_gate_batch: GpuTensor, - /// `[N × k_top × mi]` — up output across tokens × experts. - pub pb_moe_up_batch: GpuTensor, - /// `[N × k_top × mi]` — gelu_tanh(gate) * up. - pub pb_moe_hidden_batch: GpuTensor, - /// `[N × dim]` — accumulator for MoE branch output per token. - pub pb_moe_cur_moe: GpuTensor, - /// `[N × dim]` — post_norm_1(ffn_out) per token (dense branch). - pub pb_moe_cur_mlp: GpuTensor, - /// `[N × dim]` — residual stream per token across the prefill batch. - pub pb_residual: GpuTensor, - /// `[N × dim]` — post-rmsnorm input for batched projections. - pub pb_tmp: GpuTensor, - /// `[N × max_q_dim]` — batched Q projection output (sized for full layer = n_heads * full_head_dim). - pub pb_q: GpuTensor, - /// `[N × max_q_dim]` — batched attention output, kept separate from `pb_q` - /// so batched flash attention never reads and writes the same buffer. - pub pb_attn_q: GpuTensor, - /// `[N × n_heads × max_tiles_full × (2+full_head_dim)]` — batch-sized flash - /// partials for batched masked prefill attention. The single-query - /// `flash_partials` overflows with batch_size>1. - pub pb_flash_partials: GpuTensor, - /// `[N × max_kv_dim]` — batched K projection output. - pub pb_k: GpuTensor, - /// `[N × max_kv_dim]` — batched V projection output. - pub pb_v: GpuTensor, - /// `[N × hidden_dim]` — batched gate proj output. - pub pb_gate: GpuTensor, - /// `[N × hidden_dim]` — batched up proj output. - pub pb_up: GpuTensor, - /// `[N × hidden_dim]` — batched gelu(gate)*up. - pub pb_ffn_hidden: GpuTensor, - /// `[N]` — i32 position buffer for batched RoPE. - pub pb_positions: GpuTensor, - /// `[MAX_BATCH × maxK]` BF16 — persistent staging for BF16 MFMA prefill. - /// Sized once as `MAX_PREFILL_BATCH * max(dim, hidden_dim)` BF16 elements - /// (~5.5 MB on 31B). Used by `run_prefill_gemm_inner`'s BF16 arm via - /// `gpu.convert_f32_to_bf16(x, &pb_bf16_view, batch*k)`. Reused across - /// q/k/v and gate+up groups (hoisted once per shared x). - pub pb_bf16: GpuTensor, /// `[k_top × mi]` — gelu_tanh(gate)*up batched over k_top experts. pub moe_expert_hidden_batch: GpuTensor, } @@ -2398,8 +1975,8 @@ pub struct Gemma4Scratch { impl Gemma4Scratch { /// `max_seq` is the sole allocation authority — MUST equal `LoadCtx::max_seq` /// and the `KvCache::max_seq` / `physical_cap` for the paired caches. - /// Both `flash_partials` and `pb_flash_partials` are sized from this single - /// value; the `HIPFIRE_KV_SEQ` env var is no longer consulted. + /// `flash_partials` is sized from this single value; the `HIPFIRE_KV_SEQ` + /// env var is no longer consulted. pub fn new(gpu: &mut Gpu, config: &Gemma4Config, max_seq: usize) -> HipResult { Self::new_with_alloc( gpu, @@ -2552,47 +2129,6 @@ impl Gemma4Scratch { let i_moe_expert_up_batch = alloc!(&[k_top * mi], DType::F32); let i_moe_expert_hidden_batch = alloc!(&[k_top * mi], DType::F32); - // Prefill-batch scratch (N tokens at once). Larger batches expose - // more concurrent GPU work — total batch scratch ≈ N*0.16 MB. - let i_pb_attn_out = alloc!(&[GEMMA4_MAX_PREFILL_BATCH, dim], DType::F32); - let i_pb_ffn_out = alloc!(&[GEMMA4_MAX_PREFILL_BATCH, dim], DType::F32); - let i_pb_moe_pre2 = alloc!(&[GEMMA4_MAX_PREFILL_BATCH, dim], DType::F32); - let i_pb_moe_pre2_rot = alloc!(&[GEMMA4_MAX_PREFILL_BATCH, dim], DType::F32); - let i_pb_moe_router_in = alloc!(&[GEMMA4_MAX_PREFILL_BATCH, dim], DType::F32); - let i_pb_moe_router_logits = alloc!(&[GEMMA4_MAX_PREFILL_BATCH, n_exp], DType::F32); - let i_pb_moe_topk_indices = alloc!(&[GEMMA4_MAX_PREFILL_BATCH, k_top], DType::F32); - let i_pb_moe_topk_weights = alloc!(&[GEMMA4_MAX_PREFILL_BATCH, k_top], DType::F32); - // Routing-bucket scratch (Phase B). expert_offsets has n_exp+1 entries. - // expert_token_list has one entry per (token, krank) pair = N × k_top. - let i_pb_moe_expert_offsets = alloc!(&[n_exp + 1], DType::F32); // i32-typed slots - let i_pb_moe_expert_token_list = alloc!(&[GEMMA4_MAX_PREFILL_BATCH, k_top], DType::F32); - let i_pb_moe_gate_batch = alloc!(&[GEMMA4_MAX_PREFILL_BATCH, k_top * mi], DType::F32); - let i_pb_moe_up_batch = alloc!(&[GEMMA4_MAX_PREFILL_BATCH, k_top * mi], DType::F32); - let i_pb_moe_hidden_batch = alloc!(&[GEMMA4_MAX_PREFILL_BATCH, k_top * mi], DType::F32); - let i_pb_moe_cur_moe = alloc!(&[GEMMA4_MAX_PREFILL_BATCH, dim], DType::F32); - let i_pb_moe_cur_mlp = alloc!(&[GEMMA4_MAX_PREFILL_BATCH, dim], DType::F32); - let i_pb_residual = alloc!(&[GEMMA4_MAX_PREFILL_BATCH, dim], DType::F32); - let i_pb_tmp = alloc!(&[GEMMA4_MAX_PREFILL_BATCH, dim], DType::F32); - // Sized for max across sliding/full per-token vector dims. - // q_dim_max = n_heads * max(sliding_head_dim, full_head_dim) - let q_dim_max = config.n_heads * config.sliding_head_dim.max(config.full_head_dim); - let kv_dim_max = (config.sliding_n_kv_heads * config.sliding_head_dim) - .max(config.full_n_kv_heads * config.full_head_dim); - let i_pb_q = alloc!(&[GEMMA4_MAX_PREFILL_BATCH, q_dim_max], DType::F32); - let i_pb_attn_q = alloc!(&[GEMMA4_MAX_PREFILL_BATCH, q_dim_max], DType::F32); - let pb_flash_partials_sz = - gemma4_pb_flash_partials_len(max_seq, config.n_heads, config.full_head_dim); - let i_pb_flash_partials = alloc!(&[pb_flash_partials_sz], DType::F32); - let i_pb_k = alloc!(&[GEMMA4_MAX_PREFILL_BATCH, kv_dim_max], DType::F32); - let i_pb_v = alloc!(&[GEMMA4_MAX_PREFILL_BATCH, kv_dim_max], DType::F32); - let i_pb_gate = alloc!(&[GEMMA4_MAX_PREFILL_BATCH, config.hidden_dim], DType::F32); - let i_pb_up = alloc!(&[GEMMA4_MAX_PREFILL_BATCH, config.hidden_dim], DType::F32); - let i_pb_ffn_hidden = alloc!(&[GEMMA4_MAX_PREFILL_BATCH, config.hidden_dim], DType::F32); - let i_pb_positions = alloc!(&[GEMMA4_MAX_PREFILL_BATCH], DType::F32); // i32 packed in f32 slots - // BF16 staging for calibration MFMA: persistent, sized once. - let max_k_bf16 = config.dim.max(config.hidden_dim); - let i_pb_bf16 = alloc!(&[GEMMA4_MAX_PREFILL_BATCH * max_k_bf16], DType::BF16); - let scratch = Gemma4Scratch { x: take!(i_x), residual: take!(i_residual), @@ -2626,34 +2162,6 @@ impl Gemma4Scratch { moe_expert_gate_batch: take!(i_moe_expert_gate_batch), moe_expert_up_batch: take!(i_moe_expert_up_batch), moe_expert_hidden_batch: take!(i_moe_expert_hidden_batch), - max_prefill_batch: GEMMA4_MAX_PREFILL_BATCH, - pb_attn_out: take!(i_pb_attn_out), - pb_ffn_out: take!(i_pb_ffn_out), - pb_moe_pre2: take!(i_pb_moe_pre2), - pb_moe_pre2_rot: take!(i_pb_moe_pre2_rot), - pb_moe_router_in: take!(i_pb_moe_router_in), - pb_moe_router_logits: take!(i_pb_moe_router_logits), - pb_moe_topk_indices: take!(i_pb_moe_topk_indices), - pb_moe_topk_weights: take!(i_pb_moe_topk_weights), - pb_moe_expert_offsets: take!(i_pb_moe_expert_offsets), - pb_moe_expert_token_list: take!(i_pb_moe_expert_token_list), - pb_moe_gate_batch: take!(i_pb_moe_gate_batch), - pb_moe_up_batch: take!(i_pb_moe_up_batch), - pb_moe_hidden_batch: take!(i_pb_moe_hidden_batch), - pb_moe_cur_moe: take!(i_pb_moe_cur_moe), - pb_moe_cur_mlp: take!(i_pb_moe_cur_mlp), - pb_residual: take!(i_pb_residual), - pb_tmp: take!(i_pb_tmp), - pb_q: take!(i_pb_q), - pb_attn_q: take!(i_pb_attn_q), - pb_flash_partials: take!(i_pb_flash_partials), - pb_k: take!(i_pb_k), - pb_v: take!(i_pb_v), - pb_gate: take!(i_pb_gate), - pb_up: take!(i_pb_up), - pb_ffn_hidden: take!(i_pb_ffn_hidden), - pb_positions: take!(i_pb_positions), - pb_bf16: take!(i_pb_bf16), }; debug_assert!( slots.iter().all(|s| s.is_none()) && pos_slot.is_none(), @@ -2700,39 +2208,44 @@ impl Gemma4Scratch { let _ = gpu.free_tensor(self.moe_expert_gate_batch); let _ = gpu.free_tensor(self.moe_expert_up_batch); let _ = gpu.free_tensor(self.moe_expert_hidden_batch); - let _ = gpu.free_tensor(self.pb_attn_out); - let _ = gpu.free_tensor(self.pb_ffn_out); - let _ = gpu.free_tensor(self.pb_moe_pre2); - let _ = gpu.free_tensor(self.pb_moe_pre2_rot); - let _ = gpu.free_tensor(self.pb_moe_router_in); - let _ = gpu.free_tensor(self.pb_moe_router_logits); - let _ = gpu.free_tensor(self.pb_moe_topk_indices); - let _ = gpu.free_tensor(self.pb_moe_topk_weights); - let _ = gpu.free_tensor(self.pb_moe_expert_offsets); - let _ = gpu.free_tensor(self.pb_moe_expert_token_list); - let _ = gpu.free_tensor(self.pb_moe_gate_batch); - let _ = gpu.free_tensor(self.pb_moe_up_batch); - let _ = gpu.free_tensor(self.pb_moe_hidden_batch); - let _ = gpu.free_tensor(self.pb_moe_cur_moe); - let _ = gpu.free_tensor(self.pb_moe_cur_mlp); - let _ = gpu.free_tensor(self.pb_residual); - let _ = gpu.free_tensor(self.pb_tmp); - let _ = gpu.free_tensor(self.pb_q); - let _ = gpu.free_tensor(self.pb_attn_q); - let _ = gpu.free_tensor(self.pb_flash_partials); - let _ = gpu.free_tensor(self.pb_k); - let _ = gpu.free_tensor(self.pb_v); - let _ = gpu.free_tensor(self.pb_gate); - let _ = gpu.free_tensor(self.pb_up); - let _ = gpu.free_tensor(self.pb_ffn_hidden); - let _ = gpu.free_tensor(self.pb_positions); - let _ = gpu.free_tensor(self.pb_bf16); } } #[cfg(test)] mod scratch_geometry_tests { use super::*; + #[test] + fn flash_partials_geometry_matches_formula() { + // Single tile edge + assert_eq!(gemma4_flash_partials_len(128, 32, 512), 32 * 1 * 514); + assert_eq!(gemma4_flash_partials_len(129, 32, 512), 32 * 2 * 514); + assert_eq!(gemma4_flash_partials_len(256, 32, 512), 32 * 2 * 514); + // 32k baseline (FALLBACK_KV_SEQ before fix) + assert_eq!(gemma4_flash_partials_len(32768, 32, 512), 32 * 256 * 514); + assert_eq!(gemma4_flash_partials_len(32768, 32, 512), 4_210_688); + // 131072 must be exactly 4× the 32768 geometry (no overflow, no env var) + assert_eq!(gemma4_flash_partials_len(131072, 32, 512), 32 * 1024 * 514); + assert_eq!(gemma4_flash_partials_len(131072, 32, 512), 16_842_752); + assert_eq!( + gemma4_flash_partials_len(131072, 32, 512), + 4 * gemma4_flash_partials_len(32768, 32, 512) + ); + } + + #[test] + fn flash_scales_invariant_to_tile_rounding() { + // Non-multiple of 128 must ceil + let tiles_32769 = 32769usize.div_ceil(GEMMA4_FLASH_TILE); + assert_eq!(tiles_32769, 257); + assert_eq!(gemma4_flash_partials_len(32769, 32, 512), 32 * 257 * 514); + // 131071 is one short of 131072 -> still 1024 tiles (ceil) + assert_eq!(131071usize.div_ceil(GEMMA4_FLASH_TILE), 1024); + assert_eq!( + gemma4_flash_partials_len(131071, 32, 512), + gemma4_flash_partials_len(131072, 32, 512) + ); + } + fn dummy_cfg_31b() -> Gemma4Config { // Minimal config mirroring 31B/26B shapes: n_heads=32, full_head_dim=512 Gemma4Config { @@ -2772,61 +2285,6 @@ mod scratch_geometry_tests { } } - #[test] - fn flash_partials_geometry_matches_formula() { - // Single tile edge - assert_eq!(gemma4_flash_partials_len(128, 32, 512), 32 * 1 * 514); - assert_eq!(gemma4_flash_partials_len(129, 32, 512), 32 * 2 * 514); - assert_eq!(gemma4_flash_partials_len(256, 32, 512), 32 * 2 * 514); - // 32k baseline (FALLBACK_KV_SEQ before fix) - assert_eq!(gemma4_flash_partials_len(32768, 32, 512), 32 * 256 * 514); - assert_eq!(gemma4_flash_partials_len(32768, 32, 512), 4_210_688); - // 131072 must be exactly 4× the 32768 geometry (no overflow, no env var) - assert_eq!(gemma4_flash_partials_len(131072, 32, 512), 32 * 1024 * 514); - assert_eq!(gemma4_flash_partials_len(131072, 32, 512), 16_842_752); - assert_eq!( - gemma4_flash_partials_len(131072, 32, 512), - 4 * gemma4_flash_partials_len(32768, 32, 512) - ); - } - - #[test] - fn pb_flash_is_batch_scaled_and_single_authority() { - for &max_seq in &[128usize, 1024, 32768, 131072] { - let single = gemma4_flash_partials_len(max_seq, 32, 512); - let batched = gemma4_pb_flash_partials_len(max_seq, 32, 512); - assert_eq!(batched, GEMMA4_MAX_PREFILL_BATCH * single); - // PB via config helper shares the same max_seq authority - let cfg = dummy_cfg_31b(); - assert_eq!( - gemma4_pb_flash_partials_len_for_config(max_seq, &cfg), - batched - ); - assert_eq!(gemma4_flash_partials_len_for_config(max_seq, &cfg), single); - } - // Concrete 131072 batched size without allocating GPU memory - assert_eq!( - gemma4_pb_flash_partials_len(131072, 32, 512), - 128 * 16_842_752 - ); - assert_eq!(gemma4_pb_flash_partials_len(131072, 32, 512), 2_155_872_256); - assert_eq!(gemma4_pb_flash_partials_len(32768, 32, 512), 538_968_064); - } - - #[test] - fn flash_scales_invariant_to_tile_rounding() { - // Non-multiple of 128 must ceil - let tiles_32769 = 32769usize.div_ceil(GEMMA4_FLASH_TILE); - assert_eq!(tiles_32769, 257); - assert_eq!(gemma4_flash_partials_len(32769, 32, 512), 32 * 257 * 514); - // 131071 is one short of 131072 -> still 1024 tiles (ceil) - assert_eq!(131071usize.div_ceil(GEMMA4_FLASH_TILE), 1024); - assert_eq!( - gemma4_flash_partials_len(131071, 32, 512), - gemma4_flash_partials_len(131072, 32, 512) - ); - } - #[test] fn decode_partials_cover_q8_decode_tile() { let cfg = dummy_cfg_31b(); @@ -3358,477 +2816,13 @@ mod load_rollback_tests { // ─── Forward pass ─────────────────────────────────────────────────────── -/// Indexed single-token MoE expert phase shared by decode and batched prefill. -/// -/// Sequence matches the historical `apply_moe_branch` fast arm: -/// gate_up (dtype-branched, MQ4 rotates `pre2` → `pre2_rot` once) → -/// GELU-tanh → mul → scaled indexed down into `cur_moe`. -/// -/// `cur_moe` write semantics by down dtype: -/// - HFQ4G128: kernel ASSIGNS the residual row (pre-zero harmless) -/// - Q8_0: kernel atomicAdds into the residual row (caller MUST zero) -/// -/// Caller always zeroes `cur_moe` before calling. Views are non-owning; this -/// helper never allocates, copies, or builds `sub_offset` slices. -#[allow(clippy::too_many_arguments)] -fn moe_token_indexed( - gpu: &mut Gpu, - moe: &MoeLayerExtras, - gate_up_dtype: DType, - down_is_q8: bool, - pre2: &GpuTensor, - pre2_rot: &GpuTensor, - topk_idx: &GpuTensor, - topk_wt: &GpuTensor, - gate: &GpuTensor, - up: &GpuTensor, - hidden: &GpuTensor, - cur_moe: &GpuTensor, - dim: usize, - mi: usize, - k_top: usize, -) -> HipResult<()> { - // Indexed gate_up: 8 fused GEMVs reading expert IDs from device. - // y_gate: [k_top × mi], y_up: [k_top × mi] - if gate_up_dtype == DType::MQ4G256 { - // MQ4G256 needs FWHT-rotated input. - gpu.rotate_x_mq(pre2, pre2_rot, dim)?; - gpu.gemv_mq4g256_moe_gate_up_k8_indexed( - &moe.experts_gate_up_ptrs, - topk_idx, - pre2_rot, - gate, - up, - 2 * mi, - dim, - )?; - } else if gate_up_dtype == DType::HFQ4G256 || gate_up_dtype == DType::HFQ6G256 { - run_uniform_moe_gate_up( - gpu, - gate_up_dtype, - &moe.experts_gate_up_ptrs, - topk_idx, - pre2, - gate, - up, - 2 * mi, - dim, - k_top, - ) - .map_err(|e| hip_bridge::HipError::new(0, &e.to_string()))?; - } else { - // Q8_0 — no rotation needed. - gpu.gemv_q8_0_moe_gate_up_k8_indexed( - &moe.experts_gate_up_ptrs, - topk_idx, - pre2, - gate, - up, - 2 * mi, - dim, - )?; - } - - // Batched gelu_tanh + mul over [k_top × mi]. - gpu.gelu_tanh_f32(gate, hidden, k_top * mi)?; - gpu.mul_f32(hidden, up, hidden)?; - - // Indexed down + scaled residual: 8 fused GEMVs. Quant variant picked by - // the down weight format. HFQ4G128 assigns; Q8_0 atomicAdds (caller-zeroed). - if down_is_q8 { - gpu.gemv_q8_0_moe_down_residual_scaled_k8_indexed( - &moe.experts_down_ptrs, - topk_idx, - topk_wt, - &moe.per_expert_scale, - hidden, - cur_moe, - dim, - mi, - )?; - } else { - // down_hfq4g128 path. - gpu.gemv_hfq4g128_moe_down_residual_scaled_k8_indexed( - &moe.experts_down_ptrs, - topk_idx, - topk_wt, - &moe.per_expert_scale, - hidden, - cur_moe, - dim, - mi, - )?; - } - Ok(()) -} - -/// Re-point a non-owning F32 row view's `buf` at `owner[offset_elems .. +len]`. -/// Shape stays as constructed (no `Vec` alloc). View must not outlive owner. -#[inline] -fn moe_repoint_f32_row( - view: &mut GpuTensor, - owner: &GpuTensor, - offset_elems: usize, - len_elems: usize, -) { - let byte_off = offset_elems - .checked_mul(4) - .expect("moe row view offset overflow"); - let byte_len = len_elems - .checked_mul(4) - .expect("moe row view length overflow"); - let ptr = unsafe { (owner.buf.as_ptr() as *mut u8).add(byte_off) as *mut std::ffi::c_void }; - view.buf = unsafe { hip_bridge::DeviceBuffer::from_raw(ptr, byte_len) }; -} - -/// Batched MoE branch — processes N tokens at once. Mirrors `apply_moe_branch` -/// but every per-token tensor becomes a `[N × *]` row-major batch and every -/// kernel call uses the `_batched` variant. -/// -/// Inputs (live in `scratch.pb_*`): -/// - pb_attn_out[N × dim]: post-attention output (input to pre_norm_2 and router) -/// - pb_ffn_out [N × dim]: dense FFN output (input to post_norm_1) -/// -/// Output: -/// - pb_moe_pre2[N × dim] holds the post-FFN-layernorm(cur_mlp + cur_moe) — -/// i.e. what the per-token path writes to `scratch.tmp`. Caller adds it -/// to the per-token residual + applies the layer scalar. -/// -/// Shares the single-token indexed sequence (`moe_token_indexed`) per batch -/// row via non-owning views. Unsupported quant mixes retain the historical -/// D2H + per-expert CPU loop (slow but correct). -fn apply_moe_branch_batched( - gpu: &mut Gpu, - config: &Gemma4Config, - scratch: &Gemma4Scratch, - moe: &MoeLayerExtras, - post_ffn_norm: &GpuTensor, - n_batch: usize, -) -> HipResult<()> { - let dim = config.dim; - let mi = config.moe_intermediate_size; - let n_exp = config.num_experts; - let k_top = config.top_k_experts; - if k_top != 8 { - return Err(hip_bridge::HipError::new( - 0, - &format!("MoE top_k_experts={k_top} unsupported (kernel hardcoded to 8)"), - )); - } - debug_assert!( - n_batch <= scratch.max_prefill_batch, - "n_batch={n_batch} > MAX_PREFILL_BATCH={}", - scratch.max_prefill_batch - ); - - let first = &moe.experts[0]; - - // 1) cur_mlp_batch = post_feedforward_layernorm_1(pb_ffn_out) - gpu.rmsnorm_batched( - &scratch.pb_ffn_out, - &moe.post_feedforward_layernorm_1, - &scratch.pb_moe_cur_mlp, - n_batch, - dim, - config.norm_eps, - )?; - - // 2) pre2_batch = pre_feedforward_layernorm_2(pb_attn_out) - gpu.rmsnorm_batched( - &scratch.pb_attn_out, - &moe.pre_feedforward_layernorm_2, - &scratch.pb_moe_pre2, - n_batch, - dim, - config.norm_eps, - )?; - - // 3) router_in = rmsnorm(attn_out, router_scale) / sqrt(dim) - gpu.rmsnorm_batched( - &scratch.pb_attn_out, - &moe.router_scale, - &scratch.pb_moe_router_in, - n_batch, - dim, - config.norm_eps, - )?; - gpu.scale_f32(&scratch.pb_moe_router_in, 1.0 / (dim as f32).sqrt())?; - // scale_f32 operates on the full tensor — for [N × dim] this scales - // every element, which is what we want (each row is divided by sqrt(dim)). - - // 4) Router GEMM → logits [N × n_exp] - run_prefill_gemm( - gpu, - &moe.router_proj, - &scratch.pb_moe_router_in, - &scratch.pb_moe_router_logits, - n_batch, - Some(&scratch.pb_bf16), - )?; - - // 5) Batched top-K softmax + renorm. - gpu.moe_softmax_topk_renorm_k8_batched( - &scratch.pb_moe_router_logits, - &scratch.pb_moe_topk_indices, - &scratch.pb_moe_topk_weights, - n_exp, - true, - n_batch, - )?; - - // 5b) Phase B (opt-in): build per-expert buckets so the bucketed GEMV - // can reuse the weight tile across all tokens routed to the same expert. - // Measured -5.7% prefill regression on Gemma 4 26B-A4B-it / gfx1201 at - // N=128 batch (132 → 124 tok/s): the bucketed kernel pre-loads the - // weight tile into ~48 VGPRs per lane (sc/zp/pk × 16 groups) which - // cuts occupancy below the indexed_batched kernel that streams weights - // in the inner loop, and the serialized loop over bucket tokens loses - // more parallelism than is recovered from launch-overhead savings - // (180k blocks vs 1.4M). Default OFF; kept behind opt-in for future - // tuning (smaller MAX_GROUPS, LDS staging, fewer launch_bounds waves). - let use_bucketed = hipfire_config::developer_var("HIPFIRE_MOE_BUCKETED") - .ok() - .map(|v| v == "1") - .unwrap_or(false); - if use_bucketed { - gpu.moe_bucket_build( - &scratch.pb_moe_topk_indices, - &scratch.pb_moe_expert_offsets, - &scratch.pb_moe_expert_token_list, - n_batch, - k_top, - n_exp, - )?; - } - - // 6) Expert phase. Fast path keeps top-K device-resident and shares the - // single-token indexed sequence via per-row non-owning views. Legacy path - // retains the historical D2H + per-expert CPU loop for unsupported mixes. - let gate_mq4 = first.gate_up_proj.gpu_dtype == DType::MQ4G256; - let gate_hfq4 = first.gate_up_proj.gpu_dtype == DType::HFQ4G256; - let gate_hfq6 = first.gate_up_proj.gpu_dtype == DType::HFQ6G256; - let gate_q8 = first.gate_up_proj.gpu_dtype == DType::Q8_0; - let down_q8 = first.down_proj.gpu_dtype == DType::Q8_0; - let down_hfq4g128 = first.down_proj.gpu_dtype == DType::HFQ4G128; - let fast = (gate_mq4 || gate_hfq4 || gate_hfq6 || gate_q8) && (down_q8 || down_hfq4g128); - let dim_bytes = dim * 4; - let hid_elems = k_top * mi; - - // Zero cur_moe_batch accumulator once (required Q8 atomicAdd; harmless HFQ assign). - if let Some(s) = gpu.active_stream.as_ref() { - gpu.hip - .memset_async(&scratch.pb_moe_cur_moe.buf, 0, n_batch * dim_bytes, s)?; - } else { - gpu.hip - .memset(&scratch.pb_moe_cur_moe.buf, 0, n_batch * dim_bytes)?; - } - - if fast { - // Hoist 8 mutable row-view slots once (one shape Vec each). Per token only - // re-points `.buf` via DeviceBuffer::from_raw — no per-row shape alloc. - let mut v_pre2 = scratch.pb_moe_pre2.sub_offset(0, dim); - let mut v_pre2_rot = scratch.pb_moe_pre2_rot.sub_offset(0, dim); - let mut v_idx = scratch.pb_moe_topk_indices.sub_offset(0, k_top); - let mut v_wt = scratch.pb_moe_topk_weights.sub_offset(0, k_top); - let mut v_gate = scratch.pb_moe_gate_batch.sub_offset(0, hid_elems); - let mut v_up = scratch.pb_moe_up_batch.sub_offset(0, hid_elems); - let mut v_hidden = scratch.pb_moe_hidden_batch.sub_offset(0, hid_elems); - let mut v_cur = scratch.pb_moe_cur_moe.sub_offset(0, dim); - - for b in 0..n_batch { - moe_repoint_f32_row(&mut v_pre2, &scratch.pb_moe_pre2, b * dim, dim); - moe_repoint_f32_row(&mut v_pre2_rot, &scratch.pb_moe_pre2_rot, b * dim, dim); - moe_repoint_f32_row(&mut v_idx, &scratch.pb_moe_topk_indices, b * k_top, k_top); - moe_repoint_f32_row(&mut v_wt, &scratch.pb_moe_topk_weights, b * k_top, k_top); - moe_repoint_f32_row( - &mut v_gate, - &scratch.pb_moe_gate_batch, - b * hid_elems, - hid_elems, - ); - moe_repoint_f32_row( - &mut v_up, - &scratch.pb_moe_up_batch, - b * hid_elems, - hid_elems, - ); - moe_repoint_f32_row( - &mut v_hidden, - &scratch.pb_moe_hidden_batch, - b * hid_elems, - hid_elems, - ); - moe_repoint_f32_row(&mut v_cur, &scratch.pb_moe_cur_moe, b * dim, dim); - - // Helper rotates MQ row once into v_pre2_rot; no batched pre2 rotate. - moe_token_indexed( - gpu, - moe, - first.gate_up_proj.gpu_dtype, - down_q8, - &v_pre2, - &v_pre2_rot, - &v_idx, - &v_wt, - &v_gate, - &v_up, - &v_hidden, - &v_cur, - dim, - mi, - k_top, - )?; - } - } else { - // ── Legacy CPU per-expert path (quant mix doesn't match the fast - // kernels). D2H top-K + per-token expert loop; slow but correct. ── - let topk_idx_host = gpu.download_f32(&scratch.pb_moe_topk_indices)?; - let topk_wt_host = gpu.download_f32(&scratch.pb_moe_topk_weights)?; - let topk_indices_batch: Vec> = (0..n_batch) - .map(|b| { - unsafe { - std::slice::from_raw_parts( - topk_idx_host.as_ptr().add(b * k_top) as *const i32, - k_top, - ) - } - .iter() - .map(|&i| i as usize) - .collect() - }) - .collect(); - let topk_weights_batch: Vec> = (0..n_batch) - .map(|b| topk_wt_host[b * k_top..(b + 1) * k_top].to_vec()) - .collect(); - - for b in 0..n_batch { - for ki in 0..k_top { - let e = topk_indices_batch[b][ki]; - let weight = topk_weights_batch[b][ki] * moe.per_expert_scale_host[e]; - let expert = &moe.experts[e]; - - // Copy this token's pre2 row into scratch.moe_pre2 - if let Some(s) = gpu.active_stream.as_ref() { - gpu.hip.memcpy_dtod_async_at( - &scratch.moe_pre2.buf, - 0, - &scratch.pb_moe_pre2.buf, - b * dim_bytes, - dim_bytes, - s, - )?; - } else { - gpu.hip.memcpy_dtod_at( - &scratch.moe_pre2.buf, - 0, - &scratch.pb_moe_pre2.buf, - b * dim_bytes, - dim_bytes, - )?; - } - - // gate_up = expert.gate_up_proj @ pre2 - weight_gemv( - gpu, - &expert.gate_up_proj, - &scratch.moe_pre2, - &scratch.moe_expert_gate_up, - )?; - let gate = scratch.moe_expert_gate_up.sub_offset(0, mi); - let up = scratch.moe_expert_gate_up.sub_offset(mi, mi); - // hidden = gelu_tanh(gate) * up - gpu.gelu_tanh_f32(&gate, &scratch.moe_expert_hidden, mi)?; - gpu.mul_f32(&scratch.moe_expert_hidden, &up, &scratch.moe_expert_hidden)?; - // expert_out = expert.down_proj @ hidden - weight_gemv( - gpu, - &expert.down_proj, - &scratch.moe_expert_hidden, - &scratch.moe_expert_out, - )?; - // scaled_add into the correct row of pb_moe_cur_moe - if let Some(s) = gpu.active_stream.as_ref() { - gpu.hip.memcpy_dtod_async_at( - &scratch.tmp.buf, - 0, - &scratch.pb_moe_cur_moe.buf, - b * dim_bytes, - dim_bytes, - s, - )?; - } else { - gpu.hip.memcpy_dtod_at( - &scratch.tmp.buf, - 0, - &scratch.pb_moe_cur_moe.buf, - b * dim_bytes, - dim_bytes, - )?; - } - gpu.scaled_add_inplace_cpu_scalar_f32( - &scratch.tmp, - &scratch.moe_expert_out, - weight, - )?; - if let Some(s) = gpu.active_stream.as_ref() { - gpu.hip.memcpy_dtod_async_at( - &scratch.pb_moe_cur_moe.buf, - b * dim_bytes, - &scratch.tmp.buf, - 0, - dim_bytes, - s, - )?; - } else { - gpu.hip.memcpy_dtod_at( - &scratch.pb_moe_cur_moe.buf, - b * dim_bytes, - &scratch.tmp.buf, - 0, - dim_bytes, - )?; - } - } - } - } - - // 11) post_feedforward_layernorm_2(cur_moe) in-place batched. - gpu.rmsnorm_batched( - &scratch.pb_moe_cur_moe, - &moe.post_feedforward_layernorm_2, - &scratch.pb_moe_cur_moe, - n_batch, - dim, - config.norm_eps, - )?; - - // 12) combined = cur_mlp + cur_moe → reuse pb_moe_pre2 as the combined buffer. - gpu.add_f32( - &scratch.pb_moe_cur_mlp, - &scratch.pb_moe_cur_moe, - &scratch.pb_moe_pre2, - )?; - - // 13) post_feedforward_layernorm(combined) → batched, in-place. - gpu.rmsnorm_batched( - &scratch.pb_moe_pre2, - post_ffn_norm, - &scratch.pb_moe_pre2, - n_batch, - dim, - config.norm_eps, - )?; - - Ok(()) -} - /// Fingerprint the allocations and constants baked into the lowered AR graph. /// Position is deliberately absent: RoPE, KV writes and attention read pos_buf. fn lowered_graph_binding( weights: &Gemma4Weights, config: &Gemma4Config, - kv_sliding: &llama::KvCache, - kv_full: &llama::KvCache, + kv_sliding: &hipfire_runtime::llama::KvCache, + kv_full: &hipfire_runtime::llama::KvCache, s: &Gemma4Scratch, ) -> u64 { let mut h = 0xcbf2_9ce4_8422_2325u64; @@ -3960,6 +2954,7 @@ fn lowered_graph_binding( /// /// Precondition: the loader initializes `scratch.v_norm_ones_full` once before /// the first forward call. +#[allow(clippy::too_many_arguments)] pub fn forward_scratch( gpu: &mut Gpu, weights: &Gemma4Weights, @@ -3977,26 +2972,7 @@ pub fn forward_scratch( // call). The embed lookup fills scratch.x; the rest of the forward // (which is graph-capturable now that the MoE branch has no D2H syncs) // reads from scratch.x. Same split Qwen35 uses for its captured path. - match weights.embd_format { - EmbeddingFormat::HFQ4G256 => { - gpu.embedding_lookup_hfq4g256(&weights.embed_tokens, &scratch.x, token, dim)? - } - EmbeddingFormat::HFQ4G128 => { - gpu.embedding_lookup_hfq4g128(&weights.embed_tokens, &scratch.x, token, dim)? - } - EmbeddingFormat::Q8_0 => { - gpu.embedding_lookup_q8(&weights.embed_tokens, &scratch.x, token, dim)? - } - EmbeddingFormat::F32 => { - gpu.embedding_lookup(&weights.embed_tokens, &scratch.x, token, dim)? - } - _ => { - return Err(hip_bridge::HipError::new( - 0, - "unsupported Gemma 4 embed format", - )); - } - } + embed(gpu, weights, &scratch.x, token, dim)?; gpu.scale_f32(&scratch.x, config.embed_scale)?; // Same explicit per-model override as the hand path; HIPFIRE_GRAPH also @@ -4121,9 +3097,7 @@ fn forward_scratch_inner( kv_full: &mut hipfire_runtime::llama::KvCache, scratch: &Gemma4Scratch, ) -> HipResult<()> { - use crate::program::{ - head, Geometry, LayerKv, LayerRefs, LayerScratch, ProgramBinding, Resident, - }; + use crate::program::{head, Geometry, LayerScratch, ProgramBinding, Resident}; let hip = |e: String| hip_bridge::HipError::new(0, &e); let geo = Geometry::lowered(config); let binding = ProgramBinding { @@ -4140,25 +3114,8 @@ fn forward_scratch_inner( positions: None, scratch: LayerScratch::lowered(scratch, config.enable_moe_block), }; - // One cache slot per layer of each type, in layer order. - let mut slots = [0usize; 2]; let mut steps = Vec::with_capacity(3 * config.n_layers + 3); - for (layer_idx, layer) in weights.layers.iter().enumerate() { - let class = matches!(layer, LayerWeights::Full(_)) as usize; - let kv = LayerKv { - slot: slots[class], - writes: true, - }; - slots[class] += 1; - binding - .layer( - layer_idx, - LayerRefs::lowered(layer, config.hidden_dim), - kv, - &mut steps, - ) - .map_err(hip)?; - } + layer_steps(&binding, weights, config, &mut steps).map_err(hip)?; let lm_head = weights.lm_head.dispatch_ref(); head( &geo, @@ -4173,84 +3130,60 @@ fn forward_scratch_inner( execute_steps(gpu, &ctx, &steps).map_err(|e| hip(format!("gemma4: {e}"))) } -/// Single sliding-window attention layer. -/// -/// Order matches HF modeling_gemma4.py::Gemma4TextDecoderLayer + -/// Gemma4TextAttention (sliding branch): -/// residual = x -/// x = input_layernorm(x) — RMSNorm (sandwich pre-attn) -/// q = q_proj(x); q = q_norm(q) — RMSNorm over head_dim=256 -/// k = k_proj(x); k = k_norm(k) -/// v = v_proj(x); v = v_norm(v) — no_scale RMSNorm (ones buffer) -/// RoPE(q, k) with rotate_half, theta=10000, full head_dim=256 -/// write K, V to KV cache at position `pos` -/// attn = flash_attention(q, K, V, window_size=1024, scale=1.0 effective) -/// x = o_proj(attn) -/// x = post_attention_layernorm(x) — RMSNorm (sandwich post-attn) -/// x = residual + x -/// residual = x -/// x = pre_feedforward_layernorm(x) — RMSNorm (sandwich pre-FFN) -/// gate = gate_proj(x); up = up_proj(x) -/// ffn = gelu_pytorch_tanh(gate) * up — SwiGLU -/// x = down_proj(ffn) -/// x = post_feedforward_layernorm(x) — RMSNorm (sandwich post-FFN) -/// x = residual + x -/// x = x * layer_scalar — learned per-layer scalar -/// -/// Gemma 4 attention uses `scaling=1.0` in HF (see modeling_gemma4.py line 1143). -/// Our flash kernels bake in `scale = 1/sqrt(head_dim)`; we compensate by -/// pre-scaling Q by sqrt(head_dim) so the effective scale is 1.0. -// Sliding-window kernel correctness: the 4 attention_flash_*_window dispatch -// sites here use the kernels modified by commit b608c5f (sliding-window arg -// threaded through 7 attention_flash_*.hip files). Static review of the diff -// is clean — per-position out_of_window predicate writes -1e30 to scores; -// early-tile-skip writes {-1e30, 0, zeros} so reduce sees no contribution; -// all 32 lanes participate in __shfl_xor; window_size=0 = byte-identical -// for Qwen 3.5/3.6. Runtime NRMSE validation: verify_against_torch.rs Phase 2 -// (D2.5) — per-layer attention NRMSE against the bf16 PyTorch reference dump. - -/// Token-batched prefill. Processes up to `scratch.max_prefill_batch` tokens -/// at once, amortizing MoE-branch launch overhead across the batch via the -/// batched-indexed kernels (router GEMM, top-K, gate_up, gelu*mul, down). -/// -/// Per-token operations (attention, dense projections, RoPE, KV writes) -/// stay per-token in V1 — no batched flash-prefill kernel for asym3 -/// sliding window. The win comes from collapsing N per-token MoE calls -/// (~10 launches each) into one batched MoE call per layer. -/// -/// On exit: `kv_sliding` / `kv_full` are updated for every token. The KV -/// cache is the source of truth for downstream sampling; this function -/// does NOT compute logits (the caller should follow with a per-token -/// `forward_scratch` for the LAST token if it needs logits). -pub fn forward_prefill_batch( +/// `dst = embed(token)`, unscaled. +fn embed( gpu: &mut Gpu, weights: &Gemma4Weights, - config: &Gemma4Config, - tokens: &[u32], - start_pos: usize, - kv_sliding: &mut hipfire_runtime::llama::KvCache, - kv_full: &mut hipfire_runtime::llama::KvCache, - scratch: &Gemma4Scratch, + dst: &GpuTensor, + token: u32, + dim: usize, ) -> HipResult<()> { - // v2 — batched dense projections + batched MoE. The +55% prefill win - // from 521161f8 is back: the regressing bug was in `gemm_hfq4g128`'s - // partial-trailing-group handling (used floor instead of ceil for - // groups_per_row, silently dropped 64 input dims at K=2112 dense - // down_proj → garbage). Fixed in kernels/src/gemm_hfq4g128.hip. - forward_prefill_batch_v2( - gpu, weights, config, tokens, start_pos, kv_sliding, kv_full, scratch, - ) + let table = &weights.embed_tokens; + match weights.embd_format { + EmbeddingFormat::HFQ4G256 => gpu.embedding_lookup_hfq4g256(table, dst, token, dim), + EmbeddingFormat::HFQ4G128 => gpu.embedding_lookup_hfq4g128(table, dst, token, dim), + EmbeddingFormat::Q8_0 => gpu.embedding_lookup_q8(table, dst, token, dim), + EmbeddingFormat::F32 => gpu.embedding_lookup(table, dst, token, dim), + _ => Err(hip_bridge::HipError::new( + 0, + "unsupported Gemma 4 embed format", + )), + } } -/// V2 batched prefill — batches both the dense projections AND the MoE -/// branch across N tokens. Per-token RoPE and attention stay sequential -/// because the flash kernels haven't been ported to prefill-batched layout. -/// -/// Fixed 2026-05-19 by the gemm_hfq4g128 partial-trailing-group fix -/// (kernels/src/gemm_hfq4g128.hip). The v2 dispatch path is byte-identical -/// to v1 through attention; the divergence was the missing 64-element -/// partial group in the dense down_proj GEMM at K=2112. -fn forward_prefill_batch_v2( +/// Every layer's steps, one cache slot per layer of each type in layer order. +fn layer_steps<'a>( + binding: &crate::program::ProgramBinding<'a>, + weights: &'a Gemma4Weights, + config: &Gemma4Config, + steps: &mut Vec>, +) -> Result<(), String> { + use crate::program::{LayerKv, LayerRefs}; + let mut slots = [0usize; 2]; + for (layer_idx, layer) in weights.layers.iter().enumerate() { + let class = matches!(layer, LayerWeights::Full(_)) as usize; + let kv = LayerKv { + slot: slots[class], + writes: true, + }; + slots[class] += 1; + binding.layer( + layer_idx, + LayerRefs::lowered(layer, config.hidden_dim), + kv, + steps, + )?; + } + Ok(()) +} + +/// Batched prefill: runs `tokens` from `start_pos` through the declarative +/// layer program, `rows > 1` at a time, writing both KV caches (which must be +/// Q8). Computes no logits; follow with [`forward_scratch`] on the last token +/// when they are needed. The calibration collector, when armed, sees every +/// projection input. +#[allow(clippy::too_many_arguments)] +pub fn forward_prefill_batch( gpu: &mut Gpu, weights: &Gemma4Weights, config: &Gemma4Config, @@ -4260,1146 +3193,87 @@ fn forward_prefill_batch_v2( kv_full: &mut hipfire_runtime::llama::KvCache, scratch: &Gemma4Scratch, ) -> HipResult<()> { - let n_batch = tokens.len(); - if n_batch == 0 { - return Ok(()); - } - if n_batch > scratch.max_prefill_batch { - return Err(hip_bridge::HipError::new( - 0, - &format!( - "forward_prefill_batch: n_batch={n_batch} > max_prefill_batch={}", - scratch.max_prefill_batch - ), - )); - } - let dim = config.dim; - let dim_bytes = dim * 4; - - // Upload positions array [start_pos, start_pos+1, ..., start_pos+N-1] for batched RoPE. - let pos_array: Vec = (0..n_batch).map(|i| (start_pos + i) as i32).collect(); - let pos_bytes: Vec = pos_array.iter().flat_map(|p| p.to_ne_bytes()).collect(); - gpu.hip.memcpy_htod(&scratch.pb_positions.buf, &pos_bytes)?; - - // Step 1: per-token embed lookup into pb_residual. - for (i, &tok) in tokens.iter().enumerate() { - match weights.embd_format { - EmbeddingFormat::HFQ4G256 => { - gpu.embedding_lookup_hfq4g256(&weights.embed_tokens, &scratch.x, tok, dim)? - } - EmbeddingFormat::HFQ4G128 => { - gpu.embedding_lookup_hfq4g128(&weights.embed_tokens, &scratch.x, tok, dim)? - } - EmbeddingFormat::Q8_0 => { - gpu.embedding_lookup_q8(&weights.embed_tokens, &scratch.x, tok, dim)? - } - EmbeddingFormat::F32 => { - gpu.embedding_lookup(&weights.embed_tokens, &scratch.x, tok, dim)? - } - _ => { - return Err(hip_bridge::HipError::new( - 0, - "unsupported Gemma 4 embed format", - )); - } - } - gpu.scale_f32(&scratch.x, config.embed_scale)?; - if let Some(_s) = gpu.active_stream.as_ref() { - gpu.hip.memcpy_dtod_async_at( - &scratch.pb_residual.buf, - i * dim_bytes, - &scratch.x.buf, - 0, - dim_bytes, - _s, - )?; - } else { - gpu.hip.memcpy_dtod_at( - &scratch.pb_residual.buf, - i * dim_bytes, - &scratch.x.buf, - 0, - dim_bytes, - )?; - } - } - - // Step 2: layer loop. - let mut sliding_kv_idx = 0usize; - let mut full_kv_idx = 0usize; - for (layer_idx, layer_type) in config.layer_types.iter().copied().enumerate() { - // Invalidate the pointer-keyed F16 conversion cache between layers. - // The same activation buffer (pb_tmp) is reused with different contents - // each layer; without invalidation, ensure_fp16_x would skip the - // F32→F16 conversion and serve stale data from the previous layer. - gpu.invalidate_fp16_cache(); - let (layer_scalar, post_ffn_norm_ref, moe_opt) = - match (layer_type, &weights.layers[layer_idx]) { - (LayerType::Sliding, LayerWeights::Sliding(lw)) => ( - lw.layer_scalar_host, - &lw.post_feedforward_layernorm, - lw.moe.as_ref(), - ), - (LayerType::Full, LayerWeights::Full(lw)) => ( - lw.layer_scalar_host, - &lw.post_feedforward_layernorm, - lw.moe.as_ref(), - ), - _ => { - return Err(hip_bridge::HipError::new( - 0, - &format!("layer {layer_idx} type/weights mismatch"), - )); - } - }; - - // ── 2A: batched pre-attn rmsnorm + Q/K/V projections + Q/K/V norms + RoPE. ── - match (layer_type, &weights.layers[layer_idx]) { - (LayerType::Sliding, LayerWeights::Sliding(lw)) => { - let head_dim = config.sliding_head_dim; - let n_heads = config.n_heads; - let n_kv = config.sliding_n_kv_heads; - let q_dim = n_heads * head_dim; - let kv_dim = n_kv * head_dim; - let kv_dim_bytes = kv_dim * 4; - let q_dim_bytes = q_dim * 4; - - let _dump_on = layer_idx == 0 && start_pos == 0; - if _dump_on { - dbg_dump( - gpu, - "[v2] L0 input pb_residual[0]", - &scratch.pb_residual, - dim, - ); - } - gpu.rmsnorm_batched( - &scratch.pb_residual, - &lw.input_layernorm, - &scratch.pb_tmp, - n_batch, - dim, - config.norm_eps, - )?; - if _dump_on { - dbg_dump(gpu, "[v2] L0 after input_norm", &scratch.pb_tmp, dim); - } - // BF16 hoist: one F32->BF16 for the shared pb_tmp -> q/k/v group. - // Capture the shared F32 source (pb_tmp) once per constituent - // weight before staging — same src, per-weight k. The inner - // run_prefill_gemm_inner skips capture for BF16 x (hoisted), - // so this hoist-site capture restores coverage. Zero-cost when unarmed. - if lw.q_proj.gpu_dtype == DType::BF16 { - let nelems = n_batch * dim; - let bf16_x = scratch.pb_bf16.sub_offset(0, nelems); - gpu.maybe_capture_activation( - &lw.q_proj.buf, - &scratch.pb_tmp, - n_batch, - lw.q_proj.k, - ); - gpu.maybe_capture_activation( - &lw.k_proj.buf, - &scratch.pb_tmp, - n_batch, - lw.k_proj.k, - ); - gpu.maybe_capture_activation( - &lw.v_proj.buf, - &scratch.pb_tmp, - n_batch, - lw.v_proj.k, - ); - gpu.convert_f32_to_bf16(&scratch.pb_tmp, &bf16_x, nelems)?; - run_prefill_gemm( - gpu, - &lw.q_proj, - &bf16_x, - &scratch.pb_q, - n_batch, - Some(&scratch.pb_bf16), - )?; - run_prefill_gemm( - gpu, - &lw.k_proj, - &bf16_x, - &scratch.pb_k, - n_batch, - Some(&scratch.pb_bf16), - )?; - run_prefill_gemm( - gpu, - &lw.v_proj, - &bf16_x, - &scratch.pb_v, - n_batch, - Some(&scratch.pb_bf16), - )?; - } else { - run_prefill_gemm( - gpu, - &lw.q_proj, - &scratch.pb_tmp, - &scratch.pb_q, - n_batch, - Some(&scratch.pb_bf16), - )?; - run_prefill_gemm( - gpu, - &lw.k_proj, - &scratch.pb_tmp, - &scratch.pb_k, - n_batch, - Some(&scratch.pb_bf16), - )?; - run_prefill_gemm( - gpu, - &lw.v_proj, - &scratch.pb_tmp, - &scratch.pb_v, - n_batch, - Some(&scratch.pb_bf16), - )?; - } - if _dump_on { - dbg_dump(gpu, "[v2] L0 after q_proj", &scratch.pb_q, q_dim); - dbg_dump(gpu, "[v2] L0 after k_proj", &scratch.pb_k, kv_dim); - dbg_dump(gpu, "[v2] L0 after v_proj", &scratch.pb_v, kv_dim); - } - gpu.rmsnorm_batched( - &scratch.pb_q, - &lw.q_norm, - &scratch.pb_q, - n_batch * n_heads, - head_dim, - config.norm_eps, - )?; - gpu.rmsnorm_batched( - &scratch.pb_k, - &lw.k_norm, - &scratch.pb_k, - n_batch * n_kv, - head_dim, - config.norm_eps, - )?; - gpu.rmsnorm_batched( - &scratch.pb_v, - &scratch.v_norm_ones_full, - &scratch.pb_v, - n_batch * n_kv, - head_dim, - config.norm_eps, - )?; - if _dump_on { - dbg_dump(gpu, "[v2] L0 after q_norm", &scratch.pb_q, q_dim); - dbg_dump(gpu, "[v2] L0 after k_norm", &scratch.pb_k, kv_dim); - dbg_dump(gpu, "[v2] L0 after v_norm", &scratch.pb_v, kv_dim); - } - gpu.scale_f32(&scratch.pb_q, (head_dim as f32).sqrt())?; - if _dump_on { - dbg_dump(gpu, "[v2] L0 after scale_q", &scratch.pb_q, q_dim); - } - gpu.rope_batched_f32( - &scratch.pb_q, - &scratch.pb_k, - &scratch.pb_positions, - n_heads, - n_kv, - head_dim, - config.sliding_rope_theta, - n_batch, - )?; - if _dump_on { - dbg_dump(gpu, "[v2] L0 after rope_q", &scratch.pb_q, q_dim); - dbg_dump(gpu, "[v2] L0 after rope_k", &scratch.pb_k, kv_dim); - } - - // Per-token KV write + attention (batched projections, per-token - // attention). The batched q8 attention path has a bug that - // corrupts KV when cache_capacity > 0 (ring-buffer mode), so - // we run attention per-token through the proven decode-path - // Step::Attend with batch_size=1. This is only 1.7% of GPU - // time per the rocprof profile, so the overhead is minimal. - let sliding_cap = config.sliding_window as u32; - for i in 0..n_batch { - let pos = start_pos + i; - if let Some(stream) = gpu.active_stream.as_ref() { - gpu.hip - .stream_write_value32(stream, &scratch.pos_buf, pos as u32, 0)?; - } else { - gpu.hip - .memcpy_htod(&scratch.pos_buf, &(pos as i32).to_ne_bytes())?; - } - if let Some(_s) = gpu.active_stream.as_ref() { - gpu.hip.memcpy_dtod_async_at( - &scratch.q.buf, - 0, - &scratch.pb_q.buf, - i * q_dim_bytes, - q_dim_bytes, - _s, - )?; - gpu.hip.memcpy_dtod_async_at( - &scratch.k.buf, - 0, - &scratch.pb_k.buf, - i * kv_dim_bytes, - kv_dim_bytes, - _s, - )?; - gpu.hip.memcpy_dtod_async_at( - &scratch.v.buf, - 0, - &scratch.pb_v.buf, - i * kv_dim_bytes, - kv_dim_bytes, - _s, - )?; - } else { - gpu.hip.memcpy_dtod_at( - &scratch.q.buf, - 0, - &scratch.pb_q.buf, - i * q_dim_bytes, - q_dim_bytes, - )?; - gpu.hip.memcpy_dtod_at( - &scratch.k.buf, - 0, - &scratch.pb_k.buf, - i * kv_dim_bytes, - kv_dim_bytes, - )?; - gpu.hip.memcpy_dtod_at( - &scratch.v.buf, - 0, - &scratch.pb_v.buf, - i * kv_dim_bytes, - kv_dim_bytes, - )?; - } - let tier_inputs = KvTierInputs { - quant_asym4: kv_sliding.quant_asym4, - quant_asym3: kv_sliding.quant_asym3, - quant_asym2: kv_sliding.quant_asym2, - quant_q8: kv_sliding.quant_q8, - quant_fwht: kv_sliding.quant_fwht, - quant_hfq4: false, - quant_q4: false, - quant_int8: false, - quant_hfq8: false, - quant_bf16: false, - quant_fp8: false, - f32_policy: hipfire_dispatch::families::kv_tier::F32AttnPolicy::Simple, - v_mode_bits: kv_sliding.v_mode_bits(), - pos, - flash_mode: 2, - capture_mode: gpu.graphs.capture_mode, - batch_size: 1, - is_tree: false, - is_boundary: false, - q8_windowed: true, - window: sliding_cap as i32, - }; - let plan = KvTierPlan::derive(tier_inputs) - .map_err(|e| hip_bridge::HipError::new(0, &format!("{:?}", e)))?; - let io = AttnParams { - q: &scratch.q, - k: &scratch.k, - v: &scratch.v, - k_cache: &kv_sliding.k_gpu[sliding_kv_idx], - v_cache: &kv_sliding.v_gpu[sliding_kv_idx], - k_scales: None, - v_scales: None, - pos_buf: &scratch.pos_buf, - pos, - positions: None, - n_heads, - n_kv_heads: n_kv, - head_dim, - physical_cap: kv_sliding.physical_cap, - batch_size: 1, - max_ctx_len: 0, - flash_partials: Some(&scratch.flash_partials), - givens_cos: kv_sliding.givens_cos.as_ref(), - givens_sin: kv_sliding.givens_sin.as_ref(), - tree_bias: None, - block_start: 0, - block_cols: 0, - output_gate: None, - output_awq_scale: None, - output: &scratch.attn_out, - }; - let ctx = DispatchCtx::new(gpu); - execute_steps(gpu, &ctx, &[Step::Attend { plan, io }]) - .map_err(|e| hip_bridge::HipError::new(0, &e.to_string()))?; - // Copy attention output back to batched buffer. - if let Some(_s) = gpu.active_stream.as_ref() { - gpu.hip.memcpy_dtod_async_at( - &scratch.pb_q.buf, - i * q_dim_bytes, - &scratch.attn_out.buf, - 0, - q_dim_bytes, - _s, - )?; - } else { - gpu.hip.memcpy_dtod_at( - &scratch.pb_q.buf, - i * q_dim_bytes, - &scratch.attn_out.buf, - 0, - q_dim_bytes, - )?; - } - } - sliding_kv_idx += 1; - if _dump_on { - dbg_dump(gpu, "[v2] L0 after attention (pb_q)", &scratch.pb_q, q_dim); - } - - run_prefill_gemm( - gpu, - &lw.o_proj, - &scratch.pb_q, - &scratch.pb_attn_out, - n_batch, - Some(&scratch.pb_bf16), - )?; - if _dump_on { - dbg_dump(gpu, "[v2] L0 after o_proj", &scratch.pb_attn_out, dim); - } - gpu.rmsnorm_batched( - &scratch.pb_attn_out, - &lw.post_attention_layernorm, - &scratch.pb_attn_out, - n_batch, - dim, - config.norm_eps, - )?; - if _dump_on { - dbg_dump( - gpu, - "[v2] L0 after post_attn_norm", - &scratch.pb_attn_out, - dim, - ); - } - gpu.add_inplace_f32(&scratch.pb_residual, &scratch.pb_attn_out)?; - if _dump_on { - dbg_dump( - gpu, - "[v2] L0 after attn_residual", - &scratch.pb_residual, - dim, - ); - } - } - (LayerType::Full, LayerWeights::Full(lw)) => { - let head_dim = config.full_head_dim; - let n_heads = config.n_heads; - let n_kv = config.full_n_kv_heads; - let q_dim = n_heads * head_dim; - let kv_dim = n_kv * head_dim; - let kv_dim_bytes = kv_dim * 4; - let q_dim_bytes = q_dim * 4; - - gpu.rmsnorm_batched( - &scratch.pb_residual, - &lw.input_layernorm, - &scratch.pb_tmp, - n_batch, - dim, - config.norm_eps, - )?; - // BF16 hoist for shared pb_tmp -> q/k (full). Capture shared F32 - // pb_tmp once per weight before staging. Gate+up hoisted separately below. - if lw.q_proj.gpu_dtype == DType::BF16 { - let nelems = n_batch * dim; - let bf16_x = scratch.pb_bf16.sub_offset(0, nelems); - gpu.maybe_capture_activation( - &lw.q_proj.buf, - &scratch.pb_tmp, - n_batch, - lw.q_proj.k, - ); - gpu.maybe_capture_activation( - &lw.k_proj.buf, - &scratch.pb_tmp, - n_batch, - lw.k_proj.k, - ); - gpu.convert_f32_to_bf16(&scratch.pb_tmp, &bf16_x, nelems)?; - run_prefill_gemm( - gpu, - &lw.q_proj, - &bf16_x, - &scratch.pb_q, - n_batch, - Some(&scratch.pb_bf16), - )?; - run_prefill_gemm( - gpu, - &lw.k_proj, - &bf16_x, - &scratch.pb_k, - n_batch, - Some(&scratch.pb_bf16), - )?; - } else { - run_prefill_gemm( - gpu, - &lw.q_proj, - &scratch.pb_tmp, - &scratch.pb_q, - n_batch, - Some(&scratch.pb_bf16), - )?; - run_prefill_gemm( - gpu, - &lw.k_proj, - &scratch.pb_tmp, - &scratch.pb_k, - n_batch, - Some(&scratch.pb_bf16), - )?; - } - if let Some(s) = gpu.active_stream.as_ref() { - gpu.hip.memcpy_dtod_async_at( - &scratch.pb_v.buf, - 0, - &scratch.pb_k.buf, - 0, - n_batch * kv_dim_bytes, - s, - )?; - } else { - gpu.hip.memcpy_dtod_at( - &scratch.pb_v.buf, - 0, - &scratch.pb_k.buf, - 0, - n_batch * kv_dim_bytes, - )?; - } - gpu.rmsnorm_batched( - &scratch.pb_q, - &lw.q_norm, - &scratch.pb_q, - n_batch * n_heads, - head_dim, - config.norm_eps, - )?; - gpu.rmsnorm_batched( - &scratch.pb_k, - &lw.k_norm, - &scratch.pb_k, - n_batch * n_kv, - head_dim, - config.norm_eps, - )?; - gpu.rmsnorm_batched( - &scratch.pb_v, - &scratch.v_norm_ones_full, - &scratch.pb_v, - n_batch * n_kv, - head_dim, - config.norm_eps, - )?; - gpu.scale_f32(&scratch.pb_q, (head_dim as f32).sqrt())?; - let n_rot_pairs = - ((head_dim as f32) * config.full_partial_rotary_factor * 0.5) as usize; - // Per-token partial-halved RoPE (no batched variant yet for this shape). - for i in 0..n_batch { - let pos = start_pos + i; - if let Some(stream) = gpu.active_stream.as_ref() { - gpu.hip - .stream_write_value32(stream, &scratch.pos_buf, pos as u32, 0)?; - } else { - let pos_i32 = pos as i32; - gpu.hip - .memcpy_htod(&scratch.pos_buf, &pos_i32.to_ne_bytes())?; - } - if let Some(_s) = gpu.active_stream.as_ref() { - gpu.hip.memcpy_dtod_async_at( - &scratch.q.buf, - 0, - &scratch.pb_q.buf, - i * q_dim_bytes, - q_dim_bytes, - _s, - )?; - } else { - gpu.hip.memcpy_dtod_at( - &scratch.q.buf, - 0, - &scratch.pb_q.buf, - i * q_dim_bytes, - q_dim_bytes, - )?; - } - if let Some(_s) = gpu.active_stream.as_ref() { - gpu.hip.memcpy_dtod_async_at( - &scratch.k.buf, - 0, - &scratch.pb_k.buf, - i * kv_dim_bytes, - kv_dim_bytes, - _s, - )?; - } else { - gpu.hip.memcpy_dtod_at( - &scratch.k.buf, - 0, - &scratch.pb_k.buf, - i * kv_dim_bytes, - kv_dim_bytes, - )?; - } - gpu.rope_partial_halved_f32( - &scratch.q, - &scratch.k, - &scratch.pos_buf, - n_heads, - n_kv, - head_dim, - n_rot_pairs, - config.full_rope_theta, - )?; - if let Some(_s) = gpu.active_stream.as_ref() { - gpu.hip.memcpy_dtod_async_at( - &scratch.pb_q.buf, - i * q_dim_bytes, - &scratch.q.buf, - 0, - q_dim_bytes, - _s, - )?; - } else { - gpu.hip.memcpy_dtod_at( - &scratch.pb_q.buf, - i * q_dim_bytes, - &scratch.q.buf, - 0, - q_dim_bytes, - )?; - } - if let Some(_s) = gpu.active_stream.as_ref() { - gpu.hip.memcpy_dtod_async_at( - &scratch.pb_k.buf, - i * kv_dim_bytes, - &scratch.k.buf, - 0, - kv_dim_bytes, - _s, - )?; - } else { - gpu.hip.memcpy_dtod_at( - &scratch.pb_k.buf, - i * kv_dim_bytes, - &scratch.k.buf, - 0, - kv_dim_bytes, - )?; - } - } - - // Batched full-layer attention (hd=512): all N tokens in ONE - // masked flash call instead of a per-token loop. Full layers are - // non-ring (cache_capacity=0) and non-windowed (window_size=0); - // the batched-masked kernel (q8 or legacy asym3) derives each query's causal - // bound from pb_positions, so this reproduces the per-token - // result with 1 launch/layer instead of N. (Sliding layers stay - // per-token until the q8 batched ring kernel lands — A.2.) - // Proportional RoPE was already applied to pb_q/pb_k above. - let _ = (q_dim_bytes, kv_dim_bytes); // (no per-token staging now) - let tier_inputs = KvTierInputs { - quant_asym4: kv_full.quant_asym4, - quant_asym3: kv_full.quant_asym3, - quant_asym2: kv_full.quant_asym2, - quant_q8: kv_full.quant_q8, - quant_fwht: kv_full.quant_fwht, - quant_hfq4: false, - quant_q4: false, - quant_int8: false, - quant_hfq8: false, - quant_bf16: false, - quant_fp8: false, - f32_policy: hipfire_dispatch::families::kv_tier::F32AttnPolicy::Simple, - v_mode_bits: kv_full.v_mode_bits(), - pos: start_pos + n_batch - 1, - flash_mode: 2, - capture_mode: gpu.graphs.capture_mode, - batch_size: n_batch, - is_tree: false, - is_boundary: false, - q8_windowed: true, - window: 0, - }; - let plan = KvTierPlan::derive(tier_inputs) - .map_err(|e| hip_bridge::HipError::new(0, &format!("{:?}", e)))?; - let io = AttnParams { - q: &scratch.pb_q, - k: &scratch.pb_k, - v: &scratch.pb_v, // k_eq_v model: V = pre-k_norm k_proj out, v-normed (NOT k-normed/roped pb_k) - k_cache: &kv_full.k_gpu[full_kv_idx], - v_cache: &kv_full.v_gpu[full_kv_idx], - k_scales: None, - v_scales: None, - pos_buf: &scratch.pos_buf, - pos: start_pos + n_batch - 1, - positions: Some(&scratch.pb_positions), - n_heads, - n_kv_heads: n_kv, - head_dim, - physical_cap: kv_full.max_seq, - batch_size: n_batch, - max_ctx_len: start_pos + n_batch, - flash_partials: Some(&scratch.pb_flash_partials), - givens_cos: kv_full.givens_cos.as_ref(), - givens_sin: kv_full.givens_sin.as_ref(), - tree_bias: None, - block_start: 0, - block_cols: 0, - output_gate: None, - output_awq_scale: None, - output: &scratch.pb_attn_q, - }; - let ctx = DispatchCtx::new(gpu); - execute_steps(gpu, &ctx, &[Step::Attend { plan, io }]) - .map_err(|e| hip_bridge::HipError::new(0, &e.to_string()))?; - if hipfire_config::developer_var("HIPFIRE_GEMMA4_ATTN_VERIFY") - .ok() - .as_deref() - == Some("1") - { - let batched_out = gpu.download_f32(&scratch.pb_attn_q)?; - for i in 0..n_batch { - let pos = start_pos + i; - gpu.hip - .memcpy_htod(&scratch.pos_buf, &(pos as i32).to_ne_bytes())?; - gpu.hip.memcpy_dtod_at( - &scratch.q.buf, - 0, - &scratch.pb_q.buf, - i * q_dim_bytes, - q_dim_bytes, - )?; - gpu.hip.memcpy_dtod_at( - &scratch.k.buf, - 0, - &scratch.pb_k.buf, - i * kv_dim_bytes, - kv_dim_bytes, - )?; - gpu.hip.memcpy_dtod_at( - &scratch.v.buf, - 0, - &scratch.pb_v.buf, - i * kv_dim_bytes, - kv_dim_bytes, - )?; - let ti = KvTierInputs { - quant_asym4: kv_full.quant_asym4, - quant_asym3: kv_full.quant_asym3, - quant_asym2: kv_full.quant_asym2, - quant_q8: kv_full.quant_q8, - quant_fwht: kv_full.quant_fwht, - quant_hfq4: false, - quant_q4: false, - quant_int8: false, - quant_hfq8: false, - quant_bf16: false, - quant_fp8: false, - f32_policy: hipfire_dispatch::families::kv_tier::F32AttnPolicy::Simple, - v_mode_bits: kv_full.v_mode_bits(), - pos, - flash_mode: 2, - capture_mode: gpu.graphs.capture_mode, - batch_size: 1, - is_tree: false, - is_boundary: false, - q8_windowed: false, - window: 0, - }; - let p1 = KvTierPlan::derive(ti) - .map_err(|e| hip_bridge::HipError::new(0, &format!("{:?}", e)))?; - let io1 = AttnParams { - q: &scratch.q, - k: &scratch.k, - v: &scratch.v, - k_cache: &kv_full.k_gpu[full_kv_idx], - v_cache: &kv_full.v_gpu[full_kv_idx], - k_scales: None, - v_scales: None, - pos_buf: &scratch.pos_buf, - pos, - positions: None, - n_heads, - n_kv_heads: n_kv, - head_dim, - physical_cap: kv_full.max_seq, - batch_size: 1, - max_ctx_len: 0, - flash_partials: Some(&scratch.flash_partials), - givens_cos: kv_full.givens_cos.as_ref(), - givens_sin: kv_full.givens_sin.as_ref(), - tree_bias: None, - block_start: 0, - block_cols: 0, - output_gate: None, - output_awq_scale: None, - output: &scratch.attn_out, - }; - let c1 = DispatchCtx::new(gpu); - execute_steps(gpu, &c1, &[Step::Attend { plan: p1, io: io1 }]) - .map_err(|e| hip_bridge::HipError::new(0, &e.to_string()))?; - let single = gpu.download_f32(&scratch.attn_out)?; - let row = &batched_out[i * q_dim..(i + 1) * q_dim]; - let mut worst = 0f32; - let mut wi = 0usize; - for j in 0..q_dim { - let d = (single[j] - row[j]).abs(); - if d > worst { - worst = d; - wi = j; - } - } - eprintln!( - "[attn-verify] L{layer_idx} start={start_pos} tok={i} pos={pos} worst={worst:.5} at {wi} single={:.4} batched={:.4}", - single[wi], row[wi] - ); - } - } - full_kv_idx += 1; - - run_prefill_gemm( - gpu, - &lw.o_proj, - &scratch.pb_attn_q, - &scratch.pb_attn_out, - n_batch, - Some(&scratch.pb_bf16), - )?; - gpu.rmsnorm_batched( - &scratch.pb_attn_out, - &lw.post_attention_layernorm, - &scratch.pb_attn_out, - n_batch, - dim, - config.norm_eps, - )?; - gpu.add_inplace_f32(&scratch.pb_residual, &scratch.pb_attn_out)?; - } - _ => unreachable!(), - } - - // Snapshot post-attn residual for MoE input. - if let Some(s) = gpu.active_stream.as_ref() { - gpu.hip.memcpy_dtod_async_at( - &scratch.pb_attn_out.buf, - 0, - &scratch.pb_residual.buf, - 0, - n_batch * dim_bytes, - s, + let hip = |e: String| hip_bridge::HipError::new(0, &e); + let rows_max = crate::gemma4::GEMMA4_FORWARD_BATCH_MAX; + for (i, chunk) in tokens.chunks(rows_max).enumerate() { + let pos = start_pos + i * rows_max; + if let [token] = chunk { + forward_scratch( + gpu, weights, config, *token, pos, kv_sliding, kv_full, scratch, )?; } else { - gpu.hip.memcpy_dtod_at( - &scratch.pb_attn_out.buf, - 0, - &scratch.pb_residual.buf, - 0, - n_batch * dim_bytes, - )?; - } - - // Batched pre-FFN rmsnorm + dense FFN. - let _dump_ffn = layer_idx == 0 && start_pos == 0; - match (layer_type, &weights.layers[layer_idx]) { - (LayerType::Sliding, LayerWeights::Sliding(lw)) => { - gpu.rmsnorm_batched( - &scratch.pb_residual, - &lw.pre_feedforward_layernorm, - &scratch.pb_tmp, - n_batch, - dim, - config.norm_eps, - )?; - if _dump_ffn { - dbg_dump(gpu, "[v2] L0 after pre_ffn_norm", &scratch.pb_tmp, dim); - } - // BF16 hoist for shared pb_tmp -> gate+up. Capture shared F32 - // pb_tmp once per weight before staging. - if lw.gate_proj.gpu_dtype == DType::BF16 { - let nelems = n_batch * dim; - let bf16_x = scratch.pb_bf16.sub_offset(0, nelems); - gpu.maybe_capture_activation( - &lw.gate_proj.buf, - &scratch.pb_tmp, - n_batch, - lw.gate_proj.k, - ); - gpu.maybe_capture_activation( - &lw.up_proj.buf, - &scratch.pb_tmp, - n_batch, - lw.up_proj.k, - ); - gpu.convert_f32_to_bf16(&scratch.pb_tmp, &bf16_x, nelems)?; - run_prefill_gemm( - gpu, - &lw.gate_proj, - &bf16_x, - &scratch.pb_gate, - n_batch, - Some(&scratch.pb_bf16), - )?; - run_prefill_gemm( - gpu, - &lw.up_proj, - &bf16_x, - &scratch.pb_up, - n_batch, - Some(&scratch.pb_bf16), - )?; - } else { - run_prefill_gemm( - gpu, - &lw.gate_proj, - &scratch.pb_tmp, - &scratch.pb_gate, - n_batch, - Some(&scratch.pb_bf16), - )?; - run_prefill_gemm( - gpu, - &lw.up_proj, - &scratch.pb_tmp, - &scratch.pb_up, - n_batch, - Some(&scratch.pb_bf16), - )?; - } - if _dump_ffn { - dbg_dump( - gpu, - "[v2] L0 after gate_proj", - &scratch.pb_gate, - config.hidden_dim, - ); - dbg_dump( - gpu, - "[v2] L0 after up_proj", - &scratch.pb_up, - config.hidden_dim, - ); - } - gpu.gelu_tanh_f32( - &scratch.pb_gate, - &scratch.pb_ffn_hidden, - n_batch * config.hidden_dim, - )?; - gpu.mul_f32( - &scratch.pb_ffn_hidden, - &scratch.pb_up, - &scratch.pb_ffn_hidden, - )?; - if _dump_ffn { - dbg_dump( - gpu, - "[v2] L0 after gelu*up", - &scratch.pb_ffn_hidden, - config.hidden_dim, - ); - } - run_prefill_gemm( - gpu, - &lw.down_proj, - &scratch.pb_ffn_hidden, - &scratch.pb_ffn_out, - n_batch, - Some(&scratch.pb_bf16), - )?; - if _dump_ffn { - dbg_dump(gpu, "[v2] L0 after down_proj", &scratch.pb_ffn_out, dim); - } - } - (LayerType::Full, LayerWeights::Full(lw)) => { - gpu.rmsnorm_batched( - &scratch.pb_residual, - &lw.pre_feedforward_layernorm, - &scratch.pb_tmp, - n_batch, - dim, - config.norm_eps, - )?; - // BF16 hoist for shared pb_tmp -> gate+up (full). Capture shared - // F32 pb_tmp once per weight before staging. - if lw.gate_proj.gpu_dtype == DType::BF16 { - let nelems = n_batch * dim; - let bf16_x = scratch.pb_bf16.sub_offset(0, nelems); - gpu.maybe_capture_activation( - &lw.gate_proj.buf, - &scratch.pb_tmp, - n_batch, - lw.gate_proj.k, - ); - gpu.maybe_capture_activation( - &lw.up_proj.buf, - &scratch.pb_tmp, - n_batch, - lw.up_proj.k, - ); - gpu.convert_f32_to_bf16(&scratch.pb_tmp, &bf16_x, nelems)?; - run_prefill_gemm( - gpu, - &lw.gate_proj, - &bf16_x, - &scratch.pb_gate, - n_batch, - Some(&scratch.pb_bf16), - )?; - run_prefill_gemm( - gpu, - &lw.up_proj, - &bf16_x, - &scratch.pb_up, - n_batch, - Some(&scratch.pb_bf16), - )?; - } else { - run_prefill_gemm( - gpu, - &lw.gate_proj, - &scratch.pb_tmp, - &scratch.pb_gate, - n_batch, - Some(&scratch.pb_bf16), - )?; - run_prefill_gemm( - gpu, - &lw.up_proj, - &scratch.pb_tmp, - &scratch.pb_up, - n_batch, - Some(&scratch.pb_bf16), - )?; - } - gpu.gelu_tanh_f32( - &scratch.pb_gate, - &scratch.pb_ffn_hidden, - n_batch * config.hidden_dim, - )?; - gpu.mul_f32( - &scratch.pb_ffn_hidden, - &scratch.pb_up, - &scratch.pb_ffn_hidden, - )?; - run_prefill_gemm( - gpu, - &lw.down_proj, - &scratch.pb_ffn_hidden, - &scratch.pb_ffn_out, - n_batch, - Some(&scratch.pb_bf16), - )?; - } - _ => unreachable!(), - } - - // MoE branch (or dense fallback). HIPFIRE_MOE_BYPASS=1 forces dense - // path even on MoE layers (parity with v1 — used to isolate whether - // a regression lives in apply_moe_branch_batched vs the dense path). - let moe_bypass = hipfire_config::developer_var("HIPFIRE_MOE_BYPASS") - .ok() - .as_deref() - == Some("1"); - match (moe_opt, moe_bypass) { - (Some(moe), false) => { - apply_moe_branch_batched(gpu, config, scratch, moe, post_ffn_norm_ref, n_batch)?; - gpu.add_inplace_f32(&scratch.pb_residual, &scratch.pb_moe_pre2)?; - gpu.scale_f32(&scratch.pb_residual, layer_scalar)?; - } - _ => { - gpu.rmsnorm_batched( - &scratch.pb_ffn_out, - post_ffn_norm_ref, - &scratch.pb_ffn_out, - n_batch, - dim, - config.norm_eps, - )?; - gpu.add_inplace_f32(&scratch.pb_residual, &scratch.pb_ffn_out)?; - gpu.scale_f32(&scratch.pb_residual, layer_scalar)?; - } - } - if hipfire_config::developer_var("HIPFIRE_GEMMA4_DUMP") - .ok() - .as_deref() - == Some("1") - { - let data = gpu.download_f32(&scratch.pb_residual).unwrap_or_default(); - let last = &data[(n_batch - 1) * dim..n_batch * dim]; - let sum: f64 = last.iter().map(|&v| v as f64).sum(); - let min = last.iter().fold(f32::INFINITY, |a, &b| a.min(b)); - let max = last.iter().fold(f32::NEG_INFINITY, |a, &b| a.max(b)); - eprintln!( - "[v2 diag] L{layer_idx} {:?} last-tok hidden: first4={:?} sum={sum:.4e} min={min:.4} max={max:.4}", - layer_type, - &last[..4.min(last.len())] - ); + prefill_rows( + gpu, weights, config, chunk, pos, kv_sliding, kv_full, scratch, + ) + .map_err(hip)?; } } + Ok(()) +} - // ── Final logits: extract last token, run final_norm + lm_head + softcap ── - // The decode loop samples from scratch.logits, which must be filled - // with the last token's logits. Copy last position from pb_residual - // to scratch.tmp, run final_norm, lm_head, and softcap. - let last_offset = (n_batch - 1) * dim * 4; - if let Some(_s) = gpu.active_stream.as_ref() { - gpu.hip.memcpy_dtod_async_at( - &scratch.tmp.buf, - 0, - &scratch.pb_residual.buf, - last_offset, - dim * 4, - _s, - )?; - } else { - gpu.hip.memcpy_dtod_at( - &scratch.tmp.buf, - 0, - &scratch.pb_residual.buf, - last_offset, - dim * 4, - )?; +#[allow(clippy::too_many_arguments)] +fn prefill_rows( + gpu: &mut Gpu, + weights: &Gemma4Weights, + config: &Gemma4Config, + tokens: &[u32], + start_pos: usize, + kv_sliding: &hipfire_runtime::llama::KvCache, + kv_full: &hipfire_runtime::llama::KvCache, + scratch: &Gemma4Scratch, +) -> Result<(), String> { + use crate::program::{ + routed_scratch, Geometry, ProgramBinding, Resident, RowActivations, RowBuffers, RowWidths, + }; + let rows = tokens.len(); + let dim = config.dim; + let end = start_pos + rows; + if end > kv_sliding.physical_cap || end > kv_full.physical_cap { + return Err(format!( + "gemma4 prefill: positions [{start_pos}, {end}) exceed the KV capacity" + )); } - gpu.rmsnorm_f32( - &scratch.tmp, - &weights.final_norm, - &scratch.tmp, - config.norm_eps, + let geo = Geometry::lowered(config); + let head_dim = config.sliding_head_dim.max(config.full_head_dim); + let widths = RowWidths { + dim, + max_q: config.n_heads * head_dim, + max_kv: (config.sliding_n_kv_heads * config.sliding_head_dim) + .max(config.full_n_kv_heads * config.full_head_dim), + ffn: config.hidden_dim, + }; + let mut bufs = RowBuffers::new(gpu); + let acts = RowActivations::new(&mut bufs, gpu, rows, start_pos, widths)?; + let cap = kv_sliding.physical_cap.max(kv_full.physical_cap); + let flash_partials = bufs.alloc( + crate::program::q8_flash_partials_len(gpu, &geo, cap, rows), + "flash_partials", )?; - let ctx = DispatchCtx::new(gpu); - let wr_lm = weights.lm_head.dispatch_ref(); - execute_steps( - gpu, - &ctx, - &[Step::Gemv { - w: &wr_lm, - input: GemvInput::Raw(&scratch.tmp), - out: &scratch.logits, - }], - ) - .map_err(|e| hip_bridge::HipError::new(0, &e.to_string()))?; - if config.final_logit_softcapping > 0.0 { - gpu.logit_softcap_f32( - &scratch.logits, - config.vocab_size, - config.final_logit_softcapping, - )?; + for (row, &token) in tokens.iter().enumerate() { + embed(gpu, weights, &acts.x.sub_offset(row * dim, dim), token, dim) + .map_err(|e| format!("gemma4 prefill embed: {e:?}"))?; } - - Ok(()) + gpu.scale_f32(&acts.x, config.embed_scale) + .map_err(|e| format!("gemma4 prefill embed scale: {e:?}"))?; + let binding = ProgramBinding { + geo, + resident: Resident { + kv_sliding, + kv_full, + pos_buf: &scratch.pos_buf, + v_norm_ones: &scratch.v_norm_ones_full, + flash_partials: &flash_partials, + }, + rows, + position: start_pos, + positions: Some(&acts.positions), + scratch: acts.scratch( + None, + config.enable_moe_block.then(|| routed_scratch(scratch)), + ), + }; + let mut steps = Vec::with_capacity(3 * config.n_layers); + layer_steps(&binding, weights, config, &mut steps)?; + let ctx = DispatchCtx::new(gpu); + execute_steps(gpu, &ctx, &steps).map_err(|e| format!("gemma4 prefill: {e}")) } diff --git a/crates/hipfire-arch-gemma4/src/program.rs b/crates/hipfire-arch-gemma4/src/program.rs index 5db64ece14..f430aba391 100644 --- a/crates/hipfire-arch-gemma4/src/program.rs +++ b/crates/hipfire-arch-gemma4/src/program.rs @@ -27,7 +27,7 @@ use hipfire_dispatch::pipeline::sandwich::{ use hipfire_dispatch::pipeline::{GemvInput, Step}; use hipfire_dispatch::types::RotationPlan; use hipfire_runtime::llama::{KvCache, KvCacheExt, WeightTensor}; -use rdna_compute::GpuTensor; +use rdna_compute::{DType, Gpu, GpuTensor}; /// Attention geometry of one layer type. #[derive(Clone, Copy)] @@ -435,19 +435,182 @@ impl<'a> LayerScratch<'a> { act: &s.ffn_hidden, mlp_out: &s.ffn_out, ple: None, - moe: moe.then_some(RoutedScratch { - input: &s.moe_pre2, - input_rot: &s.moe_pre2_rot, - router_in: &s.moe_router_in, - router_logits: &s.moe_router_logits, - topk_indices: &s.moe_topk_indices, - topk_weights: &s.moe_topk_weights, - gate: &s.moe_expert_gate_batch, - up: &s.moe_expert_up_batch, - act: &s.moe_expert_hidden_batch, - out: &s.moe_cur_moe, - dense_normed: &s.moe_cur_mlp, - }), + moe: moe.then(|| routed_scratch(s)), + } + } +} + +/// Flash-attention partials for `rows` query rows over a Q8 cache of +/// `max_seq` positions, sized for the larger attention geometry. +pub(crate) fn q8_flash_partials_len( + gpu: &Gpu, + geo: &Geometry, + max_seq: usize, + rows: usize, +) -> usize { + [geo.sliding, geo.full] + .into_iter() + .map(|a| { + let tile = rdna_compute::attention::q8_flash_tile_size( + &gpu.arch, + geo.n_heads, + a.n_kv_heads, + a.head_dim, + max_seq, + ); + rows * geo.n_heads * max_seq.div_ceil(tile) * (2 + a.head_dim) + }) + .max() + .unwrap_or(0) +} + +/// The lowered state's single-row routed-expert scratch; batched forwards +/// route row by row through it. +pub(crate) fn routed_scratch(s: &lowered::Gemma4Scratch) -> RoutedScratch<'_> { + RoutedScratch { + input: &s.moe_pre2, + input_rot: &s.moe_pre2_rot, + router_in: &s.moe_router_in, + router_logits: &s.moe_router_logits, + topk_indices: &s.moe_topk_indices, + topk_weights: &s.moe_topk_weights, + gate: &s.moe_expert_gate_batch, + up: &s.moe_expert_up_batch, + act: &s.moe_expert_hidden_batch, + out: &s.moe_cur_moe, + dense_normed: &s.moe_cur_mlp, + } +} + +/// Owner of one batched forward's buffers; frees them on drop. +pub(crate) struct RowBuffers { + gpu: *mut Gpu, + tensors: Vec, +} + +impl RowBuffers { + pub fn new(gpu: &mut Gpu) -> Self { + Self { + gpu, + tensors: Vec::with_capacity(24), + } + } + + /// An F32 `[n]` buffer owned until drop; the returned view aliases it. + pub fn alloc(&mut self, n: usize, label: &str) -> Result { + // SAFETY: a RowBuffers lives inside one forward call and the Gpu + // reference it was built from outlives it. + let gpu = unsafe { &mut *self.gpu }; + let tensor = gpu + .alloc_tensor(&[n], DType::F32) + .map_err(|e| format!("gemma4 batch alloc {label}: {e:?}"))?; + let view = GpuTensor { + // SAFETY: the owning tensor stays in this ledger until every view + // has dropped and the ledger frees it exactly once. + buf: unsafe { tensor.buf.alias() }, + shape: tensor.shape.clone(), + dtype: tensor.dtype, + }; + self.tensors.push(tensor); + Ok(view) + } +} + +impl Drop for RowBuffers { + fn drop(&mut self) { + // SAFETY: see new(). + let gpu = unsafe { &mut *self.gpu }; + for tensor in self.tensors.drain(..) { + gpu.free_tensor(tensor).ok(); + } + } +} + +/// Activations of one `rows`-row forward at positions +/// `[start_pos, start_pos + rows)`. +pub(crate) struct RowActivations { + pub x: GpuTensor, + residual: GpuTensor, + normed: GpuTensor, + /// FWHT scratch for attention projections (`o_proj` reads `max_q`). + pub attn_rot: GpuTensor, + mlp_rot: GpuTensor, + q: GpuTensor, + k: GpuTensor, + v: GpuTensor, + attn_out: GpuTensor, + gate: GpuTensor, + up: GpuTensor, + act: GpuTensor, + mlp_out: GpuTensor, + /// `[rows]` i32 absolute positions. + pub positions: GpuTensor, +} + +/// Widest activation rows a layer program touches. +#[derive(Clone, Copy)] +pub(crate) struct RowWidths { + pub dim: usize, + pub max_q: usize, + pub max_kv: usize, + pub ffn: usize, +} + +impl RowActivations { + pub fn new( + bufs: &mut RowBuffers, + gpu: &mut Gpu, + rows: usize, + start_pos: usize, + w: RowWidths, + ) -> Result { + let mut alloc = |n: usize, label: &str| bufs.alloc(rows * n, label); + let acts = Self { + x: alloc(w.dim, "x")?, + residual: alloc(w.dim, "residual")?, + normed: alloc(w.dim, "normed")?, + attn_rot: alloc(w.max_q.max(w.dim), "attn_rot")?, + mlp_rot: alloc(w.ffn.max(w.dim), "mlp_rot")?, + q: alloc(w.max_q, "q")?, + k: alloc(w.max_kv, "k")?, + v: alloc(w.max_kv, "v")?, + attn_out: alloc(w.max_q, "attn_out")?, + gate: alloc(w.ffn, "gate")?, + up: alloc(w.ffn, "up")?, + act: alloc(w.ffn, "act")?, + mlp_out: alloc(w.dim, "mlp_out")?, + positions: alloc(1, "positions")?, + }; + let bytes: Vec = (start_pos..start_pos + rows) + .flat_map(|p| (p as i32).to_ne_bytes()) + .collect(); + gpu.hip + .memcpy_htod(&acts.positions.buf, &bytes) + .map_err(|e| format!("gemma4 batch htod positions: {e:?}"))?; + Ok(acts) + } + + pub fn scratch<'a>( + &'a self, + ple: Option>, + moe: Option>, + ) -> LayerScratch<'a> { + LayerScratch { + x: &self.x, + residual: &self.residual, + normed: &self.normed, + attn_rot: &self.attn_rot, + mlp_rot: &self.mlp_rot, + q: &self.q, + k: &self.k, + v: &self.v, + attn_out: &self.attn_out, + gate: &self.gate, + up: &self.up, + act: &self.act, + mlp_out: &self.mlp_out, + ple, + moe, } } } diff --git a/crates/hipfire-dispatch/src/pipeline/sandwich.rs b/crates/hipfire-dispatch/src/pipeline/sandwich.rs index f1851b8f0f..61d14ae931 100644 --- a/crates/hipfire-dispatch/src/pipeline/sandwich.rs +++ b/crates/hipfire-dispatch/src/pipeline/sandwich.rs @@ -57,6 +57,15 @@ fn batched_fusion_arch(arch: &str) -> bool { matches!(arch, "gfx1100" | "gfx1201") } +/// True while a calibration collector is armed. Routes that never +/// materialize a projection's unrotated input (fused norm+FWHT) step aside so +/// every projection reaches its `maybe_capture_activation` tap. +fn capturing(gpu: &Gpu) -> bool { + gpu.active_capture.is_some() +} + +/// `y = W x` for one row. Taps the calibration collector with the unrotated +/// input. fn gemv( gpu: &mut Gpu, ctx: &DispatchCtx, @@ -64,6 +73,10 @@ fn gemv( x: &GpuTensor, y: &GpuTensor, ) -> Result<(), DispatchError> { + if w.dtype == DType::BF16 { + return gemm_rows(gpu, ctx, w, x, y, x, 1); + } + gpu.maybe_capture_activation(w.buf, x, 1, w.k); GEMV.run_auto(ctx, gpu, w, x, y) } @@ -105,15 +118,26 @@ fn rotate_x_mq_rows( } /// `y[rows, m] = x[rows, k] · Wᵀ`. MagnumQuant weights rotate `x` into -/// `x_rot` (`[rows, k]`) first. +/// `x_rot` (`[rows, k]`) first. Formats without a batched kernel run one +/// GEMV per row. Taps the calibration collector with the unrotated input. pub fn gemm_rows( gpu: &mut Gpu, + ctx: &DispatchCtx, w: &WeightRef, x: &GpuTensor, y: &GpuTensor, x_rot: &GpuTensor, rows: usize, ) -> Result<(), DispatchError> { + if !supports_batched_projection(w.dtype) { + for row in 0..rows { + let x_row = x.sub_offset(row * w.k, w.k); + let y_row = y.sub_offset(row * w.m, w.m); + gemv(gpu, ctx, w, &x_row, &y_row)?; + } + return Ok(()); + } + gpu.maybe_capture_activation(w.buf, x, rows, w.k); match w.dtype { DType::F32 => gpu .gemm_f32_batched(w.buf, x, y, w.m, w.k, rows) @@ -134,9 +158,33 @@ pub fn gemm_rows( gpu.gemm_mq6g256_batched_lmhead(w.buf, x_rot, y, w.m, w.k, rows) .map_err(hip) } - other => Err(DispatchError::Hip(format!( - "sandwich: dtype {other:?} has no batched projection kernel" - ))), + // BF16 teachers (calibration): stage the F32 input to BF16 for the + // MFMA GEMM, which resolves only where the arch has one. + // ponytail: per-call staging buffer; hold one in scratch if BF16 + // calibration ever becomes throughput-bound. + DType::BF16 => { + let n = rows * w.k; + let staged = gpu.alloc_tensor(&[n], DType::BF16).map_err(hip)?; + let run = gpu + .convert_f32_to_bf16(x, &staged, n) + .map_err(hip) + .and_then(|_| { + GEMM.run_key( + KernelKey::GemmBf16Mfma, + ctx, + gpu, + &GemmParams { + w, + x: &staged, + y, + batch_size: rows, + }, + ) + }); + gpu.free_tensor(staged).map_err(hip)?; + run + } + other => unreachable!("{other:?} is not a batched projection format"), } } @@ -144,7 +192,7 @@ pub fn gemm_rows( pub fn supports_batched_projection(dtype: DType) -> bool { matches!( dtype, - DType::F32 | DType::Q8_0 | DType::MQ4G256 | DType::HFQ4G256 | DType::MQ6G256 + DType::F32 | DType::BF16 | DType::Q8_0 | DType::MQ4G256 | DType::HFQ4G256 | DType::MQ6G256 ) } @@ -170,7 +218,7 @@ fn fused_projection( m: &m, k: weights[0].k, rot_scratch: &[], - batch_size: Some(rows), + batch_size: (rows > 1).then_some(rows), }, ) } @@ -377,7 +425,7 @@ impl SandwichAttentionOp<'_> { .iter() .flatten() .all(|(w, _)| w.dtype == DType::MQ4G256 && w.awq_scale.is_none()); - if all_mq4 { + if all_mq4 && !capturing(gpu) { gpu.fused_rmsnorm_rotate_mq(s.x, self.input_norm, self.x_rot, s.hidden, s.eps) .map_err(hip)?; for (w, out) in self.projections().iter().flatten() { @@ -387,17 +435,15 @@ impl SandwichAttentionOp<'_> { s.norm(gpu, s.x, self.input_norm, s.normed)?; match self.wk() { Some(wk) if self.wq.dtype == DType::Q8_0 && wk.dtype == DType::Q8_0 => { - gpu.fused_gate_up_q8_0( - self.wq.buf, - wk.buf, + fused_projection( + gpu, + ctx, + KernelKey::FusedGateUpQ8_0, + &[&self.wq, wk], s.normed, - self.q, - self.k, - self.wq.m, - wk.m, - self.wq.k, - ) - .map_err(hip)?; + &[self.q, self.k], + 1, + )?; } Some(wk) => { gemv(gpu, ctx, &self.wq, s.normed, self.q)?; @@ -446,12 +492,12 @@ impl SandwichAttentionOp<'_> { )?; } (wk, wv) => { - gemm_rows(gpu, &self.wq, s.normed, self.q, self.x_rot, rows)?; + gemm_rows(gpu, ctx, &self.wq, s.normed, self.q, self.x_rot, rows)?; if let Some(wk) = wk { - gemm_rows(gpu, wk, s.normed, self.k, self.x_rot, rows)?; + gemm_rows(gpu, ctx, wk, s.normed, self.k, self.x_rot, rows)?; } if let (Some(_), Some(wv)) = (wk, wv) { - gemm_rows(gpu, wv, s.normed, self.v, self.x_rot, rows)?; + gemm_rows(gpu, ctx, wv, s.normed, self.v, self.x_rot, rows)?; v_projected = true; } } @@ -683,7 +729,7 @@ pub fn execute_sandwich_attention( if s.rows == 1 { gemv(gpu, ctx, &op.wo, op.attn_out, s.normed)?; } else { - gemm_rows(gpu, &op.wo, op.attn_out, s.normed, op.x_rot, s.rows)?; + gemm_rows(gpu, ctx, &op.wo, op.attn_out, s.normed, op.x_rot, s.rows)?; } s.post_norm_residual(gpu, s.normed, op.post_norm) } @@ -745,10 +791,13 @@ fn dense_mlp( rows, )?; } else { - gemm_rows(gpu, &op.w_gate, s.normed, op.gate, op.x_rot, rows)?; - gemm_rows(gpu, &op.w_up, s.normed, op.up, op.x_rot, rows)?; + gemm_rows(gpu, ctx, &op.w_gate, s.normed, op.gate, op.x_rot, rows)?; + gemm_rows(gpu, ctx, &op.w_up, s.normed, op.up, op.x_rot, rows)?; } - } else if op.w_gate.dtype == DType::MQ4G256 && op.w_up.dtype == DType::MQ4G256 { + } else if op.w_gate.dtype == DType::MQ4G256 + && op.w_up.dtype == DType::MQ4G256 + && !capturing(gpu) + { gpu.fused_rmsnorm_rotate_mq(s.x, op.pre_norm, op.x_rot, s.hidden, s.eps) .map_err(hip)?; gpu.fused_gate_up_hfq4g256( @@ -777,7 +826,7 @@ fn dense_mlp( if rows == 1 { gemv(gpu, ctx, &op.w_down, op.act, op.out)?; } else { - gemm_rows(gpu, &op.w_down, op.act, op.out, op.x_rot, rows)?; + gemm_rows(gpu, ctx, &op.w_down, op.act, op.out, op.x_rot, rows)?; } Ok(()) } @@ -858,12 +907,6 @@ pub fn execute_parallel_moe_mlp( let s = mlp.stream; let e = &op.experts; let r = &op.scratch; - if s.rows != 1 { - return Err(DispatchError::Hip(format!( - "parallel MoE MLP rows={} has no batched executor", - s.rows - ))); - } if e.top_k != 8 { return Err(DispatchError::Hip(format!( "parallel MoE MLP: top_k={} unsupported (indexed kernels are k=8)", @@ -877,11 +920,42 @@ pub fn execute_parallel_moe_mlp( ))); } dense_mlp(gpu, ctx, mlp)?; - s.norm(gpu, mlp.out, op.dense_post_norm, r.dense_normed)?; + if s.rows == 1 { + s.norm(gpu, mlp.out, op.dense_post_norm, r.dense_normed)?; + routed_row(gpu, ctx, e, r, s.residual, s.hidden, s.eps)?; + gpu.add_f32(r.dense_normed, r.out, s.normed).map_err(hip)?; + s.norm(gpu, s.normed, mlp.post_norm, s.normed)?; + } else { + // The indexed expert kernels are single-row: route row by row and + // accumulate into the normed dense output. + s.norm(gpu, mlp.out, op.dense_post_norm, mlp.out)?; + for row in 0..s.rows { + let residual = s.residual.sub_offset(row * s.hidden, s.hidden); + routed_row(gpu, ctx, e, r, &residual, s.hidden, s.eps)?; + let out = mlp.out.sub_offset(row * s.hidden, s.hidden); + gpu.add_inplace_f32(&out, r.out).map_err(hip)?; + } + s.norm(gpu, mlp.out, mlp.post_norm, s.normed)?; + } + s.restore_residual(gpu)?; + gpu.add_inplace_f32(s.x, s.normed).map_err(hip) +} - // Router over the post-attention residual stream. - s.norm(gpu, s.residual, e.pre_norm, r.input)?; - s.norm(gpu, s.residual, e.router_norm, r.router_in)?; +/// `r.out = experts.post_norm(moe(pre_norm(residual)))` for one row; the +/// router reads `router_norm(residual) * router_input_scale`. +fn routed_row( + gpu: &mut Gpu, + ctx: &DispatchCtx, + e: &RoutedExperts<'_>, + r: &RoutedScratch<'_>, + residual: &GpuTensor, + hidden: usize, + eps: f32, +) -> Result<(), DispatchError> { + gpu.rmsnorm_f32(residual, e.pre_norm, r.input, eps) + .map_err(hip)?; + gpu.rmsnorm_f32(residual, e.router_norm, r.router_in, eps) + .map_err(hip)?; scale(gpu, r.router_in, e.router_input_scale)?; gemv(gpu, ctx, &e.router, r.router_in, r.router_logits)?; gpu.moe_softmax_topk_renorm_k8( @@ -895,7 +969,7 @@ pub fn execute_parallel_moe_mlp( // Q8 down accumulates atomically; HFQ4G128 down assigns. gpu.zero_f32(r.out).map_err(hip)?; - let (hidden, mi, k) = (s.hidden, e.hidden_dim, e.top_k); + let (mi, k) = (e.hidden_dim, e.top_k); match e.gate_up_dtype { DType::MQ4G256 => { gpu.rotate_x_mq(r.input, r.input_rot, hidden).map_err(hip)?; @@ -960,12 +1034,7 @@ pub fn execute_parallel_moe_mlp( ) } .map_err(hip)?; - s.norm(gpu, r.out, e.post_norm, r.out)?; - - gpu.add_f32(r.dense_normed, r.out, s.normed).map_err(hip)?; - s.norm(gpu, s.normed, mlp.post_norm, s.normed)?; - s.restore_residual(gpu)?; - gpu.add_inplace_f32(s.x, s.normed).map_err(hip) + gpu.rmsnorm_f32(r.out, e.post_norm, r.out, eps).map_err(hip) } /// Per-layer-input branch: diff --git a/crates/hipfire-loader/src/admission.rs b/crates/hipfire-loader/src/admission.rs index 26283ffb9a..15c87bd1dc 100644 --- a/crates/hipfire-loader/src/admission.rs +++ b/crates/hipfire-loader/src/admission.rs @@ -1395,7 +1395,7 @@ pub fn admit_source_with_options( // Gemma 4 lowered min-context: refuse before teardown/alloc so a // small max_seq leaves the prior model serving. Eager stays exempt. if matches!(arch_id, 13 | 22) { - let use_lowered = hipfire_arch_gemma4::gemma4_source_uses_lowered(hfq, false); + let use_lowered = hipfire_arch_gemma4::gemma4_source_uses_lowered(hfq); hipfire_arch_gemma4::gemma4_context_admission(effective_seq, use_lowered)?; } } diff --git a/crates/hipfire-loader/src/carriers.rs b/crates/hipfire-loader/src/carriers.rs index a60cc770e9..18010b4119 100644 --- a/crates/hipfire-loader/src/carriers.rs +++ b/crates/hipfire-loader/src/carriers.rs @@ -3464,15 +3464,6 @@ impl Carrier for MapleCarrier { // ─── Gemma4Carrier ─────────────────────────────────────────────────── -fn gemma4_use_lowered( - enable_moe_block: bool, - want_batched: bool, - has_drafter: bool, - is_e_series: bool, -) -> bool { - enable_moe_block || (want_batched && !has_drafter && !is_e_series) -} - fn gemma4_validate_drafter_route(is_e_series: bool, has_drafter: bool) -> Result<(), String> { if is_e_series && has_drafter { return Err( @@ -4403,19 +4394,7 @@ impl Carrier for FluxDiffusionCarrier { #[cfg(test)] mod gemma4_route_tests { - use super::{gemma4_use_lowered, gemma4_validate_drafter_route}; - - #[test] - fn e_series_never_enters_dense_lowered_prefill() { - assert!(!gemma4_use_lowered(false, true, false, true)); - } - - #[test] - fn dense_opt_in_and_moe_keep_existing_routes() { - assert!(gemma4_use_lowered(false, true, false, false)); - assert!(!gemma4_use_lowered(false, true, true, false)); - assert!(gemma4_use_lowered(true, false, false, false)); - } + use super::gemma4_validate_drafter_route; #[test] fn e_series_drafter_fails_closed() { diff --git a/crates/hipfire-loader/src/lib.rs b/crates/hipfire-loader/src/lib.rs index c1a45a3229..c5a2a11028 100644 --- a/crates/hipfire-loader/src/lib.rs +++ b/crates/hipfire-loader/src/lib.rs @@ -978,12 +978,6 @@ pub fn gemma4_eagle_spec_len(spec: Option) -> Result { } } -/// Env opt-in for the gemma4 batched/WMMA prefill -/// (`HIPFIRE_BATCHED_PREFILL=1` / `HIPFIRE_WMMA_PREFILL=1`). -pub fn gemma4_batched_prefill_optin(_gpu: &Gpu) -> bool { - gemma4::lowered::batched_prefill_enabled() || gemma4::lowered::wmma_prefill_enabled() -} - // ─── LoadedModel ────────────────────────────────────────────────────── pub struct LoadedModel { diff --git a/crates/hipfire-runtime/examples/calib_sweep.rs b/crates/hipfire-runtime/examples/calib_sweep.rs index ed3141dc8b..9011fd8eae 100755 --- a/crates/hipfire-runtime/examples/calib_sweep.rs +++ b/crates/hipfire-runtime/examples/calib_sweep.rs @@ -6,7 +6,7 @@ // // Taps: // TAP1 llama.rs:1460 weight_gemm → dispatch.rs:2848 (llama 0/1, prefill_forward:1508) -// TAP2 families/gemm.rs:146 GemmFamily::run_key (gemma4 lowered.rs:173, glimmer forward.rs:80, qwen35 migrated sites qwen35.rs:12877/12928/13230) +// TAP2 families/gemm.rs:146 GemmFamily::run_key (glimmer forward.rs:80, qwen35 migrated sites qwen35.rs:12877/12928/13230) // TAP3 forward_batch.rs:79 batched_proj (lfm2 dense, single chokepoint) use hipfire_runtime::calibration::{collect, collect_grouped, CalibForward}; @@ -497,9 +497,6 @@ fn choose_syrk(gpu: &rdna_compute::Gpu, mode: &str) -> (&'static str, bool) { } } } -fn gemma_env_report() -> String { - format!("HIPFIRE_BATCHED_PREFILL={} HIPFIRE_WMMA_PREFILL={} — gemma4 batched prefill behind both 1 (lowered.rs:42-53); if off → per-token GEMV loop n=1, tap still fires but batching win absent", std::env::var("HIPFIRE_BATCHED_PREFILL").ok().as_deref().unwrap_or(""), std::env::var("HIPFIRE_WMMA_PREFILL").ok().as_deref().unwrap_or("")) -} fn check_identity(output: &Path, n_layers: usize, prefix: &str) -> Result<(), String> { let mut hfq = HfqFile::open(output).map_err(|e| format!("open {} identity: {e}", output.display()))?; @@ -750,34 +747,22 @@ fn run_gemma_batched( kv_full: &mut hipfire_runtime::llama::KvCache, scratch: &hipfire_arch_gemma4::lowered::Gemma4Scratch, ) -> Result<(), String> { - let seq = seq.max(1); - for (si, chunk) in toks.chunks(seq).enumerate() { + for (si, chunk) in toks.chunks(seq.max(1)).enumerate() { kv_sliding .clear_gpu(gpu) .map_err(|e| format!("gemma kv sliding clear {si}: {e:?}"))?; kv_full .clear_gpu(gpu) .map_err(|e| format!("gemma kv full clear {si}: {e:?}"))?; - // forward_prefill_batch has MAX_PREFILL_BATCH=128 hard cap; chunk further - let max_b = scratch.max_prefill_batch; - let mut offset = 0usize; - let mut pos = 0usize; - while offset < chunk.len() { - let end = (offset + max_b).min(chunk.len()); - let sub = &chunk[offset..end]; - hipfire_arch_gemma4::lowered::forward_prefill_batch( - gpu, weights, cfg, sub, pos, kv_sliding, kv_full, scratch, + hipfire_arch_gemma4::lowered::forward_prefill_batch( + gpu, weights, cfg, chunk, 0, kv_sliding, kv_full, scratch, + ) + .map_err(|e| { + format!( + "gemma forward_prefill_batch seq {si} B {}: {e:?}", + chunk.len() ) - .map_err(|e| { - format!( - "gemma forward_prefill_batch seq {si} sub {} B {}: {e:?}", - offset, - sub.len() - ) - })?; - pos += sub.len(); - offset = end; - } + })?; } Ok(()) } @@ -1287,14 +1272,6 @@ fn main() { } eprintln!("note: SYRK helper try_rocblas_syrk mirrors gemm.rs:24956 A=B=X m=n=K k=N beta=1.0; gpu.rocblas pub dispatch.rs:523 reachable; auto fallback logged; collector kernel calib_hessian_outer_f32 per constraint."); let _ = &try_rocblas_syrk; - if gemma { - let bp = std::env::var("HIPFIRE_BATCHED_PREFILL").ok(); - let wp = std::env::var("HIPFIRE_WMMA_PREFILL").ok(); - eprintln!("gemma4 env gate: {} ", gemma_env_report()); - if bp.as_deref() != Some("1") || wp.as_deref() != Some("1") { - eprintln!("WARN: gemma4 batched prefill requires HIPFIRE_BATCHED_PREFILL=1 and HIPFIRE_WMMA_PREFILL=1 (lowered.rs:42-53, both default OFF). Current {} {} — shim will still capture (tap families/gemm.rs:146 fires per-token GEMV n=1) but you lose ~50x: the 2048-token chunk will issue 2048 GEMVs trapping at n=1 (~0.02 TFLOP) instead of one MFMA batched GEMM (~1 TFLOP). Set both to 1 before the run that produces the .calib.hfq you will feed to --ldlq.", gemma_env_report(), ""); - } - } if llama { let cfg = hipfire_runtime::hfq::config_from_hfq(&hfq) @@ -1631,7 +1608,8 @@ fn main() { eprintln!("HIPFIRE_NORMALIZE_PROMPT=0 HIPFIRE_GRAPH=0 HIPFIRE_GRAPH_MOE=0"); eprintln!("batched: Qwen35Scratch [seq_len+16] kv_max={} DeltaNetState Q8 forward_prefill_batch MFMA n=batch_size; capture tap families/gemm.rs:146 (not weight_gemm)", kv_max); } else if gemma { - // Gemma4 — WIRED (tap2). Uses forward_prefill_batch (lowered.rs:2581) → run_prefill_gemm (71) → GemmFamily::run_key:119 → tap families/gemm.rs:146 + // Gemma4 — WIRED. forward_prefill_batch runs the declarative layer program; the + // dispatch sandwich projections (gemm_rows / gemv) tap every projection input. let cfg = hipfire_arch_gemma4::lowered::config_from_hfq(&hfq) .unwrap_or_else(|| panic!("gemma4 cfg: missing/invalid config in HFQ")); eprintln!( @@ -1686,12 +1664,19 @@ fn main() { kv_max, ) .unwrap_or_else(|e| panic!("gemma full kv alloc q8 {kv_max}: {e:?}")); - // Honesty guard: forward_prefill_batch must exist and tap must fire; if layer count mismatch or unsupported embed format, run will Err and we exit 1 (honest) — never emit partial hfq. + // Honesty guard: the tap must fire; if layer count mismatch or unsupported embed format, run will Err and we exit 1 (honest) — never emit partial hfq. let lpp = lpp.max(1); let peak = estimate_peak_gemma(&weights, 0, lpp.min(cfg.n_layers)); - eprintln!("grouped n_layers={} lpp={} n_groups={} peak≈{:.1} MB kv_max={} (seq_len+16) scratch.max_prefill_batch={}", cfg.n_layers, lpp, (cfg.n_layers+lpp-1)/lpp, peak as f64/1_048_576.0, kv_max, scratch.max_prefill_batch); + eprintln!( + "grouped n_layers={} lpp={} n_groups={} peak≈{:.1} MB kv_max={} (seq_len+16)", + cfg.n_layers, + lpp, + (cfg.n_layers + lpp - 1) / lpp, + peak as f64 / 1_048_576.0, + kv_max + ); let batch_actual = seq_len.min(n_tok.max(1)); - eprintln!("batching: gemma4::lowered::forward_prefill_batch:2581 → run_prefill_gemm:71 → GemmFamily::run_key:119 → tap families/gemm.rs:146 batch={} KV [seq_len+16] scratch {kv_max} ; capture_names over LayerWeights::Sliding/Full (q/k/v/o, gate/up/down) at lowered.rs:379/414", batch_actual); + eprintln!("batching: gemma4::lowered::forward_prefill_batch → step program → sandwich gemm_rows tap, batch={} KV [seq_len+16]", batch_actual); let t0 = std::time::Instant::now(); let grouped = lpp < cfg.n_layers; let calib_prefix_grouped = calib_prefix.clone(); @@ -1699,7 +1684,7 @@ fn main() { let summary = if grouped { let tc = toks.clone(); collect_grouped(&mut gpu, arch, cfg.n_layers, lpp, Vec::new(), Path::new(&output), - &[("source_model", serde_json::json!(model.clone())), ("corpus", serde_json::json!(corpus.clone())), ("corpus_md5", serde_json::json!(md5.clone())), ("n_calib_tokens", serde_json::json!(n_tok)), ("source_arch_id", serde_json::json!(arch)), ("seq_len", serde_json::json!(seq_len)), ("batch_size", serde_json::json!(batch_actual)), ("syrk_chosen", serde_json::json!(syrk_chosen)), ("syrk_mode", serde_json::json!(syrk_mode.clone())), ("layers_per_pass", serde_json::json!(lpp)), ("batches", serde_json::json!(n_seqs)), ("kv_max", serde_json::json!(kv_max)), ("calib_driver", serde_json::json!("calib_sweep gemma4 forward_prefill_batch:2581 → run_prefill_gemm:71 tap2"))], + &[("source_model", serde_json::json!(model.clone())), ("corpus", serde_json::json!(corpus.clone())), ("corpus_md5", serde_json::json!(md5.clone())), ("n_calib_tokens", serde_json::json!(n_tok)), ("source_arch_id", serde_json::json!(arch)), ("seq_len", serde_json::json!(seq_len)), ("batch_size", serde_json::json!(batch_actual)), ("syrk_chosen", serde_json::json!(syrk_chosen)), ("syrk_mode", serde_json::json!(syrk_mode.clone())), ("layers_per_pass", serde_json::json!(lpp)), ("batches", serde_json::json!(n_seqs)), ("kv_max", serde_json::json!(kv_max)), ("calib_driver", serde_json::json!("calib_sweep gemma4 forward_prefill_batch (step program, sandwich projection tap)"))], |s,e| build_capture_gemma(&weights,s,e, &calib_prefix_grouped), |gpu,_| { run_gemma_batched(gpu,&weights,&cfg,&tc,seq_len,&mut kv_sliding,&mut kv_full,&scratch)?; Ok(CalibForward::default()) } ).unwrap_or_else(|e| { eprintln!("collect_grouped gemma: {e}"); std::process::exit(1); }) @@ -1707,7 +1692,7 @@ fn main() { let cap = build_capture_gemma(&weights, 0, cfg.n_layers, &calib_prefix_single); let tc = toks.clone(); collect(&mut gpu, arch, cap, Vec::new(), Path::new(&output), - &[("source_model", serde_json::json!(model.clone())), ("corpus", serde_json::json!(corpus.clone())), ("corpus_md5", serde_json::json!(md5.clone())), ("n_calib_tokens", serde_json::json!(n_tok)), ("source_arch_id", serde_json::json!(arch)), ("seq_len", serde_json::json!(seq_len)), ("batch_size", serde_json::json!(batch_actual)), ("syrk_chosen", serde_json::json!(syrk_chosen)), ("syrk_mode", serde_json::json!(syrk_mode.clone())), ("layers_per_pass", serde_json::json!(cfg.n_layers)), ("batches", serde_json::json!(n_seqs)), ("kv_max", serde_json::json!(kv_max)), ("calib_driver", serde_json::json!("calib_sweep gemma4 forward_prefill_batch:2581 → run_prefill_gemm:71 tap2"))], + &[("source_model", serde_json::json!(model.clone())), ("corpus", serde_json::json!(corpus.clone())), ("corpus_md5", serde_json::json!(md5.clone())), ("n_calib_tokens", serde_json::json!(n_tok)), ("source_arch_id", serde_json::json!(arch)), ("seq_len", serde_json::json!(seq_len)), ("batch_size", serde_json::json!(batch_actual)), ("syrk_chosen", serde_json::json!(syrk_chosen)), ("syrk_mode", serde_json::json!(syrk_mode.clone())), ("layers_per_pass", serde_json::json!(cfg.n_layers)), ("batches", serde_json::json!(n_seqs)), ("kv_max", serde_json::json!(kv_max)), ("calib_driver", serde_json::json!("calib_sweep gemma4 forward_prefill_batch (step program, sandwich projection tap)"))], |gpu| { run_gemma_batched(gpu,&weights,&cfg,&tc,seq_len,&mut kv_sliding,&mut kv_full,&scratch)?; Ok(CalibForward::default()) } ).unwrap_or_else(|e| { eprintln!("collect gemma: {e}"); std::process::exit(1); }) }; @@ -1758,7 +1743,7 @@ fn main() { cfg.n_layers, cfg.dim, cfg.hidden_dim, kv_max, lpp ); eprintln!("HIPFIRE_NORMALIZE_PROMPT=0 HIPFIRE_GRAPH=0 HIPFIRE_GEMMA4_GRAPH=0"); - eprintln!("batched: Gemma4Scratch [seq_len+16] scratch.max_prefill_batch=128 chunked prefill; KV ONCE seq_len+16"); + eprintln!("batched: step-program prefill (64-row chunks); KV ONCE seq_len+16"); } else if glimmer { let cfg = hipfire_arch_muse_glimmer::config::GlimmerConfig::from_hfq(&hfq) .unwrap_or_else(|e| panic!("glimmer cfg: {e}")); diff --git a/crates/hipfire-runtime/examples/eval_hipfire.rs b/crates/hipfire-runtime/examples/eval_hipfire.rs index 05afe05959..1012b625f3 100644 --- a/crates/hipfire-runtime/examples/eval_hipfire.rs +++ b/crates/hipfire-runtime/examples/eval_hipfire.rs @@ -39,7 +39,7 @@ //! Multi-arch dispatch (Tier-1): mirrors `calib_sweep.rs` arch detection //! llama 0|1, qwen35 5|6, gemma4 13, glimmer 14, lfm2 11. //! Per-arch chunking mirrors calib_sweep: -//! gemma4 sub-chunks to 128 (scratch.max_prefill_batch) and needs q8 KV +//! gemma4 uses forward_prefill_batch (step program, 64-row chunks) with q8 KV //! glimmer chunks 192 (glimmer_prefill_chunk_size / prefill_with_capture) //! lfm2 uses forward_decode_batch_lfm with explicit positions, q8 KV //! qwen35 uses forward_prefill_batch with q8 KV at seq_len+16 @@ -96,7 +96,8 @@ fn main() { let v = argv[i + 1].clone(); if !matches!( v.as_str(), - "q8" | "fp8" | "bf16" + "q8" | "fp8" + | "bf16" | "asym2" | "asym3" | "asym4" @@ -489,7 +490,7 @@ fn main() { kv_max, ) .unwrap() - }, + } "bf16" => { // Quality control: unscaled bf16 K/V (stride D*2 per head, // no scales). Same exact admission as fp8 (gfx1201, dense @@ -532,7 +533,7 @@ fn main() { kv_max, ) .unwrap() - }, + } "asym4" => KvCache::new_gpu_asym4( &mut gpu, config.n_layers, @@ -606,8 +607,11 @@ fn main() { let hidden_buf = if args.scoring_mode == "prefill" { Some( - gpu.alloc_tensor(&[scored_per_chunk + usize::from(pad_scoring_tail), config.dim], DType::F32) - .expect("alloc hidden_buf"), + gpu.alloc_tensor( + &[scored_per_chunk + usize::from(pad_scoring_tail), config.dim], + DType::F32, + ) + .expect("alloc hidden_buf"), ) } else { None @@ -947,13 +951,12 @@ fn main() { p99_kld_per_seq.push(p99); mean_nll_per_seq.push(mean_nll); } + } else if is_gemma { // ----- gemma4 13 — WIRED ----- // Quant dtypes admitted: MQ4G256 / MG4G256 (campaign mq4, quant_type 13/30 alias), // HFQ4G256, HFQ4G128, HFQ6G256, HFQ2/3, Q4K, MQ8G256, MQ6G256, MQ3G256, MQ2G256, // plus F32/BF16 teachers (BF16 kept as BF16 when HIPFIRE_CALIB_BF16=1 else widened to F32). - // See lowered.rs:load_gemma4_weight match (lines 696-736). - // KV must be q8 (KvCache::new_gpu_q8): F32 cache fails with "no implementation for KvWriteF32". - // Prefill sub-chunks to scratch.max_prefill_batch (128) — mirrors calib_sweep gemma arm. + // See lowered.rs:load_gemma4_weight. Batched prefill needs a Q8 KV cache. use hipfire_arch_gemma4::lowered as gemma4; if args.model.is_dir() { eprintln!("eval_hipfire: gemma4 dir not yet wired — use HFQ for arch 13"); @@ -1012,28 +1015,18 @@ fn main() { let mut chunk_nll_count: usize = 0; if args.scoring_mode == "prefill" { - // Prefill prefix with batched forward chunked to 128 (max_prefill_batch), then per-token for scored window. - // This mirrors calib_sweep run_gemma_batched: while offset Date: Sat, 3 Oct 2026 09:24:11 +0200 Subject: [PATCH 06/13] docs(gemma4): record calibration-tool port commit (cherry picked from commit 1bb6ba3378c7d4a02a4d07b1635fca2302108967) --- docs/design/gemma4-declarative-program.md | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/docs/design/gemma4-declarative-program.md b/docs/design/gemma4-declarative-program.md index a9e0c48c60..080b0d3d73 100644 --- a/docs/design/gemma4-declarative-program.md +++ b/docs/design/gemma4-declarative-program.md @@ -118,7 +118,7 @@ batched kernel one GEMV per row. | Eager decode + batched prefill/verify + E-series PLE as `[SandwichAttention, SandwichMlp, PerLayerInput?, Scale?]` (`program.rs`, dispatch `sandwich.rs`); hand arms deleted | landed; bitwise | `30b55f753` | | MoE decode (`ParallelMoeMlp`), lowered decode on the shared program; super-op facade (`lower_variant`, `Gemma4Bindings`) and lowered hand arms deleted; unsupported expert formats refuse at load | landed; 26B byte-identical with the two non-bitwise fusions off, coherent with them on | `287e67d5b` | | EAGLE draft head as one step list (`Gemv` pre-projection, query-only `[SandwichAttention, SandwichMlp, Scale?]` blocks over the target's last slot, norm, `lm_head`, post-projection) | landed; EAGLE text and per-prompt tau identical | `287e67d5b` | -| Calibration tools (`calib_sweep`, `eval_hipfire`, `prefill_parity_gemma4`) on the batched program; dispatch-owned calibration taps; batched `ParallelMoeMlp`; old batched prefill, `HIPFIRE_BATCHED_PREFILL`/`HIPFIRE_WMMA_PREFILL` deleted | landed; tool parity above | (this commit) | +| Calibration tools (`calib_sweep`, `eval_hipfire`, `prefill_parity_gemma4`) on the batched program; dispatch-owned calibration taps; batched `ParallelMoeMlp`; old batched prefill, `HIPFIRE_BATCHED_PREFILL`/`HIPFIRE_WMMA_PREFILL` deleted | landed; tool parity above | `e89080332` | ## Remaining From b095b03efc5152540702d4ad96aa25b6a35bcea3 Mon Sep 17 00:00:00 2001 From: Bjoern Agent Date: Sat, 3 Oct 2026 11:55:04 +0200 Subject: [PATCH 07/13] docs(gemma4): record BF16-reference KLD, pre-port vs declarative tools (cherry picked from commit 399e84be5dd863354a5e9f2bf0670e0717456bb9) --- docs/design/gemma4-declarative-program.md | 6 ++++++ 1 file changed, 6 insertions(+) diff --git a/docs/design/gemma4-declarative-program.md b/docs/design/gemma4-declarative-program.md index 080b0d3d73..764c3d41b9 100644 --- a/docs/design/gemma4-declarative-program.md +++ b/docs/design/gemma4-declarative-program.md @@ -110,6 +110,12 @@ batched kernel one GEMV per row. still hold: 12B/E2B/E4B tokens and top-16 logits bitwise, EAGLE tau unchanged, 26B text identical to the first port, Redline shadow exact with the same tape hash `249f24e5f39d425d`. +- KLD against a BF16 reference (`google/gemma-4-12B-it` `--format oracle`, + pre-port `build_kld_ref_native_gemma4`, WT2 test, 8 x 512, top-256, ref md5 + `ea5e297378c7ff692013af63d8411323`), candidate `gemma4-12b.mq4`, Q8 KV: + per-token scoring 0.7338 pre-port vs 0.7335 now; prefill scoring 0.8150 vs + 0.8242 (prefill wall 160 s vs 110 s). The pre-port `eval_hipfire` needed + the Gemma-branch fix to score at all. ## Progress From c582e1a47596da29ba435333180f513226923dbd Mon Sep 17 00:00:00 2001 From: Bjoern Agent Date: Thu, 8 Oct 2026 09:22:08 +0200 Subject: [PATCH 08/13] fix(railgun-cert): recording inventory follows the Gemma 4 declarative port The port deletes the Gemma 4 lowered hand arms, so their SITES rows go. The sandwich executor's calibration-collector check was named `capturing`, which the inventory scanner reads as the graph-slot capture predicate; it is renamed `calibrating`. The design doc records the rebased-on-beta evidence. (Gemma half of 6e89ab856 and ae9a02f51 from #813.) --- .../hipfire-dispatch/src/pipeline/sandwich.rs | 6 +-- crates/railgun-cert/src/recording.rs | 4 -- docs/design/gemma4-declarative-program.md | 41 ++++++++++++++----- 3 files changed, 34 insertions(+), 17 deletions(-) diff --git a/crates/hipfire-dispatch/src/pipeline/sandwich.rs b/crates/hipfire-dispatch/src/pipeline/sandwich.rs index 61d14ae931..c8ef0fb6d1 100644 --- a/crates/hipfire-dispatch/src/pipeline/sandwich.rs +++ b/crates/hipfire-dispatch/src/pipeline/sandwich.rs @@ -60,7 +60,7 @@ fn batched_fusion_arch(arch: &str) -> bool { /// True while a calibration collector is armed. Routes that never /// materialize a projection's unrotated input (fused norm+FWHT) step aside so /// every projection reaches its `maybe_capture_activation` tap. -fn capturing(gpu: &Gpu) -> bool { +fn calibrating(gpu: &Gpu) -> bool { gpu.active_capture.is_some() } @@ -425,7 +425,7 @@ impl SandwichAttentionOp<'_> { .iter() .flatten() .all(|(w, _)| w.dtype == DType::MQ4G256 && w.awq_scale.is_none()); - if all_mq4 && !capturing(gpu) { + if all_mq4 && !calibrating(gpu) { gpu.fused_rmsnorm_rotate_mq(s.x, self.input_norm, self.x_rot, s.hidden, s.eps) .map_err(hip)?; for (w, out) in self.projections().iter().flatten() { @@ -796,7 +796,7 @@ fn dense_mlp( } } else if op.w_gate.dtype == DType::MQ4G256 && op.w_up.dtype == DType::MQ4G256 - && !capturing(gpu) + && !calibrating(gpu) { gpu.fused_rmsnorm_rotate_mq(s.x, op.pre_norm, op.x_rot, s.hidden, s.eps) .map_err(hip)?; diff --git a/crates/railgun-cert/src/recording.rs b/crates/railgun-cert/src/recording.rs index e00d9b28fe..44ad16ac27 100644 --- a/crates/railgun-cert/src/recording.rs +++ b/crates/railgun-cert/src/recording.rs @@ -1159,11 +1159,7 @@ pub const SITES: &[Site] = &[ Site { file: "hipfire-arch-deepseek4/src/spec_impl.rs", function: "dspark_verify_forward", occurrences: 2, decisions: &["dspark-capture-safe-verify-body", "graph-capture-lifecycle"] }, Site { file: "hipfire-arch-deepseek4/src/spec_impl.rs", function: "redline_dspark_verify_direct", occurrences: 2, decisions: &["dspark-capture-safe-verify-body"] }, Site { file: "hipfire-arch-deepseek4/src/spec_impl.rs", function: "redline_dspark_verify_pm4", occurrences: 1, decisions: &["dspark-capture-safe-verify-body"] }, - Site { file: "hipfire-arch-gemma4/src/lowered.rs", function: "forward_prefill_batch_v2", occurrences: 6, decisions: &["kv-tier-capture-forces-flash"] }, Site { file: "hipfire-arch-gemma4/src/lowered.rs", function: "forward_scratch", occurrences: 3, decisions: &["graph-capture-lifecycle", "replay-recorder-lifecycle"] }, - Site { file: "hipfire-arch-gemma4/src/lowered.rs", function: "full_layer_decode_impl", occurrences: 2, decisions: &["kv-tier-capture-forces-flash"] }, - Site { file: "hipfire-arch-gemma4/src/lowered.rs", function: "run_attend", occurrences: 4, decisions: &["kv-tier-capture-forces-flash"] }, - Site { file: "hipfire-arch-gemma4/src/lowered.rs", function: "sliding_layer_decode_impl", occurrences: 2, decisions: &["kv-tier-capture-forces-flash"] }, Site { file: "hipfire-arch-lfm2moe/src/forward.rs", function: "decode_step", occurrences: 1, decisions: &["replay-enabled-process-route"] }, Site { file: "hipfire-arch-lfm2moe/src/forward.rs", function: "decode_step_with_retained_replay", occurrences: 3, decisions: &["replay-recorder-lifecycle"] }, Site { file: "hipfire-arch-llama/src/arch.rs", function: "forward_scratch_layers", occurrences: 2, decisions: &["kv-tier-capture-forces-flash"] }, diff --git a/docs/design/gemma4-declarative-program.md b/docs/design/gemma4-declarative-program.md index 764c3d41b9..869ef083bb 100644 --- a/docs/design/gemma4-declarative-program.md +++ b/docs/design/gemma4-declarative-program.md @@ -1,7 +1,7 @@ # Gemma 4 declarative layer program (#666 G6) -Status: in progress on `feat/gemma-declarative` (cut from #795 -`feat/qwen35-declarative` at `d101a15d0`). +Status: landed on `feat/gemma-declarative` (#813), rebased onto `beta` with +the #795 Qwen3.5 port it builds on. ## Goal @@ -23,14 +23,15 @@ The #397 super-op facade (`lowered.rs` `lower_variant` / `Gemma4Bindings`) and the hand-written decode arms in `forward.rs` and `lowered.rs` are deleted. Gemma 4 has no EP/TP route, so nothing keeps the super-op program alive for it. -## Today +## Before the port -Two parallel stacks, each with its own config, weights and loader: +Two parallel stacks, each with its own config, weights and loader (both stay; +they now bind the same program): | Stack | Fixtures | KV | Notes | |---|---|---|---| | eager (`config.rs`, `gemma4.rs`, `forward.rs`) | 12B dense (default), E2B/E4B, EAGLE target | Q8 both tiers, `physical_cap = max_seq` | fused qk-norm+RoPE and fused post-norm+residual; hipGraph decode; `forward_batch{,_spec}` B≤64 | -| lowered (`lowered.rs`) | 26B-A4B MoE (always); 12B with `HIPFIRE_{BATCHED,WMMA}_PREFILL=1` | sliding Q8 ring (`window`), full asym3 hd512 | super-op facade; per-token serve prefill; `forward_prefill_batch` only for calibration tools | +| lowered (`lowered.rs`) | 26B-A4B MoE (always); 12B with `HIPFIRE_{BATCHED,WMMA}_PREFILL=1` (removed) | sliding Q8 ring (`window`), full tier Q8 (beta default; asym3 hd512 under `legacy-asym3`) | super-op facade (removed); per-token serve prefill; `forward_prefill_batch` only for calibration tools | ## Vocabulary @@ -39,7 +40,7 @@ New model-neutral ops live in `hipfire-dispatch/src/pipeline/sandwich.rs` | Step | Semantics | Route choice owned by dispatch | |---|---|---| -| `SandwichAttention` | `x = residual + post_norm(o_proj(attend(rope(norm_h(q)), rope(norm_h(k)), norm(v))))` on `input_norm(x)`. Optional V projection (`None` = K=V from pre-norm K); optional weightless V norm; Q prescale; **explicit `Rope { kind: RotateHalf \| PartialHalved { rot_pairs }, theta }`**; KV `write: Option` (`None` = read-only shared slot, used by E-series KV sharing and the drafter); `window` | fused norm+FWHT-rotate + prerotated GEMVs (MQ4), fused Q8 q+k, fused qk-norm+RoPE, fused post-norm+residual, KV tier via `KvTierPlan` | +| `SandwichAttention` | `x = residual + post_norm(o_proj(attend(rope(norm_h(q)), rope(norm_h(k)), norm(v))))` on `input_norm(x)`. Optional V projection (`None` = K=V from pre-norm K); optional weightless V norm; Q prescale; **explicit `Rope { kind: RotateHalf \| PartialHalved { rot_pairs }, theta }`**; `kv_proj: Option` (`None` = query-only attend over a shared slot, no cache write: E-series KV sharing and the drafter); `window` | fused norm+FWHT-rotate + prerotated GEMVs (MQ4), fused Q8 q+k, fused qk-norm+RoPE, fused post-norm+residual, KV tier via `KvTierPlan` | | `SandwichMlp` | `x = residual + post_norm(down(act(gate(n)) * up(n)))`, `n = pre_norm(x)`, `act ∈ {GeluTanh}` | fused norm+rotate + fused gate/up (MQ4), fused post-norm+residual | | `ParallelMoeMlp` | Gemma 4 MoE block: dense GeGLU MLP ∥ routed experts (router input `rms(x)·router_scale/√dim`, softmax top-k renorm, per-expert scale, GeGLU), `post_ffn_norm(norm1(mlp) + norm2(moe))`, residual add | expert dtype routes (MQ4 / HFQ4 / Q8 indexed) | | `PerLayerInput` | E-series PLE branch: `x += post_norm(proj(gelu(gate(x)) * ple[layer]))` | batched strided PLE kernel | @@ -57,10 +58,10 @@ Projections, the final norm and the LM head reuse `Gemv` and same program. The token-level PLE staging and embedding stay in the arch crate, outside the captured body. 2. **MoE.** `ParallelMoeMlp` (`rows == 1` and `rows > 1`); 26B-A4B runs the - shared program on its ring/asym3 KV. -3. **EAGLE/MTP drafter.** The drafter block is `[SandwichAttention(q-only, - write: None, target KV), SandwichMlp, Scale]`; pre/post projection and head - are `Gemv` steps. + shared program on its sliding ring and full-tier KV. +3. **EAGLE/MTP drafter.** The drafter block is `[SandwichAttention(query-only, + kv_proj: None, target KV), SandwichMlp, Scale]`; pre/post projection and + head are `Gemv` steps. 4. **Deletion.** `lowered.rs` super-op facade and hand arms; calibration tools move to the shared batched program. @@ -116,6 +117,26 @@ batched kernel one GEMV per row. per-token scoring 0.7338 pre-port vs 0.7335 now; prefill scoring 0.8150 vs 0.8242 (prefill wall 160 s vs 110 s). The pre-port `eval_hipfire` needed the Gemma-branch fix to score at all. +- Rebased onto `beta` `e268a0798` (gfx1151, beta daemon md5 + `94070f1ce9ce722fcbf4dcf7cc87da58` vs rebased daemon + `ea7b9af9f33b242a81a400486f169f02`), five greedy prompts per row: + 12B AR, 12B prefill batch 64, 12B `HIPFIRE_GEMMA4_EAGLE=1` (strict), E2B, + E4B, E2B/E4B prefill batch 64 have byte-identical token streams and top-16 + logit traces; EAGLE (now opt-in on beta via `HIPFIRE_GEMMA4_EAGLE=1`) is + byte-identical with per-prompt tau 3.368/3.459/2.667/3.447/3.514 on both; + 26B-A4B q8-experts: 3 of 5 identical, 2 diverge late (chars 295 and 322), + coherent; `gemma4-26b-a4b.mq4` (MQ6G256/HFQ4G128 experts) produces garbage + on beta and refuses at load here. `serve_harness.py --thinking off + --sampling greedy`: 12B battery and chain, E4B battery and 26B battery + `transcript_byte_identical=true` against beta. 26B Redline `--pm4 + --skip-prefill` shadow exact on both (beta 1263 launches, hash + `9f57d0ca48889bda`; here 1083, hash `f2675ac01238613c`; both with beta's + default Q8 full tier), decode 38.1 vs 38.2 tok/s. `prefill_parity_gemma4` + 725 tokens: 12B and 26B batched vs per-token same argmax and identical + 24-token continuation (12B 4.2 s vs 33.8 s per-token, 26B 8.2 s vs 19.2 s). + `calib_sweep` 12B coverage 328/328, consistency 0, identity PASS. + `eval_hipfire` per-token 0.733489 and prefill 0.824203: both `.kldseq` + files byte-identical to the pre-rebase branch. ## Progress From d379cc29187220fc2d66373f369545fc99547c04 Mon Sep 17 00:00:00 2001 From: Bjoern Agent Date: Sat, 3 Oct 2026 21:24:15 +0200 Subject: [PATCH 09/13] fix(gemma4): run 26B-A4B MQ4G256V2/MQ6G256 experts; refuse cross-row HFQ4-G128 The local gemma4-26b-a4b.mq4 produced garbage on beta for two reasons: - Its routed-expert down_proj is HFQ4-G128 with K = 704 and was quantized before b4846285e, which packed groups across rows. Every per-row kernel reads such a tensor as garbage. The lowered loader now refuses HFQ4-G128 tensors whose byte size is the cross-row layout, naming the cause. - 4 of its 30 layers carry MQ6G256 expert gate/up, which had no indexed route. MQ6G256 is the HFQ6G256 container over a FWHT-rotated input. The current quantizer emits MQ4G256V2 (qt 44) for the remaining layers, which the lowered loader did not map at all. Routed experts now accept MQ4G256V2 and MQ6G256 gate/up (rotate, then the shared indexed kernel), the loader maps qt 44, and the hipGraph eligibility check uses RoutedExperts::supports instead of a second format list. Verified on gfx1151 with a requantization of google/gemma-4-26B-A4B-it: coherent greedy text, serve battery clean, Redline --pm4 shadow exact (1202 launches), hipGraph decode equal to direct decode, prefill_parity same continuation. Q8-expert 26B and 12B token streams are unchanged. (cherry picked from commit 5ff8e28405c07338854d084e05a64b0710b1de1b) --- crates/hipfire-arch-gemma4/src/carrier.rs | 2 +- crates/hipfire-arch-gemma4/src/lowered.rs | 61 ++++++++++++++++++- .../hipfire-dispatch/src/pipeline/sandwich.rs | 23 ++++++- docs/design/gemma4-declarative-program.md | 11 +++- 4 files changed, 90 insertions(+), 7 deletions(-) diff --git a/crates/hipfire-arch-gemma4/src/carrier.rs b/crates/hipfire-arch-gemma4/src/carrier.rs index 97a0363f24..357bd7cb20 100644 --- a/crates/hipfire-arch-gemma4/src/carrier.rs +++ b/crates/hipfire-arch-gemma4/src/carrier.rs @@ -103,7 +103,7 @@ fn unsupported_expert_formats(weights: &lowered::Gemma4Weights) -> Option HipResult<()> { + let rows = m * k.div_ceil(128) * 72; + if bytes == rows { + return Ok(()); + } + let reason = if bytes == (m * k).div_ceil(128) * 72 { + "packs HFQ4-G128 groups across rows (quantized before hipfire-quantize \ + b4846285e); requantize the model" + } else { + "has an unexpected byte size" + }; + Err(hip_bridge::HipError::new( + 0, + &format!("{name} [{m} x {k}] {reason} ({bytes} bytes, expected {rows})"), + )) +} + /// Load a quantized projection weight. Mirrors qwen35::load_weight_tensor_raw /// but uses the Gemma 4 tensor-name convention (`model.language_model.`). fn load_gemma4_weight( @@ -794,6 +815,7 @@ fn load_gemma4_weight_impl( 11 => DType::HFQ3G256, 12 => DType::HFQ3G128, 13 => DType::MQ4G256, + 44 => DType::MQ4G256V2, 14 => DType::MQ8G256, 15 => DType::MQ6G256, 17 => DType::MQ3G256, @@ -811,6 +833,9 @@ fn load_gemma4_weight_impl( )); } }; + if dtype == DType::HFQ4G128 { + check_hfq4g128_rows(name, m, k, data.len())?; + } let buf = upload_pooled_bytes(gpu, data, &[data.len()])?; // Fault seam: fail after the primary owns its buffer but before the // sidecar attaches. Rollback here frees the primary directly — the outer @@ -1014,6 +1039,7 @@ fn load_moe_layer_extras( 12 => DType::HFQ3G128, // MQ4G256 (13) and MG4G256 (30) share dispatch. 13 | 30 => DType::MQ4G256, + 44 => DType::MQ4G256V2, 14 => DType::MQ8G256, 15 => DType::MQ6G256, 17 => DType::MQ3G256, @@ -1025,6 +1051,21 @@ fn load_moe_layer_extras( )); } }; + if dtype == DType::HFQ4G128 { + let (m, k) = match first_info.shape[..] { + [m, k] => (m as usize, k as usize), + _ => { + return Err(hip_bridge::HipError::new( + 0, + &format!( + "{first_name}: expected a 2-D shape, got {:?}", + first_info.shape + ), + )) + } + }; + check_hfq4g128_rows(&first_name, m, k, bytes_per_expert)?; + } // Concat all experts' bytes into one CPU buffer, upload once. let mut concat = Vec::with_capacity(bytes_per_expert * n_exp); concat.extend_from_slice(first_data); @@ -2285,6 +2326,18 @@ mod scratch_geometry_tests { } } + #[test] + fn hfq4g128_rows_must_start_fresh_groups() { + // K = 704: 5.5 groups per row. The 2-D packer pads each row to 6 groups. + assert!(check_hfq4g128_rows("w", 2816, 704, 2816 * 6 * 72).is_ok()); + // The pre-b4846285e flat packer straddles rows: refused with its cause. + let err = check_hfq4g128_rows("w", 2816, 704, 15488 * 72).unwrap_err(); + assert!(err.to_string().contains("across rows"), "{err}"); + assert!(check_hfq4g128_rows("w", 2816, 704, 2816 * 6 * 72 - 72).is_err()); + // Group-aligned K: both packers agree. + assert!(check_hfq4g128_rows("w", 4, 256, 4 * 2 * 72).is_ok()); + } + #[test] fn decode_partials_cover_q8_decode_tile() { let cfg = dummy_cfg_31b(); @@ -3001,9 +3054,11 @@ pub fn forward_scratch( LayerWeights::Full(lw) => &lw.moe, }; moe.as_ref().map_or(true, |m| m.experts.first().is_some_and(|e| { - matches!(e.gate_up_proj.gpu_dtype, - DType::MQ4G256 | DType::HFQ4G256 | DType::HFQ6G256 | DType::Q8_0) - && e.down_proj.gpu_dtype == DType::HFQ4G128 + e.down_proj.gpu_dtype == DType::HFQ4G128 + && hipfire_dispatch::pipeline::sandwich::RoutedExperts::supports( + e.gate_up_proj.gpu_dtype, + e.down_proj.gpu_dtype, + ) })) }) && hipfire_config::developer_var("HIPFIRE_GEMMA4_DUMP").as_deref() != Ok("1"); diff --git a/crates/hipfire-dispatch/src/pipeline/sandwich.rs b/crates/hipfire-dispatch/src/pipeline/sandwich.rs index c8ef0fb6d1..4169218f62 100644 --- a/crates/hipfire-dispatch/src/pipeline/sandwich.rs +++ b/crates/hipfire-dispatch/src/pipeline/sandwich.rs @@ -861,7 +861,12 @@ impl RoutedExperts<'_> { pub fn supports(gate_up: DType, down: DType) -> bool { matches!( gate_up, - DType::MQ4G256 | DType::HFQ4G256 | DType::HFQ6G256 | DType::Q8_0 + DType::MQ4G256 + | DType::MQ4G256V2 + | DType::MQ6G256 + | DType::HFQ4G256 + | DType::HFQ6G256 + | DType::Q8_0 ) && matches!(down, DType::Q8_0 | DType::HFQ4G128) } } @@ -984,6 +989,22 @@ fn routed_row( ) .map_err(hip)?; } + // FWHT-rotated input; MQ6G256 reads it with the HFQ6G256 kernel. + DType::MQ4G256V2 | DType::MQ6G256 => { + gpu.rotate_x_mq(r.input, r.input_rot, hidden).map_err(hip)?; + crate::pipeline::run_uniform_moe_gate_up( + gpu, + e.gate_up_dtype, + e.gate_up_ptrs, + r.topk_indices, + r.input_rot, + r.gate, + r.up, + 2 * mi, + hidden, + k, + )? + } DType::HFQ4G256 | DType::HFQ6G256 => crate::pipeline::run_uniform_moe_gate_up( gpu, e.gate_up_dtype, diff --git a/docs/design/gemma4-declarative-program.md b/docs/design/gemma4-declarative-program.md index 869ef083bb..04a685cb91 100644 --- a/docs/design/gemma4-declarative-program.md +++ b/docs/design/gemma4-declarative-program.md @@ -125,8 +125,15 @@ batched kernel one GEMV per row. logit traces; EAGLE (now opt-in on beta via `HIPFIRE_GEMMA4_EAGLE=1`) is byte-identical with per-prompt tau 3.368/3.459/2.667/3.447/3.514 on both; 26B-A4B q8-experts: 3 of 5 identical, 2 diverge late (chars 295 and 322), - coherent; `gemma4-26b-a4b.mq4` (MQ6G256/HFQ4G128 experts) produces garbage - on beta and refuses at load here. `serve_harness.py --thinking off + coherent. The local `gemma4-26b-a4b.mq4` produces garbage on beta for two + reasons: its HFQ4-G128 expert `down_proj` (K = 704) was packed across rows + by the quantizer before `b4846285e`, and its 4 MQ6G256 expert layers had no + indexed gate/up route. Such files now refuse at load. A requantization with + the current quantizer (MQ4G256V2/MQ6G256 gate/up, row-padded HFQ4-G128 down) + loads and runs: coherent greedy text, serve battery with no runaway, empty or + attractor turns, Redline `--pm4` shadow exact (1202 launches), hipGraph + decode identical to direct decode, `prefill_parity_gemma4` same argmax and + continuation. `serve_harness.py --thinking off --sampling greedy`: 12B battery and chain, E4B battery and 26B battery `transcript_byte_identical=true` against beta. 26B Redline `--pm4 --skip-prefill` shadow exact on both (beta 1263 launches, hash From 729ee4b62d3a2596215158ef80e7c989b2ea81de Mon Sep 17 00:00:00 2001 From: Bjoern Agent Date: Thu, 8 Oct 2026 09:23:36 +0200 Subject: [PATCH 10/13] docs: rebase bookkeeping for the declarative Gemma 4 port on beta - CHANGELOG entry goes to Unreleased. - Dispatch-bypass ledger: hipfire-arch-gemma4 2/4 -> debt 2/0. - Regenerated env-var lifecycle inventory and crate maps; rustfmt on the Rust files this branch changes against beta (recording.rs is a data table beta keeps unformatted). (Gemma half of adefef614 and 2f9770ee6 from #813.) --- CHANGELOG.md | 30 +++- crates/hipfire-arch-gemma4/map.md | 5 +- crates/hipfire-arch-gemma4/src/carrier.rs | 6 +- crates/hipfire-arch-gemma4/src/lowered.rs | 151 +++++++++++++----- .../src/families/attention.rs | 3 +- crates/hipfire-dispatch/src/pipeline/steps.rs | 13 +- crates/hipfire-loader/src/lib.rs | 15 +- docs/governance/debt-dispatch-bypass.txt | 2 +- 8 files changed, 168 insertions(+), 57 deletions(-) diff --git a/CHANGELOG.md b/CHANGELOG.md index b3a7add27e..3e4b0c9f31 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -129,7 +129,35 @@ - **GDN replay route no longer depends on memory pressure (exactness fix):** the gfx1201 multi-layer GDN tape replay (and the D8 snapshot-source tables) armed lazily at a tape's first accept, and an allocation failure there silently moved that tape to the per-layer replay for life, which commits different DeltaNet bytes on the production model. `GdnTape::new_for_config` now arms the tables when the tape is built (`GdnTape::new_for_dflash` adds the D8 tables), so an allocation failure refuses the lane or request at provision with a logged error; an admitted replay never falls back. The replay's q/k/v/out scratch is now a `Gpu`-owned map keyed by scratch shape (layers, rows, value width): every armed tape of one shape shares one refcounted buffer, released when no tape or replay graph references it and drained on model unload, so eager arming costs a lane only its pointer tables instead of ~75 MB on the 27B. The serve engine's batched verify tape, which is repaired per layer and never replayed, is built with `GdnTape::new_capture_only`. `HIPFIRE_GDN_REPLAY_ML_OFF=1` stays an explicit diagnostic opt-in and is documented as not byte-identical. `GpuPool::alloc` now returns its free lists to HIP and retries once when `hipMalloc` runs out of memory (freed scratch parked in the pool is invisible to HIP; 9 GiB in the reproduction). `DeltaNetSnapshot::new_for` frees partial buffers on a mid-way failure, and `take_dn_checkpoint` logs a skipped checkpoint and frees the snapshot when its save fails. Open: the per-layer `replay_gdn_inner` is still not byte-identical to the multi-layer replay. - **VMM maps reclaim pooled memory:** `hipMemCreate` cannot see buffers parked in `GpuPool` either (12.6–13.7 GiB pooled with 0.1–1.4 GiB reported free in the 4-lane MTP state oracle, which failed its 32K lanes with `hipMemCreate: out of memory`). A whole-segment VMM map (`alloc_vmm_tensor`, `grow_vmm_tensor`) that runs out of memory now returns the pool to HIP and retries once; a granule map (`grow_vmm_tensor_granules`, prefix forks), whose failure aborts its arena, returns the pool first when HIP's free figure cannot cover it plus 64 MiB. Both log a `GpuPool: VMM map ...` line. - **Removed config keys warn instead of failing the load:** 0.4.1 removed `kernel.mw16` (shipped in 0.3.1 and 0.4.0) from the schema outright, so a `config.toml` that still set it made every command fail with `unknown configuration key 'kernel.mw16'`. Removed keys that shipped now live in `hipfire_config::RETIRED_CONFIG_KEYS`; through their deprecation window (`kernel.mw16`: until 0.5.0) the load drops them and run/serve/chat and `hipfire config show` print `warning: ignored kernel.mw16 in : removed in 0.4.1 (…)`. Keys that never shipped (such as `kernel.gfx11_iu4_swizzle`, which only existed on an experiment branch) and typos still fail the load. See [`docs/CONFIG.md`](docs/CONFIG.md#lifecycle-status). - +- Gemma 4 layers run as engine `Step`s: 12B dense, E2B/E4B, the 26B-A4B MoE + and the EAGLE draft head (`gemma-4-*-assistant`) all execute + `[SandwichAttention, SandwichMlp | ParallelMoeMlp, PerLayerInput?, Scale?]` + from `hipfire_dispatch`; dense and E-series batched prefill and EAGLE verify + run the same program with `rows > 1`. The Gemma 4 super-op path and its + `HIPFIRE_FORWARD_LOWERED` hand arms are gone. + - 12B, E2B and E4B decode, batched prefill and EAGLE verify are + byte-identical, logits included; EAGLE tau is unchanged. + - The 26B-A4B now takes the 12B's fused qk-norm+RoPE and post-norm+residual + kernels. Greedy text stays coherent and can diverge late. + - `calib_sweep`, `eval_hipfire` and `prefill_parity_gemma4` prefill Gemma 4 + through the same program (`lowered::forward_prefill_batch`, 64-row + chunks, Q8 KV). The dispatch sandwich projections feed the calibration + collector, so any architecture on these steps calibrates without + per-architecture taps. Gemma calibration and KLD numbers move: the tools + now use the serve kernels. The 26B-A4B MoE block runs batched too. + - `HIPFIRE_BATCHED_PREFILL` and `HIPFIRE_WMMA_PREFILL` are removed; dense + Gemma 4 always loads on the eager stack, MoE on the lowered one. + - The `HIPFIRE_GEMMA4_FUSED_{FFN,QK,QK_ROPE,POSTNORM,ATTN_NORM,PROJ}` developer + switches are removed; the fused routes are always on. The debug switches + `HIPFIRE_GEMMA4_{BASELINE_ATTN,ATTN_VERIFY,GEMM_VERIFY}` and + `HIPFIRE_MOE_{BYPASS,BUCKETED}` are removed with their hand-written paths. + - The 26B-A4B lowered stack loads MQ4G256V2 weights and runs MQ4G256V2 and + MQ6G256 routed experts on the indexed kernels, which is what the current + quantizer emits for it. A MoE checkpoint whose expert formats have no + indexed kernel pair now refuses to load instead of running a host-side + expert loop, and so does one whose HFQ4-G128 expert `down_proj` (K = 704) + was packed across rows by a quantizer before `b4846285e`; both produced + garbage. Requantize such files. ## v0.4.1.1 — release draft diff --git a/crates/hipfire-arch-gemma4/map.md b/crates/hipfire-arch-gemma4/map.md index 8e25e05498..e93e7bd6fc 100644 --- a/crates/hipfire-arch-gemma4/map.md +++ b/crates/hipfire-arch-gemma4/map.md @@ -22,14 +22,15 @@ _Generated by `scripts/check-crate-maps.py` from the tree — do not edit inside ### Modules and exported symbols - [`src/arch.rs`](src/arch.rs): `ARCH_ID`, `Gemma4` -- [`src/carrier.rs`](src/carrier.rs): `Gemma4Bundle`, `Gemma4EagerBundle`, `Gemma4LoweredBundle`, `gemma4_context_admission`, `gemma4_source_uses_lowered`, `gemma4_use_lowered`, `load_gemma4_bundle` +- [`src/carrier.rs`](src/carrier.rs): `Gemma4Bundle`, `Gemma4EagerBundle`, `Gemma4LoweredBundle`, `gemma4_context_admission`, `gemma4_source_uses_lowered`, `load_gemma4_bundle` - [`src/config.rs`](src/config.rs): `Gemma4Config`, `Gemma4ESeriesVariant`, `LayerType`, `RopeType`, `e_series_variant`, `ffn_hidden_dim_for_layer`, `first_kv_shared_layer_idx`, `from_hfq`, `from_metadata_json`, `is_kv_shared_layer`, `kv_shared_source_layer_idx`, `max_ffn_hidden_dim`, `max_head_dim`, `max_kv_dim`, `max_q_dim`, `n_full_kv_slots`, `n_full_layers`, `n_sliding_kv_slots`, `n_sliding_layers` - [`src/drafter.rs`](src/drafter.rs): `DRAFTER_ARCH_ID`, `DrafterLayerWeights`, `DrafterStepOut`, `Gemma4DrafterConfig`, `Gemma4DrafterScratch`, `Gemma4DrafterWeights`, `drafter_step`, `drafter_step_from_concat`, `free_gpu`, `from_hfq`, `load`, `max_head_dim`, `max_q_dim`, `new`, `pre_proj_in` - [`src/forward.rs`](src/forward.rs): `decode_step`, `decode_step_capture`, `decode_step_with_graph`, `forward_batch`, `forward_batch_spec`, `supports_batched_prefill` - [`src/gemma4.rs`](src/gemma4.rs): `FullLayerWeights`, `Gemma4State`, `Gemma4Weights`, `LayerWeights`, `PerLayerBranchWeights`, `PerLayerInputWeights`, `SlidingLayerWeights`, `free_gpu`, `load`, `new`, `new_with_fwht3_max_seq`, `new_with_max_seq`, `reset` - [`src/gemma4_vision.rs`](src/gemma4_vision.rs): `Gemma4VisionConfig`, `Gemma4VisionScratch`, `Gemma4VisionWeights` - [`src/lib.rs`](src/lib.rs): `arch`, `carrier`, `config`, `drafter`, `forward`, `gemma4`, `lowered`, `speculative` -- [`src/lowered.rs`](src/lowered.rs): `ATTEND_FULL`, `ATTEND_SLIDING`, `FullLayerWeights`, `GEMMA4_FLASH_TILE`, `GEMMA4_MAX_PREFILL_BATCH`, `Gemma4Config`, `Gemma4Scratch`, `Gemma4Weights`, `LayerType`, `LayerWeights`, `MOE_BRANCH`, `MoeExpertWeights`, `MoeLayerExtras`, `NORM_INPUT`, `NORM_POST_ATTN`, `NORM_POST_FFN`, `NORM_PRE_FFN`, `PROJ_DOWN`, `PROJ_GATE_UP`, `PROJ_K_FULL`, `PROJ_O`, `PROJ_QK_SLIDING`, `PROJ_Q_FULL`, `PROJ_V_SLIDING`, `RESID_POST_ATTN`, `RESID_POST_FFN`, `RopeType`, `SlidingLayerWeights`, `batched_prefill_enabled`, `config_from_hfq`, `forward_prefill_batch`, `forward_scratch`, `free_gpu`, `gemma4_decode_flash_partials_len`, `gemma4_flash_partials_len`, `gemma4_flash_partials_len_for_config`, `gemma4_pb_flash_partials_len`, `gemma4_pb_flash_partials_len_for_config`, `init_scratch_constants`, `load_weights`, `new`, `wmma_prefill_enabled` +- [`src/lowered.rs`](src/lowered.rs): `FullLayerWeights`, `GEMMA4_FLASH_TILE`, `Gemma4Config`, `Gemma4Scratch`, `Gemma4Weights`, `LayerType`, `LayerWeights`, `MoeExpertWeights`, `MoeLayerExtras`, `RopeType`, `SlidingLayerWeights`, `config_from_hfq`, `forward_prefill_batch`, `forward_scratch`, `free_gpu`, `gemma4_decode_flash_partials_len`, `gemma4_flash_partials_len`, `gemma4_flash_partials_len_for_config`, `init_scratch_constants`, `load_weights`, `new` +- [`src/program.rs`](src/program.rs): `alloc`, `drafter`, `eager`, `layer`, `lowered`, `new`, `scratch` - [`src/speculative.rs`](src/speculative.rs): `Gemma4SpecScratch`, `SpecStepOut`, `free`, `new`, `set_seed_hidden_from`, `spec_step_gemma4_eagle` ### Dependencies (from `Cargo.toml`) diff --git a/crates/hipfire-arch-gemma4/src/carrier.rs b/crates/hipfire-arch-gemma4/src/carrier.rs index 357bd7cb20..2514a70971 100644 --- a/crates/hipfire-arch-gemma4/src/carrier.rs +++ b/crates/hipfire-arch-gemma4/src/carrier.rs @@ -337,7 +337,11 @@ mod tests { fn lowered_full_kv_mode_follows_kv_cache() { use hipfire_runtime::kv_mode::KvMode; for raw in ["", "auto", "q8"] { - assert_eq!(resolve_lowered_full_kv_mode(raw), Ok((KvMode::Q8, None)), "{raw:?}"); + assert_eq!( + resolve_lowered_full_kv_mode(raw), + Ok((KvMode::Q8, None)), + "{raw:?}" + ); } assert_eq!( resolve_lowered_full_kv_mode("legacy-asym3"), diff --git a/crates/hipfire-arch-gemma4/src/lowered.rs b/crates/hipfire-arch-gemma4/src/lowered.rs index 78acce64dd..bf28f873ac 100644 --- a/crates/hipfire-arch-gemma4/src/lowered.rs +++ b/crates/hipfire-arch-gemma4/src/lowered.rs @@ -1930,7 +1930,11 @@ pub fn gemma4_flash_partials_len_for_config(max_seq: usize, config: &Gemma4Confi /// and the Q8 decode tile (`q8_flash_tile_size`) for the full tier at `max_seq` /// and the sliding ring at `min(sliding_window, max_seq)`. The Q8 tile can be /// smaller than 128 (tile32 on gfx1100 up to 8K), which needs more partials. -pub fn gemma4_decode_flash_partials_len(arch: &str, config: &Gemma4Config, max_seq: usize) -> usize { +pub fn gemma4_decode_flash_partials_len( + arch: &str, + config: &Gemma4Config, + max_seq: usize, +) -> usize { let q8_len = |n_kv_heads: usize, head_dim: usize, cap: usize| { let tile = rdna_compute::attention::q8_flash_tile_size( arch, @@ -1943,8 +1947,16 @@ pub fn gemma4_decode_flash_partials_len(arch: &str, config: &Gemma4Config, max_s }; let sliding_cap = config.sliding_window.min(max_seq); gemma4_flash_partials_len_for_config(max_seq, config) - .max(q8_len(config.full_n_kv_heads, config.full_head_dim, max_seq)) - .max(q8_len(config.sliding_n_kv_heads, config.sliding_head_dim, sliding_cap)) + .max(q8_len( + config.full_n_kv_heads, + config.full_head_dim, + max_seq, + )) + .max(q8_len( + config.sliding_n_kv_heads, + config.sliding_head_dim, + sliding_cap, + )) } /// Per-decode scratch, sized once at model-load time against the MAX of @@ -2913,17 +2925,30 @@ fn lowered_graph_binding( projection!(weights.lm_head); mix(weights.embd_format as u64); for v in [ - config.dim, config.n_layers, config.vocab_size, config.n_heads, - config.sliding_head_dim, config.sliding_n_kv_heads, config.sliding_window, - config.full_head_dim, config.full_n_kv_heads, config.hidden_dim, - config.enable_moe_block as usize, config.moe_intermediate_size, - config.num_experts, config.top_k_experts, config.attention_k_eq_v as usize, + config.dim, + config.n_layers, + config.vocab_size, + config.n_heads, + config.sliding_head_dim, + config.sliding_n_kv_heads, + config.sliding_window, + config.full_head_dim, + config.full_n_kv_heads, + config.hidden_dim, + config.enable_moe_block as usize, + config.moe_intermediate_size, + config.num_experts, + config.top_k_experts, + config.attention_k_eq_v as usize, ] { mix(v as u64); } for v in [ - config.norm_eps, config.sliding_rope_theta, config.full_rope_theta, - config.full_partial_rotary_factor, config.final_logit_softcapping, + config.norm_eps, + config.sliding_rope_theta, + config.full_rope_theta, + config.full_partial_rotary_factor, + config.final_logit_softcapping, ] { mix(v.to_bits() as u64); } @@ -2934,9 +2959,13 @@ fn lowered_graph_binding( ($lw:expr) => {{ let lw = $lw; tensors!( - lw.input_layernorm, lw.post_attention_layernorm, - lw.pre_feedforward_layernorm, lw.post_feedforward_layernorm, - lw.layer_scalar, lw.q_norm, lw.k_norm, + lw.input_layernorm, + lw.post_attention_layernorm, + lw.pre_feedforward_layernorm, + lw.post_feedforward_layernorm, + lw.layer_scalar, + lw.q_norm, + lw.k_norm, ); mix(lw.layer_scalar_host.to_bits() as u64); projection!(lw.q_proj); @@ -2949,10 +2978,15 @@ fn lowered_graph_binding( if let Some(m) = &lw.moe { projection!(m.router_proj); tensors!( - m.router_scale, m.per_expert_scale, - m.pre_feedforward_layernorm_2, m.post_feedforward_layernorm_1, - m.post_feedforward_layernorm_2, m.experts_gate_up_pool, - m.experts_down_pool, m.experts_gate_up_ptrs, m.experts_down_ptrs, + m.router_scale, + m.per_expert_scale, + m.pre_feedforward_layernorm_2, + m.post_feedforward_layernorm_1, + m.post_feedforward_layernorm_2, + m.experts_gate_up_pool, + m.experts_down_pool, + m.experts_gate_up_ptrs, + m.experts_down_ptrs, ); if let Some(first) = m.experts.first() { projection!(first.gate_up_proj); @@ -2976,9 +3010,14 @@ fn lowered_graph_binding( } for kv in [kv_sliding, kv_full] { for v in [ - kv.max_seq, kv.physical_cap, kv.compact_offset, - kv.quant_q8 as usize, kv.quant_asym3 as usize, kv.quant_fwht as usize, - kv.quant_asym4 as usize, kv.quant_asym2 as usize, + kv.max_seq, + kv.physical_cap, + kv.compact_offset, + kv.quant_q8 as usize, + kv.quant_asym3 as usize, + kv.quant_fwht as usize, + kv.quant_asym4 as usize, + kv.quant_asym2 as usize, kv.v_mode_bits() as usize, ] { mix(v as u64); @@ -2992,13 +3031,34 @@ fn lowered_graph_binding( } mix(s.pos_buf.as_ptr() as u64); tensors!( - s.x, s.residual, s.tmp, s.q, s.k, s.v, s.attn_out, - s.gate_ffn, s.up_ffn, s.ffn_hidden, s.ffn_out, s.logits, - s.flash_partials, s.v_norm_ones_full, s.moe_cur_mlp, s.moe_pre2, - s.moe_router_in, s.moe_router_logits, s.moe_topk_indices, - s.moe_topk_weights, s.moe_cur_moe, s.moe_expert_gate_up, - s.moe_expert_hidden, s.moe_expert_out, s.moe_pre2_rot, - s.moe_expert_gate_batch, s.moe_expert_up_batch, s.moe_expert_hidden_batch, + s.x, + s.residual, + s.tmp, + s.q, + s.k, + s.v, + s.attn_out, + s.gate_ffn, + s.up_ffn, + s.ffn_hidden, + s.ffn_out, + s.logits, + s.flash_partials, + s.v_norm_ones_full, + s.moe_cur_mlp, + s.moe_pre2, + s.moe_router_in, + s.moe_router_logits, + s.moe_topk_indices, + s.moe_topk_weights, + s.moe_cur_moe, + s.moe_expert_gate_up, + s.moe_expert_hidden, + s.moe_expert_out, + s.moe_pre2_rot, + s.moe_expert_gate_batch, + s.moe_expert_up_batch, + s.moe_expert_hidden_batch, ); h.max(1) } @@ -3083,31 +3143,38 @@ pub fn forward_scratch( && gpu.graphs.graph_exec.is_some() { gpu.hip.stream_write_value32( - gpu.active_stream.as_ref().unwrap(), &scratch.pos_buf, pos as u32, 0, + gpu.active_stream.as_ref().unwrap(), + &scratch.pos_buf, + pos as u32, + 0, )?; - if let Err(e) = gpu.graphs.graph_launch( - &gpu.hip, gpu.device_id, gpu.active_stream.as_ref().unwrap(), - ) { + if let Err(e) = + gpu.graphs + .graph_launch(&gpu.hip, gpu.device_id, gpu.active_stream.as_ref().unwrap()) + { gpu.graphs.graph_destroy(&gpu.hip, gpu.device_id); return Err(e); } } else if use_graph && !gpu.graphs.ar_forward_kernel_dirty { - gpu.hip.memcpy_htod(&scratch.pos_buf, &(pos as i32).to_ne_bytes())?; + gpu.hip + .memcpy_htod(&scratch.pos_buf, &(pos as i32).to_ne_bytes())?; gpu.graphs.drop_captured_graph(&gpu.hip, gpu.device_id); if let Err(e) = gpu.graphs.begin_graph_capture( - &gpu.hip, gpu.device_id, gpu.active_stream.as_ref().unwrap(), + &gpu.hip, + gpu.device_id, + gpu.active_stream.as_ref().unwrap(), ) { gpu.graphs.capture_mode = false; gpu.graphs.graph_destroy(&gpu.hip, gpu.device_id); return Err(e); } - let body = forward_scratch_inner( - gpu, weights, config, pos, kv_sliding, kv_full, scratch, - ); + let body = forward_scratch_inner(gpu, weights, config, pos, kv_sliding, kv_full, scratch); // Always close capture, including a failing body, before dropping its // blobs. Captured work is never launched when either phase fails. let end = gpu.graphs.end_graph_capture( - &gpu.hip, gpu.device_id, gpu.active_stream.as_ref().unwrap(), + &gpu.hip, + gpu.device_id, + gpu.active_stream.as_ref().unwrap(), ); if let Err(e) = body.and(end) { gpu.graphs.capture_mode = false; @@ -3115,9 +3182,10 @@ pub fn forward_scratch( return Err(e); } gpu.graphs.ar_forward_binding = binding; - if let Err(e) = gpu.graphs.graph_launch( - &gpu.hip, gpu.device_id, gpu.active_stream.as_ref().unwrap(), - ) { + if let Err(e) = + gpu.graphs + .graph_launch(&gpu.hip, gpu.device_id, gpu.active_stream.as_ref().unwrap()) + { gpu.graphs.graph_destroy(&gpu.hip, gpu.device_id); return Err(e); } @@ -3130,7 +3198,8 @@ pub fn forward_scratch( if !use_graph && gpu.graphs.graph_exec.is_some() { gpu.graphs.graph_destroy(&gpu.hip, gpu.device_id); } - gpu.hip.memcpy_htod(&scratch.pos_buf, &(pos as i32).to_ne_bytes())?; + gpu.hip + .memcpy_htod(&scratch.pos_buf, &(pos as i32).to_ne_bytes())?; forward_scratch_inner(gpu, weights, config, pos, kv_sliding, kv_full, scratch)?; // A successful direct call JITs the same body before capture. if use_graph { diff --git a/crates/hipfire-dispatch/src/families/attention.rs b/crates/hipfire-dispatch/src/families/attention.rs index 55df5581b8..7e33f2a936 100644 --- a/crates/hipfire-dispatch/src/families/attention.rs +++ b/crates/hipfire-dispatch/src/families/attention.rs @@ -2163,7 +2163,8 @@ fn dispatch_attend( && (64..=8192).contains(&io.batch_size) && (io.batch_size <= 512 || io.batch_size % 512 == 0) && gpu.fa2_gfx11_ctx_admitted(io.max_ctx_len) - && io.max_ctx_len + && io + .max_ctx_len .checked_mul(4 * (256 / 32) * 34) .is_some_and(|bytes| { io.k_cache.buf.size() >= bytes && io.v_cache.buf.size() >= bytes diff --git a/crates/hipfire-dispatch/src/pipeline/steps.rs b/crates/hipfire-dispatch/src/pipeline/steps.rs index 5724352d59..8d4a73e673 100644 --- a/crates/hipfire-dispatch/src/pipeline/steps.rs +++ b/crates/hipfire-dispatch/src/pipeline/steps.rs @@ -942,12 +942,13 @@ pub fn execute_validated_steps_with_moe_hooks<'a>( && read.input.buf.as_ptr() == write.output.buf.as_ptr() && read.rows == write.rows => { - paired_hyper_write(steps, i + 4, read).filter(|(_, next_write)| { - crate::pipeline::layer_ops::hyper_read_prenorm_applies( - gpu, read, next_write, - ) - }) - .map(|(j, next_write)| (j, read, next_write)) + paired_hyper_write(steps, i + 4, read) + .filter(|(_, next_write)| { + crate::pipeline::layer_ops::hyper_read_prenorm_applies( + gpu, read, next_write, + ) + }) + .map(|(j, next_write)| (j, read, next_write)) } _ => None, }; diff --git a/crates/hipfire-loader/src/lib.rs b/crates/hipfire-loader/src/lib.rs index c5a2a11028..75acad3c82 100644 --- a/crates/hipfire-loader/src/lib.rs +++ b/crates/hipfire-loader/src/lib.rs @@ -4579,7 +4579,9 @@ fn load_model_tp_qwen35_dense( } else { errors.join("; ") }; - return Err(format!("MTP head required (mtp=on) but not loaded: {reason}")); + return Err(format!( + "MTP head required (mtp=on) but not loaded: {reason}" + )); } head } else { @@ -5186,9 +5188,14 @@ mod ep_admission_tests { let before = active.request(); let mut effects = LoadEffects::default(); let refusal = attempt_candidate_swap( - &candidate, 1, admission::KvBackendRequest::Explicit(hipfire_runtime::kv_backend::KvBackend::Vmm), - "gfx1100", &mut active, &mut effects, - ).unwrap_err(); + &candidate, + 1, + admission::KvBackendRequest::Explicit(hipfire_runtime::kv_backend::KvBackend::Vmm), + "gfx1100", + &mut active, + &mut effects, + ) + .unwrap_err(); assert!(refusal.contains("vmm") && refusal.contains("unsupported")); assert_eq!(effects, LoadEffects::default()); assert_eq!(active.request(), before); diff --git a/docs/governance/debt-dispatch-bypass.txt b/docs/governance/debt-dispatch-bypass.txt index 552d9e09e1..084b05a766 100644 --- a/docs/governance/debt-dispatch-bypass.txt +++ b/docs/governance/debt-dispatch-bypass.txt @@ -34,7 +34,7 @@ hipfire-arch-qwen35 partial 184 191 migrate last; MTP routed MoE now dispatch-layer owned; +10 = DFlash fp16-X GEMM entries (S3/S4 launch fusion); -20 = MQ v2 lm_head match deduped in speculative.rs; +7 = MQ4-Lloyd (qt=52) FP8-LUT prefill arms (gemm_*_fp8_lloyd) pending registry keys with LUT plumbing (gfx11 MMQ-LUT PR); +15 = MQ4-Lloyd gfx11 MMQ-LUT prefill arms (gemm_*_mmq_lloyd), same registry-key debt as the FP8-LUT arms; +40 = speculation-decode additions recorded by the SCS hygiene pass (was carried unrecorded on the beta integration base); +1 from the SCS port PR series (landed with the SCS fold); +2 landed on beta since this ledger was frozen (recorded on rebase, not new debt); -5 = sealed MoE prefill moved into dispatch (PR #755 via #774); +4 = gfx1151 A4-fusion gate/up + down-sidecar arms (gemm_gate_up_silu_a4_*, gemm_*_iu4_down_prepared) beside the Iu4H direct calls they mirror, same registry-key debt hipfire-arch-deepseek4 debt 19 0 cheapest full conversion; no registry use to reconcile -hipfire-arch-gemma4 partial 2 4 +hipfire-arch-gemma4 debt 2 0 -15 = Gemma 4 layers run as hipfire-dispatch sandwich Steps (#813); the 2 left are the PLE Q8 projection GEMMs (forward.rs) hipfire-arch-lfm2moe debt 17 0 cheapest full conversion; no registry use to reconcile hipfire-arch-minimax debt 12 0 hipfire-arch-muse-glimmer partial 12 17 From 95035c904fc51c596c7b39a6662959517811eebc Mon Sep 17 00:00:00 2001 From: Bjoern Agent Date: Thu, 8 Oct 2026 14:13:36 +0200 Subject: [PATCH 11/13] refactor(gemma4): one weight stack, engine session cache, generic EAGLE seams Bring Gemma 4 onto the architecture Qwen4 introduced. One stack. The 26B-A4B MoE loads through the same config, weights, state and forward as the 12B and E-series; lowered.rs is deleted. Its sliding cache is position-indexed (no 1024-row ring), so batched prefill reaches it on Q8/Q8 loads. Decode hipGraph stays off when an expert down_proj is Q8_0 (atomic sum, not replay-exact) or a launch recorder is armed. Session cache. Gemma4Bundle (bundle.rs) implements SessionState: Q8 K/V rows of every cache slot as RowStreams, scoped by prefill width, a boundary every 1024 prompt tokens. The route-generic live/plan/begin/ commit logic moves from the Qwen4 bundle into hipfire_runtime::session_cache::SessionDriver; Qwen4 and Gemma 4 both use it. Each Gemma 4 request now plans against its own prompt; before, it continued the previous request's KV in the same process. Spec seams. Gemma4Drafter is an MtpDrafter behind MtpSpeculator, the bundle a SpecTarget, GemmaSpecEmit the emitter, and generate_spec the loop. The draft head is one step list ending in DraftHead. The inline EAGLE loop, speculative.rs and the infer_gemma4_spec example are deleted. EAGLE stays opt-in (HIPFIRE_GEMMA4_EAGLE=1). Evidence (gfx1151; PR-head daemon 83c3ca2c175e0ac66b47ce1802e8f815, this daemon 7181009cd7a133fc40592ea7b8f1601e; daemon reset between prompts on both arms): - 12B, 12B pb64, E2B, E4B, E4B pb64: tokens and top-16 logit traces bitwise. 26B q8-experts and requant MQ4: text identical; decode 36.1 vs 36.4 and 43.2 vs 43.9 tok/s. - serve_harness --compare-transcript: 12B battery+chain, E4B battery, 26B battery identical. - prefill_parity_gemma4 12B and 26B: logits fnv identical to the pre-port run; eval_hipfire per-token/prefill .kldseq byte-identical; calib_sweep 328/328, consistent, identity PASS. - EAGLE spec 2/3/5: 4 of 5 prompts equal AR. p1 leaves AR at step 91, AR top-2 margin 0.021 logits; batched verify logits differ from decode by ~1e-2 from the first verify row, so equality holds only above that margin. --- CHANGELOG.md | 48 +- crates/hipfire-arch-gemma4/Cargo.toml | 4 - .../examples/infer_gemma4_spec.rs | 403 -- .../examples/prefill_parity_gemma4.rs | 96 +- .../examples/verify_drafter_gemma4.rs | 76 +- crates/hipfire-arch-gemma4/map.md | 21 +- crates/hipfire-arch-gemma4/src/bundle.rs | 393 ++ crates/hipfire-arch-gemma4/src/carrier.rs | 406 +- crates/hipfire-arch-gemma4/src/config.rs | 37 +- crates/hipfire-arch-gemma4/src/drafter.rs | 510 ++- crates/hipfire-arch-gemma4/src/emit.rs | 641 ++++ crates/hipfire-arch-gemma4/src/forward.rs | 231 +- crates/hipfire-arch-gemma4/src/gemma4.rs | 1722 +++++---- crates/hipfire-arch-gemma4/src/lib.rs | 31 +- crates/hipfire-arch-gemma4/src/lowered.rs | 3403 ----------------- crates/hipfire-arch-gemma4/src/mtp.rs | 385 ++ crates/hipfire-arch-gemma4/src/program.rs | 172 +- crates/hipfire-arch-gemma4/src/speculative.rs | 252 -- crates/hipfire-arch-qwen4/src/bundle.rs | 492 +-- crates/hipfire-daemon/src/main.rs | 11 +- crates/hipfire-dispatch/map.md | 7 +- crates/hipfire-dispatch/src/families/gemm.rs | 2 +- .../src/pipeline/layer_ops.rs | 13 +- crates/hipfire-generate/map.md | 2 +- crates/hipfire-generate/src/ar.rs | 53 +- crates/hipfire-generate/src/dense.rs | 1385 +------ crates/hipfire-generate/src/redline.rs | 187 +- crates/hipfire-loader/map.md | 2 +- crates/hipfire-loader/src/admission.rs | 6 - crates/hipfire-loader/src/carriers.rs | 315 +- crates/hipfire-loader/src/lib.rs | 127 +- .../examples/build_kld_ref_native_gemma4.rs | 60 +- .../hipfire-runtime/examples/calib_sweep.rs | 96 +- .../examples/debug_gemma4_attention.rs | 63 +- .../hipfire-runtime/examples/eval_hipfire.rs | 84 +- .../examples/eval_hipfire_gemma4.rs | 60 +- .../hipfire-runtime/examples/gemma4_oracle.rs | 100 +- crates/hipfire-runtime/map.md | 4 +- crates/hipfire-runtime/src/kv_mode.rs | 22 +- crates/hipfire-runtime/src/session_cache.rs | 375 +- crates/railgun-cert/src/recording.rs | 5 +- .../examples/build_kld_ref_native.rs | 58 +- docs/CONFIG.md | 11 +- docs/MODELS.md | 7 +- docs/architecture-ids.md | 2 +- docs/design/gemma4-declarative-program.md | 83 +- docs/governance/debt-dispatch-bypass.txt | 2 +- scripts/oracle_gemma4_26b.py | 2 +- 48 files changed, 4291 insertions(+), 8176 deletions(-) delete mode 100644 crates/hipfire-arch-gemma4/examples/infer_gemma4_spec.rs create mode 100644 crates/hipfire-arch-gemma4/src/bundle.rs create mode 100644 crates/hipfire-arch-gemma4/src/emit.rs delete mode 100644 crates/hipfire-arch-gemma4/src/lowered.rs create mode 100644 crates/hipfire-arch-gemma4/src/mtp.rs delete mode 100644 crates/hipfire-arch-gemma4/src/speculative.rs diff --git a/CHANGELOG.md b/CHANGELOG.md index 3e4b0c9f31..db498831c5 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -140,24 +140,52 @@ - The 26B-A4B now takes the 12B's fused qk-norm+RoPE and post-norm+residual kernels. Greedy text stays coherent and can diverge late. - `calib_sweep`, `eval_hipfire` and `prefill_parity_gemma4` prefill Gemma 4 - through the same program (`lowered::forward_prefill_batch`, 64-row + through the same program (`forward::forward_prefill_batch`, 64-row chunks, Q8 KV). The dispatch sandwich projections feed the calibration collector, so any architecture on these steps calibrates without per-architecture taps. Gemma calibration and KLD numbers move: the tools now use the serve kernels. The 26B-A4B MoE block runs batched too. - - `HIPFIRE_BATCHED_PREFILL` and `HIPFIRE_WMMA_PREFILL` are removed; dense - Gemma 4 always loads on the eager stack, MoE on the lowered one. + - `HIPFIRE_BATCHED_PREFILL` and `HIPFIRE_WMMA_PREFILL` are removed. - The `HIPFIRE_GEMMA4_FUSED_{FFN,QK,QK_ROPE,POSTNORM,ATTN_NORM,PROJ}` developer switches are removed; the fused routes are always on. The debug switches `HIPFIRE_GEMMA4_{BASELINE_ATTN,ATTN_VERIFY,GEMM_VERIFY}` and `HIPFIRE_MOE_{BYPASS,BUCKETED}` are removed with their hand-written paths. - - The 26B-A4B lowered stack loads MQ4G256V2 weights and runs MQ4G256V2 and - MQ6G256 routed experts on the indexed kernels, which is what the current - quantizer emits for it. A MoE checkpoint whose expert formats have no - indexed kernel pair now refuses to load instead of running a host-side - expert loop, and so does one whose HFQ4-G128 expert `down_proj` (K = 704) - was packed across rows by a quantizer before `b4846285e`; both produced - garbage. Requantize such files. + - The 26B-A4B loads MQ4G256V2 weights and runs MQ4G256V2 and MQ6G256 routed + experts on the indexed kernels, which is what the current quantizer emits + for it. A MoE checkpoint whose expert formats have no indexed kernel pair + now refuses to load instead of running a host-side expert loop, and so + does one whose HFQ4-G128 expert `down_proj` (K = 704) was packed across + rows by a quantizer before `b4846285e`; both produced garbage. Requantize + such files. +- Gemma 4 runs on one weight stack: the 26B-A4B MoE loads through the same + config, weights, state and forward as the 12B and E-series, and the separate + `lowered.rs` stack is gone. Its sliding cache is position-indexed like the + 12B's (no 1024-row ring), so batched prefill (default on gfx1100/gfx1201) + now reaches the 26B on Q8/Q8 loads. + Its hipGraph decode is off when expert `down_proj` is Q8_0, whose atomic + sum is not replay-exact. Text on the q8-experts and requantized MQ4 26B + fixtures is unchanged from the previous stack. +- Gemma 4 joins the engine session cache (`memory.session_cache_bytes`): + prefill snapshots the Q8 KV rows every 1024 prompt tokens and continues the + previous turn's live state in place when the new prompt extends it. A + repeated long prompt restores the snapshot instead of re-prefilling it. + Each request now starts from its own prompt; previously a Gemma 4 request + continued the previous request's KV in the same process until the context + filled, so its output depended on what ran before it. The route-generic + live/plan/begin/commit logic moved from Qwen4 into + `hipfire_runtime::session_cache::SessionDriver`, which Qwen4 and Gemma 4 + both use. +- Gemma 4 EAGLE (`params.drafter`, opt-in `HIPFIRE_GEMMA4_EAGLE=1`) runs on + the generic speculative seams: `Gemma4Drafter` is an `MtpDrafter` behind + `MtpSpeculator`, the bundle is a `SpecTarget`, `GemmaSpecEmit` routes the + thought channel, and `generate_spec` drives the loop. The inline EAGLE loop + and `speculative.rs` are deleted, along with the `infer_gemma4_spec` example. + The draft head is one step list ending in the engine `DraftHead`. The first + token now comes from prefill, and each window pairs its seed with the hidden + that predicted it, so per-prompt tau moves slightly. Batched verify is not + bitwise-equal to single-row decode, so greedy EAGLE can leave the AR stream + where two logits are nearly tied. Measured on gfx1151: one of five fixture + prompts at a 0.021-logit margin. ## v0.4.1.1 — release draft diff --git a/crates/hipfire-arch-gemma4/Cargo.toml b/crates/hipfire-arch-gemma4/Cargo.toml index 5295942cd5..3a0d3513bd 100644 --- a/crates/hipfire-arch-gemma4/Cargo.toml +++ b/crates/hipfire-arch-gemma4/Cargo.toml @@ -48,10 +48,6 @@ required-features = ["lab"] name = "hfq4_gemm_parity" required-features = ["lab"] -[[example]] -name = "infer_gemma4_spec" -required-features = ["lab"] - [[example]] name = "prefill_parity_gemma4" required-features = ["lab"] diff --git a/crates/hipfire-arch-gemma4/examples/infer_gemma4_spec.rs b/crates/hipfire-arch-gemma4/examples/infer_gemma4_spec.rs deleted file mode 100644 index e4f941c61e..0000000000 --- a/crates/hipfire-arch-gemma4/examples/infer_gemma4_spec.rs +++ /dev/null @@ -1,403 +0,0 @@ -// SPDX-License-Identifier: Apache-2.0 -// Copyright (c) 2026 Kaden Schutt -// hipfire — see LICENSE and NOTICE in the project root. - -//! Gemma 4 EAGLE speculative-decode driver — the capstone that turns the -//! validated drafter + batched verify into actual speedup. -//! -//! Loads the arch-13 TARGET HFQ + the arch-22 DRAFTER HFQ, prefills per-token -//! via `decode_step` (capturing the last position's post-`model.norm` hidden as -//! the drafter seed), then loops `spec_step_gemma4_eagle` until EOS or `--max`. -//! Reports the generated text, mean committed-per-round (τ), tok/s, and the -//! total number of draft rounds. -//! -//! Correctness mode (`--check-eager`): also runs the eager `decode_step` greedy -//! AR sequence over the same prompt and prints both token-id streams + a diff. -//! With temp 0 the spec-decode committed sequence is provably the target's -//! greedy AR sequence, so the two MUST be byte-identical. -//! -//! Measured (hiptrx gfx1201, gemma4-12B, code prompt, max 256): -//! - Q8 target: spec == eager BYTE-IDENTICAL, deterministic, τ≈3.1–4.5 over -//! draft-len 3..6. (Q8 spec is correctness-proven but lm_head-bound, so -//! ~13 tok/s < the 34.8 tok/s Q8 AR baseline — the per-position full-vocab -//! lm_head dominates each round on the bigger Q8 weights.) -//! - MQ4-attn target (gemma4-12b-mq4attn): spec == eager BYTE-IDENTICAL for -//! draft-len ≤ 4; **80.5 tok/s @ draft-len 3 vs 50.1 tok/s AR = +61%**, -//! τ≈3.08. draft-len ≥ 5 currently hits a PRE-EXISTING, non-deterministic -//! batched-MQ4 `forward_batch` fault (illegal access) on gfx12 — reproduced -//! standalone on the UNMODIFIED `forward_batch` (`verify_batch_gemma4 -//! --bs 16`), so it is a batched-MQ4-GEMM kernel limitation, NOT a spec-loop -//! bug. The spec loop itself is verified correct on the deterministic Q8 -//! path. See the warmup note below. -//! - UNION target (gemma4-12b-it.union.mq4: Q8 attn + Q8 tied lm_head, MQ4 -//! FFN), 2026-06-10, after `proj_gemm_batched` Q8 routed through -//! `gemm_q8_0_batched_chunked` (WMMA on gfx12): spec == eager -//! BYTE-IDENTICAL for draft-len 3..=5; 86-93 tok/s @ dl3 and ~106 tok/s -//! @ dl5 vs 51.0 tok/s AR (1.7-2.1x). Before that routing fix the scalar -//! `gemm_q8_0_batched` cost 334 us/call x 184 attn projections per -//! verify (~61 ms/round, 61% of GPU time per rocprofv3) -> 32 tok/s = -//! 0.63x AR. The dl >= 5 fault note above predates the fix and did NOT -//! reproduce at dl5 (b=6) on the union target across repeated -//! --check-eager runs; post-fix `verify_batch_gemma4 --bs 8,16,32`: -//! B=8 PASS, B=16 PASS (the old illegal-access fault is gone), B=32 -//! argmax near-tie flip at cosine 0.9999 (numerical, no fault) — all -//! far above the daemon's dl <= 5 (b <= 6) envelope. -//! -//! Usage: -//! infer_gemma4_spec --model --draft-model \ -//! [--prompt ] [--token-ids 2,...] [--draft-len N] [--max M] \ -//! [--temperature T(0 only)] [--check-eager] - -#[cfg(not(feature = "deltanet"))] -fn main() { - eprintln!("build with --features deltanet"); -} - -#[cfg(feature = "deltanet")] -fn main() { - use hipfire_arch_gemma4::config::Gemma4Config; - use hipfire_arch_gemma4::drafter::{ - Gemma4DrafterConfig, Gemma4DrafterScratch, Gemma4DrafterWeights, - }; - use hipfire_arch_gemma4::forward::{decode_step, forward_batch}; - use hipfire_arch_gemma4::gemma4::{Gemma4State, Gemma4Weights}; - use hipfire_arch_gemma4::speculative::{spec_step_gemma4_eagle, Gemma4SpecScratch}; - use hipfire_runtime::hfq::HfqFile; - use hipfire_runtime::tokenizer::Tokenizer; - use std::path::PathBuf; - - let argv: Vec = std::env::args().collect(); - let mut model: Option = None; - let mut draft_model: Option = None; - let mut prompt = "The capital of France is".to_string(); - let mut token_ids: Option> = None; - let mut draft_len: usize = 4; - let mut max: usize = 256; - let mut temperature: f32 = 0.0; - let mut check_eager = false; - let mut i = 1; - while i < argv.len() { - match argv[i].as_str() { - "--model" => { - model = Some(PathBuf::from(&argv[i + 1])); - i += 2; - } - "--draft-model" => { - draft_model = Some(PathBuf::from(&argv[i + 1])); - i += 2; - } - "--prompt" => { - prompt = argv[i + 1].clone(); - i += 2; - } - "--prompt-file" => { - prompt = std::fs::read_to_string(&argv[i + 1]) - .expect("--prompt-file: read"); - i += 2; - } - "--token-ids" => { - token_ids = Some( - argv[i + 1] - .split(',') - .filter(|s| !s.is_empty()) - .map(|s| s.trim().parse::().expect("--token-ids")) - .collect(), - ); - i += 2; - } - "--draft-len" => { - draft_len = argv[i + 1].parse().expect("--draft-len"); - i += 2; - } - "--max" => { - max = argv[i + 1].parse().expect("--max"); - i += 2; - } - "--temperature" => { - temperature = argv[i + 1].parse().expect("--temperature"); - i += 2; - } - "--check-eager" => { - check_eager = true; - i += 1; - } - other => { - eprintln!("unknown arg {other}"); - std::process::exit(1); - } - } - } - let model = model.expect("--model required"); - let draft_model = draft_model.expect("--draft-model required"); - assert!(draft_len >= 1, "--draft-len must be ≥ 1"); - assert_eq!(temperature, 0.0, "only greedy (temp 0) is supported"); - - let mut gpu = rdna_compute::Gpu::init().expect("gpu init"); - - // ── Target ── - let hfq = HfqFile::open(&model).expect("open target"); - let cfg = Gemma4Config::from_hfq(&hfq).expect("target config"); - eprintln!( - "target dim={} layers={} sliding={} full={} vocab={} softcap={}", - cfg.dim, - cfg.n_layers, - cfg.n_sliding_layers(), - cfg.n_full_layers(), - cfg.vocab_size, - cfg.final_logit_softcapping - ); - let tok = if token_ids.is_none() { - Some(Tokenizer::from_hfq_metadata(&hfq.metadata_json).expect("tokenizer")) - } else { - None - }; - let t_load = std::time::Instant::now(); - let weights = Gemma4Weights::load(&hfq, &cfg, &mut gpu).expect("target weights"); - - // ── Drafter ── - let dhfq = HfqFile::open(&draft_model).expect("open drafter"); - let dcfg = Gemma4DrafterConfig::from_hfq(&dhfq).expect("drafter config"); - assert_eq!( - dcfg.backbone_hidden, cfg.dim, - "drafter backbone must equal target hidden" - ); - eprintln!( - "drafter hidden={} layers={} backbone={} vocab={}", - dcfg.hidden, dcfg.n_layers, dcfg.backbone_hidden, dcfg.vocab_size - ); - let dweights = Gemma4DrafterWeights::load(&dhfq, &dcfg, &mut gpu).expect("drafter weights"); - let mut dscratch = Gemma4DrafterScratch::new(&mut gpu, &dcfg).expect("drafter scratch"); - eprintln!("loaded target+drafter in {:.1}s", t_load.elapsed().as_secs_f64()); - - // ── Prompt → tokens (prepend BOS) ── - let mut prompt_ids = match token_ids { - Some(ids) => ids, - None => tok.as_ref().unwrap().encode(&prompt), - }; - if prompt_ids.first() != Some(&cfg.bos_token) { - prompt_ids.insert(0, cfg.bos_token); - } - eprintln!("prompt {:?} → {} tokens", prompt, prompt_ids.len()); - let max_seq = prompt_ids.len() + max + 16; - - let argmax = |v: &[f32]| -> u32 { - let mut bi = 0u32; - let mut bv = f32::NEG_INFINITY; - for (k, &x) in v.iter().enumerate() { - if x > bv { - bv = x; - bi = k as u32; - } - } - bi - }; - let is_eos = |t: u32| matches!(t, 1 | 106); - - // ════════════════════════════════════════════════════════════════════ - // Optional eager reference (for the correctness gate). - // ════════════════════════════════════════════════════════════════════ - let eager_gen: Option> = if check_eager { - let mut st = Gemma4State::new_with_max_seq(&mut gpu, &cfg, max_seq).expect("eager state"); - let mut logits = Vec::new(); - for (pos, &t) in prompt_ids.iter().enumerate() { - logits = - decode_step(&cfg, &weights, &mut st, &mut gpu, t, pos as u32).expect("eager prefill"); - } - let mut gen = Vec::new(); - let mut pos = prompt_ids.len(); - for _ in 0..max { - let next = argmax(&logits); - if is_eos(next) { - break; - } - gen.push(next); - logits = decode_step(&cfg, &weights, &mut st, &mut gpu, next, pos as u32) - .expect("eager decode"); - pos += 1; - } - Some(gen) - } else { - None - }; - - // ════════════════════════════════════════════════════════════════════ - // Spec-decode run. - // ════════════════════════════════════════════════════════════════════ - let mut state = Gemma4State::new_with_max_seq(&mut gpu, &cfg, max_seq).expect("spec state"); - let mut sp = Gemma4SpecScratch::new(&mut gpu, &cfg, draft_len).expect("spec scratch"); - - // ── Prefill per-token via `decode_step` (correctness-first; matches the - // validated drafter harness). The seed TOKEN is the last prompt token - // (committed at position plen-1, KV present); the first generated token - // is produced by the first spec round's verify (argmax_per_pos[0]), - // exactly as eager's first decode does — so we do NOT pre-emit it. The - // drafter's step-0 seed hidden = the target's post-`model.norm` hidden - // left in `state.tmp` after the last `decode_step` (the lm_head input). - // - // NOTE: batched prefill via `forward_batch` is NOT used here. gemma4's - // batched MQ4 path faults at large B (b≥16) — a pre-existing - // `forward_batch` limitation surfaced by `verify_batch_gemma4 --bs 16+`, - // orthogonal to this spec wiring (the verify block stays at draft_len+1 - // ≤ 7, well within the proven b≤8 range). Per-token prefill sidesteps it. - let plen = prompt_ids.len(); - let t_pf = std::time::Instant::now(); - for (pos, &t) in prompt_ids.iter().enumerate() { - decode_step(&cfg, &weights, &mut state, &mut gpu, t, pos as u32).expect("prefill"); - } - // Seed hidden = post-`model.norm` hidden of the last prompt position. - sp.set_seed_hidden_from(&gpu, &state.tmp).expect("seed hidden"); - eprintln!("prefill {} tok in {:.2}s", plen, t_pf.elapsed().as_secs_f64()); - - // ── Warm the batched verify path on a disposable state before the spec - // loop. gemma4's b>1 batched MQ4 FFN GEMM lazily initializes scratch on - // first use and faults (illegal access) if its very first call is the one - // inside the loop (a pre-existing `forward_batch` fragility, unrelated to - // this spec wiring — see report). Priming with a throwaway `forward_batch` - // at b=1 THEN at the real block size (draft_len+1) initializes both the - // scalar and batched-WMMA scratch. (Q8 targets: cheap no-op.) Plain - // `forward_batch` (scalar last-row lm_head) is used — NOT - // `forward_batch_spec` — because the spec verify's per-position lm_head is - // the scalar per-row `weight_gemv` (already warm from prefill), so only - // the batched FFN GEMM needs priming. ── - { - // Prime the batched verify path (b=1 then the real block size) on a - // disposable state. gemma4's b>1 batched MQ4 FFN GEMM lazily initializes - // scratch on first use and can fault if its first call is the one inside - // the loop. (Q8 targets: cheap no-op.) This is a pre-existing - // `forward_batch` fragility, unrelated to this spec wiring — see report. - // For MQ4-attn targets, set HIPFIRE_FP16=0 to route the batched MQ4 GEMM - // through the deterministic scalar kernel (the gfx12 fp16-WMMA batched - // path is non-deterministically unstable at this output width). - let warm_b = draft_len + 1; - let mut warm_state = - Gemma4State::new_with_max_seq(&mut gpu, &cfg, warm_b + 4).expect("warm state"); - let _ = forward_batch(&cfg, &weights, &mut warm_state, &mut gpu, &[cfg.bos_token], 0); - gpu.hip.device_synchronize().ok(); - let mut warm2 = - Gemma4State::new_with_max_seq(&mut gpu, &cfg, warm_b + 4).expect("warm state 2"); - let _ = forward_batch( - &cfg, - &weights, - &mut warm2, - &mut gpu, - &vec![cfg.bos_token; warm_b], - 0, - ); - gpu.hip.device_synchronize().ok(); - } - - // Seed = last prompt token (at position plen-1). The first spec round's - // verify re-writes its KV at L-1 and predicts the first generated token. - let mut seed_token = *prompt_ids.last().unwrap(); - let mut gen: Vec = Vec::new(); - let mut rounds = 0usize; - let mut total_accepted = 0usize; // accept_len summed (excludes per-round bonus) - let mut stop = false; - - let t0 = std::time::Instant::now(); - while !stop && gen.len() < max { - let committed_len = state.n_tokens; // L = positions filled = seed at L-1 - // Guard the KV/seq bound: block occupies [L-1, L-1 + draft_len + 1). - if committed_len + draft_len + 1 >= max_seq { - break; - } - let spec = spec_step_gemma4_eagle( - &mut gpu, - &weights, - &cfg, - &mut state, - &dweights, - &dcfg, - &mut dscratch, - &mut sp, - seed_token, - committed_len, - draft_len, - temperature, - ) - .expect("spec step"); - rounds += 1; - total_accepted += spec.accept_len; - - // Emit committed tokens (accepted drafts ++ bonus). Stop at EOS / max. - for &t in &spec.committed { - if is_eos(t) { - stop = true; - break; - } - gen.push(t); - if gen.len() >= max { - stop = true; - break; - } - } - // Next round's seed = this round's bonus (= last committed). Its hidden - // is already staged in sp.seed_hidden by spec_step. - seed_token = spec.next_seed_token; - } - let dt = t0.elapsed().as_secs_f64(); - - // τ = mean tokens committed per round (accepted drafts + 1 bonus each). - let committed_total = total_accepted + rounds; // each round commits accept_len + 1 - let tau = if rounds > 0 { - committed_total as f64 / rounds as f64 - } else { - 0.0 - }; - let toks_per_s = if dt > 0.0 { gen.len() as f64 / dt } else { 0.0 }; - - eprintln!( - "spec-decode: {} tok in {:.2}s ({:.1} tok/s) | rounds={} τ={:.3} draft_len={}", - gen.len(), - dt, - toks_per_s, - rounds, - tau, - draft_len - ); - - match &tok { - Some(t) => println!( - "=== PROMPT ===\n{prompt}\n=== GENERATION (spec, draft_len={draft_len}) ===\n{}", - t.decode(&gen) - ), - None => println!("=== GENERATION token ids ===\n{:?}", gen), - } - - // ── Correctness gate ── - if let Some(eager) = eager_gen { - let identical = eager == gen; - println!("\n=== CORRECTNESS GATE (spec vs eager greedy AR) ==="); - println!("eager ({} tok): {:?}", eager.len(), &eager[..eager.len().min(64)]); - println!("spec ({} tok): {:?}", gen.len(), &gen[..gen.len().min(64)]); - if identical { - println!("=> IDENTICAL ✓"); - } else { - // Find first divergence. - let mut div = None; - for k in 0..eager.len().min(gen.len()) { - if eager[k] != gen[k] { - div = Some(k); - break; - } - } - match div { - Some(k) => println!( - "=> DIVERGE at index {k}: eager={} spec={} (eager_len={} spec_len={})", - eager[k], - gen[k], - eager.len(), - gen.len() - ), - None => println!( - "=> PREFIX-MATCH but length differs: eager_len={} spec_len={}", - eager.len(), - gen.len() - ), - } - std::process::exit(1); - } - } - - eprintln!("token ids: {:?}", &gen[..gen.len().min(60)]); -} diff --git a/crates/hipfire-arch-gemma4/examples/prefill_parity_gemma4.rs b/crates/hipfire-arch-gemma4/examples/prefill_parity_gemma4.rs index b9040366dd..8a8395264e 100644 --- a/crates/hipfire-arch-gemma4/examples/prefill_parity_gemma4.rs +++ b/crates/hipfire-arch-gemma4/examples/prefill_parity_gemma4.rs @@ -2,14 +2,13 @@ // Copyright (c) 2026 Kaden Schutt // hipfire — see LICENSE and NOTICE in the project root. -//! Batched-vs-per-token prefill parity harness for the gemma4 lowered path. +//! Batched-vs-per-token prefill parity harness for Gemma 4. //! -//! Runs the SAME prompt through (A) per-token `forward_scratch` and +//! Runs the SAME prompt through (A) per-token `decode_step` and //! (B) `forward_prefill_batch` over all but the last token plus one -//! `forward_scratch`, with fresh KV each; compares the prompt's last logits +//! `decode_step`, with fresh state each; compares the prompt's last logits //! (max-abs diff, argmax) and an N-token greedy continuation (per-token decode -//! from both prefill states). Uses unwindowed Q8 KV for both attention tiers, -//! the layout batched prefill writes. +//! from both prefill states). //! //! Usage: //! prefill_parity_gemma4 --model [--prompt ] [--decode N] @@ -21,9 +20,8 @@ fn main() { #[cfg(feature = "deltanet")] fn main() { - use hipfire_arch_gemma4::lowered; + use hipfire_arch_gemma4::{forward, Gemma4Config, Gemma4State, Gemma4Weights}; use hipfire_runtime::hfq::HfqFile; - use hipfire_runtime::llama::KvCache; use hipfire_runtime::tokenizer::Tokenizer; use std::path::PathBuf; @@ -61,17 +59,15 @@ fn main() { let mut gpu = rdna_compute::Gpu::init().expect("gpu init"); eprintln!("arch = {}", gpu.arch); - let mut hfq = HfqFile::open(&model).expect("open model"); - let cfg = lowered::config_from_hfq(&hfq).expect("lowered config"); + let hfq = HfqFile::open(&model).expect("open model"); + let cfg = Gemma4Config::from_hfq(&hfq).expect("config"); let tok = Tokenizer::from_hfq_metadata(&hfq.metadata_json).expect("tokenizer"); - let weights = lowered::load_weights(&mut hfq, &cfg, &mut gpu).expect("weights"); + let weights = Gemma4Weights::load(&hfq, &cfg, &mut gpu).expect("weights"); let mut ids = tok.encode(&prompt); if ids.first() != Some(&cfg.bos_token) { ids.insert(0, cfg.bos_token); } let max_seq = (ids.len() + decode_n + 16).max(cfg.sliding_window + 1); - let scratch = lowered::Gemma4Scratch::new(&mut gpu, &cfg, max_seq).expect("scratch"); - lowered::init_scratch_constants(&mut gpu, &scratch, cfg.full_head_dim).expect("scratch consts"); eprintln!("prompt tokens = {}", ids.len()); @@ -96,69 +92,26 @@ fn main() { }; let mut run = |label: &str, batched: bool| -> (Vec, Vec) { - let mut kv_sliding = KvCache::new_gpu_q8( - &mut gpu, - cfg.n_layers, - cfg.sliding_n_kv_heads, - cfg.sliding_head_dim, - max_seq, - ) - .expect("kv sliding"); - let mut kv_full = KvCache::new_gpu_q8( - &mut gpu, - cfg.n_layers, - cfg.full_n_kv_heads, - cfg.full_head_dim, - max_seq, - ) - .expect("kv full"); + let mut state = Gemma4State::new_with_max_seq(&mut gpu, &cfg, max_seq).expect("state"); let t0 = std::time::Instant::now(); if batched { let last = ids.len() - 1; - lowered::forward_prefill_batch( - &mut gpu, - &weights, - &cfg, - &ids[..last], - 0, - &mut kv_sliding, - &mut kv_full, - &scratch, - ) - .expect("batched prefill"); - lowered::forward_scratch( - &mut gpu, - &weights, - &cfg, - ids[last], - last, - &mut kv_sliding, - &mut kv_full, - &scratch, - ) - .expect("last prompt token"); + forward::forward_prefill_batch(&cfg, &weights, &mut state, &mut gpu, &ids[..last], 0) + .expect("batched prefill"); + forward::decode_step(&cfg, &weights, &mut state, &mut gpu, ids[last], last as u32) + .expect("last prompt token"); } else { for (p, &t) in ids.iter().enumerate() { - lowered::forward_scratch( - &mut gpu, - &weights, - &cfg, - t, - p, - &mut kv_sliding, - &mut kv_full, - &scratch, - ) - .expect("per-token prefill"); + forward::decode_step(&cfg, &weights, &mut state, &mut gpu, t, p as u32) + .expect("per-token prefill"); } } - let _ = gpu.download_f32(&scratch.logits); eprintln!( "[{label}] prefill {} tok in {:.3}s", ids.len(), t0.elapsed().as_secs_f64() ); - let logits = gpu.download_f32(&scratch.logits).expect("logits dl"); + let logits = gpu.download_f32(&state.logits).expect("logits dl"); assert!( logits.iter().all(|v| v.is_finite()), "non-finite {label} logits" @@ -177,25 +130,14 @@ fn main() { let mut next = am as u32; for _ in 0..decode_n { cont.push(next); - lowered::forward_scratch( - &mut gpu, - &weights, - &cfg, - next, - pos, - &mut kv_sliding, - &mut kv_full, - &scratch, - ) - .expect("decode"); - let l = gpu.download_f32(&scratch.logits).expect("dl"); + let l = forward::decode_step(&cfg, &weights, &mut state, &mut gpu, next, pos as u32) + .expect("decode"); next = argmax(&l).0 as u32; pos += 1; } eprintln!("[{label}] cont ids: {:?}", cont); eprintln!("[{label}] cont text: {:?}", tok.decode(&cont)); - kv_sliding.free_gpu(&mut gpu).expect("free sliding KV"); - kv_full.free_gpu(&mut gpu).expect("free full KV"); + state.free_gpu(&mut gpu); (logits, cont) }; diff --git a/crates/hipfire-arch-gemma4/examples/verify_drafter_gemma4.rs b/crates/hipfire-arch-gemma4/examples/verify_drafter_gemma4.rs index 9c142ed59a..0514f86bf7 100644 --- a/crates/hipfire-arch-gemma4/examples/verify_drafter_gemma4.rs +++ b/crates/hipfire-arch-gemma4/examples/verify_drafter_gemma4.rs @@ -21,7 +21,7 @@ //! validated arch-13 target only reaches ~0.88 min-cos by layer 5 at this //! tiny random-weight scale). Without `--inject-hf` the drafter consumes //! hipfire's own (drifted) target hidden+KV — a looser, end-to-end check. -//! 4. Run `drafter_step` for N steps, mirroring the HF candidate generator: +//! 4. Run `draft_step` for N steps, mirroring the HF candidate generator: //! step 0 → prev = last prompt token, hidden = target last-hidden; //! step i>0 → prev = prev step argmax, hidden = prev step post_proj. //! The drafter reads the target's LAST sliding + LAST full KV slot directly @@ -44,8 +44,8 @@ fn main() { fn main() { use hipfire_arch_gemma4::config::Gemma4Config; use hipfire_arch_gemma4::drafter::{ - drafter_step, drafter_step_from_concat, Gemma4DrafterConfig, Gemma4DrafterScratch, - Gemma4DrafterWeights, + draft_step, drafter_step_from_concat, DrafterStepOut, Gemma4DrafterConfig, + Gemma4DrafterScratch, Gemma4DrafterWeights, }; use hipfire_arch_gemma4::forward::decode_step; use hipfire_arch_gemma4::gemma4::{Gemma4State, Gemma4Weights}; @@ -219,10 +219,14 @@ fn main() { dcfg.backbone_hidden, dcfg.vocab_size, dcfg.final_logit_softcapping, dcfg.num_centroids, dcfg.centroid_top_k, ); - assert_eq!(dcfg.backbone_hidden, tcfg.dim, "drafter backbone must equal target hidden"); + assert_eq!( + dcfg.backbone_hidden, tcfg.dim, + "drafter backbone must equal target hidden" + ); let drafter_weights = Gemma4DrafterWeights::load(&drafter_hfq, &dcfg, &mut gpu).expect("drafter weights"); - let mut dscratch = Gemma4DrafterScratch::new(&mut gpu, &dcfg).expect("drafter scratch"); + let mut dscratch = + Gemma4DrafterScratch::new(&mut gpu, &dcfg, &drafter_weights).expect("drafter scratch"); // ── Run the target to populate its KV caches ── let max_seq = n_ctx + 8; @@ -246,9 +250,7 @@ fn main() { }; let tcos = cosine(&target_normed, &target_last_hidden_hf); - eprintln!( - "[sanity] hipfire target normed-final-hidden vs HF cosine = {tcos:.6}" - ); + eprintln!("[sanity] hipfire target normed-final-hidden vs HF cosine = {tcos:.6}"); // ── Optional: inject HF-exact KV into the target's last slots ── if inject_hf { @@ -269,10 +271,22 @@ fn main() { std::slice::from_raw_parts(ph.as_ptr() as *const u8, 4) }) .expect("pos htod"); - gpu.kv_cache_write_q8_0(&tstate.kv_sliding.k_gpu[s_slot], &kt, &pos_buf, n_kv_s, hd_s) - .expect("inject sliding k"); - gpu.kv_cache_write_q8_0(&tstate.kv_sliding.v_gpu[s_slot], &vt, &pos_buf, n_kv_s, hd_s) - .expect("inject sliding v"); + gpu.kv_cache_write_q8_0( + &tstate.kv_sliding.k_gpu[s_slot], + &kt, + &pos_buf, + n_kv_s, + hd_s, + ) + .expect("inject sliding k"); + gpu.kv_cache_write_q8_0( + &tstate.kv_sliding.v_gpu[s_slot], + &vt, + &pos_buf, + n_kv_s, + hd_s, + ) + .expect("inject sliding v"); } // full K/V { @@ -301,7 +315,11 @@ fn main() { } else { target_normed.clone() }; - let mut hidden_backbone = gpu.upload_f32(&step0_hidden, &[bb]).expect("upload step0 hidden"); + let mut hidden_backbone = gpu + .upload_f32(&step0_hidden, &[bb]) + .expect("upload step0 hidden"); + // The step-0 hidden half stays on device across steps unless re-injected. + let mut stage_hidden = true; // ── Drafter loop ── let last_prompt_token = *tokens.last().unwrap(); @@ -325,6 +343,7 @@ fn main() { hidden_backbone = gpu .upload_f32(&o.pre_in_hidden, &[bb]) .expect("upload hf hidden"); + stage_hidden = true; } else if !inject_pre_in && prev_token != o.prev_token { eprintln!( "[warn] step {step}: hipfire prev_token {prev_token} != oracle {} — feeding hipfire's", @@ -337,6 +356,7 @@ fn main() { &drafter_weights, &dcfg, &mut dscratch, + &target_weights, &tstate, &tcfg, &o.pre_in, @@ -344,7 +364,13 @@ fn main() { ) .expect("drafter_step_from_concat") } else { - drafter_step( + if stage_hidden { + dscratch + .begin_round(&mut gpu, &hidden_backbone, query_pos) + .expect("begin_round"); + stage_hidden = false; + } + let argmax = draft_step( &mut gpu, &drafter_weights, &dcfg, @@ -353,10 +379,19 @@ fn main() { &tstate, &tcfg, prev_token, - &hidden_backbone, query_pos, ) - .expect("drafter_step") + .expect("draft_step"); + DrafterStepOut { + argmax, + post_proj_hidden: gpu + .download_f32(&dscratch.concat.sub_offset(bb, bb)) + .expect("download post_proj"), + normed_hidden: gpu.download_f32(&dscratch.normed).expect("download normed"), + logits: gpu + .download_f32(dscratch.head.logits()) + .expect("download logits"), + } }; let _ = &o.pre_out; @@ -379,19 +414,12 @@ fn main() { ); prev_token = out.argmax; - hidden_backbone = gpu - .upload_f32(&out.post_proj_hidden, &[bb]) - .expect("upload next hidden"); } // Gate: per-step isolation requires EVERY step ≥0.99; the realistic // feedback chain only requires step 0 (drift downstream is corrected by the // target's verify — the drafter affects acceptance τ, not correctness). - let overall = if feedback { - step0_pass - } else { - all_pass - }; + let overall = if feedback { step0_pass } else { all_pass }; println!( "\n=== OVERALL ({mode}): {} ===", if overall { "PASS" } else { "FAIL" } diff --git a/crates/hipfire-arch-gemma4/map.md b/crates/hipfire-arch-gemma4/map.md index e93e7bd6fc..c3d0376ed4 100644 --- a/crates/hipfire-arch-gemma4/map.md +++ b/crates/hipfire-arch-gemma4/map.md @@ -1,11 +1,11 @@ # hipfire-arch-gemma4 — map > **Status:** `production` (vocabulary owned by [`docs/GLOSSARY.md`](../../docs/GLOSSARY.md)). -> **Layer:** Arch forward — hangs off `hipfire-loader` via `Carrier` (`arch_id` for the dense 12B port); see the layering table in [`docs/ARCHITECTURE.md`](../../docs/ARCHITECTURE.md). +> **Layer:** Arch forward — hangs off `hipfire-loader` via `Carrier` (`arch_id` 13; the arch-22 EAGLE drafter rides it); see the layering table in [`docs/ARCHITECTURE.md`](../../docs/ARCHITECTURE.md). ## Purpose -Gemma 4 12B dense text (and strictly admitted E2B/E4B text variants): hybrid attention (5 sliding window-1024 layers per 1 full layer with head_dim 256/512, K=V sharing on full), sandwich RMSNorm with per-layer `layer_scalar`, attention scale 1.0 (Q pre-scaled so kernel's `1/√d` cancels), `logit_softcap` = 30, `softcap_scale` 60. Verified against HF Gemma4 F32. The `//!` docs in [`src/lib.rs`](src/lib.rs) enumerate the distinguishing features vs plain pre-norm. +Gemma 4 text on one weight stack, state and layer program: the 12B dense, strictly admitted E2B/E4B, and the 26B-A4B MoE. Hybrid attention (5 sliding window-1024 layers per 1 full layer with head_dim 256/512, K=V sharing on full), sandwich RMSNorm with per-layer `layer_scalar`, attention scale 1.0 (Q pre-scaled so kernel's `1/√d` cancels), `logit_softcap` = 30, `softcap_scale` 60. Verified against HF Gemma4 F32. The `//!` docs in [`src/lib.rs`](src/lib.rs) enumerate the distinguishing features vs plain pre-norm. `bundle.rs` owns the engine session cache (`SessionState`); `mtp.rs` puts the EAGLE drafter on the generic `MtpDrafter`/`SpecTarget` seams; `emit.rs` routes thought-channel output for AR and spec. ## Gotchas @@ -22,16 +22,17 @@ _Generated by `scripts/check-crate-maps.py` from the tree — do not edit inside ### Modules and exported symbols - [`src/arch.rs`](src/arch.rs): `ARCH_ID`, `Gemma4` -- [`src/carrier.rs`](src/carrier.rs): `Gemma4Bundle`, `Gemma4EagerBundle`, `Gemma4LoweredBundle`, `gemma4_context_admission`, `gemma4_source_uses_lowered`, `load_gemma4_bundle` +- [`src/bundle.rs`](src/bundle.rs): `Gemma4Bundle`, `SESSION_STRIDE`, `attach_session_cache`, `new`, `prefill`, `prefill_rows_for_arch`, `session_cache`, `session_snapshot_bytes` +- [`src/carrier.rs`](src/carrier.rs): `load_gemma4_bundle` - [`src/config.rs`](src/config.rs): `Gemma4Config`, `Gemma4ESeriesVariant`, `LayerType`, `RopeType`, `e_series_variant`, `ffn_hidden_dim_for_layer`, `first_kv_shared_layer_idx`, `from_hfq`, `from_metadata_json`, `is_kv_shared_layer`, `kv_shared_source_layer_idx`, `max_ffn_hidden_dim`, `max_head_dim`, `max_kv_dim`, `max_q_dim`, `n_full_kv_slots`, `n_full_layers`, `n_sliding_kv_slots`, `n_sliding_layers` -- [`src/drafter.rs`](src/drafter.rs): `DRAFTER_ARCH_ID`, `DrafterLayerWeights`, `DrafterStepOut`, `Gemma4DrafterConfig`, `Gemma4DrafterScratch`, `Gemma4DrafterWeights`, `drafter_step`, `drafter_step_from_concat`, `free_gpu`, `from_hfq`, `load`, `max_head_dim`, `max_q_dim`, `new`, `pre_proj_in` -- [`src/forward.rs`](src/forward.rs): `decode_step`, `decode_step_capture`, `decode_step_with_graph`, `forward_batch`, `forward_batch_spec`, `supports_batched_prefill` -- [`src/gemma4.rs`](src/gemma4.rs): `FullLayerWeights`, `Gemma4State`, `Gemma4Weights`, `LayerWeights`, `PerLayerBranchWeights`, `PerLayerInputWeights`, `SlidingLayerWeights`, `free_gpu`, `load`, `new`, `new_with_fwht3_max_seq`, `new_with_max_seq`, `reset` +- [`src/drafter.rs`](src/drafter.rs): `DRAFTER_ARCH_ID`, `DrafterLayerWeights`, `DrafterStepOut`, `Gemma4DrafterConfig`, `Gemma4DrafterScratch`, `Gemma4DrafterWeights`, `begin_round`, `draft_step`, `drafter_step_from_concat`, `free_gpu`, `from_hfq`, `load`, `max_head_dim`, `max_q_dim`, `new`, `pre_proj_in`, `target_embedding_table` +- [`src/emit.rs`](src/emit.rs): `GemmaChannel`, `GemmaEmit`, `GemmaSpecEmit`, `GemmaThoughtRouter`, `flush`, `gemma_is_marker_prefix`, `gemma_longest_marker_suffix`, `new`, `push` +- [`src/forward.rs`](src/forward.rs): `decode_body`, `decode_step`, `decode_step_capture`, `decode_step_with_graph`, `forward_batch`, `forward_batch_spec`, `forward_prefill_batch`, `prepare_token`, `supports_batched_prefill` +- [`src/gemma4.rs`](src/gemma4.rs): `FullKvTier`, `FullLayerWeights`, `Gemma4State`, `Gemma4Weights`, `LayerWeights`, `MoeLayerExtras`, `MoeScratch`, `PerLayerBranchWeights`, `PerLayerInputWeights`, `SlidingLayerWeights`, `alloc`, `free_gpu`, `kv`, `load`, `moe`, `new`, `new_with_full_tier`, `new_with_fwht3_max_seq`, `new_with_max_seq`, `reset`, `rollback`, `scratch_tensors`, `zeros` - [`src/gemma4_vision.rs`](src/gemma4_vision.rs): `Gemma4VisionConfig`, `Gemma4VisionScratch`, `Gemma4VisionWeights` -- [`src/lib.rs`](src/lib.rs): `arch`, `carrier`, `config`, `drafter`, `forward`, `gemma4`, `lowered`, `speculative` -- [`src/lowered.rs`](src/lowered.rs): `FullLayerWeights`, `GEMMA4_FLASH_TILE`, `Gemma4Config`, `Gemma4Scratch`, `Gemma4Weights`, `LayerType`, `LayerWeights`, `MoeExpertWeights`, `MoeLayerExtras`, `RopeType`, `SlidingLayerWeights`, `config_from_hfq`, `forward_prefill_batch`, `forward_scratch`, `free_gpu`, `gemma4_decode_flash_partials_len`, `gemma4_flash_partials_len`, `gemma4_flash_partials_len_for_config`, `init_scratch_constants`, `load_weights`, `new` -- [`src/program.rs`](src/program.rs): `alloc`, `drafter`, `eager`, `layer`, `lowered`, `new`, `scratch` -- [`src/speculative.rs`](src/speculative.rs): `Gemma4SpecScratch`, `SpecStepOut`, `free`, `new`, `set_seed_hidden_from`, `spec_step_gemma4_eagle` +- [`src/lib.rs`](src/lib.rs): `arch`, `bundle`, `carrier`, `config`, `drafter`, `emit`, `forward`, `gemma4`, `mtp` +- [`src/mtp.rs`](src/mtp.rs): `Gemma4Drafter`, `config`, `load` +- [`src/program.rs`](src/program.rs): `alloc`, `drafter`, `eager`, `layer`, `new`, `scratch` ### Dependencies (from `Cargo.toml`) diff --git a/crates/hipfire-arch-gemma4/src/bundle.rs b/crates/hipfire-arch-gemma4/src/bundle.rs new file mode 100644 index 0000000000..a76c1d094e --- /dev/null +++ b/crates/hipfire-arch-gemma4/src/bundle.rs @@ -0,0 +1,393 @@ +// SPDX-License-Identifier: MIT OR Apache-2.0 +// Copyright (c) 2026 Björn Bösel +// hipfire — see LICENSE and NOTICE in the project root. + +//! Published Gemma 4 owner: config, weights, state and the engine session +//! cache. +//! +//! The KV caches are the whole request state (Gemma 4 is pure attention, and +//! the drafter keeps nothing of its own), so a session snapshot is the +//! append-only KV rows of every cache slot up to its position. Rows are +//! position-indexed and never rewritten below the position, so a snapshot +//! stores only the rows above its parent. + +use crate::config::Gemma4Config; +use crate::forward; +use crate::gemma4::{Gemma4State, Gemma4Weights}; +use hipfire_runtime::llama::KvCache; +use hipfire_runtime::session_cache::{ + stride_boundaries, RowStream, SessionCache, SessionDriver, SessionRoute, SessionState, + SnapshotParts, StateLayout, +}; +use rdna_compute::Gpu; + +/// Snapshot boundary stride in prompt tokens (a multiple of the 64-row +/// prefill chunk, so restored prefills keep the cold chunk grid). +pub const SESSION_STRIDE: usize = 1024; + +/// Prefill rows per forward on `gpu_arch`: 64-row batches where batched +/// prefill is validated (`HIPFIRE_GEMMA4_PREFILL_BATCH` overrides), else +/// per-token decode. +pub fn prefill_rows_for_arch(gpu_arch: &str, requested: Option) -> usize { + requested + .unwrap_or(if matches!(gpu_arch, "gfx1100" | "gfx1201") { + 64 + } else { + 1 + }) + .clamp(1, 64) +} + +/// Published Gemma 4 owner (arch_id 13). +pub struct Gemma4Bundle { + pub config: Gemma4Config, + pub weights: Gemma4Weights, + pub state: Gemma4State, + /// End-of-turn token resolved from the tokenizer at load. + pub eos_tok: u32, + /// Rows per prefill forward (1 = per-token decode). + prefill_rows: usize, + /// Prefill snapshot cache shared by every session plus this + /// conversation's committed live state; `None` = off. + session: Option, +} + +/// Bytes of one position of one Q8_0 cache slot (K or V): 34-byte blocks of +/// 32 elements per head. +fn q8_row_bytes(kv: &KvCache) -> usize { + kv.n_kv_heads * (kv.head_dim / 32) * 34 +} + +impl Gemma4Bundle { + pub fn new( + config: Gemma4Config, + weights: Gemma4Weights, + state: Gemma4State, + gpu: &Gpu, + ) -> Self { + let requested = hipfire_config::developer_var("HIPFIRE_GEMMA4_PREFILL_BATCH") + .ok() + .and_then(|value| value.parse::().ok()); + let mut prefill_rows = prefill_rows_for_arch(&gpu.arch, requested); + if prefill_rows > 1 && !forward::supports_batched_prefill(&weights, &state) { + eprintln!( + " gemma4: batched prefill is unavailable for this load; using per-token prefill" + ); + prefill_rows = 1; + } + Self { + eos_tok: config.eos_token, + config, + weights, + state, + prefill_rows, + session: None, + } + } + + /// Attach the session cache: prefill then restores and captures + /// [`SESSION_STRIDE`]-token snapshots through it. + pub fn attach_session_cache(&mut self, cache: SessionCache) { + self.session = Some(SessionDriver::new(cache)); + } + + /// The attached session cache, if any. + pub fn session_cache(&self) -> Option<&SessionCache> { + self.session.as_ref().map(SessionDriver::cache) + } + + /// Bytes of one [`SESSION_STRIDE`]-token delta snapshot. + pub fn session_snapshot_bytes(&self) -> usize { + let row = |kv: &KvCache| 2 * kv.k_gpu.len() * q8_row_bytes(kv); + SESSION_STRIDE * (row(&self.state.kv_sliding) + row(&self.state.kv_full)) + } + + /// Run `f` with the driver taken out, so it can drive `self` as its + /// [`SessionState`]. `None` without a cache. + fn with_session(&mut self, f: impl FnOnce(&mut SessionDriver, &mut Self) -> R) -> Option { + let mut session = self.session.take()?; + let result = f(&mut session, self); + self.session = Some(session); + Some(result) + } + + /// Prefill the canonical `prompt` on `route`, skipping the `reused` + /// tokens `ArchModel::session_plan` returned: continue the live state, + /// restore a snapshot, or start cold, then run the rest in prefill + /// chunks that stop at every snapshot boundary. Returns the last prompt + /// token's logits. + pub fn prefill( + &mut self, + gpu: &mut Gpu, + prompt: &[u32], + reused: usize, + route: SessionRoute, + ) -> Result, String> { + if reused >= prompt.len() { + return Err(format!( + "gemma4 prefill: reusing {reused} of {} prompt tokens leaves nothing to compute", + prompt.len() + )); + } + match self.with_session(|session, bundle| { + session.begin(gpu, bundle, prompt, route, reused, false) + }) { + Some(result) => { + result?; + } + None if reused == 0 => self.state.reset(), + None => { + return Err(format!( + "gemma4 session cache is not attached; cannot reuse {reused} tokens" + )) + } + } + let mut pos = reused; + loop { + let boundary = self.session.as_ref().and_then(SessionDriver::next_boundary); + let stop = boundary.unwrap_or(prompt.len()).min(prompt.len()); + if stop == prompt.len() { + let logits = self.forward_rows(gpu, &prompt[pos..], pos)?; + if boundary == Some(stop) { + self.with_session(|session, bundle| session.at_boundary(gpu, bundle, prompt)) + .unwrap_or(Ok(()))?; + } + return Ok(logits); + } + self.forward_rows_silent(gpu, &prompt[pos..stop], pos)?; + self.with_session(|session, bundle| session.at_boundary(gpu, bundle, &prompt[..stop])) + .unwrap_or(Ok(()))?; + pos = stop; + } + } + + /// `tokens` at `[start, start + len)` in prefill chunks; returns the last + /// token's logits. + pub(crate) fn forward_rows( + &mut self, + gpu: &mut Gpu, + tokens: &[u32], + start: usize, + ) -> Result, String> { + let width = self.prefill_rows; + let mut logits = Vec::new(); + for (i, chunk) in tokens.chunks(width).enumerate() { + let pos = start + i * width; + logits = if let [token] = chunk { + forward::decode_step( + &self.config, + &self.weights, + &mut self.state, + gpu, + *token, + pos as u32, + )? + } else { + forward::forward_batch( + &self.config, + &self.weights, + &mut self.state, + gpu, + chunk, + pos, + )? + }; + } + Ok(logits) + } + + /// As [`Self::forward_rows`] but without the logits of batched chunks. + fn forward_rows_silent( + &mut self, + gpu: &mut Gpu, + tokens: &[u32], + start: usize, + ) -> Result<(), String> { + if self.prefill_rows == 1 { + return self.forward_rows(gpu, tokens, start).map(drop); + } + forward::forward_prefill_batch( + &self.config, + &self.weights, + &mut self.state, + gpu, + tokens, + start, + ) + } + + fn layout(&self, rows: usize) -> StateLayout<'_> { + let mut streams = Vec::new(); + for kv in [&self.state.kv_sliding, &self.state.kv_full] { + let row_bytes = q8_row_bytes(kv); + for tensor in kv.k_gpu.iter().chain(&kv.v_gpu) { + streams.push(RowStream { + buf: &tensor.buf, + row_bytes, + rows, + }); + } + } + StateLayout { + fixed: Vec::new(), + rows: streams, + } + } +} + +fn meta_position(meta: &[u8]) -> Result { + let bytes: [u8; 8] = meta + .try_into() + .map_err(|_| "gemma4 session snapshot shape mismatch".to_string())?; + Ok(u64::from_le_bytes(bytes) as usize) +} + +impl SessionState for Gemma4Bundle { + /// One scope for both routes: the drafter keeps no state and the target + /// prefills the same way on either. The prefill width is part of it + /// because batched and per-token prefill write different KV bits. + fn snapshot_scope(&self, _route: SessionRoute) -> Option { + (self.state.kv_sliding.quant_q8 && self.state.kv_full.quant_q8) + .then(|| format!("gemma4/rows{}", self.prefill_rows)) + } + + fn snapshot_boundaries(&self, prompt: &[u32], after: usize, up_to: usize) -> Vec { + stride_boundaries(SESSION_STRIDE, prompt, None, after, up_to) + } + + fn snapshot_parts( + &mut self, + _route: SessionRoute, + position: usize, + ) -> Result, String> { + if self.state.n_tokens != position { + return Err(format!( + "gemma4 session capture: state at {}, not {position}", + self.state.n_tokens + )); + } + Ok(SnapshotParts { + meta: (position as u64).to_le_bytes().to_vec(), + layout: self.layout(position), + }) + } + + fn restore_parts( + &mut self, + _gpu: &mut Gpu, + _route: SessionRoute, + meta: &[u8], + ) -> Result, String> { + let position = meta_position(meta)?; + if position >= self.state.max_seq { + return Err(format!( + "gemma4 session restore: position {position} exceeds max_seq {}", + self.state.max_seq + )); + } + Ok(self.layout(position)) + } + + fn finish_restore( + &mut self, + _gpu: &mut Gpu, + _route: SessionRoute, + meta: &[u8], + ) -> Result<(), String> { + self.state.n_tokens = meta_position(meta)?; + Ok(()) + } + + fn live_position(&self, _route: SessionRoute) -> Option { + Some(self.state.n_tokens) + } + + fn growth_reserve_bytes(&self, _gpu: &Gpu, _through: usize) -> u64 { + // Both caches are allocated for the admitted context at load. + 0 + } + + fn reset(&mut self, _gpu: &mut Gpu) -> Result<(), String> { + self.state.reset(); + Ok(()) + } +} + +impl hipfire_runtime::arch_model::ArchModel for Gemma4Bundle { + fn dim(&self) -> usize { + self.config.dim + } + fn n_layers(&self) -> usize { + self.config.n_layers + } + fn vocab_size(&self) -> usize { + self.config.vocab_size + } + fn arch_key(&self) -> &'static str { + "gemma4" + } + fn kv_cache_mut(&mut self) -> Option<&mut KvCache> { + // Two caches (sliding + full) and no basis for preferring one. + None + } + fn reset_session_state(&mut self, _gpu: &mut Gpu) -> Result<(), String> { + self.state.reset(); + if let Some(session) = self.session.as_mut() { + session.forget_live(); + } + Ok(()) + } + fn session_cache_attached(&self) -> bool { + self.session.is_some() + } + fn session_plan(&self, prompt: &[u32], route: SessionRoute) -> usize { + self.session + .as_ref() + .map_or(0, |session| session.plan(self, prompt, route)) + } + fn session_commit(&mut self) { + if let Some(session) = self.session.as_mut() { + session.commit(); + } + } + fn session_commit_live(&mut self, consumed: &[u32]) { + if let Some(mut session) = self.session.take() { + session.commit_live(self, consumed); + self.session = Some(session); + } + } + fn free_gpu(self: Box, gpu: &mut Gpu) { + let b = *self; + if let Some(mut session) = b.session { + session.clear(gpu); + } + b.state.free_gpu(gpu); + b.weights.free_gpu(gpu); + } +} + +#[cfg(test)] +mod tests { + use super::prefill_rows_for_arch; + + #[test] + fn validated_arches_default_to_full_wmma_batch_tile_group() { + assert_eq!(prefill_rows_for_arch("gfx1100", None), 64); + assert_eq!(prefill_rows_for_arch("gfx1201", None), 64); + } + + #[test] + fn unvalidated_arches_keep_the_sequential_default() { + for arch in ["gfx1151", "gfx1200", "gfx942"] { + assert_eq!(prefill_rows_for_arch(arch, None), 1); + } + } + + #[test] + fn explicit_override_wins_and_is_clamped() { + assert_eq!(prefill_rows_for_arch("gfx1100", Some(8)), 8); + assert_eq!(prefill_rows_for_arch("gfx1201", Some(32)), 32); + assert_eq!(prefill_rows_for_arch("gfx1100", Some(0)), 1); + assert_eq!(prefill_rows_for_arch("gfx1100", Some(128)), 64); + } +} diff --git a/crates/hipfire-arch-gemma4/src/carrier.rs b/crates/hipfire-arch-gemma4/src/carrier.rs index 2514a70971..87aed08de2 100644 --- a/crates/hipfire-arch-gemma4/src/carrier.rs +++ b/crates/hipfire-arch-gemma4/src/carrier.rs @@ -4,42 +4,16 @@ //! Gemma 4 carrier bundle loader — HFQ path. //! -//! Verbatim relocation of the model-loading work from -//! `hipfire-loader/src/carriers.rs::Gemma4Carrier::load`. The loader retains -//! `LoadedModel` assembly, `SourceMeta`/`resolve_source_meta`, chat-template -//! and tokenizer handling, `Gemma4EagleState` side-car load, and -//! `spec_build::build_speculator`. This module owns the GPU bundle construction -//! (lowered vs eager decision, weight/state/KV allocation) with error strings -//! byte-identical to the prior inline block. +//! The loader retains `LoadedModel` assembly, `SourceMeta`/`resolve_source_meta`, +//! chat-template and tokenizer handling, the session cache attach and the +//! EAGLE speculator (`crate::mtp::Gemma4Drafter`). This module owns the GPU bundle +//! construction: config admission, weights, state and both KV caches, for +//! every Gemma 4 variant (dense, E-series, MoE). +use crate::bundle::Gemma4Bundle; use crate::config::Gemma4Config; -use crate::gemma4::{Gemma4State, Gemma4Weights}; -use crate::lowered; -use hipfire_runtime::llama::KvCache; +use crate::gemma4::{FullKvTier, Gemma4State, Gemma4Weights}; use hipfire_runtime::loader_api::{LoadCtx, ModelSource}; -use rdna_compute::Gpu; - -// ─── Helpers moved verbatim from carriers.rs ───────────────────────────── - -/// Pure context admission for Gemma 4. Refuses the lowered route when -/// `max_seq < 128`. Eager loads are unrestricted (no eager assert exists). -/// -/// Callers compute `use_lowered` via [`gemma4_source_uses_lowered`] so the -/// refusal string is byte-identical at admission and carrier layers. -pub fn gemma4_context_admission(max_seq: usize, use_lowered: bool) -> Result<(), String> { - if use_lowered && max_seq < 128 { - return Err(format!( - "gemma4 lowered path requires max_seq >= 128 (got {max_seq})" - )); - } - Ok(()) -} - -/// MoE checkpoints run on the lowered stack; every other Gemma 4 load is -/// eager. -pub fn gemma4_source_uses_lowered(hfq: &hipfire_runtime::hfq::HfqFile) -> bool { - lowered::config_from_hfq(hfq).is_some_and(|cfg| cfg.enable_moe_block) -} fn gemma4_validate_drafter_route(is_e_series: bool, has_drafter: bool) -> Result<(), String> { if is_e_series && has_drafter { @@ -51,55 +25,13 @@ fn gemma4_validate_drafter_route(is_e_series: bool, has_drafter: bool) -> Result Ok(()) } -// ─── Bundle types ───────────────────────────────────────────────────────── - -pub struct Gemma4EagerBundle { - pub config: Gemma4Config, - pub weights: Gemma4Weights, - pub state: Gemma4State, -} - -pub struct Gemma4LoweredBundle { - pub config: lowered::Gemma4Config, - pub weights: lowered::Gemma4Weights, - pub scratch: lowered::Gemma4Scratch, - pub kv_sliding: KvCache, - pub kv_full: KvCache, -} - -pub enum Gemma4Bundle { - Eager(Gemma4EagerBundle), - Lowered(Gemma4LoweredBundle), -} - -fn lowered_kv_layer_counts(layer_types: &[lowered::LayerType]) -> (usize, usize) { - layer_types - .iter() - .fold((0, 0), |(sliding, full), layer_type| match layer_type { - lowered::LayerType::Sliding => (sliding + 1, full), - lowered::LayerType::Full => (sliding, full + 1), - }) -} - -/// Preserve the primary operation error; append cleanup failure context when present. -fn append_cleanup_context(op_err: String, cleanup: Result<(), String>) -> String { - match cleanup { - Ok(()) => op_err, - Err(c) => format!("{op_err}; cleanup also failed: {c}"), - } -} - /// Refuse MoE checkpoints whose expert formats have no indexed device kernel /// pair (the routed-expert step is capture-safe and never falls back to a /// host loop). -fn unsupported_expert_formats(weights: &lowered::Gemma4Weights) -> Option { +fn unsupported_expert_formats(weights: &Gemma4Weights) -> Option { weights.layers.iter().find_map(|layer| { - let moe = match layer { - lowered::LayerWeights::Sliding(l) => l.moe.as_ref(), - lowered::LayerWeights::Full(l) => l.moe.as_ref(), - }?; - let expert = moe.experts.first()?; - let (gate_up, down) = (expert.gate_up_proj.gpu_dtype, expert.down_proj.gpu_dtype); + let moe = layer.moe()?; + let (gate_up, down) = (moe.gate_up_dtype, moe.down_dtype); (!hipfire_dispatch::pipeline::sandwich::RoutedExperts::supports(gate_up, down)).then(|| { format!( "gemma4 MoE: expert formats gate_up={gate_up:?} down={down:?} have no indexed \ @@ -110,63 +42,24 @@ fn unsupported_expert_formats(weights: &lowered::Gemma4Weights) -> Option Result<(), String> { - let mut first: Option = None; - if let Err(e) = kv_sliding.free_gpu(gpu) { - first = Some(e.to_string()); - } - scratch.free_gpu(gpu); - free_lowered_weights(weights, gpu); - match first { - Some(e) => Err(e), - None => Ok(()), - } -} - -/// Resolve the lowered full-attention KV tier from the raw `kv_cache` value -/// through [`GEMMA4_LOWERED_FULL_POLICY`](hipfire_runtime::kv_mode::GEMMA4_LOWERED_FULL_POLICY): +/// Resolve the full-attention KV tier from the raw `kv_cache` value through +/// [`GEMMA4_FULL_POLICY`](hipfire_runtime::kv_mode::GEMMA4_FULL_POLICY): /// `auto`/`q8` → Q8, `legacy-asym3` → the previous Givens asym3 tier. /// Unsupported names fall back to Q8 with the returned warning; fp8/bf16 fail -/// closed (no Gemma4 lowered attend site admits a native tier). -fn resolve_lowered_full_kv_mode( - mode_raw: &str, -) -> Result<(hipfire_runtime::kv_mode::KvMode, Option<&'static str>), String> { +/// closed (no Gemma4 attend site admits a native tier). +fn resolve_full_kv_tier(mode_raw: &str) -> Result<(FullKvTier, Option<&'static str>), String> { use hipfire_runtime::kv_mode::{self, KvMode}; - let rr = kv_mode::resolve(mode_raw, &kv_mode::GEMMA4_LOWERED_FULL_POLICY); + let rr = kv_mode::resolve(mode_raw, &kv_mode::GEMMA4_FULL_POLICY); match rr.mode { - KvMode::Q8 | KvMode::Asym3 => Ok((rr.mode, rr.warning)), + KvMode::Q8 => Ok((FullKvTier::Q8, rr.warning)), + KvMode::Asym3 => Ok((FullKvTier::LegacyAsym3, rr.warning)), other => Err(format!( - "gemma4 (lowered): kv_cache={mode_raw} ({other:?}) has no full-attention KV tier; use auto, q8 or legacy-asym3" + "gemma4: kv_cache={mode_raw} ({other:?}) has no full-attention KV tier; use auto, q8 or legacy-asym3" )), } } -// ─── Bundle load ────────────────────────────────────────────────────────── - /// Build the Gemma 4 GPU bundle from an HFQ source. -/// -/// `ModelSource::Dir` returns the same error string the carrier previously -/// emitted inline. HFQ path: lowered (MoE) / eager selection, E-series -/// validation, weight/state/KV allocation, and the `eprintln!` diagnostics -/// for the chosen path. pub fn load_gemma4_bundle(src: ModelSource, ctx: &mut LoadCtx) -> Result { if ctx.kv_backend != hipfire_runtime::kv_backend::KvBackend::Legacy { return Err("gemma4: sliding/full KV owners require legacy backend".into()); @@ -177,247 +70,76 @@ pub fn load_gemma4_bundle(src: ModelSource, ctx: &mut LoadCtx) -> Result "legacy-asym3", - _ => "q8", - }; - let (n_sliding_layers, n_full_layers) = lowered_kv_layer_counts(&lcfg.layer_types); - let mut hfq2 = hfq; - // Construction order: weights → scratch → constants → sliding KV → full KV. - // On any later error free every completed earlier owner in reverse. - let weights = lowered::load_weights(&mut hfq2, &lcfg, ctx.gpu) - .map_err(|e| format!("gemma4 (lowered) load_weights: {e:?}"))?; - if let Some(e) = unsupported_expert_formats(&weights) { - free_lowered_weights(weights, ctx.gpu); - return Err(e); - } - let scratch = match lowered::Gemma4Scratch::new(ctx.gpu, &lcfg, ctx.max_seq) { - Ok(v) => v, - Err(e) => { - free_lowered_weights(weights, ctx.gpu); - return Err(format!("gemma4 (lowered) scratch: {e:?}")); - } - }; - if let Err(e) = lowered::init_scratch_constants(ctx.gpu, &scratch, lcfg.full_head_dim) { - free_lowered_scratch_and_weights(scratch, weights, ctx.gpu); - return Err(format!("gemma4 (lowered) init_scratch_constants: {e:?}")); - } - // Physical ring is min(window, max_seq); logical full/scratch stay max_seq. - let sliding_cap = lcfg.sliding_window.min(ctx.max_seq); - let kv_sliding = match KvCache::new_gpu_q8_capped( - ctx.gpu, - n_sliding_layers, - lcfg.sliding_n_kv_heads, - lcfg.sliding_head_dim, - ctx.max_seq, - sliding_cap, - ) { - Ok(v) => v, - Err(e) => { - free_lowered_scratch_and_weights(scratch, weights, ctx.gpu); - return Err(format!( - "gemma4 (lowered) sliding KV alloc (q8 ring): {e:?}" - )); - } - }; - // Full tier: q8 by default, legacy Givens asym3 on explicit opt-out. - // Logical limit is max_seq (no ring). - let kv_full_alloc = if full_kv_mode == hipfire_runtime::kv_mode::KvMode::Asym3 { - KvCache::new_gpu_asym3_gemma4( - ctx.gpu, - n_full_layers, - lcfg.full_n_kv_heads, - lcfg.full_head_dim, - ctx.max_seq, - ) - } else { - KvCache::new_gpu_q8( - ctx.gpu, - n_full_layers, - lcfg.full_n_kv_heads, - lcfg.full_head_dim, - ctx.max_seq, - ) - }; - let kv_full = match kv_full_alloc { - Ok(v) => v, - Err(e) => { - let cleanup = - free_lowered_sliding_scratch_weights(kv_sliding, scratch, weights, ctx.gpu); - return Err(append_cleanup_context( - format!("gemma4 (lowered) full KV alloc ({full_kv_label}): {e:?}"), - cleanup, - )); - } - }; eprintln!( - " gemma4 lowered path: moe (sliding q8-ring + full {full_kv_label} KV; kv_cache={mode_raw})" + " gemma4 E-series: {:?} (PLE + shared KV)", + config.e_series_variant()? ); - return Ok(Gemma4Bundle::Lowered(Gemma4LoweredBundle { - config: lcfg, - weights, - scratch, - kv_sliding, - kv_full, - })); } - // ── Eager dense / E-series path ── - let config = match eager_config { - Some(c) => c, - None => Gemma4Config::from_hfq(&hfq)?, - }; - if is_e_series { + gemma4_validate_drafter_route(is_e_series, ctx.gemma4_drafter_path.is_some())?; + + // Pure CPU: resolve the full-attention tier before any device upload. + let mode_raw = ctx + .kv_mode_override + .filter(|s| !s.is_empty()) + .map(str::to_string) + .unwrap_or_else(|| hipfire_runtime::config::get().kv_mode.clone()); + let (tier, warning) = resolve_full_kv_tier(&mode_raw)?; + if let Some(w) = warning { eprintln!( - " gemma4 E-series eager path: {:?} (PLE + shared KV)", - config.e_series_variant()? + " KV cache: {w} (site {})", + hipfire_runtime::kv_mode::GEMMA4_FULL_POLICY.site ); } + let weights = Gemma4Weights::load(&hfq, &config, ctx.gpu)?; - let state = Gemma4State::new_with_max_seq(ctx.gpu, &config, ctx.max_seq) - .map_err(|e| format!("gemma4: Gemma4State::new_with_max_seq failed: {e}"))?; - let _ = &weights; - Ok(Gemma4Bundle::Eager(Gemma4EagerBundle { - config, - weights, - state, - })) + if let Some(e) = unsupported_expert_formats(&weights) { + weights.free_gpu(ctx.gpu); + return Err(e); + } + let state = match Gemma4State::new_with_full_tier(ctx.gpu, &config, ctx.max_seq, tier) { + Ok(state) => state, + Err(e) => { + weights.free_gpu(ctx.gpu); + return Err(format!("gemma4: state: {e}")); + } + }; + eprintln!( + " gemma4: {} (sliding q8 + full {tier:?} KV; kv_cache={mode_raw})", + if config.enable_moe_block { + "moe" + } else { + "dense" + } + ); + Ok(Gemma4Bundle::new(config, weights, state, ctx.gpu)) } -// Alias for task's naming convention if callers use `load_bundle`. -pub use load_gemma4_bundle as load_bundle; - #[cfg(test)] mod tests { use super::*; #[test] - fn lowered_full_kv_mode_follows_kv_cache() { - use hipfire_runtime::kv_mode::KvMode; + fn full_kv_tier_follows_kv_cache() { for raw in ["", "auto", "q8"] { assert_eq!( - resolve_lowered_full_kv_mode(raw), - Ok((KvMode::Q8, None)), + resolve_full_kv_tier(raw), + Ok((FullKvTier::Q8, None)), "{raw:?}" ); } assert_eq!( - resolve_lowered_full_kv_mode("legacy-asym3"), - Ok((KvMode::Asym3, None)) + resolve_full_kv_tier("legacy-asym3"), + Ok((FullKvTier::LegacyAsym3, None)) ); // Bare asym3 is the Qwen FWHT alias: no hd512 FWHT tier, so q8 + warning. - let (mode, warning) = resolve_lowered_full_kv_mode("asym3").unwrap(); - assert_eq!(mode, KvMode::Q8); + let (tier, warning) = resolve_full_kv_tier("asym3").unwrap(); + assert_eq!(tier, FullKvTier::Q8); assert!(warning.is_some()); // Native presets fail closed instead of silently becoming q8. - assert!(resolve_lowered_full_kv_mode("fp8").is_err()); - assert!(resolve_lowered_full_kv_mode("bf16").is_err()); - } - - #[test] - fn lowered_kv_counts_follow_attention_layer_types() { - let layer_types = (0..48) - .map(|layer_idx| { - if (layer_idx + 1) % 6 == 0 { - lowered::LayerType::Full - } else { - lowered::LayerType::Sliding - } - }) - .collect::>(); - assert_eq!(lowered_kv_layer_counts(&layer_types), (40, 8)); - } - - #[test] - fn scratch_geometry_uses_single_max_seq_authority() { - // No HIPFIRE_KV_SEQ env var should affect geometry; both partials - // and KV are sized from the same ctx.max_seq. Verify the pure helpers - // that the loader now uses. - let max_seq_small = 32768usize; - let max_seq_large = 131072usize; - let n_heads = 32usize; - let full_hd = 512usize; - let s_small = lowered::gemma4_flash_partials_len(max_seq_small, n_heads, full_hd); - let s_large = lowered::gemma4_flash_partials_len(max_seq_large, n_heads, full_hd); - assert_eq!(s_small, 4_210_688); - assert_eq!(s_large, 16_842_752); - assert_eq!(s_large, 4 * s_small); - } - - #[test] - fn scratch_geometry_has_no_env_mismatch() { - // Simulate that an old env var could earlier cause mismatch between - // KV (ctx.max_seq) and scratch (HIPFIRE_KV_SEQ). After the fix both - // derive from the same max_seq, so the arithmetic must be identical - // for any max_seq value, including 131072 without GPU alloc. - for &max_seq in &[8192usize, 32768, 65536, 131072] { - let n_heads = 32; - let hd = 512; - let tiles = max_seq.div_ceil(lowered::GEMMA4_FLASH_TILE); - let expected = n_heads * tiles * (2 + hd); - assert_eq!( - lowered::gemma4_flash_partials_len(max_seq, n_heads, hd), - expected - ); - } - } - - #[test] - fn context_admission_refuses_lowered_below_floor() { - assert!(gemma4_context_admission(64, true).is_err()); - assert!(gemma4_context_admission(127, true).is_err()); - assert!(gemma4_context_admission(128, true).is_ok()); - assert!(gemma4_context_admission(512, true).is_ok()); - } - - #[test] - fn context_admission_eager_exempt_at_any_seq() { - // Eager has no min-context assert; small contexts stay admitted. - assert!(gemma4_context_admission(1, false).is_ok()); - assert!(gemma4_context_admission(64, false).is_ok()); - assert!(gemma4_context_admission(127, false).is_ok()); - assert!(gemma4_context_admission(128, false).is_ok()); + assert!(resolve_full_kv_tier("fp8").is_err()); + assert!(resolve_full_kv_tier("bf16").is_err()); } } diff --git a/crates/hipfire-arch-gemma4/src/config.rs b/crates/hipfire-arch-gemma4/src/config.rs index 4d10908062..a112172373 100644 --- a/crates/hipfire-arch-gemma4/src/config.rs +++ b/crates/hipfire-arch-gemma4/src/config.rs @@ -2,12 +2,12 @@ // Copyright (c) 2026 Kaden Schutt // hipfire — see LICENSE and NOTICE in the project root. -//! Gemma 4 dense (text-only) config. +//! Gemma 4 text config: dense 12B, E2B/E4B, and the 26B-A4B MoE. //! -//! Parsed from the HFQ `metadata_json.config.text_config.*` envelope. Stage-1 -//! is dense text-only: MoE and multimodal towers are out of scope, while dense -//! Gemma4 features such as KV sharing, double-wide MLP layers, and per-layer -//! inputs are represented in the config so the forward path can opt in safely. +//! Parsed from the HFQ `metadata_json.config.text_config.*` envelope. KV +//! sharing, double-wide MLP layers, per-layer inputs and the parallel MoE +//! block are all represented here so one weight stack and forward serve every +//! variant. Multimodal towers remain out of scope. // ─── Layer / RoPE types ───────────────────────────────────────────────── @@ -87,6 +87,13 @@ pub struct Gemma4Config { // the last non-shared layer of the same attention type for K/V. pub num_kv_shared_layers: usize, + // MoE (26B-A4B). When set, every layer runs routed experts in parallel + // with its dense FFN; the two outputs meet before the post-FFN norm. + pub enable_moe_block: bool, + pub moe_intermediate_size: usize, + pub num_experts: usize, + pub top_k_experts: usize, + // Output pub final_logit_softcapping: f32, // 30.0 — tanh(x/30)*30 pub tie_word_embeddings: bool, // true — lm_head aliases embed_tokens @@ -210,6 +217,18 @@ impl Gemma4Config { let vocab_size_per_layer_input = getu(tc, "vocab_size_per_layer_input").unwrap_or(vocab_size as u64) as usize; let num_kv_shared_layers = getu(tc, "num_kv_shared_layers").unwrap_or(0) as usize; + let enable_moe_block = getb(tc, "enable_moe_block").unwrap_or(false); + let moe_intermediate_size = getu(tc, "moe_intermediate_size").unwrap_or(0) as usize; + let num_experts = getu(tc, "num_experts").unwrap_or(0) as usize; + let top_k_experts = getu(tc, "top_k_experts").unwrap_or(0) as usize; + if enable_moe_block + && (moe_intermediate_size == 0 || num_experts == 0 || top_k_experts == 0) + { + return Err(format!( + "gemma4: enable_moe_block needs nonzero moe_intermediate_size/num_experts/top_k_experts \ + (got {moe_intermediate_size}/{num_experts}/{top_k_experts})" + )); + } let final_logit_softcapping = getf(tc, "final_logit_softcapping").unwrap_or(0.0) as f32; let tie_word_embeddings = getb(tc, "tie_word_embeddings") @@ -270,6 +289,10 @@ impl Gemma4Config { hidden_size_per_layer_input, vocab_size_per_layer_input, num_kv_shared_layers, + enable_moe_block, + moe_intermediate_size, + num_experts, + top_k_experts, final_logit_softcapping, tie_word_embeddings, embed_scale, @@ -553,6 +576,10 @@ mod tests { hidden_size_per_layer_input: 256, vocab_size_per_layer_input: 262144, num_kv_shared_layers: 20, + enable_moe_block: false, + moe_intermediate_size: 0, + num_experts: 0, + top_k_experts: 0, final_logit_softcapping: 30.0, tie_word_embeddings: true, embed_scale: (1536.0f32).sqrt(), diff --git a/crates/hipfire-arch-gemma4/src/drafter.rs b/crates/hipfire-arch-gemma4/src/drafter.rs index 99d8f460bf..182c5cdd8f 100644 --- a/crates/hipfire-arch-gemma4/src/drafter.rs +++ b/crates/hipfire-arch-gemma4/src/drafter.rs @@ -39,14 +39,18 @@ //! forward at a FIXED position. `use_ordered_embeddings`/centroid masking is //! disabled (the shipped head ties lm_head directly). //! -//! A draft step is one declarative step list: `Gemv(pre_projection)`, one +//! A draft step is one declarative step list: the target-embedding lookup and +//! its √backbone scale into the first concat half, `Gemv(pre_projection)`, one //! query-only `[SandwichAttention, SandwichMlp, Scale?]` block per layer over //! the target's last cache slot of the matching type (`crate::program`, the -//! same ops the target runs), the final norm, the tied `lm_head` and -//! `post_projection`. The target's weights and state are READ-ONLY here. +//! same ops the target runs), the final norm, and `post_projection` straight +//! into the second concat half (the next step's hidden input). The shared +//! `DraftHead` ranks the tied `lm_head` on the GPU and returns the draft. The +//! target's weights and state are READ-ONLY here. use crate::config::{Gemma4Config, LayerType, RopeType}; use crate::gemma4::{Gemma4State, Gemma4Weights}; +use hipfire_dispatch::pipeline::{DraftHead, DraftHeadLayout, DraftHeadPolicy}; use hipfire_runtime::hfq::{load_awq_scale, HfqFile}; use hipfire_runtime::llama::{f16_to_f32, WeightTensor}; use rdna_compute::{DType, Gpu, GpuTensor}; @@ -435,7 +439,6 @@ pub struct Gemma4DrafterWeights { pub post_projection: WeightTensor, // [backbone_hidden, hidden] pub layers: Vec, - embd_q8: bool, } impl Gemma4DrafterWeights { @@ -450,12 +453,11 @@ impl Gemma4DrafterWeights { let (embed_info, embed_data) = hfq .tensor_data(embed_name) .ok_or_else(|| "gemma4-drafter: embed_tokens not found".to_string())?; - let (embed_tokens, embd_dtype, embd_q8) = match embed_info.quant_type { + let (embed_tokens, embd_dtype) = match embed_info.quant_type { 3 => ( gpu.upload_raw(embed_data, &[embed_data.len()]) .map_err(|e| format!("gemma4-drafter: upload embed: {e:?}"))?, DType::Q8_0, - true, ), 1 => { let f32_data: Vec = embed_data @@ -466,7 +468,6 @@ impl Gemma4DrafterWeights { gpu.upload_f32(&f32_data, &[cfg.vocab_size, hidden]) .map_err(|e| format!("gemma4-drafter: upload embed f32: {e:?}"))?, DType::F32, - false, ) } 16 => { @@ -478,7 +479,6 @@ impl Gemma4DrafterWeights { gpu.upload_f32(&f32_data, &[cfg.vocab_size, hidden]) .map_err(|e| format!("gemma4-drafter: upload embed f32: {e:?}"))?, DType::F32, - false, ) } qt => return Err(format!("gemma4-drafter: unsupported embed quant_type {qt}")), @@ -606,7 +606,6 @@ impl Gemma4DrafterWeights { pre_projection, post_projection, layers, - embd_q8, }) } @@ -642,12 +641,14 @@ impl Gemma4DrafterWeights { /// Per-draft-step GPU scratch. No KV cache (the drafter reads the target's). pub struct Gemma4DrafterScratch { - pub concat: GpuTensor, // [2·backbone_hidden] = 7680 pre-proj input - pub embed_half: GpuTensor, // [backbone_hidden] = 3840 target-embed row scratch - pub x: GpuTensor, // [hidden] = 1024 residual stream - pub residual: GpuTensor, // [hidden] - pub tmp: GpuTensor, // [hidden] norm / o_proj scratch - /// FWHT rotation scratch for MagnumQuant projections, `[max_q_dim]`. + /// `[2·backbone_hidden]`: target embedding half ‖ hidden half. The hidden + /// half holds the target hidden for a round's first step, then each + /// step's `post_projection` output for the next. + pub concat: GpuTensor, + pub x: GpuTensor, // [hidden] residual stream + pub residual: GpuTensor, // [hidden] + pub tmp: GpuTensor, // [hidden] norm / o_proj scratch + /// FWHT rotation scratch for MagnumQuant projections, `[max(q, hidden, bb)]`. pub x_rot: GpuTensor, pub q: GpuTensor, // [max_q_dim] pub attn_out: GpuTensor, // [max_q_dim] @@ -656,202 +657,239 @@ pub struct Gemma4DrafterScratch { pub ffn_hidden: GpuTensor, // [hidden_dim] pub ffn_out: GpuTensor, // [hidden] pub normed: GpuTensor, // [hidden] final-norm output (drives lm_head + post_proj) - pub post_proj: GpuTensor, // [backbone_hidden] post_projection output - pub logits: GpuTensor, // [vocab] + /// `[1]` i32 input token of the step. + pub token_ids: GpuTensor, /// device i32 query position scalar (constant across a round). pub pos_buf: hip_bridge::DeviceBuffer, - /// Host copy of the query position staged in `pos_buf`. - pub query_pos: usize, + /// Ranks the tied `lm_head` and returns the draft token. + pub head: DraftHead, } impl Gemma4DrafterScratch { - pub fn new(gpu: &mut Gpu, cfg: &Gemma4DrafterConfig) -> Result { - let alloc = |g: &mut Gpu, n: usize, label: &str| -> Result { - g.zeros(&[n], DType::F32) - .map_err(|e| format!("gemma4-drafter: alloc {label}: {e:?}")) + pub fn new( + gpu: &mut Gpu, + cfg: &Gemma4DrafterConfig, + weights: &Gemma4DrafterWeights, + ) -> Result { + let mut owned: Vec = Vec::new(); + let mut alloc = |g: &mut Gpu, n: usize, label: &str| -> Result { + let t = g + .zeros(&[n], DType::F32) + .map_err(|e| format!("gemma4-drafter: alloc {label}: {e:?}"))?; + // SAFETY: rollback-only duplicate owner; see `crate::gemma4`. + owned.push(GpuTensor { + buf: unsafe { std::ptr::read(&t.buf) }, + shape: t.shape.clone(), + dtype: t.dtype, + }); + Ok(t) + }; + let built = (|| -> Result<_, String> { + let concat = alloc(gpu, cfg.pre_proj_in(), "concat")?; + let x = alloc(gpu, cfg.hidden, "x")?; + let residual = alloc(gpu, cfg.hidden, "residual")?; + let tmp = alloc(gpu, cfg.hidden, "tmp")?; + let rot = cfg.max_q_dim().max(cfg.hidden).max(cfg.backbone_hidden); + let x_rot = alloc(gpu, rot, "x_rot")?; + let q = alloc(gpu, cfg.max_q_dim(), "q")?; + let attn_out = alloc(gpu, cfg.max_q_dim(), "attn_out")?; + let gate_ffn = alloc(gpu, cfg.hidden_dim, "gate_ffn")?; + let up_ffn = alloc(gpu, cfg.hidden_dim, "up_ffn")?; + let ffn_hidden = alloc(gpu, cfg.hidden_dim, "ffn_hidden")?; + let ffn_out = alloc(gpu, cfg.hidden, "ffn_out")?; + let normed = alloc(gpu, cfg.hidden, "normed")?; + let token_ids = alloc(gpu, 1, "token_ids")?; + Ok(( + concat, x, residual, tmp, x_rot, q, attn_out, gate_ffn, up_ffn, ffn_hidden, + ffn_out, normed, token_ids, + )) + })(); + let rollback = |gpu: &mut Gpu, owned: Vec| { + for t in owned { + let _ = gpu.free_tensor(t); + } + }; + let ( + concat, + x, + residual, + tmp, + x_rot, + q, + attn_out, + gate_ffn, + up_ffn, + ffn_hidden, + ffn_out, + normed, + token_ids, + ) = match built { + Ok(v) => v, + Err(e) => { + rollback(gpu, owned); + return Err(e); + } + }; + let layout = DraftHeadLayout { + vocab: cfg.vocab_size, + hidden: cfg.hidden, + front: 0, + special: cfg.vocab_size, + full_hold: 0, + }; + let policy = DraftHeadPolicy { + copy: None, + rescore: false, + }; + let head = match DraftHead::new(gpu, &weights.lm_head.buf, layout, policy) { + Ok(head) => head, + Err(e) => { + rollback(gpu, owned); + return Err(format!("gemma4-drafter: draft head: {e}")); + } + }; + let pos_buf = match gpu.hip.malloc(4) { + Ok(buf) => buf, + Err(e) => { + let _ = head.free_gpu(gpu); + rollback(gpu, owned); + return Err(format!("gemma4-drafter: pos_buf malloc: {e:?}")); + } }; - let pos_buf = gpu - .hip - .malloc(4) - .map_err(|e| format!("gemma4-drafter: pos_buf malloc: {e:?}"))?; Ok(Gemma4DrafterScratch { - concat: alloc(gpu, cfg.pre_proj_in(), "concat")?, - embed_half: alloc(gpu, cfg.backbone_hidden, "embed_half")?, - x: alloc(gpu, cfg.hidden, "x")?, - residual: alloc(gpu, cfg.hidden, "residual")?, - tmp: alloc(gpu, cfg.hidden, "tmp")?, - x_rot: alloc(gpu, cfg.max_q_dim().max(cfg.hidden), "x_rot")?, - q: alloc(gpu, cfg.max_q_dim(), "q")?, - attn_out: alloc(gpu, cfg.max_q_dim(), "attn_out")?, - gate_ffn: alloc(gpu, cfg.hidden_dim, "gate_ffn")?, - up_ffn: alloc(gpu, cfg.hidden_dim, "up_ffn")?, - ffn_hidden: alloc(gpu, cfg.hidden_dim, "ffn_hidden")?, - ffn_out: alloc(gpu, cfg.hidden, "ffn_out")?, - normed: alloc(gpu, cfg.hidden, "normed")?, - post_proj: alloc(gpu, cfg.backbone_hidden, "post_proj")?, - logits: alloc(gpu, cfg.vocab_size, "logits")?, + concat, + x, + residual, + tmp, + x_rot, + q, + attn_out, + gate_ffn, + up_ffn, + ffn_hidden, + ffn_out, + normed, + token_ids, pos_buf, - query_pos: 0, + head, }) } - /// Return all scratch buffers to the pool. Consumes self. The raw - /// `pos_buf` DeviceBuffer is freed directly (it never went through the - /// tensor pool). + /// Return all scratch buffers to the pool. Consumes self. pub fn free_gpu(self, gpu: &mut Gpu) { - let _ = gpu.free_tensor(self.concat); - let _ = gpu.free_tensor(self.embed_half); - let _ = gpu.free_tensor(self.x); - let _ = gpu.free_tensor(self.residual); - let _ = gpu.free_tensor(self.tmp); - let _ = gpu.free_tensor(self.x_rot); - let _ = gpu.free_tensor(self.q); - let _ = gpu.free_tensor(self.attn_out); - let _ = gpu.free_tensor(self.gate_ffn); - let _ = gpu.free_tensor(self.up_ffn); - let _ = gpu.free_tensor(self.ffn_hidden); - let _ = gpu.free_tensor(self.ffn_out); - let _ = gpu.free_tensor(self.normed); - let _ = gpu.free_tensor(self.post_proj); - let _ = gpu.free_tensor(self.logits); + for t in [ + self.concat, + self.x, + self.residual, + self.tmp, + self.x_rot, + self.q, + self.attn_out, + self.gate_ffn, + self.up_ffn, + self.ffn_hidden, + self.ffn_out, + self.normed, + self.token_ids, + ] { + let _ = gpu.free_tensor(t); + } + let _ = self.head.free_gpu(gpu); let _ = gpu.hip.free(self.pos_buf); } -} -// ─── One draft step ───────────────────────────────────────────────────────── - -/// Result of one drafter step: the argmax draft token + its post-projected -/// hidden half (3840), which becomes the `hidden_3840` half of the NEXT step's -/// concat. Also exposes the final normed hidden + full logits for validation. -pub struct DrafterStepOut { - pub argmax: u32, - /// post_projection(normed) — [backbone_hidden]; the next step's hidden half. - pub post_proj_hidden: Vec, - /// final-norm output [hidden] (drives lm_head + post_proj) — for oracle. - pub normed_hidden: Vec, - /// full logits [vocab] — for oracle cosine. - pub logits: Vec, -} - -/// One drafter step (single-token forward at a FIXED `query_pos`). -/// -/// * `dw` / `dcfg` / `ds` — drafter weights / config / scratch. -/// * `target_weights` — for the target embed table (3840-dim) + √backbone scale. -/// * `target_state` — for the target's last sliding / last full KV slot. -/// * `prev_token` — token whose target-embed forms the concat's [0:3840] half. -/// * `hidden_backbone`— the [3840] hidden half (target's NORMED final hidden on -/// step 0, else the previous step's `post_proj_hidden`). -/// * `query_pos` — CONSTANT across the round = committed seq_len − 1. -/// -/// Reads `target_weights` + `target_state` ONLY; writes nothing back. -#[allow(clippy::too_many_arguments)] -pub fn drafter_step( - gpu: &mut Gpu, - dw: &Gemma4DrafterWeights, - dcfg: &Gemma4DrafterConfig, - ds: &mut Gemma4DrafterScratch, - target_weights: &Gemma4Weights, - target_state: &Gemma4State, - target_cfg: &Gemma4Config, - prev_token: u32, - hidden_backbone: &GpuTensor, - query_pos: usize, -) -> Result { - let bb = dcfg.backbone_hidden; - - // Set the device position scalar = query_pos (drives RoPE + attention range). - { - let pos_host = [query_pos as i32]; - let pos_bytes = unsafe { std::slice::from_raw_parts(pos_host.as_ptr() as *const u8, 4) }; - gpu.memcpy_htod_auto(&ds.pos_buf, pos_bytes) + /// Start a round at the constant `query_pos` with the target `hidden` + /// (`[backbone_hidden]`) as the first step's hidden half. + pub fn begin_round( + &self, + gpu: &mut Gpu, + hidden: &GpuTensor, + query_pos: usize, + ) -> Result<(), String> { + let bb = self.concat.numel() / 2; + gpu.memcpy_htod_auto(&self.pos_buf, &(query_pos as i32).to_ne_bytes()) .map_err(|e| format!("gemma4-drafter: htod pos: {e:?}"))?; - ds.query_pos = query_pos; + gpu.hip + .memcpy_dtod_at(&self.concat.buf, bb * 4, &hidden.buf, 0, bb * 4) + .map_err(|e| format!("gemma4-drafter: stage hidden half: {e:?}")) } +} - // ── (a) concat = [ target_embed(prev)·√bb ‖ hidden_backbone ] (7680) ── - // target_embed half: look up the TARGET embed row (bb-dim), scale by √bb. - // (The target's ScaledWordEmbedding bakes ·√bb into its forward, which the - // candidate generator invokes; we replicate that scale here.) +/// The target embedding table as a format-tagged tensor the embedding step +/// reads; `None` for formats without a batched lookup (F32 oracle tables). +pub fn target_embedding_table(target: &Gemma4Weights) -> Option { use hipfire_runtime::llama::EmbeddingFormat; - match target_weights.embd_format { - EmbeddingFormat::HFQ4G256 => gpu - .embedding_lookup_hfq4g256(&target_weights.embed_tokens, &ds.embed_half, prev_token, bb) - .map_err(|e| format!("gemma4-drafter: target embed hfq4g256: {e:?}"))?, - EmbeddingFormat::HFQ4G128 => gpu - .embedding_lookup_hfq4g128(&target_weights.embed_tokens, &ds.embed_half, prev_token, bb) - .map_err(|e| format!("gemma4-drafter: target embed hfq4g128: {e:?}"))?, - EmbeddingFormat::Q8_0 => gpu - .embedding_lookup_q8(&target_weights.embed_tokens, &ds.embed_half, prev_token, bb) - .map_err(|e| format!("gemma4-drafter: target embed q8: {e:?}"))?, - EmbeddingFormat::F32 => gpu - .embedding_lookup(&target_weights.embed_tokens, &ds.embed_half, prev_token, bb) - .map_err(|e| format!("gemma4-drafter: target embed f32: {e:?}"))?, - EmbeddingFormat::Q4K => { - return Err("gemma4-drafter: Q4K target embed unsupported".to_string()) - } - } - let _ = target_cfg; // embed scale uses bb (= target hidden), not target_cfg directly - gpu.scale_f32(&ds.embed_half, (bb as f32).sqrt()) - .map_err(|e| format!("gemma4-drafter: embed scale: {e:?}"))?; - // concat[0:bb] = embed_half; concat[bb:2bb] = hidden_backbone (no scale). - gpu.memcpy_dtod_at_auto(&ds.concat.buf, 0, &ds.embed_half.buf, 0, bb * 4) - .map_err(|e| format!("gemma4-drafter: concat embed half: {e:?}"))?; - gpu.memcpy_dtod_at_auto(&ds.concat.buf, bb * 4, &hidden_backbone.buf, 0, bb * 4) - .map_err(|e| format!("gemma4-drafter: concat hidden half: {e:?}"))?; - - // ── (b)-(f) pre_projection → layers → norm → lm_head/post_proj ── - // pos_buf is already staged above; concat is built into ds.concat. - drafter_step_from_concat_inner(gpu, dw, dcfg, ds, target_state, target_cfg) + let dtype = match target.embd_format { + EmbeddingFormat::Q8_0 => DType::Q8_0, + EmbeddingFormat::HFQ4G256 => DType::HFQ4G256, + EmbeddingFormat::HFQ4G128 => DType::HFQ4G128, + EmbeddingFormat::F32 | EmbeddingFormat::Q4K => return None, + }; + Some(GpuTensor { + // SAFETY: a view of the target-owned table for one step list. + buf: unsafe { target.embed_tokens.buf.alias() }, + shape: target.embed_tokens.shape.clone(), + dtype, + }) } -/// Run (b)-(f) of a draft step from a PRE-BUILT concat in `ds.concat` and a -/// PRE-STAGED `ds.pos_buf`. Used by `drafter_step` (after it builds the concat) -/// and by the validation harness (which injects an HF-exact concat to isolate -/// the backbone + heads from the embed/feedback path). +// ─── One draft step ───────────────────────────────────────────────────────── + +/// One draft step from `prev_token` at the round's query position: reads the +/// hidden half staged in `ds.concat`, leaves the next step's hidden half +/// there, and returns the draft token. Reads `target_weights` and +/// `target_state` only. #[allow(clippy::too_many_arguments)] -pub fn drafter_step_from_concat( +pub fn draft_step( gpu: &mut Gpu, dw: &Gemma4DrafterWeights, dcfg: &Gemma4DrafterConfig, ds: &mut Gemma4DrafterScratch, + target_weights: &Gemma4Weights, target_state: &Gemma4State, target_cfg: &Gemma4Config, - concat: &[f32], + prev_token: u32, query_pos: usize, -) -> Result { - { - let pos_host = [query_pos as i32]; - let pos_bytes = unsafe { std::slice::from_raw_parts(pos_host.as_ptr() as *const u8, 4) }; - gpu.memcpy_htod_auto(&ds.pos_buf, pos_bytes) - .map_err(|e| format!("gemma4-drafter: htod pos: {e:?}"))?; - ds.query_pos = query_pos; - } - if concat.len() != dcfg.pre_proj_in() { - return Err(format!( - "gemma4-drafter: concat len {} != {}", - concat.len(), - dcfg.pre_proj_in() - )); - } - gpu.hip - .memcpy_htod(&ds.concat.buf, unsafe { - std::slice::from_raw_parts(concat.as_ptr() as *const u8, concat.len() * 4) - }) - .map_err(|e| format!("gemma4-drafter: upload concat: {e:?}"))?; - drafter_step_from_concat_inner(gpu, dw, dcfg, ds, target_state, target_cfg) +) -> Result { + gpu.memcpy_htod_auto(&ds.token_ids.buf, &(prev_token as i32).to_ne_bytes()) + .map_err(|e| format!("gemma4-drafter: htod token: {e:?}"))?; + run_step( + gpu, + dw, + dcfg, + ds, + target_weights, + target_state, + target_cfg, + query_pos, + true, + )?; + ds.head + .draft(gpu, &dw.lm_head.buf, &ds.normed) + .map_err(|e| format!("gemma4-drafter: draft head: {e}")) } -fn drafter_step_from_concat_inner( +/// The step list. With `embed`, the first concat half is the target +/// embedding of `ds.token_ids` (× √backbone); without, the caller staged the +/// whole concat (validation harness). +#[allow(clippy::too_many_arguments)] +fn run_step( gpu: &mut Gpu, dw: &Gemma4DrafterWeights, dcfg: &Gemma4DrafterConfig, ds: &mut Gemma4DrafterScratch, + target_weights: &Gemma4Weights, target_state: &Gemma4State, target_cfg: &Gemma4Config, -) -> Result { + pos: usize, + embed: bool, +) -> Result<(), String> { use crate::program::{Geometry, LayerKv, LayerRefs, LayerScratch, ProgramBinding, Resident}; + use hipfire_dispatch::pipeline::sandwich::ScaleOp; + use hipfire_dispatch::pipeline::EmbeddingOp; use hipfire_dispatch::pipeline::{execute_steps, GemvInput, Step}; + let bb = dcfg.backbone_hidden; // Each block attends the target's LAST cache slot of its type. let last_slot = |n: usize, kind: &str| { n.checked_sub(1) @@ -859,7 +897,10 @@ fn drafter_step_from_concat_inner( }; let sliding_slot = last_slot(target_cfg.n_sliding_layers(), "sliding")?; let full_slot = last_slot(target_cfg.n_full_layers(), "full")?; - let pos = ds.query_pos; + let table = target_embedding_table(target_weights) + .ok_or("gemma4-drafter: target embedding format has no batched lookup")?; + let embed_half = ds.concat.sub_offset(0, bb); + let hidden_half = ds.concat.sub_offset(bb, bb); let binding = ProgramBinding { geo: Geometry::drafter(dcfg), resident: Resident { @@ -891,12 +932,24 @@ fn drafter_step_from_concat_inner( }, }; - // pre_projection → blocks → final norm → tied lm_head (no softcap) and - // post_projection, one step list. let pre = dw.pre_projection.dispatch_ref(); - let lm_head = dw.lm_head.dispatch_ref(); let post = dw.post_projection.dispatch_ref(); - let mut steps = Vec::with_capacity(3 * dcfg.n_layers + 4); + let mut steps = Vec::with_capacity(3 * dcfg.n_layers + 5); + if embed { + steps.push(Step::Embed(EmbeddingOp { + table: &table, + rotated: &ds.x_rot, + token_ids: &ds.token_ids, + output: &embed_half, + rows: 1, + dim: bb, + })); + // The target's ScaledWordEmbedding bakes ·√backbone into its forward. + steps.push(Step::Scale(ScaleOp { + x: &embed_half, + factor: (bb as f32).sqrt(), + })); + } steps.push(Step::Gemv { w: &pre, input: GemvInput::Raw(&ds.concat), @@ -928,47 +981,82 @@ fn drafter_step_from_concat_inner( eps: dcfg.norm_eps, rotation: hipfire_dispatch::types::RotationPlan::None, }); - steps.push(Step::Gemv { - w: &lm_head, - input: GemvInput::Raw(&ds.normed), - out: &ds.logits, - }); + // The next step's hidden half; pre_projection already consumed this one. steps.push(Step::Gemv { w: &post, input: GemvInput::Raw(&ds.normed), - out: &ds.post_proj, + out: &hidden_half, }); let ctx = hipfire_dispatch::context::DispatchCtx::new(gpu); - execute_steps(gpu, &ctx, &steps).map_err(|e| format!("gemma4-drafter: {e}"))?; - drop(steps); - - let logits = gpu - .download_f32(&ds.logits) - .map_err(|e| format!("gemma4-drafter: download logits: {e:?}"))?; - let argmax = argmax_f32(&logits); - let post_proj_hidden = gpu - .download_f32(&ds.post_proj) - .map_err(|e| format!("gemma4-drafter: download post_proj: {e:?}"))?; - let normed_hidden = gpu - .download_f32(&ds.normed) - .map_err(|e| format!("gemma4-drafter: download normed: {e:?}"))?; - let _ = dw.embd_q8; - Ok(DrafterStepOut { - argmax, - post_proj_hidden, - normed_hidden, - logits, - }) + execute_steps(gpu, &ctx, &steps).map_err(|e| format!("gemma4-drafter: {e}")) } -fn argmax_f32(v: &[f32]) -> u32 { - let mut bi = 0u32; - let mut bv = f32::NEG_INFINITY; - for (i, &x) in v.iter().enumerate() { - if x > bv { - bv = x; - bi = i as u32; - } +/// Oracle outputs of one validation step: the draft, its full logits, the +/// final normed hidden and the post-projected hidden half. +pub struct DrafterStepOut { + pub argmax: u32, + pub post_proj_hidden: Vec, + pub normed_hidden: Vec, + pub logits: Vec, +} + +/// Validation entry: run one step from a PRE-BUILT `concat` (HF-exact inputs +/// isolate the backbone and heads from the embedding/feedback path) at +/// `query_pos` and download every intermediate the oracle compares. +#[allow(clippy::too_many_arguments)] +pub fn drafter_step_from_concat( + gpu: &mut Gpu, + dw: &Gemma4DrafterWeights, + dcfg: &Gemma4DrafterConfig, + ds: &mut Gemma4DrafterScratch, + target_weights: &Gemma4Weights, + target_state: &Gemma4State, + target_cfg: &Gemma4Config, + concat: &[f32], + query_pos: usize, +) -> Result { + if concat.len() != dcfg.pre_proj_in() { + return Err(format!( + "gemma4-drafter: concat len {} != {}", + concat.len(), + dcfg.pre_proj_in() + )); } - bi + gpu.memcpy_htod_auto(&ds.pos_buf, &(query_pos as i32).to_ne_bytes()) + .map_err(|e| format!("gemma4-drafter: htod pos: {e:?}"))?; + gpu.hip + .memcpy_htod(&ds.concat.buf, unsafe { + std::slice::from_raw_parts(concat.as_ptr() as *const u8, concat.len() * 4) + }) + .map_err(|e| format!("gemma4-drafter: upload concat: {e:?}"))?; + run_step( + gpu, + dw, + dcfg, + ds, + target_weights, + target_state, + target_cfg, + query_pos, + false, + )?; + let argmax = ds + .head + .draft(gpu, &dw.lm_head.buf, &ds.normed) + .map_err(|e| format!("gemma4-drafter: draft head: {e}"))?; + let download = |gpu: &mut Gpu, t: &GpuTensor, what: &str| { + gpu.download_f32(t) + .map_err(|e| format!("gemma4-drafter: download {what}: {e:?}")) + }; + Ok(DrafterStepOut { + argmax, + post_proj_hidden: download( + gpu, + &ds.concat + .sub_offset(dcfg.backbone_hidden, dcfg.backbone_hidden), + "post_proj", + )?, + normed_hidden: download(gpu, &ds.normed, "normed")?, + logits: download(gpu, ds.head.logits(), "logits")?, + }) } diff --git a/crates/hipfire-arch-gemma4/src/emit.rs b/crates/hipfire-arch-gemma4/src/emit.rs new file mode 100644 index 0000000000..c285c3df8e --- /dev/null +++ b/crates/hipfire-arch-gemma4/src/emit.rs @@ -0,0 +1,641 @@ +// SPDX-License-Identifier: Apache-2.0 +// Copyright (c) 2026 Kaden Schutt +// hipfire — see LICENSE and NOTICE in the project root. + +//! Gemma 4 output routing: the thought-channel router that splits decoded +//! text into reasoning and visible answer, shared by the AR loop and the +//! speculative emitter. + +use hipfire_runtime::prompt_frame::ThinkMode; +use hipfire_runtime::spec::{ + ClientEvent, EmitOutcome, FinishSummary, SpecEmit, SpecEmitCtx, StopReason, +}; +use hipfire_runtime::tokenizer::{TokenTextStream, Tokenizer}; + +/// Gemma thought-channel router (minimal, follows Glimmer patterns and checked-in Gemma template). +/// +/// Gemma4 template uses `<|channel>thought\n` ... `\n` to wrap reasoning. +/// When `enable_thinking` is false the Jinja prompt already emits a closed empty +/// `thought` channel (`<|channel>thought\n`), so the router starts in +/// Answer mode. When true it starts awaiting the thought opening. Content inside +/// the thought channel is emitted as semantic `reasoning` events; +/// content outside is visible answer tokens. `max_think_tokens` is an orthogonal +/// force-close cap (0 = uncapped) that, when reached, flushes pending reasoning +/// and transitions to Answer, mirroring Glimmer's `max_think_tokens` handling. +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +pub enum GemmaChannel { + AwaitingThought, + Reasoning, + Answer, +} + +#[derive(Debug, Clone, PartialEq, Eq)] +pub enum GemmaEmit { + Reasoning(String), + Token(String), +} + +pub struct GemmaThoughtRouter { + pub state: GemmaChannel, + pub pending: String, + pub reasoning_tokens: usize, + pub max_think_tokens: usize, + pub enable_thinking: bool, + pub just_forced: bool, +} + +impl GemmaThoughtRouter { + pub fn new(enable_thinking: bool, max_think_tokens: usize) -> Self { + let state = if enable_thinking { + GemmaChannel::AwaitingThought + } else { + GemmaChannel::Answer + }; + Self { + state, + pending: String::new(), + reasoning_tokens: 0, + max_think_tokens, + enable_thinking, + just_forced: false, + } + } + + pub fn push(&mut self, frag: &str) -> (Vec, bool) { + if frag.is_empty() { + return (Vec::new(), false); + } + self.pending.push_str(frag); + let mut out = Vec::new(); + let entered_as_reasoning = self.state == GemmaChannel::Reasoning; + loop { + match self.state { + GemmaChannel::AwaitingThought => { + const CHANNEL_OPEN: &str = "<|channel>"; + const THOUGHT_OPEN: &str = "<|channel>thought"; + + // A thought channel is optional and may be split across + // decoded fragments. Hold only while the bytes seen so far + // can still become the canonical opening header. + if THOUGHT_OPEN.starts_with(&self.pending) { + break; + } + if self.pending.starts_with(THOUGHT_OPEN) { + let mut header_end = THOUGHT_OPEN.len(); + if self.pending[header_end..].starts_with('\n') { + header_end += 1; + } + self.pending.drain(..header_end); + self.state = GemmaChannel::Reasoning; + continue; + } + + // The response schema makes the thought header optional. + // Once the buffered bytes cannot form that header, route + // them as answer content. Some Gemma4 checkpoints emit an + // orphan `<|channel>` before an otherwise valid answer; + // consume that control token without dropping its payload. + if self.pending.starts_with(CHANNEL_OPEN) { + self.pending.drain(..CHANNEL_OPEN.len()); + if self.pending.starts_with('\n') { + self.pending.drain(..1); + } + } + self.state = GemmaChannel::Answer; + continue; + } + GemmaChannel::Reasoning => { + if let Some(pos) = self.pending.find("") { + if pos > 0 { + let text = self.pending[..pos].to_string(); + self.pending.drain(..pos); + if !text.is_empty() { + out.push(GemmaEmit::Reasoning(text)); + } + continue; + } + let end = "".len(); + self.pending.drain(..end); + if self.pending.starts_with('\n') { + self.pending.drain(..1); + } + self.state = GemmaChannel::Answer; + continue; + } else { + let hold = gemma_longest_marker_suffix(&self.pending); + let emit_len = self.pending.len().saturating_sub(hold); + if emit_len > 0 { + let text = self.pending[..emit_len].to_string(); + self.pending.drain(..emit_len); + if !text.is_empty() { + out.push(GemmaEmit::Reasoning(text)); + } + } + break; + } + } + GemmaChannel::Answer => { + // Answer must chunk-safely strip any Gemma channel + // control markers, including canonical <|channel|> + // forms, orphan <|channel> variants, and markers + // arriving after a forced max-think transition. + // Preserve payload before/after each marker and hold + // a suffix that could still become a marker. + const ANSWER_MARKERS: &[&str] = &[ + "<|channel>thought", + "<|channel>", + "<|channel|>", + "", + "<|turn>", + "", + ]; + loop { + let hold = gemma_longest_marker_suffix(&self.pending); + let search_len = self.pending.len().saturating_sub(hold); + let searchable = &self.pending[..search_len]; + let mut best_pos: Option = None; + let mut best_len = 0usize; + for &m in ANSWER_MARKERS { + if let Some(pos) = searchable.find(m) { + if best_pos.is_none() + || pos < best_pos.unwrap() + || (pos == best_pos.unwrap() && m.len() > best_len) + { + best_pos = Some(pos); + best_len = m.len(); + } + } + } + if let Some(pos) = best_pos { + if pos > 0 { + let text = self.pending[..pos].to_string(); + self.pending.drain(..pos); + if !text.is_empty() { + out.push(GemmaEmit::Token(text)); + } + continue; + } + self.pending.drain(..best_len); + if self.pending.starts_with('\n') { + self.pending.drain(..1); + } + if self.pending.is_empty() { + break; + } + continue; + } + if search_len > 0 { + let text = self.pending[..search_len].to_string(); + self.pending.drain(..search_len); + if !text.is_empty() { + out.push(GemmaEmit::Token(text)); + } + } + break; + } + break; + } + } + } + let emitted_reasoning = out.iter().any(|e| matches!(e, GemmaEmit::Reasoning(_))); + let should_count = entered_as_reasoning || emitted_reasoning; + if should_count + && self.max_think_tokens != 0 + && self.reasoning_tokens < self.max_think_tokens + { + self.reasoning_tokens += 1; + if self.reasoning_tokens >= self.max_think_tokens + && self.state == GemmaChannel::Reasoning + { + if !self.pending.is_empty() { + let tail = std::mem::take(&mut self.pending); + if !gemma_is_marker_prefix(&tail) && !tail.is_empty() { + out.push(GemmaEmit::Reasoning(tail)); + } + } + self.state = GemmaChannel::Answer; + self.just_forced = true; + } + } + (out, false) + } + + pub fn flush(&mut self) -> Vec { + if self.pending.is_empty() { + return Vec::new(); + } + if self.state == GemmaChannel::AwaitingThought && gemma_is_marker_prefix(&self.pending) { + self.pending.clear(); + return Vec::new(); + } + let text = std::mem::take(&mut self.pending); + if gemma_is_marker_prefix(&text) { + return Vec::new(); + } + vec![if self.state == GemmaChannel::Reasoning { + GemmaEmit::Reasoning(text) + } else { + GemmaEmit::Token(text) + }] + } +} + +pub fn gemma_is_marker_prefix(s: &str) -> bool { + const MARKERS: &[&str] = &[ + "<|channel>thought", + "<|channel>", + "<|channel|>", + "", + "<|turn>", + "", + ]; + MARKERS.iter().any(|m| m.starts_with(s)) +} + +/// [`SpecEmit`] for Gemma 4: the AR loop's emission (incremental decode +/// through the [`GemmaThoughtRouter`]) over speculatively committed tokens. +/// A stop token ends the turn uncommitted, as in AR. +pub struct GemmaSpecEmit<'a> { + tokenizer: &'a Tokenizer, + stops: Vec, + router: GemmaThoughtRouter, + text: TokenTextStream, + committed: usize, +} + +impl<'a> GemmaSpecEmit<'a> { + /// Stops on the target EOS, `im_end` (the config EOS) and Gemma's + /// `` id 106; thinking is on unless `think_mode` is + /// `NonThink`. + pub fn new(ctx: SpecEmitCtx<'a>) -> Self { + let mut stops = vec![ctx.eos, 106]; + stops.extend(ctx.im_end); + Self { + tokenizer: ctx.tokenizer, + stops, + router: GemmaThoughtRouter::new( + !matches!(ctx.think_mode, ThinkMode::NonThink), + ctx.max_think, + ), + text: TokenTextStream::new(), + committed: 0, + } + } +} + +fn client_events(emits: Vec, events: &mut Vec) { + events.extend(emits.into_iter().map(|e| match e { + GemmaEmit::Reasoning(text) => ClientEvent::Reasoning(text), + GemmaEmit::Token(text) => ClientEvent::Token(text), + })); +} + +impl SpecEmit for GemmaSpecEmit<'_> { + fn begin(&mut self, first_token: u32) -> EmitOutcome { + self.observe(first_token) + } + + fn observe(&mut self, token: u32) -> EmitOutcome { + if self.stops.contains(&token) { + return EmitOutcome { + events: Vec::new(), + stop: Some(StopReason::Eos), + }; + } + let frag = self.text.push(self.tokenizer, token); + let mut events = Vec::new(); + client_events(self.router.push(&frag).0, &mut events); + events.push(ClientEvent::Committed { + id: token, + idx: self.committed, + }); + self.committed += 1; + EmitOutcome { events, stop: None } + } + + fn finish(mut self: Box) -> FinishSummary { + let tail = self.text.flush(); + let mut events = Vec::new(); + client_events(self.router.push(&tail).0, &mut events); + client_events(self.router.flush(), &mut events); + FinishSummary { + events, + finish_reason: "stop", + ..FinishSummary::default() + } + } +} + +#[cfg(test)] +mod gemma_thought_router_tests { + use super::{gemma_is_marker_prefix, GemmaChannel, GemmaEmit, GemmaThoughtRouter}; + + fn route(enable_thinking: bool, chunks: &[&str]) -> (String, String, GemmaChannel) { + let mut router = GemmaThoughtRouter::new(enable_thinking, 0); + let mut visible = String::new(); + let mut reasoning = String::new(); + for chunk in chunks { + for event in router.push(chunk).0 { + match event { + GemmaEmit::Reasoning(text) => reasoning.push_str(&text), + GemmaEmit::Token(text) => visible.push_str(&text), + } + } + } + for event in router.flush() { + match event { + GemmaEmit::Reasoning(text) => reasoning.push_str(&text), + GemmaEmit::Token(text) => visible.push_str(&text), + } + } + (visible, reasoning, router.state) + } + + #[test] + fn gemma_router_routes_canonical_thought_then_answer() { + let (visible, reasoning, state) = route( + true, + &["<|channel>", "thought", "\nplan\nanswer"], + ); + assert_eq!(reasoning, "plan"); + assert_eq!(visible, "answer"); + assert_eq!(state, GemmaChannel::Answer); + } + + #[test] + fn gemma_router_recovers_orphan_channel_before_answer() { + let (visible, reasoning, state) = + route(true, &["<|channel>", "\n", "```python\nprint('ok')\n```"]); + assert_eq!(visible, "```python\nprint('ok')\n```"); + assert!(reasoning.is_empty()); + assert_eq!(state, GemmaChannel::Answer); + } + + #[test] + fn gemma_router_orphan_channel_is_chunk_boundary_invariant() { + let full = "<|channel>\nanswer"; + let expected = route(true, &[full]); + for split in 1..full.len() { + if full.is_char_boundary(split) { + assert_eq!(route(true, &[&full[..split], &full[split..]]), expected); + } + } + } + + #[test] + fn gemma_router_thinking_request_can_emit_direct_answer() { + let (visible, reasoning, state) = route(true, &["direct answer"]); + assert_eq!(visible, "direct answer"); + assert!(reasoning.is_empty()); + assert_eq!(state, GemmaChannel::Answer); + } + + #[test] + fn gemma_router_drops_only_an_unfinished_control_marker_at_eos() { + let (visible, reasoning, state) = route(true, &["<|chan"]); + assert!(visible.is_empty()); + assert!(reasoning.is_empty()); + assert_eq!(state, GemmaChannel::AwaitingThought); + } + + #[test] + fn gemma_marker_prefix_does_not_classify_marker_plus_payload() { + assert!(gemma_is_marker_prefix("<|chan")); + assert!(gemma_is_marker_prefix("<|channel>")); + assert!(!gemma_is_marker_prefix("<|channel>\nanswer")); + } + + fn route_with_cap( + enable_thinking: bool, + max_think_tokens: usize, + chunks: &[&str], + ) -> (String, String, GemmaChannel) { + let mut router = GemmaThoughtRouter::new(enable_thinking, max_think_tokens); + let mut visible = String::new(); + let mut reasoning = String::new(); + for chunk in chunks { + for event in router.push(chunk).0 { + match event { + GemmaEmit::Reasoning(text) => reasoning.push_str(&text), + GemmaEmit::Token(text) => visible.push_str(&text), + } + } + } + for event in router.flush() { + match event { + GemmaEmit::Reasoning(text) => reasoning.push_str(&text), + GemmaEmit::Token(text) => visible.push_str(&text), + } + } + (visible, reasoning, router.state) + } + + fn assert_no_markers(s: &str) { + for m in &[ + "<|channel>thought", + "<|channel>", + "<|channel|>", + "", + "<|turn>", + "", + ] { + assert!(!s.contains(m), "visible leaked marker {:?} in {:?}", m, s); + } + } + + #[test] + fn gemma_router_thinking_off_strips_all_channel_markers_every_split() { + // Thinking-off starts in Answer; every channel/turn marker must be + // stripped chunk-safely regardless of split. + let full = "pre<|channel>midpost<|channel|>inner<|channel>thought\nXtail<|turn>endafter"; + let expected = route_with_cap(false, 0, &[full]); + assert_no_markers(&expected.0); + // Adjacent payload must survive: markers stripped, text joined. + assert_eq!(expected.0, "premidpostinnerXtailendafter"); + for split in 1..full.len() { + if !full.is_char_boundary(split) { + continue; + } + let got = route_with_cap(false, 0, &[&full[..split], &full[split..]]); + assert_eq!(got, expected, "mismatch at split {}", split); + assert_no_markers(&got.0); + } + // Also verify orphan <|channel> with newline framing is stripped. + let full2 = "<|channel>\nanswer"; + let exp2 = route_with_cap(false, 0, &[full2]); + assert_eq!(exp2.0, "answer"); + for split in 1..full2.len() { + if !full2.is_char_boundary(split) { + continue; + } + assert_eq!( + route_with_cap(false, 0, &[&full2[..split], &full2[split..]]), + exp2 + ); + } + // Canonical <|channel|> in thinking-off must also be stripped. + let full3 = "A<|channel|>B"; + let exp3 = route_with_cap(false, 0, &[full3]); + assert_eq!(exp3.0, "AB"); + assert_no_markers(&exp3.0); + for split in 1..full3.len() { + if !full3.is_char_boundary(split) { + continue; + } + assert_eq!( + route_with_cap(false, 0, &[&full3[..split], &full3[split..]]), + exp3 + ); + } + } + + #[test] + fn gemma_router_forced_close_strips_markers_after_transition_every_split() { + // Force max-think after one reasoning push, then ensure every + // subsequent channel/turn marker in Answer is stripped at every + // split, preserving adjacent payload. + let pre = "<|channel>thought\nAAA"; + let post = "BBB<|channel>CCCDDD<|channel|>EEE<|turn>FFFGGG"; + let full = format!("{}{}", pre, post); + // full = "<|channel>thought\nAAABBB<|channel>CCCDDD<|channel|>EEE<|turn>FFFGGG" + // With max_think=1 the router forces to Answer after the first + // reasoning push; the trailing that closes thought and + // all markers inside post must be stripped, not leaked. + let expected = route_with_cap(true, 1, &[&full]); + assert!(expected.1.contains("AAA") || expected.0.contains("AAA")); + assert_no_markers(&expected.0); + // Answer payload should be the post text with markers removed. + // Post without markers: "BBBCCCDDDEEEFFFGGG" + assert_eq!(expected.0, "BBBCCCDDDEEEFFFGGG"); + // For split invariance after forced close, keep the reasoning header + // as one chunk and only split the post payload. Splitting the header + // itself changes per-push reasoning counting and is not required to + // be invariant for this test. + for split in 0..=post.len() { + if split != 0 && !post.is_char_boundary(split) { + continue; + } + let got = if split == 0 || split == post.len() { + route_with_cap(true, 1, &[&full]) + } else { + let c1 = &post[..split]; + let c2 = &post[split..]; + route_with_cap(true, 1, &[pre, c1, c2]) + }; + assert_eq!(got.0, expected.0, "forced mismatch at post split {}", split); + assert_no_markers(&got.0); + } + // Also test forced transition where pending marker is split across + // the forced boundary: reasoning chunk ends with partial marker prefix. + let full2 = "<|channel>thought\nRRXX<|channel>YY"; + let exp2 = route_with_cap(true, 1, &[full2]); + assert_no_markers(&exp2.0); + // Split only the post part after the forced close to keep reasoning counting stable + let pre2 = "<|channel>thought\nRR"; + let post2 = "XX<|channel>YY"; + let exp2_post = route_with_cap(true, 1, &[full2]); + for split in 0..=post2.len() { + if split != 0 && !post2.is_char_boundary(split) { + continue; + } + let got = if split == 0 || split == post2.len() { + route_with_cap(true, 1, &[full2]) + } else { + route_with_cap(true, 1, &[pre2, &post2[..split], &post2[split..]]) + }; + assert_eq!( + got.0, exp2.0, + "forced split2 mismatch at post split {}", + split + ); + assert_no_markers(&got.0); + } + // Verify that a marker arriving strictly after forced transition + // as a separate push is still stripped at every internal split. + for payload in &[ + "helloworld", + "hello<|channel>world", + "hello<|channel|>world", + "helloworld", + "hello<|turn>world", + ] { + let expected_payload = (*payload) + .replace("<|channel>thought", "") + .replace("<|channel>", "") + .replace("<|channel|>", "") + .replace("", "") + .replace("<|turn>", "") + .replace("", ""); + for split in 0..=payload.len() { + if split != 0 && !payload.is_char_boundary(split) { + continue; + } + // Build payload split after forced transition. + let full = if split == 0 || split == payload.len() { + format!("<|channel>thought\nZ{}", payload) + } else { + // Simulate payload split across two pushes after forced. + // Create a single concatenated string and test split invariance + // via the full-string split test already done; here just + // verify the payload alone after forced. + let c1 = &payload[..split]; + let c2 = &payload[split..]; + let mut rr = GemmaThoughtRouter::new(true, 1); + let _ = rr.push("<|channel>thought\nZ"); + let mut vis = String::new(); + for ev in rr.push("").0 { + if let GemmaEmit::Token(t) = ev { + vis.push_str(&t); + } + } + for ev in rr.push(c1).0 { + if let GemmaEmit::Token(t) = ev { + vis.push_str(&t); + } + } + for ev in rr.push(c2).0 { + if let GemmaEmit::Token(t) = ev { + vis.push_str(&t); + } + } + for ev in rr.flush() { + if let GemmaEmit::Token(t) = ev { + vis.push_str(&t); + } + } + assert_eq!( + vis, expected_payload, + "payload {:?} split {}", + payload, split + ); + assert_no_markers(&vis); + continue; + }; + let got = route_with_cap(true, 1, &[&full]); + assert!(got.0.contains(&expected_payload) || got.0 == expected_payload); + assert_no_markers(&got.0); + } + } + } +} + +pub fn gemma_longest_marker_suffix(s: &str) -> usize { + const MARKERS: &[&str] = &[ + "<|channel>thought", + "<|channel>", + "<|channel|>", + "", + "<|turn>", + "", + ]; + for len in (1..=s.len()).rev() { + if !s.is_char_boundary(s.len() - len) { + continue; + } + let suffix = &s[s.len() - len..]; + if MARKERS.iter().any(|m| m.starts_with(suffix)) { + return len; + } + } + 0 +} diff --git a/crates/hipfire-arch-gemma4/src/forward.rs b/crates/hipfire-arch-gemma4/src/forward.rs index f215a8772a..6c0cafa599 100644 --- a/crates/hipfire-arch-gemma4/src/forward.rs +++ b/crates/hipfire-arch-gemma4/src/forward.rs @@ -64,16 +64,6 @@ use hipfire_dispatch::pipeline::execute_steps; use hipfire_runtime::llama::{weight_gemv, WeightTensor}; use rdna_compute::{DType, Gpu, GpuTensor}; -/// Greedy EAGLE promises byte-identical target decisions. Its batched verify -/// must therefore use the same arithmetic family as eager decode rather than -/// numerically-close fused/WMMA variants whose small drift accumulates in KV. -fn eagle_strict_enabled() -> bool { - hipfire_config::developer_var("HIPFIRE_GEMMA4_EAGLE") - .ok() - .as_deref() - == Some("1") -} - /// Decode one token (eager); returns the full logits vector. Used for prefill, /// the warm pass, and as the `HIPFIRE_GEMMA4_GRAPH=0` fallback. pub fn decode_step( @@ -85,11 +75,61 @@ pub fn decode_step( position: u32, ) -> Result, String> { prepare_token_inputs(cfg, weights, state, gpu, token_id)?; + stage_position(state, gpu, position)?; decode_step_body(cfg, weights, state, gpu, position, None)?; gpu.download_f32(&state.logits) .map_err(|e| format!("gemma4: download logits: {e:?}")) } +/// Stage `position` into the device scalar every position-reading kernel +/// uses. Inside a hipGraph capture the copy is recorded and re-reads +/// `state.pos_host` on replay. +fn stage_position(state: &mut Gemma4State, gpu: &mut Gpu, position: u32) -> Result<(), String> { + state.pos_host[0] = position as i32; + let pos_bytes = unsafe { std::slice::from_raw_parts(state.pos_host.as_ptr() as *const u8, 4) }; + gpu.memcpy_htod_auto(&state.pos_buf.buf, pos_bytes) + .map_err(|e| format!("gemma4: htod pos: {e:?}")) +} + +/// The token-dependent prologue of one decode step: embedding (× √dim) into +/// `state.x` and, on E-series checkpoints, the per-layer inputs. +pub fn prepare_token( + cfg: &Gemma4Config, + weights: &Gemma4Weights, + state: &mut Gemma4State, + gpu: &mut Gpu, + token_id: u32, +) -> Result<(), String> { + prepare_token_inputs(cfg, weights, state, gpu, token_id) +} + +/// The position-independent body of one decode step (every layer, final +/// norm, LM head, softcap into `state.logits`). It reads the token from +/// `state.x` ([`prepare_token`]) and the position from `state.pos_buf`, which +/// the caller staged; it issues no copies, so a launch recorder sees the +/// whole step. +pub fn decode_body( + cfg: &Gemma4Config, + weights: &Gemma4Weights, + state: &mut Gemma4State, + gpu: &mut Gpu, + position: u32, +) -> Result<(), String> { + decode_step_body(cfg, weights, state, gpu, position, None) +} + +/// Whether replaying a captured decode reproduces direct decode exactly. The +/// Q8_0 expert `down_proj` accumulates with atomics, whose scheduling-order +/// sum is not exact; a launch recorder must see direct launches. +fn graph_exact(weights: &Gemma4Weights, gpu: &Gpu) -> bool { + !gpu.replay.is_recording() + && weights + .layers + .iter() + .filter_map(LayerWeights::moe) + .all(|moe| moe.down_dtype == DType::HFQ4G128) +} + /// Decode one token, appending each layer's post-residual hidden state (pre /// final-norm) to `capture[layer]` — used by the oracle dumper. Eager only /// (the per-layer D2H downloads are incompatible with graph capture). @@ -103,6 +143,7 @@ pub fn decode_step_capture( capture: &mut [Vec], ) -> Result<(), String> { prepare_token_inputs(cfg, weights, state, gpu, token_id)?; + stage_position(state, gpu, position)?; decode_step_body(cfg, weights, state, gpu, position, Some(capture)) } @@ -144,7 +185,7 @@ pub fn decode_step_with_graph( // The captured path is the default. Set HIPFIRE_GEMMA4_GRAPH=0 to retain // the eager fallback for diagnostics. let graph_on = env_override.unwrap_or(true); - if !graph_on { + if !graph_on || !graph_exact(weights, gpu) { return decode_step(cfg, weights, state, gpu, token_id, position); } @@ -173,16 +214,22 @@ pub fn decode_step_with_graph( if gpu.graphs.graph_exec.is_none() { // ── Capture phase ────────────────────────────────────────────── - // decode_step_body stages pos_host → pos_buf via memcpy_htod_auto - // INSIDE the capture, so the recorded memcpy node re-reads pos_host - // on each replay. + // The position copy is captured too, so the recorded memcpy node + // re-reads pos_host on each replay. gpu.graphs .begin_graph_capture(&gpu.hip, gpu.device_id, gpu.active_stream.as_ref().unwrap()) .map_err(|e| format!("gemma4 begin_graph_capture: {e:?}"))?; - decode_step_body(cfg, weights, state, gpu, position, None)?; - gpu.graphs + let body = stage_position(state, gpu, position) + .and_then(|()| decode_step_body(cfg, weights, state, gpu, position, None)); + let end = gpu + .graphs .end_graph_capture(&gpu.hip, gpu.device_id, gpu.active_stream.as_ref().unwrap()) - .map_err(|e| format!("gemma4 end_graph_capture: {e:?}"))?; + .map_err(|e| format!("gemma4 end_graph_capture: {e:?}")); + if let Err(e) = body.and(end) { + gpu.graphs.capture_mode = false; + gpu.graphs.graph_destroy(&gpu.hip, gpu.device_id); + return Err(e); + } // Captured kernels were RECORDED, not run — launch once so this token's // logits actually get produced. gpu.graphs @@ -483,7 +530,7 @@ fn eager_resident(state: &Gemma4State) -> Resident<'_> { Resident { kv_sliding: &state.kv_sliding, kv_full: &state.kv_full, - pos_buf: &state.pos_buf, + pos_buf: &state.pos_buf.buf, v_norm_ones: &state.v_norm_ones, flash_partials: &state.q8_flash_partials, } @@ -504,16 +551,7 @@ fn decode_step_body( )); } - // Device position scalar (i32). Staged from the heap-stable `state.pos_host` - // so the captured memcpy re-reads it on replay (memcpy_htod_auto → async on - // the capture stream when capturing). - state.pos_host[0] = position as i32; - { - let pos_bytes = - unsafe { std::slice::from_raw_parts(state.pos_host.as_ptr() as *const u8, 4) }; - gpu.memcpy_htod_auto(&state.pos_buf, pos_bytes) - .map_err(|e| format!("gemma4: htod pos: {e:?}"))?; - } + // The caller staged `position` into `state.pos_buf`. // Per-layer program: [SandwichAttention, SandwichMlp, PerLayerInput?, Scale?]. let ctx = DispatchCtx::new(gpu); @@ -583,32 +621,35 @@ fn decode_step_body( // ADDITIVE: does not touch `decode_step` / `decode_step_with_graph`. The eager // path is unchanged. -/// Whether every projection used by batched prefill has a matching kernel. -/// Other loadable quantizations must remain on eager prefill. -pub fn supports_batched_prefill(weights: &Gemma4Weights) -> bool { +/// Whether every projection used by batched prefill has a matching kernel +/// and both caches are Q8 (the batched attention writes Q8 rows). Other +/// loads remain on eager prefill. +pub fn supports_batched_prefill(weights: &Gemma4Weights, state: &Gemma4State) -> bool { let supports = |weight: &WeightTensor| { hipfire_dispatch::pipeline::sandwich::supports_batched_projection(weight.gpu_dtype) }; - weights.layers.iter().all(|layer| match layer { - LayerWeights::Sliding(layer) => { - supports(&layer.q_proj) - && supports(&layer.k_proj) - && supports(&layer.v_proj) - && supports(&layer.o_proj) - && supports(&layer.gate_proj) - && supports(&layer.up_proj) - && supports(&layer.down_proj) - } - LayerWeights::Full(layer) => { - supports(&layer.q_proj) - && supports(&layer.k_proj) - && layer.v_proj.as_ref().map_or(true, supports) - && supports(&layer.o_proj) - && supports(&layer.gate_proj) - && supports(&layer.up_proj) - && supports(&layer.down_proj) - } - }) + state.kv_sliding.quant_q8 + && state.kv_full.quant_q8 + && weights.layers.iter().all(|layer| match layer { + LayerWeights::Sliding(layer) => { + supports(&layer.q_proj) + && supports(&layer.k_proj) + && supports(&layer.v_proj) + && supports(&layer.o_proj) + && supports(&layer.gate_proj) + && supports(&layer.up_proj) + && supports(&layer.down_proj) + } + LayerWeights::Full(layer) => { + supports(&layer.q_proj) + && supports(&layer.k_proj) + && layer.v_proj.as_ref().map_or(true, supports) + && supports(&layer.o_proj) + && supports(&layer.gate_proj) + && supports(&layer.up_proj) + && supports(&layer.down_proj) + } + }) } /// argmax over a logits row (spec-decode greedy per-position prediction). @@ -682,6 +723,57 @@ pub fn forward_batch_spec( start_pos: usize, per_token_hidden_out: Option<&GpuTensor>, per_pos_argmax_out: Option<&mut Vec>, +) -> Result, String> { + forward_rows( + cfg, + weights, + state, + gpu, + tokens, + start_pos, + per_token_hidden_out, + per_pos_argmax_out, + true, + ) +} + +/// Prefill `tokens` from `start_pos` through the batched program in chunks of +/// [`GEMMA4_FORWARD_BATCH_MAX`] rows, writing both KV caches and computing no +/// logits. A calibration collector, when armed, sees every projection input. +pub fn forward_prefill_batch( + cfg: &Gemma4Config, + weights: &Gemma4Weights, + state: &mut Gemma4State, + gpu: &mut Gpu, + tokens: &[u32], + start_pos: usize, +) -> Result<(), String> { + for (i, chunk) in tokens.chunks(GEMMA4_FORWARD_BATCH_MAX).enumerate() { + let pos = start_pos + i * GEMMA4_FORWARD_BATCH_MAX; + // One row runs the decode route, which reads the staged position. + if let [token] = chunk { + decode_step(cfg, weights, state, gpu, *token, pos as u32)?; + continue; + } + forward_rows(cfg, weights, state, gpu, chunk, pos, None, None, false)?; + } + Ok(()) +} + +/// The batched program over `tokens` at `[start_pos, start_pos + B)`; with +/// `last_logits` the last row's logits are computed, downloaded and returned +/// (else an empty vector). +#[allow(clippy::too_many_arguments)] +fn forward_rows( + cfg: &Gemma4Config, + weights: &Gemma4Weights, + state: &mut Gemma4State, + gpu: &mut Gpu, + tokens: &[u32], + start_pos: usize, + per_token_hidden_out: Option<&GpuTensor>, + per_pos_argmax_out: Option<&mut Vec>, + last_logits: bool, ) -> Result, String> { let b = tokens.len(); if b == 0 { @@ -701,6 +793,9 @@ pub fn forward_batch_spec( let ple_packed = cfg.n_layers * ple_dim; // seq_len after this batch = absolute positions [start_pos, start_pos+B). checked_batch_seq_len(start_pos, b, state.max_seq)?; + if !(state.kv_sliding.quant_q8 && state.kv_full.quant_q8) { + return Err("gemma4 forward_batch: batched attention needs Q8 caches".to_string()); + } let mut bufs = RowBuffers::new(gpu); let widths = RowWidths { @@ -803,7 +898,7 @@ pub fn forward_batch_spec( rows: b, position: start_pos, positions: Some(&acts.positions), - scratch: acts.scratch(ple, None), + scratch: acts.scratch(ple, state.moe.as_ref().map(crate::program::routed_scratch)), }; let mut steps = Vec::with_capacity(4 * cfg.n_layers); for (layer_idx, layer) in weights.layers.iter().enumerate() { @@ -868,25 +963,16 @@ pub fn forward_batch_spec( // (tanh-scaled), so argmax(softcap(z)) == argmax(z). The accept // decision is an argmax, so this is bit-exact in the decision. let logits_b = bufs.alloc(b * vocab, "spec_logits_b")?; - let lm_result = if eagle_strict_enabled() { - gpu.gemm_q8_0_batched( - &weights.lm_head.buf, - normed_hidden, - &logits_b, - weights.lm_head.m, - weights.lm_head.k, - b, - ) - } else { - gpu.gemm_q8_0_batched_chunked( - &weights.lm_head.buf, - normed_hidden, - &logits_b, - weights.lm_head.m, - weights.lm_head.k, - b, - ) - }; + // Not `_chunked`: verify stays on the scalar kernel rather + // than the gfx12 WMMA variant, whose drift accumulates in KV. + let lm_result = gpu.gemm_q8_0_batched( + &weights.lm_head.buf, + normed_hidden, + &logits_b, + weights.lm_head.m, + weights.lm_head.k, + b, + ); lm_result .map_err(|e| format!("gemma4 forward_batch_spec batched lm_head: {e:?}"))?; // GPU per-row argmax over [B, vocab]; only B indices land on PCIe. @@ -930,6 +1016,9 @@ pub fn forward_batch_spec( } } + if !last_logits { + return Ok(Vec::new()); + } // ── Final RMSNorm + tied lm_head on the LAST row only (verify needs the // last position's logits). ── let x_last = bufs.alloc(dim, "x_last")?; diff --git a/crates/hipfire-arch-gemma4/src/gemma4.rs b/crates/hipfire-arch-gemma4/src/gemma4.rs index a289235fa8..b3943f41fa 100644 --- a/crates/hipfire-arch-gemma4/src/gemma4.rs +++ b/crates/hipfire-arch-gemma4/src/gemma4.rs @@ -2,23 +2,125 @@ // Copyright (c) 2026 Kaden Schutt // hipfire — see LICENSE and NOTICE in the project root. -//! Gemma 4 dense (text-only) weights + per-decode state. -//! -//! The dense 12B path remains the default. E2B/E4B add per-layer embeddings, -//! KV sharing, and (for E2B) double-wide MLP tails behind the strict topology -//! contract in [`Gemma4Config`]. MoE and multimodal towers remain out of scope. +//! Gemma 4 weights + per-decode state, one stack for every variant: the +//! dense 12B, E2B/E4B (per-layer embeddings, KV sharing, double-wide MLP +//! tails behind the strict topology contract in [`Gemma4Config`]) and the +//! 26B-A4B MoE (routed experts in parallel with every dense FFN). Multimodal +//! towers remain out of scope. use crate::config::{Gemma4Config, LayerType}; use hipfire_runtime::hfq::HfqFile; use hipfire_runtime::llama::{f16_to_f32, EmbeddingFormat, KvCache, WeightTensor}; use hipfire_runtime::weight_backend::load_embedding; +use hipfire_runtime::weight_store::upload_pooled_bytes; use rdna_compute::{DType, Gpu, GpuTensor}; pub(crate) const GEMMA4_FORWARD_BATCH_MAX: usize = 64; +// ───────────────────────── Allocation ledger ───────────────────────── + +/// Every device allocation of one load, as second owner handles. A failed +/// load frees them all; a successful one forgets them (the published owners +/// hold the bytes). Neither `GpuTensor` nor `DeviceBuffer` has a `Drop`, so a +/// forgotten or dropped handle frees nothing. +pub(crate) struct Ledger<'g> { + pub gpu: &'g mut Gpu, + owners: Vec, + ops: usize, + /// Fail the n-th allocation (rollback tests). + fail_at: Option, +} + +/// A second owner handle of `tensor`'s allocation (a non-owning alias would +/// be refused by `free_tensor`). +fn owner_handle(tensor: &GpuTensor) -> GpuTensor { + GpuTensor { + // SAFETY: `DeviceBuffer` is plain data without `Drop`; exactly one of + // the two handles is ever freed — the ledger's on rollback (the + // partially built owners are then dropped unfreed), else the + // published owner's. + buf: unsafe { std::ptr::read(&tensor.buf) }, + shape: tensor.shape.clone(), + dtype: tensor.dtype, + } +} + +impl<'g> Ledger<'g> { + pub fn new(gpu: &'g mut Gpu, fail_at: Option) -> Self { + Self { + gpu, + owners: Vec::new(), + ops: 0, + fail_at, + } + } + + /// Run one allocation and record its owner. + pub fn alloc( + &mut self, + label: &str, + f: impl FnOnce(&mut Gpu) -> hip_bridge::HipResult, + ) -> Result { + self.ops += 1; + if self.fail_at == Some(self.ops) { + return Err(format!( + "gemma4: injected load fault at {label} (op {})", + self.ops + )); + } + let tensor = f(self.gpu).map_err(|e| format!("gemma4: alloc {label}: {e:?}"))?; + self.owners.push(owner_handle(&tensor)); + Ok(tensor) + } + + /// A zeroed F32 `[n]` tensor. + pub fn zeros(&mut self, n: usize, label: &str) -> Result { + self.alloc(label, |g| g.zeros(&[n], DType::F32)) + } + + /// Build a KV cache and record every tensor it owns. + pub fn kv( + &mut self, + label: &str, + f: impl FnOnce(&mut Gpu) -> hip_bridge::HipResult, + ) -> Result { + self.ops += 1; + if self.fail_at == Some(self.ops) { + return Err(format!( + "gemma4: injected load fault at {label} (op {})", + self.ops + )); + } + let kv = f(self.gpu).map_err(|e| format!("gemma4: {label}: {e:?}"))?; + let tensors = kv + .k_gpu + .iter() + .chain(&kv.v_gpu) + .chain(&kv.k_scales) + .chain(&kv.v_scales) + .chain(&kv.givens_cos) + .chain(&kv.givens_sin); + self.owners.extend(tensors.map(owner_handle)); + Ok(kv) + } + + /// Free every recorded owner and hand the GPU back. + pub fn rollback(self) -> &'g mut Gpu { + for tensor in self.owners.into_iter().rev() { + let _ = self.gpu.free_tensor(tensor); + } + self.gpu + } +} + // ───────────────────────── HFQ load helpers ───────────────────────── -/// Decode a shape-[n] F16/F32 tensor into an F32 host Vec. +fn f32_bytes(values: &[f32]) -> &[u8] { + // SAFETY: f32 has no padding and any byte pattern is a valid u8. + unsafe { std::slice::from_raw_parts(values.as_ptr() as *const u8, values.len() * 4) } +} + +/// Decode a shape-[n] F16/F32/BF16 tensor into an F32 host Vec. fn load_f32_vec(hfq: &HfqFile, name: &str, expected_n: usize) -> Result, String> { let (info, data) = hfq .tensor_data(name) @@ -47,12 +149,12 @@ fn load_f32_vec(hfq: &HfqFile, name: &str, expected_n: usize) -> Result Ok(f32_data) } -/// Load a Gemma 4 RMSNorm weight — plain `x * w` form (NO +1 shift), unless the -/// `norm_plus_one` config toggle is set (`HIPFIRE_GEMMA4_NORM_PLUS_ONE=1`), in -/// which case the Gemma-2/3 `x * (1 + w)` convention is baked at load time. +/// Load a Gemma 4 RMSNorm weight — plain `x * w` form (NO +1 shift), unless +/// `plus_one` (`HIPFIRE_GEMMA4_NORM_PLUS_ONE=1`) bakes the Gemma-2/3 +/// `x * (1 + w)` convention at load time. fn load_norm( + l: &mut Ledger, hfq: &HfqFile, - gpu: &mut Gpu, name: &str, dim: usize, plus_one: bool, @@ -63,13 +165,12 @@ fn load_norm( *v += 1.0; } } - gpu.upload_f32(&f32_data, &[dim]) - .map_err(|e| format!("gemma4: upload norm {name}: {e:?}")) + l.alloc(name, |g| g.upload_f32(&f32_data, &[dim])) } /// Load the learned per-layer `layer_scalar` (shape-[1]); returns the host f32 -/// value so decode can call `gpu.scale_f32(x, scalar)` with no D2H round-trip. -/// Returns `1.0` (no-op) when the tensor is absent — the 12B may not ship it. +/// value so decode can scale with no D2H round-trip. Returns `1.0` (no-op) +/// when the tensor is absent — the 12B may not ship it. fn load_layer_scalar(hfq: &HfqFile, name: &str) -> f32 { match load_f32_vec(hfq, name, 1) { Ok(v) => v[0], @@ -77,91 +178,144 @@ fn load_layer_scalar(hfq: &HfqFile, name: &str) -> f32 { } } -/// quant_type → DType mapping for Gemma 4 projection weights. Mirrors the old -/// branch's `load_gemma4_weight`. F16 is dequantized to F32 on upload. -fn load_wt( - hfq: &HfqFile, - gpu: &mut Gpu, - name: &str, - m: usize, - k: usize, -) -> Result { - let (info, data) = hfq - .tensor_data(name) - .ok_or_else(|| format!("gemma4: tensor not found: {name}"))?; - if info.quant_type == 1 || info.quant_type == 16 { - // F16/BF16 → upload as F32. - let f32_data: Vec = if info.quant_type == 16 { - data.chunks_exact(2) - .map(|c| f32::from_bits((u16::from_le_bytes([c[0], c[1]]) as u32) << 16)) - .collect() - } else { - data.chunks_exact(2) - .map(|c| f16_to_f32(u16::from_le_bytes([c[0], c[1]]))) - .collect() - }; - let bytes: &[u8] = unsafe { - std::slice::from_raw_parts(f32_data.as_ptr() as *const u8, f32_data.len() * 4) - }; - let buf = gpu - .upload_raw(bytes, &[m, k]) - .map_err(|e| format!("gemma4: upload F32 {name}: {e:?}"))?; - return Ok(WeightTensor { - buf, - gpu_dtype: DType::F32, - m, - k, - row_stride: 0, - paro: None, - awq_scale: None, - lloyd_lut_e4m3: None, - lloyd_lut_f16: None, - lloyd_lut_c16: None, - }); - } - let dtype = match info.quant_type { - 2 => { - let buf = gpu - .upload_raw(data, &[m, k]) - .map_err(|e| format!("gemma4: upload F32 {name}: {e:?}"))?; - return Ok(WeightTensor { - buf, - gpu_dtype: DType::F32, - m, - k, - row_stride: 0, - paro: None, - awq_scale: None, - lloyd_lut_e4m3: None, - lloyd_lut_f16: None, - lloyd_lut_c16: None, - }); - } +/// Quantized projection formats. 19 and 30 are Lloyd formats now (the old +/// MG4 aliases), which no Gemma 4 kernel reads, so they refuse. +fn quantized_dtype(quant_type: u8) -> Option { + Some(match quant_type { 3 => DType::Q8_0, 4 => DType::Q4K, 6 => DType::HFQ4G256, 7 => DType::HFQ4G128, 8 => DType::HFQ6G256, 9 => DType::HFQ2G256, + 10 => DType::HFQ2G128, 11 => DType::HFQ3G256, + 12 => DType::HFQ3G128, 13 => DType::MQ4G256, + 14 => DType::MQ8G256, 15 => DType::MQ6G256, 17 => DType::MQ3G256, - // MG4G256 (qt=19) reads back as MQ4G256: identical binary layout. - 19 => DType::MQ4G256, - qt => return Err(format!("gemma4: unsupported quant_type {qt} for {name}")), + 18 => DType::MQ2G256, + 44 => DType::MQ4G256V2, + _ => return None, + }) +} + +/// HFQ4-G128 rows each start a fresh 72-byte group; a trailing partial group +/// is padded. Files quantized before the 2-D packer (`b4846285e`) packed +/// groups across rows whenever `K % 128 != 0` (Gemma 4 expert `down_proj`, +/// K = 704), which every per-row kernel reads as garbage. Refuse them. +fn check_hfq4g128_rows(name: &str, m: usize, k: usize, bytes: usize) -> Result<(), String> { + let rows = m * k.div_ceil(128) * 72; + if bytes == rows { + return Ok(()); + } + let reason = if bytes == (m * k).div_ceil(128) * 72 { + "packs HFQ4-G128 groups across rows (quantized before hipfire-quantize \ + b4846285e); requantize the model" + } else { + "has an unexpected byte size" + }; + Err(format!( + "{name} [{m} x {k}] {reason} ({bytes} bytes, expected {rows})" + )) +} + +/// AWQ per-channel scale sidecar (`.awq_scale.weight`, F16 `[k]`). +/// Absent or malformed sidecars are skipped; a present, well-formed one whose +/// upload fails fails the load instead of silently computing `(W·s)·x`. +fn load_awq_scale( + l: &mut Ledger, + hfq: &HfqFile, + weight_name: &str, + k: usize, +) -> Result, String> { + let sidecar_name = match weight_name.strip_suffix(".weight") { + Some(stem) => format!("{stem}.awq_scale.weight"), + None => format!("{weight_name}.awq_scale.weight"), + }; + let Some((info, data)) = hfq.tensor_data_vec(&sidecar_name) else { + return Ok(None); + }; + if info.quant_type != 1 || info.shape.len() != 1 || info.shape[0] as usize != k { + eprintln!( + "warning: AWQ sidecar {sidecar_name} (qt={}, shape {:?}) is not F16 [{k}]; skipping", + info.quant_type, info.shape + ); + return Ok(None); + } + let scale: Vec = data + .chunks_exact(2) + .map(|c| f16_to_f32(u16::from_le_bytes([c[0], c[1]]))) + .collect(); + let bytes = f32_bytes(&scale); + l.alloc(&sidecar_name, |g| { + upload_pooled_bytes(g, bytes, &[bytes.len()]) + }) + .map(Some) +} + +/// Load one `[m, k]` projection weight (plus its AWQ sidecar). +fn load_wt( + l: &mut Ledger, + hfq: &HfqFile, + name: &str, + m: usize, + k: usize, +) -> Result { + let (info, data) = hfq + .tensor_data(name) + .ok_or_else(|| format!("gemma4: tensor not found: {name}"))?; + let widen = |bf16: bool| -> Vec { + data.chunks_exact(2) + .map(|c| { + let bits = u16::from_le_bytes([c[0], c[1]]); + if bf16 { + f32::from_bits((bits as u32) << 16) + } else { + f16_to_f32(bits) + } + }) + .collect() + }; + let (buf, gpu_dtype) = match info.quant_type { + // BF16 stays BF16 for the calibration teacher's MFMA GEMM. + 16 if rdna_compute::calib_force_bf16() => ( + l.alloc(name, |g| upload_pooled_bytes(g, data, &[m, k]))?, + DType::BF16, + ), + qt @ (1 | 16) => { + let values = widen(qt == 16); + let bytes = f32_bytes(&values); + ( + l.alloc(name, |g| upload_pooled_bytes(g, bytes, &[m, k]))?, + DType::F32, + ) + } + 2 => ( + l.alloc(name, |g| upload_pooled_bytes(g, data, &[m, k]))?, + DType::F32, + ), + qt => { + let dtype = quantized_dtype(qt) + .ok_or_else(|| format!("gemma4: unsupported quant_type {qt} for {name}"))?; + if dtype == DType::HFQ4G128 { + check_hfq4g128_rows(name, m, k, data.len())?; + } + ( + l.alloc(name, |g| upload_pooled_bytes(g, data, &[data.len()]))?, + dtype, + ) + } }; - let buf = gpu - .upload_raw(data, &[data.len()]) - .map_err(|e| format!("gemma4: upload {name}: {e:?}"))?; - let awq_scale = if dtype.supports_awq_sidecar() { - hipfire_runtime::hfq::load_awq_scale(hfq, gpu, name, k) + let awq_scale = if gpu_dtype.supports_awq_sidecar() { + load_awq_scale(l, hfq, name, k)? } else { None }; Ok(WeightTensor { buf, - gpu_dtype: dtype, + gpu_dtype, m, k, row_stride: 0, @@ -173,6 +327,28 @@ fn load_wt( }) } +/// A tied weight view over `table` (no ownership: the table frees the bytes). +fn tied_head(table: &GpuTensor, format: EmbeddingFormat, m: usize, k: usize) -> WeightTensor { + let dtype = hipfire_runtime::weight_backend::embedding_format_dtype(format); + WeightTensor { + buf: GpuTensor { + // SAFETY: `embed_tokens` owns the allocation and outlives the head. + buf: unsafe { table.buf.alias() }, + shape: table.shape.clone(), + dtype, + }, + gpu_dtype: dtype, + m, + k, + row_stride: 0, + paro: None, + awq_scale: None, + lloyd_lut_e4m3: None, + lloyd_lut_f16: None, + lloyd_lut_c16: None, + } +} + // ──────────────────────────── Weights ──────────────────────────── /// Per-layer weights for a SLIDING layer (head_dim 256, full RoPE, own v_proj). @@ -195,13 +371,14 @@ pub struct SlidingLayerWeights { pub down_proj: WeightTensor, pub ffn_hidden_dim: usize, pub per_layer: Option, + pub moe: Option, } /// Per-layer weights for a FULL layer (head_dim 512, partial RoPE). /// -/// `attention_k_eq_v == true` (12B): no `v_proj`; V is the PRE-k_norm output of -/// k_proj, renormed by a weight-less (ones) RMSNorm. When false (other Gemma 4 -/// variants), a separate `v_proj` is loaded. +/// `attention_k_eq_v == true` (12B, 26B): no `v_proj`; V is the PRE-k_norm +/// output of k_proj, renormed by a weight-less (ones) RMSNorm. When false +/// (E-series), a separate `v_proj` is loaded. pub struct FullLayerWeights { pub input_layernorm: GpuTensor, pub post_attention_layernorm: GpuTensor, @@ -221,6 +398,7 @@ pub struct FullLayerWeights { pub down_proj: WeightTensor, pub ffn_hidden_dim: usize, pub per_layer: Option, + pub moe: Option, } pub enum LayerWeights { @@ -228,6 +406,15 @@ pub enum LayerWeights { Full(FullLayerWeights), } +impl LayerWeights { + pub fn moe(&self) -> Option<&MoeLayerExtras> { + match self { + LayerWeights::Sliding(l) => l.moe.as_ref(), + LayerWeights::Full(l) => l.moe.as_ref(), + } + } +} + /// Per-layer residual branch driven by the token-level PLE source. pub struct PerLayerBranchWeights { pub input_gate: WeightTensor, @@ -243,9 +430,146 @@ pub struct PerLayerInputWeights { pub projection_norm: GpuTensor, } +/// The routed-expert branch of one MoE layer (26B-A4B). Each expert kind lives +/// in one pooled allocation (one hipMalloc per kind and layer instead of one +/// per expert, which fragmented the heap); the indexed kernels find expert +/// `e` through the `[n_experts]` u64 device pointer tables. +pub struct MoeLayerExtras { + /// `[n_experts, dim]` router projection. + pub router_proj: WeightTensor, + /// `[dim]` multiplicative router-input scale (`router.scale`). + pub router_scale: GpuTensor, + /// `[n_experts]` per-expert post-`down_proj` scale. + pub per_expert_scale: GpuTensor, + /// `[dim]` RMSNorm on the expert input. + pub pre_feedforward_layernorm_2: GpuTensor, + /// `[dim]` RMSNorm on the dense FFN output before the sum. + pub post_feedforward_layernorm_1: GpuTensor, + /// `[dim]` RMSNorm on the expert output before the sum. + pub post_feedforward_layernorm_2: GpuTensor, + pub n_experts: usize, + pub gate_up_dtype: DType, + pub down_dtype: DType, + pub experts_gate_up_pool: GpuTensor, + pub experts_down_pool: GpuTensor, + pub experts_gate_up_ptrs: GpuTensor, + pub experts_down_ptrs: GpuTensor, +} + +/// Every `{p}.experts.{x}.{kind}.weight` of one layer in a single pooled +/// allocation; returns `(pool, dtype, bytes_per_expert)`. +fn load_expert_pool( + l: &mut Ledger, + hfq: &HfqFile, + p: &str, + kind: &str, + n_experts: usize, +) -> Result<(GpuTensor, DType, usize), String> { + let first_name = format!("{p}.experts.0.{kind}.weight"); + let (first_info, first_data) = hfq + .tensor_data(&first_name) + .ok_or_else(|| format!("gemma4: MoE expert tensor not found: {first_name}"))?; + let bytes_per_expert = first_data.len(); + let dtype = quantized_dtype(first_info.quant_type).ok_or_else(|| { + format!( + "gemma4: unsupported MoE expert quant_type {} for {first_name}", + first_info.quant_type + ) + })?; + if dtype == DType::HFQ4G128 { + let [m, k] = first_info.shape[..] else { + return Err(format!( + "{first_name}: expected a 2-D shape, got {:?}", + first_info.shape + )); + }; + check_hfq4g128_rows(&first_name, m as usize, k as usize, bytes_per_expert)?; + } + let mut concat = Vec::with_capacity(bytes_per_expert * n_experts); + concat.extend_from_slice(first_data); + for x in 1..n_experts { + let name = format!("{p}.experts.{x}.{kind}.weight"); + let (info, data) = hfq + .tensor_data(&name) + .ok_or_else(|| format!("gemma4: MoE expert tensor not found: {name}"))?; + if data.len() != bytes_per_expert || info.quant_type != first_info.quant_type { + return Err(format!( + "gemma4: MoE expert {name} (qt={}, {} bytes) differs from expert 0 (qt={}, {bytes_per_expert} bytes)", + info.quant_type, + data.len(), + first_info.quant_type + )); + } + concat.extend_from_slice(data); + } + let pool = l.alloc(&format!("{p} {kind} pool"), |g| { + upload_pooled_bytes(g, &concat, &[concat.len()]) + })?; + Ok((pool, dtype, bytes_per_expert)) +} + +fn load_moe_layer_extras( + l: &mut Ledger, + hfq: &HfqFile, + cfg: &Gemma4Config, + p: &str, + plus_one: bool, +) -> Result { + let n_exp = cfg.num_experts; + let dim = cfg.dim; + let router_proj = load_wt(l, hfq, &format!("{p}.router.proj.weight"), n_exp, dim)?; + // `router.scale` and `router.per_expert_scale` ship WITHOUT the `.weight` + // suffix in HF's 26B-A4B safetensors. Neither is an RMSNorm weight. + let router_scale = load_norm(l, hfq, &format!("{p}.router.scale"), dim, false)?; + let per_expert = load_f32_vec(hfq, &format!("{p}.router.per_expert_scale"), n_exp)?; + let per_expert_bytes = f32_bytes(&per_expert); + let per_expert_scale = l.alloc("per_expert_scale", |g| { + upload_pooled_bytes(g, per_expert_bytes, &[n_exp]) + })?; + let norm = + |l: &mut Ledger, rel: &str| load_norm(l, hfq, &format!("{p}.{rel}.weight"), dim, plus_one); + let pre_feedforward_layernorm_2 = norm(l, "pre_feedforward_layernorm_2")?; + let post_feedforward_layernorm_1 = norm(l, "post_feedforward_layernorm_1")?; + let post_feedforward_layernorm_2 = norm(l, "post_feedforward_layernorm_2")?; + let (gate_up_pool, gate_up_dtype, gate_up_bpe) = + load_expert_pool(l, hfq, p, "gate_up_proj", n_exp)?; + let (down_pool, down_dtype, down_bpe) = load_expert_pool(l, hfq, p, "down_proj", n_exp)?; + // Pool allocations never move, so the device pointer tables stay valid + // for the model's lifetime. + let ptrs = |pool: &GpuTensor, bpe: usize| -> Vec { + (0..n_exp) + .flat_map(|x| (pool.buf.as_ptr() as u64 + (x * bpe) as u64).to_ne_bytes()) + .collect() + }; + let gate_up_ptr_bytes = ptrs(&gate_up_pool, gate_up_bpe); + let down_ptr_bytes = ptrs(&down_pool, down_bpe); + // Each u64 occupies two f32 slots; the kernels cast the buffer to u64*. + let experts_gate_up_ptrs = l.alloc("gate_up ptrs", |g| { + upload_pooled_bytes(g, &gate_up_ptr_bytes, &[n_exp * 2]) + })?; + let experts_down_ptrs = l.alloc("down ptrs", |g| { + upload_pooled_bytes(g, &down_ptr_bytes, &[n_exp * 2]) + })?; + Ok(MoeLayerExtras { + router_proj, + router_scale, + per_expert_scale, + pre_feedforward_layernorm_2, + post_feedforward_layernorm_1, + post_feedforward_layernorm_2, + n_experts: n_exp, + gate_up_dtype, + down_dtype, + experts_gate_up_pool: gate_up_pool, + experts_down_pool: down_pool, + experts_gate_up_ptrs, + experts_down_ptrs, + }) +} + fn load_per_layer_branch( + l: &mut Ledger, hfq: &HfqFile, - gpu: &mut Gpu, cfg: &Gemma4Config, prefix: &str, plus_one: bool, @@ -256,22 +580,22 @@ fn load_per_layer_branch( } Ok(Some(PerLayerBranchWeights { input_gate: load_wt( + l, hfq, - gpu, &format!("{prefix}.per_layer_input_gate.weight"), ple_dim, cfg.dim, )?, projection: load_wt( + l, hfq, - gpu, &format!("{prefix}.per_layer_projection.weight"), cfg.dim, ple_dim, )?, post_input_norm: load_norm( + l, hfq, - gpu, &format!("{prefix}.post_per_layer_input_norm.weight"), cfg.dim, plus_one, @@ -279,6 +603,96 @@ fn load_per_layer_branch( })) } +fn load_layer( + l: &mut Ledger, + hfq: &HfqFile, + cfg: &Gemma4Config, + name_root: &str, + i: usize, +) -> Result { + let dim = cfg.dim; + let plus_one = cfg.norm_plus_one; + let p = format!("{name_root}.layers.{i}"); + let ffn_hd = cfg.ffn_hidden_dim_for_layer(i); + let (hd, n_kv) = match cfg.layer_types[i] { + LayerType::Sliding => (cfg.sliding_head_dim, cfg.sliding_n_kv_heads), + LayerType::Full => (cfg.full_head_dim, cfg.full_n_kv_heads), + }; + let kv_dim = n_kv * hd; + let q_dim = cfg.n_heads * hd; + let norm = |l: &mut Ledger, rel: &str, n: usize| { + load_norm(l, hfq, &format!("{p}.{rel}.weight"), n, plus_one) + }; + let wt = |l: &mut Ledger, rel: &str, m: usize, k: usize| { + load_wt(l, hfq, &format!("{p}.{rel}.weight"), m, k) + }; + let input_layernorm = norm(l, "input_layernorm", dim)?; + let post_attention_layernorm = norm(l, "post_attention_layernorm", dim)?; + let pre_feedforward_layernorm = norm(l, "pre_feedforward_layernorm", dim)?; + let post_feedforward_layernorm = norm(l, "post_feedforward_layernorm", dim)?; + let layer_scalar_host = load_layer_scalar(hfq, &format!("{p}.layer_scalar")); + let q_proj = wt(l, "self_attn.q_proj", q_dim, dim)?; + let k_proj = wt(l, "self_attn.k_proj", kv_dim, dim)?; + let has_v = cfg.layer_types[i] == LayerType::Sliding || !cfg.attention_k_eq_v; + let v_proj = if has_v { + Some(wt(l, "self_attn.v_proj", kv_dim, dim)?) + } else { + None + }; + let o_proj = wt(l, "self_attn.o_proj", dim, q_dim)?; + let q_norm = norm(l, "self_attn.q_norm", hd)?; + let k_norm = norm(l, "self_attn.k_norm", hd)?; + let gate_proj = wt(l, "mlp.gate_proj", ffn_hd, dim)?; + let up_proj = wt(l, "mlp.up_proj", ffn_hd, dim)?; + let down_proj = wt(l, "mlp.down_proj", dim, ffn_hd)?; + let per_layer = load_per_layer_branch(l, hfq, cfg, &p, plus_one)?; + let moe = if cfg.enable_moe_block { + Some(load_moe_layer_extras(l, hfq, cfg, &p, plus_one)?) + } else { + None + }; + Ok(match (cfg.layer_types[i], v_proj) { + (LayerType::Sliding, Some(v_proj)) => LayerWeights::Sliding(SlidingLayerWeights { + input_layernorm, + post_attention_layernorm, + pre_feedforward_layernorm, + post_feedforward_layernorm, + layer_scalar_host, + q_proj, + k_proj, + v_proj, + o_proj, + q_norm, + k_norm, + gate_proj, + up_proj, + down_proj, + ffn_hidden_dim: ffn_hd, + per_layer, + moe, + }), + (_, v_proj) => LayerWeights::Full(FullLayerWeights { + input_layernorm, + post_attention_layernorm, + pre_feedforward_layernorm, + post_feedforward_layernorm, + layer_scalar_host, + q_proj, + k_proj, + v_proj, + o_proj, + q_norm, + k_norm, + gate_proj, + up_proj, + down_proj, + ffn_hidden_dim: ffn_hd, + per_layer, + moe, + }), + }) +} + pub struct Gemma4Weights { /// Token embedding [vocab, dim]; aliased as lm_head when tied. pub embed_tokens: GpuTensor, @@ -292,16 +706,35 @@ pub struct Gemma4Weights { impl Gemma4Weights { pub fn load(hfq: &HfqFile, cfg: &Gemma4Config, gpu: &mut Gpu) -> Result { + Self::load_with_fault(hfq, cfg, gpu, None) + } + + /// [`Self::load`] that fails its `fail_at`-th allocation: every staged + /// owner, completed layers and AWQ sidecars included, is freed. + pub(crate) fn load_with_fault( + hfq: &HfqFile, + cfg: &Gemma4Config, + gpu: &mut Gpu, + fail_at: Option, + ) -> Result { if cfg.hidden_size_per_layer_input != 0 || cfg.num_kv_shared_layers != 0 { cfg.e_series_variant()?; } + let mut l = Ledger::new(gpu, fail_at); + match Self::load_into(&mut l, hfq, cfg) { + Ok(weights) => Ok(weights), + Err(e) => { + l.rollback(); + Err(e) + } + } + } + + fn load_into(l: &mut Ledger, hfq: &HfqFile, cfg: &Gemma4Config) -> Result { let dim = cfg.dim; let plus_one = cfg.norm_plus_one; - - let name_roots = ["model.language_model", "language_model", "model"]; - let name_root = name_roots - .iter() - .copied() + let name_root = ["model.language_model", "language_model", "model"] + .into_iter() .find(|root| { hfq.tensor_data(&format!("{root}.embed_tokens.weight")) .is_some() @@ -311,94 +744,55 @@ impl Gemma4Weights { .to_string() })?; - // ── Embedding ────────────────────────────────────────────────────── - let embed_name = format!("{name_root}.embed_tokens.weight"); let (embed_info, embed_data) = hfq - .tensor_data(&embed_name) + .tensor_data(&format!("{name_root}.embed_tokens.weight")) .ok_or_else(|| "gemma4: embed_tokens not found in HFQ".to_string())?; - let (embed_tokens, embd_format) = match embed_info.quant_type { - 3 => ( - gpu.upload_raw(embed_data, &[embed_data.len()]) - .map_err(|e| format!("gemma4: upload embed: {e:?}"))?, - EmbeddingFormat::Q8_0, - ), - 6 => ( - gpu.upload_raw(embed_data, &[embed_data.len()]) - .map_err(|e| format!("gemma4: upload embed: {e:?}"))?, - EmbeddingFormat::HFQ4G256, - ), - 7 => ( - gpu.upload_raw(embed_data, &[embed_data.len()]) - .map_err(|e| format!("gemma4: upload embed: {e:?}"))?, - EmbeddingFormat::HFQ4G128, - ), - 1 | 2 | 16 => { - load_embedding(gpu, embed_info.quant_type, embed_data, cfg.vocab_size, dim) - .map_err(|e| format!("gemma4: upload embed: {e:?}"))? - } - qt => return Err(format!("gemma4: unsupported embed quant_type {qt}")), - }; - - // Tied lm_head: WeightTensor aliasing the embed allocation. free path - // skips it (embed owns the bytes). - let lm_head = { - let alias_buf = unsafe { embed_tokens.buf.alias() }; - let dtype = match embd_format { - EmbeddingFormat::Q8_0 => DType::Q8_0, - EmbeddingFormat::HFQ4G256 => DType::HFQ4G256, - EmbeddingFormat::HFQ4G128 => DType::HFQ4G128, - EmbeddingFormat::F32 => DType::F32, - EmbeddingFormat::Q4K => DType::Q4K, - }; - let alias_tensor = GpuTensor { - buf: alias_buf, - shape: embed_tokens.shape.clone(), - dtype, - }; - WeightTensor { - buf: alias_tensor, - gpu_dtype: dtype, - m: cfg.vocab_size, - k: dim, - row_stride: 0, - paro: None, - awq_scale: None, - lloyd_lut_e4m3: None, - lloyd_lut_f16: None, - lloyd_lut_c16: None, - } - }; - - let final_norm = load_norm(hfq, gpu, &format!("{name_root}.norm.weight"), dim, plus_one)?; + let mut embd_format = EmbeddingFormat::F32; + let embed_tokens = l.alloc("embed_tokens", |g| { + load_embedding(g, embed_info.quant_type, embed_data, cfg.vocab_size, dim).map( + |(table, format)| { + embd_format = format; + table + }, + ) + })?; + let lm_head = tied_head(&embed_tokens, embd_format, cfg.vocab_size, dim); + let final_norm = load_norm(l, hfq, &format!("{name_root}.norm.weight"), dim, plus_one)?; let per_layer_input = if cfg.hidden_size_per_layer_input != 0 { let ple_dim = cfg.hidden_size_per_layer_input; let packed_dim = cfg.n_layers * ple_dim; - let ple_embed_name = format!("{name_root}.embed_tokens_per_layer.weight"); + let name = format!("{name_root}.embed_tokens_per_layer.weight"); let (info, data) = hfq - .tensor_data(&ple_embed_name) - .ok_or_else(|| format!("gemma4: tensor not found: {ple_embed_name}"))?; - let (embed_tokens, embd_format) = load_embedding( - gpu, - info.quant_type, - data, - cfg.vocab_size_per_layer_input, - packed_dim, - ) - .map_err(|e| format!("gemma4: load embed_tokens_per_layer: {e:?}"))?; + .tensor_data(&name) + .ok_or_else(|| format!("gemma4: tensor not found: {name}"))?; + let mut ple_format = EmbeddingFormat::F32; + let embed_tokens = l.alloc(&name, |g| { + load_embedding( + g, + info.quant_type, + data, + cfg.vocab_size_per_layer_input, + packed_dim, + ) + .map(|(table, format)| { + ple_format = format; + table + }) + })?; Some(PerLayerInputWeights { embed_tokens, - embd_format, + embd_format: ple_format, model_projection: load_wt( + l, hfq, - gpu, &format!("{name_root}.per_layer_model_projection.weight"), packed_dim, dim, )?, projection_norm: load_norm( + l, hfq, - gpu, &format!("{name_root}.per_layer_projection_norm.weight"), ple_dim, plus_one, @@ -408,223 +802,9 @@ impl Gemma4Weights { None }; - // ── Layers ───────────────────────────────────────────────────────── - let mut layers = Vec::with_capacity(cfg.n_layers); - for i in 0..cfg.n_layers { - let p = format!("{name_root}.layers.{i}"); - let layer_scalar_host = load_layer_scalar(hfq, &format!("{p}.layer_scalar")); - let ffn_hd = cfg.ffn_hidden_dim_for_layer(i); - match cfg.layer_types[i] { - LayerType::Sliding => { - let hd = cfg.sliding_head_dim; - let kv_dim = cfg.sliding_n_kv_heads * hd; - let q_dim = cfg.n_heads * hd; - layers.push(LayerWeights::Sliding(SlidingLayerWeights { - input_layernorm: load_norm( - hfq, - gpu, - &format!("{p}.input_layernorm.weight"), - dim, - plus_one, - )?, - post_attention_layernorm: load_norm( - hfq, - gpu, - &format!("{p}.post_attention_layernorm.weight"), - dim, - plus_one, - )?, - pre_feedforward_layernorm: load_norm( - hfq, - gpu, - &format!("{p}.pre_feedforward_layernorm.weight"), - dim, - plus_one, - )?, - post_feedforward_layernorm: load_norm( - hfq, - gpu, - &format!("{p}.post_feedforward_layernorm.weight"), - dim, - plus_one, - )?, - layer_scalar_host, - q_proj: load_wt( - hfq, - gpu, - &format!("{p}.self_attn.q_proj.weight"), - q_dim, - dim, - )?, - k_proj: load_wt( - hfq, - gpu, - &format!("{p}.self_attn.k_proj.weight"), - kv_dim, - dim, - )?, - v_proj: load_wt( - hfq, - gpu, - &format!("{p}.self_attn.v_proj.weight"), - kv_dim, - dim, - )?, - o_proj: load_wt( - hfq, - gpu, - &format!("{p}.self_attn.o_proj.weight"), - dim, - q_dim, - )?, - q_norm: load_norm( - hfq, - gpu, - &format!("{p}.self_attn.q_norm.weight"), - hd, - plus_one, - )?, - k_norm: load_norm( - hfq, - gpu, - &format!("{p}.self_attn.k_norm.weight"), - hd, - plus_one, - )?, - gate_proj: load_wt( - hfq, - gpu, - &format!("{p}.mlp.gate_proj.weight"), - ffn_hd, - dim, - )?, - up_proj: load_wt( - hfq, - gpu, - &format!("{p}.mlp.up_proj.weight"), - ffn_hd, - dim, - )?, - down_proj: load_wt( - hfq, - gpu, - &format!("{p}.mlp.down_proj.weight"), - dim, - ffn_hd, - )?, - ffn_hidden_dim: ffn_hd, - per_layer: load_per_layer_branch(hfq, gpu, cfg, &p, plus_one)?, - })); - } - LayerType::Full => { - let hd = cfg.full_head_dim; - let kv_dim = cfg.full_n_kv_heads * hd; - let q_dim = cfg.n_heads * hd; - let v_proj = if cfg.attention_k_eq_v { - None - } else { - Some(load_wt( - hfq, - gpu, - &format!("{p}.self_attn.v_proj.weight"), - kv_dim, - dim, - )?) - }; - layers.push(LayerWeights::Full(FullLayerWeights { - input_layernorm: load_norm( - hfq, - gpu, - &format!("{p}.input_layernorm.weight"), - dim, - plus_one, - )?, - post_attention_layernorm: load_norm( - hfq, - gpu, - &format!("{p}.post_attention_layernorm.weight"), - dim, - plus_one, - )?, - pre_feedforward_layernorm: load_norm( - hfq, - gpu, - &format!("{p}.pre_feedforward_layernorm.weight"), - dim, - plus_one, - )?, - post_feedforward_layernorm: load_norm( - hfq, - gpu, - &format!("{p}.post_feedforward_layernorm.weight"), - dim, - plus_one, - )?, - layer_scalar_host, - q_proj: load_wt( - hfq, - gpu, - &format!("{p}.self_attn.q_proj.weight"), - q_dim, - dim, - )?, - k_proj: load_wt( - hfq, - gpu, - &format!("{p}.self_attn.k_proj.weight"), - kv_dim, - dim, - )?, - v_proj, - o_proj: load_wt( - hfq, - gpu, - &format!("{p}.self_attn.o_proj.weight"), - dim, - q_dim, - )?, - q_norm: load_norm( - hfq, - gpu, - &format!("{p}.self_attn.q_norm.weight"), - hd, - plus_one, - )?, - k_norm: load_norm( - hfq, - gpu, - &format!("{p}.self_attn.k_norm.weight"), - hd, - plus_one, - )?, - gate_proj: load_wt( - hfq, - gpu, - &format!("{p}.mlp.gate_proj.weight"), - ffn_hd, - dim, - )?, - up_proj: load_wt( - hfq, - gpu, - &format!("{p}.mlp.up_proj.weight"), - ffn_hd, - dim, - )?, - down_proj: load_wt( - hfq, - gpu, - &format!("{p}.mlp.down_proj.weight"), - dim, - ffn_hd, - )?, - ffn_hidden_dim: ffn_hd, - per_layer: load_per_layer_branch(hfq, gpu, cfg, &p, plus_one)?, - })); - } - } - } - + let layers = (0..cfg.n_layers) + .map(|i| load_layer(l, hfq, cfg, name_root, i)) + .collect::, _>>()?; Ok(Gemma4Weights { embed_tokens, embd_format, @@ -635,13 +815,10 @@ impl Gemma4Weights { }) } - /// Return all GPU weight buffers to the pool (drained on unload by the - /// daemon's `unload_model`). Consumes self. + /// Return all GPU weight buffers to the pool. Consumes self. /// - /// `lm_head` is NOT freed here: it is always an alias of `embed_tokens`' - /// allocation (tied embeddings — see the alias construction in `load`). - /// `embed_tokens` owns the bytes and is freed exactly once below; dropping - /// the alias is a no-op (`DeviceBuffer` has no `Drop`). + /// `lm_head` is NOT freed here: it aliases `embed_tokens`' allocation + /// (tied embeddings), which is freed exactly once below. pub fn free_gpu(self, gpu: &mut Gpu) { let _ = gpu.free_tensor(self.embed_tokens); let _ = gpu.free_tensor(self.final_norm); @@ -650,49 +827,84 @@ impl Gemma4Weights { ple.model_projection.free_all(gpu); let _ = gpu.free_tensor(ple.projection_norm); } + let free_branch = |gpu: &mut Gpu, ple: Option| { + if let Some(ple) = ple { + ple.input_gate.free_all(gpu); + ple.projection.free_all(gpu); + let _ = gpu.free_tensor(ple.post_input_norm); + } + }; + let free_moe = |gpu: &mut Gpu, moe: Option| { + if let Some(moe) = moe { + moe.router_proj.free_all(gpu); + for t in [ + moe.router_scale, + moe.per_expert_scale, + moe.pre_feedforward_layernorm_2, + moe.post_feedforward_layernorm_1, + moe.post_feedforward_layernorm_2, + moe.experts_gate_up_pool, + moe.experts_down_pool, + moe.experts_gate_up_ptrs, + moe.experts_down_ptrs, + ] { + let _ = gpu.free_tensor(t); + } + } + }; for l in self.layers { match l { LayerWeights::Sliding(l) => { - let _ = gpu.free_tensor(l.input_layernorm); - let _ = gpu.free_tensor(l.post_attention_layernorm); - let _ = gpu.free_tensor(l.pre_feedforward_layernorm); - let _ = gpu.free_tensor(l.post_feedforward_layernorm); - l.q_proj.free_all(gpu); - l.k_proj.free_all(gpu); - l.v_proj.free_all(gpu); - l.o_proj.free_all(gpu); - let _ = gpu.free_tensor(l.q_norm); - let _ = gpu.free_tensor(l.k_norm); - l.gate_proj.free_all(gpu); - l.up_proj.free_all(gpu); - l.down_proj.free_all(gpu); - if let Some(ple) = l.per_layer { - ple.input_gate.free_all(gpu); - ple.projection.free_all(gpu); - let _ = gpu.free_tensor(ple.post_input_norm); + for t in [ + l.input_layernorm, + l.post_attention_layernorm, + l.pre_feedforward_layernorm, + l.post_feedforward_layernorm, + l.q_norm, + l.k_norm, + ] { + let _ = gpu.free_tensor(t); + } + for w in [ + l.q_proj, + l.k_proj, + l.v_proj, + l.o_proj, + l.gate_proj, + l.up_proj, + l.down_proj, + ] { + w.free_all(gpu); } + free_branch(gpu, l.per_layer); + free_moe(gpu, l.moe); } LayerWeights::Full(l) => { - let _ = gpu.free_tensor(l.input_layernorm); - let _ = gpu.free_tensor(l.post_attention_layernorm); - let _ = gpu.free_tensor(l.pre_feedforward_layernorm); - let _ = gpu.free_tensor(l.post_feedforward_layernorm); - l.q_proj.free_all(gpu); - l.k_proj.free_all(gpu); - if let Some(v) = l.v_proj { - v.free_all(gpu); + for t in [ + l.input_layernorm, + l.post_attention_layernorm, + l.pre_feedforward_layernorm, + l.post_feedforward_layernorm, + l.q_norm, + l.k_norm, + ] { + let _ = gpu.free_tensor(t); } - l.o_proj.free_all(gpu); - let _ = gpu.free_tensor(l.q_norm); - let _ = gpu.free_tensor(l.k_norm); - l.gate_proj.free_all(gpu); - l.up_proj.free_all(gpu); - l.down_proj.free_all(gpu); - if let Some(ple) = l.per_layer { - ple.input_gate.free_all(gpu); - ple.projection.free_all(gpu); - let _ = gpu.free_tensor(ple.post_input_norm); + for w in [ + l.q_proj, + l.k_proj, + l.o_proj, + l.gate_proj, + l.up_proj, + l.down_proj, + ] + .into_iter() + .chain(l.v_proj) + { + w.free_all(gpu); } + free_branch(gpu, l.per_layer); + free_moe(gpu, l.moe); } } } @@ -701,20 +913,86 @@ impl Gemma4Weights { // ──────────────────────────── State ──────────────────────────── +/// Storage of the full-attention KV tier. The sliding tier is always Q8. +#[derive(Clone, Copy, Debug, PartialEq, Eq)] +pub enum FullKvTier { + Q8, + /// FWHT-512 3-bit K + Q8_0 V (developer examples only). + Fwht3, + /// Deprecated Givens asym3 (`kv_cache=legacy-asym3`). + LegacyAsym3, +} + +/// Single-row routed-expert scratch of a MoE checkpoint; batched forwards +/// route row by row through it. +pub struct MoeScratch { + pub cur_mlp: GpuTensor, // [dim] post_norm_1(dense FFN out) + pub pre2: GpuTensor, // [dim] expert input + pub pre2_rot: GpuTensor, // [dim] FWHT-rotated expert input + pub router_in: GpuTensor, // [dim] + pub router_logits: GpuTensor, // [n_experts] + pub topk_indices: GpuTensor, // [top_k] i32 in f32 slots + pub topk_weights: GpuTensor, // [top_k] + pub cur_moe: GpuTensor, // [dim] expert accumulator + pub gate_batch: GpuTensor, // [top_k * mi] + pub up_batch: GpuTensor, // [top_k * mi] + pub hidden_batch: GpuTensor, // [top_k * mi] +} + +impl MoeScratch { + fn new(l: &mut Ledger, cfg: &Gemma4Config) -> Result { + let dim = cfg.dim; + let k_mi = cfg.top_k_experts * cfg.moe_intermediate_size; + Ok(Self { + cur_mlp: l.zeros(dim, "moe cur_mlp")?, + pre2: l.zeros(dim, "moe pre2")?, + pre2_rot: l.zeros(dim, "moe pre2_rot")?, + router_in: l.zeros(dim, "moe router_in")?, + router_logits: l.zeros(cfg.num_experts, "moe router_logits")?, + topk_indices: l.zeros(cfg.top_k_experts, "moe topk_indices")?, + topk_weights: l.zeros(cfg.top_k_experts, "moe topk_weights")?, + cur_moe: l.zeros(dim, "moe cur_moe")?, + gate_batch: l.zeros(k_mi, "moe gate_batch")?, + up_batch: l.zeros(k_mi, "moe up_batch")?, + hidden_batch: l.zeros(k_mi, "moe hidden_batch")?, + }) + } + + fn tensors(self) -> [GpuTensor; 11] { + [ + self.cur_mlp, + self.pre2, + self.pre2_rot, + self.router_in, + self.router_logits, + self.topk_indices, + self.topk_weights, + self.cur_moe, + self.gate_batch, + self.up_batch, + self.hidden_batch, + ] + } +} + /// Per-decode GPU scratch + the two KV caches (sliding + full). /// /// Scratch is sized once against the MAX of sliding/full attention dims so a /// single buffer set works across both layer types. `kv_slot_for_layer` maps -/// each layer ordinal to its slot within the per-type cache. +/// each layer ordinal to its slot within the per-type cache. Both caches are +/// position-indexed (`physical_cap == max_seq`, no ring): a row below the +/// current position is never rewritten, which speculative verify and session +/// snapshots rely on. pub struct Gemma4State { - /// Sliding-window KV cache (head_dim 256), one slot per sliding layer. + /// Sliding-window KV cache (Q8), one slot per sliding layer. pub kv_sliding: KvCache, - /// Full-attention KV cache (head_dim 512), one slot per full layer. + /// Full-attention KV cache, one slot per full layer. pub kv_full: KvCache, /// Per-layer slot index into the matching per-type cache. pub kv_slot_for_layer: Vec, - pub pos_buf: hip_bridge::DeviceBuffer, // device i32 position scalar + /// Device i32 position scalar. + pub pos_buf: GpuTensor, /// Stable host source for the device position scalar. The hipGraph decode /// path captures a `memcpy_htod_auto` from these bytes; the captured node /// re-reads this heap-stable `Box` on every replay (see @@ -732,8 +1010,7 @@ pub struct Gemma4State { pub x: GpuTensor, // [dim] pub residual: GpuTensor, // [dim] pub tmp: GpuTensor, // [dim] norm / o_proj scratch - /// FWHT-rotated rmsnorm output for the fused MQ4 FFN path - /// (`fused_rmsnorm_rotate_mq` → `fused_gate_up_hfq4g256`). [dim] + /// FWHT-rotated rmsnorm output for MagnumQuant projections. [max(dim, q)] pub tmp_rot: GpuTensor, // attention scratch (sized to max over layer types) @@ -741,7 +1018,7 @@ pub struct Gemma4State { pub k: GpuTensor, // [max_kv_dim] pub v: GpuTensor, // [max_kv_dim] pub attn_out: GpuTensor, // [max_q_dim] - /// Shared tile/reduce workspace for eager and batched Q8 attention. Sized + /// Shared tile/reduce workspace for eager and batched attention. Sized /// for the larger attention geometry and the admitted forward batch cap. pub q8_flash_partials: GpuTensor, @@ -754,13 +1031,16 @@ pub struct Gemma4State { pub ffn_hidden: GpuTensor, // [hidden_dim] pub ffn_out: GpuTensor, // [dim] - // Per-layer input (PLE) scratch. Dense 12B leaves these absent. + // Per-layer input (PLE) scratch. Dense and MoE checkpoints leave these absent. pub ple_token_inputs: Option, // [n_layers * ple_dim] pub ple_projection_all: Option, // [n_layers * ple_dim] pub ple_gate: Option, // [ple_dim] pub ple_hidden: Option, // [ple_dim] pub ple_out: Option, // [dim] + /// Routed-expert scratch; MoE checkpoints only. + pub moe: Option, + // head pub logits: GpuTensor, // [vocab] } @@ -799,13 +1079,10 @@ fn build_kv_slot_map(cfg: &Gemma4Config) -> Result, String> { Ok(slots) } -fn q8_flash_partials_len(gpu: &Gpu, cfg: &Gemma4Config, max_seq: usize) -> usize { - crate::program::q8_flash_partials_len( - gpu, - &crate::program::Geometry::eager(cfg), - max_seq, - GEMMA4_FORWARD_BATCH_MAX, - ) +/// Single-query asym3 hd512 flash partials: `n_heads * ceil(max_seq/128) * +/// (2 + head_dim)` floats at the asym3 kernel's fixed 128-row tile. +fn asym3_flash_partials_len(cfg: &Gemma4Config, max_seq: usize) -> usize { + cfg.n_heads * max_seq.div_ceil(128) * (2 + cfg.full_head_dim) } impl Gemma4State { @@ -820,176 +1097,115 @@ impl Gemma4State { cfg: &Gemma4Config, max_seq: usize, ) -> Result { - let dim = cfg.dim; + Self::new_with_full_tier(gpu, cfg, max_seq, FullKvTier::Q8) + } + pub fn new_with_fwht3_max_seq( + gpu: &mut Gpu, + cfg: &Gemma4Config, + max_seq: usize, + ) -> Result { + Self::new_with_full_tier(gpu, cfg, max_seq, FullKvTier::Fwht3) + } + + pub fn new_with_full_tier( + gpu: &mut Gpu, + cfg: &Gemma4Config, + max_seq: usize, + tier: FullKvTier, + ) -> Result { // FWHT sign LUT must exist before any fused_rmsnorm_rotate_mq / // fused_gate_up_hfq4g256 launch (the MQ4 fused FFN path). gpu.ensure_mq_signs() .map_err(|e| format!("gemma4: ensure_mq_signs: {e:?}"))?; - - // Two Q8 KV caches: one slot per layer of the matching type. - let kv_sliding = KvCache::new_gpu_q8( - gpu, - cfg.n_sliding_kv_slots(), - cfg.sliding_n_kv_heads, - cfg.sliding_head_dim, - max_seq, - ) - .map_err(|e| format!("gemma4: sliding kv cache: {e:?}"))?; - let kv_full = KvCache::new_gpu_q8( - gpu, - cfg.n_full_kv_slots(), - cfg.full_n_kv_heads, - cfg.full_head_dim, - max_seq, - ) - .map_err(|e| format!("gemma4: full kv cache: {e:?}"))?; - let kv_slot_for_layer = build_kv_slot_map(cfg)?; - - let pos_buf = gpu - .hip - .malloc(4) - .map_err(|e| format!("gemma4: pos_buf malloc: {e:?}"))?; - - let alloc = |g: &mut Gpu, n: usize, label: &str| -> Result { - g.zeros(&[n], DType::F32) - .map_err(|e| format!("gemma4: alloc {label}: {e:?}")) - }; - - let max_q = cfg.max_q_dim(); - let max_kv = cfg.max_kv_dim(); - let max_hd = cfg.max_head_dim(); - let ple_dim = cfg.hidden_size_per_layer_input; - let ple_packed = cfg.n_layers * ple_dim; - let q8_flash_partials_len = q8_flash_partials_len(gpu, cfg, max_seq); - - // Ones-filled weight buffer for the weight-less V RMSNorm. - let v_norm_ones = alloc(gpu, max_hd, "v_norm_ones")?; - { - let ones: Vec = vec![1.0; max_hd]; - let bytes: &[u8] = - unsafe { std::slice::from_raw_parts(ones.as_ptr() as *const u8, ones.len() * 4) }; - gpu.hip - .memcpy_htod(&v_norm_ones.buf, bytes) - .map_err(|e| format!("gemma4: init v_norm_ones: {e:?}"))?; + let mut l = Ledger::new(gpu, None); + match Self::alloc(&mut l, cfg, max_seq, tier, kv_slot_for_layer) { + Ok(state) => Ok(state), + Err(e) => { + l.rollback(); + Err(e) + } } - - Ok(Gemma4State { - kv_sliding, - kv_full, - kv_slot_for_layer, - pos_buf, - pos_host: vec![0i32; 1].into_boxed_slice(), - max_seq, - n_tokens: 0, - ar_warmed_up: false, - x: alloc(gpu, dim, "x")?, - residual: alloc(gpu, dim, "residual")?, - tmp: alloc(gpu, dim, "tmp")?, - tmp_rot: alloc(gpu, dim, "tmp_rot")?, - q: alloc(gpu, max_q, "q")?, - k: alloc(gpu, max_kv, "k")?, - v: alloc(gpu, max_kv, "v")?, - attn_out: alloc(gpu, max_q, "attn_out")?, - q8_flash_partials: alloc(gpu, q8_flash_partials_len, "q8_flash_partials")?, - v_norm_ones, - gate_ffn: alloc(gpu, cfg.max_ffn_hidden_dim(), "gate_ffn")?, - up_ffn: alloc(gpu, cfg.max_ffn_hidden_dim(), "up_ffn")?, - ffn_hidden: alloc(gpu, cfg.max_ffn_hidden_dim(), "ffn_hidden")?, - ffn_out: alloc(gpu, dim, "ffn_out")?, - ple_token_inputs: if ple_dim != 0 { - Some(alloc(gpu, ple_packed, "ple_token_inputs")?) - } else { - None - }, - ple_projection_all: if ple_dim != 0 { - Some(alloc(gpu, ple_packed, "ple_projection_all")?) - } else { - None - }, - ple_gate: if ple_dim != 0 { - Some(alloc(gpu, ple_dim, "ple_gate")?) - } else { - None - }, - ple_hidden: if ple_dim != 0 { - Some(alloc(gpu, ple_dim, "ple_hidden")?) - } else { - None - }, - ple_out: if ple_dim != 0 { - Some(alloc(gpu, dim, "ple_out")?) - } else { - None - }, - logits: alloc(gpu, cfg.vocab_size, "logits")?, - }) } - /// Like [] but replaces the full-layer KV cache with - /// FWHT-512 3-bit K + Q8_0 V (fwht3). The sliding-layer KV cache stays - /// Q8_0 (sliding layers use head_dim=256, which has the same byte layout). - pub fn new_with_fwht3_max_seq( - gpu: &mut Gpu, + fn alloc( + l: &mut Ledger, cfg: &Gemma4Config, max_seq: usize, + tier: FullKvTier, + kv_slot_for_layer: Vec, ) -> Result { + let kv_sliding = l.kv("sliding kv cache", |g| { + KvCache::new_gpu_q8( + g, + cfg.n_sliding_kv_slots(), + cfg.sliding_n_kv_heads, + cfg.sliding_head_dim, + max_seq, + ) + })?; + let n_full = cfg.n_full_kv_slots(); + let (heads, hd) = (cfg.full_n_kv_heads, cfg.full_head_dim); + let kv_full = l.kv("full kv cache", |g| match tier { + FullKvTier::Q8 => KvCache::new_gpu_q8(g, n_full, heads, hd, max_seq), + FullKvTier::Fwht3 => KvCache::new_gpu_fwht3_capped_filtered_gemma4( + g, + &vec![true; n_full], + heads, + hd, + max_seq, + max_seq, + ), + FullKvTier::LegacyAsym3 => KvCache::new_gpu_asym3_gemma4(g, n_full, heads, hd, max_seq), + })?; let dim = cfg.dim; - gpu.ensure_mq_signs() - .map_err(|e| format!("gemma4-fwht3: ensure_mq_signs: {e:?}"))?; - - // Sliding layers: Q8 (hd=256, ring-capped to sliding_window). - let kv_sliding = KvCache::new_gpu_q8( - gpu, - cfg.n_sliding_kv_slots(), - cfg.sliding_n_kv_heads, - cfg.sliding_head_dim, - max_seq, - ) - .map_err(|e| format!("gemma4-fwht3: sliding kv cache: {e:?}"))?; - - // Full layers: fwht3 hd=512 (K = FWHT-512 3-bit, V = Q8_0). - let all_true: Vec = vec![true; cfg.n_full_kv_slots()]; - let kv_full = KvCache::new_gpu_fwht3_capped_filtered_gemma4( - gpu, - &all_true, - cfg.full_n_kv_heads, - cfg.full_head_dim, - max_seq, - max_seq, - ) - .map_err(|e| format!("gemma4-fwht3: full kv cache: {e:?}"))?; - - let kv_slot_for_layer = build_kv_slot_map(cfg)?; - - let pos_buf = gpu - .hip - .malloc(4) - .map_err(|e| format!("gemma4-fwht3: pos_buf malloc: {e:?}"))?; - - let alloc = |g: &mut Gpu, n: usize, label: &str| -> Result { - g.zeros(&[n], DType::F32) - .map_err(|e| format!("gemma4-fwht3: alloc {label}: {e:?}")) - }; - let max_q = cfg.max_q_dim(); let max_kv = cfg.max_kv_dim(); let max_hd = cfg.max_head_dim(); let ple_dim = cfg.hidden_size_per_layer_input; let ple_packed = cfg.n_layers * ple_dim; - let q8_flash_partials_len = q8_flash_partials_len(gpu, cfg, max_seq); - - let v_norm_ones = alloc(gpu, max_hd, "v_norm_ones")?; - { - let ones: Vec = vec![1.0; max_hd]; - let bytes: &[u8] = - unsafe { std::slice::from_raw_parts(ones.as_ptr() as *const u8, ones.len() * 4) }; - gpu.hip - .memcpy_htod(&v_norm_ones.buf, bytes) - .map_err(|e| format!("gemma4-fwht3: init v_norm_ones: {e:?}"))?; + let mut partials = crate::program::q8_flash_partials_len( + l.gpu, + &crate::program::Geometry::eager(cfg), + max_seq, + GEMMA4_FORWARD_BATCH_MAX, + ); + if tier != FullKvTier::Q8 { + partials = partials.max(asym3_flash_partials_len(cfg, max_seq)); } + let pos_buf = l.alloc("pos_buf", |g| g.zeros(&[1], DType::F32))?; + let v_norm_ones = l.alloc("v_norm_ones", |g| { + g.upload_f32(&vec![1.0; max_hd], &[max_hd]) + })?; + let x = l.zeros(dim, "x")?; + let residual = l.zeros(dim, "residual")?; + let tmp = l.zeros(dim, "tmp")?; + let tmp_rot = l.zeros(dim.max(max_q), "tmp_rot")?; + let q = l.zeros(max_q, "q")?; + let k = l.zeros(max_kv, "k")?; + let v = l.zeros(max_kv, "v")?; + let attn_out = l.zeros(max_q, "attn_out")?; + let q8_flash_partials = l.zeros(partials, "flash_partials")?; + let gate_ffn = l.zeros(cfg.max_ffn_hidden_dim(), "gate_ffn")?; + let up_ffn = l.zeros(cfg.max_ffn_hidden_dim(), "up_ffn")?; + let ffn_hidden = l.zeros(cfg.max_ffn_hidden_dim(), "ffn_hidden")?; + let ffn_out = l.zeros(dim, "ffn_out")?; + let mut ple = |n: usize, label: &str| -> Result, String> { + (ple_dim != 0).then(|| l.zeros(n, label)).transpose() + }; + let ple_token_inputs = ple(ple_packed, "ple_token_inputs")?; + let ple_projection_all = ple(ple_packed, "ple_projection_all")?; + let ple_gate = ple(ple_dim, "ple_gate")?; + let ple_hidden = ple(ple_dim, "ple_hidden")?; + let ple_out = ple(dim, "ple_out")?; + let moe = if cfg.enable_moe_block { + Some(MoeScratch::new(l, cfg)?) + } else { + None + }; + let logits = l.zeros(cfg.vocab_size, "logits")?; Ok(Gemma4State { kv_sliding, kv_full, @@ -999,46 +1215,27 @@ impl Gemma4State { max_seq, n_tokens: 0, ar_warmed_up: false, - x: alloc(gpu, dim, "x")?, - residual: alloc(gpu, dim, "residual")?, - tmp: alloc(gpu, dim, "tmp")?, - tmp_rot: alloc(gpu, dim, "tmp_rot")?, - q: alloc(gpu, max_q, "q")?, - k: alloc(gpu, max_kv, "k")?, - v: alloc(gpu, max_kv, "v")?, - attn_out: alloc(gpu, max_q, "attn_out")?, - q8_flash_partials: alloc(gpu, q8_flash_partials_len, "q8_flash_partials")?, + x, + residual, + tmp, + tmp_rot, + q, + k, + v, + attn_out, + q8_flash_partials, v_norm_ones, - gate_ffn: alloc(gpu, cfg.max_ffn_hidden_dim(), "gate_ffn")?, - up_ffn: alloc(gpu, cfg.max_ffn_hidden_dim(), "up_ffn")?, - ffn_hidden: alloc(gpu, cfg.max_ffn_hidden_dim(), "ffn_hidden")?, - ffn_out: alloc(gpu, dim, "ffn_out")?, - ple_token_inputs: if ple_dim != 0 { - Some(alloc(gpu, ple_packed, "ple_token_inputs")?) - } else { - None - }, - ple_projection_all: if ple_dim != 0 { - Some(alloc(gpu, ple_packed, "ple_projection_all")?) - } else { - None - }, - ple_gate: if ple_dim != 0 { - Some(alloc(gpu, ple_dim, "ple_gate")?) - } else { - None - }, - ple_hidden: if ple_dim != 0 { - Some(alloc(gpu, ple_dim, "ple_hidden")?) - } else { - None - }, - ple_out: if ple_dim != 0 { - Some(alloc(gpu, dim, "ple_out")?) - } else { - None - }, - logits: alloc(gpu, cfg.vocab_size, "logits")?, + gate_ffn, + up_ffn, + ffn_hidden, + ffn_out, + ple_token_inputs, + ple_projection_all, + ple_gate, + ple_hidden, + ple_out, + moe, + logits, }) } @@ -1046,15 +1243,60 @@ impl Gemma4State { self.n_tokens = 0; } + /// Every scratch buffer (not the KV caches), for whole-surface resets. + pub fn scratch_tensors(&self) -> Vec<&GpuTensor> { + let mut out = vec![ + &self.x, + &self.residual, + &self.tmp, + &self.tmp_rot, + &self.q, + &self.k, + &self.v, + &self.attn_out, + &self.q8_flash_partials, + &self.gate_ffn, + &self.up_ffn, + &self.ffn_hidden, + &self.ffn_out, + &self.logits, + ]; + out.extend( + [ + &self.ple_token_inputs, + &self.ple_projection_all, + &self.ple_gate, + &self.ple_hidden, + &self.ple_out, + ] + .into_iter() + .flatten(), + ); + if let Some(m) = &self.moe { + out.extend([ + &m.cur_mlp, + &m.pre2, + &m.pre2_rot, + &m.router_in, + &m.router_logits, + &m.topk_indices, + &m.topk_weights, + &m.cur_moe, + &m.gate_batch, + &m.up_batch, + &m.hidden_batch, + ]); + } + out + } + /// Return all GPU state buffers (both KV caches, the device position - /// scalar, and the per-decode scratch tensors) to the pool. Consumes - /// self. Caller follows with `gpu.drain_pool()` (the daemon's - /// `unload_model` already does). + /// scalar, and the per-decode scratch tensors) to the pool. Consumes self. pub fn free_gpu(self, gpu: &mut Gpu) { let _ = self.kv_sliding.free_gpu(gpu); let _ = self.kv_full.free_gpu(gpu); - let _ = gpu.hip.free(self.pos_buf); - for t in [ + let tensors = [ + self.pos_buf, self.x, self.residual, self.tmp, @@ -1070,19 +1312,219 @@ impl Gemma4State { self.ffn_hidden, self.ffn_out, self.logits, - ] { - let _ = gpu.free_tensor(t); - } - for t in [ + ]; + let ple = [ self.ple_token_inputs, self.ple_projection_all, self.ple_gate, self.ple_hidden, self.ple_out, - ] { - if let Some(t) = t { - let _ = gpu.free_tensor(t); + ]; + for t in tensors + .into_iter() + .chain(ple.into_iter().flatten()) + .chain(self.moe.into_iter().flat_map(MoeScratch::tensors)) + { + let _ = gpu.free_tensor(t); + } + } +} + +#[cfg(test)] +mod load_rollback_tests { + use super::*; + use hipfire_runtime::hfq::{write_hfqm_package_mem, HfqMemTensor}; + + fn tiny_config(moe: bool) -> Gemma4Config { + let meta = serde_json::json!({"config": {"text_config": { + "hidden_size": 16, "num_hidden_layers": 2, "vocab_size": 16, + "num_attention_heads": 2, "head_dim": 8, "num_key_value_heads": 2, + "global_head_dim": 8, "num_global_key_value_heads": 2, + "attention_k_eq_v": true, "intermediate_size": 32, + "layer_types": ["sliding_attention", "full_attention"], + "enable_moe_block": moe, "moe_intermediate_size": 8, + "num_experts": 2, "top_k_experts": 2, + }}}); + Gemma4Config::from_metadata_json(&meta.to_string()).expect("tiny config") + } + + fn mem(name: String, quant_type: u8, shape: Vec, data: Vec) -> HfqMemTensor { + HfqMemTensor { + name, + quant_type, + shape, + group_size: 0, + data, + } + } + + fn ones(n: usize) -> Vec { + f32_bytes(&vec![1.0f32; n]).to_vec() + } + + /// An MQ4G256 blob plus its F16 AWQ sidecar, so rollback must reclaim + /// sidecars too. + fn proj(name: String, m: usize, k: usize) -> Vec { + let awq = name.replace(".weight", ".awq_scale.weight"); + vec![ + mem(name, 13, vec![m as u32, k as u32], vec![0u8; 256]), + mem(awq, 1, vec![k as u32], vec![0u8; k * 2]), + ] + } + + fn fixture(cfg: &Gemma4Config) -> Vec { + let root = "model.language_model"; + let dim = cfg.dim; + let mut v = vec![ + mem( + format!("{root}.embed_tokens.weight"), + 2, + vec![cfg.vocab_size as u32, dim as u32], + ones(cfg.vocab_size * dim), + ), + mem( + format!("{root}.norm.weight"), + 2, + vec![dim as u32], + ones(dim), + ), + ]; + for (i, ty) in cfg.layer_types.iter().enumerate() { + let p = format!("{root}.layers.{i}"); + for rel in [ + "input_layernorm", + "post_attention_layernorm", + "pre_feedforward_layernorm", + "post_feedforward_layernorm", + ] { + v.push(mem( + format!("{p}.{rel}.weight"), + 2, + vec![dim as u32], + ones(dim), + )); + } + let (hd, kv) = match ty { + LayerType::Sliding => (cfg.sliding_head_dim, cfg.sliding_n_kv_heads), + LayerType::Full => (cfg.full_head_dim, cfg.full_n_kv_heads), + }; + let (q_dim, kv_dim) = (cfg.n_heads * hd, kv * hd); + v.extend(proj(format!("{p}.self_attn.q_proj.weight"), q_dim, dim)); + v.extend(proj(format!("{p}.self_attn.k_proj.weight"), kv_dim, dim)); + if *ty == LayerType::Sliding { + v.extend(proj(format!("{p}.self_attn.v_proj.weight"), kv_dim, dim)); + } + v.extend(proj(format!("{p}.self_attn.o_proj.weight"), dim, q_dim)); + v.extend(proj( + format!("{p}.mlp.gate_proj.weight"), + cfg.hidden_dim, + dim, + )); + v.extend(proj(format!("{p}.mlp.up_proj.weight"), cfg.hidden_dim, dim)); + v.extend(proj( + format!("{p}.mlp.down_proj.weight"), + dim, + cfg.hidden_dim, + )); + for rel in ["q_norm", "k_norm"] { + v.push(mem( + format!("{p}.self_attn.{rel}.weight"), + 2, + vec![hd as u32], + ones(hd), + )); + } + if cfg.enable_moe_block { + let n = cfg.num_experts; + v.extend(proj(format!("{p}.router.proj.weight"), n, dim)); + v.push(mem( + format!("{p}.router.scale"), + 2, + vec![dim as u32], + ones(dim), + )); + v.push(mem( + format!("{p}.router.per_expert_scale"), + 2, + vec![n as u32], + ones(n), + )); + for rel in [ + "pre_feedforward_layernorm_2", + "post_feedforward_layernorm_1", + "post_feedforward_layernorm_2", + ] { + v.push(mem( + format!("{p}.{rel}.weight"), + 2, + vec![dim as u32], + ones(dim), + )); + } + for x in 0..n { + for kind in ["gate_up_proj", "down_proj"] { + v.push(mem( + format!("{p}.experts.{x}.{kind}.weight"), + 3, + vec![64], + vec![0u8; 64], + )); + } + } } } + v + } + + /// Fail every allocation of a load in turn: each failure is the injected + /// one and returns the pool to its warm level (completed layers, MoE pools + /// and AWQ sidecars included); the first clean load then frees everything. + fn sweep(tag: &str, moe: bool) { + let mut gpu = match Gpu::init() { + Ok(gpu) => gpu, + Err(e) => return eprintln!("skip: no GPU ({e:?})"), + }; + let cfg = tiny_config(moe); + let path = + std::env::temp_dir().join(format!("gemma4_ledger_{tag}_{}.hfq", std::process::id())); + write_hfqm_package_mem(&path, 13, "{}", &fixture(&cfg)).expect("write fixture"); + let hfq = HfqFile::open(&path).expect("open fixture"); + for _ in 0..2 { + Gemma4Weights::load(&hfq, &cfg, &mut gpu) + .expect("warm load") + .free_gpu(&mut gpu); + } + let warm = gpu.pool_stats().0; + let mut fail_at = 1; + loop { + match Gemma4Weights::load_with_fault(&hfq, &cfg, &mut gpu, Some(fail_at)) { + Ok(weights) => { + assert!(moe == weights.layers[0].moe().is_some()); + weights.free_gpu(&mut gpu); + break; + } + Err(e) => { + assert!(e.contains("injected"), "fail_at={fail_at}: {e}"); + assert_eq!(gpu.pool_stats().0, warm, "fail_at={fail_at} leaked"); + } + } + fail_at += 1; + } + assert!(fail_at > 10, "the sweep must cover every staged owner"); + assert_eq!(gpu.pool_stats().0, warm, "clean load leaked"); + gpu.drain_pool(); + let _ = std::fs::remove_file(&path); + } + + #[test] + #[ignore = "requires an AMD GPU; exercises ledger rollback across layers and AWQ sidecars"] + fn dense_load_failure_reclaims_every_owner() { + sweep("dense", false); + } + + #[test] + #[ignore = "requires an AMD GPU; exercises MoE pool / pointer-table rollback"] + fn moe_load_failure_reclaims_every_owner() { + sweep("moe", true); } } diff --git a/crates/hipfire-arch-gemma4/src/lib.rs b/crates/hipfire-arch-gemma4/src/lib.rs index 0dccd62756..2f1c7df0de 100644 --- a/crates/hipfire-arch-gemma4/src/lib.rs +++ b/crates/hipfire-arch-gemma4/src/lib.rs @@ -2,10 +2,9 @@ // Copyright (c) 2026 Kaden Schutt, Kevin Read // hipfire — see LICENSE and NOTICE in the project root. -//! Gemma 4 dense text plus strictly admitted E2B/E4B text variants for hipfire. -//! -//! A DENSE 12B port of the Gemma 4 architecture. Distinguishing features vs a -//! plain pre-norm transformer: +//! Gemma 4 text models for hipfire: the dense 12B, the strictly admitted +//! E2B/E4B variants and the 26B-A4B MoE, on one weight stack and one layer +//! program. Distinguishing features vs a plain pre-norm transformer: //! - **Hybrid attention**: 5 sliding-window layers (head_dim 256, window //! 1024, RoPE θ=10000 full rotate-half) per 1 full layer (head_dim 512, //! partial proportional RoPE θ=1e6 over the first 25% of head_dim). The @@ -19,34 +18,34 @@ //! - **gelu_pytorch_tanh** SwiGLU MLP; **embed scaled by √hidden_size**; //! **tied LM head**; **final logit softcap** `tanh(x/30)*30`. //! -//! DROPPED from the broader Gemma 4 family (the dense 12B uses none of these): -//! MoE blocks, the E-series per-layer-embedding / KV-sharing-layer / -//! double-wide-MLP machinery, and vision/audio. +//! - **E-series**: per-layer embeddings, KV-sharing layers, double-wide MLP. +//! - **MoE** (26B-A4B): routed experts in parallel with every dense FFN. +//! +//! Vision/audio towers are out of scope. //! //! arch_id = 13 (see docs/architecture-ids.md). Reuses existing kernels: //! `rope_f32` / `rope_partial_halved_f32`, `attention_q8_0_kv_swa`, //! `gelu_tanh_f32`, `logit_softcap_f32`, plus the shared GEMV path. pub mod arch; +pub mod bundle; pub mod carrier; pub mod config; pub mod drafter; +pub mod emit; pub mod forward; pub mod gemma4; -pub mod lowered; +pub mod mtp; mod program; -pub mod speculative; -pub use carrier::{ - gemma4_context_admission, gemma4_source_uses_lowered, load_gemma4_bundle, Gemma4Bundle, - Gemma4EagerBundle, Gemma4LoweredBundle, -}; +pub use bundle::Gemma4Bundle; +pub use carrier::load_gemma4_bundle; pub use arch::{Gemma4, ARCH_ID}; pub use config::{Gemma4Config, LayerType, RopeType}; pub use drafter::{ - drafter_step, drafter_step_from_concat, DrafterLayerWeights, DrafterStepOut, - Gemma4DrafterConfig, Gemma4DrafterScratch, Gemma4DrafterWeights, DRAFTER_ARCH_ID, + draft_step, drafter_step_from_concat, DrafterLayerWeights, DrafterStepOut, Gemma4DrafterConfig, + Gemma4DrafterScratch, Gemma4DrafterWeights, DRAFTER_ARCH_ID, }; pub use forward::{forward_batch, forward_batch_spec}; pub use gemma4::{FullLayerWeights, Gemma4State, Gemma4Weights, LayerWeights, SlidingLayerWeights}; -pub use speculative::{spec_step_gemma4_eagle, Gemma4SpecScratch, SpecStepOut}; +pub use mtp::Gemma4Drafter; diff --git a/crates/hipfire-arch-gemma4/src/lowered.rs b/crates/hipfire-arch-gemma4/src/lowered.rs deleted file mode 100644 index bf28f873ac..0000000000 --- a/crates/hipfire-arch-gemma4/src/lowered.rs +++ /dev/null @@ -1,3403 +0,0 @@ -// SPDX-License-Identifier: MIT OR Apache-2.0 -// Copyright (c) 2026 Kaden Schutt, Kate, Kevin Read -// hipfire — see LICENSE and NOTICE in the project root. - -//! Gemma 4 model: hybrid sliding-window + full attention, dense FFN (SwiGLU + gelu_pytorch_tanh). -//! -//! Architectural features vs. Qwen3.5: -//! • Sliding-window attention on 5 of every 6 layers (window=1024). -//! • Full attention layers use head_dim=512 (global_head_dim) with -//! attention_k_eq_v: V is the pre-k_norm output of k_proj (no v_proj). -//! • Partial proportional RoPE on full layers (first 64 of 512 dims rotate, -//! rope_theta=1e6; sliding uses default RoPE with theta=10000). -//! • Sandwich RMSNorm: input + post-attn + pre-FFN + post-FFN per layer, -//! plus a learned per-layer `layer_scalar [1]` at layer end. -//! • Attention scale = 1.0 (not 1/√d); Q/K norms absorb scaling. -//! • Final logit softcap: `tanh(logits/30) * 30` before sampling. -//! • MLP: SwiGLU with `gelu_pytorch_tanh` activation. -//! • Tied LM head (embed_tokens.weight aliased). -//! • Embed scale: sqrt(hidden_size) multiplied onto every embedding row lookup. - -use hip_bridge::HipResult; -use hipfire_dispatch::context::DispatchCtx; -use hipfire_dispatch::pipeline::execute_steps; -use hipfire_runtime::hfq::HfqFile; -use hipfire_runtime::llama::{f16_to_f32, EmbeddingFormat, WeightTensor}; -use hipfire_runtime::weight_store::upload_pooled_bytes; -use rdna_compute::{DType, Gpu, GpuTensor}; - -// ─── Config ───────────────────────────────────────────────────────────── - -#[derive(Debug, Clone, Copy, PartialEq)] -pub enum LayerType { - /// Sliding-window causal attention (window=1024 on 31B). - Sliding, - /// Full causal attention (global). - Full, -} - -#[derive(Debug, Clone, Copy, PartialEq)] -pub enum RopeType { - /// Standard RoPE: all head_dim positions rotate. - Default, - /// Proportional RoPE (Gemma 4 full layers): only the first - /// `partial_rotary_factor × head_dim` positions rotate; rest are NoPE. - Proportional, -} - -#[derive(Debug, Clone)] -pub struct Gemma4Config { - // Common - pub dim: usize, // hidden_size, e.g. 5376 on 31B - pub n_layers: usize, // 60 on 31B - pub vocab_size: usize, // 262144 on Gemma 4 - pub norm_eps: f32, // 1e-6 - pub bos_token: u32, // 2 - pub eos_token: u32, // 1 - pub pad_token: u32, // 0 - - // Attention heads (same count for sliding + full) - pub n_heads: usize, // 32 on 31B - - // Sliding-window attention - pub sliding_head_dim: usize, // 256 on 31B - pub sliding_n_kv_heads: usize, // 16 on 31B - pub sliding_rope_theta: f32, // 10000.0 - pub sliding_window: usize, // 1024 - - // Full attention (global) - pub full_head_dim: usize, // 512 on 31B (= global_head_dim) - pub full_n_kv_heads: usize, // 4 on 31B - pub full_rope_theta: f32, // 1_000_000.0 - pub full_rope_type: RopeType, // Proportional on 31B - pub full_partial_rotary_factor: f32, // 0.25 - pub attention_k_eq_v: bool, // true on 31B — V = pre-k_norm output - - // FFN (SwiGLU, gelu_pytorch_tanh) - pub hidden_dim: usize, // intermediate_size = 21504 on 31B - - // MoE (26B-A4B). enable_moe_block=true → every layer carries a parallel - // MoE branch whose output sums with the standard SwiGLU output before - // the post_feedforward_layernorm. Zero on dense models (31B). - pub enable_moe_block: bool, // true on 26B-A4B - pub moe_intermediate_size: usize, // 704 on 26B-A4B (per-expert FFN hidden) - pub num_experts: usize, // 128 on 26B-A4B - pub top_k_experts: usize, // 8 on 26B-A4B (kernel hardcoded to 8) - - // Output - pub final_logit_softcapping: f32, // 30.0 — tanh(x/30)*30 - pub tie_word_embeddings: bool, // true — lm_head aliases embed_tokens - pub embed_scale: f32, // sqrt(dim), applied at embed lookup - - // Per-layer dispatch (len == n_layers) - pub layer_types: Vec, - - // Vision integration (present even on text-only 31B since config ships it) - pub has_vision: bool, - pub image_token_id: u32, // 258880 - pub boi_token_id: u32, // 255999 - pub eoi_token_id: u32, // 258882 - pub audio_token_id: u32, // 258881 (reserved, unused on dense 31B) - pub video_token_id: u32, // 258884 (reserved) -} - -pub fn config_from_hfq(hfq: &HfqFile) -> Option { - let meta: serde_json::Value = serde_json::from_str(&hfq.metadata_json).ok()?; - let config = meta.get("config")?; - let tc = config.get("text_config").unwrap_or(config); - - let dim = tc.get("hidden_size")?.as_u64()? as usize; - let n_layers = tc.get("num_hidden_layers")?.as_u64()? as usize; - let vocab_size = tc.get("vocab_size")?.as_u64()? as usize; - let norm_eps = tc - .get("rms_norm_eps") - .and_then(|v| v.as_f64()) - .unwrap_or(1e-6) as f32; - let bos_token = tc.get("bos_token_id").and_then(|v| v.as_u64()).unwrap_or(2) as u32; - let eos_token = tc.get("eos_token_id").and_then(|v| v.as_u64()).unwrap_or(1) as u32; - let pad_token = tc.get("pad_token_id").and_then(|v| v.as_u64()).unwrap_or(0) as u32; - - let n_heads = tc.get("num_attention_heads")?.as_u64()? as usize; - - // Sliding attention params - let sliding_head_dim = tc - .get("head_dim") - .and_then(|v| v.as_u64()) - .map(|v| v as usize) - .unwrap_or(dim / n_heads); - let sliding_n_kv_heads = tc - .get("num_key_value_heads") - .and_then(|v| v.as_u64()) - .unwrap_or(n_heads as u64) as usize; - let sliding_window = tc - .get("sliding_window") - .and_then(|v| v.as_u64()) - .unwrap_or(1024) as usize; - - // Full attention params (may differ from sliding) - let full_head_dim = tc - .get("global_head_dim") - .and_then(|v| v.as_u64()) - .map(|v| v as usize) - .unwrap_or(sliding_head_dim); - let full_n_kv_heads = tc - .get("num_global_key_value_heads") - .and_then(|v| v.as_u64()) - .unwrap_or(sliding_n_kv_heads as u64) as usize; - let attention_k_eq_v = tc - .get("attention_k_eq_v") - .and_then(|v| v.as_bool()) - .unwrap_or(false); - - // rope_parameters is a dict with "sliding_attention" and "full_attention" sub-dicts - // per the Gemma 4 config schema. Parse both independently. - let rope_params = tc.get("rope_parameters"); - let sliding_rope = rope_params.and_then(|r| r.get("sliding_attention")); - let full_rope = rope_params.and_then(|r| r.get("full_attention")); - - let sliding_rope_theta = sliding_rope - .and_then(|r| r.get("rope_theta")) - .and_then(|v| v.as_f64()) - .unwrap_or(10_000.0) as f32; - let full_rope_theta = full_rope - .and_then(|r| r.get("rope_theta")) - .and_then(|v| v.as_f64()) - .unwrap_or(1_000_000.0) as f32; - let full_rope_type = match full_rope - .and_then(|r| r.get("rope_type")) - .and_then(|v| v.as_str()) - { - Some("proportional") => RopeType::Proportional, - _ => RopeType::Default, - }; - let full_partial_rotary_factor = full_rope - .and_then(|r| r.get("partial_rotary_factor")) - .and_then(|v| v.as_f64()) - .unwrap_or(1.0) as f32; - - let hidden_dim = tc.get("intermediate_size")?.as_u64()? as usize; - - // MoE config (26B-A4B). Absent / false on dense models (31B). - let enable_moe_block = tc - .get("enable_moe_block") - .and_then(|v| v.as_bool()) - .unwrap_or(false); - let moe_intermediate_size = tc - .get("moe_intermediate_size") - .and_then(|v| v.as_u64()) - .unwrap_or(0) as usize; - let num_experts = tc.get("num_experts").and_then(|v| v.as_u64()).unwrap_or(0) as usize; - let top_k_experts = tc - .get("top_k_experts") - .and_then(|v| v.as_u64()) - .unwrap_or(0) as usize; - - let final_logit_softcapping = tc - .get("final_logit_softcapping") - .and_then(|v| v.as_f64()) - .unwrap_or(0.0) as f32; - let tie_word_embeddings = tc - .get("tie_word_embeddings") - .and_then(|v| v.as_bool()) - .or_else(|| config.get("tie_word_embeddings").and_then(|v| v.as_bool())) - .unwrap_or(true); - - let embed_scale = (dim as f32).sqrt(); - - let layer_types: Vec = tc - .get("layer_types") - .and_then(|v| v.as_array()) - .map(|arr| { - arr.iter() - .map(|v| match v.as_str().unwrap_or("sliding_attention") { - "full_attention" => LayerType::Full, - _ => LayerType::Sliding, - }) - .collect() - }) - .unwrap_or_else(|| vec![LayerType::Sliding; n_layers]); - - // Multimodal token IDs (top-level in config, not under text_config) - let has_vision = config - .get("vision_config") - .map(|v| !v.is_null()) - .unwrap_or(false); - let image_token_id = config - .get("image_token_id") - .and_then(|v| v.as_u64()) - .unwrap_or(258880) as u32; - let boi_token_id = config - .get("boi_token_id") - .and_then(|v| v.as_u64()) - .unwrap_or(255999) as u32; - let eoi_token_id = config - .get("eoi_token_id") - .and_then(|v| v.as_u64()) - .unwrap_or(258882) as u32; - let audio_token_id = config - .get("audio_token_id") - .and_then(|v| v.as_u64()) - .unwrap_or(258881) as u32; - let video_token_id = config - .get("video_token_id") - .and_then(|v| v.as_u64()) - .unwrap_or(258884) as u32; - - Some(Gemma4Config { - dim, - n_layers, - vocab_size, - norm_eps, - bos_token, - eos_token, - pad_token, - n_heads, - sliding_head_dim, - sliding_n_kv_heads, - sliding_rope_theta, - sliding_window, - full_head_dim, - full_n_kv_heads, - full_rope_theta, - full_rope_type, - full_partial_rotary_factor, - attention_k_eq_v, - hidden_dim, - enable_moe_block, - moe_intermediate_size, - num_experts, - top_k_experts, - final_logit_softcapping, - tie_word_embeddings, - embed_scale, - layer_types, - has_vision, - image_token_id, - boi_token_id, - eoi_token_id, - audio_token_id, - video_token_id, - }) -} - -// ─── Weights ──────────────────────────────────────────────────────────── - -/// Per-layer weights for a SLIDING layer (head_dim=256, 16 KV heads, full RoPE). -pub struct SlidingLayerWeights { - pub input_layernorm: GpuTensor, // [dim] - pub post_attention_layernorm: GpuTensor, // [dim] - pub pre_feedforward_layernorm: GpuTensor, // [dim] - pub post_feedforward_layernorm: GpuTensor, // [dim] - pub layer_scalar: GpuTensor, // [1] - /// Host-side mirror of layer_scalar. Populated at load time so decode can - /// call `gpu.scale_f32(x, layer_scalar_host)` without a D2H round-trip. - pub layer_scalar_host: f32, - - // Attention (sliding — head_dim=256) - pub q_proj: WeightTensor, // [n_heads * 256, dim] - pub k_proj: WeightTensor, // [16 * 256, dim] - pub v_proj: WeightTensor, // [16 * 256, dim] - pub o_proj: WeightTensor, // [dim, n_heads * 256] - pub q_norm: GpuTensor, // [256] - pub k_norm: GpuTensor, // [256] - - // MLP (SwiGLU) - pub gate_proj: WeightTensor, // [hidden_dim, dim] - pub up_proj: WeightTensor, // [hidden_dim, dim] - pub down_proj: WeightTensor, // [dim, hidden_dim] - - // MoE branch — Some on 26B-A4B (every layer is MoE), None on dense models. - pub moe: Option, -} - -/// Per-layer weights for a FULL layer (head_dim=512, 4 KV heads, K=V shared). -/// -/// Note: no `v_proj` — V is the pre-k_norm output of k_proj, renormed by -/// weight-less `v_norm`. No `v_norm` tensor either (no_scale — the `with_scale=False` -/// RMSNorm applies only the divide, no learned gain). We reuse the existing -/// rmsnorm kernel with a ones-filled `v_norm_ones` buffer (shared across -/// full-attn layers) to preserve the no-scale semantics. -pub struct FullLayerWeights { - pub input_layernorm: GpuTensor, - pub post_attention_layernorm: GpuTensor, - pub pre_feedforward_layernorm: GpuTensor, - pub post_feedforward_layernorm: GpuTensor, - pub layer_scalar: GpuTensor, - /// Host-side mirror of layer_scalar. See SlidingLayerWeights for rationale. - pub layer_scalar_host: f32, - - // Attention (full — head_dim=512, K=V) - pub q_proj: WeightTensor, // [n_heads * 512, dim] - pub k_proj: WeightTensor, // [4 * 512, dim] - // no v_proj — V = pre-k_norm output of k_proj - pub o_proj: WeightTensor, // [dim, n_heads * 512] - pub q_norm: GpuTensor, // [512] - pub k_norm: GpuTensor, // [512] - // no v_norm weight — v_norm is no-scale (divide only) - - // MLP (SwiGLU, same shape as sliding) - pub gate_proj: WeightTensor, - pub up_proj: WeightTensor, - pub down_proj: WeightTensor, - - // MoE branch — Some on 26B-A4B (every layer is MoE), None on dense models. - pub moe: Option, -} - -/// Per-expert FFN weights for a single MoE expert. 128 of these per layer -/// on 26B-A4B; views into the per-layer pool allocation (so `free_gpu` -/// doesn't free these — the pool owns the bytes). -pub struct MoeExpertWeights { - /// `[2 * moe_intermediate, dim]` — gate + up fused. Rows [0, mi) are - /// gate; rows [mi, 2*mi) are up. Quantized as MQ4G256 / MG4G256 when - /// dim is 256-aligned (it is on 26B-A4B: dim=2816). - pub gate_up_proj: WeightTensor, - /// `[dim, moe_intermediate]` — projects per-expert FFN hidden back to dim. - /// On 26B-A4B, mi=704 isn't 256-aligned so this drops to Q8_0 via the - /// quantizer fallback chain. - pub down_proj: WeightTensor, -} - -/// MoE branch weights for a Gemma 4 MoE layer (26B-A4B). Present on every -/// layer when `config.enable_moe_block` is set. The branch adds a parallel -/// FFN computation alongside the standard SwiGLU; outputs are summed via -/// sandwich norms then a final post_feedforward_layernorm closes the layer. -pub struct MoeLayerExtras { - /// `[n_experts, dim]` — projects router input to expert logits. - pub router_proj: WeightTensor, - /// `[dim]` — multiplicative scale on router input (`router.scale` in HF). - pub router_scale: GpuTensor, - /// `[n_experts]` — per-expert post-`down_proj` scale (`router.per_expert_scale`). - pub per_expert_scale: GpuTensor, - /// Host mirror of `per_expert_scale` for fast top-K weight composition. - pub per_expert_scale_host: Vec, - /// `[dim]` — RMSNorm applied to attn_out before the MoE branch. - pub pre_feedforward_layernorm_2: GpuTensor, - /// `[dim]` — RMSNorm applied to cur_mlp (standard SwiGLU output) BEFORE summing. - pub post_feedforward_layernorm_1: GpuTensor, - /// `[dim]` — RMSNorm applied to cur_moe (MoE branch output) BEFORE summing. - pub post_feedforward_layernorm_2: GpuTensor, - /// Pool allocation for all gate_up tensors. Per-expert WeightTensors - /// alias into this; `free_gpu` frees the pool, not each WeightTensor. - pub experts_gate_up_pool: GpuTensor, - /// Pool allocation for all down tensors. Same aliasing. - pub experts_down_pool: GpuTensor, - /// Per-expert views into the pools above. - pub experts: Vec, - /// `[n_exp]` u64 device pointers — one per expert's gate_up weight - /// base. Built once at load by reading each expert's pool sub-view - /// pointer. The indexed MoE kernels read - /// `expert_ptrs[topk_indices[krank]]` to locate the active expert - /// weight WITHOUT a D2H sync. - pub experts_gate_up_ptrs: GpuTensor, - /// `[n_exp]` u64 device pointers for each expert's down weight base. - pub experts_down_ptrs: GpuTensor, -} - -pub enum LayerWeights { - Sliding(SlidingLayerWeights), - Full(FullLayerWeights), -} - -pub struct Gemma4Weights { - /// Token embedding [vocab_size, dim], Q8F16 to keep the 262144×5376 table manageable. - /// Aliased as lm_head when tie_word_embeddings is true. - pub embed_tokens: GpuTensor, - /// Embed/LM-head format tag for dispatch. - pub embd_format: EmbeddingFormat, - /// LM-head projection (shares bytes with embed_tokens when tied). - pub lm_head: WeightTensor, - /// Model-final RMSNorm scale [dim]. - pub final_norm: GpuTensor, - /// Per-layer weights indexed by layer ordinal. - pub layers: Vec, -} - -impl Gemma4Weights { - pub fn free_gpu(self, gpu: &mut Gpu) { - let _ = gpu.free_tensor(self.embed_tokens); - let _ = gpu.free_tensor(self.final_norm); - // lm_head aliases embed_tokens (tied weights) — a Borrowed view, - // dropped here, never freed. - for l in self.layers { - Self::free_layer(gpu, l); - } - } - - /// Reclaim every GPU owner of one completed layer. Norm/scalar tensors go - /// through `free_tensor`; projection weights go through - /// `WeightTensor::free_all` so attached AWQ sidecars are reclaimed too. - /// MoE expert views (`sub_offset` into the pools) are dropped, never - /// freed — the pools own those bytes. Shared with load-time rollback so a - /// late load failure reclaims completed layers exactly like unload does. - fn free_layer(gpu: &mut Gpu, layer: LayerWeights) { - match layer { - LayerWeights::Sliding(s) => { - for t in [ - s.input_layernorm, - s.post_attention_layernorm, - s.pre_feedforward_layernorm, - s.post_feedforward_layernorm, - s.layer_scalar, - s.q_norm, - s.k_norm, - ] { - let _ = gpu.free_tensor(t); - } - for w in [ - s.q_proj, - s.k_proj, - s.v_proj, - s.o_proj, - s.gate_proj, - s.up_proj, - s.down_proj, - ] { - w.free_all(gpu); - } - if let Some(moe) = s.moe { - Self::free_moe(gpu, moe); - } - } - LayerWeights::Full(f) => { - for t in [ - f.input_layernorm, - f.post_attention_layernorm, - f.pre_feedforward_layernorm, - f.post_feedforward_layernorm, - f.layer_scalar, - f.q_norm, - f.k_norm, - ] { - let _ = gpu.free_tensor(t); - } - for w in [ - f.q_proj, - f.k_proj, - f.o_proj, - f.gate_proj, - f.up_proj, - f.down_proj, - ] { - w.free_all(gpu); - } - if let Some(moe) = f.moe { - Self::free_moe(gpu, moe); - } - } - } - } - - fn free_moe(gpu: &mut Gpu, moe: MoeLayerExtras) { - // router_proj is a real owner (buffer + optional AWQ sidecar). - moe.router_proj.free_all(gpu); - for t in [ - moe.router_scale, - moe.per_expert_scale, - moe.pre_feedforward_layernorm_2, - moe.post_feedforward_layernorm_1, - moe.post_feedforward_layernorm_2, - ] { - let _ = gpu.free_tensor(t); - } - // Pool + pointer-table owners. Per-expert WeightTensors alias into the - // pools via sub_offset — Borrowed views, dropped (never freed) with - // `moe.experts` at function end. - for t in [ - moe.experts_gate_up_pool, - moe.experts_down_pool, - moe.experts_gate_up_ptrs, - moe.experts_down_ptrs, - ] { - let _ = gpu.free_tensor(t); - } - } -} -// ─── Load-time fault seam ─────────────────────────────────────────────── -// Private injectable failure points for GPU failure/retry regressions. -// Production passes `None` (one untaken branch per stage, zero behavior -// change); in-file tests pass `Some` to fail the n-th checked GPU upload so -// rollback of every staged owner is exercised. This is the lowered analogue -// of qwen35's `new_opt_with_alloc` counting-allocator seam, shaped as a -// check-hook because gemma4 leaf uploads are heterogeneous (upload_f32 / -// pooled-bytes) and don't funnel through one allocator closure. No env knob, -// no global allocation sweep: rollback frees only slot-staged owners plus -// completed layers, never aliases (lm_head, expert views) or GPU-global -// caches (mq signs/rotation scratch owned by `Gpu`). -struct AllocFaults { - calls: usize, - fail_at: Option, -} - -impl AllocFaults { - fn check(&mut self, stage: &'static str) -> HipResult<()> { - self.calls += 1; - if self.fail_at == Some(self.calls) { - return Err(hip_bridge::HipError::new( - 0, - &format!("injected gemma4 load fault at {stage} (op {})", self.calls), - )); - } - Ok(()) - } -} - -/// Run one staged fault check: fail this GPU-owning step when a test asked -/// for it. `None` (production) is a no-op. -fn fault_check(fault: &mut Option<&mut AllocFaults>, stage: &'static str) -> HipResult<()> { - if let Some(f) = fault { - f.check(stage)?; - } - Ok(()) -} - -// ─── Loading helpers ─────────────────────────────────────────────────── - -/// Decode a shape-[n] F16 or F32 tensor from HFQ into an F32 host Vec. -fn load_f32_vec(hfq: &HfqFile, name: &str, expected_n: usize) -> HipResult> { - let (info, data) = hfq - .tensor_data(name) - .ok_or_else(|| hip_bridge::HipError::new(0, &format!("tensor not found: {name}")))?; - let n: usize = info.shape.iter().map(|&s| s as usize).product(); - if n != expected_n { - return Err(hip_bridge::HipError::new( - 0, - &format!("shape mismatch for {name}: expected {expected_n}, got {n}"), - )); - } - if hipfire_config::developer_var("HIPFIRE_GEMMA4_DUMP") - .ok() - .as_deref() - == Some("1") - && data.len() <= 4 - && name.contains("layer_scalar") - { - eprintln!( - "[gemma4] load_f32_vec({name}): qt={}, shape={:?}, raw_bytes={:02x?}", - info.quant_type, - info.shape, - &data[..data.len().min(4)] - ); - } - let f32_data = match info.quant_type { - 1 => data - .chunks_exact(2) - .map(|c| f16_to_f32(u16::from_le_bytes([c[0], c[1]]))) - .collect(), - 2 => data - .chunks_exact(4) - .map(|c| f32::from_le_bytes([c[0], c[1], c[2], c[3]])) - .collect(), - 16 => data - .chunks_exact(2) - .map(|c| f32::from_bits((u16::from_le_bytes([c[0], c[1]]) as u32) << 16)) - .collect(), - qt => { - return Err(hip_bridge::HipError::new( - 0, - &format!("expected F16/F32 for {name}, got qt={qt}"), - )); - } - }; - Ok(f32_data) -} - -/// Load a Gemma 4 RMSNorm weight — `x * weight` form, NO +1 shift. -/// -/// Distinct from qwen35::load_norm_weight which shifts by +1 for HF Gemma -/// 2/3-style `x * (1 + weight)`. Gemma 4 uses plain `x * weight` with weights -/// initialized to 1.0 (see modeling_gemma4.py::Gemma4RMSNorm line 157). -fn load_gemma4_norm(hfq: &HfqFile, gpu: &mut Gpu, name: &str, dim: usize) -> HipResult { - let f32_data = load_f32_vec(hfq, name, dim)?; - gpu.upload_f32(&f32_data, &[dim]) -} - -/// Load a 256-element head-dim Q/K RMSNorm weight. Same semantics as -/// `load_gemma4_norm` but scoped to the attention head_dim (256 on sliding, -/// 512 on full). -fn load_gemma4_head_norm( - hfq: &HfqFile, - gpu: &mut Gpu, - name: &str, - head_dim: usize, -) -> HipResult { - load_gemma4_norm(hfq, gpu, name, head_dim) -} - -/// Load the per-layer `layer_scalar` — shape-[1] BF16/F16 tensor — returning -/// both a GPU-resident [1]-tensor (for potential batched use) and its host-side -/// f32 value (used by the decode path to call `scale_f32(x, cpu_scalar)`). -fn load_layer_scalar(hfq: &HfqFile, gpu: &mut Gpu, name: &str) -> HipResult<(GpuTensor, f32)> { - let data = load_f32_vec(hfq, name, 1)?; - let host_val = data[0]; - let gpu_tensor = gpu.upload_f32(&data, &[1])?; - Ok((gpu_tensor, host_val)) -} -/// Load an AWQ per-channel scale sidecar (`.awq_scale.weight`, F16 -/// [k]) through the pooled uploader. Absence or malformed sidecars return -/// None exactly like `hfq::load_awq_scale`; a present, well-formed sidecar -/// whose allocation/copy fails propagates the error into staged rollback -/// instead of silently dropping the scale (which would compute `(W·s)·x`). -fn load_gemma4_awq_scale( - hfq: &HfqFile, - gpu: &mut Gpu, - weight_name: &str, - k: usize, -) -> HipResult> { - let sidecar_name = match weight_name.strip_suffix(".weight") { - Some(stem) => format!("{stem}.awq_scale.weight"), - None => format!("{weight_name}.awq_scale.weight"), - }; - let Some((sc_info, sc_data)) = hfq.tensor_data_vec(&sidecar_name) else { - return Ok(None); - }; - if sc_info.quant_type != 1 { - eprintln!( - "warning: AWQ sidecar {sidecar_name} has quant_type={} (expected 1=F16); skipping", - sc_info.quant_type - ); - return Ok(None); - } - if sc_info.shape.len() != 1 || sc_info.shape[0] as usize != k { - eprintln!( - "warning: AWQ sidecar {sidecar_name} shape mismatch ({:?} vs expected [{}]); skipping", - sc_info.shape, k - ); - return Ok(None); - } - let f32_data: Vec = sc_data - .chunks_exact(2) - .map(|c| f16_to_f32(u16::from_le_bytes([c[0], c[1]]))) - .collect(); - let f32_bytes: Vec = f32_data.iter().flat_map(|v| v.to_le_bytes()).collect(); - let t = upload_pooled_bytes(gpu, &f32_bytes, &[f32_bytes.len()])?; - Ok(Some(t)) -} - -/// HFQ4-G128 rows each start a fresh 72-byte group; a trailing partial group -/// is padded. Files quantized before the 2-D packer (`b4846285e`) packed -/// groups across rows whenever `K % 128 != 0` (Gemma 4 expert `down_proj`, -/// K = 704), which every per-row kernel reads as garbage. Refuse them. -fn check_hfq4g128_rows(name: &str, m: usize, k: usize, bytes: usize) -> HipResult<()> { - let rows = m * k.div_ceil(128) * 72; - if bytes == rows { - return Ok(()); - } - let reason = if bytes == (m * k).div_ceil(128) * 72 { - "packs HFQ4-G128 groups across rows (quantized before hipfire-quantize \ - b4846285e); requantize the model" - } else { - "has an unexpected byte size" - }; - Err(hip_bridge::HipError::new( - 0, - &format!("{name} [{m} x {k}] {reason} ({bytes} bytes, expected {rows})"), - )) -} - -/// Load a quantized projection weight. Mirrors qwen35::load_weight_tensor_raw -/// but uses the Gemma 4 tensor-name convention (`model.language_model.`). -fn load_gemma4_weight( - hfq: &HfqFile, - gpu: &mut Gpu, - name: &str, - m: usize, - k: usize, -) -> HipResult { - load_gemma4_weight_impl(hfq, gpu, name, m, k, None) -} - -/// Injectable-fault twin of [`load_gemma4_weight`]: tests fail the sidecar -/// stage after the primary buffer is owned, proving the primary is reclaimed -/// even though the outer staged rollback never sees a half-built WeightTensor. -fn load_gemma4_weight_impl( - hfq: &HfqFile, - gpu: &mut Gpu, - name: &str, - m: usize, - k: usize, - mut fault: Option<&mut AllocFaults>, -) -> HipResult { - let (info, data) = hfq - .tensor_data(name) - .ok_or_else(|| hip_bridge::HipError::new(0, &format!("tensor not found: {name}")))?; - let dtype = match info.quant_type { - 1 => { - // F16 → upload as f32 - let f32_data: Vec = data - .chunks_exact(2) - .map(|c| f16_to_f32(u16::from_le_bytes([c[0], c[1]]))) - .collect(); - let bytes: &[u8] = unsafe { - std::slice::from_raw_parts(f32_data.as_ptr() as *const u8, f32_data.len() * 4) - }; - let buf = upload_pooled_bytes(gpu, bytes, &[m, k])?; - return Ok(WeightTensor { - buf, - gpu_dtype: DType::F32, - m, - k, - row_stride: 0, - awq_scale: None, - lloyd_lut_e4m3: None, - lloyd_lut_f16: None, - lloyd_lut_c16: None, - paro: None, - }); - } - 16 => { - if rdna_compute::calib_force_bf16() { - // Native BF16 teacher — keep raw 2-byte payload as BF16 for MFMA. - // Otherwise the batched GEMM would land on the scalar F32 kernel. - let buf = upload_pooled_bytes(gpu, data, &[m, k])?; - return Ok(WeightTensor { - buf, - gpu_dtype: DType::BF16, - m, - k, - row_stride: 0, - awq_scale: None, - lloyd_lut_e4m3: None, - lloyd_lut_f16: None, - lloyd_lut_c16: None, - paro: None, - }); - } - // Default: BF16 → widen to F32 (shift, not f16 decode) - let f32_data: Vec = data - .chunks_exact(2) - .map(|c| f32::from_bits((u16::from_le_bytes([c[0], c[1]]) as u32) << 16)) - .collect(); - let bytes: &[u8] = unsafe { - std::slice::from_raw_parts(f32_data.as_ptr() as *const u8, f32_data.len() * 4) - }; - let buf = upload_pooled_bytes(gpu, bytes, &[m, k])?; - return Ok(WeightTensor { - buf, - gpu_dtype: DType::F32, - m, - k, - row_stride: 0, - awq_scale: None, - lloyd_lut_e4m3: None, - lloyd_lut_f16: None, - lloyd_lut_c16: None, - paro: None, - }); - } - 2 => { - // F32 raw (oracle / --format f32 passthrough .hfq) — upload as-is. - let buf = upload_pooled_bytes(gpu, data, &[m, k])?; - return Ok(WeightTensor { - buf, - gpu_dtype: DType::F32, - m, - k, - row_stride: 0, - awq_scale: None, - lloyd_lut_e4m3: None, - lloyd_lut_f16: None, - lloyd_lut_c16: None, - paro: None, - }); - } - // Q8F16/Q8_0 projections are emitted when a matrix is not eligible - // for the requested grouped format. The lowered forward path already - // dispatches DType::Q8_0; keep its loader aligned with the eager and - // drafter loaders instead of rejecting a valid quantizer fallback. - 3 => DType::Q8_0, - 4 => DType::Q4K, - 6 => DType::HFQ4G256, - 7 => DType::HFQ4G128, - 8 => DType::HFQ6G256, - 9 => DType::HFQ2G256, - 10 => DType::HFQ2G128, - 11 => DType::HFQ3G256, - 12 => DType::HFQ3G128, - 13 => DType::MQ4G256, - 44 => DType::MQ4G256V2, - 14 => DType::MQ8G256, - 15 => DType::MQ6G256, - 17 => DType::MQ3G256, - 18 => DType::MQ2G256, - // MG4-G256 — Magnum-Gemma 4-bit. Same binary layout as MQ4G256 (136 B/group), - // differs only in calibration policy at quant time. Alias to MQ4G256 so the - // existing GEMV path handles it without a kernel change. ID was 19 on - // origin/gemma4 pre-rebase; reassigned to 30 because master shipped - // MQ2G256Lloyd at 19. - 30 => DType::MQ4G256, - qt => { - return Err(hip_bridge::HipError::new( - 0, - &format!("unsupported quant_type {qt} for {name}"), - )); - } - }; - if dtype == DType::HFQ4G128 { - check_hfq4g128_rows(name, m, k, data.len())?; - } - let buf = upload_pooled_bytes(gpu, data, &[data.len()])?; - // Fault seam: fail after the primary owns its buffer but before the - // sidecar attaches. Rollback here frees the primary directly — the outer - // transaction never received it. - if let Err(e) = fault_check(&mut fault, "awq_sidecar") { - let _ = gpu.free_tensor(buf); - return Err(e); - } - let awq_scale = if dtype.supports_awq_sidecar() { - match load_gemma4_awq_scale(hfq, gpu, name, k) { - Ok(s) => s, - Err(e) => { - let _ = gpu.free_tensor(buf); - return Err(e); - } - } - } else { - None - }; - Ok(WeightTensor { - buf, - gpu_dtype: dtype, - m, - k, - row_stride: 0, - awq_scale, - lloyd_lut_e4m3: None, - lloyd_lut_f16: None, - lloyd_lut_c16: None, - paro: None, - }) -} - -/// Load the MoE branch weights for a single Gemma 4 MoE layer (26B-A4B). -/// Builds 128 expert WeightTensors aliased into per-layer gate_up / down -/// pool allocations. Pooling avoids the small-allocation HIP fragmentation -/// that OOM'd at layer 25 on the pre-pool path (origin/gemma4 commit log). -fn load_moe_layer_extras( - hfq: &HfqFile, - gpu: &mut Gpu, - p: &str, - config: &Gemma4Config, - mut fault: Option<&mut AllocFaults>, -) -> HipResult { - let n_exp = config.num_experts; - let dim = config.dim; - let mi = config.moe_intermediate_size; - - // Staged owners: every successful GPU allocation lands in a slot the - // moment it succeeds. Per-expert WeightTensors are sub_offset views into - // the pools — never staged, never freed. On any error the fail! arm frees - // every staged owner (weights via free_all so AWQ sidecars go too). - let mut router_proj_opt: Option = None; - let mut router_scale_opt: Option = None; - let mut per_expert_scale_opt: Option = None; - let mut per_expert_scale_host: Vec = Vec::new(); - let mut pre2_opt: Option = None; - let mut post1_opt: Option = None; - let mut post2_opt: Option = None; - let mut gate_up_pool_opt: Option = None; - let mut down_pool_opt: Option = None; - let mut gate_ptrs_opt: Option = None; - let mut down_ptrs_opt: Option = None; - macro_rules! fail { - ($e:expr) => {{ - if let Some(t) = down_ptrs_opt.take() { - let _ = gpu.free_tensor(t); - } - if let Some(t) = gate_ptrs_opt.take() { - let _ = gpu.free_tensor(t); - } - if let Some(t) = down_pool_opt.take() { - let _ = gpu.free_tensor(t); - } - if let Some(t) = gate_up_pool_opt.take() { - let _ = gpu.free_tensor(t); - } - for t in [ - post2_opt.take(), - post1_opt.take(), - pre2_opt.take(), - per_expert_scale_opt.take(), - router_scale_opt.take(), - ] - .into_iter() - .flatten() - { - let _ = gpu.free_tensor(t); - } - if let Some(w) = router_proj_opt.take() { - w.free_all(gpu); - } - return Err($e); - }}; - } - macro_rules! check { - ($stage:expr) => { - if let Err(e) = fault_check(&mut fault, $stage) { - fail!(e); - } - }; - } - macro_rules! stage { - ($slot:ident, $stage:expr, $val:expr) => {{ - check!($stage); - match $val { - Ok(v) => { - $slot = Some(v); - } - Err(e) => fail!(e), - } - }}; - } - - stage!( - router_proj_opt, - "moe_router_proj", - load_gemma4_weight(hfq, gpu, &format!("{p}.router.proj.weight"), n_exp, dim) - ); - // NOTE: `router.scale` and `router.per_expert_scale` ship WITHOUT the - // `.weight` suffix in HF's 26B-A4B safetensors (so `should_quantize` - // returns false → stored as F16). Loader uses bare paths. - stage!( - router_scale_opt, - "moe_router_scale", - load_gemma4_norm(hfq, gpu, &format!("{p}.router.scale"), dim) - ); - check!("moe_per_expert_scale_host"); - match load_f32_vec(hfq, &format!("{p}.router.per_expert_scale"), n_exp) { - Ok(v) => { - per_expert_scale_host = v; - } - Err(e) => fail!(e), - } - check!("moe_per_expert_scale"); - match (|| -> HipResult { - let bytes: &[u8] = unsafe { - std::slice::from_raw_parts( - per_expert_scale_host.as_ptr() as *const u8, - per_expert_scale_host.len() * 4, - ) - }; - upload_pooled_bytes(gpu, bytes, &[n_exp]) - })() { - Ok(t) => { - per_expert_scale_opt = Some(t); - } - Err(e) => fail!(e), - } - stage!( - pre2_opt, - "moe_pre_norm2", - load_gemma4_norm( - hfq, - gpu, - &format!("{p}.pre_feedforward_layernorm_2.weight"), - dim, - ) - ); - stage!( - post1_opt, - "moe_post_norm1", - load_gemma4_norm( - hfq, - gpu, - &format!("{p}.post_feedforward_layernorm_1.weight"), - dim, - ) - ); - stage!( - post2_opt, - "moe_post_norm2", - load_gemma4_norm( - hfq, - gpu, - &format!("{p}.post_feedforward_layernorm_2.weight"), - dim, - ) - ); - - // Pool all `n_experts` weights of one kind into a single GPU allocation. - // 128 experts × 2 kinds × 30 layers = 7680 separate hipMalloc on the - // unpooled path fragmented the HIP heap and OOM'd at layer ~25 even - // when total memory fit. Pool collapses that to 60 allocs. - let load_pool = |gpu: &mut Gpu, base: &str| -> HipResult<(GpuTensor, DType, usize)> { - // First pass: read first expert to learn quant_type + bytes-per-expert. - let first_name = format!("{p}.experts.0.{base}.weight"); - let (first_info, first_data) = hfq.tensor_data(&first_name).ok_or_else(|| { - hip_bridge::HipError::new(0, &format!("MoE expert tensor not found: {first_name}")) - })?; - let bytes_per_expert = first_data.len(); - let dtype = match first_info.quant_type { - 3 => DType::Q8_0, - 4 => DType::Q4K, - 6 => DType::HFQ4G256, - 7 => DType::HFQ4G128, - 8 => DType::HFQ6G256, - 9 => DType::HFQ2G256, - 10 => DType::HFQ2G128, - 11 => DType::HFQ3G256, - 12 => DType::HFQ3G128, - // MQ4G256 (13) and MG4G256 (30) share dispatch. - 13 | 30 => DType::MQ4G256, - 44 => DType::MQ4G256V2, - 14 => DType::MQ8G256, - 15 => DType::MQ6G256, - 17 => DType::MQ3G256, - 18 => DType::MQ2G256, - qt => { - return Err(hip_bridge::HipError::new( - 0, - &format!("unsupported MoE expert quant_type {qt} for {first_name}"), - )); - } - }; - if dtype == DType::HFQ4G128 { - let (m, k) = match first_info.shape[..] { - [m, k] => (m as usize, k as usize), - _ => { - return Err(hip_bridge::HipError::new( - 0, - &format!( - "{first_name}: expected a 2-D shape, got {:?}", - first_info.shape - ), - )) - } - }; - check_hfq4g128_rows(&first_name, m, k, bytes_per_expert)?; - } - // Concat all experts' bytes into one CPU buffer, upload once. - let mut concat = Vec::with_capacity(bytes_per_expert * n_exp); - concat.extend_from_slice(first_data); - for x in 1..n_exp { - let name = format!("{p}.experts.{x}.{base}.weight"); - let (info, data) = hfq.tensor_data(&name).ok_or_else(|| { - hip_bridge::HipError::new(0, &format!("MoE expert tensor not found: {name}")) - })?; - if data.len() != bytes_per_expert { - return Err(hip_bridge::HipError::new( - 0, - &format!( - "MoE expert {name} byte size mismatch ({} vs {bytes_per_expert})", - data.len() - ), - )); - } - if info.quant_type != first_info.quant_type { - return Err(hip_bridge::HipError::new( - 0, - &format!( - "MoE expert {name} quant_type mismatch ({} vs {})", - info.quant_type, first_info.quant_type - ), - )); - } - concat.extend_from_slice(data); - } - let pool = upload_pooled_bytes(gpu, &concat, &[concat.len()])?; - Ok((pool, dtype, bytes_per_expert)) - }; - - check!("moe_gate_up_pool"); - let (gate_up_pool, gate_up_dtype, gate_up_bpe) = match load_pool(gpu, "gate_up_proj") { - Ok(v) => v, - Err(e) => fail!(e), - }; - gate_up_pool_opt = Some(gate_up_pool); - check!("moe_down_pool"); - let (down_pool, down_dtype, down_bpe) = match load_pool(gpu, "down_proj") { - Ok(v) => v, - Err(e) => fail!(e), - }; - down_pool_opt = Some(down_pool); - - let mut experts = Vec::with_capacity(n_exp); - for x in 0..n_exp { - let gu_view = gate_up_pool_opt - .as_ref() - .expect("gemma4 load: gate-up pool staged before views") - .sub_offset(x * gate_up_bpe, gate_up_bpe); - let dn_view = down_pool_opt - .as_ref() - .expect("gemma4 load: down pool staged before views") - .sub_offset(x * down_bpe, down_bpe); - experts.push(MoeExpertWeights { - gate_up_proj: WeightTensor { - buf: gu_view, - gpu_dtype: gate_up_dtype, - m: 2 * mi, - k: dim, - row_stride: 0, - awq_scale: None, - lloyd_lut_e4m3: None, - lloyd_lut_f16: None, - lloyd_lut_c16: None, - paro: None, - }, - down_proj: WeightTensor { - buf: dn_view, - gpu_dtype: down_dtype, - m: dim, - k: mi, - row_stride: 0, - awq_scale: None, - lloyd_lut_e4m3: None, - lloyd_lut_f16: None, - lloyd_lut_c16: None, - paro: None, - }, - }); - } - - // Build [n_exp] device tensors of u64 weight-base pointers — one for - // each pool. The indexed MoE kernels read these as - // `expert_ptrs[topk_indices[krank]]`, eliminating the per-token D2H - // sync of the legacy CPU per-expert loop. Pointers are stable for the - // model's lifetime (pool allocations don't move). - let gate_up_ptr_u64: Vec = experts - .iter() - .map(|e| e.gate_up_proj.buf.buf.as_ptr() as u64) - .collect(); - let down_ptr_u64: Vec = experts - .iter() - .map(|e| e.down_proj.buf.buf.as_ptr() as u64) - .collect(); - let gate_up_ptr_bytes: Vec = gate_up_ptr_u64 - .iter() - .flat_map(|p| p.to_ne_bytes()) - .collect(); - let down_ptr_bytes: Vec = down_ptr_u64.iter().flat_map(|p| p.to_ne_bytes()).collect(); - // Each u64 = 8 bytes = 2 f32 slots. The tensor sees [n_exp * 2] - // f32 entries; the kernel casts the backing buffer to u64* itself. - check!("moe_gate_up_ptrs"); - match upload_pooled_bytes(gpu, &gate_up_ptr_bytes, &[n_exp * 2]) { - Ok(t) => { - gate_ptrs_opt = Some(t); - } - Err(e) => fail!(e), - } - check!("moe_down_ptrs"); - match upload_pooled_bytes(gpu, &down_ptr_bytes, &[n_exp * 2]) { - Ok(t) => { - down_ptrs_opt = Some(t); - } - Err(e) => fail!(e), - } - - Ok(MoeLayerExtras { - router_proj: router_proj_opt - .take() - .expect("gemma4 load: router_proj staged once"), - router_scale: router_scale_opt - .take() - .expect("gemma4 load: router_scale staged once"), - per_expert_scale: per_expert_scale_opt - .take() - .expect("gemma4 load: per_expert_scale staged once"), - per_expert_scale_host, - pre_feedforward_layernorm_2: pre2_opt.take().expect("gemma4 load: pre_norm2 staged once"), - post_feedforward_layernorm_1: post1_opt - .take() - .expect("gemma4 load: post_norm1 staged once"), - post_feedforward_layernorm_2: post2_opt - .take() - .expect("gemma4 load: post_norm2 staged once"), - experts_gate_up_pool: gate_up_pool_opt - .take() - .expect("gemma4 load: gate-up pool staged once"), - experts_down_pool: down_pool_opt - .take() - .expect("gemma4 load: down pool staged once"), - experts, - experts_gate_up_ptrs: gate_ptrs_opt - .take() - .expect("gemma4 load: gate-up ptrs staged once"), - experts_down_ptrs: down_ptrs_opt - .take() - .expect("gemma4 load: down ptrs staged once"), - }) -} - -/// Load Gemma 4 text model weights from an HFQ file. -/// -/// Design notes: -/// - `lm_head` aliases the `embed_tokens` GPU bytes (tied weights). We upload -/// the embed data once and create a second WeightTensor whose DeviceBuffer -/// points at the same allocation via `buf.alias()`. `Gemma4Weights::free_gpu` -/// skips freeing the LM head to avoid a double-free. -/// - Vision tensors are skipped here — Phase 7 `gemma4_vision::load_weights` -/// picks those up from the same HFQ file in a separate pass. -/// - The `v_norm_ones_full` ones-filled scratch buffer is populated here so -/// the forward pass never has to manage one-time init state. -/// - Transactional construction: every successful GPU allocation lands in a -/// slot the moment it succeeds (embed/final-norm at top level, each leaf -/// inside [`load_single_layer`], pools/ptr tables inside -/// `load_moe_layer_extras`). Any failure frees all staged owners — weights -/// via `WeightTensor::free_all` so AWQ sidecars go too — plus every -/// completed layer. Aliases (`lm_head`, expert views) are dropped, never -/// freed. Success takes each slot exactly once into the returned structs. -pub fn load_weights( - hfq: &mut HfqFile, - config: &Gemma4Config, - gpu: &mut Gpu, -) -> HipResult { - load_weights_impl(hfq, config, gpu, None) -} - -/// Injectable-fault twin of [`load_weights`] for GPU failure/retry -/// regressions. Production passes `None`; in-file tests pass `Some` to fail -/// the n-th staged step and prove every nested owner rolls back. -fn load_weights_impl( - hfq: &mut HfqFile, - config: &Gemma4Config, - gpu: &mut Gpu, - mut fault: Option<&mut AllocFaults>, -) -> HipResult { - let mut embed_opt: Option = None; - let mut embd_format_opt: Option = None; - // Borrowed alias of embed_tokens — staged for take-once publish, dropped - // (never freed) on rollback. - let mut lm_head_opt: Option = None; - let mut final_norm_opt: Option = None; - let mut layers: Vec = Vec::with_capacity(config.n_layers); - macro_rules! fail { - ($e:expr) => {{ - for layer in layers.drain(..) { - Gemma4Weights::free_layer(gpu, layer); - } - if let Some(t) = final_norm_opt.take() { - let _ = gpu.free_tensor(t); - } - let _ = lm_head_opt.take(); - if let Some(t) = embed_opt.take() { - let _ = gpu.free_tensor(t); - } - return Err($e); - }}; - } - macro_rules! check { - ($stage:expr) => { - if let Err(e) = fault_check(&mut fault, $stage) { - fail!(e); - } - }; - } - macro_rules! stage { - ($slot:ident, $stage:expr, $val:expr) => {{ - check!($stage); - match $val { - Ok(v) => { - $slot = Some(v); - } - Err(e) => fail!(e), - } - }}; - } - - eprintln!("gemma4: loading embed_tokens..."); - // Nothing staged yet: plain `?` cannot leak here. - fault_check(&mut fault, "embed")?; - let embed_name = "model.language_model.embed_tokens.weight"; - let (embed_info, embed_data) = hfq - .tensor_data(embed_name) - .ok_or_else(|| hip_bridge::HipError::new(0, "embed_tokens not found in HFQ"))?; - let (embed_tokens, embd_format) = match embed_info.quant_type { - 3 => { - eprintln!(" (Q8_0 / Q8F16, {} MB)", embed_data.len() / 1_000_000); - ( - upload_pooled_bytes(gpu, embed_data, &[embed_data.len()])?, - EmbeddingFormat::Q8_0, - ) - } - 6 => { - eprintln!(" (HFQ4-G256, {} MB)", embed_data.len() / 1_000_000); - ( - upload_pooled_bytes(gpu, embed_data, &[embed_data.len()])?, - EmbeddingFormat::HFQ4G256, - ) - } - 7 => { - eprintln!(" (HFQ4-G128, {} MB)", embed_data.len() / 1_000_000); - ( - upload_pooled_bytes(gpu, embed_data, &[embed_data.len()])?, - EmbeddingFormat::HFQ4G128, - ) - } - 1 => { - eprintln!(" (F16 → F32)"); - let f32_data: Vec = embed_data - .chunks_exact(2) - .map(|c| f16_to_f32(u16::from_le_bytes([c[0], c[1]]))) - .collect(); - ( - gpu.upload_f32(&f32_data, &[config.vocab_size, config.dim])?, - EmbeddingFormat::F32, - ) - } - 16 => { - eprintln!(" (BF16 → F32)"); - let f32_data: Vec = embed_data - .chunks_exact(2) - .map(|c| f32::from_bits((u16::from_le_bytes([c[0], c[1]]) as u32) << 16)) - .collect(); - ( - gpu.upload_f32(&f32_data, &[config.vocab_size, config.dim])?, - EmbeddingFormat::F32, - ) - } - 2 => { - // F32 raw (oracle / --format f32 passthrough .hfq). - eprintln!(" (F32 raw, {} MB)", embed_data.len() / 1_000_000); - let f32_data: Vec = embed_data - .chunks_exact(4) - .map(|c| f32::from_le_bytes([c[0], c[1], c[2], c[3]])) - .collect(); - ( - gpu.upload_f32(&f32_data, &[config.vocab_size, config.dim])?, - EmbeddingFormat::F32, - ) - } - qt => { - return Err(hip_bridge::HipError::new( - 0, - &format!("unsupported embed quant_type {qt}"), - )); - } - }; - embed_opt = Some(embed_tokens); - embd_format_opt = Some(embd_format); - - // Tied LM head: WeightTensor whose buffer aliases the embed allocation. - // Rollback and free_gpu skip freeing this — embed_tokens owns the bytes. - check!("lm_head"); - let lm_head = { - let alias_buf = unsafe { - embed_opt - .as_ref() - .expect("gemma4 load: embed staged") - .buf - .alias() - }; - let dtype = match embd_format_opt - .as_ref() - .expect("gemma4 load: embed format staged") - { - EmbeddingFormat::Q8_0 => DType::Q8_0, - EmbeddingFormat::HFQ4G256 => DType::HFQ4G256, - EmbeddingFormat::HFQ4G128 => DType::HFQ4G128, - EmbeddingFormat::F32 => DType::F32, - EmbeddingFormat::Q4K => DType::Q4K, - }; - let alias_tensor = GpuTensor { - buf: alias_buf, - shape: embed_opt - .as_ref() - .expect("gemma4 load: embed staged") - .shape - .clone(), - dtype, - }; - WeightTensor { - buf: alias_tensor, - gpu_dtype: dtype, - m: config.vocab_size, - k: config.dim, - row_stride: 0, - awq_scale: None, - lloyd_lut_e4m3: None, - lloyd_lut_f16: None, - lloyd_lut_c16: None, - paro: None, - } - }; - lm_head_opt = Some(lm_head); - - eprintln!("gemma4: loading final norm..."); - stage!( - final_norm_opt, - "final_norm", - load_gemma4_norm(hfq, gpu, "model.language_model.norm.weight", config.dim) - ); - - eprintln!("gemma4: loading {} layers...", config.n_layers); - for i in 0..config.n_layers { - check!("layer"); - match load_single_layer(hfq, gpu, config, i, fault.as_deref_mut()) { - Ok(layer) => layers.push(layer), - Err(e) => fail!(e), - } - } - eprintln!("gemma4: loaded all {} layers", config.n_layers); - - Ok(Gemma4Weights { - embed_tokens: embed_opt - .take() - .expect("gemma4 load: embed_tokens staged once"), - embd_format: embd_format_opt - .take() - .expect("gemma4 load: embed format staged once"), - lm_head: lm_head_opt - .take() - .expect("gemma4 load: lm_head staged once"), - final_norm: final_norm_opt - .take() - .expect("gemma4 load: final_norm staged once"), - layers, - }) -} - -/// Build one decoder layer with staged-owner rollback: every successful leaf -/// (norms, scalar, projections with their AWQ sidecars, MoE pools/ptr -/// tables) lands in a slot immediately; any later failure frees all staged -/// owners and the caller frees nothing further for this layer. Expert views -/// stay inside a successfully built `MoeLayerExtras` and are never freed. -fn load_single_layer( - hfq: &HfqFile, - gpu: &mut Gpu, - config: &Gemma4Config, - i: usize, - mut fault: Option<&mut AllocFaults>, -) -> HipResult { - let p = format!("model.language_model.layers.{i}"); - // Staged-owner slots, declared ahead of the macros so the fail!/check!/stage! - // bodies resolve them lexically (macro_rules hygiene: bare identifiers in a - // macro body resolve at the macro definition site). Both layer arms share - // these bindings — only one arm executes per call. v_opt stays None on full - // layers, which reuse k_proj's pre-norm output as V. - let mut scalar_opt: Option = None; - let mut scalar_host_opt: Option = None; - let mut moe_opt: Option = None; - let mut input_opt: Option = None; - let mut post_attn_opt: Option = None; - let mut pre_ffn_opt: Option = None; - let mut post_ffn_opt: Option = None; - let mut q_opt: Option = None; - let mut k_opt: Option = None; - let mut v_opt: Option = None; - let mut o_opt: Option = None; - let mut qn_opt: Option = None; - let mut kn_opt: Option = None; - let mut gate_opt: Option = None; - let mut up_opt: Option = None; - let mut down_opt: Option = None; - macro_rules! fail { - ($e:expr) => {{ - if let Some(m) = moe_opt.take() { - Gemma4Weights::free_moe(gpu, m); - } - for w in [ - q_opt.take(), - k_opt.take(), - v_opt.take(), - o_opt.take(), - gate_opt.take(), - up_opt.take(), - down_opt.take(), - ] - .into_iter() - .flatten() - { - w.free_all(gpu); - } - for t in [ - input_opt.take(), - post_attn_opt.take(), - pre_ffn_opt.take(), - post_ffn_opt.take(), - scalar_opt.take(), - qn_opt.take(), - kn_opt.take(), - ] - .into_iter() - .flatten() - { - let _ = gpu.free_tensor(t); - } - return Err($e); - }}; - } - macro_rules! check { - ($stage:expr) => { - if let Err(e) = fault_check(&mut fault, $stage) { - fail!(e); - } - }; - } - macro_rules! stage { - ($slot:ident, $stage:expr, $val:expr) => {{ - check!($stage); - match $val { - Ok(v) => { - $slot = Some(v); - } - Err(e) => fail!(e), - } - }}; - } - match config.layer_types[i] { - LayerType::Sliding => { - let hd = config.sliding_head_dim; - let kv_dim = config.sliding_n_kv_heads * hd; - let q_dim = config.n_heads * hd; - check!("layer_scalar"); - match load_layer_scalar(hfq, gpu, &format!("{p}.layer_scalar")) { - Ok((t, h)) => { - scalar_opt = Some(t); - scalar_host_opt = Some(h); - } - Err(e) => fail!(e), - } - if i == 0 { - eprintln!( - "[gemma4] L0 sliding layer_scalar = {}", - scalar_host_opt - .as_ref() - .expect("gemma4 load: scalar staged") - ); - } - check!("moe"); - moe_opt = if config.enable_moe_block { - match load_moe_layer_extras(hfq, gpu, &p, config, fault.as_deref_mut()) { - Ok(m) => Some(m), - Err(e) => fail!(e), - } - } else { - None - }; - stage!( - input_opt, - "input_layernorm", - load_gemma4_norm(hfq, gpu, &format!("{p}.input_layernorm.weight"), config.dim,) - ); - stage!( - post_attn_opt, - "post_attention_layernorm", - load_gemma4_norm( - hfq, - gpu, - &format!("{p}.post_attention_layernorm.weight"), - config.dim, - ) - ); - stage!( - pre_ffn_opt, - "pre_feedforward_layernorm", - load_gemma4_norm( - hfq, - gpu, - &format!("{p}.pre_feedforward_layernorm.weight"), - config.dim, - ) - ); - stage!( - post_ffn_opt, - "post_feedforward_layernorm", - load_gemma4_norm( - hfq, - gpu, - &format!("{p}.post_feedforward_layernorm.weight"), - config.dim, - ) - ); - stage!( - q_opt, - "q_proj", - load_gemma4_weight( - hfq, - gpu, - &format!("{p}.self_attn.q_proj.weight"), - q_dim, - config.dim, - ) - ); - stage!( - k_opt, - "k_proj", - load_gemma4_weight( - hfq, - gpu, - &format!("{p}.self_attn.k_proj.weight"), - kv_dim, - config.dim, - ) - ); - stage!( - v_opt, - "v_proj", - load_gemma4_weight( - hfq, - gpu, - &format!("{p}.self_attn.v_proj.weight"), - kv_dim, - config.dim, - ) - ); - stage!( - o_opt, - "o_proj", - load_gemma4_weight( - hfq, - gpu, - &format!("{p}.self_attn.o_proj.weight"), - config.dim, - q_dim, - ) - ); - stage!( - qn_opt, - "q_norm", - load_gemma4_head_norm(hfq, gpu, &format!("{p}.self_attn.q_norm.weight"), hd,) - ); - stage!( - kn_opt, - "k_norm", - load_gemma4_head_norm(hfq, gpu, &format!("{p}.self_attn.k_norm.weight"), hd,) - ); - stage!( - gate_opt, - "gate_proj", - load_gemma4_weight( - hfq, - gpu, - &format!("{p}.mlp.gate_proj.weight"), - config.hidden_dim, - config.dim, - ) - ); - stage!( - up_opt, - "up_proj", - load_gemma4_weight( - hfq, - gpu, - &format!("{p}.mlp.up_proj.weight"), - config.hidden_dim, - config.dim, - ) - ); - stage!( - down_opt, - "down_proj", - load_gemma4_weight( - hfq, - gpu, - &format!("{p}.mlp.down_proj.weight"), - config.dim, - config.hidden_dim, - ) - ); - Ok(LayerWeights::Sliding(SlidingLayerWeights { - input_layernorm: input_opt.take().expect("gemma4 load: input staged once"), - post_attention_layernorm: post_attn_opt - .take() - .expect("gemma4 load: post-attn staged once"), - pre_feedforward_layernorm: pre_ffn_opt - .take() - .expect("gemma4 load: pre-ffn staged once"), - post_feedforward_layernorm: post_ffn_opt - .take() - .expect("gemma4 load: post-ffn staged once"), - layer_scalar: scalar_opt.take().expect("gemma4 load: scalar staged once"), - layer_scalar_host: scalar_host_opt - .take() - .expect("gemma4 load: scalar host staged"), - q_proj: q_opt.take().expect("gemma4 load: q_proj staged once"), - k_proj: k_opt.take().expect("gemma4 load: k_proj staged once"), - v_proj: v_opt.take().expect("gemma4 load: v_proj staged once"), - o_proj: o_opt.take().expect("gemma4 load: o_proj staged once"), - q_norm: qn_opt.take().expect("gemma4 load: q_norm staged once"), - k_norm: kn_opt.take().expect("gemma4 load: k_norm staged once"), - gate_proj: gate_opt.take().expect("gemma4 load: gate staged once"), - up_proj: up_opt.take().expect("gemma4 load: up staged once"), - down_proj: down_opt.take().expect("gemma4 load: down staged once"), - moe: moe_opt.take(), - })) - } - LayerType::Full => { - let hd = config.full_head_dim; - let kv_dim = config.full_n_kv_heads * hd; - let q_dim = config.n_heads * hd; - // v_opt (declared above) stays None here: full layers reuse k_proj's - // pre-norm output as V, so the slot stages and frees nothing. - check!("layer_scalar"); - match load_layer_scalar(hfq, gpu, &format!("{p}.layer_scalar")) { - Ok((t, h)) => { - scalar_opt = Some(t); - scalar_host_opt = Some(h); - } - Err(e) => fail!(e), - } - if i <= 6 { - eprintln!( - "[gemma4] L{i} full layer_scalar = {}", - scalar_host_opt - .as_ref() - .expect("gemma4 load: scalar staged") - ); - } - check!("moe"); - moe_opt = if config.enable_moe_block { - match load_moe_layer_extras(hfq, gpu, &p, config, fault.as_deref_mut()) { - Ok(m) => Some(m), - Err(e) => fail!(e), - } - } else { - None - }; - stage!( - input_opt, - "input_layernorm", - load_gemma4_norm(hfq, gpu, &format!("{p}.input_layernorm.weight"), config.dim,) - ); - stage!( - post_attn_opt, - "post_attention_layernorm", - load_gemma4_norm( - hfq, - gpu, - &format!("{p}.post_attention_layernorm.weight"), - config.dim, - ) - ); - stage!( - pre_ffn_opt, - "pre_feedforward_layernorm", - load_gemma4_norm( - hfq, - gpu, - &format!("{p}.pre_feedforward_layernorm.weight"), - config.dim, - ) - ); - stage!( - post_ffn_opt, - "post_feedforward_layernorm", - load_gemma4_norm( - hfq, - gpu, - &format!("{p}.post_feedforward_layernorm.weight"), - config.dim, - ) - ); - stage!( - q_opt, - "q_proj", - load_gemma4_weight( - hfq, - gpu, - &format!("{p}.self_attn.q_proj.weight"), - q_dim, - config.dim, - ) - ); - stage!( - k_opt, - "k_proj", - load_gemma4_weight( - hfq, - gpu, - &format!("{p}.self_attn.k_proj.weight"), - kv_dim, - config.dim, - ) - ); - // no v_proj on full layers — V reuses k_proj's pre-norm output. - stage!( - o_opt, - "o_proj", - load_gemma4_weight( - hfq, - gpu, - &format!("{p}.self_attn.o_proj.weight"), - config.dim, - q_dim, - ) - ); - stage!( - qn_opt, - "q_norm", - load_gemma4_head_norm(hfq, gpu, &format!("{p}.self_attn.q_norm.weight"), hd,) - ); - stage!( - kn_opt, - "k_norm", - load_gemma4_head_norm(hfq, gpu, &format!("{p}.self_attn.k_norm.weight"), hd,) - ); - // no v_norm weight — v_norm is no-scale (ones buffer passed at decode time). - stage!( - gate_opt, - "gate_proj", - load_gemma4_weight( - hfq, - gpu, - &format!("{p}.mlp.gate_proj.weight"), - config.hidden_dim, - config.dim, - ) - ); - stage!( - up_opt, - "up_proj", - load_gemma4_weight( - hfq, - gpu, - &format!("{p}.mlp.up_proj.weight"), - config.hidden_dim, - config.dim, - ) - ); - stage!( - down_opt, - "down_proj", - load_gemma4_weight( - hfq, - gpu, - &format!("{p}.mlp.down_proj.weight"), - config.dim, - config.hidden_dim, - ) - ); - // Full layers never load v_proj: the slot stays None and frees nothing. - debug_assert!(v_opt.is_none(), "gemma4 load: full layers stage no v_proj"); - Ok(LayerWeights::Full(FullLayerWeights { - input_layernorm: input_opt.take().expect("gemma4 load: input staged once"), - post_attention_layernorm: post_attn_opt - .take() - .expect("gemma4 load: post-attn staged once"), - pre_feedforward_layernorm: pre_ffn_opt - .take() - .expect("gemma4 load: pre-ffn staged once"), - post_feedforward_layernorm: post_ffn_opt - .take() - .expect("gemma4 load: post-ffn staged once"), - layer_scalar: scalar_opt.take().expect("gemma4 load: scalar staged once"), - layer_scalar_host: scalar_host_opt - .take() - .expect("gemma4 load: scalar host staged"), - q_proj: q_opt.take().expect("gemma4 load: q_proj staged once"), - k_proj: k_opt.take().expect("gemma4 load: k_proj staged once"), - // no v_proj — V = pre-k_norm output of k_proj - o_proj: o_opt.take().expect("gemma4 load: o_proj staged once"), - q_norm: qn_opt.take().expect("gemma4 load: q_norm staged once"), - k_norm: kn_opt.take().expect("gemma4 load: k_norm staged once"), - gate_proj: gate_opt.take().expect("gemma4 load: gate staged once"), - up_proj: up_opt.take().expect("gemma4 load: up staged once"), - down_proj: down_opt.take().expect("gemma4 load: down staged once"), - moe: moe_opt.take(), - })) - } - } -} - -/// One-time init for the scratch buffers that must hold a constant value -/// across forward passes (notably the ones-filled `v_norm_ones_full`). -/// Call once after `Gemma4Scratch::new` before the first forward pass. -pub fn init_scratch_constants( - gpu: &mut Gpu, - scratch: &Gemma4Scratch, - full_head_dim: usize, -) -> HipResult<()> { - let ones: Vec = vec![1.0; full_head_dim]; - let bytes: &[u8] = - unsafe { std::slice::from_raw_parts(ones.as_ptr() as *const u8, ones.len() * 4) }; - gpu.hip.memcpy_htod(&scratch.v_norm_ones_full.buf, bytes)?; - Ok(()) -} - -// ─── Scratch ──────────────────────────────────────────────────────────── - -use hip_bridge::DeviceBuffer; -/// Flash tile size for gemma4 lowered path. Matches the HIP partition kernel's -/// `TILE_SIZE = 128`. -pub const GEMMA4_FLASH_TILE: usize = 128; - -/// Pure geometry: single-query flash partial length for `max_seq`. -/// `n_heads * ceil(max_seq / TILE) * (2 + head_dim)` floats. -#[inline] -pub fn gemma4_flash_partials_len(max_seq: usize, n_heads: usize, full_head_dim: usize) -> usize { - let tiles = max_seq.div_ceil(GEMMA4_FLASH_TILE); - n_heads * tiles * (2 + full_head_dim) -} - -/// Convenience for `Gemma4Config`. -#[inline] -pub fn gemma4_flash_partials_len_for_config(max_seq: usize, config: &Gemma4Config) -> usize { - gemma4_flash_partials_len(max_seq, config.n_heads, config.full_head_dim) -} - -/// Single-query flash partial length that covers every decode attend tier the -/// lowered path can run on `arch`: the asym3 hd512 tile ([`GEMMA4_FLASH_TILE`]) -/// and the Q8 decode tile (`q8_flash_tile_size`) for the full tier at `max_seq` -/// and the sliding ring at `min(sliding_window, max_seq)`. The Q8 tile can be -/// smaller than 128 (tile32 on gfx1100 up to 8K), which needs more partials. -pub fn gemma4_decode_flash_partials_len( - arch: &str, - config: &Gemma4Config, - max_seq: usize, -) -> usize { - let q8_len = |n_kv_heads: usize, head_dim: usize, cap: usize| { - let tile = rdna_compute::attention::q8_flash_tile_size( - arch, - config.n_heads, - n_kv_heads, - head_dim, - cap, - ); - config.n_heads * cap.div_ceil(tile) * (2 + head_dim) - }; - let sliding_cap = config.sliding_window.min(max_seq); - gemma4_flash_partials_len_for_config(max_seq, config) - .max(q8_len( - config.full_n_kv_heads, - config.full_head_dim, - max_seq, - )) - .max(q8_len( - config.sliding_n_kv_heads, - config.sliding_head_dim, - sliding_cap, - )) -} - -/// Per-decode scratch, sized once at model-load time against the MAX of -/// sliding and full attention dimensions so a single buffer works across -/// layer types. 31B target shapes: sliding Q=[32*256]=8192, full Q=[32*512]=16384 -/// → size Q at 16384. Sliding KV=[16*256]=4096, full KV=[4*512]=2048 → size at 4096. -pub struct Gemma4Scratch { - pub x: GpuTensor, // [dim] — hidden state - pub residual: GpuTensor, // [dim] — saved for sandwich residual - pub tmp: GpuTensor, // [dim] — norm output scratch - /// `[dim]` FWHT-rotated normed activation shared by MagnumQuant projections. - pub x_rot: GpuTensor, - - /// Position buffer (single i32 on device, updated per decode step). - pub pos_buf: DeviceBuffer, - - // Attention scratch — sized for max(sliding, full) - pub q: GpuTensor, // [max(n_heads*head_dim_sliding, n_heads*head_dim_full)] - pub k: GpuTensor, // [max(n_kv_heads*head_dim for each layer type)] - pub v: GpuTensor, // [same as k] - pub attn_out: GpuTensor, // [same as q] - - // MLP scratch - pub gate_ffn: GpuTensor, // [hidden_dim] - pub up_ffn: GpuTensor, // [hidden_dim] - pub ffn_hidden: GpuTensor, // [hidden_dim] - pub ffn_out: GpuTensor, // [dim] - - // Output - pub logits: GpuTensor, // [vocab_size] - pub sample_buf: GpuTensor, // [2] — (token_id, new_rng_state) for GPU sampling - pub repeat_buf: GpuTensor, // [1024] — rolling window for repeat penalty - - // Flash attention tile partials. Sized by `gemma4_decode_flash_partials_len` - // for the LARGER of the cache shapes: full-attn head_dim=512 at the asym3 - // tile (128) or the arch's Q8 decode tile, sliding head_dim=256 over the ring. - pub flash_partials: GpuTensor, - - // No-scale v_norm ones buffer (full-attn layers compute v_norm without - // a learned weight — we pass this ones-filled tensor to the existing - // rmsnorm kernel to get no-scale RMS semantics). - pub v_norm_ones_full: GpuTensor, // [full_head_dim] - - // ── MoE scratch (26B-A4B only). Zero-sized on dense models. ───────── - pub moe_cur_mlp: GpuTensor, // [dim] — rmsnorm(ffn_out, post_norm_1) - pub moe_pre2: GpuTensor, // [dim] — rmsnorm(attn_out, pre_norm_2) - pub moe_router_in: GpuTensor, // [dim] — router input (post-rmsnorm + scale) - pub moe_router_logits: GpuTensor, // [n_experts] - pub moe_topk_indices: GpuTensor, // [top_k_experts] — i32 packed in f32 slots - pub moe_topk_weights: GpuTensor, // [top_k_experts] - pub moe_cur_moe: GpuTensor, // [dim] — accumulator across top-K experts - pub moe_expert_gate_up: GpuTensor, // [2 * moe_intermediate_size] - pub moe_expert_hidden: GpuTensor, // [moe_intermediate_size] — gelu(gate) * up - pub moe_expert_out: GpuTensor, // [dim] — single expert's down_proj output - - // ── Indexed MoE batched scratch (k_top=8 hardcoded by kernel). ────── - // These back the device-side fused path that replaces the 8-iteration - // per-expert CPU loop. Only allocated when the MoE branch is enabled. - /// `[dim]` — moe_pre2 after one FWHT pass (MQ4 gate_up expects pre-rotated x). - pub moe_pre2_rot: GpuTensor, - /// `[k_top × 2 × mi]` — fused gate+up output, one row per top-K rank. - pub moe_expert_gate_batch: GpuTensor, // [k_top × mi] - pub moe_expert_up_batch: GpuTensor, // [k_top × mi] - - /// `[k_top × mi]` — gelu_tanh(gate)*up batched over k_top experts. - pub moe_expert_hidden_batch: GpuTensor, -} - -impl Gemma4Scratch { - /// `max_seq` is the sole allocation authority — MUST equal `LoadCtx::max_seq` - /// and the `KvCache::max_seq` / `physical_cap` for the paired caches. - /// `flash_partials` is sized from this single value; the `HIPFIRE_KV_SEQ` - /// env var is no longer consulted. - pub fn new(gpu: &mut Gpu, config: &Gemma4Config, max_seq: usize) -> HipResult { - Self::new_with_alloc( - gpu, - config, - max_seq, - |gpu, shape, dtype| gpu.zeros(shape, dtype), - |gpu| gpu.hip.malloc(4), - ) - } - - /// Injectable-allocator twin of [`new`] for GPU failure/retry regressions: - /// tests fail the n-th allocation and prove every staged owner (tensors + - /// `pos_buf`) rolls back into the pool for immediate retry. Mirrors - /// qwen35's `new_opt_with_alloc` seam. No env knob, no global sweep. - fn new_with_alloc( - gpu: &mut Gpu, - config: &Gemma4Config, - max_seq: usize, - mut alloc: impl FnMut(&mut Gpu, &[usize], DType) -> HipResult, - mut alloc_pos: impl FnMut(&mut Gpu) -> HipResult, - ) -> HipResult { - // Library code must not abort hosts: reject tiny contexts as an error - // BEFORE any allocation (admission refuses these up front; direct - // `new` callers in examples/tools get a clean Err, not a panic). - if max_seq < GEMMA4_FLASH_TILE { - return Err(hip_bridge::HipError::new( - 0, - &format!( - "gemma4 scratch: max_seq {max_seq} too small (minimum one flash tile = 128)" - ), - )); - } - let dim = config.dim; - let q_dim = - (config.n_heads * config.sliding_head_dim).max(config.n_heads * config.full_head_dim); - let kv_dim = (config.sliding_n_kv_heads * config.sliding_head_dim) - .max(config.full_n_kv_heads * config.full_head_dim); - - // Transactional construction: slots own every tensor until take-once - // publish, pos_slot owns the raw position buffer. Any failure - // reverse-drains all staged owners (GpuTensor/DeviceBuffer have no - // Drop that could release device memory for us). - let mut slots: Vec> = Vec::with_capacity(64); - let mut pos_slot: Option = None; - macro_rules! rollback { - () => {{ - while let Some(slot) = slots.pop() { - if let Some(t) = slot { - let _ = gpu.free_tensor(t); - } - } - if let Some(pos) = pos_slot.take() { - let _ = gpu.hip.free(pos); - } - }}; - } - macro_rules! alloc { - ($shape:expr, $dt:expr) => {{ - match alloc(gpu, $shape, $dt) { - Ok(t) => { - slots.push(Some(t)); - slots.len() - 1 - } - Err(e) => { - rollback!(); - return Err(e); - } - } - }}; - } - macro_rules! take { - ($i:expr) => { - slots[$i].take().expect("gemma4 scratch slot taken twice") - }; - } - - let i_x = alloc!(&[dim], DType::F32); - let i_residual = alloc!(&[dim], DType::F32); - let i_tmp = alloc!(&[dim], DType::F32); - let i_x_rot = alloc!(&[dim], DType::F32); - - match alloc_pos(gpu) { - Ok(pos) => { - pos_slot = Some(pos); - } - Err(e) => { - rollback!(); - return Err(e); - } - } - - let i_q = alloc!(&[q_dim], DType::F32); - let i_k = alloc!(&[kv_dim], DType::F32); - let i_v = alloc!(&[kv_dim], DType::F32); - let i_attn_out = alloc!(&[q_dim], DType::F32); - - let i_gate_ffn = alloc!(&[config.hidden_dim], DType::F32); - let i_up_ffn = alloc!(&[config.hidden_dim], DType::F32); - let i_ffn_hidden = alloc!(&[config.hidden_dim], DType::F32); - let i_ffn_out = alloc!(&[dim], DType::F32); - let i_logits = alloc!(&[config.vocab_size], DType::F32); - let i_sample_buf = alloc!(&[2], DType::F32); - let i_repeat_buf = alloc!(&[1024], DType::F32); - - // Flash partials sizing. Per-head × max_tiles × (2 + head_dim) floats, - // covering the full tier (head_dim=512, stride 514) in either KV format - // and the sliding q8 ring (256, stride 258) at this arch's Q8 decode - // tile. `max_seq` is the single authority shared with both KV caches — - // no independent `HIPFIRE_KV_SEQ` env var. - let flash_partials_sz = gemma4_decode_flash_partials_len(&gpu.arch, config, max_seq); - let i_flash_partials = alloc!(&[flash_partials_sz], DType::F32); - - // (Note 2026-05-19): removed the precomputed sliding/full cos+sin - // tables that were allocated here but never read by any kernel. - // `rope_f32` and `rope_partial_halved_f32` compute cos/sin inline - // from `pos_buf[0]` + `freq_base`; the lookup-table path was wired - // but never finished. The tables ate 4 * max_kv_seq * head_dim - // floats — 768 MB at max_kv_seq=131072 — which pushed the full KV - // cache (~970 MB at 128k) out of VRAM into GTT (PCIe-paged system - // RAM), causing attention reads to take a slow path. - - // v_norm ones — populated on first use in the forward pass. - // Allocated up to the LARGER of sliding_head_dim and full_head_dim - // since both sliding and full apply no-scale v_norm (the post-rebase - // fix added v_norm to sliding_layer_decode; sliding head_dim=256, - // full head_dim=512 → max=512 covers both). - let v_norm_max = config.sliding_head_dim.max(config.full_head_dim); - let i_v_norm_ones_full = alloc!(&[v_norm_max], DType::F32); - - // MoE scratch. Allocated unconditionally because the buffers are tiny - // relative to the model; zero-sized on dense models would just complicate - // the dispatch path. Sized for 26B-A4B: n_experts=128, top_k=8, mi=704. - let n_exp = config.num_experts.max(1); - let mi = config.moe_intermediate_size.max(1); - let k_top = config.top_k_experts.max(1); - let i_moe_cur_mlp = alloc!(&[dim], DType::F32); - let i_moe_pre2 = alloc!(&[dim], DType::F32); - let i_moe_router_in = alloc!(&[dim], DType::F32); - let i_moe_router_logits = alloc!(&[n_exp], DType::F32); - let i_moe_topk_indices = alloc!(&[k_top], DType::F32); - let i_moe_topk_weights = alloc!(&[k_top], DType::F32); - let i_moe_cur_moe = alloc!(&[dim], DType::F32); - let i_moe_expert_gate_up = alloc!(&[2 * mi], DType::F32); - let i_moe_expert_hidden = alloc!(&[mi], DType::F32); - let i_moe_expert_out = alloc!(&[dim], DType::F32); - - // Indexed-MoE scratch (k_top fixed at 8 by the kernel). - let i_moe_pre2_rot = alloc!(&[dim], DType::F32); - let i_moe_expert_gate_batch = alloc!(&[k_top * mi], DType::F32); - let i_moe_expert_up_batch = alloc!(&[k_top * mi], DType::F32); - let i_moe_expert_hidden_batch = alloc!(&[k_top * mi], DType::F32); - - let scratch = Gemma4Scratch { - x: take!(i_x), - residual: take!(i_residual), - tmp: take!(i_tmp), - x_rot: take!(i_x_rot), - pos_buf: pos_slot.take().expect("gemma4 scratch pos_buf staged once"), - q: take!(i_q), - k: take!(i_k), - v: take!(i_v), - attn_out: take!(i_attn_out), - gate_ffn: take!(i_gate_ffn), - up_ffn: take!(i_up_ffn), - ffn_hidden: take!(i_ffn_hidden), - ffn_out: take!(i_ffn_out), - logits: take!(i_logits), - sample_buf: take!(i_sample_buf), - repeat_buf: take!(i_repeat_buf), - flash_partials: take!(i_flash_partials), - v_norm_ones_full: take!(i_v_norm_ones_full), - moe_cur_mlp: take!(i_moe_cur_mlp), - moe_pre2: take!(i_moe_pre2), - moe_router_in: take!(i_moe_router_in), - moe_router_logits: take!(i_moe_router_logits), - moe_topk_indices: take!(i_moe_topk_indices), - moe_topk_weights: take!(i_moe_topk_weights), - moe_cur_moe: take!(i_moe_cur_moe), - moe_expert_gate_up: take!(i_moe_expert_gate_up), - moe_expert_hidden: take!(i_moe_expert_hidden), - moe_expert_out: take!(i_moe_expert_out), - moe_pre2_rot: take!(i_moe_pre2_rot), - moe_expert_gate_batch: take!(i_moe_expert_gate_batch), - moe_expert_up_batch: take!(i_moe_expert_up_batch), - moe_expert_hidden_batch: take!(i_moe_expert_hidden_batch), - }; - debug_assert!( - slots.iter().all(|s| s.is_none()) && pos_slot.is_none(), - "gemma4 scratch: every staged owner published exactly once" - ); - Ok(scratch) - } - - /// Release every GPU allocation owned by this scratch. Mirrors the - /// Qwen35Scratch / LlamaScratch pattern so `unload_model` in the daemon - /// can reclaim VRAM on idle eviction. - pub fn free_gpu(self, gpu: &mut Gpu) { - let _ = gpu.free_tensor(self.x); - let _ = gpu.free_tensor(self.residual); - let _ = gpu.free_tensor(self.tmp); - let _ = gpu.free_tensor(self.x_rot); - // pos_buf is a raw DeviceBuffer (not a GpuTensor): free it explicitly — - // DeviceBuffer has no Drop-side free. Eager precedent: gemma4.rs free_gpu. - let _ = gpu.hip.free(self.pos_buf); - let _ = gpu.free_tensor(self.q); - let _ = gpu.free_tensor(self.k); - let _ = gpu.free_tensor(self.v); - let _ = gpu.free_tensor(self.attn_out); - let _ = gpu.free_tensor(self.gate_ffn); - let _ = gpu.free_tensor(self.up_ffn); - let _ = gpu.free_tensor(self.ffn_hidden); - let _ = gpu.free_tensor(self.ffn_out); - let _ = gpu.free_tensor(self.logits); - let _ = gpu.free_tensor(self.sample_buf); - let _ = gpu.free_tensor(self.repeat_buf); - let _ = gpu.free_tensor(self.flash_partials); - let _ = gpu.free_tensor(self.v_norm_ones_full); - let _ = gpu.free_tensor(self.moe_cur_mlp); - let _ = gpu.free_tensor(self.moe_pre2); - let _ = gpu.free_tensor(self.moe_router_in); - let _ = gpu.free_tensor(self.moe_router_logits); - let _ = gpu.free_tensor(self.moe_topk_indices); - let _ = gpu.free_tensor(self.moe_topk_weights); - let _ = gpu.free_tensor(self.moe_cur_moe); - let _ = gpu.free_tensor(self.moe_expert_gate_up); - let _ = gpu.free_tensor(self.moe_expert_hidden); - let _ = gpu.free_tensor(self.moe_expert_out); - let _ = gpu.free_tensor(self.moe_pre2_rot); - let _ = gpu.free_tensor(self.moe_expert_gate_batch); - let _ = gpu.free_tensor(self.moe_expert_up_batch); - let _ = gpu.free_tensor(self.moe_expert_hidden_batch); - } -} -#[cfg(test)] -mod scratch_geometry_tests { - use super::*; - - #[test] - fn flash_partials_geometry_matches_formula() { - // Single tile edge - assert_eq!(gemma4_flash_partials_len(128, 32, 512), 32 * 1 * 514); - assert_eq!(gemma4_flash_partials_len(129, 32, 512), 32 * 2 * 514); - assert_eq!(gemma4_flash_partials_len(256, 32, 512), 32 * 2 * 514); - // 32k baseline (FALLBACK_KV_SEQ before fix) - assert_eq!(gemma4_flash_partials_len(32768, 32, 512), 32 * 256 * 514); - assert_eq!(gemma4_flash_partials_len(32768, 32, 512), 4_210_688); - // 131072 must be exactly 4× the 32768 geometry (no overflow, no env var) - assert_eq!(gemma4_flash_partials_len(131072, 32, 512), 32 * 1024 * 514); - assert_eq!(gemma4_flash_partials_len(131072, 32, 512), 16_842_752); - assert_eq!( - gemma4_flash_partials_len(131072, 32, 512), - 4 * gemma4_flash_partials_len(32768, 32, 512) - ); - } - - #[test] - fn flash_scales_invariant_to_tile_rounding() { - // Non-multiple of 128 must ceil - let tiles_32769 = 32769usize.div_ceil(GEMMA4_FLASH_TILE); - assert_eq!(tiles_32769, 257); - assert_eq!(gemma4_flash_partials_len(32769, 32, 512), 32 * 257 * 514); - // 131071 is one short of 131072 -> still 1024 tiles (ceil) - assert_eq!(131071usize.div_ceil(GEMMA4_FLASH_TILE), 1024); - assert_eq!( - gemma4_flash_partials_len(131071, 32, 512), - gemma4_flash_partials_len(131072, 32, 512) - ); - } - - fn dummy_cfg_31b() -> Gemma4Config { - // Minimal config mirroring 31B/26B shapes: n_heads=32, full_head_dim=512 - Gemma4Config { - dim: 5376, - n_layers: 40, - vocab_size: 262144, - norm_eps: 1e-6, - bos_token: 2, - eos_token: 1, - pad_token: 0, - n_heads: 32, - sliding_head_dim: 256, - sliding_n_kv_heads: 16, - sliding_rope_theta: 10000.0, - sliding_window: 1024, - full_head_dim: 512, - full_n_kv_heads: 4, - full_rope_theta: 1_000_000.0, - full_rope_type: RopeType::Proportional, - full_partial_rotary_factor: 0.25, - attention_k_eq_v: true, - hidden_dim: 21504, - enable_moe_block: false, - moe_intermediate_size: 704, - num_experts: 128, - top_k_experts: 8, - final_logit_softcapping: 30.0, - tie_word_embeddings: true, - embed_scale: (5376 as f32).sqrt(), - layer_types: vec![LayerType::Sliding; 40], - has_vision: false, - image_token_id: 258880, - boi_token_id: 255999, - eoi_token_id: 258882, - audio_token_id: 258881, - video_token_id: 258884, - } - } - - #[test] - fn hfq4g128_rows_must_start_fresh_groups() { - // K = 704: 5.5 groups per row. The 2-D packer pads each row to 6 groups. - assert!(check_hfq4g128_rows("w", 2816, 704, 2816 * 6 * 72).is_ok()); - // The pre-b4846285e flat packer straddles rows: refused with its cause. - let err = check_hfq4g128_rows("w", 2816, 704, 15488 * 72).unwrap_err(); - assert!(err.to_string().contains("across rows"), "{err}"); - assert!(check_hfq4g128_rows("w", 2816, 704, 2816 * 6 * 72 - 72).is_err()); - // Group-aligned K: both packers agree. - assert!(check_hfq4g128_rows("w", 4, 256, 4 * 2 * 72).is_ok()); - } - - #[test] - fn decode_partials_cover_q8_decode_tile() { - let cfg = dummy_cfg_31b(); - // gfx1100 decodes Q8 with tile32 up to 8K: 4x the tile-128 partials. - assert_eq!( - gemma4_decode_flash_partials_len("gfx1100", &cfg, 8192), - 32 * (8192 / 32) * 514 - ); - assert_eq!( - gemma4_decode_flash_partials_len("gfx1100", &cfg, 8192), - 4 * gemma4_flash_partials_len_for_config(8192, &cfg) - ); - // Tile128 arches keep the tile-128 geometry. - for arch in ["gfx1201", "gfx1151"] { - for max_seq in [128usize, 8192, 32768] { - assert_eq!( - gemma4_decode_flash_partials_len(arch, &cfg, max_seq), - gemma4_flash_partials_len_for_config(max_seq, &cfg) - ); - } - } - // Past 8K gfx1100 is tile128 too. - assert_eq!( - gemma4_decode_flash_partials_len("gfx1100", &cfg, 32768), - gemma4_flash_partials_len_for_config(32768, &cfg) - ); - // Small contexts: the sliding q8 ring (tile32) must fit as well. - let small = gemma4_decode_flash_partials_len("gfx1100", &cfg, 1000); - assert!(small >= 32 * 1000usize.div_ceil(32) * 258); - assert!(small >= 32 * 1000usize.div_ceil(32) * 514); - } -} -#[cfg(test)] -mod load_rollback_tests { - use super::*; - use hipfire_runtime::hfq::{write_hfqm_package_mem, HfqFile, HfqMemTensor}; - use std::cell::Cell; - - fn try_gpu() -> Option { - match Gpu::init() { - Ok(g) => Some(g), - Err(e) => { - eprintln!("skip: no GPU ({e:?})"); - None - } - } - } - - fn tiny_config(layers: Vec, moe: bool) -> Gemma4Config { - Gemma4Config { - dim: 16, - n_layers: layers.len(), - vocab_size: 16, - norm_eps: 1e-6, - bos_token: 2, - eos_token: 1, - pad_token: 0, - n_heads: 2, - sliding_head_dim: 8, - sliding_n_kv_heads: 2, - sliding_rope_theta: 10_000.0, - sliding_window: 32, - full_head_dim: 8, - full_n_kv_heads: 2, - full_rope_theta: 1_000_000.0, - full_rope_type: RopeType::Proportional, - full_partial_rotary_factor: 0.25, - attention_k_eq_v: true, - hidden_dim: 32, - enable_moe_block: moe, - moe_intermediate_size: 8, - num_experts: if moe { 2 } else { 0 }, - top_k_experts: if moe { 2 } else { 0 }, - final_logit_softcapping: 30.0, - tie_word_embeddings: true, - embed_scale: 4.0, - layer_types: layers, - has_vision: false, - image_token_id: 0, - boi_token_id: 0, - eoi_token_id: 0, - audio_token_id: 0, - video_token_id: 0, - } - } - - fn mem(name: String, quant_type: u8, shape: Vec, data: Vec) -> HfqMemTensor { - HfqMemTensor { - name, - quant_type, - shape, - group_size: 0, - data, - } - } - - fn f32_bytes(vals: &[f32]) -> Vec { - vals.iter().flat_map(|v| v.to_le_bytes()).collect() - } - - fn ones_f32(n: usize) -> Vec { - f32_bytes(&vec![1.0f32; n]) - } - - /// Dense projections as MQ4G256 blobs (quant_type 13) with F16 AWQ sidecars - /// so rollback/unload must reclaim sidecars too — the pre-fix free_gpu - /// leaked them via free_tensor-on-buf. - fn proj_entries(p: &str, rel: &str, m: usize, k: usize) -> Vec { - let name = format!("{p}.{rel}"); - let awq_name = match name.strip_suffix(".weight") { - Some(stem) => format!("{stem}.awq_scale.weight"), - None => format!("{name}.awq_scale.weight"), - }; - vec![ - mem(name, 13, vec![m as u32, k as u32], vec![0u8; 256]), - mem(awq_name, 1, vec![k as u32], vec![0u8; k * 2]), - ] - } - - fn dense_layer_entries(cfg: &Gemma4Config, i: usize) -> Vec { - let p = format!("model.language_model.layers.{i}"); - let mut v = vec![ - mem(format!("{p}.layer_scalar"), 2, vec![1], f32_bytes(&[1.0])), - mem( - format!("{p}.input_layernorm.weight"), - 2, - vec![cfg.dim as u32], - ones_f32(cfg.dim), - ), - mem( - format!("{p}.post_attention_layernorm.weight"), - 2, - vec![cfg.dim as u32], - ones_f32(cfg.dim), - ), - mem( - format!("{p}.pre_feedforward_layernorm.weight"), - 2, - vec![cfg.dim as u32], - ones_f32(cfg.dim), - ), - mem( - format!("{p}.post_feedforward_layernorm.weight"), - 2, - vec![cfg.dim as u32], - ones_f32(cfg.dim), - ), - ]; - let (hd, n_kv, has_v) = match cfg.layer_types[i] { - LayerType::Sliding => (cfg.sliding_head_dim, cfg.sliding_n_kv_heads, true), - LayerType::Full => (cfg.full_head_dim, cfg.full_n_kv_heads, false), - }; - let kv_dim = n_kv * hd; - let q_dim = cfg.n_heads * hd; - let mut projs = vec![ - ("self_attn.q_proj.weight", q_dim, cfg.dim), - ("self_attn.k_proj.weight", kv_dim, cfg.dim), - ]; - if has_v { - projs.push(("self_attn.v_proj.weight", kv_dim, cfg.dim)); - } - projs.extend([ - ("self_attn.o_proj.weight", cfg.dim, q_dim), - ("mlp.gate_proj.weight", cfg.hidden_dim, cfg.dim), - ("mlp.up_proj.weight", cfg.hidden_dim, cfg.dim), - ("mlp.down_proj.weight", cfg.dim, cfg.hidden_dim), - ]); - for (rel, m, k) in projs { - v.extend(proj_entries(&p, rel, m, k)); - } - v.push(mem( - format!("{p}.self_attn.q_norm.weight"), - 2, - vec![hd as u32], - ones_f32(hd), - )); - v.push(mem( - format!("{p}.self_attn.k_norm.weight"), - 2, - vec![hd as u32], - ones_f32(hd), - )); - v - } - - fn moe_entries(cfg: &Gemma4Config, i: usize) -> Vec { - let p = format!("model.language_model.layers.{i}"); - let n_exp = cfg.num_experts; - let dim = cfg.dim; - let mut v = proj_entries(&p, "router.proj.weight", n_exp, dim); - v.push(mem( - format!("{p}.router.scale"), - 2, - vec![dim as u32], - ones_f32(dim), - )); - v.push(mem( - format!("{p}.router.per_expert_scale"), - 2, - vec![n_exp as u32], - ones_f32(n_exp), - )); - for rel in [ - "pre_feedforward_layernorm_2.weight", - "post_feedforward_layernorm_1.weight", - "post_feedforward_layernorm_2.weight", - ] { - v.push(mem( - format!("{p}.{rel}"), - 2, - vec![dim as u32], - ones_f32(dim), - )); - } - // Q8_0 expert blobs: content is never executed, only uploaded into - // the pools; equal sizes keep the pool concat path happy. - for x in 0..n_exp { - v.push(mem( - format!("{p}.experts.{x}.gate_up_proj.weight"), - 3, - vec![64], - vec![0u8; 64], - )); - v.push(mem( - format!("{p}.experts.{x}.down_proj.weight"), - 3, - vec![64], - vec![0u8; 64], - )); - } - v - } - - fn fixture_tensors(cfg: &Gemma4Config) -> Vec { - let mut v = vec![ - mem( - "model.language_model.embed_tokens.weight".to_string(), - 2, - vec![cfg.vocab_size as u32, cfg.dim as u32], - ones_f32(cfg.vocab_size * cfg.dim), - ), - mem( - "model.language_model.norm.weight".to_string(), - 2, - vec![cfg.dim as u32], - ones_f32(cfg.dim), - ), - ]; - for i in 0..cfg.n_layers { - v.extend(dense_layer_entries(cfg, i)); - if cfg.enable_moe_block { - v.extend(moe_entries(cfg, i)); - } - } - v - } - - fn open_fixture(tensors: Vec, tag: &str) -> (HfqFile, std::path::PathBuf) { - let path = - std::env::temp_dir().join(format!("gemma4_rollback_{tag}_{}.hfq", std::process::id())); - write_hfqm_package_mem(&path, 13, "{}", &tensors).expect("write fixture hfq"); - let hfq = HfqFile::open(&path).expect("open fixture hfq"); - (hfq, path) - } - - fn expect_injected(err: hip_bridge::HipError, fail_at: usize) { - let msg = format!("{err:?}"); - assert!( - msg.contains("injected gemma4"), - "fail_at={fail_at}: expected the injected fault, got {msg}" - ); - } - - #[test] - #[ignore = "requires an AMD GPU; proves tiny max_seq is a rollback-safe Err, not a host abort"] - fn scratch_rejects_small_max_seq_before_allocating() { - let Some(mut gpu) = try_gpu() else { return }; - let cfg = tiny_config(vec![LayerType::Sliding], false); - let before = gpu.pool_stats(); - for max_seq in [0usize, 1, 64, 127] { - match Gemma4Scratch::new(&mut gpu, &cfg, max_seq) { - Ok(_) => panic!("max_seq={max_seq} must fail"), - Err(e) => assert!( - format!("{e:?}").contains("too small"), - "unexpected error for max_seq={max_seq}: {e:?}" - ), - } - } - assert_eq!( - gpu.pool_stats(), - before, - "rejected scratch must not allocate" - ); - Gemma4Scratch::new(&mut gpu, &cfg, 128) - .expect("floor value loads") - .free_gpu(&mut gpu); - gpu.drain_pool(); - } - - #[test] - #[ignore = "requires an AMD GPU; exercises real allocation rollback and retry"] - fn scratch_alloc_failure_at_any_point_reclaims_all_and_retry_succeeds() { - let Some(mut gpu) = try_gpu() else { return }; - let cfg = tiny_config(vec![LayerType::Sliding], false); - // Two warm create/free cycles settle the pool: on the first reuse pass - // a larger pooled block can satisfy a smaller same-bucket request - // (LIFO pop), orphaning the smaller block to hip.free plus one fresh - // alloc. After that the size distribution is stable, so the baseline - // below proves rollback reclaims rather than warmup noise. - for _ in 0..2 { - Gemma4Scratch::new(&mut gpu, &cfg, 128) - .expect("warm scratch") - .free_gpu(&mut gpu); - } - let fresh = gpu.pool_stats().0; - // Count staged allocations in one clean pass; the sweep below fails - // each observed allocation in turn, so no count is pinned here. - let total = { - let n = Cell::new(0usize); - let s = Gemma4Scratch::new_with_alloc( - &mut gpu, - &cfg, - 128, - |g, shape, dt| { - n.set(n.get() + 1); - g.alloc_tensor(shape, dt) - }, - |g| { - n.set(n.get() + 1); - g.hip.malloc(4) - }, - ) - .expect("counting pass"); - let total = n.get(); - s.free_gpu(&mut gpu); - total - }; - assert_eq!(gpu.pool_stats().0, fresh); - // Fail every staged allocation in turn — including the pos_buf slot - // and the last pb_bf16 tail. Pool-stat plateau observes pooled - // owners exactly; the 4-byte pos_buf is covered by the explicit - // hip.free arm plus retry success. - for fail_at in 1..=total { - let n = Cell::new(0usize); - let fail = |n: &Cell| { - n.set(n.get() + 1); - n.get() == fail_at - }; - let r = Gemma4Scratch::new_with_alloc( - &mut gpu, - &cfg, - 128, - |g, shape, dt| { - if fail(&n) { - Err(hip_bridge::HipError::new(0, "injected scratch alloc fault")) - } else { - g.alloc_tensor(shape, dt) - } - }, - |g| { - if fail(&n) { - Err(hip_bridge::HipError::new(0, "injected scratch alloc fault")) - } else { - g.hip.malloc(4) - } - }, - ); - match r { - Ok(_) => panic!("fail_at={fail_at} must fail"), - Err(_) => assert_eq!( - gpu.pool_stats().0, - fresh, - "fail_at={fail_at} leaked pooled owners" - ), - } - } - Gemma4Scratch::new(&mut gpu, &cfg, 128) - .expect("immediate retry after allocation failure") - .free_gpu(&mut gpu); - assert_eq!( - gpu.pool_stats().0, - fresh, - "retry leaked instead of reusing the warm pool" - ); - gpu.drain_pool(); - } - - /// Exhaustive fail-point sweep over one fixture: fail every staged step, - /// prove each failure is the injected one and the pool plateaus, then - /// prove the public path retries clean and unloads (sidecars included). - fn sweep_fixture(tag: &str, cfg: &Gemma4Config) { - let Some(mut gpu) = try_gpu() else { return }; - let (mut hfq, path) = open_fixture(fixture_tensors(cfg), tag); - load_weights(&mut hfq, cfg, &mut gpu) - .expect("warm load") - .free_gpu(&mut gpu); - let fresh = gpu.pool_stats().0; - let mut probe = AllocFaults { - calls: 0, - fail_at: None, - }; - match load_weights_impl(&mut hfq, cfg, &mut gpu, Some(&mut probe)) { - Ok(w) => w.free_gpu(&mut gpu), - Err(e) => panic!("counting pass must succeed: {e:?}"), - } - let total = probe.calls; - assert_eq!(gpu.pool_stats().0, fresh); - for fail_at in 1..=total { - let mut f = AllocFaults { - calls: 0, - fail_at: Some(fail_at), - }; - match load_weights_impl(&mut hfq, cfg, &mut gpu, Some(&mut f)) { - Ok(_) => panic!("fail_at={fail_at} must fail"), - Err(e) => { - expect_injected(e, fail_at); - assert_eq!( - gpu.pool_stats().0, - fresh, - "fail_at={fail_at} leaked pooled owners" - ); - } - } - } - // An unfired hook loads clean: success takes every slot exactly once - // (any double-take panics on the take().expect publish). - let mut never = AllocFaults { - calls: 0, - fail_at: Some(total + 100), - }; - match load_weights_impl(&mut hfq, cfg, &mut gpu, Some(&mut never)) { - Ok(w) => { - assert_eq!(w.layers.len(), cfg.n_layers); - w.free_gpu(&mut gpu); - } - Err(e) => panic!("unfired hook must load clean: {e:?}"), - } - // Public-path retry reuses the warm pool, then unloads everything. - load_weights(&mut hfq, cfg, &mut gpu) - .expect("immediate retry after load failure") - .free_gpu(&mut gpu); - assert_eq!( - gpu.pool_stats().0, - fresh, - "retry leaked instead of reusing the warm pool" - ); - gpu.drain_pool(); - let _ = std::fs::remove_file(&path); - } - - #[test] - #[ignore = "requires an AMD GPU; exercises staged-owner rollback across layers and AWQ sidecars"] - fn load_weights_mid_layer_failure_reclaims_completed_layers_and_retry_succeeds() { - let cfg = tiny_config(vec![LayerType::Sliding, LayerType::Full], false); - sweep_fixture("dense", &cfg); - } - - #[test] - #[ignore = "requires an AMD GPU; exercises MoE pool/pointer-table rollback and retry"] - fn load_weights_moe_pool_failure_reclaims_and_retry_succeeds() { - let cfg = tiny_config(vec![LayerType::Sliding], true); - sweep_fixture("moe", &cfg); - } - - /// Bounded sidecar-stage regression: fail after the primary buffer is - /// owned but before a present, valid AWQ sidecar attaches. The outer - /// staged sweep never covers this nested seam (the WeightTensor is never - /// built), so the leaf itself must free the primary. Retry must attach - /// the sidecar and reclaim everything. - #[test] - #[ignore = "requires an AMD GPU; exercises sidecar-stage primary reclaim and retry"] - fn load_weight_sidecar_failure_reclaims_primary_and_retry_succeeds() { - let Some(mut gpu) = try_gpu() else { return }; - let tensors = vec![ - mem("test.weight".to_string(), 13, vec![8, 16], vec![0u8; 256]), - mem( - "test.awq_scale.weight".to_string(), - 1, - vec![16], - vec![0u8; 32], - ), - ]; - let (hfq, path) = open_fixture(tensors, "sidecar"); - // Warm through the production wrapper: sidecar attaches, then unload. - match load_gemma4_weight(&hfq, &mut gpu, "test.weight", 8, 16) { - Ok(w) => { - assert!( - w.awq_scale.is_some(), - "fixture sidecar must attach on the clean path" - ); - w.free_all(&mut gpu); - } - Err(e) => panic!("warm sidecar load must succeed: {e:?}"), - } - let fresh = gpu.pool_stats().0; - // Fail at the sidecar stage: primary owned, sidecar never attached. - let mut f = AllocFaults { - calls: 0, - fail_at: Some(1), - }; - match load_gemma4_weight_impl(&hfq, &mut gpu, "test.weight", 8, 16, Some(&mut f)) { - Ok(_) => panic!("sidecar-stage fault must fail"), - Err(e) => { - expect_injected(e, 1); - assert_eq!( - gpu.pool_stats().0, - fresh, - "sidecar-stage failure leaked the primary buffer" - ); - } - } - // Retry through the production wrapper reuses the warm pool. - match load_gemma4_weight(&hfq, &mut gpu, "test.weight", 8, 16) { - Ok(w) => { - assert!(w.awq_scale.is_some(), "retry must attach the sidecar"); - w.free_all(&mut gpu); - } - Err(e) => panic!("retry after sidecar failure must succeed: {e:?}"), - } - assert_eq!( - gpu.pool_stats().0, - fresh, - "retry leaked instead of reusing the warm pool" - ); - gpu.drain_pool(); - let _ = std::fs::remove_file(&path); - } -} - -// ─── Forward pass ─────────────────────────────────────────────────────── - -/// Fingerprint the allocations and constants baked into the lowered AR graph. -/// Position is deliberately absent: RoPE, KV writes and attention read pos_buf. -fn lowered_graph_binding( - weights: &Gemma4Weights, - config: &Gemma4Config, - kv_sliding: &hipfire_runtime::llama::KvCache, - kv_full: &hipfire_runtime::llama::KvCache, - s: &Gemma4Scratch, -) -> u64 { - let mut h = 0xcbf2_9ce4_8422_2325u64; - let mut mix = |v: u64| h = (h ^ v).wrapping_mul(0x100_0000_01b3); - macro_rules! tensors { - ($($t:expr),* $(,)?) => { - $( - mix($t.buf.as_ptr() as u64); - mix($t.numel() as u64); - mix($t.dtype as u64); - )* - }; - } - macro_rules! projection { - ($w:expr) => {{ - let w = &$w; - tensors!(w.buf); - for v in [w.m, w.k, w.row_stride, w.gpu_dtype as usize] { - mix(v as u64); - } - if let Some(p) = &w.paro { - tensors!(p.pairs, p.theta, p.channel_scales); - mix(p.krot as u64); - mix(p.group_size as u64); - } - if let Some(t) = &w.awq_scale { - tensors!(t); - } - for v in w.lloyd_lut_f16.iter().flatten() { - mix(*v as u64); - } - }}; - } - tensors!(weights.embed_tokens, weights.final_norm); - projection!(weights.lm_head); - mix(weights.embd_format as u64); - for v in [ - config.dim, - config.n_layers, - config.vocab_size, - config.n_heads, - config.sliding_head_dim, - config.sliding_n_kv_heads, - config.sliding_window, - config.full_head_dim, - config.full_n_kv_heads, - config.hidden_dim, - config.enable_moe_block as usize, - config.moe_intermediate_size, - config.num_experts, - config.top_k_experts, - config.attention_k_eq_v as usize, - ] { - mix(v as u64); - } - for v in [ - config.norm_eps, - config.sliding_rope_theta, - config.full_rope_theta, - config.full_partial_rotary_factor, - config.final_logit_softcapping, - ] { - mix(v.to_bits() as u64); - } - for ty in &config.layer_types { - mix(matches!(ty, LayerType::Full) as u64); - } - macro_rules! layer { - ($lw:expr) => {{ - let lw = $lw; - tensors!( - lw.input_layernorm, - lw.post_attention_layernorm, - lw.pre_feedforward_layernorm, - lw.post_feedforward_layernorm, - lw.layer_scalar, - lw.q_norm, - lw.k_norm, - ); - mix(lw.layer_scalar_host.to_bits() as u64); - projection!(lw.q_proj); - projection!(lw.k_proj); - projection!(lw.o_proj); - projection!(lw.gate_proj); - projection!(lw.up_proj); - projection!(lw.down_proj); - mix(lw.moe.is_some() as u64); - if let Some(m) = &lw.moe { - projection!(m.router_proj); - tensors!( - m.router_scale, - m.per_expert_scale, - m.pre_feedforward_layernorm_2, - m.post_feedforward_layernorm_1, - m.post_feedforward_layernorm_2, - m.experts_gate_up_pool, - m.experts_down_pool, - m.experts_gate_up_ptrs, - m.experts_down_ptrs, - ); - if let Some(first) = m.experts.first() { - projection!(first.gate_up_proj); - projection!(first.down_proj); - } - } - }}; - } - for lw in &weights.layers { - match lw { - LayerWeights::Sliding(lw) => { - mix(0); - layer!(lw); - projection!(lw.v_proj); - } - LayerWeights::Full(lw) => { - mix(1); - layer!(lw); - } - } - } - for kv in [kv_sliding, kv_full] { - for v in [ - kv.max_seq, - kv.physical_cap, - kv.compact_offset, - kv.quant_q8 as usize, - kv.quant_asym3 as usize, - kv.quant_fwht as usize, - kv.quant_asym4 as usize, - kv.quant_asym2 as usize, - kv.v_mode_bits() as usize, - ] { - mix(v as u64); - } - for t in kv.k_gpu.iter().chain(&kv.v_gpu) { - tensors!(t); - } - for t in kv.givens_cos.iter().chain(&kv.givens_sin) { - tensors!(t); - } - } - mix(s.pos_buf.as_ptr() as u64); - tensors!( - s.x, - s.residual, - s.tmp, - s.q, - s.k, - s.v, - s.attn_out, - s.gate_ffn, - s.up_ffn, - s.ffn_hidden, - s.ffn_out, - s.logits, - s.flash_partials, - s.v_norm_ones_full, - s.moe_cur_mlp, - s.moe_pre2, - s.moe_router_in, - s.moe_router_logits, - s.moe_topk_indices, - s.moe_topk_weights, - s.moe_cur_moe, - s.moe_expert_gate_up, - s.moe_expert_hidden, - s.moe_expert_out, - s.moe_pre2_rot, - s.moe_expert_gate_batch, - s.moe_expert_up_batch, - s.moe_expert_hidden_batch, - ); - h.max(1) -} - -/// Single-token decode. Phase 3 implementation. -/// -/// Precondition: the loader initializes `scratch.v_norm_ones_full` once before -/// the first forward call. -#[allow(clippy::too_many_arguments)] -pub fn forward_scratch( - gpu: &mut Gpu, - weights: &Gemma4Weights, - config: &Gemma4Config, - token: u32, - pos: usize, - kv_sliding: &mut hipfire_runtime::llama::KvCache, - kv_full: &mut hipfire_runtime::llama::KvCache, - scratch: &Gemma4Scratch, -) -> HipResult<()> { - let dim = config.dim; - - // 1) Embedding lookup + sqrt(dim) scale. - // ALWAYS direct — the captured graph can't bake in `token` (varies per - // call). The embed lookup fills scratch.x; the rest of the forward - // (which is graph-capturable now that the MoE branch has no D2H syncs) - // reads from scratch.x. Same split Qwen35 uses for its captured path. - embed(gpu, weights, &scratch.x, token, dim)?; - gpu.scale_f32(&scratch.x, config.embed_scale)?; - - // Same explicit per-model override as the hand path; HIPFIRE_GRAPH also - // controls this carrier. With both switches unset, capture/replay is the - // measured default on exact gfx1201 only; every other arch stays opt-in. - static GRAPH_ENV: std::sync::LazyLock> = std::sync::LazyLock::new(|| { - let parse = |name| match hipfire_config::developer_var(name).ok().as_deref() { - Some("0") => Some(false), - Some("1") => Some(true), - _ => None, - }; - parse("HIPFIRE_GEMMA4_GRAPH").or_else(|| parse("HIPFIRE_GRAPH")) - }); - let graph_on = (*GRAPH_ENV).unwrap_or(gpu.arch == "gfx1201"); - let eligible = std::mem::replace(&mut gpu.graphs.ar_graph_eligible, true); - let use_graph = graph_on - && eligible - && !gpu.replay.is_recording() - && kv_sliding.compact_offset == 0 - && kv_full.compact_offset == 0 - // The CPU expert fallback downloads top-K inside the body. Q8 expert - // down uses atomicAdd, whose scheduling-dependent sum is not exact. - && weights.layers.iter().all(|layer| { - let moe = match layer { - LayerWeights::Sliding(lw) => &lw.moe, - LayerWeights::Full(lw) => &lw.moe, - }; - moe.as_ref().map_or(true, |m| m.experts.first().is_some_and(|e| { - e.down_proj.gpu_dtype == DType::HFQ4G128 - && hipfire_dispatch::pipeline::sandwich::RoutedExperts::supports( - e.gate_up_proj.gpu_dtype, - e.down_proj.gpu_dtype, - ) - })) - }) - && hipfire_config::developer_var("HIPFIRE_GEMMA4_DUMP").as_deref() != Ok("1"); - let binding = if use_graph { - lowered_graph_binding(weights, config, kv_sliding, kv_full, scratch) - } else { - 0 - }; - if use_graph && gpu.graphs.ar_forward_binding != binding { - gpu.graphs.graph_destroy(&gpu.hip, gpu.device_id); - gpu.graphs.ar_forward_binding = binding; - } - if use_graph && gpu.active_stream.is_none() { - // Embedding above may have run on the default stream. Complete it - // before switching streams for capture/replay. - gpu.hip.device_synchronize()?; - gpu.active_stream = Some(gpu.hip.stream_create()?); - } - if use_graph - && gpu.graphs.ar_forward_replay_enabled - && !gpu.graphs.ar_forward_kernel_dirty - && gpu.graphs.graph_exec.is_some() - { - gpu.hip.stream_write_value32( - gpu.active_stream.as_ref().unwrap(), - &scratch.pos_buf, - pos as u32, - 0, - )?; - if let Err(e) = - gpu.graphs - .graph_launch(&gpu.hip, gpu.device_id, gpu.active_stream.as_ref().unwrap()) - { - gpu.graphs.graph_destroy(&gpu.hip, gpu.device_id); - return Err(e); - } - } else if use_graph && !gpu.graphs.ar_forward_kernel_dirty { - gpu.hip - .memcpy_htod(&scratch.pos_buf, &(pos as i32).to_ne_bytes())?; - gpu.graphs.drop_captured_graph(&gpu.hip, gpu.device_id); - if let Err(e) = gpu.graphs.begin_graph_capture( - &gpu.hip, - gpu.device_id, - gpu.active_stream.as_ref().unwrap(), - ) { - gpu.graphs.capture_mode = false; - gpu.graphs.graph_destroy(&gpu.hip, gpu.device_id); - return Err(e); - } - let body = forward_scratch_inner(gpu, weights, config, pos, kv_sliding, kv_full, scratch); - // Always close capture, including a failing body, before dropping its - // blobs. Captured work is never launched when either phase fails. - let end = gpu.graphs.end_graph_capture( - &gpu.hip, - gpu.device_id, - gpu.active_stream.as_ref().unwrap(), - ); - if let Err(e) = body.and(end) { - gpu.graphs.capture_mode = false; - gpu.graphs.graph_destroy(&gpu.hip, gpu.device_id); - return Err(e); - } - gpu.graphs.ar_forward_binding = binding; - if let Err(e) = - gpu.graphs - .graph_launch(&gpu.hip, gpu.device_id, gpu.active_stream.as_ref().unwrap()) - { - gpu.graphs.graph_destroy(&gpu.hip, gpu.device_id); - return Err(e); - } - gpu.graphs.ar_forward_replay_enabled = true; - eprintln!( - "[gemma4 lowered hipGraph] captured {} blobs, instantiated", - gpu.graphs.ar_forward_blobs.len(), - ); - } else { - if !use_graph && gpu.graphs.graph_exec.is_some() { - gpu.graphs.graph_destroy(&gpu.hip, gpu.device_id); - } - gpu.hip - .memcpy_htod(&scratch.pos_buf, &(pos as i32).to_ne_bytes())?; - forward_scratch_inner(gpu, weights, config, pos, kv_sliding, kv_full, scratch)?; - // A successful direct call JITs the same body before capture. - if use_graph { - gpu.graphs.ar_forward_kernel_dirty = false; - } - } - Ok(()) -} - -/// Graph-capturable body of `forward_scratch`. Caller is responsible for: -/// - embed lookup + embed scale (fills scratch.x — varies per call) -/// - pos_buf update (htod for warmup/capture; stream_write_value32 for replay) -fn forward_scratch_inner( - gpu: &mut Gpu, - weights: &Gemma4Weights, - config: &Gemma4Config, - pos: usize, - kv_sliding: &mut hipfire_runtime::llama::KvCache, - kv_full: &mut hipfire_runtime::llama::KvCache, - scratch: &Gemma4Scratch, -) -> HipResult<()> { - use crate::program::{head, Geometry, LayerScratch, ProgramBinding, Resident}; - let hip = |e: String| hip_bridge::HipError::new(0, &e); - let geo = Geometry::lowered(config); - let binding = ProgramBinding { - geo, - resident: Resident { - kv_sliding, - kv_full, - pos_buf: &scratch.pos_buf, - v_norm_ones: &scratch.v_norm_ones_full, - flash_partials: &scratch.flash_partials, - }, - rows: 1, - position: pos, - positions: None, - scratch: LayerScratch::lowered(scratch, config.enable_moe_block), - }; - let mut steps = Vec::with_capacity(3 * config.n_layers + 3); - layer_steps(&binding, weights, config, &mut steps).map_err(hip)?; - let lm_head = weights.lm_head.dispatch_ref(); - head( - &geo, - &weights.final_norm, - &lm_head, - &scratch.x, - &scratch.tmp, - &scratch.logits, - &mut steps, - ); - let ctx = DispatchCtx::new(gpu); - execute_steps(gpu, &ctx, &steps).map_err(|e| hip(format!("gemma4: {e}"))) -} - -/// `dst = embed(token)`, unscaled. -fn embed( - gpu: &mut Gpu, - weights: &Gemma4Weights, - dst: &GpuTensor, - token: u32, - dim: usize, -) -> HipResult<()> { - let table = &weights.embed_tokens; - match weights.embd_format { - EmbeddingFormat::HFQ4G256 => gpu.embedding_lookup_hfq4g256(table, dst, token, dim), - EmbeddingFormat::HFQ4G128 => gpu.embedding_lookup_hfq4g128(table, dst, token, dim), - EmbeddingFormat::Q8_0 => gpu.embedding_lookup_q8(table, dst, token, dim), - EmbeddingFormat::F32 => gpu.embedding_lookup(table, dst, token, dim), - _ => Err(hip_bridge::HipError::new( - 0, - "unsupported Gemma 4 embed format", - )), - } -} - -/// Every layer's steps, one cache slot per layer of each type in layer order. -fn layer_steps<'a>( - binding: &crate::program::ProgramBinding<'a>, - weights: &'a Gemma4Weights, - config: &Gemma4Config, - steps: &mut Vec>, -) -> Result<(), String> { - use crate::program::{LayerKv, LayerRefs}; - let mut slots = [0usize; 2]; - for (layer_idx, layer) in weights.layers.iter().enumerate() { - let class = matches!(layer, LayerWeights::Full(_)) as usize; - let kv = LayerKv { - slot: slots[class], - writes: true, - }; - slots[class] += 1; - binding.layer( - layer_idx, - LayerRefs::lowered(layer, config.hidden_dim), - kv, - steps, - )?; - } - Ok(()) -} - -/// Batched prefill: runs `tokens` from `start_pos` through the declarative -/// layer program, `rows > 1` at a time, writing both KV caches (which must be -/// Q8). Computes no logits; follow with [`forward_scratch`] on the last token -/// when they are needed. The calibration collector, when armed, sees every -/// projection input. -#[allow(clippy::too_many_arguments)] -pub fn forward_prefill_batch( - gpu: &mut Gpu, - weights: &Gemma4Weights, - config: &Gemma4Config, - tokens: &[u32], - start_pos: usize, - kv_sliding: &mut hipfire_runtime::llama::KvCache, - kv_full: &mut hipfire_runtime::llama::KvCache, - scratch: &Gemma4Scratch, -) -> HipResult<()> { - let hip = |e: String| hip_bridge::HipError::new(0, &e); - let rows_max = crate::gemma4::GEMMA4_FORWARD_BATCH_MAX; - for (i, chunk) in tokens.chunks(rows_max).enumerate() { - let pos = start_pos + i * rows_max; - if let [token] = chunk { - forward_scratch( - gpu, weights, config, *token, pos, kv_sliding, kv_full, scratch, - )?; - } else { - prefill_rows( - gpu, weights, config, chunk, pos, kv_sliding, kv_full, scratch, - ) - .map_err(hip)?; - } - } - Ok(()) -} - -#[allow(clippy::too_many_arguments)] -fn prefill_rows( - gpu: &mut Gpu, - weights: &Gemma4Weights, - config: &Gemma4Config, - tokens: &[u32], - start_pos: usize, - kv_sliding: &hipfire_runtime::llama::KvCache, - kv_full: &hipfire_runtime::llama::KvCache, - scratch: &Gemma4Scratch, -) -> Result<(), String> { - use crate::program::{ - routed_scratch, Geometry, ProgramBinding, Resident, RowActivations, RowBuffers, RowWidths, - }; - let rows = tokens.len(); - let dim = config.dim; - let end = start_pos + rows; - if end > kv_sliding.physical_cap || end > kv_full.physical_cap { - return Err(format!( - "gemma4 prefill: positions [{start_pos}, {end}) exceed the KV capacity" - )); - } - let geo = Geometry::lowered(config); - let head_dim = config.sliding_head_dim.max(config.full_head_dim); - let widths = RowWidths { - dim, - max_q: config.n_heads * head_dim, - max_kv: (config.sliding_n_kv_heads * config.sliding_head_dim) - .max(config.full_n_kv_heads * config.full_head_dim), - ffn: config.hidden_dim, - }; - let mut bufs = RowBuffers::new(gpu); - let acts = RowActivations::new(&mut bufs, gpu, rows, start_pos, widths)?; - let cap = kv_sliding.physical_cap.max(kv_full.physical_cap); - let flash_partials = bufs.alloc( - crate::program::q8_flash_partials_len(gpu, &geo, cap, rows), - "flash_partials", - )?; - for (row, &token) in tokens.iter().enumerate() { - embed(gpu, weights, &acts.x.sub_offset(row * dim, dim), token, dim) - .map_err(|e| format!("gemma4 prefill embed: {e:?}"))?; - } - gpu.scale_f32(&acts.x, config.embed_scale) - .map_err(|e| format!("gemma4 prefill embed scale: {e:?}"))?; - let binding = ProgramBinding { - geo, - resident: Resident { - kv_sliding, - kv_full, - pos_buf: &scratch.pos_buf, - v_norm_ones: &scratch.v_norm_ones_full, - flash_partials: &flash_partials, - }, - rows, - position: start_pos, - positions: Some(&acts.positions), - scratch: acts.scratch( - None, - config.enable_moe_block.then(|| routed_scratch(scratch)), - ), - }; - let mut steps = Vec::with_capacity(3 * config.n_layers); - layer_steps(&binding, weights, config, &mut steps)?; - let ctx = DispatchCtx::new(gpu); - execute_steps(gpu, &ctx, &steps).map_err(|e| format!("gemma4 prefill: {e}")) -} diff --git a/crates/hipfire-arch-gemma4/src/mtp.rs b/crates/hipfire-arch-gemma4/src/mtp.rs new file mode 100644 index 0000000000..ead0ee853a --- /dev/null +++ b/crates/hipfire-arch-gemma4/src/mtp.rs @@ -0,0 +1,385 @@ +// SPDX-License-Identifier: MIT OR Apache-2.0 +// Copyright (c) 2026 Björn Bösel +// hipfire — see LICENSE and NOTICE in the project root. + +//! Gemma 4 speculative decode on the engine seam: the arch-22 EAGLE drafter +//! as an [`MtpDrafter`] (driven by the generic `MtpSpeculator`) and +//! [`SpecTarget`] for [`Gemma4Bundle`]. +//! +//! A window at `position` holds the pending seed (emitted, not yet in the +//! KV) and the target's post-norm hidden of `position - 1`, the hidden that +//! predicted it. The drafter drafts `k` tokens at the constant query +//! position `position`; one batched verify writes `[seed, drafts..]` at +//! `position..` and returns per-row argmax and post-norm hidden; the greedy +//! accept rule commits the accepted drafts plus the bonus. Gemma 4 KV is +//! position-indexed, so rejected rows need no rollback: the next verify +//! overwrites them. Greedy only. + +use crate::bundle::Gemma4Bundle; +use crate::drafter::{ + draft_step, target_embedding_table, Gemma4DrafterConfig, Gemma4DrafterScratch, + Gemma4DrafterWeights, DRAFTER_ARCH_ID, +}; +use crate::forward::{decode_step, forward_batch_spec}; +use hipfire_runtime::hfq::HfqFile; +use hipfire_runtime::llama::argmax; +use hipfire_runtime::session_cache::SessionRoute; +use hipfire_runtime::spec::{ + accept_greedy_prefix, MtpDrafter, MtpWindow, SpecAdvance, SpecGrammar, SpecScratch, SpecTarget, +}; +use rdna_compute::{DType, Gpu, GpuTensor}; + +fn gemma4(target: &mut dyn SpecTarget) -> Result<&mut Gemma4Bundle, String> { + target + .as_any_mut() + .downcast_mut::() + .ok_or_else(|| "gemma4 drafter: target is not a Gemma 4 bundle".to_string()) +} + +/// The arch-22 EAGLE drafter riding a Gemma 4 target. +pub struct Gemma4Drafter { + config: Gemma4DrafterConfig, + weights: Gemma4DrafterWeights, + scratch: Gemma4DrafterScratch, + /// Target post-norm hidden of the position before the pending seed. + pending_hidden: GpuTensor, + /// Per-row post-norm hidden of a verify block, `[(k + 1) · dim]`. + verify_hidden: GpuTensor, + k: usize, + ctx_capacity: usize, +} + +impl Gemma4Drafter { + /// Load the drafter at `path` for `target`, drafting `draft_len` tokens + /// per window (the loader admits 1..=5). + pub fn load( + path: &str, + draft_len: usize, + target: &Gemma4Bundle, + gpu: &mut Gpu, + ) -> Result { + let hfq = HfqFile::open(std::path::Path::new(path)) + .map_err(|e| format!("open gemma4 drafter: {e}"))?; + if hfq.arch_id != DRAFTER_ARCH_ID { + return Err(format!( + "gemma4 EAGLE drafter must be arch_id={DRAFTER_ARCH_ID} (gemma4_unified_assistant); got arch_id={} — a DFlash draft goes in params.draft on qwen3.5 targets, not params.drafter", + hfq.arch_id + )); + } + let config = Gemma4DrafterConfig::from_hfq(&hfq)?; + let dim = target.config.dim; + if config.backbone_hidden != dim { + return Err(format!( + "drafter backbone_hidden ({}) != target hidden ({dim}) — this drafter was trained against a different target width", + config.backbone_hidden + )); + } + if target_embedding_table(&target.weights).is_none() { + return Err(format!( + "gemma4 drafter: target embedding format {:?} has no batched lookup", + target.weights.embd_format + )); + } + let weights = Gemma4DrafterWeights::load(&hfq, &config, gpu)?; + let scratch = match Gemma4DrafterScratch::new(gpu, &config, &weights) { + Ok(scratch) => scratch, + Err(e) => { + weights.free_gpu(gpu); + return Err(e); + } + }; + let hidden = gpu.zeros(&[dim], DType::F32).and_then(|pending| { + match gpu.zeros(&[(draft_len + 1) * dim], DType::F32) { + Ok(verify) => Ok((pending, verify)), + Err(e) => { + let _ = gpu.free_tensor(pending); + Err(e) + } + } + }); + let (pending_hidden, verify_hidden) = match hidden { + Ok(v) => v, + Err(e) => { + scratch.free_gpu(gpu); + weights.free_gpu(gpu); + return Err(format!("gemma4 drafter: hidden buffers: {e:?}")); + } + }; + Ok(Self { + config, + weights, + scratch, + pending_hidden, + verify_hidden, + k: draft_len, + ctx_capacity: target.state.max_seq, + }) + } + + pub fn config(&self) -> &Gemma4DrafterConfig { + &self.config + } + + /// The target's post-norm hidden of its last processed position becomes + /// the pending seed's conditioning hidden. + fn take_target_hidden(&self, gpu: &mut Gpu, bundle: &Gemma4Bundle) -> Result<(), String> { + let bytes = bundle.config.dim * 4; + gpu.hip + .memcpy_dtod_at(&self.pending_hidden.buf, 0, &bundle.state.tmp.buf, 0, bytes) + .map_err(|e| format!("gemma4 drafter: capture hidden: {e:?}")) + } +} + +impl MtpDrafter for Gemma4Drafter { + fn mtp_prefill( + &mut self, + gpu: &mut Gpu, + target: &mut dyn SpecTarget, + prompt_tokens: &[u32], + _fill_tokens: &[u32], + start_pos: usize, + _cache_hit: bool, + _abort: &dyn Fn() -> bool, + ) -> Result { + let bundle = gemma4(target)?; + let logits = bundle.prefill(gpu, prompt_tokens, start_pos, SessionRoute::Mtp)?; + self.take_target_hidden(gpu, bundle)?; + Ok(argmax(&logits)) + } + + fn mtp_step( + &mut self, + gpu: &mut Gpu, + target: &mut dyn SpecTarget, + position: usize, + seed: u32, + _emitted: &[u32], + k: usize, + eos: u32, + _grammar: Option<&mut dyn SpecGrammar>, + ) -> Result { + let bundle = gemma4(target)?; + if bundle.state.n_tokens != position { + return Err(format!( + "gemma4 drafter: window at {position} but the target is at {}", + bundle.state.n_tokens + )); + } + let dim = bundle.config.dim; + if k == 0 { + // The seed alone: one ordinary decode step. + let logits = decode_step( + &bundle.config, + &bundle.weights, + &mut bundle.state, + gpu, + seed, + position as u32, + )?; + self.take_target_hidden(gpu, bundle)?; + return Ok(MtpWindow { + committed: vec![argmax(&logits)], + accepted: 0, + drafts_generated: 0, + }); + } + let k = k.min(self.k); + self.scratch + .begin_round(gpu, &self.pending_hidden, position)?; + let mut block = Vec::with_capacity(k + 1); + block.push(seed); + for _ in 0..k { + let prev = *block.last().expect("seeded block"); + block.push(draft_step( + gpu, + &self.weights, + &self.config, + &mut self.scratch, + &bundle.weights, + &bundle.state, + &bundle.config, + prev, + position, + )?); + } + let hidden = self.verify_hidden.sub_offset(0, block.len() * dim); + let mut picks = Vec::with_capacity(block.len()); + forward_batch_spec( + &bundle.config, + &bundle.weights, + &mut bundle.state, + gpu, + &block, + position, + Some(&hidden), + Some(&mut picks), + )?; + let accept = accept_greedy_prefix(&block[1..], &picks, Some(eos)); + let committed = accept.committed.len(); + // The seed and every committed token but the last are in the KV; the + // last stays pending, conditioned on the hidden of the row before it. + bundle.state.n_tokens = position + committed; + let row = hidden.sub_offset((committed - 1) * dim, dim); + gpu.hip + .memcpy_dtod_at(&self.pending_hidden.buf, 0, &row.buf, 0, dim * 4) + .map_err(|e| format!("gemma4 drafter: capture verify hidden: {e:?}"))?; + Ok(MtpWindow { + committed: accept.committed, + accepted: accept.accepted, + drafts_generated: k, + }) + } + + fn mtp_forced_advance( + &mut self, + gpu: &mut Gpu, + target: &mut dyn SpecTarget, + tokens: &[u32], + start_pos: usize, + _abort: &dyn Fn() -> bool, + ) -> Result { + let bundle = gemma4(target)?; + if !tokens.is_empty() { + bundle.forward_rows(gpu, tokens, start_pos)?; + self.take_target_hidden(gpu, bundle)?; + } + Ok(true) + } + + fn mtp_reset(&mut self, _gpu: &mut Gpu) -> Result<(), String> { + // Stateless between windows: prefill restages the conditioning hidden. + Ok(()) + } + + fn mtp_free(self: Box, gpu: &mut Gpu) { + let d = *self; + let _ = gpu.free_tensor(d.pending_hidden); + let _ = gpu.free_tensor(d.verify_hidden); + d.scratch.free_gpu(gpu); + d.weights.free_gpu(gpu); + } + + fn k(&self) -> usize { + self.k + } + + fn ctx_capacity(&self) -> usize { + self.ctx_capacity + } + + fn requires_greedy(&self) -> bool { + true + } + + fn name(&self) -> &'static str { + "eagle" + } +} + +/// Gemma 4 verify needs no scratch of its own: KV is position-indexed and +/// nothing is snapshotted before a verify. +struct NoScratch; + +impl SpecScratch for NoScratch { + fn as_any_mut(&mut self) -> &mut dyn std::any::Any { + self + } + fn free(self: Box, _gpu: &mut Gpu) {} +} + +impl SpecTarget for Gemma4Bundle { + fn as_any_mut(&mut self) -> &mut dyn std::any::Any { + self + } + + fn reset_recurrent(&mut self, gpu: &mut Gpu) -> Result<(), String> { + hipfire_runtime::arch_model::ArchModel::reset_session_state(self, gpu) + } + + fn new_spec_scratch( + &mut self, + _gpu: &mut Gpu, + _block_size: usize, + ) -> Result, String> { + Ok(Box::new(NoScratch)) + } + + fn spec_advance( + &mut self, + gpu: &mut Gpu, + tokens: &[u32], + start_pos: usize, + reset: bool, + _abort: &dyn Fn() -> bool, + _hidden_out: Option<&mut Vec>, + ) -> Result { + if reset { + self.state.reset(); + } + if self.state.n_tokens != start_pos { + return Err(format!( + "gemma4 spec_advance from {start_pos} but the state is at {}", + self.state.n_tokens + )); + } + let logits = self.forward_rows(gpu, tokens, start_pos)?; + Ok(SpecAdvance::Ready { + last_argmax: argmax(&logits), + last_logits: Some(logits), + }) + } + + fn verify_block( + &mut self, + gpu: &mut Gpu, + block: &[u32], + position: usize, + _scratch: &mut dyn SpecScratch, + _hidden_out: Option<&mut Vec>, + ) -> Result, String> { + if let [token] = block { + let logits = decode_step( + &self.config, + &self.weights, + &mut self.state, + gpu, + *token, + position as u32, + )?; + return Ok(vec![argmax(&logits)]); + } + let mut picks = Vec::with_capacity(block.len()); + forward_batch_spec( + &self.config, + &self.weights, + &mut self.state, + gpu, + block, + position, + None, + Some(&mut picks), + )?; + Ok(picks) + } + + fn commit_prefix( + &mut self, + _gpu: &mut Gpu, + _block: &[u32], + accept_len: usize, + position: usize, + _scratch: &mut dyn SpecScratch, + ) -> Result<(), String> { + // Rows past the committed prefix are overwritten by the next verify. + self.state.n_tokens = position + accept_len + 1; + Ok(()) + } + + fn eos_token(&self) -> u32 { + self.eos_tok + } + + fn ctx_capacity(&self) -> usize { + self.state.max_seq + } +} diff --git a/crates/hipfire-arch-gemma4/src/program.rs b/crates/hipfire-arch-gemma4/src/program.rs index f430aba391..cebe0c55b6 100644 --- a/crates/hipfire-arch-gemma4/src/program.rs +++ b/crates/hipfire-arch-gemma4/src/program.rs @@ -12,11 +12,10 @@ //! and scratch to the shared `hipfire_dispatch::pipeline::sandwich` //! operations; route selection and executor bodies live in dispatch. The same //! program serves single-token decode (`rows == 1`) and batched prefill / -//! verify (`rows > 1`), for both weight stacks (`gemma4` and `lowered`). +//! verify (`rows > 1`) of every Gemma 4 variant. use crate::config::{Gemma4Config, LayerType, RopeType}; -use crate::gemma4::{self, PerLayerBranchWeights}; -use crate::lowered; +use crate::gemma4::{self, MoeLayerExtras, MoeScratch, PerLayerBranchWeights}; use hip_bridge::DeviceBuffer; use hipfire_dispatch::families::gemv::WeightRef; use hipfire_dispatch::pipeline::sandwich::{ @@ -39,7 +38,7 @@ pub(crate) struct AttnGeometry { pub rope: Rope, } -/// Model geometry the program reads; built from either config type. +/// Model geometry the program reads. #[derive(Clone, Copy)] pub(crate) struct Geometry { pub dim: usize, @@ -78,8 +77,16 @@ impl Geometry { vocab: cfg.vocab_size, softcap: cfg.final_logit_softcapping, ple_width: cfg.hidden_size_per_layer_input, - moe_top_k: 0, - moe_hidden: 0, + moe_top_k: if cfg.enable_moe_block { + cfg.top_k_experts + } else { + 0 + }, + moe_hidden: if cfg.enable_moe_block { + cfg.moe_intermediate_size + } else { + 0 + }, sliding: AttnGeometry { head_dim: cfg.sliding_head_dim, n_kv_heads: cfg.sliding_n_kv_heads, @@ -138,44 +145,6 @@ impl Geometry { }, } } - - pub fn lowered(cfg: &lowered::Gemma4Config) -> Self { - let rope_type = match cfg.full_rope_type { - lowered::RopeType::Default => RopeType::Default, - lowered::RopeType::Proportional => RopeType::Proportional, - }; - Self { - dim: cfg.dim, - eps: cfg.norm_eps, - n_heads: cfg.n_heads, - n_layers: cfg.n_layers, - vocab: cfg.vocab_size, - softcap: cfg.final_logit_softcapping, - ple_width: 0, - moe_top_k: cfg.top_k_experts, - moe_hidden: cfg.moe_intermediate_size, - sliding: AttnGeometry { - head_dim: cfg.sliding_head_dim, - n_kv_heads: cfg.sliding_n_kv_heads, - window: cfg.sliding_window, - rope: Rope { - kind: RopeKind::RotateHalf, - theta: cfg.sliding_rope_theta, - }, - }, - full: AttnGeometry { - head_dim: cfg.full_head_dim, - n_kv_heads: cfg.full_n_kv_heads, - window: 0, - rope: full_rope( - cfg.full_head_dim, - rope_type, - cfg.full_partial_rotary_factor, - cfg.full_rope_theta, - ), - }, - } - } } #[derive(Clone, Copy)] @@ -204,7 +173,7 @@ pub(crate) struct LayerRefs<'a> { down: &'a WeightTensor, ffn_hidden_dim: usize, per_layer: Option<&'a PerLayerBranchWeights>, - moe: Option<&'a lowered::MoeLayerExtras>, + moe: Option<&'a MoeLayerExtras>, } impl<'a> LayerRefs<'a> { @@ -230,7 +199,7 @@ impl<'a> LayerRefs<'a> { down: &l.down_proj, ffn_hidden_dim: l.ffn_hidden_dim, per_layer: l.per_layer.as_ref(), - moe: None, + moe: l.moe.as_ref(), }, gemma4::LayerWeights::Full(l) => Self { layer_type: LayerType::Full, @@ -252,56 +221,6 @@ impl<'a> LayerRefs<'a> { down: &l.down_proj, ffn_hidden_dim: l.ffn_hidden_dim, per_layer: l.per_layer.as_ref(), - moe: None, - }, - } - } - - /// `hidden_dim` is the lowered config's dense FFN width. - pub fn lowered(layer: &'a lowered::LayerWeights, hidden_dim: usize) -> Self { - match layer { - lowered::LayerWeights::Sliding(l) => Self { - layer_type: LayerType::Sliding, - input_norm: &l.input_layernorm, - post_attn_norm: &l.post_attention_layernorm, - pre_ffn_norm: &l.pre_feedforward_layernorm, - post_ffn_norm: &l.post_feedforward_layernorm, - layer_scalar: l.layer_scalar_host, - q: &l.q_proj, - q_norm: &l.q_norm, - kv: Some(LayerKvWeights { - k: &l.k_proj, - v: Some(&l.v_proj), - k_norm: &l.k_norm, - }), - o: &l.o_proj, - gate: &l.gate_proj, - up: &l.up_proj, - down: &l.down_proj, - ffn_hidden_dim: hidden_dim, - per_layer: None, - moe: l.moe.as_ref(), - }, - lowered::LayerWeights::Full(l) => Self { - layer_type: LayerType::Full, - input_norm: &l.input_layernorm, - post_attn_norm: &l.post_attention_layernorm, - pre_ffn_norm: &l.pre_feedforward_layernorm, - post_ffn_norm: &l.post_feedforward_layernorm, - layer_scalar: l.layer_scalar_host, - q: &l.q_proj, - q_norm: &l.q_norm, - kv: Some(LayerKvWeights { - k: &l.k_proj, - v: None, - k_norm: &l.k_norm, - }), - o: &l.o_proj, - gate: &l.gate_proj, - up: &l.up_proj, - down: &l.down_proj, - ffn_hidden_dim: hidden_dim, - per_layer: None, moe: l.moe.as_ref(), }, } @@ -414,28 +333,7 @@ impl<'a> LayerScratch<'a> { act: &state.ffn_hidden, mlp_out: &state.ffn_out, ple, - moe: None, - } - } - - /// The resident single-row scratch of the lowered state. - pub fn lowered(s: &'a lowered::Gemma4Scratch, moe: bool) -> Self { - Self { - x: &s.x, - residual: &s.residual, - normed: &s.tmp, - attn_rot: &s.x_rot, - mlp_rot: &s.x_rot, - q: &s.q, - k: &s.k, - v: &s.v, - attn_out: &s.attn_out, - gate: &s.gate_ffn, - up: &s.up_ffn, - act: &s.ffn_hidden, - mlp_out: &s.ffn_out, - ple: None, - moe: moe.then(|| routed_scratch(s)), + moe: state.moe.as_ref().map(routed_scratch), } } } @@ -464,21 +362,21 @@ pub(crate) fn q8_flash_partials_len( .unwrap_or(0) } -/// The lowered state's single-row routed-expert scratch; batched forwards -/// route row by row through it. -pub(crate) fn routed_scratch(s: &lowered::Gemma4Scratch) -> RoutedScratch<'_> { +/// The state's single-row routed-expert scratch; batched forwards route row +/// by row through it. +pub(crate) fn routed_scratch(s: &MoeScratch) -> RoutedScratch<'_> { RoutedScratch { - input: &s.moe_pre2, - input_rot: &s.moe_pre2_rot, - router_in: &s.moe_router_in, - router_logits: &s.moe_router_logits, - topk_indices: &s.moe_topk_indices, - topk_weights: &s.moe_topk_weights, - gate: &s.moe_expert_gate_batch, - up: &s.moe_expert_up_batch, - act: &s.moe_expert_hidden_batch, - out: &s.moe_cur_moe, - dense_normed: &s.moe_cur_mlp, + input: &s.pre2, + input_rot: &s.pre2_rot, + router_in: &s.router_in, + router_logits: &s.router_logits, + topk_indices: &s.topk_indices, + topk_weights: &s.topk_weights, + gate: &s.gate_batch, + up: &s.up_batch, + act: &s.hidden_batch, + out: &s.cur_moe, + dense_normed: &s.cur_mlp, } } @@ -714,10 +612,6 @@ impl<'a> ProgramBinding<'a> { .moe .as_ref() .ok_or_else(|| format!("gemma4 layer {layer_idx}: missing MoE scratch"))?; - let expert = moe - .experts - .first() - .ok_or_else(|| format!("gemma4 layer {layer_idx}: MoE layer has no experts"))?; steps.push(Step::ParallelMoeMlp(ParallelMoeMlpOp { mlp, dense_post_norm: &moe.post_feedforward_layernorm_1, @@ -726,11 +620,11 @@ impl<'a> ProgramBinding<'a> { router_norm: &moe.router_scale, router_input_scale: 1.0 / (geo.dim as f32).sqrt(), router: moe.router_proj.dispatch_ref(), - n_experts: moe.experts.len(), + n_experts: moe.n_experts, top_k: geo.moe_top_k, hidden_dim: geo.moe_hidden, - gate_up_dtype: expert.gate_up_proj.gpu_dtype, - down_dtype: expert.down_proj.gpu_dtype, + gate_up_dtype: moe.gate_up_dtype, + down_dtype: moe.down_dtype, gate_up_ptrs: &moe.experts_gate_up_ptrs, down_ptrs: &moe.experts_down_ptrs, per_expert_scale: &moe.per_expert_scale, diff --git a/crates/hipfire-arch-gemma4/src/speculative.rs b/crates/hipfire-arch-gemma4/src/speculative.rs deleted file mode 100644 index 009c07cbbb..0000000000 --- a/crates/hipfire-arch-gemma4/src/speculative.rs +++ /dev/null @@ -1,252 +0,0 @@ -// SPDX-License-Identifier: Apache-2.0 -// Copyright (c) 2026 Kaden Schutt -// hipfire — see LICENSE and NOTICE in the project root. - -//! Gemma 4 EAGLE speculative-decode step (`spec_step_gemma4_eagle`). -//! -//! Ties the validated pieces together into one draft→verify→accept→commit -//! cycle: -//! - the EAGLE drafter (`crate::drafter::drafter_step`, per-step cosine -//! 0.994–0.9996 vs HF) drafts N tokens at a CONSTANT query position; and -//! - the batched verify forward (`crate::forward::forward_batch_spec`, -//! argmax-identical to B sequential `decode_step`) verifies the block in one -//! pass and exposes per-position argmax + per-position post-`model.norm` -//! hidden. -//! -//! Structure mirrors qwen35's `spec_step_dflash` (draft loop, greedy accept -//! indexing, bonus token, next-seed-hidden capture from the verify row) but -//! gemma4 is a PURE-ATTENTION model — there is NO DeltaNet / recurrent state, so -//! all of dflash's `DeltaNetSnapshot` / `GdnTape` / replay machinery is dropped. -//! The only spec-decode state that matters is the KV cache, and the two gemma4 -//! KV caches are absolute-position (each write is keyed by `pos_buf` into a -//! `physical_cap`-sized slot), so rejected drafts' KV slots are overwritten -//! losslessly when the next round's verify re-occupies those positions from the -//! new committed length. No explicit truncation is needed. -//! -//! ## The accept invariant (greedy, temp 0) -//! -//! With greedy verify the committed sequence is provably the target's greedy -//! AR sequence: `argmax_per_pos[i]` is the target's argmax AFTER block position -//! `i`, so accepting the longest prefix where `argmax_per_pos[i] == block[i+1]` -//! and emitting `bonus = argmax_per_pos[accept_len]` yields exactly what eager -//! `decode_step` would have produced. Any divergence from eager is a bug. - -use crate::config::Gemma4Config; -use crate::drafter::{drafter_step, Gemma4DrafterConfig, Gemma4DrafterScratch, Gemma4DrafterWeights}; -use crate::forward::forward_batch_spec; -use crate::gemma4::{Gemma4State, Gemma4Weights}; -use rdna_compute::{DType, Gpu, GpuTensor}; - -/// Result of one EAGLE spec step. -pub struct SpecStepOut { - /// Tokens committed this round: accepted drafts ++ [bonus]. Length ≥ 1. - pub committed: Vec, - /// Number of drafts accepted (0..=N). `committed.len() == accept_len + 1`. - pub accept_len: usize, - /// The bonus token (= next round's seed token). Last element of `committed`. - pub next_seed_token: u32, - /// Drafts produced this round (for diagnostics). Length N. - pub drafts: Vec, -} - -/// Reusable GPU scratch for the EAGLE spec loop. Holds the seed-hidden buffer -/// (post-`model.norm`, [dim]) that seeds each round, a buffer for the drafter's -/// per-step post-projection feedback ([dim]), and the [N+1, dim] verify-hidden -/// buffer that `forward_batch_spec` fills (so the next seed hidden is captured -/// from a verify row with no extra forward). -pub struct Gemma4SpecScratch { - /// Last committed position's post-`model.norm` hidden — seeds the round. [dim] - pub seed_hidden: GpuTensor, - /// Drafter per-step hidden feedback (post_proj of the previous draft). [dim] - pub draft_hidden: GpuTensor, - /// Verify per-position post-`model.norm` hidden ([N_max+1, dim]); the next - /// round's seed hidden is copied out of the accepted-bonus row. - pub verify_hidden: GpuTensor, - dim: usize, -} - -impl Gemma4SpecScratch { - /// `max_draft` = the largest `draft_len` that will be passed to - /// `spec_step_gemma4_eagle` (the verify block is `max_draft + 1` rows). - pub fn new(gpu: &mut Gpu, cfg: &Gemma4Config, max_draft: usize) -> Result { - let dim = cfg.dim; - let alloc = |g: &mut Gpu, n: usize, label: &str| -> Result { - g.zeros(&[n], DType::F32) - .map_err(|e| format!("gemma4-spec: alloc {label}: {e:?}")) - }; - Ok(Gemma4SpecScratch { - seed_hidden: alloc(gpu, dim, "seed_hidden")?, - draft_hidden: alloc(gpu, dim, "draft_hidden")?, - verify_hidden: alloc(gpu, (max_draft + 1) * dim, "verify_hidden")?, - dim, - }) - } - - /// Seed the cycle from the prefill's last post-`model.norm` hidden (the - /// lm_head input of the last prefilled token). For `forward_batch` prefill, - /// pass the per-token-hidden row of the last position; for eager prefill, - /// pass `state.tmp` after the final `decode_step`. - pub fn set_seed_hidden_from(&self, gpu: &Gpu, src: &GpuTensor) -> Result<(), String> { - gpu.hip - .memcpy_dtod_at(&self.seed_hidden.buf, 0, &src.buf, 0, self.dim * 4) - .map_err(|e| format!("gemma4-spec: seed_hidden copy: {e:?}")) - } - - pub fn free(self, gpu: &mut Gpu) { - let _ = gpu.free_tensor(self.seed_hidden); - let _ = gpu.free_tensor(self.draft_hidden); - let _ = gpu.free_tensor(self.verify_hidden); - } -} - -/// One EAGLE spec-decode step. -/// -/// * `gpu` — device handle. -/// * `target_w` / `target_cfg` / `target_state` — the arch-13 target (read + -/// KV-written by verify). `target_state.n_tokens` is updated to the new -/// committed length on success. -/// * `dw` / `dcfg` / `ds` — the arch-22 drafter (weights / config / scratch), -/// read-only against the target. -/// * `sp` — reusable spec scratch. `sp.seed_hidden` MUST hold the last committed -/// position's post-`model.norm` hidden on entry (set by the previous round's -/// capture, or by `set_seed_hidden_from` after prefill). On exit it holds the -/// NEXT round's seed hidden (the accepted-bonus row's hidden). -/// * `seed_token` — the last committed token (whose KV sits at position L-1). -/// * `committed_len` — L, the number of committed tokens (= the next position). -/// * `draft_len` — N drafts to speculate (verify block is N+1 tokens). Must be -/// ≤ the `max_draft` the scratch was sized for and keep `L-1+N+1 ≤ max_seq`. -/// * `temp` — 0.0 = greedy (the only validated mode here). temp > 0 is rejected. -/// -/// Returns the committed tokens (accepted drafts ++ bonus) and the next seed -/// token. Side effects: verify writes both KV caches for [L-1, L+N) and the fn -/// sets `target_state.n_tokens = L + accept_len + 1`. -#[allow(clippy::too_many_arguments)] -pub fn spec_step_gemma4_eagle( - gpu: &mut Gpu, - target_w: &Gemma4Weights, - target_cfg: &Gemma4Config, - target_state: &mut Gemma4State, - dw: &Gemma4DrafterWeights, - dcfg: &Gemma4DrafterConfig, - ds: &mut Gemma4DrafterScratch, - sp: &mut Gemma4SpecScratch, - seed_token: u32, - committed_len: usize, - draft_len: usize, - temp: f32, -) -> Result { - if temp != 0.0 { - return Err("gemma4-spec: only greedy (temp=0) is supported".to_string()); - } - if draft_len == 0 { - return Err("gemma4-spec: draft_len must be ≥ 1".to_string()); - } - let l = committed_len; - if l == 0 { - return Err("gemma4-spec: committed_len must be ≥ 1 (need a seed at L-1)".to_string()); - } - let n = draft_len; - let dim = target_cfg.dim; - // The drafter queries the target's shared KV at the CONSTANT position L-1 - // across the whole round (per the candidate-generator semantics). - let query_pos = l - 1; - - // ── 1. Draft N tokens. step0: prev = seed_token, h = sp.seed_hidden. ── - let mut drafts: Vec = Vec::with_capacity(n); - let mut prev_token = seed_token; - // The hidden half fed to the drafter: round-0 = the target's last-committed - // normed hidden (sp.seed_hidden), subsequent steps = the previous step's - // post_proj. We keep the live hidden in `sp.draft_hidden` and point at it. - gpu.hip - .memcpy_dtod_at(&sp.draft_hidden.buf, 0, &sp.seed_hidden.buf, 0, dim * 4) - .map_err(|e| format!("gemma4-spec: seed→draft_hidden copy: {e:?}"))?; - for _j in 0..n { - let out = drafter_step( - gpu, - dw, - dcfg, - ds, - target_w, - target_state, - target_cfg, - prev_token, - &sp.draft_hidden, - query_pos, - )?; - drafts.push(out.argmax); - prev_token = out.argmax; - // Feed this step's post_proj as the next step's hidden half. - gpu.hip - .memcpy_htod(&sp.draft_hidden.buf, unsafe { - std::slice::from_raw_parts( - out.post_proj_hidden.as_ptr() as *const u8, - out.post_proj_hidden.len() * 4, - ) - }) - .map_err(|e| format!("gemma4-spec: draft post_proj htod: {e:?}"))?; - } - - // ── 2. Verify. block = [seed_token, d_0..d_{N-1}] at start_pos = L-1. ── - // forward_batch_spec writes the seed's KV at L-1 (re-write, identical) and - // each draft d_j at L+j, and returns per-position argmax + per-position - // post-norm hidden. Both KV caches are written for [L-1, L+N). - let mut block: Vec = Vec::with_capacity(n + 1); - block.push(seed_token); - block.extend_from_slice(&drafts); - let block_len = block.len(); // N+1 - - let hidden_out = sp.verify_hidden.sub_offset(0, block_len * dim); - let mut argmax_per_pos: Vec = Vec::with_capacity(block_len); - let _last_logits = forward_batch_spec( - target_cfg, - target_w, - target_state, - gpu, - &block, - query_pos, // start_pos = L-1 - Some(&hidden_out), - Some(&mut argmax_per_pos), - )?; - debug_assert_eq!(argmax_per_pos.len(), block_len); - - // ── 3. Greedy accept. ── - // argmax_per_pos[i] = target's prediction AFTER block position i. - // accept d_i (= block[i+1]) iff argmax_per_pos[i] == block[i+1]. - // bonus = argmax_per_pos[accept_len]. - let mut accept_len = 0usize; - for i in 0..n { - if argmax_per_pos[i] == block[i + 1] { - accept_len += 1; - } else { - break; - } - } - let bonus = argmax_per_pos[accept_len]; - - // committed = block[1..1+accept_len] (accepted drafts) ++ [bonus]. - let mut committed: Vec = Vec::with_capacity(accept_len + 1); - committed.extend_from_slice(&block[1..1 + accept_len]); - committed.push(bonus); - - // ── 4. Commit. ── - // Accepted drafts occupy positions L..L+accept_len-1 (KV already written by - // verify). The bonus is the NEW token at L+accept_len — its KV is NOT yet - // written (it wasn't in the block); it becomes the next round's seed and is - // (re)written at the top of the next verify. Rejected drafts at - // L+accept_len..L+N-1 keep stale KV, overwritten next round. - target_state.n_tokens = l + accept_len + 1; - - // Next seed hidden = the verify row whose argmax produced the bonus = row - // `accept_len` of the per-position post-norm hidden buffer. - let bonus_row = sp.verify_hidden.sub_offset(accept_len * dim, dim); - gpu.hip - .memcpy_dtod_at(&sp.seed_hidden.buf, 0, &bonus_row.buf, 0, dim * 4) - .map_err(|e| format!("gemma4-spec: next seed_hidden capture: {e:?}"))?; - - Ok(SpecStepOut { - committed, - accept_len, - next_seed_token: bonus, - drafts, - }) -} diff --git a/crates/hipfire-arch-qwen4/src/bundle.rs b/crates/hipfire-arch-qwen4/src/bundle.rs index 87a484cbca..c61bc142a5 100644 --- a/crates/hipfire-arch-qwen4/src/bundle.rs +++ b/crates/hipfire-arch-qwen4/src/bundle.rs @@ -31,7 +31,8 @@ use hipfire_runtime::sampler::{ apply_logit_policy_candidates_cpu, apply_logit_policy_cpu, PenaltyTable, SamplerConfig, }; use hipfire_runtime::session_cache::{ - SessionCache, SessionRoute, SessionState, SnapshotParts, StateLayout, + stride_boundaries, SessionCache, SessionDriver, SessionRoute, SessionState, SnapshotParts, + StateLayout, }; use hipfire_runtime::spec_sampling::{SampleSpec, SparseDist}; use hipfire_runtime::weight_manifest::{WeightEntry, WeightResidency}; @@ -85,113 +86,25 @@ impl AttachedWeightStore { } } -/// Host record of the committed live state: the consumed history it holds and -/// the route that produced it. A live continuation is *session-exact* (its -/// state descends from decode, as in a continuous conversation), not -/// cold-exact, so it is only resumed in place and never captured into the -/// cross-session snapshot cache, except by message-end snapshots -/// (`HIPFIRE_QWEN4_TURN_SNAPSHOTS`), which live in their own `/turns` scope. -#[derive(Clone, Debug, PartialEq, Eq)] -struct Qwen4LiveRecord { - tokens: Vec, - route: SessionRoute, -} - -/// `Some(L)` when the live state recorded in `record` can serve `prompt` on -/// `route`: it holds exactly `L = record.tokens.len()` tokens (`0 < L < -/// prompt.len()`, so at least the last prompt token is computed), those tokens -/// are `prompt[..L]`, no verify row capture is armed, the target is exactly at -/// `L` and, for native MTP, so is the head (a retired head lags and cannot -/// continue). Plain data, so the decision runs without a GPU. -fn live_start( - record: Option<&Qwen4LiveRecord>, - prompt: &[u32], +/// Where every owner of `route` sits when the live state can continue in +/// place: the target position, for native MTP only when the head sits there +/// too (a retired head lags, a missing one cannot continue), and never while +/// a verify row capture is armed. Plain data, so it runs without a GPU. +fn live_position( route: SessionRoute, target_position: usize, head_position: Option, row_capture_armed: bool, ) -> Option { - let record = record?; - let end = record.tokens.len(); - (record.route == route - && end > 0 - && end < prompt.len() - && prompt[..end] == record.tokens[..] - && !row_capture_armed - && target_position == end - && (route == SessionRoute::Ar || head_position == Some(end))) - .then_some(end) -} - -/// Whether the state left by the turn that began with `turn_prompt` on `route` -/// is the live state after `consumed`, the full host history the client -/// committed: `consumed` extends the turn's prompt, the target is exactly at -/// its end (and for native MTP so is the head) and no row capture is armed. -/// An empty `consumed` (unrepaired terminal) is never live. -fn live_after_commit( - turn_prompt: &[u32], - route: SessionRoute, - consumed: &[u32], - target_position: usize, - head_position: Option, - row_capture_armed: bool, -) -> bool { - let prompt_len = turn_prompt.len(); - prompt_len > 0 - && consumed.len() >= prompt_len - && consumed[..prompt_len] == *turn_prompt - && !row_capture_armed - && target_position == consumed.len() - && (route == SessionRoute::Ar || head_position == Some(consumed.len())) -} - -/// Prompt tokens the next prefill skips: the live state when it reaches at -/// least as far as the longest cached snapshot (a tie prefers live: no copy), -/// else that snapshot. -fn plan_start(snapshot: usize, live: Option) -> usize { - match live { - Some(live) if live >= snapshot => live, - _ => snapshot, - } + (!row_capture_armed + && (route == SessionRoute::Ar || head_position == Some(target_position))) + .then_some(target_position) } /// Environment flag that adds message-end snapshot boundaries (P1b), read once /// when the bundle is assembled. Off by default. const TURN_SNAPSHOTS_ENV: &str = "HIPFIRE_QWEN4_TURN_SNAPSHOTS"; -/// Positions just after every `im_end` token of `prompt` (a message end: -/// user, tool, system or assistant) in `(after, up_to]`, ascending. -fn message_end_boundaries(prompt: &[u32], im_end: u32, after: usize, up_to: usize) -> Vec { - let end = up_to.min(prompt.len()); - if after >= end { - return Vec::new(); - } - (after..end) - .filter(|&i| prompt[i] == im_end) - .map(|i| i + 1) - .collect() -} - -/// Ascending, deduplicated snapshot boundaries in `(after, up_to]`: multiples -/// of `chunk`, plus message ends when `im_end` is `Some`. -fn session_boundaries( - chunk: usize, - prompt: &[u32], - im_end: Option, - after: usize, - up_to: usize, -) -> Vec { - let mut boundaries: Vec = (after / chunk + 1..=up_to / chunk) - .map(|k| k * chunk) - .collect(); - if let Some(im_end) = im_end { - boundaries.extend(message_end_boundaries(prompt, im_end, after, up_to)); - boundaries.sort_unstable(); - boundaries.dedup(); - } - boundaries -} - /// Published Qwen4 architecture owner. pub struct Qwen4Bundle { pub config: Qwen4Config, @@ -221,9 +134,11 @@ pub struct Qwen4Bundle { /// Staging of the GPU penalty prepass (`apply_penalty_table`), allocated /// by the first penalized row and kept until unload. penalty_stage: Option, - /// Prefill snapshot cache shared by every session (`attach_session_cache`); - /// `None` = off. - session: Option, + /// Prefill snapshot cache shared by every session plus this + /// conversation's committed live state (`attach_session_cache`); `None` = + /// off. Message-end snapshots of a live continuation live in their own + /// `/turns` scope (`HIPFIRE_QWEN4_TURN_SNAPSHOTS`). + session: Option, /// PLE hot-row store and its worker thread (`attach_ple_hot_store`); /// `None` = off. Dropping it stops and joins the worker; `free_gpu` does /// so before the row store unloads. @@ -235,13 +150,6 @@ pub struct Qwen4Bundle { /// Where the PLE rows come from: the model file itself or the sidecar /// whose descriptors the row store was built from. ple_source: Qwen4PleRows, - /// Committed live state kept for a continuation of the same conversation - /// (`session_commit_live`); `None` = none. Cleared by `reset`, - /// `session_clear`, `session_commit` and `session_begin` (it is consumed - /// while a request is in flight). - live: Option, - /// Prompt and route of the request in flight (set by `session_begin`). - turn: Option<(Vec, SessionRoute)>, /// `HIPFIRE_QWEN4_TURN_SNAPSHOTS` (read once at assembly): also snapshot /// at message ends. Active only with `turn_end_token`. turn_snapshots: bool, @@ -354,7 +262,7 @@ impl SessionState for Qwen4Bundle { let Some(chunk) = self.spec_chunk_rows() else { return Vec::new(); }; - session_boundaries(chunk, prompt, self.turn_end_token_active(), after, up_to) + stride_boundaries(chunk, prompt, self.turn_end_token_active(), after, up_to) } fn snapshot_parts( @@ -433,11 +341,16 @@ impl SessionState for Qwen4Bundle { .map_err(|e| e.to_string()) } - fn growth_reserve_bytes(&self, gpu: &Gpu) -> u64 { - let through = self - .turn - .as_ref() - .map_or(self.state.position, |(prompt, _)| prompt.len()); + fn live_position(&self, route: SessionRoute) -> Option { + live_position( + route, + self.state.position, + self.mtp.as_ref().map(Qwen4MtpGpu::position), + self.state.row_capture_armed, + ) + } + + fn growth_reserve_bytes(&self, gpu: &Gpu, through: usize) -> u64 { let (context, gather) = self.context_growth_bytes(gpu, through); context.saturating_add(gather) } @@ -599,8 +512,6 @@ impl Qwen4Bundle { spec_host_top1: Vec::new(), penalty_stage: None, session: None, - live: None, - turn: None, turn_snapshots: hipfire_config::developer_bool(TURN_SNAPSHOTS_ENV, false), turn_end_token: None, state_format, @@ -1969,8 +1880,9 @@ impl Qwen4Bundle { /// Reset every owner for a cold prefill. pub fn reset(&mut self, gpu: &mut Gpu) -> Result<(), BundleError> { - self.live = None; - self.turn = None; + if let Some(session) = self.session.as_mut() { + session.forget_live(); + } self.invalidate_ple_epoch()?; self.state.reset(gpu).map_err(BundleError::State)?; if let Some(mtp) = self.mtp.as_mut() { @@ -2081,55 +1993,35 @@ impl Qwen4Bundle { /// Attach the session cache (`hipfire_runtime::session_cache`): prefill /// then restores and captures whole-chunk snapshots through it. pub fn attach_session_cache(&mut self, cache: SessionCache) { - self.session = Some(cache); + self.session = Some(SessionDriver::new(cache)); } /// The attached session cache, if any. pub fn session_cache(&self) -> Option<&SessionCache> { - self.session.as_ref() + self.session.as_ref().map(SessionDriver::cache) } /// Drop every session snapshot, pending ones included, releasing their - /// device buffers. No-op without a cache. + /// device buffers, and the live state. No-op without a cache. pub fn session_clear(&mut self, gpu: &mut Gpu) { - self.live = None; - self.turn = None; - if let Some(cache) = self.session.as_mut() { - cache.clear(gpu); + if let Some(session) = self.session.as_mut() { + session.clear(gpu); } } - /// Run `f` with the cache taken out, so it can drive `self` as its + /// Run `f` with the driver taken out, so it can drive `self` as its /// [`SessionState`]. `None` without a cache. - fn with_session(&mut self, f: impl FnOnce(&mut SessionCache, &mut Self) -> R) -> Option { - let mut cache = self.session.take()?; - let result = f(&mut cache, self); - self.session = Some(cache); + fn with_session(&mut self, f: impl FnOnce(&mut SessionDriver, &mut Self) -> R) -> Option { + let mut session = self.session.take()?; + let result = f(&mut session, self); + self.session = Some(session); Some(result) } - /// The live state's end when it can serve `prompt` on `route` now. - fn live_hit(&self, prompt: &[u32], route: SessionRoute) -> Option { - live_start( - self.live.as_ref(), - prompt, - route, - self.state.position, - self.mtp.as_ref().map(Qwen4MtpGpu::position), - self.state.row_capture_armed, - ) - } - - /// Start a prefill of `prompt` on `route` that skips `reused` tokens: - /// - live: `reused` is the end of the committed live state (this - /// conversation's previous turn); every owner and position is kept, no - /// snapshot is restored, and nothing is captured unless message-end - /// snapshots are active (then the boundaries above `reused` are, in the - /// `/turns` scope); - /// - snapshot: restore the `reused`-token snapshot the cache planned; - /// - `reused == 0`: cold start. - /// - /// The live record is consumed while the request is in flight. + /// Start a prefill of `prompt` on `route` that skips `reused` tokens + /// ([`SessionDriver::begin`]): a live continuation of this conversation's + /// previous turn, a snapshot restore, or a cold start. Message-end + /// boundaries of a live continuation are captured when active. pub(crate) fn session_begin( &mut self, gpu: &mut Gpu, @@ -2139,35 +2031,20 @@ impl Qwen4Bundle { ) -> Result<(), BundleError> { // A lookahead left by an aborted request names another prompt's rows. self.set_ple_lookahead(&[]); - let live = - reused > 0 && self.session.is_some() && self.live_hit(prompt, route) == Some(reused); - self.live = None; - self.turn = None; let capture = self.turn_snapshots_active(); - let result = if live { + match self.with_session(|session, bundle| { + session.begin(gpu, bundle, prompt, route, reused, capture) + }) { // The kept state continues, so request-local PLE work from the // previous request must not publish into it. - self.invalidate_ple_epoch().map(|()| { - self.with_session(|cache, bundle| { - cache.begin_live(gpu, bundle, prompt, route, reused, capture) - }); - }) - } else { - match self.with_session(|cache, bundle| cache.begin(gpu, bundle, prompt, route, reused)) - { - Some(result) => result.map_err(BundleError::Forward), - None if reused == 0 => self.reset(gpu), - None => Err(BundleError::Forward(format!( - "Qwen4 session cache is not attached; cannot reuse {reused} tokens" - ))), - } - }; - if result.is_ok() { - self.turn = Some((prompt.to_vec(), route)); - } else { - self.live = None; + Some(Ok(true)) => self.invalidate_ple_epoch(), + Some(Ok(false)) => Ok(()), + Some(Err(error)) => Err(BundleError::Forward(error)), + None if reused == 0 => self.reset(gpu), + None => Err(BundleError::Forward(format!( + "Qwen4 session cache is not attached; cannot reuse {reused} tokens" + ))), } - result } /// Next absolute position the running prefill must stop at and report. @@ -2181,7 +2058,7 @@ impl Qwen4Bundle { gpu: &mut Gpu, prefix: &[u32], ) -> Result<(), BundleError> { - self.with_session(|cache, bundle| cache.at_boundary(gpu, bundle, prefix)) + self.with_session(|session, bundle| session.at_boundary(gpu, bundle, prefix)) .unwrap_or(Ok(())) .map_err(BundleError::Forward) } @@ -2841,20 +2718,15 @@ impl hipfire_runtime::arch_model::ArchModel for Qwen4Bundle { } fn session_plan(&self, prompt: &[u32], route: SessionRoute) -> usize { - self.session.as_ref().map_or(0, |cache| { - plan_start( - cache.plan(self, prompt, route), - self.live_hit(prompt, route), - ) - }) + self.session + .as_ref() + .map_or(0, |session| session.plan(self, prompt, route)) } fn session_commit(&mut self) { - if let Some(cache) = self.session.as_mut() { - cache.commit(); + if let Some(session) = self.session.as_mut() { + session.commit(); } - self.live = None; - self.turn = None; } /// Publish as [`Self::session_commit`] and keep the live state when @@ -2864,23 +2736,10 @@ impl hipfire_runtime::arch_model::ArchModel for Qwen4Bundle { /// decode lineage); the snapshots just published stay cold-exact because a /// live turn captures none unless message-end snapshots are active. fn session_commit_live(&mut self, consumed: &[u32]) { - let Some(cache) = self.session.as_mut() else { - self.live = None; - self.turn = None; - return; - }; - cache.commit(); - let position = self.state.position; - let head = self.mtp.as_ref().map(Qwen4MtpGpu::position); - let armed = self.state.row_capture_armed; - self.live = self.turn.take().and_then(|(prompt, route)| { - live_after_commit(&prompt, route, consumed, position, head, armed).then(|| { - Qwen4LiveRecord { - tokens: consumed.to_vec(), - route, - } - }) - }); + if let Some(mut session) = self.session.take() { + session.commit_live(self, consumed); + self.session = Some(session); + } } fn free_gpu(self: Box, gpu: &mut Gpu) { @@ -3577,218 +3436,19 @@ mod tests { const AR: SessionRoute = SessionRoute::Ar; const MTP: SessionRoute = SessionRoute::Mtp; - fn record(tokens: &[u32], route: SessionRoute) -> Qwen4LiveRecord { - Qwen4LiveRecord { - tokens: tokens.to_vec(), - route, - } - } - - const END: u32 = 7; - - #[test] - fn message_end_boundaries_follow_every_im_end() { - // None. - assert!(message_end_boundaries(&[1, 2, 3], END, 0, 3).is_empty()); - assert!(message_end_boundaries(&[], END, 0, 0).is_empty()); - // Several: the position just after each im_end. - let prompt = [1, END, 2, 3, END, 4, END]; - assert_eq!(message_end_boundaries(&prompt, END, 0, 7), [2, 5, 7]); - // `after` is exclusive, `up_to` inclusive. - assert_eq!(message_end_boundaries(&prompt, END, 2, 7), [5, 7]); - assert_eq!(message_end_boundaries(&prompt, END, 1, 7), [2, 5, 7]); - assert_eq!(message_end_boundaries(&prompt, END, 0, 5), [2, 5]); - assert_eq!(message_end_boundaries(&prompt, END, 0, 4), [2]); - assert_eq!(message_end_boundaries(&prompt, END, 5, 7), [7]); - assert!(message_end_boundaries(&prompt, END, 7, 7).is_empty()); - assert!(message_end_boundaries(&prompt, END, 3, 4).is_empty()); - // Consecutive im_end tokens each end a message. - assert_eq!( - message_end_boundaries(&[END, END, 1, END], END, 0, 4), - [1, 2, 4] - ); - // im_end at the last position ends at the prompt length. - assert_eq!(message_end_boundaries(&[1, 2, END], END, 0, 3), [3]); - // `up_to` beyond the prompt never reads past it. - assert_eq!(message_end_boundaries(&[1, END], END, 0, 9), [2]); - } - - #[test] - fn session_boundaries_union_chunks_and_message_ends() { - let mut prompt = vec![1u32; 20]; - prompt[2] = END; // ends at 3 - prompt[7] = END; // ends at 8: a chunk multiple - prompt[13] = END; // ends at 14 - // Chunk multiples only without a token (flag off or token unknown). - assert_eq!( - session_boundaries(4, &prompt, None, 0, 20), - [4, 8, 12, 16, 20] - ); - // Union, sorted and deduplicated (8 is both). - assert_eq!( - session_boundaries(4, &prompt, Some(END), 0, 20), - [3, 4, 8, 12, 14, 16, 20] - ); - // The `(after, up_to]` window applies to both sources. - assert_eq!( - session_boundaries(4, &prompt, Some(END), 3, 14), - [4, 8, 12, 14] - ); - assert_eq!(session_boundaries(4, &prompt, None, 3, 14), [4, 8, 12]); - // Message ends below one chunk are boundaries on their own. - assert_eq!( - session_boundaries(64, &prompt, Some(END), 0, 20), - [3, 8, 14] - ); - assert!(session_boundaries(64, &prompt, None, 0, 20).is_empty()); - // A prompt without the token matches the chunk-only schedule. - assert_eq!( - session_boundaries(4, &[1; 20], Some(END), 0, 20), - session_boundaries(4, &[1; 20], None, 0, 20) - ); - } - - #[test] - fn live_start_requires_a_strict_extension_of_the_record() { - let rec = record(&[1, 2, 3], AR); - let at = |prompt: &[u32]| live_start(Some(&rec), prompt, AR, 3, None, false); - assert_eq!(at(&[1, 2, 3, 4]), Some(3)); - assert_eq!(at(&[1, 2, 3, 4, 5, 6]), Some(3)); - // Diverges inside the record. - assert_eq!(at(&[1, 9, 3, 4]), None); - assert_eq!(at(&[9, 2, 3, 4]), None); - // L == prompt.len() leaves no suffix; a shorter prompt is no extension. - assert_eq!(at(&[1, 2, 3]), None); - assert_eq!(at(&[1, 2]), None); - assert_eq!(at(&[]), None); - // No record, or an empty one (L == 0). - assert_eq!(live_start(None, &[1, 2, 3, 4], AR, 3, None, false), None); - let empty = record(&[], AR); - assert_eq!(live_start(Some(&empty), &[1, 2], AR, 0, None, false), None); - } - - #[test] - fn live_start_checks_route_positions_and_row_capture() { - let prompt = [1, 2, 3, 4]; - let ar = record(&[1, 2, 3], AR); - let mtp = record(&[1, 2, 3], MTP); - // Route mismatch, both ways. - assert_eq!(live_start(Some(&ar), &prompt, MTP, 3, Some(3), false), None); - assert_eq!(live_start(Some(&mtp), &prompt, AR, 3, Some(3), false), None); - // Target moved off the record's end. - assert_eq!(live_start(Some(&ar), &prompt, AR, 2, None, false), None); - assert_eq!(live_start(Some(&ar), &prompt, AR, 4, None, false), None); - // A verify row capture in flight. - assert_eq!(live_start(Some(&ar), &prompt, AR, 3, None, true), None); - assert_eq!(live_start(Some(&mtp), &prompt, MTP, 3, Some(3), true), None); - // AR ignores the head; MTP needs it exactly at L (floor retirement - // leaves it behind, a missing head cannot continue). - assert_eq!( - live_start(Some(&ar), &prompt, AR, 3, Some(1), false), - Some(3) - ); - assert_eq!( - live_start(Some(&mtp), &prompt, MTP, 3, Some(3), false), - Some(3) - ); - assert_eq!( - live_start(Some(&mtp), &prompt, MTP, 3, Some(2), false), - None - ); - assert_eq!( - live_start(Some(&mtp), &prompt, MTP, 3, Some(4), false), - None - ); - assert_eq!(live_start(Some(&mtp), &prompt, MTP, 3, None, false), None); - } - - #[test] - fn live_after_commit_needs_consumed_to_extend_the_turn_prompt() { - let prompt = [1, 2, 3]; - let commit = |consumed: &[u32], pos: usize| { - live_after_commit(&prompt, AR, consumed, pos, None, false) - }; - assert!(commit(&[1, 2, 3, 7, 8], 5)); - // Nothing generated: consumed == prompt. - assert!(commit(&[1, 2, 3], 3)); - // Unrepaired terminal: the host history is empty. - assert!(!commit(&[], 0)); - assert!(!commit(&[], 3)); - // Shorter than the prompt, or not extending it. - assert!(!commit(&[1, 2], 2)); - assert!(!commit(&[1, 9, 3, 7], 4)); - // Device position not at the end of the consumed history. - assert!(!commit(&[1, 2, 3, 7, 8], 4)); - assert!(!commit(&[1, 2, 3, 7, 8], 6)); - // A turn without a prompt is never live. - assert!(!live_after_commit(&[], AR, &[1], 1, None, false)); - } - - #[test] - fn live_after_commit_checks_head_and_row_capture() { - let prompt = [1, 2, 3]; - let consumed = [1, 2, 3, 7, 8]; - assert!(live_after_commit( - &prompt, - MTP, - &consumed, - 5, - Some(5), - false - )); - // Floor retirement left the head behind (or absent). - assert!(!live_after_commit( - &prompt, - MTP, - &consumed, - 5, - Some(3), - false - )); - assert!(!live_after_commit(&prompt, MTP, &consumed, 5, None, false)); - // AR does not care about the head. - assert!(live_after_commit(&prompt, AR, &consumed, 5, Some(3), false)); - assert!(live_after_commit(&prompt, AR, &consumed, 5, None, false)); - // Row capture armed. - assert!(!live_after_commit(&prompt, AR, &consumed, 5, None, true)); - assert!(!live_after_commit( - &prompt, - MTP, - &consumed, - 5, - Some(5), - true - )); - } - - #[test] - fn plan_prefers_live_on_a_tie_and_the_longer_source_otherwise() { - // No live state: the snapshot (or a miss). - assert_eq!(plan_start(4096, None), 4096); - assert_eq!(plan_start(0, None), 0); - // Live alone. - assert_eq!(plan_start(0, Some(300)), 300); - // Tie prefers live (no restore copy). - assert_eq!(plan_start(4096, Some(4096)), 4096); - // The longer source wins either way. - assert_eq!(plan_start(4096, Some(9000)), 9000); - assert_eq!(plan_start(8192, Some(300)), 8192); - } - #[test] - fn live_hit_then_commit_round_trips_through_the_pure_helpers() { - // Turn 1: cold prompt of 3, decode two tokens, commit. - let turn1 = [1u32, 2, 3]; - let consumed1 = [1u32, 2, 3, 7, 8]; - assert!(live_after_commit(&turn1, AR, &consumed1, 5, None, false)); - let rec = record(&consumed1, AR); - // Turn 2 re-renders the history and appends the next user turn. - let turn2 = [1u32, 2, 3, 7, 8, 11, 12]; - let start = live_start(Some(&rec), &turn2, AR, 5, None, false); - assert_eq!(start, Some(5)); - assert_eq!(plan_start(0, start), 5); - // A re-rendered history that rewrote the assistant turn diverges. - let rewritten = [1u32, 2, 3, 70, 8, 11, 12]; - assert_eq!(live_start(Some(&rec), &rewritten, AR, 5, None, false), None); + fn live_position_checks_the_head_and_row_capture() { + // AR continues at the target and ignores the head. + assert_eq!(live_position(AR, 3, None, false), Some(3)); + assert_eq!(live_position(AR, 3, Some(1), false), Some(3)); + // MTP needs the head exactly at the target (floor retirement leaves + // it behind, a missing head cannot continue). + assert_eq!(live_position(MTP, 3, Some(3), false), Some(3)); + assert_eq!(live_position(MTP, 3, Some(2), false), None); + assert_eq!(live_position(MTP, 3, Some(4), false), None); + assert_eq!(live_position(MTP, 3, None, false), None); + // A verify row capture in flight never continues. + assert_eq!(live_position(AR, 3, None, true), None); + assert_eq!(live_position(MTP, 3, Some(3), true), None); } } diff --git a/crates/hipfire-daemon/src/main.rs b/crates/hipfire-daemon/src/main.rs index 3abf9f6b80..7af2c786be 100644 --- a/crates/hipfire-daemon/src/main.rs +++ b/crates/hipfire-daemon/src/main.rs @@ -2500,16 +2500,11 @@ fn main() { batch_scheduler = staged_batch_scheduler; vmm_batch = if staging.vmm { parsed_vmm_batch } else { None }; // `cache_capable` is the daemon's prompt-cache source of truth. - // arch_id 13 (gemma4) is intentionally ABSENT: hipfire_generate::dense::generate_gemma4 has - // no LCP prefix-cache block and always cold-prefills the full - // Jinja-rendered prompt. Enabling the cache would corrupt KV - // slot offsets after turn 1 (stale prefix reuse). Wire when - // hipfire_generate::dense::generate_gemma4 gains an LCP block matching other archs. - // Qwen4 (16) only with its attached session cache - // (`ArchModel::session_cache_attached`). + // Gemma 4 (13) and Qwen4 (16) only with their attached + // session cache (`ArchModel::session_cache_attached`). let cache_capable = match m.arch_id { 5 | 6 | 9 | 10 | 12 | 14 => true, - 16 => m.state.as_deref().is_some_and(|state| state.session_cache_attached()), + 13 | 16 => m.state.as_deref().is_some_and(|state| state.session_cache_attached()), _ => false, }; let retry_reset_eligible = model_retry_reset_eligible(m.arch_id); diff --git a/crates/hipfire-dispatch/map.md b/crates/hipfire-dispatch/map.md index 1b4f7b77e6..094bbbdf49 100644 --- a/crates/hipfire-dispatch/map.md +++ b/crates/hipfire-dispatch/map.md @@ -24,7 +24,7 @@ _Generated by `scripts/check-crate-maps.py` from the tree — do not edit inside - [`src/context.rs`](src/context.rs): `DispatchCtx`, `DispatchWorkload`, `device_id`, `for_test`, `for_test_device`, `new`, `with_verify_tile_attend`, `with_workload` - [`src/coverage_tests.rs`](src/coverage_tests.rs): — - [`src/cpu_exec.rs`](src/cpu_exec.rs): `cpu_exec_counters`, `cpu_exec_enabled`, `cpu_exec_redline_conflict`, `cpu_offload_active`, `cpu_quant_for`, `host_mapped_cpu_capable`, `log_capture_disabled_once`, `reject_cpu_exec_under_redline`, `run_host_mapped_gemv`, `run_host_mapped_gemv_residual` -- [`src/families/attention.rs`](src/families/attention.rs): `AttentionFamily`, `AttnParams`, `FullAttnParams`, `new`, `positions`, `registry`, `resolve`, `run_attention`, `run_full_attention` +- [`src/families/attention.rs`](src/families/attention.rs): `AttentionFamily`, `AttnParams`, `FullAttnParams`, `new`, `positions`, `registry`, `resolve`, `run_attend_only`, `run_attention`, `run_full_attention` - [`src/families/fused_qkv.rs`](src/families/fused_qkv.rs): `FusedQkvBiasParams`, `FusedQkvFamily`, `FusedQkvParams`, `fused_gate_up_key_for`, `new`, `registry`, `resolve`, `run`, `run_with_qwen2_bias` - [`src/families/gemm.rs`](src/families/gemm.rs): `GemmFamily`, `GemmParams`, `new`, `registry`, `residual_gemm_key_for`, `resolve`, `run`, `run_key`, `run_q8_projection_fresh` - [`src/families/gemv.rs`](src/families/gemv.rs): `GemvFamily`, `GemvParams`, `GivensRef`, `RotInput`, `RotateInputs`, `RotatedActivation`, `RotationTag`, `WeightRef`, `buf`, `check_rotation_tag`, `into_buf`, `new`, `registry`, `resolve`, `rotate`, `run`, `run_auto`, `run_input`, `select_rotation_variant`, `tag` @@ -45,13 +45,14 @@ _Generated by `scripts/check-crate-maps.py` from the tree — do not edit inside - [`src/ops/verify_twins.rs`](src/ops/verify_twins.rs): `attention_fp8_segs`, `rdna_compute`, `stage_attention_fp8_segs` - [`src/pipeline/draft_head.rs`](src/pipeline/draft_head.rs): `DraftHead`, `DraftHeadLayout`, `DraftHeadPolicy`, `DraftHeadRequestState`, `device_bytes`, `draft`, `draft_enqueue`, `draft_finish`, `free_gpu`, `logits`, `margin`, `new`, `observe`, `parse`, `request_state`, `rescored_candidates`, `reset_request_state`, `set_request_state` - [`src/pipeline/layer_ops.rs`](src/pipeline/layer_ops.rs): `BroadcastAddOp`, `ClearOp`, `EmbeddingOp`, `GatedDeltaNetOp`, `GdnRowCapture`, `GroupedDepthwiseOp`, `HyperNormOp`, `HyperReadOp`, `HyperWriteOp`, `IndexedAttentionMode`, `IndexedAttentionOp`, `IndexedAttentionState`, `MixerPhase`, `ProjectOp`, `QsaProjectionHook`, `execute_argmax`, `execute_broadcast_add`, `execute_clear`, `execute_embedding`, `execute_final_hyper`, `execute_gated_delta_net`, `execute_gated_delta_net_hc`, `execute_grouped_depthwise`, `execute_hyper_norm`, `execute_hyper_read`, `execute_hyper_read_paired`, `execute_hyper_read_prenormed`, `execute_hyper_write`, `execute_hyper_write_pregated`, `execute_hyper_write_then_read`, `execute_indexed_attention`, `execute_indexed_attention_hc`, `execute_lm_head`, `execute_project`, `hyper_gate_quarters`, `hyper_read_pairs_write`, `hyper_read_prenorm_applies`, `next_lengths`, `project_weight`, `project_weights`, `qsa_projection_hook_installed`, `reset_qsa_projection_slot`, `set_qsa_projection_hook`, `validate_for_gpu`, `validate_lm_head` -- [`src/pipeline/mod.rs`](src/pipeline/mod.rs): `LinearParams`, `Pipeline`, `PipelineParams`, `can_satisfy`, `check_moe_decode_batch_size`, `check_moe_decode_supported`, `dispatch_fused`, `draft_head`, `dump_hidden_localize`, `execute_pipeline`, `layer_ops`, `moe_dump`, `moe_program`, `new`, `run_moe_decode_bias_aware`, `run_moe_decode_selected`, `run_moe_prefill_bias_aware`, `run_uniform_moe_down_expanded`, `run_uniform_moe_gate_up`, `sealed_moe`, `steps`, `superop` +- [`src/pipeline/mod.rs`](src/pipeline/mod.rs): `LinearParams`, `Pipeline`, `PipelineParams`, `can_satisfy`, `check_moe_decode_batch_size`, `check_moe_decode_supported`, `dispatch_fused`, `draft_head`, `dump_hidden_localize`, `execute_pipeline`, `layer_ops`, `moe_dump`, `moe_program`, `new`, `run_moe_decode_bias_aware`, `run_moe_decode_selected`, `run_moe_prefill_bias_aware`, `run_uniform_moe_down_expanded`, `run_uniform_moe_gate_up`, `sandwich`, `sealed_moe`, `steps`, `superop` - [`src/pipeline/moe_dump.rs`](src/pipeline/moe_dump.rs): `FoldInfo`, `FoldPending`, `SharedPending`, `enabled`, `final_residual`, `final_residual_deferred`, `fold_begin`, `fold_end`, `set_context`, `shared_begin`, `shared_end` - [`src/pipeline/moe_program.rs`](src/pipeline/moe_program.rs): `MoeStage`, `SealedMoeOp` - [`src/pipeline/qt44_qt53_prefill.rs`](src/pipeline/qt44_qt53_prefill.rs): — +- [`src/pipeline/sandwich.rs`](src/pipeline/sandwich.rs): `Activation`, `KvProjection`, `ParallelMoeMlpOp`, `PerLayerInputOp`, `Rope`, `RopeKind`, `RoutedExperts`, `RoutedScratch`, `SandwichAttentionOp`, `SandwichKv`, `SandwichMlpOp`, `SandwichStream`, `ScaleOp`, `SoftcapOp`, `execute_parallel_moe_mlp`, `execute_per_layer_input`, `execute_sandwich_attention`, `execute_sandwich_mlp`, `execute_scale`, `execute_softcap`, `gemm_rows`, `supports`, `supports_batched_projection` - [`src/pipeline/sealed_moe.rs`](src/pipeline/sealed_moe.rs): `ActivationIdentity`, `ActivationInput`, `BoundMoeExperts`, `CompactLiveWeight`, `ContractAssignment`, `ContractAxis`, `ContractCollectiveHint`, `ContractCollectiveRow`, `ContractParallelism`, `ExpertBindingCache`, `ExpertExecutionContract`, `ExpertMetadata`, `ExpertResource`, `ExpertResources`, `ExpertTable`, `MoeContribution`, `MoePrefillRouteProducerProof`, `MoeProtocol`, `MoeRouteProducerProof`, `MoeRouterInput`, `MoeSharedContribution`, `PrefillRouteMode`, `ROOT_ROUTED_EP_EXECUTION`, `ResourceAlias`, `RetainedBodyAction`, `SealedMoeCall`, `SpecializedRouteAdmission`, `activation`, `alias`, `alignment`, `arm_fault_after_expert_mutation`, `assignment`, `basis`, `bind_live`, `bind_live_compact`, `byte_len`, `byte_offset`, `checked_grouped_m_total_bound`, `checked_tp_local_shapes`, `collective_rows`, `contract_id`, `contribution`, `decode_params`, `down`, `dtype`, `encoded_bytes`, `ep_route_producer_proof`, `execute_ep_slot_combine`, `execution`, `execution_contract`, `expert`, `expert_weights_match_metadata`, `experts`, `fingerprint`, `from_cache`, `fused`, `gate`, `gate_up`, `global_id`, `group`, `input`, `invocation`, `is_fused`, `is_live_bound`, `is_root_routed_ep`, `layer`, `local_expert_ids`, `local_rank`, `local_slot`, `local_slots`, `mapping_fingerprint`, `mesh_epoch`, `n_experts`, `new`, `owner`, `owner_rank`, `owner_ranks`, `parallelism`, `physical_device`, `physical_devices`, `prefill_params`, `preflight_decode`, `preflight_prefill`, `prepare_binding`, `protocol`, `rank_count`, `records`, `resources`, `retained_body_action`, `rotation`, `router_input`, `row_stride`, `seal_decode`, `seal_ep_routed_contrib`, `seal_prefill`, `seal_prefill_ep`, `separate`, `shape`, `shared_contribution`, `sidecars`, `source_fingerprint`, `source_name`, `specialized_sealed_moe_retained_admission`, `take_fault_after_expert_mutation`, `up`, `with_alias`, `with_execution_contract`, `with_sidecars` - [`src/pipeline/steps.rs`](src/pipeline/steps.rs): `FusedPattern`, `GemvInput`, `MoeStepHooks`, `Step`, `execute_steps`, `execute_validated_steps`, `execute_validated_steps_with_moe_hooks`, `match_prefix`, `validate_steps` -- [`src/pipeline/superop.rs`](src/pipeline/superop.rs): `ActFlavor`, `AttnFlavor`, `EpMoeCombineMode`, `EpMoeRouteView`, `EscapeKind`, `ForwardBindings`, `LayerProgram`, `LoweredForward`, `OpBinding`, `OpFlavor`, `RopeFlavor`, `ScratchSlot`, `SuperOp`, `SuperOpKind`, `WeightSlot`, `contract_id`, `dispatch_super_op`, `execution_contract`, `lower_layer`, `new`, `run_layer_program` +- [`src/pipeline/superop.rs`](src/pipeline/superop.rs): `EpMoeCombineMode`, `EpMoeRouteView`, `EscapeKind`, `ForwardBindings`, `LayerProgram`, `LoweredForward`, `OpBinding`, `OpFlavor`, `ScratchSlot`, `SuperOp`, `SuperOpKind`, `WeightSlot`, `contract_id`, `dispatch_super_op`, `execution_contract`, `lower_layer`, `new`, `run_layer_program` - [`src/resource/mod.rs`](src/resource/mod.rs): `ResourceManager`, `for_test`, `new` - [`src/tables/attention_table.rs`](src/tables/attention_table.rs): `populate` - [`src/tables/fused_qkv_table.rs`](src/tables/fused_qkv_table.rs): `populate` diff --git a/crates/hipfire-dispatch/src/families/gemm.rs b/crates/hipfire-dispatch/src/families/gemm.rs index 54ce25477a..2e54d5b8a2 100644 --- a/crates/hipfire-dispatch/src/families/gemm.rs +++ b/crates/hipfire-dispatch/src/families/gemm.rs @@ -317,7 +317,7 @@ impl GemmFamily { // Calibration tap. This is the batched chokepoint for every arch that // migrated off `llama::weight_gemm` onto dispatcher-entry keys — - // gemma4 (`lowered.rs:173`), muse-glimmer (`forward.rs:80`) and + // gemma4 (the dispatch sandwich projections), muse-glimmer (`forward.rs:80`) and // qwen35's migrated prefill sites. Without it those arches capture // NOTHING on the batched path, because `weight_gemm`'s tap never runs // for them. `batch_size` is the real row count, so one prefill call diff --git a/crates/hipfire-dispatch/src/pipeline/layer_ops.rs b/crates/hipfire-dispatch/src/pipeline/layer_ops.rs index 6f19bd90a1..d13c6bb5ba 100644 --- a/crates/hipfire-dispatch/src/pipeline/layer_ops.rs +++ b/crates/hipfire-dispatch/src/pipeline/layer_ops.rs @@ -621,8 +621,7 @@ fn project_rotated( // dequantised weights to F16. The trunk is Q8 precisely because these // projections write straight into the residual stream, and decode reads // them with the F32 `gemv_q8_0` — an F16-rounded prefill would also - // disagree with decode. Same reasoning as gemma4's `lowered` path, which - // maps Q8_0 to `GemmQ8_0BatchedF32Chunked` for the same reason. + // disagree with decode. (DType::Q8_0, true) => { gpu.gemm_q8_0_batched_f32_chunked(weight.buf, x, output, weight.m, weight.k, rows) } @@ -3145,6 +3144,8 @@ enum EmbeddingPath { Mq4V2, Bf16, Q8, + Hfq4G256, + Hfq4G128, } fn embedding_path(dtype: DType) -> Option { @@ -3152,6 +3153,8 @@ fn embedding_path(dtype: DType) -> Option { DType::MQ4G256V2 | DType::MQ4G128V2 => Some(EmbeddingPath::Mq4V2), DType::BF16 => Some(EmbeddingPath::Bf16), DType::Q8_0 => Some(EmbeddingPath::Q8), + DType::HFQ4G256 => Some(EmbeddingPath::Hfq4G256), + DType::HFQ4G128 => Some(EmbeddingPath::Hfq4G128), _ => None, } } @@ -3197,6 +3200,12 @@ pub fn execute_embedding(gpu: &mut Gpu, op: &EmbeddingOp<'_>) -> Result<(), Disp gpu.embedding_lookup_bf16_batched(table, output, ids, op.rows, op.dim) } EmbeddingPath::Q8 => gpu.embedding_lookup_q8_batched(table, output, ids, op.rows, op.dim), + EmbeddingPath::Hfq4G256 => { + gpu.embedding_lookup_hfq4g256_batched(table, output, ids, op.rows, op.dim) + } + EmbeddingPath::Hfq4G128 => { + gpu.embedding_lookup_hfq4g128_batched(table, output, ids, op.rows, op.dim) + } }) } diff --git a/crates/hipfire-generate/map.md b/crates/hipfire-generate/map.md index 60abc1cb51..33934e6439 100644 --- a/crates/hipfire-generate/map.md +++ b/crates/hipfire-generate/map.md @@ -24,7 +24,7 @@ _Generated by `scripts/check-crate-maps.py` from the tree — do not edit inside - [`src/ar.rs`](src/ar.rs): `ALL`, `GenerationRoute`, `GenerationRouteAdapter`, `GenerationRouteInputs`, `GenerationRouteScope`, `QwenArCacheAction`, `QwenArForwardFailAction`, `QwenArRawCommitDisposition`, `QwenArRouteFinish`, `QwenArSemanticProducer`, `QwenArTerminalCause`, `QwenArThinkBudgetAction`, `RouteStartAdapter`, `RouteTerminal`, `RouteTerminalAdapter`, `RouteTerminalEvent`, `active_generation_route`, `arm_fault_after_prefill`, `ckpt_interval`, `ckpt_max`, `ckpt_resume_enabled`, `commit_and_classify`, `commit_and_observe`, `commit_raw`, `deepseek4_spec_requested`, `deepseek4_spec_requested_from_policy`, `emit_active_route_cancel`, `emit_active_route_done`, `emit_active_route_done_value`, `emit_active_route_error`, `emit_cancel`, `emit_done`, `emit_error`, `emit_generation_cancel`, `emit_generation_done`, `emit_generation_done_value`, `emit_generation_error`, `emit_generation_start`, `emit_qwen_ar_done`, `emit_start`, `emit_start_with`, `emit_terminal`, `enter`, `finish`, `fn`, `generate`, `generate_ar_with_forward`, `generation_route_adapter`, `llama_prefill_sample_seed`, `llama_qwen3_batched_prefill_eligible`, `model_retry_reset_eligible`, `new`, `new_with_tool_protocol`, `pflash_done_fragment`, `qwen_ar_apply_cache_action`, `qwen_ar_cache_action`, `qwen_ar_done_value`, `qwen_ar_drain_pending_into_router`, `qwen_ar_evict`, `qwen_ar_eviction_fail_closed`, `qwen_ar_eviction_prefill_chunk_limit`, `qwen_ar_finish_route`, `qwen_ar_forward_fail_action`, `qwen_ar_forward_fail_message`, `qwen_ar_observe_and_route`, `qwen_ar_raw_commit_token`, `qwen_ar_route_filter_text`, `qwen_ar_route_think_events`, `qwen_ar_think_budget_step`, `reset_core_arch_key`, `resolve`, `select_generation_route`, `set_generation_route`, `stop_or_reject`, `take_fault_after_prefill`, `truncate_checkpoints`, `visible`, `with_stop`, `write_error` - [`src/batch.rs`](src/batch.rs): `PromotedDecode`, `PromotedRequest`, `PromotionPermit`, `Qwen4PromotedRequest`, `arm_promotion`, `attach_qwen_ep_batch_receipt_evidence`, `batch_request_eligible_for_route`, `drive_lfm_continuous_batch`, `drive_promoted`, `drive_qwen35_ep_continuous_batch`, `drive_qwen_continuous_batch`, `drive_qwen_vmm_continuous_batch`, `drive_staged_continuous_batch`, `emit_uncorrelated_error`, `is_batch_request_eligible`, `is_qwen_ep_batch_request_eligible`, `is_vmm_batch_request_eligible`, `lfm_prefill_cancellable_or_fallback`, `set_promotion_permit`, `take_promoted`, `take_promoted_qwen4`, `vmm_loaded_ack`, `vmm_request_fits_batch`, `vmm_route_active`, `vmm_staging_request` - [`src/common.rs`](src/common.rs): `DEEPSEEK4_REASONING_HIGH_PREFIX`, `DEEPSEEK4_REASONING_MAX_PREFIX`, `Ds4ArEpRouteTerminal`, `Ds4CacheAction`, `Ds4ClientCommitEffects`, `Ds4MalformedTerminalAction`, `FIT_MAX_TOKENS_MARGIN`, `FitMaxTokensGuard`, `RequestFaultGuard`, `RollbackEpilogue`, `SpecEmitRequest`, `SpecRun`, `arm_dots_fault`, `arm_eviction_fault`, `arm_from_request`, `arm_generation_fault_after_first_decode`, `arm_generation_fault_after_prefill`, `arm_generation_faults_from_request`, `arm_vision_fault`, `asst_turn_fingerprint`, `attest_rollback_steps`, `build_deepseek4_dsml_prompt`, `deepseek4_reasoning_prefix`, `disarm_generation_faults`, `ds4_absorb_stream_event`, `ds4_apply_cache_action`, `ds4_ar_ep_cache_action`, `ds4_ar_ep_finish_route`, `ds4_client_commit_effects`, `ds4_ep_abort_wire_events`, `ds4_gen_start_contract_version`, `ds4_malformed_terminal_action`, `ds4_spec_finish_route`, `ds4_stream_event_wireable`, `emit_committed_event`, `emit_ds4_malformed_action`, `emit_ds4_malformed_terminal`, `emit_fail_closed_error`, `emit_fail_closed_error_for_route`, `emit_qwen_dflash_done_terminal`, `emit_spec_cancel_after_rollback`, `emit_stream_event`, `fail_closed_device_sync`, `fail_closed_epilogue_after_sync`, `fail_closed_invalidate_graphs_and_replay`, `fail_closed_reset_target_and_spec`, `finish_summary_held_tool_calls`, `fit_max_tokens`, `fit_max_tokens_if`, `free_checkpoints`, `gen_start_contract_version_for_arch`, `length_or_stop`, `maybe_inject_fault_after_first_decode_dflash`, `maybe_inject_fault_after_prefill_ar`, `maybe_inject_fault_after_prefill_dflash`, `normalize_asst_turn_for_fingerprint`, `peek_dots_fault`, `peek_vision_fault`, `production_fail_closed_rollback`, `production_fail_closed_rollback_live`, `production_pp_fail_closed_rollback`, `push_reset_err`, `qwen_dflash_epilogue_after_spec_run`, `qwen_dflash_hit_length_cap`, `reset_mesh_request_state`, `reset_qwen35_recurrent`, `reset_qwen4_recurrent`, `set`, `spec_ctx_request_fits`, `strip_think_for_fingerprint`, `take_eviction_fault`, `take_generation_fault_after_first_decode`, `take_generation_fault_after_prefill`, `token_logprob_fields` -- [`src/dense.rs`](src/dense.rs): `Ds4SpecWireTerminal`, `GLIMMER_END_OF_TEXT_ID`, `GLIMMER_EOM_ID`, `GLIMMER_EOT_ID`, `GLIMMER_MESSAGE_ID`, `GLIMMER_PROFIT_BAD_NUM`, `GLIMMER_PROFIT_BAD_TO_RETIRE`, `GLIMMER_PROFIT_GOOD_NUM`, `GLIMMER_PROFIT_RATIO_DEN`, `GLIMMER_PROFIT_WARMUP_WINDOWS`, `GLIMMER_PROFIT_WINDOWS_PER_EVAL`, `GLIMMER_START_ID`, `GemmaChannel`, `GemmaEmit`, `GemmaThoughtRouter`, `GlimmerChannel`, `GlimmerChannelRecorder`, `GlimmerEmit`, `GlimmerHarmonyRouter`, `GlimmerMirrorAction`, `GlimmerMirrorReconcile`, `GlimmerOpenBody`, `GlimmerProfitGuardStatus`, `GlimmerProfitLedger`, `GlimmerProfitProbeKind`, `GlimmerRecordRefusal`, `GlimmerRecordedRecipient`, `GlimmerRecorderState`, `GlimmerRouteError`, `GlimmerSpecMode`, `GlimmerSpecProfitGuard`, `MapleThoughtRouter`, `bad_evaluations`, `dense_fail_closed_error`, `ds4_ar_client_abort`, `ds4_cache_action`, `ds4_heterogeneous_client_abort`, `ds4_heterogeneous_fail_closed_error`, `ds4_spec_wire_terminal`, `eligible_windows`, `emit_active_attempt_error`, `emit_dense_token_text`, `enabled`, `evaluations`, `flush`, `gemma_is_marker_prefix`, `gemma_longest_marker_suffix`, `generate_cohere2moe`, `generate_deepseek4`, `generate_deepseek4_heterogeneous`, `generate_deepseek4_spec`, `generate_gemma4`, `generate_gemma4_lowered`, `generate_lfm2moe`, `generate_maple`, `generate_minimax`, `generate_muse_glimmer`, `generate_qwen2`, `glimmer_ar_forward_pending`, `glimmer_commit_terminal`, `glimmer_longest_marker_suffix`, `glimmer_mirror_action`, `glimmer_profit_guard_status_label`, `glimmer_profit_ledger_after_bonus_decode`, `glimmer_profit_ledger_post_window`, `glimmer_profit_ledger_route_prediction`, `glimmer_reasoning_strength`, `glimmer_spec_admission`, `glimmer_store_cached_turn`, `glimmer_tap_layers`, `glimmer_turn_key`, `into_cached_turn`, `is_marker_prefix`, `is_retired`, `just_forced`, `last_ar_probe_ns`, `last_productive`, `last_ratio_milli`, `last_spec_ns`, `mark_forced_reasoning_close`, `new`, `normalize_glimmer_tool_arguments`, `observe`, `observe_full_window`, `observe_probe`, `parse_glimmer_atem`, `pending_probe`, `prepare_glimmer_onyx_history`, `push`, `reconcile_glimmer_mirror`, `retire_cycle`, `retire_evaluation`, `status`, `thinking_enabled`, `write_error_envelope` +- [`src/dense.rs`](src/dense.rs): `Ds4SpecWireTerminal`, `GLIMMER_END_OF_TEXT_ID`, `GLIMMER_EOM_ID`, `GLIMMER_EOT_ID`, `GLIMMER_MESSAGE_ID`, `GLIMMER_PROFIT_BAD_NUM`, `GLIMMER_PROFIT_BAD_TO_RETIRE`, `GLIMMER_PROFIT_GOOD_NUM`, `GLIMMER_PROFIT_RATIO_DEN`, `GLIMMER_PROFIT_WARMUP_WINDOWS`, `GLIMMER_PROFIT_WINDOWS_PER_EVAL`, `GLIMMER_START_ID`, `GlimmerChannel`, `GlimmerChannelRecorder`, `GlimmerEmit`, `GlimmerHarmonyRouter`, `GlimmerMirrorAction`, `GlimmerMirrorReconcile`, `GlimmerOpenBody`, `GlimmerProfitGuardStatus`, `GlimmerProfitLedger`, `GlimmerProfitProbeKind`, `GlimmerRecordRefusal`, `GlimmerRecordedRecipient`, `GlimmerRecorderState`, `GlimmerRouteError`, `GlimmerSpecMode`, `GlimmerSpecProfitGuard`, `MapleThoughtRouter`, `bad_evaluations`, `dense_fail_closed_error`, `ds4_ar_client_abort`, `ds4_cache_action`, `ds4_heterogeneous_client_abort`, `ds4_heterogeneous_fail_closed_error`, `ds4_spec_wire_terminal`, `eligible_windows`, `emit_active_attempt_error`, `emit_dense_token_text`, `enabled`, `evaluations`, `flush`, `generate_cohere2moe`, `generate_deepseek4`, `generate_deepseek4_heterogeneous`, `generate_deepseek4_spec`, `generate_gemma4`, `generate_lfm2moe`, `generate_maple`, `generate_minimax`, `generate_muse_glimmer`, `generate_qwen2`, `glimmer_ar_forward_pending`, `glimmer_commit_terminal`, `glimmer_longest_marker_suffix`, `glimmer_mirror_action`, `glimmer_profit_guard_status_label`, `glimmer_profit_ledger_after_bonus_decode`, `glimmer_profit_ledger_post_window`, `glimmer_profit_ledger_route_prediction`, `glimmer_reasoning_strength`, `glimmer_spec_admission`, `glimmer_store_cached_turn`, `glimmer_tap_layers`, `glimmer_turn_key`, `into_cached_turn`, `is_marker_prefix`, `is_retired`, `just_forced`, `last_ar_probe_ns`, `last_productive`, `last_ratio_milli`, `last_spec_ns`, `mark_forced_reasoning_close`, `new`, `normalize_glimmer_tool_arguments`, `observe`, `observe_full_window`, `observe_probe`, `parse_glimmer_atem`, `pending_probe`, `prepare_glimmer_onyx_history`, `push`, `reconcile_glimmer_mirror`, `retire_cycle`, `retire_evaluation`, `status`, `thinking_enabled`, `write_error_envelope` - [`src/fn_batch.rs`](src/fn_batch.rs): `drive_qwen4_lane_batch` - [`src/img.rs`](src/img.rs): `generate_img` - [`src/lib.rs`](src/lib.rs): `ar`, `batch`, `common`, `dense`, `fn_batch`, `img`, `qwen`, `redline`, `vision`, `vmm_conv` diff --git a/crates/hipfire-generate/src/ar.rs b/crates/hipfire-generate/src/ar.rs index 2fa39cfef8..11a32870b8 100644 --- a/crates/hipfire-generate/src/ar.rs +++ b/crates/hipfire-generate/src/ar.rs @@ -2448,37 +2448,6 @@ pub fn generate( match hipfire_loader::generation_early_route(m.arch_id) { Some(hipfire_loader::GenerationEarlyRoute::Gemma4) => { - // The loader publishes one of two mutually-exclusive Gemma4 states: - // eager dense (ModelState::Gemma4) and lowered/MoE - // (ModelState::Gemma4Lowered). Lowered models are served by - // generate_gemma4_lowered below; eager models continue through - // generate_gemma4. - if m.gemma4_lowered_mut().is_some() { - crate::dense::generate_gemma4_lowered( - m, - gpu, - stdout, - id, - prompt, - system_prompt, - temp, - top_p, - top_k, - min_p, - max_tokens, - repeat_penalty, - repeat_window, - presence_penalty, - frequency_penalty, - max_think_tokens, - enable_thinking, - tools, - messages_history, - logprobs_top_k, - request_seed, - ); - return; - } let _ = ( budget_alert_at_tok, budget_alert_text, @@ -2487,16 +2456,8 @@ pub fn generate( pflash_cfg, think_mode, user_explicit_sampling, - top_k, - min_p, cactus_delta, ); - let _ = ( - repeat_penalty, - repeat_window, - presence_penalty, - frequency_penalty, - ); crate::dense::generate_gemma4( m, gpu, @@ -2504,8 +2465,18 @@ pub fn generate( id, prompt, system_prompt, - temp, - top_p, + SamplerConfig { + temperature: temp, + top_p, + repeat_penalty, + repeat_window, + presence_penalty, + frequency_penalty, + blocked_tokens: Vec::new(), + top_k, + min_p, + }, + request_seed, max_tokens, max_think_tokens, enable_thinking, diff --git a/crates/hipfire-generate/src/dense.rs b/crates/hipfire-generate/src/dense.rs index 3a4bc3381f..c0ed0bebc4 100644 --- a/crates/hipfire-generate/src/dense.rs +++ b/crates/hipfire-generate/src/dense.rs @@ -9,6 +9,7 @@ use hipfire_arch_cohere2moe as cohere2moe; use hipfire_arch_deepseek4 as deepseek4; use hipfire_arch_gemma4 as gemma4; +use gemma4::emit::{GemmaEmit, GemmaThoughtRouter}; use hipfire_arch_lfm2moe as lfm2moe; use hipfire_arch_minimax as minimax; use hipfire_arch_muse_glimmer as glimmer; @@ -1979,46 +1980,6 @@ pub fn generate_deepseek4_heterogeneous( "[req {id}] drafter=ar-heterogeneous tau=1.00 tok/s={tok_s:.1} decode ({generated} tok)" ); } -/// Gemma 4 dense text (arch_id=13) eager AR path. -/// -/// Ported from `origin/feat/gemma4-union` and rehomed onto the beta -/// `ModelState::Gemma4(bundle)` design: reads `bundle.config` / -/// `bundle.weights` / `bundle.state` / `bundle.eos_tok` instead of the PR's -/// `m.gemma4_*` Option fields. Same shape as `generate_lfm2moe` (prefill loop, -/// decode loop, JSONL `token` / `done` events) with four gemma4-specifics: -/// -/// 1. Prompt build goes through the model's chat template (Jinja), rendered -/// with an EXPLICIT `bos_token: Some("")` — gemma4's tokenizer -/// decodes `` to id 2 but the raw-encode fallback doesn't prepend -/// it, so the BOS guard below handles the fallback path. -/// 2. Prefill is per-token eager `gemma4::forward::decode_step`; decode is -/// `gemma4::forward::decode_step_with_graph` (hipGraph, default OFF). -/// 3. Sampling is host-side from the returned logits vector via -/// `deepseek4::sampling::sample_token`. -/// 4. Stop set: `` (config.eos_token, id 1) UNION `bundle.eos_tok` -/// (loader-resolved `` / ``) UNION documented 106 — -/// the HF `eos_token_id` list is `[1, 106]`; parsing as a scalar drops -/// 106 and decode loops `` forever. -/// -/// Prompt-cache: arch 13 is intentionally ABSENT from the `cache_capable` -/// allowlist (see load handler) — this path has no LCP prefix-cache block and -/// always cold-prefills the full Jinja render. Enabling cache would corrupt -/// KV slot offsets after turn 1. -/// -/// EAGLE spec-decode (arch-22 drafter via `params.drafter`) is wired when -/// `bundle.eagle.is_some()` and `temp <= 1e-6` — greedy-only accept rule, -/// same contract as DFlash. `HIPFIRE_GEMMA4_EAGLE=0` opts out. -#[inline] -fn gemma4_prefill_batch_for_arch(gpu_arch: &str, requested: Option) -> usize { - requested - .unwrap_or(if matches!(gpu_arch, "gfx1100" | "gfx1201") { - 64 - } else { - 1 - }) - .clamp(1, 64) -} - struct Gemma4LogitTraceConfig { dir: PathBuf, top_k: usize, @@ -2190,318 +2151,36 @@ fn trace_gemma4_logits( } } -#[cfg(test)] -mod gemma4_prefill_batch_tests { - use super::gemma4_prefill_batch_for_arch; - - #[test] - fn validated_arches_default_to_full_wmma_batch_tile_group() { - assert_eq!(gemma4_prefill_batch_for_arch("gfx1100", None), 64); - assert_eq!(gemma4_prefill_batch_for_arch("gfx1201", None), 64); - } - - #[test] - fn unvalidated_arches_keep_the_sequential_default() { - for arch in ["gfx1151", "gfx1200", "gfx942"] { - assert_eq!(gemma4_prefill_batch_for_arch(arch, None), 1); - } - } - - #[test] - fn explicit_override_wins_and_is_clamped() { - assert_eq!(gemma4_prefill_batch_for_arch("gfx1100", Some(8)), 8); - assert_eq!(gemma4_prefill_batch_for_arch("gfx1201", Some(32)), 32); - assert_eq!(gemma4_prefill_batch_for_arch("gfx1100", Some(0)), 1); - assert_eq!(gemma4_prefill_batch_for_arch("gfx1100", Some(128)), 64); - } -} - -#[allow(clippy::too_many_arguments)] -pub fn generate_gemma4_lowered( - m: &mut LoadedModel, - gpu: &mut rdna_compute::Gpu, - stdout: &mut std::io::Stdout, - id: &str, - prompt: &str, - system_prompt: Option<&str>, - temp: f32, - top_p: f32, - top_k: Option, - min_p: Option, - max_tokens: usize, - repeat_penalty: f32, - repeat_window: usize, - presence_penalty: f32, - frequency_penalty: f32, - max_think_tokens: usize, - enable_thinking: bool, - tools: Option<&[serde_json::Value]>, - messages_history: Option<&[hipfire_runtime::prompt_frame::Message]>, - logprobs_top_k: Option, - request_seed: u32, -) { - if m.tokenizer.is_none() { - emit_error_with_id(stdout, id, "tokenizer not loaded"); - return; - } - crate::ar::emit_generation_start( - crate::ar::active_generation_route().unwrap_or(crate::ar::GenerationRoute::Unknown), - stdout, - id, - false, - ); - - let Some(bundle_ref) = m.gemma4_lowered_mut() else { - emit_error_with_id(stdout, id, "gemma4 lowered bundle missing"); - return; - }; - let bos_tok = bundle_ref.config.bos_token; - let cfg_eos_tok = bundle_ref.config.eos_token; - let bundle = bundle_ref as *mut hipfire_loader::Gemma4LoweredBundle; - - let prompt_ids: Vec = { - let tokenizer = m.tokenizer.as_ref().unwrap(); - let try_jinja = hipfire_config::developer_var("HIPFIRE_JINJA_CHAT") - .ok() - .as_deref() - != Some("0") - && m.chat_template.is_some(); - let mut ids = if try_jinja { - let frame = hipfire_runtime::prompt_frame::JinjaChatFrame { - tokenizer, - template: m.chat_template.as_ref().unwrap(), - system: system_prompt, - user: prompt, - enable_thinking, - bos_token: Some(""), - reasoning_strength: None, - reasoning_effort: None, - }; - let rendered = if tools.is_some() || messages_history.is_some() { - let synthesized; - let history = match messages_history { - Some(history) => history, - None => { - let mut messages = Vec::new(); - if let Some(system) = system_prompt { - messages.push(hipfire_runtime::prompt_frame::Message { - role: hipfire_runtime::prompt_frame::Role::System, - content: system.to_owned(), - reasoning_content: None, - name: None, - rendered_name: None, - tool_calls: Vec::new(), - tool_call_id: None, - tool_plan: String::new(), - }); - } - messages.push(hipfire_runtime::prompt_frame::Message { - role: hipfire_runtime::prompt_frame::Role::User, - content: prompt.to_owned(), - reasoning_content: None, - name: None, - rendered_name: None, - tool_calls: Vec::new(), - tool_call_id: None, - tool_plan: String::new(), - }); - synthesized = messages; - &synthesized - } - }; - frame.render_messages(history, tools, None) - } else { - frame.render() - }; - match rendered { - Ok(rendered) => tokenizer.encode(&rendered), - Err(error) => { - eprintln!("[daemon] jinja render failed in Gemma4 lowered path ({error}); using raw prompt"); - tokenizer.encode(prompt) - } - } - } else { - tokenizer.encode(prompt) - }; - if ids.first() != Some(&bos_tok) { - ids.insert(0, bos_tok); - } - ids - }; - - if prompt_ids.is_empty() { - emit_error_with_id(stdout, id, "empty prompt after tokenize"); - return; - } - if prompt_ids.len() + max_tokens > m.max_seq { - emit_error_with_id( - stdout, - id, - format!( - "gemma4 lowered request needs {} KV positions but max_seq is {}", - prompt_ids.len() + max_tokens, - m.max_seq - ), - ); - return; - } - - // The lowered route does not yet publish a prompt-cache contract. Rebuild - // the full Jinja frame from position zero so stale KV can never leak across - // requests; absolute-position writes overwrite every row that is observed. - m.seq_pos = 0; - m.conversation_tokens.clear(); - let t0 = Instant::now(); - for (pos, &token) in prompt_ids.iter().enumerate() { - // Mirror the hand carrier: prompt prefill is eager; only AR decode - // may warm up, capture and replay its single-token body. - gpu.graphs.ar_graph_eligible = false; - let result = unsafe { - gemma4::lowered::forward_scratch( - gpu, - &(*bundle).weights, - &(*bundle).config, - token, - pos, - &mut (*bundle).kv_sliding, - &mut (*bundle).kv_full, - &(*bundle).scratch, - ) - }; - if let Err(error) = result { - emit_error_with_id( - stdout, - id, - format!("gemma4 lowered prefill failed: {error:?}"), - ); - return; - } - } - m.conversation_tokens.extend_from_slice(&prompt_ids); - m.seq_pos = prompt_ids.len(); - let prefill_ms = t0.elapsed().as_millis(); - - let stop_set = unsafe { [cfg_eos_tok, (*bundle).eos_tok, 106] }; - let sampler_cfg = hipfire_runtime::sampler::SamplerConfig { - temperature: temp, - top_p, - repeat_penalty, - repeat_window, - presence_penalty, - frequency_penalty, - blocked_tokens: Vec::new(), - top_k, - min_p, - }; - let mut rng_state = request_seed; - let mut router = GemmaThoughtRouter::new(enable_thinking, max_think_tokens); - let mut generated = 0usize; - let mut ttft_ms = None; - let decode_t0 = Instant::now(); - let mut text_stream = TokenTextStream::new(); - - while generated < max_tokens { - let next = unsafe { - hipfire_runtime::sampler::sample( - gpu, - &(*bundle).scratch.logits, - &(*bundle).scratch.sample_buf, - &(*bundle).scratch.repeat_buf, - (*bundle).config.vocab_size, - &m.conversation_tokens, - &sampler_cfg, - &mut rng_state, - ) - }; - if stop_set.contains(&next) { - break; - } - if ttft_ms.is_none() { - ttft_ms = Some(t0.elapsed().as_secs_f64() * 1000.0); - } - let frag = text_stream.push(m.tokenizer.as_ref().unwrap(), next); - let host_logits = if logprobs_top_k.is_some() { - unsafe { gpu.download_f32(&(*bundle).scratch.logits).ok() } - } else { - None - }; - for event in router.push(&frag).0 { - match event { - GemmaEmit::Reasoning(text) => emit_reasoning_token(stdout, id, &text), - GemmaEmit::Token(text) => { - let mut envelope = serde_json::json!({ - "type": "token", "id": id, "text": text, - "attempt_id": active_attempt_id(), - }); - if let Some(logits) = host_logits.as_ref() { - if let Some((logprob, top)) = crate::common::token_logprob_fields( - logits, - next, - logprobs_top_k, - m.tokenizer.as_ref().unwrap(), - ) { - envelope["logprob"] = serde_json::json!(logprob); - envelope["top_logprobs"] = top; - } - } - let _ = writeln!(stdout, "{envelope}"); - let _ = stdout.flush(); - } - } - } - m.conversation_tokens.push(next); - generated += 1; - let pos = m.seq_pos; - let result = unsafe { - gemma4::lowered::forward_scratch( - gpu, - &(*bundle).weights, - &(*bundle).config, - next, - pos, - &mut (*bundle).kv_sliding, - &mut (*bundle).kv_full, - &(*bundle).scratch, - ) - }; - if let Err(error) = result { - emit_error_with_id( - stdout, - id, - format!("gemma4 lowered decode failed: {error:?}"), - ); - return; - } - m.seq_pos += 1; - } - let tail = text_stream.flush(); - for event in router.push(&tail).0.into_iter().chain(router.flush()) { - match event { - GemmaEmit::Reasoning(text) => emit_reasoning_token(stdout, id, &text), - GemmaEmit::Token(text) => emit_visible_token(stdout, id, &text), - } - } - let decode_ms = decode_t0.elapsed().as_millis().max(1); - let total_ms = t0.elapsed().as_millis().max(1); - let decode_tok_s = generated as f64 * 1000.0 / decode_ms as f64; - let prefill_tok_s = prompt_ids.len() as f64 * 1000.0 / prefill_ms.max(1) as f64; - let _ = writeln!( - stdout, - r#"{{"type":"done","id":"{}","tokens":{},"tok_s":{:.2},"prefill_tokens":{},"prefill_ms":{},"prefill_tok_s":{:.2},"decode_tok_s":{:.2},"ttft_ms":{:.3},"total_ms":{},"attempt_id":{}}}"#, - id, - generated, - decode_tok_s, - prompt_ids.len(), - prefill_ms, - prefill_tok_s, - decode_tok_s, - ttft_ms.unwrap_or(total_ms as f64), - total_ms, - active_attempt_id(), - ); - let _ = stdout.flush(); -} - +/// Gemma 4 (arch_id=13) AR path for every variant. +/// +/// Ported from `origin/feat/gemma4-union` and rehomed onto the beta +/// `ModelState::Gemma4(bundle)` design: reads `bundle.config` / +/// `bundle.weights` / `bundle.state` / `bundle.eos_tok` instead of the PR's +/// `m.gemma4_*` Option fields. Same shape as `generate_lfm2moe` (prefill loop, +/// decode loop, JSONL `token` / `done` events) with four gemma4-specifics: +/// +/// 1. Prompt build goes through the model's chat template (Jinja), rendered +/// with an EXPLICIT `bos_token: Some("")` — gemma4's tokenizer +/// decodes `` to id 2 but the raw-encode fallback doesn't prepend +/// it, so the BOS guard below handles the fallback path. +/// 2. Prefill is batched where the arch and formats admit it, else +/// per-token `gemma4::forward::decode_step`; decode is +/// `gemma4::forward::decode_step_with_graph`. +/// 3. Sampling is host-side from the returned logits vector via +/// `sampler::sample_cpu` with the request's full sampler contract. +/// 4. Stop set: `` (config.eos_token, id 1) UNION `bundle.eos_tok` +/// (loader-resolved `` / ``) UNION documented 106 — +/// the HF `eos_token_id` list is `[1, 106]`; parsing as a scalar drops +/// 106 and decode loops `` forever. +/// +/// Prompt cache: `Gemma4Bundle::prefill` plans the canonical render against +/// the engine session cache (this conversation's live state or a cold-exact +/// snapshot) and the turn commits after `done`; the daemon reports arch 13 +/// `cache_capable` when a cache is attached. +/// +/// EAGLE spec-decode (arch-22 drafter via `params.drafter`, opt-in with +/// `HIPFIRE_GEMMA4_EAGLE=1`) runs greedy requests through the generic +/// `generate_spec` core instead ([`generate_gemma4_spec`]). #[allow(clippy::too_many_arguments)] pub fn generate_gemma4( m: &mut LoadedModel, @@ -2510,8 +2189,8 @@ pub fn generate_gemma4( id: &str, prompt: &str, system_prompt: Option<&str>, - temp: f32, - top_p: f32, + sampler_cfg: hipfire_runtime::sampler::SamplerConfig, + request_seed: u32, max_tokens: usize, max_think_tokens: usize, enable_thinking: bool, @@ -2522,6 +2201,7 @@ pub fn generate_gemma4( logprobs_top_k: Option, ) { // Gemma4 thinking is explicit boolean (default off); max_think_tokens is orthogonal cap. + let temp = sampler_cfg.temperature; if m.tokenizer.is_none() { emit_error_with_id(stdout, id, "tokenizer not loaded"); @@ -2545,7 +2225,7 @@ pub fn generate_gemma4( emit_error_with_id( stdout, id, - "gemma4 bundle missing on arch_id=13 generate (eager dense only; EAGLE/lowered not yet wired)", + "gemma4 bundle missing on arch_id=13 generate", ); return; }; @@ -2645,18 +2325,8 @@ pub fn generate_gemma4( s }; - // Capacity guard. No eviction on arch_id=13 — reset the KV cursors when - // the requested run would overflow the physical cache. - let overflow = { bundle.state.n_tokens + prompt_ids.len() + max_tokens > bundle.state.max_seq }; - if overflow { - let (n, cap) = (bundle.state.n_tokens, bundle.state.max_seq); - eprintln!("[daemon] arch_id=13 context full ({n}/{cap}) — resetting Gemma4State"); - bundle.state.reset(); - m.seq_pos = 0; - m.conversation_tokens.clear(); - } - // Hard refusal: even from a cold cache the prompt alone must fit (KV - // writes at pos >= max_seq would be out of bounds). + // Hard refusal: the prompt alone must fit (KV writes at pos >= max_seq + // would be out of bounds). let cache_cap = bundle.state.max_seq; if prompt_ids.len() >= cache_cap { emit_error_with_id( @@ -2670,291 +2340,65 @@ pub fn generate_gemma4( return; } - let t0 = Instant::now(); - - // ── Prefill. `forward_batch` preserves the eager result/KV contract while - // reading projection weights once for a small token block. Keep the path - // opt-in until long-context parity is certified on each supported board. - let mut last_logits: Vec = Vec::new(); - { - let requested_prefill_batch = hipfire_config::developer_var("HIPFIRE_GEMMA4_PREFILL_BATCH") - .ok() - .and_then(|value| value.parse::().ok()); - let resolved_prefill_batch = - gemma4_prefill_batch_for_arch(&gpu.arch, requested_prefill_batch); - let prefill_batch = if resolved_prefill_batch > 1 - && !gemma4::forward::supports_batched_prefill(&bundle.weights) - { - eprintln!( - "[daemon] Gemma4 batched prefill is unavailable for this weight format; using eager prefill" - ); - 1 - } else { - resolved_prefill_batch - }; - let mut offset = 0usize; - while offset < prompt_ids.len() { - let width = prefill_batch.min(prompt_ids.len() - offset); - let start_pos = bundle.state.n_tokens; - let result = if width == 1 { - gemma4::forward::decode_step( - &bundle.config, - &bundle.weights, - &mut bundle.state, - gpu, - prompt_ids[offset], - start_pos as u32, - ) - } else { - gemma4::forward::forward_batch( - &bundle.config, - &bundle.weights, - &mut bundle.state, - gpu, - &prompt_ids[offset..offset + width], - start_pos, - ) - }; - match result { - Ok(logits) => last_logits = logits, - Err(e) => { - emit_error_with_id(stdout, id, format!("gemma4 prefill failed: {e:?}")); - return; - } - } - offset += width; - } - } - for &tok in &prompt_ids { - m.conversation_tokens.push(tok); - } - let prefill_ms = t0.elapsed().as_millis(); - let mut gemma_router = GemmaThoughtRouter::new(enable_thinking, max_think_tokens); - // Multi-byte characters span tokens: decode incrementally, never per token. - let mut text_stream = TokenTextStream::new(); - - // ── EAGLE spec-decode fast path (arch-22 drafter loaded; greedy only) ── - // - // Mirrors the DFlash dispatch contract: the accept rule is greedy-argmax, - // so the committed stream is PROVABLY the target's greedy AR sequence — - // the same tokens the AR loop below would emit at temp 0. That invariant - // is the gate `infer_gemma4_spec --check-eager` validates byte-for-byte - // (spec == eager on hiptrx/gfx1201, Q8 and MQ4-attn targets, dl ≤ 4). - // temp > 0 falls through to the AR sampling loop, exactly like DFlash. - // GATED OFF: EAGLE currently violates its own correctness contract. - // - // The accept rule is greedy-argmax, so at temp 0 the committed stream is - // supposed to be PROVABLY the target's greedy AR sequence — byte-identical - // to the AR loop below. Measured on gfx1201, 12B-it, temp 0, 48 tokens, - // same prompt, on TWO different quantizations: - // - // uniform MQ4 AR 52.81 tok/s "To understand quantum computing, you - // first have to understand how a normal - // computer works..." - // EAGLE 44.28 tok/s "To explain it, we have to look at the - // \"rules\" of the world we see..." - // tau 1.531, rounds 32 — NOT byte-identical, and 0.84x - // - // Q8 diverges the same way at tau 1.912, 0.73x - // - // Divergence begins at the first committed token, which points at the seed - // hidden / first-round verify rather than at acceptance rate. Reproducing - // across two quants rules out a quant-specific kernel issue. - // - // Until parity is proven it must not run: wrong tokens that look fluent are - // worse than a refusal, and it is slower anyway. Opt in for debugging with - // HIPFIRE_GEMMA4_EAGLE=1; HIPFIRE_GEMMA4_EAGLE=0 remains an explicit off. - // - // Separately: a K-map-promoted target cannot run this path at all — - // `proj_gemm_batched` handles only Q8_0 and MQ4G256/HFQ4G256, so the - // mode-3 `Promote6` on v_proj yields MQ6G256 and the verify fails loud with - // "dtype MQ6G256 has no batched proj kernel". A uniform (--no-kmap) target - // avoids that, which is how the numbers above were taken. - let eagle_active = bundle.eagle.is_some() - && temp <= 1e-6 + // Greedy, penalty-free requests that fit one verify block past their + // budget go to the EAGLE speculator. It commits the target's batched + // verify argmax, which can differ from AR decode on near-tied logits. + let spec = temp <= 1e-6 + && sampler_cfg.repeat_penalty == 1.0 + && sampler_cfg.presence_penalty == 0.0 + && sampler_cfg.frequency_penalty == 0.0 + && logprobs_top_k.is_none() && hipfire_config::developer_var("HIPFIRE_GEMMA4_EAGLE") .ok() .as_deref() - == Some("1"); - if eagle_active { - let draft_len = bundle.eagle.as_ref().unwrap().draft_len; - // Seed hidden = post-`model.norm` hidden of the last prompt position - // (left in `state.tmp` by the final prefill `decode_step` — the - // lm_head input). The seed TOKEN is the last prompt token; its KV is - // re-written (identically) at the top of the first verify. The first - // generated token comes out of round 1's verify (argmax_per_pos[0]), - // exactly as the AR loop's first sample from `last_logits` would — - // so `last_logits` is intentionally unused on this path. - { - let eagle = bundle.eagle.as_ref().unwrap(); - if let Err(e) = eagle - .spec_scratch - .set_seed_hidden_from(gpu, &bundle.state.tmp) - { - emit_error_with_id(stdout, id, format!("gemma4 eagle seed hidden: {e}")); - return; - } - } - let prefill_end = bundle.state.n_tokens; - let mut seed_token = *prompt_ids.last().unwrap(); - let mut generated_count = 0usize; - let mut rounds = 0usize; - let mut total_accepted = 0usize; - let mut stop = false; - let mut ttft_ms: Option = None; - let decode_t0 = Instant::now(); - while !stop && generated_count < max_tokens { - let committed_len = bundle.state.n_tokens; - // KV/seq bound: the verify block occupies [L-1, L-1+draft_len+1). - if committed_len + draft_len + 1 >= cache_cap { - break; - } - let spec = { - // Split borrows: config/weights immutably, state+eagle mutably. - // Use raw pointers to avoid borrow-checker overlap on `bundle`. - let bundle_ptr = bundle as *mut hipfire_loader::Gemma4Bundle; - unsafe { - let cfg = &(*bundle_ptr).config; - let weights = &(*bundle_ptr).weights; - let state = &mut (*bundle_ptr).state; - let eagle = (*bundle_ptr).eagle.as_mut().unwrap(); - gemma4::speculative::spec_step_gemma4_eagle( - gpu, - weights, - cfg, - state, - &eagle.drafter_weights, - &eagle.drafter_config, - &mut eagle.drafter_scratch, - &mut eagle.spec_scratch, - seed_token, - committed_len, - draft_len, - 0.0, - ) - } - }; - let spec = match spec { - Ok(s) => s, - Err(e) => { - emit_error_with_id(stdout, id, format!("gemma4 eagle spec step failed: {e}")); - return; - } - }; - rounds += 1; - total_accepted += spec.accept_len; - // Emit the committed tokens (accepted drafts ++ bonus); stop at - // EOS / max_tokens — identical to what the AR loop would commit. - for &t in &spec.committed { - if stop_set.contains(&t) { - stop = true; - break; - } - if ttft_ms.is_none() { - ttft_ms = Some(t0.elapsed().as_secs_f64() * 1000.0); - } - let frag = text_stream.push(m.tokenizer.as_ref().unwrap(), t); - let (emits, _) = gemma_router.push(&frag); - for ev in emits { - match ev { - GemmaEmit::Reasoning(text) => emit_reasoning_token(stdout, id, &text), - GemmaEmit::Token(text) => emit_visible_token(stdout, id, &text), - } - } - m.conversation_tokens.push(t); - generated_count += 1; - if generated_count >= max_tokens { - stop = true; - break; - } - } - // Next round's seed = this round's bonus; its hidden is already - // staged in spec_scratch.seed_hidden by spec_step. - seed_token = spec.next_seed_token; - // If we stopped on EOS mid-block, spec.n_tokens already counts the - // committed-but-not-emitted tail — the cursor settle below re-anchors - // to emitted extent. - if stop { - break; - } - } - for ev in gemma_router.push(&text_stream.flush()).0 { - match ev { - GemmaEmit::Reasoning(text) => emit_reasoning_token(stdout, id, &text), - GemmaEmit::Token(text) => emit_visible_token(stdout, id, &text), - } - } - - // ── Cursor settle. `spec_step` leaves n_tokens = L+accept_len+1 with - // the final bonus's KV slot unwritten, and an EOS mid-block leaves - // committed-but-not-emitted tokens counted. Re-anchor the cursor to - // the EMITTED extent and re-forward the last emitted token at its - // slot so every position < n_tokens carries valid KV — the same - // invariant the AR loop leaves behind. (KV writes are absolute- - // position keyed and deterministic, so the re-write is identical to - // a fresh forward — the same property the verify itself relies on - // when it re-writes the seed's KV each round.) ── - if generated_count > 0 { - let last_tok = *m.conversation_tokens.last().unwrap(); - let last_pos = (prefill_end + generated_count - 1) as u32; - let bundle_ptr = bundle as *mut hipfire_loader::Gemma4Bundle; - let settle_res = unsafe { - let cfg = &(*bundle_ptr).config; - let weights = &(*bundle_ptr).weights; - let state = &mut (*bundle_ptr).state; - gemma4::forward::decode_step(cfg, weights, state, gpu, last_tok, last_pos) - }; - if let Err(e) = settle_res { - eprintln!("[daemon] gemma4 eagle cursor settle failed: {e:?}"); - bundle.state.n_tokens = prefill_end + generated_count; - } - } else { - bundle.state.n_tokens = prefill_end; - } - m.seq_pos = bundle.state.n_tokens; - - let decode_ms = decode_t0.elapsed().as_millis().max(1); - let total_ms = t0.elapsed().as_millis().max(1); - let tok_s = if generated_count > 0 { - (generated_count as f64 * 1000.0) / decode_ms as f64 - } else { - 0.0 - }; - let prefill_tok_s = prompt_ids.len() as f64 * 1000.0 / prefill_ms.max(1) as f64; - let ttft_ms = ttft_ms.unwrap_or(total_ms as f64); - // τ = mean tokens committed per round (accepted drafts + 1 bonus). - let tau = if rounds > 0 { - (total_accepted + rounds) as f64 / rounds as f64 - } else { - 0.0 - }; - let _ = writeln!( + == Some("1") + && m.speculator.as_ref().is_some_and(|s| { + spec_ctx_request_fits(prompt_ids.len(), max_tokens, s.block_size(), cache_cap) + }); + let route = if spec { + hipfire_runtime::session_cache::SessionRoute::Mtp + } else { + hipfire_runtime::session_cache::SessionRoute::Ar + }; + // Plan against the canonical render: this conversation's live state or + // a cached snapshot. Prefill restores it and computes only the rest. + let reused = hipfire_runtime::arch_model::ArchModel::session_plan(&*bundle, &prompt_ids, route); + if spec { + generate_gemma4_spec( + m, + gpu, stdout, - r#"{{"type":"done","id":"{}","tokens":{},"tok_s":{:.2},"prefill_tokens":{},"prefill_ms":{},"prefill_tok_s":{:.2},"decode_tok_s":{:.2},"ttft_ms":{:.3},"total_ms":{},"spec":"gemma4_eagle","rounds":{},"tau":{:.3},"draft_len":{},"attempt_id":{}}}"#, id, - generated_count, - tok_s, - prompt_ids.len(), - prefill_ms, - prefill_tok_s, - tok_s, - ttft_ms, - total_ms, - rounds, - tau, - draft_len, - active_attempt_id(), + prompt_ids, + reused, + max_tokens, + max_think_tokens, + enable_thinking, + cfg_eos_tok, ); - let _ = stdout.flush(); return; } + m.seq_pos = 0; + m.conversation_tokens.clear(); + let t0 = Instant::now(); + let mut last_logits = match bundle.prefill(gpu, &prompt_ids, reused, route) { + Ok(logits) => logits, + Err(e) => { + emit_error_with_id(stdout, id, format!("gemma4 prefill failed: {e}")); + return; + } + }; + m.conversation_tokens.extend_from_slice(&prompt_ids); + let computed_prompt = prompt_ids.len() - reused; + let prefill_ms = t0.elapsed().as_millis(); + let mut gemma_router = GemmaThoughtRouter::new(enable_thinking, max_think_tokens); + // Multi-byte characters span tokens: decode incrementally, never per token. + let mut text_stream = TokenTextStream::new(); - // ── Decode loop. Sample host-side from the running logits vector. ── - let seed = std::time::SystemTime::now() - .duration_since(std::time::UNIX_EPOCH) - .map(|d| d.as_nanos() as u64) - .unwrap_or(0x9E3779B97F4A7C15); - let mut rng = deepseek4::sampling::Xorshift::new(seed); + // ── Decode loop. Sample host-side from the running logits vector with the + // request's full sampler contract (penalties over the conversation). ── + hipfire_runtime::llama::reset_cpu_sampler_rng(request_seed); + let mut sample_logits: Vec = Vec::with_capacity(last_logits.len()); let mut generated_count: usize = 0; let mut ttft_ms: Option = None; @@ -2963,7 +2407,13 @@ pub fn generate_gemma4( if generated_count >= max_tokens { break; } - let next_tok = deepseek4::sampling::sample_token(&last_logits, temp, 0, top_p, &mut rng); + sample_logits.clear(); + sample_logits.extend_from_slice(&last_logits); + let next_tok = hipfire_runtime::sampler::sample_cpu( + &mut sample_logits, + &m.conversation_tokens, + &sampler_cfg, + ); trace_gemma4_logits(gpu, id, generated_count, next_tok, &last_logits); if stop_set.contains(&next_tok) { break; @@ -3052,15 +2502,16 @@ pub fn generate_gemma4( } else { 0.0 }; - let prefill_tok_s = prompt_ids.len() as f64 * 1000.0 / prefill_ms.max(1) as f64; + let prefill_tok_s = computed_prompt as f64 * 1000.0 / prefill_ms.max(1) as f64; let ttft_ms = ttft_ms.unwrap_or(total_ms as f64); let _ = writeln!( stdout, - r#"{{"type":"done","id":"{}","tokens":{},"tok_s":{:.2},"prefill_tokens":{},"prefill_ms":{},"prefill_tok_s":{:.2},"decode_tok_s":{:.2},"ttft_ms":{:.3},"total_ms":{},"attempt_id":{}}}"#, + r#"{{"type":"done","id":"{}","tokens":{},"tok_s":{:.2},"prefill_tokens":{},"cached_tokens":{},"prefill_ms":{},"prefill_tok_s":{:.2},"decode_tok_s":{:.2},"ttft_ms":{:.3},"total_ms":{},"attempt_id":{}}}"#, id, generated_count, tok_s, - prompt_ids.len(), + computed_prompt, + reused, prefill_ms, prefill_tok_s, tok_s, @@ -3069,7 +2520,93 @@ pub fn generate_gemma4( active_attempt_id(), ); let _ = stdout.flush(); + hipfire_runtime::arch_model::ArchModel::session_commit_live(bundle, &m.conversation_tokens); +} + +/// Gemma 4 EAGLE spec decode on the generic `generate_spec` core: the +/// `MtpSpeculator` in `m.speculator`, the bundle as +/// `SpecTarget` and `GemmaSpecEmit`. `reused` prompt tokens come from the +/// session plan; greedy only. Writes the Gemma `done` envelope and commits +/// the session's live state. +#[allow(clippy::too_many_arguments)] +fn generate_gemma4_spec( + m: &mut LoadedModel, + gpu: &mut rdna_compute::Gpu, + stdout: &mut std::io::Stdout, + id: &str, + prompt_ids: Vec, + reused: usize, + max_tokens: usize, + max_think_tokens: usize, + enable_thinking: bool, + cfg_eos_tok: u32, +) { + // Host mirrors match the reused prefix; the bundle owns device state. + m.seq_pos = reused; + m.conversation_tokens.clear(); + m.conversation_tokens + .extend_from_slice(&prompt_ids[..reused]); + let draft_len = m.speculator.as_ref().map_or(0, |s| s.block_size()); + let suffix = prompt_ids[reused..].to_vec(); + let Some(run) = crate::qwen::generate_spec( + m, + gpu, + stdout, + id, + prompt_ids, + suffix, + reused, + reused > 0, + None, + max_tokens, + SpecEmitRequest { + im_end: Some(cfg_eos_tok), + tools: None, + enable_grammar: false, + stop: Vec::new(), + max_think: max_think_tokens, + assistant_prefix: AssistantPrefix::Plain, + think_mode: if enable_thinking { + ThinkMode::High + } else { + ThinkMode::NonThink + }, + decoded_vocab: None, + }, + 0.0, + ) else { + // generate_spec already wrote the terminal event. + return; + }; + let per_s = |n: usize, s: f64| if s > 0.0 { n as f64 / s } else { 0.0 }; + let tok_s = per_s(run.generated, run.decode_s); + // τ = mean tokens committed per window (accepted drafts + 1 bonus). + let tau = per_s(run.spec_accepted + run.spec_cycles, run.spec_cycles as f64); + let prefill_ms = run.prefill_s * 1000.0; + let _ = writeln!( + stdout, + r#"{{"type":"done","id":"{}","tokens":{},"tok_s":{:.2},"prefill_tokens":{},"cached_tokens":{},"prefill_ms":{},"prefill_tok_s":{:.2},"decode_tok_s":{:.2},"ttft_ms":{:.3},"total_ms":{},"spec":"gemma4_eagle","rounds":{},"tau":{:.3},"draft_len":{},"attempt_id":{}}}"#, + id, + run.generated, + tok_s, + run.prefill_tokens_len, + reused, + prefill_ms as u128, + per_s(run.prefill_tokens_len, run.prefill_s), + tok_s, + prefill_ms, + (run.total_s * 1000.0) as u128, + run.spec_cycles, + tau, + draft_len, + active_attempt_id(), + ); + let _ = stdout.flush(); + if let Some(state) = m.state.as_deref_mut() { + state.session_commit_live(&m.conversation_tokens); + } } + /// Muse Glimmer dense text (arch_id=14) eager AR path. /// /// Mirrors `generate_gemma4` shape (prefill loop / decode loop, JSONL @@ -3351,560 +2888,6 @@ impl GlimmerHarmonyRouter { } } -/// Gemma thought-channel router (minimal, follows Glimmer patterns and checked-in Gemma template). -/// -/// Gemma4 template uses `<|channel>thought\n` ... `\n` to wrap reasoning. -/// When `enable_thinking` is false the Jinja prompt already emits a closed empty -/// `thought` channel (`<|channel>thought\n`), so the router starts in -/// Answer mode. When true it starts awaiting the thought opening. Content inside -/// the thought channel is emitted as semantic `reasoning` events; -/// content outside is visible answer tokens. `max_think_tokens` is an orthogonal -/// force-close cap (0 = uncapped) that, when reached, flushes pending reasoning -/// and transitions to Answer, mirroring Glimmer's `max_think_tokens` handling. -#[derive(Debug, Clone, Copy, PartialEq, Eq)] -pub enum GemmaChannel { - AwaitingThought, - Reasoning, - Answer, -} - -#[derive(Debug, Clone, PartialEq, Eq)] -pub enum GemmaEmit { - Reasoning(String), - Token(String), -} - -pub struct GemmaThoughtRouter { - pub state: GemmaChannel, - pub pending: String, - pub reasoning_tokens: usize, - pub max_think_tokens: usize, - pub enable_thinking: bool, - pub just_forced: bool, -} - -impl GemmaThoughtRouter { - pub fn new(enable_thinking: bool, max_think_tokens: usize) -> Self { - let state = if enable_thinking { - GemmaChannel::AwaitingThought - } else { - GemmaChannel::Answer - }; - Self { - state, - pending: String::new(), - reasoning_tokens: 0, - max_think_tokens, - enable_thinking, - just_forced: false, - } - } - - pub fn push(&mut self, frag: &str) -> (Vec, bool) { - if frag.is_empty() { - return (Vec::new(), false); - } - self.pending.push_str(frag); - let mut out = Vec::new(); - let entered_as_reasoning = self.state == GemmaChannel::Reasoning; - loop { - match self.state { - GemmaChannel::AwaitingThought => { - const CHANNEL_OPEN: &str = "<|channel>"; - const THOUGHT_OPEN: &str = "<|channel>thought"; - - // A thought channel is optional and may be split across - // decoded fragments. Hold only while the bytes seen so far - // can still become the canonical opening header. - if THOUGHT_OPEN.starts_with(&self.pending) { - break; - } - if self.pending.starts_with(THOUGHT_OPEN) { - let mut header_end = THOUGHT_OPEN.len(); - if self.pending[header_end..].starts_with('\n') { - header_end += 1; - } - self.pending.drain(..header_end); - self.state = GemmaChannel::Reasoning; - continue; - } - - // The response schema makes the thought header optional. - // Once the buffered bytes cannot form that header, route - // them as answer content. Some Gemma4 checkpoints emit an - // orphan `<|channel>` before an otherwise valid answer; - // consume that control token without dropping its payload. - if self.pending.starts_with(CHANNEL_OPEN) { - self.pending.drain(..CHANNEL_OPEN.len()); - if self.pending.starts_with('\n') { - self.pending.drain(..1); - } - } - self.state = GemmaChannel::Answer; - continue; - } - GemmaChannel::Reasoning => { - if let Some(pos) = self.pending.find("") { - if pos > 0 { - let text = self.pending[..pos].to_string(); - self.pending.drain(..pos); - if !text.is_empty() { - out.push(GemmaEmit::Reasoning(text)); - } - continue; - } - let end = "".len(); - self.pending.drain(..end); - if self.pending.starts_with('\n') { - self.pending.drain(..1); - } - self.state = GemmaChannel::Answer; - continue; - } else { - let hold = gemma_longest_marker_suffix(&self.pending); - let emit_len = self.pending.len().saturating_sub(hold); - if emit_len > 0 { - let text = self.pending[..emit_len].to_string(); - self.pending.drain(..emit_len); - if !text.is_empty() { - out.push(GemmaEmit::Reasoning(text)); - } - } - break; - } - } - GemmaChannel::Answer => { - // Answer must chunk-safely strip any Gemma channel - // control markers, including canonical <|channel|> - // forms, orphan <|channel> variants, and markers - // arriving after a forced max-think transition. - // Preserve payload before/after each marker and hold - // a suffix that could still become a marker. - const ANSWER_MARKERS: &[&str] = &[ - "<|channel>thought", - "<|channel>", - "<|channel|>", - "", - "<|turn>", - "", - ]; - loop { - let hold = gemma_longest_marker_suffix(&self.pending); - let search_len = self.pending.len().saturating_sub(hold); - let searchable = &self.pending[..search_len]; - let mut best_pos: Option = None; - let mut best_len = 0usize; - for &m in ANSWER_MARKERS { - if let Some(pos) = searchable.find(m) { - if best_pos.is_none() - || pos < best_pos.unwrap() - || (pos == best_pos.unwrap() && m.len() > best_len) - { - best_pos = Some(pos); - best_len = m.len(); - } - } - } - if let Some(pos) = best_pos { - if pos > 0 { - let text = self.pending[..pos].to_string(); - self.pending.drain(..pos); - if !text.is_empty() { - out.push(GemmaEmit::Token(text)); - } - continue; - } - self.pending.drain(..best_len); - if self.pending.starts_with('\n') { - self.pending.drain(..1); - } - if self.pending.is_empty() { - break; - } - continue; - } - if search_len > 0 { - let text = self.pending[..search_len].to_string(); - self.pending.drain(..search_len); - if !text.is_empty() { - out.push(GemmaEmit::Token(text)); - } - } - break; - } - break; - } - } - } - let emitted_reasoning = out.iter().any(|e| matches!(e, GemmaEmit::Reasoning(_))); - let should_count = entered_as_reasoning || emitted_reasoning; - if should_count - && self.max_think_tokens != 0 - && self.reasoning_tokens < self.max_think_tokens - { - self.reasoning_tokens += 1; - if self.reasoning_tokens >= self.max_think_tokens - && self.state == GemmaChannel::Reasoning - { - if !self.pending.is_empty() { - let tail = std::mem::take(&mut self.pending); - if !gemma_is_marker_prefix(&tail) && !tail.is_empty() { - out.push(GemmaEmit::Reasoning(tail)); - } - } - self.state = GemmaChannel::Answer; - self.just_forced = true; - } - } - (out, false) - } - - pub fn flush(&mut self) -> Vec { - if self.pending.is_empty() { - return Vec::new(); - } - if self.state == GemmaChannel::AwaitingThought && gemma_is_marker_prefix(&self.pending) { - self.pending.clear(); - return Vec::new(); - } - let text = std::mem::take(&mut self.pending); - if gemma_is_marker_prefix(&text) { - return Vec::new(); - } - vec![if self.state == GemmaChannel::Reasoning { - GemmaEmit::Reasoning(text) - } else { - GemmaEmit::Token(text) - }] - } -} - -pub fn gemma_is_marker_prefix(s: &str) -> bool { - const MARKERS: &[&str] = &[ - "<|channel>thought", - "<|channel>", - "<|channel|>", - "", - "<|turn>", - "", - ]; - MARKERS.iter().any(|m| m.starts_with(s)) -} - -#[cfg(test)] -mod gemma_thought_router_tests { - use super::{gemma_is_marker_prefix, GemmaChannel, GemmaEmit, GemmaThoughtRouter}; - - fn route(enable_thinking: bool, chunks: &[&str]) -> (String, String, GemmaChannel) { - let mut router = GemmaThoughtRouter::new(enable_thinking, 0); - let mut visible = String::new(); - let mut reasoning = String::new(); - for chunk in chunks { - for event in router.push(chunk).0 { - match event { - GemmaEmit::Reasoning(text) => reasoning.push_str(&text), - GemmaEmit::Token(text) => visible.push_str(&text), - } - } - } - for event in router.flush() { - match event { - GemmaEmit::Reasoning(text) => reasoning.push_str(&text), - GemmaEmit::Token(text) => visible.push_str(&text), - } - } - (visible, reasoning, router.state) - } - - #[test] - fn gemma_router_routes_canonical_thought_then_answer() { - let (visible, reasoning, state) = route( - true, - &["<|channel>", "thought", "\nplan\nanswer"], - ); - assert_eq!(reasoning, "plan"); - assert_eq!(visible, "answer"); - assert_eq!(state, GemmaChannel::Answer); - } - - #[test] - fn gemma_router_recovers_orphan_channel_before_answer() { - let (visible, reasoning, state) = - route(true, &["<|channel>", "\n", "```python\nprint('ok')\n```"]); - assert_eq!(visible, "```python\nprint('ok')\n```"); - assert!(reasoning.is_empty()); - assert_eq!(state, GemmaChannel::Answer); - } - - #[test] - fn gemma_router_orphan_channel_is_chunk_boundary_invariant() { - let full = "<|channel>\nanswer"; - let expected = route(true, &[full]); - for split in 1..full.len() { - if full.is_char_boundary(split) { - assert_eq!(route(true, &[&full[..split], &full[split..]]), expected); - } - } - } - - #[test] - fn gemma_router_thinking_request_can_emit_direct_answer() { - let (visible, reasoning, state) = route(true, &["direct answer"]); - assert_eq!(visible, "direct answer"); - assert!(reasoning.is_empty()); - assert_eq!(state, GemmaChannel::Answer); - } - - #[test] - fn gemma_router_drops_only_an_unfinished_control_marker_at_eos() { - let (visible, reasoning, state) = route(true, &["<|chan"]); - assert!(visible.is_empty()); - assert!(reasoning.is_empty()); - assert_eq!(state, GemmaChannel::AwaitingThought); - } - - #[test] - fn gemma_marker_prefix_does_not_classify_marker_plus_payload() { - assert!(gemma_is_marker_prefix("<|chan")); - assert!(gemma_is_marker_prefix("<|channel>")); - assert!(!gemma_is_marker_prefix("<|channel>\nanswer")); - } - - fn route_with_cap( - enable_thinking: bool, - max_think_tokens: usize, - chunks: &[&str], - ) -> (String, String, GemmaChannel) { - let mut router = GemmaThoughtRouter::new(enable_thinking, max_think_tokens); - let mut visible = String::new(); - let mut reasoning = String::new(); - for chunk in chunks { - for event in router.push(chunk).0 { - match event { - GemmaEmit::Reasoning(text) => reasoning.push_str(&text), - GemmaEmit::Token(text) => visible.push_str(&text), - } - } - } - for event in router.flush() { - match event { - GemmaEmit::Reasoning(text) => reasoning.push_str(&text), - GemmaEmit::Token(text) => visible.push_str(&text), - } - } - (visible, reasoning, router.state) - } - - fn assert_no_markers(s: &str) { - for m in &[ - "<|channel>thought", - "<|channel>", - "<|channel|>", - "", - "<|turn>", - "", - ] { - assert!(!s.contains(m), "visible leaked marker {:?} in {:?}", m, s); - } - } - - #[test] - fn gemma_router_thinking_off_strips_all_channel_markers_every_split() { - // Thinking-off starts in Answer; every channel/turn marker must be - // stripped chunk-safely regardless of split. - let full = "pre<|channel>midpost<|channel|>inner<|channel>thought\nXtail<|turn>endafter"; - let expected = route_with_cap(false, 0, &[full]); - assert_no_markers(&expected.0); - // Adjacent payload must survive: markers stripped, text joined. - assert_eq!(expected.0, "premidpostinnerXtailendafter"); - for split in 1..full.len() { - if !full.is_char_boundary(split) { - continue; - } - let got = route_with_cap(false, 0, &[&full[..split], &full[split..]]); - assert_eq!(got, expected, "mismatch at split {}", split); - assert_no_markers(&got.0); - } - // Also verify orphan <|channel> with newline framing is stripped. - let full2 = "<|channel>\nanswer"; - let exp2 = route_with_cap(false, 0, &[full2]); - assert_eq!(exp2.0, "answer"); - for split in 1..full2.len() { - if !full2.is_char_boundary(split) { - continue; - } - assert_eq!( - route_with_cap(false, 0, &[&full2[..split], &full2[split..]]), - exp2 - ); - } - // Canonical <|channel|> in thinking-off must also be stripped. - let full3 = "A<|channel|>B"; - let exp3 = route_with_cap(false, 0, &[full3]); - assert_eq!(exp3.0, "AB"); - assert_no_markers(&exp3.0); - for split in 1..full3.len() { - if !full3.is_char_boundary(split) { - continue; - } - assert_eq!( - route_with_cap(false, 0, &[&full3[..split], &full3[split..]]), - exp3 - ); - } - } - - #[test] - fn gemma_router_forced_close_strips_markers_after_transition_every_split() { - // Force max-think after one reasoning push, then ensure every - // subsequent channel/turn marker in Answer is stripped at every - // split, preserving adjacent payload. - let pre = "<|channel>thought\nAAA"; - let post = "BBB<|channel>CCCDDD<|channel|>EEE<|turn>FFFGGG"; - let full = format!("{}{}", pre, post); - // full = "<|channel>thought\nAAABBB<|channel>CCCDDD<|channel|>EEE<|turn>FFFGGG" - // With max_think=1 the router forces to Answer after the first - // reasoning push; the trailing that closes thought and - // all markers inside post must be stripped, not leaked. - let expected = route_with_cap(true, 1, &[&full]); - assert!(expected.1.contains("AAA") || expected.0.contains("AAA")); - assert_no_markers(&expected.0); - // Answer payload should be the post text with markers removed. - // Post without markers: "BBBCCCDDDEEEFFFGGG" - assert_eq!(expected.0, "BBBCCCDDDEEEFFFGGG"); - // For split invariance after forced close, keep the reasoning header - // as one chunk and only split the post payload. Splitting the header - // itself changes per-push reasoning counting and is not required to - // be invariant for this test. - for split in 0..=post.len() { - if split != 0 && !post.is_char_boundary(split) { - continue; - } - let got = if split == 0 || split == post.len() { - route_with_cap(true, 1, &[&full]) - } else { - let c1 = &post[..split]; - let c2 = &post[split..]; - route_with_cap(true, 1, &[pre, c1, c2]) - }; - assert_eq!(got.0, expected.0, "forced mismatch at post split {}", split); - assert_no_markers(&got.0); - } - // Also test forced transition where pending marker is split across - // the forced boundary: reasoning chunk ends with partial marker prefix. - let full2 = "<|channel>thought\nRRXX<|channel>YY"; - let exp2 = route_with_cap(true, 1, &[full2]); - assert_no_markers(&exp2.0); - // Split only the post part after the forced close to keep reasoning counting stable - let pre2 = "<|channel>thought\nRR"; - let post2 = "XX<|channel>YY"; - let exp2_post = route_with_cap(true, 1, &[full2]); - for split in 0..=post2.len() { - if split != 0 && !post2.is_char_boundary(split) { - continue; - } - let got = if split == 0 || split == post2.len() { - route_with_cap(true, 1, &[full2]) - } else { - route_with_cap(true, 1, &[pre2, &post2[..split], &post2[split..]]) - }; - assert_eq!( - got.0, exp2.0, - "forced split2 mismatch at post split {}", - split - ); - assert_no_markers(&got.0); - } - // Verify that a marker arriving strictly after forced transition - // as a separate push is still stripped at every internal split. - for payload in &[ - "helloworld", - "hello<|channel>world", - "hello<|channel|>world", - "helloworld", - "hello<|turn>world", - ] { - let expected_payload = (*payload) - .replace("<|channel>thought", "") - .replace("<|channel>", "") - .replace("<|channel|>", "") - .replace("", "") - .replace("<|turn>", "") - .replace("", ""); - for split in 0..=payload.len() { - if split != 0 && !payload.is_char_boundary(split) { - continue; - } - // Build payload split after forced transition. - let full = if split == 0 || split == payload.len() { - format!("<|channel>thought\nZ{}", payload) - } else { - // Simulate payload split across two pushes after forced. - // Create a single concatenated string and test split invariance - // via the full-string split test already done; here just - // verify the payload alone after forced. - let c1 = &payload[..split]; - let c2 = &payload[split..]; - let mut rr = GemmaThoughtRouter::new(true, 1); - let _ = rr.push("<|channel>thought\nZ"); - let mut vis = String::new(); - for ev in rr.push("").0 { - if let GemmaEmit::Token(t) = ev { - vis.push_str(&t); - } - } - for ev in rr.push(c1).0 { - if let GemmaEmit::Token(t) = ev { - vis.push_str(&t); - } - } - for ev in rr.push(c2).0 { - if let GemmaEmit::Token(t) = ev { - vis.push_str(&t); - } - } - for ev in rr.flush() { - if let GemmaEmit::Token(t) = ev { - vis.push_str(&t); - } - } - assert_eq!( - vis, expected_payload, - "payload {:?} split {}", - payload, split - ); - assert_no_markers(&vis); - continue; - }; - let got = route_with_cap(true, 1, &[&full]); - assert!(got.0.contains(&expected_payload) || got.0 == expected_payload); - assert_no_markers(&got.0); - } - } - } -} - -pub fn gemma_longest_marker_suffix(s: &str) -> usize { - const MARKERS: &[&str] = &[ - "<|channel>thought", - "<|channel>", - "<|channel|>", - "", - "<|turn>", - "", - ]; - for len in (1..=s.len()).rev() { - if !s.is_char_boundary(s.len() - len) { - continue; - } - let suffix = &s[s.len() - len..]; - if MARKERS.iter().any(|m| m.starts_with(suffix)) { - return len; - } - } - 0 -} - pub fn is_marker_prefix(s: &str) -> bool { const MARKERS: &[&str] = &[ "<|start|>", diff --git a/crates/hipfire-generate/src/redline.rs b/crates/hipfire-generate/src/redline.rs index 001581c16a..c9414af24d 100644 --- a/crates/hipfire-generate/src/redline.rs +++ b/crates/hipfire-generate/src/redline.rs @@ -288,7 +288,7 @@ fn redline_append_mapped( fn redline_gemma4_snapshot( gpu: &rdna_compute::Gpu, - bundle: &hipfire_loader::Gemma4LoweredBundle, + bundle: &hipfire_loader::Gemma4Bundle, _position: usize, ) -> Result { fn append_kv( @@ -317,8 +317,9 @@ fn redline_gemma4_snapshot( Ok(()) } + let state = &bundle.state; let mut logits = Vec::new(); - redline_append_buffer(gpu, &mut logits, &bundle.scratch.logits.buf)?; + redline_append_buffer(gpu, &mut logits, &state.logits.buf)?; let mut sliding_kv = Vec::new(); let mut sliding_kv_regions = Vec::new(); append_kv( @@ -326,7 +327,7 @@ fn redline_gemma4_snapshot( &mut sliding_kv, &mut sliding_kv_regions, "sliding", - &bundle.kv_sliding, + &state.kv_sliding, )?; let mut full_kv = Vec::new(); let mut full_kv_regions = Vec::new(); @@ -335,84 +336,59 @@ fn redline_gemma4_snapshot( &mut full_kv, &mut full_kv_regions, "full", - &bundle.kv_full, + &state.kv_full, )?; - macro_rules! snapshot_buffer { - ($buffer:expr) => {{ - let mut bytes = Vec::new(); - redline_append_buffer(gpu, &mut bytes, $buffer)?; - bytes - }}; - } + let snapshot = |tensor: Option<&rdna_compute::GpuTensor>| -> Result, String> { + let mut bytes = Vec::new(); + if let Some(tensor) = tensor { + redline_append_buffer(gpu, &mut bytes, &tensor.buf)?; + } + Ok(bytes) + }; + let moe = state.moe.as_ref(); Ok(RedlineGemma4Snapshot { logits, sliding_kv, full_kv, sliding_kv_regions, full_kv_regions, - scratch_x: snapshot_buffer!(&bundle.scratch.x.buf), - scratch_tmp: snapshot_buffer!(&bundle.scratch.tmp.buf), - scratch_q: snapshot_buffer!(&bundle.scratch.q.buf), - scratch_k: snapshot_buffer!(&bundle.scratch.k.buf), - scratch_v: snapshot_buffer!(&bundle.scratch.v.buf), - scratch_attn_out: snapshot_buffer!(&bundle.scratch.attn_out.buf), - scratch_residual: snapshot_buffer!(&bundle.scratch.residual.buf), - scratch_moe_cur_mlp: snapshot_buffer!(&bundle.scratch.moe_cur_mlp.buf), - scratch_moe_cur_moe: snapshot_buffer!(&bundle.scratch.moe_cur_moe.buf), - scratch_gate_ffn: snapshot_buffer!(&bundle.scratch.gate_ffn.buf), - scratch_up_ffn: snapshot_buffer!(&bundle.scratch.up_ffn.buf), - scratch_ffn_hidden: snapshot_buffer!(&bundle.scratch.ffn_hidden.buf), - scratch_ffn_out: snapshot_buffer!(&bundle.scratch.ffn_out.buf), + scratch_x: snapshot(Some(&state.x))?, + scratch_tmp: snapshot(Some(&state.tmp))?, + scratch_q: snapshot(Some(&state.q))?, + scratch_k: snapshot(Some(&state.k))?, + scratch_v: snapshot(Some(&state.v))?, + scratch_attn_out: snapshot(Some(&state.attn_out))?, + scratch_residual: snapshot(Some(&state.residual))?, + scratch_moe_cur_mlp: snapshot(moe.map(|m| &m.cur_mlp))?, + scratch_moe_cur_moe: snapshot(moe.map(|m| &m.cur_moe))?, + scratch_gate_ffn: snapshot(Some(&state.gate_ffn))?, + scratch_up_ffn: snapshot(Some(&state.up_ffn))?, + scratch_ffn_hidden: snapshot(Some(&state.ffn_hidden))?, + scratch_ffn_out: snapshot(Some(&state.ffn_out))?, }) } fn redline_reset_gemma4( gpu: &mut rdna_compute::Gpu, - bundle: &mut hipfire_loader::Gemma4LoweredBundle, + bundle: &mut hipfire_loader::Gemma4Bundle, ) -> Result<(), String> { - bundle + let state = &mut bundle.state; + state .kv_sliding .clear_gpu(gpu) .map_err(|error| error.to_string())?; - bundle + state .kv_full .clear_gpu(gpu) .map_err(|error| error.to_string())?; - bundle.kv_sliding.compact_offset = 0; - bundle.kv_full.compact_offset = 0; + state.kv_sliding.compact_offset = 0; + state.kv_full.compact_offset = 0; // The three oracle arms must begin with identical scratch as well as KV. - // Reset the complete lowered scratch surface so inactive and tail lanes - // are deterministic across the HIP, kernarg-blob, and retained arms. - for buffer in [ - &bundle.scratch.x.buf, - &bundle.scratch.residual.buf, - &bundle.scratch.tmp.buf, - &bundle.scratch.q.buf, - &bundle.scratch.k.buf, - &bundle.scratch.v.buf, - &bundle.scratch.attn_out.buf, - &bundle.scratch.gate_ffn.buf, - &bundle.scratch.up_ffn.buf, - &bundle.scratch.ffn_hidden.buf, - &bundle.scratch.ffn_out.buf, - &bundle.scratch.logits.buf, - &bundle.scratch.flash_partials.buf, - &bundle.scratch.moe_cur_mlp.buf, - &bundle.scratch.moe_pre2.buf, - &bundle.scratch.moe_router_in.buf, - &bundle.scratch.moe_router_logits.buf, - &bundle.scratch.moe_topk_indices.buf, - &bundle.scratch.moe_topk_weights.buf, - &bundle.scratch.moe_cur_moe.buf, - &bundle.scratch.moe_expert_gate_up.buf, - &bundle.scratch.moe_expert_hidden.buf, - &bundle.scratch.moe_expert_out.buf, - &bundle.scratch.moe_pre2_rot.buf, - &bundle.scratch.moe_expert_gate_batch.buf, - &bundle.scratch.moe_expert_up_batch.buf, - ] { + // Reset the complete scratch surface so inactive and tail lanes are + // deterministic across the HIP, kernarg-blob, and retained arms. + for tensor in state.scratch_tensors() { gpu.hip - .memset(buffer, 0, buffer.size()) + .memset(&tensor.buf, 0, tensor.buf.size()) .map_err(|error| error.to_string())?; } gpu.invalidate_graph_state(); @@ -423,67 +399,41 @@ fn redline_reset_gemma4( fn redline_prime_gemma4( gpu: &mut rdna_compute::Gpu, - bundle: &mut hipfire_loader::Gemma4LoweredBundle, + bundle: &mut hipfire_loader::Gemma4Bundle, context: usize, ) -> Result<(), String> { for i in 0..context { - hipfire_arch_gemma4::lowered::forward_scratch( - gpu, - &bundle.weights, + hipfire_arch_gemma4::forward::decode_step( &bundle.config, + &bundle.weights, + &mut bundle.state, + gpu, 10 + (i as u32 % 1000), - i, - &mut bundle.kv_sliding, - &mut bundle.kv_full, - &bundle.scratch, - ) - .map_err(|error| error.to_string())?; + i as u32, + )?; } gpu.hip .device_synchronize() .map_err(|error| error.to_string()) } +/// Stage `token`'s embedding and `position` exactly as a decode step does, +/// outside the recorded body. fn redline_prepare_gemma4( gpu: &mut rdna_compute::Gpu, - bundle: &hipfire_loader::Gemma4LoweredBundle, + bundle: &mut hipfire_loader::Gemma4Bundle, token: u32, position: usize, ) -> Result<(), String> { - use hipfire_runtime::llama::EmbeddingFormat; - - match bundle.weights.embd_format { - EmbeddingFormat::HFQ4G256 => gpu.embedding_lookup_hfq4g256( - &bundle.weights.embed_tokens, - &bundle.scratch.x, - token, - bundle.config.dim, - ), - EmbeddingFormat::HFQ4G128 => gpu.embedding_lookup_hfq4g128( - &bundle.weights.embed_tokens, - &bundle.scratch.x, - token, - bundle.config.dim, - ), - EmbeddingFormat::Q8_0 => gpu.embedding_lookup_q8( - &bundle.weights.embed_tokens, - &bundle.scratch.x, - token, - bundle.config.dim, - ), - EmbeddingFormat::F32 => gpu.embedding_lookup( - &bundle.weights.embed_tokens, - &bundle.scratch.x, - token, - bundle.config.dim, - ), - _ => return Err("unsupported Gemma4 Redline embedding format".into()), - } - .map_err(|error| error.to_string())?; - gpu.scale_f32(&bundle.scratch.x, bundle.config.embed_scale) - .map_err(|error| error.to_string())?; + hipfire_arch_gemma4::forward::prepare_token( + &bundle.config, + &bundle.weights, + &mut bundle.state, + gpu, + token, + )?; gpu.hip - .memcpy_htod(&bundle.scratch.pos_buf, &(position as i32).to_ne_bytes()) + .memcpy_htod(&bundle.state.pos_buf.buf, &(position as i32).to_ne_bytes()) .map_err(|error| error.to_string()) } @@ -4209,8 +4159,8 @@ fn redline_shadow_gemma4( let replay_arm = (|| -> Result<(RedlineGemma4Snapshot, f64, f64), String> { let bundle = loaded - .gemma4_lowered_mut() - .ok_or("Gemma4 Redline shadow requires lowered state")?; + .gemma4_mut() + .ok_or("Gemma4 Redline shadow requires Gemma 4 state")?; redline_reset_gemma4(gpu, bundle)?; redline_prime_gemma4(gpu, bundle, context)?; let started = Instant::now(); @@ -4256,8 +4206,8 @@ fn redline_shadow_gemma4( let blob_snapshot = (|| -> Result { let bundle = loaded - .gemma4_lowered_mut() - .ok_or("Gemma4 Redline blob oracle requires lowered state")?; + .gemma4_mut() + .ok_or("Gemma4 Redline blob oracle requires Gemma 4 state")?; redline_reset_gemma4(gpu, bundle)?; redline_prime_gemma4(gpu, bundle, context)?; for i in 0..iterations { @@ -4273,24 +4223,21 @@ fn redline_shadow_gemma4( let hip_arm = (|| -> Result<(RedlineGemma4Snapshot, f64), String> { let bundle = loaded - .gemma4_lowered_mut() - .ok_or("Gemma4 Redline HIP oracle requires lowered state")?; + .gemma4_mut() + .ok_or("Gemma4 Redline HIP oracle requires Gemma 4 state")?; redline_reset_gemma4(gpu, bundle)?; redline_prime_gemma4(gpu, bundle, context)?; let started = Instant::now(); for i in 0..iterations { redline_prepare_gemma4(gpu, bundle, 101 + (i as u32 % 1000), position(i))?; - hipfire_arch_gemma4::lowered::forward_scratch( - gpu, - &bundle.weights, + hipfire_arch_gemma4::forward::decode_step( &bundle.config, + &bundle.weights, + &mut bundle.state, + gpu, 101 + (i as u32 % 1000), - position(i), - &mut bundle.kv_sliding, - &mut bundle.kv_full, - &bundle.scratch, - ) - .map_err(|error| error.to_string())?; + position(i) as u32, + )?; } gpu.hip .device_synchronize() @@ -4389,7 +4336,7 @@ pub fn handle_redline_shadow( loaded.pp == 1 && loaded.ep.is_none() && loaded.state.as_ref().is_some_and(|state| { - (state.as_ref() as &dyn Any).is::() + (state.as_ref() as &dyn Any).is::() }) }) { let loaded = model.as_mut().expect("Gemma4 retained route checked"); diff --git a/crates/hipfire-loader/map.md b/crates/hipfire-loader/map.md index fd3911c2e5..d6420922db 100644 --- a/crates/hipfire-loader/map.md +++ b/crates/hipfire-loader/map.md @@ -29,7 +29,7 @@ _Generated by `scripts/check-crate-maps.py` from the tree — do not edit inside - [`src/admission.rs`](src/admission.rs): `EffectiveTopology`, `KvBackendHints`, `KvBackendRequest`, `QWEN4_MEMORY_HEADROOM_BYTES`, `Qwen4MemoryAdmission`, `Qwen4SourceAdmission`, `SequenceResolution`, `SourceAdmission`, `SourceAdmissionOptions`, `admit_qwen4_source`, `admit_source`, `admit_source_with_options`, `classify_vision`, `fn`, `from_override`, `growth`, `legacy_warning`, `ple_progress_printer`, `qwen4_kv_adaptive_requested`, `qwen4_memory_admission`, `qwen4_request_admission`, `qwen_default_q8_enabled`, `resolve_qwen4_ple_rows` - [`src/batch_staging.rs`](src/batch_staging.rs): `BatchStaging`, `VmmStagingRequest`, `qwen_batch_weight_formats_supported`, `stage_continuous_batch` - [`src/carriers.rs`](src/carriers.rs): `Cohere2MoeCarrier`, `Deepseek4Carrier`, `DotsOcrCarrier`, `FluxDiffusionCarrier`, `Gemma4Carrier`, `Lfm2MoeCarrier`, `LlamaCarrier`, `MapleCarrier`, `MinimaxCarrier`, `MuseGlimmerCarrier`, `Qwen2Carrier`, `Qwen35Carrier`, `Qwen4Carrier`, `prepare_host_memory_for` -- [`src/lib.rs`](src/lib.rs): `AsstTurnCache`, `BenchDecodeRoute`, `Carrier`, `ContinuousBatchRoute`, `Deepseek4HeterogeneousBundle`, `EpArch`, `EpEosRoute`, `EpPromptRoute`, `EpState`, `Eviction`, `GEMMA4_EAGLE_DRAFT_LEN`, `Gemma4Bundle`, `Gemma4EagleState`, `Gemma4LoweredBundle`, `GenerationEarlyRoute`, `GlimmerDrafterBundle`, `ImgRoute`, `LoadedModel`, `MuseGlimmerBundle`, `VisionRoute`, `ack_dims`, `admission`, `arch_label`, `as_arch_model`, `as_arch_model_mut`, `batch_staging`, `bench_decode_route`, `beta`, `budget`, `can_rewind_session_to`, `capture_row_elems`, `carrier_for`, `carriers`, `clear`, `cohere2moe`, `cohere2moe_mut`, `contains_key`, `continuous_batch_route`, `deepseek4`, `deepseek4_heterogeneous_mut`, `deepseek4_mut`, `device_hidden_capture_enabled`, `device_hidden_log_end`, `dots_ocr`, `dots_ocr_mut`, `drafter_ctx_capacity`, `ensure_vmm_ready_for_load`, `ep_admission`, `ep_eos_route`, `ep_prompt_route`, `flux_pipe`, `flux_pipe_mut`, `free_gpu`, `gemma4`, `gemma4_batched_prefill_optin`, `gemma4_eagle_spec_len`, `gemma4_lowered_mut`, `gemma4_mut`, `generation_early_route`, `get`, `has_vision_encoder`, `host_capture_rows`, `img_route`, `insert`, `is_diffusion`, `is_empty`, `lfm2_vision`, `lfm2moe`, `lfm2moe_mut`, `llama`, `llama_mut`, `load_admitted_with_gemma4_drafter`, `load_model`, `load_model_ep`, `load_model_ep_admitted`, `load_model_ep_with_compressor_cache`, `load_model_ep_with_kv_mode`, `load_model_with_gemma4_drafter`, `load_model_with_kv_backend`, `maple`, `maple_mut`, `maybe_evict`, `minimax`, `minimax_mut`, `model_dims`, `muse_glimmer`, `muse_glimmer_mut`, `n_tokens`, `new_from_env`, `qwen2_mut`, `qwen35`, `qwen35_ep_moe_topology_refusal`, `qwen35_mut`, `qwen4`, `qwen4_mut`, `reset_session_state`, `rewind_session_to`, `rewrite_muse_glimmer_onyx_template`, `skeleton`, `skeleton_pp`, `spec_build`, `touch_mru`, `unload_model`, `vision_config`, `vision_config_mut`, `vision_route`, `vision_weights`, `vision_weights_mut` +- [`src/lib.rs`](src/lib.rs): `AsstTurnCache`, `BenchDecodeRoute`, `Carrier`, `ContinuousBatchRoute`, `Deepseek4HeterogeneousBundle`, `EpArch`, `EpEosRoute`, `EpPromptRoute`, `EpState`, `Eviction`, `GEMMA4_EAGLE_DRAFT_LEN`, `GenerationEarlyRoute`, `GlimmerDrafterBundle`, `ImgRoute`, `LoadedModel`, `MuseGlimmerBundle`, `VisionRoute`, `ack_dims`, `admission`, `arch_label`, `as_arch_model`, `as_arch_model_mut`, `batch_staging`, `bench_decode_route`, `beta`, `budget`, `can_rewind_session_to`, `capture_row_elems`, `carrier_for`, `carriers`, `clear`, `cohere2moe`, `cohere2moe_mut`, `contains_key`, `continuous_batch_route`, `deepseek4`, `deepseek4_heterogeneous_mut`, `deepseek4_mut`, `device_hidden_capture_enabled`, `device_hidden_log_end`, `dots_ocr`, `dots_ocr_mut`, `drafter_ctx_capacity`, `ensure_vmm_ready_for_load`, `ep_admission`, `ep_eos_route`, `ep_prompt_route`, `flux_pipe`, `flux_pipe_mut`, `free_gpu`, `gemma4`, `gemma4_eagle_spec_len`, `gemma4_mut`, `generation_early_route`, `get`, `has_vision_encoder`, `host_capture_rows`, `img_route`, `insert`, `is_diffusion`, `is_empty`, `lfm2_vision`, `lfm2moe`, `lfm2moe_mut`, `llama`, `llama_mut`, `load_admitted_with_gemma4_drafter`, `load_model`, `load_model_ep`, `load_model_ep_admitted`, `load_model_ep_with_compressor_cache`, `load_model_ep_with_kv_mode`, `load_model_with_gemma4_drafter`, `load_model_with_kv_backend`, `maple`, `maple_mut`, `maybe_evict`, `minimax`, `minimax_mut`, `model_dims`, `muse_glimmer`, `muse_glimmer_mut`, `n_tokens`, `new_from_env`, `qwen2_mut`, `qwen35`, `qwen35_ep_moe_topology_refusal`, `qwen35_mut`, `qwen4`, `qwen4_mut`, `reset_session_state`, `rewind_session_to`, `rewrite_muse_glimmer_onyx_template`, `skeleton`, `skeleton_pp`, `spec_build`, `touch_mru`, `unload_model`, `vision_config`, `vision_config_mut`, `vision_route`, `vision_weights`, `vision_weights_mut` - [`src/spec_build.rs`](src/spec_build.rs): `Qwen35DenseTpGuard`, `Qwen35SlotGuard`, `build_speculator`, `is_qwen35_dense_tp`, `model_slot`, `take` ### Dependencies (from `Cargo.toml`) diff --git a/crates/hipfire-loader/src/admission.rs b/crates/hipfire-loader/src/admission.rs index 15c87bd1dc..48dec1f2b6 100644 --- a/crates/hipfire-loader/src/admission.rs +++ b/crates/hipfire-loader/src/admission.rs @@ -1392,12 +1392,6 @@ pub fn admit_source_with_options( )); } } - // Gemma 4 lowered min-context: refuse before teardown/alloc so a - // small max_seq leaves the prior model serving. Eager stays exempt. - if matches!(arch_id, 13 | 22) { - let use_lowered = hipfire_arch_gemma4::gemma4_source_uses_lowered(hfq); - hipfire_arch_gemma4::gemma4_context_admission(effective_seq, use_lowered)?; - } } // Head overlay (`params.head`): validated AND attached to the retained // base here, so the admitted source is already effective and loading diff --git a/crates/hipfire-loader/src/carriers.rs b/crates/hipfire-loader/src/carriers.rs index 18010b4119..b7e3cffdb9 100644 --- a/crates/hipfire-loader/src/carriers.rs +++ b/crates/hipfire-loader/src/carriers.rs @@ -3464,16 +3464,6 @@ impl Carrier for MapleCarrier { // ─── Gemma4Carrier ─────────────────────────────────────────────────── -fn gemma4_validate_drafter_route(is_e_series: bool, has_drafter: bool) -> Result<(), String> { - if is_e_series && has_drafter { - return Err( - "gemma4: E2B/E4B EAGLE spec-decode is not yet supported; load the E-series target without params.drafter" - .into(), - ); - } - Ok(()) -} - pub struct Gemma4Carrier; impl Carrier for Gemma4Carrier { fn name(&self) -> &'static str { @@ -3481,16 +3471,22 @@ impl Carrier for Gemma4Carrier { } fn spec_target_guard<'m>( &self, - _state: &'m mut Option>, + state: &'m mut Option>, _model_path: &str, ) -> Result, String> { - Err("gemma4: spec decode not yet wired (AR-only)".into()) + match state + .as_mut() + .and_then(|s| (s.as_mut() as &mut dyn Any).downcast_mut::()) + { + Some(bundle) => Ok(Box::new(InPlaceGuard { bundle })), + _ => Err("gemma4: spec target state mismatch".into()), + } } fn make_spec_emitter<'a>( &self, - _ctx: SpecEmitCtx<'a>, + ctx: SpecEmitCtx<'a>, ) -> Result, String> { - Err("gemma4: spec emitter not yet wired".into()) + Ok(Box::new(hipfire_arch_gemma4::emit::GemmaSpecEmit::new(ctx))) } fn claims_arch_id(&self, arch_id: u32, _is_dir: bool) -> bool { // 13 = gemma4_text (primary), 22 = gemma4_unified_assistant (EAGLE drafter sidecar). @@ -3546,33 +3542,14 @@ impl Carrier for Gemma4Carrier { tok, i as u32, ) { - *prefill_err = Some(format!("gemma4 eager bench prefill failed: {error:?}")); - return Some(false); - } - } - return Some(true); - } - - if let Some(bundle) = m.gemma4_lowered_mut() { - for (i, &tok) in synthetic.iter().enumerate() { - if let Err(error) = hipfire_arch_gemma4::lowered::forward_scratch( - gpu, - &bundle.weights, - &bundle.config, - tok, - i, - &mut bundle.kv_sliding, - &mut bundle.kv_full, - &bundle.scratch, - ) { - *prefill_err = Some(format!("gemma4 lowered bench prefill failed: {error:?}")); + *prefill_err = Some(format!("gemma4 bench prefill failed: {error:?}")); return Some(false); } } return Some(true); } - *prefill_err = Some("gemma4 bench prefill missing eager/lowered state".into()); + *prefill_err = Some("gemma4 bench prefill missing state".into()); Some(false) } fn bench_decode_prime( @@ -3611,34 +3588,14 @@ impl Carrier for Gemma4Carrier { token, (context + i) as u32, ) { - *decode_err = Some(format!("gemma4 eager bench decode failed: {error:?}")); - return Some(false); - } - } - return Some(true); - } - - if let Some(bundle) = m.gemma4_lowered_mut() { - for i in 0..iterations { - let token = 101 + (i as u32 % 1000); - if let Err(error) = hipfire_arch_gemma4::lowered::forward_scratch( - gpu, - &bundle.weights, - &bundle.config, - token, - context + i, - &mut bundle.kv_sliding, - &mut bundle.kv_full, - &bundle.scratch, - ) { - *decode_err = Some(format!("gemma4 lowered bench decode failed: {error:?}")); + *decode_err = Some(format!("gemma4 bench decode failed: {error:?}")); return Some(false); } } return Some(true); } - *decode_err = Some("gemma4 bench decode missing eager/lowered state".into()); + *decode_err = Some("gemma4 bench decode missing state".into()); Some(false) } fn load(&self, src: ModelSource, ctx: &mut LoadCtx) -> Result { @@ -3650,175 +3607,75 @@ impl Carrier for Gemma4Carrier { } dir_diag(&src); let meta = resolve_source_meta(&src, ctx.path)?; - let bundle = hipfire_arch_gemma4::load_gemma4_bundle(src, ctx)?; - match bundle { - hipfire_arch_gemma4::Gemma4Bundle::Lowered(l) => { - let eos_tok = resolve_eos_tok( - &meta.tokenizer, - &["", "", "", "<|im_end|>"], - ); - let speculator = crate::spec_build::build_speculator( - meta.arch_id, - None, - None, - true, - ctx.max_seq, - ctx.spec, - ); - Ok(LoadedModel { - state: Some(Box::new(crate::Gemma4LoweredBundle { - config: l.config, - weights: l.weights, - scratch: l.scratch, - kv_sliding: l.kv_sliding, - kv_full: l.kv_full, - eos_tok, - })), - speculator, - ..LoadedModel::skeleton( - meta.arch_id, - meta.tokenizer, - ctx.max_seq, - ctx.max_seq, - ctx.path.to_string(), - meta.chat_template, - ) - }) - } - hipfire_arch_gemma4::Gemma4Bundle::Eager(e) => { - let eos_tok = resolve_eos_tok( - &meta.tokenizer, - &["", "", "", "<|im_end|>"], - ); - let _ = &e.weights; - // Optional EAGLE drafter (arch-22) — populated only when - // `gemma4_drafter_path` is Some. Validates draft_len 1..=5, - // arch_id 22, and backbone_hidden == target dim. On failure - // logs and falls back to AR-only (mirrors PR's contract) to - // avoid hard failing a valid target model due to a bad sidecar. - let eagle = if let Some(dp) = ctx.gemma4_drafter_path { - let draft_len = crate::gemma4_eagle_spec_len(Some(ctx.gemma4_draft_len as u64)) - .map_err(|e| format!("gemma4 drafter spec_len: {e}"))?; - match load_gemma4_eagle_state(dp, draft_len, &e.config, &e.weights, ctx.gpu) { - Ok(st) => { - eprintln!( - " gemma4 EAGLE drafter loaded: {} (layers={}, hidden={}, draft_len={})", - dp, st.drafter_config.n_layers, st.drafter_config.hidden, st.draft_len, - ); - Some(st) - } - Err(e) => { - eprintln!( - " gemma4 EAGLE drafter load failed ({}): {} — falling back to AR only", - dp, e - ); - None - } - } - } else { - None - }; - let speculator = crate::spec_build::build_speculator( - meta.arch_id, - None, - None, - true, - ctx.max_seq, - ctx.spec, - ); - Ok(LoadedModel { - state: Some(Box::new(crate::Gemma4Bundle { - config: e.config, - weights: e.weights, - state: e.state, - eos_tok, - eagle, - })), - speculator, - ..LoadedModel::skeleton( - meta.arch_id, - meta.tokenizer, - ctx.max_seq, - ctx.max_seq, - ctx.path.to_string(), - meta.chat_template, - ) - }) - } + let session_domain = match &src { + ModelSource::Hfq(hfq) => Some(hipfire_runtime::serve_contract::CacheDomain::for_model( + hfq, + &meta.tokenizer, + meta.chat_template.as_deref(), + "gemma4", + ctx.gpu.device_id, + )), + ModelSource::Dir(_) => None, + }; + let mut bundle = hipfire_arch_gemma4::load_gemma4_bundle(src, ctx)?; + let session_budget = hipfire_config::memory::session_cache_bytes(); + let snapshots = hipfire_runtime::session_cache::SessionState::snapshot_scope( + &bundle, + hipfire_runtime::session_cache::SessionRoute::Ar, + ) + .is_some(); + if let Some(domain) = session_domain.filter(|_| session_budget > 0 && snapshots) { + bundle.attach_session_cache(hipfire_runtime::session_cache::SessionCache::new( + domain, + session_budget, + )); + eprintln!( + " gemma4 session cache: {} MiB budget, {} MiB per {}-token snapshot", + session_budget >> 20, + bundle.session_snapshot_bytes() >> 20, + hipfire_arch_gemma4::bundle::SESSION_STRIDE, + ); } - } -} - -fn load_gemma4_eagle_state( - drafter_path: &str, - draft_len: usize, - target_cfg: &hipfire_arch_gemma4::config::Gemma4Config, - target_weights: &hipfire_arch_gemma4::gemma4::Gemma4Weights, - gpu: &mut rdna_compute::Gpu, -) -> Result { - use std::path::Path; - let dhfq = hipfire_runtime::hfq::HfqFile::open(Path::new(drafter_path)) - .map_err(|e| format!("open gemma4 drafter: {e}"))?; - if dhfq.arch_id != 22 { - return Err(format!( - "gemma4 EAGLE drafter must be arch_id=22 (gemma4_unified_assistant); got arch_id={} — a DFlash draft goes in params.draft on qwen3.5 targets, not params.drafter", - dhfq.arch_id - )); - } - let dcfg = hipfire_arch_gemma4::drafter::Gemma4DrafterConfig::from_hfq(&dhfq)?; - if dcfg.backbone_hidden != target_cfg.dim { - return Err(format!( - "drafter backbone_hidden ({}) != target hidden ({}) — this drafter was trained against a different target width", - dcfg.backbone_hidden, target_cfg.dim - )); - } - let drafter_weights = - hipfire_arch_gemma4::drafter::Gemma4DrafterWeights::load(&dhfq, &dcfg, gpu)?; - let drafter_scratch = hipfire_arch_gemma4::drafter::Gemma4DrafterScratch::new(gpu, &dcfg) - .map_err(|e| format!("gemma4 drafter scratch: {e}"))?; - let spec_scratch = - hipfire_arch_gemma4::speculative::Gemma4SpecScratch::new(gpu, target_cfg, draft_len) - .map_err(|e| format!("gemma4 spec scratch: {e}"))?; - // Prime the batched verify path (b=1 then real block size) on disposable - // throwaway states — mirrors PR's warmup to ensure kernels are compiled. - let warm_b = draft_len + 1; - { - let mut warm = - hipfire_arch_gemma4::gemma4::Gemma4State::new_with_max_seq(gpu, target_cfg, warm_b + 4) - .map_err(|e| format!("gemma4-eagle: warm state: {e}"))?; - let _ = hipfire_arch_gemma4::forward::forward_batch( - target_cfg, - target_weights, - &mut warm, - gpu, - &[target_cfg.bos_token], - 0, + bundle.eos_tok = resolve_eos_tok( + &meta.tokenizer, + &["", "", "", "<|im_end|>"], ); - let _ = gpu.hip.device_synchronize(); - warm.free_gpu(gpu); + // Optional EAGLE drafter (arch-22) as the request speculator. A bad + // sidecar logs and leaves the target AR-only. + let mut speculator = None; + if let Some(dp) = ctx.gemma4_drafter_path { + let draft_len = crate::gemma4_eagle_spec_len(Some(ctx.gemma4_draft_len as u64)) + .map_err(|e| format!("gemma4 drafter spec_len: {e}"))?; + match hipfire_arch_gemma4::Gemma4Drafter::load(dp, draft_len, &bundle, ctx.gpu) { + Ok(drafter) => { + eprintln!( + " gemma4 EAGLE drafter loaded: {} (layers={}, hidden={}, draft_len={draft_len})", + dp, + drafter.config().n_layers, + drafter.config().hidden, + ); + speculator = Some(Box::new(hipfire_runtime::spec::MtpSpeculator::new(drafter)) + as Box); + } + Err(e) => eprintln!( + " gemma4 EAGLE drafter load failed ({}): {} — falling back to AR only", + dp, e + ), + } + } + Ok(LoadedModel { + state: Some(Box::new(bundle)), + speculator, + ..LoadedModel::skeleton( + meta.arch_id, + meta.tokenizer, + ctx.max_seq, + ctx.max_seq, + ctx.path.to_string(), + meta.chat_template, + ) + }) } - { - let mut warm = - hipfire_arch_gemma4::gemma4::Gemma4State::new_with_max_seq(gpu, target_cfg, warm_b + 4) - .map_err(|e| format!("gemma4-eagle: warm state 2: {e}"))?; - let _ = hipfire_arch_gemma4::forward::forward_batch( - target_cfg, - target_weights, - &mut warm, - gpu, - &vec![target_cfg.bos_token; warm_b], - 0, - ); - let _ = gpu.hip.device_synchronize(); - warm.free_gpu(gpu); - } - Ok(crate::Gemma4EagleState { - drafter_config: dcfg, - drafter_weights, - drafter_scratch, - spec_scratch, - draft_len, - }) } // ─── MuseGlimmerCarrier ──────────────────────────────────────────────── @@ -4392,18 +4249,6 @@ impl Carrier for FluxDiffusionCarrier { } } -#[cfg(test)] -mod gemma4_route_tests { - use super::gemma4_validate_drafter_route; - - #[test] - fn e_series_drafter_fails_closed() { - assert!(gemma4_validate_drafter_route(true, true).is_err()); - assert!(gemma4_validate_drafter_route(true, false).is_ok()); - assert!(gemma4_validate_drafter_route(false, true).is_ok()); - } -} - #[cfg(test)] mod qwen4_source_policy_tests { use super::qwen4_use_range_payload; diff --git a/crates/hipfire-loader/src/lib.rs b/crates/hipfire-loader/src/lib.rs index 75acad3c82..666934155d 100644 --- a/crates/hipfire-loader/src/lib.rs +++ b/crates/hipfire-loader/src/lib.rs @@ -622,73 +622,10 @@ impl hipfire_runtime::arch_model::ArchModel for MuseGlimmerBundle { } } -/// `ArchModel` for the three bundles the loader defines itself. +/// `ArchModel` for the bundle the loader defines itself. /// -/// `Gemma4Bundle`, `Gemma4LoweredBundle` and `Deepseek4HeterogeneousBundle` are -/// declared here rather than in their arch crates, so — orphan rule — the impls -/// must be here too. Each `free_gpu` mirrors its `unload_model` arm exactly; -/// freeing less leaks, freeing more double-frees. -impl hipfire_runtime::arch_model::ArchModel for Gemma4Bundle { - fn dim(&self) -> usize { - self.config.dim - } - fn n_layers(&self) -> usize { - self.config.n_layers - } - fn vocab_size(&self) -> usize { - self.config.vocab_size - } - fn arch_key(&self) -> &'static str { - "gemma4" - } - fn kv_cache_mut(&mut self) -> Option<&mut hipfire_runtime::llama::KvCache> { - None - } - fn reset_session_state(&mut self, _gpu: &mut rdna_compute::Gpu) -> Result<(), String> { - // Mirrors daemon reset arms (main.rs:3336 / 3387): `bundle.state.reset()`. - self.state.reset(); - Ok(()) - } - fn free_gpu(self: Box, gpu: &mut rdna_compute::Gpu) { - let b = *self; - if let Some(eagle) = b.eagle { - eagle.spec_scratch.free(gpu); - eagle.drafter_scratch.free_gpu(gpu); - eagle.drafter_weights.free_gpu(gpu); - } - b.state.free_gpu(gpu); - b.weights.free_gpu(gpu); - } -} - -impl hipfire_runtime::arch_model::ArchModel for Gemma4LoweredBundle { - fn dim(&self) -> usize { - self.config.dim - } - fn n_layers(&self) -> usize { - self.config.n_layers - } - fn vocab_size(&self) -> usize { - self.config.vocab_size - } - fn arch_key(&self) -> &'static str { - "gemma4" - } - fn kv_cache_mut(&mut self) -> Option<&mut hipfire_runtime::llama::KvCache> { - // Two caches (q8 sliding + q8/legacy-asym3 full) and no basis for preferring - // one, so expose neither rather than silently picking. - None - } - fn free_gpu(self: Box, gpu: &mut rdna_compute::Gpu) { - let b = *self; - // Reverse of lowered load construction: full → sliding → scratch → weights. - let _ = b.kv_full.free_gpu(gpu); - let _ = b.kv_sliding.free_gpu(gpu); - b.scratch.free_gpu(gpu); - b.weights.free_gpu(gpu); - } -} - +/// `Deepseek4HeterogeneousBundle` is declared here rather than in its arch +/// crate, so — orphan rule — the impl must be here too. impl hipfire_runtime::arch_model::ArchModel for Deepseek4HeterogeneousBundle { fn dim(&self) -> usize { self.model.config.hidden_size @@ -740,57 +677,11 @@ pub use minimax::MiniMaxBundle; /// Field-identical to the prior loader-local struct. pub use cohere2moe::Cohere2MoeBundle; -/// Gemma 4 EAGLE speculative-decode scratch state (arch-22 drafter riding an -/// arch-13 target). Populated only when `LoadCtx::gemma4_drafter_path` is -/// `Some`. The drafter has NO KV cache of its own — it queries the target's -/// KV at a constant position, so the only per-session state is the target's -/// `Gemma4State` plus this scratch. Mirrors the PR's `Gemma4EagleState`. -pub struct Gemma4EagleState { - pub drafter_config: gemma4::drafter::Gemma4DrafterConfig, - pub drafter_weights: gemma4::drafter::Gemma4DrafterWeights, - pub drafter_scratch: gemma4::drafter::Gemma4DrafterScratch, - /// Reusable seed/draft/verify hidden buffers for `spec_step_gemma4_eagle`, - /// sized for `draft_len` at load time. - pub spec_scratch: gemma4::speculative::Gemma4SpecScratch, - /// Drafts per round (verify block = draft_len + 1). - pub draft_len: usize, -} - -/// Gemma 4 dense text (arch_id=13) GPU bundle — eager dense path. Re-exported -/// from the arch crate, which owns config/weights/state so `impl SpecTarget` -/// can live there when needed. The `eagle` field holds the optional EAGLE -/// drafter state (arch-22) when `gemma4_drafter_path` was supplied; `None` -/// is AR-only. This keeps one `ModelState::Gemma4` variant for the dense -/// path, matching beta's pattern for other arches. -pub struct Gemma4Bundle { - pub config: gemma4::config::Gemma4Config, - pub weights: gemma4::gemma4::Gemma4Weights, - pub state: gemma4::gemma4::Gemma4State, - pub eos_tok: u32, - pub eagle: Option, -} - -/// Lowered / MoE Gemma 4 execution bundle (arch_id=13 26B-A4B + opt-in -/// batched/WMMA dense prefill). Uses `lowered::{Gemma4Config, Gemma4Weights, -/// Gemma4Scratch}` plus TWO `hipfire_runtime::llama::KvCache`s (q8 -/// ring-buffered sliding + full tier from `kv_cache`: q8, or legacy asym3 on -/// `legacy-asym3`) and `eos_tok`. Mutually exclusive -/// with `Gemma4Bundle` (eager) via the `ModelState` enum — a given -/// `LoadedModel` populates exactly one of the two variants. -/// Chose second `ModelState` variant over enum-inside-bundle to keep -/// `Gemma4Bundle`'s struct shape stable for the existing eager AR path that -/// `Gemma4Generate` already matches as `Some(ModelState::Gemma4(bundle))`. -pub struct Gemma4LoweredBundle { - pub config: gemma4::lowered::Gemma4Config, - pub weights: gemma4::lowered::Gemma4Weights, - pub scratch: gemma4::lowered::Gemma4Scratch, - pub kv_sliding: llama::KvCache, - pub kv_full: llama::KvCache, - pub eos_tok: u32, -} +/// Gemma 4 (arch_id=13) GPU bundle, owned by its arch crate (which also +/// implements `ArchModel` for it). +pub use gemma4::Gemma4Bundle; /// Muse Glimmer 30B dense text (arch_id=14) GPU bundle — eager dense path. -/// Mirrors Gemma4Bundle exactly (per cross-agent contract). pub struct MuseGlimmerBundle { pub config: glimmer::config::GlimmerConfig, pub weights: glimmer::glimmer::GlimmerWeights, @@ -1160,12 +1051,6 @@ impl LoadedModel { .and_then(|s| (s as &mut dyn Any).downcast_mut::()) } - pub fn gemma4_lowered_mut(&mut self) -> Option<&mut Gemma4LoweredBundle> { - self.state - .as_deref_mut() - .and_then(|s| (s as &mut dyn Any).downcast_mut::()) - } - pub fn muse_glimmer(&self) -> Option<&MuseGlimmerBundle> { self.state .as_deref() diff --git a/crates/hipfire-runtime/examples/build_kld_ref_native_gemma4.rs b/crates/hipfire-runtime/examples/build_kld_ref_native_gemma4.rs index ef00f10f28..1aac84ac85 100644 --- a/crates/hipfire-runtime/examples/build_kld_ref_native_gemma4.rs +++ b/crates/hipfire-runtime/examples/build_kld_ref_native_gemma4.rs @@ -10,11 +10,10 @@ //! format that `eval_hipfire_gemma4` consumes. //! //! Differences from the qwen35 arm, all gemma4-architectural: -//! * forward is `hipfire_arch_gemma4::lowered::forward_scratch` with the -//! dual KV (sliding F32 `new_gpu` + full asym3 `new_gpu_asym3`) exactly -//! as `gemma4_oracle.rs` (the HF-parity-validated config). The KV setup -//! is IDENTICAL between this ref builder and the candidate eval, so KV -//! noise cancels and the KLD isolates weight precision. +//! * forward is `hipfire_arch_gemma4::forward::decode_step` with Q8 KV on +//! both tiers. The KV setup is IDENTICAL between this ref builder and the +//! candidate eval, so KV noise cancels and the KLD isolates weight +//! precision. //! * gemma has add_bos=true: each n_ctx chunk is materialized as //! [BOS] + (n_ctx-1) corpus tokens (llama-perplexity add_bos semantics). //! The materialized chunk tokens (BOS included) are what's written to the @@ -31,9 +30,8 @@ //! --output -f32-native.kldref.bin [--max-chunks N] fn main() { - use hipfire_arch_gemma4::lowered::{self as gemma4, Gemma4Scratch}; + use hipfire_arch_gemma4::{forward, Gemma4Config, Gemma4State, Gemma4Weights}; use hipfire_runtime::hfq::HfqFile; - use hipfire_runtime::llama::KvCache; use std::cmp::Ordering; use std::fs::File; use std::io::{BufWriter, Write}; @@ -100,8 +98,8 @@ fn main() { } // -------- load oracle model + tokenizer -------- - let mut hfq = HfqFile::open(&model).expect("open oracle model"); - let config = gemma4::config_from_hfq(&hfq).expect("read config"); + let hfq = HfqFile::open(&model).expect("open oracle model"); + let config = Gemma4Config::from_hfq(&hfq).expect("read config"); let tokenizer = hipfire_runtime::tokenizer::Tokenizer::from_hfq_metadata(&hfq.metadata_json) .expect("tokenizer"); let mut gpu = rdna_compute::Gpu::init().expect("gpu init"); @@ -110,7 +108,7 @@ fn main() { gpu.arch, model.display() ); - let weights = gemma4::load_weights(&mut hfq, &config, &mut gpu).expect("load weights"); + let weights = Gemma4Weights::load(&hfq, &config, &mut gpu).expect("load weights"); eprintln!( "loaded {} layers, vocab={}, n_ctx={}, top_k={}, bos={}", config.n_layers, config.vocab_size, n_ctx, top_k, config.bos_token @@ -163,31 +161,10 @@ fn main() { out.write_all(&t.to_le_bytes()).unwrap(); } - // -------- scratch + dual KV (gemma4_oracle config: sliding F32, full asym3) -------- + // -------- state: Q8 KV on both tiers (the unified state has no F32 KV + // tier; the forward refuses F32 cache writes) -------- let kv_max = n_ctx + 16; - let scratch = Gemma4Scratch::new(&mut gpu, &config, kv_max).expect("scratch"); - gemma4::init_scratch_constants(&mut gpu, &scratch, config.full_head_dim) - .expect("init_scratch_constants"); - let mut kv_sliding = KvCache::new_gpu( - &mut gpu, - config.n_layers, - config.sliding_n_kv_heads, - config.sliding_head_dim, - kv_max, - ) - .expect("kv sliding alloc"); - // FULL KV = F32, NOT asym3: on gfx942/CDNA the asym3 full-KV path is - // catastrophically wrong (grows with depth; PPL 3826 at 512 ctx) while - // F32 full-KV is HF-EXACT (top-5 logits match HF to 1e-4 at 128 ids, - // 2026-06-10). F32 both sides also removes the shared KV-noise floor. - let mut kv_full = KvCache::new_gpu( - &mut gpu, - config.n_layers, - config.full_n_kv_heads, - config.full_head_dim, - kv_max, - ) - .expect("kv full alloc"); + let mut state = Gemma4State::new_with_max_seq(&mut gpu, &config, kv_max).expect("state"); // -------- per-chunk forward + top-K reduce -------- let k = top_k; @@ -203,21 +180,18 @@ fn main() { // eval relies on; gemma4 has no recurrent state to reset). let chunk = &tokens[c * n_ctx..(c + 1) * n_ctx]; for pos in 0..(n_ctx - 1) { - gemma4::forward_scratch( - &mut gpu, - &weights, + let cand_logits = forward::decode_step( &config, + &weights, + &mut state, + &mut gpu, chunk[pos], - pos, - &mut kv_sliding, - &mut kv_full, - &scratch, + pos as u32, ) - .expect("forward_scratch"); + .expect("decode_step"); if pos < scoring_start { continue; } - let cand_logits = gpu.download_f32(&scratch.logits).expect("download logits"); let cand_logits = &cand_logits[..config.vocab_size.min(cand_logits.len())]; // Convert logits -> full log-prob vector (fp64 log-softmax). diff --git a/crates/hipfire-runtime/examples/calib_sweep.rs b/crates/hipfire-runtime/examples/calib_sweep.rs index 9011fd8eae..38d30a0bb9 100755 --- a/crates/hipfire-runtime/examples/calib_sweep.rs +++ b/crates/hipfire-runtime/examples/calib_sweep.rs @@ -172,7 +172,7 @@ fn build_capture_llama( m } fn build_capture_gemma( - weights: &hipfire_arch_gemma4::lowered::Gemma4Weights, + weights: &hipfire_arch_gemma4::Gemma4Weights, start: usize, end: usize, prefix: &str, @@ -181,7 +181,7 @@ fn build_capture_gemma( for i in start..end.min(weights.layers.len()) { let p = format!("{prefix}layers.{i}"); match &weights.layers[i] { - hipfire_arch_gemma4::lowered::LayerWeights::Sliding(s) => { + hipfire_arch_gemma4::LayerWeights::Sliding(s) => { m.insert( s.q_proj.buf.buf.as_ptr() as usize, format!("{p}.self_attn.q_proj.weight"), @@ -211,7 +211,7 @@ fn build_capture_gemma( format!("{p}.mlp.down_proj.weight"), ); } - hipfire_arch_gemma4::lowered::LayerWeights::Full(f) => { + hipfire_arch_gemma4::LayerWeights::Full(f) => { m.insert( f.q_proj.buf.buf.as_ptr() as usize, format!("{p}.self_attn.q_proj.weight"), @@ -220,7 +220,13 @@ fn build_capture_gemma( f.k_proj.buf.buf.as_ptr() as usize, format!("{p}.self_attn.k_proj.weight"), ); - // no v_proj on Full — V reuses k_proj pre-norm + // Full layers with attention_k_eq_v reuse k_proj's pre-norm output as V. + if let Some(v) = &f.v_proj { + m.insert( + v.buf.buf.as_ptr() as usize, + format!("{p}.self_attn.v_proj.weight"), + ); + } m.insert( f.o_proj.buf.buf.as_ptr() as usize, format!("{p}.self_attn.o_proj.weight"), @@ -352,14 +358,14 @@ fn build_capture_lfm( m } fn estimate_peak_gemma( - weights: &hipfire_arch_gemma4::lowered::Gemma4Weights, + weights: &hipfire_arch_gemma4::Gemma4Weights, start: usize, end: usize, ) -> u64 { let mut b = 0u64; for i in start..end.min(weights.layers.len()) { let ks: Vec = match &weights.layers[i] { - hipfire_arch_gemma4::lowered::LayerWeights::Sliding(s) => vec![ + hipfire_arch_gemma4::LayerWeights::Sliding(s) => vec![ s.q_proj.k, s.k_proj.k, s.v_proj.k, @@ -368,14 +374,18 @@ fn estimate_peak_gemma( s.up_proj.k, s.down_proj.k, ], - hipfire_arch_gemma4::lowered::LayerWeights::Full(f) => vec![ - f.q_proj.k, - f.k_proj.k, - f.o_proj.k, - f.gate_proj.k, - f.up_proj.k, - f.down_proj.k, - ], + hipfire_arch_gemma4::LayerWeights::Full(f) => { + let mut ks = vec![ + f.q_proj.k, + f.k_proj.k, + f.o_proj.k, + f.gate_proj.k, + f.up_proj.k, + f.down_proj.k, + ]; + ks.extend(f.v_proj.as_ref().map(|v| v.k)); + ks + } }; for k in ks { b += compact_hessian_bytes(k) + (k * 4) as u64; @@ -739,24 +749,22 @@ fn run_llama_batched( } fn run_gemma_batched( gpu: &mut rdna_compute::Gpu, - weights: &hipfire_arch_gemma4::lowered::Gemma4Weights, - cfg: &hipfire_arch_gemma4::lowered::Gemma4Config, + weights: &hipfire_arch_gemma4::Gemma4Weights, + cfg: &hipfire_arch_gemma4::Gemma4Config, toks: &[u32], seq: usize, - kv_sliding: &mut hipfire_runtime::llama::KvCache, - kv_full: &mut hipfire_runtime::llama::KvCache, - scratch: &hipfire_arch_gemma4::lowered::Gemma4Scratch, + state: &mut hipfire_arch_gemma4::Gemma4State, ) -> Result<(), String> { for (si, chunk) in toks.chunks(seq.max(1)).enumerate() { - kv_sliding + state + .kv_sliding .clear_gpu(gpu) .map_err(|e| format!("gemma kv sliding clear {si}: {e:?}"))?; - kv_full + state + .kv_full .clear_gpu(gpu) .map_err(|e| format!("gemma kv full clear {si}: {e:?}"))?; - hipfire_arch_gemma4::lowered::forward_prefill_batch( - gpu, weights, cfg, chunk, 0, kv_sliding, kv_full, scratch, - ) + hipfire_arch_gemma4::forward::forward_prefill_batch(cfg, weights, state, gpu, chunk, 0) .map_err(|e| { format!( "gemma forward_prefill_batch seq {si} B {}: {e:?}", @@ -1610,8 +1618,8 @@ fn main() { } else if gemma { // Gemma4 — WIRED. forward_prefill_batch runs the declarative layer program; the // dispatch sandwich projections (gemm_rows / gemv) tap every projection input. - let cfg = hipfire_arch_gemma4::lowered::config_from_hfq(&hfq) - .unwrap_or_else(|| panic!("gemma4 cfg: missing/invalid config in HFQ")); + let cfg = hipfire_arch_gemma4::Gemma4Config::from_hfq(&hfq) + .unwrap_or_else(|e| panic!("gemma4 cfg: {e}")); eprintln!( "gemma4 arch=13 n_layers={} dim={} vocab={} hidden={} layer_types={:?} enable_moe={}", cfg.n_layers, @@ -1626,7 +1634,7 @@ fn main() { eprintln!("If this is the dense Gemma4-12B/27B (no MoE), enable_moe_block should be false — check the model is the dense variant."); std::process::exit(2); } - let mut weights = hipfire_arch_gemma4::lowered::load_weights(&mut hfq, &cfg, &mut gpu) + let weights = hipfire_arch_gemma4::Gemma4Weights::load(&hfq, &cfg, &mut gpu) .unwrap_or_else(|e| panic!("gemma4 weights: {e}")); { let all_names: Vec = @@ -1639,31 +1647,11 @@ fn main() { } drop(hfq); let kv_max = seq_len + 16; - let scratch = hipfire_arch_gemma4::lowered::Gemma4Scratch::new(&mut gpu, &cfg, kv_max) - .unwrap_or_else(|e| panic!("gemma4 scratch: {e:?}")); - hipfire_arch_gemma4::lowered::init_scratch_constants(&mut gpu, &scratch, cfg.full_head_dim) - .unwrap_or_else(|e| panic!("gemma4 init_scratch_constants: {e:?}")); - // Q8 KV, matching the qwen35 arm above. `new_gpu` allocates an F32 cache - // and gemma4's batched prefill has no `KvWriteF32` kernel registered, so - // an F32 cache fails at the first layer with "no implementation for - // KvWriteF32". KV dtype does not affect what is captured — the tap reads + // Q8 KV on both tiers: gemma4's batched prefill has no `KvWriteF32` + // kernel. KV dtype does not affect what is captured — the tap reads // the projection INPUT activations, not cache contents. - let mut kv_sliding = hipfire_runtime::llama::KvCache::new_gpu_q8( - &mut gpu, - cfg.n_layers, - cfg.sliding_n_kv_heads, - cfg.sliding_head_dim, - kv_max, - ) - .unwrap_or_else(|e| panic!("gemma sliding kv alloc q8 {kv_max}: {e:?}")); - let mut kv_full = hipfire_runtime::llama::KvCache::new_gpu_q8( - &mut gpu, - cfg.n_layers, - cfg.full_n_kv_heads, - cfg.full_head_dim, - kv_max, - ) - .unwrap_or_else(|e| panic!("gemma full kv alloc q8 {kv_max}: {e:?}")); + let mut state = hipfire_arch_gemma4::Gemma4State::new_with_max_seq(&mut gpu, &cfg, kv_max) + .unwrap_or_else(|e| panic!("gemma4 state {kv_max}: {e}")); // Honesty guard: the tap must fire; if layer count mismatch or unsupported embed format, run will Err and we exit 1 (honest) — never emit partial hfq. let lpp = lpp.max(1); let peak = estimate_peak_gemma(&weights, 0, lpp.min(cfg.n_layers)); @@ -1676,7 +1664,7 @@ fn main() { kv_max ); let batch_actual = seq_len.min(n_tok.max(1)); - eprintln!("batching: gemma4::lowered::forward_prefill_batch → step program → sandwich gemm_rows tap, batch={} KV [seq_len+16]", batch_actual); + eprintln!("batching: gemma4::forward::forward_prefill_batch → step program → sandwich gemm_rows tap, batch={} KV [seq_len+16]", batch_actual); let t0 = std::time::Instant::now(); let grouped = lpp < cfg.n_layers; let calib_prefix_grouped = calib_prefix.clone(); @@ -1686,14 +1674,14 @@ fn main() { collect_grouped(&mut gpu, arch, cfg.n_layers, lpp, Vec::new(), Path::new(&output), &[("source_model", serde_json::json!(model.clone())), ("corpus", serde_json::json!(corpus.clone())), ("corpus_md5", serde_json::json!(md5.clone())), ("n_calib_tokens", serde_json::json!(n_tok)), ("source_arch_id", serde_json::json!(arch)), ("seq_len", serde_json::json!(seq_len)), ("batch_size", serde_json::json!(batch_actual)), ("syrk_chosen", serde_json::json!(syrk_chosen)), ("syrk_mode", serde_json::json!(syrk_mode.clone())), ("layers_per_pass", serde_json::json!(lpp)), ("batches", serde_json::json!(n_seqs)), ("kv_max", serde_json::json!(kv_max)), ("calib_driver", serde_json::json!("calib_sweep gemma4 forward_prefill_batch (step program, sandwich projection tap)"))], |s,e| build_capture_gemma(&weights,s,e, &calib_prefix_grouped), - |gpu,_| { run_gemma_batched(gpu,&weights,&cfg,&tc,seq_len,&mut kv_sliding,&mut kv_full,&scratch)?; Ok(CalibForward::default()) } + |gpu,_| { run_gemma_batched(gpu,&weights,&cfg,&tc,seq_len,&mut state)?; Ok(CalibForward::default()) } ).unwrap_or_else(|e| { eprintln!("collect_grouped gemma: {e}"); std::process::exit(1); }) } else { let cap = build_capture_gemma(&weights, 0, cfg.n_layers, &calib_prefix_single); let tc = toks.clone(); collect(&mut gpu, arch, cap, Vec::new(), Path::new(&output), &[("source_model", serde_json::json!(model.clone())), ("corpus", serde_json::json!(corpus.clone())), ("corpus_md5", serde_json::json!(md5.clone())), ("n_calib_tokens", serde_json::json!(n_tok)), ("source_arch_id", serde_json::json!(arch)), ("seq_len", serde_json::json!(seq_len)), ("batch_size", serde_json::json!(batch_actual)), ("syrk_chosen", serde_json::json!(syrk_chosen)), ("syrk_mode", serde_json::json!(syrk_mode.clone())), ("layers_per_pass", serde_json::json!(cfg.n_layers)), ("batches", serde_json::json!(n_seqs)), ("kv_max", serde_json::json!(kv_max)), ("calib_driver", serde_json::json!("calib_sweep gemma4 forward_prefill_batch (step program, sandwich projection tap)"))], - |gpu| { run_gemma_batched(gpu,&weights,&cfg,&tc,seq_len,&mut kv_sliding,&mut kv_full,&scratch)?; Ok(CalibForward::default()) } + |gpu| { run_gemma_batched(gpu,&weights,&cfg,&tc,seq_len,&mut state)?; Ok(CalibForward::default()) } ).unwrap_or_else(|e| { eprintln!("collect gemma: {e}"); std::process::exit(1); }) }; // prevent weights moved in closure from being dropped before summary diff --git a/crates/hipfire-runtime/examples/debug_gemma4_attention.rs b/crates/hipfire-runtime/examples/debug_gemma4_attention.rs index d855a8037f..90a4c4172d 100644 --- a/crates/hipfire-runtime/examples/debug_gemma4_attention.rs +++ b/crates/hipfire-runtime/examples/debug_gemma4_attention.rs @@ -2,68 +2,27 @@ // Copyright (c) 2026 Kevin Read // hipfire — see LICENSE and NOTICE in the project root. -use hipfire_arch_gemma4::lowered::{self as gemma4, Gemma4Scratch}; +use hipfire_arch_gemma4::gemma4::FullKvTier; +use hipfire_arch_gemma4::{forward, Gemma4Config, Gemma4State, Gemma4Weights}; use hipfire_runtime::hfq::HfqFile; -use hipfire_runtime::llama; use rdna_compute::Gpu; use std::path::Path; fn main() { let path = Path::new("/local/models/google/gemma-4-12B-it.hfq"); let mut gpu = Gpu::init().expect("failed to init GPU"); - let mut hfq = HfqFile::open(path).expect("failed to open HFQ"); + let hfq = HfqFile::open(path).expect("failed to open HFQ"); - let config = gemma4::config_from_hfq(&hfq).expect("failed config"); - let weights = gemma4::load_weights(&mut hfq, &config, &mut gpu).expect("failed weights"); - let scratch = Gemma4Scratch::new(&mut gpu, &config, 2048).expect("failed scratch"); - gemma4::init_scratch_constants(&mut gpu, &scratch, config.full_head_dim) - .expect("failed init scratch"); - - // Allocate caches - let mut kv_sliding = llama::KvCache::new_gpu_asym3( - &mut gpu, - config.n_layers, - config.sliding_n_kv_heads, - config.sliding_head_dim, - config.sliding_window, - ) - .expect("failed sliding cache"); - - let mut kv_full = llama::KvCache::new_gpu_asym3( - &mut gpu, - config.n_layers, - config.full_n_kv_heads, - config.full_head_dim, - 2048, - ) - .expect("failed full cache"); - - // Set DUMP env var to 1 programmatically - std::env::set_var("HIPFIRE_GEMMA4_DUMP", "1"); + let config = Gemma4Config::from_hfq(&hfq).expect("failed config"); + let weights = Gemma4Weights::load(&hfq, &config, &mut gpu).expect("failed weights"); + // Sliding tier is always Q8; the full tier keeps the asym3 cache this probe used. + let mut state = Gemma4State::new_with_full_tier(&mut gpu, &config, 2048, FullKvTier::LegacyAsym3) + .expect("failed state"); println!("Running Layer 0 forward with BOS (2)..."); - gemma4::forward_scratch( - &mut gpu, - &weights, - &config, - 2, - 0, - &mut kv_sliding, - &mut kv_full, - &scratch, - ) - .expect("BOS forward failed"); + forward::decode_step(&config, &weights, &mut state, &mut gpu, 2, 0).expect("BOS forward failed"); println!("Running Layer 0 forward with 'Hello' (9259)..."); - gemma4::forward_scratch( - &mut gpu, - &weights, - &config, - 9259, - 1, - &mut kv_sliding, - &mut kv_full, - &scratch, - ) - .expect("Hello forward failed"); + forward::decode_step(&config, &weights, &mut state, &mut gpu, 9259, 1) + .expect("Hello forward failed"); } diff --git a/crates/hipfire-runtime/examples/eval_hipfire.rs b/crates/hipfire-runtime/examples/eval_hipfire.rs index 1012b625f3..b6cbaa1ec7 100644 --- a/crates/hipfire-runtime/examples/eval_hipfire.rs +++ b/crates/hipfire-runtime/examples/eval_hipfire.rs @@ -39,7 +39,7 @@ //! Multi-arch dispatch (Tier-1): mirrors `calib_sweep.rs` arch detection //! llama 0|1, qwen35 5|6, gemma4 13, glimmer 14, lfm2 11. //! Per-arch chunking mirrors calib_sweep: -//! gemma4 uses forward_prefill_batch (step program, 64-row chunks) with q8 KV +//! gemma4 uses forward::forward_prefill_batch (step program, 64-row chunks) with q8 KV //! glimmer chunks 192 (glimmer_prefill_chunk_size / prefill_with_capture) //! lfm2 uses forward_decode_batch_lfm with explicit positions, q8 KV //! qwen35 uses forward_prefill_batch with q8 KV at seq_len+16 @@ -953,17 +953,17 @@ fn main() { } } else if is_gemma { // ----- gemma4 13 — WIRED ----- - // Quant dtypes admitted: MQ4G256 / MG4G256 (campaign mq4, quant_type 13/30 alias), - // HFQ4G256, HFQ4G128, HFQ6G256, HFQ2/3, Q4K, MQ8G256, MQ6G256, MQ3G256, MQ2G256, - // plus F32/BF16 teachers (BF16 kept as BF16 when HIPFIRE_CALIB_BF16=1 else widened to F32). - // See lowered.rs:load_gemma4_weight. Batched prefill needs a Q8 KV cache. - use hipfire_arch_gemma4::lowered as gemma4; + // Quant dtypes admitted: MQ4G256(V2), HFQ4G256, HFQ4G128, HFQ6G256, HFQ2/3, + // Q4K, MQ8G256, MQ6G256, MQ3G256, MQ2G256, plus F32/BF16 teachers (BF16 + // kept as BF16 when HIPFIRE_CALIB_BF16=1 else widened to F32). See + // gemma4.rs:load_wt. Batched prefill needs a Q8 KV cache. + use hipfire_arch_gemma4::{forward, Gemma4Config, Gemma4State, Gemma4Weights}; if args.model.is_dir() { eprintln!("eval_hipfire: gemma4 dir not yet wired — use HFQ for arch 13"); std::process::exit(1); } - let mut hfq = HfqFile::open(&args.model).expect("open model"); - let cfg = gemma4::config_from_hfq(&hfq).unwrap_or_else(|| panic!("gemma4 config missing")); + let hfq = HfqFile::open(&args.model).expect("open model"); + let cfg = Gemma4Config::from_hfq(&hfq).expect("gemma4 config"); if ref_n_vocab != cfg.vocab_size { eprintln!( "vocab mismatch: ref says {ref_n_vocab}, model says {}", @@ -978,37 +978,19 @@ fn main() { .chunks_exact(4) .map(|b| u32::from_le_bytes(b.try_into().unwrap())) .collect(); - let mut weights = - gemma4::load_weights(&mut hfq, &cfg, &mut gpu).expect("gemma4 load weights"); + let weights = Gemma4Weights::load(&hfq, &cfg, &mut gpu).expect("gemma4 load weights"); let kv_max = n_ctx + 16; - let scratch = gemma4::Gemma4Scratch::new(&mut gpu, &cfg, kv_max).expect("gemma scratch"); - gemma4::init_scratch_constants(&mut gpu, &scratch, cfg.full_head_dim) - .expect("gemma init scratch"); // Q8 KV — mirrors calib_sweep: gemma sliding+full both q8 (lines 853-854) if args.kv_mode != "q8" { eprintln!("eval_hipfire: gemma4 requires --kv-mode q8 (got {}); F32 cache has no KvWriteF32 kernel — refusing", args.kv_mode); std::process::exit(1); } - let mut kv_sliding = KvCache::new_gpu_q8( - &mut gpu, - cfg.n_layers, - cfg.sliding_n_kv_heads, - cfg.sliding_head_dim, - kv_max, - ) - .expect("gemma sliding kv q8"); - let mut kv_full = KvCache::new_gpu_q8( - &mut gpu, - cfg.n_layers, - cfg.full_n_kv_heads, - cfg.full_head_dim, - kv_max, - ) - .expect("gemma full kv q8"); + let mut state = + Gemma4State::new_with_max_seq(&mut gpu, &cfg, kv_max).expect("gemma4 state"); for c in 0..effective_n_chunk { - kv_sliding.clear_gpu(&mut gpu).expect("kv sliding clear"); - kv_full.clear_gpu(&mut gpu).expect("kv full clear"); + state.kv_sliding.clear_gpu(&mut gpu).expect("kv sliding clear"); + state.kv_full.clear_gpu(&mut gpu).expect("kv full clear"); let chunk_tokens = &tokens[c * n_ctx..(c + 1) * n_ctx]; let mut chunk_klds: Vec = Vec::with_capacity(scored_per_chunk); let mut chunk_nll_sum: f64 = 0.0; @@ -1016,32 +998,27 @@ fn main() { if args.scoring_mode == "prefill" { // Batched prefill of the prefix, then per-token for the scored window. - gemma4::forward_prefill_batch( - &mut gpu, - &weights, + forward::forward_prefill_batch( &cfg, + &weights, + &mut state, + &mut gpu, &chunk_tokens[..scoring_start], 0, - &mut kv_sliding, - &mut kv_full, - &scratch, ) .expect("gemma prefill prefix"); for pos in scoring_start..(n_ctx - 1) { - gemma4::forward_scratch( - &mut gpu, - &weights, + let cand_logits = forward::decode_step( &cfg, + &weights, + &mut state, + &mut gpu, chunk_tokens[pos], - pos, - &mut kv_sliding, - &mut kv_full, - &scratch, + pos as u32, ) - .expect("gemma forward_scratch"); + .expect("gemma decode_step"); let actual_next = chunk_tokens[pos + 1] as usize; ref_in.read_exact(&mut block_buf).expect("read ref block"); - let cand_logits = gpu.download_f32(&scratch.logits).expect("download logits"); let (kld, nll) = kld_from_block(&cand_logits, &block_buf, top_k, actual_next); chunk_klds.push(kld); if let Some(n) = nll { @@ -1066,23 +1043,20 @@ fn main() { } } else { for pos in 0..(n_ctx - 1) { - gemma4::forward_scratch( - &mut gpu, - &weights, + let cand_logits = forward::decode_step( &cfg, + &weights, + &mut state, + &mut gpu, chunk_tokens[pos], - pos, - &mut kv_sliding, - &mut kv_full, - &scratch, + pos as u32, ) - .expect("gemma forward_scratch"); + .expect("gemma decode_step"); if pos < scoring_start { continue; } let actual_next = chunk_tokens[pos + 1] as usize; ref_in.read_exact(&mut block_buf).expect("read ref block"); - let cand_logits = gpu.download_f32(&scratch.logits).expect("download logits"); let (kld, nll) = kld_from_block(&cand_logits, &block_buf, top_k, actual_next); chunk_klds.push(kld); if let Some(n) = nll { diff --git a/crates/hipfire-runtime/examples/eval_hipfire_gemma4.rs b/crates/hipfire-runtime/examples/eval_hipfire_gemma4.rs index 68df27ee88..3bb159cfcd 100644 --- a/crates/hipfire-runtime/examples/eval_hipfire_gemma4.rs +++ b/crates/hipfire-runtime/examples/eval_hipfire_gemma4.rs @@ -7,11 +7,10 @@ //! `build_kld_ref_native_gemma4`). //! //! Mirrors `eval_hipfire` (qwen35) with gemma4-architectural swaps: -//! * forward is `hipfire_arch_gemma4::lowered::forward_scratch`, per-token -//! scoring mode only (no gemma4 batched-prefill capture path wired here). -//! * dual KV: sliding F32 `new_gpu` + full asym3 `new_gpu_asym3` — the -//! IDENTICAL config the ref builder used (gemma4_oracle's HF-parity -//! setup), so KV noise cancels and KLD isolates weight precision. +//! * forward is `hipfire_arch_gemma4::forward::decode_step`, per-token +//! scoring mode only. +//! * Q8 KV on both tiers — the IDENTICAL config the ref builder uses, so KV +//! noise cancels and KLD isolates weight precision. //! * additionally reports top-1 agreement (candidate argmax == ref top-1). //! //! Output: HFKSEQ v2 — per-sequence (mean, p99, mean_nll) fp64 triples. @@ -21,9 +20,8 @@ //! --output [--max-chunks N] fn main() { - use hipfire_arch_gemma4::lowered::{self as gemma4, Gemma4Scratch}; + use hipfire_arch_gemma4::{forward, Gemma4Config, Gemma4State, Gemma4Weights}; use hipfire_runtime::hfq::HfqFile; - use hipfire_runtime::llama::KvCache; use std::fs::File; use std::io::{BufReader, BufWriter, Read, Write}; use std::path::PathBuf; @@ -83,9 +81,9 @@ fn main() { gpu.arch, model.display() ); - let mut hfq = HfqFile::open(&model).expect("open model"); - let config = gemma4::config_from_hfq(&hfq).expect("read config"); - let weights = gemma4::load_weights(&mut hfq, &config, &mut gpu).expect("load weights"); + let hfq = HfqFile::open(&model).expect("open model"); + let config = Gemma4Config::from_hfq(&hfq).expect("read config"); + let weights = Gemma4Weights::load(&hfq, &config, &mut gpu).expect("load weights"); // -------- read reference (HFKLDR β) header + tokens -------- let ref_file = File::open(&ref_path).expect("open ref"); @@ -140,31 +138,10 @@ fn main() { .map(|b| u32::from_le_bytes(b.try_into().unwrap())) .collect(); - // -------- scratch + dual KV (identical to ref builder) -------- + // -------- state: Q8 KV on both tiers (the unified state has no F32 KV + // tier; the forward refuses F32 cache writes) -------- let kv_max = n_ctx + 16; - let scratch = Gemma4Scratch::new(&mut gpu, &config, kv_max).expect("scratch"); - gemma4::init_scratch_constants(&mut gpu, &scratch, config.full_head_dim) - .expect("init_scratch_constants"); - let mut kv_sliding = KvCache::new_gpu( - &mut gpu, - config.n_layers, - config.sliding_n_kv_heads, - config.sliding_head_dim, - kv_max, - ) - .expect("kv sliding alloc"); - // FULL KV = F32, NOT asym3: on gfx942/CDNA the asym3 full-KV path is - // catastrophically wrong (grows with depth; PPL 3826 at 512 ctx) while - // F32 full-KV is HF-EXACT (top-5 logits match HF to 1e-4 at 128 ids, - // 2026-06-10). F32 both sides also removes the shared KV-noise floor. - let mut kv_full = KvCache::new_gpu( - &mut gpu, - config.n_layers, - config.full_n_kv_heads, - config.full_head_dim, - kv_max, - ) - .expect("kv full alloc"); + let mut state = Gemma4State::new_with_max_seq(&mut gpu, &config, kv_max).expect("state"); // -------- per-chunk loop -------- let mut mean_kld_per_seq: Vec = Vec::with_capacity(effective_n_chunk); @@ -185,17 +162,15 @@ fn main() { let mut chunk_nll_count: usize = 0; for pos in 0..(n_ctx - 1) { - gemma4::forward_scratch( - &mut gpu, - &weights, + let cand_logits = forward::decode_step( &config, + &weights, + &mut state, + &mut gpu, chunk_tokens[pos], - pos, - &mut kv_sliding, - &mut kv_full, - &scratch, + pos as u32, ) - .expect("forward_scratch"); + .expect("decode_step"); if pos < scoring_start { continue; } @@ -221,7 +196,6 @@ fn main() { let sum_p_residual = f32::from_le_bytes(block_buf[resid_off..resid_off + 4].try_into().unwrap()); - let cand_logits = gpu.download_f32(&scratch.logits).expect("download logits"); let cand_logits = &cand_logits[..config.vocab_size.min(cand_logits.len())]; // Candidate log-Z (fp64) + argmax for top-1 agreement. diff --git a/crates/hipfire-runtime/examples/gemma4_oracle.rs b/crates/hipfire-runtime/examples/gemma4_oracle.rs index e29424054e..a397079cd4 100644 --- a/crates/hipfire-runtime/examples/gemma4_oracle.rs +++ b/crates/hipfire-runtime/examples/gemma4_oracle.rs @@ -3,29 +3,20 @@ // hipfire — see LICENSE and NOTICE in the project root. //! Gemma-4 logit oracle. //! -//! Runs `gemma4::forward_scratch` tokenwise over a token-id list and dumps the +//! Runs `forward::decode_step` tokenwise over a token-id list and dumps the //! final-position top-k logits as JSON. Token IDs are read from a file so the //! HF reference (`scripts/oracle_gemma4.py`) and hipfire compare BYTE-IDENTICAL //! inputs (no tokenizer differences enter). //! -//! KV geometry mirrors the production lowered bundle (`load_gemma4_bundle`): -//! sliding Q8 ring over the actual sliding layers with physical capacity -//! `min(sliding_window, max_seq)`, full asym3 over the actual full layers with -//! capacity `max_seq`, scratch sized from the same `max_seq`. This example -//! still bypasses the production admission/control plane — it calls -//! `forward_scratch` directly with no daemon sliding-window guard. +//! The state mirrors the production bundle: Q8 sliding tier and an asym3 full +//! tier (Q8 with `HIPFIRE_ORACLE_KV_Q8=1`), both position-indexed over +//! `max_seq`. This example bypasses the production admission/control plane. //! //! Usage: gemma4_oracle [out.json] -//! -//! Common short-token case (fewer than 128 ids, e.g. the 15-token parity -//! prompt): `max_seq` floors at 128 (one flash tile) and both caches cover -//! every requested position, so no ring wrap occurs. Rollover case: only when -//! the id count exceeds the sliding physical capacity (`min(sliding_window, -//! max_seq)`) do later positions wrap the Q8 ring; shorter runs never wrap. -use hipfire_arch_gemma4::lowered::{self as gemma4, Gemma4Scratch}; +use hipfire_arch_gemma4::gemma4::FullKvTier; +use hipfire_arch_gemma4::{forward, Gemma4Config, Gemma4State, Gemma4Weights}; use hipfire_runtime::hfq::HfqFile; -use hipfire_runtime::llama::KvCache; use std::path::Path; fn main() { @@ -44,78 +35,25 @@ fn main() { assert!(!ids.is_empty(), "no ids"); let mut gpu = rdna_compute::Gpu::init().expect("gpu init"); - let mut hfq = HfqFile::open(model_path).expect("open model"); - let config = gemma4::config_from_hfq(&hfq).expect("config_from_hfq"); - let weights = gemma4::load_weights(&mut hfq, &config, &mut gpu).expect("load_weights"); - // Logical authority: cover every requested position, floored at one flash - // tile (128) so the scratch partials stay scratch-compatible. + let hfq = HfqFile::open(model_path).expect("open model"); + let config = Gemma4Config::from_hfq(&hfq).expect("config"); + let weights = Gemma4Weights::load(&hfq, &config, &mut gpu).expect("load weights"); let max_seq = ids.len().max(128); - // Actual per-type layer counts, mirroring the production bundle — NOT - // config.n_layers for both caches. - let (n_sliding_layers, n_full_layers) = - config - .layer_types - .iter() - .fold((0, 0), |(sliding, full), layer_type| match layer_type { - gemma4::LayerType::Sliding => (sliding + 1, full), - gemma4::LayerType::Full => (sliding, full + 1), - }); - // Bounded sliding physical capacity: the ring holds at most - // min(sliding_window, max_seq) slots (production passes sliding_window - // against a much larger daemon max_seq; here max_seq can be smaller). - let sliding_cap = config.sliding_window.min(max_seq); - eprintln!( - "oracle: max_seq={max_seq} sliding=q8-ring({n_sliding_layers} layers,cap={sliding_cap}) full=asym3({n_full_layers} layers,cap={max_seq}) wrap={}", - ids.len() > sliding_cap - ); - let mut scratch = Gemma4Scratch::new(&mut gpu, &config, max_seq).expect("scratch"); - gemma4::init_scratch_constants(&mut gpu, &scratch, config.full_head_dim) - .expect("init_scratch_constants"); - let mut kv_sliding = KvCache::new_gpu_q8_capped( - &mut gpu, - n_sliding_layers, - config.sliding_n_kv_heads, - config.sliding_head_dim, - max_seq, - sliding_cap, - ) - .expect("kv sliding alloc"); - let mut kv_full = if std::env::var("HIPFIRE_ORACLE_KV_F32").ok().as_deref() == Some("1") { - eprintln!("oracle: FULL KV = F32 (HIPFIRE_ORACLE_KV_F32=1)"); - KvCache::new_gpu( - &mut gpu, - n_full_layers, - config.full_n_kv_heads, - config.full_head_dim, - max_seq, - ) - .expect("kv full alloc f32") + // The unified state has no F32 KV tier; Q8 is the closest exact-write tier. + let tier = if std::env::var("HIPFIRE_ORACLE_KV_Q8").ok().as_deref() == Some("1") { + FullKvTier::Q8 } else { - KvCache::new_gpu_asym3_gemma4( - &mut gpu, - n_full_layers, - config.full_n_kv_heads, - config.full_head_dim, - max_seq, - ) - .expect("kv full alloc") + FullKvTier::LegacyAsym3 }; + eprintln!("oracle: max_seq={max_seq} sliding=q8 full={tier:?}"); + let mut state = + Gemma4State::new_with_full_tier(&mut gpu, &config, max_seq, tier).expect("state"); + let mut logits = Vec::new(); for (i, &tok) in ids.iter().enumerate() { - gemma4::forward_scratch( - &mut gpu, - &weights, - &config, - tok, - i, - &mut kv_sliding, - &mut kv_full, - &mut scratch, - ) - .unwrap_or_else(|e| panic!("forward_scratch at pos {i}: {e:?}")); + logits = forward::decode_step(&config, &weights, &mut state, &mut gpu, tok, i as u32) + .unwrap_or_else(|e| panic!("decode_step at pos {i}: {e:?}")); } - - let logits = gpu.download_f32(&scratch.logits).expect("download logits"); let vocab = config.vocab_size.min(logits.len()); let mut idx: Vec = (0..vocab).collect(); idx.sort_by(|&a, &b| logits[b].partial_cmp(&logits[a]).unwrap()); diff --git a/crates/hipfire-runtime/map.md b/crates/hipfire-runtime/map.md index 7e830a2698..83c2732124 100644 --- a/crates/hipfire-runtime/map.md +++ b/crates/hipfire-runtime/map.md @@ -58,7 +58,7 @@ _Generated by `scripts/check-crate-maps.py` from the tree — do not edit inside - [`src/imagedec.rs`](src/imagedec.rs): `decode_dynamic`, `decode_dynamic_path`, `decode_rgb8`, `decode_rgb8_path`, `is_jpeg`, `probe_dimensions`, `probe_dimensions_path` - [`src/kv_adaptive.rs`](src/kv_adaptive.rs): `KMode`, `KvAdaptive`, `Preset`, `Step`, `bits`, `bytes_per_head`, `cap_min`, `configure_eviction_handoff`, `current_cap`, `eviction_gate`, `from_preset`, `handoff_at`, `handoff_complete`, `is_poisoned`, `k_buf_bytes_per_layer`, `maybe_downshift`, `new`, `poison`, `poison_reason`, `reset`, `reset_with_cache`, `rot_width`, `v_buf_bytes_per_layer`, `v_bytes_per_head` - [`src/kv_backend.rs`](src/kv_backend.rs): `saddle_core` -- [`src/kv_mode.rs`](src/kv_mode.rs): `DEPRECATED_KV_NAMES`, `DIR_SAFETENSORS_POLICY`, `GEMMA4_LOWERED_FULL_POLICY`, `HFQ_Q8_ONLY_POLICY`, `KvModePolicy`, `KvPair`, `KvPairError`, `LLAMA_HFQ_POLICY`, `MAPLE_POLICY`, `QWEN35_HFQ_POLICY`, `QWEN35_PARO_POLICY`, `QWEN35_PP_POLICY`, `QWEN35_SLOTS_POLICY`, `QWEN35_TP_POLICY`, `QWEN4_POLICY`, `ResolveResult`, `deprecated_kv_name_warning`, `k`, `parse_qwen_k_name`, `parse_qwen_v_name`, `qwen35_native_eligible`, `qwen35_policy_for_native`, `qwen_k_display_name`, `qwen_v_display_name`, `resolve`, `resolve_kv_pair`, `resolve_qwen4`, `saddle_core`, `v` +- [`src/kv_mode.rs`](src/kv_mode.rs): `DEPRECATED_KV_NAMES`, `DIR_SAFETENSORS_POLICY`, `GEMMA4_FULL_POLICY`, `HFQ_Q8_ONLY_POLICY`, `KvModePolicy`, `KvPair`, `KvPairError`, `LLAMA_HFQ_POLICY`, `MAPLE_POLICY`, `QWEN35_HFQ_POLICY`, `QWEN35_PARO_POLICY`, `QWEN35_PP_POLICY`, `QWEN35_SLOTS_POLICY`, `QWEN35_TP_POLICY`, `QWEN4_POLICY`, `ResolveResult`, `deprecated_kv_name_warning`, `k`, `parse_qwen_k_name`, `parse_qwen_v_name`, `qwen35_native_eligible`, `qwen35_policy_for_native`, `qwen_k_display_name`, `qwen_v_display_name`, `resolve`, `resolve_kv_pair`, `resolve_qwen4`, `saddle_core`, `v` - [`src/lib.rs`](src/lib.rs): `admission`, `arch`, `arch_mapping`, `arch_model`, `arch_spec`, `augmentor`, `batch_sampler`, `bf16_loader`, `cache_plan`, `calibration`, `cask`, `chatml`, `checkpoint_pool`, `config`, `cpu_router`, `crate`, `ddtree`, `device_mesh`, `dflash`, `dflash_adaptive_block`, `dflash_generic`, `dspark_block_controller`, `dspark_core`, `emit_text`, `eos_filter`, `ep`, `eval_common`, `external_rows`, `gguf`, `hfq`, `hfq_parallel`, `imagedec`, `kv_adaptive`, `kv_backend`, `kv_mode`, `llama`, `llama_spec`, `lloyd_lut`, `loader_api`, `loop_guard`, `model_load`, `model_source`, `multi_gpu`, `ngram_mod`, `paro`, `ple_convert`, `ple_hot_pack`, `ple_hot_store`, `ple_prewarm`, `ple_q4`, `ple_sidecar`, `prefix`, `prefix_index`, `prompt_frame`, `reset_core`, `safetensors_source`, `sampler`, `scheduler`, `sealed_moe`, `semantic`, `serve`, `serve_contract`, `serve_fairness`, `serve_wait`, `session_cache`, `session_table`, `sidecar`, `slot_batch`, `spec`, `spec_ngram`, `spec_sampling`, `stop_sequence`, `swap`, `tokenizer`, `tool_call`, `tp_shard`, `triattn`, `weight_backend`, `weight_host`, `weight_manifest`, `weight_pager`, `weight_store` - [`src/llama.rs`](src/llama.rs): `EmbeddingFormat`, `ForwardScratch`, `HiddenCaptureSink`, `KvCacheExt`, `KvTarget`, `LayerWeights`, `LlamaConfig`, `LlamaWeights`, `ModelArch`, `PREFILL_MAX_BATCH`, `ParoRotation`, `PrefillBatchScratch`, `SamplingConfig`, `TreeMaskRef`, `WeightTensor`, `apply_ngram_block`, `apply_repeat_penalty`, `apply_repeat_penalty_candidates`, `apply_rope_cpu_pub`, `apply_special_token_attractor_block`, `apply_unclosed_attractor_block`, `ar_candidate_cap`, `argmax`, `attention_family`, `attention_flash_mode`, `attention_flash_mode_for`, `convert_q4k_to_q4f16_g32`, `convert_q4k_to_q4f16_g64`, `dequantize_q4_0`, `dequantize_q4_k`, `dequantize_q6_k`, `dequantize_q8_0`, `dispatch_ref`, `embedding_lookup_dispatch`, `f16_to_f32`, `f32_to_f16`, `forward`, `forward_early_exit`, `forward_prefill_batch`, `forward_prefill_batch_capture`, `forward_prefill_batch_chunk_captured`, `forward_prefill_batch_tree`, `forward_sample`, `forward_scratch`, `forward_scratch_compute`, `forward_scratch_compute_capture`, `forward_scratch_embed`, `forward_scratch_layers`, `free_all`, `free_gpu`, `free_metadata_only`, `from_gguf`, `fused_qkv_family`, `fused_rmsnorm_rotate_for_mq`, `fused_rmsnorm_rotate_mq_batched_for`, `fused_rmsnorm_rotate_mq_f16_batched_for`, `fused_silu_mul_rotate_mq_batched_for`, `fused_silu_mul_rotate_mq_f16_batched_for`, `fused_silu_mul_rotate_mq_for`, `gemm_family`, `gemv_family`, `hipfire_dispatch`, `is_batchable_la`, `load_weights`, `moe_family`, `mqv2_gfx11_wmma_enabled_from_env`, `mqv2_wmma_batchable`, `new`, `new_with_max_seq`, `prefill_forward`, `reset_cpu_sampler_rng`, `rotate_x_for_mq`, `rotate_x_mq_128_for`, `rotate_x_mq_batched_for`, `rotate_x_mq_for`, `rotate_x_paro_for`, `saddle_core`, `sample_full_dist`, `sample_top_k_p`, `sample_top_p`, `sample_top_p_from_candidates`, `sampler_rng_restore`, `sampler_rng_snapshot`, `simple`, `text_thinking`, `upload_prefill_batch_inputs`, `vl_thinking`, `weight_gemm`, `weight_gemv`, `weight_gemv_prerotated`, `weight_gemv_residual`, `weight_gemv_swiglu_residual` - [`src/llama_spec.rs`](src/llama_spec.rs): `lm_head_logits_n_rows`, `verify_block_argmax`, `verify_block_argmax_capture_gpu`, `verify_block_logits`, `verify_block_sampled_capture_gpu`, `verify_tree_logits` @@ -89,7 +89,7 @@ _Generated by `scripts/check-crate-maps.py` from the tree — do not edit inside - [`src/serve_contract.rs`](src/serve_contract.rs): `ArchPolicy`, `CacheDomain`, `CanonicalError`, `CheckpointId`, `CommitBoundary`, `DeviceTopology`, `Digest`, `DrafterDecision`, `KvLayout`, `LastTokenHandling`, `MissReason`, `NONE`, `PrefixLookup`, `PrefixLookupResult`, `PublishLease`, `ReleaseDisposition`, `ReservationError`, `ResumeBundle`, `ResumePlan`, `ResumePlanError`, `SharingNamespace`, `StepNeeds`, `StepReservation`, `StepTicket`, `TemplateIdentity`, `TokenizerIdentity`, `fits`, `for_model`, `from_canonical_bytes`, `is_none`, `is_some`, `new`, `scoped`, `sha256_file`, `sha256_len_prefixed`, `to_canonical_bytes`, `total_rows` - [`src/serve_fairness.rs`](src/serve_fairness.rs): `FairQueue`, `FairRequest`, `FairnessError`, `Grant`, `Selection`, `admit`, `get`, `id`, `is_empty`, `len`, `new`, `remove`, `rows`, `select`, `set_needs`, `set_uncached_prefill`, `skip_round_complete`, `tick` - [`src/serve_wait.rs`](src/serve_wait.rs): `WaitError`, `WaitQueue`, `Waiter`, `deadline`, `expire`, `max_bytes`, `max_count`, `new`, `pop_ready`, `queued_bytes`, `queued_count`, `remove`, `timeout_ticks`, `try_enqueue` -- [`src/session_cache.rs`](src/session_cache.rs): `RowStream`, `SessionCache`, `SessionRoute`, `SessionState`, `SnapshotParts`, `StateLayout`, `StatePart`, `at_boundary`, `begin`, `begin_live`, `clear`, `commit`, `new`, `next_boundary`, `plan`, `release_for`, `stored_bytes` +- [`src/session_cache.rs`](src/session_cache.rs): `RowStream`, `SessionCache`, `SessionDriver`, `SessionRoute`, `SessionState`, `SnapshotParts`, `StateLayout`, `StatePart`, `at_boundary`, `begin`, `begin_live`, `cache`, `clear`, `commit`, `commit_live`, `forget_live`, `new`, `next_boundary`, `plan`, `release_for`, `stored_bytes`, `stride_boundaries` - [`src/session_table.rs`](src/session_table.rs): `Residency`, `Session`, `SessionId`, `SessionTable`, `active`, `begin_turn`, `close`, `confirm_reentry`, `find_continuation`, `get`, `get_mut`, `is_admitted`, `iter`, `lru_idle_victim`, `mark_cold`, `mark_resident`, `mark_swapped`, `open`, `readmit`, `resident`, `resize_grant`, `touch` - [`src/sidecar.rs`](src/sidecar.rs): `SidecarExt`, `find_vl_sidecar`, `resolve_vl_sidecar`, `sidecar_candidates`, `vl_sidecar_candidates` - [`src/slot_batch.rs`](src/slot_batch.rs): `BatchPlanner`, `BatchStepPlan`, `RequestAdvance`, `RequestEpoch`, `RequestRows`, `RequestStepKind`, `RowRange`, `SlotBatch`, `StepOutput`, `build`, `discard`, `fn`, `is_empty`, `new`, `plan_step`, `publish`, `total_rows` diff --git a/crates/hipfire-runtime/src/kv_mode.rs b/crates/hipfire-runtime/src/kv_mode.rs index a21426e5ca..e955ea5909 100644 --- a/crates/hipfire-runtime/src/kv_mode.rs +++ b/crates/hipfire-runtime/src/kv_mode.rs @@ -412,7 +412,7 @@ pub const HFQ_Q8_ONLY_POLICY: KvModePolicy = KvModePolicy { default: Q8, }; -/// Gemma4 lowered full-attention tier. Names use the shared Qwen table, so +/// Gemma4 full-attention tier. Names use the shared Qwen table, so /// bare `asym3`/`turbo3` mean FWHT (no hd512 FWHT tier here → q8 + warn) and /// the legacy Givens tier needs the explicit `legacy-asym3` opt-out. /// `fp8`/`bf16` are carried forward so the carrier fails closed on them. @@ -425,14 +425,14 @@ fn normalize_gemma4(raw: &str) -> Option { } } -/// Gemma4 lowered carrier, full-attention tier (the sliding tier is -/// always the q8 ring). `auto` follows the shared policy: fp8 only where the -/// site admits Fp8, else q8. Fp8 is not admitted: every lowered Gemma4 attend -/// site builds its tier plan with `quant_fp8: false`, and the batched full -/// prefill uses the windowed Q8 plan, which refuses fp8. So `auto` is q8 on -/// every arch; `legacy-asym3` restores the previous Givens asym3 tier. -pub const GEMMA4_LOWERED_FULL_POLICY: KvModePolicy = KvModePolicy { - site: "gemma4-lowered-full", +/// Gemma4 carrier, full-attention tier (the sliding tier is always q8). +/// `auto` follows the shared policy: fp8 only where the site admits Fp8, else +/// q8. Fp8 is not admitted: every Gemma4 attend site builds its tier plan +/// with `quant_fp8: false`, and the batched prefill uses the Q8 plan, which +/// refuses fp8. So `auto` is q8 on every arch; `legacy-asym3` restores the +/// previous Givens asym3 tier. +pub const GEMMA4_FULL_POLICY: KvModePolicy = KvModePolicy { + site: "gemma4-full", normalize_alias: normalize_gemma4, accepted: &[Q8, Asym3], default: Q8, @@ -1272,8 +1272,8 @@ mod tests { } #[test] - fn gemma4_lowered_full_policy() { - let p = &GEMMA4_LOWERED_FULL_POLICY; + fn gemma4_full_policy() { + let p = &GEMMA4_FULL_POLICY; for raw in ["", "auto", "q8"] { let r = resolve(raw, p); assert_eq!(r.mode, Q8, "{raw:?}"); diff --git a/crates/hipfire-runtime/src/session_cache.rs b/crates/hipfire-runtime/src/session_cache.rs index 32c9267d3e..7f4b9c90b9 100644 --- a/crates/hipfire-runtime/src/session_cache.rs +++ b/crates/hipfire-runtime/src/session_cache.rs @@ -19,6 +19,12 @@ //! [`SessionCache::begin_live`], which arms boundaries only when asked to //! capture. //! +//! [`SessionDriver`] is what an architecture's bundle holds: the cache plus +//! the committed live state of the current conversation. It plans the longer +//! of the live state and a cached snapshot, starts each prefill as a live +//! continuation, a restore or a cold start, and keeps the live state after a +//! committed turn whose consumed history extends the turn's prompt. +//! //! Snapshots are deltas: one stores the fixed (overwritten-in-place) state //! whole, but of each append-only [`RowStream`] only the rows above its //! parent, the deepest snapshot of the same prefix that was present when it @@ -116,10 +122,14 @@ pub trait SessionState { route: SessionRoute, meta: &[u8], ) -> Result<(), String>; + /// The position every owner of `route` sits at when the live state can + /// continue in place; `None` when it cannot (an owner lags, or verify + /// state is armed). + fn live_position(&self, route: SessionRoute) -> Option; /// Device bytes the live state may still map or allocate before it - /// reaches the end of the prompt in flight (the cache leaves them free: - /// it never takes memory the context needs). - fn growth_reserve_bytes(&self, gpu: &Gpu) -> u64; + /// holds `through` tokens, the end of the prompt in flight (the cache + /// leaves them free: it never takes memory the context needs). + fn growth_reserve_bytes(&self, gpu: &Gpu, through: usize) -> u64; /// Cold-start the live state. fn reset(&mut self, gpu: &mut Gpu) -> Result<(), String>; } @@ -182,6 +192,8 @@ struct Turn { domain: CacheDomain, route: SessionRoute, boundaries: VecDeque, + /// Length of the prompt in flight. + end: usize, } pub struct SessionCache { @@ -346,6 +358,7 @@ impl SessionCache { .into(), domain, route, + end: prompt.len(), }); Ok(()) } @@ -394,6 +407,7 @@ impl SessionCache { .into(), domain: self.domain.scoped(&scope), route, + end: prompt.len(), }); } @@ -514,12 +528,12 @@ impl SessionCache { )); } turn.boundaries.pop_front(); - let (domain, route, p) = (turn.domain.clone(), turn.route, prefix.len()); + let (domain, route, p, end) = (turn.domain.clone(), turn.route, prefix.len(), turn.end); let fp = prefix_fingerprint(prefix); if self.pool.contains(&domain, p as u64, fp) { return Ok(()); } - let growth = state.growth_reserve_bytes(gpu); + let growth = state.growth_reserve_bytes(gpu, end); let ancestors: Vec = state .snapshot_boundaries(prefix, 0, p - 1) .into_iter() @@ -738,6 +752,216 @@ impl SessionCache { } } +/// Ascending snapshot boundaries in `(after, up_to]`: multiples of `stride`, +/// plus the position after every `turn_end` token of `prompt` (message ends) +/// when it is `Some`. +pub fn stride_boundaries( + stride: usize, + prompt: &[u32], + turn_end: Option, + after: usize, + up_to: usize, +) -> Vec { + let mut boundaries: Vec = (after / stride + 1..=up_to / stride) + .map(|k| k * stride) + .collect(); + if let Some(turn_end) = turn_end { + let end = up_to.min(prompt.len()); + boundaries.extend( + (after.min(end)..end) + .filter(|&i| prompt[i] == turn_end) + .map(|i| i + 1), + ); + boundaries.sort_unstable(); + boundaries.dedup(); + } + boundaries +} + +/// Host record of the committed live state: the consumed history it holds and +/// the route that produced it. A live continuation is *session-exact* (its +/// state descends from decode, as in a continuous conversation), not +/// cold-exact, so it is only resumed in place and never captured into the +/// cross-session cache, except by boundaries the architecture scopes apart. +#[derive(Clone, Debug, PartialEq, Eq)] +struct LiveRecord { + tokens: Vec, + route: SessionRoute, +} + +/// `Some(L)` when the live state in `record` can serve `prompt` on `route`: +/// it holds exactly `L = record.tokens.len()` tokens (`0 < L < prompt.len()`, +/// so at least the last prompt token is computed), those tokens are +/// `prompt[..L]`, and every owner sits at `L`. +fn live_start( + record: Option<&LiveRecord>, + prompt: &[u32], + route: SessionRoute, + live_position: Option, +) -> Option { + let record = record?; + let end = record.tokens.len(); + (record.route == route + && end > 0 + && end < prompt.len() + && prompt[..end] == record.tokens[..] + && live_position == Some(end)) + .then_some(end) +} + +/// Whether the state left by the turn that began with `turn_prompt` is the +/// live state after `consumed`, the full host history the client committed: +/// `consumed` extends the turn's prompt and every owner sits at its end. An +/// empty `consumed` (unrepaired terminal) is never live. +fn live_after_commit(turn_prompt: &[u32], consumed: &[u32], live_position: Option) -> bool { + let prompt_len = turn_prompt.len(); + prompt_len > 0 + && consumed.len() >= prompt_len + && consumed[..prompt_len] == *turn_prompt + && live_position == Some(consumed.len()) +} + +/// Prompt tokens the next prefill skips: the live state when it reaches at +/// least as far as the longest cached snapshot (a tie prefers live: no copy), +/// else that snapshot. +fn plan_start(snapshot: usize, live: Option) -> usize { + match live { + Some(live) if live >= snapshot => live, + _ => snapshot, + } +} + +/// The session machinery an architecture's bundle holds: the shared snapshot +/// cache plus the committed live state of the current conversation. +pub struct SessionDriver { + cache: SessionCache, + /// Committed live state kept for a continuation of the same conversation; + /// consumed while a request is in flight. + live: Option, + /// Prompt and route of the request in flight. + turn: Option<(Vec, SessionRoute)>, +} + +impl SessionDriver { + pub fn new(cache: SessionCache) -> Self { + Self { + cache, + live: None, + turn: None, + } + } + + pub fn cache(&self) -> &SessionCache { + &self.cache + } + + /// [`SessionCache::release_for`]. + pub fn release_for(&mut self, gpu: &mut Gpu, need: u64) -> u64 { + self.cache.release_for(gpu, need) + } + + fn live_hit( + &self, + state: &dyn SessionState, + prompt: &[u32], + route: SessionRoute, + ) -> Option { + live_start( + self.live.as_ref(), + prompt, + route, + state.live_position(route), + ) + } + + /// Prompt tokens the next prefill of `prompt` on `route` can skip: the + /// longer of the live state and a cached snapshot; 0 on a miss. + pub fn plan(&self, state: &dyn SessionState, prompt: &[u32], route: SessionRoute) -> usize { + plan_start( + self.cache.plan(state, prompt, route), + self.live_hit(state, prompt, route), + ) + } + + /// Start a prefill of `prompt` on `route` that skips the `reused` tokens + /// [`Self::plan`] returned: + /// - live: `reused` is the end of the committed live state; every owner + /// is kept, nothing is restored, and boundaries above `reused` are + /// captured only with `capture_live`; + /// - snapshot: restore the `reused`-token snapshot; + /// - `reused == 0`: cold start. + /// + /// Returns whether the prefill continues the live state. + pub fn begin( + &mut self, + gpu: &mut Gpu, + state: &mut dyn SessionState, + prompt: &[u32], + route: SessionRoute, + reused: usize, + capture_live: bool, + ) -> Result { + let live = reused > 0 && self.live_hit(state, prompt, route) == Some(reused); + self.live = None; + self.turn = None; + if live { + self.cache + .begin_live(gpu, state, prompt, route, reused, capture_live); + } else { + self.cache.begin(gpu, state, prompt, route, reused)?; + } + self.turn = Some((prompt.to_vec(), route)); + Ok(live) + } + + /// The next position the running prefill must stop at and report. + pub fn next_boundary(&self) -> Option { + self.cache.next_boundary() + } + + /// Report that every owner consumed exactly `prefix` (the boundary). + pub fn at_boundary( + &mut self, + gpu: &mut Gpu, + state: &mut dyn SessionState, + prefix: &[u32], + ) -> Result<(), String> { + self.cache.at_boundary(gpu, state, prefix) + } + + /// Publish this turn's snapshots and drop the live state. + pub fn commit(&mut self) { + self.cache.commit(); + self.forget_live(); + } + + /// Publish this turn's snapshots and keep the live state when `consumed` + /// (the host history the client committed) extends this turn's prompt and + /// every owner sits exactly at its end. The next turn of the conversation + /// then prefills only the suffix in place. + pub fn commit_live(&mut self, state: &dyn SessionState, consumed: &[u32]) { + self.cache.commit(); + self.live = self.turn.take().and_then(|(prompt, route)| { + live_after_commit(&prompt, consumed, state.live_position(route)).then(|| LiveRecord { + tokens: consumed.to_vec(), + route, + }) + }); + } + + /// Forget the live state (the owners were reset). + pub fn forget_live(&mut self) { + self.live = None; + self.turn = None; + } + + /// Drop every snapshot, pending ones included, and the live state. + pub fn clear(&mut self, gpu: &mut Gpu) { + self.forget_live(); + self.cache.clear(gpu); + } +} + #[cfg(test)] mod tests { use super::*; @@ -841,7 +1065,10 @@ mod tests { self.position = u64::from_le_bytes(meta.try_into().unwrap()) as usize; Ok(()) } - fn growth_reserve_bytes(&self, _gpu: &Gpu) -> u64 { + fn live_position(&self, _route: SessionRoute) -> Option { + Some(self.position) + } + fn growth_reserve_bytes(&self, _gpu: &Gpu, _through: usize) -> u64 { 0 } fn reset(&mut self, gpu: &mut Gpu) -> Result<(), String> { @@ -1133,4 +1360,140 @@ mod tests { gpu.free_tensor(toy.fixed).unwrap(); gpu.free_tensor(toy.rows).unwrap(); } + + const AR: SessionRoute = SessionRoute::Ar; + const MTP: SessionRoute = SessionRoute::Mtp; + const END: u32 = 7; + + fn record(tokens: &[u32], route: SessionRoute) -> LiveRecord { + LiveRecord { + tokens: tokens.to_vec(), + route, + } + } + + #[test] + fn stride_boundaries_union_strides_and_message_ends() { + let mut prompt = vec![1u32; 20]; + prompt[2] = END; // ends at 3 + prompt[7] = END; // ends at 8: a stride multiple + prompt[13] = END; // ends at 14 + assert_eq!(stride_boundaries(4, &prompt, None, 0, 20), [4, 8, 12, 16, 20]); + // Union, sorted and deduplicated (8 is both). + assert_eq!( + stride_boundaries(4, &prompt, Some(END), 0, 20), + [3, 4, 8, 12, 14, 16, 20] + ); + // `after` is exclusive and `up_to` inclusive for both sources. + assert_eq!(stride_boundaries(4, &prompt, Some(END), 3, 14), [4, 8, 12, 14]); + assert_eq!(stride_boundaries(4, &prompt, None, 3, 14), [4, 8, 12]); + // Message ends below one stride are boundaries on their own. + assert_eq!(stride_boundaries(64, &prompt, Some(END), 0, 20), [3, 8, 14]); + assert!(stride_boundaries(64, &prompt, None, 0, 20).is_empty()); + // Consecutive ends each count; `up_to` past the prompt never reads past it. + assert_eq!(stride_boundaries(64, &[END, END, 1, END], Some(END), 0, 9), [1, 2, 4]); + } + + #[test] + fn live_start_requires_a_strict_extension_at_the_owner_position() { + let rec = record(&[1, 2, 3], AR); + let at = |prompt: &[u32]| live_start(Some(&rec), prompt, AR, Some(3)); + assert_eq!(at(&[1, 2, 3, 4]), Some(3)); + assert_eq!(at(&[1, 2, 3, 4, 5, 6]), Some(3)); + // Diverges inside the record. + assert_eq!(at(&[1, 9, 3, 4]), None); + // L == prompt.len() leaves no suffix; a shorter prompt is no extension. + assert_eq!(at(&[1, 2, 3]), None); + assert_eq!(at(&[1, 2]), None); + // No record, an empty one, another route, or owners off the end. + assert_eq!(live_start(None, &[1, 2, 3, 4], AR, Some(3)), None); + assert_eq!(live_start(Some(&record(&[], AR)), &[1, 2], AR, Some(0)), None); + assert_eq!(live_start(Some(&rec), &[1, 2, 3, 4], MTP, Some(3)), None); + assert_eq!(live_start(Some(&rec), &[1, 2, 3, 4], AR, Some(2)), None); + assert_eq!(live_start(Some(&rec), &[1, 2, 3, 4], AR, None), None); + } + + #[test] + fn live_after_commit_needs_consumed_to_extend_the_turn_prompt() { + let prompt = [1, 2, 3]; + let commit = |consumed: &[u32], pos: usize| live_after_commit(&prompt, consumed, Some(pos)); + assert!(commit(&[1, 2, 3, 7, 8], 5)); + // Nothing generated: consumed == prompt. + assert!(commit(&[1, 2, 3], 3)); + // Unrepaired terminal: the host history is empty. + assert!(!commit(&[], 0)); + // Shorter than the prompt, or not extending it. + assert!(!commit(&[1, 2], 2)); + assert!(!commit(&[1, 9, 3, 7], 4)); + // Owners not at the end of the consumed history, or unable to continue. + assert!(!commit(&[1, 2, 3, 7, 8], 4)); + assert!(!live_after_commit(&prompt, &[1, 2, 3, 7], None)); + // A turn without a prompt is never live. + assert!(!live_after_commit(&[], &[1], Some(1))); + } + + #[test] + fn plan_prefers_live_on_a_tie_and_the_longer_source_otherwise() { + assert_eq!(plan_start(4096, None), 4096); + assert_eq!(plan_start(0, None), 0); + assert_eq!(plan_start(0, Some(300)), 300); + assert_eq!(plan_start(4096, Some(4096)), 4096); + assert_eq!(plan_start(4096, Some(9000)), 9000); + assert_eq!(plan_start(8192, Some(300)), 8192); + } + + #[test] + fn driver_continues_live_state_and_restores_snapshots() { + let Ok(mut gpu) = Gpu::init() else { + eprintln!("skip: session cache tests require a GPU"); + return; + }; + let gpu = &mut gpu; + let fixed = gpu.alloc_tensor(&[FIXED_BYTES], DType::Raw).unwrap(); + let rows = gpu + .alloc_tensor(&[ROW_CAPACITY * ROW_BYTES], DType::Raw) + .unwrap(); + let mut toy = Toy { + fixed, + rows, + position: 0, + }; + let one = layout([FIXED_BYTES, STRIDE * ROW_BYTES]).1 as u64; + let mut driver = SessionDriver::new(SessionCache::new(domain(), 8 * one)); + let prefill = |driver: &mut SessionDriver, toy: &mut Toy, gpu: &mut Gpu, p: &[u32]| { + let reused = driver.plan(toy, p, AR); + let live = driver.begin(gpu, toy, p, AR, reused, false).unwrap(); + while let Some(b) = driver.next_boundary() { + toy.advance(gpu, &p[..b]); + driver.at_boundary(gpu, toy, &p[..b]).unwrap(); + } + toy.advance(gpu, p); + (reused, live) + }; + // Turn 1, cold; the decoded tail extends the state, commit keeps it live. + let turn1 = prompt(1, STRIDE + 5); + assert_eq!(prefill(&mut driver, &mut toy, gpu, &turn1), (0, false)); + let mut consumed = turn1.clone(); + consumed.extend([11, 12]); + toy.advance(gpu, &consumed); + driver.commit_live(&toy, &consumed); + // Turn 2 extends the consumed history: continue in place. + let mut turn2 = consumed.clone(); + turn2.extend(prompt(2, 9)); + assert_eq!( + prefill(&mut driver, &mut toy, gpu, &turn2), + (consumed.len(), true) + ); + assert_eq!(toy.bytes(gpu), (fixed_of(&turn2), rows_of(&turn2))); + // A plain commit drops the live state: another conversation sharing + // turn 1's first stride restores the snapshot. + driver.commit(); + let mut other = turn1[..STRIDE].to_vec(); + other.extend(prompt(3, 4)); + assert_eq!(prefill(&mut driver, &mut toy, gpu, &other), (STRIDE, false)); + assert_eq!(toy.bytes(gpu), (fixed_of(&other), rows_of(&other))); + driver.clear(gpu); + gpu.free_tensor(toy.fixed).unwrap(); + gpu.free_tensor(toy.rows).unwrap(); + } } diff --git a/crates/railgun-cert/src/recording.rs b/crates/railgun-cert/src/recording.rs index 44ad16ac27..a779ef09c8 100644 --- a/crates/railgun-cert/src/recording.rs +++ b/crates/railgun-cert/src/recording.rs @@ -649,7 +649,7 @@ pub const DECISIONS: &[Decision] = &[ Decision { id: "replay-recorder-lifecycle", predicates: &[IsRecording, ReplayState], - condition: "`should_auto_finalize_capture` (replay.rs:4586-4588) in qwen35 forward_scratch (forward.rs:1831), deepseek4 decode_step_with_graph (forward.rs:3642), lfm2moe decode_step_with_retained_replay (state checks, forward.rs:106-170); gemma4 lowered forward_scratch keeps its AR HipGraph off while `is_recording()` (lowered.rs)", + condition: "`should_auto_finalize_capture` (replay.rs:4586-4588) in qwen35 forward_scratch (forward.rs:1831), deepseek4 decode_step_with_graph (forward.rs:3642), lfm2moe decode_step_with_retained_replay (state checks, forward.rs:106-170); gemma4 `graph_exact` keeps its AR HipGraph off while `is_recording()` (forward.rs)", effect: Lifecycle, switches: "finish the one-shot capture after the recorded forward", kernels: &[], @@ -1159,7 +1159,8 @@ pub const SITES: &[Site] = &[ Site { file: "hipfire-arch-deepseek4/src/spec_impl.rs", function: "dspark_verify_forward", occurrences: 2, decisions: &["dspark-capture-safe-verify-body", "graph-capture-lifecycle"] }, Site { file: "hipfire-arch-deepseek4/src/spec_impl.rs", function: "redline_dspark_verify_direct", occurrences: 2, decisions: &["dspark-capture-safe-verify-body"] }, Site { file: "hipfire-arch-deepseek4/src/spec_impl.rs", function: "redline_dspark_verify_pm4", occurrences: 1, decisions: &["dspark-capture-safe-verify-body"] }, - Site { file: "hipfire-arch-gemma4/src/lowered.rs", function: "forward_scratch", occurrences: 3, decisions: &["graph-capture-lifecycle", "replay-recorder-lifecycle"] }, + Site { file: "hipfire-arch-gemma4/src/forward.rs", function: "decode_step_with_graph", occurrences: 1, decisions: &["graph-capture-lifecycle"] }, + Site { file: "hipfire-arch-gemma4/src/forward.rs", function: "graph_exact", occurrences: 1, decisions: &["replay-recorder-lifecycle"] }, Site { file: "hipfire-arch-lfm2moe/src/forward.rs", function: "decode_step", occurrences: 1, decisions: &["replay-enabled-process-route"] }, Site { file: "hipfire-arch-lfm2moe/src/forward.rs", function: "decode_step_with_retained_replay", occurrences: 3, decisions: &["replay-recorder-lifecycle"] }, Site { file: "hipfire-arch-llama/src/arch.rs", function: "forward_scratch_layers", occurrences: 2, decisions: &["kv-tier-capture-forces-flash"] }, diff --git a/crates/saddle-lab/examples/build_kld_ref_native.rs b/crates/saddle-lab/examples/build_kld_ref_native.rs index 7863592538..fefd8225cf 100644 --- a/crates/saddle-lab/examples/build_kld_ref_native.rs +++ b/crates/saddle-lab/examples/build_kld_ref_native.rs @@ -537,11 +537,10 @@ fn main() { ); let _ = Path::new("/dev/null"); } else if is_gemma { - // -------- gemma4 (arch 13) — q8 KV, sub-chunk 128, forward_prefill_batch -------- - use hipfire_arch_gemma4::lowered::{self, Gemma4Scratch}; - use hipfire_runtime::llama::KvCache; + // -------- gemma4 (arch 13) — q8 KV, per-token decode -------- + use hipfire_arch_gemma4::{forward, Gemma4Config, Gemma4State, Gemma4Weights}; - let cfg = lowered::config_from_hfq(&hfq).expect("gemma4 config_from_hfq"); + let cfg = Gemma4Config::from_hfq(&hfq).expect("gemma4 config"); eprintln!( "gemma4 arch=13 n_layers={} dim={} vocab={} hidden={} n_ctx={} top_k={}", cfg.n_layers, cfg.dim, cfg.vocab_size, cfg.hidden_dim, n_ctx, top_k @@ -550,7 +549,7 @@ fn main() { eprintln!("gemma4 MoE (26B-A4B) not supported for BF16 KLD dense reference; refusing."); std::process::exit(1); } - let weights = lowered::load_weights(&mut hfq, &cfg, &mut gpu).expect("gemma4 load_weights"); + let weights = Gemma4Weights::load(&hfq, &cfg, &mut gpu).expect("gemma4 load_weights"); eprintln!("gemma4 loaded {} layers", cfg.n_layers); if let Some(parent) = output.parent() { @@ -574,25 +573,8 @@ fn main() { } let kv_max = n_ctx + 16; - let scratch = Gemma4Scratch::new(&mut gpu, &cfg, kv_max).expect("gemma4 scratch"); - lowered::init_scratch_constants(&mut gpu, &scratch, cfg.full_head_dim) - .expect("gemma4 init_scratch_constants"); - let mut kv_sliding = KvCache::new_gpu_q8( - &mut gpu, - cfg.n_layers, - cfg.sliding_n_kv_heads, - cfg.sliding_head_dim, - kv_max, - ) - .expect("gemma sliding q8 kv"); - let mut kv_full = KvCache::new_gpu_q8( - &mut gpu, - cfg.n_layers, - cfg.full_n_kv_heads, - cfg.full_head_dim, - kv_max, - ) - .expect("gemma full q8 kv"); + let mut state = + Gemma4State::new_with_max_seq(&mut gpu, &cfg, kv_max).expect("gemma4 state"); // An F32 cache would error with `no implementation for KvWriteF32` — proven in calib_sweep. let k = top_k; @@ -603,35 +585,27 @@ fn main() { let mut scored_done = 0usize; for c in 0..n_chunk { - kv_sliding + state + .kv_sliding .clear_gpu(&mut gpu) .expect("gemma kv sliding clear"); - kv_full.clear_gpu(&mut gpu).expect("gemma kv full clear"); + state.kv_full.clear_gpu(&mut gpu).expect("gemma kv full clear"); let chunk = &tokens[c * n_ctx..(c + 1) * n_ctx]; // Per-token decode for the entire chunk, matching qwen35's proven - // per-token loop and calib_sweep's correctness reference (forward_scratch - // via Gemma4Bindings, not the batched prefill). The previous - // hybrid (batched prefill for prefix + per-token for scored window) - // left KV/position state misaligned at the 256 boundary — see - // lowered.rs forward_prefill_batch_v2 (batched Q/K/V + batched full - // attention) vs forward_scratch_inner_lowered (per-token). Using - // the single per-token path eliminates the boundary class. + // per-token loop. for pos in 0..(n_ctx - 1) { - lowered::forward_scratch( - &mut gpu, - &weights, + let cand_logits = forward::decode_step( &cfg, + &weights, + &mut state, + &mut gpu, chunk[pos], - pos, - &mut kv_sliding, - &mut kv_full, - &scratch, + pos as u32, ) - .expect("gemma forward_scratch"); + .expect("gemma decode_step"); if pos < scoring_start { continue; } - let cand_logits = gpu.download_f32(&scratch.logits).expect("download logits"); let actual_next = chunk[pos + 1] as usize; write_scored_position( &cand_logits, diff --git a/docs/CONFIG.md b/docs/CONFIG.md index cbde2ed9a2..43e6eac3c3 100644 --- a/docs/CONFIG.md +++ b/docs/CONFIG.md @@ -365,11 +365,12 @@ never inherit fp8. Registry Qwen cards leave mode as `auto` (a registry `default_kv_mode = "q8"` is not lowered into config). **Non-Qwen families:** Maple keeps BF16/BF16 (registry and direct-path -auto; explicit `--kv-mode q8` still works). Gemma4 eager stays Q8/Q8. -Gemma4 lowered keeps the sliding Q8 ring; its full-attention tier follows -`kv_cache`: `auto`/`q8` → Q8 on every arch (no lowered Gemma4 attend site -admits native FP8), `legacy-asym3` → the previous Givens Asym3 K + Q8 V, -`fp8`/`bf16` → load error, anything else → Q8 with a warning. llama HFQ, +auto; explicit `--kv-mode q8` still works). Gemma4 keeps its sliding tier +Q8 on every variant; its full-attention tier follows `kv_cache`: `auto`/`q8` +→ Q8 on every arch (no Gemma4 attend site admits native FP8), +`legacy-asym3` → the previous Givens Asym3 K + Q8 V, +`fp8`/`bf16` → load error, anything else → Q8 with a warning. Only Q8/Q8 +loads take batched prefill and the session cache. llama HFQ, MiniMax and LFM2-MoE resolve `auto` to q8. DeepSeek4 keeps its F32 compressor default (or explicit F16); V is not independently selectable there. Other carriers keep their existing site policy (llama, MiniMax and diff --git a/docs/MODELS.md b/docs/MODELS.md index f2340c8743..eeb9769619 100644 --- a/docs/MODELS.md +++ b/docs/MODELS.md @@ -217,9 +217,8 @@ Registry `recommended_settings` for LFM tags is low temperature (0.05–0.2) wit > `gemma4:` tags so that no registry row can point at a 404. See the report at > the bottom of this section for the exact repos/files to publish. -**Architecture ground truth** (`crates/hipfire-arch-gemma4/src/config.rs`, -`crates/hipfire-arch-gemma4/src/lowered.rs`, and -`crates/hipfire-arch-gemma4/src/gemma4.rs`): +**Architecture ground truth** (`crates/hipfire-arch-gemma4/src/config.rs` +and `crates/hipfire-arch-gemma4/src/gemma4.rs`): - **Hybrid 5:1 sliding:global** — every 6th layer is global (full) attention. The per-layer `layer_types` array is authoritative; on the 12B dense text @@ -237,7 +236,7 @@ Registry `recommended_settings` for LFM tags is low temperature (0.05–0.2) wit weight-less, implemented with a ones-filled scratch buffer). - **KV head counts (variant-dependent):** 12B text — `num_attention_heads = 16`, `num_key_value_heads = 8` (sliding) and `num_global_key_value_heads` defaults - to sliding when absent; 31B target (per `lowered.rs` comments) — `n_heads = 32`, + to sliding when absent; 31B target — `n_heads = 32`, `sliding_n_kv_heads = 16`, `full_n_kv_heads = 4`, `sliding_head_dim = 256`, `full_head_dim = 512`, `hidden_dim = 21504`. The 12B `hidden_dim` is `intermediate_size = 15360`, `dim = 3840`. diff --git a/docs/architecture-ids.md b/docs/architecture-ids.md index e8cd03c5c9..62f68c8e49 100644 --- a/docs/architecture-ids.md +++ b/docs/architecture-ids.md @@ -47,7 +47,7 @@ them for ordinary dispatch. |---:|---|---|---| | 20 | DFlash draft artifact (generic) | `hipfire-quantize` `dflash_convert` (`ARCH_ID_DFLASH_DRAFT`); loader draft open checks `draft_hfq.arch_id == 20` | Generic Qwen-family diffusion draft (fc/hidden_norm/norm naming, `dflash` metadata). Loaded as a draft alongside a compatible target, not as a standalone serve model. | | 21 | Qwen3.5 MTP head | `hipfire-quantize` `mtp_extract` (`ARCH_ID_QWEN35_MTP_HEAD`); consumed as `.mq4-mtp` trailer / `.mtp` sidecar on qwen35 trunks | Attached onto arch 5/6 loads (`LoadedModel.qwen35_mtp_head`), not a carrier `arch_id`. Refused before upload unless its metadata `n_embd`/`vocab_size` equal the trunk's hidden size and vocab (`mtp_head::check_mtp_head_for_trunk`). | -| 22 | Gemma4 EAGLE draft | `hipfire-arch-gemma4` `DRAFTER_ARCH_ID`; claimed by `Gemma4Carrier` alongside 13 | `model_type gemma4_unified_assistant` — hidden-state-consuming EAGLE (pre/post projections, KV-shared layers). Loaded via `params.drafter` on arch 13. **Currently gated off** (`HIPFIRE_GEMMA4_EAGLE=1` to enable): diverges from the greedy AR sequence it is required to reproduce, and is slower. | +| 22 | Gemma4 EAGLE draft | `hipfire-arch-gemma4` `DRAFTER_ARCH_ID`; claimed by `Gemma4Carrier` alongside 13 | `model_type gemma4_unified_assistant` — hidden-state-consuming EAGLE (pre/post projections, KV-shared layers). Loaded via `params.drafter` on arch 13 as an `MtpSpeculator` (`Gemma4Drafter`). **Opt-in** (`HIPFIRE_GEMMA4_EAGLE=1`): batched verify is not bitwise-equal to single-row decode, so greedy output can leave the AR sequence at near-tied logits, and it is not faster than AR on gfx1151. | | 23 | Muse Glimmer DFlash draft | `hipfire-arch-muse-glimmer` `GLIMMER_DRAFTER_ARCH_ID` | `model_type muse_glimmer_assistant` — 5-layer block-diffusion draft (encoder.fc/output_norm_enc, block 16, target_layer_ids [1,13,25,37,49], mask 201818). Loaded via `HIPFIRE_DFLASH_DRAFT` on arch 14. NOT a `dflash_convert` 20 artifact (different tensor naming: encoder.fc vs fc, separate `config` envelope). Quantizer auto-maps `muse_glimmer_assistant` → 23 (`hipfire-quantize/src/main.rs:8607`, `hipfire-runtime/src/safetensors_source.rs:271`); HFQ on disk is `arch 23`. | | 0xFF | Toy / template | `hipfire-arch-toy` | Never ship; daemon must not dispatch. No current carrier claims it. Registry disjointness tests include `0xFF` and assert at most one claimer — they do **not** assert zero claimers / hard-reserved. | | `u32::MAX` | Unclaimed dir sentinel | `safetensors_source::UNCLAIMED_ARCH_ID` | Emitted for unrecognized `model_type`; no carrier matches → fail closed. | diff --git a/docs/design/gemma4-declarative-program.md b/docs/design/gemma4-declarative-program.md index 04a685cb91..2094aac3f3 100644 --- a/docs/design/gemma4-declarative-program.md +++ b/docs/design/gemma4-declarative-program.md @@ -13,20 +13,23 @@ EAGLE/MTP drafter. The architecture crate keeps: -- weight and state ownership (one config, one weight set, KV allocation); +- weight and state ownership (one config, one weight set, KV allocation) + for every variant: 12B dense, E2B/E4B and the 26B-A4B MoE; - the per-layer program shape (`src/program.rs`: which ops, in which order, bound to which tensors); -- lifecycle: hipGraph capture, the spec-decode driver, Redline shadow, the - chunk loop. +- lifecycle: hipGraph capture, the chunk loop, the session-cache owner + (`src/bundle.rs`) and the EAGLE drafter/verify seams (`src/mtp.rs`). The + spec-decode loop itself is the engine's generic `generate_spec`. The #397 super-op facade (`lowered.rs` `lower_variant` / `Gemma4Bindings`) -and the hand-written decode arms in `forward.rs` and `lowered.rs` are deleted. -Gemma 4 has no EP/TP route, so nothing keeps the super-op program alive for it. +and the hand-written decode arms in `forward.rs` are deleted, and so is the +whole `lowered.rs` stack. Gemma 4 has no EP/TP route, so nothing keeps the +super-op program alive for it. ## Before the port -Two parallel stacks, each with its own config, weights and loader (both stay; -they now bind the same program): +Two parallel stacks, each with its own config, weights and loader (unified +since; see "Engine seams"): | Stack | Fixtures | KV | Notes | |---|---|---|---| @@ -145,6 +148,59 @@ batched kernel one GEMV per row. `eval_hipfire` per-token 0.733489 and prefill 0.824203: both `.kldseq` files byte-identical to the pre-rebase branch. +## Engine seams + +After the program port, Gemma 4 moved onto the seams Qwen4 introduced: + +- **One stack.** `config.rs` carries the MoE fields, `gemma4.rs` loads every + variant (pooled expert tensors and pointer tables for the 26B), and + `Gemma4State` holds position-indexed sliding and full caches for all of + them (the 26B's 1024-row sliding ring is gone). `lowered.rs` is deleted. + `supports_batched_prefill` admits the 26B on Q8/Q8 loads. The decode + hipGraph stays off when an expert `down_proj` is Q8_0 (its atomic sum is + not replay-exact) and while a launch recorder is armed (`graph_exact`). +- **Session cache.** `Gemma4Bundle` (`bundle.rs`) implements + `SessionState`: a snapshot is the Q8 K/V rows of every cache slot up to its + position (`RowStream`s, no fixed state), scoped `gemma4/rows{N}` by prefill + width because batched and per-token prefill write different bits. Boundaries + fall every 1024 prompt tokens. The route-generic live/plan/begin/commit + logic is `hipfire_runtime::session_cache::SessionDriver`, shared with + Qwen4. Each request now plans against its own prompt. Before this, a Gemma 4 + request continued the previous request's KV in the same process. +- **Spec seams.** `mtp.rs`: `Gemma4Drafter: MtpDrafter` (driven by the + generic `MtpSpeculator`) and `SpecTarget for Gemma4Bundle`. A window at + position `p` holds the pending seed and the target's post-norm hidden of + `p - 1`. It drafts `k` tokens at the constant query position `p` through + the step-list draft head (`Embed`, `Scale`, `Gemv` pre-projection, the + query-only blocks, norm, `DraftHead` whole-vocab argmax, `Gemv` + post-projection), verifies `[seed, drafts]` in one batched forward and + commits the greedy prefix plus bonus. KV is position-indexed, so + `commit_prefix` only moves the cursor. `emit.rs` `GemmaSpecEmit` routes + the thought channel, and `generate_gemma4_spec` wraps the engine's + `generate_spec`. The inline EAGLE loop and `speculative.rs` are deleted. + +Measured (gfx1151, PR-head daemon md5 `83c3ca2c175e0ac66b47ce1802e8f815` +vs `7cb5f3aa83272fd5283628332ad4e29d`, five greedy prompts, 128 tokens, a +daemon `reset` between prompts on both arms so neither carries the previous +request): + +- 12B AR, 12B prefill batch 64, E2B, E4B, E4B prefill batch 64: token + streams and top-16 logit traces byte-identical. +- 26B-A4B q8-experts and requantized MQ4: text identical on all five + prompts. Decode 36.1 vs 36.4 tok/s (q8) and 43.2 vs 43.9 (MQ4), measured + alone. Runs overlapping a cargo build read 10–15% lower on either binary. +- EAGLE (12B, `HIPFIRE_GEMMA4_EAGLE=1`, `spec` 2, 3 and 5): four of five + prompts equal AR. The fifth (p1) leaves AR at step 91, where AR's top-1/top-2 + margin is 0.021 logits (26.2946 vs 26.2735), the smallest in that sequence. + Batched verify logits differ from single-row decode by about 1e-2 from the + first verify row on. Strict mode does not make them bitwise-equal on + gfx1151, so greedy EAGLE equals AR only where the margin exceeds that + noise. The PR-head loop matched AR on all five prompts. It re-verified the + last prompt token, so its KV differs from this route's and the near-tie + fell the other way. Per-prompt tau at `spec` 3: 3.342/3.175/2.592/3.528/3.024 + vs PR head 3.368/3.122/2.633/3.639/2.977. The first token now comes from + prefill, and each window pairs its seed with the hidden that predicted it. + ## Progress | Piece | State | Commit | @@ -153,9 +209,16 @@ batched kernel one GEMV per row. | MoE decode (`ParallelMoeMlp`), lowered decode on the shared program; super-op facade (`lower_variant`, `Gemma4Bindings`) and lowered hand arms deleted; unsupported expert formats refuse at load | landed; 26B byte-identical with the two non-bitwise fusions off, coherent with them on | `287e67d5b` | | EAGLE draft head as one step list (`Gemv` pre-projection, query-only `[SandwichAttention, SandwichMlp, Scale?]` blocks over the target's last slot, norm, `lm_head`, post-projection) | landed; EAGLE text and per-prompt tau identical | `287e67d5b` | | Calibration tools (`calib_sweep`, `eval_hipfire`, `prefill_parity_gemma4`) on the batched program; dispatch-owned calibration taps; batched `ParallelMoeMlp`; old batched prefill, `HIPFIRE_BATCHED_PREFILL`/`HIPFIRE_WMMA_PREFILL` deleted | landed; tool parity above | `e89080332` | +| One weight stack for every variant; `lowered.rs` deleted | landed; see "Engine seams" | this change | +| Engine session cache (`SessionState`, `SessionDriver` shared with Qwen4) | landed; see "Engine seams" | this change | +| EAGLE on `MtpDrafter`/`SpecTarget`/`SpecEmit` + `generate_spec`; `DraftHead` | landed; see "Engine seams" | this change | ## Remaining -- Two weight stacks remain (`gemma4.rs` eager for dense/E-series/EAGLE, - `lowered.rs` for MoE) with duplicate config and loaders; both bind to the - same program. +- Batched verify is not bitwise-equal to single-row decode on gfx1151, so + greedy EAGLE can leave the AR stream at near-tied logits. EAGLE stays + opt-in until the verify and decode arithmetic match. +- In the session scenario a follow-up turn restored the newest snapshot + instead of extending the previous turn's live state, so its rendered + history is not a token-exact extension of that turn (assumed cause: the + reply is re-rendered from text; Qwen keeps a verbatim assistant-turn cache). diff --git a/docs/governance/debt-dispatch-bypass.txt b/docs/governance/debt-dispatch-bypass.txt index 084b05a766..bee8be2ae8 100644 --- a/docs/governance/debt-dispatch-bypass.txt +++ b/docs/governance/debt-dispatch-bypass.txt @@ -34,7 +34,7 @@ hipfire-arch-qwen35 partial 184 191 migrate last; MTP routed MoE now dispatch-layer owned; +10 = DFlash fp16-X GEMM entries (S3/S4 launch fusion); -20 = MQ v2 lm_head match deduped in speculative.rs; +7 = MQ4-Lloyd (qt=52) FP8-LUT prefill arms (gemm_*_fp8_lloyd) pending registry keys with LUT plumbing (gfx11 MMQ-LUT PR); +15 = MQ4-Lloyd gfx11 MMQ-LUT prefill arms (gemm_*_mmq_lloyd), same registry-key debt as the FP8-LUT arms; +40 = speculation-decode additions recorded by the SCS hygiene pass (was carried unrecorded on the beta integration base); +1 from the SCS port PR series (landed with the SCS fold); +2 landed on beta since this ledger was frozen (recorded on rebase, not new debt); -5 = sealed MoE prefill moved into dispatch (PR #755 via #774); +4 = gfx1151 A4-fusion gate/up + down-sidecar arms (gemm_gate_up_silu_a4_*, gemm_*_iu4_down_prepared) beside the Iu4H direct calls they mirror, same registry-key debt hipfire-arch-deepseek4 debt 19 0 cheapest full conversion; no registry use to reconcile -hipfire-arch-gemma4 debt 2 0 -15 = Gemma 4 layers run as hipfire-dispatch sandwich Steps (#813); the 2 left are the PLE Q8 projection GEMMs (forward.rs) +hipfire-arch-gemma4 debt 1 0 -15 = Gemma 4 layers run as hipfire-dispatch sandwich Steps (#813); -1 = spec verify lm_head keeps only the scalar Q8 GEMM (the env-gated chunked branch is gone); the 1 left is that lm_head (forward.rs) hipfire-arch-lfm2moe debt 17 0 cheapest full conversion; no registry use to reconcile hipfire-arch-minimax debt 12 0 hipfire-arch-muse-glimmer partial 12 17 diff --git a/scripts/oracle_gemma4_26b.py b/scripts/oracle_gemma4_26b.py index 26e79c68f4..e3d6089c09 100644 --- a/scripts/oracle_gemma4_26b.py +++ b/scripts/oracle_gemma4_26b.py @@ -8,7 +8,7 @@ intermediate tensors: post-attn hidden, post-FFN hidden, MoE branch intermediates (router logits, topk indices, expert outputs). -Compares against hipfire's HIPFIRE_GEMMA4_DUMP=1 output. +Compares against hipfire's per-layer hidden states (`dump_gemma4_hidden_states` example). Usage: # Quick test with short prompt From ecb0471505579be5aad1debb12141b2b62fc69ce Mon Sep 17 00:00:00 2001 From: Bjoern Agent Date: Fri, 9 Oct 2026 14:39:38 +0200 Subject: [PATCH 12/13] perf(gemma4): gfx1151 batched prefill by default, 26B-A4B decode and prefill kernels gfx1151 now prefills Gemma 4 in 256-row batches (GEMMA4_FORWARD_BATCH_MAX 64 -> 256; flash partials stay sized for 64 rows and launches sub-batch). gfx1100/gfx1201 keep 64, other archs stay per-token. Prefill (batched rows; new routes gated to gfx1151): - 26B-A4B routed experts sorted by expert: batched router (F32-activation Q8 GEMV; the F16 WMMA GEMM flips near-tied experts), one scatter, grouped WMMA gate/up (MQ4G256V2, MQ6G256), grouped HFQ4-G128 down (new WMMA kernel), combine. Grouped tile count rounds up to whole 16-slot tiles. - MQ4G256V2 projections take the WMMA GEMM (previously one GEMV per row); Q8 projections the 4-warp 64x64 WMMA tile (its F32 input is now converted uncached). - GQA-shared WMMA flash prefill attention, sliding window included (attention_q8_0_prefill_gqa_wmma, HIPFIRE_Q8_PREFILL_GQA_WMMA=0 opts out). - LM head on the last prefill chunk only. Decode: - Flash Q8 tiles wholly below the sliding window exit early; the reduce skips zero-weight tiles, computes the global max in parallel and loads partials eight at a time (bit-exact; P0 gfx1201 source digests repinned). - gfx1151 GQA-shared 8-wave decode tile at the Gemma shapes (HIPFIRE_GFX1151_Q8_DECODE_ATTN_GQA_GEMMA=0 opts out). - HFQ4-G128 routed down: eight rows per workgroup, all experts in flight. Evidence (gfx1151): gemma4_kernel_f64_check (f64 CPU reference) 2.8e-7.. 4.2e-7 decode attention, 1.3e-3..1.6e-3 F16 prefill attention, 1.9e-7 / 3.3e-4 MoE down. Reference-continuation NLL vs per-token prefill moves by <= 0.02 nats/token on 12B, E2B, E4B and both 26B fixtures. Redline --pm4 --skip-prefill 26B MQ4: shadow exact, 1202 launches. serve_harness battery + chain clean on 12B, E4B, 26B q8-experts, 26B MQ4. Perf in the PR description. --- CHANGELOG.md | 28 +- crates/hipfire-arch-gemma4/src/bundle.rs | 38 ++- crates/hipfire-arch-gemma4/src/forward.rs | 58 +++- crates/hipfire-arch-gemma4/src/gemma4.rs | 8 +- crates/hipfire-arch-gemma4/src/program.rs | 1 + .../hipfire-dispatch/src/pipeline/sandwich.rs | 272 ++++++++++++++- crates/rdna-compute/Cargo.toml | 4 + .../examples/gemma4_kernel_f64_check.rs | 310 ++++++++++++++++++ crates/rdna-compute/src/attention.rs | 159 ++++++++- crates/rdna-compute/src/gemm.rs | 82 ++++- crates/rdna-compute/src/kernel_registry.rs | 6 +- crates/rdna-compute/src/kernels.rs | 16 + crates/rdna-compute/src/replay.rs | 6 +- .../tests/fixtures/p0-gfx1201-sha256.tsv | 6 +- kernels/src/attention_flash_q8_0_reduce.hip | 58 +++- kernels/src/attention_flash_q8_0_tile.hip | 11 + ...tion_flash_q8_0_tile_gqa_gemma.gfx1151.hip | 232 +++++++++++++ .../attention_q8_0_prefill_gqa_wmma.gfx11.hip | 228 +++++++++++++ .../gemm_hfq4g128_moe_down_grouped_wmma.hip | 79 +++++ .../src/gemv_hfq4g128_moe_down_grouped.hip | 96 ++++++ ...28_moe_down_residual_scaled_k8_indexed.hip | 90 +++-- 21 files changed, 1698 insertions(+), 90 deletions(-) create mode 100644 crates/rdna-compute/examples/gemma4_kernel_f64_check.rs create mode 100644 kernels/src/attention_flash_q8_0_tile_gqa_gemma.gfx1151.hip create mode 100644 kernels/src/attention_q8_0_prefill_gqa_wmma.gfx11.hip create mode 100644 kernels/src/gemm_hfq4g128_moe_down_grouped_wmma.hip create mode 100644 kernels/src/gemv_hfq4g128_moe_down_grouped.hip diff --git a/CHANGELOG.md b/CHANGELOG.md index db498831c5..e478d24d12 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -140,7 +140,7 @@ - The 26B-A4B now takes the 12B's fused qk-norm+RoPE and post-norm+residual kernels. Greedy text stays coherent and can diverge late. - `calib_sweep`, `eval_hipfire` and `prefill_parity_gemma4` prefill Gemma 4 - through the same program (`forward::forward_prefill_batch`, 64-row + through the same program (`forward::forward_prefill_batch`, 256-row chunks, Q8 KV). The dispatch sandwich projections feed the calibration collector, so any architecture on these steps calibrates without per-architecture taps. Gemma calibration and KLD numbers move: the tools @@ -160,7 +160,7 @@ - Gemma 4 runs on one weight stack: the 26B-A4B MoE loads through the same config, weights, state and forward as the 12B and E-series, and the separate `lowered.rs` stack is gone. Its sliding cache is position-indexed like the - 12B's (no 1024-row ring), so batched prefill (default on gfx1100/gfx1201) + 12B's (no 1024-row ring), so batched prefill (default on gfx1100/gfx1151/gfx1201) now reaches the 26B on Q8/Q8 loads. Its hipGraph decode is off when expert `down_proj` is Q8_0, whose atomic sum is not replay-exact. Text on the q8-experts and requantized MQ4 26B @@ -186,6 +186,30 @@ bitwise-equal to single-row decode, so greedy EAGLE can leave the AR stream where two logits are nearly tied. Measured on gfx1151: one of five fixture prompts at a 0.021-logit margin. +- Gemma 4 prefills in 256-row batches by default on gfx1151 + (`HIPFIRE_GEMMA4_PREFILL_BATCH` overrides; gfx1100/gfx1201 keep 64, other + architectures stay per-token). On gfx1151 the batched rows: + - route the 26B-A4B's experts sorted by expert: one scatter, a grouped WMMA + gate/up (MQ4G256V2, MQ6G256) and a grouped HFQ4-G128 down per batch, so + each expert's weights are read once per batch. The router stays an + F32-activation Q8 GEMV: the F16 WMMA GEMM moves router logits by ~1e-3 + and flips near-tied experts; + - project MQ4G256V2 weights through WMMA (previously one GEMV per row) and + Q8 weights through the 4-warp 64x64 WMMA tile; + - attend with a GQA-shared WMMA flash kernel, sliding window included + (`HIPFIRE_Q8_PREFILL_GQA_WMMA=0` restores the tile kernels); + - compute the LM head on the last chunk only. + Greedy continuations differ from per-token prefill only by reduction order: + the reference continuation's NLL moves by at most 0.02 nats/token across + 12B, E2B, E4B and both 26B fixtures. +- Gemma 4 26B-A4B long-context decode: flash-decode tiles that lie wholly + below the sliding window exit without reading V, and the Q8 reduce skips + them and loads tile partials eight at a time (both bit-exact, for every + model on these kernels). On gfx1151 the Gemma decode shapes (head_dim 256 / + GQA 2, head_dim 512 / GQA 8) run a GQA-shared 8-wave tile + (`HIPFIRE_GFX1151_Q8_DECODE_ATTN_GQA_GEMMA=0` restores the reference tile), + and the HFQ4-G128 routed down kernel computes eight rows per workgroup + with all eight experts' loads in flight. ## v0.4.1.1 — release draft diff --git a/crates/hipfire-arch-gemma4/src/bundle.rs b/crates/hipfire-arch-gemma4/src/bundle.rs index a76c1d094e..2e8e6cf206 100644 --- a/crates/hipfire-arch-gemma4/src/bundle.rs +++ b/crates/hipfire-arch-gemma4/src/bundle.rs @@ -21,21 +21,22 @@ use hipfire_runtime::session_cache::{ }; use rdna_compute::Gpu; -/// Snapshot boundary stride in prompt tokens (a multiple of the 64-row +/// Snapshot boundary stride in prompt tokens (a multiple of the 256-row /// prefill chunk, so restored prefills keep the cold chunk grid). pub const SESSION_STRIDE: usize = 1024; -/// Prefill rows per forward on `gpu_arch`: 64-row batches where batched -/// prefill is validated (`HIPFIRE_GEMMA4_PREFILL_BATCH` overrides), else -/// per-token decode. +/// Prefill rows per forward on `gpu_arch` where batched prefill is validated +/// (`HIPFIRE_GEMMA4_PREFILL_BATCH` overrides), else per-token decode. gfx1151 +/// takes 256-row batches: its grouped MoE, Q8 and attention prefill routes +/// read each weight / KV tile once per batch. pub fn prefill_rows_for_arch(gpu_arch: &str, requested: Option) -> usize { requested - .unwrap_or(if matches!(gpu_arch, "gfx1100" | "gfx1201") { - 64 - } else { - 1 + .unwrap_or(match gpu_arch { + "gfx1151" => 256, + "gfx1100" | "gfx1201" => 64, + _ => 1, }) - .clamp(1, 64) + .clamp(1, crate::gemma4::GEMMA4_FORWARD_BATCH_MAX) } /// Published Gemma 4 owner (arch_id 13). @@ -171,6 +172,7 @@ impl Gemma4Bundle { ) -> Result, String> { let width = self.prefill_rows; let mut logits = Vec::new(); + let n_chunks = tokens.len().div_ceil(width); for (i, chunk) in tokens.chunks(width).enumerate() { let pos = start + i * width; logits = if let [token] = chunk { @@ -182,6 +184,17 @@ impl Gemma4Bundle { *token, pos as u32, )? + } else if i + 1 < n_chunks { + // Only the last chunk's logits are returned; skip the LM head. + forward::forward_prefill_batch( + &self.config, + &self.weights, + &mut self.state, + gpu, + chunk, + pos, + )?; + Vec::new() } else { forward::forward_batch( &self.config, @@ -371,14 +384,15 @@ mod tests { use super::prefill_rows_for_arch; #[test] - fn validated_arches_default_to_full_wmma_batch_tile_group() { + fn validated_arches_default_to_batched_prefill() { assert_eq!(prefill_rows_for_arch("gfx1100", None), 64); assert_eq!(prefill_rows_for_arch("gfx1201", None), 64); + assert_eq!(prefill_rows_for_arch("gfx1151", None), 256); } #[test] fn unvalidated_arches_keep_the_sequential_default() { - for arch in ["gfx1151", "gfx1200", "gfx942"] { + for arch in ["gfx1200", "gfx942"] { assert_eq!(prefill_rows_for_arch(arch, None), 1); } } @@ -388,6 +402,6 @@ mod tests { assert_eq!(prefill_rows_for_arch("gfx1100", Some(8)), 8); assert_eq!(prefill_rows_for_arch("gfx1201", Some(32)), 32); assert_eq!(prefill_rows_for_arch("gfx1100", Some(0)), 1); - assert_eq!(prefill_rows_for_arch("gfx1100", Some(128)), 64); + assert_eq!(prefill_rows_for_arch("gfx1100", Some(512)), 256); } } diff --git a/crates/hipfire-arch-gemma4/src/forward.rs b/crates/hipfire-arch-gemma4/src/forward.rs index 6c0cafa599..bfa7f1e986 100644 --- a/crates/hipfire-arch-gemma4/src/forward.rs +++ b/crates/hipfire-arch-gemma4/src/forward.rs @@ -678,7 +678,7 @@ fn checked_batch_seq_len(start_pos: usize, batch: usize, max_seq: usize) -> Resu } /// Batched verify forward. See module-level note above. `tokens.len()` = B -/// (1..=64, the `gemm_q8_0_batched` / batched-attention kernel cap). Returns +/// (1..=`GEMMA4_FORWARD_BATCH_MAX`). Returns /// the LAST token's logits. Side effect: writes both KV caches for positions /// [start_pos, start_pos+B) and sets `state.n_tokens = start_pos + B`. /// @@ -892,13 +892,67 @@ fn forward_rows( }), _ => None, }; + // Grouped routed-expert scratch for the batched rows (index buffers are + // i32 in F32-sized slots). + let routed_batch = match state.moe.as_ref() { + Some(_) if b > 1 => { + let (k, ne, mi) = ( + cfg.top_k_experts, + cfg.num_experts, + cfg.moe_intermediate_size, + ); + let m_total = hipfire_dispatch::pipeline::sandwich::RoutedBatch::m_total(b, k, ne); + let block_m = hipfire_dispatch::families::moe::MOE_GROUPED_BLOCK_M; + let mut a = |n: usize, label: &str| bufs.alloc(n, label); + Some([ + a(b * dim, "moe input")?, + a(b * dim, "moe input_rot")?, + a(b * ne, "moe router_logits")?, + a(b * k, "moe topk_indices")?, + a(b * k, "moe topk_weights")?, + a(b * k, "moe inverse_perm")?, + a(ne, "moe counts")?, + a(ne + 1, "moe offsets")?, + a(m_total, "moe sorted_slots")?, + a(m_total / block_m, "moe tile_ids")?, + a(m_total * (2 * mi).max(dim), "moe grouped rows")?, + a(b * k * mi, "moe gate")?, + a(b * k * mi, "moe up")?, + a(b * dim, "moe out")?, + ]) + } + _ => None, + }; + let routed = state.moe.as_ref().map(|m| { + let mut r = crate::program::routed_scratch(m); + r.batch = + routed_batch + .as_ref() + .map(|t| hipfire_dispatch::pipeline::sandwich::RoutedBatch { + input: &t[0], + input_rot: &t[1], + router_logits: &t[2], + topk_indices: &t[3], + topk_weights: &t[4], + inverse_perm: &t[5], + counts: &t[6], + offsets: &t[7], + sorted_slots: &t[8], + tile_ids: &t[9], + gate_up_grouped: &t[10], + gate: &t[11], + up: &t[12], + out: &t[13], + }); + r + }); let binding = ProgramBinding { geo: Geometry::eager(cfg), resident: eager_resident(state), rows: b, position: start_pos, positions: Some(&acts.positions), - scratch: acts.scratch(ple, state.moe.as_ref().map(crate::program::routed_scratch)), + scratch: acts.scratch(ple, routed), }; let mut steps = Vec::with_capacity(4 * cfg.n_layers); for (layer_idx, layer) in weights.layers.iter().enumerate() { diff --git a/crates/hipfire-arch-gemma4/src/gemma4.rs b/crates/hipfire-arch-gemma4/src/gemma4.rs index b3943f41fa..ccb7634176 100644 --- a/crates/hipfire-arch-gemma4/src/gemma4.rs +++ b/crates/hipfire-arch-gemma4/src/gemma4.rs @@ -15,7 +15,11 @@ use hipfire_runtime::weight_backend::load_embedding; use hipfire_runtime::weight_store::upload_pooled_bytes; use rdna_compute::{DType, Gpu, GpuTensor}; -pub(crate) const GEMMA4_FORWARD_BATCH_MAX: usize = 64; +/// Rows per batched forward. +pub(crate) const GEMMA4_FORWARD_BATCH_MAX: usize = 256; +/// Query rows the flash partials are sized for; batched attention launches +/// split larger batches into sub-batches of this many rows. +const FLASH_PARTIALS_ROWS: usize = 64; // ───────────────────────── Allocation ledger ───────────────────────── @@ -1169,7 +1173,7 @@ impl Gemma4State { l.gpu, &crate::program::Geometry::eager(cfg), max_seq, - GEMMA4_FORWARD_BATCH_MAX, + FLASH_PARTIALS_ROWS, ); if tier != FullKvTier::Q8 { partials = partials.max(asym3_flash_partials_len(cfg, max_seq)); diff --git a/crates/hipfire-arch-gemma4/src/program.rs b/crates/hipfire-arch-gemma4/src/program.rs index cebe0c55b6..403eb9f8e0 100644 --- a/crates/hipfire-arch-gemma4/src/program.rs +++ b/crates/hipfire-arch-gemma4/src/program.rs @@ -377,6 +377,7 @@ pub(crate) fn routed_scratch(s: &MoeScratch) -> RoutedScratch<'_> { act: &s.hidden_batch, out: &s.cur_moe, dense_normed: &s.cur_mlp, + batch: None, } } diff --git a/crates/hipfire-dispatch/src/pipeline/sandwich.rs b/crates/hipfire-dispatch/src/pipeline/sandwich.rs index 4169218f62..7d56d1a209 100644 --- a/crates/hipfire-dispatch/src/pipeline/sandwich.rs +++ b/crates/hipfire-dispatch/src/pipeline/sandwich.rs @@ -142,9 +142,43 @@ pub fn gemm_rows( DType::F32 => gpu .gemm_f32_batched(w.buf, x, y, w.m, w.k, rows) .map_err(hip), - DType::Q8_0 if *EAGLE_STRICT => gpu - .gemm_q8_0_batched(w.buf, x, y, w.m, w.k, rows) - .map_err(hip), + DType::Q8_0 if *EAGLE_STRICT => { + // The scalar kernel holds at most 64 rows. + for off in (0..rows).step_by(64) { + let n = (rows - off).min(64); + gpu.gemm_q8_0_batched( + w.buf, + &x.sub_offset(off * w.k, n * w.k), + &y.sub_offset(off * w.m, n * w.m), + w.m, + w.k, + n, + ) + .map_err(hip)?; + } + Ok(()) + } + // gfx1151: the 4-warp 64x64 tile reads each weight tile once per 64 + // rows (the 16x16 kernel re-reads it per 16 rows); a sub-64 row tail + // takes the 16x16 kernel. + DType::Q8_0 if gpu.arch_caps.is_gfx1151() && w.m % 64 == 0 && rows >= 64 => { + let full = rows / 64 * 64; + gpu.gemm_q8_0_wmma_4w(w.buf, &x.sub_offset(0, full * w.k), y, w.m, w.k, full) + .map_err(hip)?; + if full < rows { + let tail = rows - full; + gpu.gemm_q8_0_wmma( + w.buf, + &x.sub_offset(full * w.k, tail * w.k), + &y.sub_offset(full * w.m, tail * w.m), + w.m, + w.k, + tail, + ) + .map_err(hip)?; + } + Ok(()) + } DType::Q8_0 => gpu .gemm_q8_0_batched_chunked(w.buf, x, y, w.m, w.k, rows) .map_err(hip), @@ -158,6 +192,11 @@ pub fn gemm_rows( gpu.gemm_mq6g256_batched_lmhead(w.buf, x_rot, y, w.m, w.k, rows) .map_err(hip) } + DType::MQ4G256V2 => { + rotate_x_mq_rows(gpu, w, x, x_rot, rows)?; + gpu.gemm_mq4g256v2(w.buf, x_rot, y, w.m, w.k, rows) + .map_err(hip) + } // BF16 teachers (calibration): stage the F32 input to BF16 for the // MFMA GEMM, which resolves only where the arch has one. // ponytail: per-call staging buffer; hold one in scratch if BF16 @@ -192,7 +231,13 @@ pub fn gemm_rows( pub fn supports_batched_projection(dtype: DType) -> bool { matches!( dtype, - DType::F32 | DType::BF16 | DType::Q8_0 | DType::MQ4G256 | DType::HFQ4G256 | DType::MQ6G256 + DType::F32 + | DType::BF16 + | DType::Q8_0 + | DType::MQ4G256 + | DType::HFQ4G256 + | DType::MQ6G256 + | DType::MQ4G256V2 ) } @@ -656,8 +701,22 @@ impl SandwichAttentionOp<'_> { .map_err(hip)?; } // Prefill is not tree verification: the native causal kernels, and the - // windowed tile kernel for sliding layers. - if kv.window > 0 { + // windowed tile kernel for sliding layers. gfx1151 takes the + // GQA-shared WMMA flash kernel for every layer. + if gpu.attention_q8_0_prefill_gqa_wmma_admitted(n_heads, n_kv, hd) { + gpu.attention_q8_0_prefill_gqa_wmma( + self.q, + kv.k_cache, + kv.v_cache, + self.attn_out, + positions, + n_heads, + n_kv, + hd, + rows, + kv.window as i32, + ) + } else if kv.window > 0 { gpu.attention_flash_q8_0_batched_masked_windowed( self.q, kv.k_cache, @@ -869,6 +928,54 @@ impl RoutedExperts<'_> { | DType::Q8_0 ) && matches!(down, DType::Q8_0 | DType::HFQ4G128) } + + /// Formats whose batched rows take the grouped (sorted-by-expert) gate/up + /// GEMM: each expert's weights are read once per batch instead of once + /// per routed row. + pub fn grouped_rows(gate_up: DType, down: DType) -> bool { + matches!(gate_up, DType::MQ4G256V2 | DType::MQ6G256) && down == DType::HFQ4G128 + } +} + +/// Batched-rows scratch of the grouped routed-expert route (`rows` rows, +/// `slots = rows * top_k`, `m_total = slots + n_experts * MOE_GROUPED_BLOCK_M`). +/// Index buffers hold i32. +#[derive(Clone, Copy)] +pub struct RoutedBatch<'a> { + /// `[rows, hidden]`: router input, then `pre_norm(residual)`. + pub input: &'a GpuTensor, + /// `[rows, hidden]` FWHT rotation (router `x_rot` scratch first). + pub input_rot: &'a GpuTensor, + /// `[rows, n_experts]`. + pub router_logits: &'a GpuTensor, + /// `[slots]` each. + pub topk_indices: &'a GpuTensor, + pub topk_weights: &'a GpuTensor, + pub inverse_perm: &'a GpuTensor, + /// `[n_experts]`, `[n_experts + 1]`. + pub counts: &'a GpuTensor, + pub offsets: &'a GpuTensor, + /// `[m_total]`, `[m_total / MOE_GROUPED_BLOCK_M]`. + pub sorted_slots: &'a GpuTensor, + pub tile_ids: &'a GpuTensor, + /// `[m_total, max(2 * hidden_dim, hidden)]`: grouped gate/up rows, then + /// grouped down rows. + pub gate_up_grouped: &'a GpuTensor, + /// `[slots, hidden_dim]` each; `gate` ends as the activation. + pub gate: &'a GpuTensor, + pub up: &'a GpuTensor, + /// `[rows, hidden]` combined expert output. + pub out: &'a GpuTensor, +} + +impl RoutedBatch<'_> { + /// Grouped rows (with padding) for `rows` rows: a whole number of + /// 16-slot tiles, because the grouped kernels launch `ceil(m_total / 16)` + /// tiles and read one tile id each. + pub fn m_total(rows: usize, top_k: usize, n_experts: usize) -> usize { + let block = crate::families::moe::MOE_GROUPED_BLOCK_M; + (rows * top_k).next_multiple_of(block) + n_experts * block + } } /// Scratch of one routed-expert pass. @@ -889,6 +996,8 @@ pub struct RoutedScratch<'a> { pub out: &'a GpuTensor, /// `dense_post_norm(dense MLP output)`. pub dense_normed: &'a GpuTensor, + /// Grouped-route scratch for batched rows; `None` routes row by row. + pub batch: Option>, } /// Parallel dense+MoE MLP sublayer: @@ -931,14 +1040,24 @@ pub fn execute_parallel_moe_mlp( gpu.add_f32(r.dense_normed, r.out, s.normed).map_err(hip)?; s.norm(gpu, s.normed, mlp.post_norm, s.normed)?; } else { - // The indexed expert kernels are single-row: route row by row and - // accumulate into the normed dense output. s.norm(gpu, mlp.out, op.dense_post_norm, mlp.out)?; - for row in 0..s.rows { - let residual = s.residual.sub_offset(row * s.hidden, s.hidden); - routed_row(gpu, ctx, e, r, &residual, s.hidden, s.eps)?; - let out = mlp.out.sub_offset(row * s.hidden, s.hidden); - gpu.add_inplace_f32(&out, r.out).map_err(hip)?; + match r.batch.filter(|_| { + gpu.arch_caps.is_gfx1151() && RoutedExperts::grouped_rows(e.gate_up_dtype, e.down_dtype) + }) { + Some(b) => { + routed_rows_grouped(gpu, ctx, e, &b, s)?; + gpu.add_inplace_f32(mlp.out, b.out).map_err(hip)?; + } + // The indexed expert kernels are single-row: route row by row + // and accumulate into the normed dense output. + None => { + for row in 0..s.rows { + let residual = s.residual.sub_offset(row * s.hidden, s.hidden); + routed_row(gpu, ctx, e, r, &residual, s.hidden, s.eps)?; + let out = mlp.out.sub_offset(row * s.hidden, s.hidden); + gpu.add_inplace_f32(&out, r.out).map_err(hip)?; + } + } } s.norm(gpu, mlp.out, mlp.post_norm, s.normed)?; } @@ -946,6 +1065,132 @@ pub fn execute_parallel_moe_mlp( gpu.add_inplace_f32(s.x, s.normed).map_err(hip) } +/// `b.out = experts.post_norm(moe(pre_norm(residual)))` for all `s.rows` +/// rows: batched router, then the routed slots sorted by expert so each +/// expert's gate/up and down weights are read once per 16-slot tile. +fn routed_rows_grouped( + gpu: &mut Gpu, + ctx: &DispatchCtx, + e: &RoutedExperts<'_>, + b: &RoutedBatch<'_>, + s: SandwichStream<'_>, +) -> Result<(), DispatchError> { + use crate::families::moe::MOE_GROUPED_BLOCK_M; + let (rows, hidden, mi, k) = (s.rows, s.hidden, e.hidden_dim, e.top_k); + gpu.rmsnorm_batched(s.residual, e.router_norm, b.input, rows, hidden, s.eps) + .map_err(hip)?; + scale(gpu, b.input, e.router_input_scale)?; + // Expert selection is tie-sensitive: the F16-input Q8 WMMA GEMM moves + // router logits by ~1e-3, which flips near-tied experts and costs + // prefill KL 0.16 -> 1.19 on 26B-A4B. The batched Q8 GEMV keeps F32 + // activations and shares each weight load across eight rows. + if e.router.dtype == DType::Q8_0 { + gpu.gemm_q8_0_batched_wide_exact( + e.router.buf, + b.input, + b.router_logits, + e.router.m, + e.router.k, + rows, + ) + .map_err(hip)?; + } else { + gemm_rows( + gpu, + ctx, + &e.router, + b.input, + b.router_logits, + b.input_rot, + rows, + )?; + } + gpu.moe_softmax_topk_renorm_k8_batched( + b.router_logits, + b.topk_indices, + b.topk_weights, + e.n_experts, + true, + rows, + ) + .map_err(hip)?; + + gpu.rmsnorm_batched(s.residual, e.pre_norm, b.input, rows, hidden, s.eps) + .map_err(hip)?; + gpu.rotate_x_mq_batched(b.input, b.input_rot, hidden, rows) + .map_err(hip)?; + let m_total = RoutedBatch::m_total(rows, k, e.n_experts); + gpu.moe_scatter_fused_k8( + b.topk_indices, + b.counts, + b.offsets, + b.sorted_slots, + b.tile_ids, + b.inverse_perm, + rows * k, + e.n_experts, + m_total, + MOE_GROUPED_BLOCK_M, + ) + .map_err(hip)?; + super::dispatch_grouped_gemm( + gpu, + e.gate_up_dtype, + None, + e.gate_up_ptrs, + b.tile_ids, + b.sorted_slots, + b.input_rot, + b.gate_up_grouped, + 2 * mi, + hidden, + k, + m_total, + rows, + false, + false, + false, + )?; + gpu.moe_gate_up_unscatter_k8( + b.gate_up_grouped, + b.sorted_slots, + b.gate, + b.up, + mi, + k, + m_total, + ) + .map_err(hip)?; + gpu.gelu_tanh_f32(b.gate, b.gate, rows * k * mi) + .map_err(hip)?; + gpu.mul_f32(b.gate, b.up, b.gate).map_err(hip)?; + gpu.gemv_hfq4g128_moe_down_grouped( + e.down_ptrs, + b.tile_ids, + b.sorted_slots, + e.per_expert_scale, + b.gate, + b.gate_up_grouped, + hidden, + mi, + m_total, + ) + .map_err(hip)?; + gpu.zero_f32(b.out).map_err(hip)?; + gpu.moe_down_combine_grouped_k8( + b.gate_up_grouped, + b.inverse_perm, + b.topk_weights, + b.out, + hidden, + k, + rows, + ) + .map_err(hip)?; + gpu.rmsnorm_batched(b.out, e.post_norm, b.out, rows, hidden, s.eps) + .map_err(hip) +} + /// `r.out = experts.post_norm(moe(pre_norm(residual)))` for one row; the /// router reads `router_norm(residual) * router_input_scale`. fn routed_row( @@ -1182,6 +1427,7 @@ mod tests { DType::MQ4G256, DType::HFQ4G256, DType::MQ6G256, + DType::MQ4G256V2, ] { assert!(supports_batched_projection(dtype)); } diff --git a/crates/rdna-compute/Cargo.toml b/crates/rdna-compute/Cargo.toml index 8fc8c4cb59..c6bf26d967 100644 --- a/crates/rdna-compute/Cargo.toml +++ b/crates/rdna-compute/Cargo.toml @@ -899,3 +899,7 @@ required-features = ["lab"] [[example]] name = "cb_rows_twins_gate" required-features = ["deltanet"] + +[[example]] +name = "gemma4_kernel_f64_check" +required-features = ["lab"] diff --git a/crates/rdna-compute/examples/gemma4_kernel_f64_check.rs b/crates/rdna-compute/examples/gemma4_kernel_f64_check.rs new file mode 100644 index 0000000000..610b3819a5 --- /dev/null +++ b/crates/rdna-compute/examples/gemma4_kernel_f64_check.rs @@ -0,0 +1,310 @@ +// SPDX-License-Identifier: MIT OR Apache-2.0 +// Copyright (c) 2026 Björn Bösel +// hipfire — see LICENSE and NOTICE in the project root. +// +// The Gemma 4 26B-A4B gfx1151 kernels against an f64 CPU reference, at the +// model's shapes (head_dim 256 / 16q:8kv / window 1024 sliding, head_dim 512 / +// 16q:2kv full; routed experts M 2816, K 704, top-8): +// * decode Q8_0 flash attention (`attention_flash_q8_0_windowed`); run with +// HIPFIRE_GFX1151_Q8_DECODE_ATTN_GQA_GEMMA=0 to compare the reference tile; +// * batched prefill attention (`attention_q8_0_prefill_gqa_wmma`); +// * the HFQ4-G128 routed down projection, per-row and grouped; +// plus eager timings of the decode attention at 290 and 32768 keys. Exit 1 +// when a relative error exceeds its bound (1e-4 F32 routes, 2e-3 F16 WMMA). +// Run: cargo run --release -p rdna-compute --example gemma4_kernel_f64_check --features lab + +use rdna_compute::{DType, Gpu}; + +struct Rng(u64); +impl Rng { + fn next(&mut self) -> f64 { + self.0 = self.0.wrapping_mul(6364136223846793005).wrapping_add(1442695040888963407); + ((self.0 >> 11) as f64) / ((1u64 << 53) as f64) + } +} + +fn f16_bits(x: f32) -> u16 { + half_from_f32(x) +} + +// Minimal round-to-nearest f32 -> f16 for normal positive ranges used here. +fn half_from_f32(x: f32) -> u16 { + let b = x.to_bits(); + let sign = ((b >> 16) & 0x8000) as u16; + let exp = ((b >> 23) & 0xff) as i32 - 127 + 15; + let man = b & 0x7f_ffff; + let m = (man + 0x1000) >> 13; + sign | ((exp as u16) << 10) + m as u16 +} + +fn f16_to_f64(h: u16) -> f64 { + let exp = ((h >> 10) & 0x1f) as i32; + let man = (h & 0x3ff) as f64; + let v = (1.0 + man / 1024.0) * 2f64.powi(exp - 15); + if h & 0x8000 != 0 { -v } else { v } +} + +fn main() { + let mut gpu = Gpu::init().expect("gpu init"); + let mut ok = true; + for &(nh, nkv, hd, window, seq) in &[ + (16usize, 8usize, 256usize, 1024i32, 3000usize), + (16, 8, 256, 1024, 700), + (16, 2, 512, 0, 3000), + (16, 2, 512, 0, 129), + ] { + let max_seq = 4096usize; + let mut rng = Rng(seq as u64 * 31 + hd as u64); + let nb = hd / 32; + let row = nkv * nb * 34; + let mut kv = [vec![0u8; max_seq * row], vec![0u8; max_seq * row]]; + for cache in kv.iter_mut() { + for blk in cache.chunks_mut(34) { + let s = f16_bits((0.01 + 0.02 * rng.next()) as f32); + blk[0] = s as u8; + blk[1] = (s >> 8) as u8; + for c in &mut blk[2..] { + *c = ((rng.next() * 254.0) as i32 - 127) as i8 as u8; + } + } + } + let qh: Vec = (0..nh * hd).map(|_| (rng.next() * 2.0 - 1.0) as f32).collect(); + let deq = |cache: &[u8], t: usize, kh: usize, d: usize| -> f64 { + let off = t * row + (kh * nb + d / 32) * 34; + f16_to_f64(u16::from_le_bytes([cache[off], cache[off + 1]])) + * (cache[off + 2 + d % 32] as i8) as f64 + }; + let lo = if window > 0 { seq.saturating_sub(window as usize) } else { 0 }; + let scale = 1.0 / (hd as f64).sqrt(); + let mut want = vec![0f64; nh * hd]; + for h in 0..nh { + let kh = h / (nh / nkv); + let sc: Vec = (lo..seq) + .map(|t| (0..hd).map(|d| qh[h * hd + d] as f64 * deq(&kv[0], t, kh, d)).sum::() * scale) + .collect(); + let m = sc.iter().cloned().fold(f64::MIN, f64::max); + let e: Vec = sc.iter().map(|s| (s - m).exp()).collect(); + let z: f64 = e.iter().sum(); + for d in 0..hd { + want[h * hd + d] = + (lo..seq).zip(&e).map(|(t, w)| w * deq(&kv[1], t, kh, d)).sum::() / z; + } + } + let q = gpu.upload_f32(&qh, &[nh * hd]).unwrap(); + let k = gpu.upload_raw(&kv[0], &[kv[0].len()]).unwrap(); + let v = gpu.upload_raw(&kv[1], &[kv[1].len()]).unwrap(); + let out = gpu.zeros(&[nh * hd], DType::F32).unwrap(); + let pos = ((seq - 1) as i32).to_le_bytes(); + let pos_t = gpu.upload_raw(&pos, &[1]).unwrap(); + let partials = gpu.zeros(&[nh * max_seq.div_ceil(16) * (2 + hd)], DType::F32).unwrap(); + gpu.attention_flash_q8_0_windowed( + &q, &k, &v, &out, &pos_t.buf, seq, nh, nkv, hd, max_seq, &partials, window, + ) + .unwrap(); + let got = gpu.download_f32(&out).unwrap(); + let peak = want.iter().fold(0f64, |a, b| a.max(b.abs())); + let err = got.iter().zip(&want).fold(0f64, |a, (g, w)| a.max((*g as f64 - w).abs())); + let rel = err / peak; + println!("hd={hd} q:kv={nh}:{nkv} window={window} seq={seq}: max|err|/peak = {rel:.3e}"); + ok &= rel < 1e-4; + } + // Batched prefill attention (GQA WMMA route on gfx11) against f64. + for &(nh, nkv, hd, window, seq, nb) in &[ + (16usize, 8usize, 256usize, 1024i32, 3000usize, 256usize), + (16, 2, 512, 0, 3000, 200), + (16, 2, 512, 0, 40, 40), + (16, 8, 256, 1024, 70, 64), + ] { + let max_seq = 4096usize; + let mut rng = Rng(seq as u64 * 7 + hd as u64 + nb as u64); + let nbk = hd / 32; + let row = nkv * nbk * 34; + let mut kv = [vec![0u8; max_seq * row], vec![0u8; max_seq * row]]; + for cache in kv.iter_mut() { + for blk in cache.chunks_mut(34) { + let s = f16_bits((0.01 + 0.02 * rng.next()) as f32); + blk[0] = s as u8; + blk[1] = (s >> 8) as u8; + for c in &mut blk[2..] { + *c = ((rng.next() * 254.0) as i32 - 127) as i8 as u8; + } + } + } + let qh: Vec = (0..nb * nh * hd).map(|_| (rng.next() * 2.0 - 1.0) as f32 * 4.0).collect(); + let deq = |cache: &[u8], t: usize, kh: usize, d: usize| -> f64 { + let off = t * row + (kh * nbk + d / 32) * 34; + f16_to_f64(u16::from_le_bytes([cache[off], cache[off + 1]])) + * (cache[off + 2 + d % 32] as i8) as f64 + }; + let scale = 1.0 / (hd as f64).sqrt(); + let pos0 = seq - nb; + let mut want = vec![0f64; nb * nh * hd]; + for b in 0..nb { + let p = pos0 + b; + let lo = if window > 0 { (p + 1).saturating_sub(window as usize) } else { 0 }; + for h in 0..nh { + let kh = h / (nh / nkv); + let qo = (b * nh + h) * hd; + let sc: Vec = (lo..=p) + .map(|t| (0..hd).map(|d| qh[qo + d] as f64 * deq(&kv[0], t, kh, d)).sum::() * scale) + .collect(); + let m = sc.iter().cloned().fold(f64::MIN, f64::max); + let e: Vec = sc.iter().map(|s| (s - m).exp()).collect(); + let z: f64 = e.iter().sum(); + for d in 0..hd { + want[qo + d] = (lo..=p).zip(&e).map(|(t, w)| w * deq(&kv[1], t, kh, d)).sum::() / z; + } + } + } + let q = gpu.upload_f32(&qh, &[nb * nh * hd]).unwrap(); + let k = gpu.upload_raw(&kv[0], &[kv[0].len()]).unwrap(); + let v = gpu.upload_raw(&kv[1], &[kv[1].len()]).unwrap(); + let out = gpu.zeros(&[nb * nh * hd], DType::F32).unwrap(); + let pos: Vec = (pos0..seq).flat_map(|p| (p as i32).to_le_bytes()).collect(); + let pos_t = gpu.upload_raw(&pos, &[nb]).unwrap(); + gpu.attention_q8_0_prefill_gqa_wmma(&q, &k, &v, &out, &pos_t, nh, nkv, hd, nb, window).unwrap(); + let got = gpu.download_f32(&out).unwrap(); + let peak = want.iter().fold(0f64, |a, b| a.max(b.abs())); + let err = got.iter().zip(&want).fold(0f64, |a, (g, w)| a.max((*g as f64 - w).abs())); + println!("prefill attn hd={hd} q:kv={nh}:{nkv} window={window} seq={seq} rows={nb}: max|err|/peak = {:.3e}", err / peak); + ok &= err / peak < 2e-3; + } + // HFQ4-G128 routed down (M 2816, K 704, k8) against f64. + { + let (m, k, n_exp) = (2816usize, 704usize, 16usize); + let gpr = k.div_ceil(128); + let mut rng = Rng(7); + let mut experts = Vec::new(); + let mut host = Vec::new(); + for _ in 0..n_exp { + let mut bytes = vec![0u8; m * gpr * 72]; + for grp in bytes.chunks_mut(72) { + grp[0..4].copy_from_slice(&((0.01 + 0.02 * rng.next()) as f32).to_le_bytes()); + grp[4..8].copy_from_slice(&((rng.next() * 0.2 - 0.1) as f32).to_le_bytes()); + for b in &mut grp[8..] { + *b = (rng.next() * 256.0) as u8; + } + } + experts.push(gpu.upload_raw(&bytes, &[bytes.len()]).unwrap()); + host.push(bytes); + } + let ptrs: Vec = experts.iter().flat_map(|t| (t.buf.as_ptr() as u64).to_le_bytes()).collect(); + let ptrs_t = gpu.upload_raw(&ptrs, &[ptrs.len()]).unwrap(); + let idx: Vec = vec![3, 0, 15, 7, 9, 1, 12, 5]; + let idx_t = gpu.upload_raw(&idx.iter().flat_map(|v| v.to_le_bytes()).collect::>(), &[8]).unwrap(); + let wts: Vec = (0..8).map(|_| rng.next() as f32).collect(); + let pes: Vec = (0..n_exp).map(|_| (0.5 + rng.next()) as f32).collect(); + let hid: Vec = (0..8 * k).map(|_| (rng.next() * 2.0 - 1.0) as f32).collect(); + let wts_t = gpu.upload_f32(&wts, &[8]).unwrap(); + let pes_t = gpu.upload_f32(&pes, &[n_exp]).unwrap(); + let hid_t = gpu.upload_f32(&hid, &[8 * k]).unwrap(); + let out = gpu.zeros(&[m], DType::F32).unwrap(); + gpu.gemv_hfq4g128_moe_down_residual_scaled_k8_indexed( + &ptrs_t, &idx_t, &wts_t, &pes_t, &hid_t, &out, m, k, + ) + .unwrap(); + let got = gpu.download_f32(&out).unwrap(); + let mut want = vec![0f64; m]; + for (r, w) in want.iter_mut().enumerate() { + for (kr, &e) in idx.iter().enumerate() { + let b = &host[e as usize]; + let mut acc = 0f64; + for d in 0..k { + let off = (r * gpr + d / 128) * 72; + let sc = f32::from_le_bytes(b[off..off + 4].try_into().unwrap()) as f64; + let z = f32::from_le_bytes(b[off + 4..off + 8].try_into().unwrap()) as f64; + let byte = b[off + 8 + (d % 128) / 2]; + let q = if d % 2 == 0 { byte & 0xf } else { byte >> 4 } as f64; + acc += (sc * q + z) * hid[kr * k + d] as f64; + } + *w += wts[kr] as f64 * pes[e as usize] as f64 * acc; + } + } + let peak = want.iter().fold(0f64, |a, b| a.max(b.abs())); + let err = got.iter().zip(&want).fold(0f64, |a, (g, w)| a.max((*g as f64 - w).abs())); + println!("moe down hfq4g128 m={m} k={k}: max|err|/peak = {:.3e}", err / peak); + ok &= err / peak < 1e-4; + gpu.hip.device_synchronize().unwrap(); + let t = std::time::Instant::now(); + for _ in 0..50 { + gpu.gemv_hfq4g128_moe_down_residual_scaled_k8_indexed( + &ptrs_t, &idx_t, &wts_t, &pes_t, &hid_t, &out, m, k, + ) + .unwrap(); + } + gpu.hip.device_synchronize().unwrap(); + println!("time moe down: {:.1} us/call", t.elapsed().as_secs_f64() * 1e6 / 50.0); + // Grouped: 4 tiles of 16 slots; tile 1 unused, tile 3 half padding. + let tiles: Vec = vec![3, -1, 9, 15]; + let mut slots = vec![-1i32; 64]; + let n_act = 40usize; + let mut r2 = Rng(11); + for (t, &e) in tiles.iter().enumerate() { + if e < 0 { continue; } + let n = if t == 3 { 8 } else { 16 }; + for j in 0..n { slots[t * 16 + j] = ((r2.next() * n_act as f64) as i32).min(n_act as i32 - 1); } + } + let act: Vec = (0..n_act * k).map(|_| (r2.next() * 2.0 - 1.0) as f32).collect(); + let i32b = |v: &[i32]| v.iter().flat_map(|x| x.to_le_bytes()).collect::>(); + let tiles_t = gpu.upload_raw(&i32b(&tiles), &[4]).unwrap(); + let slots_t = gpu.upload_raw(&i32b(&slots), &[64]).unwrap(); + let act_t = gpu.upload_f32(&act, &[n_act * k]).unwrap(); + let y_t = gpu.zeros(&[64 * m], DType::F32).unwrap(); + gpu.gemv_hfq4g128_moe_down_grouped(&ptrs_t, &tiles_t, &slots_t, &pes_t, &act_t, &y_t, m, k, 64).unwrap(); + let y = gpu.download_f32(&y_t).unwrap(); + let (mut err, mut peak) = (0f64, 0f64); + for (gr, &sl) in slots.iter().enumerate() { + let e = tiles[gr / 16]; + if sl < 0 || e < 0 { continue; } + let bw = &host[e as usize]; + for r in 0..m { + let mut acc = 0f64; + for d in 0..k { + let off = (r * gpr + d / 128) * 72; + let sc = f32::from_le_bytes(bw[off..off + 4].try_into().unwrap()) as f64; + let z = f32::from_le_bytes(bw[off + 4..off + 8].try_into().unwrap()) as f64; + let byte = bw[off + 8 + (d % 128) / 2]; + let q = if d % 2 == 0 { byte & 0xf } else { byte >> 4 } as f64; + acc += (sc * q + z) * act[sl as usize * k + d] as f64; + } + let want = pes[e as usize] as f64 * acc; + peak = peak.max(want.abs()); + err = err.max((y[gr * m + r] as f64 - want).abs()); + } + } + println!("moe down grouped: max|err|/peak = {:.3e}", err / peak); + // gfx11 takes the F16-WMMA route (F16 weights and activations). + ok &= err / peak < 1e-3; + } + // Timing at the 26B decode shapes (max_seq 34816, as the harness). + for &(nh, nkv, hd, window) in &[(16usize, 8usize, 256usize, 1024i32), (16, 2, 512, 0)] { + let max_seq = 34816usize; + let row = nkv * hd / 32 * 34; + let k = gpu.zeros(&[max_seq * row / 4], DType::F32).unwrap(); + let v = gpu.zeros(&[max_seq * row / 4], DType::F32).unwrap(); + let q = gpu.zeros(&[nh * hd], DType::F32).unwrap(); + let out = gpu.zeros(&[nh * hd], DType::F32).unwrap(); + let partials = gpu.zeros(&[nh * max_seq.div_ceil(16) * (2 + hd)], DType::F32).unwrap(); + for seq in [290usize, 32768] { + let pos_t = gpu.upload_raw(&((seq - 1) as i32).to_le_bytes(), &[1]).unwrap(); + let mut run = |gpu: &mut Gpu| { + gpu.attention_flash_q8_0_windowed( + &q, &k, &v, &out, &pos_t.buf, seq, nh, nkv, hd, max_seq, &partials, window, + ) + .unwrap() + }; + run(&mut gpu); + gpu.hip.device_synchronize().unwrap(); + let t = std::time::Instant::now(); + for _ in 0..50 { + run(&mut gpu); + } + gpu.hip.device_synchronize().unwrap(); + println!( + "time hd={hd} window={window} seq={seq}: {:.1} us/call (tile+reduce)", + t.elapsed().as_secs_f64() * 1e6 / 50.0 + ); + } + } + std::process::exit(if ok { 0 } else { 1 }); +} diff --git a/crates/rdna-compute/src/attention.rs b/crates/rdna-compute/src/attention.rs index 77e81c3cdf..047533eca9 100644 --- a/crates/rdna-compute/src/attention.rs +++ b/crates/rdna-compute/src/attention.rs @@ -397,6 +397,36 @@ fn q8_decode_attn_gqa_gfx1151_admitted( && hipfire_config::developer_bool("HIPFIRE_GFX1151_Q8_DECODE_ATTN_GQA", true) } +/// gfx1151 Q8_0 decode attention for the Gemma 4 26B-A4B shapes through a +/// GQA-shared, 8-wave tile (`attention_flash_q8_0_tile_gqa_gemma.gfx1151.hip`): +/// head_dim 256 with GQA group 2 (sliding layers) and head_dim 512 with group +/// 8 (full layers), tile 128, any window. Same partials as the reference tile +/// and the same reduce, not bit-identical (different summation order). +/// `HIPFIRE_GFX1151_Q8_DECODE_ATTN_GQA_GEMMA=0` restores the reference tile. +fn q8_decode_attn_gqa_gemma_gfx1151( + gpu: &Gpu, + n_heads: usize, + n_kv_heads: usize, + head_dim: usize, + tile_size: usize, +) -> Option<&'static str> { + if !gpu.arch_caps.is_gfx1151() + || tile_size != 128 + || n_kv_heads == 0 + || !hipfire_config::developer_bool("HIPFIRE_GFX1151_Q8_DECODE_ATTN_GQA_GEMMA", true) + { + return None; + } + match (head_dim, n_heads / n_kv_heads, n_heads % n_kv_heads) { + (256, 2, 0) => Some("attention_flash_q8_0_tile_gqa_d256g2_gfx1151"), + (512, 8, 0) => Some("attention_flash_q8_0_tile_gqa_d512g8_gfx1151"), + _ => None, + } +} + +/// Grid-y cap of the Gemma GQA tile; its workgroups loop over the tiles. +const Q8_DECODE_GQA_GEMMA_GRID_Y_CAP: usize = 64; + /// Kernel pair of a GQA-shared decode attention route: a 13-arg tile with one /// 256-thread workgroup per (kv head, tile) and a 7-arg head-dim-split reduce. struct DecodeGqaPair { @@ -9243,6 +9273,100 @@ impl Gpu { /// method). Used by cohere2moe's `sliding_attention` layers so context beyond /// `sliding_window` clips correctly instead of running full-causal (degraded). #[allow(clippy::too_many_arguments)] + /// Whether [`Self::attention_q8_0_prefill_gqa_wmma`] admits this shape. + pub fn attention_q8_0_prefill_gqa_wmma_admitted( + &self, + n_heads: usize, + n_kv_heads: usize, + head_dim: usize, + ) -> bool { + // Measured on gfx1151 only; the kernel is gfx11 wave32 WMMA. + self.arch_caps.is_gfx1151() + && matches!(head_dim, 256 | 512) + && n_kv_heads > 0 + && n_heads % n_kv_heads == 0 + && matches!(n_heads / n_kv_heads, 1 | 2 | 4 | 8 | 16) + && hipfire_config::developer_bool("HIPFIRE_Q8_PREFILL_GQA_WMMA", true) + } + + /// Causal attention of `batch` consecutive query rows (`positions`, i32, + /// increasing) over a position-indexed Q8_0 KV cache, GQA-shared, on + /// gfx11 WMMA (see `attention_q8_0_prefill_gqa_wmma.gfx11.hip`). + /// `window > 0` limits each row to its last `window` keys. + #[allow(clippy::too_many_arguments)] + pub fn attention_q8_0_prefill_gqa_wmma( + &mut self, + q: &GpuTensor, + k_cache: &GpuTensor, + v_cache: &GpuTensor, + out: &GpuTensor, + positions: &GpuTensor, + n_heads: usize, + n_kv_heads: usize, + head_dim: usize, + batch: usize, + window: i32, + ) -> HipResult<()> { + if !self.attention_q8_0_prefill_gqa_wmma_admitted(n_heads, n_kv_heads, head_dim) { + return Err(hip_bridge::HipError::new( + 0, + &format!( + "attention_q8_0_prefill_gqa_wmma: unsupported shape heads={n_heads} kv={n_kv_heads} hd={head_dim} on {}", + self.arch + ), + )); + } + self.bind_thread()?; + let func = if head_dim == 256 { + "attention_q8_0_prefill_gqa_wmma_d256" + } else { + "attention_q8_0_prefill_gqa_wmma_d512" + }; + self.ensure_kernel( + "attention_q8_0_prefill_gqa_wmma", + kernels::ATTENTION_Q8_0_PREFILL_GQA_WMMA_GFX11_SRC, + func, + )?; + let g = n_heads / n_kv_heads; + let qp = q.buf.as_ptr(); + let kp = k_cache.buf.as_ptr(); + let vp = v_cache.buf.as_ptr(); + let op = out.buf.as_ptr(); + let pp = positions.buf.as_ptr(); + let nh = n_heads as i32; + let nkv = n_kv_heads as i32; + let b = batch as i32; + let sc = 1.0f32 / (head_dim as f32).sqrt(); + let wn = window; + let mut params: Vec<*mut c_void> = vec![ + &qp as *const _ as *mut c_void, + &kp as *const _ as *mut c_void, + &vp as *const _ as *mut c_void, + &op as *const _ as *mut c_void, + &pp as *const _ as *mut c_void, + &nh as *const _ as *mut c_void, + &nkv as *const _ as *mut c_void, + &b as *const _ as *mut c_void, + &sc as *const _ as *mut c_void, + &wn as *const _ as *mut c_void, + ]; + let grid = [(batch * g).div_ceil(16) as u32, n_kv_heads as u32, 1]; + self.launch_maybe_blob(func, grid, [128, 1, 1], 0, &mut params, || { + let mut kb = hip_bridge::KernargBlob::new(); + kb.push_ptr(qp); + kb.push_ptr(kp); + kb.push_ptr(vp); + kb.push_ptr(op); + kb.push_ptr(pp); + kb.push_i32(nh); + kb.push_i32(nkv); + kb.push_i32(b); + kb.push_f32(sc); + kb.push_i32(wn); + kb + }) + } + pub fn attention_flash_q8_0_batched_masked_windowed( &mut self, q: &GpuTensor, @@ -9539,21 +9663,34 @@ impl Gpu { } // ── Tile kernel ── + let gemma_gqa = if output_gate.is_none() { + q8_decode_attn_gqa_gemma_gfx1151(self, n_heads, n_kv_heads, head_dim, tile_size) + } else { + None + }; let gfx1151_tile_dpp = self.arch_caps.is_gfx1151() && hipfire_config::developer_var("HIPFIRE_GFX1151_ATTENTION_TILE_DPP").as_deref() == Ok("1"); - let (tile_module, tile_src) = if gfx1151_tile_dpp { + let (tile_module, tile_src, tile_fn) = if let Some(func) = gemma_gqa { + ( + "attention_flash_q8_0_tile_gqa_gemma_gfx1151", + kernels::ATTENTION_FLASH_Q8_0_TILE_GQA_GEMMA_GFX1151_SRC, + func, + ) + } else if gfx1151_tile_dpp { ( "attention_flash_q8_0_tile_dpp_gfx1151", kernels::ATTENTION_FLASH_Q8_0_TILE_DPP_GFX1151_SRC, + "attention_flash_q8_0_tile", ) } else { ( "attention_flash_q8_0_tile", kernels::ATTENTION_FLASH_Q8_0_TILE_SRC, + "attention_flash_q8_0_tile", ) }; - self.ensure_kernel(tile_module, tile_src, "attention_flash_q8_0_tile")?; + self.ensure_kernel(tile_module, tile_src, tile_fn)?; { let scale = 1.0f32 / (head_dim as f32).sqrt(); let q_ptr = q.buf.as_ptr(); @@ -9569,8 +9706,18 @@ impl Gpu { let ts = tile_size as i32; let wn = window; let es = effective_seq_arg; - let grid = [n_heads as u32, launch_tiles as u32, 1]; - let shared = ((tile_size + head_dim) * 4) as u32; + // The GQA route runs one 256-thread workgroup per (kv head, + // tile) and loops over tiles, so its grid is capped. + let (grid, block, shared) = if gemma_gqa.is_some() { + let y = launch_tiles.min(Q8_DECODE_GQA_GEMMA_GRID_Y_CAP) as u32; + ([n_kv_heads as u32, y, 1], [256, 1, 1], 0) + } else { + ( + [n_heads as u32, launch_tiles as u32, 1], + [32, 1, 1], + ((tile_size + head_dim) * 4) as u32, + ) + }; let mut params: Vec<*mut c_void> = vec![ &q_ptr as *const _ as *mut c_void, &k_ptr as *const _ as *mut c_void, @@ -9587,9 +9734,9 @@ impl Gpu { &es as *const _ as *mut c_void, ]; self.launch_maybe_blob_position_grid( - "attention_flash_q8_0_tile", + tile_fn, grid, - [32, 1, 1], + block, shared, &mut params, 1, diff --git a/crates/rdna-compute/src/gemm.rs b/crates/rdna-compute/src/gemm.rs index 151bbb6e59..18edb45e28 100644 --- a/crates/rdna-compute/src/gemm.rs +++ b/crates/rdna-compute/src/gemm.rs @@ -29409,12 +29409,14 @@ impl Gpu { kernels::GEMM_Q8_0_WMMA_4W_SRC, "gemm_q8_0_wmma_4w", )?; - // Stage F32 → F16 input if needed. + // Stage F32 → F16 input if needed. Uncached: callers reuse one + // activation buffer with new contents per layer, which the + // pointer-keyed `ensure_fp16_x` would hand back stale. let xp_owned = x.buf.as_ptr(); let mut xp = if matches!(x.dtype, DType::F16) { xp_owned } else { - self.ensure_fp16_x(x, batch_size * k)? + self.convert_fp16_x_uncached(x, batch_size * k)? }; let mut a_p = a.buf.as_ptr(); @@ -31734,7 +31736,9 @@ impl Gpu { &m_val as *const _ as *mut c_void, &k_val as *const _ as *mut c_void, ]; - self.launch_maybe_blob(KERNEL, [m as u32, 1, 1], [32, 1, 1], 0, &mut params, || { + // Eight rows per 256-thread workgroup, one wave per row. + let grid = [m.div_ceil(8) as u32, 1, 1]; + self.launch_maybe_blob(KERNEL, grid, [256, 1, 1], 0, &mut params, || { let mut b = hip_bridge::KernargBlob::new(); b.push_ptr(pp); b.push_ptr(ip); @@ -31747,6 +31751,78 @@ impl Gpu { b }) } + /// Grouped HFQ4-G128 MoE down: `y[r, :] = per_expert_scale[e] * W_e · + /// act[sorted_slots[r]]` for every grouped row `r` of the 16-slot tiles + /// `tile_ids` (the `moe_scatter_fused_k8` layout); padding rows are not + /// written. `m_total` is a multiple of 16. + #[allow(clippy::too_many_arguments)] + pub fn gemv_hfq4g128_moe_down_grouped( + &mut self, + expert_ptrs: &GpuTensor, + tile_ids: &GpuTensor, + sorted_slots: &GpuTensor, + per_expert_scale: &GpuTensor, + act: &GpuTensor, + y: &GpuTensor, + m: usize, + k: usize, + m_total: usize, + ) -> HipResult<()> { + self.bind_thread()?; + // gfx1151: wave32 WMMA (F16 activations); the scalar kernel elsewhere. + let wmma = self.arch_caps.is_gfx1151() && k % 16 == 0; + let (kernel, src, grid, block): (&str, &str, [u32; 3], u32) = if wmma { + ( + "gemm_hfq4g128_moe_down_grouped_wmma", + kernels::GEMM_HFQ4G128_MOE_DOWN_GROUPED_WMMA_SRC, + [m.div_ceil(16) as u32, (m_total / 16) as u32, 1], + 32, + ) + } else { + ( + "gemv_hfq4g128_moe_down_grouped", + kernels::GEMV_HFQ4G128_MOE_DOWN_GROUPED_SRC, + [m.div_ceil(8) as u32, (m_total / 16) as u32, 1], + 256, + ) + }; + self.ensure_kernel(kernel, src, kernel)?; + let pp = expert_ptrs.buf.as_ptr(); + let tp = tile_ids.buf.as_ptr(); + let sp = sorted_slots.buf.as_ptr(); + let ep = per_expert_scale.buf.as_ptr(); + // Uncached: one activation buffer is reused with new contents per layer. + let ap = if wmma { + self.convert_fp16_x_uncached(act, act.numel())? + } else { + act.buf.as_ptr() + }; + let yp = y.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, + &tp as *const _ as *mut c_void, + &sp as *const _ as *mut c_void, + &ep as *const _ as *mut c_void, + &ap 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(kernel, grid, [block, 1, 1], 0, &mut params, || { + let mut b = hip_bridge::KernargBlob::new(); + b.push_ptr(pp); + b.push_ptr(tp); + b.push_ptr(sp); + b.push_ptr(ep); + b.push_ptr(ap); + b.push_ptr(yp); + b.push_i32(m_val); + b.push_i32(k_val); + b + }) + } #[allow(unused_variables)] pub fn gemv_hfq4g128_moe_down_residual_scaled_k8_indexed_batched( &mut self, diff --git a/crates/rdna-compute/src/kernel_registry.rs b/crates/rdna-compute/src/kernel_registry.rs index 4bfad58f0b..cd5796594e 100644 --- a/crates/rdna-compute/src/kernel_registry.rs +++ b/crates/rdna-compute/src/kernel_registry.rs @@ -472,6 +472,7 @@ pub fn entries(arch: &str, extra_flags: &str) -> Result, Regist } if arch == "gfx1151" { add!("attention_flash_q8_0_tile_gqa_gfx1151", kernels::ATTENTION_FLASH_Q8_0_TILE_GQA_GFX1151_SRC, ["attention_flash_q8_0_tile_gqa_gfx1151"]); + add!("attention_flash_q8_0_tile_gqa_gemma_gfx1151", kernels::ATTENTION_FLASH_Q8_0_TILE_GQA_GEMMA_GFX1151_SRC, ["attention_flash_q8_0_tile_gqa_d256g2_gfx1151", "attention_flash_q8_0_tile_gqa_d512g8_gfx1151"]); add!("attention_flash_reduce_dsplit_gfx1151", kernels::ATTENTION_FLASH_REDUCE_DSPLIT_GFX1151_SRC, ["attention_flash_reduce_dsplit_gfx1151"]); add!("attention_q8_0_fa2_gqa_gfx1151", kernels::ATTENTION_Q8_0_FA2_GQA_GFX1151_SRC, ["attention_fa2_q_preconvert_gfx1151", "attention_q8_0_fa2_gqa_gfx1151"]); add!("conv1d_silu_split_qknorm_b256", kernels::CONV1D_SILU_SPLIT_QKNORM_B256_SRC, ["conv1d_silu_split_qknorm_b256"]); @@ -1795,7 +1796,10 @@ mod tests { // four `sample_top_p_parallel*` modules' penalty prepass now finds // each token's (last index, count) by an LDS sort; the per-token // update and every result bit are unchanged (sampler oracle dumps). - const REPINNED_SINCE_P0: [&str; 18] = [ + const REPINNED_SINCE_P0: [&str; 21] = [ + "attention_flash_fp8_e4m3_tile", + "attention_flash_q8_0_reduce", + "attention_flash_q8_0_tile", "conv1d_silu_split_qknorm_b256", "fused_qk_l2_norm_scale", "fused_rmsnorm_mq_rotate", diff --git a/crates/rdna-compute/src/kernels.rs b/crates/rdna-compute/src/kernels.rs index 4431c23012..1923f7f7f8 100644 --- a/crates/rdna-compute/src/kernels.rs +++ b/crates/rdna-compute/src/kernels.rs @@ -7942,6 +7942,12 @@ pub const ATTENTION_FLASH_Q8_0_REDUCE_GATED_MQ_ROTATE_AWQ_DEC_GFX1100_SRC: &str /// GQA group 6, tile 128, window 0. gfx1151 only. pub const ATTENTION_FLASH_Q8_0_TILE_GQA_GFX1151_SRC: &str = include_str!("../../../kernels/src/attention_flash_q8_0_tile_gqa.gfx1151.hip"); +/// Gemma 4 26B-A4B GQA-shared Q8_0 decode flash tiles +/// (`attention_flash_q8_0_tile_gqa_d256g2_gfx1151`, `..._d512g8_gfx1151`): same +/// 13-arg ABI and partials as [`ATTENTION_FLASH_Q8_0_TILE_SRC`], any window. +/// gfx1151 only. +pub const ATTENTION_FLASH_Q8_0_TILE_GQA_GEMMA_GFX1151_SRC: &str = + include_str!("../../../kernels/src/attention_flash_q8_0_tile_gqa_gemma.gfx1151.hip"); /// gfx1151 twin of [`ATTENTION_FLASH_REDUCE_DSPLIT_GFX1201_SRC`] /// (`attention_flash_reduce_dsplit_gfx1151`): byte-identical to /// `attention_flash_q8_0_reduce` on the same partials, one workgroup per @@ -9807,6 +9813,16 @@ pub const GEMV_Q8_0_MOE_DOWN_RESIDUAL_SCALED_K8_INDEXED_SRC: &str = include_str!("../../../kernels/src/gemv_q8_0_moe_down_residual_scaled_k8_indexed.hip"); pub const GEMV_HFQ4G128_MOE_DOWN_RESIDUAL_SCALED_K8_INDEXED_SRC: &str = include_str!("../../../kernels/src/gemv_hfq4g128_moe_down_residual_scaled_k8_indexed.hip"); +/// Grouped (sorted-by-expert, 16-slot tiles) HFQ4-G128 MoE down projection. +pub const GEMV_HFQ4G128_MOE_DOWN_GROUPED_SRC: &str = + include_str!("../../../kernels/src/gemv_hfq4g128_moe_down_grouped.hip"); +/// GQA-shared Q8_0 prefill attention on gfx11 wave32 WMMA, optional sliding +/// window (`attention_q8_0_prefill_gqa_wmma_d256` / `_d512`). +pub const ATTENTION_Q8_0_PREFILL_GQA_WMMA_GFX11_SRC: &str = + include_str!("../../../kernels/src/attention_q8_0_prefill_gqa_wmma.gfx11.hip"); +/// gfx11 wave32-WMMA twin of [`GEMV_HFQ4G128_MOE_DOWN_GROUPED_SRC`] (F16 X). +pub const GEMM_HFQ4G128_MOE_DOWN_GROUPED_WMMA_SRC: &str = + include_str!("../../../kernels/src/gemm_hfq4g128_moe_down_grouped_wmma.hip"); // ─── Gemma 4 hd512 attention + KV write kernels ───────────────────────── // (ROPE_PARTIAL_HALVED_SRC / LOGIT_SOFTCAP_SRC already defined above.) diff --git a/crates/rdna-compute/src/replay.rs b/crates/rdna-compute/src/replay.rs index ebfb9c8e17..43a6bbd5a9 100644 --- a/crates/rdna-compute/src/replay.rs +++ b/crates/rdna-compute/src/replay.rs @@ -1372,7 +1372,9 @@ fn pointer_effects(kernel: &str) -> Option> { "attention_flash_q8_0_tile" | "attention_flash_fp8_e4m3_tile_gqa_gfx1201" | "attention_flash_q8_0_tile_gqa_gfx1100" - | "attention_flash_q8_0_tile_gqa_gfx1151" => { + | "attention_flash_q8_0_tile_gqa_gfx1151" + | "attention_flash_q8_0_tile_gqa_d256g2_gfx1151" + | "attention_flash_q8_0_tile_gqa_d512g8_gfx1151" => { Some(vec![read(0), read(8), read(16), write(24), read(32)]) } "attention_flash_q8_0_reduce" @@ -1802,6 +1804,8 @@ fn expected_kernarg_bytes(kernel: &str) -> Option { | "gated_norm_mq_rotate_awq_k6144_gfx1201" | "attention_flash_q8_0_tile_gqa_gfx1100" | "attention_flash_q8_0_tile_gqa_gfx1151" + | "attention_flash_q8_0_tile_gqa_d256g2_gfx1151" + | "attention_flash_q8_0_tile_gqa_d512g8_gfx1151" | "fused_qkv_hfq4g256" | "fused_qkv_mq4g256v2" | "fused_qkv_mq4g256v2_k2048_x_buffer_gfx1100" diff --git a/crates/rdna-compute/tests/fixtures/p0-gfx1201-sha256.tsv b/crates/rdna-compute/tests/fixtures/p0-gfx1201-sha256.tsv index e54ccf5584..6720f88d28 100644 --- a/crates/rdna-compute/tests/fixtures/p0-gfx1201-sha256.tsv +++ b/crates/rdna-compute/tests/fixtures/p0-gfx1201-sha256.tsv @@ -5,9 +5,9 @@ alpha_gate 95ad80057a0ea52c318f6870d24527e6da269ab651764d71c1919edfe46693fa attention_flash_asym3_tile ea142beec2099d32d4c6c87e64dab7daa6f1d4084379f5475a8f93724f4fb91c attention_flash_asym3_tile_batched fa4aab5558cce1ce1df6d58b2d41aef5fb02a69b05b022c55cb55f9cecfe7e5d attention_flash_asym_reduce_batched c5aa6b1c05f04b38f108ea69d8c4c7e14312075db72e457e1072692eeaab8506 -attention_flash_fp8_e4m3_tile 71841378051aa32449defd37f69865e667e927c22a8aadbfbb28d07ce2f575a7 -attention_flash_q8_0_reduce a67f131323361b80c649e88fa6db50750f20be09a7a13bccaadc96587dd36aa8 -attention_flash_q8_0_tile db10a96ed529bc7b261f2f0840e85d74b3dd992f37b5e5226d72eea1cff198bf +attention_flash_fp8_e4m3_tile 0ccb71c046876f6102fea18760de2c466c98f8f69d930b17c62c364d566877b0 +attention_flash_q8_0_reduce da3ec1a9d7df02280de15e74fc0c0eb0bdd463627b3edc88881e3c5d825a0855 +attention_flash_q8_0_tile 727ad7c3ef1785c31fb699326062f10f3c25777659f9ce570b74ebc999c93d95 attention_fp8_e4m3_kv_batched 2aafbc2b321b4ecf32f10c69b9d7ccf83ef83be4694c0df73dd2bcc3ee9c2e04 attention_q8_0_flash_prefill bac4c3e44d2480e80e240c87901318438cdced7b6d96c3bcf372eb321e5744c5 attention_q8_0_kv da2bf0b122879de37427875d24d4cc3370e55e9202c1bfc9144880300e0d9458 diff --git a/kernels/src/attention_flash_q8_0_reduce.hip b/kernels/src/attention_flash_q8_0_reduce.hip index 6abe07afe8..1a07ff4c83 100644 --- a/kernels/src/attention_flash_q8_0_reduce.hip +++ b/kernels/src/attention_flash_q8_0_reduce.hip @@ -48,15 +48,25 @@ __global__ void attention_flash_q8_0_reduce( extern __shared__ float corr_lut[]; // [max_tiles] __shared__ float global_max_shared; __shared__ float global_sum_shared; + __shared__ float wave_max[32]; - // Preserve the reference kernel's increasing-tile fmax order exactly. + // Global max over tiles with a positive sum. fmaxf is exact and + // order-independent, so a parallel reduction equals the reference's + // increasing-tile serial order; it avoids a serial chain of global loads. + float local_max = -1e30f; + for (int tile = tid; tile < n_tiles; tile += blockDim.x) { + const float* p = partials + (h * max_tiles + tile) * stride; + if (p[1] > 0.0f) + local_max = fmaxf(local_max, p[0]); + } + for (int off = warpSize / 2; off > 0; off >>= 1) + local_max = fmaxf(local_max, __shfl_xor(local_max, off)); + if ((tid % warpSize) == 0) wave_max[tid / warpSize] = local_max; + __syncthreads(); if (tid == 0) { float global_max = -1e30f; - for (int tile = 0; tile < n_tiles; tile++) { - const float* p = partials + (h * max_tiles + tile) * stride; - if (p[1] > 0.0f) - global_max = fmaxf(global_max, p[0]); - } + for (int w = 0; w < (int)((blockDim.x + warpSize - 1) / warpSize); w++) + global_max = fmaxf(global_max, wave_max[w]); global_max_shared = global_max; } __syncthreads(); @@ -69,10 +79,21 @@ __global__ void attention_flash_q8_0_reduce( } __syncthreads(); - // Preserve the reference kernel's increasing-tile global-sum order. + // Preserve the reference kernel's increasing-tile global-sum order; the + // p[1] loads are issued eight at a time so they overlap. if (tid == 0) { float global_sum = 0.0f; - for (int tile = 0; tile < n_tiles; tile++) { + int tile = 0; + for (; tile + 8 <= n_tiles; tile += 8) { + float s[8]; +#pragma unroll + for (int j = 0; j < 8; j++) + s[j] = partials[(h * max_tiles + tile + j) * stride + 1]; +#pragma unroll + for (int j = 0; j < 8; j++) + global_sum += s[j] * corr_lut[tile + j]; + } + for (; tile < n_tiles; tile++) { const float* p = partials + (h * max_tiles + tile) * stride; global_sum += p[1] * corr_lut[tile]; } @@ -83,9 +104,26 @@ __global__ void attention_flash_q8_0_reduce( const float inv_sum = global_sum_shared > 0.0f ? 1.0f / global_sum_shared : 0.0f; for (int d = tid; d < head_dim; d += blockDim.x) { float g = 0.0f; - for (int tile = 0; tile < n_tiles; tile++) { + // Zero-weight tiles (empty or fully masked) add exactly 0; their V + // partials may be unwritten, so do not read them. Loads are issued + // eight tiles at a time; the adds keep increasing-tile order. + int tile = 0; + for (; tile + 8 <= n_tiles; tile += 8) { + float c[8], v[8]; +#pragma unroll + for (int j = 0; j < 8; j++) c[j] = corr_lut[tile + j]; +#pragma unroll + for (int j = 0; j < 8; j++) + v[j] = c[j] != 0.0f ? partials[(h * max_tiles + tile + j) * stride + 2 + d] : 0.0f; +#pragma unroll + for (int j = 0; j < 8; j++) + if (c[j] != 0.0f) g += v[j] * c[j]; + } + for (; tile < n_tiles; tile++) { + const float c = corr_lut[tile]; + if (c == 0.0f) continue; const float* p = partials + (h * max_tiles + tile) * stride; - g += p[2 + d] * corr_lut[tile]; + g += p[2 + d] * c; } out[h * head_dim + d] = g * inv_sum; } diff --git a/kernels/src/attention_flash_q8_0_tile.hip b/kernels/src/attention_flash_q8_0_tile.hip index 7923b056d0..888585798e 100644 --- a/kernels/src/attention_flash_q8_0_tile.hip +++ b/kernels/src/attention_flash_q8_0_tile.hip @@ -112,6 +112,17 @@ __global__ void attention_flash_q8_0_tile( const int tile_end = min(tile_start + tile_size, seq_len); if (tile_start >= seq_len) return; const int tile_len = tile_end - tile_start; + // A tile wholly below the window contributes nothing: publish an empty + // partial (sum 0, which the reduce skips) instead of streaming its V rows + // with zero weights. Same output as the full pass. + if (tile_end <= win_lo) { + if (threadIdx.x == 0) { + float* pe = partials + (h * ((max_seq + tile_size - 1) / tile_size) + tile_id) * (2 + head_dim); + pe[0] = -1e30f; + pe[1] = 0.0f; + } + return; + } const int kv_group = n_heads / n_kv_heads; const int kv_h = h / kv_group; diff --git a/kernels/src/attention_flash_q8_0_tile_gqa_gemma.gfx1151.hip b/kernels/src/attention_flash_q8_0_tile_gqa_gemma.gfx1151.hip new file mode 100644 index 0000000000..9e8b8742c1 --- /dev/null +++ b/kernels/src/attention_flash_q8_0_tile_gqa_gemma.gfx1151.hip @@ -0,0 +1,232 @@ +// SPDX-License-Identifier: MIT OR Apache-2.0 +// Copyright (c) 2026 Björn Bösel +// hipfire — see LICENSE and NOTICE in the project root. + +// GQA-shared Q8_0 decode flash tile for the Gemma 4 26B-A4B shapes (gfx1151): +// head_dim 256 / GQA group 2 (sliding layers, window > 0) and head_dim 512 / +// GQA group 8 (full layers). One 256-thread workgroup per (kv head, tile of +// 128 keys) serves every q head of its kv head, so each K/V row is read once +// instead of once per q head, and a tile is spread over 8 waves instead of one. +// +// Same 13-argument ABI, window semantics and `[n_heads, max_tiles, 2 + hd]` +// f32 partials as `attention_flash_q8_0_tile`, consumed by the unchanged +// `attention_flash_q8_0_reduce`. Not bit-identical to it: the dot products +// and V sums associate differently. Tiles wholly below the window publish an +// empty partial (sum 0), which the reduce skips. Workgroup y serves tiles +// y, y + gridDim.y, ... so the host can cap the grid. +#include + +#if defined(__HIP_DEVICE_COMPILE__) && !defined(__gfx1151__) +#error "attention_flash_q8_0_tile_gqa_gemma requires --offload-arch=gfx1151" +#endif + +#define GG_TILE 128 +#define GG_THREADS 256 +#define GG_BLK 34 + +// Unaligned 32-bit load (RDNA3 unaligned access mode), as the reference +// tile reads its codes. +__device__ __forceinline__ unsigned int gg_u32(const unsigned char* p) { + return *(const unsigned int*)p; +} + +__device__ __forceinline__ float gg_f16(const unsigned char* p) { + return (float)*(const _Float16*)p; +} + +__device__ __forceinline__ float gg_code(unsigned int w, int i) { + return (float)(int)(signed char)(w >> (8 * i)); +} + +template +__device__ __forceinline__ void gg_tile( + const float* __restrict__ q_lds, const unsigned char* __restrict__ k_cache, + const unsigned char* __restrict__ v_cache, float* __restrict__ partials, float* __restrict__ sc, + int n_kv_heads, int max_seq, int max_tiles, float scale_attn, int window, int seq_len, + int win_lo, int kv_h, int tile_id) { + constexpr int NB = HD / 32; // Q8 blocks per kv head + constexpr int NBH = NB / 2; // blocks per scoring half + constexpr int DPT = HD / GG_THREADS; // V dims per thread + const int tid = threadIdx.x; + const int lane = tid & 31; + const int wave = tid >> 5; + const int tile_start = tile_id * GG_TILE; + const int tile_len = min(GG_TILE, seq_len - tile_start); + const int h0 = kv_h * G; + const size_t row_stride = (size_t)n_kv_heads * NB * GG_BLK; + const int head_blk = kv_h * NB; + + // ── Scores: thread pair (2t, 2t+1) owns key t; each covers half of HD. + { + const int t = tid >> 1; + const int half = tid & 1; + const int tc = min(t, tile_len - 1); + const int key = tile_start + tc; + const int slot = window > 0 ? key % max_seq : key; + const unsigned char* kr = k_cache + (size_t)slot * row_stride + (size_t)(head_blk + half * NBH) * GG_BLK; + float acc[G]; +#pragma unroll + for (int g = 0; g < G; ++g) acc[g] = 0.0f; +#pragma unroll 2 + for (int b = 0; b < NBH; ++b) { + const unsigned char* blk = kr + b * GG_BLK; + const float ks = gg_f16(blk); + unsigned int w[8]; +#pragma unroll + for (int i = 0; i < 8; ++i) w[i] = gg_u32(blk + 2 + 4 * i); + float a[32]; +#pragma unroll + for (int i = 0; i < 8; ++i) { +#pragma unroll + for (int j = 0; j < 4; ++j) a[4 * i + j] = ks * gg_code(w[i], j); + } + const int d0 = (half * NBH + b) * 32; +#pragma unroll + for (int g = 0; g < G; ++g) { + const float4* qv = (const float4*)(q_lds + g * HD + d0); + float s = 0.0f; +#pragma unroll + for (int i = 0; i < 8; ++i) { + const float4 x = qv[i]; + s = __builtin_fmaf(x.x, a[4 * i + 0], s); + s = __builtin_fmaf(x.y, a[4 * i + 1], s); + s = __builtin_fmaf(x.z, a[4 * i + 2], s); + s = __builtin_fmaf(x.w, a[4 * i + 3], s); + } + acc[g] += s; + } + } +#pragma unroll + for (int g = 0; g < G; ++g) acc[g] += __shfl_xor(acc[g], 1); + if (half == 0) { + const bool valid = t < tile_len && tile_start + t >= win_lo; +#pragma unroll + for (int g = 0; g < G; ++g) sc[g * GG_TILE + t] = valid ? acc[g] * scale_attn : -INFINITY; + } + } + __syncthreads(); + + // ── Softmax: one wave per head (G <= 8 waves); lane holds 4 keys. + if (wave < G) { + float* s_h = sc + wave * GG_TILE; + float m = -1e30f; +#pragma unroll + for (int i = 0; i < GG_TILE / 32; ++i) m = fmaxf(m, s_h[lane + 32 * i]); + for (int off = 16; off > 0; off >>= 1) m = fmaxf(m, __shfl_xor(m, off)); + float sum = 0.0f; +#pragma unroll + for (int i = 0; i < GG_TILE / 32; ++i) { + const float e = expf(s_h[lane + 32 * i] - m); + s_h[lane + 32 * i] = e; + sum += e; + } + for (int off = 16; off > 0; off >>= 1) sum += __shfl_xor(sum, off); + if (lane == 0) { + float* p = partials + ((size_t)(h0 + wave) * max_tiles + tile_id) * (2 + HD); + p[0] = m; + p[1] = sum; + } + } + __syncthreads(); + + // ── V: thread owns DPT consecutive dims for all G heads. + const int d = tid * DPT; + const int vblk = head_blk + d / 32; + const int coff = 2 + (d & 31); + float out[G][DPT]; +#pragma unroll + for (int g = 0; g < G; ++g) +#pragma unroll + for (int j = 0; j < DPT; ++j) out[g][j] = 0.0f; + const int t_lo = max(0, win_lo - tile_start); +#pragma unroll 4 + for (int t = t_lo; t < tile_len; ++t) { + const int key = tile_start + t; + const int slot = window > 0 ? key % max_seq : key; + const unsigned char* vb = v_cache + (size_t)slot * row_stride + (size_t)vblk * GG_BLK; + const float vs = gg_f16(vb); + float a[DPT]; + if constexpr (DPT == 2) { + const unsigned int c = *(const unsigned short*)(vb + coff); + a[0] = vs * gg_code(c, 0); + a[1] = vs * gg_code(c, 1); + } else { + a[0] = vs * (float)(int)(signed char)vb[coff]; + } +#pragma unroll + for (int g = 0; g < G; ++g) { + const float w = sc[g * GG_TILE + t]; +#pragma unroll + for (int j = 0; j < DPT; ++j) out[g][j] = __builtin_fmaf(w, a[j], out[g][j]); + } + } +#pragma unroll + for (int g = 0; g < G; ++g) { + float* p = partials + ((size_t)(h0 + g) * max_tiles + tile_id) * (2 + HD) + 2 + d; +#pragma unroll + for (int j = 0; j < DPT; ++j) p[j] = out[g][j]; + } +} + +template +__device__ __forceinline__ void gg_kernel( + const float* __restrict__ q, const unsigned char* __restrict__ k_cache, + const unsigned char* __restrict__ v_cache, float* __restrict__ partials, + const int* __restrict__ pos_buf, int n_kv_heads, int max_seq, float scale_attn, int window) { + __shared__ __attribute__((aligned(16))) float q_lds[G * HD]; + __shared__ __attribute__((aligned(16))) float sc[G * GG_TILE]; + const int kv_h = blockIdx.x; + if (kv_h >= n_kv_heads) return; + const int logical_seq_len = pos_buf[0] + 1; + const int seq_len = window > 0 ? min(logical_seq_len, max_seq) : logical_seq_len; + const int win_lo = window > 0 ? seq_len - window : 0; + const int max_tiles = (max_seq + GG_TILE - 1) / GG_TILE; + const int n_tiles = min((seq_len + GG_TILE - 1) / GG_TILE, max_tiles); + const int tid = threadIdx.x; + const int h0 = kv_h * G; + + bool q_loaded = false; + for (int tile_id = blockIdx.y; tile_id < n_tiles; tile_id += gridDim.y) { + if ((tile_id + 1) * GG_TILE <= win_lo) { + // Wholly below the window: empty partial, skipped by the reduce. + if (tid < G) { + float* p = partials + ((size_t)(h0 + tid) * max_tiles + tile_id) * (2 + HD); + p[0] = -1e30f; + p[1] = 0.0f; + } + continue; + } + if (!q_loaded) { + for (int i = tid; i < G * HD / 4; i += GG_THREADS) + ((float4*)q_lds)[i] = ((const float4*)(q + (size_t)h0 * HD))[i]; + q_loaded = true; + } else { + __syncthreads(); // previous tile's sc readers are done + } + __syncthreads(); + gg_tile(q_lds, k_cache, v_cache, partials, sc, n_kv_heads, max_seq, max_tiles, + scale_attn, window, seq_len, win_lo, kv_h, tile_id); + } +} + +// Host admits head_dim 256 with n_heads == 2 * n_kv_heads, tile 128. +extern "C" __launch_bounds__(GG_THREADS, 1) +__global__ void attention_flash_q8_0_tile_gqa_d256g2_gfx1151( + const float* __restrict__ q, const unsigned char* __restrict__ k_cache, + const unsigned char* __restrict__ v_cache, float* __restrict__ partials, + const int* __restrict__ pos_buf, int n_heads, int n_kv_heads, int head_dim, int max_seq, + float scale_attn, int tile_size, int window, int effective_seq_len) { + (void)n_heads; (void)head_dim; (void)tile_size; (void)effective_seq_len; + gg_kernel<256, 2>(q, k_cache, v_cache, partials, pos_buf, n_kv_heads, max_seq, scale_attn, window); +} + +// Host admits head_dim 512 with n_heads == 8 * n_kv_heads, tile 128. +extern "C" __launch_bounds__(GG_THREADS, 1) +__global__ void attention_flash_q8_0_tile_gqa_d512g8_gfx1151( + const float* __restrict__ q, const unsigned char* __restrict__ k_cache, + const unsigned char* __restrict__ v_cache, float* __restrict__ partials, + const int* __restrict__ pos_buf, int n_heads, int n_kv_heads, int head_dim, int max_seq, + float scale_attn, int tile_size, int window, int effective_seq_len) { + (void)n_heads; (void)head_dim; (void)tile_size; (void)effective_seq_len; + gg_kernel<512, 8>(q, k_cache, v_cache, partials, pos_buf, n_kv_heads, max_seq, scale_attn, window); +} diff --git a/kernels/src/attention_q8_0_prefill_gqa_wmma.gfx11.hip b/kernels/src/attention_q8_0_prefill_gqa_wmma.gfx11.hip new file mode 100644 index 0000000000..3ef0735cbd --- /dev/null +++ b/kernels/src/attention_q8_0_prefill_gqa_wmma.gfx11.hip @@ -0,0 +1,228 @@ +// SPDX-License-Identifier: MIT OR Apache-2.0 +// Copyright (c) 2026 Björn Bösel +// hipfire — see LICENSE and NOTICE in the project root. +// +// Batched (prefill) causal attention over a Q8_0 KV cache on gfx11 wave32 +// WMMA, GQA-shared, optional sliding window. Flash-attention style: one pass +// over the keys with an online softmax, no partials. +// +// A workgroup (4 waves) owns one KV head and 16 "query rows": the G = n_heads +// / n_kv_heads q heads of 16 / G consecutive batch rows, so each K/V tile is +// read once for all of them. Row r of the tile is batch row 16/G * tile + +// r / G, q head G * kv_h + r % G. +// +// Per 32-key tile: +// S = Q K^T each wave covers a quarter of head_dim (Q in registers, +// K dequantized from the cache straight into the B operand); +// the four partial S tiles are summed in LDS. +// softmax online per row (scale, causal/window mask), P in f16 in LDS. +// O += P V each wave owns a quarter of head_dim of O; V is dequantized +// into LDS transposed ([dim][key]) so the B operand is +// contiguous. +// Q, K, P and V enter WMMA as f16; S, the softmax state and O are f32. +// +// Grid: [ceil(batch * G / 16), n_kv_heads]. Block: [128]. Requires +// head_dim in {256, 512}, G in {1, 2, 4, 8, 16}, K/V rows position-indexed +// (no ring), positions[] increasing within a launch. +#include + +typedef _Float16 __attribute__((ext_vector_type(16))) half16_t; +typedef float __attribute__((ext_vector_type(8))) float8_t; + +#define PG_KT 32 +#define PG_WAVES 4 +#define PG_THREADS (PG_WAVES * 32) + +__device__ __forceinline__ unsigned short pg_u16(const unsigned char* p) { + unsigned short v; + __builtin_memcpy(&v, p, 2); + return v; +} + +__device__ __forceinline__ float pg_f16(const unsigned char* p) { + return (float)__builtin_bit_cast(_Float16, pg_u16(p)); +} + +template +__device__ __forceinline__ void pg_body( + const float* __restrict__ q, const unsigned char* __restrict__ k_cache, + const unsigned char* __restrict__ v_cache, float* __restrict__ out, + const int* __restrict__ positions, int n_heads, int n_kv_heads, int batch, float scale, + int window) { + constexpr int NB = D / 32; // Q8 blocks per kv head + constexpr int DW = D / PG_WAVES; // dims per wave + constexpr int KS = DW / 16; // 16-dim k-steps per wave + __shared__ float s_part[PG_WAVES][16][PG_KT]; + __shared__ __attribute__((aligned(32))) _Float16 p_lds[16][PG_KT]; + __shared__ __attribute__((aligned(32))) _Float16 vt[D][PG_KT]; + __shared__ float alpha_lds[16]; + __shared__ float l_lds[16]; + + const int tid = threadIdx.x; + const int wave = tid >> 5; + const int lane = tid & 31; + const int kv_h = blockIdx.y; + const int g = n_heads / n_kv_heads; + const int rows_per_tile = 16 / g; + const int row0 = blockIdx.x * rows_per_tile; + const size_t row_stride = (size_t)n_kv_heads * NB * 34; + const size_t head_off = (size_t)kv_h * NB * 34; + + // Tile key range from the first and last valid batch rows. + const int last = min(row0 + rows_per_tile, batch) - 1; + const int pos_min = positions[row0]; + const int pos_max = positions[last]; + const int k_lo = window > 0 ? max(0, pos_min - window + 1) : 0; + const int k_hi = pos_max; + + // Q A-operand: lane row r = lane & 15, this wave's dims, as f16. + const int r_a = lane & 15; + const int brow_a = min(row0 + r_a / g, batch - 1); + const float* q_row = q + ((size_t)brow_a * n_heads + kv_h * g + r_a % g) * D + wave * DW; + half16_t qa[KS]; +#pragma unroll + for (int ks = 0; ks < KS; ++ks) { +#pragma unroll + for (int i = 0; i < 16; ++i) qa[ks][i] = (_Float16)q_row[ks * 16 + i]; + } + + // Softmax ownership: thread t owns row t / 8, keys 4 * (t % 8) .. +3. + const int sm_row = tid >> 3; + const int sm_k0 = (tid & 7) * 4; + const int sm_brow = row0 + sm_row / g; + const bool sm_valid = sm_brow < batch; + const int sm_pos = sm_valid ? positions[sm_brow] : -1; + float m_run = -1e30f; + float l_run = 0.0f; + + float8_t acc[KS]; +#pragma unroll + for (int c = 0; c < KS; ++c) acc[c] = float8_t{0, 0, 0, 0, 0, 0, 0, 0}; + + for (int kt0 = k_lo; kt0 <= k_hi; kt0 += PG_KT) { + // ── V tile, transposed into LDS (keys clamped into the range). + for (int p = tid; p < PG_KT * NB; p += PG_THREADS) { + const int j = p % PG_KT; + const int b = p / PG_KT; + const int t = min(kt0 + j, k_hi); + const unsigned char* blk = v_cache + (size_t)t * row_stride + head_off + b * 34; + const float vs = pg_f16(blk); +#pragma unroll + for (int i = 0; i < 32; i += 2) { + const unsigned short c2 = pg_u16(blk + 2 + i); + vt[b * 32 + i][j] = (_Float16)(vs * (float)(signed char)(c2 & 0xff)); + vt[b * 32 + i + 1][j] = (_Float16)(vs * (float)(signed char)(c2 >> 8)); + } + } + + // ── Partial S over this wave's dims, two 16-key blocks. +#pragma unroll + for (int kb = 0; kb < PG_KT / 16; ++kb) { + const int t = min(kt0 + kb * 16 + (lane & 15), k_hi); + const unsigned char* krow = k_cache + (size_t)t * row_stride + head_off; + float8_t s = {0, 0, 0, 0, 0, 0, 0, 0}; +#pragma unroll + for (int ks = 0; ks < KS; ++ks) { + const int d0 = wave * DW + ks * 16; + const unsigned char* blk = krow + (d0 / 32) * 34; + const float sc = pg_f16(blk); + const unsigned char* codes = blk + 2 + (d0 & 31); + half16_t kb16; +#pragma unroll + for (int i = 0; i < 16; i += 2) { + const unsigned short c2 = pg_u16(codes + i); + kb16[i] = (_Float16)(sc * (float)(signed char)(c2 & 0xff)); + kb16[i + 1] = (_Float16)(sc * (float)(signed char)(c2 >> 8)); + } + s = __builtin_amdgcn_wmma_f32_16x16x16_f16_w32(qa[ks], kb16, s); + } +#pragma unroll + for (int j = 0; j < 8; ++j) s_part[wave][2 * j + (lane >> 4)][kb * 16 + (lane & 15)] = s[j]; + } + __syncthreads(); + + // ── Online softmax. + { + float sv[4]; + float tmax = -1e30f; +#pragma unroll + for (int i = 0; i < 4; ++i) { + const int key = sm_k0 + i; + const int t = kt0 + key; + const bool ok = sm_valid && t <= sm_pos && (window <= 0 || t > sm_pos - window); + const float v = (s_part[0][sm_row][key] + s_part[1][sm_row][key] + + s_part[2][sm_row][key] + s_part[3][sm_row][key]) * scale; + sv[i] = ok ? v : -1e30f; + tmax = fmaxf(tmax, sv[i]); + } + tmax = fmaxf(tmax, __shfl_xor(tmax, 1)); + tmax = fmaxf(tmax, __shfl_xor(tmax, 2)); + tmax = fmaxf(tmax, __shfl_xor(tmax, 4)); + const float m_new = fmaxf(m_run, tmax); + const float alpha = __expf(m_run - m_new); + float psum = 0.0f; +#pragma unroll + for (int i = 0; i < 4; ++i) { + const float p = sv[i] > -1e29f ? __expf(sv[i] - m_new) : 0.0f; + p_lds[sm_row][sm_k0 + i] = (_Float16)p; + psum += p; + } + psum += __shfl_xor(psum, 1); + psum += __shfl_xor(psum, 2); + psum += __shfl_xor(psum, 4); + l_run = l_run * alpha + psum; + m_run = m_new; + if ((tid & 7) == 0) alpha_lds[sm_row] = alpha; + } + __syncthreads(); + + // ── Rescale O, then O += P V over this wave's dims. + { + float al[8]; +#pragma unroll + for (int j = 0; j < 8; ++j) al[j] = alpha_lds[2 * j + (lane >> 4)]; +#pragma unroll + for (int kb = 0; kb < PG_KT / 16; ++kb) { + const half16_t pa = *(const half16_t*)&p_lds[lane & 15][kb * 16]; +#pragma unroll + for (int c = 0; c < KS; ++c) { + if (kb == 0) { +#pragma unroll + for (int j = 0; j < 8; ++j) acc[c][j] *= al[j]; + } + const int d = wave * DW + c * 16 + (lane & 15); + const half16_t vb = *(const half16_t*)&vt[d][kb * 16]; + acc[c] = __builtin_amdgcn_wmma_f32_16x16x16_f16_w32(pa, vb, acc[c]); + } + } + } + __syncthreads(); + } + + if ((tid & 7) == 0) l_lds[sm_row] = l_run; + __syncthreads(); +#pragma unroll + for (int j = 0; j < 8; ++j) { + const int r = 2 * j + (lane >> 4); + const int brow = row0 + r / g; + if (brow >= batch) continue; + const float l = l_lds[r]; + const float inv = l > 0.0f ? 1.0f / l : 0.0f; + float* o = out + ((size_t)brow * n_heads + kv_h * g + r % g) * D + wave * DW; +#pragma unroll + for (int c = 0; c < KS; ++c) o[c * 16 + (lane & 15)] = acc[c][j] * inv; + } +} + +#define PG_ENTRY(NAME, D) \ + extern "C" __launch_bounds__(PG_THREADS, 1) __global__ void NAME( \ + const float* __restrict__ q, const unsigned char* __restrict__ k_cache, \ + const unsigned char* __restrict__ v_cache, float* __restrict__ out, \ + const int* __restrict__ positions, int n_heads, int n_kv_heads, int batch, float scale, \ + int window) { \ + pg_body(q, k_cache, v_cache, out, positions, n_heads, n_kv_heads, batch, scale, \ + window); \ + } + +PG_ENTRY(attention_q8_0_prefill_gqa_wmma_d256, 256) +PG_ENTRY(attention_q8_0_prefill_gqa_wmma_d512, 512) diff --git a/kernels/src/gemm_hfq4g128_moe_down_grouped_wmma.hip b/kernels/src/gemm_hfq4g128_moe_down_grouped_wmma.hip new file mode 100644 index 0000000000..5b9aa39aa4 --- /dev/null +++ b/kernels/src/gemm_hfq4g128_moe_down_grouped_wmma.hip @@ -0,0 +1,79 @@ +// SPDX-License-Identifier: MIT OR Apache-2.0 +// Copyright (c) 2026 Björn Bösel +// hipfire — see LICENSE and NOTICE in the project root. +// +// Grouped HFQ4-G128 MoE down projection on wave32 WMMA (gfx11). Same +// contract as `gemv_hfq4g128_moe_down_grouped`: 16-slot tiles sorted by +// expert (`moe_scatter_fused_k8` layout), `y[grouped_row, :] = +// per_expert_scale[e] * W_e · act[sorted_slots[grouped_row]]`, padding rows +// unwritten. Adapted from `gemm_paro_q4g128_moe_grouped_wmma_k2` for rows +// whose K is not a multiple of 128: a row holds ceil(K/128) 72-byte groups +// and the last group's tail beyond K is skipped (K % 16 == 0). Weights are +// dequantized in f32 and rounded once to f16; X is F16. +// +// Grid: [ceil(M/16), m_total/16]. Block: [32]. +#include +#include + + +typedef _Float16 __attribute__((ext_vector_type(16))) half16_t; +typedef float __attribute__((ext_vector_type(8))) float8_t; + +__launch_bounds__(32, 2) +extern "C" __global__ void gemm_hfq4g128_moe_down_grouped_wmma( + const unsigned long long* __restrict__ expert_ptrs, + const int* __restrict__ tile_ids, + const int* __restrict__ sorted_slots, + const float* __restrict__ per_expert_scale, + const _Float16* __restrict__ act, // [slots, K] + float* __restrict__ y, // [m_total, M] + int M, int K +) { + const int tid = threadIdx.x; + const int row_start = blockIdx.x * 16; + const int tile = blockIdx.y; + if (row_start >= M) return; + const int expert = tile_ids[tile]; + if (expert < 0) return; + const char* __restrict__ A = reinterpret_cast(expert_ptrs[expert]); + + const int lane = tid & 15; + const int slot = sorted_slots[tile * 16 + lane]; + const _Float16* x_base = slot >= 0 ? act + (long long)slot * K : nullptr; + const int safe_row = row_start + lane < M ? row_start + lane : M - 1; + const int groups_per_row = (K + 127) / 128; + const char* row_base = A + (long long)safe_row * groups_per_row * 72; + + float8_t acc = {0, 0, 0, 0, 0, 0, 0, 0}; + const half16_t zero_b = {0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0}; + + for (int g = 0; g < groups_per_row; g++) { + const char* gp = row_base + g * 72; + const float sc = __builtin_bit_cast(float, *(const unsigned int*)gp); + const float zp = __builtin_bit_cast(float, *(const unsigned int*)(gp + 4)); + const int kt_end = min(8, (K - g * 128) / 16); + for (int kt = 0; kt < kt_end; ++kt) { + const int k_off = kt * 16; + const unsigned int pk0 = *(const unsigned int*)(gp + 8 + k_off / 2); + const unsigned int pk1 = *(const unsigned int*)(gp + 8 + k_off / 2 + 4); + const half16_t b = x_base != nullptr + ? *(const half16_t*)(x_base + g * 128 + k_off) : zero_b; + half16_t a; +#pragma unroll + for (int i = 0; i < 8; ++i) { + a[i] = (_Float16)__builtin_fmaf(sc, (float)((pk0 >> (4 * i)) & 0xFu), zp); + a[8 + i] = (_Float16)__builtin_fmaf(sc, (float)((pk1 >> (4 * i)) & 0xFu), zp); + } + acc = __builtin_amdgcn_wmma_f32_16x16x16_f16_w32(a, b, acc); + } + } + + if (slot >= 0) { + const float es = per_expert_scale[expert]; +#pragma unroll + for (int j = 0; j < 8; j++) { + const int out_row = row_start + 2 * j + (tid >> 4); + if (out_row < M) y[(long long)(tile * 16 + lane) * M + out_row] = es * acc[j]; + } + } +} diff --git a/kernels/src/gemv_hfq4g128_moe_down_grouped.hip b/kernels/src/gemv_hfq4g128_moe_down_grouped.hip new file mode 100644 index 0000000000..2fb8989e0d --- /dev/null +++ b/kernels/src/gemv_hfq4g128_moe_down_grouped.hip @@ -0,0 +1,96 @@ +// SPDX-License-Identifier: MIT OR Apache-2.0 +// Copyright (c) 2026 Björn Bösel +// hipfire — see LICENSE and NOTICE in the project root. +// +// Grouped MoE down projection for HFQ4-G128 expert weights (batched rows). +// Routed slots arrive sorted by expert in 16-slot tiles (`moe_scatter_fused_k8` +// layout): tile t serves expert `tile_ids[t]` (-1 = unused) and grouped rows +// 16t .. 16t+15, whose activation row is `sorted_slots[row]` (-1 = padding). +// +// One 256-thread workgroup computes eight output rows (one wave each) for +// the sixteen slots of one tile, so each expert row is read once per tile +// instead of once per routed slot. A group (f32 scale, f32 zero, 64 nibble +// bytes = 128 values) is split over four lanes of 32 values. Output: +// y[grouped_row, row] = per_expert_scale[e] * (W_e · act[slot]); the mixture +// weights are applied by `moe_down_combine_grouped_k8`. +#include + +#define GD_ROWS 8 +#define GD_SLOTS 16 + +__launch_bounds__(256, 2) +extern "C" __global__ void gemv_hfq4g128_moe_down_grouped( + const unsigned long long* __restrict__ expert_ptrs, + const int* __restrict__ tile_ids, + const int* __restrict__ sorted_slots, + const float* __restrict__ per_expert_scale, + const float* __restrict__ act, // [slots, K] + float* __restrict__ y, // [m_total, M] + int M, int K +) { + const int tile = blockIdx.y; + const int expert = tile_ids[tile]; + if (expert < 0) return; + const int tid = threadIdx.x; + const int lane = tid & 31; + const int row = blockIdx.x * GD_ROWS + (tid >> 5); + if (row >= M) return; + + int slot[GD_SLOTS]; +#pragma unroll + for (int j = 0; j < GD_SLOTS; ++j) slot[j] = sorted_slots[tile * GD_SLOTS + j]; + + const int groups_per_row = (K + 127) / 128; + const char* wrow = reinterpret_cast(expert_ptrs[expert]) + + (long long)row * groups_per_row * 72; + float acc[GD_SLOTS]; +#pragma unroll + for (int j = 0; j < GD_SLOTS; ++j) acc[j] = 0.0f; + + for (int c = lane; c < groups_per_row * 4; c += 32) { + const int g = c >> 2; + const int base = g * 128 + (c & 3) * 32; + if (base >= K) continue; + const char* gptr = wrow + g * 72; + const float scale = __builtin_bit_cast(float, *(const unsigned int*)gptr); + const float zero = __builtin_bit_cast(float, *(const unsigned int*)(gptr + 4)); + const uint2* nib = (const uint2*)(gptr + 8 + (c & 3) * 16); + const uint2 w0 = nib[0]; + const uint2 w1 = nib[1]; + const unsigned int w[4] = {w0.x, w0.y, w1.x, w1.y}; + float q[32]; +#pragma unroll + for (int i = 0; i < 32; ++i) q[i] = (float)((w[i >> 3] >> (4 * (i & 7))) & 0xf); + const int len = min(32, K - base); +#pragma unroll + for (int j = 0; j < GD_SLOTS; ++j) { + if (slot[j] < 0) continue; + const float* x = act + (size_t)slot[j] * K + base; + float dot = 0.0f, xs = 0.0f; + if (len == 32 && (K & 3) == 0) { +#pragma unroll + for (int i = 0; i < 8; ++i) { + const float4 xv = *(const float4*)(x + 4 * i); + dot = __builtin_fmaf(q[4 * i + 0], xv.x, dot); + dot = __builtin_fmaf(q[4 * i + 1], xv.y, dot); + dot = __builtin_fmaf(q[4 * i + 2], xv.z, dot); + dot = __builtin_fmaf(q[4 * i + 3], xv.w, dot); + xs += (xv.x + xv.y) + (xv.z + xv.w); + } + } else { + for (int i = 0; i < len; ++i) { + dot = __builtin_fmaf(q[i], x[i], dot); + xs += x[i]; + } + } + acc[j] += __builtin_fmaf(scale, dot, zero * xs); + } + } + const float es = per_expert_scale[expert]; +#pragma unroll + for (int j = 0; j < GD_SLOTS; ++j) { + float v = acc[j]; + for (int offset = 16; offset > 0; offset >>= 1) v += __shfl_down(v, offset); + if (lane == 0 && slot[j] >= 0) y[(size_t)(tile * GD_SLOTS + j) * M + row] = es * v; + } +} diff --git a/kernels/src/gemv_hfq4g128_moe_down_residual_scaled_k8_indexed.hip b/kernels/src/gemv_hfq4g128_moe_down_residual_scaled_k8_indexed.hip index ccb0ae79cd..151772d00c 100644 --- a/kernels/src/gemv_hfq4g128_moe_down_residual_scaled_k8_indexed.hip +++ b/kernels/src/gemv_hfq4g128_moe_down_residual_scaled_k8_indexed.hip @@ -1,13 +1,21 @@ // SPDX-License-Identifier: Apache-2.0 // Copyright (c) 2026 Kaden Schutt // hipfire — see LICENSE and NOTICE in the project root. - -#include - // Device-routed MoE down projection for HFQ4-G128 expert weights. Expert // selection and mixture weights remain device-resident so the launch sequence // and kernargs are stable across generated tokens. -__launch_bounds__(32, 16) +// +// One 256-thread workgroup serves eight output rows, one wave per row. A +// group (f32 scale, f32 zero, 64 nibble bytes = 128 values) is split over +// four lanes of 32 values, so a row's groups are read with wide loads and all +// eight experts' loads are in flight together; the activations (8 * K floats, +// cache-resident) are read as float4. Each lane folds the mixture weight in +// before one wave reduction. +#include + +#define MD_ROWS 8 + +__launch_bounds__(256, 2) extern "C" __global__ void gemv_hfq4g128_moe_down_residual_scaled_k8_indexed( const unsigned long long* __restrict__ expert_ptrs, const int* __restrict__ topk_indices, @@ -17,42 +25,54 @@ extern "C" __global__ void gemv_hfq4g128_moe_down_residual_scaled_k8_indexed( float* __restrict__ x_residual, int M, int K ) { - const int row = blockIdx.x; - if (row >= M) return; const int tid = threadIdx.x; + const int lane = tid & 31; + const int row = blockIdx.x * MD_ROWS + (tid >> 5); + if (row >= M) return; + const int groups_per_row = (K + 127) / 128; + const int chunks = groups_per_row * 4; float mixed = 0.0f; - // Fixed rank order makes HIP, kernarg-blob and retained PM4 bit-exact. - #pragma unroll - for (int krank = 0; krank < 8; ++krank) { - const int expert_id = topk_indices[krank]; - const char* __restrict__ A = - reinterpret_cast(expert_ptrs[expert_id]); - const float* __restrict__ x = hidden_batch + (size_t)krank * K; - const char* row_ptr = A + (long long)row * groups_per_row * 72; - float acc = 0.0f; - for (int g = 0; g < groups_per_row; ++g) { - const char* gptr = row_ptr + g * 72; + for (int c = lane; c < chunks; c += 32) { + const int g = c >> 2; + const int base = g * 128 + (c & 3) * 32; + if (base >= K) continue; +#pragma unroll + for (int krank = 0; krank < 8; ++krank) { + const int expert_id = topk_indices[krank]; + const char* gptr = reinterpret_cast(expert_ptrs[expert_id]) + + ((long long)row * groups_per_row + g) * 72; const float scale = __builtin_bit_cast(float, *(const unsigned int*)gptr); const float zero = __builtin_bit_cast(float, *(const unsigned int*)(gptr + 4)); - const unsigned char* nibbles = (const unsigned char*)(gptr + 8); - const int base = g * 128 + tid * 4; - const int byte_off = tid * 2; - const unsigned char b0 = nibbles[byte_off]; - const unsigned char b1 = nibbles[byte_off + 1]; - const float v0 = scale * (float)(b0 & 0xf) + zero; - const float v1 = scale * (float)(b0 >> 4) + zero; - const float v2 = scale * (float)(b1 & 0xf) + zero; - const float v3 = scale * (float)(b1 >> 4) + zero; - if (base < K) acc += v0 * x[base]; - if (base + 1 < K) acc += v1 * x[base + 1]; - if (base + 2 < K) acc += v2 * x[base + 2]; - if (base + 3 < K) acc += v3 * x[base + 3]; + const uint2* nib = (const uint2*)(gptr + 8 + (c & 3) * 16); + const uint2 w0 = nib[0]; + const uint2 w1 = nib[1]; + const unsigned int w[4] = {w0.x, w0.y, w1.x, w1.y}; + const float* x = hidden_batch + (size_t)krank * K + base; + float dot = 0.0f, xs = 0.0f; + if (base + 32 <= K && (K & 3) == 0) { +#pragma unroll + for (int q = 0; q < 8; ++q) { + const float4 xv = *(const float4*)(x + 4 * q); + const unsigned int wq = w[q >> 1] >> (16 * (q & 1)); + dot = __builtin_fmaf((float)(wq & 0xf), xv.x, dot); + dot = __builtin_fmaf((float)((wq >> 4) & 0xf), xv.y, dot); + dot = __builtin_fmaf((float)((wq >> 8) & 0xf), xv.z, dot); + dot = __builtin_fmaf((float)((wq >> 12) & 0xf), xv.w, dot); + xs += (xv.x + xv.y) + (xv.z + xv.w); + } + } else { + for (int j = 0; j < 32 && base + j < K; ++j) { + const float xv = x[j]; + dot = __builtin_fmaf((float)((w[j >> 3] >> (4 * (j & 7))) & 0xf), xv, dot); + xs += xv; + } + } + const float coef = topk_weights[krank] * per_expert_scale[expert_id]; + mixed = __builtin_fmaf(coef, __builtin_fmaf(scale, dot, zero * xs), mixed); } - for (int offset = 16; offset > 0; offset >>= 1) - acc += __shfl_down(acc, offset); - if (tid == 0) - mixed += topk_weights[krank] * per_expert_scale[expert_id] * acc; } - if (tid == 0) x_residual[row] = mixed; + for (int offset = 16; offset > 0; offset >>= 1) + mixed += __shfl_down(mixed, offset); + if (lane == 0) x_residual[row] = mixed; } From c944ce666377307dce954224fdf7fbb2e6f7eae7 Mon Sep 17 00:00:00 2001 From: Bjoern Agent Date: Sat, 10 Oct 2026 08:01:14 +0200 Subject: [PATCH 13/13] docs: regenerate crate maps and env-var inventory on beta 5a0dd38fd --- crates/hipfire-dispatch/map.md | 2 +- crates/rdna-compute/map.md | 6 +++--- docs/env-vars.md | 29 +++++++++-------------------- 3 files changed, 13 insertions(+), 24 deletions(-) diff --git a/crates/hipfire-dispatch/map.md b/crates/hipfire-dispatch/map.md index 094bbbdf49..a79ddeb6c4 100644 --- a/crates/hipfire-dispatch/map.md +++ b/crates/hipfire-dispatch/map.md @@ -49,7 +49,7 @@ _Generated by `scripts/check-crate-maps.py` from the tree — do not edit inside - [`src/pipeline/moe_dump.rs`](src/pipeline/moe_dump.rs): `FoldInfo`, `FoldPending`, `SharedPending`, `enabled`, `final_residual`, `final_residual_deferred`, `fold_begin`, `fold_end`, `set_context`, `shared_begin`, `shared_end` - [`src/pipeline/moe_program.rs`](src/pipeline/moe_program.rs): `MoeStage`, `SealedMoeOp` - [`src/pipeline/qt44_qt53_prefill.rs`](src/pipeline/qt44_qt53_prefill.rs): — -- [`src/pipeline/sandwich.rs`](src/pipeline/sandwich.rs): `Activation`, `KvProjection`, `ParallelMoeMlpOp`, `PerLayerInputOp`, `Rope`, `RopeKind`, `RoutedExperts`, `RoutedScratch`, `SandwichAttentionOp`, `SandwichKv`, `SandwichMlpOp`, `SandwichStream`, `ScaleOp`, `SoftcapOp`, `execute_parallel_moe_mlp`, `execute_per_layer_input`, `execute_sandwich_attention`, `execute_sandwich_mlp`, `execute_scale`, `execute_softcap`, `gemm_rows`, `supports`, `supports_batched_projection` +- [`src/pipeline/sandwich.rs`](src/pipeline/sandwich.rs): `Activation`, `KvProjection`, `ParallelMoeMlpOp`, `PerLayerInputOp`, `Rope`, `RopeKind`, `RoutedBatch`, `RoutedExperts`, `RoutedScratch`, `SandwichAttentionOp`, `SandwichKv`, `SandwichMlpOp`, `SandwichStream`, `ScaleOp`, `SoftcapOp`, `execute_parallel_moe_mlp`, `execute_per_layer_input`, `execute_sandwich_attention`, `execute_sandwich_mlp`, `execute_scale`, `execute_softcap`, `gemm_rows`, `grouped_rows`, `m_total`, `supports`, `supports_batched_projection` - [`src/pipeline/sealed_moe.rs`](src/pipeline/sealed_moe.rs): `ActivationIdentity`, `ActivationInput`, `BoundMoeExperts`, `CompactLiveWeight`, `ContractAssignment`, `ContractAxis`, `ContractCollectiveHint`, `ContractCollectiveRow`, `ContractParallelism`, `ExpertBindingCache`, `ExpertExecutionContract`, `ExpertMetadata`, `ExpertResource`, `ExpertResources`, `ExpertTable`, `MoeContribution`, `MoePrefillRouteProducerProof`, `MoeProtocol`, `MoeRouteProducerProof`, `MoeRouterInput`, `MoeSharedContribution`, `PrefillRouteMode`, `ROOT_ROUTED_EP_EXECUTION`, `ResourceAlias`, `RetainedBodyAction`, `SealedMoeCall`, `SpecializedRouteAdmission`, `activation`, `alias`, `alignment`, `arm_fault_after_expert_mutation`, `assignment`, `basis`, `bind_live`, `bind_live_compact`, `byte_len`, `byte_offset`, `checked_grouped_m_total_bound`, `checked_tp_local_shapes`, `collective_rows`, `contract_id`, `contribution`, `decode_params`, `down`, `dtype`, `encoded_bytes`, `ep_route_producer_proof`, `execute_ep_slot_combine`, `execution`, `execution_contract`, `expert`, `expert_weights_match_metadata`, `experts`, `fingerprint`, `from_cache`, `fused`, `gate`, `gate_up`, `global_id`, `group`, `input`, `invocation`, `is_fused`, `is_live_bound`, `is_root_routed_ep`, `layer`, `local_expert_ids`, `local_rank`, `local_slot`, `local_slots`, `mapping_fingerprint`, `mesh_epoch`, `n_experts`, `new`, `owner`, `owner_rank`, `owner_ranks`, `parallelism`, `physical_device`, `physical_devices`, `prefill_params`, `preflight_decode`, `preflight_prefill`, `prepare_binding`, `protocol`, `rank_count`, `records`, `resources`, `retained_body_action`, `rotation`, `router_input`, `row_stride`, `seal_decode`, `seal_ep_routed_contrib`, `seal_prefill`, `seal_prefill_ep`, `separate`, `shape`, `shared_contribution`, `sidecars`, `source_fingerprint`, `source_name`, `specialized_sealed_moe_retained_admission`, `take_fault_after_expert_mutation`, `up`, `with_alias`, `with_execution_contract`, `with_sidecars` - [`src/pipeline/steps.rs`](src/pipeline/steps.rs): `FusedPattern`, `GemvInput`, `MoeStepHooks`, `Step`, `execute_steps`, `execute_validated_steps`, `execute_validated_steps_with_moe_hooks`, `match_prefix`, `validate_steps` - [`src/pipeline/superop.rs`](src/pipeline/superop.rs): `EpMoeCombineMode`, `EpMoeRouteView`, `EscapeKind`, `ForwardBindings`, `LayerProgram`, `LoweredForward`, `OpBinding`, `OpFlavor`, `ScratchSlot`, `SuperOp`, `SuperOpKind`, `WeightSlot`, `contract_id`, `dispatch_super_op`, `execution_contract`, `lower_layer`, `new`, `run_layer_program` diff --git a/crates/rdna-compute/map.md b/crates/rdna-compute/map.md index 28c881e7df..681fb475d1 100644 --- a/crates/rdna-compute/map.md +++ b/crates/rdna-compute/map.md @@ -22,7 +22,7 @@ _Generated by `scripts/check-crate-maps.py` from the tree — do not edit inside ### Modules and exported symbols - [`src/arch_caps.rs`](src/arch_caps.rs): `ArchCaps`, `arch`, `gemv_dp4a_enabled`, `gemv_prefetch_enabled`, `gemv_rows_default`, `has_cdna3_lds_gemv`, `has_dot2_f32_f16`, `has_gfx11_plus_simt`, `has_hfq3_dp4a`, `has_hfq3_mmq`, `has_hfq3_sdot4`, `has_hfq4_mmq`, `has_mmq`, `has_wmma`, `has_wmma_w32`, `has_wmma_w32_gfx12`, `is_cdna1`, `is_cdna3`, `is_gcn5`, `is_gcn5_wave64`, `is_gfx1010`, `is_gfx1011`, `is_gfx1012`, `is_gfx1030`, `is_gfx1031`, `is_gfx1032`, `is_gfx1100`, `is_gfx1101`, `is_gfx1102`, `is_gfx1103`, `is_gfx1150`, `is_gfx1151`, `is_gfx1152`, `is_gfx1200`, `is_gfx1201`, `is_gfx906`, `is_gfx908`, `is_gfx940`, `is_gfx941`, `is_gfx942`, `is_rdna1`, `is_rdna1p1`, `is_rdna2`, `is_rdna3`, `is_rdna35_apu`, `is_rdna3_dgpu`, `is_rdna3_discrete`, `is_rdna3p5`, `is_rdna4`, `is_wave32`, `is_wave64_native`, `new`, `note_process_gpu_arch`, `process_gpu_arch`, `should_use_mmq`, `supports_dflash_f16_residual_fusions`, `supports_dflash_hidden_scatter_fusion`, `supports_ds4_f16_compressor_cache`, `supports_mq3_lloyd_mb4`, `supports_mq4_lloyd_mb4` -- [`src/attention.rs`](src/attention.rs): `ALL`, `ATTENTION_FLASH_DECODE_VMM_ROWS_SRC`, `GFX12_Q8_FA2_MAX_CTX`, `GFX12_QUERY16_MAX_CTX`, `KV_CACHE_WRITE_FP8_E4M3_VMM_SRC`, `KV_SLOT_DESC_VMM_H`, `VerifyKv`, `VerifyWmmaGeometry`, `VerifyWmmaPv`, `VmmFlashDecodePlan`, `VmmFlashDecodeRoute`, `VmmKvFormat`, `VmmKvSide`, `apply_rope_2d_vision_f32`, `attention_bf16_kv`, `attention_bf16_kv_batched`, `attention_causal_batched`, `attention_decode_batched_history`, `attention_dflash_f32`, `attention_dflash_sliding_f32`, `attention_dflash_swa_f32`, `attention_dflash_wmma_f32`, `attention_dflash_wmma_m128_n32_f16kv_v7_f32`, `attention_dflash_wmma_m128_n32_f16kv_v7b_f32`, `attention_dflash_wmma_m32_f32`, `attention_dflash_wmma_m64_n128_f16kv_f32`, `attention_dflash_wmma_m64_n128_f16kv_v2_f32`, `attention_dflash_wmma_m64_n128_f16kv_v3_causal_f32`, `attention_dflash_wmma_m64_n128_f16kv_v3_f32`, `attention_dflash_wmma_m64_n128_f16kv_v4_causal_f32`, `attention_dflash_wmma_m64_n128_f16kv_v4_f32`, `attention_dflash_wmma_m64_n32_f16kv_v5_f32`, `attention_dflash_wmma_m64_n32_f16kv_v6_f32`, `attention_dflash_wmma_n128_f16kv_f32`, `attention_dflash_wmma_n64_f16kv_f32`, `attention_dflash_wmma_n64_f32`, `attention_f32`, `attention_flash`, `attention_flash_asym2`, `attention_flash_asym2_batched`, `attention_flash_asym2_batched_slots`, `attention_flash_asym3`, `attention_flash_asym3_batched`, `attention_flash_asym3_batched_masked`, `attention_flash_asym3_batched_masked_slots`, `attention_flash_asym3_gated_mq_rotate_awq`, `attention_flash_asym4`, `attention_flash_asym4_batched`, `attention_flash_asym4_batched_masked`, `attention_flash_asym4_batched_masked_slots`, `attention_flash_asym4_wmma_tile_batched`, `attention_flash_asym4_wmma_tile_batched_gfx12`, `attention_flash_bf16_batched_masked_windowed`, `attention_flash_bf16_batched_masked_windowed_slots`, `attention_flash_bf16_windowed`, `attention_flash_decode_vmm`, `attention_flash_f16_batched_masked_windowed`, `attention_flash_f16_batched_masked_windowed_slots`, `attention_flash_f16_windowed`, `attention_flash_fp8_e4m3_tile`, `attention_flash_fp8_e4m3_tile_batched`, `attention_flash_fwht2`, `attention_flash_fwht2_batched`, `attention_flash_fwht2_batched_slots`, `attention_flash_fwht3`, `attention_flash_fwht3_batched`, `attention_flash_fwht3_batched_masked`, `attention_flash_fwht3_batched_masked_slots`, `attention_flash_fwht4`, `attention_flash_fwht4_batched`, `attention_flash_fwht4_batched_masked`, `attention_flash_fwht4_batched_masked_slots`, `attention_flash_gqa`, `attention_flash_gqa_fused`, `attention_flash_q8_0`, `attention_flash_q8_0_batched_masked`, `attention_flash_q8_0_batched_masked_slots`, `attention_flash_q8_0_batched_masked_windowed`, `attention_flash_q8_0_gated_mq_rotate_awq_gfx1100`, `attention_flash_q8_0_gated_mq_rotate_gfx1100`, `attention_flash_q8_0_rows_masked`, `attention_flash_q8_0_windowed`, `attention_flux_best_f16kv_f32`, `attention_flux_v2_wmma_f16kv`, `attention_flux_vt_wmma_f16kv_f32`, `attention_flux_vtk_wmma_f16kv_f32`, `attention_fp8_e4m3_fa2_gqa_f16_gfx1201`, `attention_fp8_e4m3_fa2_gqa_fp8_gfx1201`, `attention_fp8_e4m3_fa2_gqa_packet_gfx1201`, `attention_fp8_e4m3_fa2_gqa_packet_split_gfx1201_bench`, `attention_fp8_e4m3_fa2_gqa_qresident_gfx1201`, `attention_fp8_e4m3_fa2_gqa_qresident_v2_gfx1201`, `attention_fp8_e4m3_fa2_gqa_qresident_v2_q8_a4epi_gfx1201`, `attention_fp8_e4m3_fa2_gqa_qresident_v2_q8_gfx1201`, `attention_fp8_e4m3_fa2_gqa_stageb_split_gfx1201_bench`, `attention_fp8_e4m3_kv`, `attention_fp8_e4m3_kv_batched`, `attention_gqa_warp`, `attention_gqa_warp_dv`, `attention_hfq4_kv`, `attention_hfq8_kv`, `attention_int8_kv`, `attention_int8c_f16_kv`, `attention_int8c_kv`, `attention_kv_batched_vmm`, `attention_partition_merge_f32_bench`, `attention_q4kv`, `attention_q8_0_fa2_gqa_fp8_gfx1201`, `attention_q8_0_fa2_gqa_fwht3k_gfx11`, `attention_q8_0_fa2_gqa_fwht3k_gfx1201`, `attention_q8_0_fa2_gqa_gfx11`, `attention_q8_0_fa2_gqa_gfx1201`, `attention_q8_0_fa2_gqa_split_gfx1201`, `attention_q8_0_flash_prefill`, `attention_q8_0_flash_prefill_slots`, `attention_q8_0_flash_prefill_wmma`, `attention_q8_0_flash_prefill_wmma_incumbent_gfx1201_bench`, `attention_q8_0_flash_prefill_wmma_partial_gfx1201_bench`, `attention_q8_0_flash_prefill_wmma_slots`, `attention_q8_0_flash_prefill_wmma_swa`, `attention_q8_0_kv`, `attention_q8_0_kv_batched`, `attention_q8_0_kv_batched_masked`, `attention_q8_0_kv_batched_masked_slots`, `attention_q8_0_kv_independent`, `attention_q8_0_kv_independent_lds_bytes`, `attention_q8_0_kv_independent_masked_windowed`, `attention_q8_0_kv_independent_max_lane_capacity`, `attention_q8_0_kv_independent_shared_mem_limit`, `attention_q8_0_kv_swa`, `attention_q8_0_kv_timed`, `attention_q8kv`, `attention_verify_wmma_with`, `bf16_row_bytes`, `bytes_per_token`, `cast_f32_to_f16_at_slot_buf`, `compressor_add_ape_batched_f32`, `compressor_add_ape_f32_buf`, `compressor_compress_aligned_batched_f32`, `compressor_overlap_concat_f32`, `compressor_ring_write_batched_f32`, `compressor_softmax_pool_f32`, `compressor_softmax_pool_f32_buf`, `compressor_softmax_pool_f32_staged_buf`, `deepseek4_attn_pos0`, `deepseek4_attn_swa`, `deepseek4_attn_swa_batched`, `deepseek4_attn_swa_batched_debug`, `deepseek4_attn_swa_buf`, `deepseek4_attn_swa_topk_batched_f32`, `deepseek4_attn_swa_topk_batched_wmma`, `deepseek4_attn_swa_topk_batched_wmma_gfx12`, `deepseek4_attn_swa_topk_direct_batched_f32`, `deepseek4_attn_swa_topk_direct_wmma`, `deepseek4_attn_swa_topk_f32_buf`, `dspark_stage_kv`, `ensure_attention_q8_0_kv_independent_lds`, `fa2_gfx11_batch_admitted`, `fa2_gfx11_ctx_admitted`, `fa2_gfx11_max_ctx`, `fa2_gqa_split_gfx1201_partials_len`, `flash_attn_ck_loaded`, `flux_attn_route_name`, `fp8_e4m3_row_bytes`, `gfx12_q8_fa2_prefill_admitted`, `hc_apply_alpha`, `hc_apply_alpha_batched`, `hc_compute_control`, `hc_compute_control_batched`, `hc_compute_control_batched_fused24_gfx1201`, `hc_compute_control_vec4_finalize`, `hc_compute_control_wmma_batched_gfx1201`, `hc_finalize_control`, `hc_finalize_input_map`, `hc_head_compute_pre`, `hc_input_map_4stream`, `hc_input_map_4stream_batched`, `hc_mix_4stream`, `hc_mix_4stream_batched`, `hc_mix_4stream_peer3_gfx1201`, `hc_mix_4stream_peer4_gfx1201`, `hc_pre_post_sigmoid_scale_f32`, `hc_sinkhorn_4x4`, `hc_sinkhorn_4x4_batched`, `hc_split_finalize_batched`, `hc_streams_init_from_embed_batched`, `hc_streams_init_from_embed_batched_bf16`, `indexer_compressed_k_score`, `indexer_kv_gather`, `indexer_relu_score_batched_f16`, `indexer_relu_score_batched_f32`, `indexer_relu_score_f16_buf`, `indexer_relu_score_f32_buf`, `indexer_relu_score_f32_buf_sharded_gfx1201`, `indexer_relu_score_wmma_batched_f16`, `indexer_relu_score_wmma_batched_f32`, `indexer_relu_score_wmma_batched_klds_gfx1151`, `indexer_relu_score_wmma_batched_sharded_gfx1201`, `indexer_top_k`, `indexer_top_k_batched`, `indexer_top_k_batched_bounded_gfx1151`, `indexer_top_k_batched_bounded_gfx1201`, `indexer_top_k_batched_buf`, `indexer_top_k_buf`, `indexer_top_k_two_stage`, `kv_cache_write`, `kv_cache_write_asym2_batched`, `kv_cache_write_asym2_batched_slots`, `kv_cache_write_asym2_fused`, `kv_cache_write_asym3_batched`, `kv_cache_write_asym3_batched_slots`, `kv_cache_write_asym3_fused`, `kv_cache_write_asym4_batched`, `kv_cache_write_asym4_batched_slots`, `kv_cache_write_asym4_fused`, `kv_cache_write_batched_vmm`, `kv_cache_write_bf16`, `kv_cache_write_bf16_batched`, `kv_cache_write_f16`, `kv_cache_write_f16_batched`, `kv_cache_write_f32_batched`, `kv_cache_write_fp8_e4m3`, `kv_cache_write_fp8_e4m3_batched`, `kv_cache_write_fwht2_batched`, `kv_cache_write_fwht2_batched_slots`, `kv_cache_write_fwht2_fused`, `kv_cache_write_fwht3_batched`, `kv_cache_write_fwht3_batched_slots`, `kv_cache_write_fwht3_fused`, `kv_cache_write_fwht3_vec`, `kv_cache_write_fwht3_vec_batched`, `kv_cache_write_fwht3_vec_batched_slots`, `kv_cache_write_fwht4_batched`, `kv_cache_write_fwht4_batched_slots`, `kv_cache_write_fwht4_fused`, `kv_cache_write_hfq4`, `kv_cache_write_hfq8`, `kv_cache_write_int8`, `kv_cache_write_int8c`, `kv_cache_write_int8c_f16`, `kv_cache_write_q4`, `kv_cache_write_q8`, `kv_cache_write_q8_0`, `kv_cache_write_q8_0_batched`, `kv_cache_write_q8_0_batched_slots`, `kv_cache_write_q8_0_independent`, `kv_cache_write_q8_0_independent_masked`, `kv_cache_write_q8_0_pair`, `kv_cache_write_q8_0_ring`, `kv_cache_write_v256_2bit_vec`, `kv_cache_write_v256_2bit_vec_batched`, `kv_cache_write_v256_4bit_vec`, `kv_cache_write_v256_4bit_vec_batched`, `kv_compact_gather`, `kv_fold_asym2`, `kv_fold_asym3`, `kv_fold_asym4`, `kv_fold_q8`, `kv_slot_desc_vmm_source`, `pflash_score_fwht2_kv`, `pflash_score_fwht3_kv`, `pflash_score_fwht4_kv`, `pflash_score_q8_kv`, `q8_flash_tile_size`, `rmsnorm_f32_staged_buf`, `rope_tail_interleaved`, `rope_tail_interleaved_batched`, `rope_tail_interleaved_h64d128r64_gfx1201`, `rope_tail_yarn_interleaved`, `rope_tail_yarn_interleaved_at_slot_buf`, `rope_tail_yarn_interleaved_batched`, `rope_tail_yarn_interleaved_staged_buf`, `rows_batched`, `state_overlap_shift_f32_buf`, `state_ring_write_f32_buf`, `swa_ring_write_batched_f32`, `swa_ring_write_batched_pos_f32`, `swa_ring_write_f32_buf`, `swa_visibility_stage_batched`, `swa_visibility_stage_batched_pos`, `tp_graph_signal_store_gfx1201`, `tp_graph_signal_wait2_gfx1201`, `tp_graph_signal_wait3_gfx1201`, `transcode_k_fwht4_to_fwht2`, `transcode_k_fwht4_to_fwht3`, `transcode_v_lloyd_down`, `transcode_v_q8_to_lloyd4`, `triattn_accumulate`, `triattn_score_asym2`, `triattn_score_asym3`, `triattn_score_asym4`, `triattn_score_fwht`, `triattn_score_q8`, `try_attention_verify_gqa`, `try_attention_verify_gqa_rows`, `try_attention_verify_wmma`, `vit_attention_f32`, `vit_attention_opt`, `vit_attention_qtiled_f32`, `vmm_flash_decode_plan`, `vmm_rows_source` +- [`src/attention.rs`](src/attention.rs): `ALL`, `ATTENTION_FLASH_DECODE_VMM_ROWS_SRC`, `GFX12_Q8_FA2_MAX_CTX`, `GFX12_QUERY16_MAX_CTX`, `KV_CACHE_WRITE_FP8_E4M3_VMM_SRC`, `KV_SLOT_DESC_VMM_H`, `VerifyKv`, `VerifyWmmaGeometry`, `VerifyWmmaPv`, `VmmFlashDecodePlan`, `VmmFlashDecodeRoute`, `VmmKvFormat`, `VmmKvSide`, `apply_rope_2d_vision_f32`, `attention_bf16_kv`, `attention_bf16_kv_batched`, `attention_causal_batched`, `attention_decode_batched_history`, `attention_dflash_f32`, `attention_dflash_sliding_f32`, `attention_dflash_swa_f32`, `attention_dflash_wmma_f32`, `attention_dflash_wmma_m128_n32_f16kv_v7_f32`, `attention_dflash_wmma_m128_n32_f16kv_v7b_f32`, `attention_dflash_wmma_m32_f32`, `attention_dflash_wmma_m64_n128_f16kv_f32`, `attention_dflash_wmma_m64_n128_f16kv_v2_f32`, `attention_dflash_wmma_m64_n128_f16kv_v3_causal_f32`, `attention_dflash_wmma_m64_n128_f16kv_v3_f32`, `attention_dflash_wmma_m64_n128_f16kv_v4_causal_f32`, `attention_dflash_wmma_m64_n128_f16kv_v4_f32`, `attention_dflash_wmma_m64_n32_f16kv_v5_f32`, `attention_dflash_wmma_m64_n32_f16kv_v6_f32`, `attention_dflash_wmma_n128_f16kv_f32`, `attention_dflash_wmma_n64_f16kv_f32`, `attention_dflash_wmma_n64_f32`, `attention_f32`, `attention_flash`, `attention_flash_asym2`, `attention_flash_asym2_batched`, `attention_flash_asym2_batched_slots`, `attention_flash_asym3`, `attention_flash_asym3_batched`, `attention_flash_asym3_batched_masked`, `attention_flash_asym3_batched_masked_slots`, `attention_flash_asym3_gated_mq_rotate_awq`, `attention_flash_asym4`, `attention_flash_asym4_batched`, `attention_flash_asym4_batched_masked`, `attention_flash_asym4_batched_masked_slots`, `attention_flash_asym4_wmma_tile_batched`, `attention_flash_asym4_wmma_tile_batched_gfx12`, `attention_flash_bf16_batched_masked_windowed`, `attention_flash_bf16_batched_masked_windowed_slots`, `attention_flash_bf16_windowed`, `attention_flash_decode_vmm`, `attention_flash_f16_batched_masked_windowed`, `attention_flash_f16_batched_masked_windowed_slots`, `attention_flash_f16_windowed`, `attention_flash_fp8_e4m3_tile`, `attention_flash_fp8_e4m3_tile_batched`, `attention_flash_fwht2`, `attention_flash_fwht2_batched`, `attention_flash_fwht2_batched_slots`, `attention_flash_fwht3`, `attention_flash_fwht3_batched`, `attention_flash_fwht3_batched_masked`, `attention_flash_fwht3_batched_masked_slots`, `attention_flash_fwht4`, `attention_flash_fwht4_batched`, `attention_flash_fwht4_batched_masked`, `attention_flash_fwht4_batched_masked_slots`, `attention_flash_gqa`, `attention_flash_gqa_fused`, `attention_flash_q8_0`, `attention_flash_q8_0_batched_masked`, `attention_flash_q8_0_batched_masked_slots`, `attention_flash_q8_0_batched_masked_windowed`, `attention_flash_q8_0_gated_mq_rotate_awq_gfx1100`, `attention_flash_q8_0_gated_mq_rotate_gfx1100`, `attention_flash_q8_0_rows_masked`, `attention_flash_q8_0_windowed`, `attention_flux_best_f16kv_f32`, `attention_flux_v2_wmma_f16kv`, `attention_flux_vt_wmma_f16kv_f32`, `attention_flux_vtk_wmma_f16kv_f32`, `attention_fp8_e4m3_fa2_gqa_f16_gfx1201`, `attention_fp8_e4m3_fa2_gqa_fp8_gfx1201`, `attention_fp8_e4m3_fa2_gqa_packet_gfx1201`, `attention_fp8_e4m3_fa2_gqa_packet_split_gfx1201_bench`, `attention_fp8_e4m3_fa2_gqa_qresident_gfx1201`, `attention_fp8_e4m3_fa2_gqa_qresident_v2_gfx1201`, `attention_fp8_e4m3_fa2_gqa_qresident_v2_q8_a4epi_gfx1201`, `attention_fp8_e4m3_fa2_gqa_qresident_v2_q8_gfx1201`, `attention_fp8_e4m3_fa2_gqa_stageb_split_gfx1201_bench`, `attention_fp8_e4m3_kv`, `attention_fp8_e4m3_kv_batched`, `attention_gqa_warp`, `attention_gqa_warp_dv`, `attention_hfq4_kv`, `attention_hfq8_kv`, `attention_int8_kv`, `attention_int8c_f16_kv`, `attention_int8c_kv`, `attention_kv_batched_vmm`, `attention_partition_merge_f32_bench`, `attention_q4kv`, `attention_q8_0_fa2_gqa_fp8_gfx1201`, `attention_q8_0_fa2_gqa_fwht3k_gfx11`, `attention_q8_0_fa2_gqa_fwht3k_gfx1201`, `attention_q8_0_fa2_gqa_gfx11`, `attention_q8_0_fa2_gqa_gfx1201`, `attention_q8_0_fa2_gqa_split_gfx1201`, `attention_q8_0_flash_prefill`, `attention_q8_0_flash_prefill_slots`, `attention_q8_0_flash_prefill_wmma`, `attention_q8_0_flash_prefill_wmma_incumbent_gfx1201_bench`, `attention_q8_0_flash_prefill_wmma_partial_gfx1201_bench`, `attention_q8_0_flash_prefill_wmma_slots`, `attention_q8_0_flash_prefill_wmma_swa`, `attention_q8_0_kv`, `attention_q8_0_kv_batched`, `attention_q8_0_kv_batched_masked`, `attention_q8_0_kv_batched_masked_slots`, `attention_q8_0_kv_independent`, `attention_q8_0_kv_independent_lds_bytes`, `attention_q8_0_kv_independent_masked_windowed`, `attention_q8_0_kv_independent_max_lane_capacity`, `attention_q8_0_kv_independent_shared_mem_limit`, `attention_q8_0_kv_swa`, `attention_q8_0_kv_timed`, `attention_q8_0_prefill_gqa_wmma`, `attention_q8_0_prefill_gqa_wmma_admitted`, `attention_q8kv`, `attention_verify_wmma_with`, `bf16_row_bytes`, `bytes_per_token`, `cast_f32_to_f16_at_slot_buf`, `compressor_add_ape_batched_f32`, `compressor_add_ape_f32_buf`, `compressor_compress_aligned_batched_f32`, `compressor_overlap_concat_f32`, `compressor_ring_write_batched_f32`, `compressor_softmax_pool_f32`, `compressor_softmax_pool_f32_buf`, `compressor_softmax_pool_f32_staged_buf`, `deepseek4_attn_pos0`, `deepseek4_attn_swa`, `deepseek4_attn_swa_batched`, `deepseek4_attn_swa_batched_debug`, `deepseek4_attn_swa_buf`, `deepseek4_attn_swa_topk_batched_f32`, `deepseek4_attn_swa_topk_batched_wmma`, `deepseek4_attn_swa_topk_batched_wmma_gfx12`, `deepseek4_attn_swa_topk_direct_batched_f32`, `deepseek4_attn_swa_topk_direct_wmma`, `deepseek4_attn_swa_topk_f32_buf`, `dspark_stage_kv`, `ensure_attention_q8_0_kv_independent_lds`, `fa2_gfx11_batch_admitted`, `fa2_gfx11_ctx_admitted`, `fa2_gfx11_max_ctx`, `fa2_gqa_split_gfx1201_partials_len`, `flash_attn_ck_loaded`, `flux_attn_route_name`, `fp8_e4m3_row_bytes`, `gfx12_q8_fa2_prefill_admitted`, `hc_apply_alpha`, `hc_apply_alpha_batched`, `hc_compute_control`, `hc_compute_control_batched`, `hc_compute_control_batched_fused24_gfx1201`, `hc_compute_control_vec4_finalize`, `hc_compute_control_wmma_batched_gfx1201`, `hc_finalize_control`, `hc_finalize_input_map`, `hc_head_compute_pre`, `hc_input_map_4stream`, `hc_input_map_4stream_batched`, `hc_mix_4stream`, `hc_mix_4stream_batched`, `hc_mix_4stream_peer3_gfx1201`, `hc_mix_4stream_peer4_gfx1201`, `hc_pre_post_sigmoid_scale_f32`, `hc_sinkhorn_4x4`, `hc_sinkhorn_4x4_batched`, `hc_split_finalize_batched`, `hc_streams_init_from_embed_batched`, `hc_streams_init_from_embed_batched_bf16`, `indexer_compressed_k_score`, `indexer_kv_gather`, `indexer_relu_score_batched_f16`, `indexer_relu_score_batched_f32`, `indexer_relu_score_f16_buf`, `indexer_relu_score_f32_buf`, `indexer_relu_score_f32_buf_sharded_gfx1201`, `indexer_relu_score_wmma_batched_f16`, `indexer_relu_score_wmma_batched_f32`, `indexer_relu_score_wmma_batched_klds_gfx1151`, `indexer_relu_score_wmma_batched_sharded_gfx1201`, `indexer_top_k`, `indexer_top_k_batched`, `indexer_top_k_batched_bounded_gfx1151`, `indexer_top_k_batched_bounded_gfx1201`, `indexer_top_k_batched_buf`, `indexer_top_k_buf`, `indexer_top_k_two_stage`, `kv_cache_write`, `kv_cache_write_asym2_batched`, `kv_cache_write_asym2_batched_slots`, `kv_cache_write_asym2_fused`, `kv_cache_write_asym3_batched`, `kv_cache_write_asym3_batched_slots`, `kv_cache_write_asym3_fused`, `kv_cache_write_asym4_batched`, `kv_cache_write_asym4_batched_slots`, `kv_cache_write_asym4_fused`, `kv_cache_write_batched_vmm`, `kv_cache_write_bf16`, `kv_cache_write_bf16_batched`, `kv_cache_write_f16`, `kv_cache_write_f16_batched`, `kv_cache_write_f32_batched`, `kv_cache_write_fp8_e4m3`, `kv_cache_write_fp8_e4m3_batched`, `kv_cache_write_fwht2_batched`, `kv_cache_write_fwht2_batched_slots`, `kv_cache_write_fwht2_fused`, `kv_cache_write_fwht3_batched`, `kv_cache_write_fwht3_batched_slots`, `kv_cache_write_fwht3_fused`, `kv_cache_write_fwht3_vec`, `kv_cache_write_fwht3_vec_batched`, `kv_cache_write_fwht3_vec_batched_slots`, `kv_cache_write_fwht4_batched`, `kv_cache_write_fwht4_batched_slots`, `kv_cache_write_fwht4_fused`, `kv_cache_write_hfq4`, `kv_cache_write_hfq8`, `kv_cache_write_int8`, `kv_cache_write_int8c`, `kv_cache_write_int8c_f16`, `kv_cache_write_q4`, `kv_cache_write_q8`, `kv_cache_write_q8_0`, `kv_cache_write_q8_0_batched`, `kv_cache_write_q8_0_batched_slots`, `kv_cache_write_q8_0_independent`, `kv_cache_write_q8_0_independent_masked`, `kv_cache_write_q8_0_pair`, `kv_cache_write_q8_0_ring`, `kv_cache_write_v256_2bit_vec`, `kv_cache_write_v256_2bit_vec_batched`, `kv_cache_write_v256_4bit_vec`, `kv_cache_write_v256_4bit_vec_batched`, `kv_compact_gather`, `kv_fold_asym2`, `kv_fold_asym3`, `kv_fold_asym4`, `kv_fold_q8`, `kv_slot_desc_vmm_source`, `pflash_score_fwht2_kv`, `pflash_score_fwht3_kv`, `pflash_score_fwht4_kv`, `pflash_score_q8_kv`, `q8_flash_tile_size`, `rmsnorm_f32_staged_buf`, `rope_tail_interleaved`, `rope_tail_interleaved_batched`, `rope_tail_interleaved_h64d128r64_gfx1201`, `rope_tail_yarn_interleaved`, `rope_tail_yarn_interleaved_at_slot_buf`, `rope_tail_yarn_interleaved_batched`, `rope_tail_yarn_interleaved_staged_buf`, `rows_batched`, `state_overlap_shift_f32_buf`, `state_ring_write_f32_buf`, `swa_ring_write_batched_f32`, `swa_ring_write_batched_pos_f32`, `swa_ring_write_f32_buf`, `swa_visibility_stage_batched`, `swa_visibility_stage_batched_pos`, `tp_graph_signal_store_gfx1201`, `tp_graph_signal_wait2_gfx1201`, `tp_graph_signal_wait3_gfx1201`, `transcode_k_fwht4_to_fwht2`, `transcode_k_fwht4_to_fwht3`, `transcode_v_lloyd_down`, `transcode_v_q8_to_lloyd4`, `triattn_accumulate`, `triattn_score_asym2`, `triattn_score_asym3`, `triattn_score_asym4`, `triattn_score_fwht`, `triattn_score_q8`, `try_attention_verify_gqa`, `try_attention_verify_gqa_rows`, `try_attention_verify_wmma`, `vit_attention_f32`, `vit_attention_opt`, `vit_attention_qtiled_f32`, `vmm_flash_decode_plan`, `vmm_rows_source` - [`src/bin/hipfire-kernel-hash.rs`](src/bin/hipfire-kernel-hash.rs): — - [`src/bin/hipfire-kernel-manifest.rs`](src/bin/hipfire-kernel-manifest.rs): — - [`src/bin/hipfire-kernel-pack.rs`](src/bin/hipfire-kernel-pack.rs): — @@ -44,7 +44,7 @@ _Generated by `scripts/check-crate-maps.py` from the tree — do not edit inside - [`src/flux_fused.rs`](src/flux_fused.rs): `gemv_f16_bias_xf32`, `layernorm_modulate`, `qk_rmsnorm_rope_flux` - [`src/fn_gdn_dense.rs`](src/fn_gdn_dense.rs): `applies`, `gate`, `run` - [`src/gap_timing.rs`](src/gap_timing.rs): `Slot`, `add_ns`, `add_since`, `begin_forward`, `enabled`, `end_forward`, `now` -- [`src/gemm.rs`](src/gemm.rs): `ADDIN`, `BF16_SCALED_ADD`, `F2GdnTargets`, `GATED`, `GELU`, `GemmEpilogue`, `HC_DOWN_GFX1201`, `LDS_EPI_TILES`, `LDS_TILE_VARIANTS`, `LOWBIT_WMMA_WAVES`, `LdsTile`, `LdsTileSplitK`, `OUT_F16`, `QWEN4_F16_WMMA_MIN_TOKENS`, `QWEN4_MOE_SYM_IU4_MIN_ROWS`, `QWEN4_MQ4_X4_MIN_N_GFX1151`, `QWEN4_MQ6_X4_PM_MIN_M_GFX1151`, `QWEN4_MQ6_X4_PM_MIN_N_GFX1151`, `QWEN4_SHARED_GU_PM_WAVES`, `SUPPORTED`, `a4_rms_fdiv_planes_wanted`, `a4_slab_active`, `arm_residual_fold`, `disarm_residual_fold`, `ensure_qwen4_moe_sym_pm_entry`, `entry`, `entry_epi`, `entry_suffix`, `flush_residual_fold`, `fn`, `fp8_f2_row_active`, `fp8_silu_h_active`, `fused_gate_up_hfq4g256`, `fused_gate_up_hfq4g256_dp4a`, `fused_gate_up_hfq4g256_mq4v2`, `fused_gate_up_hfq6g256_dp4a`, `fused_gate_up_mq2g256v2`, `fused_gate_up_mq3g256v2`, `fused_gate_up_mq4cg256`, `fused_gate_up_mq4g256v2_lloyd`, `fused_gate_up_mq5g256v2`, `fused_gate_up_mq6g256v2`, `fused_gate_up_q4k`, `fused_gate_up_q8_0`, `fused_qk_l2_norm_scale_f32`, `fused_qk_l2_norm_scale_f32_batched`, `fused_qk_l2_norm_scale_interleave_f32_batched`, `fused_qkv_hfq4g256`, `fused_qkv_hfq4g256_dp4a`, `fused_qkv_hfq4g256_mq4v2`, `fused_qkv_hfq4g256_mq4v2_generic_exact`, `fused_qkv_hfq4g256_mq4v2_k2048_x_buffer_gfx1100_exact`, `fused_qkv_hfq4g256_with_bias`, `fused_qkv_hfq4g256_with_bias_mq4v2`, `fused_qkv_hfq6g256`, `fused_qkv_hfq6g256_dp4a`, `fused_qkv_hfq6g256_with_bias`, `fused_qkv_mq2g256v2`, `fused_qkv_mq3g256v2`, `fused_qkv_mq4cg256`, `fused_qkv_mq4cg256_with_bias`, `fused_qkv_mq4g256v2_lloyd`, `fused_qkv_mq5g256v2`, `fused_qkv_mq6g256v2`, `fused_qkv_q4k`, `fused_qkv_q4k_with_bias`, `fused_qkv_q8_0`, `fused_qkv_q8_0_with_bias`, `fused_qkvza_hfq4g256`, `fused_qkvza_hfq4g256_dp4a`, `fused_qkvza_hfq4g256_mq4v2`, `fused_qkvza_hfq4g256_mq4v2_generic_exact`, `fused_qkvza_hfq4g256_mq4v2_k2048_hoist_x32_exact`, `fused_qkvza_hfq4g256_scalar_prep_gfx1100`, `fused_qkvza_hfq6g256_dp4a`, `fused_qkvza_mq2g256v2`, `fused_qkvza_mq3g256v2`, `fused_qkvza_mq4cg256`, `fused_qkvza_mq4g256v2_lloyd`, `fused_qkvza_mq5g256v2`, `fused_qkvza_mq6g256v2`, `fused_qkvza_q8_0`, `fused_sigmoid_alpha_gate_f32`, `fused_sigmoid_alpha_gate_f32_batched`, `gemm_bf16_mfma_gfx942`, `gemm_bf16_xf16_f16_wmma`, `gemm_bf16_xf16_f16_wmma_regions`, `gemm_bf16_xf32_batched`, `gemm_bf16_xf32_f16_wmma_qwen4`, `gemm_bf16_xf32_f16_wmma_qwen4_bf16st`, `gemm_bf16_xf32_f16_wmma_qwen4_groups`, `gemm_bf16_xf32_f16_wmma_qwen4_hcsd`, `gemm_bf16_xf32_f16_wmma_qwen4_hcsd_applies`, `gemm_bf16_xf32_f16_wmma_qwen4_shared_gu_pm`, `gemm_bf16_xf32_multirow`, `gemm_bq1g128_wmma`, `gemm_f16`, `gemm_f16_batched_lmhead`, `gemm_f16_bias`, `gemm_f16_tiled`, `gemm_f16_wmma`, `gemm_f16_wmma_mb4`, `gemm_f16_wmma_mb8`, `gemm_f16_x_f16_mfma_gfx942`, `gemm_f16_x_f16_wmma`, `gemm_f16_x_f16_wmma_lds`, `gemm_f16_x_f16_wmma_lds_auto`, `gemm_f16_x_f16_wmma_lds_auto_epi`, `gemm_f16_x_f16_wmma_lds_auto_epi_ld`, `gemm_f16_x_f16_wmma_lds_auto_ld`, `gemm_f16_x_f16_wmma_lds_epi`, `gemm_f16_x_f16_wmma_lds_epi_ld`, `gemm_f16_x_f16_wmma_lds_splitk`, `gemm_f16_x_f16_wmma_lds_tiled`, `gemm_f16_x_f16_wmma_lds_tiled_ld`, `gemm_f32_batched`, `gemm_f32_register_tiled`, `gemm_gate_up_hfp4g32`, `gemm_gate_up_hfp4g32_wmma`, `gemm_gate_up_hfp4g32_wmma_gfx12`, `gemm_gate_up_hfq3g256`, `gemm_gate_up_hfq3g256_dot2`, `gemm_gate_up_hfq3g256_dp4a`, `gemm_gate_up_hfq3g256_fp16`, `gemm_gate_up_hfq3g256_mmq`, `gemm_gate_up_hfq3g256_mmq_x16`, `gemm_gate_up_hfq3g256_mmq_x32`, `gemm_gate_up_hfq3g256_mmq_x32_y64`, `gemm_gate_up_hfq3g256_mmq_x32_y96`, `gemm_gate_up_hfq3g256_mmq_x8`, `gemm_gate_up_hfq3g256_wmma`, `gemm_gate_up_hfq3g256_wmma_gfx12`, `gemm_gate_up_hfq3g256_wmma_mb4`, `gemm_gate_up_hfq4g256`, `gemm_gate_up_hfq4g256_dot2`, `gemm_gate_up_hfq4g256_fp16`, `gemm_gate_up_hfq4g256_fp16_wave64`, `gemm_gate_up_hfq4g256_mmq`, `gemm_gate_up_hfq4g256_mmq_gfx906`, `gemm_gate_up_hfq4g256_mmq_gfx906_prequant`, `gemm_gate_up_hfq4g256_mmq_x16`, `gemm_gate_up_hfq4g256_mmq_x32`, `gemm_gate_up_hfq4g256_mq4v2`, `gemm_gate_up_hfq4g256_wave64_dp4a`, `gemm_gate_up_hfq4g256_wave64_dp4a_prequant`, `gemm_gate_up_hfq4g256_wmma`, `gemm_gate_up_hfq4g256_wmma_gfx12`, `gemm_gate_up_hfq4g256_wmma_gfx12_mq4v2`, `gemm_gate_up_hfq4g256_wmma_gfx12_mq4v2_fp8_bt12`, `gemm_gate_up_hfq4g256_wmma_gfx12_mq4v2_fp8_bt12_prepared`, `gemm_gate_up_hfq4g256_wmma_gfx12_mq4v2_fp8_bt12_prepared_lloyd`, `gemm_gate_up_hfq4g256_wmma_gfx12_mq4v2_fp8_lloyd`, `gemm_gate_up_hfq6g256`, `gemm_gate_up_hfq6g256_dot2`, `gemm_gate_up_hfq6g256_fp16`, `gemm_gate_up_hfq6g256_wave64_dp4a`, `gemm_gate_up_hfq6g256_wmma`, `gemm_gate_up_hfq6g256_wmma_gfx12`, `gemm_gate_up_mq2g256v2_wmma`, `gemm_gate_up_mq2g256v2_wmma_gfx11`, `gemm_gate_up_mq2g256v2_wmma_gfx12`, `gemm_gate_up_mq3g256_lloyd_wmma`, `gemm_gate_up_mq3g256_lloyd_wmma_mb4`, `gemm_gate_up_mq3g256_wmma`, `gemm_gate_up_mq3g256v2_wmma`, `gemm_gate_up_mq3g256v2_wmma_gfx11`, `gemm_gate_up_mq3g256v2_wmma_gfx12`, `gemm_gate_up_mq4cg256`, `gemm_gate_up_mq4cg256_wmma_gfx12`, `gemm_gate_up_mq4g256_lloyd_wmma`, `gemm_gate_up_mq4g256_lloyd_wmma_mb4`, `gemm_gate_up_mq4g256v2_mmq_lloyd`, `gemm_gate_up_mq4g256v2_wmma`, `gemm_gate_up_mq4g256v2_wmma_gfx1100_bt`, `gemm_gate_up_mq4g256v2_wmma_gfx1100_mw_lds`, `gemm_gate_up_mq4g256v2_wmma_gfx1151_bt`, `gemm_gate_up_mq4g256v2_wmma_iu4_prepared`, `gemm_gate_up_mq5g256v2_wmma`, `gemm_gate_up_mq5g256v2_wmma_gfx11`, `gemm_gate_up_mq5g256v2_wmma_gfx12`, `gemm_gate_up_mq6g256v2_wmma`, `gemm_gate_up_mq6g256v2_wmma_gfx11`, `gemm_gate_up_mq6g256v2_wmma_gfx12`, `gemm_gate_up_mqv2_wmma_gfx11_bt`, `gemm_gate_up_mqv2_wmma_gfx11_mw_lds`, `gemm_gate_up_q8_0_wmma`, `gemm_gate_up_q8_0_wmma_gfx12`, `gemm_gate_up_silu_a4_mq4g256v2_iu4_prepared`, `gemm_gate_up_silu_mq4g256v2_i8_prepared`, `gemm_gate_up_silu_mq4g256v2_i8_prequant`, `gemm_gate_up_silu_mq4g256v2_iu4_prepared`, `gemm_hfp4g32_moe_grouped_wmma`, `gemm_hfp4g32_residual`, `gemm_hfp4g32_residual_wmma`, `gemm_hfp4g32_residual_wmma_gfx12`, `gemm_hfq2g256v2`, `gemm_hfq2g256v2_batched_lmhead`, `gemm_hfq3g256_batched_lmhead`, `gemm_hfq3g256_moe_grouped_wmma`, `gemm_hfq3g256_residual`, `gemm_hfq3g256_residual_dot2`, `gemm_hfq3g256_residual_fp16`, `gemm_hfq3g256_residual_mmq`, `gemm_hfq3g256_residual_mmq_x16`, `gemm_hfq3g256_residual_mmq_x32`, `gemm_hfq3g256_residual_mmq_x32_y32`, `gemm_hfq3g256_residual_mmq_x32_y64`, `gemm_hfq3g256_residual_mmq_x8`, `gemm_hfq3g256_residual_wmma`, `gemm_hfq3g256_residual_wmma_gfx12`, `gemm_hfq3g256_residual_wmma_mb4`, `gemm_hfq3g256v2`, `gemm_hfq3g256v2_batched_lmhead`, `gemm_hfq4g128`, `gemm_hfq4g256`, `gemm_hfq4g256_batched_lmhead`, `gemm_hfq4g256_dp4a`, `gemm_hfq4g256_lmhead_wmma_gfx12`, `gemm_hfq4g256_mmq_set`, `gemm_hfq4g256_mmq_set_gfx906`, `gemm_hfq4g256_mmq_set_prequant`, `gemm_hfq4g256_moe_grouped_mmq_gfx1151`, `gemm_hfq4g256_moe_grouped_mmq_gfx11_dgpu`, `gemm_hfq4g256_moe_grouped_mmq_gfx12`, `gemm_hfq4g256_moe_grouped_mmq_k4_gfx1151`, `gemm_hfq4g256_moe_grouped_mmq_k4_gfx11_dgpu`, `gemm_hfq4g256_moe_grouped_mmq_k4_gfx12`, `gemm_hfq4g256_moe_grouped_mmq_k8_gfx1151`, `gemm_hfq4g256_moe_grouped_wmma_k2`, `gemm_hfq4g256_moe_grouped_wmma_k2_fp16`, `gemm_hfq4g256_residual`, `gemm_hfq4g256_residual_fp16`, `gemm_hfq4g256_residual_fp16_wave64`, `gemm_hfq4g256_residual_mfma_gfx942`, `gemm_hfq4g256_residual_mfma_v2_gfx942`, `gemm_hfq4g256_residual_mfma_v3_gfx942`, `gemm_hfq4g256_residual_mfma_v4_gfx942`, `gemm_hfq4g256_residual_mmq`, `gemm_hfq4g256_residual_mmq_gfx906`, `gemm_hfq4g256_residual_mmq_lloyd`, `gemm_hfq4g256_residual_mmq_rdna2`, `gemm_hfq4g256_residual_mmq_rdna2_auto`, `gemm_hfq4g256_residual_mmq_x16`, `gemm_hfq4g256_residual_mmq_x32`, `gemm_hfq4g256_residual_mmq_x32_y64`, `gemm_hfq4g256_residual_mq4v2`, `gemm_hfq4g256_residual_muse`, `gemm_hfq4g256_residual_muse_gfx1100`, `gemm_hfq4g256_residual_muse_gfx1100_cb`, `gemm_hfq4g256_residual_muse_gfx1100_ksplit`, `gemm_hfq4g256_residual_muse_gfx1100_lds_g256`, `gemm_hfq4g256_residual_muse_gfx1100_mw`, `gemm_hfq4g256_residual_muse_gfx1100_rm`, `gemm_hfq4g256_residual_muse_gfx1100_rm2_pipe`, `gemm_hfq4g256_residual_muse_gfx1100_rm_hb`, `gemm_hfq4g256_residual_muse_gfx1100_rm_pk`, `gemm_hfq4g256_residual_wave64_dp4a`, `gemm_hfq4g256_residual_wave64_dp4a_prequant`, `gemm_hfq4g256_residual_wmma`, `gemm_hfq4g256_residual_wmma_gfx12`, `gemm_hfq4g256_residual_wmma_gfx12_mq4v2`, `gemm_hfq4g256_residual_wmma_gfx12_mq4v2_fp8`, `gemm_hfq4g256_residual_wmma_gfx12_mq4v2_fp8_lloyd`, `gemm_hfq4g256_residual_wmma_gfx12_mq4v2_fp8_prepared`, `gemm_hfq4g256_residual_wmma_gfx12_mq4v2_fp8_prepared_lloyd`, `gemm_hfq4g256_residual_wmma_ksplit_det`, `gemm_hfq4g256_wmma`, `gemm_hfq4g256v2`, `gemm_hfq4g256v2_batched_lmhead`, `gemm_hfq5g256v2`, `gemm_hfq5g256v2_batched_lmhead`, `gemm_hfq6g256_batched_lmhead`, `gemm_hfq6g256_moe_grouped_mmq_gfx1151`, `gemm_hfq6g256_moe_grouped_wmma`, `gemm_hfq6g256_residual`, `gemm_hfq6g256_residual_fp16`, `gemm_hfq6g256_residual_wave64_dp4a`, `gemm_hfq6g256_residual_wmma`, `gemm_hfq6g256_residual_wmma_gfx12`, `gemm_hfq6g256v2`, `gemm_hfq6g256v2_batched_lmhead`, `gemm_mfp4g32_e8_moe_grouped_wmma`, `gemm_mfp4g32_e8_soa_grouped_wmma`, `gemm_mfp4g32_e8_soa_grouped_wmma_b2`, `gemm_mfp4g32_e8_soa_grouped_wmma_gfx1201_f16`, `gemm_mfp4g32_e8_soa_wmma`, `gemm_mfp4g32_e8_soa_wmma_b16`, `gemm_mfp4g32_e8_soa_wmma_b16_gfx1201_f16`, `gemm_mfp4g32_e8_soa_wmma_b2`, `gemm_mfp4g32_e8_soa_wmma_b2_gfx1201_f16`, `gemm_mfp4g32_e8_soa_wmma_b4`, `gemm_mfp4g32_e8_soa_wmma_b4_gfx1201_f16`, `gemm_mfp4g32_e8_soa_wmma_b8`, `gemm_mfp4g32_e8_soa_wmma_b8_gfx1201_f16`, `gemm_mfp4g32_e8_soa_wmma_coop16`, `gemm_mfp4g32_e8_soa_wmma_coop4`, `gemm_mfp4g32_e8_soa_wmma_gfx1201_f16`, `gemm_mfp4g32_e8g128_moe_grouped_wmma`, `gemm_mixed_moe_grouped_wmma`, `gemm_mq2g256_lloyd_moe_grouped`, `gemm_mq2g256_lloyd_moe_grouped_mfma_gfx942`, `gemm_mq2g256_lloyd_moe_grouped_mmq_gfx1030`, `gemm_mq2g256_lloyd_moe_grouped_mmq_gfx1151`, `gemm_mq2g256_lloyd_moe_grouped_mmq_perm_gfx1151`, `gemm_mq2g256_lloyd_moe_grouped_mmq_prequant_perm_gfx1151`, `gemm_mq2g256_lloyd_moe_grouped_wmma`, `gemm_mq2g256_lloyd_moe_grouped_wmma_4w_k2`, `gemm_mq2g256_lloyd_moe_grouped_wmma_4w_k2_cnd`, `gemm_mq2g256_lloyd_moe_grouped_wmma_4w_k2_mmqload`, `gemm_mq2g256_lloyd_moe_grouped_wmma_4w_k2_mmqload_nosync`, `gemm_mq2g256_lloyd_moe_grouped_wmma_4w_k2_n32`, `gemm_mq2g256_lloyd_moe_grouped_wmma_8w_k2`, `gemm_mq2g256_lloyd_moe_grouped_wmma_k2`, `gemm_mq2g256v2`, `gemm_mq2g256v2_batched_lmhead`, `gemm_mq2g256v2_residual_wmma`, `gemm_mq2g256v2_residual_wmma_gfx11`, `gemm_mq2g256v2_residual_wmma_gfx12`, `gemm_mq3g256_lloyd_moe_grouped`, `gemm_mq3g256_lloyd_moe_grouped_mmq_gfx1151`, `gemm_mq3g256_lloyd_moe_grouped_wmma`, `gemm_mq3g256_lloyd_residual_wmma`, `gemm_mq3g256_lloyd_residual_wmma_mb4`, `gemm_mq3g256_residual_wmma`, `gemm_mq3g256v2`, `gemm_mq3g256v2_batched_lmhead`, `gemm_mq3g256v2_residual_wmma`, `gemm_mq3g256v2_residual_wmma_gfx11`, `gemm_mq3g256v2_residual_wmma_gfx12`, `gemm_mq4cg256`, `gemm_mq4cg256_batched_lmhead`, `gemm_mq4cg256_residual`, `gemm_mq4cg256_residual_wmma_gfx12`, `gemm_mq4g128v2_batched`, `gemm_mq4g256_lloyd_residual_wmma`, `gemm_mq4g256_lloyd_residual_wmma_mb2`, `gemm_mq4g256_lloyd_residual_wmma_mb4`, `gemm_mq4g256v2`, `gemm_mq4g256v2_batched_lmhead`, `gemm_mq4g256v2_hcw`, `gemm_mq4g256v2_hcw_applies`, `gemm_mq4g256v2_mmq_add_prequant`, `gemm_mq4g256v2_mmq_add_prequant_i8`, `gemm_mq4g256v2_mmq_add_prequant_iu4`, `gemm_mq4g256v2_mmq_add_prequant_lloyd`, `gemm_mq4g256v2_mmq_set_prequant`, `gemm_mq4g256v2_mmq_set_prequant_i8`, `gemm_mq4g256v2_mmq_set_prequant_iu4`, `gemm_mq4g256v2_mmq_set_prequant_lloyd`, `gemm_mq4g256v2_moe_grouped_top10`, `gemm_mq4g256v2_moe_grouped_top10_bf16out`, `gemm_mq4g256v2_moe_grouped_top10_silu_bf16out`, `gemm_mq4g256v2_moe_grouped_wmma_k2`, `gemm_mq4g256v2_pm_applies`, `gemm_mq4g256v2_residual_wmma`, `gemm_mq4g256v2_residual_wmma_gfx1100_bt`, `gemm_mq4g256v2_residual_wmma_gfx1100_ksplit_lds`, `gemm_mq4g256v2_residual_wmma_gfx1100_ldsstage`, `gemm_mq4g256v2_residual_wmma_gfx1100_mw_lds`, `gemm_mq4g256v2_residual_wmma_gfx1151_bt`, `gemm_mq4g256v2_residual_wmma_i8_prepared`, `gemm_mq4g256v2_residual_wmma_iu4_down_prepared`, `gemm_mq4g256v2_residual_wmma_iu4_prepared`, `gemm_mq4g256v2_xf16`, `gemm_mq4g256v2_xf16_applies`, `gemm_mq4g256v2_xf16_regions`, `gemm_mq5g256v2`, `gemm_mq5g256v2_batched_lmhead`, `gemm_mq5g256v2_residual_wmma`, `gemm_mq5g256v2_residual_wmma_gfx11`, `gemm_mq5g256v2_residual_wmma_gfx12`, `gemm_mq6g256_batched_lmhead`, `gemm_mq6g256v2`, `gemm_mq6g256v2_batched_lmhead`, `gemm_mq6g256v2_hcw`, `gemm_mq6g256v2_hcw_applies`, `gemm_mq6g256v2_moe_grouped_wmma_k2`, `gemm_mq6g256v2_residual_wmma`, `gemm_mq6g256v2_residual_wmma_gfx11`, `gemm_mq6g256v2_residual_wmma_gfx12`, `gemm_mq6g256v2_residual_wmma_gfx12_bt`, `gemm_mq6g256v2_xf16`, `gemm_mq6g256v2_xf16_applies`, `gemm_mq6g256v2_xf16_regions`, `gemm_mqv2_residual_wmma_gfx11_bt`, `gemm_mqv2_residual_wmma_gfx11_mw_lds`, `gemm_mw16_residual_wmma_f16x`, `gemm_paro_q4g128_moe_grouped_mmq_gfx1151`, `gemm_paro_q4g128_moe_grouped_mmq_k8_gfx1151`, `gemm_paro_q4g128_moe_grouped_wmma_k2`, `gemm_q8_0_batched`, `gemm_q8_0_batched_chunked`, `gemm_q8_0_batched_f32_chunked`, `gemm_q8_0_batched_wide_exact`, `gemm_q8_0_mmq_4w_gfx1151`, `gemm_q8_0_mmq_gfx1030`, `gemm_q8_0_residual_wmma`, `gemm_q8_0_residual_wmma_gfx12`, `gemm_q8_0_wmma`, `gemm_q8_0_wmma_4w`, `gemm_q8_0_wmma_prefill_auto`, `gemm_q8_0_wmma_x64`, `gemm_qkv_hfp4g32`, `gemm_qkv_hfp4g32_wmma`, `gemm_qkv_hfp4g32_wmma_fp8_gfx12`, `gemm_qkv_hfp4g32_wmma_gfx12`, `gemm_qkv_hfq3g256`, `gemm_qkv_hfq3g256_dot2`, `gemm_qkv_hfq3g256_dp4a`, `gemm_qkv_hfq3g256_fp16`, `gemm_qkv_hfq3g256_mmq`, `gemm_qkv_hfq3g256_mmq_x16`, `gemm_qkv_hfq3g256_mmq_x32`, `gemm_qkv_hfq3g256_mmq_x8`, `gemm_qkv_hfq3g256_wmma`, `gemm_qkv_hfq3g256_wmma_gfx12`, `gemm_qkv_hfq3g256_wmma_mb4`, `gemm_qkv_hfq4g256`, `gemm_qkv_hfq4g256_dot2`, `gemm_qkv_hfq4g256_fp16`, `gemm_qkv_hfq4g256_fp16_wave64`, `gemm_qkv_hfq4g256_mmq`, `gemm_qkv_hfq4g256_mmq_gfx906`, `gemm_qkv_hfq4g256_mmq_x16`, `gemm_qkv_hfq4g256_mmq_x32`, `gemm_qkv_hfq4g256_mq4v2`, `gemm_qkv_hfq4g256_wave64_dp4a`, `gemm_qkv_hfq4g256_wave64_dp4a_prequant`, `gemm_qkv_hfq4g256_wmma`, `gemm_qkv_hfq4g256_wmma_gfx12`, `gemm_qkv_hfq4g256_wmma_gfx12_mq4v2`, `gemm_qkv_hfq4g256_wmma_gfx12_mq4v2_fp8`, `gemm_qkv_hfq4g256_wmma_gfx12_mq4v2_fp8_lloyd`, `gemm_qkv_hfq4g256_wmma_gfx12_mq4v2_fp8_prepared`, `gemm_qkv_hfq4g256_wmma_gfx12_mq4v2_fp8_prepared_lloyd`, `gemm_qkv_hfq6g256`, `gemm_qkv_hfq6g256_dot2`, `gemm_qkv_hfq6g256_fp16`, `gemm_qkv_hfq6g256_wave64_dp4a`, `gemm_qkv_hfq6g256_wmma`, `gemm_qkv_hfq6g256_wmma_gfx12`, `gemm_qkv_mq2g256v2_wmma`, `gemm_qkv_mq2g256v2_wmma_gfx11`, `gemm_qkv_mq2g256v2_wmma_gfx12`, `gemm_qkv_mq3g256_lloyd_wmma`, `gemm_qkv_mq3g256_lloyd_wmma_mb4`, `gemm_qkv_mq3g256v2_wmma`, `gemm_qkv_mq3g256v2_wmma_gfx11`, `gemm_qkv_mq3g256v2_wmma_gfx12`, `gemm_qkv_mq4cg256`, `gemm_qkv_mq4cg256_wmma_gfx12`, `gemm_qkv_mq4g256_lloyd_wmma`, `gemm_qkv_mq4g256_lloyd_wmma_mb4`, `gemm_qkv_mq4g256v2_mmq_lloyd`, `gemm_qkv_mq4g256v2_wmma`, `gemm_qkv_mq4g256v2_wmma_gfx1100_bt`, `gemm_qkv_mq4g256v2_wmma_gfx1151_bt`, `gemm_qkv_mq4g256v2_wmma_gfx1201_bt`, `gemm_qkv_mq4g256v2_wmma_iu4_prepared`, `gemm_qkv_mq5g256v2_wmma`, `gemm_qkv_mq5g256v2_wmma_gfx11`, `gemm_qkv_mq5g256v2_wmma_gfx12`, `gemm_qkv_mq6g256v2_wmma`, `gemm_qkv_mq6g256v2_wmma_gfx11`, `gemm_qkv_mq6g256v2_wmma_gfx12`, `gemm_qkv_mqv2_wmma_gfx11_bt`, `gemm_qkv_mqv2_wmma_gfx1201_bt8`, `gemm_qkv_q8_0_wmma`, `gemm_qkv_q8_0_wmma_gfx12`, `gemm_qkvza_hfp4g32`, `gemm_qkvza_hfp4g32_wmma`, `gemm_qkvza_hfp4g32_wmma_gfx12`, `gemm_qkvza_hfq3g256`, `gemm_qkvza_hfq3g256_dot2`, `gemm_qkvza_hfq3g256_fp16`, `gemm_qkvza_hfq3g256_mmq`, `gemm_qkvza_hfq3g256_mmq_x16`, `gemm_qkvza_hfq3g256_mmq_x32`, `gemm_qkvza_hfq3g256_mmq_x8`, `gemm_qkvza_hfq3g256_wmma`, `gemm_qkvza_hfq3g256_wmma_gfx12`, `gemm_qkvza_hfq3g256_wmma_mb4`, `gemm_qkvza_hfq4g256`, `gemm_qkvza_hfq4g256_dot2`, `gemm_qkvza_hfq4g256_fp16`, `gemm_qkvza_hfq4g256_fp16_wave64`, `gemm_qkvza_hfq4g256_mmq`, `gemm_qkvza_hfq4g256_mmq_x16`, `gemm_qkvza_hfq4g256_mmq_x32`, `gemm_qkvza_hfq4g256_mq4v2`, `gemm_qkvza_hfq4g256_wave64_dp4a`, `gemm_qkvza_hfq4g256_wave64_dp4a_prequant`, `gemm_qkvza_hfq4g256_wmma`, `gemm_qkvza_hfq4g256_wmma_gfx12`, `gemm_qkvza_hfq4g256_wmma_gfx12_mq4v2`, `gemm_qkvza_hfq4g256_wmma_gfx12_mq4v2_fp8`, `gemm_qkvza_hfq4g256_wmma_gfx12_mq4v2_fp8_lloyd`, `gemm_qkvza_hfq4g256_wmma_gfx12_mq4v2_fp8_prepared`, `gemm_qkvza_hfq4g256_wmma_gfx12_mq4v2_fp8_prepared_lloyd`, `gemm_qkvza_hfq6g256`, `gemm_qkvza_hfq6g256_dot2`, `gemm_qkvza_hfq6g256_fp16`, `gemm_qkvza_hfq6g256_wave64_dp4a`, `gemm_qkvza_hfq6g256_wmma`, `gemm_qkvza_hfq6g256_wmma_gfx12`, `gemm_qkvza_mq2g256v2_wmma`, `gemm_qkvza_mq2g256v2_wmma_gfx11`, `gemm_qkvza_mq2g256v2_wmma_gfx12`, `gemm_qkvza_mq3g256_lloyd_wmma`, `gemm_qkvza_mq3g256_lloyd_wmma_mb4`, `gemm_qkvza_mq3g256_wmma`, `gemm_qkvza_mq3g256v2_wmma`, `gemm_qkvza_mq3g256v2_wmma_gfx11`, `gemm_qkvza_mq3g256v2_wmma_gfx12`, `gemm_qkvza_mq4cg256`, `gemm_qkvza_mq4cg256_wmma_gfx12`, `gemm_qkvza_mq4g256_lloyd_wmma`, `gemm_qkvza_mq4g256_lloyd_wmma_mb4`, `gemm_qkvza_mq4g256v2_fp8_gdn_prepared`, `gemm_qkvza_mq4g256v2_iu4_gdn_prepared`, `gemm_qkvza_mq4g256v2_mmq_lloyd`, `gemm_qkvza_mq4g256v2_wmma`, `gemm_qkvza_mq4g256v2_wmma_gfx1100_bt`, `gemm_qkvza_mq4g256v2_wmma_gfx1151_bt`, `gemm_qkvza_mq4g256v2_wmma_iu4_fold_prepared`, `gemm_qkvza_mq4g256v2_wmma_iu4_prepared`, `gemm_qkvza_mq5g256v2_wmma`, `gemm_qkvza_mq5g256v2_wmma_gfx11`, `gemm_qkvza_mq5g256v2_wmma_gfx12`, `gemm_qkvza_mq6g256v2_wmma`, `gemm_qkvza_mq6g256v2_wmma_gfx11`, `gemm_qkvza_mq6g256v2_wmma_gfx12`, `gemm_qkvza_mqv2_wmma_gfx11_bt`, `gemm_qkvza_q8_0_wmma`, `gemm_qkvza_q8_0_wmma_gfx12`, `gemm_qkvza_tq2g128`, `gemm_qwen4_moe_down_iu4_sym`, `gemm_qwen4_moe_gate_up_silu_iu4_sym`, `gemm_qwen4_trunk_zba_iu4`, `gemm_tq2g128_wmma`, `gemv_hfq4g128_moe_down_residual_scaled_bucketed`, `gemv_hfq4g128_moe_down_residual_scaled_k8_indexed`, `gemv_hfq4g128_moe_down_residual_scaled_k8_indexed_batched`, `gemv_hfq4g256_moe_gate_up_bucketed`, `gemv_mq4g256_moe_gate_up_bucketed`, `gemv_mq4g256_moe_gate_up_k8_indexed`, `gemv_mq4g256v2_xbatch`, `gemv_q8_0_moe_down_residual_scaled_k8_indexed`, `gemv_q8_0_moe_gate_up_k8_indexed`, `iu4_gdn_fold_active`, `iu4_gdn_fold_rows`, `label`, `lds_pipe_default`, `lds_pipe_gate`, `lds_tile_for`, `mask`, `moe_bucket_build`, `moe_scatter_stable_top10`, `mq4v2_fold_betaalpha_active`, `mq4v2_fold_betaalpha_padded_m`, `qwen4_bf16_streams`, `qwen4_convert_f16_x`, `qwen4_f16_wmma_applies`, `qwen4_f16_wmma_gfx1201`, `qwen4_f16_x_scratch`, `qwen4_lds_tile_gfx1201`, `qwen4_moe_gateup_wmma_applies`, `qwen4_moe_rotate128_i4`, `qwen4_moe_rotate256_i4`, `qwen4_moe_sym_check`, `qwen4_moe_sym_down_rr`, `qwen4_moe_sym_gemm_symbols`, `qwen4_moe_sym_iu4_applies`, `qwen4_moe_sym_iu4_requested`, `qwen4_mq6_x4_gfx1201`, `qwen4_mq6_x4_regions`, `qwen4_mq6_x4_tile`, `qwen4_shared_down_bf16_epi`, `qwen4_shared_gu_pm_applies`, `qwen4_shared_gu_pm_launch`, `qwen4_trunk_iu4_applies`, `qwen4_trunk_iu4_requested`, `qwen4_trunk_sym_check`, `reserve_qwen4_moe_sym`, `rocblas_gemm_bf16_prefill`, `rocblas_gemm_hfq4_prefill`, `rocblas_gemm_hfq4_prefill_residual`, `rocblas_gemm_mfp4e8_soa_prefill_auto` +- [`src/gemm.rs`](src/gemm.rs): `ADDIN`, `BF16_SCALED_ADD`, `F2GdnTargets`, `GATED`, `GELU`, `GemmEpilogue`, `HC_DOWN_GFX1201`, `LDS_EPI_TILES`, `LDS_TILE_VARIANTS`, `LOWBIT_WMMA_WAVES`, `LdsTile`, `LdsTileSplitK`, `OUT_F16`, `QWEN4_F16_WMMA_MIN_TOKENS`, `QWEN4_MOE_SYM_IU4_MIN_ROWS`, `QWEN4_MQ4_X4_MIN_N_GFX1151`, `QWEN4_MQ6_X4_PM_MIN_M_GFX1151`, `QWEN4_MQ6_X4_PM_MIN_N_GFX1151`, `QWEN4_SHARED_GU_PM_WAVES`, `SUPPORTED`, `a4_rms_fdiv_planes_wanted`, `a4_slab_active`, `arm_residual_fold`, `disarm_residual_fold`, `ensure_qwen4_moe_sym_pm_entry`, `entry`, `entry_epi`, `entry_suffix`, `flush_residual_fold`, `fn`, `fp8_f2_row_active`, `fp8_silu_h_active`, `fused_gate_up_hfq4g256`, `fused_gate_up_hfq4g256_dp4a`, `fused_gate_up_hfq4g256_mq4v2`, `fused_gate_up_hfq6g256_dp4a`, `fused_gate_up_mq2g256v2`, `fused_gate_up_mq3g256v2`, `fused_gate_up_mq4cg256`, `fused_gate_up_mq4g256v2_lloyd`, `fused_gate_up_mq5g256v2`, `fused_gate_up_mq6g256v2`, `fused_gate_up_q4k`, `fused_gate_up_q8_0`, `fused_qk_l2_norm_scale_f32`, `fused_qk_l2_norm_scale_f32_batched`, `fused_qk_l2_norm_scale_interleave_f32_batched`, `fused_qkv_hfq4g256`, `fused_qkv_hfq4g256_dp4a`, `fused_qkv_hfq4g256_mq4v2`, `fused_qkv_hfq4g256_mq4v2_generic_exact`, `fused_qkv_hfq4g256_mq4v2_k2048_x_buffer_gfx1100_exact`, `fused_qkv_hfq4g256_with_bias`, `fused_qkv_hfq4g256_with_bias_mq4v2`, `fused_qkv_hfq6g256`, `fused_qkv_hfq6g256_dp4a`, `fused_qkv_hfq6g256_with_bias`, `fused_qkv_mq2g256v2`, `fused_qkv_mq3g256v2`, `fused_qkv_mq4cg256`, `fused_qkv_mq4cg256_with_bias`, `fused_qkv_mq4g256v2_lloyd`, `fused_qkv_mq5g256v2`, `fused_qkv_mq6g256v2`, `fused_qkv_q4k`, `fused_qkv_q4k_with_bias`, `fused_qkv_q8_0`, `fused_qkv_q8_0_with_bias`, `fused_qkvza_hfq4g256`, `fused_qkvza_hfq4g256_dp4a`, `fused_qkvza_hfq4g256_mq4v2`, `fused_qkvza_hfq4g256_mq4v2_generic_exact`, `fused_qkvza_hfq4g256_mq4v2_k2048_hoist_x32_exact`, `fused_qkvza_hfq4g256_scalar_prep_gfx1100`, `fused_qkvza_hfq6g256_dp4a`, `fused_qkvza_mq2g256v2`, `fused_qkvza_mq3g256v2`, `fused_qkvza_mq4cg256`, `fused_qkvza_mq4g256v2_lloyd`, `fused_qkvza_mq5g256v2`, `fused_qkvza_mq6g256v2`, `fused_qkvza_q8_0`, `fused_sigmoid_alpha_gate_f32`, `fused_sigmoid_alpha_gate_f32_batched`, `gemm_bf16_mfma_gfx942`, `gemm_bf16_xf16_f16_wmma`, `gemm_bf16_xf16_f16_wmma_regions`, `gemm_bf16_xf32_batched`, `gemm_bf16_xf32_f16_wmma_qwen4`, `gemm_bf16_xf32_f16_wmma_qwen4_bf16st`, `gemm_bf16_xf32_f16_wmma_qwen4_groups`, `gemm_bf16_xf32_f16_wmma_qwen4_hcsd`, `gemm_bf16_xf32_f16_wmma_qwen4_hcsd_applies`, `gemm_bf16_xf32_f16_wmma_qwen4_shared_gu_pm`, `gemm_bf16_xf32_multirow`, `gemm_bq1g128_wmma`, `gemm_f16`, `gemm_f16_batched_lmhead`, `gemm_f16_bias`, `gemm_f16_tiled`, `gemm_f16_wmma`, `gemm_f16_wmma_mb4`, `gemm_f16_wmma_mb8`, `gemm_f16_x_f16_mfma_gfx942`, `gemm_f16_x_f16_wmma`, `gemm_f16_x_f16_wmma_lds`, `gemm_f16_x_f16_wmma_lds_auto`, `gemm_f16_x_f16_wmma_lds_auto_epi`, `gemm_f16_x_f16_wmma_lds_auto_epi_ld`, `gemm_f16_x_f16_wmma_lds_auto_ld`, `gemm_f16_x_f16_wmma_lds_epi`, `gemm_f16_x_f16_wmma_lds_epi_ld`, `gemm_f16_x_f16_wmma_lds_splitk`, `gemm_f16_x_f16_wmma_lds_tiled`, `gemm_f16_x_f16_wmma_lds_tiled_ld`, `gemm_f32_batched`, `gemm_f32_register_tiled`, `gemm_gate_up_hfp4g32`, `gemm_gate_up_hfp4g32_wmma`, `gemm_gate_up_hfp4g32_wmma_gfx12`, `gemm_gate_up_hfq3g256`, `gemm_gate_up_hfq3g256_dot2`, `gemm_gate_up_hfq3g256_dp4a`, `gemm_gate_up_hfq3g256_fp16`, `gemm_gate_up_hfq3g256_mmq`, `gemm_gate_up_hfq3g256_mmq_x16`, `gemm_gate_up_hfq3g256_mmq_x32`, `gemm_gate_up_hfq3g256_mmq_x32_y64`, `gemm_gate_up_hfq3g256_mmq_x32_y96`, `gemm_gate_up_hfq3g256_mmq_x8`, `gemm_gate_up_hfq3g256_wmma`, `gemm_gate_up_hfq3g256_wmma_gfx12`, `gemm_gate_up_hfq3g256_wmma_mb4`, `gemm_gate_up_hfq4g256`, `gemm_gate_up_hfq4g256_dot2`, `gemm_gate_up_hfq4g256_fp16`, `gemm_gate_up_hfq4g256_fp16_wave64`, `gemm_gate_up_hfq4g256_mmq`, `gemm_gate_up_hfq4g256_mmq_gfx906`, `gemm_gate_up_hfq4g256_mmq_gfx906_prequant`, `gemm_gate_up_hfq4g256_mmq_x16`, `gemm_gate_up_hfq4g256_mmq_x32`, `gemm_gate_up_hfq4g256_mq4v2`, `gemm_gate_up_hfq4g256_wave64_dp4a`, `gemm_gate_up_hfq4g256_wave64_dp4a_prequant`, `gemm_gate_up_hfq4g256_wmma`, `gemm_gate_up_hfq4g256_wmma_gfx12`, `gemm_gate_up_hfq4g256_wmma_gfx12_mq4v2`, `gemm_gate_up_hfq4g256_wmma_gfx12_mq4v2_fp8_bt12`, `gemm_gate_up_hfq4g256_wmma_gfx12_mq4v2_fp8_bt12_prepared`, `gemm_gate_up_hfq4g256_wmma_gfx12_mq4v2_fp8_bt12_prepared_lloyd`, `gemm_gate_up_hfq4g256_wmma_gfx12_mq4v2_fp8_lloyd`, `gemm_gate_up_hfq6g256`, `gemm_gate_up_hfq6g256_dot2`, `gemm_gate_up_hfq6g256_fp16`, `gemm_gate_up_hfq6g256_wave64_dp4a`, `gemm_gate_up_hfq6g256_wmma`, `gemm_gate_up_hfq6g256_wmma_gfx12`, `gemm_gate_up_mq2g256v2_wmma`, `gemm_gate_up_mq2g256v2_wmma_gfx11`, `gemm_gate_up_mq2g256v2_wmma_gfx12`, `gemm_gate_up_mq3g256_lloyd_wmma`, `gemm_gate_up_mq3g256_lloyd_wmma_mb4`, `gemm_gate_up_mq3g256_wmma`, `gemm_gate_up_mq3g256v2_wmma`, `gemm_gate_up_mq3g256v2_wmma_gfx11`, `gemm_gate_up_mq3g256v2_wmma_gfx12`, `gemm_gate_up_mq4cg256`, `gemm_gate_up_mq4cg256_wmma_gfx12`, `gemm_gate_up_mq4g256_lloyd_wmma`, `gemm_gate_up_mq4g256_lloyd_wmma_mb4`, `gemm_gate_up_mq4g256v2_mmq_lloyd`, `gemm_gate_up_mq4g256v2_wmma`, `gemm_gate_up_mq4g256v2_wmma_gfx1100_bt`, `gemm_gate_up_mq4g256v2_wmma_gfx1100_mw_lds`, `gemm_gate_up_mq4g256v2_wmma_gfx1151_bt`, `gemm_gate_up_mq4g256v2_wmma_iu4_prepared`, `gemm_gate_up_mq5g256v2_wmma`, `gemm_gate_up_mq5g256v2_wmma_gfx11`, `gemm_gate_up_mq5g256v2_wmma_gfx12`, `gemm_gate_up_mq6g256v2_wmma`, `gemm_gate_up_mq6g256v2_wmma_gfx11`, `gemm_gate_up_mq6g256v2_wmma_gfx12`, `gemm_gate_up_mqv2_wmma_gfx11_bt`, `gemm_gate_up_mqv2_wmma_gfx11_mw_lds`, `gemm_gate_up_q8_0_wmma`, `gemm_gate_up_q8_0_wmma_gfx12`, `gemm_gate_up_silu_a4_mq4g256v2_iu4_prepared`, `gemm_gate_up_silu_mq4g256v2_i8_prepared`, `gemm_gate_up_silu_mq4g256v2_i8_prequant`, `gemm_gate_up_silu_mq4g256v2_iu4_prepared`, `gemm_hfp4g32_moe_grouped_wmma`, `gemm_hfp4g32_residual`, `gemm_hfp4g32_residual_wmma`, `gemm_hfp4g32_residual_wmma_gfx12`, `gemm_hfq2g256v2`, `gemm_hfq2g256v2_batched_lmhead`, `gemm_hfq3g256_batched_lmhead`, `gemm_hfq3g256_moe_grouped_wmma`, `gemm_hfq3g256_residual`, `gemm_hfq3g256_residual_dot2`, `gemm_hfq3g256_residual_fp16`, `gemm_hfq3g256_residual_mmq`, `gemm_hfq3g256_residual_mmq_x16`, `gemm_hfq3g256_residual_mmq_x32`, `gemm_hfq3g256_residual_mmq_x32_y32`, `gemm_hfq3g256_residual_mmq_x32_y64`, `gemm_hfq3g256_residual_mmq_x8`, `gemm_hfq3g256_residual_wmma`, `gemm_hfq3g256_residual_wmma_gfx12`, `gemm_hfq3g256_residual_wmma_mb4`, `gemm_hfq3g256v2`, `gemm_hfq3g256v2_batched_lmhead`, `gemm_hfq4g128`, `gemm_hfq4g256`, `gemm_hfq4g256_batched_lmhead`, `gemm_hfq4g256_dp4a`, `gemm_hfq4g256_lmhead_wmma_gfx12`, `gemm_hfq4g256_mmq_set`, `gemm_hfq4g256_mmq_set_gfx906`, `gemm_hfq4g256_mmq_set_prequant`, `gemm_hfq4g256_moe_grouped_mmq_gfx1151`, `gemm_hfq4g256_moe_grouped_mmq_gfx11_dgpu`, `gemm_hfq4g256_moe_grouped_mmq_gfx12`, `gemm_hfq4g256_moe_grouped_mmq_k4_gfx1151`, `gemm_hfq4g256_moe_grouped_mmq_k4_gfx11_dgpu`, `gemm_hfq4g256_moe_grouped_mmq_k4_gfx12`, `gemm_hfq4g256_moe_grouped_mmq_k8_gfx1151`, `gemm_hfq4g256_moe_grouped_wmma_k2`, `gemm_hfq4g256_moe_grouped_wmma_k2_fp16`, `gemm_hfq4g256_residual`, `gemm_hfq4g256_residual_fp16`, `gemm_hfq4g256_residual_fp16_wave64`, `gemm_hfq4g256_residual_mfma_gfx942`, `gemm_hfq4g256_residual_mfma_v2_gfx942`, `gemm_hfq4g256_residual_mfma_v3_gfx942`, `gemm_hfq4g256_residual_mfma_v4_gfx942`, `gemm_hfq4g256_residual_mmq`, `gemm_hfq4g256_residual_mmq_gfx906`, `gemm_hfq4g256_residual_mmq_lloyd`, `gemm_hfq4g256_residual_mmq_rdna2`, `gemm_hfq4g256_residual_mmq_rdna2_auto`, `gemm_hfq4g256_residual_mmq_x16`, `gemm_hfq4g256_residual_mmq_x32`, `gemm_hfq4g256_residual_mmq_x32_y64`, `gemm_hfq4g256_residual_mq4v2`, `gemm_hfq4g256_residual_muse`, `gemm_hfq4g256_residual_muse_gfx1100`, `gemm_hfq4g256_residual_muse_gfx1100_cb`, `gemm_hfq4g256_residual_muse_gfx1100_ksplit`, `gemm_hfq4g256_residual_muse_gfx1100_lds_g256`, `gemm_hfq4g256_residual_muse_gfx1100_mw`, `gemm_hfq4g256_residual_muse_gfx1100_rm`, `gemm_hfq4g256_residual_muse_gfx1100_rm2_pipe`, `gemm_hfq4g256_residual_muse_gfx1100_rm_hb`, `gemm_hfq4g256_residual_muse_gfx1100_rm_pk`, `gemm_hfq4g256_residual_wave64_dp4a`, `gemm_hfq4g256_residual_wave64_dp4a_prequant`, `gemm_hfq4g256_residual_wmma`, `gemm_hfq4g256_residual_wmma_gfx12`, `gemm_hfq4g256_residual_wmma_gfx12_mq4v2`, `gemm_hfq4g256_residual_wmma_gfx12_mq4v2_fp8`, `gemm_hfq4g256_residual_wmma_gfx12_mq4v2_fp8_lloyd`, `gemm_hfq4g256_residual_wmma_gfx12_mq4v2_fp8_prepared`, `gemm_hfq4g256_residual_wmma_gfx12_mq4v2_fp8_prepared_lloyd`, `gemm_hfq4g256_residual_wmma_ksplit_det`, `gemm_hfq4g256_wmma`, `gemm_hfq4g256v2`, `gemm_hfq4g256v2_batched_lmhead`, `gemm_hfq5g256v2`, `gemm_hfq5g256v2_batched_lmhead`, `gemm_hfq6g256_batched_lmhead`, `gemm_hfq6g256_moe_grouped_mmq_gfx1151`, `gemm_hfq6g256_moe_grouped_wmma`, `gemm_hfq6g256_residual`, `gemm_hfq6g256_residual_fp16`, `gemm_hfq6g256_residual_wave64_dp4a`, `gemm_hfq6g256_residual_wmma`, `gemm_hfq6g256_residual_wmma_gfx12`, `gemm_hfq6g256v2`, `gemm_hfq6g256v2_batched_lmhead`, `gemm_mfp4g32_e8_moe_grouped_wmma`, `gemm_mfp4g32_e8_soa_grouped_wmma`, `gemm_mfp4g32_e8_soa_grouped_wmma_b2`, `gemm_mfp4g32_e8_soa_grouped_wmma_gfx1201_f16`, `gemm_mfp4g32_e8_soa_wmma`, `gemm_mfp4g32_e8_soa_wmma_b16`, `gemm_mfp4g32_e8_soa_wmma_b16_gfx1201_f16`, `gemm_mfp4g32_e8_soa_wmma_b2`, `gemm_mfp4g32_e8_soa_wmma_b2_gfx1201_f16`, `gemm_mfp4g32_e8_soa_wmma_b4`, `gemm_mfp4g32_e8_soa_wmma_b4_gfx1201_f16`, `gemm_mfp4g32_e8_soa_wmma_b8`, `gemm_mfp4g32_e8_soa_wmma_b8_gfx1201_f16`, `gemm_mfp4g32_e8_soa_wmma_coop16`, `gemm_mfp4g32_e8_soa_wmma_coop4`, `gemm_mfp4g32_e8_soa_wmma_gfx1201_f16`, `gemm_mfp4g32_e8g128_moe_grouped_wmma`, `gemm_mixed_moe_grouped_wmma`, `gemm_mq2g256_lloyd_moe_grouped`, `gemm_mq2g256_lloyd_moe_grouped_mfma_gfx942`, `gemm_mq2g256_lloyd_moe_grouped_mmq_gfx1030`, `gemm_mq2g256_lloyd_moe_grouped_mmq_gfx1151`, `gemm_mq2g256_lloyd_moe_grouped_mmq_perm_gfx1151`, `gemm_mq2g256_lloyd_moe_grouped_mmq_prequant_perm_gfx1151`, `gemm_mq2g256_lloyd_moe_grouped_wmma`, `gemm_mq2g256_lloyd_moe_grouped_wmma_4w_k2`, `gemm_mq2g256_lloyd_moe_grouped_wmma_4w_k2_cnd`, `gemm_mq2g256_lloyd_moe_grouped_wmma_4w_k2_mmqload`, `gemm_mq2g256_lloyd_moe_grouped_wmma_4w_k2_mmqload_nosync`, `gemm_mq2g256_lloyd_moe_grouped_wmma_4w_k2_n32`, `gemm_mq2g256_lloyd_moe_grouped_wmma_8w_k2`, `gemm_mq2g256_lloyd_moe_grouped_wmma_k2`, `gemm_mq2g256v2`, `gemm_mq2g256v2_batched_lmhead`, `gemm_mq2g256v2_residual_wmma`, `gemm_mq2g256v2_residual_wmma_gfx11`, `gemm_mq2g256v2_residual_wmma_gfx12`, `gemm_mq3g256_lloyd_moe_grouped`, `gemm_mq3g256_lloyd_moe_grouped_mmq_gfx1151`, `gemm_mq3g256_lloyd_moe_grouped_wmma`, `gemm_mq3g256_lloyd_residual_wmma`, `gemm_mq3g256_lloyd_residual_wmma_mb4`, `gemm_mq3g256_residual_wmma`, `gemm_mq3g256v2`, `gemm_mq3g256v2_batched_lmhead`, `gemm_mq3g256v2_residual_wmma`, `gemm_mq3g256v2_residual_wmma_gfx11`, `gemm_mq3g256v2_residual_wmma_gfx12`, `gemm_mq4cg256`, `gemm_mq4cg256_batched_lmhead`, `gemm_mq4cg256_residual`, `gemm_mq4cg256_residual_wmma_gfx12`, `gemm_mq4g128v2_batched`, `gemm_mq4g256_lloyd_residual_wmma`, `gemm_mq4g256_lloyd_residual_wmma_mb2`, `gemm_mq4g256_lloyd_residual_wmma_mb4`, `gemm_mq4g256v2`, `gemm_mq4g256v2_batched_lmhead`, `gemm_mq4g256v2_hcw`, `gemm_mq4g256v2_hcw_applies`, `gemm_mq4g256v2_mmq_add_prequant`, `gemm_mq4g256v2_mmq_add_prequant_i8`, `gemm_mq4g256v2_mmq_add_prequant_iu4`, `gemm_mq4g256v2_mmq_add_prequant_lloyd`, `gemm_mq4g256v2_mmq_set_prequant`, `gemm_mq4g256v2_mmq_set_prequant_i8`, `gemm_mq4g256v2_mmq_set_prequant_iu4`, `gemm_mq4g256v2_mmq_set_prequant_lloyd`, `gemm_mq4g256v2_moe_grouped_top10`, `gemm_mq4g256v2_moe_grouped_top10_bf16out`, `gemm_mq4g256v2_moe_grouped_top10_silu_bf16out`, `gemm_mq4g256v2_moe_grouped_wmma_k2`, `gemm_mq4g256v2_pm_applies`, `gemm_mq4g256v2_residual_wmma`, `gemm_mq4g256v2_residual_wmma_gfx1100_bt`, `gemm_mq4g256v2_residual_wmma_gfx1100_ksplit_lds`, `gemm_mq4g256v2_residual_wmma_gfx1100_ldsstage`, `gemm_mq4g256v2_residual_wmma_gfx1100_mw_lds`, `gemm_mq4g256v2_residual_wmma_gfx1151_bt`, `gemm_mq4g256v2_residual_wmma_i8_prepared`, `gemm_mq4g256v2_residual_wmma_iu4_down_prepared`, `gemm_mq4g256v2_residual_wmma_iu4_prepared`, `gemm_mq4g256v2_xf16`, `gemm_mq4g256v2_xf16_applies`, `gemm_mq4g256v2_xf16_regions`, `gemm_mq5g256v2`, `gemm_mq5g256v2_batched_lmhead`, `gemm_mq5g256v2_residual_wmma`, `gemm_mq5g256v2_residual_wmma_gfx11`, `gemm_mq5g256v2_residual_wmma_gfx12`, `gemm_mq6g256_batched_lmhead`, `gemm_mq6g256v2`, `gemm_mq6g256v2_batched_lmhead`, `gemm_mq6g256v2_hcw`, `gemm_mq6g256v2_hcw_applies`, `gemm_mq6g256v2_moe_grouped_wmma_k2`, `gemm_mq6g256v2_residual_wmma`, `gemm_mq6g256v2_residual_wmma_gfx11`, `gemm_mq6g256v2_residual_wmma_gfx12`, `gemm_mq6g256v2_residual_wmma_gfx12_bt`, `gemm_mq6g256v2_xf16`, `gemm_mq6g256v2_xf16_applies`, `gemm_mq6g256v2_xf16_regions`, `gemm_mqv2_residual_wmma_gfx11_bt`, `gemm_mqv2_residual_wmma_gfx11_mw_lds`, `gemm_mw16_residual_wmma_f16x`, `gemm_paro_q4g128_moe_grouped_mmq_gfx1151`, `gemm_paro_q4g128_moe_grouped_mmq_k8_gfx1151`, `gemm_paro_q4g128_moe_grouped_wmma_k2`, `gemm_q8_0_batched`, `gemm_q8_0_batched_chunked`, `gemm_q8_0_batched_f32_chunked`, `gemm_q8_0_batched_wide_exact`, `gemm_q8_0_mmq_4w_gfx1151`, `gemm_q8_0_mmq_gfx1030`, `gemm_q8_0_residual_wmma`, `gemm_q8_0_residual_wmma_gfx12`, `gemm_q8_0_wmma`, `gemm_q8_0_wmma_4w`, `gemm_q8_0_wmma_prefill_auto`, `gemm_q8_0_wmma_x64`, `gemm_qkv_hfp4g32`, `gemm_qkv_hfp4g32_wmma`, `gemm_qkv_hfp4g32_wmma_fp8_gfx12`, `gemm_qkv_hfp4g32_wmma_gfx12`, `gemm_qkv_hfq3g256`, `gemm_qkv_hfq3g256_dot2`, `gemm_qkv_hfq3g256_dp4a`, `gemm_qkv_hfq3g256_fp16`, `gemm_qkv_hfq3g256_mmq`, `gemm_qkv_hfq3g256_mmq_x16`, `gemm_qkv_hfq3g256_mmq_x32`, `gemm_qkv_hfq3g256_mmq_x8`, `gemm_qkv_hfq3g256_wmma`, `gemm_qkv_hfq3g256_wmma_gfx12`, `gemm_qkv_hfq3g256_wmma_mb4`, `gemm_qkv_hfq4g256`, `gemm_qkv_hfq4g256_dot2`, `gemm_qkv_hfq4g256_fp16`, `gemm_qkv_hfq4g256_fp16_wave64`, `gemm_qkv_hfq4g256_mmq`, `gemm_qkv_hfq4g256_mmq_gfx906`, `gemm_qkv_hfq4g256_mmq_x16`, `gemm_qkv_hfq4g256_mmq_x32`, `gemm_qkv_hfq4g256_mq4v2`, `gemm_qkv_hfq4g256_wave64_dp4a`, `gemm_qkv_hfq4g256_wave64_dp4a_prequant`, `gemm_qkv_hfq4g256_wmma`, `gemm_qkv_hfq4g256_wmma_gfx12`, `gemm_qkv_hfq4g256_wmma_gfx12_mq4v2`, `gemm_qkv_hfq4g256_wmma_gfx12_mq4v2_fp8`, `gemm_qkv_hfq4g256_wmma_gfx12_mq4v2_fp8_lloyd`, `gemm_qkv_hfq4g256_wmma_gfx12_mq4v2_fp8_prepared`, `gemm_qkv_hfq4g256_wmma_gfx12_mq4v2_fp8_prepared_lloyd`, `gemm_qkv_hfq6g256`, `gemm_qkv_hfq6g256_dot2`, `gemm_qkv_hfq6g256_fp16`, `gemm_qkv_hfq6g256_wave64_dp4a`, `gemm_qkv_hfq6g256_wmma`, `gemm_qkv_hfq6g256_wmma_gfx12`, `gemm_qkv_mq2g256v2_wmma`, `gemm_qkv_mq2g256v2_wmma_gfx11`, `gemm_qkv_mq2g256v2_wmma_gfx12`, `gemm_qkv_mq3g256_lloyd_wmma`, `gemm_qkv_mq3g256_lloyd_wmma_mb4`, `gemm_qkv_mq3g256v2_wmma`, `gemm_qkv_mq3g256v2_wmma_gfx11`, `gemm_qkv_mq3g256v2_wmma_gfx12`, `gemm_qkv_mq4cg256`, `gemm_qkv_mq4cg256_wmma_gfx12`, `gemm_qkv_mq4g256_lloyd_wmma`, `gemm_qkv_mq4g256_lloyd_wmma_mb4`, `gemm_qkv_mq4g256v2_mmq_lloyd`, `gemm_qkv_mq4g256v2_wmma`, `gemm_qkv_mq4g256v2_wmma_gfx1100_bt`, `gemm_qkv_mq4g256v2_wmma_gfx1151_bt`, `gemm_qkv_mq4g256v2_wmma_gfx1201_bt`, `gemm_qkv_mq4g256v2_wmma_iu4_prepared`, `gemm_qkv_mq5g256v2_wmma`, `gemm_qkv_mq5g256v2_wmma_gfx11`, `gemm_qkv_mq5g256v2_wmma_gfx12`, `gemm_qkv_mq6g256v2_wmma`, `gemm_qkv_mq6g256v2_wmma_gfx11`, `gemm_qkv_mq6g256v2_wmma_gfx12`, `gemm_qkv_mqv2_wmma_gfx11_bt`, `gemm_qkv_mqv2_wmma_gfx1201_bt8`, `gemm_qkv_q8_0_wmma`, `gemm_qkv_q8_0_wmma_gfx12`, `gemm_qkvza_hfp4g32`, `gemm_qkvza_hfp4g32_wmma`, `gemm_qkvza_hfp4g32_wmma_gfx12`, `gemm_qkvza_hfq3g256`, `gemm_qkvza_hfq3g256_dot2`, `gemm_qkvza_hfq3g256_fp16`, `gemm_qkvza_hfq3g256_mmq`, `gemm_qkvza_hfq3g256_mmq_x16`, `gemm_qkvza_hfq3g256_mmq_x32`, `gemm_qkvza_hfq3g256_mmq_x8`, `gemm_qkvza_hfq3g256_wmma`, `gemm_qkvza_hfq3g256_wmma_gfx12`, `gemm_qkvza_hfq3g256_wmma_mb4`, `gemm_qkvza_hfq4g256`, `gemm_qkvza_hfq4g256_dot2`, `gemm_qkvza_hfq4g256_fp16`, `gemm_qkvza_hfq4g256_fp16_wave64`, `gemm_qkvza_hfq4g256_mmq`, `gemm_qkvza_hfq4g256_mmq_x16`, `gemm_qkvza_hfq4g256_mmq_x32`, `gemm_qkvza_hfq4g256_mq4v2`, `gemm_qkvza_hfq4g256_wave64_dp4a`, `gemm_qkvza_hfq4g256_wave64_dp4a_prequant`, `gemm_qkvza_hfq4g256_wmma`, `gemm_qkvza_hfq4g256_wmma_gfx12`, `gemm_qkvza_hfq4g256_wmma_gfx12_mq4v2`, `gemm_qkvza_hfq4g256_wmma_gfx12_mq4v2_fp8`, `gemm_qkvza_hfq4g256_wmma_gfx12_mq4v2_fp8_lloyd`, `gemm_qkvza_hfq4g256_wmma_gfx12_mq4v2_fp8_prepared`, `gemm_qkvza_hfq4g256_wmma_gfx12_mq4v2_fp8_prepared_lloyd`, `gemm_qkvza_hfq6g256`, `gemm_qkvza_hfq6g256_dot2`, `gemm_qkvza_hfq6g256_fp16`, `gemm_qkvza_hfq6g256_wave64_dp4a`, `gemm_qkvza_hfq6g256_wmma`, `gemm_qkvza_hfq6g256_wmma_gfx12`, `gemm_qkvza_mq2g256v2_wmma`, `gemm_qkvza_mq2g256v2_wmma_gfx11`, `gemm_qkvza_mq2g256v2_wmma_gfx12`, `gemm_qkvza_mq3g256_lloyd_wmma`, `gemm_qkvza_mq3g256_lloyd_wmma_mb4`, `gemm_qkvza_mq3g256_wmma`, `gemm_qkvza_mq3g256v2_wmma`, `gemm_qkvza_mq3g256v2_wmma_gfx11`, `gemm_qkvza_mq3g256v2_wmma_gfx12`, `gemm_qkvza_mq4cg256`, `gemm_qkvza_mq4cg256_wmma_gfx12`, `gemm_qkvza_mq4g256_lloyd_wmma`, `gemm_qkvza_mq4g256_lloyd_wmma_mb4`, `gemm_qkvza_mq4g256v2_fp8_gdn_prepared`, `gemm_qkvza_mq4g256v2_iu4_gdn_prepared`, `gemm_qkvza_mq4g256v2_mmq_lloyd`, `gemm_qkvza_mq4g256v2_wmma`, `gemm_qkvza_mq4g256v2_wmma_gfx1100_bt`, `gemm_qkvza_mq4g256v2_wmma_gfx1151_bt`, `gemm_qkvza_mq4g256v2_wmma_iu4_fold_prepared`, `gemm_qkvza_mq4g256v2_wmma_iu4_prepared`, `gemm_qkvza_mq5g256v2_wmma`, `gemm_qkvza_mq5g256v2_wmma_gfx11`, `gemm_qkvza_mq5g256v2_wmma_gfx12`, `gemm_qkvza_mq6g256v2_wmma`, `gemm_qkvza_mq6g256v2_wmma_gfx11`, `gemm_qkvza_mq6g256v2_wmma_gfx12`, `gemm_qkvza_mqv2_wmma_gfx11_bt`, `gemm_qkvza_q8_0_wmma`, `gemm_qkvza_q8_0_wmma_gfx12`, `gemm_qkvza_tq2g128`, `gemm_qwen4_moe_down_iu4_sym`, `gemm_qwen4_moe_gate_up_silu_iu4_sym`, `gemm_qwen4_trunk_zba_iu4`, `gemm_tq2g128_wmma`, `gemv_hfq4g128_moe_down_grouped`, `gemv_hfq4g128_moe_down_residual_scaled_bucketed`, `gemv_hfq4g128_moe_down_residual_scaled_k8_indexed`, `gemv_hfq4g128_moe_down_residual_scaled_k8_indexed_batched`, `gemv_hfq4g256_moe_gate_up_bucketed`, `gemv_mq4g256_moe_gate_up_bucketed`, `gemv_mq4g256_moe_gate_up_k8_indexed`, `gemv_mq4g256v2_xbatch`, `gemv_q8_0_moe_down_residual_scaled_k8_indexed`, `gemv_q8_0_moe_gate_up_k8_indexed`, `iu4_gdn_fold_active`, `iu4_gdn_fold_rows`, `label`, `lds_pipe_default`, `lds_pipe_gate`, `lds_tile_for`, `mask`, `moe_bucket_build`, `moe_scatter_stable_top10`, `mq4v2_fold_betaalpha_active`, `mq4v2_fold_betaalpha_padded_m`, `qwen4_bf16_streams`, `qwen4_convert_f16_x`, `qwen4_f16_wmma_applies`, `qwen4_f16_wmma_gfx1201`, `qwen4_f16_x_scratch`, `qwen4_lds_tile_gfx1201`, `qwen4_moe_gateup_wmma_applies`, `qwen4_moe_rotate128_i4`, `qwen4_moe_rotate256_i4`, `qwen4_moe_sym_check`, `qwen4_moe_sym_down_rr`, `qwen4_moe_sym_gemm_symbols`, `qwen4_moe_sym_iu4_applies`, `qwen4_moe_sym_iu4_requested`, `qwen4_mq6_x4_gfx1201`, `qwen4_mq6_x4_regions`, `qwen4_mq6_x4_tile`, `qwen4_shared_down_bf16_epi`, `qwen4_shared_gu_pm_applies`, `qwen4_shared_gu_pm_launch`, `qwen4_trunk_iu4_applies`, `qwen4_trunk_iu4_requested`, `qwen4_trunk_sym_check`, `reserve_qwen4_moe_sym`, `rocblas_gemm_bf16_prefill`, `rocblas_gemm_hfq4_prefill`, `rocblas_gemm_hfq4_prefill_residual`, `rocblas_gemm_mfp4e8_soa_prefill_auto` - [`src/gemm_fp8_verify.rs`](src/gemm_fp8_verify.rs): `Fp8VerifyOp`, `Fp8VerifyReceipt`, `Fp8VerifyTarget`, `Fp8VerifyTile`, `Fp8VerifyX`, `fn`, `fp8_verify_enter`, `fp8_verify_exit`, `fp8_verify_pack_x`, `fp8_verify_receipt`, `fp8_verify_symbol`, `fp8_verify_tile`, `gemm_gate_up_mq4g256v2_verify_fp8`, `gemm_mq4g256v2_lmhead_verify_fp8`, `gemm_mq4g256v2_residual_verify_fp8`, `gemm_mq4g256v2_verify_fp8_prepacked`, `gemm_qkv_mq4g256v2_verify_fp8`, `gemm_qkvza_mq4g256v2_verify_fp8` - [`src/gemm_vt.rs`](src/gemm_vt.rs): `gemm_gate_up_mq4g256v2_verify_exact`, `gemm_mq4g256v2_lmhead_verify_exact`, `gemm_mq4g256v2_residual_verify_exact`, `gemm_qkv_mq4g256v2_verify_exact`, `gemm_qkvza_mq4g256v2_verify_exact`, `mq4_verify_chunk_rows`, `mq4_verify_chunk_rows_for`, `mq4_verify_pm_selected`, `mq4_verify_singleton_chain`, `mq4_verify_singleton_chain_for` - [`src/gemma4_ext.rs`](src/gemma4_ext.rs): `attention_flash_asym2_window`, `attention_flash_asym3_batched_window`, `attention_flash_asym3_hd512`, `attention_flash_asym3_window`, `attention_flash_asym4_window`, `attention_flash_fwht3_hd512`, `attention_flash_q8_0_window`, `gemv_hfq4g128_moe_down_residual_scaled_bucketed`, `gemv_hfq4g128_moe_down_residual_scaled_k8_indexed`, `gemv_hfq4g128_moe_down_residual_scaled_k8_indexed_batched`, `gemv_hfq4g256_moe_gate_up_bucketed`, `gemv_mq4g256_moe_gate_up_bucketed`, `gemv_mq4g256_moe_gate_up_k8_indexed`, `gemv_q8_0_moe_down_residual_scaled_k8_indexed`, `gemv_q8_0_moe_gate_up_k8_indexed`, `kv_cache_write_asym3_hd512`, `kv_cache_write_fwht3_hd512`, `moe_bucket_build` @@ -55,7 +55,7 @@ _Generated by `scripts/check-crate-maps.py` from the tree — do not edit inside - [`src/hc_row_fold.rs`](src/hc_row_fold.rs): `HcRowFold`, `hc_row_fold_applies`, `hc_row_fold_norm_gate`, `moe_combine_order_top10` - [`src/kernel_pack.rs`](src/kernel_pack.rs): `MANIFEST_FILE`, `PACK_FORMAT`, `PackManifest`, `VerifiedPack`, `asset_stem`, `check`, `check_verified`, `code_object_version`, `crate`, `default_rocm_range`, `rocm_admitted`, `verify_dir` - [`src/kernel_registry.rs`](src/kernel_registry.rs): `ALPHA_GATE_SRC`, `CONV1D_SILU_SPLIT_SRC`, `CONV1D_SILU_SPLIT_TREE_SRC`, `CONV1D_SILU_SRC`, `FUSED_QK_L2_NORM_SCALE_SRC`, `FUSED_SIGMOID_ALPHA_GATE_SRC`, `GATED_DELTA_NET_Q8_SRC`, `GATED_DELTA_NET_Q8_TREE_SRC`, `GATED_NORM_SRC`, `KernelEntry`, `KernelSpecRoute`, `KernelTensorMeta`, `L2_NORM_SRC`, `PlannedKernel`, `QWEN35_FA_PREP_GFX1100_SRC`, `QWEN35_FA_PREP_GFX1151_SRC`, `QWEN35_FA_PREP_GFX1201_SRC`, `RAILGUN_COPY_SRC`, `ROPE_PARTIAL_HALFSPLIT_BATCHED_SRC`, `ROPE_PARTIAL_HALFSPLIT_HEADGRID_SRC`, `ROPE_PARTIAL_HALFSPLIT_SRC`, `ROPE_PARTIAL_INTERLEAVED_SRC`, `RegistryError`, `RouteKernelInput`, `RouteKernelPlan`, `SCALE_F32_SRC`, `SIGMOID_SRC`, `SUPPORTED_ARCHES`, `corpus_entries`, `crate`, `default_route_entries`, `entries`, `lookup`, `module`, `railgun_entries`, `route_entries`, `source`, `symbols`, `validate_route_plan` -- [`src/kernels.rs`](src/kernels.rs): `ACT_QUANT_FP4_UE8M0_G32_INPLACE_GFX942_SRC`, `ACT_QUANT_FP8_UE8M0_INPLACE_GFX942_SRC`, `ADD_INPLACE_SRC`, `ADD_PEER_RESIDUAL_SRC`, `ADD_SRC`, `ALPHA_GATE_SRC`, `APPLY_ROPE_2D_VISION_SRC`, `ARGMAX_BATCHED_SRC`, `ARGMAX_SRC`, `ARGMAX_TOKEN_CHAIN_SRC`, `ATTENTION_BF16_KV_BATCHED_SRC`, `ATTENTION_BF16_KV_SRC`, `ATTENTION_CAUSAL_BATCHED_SRC`, `ATTENTION_DECODE_BATCHED_HISTORY_SRC`, `ATTENTION_DFLASH_SLIDING_SRC`, `ATTENTION_DFLASH_SRC`, `ATTENTION_DFLASH_WMMA_GFX12_SRC`, `ATTENTION_DFLASH_WMMA_M128_N32_F16KV_V7B_SRC`, `ATTENTION_DFLASH_WMMA_M128_N32_F16KV_V7_SRC`, `ATTENTION_DFLASH_WMMA_M32_GFX12_SRC`, `ATTENTION_DFLASH_WMMA_M32_SRC`, `ATTENTION_DFLASH_WMMA_M64_N128_F16KV_SRC`, `ATTENTION_DFLASH_WMMA_M64_N128_F16KV_V2_SRC`, `ATTENTION_DFLASH_WMMA_M64_N128_F16KV_V3_CAUSAL_GFX12_SRC`, `ATTENTION_DFLASH_WMMA_M64_N128_F16KV_V3_CAUSAL_SRC`, `ATTENTION_DFLASH_WMMA_M64_N128_F16KV_V3_SRC`, `ATTENTION_DFLASH_WMMA_M64_N128_F16KV_V4_CAUSAL_SRC`, `ATTENTION_DFLASH_WMMA_M64_N128_F16KV_V4_SRC`, `ATTENTION_DFLASH_WMMA_M64_N32_F16KV_V5_GFX12_SRC`, `ATTENTION_DFLASH_WMMA_M64_N32_F16KV_V5_SRC`, `ATTENTION_DFLASH_WMMA_M64_N32_F16KV_V6_SRC`, `ATTENTION_DFLASH_WMMA_N128_F16KV_SRC`, `ATTENTION_DFLASH_WMMA_N64_F16KV_SRC`, `ATTENTION_DFLASH_WMMA_N64_SRC`, `ATTENTION_DFLASH_WMMA_SRC`, `ATTENTION_FLASH_ASYM2_TILE_BATCHED_SRC`, `ATTENTION_FLASH_ASYM2_TILE_SRC`, `ATTENTION_FLASH_ASYM3_TILE_BATCHED_SRC`, `ATTENTION_FLASH_ASYM3_TILE_HD512_BATCHED_SRC`, `ATTENTION_FLASH_ASYM3_TILE_HD512_SRC`, `ATTENTION_FLASH_ASYM3_TILE_SRC`, `ATTENTION_FLASH_ASYM4_TILE_BATCHED_SRC`, `ATTENTION_FLASH_ASYM4_TILE_SRC`, `ATTENTION_FLASH_ASYM4_WMMA_TILE_BATCHED_GFX12_SRC`, `ATTENTION_FLASH_ASYM4_WMMA_TILE_BATCHED_SRC`, `ATTENTION_FLASH_ASYM_REDUCE_BATCHED_SRC`, `ATTENTION_FLASH_BF16_TILE_BATCHED_SRC`, `ATTENTION_FLASH_BF16_TILE_SRC`, `ATTENTION_FLASH_F16_TILE_BATCHED_SRC`, `ATTENTION_FLASH_F16_TILE_SRC`, `ATTENTION_FLASH_FP8_E4M3_TILE_BATCHED_SRC`, `ATTENTION_FLASH_FP8_E4M3_TILE_GQA_GFX1201_SRC`, `ATTENTION_FLASH_FP8_E4M3_TILE_SRC`, `ATTENTION_FLASH_FWHT2_TILE_BATCHED_SRC`, `ATTENTION_FLASH_FWHT2_TILE_SRC`, `ATTENTION_FLASH_FWHT3_TILE_BATCHED_SRC`, `ATTENTION_FLASH_FWHT3_TILE_HD512_BATCHED_SRC`, `ATTENTION_FLASH_FWHT3_TILE_HD512_SRC`, `ATTENTION_FLASH_FWHT3_TILE_SRC`, `ATTENTION_FLASH_FWHT4_TILE_BATCHED_SRC`, `ATTENTION_FLASH_FWHT4_TILE_SRC`, `ATTENTION_FLASH_GQA_FUSED_SRC`, `ATTENTION_FLASH_GQA_SRC`, `ATTENTION_FLASH_LLOYD_REDUCE_BATCHED_SRC`, `ATTENTION_FLASH_LLOYD_REDUCE_SRC`, `ATTENTION_FLASH_Q8_0_REDUCE_GATED_MQ_ROTATE_AWQ_DEC_GFX1100_SRC`, `ATTENTION_FLASH_Q8_0_REDUCE_GATED_MQ_ROTATE_AWQ_GFX1100_SRC`, `ATTENTION_FLASH_Q8_0_REDUCE_GATED_MQ_ROTATE_AWQ_GFX1151_SRC`, `ATTENTION_FLASH_Q8_0_REDUCE_GATED_MQ_ROTATE_AWQ_GFX1201_SRC`, `ATTENTION_FLASH_Q8_0_REDUCE_GATED_MQ_ROTATE_GFX1100_SRC`, `ATTENTION_FLASH_Q8_0_REDUCE_GATED_MQ_ROTATE_GFX1151_SRC`, `ATTENTION_FLASH_Q8_0_REDUCE_GATED_MQ_ROTATE_GFX1201_SRC`, `ATTENTION_FLASH_Q8_0_REDUCE_SRC`, `ATTENTION_FLASH_Q8_0_TILE_BATCHED_SRC`, `ATTENTION_FLASH_Q8_0_TILE_DPP_GFX1151_SRC`, `ATTENTION_FLASH_Q8_0_TILE_GQA_GFX1100_SRC`, `ATTENTION_FLASH_Q8_0_TILE_GQA_GFX1151_SRC`, `ATTENTION_FLASH_Q8_0_TILE_ROWS_SRC`, `ATTENTION_FLASH_Q8_0_TILE_SRC`, `ATTENTION_FLASH_REDUCE_DSPLIT_GFX1151_SRC`, `ATTENTION_FLASH_REDUCE_DSPLIT_GFX1201_SRC`, `ATTENTION_FLASH_SRC`, `ATTENTION_FLUX_V2_WMMA_F16KV_SRC`, `ATTENTION_FLUX_VTK_WMMA_F16KV_F32_SRC`, `ATTENTION_FLUX_VT_WMMA_F16KV_F32_GFX12_SRC`, `ATTENTION_FLUX_VT_WMMA_F16KV_F32_SRC`, `ATTENTION_FP8_E4M3_FA2_GQA_F16_GFX1201_SRC`, `ATTENTION_FP8_E4M3_FA2_GQA_FP8_GFX1201_SRC`, `ATTENTION_FP8_E4M3_FA2_GQA_PACKET_GFX1201_SRC`, `ATTENTION_FP8_E4M3_FA2_GQA_QRESIDENT_GFX1201_SRC`, `ATTENTION_FP8_E4M3_FA2_GQA_QRESIDENT_V2_GFX1201_SRC`, `ATTENTION_FP8_E4M3_FA2_GQA_QRESIDENT_V2_Q8_A4EPI_GFX1201_SRC`, `ATTENTION_FP8_E4M3_FA2_GQA_QRESIDENT_V2_Q8_GFX1201_SRC`, `ATTENTION_FP8_E4M3_KV_BATCHED_SRC`, `ATTENTION_FP8_E4M3_KV_SRC`, `ATTENTION_GQA_WARP_DV_SRC`, `ATTENTION_GQA_WARP_SRC`, `ATTENTION_HFQ4_KV_SRC`, `ATTENTION_HFQ8_KV_SRC`, `ATTENTION_INT8C_F16_KV_SRC`, `ATTENTION_INT8C_KV_SRC`, `ATTENTION_INT8_KV_SRC`, `ATTENTION_PARTITION_MERGE_F32_SRC`, `ATTENTION_Q4KV_SRC`, `ATTENTION_Q8KV_SRC`, `ATTENTION_Q8_0_FA2_GQA_FP8_GFX1201_SRC`, `ATTENTION_Q8_0_FA2_GQA_FWHT3K_GFX11_SRC`, `ATTENTION_Q8_0_FA2_GQA_FWHT3K_GFX1201_SRC`, `ATTENTION_Q8_0_FA2_GQA_GFX1151_SRC`, `ATTENTION_Q8_0_FA2_GQA_GFX11_SRC`, `ATTENTION_Q8_0_FA2_GQA_GFX1201_SRC`, `ATTENTION_Q8_0_FA2_GQA_MERGE_FP8_GFX1201_SRC`, `ATTENTION_Q8_0_FA2_GQA_PARTIAL_FP8_GFX1201_SRC`, `ATTENTION_Q8_0_FLASH_PREFILL_SRC`, `ATTENTION_Q8_0_FLASH_PREFILL_WMMA_GFX12_SRC`, `ATTENTION_Q8_0_FLASH_PREFILL_WMMA_PARTIAL_GFX1201_SRC`, `ATTENTION_Q8_0_FLASH_PREFILL_WMMA_SRC`, `ATTENTION_Q8_0_FLASH_PREFILL_WMMA_SWA_GFX1100_SRC`, `ATTENTION_Q8_0_FLASH_PREFILL_WMMA_SWA_GFX1151_SRC`, `ATTENTION_Q8_0_FLASH_PREFILL_WMMA_SWA_GFX12_SRC`, `ATTENTION_Q8_0_FLASH_PREFILL_WMMA_SWA_SRC`, `ATTENTION_Q8_0_KV_BATCHED_SRC`, `ATTENTION_Q8_0_KV_SRC`, `ATTENTION_Q8_0_KV_SWA_SRC`, `ATTENTION_Q8_0_KV_TIMED_SRC`, `ATTENTION_SRC`, `ATTENTION_T5_BIAS_F32_SRC`, `ATTENTION_VERIFY_GQA_GFX1100_SRC`, `ATTENTION_VERIFY_GQA_GFX1201_SRC`, `ATTENTION_VERIFY_WMMA_GFX1151_SRC`, `BATCHED_CATEGORICAL_SAMPLE_SRC`, `BF16_ROUND_TRIP_SRC`, `BIAS_ADD_SRC`, `BLOCK_I4_128_QUANT_SRC`, `CALIB_REDUCE_SRC`, `CAST_F32_TO_F16_SRC`, `CHAIN_ACCEPT_SPEC_SRC`, `COMPRESSOR_ADD_APE_BATCHED_SRC`, `COMPRESSOR_COMPRESS_ALIGNED_BATCHED_SRC`, `COMPRESSOR_OVERLAP_CONCAT_SRC`, `COMPRESSOR_RING_WRITE_BATCHED_SRC`, `COMPRESSOR_SOFTMAX_POOL_BUF_SRC`, `COMPRESSOR_SOFTMAX_POOL_SRC`, `CONV1D_DECODE_SRC`, `CONV1D_GATED_DECODE_SRC`, `CONV1D_SILU_SPLIT_QKNORM_B128_SRC`, `CONV1D_SILU_SPLIT_QKNORM_B256_SCALAR_PREP_SRC`, `CONV1D_SILU_SPLIT_QKNORM_B256_SRC`, `CONV1D_SILU_SPLIT_QKNORM_B32_SRC`, `CONV1D_SILU_SPLIT_QKNORM_B512_SRC`, `CONV1D_SILU_SPLIT_QKNORM_B64_SRC`, `CONV1D_SILU_SPLIT_SRC`, `CONV1D_SILU_SPLIT_TREE_SRC`, `CONV1D_SILU_SRC`, `COPY_F32_BUFFER_SRC`, `COPY_ROWS_STRIDED_F32_SRC`, `CROSS_ENTROPY_LOSS_SRC`, `DDTREE_BUILD_ATTN_MASK_SRC`, `DDTREE_GUMBEL_TOPK_BATCHED_SRC`, `DDTREE_SWOR_WALK_SRC`, `DEEPSEEK4_COMPRESSOR_CACHE_F16_SRC`, `DEEPSEEK4_COMPRESSOR_COMMIT_STAGED_F16_SRC`, `DEINTERLEAVE_BATCHED_SRC`, `DEINTERLEAVE_Q_RMSNORM_BATCHED_SRC`, `DEINTERLEAVE_SRC`, `DEQUANTIZE_MFP4G32_E8_SOA_TO_F16_GFX942_SRC`, `DEQUANTIZE_MFP4G32_E8_TO_F16_SRC`, `DEQUANTIZE_MFP4G32_LLOYD_TO_F16_SRC`, `DEQUANTIZE_MFP4G32_P_TO_F16_SRC`, `DEQUANT_BQ1G128_TO_F16_SRC`, `DEQUANT_FP4_E2M1_UE8M0_G32_TO_BF16_GFX942_SRC`, `DEQUANT_FP8_E4M3_UE8M0_BLK128_TO_BF16_GFX942_SRC`, `DEQUANT_HFQ4G256_TO_F16_SRC`, `DEQUANT_TQ2G128_TO_F16_SRC`, `DSPARK_SAMPLE_ACCEPT_LAZY_SRC`, `DSPARK_STAGE_KV_SRC`, `DYNAMIC_CAUSAL_CONV_F32_SRC`, `DYNAMIC_CONV_F32_SRC`, `EMBEDDING_BF16_BATCHED_SRC`, `EMBEDDING_F16_BATCHED_SRC`, `EMBEDDING_HFQ4G128_BATCHED_SRC`, `EMBEDDING_HFQ4G128_SRC`, `EMBEDDING_HFQ4G256_BATCHED_SRC`, `EMBEDDING_HFQ4G256_SRC`, `EMBEDDING_MQ4V2_BATCHED_SRC`, `EMBEDDING_Q4K_SRC`, `EMBEDDING_Q8_BATCHED_SRC`, `EMBEDDING_Q8_SRC`, `EMBEDDING_SCATTER_EXT_SRC`, `FN_GDN_DENSE_SRC`, `FUSED_GATE_UP_HFQ4G256_DOT_PREFETCH_GFX1100_SRC`, `FUSED_GATE_UP_HFQ4G256_DOT_REFORM_GFX1100_SRC`, `FUSED_GATE_UP_HFQ4G256_K1024_GFX1201_SRC`, `FUSED_GATE_UP_HFQ4G256_LANE0_HEADERS_GFX1100_SRC`, `FUSED_GATE_UP_HFQ4G256_PAIR2_GFX1100_SRC`, `FUSED_GATE_UP_HFQ4G256_PAIR_GFX1100_SRC`, `FUSED_GATE_UP_HFQ4G256_QUAD_PREFETCH_GFX1100_SRC`, `FUSED_GATE_UP_HFQ4G256_SETPRIO_GFX1100_SRC`, `FUSED_GATE_UP_HFQ4G256_SRC`, `FUSED_GATE_UP_HFQ4G256_STAGE_X32_GFX1100_SRC`, `FUSED_GATE_UP_HFQ4G256_V2_GFX942_SRC`, `FUSED_GATE_UP_HFQ4G256_WAVE64_DP4A_SRC`, `FUSED_GATE_UP_HFQ4G256_WAVE64_SRC`, `FUSED_GATE_UP_HFQ6G256_WAVE64_DP4A_SRC`, `FUSED_GATE_UP_MFP4G32_E8_GFX1151_SRC`, `FUSED_GATE_UP_MQ2G256V2_SRC`, `FUSED_GATE_UP_MQ3G256V2_SRC`, `FUSED_GATE_UP_MQ3G256_LLOYD_GFX1100_SRC`, `FUSED_GATE_UP_MQ3G256_LLOYD_SRC`, `FUSED_GATE_UP_MQ4CG256_SRC`, `FUSED_GATE_UP_MQ4G256V2_LUT_SRC`, `FUSED_GATE_UP_MQ4G256V2_SRC`, `FUSED_GATE_UP_MQ4G256_LLOYD_GFX1100_SRC`, `FUSED_GATE_UP_MQ4G256_LLOYD_SRC`, `FUSED_GATE_UP_MQ5G256V2_SRC`, `FUSED_GATE_UP_MQ6G256V2_SRC`, `FUSED_GATE_UP_Q4K_SRC`, `FUSED_GATE_UP_Q8_0_SRC`, `FUSED_GEMMA4_QK_NORM_ROPE_SRC`, `FUSED_GLIMMER_GATE_UP_HFQ4G256_K6656_GFX1100_SRC`, `FUSED_GLIMMER_QKVG_HFQ4G256_K6656_GFX1100_SRC`, `FUSED_QKVZA_HFQ4G256_2WAVE_GFX1100_SRC`, `FUSED_QKVZA_HFQ4G256_HOIST_X32_GFX1100_SRC`, `FUSED_QKVZA_HFQ4G256_K2048_ALL_BUFFER_DLC_GFX1151_SRC`, `FUSED_QKVZA_HFQ4G256_K2048_ALL_BUFFER_GFX1151_SRC`, `FUSED_QKVZA_HFQ4G256_K2048_ALL_BUFFER_GLC_GFX1151_SRC`, `FUSED_QKVZA_HFQ4G256_K2048_ALL_BUFFER_SLC_GFX1151_SRC`, `FUSED_QKVZA_HFQ4G256_K2048_BUFFER_GFX1151_SRC`, `FUSED_QKVZA_HFQ4G256_K2048_CPOL_SLC_GFX1100_SRC`, `FUSED_QKVZA_HFQ4G256_K2048_GFX1100_SRC`, `FUSED_QKVZA_HFQ4G256_K2048_HOIST_X32_GFX1151_SRC`, `FUSED_QKVZA_HFQ4G256_K2048_HYBRID_BUFFER_GFX1151_SRC`, `FUSED_QKVZA_HFQ4G256_K2048_PAIR_BUFFER_GFX1151_SRC`, `FUSED_QKVZA_HFQ4G256_K2048_R2_BUFFER_GFX1151_SRC`, `FUSED_QKVZA_HFQ4G256_K2048_R2_GFX1100_SRC`, `FUSED_QKVZA_HFQ4G256_K2048_R4_STREAM_GFX1151_SRC`, `FUSED_QKVZA_HFQ4G256_K2048_SCALAR_PREP_GFX1100_SRC`, `FUSED_QKVZA_HFQ4G256_K2048_X_BUFFER_GFX1151_SRC`, `FUSED_QKVZA_HFQ4G256_LDSX8_BUFFER_GFX1151_SRC`, `FUSED_QKVZA_HFQ4G256_LDSX8_GFX1100_SRC`, `FUSED_QKVZA_HFQ4G256_REDUCE_CHAIN_GFX1100_SRC`, `FUSED_QKVZA_HFQ4G256_SRC`, `FUSED_QKVZA_HFQ4G256_V2_GFX942_SRC`, `FUSED_QKVZA_HFQ4G256_WAVE64_DP4A_SRC`, `FUSED_QKVZA_HFQ4G256_WAVE64_SHARE_X_GFX1151_SRC`, `FUSED_QKVZA_HFQ4G256_WAVE64_SRC`, `FUSED_QKVZA_HFQ4G256_WAVEPACK4_GFX1100_SRC`, `FUSED_QKVZA_HFQ6G256_WAVE64_DP4A_SRC`, `FUSED_QKVZA_MFP4G32_E8_GFX1151_SRC`, `FUSED_QKVZA_MQ2G256V2_SRC`, `FUSED_QKVZA_MQ3G256V2_SRC`, `FUSED_QKVZA_MQ3G256_LLOYD_GFX1100_SRC`, `FUSED_QKVZA_MQ3G256_LLOYD_SRC`, `FUSED_QKVZA_MQ4CG256_SRC`, `FUSED_QKVZA_MQ4G256V2_K2048_HOIST_X32_GFX1100_SRC`, `FUSED_QKVZA_MQ4G256V2_LUT_SRC`, `FUSED_QKVZA_MQ4G256V2_SRC`, `FUSED_QKVZA_MQ4G256_LLOYD_GFX1100_SRC`, `FUSED_QKVZA_MQ4G256_LLOYD_SRC`, `FUSED_QKVZA_MQ5G256V2_SRC`, `FUSED_QKVZA_MQ6G256V2_SRC`, `FUSED_QKVZA_Q8_0_SRC`, `FUSED_QKV_HFQ4G256_K2048_ALL_BUFFER_GFX1151_SRC`, `FUSED_QKV_HFQ4G256_K2048_ALL_BUFFER_SLC_GFX1151_SRC`, `FUSED_QKV_HFQ4G256_K2048_X_BUFFER_GFX1151_SRC`, `FUSED_QKV_HFQ4G256_QWEN2_BIAS_SRC`, `FUSED_QKV_HFQ4G256_SRC`, `FUSED_QKV_HFQ4G256_V2_GFX942_QWEN2_BIAS_SRC`, `FUSED_QKV_HFQ4G256_V2_GFX942_SRC`, `FUSED_QKV_HFQ4G256_WAVE64_DP4A_SRC`, `FUSED_QKV_HFQ4G256_WAVE64_QWEN2_BIAS_SRC`, `FUSED_QKV_HFQ4G256_WAVE64_SRC`, `FUSED_QKV_HFQ6G256_SRC`, `FUSED_QKV_HFQ6G256_WAVE64_DP4A_SRC`, `FUSED_QKV_MQ2G256V2_QWEN2_BIAS_SRC`, `FUSED_QKV_MQ2G256V2_SRC`, `FUSED_QKV_MQ3G256V2_QWEN2_BIAS_SRC`, `FUSED_QKV_MQ3G256V2_SRC`, `FUSED_QKV_MQ3G256_LLOYD_GFX1100_QWEN2_BIAS_SRC`, `FUSED_QKV_MQ3G256_LLOYD_GFX1100_SRC`, `FUSED_QKV_MQ3G256_LLOYD_QWEN2_BIAS_SRC`, `FUSED_QKV_MQ3G256_LLOYD_SRC`, `FUSED_QKV_MQ4CG256_QWEN2_BIAS_SRC`, `FUSED_QKV_MQ4CG256_SRC`, `FUSED_QKV_MQ4G256V2_K2048_X_BUFFER_GFX1100_SRC`, `FUSED_QKV_MQ4G256V2_LUT_SRC`, `FUSED_QKV_MQ4G256V2_QWEN2_BIAS_SRC`, `FUSED_QKV_MQ4G256V2_SRC`, `FUSED_QKV_MQ4G256_LLOYD_GFX1100_QWEN2_BIAS_SRC`, `FUSED_QKV_MQ4G256_LLOYD_GFX1100_SRC`, `FUSED_QKV_MQ4G256_LLOYD_QWEN2_BIAS_SRC`, `FUSED_QKV_MQ4G256_LLOYD_SRC`, `FUSED_QKV_MQ5G256V2_QWEN2_BIAS_SRC`, `FUSED_QKV_MQ5G256V2_SRC`, `FUSED_QKV_MQ6G256V2_QWEN2_BIAS_SRC`, `FUSED_QKV_MQ6G256V2_SRC`, `FUSED_QKV_Q4K_QWEN2_BIAS_SRC`, `FUSED_QKV_Q4K_SRC`, `FUSED_QKV_Q8_0_QWEN2_BIAS_SRC`, `FUSED_QKV_Q8_0_SRC`, `FUSED_QK_L2_NORM_SCALE_INTERLEAVE_F32_BATCHED_SRC`, `FUSED_QK_L2_NORM_SCALE_SRC`, `FUSED_RMSNORM_MQ_ROTATE_AWQ_FP8_GFX12_SRC`, `FUSED_RMSNORM_MQ_ROTATE_AWQ_FP8_INREG_GFX12_SRC`, `FUSED_RMSNORM_MQ_ROTATE_AWQ_FP8_INREG_SHORT_GFX12_SRC`, `FUSED_RMSNORM_MQ_ROTATE_AWQ_FP8_INREG_SHORT_SEQ_GFX12_SRC`, `FUSED_RMSNORM_MQ_ROTATE_AWQ_G12DEC_SRC`, `FUSED_RMSNORM_MQ_ROTATE_AWQ_I4_B8_GFX1151_SRC`, `FUSED_RMSNORM_MQ_ROTATE_AWQ_I4_B8_SRC`, `FUSED_RMSNORM_MQ_ROTATE_AWQ_I4_FOLD_B8_SRC`, `FUSED_RMSNORM_MQ_ROTATE_AWQ_I4_FOLD_SRC`, `FUSED_RMSNORM_MQ_ROTATE_AWQ_I4_GFX12_SRC`, `FUSED_RMSNORM_MQ_ROTATE_AWQ_I4_GFX12_V2_SLAB_FDIV_SRC`, `FUSED_RMSNORM_MQ_ROTATE_AWQ_I4_GFX12_V2_SLAB_SRC`, `FUSED_RMSNORM_MQ_ROTATE_AWQ_I4_GFX12_V2_SRC`, `FUSED_RMSNORM_MQ_ROTATE_AWQ_I4_SRC`, `FUSED_RMSNORM_MQ_ROTATE_AWQ_SRC`, `FUSED_RMSNORM_MQ_ROTATE_FP8_GFX12_SRC`, `FUSED_RMSNORM_MQ_ROTATE_FP8_INREG_GFX12_SRC`, `FUSED_RMSNORM_MQ_ROTATE_FP8_INREG_SHORT_GFX12_SRC`, `FUSED_RMSNORM_MQ_ROTATE_FP8_INREG_SHORT_SEQ_GFX12_SRC`, `FUSED_RMSNORM_MQ_ROTATE_I4_B8_GFX1151_SRC`, `FUSED_RMSNORM_MQ_ROTATE_I4_B8_SRC`, `FUSED_RMSNORM_MQ_ROTATE_I4_FOLD_B8_SRC`, `FUSED_RMSNORM_MQ_ROTATE_I4_FOLD_SRC`, `FUSED_RMSNORM_MQ_ROTATE_I4_GFX12_SRC`, `FUSED_RMSNORM_MQ_ROTATE_I4_GFX12_V2_SRC`, `FUSED_RMSNORM_MQ_ROTATE_I4_SRC`, `FUSED_RMSNORM_MQ_ROTATE_PLAIN_SRC`, `FUSED_RMSNORM_MQ_ROTATE_SRC`, `FUSED_RMSNORM_MQ_ROTATE_VECSUM_GFX1100_SRC`, `FUSED_RMSNORM_MQ_ROTATE_VECSUM_SIGN_CONST_GFX1100_SRC`, `FUSED_RMSNORM_MQ_ROTATE_VECSUM_SIGN_LDS_GFX1100_SRC`, `FUSED_RMSNORM_MQ_ROTATE_WAVEGRID_GFX1100_SRC`, `FUSED_SIGMOID_ALPHA_GATE_SRC`, `FUSED_SILU_MUL_GIVENS_ROTATE_SRC`, `FUSED_SILU_MUL_MQ_ROTATE_AWQ_FP8_GFX12_SRC`, `FUSED_SILU_MUL_MQ_ROTATE_AWQ_FP8_INREG_GFX12_SRC`, `FUSED_SILU_MUL_MQ_ROTATE_AWQ_HIN_BF16_FP8_GFX12_SRC`, `FUSED_SILU_MUL_MQ_ROTATE_AWQ_HIN_FP8_GFX12_SRC`, `FUSED_SILU_MUL_MQ_ROTATE_AWQ_I4_GFX12_SRC`, `FUSED_SILU_MUL_MQ_ROTATE_AWQ_I4_HIN_BF16_GFX12_SRC`, `FUSED_SILU_MUL_MQ_ROTATE_AWQ_I4_HIN_GFX12_SLAB_SRC`, `FUSED_SILU_MUL_MQ_ROTATE_AWQ_I4_HIN_GFX12_SLAB_TOKFAST_SRC`, `FUSED_SILU_MUL_MQ_ROTATE_AWQ_I4_HIN_GFX12_SRC`, `FUSED_SILU_MUL_MQ_ROTATE_AWQ_I4_HIN_SRC`, `FUSED_SILU_MUL_MQ_ROTATE_AWQ_I4_SRC`, `FUSED_SILU_MUL_MQ_ROTATE_AWQ_INDEXED_SRC`, `FUSED_SILU_MUL_MQ_ROTATE_AWQ_SRC`, `FUSED_SILU_MUL_MQ_ROTATE_FP8_GFX12_SRC`, `FUSED_SILU_MUL_MQ_ROTATE_FP8_INREG_GFX12_SRC`, `FUSED_SILU_MUL_MQ_ROTATE_HIN_BF16_FP8_GFX12_SRC`, `FUSED_SILU_MUL_MQ_ROTATE_HIN_FP8_GFX12_SRC`, `FUSED_SILU_MUL_MQ_ROTATE_I4_GFX12_SRC`, `FUSED_SILU_MUL_MQ_ROTATE_I4_SRC`, `FUSED_SILU_MUL_MQ_ROTATE_SRC`, `GATED_ADD_F32_SRC`, `GATED_DELTA_NET_F32_BATCH_SEQ_SRC`, `GATED_DELTA_NET_F32_CHUNKED_SRC`, `GATED_DELTA_NET_F32_TREE_SRC`, `GATED_DELTA_NET_Q4_SRC`, `GATED_DELTA_NET_Q8_COMPACT2_B12_SRC`, `GATED_DELTA_NET_Q8_COMPACT2_B16_SRC`, `GATED_DELTA_NET_Q8_COMPACT2_B2_SRC`, `GATED_DELTA_NET_Q8_COMPACT2_B4_SRC`, `GATED_DELTA_NET_Q8_COMPACT2_B8_SRC`, `GATED_DELTA_NET_Q8_COMPACT2_DPP_GFX1151_SRC`, `GATED_DELTA_NET_Q8_COMPACT2_R4X2_GFX1151_SRC`, `GATED_DELTA_NET_Q8_COMPACT2_R8_GFX1151_SRC`, `GATED_DELTA_NET_Q8_COMPACT3_B2_SRC`, `GATED_DELTA_NET_Q8_FAST_SRC`, `GATED_DELTA_NET_Q8_SRC`, `GATED_DELTA_NET_Q8_TREE_SRC`, `GATED_DELTA_NET_SRC`, `GATED_NORM_MQ_ROTATE_AWQ_FP8_GFX12_SRC`, `GATED_NORM_MQ_ROTATE_AWQ_FP8_GFX12_XBF16_SRC`, `GATED_NORM_MQ_ROTATE_AWQ_FP8_INREG_GFX12_SRC`, `GATED_NORM_MQ_ROTATE_AWQ_FP8_INREG_GFX12_XBF16_SRC`, `GATED_NORM_MQ_ROTATE_AWQ_GFX1100_SRC`, `GATED_NORM_MQ_ROTATE_AWQ_GFX1151_SRC`, `GATED_NORM_MQ_ROTATE_AWQ_GFX1201_SRC`, `GATED_NORM_MQ_ROTATE_AWQ_I4_GFX1100_V2_SRC`, `GATED_NORM_MQ_ROTATE_AWQ_I4_GFX1151_V2_SRC`, `GATED_NORM_MQ_ROTATE_AWQ_I4_GFX11_SRC`, `GATED_NORM_MQ_ROTATE_AWQ_I4_GFX12_SRC`, `GATED_NORM_MQ_ROTATE_AWQ_I4_GFX12_V2_SLAB_SRC`, `GATED_NORM_MQ_ROTATE_AWQ_I4_GFX12_V2_SRC`, `GATED_NORM_MQ_ROTATE_AWQ_I4_GFX12_V2_XBF16_SLAB_SRC`, `GATED_NORM_MQ_ROTATE_AWQ_I4_GFX12_V2_XBF16_SRC`, `GATED_NORM_MQ_ROTATE_AWQ_I4_GFX12_XBF16_SRC`, `GATED_NORM_MQ_ROTATE_FP8_GFX12_SRC`, `GATED_NORM_MQ_ROTATE_FP8_GFX12_XBF16_SRC`, `GATED_NORM_MQ_ROTATE_FP8_INREG_GFX12_SRC`, `GATED_NORM_MQ_ROTATE_FP8_INREG_GFX12_XBF16_SRC`, `GATED_NORM_MQ_ROTATE_GFX1100_SRC`, `GATED_NORM_MQ_ROTATE_GFX1151_SRC`, `GATED_NORM_MQ_ROTATE_GFX1201_SRC`, `GATED_NORM_MQ_ROTATE_I4_GFX1100_V2_SRC`, `GATED_NORM_MQ_ROTATE_I4_GFX1151_V2_SRC`, `GATED_NORM_MQ_ROTATE_I4_GFX11_SRC`, `GATED_NORM_MQ_ROTATE_I4_GFX12_SRC`, `GATED_NORM_MQ_ROTATE_I4_GFX12_V2_SRC`, `GATED_NORM_MQ_ROTATE_I4_GFX12_V2_XBF16_SRC`, `GATED_NORM_MQ_ROTATE_I4_GFX12_XBF16_SRC`, `GATED_NORM_SRC`, `GDN_CHUNK_KKT_SOLVE_BATCHED_SRC`, `GDN_CHUNK_KKT_SOLVE_GFX1100_SRC`, `GDN_CHUNK_KKT_SOLVE_SRC`, `GDN_CHUNK_PREP_FIXUP_SRC`, `GDN_CHUNK_PREP_GFX11_SRC`, `GDN_CHUNK_PREP_SRC`, `GDN_CHUNK_SCAN_BF16_MSEG_SRC`, `GDN_CHUNK_SCAN_BF16_SRC`, `GDN_CHUNK_SCAN_EMU_BF16_SRC`, `GDN_CHUNK_SCAN_EMU_F16_SRC`, `GDN_CHUNK_SCAN_GFX1151_SRC`, `GDN_CHUNK_SCAN_SRC`, `GDN_PRE_BATCHED_GFX1201_SRC`, `GELU_NEW_MUL_F32_SRC`, `GELU_TANH_SRC`, `GEMMA4_PLE_ACTIVATION_SRC`, `GEMM_BF16_MFMA_GFX942_SRC`, `GEMM_BF16_XF32_BATCHED_SRC`, `GEMM_BF16_XF32_MULTIROW_R16_GFX1151_SRC`, `GEMM_BF16_XF32_MULTIROW_SRC`, `GEMM_BQ1G128_PREFILL_SRC`, `GEMM_BQ1G128_WMMA_SRC`, `GEMM_F16_BIAS_SRC`, `GEMM_F16_SRC`, `GEMM_F16_TILED_SRC`, `GEMM_F16_WMMA_MB4_SRC`, `GEMM_F16_WMMA_MB8_GFX12_SRC`, `GEMM_F16_WMMA_MB8_SRC`, `GEMM_F16_WMMA_SRC`, `GEMM_F16_X_F16_MFMA_GFX942_SRC`, `GEMM_F16_X_F16_WMMA_GFX12_SRC`, `GEMM_F16_X_F16_WMMA_LDS256_SRC`, `GEMM_F16_X_F16_WMMA_LDS_SPLITK_SRC`, `GEMM_F16_X_F16_WMMA_LDS_SRC`, `GEMM_F16_X_F16_WMMA_SRC`, `GEMM_F32_REGISTER_TILED_SRC`, `GEMM_F32_SRC`, `GEMM_GATE_UP_HFP4G32_WMMA_GFX12_SRC`, `GEMM_GATE_UP_HFP4G32_WMMA_SRC`, `GEMM_GATE_UP_HFQ3G256_DOT2_SRC`, `GEMM_GATE_UP_HFQ3G256_DP4A_SRC`, `GEMM_GATE_UP_HFQ3G256_FP16_SRC`, `GEMM_GATE_UP_HFQ3G256_MMQ_BODY_CUH`, `GEMM_GATE_UP_HFQ3G256_MMQ_X16_SRC`, `GEMM_GATE_UP_HFQ3G256_MMQ_X32_SRC`, `GEMM_GATE_UP_HFQ3G256_MMQ_X32_Y64_SRC`, `GEMM_GATE_UP_HFQ3G256_MMQ_X32_Y96_SRC`, `GEMM_GATE_UP_HFQ3G256_MMQ_X8_SRC`, `GEMM_GATE_UP_HFQ3G256_SRC`, `GEMM_GATE_UP_HFQ3G256_WMMA_GFX12_SRC`, `GEMM_GATE_UP_HFQ3G256_WMMA_MB4_SRC`, `GEMM_GATE_UP_HFQ3G256_WMMA_SRC`, `GEMM_GATE_UP_HFQ4G256_DOT2_SRC`, `GEMM_GATE_UP_HFQ4G256_FP16_SRC`, `GEMM_GATE_UP_HFQ4G256_FP16_WAVE64_SRC`, `GEMM_GATE_UP_HFQ4G256_MMQ_BODY_CUH`, `GEMM_GATE_UP_HFQ4G256_MMQ_GFX906_BODY_CUH`, `GEMM_GATE_UP_HFQ4G256_MMQ_GFX906_X16_SRC`, `GEMM_GATE_UP_HFQ4G256_MMQ_GFX906_X16_Y64_SRC`, `GEMM_GATE_UP_HFQ4G256_MMQ_GFX906_X32_SRC`, `GEMM_GATE_UP_HFQ4G256_MMQ_GFX906_X32_Y64_SRC`, `GEMM_GATE_UP_HFQ4G256_MMQ_GFX906_X64_SRC`, `GEMM_GATE_UP_HFQ4G256_MMQ_GFX906_X8_SRC`, `GEMM_GATE_UP_HFQ4G256_MMQ_X16_SRC`, `GEMM_GATE_UP_HFQ4G256_MMQ_X32_SRC`, `GEMM_GATE_UP_HFQ4G256_SRC`, `GEMM_GATE_UP_HFQ4G256_WAVE64_DP4A_SRC`, `GEMM_GATE_UP_HFQ4G256_WMMA_2TILE_SRC`, `GEMM_GATE_UP_HFQ4G256_WMMA_GFX12_BT_SRC`, `GEMM_GATE_UP_HFQ4G256_WMMA_GFX12_SRC`, `GEMM_GATE_UP_HFQ4G256_WMMA_K4_SRC`, `GEMM_GATE_UP_HFQ4G256_WMMA_LDSCOOP_NOSYNC_SRC`, `GEMM_GATE_UP_HFQ4G256_WMMA_LDSCOOP_SRC`, `GEMM_GATE_UP_HFQ4G256_WMMA_LDSX_SRC`, `GEMM_GATE_UP_HFQ4G256_WMMA_SRC`, `GEMM_GATE_UP_HFQ6G256_DOT2_SRC`, `GEMM_GATE_UP_HFQ6G256_FP16_SRC`, `GEMM_GATE_UP_HFQ6G256_SRC`, `GEMM_GATE_UP_HFQ6G256_WAVE64_DP4A_SRC`, `GEMM_GATE_UP_HFQ6G256_WMMA_GFX12_SRC`, `GEMM_GATE_UP_HFQ6G256_WMMA_SRC`, `GEMM_GATE_UP_MQ2G256V2_WMMA_GFX12_BT_SRC`, `GEMM_GATE_UP_MQ2G256V2_WMMA_GFX12_SRC`, `GEMM_GATE_UP_MQ2G256V2_WMMA_SRC`, `GEMM_GATE_UP_MQ3G256V2_WMMA_GFX12_BT_SRC`, `GEMM_GATE_UP_MQ3G256V2_WMMA_GFX12_SRC`, `GEMM_GATE_UP_MQ3G256V2_WMMA_SRC`, `GEMM_GATE_UP_MQ3G256_LLOYD_WMMA_GFX12_SRC`, `GEMM_GATE_UP_MQ3G256_LLOYD_WMMA_MB4_NOSYNC_SRC`, `GEMM_GATE_UP_MQ3G256_LLOYD_WMMA_MB4_SRC`, `GEMM_GATE_UP_MQ3G256_LLOYD_WMMA_NOSYNC_SRC`, `GEMM_GATE_UP_MQ3G256_LLOYD_WMMA_SRC`, `GEMM_GATE_UP_MQ4CG256_WMMA_GFX12_BT_SRC`, `GEMM_GATE_UP_MQ4CG256_WMMA_GFX12_SRC`, `GEMM_GATE_UP_MQ4G256V2_WMMA_FP8_GFX12_BT4_LUT_SRC`, `GEMM_GATE_UP_MQ4G256V2_WMMA_FP8_GFX12_BT4_SRC`, `GEMM_GATE_UP_MQ4G256V2_WMMA_FP8_GFX12_BT8_LUT_SRC`, `GEMM_GATE_UP_MQ4G256V2_WMMA_FP8_GFX12_BT8_SRC`, `GEMM_GATE_UP_MQ4G256V2_WMMA_FP8_GFX12_BT_LUT_SRC`, `GEMM_GATE_UP_MQ4G256V2_WMMA_FP8_GFX12_BT_SRC`, `GEMM_GATE_UP_MQ4G256V2_WMMA_FP8_GFX12_S2BT8_LUT_SRC`, `GEMM_GATE_UP_MQ4G256V2_WMMA_FP8_GFX12_S2BT8_SRC`, `GEMM_GATE_UP_MQ4G256V2_WMMA_FP8_GFX12_V2_B128X128_SILU_H_BF16_SRC`, `GEMM_GATE_UP_MQ4G256V2_WMMA_FP8_GFX12_V2_B128X128_SILU_H_SRC`, `GEMM_GATE_UP_MQ4G256V2_WMMA_FP8_GFX12_V2_B128X128_SRC`, `GEMM_GATE_UP_MQ4G256V2_WMMA_FP8_GFX12_V2_B128X64W4_SRC`, `GEMM_GATE_UP_MQ4G256V2_WMMA_FP8_GFX12_V2_B64X256_SRC`, `GEMM_GATE_UP_MQ4G256V2_WMMA_FP8_GFX12_V2_SRC`, `GEMM_GATE_UP_MQ4G256V2_WMMA_GFX1100_LDSSTAGE_SRC`, `GEMM_GATE_UP_MQ4G256V2_WMMA_GFX1100_MW_LDS_SRC`, `GEMM_GATE_UP_MQ4G256V2_WMMA_GFX11_BT_SRC`, `GEMM_GATE_UP_MQ4G256V2_WMMA_GFX12_BT_SRC`, `GEMM_GATE_UP_MQ4G256V2_WMMA_GFX12_SRC`, `GEMM_GATE_UP_MQ4G256V2_WMMA_GFX12_VT_SRC`, `GEMM_GATE_UP_MQ4G256V2_WMMA_SRC`, `GEMM_GATE_UP_MQ4G256_LLOYD_WMMA_GFX1151_SRC`, `GEMM_GATE_UP_MQ4G256_LLOYD_WMMA_GFX12_SRC`, `GEMM_GATE_UP_MQ4G256_LLOYD_WMMA_MB4_GFX1151_SRC`, `GEMM_GATE_UP_MQ4G256_LLOYD_WMMA_MB4_NOSYNC_GFX1151_SRC`, `GEMM_GATE_UP_MQ4G256_LLOYD_WMMA_MB4_NOSYNC_SRC`, `GEMM_GATE_UP_MQ4G256_LLOYD_WMMA_MB4_SRC`, `GEMM_GATE_UP_MQ4G256_LLOYD_WMMA_NOSYNC_GFX1151_SRC`, `GEMM_GATE_UP_MQ4G256_LLOYD_WMMA_NOSYNC_SRC`, `GEMM_GATE_UP_MQ4G256_LLOYD_WMMA_SRC`, `GEMM_GATE_UP_MQ5G256V2_WMMA_GFX12_BT_SRC`, `GEMM_GATE_UP_MQ5G256V2_WMMA_GFX12_SRC`, `GEMM_GATE_UP_MQ5G256V2_WMMA_SRC`, `GEMM_GATE_UP_MQ6G256V2_WMMA_GFX12_BT_SRC`, `GEMM_GATE_UP_MQ6G256V2_WMMA_GFX12_SRC`, `GEMM_GATE_UP_MQ6G256V2_WMMA_SRC`, `GEMM_GATE_UP_Q8_0_WMMA_GFX12_SRC`, `GEMM_GATE_UP_Q8_0_WMMA_SRC`, `GEMM_HFP4G32_MOE_GROUPED_WMMA_GFX12_SRC`, `GEMM_HFP4G32_RESIDUAL_WMMA_GFX12_SRC`, `GEMM_HFP4G32_RESIDUAL_WMMA_SRC`, `GEMM_HFQ3G256_MOE_GROUPED_WMMA_GFX12_SRC`, `GEMM_HFQ3G256_RESIDUAL_DOT2_SRC`, `GEMM_HFQ3G256_RESIDUAL_FP16_SRC`, `GEMM_HFQ3G256_RESIDUAL_MMQ_BODY_CUH`, `GEMM_HFQ3G256_RESIDUAL_MMQ_X16_SRC`, `GEMM_HFQ3G256_RESIDUAL_MMQ_X32_SRC`, `GEMM_HFQ3G256_RESIDUAL_MMQ_X32_Y32_SRC`, `GEMM_HFQ3G256_RESIDUAL_MMQ_X32_Y64_SRC`, `GEMM_HFQ3G256_RESIDUAL_MMQ_X8_SRC`, `GEMM_HFQ3G256_RESIDUAL_SRC`, `GEMM_HFQ3G256_RESIDUAL_WMMA_GFX12_SRC`, `GEMM_HFQ3G256_RESIDUAL_WMMA_MB4_SRC`, `GEMM_HFQ3G256_RESIDUAL_WMMA_SRC`, `GEMM_HFQ4G128_MMQ_GFX1151_SRC`, `GEMM_HFQ4G128_SRC`, `GEMM_HFQ4G256_LMHEAD_WMMA_GFX12_SRC`, `GEMM_HFQ4G256_MOE_GROUPED_MMQ_GFX1151_SRC`, `GEMM_HFQ4G256_MOE_GROUPED_MMQ_GFX11_DGPU_SRC`, `GEMM_HFQ4G256_MOE_GROUPED_MMQ_GFX12_SRC`, `GEMM_HFQ4G256_MOE_GROUPED_MMQ_K4_GFX1151_SRC`, `GEMM_HFQ4G256_MOE_GROUPED_MMQ_K4_GFX11_DGPU_SRC`, `GEMM_HFQ4G256_MOE_GROUPED_MMQ_K4_GFX12_SRC`, `GEMM_HFQ4G256_MOE_GROUPED_MMQ_K8_GFX1151_SRC`, `GEMM_HFQ4G256_MOE_GROUPED_WMMA_GFX12_SRC`, `GEMM_HFQ4G256_MOE_GROUPED_WMMA_K2_SRC`, `GEMM_HFQ4G256_MOE_GROUPED_WMMA_M2_GFX12_SRC`, `GEMM_HFQ4G256_RESIDUAL_FP16_SRC`, `GEMM_HFQ4G256_RESIDUAL_FP16_WAVE64_SRC`, `GEMM_HFQ4G256_RESIDUAL_MFMA_GFX942_SRC`, `GEMM_HFQ4G256_RESIDUAL_MFMA_V2_GFX942_SRC`, `GEMM_HFQ4G256_RESIDUAL_MFMA_V3_GFX942_SRC`, `GEMM_HFQ4G256_RESIDUAL_MFMA_V4_GFX942_SRC`, `GEMM_HFQ4G256_RESIDUAL_MMQ_BODY_CUH`, `GEMM_HFQ4G256_RESIDUAL_MMQ_GFX12_SRC`, `GEMM_HFQ4G256_RESIDUAL_MMQ_GFX906_BODY_CUH`, `GEMM_HFQ4G256_RESIDUAL_MMQ_GFX906_X16_SRC`, `GEMM_HFQ4G256_RESIDUAL_MMQ_GFX906_X24_SRC`, `GEMM_HFQ4G256_RESIDUAL_MMQ_GFX906_X32_SRC`, `GEMM_HFQ4G256_RESIDUAL_MMQ_GFX906_X40_SRC`, `GEMM_HFQ4G256_RESIDUAL_MMQ_GFX906_X48_SRC`, `GEMM_HFQ4G256_RESIDUAL_MMQ_GFX906_X56_SRC`, `GEMM_HFQ4G256_RESIDUAL_MMQ_GFX906_X64_SRC`, `GEMM_HFQ4G256_RESIDUAL_MMQ_GFX906_X8_SRC`, `GEMM_HFQ4G256_RESIDUAL_MMQ_RDNA2_SRC`, `GEMM_HFQ4G256_RESIDUAL_MMQ_SRC`, `GEMM_HFQ4G256_RESIDUAL_MMQ_X16_SRC`, `GEMM_HFQ4G256_RESIDUAL_MMQ_X32_SRC`, `GEMM_HFQ4G256_RESIDUAL_MMQ_X32_Y64_SRC`, `GEMM_HFQ4G256_RESIDUAL_SRC`, `GEMM_HFQ4G256_RESIDUAL_WAVE64_DP4A_SRC`, `GEMM_HFQ4G256_RESIDUAL_WAVE64_SRC`, `GEMM_HFQ4G256_RESIDUAL_WMMA2_SRC`, `GEMM_HFQ4G256_RESIDUAL_WMMA_GFX1100_MUSE_BT_SRC`, `GEMM_HFQ4G256_RESIDUAL_WMMA_GFX1100_MUSE_CB_BT_SRC`, `GEMM_HFQ4G256_RESIDUAL_WMMA_GFX1100_MUSE_KS_BT_SRC`, `GEMM_HFQ4G256_RESIDUAL_WMMA_GFX1100_MUSE_LDS_G256_SRC`, `GEMM_HFQ4G256_RESIDUAL_WMMA_GFX1100_MUSE_MW_BT_SRC`, `GEMM_HFQ4G256_RESIDUAL_WMMA_GFX1100_MUSE_RM_BT_SRC`, `GEMM_HFQ4G256_RESIDUAL_WMMA_GFX1100_MUSE_RM_HB_SRC`, `GEMM_HFQ4G256_RESIDUAL_WMMA_GFX1100_MUSE_RM_PIPE_SRC`, `GEMM_HFQ4G256_RESIDUAL_WMMA_GFX1100_MUSE_RM_PK_SRC`, `GEMM_HFQ4G256_RESIDUAL_WMMA_GFX12_BT_SRC`, `GEMM_HFQ4G256_RESIDUAL_WMMA_GFX12_MUSE_SRC`, `GEMM_HFQ4G256_RESIDUAL_WMMA_GFX12_SRC`, `GEMM_HFQ4G256_RESIDUAL_WMMA_K2X32_SRC`, `GEMM_HFQ4G256_RESIDUAL_WMMA_K2_SRC`, `GEMM_HFQ4G256_RESIDUAL_WMMA_K4_SRC`, `GEMM_HFQ4G256_RESIDUAL_WMMA_KSPLIT_DET_SRC`, `GEMM_HFQ4G256_RESIDUAL_WMMA_KSPLIT_SRC`, `GEMM_HFQ4G256_RESIDUAL_WMMA_SRC`, `GEMM_HFQ4G256_SRC`, `GEMM_HFQ4G256_WAVE64_DP4A_SRC`, `GEMM_HFQ4G256_WAVE64_SRC`, `GEMM_HFQ4G256_WMMA_GFX12_SRC`, `GEMM_HFQ4G256_WMMA_SRC`, `GEMM_HFQ6G256_MOE_GROUPED_MMQ_GFX1151_SRC`, `GEMM_HFQ6G256_MOE_GROUPED_WMMA_GFX1151_SRC`, `GEMM_HFQ6G256_MOE_GROUPED_WMMA_GFX12_SRC`, `GEMM_HFQ6G256_MOE_GROUPED_WMMA_K2_SRC`, `GEMM_HFQ6G256_MOE_GROUPED_WMMA_V2_GFX12_SRC`, `GEMM_HFQ6G256_RESIDUAL_FP16_SRC`, `GEMM_HFQ6G256_RESIDUAL_SRC`, `GEMM_HFQ6G256_RESIDUAL_WAVE64_DP4A_SRC`, `GEMM_HFQ6G256_RESIDUAL_WMMA_GFX12_SRC`, `GEMM_HFQ6G256_RESIDUAL_WMMA_K2_SRC`, `GEMM_KSPLIT_DET_FINALIZE_SRC`, `GEMM_MFP4G32_E8G128_MOE_GROUPED_WMMA_GFX1151_SRC`, `GEMM_MFP4G32_E8_MOE_GROUPED_WMMA_GFX1151_SRC`, `GEMM_MFP4G32_E8_MOE_GROUPED_WMMA_GFX12_SRC`, `GEMM_MFP4G32_E8_SOA_GROUPED_WMMA_B2_GFX1151_SRC`, `GEMM_MFP4G32_E8_SOA_GROUPED_WMMA_GFX1151_SRC`, `GEMM_MFP4G32_E8_SOA_GROUPED_WMMA_GFX1201_SRC`, `GEMM_MFP4G32_E8_SOA_WMMA_B16_GFX1151_SRC`, `GEMM_MFP4G32_E8_SOA_WMMA_B16_GFX1201_SRC`, `GEMM_MFP4G32_E8_SOA_WMMA_B2_GFX1151_SRC`, `GEMM_MFP4G32_E8_SOA_WMMA_B2_GFX1201_SRC`, `GEMM_MFP4G32_E8_SOA_WMMA_B4_GFX1151_SRC`, `GEMM_MFP4G32_E8_SOA_WMMA_B4_GFX1201_SRC`, `GEMM_MFP4G32_E8_SOA_WMMA_B8_GFX1151_SRC`, `GEMM_MFP4G32_E8_SOA_WMMA_B8_GFX1201_SRC`, `GEMM_MFP4G32_E8_SOA_WMMA_COOP16_GFX1151_SRC`, `GEMM_MFP4G32_E8_SOA_WMMA_COOP4_GFX1151_SRC`, `GEMM_MFP4G32_E8_SOA_WMMA_GFX1151_SRC`, `GEMM_MFP4G32_E8_SOA_WMMA_GFX1201_SRC`, `GEMM_MIXED_MOE_GROUPED_WMMA_4W_K2_SRC`, `GEMM_MIXED_MOE_GROUPED_WMMA_GFX12_SRC`, `GEMM_MIXED_MOE_GROUPED_WMMA_K2_SRC`, `GEMM_MQ2G256V2_RESIDUAL_WMMA_GFX12_BT_SRC`, `GEMM_MQ2G256V2_RESIDUAL_WMMA_GFX12_SRC`, `GEMM_MQ2G256V2_RESIDUAL_WMMA_SRC`, `GEMM_MQ2G256_LLOYD_MOE_GROUPED_MFMA_GFX942_SRC`, `GEMM_MQ2G256_LLOYD_MOE_GROUPED_MMQ_GFX1030_SRC`, `GEMM_MQ2G256_LLOYD_MOE_GROUPED_MMQ_GFX1151_SRC`, `GEMM_MQ2G256_LLOYD_MOE_GROUPED_MMQ_PERM_GFX1151_SRC`, `GEMM_MQ2G256_LLOYD_MOE_GROUPED_MMQ_PREQUANT_PERM_GFX1151_SRC`, `GEMM_MQ2G256_LLOYD_MOE_GROUPED_WMMA_4W_K2_CND_SRC`, `GEMM_MQ2G256_LLOYD_MOE_GROUPED_WMMA_4W_K2_GFX12_SRC`, `GEMM_MQ2G256_LLOYD_MOE_GROUPED_WMMA_4W_K2_MMQLOAD_NOSYNC_SRC`, `GEMM_MQ2G256_LLOYD_MOE_GROUPED_WMMA_4W_K2_MMQLOAD_SRC`, `GEMM_MQ2G256_LLOYD_MOE_GROUPED_WMMA_4W_K2_N32_GFX12_SRC`, `GEMM_MQ2G256_LLOYD_MOE_GROUPED_WMMA_4W_K2_N32_SRC`, `GEMM_MQ2G256_LLOYD_MOE_GROUPED_WMMA_4W_K2_SRC`, `GEMM_MQ2G256_LLOYD_MOE_GROUPED_WMMA_8W_K2_SRC`, `GEMM_MQ2G256_LLOYD_MOE_GROUPED_WMMA_GFX12_SRC`, `GEMM_MQ2G256_LLOYD_MOE_GROUPED_WMMA_K2_SRC`, `GEMM_MQ3G256V2_RESIDUAL_WMMA_GFX12_BT_SRC`, `GEMM_MQ3G256V2_RESIDUAL_WMMA_GFX12_SRC`, `GEMM_MQ3G256V2_RESIDUAL_WMMA_SRC`, `GEMM_MQ3G256_LLOYD_MOE_GROUPED_MMQ_GFX1151_SRC`, `GEMM_MQ3G256_LLOYD_MOE_GROUPED_WMMA_GFX12_SRC`, `GEMM_MQ3G256_LLOYD_MOE_GROUPED_WMMA_K2_SRC`, `GEMM_MQ3G256_LLOYD_RESIDUAL_WMMA_GFX12_SRC`, `GEMM_MQ3G256_LLOYD_RESIDUAL_WMMA_MB4_SRC`, `GEMM_MQ3G256_LLOYD_RESIDUAL_WMMA_SRC`, `GEMM_MQ4CG256_RESIDUAL_WMMA_GFX12_BT_SRC`, `GEMM_MQ4CG256_RESIDUAL_WMMA_GFX12_SRC`, `GEMM_MQ4G128V2_BATCHED_SRC`, `GEMM_MQ4G128V2_MOE_GROUPED_TOP10_MULTIROW_SRC`, `GEMM_MQ4G128V2_MOE_GROUPED_TOP10_O8_R16_GFX1151_SRC`, `GEMM_MQ4G128V2_MOE_GROUPED_WMMA_GFX1151_SRC`, `GEMM_MQ4G128V2_MULTIROW_GFX1151_SRC`, `GEMM_MQ4G256V2_GATE_UP_SILU_A4_IU4_PM_V2B_GFX1151`, `GEMM_MQ4G256V2_MOE_GROUPED_TOP10_O2_R8_GFX1151_SRC`, `GEMM_MQ4G256V2_MOE_GROUPED_TOP10_O4_R8_X4_GFX1151_SRC`, `GEMM_MQ4G256V2_MOE_GROUPED_TOP10_SIMT_SRC`, `GEMM_MQ4G256V2_MOE_GROUPED_WMMA_GFX12_SRC`, `GEMM_MQ4G256V2_MOE_GROUPED_WMMA_K2_SRC`, `GEMM_MQ4G256V2_RESIDUAL_IU4_PM_GFX1100`, `GEMM_MQ4G256V2_RESIDUAL_IU4_PM_V2B_GFX1151`, `GEMM_MQ4G256V2_RESIDUAL_IU4_V2B_GFX11_SRC`, `GEMM_MQ4G256V2_RESIDUAL_IU4_V2C_GFX11_SRC`, `GEMM_MQ4G256V2_RESIDUAL_MMQ_I8_GFX12_APF0_SRC`, `GEMM_MQ4G256V2_RESIDUAL_MMQ_I8_GFX12_SRC`, `GEMM_MQ4G256V2_RESIDUAL_MMQ_IU4_GFX11_SYMFOLD_SRC`, `GEMM_MQ4G256V2_RESIDUAL_MMQ_IU4_GFX11_X5_SYMFOLD_SRC`, `GEMM_MQ4G256V2_RESIDUAL_MMQ_IU4_GFX12_B1`, `GEMM_MQ4G256V2_RESIDUAL_MMQ_IU4_GFX12_B1S`, `GEMM_MQ4G256V2_RESIDUAL_MMQ_IU4_GFX12_B1_CONTROL`, `GEMM_MQ4G256V2_RESIDUAL_MMQ_IU4_GFX12_G12R_SRC`, `GEMM_MQ4G256V2_RESIDUAL_MMQ_IU4_GFX12_SRC`, `GEMM_MQ4G256V2_RESIDUAL_MMQ_IU4_GFX12_SYMFOLD_G12R_SRC`, `GEMM_MQ4G256V2_RESIDUAL_MMQ_IU4_GFX12_SYMFOLD_SRC`, `GEMM_MQ4G256V2_RESIDUAL_MMQ_IU4_GFX12_V3_SRC`, `GEMM_MQ4G256V2_RESIDUAL_MMQ_IU4_GRIDSPEC_SRC`, `GEMM_MQ4G256V2_RESIDUAL_MMQ_IU4_GRIDSPEC_SYMFOLD_SRC`, `GEMM_MQ4G256V2_RESIDUAL_MMQ_IU4_SRC`, `GEMM_MQ4G256V2_RESIDUAL_MMQ_LUT_PERM_SRC`, `GEMM_MQ4G256V2_RESIDUAL_MMQ_LUT_SRC`, `GEMM_MQ4G256V2_RESIDUAL_MMQ_SRC`, `GEMM_MQ4G256V2_RESIDUAL_WMMA_FP8_GFX12_BT12_LUT_SRC`, `GEMM_MQ4G256V2_RESIDUAL_WMMA_FP8_GFX12_BT12_SRC`, `GEMM_MQ4G256V2_RESIDUAL_WMMA_FP8_GFX12_BT4_LUT_SRC`, `GEMM_MQ4G256V2_RESIDUAL_WMMA_FP8_GFX12_BT4_SRC`, `GEMM_MQ4G256V2_RESIDUAL_WMMA_FP8_GFX12_BT8_LUT_SRC`, `GEMM_MQ4G256V2_RESIDUAL_WMMA_FP8_GFX12_BT8_SRC`, `GEMM_MQ4G256V2_RESIDUAL_WMMA_FP8_GFX12_S2BT8_LUT_SRC`, `GEMM_MQ4G256V2_RESIDUAL_WMMA_FP8_GFX12_S2BT8_SRC`, `GEMM_MQ4G256V2_RESIDUAL_WMMA_FP8_GFX12_V2_B128X128_SRC`, `GEMM_MQ4G256V2_RESIDUAL_WMMA_FP8_GFX12_V2_B128X64W4_SRC`, `GEMM_MQ4G256V2_RESIDUAL_WMMA_FP8_GFX12_V2_B64X256_SRC`, `GEMM_MQ4G256V2_RESIDUAL_WMMA_FP8_GFX12_V2_SRC`, `GEMM_MQ4G256V2_RESIDUAL_WMMA_GFX1100_KSPLIT_LDS_SRC`, `GEMM_MQ4G256V2_RESIDUAL_WMMA_GFX1100_LDSSTAGE_SRC`, `GEMM_MQ4G256V2_RESIDUAL_WMMA_GFX1100_MW_LDS_SRC`, `GEMM_MQ4G256V2_RESIDUAL_WMMA_GFX11_BT_SRC`, `GEMM_MQ4G256V2_RESIDUAL_WMMA_GFX12_BT_SRC`, `GEMM_MQ4G256V2_RESIDUAL_WMMA_GFX12_SRC`, `GEMM_MQ4G256V2_RESIDUAL_WMMA_GFX12_VT_SRC`, `GEMM_MQ4G256V2_RESIDUAL_WMMA_SRC`, `GEMM_MQ4G256V2_WMMA_FP8_GFX12_B1`, `GEMM_MQ4G256_LLOYD_RESIDUAL_WMMA_GFX12_SRC`, `GEMM_MQ4G256_LLOYD_RESIDUAL_WMMA_MB2_SRC`, `GEMM_MQ4G256_LLOYD_RESIDUAL_WMMA_MB4_GFX1151_SRC`, `GEMM_MQ4G256_LLOYD_RESIDUAL_WMMA_MB4_SRC`, `GEMM_MQ4G256_LLOYD_RESIDUAL_WMMA_SRC`, `GEMM_MQ5G256V2_RESIDUAL_WMMA_GFX12_BT_SRC`, `GEMM_MQ5G256V2_RESIDUAL_WMMA_GFX12_SRC`, `GEMM_MQ5G256V2_RESIDUAL_WMMA_SRC`, `GEMM_MQ6G256V2_F32_ROWS_SRC`, `GEMM_MQ6G256V2_MOE_GROUPED_WMMA_GFX12_SRC`, `GEMM_MQ6G256V2_MOE_GROUPED_WMMA_K2_SRC`, `GEMM_MQ6G256V2_RESIDUAL_WMMA_GFX12_BT_SRC`, `GEMM_MQ6G256V2_RESIDUAL_WMMA_GFX12_SRC`, `GEMM_MQ6G256V2_RESIDUAL_WMMA_SRC`, `GEMM_MQ6G256_SRC`, `GEMM_MQV2_WMMA_GFX11_BT_SRC`, `GEMM_MQV2_WMMA_GFX11_MW_LDS_SRC`, `GEMM_MW16_RESIDUAL_WMMA_SRC`, `GEMM_PARO_Q4G128_MOE_GROUPED_MMQ_GFX1151_SRC`, `GEMM_PARO_Q4G128_MOE_GROUPED_MMQ_K8_GFX1151_SRC`, `GEMM_PARO_Q4G128_MOE_GROUPED_WMMA_K2_SRC`, `GEMM_Q8_0_BATCHED_SRC`, `GEMM_Q8_0_BATCHED_WIDE_EXACT_SRC`, `GEMM_Q8_0_MMQ_4W_GFX1151_SRC`, `GEMM_Q8_0_MMQ_GFX1030_SRC`, `GEMM_Q8_0_MOE_GROUPED_TOP10_SRC`, `GEMM_Q8_0_RESIDUAL_WMMA_GFX12_SRC`, `GEMM_Q8_0_RESIDUAL_WMMA_SRC`, `GEMM_Q8_0_WMMA_4W_SRC`, `GEMM_Q8_0_WMMA_GFX12_SRC`, `GEMM_Q8_0_WMMA_SRC`, `GEMM_Q8_0_WMMA_X64_SRC`, `GEMM_QKVZA_HFP4G32_WMMA_GFX12_SRC`, `GEMM_QKVZA_HFP4G32_WMMA_SRC`, `GEMM_QKVZA_HFQ3G256_DOT2_SRC`, `GEMM_QKVZA_HFQ3G256_FP16_SRC`, `GEMM_QKVZA_HFQ3G256_MMQ_BODY_CUH`, `GEMM_QKVZA_HFQ3G256_MMQ_X16_SRC`, `GEMM_QKVZA_HFQ3G256_MMQ_X32_SRC`, `GEMM_QKVZA_HFQ3G256_MMQ_X8_SRC`, `GEMM_QKVZA_HFQ3G256_SRC`, `GEMM_QKVZA_HFQ3G256_WMMA_GFX12_SRC`, `GEMM_QKVZA_HFQ3G256_WMMA_MB4_SRC`, `GEMM_QKVZA_HFQ3G256_WMMA_SRC`, `GEMM_QKVZA_HFQ4G256_DOT2_SRC`, `GEMM_QKVZA_HFQ4G256_FP16_SRC`, `GEMM_QKVZA_HFQ4G256_FP16_WAVE64_SRC`, `GEMM_QKVZA_HFQ4G256_MMQ_BODY_CUH`, `GEMM_QKVZA_HFQ4G256_MMQ_X16_SRC`, `GEMM_QKVZA_HFQ4G256_MMQ_X32_SRC`, `GEMM_QKVZA_HFQ4G256_SRC`, `GEMM_QKVZA_HFQ4G256_WAVE64_DP4A_SRC`, `GEMM_QKVZA_HFQ4G256_WAVE64_SRC`, `GEMM_QKVZA_HFQ4G256_WMMA_GFX12_BT_SRC`, `GEMM_QKVZA_HFQ4G256_WMMA_GFX12_SRC`, `GEMM_QKVZA_HFQ4G256_WMMA_SRC`, `GEMM_QKVZA_HFQ6G256_DOT2_SRC`, `GEMM_QKVZA_HFQ6G256_FP16_SRC`, `GEMM_QKVZA_HFQ6G256_SRC`, `GEMM_QKVZA_HFQ6G256_WAVE64_DP4A_SRC`, `GEMM_QKVZA_HFQ6G256_WMMA_GFX12_SRC`, `GEMM_QKVZA_HFQ6G256_WMMA_SRC`, `GEMM_QKVZA_MQ2G256V2_WMMA_GFX12_BT_SRC`, `GEMM_QKVZA_MQ2G256V2_WMMA_GFX12_SRC`, `GEMM_QKVZA_MQ2G256V2_WMMA_SRC`, `GEMM_QKVZA_MQ3G256V2_WMMA_GFX12_BT_SRC`, `GEMM_QKVZA_MQ3G256V2_WMMA_GFX12_SRC`, `GEMM_QKVZA_MQ3G256V2_WMMA_SRC`, `GEMM_QKVZA_MQ3G256_LLOYD_WMMA_GFX12_SRC`, `GEMM_QKVZA_MQ3G256_LLOYD_WMMA_MB4_SRC`, `GEMM_QKVZA_MQ3G256_LLOYD_WMMA_SRC`, `GEMM_QKVZA_MQ4CG256_WMMA_GFX12_BT_SRC`, `GEMM_QKVZA_MQ4CG256_WMMA_GFX12_SRC`, `GEMM_QKVZA_MQ4G256V2_WMMA_FP8_GFX12_BT12_LUT_SRC`, `GEMM_QKVZA_MQ4G256V2_WMMA_FP8_GFX12_BT12_SRC`, `GEMM_QKVZA_MQ4G256V2_WMMA_FP8_GFX12_BT4_LUT_SRC`, `GEMM_QKVZA_MQ4G256V2_WMMA_FP8_GFX12_BT4_SRC`, `GEMM_QKVZA_MQ4G256V2_WMMA_FP8_GFX12_BT8_LUT_SRC`, `GEMM_QKVZA_MQ4G256V2_WMMA_FP8_GFX12_BT8_SRC`, `GEMM_QKVZA_MQ4G256V2_WMMA_FP8_GFX12_S2BT8_LUT_SRC`, `GEMM_QKVZA_MQ4G256V2_WMMA_FP8_GFX12_S2BT8_SRC`, `GEMM_QKVZA_MQ4G256V2_WMMA_FP8_GFX12_V2_B128X128_SRC`, `GEMM_QKVZA_MQ4G256V2_WMMA_FP8_GFX12_V2_B128X64W4_SRC`, `GEMM_QKVZA_MQ4G256V2_WMMA_FP8_GFX12_V2_B64X256_SRC`, `GEMM_QKVZA_MQ4G256V2_WMMA_FP8_GFX12_V2_SRC`, `GEMM_QKVZA_MQ4G256V2_WMMA_GFX11_BT_SRC`, `GEMM_QKVZA_MQ4G256V2_WMMA_GFX12_BT_SRC`, `GEMM_QKVZA_MQ4G256V2_WMMA_GFX12_SRC`, `GEMM_QKVZA_MQ4G256V2_WMMA_GFX12_VT_SRC`, `GEMM_QKVZA_MQ4G256V2_WMMA_SRC`, `GEMM_QKVZA_MQ4G256_LLOYD_WMMA_GFX1151_SRC`, `GEMM_QKVZA_MQ4G256_LLOYD_WMMA_GFX12_SRC`, `GEMM_QKVZA_MQ4G256_LLOYD_WMMA_MB4_GFX1151_SRC`, `GEMM_QKVZA_MQ4G256_LLOYD_WMMA_MB4_SRC`, `GEMM_QKVZA_MQ4G256_LLOYD_WMMA_SRC`, `GEMM_QKVZA_MQ5G256V2_WMMA_GFX12_BT_SRC`, `GEMM_QKVZA_MQ5G256V2_WMMA_GFX12_SRC`, `GEMM_QKVZA_MQ5G256V2_WMMA_SRC`, `GEMM_QKVZA_MQ6G256V2_WMMA_GFX12_BT_SRC`, `GEMM_QKVZA_MQ6G256V2_WMMA_GFX12_SRC`, `GEMM_QKVZA_MQ6G256V2_WMMA_SRC`, `GEMM_QKVZA_Q8_0_WMMA_GFX12_SRC`, `GEMM_QKVZA_Q8_0_WMMA_SRC`, `GEMM_QKVZA_TQ2G128_SRC`, `GEMM_QKV_HFP4G32_WMMA_FP8_GFX12_SRC`, `GEMM_QKV_HFP4G32_WMMA_GFX12_SRC`, `GEMM_QKV_HFP4G32_WMMA_SRC`, `GEMM_QKV_HFQ3G256_DOT2_SRC`, `GEMM_QKV_HFQ3G256_DP4A_SRC`, `GEMM_QKV_HFQ3G256_FP16_SRC`, `GEMM_QKV_HFQ3G256_MMQ_BODY_CUH`, `GEMM_QKV_HFQ3G256_MMQ_X16_SRC`, `GEMM_QKV_HFQ3G256_MMQ_X32_SRC`, `GEMM_QKV_HFQ3G256_MMQ_X8_SRC`, `GEMM_QKV_HFQ3G256_SRC`, `GEMM_QKV_HFQ3G256_WMMA_GFX12_SRC`, `GEMM_QKV_HFQ3G256_WMMA_MB4_SRC`, `GEMM_QKV_HFQ3G256_WMMA_SRC`, `GEMM_QKV_HFQ4G256_DOT2_SRC`, `GEMM_QKV_HFQ4G256_FP16_SRC`, `GEMM_QKV_HFQ4G256_FP16_WAVE64_SRC`, `GEMM_QKV_HFQ4G256_MMQ_BODY_CUH`, `GEMM_QKV_HFQ4G256_MMQ_GFX906_BODY_CUH`, `GEMM_QKV_HFQ4G256_MMQ_GFX906_X16_SRC`, `GEMM_QKV_HFQ4G256_MMQ_GFX906_X32_SRC`, `GEMM_QKV_HFQ4G256_MMQ_GFX906_X64_SRC`, `GEMM_QKV_HFQ4G256_MMQ_GFX906_X8_SRC`, `GEMM_QKV_HFQ4G256_MMQ_X16_SRC`, `GEMM_QKV_HFQ4G256_MMQ_X32_SRC`, `GEMM_QKV_HFQ4G256_SRC`, `GEMM_QKV_HFQ4G256_WAVE64_DP4A_SRC`, `GEMM_QKV_HFQ4G256_WAVE64_SRC`, `GEMM_QKV_HFQ4G256_WMMA_GFX12_SRC`, `GEMM_QKV_HFQ4G256_WMMA_SRC`, `GEMM_QKV_HFQ6G256_DOT2_SRC`, `GEMM_QKV_HFQ6G256_FP16_SRC`, `GEMM_QKV_HFQ6G256_SRC`, `GEMM_QKV_HFQ6G256_WAVE64_DP4A_SRC`, `GEMM_QKV_HFQ6G256_WMMA_GFX12_SRC`, `GEMM_QKV_HFQ6G256_WMMA_SRC`, `GEMM_QKV_MQ2G256V2_WMMA_GFX12_SRC`, `GEMM_QKV_MQ2G256V2_WMMA_SRC`, `GEMM_QKV_MQ3G256V2_WMMA_GFX12_SRC`, `GEMM_QKV_MQ3G256V2_WMMA_SRC`, `GEMM_QKV_MQ3G256_LLOYD_WMMA_GFX12_SRC`, `GEMM_QKV_MQ3G256_LLOYD_WMMA_MB4_SRC`, `GEMM_QKV_MQ3G256_LLOYD_WMMA_SRC`, `GEMM_QKV_MQ4CG256_WMMA_GFX12_SRC`, `GEMM_QKV_MQ4G256V2_WMMA_FP8_GFX12_BT12_LUT_SRC`, `GEMM_QKV_MQ4G256V2_WMMA_FP8_GFX12_BT12_SRC`, `GEMM_QKV_MQ4G256V2_WMMA_FP8_GFX12_BT4_LUT_SRC`, `GEMM_QKV_MQ4G256V2_WMMA_FP8_GFX12_BT4_SRC`, `GEMM_QKV_MQ4G256V2_WMMA_FP8_GFX12_BT8_LUT_SRC`, `GEMM_QKV_MQ4G256V2_WMMA_FP8_GFX12_BT8_SRC`, `GEMM_QKV_MQ4G256V2_WMMA_FP8_GFX12_S2BT8_LUT_SRC`, `GEMM_QKV_MQ4G256V2_WMMA_FP8_GFX12_S2BT8_SRC`, `GEMM_QKV_MQ4G256V2_WMMA_FP8_GFX12_V2_B128X128_SRC`, `GEMM_QKV_MQ4G256V2_WMMA_FP8_GFX12_V2_B128X64W4_SRC`, `GEMM_QKV_MQ4G256V2_WMMA_FP8_GFX12_V2_B64X256_SRC`, `GEMM_QKV_MQ4G256V2_WMMA_FP8_GFX12_V2_SRC`, `GEMM_QKV_MQ4G256V2_WMMA_GFX11_BT_SRC`, `GEMM_QKV_MQ4G256V2_WMMA_GFX1201_BT_SRC`, `GEMM_QKV_MQ4G256V2_WMMA_GFX12_SRC`, `GEMM_QKV_MQ4G256V2_WMMA_GFX12_VT_SRC`, `GEMM_QKV_MQ4G256V2_WMMA_SRC`, `GEMM_QKV_MQ4G256_LLOYD_WMMA_GFX1151_SRC`, `GEMM_QKV_MQ4G256_LLOYD_WMMA_GFX12_SRC`, `GEMM_QKV_MQ4G256_LLOYD_WMMA_MB4_GFX1151_SRC`, `GEMM_QKV_MQ4G256_LLOYD_WMMA_MB4_SRC`, `GEMM_QKV_MQ4G256_LLOYD_WMMA_SRC`, `GEMM_QKV_MQ5G256V2_WMMA_GFX12_SRC`, `GEMM_QKV_MQ5G256V2_WMMA_SRC`, `GEMM_QKV_MQ6G256V2_WMMA_GFX12_SRC`, `GEMM_QKV_MQ6G256V2_WMMA_SRC`, `GEMM_QKV_MQV2_WMMA_GFX1201_BT_SRC`, `GEMM_QKV_Q8_0_WMMA_GFX12_SRC`, `GEMM_QKV_Q8_0_WMMA_SRC`, `GEMM_TQ2G128_PREFILL_SRC`, `GEMM_TQ2G128_WMMA_SRC`, `GEMV_BF16_XF32_SRC`, `GEMV_BQ1G128_SRC`, `GEMV_BQ1G128_XBATCH_SRC`, `GEMV_F16_BIAS_XF32_SRC`, `GEMV_F16_XF32_SRC`, `GEMV_HFP4G32_DOT2_GFX11_SRC`, `GEMV_HFP4G32_FP8_GFX12_SRC`, `GEMV_HFP4G32_GFX1100_SRC`, `GEMV_HFP4G32_SRC`, `GEMV_HFQ2G128_SRC`, `GEMV_HFQ2G256_SRC`, `GEMV_HFQ3G128_SRC`, `GEMV_HFQ3G256_GFX1100_SRC`, `GEMV_HFQ3G256_RESIDUAL_GFX1100_SRC`, `GEMV_HFQ3G256_RESIDUAL_SRC`, `GEMV_HFQ3G256_SRC`, `GEMV_HFQ4G1024_SRC`, `GEMV_HFQ4G128_MOE_DOWN_RESIDUAL_SCALED_K8_INDEXED_SRC`, `GEMV_HFQ4G128_RESIDUAL_SIGMOID_SCALED_SRC`, `GEMV_HFQ4G128_SRC`, `GEMV_HFQ4G256_GFX1030_V1_SRC`, `GEMV_HFQ4G256_GFX1030_V2_SRC`, `GEMV_HFQ4G256_GFX1030_V3_SRC`, `GEMV_HFQ4G256_GFX1030_V4_SRC`, `GEMV_HFQ4G256_GFX1030_V5_SRC`, `GEMV_HFQ4G256_GFX1100_SRC`, `GEMV_HFQ4G256_K2048_GFX1100_SRC`, `GEMV_HFQ4G256_LM_HEAD_DOT2_GFX1151_SRC`, `GEMV_HFQ4G256_LM_HEAD_R1_HYBRID_BUFFER_GFX1151_SRC`, `GEMV_HFQ4G256_MOE_DOWN_INDEXED_BATCHED_SRC`, `GEMV_HFQ4G256_MOE_DOWN_INDEXED_BATCHED_WAVE64_SRC`, `GEMV_HFQ4G256_MOE_DOWN_INDEXED_SRC`, `GEMV_HFQ4G256_MOE_DOWN_INDEXED_WAVE64_SRC`, `GEMV_HFQ4G256_MOE_DOWN_K8_INDEXED_BATCHED_EXPANDED_BUFFER_GFX1151_SRC`, `GEMV_HFQ4G256_MOE_DOWN_K8_INDEXED_BATCHED_EXPANDED_CPOL_SLC_GFX1100_SRC`, `GEMV_HFQ4G256_MOE_DOWN_K8_INDEXED_BATCHED_EXPANDED_HYBRID_BUFFER_GFX1151_SRC`, `GEMV_HFQ4G256_MOE_DOWN_K8_INDEXED_BATCHED_EXPANDED_ROW1_BUFFER_GFX1151_SRC`, `GEMV_HFQ4G256_MOE_DOWN_K8_INDEXED_BATCHED_EXPANDED_ROW2_BUFFER_GFX1151_SRC`, `GEMV_HFQ4G256_MOE_DOWN_K8_INDEXED_BATCHED_EXPANDED_ROW2_CLUSTERED_GFX1151_SRC`, `GEMV_HFQ4G256_MOE_DOWN_K8_INDEXED_BATCHED_EXPANDED_ROW8_GFX1151_SRC`, `GEMV_HFQ4G256_MOE_DOWN_K8_INDEXED_BATCHED_EXPANDED_SRC`, `GEMV_HFQ4G256_MOE_DOWN_K8_INDEXED_FUSED_ACC_SRC`, `GEMV_HFQ4G256_MOE_DOWN_K8_INDEXED_LAST_COMBINE_SRC`, `GEMV_HFQ4G256_MOE_DOWN_SRC`, `GEMV_HFQ4G256_MOE_GATE_INDEXED_K2048_GFX1151_SRC`, `GEMV_HFQ4G256_MOE_GATE_UP_INDEXED_BATCHED_SRC`, `GEMV_HFQ4G256_MOE_GATE_UP_INDEXED_BATCHED_WAVE64_SRC`, `GEMV_HFQ4G256_MOE_GATE_UP_INDEXED_CPOL_DLC_GFX1100_SRC`, `GEMV_HFQ4G256_MOE_GATE_UP_INDEXED_CPOL_GLC_GFX1100_SRC`, `GEMV_HFQ4G256_MOE_GATE_UP_INDEXED_CPOL_SLC_GFX1100_SRC`, `GEMV_HFQ4G256_MOE_GATE_UP_INDEXED_K2048_ALL_BUFFER_GFX1151_SRC`, `GEMV_HFQ4G256_MOE_GATE_UP_INDEXED_K2048_BUFFER_GFX1151_SRC`, `GEMV_HFQ4G256_MOE_GATE_UP_INDEXED_K2048_GFX1100_SRC`, `GEMV_HFQ4G256_MOE_GATE_UP_INDEXED_K2048_GFX1151_SRC`, `GEMV_HFQ4G256_MOE_GATE_UP_INDEXED_K2048_HYBRID_GFX1151_SRC`, `GEMV_HFQ4G256_MOE_GATE_UP_INDEXED_K2048_LOW_VGPR_GFX1151_SRC`, `GEMV_HFQ4G256_MOE_GATE_UP_INDEXED_K2048_PAIR_ALL_BUFFER_GFX1151_SRC`, `GEMV_HFQ4G256_MOE_GATE_UP_INDEXED_K2048_PAIR_BUFFER_GFX1151_SRC`, `GEMV_HFQ4G256_MOE_GATE_UP_INDEXED_K2048_PAIR_VGPR_GFX1151_SRC`, `GEMV_HFQ4G256_MOE_GATE_UP_INDEXED_K2048_ROUTE_ALL_BUFFER_GFX1151_SRC`, `GEMV_HFQ4G256_MOE_GATE_UP_INDEXED_K2816_CPOL_DLC_GFX1100_SRC`, `GEMV_HFQ4G256_MOE_GATE_UP_INDEXED_K2816_CPOL_GLC_GFX1100_SRC`, `GEMV_HFQ4G256_MOE_GATE_UP_INDEXED_K2816_CPOL_SLC_GFX1100_SRC`, `GEMV_HFQ4G256_MOE_GATE_UP_INDEXED_K2816_SRC`, `GEMV_HFQ4G256_MOE_GATE_UP_INDEXED_LOW_VGPR_GFX1100_SRC`, `GEMV_HFQ4G256_MOE_GATE_UP_INDEXED_PAIR_SLC_GFX1100_SRC`, `GEMV_HFQ4G256_MOE_GATE_UP_INDEXED_RANK_INTERLEAVE_GFX1100_SRC`, `GEMV_HFQ4G256_MOE_GATE_UP_INDEXED_ROWTILE_SRC`, `GEMV_HFQ4G256_MOE_GATE_UP_INDEXED_SRC`, `GEMV_HFQ4G256_MOE_GATE_UP_INDEXED_WAVE64_GFX1151_SRC`, `GEMV_HFQ4G256_MOE_GATE_UP_INDEXED_WAVE64_SRC`, `GEMV_HFQ4G256_MOE_GATE_UP_INDEXED_WG2_SRC`, `GEMV_HFQ4G256_MOE_GATE_UP_PAIRED_WAVES_K2048_GFX1151_SRC`, `GEMV_HFQ4G256_MOE_GATE_UP_PERSISTENT_RANK8_GFX1151_SRC`, `GEMV_HFQ4G256_MOE_GATE_UP_SRC`, `GEMV_HFQ4G256_MOE_NINEPATH_D3_SRC`, `GEMV_HFQ4G256_MOE_NINEPATH_D4_SRC`, `GEMV_HFQ4G256_MOE_UP_INDEXED_K2048_GFX1151_SRC`, `GEMV_HFQ4G256_MULTIROW_ALL_BUFFER_GFX1151_SRC`, `GEMV_HFQ4G256_MULTIROW_BUFFER_GFX1151_SRC`, `GEMV_HFQ4G256_MULTIROW_GFX1100_SRC`, `GEMV_HFQ4G256_MULTIROW_HYBRID_BUFFER_GFX1151_SRC`, `GEMV_HFQ4G256_MULTIROW_HYBRID_DLC_GFX1151_SRC`, `GEMV_HFQ4G256_MULTIROW_HYBRID_GLC_GFX1151_SRC`, `GEMV_HFQ4G256_MULTIROW_HYBRID_SLC_GFX1151_SRC`, `GEMV_HFQ4G256_MULTIROW_K2048_GFX1151_SRC`, `GEMV_HFQ4G256_MULTIROW_SRC`, `GEMV_HFQ4G256_RESIDUAL_BUFFER_GFX1151_SRC`, `GEMV_HFQ4G256_RESIDUAL_CPOL_RT_GFX1100_SRC`, `GEMV_HFQ4G256_RESIDUAL_CPOL_RT_LOW_GFX1100_SRC`, `GEMV_HFQ4G256_RESIDUAL_CPOL_SLC_GFX1100_SRC`, `GEMV_HFQ4G256_RESIDUAL_GFX1100_SRC`, `GEMV_HFQ4G256_RESIDUAL_GFX942_SRC`, `GEMV_HFQ4G256_RESIDUAL_HYBRID_BUFFER_GFX1151_SRC`, `GEMV_HFQ4G256_RESIDUAL_K2048_GFX1100_SRC`, `GEMV_HFQ4G256_RESIDUAL_K4096_GFX1151_SRC`, `GEMV_HFQ4G256_RESIDUAL_MULTIROW_GFX1100_SRC`, `GEMV_HFQ4G256_RESIDUAL_MULTIROW_GFX1151_SRC`, `GEMV_HFQ4G256_RESIDUAL_ROW1_GFX1151_SRC`, `GEMV_HFQ4G256_RESIDUAL_RT_LOW_GFX1151_SRC`, `GEMV_HFQ4G256_RESIDUAL_SCALED_SRC`, `GEMV_HFQ4G256_RESIDUAL_SIGMOID_BUFFER_GFX1100_SRC`, `GEMV_HFQ4G256_RESIDUAL_SIGMOID_K512_BUFFER_GFX1100_SRC`, `GEMV_HFQ4G256_RESIDUAL_SIGMOID_K512_BUFFER_GFX1151_SRC`, `GEMV_HFQ4G256_RESIDUAL_SIGMOID_K512_HOIST_X16_BUFFER_GFX1100_SRC`, `GEMV_HFQ4G256_RESIDUAL_SIGMOID_ROWS4_BUFFER_GFX1100_SRC`, `GEMV_HFQ4G256_RESIDUAL_SIGMOID_ROWS4_GFX1100_SRC`, `GEMV_HFQ4G256_RESIDUAL_SRC`, `GEMV_HFQ4G256_RESIDUAL_STAGE_X32_GFX1100_SRC`, `GEMV_HFQ4G256_RESIDUAL_V2_GFX942_SRC`, `GEMV_HFQ4G256_RESIDUAL_V3_GFX942_SRC`, `GEMV_HFQ4G256_RESIDUAL_WAVE64_GFX1151_SRC`, `GEMV_HFQ4G256_RESIDUAL_WAVE64_PREFETCH_SRC`, `GEMV_HFQ4G256_RESIDUAL_WAVE64_SRC`, `GEMV_HFQ4G256_SRC`, `GEMV_HFQ4G256_WIDE_SRC`, `GEMV_HFQ4G256_XBATCH_GEN_SRC`, `GEMV_HFQ4G256_XBATCH_SRC`, `GEMV_HFQ4G512_SRC`, `GEMV_HFQ5G256_MOE_DOWN_K8_INDEXED_BATCHED_EXPANDED_SRC`, `GEMV_HFQ5G256_MOE_GATE_UP_INDEXED_BATCHED_SRC`, `GEMV_HFQ5G256_MOE_GATE_UP_INDEXED_SRC`, `GEMV_HFQ5G256_RESIDUAL_SIGMOID_SCALED_SRC`, `GEMV_HFQ5G256_RESIDUAL_SRC`, `GEMV_HFQ5G256_RESIDUAL_WAVE64_PREFETCH_SRC`, `GEMV_HFQ5G256_RESIDUAL_WAVE64_SRC`, `GEMV_HFQ5G256_SRC`, `GEMV_HFQ6G256_MOE_DOWN_K8_INDEXED_BATCHED_EXPANDED_SRC`, `GEMV_HFQ6G256_MOE_GATE_UP_INDEXED_BATCHED_SRC`, `GEMV_HFQ6G256_MOE_GATE_UP_INDEXED_SRC`, `GEMV_HFQ6G256_RESIDUAL_SIGMOID_SCALED_SRC`, `GEMV_HFQ6G256_RESIDUAL_SRC`, `GEMV_HFQ6G256_RESIDUAL_WAVE64_PREFETCH_SRC`, `GEMV_HFQ6G256_RESIDUAL_WAVE64_SRC`, `GEMV_HFQ6G256_SRC`, `GEMV_HFQ8G256_SRC`, `GEMV_MFP3G32_E8_SRC`, `GEMV_MFP4G32_E8G128_MOE_DOWN_K8_INDEXED_BATCHED_EXPANDED_SRC`, `GEMV_MFP4G32_E8G128_MOE_GATE_UP_K8_INDEXED_BATCHED_SRC`, `GEMV_MFP4G32_E8G128_SRC`, `GEMV_MFP4G32_E8_GFX1151_SRC`, `GEMV_MFP4G32_E8_LDSX_GFX1151_SRC`, `GEMV_MFP4G32_E8_MOE_DOWN_K8_INDEXED_BATCHED_EXPANDED_GFX1151_SRC`, `GEMV_MFP4G32_E8_MOE_DOWN_K8_INDEXED_BATCHED_EXPANDED_GFX11_DGPU_SRC`, `GEMV_MFP4G32_E8_MOE_GATE_UP_K8_INDEXED_BATCHED_GFX1151_SRC`, `GEMV_MFP4G32_E8_MOE_GATE_UP_K8_INDEXED_BATCHED_GFX11_DGPU_SRC`, `GEMV_MFP4G32_E8_SOA_BATCHED_B16_GFX1151_SRC`, `GEMV_MFP4G32_E8_SOA_BATCHED_B1_GFX1151_SRC`, `GEMV_MFP4G32_E8_SOA_BATCHED_B2_GFX1151_SRC`, `GEMV_MFP4G32_E8_SOA_BATCHED_B3_GFX1151_SRC`, `GEMV_MFP4G32_E8_SOA_BATCHED_B4_GFX1151_SRC`, `GEMV_MFP4G32_E8_SOA_BATCHED_B5_GFX1151_SRC`, `GEMV_MFP4G32_E8_SOA_BATCHED_B6_GFX1151_SRC`, `GEMV_MFP4G32_E8_SOA_BATCHED_B7_GFX1151_SRC`, `GEMV_MFP4G32_E8_SOA_BATCHED_B8_GFX1151_SRC`, `GEMV_MFP4G32_E8_SOA_BATCHED_PACK_B3_GFX1151_SRC`, `GEMV_MFP4G32_E8_SOA_BATCHED_PAIR_B3_GFX1151_SRC`, `GEMV_MFP4G32_E8_SOA_FP8_MFMA_GFX942_SRC`, `GEMV_MFP4G32_E8_SOA_GFX1151_SRC`, `GEMV_MFP4G32_E8_SOA_GFX1201_SRC`, `GEMV_MFP4G32_E8_SOA_GFX942_SRC`, `GEMV_MFP4G32_E8_SOA_GROUPED_BATCHED_B16_GFX1151_SRC`, `GEMV_MFP4G32_E8_SOA_GROUPED_BATCHED_B1_GFX1151_SRC`, `GEMV_MFP4G32_E8_SOA_GROUPED_BATCHED_B2_GFX1151_SRC`, `GEMV_MFP4G32_E8_SOA_GROUPED_BATCHED_B3_GFX1151_SRC`, `GEMV_MFP4G32_E8_SOA_GROUPED_BATCHED_B4_GFX1151_SRC`, `GEMV_MFP4G32_E8_SOA_GROUPED_BATCHED_B5_GFX1151_SRC`, `GEMV_MFP4G32_E8_SOA_GROUPED_BATCHED_B6_GFX1151_SRC`, `GEMV_MFP4G32_E8_SOA_GROUPED_BATCHED_B7_GFX1151_SRC`, `GEMV_MFP4G32_E8_SOA_GROUPED_BATCHED_B8_GFX1151_SRC`, `GEMV_MFP4G32_E8_SOA_GROUPED_BUFFER_GFX1151_SRC`, `GEMV_MFP4G32_E8_SOA_GROUPED_GFX1100_SRC`, `GEMV_MFP4G32_E8_SOA_GROUPED_GFX1151_SRC`, `GEMV_MFP4G32_E8_SOA_GROUPED_GFX1201_SRC`, `GEMV_MFP4G32_E8_SOA_LATE_SCALE_GFX1201_SRC`, `GEMV_MFP4G32_E8_SOA_LUT_SRC`, `GEMV_MFP4G32_E8_SOA_MIXED_JOBS_GFX1201_SRC`, `GEMV_MFP4G32_E8_SOA_MOE_GATE_UP_K8_INDEXED_BATCHED_SRC`, `GEMV_MFP4G32_E8_SOA_PAIR_GFX942_SRC`, `GEMV_MFP4G32_E8_SOA_PREFETCH4_GFX1201_SRC`, `GEMV_MFP4G32_E8_SOA_SHARED_JOBS_GFX1100_SRC`, `GEMV_MFP4G32_E8_SOA_SHARED_PAIR_GFX1201_SRC`, `GEMV_MFP4G32_E8_SOA_SRC`, `GEMV_MFP4G32_E8_SOA_STRIP_SRC`, `GEMV_MFP4G32_E8_SOA_U4_BUFFER_CPOL0_GFX1151_SRC`, `GEMV_MFP4G32_E8_SOA_U4_BUFFER_CPOL20_GFX1151_SRC`, `GEMV_MFP4G32_E8_SOA_U4_BUFFER_CPOL22_GFX1151_SRC`, `GEMV_MFP4G32_E8_SOA_U4_BUFFER_SQRT_SOFTPLUS_GFX1151_SRC`, `GEMV_MFP4G32_E8_SOA_U4_SRC`, `GEMV_MFP4G32_E8_SOA_U8_SRC`, `GEMV_MFP4G32_E8_SOA_W128_GFX942_SRC`, `GEMV_MFP4G32_E8_SRC`, `GEMV_MFP4G32_E8_STRIP_GFX1151_SRC`, `GEMV_MFP4G32_LLOYD_SRC`, `GEMV_MFP4G32_P_SRC`, `GEMV_MIXED_MOE_DOWN_K8_INDEXED_BATCHED_EXPANDED_SRC`, `GEMV_MIXED_MOE_GATE_UP_INDEXED_BATCHED_SRC`, `GEMV_MQ2G256GL_MOE_DOWN_INDEXED_SRC`, `GEMV_MQ2G256GL_MOE_GATE_UP_INDEXED_SRC`, `GEMV_MQ2G256V2_MULTIROW_SRC`, `GEMV_MQ2G256V2_RESIDUAL_SRC`, `GEMV_MQ2G256V2_SRC`, `GEMV_MQ2G256_LLOYD_MOE_DOWN_ALLRANKS_GFX942_SRC`, `GEMV_MQ2G256_LLOYD_MOE_DOWN_EXPANDED_K4_GFX942_SRC`, `GEMV_MQ2G256_LLOYD_MOE_DOWN_EXPANDED_K4_SRC`, `GEMV_MQ2G256_LLOYD_MOE_DOWN_INDEXED_BATCHED_K4_SRC`, `GEMV_MQ2G256_LLOYD_MOE_DOWN_INDEXED_BATCHED_K8ALL_SRC`, `GEMV_MQ2G256_LLOYD_MOE_DOWN_INDEXED_GFX942_SRC`, `GEMV_MQ2G256_LLOYD_MOE_DOWN_INDEXED_MFMA_F16_GFX942_SRC`, `GEMV_MQ2G256_LLOYD_MOE_DOWN_INDEXED_MFMA_GFX942_SRC`, `GEMV_MQ2G256_LLOYD_MOE_DOWN_INDEXED_R4_SRC`, `GEMV_MQ2G256_LLOYD_MOE_DOWN_INDEXED_SRC`, `GEMV_MQ2G256_LLOYD_MOE_GATE_UP_INDEXED_BATCHED_K4096_LDS_SRC`, `GEMV_MQ2G256_LLOYD_MOE_GATE_UP_INDEXED_BATCHED_K4_SRC`, `GEMV_MQ2G256_LLOYD_MOE_GATE_UP_INDEXED_GFX942_SRC`, `GEMV_MQ2G256_LLOYD_MOE_GATE_UP_INDEXED_MFMA_F16_GFX942_SRC`, `GEMV_MQ2G256_LLOYD_MOE_GATE_UP_INDEXED_MFMA_GFX942_SRC`, `GEMV_MQ2G256_LLOYD_MOE_GATE_UP_INDEXED_SRC`, `GEMV_MQ2G256_LLOYD_SRC`, `GEMV_MQ3G256GL_MOE_DOWN_INDEXED_SRC`, `GEMV_MQ3G256GL_MOE_GATE_UP_INDEXED_SRC`, `GEMV_MQ3G256V2_MULTIROW_SRC`, `GEMV_MQ3G256V2_RESIDUAL_SRC`, `GEMV_MQ3G256V2_SRC`, `GEMV_MQ3G256_LLOYD_GFX1100_SRC`, `GEMV_MQ3G256_LLOYD_MOE_DOWN_INDEXED_BATCHED_K4_SRC`, `GEMV_MQ3G256_LLOYD_MOE_DOWN_INDEXED_R4_SRC`, `GEMV_MQ3G256_LLOYD_MOE_DOWN_INDEXED_SRC`, `GEMV_MQ3G256_LLOYD_MOE_GATE_UP_INDEXED_SRC`, `GEMV_MQ3G256_LLOYD_MOE_NINEPATH_D4_SRC`, `GEMV_MQ3G256_LLOYD_RESIDUAL_GFX1100_SRC`, `GEMV_MQ3G256_LLOYD_RESIDUAL_SRC`, `GEMV_MQ3G256_LLOYD_SRC`, `GEMV_MQ4CG256_MULTIROW_SRC`, `GEMV_MQ4CG256_RESIDUAL_SRC`, `GEMV_MQ4CG256_SRC`, `GEMV_MQ4G128V2_MOE_DOWN_TOP10_INDEXED_BATCHED_EXPANDED_SRC`, `GEMV_MQ4G128V2_SRC`, `GEMV_MQ4G128_SRC`, `GEMV_MQ4G256V2_LUT_SRC`, `GEMV_MQ4G256V2_MOE_DOWN_K8_INDEXED_BATCHED_EXPANDED_SRC`, `GEMV_MQ4G256V2_MOE_GATE_UP_K8_INDEXED_BATCHED_SRC`, `GEMV_MQ4G256V2_MOE_GATE_UP_K8_INDEXED_K2048_NOLDS_GFX1100_SRC`, `GEMV_MQ4G256V2_MOE_GATE_UP_K8_INDEXED_SRC`, `GEMV_MQ4G256V2_MOE_GATE_UP_TOP10_INDEXED_BATCHED_SRC`, `GEMV_MQ4G256V2_MOE_NINEPATH_D4_SRC`, `GEMV_MQ4G256V2_MOE_NINEPATH_RPB8_GFX1100_SRC`, `GEMV_MQ4G256V2_MULTIROW_SRC`, `GEMV_MQ4G256V2_RESIDUAL_LUT_SRC`, `GEMV_MQ4G256V2_RESIDUAL_R1_K4096_GFX1100_NOSCRATCH_SRC`, `GEMV_MQ4G256V2_RESIDUAL_ROW_SERIAL_GFX1151_SRC`, `GEMV_MQ4G256V2_RESIDUAL_SIGMOID_SCALED_K512_SRC`, `GEMV_MQ4G256V2_RESIDUAL_SRC`, `GEMV_MQ4G256V2_SRC`, `GEMV_MQ4G256V2_XBATCH_RESIDUAL_SRC`, `GEMV_MQ4G256V2_XBATCH_SRC`, `GEMV_MQ4G256_LLOYD_GFX1100_SRC`, `GEMV_MQ4G256_LLOYD_MULTIACC_DIAG_GFX1100_SRC`, `GEMV_MQ4G256_LLOYD_RESIDUAL_GFX1100_SRC`, `GEMV_MQ4G256_LLOYD_RESIDUAL_SRC`, `GEMV_MQ4G256_LLOYD_SRC`, `GEMV_MQ4G256_MOE_GATE_UP_INDEXED_K2816_SRC`, `GEMV_MQ4G256_SRC`, `GEMV_MQ5G256V2_MULTIROW_SRC`, `GEMV_MQ5G256V2_RESIDUAL_SRC`, `GEMV_MQ5G256V2_SRC`, `GEMV_MQ5G256_SRC`, `GEMV_MQ6G256V2_MOE_DOWN_K8_INDEXED_BATCHED_EXPANDED_SRC`, `GEMV_MQ6G256V2_MOE_GATE_UP_K8_INDEXED_BATCHED_SRC`, `GEMV_MQ6G256V2_MOE_GATE_UP_K8_INDEXED_SRC`, `GEMV_MQ6G256V2_MOE_NINEPATH_D4_SRC`, `GEMV_MQ6G256V2_MULTIROW_SRC`, `GEMV_MQ6G256V2_RESIDUAL_SRC`, `GEMV_MQ6G256V2_SRC`, `GEMV_MQ6G256V2_XBATCH_RESIDUAL_SRC`, `GEMV_MQ6G256V2_XBATCH_SRC`, `GEMV_MQ6G256_SRC`, `GEMV_MQ8G256_SRC`, `GEMV_PARO4G128_SRC`, `GEMV_PARO_Q4G128_MOE_DOWN_K8_INDEXED_BATCHED_SRC`, `GEMV_PARO_Q4G128_MOE_GATE_UP_INDEXED_SRC`, `GEMV_PARO_Q4G128_MOE_GATE_UP_K8_INDEXED_BATCHED_SRC`, `GEMV_Q4AS8_SRC`, `GEMV_Q4F16_G32_SRC`, `GEMV_Q4F16_G64_SRC`, `GEMV_Q4F16_G64_WIDE_SRC`, `GEMV_Q4K_SRC`, `GEMV_Q4LUT_SRC`, `GEMV_Q4WAVE_SRC`, `GEMV_Q6K_SRC`, `GEMV_Q8HFQ_SRC`, `GEMV_Q8HFQ_WIDE_SRC`, `GEMV_Q8_0_MOE_DOWN_RESIDUAL_SCALED_K8_INDEXED_SRC`, `GEMV_Q8_0_MOE_DOWN_TOP10_INDEXED_BATCHED_EXPANDED_SRC`, `GEMV_Q8_0_MOE_GATE_UP_K8_INDEXED_SRC`, `GEMV_Q8_0_SRC`, `GEMV_Q8_0_WIDE_SRC`, `GEMV_SRC`, `GEMV_TQ2G128_SRC`, `GEMV_TQ2G128_XBATCH_SRC`, `GIVENS_COMMON_SRC`, `GIVENS_ROTATE_SRC`, `GIVENS_ROTATE_TO_SRC`, `GREEDY_ACCEPT_SRC`, `HASH_ROUTER_NORMALIZE_BATCHED_SRC`, `HASH_ROUTER_NORMALIZE_BUF_SRC`, `HASH_ROUTER_NORMALIZE_SRC`, `HC_APPLY_ALPHA_BATCHED_SRC`, `HC_APPLY_ALPHA_SRC`, `HC_COMPUTE_CONTROL_BATCHED_FUSED24_GFX1201_SRC`, `HC_COMPUTE_CONTROL_BATCHED_SRC`, `HC_COMPUTE_CONTROL_SRC`, `HC_COMPUTE_CONTROL_T1024_SRC`, `HC_COMPUTE_CONTROL_WMMA_B1_GFX1201_SRC`, `HC_COMPUTE_CONTROL_WMMA_B2_GFX1201_SRC`, `HC_COMPUTE_CONTROL_WMMA_B4_GFX1201_SRC`, `HC_FINALIZE_CONTROL_SRC`, `HC_FINALIZE_INPUT_MAP_SRC`, `HC_HEAD_COMPUTE_PRE_SRC`, `HC_INPUT_MAP_BATCHED_SRC`, `HC_INPUT_MAP_SRC`, `HC_INV_RMS_BATCHED_GFX1201_SRC`, `HC_MIX_4STREAM_BATCHED_SRC`, `HC_MIX_4STREAM_PEER4_GFX1201_SRC`, `HC_MIX_4STREAM_SRC`, `HC_PRE_POST_SIGMOID_SCALE_SRC`, `HC_SINKHORN_4X4_BATCHED_SRC`, `HC_SINKHORN_4X4_SRC`, `HC_SPLIT_FINALIZE_BATCHED_SRC`, `HC_STREAMS_INIT_FROM_EMBED_BATCHED_SRC`, `HFQ4G256_DEQUANTIZE_TO_F16_SRC`, `INDEXER_COMPRESSED_K_SCORE_SRC`, `INDEXER_KV_GATHER_SRC`, `INDEXER_RELU_SCORE_BATCHED_SRC`, `INDEXER_RELU_SCORE_BUF_SHARDED_GFX1201_SRC`, `INDEXER_RELU_SCORE_BUF_SRC`, `INDEXER_RELU_SCORE_WMMA_BATCHED_GFX12_SRC`, `INDEXER_RELU_SCORE_WMMA_BATCHED_KLDS_GFX1151_SRC`, `INDEXER_RELU_SCORE_WMMA_BATCHED_SHARDED_GFX1201_SRC`, `INDEXER_RELU_SCORE_WMMA_BATCHED_SRC`, `INDEXER_TOPK_TWOSTAGE_SRC`, `INDEXER_TOP_K_BATCHED_BOUNDED_GFX1151_SRC`, `INDEXER_TOP_K_BATCHED_BOUNDED_GFX1201_SRC`, `INDEXER_TOP_K_BATCHED_SRC`, `INDEXER_TOP_K_BUF_BOUNDED_GFX1151_SRC`, `INDEXER_TOP_K_BUF_BOUNDED_GFX942_SRC`, `INDEXER_TOP_K_BUF_PARALLEL_GFX1151_SRC`, `INDEXER_TOP_K_BUF_SRC`, `INDEXER_TOP_K_BUF_UNROLLED_GFX1151_SRC`, `INDEXER_TOP_K_SRC`, `KV_CACHE_WRITE_ASYM3_Q8_PAIR_GFX1100_SRC`, `KV_CACHE_WRITE_ASYM_K_FWHT2_BATCHED_SRC`, `KV_CACHE_WRITE_ASYM_K_FWHT2_SRC`, `KV_CACHE_WRITE_ASYM_K_FWHT3_BATCHED_SRC`, `KV_CACHE_WRITE_ASYM_K_FWHT3_SRC`, `KV_CACHE_WRITE_ASYM_K_FWHT4_BATCHED_SRC`, `KV_CACHE_WRITE_ASYM_K_FWHT4_SRC`, `KV_CACHE_WRITE_ASYM_K_GIVENS2_BATCHED_SRC`, `KV_CACHE_WRITE_ASYM_K_GIVENS2_SRC`, `KV_CACHE_WRITE_ASYM_K_GIVENS3_BATCHED_SRC`, `KV_CACHE_WRITE_ASYM_K_GIVENS3_HD512_BATCHED_SRC`, `KV_CACHE_WRITE_ASYM_K_GIVENS3_HD512_SRC`, `KV_CACHE_WRITE_ASYM_K_GIVENS3_SRC`, `KV_CACHE_WRITE_ASYM_K_GIVENS4_BATCHED_SRC`, `KV_CACHE_WRITE_ASYM_K_GIVENS4_SRC`, `KV_CACHE_WRITE_BF16_SRC`, `KV_CACHE_WRITE_F16_SRC`, `KV_CACHE_WRITE_F32_BATCHED_SRC`, `KV_CACHE_WRITE_FP8_E4M3_BATCHED_SRC`, `KV_CACHE_WRITE_FP8_E4M3_SRC`, `KV_CACHE_WRITE_FWHT256_2BIT_BATCHED_SRC`, `KV_CACHE_WRITE_FWHT256_2BIT_SRC`, `KV_CACHE_WRITE_FWHT256_4BIT_BATCHED_SRC`, `KV_CACHE_WRITE_FWHT256_4BIT_SRC`, `KV_CACHE_WRITE_FWHT3_HD512_BATCHED_SRC`, `KV_CACHE_WRITE_FWHT3_HD512_SRC`, `KV_CACHE_WRITE_HFQ4_SRC`, `KV_CACHE_WRITE_HFQ8_SRC`, `KV_CACHE_WRITE_INT8C_F16_SRC`, `KV_CACHE_WRITE_INT8C_SRC`, `KV_CACHE_WRITE_INT8_SRC`, `KV_CACHE_WRITE_Q4_SRC`, `KV_CACHE_WRITE_Q8_0_BATCHED_SRC`, `KV_CACHE_WRITE_Q8_0_PAIR_GFX1100_SRC`, `KV_CACHE_WRITE_Q8_0_SRC`, `KV_CACHE_WRITE_Q8_SRC`, `KV_CACHE_WRITE_SRC`, `KV_COMPACT_GATHER_SRC`, `KV_FOLD_ASYM2_SRC`, `KV_FOLD_ASYM3_SRC`, `KV_FOLD_ASYM4_SRC`, `KV_FOLD_Q8_SRC`, `KV_SLOT_DESC_H`, `KV_SLOT_DESC_PAGED_H`, `KV_TRANSCODE_K_FWHT4_TO_FWHT2_SRC`, `KV_TRANSCODE_K_FWHT4_TO_FWHT3_SRC`, `KV_TRANSCODE_V_LLOYD_DOWN_SRC`, `KV_TRANSCODE_V_Q8_TO_LLOYD4_SRC`, `L2_NORM_SRC`, `LAYERNORM_MODULATE_F32_SRC`, `LAYERNORM_SRC`, `LOGIT_PENALTY_TABLE_SRC`, `LOGIT_SOFTCAP_SRC`, `MAX_PROB_SRC`, `MODULATE_F32_SRC`, `MOE_DOWN_COMBINE_GROUPED_K8_SRC`, `MOE_DOWN_COMBINE_GROUPED_TOP10_SRC`, `MOE_DOWN_COMBINE_K8_BATCHED_SRC`, `MOE_DOWN_COMBINE_K8_BATCHED_VEC4_SRC`, `MOE_DOWN_COMBINE_RMSNORM_MQ_ROTATE_VECSUM_GFX1100_SRC`, `MOE_DOWN_COMBINE_RMSNORM_MQ_ROTATE_VECSUM_GFX1151_SRC`, `MOE_DOWN_COMBINE_TOP10_BATCHED_SRC`, `MOE_DOWN_UNSCATTER_K8_SRC`, `MOE_GATE_UP_UNSCATTER_K8_SRC`, `MOE_GATE_UP_UNSCATTER_SILU_TOP10_SRC`, `MOE_GATE_UP_UNSCATTER_TOP10_SRC`, `MOE_ROUTER_SOFTMAX_TOP10_F32_FAST_SRC`, `MOE_ROUTER_SOFTMAX_TOP10_F32_SRC`, `MOE_ROUTER_SOFTMAX_TOPK_K8_WAVE64_EXACT_SHARED_SILU_MQ_ROTATE_SRC`, `MOE_ROUTER_SOFTMAX_TOPK_K8_WAVE64_EXACT_SRC`, `MOE_ROUTER_SOFTMAX_TOPK_K8_WAVE64_SRC`, `MOE_SCATTER_FUSED_K8_SRC`, `MOE_SCATTER_FUSED_TOP10_SRC`, `MOE_SCATTER_HISTOGRAM_K8_SRC`, `MOE_SCATTER_OFFSETS_K8_SRC`, `MOE_SCATTER_PERMUTE_K8_SRC`, `MOE_SOFTMAX_TOPK_K8_BATCHED_SRC`, `MOE_SOFTMAX_TOPK_K8_SRC`, `MOE_TOPK_RENORM_K8_BATCHED_SRC`, `MOE_TOPK_RENORM_K8_SRC`, `MOE_TOPK_RENORM_TOP10_BATCHED_SRC`, `MOE_UNSCATTER_SILU_CLAMP_K8_SRC`, `MQ4V2_FP8_FRAGMENT_REPACK_GFX1201`, `MQ4_FP8_VERIFY_PM_GFX1201`, `MQ4_ROWS_X4_PM_GFX1151`, `MQ4_ROWS_X4_PM_GFX1201`, `MQ4_VERIFY_PM_GFX1201`, `MQ_ROTATE_X_128_V2_SRC`, `MQ_ROTATE_X_AWQ_FP8_GFX12_SRC`, `MQ_ROTATE_X_AWQ_FP8_INREG_GFX12_SRC`, `MQ_ROTATE_X_AWQ_I4_GFX12_SRC`, `MQ_ROTATE_X_AWQ_I4_SRC`, `MQ_ROTATE_X_DUAL_FP8_GFX12_SRC`, `MQ_ROTATE_X_FP8_GFX12_SRC`, `MQ_ROTATE_X_FP8_INREG_GFX12_SRC`, `MQ_ROTATE_X_GFX942_SRC`, `MQ_ROTATE_X_I4_GFX12_SRC`, `MQ_ROTATE_X_I4_SRC`, `MUL_SRC`, `PACK_F16_TO_FP8_MQ4V2_GFX12_SRC`, `PACK_F32_TO_FP8_GFX12_SRC`, `PACK_F32_TO_FP8_MQ4V2_GFX12_SRC`, `PACK_FP8_VERIFY_ROWS_GFX1201_SRC`, `PFLASH_SCORE_FWHT2_KV_SRC`, `PFLASH_SCORE_FWHT3_KV_SRC`, `PFLASH_SCORE_FWHT4_KV_SRC`, `PFLASH_SCORE_Q8_KV_SRC`, `QKV_SPLIT_INTERLEAVED_SRC`, `QK_RMSNORM_ROPE_FLUX_SRC`, `QSA_GATHER_PM_GFX1151`, `QSA_GATHER_PM_GFX1201`, `QSA_SELECT_BF16_PM_GFX1151`, `QSA_SELECT_PM_GFX1151`, `QUICK_GELU_F32_SRC`, `QWEN35_FA_PREP_GFX1100_SRC`, `QWEN35_FA_PREP_GFX1151_SRC`, `QWEN35_FA_PREP_GFX1201_SRC`, `QWEN4_BF16_ROUND_TRIP_SRC`, `QWEN4_GEMM_F16_X_F16_WMMA_LDS256_SRC`, `QWEN4_GEMM_MQ4G256V2_MOE_GROUPED_WMMA_K2_SRC`, `QWEN4_GEMM_MQ6G256V2_WMMA_GFX12_X4_SRC`, `QWEN4_GEMM_MQV2_WMMA_GFX11_BT_SRC`, `QWEN4_GEMV_BF16_XF32_SRC`, `QWEN4_GEMV_MQ4G256V2_MOE_GATE_UP_K8_INDEXED_BATCHED_SRC`, `QWEN4_GEMV_MQ4G256_SRC`, `QWEN4_GEMV_MQ6G256V2_SRC`, `QWEN4_GEMV_Q8_0_SRC`, `QWEN4_GEMV_Q8_0_WIDE_SRC`, `QWEN4_HC_STREAMS_INIT_FROM_EMBED_BATCHED_SRC`, `QWEN4_LANE_ROWS_PM_GFX1151`, `QWEN4_LANE_ROWS_PM_GFX1151_SHA256`, `QWEN4_MOE_IU4_SYM_GFX1151_SRC`, `QWEN4_MOE_IU4_SYM_PM_GFX1151`, `QWEN4_MOE_IU4_SYM_PM_GFX1201`, `QWEN4_MOE_ROTATE128_I4_SRC`, `QWEN4_MOE_SCATTER_FUSED_K8_SRC`, `QWEN4_MOE_SCATTER_STABLE_TOP10_SRC`, `QWEN4_MOE_SYM_CHECK_GFX1201_SRC`, `QWEN4_MQ4_X4_PM_GFX1151`, `QWEN4_MQ6_X4_PM_GFX1151`, `QWEN4_MQ6_X4_PM_GFX1201`, `QWEN4_SHARED_GU_PM_GFX1151`, `QWEN4_SILU_MUL_SRC`, `REPEAT_INTERLEAVE_QK_BATCHED_SRC`, `REPEAT_INTERLEAVE_QK_SRC`, `REQUANT_G256_SRC`, `RMSNORM_AT_SLOT_BUF_SRC`, `RMSNORM_REDUCE_GFX1100_SRC`, `RMSNORM_REDUCE_GFX942_SRC`, `RMSNORM_RESIDUAL_ADD_SRC`, `RMSNORM_ROWSPLIT_SRC`, `RMSNORM_SRC`, `ROPE_2D_FLUX_F32_FAST_SRC`, `ROPE_2D_FLUX_F32_SRC`, `ROPE_2D_HALFSPLIT_QKV_INTERLEAVED_SRC`, `ROPE_2D_HALFSPLIT_SRC`, `ROPE_BATCHED_SRC`, `ROPE_MROPE_HALFSPLIT_BATCHED_SRC`, `ROPE_MROPE_HALFSPLIT_SRC`, `ROPE_PARTIAL_HALFSPLIT_BATCHED_SRC`, `ROPE_PARTIAL_HALFSPLIT_HEADGRID_SRC`, `ROPE_PARTIAL_HALFSPLIT_SRC`, `ROPE_PARTIAL_HALVED_BATCHED_SRC`, `ROPE_PARTIAL_HALVED_SRC`, `ROPE_PARTIAL_INTERLEAVED_BATCHED_SRC`, `ROPE_PARTIAL_INTERLEAVED_SRC`, `ROPE_SRC`, `ROPE_TAIL_INTERLEAVED_BATCHED_SRC`, `ROPE_TAIL_INTERLEAVED_H64D128R64_GFX1201_SRC`, `ROPE_TAIL_INTERLEAVED_SRC`, `ROPE_TAIL_YARN_INTERLEAVED_AT_SLOT_BUF_SRC`, `ROPE_TAIL_YARN_INTERLEAVED_BATCHED_SRC`, `ROPE_TAIL_YARN_INTERLEAVED_SRC`, `ROTATE_WITH_RMS_GFX1100_SRC`, `ROTATE_WITH_RMS_GFX942_SRC`, `ROTATE_X_MQ_AWQ_SRC`, `ROUND_H_BF16_QUALITY_GFX12_SRC`, `SAMPLE_ROWS_PF_SRC`, `SAMPLE_TOP_P_PARALLEL_SRC`, `SAMPLE_TOP_P_SRC`, `SCALED_ADD_INPLACE_SRC`, `SCALE_F32_SRC`, `SIGMOID_MUL_MQ_ROTATE_X_AWQ_FP8_GFX12_SRC`, `SIGMOID_MUL_MQ_ROTATE_X_AWQ_FP8_GIL_GFX12_SRC`, `SIGMOID_MUL_MQ_ROTATE_X_AWQ_FP8_INREG_GFX12_SRC`, `SIGMOID_MUL_MQ_ROTATE_X_AWQ_FP8_INREG_GIL_GFX12_SRC`, `SIGMOID_MUL_MQ_ROTATE_X_AWQ_I4_GFX11_SRC`, `SIGMOID_MUL_MQ_ROTATE_X_AWQ_I4_GFX12_SLAB_SRC`, `SIGMOID_MUL_MQ_ROTATE_X_AWQ_I4_GFX12_SRC`, `SIGMOID_MUL_MQ_ROTATE_X_AWQ_I4_GIL_GFX12_SLAB_SRC`, `SIGMOID_MUL_MQ_ROTATE_X_AWQ_I4_GIL_GFX12_SRC`, `SIGMOID_MUL_MQ_ROTATE_X_FP8_GFX12_SRC`, `SIGMOID_MUL_MQ_ROTATE_X_FP8_GIL_GFX12_SRC`, `SIGMOID_MUL_MQ_ROTATE_X_FP8_INREG_GFX12_SRC`, `SIGMOID_MUL_MQ_ROTATE_X_FP8_INREG_GIL_GFX12_SRC`, `SIGMOID_MUL_SRC`, `SIGMOID_SCALED_RESIDUAL_ADD_BATCHED_SRC`, `SIGMOID_SRC`, `SILU_MUL_SRC`, `SILU_SRC`, `SOFTMAX_PROB_GATHER_BATCHED_SRC`, `SOFTMAX_SRC`, `SOFTMAX_TEMP_BATCHED_SRC`, `SOFTPLUS_SRC`, `SPLIT_MQ4V2_Z_BETAALPHA_SRC`, `SQRT_SOFTPLUS_F32_SRC`, `STATE_OVERLAP_SHIFT_F32_BUF_SRC`, `STATE_RING_WRITE_F32_BUF_SRC`, `SWA_RING_WRITE_BATCHED_SRC`, `SWA_RING_WRITE_BUF_SRC`, `SWA_VISIBILITY_STAGE_BATCHED_SRC`, `TOPK8_RESCORE_MQ6G256V2_SRC`, `TOPK_LOGITS_SRC`, `TOPK_LOGSUMEXP_BATCHED_SRC`, `TP_GRAPH_SIGNAL_GFX1201_SRC`, `TRANSPOSE_SRC`, `TRIATTN_ACCUMULATE_SRC`, `TRIATTN_SCORE_ASYM2_SRC`, `TRIATTN_SCORE_ASYM3_SRC`, `TRIATTN_SCORE_ASYM4_SRC`, `TRIATTN_SCORE_FWHT_SRC`, `TRIATTN_SCORE_Q8_SRC`, `TURBO_COMMON_H`, `V4F_ATTN_POS0_SRC`, `V4F_ATTN_SWA_BATCHED_DEBUG_SRC`, `V4F_ATTN_SWA_BATCHED_SRC`, `V4F_ATTN_SWA_BUF_SRC`, `V4F_ATTN_SWA_SRC`, `V4F_ATTN_SWA_TOPK_BATCHED_SRC`, `V4F_ATTN_SWA_TOPK_BATCHED_WMMA_GFX12_SRC`, `V4F_ATTN_SWA_TOPK_BATCHED_WMMA_SRC`, `V4F_ATTN_SWA_TOPK_BUF_LARGE_SERIAL_GFX1151_SRC`, `V4F_ATTN_SWA_TOPK_BUF_SRC`, `V4F_ATTN_SWA_TOPK_BUF_XLANE_GFX1151_SRC`, `V4F_ATTN_SWA_TOPK_DIRECT_BATCHED_SRC`, `V4F_ATTN_SWA_TOPK_DIRECT_WMMA_SRC`, `V4F_CONVERT_F32_TO_F16_SRC`, `V4F_FUSED_SILU_MUL_CLAMP_MQ_ROTATE_SRC`, `V4F_MOE_TOPK_BIAS_AWARE_BATCHED_SRC`, `V4F_MOE_TOPK_BIAS_AWARE_SRC`, `V4F_SILU_MUL_CLAMP_SRC`, `V4F_TOPK_KV_GATHER_BATCHED_SRC`, `V4F_TOPK_KV_GATHER_BATCHED_TILED_GFX1201_SRC`, `V4F_TOPK_KV_GATHER_BATCHED_TILED_SHARDED_GFX1201_SRC`, `V4F_TOPK_KV_GATHER_BUF_SHARDED_GFX1201_SRC`, `V4F_TOPK_KV_GATHER_BUF_SRC`, `V4F_TOPK_KV_GATHER_IDENTITY_BATCHED_SHARDED_GFX1201_SRC`, `V4F_TOPK_KV_GATHER_IDENTITY_BATCHED_SRC`, `V4F_TOPK_KV_GATHER_IDENTITY_BUF_SHARDED_GFX1201_SRC`, `V4F_TOPK_KV_GATHER_IDENTITY_BUF_SRC`, `V4F_TOPK_KV_GATHER_TILED_GFX1151_SRC`, `VAE_ATTN_SRC`, `VAE_CONV1X1_SRC`, `VAE_CONV3X3_S2_SRC`, `VAE_CONV3X3_SRC`, `VAE_GROUPNORM_SRC`, `VAE_IM2COL_SRC`, `VAE_LAYOUT_SRC`, `VAE_UPSAMPLE2X_SRC`, `VIT_ATTENTION_OPT_SRC`, `VIT_ATTENTION_QTILED_SRC`, `VIT_ATTENTION_SRC`, `WO_PER_GROUP_BATCHED_F32_SRC`, `WO_PER_GROUP_BATCHED_HFQ4G256_SRC`, `WO_PER_GROUP_BATCHED_Q8_0_MULTIROW_SRC`, `WO_PER_GROUP_BATCHED_Q8_0_SRC`, `WO_PER_GROUP_BATCHED_Q8_0_WMMA_4W_SRC`, `ZERO_F32_SRC`, `ZERO_INACTIVE_ROWS_F32_SRC`, `fused_gate_up_mq3g256_lloyd_for_arch`, `fused_gate_up_mq4g256_lloyd_for_arch`, `fused_gate_up_mq4g256v2_k5120_gfx1100_src`, `fused_qkv_mq3g256_lloyd_for_arch`, `fused_qkv_mq3g256_lloyd_qwen2_bias_for_arch`, `fused_qkv_mq4g256_lloyd_for_arch`, `fused_qkv_mq4g256_lloyd_qwen2_bias_for_arch`, `fused_qkvza_mq3g256_lloyd_for_arch`, `fused_qkvza_mq4g256_lloyd_for_arch`, `gated_norm_mq_rotate_awq_k6144_gfx1100_src`, `gated_norm_mq_rotate_awq_k6144_gfx1201_src`, `gated_norm_mq_rotate_k6144_gfx1100_src`, `gated_norm_mq_rotate_k6144_gfx1201_src`, `gemm_f16_x_f16_wmma_lds256_src`, `gemm_gate_up_mq3g256_lloyd_wmma_for_arch`, `gemm_gate_up_mq3g256_lloyd_wmma_mb4_for_arch`, `gemm_gate_up_mq3g256_lloyd_wmma_mb4_nosync_for_arch`, `gemm_gate_up_mq3g256_lloyd_wmma_nosync_for_arch`, `gemm_gate_up_mq4g256_lloyd_wmma_for_arch`, `gemm_gate_up_mq4g256_lloyd_wmma_mb4_for_arch`, `gemm_gate_up_mq4g256_lloyd_wmma_mb4_nosync_for_arch`, `gemm_gate_up_mq4g256_lloyd_wmma_nosync_for_arch`, `gemm_mq3g256_lloyd_residual_wmma_for_arch`, `gemm_mq3g256_lloyd_residual_wmma_mb4_for_arch`, `gemm_mq4g256_lloyd_residual_wmma_for_arch`, `gemm_mq4g256_lloyd_residual_wmma_mb2_for_arch`, `gemm_mq4g256_lloyd_residual_wmma_mb4_for_arch`, `gemm_mqv2_wmma_gfx11_bt_src`, `gemm_qkv_mq3g256_lloyd_wmma_for_arch`, `gemm_qkv_mq3g256_lloyd_wmma_mb4_for_arch`, `gemm_qkv_mq4g256_lloyd_wmma_for_arch`, `gemm_qkv_mq4g256_lloyd_wmma_mb4_for_arch`, `gemm_qkvza_mq3g256_lloyd_wmma_for_arch`, `gemm_qkvza_mq3g256_lloyd_wmma_mb4_for_arch`, `gemm_qkvza_mq4g256_lloyd_wmma_for_arch`, `gemm_qkvza_mq4g256_lloyd_wmma_mb4_for_arch`, `gemv_bf16_xf32_src`, `gemv_hfp4g32_for_arch`, `gemv_hfq3g256_for_arch`, `gemv_hfq3g256_residual_for_arch`, `gemv_hfq4g256_for_arch`, `gemv_hfq4g256_residual_for_arch`, `gemv_mq3g256_lloyd_for_arch`, `gemv_mq3g256_lloyd_residual_for_arch`, `gemv_mq3g256v2_for_arch`, `gemv_mq3g256v2_residual_for_arch`, `gemv_mq4g256_lloyd_for_arch`, `gemv_mq4g256_lloyd_residual_for_arch`, `gemv_mq4g256v2_for_arch`, `gemv_mq4g256v2_moe_gate_up_k8_indexed_batched_src`, `gemv_mq4g256v2_residual_for_arch`, `gemv_mq5g256v2_for_arch`, `gemv_mq5g256v2_residual_for_arch`, `gemv_mq6g256v2_for_arch`, `gemv_mq6g256v2_residual_for_arch`, `kv_slot_desc_paged_source`, `kv_slot_desc_source`, `kv_slot_givens4_paged_source`, `moe_scatter_fused_k8_src`, `mq2g256_lloyd_moe_grouped_wmma_source`, `mq3g256_lloyd_moe_grouped_wmma_source`, `mq4g256v2_moe_grouped_wmma_source`, `mq6g256v2_moe_grouped_wmma_source`, `qwen36_27b_fa_prep_gfx1100_src`, `qwen36_27b_fa_prep_gfx1201_src`, `scheduler_profile_for_module` +- [`src/kernels.rs`](src/kernels.rs): `ACT_QUANT_FP4_UE8M0_G32_INPLACE_GFX942_SRC`, `ACT_QUANT_FP8_UE8M0_INPLACE_GFX942_SRC`, `ADD_INPLACE_SRC`, `ADD_PEER_RESIDUAL_SRC`, `ADD_SRC`, `ALPHA_GATE_SRC`, `APPLY_ROPE_2D_VISION_SRC`, `ARGMAX_BATCHED_SRC`, `ARGMAX_SRC`, `ARGMAX_TOKEN_CHAIN_SRC`, `ATTENTION_BF16_KV_BATCHED_SRC`, `ATTENTION_BF16_KV_SRC`, `ATTENTION_CAUSAL_BATCHED_SRC`, `ATTENTION_DECODE_BATCHED_HISTORY_SRC`, `ATTENTION_DFLASH_SLIDING_SRC`, `ATTENTION_DFLASH_SRC`, `ATTENTION_DFLASH_WMMA_GFX12_SRC`, `ATTENTION_DFLASH_WMMA_M128_N32_F16KV_V7B_SRC`, `ATTENTION_DFLASH_WMMA_M128_N32_F16KV_V7_SRC`, `ATTENTION_DFLASH_WMMA_M32_GFX12_SRC`, `ATTENTION_DFLASH_WMMA_M32_SRC`, `ATTENTION_DFLASH_WMMA_M64_N128_F16KV_SRC`, `ATTENTION_DFLASH_WMMA_M64_N128_F16KV_V2_SRC`, `ATTENTION_DFLASH_WMMA_M64_N128_F16KV_V3_CAUSAL_GFX12_SRC`, `ATTENTION_DFLASH_WMMA_M64_N128_F16KV_V3_CAUSAL_SRC`, `ATTENTION_DFLASH_WMMA_M64_N128_F16KV_V3_SRC`, `ATTENTION_DFLASH_WMMA_M64_N128_F16KV_V4_CAUSAL_SRC`, `ATTENTION_DFLASH_WMMA_M64_N128_F16KV_V4_SRC`, `ATTENTION_DFLASH_WMMA_M64_N32_F16KV_V5_GFX12_SRC`, `ATTENTION_DFLASH_WMMA_M64_N32_F16KV_V5_SRC`, `ATTENTION_DFLASH_WMMA_M64_N32_F16KV_V6_SRC`, `ATTENTION_DFLASH_WMMA_N128_F16KV_SRC`, `ATTENTION_DFLASH_WMMA_N64_F16KV_SRC`, `ATTENTION_DFLASH_WMMA_N64_SRC`, `ATTENTION_DFLASH_WMMA_SRC`, `ATTENTION_FLASH_ASYM2_TILE_BATCHED_SRC`, `ATTENTION_FLASH_ASYM2_TILE_SRC`, `ATTENTION_FLASH_ASYM3_TILE_BATCHED_SRC`, `ATTENTION_FLASH_ASYM3_TILE_HD512_BATCHED_SRC`, `ATTENTION_FLASH_ASYM3_TILE_HD512_SRC`, `ATTENTION_FLASH_ASYM3_TILE_SRC`, `ATTENTION_FLASH_ASYM4_TILE_BATCHED_SRC`, `ATTENTION_FLASH_ASYM4_TILE_SRC`, `ATTENTION_FLASH_ASYM4_WMMA_TILE_BATCHED_GFX12_SRC`, `ATTENTION_FLASH_ASYM4_WMMA_TILE_BATCHED_SRC`, `ATTENTION_FLASH_ASYM_REDUCE_BATCHED_SRC`, `ATTENTION_FLASH_BF16_TILE_BATCHED_SRC`, `ATTENTION_FLASH_BF16_TILE_SRC`, `ATTENTION_FLASH_F16_TILE_BATCHED_SRC`, `ATTENTION_FLASH_F16_TILE_SRC`, `ATTENTION_FLASH_FP8_E4M3_TILE_BATCHED_SRC`, `ATTENTION_FLASH_FP8_E4M3_TILE_GQA_GFX1201_SRC`, `ATTENTION_FLASH_FP8_E4M3_TILE_SRC`, `ATTENTION_FLASH_FWHT2_TILE_BATCHED_SRC`, `ATTENTION_FLASH_FWHT2_TILE_SRC`, `ATTENTION_FLASH_FWHT3_TILE_BATCHED_SRC`, `ATTENTION_FLASH_FWHT3_TILE_HD512_BATCHED_SRC`, `ATTENTION_FLASH_FWHT3_TILE_HD512_SRC`, `ATTENTION_FLASH_FWHT3_TILE_SRC`, `ATTENTION_FLASH_FWHT4_TILE_BATCHED_SRC`, `ATTENTION_FLASH_FWHT4_TILE_SRC`, `ATTENTION_FLASH_GQA_FUSED_SRC`, `ATTENTION_FLASH_GQA_SRC`, `ATTENTION_FLASH_LLOYD_REDUCE_BATCHED_SRC`, `ATTENTION_FLASH_LLOYD_REDUCE_SRC`, `ATTENTION_FLASH_Q8_0_REDUCE_GATED_MQ_ROTATE_AWQ_DEC_GFX1100_SRC`, `ATTENTION_FLASH_Q8_0_REDUCE_GATED_MQ_ROTATE_AWQ_GFX1100_SRC`, `ATTENTION_FLASH_Q8_0_REDUCE_GATED_MQ_ROTATE_AWQ_GFX1151_SRC`, `ATTENTION_FLASH_Q8_0_REDUCE_GATED_MQ_ROTATE_AWQ_GFX1201_SRC`, `ATTENTION_FLASH_Q8_0_REDUCE_GATED_MQ_ROTATE_GFX1100_SRC`, `ATTENTION_FLASH_Q8_0_REDUCE_GATED_MQ_ROTATE_GFX1151_SRC`, `ATTENTION_FLASH_Q8_0_REDUCE_GATED_MQ_ROTATE_GFX1201_SRC`, `ATTENTION_FLASH_Q8_0_REDUCE_SRC`, `ATTENTION_FLASH_Q8_0_TILE_BATCHED_SRC`, `ATTENTION_FLASH_Q8_0_TILE_DPP_GFX1151_SRC`, `ATTENTION_FLASH_Q8_0_TILE_GQA_GEMMA_GFX1151_SRC`, `ATTENTION_FLASH_Q8_0_TILE_GQA_GFX1100_SRC`, `ATTENTION_FLASH_Q8_0_TILE_GQA_GFX1151_SRC`, `ATTENTION_FLASH_Q8_0_TILE_ROWS_SRC`, `ATTENTION_FLASH_Q8_0_TILE_SRC`, `ATTENTION_FLASH_REDUCE_DSPLIT_GFX1151_SRC`, `ATTENTION_FLASH_REDUCE_DSPLIT_GFX1201_SRC`, `ATTENTION_FLASH_SRC`, `ATTENTION_FLUX_V2_WMMA_F16KV_SRC`, `ATTENTION_FLUX_VTK_WMMA_F16KV_F32_SRC`, `ATTENTION_FLUX_VT_WMMA_F16KV_F32_GFX12_SRC`, `ATTENTION_FLUX_VT_WMMA_F16KV_F32_SRC`, `ATTENTION_FP8_E4M3_FA2_GQA_F16_GFX1201_SRC`, `ATTENTION_FP8_E4M3_FA2_GQA_FP8_GFX1201_SRC`, `ATTENTION_FP8_E4M3_FA2_GQA_PACKET_GFX1201_SRC`, `ATTENTION_FP8_E4M3_FA2_GQA_QRESIDENT_GFX1201_SRC`, `ATTENTION_FP8_E4M3_FA2_GQA_QRESIDENT_V2_GFX1201_SRC`, `ATTENTION_FP8_E4M3_FA2_GQA_QRESIDENT_V2_Q8_A4EPI_GFX1201_SRC`, `ATTENTION_FP8_E4M3_FA2_GQA_QRESIDENT_V2_Q8_GFX1201_SRC`, `ATTENTION_FP8_E4M3_KV_BATCHED_SRC`, `ATTENTION_FP8_E4M3_KV_SRC`, `ATTENTION_GQA_WARP_DV_SRC`, `ATTENTION_GQA_WARP_SRC`, `ATTENTION_HFQ4_KV_SRC`, `ATTENTION_HFQ8_KV_SRC`, `ATTENTION_INT8C_F16_KV_SRC`, `ATTENTION_INT8C_KV_SRC`, `ATTENTION_INT8_KV_SRC`, `ATTENTION_PARTITION_MERGE_F32_SRC`, `ATTENTION_Q4KV_SRC`, `ATTENTION_Q8KV_SRC`, `ATTENTION_Q8_0_FA2_GQA_FP8_GFX1201_SRC`, `ATTENTION_Q8_0_FA2_GQA_FWHT3K_GFX11_SRC`, `ATTENTION_Q8_0_FA2_GQA_FWHT3K_GFX1201_SRC`, `ATTENTION_Q8_0_FA2_GQA_GFX1151_SRC`, `ATTENTION_Q8_0_FA2_GQA_GFX11_SRC`, `ATTENTION_Q8_0_FA2_GQA_GFX1201_SRC`, `ATTENTION_Q8_0_FA2_GQA_MERGE_FP8_GFX1201_SRC`, `ATTENTION_Q8_0_FA2_GQA_PARTIAL_FP8_GFX1201_SRC`, `ATTENTION_Q8_0_FLASH_PREFILL_SRC`, `ATTENTION_Q8_0_FLASH_PREFILL_WMMA_GFX12_SRC`, `ATTENTION_Q8_0_FLASH_PREFILL_WMMA_PARTIAL_GFX1201_SRC`, `ATTENTION_Q8_0_FLASH_PREFILL_WMMA_SRC`, `ATTENTION_Q8_0_FLASH_PREFILL_WMMA_SWA_GFX1100_SRC`, `ATTENTION_Q8_0_FLASH_PREFILL_WMMA_SWA_GFX1151_SRC`, `ATTENTION_Q8_0_FLASH_PREFILL_WMMA_SWA_GFX12_SRC`, `ATTENTION_Q8_0_FLASH_PREFILL_WMMA_SWA_SRC`, `ATTENTION_Q8_0_KV_BATCHED_SRC`, `ATTENTION_Q8_0_KV_SRC`, `ATTENTION_Q8_0_KV_SWA_SRC`, `ATTENTION_Q8_0_KV_TIMED_SRC`, `ATTENTION_Q8_0_PREFILL_GQA_WMMA_GFX11_SRC`, `ATTENTION_SRC`, `ATTENTION_T5_BIAS_F32_SRC`, `ATTENTION_VERIFY_GQA_GFX1100_SRC`, `ATTENTION_VERIFY_GQA_GFX1201_SRC`, `ATTENTION_VERIFY_WMMA_GFX1151_SRC`, `BATCHED_CATEGORICAL_SAMPLE_SRC`, `BF16_ROUND_TRIP_SRC`, `BIAS_ADD_SRC`, `BLOCK_I4_128_QUANT_SRC`, `CALIB_REDUCE_SRC`, `CAST_F32_TO_F16_SRC`, `CHAIN_ACCEPT_SPEC_SRC`, `COMPRESSOR_ADD_APE_BATCHED_SRC`, `COMPRESSOR_COMPRESS_ALIGNED_BATCHED_SRC`, `COMPRESSOR_OVERLAP_CONCAT_SRC`, `COMPRESSOR_RING_WRITE_BATCHED_SRC`, `COMPRESSOR_SOFTMAX_POOL_BUF_SRC`, `COMPRESSOR_SOFTMAX_POOL_SRC`, `CONV1D_DECODE_SRC`, `CONV1D_GATED_DECODE_SRC`, `CONV1D_SILU_SPLIT_QKNORM_B128_SRC`, `CONV1D_SILU_SPLIT_QKNORM_B256_SCALAR_PREP_SRC`, `CONV1D_SILU_SPLIT_QKNORM_B256_SRC`, `CONV1D_SILU_SPLIT_QKNORM_B32_SRC`, `CONV1D_SILU_SPLIT_QKNORM_B512_SRC`, `CONV1D_SILU_SPLIT_QKNORM_B64_SRC`, `CONV1D_SILU_SPLIT_SRC`, `CONV1D_SILU_SPLIT_TREE_SRC`, `CONV1D_SILU_SRC`, `COPY_F32_BUFFER_SRC`, `COPY_ROWS_STRIDED_F32_SRC`, `CROSS_ENTROPY_LOSS_SRC`, `DDTREE_BUILD_ATTN_MASK_SRC`, `DDTREE_GUMBEL_TOPK_BATCHED_SRC`, `DDTREE_SWOR_WALK_SRC`, `DEEPSEEK4_COMPRESSOR_CACHE_F16_SRC`, `DEEPSEEK4_COMPRESSOR_COMMIT_STAGED_F16_SRC`, `DEINTERLEAVE_BATCHED_SRC`, `DEINTERLEAVE_Q_RMSNORM_BATCHED_SRC`, `DEINTERLEAVE_SRC`, `DEQUANTIZE_MFP4G32_E8_SOA_TO_F16_GFX942_SRC`, `DEQUANTIZE_MFP4G32_E8_TO_F16_SRC`, `DEQUANTIZE_MFP4G32_LLOYD_TO_F16_SRC`, `DEQUANTIZE_MFP4G32_P_TO_F16_SRC`, `DEQUANT_BQ1G128_TO_F16_SRC`, `DEQUANT_FP4_E2M1_UE8M0_G32_TO_BF16_GFX942_SRC`, `DEQUANT_FP8_E4M3_UE8M0_BLK128_TO_BF16_GFX942_SRC`, `DEQUANT_HFQ4G256_TO_F16_SRC`, `DEQUANT_TQ2G128_TO_F16_SRC`, `DSPARK_SAMPLE_ACCEPT_LAZY_SRC`, `DSPARK_STAGE_KV_SRC`, `DYNAMIC_CAUSAL_CONV_F32_SRC`, `DYNAMIC_CONV_F32_SRC`, `EMBEDDING_BF16_BATCHED_SRC`, `EMBEDDING_F16_BATCHED_SRC`, `EMBEDDING_HFQ4G128_BATCHED_SRC`, `EMBEDDING_HFQ4G128_SRC`, `EMBEDDING_HFQ4G256_BATCHED_SRC`, `EMBEDDING_HFQ4G256_SRC`, `EMBEDDING_MQ4V2_BATCHED_SRC`, `EMBEDDING_Q4K_SRC`, `EMBEDDING_Q8_BATCHED_SRC`, `EMBEDDING_Q8_SRC`, `EMBEDDING_SCATTER_EXT_SRC`, `FN_GDN_DENSE_SRC`, `FUSED_GATE_UP_HFQ4G256_DOT_PREFETCH_GFX1100_SRC`, `FUSED_GATE_UP_HFQ4G256_DOT_REFORM_GFX1100_SRC`, `FUSED_GATE_UP_HFQ4G256_K1024_GFX1201_SRC`, `FUSED_GATE_UP_HFQ4G256_LANE0_HEADERS_GFX1100_SRC`, `FUSED_GATE_UP_HFQ4G256_PAIR2_GFX1100_SRC`, `FUSED_GATE_UP_HFQ4G256_PAIR_GFX1100_SRC`, `FUSED_GATE_UP_HFQ4G256_QUAD_PREFETCH_GFX1100_SRC`, `FUSED_GATE_UP_HFQ4G256_SETPRIO_GFX1100_SRC`, `FUSED_GATE_UP_HFQ4G256_SRC`, `FUSED_GATE_UP_HFQ4G256_STAGE_X32_GFX1100_SRC`, `FUSED_GATE_UP_HFQ4G256_V2_GFX942_SRC`, `FUSED_GATE_UP_HFQ4G256_WAVE64_DP4A_SRC`, `FUSED_GATE_UP_HFQ4G256_WAVE64_SRC`, `FUSED_GATE_UP_HFQ6G256_WAVE64_DP4A_SRC`, `FUSED_GATE_UP_MFP4G32_E8_GFX1151_SRC`, `FUSED_GATE_UP_MQ2G256V2_SRC`, `FUSED_GATE_UP_MQ3G256V2_SRC`, `FUSED_GATE_UP_MQ3G256_LLOYD_GFX1100_SRC`, `FUSED_GATE_UP_MQ3G256_LLOYD_SRC`, `FUSED_GATE_UP_MQ4CG256_SRC`, `FUSED_GATE_UP_MQ4G256V2_LUT_SRC`, `FUSED_GATE_UP_MQ4G256V2_SRC`, `FUSED_GATE_UP_MQ4G256_LLOYD_GFX1100_SRC`, `FUSED_GATE_UP_MQ4G256_LLOYD_SRC`, `FUSED_GATE_UP_MQ5G256V2_SRC`, `FUSED_GATE_UP_MQ6G256V2_SRC`, `FUSED_GATE_UP_Q4K_SRC`, `FUSED_GATE_UP_Q8_0_SRC`, `FUSED_GEMMA4_QK_NORM_ROPE_SRC`, `FUSED_GLIMMER_GATE_UP_HFQ4G256_K6656_GFX1100_SRC`, `FUSED_GLIMMER_QKVG_HFQ4G256_K6656_GFX1100_SRC`, `FUSED_QKVZA_HFQ4G256_2WAVE_GFX1100_SRC`, `FUSED_QKVZA_HFQ4G256_HOIST_X32_GFX1100_SRC`, `FUSED_QKVZA_HFQ4G256_K2048_ALL_BUFFER_DLC_GFX1151_SRC`, `FUSED_QKVZA_HFQ4G256_K2048_ALL_BUFFER_GFX1151_SRC`, `FUSED_QKVZA_HFQ4G256_K2048_ALL_BUFFER_GLC_GFX1151_SRC`, `FUSED_QKVZA_HFQ4G256_K2048_ALL_BUFFER_SLC_GFX1151_SRC`, `FUSED_QKVZA_HFQ4G256_K2048_BUFFER_GFX1151_SRC`, `FUSED_QKVZA_HFQ4G256_K2048_CPOL_SLC_GFX1100_SRC`, `FUSED_QKVZA_HFQ4G256_K2048_GFX1100_SRC`, `FUSED_QKVZA_HFQ4G256_K2048_HOIST_X32_GFX1151_SRC`, `FUSED_QKVZA_HFQ4G256_K2048_HYBRID_BUFFER_GFX1151_SRC`, `FUSED_QKVZA_HFQ4G256_K2048_PAIR_BUFFER_GFX1151_SRC`, `FUSED_QKVZA_HFQ4G256_K2048_R2_BUFFER_GFX1151_SRC`, `FUSED_QKVZA_HFQ4G256_K2048_R2_GFX1100_SRC`, `FUSED_QKVZA_HFQ4G256_K2048_R4_STREAM_GFX1151_SRC`, `FUSED_QKVZA_HFQ4G256_K2048_SCALAR_PREP_GFX1100_SRC`, `FUSED_QKVZA_HFQ4G256_K2048_X_BUFFER_GFX1151_SRC`, `FUSED_QKVZA_HFQ4G256_LDSX8_BUFFER_GFX1151_SRC`, `FUSED_QKVZA_HFQ4G256_LDSX8_GFX1100_SRC`, `FUSED_QKVZA_HFQ4G256_REDUCE_CHAIN_GFX1100_SRC`, `FUSED_QKVZA_HFQ4G256_SRC`, `FUSED_QKVZA_HFQ4G256_V2_GFX942_SRC`, `FUSED_QKVZA_HFQ4G256_WAVE64_DP4A_SRC`, `FUSED_QKVZA_HFQ4G256_WAVE64_SHARE_X_GFX1151_SRC`, `FUSED_QKVZA_HFQ4G256_WAVE64_SRC`, `FUSED_QKVZA_HFQ4G256_WAVEPACK4_GFX1100_SRC`, `FUSED_QKVZA_HFQ6G256_WAVE64_DP4A_SRC`, `FUSED_QKVZA_MFP4G32_E8_GFX1151_SRC`, `FUSED_QKVZA_MQ2G256V2_SRC`, `FUSED_QKVZA_MQ3G256V2_SRC`, `FUSED_QKVZA_MQ3G256_LLOYD_GFX1100_SRC`, `FUSED_QKVZA_MQ3G256_LLOYD_SRC`, `FUSED_QKVZA_MQ4CG256_SRC`, `FUSED_QKVZA_MQ4G256V2_K2048_HOIST_X32_GFX1100_SRC`, `FUSED_QKVZA_MQ4G256V2_LUT_SRC`, `FUSED_QKVZA_MQ4G256V2_SRC`, `FUSED_QKVZA_MQ4G256_LLOYD_GFX1100_SRC`, `FUSED_QKVZA_MQ4G256_LLOYD_SRC`, `FUSED_QKVZA_MQ5G256V2_SRC`, `FUSED_QKVZA_MQ6G256V2_SRC`, `FUSED_QKVZA_Q8_0_SRC`, `FUSED_QKV_HFQ4G256_K2048_ALL_BUFFER_GFX1151_SRC`, `FUSED_QKV_HFQ4G256_K2048_ALL_BUFFER_SLC_GFX1151_SRC`, `FUSED_QKV_HFQ4G256_K2048_X_BUFFER_GFX1151_SRC`, `FUSED_QKV_HFQ4G256_QWEN2_BIAS_SRC`, `FUSED_QKV_HFQ4G256_SRC`, `FUSED_QKV_HFQ4G256_V2_GFX942_QWEN2_BIAS_SRC`, `FUSED_QKV_HFQ4G256_V2_GFX942_SRC`, `FUSED_QKV_HFQ4G256_WAVE64_DP4A_SRC`, `FUSED_QKV_HFQ4G256_WAVE64_QWEN2_BIAS_SRC`, `FUSED_QKV_HFQ4G256_WAVE64_SRC`, `FUSED_QKV_HFQ6G256_SRC`, `FUSED_QKV_HFQ6G256_WAVE64_DP4A_SRC`, `FUSED_QKV_MQ2G256V2_QWEN2_BIAS_SRC`, `FUSED_QKV_MQ2G256V2_SRC`, `FUSED_QKV_MQ3G256V2_QWEN2_BIAS_SRC`, `FUSED_QKV_MQ3G256V2_SRC`, `FUSED_QKV_MQ3G256_LLOYD_GFX1100_QWEN2_BIAS_SRC`, `FUSED_QKV_MQ3G256_LLOYD_GFX1100_SRC`, `FUSED_QKV_MQ3G256_LLOYD_QWEN2_BIAS_SRC`, `FUSED_QKV_MQ3G256_LLOYD_SRC`, `FUSED_QKV_MQ4CG256_QWEN2_BIAS_SRC`, `FUSED_QKV_MQ4CG256_SRC`, `FUSED_QKV_MQ4G256V2_K2048_X_BUFFER_GFX1100_SRC`, `FUSED_QKV_MQ4G256V2_LUT_SRC`, `FUSED_QKV_MQ4G256V2_QWEN2_BIAS_SRC`, `FUSED_QKV_MQ4G256V2_SRC`, `FUSED_QKV_MQ4G256_LLOYD_GFX1100_QWEN2_BIAS_SRC`, `FUSED_QKV_MQ4G256_LLOYD_GFX1100_SRC`, `FUSED_QKV_MQ4G256_LLOYD_QWEN2_BIAS_SRC`, `FUSED_QKV_MQ4G256_LLOYD_SRC`, `FUSED_QKV_MQ5G256V2_QWEN2_BIAS_SRC`, `FUSED_QKV_MQ5G256V2_SRC`, `FUSED_QKV_MQ6G256V2_QWEN2_BIAS_SRC`, `FUSED_QKV_MQ6G256V2_SRC`, `FUSED_QKV_Q4K_QWEN2_BIAS_SRC`, `FUSED_QKV_Q4K_SRC`, `FUSED_QKV_Q8_0_QWEN2_BIAS_SRC`, `FUSED_QKV_Q8_0_SRC`, `FUSED_QK_L2_NORM_SCALE_INTERLEAVE_F32_BATCHED_SRC`, `FUSED_QK_L2_NORM_SCALE_SRC`, `FUSED_RMSNORM_MQ_ROTATE_AWQ_FP8_GFX12_SRC`, `FUSED_RMSNORM_MQ_ROTATE_AWQ_FP8_INREG_GFX12_SRC`, `FUSED_RMSNORM_MQ_ROTATE_AWQ_FP8_INREG_SHORT_GFX12_SRC`, `FUSED_RMSNORM_MQ_ROTATE_AWQ_FP8_INREG_SHORT_SEQ_GFX12_SRC`, `FUSED_RMSNORM_MQ_ROTATE_AWQ_G12DEC_SRC`, `FUSED_RMSNORM_MQ_ROTATE_AWQ_I4_B8_GFX1151_SRC`, `FUSED_RMSNORM_MQ_ROTATE_AWQ_I4_B8_SRC`, `FUSED_RMSNORM_MQ_ROTATE_AWQ_I4_FOLD_B8_SRC`, `FUSED_RMSNORM_MQ_ROTATE_AWQ_I4_FOLD_SRC`, `FUSED_RMSNORM_MQ_ROTATE_AWQ_I4_GFX12_SRC`, `FUSED_RMSNORM_MQ_ROTATE_AWQ_I4_GFX12_V2_SLAB_FDIV_SRC`, `FUSED_RMSNORM_MQ_ROTATE_AWQ_I4_GFX12_V2_SLAB_SRC`, `FUSED_RMSNORM_MQ_ROTATE_AWQ_I4_GFX12_V2_SRC`, `FUSED_RMSNORM_MQ_ROTATE_AWQ_I4_SRC`, `FUSED_RMSNORM_MQ_ROTATE_AWQ_SRC`, `FUSED_RMSNORM_MQ_ROTATE_FP8_GFX12_SRC`, `FUSED_RMSNORM_MQ_ROTATE_FP8_INREG_GFX12_SRC`, `FUSED_RMSNORM_MQ_ROTATE_FP8_INREG_SHORT_GFX12_SRC`, `FUSED_RMSNORM_MQ_ROTATE_FP8_INREG_SHORT_SEQ_GFX12_SRC`, `FUSED_RMSNORM_MQ_ROTATE_I4_B8_GFX1151_SRC`, `FUSED_RMSNORM_MQ_ROTATE_I4_B8_SRC`, `FUSED_RMSNORM_MQ_ROTATE_I4_FOLD_B8_SRC`, `FUSED_RMSNORM_MQ_ROTATE_I4_FOLD_SRC`, `FUSED_RMSNORM_MQ_ROTATE_I4_GFX12_SRC`, `FUSED_RMSNORM_MQ_ROTATE_I4_GFX12_V2_SRC`, `FUSED_RMSNORM_MQ_ROTATE_I4_SRC`, `FUSED_RMSNORM_MQ_ROTATE_PLAIN_SRC`, `FUSED_RMSNORM_MQ_ROTATE_SRC`, `FUSED_RMSNORM_MQ_ROTATE_VECSUM_GFX1100_SRC`, `FUSED_RMSNORM_MQ_ROTATE_VECSUM_SIGN_CONST_GFX1100_SRC`, `FUSED_RMSNORM_MQ_ROTATE_VECSUM_SIGN_LDS_GFX1100_SRC`, `FUSED_RMSNORM_MQ_ROTATE_WAVEGRID_GFX1100_SRC`, `FUSED_SIGMOID_ALPHA_GATE_SRC`, `FUSED_SILU_MUL_GIVENS_ROTATE_SRC`, `FUSED_SILU_MUL_MQ_ROTATE_AWQ_FP8_GFX12_SRC`, `FUSED_SILU_MUL_MQ_ROTATE_AWQ_FP8_INREG_GFX12_SRC`, `FUSED_SILU_MUL_MQ_ROTATE_AWQ_HIN_BF16_FP8_GFX12_SRC`, `FUSED_SILU_MUL_MQ_ROTATE_AWQ_HIN_FP8_GFX12_SRC`, `FUSED_SILU_MUL_MQ_ROTATE_AWQ_I4_GFX12_SRC`, `FUSED_SILU_MUL_MQ_ROTATE_AWQ_I4_HIN_BF16_GFX12_SRC`, `FUSED_SILU_MUL_MQ_ROTATE_AWQ_I4_HIN_GFX12_SLAB_SRC`, `FUSED_SILU_MUL_MQ_ROTATE_AWQ_I4_HIN_GFX12_SLAB_TOKFAST_SRC`, `FUSED_SILU_MUL_MQ_ROTATE_AWQ_I4_HIN_GFX12_SRC`, `FUSED_SILU_MUL_MQ_ROTATE_AWQ_I4_HIN_SRC`, `FUSED_SILU_MUL_MQ_ROTATE_AWQ_I4_SRC`, `FUSED_SILU_MUL_MQ_ROTATE_AWQ_INDEXED_SRC`, `FUSED_SILU_MUL_MQ_ROTATE_AWQ_SRC`, `FUSED_SILU_MUL_MQ_ROTATE_FP8_GFX12_SRC`, `FUSED_SILU_MUL_MQ_ROTATE_FP8_INREG_GFX12_SRC`, `FUSED_SILU_MUL_MQ_ROTATE_HIN_BF16_FP8_GFX12_SRC`, `FUSED_SILU_MUL_MQ_ROTATE_HIN_FP8_GFX12_SRC`, `FUSED_SILU_MUL_MQ_ROTATE_I4_GFX12_SRC`, `FUSED_SILU_MUL_MQ_ROTATE_I4_SRC`, `FUSED_SILU_MUL_MQ_ROTATE_SRC`, `GATED_ADD_F32_SRC`, `GATED_DELTA_NET_F32_BATCH_SEQ_SRC`, `GATED_DELTA_NET_F32_CHUNKED_SRC`, `GATED_DELTA_NET_F32_TREE_SRC`, `GATED_DELTA_NET_Q4_SRC`, `GATED_DELTA_NET_Q8_COMPACT2_B12_SRC`, `GATED_DELTA_NET_Q8_COMPACT2_B16_SRC`, `GATED_DELTA_NET_Q8_COMPACT2_B2_SRC`, `GATED_DELTA_NET_Q8_COMPACT2_B4_SRC`, `GATED_DELTA_NET_Q8_COMPACT2_B8_SRC`, `GATED_DELTA_NET_Q8_COMPACT2_DPP_GFX1151_SRC`, `GATED_DELTA_NET_Q8_COMPACT2_R4X2_GFX1151_SRC`, `GATED_DELTA_NET_Q8_COMPACT2_R8_GFX1151_SRC`, `GATED_DELTA_NET_Q8_COMPACT3_B2_SRC`, `GATED_DELTA_NET_Q8_FAST_SRC`, `GATED_DELTA_NET_Q8_SRC`, `GATED_DELTA_NET_Q8_TREE_SRC`, `GATED_DELTA_NET_SRC`, `GATED_NORM_MQ_ROTATE_AWQ_FP8_GFX12_SRC`, `GATED_NORM_MQ_ROTATE_AWQ_FP8_GFX12_XBF16_SRC`, `GATED_NORM_MQ_ROTATE_AWQ_FP8_INREG_GFX12_SRC`, `GATED_NORM_MQ_ROTATE_AWQ_FP8_INREG_GFX12_XBF16_SRC`, `GATED_NORM_MQ_ROTATE_AWQ_GFX1100_SRC`, `GATED_NORM_MQ_ROTATE_AWQ_GFX1151_SRC`, `GATED_NORM_MQ_ROTATE_AWQ_GFX1201_SRC`, `GATED_NORM_MQ_ROTATE_AWQ_I4_GFX1100_V2_SRC`, `GATED_NORM_MQ_ROTATE_AWQ_I4_GFX1151_V2_SRC`, `GATED_NORM_MQ_ROTATE_AWQ_I4_GFX11_SRC`, `GATED_NORM_MQ_ROTATE_AWQ_I4_GFX12_SRC`, `GATED_NORM_MQ_ROTATE_AWQ_I4_GFX12_V2_SLAB_SRC`, `GATED_NORM_MQ_ROTATE_AWQ_I4_GFX12_V2_SRC`, `GATED_NORM_MQ_ROTATE_AWQ_I4_GFX12_V2_XBF16_SLAB_SRC`, `GATED_NORM_MQ_ROTATE_AWQ_I4_GFX12_V2_XBF16_SRC`, `GATED_NORM_MQ_ROTATE_AWQ_I4_GFX12_XBF16_SRC`, `GATED_NORM_MQ_ROTATE_FP8_GFX12_SRC`, `GATED_NORM_MQ_ROTATE_FP8_GFX12_XBF16_SRC`, `GATED_NORM_MQ_ROTATE_FP8_INREG_GFX12_SRC`, `GATED_NORM_MQ_ROTATE_FP8_INREG_GFX12_XBF16_SRC`, `GATED_NORM_MQ_ROTATE_GFX1100_SRC`, `GATED_NORM_MQ_ROTATE_GFX1151_SRC`, `GATED_NORM_MQ_ROTATE_GFX1201_SRC`, `GATED_NORM_MQ_ROTATE_I4_GFX1100_V2_SRC`, `GATED_NORM_MQ_ROTATE_I4_GFX1151_V2_SRC`, `GATED_NORM_MQ_ROTATE_I4_GFX11_SRC`, `GATED_NORM_MQ_ROTATE_I4_GFX12_SRC`, `GATED_NORM_MQ_ROTATE_I4_GFX12_V2_SRC`, `GATED_NORM_MQ_ROTATE_I4_GFX12_V2_XBF16_SRC`, `GATED_NORM_MQ_ROTATE_I4_GFX12_XBF16_SRC`, `GATED_NORM_SRC`, `GDN_CHUNK_KKT_SOLVE_BATCHED_SRC`, `GDN_CHUNK_KKT_SOLVE_GFX1100_SRC`, `GDN_CHUNK_KKT_SOLVE_SRC`, `GDN_CHUNK_PREP_FIXUP_SRC`, `GDN_CHUNK_PREP_GFX11_SRC`, `GDN_CHUNK_PREP_SRC`, `GDN_CHUNK_SCAN_BF16_MSEG_SRC`, `GDN_CHUNK_SCAN_BF16_SRC`, `GDN_CHUNK_SCAN_EMU_BF16_SRC`, `GDN_CHUNK_SCAN_EMU_F16_SRC`, `GDN_CHUNK_SCAN_GFX1151_SRC`, `GDN_CHUNK_SCAN_SRC`, `GDN_PRE_BATCHED_GFX1201_SRC`, `GELU_NEW_MUL_F32_SRC`, `GELU_TANH_SRC`, `GEMMA4_PLE_ACTIVATION_SRC`, `GEMM_BF16_MFMA_GFX942_SRC`, `GEMM_BF16_XF32_BATCHED_SRC`, `GEMM_BF16_XF32_MULTIROW_R16_GFX1151_SRC`, `GEMM_BF16_XF32_MULTIROW_SRC`, `GEMM_BQ1G128_PREFILL_SRC`, `GEMM_BQ1G128_WMMA_SRC`, `GEMM_F16_BIAS_SRC`, `GEMM_F16_SRC`, `GEMM_F16_TILED_SRC`, `GEMM_F16_WMMA_MB4_SRC`, `GEMM_F16_WMMA_MB8_GFX12_SRC`, `GEMM_F16_WMMA_MB8_SRC`, `GEMM_F16_WMMA_SRC`, `GEMM_F16_X_F16_MFMA_GFX942_SRC`, `GEMM_F16_X_F16_WMMA_GFX12_SRC`, `GEMM_F16_X_F16_WMMA_LDS256_SRC`, `GEMM_F16_X_F16_WMMA_LDS_SPLITK_SRC`, `GEMM_F16_X_F16_WMMA_LDS_SRC`, `GEMM_F16_X_F16_WMMA_SRC`, `GEMM_F32_REGISTER_TILED_SRC`, `GEMM_F32_SRC`, `GEMM_GATE_UP_HFP4G32_WMMA_GFX12_SRC`, `GEMM_GATE_UP_HFP4G32_WMMA_SRC`, `GEMM_GATE_UP_HFQ3G256_DOT2_SRC`, `GEMM_GATE_UP_HFQ3G256_DP4A_SRC`, `GEMM_GATE_UP_HFQ3G256_FP16_SRC`, `GEMM_GATE_UP_HFQ3G256_MMQ_BODY_CUH`, `GEMM_GATE_UP_HFQ3G256_MMQ_X16_SRC`, `GEMM_GATE_UP_HFQ3G256_MMQ_X32_SRC`, `GEMM_GATE_UP_HFQ3G256_MMQ_X32_Y64_SRC`, `GEMM_GATE_UP_HFQ3G256_MMQ_X32_Y96_SRC`, `GEMM_GATE_UP_HFQ3G256_MMQ_X8_SRC`, `GEMM_GATE_UP_HFQ3G256_SRC`, `GEMM_GATE_UP_HFQ3G256_WMMA_GFX12_SRC`, `GEMM_GATE_UP_HFQ3G256_WMMA_MB4_SRC`, `GEMM_GATE_UP_HFQ3G256_WMMA_SRC`, `GEMM_GATE_UP_HFQ4G256_DOT2_SRC`, `GEMM_GATE_UP_HFQ4G256_FP16_SRC`, `GEMM_GATE_UP_HFQ4G256_FP16_WAVE64_SRC`, `GEMM_GATE_UP_HFQ4G256_MMQ_BODY_CUH`, `GEMM_GATE_UP_HFQ4G256_MMQ_GFX906_BODY_CUH`, `GEMM_GATE_UP_HFQ4G256_MMQ_GFX906_X16_SRC`, `GEMM_GATE_UP_HFQ4G256_MMQ_GFX906_X16_Y64_SRC`, `GEMM_GATE_UP_HFQ4G256_MMQ_GFX906_X32_SRC`, `GEMM_GATE_UP_HFQ4G256_MMQ_GFX906_X32_Y64_SRC`, `GEMM_GATE_UP_HFQ4G256_MMQ_GFX906_X64_SRC`, `GEMM_GATE_UP_HFQ4G256_MMQ_GFX906_X8_SRC`, `GEMM_GATE_UP_HFQ4G256_MMQ_X16_SRC`, `GEMM_GATE_UP_HFQ4G256_MMQ_X32_SRC`, `GEMM_GATE_UP_HFQ4G256_SRC`, `GEMM_GATE_UP_HFQ4G256_WAVE64_DP4A_SRC`, `GEMM_GATE_UP_HFQ4G256_WMMA_2TILE_SRC`, `GEMM_GATE_UP_HFQ4G256_WMMA_GFX12_BT_SRC`, `GEMM_GATE_UP_HFQ4G256_WMMA_GFX12_SRC`, `GEMM_GATE_UP_HFQ4G256_WMMA_K4_SRC`, `GEMM_GATE_UP_HFQ4G256_WMMA_LDSCOOP_NOSYNC_SRC`, `GEMM_GATE_UP_HFQ4G256_WMMA_LDSCOOP_SRC`, `GEMM_GATE_UP_HFQ4G256_WMMA_LDSX_SRC`, `GEMM_GATE_UP_HFQ4G256_WMMA_SRC`, `GEMM_GATE_UP_HFQ6G256_DOT2_SRC`, `GEMM_GATE_UP_HFQ6G256_FP16_SRC`, `GEMM_GATE_UP_HFQ6G256_SRC`, `GEMM_GATE_UP_HFQ6G256_WAVE64_DP4A_SRC`, `GEMM_GATE_UP_HFQ6G256_WMMA_GFX12_SRC`, `GEMM_GATE_UP_HFQ6G256_WMMA_SRC`, `GEMM_GATE_UP_MQ2G256V2_WMMA_GFX12_BT_SRC`, `GEMM_GATE_UP_MQ2G256V2_WMMA_GFX12_SRC`, `GEMM_GATE_UP_MQ2G256V2_WMMA_SRC`, `GEMM_GATE_UP_MQ3G256V2_WMMA_GFX12_BT_SRC`, `GEMM_GATE_UP_MQ3G256V2_WMMA_GFX12_SRC`, `GEMM_GATE_UP_MQ3G256V2_WMMA_SRC`, `GEMM_GATE_UP_MQ3G256_LLOYD_WMMA_GFX12_SRC`, `GEMM_GATE_UP_MQ3G256_LLOYD_WMMA_MB4_NOSYNC_SRC`, `GEMM_GATE_UP_MQ3G256_LLOYD_WMMA_MB4_SRC`, `GEMM_GATE_UP_MQ3G256_LLOYD_WMMA_NOSYNC_SRC`, `GEMM_GATE_UP_MQ3G256_LLOYD_WMMA_SRC`, `GEMM_GATE_UP_MQ4CG256_WMMA_GFX12_BT_SRC`, `GEMM_GATE_UP_MQ4CG256_WMMA_GFX12_SRC`, `GEMM_GATE_UP_MQ4G256V2_WMMA_FP8_GFX12_BT4_LUT_SRC`, `GEMM_GATE_UP_MQ4G256V2_WMMA_FP8_GFX12_BT4_SRC`, `GEMM_GATE_UP_MQ4G256V2_WMMA_FP8_GFX12_BT8_LUT_SRC`, `GEMM_GATE_UP_MQ4G256V2_WMMA_FP8_GFX12_BT8_SRC`, `GEMM_GATE_UP_MQ4G256V2_WMMA_FP8_GFX12_BT_LUT_SRC`, `GEMM_GATE_UP_MQ4G256V2_WMMA_FP8_GFX12_BT_SRC`, `GEMM_GATE_UP_MQ4G256V2_WMMA_FP8_GFX12_S2BT8_LUT_SRC`, `GEMM_GATE_UP_MQ4G256V2_WMMA_FP8_GFX12_S2BT8_SRC`, `GEMM_GATE_UP_MQ4G256V2_WMMA_FP8_GFX12_V2_B128X128_SILU_H_BF16_SRC`, `GEMM_GATE_UP_MQ4G256V2_WMMA_FP8_GFX12_V2_B128X128_SILU_H_SRC`, `GEMM_GATE_UP_MQ4G256V2_WMMA_FP8_GFX12_V2_B128X128_SRC`, `GEMM_GATE_UP_MQ4G256V2_WMMA_FP8_GFX12_V2_B128X64W4_SRC`, `GEMM_GATE_UP_MQ4G256V2_WMMA_FP8_GFX12_V2_B64X256_SRC`, `GEMM_GATE_UP_MQ4G256V2_WMMA_FP8_GFX12_V2_SRC`, `GEMM_GATE_UP_MQ4G256V2_WMMA_GFX1100_LDSSTAGE_SRC`, `GEMM_GATE_UP_MQ4G256V2_WMMA_GFX1100_MW_LDS_SRC`, `GEMM_GATE_UP_MQ4G256V2_WMMA_GFX11_BT_SRC`, `GEMM_GATE_UP_MQ4G256V2_WMMA_GFX12_BT_SRC`, `GEMM_GATE_UP_MQ4G256V2_WMMA_GFX12_SRC`, `GEMM_GATE_UP_MQ4G256V2_WMMA_GFX12_VT_SRC`, `GEMM_GATE_UP_MQ4G256V2_WMMA_SRC`, `GEMM_GATE_UP_MQ4G256_LLOYD_WMMA_GFX1151_SRC`, `GEMM_GATE_UP_MQ4G256_LLOYD_WMMA_GFX12_SRC`, `GEMM_GATE_UP_MQ4G256_LLOYD_WMMA_MB4_GFX1151_SRC`, `GEMM_GATE_UP_MQ4G256_LLOYD_WMMA_MB4_NOSYNC_GFX1151_SRC`, `GEMM_GATE_UP_MQ4G256_LLOYD_WMMA_MB4_NOSYNC_SRC`, `GEMM_GATE_UP_MQ4G256_LLOYD_WMMA_MB4_SRC`, `GEMM_GATE_UP_MQ4G256_LLOYD_WMMA_NOSYNC_GFX1151_SRC`, `GEMM_GATE_UP_MQ4G256_LLOYD_WMMA_NOSYNC_SRC`, `GEMM_GATE_UP_MQ4G256_LLOYD_WMMA_SRC`, `GEMM_GATE_UP_MQ5G256V2_WMMA_GFX12_BT_SRC`, `GEMM_GATE_UP_MQ5G256V2_WMMA_GFX12_SRC`, `GEMM_GATE_UP_MQ5G256V2_WMMA_SRC`, `GEMM_GATE_UP_MQ6G256V2_WMMA_GFX12_BT_SRC`, `GEMM_GATE_UP_MQ6G256V2_WMMA_GFX12_SRC`, `GEMM_GATE_UP_MQ6G256V2_WMMA_SRC`, `GEMM_GATE_UP_Q8_0_WMMA_GFX12_SRC`, `GEMM_GATE_UP_Q8_0_WMMA_SRC`, `GEMM_HFP4G32_MOE_GROUPED_WMMA_GFX12_SRC`, `GEMM_HFP4G32_RESIDUAL_WMMA_GFX12_SRC`, `GEMM_HFP4G32_RESIDUAL_WMMA_SRC`, `GEMM_HFQ3G256_MOE_GROUPED_WMMA_GFX12_SRC`, `GEMM_HFQ3G256_RESIDUAL_DOT2_SRC`, `GEMM_HFQ3G256_RESIDUAL_FP16_SRC`, `GEMM_HFQ3G256_RESIDUAL_MMQ_BODY_CUH`, `GEMM_HFQ3G256_RESIDUAL_MMQ_X16_SRC`, `GEMM_HFQ3G256_RESIDUAL_MMQ_X32_SRC`, `GEMM_HFQ3G256_RESIDUAL_MMQ_X32_Y32_SRC`, `GEMM_HFQ3G256_RESIDUAL_MMQ_X32_Y64_SRC`, `GEMM_HFQ3G256_RESIDUAL_MMQ_X8_SRC`, `GEMM_HFQ3G256_RESIDUAL_SRC`, `GEMM_HFQ3G256_RESIDUAL_WMMA_GFX12_SRC`, `GEMM_HFQ3G256_RESIDUAL_WMMA_MB4_SRC`, `GEMM_HFQ3G256_RESIDUAL_WMMA_SRC`, `GEMM_HFQ4G128_MMQ_GFX1151_SRC`, `GEMM_HFQ4G128_MOE_DOWN_GROUPED_WMMA_SRC`, `GEMM_HFQ4G128_SRC`, `GEMM_HFQ4G256_LMHEAD_WMMA_GFX12_SRC`, `GEMM_HFQ4G256_MOE_GROUPED_MMQ_GFX1151_SRC`, `GEMM_HFQ4G256_MOE_GROUPED_MMQ_GFX11_DGPU_SRC`, `GEMM_HFQ4G256_MOE_GROUPED_MMQ_GFX12_SRC`, `GEMM_HFQ4G256_MOE_GROUPED_MMQ_K4_GFX1151_SRC`, `GEMM_HFQ4G256_MOE_GROUPED_MMQ_K4_GFX11_DGPU_SRC`, `GEMM_HFQ4G256_MOE_GROUPED_MMQ_K4_GFX12_SRC`, `GEMM_HFQ4G256_MOE_GROUPED_MMQ_K8_GFX1151_SRC`, `GEMM_HFQ4G256_MOE_GROUPED_WMMA_GFX12_SRC`, `GEMM_HFQ4G256_MOE_GROUPED_WMMA_K2_SRC`, `GEMM_HFQ4G256_MOE_GROUPED_WMMA_M2_GFX12_SRC`, `GEMM_HFQ4G256_RESIDUAL_FP16_SRC`, `GEMM_HFQ4G256_RESIDUAL_FP16_WAVE64_SRC`, `GEMM_HFQ4G256_RESIDUAL_MFMA_GFX942_SRC`, `GEMM_HFQ4G256_RESIDUAL_MFMA_V2_GFX942_SRC`, `GEMM_HFQ4G256_RESIDUAL_MFMA_V3_GFX942_SRC`, `GEMM_HFQ4G256_RESIDUAL_MFMA_V4_GFX942_SRC`, `GEMM_HFQ4G256_RESIDUAL_MMQ_BODY_CUH`, `GEMM_HFQ4G256_RESIDUAL_MMQ_GFX12_SRC`, `GEMM_HFQ4G256_RESIDUAL_MMQ_GFX906_BODY_CUH`, `GEMM_HFQ4G256_RESIDUAL_MMQ_GFX906_X16_SRC`, `GEMM_HFQ4G256_RESIDUAL_MMQ_GFX906_X24_SRC`, `GEMM_HFQ4G256_RESIDUAL_MMQ_GFX906_X32_SRC`, `GEMM_HFQ4G256_RESIDUAL_MMQ_GFX906_X40_SRC`, `GEMM_HFQ4G256_RESIDUAL_MMQ_GFX906_X48_SRC`, `GEMM_HFQ4G256_RESIDUAL_MMQ_GFX906_X56_SRC`, `GEMM_HFQ4G256_RESIDUAL_MMQ_GFX906_X64_SRC`, `GEMM_HFQ4G256_RESIDUAL_MMQ_GFX906_X8_SRC`, `GEMM_HFQ4G256_RESIDUAL_MMQ_RDNA2_SRC`, `GEMM_HFQ4G256_RESIDUAL_MMQ_SRC`, `GEMM_HFQ4G256_RESIDUAL_MMQ_X16_SRC`, `GEMM_HFQ4G256_RESIDUAL_MMQ_X32_SRC`, `GEMM_HFQ4G256_RESIDUAL_MMQ_X32_Y64_SRC`, `GEMM_HFQ4G256_RESIDUAL_SRC`, `GEMM_HFQ4G256_RESIDUAL_WAVE64_DP4A_SRC`, `GEMM_HFQ4G256_RESIDUAL_WAVE64_SRC`, `GEMM_HFQ4G256_RESIDUAL_WMMA2_SRC`, `GEMM_HFQ4G256_RESIDUAL_WMMA_GFX1100_MUSE_BT_SRC`, `GEMM_HFQ4G256_RESIDUAL_WMMA_GFX1100_MUSE_CB_BT_SRC`, `GEMM_HFQ4G256_RESIDUAL_WMMA_GFX1100_MUSE_KS_BT_SRC`, `GEMM_HFQ4G256_RESIDUAL_WMMA_GFX1100_MUSE_LDS_G256_SRC`, `GEMM_HFQ4G256_RESIDUAL_WMMA_GFX1100_MUSE_MW_BT_SRC`, `GEMM_HFQ4G256_RESIDUAL_WMMA_GFX1100_MUSE_RM_BT_SRC`, `GEMM_HFQ4G256_RESIDUAL_WMMA_GFX1100_MUSE_RM_HB_SRC`, `GEMM_HFQ4G256_RESIDUAL_WMMA_GFX1100_MUSE_RM_PIPE_SRC`, `GEMM_HFQ4G256_RESIDUAL_WMMA_GFX1100_MUSE_RM_PK_SRC`, `GEMM_HFQ4G256_RESIDUAL_WMMA_GFX12_BT_SRC`, `GEMM_HFQ4G256_RESIDUAL_WMMA_GFX12_MUSE_SRC`, `GEMM_HFQ4G256_RESIDUAL_WMMA_GFX12_SRC`, `GEMM_HFQ4G256_RESIDUAL_WMMA_K2X32_SRC`, `GEMM_HFQ4G256_RESIDUAL_WMMA_K2_SRC`, `GEMM_HFQ4G256_RESIDUAL_WMMA_K4_SRC`, `GEMM_HFQ4G256_RESIDUAL_WMMA_KSPLIT_DET_SRC`, `GEMM_HFQ4G256_RESIDUAL_WMMA_KSPLIT_SRC`, `GEMM_HFQ4G256_RESIDUAL_WMMA_SRC`, `GEMM_HFQ4G256_SRC`, `GEMM_HFQ4G256_WAVE64_DP4A_SRC`, `GEMM_HFQ4G256_WAVE64_SRC`, `GEMM_HFQ4G256_WMMA_GFX12_SRC`, `GEMM_HFQ4G256_WMMA_SRC`, `GEMM_HFQ6G256_MOE_GROUPED_MMQ_GFX1151_SRC`, `GEMM_HFQ6G256_MOE_GROUPED_WMMA_GFX1151_SRC`, `GEMM_HFQ6G256_MOE_GROUPED_WMMA_GFX12_SRC`, `GEMM_HFQ6G256_MOE_GROUPED_WMMA_K2_SRC`, `GEMM_HFQ6G256_MOE_GROUPED_WMMA_V2_GFX12_SRC`, `GEMM_HFQ6G256_RESIDUAL_FP16_SRC`, `GEMM_HFQ6G256_RESIDUAL_SRC`, `GEMM_HFQ6G256_RESIDUAL_WAVE64_DP4A_SRC`, `GEMM_HFQ6G256_RESIDUAL_WMMA_GFX12_SRC`, `GEMM_HFQ6G256_RESIDUAL_WMMA_K2_SRC`, `GEMM_KSPLIT_DET_FINALIZE_SRC`, `GEMM_MFP4G32_E8G128_MOE_GROUPED_WMMA_GFX1151_SRC`, `GEMM_MFP4G32_E8_MOE_GROUPED_WMMA_GFX1151_SRC`, `GEMM_MFP4G32_E8_MOE_GROUPED_WMMA_GFX12_SRC`, `GEMM_MFP4G32_E8_SOA_GROUPED_WMMA_B2_GFX1151_SRC`, `GEMM_MFP4G32_E8_SOA_GROUPED_WMMA_GFX1151_SRC`, `GEMM_MFP4G32_E8_SOA_GROUPED_WMMA_GFX1201_SRC`, `GEMM_MFP4G32_E8_SOA_WMMA_B16_GFX1151_SRC`, `GEMM_MFP4G32_E8_SOA_WMMA_B16_GFX1201_SRC`, `GEMM_MFP4G32_E8_SOA_WMMA_B2_GFX1151_SRC`, `GEMM_MFP4G32_E8_SOA_WMMA_B2_GFX1201_SRC`, `GEMM_MFP4G32_E8_SOA_WMMA_B4_GFX1151_SRC`, `GEMM_MFP4G32_E8_SOA_WMMA_B4_GFX1201_SRC`, `GEMM_MFP4G32_E8_SOA_WMMA_B8_GFX1151_SRC`, `GEMM_MFP4G32_E8_SOA_WMMA_B8_GFX1201_SRC`, `GEMM_MFP4G32_E8_SOA_WMMA_COOP16_GFX1151_SRC`, `GEMM_MFP4G32_E8_SOA_WMMA_COOP4_GFX1151_SRC`, `GEMM_MFP4G32_E8_SOA_WMMA_GFX1151_SRC`, `GEMM_MFP4G32_E8_SOA_WMMA_GFX1201_SRC`, `GEMM_MIXED_MOE_GROUPED_WMMA_4W_K2_SRC`, `GEMM_MIXED_MOE_GROUPED_WMMA_GFX12_SRC`, `GEMM_MIXED_MOE_GROUPED_WMMA_K2_SRC`, `GEMM_MQ2G256V2_RESIDUAL_WMMA_GFX12_BT_SRC`, `GEMM_MQ2G256V2_RESIDUAL_WMMA_GFX12_SRC`, `GEMM_MQ2G256V2_RESIDUAL_WMMA_SRC`, `GEMM_MQ2G256_LLOYD_MOE_GROUPED_MFMA_GFX942_SRC`, `GEMM_MQ2G256_LLOYD_MOE_GROUPED_MMQ_GFX1030_SRC`, `GEMM_MQ2G256_LLOYD_MOE_GROUPED_MMQ_GFX1151_SRC`, `GEMM_MQ2G256_LLOYD_MOE_GROUPED_MMQ_PERM_GFX1151_SRC`, `GEMM_MQ2G256_LLOYD_MOE_GROUPED_MMQ_PREQUANT_PERM_GFX1151_SRC`, `GEMM_MQ2G256_LLOYD_MOE_GROUPED_WMMA_4W_K2_CND_SRC`, `GEMM_MQ2G256_LLOYD_MOE_GROUPED_WMMA_4W_K2_GFX12_SRC`, `GEMM_MQ2G256_LLOYD_MOE_GROUPED_WMMA_4W_K2_MMQLOAD_NOSYNC_SRC`, `GEMM_MQ2G256_LLOYD_MOE_GROUPED_WMMA_4W_K2_MMQLOAD_SRC`, `GEMM_MQ2G256_LLOYD_MOE_GROUPED_WMMA_4W_K2_N32_GFX12_SRC`, `GEMM_MQ2G256_LLOYD_MOE_GROUPED_WMMA_4W_K2_N32_SRC`, `GEMM_MQ2G256_LLOYD_MOE_GROUPED_WMMA_4W_K2_SRC`, `GEMM_MQ2G256_LLOYD_MOE_GROUPED_WMMA_8W_K2_SRC`, `GEMM_MQ2G256_LLOYD_MOE_GROUPED_WMMA_GFX12_SRC`, `GEMM_MQ2G256_LLOYD_MOE_GROUPED_WMMA_K2_SRC`, `GEMM_MQ3G256V2_RESIDUAL_WMMA_GFX12_BT_SRC`, `GEMM_MQ3G256V2_RESIDUAL_WMMA_GFX12_SRC`, `GEMM_MQ3G256V2_RESIDUAL_WMMA_SRC`, `GEMM_MQ3G256_LLOYD_MOE_GROUPED_MMQ_GFX1151_SRC`, `GEMM_MQ3G256_LLOYD_MOE_GROUPED_WMMA_GFX12_SRC`, `GEMM_MQ3G256_LLOYD_MOE_GROUPED_WMMA_K2_SRC`, `GEMM_MQ3G256_LLOYD_RESIDUAL_WMMA_GFX12_SRC`, `GEMM_MQ3G256_LLOYD_RESIDUAL_WMMA_MB4_SRC`, `GEMM_MQ3G256_LLOYD_RESIDUAL_WMMA_SRC`, `GEMM_MQ4CG256_RESIDUAL_WMMA_GFX12_BT_SRC`, `GEMM_MQ4CG256_RESIDUAL_WMMA_GFX12_SRC`, `GEMM_MQ4G128V2_BATCHED_SRC`, `GEMM_MQ4G128V2_MOE_GROUPED_TOP10_MULTIROW_SRC`, `GEMM_MQ4G128V2_MOE_GROUPED_TOP10_O8_R16_GFX1151_SRC`, `GEMM_MQ4G128V2_MOE_GROUPED_WMMA_GFX1151_SRC`, `GEMM_MQ4G128V2_MULTIROW_GFX1151_SRC`, `GEMM_MQ4G256V2_GATE_UP_SILU_A4_IU4_PM_V2B_GFX1151`, `GEMM_MQ4G256V2_MOE_GROUPED_TOP10_O2_R8_GFX1151_SRC`, `GEMM_MQ4G256V2_MOE_GROUPED_TOP10_O4_R8_X4_GFX1151_SRC`, `GEMM_MQ4G256V2_MOE_GROUPED_TOP10_SIMT_SRC`, `GEMM_MQ4G256V2_MOE_GROUPED_WMMA_GFX12_SRC`, `GEMM_MQ4G256V2_MOE_GROUPED_WMMA_K2_SRC`, `GEMM_MQ4G256V2_RESIDUAL_IU4_PM_GFX1100`, `GEMM_MQ4G256V2_RESIDUAL_IU4_PM_V2B_GFX1151`, `GEMM_MQ4G256V2_RESIDUAL_IU4_V2B_GFX11_SRC`, `GEMM_MQ4G256V2_RESIDUAL_IU4_V2C_GFX11_SRC`, `GEMM_MQ4G256V2_RESIDUAL_MMQ_I8_GFX12_APF0_SRC`, `GEMM_MQ4G256V2_RESIDUAL_MMQ_I8_GFX12_SRC`, `GEMM_MQ4G256V2_RESIDUAL_MMQ_IU4_GFX11_SYMFOLD_SRC`, `GEMM_MQ4G256V2_RESIDUAL_MMQ_IU4_GFX11_X5_SYMFOLD_SRC`, `GEMM_MQ4G256V2_RESIDUAL_MMQ_IU4_GFX12_B1`, `GEMM_MQ4G256V2_RESIDUAL_MMQ_IU4_GFX12_B1S`, `GEMM_MQ4G256V2_RESIDUAL_MMQ_IU4_GFX12_B1_CONTROL`, `GEMM_MQ4G256V2_RESIDUAL_MMQ_IU4_GFX12_G12R_SRC`, `GEMM_MQ4G256V2_RESIDUAL_MMQ_IU4_GFX12_SRC`, `GEMM_MQ4G256V2_RESIDUAL_MMQ_IU4_GFX12_SYMFOLD_G12R_SRC`, `GEMM_MQ4G256V2_RESIDUAL_MMQ_IU4_GFX12_SYMFOLD_SRC`, `GEMM_MQ4G256V2_RESIDUAL_MMQ_IU4_GFX12_V3_SRC`, `GEMM_MQ4G256V2_RESIDUAL_MMQ_IU4_GRIDSPEC_SRC`, `GEMM_MQ4G256V2_RESIDUAL_MMQ_IU4_GRIDSPEC_SYMFOLD_SRC`, `GEMM_MQ4G256V2_RESIDUAL_MMQ_IU4_SRC`, `GEMM_MQ4G256V2_RESIDUAL_MMQ_LUT_PERM_SRC`, `GEMM_MQ4G256V2_RESIDUAL_MMQ_LUT_SRC`, `GEMM_MQ4G256V2_RESIDUAL_MMQ_SRC`, `GEMM_MQ4G256V2_RESIDUAL_WMMA_FP8_GFX12_BT12_LUT_SRC`, `GEMM_MQ4G256V2_RESIDUAL_WMMA_FP8_GFX12_BT12_SRC`, `GEMM_MQ4G256V2_RESIDUAL_WMMA_FP8_GFX12_BT4_LUT_SRC`, `GEMM_MQ4G256V2_RESIDUAL_WMMA_FP8_GFX12_BT4_SRC`, `GEMM_MQ4G256V2_RESIDUAL_WMMA_FP8_GFX12_BT8_LUT_SRC`, `GEMM_MQ4G256V2_RESIDUAL_WMMA_FP8_GFX12_BT8_SRC`, `GEMM_MQ4G256V2_RESIDUAL_WMMA_FP8_GFX12_S2BT8_LUT_SRC`, `GEMM_MQ4G256V2_RESIDUAL_WMMA_FP8_GFX12_S2BT8_SRC`, `GEMM_MQ4G256V2_RESIDUAL_WMMA_FP8_GFX12_V2_B128X128_SRC`, `GEMM_MQ4G256V2_RESIDUAL_WMMA_FP8_GFX12_V2_B128X64W4_SRC`, `GEMM_MQ4G256V2_RESIDUAL_WMMA_FP8_GFX12_V2_B64X256_SRC`, `GEMM_MQ4G256V2_RESIDUAL_WMMA_FP8_GFX12_V2_SRC`, `GEMM_MQ4G256V2_RESIDUAL_WMMA_GFX1100_KSPLIT_LDS_SRC`, `GEMM_MQ4G256V2_RESIDUAL_WMMA_GFX1100_LDSSTAGE_SRC`, `GEMM_MQ4G256V2_RESIDUAL_WMMA_GFX1100_MW_LDS_SRC`, `GEMM_MQ4G256V2_RESIDUAL_WMMA_GFX11_BT_SRC`, `GEMM_MQ4G256V2_RESIDUAL_WMMA_GFX12_BT_SRC`, `GEMM_MQ4G256V2_RESIDUAL_WMMA_GFX12_SRC`, `GEMM_MQ4G256V2_RESIDUAL_WMMA_GFX12_VT_SRC`, `GEMM_MQ4G256V2_RESIDUAL_WMMA_SRC`, `GEMM_MQ4G256V2_WMMA_FP8_GFX12_B1`, `GEMM_MQ4G256_LLOYD_RESIDUAL_WMMA_GFX12_SRC`, `GEMM_MQ4G256_LLOYD_RESIDUAL_WMMA_MB2_SRC`, `GEMM_MQ4G256_LLOYD_RESIDUAL_WMMA_MB4_GFX1151_SRC`, `GEMM_MQ4G256_LLOYD_RESIDUAL_WMMA_MB4_SRC`, `GEMM_MQ4G256_LLOYD_RESIDUAL_WMMA_SRC`, `GEMM_MQ5G256V2_RESIDUAL_WMMA_GFX12_BT_SRC`, `GEMM_MQ5G256V2_RESIDUAL_WMMA_GFX12_SRC`, `GEMM_MQ5G256V2_RESIDUAL_WMMA_SRC`, `GEMM_MQ6G256V2_F32_ROWS_SRC`, `GEMM_MQ6G256V2_MOE_GROUPED_WMMA_GFX12_SRC`, `GEMM_MQ6G256V2_MOE_GROUPED_WMMA_K2_SRC`, `GEMM_MQ6G256V2_RESIDUAL_WMMA_GFX12_BT_SRC`, `GEMM_MQ6G256V2_RESIDUAL_WMMA_GFX12_SRC`, `GEMM_MQ6G256V2_RESIDUAL_WMMA_SRC`, `GEMM_MQ6G256_SRC`, `GEMM_MQV2_WMMA_GFX11_BT_SRC`, `GEMM_MQV2_WMMA_GFX11_MW_LDS_SRC`, `GEMM_MW16_RESIDUAL_WMMA_SRC`, `GEMM_PARO_Q4G128_MOE_GROUPED_MMQ_GFX1151_SRC`, `GEMM_PARO_Q4G128_MOE_GROUPED_MMQ_K8_GFX1151_SRC`, `GEMM_PARO_Q4G128_MOE_GROUPED_WMMA_K2_SRC`, `GEMM_Q8_0_BATCHED_SRC`, `GEMM_Q8_0_BATCHED_WIDE_EXACT_SRC`, `GEMM_Q8_0_MMQ_4W_GFX1151_SRC`, `GEMM_Q8_0_MMQ_GFX1030_SRC`, `GEMM_Q8_0_MOE_GROUPED_TOP10_SRC`, `GEMM_Q8_0_RESIDUAL_WMMA_GFX12_SRC`, `GEMM_Q8_0_RESIDUAL_WMMA_SRC`, `GEMM_Q8_0_WMMA_4W_SRC`, `GEMM_Q8_0_WMMA_GFX12_SRC`, `GEMM_Q8_0_WMMA_SRC`, `GEMM_Q8_0_WMMA_X64_SRC`, `GEMM_QKVZA_HFP4G32_WMMA_GFX12_SRC`, `GEMM_QKVZA_HFP4G32_WMMA_SRC`, `GEMM_QKVZA_HFQ3G256_DOT2_SRC`, `GEMM_QKVZA_HFQ3G256_FP16_SRC`, `GEMM_QKVZA_HFQ3G256_MMQ_BODY_CUH`, `GEMM_QKVZA_HFQ3G256_MMQ_X16_SRC`, `GEMM_QKVZA_HFQ3G256_MMQ_X32_SRC`, `GEMM_QKVZA_HFQ3G256_MMQ_X8_SRC`, `GEMM_QKVZA_HFQ3G256_SRC`, `GEMM_QKVZA_HFQ3G256_WMMA_GFX12_SRC`, `GEMM_QKVZA_HFQ3G256_WMMA_MB4_SRC`, `GEMM_QKVZA_HFQ3G256_WMMA_SRC`, `GEMM_QKVZA_HFQ4G256_DOT2_SRC`, `GEMM_QKVZA_HFQ4G256_FP16_SRC`, `GEMM_QKVZA_HFQ4G256_FP16_WAVE64_SRC`, `GEMM_QKVZA_HFQ4G256_MMQ_BODY_CUH`, `GEMM_QKVZA_HFQ4G256_MMQ_X16_SRC`, `GEMM_QKVZA_HFQ4G256_MMQ_X32_SRC`, `GEMM_QKVZA_HFQ4G256_SRC`, `GEMM_QKVZA_HFQ4G256_WAVE64_DP4A_SRC`, `GEMM_QKVZA_HFQ4G256_WAVE64_SRC`, `GEMM_QKVZA_HFQ4G256_WMMA_GFX12_BT_SRC`, `GEMM_QKVZA_HFQ4G256_WMMA_GFX12_SRC`, `GEMM_QKVZA_HFQ4G256_WMMA_SRC`, `GEMM_QKVZA_HFQ6G256_DOT2_SRC`, `GEMM_QKVZA_HFQ6G256_FP16_SRC`, `GEMM_QKVZA_HFQ6G256_SRC`, `GEMM_QKVZA_HFQ6G256_WAVE64_DP4A_SRC`, `GEMM_QKVZA_HFQ6G256_WMMA_GFX12_SRC`, `GEMM_QKVZA_HFQ6G256_WMMA_SRC`, `GEMM_QKVZA_MQ2G256V2_WMMA_GFX12_BT_SRC`, `GEMM_QKVZA_MQ2G256V2_WMMA_GFX12_SRC`, `GEMM_QKVZA_MQ2G256V2_WMMA_SRC`, `GEMM_QKVZA_MQ3G256V2_WMMA_GFX12_BT_SRC`, `GEMM_QKVZA_MQ3G256V2_WMMA_GFX12_SRC`, `GEMM_QKVZA_MQ3G256V2_WMMA_SRC`, `GEMM_QKVZA_MQ3G256_LLOYD_WMMA_GFX12_SRC`, `GEMM_QKVZA_MQ3G256_LLOYD_WMMA_MB4_SRC`, `GEMM_QKVZA_MQ3G256_LLOYD_WMMA_SRC`, `GEMM_QKVZA_MQ4CG256_WMMA_GFX12_BT_SRC`, `GEMM_QKVZA_MQ4CG256_WMMA_GFX12_SRC`, `GEMM_QKVZA_MQ4G256V2_WMMA_FP8_GFX12_BT12_LUT_SRC`, `GEMM_QKVZA_MQ4G256V2_WMMA_FP8_GFX12_BT12_SRC`, `GEMM_QKVZA_MQ4G256V2_WMMA_FP8_GFX12_BT4_LUT_SRC`, `GEMM_QKVZA_MQ4G256V2_WMMA_FP8_GFX12_BT4_SRC`, `GEMM_QKVZA_MQ4G256V2_WMMA_FP8_GFX12_BT8_LUT_SRC`, `GEMM_QKVZA_MQ4G256V2_WMMA_FP8_GFX12_BT8_SRC`, `GEMM_QKVZA_MQ4G256V2_WMMA_FP8_GFX12_S2BT8_LUT_SRC`, `GEMM_QKVZA_MQ4G256V2_WMMA_FP8_GFX12_S2BT8_SRC`, `GEMM_QKVZA_MQ4G256V2_WMMA_FP8_GFX12_V2_B128X128_SRC`, `GEMM_QKVZA_MQ4G256V2_WMMA_FP8_GFX12_V2_B128X64W4_SRC`, `GEMM_QKVZA_MQ4G256V2_WMMA_FP8_GFX12_V2_B64X256_SRC`, `GEMM_QKVZA_MQ4G256V2_WMMA_FP8_GFX12_V2_SRC`, `GEMM_QKVZA_MQ4G256V2_WMMA_GFX11_BT_SRC`, `GEMM_QKVZA_MQ4G256V2_WMMA_GFX12_BT_SRC`, `GEMM_QKVZA_MQ4G256V2_WMMA_GFX12_SRC`, `GEMM_QKVZA_MQ4G256V2_WMMA_GFX12_VT_SRC`, `GEMM_QKVZA_MQ4G256V2_WMMA_SRC`, `GEMM_QKVZA_MQ4G256_LLOYD_WMMA_GFX1151_SRC`, `GEMM_QKVZA_MQ4G256_LLOYD_WMMA_GFX12_SRC`, `GEMM_QKVZA_MQ4G256_LLOYD_WMMA_MB4_GFX1151_SRC`, `GEMM_QKVZA_MQ4G256_LLOYD_WMMA_MB4_SRC`, `GEMM_QKVZA_MQ4G256_LLOYD_WMMA_SRC`, `GEMM_QKVZA_MQ5G256V2_WMMA_GFX12_BT_SRC`, `GEMM_QKVZA_MQ5G256V2_WMMA_GFX12_SRC`, `GEMM_QKVZA_MQ5G256V2_WMMA_SRC`, `GEMM_QKVZA_MQ6G256V2_WMMA_GFX12_BT_SRC`, `GEMM_QKVZA_MQ6G256V2_WMMA_GFX12_SRC`, `GEMM_QKVZA_MQ6G256V2_WMMA_SRC`, `GEMM_QKVZA_Q8_0_WMMA_GFX12_SRC`, `GEMM_QKVZA_Q8_0_WMMA_SRC`, `GEMM_QKVZA_TQ2G128_SRC`, `GEMM_QKV_HFP4G32_WMMA_FP8_GFX12_SRC`, `GEMM_QKV_HFP4G32_WMMA_GFX12_SRC`, `GEMM_QKV_HFP4G32_WMMA_SRC`, `GEMM_QKV_HFQ3G256_DOT2_SRC`, `GEMM_QKV_HFQ3G256_DP4A_SRC`, `GEMM_QKV_HFQ3G256_FP16_SRC`, `GEMM_QKV_HFQ3G256_MMQ_BODY_CUH`, `GEMM_QKV_HFQ3G256_MMQ_X16_SRC`, `GEMM_QKV_HFQ3G256_MMQ_X32_SRC`, `GEMM_QKV_HFQ3G256_MMQ_X8_SRC`, `GEMM_QKV_HFQ3G256_SRC`, `GEMM_QKV_HFQ3G256_WMMA_GFX12_SRC`, `GEMM_QKV_HFQ3G256_WMMA_MB4_SRC`, `GEMM_QKV_HFQ3G256_WMMA_SRC`, `GEMM_QKV_HFQ4G256_DOT2_SRC`, `GEMM_QKV_HFQ4G256_FP16_SRC`, `GEMM_QKV_HFQ4G256_FP16_WAVE64_SRC`, `GEMM_QKV_HFQ4G256_MMQ_BODY_CUH`, `GEMM_QKV_HFQ4G256_MMQ_GFX906_BODY_CUH`, `GEMM_QKV_HFQ4G256_MMQ_GFX906_X16_SRC`, `GEMM_QKV_HFQ4G256_MMQ_GFX906_X32_SRC`, `GEMM_QKV_HFQ4G256_MMQ_GFX906_X64_SRC`, `GEMM_QKV_HFQ4G256_MMQ_GFX906_X8_SRC`, `GEMM_QKV_HFQ4G256_MMQ_X16_SRC`, `GEMM_QKV_HFQ4G256_MMQ_X32_SRC`, `GEMM_QKV_HFQ4G256_SRC`, `GEMM_QKV_HFQ4G256_WAVE64_DP4A_SRC`, `GEMM_QKV_HFQ4G256_WAVE64_SRC`, `GEMM_QKV_HFQ4G256_WMMA_GFX12_SRC`, `GEMM_QKV_HFQ4G256_WMMA_SRC`, `GEMM_QKV_HFQ6G256_DOT2_SRC`, `GEMM_QKV_HFQ6G256_FP16_SRC`, `GEMM_QKV_HFQ6G256_SRC`, `GEMM_QKV_HFQ6G256_WAVE64_DP4A_SRC`, `GEMM_QKV_HFQ6G256_WMMA_GFX12_SRC`, `GEMM_QKV_HFQ6G256_WMMA_SRC`, `GEMM_QKV_MQ2G256V2_WMMA_GFX12_SRC`, `GEMM_QKV_MQ2G256V2_WMMA_SRC`, `GEMM_QKV_MQ3G256V2_WMMA_GFX12_SRC`, `GEMM_QKV_MQ3G256V2_WMMA_SRC`, `GEMM_QKV_MQ3G256_LLOYD_WMMA_GFX12_SRC`, `GEMM_QKV_MQ3G256_LLOYD_WMMA_MB4_SRC`, `GEMM_QKV_MQ3G256_LLOYD_WMMA_SRC`, `GEMM_QKV_MQ4CG256_WMMA_GFX12_SRC`, `GEMM_QKV_MQ4G256V2_WMMA_FP8_GFX12_BT12_LUT_SRC`, `GEMM_QKV_MQ4G256V2_WMMA_FP8_GFX12_BT12_SRC`, `GEMM_QKV_MQ4G256V2_WMMA_FP8_GFX12_BT4_LUT_SRC`, `GEMM_QKV_MQ4G256V2_WMMA_FP8_GFX12_BT4_SRC`, `GEMM_QKV_MQ4G256V2_WMMA_FP8_GFX12_BT8_LUT_SRC`, `GEMM_QKV_MQ4G256V2_WMMA_FP8_GFX12_BT8_SRC`, `GEMM_QKV_MQ4G256V2_WMMA_FP8_GFX12_S2BT8_LUT_SRC`, `GEMM_QKV_MQ4G256V2_WMMA_FP8_GFX12_S2BT8_SRC`, `GEMM_QKV_MQ4G256V2_WMMA_FP8_GFX12_V2_B128X128_SRC`, `GEMM_QKV_MQ4G256V2_WMMA_FP8_GFX12_V2_B128X64W4_SRC`, `GEMM_QKV_MQ4G256V2_WMMA_FP8_GFX12_V2_B64X256_SRC`, `GEMM_QKV_MQ4G256V2_WMMA_FP8_GFX12_V2_SRC`, `GEMM_QKV_MQ4G256V2_WMMA_GFX11_BT_SRC`, `GEMM_QKV_MQ4G256V2_WMMA_GFX1201_BT_SRC`, `GEMM_QKV_MQ4G256V2_WMMA_GFX12_SRC`, `GEMM_QKV_MQ4G256V2_WMMA_GFX12_VT_SRC`, `GEMM_QKV_MQ4G256V2_WMMA_SRC`, `GEMM_QKV_MQ4G256_LLOYD_WMMA_GFX1151_SRC`, `GEMM_QKV_MQ4G256_LLOYD_WMMA_GFX12_SRC`, `GEMM_QKV_MQ4G256_LLOYD_WMMA_MB4_GFX1151_SRC`, `GEMM_QKV_MQ4G256_LLOYD_WMMA_MB4_SRC`, `GEMM_QKV_MQ4G256_LLOYD_WMMA_SRC`, `GEMM_QKV_MQ5G256V2_WMMA_GFX12_SRC`, `GEMM_QKV_MQ5G256V2_WMMA_SRC`, `GEMM_QKV_MQ6G256V2_WMMA_GFX12_SRC`, `GEMM_QKV_MQ6G256V2_WMMA_SRC`, `GEMM_QKV_MQV2_WMMA_GFX1201_BT_SRC`, `GEMM_QKV_Q8_0_WMMA_GFX12_SRC`, `GEMM_QKV_Q8_0_WMMA_SRC`, `GEMM_TQ2G128_PREFILL_SRC`, `GEMM_TQ2G128_WMMA_SRC`, `GEMV_BF16_XF32_SRC`, `GEMV_BQ1G128_SRC`, `GEMV_BQ1G128_XBATCH_SRC`, `GEMV_F16_BIAS_XF32_SRC`, `GEMV_F16_XF32_SRC`, `GEMV_HFP4G32_DOT2_GFX11_SRC`, `GEMV_HFP4G32_FP8_GFX12_SRC`, `GEMV_HFP4G32_GFX1100_SRC`, `GEMV_HFP4G32_SRC`, `GEMV_HFQ2G128_SRC`, `GEMV_HFQ2G256_SRC`, `GEMV_HFQ3G128_SRC`, `GEMV_HFQ3G256_GFX1100_SRC`, `GEMV_HFQ3G256_RESIDUAL_GFX1100_SRC`, `GEMV_HFQ3G256_RESIDUAL_SRC`, `GEMV_HFQ3G256_SRC`, `GEMV_HFQ4G1024_SRC`, `GEMV_HFQ4G128_MOE_DOWN_GROUPED_SRC`, `GEMV_HFQ4G128_MOE_DOWN_RESIDUAL_SCALED_K8_INDEXED_SRC`, `GEMV_HFQ4G128_RESIDUAL_SIGMOID_SCALED_SRC`, `GEMV_HFQ4G128_SRC`, `GEMV_HFQ4G256_GFX1030_V1_SRC`, `GEMV_HFQ4G256_GFX1030_V2_SRC`, `GEMV_HFQ4G256_GFX1030_V3_SRC`, `GEMV_HFQ4G256_GFX1030_V4_SRC`, `GEMV_HFQ4G256_GFX1030_V5_SRC`, `GEMV_HFQ4G256_GFX1100_SRC`, `GEMV_HFQ4G256_K2048_GFX1100_SRC`, `GEMV_HFQ4G256_LM_HEAD_DOT2_GFX1151_SRC`, `GEMV_HFQ4G256_LM_HEAD_R1_HYBRID_BUFFER_GFX1151_SRC`, `GEMV_HFQ4G256_MOE_DOWN_INDEXED_BATCHED_SRC`, `GEMV_HFQ4G256_MOE_DOWN_INDEXED_BATCHED_WAVE64_SRC`, `GEMV_HFQ4G256_MOE_DOWN_INDEXED_SRC`, `GEMV_HFQ4G256_MOE_DOWN_INDEXED_WAVE64_SRC`, `GEMV_HFQ4G256_MOE_DOWN_K8_INDEXED_BATCHED_EXPANDED_BUFFER_GFX1151_SRC`, `GEMV_HFQ4G256_MOE_DOWN_K8_INDEXED_BATCHED_EXPANDED_CPOL_SLC_GFX1100_SRC`, `GEMV_HFQ4G256_MOE_DOWN_K8_INDEXED_BATCHED_EXPANDED_HYBRID_BUFFER_GFX1151_SRC`, `GEMV_HFQ4G256_MOE_DOWN_K8_INDEXED_BATCHED_EXPANDED_ROW1_BUFFER_GFX1151_SRC`, `GEMV_HFQ4G256_MOE_DOWN_K8_INDEXED_BATCHED_EXPANDED_ROW2_BUFFER_GFX1151_SRC`, `GEMV_HFQ4G256_MOE_DOWN_K8_INDEXED_BATCHED_EXPANDED_ROW2_CLUSTERED_GFX1151_SRC`, `GEMV_HFQ4G256_MOE_DOWN_K8_INDEXED_BATCHED_EXPANDED_ROW8_GFX1151_SRC`, `GEMV_HFQ4G256_MOE_DOWN_K8_INDEXED_BATCHED_EXPANDED_SRC`, `GEMV_HFQ4G256_MOE_DOWN_K8_INDEXED_FUSED_ACC_SRC`, `GEMV_HFQ4G256_MOE_DOWN_K8_INDEXED_LAST_COMBINE_SRC`, `GEMV_HFQ4G256_MOE_DOWN_SRC`, `GEMV_HFQ4G256_MOE_GATE_INDEXED_K2048_GFX1151_SRC`, `GEMV_HFQ4G256_MOE_GATE_UP_INDEXED_BATCHED_SRC`, `GEMV_HFQ4G256_MOE_GATE_UP_INDEXED_BATCHED_WAVE64_SRC`, `GEMV_HFQ4G256_MOE_GATE_UP_INDEXED_CPOL_DLC_GFX1100_SRC`, `GEMV_HFQ4G256_MOE_GATE_UP_INDEXED_CPOL_GLC_GFX1100_SRC`, `GEMV_HFQ4G256_MOE_GATE_UP_INDEXED_CPOL_SLC_GFX1100_SRC`, `GEMV_HFQ4G256_MOE_GATE_UP_INDEXED_K2048_ALL_BUFFER_GFX1151_SRC`, `GEMV_HFQ4G256_MOE_GATE_UP_INDEXED_K2048_BUFFER_GFX1151_SRC`, `GEMV_HFQ4G256_MOE_GATE_UP_INDEXED_K2048_GFX1100_SRC`, `GEMV_HFQ4G256_MOE_GATE_UP_INDEXED_K2048_GFX1151_SRC`, `GEMV_HFQ4G256_MOE_GATE_UP_INDEXED_K2048_HYBRID_GFX1151_SRC`, `GEMV_HFQ4G256_MOE_GATE_UP_INDEXED_K2048_LOW_VGPR_GFX1151_SRC`, `GEMV_HFQ4G256_MOE_GATE_UP_INDEXED_K2048_PAIR_ALL_BUFFER_GFX1151_SRC`, `GEMV_HFQ4G256_MOE_GATE_UP_INDEXED_K2048_PAIR_BUFFER_GFX1151_SRC`, `GEMV_HFQ4G256_MOE_GATE_UP_INDEXED_K2048_PAIR_VGPR_GFX1151_SRC`, `GEMV_HFQ4G256_MOE_GATE_UP_INDEXED_K2048_ROUTE_ALL_BUFFER_GFX1151_SRC`, `GEMV_HFQ4G256_MOE_GATE_UP_INDEXED_K2816_CPOL_DLC_GFX1100_SRC`, `GEMV_HFQ4G256_MOE_GATE_UP_INDEXED_K2816_CPOL_GLC_GFX1100_SRC`, `GEMV_HFQ4G256_MOE_GATE_UP_INDEXED_K2816_CPOL_SLC_GFX1100_SRC`, `GEMV_HFQ4G256_MOE_GATE_UP_INDEXED_K2816_SRC`, `GEMV_HFQ4G256_MOE_GATE_UP_INDEXED_LOW_VGPR_GFX1100_SRC`, `GEMV_HFQ4G256_MOE_GATE_UP_INDEXED_PAIR_SLC_GFX1100_SRC`, `GEMV_HFQ4G256_MOE_GATE_UP_INDEXED_RANK_INTERLEAVE_GFX1100_SRC`, `GEMV_HFQ4G256_MOE_GATE_UP_INDEXED_ROWTILE_SRC`, `GEMV_HFQ4G256_MOE_GATE_UP_INDEXED_SRC`, `GEMV_HFQ4G256_MOE_GATE_UP_INDEXED_WAVE64_GFX1151_SRC`, `GEMV_HFQ4G256_MOE_GATE_UP_INDEXED_WAVE64_SRC`, `GEMV_HFQ4G256_MOE_GATE_UP_INDEXED_WG2_SRC`, `GEMV_HFQ4G256_MOE_GATE_UP_PAIRED_WAVES_K2048_GFX1151_SRC`, `GEMV_HFQ4G256_MOE_GATE_UP_PERSISTENT_RANK8_GFX1151_SRC`, `GEMV_HFQ4G256_MOE_GATE_UP_SRC`, `GEMV_HFQ4G256_MOE_NINEPATH_D3_SRC`, `GEMV_HFQ4G256_MOE_NINEPATH_D4_SRC`, `GEMV_HFQ4G256_MOE_UP_INDEXED_K2048_GFX1151_SRC`, `GEMV_HFQ4G256_MULTIROW_ALL_BUFFER_GFX1151_SRC`, `GEMV_HFQ4G256_MULTIROW_BUFFER_GFX1151_SRC`, `GEMV_HFQ4G256_MULTIROW_GFX1100_SRC`, `GEMV_HFQ4G256_MULTIROW_HYBRID_BUFFER_GFX1151_SRC`, `GEMV_HFQ4G256_MULTIROW_HYBRID_DLC_GFX1151_SRC`, `GEMV_HFQ4G256_MULTIROW_HYBRID_GLC_GFX1151_SRC`, `GEMV_HFQ4G256_MULTIROW_HYBRID_SLC_GFX1151_SRC`, `GEMV_HFQ4G256_MULTIROW_K2048_GFX1151_SRC`, `GEMV_HFQ4G256_MULTIROW_SRC`, `GEMV_HFQ4G256_RESIDUAL_BUFFER_GFX1151_SRC`, `GEMV_HFQ4G256_RESIDUAL_CPOL_RT_GFX1100_SRC`, `GEMV_HFQ4G256_RESIDUAL_CPOL_RT_LOW_GFX1100_SRC`, `GEMV_HFQ4G256_RESIDUAL_CPOL_SLC_GFX1100_SRC`, `GEMV_HFQ4G256_RESIDUAL_GFX1100_SRC`, `GEMV_HFQ4G256_RESIDUAL_GFX942_SRC`, `GEMV_HFQ4G256_RESIDUAL_HYBRID_BUFFER_GFX1151_SRC`, `GEMV_HFQ4G256_RESIDUAL_K2048_GFX1100_SRC`, `GEMV_HFQ4G256_RESIDUAL_K4096_GFX1151_SRC`, `GEMV_HFQ4G256_RESIDUAL_MULTIROW_GFX1100_SRC`, `GEMV_HFQ4G256_RESIDUAL_MULTIROW_GFX1151_SRC`, `GEMV_HFQ4G256_RESIDUAL_ROW1_GFX1151_SRC`, `GEMV_HFQ4G256_RESIDUAL_RT_LOW_GFX1151_SRC`, `GEMV_HFQ4G256_RESIDUAL_SCALED_SRC`, `GEMV_HFQ4G256_RESIDUAL_SIGMOID_BUFFER_GFX1100_SRC`, `GEMV_HFQ4G256_RESIDUAL_SIGMOID_K512_BUFFER_GFX1100_SRC`, `GEMV_HFQ4G256_RESIDUAL_SIGMOID_K512_BUFFER_GFX1151_SRC`, `GEMV_HFQ4G256_RESIDUAL_SIGMOID_K512_HOIST_X16_BUFFER_GFX1100_SRC`, `GEMV_HFQ4G256_RESIDUAL_SIGMOID_ROWS4_BUFFER_GFX1100_SRC`, `GEMV_HFQ4G256_RESIDUAL_SIGMOID_ROWS4_GFX1100_SRC`, `GEMV_HFQ4G256_RESIDUAL_SRC`, `GEMV_HFQ4G256_RESIDUAL_STAGE_X32_GFX1100_SRC`, `GEMV_HFQ4G256_RESIDUAL_V2_GFX942_SRC`, `GEMV_HFQ4G256_RESIDUAL_V3_GFX942_SRC`, `GEMV_HFQ4G256_RESIDUAL_WAVE64_GFX1151_SRC`, `GEMV_HFQ4G256_RESIDUAL_WAVE64_PREFETCH_SRC`, `GEMV_HFQ4G256_RESIDUAL_WAVE64_SRC`, `GEMV_HFQ4G256_SRC`, `GEMV_HFQ4G256_WIDE_SRC`, `GEMV_HFQ4G256_XBATCH_GEN_SRC`, `GEMV_HFQ4G256_XBATCH_SRC`, `GEMV_HFQ4G512_SRC`, `GEMV_HFQ5G256_MOE_DOWN_K8_INDEXED_BATCHED_EXPANDED_SRC`, `GEMV_HFQ5G256_MOE_GATE_UP_INDEXED_BATCHED_SRC`, `GEMV_HFQ5G256_MOE_GATE_UP_INDEXED_SRC`, `GEMV_HFQ5G256_RESIDUAL_SIGMOID_SCALED_SRC`, `GEMV_HFQ5G256_RESIDUAL_SRC`, `GEMV_HFQ5G256_RESIDUAL_WAVE64_PREFETCH_SRC`, `GEMV_HFQ5G256_RESIDUAL_WAVE64_SRC`, `GEMV_HFQ5G256_SRC`, `GEMV_HFQ6G256_MOE_DOWN_K8_INDEXED_BATCHED_EXPANDED_SRC`, `GEMV_HFQ6G256_MOE_GATE_UP_INDEXED_BATCHED_SRC`, `GEMV_HFQ6G256_MOE_GATE_UP_INDEXED_SRC`, `GEMV_HFQ6G256_RESIDUAL_SIGMOID_SCALED_SRC`, `GEMV_HFQ6G256_RESIDUAL_SRC`, `GEMV_HFQ6G256_RESIDUAL_WAVE64_PREFETCH_SRC`, `GEMV_HFQ6G256_RESIDUAL_WAVE64_SRC`, `GEMV_HFQ6G256_SRC`, `GEMV_HFQ8G256_SRC`, `GEMV_MFP3G32_E8_SRC`, `GEMV_MFP4G32_E8G128_MOE_DOWN_K8_INDEXED_BATCHED_EXPANDED_SRC`, `GEMV_MFP4G32_E8G128_MOE_GATE_UP_K8_INDEXED_BATCHED_SRC`, `GEMV_MFP4G32_E8G128_SRC`, `GEMV_MFP4G32_E8_GFX1151_SRC`, `GEMV_MFP4G32_E8_LDSX_GFX1151_SRC`, `GEMV_MFP4G32_E8_MOE_DOWN_K8_INDEXED_BATCHED_EXPANDED_GFX1151_SRC`, `GEMV_MFP4G32_E8_MOE_DOWN_K8_INDEXED_BATCHED_EXPANDED_GFX11_DGPU_SRC`, `GEMV_MFP4G32_E8_MOE_GATE_UP_K8_INDEXED_BATCHED_GFX1151_SRC`, `GEMV_MFP4G32_E8_MOE_GATE_UP_K8_INDEXED_BATCHED_GFX11_DGPU_SRC`, `GEMV_MFP4G32_E8_SOA_BATCHED_B16_GFX1151_SRC`, `GEMV_MFP4G32_E8_SOA_BATCHED_B1_GFX1151_SRC`, `GEMV_MFP4G32_E8_SOA_BATCHED_B2_GFX1151_SRC`, `GEMV_MFP4G32_E8_SOA_BATCHED_B3_GFX1151_SRC`, `GEMV_MFP4G32_E8_SOA_BATCHED_B4_GFX1151_SRC`, `GEMV_MFP4G32_E8_SOA_BATCHED_B5_GFX1151_SRC`, `GEMV_MFP4G32_E8_SOA_BATCHED_B6_GFX1151_SRC`, `GEMV_MFP4G32_E8_SOA_BATCHED_B7_GFX1151_SRC`, `GEMV_MFP4G32_E8_SOA_BATCHED_B8_GFX1151_SRC`, `GEMV_MFP4G32_E8_SOA_BATCHED_PACK_B3_GFX1151_SRC`, `GEMV_MFP4G32_E8_SOA_BATCHED_PAIR_B3_GFX1151_SRC`, `GEMV_MFP4G32_E8_SOA_FP8_MFMA_GFX942_SRC`, `GEMV_MFP4G32_E8_SOA_GFX1151_SRC`, `GEMV_MFP4G32_E8_SOA_GFX1201_SRC`, `GEMV_MFP4G32_E8_SOA_GFX942_SRC`, `GEMV_MFP4G32_E8_SOA_GROUPED_BATCHED_B16_GFX1151_SRC`, `GEMV_MFP4G32_E8_SOA_GROUPED_BATCHED_B1_GFX1151_SRC`, `GEMV_MFP4G32_E8_SOA_GROUPED_BATCHED_B2_GFX1151_SRC`, `GEMV_MFP4G32_E8_SOA_GROUPED_BATCHED_B3_GFX1151_SRC`, `GEMV_MFP4G32_E8_SOA_GROUPED_BATCHED_B4_GFX1151_SRC`, `GEMV_MFP4G32_E8_SOA_GROUPED_BATCHED_B5_GFX1151_SRC`, `GEMV_MFP4G32_E8_SOA_GROUPED_BATCHED_B6_GFX1151_SRC`, `GEMV_MFP4G32_E8_SOA_GROUPED_BATCHED_B7_GFX1151_SRC`, `GEMV_MFP4G32_E8_SOA_GROUPED_BATCHED_B8_GFX1151_SRC`, `GEMV_MFP4G32_E8_SOA_GROUPED_BUFFER_GFX1151_SRC`, `GEMV_MFP4G32_E8_SOA_GROUPED_GFX1100_SRC`, `GEMV_MFP4G32_E8_SOA_GROUPED_GFX1151_SRC`, `GEMV_MFP4G32_E8_SOA_GROUPED_GFX1201_SRC`, `GEMV_MFP4G32_E8_SOA_LATE_SCALE_GFX1201_SRC`, `GEMV_MFP4G32_E8_SOA_LUT_SRC`, `GEMV_MFP4G32_E8_SOA_MIXED_JOBS_GFX1201_SRC`, `GEMV_MFP4G32_E8_SOA_MOE_GATE_UP_K8_INDEXED_BATCHED_SRC`, `GEMV_MFP4G32_E8_SOA_PAIR_GFX942_SRC`, `GEMV_MFP4G32_E8_SOA_PREFETCH4_GFX1201_SRC`, `GEMV_MFP4G32_E8_SOA_SHARED_JOBS_GFX1100_SRC`, `GEMV_MFP4G32_E8_SOA_SHARED_PAIR_GFX1201_SRC`, `GEMV_MFP4G32_E8_SOA_SRC`, `GEMV_MFP4G32_E8_SOA_STRIP_SRC`, `GEMV_MFP4G32_E8_SOA_U4_BUFFER_CPOL0_GFX1151_SRC`, `GEMV_MFP4G32_E8_SOA_U4_BUFFER_CPOL20_GFX1151_SRC`, `GEMV_MFP4G32_E8_SOA_U4_BUFFER_CPOL22_GFX1151_SRC`, `GEMV_MFP4G32_E8_SOA_U4_BUFFER_SQRT_SOFTPLUS_GFX1151_SRC`, `GEMV_MFP4G32_E8_SOA_U4_SRC`, `GEMV_MFP4G32_E8_SOA_U8_SRC`, `GEMV_MFP4G32_E8_SOA_W128_GFX942_SRC`, `GEMV_MFP4G32_E8_SRC`, `GEMV_MFP4G32_E8_STRIP_GFX1151_SRC`, `GEMV_MFP4G32_LLOYD_SRC`, `GEMV_MFP4G32_P_SRC`, `GEMV_MIXED_MOE_DOWN_K8_INDEXED_BATCHED_EXPANDED_SRC`, `GEMV_MIXED_MOE_GATE_UP_INDEXED_BATCHED_SRC`, `GEMV_MQ2G256GL_MOE_DOWN_INDEXED_SRC`, `GEMV_MQ2G256GL_MOE_GATE_UP_INDEXED_SRC`, `GEMV_MQ2G256V2_MULTIROW_SRC`, `GEMV_MQ2G256V2_RESIDUAL_SRC`, `GEMV_MQ2G256V2_SRC`, `GEMV_MQ2G256_LLOYD_MOE_DOWN_ALLRANKS_GFX942_SRC`, `GEMV_MQ2G256_LLOYD_MOE_DOWN_EXPANDED_K4_GFX942_SRC`, `GEMV_MQ2G256_LLOYD_MOE_DOWN_EXPANDED_K4_SRC`, `GEMV_MQ2G256_LLOYD_MOE_DOWN_INDEXED_BATCHED_K4_SRC`, `GEMV_MQ2G256_LLOYD_MOE_DOWN_INDEXED_BATCHED_K8ALL_SRC`, `GEMV_MQ2G256_LLOYD_MOE_DOWN_INDEXED_GFX942_SRC`, `GEMV_MQ2G256_LLOYD_MOE_DOWN_INDEXED_MFMA_F16_GFX942_SRC`, `GEMV_MQ2G256_LLOYD_MOE_DOWN_INDEXED_MFMA_GFX942_SRC`, `GEMV_MQ2G256_LLOYD_MOE_DOWN_INDEXED_R4_SRC`, `GEMV_MQ2G256_LLOYD_MOE_DOWN_INDEXED_SRC`, `GEMV_MQ2G256_LLOYD_MOE_GATE_UP_INDEXED_BATCHED_K4096_LDS_SRC`, `GEMV_MQ2G256_LLOYD_MOE_GATE_UP_INDEXED_BATCHED_K4_SRC`, `GEMV_MQ2G256_LLOYD_MOE_GATE_UP_INDEXED_GFX942_SRC`, `GEMV_MQ2G256_LLOYD_MOE_GATE_UP_INDEXED_MFMA_F16_GFX942_SRC`, `GEMV_MQ2G256_LLOYD_MOE_GATE_UP_INDEXED_MFMA_GFX942_SRC`, `GEMV_MQ2G256_LLOYD_MOE_GATE_UP_INDEXED_SRC`, `GEMV_MQ2G256_LLOYD_SRC`, `GEMV_MQ3G256GL_MOE_DOWN_INDEXED_SRC`, `GEMV_MQ3G256GL_MOE_GATE_UP_INDEXED_SRC`, `GEMV_MQ3G256V2_MULTIROW_SRC`, `GEMV_MQ3G256V2_RESIDUAL_SRC`, `GEMV_MQ3G256V2_SRC`, `GEMV_MQ3G256_LLOYD_GFX1100_SRC`, `GEMV_MQ3G256_LLOYD_MOE_DOWN_INDEXED_BATCHED_K4_SRC`, `GEMV_MQ3G256_LLOYD_MOE_DOWN_INDEXED_R4_SRC`, `GEMV_MQ3G256_LLOYD_MOE_DOWN_INDEXED_SRC`, `GEMV_MQ3G256_LLOYD_MOE_GATE_UP_INDEXED_SRC`, `GEMV_MQ3G256_LLOYD_MOE_NINEPATH_D4_SRC`, `GEMV_MQ3G256_LLOYD_RESIDUAL_GFX1100_SRC`, `GEMV_MQ3G256_LLOYD_RESIDUAL_SRC`, `GEMV_MQ3G256_LLOYD_SRC`, `GEMV_MQ4CG256_MULTIROW_SRC`, `GEMV_MQ4CG256_RESIDUAL_SRC`, `GEMV_MQ4CG256_SRC`, `GEMV_MQ4G128V2_MOE_DOWN_TOP10_INDEXED_BATCHED_EXPANDED_SRC`, `GEMV_MQ4G128V2_SRC`, `GEMV_MQ4G128_SRC`, `GEMV_MQ4G256V2_LUT_SRC`, `GEMV_MQ4G256V2_MOE_DOWN_K8_INDEXED_BATCHED_EXPANDED_SRC`, `GEMV_MQ4G256V2_MOE_GATE_UP_K8_INDEXED_BATCHED_SRC`, `GEMV_MQ4G256V2_MOE_GATE_UP_K8_INDEXED_K2048_NOLDS_GFX1100_SRC`, `GEMV_MQ4G256V2_MOE_GATE_UP_K8_INDEXED_SRC`, `GEMV_MQ4G256V2_MOE_GATE_UP_TOP10_INDEXED_BATCHED_SRC`, `GEMV_MQ4G256V2_MOE_NINEPATH_D4_SRC`, `GEMV_MQ4G256V2_MOE_NINEPATH_RPB8_GFX1100_SRC`, `GEMV_MQ4G256V2_MULTIROW_SRC`, `GEMV_MQ4G256V2_RESIDUAL_LUT_SRC`, `GEMV_MQ4G256V2_RESIDUAL_R1_K4096_GFX1100_NOSCRATCH_SRC`, `GEMV_MQ4G256V2_RESIDUAL_ROW_SERIAL_GFX1151_SRC`, `GEMV_MQ4G256V2_RESIDUAL_SIGMOID_SCALED_K512_SRC`, `GEMV_MQ4G256V2_RESIDUAL_SRC`, `GEMV_MQ4G256V2_SRC`, `GEMV_MQ4G256V2_XBATCH_RESIDUAL_SRC`, `GEMV_MQ4G256V2_XBATCH_SRC`, `GEMV_MQ4G256_LLOYD_GFX1100_SRC`, `GEMV_MQ4G256_LLOYD_MULTIACC_DIAG_GFX1100_SRC`, `GEMV_MQ4G256_LLOYD_RESIDUAL_GFX1100_SRC`, `GEMV_MQ4G256_LLOYD_RESIDUAL_SRC`, `GEMV_MQ4G256_LLOYD_SRC`, `GEMV_MQ4G256_MOE_GATE_UP_INDEXED_K2816_SRC`, `GEMV_MQ4G256_SRC`, `GEMV_MQ5G256V2_MULTIROW_SRC`, `GEMV_MQ5G256V2_RESIDUAL_SRC`, `GEMV_MQ5G256V2_SRC`, `GEMV_MQ5G256_SRC`, `GEMV_MQ6G256V2_MOE_DOWN_K8_INDEXED_BATCHED_EXPANDED_SRC`, `GEMV_MQ6G256V2_MOE_GATE_UP_K8_INDEXED_BATCHED_SRC`, `GEMV_MQ6G256V2_MOE_GATE_UP_K8_INDEXED_SRC`, `GEMV_MQ6G256V2_MOE_NINEPATH_D4_SRC`, `GEMV_MQ6G256V2_MULTIROW_SRC`, `GEMV_MQ6G256V2_RESIDUAL_SRC`, `GEMV_MQ6G256V2_SRC`, `GEMV_MQ6G256V2_XBATCH_RESIDUAL_SRC`, `GEMV_MQ6G256V2_XBATCH_SRC`, `GEMV_MQ6G256_SRC`, `GEMV_MQ8G256_SRC`, `GEMV_PARO4G128_SRC`, `GEMV_PARO_Q4G128_MOE_DOWN_K8_INDEXED_BATCHED_SRC`, `GEMV_PARO_Q4G128_MOE_GATE_UP_INDEXED_SRC`, `GEMV_PARO_Q4G128_MOE_GATE_UP_K8_INDEXED_BATCHED_SRC`, `GEMV_Q4AS8_SRC`, `GEMV_Q4F16_G32_SRC`, `GEMV_Q4F16_G64_SRC`, `GEMV_Q4F16_G64_WIDE_SRC`, `GEMV_Q4K_SRC`, `GEMV_Q4LUT_SRC`, `GEMV_Q4WAVE_SRC`, `GEMV_Q6K_SRC`, `GEMV_Q8HFQ_SRC`, `GEMV_Q8HFQ_WIDE_SRC`, `GEMV_Q8_0_MOE_DOWN_RESIDUAL_SCALED_K8_INDEXED_SRC`, `GEMV_Q8_0_MOE_DOWN_TOP10_INDEXED_BATCHED_EXPANDED_SRC`, `GEMV_Q8_0_MOE_GATE_UP_K8_INDEXED_SRC`, `GEMV_Q8_0_SRC`, `GEMV_Q8_0_WIDE_SRC`, `GEMV_SRC`, `GEMV_TQ2G128_SRC`, `GEMV_TQ2G128_XBATCH_SRC`, `GIVENS_COMMON_SRC`, `GIVENS_ROTATE_SRC`, `GIVENS_ROTATE_TO_SRC`, `GREEDY_ACCEPT_SRC`, `HASH_ROUTER_NORMALIZE_BATCHED_SRC`, `HASH_ROUTER_NORMALIZE_BUF_SRC`, `HASH_ROUTER_NORMALIZE_SRC`, `HC_APPLY_ALPHA_BATCHED_SRC`, `HC_APPLY_ALPHA_SRC`, `HC_COMPUTE_CONTROL_BATCHED_FUSED24_GFX1201_SRC`, `HC_COMPUTE_CONTROL_BATCHED_SRC`, `HC_COMPUTE_CONTROL_SRC`, `HC_COMPUTE_CONTROL_T1024_SRC`, `HC_COMPUTE_CONTROL_WMMA_B1_GFX1201_SRC`, `HC_COMPUTE_CONTROL_WMMA_B2_GFX1201_SRC`, `HC_COMPUTE_CONTROL_WMMA_B4_GFX1201_SRC`, `HC_FINALIZE_CONTROL_SRC`, `HC_FINALIZE_INPUT_MAP_SRC`, `HC_HEAD_COMPUTE_PRE_SRC`, `HC_INPUT_MAP_BATCHED_SRC`, `HC_INPUT_MAP_SRC`, `HC_INV_RMS_BATCHED_GFX1201_SRC`, `HC_MIX_4STREAM_BATCHED_SRC`, `HC_MIX_4STREAM_PEER4_GFX1201_SRC`, `HC_MIX_4STREAM_SRC`, `HC_PRE_POST_SIGMOID_SCALE_SRC`, `HC_SINKHORN_4X4_BATCHED_SRC`, `HC_SINKHORN_4X4_SRC`, `HC_SPLIT_FINALIZE_BATCHED_SRC`, `HC_STREAMS_INIT_FROM_EMBED_BATCHED_SRC`, `HFQ4G256_DEQUANTIZE_TO_F16_SRC`, `INDEXER_COMPRESSED_K_SCORE_SRC`, `INDEXER_KV_GATHER_SRC`, `INDEXER_RELU_SCORE_BATCHED_SRC`, `INDEXER_RELU_SCORE_BUF_SHARDED_GFX1201_SRC`, `INDEXER_RELU_SCORE_BUF_SRC`, `INDEXER_RELU_SCORE_WMMA_BATCHED_GFX12_SRC`, `INDEXER_RELU_SCORE_WMMA_BATCHED_KLDS_GFX1151_SRC`, `INDEXER_RELU_SCORE_WMMA_BATCHED_SHARDED_GFX1201_SRC`, `INDEXER_RELU_SCORE_WMMA_BATCHED_SRC`, `INDEXER_TOPK_TWOSTAGE_SRC`, `INDEXER_TOP_K_BATCHED_BOUNDED_GFX1151_SRC`, `INDEXER_TOP_K_BATCHED_BOUNDED_GFX1201_SRC`, `INDEXER_TOP_K_BATCHED_SRC`, `INDEXER_TOP_K_BUF_BOUNDED_GFX1151_SRC`, `INDEXER_TOP_K_BUF_BOUNDED_GFX942_SRC`, `INDEXER_TOP_K_BUF_PARALLEL_GFX1151_SRC`, `INDEXER_TOP_K_BUF_SRC`, `INDEXER_TOP_K_BUF_UNROLLED_GFX1151_SRC`, `INDEXER_TOP_K_SRC`, `KV_CACHE_WRITE_ASYM3_Q8_PAIR_GFX1100_SRC`, `KV_CACHE_WRITE_ASYM_K_FWHT2_BATCHED_SRC`, `KV_CACHE_WRITE_ASYM_K_FWHT2_SRC`, `KV_CACHE_WRITE_ASYM_K_FWHT3_BATCHED_SRC`, `KV_CACHE_WRITE_ASYM_K_FWHT3_SRC`, `KV_CACHE_WRITE_ASYM_K_FWHT4_BATCHED_SRC`, `KV_CACHE_WRITE_ASYM_K_FWHT4_SRC`, `KV_CACHE_WRITE_ASYM_K_GIVENS2_BATCHED_SRC`, `KV_CACHE_WRITE_ASYM_K_GIVENS2_SRC`, `KV_CACHE_WRITE_ASYM_K_GIVENS3_BATCHED_SRC`, `KV_CACHE_WRITE_ASYM_K_GIVENS3_HD512_BATCHED_SRC`, `KV_CACHE_WRITE_ASYM_K_GIVENS3_HD512_SRC`, `KV_CACHE_WRITE_ASYM_K_GIVENS3_SRC`, `KV_CACHE_WRITE_ASYM_K_GIVENS4_BATCHED_SRC`, `KV_CACHE_WRITE_ASYM_K_GIVENS4_SRC`, `KV_CACHE_WRITE_BF16_SRC`, `KV_CACHE_WRITE_F16_SRC`, `KV_CACHE_WRITE_F32_BATCHED_SRC`, `KV_CACHE_WRITE_FP8_E4M3_BATCHED_SRC`, `KV_CACHE_WRITE_FP8_E4M3_SRC`, `KV_CACHE_WRITE_FWHT256_2BIT_BATCHED_SRC`, `KV_CACHE_WRITE_FWHT256_2BIT_SRC`, `KV_CACHE_WRITE_FWHT256_4BIT_BATCHED_SRC`, `KV_CACHE_WRITE_FWHT256_4BIT_SRC`, `KV_CACHE_WRITE_FWHT3_HD512_BATCHED_SRC`, `KV_CACHE_WRITE_FWHT3_HD512_SRC`, `KV_CACHE_WRITE_HFQ4_SRC`, `KV_CACHE_WRITE_HFQ8_SRC`, `KV_CACHE_WRITE_INT8C_F16_SRC`, `KV_CACHE_WRITE_INT8C_SRC`, `KV_CACHE_WRITE_INT8_SRC`, `KV_CACHE_WRITE_Q4_SRC`, `KV_CACHE_WRITE_Q8_0_BATCHED_SRC`, `KV_CACHE_WRITE_Q8_0_PAIR_GFX1100_SRC`, `KV_CACHE_WRITE_Q8_0_SRC`, `KV_CACHE_WRITE_Q8_SRC`, `KV_CACHE_WRITE_SRC`, `KV_COMPACT_GATHER_SRC`, `KV_FOLD_ASYM2_SRC`, `KV_FOLD_ASYM3_SRC`, `KV_FOLD_ASYM4_SRC`, `KV_FOLD_Q8_SRC`, `KV_SLOT_DESC_H`, `KV_SLOT_DESC_PAGED_H`, `KV_TRANSCODE_K_FWHT4_TO_FWHT2_SRC`, `KV_TRANSCODE_K_FWHT4_TO_FWHT3_SRC`, `KV_TRANSCODE_V_LLOYD_DOWN_SRC`, `KV_TRANSCODE_V_Q8_TO_LLOYD4_SRC`, `L2_NORM_SRC`, `LAYERNORM_MODULATE_F32_SRC`, `LAYERNORM_SRC`, `LOGIT_PENALTY_TABLE_SRC`, `LOGIT_SOFTCAP_SRC`, `MAX_PROB_SRC`, `MODULATE_F32_SRC`, `MOE_DOWN_COMBINE_GROUPED_K8_SRC`, `MOE_DOWN_COMBINE_GROUPED_TOP10_SRC`, `MOE_DOWN_COMBINE_K8_BATCHED_SRC`, `MOE_DOWN_COMBINE_K8_BATCHED_VEC4_SRC`, `MOE_DOWN_COMBINE_RMSNORM_MQ_ROTATE_VECSUM_GFX1100_SRC`, `MOE_DOWN_COMBINE_RMSNORM_MQ_ROTATE_VECSUM_GFX1151_SRC`, `MOE_DOWN_COMBINE_TOP10_BATCHED_SRC`, `MOE_DOWN_UNSCATTER_K8_SRC`, `MOE_GATE_UP_UNSCATTER_K8_SRC`, `MOE_GATE_UP_UNSCATTER_SILU_TOP10_SRC`, `MOE_GATE_UP_UNSCATTER_TOP10_SRC`, `MOE_ROUTER_SOFTMAX_TOP10_F32_FAST_SRC`, `MOE_ROUTER_SOFTMAX_TOP10_F32_SRC`, `MOE_ROUTER_SOFTMAX_TOPK_K8_WAVE64_EXACT_SHARED_SILU_MQ_ROTATE_SRC`, `MOE_ROUTER_SOFTMAX_TOPK_K8_WAVE64_EXACT_SRC`, `MOE_ROUTER_SOFTMAX_TOPK_K8_WAVE64_SRC`, `MOE_SCATTER_FUSED_K8_SRC`, `MOE_SCATTER_FUSED_TOP10_SRC`, `MOE_SCATTER_HISTOGRAM_K8_SRC`, `MOE_SCATTER_OFFSETS_K8_SRC`, `MOE_SCATTER_PERMUTE_K8_SRC`, `MOE_SOFTMAX_TOPK_K8_BATCHED_SRC`, `MOE_SOFTMAX_TOPK_K8_SRC`, `MOE_TOPK_RENORM_K8_BATCHED_SRC`, `MOE_TOPK_RENORM_K8_SRC`, `MOE_TOPK_RENORM_TOP10_BATCHED_SRC`, `MOE_UNSCATTER_SILU_CLAMP_K8_SRC`, `MQ4V2_FP8_FRAGMENT_REPACK_GFX1201`, `MQ4_FP8_VERIFY_PM_GFX1201`, `MQ4_ROWS_X4_PM_GFX1151`, `MQ4_ROWS_X4_PM_GFX1201`, `MQ4_VERIFY_PM_GFX1201`, `MQ_ROTATE_X_128_V2_SRC`, `MQ_ROTATE_X_AWQ_FP8_GFX12_SRC`, `MQ_ROTATE_X_AWQ_FP8_INREG_GFX12_SRC`, `MQ_ROTATE_X_AWQ_I4_GFX12_SRC`, `MQ_ROTATE_X_AWQ_I4_SRC`, `MQ_ROTATE_X_DUAL_FP8_GFX12_SRC`, `MQ_ROTATE_X_FP8_GFX12_SRC`, `MQ_ROTATE_X_FP8_INREG_GFX12_SRC`, `MQ_ROTATE_X_GFX942_SRC`, `MQ_ROTATE_X_I4_GFX12_SRC`, `MQ_ROTATE_X_I4_SRC`, `MUL_SRC`, `PACK_F16_TO_FP8_MQ4V2_GFX12_SRC`, `PACK_F32_TO_FP8_GFX12_SRC`, `PACK_F32_TO_FP8_MQ4V2_GFX12_SRC`, `PACK_FP8_VERIFY_ROWS_GFX1201_SRC`, `PFLASH_SCORE_FWHT2_KV_SRC`, `PFLASH_SCORE_FWHT3_KV_SRC`, `PFLASH_SCORE_FWHT4_KV_SRC`, `PFLASH_SCORE_Q8_KV_SRC`, `QKV_SPLIT_INTERLEAVED_SRC`, `QK_RMSNORM_ROPE_FLUX_SRC`, `QSA_GATHER_PM_GFX1151`, `QSA_GATHER_PM_GFX1201`, `QSA_SELECT_BF16_PM_GFX1151`, `QSA_SELECT_PM_GFX1151`, `QUICK_GELU_F32_SRC`, `QWEN35_FA_PREP_GFX1100_SRC`, `QWEN35_FA_PREP_GFX1151_SRC`, `QWEN35_FA_PREP_GFX1201_SRC`, `QWEN4_BF16_ROUND_TRIP_SRC`, `QWEN4_GEMM_F16_X_F16_WMMA_LDS256_SRC`, `QWEN4_GEMM_MQ4G256V2_MOE_GROUPED_WMMA_K2_SRC`, `QWEN4_GEMM_MQ6G256V2_WMMA_GFX12_X4_SRC`, `QWEN4_GEMM_MQV2_WMMA_GFX11_BT_SRC`, `QWEN4_GEMV_BF16_XF32_SRC`, `QWEN4_GEMV_MQ4G256V2_MOE_GATE_UP_K8_INDEXED_BATCHED_SRC`, `QWEN4_GEMV_MQ4G256_SRC`, `QWEN4_GEMV_MQ6G256V2_SRC`, `QWEN4_GEMV_Q8_0_SRC`, `QWEN4_GEMV_Q8_0_WIDE_SRC`, `QWEN4_HC_STREAMS_INIT_FROM_EMBED_BATCHED_SRC`, `QWEN4_LANE_ROWS_PM_GFX1151`, `QWEN4_LANE_ROWS_PM_GFX1151_SHA256`, `QWEN4_MOE_IU4_SYM_GFX1151_SRC`, `QWEN4_MOE_IU4_SYM_PM_GFX1151`, `QWEN4_MOE_IU4_SYM_PM_GFX1201`, `QWEN4_MOE_ROTATE128_I4_SRC`, `QWEN4_MOE_SCATTER_FUSED_K8_SRC`, `QWEN4_MOE_SCATTER_STABLE_TOP10_SRC`, `QWEN4_MOE_SYM_CHECK_GFX1201_SRC`, `QWEN4_MQ4_X4_PM_GFX1151`, `QWEN4_MQ6_X4_PM_GFX1151`, `QWEN4_MQ6_X4_PM_GFX1201`, `QWEN4_SHARED_GU_PM_GFX1151`, `QWEN4_SILU_MUL_SRC`, `REPEAT_INTERLEAVE_QK_BATCHED_SRC`, `REPEAT_INTERLEAVE_QK_SRC`, `REQUANT_G256_SRC`, `RMSNORM_AT_SLOT_BUF_SRC`, `RMSNORM_REDUCE_GFX1100_SRC`, `RMSNORM_REDUCE_GFX942_SRC`, `RMSNORM_RESIDUAL_ADD_SRC`, `RMSNORM_ROWSPLIT_SRC`, `RMSNORM_SRC`, `ROPE_2D_FLUX_F32_FAST_SRC`, `ROPE_2D_FLUX_F32_SRC`, `ROPE_2D_HALFSPLIT_QKV_INTERLEAVED_SRC`, `ROPE_2D_HALFSPLIT_SRC`, `ROPE_BATCHED_SRC`, `ROPE_MROPE_HALFSPLIT_BATCHED_SRC`, `ROPE_MROPE_HALFSPLIT_SRC`, `ROPE_PARTIAL_HALFSPLIT_BATCHED_SRC`, `ROPE_PARTIAL_HALFSPLIT_HEADGRID_SRC`, `ROPE_PARTIAL_HALFSPLIT_SRC`, `ROPE_PARTIAL_HALVED_BATCHED_SRC`, `ROPE_PARTIAL_HALVED_SRC`, `ROPE_PARTIAL_INTERLEAVED_BATCHED_SRC`, `ROPE_PARTIAL_INTERLEAVED_SRC`, `ROPE_SRC`, `ROPE_TAIL_INTERLEAVED_BATCHED_SRC`, `ROPE_TAIL_INTERLEAVED_H64D128R64_GFX1201_SRC`, `ROPE_TAIL_INTERLEAVED_SRC`, `ROPE_TAIL_YARN_INTERLEAVED_AT_SLOT_BUF_SRC`, `ROPE_TAIL_YARN_INTERLEAVED_BATCHED_SRC`, `ROPE_TAIL_YARN_INTERLEAVED_SRC`, `ROTATE_WITH_RMS_GFX1100_SRC`, `ROTATE_WITH_RMS_GFX942_SRC`, `ROTATE_X_MQ_AWQ_SRC`, `ROUND_H_BF16_QUALITY_GFX12_SRC`, `SAMPLE_ROWS_PF_SRC`, `SAMPLE_TOP_P_PARALLEL_SRC`, `SAMPLE_TOP_P_SRC`, `SCALED_ADD_INPLACE_SRC`, `SCALE_F32_SRC`, `SIGMOID_MUL_MQ_ROTATE_X_AWQ_FP8_GFX12_SRC`, `SIGMOID_MUL_MQ_ROTATE_X_AWQ_FP8_GIL_GFX12_SRC`, `SIGMOID_MUL_MQ_ROTATE_X_AWQ_FP8_INREG_GFX12_SRC`, `SIGMOID_MUL_MQ_ROTATE_X_AWQ_FP8_INREG_GIL_GFX12_SRC`, `SIGMOID_MUL_MQ_ROTATE_X_AWQ_I4_GFX11_SRC`, `SIGMOID_MUL_MQ_ROTATE_X_AWQ_I4_GFX12_SLAB_SRC`, `SIGMOID_MUL_MQ_ROTATE_X_AWQ_I4_GFX12_SRC`, `SIGMOID_MUL_MQ_ROTATE_X_AWQ_I4_GIL_GFX12_SLAB_SRC`, `SIGMOID_MUL_MQ_ROTATE_X_AWQ_I4_GIL_GFX12_SRC`, `SIGMOID_MUL_MQ_ROTATE_X_FP8_GFX12_SRC`, `SIGMOID_MUL_MQ_ROTATE_X_FP8_GIL_GFX12_SRC`, `SIGMOID_MUL_MQ_ROTATE_X_FP8_INREG_GFX12_SRC`, `SIGMOID_MUL_MQ_ROTATE_X_FP8_INREG_GIL_GFX12_SRC`, `SIGMOID_MUL_SRC`, `SIGMOID_SCALED_RESIDUAL_ADD_BATCHED_SRC`, `SIGMOID_SRC`, `SILU_MUL_SRC`, `SILU_SRC`, `SOFTMAX_PROB_GATHER_BATCHED_SRC`, `SOFTMAX_SRC`, `SOFTMAX_TEMP_BATCHED_SRC`, `SOFTPLUS_SRC`, `SPLIT_MQ4V2_Z_BETAALPHA_SRC`, `SQRT_SOFTPLUS_F32_SRC`, `STATE_OVERLAP_SHIFT_F32_BUF_SRC`, `STATE_RING_WRITE_F32_BUF_SRC`, `SWA_RING_WRITE_BATCHED_SRC`, `SWA_RING_WRITE_BUF_SRC`, `SWA_VISIBILITY_STAGE_BATCHED_SRC`, `TOPK8_RESCORE_MQ6G256V2_SRC`, `TOPK_LOGITS_SRC`, `TOPK_LOGSUMEXP_BATCHED_SRC`, `TP_GRAPH_SIGNAL_GFX1201_SRC`, `TRANSPOSE_SRC`, `TRIATTN_ACCUMULATE_SRC`, `TRIATTN_SCORE_ASYM2_SRC`, `TRIATTN_SCORE_ASYM3_SRC`, `TRIATTN_SCORE_ASYM4_SRC`, `TRIATTN_SCORE_FWHT_SRC`, `TRIATTN_SCORE_Q8_SRC`, `TURBO_COMMON_H`, `V4F_ATTN_POS0_SRC`, `V4F_ATTN_SWA_BATCHED_DEBUG_SRC`, `V4F_ATTN_SWA_BATCHED_SRC`, `V4F_ATTN_SWA_BUF_SRC`, `V4F_ATTN_SWA_SRC`, `V4F_ATTN_SWA_TOPK_BATCHED_SRC`, `V4F_ATTN_SWA_TOPK_BATCHED_WMMA_GFX12_SRC`, `V4F_ATTN_SWA_TOPK_BATCHED_WMMA_SRC`, `V4F_ATTN_SWA_TOPK_BUF_LARGE_SERIAL_GFX1151_SRC`, `V4F_ATTN_SWA_TOPK_BUF_SRC`, `V4F_ATTN_SWA_TOPK_BUF_XLANE_GFX1151_SRC`, `V4F_ATTN_SWA_TOPK_DIRECT_BATCHED_SRC`, `V4F_ATTN_SWA_TOPK_DIRECT_WMMA_SRC`, `V4F_CONVERT_F32_TO_F16_SRC`, `V4F_FUSED_SILU_MUL_CLAMP_MQ_ROTATE_SRC`, `V4F_MOE_TOPK_BIAS_AWARE_BATCHED_SRC`, `V4F_MOE_TOPK_BIAS_AWARE_SRC`, `V4F_SILU_MUL_CLAMP_SRC`, `V4F_TOPK_KV_GATHER_BATCHED_SRC`, `V4F_TOPK_KV_GATHER_BATCHED_TILED_GFX1201_SRC`, `V4F_TOPK_KV_GATHER_BATCHED_TILED_SHARDED_GFX1201_SRC`, `V4F_TOPK_KV_GATHER_BUF_SHARDED_GFX1201_SRC`, `V4F_TOPK_KV_GATHER_BUF_SRC`, `V4F_TOPK_KV_GATHER_IDENTITY_BATCHED_SHARDED_GFX1201_SRC`, `V4F_TOPK_KV_GATHER_IDENTITY_BATCHED_SRC`, `V4F_TOPK_KV_GATHER_IDENTITY_BUF_SHARDED_GFX1201_SRC`, `V4F_TOPK_KV_GATHER_IDENTITY_BUF_SRC`, `V4F_TOPK_KV_GATHER_TILED_GFX1151_SRC`, `VAE_ATTN_SRC`, `VAE_CONV1X1_SRC`, `VAE_CONV3X3_S2_SRC`, `VAE_CONV3X3_SRC`, `VAE_GROUPNORM_SRC`, `VAE_IM2COL_SRC`, `VAE_LAYOUT_SRC`, `VAE_UPSAMPLE2X_SRC`, `VIT_ATTENTION_OPT_SRC`, `VIT_ATTENTION_QTILED_SRC`, `VIT_ATTENTION_SRC`, `WO_PER_GROUP_BATCHED_F32_SRC`, `WO_PER_GROUP_BATCHED_HFQ4G256_SRC`, `WO_PER_GROUP_BATCHED_Q8_0_MULTIROW_SRC`, `WO_PER_GROUP_BATCHED_Q8_0_SRC`, `WO_PER_GROUP_BATCHED_Q8_0_WMMA_4W_SRC`, `ZERO_F32_SRC`, `ZERO_INACTIVE_ROWS_F32_SRC`, `fused_gate_up_mq3g256_lloyd_for_arch`, `fused_gate_up_mq4g256_lloyd_for_arch`, `fused_gate_up_mq4g256v2_k5120_gfx1100_src`, `fused_qkv_mq3g256_lloyd_for_arch`, `fused_qkv_mq3g256_lloyd_qwen2_bias_for_arch`, `fused_qkv_mq4g256_lloyd_for_arch`, `fused_qkv_mq4g256_lloyd_qwen2_bias_for_arch`, `fused_qkvza_mq3g256_lloyd_for_arch`, `fused_qkvza_mq4g256_lloyd_for_arch`, `gated_norm_mq_rotate_awq_k6144_gfx1100_src`, `gated_norm_mq_rotate_awq_k6144_gfx1201_src`, `gated_norm_mq_rotate_k6144_gfx1100_src`, `gated_norm_mq_rotate_k6144_gfx1201_src`, `gemm_f16_x_f16_wmma_lds256_src`, `gemm_gate_up_mq3g256_lloyd_wmma_for_arch`, `gemm_gate_up_mq3g256_lloyd_wmma_mb4_for_arch`, `gemm_gate_up_mq3g256_lloyd_wmma_mb4_nosync_for_arch`, `gemm_gate_up_mq3g256_lloyd_wmma_nosync_for_arch`, `gemm_gate_up_mq4g256_lloyd_wmma_for_arch`, `gemm_gate_up_mq4g256_lloyd_wmma_mb4_for_arch`, `gemm_gate_up_mq4g256_lloyd_wmma_mb4_nosync_for_arch`, `gemm_gate_up_mq4g256_lloyd_wmma_nosync_for_arch`, `gemm_mq3g256_lloyd_residual_wmma_for_arch`, `gemm_mq3g256_lloyd_residual_wmma_mb4_for_arch`, `gemm_mq4g256_lloyd_residual_wmma_for_arch`, `gemm_mq4g256_lloyd_residual_wmma_mb2_for_arch`, `gemm_mq4g256_lloyd_residual_wmma_mb4_for_arch`, `gemm_mqv2_wmma_gfx11_bt_src`, `gemm_qkv_mq3g256_lloyd_wmma_for_arch`, `gemm_qkv_mq3g256_lloyd_wmma_mb4_for_arch`, `gemm_qkv_mq4g256_lloyd_wmma_for_arch`, `gemm_qkv_mq4g256_lloyd_wmma_mb4_for_arch`, `gemm_qkvza_mq3g256_lloyd_wmma_for_arch`, `gemm_qkvza_mq3g256_lloyd_wmma_mb4_for_arch`, `gemm_qkvza_mq4g256_lloyd_wmma_for_arch`, `gemm_qkvza_mq4g256_lloyd_wmma_mb4_for_arch`, `gemv_bf16_xf32_src`, `gemv_hfp4g32_for_arch`, `gemv_hfq3g256_for_arch`, `gemv_hfq3g256_residual_for_arch`, `gemv_hfq4g256_for_arch`, `gemv_hfq4g256_residual_for_arch`, `gemv_mq3g256_lloyd_for_arch`, `gemv_mq3g256_lloyd_residual_for_arch`, `gemv_mq3g256v2_for_arch`, `gemv_mq3g256v2_residual_for_arch`, `gemv_mq4g256_lloyd_for_arch`, `gemv_mq4g256_lloyd_residual_for_arch`, `gemv_mq4g256v2_for_arch`, `gemv_mq4g256v2_moe_gate_up_k8_indexed_batched_src`, `gemv_mq4g256v2_residual_for_arch`, `gemv_mq5g256v2_for_arch`, `gemv_mq5g256v2_residual_for_arch`, `gemv_mq6g256v2_for_arch`, `gemv_mq6g256v2_residual_for_arch`, `kv_slot_desc_paged_source`, `kv_slot_desc_source`, `kv_slot_givens4_paged_source`, `moe_scatter_fused_k8_src`, `mq2g256_lloyd_moe_grouped_wmma_source`, `mq3g256_lloyd_moe_grouped_wmma_source`, `mq4g256v2_moe_grouped_wmma_source`, `mq6g256v2_moe_grouped_wmma_source`, `qwen36_27b_fa_prep_gfx1100_src`, `qwen36_27b_fa_prep_gfx1201_src`, `scheduler_profile_for_module` - [`src/kv_slots.rs`](src/kv_slots.rs): `KvSlotDesc`, `MASKED`, `R9700_VRAM_BYTES`, `VmmKvSlotDesc`, `as_bytes`, `build_arena`, `build_asym3_k_arena`, `build_tiles`, `check_position`, `half_from_f32`, `legacy_cap`, `mem_available_bytes`, `preflight_alloc`, `total_rows`, `validate_vmm_rows` - [`src/lane_rows_pm.rs`](src/lane_rows_pm.rs): `LANE_ROWS_PM_MAX`, `gemv_mq6g256v2_x4_xbatch_pm`, `gemv_mq6g256v2_xbatch_pm`, `gemv_q8_0_k320_xbatch_pm`, `gemv_q8_0_k8_xbatch_pm`, `lane_rows_pm_enabled` - [`src/lib.rs`](src/lib.rs): `arch_caps`, `attention`, `calib_force_bf16`, `cdna`, `code_object`, `compile_jobs`, `compiler`, `dflash_draft_fusion`, `dflash_gdn_pre`, `dflash_gdn_replay`, `dflash_hidden_scatter`, `dflash_state_copy`, `dispatch`, `embedding`, `feature_flags`, `flash_attn_ck`, `flux_fused`, `fn_gdn_dense`, `gap_timing`, `gemm`, `gemm_fp8_verify`, `gemv`, `graph`, `grouped_ops`, `hc_row_fold`, `hip_bridge`, `kernel_pack`, `kernel_registry`, `kernels`, `kv_slots`, `lane_rows_pm`, `moe`, `mq_f16_producers`, `mq_f16_residual_producers`, `norm`, `page_pool`, `pm_decode_twins`, `pm_xbatch`, `pool`, `prefill_invariant`, `profile`, `profile_rocprof`, `profiler`, `qwen35_fa_batch`, `railgun_check`, `rdna`, `replay`, `sampling`, `scratch`, `select_regrid`, `slot_pool`, `tensor_ops`, `text_encoder`, `trunk_mask`, `uma_memory`, `vae`, `verify_twins` diff --git a/docs/env-vars.md b/docs/env-vars.md index 07bd53b8d0..9f71f8fc99 100644 --- a/docs/env-vars.md +++ b/docs/env-vars.md @@ -264,6 +264,8 @@ Read only by the Qwen4 carrier and its kernels; no other model reads them. | `HIPFIRE_FP8_DECODE_ATTN_GQA` | Exact-gfx1201 native-fp8 **decode** attention (head_dim 256, GQA group 6, tile 128, no output gate: H2): GQA-shared flash tile (one 256-thread workgroup per kv head and 128-key tile serves its six q heads, so K/V are read once) + head-dim-split reduce (`attention_flash_fp8_e4m3_tile_gqa_gfx1201` / `attention_flash_reduce_dsplit_gfx1201`); byte-identical partials and output — default ON; `=0` restores `attention_flash_fp8_e4m3_tile` + `attention_flash_q8_0_reduce` | | `HIPFIRE_GFX1100_DECODE_ATTN_GQA` | Exact-gfx1100 Q8_0 **decode** attention with the gated AWQ MQ-rotating epilogue (head_dim 256, GQA group 6, tile 128, full causal: H2): GQA-shared flash tile (`attention_flash_q8_0_tile_gqa_gfx1100`, one 256-thread workgroup per kv head and 128-key tile, K/V read once for the six q heads, every load issued at entry) + load-ahead reduce/gate/rotate (`attention_flash_q8_0_reduce_gated_mq_rotate_awq_dec_gfx1100`, one 1024-thread workgroup per head); byte-identical partials and output — default ON; `=0` restores `attention_flash_q8_0_tile` + `attention_flash_q8_0_reduce_gated_mq_rotate_awq_gfx1100` | | `HIPFIRE_GFX1151_Q8_DECODE_ATTN_GQA` | Exact-gfx1151 Q8_0 **decode** attention (head_dim 256, GQA group 6, tile 128, full causal, no output gate: H2): GQA-shared flash tile (one 256-thread workgroup per kv head and 128-key tile serves its six q heads, so K/V are read once) + head-dim-split reduce (`attention_flash_q8_0_tile_gqa_gfx1151` / `attention_flash_reduce_dsplit_gfx1151`); byte-identical partials and output — default ON; `=0` restores `attention_flash_q8_0_tile` + `attention_flash_q8_0_reduce` | +| `HIPFIRE_GFX1151_Q8_DECODE_ATTN_GQA_GEMMA` | Exact-gfx1151 Q8_0 **decode** attention at the Gemma 4 shapes (head_dim 256 / GQA 2, head_dim 512 / GQA 8, tile 128, any window): GQA-shared 8-wave tile (`attention_flash_q8_0_tile_gqa_d256g2_gfx1151` / `_d512g8_gfx1151`) feeding the unchanged `attention_flash_q8_0_reduce`; not bit-identical to the reference tile (summation order) — default ON; `=0` restores `attention_flash_q8_0_tile` | +| `HIPFIRE_Q8_PREFILL_GQA_WMMA` | Exact-gfx1151 Q8_0 **batched prefill** attention (head_dim 256/512, GQA group 1/2/4/8/16, any window): GQA-shared WMMA flash kernel (`attention_q8_0_prefill_gqa_wmma_d256` / `_d512`), F16 operands, F32 softmax/accumulators — default ON; `=0` restores the windowed tile / batched scan kernels | | `HIPFIRE_CALIB_BF16` | Calibration-only: keep native-BF16 teachers in BF16 (`kernel.calib_force_bf16`, default off; shipped inference unaffected) | | `HIPFIRE_GFX12_MQ4V2_FP8_GATEUP` / `_RESID` / `_QKVZA` / `_QKV` | gfx1201 FP8-WMMA MQ4v2 prefill route — default ON on exact gfx1201 (widened prefill chunk 4096 via `prefill.chunk_rows`); `=0` on any one moves that family (its GEMM and its `HIPFIRE_GFX12_FP8_STREAM` producer) to the F16 path (chunk 384). `=1` forces on; launchers stay exact-gfx1201-only, so other arches are unchanged | | `HIPFIRE_GFX12_MQ4V2_FP8_SLABS` | Two-slab S2BT8 FP8 symbols by default; `=1` selects the single-slab symbols | @@ -587,7 +589,7 @@ Presence in the inventory means the token appears in source; it does **not** mea **Generation method:** token scan over tracked `*.rs`, `*.py`, and `*.sh` (`scripts/check-lifecycle.py --write`). **Columns:** variable; up to two lexical source paths; lifecycle status (see [Lifecycle status](#lifecycle-status)). -**Count:** 1489 +**Count:** 1476 | Variable | Example source path(s) | Lifecycle | |---|---|---| @@ -639,7 +641,6 @@ Presence in the inventory means the token appears in source; it does **not** mea | `HIPFIRE_BASELINE_PROMPT` | scripts/baseline_quant_smoke.sh | harness | | `HIPFIRE_BASELINE_PROMPT_MODE` | scripts/baseline_quant_smoke.sh | harness | | `HIPFIRE_BASELINE_WIDE_MARGIN` | scripts/baseline_quant_smoke.sh | harness | -| `HIPFIRE_BATCHED_PREFILL` | crates/hipfire-arch-gemma4/examples/prefill_parity_gemma4.rs, crates/hipfire-arch-gemma4/src/lowered.rs | developer | | `HIPFIRE_BENCH_AB` | benchmarks/scripts/bench_dflash_27b_gfx906.sh | harness | | `HIPFIRE_BENCH_MAX` | scripts/adaptive_b_bench.sh, scripts/qwen36_bench.sh | harness | | `HIPFIRE_BENCH_N` | crates/rdna-compute/examples/bench_indexed_moe_keystone.rs | harness | @@ -1102,18 +1103,8 @@ Presence in the inventory means the token appears in source; it does **not** mea | `HIPFIRE_GDN_STATE_SRC` | crates/rdna-compute/src/dflash_gdn_replay.rs | developer | | `HIPFIRE_GDN_TILE_ROWS` | crates/rdna-compute/src/dflash_gdn_replay.rs, crates/rdna-compute/src/kernels.rs | developer | | `HIPFIRE_GDN_WAVES_PER_BLOCK` | crates/rdna-compute/src/kernels.rs | developer | -| `HIPFIRE_GEMMA4_ATTN_VERIFY` | crates/hipfire-arch-gemma4/src/lowered.rs | developer | -| `HIPFIRE_GEMMA4_BASELINE_ATTN` | crates/hipfire-arch-gemma4/src/forward.rs | developer | | `HIPFIRE_GEMMA4_BATCHED_EMBEDDING_PREFILL` | crates/hipfire-config/src/lib.rs, crates/rdna-compute/src/feature_flags.rs | experimental | -| `HIPFIRE_GEMMA4_DUMP` | crates/hipfire-arch-gemma4/examples/prefill_parity_gemma4.rs, crates/hipfire-arch-gemma4/src/lowered.rs | developer | -| `HIPFIRE_GEMMA4_EAGLE` | crates/hipfire-arch-gemma4/src/forward.rs, crates/hipfire-generate/src/dense.rs | developer | -| `HIPFIRE_GEMMA4_FUSED_ATTN_NORM` | crates/hipfire-arch-gemma4/src/forward.rs | developer | -| `HIPFIRE_GEMMA4_FUSED_FFN` | crates/hipfire-arch-gemma4/src/forward.rs | developer | -| `HIPFIRE_GEMMA4_FUSED_POSTNORM` | crates/hipfire-arch-gemma4/src/forward.rs | developer | -| `HIPFIRE_GEMMA4_FUSED_PROJ` | crates/hipfire-arch-gemma4/src/forward.rs | developer | -| `HIPFIRE_GEMMA4_FUSED_QK` | crates/hipfire-arch-gemma4/src/forward.rs | developer | -| `HIPFIRE_GEMMA4_FUSED_QK_ROPE` | crates/hipfire-arch-gemma4/src/forward.rs | developer | -| `HIPFIRE_GEMMA4_GEMM_VERIFY` | crates/hipfire-arch-gemma4/src/lowered.rs | developer | +| `HIPFIRE_GEMMA4_EAGLE` | crates/hipfire-dispatch/src/pipeline/sandwich.rs, crates/hipfire-generate/src/dense.rs | developer | | `HIPFIRE_GEMMA4_GRAPH` | crates/hipfire-arch-gemma4/examples/infer_gemma4.rs, crates/hipfire-arch-gemma4/src/forward.rs | developer | | `HIPFIRE_GEMMA4_LOGIT_TRACE_DIR` | crates/hipfire-generate/src/dense.rs, scripts/diag-gemma4-logit-routes.sh | developer | | `HIPFIRE_GEMMA4_LOGIT_TRACE_FULL_STEPS` | crates/hipfire-generate/src/dense.rs, scripts/diag-gemma4-logit-routes.sh | developer | @@ -1123,7 +1114,7 @@ Presence in the inventory means the token appears in source; it does **not** mea | `HIPFIRE_GEMMA4_PLE_ACTIVATION_FUSED_PREFILL` | crates/hipfire-config/src/lib.rs, crates/rdna-compute/src/feature_flags.rs | experimental | | `HIPFIRE_GEMMA4_PLE_BATCHED_PREFILL` | crates/hipfire-config/src/lib.rs, crates/rdna-compute/src/feature_flags.rs | experimental | | `HIPFIRE_GEMMA4_PLE_BRANCH_BATCHED_PREFILL` | crates/hipfire-config/src/lib.rs, crates/rdna-compute/src/feature_flags.rs | experimental | -| `HIPFIRE_GEMMA4_PREFILL_BATCH` | crates/hipfire-generate/src/dense.rs, scripts/eval_gemma4_eseries.py | developer | +| `HIPFIRE_GEMMA4_PREFILL_BATCH` | crates/hipfire-arch-gemma4/src/bundle.rs, scripts/eval_gemma4_eseries.py | developer | | `HIPFIRE_GEMMA4_Q8_FUSED_PREFILL` | crates/hipfire-config/src/lib.rs, crates/rdna-compute/src/feature_flags.rs | experimental | | `HIPFIRE_GEMM_DUMP` | crates/hipfire-config/src/lib.rs, crates/rdna-compute/src/feature_flags.rs | experimental | | `HIPFIRE_GEMV_DP4A` | crates/hipfire-config/src/lib.rs, crates/rdna-compute/src/feature_flags.rs | stable | @@ -1198,6 +1189,7 @@ Presence in the inventory means the token appears in source; it does **not** mea | `HIPFIRE_GFX1151_PM4_INTERLEAVE` | crates/hipfire-config/src/lib.rs, crates/rdna-compute/src/replay.rs | experimental | | `HIPFIRE_GFX1151_PM4_RESOURCE_LIMITS` | crates/hipfire-config/src/lib.rs, crates/rdna-compute/src/replay.rs | experimental | | `HIPFIRE_GFX1151_Q8_DECODE_ATTN_GQA` | crates/rdna-compute/src/attention.rs | developer | +| `HIPFIRE_GFX1151_Q8_DECODE_ATTN_GQA_GEMMA` | crates/rdna-compute/examples/gemma4_kernel_f64_check.rs, crates/rdna-compute/src/attention.rs | developer | | `HIPFIRE_GFX1151_QKVZA_ALL_BUFFER_CPOL` | crates/rdna-compute/src/gemm.rs | developer | | `HIPFIRE_GFX1151_QKVZA_HYBRID_BUFFER` | crates/rdna-compute/src/kernels.rs | developer | | `HIPFIRE_GFX1151_QKVZA_K2048_HOIST` | crates/rdna-compute/src/gemm.rs | developer | @@ -1315,7 +1307,7 @@ Presence in the inventory means the token appears in source; it does **not** mea | `HIPFIRE_GPU_TOPK` | crates/saddle-lab/examples/infer_qwen35.rs | harness | | `HIPFIRE_GQA_CHUNK` | crates/rdna-compute/src/attention.rs | developer | | `HIPFIRE_GQA_FUSED` | crates/hipfire-arch-qwen2/src/qwen2.rs, crates/hipfire-dispatch/src/families/kv_tier.rs | developer | -| `HIPFIRE_GRAPH` | benchmarks/scripts/bench_pp_gfx906.sh, crates/hipfire-arch-gemma4/src/lowered.rs | experimental | +| `HIPFIRE_GRAPH` | benchmarks/scripts/bench_pp_gfx906.sh, crates/hipfire-arch-maple/examples/maple_decode_profile.rs | experimental | | `HIPFIRE_GRAPH_MOE` | crates/hipfire-arch-qwen35/src/qwen35/forward.rs, crates/hipfire-config/src/lib.rs | experimental | | `HIPFIRE_GRAPH_PREFILL` | crates/hipfire-runtime/examples/bench_qwen35_mq4.rs | harness | | `HIPFIRE_GRID_TOKFAST` | crates/rdna-compute/src/kernels.rs | developer | @@ -1388,7 +1380,6 @@ Presence in the inventory means the token appears in source; it does **not** mea | `HIPFIRE_KV_FP8_E4M3` | crates/rdna-compute/src/kernels.rs | developer | | `HIPFIRE_KV_MODE` | benchmarks/scripts/bench_pp_gfx906.sh, crates/hipfire-arch-qwen35/src/serve_engine.rs | stable; `asymN`/`turbo*` values deprecated, removal 0.5.0 | | `HIPFIRE_KV_PHYSICAL_CAP` | crates/hipfire-runtime/src/loader_api.rs | developer | -| `HIPFIRE_KV_SEQ` | crates/hipfire-arch-gemma4/src/carrier.rs, crates/hipfire-arch-gemma4/src/lowered.rs | developer | | `HIPFIRE_KV_SLOT_PAGED` | crates/rdna-compute/src/kernel_registry.rs, crates/rdna-compute/src/kernels.rs | developer | | `HIPFIRE_KV_V` | crates/hipfire-arch-qwen35/src/carrier.rs, crates/hipfire-loader/src/admission.rs | developer | | `HIPFIRE_LABEL` | docs/investigations/evidence/ds4-mi300x-cdna-test-fail/raw/a1-m0/04-profile-feed.py, docs/investigations/evidence/ds4-mi300x-cdna-test-fail/raw/a1-m0/04-run-profile-direct.sh | harness | @@ -1448,8 +1439,6 @@ Presence in the inventory means the token appears in source; it does **not** mea | `HIPFIRE_MODEL_STORE` | scripts/baseline_quant_smoke.sh | harness | | `HIPFIRE_MODULE_LOAD_AUDIT` | crates/hip-bridge/src/ffi.rs, crates/rdna-compute/src/scratch.rs | developer | | `HIPFIRE_MOE_AWQ` | crates/hipfire-arch-qwen35/src/qwen35/load.rs | developer | -| `HIPFIRE_MOE_BUCKETED` | crates/hipfire-arch-gemma4/src/lowered.rs | developer | -| `HIPFIRE_MOE_BYPASS` | crates/hipfire-arch-gemma4/src/lowered.rs | developer | | `HIPFIRE_MOE_CODEBOOK_BATCHED` | crates/hipfire-arch-qwen35/src/qwen35/prefill.rs, crates/hipfire-dispatch/src/families/moe.rs | developer | | `HIPFIRE_MOE_COMBINE_NEXT_RMS` | crates/hipfire-arch-qwen35/src/qwen35/forward.rs | developer | | `HIPFIRE_MOE_COMBINE_NEXT_RMS_RENORM` | crates/hipfire-arch-qwen35/src/qwen35/forward.rs | developer | @@ -1589,7 +1578,7 @@ Presence in the inventory means the token appears in source; it does **not** mea | `HIPFIRE_OOM_GUARD` | crates/hipfire-config/src/lib.rs, crates/rdna-compute/src/kv_slots.rs | stable | | `HIPFIRE_ORACLE_DIVERGE` | crates/hipfire-arch-qwen35/tests/route_oracle_mesh.rs | harness | | `HIPFIRE_ORACLE_DUMP` | crates/saddle-lab/examples/tmp_gemm_v2_oracle.rs | harness | -| `HIPFIRE_ORACLE_KV_F32` | crates/hipfire-runtime/examples/gemma4_oracle.rs | harness | +| `HIPFIRE_ORACLE_KV_Q8` | crates/hipfire-runtime/examples/gemma4_oracle.rs | harness | | `HIPFIRE_ORACLE_MAX` | scripts/seed_oracle_collect.sh | harness | | `HIPFIRE_ORACLE_STATE_FP32` | crates/hipfire-arch-qwen35/tests/route_oracle_mesh.rs | harness | | `HIPFIRE_ORNITH15_MODEL` | scripts/coherence-gate-ornith15.sh | deprecated | @@ -1656,6 +1645,7 @@ Presence in the inventory means the token appears in source; it does **not** mea | `HIPFIRE_Q8_BATCHED_LEGACY` | crates/hipfire-config/src/lib.rs, crates/rdna-compute/src/feature_flags.rs | experimental | | `HIPFIRE_Q8_CLASSES` | crates/hipfire-quantize/src/diagnostics.rs, crates/hipfire-quantize/src/model_filter.rs | developer | | `HIPFIRE_Q8_FLASH_TILE` | crates/hipfire-arch-maple/src/maple.rs, crates/rdna-compute/src/attention.rs | developer | +| `HIPFIRE_Q8_PREFILL_GQA_WMMA` | crates/rdna-compute/src/attention.rs | developer | | `HIPFIRE_Q8_PREFILL_WMMA` | crates/hipfire-arch-qwen35/src/qwen35/prefill.rs | developer | | `HIPFIRE_QA_KV_MODES` | crates/saddle-lab/examples/test_inferenceQA.rs | harness | | `HIPFIRE_QCAL_BREAKDOWN` | crates/hipfire-runtime/tests/hf_tokenizer_encode_bench.rs | harness | @@ -2071,7 +2061,6 @@ Presence in the inventory means the token appears in source; it does **not** mea | `HIPFIRE_WMMA_BATCH_TILES` | crates/hipfire-arch-qwen35/examples/cb_vmm_state_oracle.rs, crates/hipfire-arch-qwen35/src/qwen35/prefill_multi.rs | developer | | `HIPFIRE_WMMA_FA` | benchmarks/results/wmma-fa-probe-gfx1100.sh, benchmarks/results/wmma-fa-probe.sh | developer | | `HIPFIRE_WMMA_FA_MIN_BATCH` | crates/rdna-compute/src/attention.rs | developer | -| `HIPFIRE_WMMA_PREFILL` | crates/hipfire-arch-gemma4/examples/prefill_parity_gemma4.rs, crates/hipfire-arch-gemma4/src/lowered.rs | developer | | `HIPFIRE_WO_MMQ` | crates/hipfire-config/src/lib.rs, crates/rdna-compute/src/feature_flags.rs | experimental | | `HIPFIRE_WO_WMMA_VARIANT` | crates/hipfire-config/src/lib.rs, crates/hipfire-runtime/examples/test_wmma_correctness.rs | experimental | | `HIPFIRE_XDNA_A` | crates/hipfire-xdna/src/lib.rs | developer |