[XCS][ttx] optimize ResidualAddLayerNorm operater performance - #404
[XCS][ttx] optimize ResidualAddLayerNorm operater performance#404Quamaly wants to merge 1 commit 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 a new NPU-optimized Triton kernel implementation for fused residual add and layer normalization (fused_add_layernorm_tle_infer_impl) and integrates it into the TTXResidualAddLayerNorm operator. The feedback highlights two important improvements: first, ensuring input tensors are contiguous before passing them to the Triton kernel to prevent memory access bugs; second, wrapping the Triton driver device property retrieval in a try-except block with a safe fallback to avoid potential initialization crashes.
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.
| use_2d_fast_wrapper = hidden_states.dim() == 2 and residual.dim() == 2 | ||
| if use_2d_fast_wrapper: | ||
| X_2d = hidden_states | ||
| R_2d = residual | ||
| else: | ||
| X_2d = hidden_states.reshape(-1, dim) | ||
| R_2d = residual.reshape(-1, dim) |
There was a problem hiding this comment.
When hidden_states is a 2D non-contiguous tensor (e.g., after a transpose or slice), reshape(-1, dim) is a no-op and returns the same non-contiguous tensor. This causes the Triton kernel to read incorrect memory locations because it assumes a contiguous layout (column stride of 1) via X_ptr_row_block + cols_off[None, :].
To prevent silent correctness bugs or out-of-bounds memory accesses, we should:
- Ensure
use_2d_fast_wrapperonly triggers when both tensors are contiguous. - Explicitly call
.contiguous()on the reshaped tensors in the fallback path.
| use_2d_fast_wrapper = hidden_states.dim() == 2 and residual.dim() == 2 | |
| if use_2d_fast_wrapper: | |
| X_2d = hidden_states | |
| R_2d = residual | |
| else: | |
| X_2d = hidden_states.reshape(-1, dim) | |
| R_2d = residual.reshape(-1, dim) | |
| use_2d_fast_wrapper = ( | |
| hidden_states.dim() == 2 | |
| and residual.dim() == 2 | |
| and hidden_states.is_contiguous() | |
| and residual.is_contiguous() | |
| ) | |
| if use_2d_fast_wrapper: | |
| X_2d = hidden_states | |
| R_2d = residual | |
| else: | |
| X_2d = hidden_states.reshape(-1, dim).contiguous() | |
| R_2d = residual.reshape(-1, dim).contiguous() |
References
- When performing pointer arithmetic to verify tensor adjacency in memory for kernel optimizations, ensure the calculation correctly accounts for the tensor's strides and the specific dimensions being checked.
|
|
||
| block_size_m = layer_norm_fwd_heuristics({"n_cols": n_cols}) | ||
| num_row_tasks = ceil_div(n_rows, block_size_m) | ||
| num_vectorcores = triton.runtime.driver.active.utils.get_device_properties("npu")["num_vectorcore"] |
There was a problem hiding this comment.
Accessing triton.runtime.driver.active directly can raise an AttributeError if the Triton driver has not been initialized yet (e.g., if this is the first Triton operation in the process). Additionally, the dictionary returned by get_device_properties might not always contain the 'num_vectorcore' key depending on the environment or Triton version.
Wrapping this in a try-except block with a safe fallback (e.g., 32 vector cores) prevents potential runtime crashes during initialization.
| num_vectorcores = triton.runtime.driver.active.utils.get_device_properties("npu")["num_vectorcore"] | |
| try: | |
| num_vectorcores = triton.runtime.driver.active.utils.get_device_properties('npu')['num_vectorcore'] | |
| except Exception: | |
| num_vectorcores = 32 |
Quamaly
changed the title
36fa938 to
c24d7bd
Compare
1 participant
TLE ResidualAddLayerNorm