# rust-ai-core
[](https://crates.io/crates/rust-ai-core) [](https://pypi.org/project/rust-ai-core-bindings/) [](https://docs.rs/rust-ai-core) [](LICENSE-MIT)
**Foundation layer for the rust-ai ecosystem**, providing unified abstractions for device selection, error handling, configuration validation, and CubeCL interop.
rust-ai-core is the shared foundation that enables a future AI framework built on transparency, traceability, performance, ease of use, repeatability, and customization depth.
## Design Philosophy
**CUDA-first**: All operations prefer GPU execution. CPU is a fallback that emits warnings, not a silent alternative. This ensures users are aware when they're not getting optimal performance.
**Ecosystem Integration**: rust-ai-core serves as the foundation for all rust-ai crates, ensuring consistent behavior, unified error handling, and seamless interoperability across the entire stack.
## Features
- **Unified Device Selection**: CUDA-first with environment variable overrides
- **Common Error Types**: `CoreError` hierarchy shared across all crates
- **Trait Interfaces**: `ValidatableConfig`, `Quantize`, `Dequantize`, `GpuDispatchable`
- **CubeCL Interop**: Candle ↔ CubeCL tensor conversion utilities
## Installation
### Rust
```toml
[dependencies]
rust-ai-core = "0.2"
# With CUDA support
rust-ai-core = { version = "0.2", features = ["cuda"] }
```
### Python
```bash
pip install rust-ai-core-bindings
```
The Python package provides bindings for memory estimation, device detection, and dtype utilities.
## Quick Start
### Rust
```rust
use rust_ai_core::{get_device, DeviceConfig, CoreError, Result};
fn main() -> Result<()> {
// Get CUDA device with automatic fallback + warning
let device = get_device(&DeviceConfig::default())?;
// Or with environment-based configuration
let config = DeviceConfig::from_env();
let device = get_device(&config)?;
Ok(())
}
```
### Python
```python
import rust_ai_core_bindings as rac
# Memory estimation for AI training planning
batch, heads, seq_len, head_dim = 1, 32, 4096, 128
attention_bytes = rac.estimate_attention_memory(batch, heads, seq_len, head_dim, "bf16")
print(f"Attention layer: {attention_bytes / 1024**2:.1f} MB")
# Tensor memory estimation
shape = [1, 512, 4096]
tensor_bytes = rac.estimate_tensor_bytes(shape, "f32")
print(f"Tensor: {tensor_bytes / 1024**2:.1f} MB")
# Memory tracking for GPU budget management
tracker = rac.create_memory_tracker(limit_bytes=8 * 1024**3) # 8 GB limit
rac.tracker_allocate(tracker, tensor_bytes)
print(f"Current: {rac.tracker_allocated_bytes(tracker)} bytes")
print(f"Peak: {rac.tracker_peak_bytes(tracker)} bytes")
# Device detection
if rac.cuda_available():
device_info = rac.get_device_info()
print(f"Device: {device_info['name']}")
# Data type utilities
print(f"f32 size: {rac.bytes_per_dtype('f32')} bytes")
print(f"bf16 is float: {rac.is_floating_point_dtype('bf16')}")
print(f"bf16 accumulator: {rac.accumulator_dtype('bf16')}")
```
## Environment Variables
| `RUST_AI_FORCE_CPU` | Set to `1` or `true` to force CPU execution |
| `RUST_AI_CUDA_DEVICE` | CUDA device ordinal (default: 0) |
Legacy variables from individual crates are also supported:
- `AXOLOTL_FORCE_CPU`, `AXOLOTL_CUDA_DEVICE`
- `VSA_OPTIM_FORCE_CPU`, `VSA_OPTIM_CUDA_DEVICE`
## Modules
### `device`
CUDA-first device selection with fallback warnings.
```rust
use rust_ai_core::{get_device, DeviceConfig, warn_if_cpu};
// Explicit configuration
let config = DeviceConfig::new()
.with_cuda_device(0)
.with_force_cpu(false)
.with_crate_name("my-crate");
let device = get_device(&config)?;
// In hot paths, warn if on CPU
warn_if_cpu(&device, "my-crate");
```
### `error`
Common error types shared across the ecosystem.
```rust
use rust_ai_core::{CoreError, Result};
fn my_function() -> Result<()> {
// Use convenient constructors
if rank == 0 {
return Err(CoreError::invalid_config("rank must be positive"));
}
if shape_a != shape_b {
return Err(CoreError::shape_mismatch(shape_a, shape_b));
}
Ok(())
}
```
### `traits`
Common trait interfaces for configuration and GPU dispatch.
```rust
use rust_ai_core::{ValidatableConfig, GpuDispatchable, CoreError, Result};
#[derive(Clone)]
struct MyConfig {
rank: usize,
}
impl ValidatableConfig for MyConfig {
fn validate(&self) -> Result<()> {
if self.rank == 0 {
return Err(CoreError::invalid_config("rank must be > 0"));
}
Ok(())
}
}
```
### `cubecl` (feature: `cuda`)
CubeCL ↔ Candle tensor interop.
```rust
use rust_ai_core::{has_cubecl_cuda_support, candle_to_cubecl_handle, cubecl_to_candle_tensor};
if has_cubecl_cuda_support() {
let buffer = candle_to_cubecl_handle(&tensor)?;
// ... launch CubeCL kernel with buffer.bytes ...
let output = cubecl_to_candle_tensor(&output_buffer, &device)?;
}
```
## Crate Integration
All rust-ai crates depend on rust-ai-core as their foundation:
```
rust-ai-core (Foundation Layer)
│
├── trit-vsa - Ternary Vector Symbolic Architectures
├── bitnet-quantize - 1.58-bit quantization
├── peft-rs - LoRA, DoRA, AdaLoRA adapters
├── qlora-rs - 4-bit quantization + QLoRA
├── unsloth-rs - GPU-optimized transformer kernels
├── vsa-optim-rs - VSA optimizers and operations
├── axolotl-rs - High-level fine-tuning orchestration
└── tritter-accel - Ternary GPU acceleration
```
Each crate uses rust-ai-core's:
- **Device selection**: Consistent CUDA-first device logic
- **Error types**: Shared `CoreError` hierarchy with domain-specific extensions
- **Traits**: Common interfaces (`ValidatableConfig`, `Quantize`, etc.)
- **CubeCL interop**: Unified Candle ↔ CubeCL tensor conversion
## Public API Reference
### Core Types
- **`DeviceConfig`** - Configuration builder for device selection
- **`CoreError`** - Unified error type with domain-specific variants
- **`Result<T>`** - Type alias for `std::result::Result<T, CoreError>`
- **`TensorBuffer`** - Intermediate representation for Candle ↔ CubeCL conversion
### Traits
- **`ValidatableConfig`** - Configuration validation interface
```rust
trait ValidatableConfig: Clone + Send + Sync {
fn validate(&self) -> Result<()>;
}
```
- **`Quantize<Q>`** - Tensor quantization (full precision → quantized)
```rust
trait Quantize<Q>: Send + Sync {
fn quantize(&self, tensor: &Tensor, device: &Device) -> Result<Q>;
}
```
- **`Dequantize<Q>`** - Tensor dequantization (quantized → full precision)
```rust
trait Dequantize<Q>: Send + Sync {
fn dequantize(&self, quantized: &Q, device: &Device) -> Result<Tensor>;
}
```
- **`GpuDispatchable`** - GPU/CPU dispatch pattern for operations with both implementations
```rust
trait GpuDispatchable: Send + Sync {
type Input;
type Output;
fn dispatch_gpu(&self, input: &Self::Input, device: &Device) -> Result<Self::Output>;
fn dispatch_cpu(&self, input: &Self::Input, device: &Device) -> Result<Self::Output>;
fn dispatch(&self, input: &Self::Input, device: &Device) -> Result<Self::Output>;
fn gpu_available(&self) -> bool;
}
```
### Device Selection Functions
- **`get_device(config: &DeviceConfig) -> Result<Device>`** - Get device with CUDA-first fallback
- **`warn_if_cpu(device: &Device, crate_name: &str)`** - Emit one-time CPU warning
### CubeCL Interop (feature: `cuda`)
- **`has_cubecl_cuda_support() -> bool`** - Check if CubeCL CUDA runtime is available
- **`candle_to_cubecl_handle(tensor: &Tensor) -> Result<TensorBuffer>`** - Convert Candle tensor to CubeCL buffer
- **`cubecl_to_candle_tensor(buffer: &TensorBuffer, device: &Device) -> Result<Tensor>`** - Convert CubeCL buffer to Candle tensor
- **`allocate_output_buffer(shape: &[usize], dtype: DType) -> Result<TensorBuffer>`** - Pre-allocate CubeCL output buffer
## Error Handling Philosophy
rust-ai-core provides a structured error hierarchy that balances specificity with ergonomics:
```rust
use rust_ai_core::{CoreError, Result};
fn validate_shapes(a: &[usize], b: &[usize]) -> Result<()> {
if a.len() != b.len() {
return Err(CoreError::dim_mismatch(
format!("expected {} dims, got {}", a.len(), b.len())
));
}
if a != b {
return Err(CoreError::shape_mismatch(a, b));
}
Ok(())
}
```
Crates should extend `CoreError` with domain-specific variants:
```rust
#[derive(Error, Debug)]
pub enum PeftError {
#[error("adapter '{0}' not found")]
AdapterNotFound(String),
#[error(transparent)]
Core(#[from] CoreError),
}
```
## Future Framework Goals
rust-ai-core is designed to enable a future AI framework with these principles:
1. **Transparency** - Clear, understandable operations at every level
2. **Traceability** - Track what happens at each step with detailed logging
3. **Performance** - GPU-accelerated, optimized for production workloads
4. **Ease of use** - Simple high-level API with sensible defaults
5. **Repeatability** - Deterministic, reproducible results
6. **Customization depth** - Users can go as deep as they want, from high-level APIs to custom kernels
See [ARCHITECTURE.md](ARCHITECTURE.md) for design decisions and extension points.
## Python API Reference
The `rust-ai-core-bindings` package exposes the following functions:
### Memory Estimation
```python
# Estimate tensor memory
estimate_tensor_bytes(shape: list[int], dtype: str) -> int
# Estimate attention layer memory (Q, K, V + attention weights + output)
estimate_attention_memory(
batch_size: int,
num_heads: int,
seq_len: int,
head_dim: int,
dtype: str
) -> int
```
### Memory Tracking
```python
# Create a memory tracker with optional limit
create_memory_tracker(limit_bytes: int = 0) -> MemoryTracker
# Record allocation (raises if exceeds limit)
tracker_allocate(tracker: MemoryTracker, bytes: int)
# Record deallocation
tracker_deallocate(tracker: MemoryTracker, bytes: int)
# Query tracker state
tracker_allocated_bytes(tracker: MemoryTracker) -> int
tracker_peak_bytes(tracker: MemoryTracker) -> int
tracker_would_fit(tracker: MemoryTracker, bytes: int) -> bool
# Reset tracker to initial state
tracker_reset(tracker: MemoryTracker)
```
### Device Detection
```python
# Check CUDA availability
cuda_available() -> bool
# Get device information (returns dict with type, ordinal, name)
get_device_info(force_cpu: bool = False, cuda_device: int = 0) -> dict
```
### Data Type Utilities
```python
# Get bytes per element for dtype
bytes_per_dtype(dtype: str) -> int
# Check if dtype is floating point
is_floating_point_dtype(dtype: str) -> bool
# Get accumulator dtype for mixed precision
accumulator_dtype(dtype: str) -> str
```
### Supported Data Types
- `"f32"` - 32-bit float
- `"f16"` - 16-bit float
- `"bf16"` - Brain float 16
- `"f64"` - 64-bit float
- `"i64"` - 64-bit integer
- `"u32"` - 32-bit unsigned integer
- `"u8"` - 8-bit unsigned integer
### Logging
```python
# Initialize logging (call once at startup)
init_logging(level: str = "info") # debug, info, warn, error
```
## License
MIT License - see [LICENSE-MIT](LICENSE-MIT)
## Contributing
Contributions welcome! Please ensure:
- All public items have documentation
- Tests pass: `cargo test`
- Lints pass: `cargo clippy --all-targets --all-features`
- Code is formatted: `cargo fmt`