@tduncan / Nvidia-Optimization-Stack-VLLM Mirror · read-only / commits / e59ca942f5
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, 16:00 UTC

Add option to use DeepGemm contiguous grouped gemm kernel for fused MoE operations. (#13932)

Signed-off-by: Bill Nell <bnell@redhat.com>

bnellnm committed Apr 1, 2025 at 12:07 UTC e59ca942f576c78fd457f16f2029bda716c81959
6 files changed +774 -115
benchmarks/kernels/benchmark_moe.py
+61 -36
@@ -30,19 +30,18 @@ class BenchmarkConfig(TypedDict):
30 num_stages: int
31
32
33 -def benchmark_config(
34 - config: BenchmarkConfig,
35 - num_tokens: int,
36 - num_experts: int,
37 - shard_intermediate_size: int,
38 - hidden_size: int,
39 - topk: int,
40 - dtype: torch.dtype,
41 - use_fp8_w8a8: bool,
42 - use_int8_w8a16: bool,
43 - num_iters: int = 100,
44 - block_quant_shape: List[int] = None,
45 -) -> float:
33 +def benchmark_config(config: BenchmarkConfig,
34 + num_tokens: int,
35 + num_experts: int,
36 + shard_intermediate_size: int,
37 + hidden_size: int,
38 + topk: int,
39 + dtype: torch.dtype,
40 + use_fp8_w8a8: bool,
41 + use_int8_w8a16: bool,
42 + num_iters: int = 100,
43 + block_quant_shape: List[int] = None,
44 + use_deep_gemm: bool = False) -> float:
45 init_dtype = torch.float16 if use_fp8_w8a8 else dtype
46 x = torch.randn(num_tokens, hidden_size, dtype=dtype)
47 if use_int8_w8a16:
@@ -115,22 +114,41 @@ def benchmark_config(
114 def run():
115 from vllm.model_executor.layers.fused_moe import override_config
116 with override_config(config):
118 - fused_moe(
119 - x,
120 - w1,
121 - w2,
122 - input_gating,
123 - topk,
124 - renormalize=True,
125 - inplace=True,
126 - use_fp8_w8a8=use_fp8_w8a8,
127 - use_int8_w8a16=use_int8_w8a16,
128 - w1_scale=w1_scale,
129 - w2_scale=w2_scale,
130 - a1_scale=a1_scale,
131 - a2_scale=a2_scale,
132 - block_shape=block_quant_shape,
133 - )
117 + if use_deep_gemm:
118 + topk_weights, topk_ids = fused_topk(x, input_gating, topk,
119 + False)
120 + return fused_experts(
121 + x,
122 + w1,
123 + w2,
124 + topk_weights,
125 + topk_ids,
126 + inplace=True,
127 + use_fp8_w8a8=use_fp8_w8a8,
128 + w1_scale=w1_scale,
129 + w2_scale=w2_scale,
130 + a1_scale=a1_scale,
131 + a2_scale=a2_scale,
132 + block_shape=block_quant_shape,
133 + allow_deep_gemm=True,
134 + )
135 + else:
136 + fused_moe(
137 + x,
138 + w1,
139 + w2,
140 + input_gating,
141 + topk,
142 + renormalize=True,
143 + inplace=True,
144 + use_fp8_w8a8=use_fp8_w8a8,
145 + use_int8_w8a16=use_int8_w8a16,
146 + w1_scale=w1_scale,
147 + w2_scale=w2_scale,
148 + a1_scale=a1_scale,
149 + a2_scale=a2_scale,
150 + block_shape=block_quant_shape,
151 + )
152
153 # JIT compilation & warmup
154 run()
@@ -366,6 +384,7 @@ class BenchmarkWorker:
384 use_fp8_w8a8: bool,
385 use_int8_w8a16: bool,
386 block_quant_shape: List[int] = None,
387 + use_deep_gemm: bool = False,
388 ) -> tuple[dict[str, int], float]:
389 current_platform.seed_everything(self.seed)
390 dtype_str = get_config_dtype_str(dtype,
@@ -396,7 +415,8 @@ class BenchmarkWorker:
415 use_fp8_w8a8,
416 use_int8_w8a16,
417 num_iters=100,
399 - block_quant_shape=block_quant_shape)
418 + block_quant_shape=block_quant_shape,
419 + use_deep_gemm=use_deep_gemm)
420 return config, kernel_time
421
422 def tune(
@@ -411,6 +431,7 @@ class BenchmarkWorker:
431 use_int8_w8a16: bool,
432 search_space: list[dict[str, int]],
433 block_quant_shape: list[int],
434 + use_deep_gemm: bool,
435 ) -> dict[str, int]:
436 best_config = None
437 best_time = float("inf")
@@ -436,7 +457,8 @@ class BenchmarkWorker:
457 use_fp8_w8a8,
458 use_int8_w8a16,
459 num_iters=20,
439 - block_quant_shape=block_quant_shape)
460 + block_quant_shape=block_quant_shape,
461 + use_deep_gemm=use_deep_gemm)
462 except triton.runtime.autotuner.OutOfResources:
463 # Some configurations may be invalid and fail to compile.
464 continue
@@ -550,6 +572,8 @@ def main(args: argparse.Namespace):
572 else:
573 batch_sizes = [args.batch_size]
574
575 + use_deep_gemm = bool(args.use_deep_gemm)
576 +
577 ray.init()
578 num_gpus = int(ray.available_resources()["GPU"])
579 workers = [BenchmarkWorker.remote(args.seed) for _ in range(num_gpus)]
@@ -572,10 +596,10 @@ def main(args: argparse.Namespace):
596
597 start = time.time()
598 configs = _distribute(
575 - "tune",
576 - [(batch_size, E, shard_intermediate_size, hidden_size, topk, dtype,
577 - use_fp8_w8a8, use_int8_w8a16, search_space, block_quant_shape)
578 - for batch_size in batch_sizes])
599 + "tune", [(batch_size, E, shard_intermediate_size, hidden_size,
600 + topk, dtype, use_fp8_w8a8, use_int8_w8a16, search_space,
601 + block_quant_shape, use_deep_gemm)
602 + for batch_size in batch_sizes])
603 best_configs = {
604 M: sort_config(config)
605 for M, config in zip(batch_sizes, configs)
@@ -589,7 +613,7 @@ def main(args: argparse.Namespace):
613 outputs = _distribute(
614 "benchmark",
615 [(batch_size, E, shard_intermediate_size, hidden_size, topk, dtype,
592 - use_fp8_w8a8, use_int8_w8a16, block_quant_shape)
616 + use_fp8_w8a8, use_int8_w8a16, block_quant_shape, use_deep_gemm)
617 for batch_size in batch_sizes])
618
619 for batch_size, (config, kernel_time) in zip(batch_sizes, outputs):
@@ -611,6 +635,7 @@ if __name__ == "__main__":
635 type=str,
636 choices=["auto", "fp8_w8a8", "int8_w8a16"],
637 default="auto")
638 + parser.add_argument("--use-deep-gemm", action="store_true")
639 parser.add_argument("--seed", type=int, default=0)
640 parser.add_argument("--batch-size", type=int, required=False)
641 parser.add_argument("--tune", action="store_true")
tests/kernels/test_block_fp8.py
+256 -25
@@ -6,12 +6,22 @@ import itertools
6 import pytest
7 import torch
8
9 +from vllm.config import VllmConfig, set_current_vllm_config
10 from vllm.model_executor.layers.activation import SiluAndMul
11 from vllm.model_executor.layers.fused_moe import fused_moe
12 +from vllm.model_executor.layers.fused_moe.fused_moe import (
13 + deep_gemm_moe_fp8, fused_topk, moe_align_block_size)
14 from vllm.model_executor.layers.quantization.utils.fp8_utils import (
15 per_token_group_quant_fp8, w8a8_block_fp8_matmul)
16 from vllm.platforms import current_platform
17
18 +dg_available = False
19 +try:
20 + import deep_gemm
21 + dg_available = True
22 +except ImportError:
23 + pass
24 +
25 if current_platform.get_device_capability() < (9, 0):
26 pytest.skip("FP8 Triton requires CUDA 9.0 or higher",
27 allow_module_level=True)
@@ -21,17 +31,18 @@ DTYPES = [torch.bfloat16] # [torch.half, torch.bfloat16, torch.float32]
31 NUM_TOKENS = [7, 83, 2048]
32 D = [512, 4096, 5120, 13824]
33 GROUP_SIZE = [64, 128, 256, 512]
24 -M = [1, 7, 83, 512, 2048]
25 -N = [128, 512, 1024, 4096, 7748, 13824]
26 -K = [256, 4096, 5120, 3884, 13824]
34 +M = [1, 7, 8, 83, 84, 512, 2048, 4096]
35 +N = [128, 512, 1024, 4096, 7168, 7748, 13824]
36 +K = [256, 4096, 5120, 3884, 13824, 16384]
37 # Deepseek-V3's intermediate size 18432, so N is 18432*2/8=4608 at TP8
38 # and its hidden size is 7168.
29 -M_moe = [1, 7, 83, 512, 2048]
30 -N_moe = [4608] # [128, 4608, 13824]
31 -K_moe = [7168] # [256, 7168, 13824]
39 +M_moe = [1, 2, 7, 83, 128, 512, 2048]
40 +M_moe_dg = [128, 192, 512, 1335, 2048]
41 +N_moe = [128, 256, 1024, 4608] # [13824]
42 +K_moe = [256, 512, 7168] # [13824]
43 BLOCK_SIZE = [[128, 128]]
33 -E = [8, 24] # [8, 24, 128, 256]
34 -TOP_KS = [2] # [1, 2, 6]
44 +E = [2, 8, 16, 24] # [128, 256]
45 +TOP_KS = [1, 2, 6]
46 OUT_DTYPES = [torch.bfloat16] # [torch.float32, torch.half, torch.bfloat16]
47 SEEDS = [0]
48
@@ -217,11 +228,16 @@ def test_w8a8_block_fp8_matmul(M, N, K, block_size, out_dtype, seed):
228 SEEDS))
229 @torch.inference_mode()
230 def test_w8a8_block_fp8_fused_moe(M, N, K, E, topk, block_size, dtype, seed):
231 + if topk > E:
232 + pytest.skip(f"Skipping test; topk={topk} > E={E}")
233 +
234 torch.manual_seed(seed)
235 factor_for_scale = 1e-2
236 fp8_info = torch.finfo(torch.float8_e4m3fn)
237 fp8_max, fp8_min = fp8_info.max, fp8_info.min
238
239 + vllm_config = VllmConfig()
240 +
241 a = torch.randn((M, K), dtype=dtype) / 10
242
243 w1_bf16 = (torch.rand(
@@ -246,25 +262,240 @@ def test_w8a8_block_fp8_fused_moe(M, N, K, E, topk, block_size, dtype, seed):
262
263 score = torch.randn((M, E), dtype=dtype)
264
249 - out = fused_moe(
250 - a,
251 - w1,
252 - w2,
253 - score,
254 - topk,
255 - renormalize=False,
256 - use_fp8_w8a8=True,
257 - w1_scale=w1_s,
258 - w2_scale=w2_s,
259 - block_shape=block_size,
260 - )
261 - ref_out = torch_w8a8_block_fp8_moe(a, w1, w2, w1_s, w2_s, score, topk,
262 - block_size)
263 -
264 - print(f"{out.sum()=}")
265 - print(f"{ref_out.sum()=}")
265 + # Set the context to avoid lots of warning spam.
266 + with set_current_vllm_config(vllm_config):
267 + out = fused_moe(
268 + a,
269 + w1,
270 + w2,
271 + score,
272 + topk,
273 + renormalize=False,
274 + use_fp8_w8a8=True,
275 + w1_scale=w1_s,
276 + w2_scale=w2_s,
277 + block_shape=block_size,
278 + )
279 + ref_out = torch_w8a8_block_fp8_moe(a, w1, w2, w1_s, w2_s, score, topk,
280 + block_size)
281 +
282 + #print(f"{out.sum()=}")
283 + #print(f"{ref_out.sum()=}")
284 +
285 + rel_diff = (torch.mean(
286 + torch.abs(out.to(torch.float32) - ref_out.to(torch.float32))) /
287 + torch.mean(torch.abs(ref_out.to(torch.float32))))
288 + assert rel_diff < 0.03
289 +
290 +
291 +def per_block_cast_to_fp8(
292 + x: torch.Tensor,
293 + block_size_n: int = 128) -> tuple[torch.Tensor, torch.Tensor]:
294 + assert x.dim() == 2
295 + m, n = x.shape
296 + x_padded = torch.zeros(
297 + (deep_gemm.ceil_div(m, 128) * 128,
298 + deep_gemm.ceil_div(n, block_size_n) * block_size_n),
299 + dtype=x.dtype,
300 + device=x.device)
301 + x_padded[:m, :n] = x
302 + x_view = x_padded.view(-1, 128, x_padded.size(1) // 128, block_size_n)
303 + x_amax = x_view.abs().float().amax(dim=(1, 3), keepdim=True).clamp(1e-4)
304 + x_scaled = (x_view * (448.0 / x_amax)).to(torch.float8_e4m3fn)
305 + x_scaled_sub = x_scaled.view_as(x_padded)[:m, :n].contiguous()
306 + scales = (x_amax / 448.0).view(x_view.size(0), x_view.size(2))
307 + return x_scaled_sub, scales
308 +
309 +
310 +@pytest.mark.parametrize(
311 + "M,N,K,block_size,out_dtype,seed",
312 + itertools.product(M, N, K, BLOCK_SIZE, OUT_DTYPES, SEEDS))
313 +@torch.inference_mode()
314 +def test_w8a8_block_fp8_deep_gemm_matmul(M, N, K, block_size, out_dtype, seed):
315 + # only aligned sizes
316 + if M % 4 != 0 or K % 128 != 0 or N % 64 != 0:
317 + pytest.skip(f"Skipping test; invalid size {M}, {N}, {K}")
318 +
319 + torch.manual_seed(seed)
320 + fp8_info = torch.finfo(torch.float8_e4m3fn)
321 + fp8_max = fp8_info.max
322 +
323 + A_fp32 = (torch.rand(M, K, dtype=torch.float32) - 0.5) * 2 * fp8_max
324 + B_fp32 = (torch.rand(N, K, dtype=torch.float32) - 0.5) * 2 * fp8_max
325 +
326 + _, block_k = block_size[0], block_size[1]
327 +
328 + A_fp8, As_fp8 = per_token_group_quant_fp8(A_fp32, block_k)
329 + B_fp8, Bs_fp8 = per_block_cast_to_fp8(B_fp32)
330 +
331 + As = As_fp8.to(torch.float32)
332 + Bs = Bs_fp8.to(torch.float32)
333 +
334 + ref_out = native_w8a8_block_fp8_matmul(A_fp8, B_fp8, As, Bs, block_size,
335 + out_dtype)
336 +
337 + # Transpose earlier so that the testing will not trigger transposing kernels
338 + As_fp8 = deep_gemm.get_col_major_tma_aligned_tensor(As_fp8)
339 +
340 + out = torch.zeros((M, N), device='cuda', dtype=out_dtype)
341 +
342 + assert As_fp8.shape == (M, (K + 127) //
343 + 128), f"{As_fp8.shape} != {(M, (K + 127) // 128)}"
344 +
345 + deep_gemm.gemm_fp8_fp8_bf16_nt((A_fp8, As_fp8), (B_fp8, Bs_fp8), out)
346 +
347 + rel_diff = (torch.mean(
348 + torch.abs(out.to(torch.float32) - ref_out.to(torch.float32))) /
349 + torch.mean(torch.abs(ref_out.to(torch.float32))))
350 + assert rel_diff < 0.001
351 +
352 +
353 +def fp8_perm(m, idx):
354 + if torch.is_floating_point(m) and torch.finfo(m.dtype).bits == 8:
355 + return m.view(dtype=torch.uint8)[idx, ...].view(dtype=m.dtype)
356 + else:
357 + return m[idx, ...]
358 +
359 +
360 +def test_moe_permute(a, a_s, topk_ids, num_groups, topk, block_m):
361 + M, K = a.shape
362 +
363 + sorted_token_ids, m_indices, num_pad = moe_align_block_size(
364 + topk_ids, block_m, num_groups, None, pad_sorted_ids=True)
365 +
366 + num_tokens = topk * M
367 +
368 + sorted_token_ids = sorted_token_ids.clamp(max=num_tokens - 1)
369 + m_indices = torch.repeat_interleave(m_indices, block_m, dim=0)
370 + inv_perm = torch.argsort(sorted_token_ids)[:M * topk]
371 +
372 + a = fp8_perm(a, sorted_token_ids // topk)
373 + if a_s is not None:
374 + a_s = a_s[sorted_token_ids // topk]
375 +
376 + return a, a_s, m_indices, inv_perm
377 +
378 +
379 +def test_moe_unpermute(out, inv_perm, topk, K, topk_weight):
380 + M = topk_weight.shape[0]
381 + out = out[inv_perm, ...]
382 + tmp_out = out.view(-1, topk, K)
383 + return (tmp_out * topk_weight.view(M, -1, 1).to(out.dtype)).sum(dim=1)
384 +
385 +
386 +def deep_gemm_w8a8_block_fp8_moe(M, K, a, w1, w2, w1_s, w2_s, score, topk,
387 + block_shape):
388 + """Fused moe with block-wise quantization using DeepGemm grouped gemm."""
389 + num_groups = w1.shape[0]
390 + M, K = a.shape
391 + N = w2.shape[-1]
392 +
393 + topk_weight, topk_ids = fused_topk(a, score.float(), topk, False)
394 +
395 + block_m = deep_gemm.get_m_alignment_for_contiguous_layout()
396 +
397 + _, block_k = block_shape[0], block_shape[1]
398 +
399 + a_q, a_s = per_token_group_quant_fp8(a, block_m)
400 +
401 + a_q, a_s, m_indices, inv_perm = test_moe_permute(a_q, a_s, topk_ids,
402 + num_groups, topk, block_m)
403 +
404 + inter_out = torch.zeros((a_q.shape[0], N * 2),
405 + dtype=torch.bfloat16,
406 + device=a.device)
407 +
408 + deep_gemm.m_grouped_gemm_fp8_fp8_bf16_nt_contiguous((a_q, a_s), (w1, w1_s),
409 + inter_out, m_indices)
410 +
411 + act_out = SiluAndMul().forward_native(inter_out)
412 + act_out_q, act_out_s = per_token_group_quant_fp8(act_out, block_k)
413 +
414 + out = torch.zeros(a_q.shape[0], K, dtype=torch.bfloat16, device=a.device)
415 +
416 + deep_gemm.m_grouped_gemm_fp8_fp8_bf16_nt_contiguous(
417 + (act_out_q, act_out_s), (w2, w2_s), out, m_indices)
418 +
419 + final_out = test_moe_unpermute(out, inv_perm, topk, K, topk_weight)
420 +
421 + return final_out
422 +
423 +
424 +@pytest.mark.parametrize(
425 + "M,N,K,E,topk,seed",
426 + itertools.product(M_moe_dg, N_moe, K_moe, E, TOP_KS, SEEDS))
427 +@pytest.mark.skipif(not dg_available, reason="DeepGemm kernels not available.")
428 +@torch.inference_mode()
429 +def test_w8a8_block_fp8_deep_gemm_fused_moe(M, N, K, E, topk, seed):
430 +
431 + block_m = deep_gemm.get_m_alignment_for_contiguous_layout()
432 + block_size = [block_m, block_m]
433 + dtype = torch.bfloat16
434 +
435 + # only aligned sizes
436 + if (N % block_m != 0 or K % block_m != 0 or topk > E):
437 + pytest.skip(
438 + f"Skipping test; bad size m={M}, n={N}, k={K}, topk={topk}, E={E}")
439 +
440 + if (N <= 512):
441 + pytest.skip("Skipping N <= 512 until performance issues solved.")
442 +
443 + vllm_config = VllmConfig()
444 +
445 + torch.manual_seed(seed)
446 + fp8_info = torch.finfo(torch.float8_e4m3fn)
447 + fp8_max, fp8_min = fp8_info.max, fp8_info.min
448 +
449 + a = torch.randn((M, K), dtype=dtype) / 10
450 +
451 + w1_bf16 = ((torch.rand((E, 2 * N, K), dtype=torch.bfloat16) - 0.5) * 2 *
452 + fp8_max).clamp(min=fp8_min, max=fp8_max)
453 +
454 + w2_bf16 = ((torch.rand((E, K, N), dtype=torch.bfloat16) - 0.5) * 2 *
455 + fp8_max).clamp(min=fp8_min, max=fp8_max)
456 +
457 + score = torch.randn((M, E), dtype=dtype)
458 +
459 + block_n, block_k = block_size[0], block_size[1]
460 + n_tiles_w1 = ((2 * N) + block_n - 1) // block_n
461 + k_tiles_w1 = (K + block_k - 1) // block_k
462 + n_tiles_w2 = (K + block_n - 1) // block_n
463 + k_tiles_w2 = (N + block_k - 1) // block_k
464 +
465 + w1 = torch.empty_like(w1_bf16, dtype=torch.float8_e4m3fn)
466 + w2 = torch.empty_like(w2_bf16, dtype=torch.float8_e4m3fn)
467 +
468 + w1_s = torch.empty((E, n_tiles_w1, k_tiles_w1), dtype=torch.float32)
469 + w2_s = torch.empty((E, n_tiles_w2, k_tiles_w2), dtype=torch.float32)
470 +
471 + w1_s = deep_gemm.get_col_major_tma_aligned_tensor(w1_s).contiguous()
472 + w2_s = deep_gemm.get_col_major_tma_aligned_tensor(w2_s).contiguous()
473 +
474 + assert w1_s.shape == (E, (2 * N + 127) // 128, (K + 127) // 128)
475 + assert (w2.shape[-2] + block_n - 1) // block_n == w2_s.shape[-2]
476 +
477 + for i in range(E):
478 + w1[i], w1_s[i] = per_block_cast_to_fp8(w1_bf16[i])
479 + w2[i], w2_s[i] = per_block_cast_to_fp8(w2_bf16[i])
480 +
481 + # Set the context to avoid lots of warning spam.
482 + with set_current_vllm_config(vllm_config):
483 + if M >= 128:
484 + ref_out = deep_gemm_w8a8_block_fp8_moe(M, K, a, w1, w2, w1_s, w2_s,
485 + score, topk, block_size)
486 + else:
487 + ref_out = torch_w8a8_block_fp8_moe(a, w1, w2, w1_s, w2_s, score,
488 + topk, block_size)
489 +
490 + topk_weights, topk_ids = fused_topk(a, score.float(), topk, False)
491 +
492 + out = deep_gemm_moe_fp8(a, w1, w2, w1_s, w2_s, topk_weights, topk_ids)
493 +
494 + #print(f"{out.sum()=}")
495 + #print(f"{ref_out.sum()=}")
496
497 rel_diff = (torch.mean(
498 torch.abs(out.to(torch.float32) - ref_out.to(torch.float32))) /
499 torch.mean(torch.abs(ref_out.to(torch.float32))))
500 +
501 assert rel_diff < 0.03
vllm/_custom_ops.py
+1 -1
@@ -1224,7 +1224,7 @@ def moe_wna16_gemm(input: torch.Tensor, output: torch.Tensor,
1224
1225 def topk_softmax(topk_weights: torch.Tensor, topk_ids: torch.Tensor,
1226 token_expert_indicies: torch.Tensor,
1227 - gating_output: float) -> None:
1227 + gating_output: torch.Tensor) -> None:
1228 torch.ops._moe_C.topk_softmax(topk_weights, topk_ids,
1229 token_expert_indicies, gating_output)
1230
vllm/envs.py
+5
@@ -105,6 +105,7 @@ if TYPE_CHECKING:
105 VLLM_V0_USE_OUTLINES_CACHE: bool = False
106 VLLM_TPU_DISABLE_TOPK_TOPP_OPTIMIZATION: bool = False
107 VLLM_TPU_BUCKET_PADDING_GAP: int = 0
108 + VLLM_USE_DEEP_GEMM: bool = False
109
110
111 def get_default_cache_root():
@@ -687,6 +688,10 @@ environment_variables: dict[str, Callable[[], Any]] = {
688 "VLLM_TPU_BUCKET_PADDING_GAP":
689 lambda: int(os.environ["VLLM_TPU_BUCKET_PADDING_GAP"])
690 if "VLLM_TPU_BUCKET_PADDING_GAP" in os.environ else 0,
691 +
692 + # Allow use of DeepGemm kernels for fused moe ops.
693 + "VLLM_USE_DEEP_GEMM":
694 + lambda: bool(int(os.getenv("VLLM_USE_DEEP_GEMM", "0"))),
695 }
696
697 # end-env-vars-definition
vllm/model_executor/layers/fused_moe/fused_moe.py
+415 -53
@@ -1,8 +1,10 @@
1 # SPDX-License-Identifier: Apache-2.0
2 """Fused MoE kernel."""
3 import functools
4 +import importlib.util
5 import json
6 import os
7 +from math import prod
8 from typing import Any, Callable, Dict, List, Optional, Tuple
9
10 import torch
@@ -15,7 +17,7 @@ from vllm.logger import init_logger
17 from vllm.model_executor.layers.quantization.utils.fp8_utils import (
18 per_token_group_quant_fp8)
19 from vllm.platforms import current_platform
18 -from vllm.utils import direct_register_custom_op
20 +from vllm.utils import direct_register_custom_op, round_up
21
22 from .rocm_aiter_fused_moe import (is_rocm_aiter_moe_enabled,
23 rocm_aiter_fused_experts,
@@ -23,6 +25,8 @@ from .rocm_aiter_fused_moe import (is_rocm_aiter_moe_enabled,
25
26 logger = init_logger(__name__)
27
28 +has_deep_gemm = importlib.util.find_spec("deep_gemm") is not None
29 +
30
31 @triton.jit
32 def write_zeros_to_output(c_ptr, stride_cm, stride_cn, pid_n, N, offs_token,
@@ -581,7 +585,8 @@ def moe_align_block_size(
585 topk_ids: torch.Tensor,
586 block_size: int,
587 num_experts: int,
584 - expert_map: torch.Tensor = None
588 + expert_map: Optional[torch.Tensor] = None,
589 + pad_sorted_ids: bool = False
590 ) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
591 """
592 Aligns the token distribution across experts to be compatible with block
@@ -596,6 +601,8 @@ def moe_align_block_size(
601 from the global space to the local index space of the current
602 expert parallel shard. If the expert is not in the current expert
603 parallel shard, the mapping is set to -1.
604 + - pad_sorted_ids: A flag indicating whether the sorted_token_ids length
605 + should be padded to a multiple of block_size,
606
607 Returns:
608 - sorted_token_ids: A tensor containing the sorted token indices according
@@ -625,6 +632,8 @@ def moe_align_block_size(
632 by block_size for proper block matrix operations.
633 """
634 max_num_tokens_padded = topk_ids.numel() + num_experts * (block_size - 1)
635 + if pad_sorted_ids:
636 + max_num_tokens_padded = round_up(max_num_tokens_padded, block_size)
637 sorted_ids = torch.empty((max_num_tokens_padded, ),
638 dtype=torch.int32,
639 device=topk_ids.device)
@@ -667,6 +676,59 @@ def moe_align_block_size(
676 return sorted_ids, expert_ids, num_tokens_post_pad
677
678
679 +def _valid_deep_gemm(hidden_states: torch.Tensor, w1: torch.Tensor,
680 + w2: torch.Tensor,
681 + expert_map: Optional[torch.Tensor]) -> bool:
682 + """
683 + Check if the given problem size is supported by the DeepGemm grouped
684 + gemm kernel. All of M, N, K and the quantization block_shape must be
685 + aligned by `dg.get_m_alignment_for_contiguous_layout()`.
686 + """
687 + if not has_deep_gemm:
688 + return False
689 +
690 + # Lazy import to avoid CUDA initialization problems.
691 + import deep_gemm as dg
692 +
693 + # Expert maps not supported yet.
694 + if expert_map is not None:
695 + return False
696 +
697 + align = dg.get_m_alignment_for_contiguous_layout()
698 + M = hidden_states.shape[0]
699 + _, K, N = w2.shape
700 +
701 + # For now, disable DeepGemm for small N until better permute/unpermute
702 + # ops are available.
703 + if N <= 512:
704 + return False
705 +
706 + if align > M or N % align != 0 or K % align != 0:
707 + return False
708 +
709 + return (hidden_states.is_contiguous() and w1.is_contiguous()
710 + and w2.is_contiguous())
711 +
712 +
713 +def _fp8_quantize(
714 + A: torch.Tensor,
715 + A_scale: Optional[torch.Tensor],
716 + block_shape: Optional[List[int]],
717 +) -> Tuple[torch.Tensor, torch.Tensor]:
718 + """
719 + Perform fp8 quantization on the inputs. If a block_shape
720 + is provided, the output will be blocked.
721 + """
722 + if block_shape is None:
723 + A, A_scale = ops.scaled_fp8_quant(A, A_scale)
724 + else:
725 + assert len(block_shape) == 2
726 + _, block_k = block_shape[0], block_shape[1]
727 + A, A_scale = per_token_group_quant_fp8(A, block_k)
728 + assert triton.cdiv(A.shape[-1], block_k) == A_scale.shape[-1]
729 + return A, A_scale
730 +
731 +
732 def invoke_fused_moe_kernel(A: torch.Tensor,
733 B: torch.Tensor,
734 C: torch.Tensor,
@@ -691,15 +753,11 @@ def invoke_fused_moe_kernel(A: torch.Tensor,
753
754 if use_fp8_w8a8:
755 assert B_scale is not None
694 - if block_shape is None:
695 - A, A_scale = ops.scaled_fp8_quant(A, A_scale)
696 - else:
697 - assert len(block_shape) == 2
698 - block_n, block_k = block_shape[0], block_shape[1]
699 - A, A_scale = per_token_group_quant_fp8(A, block_k)
700 - assert triton.cdiv(A.shape[-1], block_k) == A_scale.shape[-1]
701 - assert triton.cdiv(B.shape[-2], block_n) == B_scale.shape[-2]
702 - assert triton.cdiv(B.shape[-1], block_k) == B_scale.shape[-1]
756 + assert (block_shape is None or triton.cdiv(B.shape[-2], block_shape[0])
757 + == B_scale.shape[-2])
758 + assert (block_shape is None or triton.cdiv(B.shape[-1], block_shape[1])
759 + == B_scale.shape[-1])
760 +
761 elif use_int8_w8a16 or use_int4_w4a16:
762 assert B_scale is not None
763 assert block_shape is None or block_shape[0] == 0
@@ -1066,7 +1124,7 @@ def fused_topk(
1124 gating_output: torch.Tensor,
1125 topk: int,
1126 renormalize: bool,
1069 -):
1127 +) -> Tuple[torch.Tensor, torch.Tensor]:
1128 assert hidden_states.shape[0] == gating_output.shape[0], (
1129 "Number of tokens mismatch")
1130
@@ -1098,14 +1156,16 @@ def fused_topk(
1156
1157 # This is used by the Deepseek-V2 and Deepseek-V3 model
1158 @torch.compile(dynamic=True, backend=current_platform.simple_compile_backend)
1101 -def grouped_topk(hidden_states: torch.Tensor,
1102 - gating_output: torch.Tensor,
1103 - topk: int,
1104 - renormalize: bool,
1105 - num_expert_group: int = 0,
1106 - topk_group: int = 0,
1107 - scoring_func: str = "softmax",
1108 - e_score_correction_bias: Optional[torch.Tensor] = None):
1159 +def grouped_topk(
1160 + hidden_states: torch.Tensor,
1161 + gating_output: torch.Tensor,
1162 + topk: int,
1163 + renormalize: bool,
1164 + num_expert_group: int = 0,
1165 + topk_group: int = 0,
1166 + scoring_func: str = "softmax",
1167 + e_score_correction_bias: Optional[torch.Tensor] = None
1168 +) -> Tuple[torch.Tensor, torch.Tensor]:
1169
1170 assert hidden_states.shape[0] == gating_output.shape[0], (
1171 "Number of tokens mismatch")
@@ -1154,10 +1214,11 @@ def grouped_topk(hidden_states: torch.Tensor,
1214 return topk_weights.to(torch.float32), topk_ids.to(torch.int32)
1215
1216
1157 -def get_config_dtype_str(dtype: torch.dtype,
1158 - use_int4_w4a16: Optional[bool] = False,
1159 - use_int8_w8a16: Optional[bool] = False,
1160 - use_fp8_w8a8: Optional[bool] = False):
1217 +def get_config_dtype_str(
1218 + dtype: torch.dtype,
1219 + use_int4_w4a16: Optional[bool] = False,
1220 + use_int8_w8a16: Optional[bool] = False,
1221 + use_fp8_w8a8: Optional[bool] = False) -> Optional[str]:
1222 if use_fp8_w8a8:
1223 return "fp8_w8a8"
1224 elif use_int8_w8a16:
@@ -1318,26 +1379,123 @@ def fused_experts(hidden_states: torch.Tensor,
1379 w2_zp: Optional[torch.Tensor] = None,
1380 a1_scale: Optional[torch.Tensor] = None,
1381 a2_scale: Optional[torch.Tensor] = None,
1321 - block_shape: Optional[List[int]] = None) -> torch.Tensor:
1322 - return dispatch_fused_experts_func(inplace)(
1323 - hidden_states=hidden_states,
1324 - w1=w1,
1325 - w2=w2,
1326 - topk_weights=topk_weights,
1327 - topk_ids=topk_ids,
1328 - activation=activation,
1329 - use_fp8_w8a8=use_fp8_w8a8,
1330 - use_int8_w8a16=use_int8_w8a16,
1331 - use_int4_w4a16=use_int4_w4a16,
1332 - global_num_experts=global_num_experts,
1333 - expert_map=expert_map,
1334 - w1_scale=w1_scale,
1335 - w2_scale=w2_scale,
1336 - w1_zp=w1_zp,
1337 - w2_zp=w2_zp,
1338 - a1_scale=a1_scale,
1339 - a2_scale=a2_scale,
1340 - block_shape=block_shape)
1382 + block_shape: Optional[List[int]] = None,
1383 + allow_deep_gemm: bool = False) -> torch.Tensor:
1384 + if (allow_deep_gemm and use_fp8_w8a8
1385 + and _valid_deep_gemm(hidden_states, w1, w2, expert_map)):
1386 + return deep_gemm_moe_fp8(
1387 + hidden_states=hidden_states,
1388 + w1=w1,
1389 + w2=w2,
1390 + topk_weights=topk_weights,
1391 + topk_ids=topk_ids,
1392 + activation=activation,
1393 + global_num_experts=global_num_experts,
1394 + expert_map=expert_map,
1395 + w1_scale=w1_scale,
1396 + w2_scale=w2_scale,
1397 + a1_scale=a1_scale,
1398 + a2_scale=a2_scale,
1399 + )
1400 + else:
1401 + return dispatch_fused_experts_func(inplace)(
1402 + hidden_states=hidden_states,
1403 + w1=w1,
1404 + w2=w2,
1405 + topk_weights=topk_weights,
1406 + topk_ids=topk_ids,
1407 + activation=activation,
1408 + use_fp8_w8a8=use_fp8_w8a8,
1409 + use_int8_w8a16=use_int8_w8a16,
1410 + use_int4_w4a16=use_int4_w4a16,
1411 + global_num_experts=global_num_experts,
1412 + expert_map=expert_map,
1413 + w1_scale=w1_scale,
1414 + w2_scale=w2_scale,
1415 + w1_zp=w1_zp,
1416 + w2_zp=w2_zp,
1417 + a1_scale=a1_scale,
1418 + a2_scale=a2_scale,
1419 + block_shape=block_shape)
1420 +
1421 +
1422 +def _fp8_perm(m: torch.Tensor, idx: torch.Tensor) -> torch.Tensor:
1423 + """
1424 + A permutation routine that works on fp8 types.
1425 + """
1426 + if torch.is_floating_point(m) and torch.finfo(m.dtype).bits == 8:
1427 + return m.view(dtype=torch.uint8)[idx, ...].view(dtype=m.dtype)
1428 + else:
1429 + return m[idx, ...]
1430 +
1431 +
1432 +def _moe_permute(
1433 + curr_hidden_states: torch.Tensor,
1434 + a1q_scale: Optional[torch.Tensor],
1435 + curr_topk_ids: torch.Tensor,
1436 + global_num_experts: int,
1437 + expert_map: Optional[torch.Tensor],
1438 + top_k_num: int,
1439 + block_m: int,
1440 +) -> Tuple[torch.Tensor, Optional[torch.Tensor], torch.Tensor, torch.Tensor,
1441 + torch.Tensor]:
1442 + """
1443 + Determine the sorted_token_ids, expert_ids for the given problem size.
1444 + Permute the hidden states and scales according to `sorted_token_ids`.
1445 + """
1446 + tokens_in_chunk, _ = curr_hidden_states.shape
1447 +
1448 + sorted_token_ids, expert_ids, num_tokens_post_padded = (
1449 + moe_align_block_size(curr_topk_ids,
1450 + block_m,
1451 + global_num_experts,
1452 + expert_map,
1453 + pad_sorted_ids=True))
1454 +
1455 + inv_perm: Optional[torch.Tensor] = None
1456 +
1457 + num_tokens = top_k_num * tokens_in_chunk
1458 + sorted_token_ids = sorted_token_ids.clamp(max=num_tokens - 1)
1459 + expert_ids = torch.repeat_interleave(expert_ids, block_m, dim=0)
1460 + inv_perm = torch.argsort(sorted_token_ids)[:num_tokens]
1461 +
1462 + # Permute according to sorted token ids.
1463 + curr_hidden_states = _fp8_perm(curr_hidden_states,
1464 + sorted_token_ids // top_k_num)
1465 +
1466 + if a1q_scale is not None:
1467 + a1q_scale = a1q_scale[sorted_token_ids // top_k_num]
1468 +
1469 + return (curr_hidden_states, a1q_scale, sorted_token_ids, expert_ids,
1470 + inv_perm)
1471 +
1472 +
1473 +def _moe_unpermute_and_reduce(
1474 + out: torch.Tensor,
1475 + curr_hidden: torch.Tensor,
1476 + inv_perm: Optional[torch.Tensor],
1477 + topk: int,
1478 + K: int,
1479 + topk_weight: torch.Tensor,
1480 +) -> None:
1481 + """
1482 + Unpermute the final result and apply topk_weights, then perform the final
1483 + reduction on the hidden states.
1484 + """
1485 + M = topk_weight.shape[0]
1486 + curr_hidden = curr_hidden[inv_perm, ...]
1487 + curr_hidden = curr_hidden.view(-1, topk, K)
1488 + curr_hidden.mul_(topk_weight.view(M, -1, 1))
1489 + ops.moe_sum(curr_hidden, out)
1490 +
1491 +
1492 +def _resize_cache(x: torch.Tensor, v: Tuple[int, ...]) -> torch.Tensor:
1493 + """
1494 + Shrink the given tensor and apply the given view to it. This is
1495 + used to resize the intermediate fused_moe caches.
1496 + """
1497 + assert prod(v) <= x.numel()
1498 + return x.flatten()[:prod(v)].view(*v)
1499
1500
1501 def fused_experts_impl(hidden_states: torch.Tensor,
@@ -1376,6 +1534,7 @@ def fused_experts_impl(hidden_states: torch.Tensor,
1534
1535 num_tokens, _ = hidden_states.shape
1536 E, N, _ = w1.shape
1537 + K = w2.shape[1]
1538 if global_num_experts == -1:
1539 global_num_experts = E
1540 top_k_num = topk_ids.shape[1]
@@ -1401,13 +1560,11 @@ def fused_experts_impl(hidden_states: torch.Tensor,
1560
1561 # We can reuse the memory between these because by the time we need
1562 # cache3, we're done with cache1
1404 - cache13 = torch.empty(M * top_k_num * max(N, w2.shape[1]),
1563 + cache13 = torch.empty(M * top_k_num * max(N, K),
1564 device=hidden_states.device,
1565 dtype=hidden_states.dtype)
1407 - intermediate_cache1 = cache13[:M * top_k_num * N].view(
1408 - (M, topk_ids.shape[1], N))
1409 - intermediate_cache3 = cache13[:M * top_k_num * w2.shape[1]].view(
1410 - (M, topk_ids.shape[1], w2.shape[1]))
1566 + intermediate_cache1 = cache13[:M * top_k_num * N].view(M, top_k_num, N)
1567 + intermediate_cache3 = cache13[:M * top_k_num * K].view(M, top_k_num, K)
1568
1569 # This needs separate memory since it's used concurrently with cache1
1570 intermediate_cache2 = torch.empty((M * top_k_num, N // 2),
@@ -1452,14 +1609,23 @@ def fused_experts_impl(hidden_states: torch.Tensor,
1609 curr_topk_ids = topk_ids[begin_chunk_idx:end_chunk_idx]
1610 curr_topk_weights = topk_weights[begin_chunk_idx:end_chunk_idx]
1611
1612 + a1q_scale: Optional[torch.Tensor] = None
1613 +
1614 + if use_fp8_w8a8:
1615 + qcurr_hidden_states, a1q_scale = _fp8_quantize(
1616 + curr_hidden_states, a1_scale, block_shape)
1617 + else:
1618 + qcurr_hidden_states = curr_hidden_states
1619 + a1q_scale = a1_scale
1620 +
1621 sorted_token_ids, expert_ids, num_tokens_post_padded = (
1622 moe_align_block_size(curr_topk_ids, config['BLOCK_SIZE_M'],
1623 global_num_experts, expert_map))
1624
1459 - invoke_fused_moe_kernel(curr_hidden_states,
1625 + invoke_fused_moe_kernel(qcurr_hidden_states,
1626 w1,
1627 intermediate_cache1,
1462 - a1_scale,
1628 + a1q_scale,
1629 w1_scale,
1630 w1_zp,
1631 curr_topk_weights,
@@ -1485,10 +1651,19 @@ def fused_experts_impl(hidden_states: torch.Tensor,
1651 else:
1652 raise ValueError(f"Unsupported FusedMoe activation: {activation}")
1653
1488 - invoke_fused_moe_kernel(intermediate_cache2,
1654 + a2q_scale: Optional[torch.Tensor] = None
1655 +
1656 + if use_fp8_w8a8:
1657 + qintermediate_cache2, a2q_scale = _fp8_quantize(
1658 + intermediate_cache2, a2_scale, block_shape)
1659 + else:
1660 + qintermediate_cache2 = intermediate_cache2
1661 + a2q_scale = a2_scale
1662 +
1663 + invoke_fused_moe_kernel(qintermediate_cache2,
1664 w2,
1665 intermediate_cache3,
1491 - a2_scale,
1666 + a2q_scale,
1667 w2_scale,
1668 w2_zp,
1669 curr_topk_weights,
@@ -1617,6 +1792,193 @@ def fused_moe(
1792 block_shape=block_shape)
1793
1794
1795 +def deep_gemm_moe_fp8(
1796 + hidden_states: torch.Tensor,
1797 + w1: torch.Tensor,
1798 + w2: torch.Tensor,
1799 + w1_scale: torch.Tensor,
1800 + w2_scale: torch.Tensor,
1801 + topk_weights: torch.Tensor,
1802 + topk_ids: torch.Tensor,
1803 + inplace: bool = False,
1804 + activation: str = "silu",
1805 + global_num_experts: int = -1,
1806 + expert_map: Optional[torch.Tensor] = None,
1807 + a1_scale: Optional[torch.Tensor] = None,
1808 + a2_scale: Optional[torch.Tensor] = None,
1809 +) -> torch.Tensor:
1810 + """
1811 + This function computes a a8w8-quantized Mixture of Experts (MoE) layer
1812 + using two sets of quantized weights, w1_q and w2_q, and top-k gating
1813 + mechanism. The matrix multiplications are implemented with DeepGemm
1814 + grouped gemm.
1815 +
1816 + Parameters:
1817 + - hidden_states (torch.Tensor): The input tensor to the MoE layer.
1818 + Shape: [M, K]
1819 + - w1 (torch.Tensor): The first set of fp8 quantized expert weights.
1820 + Shape: [num_experts, K, 2N] (the weights are passed transposed)
1821 + - w2 (torch.Tensor): The second set of fp8 quantized expert weights.
1822 + Shape: [num_experts, N, K] (the weights are passed transposed)
1823 + - w1_scale (torch.Tensor): The fp32 scale to dequantize w1_q.
1824 + Shape: [num_experts] or [num_experts, 2N]
1825 + - w2_scale (torch.Tensor): The fp32 scale to dequantize w2_q.
1826 + Shape: [num_experts] or [num_experts, K]
1827 + - topk_weights (torch.Tensor): The weights of each token->expert mapping.
1828 + - topk_ids (torch.Tensor): The token->expert mapping for topk_weights.
1829 + - inplace (bool): If True, perform the operation in-place.
1830 + Defaults to False.
1831 + - activation (str): The activation function to apply after the first
1832 + MoE layer.
1833 + - global_num_experts (int): The total number of experts in the global
1834 + expert space.
1835 + - expert_map (Optional[torch.Tensor]): A tensor mapping expert indices
1836 + from the global expert space to the local expert space of the expert
1837 + parallel shard.
1838 + - a1_scale (Optional[torch.Tensor]): The optional fp32 scale to quantize a.
1839 + Shape: scalar or [M]
1840 + - a2_scale (Optional[torch.Tensor]): The optional fp32 scale to
1841 + quantize the intermediate result between the gemms.
1842 + Shape: scalar or [M]
1843 +
1844 + Returns:
1845 + - torch.Tensor: The bfloat16 output tensor after applying the MoE layer.
1846 + """
1847 + # Lazy import to avoid CUDA initialization problems.
1848 + import deep_gemm as dg
1849 +
1850 + assert expert_map is None, "Expert maps not supported yet"
1851 +
1852 + assert hidden_states.shape[1] == w1.shape[2], "Hidden size mismatch"
1853 +
1854 + assert topk_weights.shape == topk_ids.shape, "topk shape mismatch"
1855 + assert hidden_states.is_contiguous(), "Hidden_states must be contiguous"
1856 + assert w1.stride(-1) == 1, "Stride of last dimension must be 1"
1857 + assert w2.stride(-1) == 1, "Stride of last dimension must be 1"
1858 + assert hidden_states.dtype in [
1859 + torch.float32, torch.float16, torch.bfloat16
1860 + ]
1861 + assert w1.dtype == torch.float8_e4m3fn
1862 + assert w2.dtype == torch.float8_e4m3fn
1863 + assert w1.shape[0] == w2.shape[0], "Expert number mismatch"
1864 + assert w1.shape[0] == w1_scale.shape[0], "w1 scales expert number mismatch"
1865 + assert w1.shape[0] == w2_scale.shape[0], "w2 scales expert number mismatch"
1866 + assert a1_scale is None or a1_scale.dim(
1867 + ) == 0 or a1_scale.shape[0] == 1 or a1_scale.shape[
1868 + 0] == hidden_states.shape[0], "Input scale shape mismatch"
1869 + assert a2_scale is None or a1_scale is None or a2_scale.shape == a1_scale.shape, "Intermediate scale shape mismatch" # noqa: E501
1870 +
1871 + num_tokens, _ = hidden_states.shape
1872 + E, N, _ = w1.shape
1873 + K = w2.shape[1]
1874 + if global_num_experts == -1:
1875 + global_num_experts = E
1876 + top_k_num = topk_ids.shape[1]
1877 + # We execute the fused_moe kernel in chunks to circumvent this issue:
1878 + # https://github.com/vllm-project/vllm/issues/5938
1879 + CHUNK_SIZE = envs.VLLM_FUSED_MOE_CHUNK_SIZE
1880 +
1881 + assert _valid_deep_gemm(hidden_states, w1, w2, expert_map)
1882 +
1883 + if inplace:
1884 + out_hidden_states = hidden_states
1885 + else:
1886 + out_hidden_states = torch.empty_like(hidden_states)
1887 +
1888 + block_m = dg.get_m_alignment_for_contiguous_layout()
1889 + block_shape = [block_m, block_m]
1890 +
1891 + assert w1_scale is not None
1892 + assert w2_scale is not None
1893 +
1894 + # We attempt to transpose and align offline in Fp8MoEMethod, in which
1895 + # case these calls will be nops. Otherwise, they'll be performed every
1896 + # time the layer is executed.
1897 + w1_scale = dg.get_col_major_tma_aligned_tensor(w1_scale).contiguous()
1898 + w2_scale = dg.get_col_major_tma_aligned_tensor(w2_scale).contiguous()
1899 +
1900 + M_sum = topk_ids.numel() + global_num_experts * (block_m - 1)
1901 + M_sum = round_up(M_sum, block_m)
1902 +
1903 + num_chunks = (num_tokens // CHUNK_SIZE) + 1
1904 +
1905 + # We can reuse the memory between cache1 and cache3 because by the time
1906 + # we need cache3, we're done with cache1
1907 + cache13 = torch.empty(M_sum * max(N, K),
1908 + device=hidden_states.device,
1909 + dtype=hidden_states.dtype)
1910 +
1911 + intermediate_cache1 = cache13[:M_sum * N].view(M_sum, N)
1912 + intermediate_cache2 = torch.empty((M_sum, N // 2),
1913 + device=hidden_states.device,
1914 + dtype=hidden_states.dtype)
1915 + intermediate_cache3 = cache13[:M_sum * K].view(M_sum, K)
1916 +
1917 + for chunk in range(num_chunks):
1918 + begin_chunk_idx, end_chunk_idx = (chunk * CHUNK_SIZE,
1919 + min((chunk + 1) * CHUNK_SIZE,
1920 + num_tokens))
1921 + curr_hidden_states = hidden_states[begin_chunk_idx:end_chunk_idx]
1922 + tokens_in_chunk, _ = curr_hidden_states.shape
1923 +
1924 + if tokens_in_chunk == 0:
1925 + break
1926 +
1927 + curr_topk_ids = topk_ids[begin_chunk_idx:end_chunk_idx]
1928 + curr_topk_weights = topk_weights[begin_chunk_idx:end_chunk_idx]
1929 +
1930 + a1q_scale: Optional[torch.Tensor] = None
1931 +
1932 + qcurr_hidden_states, a1q_scale = _fp8_quantize(curr_hidden_states,
1933 + a1_scale, block_shape)
1934 +
1935 + (qcurr_hidden_states, a1q_scale, sorted_token_ids, expert_ids,
1936 + inv_perm) = _moe_permute(qcurr_hidden_states, a1q_scale,
1937 + curr_topk_ids, global_num_experts,
1938 + expert_map, top_k_num, block_m)
1939 +
1940 + # Adjust the intermediate cache size and config for the last chunk.
1941 + # Note that in most cases we only have one chunk so the cache size
1942 + # and config are already set correctly and do not need to be adjusted.
1943 + if tokens_in_chunk < CHUNK_SIZE and chunk > 0:
1944 + curr_M = sorted_token_ids.numel()
1945 + intermediate_cache1 = _resize_cache(intermediate_cache1,
1946 + (curr_M, N))
1947 + intermediate_cache2 = _resize_cache(intermediate_cache2,
1948 + (curr_M, N // 2))
1949 + intermediate_cache3 = _resize_cache(intermediate_cache3,
1950 + (curr_M, K))
1951 +
1952 + dg.m_grouped_gemm_fp8_fp8_bf16_nt_contiguous(
1953 + (qcurr_hidden_states, a1q_scale), (w1, w1_scale),
1954 + intermediate_cache1, expert_ids)
1955 +
1956 + if activation == "silu":
1957 + torch.ops._C.silu_and_mul(intermediate_cache2,
1958 + intermediate_cache1.view(-1, N))
1959 + elif activation == "gelu":
1960 + torch.ops._C.gelu_and_mul(intermediate_cache2,
1961 + intermediate_cache1.view(-1, N))
1962 + else:
1963 + raise ValueError(f"Unsupported FusedMoe activation: {activation}")
1964 +
1965 + a2q_scale: Optional[torch.Tensor] = None
1966 +
1967 + qintermediate_cache2, a2q_scale = _fp8_quantize(
1968 + intermediate_cache2, a2_scale, block_shape)
1969 +
1970 + dg.m_grouped_gemm_fp8_fp8_bf16_nt_contiguous(
1971 + (qintermediate_cache2, a2q_scale), (w2, w2_scale),
1972 + intermediate_cache3, expert_ids)
1973 +
1974 + _moe_unpermute_and_reduce(
1975 + out_hidden_states[begin_chunk_idx:end_chunk_idx],
1976 + intermediate_cache3.view(*intermediate_cache3.shape), inv_perm,
1977 + top_k_num, K, curr_topk_weights)
1978 +
1979 + return out_hidden_states
1980 +
1981 +
1982 #TODO make the grouped gemm kernel consistent with scaled gemm kernel
1983 def cutlass_moe_fp8(
1984 a: torch.Tensor,
vllm/model_executor/layers/quantization/fp8.py
+36
@@ -1,5 +1,6 @@
1 # SPDX-License-Identifier: Apache-2.0
2
3 +import importlib.util
4 from typing import Any, Callable, Dict, List, Optional
5
6 import torch
@@ -37,6 +38,14 @@ ACTIVATION_SCHEMES = ["static", "dynamic"]
38
39 logger = init_logger(__name__)
40
41 +has_deep_gemm = importlib.util.find_spec("deep_gemm") is not None
42 +
43 +
44 +def _is_col_major(x: torch.Tensor) -> bool:
45 + assert x.dim() == 3
46 + b, m, n = x.shape
47 + return x.stride(0) == m * n and x.stride(1) == 1 and x.stride(2) == m
48 +
49
50 class Fp8Config(QuantizationConfig):
51 """Config class for FP8."""
@@ -424,6 +433,19 @@ class Fp8MoEMethod(FusedMoEMethodBase):
433 self.quant_config = quant_config
434 self.block_quant = self.quant_config.weight_block_size is not None
435
436 + # Check for DeepGemm support.
437 + self.allow_deep_gemm = False
438 + if envs.VLLM_USE_DEEP_GEMM:
439 + if not has_deep_gemm:
440 + logger.warning_once("Failed to import DeepGemm kernels.")
441 + elif (current_platform.is_cuda()
442 + and current_platform.has_device_capability(90)):
443 + logger.info_once("Using DeepGemm kernels for Fp8MoEMethod.")
444 + self.allow_deep_gemm = True
445 + else:
446 + logger.warning_once(
447 + "DeepGemm not supported on the current platform.")
448 +
449 def create_weights(self, layer: Module, num_experts: int, hidden_size: int,
450 intermediate_size_per_partition: int,
451 params_dtype: torch.dtype, **extra_weight_attrs):
@@ -585,6 +607,19 @@ class Fp8MoEMethod(FusedMoEMethodBase):
607 requires_grad=False)
608 layer.w2_weight = torch.nn.Parameter(shuffled_w2,
609 requires_grad=False)
610 +
611 + # DeepGemm scales need to be transposed and aligned. We try to do
612 + # it ahead of time for performance reasons.
613 + if self.allow_deep_gemm:
614 + # Lazy import to avoid CUDA initialization problems.
615 + import deep_gemm as dg
616 + if _is_col_major(layer.w13_weight_scale_inv):
617 + layer.w13_weight_scale_inv = \
618 + dg.get_col_major_tma_aligned_tensor(layer.w13_weight_scale_inv).contiguous()
619 + if _is_col_major(layer.w2_weight_scale_inv):
620 + layer.w2_weight_scale_inv = \
621 + dg.get_col_major_tma_aligned_tensor(layer.w2_weight_scale_inv).contiguous()
622 +
623 return
624
625 # If checkpoint is fp16, quantize in place.
@@ -773,6 +808,7 @@ class Fp8MoEMethod(FusedMoEMethodBase):
808 a1_scale=layer.w13_input_scale,
809 a2_scale=layer.w2_input_scale,
810 block_shape=self.quant_config.weight_block_size,
811 + allow_deep_gemm=self.allow_deep_gemm,
812 )
813
814