443 lines
12 KiB
Markdown
443 lines
12 KiB
Markdown
# RustyTorch++ System Patterns
|
|
|
|
## Architecture Overview (Phase-Aligned)
|
|
|
|
```mermaid
|
|
graph TB
|
|
subgraph "User Layer"
|
|
Python[Python API]
|
|
Rust[Rust API]
|
|
CPP[C++ API]
|
|
WASM[WebAssembly]
|
|
end
|
|
|
|
subgraph "Phase 2: Tensor & Autograd"
|
|
API[API Layer]
|
|
TE[rtx-tensor]
|
|
AG[rtx-autograd]
|
|
IR[rtx-ir Graph IR]
|
|
end
|
|
|
|
subgraph "Phase 1: Core Compiler & Runtime"
|
|
COMP[rtx-compiler]
|
|
RT[rtx-runtime]
|
|
KERNEL[rtx-kernel]
|
|
RUSTG[rustg GPU Compiler]
|
|
end
|
|
|
|
subgraph "Phase 3-4: Advanced Execution"
|
|
DIST[rtx-dist]
|
|
SYNTH[rtx-synth]
|
|
TUNE[Auto-Tuner]
|
|
end
|
|
|
|
subgraph "Phase 5: Serving"
|
|
SERVE[rtx-serve]
|
|
SCHED[Inference Scheduler]
|
|
KV[Paged KV Cache]
|
|
end
|
|
|
|
subgraph "Backend Layer"
|
|
CUDA[CUDA sm_120]
|
|
ROCM[ROCm]
|
|
METAL[Metal]
|
|
end
|
|
|
|
subgraph "Infrastructure"
|
|
EVOLVE[rtx-evolve Phase 8]
|
|
GOV[rtx-governance]
|
|
PROF[rtx-profiler]
|
|
BENCH[rtx-bench]
|
|
end
|
|
|
|
Python --> API
|
|
API --> TE
|
|
TE --> AG
|
|
AG --> IR
|
|
IR --> COMP
|
|
COMP --> RUSTG
|
|
RUSTG --> KERNEL
|
|
KERNEL --> RT
|
|
|
|
RT --> CUDA
|
|
RT --> ROCM
|
|
RT --> METAL
|
|
|
|
DIST --> RT
|
|
SYNTH --> KERNEL
|
|
SERVE --> RT
|
|
EVOLVE --> TUNE
|
|
```
|
|
|
|
## Core Design Patterns
|
|
|
|
### 1. Trait-Based Abstraction
|
|
```rust
|
|
// Core tensor trait enabling backend flexibility
|
|
pub trait Tensor: Send + Sync {
|
|
type Elem: TensorElement;
|
|
type Shape: Shape;
|
|
type Backend: Backend;
|
|
|
|
fn shape(&self) -> &Self::Shape;
|
|
fn data(&self) -> &[Self::Elem];
|
|
fn grad(&self) -> Option<&Self>;
|
|
}
|
|
```
|
|
|
|
**Rationale**: Enables compile-time polymorphism and zero-cost abstractions.
|
|
|
|
### 2. Builder Pattern for Complex Operations
|
|
```rust
|
|
// Flexible configuration without breaking changes
|
|
Conv2d::builder()
|
|
.in_channels(3)
|
|
.out_channels(64)
|
|
.kernel_size((3, 3))
|
|
.stride(1)
|
|
.padding(1)
|
|
.build()
|
|
```
|
|
|
|
**Rationale**: Maintains API stability while allowing extensibility.
|
|
|
|
### 3. Arena Allocation for Tensors
|
|
```rust
|
|
// Memory efficiency through region-based management
|
|
pub struct TensorArena {
|
|
regions: Vec<MemoryRegion>,
|
|
allocator: BumpAllocator,
|
|
}
|
|
```
|
|
|
|
**Rationale**: Reduces fragmentation and improves cache locality.
|
|
|
|
### 4. Lazy Evaluation with Graph Construction
|
|
```rust
|
|
// Build computation graph before execution
|
|
let graph = ComputeGraph::new();
|
|
let result = graph.build(|| {
|
|
let x = tensor!([1, 2, 3]);
|
|
let y = tensor!([4, 5, 6]);
|
|
x.matmul(&y)
|
|
});
|
|
graph.execute()
|
|
```
|
|
|
|
**Rationale**: Enables optimization passes and kernel fusion.
|
|
|
|
## Component Architecture (Phase-Specific)
|
|
|
|
### Phase 1: Core Compiler & Runtime
|
|
- **rtx-compiler**: IR passes, rustg lowering, golden tests
|
|
- **rtx-runtime**: Device/Stream/Event/Graph abstractions, pooled allocator
|
|
- **rtx-kernel**: Fused kernels (MLP, LayerNorm, RoPE), kernel cache
|
|
- **Key Features**: CUDA Graphs capture/replay, AMP with loss scaler
|
|
|
|
### Phase 2: Tensor API & Autograd
|
|
- **rtx-tensor**: Shape/dtype/device, views/strides, broadcasting
|
|
- **rtx-autograd**: Tape engine, backward registry, checkpointing hooks
|
|
- **rtx-ir**: Mid-level Graph IR with cost/shape metadata, serialization
|
|
- **Key Features**: DLPack interop, deterministic mode
|
|
|
|
### Phase 3: Distributed Training
|
|
- **rtx-dist**: Process groups, NCCL/RCCL collectives, rendezvous
|
|
- **Parallelism**: DP/TP/PP/SP with auto-hybrid planner
|
|
- **Sharding**: FSDP/ZeRO-style with offload hooks
|
|
- **Key Features**: Elastic recovery, WAL checkpoints
|
|
|
|
### Phase 4: Auto-Kernel Synthesis
|
|
- **rtx-synth**: Pattern library, template emitters, synthesis flow
|
|
- **Auto-Tuner**: Bandit/grid search, hardware profile DB
|
|
- **AOT Compiler**: Graph bundles with compatibility checks
|
|
- **Key Features**: Per-SKU optimization, persistent kernel cache
|
|
|
|
### Phase 5: Inference Runtime
|
|
- **rtx-serve**: Scheduler with lanes, paged KV cache, quantization
|
|
- **Decoding**: Speculative/assisted, early-exit strategies
|
|
- **APIs**: gRPC/HTTP streaming, Python/Rust SDKs
|
|
- **Key Features**: Continuous batching, vLLM-class performance
|
|
|
|
### Phase 6-10: Platform Evolution
|
|
- **Unified Graph**: ETL + Model ops in same IR
|
|
- **Self-Optimizing**: Telemetry-driven auto-tuning
|
|
- **Governance**: SBOM, signatures, provenance
|
|
- **Evolution**: Agent-in-the-loop improvements
|
|
- **Multi-Tenant**: Global routing, quotas, federation
|
|
|
|
## Critical Implementation Paths
|
|
|
|
### Forward Pass Execution
|
|
```
|
|
1. User API Call
|
|
2. Tensor Validation
|
|
3. Graph Construction
|
|
4. Optimization Pass
|
|
5. Backend Selection
|
|
6. Kernel Dispatch
|
|
7. Memory Allocation
|
|
8. Computation
|
|
9. Result Return
|
|
```
|
|
|
|
### Backward Pass Execution
|
|
```
|
|
1. Loss Computation
|
|
2. Gradient Tape Replay
|
|
3. Chain Rule Application
|
|
4. Gradient Accumulation
|
|
5. Optimizer Step
|
|
6. Parameter Update
|
|
7. Gradient Clear
|
|
```
|
|
|
|
### Distributed Training Flow
|
|
```
|
|
1. Model Replication
|
|
2. Data Sharding
|
|
3. Forward Pass (Local)
|
|
4. Gradient Computation (Local)
|
|
5. Gradient Aggregation (All-Reduce)
|
|
6. Parameter Update (Local)
|
|
7. Synchronization Barrier
|
|
```
|
|
|
|
## Memory Management Strategy
|
|
|
|
### Tensor Lifecycle
|
|
1. **Allocation**: Via arena allocator with size classes
|
|
2. **Ownership**: Rust ownership for automatic deallocation
|
|
3. **Sharing**: Arc for multi-threaded access
|
|
4. **Gradient**: Separate allocation with weak references
|
|
|
|
### GPU Memory Hierarchy
|
|
```
|
|
Host Memory (RAM)
|
|
↓ (Async Transfer)
|
|
Device Memory (VRAM)
|
|
↓ (Kernel Launch)
|
|
Shared Memory (SM)
|
|
↓ (Thread Access)
|
|
Registers
|
|
```
|
|
|
|
### Memory Optimization Techniques
|
|
- **Gradient Checkpointing**: Trade compute for memory
|
|
- **Memory Pooling**: Reuse allocations across iterations
|
|
- **Operator Fusion**: Reduce intermediate tensor allocation
|
|
- **Quantization**: Reduce precision for memory savings
|
|
|
|
## Error Handling Philosophy
|
|
|
|
### Principle: Fail Fast, Recover Gracefully
|
|
```rust
|
|
pub type Result<T> = std::result::Result<T, RustyTorchError>;
|
|
|
|
#[derive(Error, Debug)]
|
|
pub enum RustyTorchError {
|
|
#[error("Shape mismatch: expected {expected:?}, got {got:?}")]
|
|
ShapeMismatch { expected: Shape, got: Shape },
|
|
|
|
#[error("Out of memory: requested {requested} bytes")]
|
|
OutOfMemory { requested: usize },
|
|
|
|
#[error("Backend error: {0}")]
|
|
BackendError(String),
|
|
}
|
|
```
|
|
|
|
### Error Categories
|
|
1. **Compile-Time**: Shape mismatches, type errors
|
|
2. **Runtime**: OOM, device errors, numerical instability
|
|
3. **Distributed**: Network failures, node crashes
|
|
4. **User**: Invalid configurations, API misuse
|
|
|
|
## Concurrency Model
|
|
|
|
### Thread Safety Guarantees
|
|
- **Tensors**: Send + Sync for parallel processing
|
|
- **Modules**: Arc<Mutex<>> for shared state
|
|
- **Autograd**: Thread-local tape with merge capability
|
|
|
|
### Parallelism Levels
|
|
1. **Data Parallelism**: Batch dimension splitting
|
|
2. **Model Parallelism**: Layer distribution
|
|
3. **Pipeline Parallelism**: Micro-batch processing
|
|
4. **Tensor Parallelism**: Operation splitting
|
|
|
|
## Extension Points
|
|
|
|
### Custom Operators
|
|
```rust
|
|
pub trait CustomOp: Send + Sync {
|
|
fn forward(&self, inputs: &[Tensor]) -> Result<Tensor>;
|
|
fn backward(&self, grad: &Tensor) -> Result<Vec<Tensor>>;
|
|
}
|
|
```
|
|
|
|
### Backend Plugins
|
|
```rust
|
|
pub trait Backend: Send + Sync {
|
|
fn matmul(&self, a: &Tensor, b: &Tensor) -> Result<Tensor>;
|
|
fn conv2d(&self, input: &Tensor, kernel: &Tensor) -> Result<Tensor>;
|
|
// ... other operations
|
|
}
|
|
```
|
|
|
|
### Optimization Passes
|
|
```rust
|
|
pub trait OptimizationPass {
|
|
fn optimize(&self, graph: &mut ComputeGraph) -> Result<()>;
|
|
}
|
|
```
|
|
|
|
## Performance Patterns (Phase-Aligned)
|
|
|
|
### Phase 1: Foundation Performance
|
|
- **CUDA Graphs**: Capture/replay for reduced launch overhead
|
|
- **Pooled Allocator**: Arena-based with deterministic ordering
|
|
- **Stream Scheduling**: Multi-stream with dependency DAG
|
|
- **Initial Fusions**: MLP, LayerNorm, RoPE kernels
|
|
|
|
### Phase 2: Operator Performance
|
|
- **Broadcasting**: Efficient view-based operations
|
|
- **Contiguous Checks**: Non-contiguous tensor handling
|
|
- **Stable Reductions**: Deterministic sum/mean operations
|
|
- **AMP Integration**: fp16/bf16 with loss scaling
|
|
|
|
### Phase 3: Distributed Performance
|
|
- **Overlap**: Communication/compute scheduling
|
|
- **Bucketing**: Gradient aggregation optimization
|
|
- **Topology-Aware**: PCIe/NVLink/IB bandwidth optimization
|
|
- **Sharding**: Memory reduction via FSDP/ZeRO
|
|
|
|
### Phase 4: Auto-Optimization
|
|
- **Hardware Profiles**: Per-SKU latency/bandwidth DB
|
|
- **Kernel Synthesis**: Generate specialized kernels
|
|
- **Auto-Tuning**: Bandit search for optimal params
|
|
- **AOT Compilation**: Pre-compiled graph bundles
|
|
|
|
### Phase 5: Inference Optimization
|
|
- **Continuous Batching**: Dynamic batch merging
|
|
- **Paged KV Cache**: GPU/CPU/NVMe tiering
|
|
- **Quantization**: INT8/INT4/FP8 for latency
|
|
- **Speculative Decoding**: Draft-verify acceleration
|
|
|
|
## Testing Strategy
|
|
|
|
### Test Levels
|
|
1. **Unit Tests**: Individual operations
|
|
2. **Integration Tests**: Module combinations
|
|
3. **Property Tests**: Invariant verification
|
|
4. **Benchmark Tests**: Performance regression
|
|
5. **Distributed Tests**: Multi-node scenarios
|
|
|
|
### Test Patterns
|
|
```rust
|
|
#[test]
|
|
fn test_operation() {
|
|
// Arrange
|
|
let tensor = Tensor::randn([32, 64]);
|
|
|
|
// Act
|
|
let result = tensor.relu();
|
|
|
|
// Assert
|
|
assert!(result.all(|x| x >= 0.0));
|
|
}
|
|
```
|
|
|
|
## Monitoring & Observability
|
|
|
|
### Metrics Collection
|
|
- Operation latencies
|
|
- Memory usage
|
|
- GPU utilization
|
|
- Network throughput
|
|
- Cache hit rates
|
|
|
|
### Tracing Integration
|
|
```rust
|
|
#[instrument]
|
|
pub fn matmul(a: &Tensor, b: &Tensor) -> Result<Tensor> {
|
|
span!(Level::DEBUG, "matmul", shape_a = ?a.shape(), shape_b = ?b.shape());
|
|
// Implementation
|
|
}
|
|
```
|
|
|
|
### Profiling Hooks
|
|
- Pre/post operation callbacks
|
|
- Memory allocation tracking
|
|
- Kernel execution timing
|
|
- Gradient flow visualization
|
|
|
|
## Rust 2024 Edition Patterns (December 2024)
|
|
|
|
### Float Comparison Safety Pattern
|
|
**Requirement**: Rust 2024 edition requires NaN-safe float comparisons
|
|
|
|
```rust
|
|
// BEFORE (Rust 2021 - panics on NaN)
|
|
values.sort_by(|a, b| a.partial_cmp(b).unwrap());
|
|
|
|
// AFTER (Rust 2024 - NaN-safe)
|
|
values.sort_by(|a, b| a.total_cmp(b));
|
|
```
|
|
|
|
**Applied to**: 200+ files across the workspace
|
|
**Rationale**: `total_cmp()` provides a total ordering that handles NaN, -0.0, and +0.0 correctly
|
|
|
|
### Module Organization Pattern (rtx-nlg)
|
|
**Problem**: rtx-nlg had 245+ compilation errors due to missing dependencies
|
|
**Solution**: Create local modules instead of modifying shared crates
|
|
|
|
```
|
|
crates/models/rtx-nlg/src/
|
|
├── dialogue/
|
|
│ └── mod.rs # Conversational AI module (LOCAL)
|
|
├── tensor_helpers.rs # Local tensor operations (LOCAL)
|
|
└── lib.rs # Updated exports
|
|
```
|
|
|
|
**Rationale**: Keeps core crates (rtx-tensor) stable while allowing model-specific helpers
|
|
|
|
### Dependency Hygiene Pattern
|
|
**Problem**: Legacy nom 3.2.1 pulled in via unused transitive dependency
|
|
**Solution**: Regular dependency auditing
|
|
|
|
```bash
|
|
# Check for unused dependencies
|
|
cargo machete
|
|
|
|
# Check dependency tree for specific versions
|
|
cargo tree -i nom
|
|
|
|
# Remove unused dependencies
|
|
cargo rm npy # from rtx-vision-advanced
|
|
```
|
|
|
|
**Rationale**: Prevents compatibility issues with new Rust editions
|
|
|
|
### Excluded Crates Pattern
|
|
**Rationale**: Some crates require special handling
|
|
|
|
| Crate | Pattern | Reason |
|
|
|-------|---------|--------|
|
|
| `integration_tests` | Workspace exclude | References unimplemented APIs |
|
|
| `rtx-flash-metal-attention` | Workspace exclude | Platform-specific (macOS only) |
|
|
| `demos/ui/src-tauri` | Workspace exclude | Different MSRV requirements |
|
|
|
|
```toml
|
|
# Cargo.toml workspace configuration
|
|
[workspace]
|
|
exclude = [
|
|
"demos/ui/src-tauri",
|
|
"crates/training/rtx-flash-metal-attention",
|
|
"integration_tests",
|
|
]
|
|
```
|
|
|
|
---
|
|
|
|
*System Patterns Last Updated: 2025-12-16*
|
|
*Rust Edition: 2024 (Rust 1.92+)* |