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")]
struct WgpuLmHead<'a> {
output_norm: &'a [f32],
lm_head: &'a [f32],
hidden_dim: usize,
vocab_size: usize,
eps: f32,
}
#[cfg(feature = "gpu")]
impl<'a> WgpuLmHead<'a> {
fn new(model: &'a crate::gguf::OwnedQuantizedModel, lm_head: &'a [f32]) -> Self {
let config = model.config();
Self {
output_norm: model.output_norm_weight(),
lm_head,
hidden_dim: config.hidden_dim,
vocab_size: config.vocab_size,
eps: config.eps,
}
}
fn normed(&self, hidden: &[f32]) -> Vec<f32> {
let sq_sum: f32 = hidden.iter().map(|x| x * x).sum();
let rms = (sq_sum / hidden.len() as f32 + self.eps).sqrt();
hidden
.iter()
.zip(self.output_norm.iter())
.map(|(x, g)| (x / rms) * g)
.collect()
}
fn logit(&self, normed: &[f32], i: usize) -> f32 {
let row = &self.lm_head[i * self.hidden_dim..(i + 1) * self.hidden_dim];
row.iter().zip(normed.iter()).map(|(w, x)| w * x).sum()
}
fn logits(&self, hidden: &[f32]) -> Vec<f32> {
let normed = self.normed(hidden);
(0..self.vocab_size).map(|i| self.logit(&normed, i)).collect()
}
fn argmax(&self, hidden: &[f32]) -> u32 {
let normed = self.normed(hidden);
first_argmax((0..self.vocab_size).map(|i| self.logit(&normed, i)))
}
}
#[cfg(feature = "gpu")]
fn first_argmax(values: impl Iterator<Item = f32>) -> u32 {
let mut best_idx = 0u32;
let mut best_val = f32::NEG_INFINITY;
for (i, v) in values.enumerate() {
if v > best_val {
best_val = v;
best_idx = i as u32;
}
}
best_idx
}
#[cfg(feature = "gpu")]
fn wgpu_kv_caches(
model: &crate::gguf::OwnedQuantizedModel,
max_seq: usize,
) -> Vec<(Vec<f32>, Vec<f32>)> {
let config = model.config();
let kv_dim = config.num_kv_heads * (config.hidden_dim / config.num_heads);
(0..config.num_layers)
.map(|_| (Vec::with_capacity(max_seq * kv_dim), Vec::with_capacity(max_seq * kv_dim)))
.collect()
}
#[cfg(feature = "gpu")]
fn wgpu_forward_token(
model: &crate::gguf::OwnedQuantizedModel,
fwd: &trueno::backends::gpu::WgslForwardPass,
kv_caches: &mut [(Vec<f32>, Vec<f32>)],
token_id: u32,
position: usize,
) -> std::result::Result<Vec<f32>, (usize, String)> {
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| (layer_idx, e))?;
}
Ok(hidden)
}
#[cfg(feature = "gpu")]
fn wgpu_forward_pass(
model: &crate::gguf::OwnedQuantizedModel,
gpu: trueno::backends::gpu::GpuDevice,
) -> trueno::backends::gpu::WgslForwardPass {
let config = model.config();
let mut fwd = trueno::backends::gpu::WgslForwardPass::new(
gpu.device,
gpu.queue,
config.hidden_dim,
config.num_heads,
config.num_kv_heads,
config.hidden_dim / config.num_heads,
config.intermediate_dim,
);
fwd.set_rms_norm_eps(config.eps);
fwd
}
#[cfg(feature = "gpu")]
fn wgpu_lm_head_weight(weights: &[(String, Vec<f32>, usize, usize)]) -> Vec<f32> {
weights
.iter()
.find(|(n, _, _, _)| n == "lm_head")
.map(|(_, d, _, _)| d.clone())
.unwrap_or_default()
}
#[cfg(feature = "gpu")]
fn wgpu_parity_probe(
model: &crate::gguf::OwnedQuantizedModel,
fwd: &trueno::backends::gpu::WgslForwardPass,
head: &WgpuLmHead<'_>,
first_token: u32,
) -> std::result::Result<(), String> {
const MULTI_STEP_PROBE_DEFAULT: usize = 3;
let steps: 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 = steps + 1;
let mut cpu_cache =
crate::gguf::OwnedQuantizedKVCache::from_config(model.config(), probe_max_seq);
let mut probe_kv_caches = wgpu_kv_caches(model, probe_max_seq);
let mut probe_token = first_token;
for step in 0..steps {
let cpu_logits = model
.forward_single_with_cache(probe_token, &mut cpu_cache, step)
.map_err(|e| {
eprintln!(
"{}, attempting fallback: CPU probe step {} forward failed: {}",
WGPU_FALLBACK_LOG_PREFIX, step, e
);
format!("wgpu parity gate: CPU probe step {step} failed: {e}")
})?;
let hidden = wgpu_forward_token(model, fwd, &mut probe_kv_caches, probe_token, step)
.map_err(|(layer_idx, e)| {
eprintln!(
"{}, attempting fallback: wgpu probe step {} layer {} failed: {}",
WGPU_FALLBACK_LOG_PREFIX, step, layer_idx, e
);
format!("wgpu parity gate: step {step} layer {layer_idx} failed: {e}")
})?;
let cos = cpu_vs_gpu_cosine_similarity(&cpu_logits, &head.logits(&hidden));
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,
steps
);
return Err(format!(
"wgpu parity gate: cosine={cos:.6} < 0.99 at step {}/{}",
step + 1,
steps
));
}
probe_token = first_argmax(cpu_logits.iter().copied());
}
Ok(())
}
#[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 mut fwd = wgpu_forward_pass(model, gpu);
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 lm_head_f32 = wgpu_lm_head_weight(&weights);
let head = WgpuLmHead::new(model, &lm_head_f32);
wgpu_parity_probe(model, &fwd, &head, *input_tokens.first().unwrap_or(&0))
.map_err(RealizarError::InferenceError)?;
let mut kv_caches = wgpu_kv_caches(model, gen_config.max_tokens + input_tokens.len() + 16);
let output_tokens = wgpu_greedy_decode::<RealizarError>(
input_tokens,
gen_config.max_tokens,
&gen_config.stop_tokens,
|token_id, position| {
wgpu_forward_token(model, &fwd, &mut kv_caches, token_id, position).map_err(
|(layer_idx, e)| {
RealizarError::InferenceError(format!("wgpu layer {layer_idx}: {e}"))
},
)
},
|hidden| head.argmax(hidden),
)?;
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>> {
let mut cuda_model = cuda_dense_model(model, input_tokens, gen_config, verbose)?;
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 cuda_dense_model(
model: crate::gguf::OwnedQuantizedModel,
input_tokens: &[u32],
gen_config: &crate::gguf::QuantizedGenerateConfig,
verbose: bool,
) -> std::result::Result<crate::gguf::OwnedQuantizedModelCuda, 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 { .. } => {},
}
Ok(cuda_model)
}
#[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 apr_wgpu_generate_config(
config: &InferenceConfig,
model: &crate::gguf::OwnedQuantizedModel,
) -> crate::gguf::QuantizedGenerateConfig {
let cfg = model.config();
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);
gen_config
}
#[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 = MappedAprModel::from_path(&config.model_path).ok()?;
let model = crate::gguf::OwnedQuantizedModel::from_apr(&mapped).ok()?;
let gen_config = apr_wgpu_generate_config(config, &model);
let weights = match wgpu_adapter::dequant_model_weights(&model) {
Ok(w) => w,
Err(e) => return Some(Err(e)),
};
let mut fwd = wgpu_forward_pass(&model, gpu);
for (name, data, _rows, _cols) in &weights {
fwd.upload_weight(name, data);
}
let lm_head_f32 = wgpu_lm_head_weight(&weights);
let head = WgpuLmHead::new(&model, &lm_head_f32);
wgpu_parity_probe(&model, &fwd, &head, *input_tokens.first().unwrap_or(&0)).ok()?;
let model_load_ms = load_start.elapsed().as_millis() as f64;
let mut kv_caches = wgpu_kv_caches(&model, gen_config.max_tokens + input_tokens.len() + 16);
let infer_start = Instant::now();
let decoded = wgpu_greedy_decode::<RealizarError>(
input_tokens,
gen_config.max_tokens,
&gen_config.stop_tokens,
|token_id, position| {
wgpu_forward_token(&model, &fwd, &mut kv_caches, token_id, position).map_err(
|(layer_idx, e)| {
RealizarError::InferenceError(format!("wgpu layer {layer_idx}: {e}"))
},
)
},
|hidden| head.argmax(hidden),
);
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 wgpu_lm_head_tail {
use super::{first_argmax, WgpuLmHead};
#[test]
fn first_argmax_takes_the_first_maximum_and_never_a_nan() {
assert_eq!(first_argmax([1.0, 3.0, 3.0, 2.0].into_iter()), 1);
assert_eq!(first_argmax([f32::NAN, 0.5, f32::NAN, 0.25].into_iter()), 1);
assert_eq!(first_argmax([f32::NEG_INFINITY, -1.0].into_iter()), 1);
assert_eq!(first_argmax(std::iter::empty()), 0);
}
#[test]
fn the_head_norms_with_eps_and_gamma_then_projects_every_vocab_row() {
let head = WgpuLmHead {
output_norm: &[1.0, 2.0],
lm_head: &[1.0, 0.0, 0.0, 1.0, 1.0, 1.0, 0.0, 1.0],
hidden_dim: 2,
vocab_size: 4,
eps: 0.5,
};
let rms = (25.0_f32 / 2.0 + 0.5).sqrt();
let (n0, n1) = (3.0 / rms, (4.0 / rms) * 2.0);
assert_eq!(head.logits(&[3.0, 4.0]), vec![n0, n1, n0 + n1, n1]);
assert_eq!(head.argmax(&[3.0, 4.0]), 2);
let tied = WgpuLmHead {
lm_head: &[1.0, 0.0, 0.0, 1.0, 0.0, 0.0, 0.0, 1.0],
..head
};
assert_eq!(tied.argmax(&[3.0, 4.0]), 1);
}
}
#[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"
);
}
}