optimize triton ops to improve performance - #398
MichelleWu351 wants to merge 2 commits into
Conversation
![gemini-code-assist[bot]](https://avatars.githubusercontent.com/in/956858?s=60&v=4){kind=small}
There was a problem hiding this comment.
Code Review
This pull request introduces significant performance optimizations and refactorings across various Triton kernels (including fused add RMSNorm, GELU, LayerNorm, SDPA, and Vision RoPE) to improve NPU hardware utilization, such as adding single-pass paths, instruction reordering, and parallelizing vector/cube computations. The review feedback highlights critical issues that need to be addressed: marking dynamic dimensions (like n_rows and task counts) as tl.constexpr will trigger frequent recompilations and latency spikes during LLM serving; a duplicate definition of gelu_tanh_approx in gelu.py creates dead code; hardcoding a small block size in vision_rope.py degrades performance; and a potential division-by-zero in sdpa.py requires a safe division pattern.
Important
The consumer version of Gemini Code Assist on GitHub is being sunset. Starting June 18, 2026, new organization installations will be blocked, and all code review activity will officially cease on July 17, 2026.
For more details on the timeline and next steps, please review the Help Documentation.
| RSTD_row_stride: tl.constexpr, | ||
| n_rows: tl.constexpr, | ||
| n_cols: tl.constexpr, |
There was a problem hiding this comment.
Marking n_rows as tl.constexpr forces Triton to recompile the kernel whenever the number of rows (which is highly dynamic during LLM serving due to varying batch sizes and sequence lengths) changes. This can lead to severe compilation overhead and latency spikes. Since n_rows is only used for loop bounds and masking, it should be passed as a regular runtime argument.
| RSTD_row_stride: tl.constexpr, | |
| n_rows: tl.constexpr, | |
| n_cols: tl.constexpr, | |
| RSTD_row_stride: tl.constexpr, | |
| n_rows, | |
| n_cols: tl.constexpr, |
| RSTD_row_stride: tl.constexpr, | ||
| n_rows: tl.constexpr, | ||
| n_cols: tl.constexpr, |
There was a problem hiding this comment.
Marking n_rows as tl.constexpr forces Triton to recompile the kernel whenever the number of rows (which is highly dynamic during LLM serving due to varying batch sizes and sequence lengths) changes. This can lead to severe compilation overhead and latency spikes. Since n_rows is only used for loop bounds and masking, it should be passed as a regular runtime argument.
| RSTD_row_stride: tl.constexpr, | |
| n_rows: tl.constexpr, | |
| n_cols: tl.constexpr, | |
| RSTD_row_stride: tl.constexpr, | |
| n_rows, | |
| n_cols: tl.constexpr, |
| stride_row: tl.constexpr, | ||
| n_rows: tl.constexpr, | ||
| n_cols: tl.constexpr, |
There was a problem hiding this comment.
Marking n_rows as tl.constexpr forces Triton to recompile the kernel whenever the number of rows (which is highly dynamic during LLM serving due to varying batch sizes and sequence lengths) changes. This can lead to severe compilation overhead and latency spikes. Since n_rows is only used for loop bounds and masking, it should be passed as a regular runtime argument.
| stride_row: tl.constexpr, | |
| n_rows: tl.constexpr, | |
| n_cols: tl.constexpr, | |
| stride_row: tl.constexpr, | |
| n_rows, | |
| n_cols: tl.constexpr, |
| filter_value: tl.constexpr, | ||
| k: tl.constexpr, | ||
| TOTAL_TASKS, | ||
| TOTAL_TASKS: tl.constexpr, |
There was a problem hiding this comment.
Marking TOTAL_TASKS as tl.constexpr forces Triton to recompile the kernel whenever the total number of tasks (which depends on the dynamic batch size) changes. This can lead to severe compilation overhead and latency spikes. Since TOTAL_TASKS is only used as a loop bound, it should be passed as a regular runtime argument.
| TOTAL_TASKS: tl.constexpr, | |
| TOTAL_TASKS, |
| GROUP_CHUNKS: tl.constexpr, | ||
| BLOCK_SIZE: tl.constexpr, | ||
| DESCENDING: tl.constexpr, | ||
| MERGE_TOTAL_TASKS: tl.constexpr, |
There was a problem hiding this comment.
Marking MERGE_TOTAL_TASKS as tl.constexpr forces Triton to recompile the kernel whenever the total number of merge tasks (which depends on the dynamic batch size) changes. This can lead to severe compilation overhead and latency spikes. Since MERGE_TOTAL_TASKS is only used as a loop bound, it should be passed as a regular runtime argument.
| MERGE_TOTAL_TASKS: tl.constexpr, | |
| MERGE_TOTAL_TASKS, |
| block_m = block_m_1 | ||
| block_l = block_l_1 | ||
|
|
||
| block_o = block_o / block_l[:, None] |
There was a problem hiding this comment.
To prevent division by zero and potential NaN propagation when all elements in a row are masked out (which makes block_l zero), use a safe division pattern. This aligns with the repository's general rules for safe division in Triton kernels.
| block_o = block_o / block_l[:, None] | |
| safe_l = tl.where(block_l > 0, block_l, 1.0) | |
| block_o = tl.where(block_l[:, None] > 0, block_o / safe_l[:, None], 0.0) |
References
- In Triton kernels, prevent division by zero when a denominator like l_i (sum of probabilities) can be zero due to masking. Use a safe division pattern, for example by dividing by tl.where(l_i > 0, l_i, 1.0) and then setting the result to zero where l_i was zero.
| @triton.jit | ||
| def gelu_tanh_approx(x): | ||
| """GELU activation using exp (iter_2 optimization, kept unchanged).""" | ||
| c = 1.5957691216057308 # 2 * sqrt(2/π) | ||
| k = 0.07135889030264642 # c * 0.044715，预计算常量 | ||
| x_sq = x * x | ||
| inner = x * (c + k * x_sq) | ||
| return x / (1 + tl.math.exp(-inner)) |
There was a problem hiding this comment.
There are two consecutive definitions of gelu_tanh_approx (the first on lines 25-31 and the second on lines 34-41). The second definition completely overrides the first one in Python's namespace, making the first one dead code. Please remove the duplicate/overridden definition or rename them if both are intended to be kept.
| cast_to_fp32 = q.dtype != torch.float32 | ||
|
|
||
| token_block_size = _get_token_block_size(n_qh, n_kh) | ||
| token_block_size = 2 |
There was a problem hiding this comment.
Hardcoding token_block_size = 2 makes the helper function _get_token_block_size dead code and can significantly hurt memory coalescing and vectorization efficiency on the NPU. A block size of 2 in the token dimension is extremely small. Consider keeping the dynamic block size selection or using a larger power-of-two block size (e.g., 16 or 32) for better hardware utilization.
310d87d to
4ee7bcd
Compare
|
These optimizations have already been reworked and merged into hw/950-perf, and then landed in master through #454. This PR can be closed as superseded. |
2 participants
No description provided.