Skip to content

vllm.model_executor.layers.rotary_embedding.mrope_vit_setup

Functions:

  • vit_mrope_setup –

    (3, total_N, half_rot_dim) fp32 cos/sin t/h/w planes for triton_mrope.

_vit_mrope_setup_kernel(inv_freq_t_ptr, inv_freq_h_ptr, inv_freq_w_ptr, cos_out_ptr, sin_out_ptr, cu_seqlens_ptr, grids_ptr, num_tokens, num_grids, half_t, half_h, half_w, spatial_merge, BLOCK_N, BLOCK_D)

Build triton_mrope's (3, N, half_rot) cos/sin t/h/w planes in one launch: compute each token's (t, h, w) position arithmetically and the rotation in fp32 from the per-axis inv_freq. Only each axis's mrope section columns are written; the rest is left uninitialized (the consumer's per-axis masked loads read only the section columns).

Source code in vllm/model_executor/layers/rotary_embedding/mrope_vit_setup.py 
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
@triton.jit(
    do_not_specialize=["num_tokens", "num_grids"],
    do_not_specialize_on_alignment=["grids_ptr"],
)
def _vit_mrope_setup_kernel(
    inv_freq_t_ptr,
    inv_freq_h_ptr,
    inv_freq_w_ptr,
    cos_out_ptr,
    sin_out_ptr,
    cu_seqlens_ptr,  # (G + 1,) int32 token boundaries per grid
    grids_ptr,  # (G, 3) int32: t, h, w per grid
    num_tokens,
    num_grids,
    half_t: tl.constexpr,
    half_h: tl.constexpr,
    half_w: tl.constexpr,
    spatial_merge: tl.constexpr,
    BLOCK_N: tl.constexpr,
    BLOCK_D: tl.constexpr,
):
    """Build triton_mrope's (3, N, half_rot) cos/sin t/h/w planes in one
    launch: compute each token's (t, h, w) position arithmetically and the
    rotation in fp32 from the per-axis inv_freq. Only each axis's mrope
    section columns are written; the rest is left uninitialized (the
    consumer's per-axis masked loads read only the section columns)."""
    pid = tl.program_id(0)
    tok = pid * BLOCK_N + tl.arange(0, BLOCK_N)
    tok_mask = tok < num_tokens

    # Grid (segment) index per token; cu_seqlens is nondecreasing.
    seg = tl.zeros([BLOCK_N], dtype=tl.int32)
    for g in range(num_grids):
        end = tl.load(cu_seqlens_ptr + g + 1)
        seg += (tok >= end).to(tl.int32)

    seg_start = tl.load(cu_seqlens_ptr + seg, mask=tok_mask, other=0)
    grid_h = tl.load(grids_ptr + seg * 3 + 1, mask=tok_mask, other=1)
    grid_w = tl.load(grids_ptr + seg * 3 + 2, mask=tok_mask, other=1)

    local = tok - seg_start
    rem = local % (grid_h * grid_w)
    tpos = local // (grid_h * grid_w)
    # Within a frame tokens are ordered by spatial-merge block (bh, bw, mh,
    # mw), matching the reference ViT position-id permute.
    bh = rem // (spatial_merge * grid_w)
    r = rem % (spatial_merge * grid_w)
    bw = r // (spatial_merge * spatial_merge)
    r2 = r % (spatial_merge * spatial_merge)
    hpos = bh * spatial_merge + r2 // spatial_merge
    wpos = bw * spatial_merge + r2 % spatial_merge

    d = tl.arange(0, BLOCK_D)
    half_rot = half_t + half_h + half_w
    row = (tok * half_rot)[:, None] + d[None, :]
    plane = num_tokens * half_rot

    # Per-axis compute+store to cap live registers; the fp32 trig sequence
    # wants few elements per thread (hence small BLOCK_N, many warps).
    in_t = d < half_t
    inv_t = tl.load(inv_freq_t_ptr + d, mask=in_t, other=0.0)
    freq = tpos.to(tl.float32)[:, None] * inv_t[None, :]
    mask = tok_mask[:, None] & in_t[None, :]
    tl.store(cos_out_ptr + row, tl.cos(freq), mask=mask)
    tl.store(sin_out_ptr + row, tl.sin(freq), mask=mask)

    in_h = (d >= half_t) & (d < half_t + half_h)
    inv_h = tl.load(inv_freq_h_ptr + d - half_t, mask=in_h, other=0.0)
    freq = hpos.to(tl.float32)[:, None] * inv_h[None, :]
    mask = tok_mask[:, None] & in_h[None, :]
    tl.store(cos_out_ptr + plane + row, tl.cos(freq), mask=mask)
    tl.store(sin_out_ptr + plane + row, tl.sin(freq), mask=mask)

    in_w = (d >= half_t + half_h) & (d < half_rot)
    inv_w = tl.load(inv_freq_w_ptr + d - half_t - half_h, mask=in_w, other=0.0)
    freq = wpos.to(tl.float32)[:, None] * inv_w[None, :]
    mask = tok_mask[:, None] & in_w[None, :]
    tl.store(cos_out_ptr + 2 * plane + row, tl.cos(freq), mask=mask)
    tl.store(sin_out_ptr + 2 * plane + row, tl.sin(freq), mask=mask)

vit_mrope_setup(inv_freq_t, inv_freq_h, inv_freq_w, grid_thw, spatial_merge_size)

(3, total_N, half_rot_dim) fp32 cos/sin t/h/w planes for triton_mrope.

One fused kernel: positions are computed arithmetically from the grids and the rotation comes from the per-axis inv_freq, so there is no cos/sin cache to size or grow.

Parameters:

  • inv_freq_t

    (Tensor) –

    t-axis inverse frequencies, (t_dim // 2,).

  • inv_freq_h

    (Tensor) –

    h-axis inverse frequencies, (h_dim // 2,).

  • inv_freq_w

    (Tensor) –

    w-axis inverse frequencies, (w_dim // 2,).

  • grid_thw

    (list[list[int]]) –

    per-grid [t, h, w] patch counts (host metadata).

  • spatial_merge_size

    (int) –

    spatial merge factor of the ViT.

Source code in vllm/model_executor/layers/rotary_embedding/mrope_vit_setup.py 
 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
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
def vit_mrope_setup(
    inv_freq_t: torch.Tensor,
    inv_freq_h: torch.Tensor,
    inv_freq_w: torch.Tensor,
    grid_thw: list[list[int]],
    spatial_merge_size: int,
) -> tuple[torch.Tensor, torch.Tensor]:
    """(3, total_N, half_rot_dim) fp32 cos/sin t/h/w planes for triton_mrope.

    One fused kernel: positions are computed arithmetically from the grids
    and the rotation comes from the per-axis inv_freq, so there is no
    cos/sin cache to size or grow.

    Args:
        inv_freq_t: t-axis inverse frequencies, (t_dim // 2,).
        inv_freq_h: h-axis inverse frequencies, (h_dim // 2,).
        inv_freq_w: w-axis inverse frequencies, (w_dim // 2,).
        grid_thw: per-grid [t, h, w] patch counts (host metadata).
        spatial_merge_size: spatial merge factor of the ViT.

    """
    device = inv_freq_t.device
    cu = [0]
    for t, h, w in grid_thw:
        cu.append(cu[-1] + t * h * w)
    n = cu[-1]
    g = len(grid_thw)
    # grid_thw is per-request host metadata, so one H2D is unavoidable;
    # pack cu_seqlens and grids into a single pinned async copy.
    buf = async_tensor_h2d(
        cu + [d for grid in grid_thw for d in grid], device=device, dtype=torch.int32
    )
    cu_seqlens = buf[: g + 1]
    grids = buf[g + 1 :].view(g, 3)

    half_t = inv_freq_t.numel()
    half_h = inv_freq_h.numel()
    half_w = inv_freq_w.numel()
    half_rot = half_t + half_h + half_w
    cos = torch.empty(3, n, half_rot, device=device, dtype=torch.float32)
    sin = torch.empty(3, n, half_rot, device=device, dtype=torch.float32)
    _vit_mrope_setup_kernel[(triton.cdiv(n, 32),)](
        inv_freq_t,
        inv_freq_h,
        inv_freq_w,
        cos,
        sin,
        cu_seqlens,
        grids,
        n,
        g,
        half_t=half_t,
        half_h=half_h,
        half_w=half_w,
        spatial_merge=spatial_merge_size,
        BLOCK_N=32,
        BLOCK_D=triton.next_power_of_2(half_rot),
        num_warps=8,
    )
    return cos, sin