Function rotary_embedding_kernel

kernels/rotary_embedding.py:25–72

Apply rotary embeddings using interleaved (even/odd) decomposition. One program per row. Each row has head_dim elements. x1 = even-indexed elements (0, 2, 4, ...) x2 = odd-indexed elements (1, 3, 5, ...) out[2i] = x1[i] * cos[i] - x2[i] * sin[i] out[2i+1] = x1[i] * sin[

(
    X_ptr,
    COS_ptr,
    SIN_ptr,
    OUT_ptr,
    seq_len,
    head_dim,
    stride_x_row,
    stride_cos_row,
    stride_sin_row,
    stride_out_row,
    half_dim,
    BLOCK_SIZE: tl.constexpr,
)

@triton.jit
def rotary_embedding_kernel(
    X_ptr,
    COS_ptr,
    SIN_ptr,
    OUT_ptr,
    seq_len,
    head_dim,
    stride_x_row,
    stride_cos_row,
    stride_sin_row,
    stride_out_row,
    half_dim,
    BLOCK_SIZE: tl.constexpr,
):
    """
    Apply rotary embeddings using interleaved (even/odd) decomposition.

    One program per row. Each row has head_dim elements.
    x1 = even-indexed elements (0, 2, 4, ...)
    x2 = odd-indexed elements  (1, 3, 5, ...)

    out[2i]   = x1[i] * cos[i] - x2[i] * sin[i]
    out[2i+1] = x1[i] * sin[i] + x2[i] * cos[i]
    """
    row_idx = tl.program_id(0)

    col_offsets = tl.arange(0, BLOCK_SIZE)
    mask_half = col_offsets < half_dim

    # Compute pointers to even-indexed (x1) and odd-indexed (x2) elements
    # even indices: 0, 2, 4, ... => col_offsets * 2
    # odd indices:  1, 3, 5, ... => col_offsets * 2 + 1
    x_row_base = X_ptr + row_idx * stride_x_row
    x1 = tl.load(x_row_base + col_offsets * 2, mask=mask_half, other=0.0).to(tl.float32)
    x2 = tl.load(x_row_base + col_offsets * 2 + 1, mask=mask_half, other=0.0).to(tl.float32)

    # Load cos and sin (shape [n_rows, half_dim])
    cos = tl.load(COS_ptr + row_idx * stride_cos_row + col_offsets, mask=mask_half, other=1.0).to(tl.float32)
    sin = tl.load(SIN_ptr + row_idx * stride_sin_row + col_offsets, mask=mask_half, other=0.0).to(tl.float32)

    # Apply rotation
    rx1 = x1 * cos - x2 * sin
    rx2 = x1 * sin + x2 * cos

    # Store results interleaved: out[2i] = rx1[i], out[2i+1] = rx2[i]
    out_row_base = OUT_ptr + row_idx * stride_out_row
    tl.store(out_row_base + col_offsets * 2, rx1, mask=mask_half)
    tl.store(out_row_base + col_offsets * 2 + 1, rx2, mask=mask_half)


def kernel_fn(