intel / intel/torch-xpu-ops

[Inductor] Optimize scaled_mm Triton Template

Open
#4,425 0 comments 0 reactions 1 assignee Claimed by @xiaowangintel View on GitHub
hw: BMG module: inductor module: ops performance
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

Open the contributing guide

Assessment

This issue has not been assessed yet.

Get new issues in your inbox

A short digest of beginner-friendly GitHub issues.