@tduncan / Nvidia-Optimization-Stack-VLLM Mirror · read-only / commits / c6e3bba8e6

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

[Feature] Integrate new deepgemm (#19820)

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

Wentao Ye committed Jun 24, 2025 at 15:51 UTC c6e3bba8e68ef586e346a0a8f5c3af168148fd63

8 files changed +234 -268

benchmarks/kernels/benchmark_moe.py

@@ -86,6 +86,9 @@ 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]

92 if use_fp8_w8a8:

93 if block_quant_shape:

94 block_n, block_k = block_quant_shape[0], block_quant_shape[1]

benchmarks/kernels/deepgemm/benchmark_fp8_block_dense_gemm.py

+134 -186

@@ -1,13 +1,11 @@

1 # SPDX-License-Identifier: Apache-2.0

2 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project

3 -# fmt: off

4 -# ruff: noqa: E501

3 import time

7 -# Import DeepGEMM functions

8 -import deep_gemm

5 import torch

10 -from deep_gemm import calc_diff, ceil_div, get_col_major_tma_aligned_tensor

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 # Import vLLM functions

11 from vllm import _custom_ops as ops

@@ -18,107 +16,84 @@ from vllm.model_executor.layers.quantization.utils.fp8_utils import (

16 from vllm.triton_utils import triton

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 -

19 # Copied from

20 # https://github.com/deepseek-ai/DeepGEMM/blob/78cacf70d41d15d688bd493ebc85845f7f2a3d5d/tests/test_core.py#L17

36 -def per_block_cast_to_fp8(

37 - x: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor]:

21 +def per_block_cast_to_fp8_vllm(x: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor]:

22 """Convert tensor to FP8 format with per-block scaling."""

23 assert x.dim() == 2

24 m, n = x.shape

41 - x_padded = torch.zeros((ceil_div(m, 128) * 128, ceil_div(n, 128) * 128),

42 - dtype=x.dtype,

43 - device=x.device)

25 + x_padded = torch.zeros(

26 + (ceil_div(m, 128) * 128, ceil_div(n, 128) * 128), dtype=x.dtype, device=x.device

27 + )

28 x_padded[:m, :n] = x

29 x_view = x_padded.view(-1, 128, x_padded.size(1) // 128, 128)

30 x_amax = x_view.abs().float().amax(dim=(1, 3), keepdim=True).clamp(1e-4)

31 x_scaled = (x_view * (448.0 / x_amax)).to(torch.float8_e4m3fn)

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:

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:

45 """Benchmark all implementations for a specific (m, n, k) shape."""

46 if verbose:

47 print(f"\n=== Benchmarking shape: m={m}, n={n}, k={k} ===")

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

49 + A = torch.randn((m, k), device="cuda", dtype=torch.bfloat16)

50 + B = torch.randn((n, k), device="cuda", dtype=torch.bfloat16)

51 torch.cuda.synchronize()

52 C_ref = A @ B.t()

54 # Pre-quantize B for all implementations

55 # (weights can be pre-quantized offline)

56 B_deepgemm, B_scale_deepgemm = per_block_cast_to_fp8(B)

73 - B_vllm, B_scale_vllm = per_block_cast_to_fp8(B)

57 + B_vllm, B_scale_vllm = per_block_cast_to_fp8_vllm(B)

59 # Block size configuration

60 block_size = [128, 128]

62 # Pre-quantize A for all implementations

63 A_deepgemm, A_scale_deepgemm = per_token_cast_to_fp8(A)

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)

64 + C_deepgemm = (

65 + torch.empty((n, m), device="cuda", dtype=torch.bfloat16).t().contiguous()

66 + )

67 A_vllm, A_scale_vllm = per_token_group_quant_fp8(A, block_size[1])

68 A_vllm_cutlass, A_scale_vllm_cutlass = per_token_group_quant_fp8(

84 - A, block_size[1], column_major_scales=True)

69 + A, block_size[1], column_major_scales=True

70 + )

86 - # === DeepGEMM Implementation ===

72 def deepgemm_gemm():

88 - # A quantization is inside the loop as it depends on activations

89 - # A_deepgemm, A_scale_deepgemm = per_token_cast_to_fp8(A)

90 - # A_deepgemm, A_scale_deepgemm = per_token_group_quant_fp8(

91 - # A, block_size[1])

92 - # A_scale_aligned = get_col_major_tma_aligned_tensor(A_scale_deepgemm)

93 - # C_deepgemm = torch.empty((m, n), device='cuda', dtype=torch.bfloat16)

94 - deep_gemm.gemm_fp8_fp8_bf16_nt((A_deepgemm, A_scale_deepgemm),

95 - (B_deepgemm, B_scale_deepgemm),

96 - C_deepgemm)

73 + fp8_gemm_nt(

74 + (A_deepgemm, A_scale_deepgemm), (B_deepgemm, B_scale_deepgemm), C_deepgemm

75 + )

76 return C_deepgemm

99 - # === vLLM Triton Implementation ===

78 def vllm_triton_gemm():

101 - # A quantization is inside the loop as it depends on activations

102 - # A_vllm, A_scale_vllm = per_token_group_quant_fp8(A, block_size[1])

103 - return w8a8_block_fp8_matmul(A_vllm,

104 - B_vllm,

105 - A_scale_vllm,

106 - B_scale_vllm,

107 - block_size,

108 - output_dtype=torch.bfloat16)

109 -

110 - # === vLLM CUTLASS Implementation ===

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 +

88 def vllm_cutlass_gemm():

112 - # A quantization is inside the loop as it depends on activations

113 - # A_vllm_cutlass, A_scale_vllm_cutlass = per_token_group_quant_fp8(

114 - # A, block_size[1], column_major_scales=True)

115 - return ops.cutlass_scaled_mm(A_vllm_cutlass,

116 - B_vllm.T,

117 - scale_a=A_scale_vllm_cutlass,

118 - scale_b=B_scale_vllm.T,

119 - out_dtype=torch.bfloat16)

120 -

121 - # Run correctness check first

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 + )

96 +

97 if verbose:

98 print("Running correctness check...")

99 C_deepgemm = deepgemm_gemm()

@@ -133,26 +108,22 @@ def benchmark_shape(m: int,

108 print(f"DeepGEMM vs Reference difference: {deepgemm_diff:.6f}")

109 print(f"vLLM Triton vs Reference difference: {vllm_triton_diff:.6f}")

110 print(f"vLLM CUTLASS vs Reference difference: {vllm_cutlass_diff:.6f}")

136 - print("vLLM Triton vs DeepGEMM difference: "

137 - f"{calc_diff(C_vllm_triton, C_deepgemm):.6f}")

138 - print("vLLM CUTLASS vs DeepGEMM difference: "

139 - f"{calc_diff(C_vllm_cutlass, C_deepgemm):.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 + )

119

141 - # Benchmark implementations

120 implementations = {

121 "DeepGEMM": deepgemm_gemm,

122 "vLLM Triton": vllm_triton_gemm,

145 - "vLLM CUTLASS": vllm_cutlass_gemm

123 + "vLLM CUTLASS": vllm_cutlass_gemm,

124 }

125

148 - benchmark_results = {

149 - "shape": {

150 - "m": m,

151 - "n": n,

152 - "k": k

153 - },

154 - "implementations": {}

155 - }

126 + benchmark_results = {"shape": {"m": m, "n": n, "k": k}, "implementations": {}}

127

128 for name, func in implementations.items():

129 # Warmup

@@ -180,38 +151,36 @@ def benchmark_shape(m: int,

151 "tflops": tflops,

152 "gb_s": gb_s,

153 "diff": {

183 - "DeepGEMM":

184 - 0.0 if name == "DeepGEMM" else calc_diff(func(), C_deepgemm),

185 - "Reference":

186 - deepgemm_diff if name == "DeepGEMM" else

187 - (vllm_triton_diff

188 - if name == "vLLM Triton" else vllm_cutlass_diff)

189 - }

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 + },

161 }

162

163 if verbose:

193 - print(

194 - f"{name}: {avg_time_ms:.3f} ms, {tflops:.2f} TFLOPS, {gb_s:.2f} GB/s"

195 - )

164 + print(f"{name}: {avg_time_ms:.3f} ms, {tflops:.2f} TFLOPS, {gb_s:.2f} GB/s")

165

166 # Calculate speedups

167 baseline = benchmark_results["implementations"]["DeepGEMM"]["time_ms"]

168 for name, data in benchmark_results["implementations"].items():

169 if name != "DeepGEMM":

170 speedup = baseline / data["time_ms"]

202 - benchmark_results["implementations"][name][

203 - "speedup_vs_deepgemm"] = speedup

171 + benchmark_results["implementations"][name]["speedup_vs_deepgemm"] = speedup

172 if verbose:

205 - print(f"DeepGEMM is {1/speedup:.2f}x "

206 - f"{'faster' if 1/speedup > 1 else 'slower'} than {name}")

173 + print(

174 + f"DeepGEMM is {1 / speedup:.2f}x "

175 + f"{'faster' if 1 / speedup > 1 else 'slower'} than {name}"

176 + )

177

208 - vllm_triton_time = benchmark_results["implementations"]["vLLM Triton"][

209 - "time_ms"]

210 - vllm_cutlass_time = benchmark_results["implementations"]["vLLM CUTLASS"][

211 - "time_ms"]

178 + vllm_triton_time = benchmark_results["implementations"]["vLLM Triton"]["time_ms"]

179 + vllm_cutlass_time = benchmark_results["implementations"]["vLLM CUTLASS"]["time_ms"]

180 cutlass_vs_triton = vllm_triton_time / vllm_cutlass_time

213 - benchmark_results["implementations"]["vLLM CUTLASS"][

214 - "speedup_vs_triton"] = cutlass_vs_triton

181 + benchmark_results["implementations"]["vLLM CUTLASS"]["speedup_vs_triton"] = (

182 + cutlass_vs_triton

183 + )

184 if verbose:

185 print(

186 f"vLLM CUTLASS is {cutlass_vs_triton:.2f}x "

@@ -223,8 +192,7 @@ def benchmark_shape(m: int,

192

193 def format_table_row(values, widths):

194 """Format a row with specified column widths."""

226 - return "| " + " | ".join(f"{val:{w}}"

227 - for val, w in zip(values, widths)) + " |"

195 + return "| " + " | ".join(f"{val:{w}}" for val, w in zip(values, widths)) + " |"

196

197

198 def print_table(headers, rows, title=None):

@@ -232,16 +200,12 @@ def print_table(headers, rows, title=None):

200 if title:

201 print(f"\n{title}")

202

235 - # Calculate column widths based on headers and data

203 widths = [

204 max(len(str(h)), max(len(str(row[i])) for row in rows))

205 for i, h in enumerate(headers)

206 ]

207

241 - # Create separator line

208 separator = "+-" + "-+-".join("-" * w for w in widths) + "-+"

243 -

244 - # Print table

209 print(separator)

210 print(format_table_row(headers, widths))

211 print(separator)

@@ -259,44 +223,22 @@ def run_benchmarks(verbose: bool = False):

223 """Run benchmarks for a set of common shapes."""

224 print("===== STARTING FP8 GEMM BENCHMARK =====")

225

262 - # Make sure we're using the GPU

226 if not torch.cuda.is_available():

227 print("CUDA not available! Tests require GPU.")

228 return

229

267 - # Print system information

230 print(f"PyTorch version: {torch.__version__}")

231 print(f"CUDA version: {torch.version.cuda}")

232 print(f"Triton version: {triton.__version__}")

233 print(f"Using device: {torch.cuda.get_device_name()}")

234

273 - # Enable TF32 for better performance

235 torch.backends.cuda.matmul.allow_tf32 = True

236 torch.backends.cudnn.allow_tf32 = True

237

277 - # Set seeds for reproducibility

238 torch.manual_seed(42)

239 torch.cuda.manual_seed(42)

240

241 # Define benchmark shapes (m, n, k)

282 - shapes = [

283 - (8, 4096, 7168),

284 - (8, 7168, 18432),

285 - (8, 18432, 7168),

286 - (64, 4096, 7168),

287 - (64, 7168, 18432),

288 - (64, 18432, 7168),

289 - (64, 24576, 1536),

290 - (64, 32768, 512),

291 - (64, 7168, 16384),

292 - (128, 4096, 7168),

293 - (128, 7168, 18432),

294 - (128, 18432, 7168),

295 - (1024, 4096, 7168),

296 - (1024, 18432, 7168),

297 - (2048, 4096, 7168),

298 - (4096, 4096, 7168),

299 - ]

242 shapes = [

243 # (64, 2112, 7168),

244 (64, 24576, 1536),

@@ -323,7 +265,6 @@ def run_benchmarks(verbose: bool = False):

265 result = benchmark_shape(m, n, k, verbose=verbose)

266 all_results.append(result)

267

326 - # Print results in a nicely formatted table

268 print("\n===== PERFORMANCE COMPARISON =====")

269

270 # Print DeepGEMM table

@@ -332,38 +273,50 @@ def run_benchmarks(verbose: bool = False):

273 for result in all_results:

274 shape = result["shape"]

275 impl_data = result["implementations"]["DeepGEMM"]

335 - deepgemm_rows.append([

336 - shape["m"], shape["n"], shape["k"], f"{impl_data['time_us']:.1f}",

337 - f"{impl_data['tflops']:.1f}", f"{impl_data['gb_s']:.1f}"

338 - ])

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 + )

286

340 - print_table(deepgemm_headers,

341 - deepgemm_rows,

342 - title="DeepGEMM Implementation:")

287 + print_table(deepgemm_headers, deepgemm_rows, title="DeepGEMM Implementation:")

288

289 # Print vLLM Triton table

345 - triton_headers = [

346 - "m", "n", "k", "Time (μs)", "TFLOPS", "GB/s", "vs DeepGEMM"

347 - ]

290 + triton_headers = ["m", "n", "k", "Time (μs)", "TFLOPS", "GB/s", "vs DeepGEMM"]

291 triton_rows = []

292 for result in all_results:

293 shape = result["shape"]

294 impl_data = result["implementations"]["vLLM Triton"]

295 speedup = impl_data.get("speedup_vs_deepgemm", 1.0)

353 - triton_rows.append([

354 - shape["m"], shape["n"], shape["k"], f"{impl_data['time_us']:.1f}",

355 - f"{impl_data['tflops']:.1f}", f"{impl_data['gb_s']:.1f}",

356 - format_speedup(speedup)

357 - ])

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 + )

307

359 - print_table(triton_headers,

360 - triton_rows,

361 - title="vLLM Triton Implementation:")

308 + print_table(triton_headers, triton_rows, title="vLLM Triton Implementation:")

309

310 # Print vLLM CUTLASS table

311 cutlass_headers = [

365 - "m", "n", "k", "Time (μs)", "TFLOPS", "GB/s", "vs DeepGEMM",

366 - "vs Triton"

312 + "m",

313 + "n",

314 + "k",

315 + "Time (μs)",

316 + "TFLOPS",

317 + "GB/s",

318 + "vs DeepGEMM",

319 + "vs Triton",

320 ]

321 cutlass_rows = []

322 for result in all_results:

@@ -371,28 +324,27 @@ def run_benchmarks(verbose: bool = False):

324 impl_data = result["implementations"]["vLLM CUTLASS"]

325 vs_deepgemm = impl_data.get("speedup_vs_deepgemm", 1.0)

326 vs_triton = impl_data.get("speedup_vs_triton", 1.0)

374 - cutlass_rows.append([

375 - shape["m"], shape["n"], shape["k"], f"{impl_data['time_us']:.1f}",

376 - f"{impl_data['tflops']:.1f}", f"{impl_data['gb_s']:.1f}",

377 - format_speedup(vs_deepgemm),

378 - format_speedup(vs_triton)

379 - ])

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 + )

339

381 - print_table(cutlass_headers,

382 - cutlass_rows,

383 - title="vLLM CUTLASS Implementation:")

340 + print_table(cutlass_headers, cutlass_rows, title="vLLM CUTLASS Implementation:")

341

342 # Calculate and print averages

343 print("\n===== AVERAGE PERFORMANCE =====")

344

345 implementations = ["DeepGEMM", "vLLM Triton", "vLLM CUTLASS"]

346 avg_metrics = {

390 - impl: {

391 - "tflops": 0,

392 - "gb_s": 0,

393 - "time_ms": 0

394 - }

395 - for impl in implementations

347 + impl: {"tflops": 0, "gb_s": 0, "time_ms": 0} for impl in implementations

348 }

349

350 for result in all_results:

@@ -410,9 +362,9 @@ def run_benchmarks(verbose: bool = False):

362 avg_tflops = avg_metrics[impl]["tflops"] / num_shapes

363 avg_mem_bw = avg_metrics[impl]["gb_s"] / num_shapes

364 avg_time = avg_metrics[impl]["time_ms"] / num_shapes

413 - avg_rows.append([

414 - impl, f"{avg_tflops:.2f}", f"{avg_mem_bw:.2f}", f"{avg_time:.2f}"

415 - ])

365 + avg_rows.append(

366 + [impl, f"{avg_tflops:.2f}", f"{avg_mem_bw:.2f}", f"{avg_time:.2f}"]

367 + )

368

369 print_table(avg_headers, avg_rows)

370

@@ -420,21 +372,19 @@ def run_benchmarks(verbose: bool = False):

372 avg_speedups = {

373 "DeepGEMM vs vLLM Triton": 0,

374 "DeepGEMM vs vLLM CUTLASS": 0,

423 - "vLLM CUTLASS vs vLLM Triton": 0

375 + "vLLM CUTLASS vs vLLM Triton": 0,

376 }

377

378 for result in all_results:

379 deepgemm_time = result["implementations"]["DeepGEMM"]["time_ms"]

380 vllm_triton_time = result["implementations"]["vLLM Triton"]["time_ms"]

429 - vllm_cutlass_time = result["implementations"]["vLLM CUTLASS"][

430 - "time_ms"]

381 + vllm_cutlass_time = result["implementations"]["vLLM CUTLASS"]["time_ms"]

382

432 - avg_speedups[

433 - "DeepGEMM vs vLLM Triton"] += vllm_triton_time / deepgemm_time

434 - avg_speedups[

435 - "DeepGEMM vs vLLM CUTLASS"] += vllm_cutlass_time / deepgemm_time

436 - avg_speedups[

437 - "vLLM CUTLASS vs vLLM Triton"] += vllm_triton_time / vllm_cutlass_time

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 + )

388

389 print("\n===== AVERAGE SPEEDUPS =====")

390 speedup_headers = ["Comparison", "Speedup"]

@@ -446,14 +396,12 @@ def run_benchmarks(verbose: bool = False):

396

397 print_table(speedup_headers, speedup_rows)

398

449 - # Average accuracy comparison

399 print("\n===== ACCURACY COMPARISON =====")

400 avg_diff = {impl: 0 for impl in implementations}

401

402 for result in all_results:

403 for impl in implementations:

455 - avg_diff[impl] += result["implementations"][impl]["diff"][

456 - "Reference"]

404 + avg_diff[impl] += result["implementations"][impl]["diff"]["Reference"]

405

406 diff_headers = ["Implementation", "Avg Diff vs Reference"]

407 diff_rows = []

tests/kernels/moe/test_deepep_deepgemm_moe.py

+2 -21

@@ -66,25 +66,6 @@ def next_power_of_2(x):

66 return 2**math.ceil(math.log2(x))

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 -

69 def make_block_quant_fp8_weights(

70 e: int,

71 n: int,

@@ -125,8 +106,8 @@ def make_block_quant_fp8_weights(

106 assert (w2.shape[-2] + block_n - 1) // block_n == w2_s.shape[-2]

107

108 for i in range(e):

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])

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])

111

112 return w1, w2, w1_s, w2_s

113

tests/kernels/quantization/test_block_fp8.py

+14 -41

@@ -18,7 +18,8 @@ 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 - per_token_group_quant_fp8, w8a8_block_fp8_matmul)

21 + get_col_major_tma_aligned_tensor, per_token_group_quant_fp8,

22 + w8a8_block_fp8_matmul)

23 from vllm.platforms import current_platform

25 dg_available = False

@@ -263,25 +264,6 @@ def test_w8a8_block_fp8_fused_moe(M, N, K, E, topk, block_size, dtype, seed):

264 assert rel_diff < 0.03

265

266

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 -

267 @pytest.mark.parametrize(

268 "M,N,K,block_size,out_dtype,seed",

269 itertools.product(M, N, K, BLOCK_SIZE, OUT_DTYPES, SEEDS))

@@ -299,10 +281,8 @@ def test_w8a8_block_fp8_deep_gemm_matmul(M, N, K, block_size, out_dtype, seed):

281 A_fp32 = (torch.rand(M, K, dtype=torch.float32) - 0.5) * 2 * fp8_max

282 B_fp32 = (torch.rand(N, K, dtype=torch.float32) - 0.5) * 2 * fp8_max

283

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)

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)

286

287 As = As_fp8.to(torch.float32)

288 Bs = Bs_fp8.to(torch.float32)

@@ -310,15 +290,12 @@ def test_w8a8_block_fp8_deep_gemm_matmul(M, N, K, block_size, out_dtype, seed):

290 ref_out = native_w8a8_block_matmul(A_fp8, B_fp8, As, Bs, block_size,

291 out_dtype)

292

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 -

293 out = torch.zeros((M, N), device='cuda', dtype=out_dtype)

294

295 assert As_fp8.shape == (M, (K + 127) //

296 128), f"{As_fp8.shape} != {(M, (K + 127) // 128)}"

297

321 - deep_gemm.gemm_fp8_fp8_bf16_nt((A_fp8, As_fp8), (B_fp8, Bs_fp8), out)

298 + deep_gemm.fp8_gemm_nt((A_fp8, As_fp8), (B_fp8, Bs_fp8), out)

299

300 rel_diff = (torch.mean(

301 torch.abs(out.to(torch.float32) - ref_out.to(torch.float32))) /

@@ -382,16 +359,16 @@ def deep_gemm_w8a8_block_fp8_moe(M, K, a, w1, w2, w1_s, w2_s, score, topk,

359 dtype=torch.bfloat16,

360 device=a.device)

361

385 - deep_gemm.m_grouped_gemm_fp8_fp8_bf16_nt_contiguous((a_q, a_s), (w1, w1_s),

386 - inter_out, m_indices)

362 + deep_gemm.m_grouped_fp8_gemm_nt_contiguous((a_q, a_s), (w1, w1_s),

363 + inter_out, m_indices)

364

365 act_out = SiluAndMul().forward_native(inter_out)

366 act_out_q, act_out_s = per_token_group_quant_fp8(act_out, block_k)

367

368 out = torch.zeros(a_q.shape[0], K, dtype=torch.bfloat16, device=a.device)

369

393 - deep_gemm.m_grouped_gemm_fp8_fp8_bf16_nt_contiguous(

394 - (act_out_q, act_out_s), (w2, w2_s), out, m_indices)

370 + deep_gemm.m_grouped_fp8_gemm_nt_contiguous((act_out_q, act_out_s),

371 + (w2, w2_s), out, m_indices)

372

373 final_out = _moe_unpermute(out, inv_perm, topk, K, topk_weight)

374

@@ -441,15 +418,15 @@ def test_w8a8_block_fp8_deep_gemm_fused_moe(M, N, K, E, topk, seed):

418 w1_s = torch.empty((E, n_tiles_w1, k_tiles_w1), dtype=torch.float32)

419 w2_s = torch.empty((E, n_tiles_w2, k_tiles_w2), dtype=torch.float32)

420

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()

421 + w1_s = get_col_major_tma_aligned_tensor(w1_s).contiguous()

422 + w2_s = get_col_major_tma_aligned_tensor(w2_s).contiguous()

423

424 assert w1_s.shape == (E, (2 * N + 127) // 128, (K + 127) // 128)

425 assert (w2.shape[-2] + block_n - 1) // block_n == w2_s.shape[-2]

426

427 for i in range(E):

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])

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])

430

431 # Set the context to avoid lots of warning spam.

432 with set_current_vllm_config(vllm_config):

@@ -460,14 +437,10 @@ def test_w8a8_block_fp8_deep_gemm_fused_moe(M, N, K, E, topk, seed):

437 ref_out = torch_w8a8_block_fp8_moe(a, w1, w2, w1_s, w2_s, score,

438 topk, block_size)

439

463 - topk_weights, topk_ids, token_expert_indices = fused_topk(

464 - a, score.float(), topk, False)

440 + topk_weights, topk_ids, _ = fused_topk(a, score.float(), topk, False)

441

442 out = deep_gemm_moe_fp8(a, w1, w2, w1_s, w2_s, topk_weights, topk_ids)

443

468 - #print(f"{out.sum()=}")

469 - #print(f"{ref_out.sum()=}")

470 -

444 rel_diff = (torch.mean(

445 torch.abs(out.to(torch.float32) - ref_out.to(torch.float32))) /

446 torch.mean(torch.abs(ref_out.to(torch.float32))))

vllm/model_executor/layers/fused_moe/batched_deep_gemm_moe.py

+8 -11

@@ -266,19 +266,16 @@ 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 -

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)

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)

273

274 assert expert_num_tokens is not None

275 a2q, a2q_scale = silu_mul_fp8_quant_deep_gemm(workspace1,

276 expert_num_tokens)

277

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)

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)

vllm/model_executor/layers/fused_moe/deep_gemm_moe.py

+4 -5

@@ -144,8 +144,8 @@ class DeepGemmExperts(mk.FusedMoEPermuteExpertsUnpermute):

144 (M_sum, N // 2))

145 mm2_out = _resize_cache(workspace2, (M_sum, K))

146

147 - dg.m_grouped_gemm_fp8_fp8_bf16_nt_contiguous(

148 - (a1q, a1q_scale), (w1, w1_scale), mm1_out, expert_ids)

147 + dg.m_grouped_fp8_gemm_nt_contiguous((a1q, a1q_scale), (w1, w1_scale),

148 + mm1_out, expert_ids)

149

150 self.activation(activation, act_out, mm1_out.view(-1, N))

151

@@ -154,9 +154,8 @@ class DeepGemmExperts(mk.FusedMoEPermuteExpertsUnpermute):

154 self.block_shape[1],

155 column_major_scales=True,

156 out_q=quant_out)

157 -

158 - dg.m_grouped_gemm_fp8_fp8_bf16_nt_contiguous(

159 - (a2q, a2q_scale), (w2, w2_scale), mm2_out, expert_ids)

157 + dg.m_grouped_fp8_gemm_nt_contiguous((a2q, a2q_scale), (w2, w2_scale),

158 + mm2_out, expert_ids)

159

160 torch.index_select(mm2_out, 0, inv_perm, out=output)

161

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.gemm_fp8_fp8_bf16_nt((A, As), (B, Bs), C)

61 + deep_gemm.fp8_gemm_nt((A, As), (B, Bs), C)

62 return C

vllm/model_executor/layers/quantization/utils/fp8_utils.py

+68 -3

@@ -114,6 +114,10 @@ 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 +

121 # TODO fix ROCm->Triton custom path:

122 # https://github.com/vllm-project/vllm/issues/14397

123 def apply_w8a8_block_fp8_linear(

@@ -158,9 +162,6 @@ def apply_w8a8_block_fp8_linear(

162 if current_platform.is_cuda():

163 if current_platform.has_device_capability(100):

164

161 - def ceil_div(x: int, y: int) -> int:

162 - return (x + y - 1) // y

163 -

165 use_cutlass = cutlass_block_fp8_supported and (

166 ceil_div(weight.shape[0], 128) == weight_scale.shape[0]

167 and ceil_div(weight.shape[1], 128) == weight_scale.shape[1])

@@ -655,3 +656,67 @@ def w8a8_block_fp8_matmul(

656 )

657

658 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