@tduncan / Nvidia-Optimization-Stack-VLLM
Code Commits Issues Pulls Cloud Agent CI/CD
main
main
py 162 lines 4.12 KB
Raw
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)