[NPU]: remove num_stages and num_warps from NPU Triton kernels#1111
Merged
Conversation
Triton-Ascend does not support num_warps and num_stages due to hardware architecture differences; autotune only supports block size and multibuffer. Remove these parameters from all Ascend NPU kernels and kernel launches to avoid no-op / misleading config. Changes: - geglu: drop NUM_STAGES and num_warps from forward/backward kernels - embedding: drop NUM_STAGES from forward/backward kernels - fused_add_rms_norm: drop NUM_STAGES from all 4 kernels (no-tiling/tiled, forward/backward) - llama4_rope: add docstring (kernel already did not use num_stages/num_warps) - qwen2vl_mrope: drop NUM_STAGES from kernel and forward/backward launches - rms_norm: drop NUM_STAGES from all 4 kernels and launches - rope: drop NUM_STAGES from kernel and forward/backward launches - swiglu: drop NUM_STAGES and num_warps from forward/backward kernels - tvd: drop NUM_STAGES from kernel and launch
Contributor
Author
|
Hi @Tcc0403, could you please help review my code? |
3 tasks
Tcc0403
reviewed
Mar 2, 2026
Comment on lines
+20
to
+22
| """ | ||
| Triton-Ascend does not support num_warps/num_stages due to hardware differences. | ||
| """ |
Collaborator
There was a problem hiding this comment.
Thanks for the clarification! imo, we don't have to keep the reminder in every docstring, better only leave it on the npu triton kernel developement guideline somehwere. wdyt?
Contributor
Author
There was a problem hiding this comment.
I agree with your suggestion. We will consolidate everything and document it uniformly in the Ascend documentation later.
3 tasks
Tcc0403
reviewed
Mar 2, 2026
| Llama4 RoPE on Ascend NPU for interleaved complex layout: | ||
| - q/k shape: (bs, sl, n_heads, hd) | ||
| - freqs_complex_ptr: (sl, hd//2, 2) | ||
| Triton-Ascend does not support num_warps/num_stages; this kernel does not use them. |
Contributor
Author
There was a problem hiding this comment.
Sorry! I overlooked it.
Tcc0403
approved these changes
Mar 2, 2026
41 tasks
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
2 participants
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Triton-Ascend does not support num_warps and num_stages due to hardware architecture differences; autotune only supports block size and multibuffer. Remove these parameters from all Ascend NPU kernels and kernel launches to avoid no-op / misleading config.
Changes:
Hardware Type: Atlas 800I A2(32GB)
make testto ensure correctnessmake checkstyleto ensure code stylemake test-convergenceto ensure convergence