@tduncan / Nvidia-Optimization-Stack-VLLM Mirror · read-only / commits / a6c4b87fbc
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 15, 04:45 UTC

Revert "[Feature] Integrate new deepgemm (#19820)" (#20049)

Signed-off-by: yewentao256 <zhyanwentao@126.com>

Wentao Ye committed Jun 24, 2025 at 22:45 UTC a6c4b87fbcca30135ba52b9199c43aa8ee9c0400
8 files changed +257 -233
benchmarks/kernels/benchmark_moe.py
-3
@@ -86,9 +86,6 @@ def benchmark_config(
86 (num_experts, 2 * shard_intermediate_size), dtype=torch.float32
87 )
88 w2_scale = torch.randn((hidden_size, num_experts), dtype=torch.float32)
89 - if use_deep_gemm:
90 - # we use the default block shape for deepgemm
91 - block_quant_shape = [128, 128]
89 if use_fp8_w8a8:
90 if block_quant_shape:
91 block_n, block_k = block_quant_shape[0], block_quant_shape[1]
benchmarks/kernels/deepgemm/benchmark_fp8_block_dense_gemm.py
+174 -133
@@ -1,11 +1,13 @@
1 # SPDX-License-Identifier: Apache-2.0
2 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 +# fmt: off
4 +# ruff: noqa: E501
5 import time
6
7 +# Import DeepGEMM functions
8 +import deep_gemm
9 import torch
6 -from deep_gemm import fp8_gemm_nt
7 -from deep_gemm.testing.numeric import calc_diff
8 -from deep_gemm.utils.math import ceil_div, per_block_cast_to_fp8, per_token_cast_to_fp8
10 +from deep_gemm import calc_diff, ceil_div, get_col_major_tma_aligned_tensor
11
12 # Import vLLM functions
13 from vllm import _custom_ops as ops
@@ -16,84 +18,96 @@ from vllm.model_executor.layers.quantization.utils.fp8_utils import (
18 from vllm.triton_utils import triton
19
20
21 +# Copied from
22 +# https://github.com/deepseek-ai/DeepGEMM/blob/78cacf70d41d15d688bd493ebc85845f7f2a3d5d/tests/test_core.py#L9
23 +def per_token_cast_to_fp8(
24 + x: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor]:
25 + """Convert tensor to FP8 format with per-token scaling."""
26 + assert x.dim() == 2 and x.size(1) % 128 == 0
27 + m, n = x.shape
28 + x_view = x.view(m, -1, 128)
29 + x_amax = x_view.abs().float().amax(dim=2).view(m, -1).clamp(1e-4)
30 + return (x_view * (448.0 / x_amax.unsqueeze(2))).to(
31 + torch.float8_e4m3fn).view(m, n), (x_amax / 448.0).view(m, -1)
32 +
33 +
34 # Copied from
35 # https://github.com/deepseek-ai/DeepGEMM/blob/78cacf70d41d15d688bd493ebc85845f7f2a3d5d/tests/test_core.py#L17
21 -def per_block_cast_to_fp8_vllm(x: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor]:
36 +def per_block_cast_to_fp8(
37 + x: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor]:
38 """Convert tensor to FP8 format with per-block scaling."""
39 assert x.dim() == 2
40 m, n = x.shape
25 - x_padded = torch.zeros(
26 - (ceil_div(m, 128) * 128, ceil_div(n, 128) * 128), dtype=x.dtype, device=x.device
27 - )
41 + x_padded = torch.zeros((ceil_div(m, 128) * 128, ceil_div(n, 128) * 128),
42 + dtype=x.dtype,
43 + device=x.device)
44 x_padded[:m, :n] = x
45 x_view = x_padded.view(-1, 128, x_padded.size(1) // 128, 128)
46 x_amax = x_view.abs().float().amax(dim=(1, 3), keepdim=True).clamp(1e-4)
47 x_scaled = (x_view * (448.0 / x_amax)).to(torch.float8_e4m3fn)
32 - return x_scaled.view_as(x_padded)[:m, :n].contiguous(), (x_amax / 448.0).view(
33 - x_view.size(0), x_view.size(2)
34 - )
35 -
36 -
37 -def benchmark_shape(
38 - m: int,
39 - n: int,
40 - k: int,
41 - warmup: int = 100,
42 - repeat: int = 10000,
43 - verbose: bool = False,
44 -) -> dict:
48 + return x_scaled.view_as(x_padded)[:m, :n].contiguous(), (
49 + x_amax / 448.0).view(x_view.size(0), x_view.size(2))
50 +
51 +
52 +def benchmark_shape(m: int,
53 + n: int,
54 + k: int,
55 + warmup: int = 100,
56 + repeat: int = 10000,
57 + verbose: bool = False) -> dict:
58 """Benchmark all implementations for a specific (m, n, k) shape."""
59 if verbose:
60 print(f"\n=== Benchmarking shape: m={m}, n={n}, k={k} ===")
61
49 - A = torch.randn((m, k), device="cuda", dtype=torch.bfloat16)
50 - B = torch.randn((n, k), device="cuda", dtype=torch.bfloat16)
62 + # Create test tensors
63 + A = torch.randn((m, k), device='cuda', dtype=torch.bfloat16)
64 + B = torch.randn((n, k), device='cuda', dtype=torch.bfloat16)
65 +
66 + # Reference result in BF16
67 torch.cuda.synchronize()
68 C_ref = A @ B.t()
69
70 # Pre-quantize B for all implementations
71 # (weights can be pre-quantized offline)
72 B_deepgemm, B_scale_deepgemm = per_block_cast_to_fp8(B)
57 - B_vllm, B_scale_vllm = per_block_cast_to_fp8_vllm(B)
73 + B_vllm, B_scale_vllm = per_block_cast_to_fp8(B)
74
75 # Block size configuration
76 block_size = [128, 128]
77
78 # Pre-quantize A for all implementations
79 A_deepgemm, A_scale_deepgemm = per_token_cast_to_fp8(A)
64 - C_deepgemm = (
65 - torch.empty((n, m), device="cuda", dtype=torch.bfloat16).t().contiguous()
66 - )
80 + A_scale_deepgemm = get_col_major_tma_aligned_tensor(A_scale_deepgemm)
81 + C_deepgemm = torch.empty((m, n), device='cuda', dtype=torch.bfloat16)
82 A_vllm, A_scale_vllm = per_token_group_quant_fp8(A, block_size[1])
83 A_vllm_cutlass, A_scale_vllm_cutlass = per_token_group_quant_fp8(
69 - A, block_size[1], column_major_scales=True
70 - )
84 + A, block_size[1], column_major_scales=True)
85
86 + # === DeepGEMM Implementation ===
87 def deepgemm_gemm():
73 - fp8_gemm_nt(
74 - (A_deepgemm, A_scale_deepgemm), (B_deepgemm, B_scale_deepgemm), C_deepgemm
75 - )
88 + deep_gemm.gemm_fp8_fp8_bf16_nt((A_deepgemm, A_scale_deepgemm),
89 + (B_deepgemm, B_scale_deepgemm),
90 + C_deepgemm)
91 return C_deepgemm
92
93 + # === vLLM Triton Implementation ===
94 def vllm_triton_gemm():
79 - return w8a8_block_fp8_matmul(
80 - A_vllm,
81 - B_vllm,
82 - A_scale_vllm,
83 - B_scale_vllm,
84 - block_size,
85 - output_dtype=torch.bfloat16,
86 - )
87 -
95 + return w8a8_block_fp8_matmul(A_vllm,
96 + B_vllm,
97 + A_scale_vllm,
98 + B_scale_vllm,
99 + block_size,
100 + output_dtype=torch.bfloat16)
101 +
102 + # === vLLM CUTLASS Implementation ===
103 def vllm_cutlass_gemm():
89 - return ops.cutlass_scaled_mm(
90 - A_vllm_cutlass,
91 - B_vllm.T,
92 - scale_a=A_scale_vllm_cutlass,
93 - scale_b=B_scale_vllm.T,
94 - out_dtype=torch.bfloat16,
95 - )
104 + return ops.cutlass_scaled_mm(A_vllm_cutlass,
105 + B_vllm.T,
106 + scale_a=A_scale_vllm_cutlass,
107 + scale_b=B_scale_vllm.T,
108 + out_dtype=torch.bfloat16)
109
110 + # Run correctness check first
111 if verbose:
112 print("Running correctness check...")
113 C_deepgemm = deepgemm_gemm()
@@ -108,22 +122,26 @@ def benchmark_shape(
122 print(f"DeepGEMM vs Reference difference: {deepgemm_diff:.6f}")
123 print(f"vLLM Triton vs Reference difference: {vllm_triton_diff:.6f}")
124 print(f"vLLM CUTLASS vs Reference difference: {vllm_cutlass_diff:.6f}")
111 - print(
112 - "vLLM Triton vs DeepGEMM difference: "
113 - f"{calc_diff(C_vllm_triton, C_deepgemm):.6f}"
114 - )
115 - print(
116 - "vLLM CUTLASS vs DeepGEMM difference: "
117 - f"{calc_diff(C_vllm_cutlass, C_deepgemm):.6f}"
118 - )
125 + print("vLLM Triton vs DeepGEMM difference: "
126 + f"{calc_diff(C_vllm_triton, C_deepgemm):.6f}")
127 + print("vLLM CUTLASS vs DeepGEMM difference: "
128 + f"{calc_diff(C_vllm_cutlass, C_deepgemm):.6f}")
129
130 + # Benchmark implementations
131 implementations = {
132 "DeepGEMM": deepgemm_gemm,
133 "vLLM Triton": vllm_triton_gemm,
123 - "vLLM CUTLASS": vllm_cutlass_gemm,
134 + "vLLM CUTLASS": vllm_cutlass_gemm
135 }
136
126 - benchmark_results = {"shape": {"m": m, "n": n, "k": k}, "implementations": {}}
137 + benchmark_results = {
138 + "shape": {
139 + "m": m,
140 + "n": n,
141 + "k": k
142 + },
143 + "implementations": {}
144 + }
145
146 for name, func in implementations.items():
147 # Warmup
@@ -151,36 +169,38 @@ def benchmark_shape(
169 "tflops": tflops,
170 "gb_s": gb_s,
171 "diff": {
154 - "DeepGEMM": 0.0
155 - if name == "DeepGEMM"
156 - else calc_diff(func(), C_deepgemm),
157 - "Reference": deepgemm_diff
158 - if name == "DeepGEMM"
159 - else (vllm_triton_diff if name == "vLLM Triton" else vllm_cutlass_diff),
160 - },
172 + "DeepGEMM":
173 + 0.0 if name == "DeepGEMM" else calc_diff(func(), C_deepgemm),
174 + "Reference":
175 + deepgemm_diff if name == "DeepGEMM" else
176 + (vllm_triton_diff
177 + if name == "vLLM Triton" else vllm_cutlass_diff)
178 + }
179 }
180
181 if verbose:
164 - print(f"{name}: {avg_time_ms:.3f} ms, {tflops:.2f} TFLOPS, {gb_s:.2f} GB/s")
182 + print(
183 + f"{name}: {avg_time_ms:.3f} ms, {tflops:.2f} TFLOPS, {gb_s:.2f} GB/s"
184 + )
185
186 # Calculate speedups
187 baseline = benchmark_results["implementations"]["DeepGEMM"]["time_ms"]
188 for name, data in benchmark_results["implementations"].items():
189 if name != "DeepGEMM":
190 speedup = baseline / data["time_ms"]
171 - benchmark_results["implementations"][name]["speedup_vs_deepgemm"] = speedup
191 + benchmark_results["implementations"][name][
192 + "speedup_vs_deepgemm"] = speedup
193 if verbose:
173 - print(
174 - f"DeepGEMM is {1 / speedup:.2f}x "
175 - f"{'faster' if 1 / speedup > 1 else 'slower'} than {name}"
176 - )
194 + print(f"DeepGEMM is {1/speedup:.2f}x "
195 + f"{'faster' if 1/speedup > 1 else 'slower'} than {name}")
196
178 - vllm_triton_time = benchmark_results["implementations"]["vLLM Triton"]["time_ms"]
179 - vllm_cutlass_time = benchmark_results["implementations"]["vLLM CUTLASS"]["time_ms"]
197 + vllm_triton_time = benchmark_results["implementations"]["vLLM Triton"][
198 + "time_ms"]
199 + vllm_cutlass_time = benchmark_results["implementations"]["vLLM CUTLASS"][
200 + "time_ms"]
201 cutlass_vs_triton = vllm_triton_time / vllm_cutlass_time
181 - benchmark_results["implementations"]["vLLM CUTLASS"]["speedup_vs_triton"] = (
182 - cutlass_vs_triton
183 - )
202 + benchmark_results["implementations"]["vLLM CUTLASS"][
203 + "speedup_vs_triton"] = cutlass_vs_triton
204 if verbose:
205 print(
206 f"vLLM CUTLASS is {cutlass_vs_triton:.2f}x "
@@ -192,7 +212,8 @@ def benchmark_shape(
212
213 def format_table_row(values, widths):
214 """Format a row with specified column widths."""
195 - return "| " + " | ".join(f"{val:{w}}" for val, w in zip(values, widths)) + " |"
215 + return "| " + " | ".join(f"{val:{w}}"
216 + for val, w in zip(values, widths)) + " |"
217
218
219 def print_table(headers, rows, title=None):
@@ -200,12 +221,16 @@ def print_table(headers, rows, title=None):
221 if title:
222 print(f"\n{title}")
223
224 + # Calculate column widths based on headers and data
225 widths = [
226 max(len(str(h)), max(len(str(row[i])) for row in rows))
227 for i, h in enumerate(headers)
228 ]
229
230 + # Create separator line
231 separator = "+-" + "-+-".join("-" * w for w in widths) + "-+"
232 +
233 + # Print table
234 print(separator)
235 print(format_table_row(headers, widths))
236 print(separator)
@@ -223,22 +248,44 @@ def run_benchmarks(verbose: bool = False):
248 """Run benchmarks for a set of common shapes."""
249 print("===== STARTING FP8 GEMM BENCHMARK =====")
250
251 + # Make sure we're using the GPU
252 if not torch.cuda.is_available():
253 print("CUDA not available! Tests require GPU.")
254 return
255
256 + # Print system information
257 print(f"PyTorch version: {torch.__version__}")
258 print(f"CUDA version: {torch.version.cuda}")
259 print(f"Triton version: {triton.__version__}")
260 print(f"Using device: {torch.cuda.get_device_name()}")
261
262 + # Enable TF32 for better performance
263 torch.backends.cuda.matmul.allow_tf32 = True
264 torch.backends.cudnn.allow_tf32 = True
265
266 + # Set seeds for reproducibility
267 torch.manual_seed(42)
268 torch.cuda.manual_seed(42)
269
270 # Define benchmark shapes (m, n, k)
271 + shapes = [
272 + (8, 4096, 7168),
273 + (8, 7168, 18432),
274 + (8, 18432, 7168),
275 + (64, 4096, 7168),
276 + (64, 7168, 18432),
277 + (64, 18432, 7168),
278 + (64, 24576, 1536),
279 + (64, 32768, 512),
280 + (64, 7168, 16384),
281 + (128, 4096, 7168),
282 + (128, 7168, 18432),
283 + (128, 18432, 7168),
284 + (1024, 4096, 7168),
285 + (1024, 18432, 7168),
286 + (2048, 4096, 7168),
287 + (4096, 4096, 7168),
288 + ]
289 shapes = [
290 # (64, 2112, 7168),
291 (64, 24576, 1536),
@@ -265,6 +312,7 @@ def run_benchmarks(verbose: bool = False):
312 result = benchmark_shape(m, n, k, verbose=verbose)
313 all_results.append(result)
314
315 + # Print results in a nicely formatted table
316 print("\n===== PERFORMANCE COMPARISON =====")
317
318 # Print DeepGEMM table
@@ -273,50 +321,38 @@ def run_benchmarks(verbose: bool = False):
321 for result in all_results:
322 shape = result["shape"]
323 impl_data = result["implementations"]["DeepGEMM"]
276 - deepgemm_rows.append(
277 - [
278 - shape["m"],
279 - shape["n"],
280 - shape["k"],
281 - f"{impl_data['time_us']:.1f}",
282 - f"{impl_data['tflops']:.1f}",
283 - f"{impl_data['gb_s']:.1f}",
284 - ]
285 - )
324 + deepgemm_rows.append([
325 + shape["m"], shape["n"], shape["k"], f"{impl_data['time_us']:.1f}",
326 + f"{impl_data['tflops']:.1f}", f"{impl_data['gb_s']:.1f}"
327 + ])
328
287 - print_table(deepgemm_headers, deepgemm_rows, title="DeepGEMM Implementation:")
329 + print_table(deepgemm_headers,
330 + deepgemm_rows,
331 + title="DeepGEMM Implementation:")
332
333 # Print vLLM Triton table
290 - triton_headers = ["m", "n", "k", "Time (μs)", "TFLOPS", "GB/s", "vs DeepGEMM"]
334 + triton_headers = [
335 + "m", "n", "k", "Time (μs)", "TFLOPS", "GB/s", "vs DeepGEMM"
336 + ]
337 triton_rows = []
338 for result in all_results:
339 shape = result["shape"]
340 impl_data = result["implementations"]["vLLM Triton"]
341 speedup = impl_data.get("speedup_vs_deepgemm", 1.0)
296 - triton_rows.append(
297 - [
298 - shape["m"],
299 - shape["n"],
300 - shape["k"],
301 - f"{impl_data['time_us']:.1f}",
302 - f"{impl_data['tflops']:.1f}",
303 - f"{impl_data['gb_s']:.1f}",
304 - format_speedup(speedup),
305 - ]
306 - )
342 + triton_rows.append([
343 + shape["m"], shape["n"], shape["k"], f"{impl_data['time_us']:.1f}",
344 + f"{impl_data['tflops']:.1f}", f"{impl_data['gb_s']:.1f}",
345 + format_speedup(speedup)
346 + ])
347
308 - print_table(triton_headers, triton_rows, title="vLLM Triton Implementation:")
348 + print_table(triton_headers,
349 + triton_rows,
350 + title="vLLM Triton Implementation:")
351
352 # Print vLLM CUTLASS table
353 cutlass_headers = [
312 - "m",
313 - "n",
314 - "k",
315 - "Time (μs)",
316 - "TFLOPS",
317 - "GB/s",
318 - "vs DeepGEMM",
319 - "vs Triton",
354 + "m", "n", "k", "Time (μs)", "TFLOPS", "GB/s", "vs DeepGEMM",
355 + "vs Triton"
356 ]
357 cutlass_rows = []
358 for result in all_results:
@@ -324,27 +360,28 @@ def run_benchmarks(verbose: bool = False):
360 impl_data = result["implementations"]["vLLM CUTLASS"]
361 vs_deepgemm = impl_data.get("speedup_vs_deepgemm", 1.0)
362 vs_triton = impl_data.get("speedup_vs_triton", 1.0)
327 - cutlass_rows.append(
328 - [
329 - shape["m"],
330 - shape["n"],
331 - shape["k"],
332 - f"{impl_data['time_us']:.1f}",
333 - f"{impl_data['tflops']:.1f}",
334 - f"{impl_data['gb_s']:.1f}",
335 - format_speedup(vs_deepgemm),
336 - format_speedup(vs_triton),
337 - ]
338 - )
363 + cutlass_rows.append([
364 + shape["m"], shape["n"], shape["k"], f"{impl_data['time_us']:.1f}",
365 + f"{impl_data['tflops']:.1f}", f"{impl_data['gb_s']:.1f}",
366 + format_speedup(vs_deepgemm),
367 + format_speedup(vs_triton)
368 + ])
369
340 - print_table(cutlass_headers, cutlass_rows, title="vLLM CUTLASS Implementation:")
370 + print_table(cutlass_headers,
371 + cutlass_rows,
372 + title="vLLM CUTLASS Implementation:")
373
374 # Calculate and print averages
375 print("\n===== AVERAGE PERFORMANCE =====")
376
377 implementations = ["DeepGEMM", "vLLM Triton", "vLLM CUTLASS"]
378 avg_metrics = {
347 - impl: {"tflops": 0, "gb_s": 0, "time_ms": 0} for impl in implementations
379 + impl: {
380 + "tflops": 0,
381 + "gb_s": 0,
382 + "time_ms": 0
383 + }
384 + for impl in implementations
385 }
386
387 for result in all_results:
@@ -362,9 +399,9 @@ def run_benchmarks(verbose: bool = False):
399 avg_tflops = avg_metrics[impl]["tflops"] / num_shapes
400 avg_mem_bw = avg_metrics[impl]["gb_s"] / num_shapes
401 avg_time = avg_metrics[impl]["time_ms"] / num_shapes
365 - avg_rows.append(
366 - [impl, f"{avg_tflops:.2f}", f"{avg_mem_bw:.2f}", f"{avg_time:.2f}"]
367 - )
402 + avg_rows.append([
403 + impl, f"{avg_tflops:.2f}", f"{avg_mem_bw:.2f}", f"{avg_time:.2f}"
404 + ])
405
406 print_table(avg_headers, avg_rows)
407
@@ -372,19 +409,21 @@ def run_benchmarks(verbose: bool = False):
409 avg_speedups = {
410 "DeepGEMM vs vLLM Triton": 0,
411 "DeepGEMM vs vLLM CUTLASS": 0,
375 - "vLLM CUTLASS vs vLLM Triton": 0,
412 + "vLLM CUTLASS vs vLLM Triton": 0
413 }
414
415 for result in all_results:
416 deepgemm_time = result["implementations"]["DeepGEMM"]["time_ms"]
417 vllm_triton_time = result["implementations"]["vLLM Triton"]["time_ms"]
381 - vllm_cutlass_time = result["implementations"]["vLLM CUTLASS"]["time_ms"]
418 + vllm_cutlass_time = result["implementations"]["vLLM CUTLASS"][
419 + "time_ms"]
420
383 - avg_speedups["DeepGEMM vs vLLM Triton"] += vllm_triton_time / deepgemm_time
384 - avg_speedups["DeepGEMM vs vLLM CUTLASS"] += vllm_cutlass_time / deepgemm_time
385 - avg_speedups["vLLM CUTLASS vs vLLM Triton"] += (
386 - vllm_triton_time / vllm_cutlass_time
387 - )
421 + avg_speedups[
422 + "DeepGEMM vs vLLM Triton"] += vllm_triton_time / deepgemm_time
423 + avg_speedups[
424 + "DeepGEMM vs vLLM CUTLASS"] += vllm_cutlass_time / deepgemm_time
425 + avg_speedups[
426 + "vLLM CUTLASS vs vLLM Triton"] += vllm_triton_time / vllm_cutlass_time
427
428 print("\n===== AVERAGE SPEEDUPS =====")
429 speedup_headers = ["Comparison", "Speedup"]
@@ -396,12 +435,14 @@ def run_benchmarks(verbose: bool = False):
435
436 print_table(speedup_headers, speedup_rows)
437
438 + # Average accuracy comparison
439 print("\n===== ACCURACY COMPARISON =====")
440 avg_diff = {impl: 0 for impl in implementations}
441
442 for result in all_results:
443 for impl in implementations:
404 - avg_diff[impl] += result["implementations"][impl]["diff"]["Reference"]
444 + avg_diff[impl] += result["implementations"][impl]["diff"][
445 + "Reference"]
446
447 diff_headers = ["Implementation", "Avg Diff vs Reference"]
448 diff_rows = []
tests/kernels/moe/test_deepep_deepgemm_moe.py
+21 -2
@@ -66,6 +66,25 @@ def next_power_of_2(x):
66 return 2**math.ceil(math.log2(x))
67
68
69 +def per_block_cast_to_fp8(
70 + x: torch.Tensor,
71 + block_size_n: int = 128) -> tuple[torch.Tensor, torch.Tensor]:
72 + assert x.dim() == 2
73 + m, n = x.shape
74 + x_padded = torch.zeros(
75 + (deep_gemm.ceil_div(m, 128) * 128,
76 + deep_gemm.ceil_div(n, block_size_n) * block_size_n),
77 + dtype=x.dtype,
78 + device=x.device)
79 + x_padded[:m, :n] = x
80 + x_view = x_padded.view(-1, 128, x_padded.size(1) // 128, block_size_n)
81 + x_amax = x_view.abs().float().amax(dim=(1, 3), keepdim=True).clamp(1e-4)
82 + x_scaled = (x_view * (448.0 / x_amax)).to(torch.float8_e4m3fn)
83 + x_scaled_sub = x_scaled.view_as(x_padded)[:m, :n].contiguous()
84 + scales = (x_amax / 448.0).view(x_view.size(0), x_view.size(2))
85 + return x_scaled_sub, scales
86 +
87 +
88 def make_block_quant_fp8_weights(
89 e: int,
90 n: int,
@@ -106,8 +125,8 @@ def make_block_quant_fp8_weights(
125 assert (w2.shape[-2] + block_n - 1) // block_n == w2_s.shape[-2]
126
127 for i in range(e):
109 - w1[i], w1_s[i] = deep_gemm.utils.math.per_block_cast_to_fp8(w1_bf16[i])
110 - w2[i], w2_s[i] = deep_gemm.utils.math.per_block_cast_to_fp8(w2_bf16[i])
128 + w1[i], w1_s[i] = per_block_cast_to_fp8(w1_bf16[i])
129 + w2[i], w2_s[i] = per_block_cast_to_fp8(w2_bf16[i])
130
131 return w1, w2, w1_s, w2_s
132
tests/kernels/quantization/test_block_fp8.py
+41 -14
@@ -18,8 +18,7 @@ from vllm.model_executor.layers.fused_moe.fused_moe import (
18 from vllm.model_executor.layers.fused_moe.moe_align_block_size import (
19 moe_align_block_size)
20 from vllm.model_executor.layers.quantization.utils.fp8_utils import (
21 - get_col_major_tma_aligned_tensor, per_token_group_quant_fp8,
22 - w8a8_block_fp8_matmul)
21 + per_token_group_quant_fp8, w8a8_block_fp8_matmul)
22 from vllm.platforms import current_platform
23
24 dg_available = False
@@ -264,6 +263,25 @@ def test_w8a8_block_fp8_fused_moe(M, N, K, E, topk, block_size, dtype, seed):
263 assert rel_diff < 0.03
264
265
266 +def per_block_cast_to_fp8(
267 + x: torch.Tensor,
268 + block_size_n: int = 128) -> tuple[torch.Tensor, torch.Tensor]:
269 + assert x.dim() == 2
270 + m, n = x.shape
271 + x_padded = torch.zeros(
272 + (deep_gemm.ceil_div(m, 128) * 128,
273 + deep_gemm.ceil_div(n, block_size_n) * block_size_n),
274 + dtype=x.dtype,
275 + device=x.device)
276 + x_padded[:m, :n] = x
277 + x_view = x_padded.view(-1, 128, x_padded.size(1) // 128, block_size_n)
278 + x_amax = x_view.abs().float().amax(dim=(1, 3), keepdim=True).clamp(1e-4)
279 + x_scaled = (x_view * (448.0 / x_amax)).to(torch.float8_e4m3fn)
280 + x_scaled_sub = x_scaled.view_as(x_padded)[:m, :n].contiguous()
281 + scales = (x_amax / 448.0).view(x_view.size(0), x_view.size(2))
282 + return x_scaled_sub, scales
283 +
284 +
285 @pytest.mark.parametrize(
286 "M,N,K,block_size,out_dtype,seed",
287 itertools.product(M, N, K, BLOCK_SIZE, OUT_DTYPES, SEEDS))
@@ -281,8 +299,10 @@ def test_w8a8_block_fp8_deep_gemm_matmul(M, N, K, block_size, out_dtype, seed):
299 A_fp32 = (torch.rand(M, K, dtype=torch.float32) - 0.5) * 2 * fp8_max
300 B_fp32 = (torch.rand(N, K, dtype=torch.float32) - 0.5) * 2 * fp8_max
301
284 - A_fp8, As_fp8 = deep_gemm.utils.math.per_token_cast_to_fp8(A_fp32)
285 - B_fp8, Bs_fp8 = deep_gemm.utils.math.per_block_cast_to_fp8(B_fp32)
302 + _, block_k = block_size[0], block_size[1]
303 +
304 + A_fp8, As_fp8 = per_token_group_quant_fp8(A_fp32, block_k)
305 + B_fp8, Bs_fp8 = per_block_cast_to_fp8(B_fp32)
306
307 As = As_fp8.to(torch.float32)
308 Bs = Bs_fp8.to(torch.float32)
@@ -290,12 +310,15 @@ def test_w8a8_block_fp8_deep_gemm_matmul(M, N, K, block_size, out_dtype, seed):
310 ref_out = native_w8a8_block_matmul(A_fp8, B_fp8, As, Bs, block_size,
311 out_dtype)
312
313 + # Transpose earlier so that the testing will not trigger transposing kernels
314 + As_fp8 = deep_gemm.get_col_major_tma_aligned_tensor(As_fp8)
315 +
316 out = torch.zeros((M, N), device='cuda', dtype=out_dtype)
317
318 assert As_fp8.shape == (M, (K + 127) //
319 128), f"{As_fp8.shape} != {(M, (K + 127) // 128)}"
320
298 - deep_gemm.fp8_gemm_nt((A_fp8, As_fp8), (B_fp8, Bs_fp8), out)
321 + deep_gemm.gemm_fp8_fp8_bf16_nt((A_fp8, As_fp8), (B_fp8, Bs_fp8), out)
322
323 rel_diff = (torch.mean(
324 torch.abs(out.to(torch.float32) - ref_out.to(torch.float32))) /
@@ -359,16 +382,16 @@ def deep_gemm_w8a8_block_fp8_moe(M, K, a, w1, w2, w1_s, w2_s, score, topk,
382 dtype=torch.bfloat16,
383 device=a.device)
384
362 - deep_gemm.m_grouped_fp8_gemm_nt_contiguous((a_q, a_s), (w1, w1_s),
363 - inter_out, m_indices)
385 + deep_gemm.m_grouped_gemm_fp8_fp8_bf16_nt_contiguous((a_q, a_s), (w1, w1_s),
386 + inter_out, m_indices)
387
388 act_out = SiluAndMul().forward_native(inter_out)
389 act_out_q, act_out_s = per_token_group_quant_fp8(act_out, block_k)
390
391 out = torch.zeros(a_q.shape[0], K, dtype=torch.bfloat16, device=a.device)
392
370 - deep_gemm.m_grouped_fp8_gemm_nt_contiguous((act_out_q, act_out_s),
371 - (w2, w2_s), out, m_indices)
393 + deep_gemm.m_grouped_gemm_fp8_fp8_bf16_nt_contiguous(
394 + (act_out_q, act_out_s), (w2, w2_s), out, m_indices)
395
396 final_out = _moe_unpermute(out, inv_perm, topk, K, topk_weight)
397
@@ -418,15 +441,15 @@ def test_w8a8_block_fp8_deep_gemm_fused_moe(M, N, K, E, topk, seed):
441 w1_s = torch.empty((E, n_tiles_w1, k_tiles_w1), dtype=torch.float32)
442 w2_s = torch.empty((E, n_tiles_w2, k_tiles_w2), dtype=torch.float32)
443
421 - w1_s = get_col_major_tma_aligned_tensor(w1_s).contiguous()
422 - w2_s = get_col_major_tma_aligned_tensor(w2_s).contiguous()
444 + w1_s = deep_gemm.get_col_major_tma_aligned_tensor(w1_s).contiguous()
445 + w2_s = deep_gemm.get_col_major_tma_aligned_tensor(w2_s).contiguous()
446
447 assert w1_s.shape == (E, (2 * N + 127) // 128, (K + 127) // 128)
448 assert (w2.shape[-2] + block_n - 1) // block_n == w2_s.shape[-2]
449
450 for i in range(E):
428 - w1[i], w1_s[i] = deep_gemm.utils.math.per_block_cast_to_fp8(w1_bf16[i])
429 - w2[i], w2_s[i] = deep_gemm.utils.math.per_block_cast_to_fp8(w2_bf16[i])
451 + w1[i], w1_s[i] = per_block_cast_to_fp8(w1_bf16[i])
452 + w2[i], w2_s[i] = per_block_cast_to_fp8(w2_bf16[i])
453
454 # Set the context to avoid lots of warning spam.
455 with set_current_vllm_config(vllm_config):
@@ -437,10 +460,14 @@ def test_w8a8_block_fp8_deep_gemm_fused_moe(M, N, K, E, topk, seed):
460 ref_out = torch_w8a8_block_fp8_moe(a, w1, w2, w1_s, w2_s, score,
461 topk, block_size)
462
440 - topk_weights, topk_ids, _ = fused_topk(a, score.float(), topk, False)
463 + topk_weights, topk_ids, token_expert_indices = fused_topk(
464 + a, score.float(), topk, False)
465
466 out = deep_gemm_moe_fp8(a, w1, w2, w1_s, w2_s, topk_weights, topk_ids)
467
468 + #print(f"{out.sum()=}")
469 + #print(f"{ref_out.sum()=}")
470 +
471 rel_diff = (torch.mean(
472 torch.abs(out.to(torch.float32) - ref_out.to(torch.float32))) /
473 torch.mean(torch.abs(ref_out.to(torch.float32))))
vllm/model_executor/layers/fused_moe/batched_deep_gemm_moe.py
+11 -8
@@ -266,16 +266,19 @@ class BatchedDeepGemmExperts(mk.FusedMoEPermuteExpertsUnpermute):
266 # for the M expectation of each batch, correctly setting this value
267 # may lead to better performance.
268 expected_m = max_num_tokens
269 - dg.fp8_m_grouped_gemm_nt_masked((a1q, a1q_scale), (w1, w1_scale),
270 - out=workspace1,
271 - masked_m=expert_num_tokens,
272 - expected_m=expected_m)
269 +
270 + dg.m_grouped_gemm_fp8_fp8_bf16_nt_masked((a1q, a1q_scale),
271 + (w1, w1_scale),
272 + out=workspace1,
273 + masked_m=expert_num_tokens,
274 + expected_m=expected_m)
275
276 assert expert_num_tokens is not None
277 a2q, a2q_scale = silu_mul_fp8_quant_deep_gemm(workspace1,
278 expert_num_tokens)
279
278 - dg.fp8_m_grouped_gemm_nt_masked((a2q, a2q_scale), (w2, w2_scale),
279 - out=output,
280 - masked_m=expert_num_tokens,
281 - expected_m=expected_m)
280 + dg.m_grouped_gemm_fp8_fp8_bf16_nt_masked((a2q, a2q_scale),
281 + (w2, w2_scale),
282 + out=output,
283 + masked_m=expert_num_tokens,
284 + expected_m=expected_m)
vllm/model_executor/layers/fused_moe/deep_gemm_moe.py
+6 -4
@@ -143,9 +143,10 @@ class DeepGemmExperts(mk.FusedMoEPermuteExpertsUnpermute):
143 quant_out = _resize_cache(workspace13.view(dtype=torch.float8_e4m3fn),
144 (M_sum, N // 2))
145 mm2_out = _resize_cache(workspace2, (M_sum, K))
146 + # import pdb; pdb.set_trace()
147
147 - dg.m_grouped_fp8_gemm_nt_contiguous((a1q, a1q_scale), (w1, w1_scale),
148 - mm1_out, expert_ids)
148 + dg.m_grouped_gemm_fp8_fp8_bf16_nt_contiguous(
149 + (a1q, a1q_scale), (w1, w1_scale), mm1_out, expert_ids)
150
151 self.activation(activation, act_out, mm1_out.view(-1, N))
152
@@ -154,8 +155,9 @@ class DeepGemmExperts(mk.FusedMoEPermuteExpertsUnpermute):
155 self.block_shape[1],
156 column_major_scales=True,
157 out_q=quant_out)
157 - dg.m_grouped_fp8_gemm_nt_contiguous((a2q, a2q_scale), (w2, w2_scale),
158 - mm2_out, expert_ids)
158 +
159 + dg.m_grouped_gemm_fp8_fp8_bf16_nt_contiguous(
160 + (a2q, a2q_scale), (w2, w2_scale), mm2_out, expert_ids)
161
162 torch.index_select(mm2_out, 0, inv_perm, out=output)
163
vllm/model_executor/layers/quantization/deepgemm.py
+1 -1
@@ -58,7 +58,7 @@ def w8a8_block_fp8_matmul_deepgemm(
58 output_dtype)
59 # Deepgemm only supports output tensor type as bfloat16
60 assert C.dtype == torch.bfloat16
61 - deep_gemm.fp8_gemm_nt((A, As), (B, Bs), C)
61 + deep_gemm.gemm_fp8_fp8_bf16_nt((A, As), (B, Bs), C)
62 return C
63
64
vllm/model_executor/layers/quantization/utils/fp8_utils.py
+3 -68
@@ -114,10 +114,6 @@ def should_use_deepgemm(output_dtype: torch.dtype, weight: torch.Tensor):
114 and weight.shape[0] % 128 == 0 and weight.shape[1] % 128 == 0)
115
116
117 -def ceil_div(x: int, y: int) -> int:
118 - return (x + y - 1) // y
119 -
120 -
117 # TODO fix ROCm->Triton custom path:
118 # https://github.com/vllm-project/vllm/issues/14397
119 def apply_w8a8_block_fp8_linear(
@@ -162,6 +158,9 @@ def apply_w8a8_block_fp8_linear(
158 if current_platform.is_cuda():
159 if current_platform.has_device_capability(100):
160
161 + def ceil_div(x: int, y: int) -> int:
162 + return (x + y - 1) // y
163 +
164 use_cutlass = cutlass_block_fp8_supported and (
165 ceil_div(weight.shape[0], 128) == weight_scale.shape[0]
166 and ceil_div(weight.shape[1], 128) == weight_scale.shape[1])
@@ -656,67 +655,3 @@ def w8a8_block_fp8_matmul(
655 )
656
657 return C
659 -
660 -
661 -# Taken from https://github.com/deepseek-ai/DeepGEMM/blob/0c88cd01392c1073c7049a97d6328c7bba9b3947
662 -# TODO(wentao): remove this function when DeepGEMM exposes this function
663 -def get_tma_aligned_size(x: int, element_size: int) -> int:
664 - """
665 - Global memory address of TMA must be 16-byte aligned.
666 - Since we use column-major layout for the LHS scaling tensor,
667 - the M-axis of the LHS scaling tensor needs to be padded to a multiple of
668 - 16 bytes.
669 -
670 - Arguments:
671 - x: original M-axis shape of the LHS scaling tensor.
672 - element_size: element size of the LHS scaling tensor.
673 -
674 - Returns:
675 - M-axis shape of the LHS scaling tensor after padding.
676 - """
677 - tma_alignment_bytes = 16
678 - assert tma_alignment_bytes % element_size == 0
679 - alignment = tma_alignment_bytes // element_size
680 - return ceil_div(x, alignment) * alignment
681 -
682 -
683 -# Taken from https://github.com/deepseek-ai/DeepGEMM/blob/0c88cd01392c1073c7049a97d6328c7bba9b3947
684 -# TODO(wentao): remove this function when DeepGEMM exposes this function
685 -def get_col_major_tma_aligned_tensor(x: torch.Tensor) -> torch.Tensor:
686 - """
687 - Returns TMA-aligned transposed format of the input tensor. `torch.transpose`
688 - will be called if necessary.
689 - If the input tensor is already column-major layout and 16-byte aligned along
690 - the M axis (thus meets the requirement of LHS scaling tensor in
691 - DeepGEMM), this function will do nothing.
692 -
693 - Arguments:
694 - x: usually the LHS scaling tensor in GEMM.
695 -
696 - Returns:
697 - The LHS scaling tensor of TMA-aligned transposed format.
698 - """
699 - # NOTES: for the extreme performance, you may rewrite/fuse this function in
700 - # CUDA
701 - assert x.dim() in (2, 3)
702 - remove_dim = False
703 - m, n = x.shape[-2], x.shape[-1]
704 - aligned_m = get_tma_aligned_size(m, x.element_size())
705 - if x.dim() == 2:
706 - if x.stride(0) == 1 and x.stride(1) == aligned_m:
707 - return x
708 - x, remove_dim = x.unsqueeze(0), True
709 -
710 - b = x.shape[0]
711 -
712 - # The last kernel gives a column-major TMA aligned layout
713 - if x.stride(0) == aligned_m * n and x.stride(1) == 1 and x.stride(
714 - 2) == aligned_m:
715 - return x.squeeze(0) if remove_dim else x
716 -
717 - # Normal layout requires transposing
718 - aligned_x = torch.transpose(
719 - torch.empty((b, n, aligned_m), device=x.device, dtype=x.dtype), 1, 2)
720 - aligned_x[:, :m, :] = x
721 - aligned_x = aligned_x[:, :m, :]
722 - return aligned_x.squeeze(0) if remove_dim else aligned_x