@tduncan / Nvidia-Optimization-Stack-VLLM Mirror · read-only / commits / 98229db244
Fork 0

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 14, 05:30 UTC

[Kernels][DP/EP] Optimize Silu Kernel for R1 (#24054)

Signed-off-by: elvircrn <elvircrn@gmail.com>

Elvir Crnčević committed Sep 13, 2025 at 09:17 UTC 98229db2444b016d572bfd36e685960b3352d900
6 files changed +1272 -131
benchmarks/kernels/benchmark_silu_mul_fp8_quant.py
+636 -38
@@ -1,77 +1,675 @@
1 -#!/usr/bin/env python3
1 # SPDX-License-Identifier: Apache-2.0
2 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
4 -import time
3 +from collections.abc import Callable
4
5 +import matplotlib.pyplot as plt
6 +import numpy as np
7 import torch
8
9 from vllm.model_executor.layers.fused_moe.batched_deep_gemm_moe import (
9 - silu_mul_fp8_quant_deep_gemm,
10 + silu_mul_fp8_quant_deep_gemm_cuda,
11 )
12 from vllm.platforms import current_platform
13 +from vllm.triton_utils import tl, triton
14 +from vllm.utils.deep_gemm import is_deep_gemm_e8m0_used
15
16
14 -def benchmark(E, T, H, G=128, runs=50):
15 - current_platform.seed_everything(42)
16 - y = torch.randn((E, T, 2 * H), dtype=torch.bfloat16, device="cuda")
17 - tokens_per_expert = torch.randint(
18 - T // 2, T, size=(E,), dtype=torch.int32, device="cuda"
17 +@triton.jit
18 +def _silu_mul_fp8_quant_deep_gemm(
19 + # Pointers ------------------------------------------------------------
20 + input_ptr, # 16-bit activations (E, T, 2*H)
21 + y_q_ptr, # fp8 quantized activations (E, T, H)
22 + y_s_ptr, # 16-bit scales (E, T, G)
23 + counts_ptr, # int32 num tokens per expert (E)
24 + # Sizes ---------------------------------------------------------------
25 + H: tl.constexpr, # hidden dimension (per output)
26 + GROUP_SIZE: tl.constexpr, # elements per group (usually 128)
27 + # Strides for input (elements) ---------------------------------------
28 + stride_i_e,
29 + stride_i_t,
30 + stride_i_h,
31 + # Strides for y_q (elements) -----------------------------------------
32 + stride_yq_e,
33 + stride_yq_t,
34 + stride_yq_h,
35 + # Strides for y_s (elements) -----------------------------------------
36 + stride_ys_e,
37 + stride_ys_t,
38 + stride_ys_g,
39 + # Stride for counts (elements)
40 + stride_counts_e,
41 + # Numeric params ------------------------------------------------------
42 + eps: tl.constexpr,
43 + fp8_min: tl.constexpr,
44 + fp8_max: tl.constexpr,
45 + use_ue8m0: tl.constexpr,
46 + # Meta ---------------------------------------------------------------
47 + BLOCK: tl.constexpr,
48 + NUM_STAGES: tl.constexpr,
49 +):
50 + G = H // GROUP_SIZE
51 +
52 + # map program id -> (e, g)
53 + pid = tl.program_id(0)
54 + e = pid // G
55 + g = pid % G
56 +
57 + e = e.to(tl.int64)
58 + g = g.to(tl.int64)
59 +
60 + # number of valid tokens for this expert
61 + n_tokens = tl.load(counts_ptr + e * stride_counts_e).to(tl.int64)
62 +
63 + cols = tl.arange(0, BLOCK).to(tl.int64)
64 + mask = cols < BLOCK
65 +
66 + base_input_offset = e * stride_i_e + g * GROUP_SIZE * stride_i_h
67 + base_gate_offset = base_input_offset + cols * stride_i_h
68 + base_up_offset = base_input_offset + H * stride_i_h + cols * stride_i_h
69 + base_yq_offset = e * stride_yq_e + g * GROUP_SIZE * stride_yq_h + cols * stride_yq_h
70 + base_ys_offset = e * stride_ys_e + g * stride_ys_g
71 +
72 + for t in tl.range(0, n_tokens, num_stages=NUM_STAGES):
73 + gate = tl.load(
74 + input_ptr + base_gate_offset + t * stride_i_t, mask=mask, other=0.0
75 + ).to(tl.float32)
76 + up = tl.load(input_ptr + base_up_offset + t * stride_i_t, mask=mask, other=0.0)
77 +
78 + gate = gate * (1.0 / (1.0 + tl.exp(-gate)))
79 + y = gate * up
80 +
81 + y_s = tl.maximum(tl.max(tl.abs(y)), eps) / fp8_max
82 + if use_ue8m0:
83 + y_s = tl.exp2(tl.ceil(tl.log2(y_s)))
84 +
85 + y_q = tl.clamp(y / y_s, fp8_min, fp8_max).to(y_q_ptr.dtype.element_ty)
86 +
87 + tl.store(y_q_ptr + base_yq_offset + t * stride_yq_t, y_q, mask=mask)
88 + tl.store(y_s_ptr + base_ys_offset + t * stride_ys_t, y_s)
89 +
90 +
91 +def silu_mul_fp8_quant_deep_gemm_triton(
92 + y: torch.Tensor, # (E, T, 2*H)
93 + tokens_per_expert: torch.Tensor, # (E,) number of valid tokens per expert
94 + num_parallel_tokens,
95 + group_size: int = 128,
96 + eps: float = 1e-10,
97 +) -> tuple[torch.Tensor, torch.Tensor]:
98 + """Quantize silu(y[..., :H]) * y[..., H:] to FP8 with group per-token scales
99 +
100 + y has shape (E, T, 2*H). The first half of the last dimension is
101 + silu-activated, multiplied by the second half, then quantized into FP8.
102 +
103 + Returns `(y_q, y_s)` where
104 + * `y_q`: FP8 tensor, shape (E, T, H), same layout as y[..., :H]
105 + * `y_s`: FP32 tensor, shape (E, T, H // group_size), strides (T*G, 1, T)
106 + """
107 + assert y.ndim == 3, "y must be (E, T, 2*H)"
108 + E, T, H2 = y.shape
109 + assert H2 % 2 == 0, "last dim of y must be even (2*H)"
110 + H = H2 // 2
111 + G = (H + group_size - 1) // group_size
112 + assert H % group_size == 0, "H must be divisible by group_size"
113 + assert tokens_per_expert.ndim == 1 and tokens_per_expert.shape[0] == E, (
114 + "tokens_per_expert must be shape (E,)"
115 + )
116 + tokens_per_expert = tokens_per_expert.to(device=y.device, dtype=torch.int32)
117 +
118 + # allocate outputs
119 + fp8_dtype = torch.float8_e4m3fn
120 + y_q = torch.empty((E, T, H), dtype=fp8_dtype, device=y.device)
121 +
122 + # strides (elements)
123 + stride_i_e, stride_i_t, stride_i_h = y.stride()
124 + stride_yq_e, stride_yq_t, stride_yq_h = y_q.stride()
125 +
126 + # desired scale strides (elements): (T*G, 1, T)
127 + stride_ys_e = T * G
128 + stride_ys_t = 1
129 + stride_ys_g = T
130 + y_s = torch.empty_strided(
131 + (E, T, G),
132 + (stride_ys_e, stride_ys_t, stride_ys_g),
133 + dtype=torch.float32,
134 + device=y.device,
135 + )
136 +
137 + stride_cnt_e = tokens_per_expert.stride()[0]
138 +
139 + # Static grid over experts and H-groups.
140 + # A loop inside the kernel handles the token dim
141 + grid = (E * G,)
142 +
143 + f_info = torch.finfo(fp8_dtype)
144 + fp8_max = f_info.max
145 + fp8_min = f_info.min
146 +
147 + _silu_mul_fp8_quant_deep_gemm[grid](
148 + y,
149 + y_q,
150 + y_s,
151 + tokens_per_expert,
152 + H,
153 + group_size,
154 + stride_i_e,
155 + stride_i_t,
156 + stride_i_h,
157 + stride_yq_e,
158 + stride_yq_t,
159 + stride_yq_h,
160 + stride_ys_e,
161 + stride_ys_t,
162 + stride_ys_g,
163 + stride_cnt_e,
164 + eps,
165 + fp8_min,
166 + fp8_max,
167 + is_deep_gemm_e8m0_used(),
168 + BLOCK=group_size,
169 + NUM_STAGES=4,
170 + num_warps=1,
171 )
172
173 + return y_q, y_s
174 +
175 +
176 +# Parse generation strategies
177 +strategies = ["uniform", "max_t", "first_t"]
178 +
179 +
180 +def benchmark(
181 + kernel: Callable,
182 + E: int,
183 + T: int,
184 + H: int,
185 + total_tokens: int,
186 + num_parallel_tokens: int = 64,
187 + G: int = 128,
188 + runs: int = 200,
189 + num_warmups: int = 20,
190 + gen_strategy: str = "default",
191 + iterations_per_run: int = 20,
192 +):
193 + def generate_data(seed_offset=0):
194 + """Generate input data with given seed offset"""
195 + current_platform.seed_everything(42 + seed_offset)
196 + y = torch.rand((E, T, 2 * H), dtype=torch.bfloat16, device="cuda").contiguous()
197 +
198 + if gen_strategy == "uniform":
199 + r = torch.rand(size=(E,), device="cuda")
200 + r /= r.sum()
201 + r *= total_tokens
202 + tokens_per_expert = r.int()
203 + tokens_per_expert = torch.minimum(
204 + tokens_per_expert,
205 + torch.ones((E,), device=r.device, dtype=torch.int) * T,
206 + )
207 + elif gen_strategy == "max_t":
208 + tokens_per_expert = torch.empty(size=(E,), dtype=torch.int32, device="cuda")
209 + tokens_per_expert.fill_(total_tokens / E)
210 + elif gen_strategy == "first_t":
211 + tokens_per_expert = torch.zeros(size=(E,), dtype=torch.int32, device="cuda")
212 + tokens_per_expert[0] = min(T, total_tokens)
213 + else:
214 + raise ValueError(f"Unknown generation strategy: {gen_strategy}")
215 + return y, tokens_per_expert
216 +
217 + dataset_count = 4
218 + # Pre-generate different input matrices for each iteration to avoid cache effects
219 + data_sets = [generate_data(i) for i in range(dataset_count)]
220 +
221 # Warmup
22 - for _ in range(10):
23 - silu_mul_fp8_quant_deep_gemm(y, tokens_per_expert, group_size=G)
24 - torch.cuda.synchronize()
222 + y, tokens_per_expert = data_sets[0]
223 + for _ in range(num_warmups):
224 + kernel(
225 + y, tokens_per_expert, num_parallel_tokens=num_parallel_tokens, group_size=G
226 + )
227 + torch.cuda.synchronize()
228 +
229 + start_event = torch.cuda.Event(enable_timing=True)
230 + end_event = torch.cuda.Event(enable_timing=True)
231
232 # Benchmark
27 - torch.cuda.synchronize()
28 - start = time.perf_counter()
233 + latencies: list[float] = []
234 for _ in range(runs):
30 - silu_mul_fp8_quant_deep_gemm(y, tokens_per_expert, group_size=G)
31 - torch.cuda.synchronize()
235 + torch.cuda.synchronize()
236 +
237 + start_event.record()
238 + for i in range(iterations_per_run):
239 + y, tokens_per_expert = data_sets[i % dataset_count]
240 + kernel(
241 + y,
242 + tokens_per_expert,
243 + num_parallel_tokens=num_parallel_tokens,
244 + group_size=G,
245 + )
246 + end_event.record()
247 + end_event.synchronize()
248 +
249 + total_time_ms = start_event.elapsed_time(end_event)
250 + per_iter_time_ms = total_time_ms / iterations_per_run
251 + latencies.append(per_iter_time_ms)
252
33 - avg_time = (time.perf_counter() - start) / runs * 1000
253 + # Use median instead of average for better outlier handling
254 + median_time_ms = np.median(latencies)
255 + median_time_s = median_time_ms / 1000
256
35 - # Calculate actual work done (only count valid tokens)
257 + # Calculate actual work done (using first dataset for consistency)
258 + _, tokens_per_expert = data_sets[0]
259 actual_tokens = tokens_per_expert.sum().item()
260 actual_elements = actual_tokens * H
261
262 # GFLOPS: operations per element = exp + 3 muls + 1 div + quantization ops ≈ 8 ops
263 ops_per_element = 8
264 total_ops = actual_elements * ops_per_element
42 - gflops = total_ops / (avg_time / 1000) / 1e9
265 + gflops = total_ops / median_time_s / 1e9
266
267 # Memory bandwidth: bfloat16 inputs (2 bytes), fp8 output (1 byte), scales (4 bytes)
268 input_bytes = actual_tokens * 2 * H * 2 # 2*H bfloat16 inputs
269 output_bytes = actual_tokens * H * 1 # H fp8 outputs
270 scale_bytes = actual_tokens * (H // G) * 4 # scales in float32
271 total_bytes = input_bytes + output_bytes + scale_bytes
49 - memory_bw = total_bytes / (avg_time / 1000) / 1e9
272 + memory_bw = total_bytes / median_time_s / 1e9
273 +
274 + HOPPER_BANDWIDTH_TBPS = 3.35
275 + return (
276 + median_time_ms,
277 + gflops,
278 + memory_bw,
279 + (memory_bw / (HOPPER_BANDWIDTH_TBPS * 1024)) * 100,
280 + )
281 +
282 +
283 +def create_comparison_plot(
284 + ratio, cuda_times, baseline_times, config_labels, strategy_name, id
285 +):
286 + """Create a comparison plot for a specific generation strategy"""
287 + fig, ax = plt.subplots(1, 1, figsize=(16, 6))
288 +
289 + # Configure x-axis positions
290 + x = np.arange(len(config_labels))
291 + width = 0.35
292 +
293 + # Execution Time plot (lower is better)
294 + ax.bar(
295 + x - width / 2, cuda_times, width, label="CUDA Kernel", alpha=0.8, color="blue"
296 + )
297 + ax.bar(
298 + x + width / 2,
299 + baseline_times,
300 + width,
301 + label="Baseline",
302 + alpha=0.8,
303 + color="orange",
304 + )
305 +
306 + # Add speedup labels over each bar pair
307 + for i in range(len(x)):
308 + speedup = ratio[i]
309 + max_height = max(cuda_times[i], baseline_times[i])
310 + ax.text(
311 + x[i],
312 + max_height + max_height * 0.02,
313 + f"{speedup:.2f}x",
314 + ha="center",
315 + va="bottom",
316 + fontweight="bold",
317 + fontsize=9,
318 + )
319 +
320 + ax.set_xlabel("Configuration")
321 + ax.set_ylabel("% Utilization")
322 + ax.set_title(
323 + f"Memory Bandwidth Utilization (%) - {strategy_name}\n(Higher is Better)"
324 + )
325 + ax.set_xticks(x)
326 + ax.set_xticklabels(config_labels, rotation=45, ha="right")
327 + ax.legend()
328 + ax.grid(True, alpha=0.3)
329 +
330 + plt.tight_layout()
331 + return fig, ax
332
51 - return avg_time, gflops, memory_bw
333
334 +def create_combined_plot(all_results):
335 + """Create a combined plot with all strategies in one PNG"""
336 + num_strategies = len(all_results)
337 + fig, axes = plt.subplots(num_strategies, 1, figsize=(20, 6 * num_strategies))
338
339 + if num_strategies == 1:
340 + axes = [axes]
341 +
342 + for idx, (
343 + strategy_name,
344 + ratio,
345 + cuda_times,
346 + baseline_times,
347 + config_labels,
348 + ) in enumerate(all_results):
349 + ax = axes[idx]
350 +
351 + # Configure x-axis positions
352 + x = np.arange(len(config_labels))
353 + width = 0.35
354 +
355 + # Execution Time plot (lower is better)
356 + ax.bar(
357 + x - width / 2,
358 + cuda_times,
359 + width,
360 + label="CUDA Kernel",
361 + alpha=0.8,
362 + color="blue",
363 + )
364 + ax.bar(
365 + x + width / 2,
366 + baseline_times,
367 + width,
368 + label="Baseline",
369 + alpha=0.8,
370 + color="orange",
371 + )
372 +
373 + # Add speedup labels over each bar pair
374 + for i in range(len(x)):
375 + speedup = ratio[i]
376 + max_height = max(cuda_times[i], baseline_times[i])
377 + ax.text(
378 + x[i],
379 + max_height + max_height * 0.02,
380 + f"{speedup:.2f}x",
381 + ha="center",
382 + va="bottom",
383 + fontweight="bold",
384 + fontsize=9,
385 + )
386 +
387 + ax.set_xlabel("Configuration")
388 + ax.set_ylabel("% Utilization")
389 + ax.set_title(
390 + f"Memory Bandwidth Utilization (%) - {strategy_name}\n(Higher is Better)"
391 + )
392 + ax.set_xticks(x)
393 + ax.set_xticklabels(config_labels, rotation=45, ha="right")
394 + ax.legend()
395 + ax.grid(True, alpha=0.3)
396 +
397 + plt.tight_layout()
398 + filename = "../../silu_bench/silu_benchmark_combined.png"
399 + plt.savefig(filename, dpi=300, bbox_inches="tight")
400 + plt.show()
401 +
402 + return filename
403 +
404 +
405 +outer_dim = 7168
406 configs = [
55 - (8, 32, 1024),
56 - (16, 64, 2048),
57 - (32, 128, 4096),
407 # DeepSeekV3 Configs
59 - (256, 16, 7168),
60 - (256, 32, 7168),
61 - (256, 64, 7168),
62 - (256, 128, 7168),
63 - (256, 256, 7168),
64 - (256, 512, 7168),
408 + (8, 1024, 7168),
409 + # DeepSeekV3 Configs
410 + (32, 1024, 7168),
411 + # DeepSeekV3 Configs
412 (256, 1024, 7168),
413 ]
414
415 +runs = 100
416 +num_warmups = 20
417 +
418 +strategy_descriptions = {
419 + "uniform": "Uniform Random",
420 + "max_t": "Even Assignment",
421 + "first_t": "experts[0] = T, experts[1:] = 0",
422 +}
423 +
424 print(f"GPU: {torch.cuda.get_device_name()}")
69 -print(f"{'Config':<20} {'Time(ms)':<10} {'GFLOPS':<10} {'GB/s':<10}")
70 -print("-" * 50)
71 -
72 -for E, T, H in configs:
73 - try:
74 - time_ms, gflops, gbps = benchmark(E, T, H)
75 - print(f"E={E:3d},T={T:4d},H={H:4d} {time_ms:8.3f} {gflops:8.1f} {gbps:8.1f}")
76 - except Exception:
77 - print(f"E={E:3d},T={T:4d},H={H:4d} FAILED")
425 +print(f"Testing strategies: {', '.join(strategies)}")
426 +print(f"Configurations: {len(configs)} configs")
427 +
428 +all_results = []
429 +
430 +# Run benchmarks for each strategy
431 +for id, strategy in enumerate(strategies):
432 + print(f"\n{'=' * 60}")
433 + print(f"Testing strategy: {strategy_descriptions[strategy]}")
434 + print(f"{'=' * 60}")
435 +
436 + # Collect benchmark data for both algorithms
437 + config_labels = []
438 + config_x_axis = []
439 + all_cuda_results = []
440 + all_baseline_results = []
441 + all_ratios = []
442 +
443 + for E, T, H in configs:
444 + total_tokens_config = [8 * E, 16 * E, 32 * E, 64 * E, 128 * E, 256 * E]
445 + config_x_axis.append(total_tokens_config)
446 +
447 + cuda_results = []
448 + baseline_results = []
449 + ratios = []
450 +
451 + for total_tokens in total_tokens_config:
452 + config_label = f"E={E},T={T},H={H},TT={total_tokens}"
453 + config_labels.append(config_label)
454 +
455 + # CUDA kernel results
456 + time_ms_cuda, gflops, gbps, perc = benchmark(
457 + silu_mul_fp8_quant_deep_gemm_cuda,
458 + E,
459 + T,
460 + H,
461 + total_tokens,
462 + runs=runs,
463 + num_warmups=num_warmups,
464 + gen_strategy=strategy,
465 + )
466 + cuda_results.append((time_ms_cuda, gflops, gbps, perc))
467 +
468 + # Baseline results
469 + time_ms_triton, gflops, gbps, perc = benchmark(
470 + silu_mul_fp8_quant_deep_gemm_triton,
471 + E,
472 + T,
473 + H,
474 + total_tokens,
475 + runs=runs,
476 + num_warmups=num_warmups,
477 + gen_strategy=strategy,
478 + )
479 + baseline_results.append((time_ms_triton, gflops, gbps, perc))
480 + ratios.append(time_ms_triton / time_ms_cuda)
481 +
482 + print(f"Completed: {config_label}")
483 + all_cuda_results.append(cuda_results)
484 + all_baseline_results.append(baseline_results)
485 + all_ratios.append(ratios)
486 +
487 + # Store results for combined plotting
488 + all_results.append(
489 + (
490 + strategy_descriptions[strategy],
491 + all_ratios,
492 + all_cuda_results,
493 + all_baseline_results,
494 + config_labels,
495 + config_x_axis,
496 + )
497 + )
498 +
499 + # Print summary table for this strategy
500 + print(f"\nSummary Table - {strategy_descriptions[strategy]}:")
501 + print(f"{'Config':<20} {'CUDA Time(ms)':<12} {'Base Time(ms)':<12} {'Speedup':<8}")
502 + print("-" * 60)
503 +
504 + for i, (E, T, H) in enumerate(configs):
505 + speedup = baseline_results[i][0] / cuda_results[i][0]
506 + config_label = f"E={E:3d},T={T:4d},H={H:4d}"
507 + print(
508 + f"{config_label:<20} {cuda_results[i][0]:8.5f} "
509 + f"{baseline_results[i][0]:8.5f} {speedup:6.2f}x"
510 + )
511 +
512 +
513 +def create_total_tokens_plot(all_results):
514 + num_strategies = len(all_results)
515 + num_configs = len(configs)
516 +
517 + # Create side-by-side subplots: 2 columns for speedup and bandwidth percentage
518 + fig, axs = plt.subplots(
519 + num_strategies, num_configs * 2, figsize=(28, 6 * num_strategies)
520 + )
521 +
522 + # Add main title to the entire figure
523 + fig.suptitle(
524 + "Performance Analysis: Speedup vs Bandwidth Utilization (Triton & CUDA)",
525 + fontsize=16,
526 + fontweight="bold",
527 + y=0.98,
528 + )
529 +
530 + # Handle single strategy case
531 + if num_strategies == 1:
532 + axs = axs.reshape(1, -1)
533 +
534 + # Handle single config case
535 + if num_configs == 1:
536 + axs = axs.reshape(-1, 2)
537 +
538 + for strategy_idx, result in enumerate(all_results):
539 + (
540 + strategy_name,
541 + all_ratios,
542 + all_cuda_results,
543 + all_baseline_results,
544 + config_labels,
545 + config_x_axis,
546 + ) = result
547 +
548 + for config_idx in range(num_configs):
549 + # Speedup plot (left column)
550 + ax_speedup = axs[strategy_idx, config_idx * 2]
551 + # Bandwidth plot (right column)
552 + ax_bandwidth = axs[strategy_idx, config_idx * 2 + 1]
553 +
554 + E, T, H = configs[config_idx]
555 + ratios = all_ratios[config_idx]
556 + total_tokens_values = config_x_axis[config_idx]
557 +
558 + # Extract CUDA and Triton bandwidth percentages
559 + cuda_bandwidth_percentages = [
560 + result[3] for result in all_cuda_results[config_idx]
561 + ]
562 + triton_bandwidth_percentages = [
563 + result[3] for result in all_baseline_results[config_idx]
564 + ]
565 +
566 + # Plot speedup ratios vs total tokens (left plot)
567 + ax_speedup.plot(
568 + total_tokens_values, ratios, "bo-", linewidth=3, markersize=8
569 + )
570 + ax_speedup.set_title(
571 + f"{strategy_name}\nSpeedup (CUDA/Triton)\nE={E}, T={T}, H={H}",
572 + fontsize=12,
573 + fontweight="bold",
574 + )
575 + ax_speedup.set_xlabel("Total Tokens", fontweight="bold", fontsize=11)
576 + ax_speedup.set_ylabel("Speedup Ratio", fontweight="bold", fontsize=11)
577 + ax_speedup.grid(True, alpha=0.3)
578 +
579 + ax_bandwidth.plot(
580 + total_tokens_values,
581 + cuda_bandwidth_percentages,
582 + "ro-",
583 + linewidth=3,
584 + markersize=8,
585 + label="CUDA",
586 + )
587 + ax_bandwidth.plot(
588 + total_tokens_values,
589 + triton_bandwidth_percentages,
590 + "go-",
591 + linewidth=3,
592 + markersize=8,
593 + label="Triton",
594 + )
595 + ax_bandwidth.set_title(
596 + f"{strategy_name}\nBandwidth Utilization (Hopper)\nE={E}, T={T}, H={H}",
597 + fontsize=12,
598 + fontweight="bold",
599 + )
600 + ax_bandwidth.set_xlabel("Total Tokens", fontweight="bold", fontsize=11)
601 + ax_bandwidth.set_ylabel(
602 + "% of Peak Bandwidth", fontweight="bold", fontsize=11
603 + )
604 + ax_bandwidth.legend(prop={"weight": "bold"})
605 + ax_bandwidth.grid(True, alpha=0.3)
606 +
607 + # Format x-axis labels for both plots
608 + for ax in [ax_speedup, ax_bandwidth]:
609 + ax.set_xticks(total_tokens_values)
610 + ax.set_xticklabels(
611 + [
612 + f"{tt // 1000}K" if tt >= 1000 else str(tt)
613 + for tt in total_tokens_values
614 + ],
615 + fontweight="bold",
616 + )
617 + # Make tick labels bold
618 + for label in ax.get_xticklabels() + ax.get_yticklabels():
619 + label.set_fontweight("bold")
620 +
621 + # Add value labels on speedup points
622 + for x, y in zip(total_tokens_values, ratios):
623 + ax_speedup.annotate(
624 + f"{y:.2f}x",
625 + (x, y),
626 + textcoords="offset points",
627 + xytext=(0, 12),
628 + ha="center",
629 + fontsize=10,
630 + fontweight="bold",
631 + bbox=dict(boxstyle="round,pad=0.3", facecolor="white", alpha=0.7),
632 + )
633 +
634 + # Add value labels on CUDA bandwidth points
635 + for x, y in zip(total_tokens_values, cuda_bandwidth_percentages):
636 + ax_bandwidth.annotate(
637 + f"{y:.1f}%",
638 + (x, y),
639 + textcoords="offset points",
640 + xytext=(0, 12),
641 + ha="center",
642 + fontsize=9,
643 + fontweight="bold",
644 + bbox=dict(boxstyle="round,pad=0.2", facecolor="red", alpha=0.3),
645 + )
646 +
647 + # Add value labels on Triton bandwidth points
648 + for x, y in zip(total_tokens_values, triton_bandwidth_percentages):
649 + ax_bandwidth.annotate(
650 + f"{y:.1f}%",
651 + (x, y),
652 + textcoords="offset points",
653 + xytext=(0, -15),
654 + ha="center",
655 + fontsize=9,
656 + fontweight="bold",
657 + bbox=dict(boxstyle="round,pad=0.2", facecolor="green", alpha=0.3),
658 + )
659 +
660 + plt.tight_layout()
661 + plt.subplots_adjust(top=0.93) # Make room for main title
662 + filename = "silu_benchmark_total_tokens.png"
663 + plt.savefig(filename, dpi=300, bbox_inches="tight")
664 + plt.show()
665 +
666 + return filename
667 +
668 +
669 +# Create combined plot with all strategies
670 +combined_plot_filename = create_total_tokens_plot(all_results)
671 +
672 +print(f"\n{'=' * 60}")
673 +print("Benchmark Complete!")
674 +print(f"Generated combined plot: {combined_plot_filename}")
675 +print(f"{'=' * 60}")
csrc/ops.h
+6
@@ -133,6 +133,12 @@ void silu_and_mul_nvfp4_quant(torch::Tensor& out,
133 torch::Tensor& input,
134 torch::Tensor& input_global_scale);
135 #endif
136 +void silu_mul_fp8_quant_deep_gemm_cuda(
137 + const at::Tensor& input, // (E, T, 2*H)
138 + const at::Tensor& counts, // (E)
139 + at::Tensor& y_q, // (E, T, H) [OUT]
140 + at::Tensor& y_s, // (E, T, H//group_size) [OUT]
141 + int64_t group_size, bool use_ue8m0, int64_t num_parallel_tokens);
142
143 void mul_and_silu(torch::Tensor& out, torch::Tensor& input);
144
csrc/quantization/activation_kernels.cu
+465
@@ -9,6 +9,26 @@
9
10 #include "quantization/fp8/common.cuh"
11
12 +#include <c10/util/Float8_e4m3fn.h>
13 +
14 +#ifndef USE_ROCM
15 + #include <cuda_bf16.h>
16 + #include <cuda_fp16.h>
17 + #include <cuda_fp8.h>
18 +#else
19 + #include <hip/hip_bf16.h>
20 + #include <hip/hip_fp16.h>
21 + #include <hip/hip_fp8.h>
22 +
23 +typedef __hip_bfloat162 __nv_bfloat162;
24 +typedef __hip_bfloat16 __nv_bfloat16;
25 +typedef __hip_bfloat16_raw __nv_bfloat16_raw;
26 +
27 +typedef __hip_fp8_e4m3 __nv_fp8_e4m3;
28 +typedef __hip_fp8x4_e4m3 __nv_fp8x4_e4m3;
29 +#endif
30 +
31 +#include "core/registration.h"
32 namespace vllm {
33
34 template <typename T>
@@ -87,6 +107,337 @@ __global__ void act_and_mul_quant_kernel(
107 }
108 }
109 }
110 +
111 +__device__ __forceinline__ float silu(float x) {
112 + return (__fdividef(x, (1.f + expf(-x))));
113 +}
114 +
115 +__device__ __forceinline__ float2 silu2(float2 x) {
116 + return make_float2(silu(x.x), silu(x.y));
117 +}
118 +
119 +#ifndef USE_ROCM
120 +__device__ __forceinline__ float warp_max(float v) {
121 + static constexpr unsigned FULL_MASK = 0xffffffffu;
122 + for (int offset = 1; offset < WARP_SIZE; offset *= 2) {
123 + v = fmaxf(v, __shfl_xor_sync(FULL_MASK, v, offset));
124 + }
125 + return v;
126 +}
127 +
128 +__device__ __forceinline__ __nv_bfloat16 warp_max(__nv_bfloat16 v) {
129 + static constexpr unsigned FULL_MASK = 0xffffffffu;
130 + for (int offset = 1; offset < WARP_SIZE; offset *= 2) {
131 + v = __hmax(v, __shfl_xor_sync(FULL_MASK, v, offset));
132 + }
133 + return v;
134 +}
135 +#endif
136 +
137 +template <typename T, typename U>
138 +__device__ __forceinline__ void cp_async4(T* _smem_ptr, const U* _glob_ptr) {
139 +#if __CUDACC_VER_MAJOR__ >= 11 && __CUDA_ARCH__ >= 800
140 + auto smem_ptr = reinterpret_cast<void*>(_smem_ptr);
141 + auto glob_ptr = reinterpret_cast<const void*>(_glob_ptr);
142 + const int BYTES = 16;
143 + uint32_t smem = static_cast<uint32_t>(__cvta_generic_to_shared(smem_ptr));
144 + asm volatile(
145 + "{\n"
146 + " cp.async.cg.shared.global [%0], [%1], %2;\n"
147 + "}\n" ::"r"(smem),
148 + "l"(glob_ptr), "n"(BYTES));
149 +#else
150 + _smem_ptr[0] = _glob_ptr[0];
151 +#endif
152 +}
153 +
154 +__device__ __forceinline__ void cp_async_fence() {
155 +#if __CUDACC_VER_MAJOR__ >= 11 && __CUDA_ARCH__ >= 800
156 + asm volatile("cp.async.commit_group;\n" ::);
157 +#else
158 +#endif
159 +}
160 +
161 +template <int N>
162 +__device__ __forceinline__ void cp_async_wait() {
163 +#if __CUDACC_VER_MAJOR__ >= 11 && __CUDA_ARCH__ >= 800
164 + asm volatile("cp.async.wait_group %0;\n" ::"n"(N));
165 +#else
166 +#endif
167 +}
168 +
169 +template <>
170 +__device__ __forceinline__ void cp_async_wait<0>() {
171 +#if __CUDACC_VER_MAJOR__ >= 11 && __CUDA_ARCH__ >= 800
172 + asm volatile("cp.async.wait_all;\n" ::);
173 +#else
174 +#endif
175 +}
176 +
177 +__device__ __forceinline__ float clip(float v, float mmin, float mmax) {
178 +#if __CUDACC_VER_MAJOR__ >= 11 && __CUDA_ARCH__ >= 800
179 + return fminf(mmax, fmaxf(v, mmin));
180 +#else
181 +#endif
182 +}
183 +
184 +__device__ __forceinline__ __nv_bfloat16 clip(__nv_bfloat16 v,
185 + __nv_bfloat16 mmin,
186 + __nv_bfloat16 mmax) {
187 + return __hmin(mmax, __hmax(v, mmin));
188 +}
189 +
190 +__device__ __forceinline__ __nv_bfloat162 clip(__nv_bfloat162 v,
191 + __nv_bfloat162 mmin,
192 + __nv_bfloat162 mmax) {
193 + return __hmin2(mmax, __hmax2(v, mmin));
194 +}
195 +
196 +// We use the following values for fp8 min/max:
197 +// __nv_fp8_e4m3 = (-448, +448)
198 +// __nv_fp8_e4m3uz = (-240.0, +240.0)
199 +// It is currently assumed that only
200 +template <class T>
201 +constexpr __nv_bfloat16 get_fp8_max() {
202 + static_assert(std::is_same_v<T, c10::Float8_e4m3fn> ||
203 + std::is_same_v<T, c10::Float8_e4m3fnuz>);
204 + if constexpr (std::is_same_v<T, c10::Float8_e4m3fn>) {
205 + return __nv_bfloat16(__nv_bfloat16_raw{.x = 17376});
206 + } else {
207 + return __nv_bfloat16(__nv_bfloat16_raw{.x = 17264});
208 + }
209 +}
210 +
211 +template <class T>
212 +constexpr __nv_bfloat16 get_fp8_min() {
213 + static_assert(std::is_same_v<T, c10::Float8_e4m3fn> ||
214 + std::is_same_v<T, c10::Float8_e4m3fnuz>);
215 + if constexpr (std::is_same_v<T, c10::Float8_e4m3fn>) {
216 + return __nv_bfloat16(__nv_bfloat16_raw{.x = 50144});
217 + } else {
218 + return __nv_bfloat16(__nv_bfloat16_raw{.x = 50032});
219 + }
220 +}
221 +#ifndef USE_ROCM
222 +template <typename fp8_type, int32_t NUM_WARPS, typename Idx_t,
223 + int NUM_PARALLEL_TOKENS, bool USE_UE8M0, int GROUP_SIZE = 128,
224 + int NUM_STAGES = 3>
225 +__global__ void silu_mul_fp8_quant_deep_gemm_kernel(
226 + const __nv_bfloat16* __restrict__ _input, fp8_type* __restrict__ _y_q,
227 + float* __restrict__ _y_s, const int32_t* __restrict__ counts,
228 +
229 + // sizes
230 + int H, int G,
231 +
232 + // strides (in elements)
233 + Idx_t stride_i_e, Idx_t stride_i_t, Idx_t stride_i_h, Idx_t stride_yq_e,
234 + Idx_t stride_yq_t, Idx_t stride_yq_h, Idx_t stride_ys_e, Idx_t stride_ys_t,
235 + Idx_t stride_ys_g, Idx_t stride_counts_e) {
236 + static constexpr __nv_bfloat16 fp8_min = get_fp8_min<fp8_type>();
237 + static constexpr __nv_bfloat16 fp8_max = get_fp8_max<fp8_type>();
238 + // We assign EPS with its 16-bit unsigned counterpart to allow constexpr.
239 + static constexpr __nv_bfloat16 EPS = (__nv_bfloat16_raw{.x = 11996});
240 +
241 + // We pack 8 16-bit bfloat16 values into a 128-bit __int128_t.
242 + static constexpr int32_t BFLOAT16_PER_GROUP = 8;
243 +
244 + // We split the shared memory in half, corresponding to gate and up matrices:
245 + // [...gate_i, ...up_i] where 0 <= i < stages.
246 + static constexpr int32_t S_NUM_128 =
247 + 2u * (GROUP_SIZE / BFLOAT16_PER_GROUP) * NUM_WARPS * NUM_STAGES;
248 + static constexpr auto THREAD_COUNT = NUM_WARPS * WARP_SIZE;
249 + static constexpr int HALF_THREAD_COUNT = THREAD_COUNT / 2;
250 + static constexpr int32_t S_NUM_64 = S_NUM_128 * 2;
251 + __shared__ __int128_t __align__(16) s_buff_128[S_NUM_128];
252 +
253 + const int32_t tid = threadIdx.x;
254 + const int32_t warp_id = tid / WARP_SIZE;
255 + const int32_t lane_id = tid % WARP_SIZE;
256 +
257 + auto s_buff_compute_32 = reinterpret_cast<__nv_bfloat162*>(s_buff_128);
258 +
259 + // block handles one (expert e, group g)
260 + int32_t pid = blockIdx.x;
261 + int32_t e = pid / G;
262 + int32_t g = pid % G;
263 +
264 + const int32_t n_tokens = counts[e * stride_counts_e];
265 +
266 + if (!n_tokens) {
267 + return; // Exit ASAP.
268 + }
269 +
270 + const Idx_t stride_i_t_128 = stride_i_t / 8u;
271 +
272 + int32_t n_tokens_lower, n_tokens_upper;
273 +
274 + // Each block i iterates over tokens of a slice of n_tokens =
275 + // expert_counts[i], with the size of chunk being
276 + // (n_tokens / NUM_PARALLEL_TOKENS) + residual, instead of
277 + // updiv(n_tokens, NUM_PARALLEL_TOKENS) for better scheduling.
278 + if (n_tokens < NUM_PARALLEL_TOKENS && blockIdx.y < n_tokens) {
279 + // Specialize this, but can be likely fused.
280 + if (blockIdx.y >= NUM_PARALLEL_TOKENS) {
281 + return;
282 + }
283 + n_tokens_lower = blockIdx.y;
284 + n_tokens_upper = blockIdx.y + 1;
285 + } else {
286 + auto chunk_size = n_tokens / NUM_PARALLEL_TOKENS;
287 + auto residual = n_tokens - chunk_size * NUM_PARALLEL_TOKENS;
288 + auto calc_id = [&](int32_t id) {
289 + if (id < residual) {
290 + return min(n_tokens, id * (chunk_size + 1));
291 + } else {
292 + return min(n_tokens, id * chunk_size + residual);
293 + }
294 + };
295 + n_tokens_lower = calc_id(blockIdx.y);
296 + n_tokens_upper = calc_id(blockIdx.y + 1);
297 + }
298 +
299 + if (n_tokens_lower >= n_tokens_upper) {
300 + return;
301 + }
302 +
303 + // We do calculations here, using constexpr wherever possible.
304 + const Idx_t base_i = e * stride_i_e + NUM_WARPS * g * GROUP_SIZE * stride_i_h;
305 + const Idx_t base_ys = e * stride_ys_e + NUM_WARPS * g * stride_ys_g;
306 + const Idx_t base_yq =
307 + e * stride_yq_e + NUM_WARPS * g * GROUP_SIZE * stride_yq_h;
308 + Idx_t gate_off_128 = (base_i / static_cast<Idx_t>(8u));
309 + auto input_128_ptr = reinterpret_cast<const __int128_t*>(_input);
310 + auto gate_128_ptr = input_128_ptr + gate_off_128 + (tid % HALF_THREAD_COUNT) +
311 + stride_i_t_128 * n_tokens_lower;
312 + auto up_128_ptr = gate_128_ptr + (H * stride_i_h) / 8u;
313 + auto y_s_ptr =
314 + _y_s + base_ys + warp_id * stride_ys_g + n_tokens_lower * stride_ys_t;
315 + auto y_q_ptr = _y_q + base_yq + warp_id * GROUP_SIZE +
316 + stride_yq_t * n_tokens_lower + 4 * lane_id;
317 + int32_t t_load = n_tokens_lower, load_stage_id = 0;
318 + auto s_buff_gate_load_128 = s_buff_128 + (tid % HALF_THREAD_COUNT);
319 + auto s_buff_up_load_128 = s_buff_gate_load_128 + S_NUM_128 / 2u;
320 + int32_t stage_offset{};
321 +
322 + static constexpr int32_t LOAD_STAGE_SIZE = (NUM_WARPS * WARP_SIZE / 2);
323 + static constexpr int32_t LOAD_STAGE_MOD =
324 + NUM_STAGES * (NUM_WARPS * WARP_SIZE / 2);
325 +
326 + // Two halves of all threads in a block conduct global loads for gate and up,
327 + // repsectively.
328 + auto load_and_advance_y_pred = [&] {
329 + if (t_load < n_tokens_upper) {
330 + auto s_gate_stage_128_staged_ptr = s_buff_gate_load_128 + stage_offset;
331 + auto s_up_stage_128_staged_ptr = s_buff_up_load_128 + stage_offset;
332 +
333 + // It is very important that LOAD_STAGE_SIZE is constexpr to avoid
334 + // unnecessary ALU ops.
335 + stage_offset += LOAD_STAGE_SIZE;
336 + stage_offset %= LOAD_STAGE_MOD;
337 +
338 + if (tid < HALF_THREAD_COUNT) {
339 + cp_async4(s_gate_stage_128_staged_ptr, gate_128_ptr);
340 + gate_128_ptr += stride_i_t_128;
341 + } else {
342 + cp_async4(s_up_stage_128_staged_ptr, up_128_ptr);
343 + up_128_ptr += stride_i_t_128;
344 + }
345 + ++t_load;
346 + ++load_stage_id;
347 + }
348 + // We fence even if there is nothing to load to simplify pipelining.
349 + cp_async_fence();
350 + };
351 +
352 + #pragma unroll
353 + for (int i = 0; i < NUM_STAGES - 1; i++) {
354 + load_and_advance_y_pred();
355 + }
356 +
357 + __int64_t* s_gate_ptr = reinterpret_cast<__int64_t*>(
358 + s_buff_compute_32 + warp_id * (GROUP_SIZE / 2)) +
359 + lane_id;
360 + __int64_t* s_up_ptr = s_gate_ptr + S_NUM_64 / 2;
361 +
362 + static constexpr int32_t STAGE_SIZE = (GROUP_SIZE * NUM_WARPS) / 4u;
363 + static constexpr int32_t STAGE_MOD = STAGE_SIZE * NUM_STAGES;
364 +
365 + int32_t compute_pipeline_offset_64 = 0;
366 +
367 + for (int32_t t = n_tokens_lower; t < n_tokens_upper; ++t) {
368 + __nv_bfloat16 y_max_bf16 = EPS;
369 + __nv_bfloat162 results_bf162[2];
370 +
371 + cp_async_wait<NUM_STAGES - 2>();
372 + __syncthreads();
373 +
374 + // We double-buffer pipelined loads so that the next load will
375 + // concurrently run with compute without overwrites.
376 + load_and_advance_y_pred();
377 +
378 + auto s_gate_compute_64 = s_gate_ptr + compute_pipeline_offset_64;
379 + auto s_up_compute_64 = s_up_ptr + compute_pipeline_offset_64;
380 +
381 + // STAGE_SIZE must also be constexpr!
382 + compute_pipeline_offset_64 += STAGE_SIZE;
383 + compute_pipeline_offset_64 %= STAGE_MOD;
384 +
385 + // Each thread loads (gate/up) 2X 4X bfloat16 values into registers.
386 + __int64_t gate64 = *s_gate_compute_64;
387 + __nv_bfloat162* s_gate_compute_32 =
388 + reinterpret_cast<__nv_bfloat162*>(&gate64);
389 +
390 + __int64_t up64 = *s_up_compute_64;
391 + __nv_bfloat162* s_up_compute_32 = reinterpret_cast<__nv_bfloat162*>(&up64);
392 +
393 + #pragma unroll
394 + for (int i = 0; i < 2; i++) {
395 + // For silu, we make sure that div is emitted.
396 + float2 gate = silu2(__bfloat1622float2(s_gate_compute_32[i]));
397 + results_bf162[i] = __float22bfloat162_rn(gate);
398 + }
399 +
400 + #pragma unroll
401 + for (int i = 0; i < 2; i++) {
402 + results_bf162[i] = __hmul2(results_bf162[i], s_up_compute_32[i]);
403 + }
404 +
405 + auto _y_max2 =
406 + __hmax2(__habs2(results_bf162[0]), __habs2(results_bf162[1]));
407 +
408 + y_max_bf16 = __hmax(_y_max2.x, _y_max2.y);
409 +
410 + // An entire group is assigned to a single warp, so a simple warp reduce
411 + // is used.
412 + __nv_bfloat16 y_s = warp_max(y_max_bf16) / fp8_max;
413 +
414 + if constexpr (USE_UE8M0) {
415 + y_s = hexp2(hceil(hlog2(y_s)));
416 + }
417 +
418 + auto inv_y = __float2bfloat16_rn(1.f) / y_s;
419 +
420 + auto y_s2 = make_bfloat162(inv_y, inv_y);
421 +
422 + #pragma unroll
423 + for (int32_t i = 0; i < 2; ++i) {
424 + results_bf162[i] =
425 + clip(__hmul2(results_bf162[i], y_s2), __bfloat162bfloat162(fp8_min),
426 + __bfloat162bfloat162(fp8_max));
427 + }
428 +
429 + auto fp8x4 = __nv_fp8x4_e4m3(results_bf162[0], results_bf162[1]);
430 + *reinterpret_cast<__nv_fp8x4_e4m3*>(y_q_ptr) = fp8x4;
431 + y_q_ptr += stride_yq_t;
432 +
433 + if (lane_id == 0) {
434 + *y_s_ptr = y_s;
435 + y_s_ptr += stride_ys_t;
436 + }
437 + }
438 +}
439 +#endif
440 +
441 } // namespace vllm
442
443 // Launch activation, gating, and quantize kernel.
@@ -119,3 +470,117 @@ void silu_and_mul_quant(torch::Tensor& out, // [..., d]
470 TORCH_CHECK(input.size(-1) % 2 == 0);
471 LAUNCH_ACTIVATION_GATE_KERNEL(vllm::silu_kernel);
472 }
473 +
474 +void silu_mul_fp8_quant_deep_gemm_cuda(
475 + const at::Tensor& input, // (E, T, 2*H)
476 + const at::Tensor& counts, // (E)
477 + at::Tensor& y_q, // (E, T, H) [OUT]
478 + at::Tensor& y_s, // (E, T, H//group_size) [OUT]
479 + int64_t group_size, bool use_ue8m0, int64_t num_parallel_tokens) {
480 +#ifndef USE_ROCM
481 + // This kernel relies heavily on cp.async and fp8 support.
482 + // This kernel currently only supports H % 128 == 0 and assumes a
483 + // fixed GROUP_SIZE of 128.
484 + TORCH_CHECK(input.dtype() == torch::kBFloat16);
485 + TORCH_CHECK(y_q.dtype() == torch::kFloat8_e4m3fn ||
486 + y_q.dtype() == torch::kFloat8_e4m3fnuz);
487 + TORCH_CHECK(y_s.dtype() == torch::kFloat32);
488 + TORCH_CHECK(input.size(-1) % 256 == 0);
489 +
490 + // Check that num_parallel_tokens is of power of 2 and between 1 and 64.
491 + TORCH_CHECK(1 <= num_parallel_tokens && num_parallel_tokens <= 64);
492 + TORCH_CHECK(!(num_parallel_tokens & (num_parallel_tokens - 1)));
493 +
494 + using Idx_t = int64_t;
495 +
496 + Idx_t E = input.size(0);
497 + Idx_t T = input.size(1);
498 + Idx_t H = input.size(2) / 2;
499 + Idx_t stride_i_e = input.stride(0);
500 + Idx_t stride_i_t = input.stride(1);
501 + Idx_t stride_i_h = input.stride(2);
502 + Idx_t stride_yq_e = y_q.stride(0);
503 + Idx_t stride_yq_t = y_q.stride(1);
504 + Idx_t stride_yq_h = y_q.stride(2);
505 + Idx_t stride_ys_e = y_s.stride(0);
506 + Idx_t stride_ys_t = y_s.stride(1);
507 + Idx_t stride_ys_g = y_s.stride(2);
508 +
509 + Idx_t stride_counts_e = counts.stride(0);
510 +
511 + static constexpr int GROUP_SIZE = 128;
512 +
513 + #define KERNEL_FN \
514 + if (use_ue8m0) { \
515 + vllm::silu_mul_fp8_quant_deep_gemm_kernel<fp8_t, NUM_WARPS, Idx_t, \
516 + NUM_PARALLEL_TOKENS, true> \
517 + <<<grid, block, 0, stream>>>( \
518 + reinterpret_cast<__nv_bfloat16*>(input.data_ptr()), \
519 + (fp8_t*)y_q.data_ptr(), y_s.data_ptr<float>(), \
520 + reinterpret_cast<int32_t*>(counts.data_ptr<int>()), H, G, \
521 + stride_i_e, stride_i_t, stride_i_h, stride_yq_e, stride_yq_t, \
522 + stride_yq_h, stride_ys_e, stride_ys_t, stride_ys_g, \
523 + stride_counts_e); \
524 + } else { \
525 + vllm::silu_mul_fp8_quant_deep_gemm_kernel<fp8_t, NUM_WARPS, Idx_t, \
526 + NUM_PARALLEL_TOKENS, false> \
527 + <<<grid, block, 0, stream>>>( \
528 + reinterpret_cast<__nv_bfloat16*>(input.data_ptr()), \
529 + (fp8_t*)y_q.data_ptr(), y_s.data_ptr<float>(), \
530 + reinterpret_cast<int32_t*>(counts.data_ptr<int>()), H, G, \
531 + stride_i_e, stride_i_t, stride_i_h, stride_yq_e, stride_yq_t, \
532 + stride_yq_h, stride_ys_e, stride_ys_t, stride_ys_g, \
533 + stride_counts_e); \
534 + }
535 +
536 + #define KERNEL_CALL_H \
537 + if (H % (4 * GROUP_SIZE) == 0) { \
538 + static constexpr int NUM_WARPS = 4; \
539 + populate_launch_params(NUM_WARPS, NUM_PARALLEL_TOKENS); \
540 + KERNEL_FN \
541 + } else { \
542 + static constexpr int NUM_WARPS = 1; \
543 + populate_launch_params(NUM_WARPS, NUM_PARALLEL_TOKENS); \
544 + KERNEL_FN \
545 + }
546 +
547 + #define KERNEL_CALL_TOP_LEVEL \
548 + if (num_parallel_tokens == 1) { \
549 + static constexpr int NUM_PARALLEL_TOKENS = 1; \
550 + KERNEL_CALL_H \
551 + } else if (num_parallel_tokens == 2) { \
552 + static constexpr int NUM_PARALLEL_TOKENS = 2; \
553 + KERNEL_CALL_H \
554 + } else if (num_parallel_tokens == 4) { \
555 + static constexpr int NUM_PARALLEL_TOKENS = 4; \
556 + KERNEL_CALL_H \
557 + } else if (num_parallel_tokens == 8) { \
558 + static constexpr int NUM_PARALLEL_TOKENS = 8; \
559 + KERNEL_CALL_H \
560 + } else if (num_parallel_tokens == 16) { \
561 + static constexpr int NUM_PARALLEL_TOKENS = 16; \
562 + KERNEL_CALL_H \
563 + } else if (num_parallel_tokens == 32) { \
564 + static constexpr int NUM_PARALLEL_TOKENS = 32; \
565 + KERNEL_CALL_H \
566 + } else if (num_parallel_tokens == 64) { \
567 + static constexpr int NUM_PARALLEL_TOKENS = 64; \
568 + KERNEL_CALL_H \
569 + }
570 +
571 + Idx_t G;
572 + dim3 block, grid;
573 + auto populate_launch_params = [&](int num_warps, int _num_parallel_tokens) {
574 + G = H / Idx_t(group_size * num_warps);
575 + grid = dim3(E * G, _num_parallel_tokens);
576 + block = dim3(num_warps * WARP_SIZE);
577 + };
578 +
579 + const cudaStream_t stream = at::cuda::getCurrentCUDAStream();
580 + const at::cuda::OptionalCUDAGuard device_guard(device_of(input));
581 + VLLM_DISPATCH_FP8_TYPES(y_q.scalar_type(),
582 + "silu_mul_fp8_quant_deep_gemm_kernel",
583 + [&] { KERNEL_CALL_TOP_LEVEL });
584 +
585 +#endif
586 +}
csrc/torch_bindings.cpp
+7
@@ -32,6 +32,13 @@ TORCH_LIBRARY_EXPAND(TORCH_EXTENSION_NAME, ops) {
32 #define stride_tag
33 #endif
34
35 + ops.def(
36 + "silu_mul_fp8_quant_deep_gemm_cuda(Tensor input, Tensor counts, Tensor! "
37 + "y_q, Tensor! y_s, int group_size, "
38 + "bool use_ue8m0, int num_parallel_tokens) -> ()");
39 + ops.impl("silu_mul_fp8_quant_deep_gemm_cuda", torch::kCUDA,
40 + &silu_mul_fp8_quant_deep_gemm_cuda);
41 +
42 ops.def("weak_ref_tensor(Tensor input) -> Tensor");
43 ops.impl("weak_ref_tensor", torch::kCUDA, &weak_ref_tensor);
44
tests/kernels/moe/test_silu_mul_fp8_quant_deep_gemm.py
+71 -33
@@ -5,28 +5,52 @@ import pytest
5 import torch
6
7 from vllm.model_executor.layers.fused_moe.batched_deep_gemm_moe import (
8 - silu_mul_fp8_quant_deep_gemm)
8 + silu_mul_fp8_quant_deep_gemm_cuda)
9 from vllm.platforms import current_platform
10 +from vllm.utils import cdiv
11 +
12 +fp8_dtype = torch.float8_e4m3fn
13
11 -# (E, T, H, group_size, seed)
14 CASES = [
13 - (1, 1, 128, 64, 0),
14 - (1, 4, 128, 128, 0),
15 - (2, 4, 256, 128, 0),
16 - (32, 64, 256, 128, 0),
17 - (17, 31, 768, 128, 0),
15 + (1, 1, 128, fp8_dtype),
16 + (1, 4, 128, fp8_dtype),
17 + (2, 4, 256, fp8_dtype),
18 + (32, 64, 256, fp8_dtype),
19 + (17, 31, 768, fp8_dtype),
20 + (1, 1, 128 * 1, fp8_dtype),
21 + (1, 1, 128 * 2, fp8_dtype),
22 + (1, 1, 128 * 3, fp8_dtype),
23 + (1, 1, 128 * 4, fp8_dtype),
24 + (8, 16, 128 * 1, fp8_dtype),
25 + (8, 16, 128 * 2, fp8_dtype),
26 + (8, 16, 128 * 3, fp8_dtype),
27 + (8, 16, 128 * 4, fp8_dtype),
28 + (8, 64, 7168, fp8_dtype),
29 + (8, 128, 7168, fp8_dtype),
30 + (8, 256, 7168, fp8_dtype),
31 + (8, 512, 7168, fp8_dtype),
32 + (8, 1024, 7168, fp8_dtype),
33 + (256, 8, 7168, fp8_dtype),
34 + (256, 16, 7168, fp8_dtype),
35 + (256, 32, 7168, fp8_dtype),
36 + (256, 64, 7168, fp8_dtype),
37 +
38 + # Only add a few fnuz tests to help with long CI times.
39 + (8, 512, 7168, torch.float8_e4m3fnuz),
40 + (8, 1024, 7168, torch.float8_e4m3fnuz),
41 ]
42
43
21 -@pytest.mark.parametrize("E,T,H,group_size,seed", CASES)
44 +@pytest.mark.parametrize("E,T,H,fp8_type", CASES)
45 @torch.inference_mode()
23 -def test_silu_mul_fp8_quant_deep_gemm(E, T, H, group_size, seed):
24 - current_platform.seed_everything(seed)
46 +def test_silu_mul_fp8_quant_deep_gemm(E, T, H, fp8_type):
47 + group_size = 128
48 + current_platform.seed_everything(42)
49
50 # Input tensor of shape (E, T, 2*H)
51 y = torch.randn((E, T, 2 * H), dtype=torch.bfloat16, device="cuda")
52 tokens_per_expert = torch.randint(
29 - low=0,
53 + low=T // 2,
54 high=T,
55 size=(E, ),
56 dtype=torch.int32,
@@ -34,45 +58,59 @@ def test_silu_mul_fp8_quant_deep_gemm(E, T, H, group_size, seed):
58 )
59
60 # Run the Triton kernel
37 - y_q, y_s = silu_mul_fp8_quant_deep_gemm(y,
38 - tokens_per_expert,
39 - group_size=group_size,
40 - eps=1e-10)
61 + y_q, y_s = silu_mul_fp8_quant_deep_gemm_cuda(y,
62 + tokens_per_expert,
63 + group_size=group_size)
64
42 - # Reference implementation
43 - fp8_info = torch.finfo(torch.float8_e4m3fn)
65 + torch.cuda.synchronize()
66 + fp8_info = torch.finfo(fp8_dtype)
67 fp8_max = fp8_info.max
68 fp8_min = fp8_info.min
69 eps = 1e-10
70
48 - # Compute silu activation and elementwise multiplication
49 - y1 = y[..., :H]
71 + y1 = y[..., :H].float()
72 y2 = y[..., H:]
73 silu_x = y1 * torch.sigmoid(y1)
74 merged = silu_x * y2
75
54 - # Compute reference scales and quantized output, skipping padded tokens
76 for e in range(E):
77 nt = tokens_per_expert[e].item()
57 - ref_s = torch.empty((T, H // group_size),
78 + ref_s = torch.empty((T, cdiv(H, group_size)),
79 dtype=torch.float32,
80 device="cuda")
60 - ref_q = torch.empty((T, H), dtype=torch.float8_e4m3fn, device="cuda")
81 + ref_q = torch.empty((T, H), dtype=fp8_dtype, device="cuda")
82 +
83 for t in range(nt):
62 - data = merged[e, t]
63 - data_grp = data.view(H // group_size, group_size)
64 - amax = data_grp.abs().amax(dim=1).clamp(min=eps)
65 - scale = amax / fp8_max
84 + data = merged[e, t].float()
85 + ref_q_row = torch.empty_like(data)
86 +
87 + # process full groups
88 + n_full_groups = H // group_size
89 + if n_full_groups > 0:
90 + data_grp = data[:n_full_groups * group_size].view(
91 + n_full_groups, group_size)
92 + amax = data_grp.abs().amax(dim=1).clamp(min=eps)
93 + scale = amax / fp8_max
94 + scaled = data[:n_full_groups *
95 + group_size] / scale.repeat_interleave(group_size)
96 + ref_q_row[:n_full_groups * group_size] = scaled.clamp(
97 + fp8_min, fp8_max).to(fp8_dtype)
98 + ref_s[t, :n_full_groups] = scale
99
67 - scaled = data / scale.repeat_interleave(group_size)
68 - clamped = scaled.clamp(fp8_min, fp8_max)
69 - q = clamped.to(torch.float8_e4m3fn)
100 + # process remainder group
101 + rem = H % group_size
102 + if rem > 0:
103 + data_rem = data[-rem:]
104 + amax = data_rem.abs().amax().clamp(min=eps)
105 + scale = amax / fp8_max
106 + scaled = data_rem / scale
107 + ref_q_row[-rem:] = scaled.clamp(fp8_min, fp8_max).to(fp8_dtype)
108 + ref_s[t, -1] = scale
109
71 - ref_s[t] = scale
72 - ref_q[t] = q
110 + ref_q[t] = ref_q_row
111
74 - y_se = y_s[e]
75 - y_qe = y_q[e]
112 + y_se = y_s[e].float()
113 + y_qe = y_q[e].float()
114
115 torch.testing.assert_close(y_se[:nt], ref_s[:nt], atol=1e-4, rtol=1e-2)
116 torch.testing.assert_close(
vllm/model_executor/layers/fused_moe/batched_deep_gemm_moe.py
+87 -60
@@ -1,5 +1,6 @@
1 # SPDX-License-Identifier: Apache-2.0
2 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 +from math import log2
4 from typing import Optional
5
6 import torch
@@ -10,6 +11,7 @@ from vllm.model_executor.layers.fused_moe.config import FusedMoEQuantConfig
11 from vllm.model_executor.layers.fused_moe.topk_weight_and_reduce import (
12 TopKWeightAndReduceDelegate)
13 from vllm.model_executor.layers.fused_moe.utils import _resize_cache
14 +from vllm.platforms import current_platform
15 from vllm.triton_utils import tl, triton
16 from vllm.utils.deep_gemm import (fp8_m_grouped_gemm_nt_masked,
17 is_deep_gemm_e8m0_used)
@@ -24,35 +26,28 @@ def _silu_mul_fp8_quant_deep_gemm(
26 y_q_ptr, # fp8 quantized activations (E, T, H)
27 y_s_ptr, # 16-bit scales (E, T, G)
28 counts_ptr, # int32 num tokens per expert (E)
27 -
29 # Sizes ---------------------------------------------------------------
30 H: tl.constexpr, # hidden dimension (per output)
31 GROUP_SIZE: tl.constexpr, # elements per group (usually 128)
31 -
32 # Strides for input (elements) ---------------------------------------
33 stride_i_e,
34 stride_i_t,
35 stride_i_h,
36 -
36 # Strides for y_q (elements) -----------------------------------------
37 stride_yq_e,
38 stride_yq_t,
39 stride_yq_h,
41 -
40 # Strides for y_s (elements) -----------------------------------------
41 stride_ys_e,
42 stride_ys_t,
43 stride_ys_g,
46 -
44 # Stride for counts (elements)
45 stride_counts_e,
49 -
46 # Numeric params ------------------------------------------------------
47 eps: tl.constexpr,
48 fp8_min: tl.constexpr,
49 fp8_max: tl.constexpr,
50 use_ue8m0: tl.constexpr,
55 -
51 # Meta ---------------------------------------------------------------
52 BLOCK: tl.constexpr,
53 NUM_STAGES: tl.constexpr,
@@ -101,17 +96,15 @@ def _silu_mul_fp8_quant_deep_gemm(
96 tl.store(y_s_ptr + base_ys_offset + t * stride_ys_t, y_s)
97
98
104 -def silu_mul_fp8_quant_deep_gemm(
99 +def silu_mul_fp8_quant_deep_gemm_cuda(
100 y: torch.Tensor, # (E, T, 2*H)
101 tokens_per_expert: torch.Tensor, # (E,) number of valid tokens per expert
102 + num_parallel_tokens=16,
103 group_size: int = 128,
108 - eps: float = 1e-10,
104 ) -> tuple[torch.Tensor, torch.Tensor]:
105 """Quantize silu(y[..., :H]) * y[..., H:] to FP8 with group per-token scales
111 -
112 - y has shape (E, T, 2*H). The first half of the last dimension is
106 + y has shape (E, T, 2*H). The first half of the last dimension is
107 silu-activated, multiplied by the second half, then quantized into FP8.
114 -
108 Returns `(y_q, y_s)` where
109 * `y_q`: FP8 tensor, shape (E, T, H), same layout as y[..., :H]
110 * `y_s`: FP32 tensor, shape (E, T, H // group_size), strides (T*G, 1, T)
@@ -120,22 +113,17 @@ def silu_mul_fp8_quant_deep_gemm(
113 E, T, H2 = y.shape
114 assert H2 % 2 == 0, "last dim of y must be even (2*H)"
115 H = H2 // 2
123 - G = H // group_size
124 - assert H % group_size == 0, "H must be divisible by group_size"
125 - assert tokens_per_expert.ndim == 1 and tokens_per_expert.shape[0] == E, \
126 - "tokens_per_expert must be shape (E,)"
116 + G = (H + group_size - 1) // group_size
117 + assert H % 8 == 0, "H must be divisible by 8"
118 + assert group_size == 128, "H must be divisible by 8"
119 + assert tokens_per_expert.ndim == 1 and tokens_per_expert.shape[0] == E
120 +
121 tokens_per_expert = tokens_per_expert.to(device=y.device,
122 dtype=torch.int32)
123
130 - # allocate outputs
124 fp8_dtype = torch.float8_e4m3fn
125 y_q = torch.empty((E, T, H), dtype=fp8_dtype, device=y.device)
126
134 - # strides (elements)
135 - stride_i_e, stride_i_t, stride_i_h = y.stride()
136 - stride_yq_e, stride_yq_t, stride_yq_h = y_q.stride()
137 -
138 - # desired scale strides (elements): (T*G, 1, T)
127 stride_ys_e = T * G
128 stride_ys_t = 1
129 stride_ys_g = T
@@ -144,47 +132,86 @@ def silu_mul_fp8_quant_deep_gemm(
132 dtype=torch.float32,
133 device=y.device)
134
147 - stride_cnt_e = tokens_per_expert.stride()[0]
148 -
149 - # Static grid over experts and H-groups.
150 - # A loop inside the kernel handles the token dim
151 - grid = (E * G, )
152 -
153 - f_info = torch.finfo(fp8_dtype)
154 - fp8_max = f_info.max
155 - fp8_min = f_info.min
156 -
157 - _silu_mul_fp8_quant_deep_gemm[grid](
158 - y,
159 - y_q,
160 - y_s,
161 - tokens_per_expert,
162 - H,
163 - group_size,
164 - stride_i_e,
165 - stride_i_t,
166 - stride_i_h,
167 - stride_yq_e,
168 - stride_yq_t,
169 - stride_yq_h,
170 - stride_ys_e,
171 - stride_ys_t,
172 - stride_ys_g,
173 - stride_cnt_e,
174 - eps,
175 - fp8_min,
176 - fp8_max,
177 - is_deep_gemm_e8m0_used(),
178 - BLOCK=group_size,
179 - NUM_STAGES=4,
180 - num_warps=1,
181 - )
135 + use_ue8m0 = is_deep_gemm_e8m0_used()
136 +
137 + if E <= 16:
138 + max_empirical_parallelism = 64
139 + elif E <= 32:
140 + max_empirical_parallelism = 16
141 + else:
142 + max_empirical_parallelism = 4
143 +
144 + # We never want to launch more than Tx number of threads
145 + # This computes the clip.
146 + num_parallel_tokens = max(
147 + 1,
148 + min(max_empirical_parallelism, 2**int(log2(min(num_parallel_tokens,
149 + T)))))
150 + cuda_arch = current_platform.get_device_capability(
151 + device_id=y.device.index).to_int()
152 +
153 + if cuda_arch >= 80:
154 + torch.ops._C.silu_mul_fp8_quant_deep_gemm_cuda(y, tokens_per_expert,
155 + y_q, y_s, group_size,
156 + use_ue8m0,
157 + num_parallel_tokens)
158 + else:
159 + # Default to triton if not on cuda or if arch is too old
160 + y_q = torch.empty((E, T, H), dtype=fp8_dtype, device=y.device)
161 +
162 + stride_cnt_e = tokens_per_expert.stride()[0]
163 +
164 + # Static grid over experts and H-groups.
165 + # A loop inside the kernel handles the token dim
166 + grid = (E * G, )
167 + # strides (elements)
168 + stride_i_e, stride_i_t, stride_i_h = y.stride()
169 + stride_yq_e, stride_yq_t, stride_yq_h = y_q.stride()
170 +
171 + # desired scale strides (elements): (T*G, 1, T)
172 + stride_ys_e = T * G
173 + stride_ys_t = 1
174 + stride_ys_g = T
175 + y_s = torch.empty_strided(
176 + (E, T, G),
177 + (stride_ys_e, stride_ys_t, stride_ys_g),
178 + dtype=torch.float32,
179 + device=y.device,
180 + )
181 + f_info = torch.finfo(fp8_dtype)
182 + fp8_max = f_info.max
183 + fp8_min = f_info.min
184 + eps: float = 1e-10
185 + _silu_mul_fp8_quant_deep_gemm[grid](
186 + y,
187 + y_q,
188 + y_s,
189 + tokens_per_expert,
190 + H,
191 + group_size,
192 + stride_i_e,
193 + stride_i_t,
194 + stride_i_h,
195 + stride_yq_e,
196 + stride_yq_t,
197 + stride_yq_h,
198 + stride_ys_e,
199 + stride_ys_t,
200 + stride_ys_g,
201 + stride_cnt_e,
202 + eps,
203 + fp8_min,
204 + fp8_max,
205 + is_deep_gemm_e8m0_used(),
206 + BLOCK=group_size,
207 + NUM_STAGES=4,
208 + num_warps=1,
209 + )
210
211 return y_q, y_s
212
213
214 class BatchedDeepGemmExperts(mk.FusedMoEPermuteExpertsUnpermute):
187 -
215 # The Deep Gemm kernels only support block size of 128
216 DEEPGEMM_BLOCK_SHAPE: list[int] = [128, 128]
217
@@ -297,8 +324,8 @@ class BatchedDeepGemmExperts(mk.FusedMoEPermuteExpertsUnpermute):
324 fp8_m_grouped_gemm_nt_masked((a1q, a1q_scale), (w1, w1_scale),
325 workspace1, expert_num_tokens, expected_m)
326
300 - a2q, a2q_scale = silu_mul_fp8_quant_deep_gemm(workspace1,
301 - expert_num_tokens)
327 + a2q, a2q_scale = silu_mul_fp8_quant_deep_gemm_cuda(
328 + workspace1, expert_num_tokens)
329
330 fp8_m_grouped_gemm_nt_masked((a2q, a2q_scale), (w2, w2_scale), output,
331 expert_num_tokens, expected_m)