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

[truong-v]

Is there an existing issue for this bug?

• I have searched the existing issues

The bug has not been fixed in the latest main branch

• I have checked the latest main branch

Do you feel comfortable sharing a concise (minimal) script that reproduces the error? :)

No, I prefer not to share.

🐛 Describe the bug

The fused_rotary_emb Triton kernel in colossalai/kernel/triton/fused_rotary_embedding.py computes two of its four rotary outputs from the wrong input tensor. The kernel body is:

out_q0 = q_0 * cos - q_1 * sin   # correct (Q real half)
out_q1 = k_0 * sin + k_1 * cos   # BUG: reads K, should read Q
out_k0 = q_0 * cos - q_1 * sin   # BUG: reads Q, should read K
out_k1 = k_0 * sin + k_1 * cos   # correct (K imaginary half)

Environment

Python 3.10, Pytorch 2.5.0, CUDA 12.4, Triton 3.1.0