edd7ceb0f2
Wave 0 (NVTX Instrumentation): - Add ml-core::nvtx module with NvtxRange RAII guard (runtime dlopen, zero overhead when absent) - Instrument 10 CUDA pipeline hot paths: experience collector, backtest evaluator, PPO collector, statistics, training guard, monitoring, replay buffer Wave 1 (Low-Effort H100 Optimizations): - L2 cache persistence: pin DQN weights (~23MB BF16) in H100's 50MB L2 via cudaCtxSetLimit(CU_LIMIT_PERSISTING_L2_CACHE_SIZE) — 3.5x effective bandwidth - Dynamic shared memory: GPU-aware tile sizing (228KB H100, 164KB A100, 100KB RTX) via CU_FUNC_ATTRIBUTE_MAX_DYNAMIC_SHARED_SIZE_BYTES opt-in - Async double buffer: sync_staging() with CUDA stream synchronization before swap Validation: 0 clippy errors, 1629 tests passed (308+410+911), 0 gpu-hotpath violations Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com>
101 lines
2.8 KiB
Rust
101 lines
2.8 KiB
Rust
//! GPU device management for ML training
|
|
//!
|
|
//! Provides centralized device selection, GPU capability detection,
|
|
//! and per-model batch size optimization.
|
|
|
|
pub mod capabilities;
|
|
pub mod l2_cache;
|
|
pub mod memory_profile;
|
|
|
|
use candle_core::Device;
|
|
use serde::{Deserialize, Serialize};
|
|
|
|
use crate::MLError;
|
|
|
|
/// Device configuration for ML training and inference.
|
|
///
|
|
/// Use `Auto` for production -- it detects CUDA if available, falls back to CPU.
|
|
/// Use `Cuda(id)` to pin to a specific GPU. Use `Cpu` to force CPU mode.
|
|
#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)]
|
|
pub enum DeviceConfig {
|
|
Cpu,
|
|
Cuda(usize),
|
|
/// Auto-detect: CUDA device 0 if available, else CPU (silent fallback -- never errors).
|
|
Auto,
|
|
}
|
|
|
|
impl DeviceConfig {
|
|
/// Resolve this config into a concrete candle `Device`.
|
|
///
|
|
/// - `Cpu` → always `Device::Cpu`
|
|
/// - `Cuda(id)` → `Device::new_cuda(id)`, errors if unavailable
|
|
/// - `Auto` → CUDA device 0 if available, else CPU
|
|
pub fn resolve(&self) -> Result<Device, MLError> {
|
|
match self {
|
|
DeviceConfig::Cpu => Ok(Device::Cpu),
|
|
DeviceConfig::Cuda(id) => Device::new_cuda(*id).map_err(|e| {
|
|
MLError::ConfigError(format!(
|
|
"CUDA device {} required but unavailable: {}",
|
|
id, e
|
|
))
|
|
}),
|
|
DeviceConfig::Auto => {
|
|
match Device::new_cuda(0) {
|
|
Ok(dev) => Ok(dev),
|
|
Err(_) => Ok(Device::Cpu),
|
|
}
|
|
}
|
|
}
|
|
}
|
|
|
|
/// Returns true if this config will attempt to use a GPU.
|
|
pub fn is_gpu(&self) -> bool {
|
|
matches!(self, DeviceConfig::Cuda(_) | DeviceConfig::Auto)
|
|
}
|
|
}
|
|
|
|
impl Default for DeviceConfig {
|
|
fn default() -> Self {
|
|
DeviceConfig::Auto
|
|
}
|
|
}
|
|
|
|
#[cfg(test)]
|
|
mod tests {
|
|
use super::*;
|
|
|
|
#[test]
|
|
fn test_cpu_always_resolves() {
|
|
let device = DeviceConfig::Cpu.resolve().unwrap();
|
|
assert!(!device.is_cuda());
|
|
}
|
|
|
|
#[test]
|
|
fn test_auto_resolves_without_error() {
|
|
let device = DeviceConfig::Auto.resolve().unwrap();
|
|
let _ = device;
|
|
}
|
|
|
|
#[test]
|
|
fn test_is_gpu() {
|
|
assert!(!DeviceConfig::Cpu.is_gpu());
|
|
assert!(DeviceConfig::Cuda(0).is_gpu());
|
|
assert!(DeviceConfig::Auto.is_gpu());
|
|
}
|
|
|
|
#[test]
|
|
fn test_default_is_auto() {
|
|
assert_eq!(DeviceConfig::default(), DeviceConfig::Auto);
|
|
}
|
|
|
|
#[test]
|
|
fn test_serde_roundtrip() {
|
|
let configs = vec![DeviceConfig::Cpu, DeviceConfig::Cuda(0), DeviceConfig::Auto];
|
|
for config in configs {
|
|
let json = serde_json::to_string(&config).unwrap();
|
|
let deserialized: DeviceConfig = serde_json::from_str(&json).unwrap();
|
|
assert_eq!(config, deserialized);
|
|
}
|
|
}
|
|
}
|