Skip to content

vllm.models.kimi_k3.nvidia.index_rope

Triton RoPE for the Kimi K3.1 DSA indexer.

Applies rope to the leading rope_dim dims of each indexer head, using the leading rope_dim // 2 frequencies sliced from the main attention's cos_sin_cache (training reuses the leading freqs of the main freqs_cis for the indexer). Compute is fp32 with a cast back to the input dtype. dim >= rope_dim is left untouched (in-place).

Two pairing layouts (same frequencies, different dim pairing): * neox half-split (interleave=False, default): dim i pairs with dim i+half; * GPT-J interleave (interleave=True, GLM native weight layout): dim 2i pairs with dim 2i+1, cos/sin broadcast per pair.

Functions:

  • index_rope –

    In-place RoPE on the leading rope_dim dims of each indexer head.

index_rope(index_q, index_k, positions, cos_sin_cache, rope_dim, interleave)

In-place RoPE on the leading rope_dim dims of each indexer head.

fp32 compute, neox half-split / GPT-J interleaved pairing; dims >= rope_dim are untouched, q/k done in a single launch. Replaces the former handwritten torch composition (~20 small kernels in eager).

Source code in vllm/models/kimi_k3/nvidia/index_rope.py 
 82
 83
 84
 85
 86
 87
 88
 89
 90
 91
 92
 93
 94
 95
 96
 97
 98
 99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
def index_rope(
    index_q: torch.Tensor,  # [T, H, D]
    index_k: torch.Tensor,  # [T, D]
    positions: torch.Tensor,  # [T] int64
    cos_sin_cache: torch.Tensor,  # [max_pos, rotary_dim]
    rope_dim: int,
    interleave: bool,
) -> None:
    """In-place RoPE on the leading rope_dim dims of each indexer head.

    fp32 compute, neox half-split / GPT-J interleaved pairing; dims >=
    rope_dim are untouched, q/k done in a single launch. Replaces the former
    handwritten torch composition (~20 small kernels in eager).
    """
    T, H, D = index_q.shape
    assert index_q.stride(2) == 1 and index_k.stride(1) == 1
    assert index_k.shape[0] == T and index_k.shape[1] == D
    assert cos_sin_cache.stride(1) == 1
    half = rope_dim // 2
    cache_half = cos_sin_cache.shape[1] // 2
    index_rope_kernel[(T * (H + 1),)](
        index_q,
        index_k,
        positions,
        cos_sin_cache,
        H,
        index_q.stride(0),
        index_q.stride(1),
        index_k.stride(0),
        positions.stride(0),
        cos_sin_cache.stride(0),
        HALF=half,
        CACHE_HALF=cache_half,
        BLOCK_HALF=triton.next_power_of_2(half),
        INTERLEAVE=interleave,
        num_warps=1 if half <= 32 else 2,
    )