Spaces:
Running on Zero
Running on Zero
File size: 4,363 Bytes
7755bd1 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 | import torch
import triton
import triton.language as tl
@triton.jit
def _causal_rope_apply_kernel(
out_ptr,
x_ptr,
freqs_real_ptr,
grid_sizes_ptr,
num_heads: tl.constexpr,
head_dim: tl.constexpr,
seq_stride: tl.constexpr,
t_dim: tl.constexpr,
h_dim: tl.constexpr,
start_frame: tl.constexpr,
conjugate: tl.constexpr,
BLOCK_SIZE: tl.constexpr,
):
row_idx = tl.program_id(0)
batch_idx = row_idx // (seq_stride * num_heads)
seq_head_idx = row_idx - batch_idx * seq_stride * num_heads
token_idx = seq_head_idx // num_heads
f = tl.load(grid_sizes_ptr + batch_idx * 3)
h = tl.load(grid_sizes_ptr + batch_idx * 3 + 1)
w = tl.load(grid_sizes_ptr + batch_idx * 3 + 2)
valid_seq_len = f * h * w
half_head_dim: tl.constexpr = head_dim // 2
col_offsets = tl.arange(0, BLOCK_SIZE)
mask = col_offsets < half_head_dim
x_row_ptr = x_ptr + row_idx * head_dim
out_row_ptr = out_ptr + row_idx * head_dim
x_real = tl.load(x_row_ptr + col_offsets * 2, mask=mask, other=0.0)
x_imag = tl.load(x_row_ptr + col_offsets * 2 + 1, mask=mask, other=0.0)
is_valid = token_idx < valid_seq_len
hw = h * w
t = token_idx // hw
rem = token_idx - t * hw
y = rem // w
z = rem - y * w
freq_row = tl.where(col_offsets < t_dim, start_frame + t, tl.where(col_offsets < t_dim + h_dim, y, z))
freq_offset = (freq_row * half_head_dim + col_offsets) * 2
load_mask = mask & is_valid
cos_vals = tl.load(freqs_real_ptr + freq_offset, mask=load_mask, other=1.0)
sin_vals = tl.load(freqs_real_ptr + freq_offset + 1, mask=load_mask, other=0.0)
if conjugate:
sin_vals = -sin_vals
rotated_real = x_real.to(tl.float32) * cos_vals.to(tl.float32) - x_imag.to(tl.float32) * sin_vals.to(tl.float32)
rotated_imag = x_real.to(tl.float32) * sin_vals.to(tl.float32) + x_imag.to(tl.float32) * cos_vals.to(tl.float32)
out_real = tl.where(is_valid, rotated_real, x_real)
out_imag = tl.where(is_valid, rotated_imag, x_imag)
tl.store(out_row_ptr + col_offsets * 2, out_real, mask=mask)
tl.store(out_row_ptr + col_offsets * 2 + 1, out_imag, mask=mask)
class _CausalRoPEApplyTriton(torch.autograd.Function):
@staticmethod
def forward(ctx, x, grid_sizes, freqs, start_frame):
b, seq_len, num_heads, head_dim = x.shape
c = head_dim // 2
t_dim = c - 2 * (c // 3)
h_dim = c // 3
x_contig = x.contiguous()
out = torch.empty_like(x_contig)
grid_sizes = grid_sizes.to(device=x.device, dtype=torch.int64, non_blocking=True).contiguous()
freqs_real = torch.view_as_real(freqs.to(device=x.device))
block_size = triton.next_power_of_2(c)
_causal_rope_apply_kernel[(b * seq_len * num_heads,)](
out,
x_contig,
freqs_real,
grid_sizes,
num_heads,
head_dim,
seq_len,
t_dim,
h_dim,
start_frame,
False,
BLOCK_SIZE=block_size,
num_warps=1,
num_stages=2,
)
ctx.save_for_backward(grid_sizes, freqs_real)
ctx.num_heads = num_heads
ctx.head_dim = head_dim
ctx.seq_len = seq_len
ctx.t_dim = t_dim
ctx.h_dim = h_dim
ctx.start_frame = start_frame
ctx.block_size = block_size
ctx.input_shape = x.shape
return out
@staticmethod
def backward(ctx, grad_output):
grid_sizes, freqs_real = ctx.saved_tensors
grad_contig = grad_output.contiguous()
grad_x = torch.empty_like(grad_contig)
b, seq_len, num_heads, _ = grad_contig.shape
_causal_rope_apply_kernel[(b * seq_len * num_heads,)](
grad_x,
grad_contig,
freqs_real,
grid_sizes,
ctx.num_heads,
ctx.head_dim,
ctx.seq_len,
ctx.t_dim,
ctx.h_dim,
ctx.start_frame,
True,
BLOCK_SIZE=ctx.block_size,
num_warps=1,
num_stages=2,
)
return grad_x.reshape(ctx.input_shape), None, None, None
def causal_rope_apply_triton(x, grid_sizes, freqs, start_frame=0):
return _CausalRoPEApplyTriton.apply(x, grid_sizes, freqs, start_frame).type_as(x)
|