github
/
ColossalAI
mirror of https://github.com/hpcaitech/ColossalAI
1
0
Fork 0
Code Issues Projects Releases Wiki Activity

[NFC] Fix code factors on inference triton kernels (#5743)

pull/5744/head
Yuanheng Zhao 2024-05-21 22:12:15 +08:00 committed by GitHub
parent c2c8c9cf17
commit bd38fe6b91
No known key found for this signature in database
GPG Key ID: B5690EEEBB952194
1 changed files with 11 additions and 11 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

22
colossalai/kernel/triton/flash_decoding.py
View File

@ -111,10 +111,10 @@ def _flash_decoding_fwd_kernel(
m = tl.max(S_ij, 0)
S_ij -= m
p_ij_hat = tl.exp(S_ij)
l = tl.sum(p_ij_hat, 0)
l_i = tl.sum(p_ij_hat, 0)
p_ij_hat = p_ij_hat.to(v_cur_block.type.element_ty)
acc += tl.sum(v_cur_block * p_ij_hat[:, None], 0)
acc = acc / l
acc = acc / l_i
offsets_mid_o = (
cur_token_idx * stride_mid_ot
@ -126,8 +126,8 @@ def _flash_decoding_fwd_kernel(
offsets_mid_o_lse = (
cur_token_idx * stride_mid_o_lset + cur_head_idx * stride_mid_o_lseh + block_start_kv * stride_mid_o_lseb
)
# logsumexp L^(j) = m^(j) + log(l^(j))
tl.store(mid_o_lse + offsets_mid_o_lse, m + tl.log(l))
# logsumexp l_i^(j) = m^(j) + log(l_i^(j))
tl.store(mid_o_lse + offsets_mid_o_lse, m + tl.log(l_i))
# Triton 2.1.0
@ -234,10 +234,10 @@ def _alibi_flash_decoding_fwd_kernel(
m = tl.max(S_ij, 0)
S_ij -= m
p_ij_hat = tl.exp(S_ij)
l = tl.sum(p_ij_hat, 0)
l_i = tl.sum(p_ij_hat, 0)
p_ij_hat = p_ij_hat.to(v_cur_block.type.element_ty)
acc += tl.sum(v_cur_block * p_ij_hat[:, None], 0)
acc = acc / l
acc = acc / l_i
offsets_mid_o = (
cur_token_idx * stride_mid_ot
@ -249,8 +249,8 @@ def _alibi_flash_decoding_fwd_kernel(
offsets_mid_o_lse = (
cur_token_idx * stride_mid_o_lset + cur_head_idx * stride_mid_o_lseh + block_start_kv * stride_mid_o_lseb
)
# logsumexp L^(j) = m^(j) + log(l^(j))
tl.store(mid_o_lse + offsets_mid_o_lse, m + tl.log(l))
# logsumexp l_i^(j) = m^(j) + log(l_i^(j))
tl.store(mid_o_lse + offsets_mid_o_lse, m + tl.log(l_i))
# Triton 2.1.0
@ -290,7 +290,7 @@ def _flash_decoding_fwd_reduce_kernel(
# BLOCK_KV == BLOCK_SIZE for now. We might want to decrease the number of blocks of kv splitted.
kv_split_num = (cur_kv_seq_len + BLOCK_KV - 1) // BLOCK_KV
m_i = float("-inf") # max logic
l = 0.0 # sum exp
l_i = 0.0 # sum exp
acc = tl.zeros([HEAD_DIM], dtype=tl.float32)
offsets_mid_o = cur_token_idx * stride_mid_ot + cur_head_idx * stride_mid_oh + offsets_dmodel
@ -304,10 +304,10 @@ def _flash_decoding_fwd_reduce_kernel(
lse -= m_ij
exp_logic = tl.exp(lse)
acc += exp_logic * mid_o_block
l = scale * l + exp_logic
l_i = scale * l_i + exp_logic
m_i = m_ij
acc = acc / l
acc = acc / l_i
offsets_O = cur_token_idx * stride_ot + cur_head_idx * stride_oh + offsets_dmodel
tl.store(O + offsets_O, acc.to(O.type.element_ty))
return