Optimize main.py for inference efficiency and GPU throughput (torch.compile, memory tuning, warp alignment)
This pull request introduces targeted performance improvements to main.py to support high-throughput and enterprise-scale inference workloads. The changes are focused, measurable, and maintain backward compatibility with the current CLI interface.
**Key Improvements:**
Enabled torch.compile() for optimized graph execution paths on supported backends (e.g., TorchInductor)
Fused tensor operations and removed non-essential memory copies
Enabled pin_memory=True and non_blocking=True in data movement paths for efficient host-to-device transfers
Aligned token padding to 32-token warp boundaries to reduce divergence and improve occupancy
Added kernel launch timing and memory tracking support (optional CLI flags can be added for runtime profiling)
**Impact:**
These changes aim to reduce kernel launch overhead, minimize unnecessary device transfers, and align computation with optimal CUDA occupancy patterns. Benchmarks (to be submitted separately) show:
Improved throughput at batch sizes 1–4
Reduced memory consumption during interactive sessions
**Validation:**
Verified outputs against baseline for consistency
Tested in single-GPU interactive and torchrun distributed settings
Model behavior preserved across transformer and mamba types
合并状态:未合并 1 条评论