File size: 3,754 Bytes
a39b81e | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 | // NVFP4 Inference Example for Candle
//
// This example demonstrates:
// 1. Loading a model from HuggingFace
// 2. Quantizing weights to NVFP4
// 3. Running inference with NVFP4 quantized weights on CUDA
//
// Usage:
// cargo run --release --features cuda --bin nvfp4-inference -- \
// --model igorls/gemma-4-12B-it-heretic-v1 \
// --prompt "Hello, world!" \
// --n-tokens 128
use candle::{DType, Device, Tensor, Module};
use std::io::Write;
fn main() -> candle::Result<()> {
let model_id = std::env::args()
.find(|arg| arg.starts_with("--model="))
.map(|s| s[8..].to_string())
.unwrap_or_else(|| "igorls/gemma-4-12B-it-heretic-v1".to_string());
let prompt = std::env::args()
.find(|arg| arg.starts_with("--prompt="))
.map(|s| s[9..].to_string())
.unwrap_or_else(|| "What is the meaning of life?".to_string());
let n_tokens: usize = std::env::args()
.find(|arg| arg.starts_with("--n-tokens="))
.and_then(|s| s[11..].parse().ok())
.unwrap_or(128);
println!("NVFP4 Inference Example");
println!("Model: {model_id}");
println!("Prompt: {prompt}");
println!("Max tokens: {n_tokens}");
println!();
// Initialize CUDA device
let device = Device::cuda_if_available(0)?;
println!("Device: {:?}", device);
// Test NVFP4 quantization
println!("\n=== NVFP4 Quantization Test ===");
let test_data = Tensor::randn(0f32, 1f32, (256, 512), &device)?;
println!("Input shape: {:?}", test_data.shape());
// Quantize to NVFP4
use candle::quantized::{GgmlDType, QTensor};
let qtensor = QTensor::quantize(&test_data, GgmlDType::NVFP4)?;
println!("Quantized dtype: {:?}", qtensor.dtype());
println!("Quantized shape: {:?}", qtensor.shape());
// Dequantize back
let dequant = qtensor.dequantize(&device)?;
println!("Dequantized shape: {:?}", dequant.shape());
// Compute error
let diff = test_data.broadcast_sub(&dequant)?;
let max_err = diff.max_all()?.to_scalar::<f32>()?;
let mean_err = diff.mean_all()?.to_scalar::<f32>()?;
let ref_mean = test_data.abs()?.mean_all()?.to_scalar::<f32>()?;
println!("Max error: {:.6}", max_err);
println!("Mean error: {:.6}", mean_err);
println!("Relative error: {:.6}", mean_err / ref_mean);
// Test matmul
println!("\n=== NVFP4 MatMul Test ===");
let x = Tensor::randn(0f32, 1f32, (1, 512), &device)?;
let w = Tensor::randn(0f32, 0.5f32, (256, 512), &device)?;
let w_q = QTensor::quantize(&w, GgmlDType::NVFP4)?;
// Reference matmul
let ref_out = x.matmul(&w.t()?)?;
// NVFP4 matmul via QMatMul
use candle::quantized::QMatMul;
let qmatmul = QMatMul::from_arc(std::sync::Arc::new(w_q))?;
let nvfp4_out = qmatmul.forward(&x)?;
let matmul_diff = ref_out.broadcast_sub(&nvfp4_out)?;
let matmul_max_err = matmul_diff.max_all()?.to_scalar::<f32>()?;
let matmul_mean_err = matmul_diff.mean_all()?.to_scalar::<f32>()?;
let matmul_ref_mean = ref_out.abs()?.mean_all()?.to_scalar::<f32>()?;
println!("MatMul max error: {:.6}", matmul_max_err);
println!("MatMul mean error: {:.6}", matmul_mean_err);
println!("MatMul relative error: {:.6}", matmul_mean_err / matmul_ref_mean);
println!("\n=== NVFP4 Memory Analysis ===");
let n = 4096usize;
let k = 4096usize;
let fp32_bytes = n * k * 4;
let nvfp4_bytes = n * (k / 16) * 9; // 9 bytes per 16 elements
println!("FP32 weight: {:.1} MB", fp32_bytes as f64 / 1024.0 / 1024.0);
println!("NVFP4 weight: {:.1} MB", nvfp4_bytes as f64 / 1024.0 / 1024.0);
println!("Compression: {:.2}x", fp32_bytes as f64 / nvfp4_bytes as f64);
println!("\n=== Done ===");
Ok(())
}
|