Optimize main.py for inference efficiency and GPU throughput (torch.compile, memory tuning, warp alignment)

[abdullatifcodes]

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

[abdullatifcodes]

Request for Feedback

This PR focuses on low-level inference performance in 'main.py', especially GPU-bound optimization.

• Are these changes compatible with current production expectations?
• Would you prefer flags for enabling torch.compile/memory tracking, or keep always-on?
• Any blockers to deeper integration with Mistral’s benchmark harness or profiling toolchain?

Happy to iterate based on review. Thanks for your time and guidance!