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, 23:45 UTC
[Refactor] Refactor MOE NVFP4 Code Base: ModelOpt + Compressed Tensor (#21631)
Signed-off-by: yewentao256 <zhyanwentao@126.com>
Wentao Ye committed
Jul 27, 2025 at 08:25 UTC
bda9d0535fa7d47169c4b50c6169e59a670aeb92
6 files changed
+75
-106
tests/quantization/test_compressed_tensors.py
+3
-1
@@ -17,7 +17,9 @@ from vllm.model_executor.layers.quantization.compressed_tensors.compressed_tenso
17
CompressedTensorsW4A4Fp4, CompressedTensorsW4A16Fp4,
18
CompressedTensorsW4A16Sparse24, CompressedTensorsW8A8Fp8,
19
CompressedTensorsW8A8Int8, CompressedTensorsW8A16Fp8,
20
- CompressedTensorsWNA16, cutlass_fp4_supported)
20
+ CompressedTensorsWNA16)
21
+from vllm.model_executor.layers.quantization.utils.quant_utils import (
22
+ cutlass_fp4_supported)
23
from vllm.model_executor.layers.quantization.utils.w8a8_utils import (
24
sparse_cutlass_supported)
25
from vllm.platforms import current_platform
vllm/model_executor/layers/quantization/compressed_tensors/compressed_tensors.py
+1
-1
@@ -33,7 +33,7 @@ from vllm.model_executor.layers.quantization.compressed_tensors.utils import (
33
find_matched_target, is_activation_quantization_format,
34
should_ignore_layer)
35
from vllm.model_executor.layers.quantization.kv_cache import BaseKVCacheMethod
36
-from vllm.model_executor.layers.quantization.utils.nvfp4_emulation_utils import ( # noqa: E501
36
+from vllm.model_executor.layers.quantization.utils.quant_utils import (
37
cutlass_fp4_supported)
38
from vllm.platforms import current_platform
39
vllm/model_executor/layers/quantization/compressed_tensors/compressed_tensors_moe.py
+8
-31
@@ -27,8 +27,8 @@ from vllm.model_executor.layers.quantization.utils.marlin_utils_fp4 import (
27
prepare_moe_fp4_layer_for_marlin)
28
from vllm.model_executor.layers.quantization.utils.marlin_utils_fp8 import (
29
prepare_moe_fp8_layer_for_marlin)
30
-from vllm.model_executor.layers.quantization.utils.nvfp4_emulation_utils import ( # noqa: E501
31
- cutlass_fp4_supported)
30
+from vllm.model_executor.layers.quantization.utils.quant_utils import (
31
+ cutlass_fp4_supported, swizzle_blockscale)
32
from vllm.model_executor.layers.quantization.utils.w8a8_utils import (
33
all_close_1d, normalize_e4m3fn_to_e4m3fnuz, per_tensor_dequantize)
34
from vllm.model_executor.utils import set_weight_attrs
@@ -193,29 +193,6 @@ class CompressedTensorsW4A4MoeMethod(CompressedTensorsMoEMethod):
193
{"quant_method": FusedMoeWeightScaleSupported.TENSOR.value})
194
set_weight_attrs(w2_input_scale, extra_weight_attrs)
195
196
- def swizzle_blockscale(self, scale: torch.tensor):
197
- assert (scale.dtype == torch.float8_e4m3fn)
198
- # Pad and blockwise interleave weight_scale
199
- scale_ndim = scale.ndim
200
- if scale.ndim == 2:
201
- scale = scale.unsqueeze(0)
202
- assert scale.ndim == 3
203
- B, M, K = scale.shape
204
- round_up_multiple = lambda x, m: (x + m - 1) // m * m
205
- M_padded = round_up_multiple(M, 128)
206
- K_padded = round_up_multiple(K, 4)
207
- padded_scale = torch.zeros((B, M_padded, K_padded), dtype=scale.dtype)
208
- padded_scale[:B, :M, :K] = scale
209
- batches, rows, cols = padded_scale.shape
210
- assert rows % 128 == 0
211
- assert cols % 4 == 0
212
- padded_scale = padded_scale.reshape(batches, rows // 128, 4, 32,
213
- cols // 4, 4)
214
- swizzled_scale = padded_scale.permute((0, 1, 4, 3, 2, 5))
215
- swizzled_scale = swizzled_scale.contiguous().cuda()
216
- return (swizzled_scale.reshape(M, K)
217
- if scale_ndim == 2 else swizzled_scale.reshape(B, M, K))
218
-
196
def process_weights_after_loading(self, layer: torch.nn.Module) -> None:
197
198
# From packed to weight
@@ -243,13 +220,13 @@ class CompressedTensorsW4A4MoeMethod(CompressedTensorsMoEMethod):
220
return
221
222
# swizzle weight scales
246
- layer.w13_blockscale_swizzled = torch.nn.Parameter(
247
- self.swizzle_blockscale(layer.w13_weight_scale),
248
- requires_grad=False)
223
+ layer.w13_blockscale_swizzled = torch.nn.Parameter(swizzle_blockscale(
224
+ layer.w13_weight_scale),
225
+ requires_grad=False)
226
250
- layer.w2_blockscale_swizzled = torch.nn.Parameter(
251
- self.swizzle_blockscale(layer.w2_weight_scale),
252
- requires_grad=False)
227
+ layer.w2_blockscale_swizzled = torch.nn.Parameter(swizzle_blockscale(
228
+ layer.w2_weight_scale),
229
+ requires_grad=False)
230
231
# w13
232
w13_input_global_scale = layer.w13_input_global_scale.max(
vllm/model_executor/layers/quantization/modelopt.py
+5
-61
@@ -9,8 +9,7 @@ from torch.nn.parameter import Parameter
9
10
import vllm.envs as envs
11
import vllm.model_executor.layers.fused_moe.modular_kernel as mk
12
-from vllm._custom_ops import (cutlass_scaled_fp4_mm,
13
- cutlass_scaled_mm_supports_fp4, scaled_fp4_quant)
12
+from vllm._custom_ops import cutlass_scaled_fp4_mm, scaled_fp4_quant
13
from vllm.distributed import get_ep_group
14
from vllm.logger import init_logger
15
from vllm.model_executor.layers.fused_moe.config import FusedMoEParallelConfig
@@ -28,7 +27,7 @@ from vllm.model_executor.layers.quantization.utils.marlin_utils_fp4 import (
27
apply_fp4_marlin_linear, is_fp4_marlin_supported,
28
prepare_fp4_layer_for_marlin, prepare_moe_fp4_layer_for_marlin)
29
from vllm.model_executor.layers.quantization.utils.quant_utils import (
31
- GroupShape, is_layer_skipped)
30
+ GroupShape, cutlass_fp4_supported, is_layer_skipped, swizzle_blockscale)
31
from vllm.model_executor.layers.quantization.utils.w8a8_utils import (
32
Fp8LinearOp, requantize_with_max_scale)
33
from vllm.model_executor.parameter import (ModelWeightParameter,
@@ -667,14 +666,6 @@ class ModelOptNvFp4Config(QuantizationConfig):
666
return None
667
668
670
-def cutlass_fp4_supported() -> bool:
671
- if not current_platform.is_cuda():
672
- return False
673
- capability_tuple = current_platform.get_device_capability()
674
- capability = -1 if capability_tuple is None else capability_tuple.to_int()
675
- return cutlass_scaled_mm_supports_fp4(capability)
676
-
677
-
669
class ModelOptFp8KVCacheMethod(BaseKVCacheMethod):
670
"""
671
Supports loading kv-cache scaling factors from FP8 checkpoints.
@@ -772,29 +763,6 @@ class ModelOptNvFp4LinearMethod(LinearMethodBase):
763
764
layer.register_parameter("weight_scale", weight_scale)
765
775
- def swizzle_blockscale(self, scale: torch.tensor):
776
- assert (scale.dtype == torch.float8_e4m3fn)
777
- # Pad and blockwise interleave weight_scale
778
- scale_ndim = scale.ndim
779
- if scale.ndim == 2:
780
- scale = scale.unsqueeze(0)
781
- assert scale.ndim == 3
782
- B, M, K = scale.shape
783
- round_up_multiple = lambda x, m: (x + m - 1) // m * m
784
- M_padded = round_up_multiple(M, 128)
785
- K_padded = round_up_multiple(K, 4)
786
- padded_scale = torch.zeros((B, M_padded, K_padded), dtype=scale.dtype)
787
- padded_scale[:B, :M, :K] = scale
788
- batches, rows, cols = padded_scale.shape
789
- assert rows % 128 == 0
790
- assert cols % 4 == 0
791
- padded_scale = padded_scale.reshape(batches, rows // 128, 4, 32,
792
- cols // 4, 4)
793
- swizzled_scale = padded_scale.permute((0, 1, 4, 3, 2, 5))
794
- swizzled_scale = swizzled_scale.contiguous().cuda()
795
- return (swizzled_scale.reshape(M, K)
796
- if scale_ndim == 2 else swizzled_scale.reshape(B, M, K))
797
-
766
def process_weights_after_loading(self, layer: Module) -> None:
767
768
# global scales:
@@ -814,7 +782,7 @@ class ModelOptNvFp4LinearMethod(LinearMethodBase):
782
"Expected weight_scale.dim(1) to be divisible by 16")
783
assert (layer.weight_scale.dtype == torch.float8_e4m3fn), (
784
"Weight Block scale must be represented as FP8-E4M3")
817
- swizzled_weight_scale = self.swizzle_blockscale(layer.weight_scale)
785
+ swizzled_weight_scale = swizzle_blockscale(layer.weight_scale)
786
787
layer.weight_scale_swizzled = Parameter(swizzled_weight_scale,
788
requires_grad=False)
@@ -1060,29 +1028,6 @@ class ModelOptNvFp4FusedMoE(FusedMoEMethodBase):
1028
weight_loader=weight_loader)
1029
layer.register_parameter("w2_input_scale", w2_input_scale)
1030
1063
- def swizzle_blockscale(self, scale: torch.tensor):
1064
- assert (scale.dtype == torch.float8_e4m3fn)
1065
- # Pad and blockwise interleave weight_scale
1066
- scale_ndim = scale.ndim
1067
- if scale.ndim == 2:
1068
- scale = scale.unsqueeze(0)
1069
- assert scale.ndim == 3
1070
- B, M, K = scale.shape
1071
- round_up_multiple = lambda x, m: (x + m - 1) // m * m
1072
- M_padded = round_up_multiple(M, 128)
1073
- K_padded = round_up_multiple(K, 4)
1074
- padded_scale = torch.zeros((B, M_padded, K_padded), dtype=scale.dtype)
1075
- padded_scale[:B, :M, :K] = scale
1076
- batches, rows, cols = padded_scale.shape
1077
- assert rows % 128 == 0
1078
- assert cols % 4 == 0
1079
- padded_scale = padded_scale.reshape(batches, rows // 128, 4, 32,
1080
- cols // 4, 4)
1081
- swizzled_scale = padded_scale.permute((0, 1, 4, 3, 2, 5))
1082
- swizzled_scale = swizzled_scale.contiguous().cuda()
1083
- return (swizzled_scale.reshape(M, K)
1084
- if scale_ndim == 2 else swizzled_scale.reshape(B, M, K))
1085
-
1031
def process_weights_after_loading(self, layer: torch.nn.Module) -> None:
1032
# GEMM 1
1033
# The FlashInfer Cutlass fused MoE kernel expects the combined weights
@@ -1128,8 +1073,7 @@ class ModelOptNvFp4FusedMoE(FusedMoEMethodBase):
1073
"Expected weight_scale.dim(1) to be divisible by 16")
1074
assert (layer.w13_weight_scale.dtype == torch.float8_e4m3fn), (
1075
"Weight Blockscale must be represented as FP8-E4M3")
1131
- w13_blockscale_swizzled = self.swizzle_blockscale(
1132
- layer.w13_weight_scale)
1076
+ w13_blockscale_swizzled = swizzle_blockscale(layer.w13_weight_scale)
1077
1078
layer.w13_blockscale_swizzled = Parameter(w13_blockscale_swizzled,
1079
requires_grad=False)
@@ -1151,7 +1095,7 @@ class ModelOptNvFp4FusedMoE(FusedMoEMethodBase):
1095
"Expected weight_scale.dim(1) to be divisible by 16")
1096
assert (layer.w2_weight_scale.dtype == torch.float8_e4m3fn), (
1097
"Weight Blockscale must be represented as FP8-E4M3")
1154
- w2_blockscale_swizzled = self.swizzle_blockscale(layer.w2_weight_scale)
1098
+ w2_blockscale_swizzled = swizzle_blockscale(layer.w2_weight_scale)
1099
1100
layer.w2_blockscale_swizzled = Parameter(w2_blockscale_swizzled,
1101
requires_grad=False)
vllm/model_executor/layers/quantization/utils/nvfp4_emulation_utils.py
+3
-12
@@ -2,13 +2,12 @@
2
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
import torch
4
5
-from vllm._custom_ops import cutlass_scaled_mm_supports_fp4
6
-from vllm.platforms import current_platform
5
from vllm.scalar_type import scalar_types
6
7
__all__ = [
10
- "break_fp4_bytes", "dequantize_to_dtype", "ref_nvfp4_quant",
11
- "cutlass_fp4_supported"
8
+ "break_fp4_bytes",
9
+ "dequantize_to_dtype",
10
+ "ref_nvfp4_quant",
11
]
12
13
FLOAT4_E2M1_MAX = scalar_types.float4_e2m1f.max()
@@ -17,14 +16,6 @@ kE2M1ToFloat = torch.tensor([0., 0.5, 1., 1.5, 2., 3., 4., 6.],
16
dtype=torch.float32)
17
18
20
-def cutlass_fp4_supported() -> bool:
21
- if not current_platform.is_cuda():
22
- return False
23
- capability_tuple = current_platform.get_device_capability()
24
- capability = -1 if capability_tuple is None else capability_tuple.to_int()
25
- return cutlass_scaled_mm_supports_fp4(capability)
26
-
27
-
19
def break_fp4_bytes(a, dtype):
20
assert a.dtype == torch.uint8
21
m, n = a.shape
vllm/model_executor/layers/quantization/utils/quant_utils.py
+55
@@ -8,8 +8,10 @@ from typing import ClassVar, NamedTuple, Optional
8
import numpy
9
import torch
10
11
+from vllm._custom_ops import cutlass_scaled_mm_supports_fp4
12
from vllm.model_executor.layers.quantization.qqq import (
13
MARLIN_QQQ_SUPPORTED_NUM_BITS)
14
+from vllm.platforms import current_platform
15
from vllm.scalar_type import ScalarType, scalar_types
16
17
@@ -592,3 +594,56 @@ def awq_pack(
594
q_w = q_w.reshape((-1, size_n)).contiguous()
595
596
return pack_cols(q_w, num_bits, size_k, size_n)
597
+
598
+
599
+def swizzle_blockscale(scale: torch.Tensor) -> torch.Tensor:
600
+ """
601
+ Pad and block-interleave the FP4 block-scales so that they match the data
602
+ layout expected by the CUTLASS / FlashInfer kernels.
603
+
604
+ Parameters
605
+ ----------
606
+ scale: torch.Tensor
607
+
608
+ Returns
609
+ -------
610
+ torch.Tensor
611
+ The swizzled tensor with the same logical shape as *scale*.
612
+ """
613
+ assert scale.dtype == torch.float8_e4m3fn, (
614
+ "swizzle_blockscale expects the input tensor to be in "
615
+ "torch.float8_e4m3fn format.")
616
+
617
+ scale_ndim = scale.ndim
618
+ if scale_ndim == 2:
619
+ scale = scale.unsqueeze(0) # (1, M, K)
620
+ assert scale.ndim == 3, "Expected a 2-D or 3-D tensor for block scales."
621
+
622
+ B, M, K = scale.shape
623
+
624
+ def _round_up(x: int, m: int) -> int:
625
+ return (x + m - 1) // m * m
626
+
627
+ M_padded = _round_up(M, 128)
628
+ K_padded = _round_up(K, 4)
629
+
630
+ padded = torch.zeros((B, M_padded, K_padded),
631
+ dtype=scale.dtype,
632
+ device=scale.device)
633
+ padded[:B, :M, :K] = scale
634
+
635
+ # Reshape / permute to the layout required by the kernel.
636
+ padded = padded.reshape(B, M_padded // 128, 4, 32, K_padded // 4, 4)
637
+ swizzled = padded.permute(0, 1, 4, 3, 2, 5).contiguous().cuda()
638
+
639
+ if scale_ndim == 2:
640
+ return swizzled.reshape(M, K)
641
+ return swizzled.reshape(B, M, K)
642
+
643
+
644
+def cutlass_fp4_supported() -> bool:
645
+ if not current_platform.is_cuda():
646
+ return False
647
+ capability_tuple = current_platform.get_device_capability()
648
+ capability = -1 if capability_tuple is None else capability_tuple.to_int()
649
+ return cutlass_scaled_mm_supports_fp4(capability)