rope_strided

Function rope_strided 

Source
pub fn rope_strided(
    gpu: &dyn GpuBackend,
    kernel: KernelHandle,
    q: DevicePtr,
    k: DevicePtr,
    positions: DevicePtr,
    num_tokens: u32,
    num_q_heads: u32,
    num_kv_heads: u32,
    head_dim: u32,
    rotary_dim: u32,
    theta: f32,
    q_row_stride: u32,
    k_row_stride: u32,
    stream: u64,
) -> Result<()>
Expand description

RoPE: apply rotary position embeddings to Q and K in-place.

Kernel: rope_forward(Q, K, positions, seq_len, num_q_heads, num_kv_heads, head_dim, rotary_dim, theta) Grid: (num_q_heads + num_kv_heads, ceil(seq_len/4), 1) Block: (128, 1, 1)

positions must be a device pointer to a u32[seq_len] array. Strided RoPE: rotates ALL num_tokens rows in ONE launch.

rope above derives each row’s address from a PACKED layout (num_*_heads * head_dim between tokens). The multi-seq decode buffer is not packed — Q and K live inside one interleaved [Q|K|V|gate] block whose rows sit per_seq_qkv apart — so that path was calling rope once per sequence with seq_len = 1: 258 launches/step at 4.6 us = 1.18 ms across the 16 attention layers.

Bit-identical to n packed launches: same math, same ordering, only the row address differs. Passing the packed strides reproduces rope exactly.