Как реализовать и протестировать встраивания RoPE в Triton? ⇐ Python

[Anonymous]

Из статьи RoFormer были представлены вложения RoPE.
https://arxiv.org/abs/2104.09864


И формулировка:


Если бы мы написали код в Triton, правильная ли это реализация для встраивания RoPe?

Код: Выделить всё

import torch
import torch.nn as nn
import triton
import triton.language as tl

# 1. Rotary Positional Embeddings (RoPE)
@triton.jit
def rope_kernel(x_ptr, max_len, dim, stride, **meta):
batch_size = tl.program_id(0)
seq_len = tl.program_id(1)

# Load the input tensor
x = tl.load(x_ptr + (batch_size * stride) + (seq_len * dim))

# Precompute frequency terms
freq_indices = tl.arange(0, dim // 2)
freqs = tl.arange(0, max_len)[:, None] * (1.0 / (10000 ** (freq_indices / dim)))

# Compute sine and cosine for RoPE
sin_freqs = tl.sin(freqs)
cos_freqs = tl.cos(freqs)

# Efficiently apply RoPE to the input
for i in range(0, dim, 2):
x_i = x[i:i+2]  # Load two elements at once
x[i] = x_i[0] * cos_freqs[seq_len] - x_i[1] * sin_freqs[seq_len]
x[i+1] = x_i[0] * sin_freqs[seq_len] + x_i[1] * cos_freqs[seq_len]

# Store the output
tl.store(x_ptr + (batch_size * stride) + (seq_len * dim), x)

Как проверить, работают ли ядра тритона правильно, как это определено формулировкой встраивания веревок?


Подробнее здесь: https://stackoverflow.com/questions/791 ... -in-triton