Jump to content

Connect SuperML | Leeroopedia MCP: Equip your AI agents with best practices, code verification, and debugging knowledge. Powered by Leeroo — building Organizational Superintelligence. Contact us at founders@leeroo.com.

Implementation:Turboderp org Exllamav2 Ext RoPE

From Leeroopedia
Revision as of 14:01, 16 February 2026 by Admin (talk | contribs) (Auto-imported from implementations/Turboderp_org_Exllamav2_Ext_RoPE.md)
(diff) ← Older revision | Latest revision (diff) | Newer revision → (diff)
Knowledge Sources
Domains Positional_Encoding, CUDA, C_Extension
Last Updated 2026-02-15 00:00 GMT

Overview

C++ extension implementing in-place Rotary Position Embedding (RoPE) application on query/key tensors and multi-dimensional position ID generation for vision-language models (mRoPE).

Description

ext_rope.cpp provides two functions for positional encoding:

rope_(x, sin, cos, past_len, num_heads, head_dim, offsets, neox_style) applies RoPE in-place to a tensor of query or key states. The function:

  • Validates that x, sin, and cos tensors are FP16 and that sin/cos tables have matching sizes.
  • Computes rows_per_batch from the tensor dimensions to handle multi-head layouts.
  • Supports optional integer offsets per batch element (for variable-length sequences in batched inference); when the offsets tensor is on a meta device, it is treated as NULL.
  • The neox_style flag selects between the original GPT-NeoX rotation pattern (interleaved real/imaginary) and the standard sequential pattern.
  • Dispatches to rope_cuda which applies the rotation on the GPU stream.

gen_mrope_pos_ids(mrope_pos_ids, ids, merge_size, spans, grids) generates multi-dimensional position IDs for Multimodal RoPE (mRoPE), used in vision-language models like Qwen2-VL. The function:

  • Maintains three output channels: temporal (t), height (h), and width (w) position IDs.
  • Iterates over input token IDs, checking if each falls within a vision embedding span.
  • For vision tokens, computes 3D grid coordinates (t, h, w) based on the image/video grid dimensions and merge_size (spatial token merging factor).
  • For text tokens, increments a monotonic base counter shared across all three channels.
  • Returns the next base temporal position for continuation.

Usage

rope_ is called during the attention forward pass (phase 1) to apply positional information to query and key tensors after projection. It is invoked by both the standard q_attn_forward_1 path and by Python-side attention implementations.

gen_mrope_pos_ids is called during input preparation for multimodal models to generate the 3-channel position ID tensor that controls how RoPE frequencies are assigned to text vs. image/video tokens.

Code Reference

Source Location

Signature

void rope_(
    torch::Tensor x,
    torch::Tensor sin,
    torch::Tensor cos,
    int past_len,
    int num_heads,
    int head_dim,
    torch::Tensor offsets,
    bool neox_style
);

int64_t gen_mrope_pos_ids(
    torch::Tensor mrope_pos_ids,
    torch::Tensor ids,
    int merge_size,
    const std::vector<std::tuple<int64_t, int64_t>>& spans,
    const std::vector<std::tuple<int64_t, int64_t, int64_t>>& grids
);

Import

from exllamav2 import exllamav2_ext as ext_c

# Apply RoPE in-place
ext_c.rope_(x, sin, cos, past_len, num_heads, head_dim, offsets, neox_style)

# Generate mRoPE position IDs
next_base = ext_c.gen_mrope_pos_ids(mrope_pos_ids, ids, merge_size, spans, grids)

I/O Contract

Function Parameter Type Direction Description
rope_ x Tensor (FP16) in/out Query or key tensor, modified in-place with rotary embeddings
rope_ sin Tensor (FP16) in Precomputed sine table for rotary frequencies
rope_ cos Tensor (FP16) in Precomputed cosine table for rotary frequencies
rope_ past_len int in Number of past tokens (for position offset in KV cache)
rope_ num_heads int in Number of attention heads
rope_ head_dim int in Dimension per attention head
rope_ offsets Tensor (int32) or meta in Per-batch position offsets; meta tensor = no offsets
rope_ neox_style bool in True for GPT-NeoX interleaved rotation; false for sequential
gen_mrope_pos_ids mrope_pos_ids Tensor (int64) out Output shape (3, max_length): t/h/w position channels
gen_mrope_pos_ids ids Tensor (int64) in Input token IDs, shape (in_length,)
gen_mrope_pos_ids merge_size int in Spatial merge factor for vision tokens
gen_mrope_pos_ids spans vector<tuple<int64, int64>> in Vision embedding span ranges (start, end)
gen_mrope_pos_ids grids vector<tuple<int64, int64, int64>> in Per-span grid dimensions (t, h, w)
Function Return Type Description
rope_ -- void Tensor modified in-place
gen_mrope_pos_ids next_base_t int64_t Next available temporal base position for continuation

Usage Examples

import torch
from exllamav2 import exllamav2_ext as ext_c

# Apply RoPE to query tensor (standard sequential style)
q = torch.randn(1, 32, 128, dtype=torch.float16, device="cuda")  # (batch, heads*seq, head_dim)
sin = precomputed_sin  # shape: (max_seq, head_dim)
cos = precomputed_cos  # shape: (max_seq, head_dim)
offsets = torch.empty(0, device="meta")  # no per-batch offsets
ext_c.rope_(q, sin, cos, past_len=0, num_heads=32, head_dim=128, offsets=offsets, neox_style=False)

# Generate mRoPE position IDs for a vision-language input
mrope_pos_ids = torch.zeros(3, 2048, dtype=torch.int64)
ids = input_token_ids  # int64 tensor
spans = [(1024, 1280)]  # vision token span
grids = [(1, 14, 14)]   # 1 frame, 14x14 spatial grid
next_pos = ext_c.gen_mrope_pos_ids(mrope_pos_ids, ids, merge_size=2, spans=spans, grids=grids)

Related Pages

Page Connections

Double-click a node to navigate. Hold to expand connections.
Principle
Implementation
Heuristic
Environment