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

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 14, 05:30 UTC

[Kernels] Support blocked fp8 quantization for compressed tensors MoE (#25219)

Signed-off-by: Bill Nell <bnell@redhat.com> Co-authored-by: Michael Goin <mgoin64@gmail.com>

bnellnm committed Sep 23, 2025 at 12:11 UTC f11e3c516be3d88733ea4b0c79f47e33cc319197

4 files changed +112 -29

vllm/model_executor/layers/quantization/compressed_tensors/compressed_tensors_moe.py

+96 -13

@@ -13,6 +13,7 @@ from compressed_tensors.quantization import (ActivationOrdering,

13 import vllm.envs as envs

14 import vllm.model_executor.layers.fused_moe.modular_kernel as mk

15 from vllm import _custom_ops as ops

16 +from vllm.distributed import get_tensor_model_parallel_world_size

17 from vllm.logger import init_logger

18 from vllm.model_executor.layers.fused_moe import (

19 FusedMoE, FusedMoEActivationFormat, FusedMoEConfig, FusedMoEMethodBase,

@@ -31,6 +32,9 @@ from vllm.model_executor.layers.quantization.utils import replace_parameter

32 from vllm.model_executor.layers.quantization.utils.flashinfer_fp4_moe import (

33 build_flashinfer_fp4_cutlass_moe_prepare_finalize, reorder_w1w3_to_w3w1,

34 select_nvfp4_gemm_impl)

35 +from vllm.model_executor.layers.quantization.utils.fp8_utils import (

36 + expert_weight_is_col_major, get_col_major_tma_aligned_tensor,

37 + requant_weight_ue8m0_inplace)

38 from vllm.model_executor.layers.quantization.utils.marlin_utils import (

39 check_moe_marlin_supports_layer, marlin_make_workspace_new,

40 marlin_moe_permute_scales)

@@ -45,6 +49,7 @@ from vllm.model_executor.layers.quantization.utils.w8a8_utils import (

49 from vllm.model_executor.utils import set_weight_attrs

50 from vllm.platforms import current_platform

51 from vllm.scalar_type import scalar_types

52 +from vllm.utils.deep_gemm import is_deep_gemm_e8m0_used

54 logger = init_logger(__name__)

@@ -505,10 +510,12 @@ class CompressedTensorsW8A8Fp8MoEMethod(CompressedTensorsMoEMethod):

510 self.weight_quant.strategy == QuantizationStrategy.CHANNEL

511 and self.input_quant.strategy == QuantizationStrategy.TOKEN)

512 if not (per_tensor or per_channel):

508 - raise ValueError(

509 - "For FP8 Fused MoE layers, we require per tensor "

510 - "or channelwise, dynamic per token quantization. Found "

511 - f"{self.weight_quant}, {self.input_quant}")

513 + assert self.weight_quant.strategy == QuantizationStrategy.BLOCK

514 + self.weight_block_size = self.weight_quant.block_structure

515 + assert self.weight_quant.dynamic is not None

516 + else:

517 + self.weight_block_size = None

518 + self.block_quant = self.weight_block_size is not None

519

520 self.static_input_scales = not self.input_quant.dynamic

521 if self.static_input_scales and per_channel:

@@ -519,7 +526,8 @@ class CompressedTensorsW8A8Fp8MoEMethod(CompressedTensorsMoEMethod):

526 # For GPUs that lack FP8 hardware support, we can leverage the Marlin

527 # kernel for fast weight-only FP8 quantization

528 self.use_marlin = (not current_platform.has_device_capability(89)

522 - or envs.VLLM_TEST_FORCE_FP8_MARLIN)

529 + or envs.VLLM_TEST_FORCE_FP8_MARLIN

530 + and not self.block_quant)

531 # Disable marlin for rocm

532 if current_platform.is_rocm():

533 self.use_marlin = False

@@ -531,8 +539,9 @@ class CompressedTensorsW8A8Fp8MoEMethod(CompressedTensorsMoEMethod):

539 # cutlass path

540 self.is_fp8_w8a8_sm100 = quant_config._is_fp8_w8a8_sm100(

541 self.weight_quant, self.input_quant)

534 - self.use_cutlass = (quant_config._is_fp8_w8a8_sm90(

535 - self.weight_quant, self.input_quant) or self.is_fp8_w8a8_sm100)

542 + self.use_cutlass = not self.block_quant and (

543 + quant_config._is_fp8_w8a8_sm90(self.weight_quant, self.input_quant)

544 + or self.is_fp8_w8a8_sm100)

545 self.disable_expert_map = False

546

547 def create_weights(self, layer: torch.nn.Module, num_experts: int,

@@ -547,6 +556,31 @@ class CompressedTensorsW8A8Fp8MoEMethod(CompressedTensorsMoEMethod):

556

557 params_dtype = torch.float8_e4m3fn

558

559 + if self.block_quant:

560 + assert self.weight_block_size is not None

561 + layer.weight_block_size = self.weight_block_size

562 + tp_size = get_tensor_model_parallel_world_size()

563 + block_n, block_k = (

564 + self.weight_block_size[0],

565 + self.weight_block_size[1],

566 + )

567 + # NOTE: To ensure proper alignment of the block-wise quantization

568 + # scales, the output_size of the weights for both the gate and up

569 + # layers must be divisible by block_n.

570 + # Required by column parallel or enabling merged weights

571 + if intermediate_size_per_partition % block_n != 0:

572 + raise ValueError(

573 + f"The output_size of gate's and up's weight = "

574 + f"{intermediate_size_per_partition} is not divisible by "

575 + f"weight quantization block_n = {block_n}.")

576 + if (tp_size > 1

577 + and intermediate_size_per_partition % block_k != 0):

578 + # Required by row parallel

579 + raise ValueError(

580 + f"The input_size of down's weight = "

581 + f"{intermediate_size_per_partition} is not divisible by "

582 + f"weight quantization block_k = {block_k}.")

583 +

584 # WEIGHTS

585 w13_weight = torch.nn.Parameter(torch.empty(

586 num_experts,

@@ -602,6 +636,27 @@ class CompressedTensorsW8A8Fp8MoEMethod(CompressedTensorsMoEMethod):

636 set_weight_attrs(w13_weight_scale, extra_weight_attrs)

637 set_weight_attrs(w2_weight_scale, extra_weight_attrs)

638

639 + elif self.weight_quant.strategy == QuantizationStrategy.BLOCK:

640 + w13_weight_scale = torch.nn.Parameter(torch.ones(

641 + num_experts,

642 + 2 *

643 + ((intermediate_size_per_partition + block_n - 1) // block_n),

644 + (hidden_size + block_k - 1) // block_k,

645 + dtype=torch.float32),

646 + requires_grad=False)

647 + layer.register_parameter("w13_weight_scale", w13_weight_scale)

648 + w2_weight_scale = torch.nn.Parameter(torch.ones(

649 + num_experts, (hidden_size + block_n - 1) // block_n,

650 + (intermediate_size_per_partition + block_k - 1) // block_k,

651 + dtype=torch.float32),

652 + requires_grad=False)

653 + layer.register_parameter("w2_weight_scale", w2_weight_scale)

654 + # Add PER-CHANNEL quantization for FusedMoE.weight_loader.

655 + extra_weight_attrs.update(

656 + {"quant_method": FusedMoeWeightScaleSupported.BLOCK.value})

657 + set_weight_attrs(w13_weight_scale, extra_weight_attrs)

658 + set_weight_attrs(w2_weight_scale, extra_weight_attrs)

659 +

660 # INPUT_SCALES

661 if self.static_input_scales:

662 w13_input_scale = torch.nn.Parameter(torch.ones(

@@ -706,6 +761,7 @@ class CompressedTensorsW8A8Fp8MoEMethod(CompressedTensorsMoEMethod):

761 del layer.w2_input_scale

762

763 if self.use_cutlass:

764 + assert self.weight_quant.strategy != QuantizationStrategy.BLOCK

765 device = layer.w13_weight.device

766 # ab_strides1 and c_strides2 are the same

767 self.ab_strides1_c_strides2 = torch.full(

@@ -724,6 +780,29 @@ class CompressedTensorsW8A8Fp8MoEMethod(CompressedTensorsMoEMethod):

780 device=device,

781 dtype=torch.int64)

782

783 + if is_deep_gemm_e8m0_used() and self.block_quant:

784 + assert layer.weight_block_size is not None

785 + # Re-quantise the expert weights so their scales are UE8M0.

786 + block_sz = tuple(layer.weight_block_size)

787 + requant_weight_ue8m0_inplace(

788 + layer.w13_weight.data,

789 + layer.w13_weight_scale.data,

790 + block_sz,

791 + )

792 + requant_weight_ue8m0_inplace(

793 + layer.w2_weight.data,

794 + layer.w2_weight_scale.data,

795 + block_sz,

796 + )

797 +

798 + # Ensure column-major TMA alignment expected by DeepGEMM.

799 + if expert_weight_is_col_major(layer.w13_weight_scale):

800 + layer.w13_weight_scale = get_col_major_tma_aligned_tensor(

801 + layer.w13_weight_scale)

802 + if expert_weight_is_col_major(layer.w2_weight_scale):

803 + layer.w2_weight_scale = get_col_major_tma_aligned_tensor(

804 + layer.w2_weight_scale)

805 +

806 def maybe_make_prepare_finalize(

807 self) -> Optional[mk.FusedMoEPrepareAndFinalize]:

808 if self.use_marlin or self.rocm_aiter_moe_enabled:

@@ -777,9 +856,10 @@ class CompressedTensorsW8A8Fp8MoEMethod(CompressedTensorsMoEMethod):

856 return experts

857

858 # triton path

780 - from vllm.model_executor.layers.fused_moe import TritonExperts

781 - from vllm.model_executor.layers.fused_moe.fused_batched_moe import (

782 - BatchedTritonExperts)

859 + from vllm.model_executor.layers.fused_moe.batched_triton_or_deep_gemm_moe import ( # noqa: E501

860 + BatchedTritonOrDeepGemmExperts)

861 + from vllm.model_executor.layers.fused_moe.triton_deep_gemm_moe import (

862 + TritonOrDeepGemmExperts)

863

864 assert not self.rocm_aiter_moe_enabled and not self.use_marlin

865

@@ -790,14 +870,16 @@ class CompressedTensorsW8A8Fp8MoEMethod(CompressedTensorsMoEMethod):

870 assert max_num_tokens_per_rank is not None

871

872 logger.debug("BatchedTritonExperts(%s)", self.__class__.__name__)

793 - return BatchedTritonExperts(

873 + return BatchedTritonOrDeepGemmExperts(

874 max_num_tokens=max_num_tokens_per_rank,

875 num_dispatchers=prepare_finalize.num_dispatchers(),

876 quant_config=self.moe_quant_config,

877 )

878 else:

799 - logger.debug("TritonExperts(%s)", self.__class__.__name__)

800 - return TritonExperts(self.moe_quant_config)

879 + logger.debug("TritonOrDeepGemmExperts(%s)",

880 + self.__class__.__name__)

881 + return TritonOrDeepGemmExperts(self.moe_quant_config,

882 + allow_deep_gemm=True)

883

884 def get_fused_moe_quant_config(

885 self, layer: torch.nn.Module) -> Optional[FusedMoEQuantConfig]:

@@ -816,6 +898,7 @@ class CompressedTensorsW8A8Fp8MoEMethod(CompressedTensorsMoEMethod):

898 a2_scale=layer.w2_input_scale,

899 per_act_token_quant=per_act_token,

900 per_out_ch_quant=per_channel_quant,

901 + block_shape=layer.weight_block_size,

902 )

903

904 def apply(

vllm/model_executor/layers/quantization/fp8.py

+8 -14

@@ -33,10 +33,10 @@ from vllm.model_executor.layers.quantization.utils.flashinfer_utils import (

33 from vllm.model_executor.layers.quantization.utils.fp8_utils import (

34 apply_fp8_block_linear, check_aiter_fp8_linear_support,

35 create_fp8_input_scale, create_fp8_scale_parameter,

36 - create_fp8_weight_parameter, get_col_major_tma_aligned_tensor,

37 - maybe_post_process_fp8_weight_block, process_fp8_weight_block_strategy,

38 - process_fp8_weight_tensor_strategy, requant_weight_ue8m0_inplace,

39 - validate_fp8_block_shape)

36 + create_fp8_weight_parameter, expert_weight_is_col_major,

37 + get_col_major_tma_aligned_tensor, maybe_post_process_fp8_weight_block,

38 + process_fp8_weight_block_strategy, process_fp8_weight_tensor_strategy,

39 + requant_weight_ue8m0_inplace, validate_fp8_block_shape)

40 from vllm.model_executor.layers.quantization.utils.marlin_utils_fp8 import (

41 apply_fp8_marlin_linear, prepare_fp8_layer_for_marlin,

42 prepare_moe_fp8_layer_for_marlin)

@@ -64,12 +64,6 @@ ACTIVATION_SCHEMES = ["static", "dynamic"]

64 logger = init_logger(__name__)

67 -def _is_col_major(x: torch.Tensor) -> bool:

68 - assert x.dim() == 3

69 - b, m, n = x.shape

70 - return x.stride(0) == m * n and x.stride(1) == 1 and x.stride(2) == m

71 -

72 -

67 class Fp8Config(QuantizationConfig):

68 """Config class for FP8."""

@@ -660,10 +654,10 @@ class Fp8MoEMethod(FusedMoEMethodBase):

654 # DeepGemm scales need to be transposed and aligned. We try to do

655 # it ahead of time for performance reasons.

656 if self.allow_deep_gemm and not is_deep_gemm_e8m0_used():

663 - if _is_col_major(layer.w13_weight_scale_inv):

657 + if expert_weight_is_col_major(layer.w13_weight_scale_inv):

658 layer.w13_weight_scale_inv = \

659 get_col_major_tma_aligned_tensor(layer.w13_weight_scale_inv)

666 - if _is_col_major(layer.w2_weight_scale_inv):

660 + if expert_weight_is_col_major(layer.w2_weight_scale_inv):

661 layer.w2_weight_scale_inv = \

662 get_col_major_tma_aligned_tensor(layer.w2_weight_scale_inv)

663

@@ -811,10 +805,10 @@ class Fp8MoEMethod(FusedMoEMethodBase):

805 )

806

807 # Ensure column-major TMA alignment expected by DeepGEMM.

814 - if _is_col_major(layer.w13_weight_scale_inv):

808 + if expert_weight_is_col_major(layer.w13_weight_scale_inv):

809 layer.w13_weight_scale_inv = get_col_major_tma_aligned_tensor(

810 layer.w13_weight_scale_inv)

817 - if _is_col_major(layer.w2_weight_scale_inv):

811 + if expert_weight_is_col_major(layer.w2_weight_scale_inv):

812 layer.w2_weight_scale_inv = get_col_major_tma_aligned_tensor(

813 layer.w2_weight_scale_inv)

814

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

@@ -1014,3 +1014,9 @@ def apply_fp8_block_linear(layer: torch.nn.Module, input: torch.Tensor,

1014 cutlass_block_fp8_supported=cutlass_block_fp8_supported,

1015 use_aiter_and_is_supported=use_aiter_and_is_supported,

1016 )

1017 +

1018 +

1019 +def expert_weight_is_col_major(x: torch.Tensor) -> bool:

1020 + assert x.dim() == 3

1021 + b, m, n = x.shape

1022 + return x.stride(0) == m * n and x.stride(1) == 1 and x.stride(2) == m

vllm/model_executor/warmup/deep_gemm_warmup.py

+2 -2

@@ -53,9 +53,9 @@ def _extract_data_from_fused_moe_module(

53 """

54 assert isinstance(m, FusedMoE)

55 w13 = m.w13_weight

56 - w13_s = m.w13_weight_scale_inv

56 + w13_s = getattr(m, "w13_weight_scale_inv", m.w13_weight_scale)

57 w2 = m.w2_weight

58 - w2_s = m.w2_weight_scale_inv

58 + w2_s = getattr(m, "w2_weight_scale_inv", m.w2_weight_scale)

59 num_topk = m.top_k

61 assert isinstance(w13, torch.Tensor)