docs: add CLAUDE.md with architecture and integration reference
Co-Authored-By: Claude Sonnet 4.6 <[email protected]>
This commit is contained in:
Omar Sobh
co-authored by
Claude Sonnet 4.6
Claude Sonnet 4.6
parent
d52d359f52
commit
bc88a14fa1
@@ -0,0 +1,138 @@
|
|||||||
|
# rustytorch
|
||||||
|
|
||||||
|
> GPU-accelerated ML framework in pure Rust — full PyTorch-equivalent with CUDA/Metal/ROCm/WebGPU backends, Flash Attention, Mixture of Experts, speculative decoding, federated learning, and domain-specific stacks for medical imaging, neuroimaging, and scientific computing.
|
||||||
|
|
||||||
|
## Problems It Solves
|
||||||
|
|
||||||
|
- No production-ready Rust ML framework matches PyTorch's feature breadth with multi-backend GPU support — rustytorch fills this gap with memory safety and zero GIL limitations
|
||||||
|
- Flash Attention for 128K+ token context windows requires memory-efficient attention kernels not available in generic frameworks
|
||||||
|
- Speculative decoding (2–3× inference speedup) requires tight integration between draft and verification models — implemented as a first-class feature
|
||||||
|
- Domain-specific workloads (DICOM medical imaging, EEG/MEG neuroimaging, PDE-solving via Neural Operators) require specialized data pipelines and architectures beyond generic ML frameworks
|
||||||
|
- Distributed training across multi-GPU clusters requires NCCL-level collective communication integrated directly into the training loop
|
||||||
|
|
||||||
|
## What It's Comprised Of
|
||||||
|
|
||||||
|
108 crates organized across 8 layers:
|
||||||
|
|
||||||
|
### Layer 1: Core Infrastructure (26 crates)
|
||||||
|
| Crate | Role |
|
||||||
|
|-------|------|
|
||||||
|
| `rtx-backend` | Multi-backend abstraction (CUDA, Metal, ROCm, SYCL, WebGPU, CPU) |
|
||||||
|
| `rtx-tensor` | Zero-copy GPU-native tensor primitives, PyTorch API compatible |
|
||||||
|
| `rtx-autograd` | Tape-based automatic differentiation with Hessian-vector products |
|
||||||
|
| `rtx-memory` | Memory pooling, NUMA awareness, huge pages |
|
||||||
|
| `rtx-kernel` | GPU kernel compilation and execution (CUDA PTX + Metal MSL) |
|
||||||
|
| `rtx-flash-attention` | Memory-efficient Flash Attention for 128K+ token contexts |
|
||||||
|
| `rtx-lora` | Low-Rank Adapter (LoRA/QLoRA) fine-tuning |
|
||||||
|
| `rtx-moe` | Mixture of Experts (SwitchTransformer, ExpertChoice, TokenChoice routing) |
|
||||||
|
| `rtx-mamba` | Mamba/State Space Models (O(n) attention alternative, selective scan) |
|
||||||
|
| `rtx-speculative-decoding` | Speculative decoding for 2–3× inference speedup |
|
||||||
|
| `rtx-fusion` | Kernel fusion DSL eliminating intermediate tensor writes |
|
||||||
|
| `rtx-cubecl` | CubeCL integration for Rust → GPU kernel synthesis |
|
||||||
|
| `rtx-bindings` | PyO3 Python bindings, C FFI, ONNX/DLPack interop |
|
||||||
|
| `rtx-macros` | Derive macros (`Module`, `Config`) for ergonomic neural network API |
|
||||||
|
|
||||||
|
### Layer 2: Training & Optimization (16 crates)
|
||||||
|
| Crate | Role |
|
||||||
|
|-------|------|
|
||||||
|
| `rtx-transformers` | Complete transformer: MQA, GQA, SwiGLU, RoPE, ALiBi positional encodings |
|
||||||
|
| `rtx-distributed` | Multi-GPU/multi-node training with NCCL/RCCL, elastic recovery |
|
||||||
|
| `rtx-rl` | Reinforcement Learning: PPO, RLHF, reward modeling |
|
||||||
|
| `rtx-compress` | Model compression: pruning, distillation, AWQ/GPTQ/SmoothQuant quantization |
|
||||||
|
| `rtx-nas` | Neural Architecture Search: DARTS, PC-DARTS, FairNAS |
|
||||||
|
| `rtx-federated` | Federated learning for 100,000+ edge devices |
|
||||||
|
| `rtx-auto` | AutoML with hyperparameter optimization |
|
||||||
|
| `rtx-evolution` | Evolutionary optimization strategies |
|
||||||
|
|
||||||
|
### Layer 3: Model Architectures (10 crates)
|
||||||
|
| Crate | Role |
|
||||||
|
|-------|------|
|
||||||
|
| `rtx-vision` | Vision Transformers (ViT, ConvNeXt, patch embeddings) |
|
||||||
|
| `rtx-multimodal` | Vision-language models (CLIP-style, audio transformers) |
|
||||||
|
| `rtx-diffuse` | Diffusion models (DDIM, classifier-free guidance) |
|
||||||
|
| `rtx-timeseries` | Time series: ARIMA, Prophet, state-space models |
|
||||||
|
| `rtx-nlg` | Natural language generation, GPT-style language models |
|
||||||
|
|
||||||
|
### Layer 4: Production & Deployment (9 crates)
|
||||||
|
| Crate | Role |
|
||||||
|
|-------|------|
|
||||||
|
| `rtx-serving-api` | HTTP/gRPC model serving with load balancing |
|
||||||
|
| `rtx-inference` | High-performance inference: continuous batching, paged KV cache |
|
||||||
|
| `rtx-streaming` | Real-time ML inference with <1ms latency |
|
||||||
|
| `rtx-wasm-inference` | WebAssembly inference for browser deployment |
|
||||||
|
| `rtx-onnx` | ONNX model import/export |
|
||||||
|
|
||||||
|
### Layer 5: Specialized Domains (29 crates)
|
||||||
|
| Domain | Key Crates |
|
||||||
|
|--------|-----------|
|
||||||
|
| Medical imaging | `rtx-medical-core`, `rtx-mri2fe`, `rtx-registration`, `rtx-segmentation` (DICOM, NIfTI) |
|
||||||
|
| Neuroimaging | `rtx-neuro` suite (16 crates): EEG/MEG, LSL real-time, source localization, GNN connectivity |
|
||||||
|
| Physics-informed | `rtx-neural-operator` (FNO, DeepONet), `rtx-piddm`, `rtx-digital-twin` |
|
||||||
|
| Scientific ML | `rtx-fea` (FEM), `rtx-cfd` (CFD), `rtx-science` |
|
||||||
|
| Interpretability | `rtx-interpret` (SAE, attribution, mechanistic interpretability) |
|
||||||
|
|
||||||
|
### Meta-Crates (3 user-facing bundles)
|
||||||
|
- `rtx-core` — Essential functionality
|
||||||
|
- `rtx-training` — Complete training stack
|
||||||
|
- `rtx-inference-stack` — Production inference
|
||||||
|
|
||||||
|
## Key APIs
|
||||||
|
|
||||||
|
```rust
|
||||||
|
// Tensor API (PyTorch-like)
|
||||||
|
let x = Tensor::randn([batch, seq_len, hidden], &device)?;
|
||||||
|
let output = transformer.forward(&x)?;
|
||||||
|
let loss = cross_entropy(&output, &labels)?;
|
||||||
|
loss.backward()?;
|
||||||
|
optimizer.step()?;
|
||||||
|
|
||||||
|
// Distributed training
|
||||||
|
let trainer = DistributedTrainer::new(model, num_gpus)?;
|
||||||
|
trainer.train(&dataset)?;
|
||||||
|
|
||||||
|
// Flash Attention
|
||||||
|
let attn = FlashAttention::new(head_dim, num_heads)?;
|
||||||
|
let out = attn.forward(&q, &k, &v, causal=true)?;
|
||||||
|
|
||||||
|
// Speculative decoding
|
||||||
|
let engine = SpeculativeEngine::new(draft_model, target_model)?;
|
||||||
|
let tokens = engine.generate(&prompt, max_tokens=512)?;
|
||||||
|
```
|
||||||
|
|
||||||
|
## GPU Backend Support
|
||||||
|
|
||||||
|
| Backend | Hardware | Feature Flag |
|
||||||
|
|---------|----------|-------------|
|
||||||
|
| CUDA | NVIDIA (RTX 4090/5090, A100, H100) | `cuda` |
|
||||||
|
| Metal | Apple M-series (M1/M2/M3/M4/M5) | `metal` |
|
||||||
|
| ROCm | AMD RDNA 2/3 (RX 7900 XTX) | `rocm` |
|
||||||
|
| SYCL | Intel Arc (A770/A750) | `sycl` |
|
||||||
|
| WebGPU | Browser (WASM) | `webgpu` |
|
||||||
|
| CPU | x86/ARM with MKL | default |
|
||||||
|
|
||||||
|
## Novel Features
|
||||||
|
|
||||||
|
- **Speculative decoding** (rtx-speculative-decoding): 2–3× inference speedup via draft model + target verification — GPU-native Rust implementation
|
||||||
|
- **Kernel fusion DSL** (rtx-fusion): Eliminates intermediate tensor writes, 50%+ bandwidth savings
|
||||||
|
- **Neural Operators** (rtx-neural-operator): FNO and DeepONet for mesh-agnostic PDE solving
|
||||||
|
- **Federated learning at scale** (rtx-federated): 100K+ edge device coordination
|
||||||
|
- **MRI→FEM pipeline** (rtx-mri2fe): Medical image to finite element mesh for biomechanical simulation
|
||||||
|
|
||||||
|
## Platform Dependencies
|
||||||
|
|
||||||
|
- **RNCCL** — AllReduce collective operations for distributed training gradient aggregation
|
||||||
|
- **rustG** — GPU-accelerated compilation toolchain for kernel synthesis
|
||||||
|
|
||||||
|
## Feeds Into
|
||||||
|
|
||||||
|
- **QPUDIDP** — rustytorch provides the ML training infrastructure for surrogate model training
|
||||||
|
- **qstar** — the Python qstar surrogate models could be retrained using rustytorch
|
||||||
|
- Jupyter notebooks via PyO3 Python bindings
|
||||||
|
|
||||||
|
## Platform Role
|
||||||
|
|
||||||
|
**Layer 8 — ML Infrastructure.** rustytorch is the ML framework that powers the ML-driven components of the QuantumRedClaw platform. QPUDIDP's surrogate models (MLP, normalizing flows, MC Dropout) and the HyperNQ neural decoder in qstar-rs ultimately depend on GPU compute infrastructure. rustytorch provides that infrastructure in pure Rust, matching the platform's no-C/C++ philosophy.
|
||||||
|
|
||||||
|
## Current State
|
||||||
|
|
||||||
|
Production-ready. 108 crates. Full transformer stack, Flash Attention, MoE, speculative decoding, federated learning, NAS, medical/neuroimaging/scientific domain stacks all implemented. 30+ demo applications included.
|
||||||
Reference in New Issue
Block a user