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

[Performance][B200] silu_mul_quant: pack scales in int32 (#28358)

Signed-off-by: Varun Sundar Rabindranath <vsundarr@redhat.com> Co-authored-by: Varun Sundar Rabindranath <vsundarr@redhat.com>

Varun Sundar Rabindranath committed Nov 13, 2025 at 13:16 UTC fe1cd7704ddd3266ddc97181ab24a167b3c9223c
7 files changed +461 -146
csrc/quantization/activation_kernels.cu
+121 -43
@@ -279,17 +279,17 @@ __device__ __forceinline__ void token_bounds(int32_t n_tokens,
279 }
280
281 template <int BLOCK_COUNT, int SMEM_SIZE_BYTES_Y, typename fp8_type,
282 - int THREADS, typename Idx_t, bool USE_UE8M0, int GROUP_SIZE = 128,
283 - int NUM_STAGES = 3>
282 + typename scale_t, int THREADS, typename Idx_t, bool CEIL_UE8M0,
283 + int GROUP_SIZE = 128, int NUM_STAGES = 3>
284 __global__ void silu_mul_fp8_quant_deep_gemm_kernel(
285 const __nv_bfloat16* __restrict__ _input, fp8_type* __restrict__ _y_q,
286 - float* __restrict__ _y_s, const int32_t* __restrict__ tokens_per_expert,
286 + scale_t* __restrict__ _y_s, const int32_t* __restrict__ tokens_per_expert,
287 // sizes
288 Idx_t E, Idx_t T, Idx_t H,
289 // strides (in elements)
290 Idx_t stride_i_e, Idx_t stride_i_t, Idx_t stride_i_h, Idx_t stride_yq_e,
291 Idx_t stride_yq_t, Idx_t stride_yq_h, Idx_t stride_ys_e, Idx_t stride_ys_t,
292 - Idx_t stride_ys_g, Idx_t stride_counts_e) {
292 + Idx_t stride_ys_g, Idx_t stride_ys_p, Idx_t stride_counts_e) {
293 #ifndef USE_ROCM
294 static constexpr int NUM_WARPS = THREADS / WARP_SIZE;
295
@@ -466,9 +466,22 @@ __global__ void silu_mul_fp8_quant_deep_gemm_kernel(
466
467 __nv_fp8x4_e4m3* y_q_base_ptr =
468 reinterpret_cast<__nv_fp8x4_e4m3*>(_y_q) + lane_id;
469 - auto y_scale_base_ptr = _y_s + warp_position_scales * stride_ys_g;
469 +
470 + Idx_t scale_group_offset = 0;
471 + if constexpr (std::is_same<scale_t, uint8_t>::value) {
472 + // packed int32_t format
473 + int pack_id = warp_position_scales / 4;
474 + int scale_in_pack = warp_position_scales % 4;
475 + scale_group_offset = pack_id * stride_ys_p + scale_in_pack * stride_ys_g;
476 + } else {
477 + scale_group_offset = warp_position_scales * stride_ys_g;
478 + }
479 +
480 + scale_t* const y_scale_base_ptr = _y_s + scale_group_offset;
481
482 for (auto j = tokens_lower; j < tokens_upper; j++) {
483 + int current_group_id = warp_position_scales; // Running count of which
484 + // group is being processed
485 const Idx_t base_ys = expert_id * stride_ys_e;
486 auto y_s_ptr = y_scale_base_ptr + base_ys + token_offset * stride_ys_t;
487 __nv_fp8x4_e4m3* y_q_ptr =
@@ -509,7 +522,7 @@ __global__ void silu_mul_fp8_quant_deep_gemm_kernel(
522
523 __nv_bfloat16 y_s = __hmul(warp_max(_y_max2.x), fp8_inv);
524
512 - if constexpr (USE_UE8M0) {
525 + if constexpr (CEIL_UE8M0) {
526 y_s = hexp2(hceil(hlog2(y_s)));
527 }
528
@@ -527,8 +540,24 @@ __global__ void silu_mul_fp8_quant_deep_gemm_kernel(
540 y_q_ptr += WARP_SIZE * stride_yq_h;
541
542 if (!lane_id) {
530 - *y_s_ptr = y_s;
531 - y_s_ptr += stride_ys_g;
543 + // Store scales.
544 + if constexpr (std::is_same<scale_t, uint8_t>::value) {
545 + // Packed UE8MO format. Remove Mantissa.
546 + *y_s_ptr = reinterpret_cast<int16_t&>(y_s) >> 7;
547 +
548 + bool const jump_pack = (current_group_id + 1) % 4 == 0;
549 + // Minus 3 because we need to get to the first group in the
550 + // next pack.
551 + y_s_ptr += jump_pack ? (stride_ys_p - 3) : stride_ys_g;
552 +
553 + } else {
554 + // float32 format
555 + static_assert(std::is_same<scale_t, float>::value);
556 + *y_s_ptr = y_s;
557 + y_s_ptr += stride_ys_g;
558 + }
559 +
560 + current_group_id += 1;
561 }
562 }
563 }
@@ -573,7 +602,7 @@ void persistent_masked_m_silu_mul_quant(
602 const at::Tensor& tokens_per_expert, // (E)
603 at::Tensor& y_q, // (E, T, H) [OUT]
604 at::Tensor& y_s, // (E, T, H//group_size) [OUT]
576 - bool use_ue8m0) {
605 + bool cast_scale_ue8m0) {
606 #ifndef USE_ROCM
607
608 // This kernel currently only supports H % 128 == 0 and assumes a
@@ -583,9 +612,12 @@ void persistent_masked_m_silu_mul_quant(
612 TORCH_CHECK(input.dtype() == torch::kBFloat16);
613 TORCH_CHECK(y_q.dtype() == torch::kFloat8_e4m3fn ||
614 y_q.dtype() == torch::kFloat8_e4m3fnuz);
586 - TORCH_CHECK(y_s.dtype() == torch::kFloat32);
615 TORCH_CHECK(input.size(-1) % (GROUP_SIZE * 2) == 0);
616
617 + bool const is_packed_ue8m0 =
618 + (y_s.dtype() == torch::kInt32 && cast_scale_ue8m0);
619 + TORCH_CHECK(y_s.dtype() == torch::kFloat32 || is_packed_ue8m0);
620 +
621 using Idx_t = int64_t;
622
623 Idx_t E = input.size(0);
@@ -597,15 +629,18 @@ void persistent_masked_m_silu_mul_quant(
629 Idx_t stride_yq_e = y_q.stride(0);
630 Idx_t stride_yq_t = y_q.stride(1);
631 Idx_t stride_yq_h = y_q.stride(2);
600 - Idx_t stride_ys_e = y_s.stride(0);
601 - Idx_t stride_ys_t = y_s.stride(1);
602 - Idx_t stride_ys_g = y_s.stride(2);
632
633 Idx_t stride_counts_e = tokens_per_expert.stride(0);
634
635 + int const NUM_GROUPS = H / GROUP_SIZE;
636 +
637 const cudaStream_t stream = at::cuda::getCurrentCUDAStream();
638
608 - #define KERNEL(BLOCK_COUNT, USE_UE8M0, THREAD_COUNT, STAGES) \
639 + // TODO: Get this from cuda_arch ?
640 + static constexpr int SILU_V2_BLOCK_COUNT = 132 * 32;
641 +
642 + #define KERNEL(BLOCK_COUNT, scale_t, STRIDE_YS_E, STRIDE_YS_T, STRIDE_YS_G, \
643 + STRIDE_YS_P, CEIL_UE8M0, THREAD_COUNT, STAGES) \
644 static constexpr int NUM_WARPS = THREAD_COUNT / WARP_SIZE; \
645 int sms = SILU_V2_BLOCK_COUNT; \
646 static constexpr int max_shared_mem_bytes = \
@@ -615,43 +650,86 @@ void persistent_masked_m_silu_mul_quant(
650 VLLM_DISPATCH_FP8_TYPES( \
651 y_q.scalar_type(), "silu_mul_fp8_quant_deep_gemm_kernel", [&] { \
652 vllm::silu_mul_fp8_quant_deep_gemm_kernel< \
618 - BLOCK_COUNT, max_shared_mem_bytes, fp8_t, THREAD_COUNT, Idx_t, \
619 - USE_UE8M0, GROUP_SIZE, STAGES> \
653 + BLOCK_COUNT, max_shared_mem_bytes, fp8_t, scale_t, THREAD_COUNT, \
654 + Idx_t, CEIL_UE8M0, GROUP_SIZE, STAGES> \
655 <<<grid, block, max_shared_mem_bytes + (E + 1) * 16, stream>>>( \
656 reinterpret_cast<__nv_bfloat16*>(input.data_ptr()), \
622 - (fp8_t*)y_q.data_ptr(), y_s.data_ptr<float>(), \
657 + (fp8_t*)y_q.data_ptr(), \
658 + reinterpret_cast<scale_t*>(y_s.data_ptr()), \
659 reinterpret_cast<int32_t*>(tokens_per_expert.data_ptr()), E, \
660 T, H, stride_i_e, stride_i_t, stride_i_h, stride_yq_e, \
625 - stride_yq_t, stride_yq_h, stride_ys_e, stride_ys_t, \
626 - stride_ys_g, stride_counts_e); \
661 + stride_yq_t, stride_yq_h, STRIDE_YS_E, STRIDE_YS_T, \
662 + STRIDE_YS_G, STRIDE_YS_P, stride_counts_e); \
663 });
664
629 - static constexpr int SILU_V2_BLOCK_COUNT = 132 * 32;
630 -
631 - int const NUM_GROUPS = H / GROUP_SIZE;
632 - if (!use_ue8m0) {
633 - if (H >= 4096 && (NUM_GROUPS % 8 == 0)) {
634 - /* 8 warps config */
635 - static constexpr int NUM_STAGES = 4;
636 - static constexpr int THREAD_COUNT = 256;
637 - KERNEL(SILU_V2_BLOCK_COUNT, false, THREAD_COUNT, NUM_STAGES);
638 - } else {
639 - /* 1 warp config */
640 - static constexpr int THREAD_COUNT = 32;
641 - KERNEL(SILU_V2_BLOCK_COUNT, false, THREAD_COUNT, 2);
642 - }
643 - } else {
644 - if (H >= 4096 && (NUM_GROUPS % 8 == 0)) {
645 - /* 8 warps config */
646 - static constexpr int NUM_STAGES = 4;
647 - static constexpr int THREAD_COUNT = 256;
648 - KERNEL(SILU_V2_BLOCK_COUNT, true, THREAD_COUNT, NUM_STAGES);
649 - } else {
650 - /* 1 warp config */
651 - static constexpr int THREAD_COUNT = 32;
652 - KERNEL(SILU_V2_BLOCK_COUNT, true, THREAD_COUNT, 2);
665 + #define LAUNCH_ON_H(scale_t, STRIDE_YS_E, STRIDE_YS_T, STRIDE_YS_G, \
666 + STRIDE_YS_P, CEIL_UE8M0) \
667 + if (H >= 4096 && (NUM_GROUPS % 8) == 0) { \
668 + /* 8 warp config */ \
669 + static constexpr int NUM_STAGES = 4; \
670 + static constexpr int THREAD_COUNT = 256; \
671 + KERNEL(SILU_V2_BLOCK_COUNT, scale_t, STRIDE_YS_E, STRIDE_YS_T, \
672 + STRIDE_YS_G, STRIDE_YS_P, CEIL_UE8M0, THREAD_COUNT, NUM_STAGES); \
673 + } else { \
674 + /* 1 warp config */ \
675 + static constexpr int THREAD_COUNT = 32; \
676 + KERNEL(SILU_V2_BLOCK_COUNT, scale_t, STRIDE_YS_E, STRIDE_YS_T, \
677 + STRIDE_YS_G, STRIDE_YS_P, CEIL_UE8M0, THREAD_COUNT, 2); \
678 }
679 +
680 + Idx_t stride_ys_e = y_s.stride(0);
681 + Idx_t stride_ys_t = y_s.stride(1);
682 + Idx_t stride_ys_g = y_s.stride(2);
683 + Idx_t stride_ys_p = 0;
684 + if (!cast_scale_ue8m0) {
685 + TORCH_CHECK(!is_packed_ue8m0);
686 + LAUNCH_ON_H(float, stride_ys_e, stride_ys_t, stride_ys_g, stride_ys_p,
687 + false);
688 + return;
689 + }
690 +
691 + if (!is_packed_ue8m0) {
692 + // UE8M0 but not packed
693 + LAUNCH_ON_H(float, stride_ys_e, stride_ys_t, stride_ys_g, stride_ys_p,
694 + true);
695 + return;
696 }
697
698 + TORCH_CHECK(cast_scale_ue8m0 && is_packed_ue8m0);
699 + TORCH_CHECK(y_s.dtype() == torch::kInt32);
700 +
701 + // Int32 packed ue8m0 scales tensor.
702 + // Let E, T, G be the number to experts, number of tokens and number of groups
703 + // respectively. Let, E = 2, T = 4, G = 6, in this case the int32 scales
704 + // tensor are of shape [1, 4, 2] and stride [8, 1, 4]. The scales are expected
705 + // to be arranged as follows,
706 + // [[T0G0-T0G1-T0G2-T0G3, T0G4-T0G5-X-X,],
707 + // [T1G0-T1G1-T1G2-T1G3, T1G4-T1G5-X-X,]
708 + // [T2G0-T2G1-T2G2-T2G3, T2G4-T2G5-X-X,]
709 + // [T3G0-T3G1-T3G2-T3G3, T3G4-T3G5-X-X,]]
710 + // where, TxGy is the scale ue8m0 scale value of Token x, Group y.
711 + //
712 + // In memory (in bytes) the scale values are arranged as,
713 + // [T0G0, T0G1, T0G2, T0G3, T1G0, T1G2, T1G3, T1G4, T2G0, T2G1, T2G3, T2G4,
714 + // T3G0, T3G1, T3G2, T3G3, T0G4, T0G5, X, X, T1G4, T1G5, X, X, T2G4, T2G5,
715 + // X, X, T3G4, T3G5, X, X]
716 + //
717 + // An Int32 tensor of size [1, 4, 2] and stride [8, 1, 4] can be represented
718 + // as an uint8 tensor of shape [1, 2, 4, 4] and stride [32, 16, 4, 1]. In
719 + // english, ignoring the Experts dimension, the original int32 tensor is
720 + // simply treated as two packed [4, 4] uint8 tensor (or two [4, 1] int32
721 + // tensor). The following strides setting reflects this change. Caveat: This
722 + // means that the G dimension is no longer contiguous. i.e. Note that to move
723 + // from G3 to G4, we need to jump along the packing dimension. The kernel
724 + // handles this case.
725 +
726 + stride_ys_e *= sizeof(int32_t);
727 + stride_ys_p = T * sizeof(int32_t); // Packing dimension
728 + stride_ys_t = sizeof(int32_t);
729 + stride_ys_g = 1;
730 +
731 + LAUNCH_ON_H(uint8_t, stride_ys_e, stride_ys_t, stride_ys_g, stride_ys_p,
732 + true);
733 +
734 #endif
735 }
tests/conftest.py
+13
@@ -1384,3 +1384,16 @@ def image_urls(request, local_asset_server) -> list[str]:
1384 """Indirect fixture: takes a list of names, returns list of full URLs."""
1385 names: list[str] = request.param
1386 return [local_asset_server.url_for(name) for name in names]
1387 +
1388 +
1389 +@pytest.fixture
1390 +def disable_deepgemm_ue8m0(monkeypatch):
1391 + from vllm.utils.deep_gemm import is_deep_gemm_e8m0_used
1392 +
1393 + with monkeypatch.context() as monkeypatch_ctx:
1394 + monkeypatch_ctx.setenv("VLLM_USE_DEEP_GEMM_E8M0", "0")
1395 + is_deep_gemm_e8m0_used.cache_clear()
1396 + yield
1397 + # Clear cache so the next time it is used it is processed with the
1398 + # default VLLM_USE_DEEP_GEMM_E8M0 setting.
1399 + is_deep_gemm_e8m0_used.cache_clear()
tests/kernels/moe/test_deepep_deepgemm_moe.py
+9 -9
@@ -21,7 +21,11 @@ from vllm.model_executor.layers.fused_moe.config import (
21 from vllm.model_executor.layers.fused_moe.fused_moe import fused_experts
22 from vllm.model_executor.layers.fused_moe.modular_kernel import FusedMoEModularKernel
23 from vllm.platforms import current_platform
24 -from vllm.utils.deep_gemm import is_deep_gemm_e8m0_used, is_deep_gemm_supported
24 +from vllm.utils.deep_gemm import (
25 + get_mk_alignment_for_contiguous_layout,
26 + is_deep_gemm_e8m0_used,
27 + is_deep_gemm_supported,
28 +)
29 from vllm.utils.import_utils import has_deep_ep, has_deep_gemm
30
31 from ...utils import multi_gpu_test
@@ -413,19 +417,16 @@ NUM_EXPERTS = [32]
417 @multi_gpu_test(num_gpus=2)
418 @requires_deep_ep
419 @requires_deep_gemm
416 -@pytest.mark.skipif(
417 - is_deep_gemm_e8m0_used(), reason="Skipping test for Blackwell DeepGEMM"
418 -)
420 def test_ht_deepep_deepgemm_moe(
421 mnk: tuple[int, int, int],
422 num_experts: int,
423 topk: int,
424 world_dp_size: tuple[int, int],
425 + disable_deepgemm_ue8m0,
426 ):
427 """
428 Tests for High-Throughput DeepEP + DeepGemm integration.
429 """
428 - import deep_gemm
430
431 m, n, k = mnk
432 current_platform.seed_everything(7)
@@ -433,7 +434,7 @@ def test_ht_deepep_deepgemm_moe(
434 if topk > num_experts:
435 pytest.skip(f"Skipping test: topk={topk} > E={num_experts}")
436
436 - block_m = deep_gemm.get_m_alignment_for_contiguous_layout()
437 + block_m = get_mk_alignment_for_contiguous_layout()[0]
438 block_size = [block_m, block_m]
439
440 world_size, dp_size = world_dp_size
@@ -487,9 +488,6 @@ USE_FP8_DISPATCH = [False]
488 @multi_gpu_test(num_gpus=2)
489 @requires_deep_ep
490 @requires_deep_gemm
490 -@pytest.mark.skipif(
491 - is_deep_gemm_e8m0_used(), reason="Skipping test for Blackwell DeepGEMM"
492 -)
491 def test_ll_deepep_deepgemm_moe(
492 mnk: tuple[int, int, int],
493 num_experts: int,
@@ -497,10 +495,12 @@ def test_ll_deepep_deepgemm_moe(
495 use_fp8_dispatch: bool,
496 block_size: list[int],
497 world_dp_size: tuple[int, int],
498 + disable_deepgemm_ue8m0,
499 ):
500 """
501 Tests for Low-Latency DeepEP + DeepGemm integration.
502 """
503 + assert not is_deep_gemm_e8m0_used()
504
505 m, n, k = mnk
506 current_platform.seed_everything(7)
tests/kernels/moe/test_deepep_moe.py
+1 -1
@@ -294,7 +294,7 @@ def torch_moe_impl(
294 # blockwise quant and de-quant.
295 assert not per_act_token_quant
296 a = test_tensors.rank_tokens
297 - aq, aq_scale = per_token_group_quant_fp8(a, 128)
297 + aq, aq_scale = per_token_group_quant_fp8(a, 128, use_ue8m0=False)
298 a = (
299 (aq.view(-1, 128).to(torch.float32) * aq_scale.view(-1, 1))
300 .view(a.shape)
tests/kernels/moe/test_silu_mul_fp8_quant_deep_gemm.py
+237 -74
@@ -1,6 +1,9 @@
1 # SPDX-License-Identifier: Apache-2.0
2 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
4 +
5 +import random
6 +
7 import pytest
8 import torch
9
@@ -8,27 +11,30 @@ from vllm.model_executor.layers.fused_moe.batched_deep_gemm_moe import (
11 persistent_masked_m_silu_mul_quant,
12 )
13 from vllm.platforms import current_platform
11 -from vllm.utils.math_utils import cdiv
14 +from vllm.utils.deep_gemm import DeepGemmQuantScaleFMT, has_deep_gemm
15 +from vllm.utils.math_utils import cdiv, round_up
16
17 fp8_dtype = torch.float8_e4m3fn
18
19 CASES = [
20 (1, 1, 128, fp8_dtype),
17 - (1, 4, 128, fp8_dtype),
18 - (2, 4, 256, fp8_dtype),
19 - (32, 64, 256, fp8_dtype),
20 - (17, 31, 768, fp8_dtype),
21 - (1, 1, 128 * 1, fp8_dtype),
22 - (1, 1, 128 * 3, fp8_dtype),
23 - (1, 1, 128 * 4, fp8_dtype),
24 - (8, 16, 128 * 1, fp8_dtype),
25 - (8, 16, 128 * 2, fp8_dtype),
26 - (8, 16, 128 * 3, fp8_dtype),
21 + (1, 4, 128 * 1, fp8_dtype),
22 + (2, 4, 128 * 2, fp8_dtype),
23 + (1, 4, 128 * 3, fp8_dtype),
24 + (8, 16, 128 * 4, fp8_dtype),
25 + (8, 16, 128 * 5, fp8_dtype),
26 + (8, 16, 128 * 6, fp8_dtype),
27 + (8, 16, 128 * 7, fp8_dtype),
28 + (8, 16, 128 * 8, fp8_dtype),
29 + (8, 16, 128 * 9, fp8_dtype),
30 (8, 64, 7168, fp8_dtype),
31 (8, 128, 128 * 33, fp8_dtype),
32 + (1, 4, 128 * 10, fp8_dtype),
33 (8, 128, 7168, fp8_dtype),
34 (8, 512, 7168, fp8_dtype),
35 (8, 1024, 7168, fp8_dtype),
36 + (17, 31, 768, fp8_dtype),
37 + (32, 64, 256, fp8_dtype),
38 (256, 8, 7168, fp8_dtype),
39 (256, 32, 7168, fp8_dtype),
40 (256, 64, 7168, fp8_dtype),
@@ -38,14 +44,159 @@ CASES = [
44 ]
45
46
47 +def as_uint8(x) -> torch.Tensor:
48 + return (
49 + torch.empty(x.shape, dtype=x.dtype, device=x.device).copy_(x).view(torch.uint8)
50 + )
51 +
52 +
53 +def silu(x: torch.Tensor) -> torch.Tensor:
54 + one_f32 = torch.tensor([1.0], device=x.device, dtype=torch.float32)
55 + x_f32 = x.to(torch.float32)
56 + act_f32 = x_f32 / (one_f32 + torch.exp(-x_f32))
57 + assert act_f32.dtype == torch.float32
58 + return act_f32.to(torch.bfloat16)
59 +
60 +
61 +def do_quant(x: torch.Tensor, group_size: int, ceil_ue8m0: bool):
62 + eps_bf16 = torch.tensor([1e-10], device=x.device, dtype=torch.bfloat16)
63 + one_bf16 = torch.tensor([1.0], device=x.device, dtype=torch.bfloat16)
64 + fp8_max_bf16 = torch.tensor(
65 + [torch.finfo(fp8_dtype).max], device=x.device, dtype=torch.bfloat16
66 + )
67 + fp8_min_bf16 = torch.tensor(
68 + [torch.finfo(fp8_dtype).min], device=x.device, dtype=torch.bfloat16
69 + )
70 + fp8_max_inv = one_bf16 / fp8_max_bf16
71 + assert fp8_max_inv.dtype == torch.bfloat16
72 +
73 + assert x.size(-1) % group_size == 0
74 + num_groups = x.numel() // group_size
75 + x_og_shape = x.shape
76 +
77 + x = x.to(torch.bfloat16)
78 + x = x.view((-1, group_size))
79 + amax = x.abs().amax(dim=1).clamp(min=eps_bf16)
80 + assert amax.dtype == torch.bfloat16
81 + s = amax * fp8_max_inv
82 +
83 + if ceil_ue8m0:
84 + s = torch.exp2(
85 + torch.ceil(torch.log2(s).to(torch.bfloat16)).to(torch.bfloat16)
86 + ).to(torch.bfloat16)
87 +
88 + inv_s = one_bf16 / s
89 + inv_s = inv_s.view((num_groups, 1))
90 + xq = torch.clamp(x * inv_s, min=fp8_min_bf16.item(), max=fp8_max_bf16.item()).to(
91 + fp8_dtype
92 + )
93 +
94 + xq = xq.view(x_og_shape)
95 + xs = s.view((-1, xq.size(-1) // group_size))
96 + return xq, xs
97 +
98 +
99 +def silu_mul_quant(
100 + gate: torch.Tensor, up: torch.Tensor, group_size: int, ceil_ue8m0: bool
101 +) -> tuple[torch.Tensor, torch.Tensor]:
102 + assert gate.size(-1) % group_size == 0
103 + assert up.size(-1) % group_size == 0
104 +
105 + assert gate.dtype == torch.bfloat16
106 + assert up.dtype == torch.bfloat16
107 +
108 + act_bf16 = silu(gate)
109 + assert act_bf16.dtype == torch.bfloat16
110 +
111 + # act & mul
112 + a_m = act_bf16 * up
113 + assert a_m.dtype == torch.bfloat16
114 +
115 + q, s = do_quant(a_m, group_size, ceil_ue8m0)
116 + return q, s
117 +
118 +
119 +def pack_scales(x: torch.Tensor, tokens_per_expert: torch.Tensor) -> torch.Tensor:
120 + """
121 + pack float32 scales into a int32 tensor
122 + """
123 + assert x.dtype == torch.float32
124 + E, T, G = x.size()
125 +
126 + # Add i32_padding here so we can view it as a i32 tensor later on.
127 + i32_padding = round_up(G, 4) - G
128 + ref_s_i8 = torch.empty((E, T, G + i32_padding), dtype=torch.uint8, device="cuda")
129 + for e in range(E):
130 + nt = tokens_per_expert[e].item()
131 + ref_s_i8[e, :nt, :G] = x[e, :nt].view(torch.int32) >> 23
132 +
133 + ref_s_i32 = ref_s_i8.view(torch.int32)
134 +
135 + return ref_s_i32
136 +
137 +
138 +def ref_with_scale_fmt(
139 + E: int,
140 + T: int,
141 + H: int,
142 + group_size: int,
143 + tokens_per_expert: torch.Tensor,
144 + gate: torch.Tensor,
145 + up: torch.Tensor,
146 + scale_fmt: DeepGemmQuantScaleFMT,
147 +) -> tuple[torch.Tensor, torch.Tensor]:
148 + """
149 + The precision types of the operations triggered by this function
150 + match closely with the kernel implementation so we compare more
151 + accurately.
152 + """
153 + scale_dtype = (
154 + torch.int32 if scale_fmt == DeepGemmQuantScaleFMT.UE8M0 else torch.float32
155 + )
156 + ceil_ue8m0 = scale_fmt in [
157 + DeepGemmQuantScaleFMT.UE8M0,
158 + DeepGemmQuantScaleFMT.FLOAT32_CEIL_UE8M0,
159 + ]
160 +
161 + ref_q = torch.empty((E, T, H), dtype=fp8_dtype, device="cuda")
162 + ref_s_f32 = torch.empty(
163 + (E, T, cdiv(H, group_size)), dtype=torch.float32, device="cuda"
164 + )
165 +
166 + for e in range(E):
167 + nt = tokens_per_expert[e].item()
168 + if nt == 0:
169 + continue
170 + ref_q[e, :nt], ref_s_f32[e, :nt] = silu_mul_quant(
171 + gate[e, :nt], up[e, :nt], group_size, ceil_ue8m0=ceil_ue8m0
172 + )
173 +
174 + if scale_dtype == torch.float32:
175 + return ref_q, ref_s_f32
176 +
177 + assert scale_dtype == torch.int32
178 + return ref_q, pack_scales(ref_s_f32, tokens_per_expert)
179 +
180 +
181 +def token_random(E, T, H2, tokens_per_expert):
182 + """
183 + Initialize each token in a random range so we test a range of
184 + scale values.
185 + """
186 + y = torch.empty((E, T, H2), dtype=torch.bfloat16, device="cuda")
187 + for e in range(E):
188 + for t in range(tokens_per_expert[e].item()):
189 + exp = random.choice(range(1, 20))
190 + y[e, t].uniform_(-(2**exp), 2**exp)
191 + return y
192 +
193 +
194 @pytest.mark.parametrize("E,T,H,fp8_type", CASES)
195 @torch.inference_mode()
43 -def test_silu_mul_fp8_quant_deep_gemm(E, T, H, fp8_type):
196 +def test_silu_mul_fp8_quant_deep_gemm(E: int, T: int, H: int, fp8_type: torch.dtype):
197 group_size = 128
198 current_platform.seed_everything(42)
199
47 - # Input tensor of shape (E, T, 2*H)
48 - y = torch.randn((E, T, 2 * H), dtype=torch.bfloat16, device="cuda")
200 tokens_per_expert = torch.randint(
201 low=0,
202 high=T,
@@ -54,71 +205,83 @@ def test_silu_mul_fp8_quant_deep_gemm(E, T, H, fp8_type):
205 device="cuda",
206 )
207
57 - # Run the SiLU V2 kernel
58 - # TODO (varun): use_e8m0 is set to false as the reference impl does
59 - # not handle that case.
60 - y_q, y_s = persistent_masked_m_silu_mul_quant(
61 - y, tokens_per_expert, group_size=group_size, use_ue8m0=False
62 - )
208 + # Input tensor of shape (E, T, 2*H)
209 + y = token_random(E, T, 2 * H, tokens_per_expert)
210
64 - torch.cuda.synchronize()
65 - fp8_info = torch.finfo(fp8_dtype)
66 - fp8_max = fp8_info.max
67 - fp8_min = fp8_info.min
68 - eps = 1e-10
211 + gate = y[..., :H].to(torch.bfloat16)
212 + up = y[..., H:].to(torch.bfloat16)
213
70 - y1 = y[..., :H].float()
71 - y2 = y[..., H:]
72 - silu_x = y1 * torch.sigmoid(y1)
73 - merged = silu_x * y2
214 + scale_fmts = [
215 + DeepGemmQuantScaleFMT.FLOAT32,
216 + DeepGemmQuantScaleFMT.FLOAT32_CEIL_UE8M0,
217 + DeepGemmQuantScaleFMT.UE8M0,
218 + ]
219
75 - for e in range(E):
76 - nt = tokens_per_expert[e].item()
77 - ref_s = torch.empty(
78 - (T, cdiv(H, group_size)), dtype=torch.float32, device="cuda"
220 + # Run the SiLU V2 kernel
221 + for scale_fmt in scale_fmts:
222 + y_q, y_s = persistent_masked_m_silu_mul_quant(
223 + y,
224 + tokens_per_expert,
225 + group_size=group_size,
226 + quant_scale_fmt=scale_fmt,
227 )
80 - ref_q = torch.empty((T, H), dtype=fp8_dtype, device="cuda")
228
82 - for t in range(nt):
83 - data = merged[e, t].float()
84 - ref_q_row = torch.empty_like(data)
229 + ref_y_q, ref_y_s = ref_with_scale_fmt(
230 + E, T, H, group_size, tokens_per_expert, gate, up, scale_fmt=scale_fmt
231 + )
232
86 - # process full groups
87 - n_full_groups = H // group_size
88 - if n_full_groups > 0:
89 - data_grp = data[: n_full_groups * group_size].view(
90 - n_full_groups, group_size
91 - )
92 - amax = data_grp.abs().amax(dim=1).clamp(min=eps)
93 - scale = amax / fp8_max
94 - scaled = data[: n_full_groups * group_size] / scale.repeat_interleave(
95 - group_size
96 - )
97 - ref_q_row[: n_full_groups * group_size] = scaled.clamp(
98 - fp8_min, fp8_max
99 - ).to(fp8_dtype)
100 - ref_s[t, :n_full_groups] = scale
101 -
102 - # process remainder group
103 - rem = H % group_size
104 - if rem > 0:
105 - data_rem = data[-rem:]
106 - amax = data_rem.abs().amax().clamp(min=eps)
107 - scale = amax / fp8_max
108 - scaled = data_rem / scale
109 - ref_q_row[-rem:] = scaled.clamp(fp8_min, fp8_max).to(fp8_dtype)
110 - ref_s[t, -1] = scale
111 -
112 - ref_q[t] = ref_q_row
113 -
114 - y_se = y_s[e].float()
115 - y_qe = y_q[e].float()
116 -
117 - torch.testing.assert_close(
118 - y_qe[:nt].to(torch.float32),
119 - ref_q[:nt].to(torch.float32),
120 - atol=2,
121 - rtol=2e-1,
233 + # deepgemm scales transform
234 + dg_scales = None
235 + if (
236 + has_deep_gemm()
237 + and current_platform.has_device_capability(100)
238 + and scale_fmt == DeepGemmQuantScaleFMT.UE8M0
239 + ):
240 + from deep_gemm import transform_sf_into_required_layout
241 +
242 + _q, _s = ref_with_scale_fmt(
243 + E,
244 + T,
245 + H,
246 + group_size,
247 + tokens_per_expert,
248 + gate,
249 + up,
250 + scale_fmt=DeepGemmQuantScaleFMT.FLOAT32_CEIL_UE8M0,
251 + )
252 + dg_scales = transform_sf_into_required_layout(
253 + sf=_s,
254 + mn=_q.size(1),
255 + k=_q.size(2),
256 + recipe=(1, 128, 128),
257 + num_groups=_q.size(0),
258 + is_sfa=True,
259 + )
260 +
261 + expected_scale_dtype = (
262 + torch.int32 if scale_fmt == DeepGemmQuantScaleFMT.UE8M0 else torch.float32
263 )
264 + assert y_s.dtype == expected_scale_dtype
265 + assert ref_y_s.dtype == expected_scale_dtype
266
124 - torch.testing.assert_close(y_se[:nt], ref_s[:nt], atol=1e-4, rtol=1e-2)
267 + for e in range(E):
268 + nt = tokens_per_expert[e].item()
269 +
270 + torch.testing.assert_close(
271 + y_q[e, :nt].to(torch.float32),
272 + ref_y_q[e, :nt].to(torch.float32),
273 + )
274 +
275 + if scale_fmt == DeepGemmQuantScaleFMT.UE8M0:
276 + G = H // group_size
277 + y_s_sliced = as_uint8(y_s[e])
278 + ref_s_sliced = as_uint8(ref_y_s[e])
279 + torch.testing.assert_close(y_s_sliced[:nt, :G], ref_s_sliced[:nt, :G])
280 + if dg_scales is not None:
281 + dg_sliced = as_uint8(dg_scales[e])
282 + torch.testing.assert_close(y_s_sliced[:nt, :G], dg_sliced[:nt, :G])
283 + else:
284 + torch.testing.assert_close(
285 + y_s[e, :nt],
286 + ref_y_s[e, :nt],
287 + )
vllm/model_executor/layers/fused_moe/batched_deep_gemm_moe.py
+57 -19
@@ -1,6 +1,7 @@
1 # SPDX-License-Identifier: Apache-2.0
2 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
4 +
5 import torch
6
7 import vllm.model_executor.layers.fused_moe.modular_kernel as mk
@@ -13,14 +14,33 @@ from vllm.model_executor.layers.fused_moe.utils import _resize_cache
14 from vllm.platforms import current_platform
15 from vllm.triton_utils import tl, triton
16 from vllm.utils.deep_gemm import (
17 + DeepGemmQuantScaleFMT,
18 fp8_m_grouped_gemm_nt_masked,
19 get_mk_alignment_for_contiguous_layout,
20 is_deep_gemm_e8m0_used,
21 )
22 +from vllm.utils.math_utils import cdiv
23
24 logger = init_logger(__name__)
25
26
27 +def scales_shape_stride_dtype(
28 + E: int, T: int, G: int, quant_scale_fmt: DeepGemmQuantScaleFMT
29 +) -> tuple[tuple[int, ...], tuple[int, ...], torch.dtype]:
30 + shape = (E, T, G)
31 + strides = (T * G, 1, T)
32 + if quant_scale_fmt in [
33 + DeepGemmQuantScaleFMT.FLOAT32,
34 + DeepGemmQuantScaleFMT.FLOAT32_CEIL_UE8M0,
35 + ]:
36 + return shape, strides, torch.float32
37 +
38 + assert quant_scale_fmt == DeepGemmQuantScaleFMT.UE8M0
39 + shape = (E, T, cdiv(G, 4))
40 + strides = (T * cdiv(G, 4), 1, T)
41 + return shape, strides, torch.int32
42 +
43 +
44 @triton.jit
45 def _silu_mul_fp8_quant_deep_gemm(
46 # Pointers ------------------------------------------------------------
@@ -49,7 +69,7 @@ def _silu_mul_fp8_quant_deep_gemm(
69 eps: tl.constexpr,
70 fp8_min: tl.constexpr,
71 fp8_max: tl.constexpr,
52 - use_ue8m0: tl.constexpr,
72 + ceil_ue8m0: tl.constexpr,
73 # Meta ---------------------------------------------------------------
74 BLOCK: tl.constexpr,
75 NUM_STAGES: tl.constexpr,
@@ -86,7 +106,7 @@ def _silu_mul_fp8_quant_deep_gemm(
106 y = gate * up
107
108 y_s = tl.maximum(tl.max(tl.abs(y)), eps) / fp8_max
89 - if use_ue8m0:
109 + if ceil_ue8m0:
110 y_s = tl.exp2(tl.ceil(tl.log2(y_s)))
111
112 y_q = tl.clamp(y / y_s, fp8_min, fp8_max).to(y_q_ptr.dtype.element_ty)
@@ -100,7 +120,7 @@ def persistent_masked_m_silu_mul_quant(
120 tokens_per_expert: torch.Tensor, # (E,) number of valid tokens per expert
121 num_parallel_tokens=16,
122 group_size: int = 128,
103 - use_ue8m0: bool | None = None,
123 + quant_scale_fmt: DeepGemmQuantScaleFMT = DeepGemmQuantScaleFMT.FLOAT32,
124 ) -> tuple[torch.Tensor, torch.Tensor]:
125 """Quantize silu(y[..., :H]) * y[..., H:] to FP8 with group per-token scales
126 y has shape (E, T, 2*H). The first half of the last dimension is
@@ -137,7 +157,13 @@ def persistent_masked_m_silu_mul_quant(
157
158 Returns `(y_q, y_s)` where
159 * `y_q`: FP8 tensor, shape (E, T, H), same layout as y[..., :H]
140 - * `y_s`: FP32 tensor, shape (E, T, H // group_size), strides (T*G, 1, T)
160 + * `y_s` depends on quant_scale_fmt,
161 + - quant_scale_fmt == FLOAT32,
162 + `y_s`: FP32 tensor, shape (E, T, H // group_size), strides (T*G, 1, T)
163 + - quant_scale_fmt == E8M0,
164 + `y_s`: Int32 tensor, shape (E, T, H // group_size // 4), strides (T*G, 1, T)
165 + - quant_scale_fmt == E8M0_FLOAT32_SPARSE
166 + `y_s`: FP32 tensor, shape (E, T, H // group_size), strides (T*G, 1, T)
167 Let NUM_WARPS be the number of warps in a single thread block and
168 `GROUP_SIZE = 128` be the size of the quantization group.
169 """
@@ -155,17 +181,18 @@ def persistent_masked_m_silu_mul_quant(
181 fp8_dtype = torch.float8_e4m3fn
182 y_q = torch.empty((E, T, H), dtype=fp8_dtype, device=y.device)
183
158 - stride_ys_e = T * G
159 - stride_ys_t = 1
160 - stride_ys_g = T
184 + ys_shape, ys_strides, ys_dtype = scales_shape_stride_dtype(E, T, G, quant_scale_fmt)
185 y_s = torch.empty_strided(
162 - (E, T, G),
163 - (stride_ys_e, stride_ys_t, stride_ys_g),
164 - dtype=torch.float32,
186 + ys_shape,
187 + ys_strides,
188 + dtype=ys_dtype,
189 device=y.device,
190 )
191
168 - use_ue8m0 = use_ue8m0 if use_ue8m0 is not None else is_deep_gemm_e8m0_used()
192 + ceil_ue8m0 = quant_scale_fmt in [
193 + DeepGemmQuantScaleFMT.FLOAT32_CEIL_UE8M0,
194 + DeepGemmQuantScaleFMT.UE8M0,
195 + ]
196
197 cuda_arch = current_platform.get_device_capability(
198 device_id=y.device.index
@@ -173,7 +200,7 @@ def persistent_masked_m_silu_mul_quant(
200
201 if cuda_arch >= 80:
202 torch.ops._C.persistent_masked_m_silu_mul_quant(
176 - y, tokens_per_expert, y_q, y_s, use_ue8m0
203 + y, tokens_per_expert, y_q, y_s, ceil_ue8m0
204 )
205 else:
206 stride_cnt_e = tokens_per_expert.stride()[0]
@@ -189,6 +216,10 @@ def persistent_masked_m_silu_mul_quant(
216 fp8_max = f_info.max
217 fp8_min = f_info.min
218 eps: float = 1e-10
219 + assert y_s.dtype == torch.float32, (
220 + "_silu_mul_fp8_quant_deep_gemm does"
221 + "not support {y_s.dtype} scales. Only torch.float32 supported."
222 + )
223 _silu_mul_fp8_quant_deep_gemm[grid](
224 y,
225 y_q,
@@ -202,14 +233,14 @@ def persistent_masked_m_silu_mul_quant(
233 stride_yq_e,
234 stride_yq_t,
235 stride_yq_h,
205 - stride_ys_e,
206 - stride_ys_t,
207 - stride_ys_g,
236 + ys_strides[0],
237 + ys_strides[1],
238 + ys_strides[2],
239 stride_cnt_e,
240 eps,
241 fp8_min,
242 fp8_max,
212 - is_deep_gemm_e8m0_used(),
243 + ceil_ue8m0,
244 BLOCK=group_size,
245 NUM_STAGES=4,
246 num_warps=1,
@@ -255,7 +286,7 @@ class BatchedDeepGemmExperts(mk.FusedMoEPermuteExpertsUnpermute):
286 """
287 DeepGemm supports packed ue8m0 activation scales format in devices == sm100
288 """
258 - return current_platform.is_device_capability(100)
289 + return is_deep_gemm_e8m0_used() and current_platform.is_device_capability(100)
290
291 def finalize_weight_and_reduce_impl(self) -> mk.TopKWeightAndReduce:
292 # Let PrepareAndFinalize::finalize() decide the impl.
@@ -329,10 +360,17 @@ class BatchedDeepGemmExperts(mk.FusedMoEPermuteExpertsUnpermute):
360 expected_m,
361 )
362
363 + quant_scale_fmt = DeepGemmQuantScaleFMT.from_oracle()
364 a2q, a2q_scale = persistent_masked_m_silu_mul_quant(
333 - workspace1, expert_num_tokens
365 + workspace1,
366 + expert_num_tokens,
367 + quant_scale_fmt=quant_scale_fmt,
368 )
369
370 fp8_m_grouped_gemm_nt_masked(
337 - (a2q, a2q_scale), (w2, self.w2_scale), output, expert_num_tokens, expected_m
371 + (a2q, a2q_scale),
372 + (w2, self.w2_scale),
373 + output,
374 + expert_num_tokens,
375 + expected_m,
376 )
vllm/utils/deep_gemm.py
+23
@@ -9,6 +9,7 @@ import functools
9 import importlib
10 import os
11 from collections.abc import Callable
12 +from enum import Enum
13 from typing import Any, NoReturn
14
15 import torch
@@ -20,6 +21,28 @@ from vllm.utils.import_utils import has_deep_gemm
21 from vllm.utils.math_utils import cdiv
22
23
24 +class DeepGemmQuantScaleFMT(Enum):
25 + # Float32 scales in Float32 tensor
26 + FLOAT32 = 0
27 + # Compute float32 scales and ceil the scales to UE8M0.
28 + # Keep the scales in Float32 tensor.
29 + FLOAT32_CEIL_UE8M0 = 1
30 + # Compute float32 scales and ceil the scales to UE8M0.
31 + # Pack the scales into a int32 tensor where each int32
32 + # element contains 4 scale values.
33 + UE8M0 = 2
34 +
35 + @staticmethod
36 + def from_oracle() -> "DeepGemmQuantScaleFMT":
37 + if not is_deep_gemm_e8m0_used():
38 + return DeepGemmQuantScaleFMT.FLOAT32
39 + return (
40 + DeepGemmQuantScaleFMT.UE8M0
41 + if current_platform.is_device_capability(100)
42 + else DeepGemmQuantScaleFMT.FLOAT32_CEIL_UE8M0
43 + )
44 +
45 +
46 @functools.cache
47 def is_deep_gemm_supported() -> bool:
48 """Return `True` if DeepGEMM is supported on the current platform.