Skip to content

fix(triton): uncross Q/K halves in fused_rotary_emb#6429

Open
Solaris-star wants to merge 1 commit into
hpcaitech:mainfrom
Solaris-star:fix/6428-fused-rotary-qk-crosswire
Open

fix(triton): uncross Q/K halves in fused_rotary_emb#6429
Solaris-star wants to merge 1 commit into
hpcaitech:mainfrom
Solaris-star:fix/6428-fused-rotary-qk-crosswire

Conversation

@Solaris-star

Copy link
Copy Markdown

Summary

In fused_rotary_emb, two of the four rotary outputs were computed from the wrong tensor:

out_q0 = q_0 * cos - q_1 * sin   # correct
out_q1 = k_0 * sin + k_1 * cos   # was K, should be Q
out_k0 = q_0 * cos - q_1 * sin   # was Q, should be K
out_k1 = k_0 * sin + k_1 * cos   # correct

Correct RoPE half-rotation is:

out_x0 = x_0 * cos - x_1 * sin
out_x1 = x_0 * sin + x_1 * cos

Also fixes off_k1 using q_token_stride instead of k_token_stride.

Linked Issue

Closes #6428

out_q1 was computed from K and out_k0 from Q, so half of each tensor
got the other tensor's rotary transform. Also use k_token_stride for
off_k1 (was accidentally q_token_stride).

Fixes hpcaitech#6428

Signed-off-by: Solaris-star <820622658@qq.com>
@Solaris-star
Solaris-star requested a review from a team as a code owner July 20, 2026 15:32
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

[BUG]: Q/K cross-wiring in fused_rotary_emb Triton kernel

1 participant