[Inductor] Optimize scaled_mm Triton Template
- Dominant language
- Python
- Stars
- 113
- Forks
- 128
- Avg merge
- 5d 9h
- Merged PRs (30d)
- 112
Description
### 🚀 The feature, motivation and pitch
## Summary
This issue tracks ongoing optimizations to the Triton template implementation of `scaled_mm` in PyTorch Inductor.
The current Triton template leaves room for performance improvements on XPU. This tracking issue serves as a central place to organize optimization efforts and link the corresponding PRs.
## Goals
- [x] Validate Triton kernel correctness
- [ ] Improve `scaled_mm` kernel performance
- [x] Use 2D block prefetch and load for matrix B
- [x] Enable SIMD32 code generation. Triton currently generates SIMD16 instructions by default, which limits performance on some workload.
- [ ] Investigate unexpected MOV instructions between 2D Block Load and DPAS
- [ ] Investigate unexpected initialization instructions before 2D Block Load
## Planned Optimizations
### Use 2D block prefetch and load for matrix B
1. Created a Triton issue(https://github.com/intel/intel-xpu-backend-for-triton/issues/7441) to track the performance issue.
2. Optimized the FP8 tensor-wise scaled_mm Triton template by leveraging tl.make_tensor_descriptor for matrix loading, enabling 2D Block Load and prefetch. The optimization improved the kernel performance for (M=512, N=512, K=256, no bias), reducing execution latency from 70 μs to 30 μs (~57% reduction).
### Enable SIMD32 code generation. Triton currently generates SIMD16 instructions by default, which limits performance on some workload.
Triton currently generates SIMD16 instructions by default. Performance analysis shows that, for the `(M=512, N=512, K=256, no bias)` workload, using TRITON_INTEL_ENABLE_DPAS_FOR_WARP_SIZE_32=1 to generate SIMD32 instructions can improve kernel performance by approximately 20%.
### Investigate unexpected MOV instructions between 2D Block Load and DPAS
1. Created a Triton issue(https://github.com/intel/intel-xpu-backend-for-triton/issues/7561) to track the performance issue.
### Investigate unexpected initialization instructions before 2D Block Load
1. Created a Triton issue(https://github.com/intel/intel-xpu-backend-for-triton/issues/7626) to track the performance issue.
Contributor guide
Assessment
This issue has not been assessed yet.