pub(crate) const CUDA_FALLBACK_LOG_PREFIX: &str = "warning: GPU (CUDA) path rejected";
pub(crate) const WGPU_FALLBACK_LOG_PREFIX: &str = "warning: GPU (wgpu) path rejected";
pub(crate) fn cpu_vs_gpu_cosine_similarity(a: &[f32], b: &[f32]) -> f32 {
if a.len() != b.len() || a.is_empty() {
return 0.0;
}
let mut dot: f64 = 0.0;
let mut norm_a: f64 = 0.0;
let mut norm_b: f64 = 0.0;
for (x, y) in a.iter().zip(b.iter()) {
let x = f64::from(*x);
let y = f64::from(*y);
dot += x * y;
norm_a += x * x;
norm_b += y * y;
}
let denom = norm_a.sqrt() * norm_b.sqrt();
if denom < 1e-12 {
0.0
} else {
(dot / denom) as f32
}
}
pub const WGPU_SAMPLING_NOTICE: &str = "[wgpu: the wgpu decoder is greedy-only; \
sampling (--temperature > 0 with --top-k != 1) runs on the CPU (#3760)]";
#[cfg(feature = "gpu")]
fn wgpu_can_serve(temperature: f32, top_k: usize) -> bool {
if crate::sampling::is_greedy(temperature, top_k) {
return true;
}
eprintln!("{WGPU_SAMPLING_NOTICE}");
false
}
#[cfg(feature = "gpu")]
#[must_use]
pub(crate) fn wgpu_fallback_allowed(
no_gpu: bool,
accel_forced: bool,
has_legacy_quant: bool,
) -> bool {
!no_gpu && accel_forced && !has_legacy_quant
}
#[cfg_attr(not(feature = "gpu"), allow(dead_code))]
fn wgpu_greedy_decode<E>(
prompt: &[u32],
max_tokens: usize,
stop_tokens: &[u32],
mut forward: impl FnMut(u32, usize) -> std::result::Result<Vec<f32>, E>,
mut pick: impl FnMut(&[f32]) -> u32,
) -> std::result::Result<Vec<u32>, E> {
let mut output_tokens = prompt.to_vec();
let Some((_, head)) = prompt.split_last().filter(|_| max_tokens > 0) else {
return Ok(output_tokens);
};
for (position, &token) in head.iter().enumerate() {
forward(token, position)?;
}
for _ in 0..max_tokens {
let position = output_tokens.len() - 1;
let hidden = forward(output_tokens[position], position)?;
let next = pick(&hidden);
output_tokens.push(next);
if crate::sampling::is_stop(next, stop_tokens) {
break;
}
}
Ok(output_tokens)
}
#[cfg(feature = "gpu")]
fn try_wgpu_generate(
model: &crate::gguf::OwnedQuantizedModel,
input_tokens: &[u32],
gen_config: &crate::gguf::QuantizedGenerateConfig,
verbose: bool,
) -> Result<(Vec<u32>, bool)> {
use crate::gpu::adapters::wgpu_adapter;
if !trueno::backends::gpu::GpuDevice::is_available() {
return Err(RealizarError::InferenceError("wgpu not available".into()));
}
let gpu = trueno::backends::gpu::GpuDevice::new()
.map_err(|e| RealizarError::InferenceError(format!("wgpu init: {e}")))?;
let _ = verbose;
eprintln!("Backend: wgpu (Vulkan)");
let config = model.config();
let hidden_dim = config.hidden_dim;
let num_layers = config.num_layers;
let num_heads = config.num_heads;
let num_kv_heads = config.num_kv_heads;
let head_dim = hidden_dim / num_heads;
let intermediate_dim = config.intermediate_dim;
let vocab_size = config.vocab_size;
let eps = config.eps;
let kv_dim = num_kv_heads * head_dim;
let mut fwd = trueno::backends::gpu::WgslForwardPass::new(
gpu.device, gpu.queue,
hidden_dim, num_heads, num_kv_heads, head_dim, intermediate_dim,
);
fwd.set_rms_norm_eps(eps);
let raw_q4k = wgpu_adapter::raw_q4k_weights(model);
let q4k_names: std::collections::HashSet<String> =
raw_q4k.iter().map(|(n, _, _, _)| n.clone()).collect();
for (name, data, _rows, _cols) in &raw_q4k {
fwd.upload_q4k_weight(name, data);
}
let weights = wgpu_adapter::dequant_model_weights_except(model, &q4k_names)?;
for (name, data, _rows, _cols) in &weights {
fwd.upload_weight(name, data);
}
let output_norm = model.output_norm_weight();
let lm_head_f32: Vec<f32> = weights.iter()
.find(|(n, _, _, _)| n == "lm_head")
.map(|(_, d, _, _)| d.clone())
.unwrap_or_default();
let max_seq = gen_config.max_tokens + input_tokens.len() + 16;
let mut kv_caches: Vec<(Vec<f32>, Vec<f32>)> = (0..num_layers)
.map(|_| (Vec::with_capacity(max_seq * kv_dim), Vec::with_capacity(max_seq * kv_dim)))
.collect();
{
const MULTI_STEP_PROBE_DEFAULT: usize = 3;
let multi_step_probe: usize = std::env::var("APR_WGPU_PARITY_STEPS")
.ok()
.and_then(|s| s.parse::<usize>().ok())
.filter(|&n| (1..=16).contains(&n))
.unwrap_or(MULTI_STEP_PROBE_DEFAULT);
let probe_max_seq = multi_step_probe + 1;
let mut cpu_cache = crate::gguf::OwnedQuantizedKVCache::from_config(&config, probe_max_seq);
let mut probe_kv_caches: Vec<(Vec<f32>, Vec<f32>)> = (0..num_layers)
.map(|_| (Vec::with_capacity(probe_max_seq * kv_dim), Vec::with_capacity(probe_max_seq * kv_dim)))
.collect();
let mut probe_token = *input_tokens.first().unwrap_or(&0);
for probe_step in 0..multi_step_probe {
let cpu_logits = match model.forward_single_with_cache(probe_token, &mut cpu_cache, probe_step) {
Ok(l) => l,
Err(e) => {
eprintln!(
"{}, attempting fallback: CPU probe step {} forward failed: {}",
WGPU_FALLBACK_LOG_PREFIX, probe_step, e
);
return Err(RealizarError::InferenceError(format!("wgpu parity gate: CPU probe step {probe_step} failed: {e}")));
}
};
let mut hidden = model.embed(&[probe_token]);
for layer_idx in 0..num_layers {
let prefix = format!("layer.{layer_idx}");
let (ref mut kv_k, ref mut kv_v) = probe_kv_caches[layer_idx];
if let Err(e) = fwd.forward_layer(&mut hidden, &prefix, probe_step, kv_k, kv_v) {
eprintln!(
"{}, attempting fallback: wgpu probe step {} layer {} failed: {}",
WGPU_FALLBACK_LOG_PREFIX, probe_step, layer_idx, e
);
return Err(RealizarError::InferenceError(format!("wgpu parity gate: step {probe_step} layer {layer_idx} failed: {e}")));
}
}
let sq_sum: f32 = hidden.iter().map(|x| x * x).sum();
let rms = (sq_sum / hidden.len() as f32 + eps).sqrt();
let normed: Vec<f32> = hidden
.iter()
.zip(output_norm.iter())
.map(|(x, g)| (x / rms) * g)
.collect();
let mut wgpu_logits = vec![0.0_f32; vocab_size];
for i in 0..vocab_size {
let row = &lm_head_f32[i * hidden_dim..(i + 1) * hidden_dim];
wgpu_logits[i] = row.iter().zip(normed.iter()).map(|(w, x)| w * x).sum();
}
let cos = cpu_vs_gpu_cosine_similarity(&cpu_logits, &wgpu_logits);
if !(cos.is_finite() && cos >= 0.99) {
eprintln!(
"{}, attempting fallback: cosine vs CPU = {:.6} (< 0.99) at step {}/{}",
WGPU_FALLBACK_LOG_PREFIX, cos, probe_step + 1, multi_step_probe
);
return Err(RealizarError::InferenceError(format!(
"wgpu parity gate: cosine={cos:.6} < 0.99 at step {}/{}",
probe_step + 1, multi_step_probe
)));
}
let mut best_idx: u32 = 0;
let mut best_val = f32::NEG_INFINITY;
for (i, &v) in cpu_logits.iter().enumerate() {
if v > best_val {
best_val = v;
best_idx = i as u32;
}
}
probe_token = best_idx;
}
}
let output_tokens = wgpu_greedy_decode::<RealizarError>(
input_tokens,
gen_config.max_tokens,
&gen_config.stop_tokens,
|token_id, position| {
let mut hidden = model.embed(&[token_id]);
for (layer_idx, (kv_k, kv_v)) in kv_caches.iter_mut().enumerate() {
let prefix = format!("layer.{layer_idx}");
fwd.forward_layer(&mut hidden, &prefix, position, kv_k, kv_v)
.map_err(|e| RealizarError::InferenceError(format!("wgpu layer {layer_idx}: {e}")))?;
}
Ok(hidden)
},
|hidden| {
let sq_sum: f32 = hidden.iter().map(|x| x * x).sum();
let rms = (sq_sum / hidden.len() as f32 + eps).sqrt();
let normed: Vec<f32> = hidden.iter().zip(output_norm.iter())
.map(|(x, g)| (x / rms) * g)
.collect();
let mut best_idx = 0u32;
let mut best_val = f32::NEG_INFINITY;
for i in 0..vocab_size {
let row = &lm_head_f32[i * hidden_dim..(i + 1) * hidden_dim];
let logit: f32 = row.iter().zip(normed.iter()).map(|(w, x)| w * x).sum();
if logit > best_val {
best_val = logit;
best_idx = i as u32;
}
}
best_idx
},
)?;
Ok((output_tokens, true)) }
#[cfg(feature = "cuda")]
fn try_gguf_gpu_generate(
model: crate::gguf::OwnedQuantizedModel,
input_tokens: &[u32],
gen_config: &crate::gguf::QuantizedGenerateConfig,
verbose: bool,
) -> std::result::Result<Result<(Vec<u32>, bool)>, Box<crate::gguf::OwnedQuantizedModel>> {
use crate::gguf::OwnedQuantizedModelCuda;
let kv_len = device_kv_len(&model, input_tokens.len(), gen_config.max_tokens);
let mut cuda_model = match OwnedQuantizedModelCuda::with_max_seq_len(model, 0, kv_len) {
Ok(m) => m,
Err(e) => {
if verbose {
eprintln!("Backend: CPU (GPU unavailable: {})", e);
}
return Err(Box::new(e.into_model()));
},
};
if verbose {
eprintln!(
"Backend: GPU ({}, {} MB VRAM)",
cuda_model.device_name(),
cuda_model.vram_mb()
);
}
match validate_gpu_first_token(&mut cuda_model, gen_config, input_tokens) {
F2Outcome::Mismatch => {
return Err(Box::new(cuda_model.into_model()));
},
F2Outcome::NotMeasured { reason } => {
eprintln!("[GH-480] F2 validation NOT MEASURED — {reason}. GPU output is UNVALIDATED (#3973)");
},
F2Outcome::Validated { .. } => {},
}
mark_generation_start(); if gen_config.trace {
let result = cuda_model
.generate_gpu_resident(input_tokens, gen_config)
.map(|tokens| (tokens, true))
.map_err(|e| RealizarError::InferenceError(format!("GPU generation failed: {}", e)));
return Ok(result);
}
let mut session = crate::gguf::dense_session::DenseSession::new(
crate::gguf::dense_session::DenseForward::cuda(cuda_model),
);
Ok(crate::gguf::dense_session::dense_turn(
&mut session,
input_tokens,
gen_config,
))
}
#[cfg(feature = "cuda")]
fn device_kv_len(
model: &crate::gguf::OwnedQuantizedModel,
prompt_len: usize,
max_tokens: usize,
) -> usize {
prompt_len
.saturating_add(max_tokens)
.min(model.config.context_length)
.max(2048)
}
#[allow(unused_variables)] fn run_gguf_generate(
model: crate::gguf::OwnedQuantizedModel,
input_tokens: &[u32],
gen_config: &crate::gguf::QuantizedGenerateConfig,
config: &InferenceConfig,
) -> Result<(Vec<u32>, bool, bool)> {
let canonical_arch =
crate::tensor_names::normalize_architecture(&model.config.architecture);
if canonical_arch == "qwen3_moe" {
return Err(RealizarError::UnsupportedOperation {
operation: "moe_forward_dispatch".to_string(),
reason: format!(
"Architecture '{}' (canonical 'qwen3_moe') uses Mixture-of-Experts FFN. \
Load step succeeded via QuantizedGGUFTransformer::from_gguf_for_moe (M32c.2) \
with all 4 contract-declared MoE tensors per layer present, but the \
forward dispatch is not yet wired to moe_forward_token in \
gpu/scheduler/moe_dispatch.rs. Tracked under contract qwen3-moe-forward-v1 \
(M32 staged plan: M32a/b/c.1/c.2 SHIPPED; M32c.2.1 forward-refusal \
IN PROGRESS; M32c.2.2 forward-wiring + M32d numerical parity PENDING). \
See contracts/qwen3-moe-forward-v1.yaml.",
model.config.architecture
),
});
}
let has_legacy_quant = model_has_legacy_quant(&model);
#[allow(unused_mut)]
let mut gpu_attempted = false;
#[cfg(feature = "cuda")]
let model = if !config.no_gpu && !has_legacy_quant {
gpu_attempted = true;
match try_gguf_gpu_generate(model, input_tokens, gen_config, config.verbose) {
Ok(result) => return result.map(|(t, u)| (t, u, true)),
Err(returned_model) => *returned_model, }
} else {
model
};
#[cfg(feature = "gpu")]
if wgpu_fallback_allowed(config.no_gpu, config.accel_forced, has_legacy_quant)
&& wgpu_can_serve(gen_config.temperature, gen_config.top_k)
{
gpu_attempted = true;
match try_wgpu_generate(&model, input_tokens, gen_config, config.verbose) {
Ok((t, u)) => return Ok((t, u, true)),
Err(e) => {
if config.verbose {
eprintln!("Backend: CPU (wgpu unavailable: {})", e);
}
}
}
}
log_cpu_backend(config.verbose, has_legacy_quant);
mark_generation_start(); let tokens = if gen_config.trace {
model.generate_with_cache(input_tokens, gen_config)
} else {
let mut session = crate::gguf::dense_session::DenseSession::new(
crate::gguf::dense_session::DenseForward::cpu(std::sync::Arc::new(model)),
);
crate::gguf::dense_session::dense_turn(&mut session, input_tokens, gen_config)
.map(|(tokens, _)| tokens)
}
.map_err(|e| RealizarError::InferenceError(format!("CPU generation failed: {}", e)))?;
Ok((tokens, false, gpu_attempted))
}
fn run_apr_inference(
config: &InferenceConfig,
prepared: &PreparedTokens,
) -> Result<InferenceResult> {
if config.verbose {
eprintln!("Loading APR model: {}", config.model_path.display());
}
let load_start = Instant::now();
let input_tokens = prepared.tokens();
let input_token_count = prepared.input_count();
#[cfg(feature = "cuda")]
if !config.no_gpu {
if let Some(result) =
try_apr_cuda_inference(config, input_tokens, input_token_count, load_start)
{
return result;
}
}
#[cfg(feature = "gpu")]
if wgpu_fallback_allowed(config.no_gpu, config.accel_forced, false)
&& wgpu_can_serve(config.temperature, config.top_k)
{
match try_apr_wgpu_inference(config, input_tokens, input_token_count, load_start) {
Some(Ok(result)) => return Ok(result),
Some(Err(e)) => {
if config.verbose {
eprintln!("Backend: CPU (wgpu failed: {})", e);
}
}
None => {
if config.verbose {
eprintln!("Backend: CPU (wgpu not available)");
}
}
}
}
run_apr_cpu_inference(config, input_tokens, input_token_count, load_start)
}
#[cfg(feature = "gpu")]
fn try_apr_wgpu_inference(
config: &InferenceConfig,
input_tokens: &[u32],
input_token_count: usize,
load_start: Instant,
) -> Option<Result<InferenceResult>> {
use crate::apr::MappedAprModel;
use crate::gpu::adapters::wgpu_adapter;
use trueno::backends::gpu::GpuDevice;
if !GpuDevice::is_available() {
return None;
}
let gpu = match GpuDevice::new() {
Ok(g) => g,
Err(e) => {
eprintln!("{}, attempting fallback: {}", WGPU_FALLBACK_LOG_PREFIX, e);
eprintln!("[GH-559] wgpu init failed: {}", e);
return None;
}
};
eprintln!("Backend: wgpu (Vulkan)");
let mapped = match MappedAprModel::from_path(&config.model_path) {
Ok(m) => m,
Err(_) => return None,
};
let model = match crate::gguf::OwnedQuantizedModel::from_apr(&mapped) {
Ok(m) => m,
Err(_) => return None,
};
let cfg = model.config();
let hidden_dim = cfg.hidden_dim;
let num_layers = cfg.num_layers;
let num_heads = cfg.num_heads;
let num_kv_heads = cfg.num_kv_heads;
let head_dim = hidden_dim / num_heads;
let intermediate_dim = cfg.intermediate_dim;
let vocab_size = cfg.vocab_size;
let eps = cfg.eps;
let kv_dim = num_kv_heads * head_dim;
let mut stop_toks: Vec<u32> = cfg.eos_token_id.into_iter().collect();
let extra = crate::infer::resolve_apr_stop_tokens(
cfg.eos_token_id, &[], &config.model_path,
);
for t in &extra {
if !stop_toks.contains(t) { stop_toks.push(*t); }
}
let mut gen_config = crate::gguf::QuantizedGenerateConfig {
max_tokens: config.max_tokens,
stop_tokens: stop_toks,
trace: config.trace,
..Default::default()
};
config.apply_sampling_to(&mut gen_config);
let weights = match wgpu_adapter::dequant_model_weights(&model) {
Ok(w) => w,
Err(e) => return Some(Err(e)),
};
let mut fwd = trueno::backends::gpu::WgslForwardPass::new(
gpu.device, gpu.queue,
hidden_dim, num_heads, num_kv_heads, head_dim, intermediate_dim,
);
fwd.set_rms_norm_eps(eps);
for (name, data, _rows, _cols) in &weights {
fwd.upload_weight(name, data);
}
let output_norm = model.output_norm_weight();
let lm_head_f32: Vec<f32> = weights.iter()
.find(|(n, _, _, _)| n == "lm_head")
.map(|(_, d, _, _)| d.clone())
.unwrap_or_default();
let max_seq = gen_config.max_tokens + input_tokens.len() + 16;
let mut kv_caches: Vec<(Vec<f32>, Vec<f32>)> = (0..num_layers)
.map(|_| (Vec::with_capacity(max_seq * kv_dim), Vec::with_capacity(max_seq * kv_dim)))
.collect();
{
const MULTI_STEP_PROBE_DEFAULT: usize = 3;
let multi_step_probe: usize = std::env::var("APR_WGPU_PARITY_STEPS")
.ok()
.and_then(|s| s.parse::<usize>().ok())
.filter(|&n| (1..=16).contains(&n))
.unwrap_or(MULTI_STEP_PROBE_DEFAULT);
let probe_max_seq = multi_step_probe + 1;
let mut cpu_cache = crate::gguf::OwnedQuantizedKVCache::from_config(cfg, probe_max_seq);
let mut probe_kv_caches: Vec<(Vec<f32>, Vec<f32>)> = (0..num_layers)
.map(|_| (Vec::with_capacity(probe_max_seq * kv_dim), Vec::with_capacity(probe_max_seq * kv_dim)))
.collect();
let mut probe_token = *input_tokens.first().unwrap_or(&0);
for step in 0..multi_step_probe {
let cpu_logits = match model.forward_single_with_cache(probe_token, &mut cpu_cache, step) {
Ok(l) => l,
Err(e) => {
eprintln!(
"{}, attempting fallback: CPU probe step {} forward failed: {}",
WGPU_FALLBACK_LOG_PREFIX, step, e
);
return None;
}
};
let mut hidden = model.embed(&[probe_token]);
for layer_idx in 0..num_layers {
let prefix = format!("layer.{layer_idx}");
let (ref mut kv_k, ref mut kv_v) = probe_kv_caches[layer_idx];
if let Err(e) = fwd.forward_layer(&mut hidden, &prefix, step, kv_k, kv_v) {
eprintln!(
"{}, attempting fallback: wgpu probe step {} layer {} failed: {}",
WGPU_FALLBACK_LOG_PREFIX, step, layer_idx, e
);
return None;
}
}
let sq_sum: f32 = hidden.iter().map(|x| x * x).sum();
let rms = (sq_sum / hidden.len() as f32 + eps).sqrt();
let normed: Vec<f32> = hidden
.iter()
.zip(output_norm.iter())
.map(|(x, g)| (x / rms) * g)
.collect();
let mut wgpu_logits = vec![0.0_f32; vocab_size];
for i in 0..vocab_size {
let row = &lm_head_f32[i * hidden_dim..(i + 1) * hidden_dim];
wgpu_logits[i] = row.iter().zip(normed.iter()).map(|(w, x)| w * x).sum();
}
let cos = cpu_vs_gpu_cosine_similarity(&cpu_logits, &wgpu_logits);
if !(cos.is_finite() && cos >= 0.99) {
eprintln!(
"{}, attempting fallback: cosine vs CPU = {:.6} (< 0.99) at step {}/{}",
WGPU_FALLBACK_LOG_PREFIX, cos, step + 1, multi_step_probe
);
return None;
}
let mut best_idx: u32 = 0;
let mut best_val = f32::NEG_INFINITY;
for (i, &v) in cpu_logits.iter().enumerate() {
if v > best_val {
best_val = v;
best_idx = i as u32;
}
}
probe_token = best_idx;
}
}
let model_load_ms = load_start.elapsed().as_millis() as f64;
let infer_start = Instant::now();
let decoded = wgpu_greedy_decode::<RealizarError>(
input_tokens,
gen_config.max_tokens,
&gen_config.stop_tokens,
|token_id, position| {
let mut hidden = model.embed(&[token_id]);
for (layer_idx, (kv_k, kv_v)) in kv_caches.iter_mut().enumerate() {
let prefix = format!("layer.{layer_idx}");
fwd.forward_layer(&mut hidden, &prefix, position, kv_k, kv_v)
.map_err(|e| RealizarError::InferenceError(format!("wgpu layer {layer_idx}: {e}")))?;
}
Ok(hidden)
},
|hidden| {
let sq_sum: f32 = hidden.iter().map(|x| x * x).sum();
let rms = (sq_sum / hidden.len() as f32 + eps).sqrt();
let normed: Vec<f32> = hidden.iter().zip(output_norm.iter())
.map(|(x, g)| (x / rms) * g)
.collect();
let mut best_idx = 0u32;
let mut best_val = f32::NEG_INFINITY;
for i in 0..vocab_size {
let row = &lm_head_f32[i * hidden_dim..(i + 1) * hidden_dim];
let logit: f32 = row.iter().zip(normed.iter()).map(|(w, x)| w * x).sum();
if logit > best_val {
best_val = logit;
best_idx = i as u32;
}
}
best_idx
},
);
let output_tokens = match decoded {
Ok(t) => t,
Err(e) => return Some(Err(e)),
};
let inference_ms = infer_start.elapsed().as_millis() as f64;
let tokens_generated = output_tokens.len() - input_token_count;
let text = crate::infer::decode_apr_tokens(&config.model_path, &output_tokens[input_token_count..]);
Some(Ok(InferenceResult {
text,
tokens: output_tokens,
input_token_count,
generated_token_count: tokens_generated,
inference_ms,
load_ms: model_load_ms,
tok_per_sec: if inference_ms > 0.0 { tokens_generated as f64 / (inference_ms / 1000.0) } else { 0.0 },
generation_ms: Some(inference_ms), format: "APR".to_string(),
used_gpu: true,
gpu_attempted: true,
}))
}
fn apr_arch_to_template_hint(apr_arch: &str, _model_name: &str) -> &'static str {
crate::tensor_names::normalize_architecture(apr_arch)
}
#[cfg(feature = "cuda")]
struct AprCudaModelInfo {
arch: String,
num_layers: usize,
vocab_size: usize,
hidden_dim: usize,
}
#[cfg(feature = "cuda")]
fn apr_cuda_decline_notice(model: &crate::gguf::OwnedQuantizedModel) -> Option<String> {
let qtype = model.first_gpu_unsupported_quant()?;
Some(format!(
"{CUDA_FALLBACK_LOG_PREFIX}: no verified GPU kernel for quantization type {qtype} — \
CUDA declined this model before wgpu or CPU was tried. Any backend error after this \
line is downstream of THIS decision, not its cause. Convert with \
`apr convert --quantize fp16` (type 1 is GPU-eligible), or run with --no-gpu (#3908)."
))
}
#[cfg(feature = "cuda")]
fn load_apr_cuda_model(
model_path: &std::path::Path,
verbose: bool,
) -> Option<(crate::gguf::OwnedQuantizedModelCuda, AprCudaModelInfo)> {
use crate::apr::MappedAprModel;
use crate::gguf::{OwnedQuantizedModel, OwnedQuantizedModelCuda};
let mapped = MappedAprModel::from_path(model_path).map_err(|e| {
if verbose { eprintln!("[APR-CUDA] MappedAprModel::from_path failed: {}", e); }
}).ok()?;
let model = OwnedQuantizedModel::from_apr(&mapped).map_err(|e| {
if verbose { eprintln!("[APR-CUDA] OwnedQuantizedModel::from_apr failed: {}", e); }
}).ok()?;
if let Some(notice) = apr_cuda_decline_notice(&model) {
eprintln!("{notice}");
return None;
}
let info = AprCudaModelInfo {
arch: model.config.architecture.clone(),
num_layers: model.config.num_layers,
vocab_size: model.config.vocab_size,
hidden_dim: model.config.hidden_dim,
};
let cuda_model = OwnedQuantizedModelCuda::with_max_seq_len(model, 0, 2048).map_err(|e| {
eprintln!("{}, attempting fallback: {}", CUDA_FALLBACK_LOG_PREFIX, e);
}).ok()?;
Some((cuda_model, info))
}
#[cfg(feature = "cuda")]
fn log_apr_cuda_info(
info: &AprCudaModelInfo,
cuda_model: &crate::gguf::OwnedQuantizedModelCuda,
load_ms: f64,
) {
eprintln!(
"Architecture: {} ({} layers, vocab_size={})",
info.arch, info.num_layers, info.vocab_size
);
let m = cuda_model.model();
eprintln!(
"Config: hidden_size={}, quant={}, backend=CUDA+KVCache, threads=1 (GPU)",
info.hidden_dim,
body_quant_label(&model_body_qtypes(m), m.lm_head_weight.qtype)
);
eprintln!("Model loaded in {:.1}ms", load_ms);
eprintln!(
"Backend: GPU ({}, {} MB VRAM)",
cuda_model.device_name(),
cuda_model.vram_mb()
);
}
#[cfg(feature = "cuda")]
fn try_apr_cuda_inference(
config: &InferenceConfig,
input_tokens: &[u32],
input_token_count: usize,
load_start: Instant,
) -> Option<Result<InferenceResult>> {
use crate::gguf::QuantizedGenerateConfig;
let (mut cuda_model, info) = load_apr_cuda_model(&config.model_path, config.verbose)?;
let load_ms = load_start.elapsed().as_secs_f64() * 1000.0;
if config.verbose {
log_apr_cuda_info(&info, &cuda_model, load_ms);
}
eprintln!("[GH-480-TRACE] try_apr_cuda_inference: model loaded OK, about to resolve stop tokens");
let stop_tokens = resolve_apr_stop_tokens(
cuda_model.model().config.eos_token_id,
&config.stop_tokens,
&config.model_path,
);
let mut gen_config = QuantizedGenerateConfig {
max_tokens: config.max_tokens,
stop_tokens,
trace: config.trace,
..Default::default()
};
config.apply_sampling_to(&mut gen_config);
eprintln!("[GH-480] F2 validation starting...");
let f2 = validate_gpu_first_token(&mut cuda_model, &gen_config, input_tokens);
eprintln!("{}", f2_status_line(&f2));
if f2 == F2Outcome::Mismatch {
return None;
}
let infer_start = Instant::now();
let tokens = match cuda_model.generate_gpu_resident(input_tokens, &gen_config) {
Ok(t) => t,
Err(e) => {
let msg = e.to_string();
eprintln!("[GH-480] generate_gpu_resident FAILED: {msg}");
if msg.contains("not supported") || msg.contains("architecture") {
if config.verbose {
eprintln!("[APR-CUDA] GPU-resident not supported, falling back to CPU: {msg}");
}
return None;
}
return Some(Err(RealizarError::InferenceError(format!(
"GPU generation failed: {}",
e
))));
},
};
let inference_ms = infer_start.elapsed().as_secs_f64() * 1000.0;
let generated_tokens = &tokens[input_token_count..];
let text = decode_apr_tokens(&config.model_path, generated_tokens);
let generated_token_count = generated_tokens.len();
Some(Ok(InferenceResult {
text,
tokens,
input_token_count,
generated_token_count,
inference_ms,
tok_per_sec: tok_per_sec(generated_token_count, inference_ms),
generation_ms: Some(inference_ms), load_ms,
format: "APR".to_string(),
used_gpu: true,
gpu_attempted: true,
}))
}
fn run_apr_cpu_inference(
config: &InferenceConfig,
input_tokens: &[u32],
input_token_count: usize,
load_start: Instant,
) -> Result<InferenceResult> {
run_apr_quantized_cpu_inference(config, input_tokens, input_token_count, load_start)
}
fn run_apr_quantized_cpu_inference(
config: &InferenceConfig,
input_tokens: &[u32],
input_token_count: usize,
load_start: Instant,
) -> Result<InferenceResult> {
use crate::apr::MappedAprModel;
use crate::gguf::{OwnedQuantizedModel, QuantizedGenerateConfig};
let mapped = MappedAprModel::from_path(&config.model_path)?;
let model = OwnedQuantizedModel::from_apr(&mapped)?;
let load_ms = load_start.elapsed().as_secs_f64() * 1000.0;
if config.verbose {
eprintln!(
"Architecture: {} ({} layers, vocab_size={})",
model.config.architecture, model.config.num_layers, model.config.vocab_size
);
eprintln!(
"Config: hidden_size={}, quant={} (OwnedQuantizedModel CPU), threads={}",
model.config.hidden_dim,
body_quant_label(&model_body_qtypes(&model), model.lm_head_weight.qtype),
rayon::current_num_threads()
);
eprintln!("Model loaded in {:.1}ms", load_ms);
eprintln!("Backend: CPU (OwnedQuantizedModel fallback for non-LLaMA arch)");
}
let stop_tokens = resolve_apr_stop_tokens(
model.config.eos_token_id,
&config.stop_tokens,
&config.model_path,
);
let mut gen_config = QuantizedGenerateConfig {
max_tokens: config.max_tokens,
stop_tokens,
trace: config.trace,
..Default::default()
};
config.apply_sampling_to(&mut gen_config);
let infer_start = Instant::now();
let tokens = model.generate_with_cache(input_tokens, &gen_config)?;
let inference_ms = infer_start.elapsed().as_secs_f64() * 1000.0;
let generated_tokens = &tokens[input_token_count..];
let text = decode_apr_tokens(&config.model_path, generated_tokens);
let generated_token_count = generated_tokens.len();
Ok(InferenceResult {
text,
tokens,
input_token_count,
generated_token_count,
inference_ms,
tok_per_sec: tok_per_sec(generated_token_count, inference_ms),
generation_ms: Some(inference_ms), load_ms,
format: "APR".to_string(),
used_gpu: false,
gpu_attempted: false,
})
}
fn resolve_apr_stop_tokens(
model_eos: Option<u32>,
caller_stop_tokens: &[u32],
model_path: &std::path::Path,
) -> Vec<u32> {
let mut tokens: Vec<u32> = model_eos.into_iter().collect();
for &t in caller_stop_tokens {
if !tokens.contains(&t) {
tokens.push(t);
}
}
if tokens.is_empty() {
tokens = resolve_stop_tokens_from_tokenizer(model_path);
}
tokens
}
fn resolve_stop_tokens_from_tokenizer(model_path: &std::path::Path) -> Vec<u32> {
let tokenizer = match crate::apr::AprV2Model::load_tokenizer(model_path) {
Some(t) => t,
None => return Vec::new(),
};
let mut tokens: Vec<u32> = tokenizer.eos_id.into_iter().collect();
for marker in &["<|im_end|>", "<|endoftext|>"] {
let id = tokenizer
.special_tokens
.get(*marker)
.or_else(|| tokenizer.token_to_id.get(*marker));
if let Some(&id) = id {
if !tokens.contains(&id) {
tokens.push(id);
}
}
}
tokens
}
fn decode_apr_tokens(model_path: &std::path::Path, tokens: &[u32]) -> String {
use crate::apr::AprV2Model;
let text = if let Some(tokenizer) = AprV2Model::load_tokenizer(model_path) {
tokenizer.decode(tokens)
} else if let Some(tokenizer) = find_fallback_tokenizer(model_path) {
tokenizer.decode(tokens)
} else {
format!("[{} tokens generated, tokenizer not found]", tokens.len())
};
clean_model_output(&text)
}
fn tok_per_sec(count: usize, ms: f64) -> f64 {
if ms > 0.0 {
count as f64 / (ms / 1000.0)
} else {
0.0
}
}
pub const SAFETENSORS_CUDA_SAMPLING_NOTICE: &str = "[safetensors: the CUDA decoder is greedy-only; \
sampling (--temperature > 0 with --top-k != 1) runs on the CPU (#3760)]";
fn run_safetensors_inference(
config: &InferenceConfig,
prepared: &PreparedTokens,
) -> Result<InferenceResult> {
if config.verbose {
eprintln!("Loading SafeTensors model: {}", config.model_path.display());
}
let input_tokens = prepared.tokens().to_vec();
let input_token_count = prepared.input_count();
#[cfg(feature = "cuda")]
if !config.no_gpu {
if crate::sampling::is_greedy(config.temperature, config.top_k) {
if let Some(result) =
try_safetensors_cuda_inference(config, &input_tokens, input_token_count)
{
return result;
}
} else {
eprintln!("{SAFETENSORS_CUDA_SAMPLING_NOTICE}");
}
}
run_safetensors_cpu_inference(config, &input_tokens, input_token_count)
}
#[cfg(feature = "cuda")]
fn try_safetensors_cuda_inference(
config: &InferenceConfig,
input_tokens: &[u32],
input_token_count: usize,
) -> Option<Result<InferenceResult>> {
use crate::safetensors_cuda::SafeTensorsCudaModel;
let load_start = Instant::now();
let mut cuda_model = match SafeTensorsCudaModel::load(&config.model_path, 0) {
Ok(m) => m,
Err(e) => {
if config.verbose {
eprintln!("Backend: CPU (GPU init failed: {})", e);
}
return None;
},
};
let load_ms = load_start.elapsed().as_secs_f64() * 1000.0;
if config.verbose {
eprintln!(
"Architecture: SafeTensors ({} layers, vocab_size={})",
cuda_model.config().num_layers,
cuda_model.config().vocab_size
);
eprintln!(
"Config: hidden_size={}, context_length={}, quant={}, threads=1 (GPU)",
cuda_model.config().hidden_dim,
cuda_model.config().context_length,
safetensors_quant_label(&config.model_path)
);
eprintln!("Model loaded in {:.1}ms", load_ms);
eprintln!(
"Backend: GPU ({}, {} MB VRAM)",
cuda_model.device_name(),
cuda_model.vram_mb()
);
}
let infer_start = Instant::now();
let eos_id = cuda_model.config().eos_token_id.unwrap_or(0);
let tokens = match cuda_model.generate(input_tokens, config.max_tokens, eos_id) {
Ok(t) => t,
Err(e) => {
return Some(Err(RealizarError::InferenceError(format!(
"GPU generation failed: {}",
e
))))
},
};
let inference_ms = infer_start.elapsed().as_secs_f64() * 1000.0;
let generated_tokens = &tokens[input_token_count..];
let text = decode_apr_tokens(&config.model_path, generated_tokens);
let generated_token_count = generated_tokens.len();
Some(Ok(InferenceResult {
text,
tokens,
input_token_count,
generated_token_count,
inference_ms,
tok_per_sec: tok_per_sec(generated_token_count, inference_ms),
generation_ms: Some(inference_ms), load_ms,
format: "SafeTensors".to_string(),
used_gpu: true,
gpu_attempted: true,
}))
}
#[cfg(test)]
mod tests {
use super::{wgpu_greedy_decode, CUDA_FALLBACK_LOG_PREFIX, WGPU_FALLBACK_LOG_PREFIX};
fn decode_trace(prompt: &[u32], max_tokens: usize, stop: &[u32]) -> (Vec<u32>, Vec<(u32, usize)>) {
let mut fed = Vec::new();
let out = wgpu_greedy_decode::<()>(
prompt,
max_tokens,
stop,
|t, p| {
fed.push((t, p));
Ok(vec![p as f32])
},
|h| 100 + h[0] as u32,
)
.expect("the fake forward never fails");
(out, fed)
}
#[test]
fn wgpu_decode_prefills_every_prompt_position_before_decoding() {
let (out, fed) = decode_trace(&[7, 8, 9], 2, &[]);
assert_eq!(
fed,
[(7, 0), (8, 1), (9, 2), (102, 3)],
"each prompt token enters the cache at its own position, then decode continues"
);
assert_eq!(out, [7, 8, 9, 102, 103]);
}
#[test]
fn wgpu_decode_stops_and_handles_edges() {
let (out, fed) = decode_trace(&[5], 4, &[101]);
assert_eq!(out, [5, 100, 101], "the stop token is kept, then decode ends");
assert_eq!(fed, [(5, 0), (100, 1)]);
assert_eq!(decode_trace(&[], 4, &[]).0, Vec::<u32>::new(), "empty prompt");
assert!(decode_trace(&[1, 2], 0, &[]).1.is_empty(), "max_tokens 0 runs no forward");
let err = wgpu_greedy_decode(&[1, 2], 1, &[], |_, p| if p == 0 { Err("boom") } else { Ok(vec![]) }, |_| 0);
assert_eq!(err, Err("boom"), "a prefill failure is returned, not skipped");
}
fn assert_reads_as_a_warning_to_a_human(prefix: &str, backend: &str) {
assert!(
!prefix.contains("apr-cpu-vs-gpu-output-parity-v1"),
"the fallback message addresses the user in a contract ID: {prefix}"
);
assert!(
!prefix.starts_with('['),
"the fallback message opens with a bracketed internal tag: {prefix}"
);
assert!(
prefix.starts_with("warning:"),
"the fallback message must announce itself as a warning: {prefix}"
);
assert!(
prefix.contains(backend),
"fallback message must say which backend was rejected; got: {prefix}"
);
assert!(
prefix.ends_with("path rejected"),
"fallback message must say the path was rejected; got: {prefix}"
);
}
#[test]
fn cuda_fallback_log_prefix_warns_the_user_in_prose() {
assert_reads_as_a_warning_to_a_human(CUDA_FALLBACK_LOG_PREFIX, "CUDA");
}
#[test]
fn wgpu_fallback_log_prefix_warns_the_user_in_prose() {
assert_reads_as_a_warning_to_a_human(WGPU_FALLBACK_LOG_PREFIX, "wgpu");
}
#[test]
fn cuda_and_wgpu_fallback_log_prefixes_share_their_shape() {
for prefix in [CUDA_FALLBACK_LOG_PREFIX, WGPU_FALLBACK_LOG_PREFIX] {
assert!(prefix.starts_with("warning: GPU ("), "{prefix}");
assert!(prefix.ends_with(") path rejected"), "{prefix}");
}
assert_ne!(
CUDA_FALLBACK_LOG_PREFIX, WGPU_FALLBACK_LOG_PREFIX,
"the two hops must remain distinguishable"
);
}
#[test]
fn cpu_vs_gpu_cosine_similarity_parallel_returns_one() {
let a = vec![1.0_f32, 2.0, 3.0, 4.0];
let b = a.clone();
let cos = super::cpu_vs_gpu_cosine_similarity(&a, &b);
assert!(
(cos - 1.0).abs() < 1e-6,
"parallel vectors must yield cosine 1.0, got {cos}"
);
}
#[test]
fn cpu_vs_gpu_cosine_similarity_orthogonal_returns_zero() {
let a = vec![1.0_f32, 0.0, 0.0, 0.0];
let b = vec![0.0_f32, 1.0, 0.0, 0.0];
let cos = super::cpu_vs_gpu_cosine_similarity(&a, &b);
assert!(
cos.abs() < 1e-6,
"orthogonal vectors must yield cosine 0.0, got {cos}"
);
}
#[test]
fn cpu_vs_gpu_cosine_similarity_fails_closed() {
let zero = vec![0.0_f32; 4];
let nonzero = vec![1.0_f32, 2.0, 3.0, 4.0];
assert_eq!(
super::cpu_vs_gpu_cosine_similarity(&zero, &nonzero),
0.0,
"zero-norm input must fail closed"
);
let short = vec![1.0_f32, 2.0];
let long = vec![1.0_f32, 2.0, 3.0, 4.0];
assert_eq!(
super::cpu_vs_gpu_cosine_similarity(&short, &long),
0.0,
"length mismatch must fail closed"
);
let empty: Vec<f32> = Vec::new();
assert_eq!(
super::cpu_vs_gpu_cosine_similarity(&empty, &empty),
0.0,
"empty input must fail closed"
);
}
}
#[cfg(all(test, feature = "gpu"))]
mod pmat3757_wgpu_attempt_gate {
use super::wgpu_fallback_allowed;
const CASES: &[(bool, bool, bool, bool, &str)] = &[
(
false, false, false, false,
"bare `apr run model.gguf` — the #3757 defect: on a default \
(non-cuda) install this dequantized 1726.8 MB to F32, failed wgpu's \
cpu-parity gate at cosine 0.9554 and fell back to CPU, costing \
7607 ms against --no-gpu's 3035 ms for the identical answer",
),
(
false, true, false, true,
"`--backend wgpu` — an explicit request is still served, and still \
refuses to report a fallback as success (rc=14)",
),
(
true, true, false, false,
"`--no-gpu --backend wgpu` — an explicit opt-out wins over an \
explicit request",
),
(
true, false, false, false,
"`--no-gpu` — nothing to attempt",
),
(
false, true, true, false,
"a legacy-quant model with `--gpu`: no GPU kernel exists for it, so \
the attempt would dequantize and fail",
),
];
#[test]
fn bare_apr_run_does_not_attempt_wgpu() {
let (no_gpu, accel_forced, legacy, _, why) = CASES[0];
assert!(
!wgpu_fallback_allowed(no_gpu, accel_forced, legacy),
"#3757 REGRESSION: {why}"
);
}
#[test]
fn an_explicit_accelerator_request_is_still_served() {
let (no_gpu, accel_forced, legacy, _, why) = CASES[1];
assert!(
wgpu_fallback_allowed(no_gpu, accel_forced, legacy),
"#3757 OVER-CORRECTION: the fix removed the backend instead of \
making it opt-in. {why}"
);
}
#[test]
fn the_whole_attempt_table_holds() {
let wrong: Vec<String> = CASES
.iter()
.filter(|(n, a, l, want, _)| wgpu_fallback_allowed(*n, *a, *l) != *want)
.map(|(n, a, l, want, why)| {
format!(
"\n - no_gpu={n} accel_forced={a} has_legacy_quant={l}: \
expected allowed={want}, got {}. {why}",
!*want
)
})
.collect();
assert!(
wrong.is_empty(),
"{} of {} wgpu-attempt cases are wrong:{}",
wrong.len(),
CASES.len(),
wrong.join("")
);
}
#[cfg(feature = "cuda")]
#[test]
fn apr_cuda_decline_names_the_type_and_is_silent_when_eligible() {
let cfg = crate::gguf::GGUFConfig {
architecture: "test".to_string(),
constraints: crate::gguf::ArchConstraints::from_architecture("test"),
hidden_dim: 64,
intermediate_dim: 128,
num_layers: 1,
num_heads: 4,
num_kv_heads: 4,
vocab_size: 100,
context_length: 256,
rope_theta: 10000.0,
eps: 1e-5,
rope_type: 0,
explicit_head_dim: None,
query_pre_attn_scalar: None,
bos_token_id: None,
eos_token_id: None,
};
let eligible = crate::gguf::test_helpers::create_test_model_with_config(&cfg);
let mut bad = crate::gguf::test_helpers::create_test_model_with_config(&cfg);
bad.lm_head_weight.qtype = 29;
let notice = super::apr_cuda_decline_notice(&bad)
.expect("a model whose lm_head has no verified GPU kernel owes the user a notice");
assert!(notice.contains("29"), "the notice must NAME the declining type: {notice}");
assert!(
notice.starts_with(super::CUDA_FALLBACK_LOG_PREFIX),
"the notice must announce which backend was rejected: {notice}"
);
assert!(
notice.contains("downstream of THIS decision"),
"the notice must say later backend errors are downstream, since misattributing \
them to wgpu is the defect it exists to prevent: {notice}"
);
assert!(
super::apr_cuda_decline_notice(&eligible).is_none(),
"a fully GPU-eligible model must produce no decline notice"
);
}
#[cfg(feature = "cuda")]
#[test]
fn the_decline_notice_is_emitted_unconditionally_not_behind_verbose() {
let src = std::fs::read_to_string(concat!(
env!("CARGO_MANIFEST_DIR"),
"/src/infer/gguf_gpu_generate.rs"
))
.expect("own source readable");
let at = src
.find("if let Some(notice) = apr_cuda_decline_notice(&model)")
.expect("the quant-whitelist exit must consult apr_cuda_decline_notice");
let tail = &src[at..at + 220];
assert!(
tail.contains("eprintln!(\"{notice}\")"),
"the quant-whitelist exit must PRINT the notice before returning None — a silent \
decline here is #3908, and the sibling exit's own comment says a CUDA rejection \
MUST be visible without --verbose"
);
assert!(
!tail.contains("if verbose"),
"the decline notice must not be behind --verbose: the user who needs it is the \
one who did not pass it"
);
}
}