@tduncan / Nvidia-Optimization-Stack-VLLM
Code Commits Issues Pulls Cloud Agent CI/CD
main
main
py 465 lines 12.8 KB
Raw
1 # SPDX-License-Identifier: Apache-2.0
2 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 """Benchmark the fused MoE-LoRA fast path (one-shot) vs two-kernel baseline.
4
5 The "one_shot" provider goes through `vllm.lora.ops.triton_ops.fused_moe_lora`
6 which dispatches to the single-kernel one-shot implementation when
7 fully_sharded=False (the prefill default).
8
9 The "two_kernel" provider drives `fused_moe_lora_shrink` + `fused_moe_lora_expand`
10 directly, bypassing the dispatch and matching the legacy two-kernel path's
11 work distribution. This isolates the win from kernel fusion.
12
13 Run:
14 .venv/bin/python -m benchmarks.kernels.benchmark_fused_moe_lora_one_shot
15 .venv/bin/python -m benchmarks.kernels.benchmark_fused_moe_lora_one_shot \\
16 --model qwen3moe
17 """
18
19 from __future__ import annotations
20
21 import argparse
22 import os
23 import random
24
25 import torch
26
27 from vllm import _custom_ops as ops
28 from vllm.lora.ops.triton_ops import (
29 fused_moe_lora,
30 fused_moe_lora_expand,
31 fused_moe_lora_shrink,
32 )
33 from vllm.triton_utils import triton
34
35 DTYPE = torch.bfloat16
36 DEVICE = "cuda"
37
38
39 # ----- input fabrication -----------------------------------------------------
40
41
42 def _round_up(x: int, base: int) -> int:
43 return ((x + base - 1) // base) * base
44
45
46 def _ceildiv(x: int, y: int) -> int:
47 return (x + y - 1) // y
48
49
50 def _assign_loras(num_tokens: int, num_sequences: int, max_loras: int) -> torch.Tensor:
51 tokens_per_seq = num_tokens // num_sequences
52 rem = num_tokens % num_sequences
53 out = torch.empty(num_tokens, dtype=torch.int32)
54 start = 0
55 for i in range(num_sequences):
56 end = start + tokens_per_seq + (1 if i < rem else 0)
57 out[start:end] = random.randint(0, max_loras - 1)
58 start = end
59 return out
60
61
62 def _assign_experts(num_tokens: int, num_experts: int, top_k: int):
63 expert_indices = torch.empty((num_tokens, top_k), dtype=torch.int32)
64 for i in range(num_tokens):
65 expert_indices[i] = torch.randperm(num_experts)[:top_k]
66 weights = torch.rand((num_tokens, top_k), dtype=torch.float32)
67 weights = weights / weights.sum(dim=1, keepdim=True)
68 return expert_indices, weights
69
70
71 def _make_inputs(
72 M: int,
73 K: int,
74 N_per_slice: int,
75 rank: int,
76 num_experts: int,
77 top_k: int,
78 max_loras: int,
79 num_slices: int,
80 block_size_m: int,
81 ):
82 """Mirrors the production caller's tensor layout."""
83 torch.manual_seed(0)
84 random.seed(0)
85
86 num_sequences = max(1, min(M, 8))
87 topk_ids_cpu, topk_weights_cpu = _assign_experts(M, num_experts, top_k)
88 token_lora_cpu = _assign_loras(M, num_sequences, max_loras)
89 lora_ids_cpu = torch.full((max_loras + 1,), -1, dtype=torch.int32)
90 uniq = torch.unique(token_lora_cpu, sorted=True)
91 lora_ids_cpu[: uniq.size(0)].copy_(uniq)
92
93 topk_ids = topk_ids_cpu.to(DEVICE)
94 topk_weights = topk_weights_cpu.to(device=DEVICE, dtype=DTYPE)
95 token_lora_mapping = token_lora_cpu.to(DEVICE)
96 lora_ids = lora_ids_cpu.to(DEVICE)
97 adapter_enabled = torch.ones(max_loras + 1, dtype=torch.int32, device=DEVICE)
98
99 lora_a = [
100 torch.randn((max_loras, num_experts, rank, K), dtype=DTYPE, device=DEVICE)
101 / max(K, 1) ** 0.5
102 for _ in range(num_slices)
103 ]
104 lora_b = [
105 torch.randn(
106 (max_loras, num_experts, N_per_slice, rank),
107 dtype=DTYPE,
108 device=DEVICE,
109 )
110 / max(rank, 1) ** 0.5
111 for _ in range(num_slices)
112 ]
113 hidden = torch.randn((M, K), dtype=DTYPE, device=DEVICE)
114 out_template = torch.zeros(
115 (M, top_k, num_slices * N_per_slice), dtype=DTYPE, device=DEVICE
116 )
117
118 # Sorted-path metadata (the prefill default).
119 max_pad = topk_ids.numel() + num_experts * (block_size_m - 1)
120 max_pad = _round_up(max_pad, block_size_m)
121 max_blocks = _ceildiv(max_pad, block_size_m)
122 sorted_token_ids = torch.empty(
123 (max_loras * max_pad,), dtype=torch.int32, device=DEVICE
124 )
125 expert_ids = torch.empty(
126 (max_loras * max_blocks,), dtype=torch.int32, device=DEVICE
127 )
128 num_post = torch.empty((max_loras,), dtype=torch.int32, device=DEVICE)
129 ops.moe_lora_align_block_size(
130 topk_ids,
131 token_lora_mapping,
132 num_experts,
133 block_size_m,
134 max_loras,
135 max_pad,
136 max_blocks,
137 sorted_token_ids,
138 expert_ids,
139 num_post,
140 adapter_enabled,
141 lora_ids,
142 )
143 expert_ids = expert_ids.view(max_loras, -1).contiguous()
144 sorted_token_ids = sorted_token_ids.view(max_loras, -1).contiguous()
145 num_active = torch.tensor([max_loras + 1], dtype=torch.int32, device="cpu")
146
147 return dict(
148 hidden=hidden,
149 lora_a=lora_a,
150 lora_b=lora_b,
151 topk_weights=topk_weights,
152 sorted_token_ids=sorted_token_ids,
153 expert_ids=expert_ids,
154 num_post=num_post,
155 token_lora_mapping=token_lora_mapping,
156 lora_ids=lora_ids,
157 num_active=num_active,
158 adapter_enabled=adapter_enabled,
159 out_template=out_template,
160 # bookkeeping
161 M=M,
162 K=K,
163 N_per_slice=N_per_slice,
164 rank=rank,
165 num_experts=num_experts,
166 top_k=top_k,
167 max_loras=max_loras,
168 num_slices=num_slices,
169 block_size_m=block_size_m,
170 )
171
172
173 # ----- providers -------------------------------------------------------------
174
175
176 def _run_one_shot(inp: dict):
177 """Drive `fused_moe_lora` with fully_sharded=False -> one-shot fast path."""
178 out = inp["out_template"].clone()
179 fused_moe_lora(
180 out,
181 inp["hidden"],
182 inp["lora_a"],
183 inp["lora_b"],
184 inp["topk_weights"],
185 inp["sorted_token_ids"],
186 inp["expert_ids"],
187 inp["num_post"],
188 inp["token_lora_mapping"],
189 inp["rank"],
190 inp["top_k"],
191 inp["lora_ids"],
192 inp["num_active"],
193 inp["adapter_enabled"],
194 inp["block_size_m"],
195 64,
196 32,
197 8,
198 4,
199 3,
200 1,
201 inp["block_size_m"],
202 64,
203 32,
204 8,
205 4,
206 3,
207 1,
208 False,
209 False,
210 0,
211 )
212 return out
213
214
215 def _run_two_kernel(inp: dict):
216 """Drive `fused_moe_lora_shrink` + `fused_moe_lora_expand` directly,
217 bypassing the dispatch. Matches the legacy two-kernel work distribution.
218 """
219 M = inp["M"]
220 top_k = inp["top_k"]
221 rank = inp["rank"]
222 num_slices = inp["num_slices"]
223 N_per_slice = inp["N_per_slice"]
224 K = inp["K"]
225 num_experts = inp["num_experts"]
226 block_m = inp["block_size_m"]
227
228 intermediate = torch.zeros((num_slices, M, top_k, rank), dtype=DTYPE, device=DEVICE)
229 out = inp["out_template"].clone()
230 EM = inp["sorted_token_ids"].shape[1]
231 num_tokens = M * top_k
232
233 fused_moe_lora_shrink(
234 intermediate,
235 inp["hidden"],
236 inp["lora_a"],
237 inp["topk_weights"],
238 inp["sorted_token_ids"],
239 inp["expert_ids"],
240 inp["num_post"],
241 inp["token_lora_mapping"],
242 top_k,
243 inp["lora_ids"],
244 inp["adapter_enabled"],
245 torch.device(DEVICE),
246 rank,
247 M,
248 EM,
249 K,
250 num_tokens,
251 num_experts,
252 num_slices,
253 block_m,
254 64,
255 32,
256 8,
257 4,
258 3,
259 1,
260 inp["num_active"],
261 False,
262 )
263 fused_moe_lora_expand(
264 out,
265 intermediate,
266 inp["lora_b"],
267 inp["topk_weights"],
268 inp["sorted_token_ids"],
269 inp["expert_ids"],
270 inp["num_post"],
271 inp["token_lora_mapping"],
272 top_k,
273 inp["lora_ids"],
274 inp["adapter_enabled"],
275 torch.device(DEVICE),
276 rank,
277 M,
278 EM,
279 K,
280 num_tokens,
281 num_experts,
282 num_slices,
283 rank,
284 N_per_slice,
285 block_m,
286 64,
287 32,
288 8,
289 4,
290 3,
291 1,
292 inp["num_active"],
293 False,
294 0,
295 )
296 return out
297
298
299 PROVIDER_FNS = {
300 "one_shot": _run_one_shot,
301 "two_kernel": _run_two_kernel,
302 }
303
304
305 # ----- model presets ---------------------------------------------------------
306
307
308 MODEL_PRESETS: dict[str, dict] = {
309 # Mixtral-8x7B style: E=8, top_k=2, hidden=4096, intermediate=14336
310 "mixtral": dict(
311 K=4096,
312 N_per_slice=7168,
313 num_experts=8,
314 top_k=2,
315 max_loras=4,
316 num_slices=2,
317 block_size_m=64,
318 ),
319 # Qwen3-MoE / DeepSeek-V2 style: E=64, top_k=8, hidden=2048, inter=1408
320 "qwen3moe": dict(
321 K=2048,
322 N_per_slice=1408,
323 num_experts=64,
324 top_k=8,
325 max_loras=4,
326 num_slices=2,
327 block_size_m=64,
328 ),
329 # GLM-5.1 (zai-org/GLM-5.1-FP8): E=256, top_k=8, hidden=6144,
330 # moe_intermediate=2048
331 "glm5_1": dict(
332 K=6144,
333 N_per_slice=2048,
334 num_experts=256,
335 top_k=8,
336 max_loras=4,
337 num_slices=2,
338 block_size_m=64,
339 ),
340 }
341
342
343 M_RANGE = [16, 64, 256, 1024, 4096, 16384]
344 RANK_RANGE = [8, 16, 32, 64]
345
346
347 def get_benchmark(model: str, max_loras: int | None = None):
348 preset = dict(MODEL_PRESETS[model])
349 if max_loras is not None:
350 preset["max_loras"] = max_loras
351
352 @triton.testing.perf_report(
353 triton.testing.Benchmark(
354 x_names=["M", "rank"],
355 x_vals=[(M, R) for M in M_RANGE for R in RANK_RANGE],
356 line_arg="provider",
357 line_vals=list(PROVIDER_FNS.keys()),
358 line_names=["one_shot (fused)", "two_kernel (legacy)"],
359 styles=[("red", "-"), ("blue", "-")],
360 ylabel="ms",
361 plot_name=f"fused_moe_lora-{model}-loras{preset['max_loras']}",
362 args={"preset": preset},
363 )
364 )
365 def benchmark(M, rank, provider, preset):
366 inp = _make_inputs(
367 M=M,
368 K=preset["K"],
369 N_per_slice=preset["N_per_slice"],
370 rank=rank,
371 num_experts=preset["num_experts"],
372 top_k=preset["top_k"],
373 max_loras=preset["max_loras"],
374 num_slices=preset["num_slices"],
375 block_size_m=preset["block_size_m"],
376 )
377 fn = PROVIDER_FNS[provider]
378 quantiles = [0.5, 0.2, 0.8]
379 ms, min_ms, max_ms = triton.testing.do_bench(
380 lambda: fn(inp), quantiles=quantiles
381 )
382 return ms, max_ms, min_ms
383
384 return benchmark
385
386
387 # ----- correctness sanity ---------------------------------------------------
388
389
390 def calculate_diff(model: str, M: int, rank: int, max_loras: int | None = None):
391 preset = dict(MODEL_PRESETS[model])
392 if max_loras is not None:
393 preset["max_loras"] = max_loras
394 inp = _make_inputs(
395 M=M,
396 K=preset["K"],
397 N_per_slice=preset["N_per_slice"],
398 rank=rank,
399 num_experts=preset["num_experts"],
400 top_k=preset["top_k"],
401 max_loras=preset["max_loras"],
402 num_slices=preset["num_slices"],
403 block_size_m=preset["block_size_m"],
404 )
405 out_one = _run_one_shot(inp)
406 out_two = _run_two_kernel(inp)
407 max_abs = (out_one.float() - out_two.float()).abs().max().item()
408 print(
409 f" model={model:<9} M={M:<6} rank={rank:<3} "
410 f"max|one_shot - two_kernel|={max_abs:.4g} "
411 f"ref|max|={out_two.float().abs().max().item():.3g}"
412 )
413 if max_abs <= 5e-2:
414 print(" ✅ outputs match within bf16 tolerance")
415 else:
416 print(" ❌ outputs differ beyond expected bf16 noise")
417
418
419 # ----- main ------------------------------------------------------------------
420
421
422 if __name__ == "__main__":
423 parser = argparse.ArgumentParser()
424 parser.add_argument(
425 "--model",
426 type=str,
427 default="mixtral",
428 choices=list(MODEL_PRESETS.keys()),
429 help="Model preset to sweep",
430 )
431 parser.add_argument(
432 "--save-path",
433 type=str,
434 default="./configs/fused_moe_lora_one_shot/",
435 help="Directory to save benchmark results",
436 )
437 parser.add_argument(
438 "--check-only",
439 action="store_true",
440 help="Run correctness sanity check only, no perf sweep",
441 )
442 parser.add_argument(
443 "--max-loras",
444 type=int,
445 default=None,
446 help="Override max_loras in the model preset (number of LoRA adapters "
447 "active in the batch). Defaults to the preset's value.",
448 )
449 args = parser.parse_args()
450
451 print(f"Correctness check ({args.model}):")
452 calculate_diff(args.model, M=256, rank=32, max_loras=args.max_loras)
453 if args.check_only:
454 raise SystemExit(0)
455
456 effective_max_loras = (
457 args.max_loras
458 if args.max_loras is not None
459 else MODEL_PRESETS[args.model]["max_loras"]
460 )
461 print(f"\nGPU: {torch.cuda.get_device_name()}")
462 print(f"Model preset: {args.model} max_loras={effective_max_loras}\n")
463 benchmark = get_benchmark(args.model, max_loras=args.max_loras)
464 os.makedirs(args.save_path, exist_ok=True)
465 benchmark.run(print_data=True, save_path=args.save_path)