From a4ccc8eb6cd2ffd3fe7ffa936467e79b2d2c0172 Mon Sep 17 00:00:00 2001 From: RunguoLi Date: Mon, 21 Sep 2026 15:58:51 -0500 Subject: [PATCH 1/2] [kernel] fix Q/K cross-wiring and OOB accesses in fused_rotary_emb triton kernel Fixes #6428. The `fused_rotary_emb` Triton kernel had several defects, and its unit test was permanently skipped and never asserted, so none of them were caught: - out_q1 was computed from k and out_k0 from q, so the second half of q and the first half of k received the other tensor's rotation. - off_k1 used q_token_stride instead of k_token_stride, giving wrong addresses whenever k is not laid out like q (e.g. a view into packed qkv). - K_HEAD_NUM was taken from q.shape[1], so with GQA/MQA the kernel read and wrote past the end of k's heads. - cumsum_lengths was loaded unmasked with next_power_of_2(total_tokens) elements instead of num_seqs, reading past the buffer and producing wrong in-sequence positions. Rewrite the test to compare against a torch reference, actually assert, and cover GQA, a non-power-of-two number of sequences, head_dim 64/128 and q/k as non-contiguous views of a fused qkv tensor (checking v is untouched). --- .../kernel/triton/fused_rotary_embedding.py | 18 +++-- .../triton/test_fused_rotary_embedding.py | 72 ++++++++++++------- 2 files changed, 58 insertions(+), 32 deletions(-) diff --git a/colossalai/kernel/triton/fused_rotary_embedding.py b/colossalai/kernel/triton/fused_rotary_embedding.py index cf2a70f7b64e..70acdd9a36f2 100644 --- a/colossalai/kernel/triton/fused_rotary_embedding.py +++ b/colossalai/kernel/triton/fused_rotary_embedding.py @@ -18,6 +18,7 @@ def fused_rotary_emb( cos_token_stride, cos_dim_stride, q_total_tokens, + num_seqs, Q_HEAD_NUM: tl.constexpr, K_HEAD_NUM: tl.constexpr, HEAD_DIM: tl.constexpr, @@ -31,7 +32,8 @@ def fused_rotary_emb( idx = block_group_index * BLOCK_SIZE + group_token_index # original seq_idx and pos - cumsum_lens = tl.load(cumsum_lengths + tl.arange(0, N_ELEMENTS)) + seq_range = tl.arange(0, N_ELEMENTS) + cumsum_lens = tl.load(cumsum_lengths + seq_range, mask=seq_range < num_seqs, other=0) ori_seq_idx = idx - tl.max(tl.where(cumsum_lens <= idx, cumsum_lens, 0)) cos = tl.load( cos_cache + ori_seq_idx * cos_token_stride + tl.arange(0, HEAD_DIM // 2) * cos_dim_stride @@ -59,7 +61,7 @@ def fused_rotary_emb( + dim_range0[None, None, :] * head_dim_stride ) off_k1 = ( - idx * q_token_stride + idx * k_token_stride + cur_head_range[None, :, None] * k_head_stride + dim_range1[None, None, :] * head_dim_stride ) @@ -89,9 +91,9 @@ def fused_rotary_emb( ) out_q0 = q_0 * cos - q_1 * sin - out_q1 = k_0 * sin + k_1 * cos + out_q1 = q_0 * sin + q_1 * cos - out_k0 = q_0 * cos - q_1 * sin + out_k0 = k_0 * cos - k_1 * sin out_k1 = k_0 * sin + k_1 * cos # concat tl.store( @@ -127,7 +129,7 @@ def fused_rotary_embedding( """ Args: q: query tensor, [total_tokens, head_num, head_dim] - k: key tensor, [total_tokens, head_num, head_dim] + k: key tensor, [total_tokens, kv_head_num, head_dim], kv_head_num <= head_num cos: cosine for rotary embedding, [max_position_len, head_dim] sin: sine for rotary embedding, [max_position_len, head_dim] lengths [num_seqs] @@ -152,7 +154,8 @@ def fused_rotary_embedding( k_token_stride = k.stride(0) k_head_stride = k.stride(1) - k_head_num = q.shape[1] + k_head_num = k.shape[1] + num_seqs = lengths.numel() cos_token_stride = cos.stride(0) cos_dim_stride = cos.stride(1) @@ -171,11 +174,12 @@ def fused_rotary_embedding( cos_token_stride, cos_dim_stride, q_total_tokens, + num_seqs, Q_HEAD_NUM=q_head_num, K_HEAD_NUM=k_head_num, HEAD_DIM=head_dim, BLOCK_HEAD=BLOCK_HEAD, BLOCK_SIZE=BLOCK_SIZE, - N_ELEMENTS=triton.next_power_of_2(q_total_tokens), + N_ELEMENTS=triton.next_power_of_2(num_seqs), num_warps=num_warps, ) diff --git a/tests/test_infer/test_kernels/triton/test_fused_rotary_embedding.py b/tests/test_infer/test_kernels/triton/test_fused_rotary_embedding.py index b69f35740d92..ef7af4aea3e8 100644 --- a/tests/test_infer/test_kernels/triton/test_fused_rotary_embedding.py +++ b/tests/test_infer/test_kernels/triton/test_fused_rotary_embedding.py @@ -1,12 +1,8 @@ -from copy import deepcopy - import pytest import torch from packaging import version from colossalai.kernel.triton.fused_rotary_embedding import fused_rotary_embedding -from colossalai.kernel.triton.no_pad_rotary_embedding import rotary_embedding -from colossalai.kernel.triton.rotary_cache_copy import get_xine_cache try: import triton # noqa @@ -19,33 +15,59 @@ TRITON_CUDA_SUPPORT = version.parse(torch.version.cuda) > version.parse("11.4") -@pytest.mark.skip(reason="cuda error") +def torch_rotary_emb(x, cos, sin): + dim = x.shape[-1] + x0 = x[:, :, : dim // 2] + x1 = x[:, :, dim // 2 :] + cos = cos[:, None, :] + sin = sin[:, None, :] + return torch.cat((x0 * cos - x1 * sin, x0 * sin + x1 * cos), dim=-1) + + @pytest.mark.skipif(not (HAS_TRITON and TRITON_CUDA_SUPPORT), reason="requires triton") -def test_fused_rotary_emb(): - num_tokens = 20 - num_kv_heads = 32 - head_dim = 64 +@pytest.mark.parametrize("lengths", [[3, 4, 6, 7], [5, 1, 17]]) +@pytest.mark.parametrize("num_heads, num_kv_heads", [(32, 32), (32, 8)]) +@pytest.mark.parametrize("head_dim", [64, 128]) +@pytest.mark.parametrize("fused_qkv", [False, True]) +def test_fused_rotary_emb(lengths, num_heads, num_kv_heads, head_dim, fused_qkv): + torch.manual_seed(123) dtype = torch.float32 - q_shape = (num_tokens, num_kv_heads, head_dim) - q = -2.3 + 0.5 * torch.randn(q_shape, dtype=dtype, device="cuda") - q_copy = deepcopy(q) + device = "cuda" + num_tokens = sum(lengths) + + if fused_qkv: + # q and k are non-contiguous views of a packed qkv tensor, so their token strides differ + # from their own shapes (and from each other when num_kv_heads < num_heads). + qkv = torch.randn(num_tokens, num_heads + 2 * num_kv_heads, head_dim, dtype=dtype, device=device) + q = qkv[:, :num_heads] + k = qkv[:, num_heads : num_heads + num_kv_heads] + else: + q = torch.randn(num_tokens, num_heads, head_dim, dtype=dtype, device=device) + k = torch.randn(num_tokens, num_kv_heads, head_dim, dtype=dtype, device=device) + q_ref, k_ref = q.clone(), k.clone() + + cos_cache = torch.randn(1024, head_dim, dtype=dtype, device=device) + sin_cache = torch.randn(1024, head_dim, dtype=dtype, device=device) + lengths = torch.tensor(lengths, dtype=torch.int32, device=device) - k_shape = (num_tokens, num_kv_heads, head_dim) - k = -2.3 + 0.5 * torch.randn(k_shape, dtype=dtype, device="cuda") - k_copy = deepcopy(k) + # each token is rotated by its position within its own sequence + positions = torch.cat([torch.arange(n, device=device) for n in lengths.tolist()]) + cos = cos_cache[positions, : head_dim // 2] + sin = sin_cache[positions, : head_dim // 2] + q_ref = torch_rotary_emb(q_ref, cos, sin) + k_ref = torch_rotary_emb(k_ref, cos, sin) - cos_shape = (1024, head_dim) - lengths = torch.tensor([3, 4, 6, 7], device="cuda") - cos_cache = -1.2 + 0.5 * torch.randn(cos_shape, dtype=dtype, device="cuda") - sin_cache = -2.0 + 0.5 * torch.randn(cos_shape, dtype=dtype, device="cuda") + if fused_qkv: + v_before = qkv[:, num_heads + num_kv_heads :].clone() - cos, sin = get_xine_cache(lengths, cos_cache[:, : head_dim // 2], sin_cache[:, : head_dim // 2]) + fused_rotary_embedding(q, k, cos_cache, sin_cache, lengths) - rotary_embedding(q, k, cos, sin) - fused_rotary_embedding(q_copy, k_copy, cos_cache, sin_cache, lengths) - torch.allclose(q, q_copy) - torch.allclose(k, k_copy) + torch.testing.assert_close(q, q_ref) + torch.testing.assert_close(k, k_ref) + if fused_qkv: + # rotating k must not spill into the neighbouring v heads + torch.testing.assert_close(qkv[:, num_heads + num_kv_heads :], v_before) if __name__ == "__main__": - test_fused_rotary_emb() + test_fused_rotary_emb([3, 4, 6, 7], 32, 8, 64, True) From 094afb38e13c7d13977f024e72fec2d36b2f9e91 Mon Sep 17 00:00:00 2001 From: "pre-commit-ci[bot]" <66853113+pre-commit-ci[bot]@users.noreply.github.com> Date: Wed, 14 Oct 2026 00:57:17 +0000 Subject: [PATCH 2/2] [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci --- tests/test_checkpoint_io/test_gemini_checkpoint_io.py | 2 +- tests/test_zero/test_gemini/test_optim.py | 2 +- 2 files changed, 2 insertions(+), 2 deletions(-) diff --git a/tests/test_checkpoint_io/test_gemini_checkpoint_io.py b/tests/test_checkpoint_io/test_gemini_checkpoint_io.py index ba60de0bd1c1..3649b18ec8bb 100644 --- a/tests/test_checkpoint_io/test_gemini_checkpoint_io.py +++ b/tests/test_checkpoint_io/test_gemini_checkpoint_io.py @@ -23,7 +23,7 @@ ASYNC_MODES = [False, True] if HAS_TENSORNVME else [False] try: - import apex.amp + pass APEX_AVAILABLE = True except ImportError: diff --git a/tests/test_zero/test_gemini/test_optim.py b/tests/test_zero/test_gemini/test_optim.py index b89c591882d6..a0970e43db9e 100644 --- a/tests/test_zero/test_gemini/test_optim.py +++ b/tests/test_zero/test_gemini/test_optim.py @@ -15,7 +15,7 @@ from tests.kit.model_zoo import model_zoo, run_fwd_bwd try: - import apex.amp + pass APEX_AMP_AVAILABLE = True except ImportError: