Skip to content

vllm.model_executor.layers.rotary_embedding.mrope_vit_setup

Classes:

  • VitMropeSetup –

    Builds triton_mrope's (3, N, half_rot) cos/sin t/h/w planes.

Functions:

VitMropeSetup

Bases: PluggableLayer

Builds triton_mrope's (3, N, half_rot) cos/sin t/h/w planes.

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.

Stateless: the caller's inv_freq buffers are passed per call. OOT platforms replace the whole layer via @VitMropeSetup.register_oot with a subclass overriding :meth:forward (e.g. to retune the kernel tiling for their hardware).

Methods:

  • forward –

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

Source code in vllm/model_executor/layers/rotary_embedding/mrope_vit_setup.py 
 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
154
155
156
157
158
159
160
161
162
163
164
165
166
167
@PluggableLayer.register("vit_mrope_setup")
class VitMropeSetup(PluggableLayer):
    """Builds triton_mrope's (3, N, half_rot) cos/sin t/h/w planes.

    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.

    Stateless: the caller's inv_freq buffers are passed per call. OOT
    platforms replace the whole layer via ``@VitMropeSetup.register_oot``
    with a subclass overriding :meth:`forward` (e.g. to retune the kernel
    tiling for their hardware).
    """

    def forward(
        self,
        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.

        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

forward(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.

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 
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
154
155
156
157
158
159
160
161
162
163
164
165
166
167
def forward(
    self,
    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.

    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

_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 
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
92
@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)

Compatibility wrapper around :class:VitMropeSetup.

Source code in vllm/model_executor/layers/rotary_embedding/mrope_vit_setup.py 
170
171
172
173
174
175
176
177
178
179
180
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]:
    """Compatibility wrapper around :class:`VitMropeSetup`."""
    return VitMropeSetup()(
        inv_freq_t, inv_freq_h, inv_freq_w, grid_thw, spatial_merge_size
    )