287 lines
9.6 KiB
Rust
287 lines
9.6 KiB
Rust
//! LLM Inference Demo - Speculative Decoding & Advanced Features
|
|
//!
|
|
//! This example demonstrates RustyTorch++ LLM inference capabilities:
|
|
//! - Speculative decoding with 6 strategies
|
|
//! - KV-cache with entropy-guided eviction
|
|
//! - Continuous batching
|
|
//! - Quantization (INT4, INT8, FP8)
|
|
//!
|
|
//! Run with: cargo run --example llm_inference_demo
|
|
|
|
use std::time::{Duration, Instant};
|
|
|
|
fn main() {
|
|
println!("=== RustyTorch++ LLM Inference Demo ===\n");
|
|
|
|
// Demo 1: Speculative Decoding Strategies
|
|
demo_speculative_decoding();
|
|
|
|
// Demo 2: KV-Cache Configuration
|
|
demo_kv_cache();
|
|
|
|
// Demo 3: Continuous Batching
|
|
demo_continuous_batching();
|
|
|
|
// Demo 4: Quantization
|
|
demo_quantization();
|
|
|
|
// Demo 5: Complete Inference Pipeline
|
|
demo_complete_pipeline();
|
|
|
|
println!("\n=== Demo Complete ===");
|
|
}
|
|
|
|
/// Demo 1: Speculative Decoding - 6 Strategies for 2-4x Speedup
|
|
fn demo_speculative_decoding() {
|
|
println!("--- Demo 1: Speculative Decoding (6 Strategies) ---\n");
|
|
|
|
println!("RustyTorch++ supports 6 speculative decoding strategies:");
|
|
println!();
|
|
|
|
// Strategy 1: Small Draft Model
|
|
println!("1. SmallModel - Use a smaller model as draft");
|
|
println!(r#"
|
|
let config = SpeculativeConfig::new()
|
|
.with_strategy(SpeculativeStrategy::SmallModel {
|
|
draft_model_id: "TinyLlama/TinyLlama-1.1B-Chat-v1.0".to_string(),
|
|
speculative_tokens: 5,
|
|
})
|
|
.with_acceptance_threshold(0.7);
|
|
|
|
let engine = SpeculativeEngine::new(target_model, config).await?;
|
|
let output = engine.generate(&prompt, max_tokens).await?;
|
|
println!("Speedup: {:.2}x", output.metrics.speedup_ratio);
|
|
"#);
|
|
|
|
// Strategy 2: Medusa Multi-Head
|
|
println!("2. Medusa - Multiple prediction heads for parallel drafting");
|
|
println!(r#"
|
|
let config = SpeculativeConfig::new()
|
|
.with_strategy(SpeculativeStrategy::Medusa {
|
|
num_heads: 4,
|
|
tokens_per_head: 3,
|
|
});
|
|
"#);
|
|
|
|
// Strategy 3: EAGLE
|
|
println!("3. EAGLE - Feature-level draft speculation");
|
|
println!(r#"
|
|
let config = SpeculativeConfig::new()
|
|
.with_strategy(SpeculativeStrategy::EAGLE {
|
|
feature_dim: 2048,
|
|
draft_depth: 2,
|
|
});
|
|
"#);
|
|
|
|
// Strategy 4: Self-Speculative
|
|
println!("4. SelfSpeculative - Layer skipping (no extra model needed)");
|
|
println!(r#"
|
|
let config = SpeculativeConfig::new()
|
|
.with_strategy(SpeculativeStrategy::SelfSpeculative {
|
|
exit_layers: vec![8, 16, 24], // Early exit points
|
|
confidence_threshold: 0.9,
|
|
});
|
|
"#);
|
|
|
|
// Strategy 5: Lookahead
|
|
println!("5. Lookahead - N-gram based prediction");
|
|
println!(r#"
|
|
let config = SpeculativeConfig::new()
|
|
.with_strategy(SpeculativeStrategy::Lookahead {
|
|
n_gram_size: 4,
|
|
max_lookahead: 8,
|
|
});
|
|
"#);
|
|
|
|
// Strategy 6: Tree Attention
|
|
println!("6. TreeAttention - Tree-structured token verification");
|
|
println!(r#"
|
|
let config = SpeculativeConfig::new()
|
|
.with_strategy(SpeculativeStrategy::TreeAttention {
|
|
tree_depth: 3,
|
|
branch_factor: 2,
|
|
});
|
|
"#);
|
|
|
|
println!("\nExpected speedups:");
|
|
println!(" | Strategy | Speedup | Memory | Best For |");
|
|
println!(" |-----------------|---------|--------|---------------------|");
|
|
println!(" | SmallModel | 2-3x | +40% | General use |");
|
|
println!(" | Medusa | 2-4x | +20% | High acceptance |");
|
|
println!(" | EAGLE | 2-3x | +15% | Long sequences |");
|
|
println!(" | SelfSpeculative | 1.5-2x | +0% | Memory constrained |");
|
|
println!(" | Lookahead | 1.3-2x | +5% | Repetitive text |");
|
|
println!(" | TreeAttention | 2-3x | +10% | Batch verification |");
|
|
}
|
|
|
|
/// Demo 2: KV-Cache with Entropy-Guided Eviction
|
|
fn demo_kv_cache() {
|
|
println!("\n--- Demo 2: KV-Cache (Entropy-Guided Eviction) ---\n");
|
|
|
|
println!("RustyTorch++ unique feature: entropy-guided cache eviction");
|
|
println!();
|
|
|
|
println!(r#"
|
|
// Configure KV-cache with entropy-based eviction
|
|
let cache_config = KVCacheConfig::new()
|
|
.with_max_tokens(32768) // Max cache size
|
|
.with_eviction_policy(EvictionPolicy::EntropyGuided {
|
|
entropy_threshold: 0.3, // Evict low-entropy tokens
|
|
min_retention: 0.5, // Keep at least 50% of cache
|
|
})
|
|
.with_memory_budget(8 * 1024 * 1024 * 1024) // 8GB budget
|
|
.with_quantization(CacheQuantization::FP8); // Compress cache
|
|
|
|
// Other eviction policies:
|
|
// - EvictionPolicy::LRU - Least recently used
|
|
// - EvictionPolicy::Sliding - Sliding window
|
|
// - EvictionPolicy::Attention - Low attention score eviction
|
|
// - EvictionPolicy::EntropyGuided - RustyTorch++ exclusive
|
|
"#);
|
|
|
|
println!("Entropy-guided benefits:");
|
|
println!(" - Keeps high-information tokens (proper nouns, numbers)");
|
|
println!(" - Evicts low-information tokens (stopwords, punctuation)");
|
|
println!(" - 30-50% memory reduction with <1% perplexity increase");
|
|
}
|
|
|
|
/// Demo 3: Continuous Batching with SLA Awareness
|
|
fn demo_continuous_batching() {
|
|
println!("\n--- Demo 3: Continuous Batching (SLA-Aware) ---\n");
|
|
|
|
println!(r#"
|
|
// Configure continuous batching scheduler
|
|
let scheduler = ContinuousBatchingScheduler::new()
|
|
.with_max_batch_size(64)
|
|
.with_max_tokens_per_batch(8192)
|
|
.with_scheduling_policy(SchedulingPolicy::SLAAware {
|
|
ttft_target_ms: 100, // Time to first token target
|
|
itl_target_ms: 20, // Inter-token latency target
|
|
preemption_enabled: true, // Allow request preemption
|
|
})
|
|
.with_priority_levels(3); // Support high/medium/low priority
|
|
|
|
// Request priorities
|
|
let high_priority = RequestConfig::new()
|
|
.with_priority(Priority::High)
|
|
.with_max_latency_ms(50);
|
|
|
|
let normal = RequestConfig::new()
|
|
.with_priority(Priority::Normal);
|
|
|
|
// Add requests to batch
|
|
scheduler.add_request(prompt1, high_priority).await?;
|
|
scheduler.add_request(prompt2, normal).await?;
|
|
|
|
// Process batch - high priority requests get preference
|
|
let results = scheduler.process_batch().await?;
|
|
"#);
|
|
|
|
println!("SLA-aware scheduling features:");
|
|
println!(" - Priority-based request ordering");
|
|
println!(" - Automatic request preemption");
|
|
println!(" - Chunked prefill for reduced TTFT");
|
|
println!(" - Dynamic batch size adjustment");
|
|
}
|
|
|
|
/// Demo 4: Quantization Options
|
|
fn demo_quantization() {
|
|
println!("\n--- Demo 4: Quantization ---\n");
|
|
|
|
println!("Supported quantization formats:\n");
|
|
println!(" | Format | Bits | Memory | Accuracy | Hardware |");
|
|
println!(" |-----------|------|--------|----------|---------------|");
|
|
println!(" | FP16 | 16 | 2x | <0.1% | All |");
|
|
println!(" | BF16 | 16 | 2x | <0.2% | Ampere+ |");
|
|
println!(" | FP8 E4M3 | 8 | 4x | <0.5% | Ada/Hopper |");
|
|
println!(" | FP8 E5M2 | 8 | 4x | <0.8% | Ada/Hopper |");
|
|
println!(" | INT8 | 8 | 4x | <1% | All |");
|
|
println!(" | INT4 | 4 | 8x | <3% | All |");
|
|
|
|
println!(r#"
|
|
|
|
// Quantize a model
|
|
use rtx_inference::{{Quantizer, QuantizationScheme}};
|
|
|
|
// Weight-only quantization (INT4)
|
|
let quantizer = Quantizer::new(QuantizationScheme::INT4)
|
|
.with_group_size(128)
|
|
.with_calibration_data(&samples);
|
|
|
|
let quantized = quantizer.quantize(&model)?;
|
|
|
|
// GPTQ quantization
|
|
use rtx_compress::GPTQ;
|
|
let gptq = GPTQ::new(bits: 4, group_size: 128);
|
|
let model = gptq.quantize(&model, &calibration_data)?;
|
|
|
|
// AWQ quantization
|
|
use rtx_compress::AWQ;
|
|
let awq = AWQ::new(bits: 4);
|
|
let model = awq.quantize(&model, &calibration_data)?;
|
|
"#);
|
|
}
|
|
|
|
/// Demo 5: Complete Inference Pipeline
|
|
fn demo_complete_pipeline() {
|
|
println!("\n--- Demo 5: Complete Inference Pipeline ---\n");
|
|
|
|
println!(r#"
|
|
// Complete LLM inference pipeline example
|
|
use rtx_hub::{{load, LoadConfig, RTXDType}};
|
|
use rtx_inference::{{
|
|
InferenceEngine, InferenceConfig,
|
|
SpeculativeConfig, SpeculativeStrategy,
|
|
KVCacheConfig, EvictionPolicy,
|
|
Quantizer, QuantizationScheme,
|
|
}};
|
|
|
|
#[tokio::main]
|
|
async fn main() -> anyhow::Result<()> {{
|
|
// 1. Load model from HuggingFace Hub
|
|
let config = LoadConfig::new()
|
|
.with_dtype(RTXDType::BF16)
|
|
.with_device("cuda:0");
|
|
|
|
let model = rtx_hub::load_with_config("meta-llama/Llama-3.2-3B", config).await?;
|
|
println!("Loaded: {{}} ({{}} layers)", model.model_id, model.num_layers());
|
|
|
|
// 2. Optional: Quantize for faster inference
|
|
let quantizer = Quantizer::new(QuantizationScheme::INT8);
|
|
let model = quantizer.quantize(&model)?;
|
|
|
|
// 3. Configure KV-cache with entropy eviction
|
|
let cache_config = KVCacheConfig::new()
|
|
.with_max_tokens(16384)
|
|
.with_eviction_policy(EvictionPolicy::EntropyGuided {{
|
|
entropy_threshold: 0.3,
|
|
min_retention: 0.5,
|
|
}});
|
|
|
|
// 4. Configure speculative decoding
|
|
let spec_config = SpeculativeConfig::new()
|
|
.with_strategy(SpeculativeStrategy::Medusa {{
|
|
num_heads: 4,
|
|
tokens_per_head: 3,
|
|
}})
|
|
.with_acceptance_threshold(0.7);
|
|
|
|
// 5. Create inference engine
|
|
let engine = InferenceEngine::new(model)
|
|
.with_kv_cache(cache_config)
|
|
.with_speculative(spec_config)
|
|
.build()?;
|
|
|
|
// 6. Generate text
|
|
let prompt = "Explain quantum computing in simple terms:";
|
|
let output = engine.generate(prompt, 256).await?;
|
|
|
|
println!("Generated: {{}}", output.text);
|
|
println!("Tokens/sec: {{:.1}}", output.metrics.tokens_per_second);
|
|
println!("Speedup: {{:.2}}x (speculative)", output.metrics.speedup_ratio);
|
|
|
|
Ok(())
|
|
}}
|
|
"#);
|
|
}
|