rustytorch/memory-bank/systemPatterns.md

RustyTorch++ System Patterns

Architecture Overview (Phase-Aligned)

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

// 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

// 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

// 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

// 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

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

pub trait CustomOp: Send + Sync {
    fn forward(&self, inputs: &[Tensor]) -> Result<Tensor>;
    fn backward(&self, grad: &Tensor) -> Result<Vec<Tensor>>;
}

Backend Plugins

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

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

#[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

#[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

// 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

# 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

# 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+)