osobhandClaude Opus 4.6
02d382d5f6
Consistent formatting pass: line wrapping, import sorting, trailing whitespace removal, let-chain indentation, merged derive attributes, and unsafe block reformatting. Co-Authored-By: Claude Opus 4.6 (1M context) <[email protected]>
728 lines
20 KiB
Rust
728 lines
20 KiB
Rust
//! # RustyTorch++ WASM Inference Runtime
|
|
//!
|
|
//! WebAssembly inference runtime for deploying ML models in browsers and Node.js.
|
|
//!
|
|
//! ## Features
|
|
//!
|
|
//! - **Browser deployment**: Run inference directly in web browsers
|
|
//! - **Node.js support**: Server-side WASM with Wasmtime/Wasmer
|
|
//! - **Zero dependencies**: No Python, CUDA, or native libraries required
|
|
//! - **Streaming inference**: Token-by-token generation for LLMs
|
|
//! - **Quantization**: INT8/INT4 models for reduced memory footprint
|
|
//!
|
|
//! ## Usage (JavaScript/TypeScript)
|
|
//!
|
|
//! ```javascript
|
|
//! import init, { WasmInferenceEngine, InferenceConfig } from 'rtx-wasm-inference';
|
|
//!
|
|
//! async function main() {
|
|
//! await init();
|
|
//!
|
|
//! const config = InferenceConfig.default();
|
|
//! const engine = await WasmInferenceEngine.new(config);
|
|
//!
|
|
//! const result = await engine.infer("Hello, world!");
|
|
//! console.log(result.text);
|
|
//! }
|
|
//! ```
|
|
|
|
#![warn(missing_docs)]
|
|
#![allow(clippy::new_without_default)]
|
|
|
|
use serde::{Deserialize, Serialize};
|
|
use std::collections::HashMap;
|
|
use wasm_bindgen::prelude::*;
|
|
|
|
pub mod model;
|
|
pub mod quantization;
|
|
pub mod runtime;
|
|
pub mod tensor;
|
|
pub mod tensor_core;
|
|
pub mod tokenizer;
|
|
pub mod webgpu;
|
|
|
|
// Re-export submodules
|
|
pub use model::*;
|
|
pub use quantization::*;
|
|
pub use runtime::*;
|
|
pub use tensor::*;
|
|
pub use tokenizer::*;
|
|
pub use webgpu::*;
|
|
|
|
// Re-export core tensor types (no_std compatible)
|
|
pub use tensor_core::backend::{
|
|
BackendLimits, BackendPreference, BackendType, UnifiedBackend, available_backends,
|
|
};
|
|
pub use tensor_core::simd::{SIMD_AVAILABLE, SIMD_WIDTH};
|
|
pub use tensor_core::wasm_api::{
|
|
WasmBackend, WasmTensor, get_backend_info, get_simd_width, is_simd_available,
|
|
};
|
|
pub use tensor_core::{DType, TensorCore, TensorError};
|
|
|
|
/// Initialize panic hook for better error messages in browser console
|
|
#[wasm_bindgen(start)]
|
|
pub fn init_panic_hook() {
|
|
#[cfg(feature = "console_error_panic_hook")]
|
|
console_error_panic_hook::set_once();
|
|
}
|
|
|
|
/// Compute backend for WASM inference
|
|
#[wasm_bindgen]
|
|
#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize)]
|
|
pub enum ComputeBackend {
|
|
/// CPU-only execution (always available)
|
|
Cpu,
|
|
/// WebGPU acceleration (requires browser support)
|
|
WebGpu,
|
|
/// Automatic selection (prefer GPU if available)
|
|
Auto,
|
|
}
|
|
|
|
impl Default for ComputeBackend {
|
|
fn default() -> Self {
|
|
Self::Auto
|
|
}
|
|
}
|
|
|
|
/// Configuration for WASM inference engine
|
|
#[wasm_bindgen]
|
|
#[derive(Debug, Clone, Serialize, Deserialize)]
|
|
pub struct InferenceConfig {
|
|
/// Maximum sequence length
|
|
max_seq_len: usize,
|
|
/// Maximum batch size
|
|
max_batch_size: usize,
|
|
/// Use quantized model
|
|
use_quantization: bool,
|
|
/// Quantization bits (4 or 8)
|
|
quantization_bits: u8,
|
|
/// Enable KV-cache
|
|
use_kv_cache: bool,
|
|
/// Temperature for sampling
|
|
temperature: f32,
|
|
/// Top-p sampling threshold
|
|
top_p: f32,
|
|
/// Top-k sampling limit
|
|
top_k: usize,
|
|
/// Enable streaming output
|
|
streaming: bool,
|
|
/// Compute backend selection
|
|
compute_backend: ComputeBackend,
|
|
/// WebGPU workgroup size (for tuning)
|
|
webgpu_workgroup_size: u32,
|
|
}
|
|
|
|
#[wasm_bindgen]
|
|
impl InferenceConfig {
|
|
/// Create default configuration
|
|
#[wasm_bindgen(constructor)]
|
|
pub fn new() -> Self {
|
|
Self::default()
|
|
}
|
|
|
|
/// Create configuration optimized for speed
|
|
#[wasm_bindgen]
|
|
pub fn fast() -> Self {
|
|
Self {
|
|
max_seq_len: 512,
|
|
max_batch_size: 1,
|
|
use_quantization: true,
|
|
quantization_bits: 4,
|
|
use_kv_cache: true,
|
|
temperature: 0.7,
|
|
top_p: 0.9,
|
|
top_k: 40,
|
|
streaming: true,
|
|
compute_backend: ComputeBackend::Auto,
|
|
webgpu_workgroup_size: 256,
|
|
}
|
|
}
|
|
|
|
/// Create configuration optimized for quality
|
|
#[wasm_bindgen]
|
|
pub fn quality() -> Self {
|
|
Self {
|
|
max_seq_len: 2048,
|
|
max_batch_size: 1,
|
|
use_quantization: false,
|
|
quantization_bits: 8,
|
|
use_kv_cache: true,
|
|
temperature: 0.8,
|
|
top_p: 0.95,
|
|
top_k: 50,
|
|
streaming: false,
|
|
compute_backend: ComputeBackend::Auto,
|
|
webgpu_workgroup_size: 256,
|
|
}
|
|
}
|
|
|
|
/// Create configuration optimized for WebGPU
|
|
#[wasm_bindgen]
|
|
pub fn webgpu_optimized() -> Self {
|
|
Self {
|
|
max_seq_len: 2048,
|
|
max_batch_size: 1,
|
|
use_quantization: true,
|
|
quantization_bits: 8,
|
|
use_kv_cache: true,
|
|
temperature: 0.7,
|
|
top_p: 0.9,
|
|
top_k: 50,
|
|
streaming: true,
|
|
compute_backend: ComputeBackend::WebGpu,
|
|
webgpu_workgroup_size: 256,
|
|
}
|
|
}
|
|
|
|
/// Create configuration for CPU-only execution
|
|
#[wasm_bindgen]
|
|
pub fn cpu_only() -> Self {
|
|
Self {
|
|
compute_backend: ComputeBackend::Cpu,
|
|
..Self::default()
|
|
}
|
|
}
|
|
|
|
/// Set maximum sequence length
|
|
#[wasm_bindgen]
|
|
pub fn with_max_seq_len(mut self, len: usize) -> Self {
|
|
self.max_seq_len = len;
|
|
self
|
|
}
|
|
|
|
/// Set temperature
|
|
#[wasm_bindgen]
|
|
pub fn with_temperature(mut self, temp: f32) -> Self {
|
|
self.temperature = temp;
|
|
self
|
|
}
|
|
|
|
/// Enable streaming
|
|
#[wasm_bindgen]
|
|
pub fn with_streaming(mut self, enabled: bool) -> Self {
|
|
self.streaming = enabled;
|
|
self
|
|
}
|
|
|
|
/// Enable quantization
|
|
#[wasm_bindgen]
|
|
pub fn with_quantization(mut self, bits: u8) -> Self {
|
|
self.use_quantization = true;
|
|
self.quantization_bits = bits;
|
|
self
|
|
}
|
|
|
|
/// Set compute backend
|
|
#[wasm_bindgen]
|
|
pub fn with_compute_backend(mut self, backend: ComputeBackend) -> Self {
|
|
self.compute_backend = backend;
|
|
self
|
|
}
|
|
|
|
/// Set WebGPU workgroup size
|
|
#[wasm_bindgen]
|
|
pub fn with_webgpu_workgroup_size(mut self, size: u32) -> Self {
|
|
self.webgpu_workgroup_size = size;
|
|
self
|
|
}
|
|
|
|
/// Get compute backend
|
|
#[wasm_bindgen(getter)]
|
|
pub fn compute_backend(&self) -> ComputeBackend {
|
|
self.compute_backend
|
|
}
|
|
|
|
/// Get WebGPU workgroup size
|
|
#[wasm_bindgen(getter)]
|
|
pub fn webgpu_workgroup_size(&self) -> u32 {
|
|
self.webgpu_workgroup_size
|
|
}
|
|
}
|
|
|
|
impl Default for InferenceConfig {
|
|
fn default() -> Self {
|
|
Self {
|
|
max_seq_len: 1024,
|
|
max_batch_size: 1,
|
|
use_quantization: true,
|
|
quantization_bits: 8,
|
|
use_kv_cache: true,
|
|
temperature: 0.7,
|
|
top_p: 0.9,
|
|
top_k: 50,
|
|
streaming: false,
|
|
compute_backend: ComputeBackend::Auto,
|
|
webgpu_workgroup_size: 256,
|
|
}
|
|
}
|
|
}
|
|
|
|
/// Result from inference
|
|
#[wasm_bindgen]
|
|
#[derive(Debug, Clone, Serialize, Deserialize)]
|
|
pub struct InferenceResult {
|
|
/// Generated text
|
|
text: String,
|
|
/// Number of tokens generated
|
|
tokens_generated: usize,
|
|
/// Inference time in milliseconds
|
|
inference_time_ms: f64,
|
|
/// Tokens per second
|
|
tokens_per_sec: f64,
|
|
/// Finish reason
|
|
finish_reason: String,
|
|
}
|
|
|
|
#[wasm_bindgen]
|
|
impl InferenceResult {
|
|
/// Get the generated text
|
|
#[wasm_bindgen(getter)]
|
|
pub fn text(&self) -> String {
|
|
self.text.clone()
|
|
}
|
|
|
|
/// Get the number of tokens generated
|
|
#[wasm_bindgen(getter)]
|
|
pub fn tokens_generated(&self) -> usize {
|
|
self.tokens_generated
|
|
}
|
|
|
|
/// Get the inference time in milliseconds
|
|
#[wasm_bindgen(getter)]
|
|
pub fn inference_time_ms(&self) -> f64 {
|
|
self.inference_time_ms
|
|
}
|
|
|
|
/// Get tokens per second
|
|
#[wasm_bindgen(getter)]
|
|
pub fn tokens_per_sec(&self) -> f64 {
|
|
self.tokens_per_sec
|
|
}
|
|
|
|
/// Get the finish reason
|
|
#[wasm_bindgen(getter)]
|
|
pub fn finish_reason(&self) -> String {
|
|
self.finish_reason.clone()
|
|
}
|
|
|
|
/// Convert to JSON string
|
|
#[wasm_bindgen]
|
|
pub fn to_json(&self) -> Result<String, JsError> {
|
|
serde_json::to_string(self).map_err(|e| JsError::new(&e.to_string()))
|
|
}
|
|
}
|
|
|
|
/// Main WASM inference engine
|
|
#[wasm_bindgen]
|
|
pub struct WasmInferenceEngine {
|
|
config: InferenceConfig,
|
|
runtime_info: WasmRuntimeInfo,
|
|
model: Option<WasmModel>,
|
|
tokenizer: Option<WasmTokenizer>,
|
|
kv_cache: Option<WasmKvCache>,
|
|
stats: EngineStats,
|
|
}
|
|
|
|
/// Engine statistics
|
|
#[derive(Debug, Clone, Default)]
|
|
struct EngineStats {
|
|
total_requests: u64,
|
|
total_tokens: u64,
|
|
total_time_ms: f64,
|
|
}
|
|
|
|
#[wasm_bindgen]
|
|
impl WasmInferenceEngine {
|
|
/// Create a new inference engine
|
|
#[wasm_bindgen(constructor)]
|
|
pub fn new(config: InferenceConfig) -> Result<WasmInferenceEngine, JsError> {
|
|
let runtime_info = WasmRuntimeInfo::detect();
|
|
|
|
Ok(Self {
|
|
config,
|
|
runtime_info,
|
|
model: None,
|
|
tokenizer: None,
|
|
kv_cache: None,
|
|
stats: EngineStats::default(),
|
|
})
|
|
}
|
|
|
|
/// Load a model from bytes
|
|
#[wasm_bindgen]
|
|
pub async fn load_model(&mut self, model_bytes: &[u8]) -> Result<(), JsError> {
|
|
let model = WasmModel::from_bytes(model_bytes, self.config.use_quantization)?;
|
|
self.model = Some(model);
|
|
|
|
if self.config.use_kv_cache {
|
|
self.kv_cache = Some(WasmKvCache::new(
|
|
self.config.max_seq_len,
|
|
32, // num_layers
|
|
32, // num_heads
|
|
128, // head_dim
|
|
));
|
|
}
|
|
|
|
Ok(())
|
|
}
|
|
|
|
/// Load tokenizer from JSON
|
|
#[wasm_bindgen]
|
|
pub fn load_tokenizer(&mut self, tokenizer_json: &str) -> Result<(), JsError> {
|
|
let tokenizer = WasmTokenizer::from_json(tokenizer_json)?;
|
|
self.tokenizer = Some(tokenizer);
|
|
Ok(())
|
|
}
|
|
|
|
/// Run inference on input text
|
|
#[wasm_bindgen]
|
|
pub async fn infer(
|
|
&mut self,
|
|
input: &str,
|
|
max_tokens: Option<usize>,
|
|
) -> Result<InferenceResult, JsError> {
|
|
let start = self.runtime_info.performance_now();
|
|
let max_tokens = max_tokens.unwrap_or(256);
|
|
|
|
// Tokenize input
|
|
let tokenizer = self
|
|
.tokenizer
|
|
.as_ref()
|
|
.ok_or_else(|| JsError::new("Tokenizer not loaded"))?;
|
|
let input_ids = tokenizer.encode(input)?;
|
|
|
|
// Run model inference
|
|
let model = self
|
|
.model
|
|
.as_ref()
|
|
.ok_or_else(|| JsError::new("Model not loaded"))?;
|
|
|
|
let mut output_ids = input_ids.clone();
|
|
let mut generated_tokens = 0;
|
|
|
|
for _ in 0..max_tokens {
|
|
// Get next token prediction
|
|
let logits = model.forward(&output_ids, self.kv_cache.as_mut())?;
|
|
|
|
// Sample next token
|
|
let next_token = self.sample_token(&logits)?;
|
|
|
|
// Check for EOS
|
|
if next_token == tokenizer.eos_token_id() {
|
|
break;
|
|
}
|
|
|
|
output_ids.push(next_token);
|
|
generated_tokens += 1;
|
|
|
|
// Update KV cache position
|
|
if let Some(cache) = &mut self.kv_cache {
|
|
cache.step();
|
|
}
|
|
}
|
|
|
|
// Decode output
|
|
let output_text = tokenizer.decode(&output_ids[input_ids.len()..])?;
|
|
|
|
let end = self.runtime_info.performance_now();
|
|
let inference_time_ms = end - start;
|
|
let tokens_per_sec = if inference_time_ms > 0.0 {
|
|
(generated_tokens as f64 / inference_time_ms) * 1000.0
|
|
} else {
|
|
0.0
|
|
};
|
|
|
|
// Update stats
|
|
self.stats.total_requests += 1;
|
|
self.stats.total_tokens += generated_tokens as u64;
|
|
self.stats.total_time_ms += inference_time_ms;
|
|
|
|
Ok(InferenceResult {
|
|
text: output_text,
|
|
tokens_generated: generated_tokens,
|
|
inference_time_ms,
|
|
tokens_per_sec,
|
|
finish_reason: if generated_tokens < max_tokens {
|
|
"stop".to_string()
|
|
} else {
|
|
"length".to_string()
|
|
},
|
|
})
|
|
}
|
|
|
|
/// Sample a token from logits
|
|
fn sample_token(&self, logits: &[f32]) -> Result<u32, JsError> {
|
|
if logits.is_empty() {
|
|
return Err(JsError::new("Empty logits"));
|
|
}
|
|
|
|
// Apply temperature
|
|
let temp = self.config.temperature;
|
|
let scaled: Vec<f32> = logits.iter().map(|&x| x / temp).collect();
|
|
|
|
// Softmax
|
|
let max_val = scaled.iter().copied().fold(f32::NEG_INFINITY, f32::max);
|
|
let exp_sum: f32 = scaled.iter().map(|&x| (x - max_val).exp()).sum();
|
|
let probs: Vec<f32> = scaled
|
|
.iter()
|
|
.map(|&x| (x - max_val).exp() / exp_sum)
|
|
.collect();
|
|
|
|
// Top-p sampling
|
|
let mut sorted_indices: Vec<usize> = (0..probs.len()).collect();
|
|
sorted_indices.sort_by(|&a, &b| probs[b].partial_cmp(&probs[a]).unwrap());
|
|
|
|
let mut cumsum = 0.0;
|
|
let mut cutoff_idx = sorted_indices.len();
|
|
for (i, &idx) in sorted_indices.iter().enumerate() {
|
|
cumsum += probs[idx];
|
|
if cumsum >= self.config.top_p {
|
|
cutoff_idx = i + 1;
|
|
break;
|
|
}
|
|
}
|
|
|
|
// Apply top-k
|
|
cutoff_idx = cutoff_idx.min(self.config.top_k);
|
|
|
|
// Renormalize
|
|
let valid_indices = &sorted_indices[..cutoff_idx];
|
|
let valid_sum: f32 = valid_indices.iter().map(|&i| probs[i]).sum();
|
|
let normalized: Vec<f32> = valid_indices
|
|
.iter()
|
|
.map(|&i| probs[i] / valid_sum)
|
|
.collect();
|
|
|
|
// Sample
|
|
let r: f32 = rand::random();
|
|
let mut cumsum = 0.0;
|
|
for (i, &prob) in normalized.iter().enumerate() {
|
|
cumsum += prob;
|
|
if r <= cumsum {
|
|
return Ok(valid_indices[i] as u32);
|
|
}
|
|
}
|
|
|
|
Ok(valid_indices[0] as u32)
|
|
}
|
|
|
|
/// Get runtime information
|
|
#[wasm_bindgen]
|
|
pub fn runtime_info(&self) -> WasmRuntimeInfo {
|
|
self.runtime_info.clone()
|
|
}
|
|
|
|
/// Get engine statistics
|
|
#[wasm_bindgen]
|
|
pub fn stats(&self) -> Result<JsValue, JsError> {
|
|
let stats = HashMap::from([
|
|
(
|
|
"total_requests".to_string(),
|
|
self.stats.total_requests as f64,
|
|
),
|
|
("total_tokens".to_string(), self.stats.total_tokens as f64),
|
|
("total_time_ms".to_string(), self.stats.total_time_ms),
|
|
(
|
|
"avg_tokens_per_sec".to_string(),
|
|
if self.stats.total_time_ms > 0.0 {
|
|
(self.stats.total_tokens as f64 / self.stats.total_time_ms) * 1000.0
|
|
} else {
|
|
0.0
|
|
},
|
|
),
|
|
]);
|
|
serde_wasm_bindgen::to_value(&stats).map_err(|e| JsError::new(&e.to_string()))
|
|
}
|
|
|
|
/// Clear KV cache
|
|
#[wasm_bindgen]
|
|
pub fn clear_cache(&mut self) {
|
|
if let Some(cache) = &mut self.kv_cache {
|
|
cache.clear();
|
|
}
|
|
}
|
|
|
|
/// Check if model is loaded
|
|
#[wasm_bindgen]
|
|
pub fn is_model_loaded(&self) -> bool {
|
|
self.model.is_some()
|
|
}
|
|
|
|
/// Check if tokenizer is loaded
|
|
#[wasm_bindgen]
|
|
pub fn is_tokenizer_loaded(&self) -> bool {
|
|
self.tokenizer.is_some()
|
|
}
|
|
|
|
/// Get memory usage estimate
|
|
#[wasm_bindgen]
|
|
pub fn memory_usage(&self) -> usize {
|
|
let model_mem = self
|
|
.model
|
|
.as_ref()
|
|
.map_or(0, model::WasmModel::memory_usage);
|
|
let cache_mem = self
|
|
.kv_cache
|
|
.as_ref()
|
|
.map_or(0, tensor::WasmKvCache::memory_usage);
|
|
let tokenizer_mem = self
|
|
.tokenizer
|
|
.as_ref()
|
|
.map_or(0, tokenizer::WasmTokenizer::memory_usage);
|
|
model_mem + cache_mem + tokenizer_mem
|
|
}
|
|
}
|
|
|
|
/// Version information
|
|
#[wasm_bindgen]
|
|
pub fn version() -> String {
|
|
env!("CARGO_PKG_VERSION").to_string()
|
|
}
|
|
|
|
/// Check if SIMD is supported
|
|
#[wasm_bindgen]
|
|
pub fn simd_supported() -> bool {
|
|
#[cfg(feature = "simd")]
|
|
{
|
|
true
|
|
}
|
|
#[cfg(not(feature = "simd"))]
|
|
{
|
|
false
|
|
}
|
|
}
|
|
|
|
/// Check if threading is supported
|
|
#[wasm_bindgen]
|
|
pub fn threads_supported() -> bool {
|
|
#[cfg(feature = "threads")]
|
|
{
|
|
true
|
|
}
|
|
#[cfg(not(feature = "threads"))]
|
|
{
|
|
false
|
|
}
|
|
}
|
|
|
|
/// Check if WebGPU is supported in the current browser
|
|
#[wasm_bindgen]
|
|
pub fn webgpu_supported() -> bool {
|
|
WebGPUContext::is_supported()
|
|
}
|
|
|
|
/// Get recommended compute backend based on browser capabilities
|
|
#[wasm_bindgen]
|
|
pub fn recommended_backend() -> ComputeBackend {
|
|
if webgpu_supported() {
|
|
ComputeBackend::WebGpu
|
|
} else {
|
|
ComputeBackend::Cpu
|
|
}
|
|
}
|
|
|
|
/// Get browser GPU capabilities summary
|
|
#[wasm_bindgen]
|
|
pub async fn get_gpu_capabilities() -> Result<JsValue, JsError> {
|
|
if !webgpu_supported() {
|
|
let info: HashMap<String, String> = HashMap::from([
|
|
("status".to_string(), "not_supported".to_string()),
|
|
(
|
|
"message".to_string(),
|
|
"WebGPU is not available in this browser".to_string(),
|
|
),
|
|
]);
|
|
return serde_wasm_bindgen::to_value(&info).map_err(|e| JsError::new(&e.to_string()));
|
|
}
|
|
|
|
let ctx = WebGPUContext::new().await?;
|
|
|
|
let mut info: HashMap<String, String> = HashMap::new();
|
|
info.insert("status".to_string(), format!("{:?}", ctx.status()));
|
|
|
|
if let Some(adapter_info) = ctx.adapter_info() {
|
|
info.insert("vendor".to_string(), adapter_info.vendor);
|
|
info.insert("architecture".to_string(), adapter_info.architecture);
|
|
info.insert("device".to_string(), adapter_info.device);
|
|
info.insert("description".to_string(), adapter_info.description);
|
|
info.insert(
|
|
"is_fallback".to_string(),
|
|
adapter_info.is_fallback.to_string(),
|
|
);
|
|
}
|
|
|
|
let limits = ctx.limits();
|
|
info.insert(
|
|
"max_buffer_size_mb".to_string(),
|
|
(limits.max_buffer_size / (1024 * 1024)).to_string(),
|
|
);
|
|
info.insert(
|
|
"max_workgroup_size_x".to_string(),
|
|
limits.max_compute_workgroup_size_x.to_string(),
|
|
);
|
|
|
|
serde_wasm_bindgen::to_value(&info).map_err(|e| JsError::new(&e.to_string()))
|
|
}
|
|
|
|
#[cfg(test)]
|
|
mod tests {
|
|
use super::*;
|
|
|
|
#[test]
|
|
fn test_config_default() {
|
|
let config = InferenceConfig::default();
|
|
assert_eq!(config.max_seq_len, 1024);
|
|
assert!(config.use_kv_cache);
|
|
}
|
|
|
|
#[test]
|
|
fn test_config_fast() {
|
|
let config = InferenceConfig::fast();
|
|
assert_eq!(config.max_seq_len, 512);
|
|
assert_eq!(config.quantization_bits, 4);
|
|
assert!(config.streaming);
|
|
}
|
|
|
|
#[test]
|
|
fn test_config_builder() {
|
|
let config = InferenceConfig::new()
|
|
.with_max_seq_len(2048)
|
|
.with_temperature(0.5)
|
|
.with_streaming(true);
|
|
|
|
assert_eq!(config.max_seq_len, 2048);
|
|
assert_eq!(config.temperature, 0.5);
|
|
assert!(config.streaming);
|
|
}
|
|
|
|
#[test]
|
|
fn test_config_webgpu_optimized() {
|
|
let config = InferenceConfig::webgpu_optimized();
|
|
assert_eq!(config.compute_backend, ComputeBackend::WebGpu);
|
|
assert_eq!(config.webgpu_workgroup_size, 256);
|
|
assert!(config.streaming);
|
|
}
|
|
|
|
#[test]
|
|
fn test_config_cpu_only() {
|
|
let config = InferenceConfig::cpu_only();
|
|
assert_eq!(config.compute_backend, ComputeBackend::Cpu);
|
|
}
|
|
|
|
#[test]
|
|
fn test_config_with_compute_backend() {
|
|
let config = InferenceConfig::new()
|
|
.with_compute_backend(ComputeBackend::WebGpu)
|
|
.with_webgpu_workgroup_size(128);
|
|
|
|
assert_eq!(config.compute_backend, ComputeBackend::WebGpu);
|
|
assert_eq!(config.webgpu_workgroup_size, 128);
|
|
}
|
|
|
|
#[test]
|
|
fn test_compute_backend_default() {
|
|
let backend = ComputeBackend::default();
|
|
assert_eq!(backend, ComputeBackend::Auto);
|
|
}
|
|
}
|