[NPU] Add Llama4_rope support on NPU - #1035
Merged
Merged
Conversation
Contributor
Author
|
Could you please help review this PR? @noemotiovon @TianHao324 |
Tcc0403
reviewed
Jan 21, 2026
Contributor
|
Thank you for your contribution! In the current tiling strategy, UB usage is computed based on the kernel implementation. If the UB usage exceeds the limit, we apply further tiling. With this approach, even when UB does not overflow, it may not be fully utilized. A better approach might be to derive the block size from the available UB capacity and try to utilize the UB as much as possible. Based on this idea, we may need to adjust both the tiling strategy and the kernel implementation, although I am still exploring how to realize this in practice. |
Contributor
|
At the moment, |
Contributor
Author
Thanks for the suggestion. A naive implementation with a hardcoded |
Tcc0403
requested changes
Jan 22, 2026
Tcc0403
left a comment
Tcc0403
left a comment
Collaborator
There was a problem hiding this comment.
Some nit changes about api and doc. tl.split exploration can be a follow-up PR.
Comment on lines
+113
to
+117
| # Interleaved offsets within a single head: [real0, imag0, real1, imag1, ...] | ||
| head_ptr = q_base + qh_idx[:, None] * q_head_stride | ||
| base = d_idx[None, :] * 2 | ||
| q_real = tl.load(head_ptr + base, mask=block_mask, other=0.0) | ||
| q_imag = tl.load(head_ptr + base + 1, mask=block_mask, other=0.0) |
Collaborator
There was a problem hiding this comment.
No need to be in this PR, but we can try tl.split in follow-up PR for potential improvement.
https://triton-lang.org/main/python-api/generated/triton.language.split.html#triton.language.split
Contributor
Author
There was a problem hiding this comment.
I have tried with tl.split, and is this what you expected.
Tcc0403
reviewed
Jan 23, 2026
Comment on lines
+108
to
+130
| for qh_block in range(0, n_qh, BLOCK_Q): | ||
| qh_idx = tl.arange(0, BLOCK_Q) + qh_block | ||
| qh_mask = qh_idx < n_qh | ||
| block_mask = qh_mask[:, None] & d_mask[None, :] | ||
|
|
||
| head_ptr = q_base + qh_idx[:, None] * q_head_stride | ||
| base = d_idx[None, :] * 2 | ||
|
|
||
| lane = tl.arange(0, 2)[None, None, :] | ||
|
|
||
| q_pair = tl.load( | ||
| head_ptr[:, :, None] + base[:, :, None] + lane, | ||
| mask=block_mask[:, :, None], | ||
| other=0.0, | ||
| ) | ||
|
|
||
| q_real, q_imag = tl.split(q_pair) | ||
|
|
||
| new_real = tl.math.fma(q_real, freqs_real, -(q_imag * freqs_imag)) | ||
| new_imag = tl.math.fma(q_real, freqs_imag, q_imag * freqs_real) | ||
|
|
||
| tl.store(head_ptr + base, new_real, mask=block_mask) | ||
| tl.store(head_ptr + base + 1, new_imag, mask=block_mask) |
Collaborator
There was a problem hiding this comment.
Yes kind of, but I would also leverage tl.reshape, tl.split, tl.interleave to simplify the code while ensuring efficient memory access.
Here's what I would imagine for code structure: (modified from src/liger_kernel/ops/llama4_rope.py)
# refactor to distinguish freq and qk access pattern
# freq original shape: (sl, hd // 2)
# freq tile shape: (padded_hd // 2) <- handle non-power-of-2 hidden size edge case
freq_idx = tl.arange(padded_hd // 2)
freq_mask = freq_idx < hd // 2
# q/k original shape: (bs, sl, n_heads, hd)
# q/k tile shape: (BLOCK_Q, padded_hd) <- interleaved complex last-dim layout
hd_idx = tl.arange(padded_hd)
hd_mask = hd_idx < hd
for qh_block in range(0, n_qh, BLOCK_Q):
# Load real, imag pairs from q_head
qh_idx = tl.arange(0, BLOCK_Q) + qh_block
qh_mask = qh_idx < n_qh
block_mask = qh_mask[:, None] & hd_mask[None, :]
# tile shape: (Q_BLOCK, padded_hd)
q_pair = tl.load(
q_base + qh_idx * q_head_stride + hd_idx,
mask=block_mask,
other=0.0,
)
# reshape: (Q_BLOCK, padded_hd) -> (Q_BLOCK, padded_hd // 2, 2)
# I'm not sure if `can_reorder` would affect the correctness and performance, need investigation
q_pair = q_pair.reshape(Q_BLOCK, padded_hd // 2, 2, can_reorder=True)
q_real, q_imag = tl.split(q_pair)
new_real = tl.math.fma(q_real, freqs_real, -(q_imag * freqs_imag))
new_imag = tl.math.fma(q_real, freqs_imag, q_imag * freqs_real)
# interleave new_q_real, new_q_imag to reconstruct new_q for 1 coalesced store as well
new_q_pair = tl.interleave(new_real, new_imag)
tl.store(q_base + qh_idx * q_head_stride + hd_idx, new_q_pair, mask=block_mask)
# Same for k
...note: I don't have npu access, so I can't gaurantee the correctness and performance gain. The code probably needs some modification to make it work on your machine.
Contributor
Author
There was a problem hiding this comment.
I really appreciate your professional and dedicated guidance. This part has already been refactored and tested.
Comment on lines
+8
to
+47
| def _prepare_freqs(freqs_cis: torch.Tensor, seq_len: int, head_dim_half: int): | ||
| """ | ||
| Canonicalize freqs to (seq_len, head_dim_half) real/imag tensors. | ||
|
|
||
| Supports: | ||
| - complex freqs: (..., head_dim_half) complex -> real/imag | ||
| - packed freqs: (..., 2*head_dim_half) real -> split into real/imag | ||
| """ | ||
| if freqs_cis.is_complex(): | ||
| freqs_real = freqs_cis.real | ||
| freqs_imag = freqs_cis.imag | ||
| else: | ||
| if freqs_cis.shape[-1] == 2 * head_dim_half: | ||
| freqs_real = freqs_cis[..., :head_dim_half] | ||
| freqs_imag = freqs_cis[..., head_dim_half:] | ||
| else: | ||
| raise ValueError( | ||
| f"Unexpected freqs_cis shape for non-complex input: {freqs_cis.shape}, " | ||
| f"expected last dim = {2 * head_dim_half}" | ||
| ) | ||
|
|
||
| if freqs_real.shape[-1] != head_dim_half: | ||
| raise ValueError(f"Unexpected last dim for freqs: {freqs_real.shape[-1]} (expected {head_dim_half})") | ||
|
|
||
| # Flatten leading dims -> (N, head_dim_half) | ||
| freqs_real = freqs_real.reshape(-1, head_dim_half) | ||
| freqs_imag = freqs_imag.reshape(-1, head_dim_half) | ||
|
|
||
| # Broadcast/slice to (seq_len, head_dim_half) | ||
| if freqs_real.shape[0] < seq_len: | ||
| if freqs_real.shape[0] == 1: | ||
| freqs_real = freqs_real.expand(seq_len, -1) | ||
| freqs_imag = freqs_imag.expand(seq_len, -1) | ||
| else: | ||
| raise ValueError(f"Insufficient rows in freqs: {freqs_real.shape[0]} < seq_len={seq_len}") | ||
| elif freqs_real.shape[0] > seq_len: | ||
| freqs_real = freqs_real[:seq_len] | ||
| freqs_imag = freqs_imag[:seq_len] | ||
|
|
||
| return freqs_real, freqs_imag |
Collaborator
There was a problem hiding this comment.
By the way, _prepare_freqs should be able to fuse into kernel similarly to cut host-side .reshape() cost.
However, I strongly suggest putting such change in another PR as a follow-up, so we can have a workable baseline and optimize upon it.
Contributor
Author
There was a problem hiding this comment.
I completely agree — this can be implemented as another PR after the current one is merged.
Tcc0403
reviewed
Jan 26, 2026
Comment on lines
+101
to
+102
| hd_idx = tl.arange(0, hd) | ||
| hd_mask = hd_idx < (hd) |
Collaborator
There was a problem hiding this comment.
Does this kernel work where hidden_size is not a power of two?
Contributor
Author
There was a problem hiding this comment.
Yes, it does. There is no requirement for hidden_size to be padded to a power of two.
Collaborator
There was a problem hiding this comment.
In triton-ascend tl.arange(), I thought start and end still have to be power of two. If it works without any issues, we can ignore it.
lowdy1
force-pushed
the
llama4rope_npu
branch
from
January 27, 2026 02:58
88ec9b6 to
851ff9b
Compare
Tcc0403
approved these changes
Jan 27, 2026
Tcc0403
left a comment
Tcc0403
left a comment
Collaborator
There was a problem hiding this comment.
Thank you! Feel free create an issue regarding frequencies fusion and work on it
Contributor
Author
For sure. But it seems that we can not pass a torch.complex tensor( |
Collaborator
Try torch.vew_as_real() before passing it into a triton code? It's more of an exploration of improvement. It's totally fine if it turns out not available for this approach. |
3 tasks
Tcc0403
pushed a commit
that referenced
this pull request
Feb 9, 2026
## Summary This PR is a descendant of #1035 It removes `_prepare_freqs` for simplicity and directly uses a single `freq_complex_ptr `for llama4_rope frequencies inside the Triton kernel. By avoiding extra preprocessing and reducing load, this approach simplifies the code path and improves performance. Benchmark results show better performance compared to the previous implementation. ## Testing Done Test done with `python -m pytest ./test/transformers/test_llama4_rope.py -v` Verified on Atlas 800I A2 - [ ] run `make test` to ensure correctness - [x] run `make checkstyle` to ensure code style - [ ] run `make test-convergence` to ensure convergence
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
3 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.
Summary
This PR implements a fully executable Llama4 RoPE operator for Ascend NPU.
Testing Done
Verified on Atlas 800I A2(32G)
make testto ensure correctnessmake checkstyleto ensure code stylemake test-convergenceto ensure convergence