ceos/unsloth
1
0
Fork 0
Code Issues Pull requests Projects Releases Packages Wiki Activity Actions

Update rope_embedding.py

Browse source
This commit is contained in:
Daniel Han-Chen 2024-03-15 22:35:44 +11:00
commit 51c2484ffd
1 changed files with 20 additions and 30 deletions
Download patch file Download diff file Expand all files Collapse all files

50
unsloth/kernels/rope_embedding.py
View file

@ -24,7 +24,7 @@ def _rope_embedding(
Q, Q_row_stride, Q, Q_row_stride,
cos, cos_row_stride, cos, cos_row_stride,
sin, sin_row_stride, sin, sin_row_stride,
seqlen, head_dim, group_size, n_heads, seqlen, head_dim,
BACKWARD_PASS: tl.constexpr, BACKWARD_PASS: tl.constexpr,
BLOCK_SIZE : tl.constexpr, BLOCK_SIZE : tl.constexpr,
): ):
@ -34,7 +34,7 @@ def _rope_embedding(
See our blog post for more info See our blog post for more info
""" """
row_position = tl.program_id(0) row_position = tl.program_id(0)
group_head_position = tl.program_id(1) head_position = tl.program_id(1)
col_offsets = tl.arange(0, BLOCK_SIZE) col_offsets = tl.arange(0, BLOCK_SIZE)
half_head_dim = head_dim // 2 half_head_dim = head_dim // 2
mask = col_offsets < half_head_dim mask = col_offsets < half_head_dim
@ -44,25 +44,23 @@ def _rope_embedding(
cos1 = tl.load(cos + (row_position % seqlen)*cos_row_stride + \ cos1 = tl.load(cos + (row_position % seqlen)*cos_row_stride + \
half_head_dim*0 + col_offsets, mask = mask, other = 0) half_head_dim*0 + col_offsets, mask = mask, other = 0)
# For Gemma - sometimes RoPE must be done in float32 and not bfloat16
Q1 = tl.load(Q + row_position*Q_row_stride + head_position*head_dim + \
half_head_dim*0 + col_offsets, mask = mask, other = 0).to(sin1.dtype)
Q2 = tl.load(Q + row_position*Q_row_stride + head_position*head_dim + \
half_head_dim*1 + col_offsets, mask = mask, other = 0).to(sin1.dtype)
if BACKWARD_PASS: if BACKWARD_PASS:
# See our blog post for more info. # See our blog post for more info.
sin1 = -sin1 sin1 = -sin1
pass pass
head_start = group_head_position * group_size tl.store(Q + row_position*Q_row_stride + head_position*head_dim + \
head_end = tl.math.min((head_start + group_size), n_heads) half_head_dim*0 + col_offsets,
Q1*cos1 - Q2*sin1, mask = mask)
for i in range(head_start, head_end): tl.store(Q + row_position*Q_row_stride + head_position*head_dim + \
offs_q1 = row_position * Q_row_stride + i * head_dim + col_offsets half_head_dim*1 + col_offsets,
offs_q2 = row_position * Q_row_stride + i * head_dim + col_offsets + half_head_dim Q2*cos1 + Q1*sin1, mask = mask)
# For Gemma - sometimes RoPE must be done in float32 and not bfloat16
Q1 = tl.load(Q + offs_q1, mask = mask, other = 0).to(sin1.dtype)
Q2 = tl.load(Q + offs_q2, mask = mask, other = 0).to(sin1.dtype)
tl.store(Q + offs_q1, Q1*cos1 - Q2*sin1, mask = mask)
tl.store(Q + offs_q2, Q2*cos1 + Q1*sin1, mask = mask)
pass
pass pass
@ -77,16 +75,12 @@ class Fast_RoPE_Embedding(torch.autograd.Function):
# [TODO] Changing blocksize to head_dim//2 seems to have # [TODO] Changing blocksize to head_dim//2 seems to have
# some concurrency / un-deterministic issues. # some concurrency / un-deterministic issues.
BLOCK_SIZE, num_warps = calculate_settings(head_dim//2) # (head_dim//2) BLOCK_SIZE, num_warps = calculate_settings(head_dim) # (head_dim//2)
group_size = 4 # 4 or 8, too large group_size can hurt performance. _rope_embedding[(n_rows, n_heads,)](
n_groups = triton.cdiv(n_heads, group_size)
grid = (n_rows, n_groups, )
_rope_embedding[grid](
Q, Q.stride(0), Q, Q.stride(0),
cos, cos.stride(0), cos, cos.stride(0),
sin, sin.stride(0), sin, sin.stride(0),
seq_len, head_dim, group_size, n_heads, seq_len, head_dim,
BACKWARD_PASS = False, BACKWARD_PASS = False,
BLOCK_SIZE = BLOCK_SIZE, BLOCK_SIZE = BLOCK_SIZE,
num_warps = num_warps, num_warps = num_warps,
@ -108,15 +102,11 @@ class Fast_RoPE_Embedding(torch.autograd.Function):
cos = ctx.cos cos = ctx.cos
sin = ctx.sin sin = ctx.sin
group_size = 4 # 4 or 8, too large group_size can hurt performance. _rope_embedding[(n_rows, n_heads,)](
n_groups = triton.cdiv(n_heads, group_size)
grid = (n_rows, n_groups, )
_rope_embedding[grid](
dY, dY .stride(0), dY, dY .stride(0),
cos, cos.stride(0), cos, cos.stride(0),
sin, sin.stride(0), sin, sin.stride(0),
seq_len, head_dim, group_size, n_heads, seq_len, head_dim,
BACKWARD_PASS = True, BACKWARD_PASS = True,
BLOCK_SIZE = ctx.BLOCK_SIZE, BLOCK_SIZE = ctx.BLOCK_SIZE,
num_warps = ctx.num_warps, num_warps = ctx.num_warps,
@ -174,4 +164,4 @@ def inplace_rope_embedding(Q, K, cos, sin, position_ids):
Q = Slow_RoPE_Embedding.apply(Q, cos, sin, position_ids) Q = Slow_RoPE_Embedding.apply(Q, cos, sin, position_ids)
K = Slow_RoPE_Embedding.apply(K, cos, sin, position_ids) K = Slow_RoPE_Embedding.apply(K, cos, sin, position_ids)
return Q, K return Q, K
pass pass