RoPE Operators¶

本页暂无中文版。以下为英文原文。

Every op on this page is used the same way: construct it once, then call it. The constructor takes what the kernel is compiled with; the call takes the tensors. Both are documented under each op — __init__ and forward, where forward is what runs when you call op(...).

Base frequencies¶

One op serves both rotation conventions: pass rope_layout as "neox" to rotate the two halves of a head against each other, or "interleaved" to rotate each adjacent pair.

tileops.rope.RoPEFwdOp ¶

RoPE with standard theta frequencies, in either rotation convention.

rope_layout="neox" splits a head at its midpoint and rotates the halves against each other; "interleaved" rotates each adjacent pair. The frequency table is the same either way, so the two differ only in which elements pair up.

References: GPT-NeoX / HuggingFace transformers RotaryEmbedding (NeoX); Su et al., "RoFormer: Enhanced Transformer with Rotary Position Embedding" (interleaved).

init ¶

__init__(
    rope_layout="neox",
    input_layout="1d",
    base=10000.0,
    *,
    target=None
)

Build the op. Shapes and dtype are taken from the first call.

Parameters:

rope_layout (str, default: 'neox' ) –

"neox" to rotate the two halves of a head against each other, "interleaved" to rotate each adjacent pair.
input_layout (str, default: '1d' ) –

"1d" for [seq_len, head_dim] or "2d" for [batch, seq_len, num_heads, head_dim].
base (float, default: 10000.0 ) –

Frequency base (default 10000).
target (Target, default: None ) –

Which set of kernels serves this op — a target name, BUILTIN for the in-tree kernels, or None to decide from the input device.

forward ¶

forward(
    x,
)

Apply RoPE rotation using internally computed cos/sin tables.

Parameters:

x (Tensor) –

Input tensor. Shape depends on the input layout: - 1D: (seq_len, head_dim) - 2D: (batch, seq_len, num_heads, head_dim)

Returns:

Tensor –

Rotated output tensor with same shape as x.

tileops.rope.RoPENeoxPositionIdsFwdOp ¶

GPT-NeoX style RoPE for packed THD tensors with explicit positions.

The first rotary_dim columns of each head rotate (all of them when rotary_dim is None) and the rest are copied.

init ¶

__init__(
    max_position,
    base=10000.0,
    rotary_dim=None,
    *,
    target=None
)

Build the op. Shapes and dtype are taken from the first call.

Parameters:

max_position (int) –

Manifest params.max_position, int.
base (float, default: 10000.0 ) –

Manifest params.base, float, default 10000.0.
rotary_dim (Optional[int], default: None ) –

Manifest params.rotary_dim, int | None, default None.
target (Target, default: None ) –

Which set of kernels serves this op — a target name, BUILTIN for the in-tree kernels, or None to decide from the input device.

forward ¶

forward(
    x,
    position_ids,
)

Run the op on the inputs the manifest declares.

Parameters:

x (Tensor) –

[tokens, heads, head_dim], dtype float16 | bfloat16 | float32.
position_ids (Tensor) –

[tokens], dtype int32 | int64, each in [0, max_position).

Returns:

Tensor –

output, shaped as x.

Scaled frequencies¶

tileops.rope.RoPELlama31FwdOp ¶

Llama 3.1 RoPE op with piecewise frequency scaling.

Computes cos/sin tables at construction using Llama 3.1 piecewise-scaled frequencies based on wavelength thresholds.

Reference: Meta Llama 3.1 model implementation.

init ¶

__init__(
    input_layout="1d",
    base=10000.0,
    scale_factor=8.0,
    low_freq_factor=1.0,
    high_freq_factor=4.0,
    original_max_position=8192,
    *,
    target=None
)

Build the op. Shapes and dtype are taken from the first call.

Parameters:

input_layout (str, default: '1d' ) –

"1d" or "2d".
base (float, default: 10000.0 ) –

Frequency base (default 10000).
scale_factor (float, default: 8.0 ) –

Scaling factor for low frequencies (default 8.0).
low_freq_factor (float, default: 1.0 ) –

Low-frequency wavelen threshold (default 1.0).
high_freq_factor (float, default: 4.0 ) –

High-frequency wavelen threshold (default 4.0).
original_max_position (int, default: 8192 ) –

Original max position (default 8192).
target (Target, default: None ) –

Which set of kernels serves this op — a target name, BUILTIN for the in-tree kernels, or None to decide from the input device.

forward ¶

forward(
    x,
)

Apply RoPE rotation using internally computed cos/sin tables.

Parameters:

x (Tensor) –

Input tensor. Shape depends on the input layout: - 1D: (seq_len, head_dim) - 2D: (batch, seq_len, num_heads, head_dim)

Returns:

Tensor –

Rotated output tensor with same shape as x.

tileops.rope.YaRNFwdOp ¶

YaRN RoPE op with linear-ramp frequency interpolation.

Computes cos/sin tables at construction using YaRN linear-ramp interpolation between scaled and original frequencies.

Reference: Peng et al., "YaRN: Efficient Context Window Extension of LLMs".

init ¶

__init__(
    input_layout="1d",
    base=10000.0,
    scale=16.0,
    original_max_position=4096,
    beta_fast=32.0,
    beta_slow=1.0,
    attn_factor=1.0,
    *,
    target=None
)

Build the op. Shapes and dtype are taken from the first call.

Parameters:

input_layout (str, default: '1d' ) –

"1d" or "2d".
base (float, default: 10000.0 ) –

Frequency base (default 10000).
scale (float, default: 16.0 ) –

Context extension scale (default 16.0).
original_max_position (int, default: 4096 ) –

Original max position (default 4096).
beta_fast (float, default: 32.0 ) –

Fast decay boundary (default 32.0).
beta_slow (float, default: 1.0 ) –

Slow decay boundary (default 1.0).
attn_factor (float, default: 1.0 ) –

Attention scaling factor (default 1.0).
target (Target, default: None ) –

Which set of kernels serves this op — a target name, BUILTIN for the in-tree kernels, or None to decide from the input device.

forward ¶

forward(
    x,
)

Apply RoPE rotation using internally computed cos/sin tables.

Parameters:

x (Tensor) –

Input tensor. Shape depends on the input layout: - 1D: (seq_len, head_dim) - 2D: (batch, seq_len, num_heads, head_dim)

Returns:

Tensor –

Rotated output tensor with same shape as x.

tileops.rope.LongRoPEFwdOp ¶

LongRoPE op with per-dimension frequency rescaling.

Computes cos/sin tables at construction using per-dimension rescale factors (ext_factors) that multiply the divisor, plus a scale-dependent amplitude factor applied to cos/sin output. Construction rejects a zero or NaN rescale factor, since the frequency divides by it.

Reference: TVM rope_freq_longrope in position_embedding.py; Ding et al., "LongRoPE: Extending LLM Context Window Beyond 2M Tokens".

init ¶

__init__(
    input_layout="1d",
    base=10000.0,
    rescale_factors=None,
    max_position_embeddings=4096,
    original_max_position_embeddings=4096,
    *,
    target=None
)

Build the op. Shapes and dtype are taken from the first call.

Parameters:

input_layout (str, default: '1d' ) –

"1d" or "2d".
base (float, default: 10000.0 ) –

Frequency base (default 10000).
rescale_factors (Optional[Tensor], default: None ) –

Per-dimension rescale factors (ext_factors) of shape (head_dim // 2,). These multiply the divisor.
max_position_embeddings (int, default: 4096 ) –

Extended max position length (default 4096).
original_max_position_embeddings (int, default: 4096 ) –

Original max position length (default 4096).
target (Target, default: None ) –

Which set of kernels serves this op — a target name, BUILTIN for the in-tree kernels, or None to decide from the input device.

Raises:

ValueError –

A rescale factor is zero or NaN; the frequency divides by it.

forward ¶

forward(
    x,
)

Apply RoPE rotation using internally computed cos/sin tables.

Parameters:

x (Tensor) –

Input tensor. Shape depends on the input layout: - 1D: (seq_len, head_dim) - 2D: (batch, seq_len, num_heads, head_dim)

Returns:

Tensor –

Rotated output tensor with same shape as x.

RoPE Operators¶

Base frequencies¶

tileops.rope.RoPEFwdOp ¶

__init__ ¶

forward ¶

tileops.rope.RoPENeoxPositionIdsFwdOp ¶

__init__ ¶

forward ¶

Scaled frequencies¶

tileops.rope.RoPELlama31FwdOp ¶

__init__ ¶

forward ¶

tileops.rope.YaRNFwdOp ¶

__init__ ¶

forward ¶

tileops.rope.LongRoPEFwdOp ¶

__init__ ¶

forward ¶

init ¶

init ¶

init ¶

init ¶

init ¶