# RustyTorch++ Next Steps - December 28, 2025

## Completed Today

### TIER 0 (Critical Blockers) - DONE
- [x] Fix gradient computation in `rtx-autograd/src/func.rs`
- [x] Implement DTensor communication in `rtx-distributed/src/dtensor.rs`
- [x] Implement compiled backward ops in `rtx-autograd/src/compiled.rs`

### TIER 1 (Production Training) - DONE
- [x] INT4 quantized matmul kernels (`rtx-tensor/src/mobile/quantized.rs`)
- [x] cuSPARSELt FFI bindings (`rtx-tensor/src/cusparelt/mod.rs`)
- [x] 1F1B Pipeline executor with StageModule (`rtx-distributed/src/pipeline_parallel.rs`)

**Commits:**
- `c52ba928` - TIER 1 implementation (rebased to `744a63fa`)

---

## TIER 2 (Feature Completeness) - DONE

### 1. FlashAttention CPU/Metal Backward (~400 lines) - DONE
**Files:** `crates/training/rtx-flash-attention/src/core.rs`, `src/lib.rs`
- [x] Implement CPU backward pass (`naive_attention_backward`)
- [x] Implement Metal backward pass (wired trait impl to kernel)
- [x] Add gradient tests (30 tests pass)

### 2. Nested vmap Support (~300 lines) - DONE
**Files:** `crates/core/rtx-autograd/src/vmap.rs`
- [x] Track batch dimensions through nested calls (`BatchDimStack`)
- [x] Handle dimension collapsing/expansion (`MultiBatchedVariable`)
- [x] Add level tracking with RAII guards (`VmapLevelGuard`)
- [x] Tests pass (12 vmap tests)

### 3. Hessian-Vector Product (hvp) (~200 lines) - DONE
**Files:** `crates/core/rtx-autograd/src/func.rs`
- [x] Implement `hvp(f, primals, tangents)` with forward-over-reverse mode
- [x] Implement `hvp_finite_diff` for validation
- [x] Implement `vhp` (vector-Hessian product)
- [x] Fixed f32 precision issues (use f64 intermediate, larger epsilon 1e-3)
- [x] Add numerical validation tests (7 hvp tests pass)

### 4. Dynamic Shape Guards (~500 lines) - DONE
**Files:** `crates/specialized/rtx-synthesis/src/aot_impl/shape_guards.rs`, `src/aot.rs`
- [x] Implement shape guard generation (`generate_guards_from_operations`, `generate_symbolic_guards`)
- [x] Add recompilation triggers for shape changes (`GuardCheckResult`, `GuardFailure`)
- [x] Cache compiled graphs by shape signature (`ShapeGuardManager`, `ShapeSignature`)
- [x] Shape dimension types: Concrete, Symbolic, Bounded, Dynamic
- [x] Shape bucketing for cache efficiency
- [x] 20 tests passing

### 5. Advanced Quantization (AWQ, GPTQ, SmoothQuant) (~1050 lines) - DONE
**Files:** `crates/training/rtx-compress/src/quantization/advanced.rs`
- [x] AWQ (Activation-aware Weight Quantization) - per-channel scales, group quantization
- [x] GPTQ (Accurate Post-Training Quantization) - Hessian-based, block-wise quantization
- [x] SmoothQuant (migration difficulty from activations to weights) - configurable alpha
- [x] QuantizedTensorData dequantization support
- [x] 9 tests passing

---

## TIER 3 (Nice to Have) - DONE

### 1. Distributed Checkpoint (DCP) (~850 lines) - DONE
**Files:** `crates/training/rtx-distributed/src/dcp.rs`
- [x] Async checkpointing (AsyncSaveHandle with progress tracking)
- [x] Sharded state dict save/load (per-rank parallel I/O)
- [x] Resumption from partial checkpoints (configurable min_shards_for_partial)
- [x] World size change handling (shard redistribution)
- [x] Atomic writes with fsync
- [x] 11 tests passing

### 2. Context Parallel (~550 lines) - DONE
**Files:** `crates/training/rtx-distributed/src/context_parallel.rs`
- [x] Sequence dimension parallelism (SequenceShardInfo with even/uneven splits)
- [x] Ring attention integration (RingAttentionState, ring_attention method)
- [x] Long-context training support (max 128K tokens)
- [x] KV cache distribution across CP ranks
- [x] Async KV prefetch configuration
- [x] 11 tests passing

### 3. Autograd Profiler (~800 lines) - DONE
**Files:** `crates/core/rtx-autograd/src/profiler.rs`
- [x] Operation timing (ProfiledEvent with duration tracking)
- [x] Memory tracking per op (MemorySnapshot, allocation/deallocation)
- [x] Gradient flow visualization (GradientFlow, DOT export)
- [x] Bottleneck detection (BottleneckInfo, severity analysis)
- [x] Chrome trace export for visualization
- [x] RAII RecordGuard for scoped profiling
- [x] 12 tests passing

### 4. SDPA Backend Auto-Selection (~700 lines) - DONE
**Files:** `crates/training/rtx-flash-attention/src/backend_selector.rs`
- [x] Automatic FlashAttention vs Math vs Memory-efficient selection
- [x] Hardware detection for optimal backend (HardwareCapabilities)
- [x] Fallback chain management (alternatives ranking)
- [x] Performance-based scoring (sequence length, memory, speedup)
- [x] Auto-tuning with performance history
- [x] Debug mode and preferred backend options
- [x] 12 tests passing

---

## Known Issues to Address

### Autograd
- `hessian()` is a stub - depends on working `grad()`
- `functionalize` and `linearize` not implemented
- Anomaly detection not implemented

### Distributed
- Process group caching missing in DeviceMesh
- Async checkpointing not implemented
- FSDP2 communication ops need real NCCL wiring

### Quantization
- MX GPU kernels fall back to CPU (no CUDA kernels yet)
- QLinear/QConv quantized ops not implemented
- FP8 training pipeline incomplete (no loss scaling)

### Sparsity
- cuSPARSELt FFI bindings exist but no actual CUDA library calls
- Sparse backward (gradients through sparse tensors) not implemented
- Movement pruning not implemented

---

## Performance Benchmarks Needed

1. INT4 matmul vs FP16 baseline
2. 2:4 sparsity speedup on A100/H100
3. Pipeline parallelism scaling (1→8 GPUs)
4. FSDP2 memory reduction vs FSDP1

---

## Documentation TODO

- [ ] Add PyTorch migration guide
- [ ] Create quickstart for each domain (vision, NLP, etc.)
- [ ] Document INT4/FP8 quantization workflows
- [ ] Add model zoo with pre-trained weights