Skip to main content

Module rope

Module rope 

Source
Expand description

Rotary Position Embedding (RoPE) GPU dispatch.

Applies rotation in pairs of elements using cos/sin of position * theta^(-2i/d). Gemma 4 uses theta=10000 for sliding attention layers and theta=1000000 for global attention layers.

Statics§

ROPE_SHADER_SOURCE
MSL source for the RoPE kernels (embedded at compile time).

Functions§

dispatch_rope
Dispatch a RoPE operation on the GPU.
dispatch_rope_neox_bf16
Dispatch a Neox/split-convention RoPE operation on the GPU (bf16 only).
dispatch_rope_neox_f32
Dispatch a Neox/split-convention RoPE operation on the GPU (f32) with optional freq_factors support.
register
Register RoPE shader sources with the given kernel registry.