github/ColossalAI
1
0
Fork 0
mirror of https://github.com/hpcaitech/ColossalAI.git synced 2025-09-02 17:46:42 +00:00
Code Issues Packages Projects Releases Wiki Activity

fix (#5311)

Browse Source
This commit is contained in:
Jianghai
2024-01-26 15:02:12 +08:00
committed by GitHub
parent 4f28cb43c0
commit 7ddd8b37f0
4 changed files with 149 additions and 75 deletions
Download Patch File Download Diff File Expand all files Collapse all files

2
colossalai/kernel/triton/fused_rotary_embedding.py
View File

@@ -136,7 +136,7 @@ def fused_rotary_embedding(
q_total_tokens, q_head_num, head_dim = q.shape
assert q.size(0) == k.size(0)
BLOCK_HEAD = 4
BLOCK_SIZE = 16
BLOCK_SIZE = 8
cumsum_lens = torch.cumsum(lengths, dim=0)
grid = (triton.cdiv(q_head_num, BLOCK_HEAD), triton.cdiv(q_total_tokens, BLOCK_SIZE), BLOCK_SIZE)