[kernel]: fix Q/K cross-wiring and out-of-bounds accesses in fused_rotary_emb triton kernel

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,
     )

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:

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)

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: