linkedin/Liger-Kernel
1
0
Fork 0
mirror of https://github.com/linkedin/Liger-Kernel.git synced 2026-10-12 02:17:47 +00:00
Code Issues Packages Projects Releases Wiki Activity

[NPU] Implement a conservative llama4_rope (#1143)

Browse Source 
## Summary
In recent versions of triton-ascend, the previous implementation of the
`llama4_rope` kernel may fail to compile. The issue appears to be
related to the use of `tl.interleave()` in the kernel implementation. So
this PR introduces a more conservative implementation of the kernel.

Additionally, the current `compute_default_tiling_strategy` is not fully
accurate for this kernel due to some two-dimensional tensors. As a
temporary workaround, the tiling configuration has been adjusted to a
larger and more conservative setting to ensure correct compilation and
execution.

## Testing Done
`python -m pytest ./test/transformers/test_llama4_rope.py -v`
`python ./benchmark/scripts/benchmark_llama4_rope.py`

Device: Atlas A2 800I

- Hardware Type: <BLANK>
- [x] run `make test` to ensure correctness
- [x] run `make checkstyle` to ensure code style
- [ ] run `make test-convergence` to ensure convergence
This commit is contained in:
lowdy1
2026-03-12 22:25:59 +00:00
committed by GitHub
parent a01a0020b4
commit 8e7641864a
1 changed files with 50 additions and 6 deletions
Download Patch File Download Diff File Expand all files Collapse all files
src/liger_kernel/ops/backends/_ascend/ops/llama4_rope.py
+50 -6
View File
@@ -83,9 +83,29 @@ def _triton_llama4_rope_npu(
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)
new_q_pair = tl.interleave(new_real, new_imag)
tl.store(head_ptr + hd_idx[None, :], new_q_pair, mask=block_mask)
pair_idx = tl.arange(0, hd // 2)
real_idx = pair_idx * 2
imag_idx = pair_idx * 2 + 1
pair_mask = pair_idx < (hd // 2)
real_mask = qh_mask[:, None] & pair_mask[None, :]
imag_mask = qh_mask[:, None] & pair_mask[None, :]
# store real
tl.store(
head_ptr + real_idx[None, :],
new_real,
mask=real_mask,
)
# store imag
tl.store(
head_ptr + imag_idx[None, :],
new_imag,
mask=imag_mask,
)
# K heads (chunked for UB)
for kh_block in range(0, n_kh, BLOCK_K):
@@ -106,9 +126,29 @@ def _triton_llama4_rope_npu(
new_real = tl.math.fma(k_real, freqs_real, -(k_imag * freqs_imag))
new_imag = tl.math.fma(k_real, freqs_imag, k_imag * freqs_real)
new_k_pair = tl.interleave(new_real, new_imag)
tl.store(head_ptr + hd_idx[None, :], new_k_pair, mask=block_mask)
pair_idx = tl.arange(0, hd // 2)
real_idx = pair_idx * 2
imag_idx = pair_idx * 2 + 1
pair_mask = pair_idx < (hd // 2)
real_mask = kh_mask[:, None] & pair_mask[None, :]
imag_mask = kh_mask[:, None] & pair_mask[None, :]
# store real
tl.store(
head_ptr + real_idx[None, :],
new_real,
mask=real_mask,
)
# store imag
tl.store(
head_ptr + imag_idx[None, :],
new_imag,
mask=imag_mask,
)
def llama4_rope_forward(q, k, freqs_cis):
@@ -139,7 +179,7 @@ def llama4_rope_forward(q, k, freqs_cis):
tile_shapes = compute_default_tiling_strategy(
safety_margin=0.90,
dtype_size=dtype_size,
memory_multiplier=12.0,
memory_multiplier=20.0,
shapes=shapes,
tiling_dims=(0, 0),
)
@@ -148,6 +188,8 @@ def llama4_rope_forward(q, k, freqs_cis):
q_tile_shape, k_tile_shape = tile_shapes
BLOCK_Q, _ = q_tile_shape
BLOCK_K, _ = k_tile_shape
BLOCK_Q = max(BLOCK_Q, 2)
BLOCK_K = max(BLOCK_K, 2)
else:
BLOCK_Q = triton.next_power_of_2(n_qh)
BLOCK_K = triton.next_power_of_2(n_kh)
@@ -207,7 +249,7 @@ def llama4_rope_backward(dq, dk, freqs_cis):
tile_shapes = compute_default_tiling_strategy(
safety_margin=0.90,
dtype_size=dtype_size,
memory_multiplier=12.0,
memory_multiplier=20.0,
shapes=shapes,
tiling_dims=(0, 0),
)
@@ -216,6 +258,8 @@ def llama4_rope_backward(dq, dk, freqs_cis):
q_tile_shape, k_tile_shape = tile_shapes
BLOCK_Q, _ = q_tile_shape
BLOCK_K, _ = k_tile_shape
BLOCK_Q = max(BLOCK_Q, 2)
BLOCK_K = max(BLOCK_K, 2)
else:
BLOCK_Q = triton.next_power_of_2(n_qh)
BLOCK_K = triton.next_power_of_2(n_kh)