github/ColossalAI
1
0
Fork 0
mirror of https://github.com/hpcaitech/ColossalAI.git synced 2025-04-28 19:55:29 +00:00
Code Issues Packages Projects Releases Wiki Activity

[Fix/Inference]Fix CUDA Rotary Rmbedding GQA (#5623)

Browse Source 
* fix rotary embedding GQA

* change test_rotary_embdding_unpad.py KH
This commit is contained in:
yuehuayingxueluo 2024-04-23 13:44:49 +08:00 committed by GitHub
parent 5d4c1fe8f5
commit 12f10d5b0b
No known key found for this signature in database
GPG Key ID: B5690EEEBB952194
2 changed files with 9 additions and 8 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

4
extensions/csrc/cuda/fused_rotary_emb_and_cache_kernel.cu
View File

@ -115,7 +115,7 @@ __device__ void apply_k_rotary_emb_compute(
(head_offset % shard_block_size) / VecSize;
const int64_t addr_offset =
token_id * key_stride + (i / half_head_dim) * head_dim + head_offset;
const int64_t target_id = block_id * head_num * head_dim * block_size +
const int64_t target_id = block_id * kv_head_num * head_dim * block_size +
(i / half_head_dim) * block_size * head_dim +
block_offset * head_dim + head_offset;
@ -137,7 +137,7 @@ __device__ void apply_k_rotary_emb_compute(
// apply value memcopy
apply_kv_memcopy<scalar_t, VecSize>(
value, value_cache, value_stride, token_id, block_id, head_num * head_dim,
value, value_cache, value_stride, token_id, block_id, kv_head_num * head_dim,
block_size, block_offset, head_dim, half_head_dim);
}

13
tests/test_infer/test_ops/cuda/test_rotary_embdding_unpad.py
View File

@ -21,9 +21,10 @@ def numpy_allclose(x, y, rtol, atol):
@pytest.mark.parametrize("BATCH_SIZE", [4])
@pytest.mark.parametrize("SEQ_LEN", [64])
@pytest.mark.parametrize("H", [32])
@pytest.mark.parametrize("K_H", [16, 32])
@pytest.mark.parametrize("D", [64])
@pytest.mark.parametrize("dtype", [torch.float16, torch.float32])
def test_rotary_emb(BATCH_SIZE, SEQ_LEN, H, D, dtype):
def test_rotary_emb(BATCH_SIZE, SEQ_LEN, H, K_H, D, dtype):
torch.manual_seed(10)
TOTAL_TOKENS = BATCH_SIZE * SEQ_LEN
# our crafted op equals to Transformers
@ -43,12 +44,12 @@ def test_rotary_emb(BATCH_SIZE, SEQ_LEN, H, D, dtype):
max_blocks_per_sequence = (TOTAL_TOKENS + block_size - 1) // block_size
q_shape = (TOTAL_TOKENS, H, D)
q = -2.3 + 0.5 * torch.randn(q_shape, dtype=dtype, device="cuda")
k_shape = (TOTAL_TOKENS, H, D)
k_shape = (TOTAL_TOKENS, K_H, D)
k = -2.3 + 0.5 * torch.randn(k_shape, dtype=dtype, device="cuda")
cos_shape = (TOTAL_TOKENS, D // 2)
cos = -1.2 + 0.5 * torch.randn(cos_shape, dtype=dtype, device="cuda")
sin = -2.0 + 0.5 * torch.randn(cos_shape, dtype=dtype, device="cuda")
cache_shape = (BATCH_SIZE * max_blocks_per_sequence, H, block_size, D)
cache_shape = (BATCH_SIZE * max_blocks_per_sequence, K_H, block_size, D)
k_cache = torch.zeros(size=cache_shape, dtype=dtype, device="cuda")
v = torch.randn_like(k)
v_cache = torch.zeros_like(k_cache)
@ -56,8 +57,8 @@ def test_rotary_emb(BATCH_SIZE, SEQ_LEN, H, D, dtype):
block_tables = mock_alloc_block_table_and_kvcache_v2(
k, v, k_cache, v_cache, past_kv_seq_lengths, BATCH_SIZE, max_blocks_per_sequence, block_size
)
new_k = torch.randn((BATCH_SIZE, H, D), dtype=dtype, device="cuda")
new_q = torch.randn_like(new_k)
new_k = torch.randn((BATCH_SIZE, K_H, D), dtype=dtype, device="cuda")
new_q = torch.randn((BATCH_SIZE, H, D), dtype=dtype, device="cuda")
new_v = torch.randn_like(new_k)
kv_seq_lengths = past_kv_seq_lengths + 1
@ -123,4 +124,4 @@ def test_rotary_emb(BATCH_SIZE, SEQ_LEN, H, D, dtype):
if __name__ == "__main__":
test_rotary_emb(16, 64, 4, 128, torch.float16)
test_rotary_emb(16, 64, 32, 16, 128, torch.float16)