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.