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
[Perf] Small optimizations for silu_mul_fp8_quant_deep_gemm (#23265)
Signed-off-by: mgoin <mgoin64@gmail.com>
Michael Goin committed
Aug 21, 2025 at 17:56 UTC
3bbe11cc136373bd4f6c12912dc094dba086fa11
3 files changed
+107
-32
benchmarks/kernels/benchmark_silu_mul_fp8_quant.py
new
+77
@@ -0,0 +1,77 @@
1
+#!/usr/bin/env python3
2
+# SPDX-License-Identifier: Apache-2.0
3
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
4
+import time
5
+
6
+import torch
7
+
8
+from vllm.model_executor.layers.fused_moe.batched_deep_gemm_moe import (
9
+ silu_mul_fp8_quant_deep_gemm,
10
+)
11
+from vllm.platforms import current_platform
12
+
13
+
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"
19
+ )
20
+
21
+ # Warmup
22
+ for _ in range(10):
23
+ silu_mul_fp8_quant_deep_gemm(y, tokens_per_expert, group_size=G)
24
+ torch.cuda.synchronize()
25
+
26
+ # Benchmark
27
+ torch.cuda.synchronize()
28
+ start = time.perf_counter()
29
+ for _ in range(runs):
30
+ silu_mul_fp8_quant_deep_gemm(y, tokens_per_expert, group_size=G)
31
+ torch.cuda.synchronize()
32
+
33
+ avg_time = (time.perf_counter() - start) / runs * 1000
34
+
35
+ # Calculate actual work done (only count valid tokens)
36
+ actual_tokens = tokens_per_expert.sum().item()
37
+ actual_elements = actual_tokens * H
38
+
39
+ # GFLOPS: operations per element = exp + 3 muls + 1 div + quantization ops ≈ 8 ops
40
+ ops_per_element = 8
41
+ total_ops = actual_elements * ops_per_element
42
+ gflops = total_ops / (avg_time / 1000) / 1e9
43
+
44
+ # Memory bandwidth: bfloat16 inputs (2 bytes), fp8 output (1 byte), scales (4 bytes)
45
+ input_bytes = actual_tokens * 2 * H * 2 # 2*H bfloat16 inputs
46
+ output_bytes = actual_tokens * H * 1 # H fp8 outputs
47
+ scale_bytes = actual_tokens * (H // G) * 4 # scales in float32
48
+ total_bytes = input_bytes + output_bytes + scale_bytes
49
+ memory_bw = total_bytes / (avg_time / 1000) / 1e9
50
+
51
+ return avg_time, gflops, memory_bw
52
+
53
+
54
+configs = [
55
+ (8, 32, 1024),
56
+ (16, 64, 2048),
57
+ (32, 128, 4096),
58
+ # 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),
65
+ (256, 1024, 7168),
66
+]
67
+
68
+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")
tests/kernels/moe/test_silu_mul_fp8_quant_deep_gemm.py
+2
-2
@@ -24,7 +24,7 @@ def test_silu_mul_fp8_quant_deep_gemm(E, T, H, group_size, seed):
24
current_platform.seed_everything(seed)
25
26
# Input tensor of shape (E, T, 2*H)
27
- y = torch.randn((E, T, 2 * H), dtype=torch.float32, device="cuda")
27
+ y = torch.randn((E, T, 2 * H), dtype=torch.bfloat16, device="cuda")
28
tokens_per_expert = torch.randint(
29
low=0,
30
high=T,
@@ -74,7 +74,7 @@ def test_silu_mul_fp8_quant_deep_gemm(E, T, H, group_size, seed):
74
y_se = y_s[e]
75
y_qe = y_q[e]
76
77
- torch.testing.assert_close(y_se[:nt], ref_s[:nt])
77
+ torch.testing.assert_close(y_se[:nt], ref_s[:nt], atol=1e-4, rtol=1e-2)
78
torch.testing.assert_close(
79
y_qe[:nt].to(torch.float32),
80
ref_q[:nt].to(torch.float32),
vllm/model_executor/layers/fused_moe/batched_deep_gemm_moe.py
+28
-30
@@ -70,53 +70,51 @@ def _silu_mul_fp8_quant_deep_gemm(
70
# number of valid tokens for this expert
71
n_tokens = tl.load(counts_ptr + e * stride_counts_e).to(tl.int64)
72
73
- cols = tl.arange(0, BLOCK)
74
- cols = cols.to(tl.int64)
75
- mask_h = cols < BLOCK
73
+ cols = tl.arange(0, BLOCK).to(tl.int64)
74
+ mask = cols < BLOCK
75
+
76
+ base_input_offset = e * stride_i_e + g * GROUP_SIZE * stride_i_h
77
+ base_gate_offset = base_input_offset + cols * stride_i_h
78
+ base_up_offset = base_input_offset + H * stride_i_h + cols * stride_i_h
79
+ base_yq_offset = (e * stride_yq_e + g * GROUP_SIZE * stride_yq_h +
80
+ cols * stride_yq_h)
81
+ base_ys_offset = e * stride_ys_e + g * stride_ys_g
82
83
for t in tl.range(0, n_tokens, num_stages=NUM_STAGES):
78
- base_i_offset = (e * stride_i_e + t * stride_i_t +
79
- g * GROUP_SIZE * stride_i_h)
80
- base_yq_offset = (e * stride_yq_e + t * stride_yq_t +
81
- g * GROUP_SIZE * stride_yq_h)
82
- base_ys_offset = e * stride_ys_e + t * stride_ys_t + g * stride_ys_g
83
-
84
- mask = mask_h
85
- x = tl.load(input_ptr + base_i_offset + cols * stride_i_h,
86
- mask=mask,
87
- other=0.0).to(tl.float32)
88
- y2 = tl.load(input_ptr + base_i_offset + H * stride_i_h +
89
- cols * stride_i_h,
84
+ gate = tl.load(input_ptr + base_gate_offset + t * stride_i_t,
85
+ mask=mask,
86
+ other=0.0).to(tl.float32)
87
+ up = tl.load(input_ptr + base_up_offset + t * stride_i_t,
88
mask=mask,
91
- other=0.0).to(tl.float32)
89
+ other=0.0)
90
+
91
+ gate = gate * (1.0 / (1.0 + tl.exp(-gate)))
92
+ y = gate * up
93
93
- x = x * (1.0 / (1.0 + tl.exp(-x)))
94
- y = x * y2
94
+ y_s = tl.maximum(tl.max(tl.abs(y)), eps) / fp8_max
95
+ if use_ue8m0:
96
+ y_s = tl.exp2(tl.ceil(tl.log2(y_s)))
97
96
- _absmax = tl.maximum(tl.max(tl.abs(y)), eps)
97
- scale_raw = _absmax / fp8_max
98
- y_s = tl.math.exp2(tl.ceil(
99
- tl.log2(scale_raw))) if use_ue8m0 else scale_raw
98
y_q = tl.clamp(y / y_s, fp8_min, fp8_max).to(y_q_ptr.dtype.element_ty)
99
102
- tl.store(y_q_ptr + base_yq_offset + cols * stride_yq_h, y_q, mask=mask)
103
- tl.store(y_s_ptr + base_ys_offset, y_s)
100
+ tl.store(y_q_ptr + base_yq_offset + t * stride_yq_t, y_q, mask=mask)
101
+ tl.store(y_s_ptr + base_ys_offset + t * stride_ys_t, y_s)
102
103
104
def silu_mul_fp8_quant_deep_gemm(
107
- y: torch.Tensor, # (E, T, 2*H) float32
105
+ y: torch.Tensor, # (E, T, 2*H)
106
tokens_per_expert: torch.Tensor, # (E,) number of valid tokens per expert
107
group_size: int = 128,
108
eps: float = 1e-10,
111
-):
109
+) -> tuple[torch.Tensor, torch.Tensor]:
110
"""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
113
silu-activated, multiplied by the second half, then quantized into FP8.
114
115
Returns `(y_q, y_s)` where
118
- * `y_q` is the FP8 tensor of shape `(E, T, H)`, same layout as `y[..., :H]`.
119
- * `y_s` has shape `(E, T, H // group_size)` and strides `(T*G, 1, T)`
116
+ * `y_q`: FP8 tensor, shape (E, T, H), same layout as y[..., :H]
117
+ * `y_s`: FP32 tensor, shape (E, T, H // group_size), strides (T*G, 1, T)
118
"""
119
assert y.ndim == 3, "y must be (E, T, 2*H)"
120
E, T, H2 = y.shape
@@ -148,7 +146,7 @@ def silu_mul_fp8_quant_deep_gemm(
146
147
stride_cnt_e = tokens_per_expert.stride()[0]
148
151
- # static grid over experts and H-groups.
149
+ # Static grid over experts and H-groups.
150
# A loop inside the kernel handles the token dim
151
grid = (E * G, )
152
@@ -178,7 +176,7 @@ def silu_mul_fp8_quant_deep_gemm(
176
fp8_max,
177
is_blackwell_deep_gemm_e8m0_used(),
178
BLOCK=group_size,
181
- NUM_STAGES=8,
179
+ NUM_STAGES=4,
180
num_warps=1,
181
)
182