chore: import upstream snapshot with attribution
This commit is contained in:
wehub-resource-sync
@@ -0,0 +1,302 @@
|
||||
import time
|
||||
from typing import Optional
|
||||
|
||||
import torch
|
||||
import triton
|
||||
import triton.language as tl
|
||||
|
||||
torch.backends.cudnn.allow_tf32 = True
|
||||
|
||||
@triton.jit
|
||||
def _fwd_recurrence(
|
||||
S, d,
|
||||
O,
|
||||
NUM_HEAD, NUM_BLOCK,
|
||||
D_MODEL_K: tl.constexpr, D_MODEL_V: tl.constexpr,
|
||||
BLOCK_MODEL_K: tl.constexpr, BLOCK_MODEL_V: tl.constexpr,
|
||||
last_kv: Optional[tl.tensor]
|
||||
):
|
||||
offset_bh = tl.program_id(0)
|
||||
offset_d = tl.program_id(1)
|
||||
offset_s = tl.program_id(2)
|
||||
|
||||
S = S + offset_bh * NUM_BLOCK * D_MODEL_K * D_MODEL_V + offset_d * D_MODEL_V * BLOCK_MODEL_K + tl.arange(0, BLOCK_MODEL_K)[:, None] * D_MODEL_V + offset_s * BLOCK_MODEL_V + tl.arange(0, BLOCK_MODEL_V)[None, :]
|
||||
|
||||
O = O + offset_bh * NUM_BLOCK * D_MODEL_K * D_MODEL_V + offset_d * D_MODEL_V * BLOCK_MODEL_K + tl.arange(0, BLOCK_MODEL_K)[:, None] * D_MODEL_V + offset_s * BLOCK_MODEL_V + tl.arange(0, BLOCK_MODEL_V)[None, :]
|
||||
|
||||
if last_kv is not None:
|
||||
last_kv = last_kv + offset_bh * D_MODEL_K * D_MODEL_V + offset_d * D_MODEL_V * BLOCK_MODEL_K + tl.arange(0, BLOCK_MODEL_K)[:, None] * D_MODEL_V + offset_s * BLOCK_MODEL_V + tl.arange(0, BLOCK_MODEL_V)[None, :]
|
||||
acc = tl.load(last_kv).to(tl.float32)
|
||||
else:
|
||||
acc = tl.zeros([BLOCK_MODEL_K, BLOCK_MODEL_V], dtype=tl.float32)
|
||||
|
||||
tl.store(O, acc.to(O.dtype.element_ty))
|
||||
O += D_MODEL_K * D_MODEL_V
|
||||
d = d + offset_bh * NUM_BLOCK
|
||||
for i in range(NUM_BLOCK-1):
|
||||
d_i = tl.load(d)
|
||||
S_i = tl.load(S)
|
||||
acc = acc * d_i + S_i
|
||||
tl.store(O, acc.to(O.dtype.element_ty))
|
||||
d += 1
|
||||
S += D_MODEL_K * D_MODEL_V
|
||||
O += D_MODEL_K * D_MODEL_V
|
||||
|
||||
|
||||
## NUM_SPLIT_K/V. K/V dimension split into NUM_SPLIT_K/V parts with equal size BLOCK_MODEL
|
||||
@triton.jit
|
||||
def _bwd_recurrence(
|
||||
S, d,
|
||||
DI, DG, DL, DS,
|
||||
NUM_HEAD, NUM_BLOCK,
|
||||
D_MODEL_K: tl.constexpr, D_MODEL_V: tl.constexpr,
|
||||
BLOCK_MODEL_K: tl.constexpr, BLOCK_MODEL_V: tl.constexpr,
|
||||
|
||||
):
|
||||
offset_bh = tl.program_id(0)
|
||||
offset_d = tl.program_id(1)
|
||||
offset_s = tl.program_id(2)
|
||||
|
||||
# offset_h = offset_bh % NUM_HEAD
|
||||
NUM_K = D_MODEL_K // BLOCK_MODEL_K
|
||||
NUM_V = D_MODEL_V // BLOCK_MODEL_V
|
||||
# skip the last chunk because it is never used
|
||||
S = S + offset_bh * NUM_BLOCK * D_MODEL_K * D_MODEL_V + offset_d * D_MODEL_V * BLOCK_MODEL_K + tl.arange(0, BLOCK_MODEL_K)[:, None] * D_MODEL_V + offset_s * BLOCK_MODEL_V + tl.arange(0, BLOCK_MODEL_V)[None, :] + (NUM_BLOCK - 2) * D_MODEL_K * D_MODEL_V
|
||||
|
||||
DI = DI + offset_bh * NUM_BLOCK * D_MODEL_K * D_MODEL_V + offset_d * D_MODEL_V * BLOCK_MODEL_K + tl.arange(0, BLOCK_MODEL_K)[:, None] * D_MODEL_V + offset_s * BLOCK_MODEL_V + tl.arange(0, BLOCK_MODEL_V)[None, :] + (NUM_BLOCK - 2) * D_MODEL_K * D_MODEL_V
|
||||
|
||||
# start from the last chunk
|
||||
DS = DS + offset_bh * NUM_BLOCK * D_MODEL_K * D_MODEL_V + offset_d * D_MODEL_V * BLOCK_MODEL_K + tl.arange(0, BLOCK_MODEL_K)[:, None] * D_MODEL_V + offset_s * BLOCK_MODEL_V + tl.arange(0, BLOCK_MODEL_V)[None, :] + (NUM_BLOCK - 1) * D_MODEL_K * D_MODEL_V
|
||||
|
||||
DG = DG + offset_bh * NUM_BLOCK * NUM_K * NUM_V + offset_d * NUM_V + offset_s + (NUM_BLOCK - 2) * NUM_K * NUM_V
|
||||
|
||||
d = d + offset_bh * NUM_BLOCK + (NUM_BLOCK - 1)
|
||||
|
||||
Dacc = tl.zeros([BLOCK_MODEL_K, BLOCK_MODEL_V], dtype=tl.float32)
|
||||
|
||||
# ignore the first chunk
|
||||
for i in range(NUM_BLOCK - 1):
|
||||
S_i = tl.load(S)
|
||||
DS_i = tl.load(DS)
|
||||
d_i = tl.load(d)
|
||||
Dacc = Dacc * d_i + DS_i
|
||||
DG_i = tl.sum(Dacc * S_i.to(tl.float32))
|
||||
|
||||
tl.store(DG, DG_i.to(DG.dtype.element_ty))
|
||||
tl.store(DI, Dacc.to(DI.dtype.element_ty))
|
||||
|
||||
S -= D_MODEL_K * D_MODEL_V
|
||||
DI -= D_MODEL_K * D_MODEL_V
|
||||
DS -= D_MODEL_K * D_MODEL_V
|
||||
DG -= NUM_K * NUM_V
|
||||
d -= 1
|
||||
|
||||
DL = DL + offset_bh * D_MODEL_K * D_MODEL_V + offset_d * D_MODEL_V * BLOCK_MODEL_K + tl.arange(0, BLOCK_MODEL_K)[:, None] * D_MODEL_V + offset_s * BLOCK_MODEL_V + tl.arange(0, BLOCK_MODEL_V)[None, :]
|
||||
DS_i = tl.load(DS)
|
||||
d_i = tl.load(d)
|
||||
Dacc = Dacc * d_i + DS_i
|
||||
tl.store(DL, Dacc.to(DL.dtype.element_ty))
|
||||
|
||||
class ChunkGateRecurrent(torch.autograd.Function):
|
||||
@staticmethod
|
||||
def forward(ctx, kv, cross_decay, last_kv=None):
|
||||
cross_decay = cross_decay.contiguous()
|
||||
kv = kv.contiguous()
|
||||
|
||||
B, H, N, D_k, D_v = kv.shape
|
||||
output = torch.empty_like(kv)
|
||||
BLOCK_MODEL_K = 64
|
||||
BLOCK_MODEL_V = 16
|
||||
|
||||
assert D_k % BLOCK_MODEL_K == 0
|
||||
assert D_v % BLOCK_MODEL_V == 0
|
||||
|
||||
grid = (B*H, D_k//BLOCK_MODEL_K, D_v//BLOCK_MODEL_V)
|
||||
ctx.grid = grid
|
||||
ctx.have_last_kv = last_kv is not None
|
||||
ctx.BLOCK_MODEL_K = BLOCK_MODEL_K
|
||||
ctx.BLOCK_MODEL_V = BLOCK_MODEL_V
|
||||
|
||||
_fwd_recurrence[grid](
|
||||
kv,
|
||||
cross_decay,
|
||||
output,
|
||||
D_MODEL_K=D_k, D_MODEL_V=D_v,
|
||||
NUM_BLOCK=N, NUM_HEAD=H,
|
||||
BLOCK_MODEL_K=BLOCK_MODEL_K,
|
||||
BLOCK_MODEL_V=BLOCK_MODEL_V,
|
||||
last_kv=last_kv
|
||||
)
|
||||
|
||||
ctx.save_for_backward(output, cross_decay)
|
||||
return output
|
||||
|
||||
@staticmethod
|
||||
def backward(ctx, DO):
|
||||
DO = DO.contiguous()
|
||||
|
||||
output, cross_decay = ctx.saved_tensors
|
||||
|
||||
B, H, N, D_k, D_v = output.shape
|
||||
|
||||
BLOCK_MODEL_K = 64
|
||||
BLOCK_MODEL_V = 16
|
||||
|
||||
grid = (B*H, D_k//BLOCK_MODEL_K, D_v//BLOCK_MODEL_V)
|
||||
|
||||
DI = torch.empty_like(DO)
|
||||
DG = torch.empty(B*H, N, D_k//BLOCK_MODEL_K, D_v//BLOCK_MODEL_V, device=cross_decay.device, dtype=cross_decay.dtype)
|
||||
DL = torch.empty(B, H, D_k, D_v, device=output.device, dtype=output.dtype)
|
||||
_bwd_recurrence[grid](
|
||||
output, cross_decay,
|
||||
DI, DG, DL, DO,
|
||||
NUM_HEAD=H, NUM_BLOCK = N,
|
||||
D_MODEL_K = D_k,
|
||||
D_MODEL_V = D_v,
|
||||
BLOCK_MODEL_K=BLOCK_MODEL_K,
|
||||
BLOCK_MODEL_V=BLOCK_MODEL_V,
|
||||
)
|
||||
|
||||
DI[:, :, -1] = 0
|
||||
DG[:, -1] = 0
|
||||
DG = DG.view(B, H, N, -1).sum(dim=-1)
|
||||
return DI, DG, DL if ctx.have_last_kv else None
|
||||
|
||||
def cross_chunk(q, k, v, g, last_hidden_state=None):
|
||||
kv = k.transpose(-1, -2) @ (v * (-g + g[:, :, :, -1, None]).exp()[..., None].to(v.dtype))
|
||||
cross_decay = g[:, :, :, -1].exp().to(kv.dtype)
|
||||
S = chunk_gate_recurrent(kv, cross_decay, last_hidden_state)
|
||||
cross = (q * g[..., None].exp().to(q.dtype)) @ S
|
||||
return cross
|
||||
|
||||
@torch.compile
|
||||
def inner_chunk(q, k, v, g):
|
||||
attn = q @ k.transpose(-1, -2)
|
||||
causal_mask = torch.full([q.shape[-2], q.shape[-2]], float("-inf"), device=q.device).triu(1).type_as(q)
|
||||
attn = attn * (g[..., None] - g[..., None, :] + causal_mask).exp().to(attn.dtype)
|
||||
inner = attn @ v
|
||||
return inner
|
||||
|
||||
def chunk_gate_retention(q, k, v, g, chunk_size=64, last_hidden_state=None):
|
||||
bsz, num_head, tgt_len, key_dim = q.shape
|
||||
head_dim = v.shape[-1]
|
||||
num_chunk = tgt_len // chunk_size
|
||||
q = q.view(bsz, num_head, num_chunk, chunk_size, key_dim)
|
||||
k = k.view(bsz, num_head, num_chunk, chunk_size, key_dim) * (key_dim ** -0.5)
|
||||
v = v.view(bsz, num_head, num_chunk, chunk_size, head_dim)
|
||||
g = g.view(bsz, num_head, num_chunk, chunk_size)
|
||||
g = g.float().cumsum(-1)
|
||||
cross = cross_chunk(q, k, v, g, last_hidden_state=last_hidden_state)
|
||||
inner = inner_chunk(q, k, v, g)
|
||||
o = cross + inner
|
||||
return o.view(bsz, num_head, tgt_len, head_dim)
|
||||
|
||||
# for long sequence parallelism
|
||||
def hier_chunk_gate_retention(q, k, v, g, chunk_size=64, hier_chunk_size=16384):
|
||||
bsz, num_head, tgt_len, key_dim = q.shape
|
||||
head_dim = v.shape[-1]
|
||||
num_hier_chunk = tgt_len // hier_chunk_size
|
||||
assert tgt_len == num_hier_chunk * hier_chunk_size
|
||||
|
||||
q = q.view(bsz, num_head, num_hier_chunk, hier_chunk_size, key_dim)
|
||||
k = k.view(bsz, num_head, num_hier_chunk, hier_chunk_size, key_dim)
|
||||
v = v.view(bsz, num_head, num_hier_chunk, hier_chunk_size, head_dim)
|
||||
g = g.view(bsz, num_head, num_hier_chunk, hier_chunk_size)
|
||||
hier_cross = cross_chunk(q, k * (key_dim ** -0.5), v, g.float().cumsum(-1)).view(bsz, num_head, tgt_len, head_dim)
|
||||
|
||||
qi = q.transpose(1, 2).reshape(bsz * num_hier_chunk, num_head, hier_chunk_size, key_dim)
|
||||
ki = k.transpose(1, 2).reshape(bsz * num_hier_chunk, num_head, hier_chunk_size, key_dim)
|
||||
vi = v.transpose(1, 2).reshape(bsz * num_hier_chunk, num_head, hier_chunk_size, head_dim)
|
||||
gi = g.transpose(1, 2).reshape(bsz * num_hier_chunk, num_head, hier_chunk_size)
|
||||
inner_cross = chunk_gate_retention(qi, ki, vi, gi, chunk_size)
|
||||
|
||||
inner_cross = inner_cross.view(bsz, num_hier_chunk, num_head, hier_chunk_size, head_dim).transpose(1, 2).reshape(bsz, num_head, tgt_len, head_dim)
|
||||
o = hier_cross + inner_cross
|
||||
return o
|
||||
|
||||
def recurrent_gate_retention(q, k, v, g, incremental_state):
|
||||
bsz, num_head, _, key_dim = q.shape
|
||||
k *= key_dim ** -0.5
|
||||
g = g.view(bsz, num_head, 1, 1).float().exp()
|
||||
kv = k.transpose(-1, -2) * v
|
||||
if "last_hidden_state" in incremental_state:
|
||||
prev_kv = incremental_state["last_hidden_state"]
|
||||
kv += prev_kv * g.to(prev_kv.dtype)
|
||||
|
||||
incremental_state["last_hidden_state"] = kv
|
||||
o = q @ kv
|
||||
return o
|
||||
|
||||
def parallel_gate_retention(q, k, v, g):
|
||||
k = k * (q.shape[-1] ** -0.5)
|
||||
causal_mask = torch.full([q.shape[-2], q.shape[-2]], float("-inf"), device=q.device).triu(1).type_as(q)
|
||||
g = g.float().cumsum(-1)
|
||||
mask = g[..., None] - g[..., None, :] + causal_mask
|
||||
mask = mask.exp()
|
||||
|
||||
attn = q @ k.transpose(-1, -2)
|
||||
attn = attn * mask.to(attn.dtype)
|
||||
o = attn @ v
|
||||
return o
|
||||
|
||||
def naive_kv_recurrent(kv, cross_decay, last_kv=None):
|
||||
BSZ, NUM_HEAD, NUM_BLOCK, D_MODEL_K, D_MODEL_V = kv.shape
|
||||
kv_recurrent = []
|
||||
kv_state = torch.zeros(BSZ, NUM_HEAD, D_MODEL_K, D_MODEL_V, dtype=kv.dtype, device="cuda") if last_kv is None else last_kv
|
||||
# accumulate kv by loop
|
||||
for i in range(NUM_BLOCK):
|
||||
kv_recurrent.append(kv_state)
|
||||
kv_state = kv_state * cross_decay[:, :, i, None, None] + kv[:, :, i]
|
||||
|
||||
kv_recurrent = torch.stack(kv_recurrent, dim=2)
|
||||
return kv_recurrent
|
||||
|
||||
chunk_gate_recurrent = ChunkGateRecurrent.apply
|
||||
|
||||
def main():
|
||||
BSZ = 4
|
||||
NUM_HEAD = 4
|
||||
NUM_BLOCK = 16
|
||||
D_MODEL_K = 256
|
||||
D_MODEL_V = 432
|
||||
dtype = torch.float16
|
||||
kv = torch.randn(BSZ, NUM_HEAD, NUM_BLOCK, D_MODEL_K, D_MODEL_V, dtype=dtype, device="cuda")
|
||||
last_kv = torch.randn(BSZ, NUM_HEAD, D_MODEL_K, D_MODEL_V, dtype=dtype, device="cuda")
|
||||
kv_triton = kv.clone().detach()
|
||||
last_kv_triton = last_kv.clone().detach()
|
||||
cross_decay = torch.randn(BSZ, NUM_HEAD, NUM_BLOCK, dtype=dtype, device="cuda")
|
||||
cross_decay = torch.sigmoid(cross_decay)
|
||||
cross_decay_triton = cross_decay.clone().detach()
|
||||
grad_weight = torch.randn(BSZ, NUM_HEAD, NUM_BLOCK, D_MODEL_K, D_MODEL_V, dtype=dtype, device="cuda")
|
||||
kv.requires_grad = True
|
||||
kv_triton.requires_grad = True
|
||||
last_kv.requires_grad = True
|
||||
last_kv_triton.requires_grad = True
|
||||
cross_decay.requires_grad = True
|
||||
cross_decay_triton.requires_grad = True
|
||||
|
||||
start = time.time()
|
||||
kv_recurrent = naive_kv_recurrent(kv, cross_decay, last_kv)
|
||||
kv_recurrent.mul(grad_weight).sum().backward()
|
||||
print("naive time:", time.time() - start)
|
||||
|
||||
start = time.time()
|
||||
kv_recurrent_triton = chunk_gate_recurrent(kv_triton, cross_decay_triton, last_kv_triton)
|
||||
kv_recurrent_triton.mul(grad_weight).sum().backward()
|
||||
print("triton time:", time.time() - start)
|
||||
|
||||
print(torch.allclose(kv_recurrent, kv_recurrent_triton, atol=1e-3))
|
||||
print((kv_recurrent - kv_recurrent_triton).abs().max(), (kv_recurrent - kv_recurrent_triton).abs().mean())
|
||||
|
||||
print(torch.allclose(kv.grad, kv_triton.grad, atol=1e-3))
|
||||
print((kv.grad - kv_triton.grad).abs().max(), (kv.grad - kv_triton.grad).abs().mean())
|
||||
|
||||
print(torch.allclose(last_kv.grad, last_kv_triton.grad, atol=1e-3))
|
||||
print((last_kv.grad - last_kv_triton.grad).abs().max(), (last_kv.grad - last_kv_triton.grad).abs().mean())
|
||||
|
||||
print(torch.allclose(cross_decay.grad, cross_decay_triton.grad, atol=1e-3))
|
||||
print((cross_decay.grad - cross_decay_triton.grad).abs().max(), (cross_decay.grad - cross_decay_triton.grad).abs().mean())
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
main()
|
||||
@@ -0,0 +1,332 @@
|
||||
# Copyright (c) 2023, Tri Dao.
|
||||
|
||||
from typing import Optional, Union
|
||||
|
||||
import torch
|
||||
|
||||
import triton
|
||||
import triton.language as tl
|
||||
|
||||
|
||||
# @triton.autotune(
|
||||
# configs=[
|
||||
# triton.Config({"BLOCK_M": 2}),
|
||||
# triton.Config({"BLOCK_M": 4}),
|
||||
# triton.Config({"BLOCK_M": 8}),
|
||||
# triton.Config({"BLOCK_M": 16}),
|
||||
# ],
|
||||
# key=["CACHE_KEY_SEQLEN", "BLOCK_K", "INTERLEAVED"],
|
||||
# )
|
||||
@triton.jit
|
||||
def rotary_kernel(
|
||||
OUT, # Pointers to matrices
|
||||
X,
|
||||
COS,
|
||||
SIN,
|
||||
CU_SEQLENS,
|
||||
SEQLEN_OFFSETS, # this could be int or a pointer
|
||||
# Matrix dimensions
|
||||
seqlen,
|
||||
nheads,
|
||||
rotary_dim,
|
||||
seqlen_ro,
|
||||
CACHE_KEY_SEQLEN,
|
||||
# strides
|
||||
stride_out_batch,
|
||||
stride_out_seqlen,
|
||||
stride_out_nheads,
|
||||
stride_out_headdim,
|
||||
stride_x_batch,
|
||||
stride_x_seqlen,
|
||||
stride_x_nheads,
|
||||
stride_x_headdim,
|
||||
# Meta-parameters
|
||||
BLOCK_K: tl.constexpr,
|
||||
IS_SEQLEN_OFFSETS_TENSOR: tl.constexpr,
|
||||
IS_VARLEN: tl.constexpr,
|
||||
INTERLEAVED: tl.constexpr,
|
||||
CONJUGATE: tl.constexpr,
|
||||
BLOCK_M: tl.constexpr,
|
||||
):
|
||||
pid_m = tl.program_id(axis=0)
|
||||
pid_batch = tl.program_id(axis=1)
|
||||
pid_head = tl.program_id(axis=2)
|
||||
rotary_dim_half = rotary_dim // 2
|
||||
|
||||
if not IS_VARLEN:
|
||||
X = X + pid_batch * stride_x_batch + pid_head * stride_x_nheads
|
||||
OUT = OUT + pid_batch * stride_out_batch + pid_head * stride_out_nheads
|
||||
else:
|
||||
start_idx = tl.load(CU_SEQLENS + pid_batch)
|
||||
seqlen = tl.load(CU_SEQLENS + pid_batch + 1) - start_idx
|
||||
X = X + start_idx * stride_x_seqlen + pid_head * stride_x_nheads
|
||||
OUT = OUT + start_idx * stride_out_seqlen + pid_head * stride_out_nheads
|
||||
|
||||
if pid_m * BLOCK_M >= seqlen:
|
||||
return
|
||||
rm = pid_m * BLOCK_M + tl.arange(0, BLOCK_M)
|
||||
if not IS_SEQLEN_OFFSETS_TENSOR:
|
||||
rm_cs = rm + SEQLEN_OFFSETS
|
||||
else:
|
||||
rm_cs = rm + tl.load(SEQLEN_OFFSETS + pid_batch)
|
||||
rk = tl.arange(0, BLOCK_K)
|
||||
rk_half = tl.arange(0, BLOCK_K // 2)
|
||||
|
||||
if not INTERLEAVED:
|
||||
# Load the 1st and 2nd halves of X, do calculation, then store to 1st and 2nd halves of OUT
|
||||
X = X + (rm[:, None] * stride_x_seqlen + rk_half[None, :] * stride_x_headdim)
|
||||
COS = COS + (rm_cs[:, None] * rotary_dim_half + rk_half[None, :])
|
||||
SIN = SIN + (rm_cs[:, None] * rotary_dim_half + rk_half[None, :])
|
||||
cos = tl.load(
|
||||
COS, mask=(rm_cs[:, None] < seqlen_ro) & (rk_half[None, :] < rotary_dim_half), other=1.0
|
||||
).to(tl.float32)
|
||||
sin = tl.load(
|
||||
SIN, mask=(rm_cs[:, None] < seqlen_ro) & (rk_half[None, :] < rotary_dim_half), other=0.0
|
||||
).to(tl.float32)
|
||||
x0 = tl.load(
|
||||
X, mask=(rm[:, None] < seqlen) & (rk_half[None, :] < rotary_dim_half), other=0.0
|
||||
).to(tl.float32)
|
||||
x1 = tl.load(
|
||||
X + rotary_dim_half * stride_x_headdim,
|
||||
mask=(rm[:, None] < seqlen) & (rk_half[None, :] < rotary_dim_half),
|
||||
other=0.0,
|
||||
).to(tl.float32)
|
||||
if CONJUGATE:
|
||||
sin = -sin
|
||||
o0 = x0 * cos - x1 * sin
|
||||
o1 = x0 * sin + x1 * cos
|
||||
# write back result
|
||||
OUT = OUT + (rm[:, None] * stride_out_seqlen + rk_half[None, :] * stride_out_headdim)
|
||||
tl.store(OUT, o0, mask=(rm[:, None] < seqlen) & (rk_half[None, :] < rotary_dim_half))
|
||||
tl.store(
|
||||
OUT + rotary_dim_half * stride_out_headdim,
|
||||
o1,
|
||||
mask=(rm[:, None] < seqlen) & (rk_half[None, :] < rotary_dim_half),
|
||||
)
|
||||
else:
|
||||
# We don't want to load X[0, 2, 4, ...] and X[1, 3, 5, ...] separately since both are slow.
|
||||
# Instead, we load x0 = X[0, 1, 2, 3, ...] and x1 = X[1, 0, 3, 2, ...].
|
||||
# Loading x0 will be fast but x1 will be slow.
|
||||
# Then we load cos = COS[0, 0, 1, 1, ...] and sin = SIN[0, 0, 1, 1, ...].
|
||||
# Then we do the calculation and use tl.where to pick put the right outputs for the even
|
||||
# and for the odd indices.
|
||||
rk_swap = rk + ((rk + 1) % 2) * 2 - 1 # 1, 0, 3, 2, 5, 4, ...
|
||||
rk_repeat = tl.arange(0, BLOCK_K) // 2
|
||||
X0 = X + (rm[:, None] * stride_x_seqlen + rk[None, :] * stride_x_headdim)
|
||||
X1 = X + (rm[:, None] * stride_x_seqlen + rk_swap[None, :] * stride_x_headdim)
|
||||
COS = COS + (rm_cs[:, None] * rotary_dim_half + rk_repeat[None, :])
|
||||
SIN = SIN + (rm_cs[:, None] * rotary_dim_half + rk_repeat[None, :])
|
||||
cos = tl.load(
|
||||
COS,
|
||||
mask=(rm_cs[:, None] < seqlen_ro) & (rk_repeat[None, :] < rotary_dim_half),
|
||||
other=1.0,
|
||||
).to(tl.float32)
|
||||
sin = tl.load(
|
||||
SIN,
|
||||
mask=(rm_cs[:, None] < seqlen_ro) & (rk_repeat[None, :] < rotary_dim_half),
|
||||
other=0.0,
|
||||
).to(tl.float32)
|
||||
x0 = tl.load(X0, mask=(rm[:, None] < seqlen) & (rk[None, :] < rotary_dim), other=0.0).to(
|
||||
tl.float32
|
||||
)
|
||||
x1 = tl.load(
|
||||
X1, mask=(rm[:, None] < seqlen) & (rk_swap[None, :] < rotary_dim), other=0.0
|
||||
).to(tl.float32)
|
||||
if CONJUGATE:
|
||||
sin = -sin
|
||||
x0_cos = x0 * cos
|
||||
x1_sin = x1 * sin
|
||||
out = tl.where(rk[None, :] % 2 == 0, x0_cos - x1_sin, x0_cos + x1_sin)
|
||||
OUT = OUT + (rm[:, None] * stride_out_seqlen + rk[None, :] * stride_out_headdim)
|
||||
tl.store(OUT, out, mask=(rm[:, None] < seqlen) & (rk[None, :] < rotary_dim))
|
||||
|
||||
|
||||
def apply_rotary(
|
||||
x: torch.Tensor,
|
||||
cos: torch.Tensor,
|
||||
sin: torch.Tensor,
|
||||
seqlen_offsets: Union[int, torch.Tensor] = 0,
|
||||
cu_seqlens: Optional[torch.Tensor] = None,
|
||||
max_seqlen: Optional[int] = None,
|
||||
interleaved=False,
|
||||
inplace=False,
|
||||
conjugate=False,
|
||||
) -> torch.Tensor:
|
||||
"""
|
||||
Arguments:
|
||||
x: (batch, seqlen, nheads, headdim) if cu_seqlens is None
|
||||
else (total_seqlen, nheads, headdim).
|
||||
cos: (seqlen_ro, rotary_dim / 2)
|
||||
sin: (seqlen_ro, rotary_dim / 2)
|
||||
seqlen_offsets: integer or integer tensor of size (batch,)
|
||||
cu_seqlens: (batch + 1,) or None
|
||||
max_seqlen: int
|
||||
Returns:
|
||||
y: (batch, seqlen, nheads, headdim)
|
||||
"""
|
||||
is_varlen = cu_seqlens is not None
|
||||
if not is_varlen:
|
||||
batch, seqlen, nheads, headdim = x.shape
|
||||
else:
|
||||
assert max_seqlen is not None, "If cu_seqlens is passed in, then max_seqlen must be passed"
|
||||
total_seqlen, nheads, headdim = x.shape
|
||||
batch_p_1 = cu_seqlens.shape[0]
|
||||
batch = batch_p_1 - 1
|
||||
seqlen = max_seqlen
|
||||
seqlen_ro, rotary_dim = cos.shape
|
||||
assert sin.shape == cos.shape
|
||||
rotary_dim *= 2
|
||||
assert rotary_dim <= headdim, "rotary_dim must be <= headdim"
|
||||
assert headdim <= 256, "Only support headdim <= 256"
|
||||
assert seqlen_ro >= seqlen, "seqlen_ro must be >= seqlen"
|
||||
|
||||
assert (
|
||||
cos.dtype == sin.dtype
|
||||
), f"cos and sin must have the same dtype, got {cos.dtype} and {sin.dtype}"
|
||||
assert (
|
||||
x.dtype == cos.dtype
|
||||
), f"Input and cos/sin must have the same dtype, got {x.dtype} and {cos.dtype}"
|
||||
|
||||
cos, sin = cos.contiguous(), sin.contiguous()
|
||||
if isinstance(seqlen_offsets, torch.Tensor):
|
||||
assert seqlen_offsets.shape == (batch,)
|
||||
assert seqlen_offsets.dtype in [torch.int32, torch.int64]
|
||||
seqlen_offsets = seqlen_offsets.contiguous()
|
||||
else:
|
||||
assert seqlen_offsets + seqlen <= seqlen_ro
|
||||
|
||||
output = torch.empty_like(x) if not inplace else x
|
||||
if rotary_dim < headdim and not inplace:
|
||||
output[..., rotary_dim:].copy_(x[..., rotary_dim:])
|
||||
|
||||
BLOCK_K = (
|
||||
32
|
||||
if rotary_dim <= 32
|
||||
else (64 if rotary_dim <= 64 else (128 if rotary_dim <= 128 else 256))
|
||||
)
|
||||
grid = lambda META: (triton.cdiv(seqlen, META["BLOCK_M"]), batch, nheads) # noqa
|
||||
BLOCK_M = 4 if interleaved else (8 if rotary_dim <= 64 else 4)
|
||||
|
||||
# Need this, otherwise Triton tries to launch from cuda:0 and we get
|
||||
# ValueError: Pointer argument (at 0) cannot be accessed from Triton (cpu tensor?)
|
||||
with torch.cuda.device(x.device.index):
|
||||
rotary_kernel[grid](
|
||||
output, # data ptrs
|
||||
x,
|
||||
cos,
|
||||
sin,
|
||||
cu_seqlens,
|
||||
seqlen_offsets,
|
||||
seqlen, # shapes
|
||||
nheads,
|
||||
rotary_dim,
|
||||
seqlen_ro,
|
||||
seqlen // 128, # key for triton cache (limit number of compilations)
|
||||
output.stride(0) if not is_varlen else 0, # batch_strides if not varlen else 0
|
||||
output.stride(-3), # seqlen_stride or total_seqlen_stride
|
||||
output.stride(-2), # nheads_stride
|
||||
output.stride(-1), # headdim_stride
|
||||
x.stride(0) if not is_varlen else 0, # batch_strides if not varlen else 0
|
||||
x.stride(-3), # seqlen stride or total_seqlen_stride
|
||||
x.stride(-2), # nheads stride
|
||||
x.stride(-1), # headdim stride
|
||||
BLOCK_K,
|
||||
isinstance(seqlen_offsets, torch.Tensor),
|
||||
is_varlen,
|
||||
interleaved,
|
||||
conjugate,
|
||||
BLOCK_M,
|
||||
)
|
||||
return output
|
||||
|
||||
class ApplyRotaryEmb(torch.autograd.Function):
|
||||
@staticmethod
|
||||
def forward(
|
||||
ctx,
|
||||
x,
|
||||
cos,
|
||||
sin,
|
||||
interleaved=False,
|
||||
inplace=False,
|
||||
seqlen_offsets: Union[int, torch.Tensor] = 0,
|
||||
cu_seqlens: Optional[torch.Tensor] = None,
|
||||
max_seqlen: Optional[int] = None,
|
||||
):
|
||||
out = apply_rotary(
|
||||
x,
|
||||
cos,
|
||||
sin,
|
||||
seqlen_offsets=seqlen_offsets,
|
||||
cu_seqlens=cu_seqlens,
|
||||
max_seqlen=max_seqlen,
|
||||
interleaved=interleaved,
|
||||
inplace=inplace,
|
||||
)
|
||||
if isinstance(seqlen_offsets, int):
|
||||
# Can't save int with save_for_backward
|
||||
ctx.save_for_backward(cos, sin, cu_seqlens)
|
||||
ctx.seqlen_offsets = seqlen_offsets
|
||||
else:
|
||||
ctx.save_for_backward(cos, sin, cu_seqlens, seqlen_offsets)
|
||||
ctx.seqlen_offsets = None
|
||||
ctx.interleaved = interleaved
|
||||
ctx.inplace = inplace
|
||||
ctx.max_seqlen = max_seqlen
|
||||
return out if not inplace else x
|
||||
|
||||
@staticmethod
|
||||
def backward(ctx, do):
|
||||
seqlen_offsets = ctx.seqlen_offsets
|
||||
if seqlen_offsets is None:
|
||||
cos, sin, cu_seqlens, seqlen_offsets = ctx.saved_tensors
|
||||
else:
|
||||
cos, sin, cu_seqlens = ctx.saved_tensors
|
||||
# TD [2023-09-02]: For some reason Triton (2.0.0.post1) errors with
|
||||
# "[CUDA]: invalid device context", and cloning makes it work. Idk why. Triton 2.1.0 works.
|
||||
if not ctx.interleaved and not ctx.inplace:
|
||||
do = do.clone()
|
||||
dx = apply_rotary(
|
||||
do,
|
||||
cos,
|
||||
sin,
|
||||
seqlen_offsets=seqlen_offsets,
|
||||
cu_seqlens=cu_seqlens,
|
||||
max_seqlen=ctx.max_seqlen,
|
||||
interleaved=ctx.interleaved,
|
||||
inplace=ctx.inplace,
|
||||
conjugate=True,
|
||||
)
|
||||
return dx, None, None, None, None, None, None, None
|
||||
|
||||
|
||||
def apply_rotary_emb(
|
||||
x,
|
||||
cos,
|
||||
sin,
|
||||
interleaved=False,
|
||||
inplace=False,
|
||||
seqlen_offsets: Union[int, torch.Tensor] = 0,
|
||||
cu_seqlens: Optional[torch.Tensor] = None,
|
||||
max_seqlen: Optional[int] = None,
|
||||
):
|
||||
"""
|
||||
Arguments:
|
||||
x: (batch_size, seqlen, nheads, headdim) if cu_seqlens is None
|
||||
else (total_seqlen, nheads, headdim)
|
||||
cos, sin: (seqlen_rotary, rotary_dim / 2)
|
||||
interleaved: if True, rotate pairs of even and odd dimensions (GPT-J style) instead
|
||||
of 1st half and 2nd half (GPT-NeoX style).
|
||||
inplace: if True, apply rotary embedding in-place.
|
||||
seqlen_offsets: (batch_size,) or int. Each sequence in x is shifted by this amount.
|
||||
Most commonly used in inference when we have KV cache.
|
||||
cu_seqlens: (batch + 1,) or None
|
||||
max_seqlen: int
|
||||
Return:
|
||||
out: (batch_size, seqlen, nheads, headdim) if cu_seqlens is None
|
||||
else (total_seqlen, nheads, headdim)
|
||||
rotary_dim must be <= headdim
|
||||
Apply rotary embedding to the first rotary_dim of x.
|
||||
"""
|
||||
return ApplyRotaryEmb.apply(
|
||||
x, cos, sin, interleaved, inplace, seqlen_offsets, cu_seqlens, max_seqlen
|
||||
)
|
||||
@@ -0,0 +1,32 @@
|
||||
import torch
|
||||
|
||||
|
||||
swiglu_fwd_codestring = """
|
||||
template <typename T> T swiglu_fwd(T x, T y) {
|
||||
return float(x) * float(y) / (1.0f + ::exp(-float(x)));
|
||||
}
|
||||
"""
|
||||
swiglu_bwd_codestring = """
|
||||
template <typename T> T swiglu_bwd(T x, T y, T g, T& dx, T& dy) {
|
||||
float x_sigmoid = 1.0f / (1.0f + ::exp(-float(x)));
|
||||
dx = x_sigmoid * (1 + float(x) * (1.0f - x_sigmoid)) * float(g) * float(y);
|
||||
dy = float(x) * x_sigmoid * float(g);
|
||||
}
|
||||
"""
|
||||
swiglu_fwd = torch.cuda.jiterator._create_jit_fn(swiglu_fwd_codestring)
|
||||
swiglu_bwd = torch.cuda.jiterator._create_multi_output_jit_fn(swiglu_bwd_codestring, num_outputs=2)
|
||||
|
||||
|
||||
class SwiGLUFunction(torch.autograd.Function):
|
||||
|
||||
@staticmethod
|
||||
def forward(ctx, x, y):
|
||||
ctx.save_for_backward(x, y)
|
||||
return swiglu_fwd(x, y)
|
||||
|
||||
@staticmethod
|
||||
def backward(ctx, dout):
|
||||
x, y = ctx.saved_tensors
|
||||
return swiglu_bwd(x, y, dout)
|
||||
|
||||
swiglu = SwiGLUFunction.apply
|
||||
Reference in New Issue
Block a user