rustytorch/next_12-28-2025.md

RustyTorch++ Next Steps - December 28, 2025

Completed Today

TIER 0 (Critical Blockers) - DONE

• Fix gradient computation in rtx-autograd/src/func.rs
• Implement DTensor communication in rtx-distributed/src/dtensor.rs
• Implement compiled backward ops in rtx-autograd/src/compiled.rs

TIER 1 (Production Training) - DONE

• INT4 quantized matmul kernels (rtx-tensor/src/mobile/quantized.rs)
• cuSPARSELt FFI bindings (rtx-tensor/src/cusparelt/mod.rs)
• 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

• Implement CPU backward pass (naive_attention_backward)
• Implement Metal backward pass (wired trait impl to kernel)
• Add gradient tests (30 tests pass)

2. Nested vmap Support (~300 lines) - DONE

Files: crates/core/rtx-autograd/src/vmap.rs

• Track batch dimensions through nested calls (BatchDimStack)
• Handle dimension collapsing/expansion (MultiBatchedVariable)
• Add level tracking with RAII guards (VmapLevelGuard)
• Tests pass (12 vmap tests)

3. Hessian-Vector Product (hvp) (~200 lines) - DONE

Files: crates/core/rtx-autograd/src/func.rs

• Implement hvp(f, primals, tangents) with forward-over-reverse mode
• Implement hvp_finite_diff for validation
• Implement vhp (vector-Hessian product)
• Fixed f32 precision issues (use f64 intermediate, larger epsilon 1e-3)
• 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

• Implement shape guard generation (generate_guards_from_operations, generate_symbolic_guards)
• Add recompilation triggers for shape changes (GuardCheckResult, GuardFailure)
• Cache compiled graphs by shape signature (ShapeGuardManager, ShapeSignature)
• Shape dimension types: Concrete, Symbolic, Bounded, Dynamic
• Shape bucketing for cache efficiency
• 20 tests passing

5. Advanced Quantization (AWQ, GPTQ, SmoothQuant) (~1050 lines) - DONE

Files: crates/training/rtx-compress/src/quantization/advanced.rs

• AWQ (Activation-aware Weight Quantization) - per-channel scales, group quantization
• GPTQ (Accurate Post-Training Quantization) - Hessian-based, block-wise quantization
• SmoothQuant (migration difficulty from activations to weights) - configurable alpha
• QuantizedTensorData dequantization support
• 9 tests passing

---

TIER 3 (Nice to Have) - DONE

1. Distributed Checkpoint (DCP) (~850 lines) - DONE

Files: crates/training/rtx-distributed/src/dcp.rs

• Async checkpointing (AsyncSaveHandle with progress tracking)
• Sharded state dict save/load (per-rank parallel I/O)
• Resumption from partial checkpoints (configurable min_shards_for_partial)
• World size change handling (shard redistribution)
• Atomic writes with fsync
• 11 tests passing

2. Context Parallel (~550 lines) - DONE

Files: crates/training/rtx-distributed/src/context_parallel.rs

• Sequence dimension parallelism (SequenceShardInfo with even/uneven splits)
• Ring attention integration (RingAttentionState, ring_attention method)
• Long-context training support (max 128K tokens)
• KV cache distribution across CP ranks
• Async KV prefetch configuration
• 11 tests passing

3. Autograd Profiler (~800 lines) - DONE

Files: crates/core/rtx-autograd/src/profiler.rs

• Operation timing (ProfiledEvent with duration tracking)
• Memory tracking per op (MemorySnapshot, allocation/deallocation)
• Gradient flow visualization (GradientFlow, DOT export)
• Bottleneck detection (BottleneckInfo, severity analysis)
• Chrome trace export for visualization
• RAII RecordGuard for scoped profiling
• 12 tests passing

4. SDPA Backend Auto-Selection (~700 lines) - DONE

Files: crates/training/rtx-flash-attention/src/backend_selector.rs

• Automatic FlashAttention vs Math vs Memory-efficient selection
• Hardware detection for optimal backend (HardwareCapabilities)
• Fallback chain management (alternatives ranking)
• Performance-based scoring (sequence length, memory, speedup)
• Auto-tuning with performance history
• Debug mode and preferred backend options
• 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