TinyTorch/docs/optimization-modules-implementation-plan.md

TinyTorch Optimization Modules Implementation Plan

Modules 15-20: Clean, Minimal, Production-Ready

Based on PyTorch expert review - focusing on MUST HAVE features only.

---

Module 15: Acceleration ✅

Status: Already well-structured
Focus: Backend optimization with clear pedagogical progression

MUST HAVE Implementation

# 1. Educational baseline (show the journey)
def matmul_naive(A, B):  # From Module 2
def matmul_blocked(A, B):  # Cache-friendly
def matmul_numpy(A, B):  # Library backend

# 2. OptimizedBackend class
class OptimizedBackend:
    def dispatch(self, op, *args):
        # Smart operation routing
        
# 3. Performance comparison
# Show 10-100x differences between implementations

Key Learning

• Why cache-friendly matters (memory hierarchy)
• When to use optimized libraries vs custom code
• Backend dispatch patterns (like PyTorch)

---

Module 16: Quantization 🔧

Status: Needs content migration from Module 17
Focus: INT8 post-training quantization for CNNs

MUST HAVE Implementation

# 1. Simple INT8 quantization
class INT8Quantizer:
    def quantize_weights(self, weights, calibration_data):
        # Compute scale and zero point
        # Convert FP32 → INT8
        
# 2. Calibration approach
def calibrate(model, calibration_dataset):
    # Run representative data
    # Collect statistics
    # Compute optimal quantization params

# 3. Quantized operations
class QuantizedConv2d:
    # INT8 convolution implementation
    
# 4. Accuracy comparison
# Show <1% accuracy loss with 4x speedup

Key Learning

• Numerical precision trade-offs
• Why INT8 works for inference
• Calibration vs training-time quantization

---

Module 17: Compression (Pruning) 🔧

Status: Needs new implementation
Focus: Magnitude-based pruning for all architectures

MUST HAVE Implementation

# 1. Magnitude-based pruning
class MagnitudePruner:
    def prune(self, weights, sparsity=0.7):
        # Remove 70% smallest weights
        
# 2. Structured pruning for CNNs
def prune_conv_filters(conv_layer, sparsity=0.5):
    # Remove entire filters
    # Maintain conv structure

# 3. Sparse operations
class SparseLinear:
    # Efficient sparse matrix multiply
    
# 4. Accuracy tracking
# Show 70% sparsity with <2% accuracy loss

Key Learning

• Neural network redundancy
• Structured vs unstructured pruning
• When pruning fails (critical connections)

---

Module 18: Caching (KV Cache) ✅

Status: Well-scoped
Focus: KV caching for transformer autoregressive generation

MUST HAVE Implementation

# 1. KV Cache implementation
class KVCache:
    def __init__(self, max_seq_len, n_heads, head_dim):
        self.cache = {}
    
    def update(self, layer, key, value, position):
        # Store computed K,V
        
    def get(self, layer, positions):
        # Retrieve cached K,V

# 2. Modified attention with cache
class CachedAttention:
    def forward(self, x, past_kv=None):
        # Use cached values for past positions
        # Only compute new position
        
# 3. Performance demonstration
# Show O(N²) → O(N) speedup for generation

Key Learning

• Memory-compute trade-offs
• Incremental computation patterns
• Why caching matters for production inference

CRITICAL: Module 14 Transformer must be updated

# Module 14 needs this change:
class TransformerBlock:
    def forward(self, x, past_kv=None):  # ADD THIS PARAMETER
        # Support for KV caching

---

Module 19: Profiling 🔧

Status: Needs complete rewrite (currently autotuning)
Focus: Build measurement infrastructure for Module 20

MUST HAVE Implementation

# 1. Timer with statistical rigor
class Timer:
    def measure(self, func, warmup=3, runs=100):
        # Warmup runs
        # Statistical sampling
        # Return percentiles (p50, p95, p99)

# 2. Memory profiler
class MemoryProfiler:
    def profile(self, func):
        # Track allocations
        # Measure peak usage
        # Identify leaks

# 3. FLOP counter
class FLOPCounter:
    def count_ops(self, model, input):
        # Count arithmetic operations
        # Identify compute bottlenecks

# 4. Profiler context manager
class ProfilerContext:
    def __enter__(self):
        # Start profiling
    def __exit__(self):
        # Generate report

Key Learning

• Importance of warmup and statistics
• Memory vs compute bottlenecks
• How to measure, not guess

---

Module 20: Benchmarking (Competition) 🎯

Status: Needs focus on competition, not infrastructure
Focus: TinyMLPerf Olympics using Module 19 profiler

MUST HAVE Implementation

# 1. Standard benchmark models
class TinyMLPerf:
    MLP_SPRINT = load_model('benchmarks/mlp.pkl')
    CNN_MARATHON = load_model('benchmarks/cnn.pkl')
    TRANSFORMER_DECATHLON = load_model('benchmarks/transformer.pkl')

# 2. Benchmark harness using Module 19
def benchmark_model(model, profiler):
    with profiler:
        # Measure inference speed
        # Measure training speed
        # Measure memory usage
    return profiler.get_results()

# 3. Relative scoring (hardware-independent)
def compute_speedup(baseline, optimized):
    # Compare against vanilla TinyTorch
    # Return improvement ratios

# 4. Competition submission
class CompetitionSubmission:
    def validate(self):
        # Check all optimizations work
    def compute_score(self):
        # Weight different metrics
    def submit_to_leaderboard(self):
        # Update rankings

Key Learning

• Fair benchmarking methodology
• Reproducible performance measurement
• Real-world optimization strategies

---

Implementation Priority & Dependencies

Must Complete First

1. Module 14 Update: Add past_kv parameter to transformers
2. Module 16 Fix: Move quantization content from Module 17
3. Module 19 Rewrite: Replace autotuning with profiling

Development Order

1. Module 15 (Acceleration) - Already good, minor polish
2. Module 16 (Quantization) - Move content, implement INT8
3. Module 17 (Compression) - New pruning implementation
4. Module 18 (Caching) - KV cache implementation
5. Module 19 (Profiling) - Complete rewrite needed
6. Module 20 (Benchmarking) - Use Module 19 profiler

Critical Cross-Module Dependencies

• Module 14 → 18: Transformer must support KV caching
• Module 19 → 20: Profiler used in benchmarking
• Module 15-18 → 20: All optimizations tested in competition

---

Success Metrics

Each module is successful when students can:

1. Module 15: Achieve 10-100x speedup with backend optimization
2. Module 16: Quantize CNN to INT8 with <1% accuracy loss
3. Module 17: Prune 70% of parameters with <2% accuracy loss
4. Module 18: Speed up transformer generation by 5-10x with KV cache
5. Module 19: Profile and identify bottlenecks in any model
6. Module 20: Submit competition entry showing cumulative speedup

---

Common Pitfalls to Avoid

❌ Don't: Try to cover every optimization technique
✅ Do: Focus on 3-4 techniques done well

❌ Don't: Hide implementation details
✅ Do: Show clear before/after performance

❌ Don't: Make competition about absolute performance
✅ Do: Focus on relative improvement and learning

❌ Don't: Mix concepts (e.g., quantization with memory optimization)
✅ Do: One clear concept per module

---

Next Steps

1. Fix Module 14 transformer to support KV caching
2. Move quantization content to Module 16
3. Launch parallel development of Modules 15-19
4. Module 20 development after Module 19 is complete