A high-throughput and memory-efficient inference and serving engine for LLMs
Mirrored from https://github.com/tduncan-lum/Nvidia-Optimization-Stack---VLLM.git
Synced Nov 5, 16:00 UTC
[Perf] triton bilinear_pos_embed kernel for ViT (#37948)
Signed-off-by: Zhanda Zhu <zhandazhu@gmail.com>
Zhanda Zhu committed
Apr 1, 2026 at 04:52 UTC
c75a31382418111b2d4cde4b67dd6352eb980bc9
3 files changed
+491
-54
benchmarks/kernels/benchmark_vit_bilinear_pos_embed.py
new
+162
@@ -0,0 +1,162 @@
1
+# SPDX-License-Identifier: Apache-2.0
2
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
+
4
+# Benchmarks the fused Triton bilinear position-embedding kernel against
5
+# the pure-PyTorch (native) implementation used in Qwen3-VL ViT models.
6
+#
7
+# == Usage Examples ==
8
+#
9
+# Default benchmark:
10
+# python3 benchmark_vit_bilinear_pos_embed.py
11
+#
12
+# Custom parameters:
13
+# python3 benchmark_vit_bilinear_pos_embed.py --hidden-dim 1152 \
14
+# --num-grid-per-side 48 --save-path ./configs/vit_pos_embed/
15
+
16
+import itertools
17
+
18
+import torch
19
+
20
+from vllm.model_executor.models.qwen3_vl import (
21
+ pos_embed_interpolate_native,
22
+ triton_pos_embed_interpolate,
23
+)
24
+from vllm.triton_utils import HAS_TRITON, triton
25
+from vllm.utils.argparse_utils import FlexibleArgumentParser
26
+
27
+# (h, w) configurations to benchmark
28
+h_w_configs = [
29
+ (16, 16),
30
+ (32, 32),
31
+ (48, 48),
32
+ (64, 64),
33
+ (128, 128),
34
+ (32, 48),
35
+ (60, 80),
36
+]
37
+
38
+# Temporal dimensions
39
+t_range = [1]
40
+
41
+configs = list(itertools.product(t_range, h_w_configs))
42
+
43
+
44
+def get_benchmark(
45
+ num_grid_per_side: int,
46
+ spatial_merge_size: int,
47
+ hidden_dim: int,
48
+ dtype: torch.dtype,
49
+ device: str,
50
+):
51
+ @triton.testing.perf_report(
52
+ triton.testing.Benchmark(
53
+ x_names=["t", "h_w"],
54
+ x_vals=[list(_) for _ in configs],
55
+ line_arg="provider",
56
+ line_vals=["native", "triton"],
57
+ line_names=["Native (PyTorch)", "Triton"],
58
+ styles=[("blue", "-"), ("red", "-")],
59
+ ylabel="us",
60
+ plot_name=(
61
+ f"vit-bilinear-pos-embed-"
62
+ f"grid{num_grid_per_side}-"
63
+ f"dim{hidden_dim}-"
64
+ f"{dtype}"
65
+ ),
66
+ args={},
67
+ )
68
+ )
69
+ def benchmark(t, h_w, provider):
70
+ h, w = h_w
71
+
72
+ torch.manual_seed(42)
73
+ embed_weight = (
74
+ torch.randn(
75
+ num_grid_per_side * num_grid_per_side,
76
+ hidden_dim,
77
+ device=device,
78
+ dtype=dtype,
79
+ )
80
+ * 0.25
81
+ )
82
+
83
+ quantiles = [0.5, 0.2, 0.8]
84
+
85
+ if provider == "native":
86
+ ms, min_ms, max_ms = triton.testing.do_bench(
87
+ lambda: pos_embed_interpolate_native(
88
+ embed_weight,
89
+ t,
90
+ h,
91
+ w,
92
+ num_grid_per_side,
93
+ spatial_merge_size,
94
+ dtype,
95
+ ),
96
+ quantiles=quantiles,
97
+ )
98
+ else:
99
+ assert HAS_TRITON, "Triton not available"
100
+ ms, min_ms, max_ms = triton.testing.do_bench(
101
+ lambda: triton_pos_embed_interpolate(
102
+ embed_weight,
103
+ t,
104
+ h,
105
+ w,
106
+ num_grid_per_side,
107
+ spatial_merge_size,
108
+ dtype,
109
+ ),
110
+ quantiles=quantiles,
111
+ )
112
+
113
+ return 1000 * ms, 1000 * max_ms, 1000 * min_ms
114
+
115
+ return benchmark
116
+
117
+
118
+if __name__ == "__main__":
119
+ parser = FlexibleArgumentParser(
120
+ description="Benchmark bilinear position embedding interpolation."
121
+ )
122
+ parser.add_argument(
123
+ "--num-grid-per-side",
124
+ type=int,
125
+ default=48,
126
+ help="Position embedding grid size (default: 48 for Qwen3-VL)",
127
+ )
128
+ parser.add_argument(
129
+ "--spatial-merge-size",
130
+ type=int,
131
+ default=2,
132
+ help="Spatial merge size (default: 2)",
133
+ )
134
+ parser.add_argument(
135
+ "--hidden-dim",
136
+ type=int,
137
+ default=1152,
138
+ help="Embedding hidden dimension (default: 1152 for Qwen3-VL)",
139
+ )
140
+ parser.add_argument(
141
+ "--device",
142
+ type=str,
143
+ choices=["cuda:0", "cuda:1"],
144
+ default="cuda:0",
145
+ )
146
+ parser.add_argument(
147
+ "--save-path",
148
+ type=str,
149
+ default="./vit_pos_embed/",
150
+ )
151
+ args = parser.parse_args()
152
+
153
+ dtype = torch.bfloat16
154
+
155
+ bench = get_benchmark(
156
+ args.num_grid_per_side,
157
+ args.spatial_merge_size,
158
+ args.hidden_dim,
159
+ dtype,
160
+ args.device,
161
+ )
162
+ bench.run(print_data=True, save_path=args.save_path)
tests/kernels/core/test_vit_bilinear_pos_embed.py
new
+120
@@ -0,0 +1,120 @@
1
+# SPDX-License-Identifier: Apache-2.0
2
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
+"""Accuracy tests for the fused Triton bilinear position-embedding kernel.
4
+
5
+Compares ``triton_pos_embed_interpolate`` against the pure-PyTorch
6
+``pos_embed_interpolate_native`` across a variety of grid shapes and dtypes.
7
+"""
8
+
9
+import pytest
10
+import torch
11
+
12
+from vllm.triton_utils import HAS_TRITON
13
+
14
+if HAS_TRITON:
15
+ from vllm.model_executor.models.qwen3_vl import (
16
+ pos_embed_interpolate_native,
17
+ triton_pos_embed_interpolate,
18
+ )
19
+
20
+
21
+DTYPES = [torch.float32, torch.bfloat16]
22
+# Qwen3-VL default
23
+NUM_GRID_PER_SIDE = 48
24
+SPATIAL_MERGE_SIZE = 2
25
+HIDDEN_DIM = 1152
26
+
27
+# 4 square + 4 non-square grids (h, w divisible by spatial_merge_size=2)
28
+SQUARE_GRIDS = [(1, 4, 4), (1, 16, 16), (1, 32, 32), (1, 48, 48)]
29
+NON_SQUARE_GRIDS = [(1, 8, 16), (1, 14, 20), (1, 32, 48), (1, 60, 80)]
30
+ALL_GRIDS = SQUARE_GRIDS + NON_SQUARE_GRIDS
31
+
32
+
33
+@pytest.mark.skipif(not HAS_TRITON, reason="Triton not available")
34
+@pytest.mark.parametrize("dtype", DTYPES, ids=lambda d: str(d).split(".")[-1])
35
+@pytest.mark.parametrize(
36
+ "grid_thw",
37
+ ALL_GRIDS,
38
+ ids=[f"{t}x{h}x{w}" for t, h, w in ALL_GRIDS],
39
+)
40
+def test_triton_matches_native(
41
+ grid_thw: tuple[int, int, int],
42
+ dtype: torch.dtype,
43
+) -> None:
44
+ """Triton kernel output must match the native PyTorch implementation."""
45
+ t, h, w = grid_thw
46
+ device = "cuda"
47
+
48
+ # Scale to match real Qwen3-VL pos_embed weight distribution (std~0.23).
49
+ torch.manual_seed(42)
50
+ embed_weight = (
51
+ torch.randn(
52
+ NUM_GRID_PER_SIDE * NUM_GRID_PER_SIDE,
53
+ HIDDEN_DIM,
54
+ device=device,
55
+ dtype=dtype,
56
+ )
57
+ * 0.25
58
+ )
59
+
60
+ native_out = pos_embed_interpolate_native(
61
+ embed_weight, t, h, w, NUM_GRID_PER_SIDE, SPATIAL_MERGE_SIZE, dtype
62
+ )
63
+ triton_out = triton_pos_embed_interpolate(
64
+ embed_weight, t, h, w, NUM_GRID_PER_SIDE, SPATIAL_MERGE_SIZE, dtype
65
+ )
66
+
67
+ assert native_out.shape == triton_out.shape, (
68
+ f"Shape mismatch: native {native_out.shape} vs triton {triton_out.shape}"
69
+ )
70
+
71
+ # Small numerical differences arise from the precomputed h/w_scale
72
+ # in the triton kernel vs torch.linspace in the native path, which can
73
+ # cause single-ULP output differences
74
+ # in a handful of elements.
75
+ atol = {torch.float32: 5e-5, torch.bfloat16: 1e-2}[dtype]
76
+ rtol = {torch.float32: 1e-5, torch.bfloat16: 1e-2}[dtype]
77
+ torch.testing.assert_close(triton_out, native_out, atol=atol, rtol=rtol)
78
+
79
+
80
+@pytest.mark.skipif(not HAS_TRITON, reason="Triton not available")
81
+@pytest.mark.parametrize("dtype", DTYPES, ids=lambda d: str(d).split(".")[-1])
82
+def test_temporal_repeat(dtype: torch.dtype) -> None:
83
+ """Verify temporal dimension t > 1 correctly repeats the spatial pattern."""
84
+ device = "cuda"
85
+ h, w = 16, 16
86
+ t_single, t_multi = 1, 3
87
+
88
+ # Scale to match real Qwen3-VL pos_embed weight distribution (std~0.23).
89
+ torch.manual_seed(42)
90
+ embed_weight = (
91
+ torch.randn(
92
+ NUM_GRID_PER_SIDE * NUM_GRID_PER_SIDE,
93
+ HIDDEN_DIM,
94
+ device=device,
95
+ dtype=dtype,
96
+ )
97
+ * 0.25
98
+ )
99
+
100
+ out_single = triton_pos_embed_interpolate(
101
+ embed_weight,
102
+ t_single,
103
+ h,
104
+ w,
105
+ NUM_GRID_PER_SIDE,
106
+ SPATIAL_MERGE_SIZE,
107
+ dtype,
108
+ )
109
+ out_multi = triton_pos_embed_interpolate(
110
+ embed_weight,
111
+ t_multi,
112
+ h,
113
+ w,
114
+ NUM_GRID_PER_SIDE,
115
+ SPATIAL_MERGE_SIZE,
116
+ dtype,
117
+ )
118
+
119
+ expected = out_single.repeat(t_multi, 1)
120
+ torch.testing.assert_close(out_multi, expected, atol=0, rtol=0)
vllm/model_executor/models/qwen3_vl.py
+209
-54
@@ -96,6 +96,7 @@ from vllm.multimodal.processing import (
96
from vllm.sequence import IntermediateTensors
97
from vllm.tokenizers.protocol import TokenizerLike
98
from vllm.tokenizers.registry import cached_tokenizer_from_config
99
+from vllm.triton_utils import HAS_TRITON, tl, triton
100
from vllm.utils.collection_utils import is_list_of
101
from vllm.utils.math_utils import round_up
102
@@ -145,6 +146,201 @@ logger = init_logger(__name__)
146
# of the maximum size.
147
DUMMY_VIDEO_NUM_FRAMES = 2048
148
149
+# ---------------------------------------------------------------------------
150
+# Triton kernel: fused bilinear position-embedding interpolation
151
+# ---------------------------------------------------------------------------
152
+# Replaces many small eager-mode CUDA kernels with a single launch.
153
+# The spatial-merge reorder is baked into the index math so the output
154
+# is ready to be added to the patch embeddings directly.
155
+# ---------------------------------------------------------------------------
156
+
157
+if HAS_TRITON:
158
+
159
+ @triton.jit
160
+ def _bilinear_pos_embed_kernel(
161
+ embed_ptr,
162
+ output_ptr,
163
+ H,
164
+ W,
165
+ h_scale,
166
+ w_scale,
167
+ NUM_GRID: tl.constexpr,
168
+ M_SIZE: tl.constexpr,
169
+ HIDDEN_DIM: tl.constexpr,
170
+ BLOCK_D: tl.constexpr,
171
+ ):
172
+ """Fused bilinear pos-embed interpolation with spatial-merge reorder."""
173
+ pid = tl.program_id(0)
174
+ total_spatial = H * W
175
+ spatial_idx = pid % total_spatial
176
+
177
+ num_blocks_w = W // M_SIZE
178
+ block_idx = spatial_idx // (M_SIZE * M_SIZE)
179
+ local_idx = spatial_idx % (M_SIZE * M_SIZE)
180
+ br = block_idx // num_blocks_w
181
+ bc = block_idx % num_blocks_w
182
+ lr = local_idx // M_SIZE
183
+ lc = local_idx % M_SIZE
184
+ row = br * M_SIZE + lr
185
+ col = bc * M_SIZE + lc
186
+
187
+ h_frac = row.to(tl.float32) * h_scale
188
+ w_frac = col.to(tl.float32) * w_scale
189
+
190
+ hf = tl.math.floor(h_frac).to(tl.int32)
191
+ wf = tl.math.floor(w_frac).to(tl.int32)
192
+ hc = tl.minimum(hf + 1, NUM_GRID - 1)
193
+ wc = tl.minimum(wf + 1, NUM_GRID - 1)
194
+
195
+ dh = h_frac - hf.to(tl.float32)
196
+ dw = w_frac - wf.to(tl.float32)
197
+ w11 = dh * dw
198
+ w10 = dh - w11
199
+ w01 = dw - w11
200
+ w00 = 1.0 - dh - w01
201
+
202
+ off00 = (hf * NUM_GRID + wf) * HIDDEN_DIM
203
+ off01 = (hf * NUM_GRID + wc) * HIDDEN_DIM
204
+ off10 = (hc * NUM_GRID + wf) * HIDDEN_DIM
205
+ off11 = (hc * NUM_GRID + wc) * HIDDEN_DIM
206
+ out_off = pid * HIDDEN_DIM
207
+
208
+ # Cast weights to output dtype so the multiply-accumulate stays
209
+ # in the same precision as the native PyTorch implementation.
210
+ out_dtype = output_ptr.dtype.element_ty
211
+ w00_c = w00.to(out_dtype)
212
+ w01_c = w01.to(out_dtype)
213
+ w10_c = w10.to(out_dtype)
214
+ w11_c = w11.to(out_dtype)
215
+
216
+ for d in tl.range(0, HIDDEN_DIM, BLOCK_D):
217
+ cols = d + tl.arange(0, BLOCK_D)
218
+ mask = cols < HIDDEN_DIM
219
+
220
+ e00 = tl.load(embed_ptr + off00 + cols, mask=mask)
221
+ e01 = tl.load(embed_ptr + off01 + cols, mask=mask)
222
+ e10 = tl.load(embed_ptr + off10 + cols, mask=mask)
223
+ e11 = tl.load(embed_ptr + off11 + cols, mask=mask)
224
+
225
+ val = w00_c * e00 + w01_c * e01 + w10_c * e10 + w11_c * e11
226
+
227
+ tl.store(output_ptr + out_off + cols, val, mask=mask)
228
+
229
+ def triton_pos_embed_interpolate(
230
+ embed_weight: torch.Tensor,
231
+ t: int,
232
+ h: int,
233
+ w: int,
234
+ num_grid_per_side: int,
235
+ m_size: int,
236
+ dtype: torch.dtype,
237
+ ) -> torch.Tensor:
238
+ """Launch the fused Triton kernel for one (t,h,w) grid.
239
+
240
+ Returns a tensor of shape ``(t * h * w, hidden_dim)`` with the
241
+ bilinearly-interpolated position embeddings in spatial-merge order.
242
+ """
243
+ assert h % m_size == 0 and w % m_size == 0, (
244
+ f"h={h} and w={w} must be divisible by m_size={m_size}"
245
+ )
246
+ hidden_dim = embed_weight.shape[1]
247
+ total_out = t * h * w
248
+ output = torch.empty(
249
+ total_out,
250
+ hidden_dim,
251
+ device=embed_weight.device,
252
+ dtype=dtype,
253
+ )
254
+
255
+ h_scale = float(num_grid_per_side - 1) / float(h - 1) if h > 1 else 0.0
256
+ w_scale = float(num_grid_per_side - 1) / float(w - 1) if w > 1 else 0.0
257
+
258
+ BLOCK_D = triton.next_power_of_2(hidden_dim)
259
+
260
+ _bilinear_pos_embed_kernel[(total_out,)](
261
+ embed_weight,
262
+ output,
263
+ h,
264
+ w,
265
+ h_scale,
266
+ w_scale,
267
+ num_grid_per_side,
268
+ m_size,
269
+ hidden_dim,
270
+ BLOCK_D,
271
+ )
272
+ return output
273
+
274
+
275
+def pos_embed_interpolate_native(
276
+ embed_weight: torch.Tensor,
277
+ t: int,
278
+ h: int,
279
+ w: int,
280
+ num_grid_per_side: int,
281
+ m_size: int,
282
+ dtype: torch.dtype,
283
+) -> torch.Tensor:
284
+ """Eager PyTorch bilinear position-embedding interpolation.
285
+
286
+ Returns a tensor of shape ``(t * h * w, hidden_dim)`` with the
287
+ bilinearly-interpolated position embeddings in spatial-merge order.
288
+ """
289
+ assert h % m_size == 0 and w % m_size == 0, (
290
+ f"h={h} and w={w} must be divisible by m_size={m_size}"
291
+ )
292
+ hidden_dim = embed_weight.shape[1]
293
+ device = embed_weight.device
294
+
295
+ h_idxs = torch.linspace(
296
+ 0,
297
+ num_grid_per_side - 1,
298
+ h,
299
+ dtype=torch.float32,
300
+ device=device,
301
+ )
302
+ w_idxs = torch.linspace(
303
+ 0,
304
+ num_grid_per_side - 1,
305
+ w,
306
+ dtype=torch.float32,
307
+ device=device,
308
+ )
309
+
310
+ h_floor = h_idxs.to(torch.long)
311
+ w_floor = w_idxs.to(torch.long)
312
+ h_ceil = torch.clamp(h_floor + 1, max=num_grid_per_side - 1)
313
+ w_ceil = torch.clamp(w_floor + 1, max=num_grid_per_side - 1)
314
+
315
+ dh = h_idxs - h_floor
316
+ dw = w_idxs - w_floor
317
+
318
+ dh_grid, dw_grid = torch.meshgrid(dh, dw, indexing="ij")
319
+ h_floor_grid, w_floor_grid = torch.meshgrid(h_floor, w_floor, indexing="ij")
320
+ h_ceil_grid, w_ceil_grid = torch.meshgrid(h_ceil, w_ceil, indexing="ij")
321
+
322
+ w11 = dh_grid * dw_grid
323
+ w10 = dh_grid - w11
324
+ w01 = dw_grid - w11
325
+ w00 = 1 - dh_grid - w01
326
+
327
+ h_grid = torch.stack([h_floor_grid, h_floor_grid, h_ceil_grid, h_ceil_grid])
328
+ w_grid = torch.stack([w_floor_grid, w_ceil_grid, w_floor_grid, w_ceil_grid])
329
+ h_grid_idx = h_grid * num_grid_per_side
330
+
331
+ indices = (h_grid_idx + w_grid).reshape(4, -1)
332
+ weights = torch.stack([w00, w01, w10, w11], dim=0).reshape(4, -1, 1)
333
+ weights = weights.to(dtype=dtype)
334
+
335
+ embeds = embed_weight[indices]
336
+ embeds *= weights
337
+ combined = embeds.sum(dim=0)
338
+
339
+ combined = combined.reshape(h // m_size, m_size, w // m_size, m_size, hidden_dim)
340
+ combined = combined.permute(0, 2, 1, 3, 4).reshape(1, -1, hidden_dim)
341
+ repeated = combined.expand(t, -1, -1).reshape(-1, hidden_dim)
342
+ return repeated.to(dtype=dtype)
343
+
344
345
class Qwen3_VisionPatchEmbed(nn.Module):
346
def __init__(
@@ -470,63 +666,22 @@ class Qwen3_VisionTransformer(nn.Module):
666
return cos_combined, sin_combined
667
668
def fast_pos_embed_interpolate(self, grid_thw: list[list[int]]) -> torch.Tensor:
473
- num_grid_per_side = self.num_grid_per_side
474
- m_size = self.spatial_merge_size
475
- hidden_dim = self.pos_embed.embedding_dim
476
-
669
+ interpolate_fn = (
670
+ triton_pos_embed_interpolate if HAS_TRITON else pos_embed_interpolate_native
671
+ )
672
outputs = []
673
for t, h, w in grid_thw:
479
- h_idxs = torch.linspace(
480
- 0, num_grid_per_side - 1, h, dtype=torch.float32, device=self.device
481
- )
482
- w_idxs = torch.linspace(
483
- 0, num_grid_per_side - 1, w, dtype=torch.float32, device=self.device
484
- )
485
-
486
- h_floor = h_idxs.to(torch.long)
487
- w_floor = w_idxs.to(torch.long)
488
- h_ceil = torch.clamp(h_floor + 1, max=num_grid_per_side - 1)
489
- w_ceil = torch.clamp(w_floor + 1, max=num_grid_per_side - 1)
490
-
491
- dh = h_idxs - h_floor
492
- dw = w_idxs - w_floor
493
-
494
- # Create meshgrid view for all h, w vars
495
- dh_grid, dw_grid = torch.meshgrid(dh, dw, indexing="ij")
496
- h_floor_grid, w_floor_grid = torch.meshgrid(h_floor, w_floor, indexing="ij")
497
- h_ceil_grid, w_ceil_grid = torch.meshgrid(h_ceil, w_ceil, indexing="ij")
498
-
499
- # original computation of weights
500
- # w00 = (1 - dh_grid) * (1 - dw_grid)
501
- # w01 = (1 - dh_grid) * dw_grid
502
- # w10 = dh_grid * (1 - dw_grid)
503
- # w11 = dh_grid * dw_grid
504
- # we reuse w11 here to avoid duplicate
505
- # dh_grid * dw_grid computation
506
- w11 = dh_grid * dw_grid
507
- w10 = dh_grid - w11
508
- w01 = dw_grid - w11
509
- w00 = 1 - dh_grid - w01
510
-
511
- h_grid = torch.stack([h_floor_grid, h_floor_grid, h_ceil_grid, h_ceil_grid])
512
- w_grid = torch.stack([w_floor_grid, w_ceil_grid, w_floor_grid, w_ceil_grid])
513
- h_grid_idx = h_grid * num_grid_per_side
514
-
515
- indices = (h_grid_idx + w_grid).reshape(4, -1)
516
- weights = torch.stack([w00, w01, w10, w11], dim=0).reshape(4, -1, 1)
517
- weights = weights.to(dtype=self.dtype)
518
-
519
- embeds = self.pos_embed(indices)
520
- embeds *= weights
521
- combined = embeds.sum(dim=0)
522
-
523
- combined = combined.reshape(
524
- h // m_size, m_size, w // m_size, m_size, hidden_dim
674
+ outputs.append(
675
+ interpolate_fn(
676
+ self.pos_embed.weight,
677
+ t,
678
+ h,
679
+ w,
680
+ self.num_grid_per_side,
681
+ self.spatial_merge_size,
682
+ self.dtype,
683
+ )
684
)
526
- combined = combined.permute(0, 2, 1, 3, 4).reshape(1, -1, hidden_dim)
527
- repeated = combined.expand(t, -1, -1).reshape(-1, hidden_dim)
528
- outputs.append(repeated)
529
-
685
return torch.cat(outputs, dim=0)
686
687
def prepare_encoder_metadata(