[Rocm] Fix to the rocm_mla_decode_rope.py returning random result (#3898)

8b681d77 · Tianxing Wu · GitHub · 194eea17 · 8b681d77
Unverified Commit 8b681d77 authored Feb 27, 2025 by Tianxing Wu Committed by GitHub Feb 26, 2025
Hide whitespace changes
Inline Side-by-side

Showing with 1 addition and 1 deletion

python/sglang/srt/layers/attention/triton_ops/rocm_mla_decode_rope.py ...g/srt/layers/attention/triton_ops/rocm_mla_decode_rope.py +1 -1

No files found.
--- a/python/sglang/srt/layers/attention/triton_ops/rocm_mla_decode_rope.py
+++ b/python/sglang/srt/layers/attention/triton_ops/rocm_mla_decode_rope.py
@@ -230,7 +230,7 @@ def _fwd_grouped_kernel_stage1_rope(
                other=0.0,
            )  # positional embedding part of keys

-            if USE_ROPE and start_n >= cur_batch_seq_len - BLOCK_N:
+            if (USE_ROPE and LAST_SPLIT) and start_n >= cur_batch_seq_len - BLOCK_N:
                k_pe = tl.where(
                    offs_n[None, :] != (split_kv_end - 1),
                    k_pe,