TinyTorch/docs/optimization-modules-tasks-remaining.md

Optimization Modules - Tasks Remaining

🚨 Critical Fixes Required

Module 14: Transformer Update

• Add past_key_value parameter to TransformerBlock.forward()
• Add past_key_value parameter to MultiHeadAttention.forward()
• Test that transformer still works without KV cache (backward compatibility)

Module 16: Content Migration

• Move quantization implementation from 17_quantization/quantization_dev.py to 16_quantization/
• Delete old memory content from 16_quantization/memory_dev.py
• Ensure INT8 quantization focuses on CNNs

Module 19: Complete Rewrite

• Delete autotuning content from 19_profiling/autotuning_dev.py
• Implement Timer, MemoryProfiler, FLOPCounter, ProfilerContext
• Export as tinytorch.profiling

---

📝 Module Development Tasks

Module 15: Acceleration (Minor Updates)

• Core implementation exists
• Add performance comparison visualization
• Add cache hierarchy explanation
• Test with MLP, CNN, and Transformer

Module 16: Quantization (Major Development)

• Implement INT8Quantizer class
• Build calibration dataset approach
• Create QuantizedConv2d implementation
• Add accuracy comparison tests
• Show 4x speedup with <1% accuracy loss

Module 17: Compression (New Implementation)

• Implement MagnitudePruner class
• Build structured pruning for CNN filters
• Create SparseLinear for efficient sparse ops
• Add pruning schedule (gradual vs one-shot)
• Demonstrate 70% sparsity with <2% accuracy loss

Module 18: Caching (New Implementation)

• Implement KVCache class
• Create CachedAttention module
• Update generate() method to use cache
• Show O(N²) → O(N) speedup
• Add memory growth analysis

Module 19: Profiling (Complete Rewrite)

• Build Timer with warmup and percentiles
• Implement MemoryProfiler with peak tracking
• Create FLOPCounter for operation counting
• Build ProfilerContext manager
• Add bottleneck identification tools

Module 20: Benchmarking (New Implementation)

• Create benchmarks/tinymlperf/ directory
• Build TinyMLPerf benchmark suite
• Implement hardware-independent scoring
• Create competition submission system
• Build leaderboard tracking

---

🔗 Cross-Module Integration

Dependencies to Resolve

1. Module 14 → 18: Transformer must support KV caching
2. Module 19 → 20: Profiler must be complete before benchmarking
3. Module 15-18 → 20: All optimizations must be testable in benchmarks

Testing Requirements

• Each module must have standalone tests
• Integration test: All optimizations work together
• Performance regression tests
• Accuracy preservation tests

---

📊 Success Criteria

Module Completion Checklist

• Module 15: 10-100x speedup demonstrated
• Module 16: INT8 quantization working with CNNs
• Module 17: 70% pruning achieved
• Module 18: KV cache speeds up generation 5-10x
• Module 19: Profiler accurately measures all metrics
• Module 20: Competition framework functional

Documentation Requirements

• Each module has complete README
• Connection to previous module explained
• Performance improvements documented
• Common pitfalls section included

---

🚀 Launch Plan

Phase 1: Critical Fixes (Do First)

1. Update Module 14 transformer for KV caching
2. Move quantization content to correct module
3. Clear out incorrect content from modules

Phase 2: Parallel Development (5 Agents)

Launch 5 parallel agents to develop:

• Agent 1: Module 15 (Acceleration) - Polish existing
• Agent 2: Module 16 (Quantization) - Major development
• Agent 3: Module 17 (Compression) - New implementation
• Agent 4: Module 18 (Caching) - New implementation
• Agent 5: Module 19 (Profiling) - Complete rewrite

Phase 3: Final Module (After Phase 2)

• Module 20 (Benchmarking) - Requires Module 19 completion

Phase 4: Integration Testing

• Test all optimizations together
• Verify cumulative speedups
• Ensure no conflicts between optimizations

---

⏰ Time Estimates

Quick Tasks (< 1 hour each)

• Module 14 transformer update
• Module 15 polish
• Directory/file cleanup

Medium Tasks (2-4 hours each)

• Module 16 quantization
• Module 17 compression
• Module 18 caching

Large Tasks (4-8 hours)

• Module 19 profiling (complete rewrite)
• Module 20 benchmarking
• Integration testing

Total Estimated Time: 20-30 hours of development