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

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 2, 09:45 UTC

[MoE Refactor] Migrate MoeWNA16Method quantization to MK oracle (#42647)

Signed-off-by: Bill Nell <bnell@redhat.com> Co-authored-by: Claude <noreply@anthropic.com>

bnellnm committed May 29, 2026 at 17:19 UTC 106aa92f04c0d1c3c37947fd9d4921530562a961

8 files changed +332 -97

vllm/model_executor/layers/fused_moe/experts/triton_moe.py

+30 -20

@@ -43,7 +43,10 @@ from vllm.model_executor.layers.quantization.utils.quant_utils import (

43 kFp8Static128BlockSym,

44 kFp8StaticChannelSym,

45 kFp8StaticTensorSym,

46 + kInt4Static,

47 + kInt4Static32,

48 kInt8DynamicTokenSym,

49 + kInt8Static,

50 kInt8StaticChannelSym,

51 )

52 from vllm.platforms import current_platform

@@ -459,40 +462,45 @@ class TritonExperts(LoRAExpertsMixin, mk.FusedMoEExpertsModular):

462 class TritonWNA16Experts(TritonExperts):

463 @staticmethod

464 def _supports_current_device() -> bool:

462 - raise NotImplementedError(

463 - "TritonWNA16Experts is not yet used by an Oracle. "

464 - "This method should not be called."

465 - )

465 + return current_platform.is_cuda_alike() or current_platform.is_xpu()

466

467 @staticmethod

468 def _supports_no_act_and_mul() -> bool:

469 - raise NotImplementedError(

470 - "TritonWNA16Experts is not yet used by an Oracle. "

471 - "This method should not be called."

472 - )

469 + return True

470

471 @staticmethod

472 def _supports_quant_scheme(

473 weight_key: QuantKey | None,

474 activation_key: QuantKey | None,

475 ) -> bool:

479 - raise NotImplementedError(

480 - "TritonWNA16Experts is not yet used by an Oracle. "

481 - "This method should not be called."

482 - )

476 + SUPPORTED_W = [

477 + kInt4Static,

478 + kInt8Static,

479 + kInt4Static32,

480 + # other group sizes?

481 + ]

482 + return weight_key in SUPPORTED_W

483

484 @staticmethod

485 def _supports_activation(activation: MoEActivation) -> bool:

486 - raise NotImplementedError(

487 - "TritonWNA16Experts is not yet used by an Oracle. "

488 - "This method should not be called."

489 - )

486 + return activation in [

487 + MoEActivation.SILU,

488 + MoEActivation.GELU,

489 + MoEActivation.GELU_TANH,

490 + MoEActivation.SWIGLUOAI,

491 + MoEActivation.SWIGLUSTEP,

492 + MoEActivation.SILU_NO_MUL,

493 + MoEActivation.GELU_NO_MUL,

494 + MoEActivation.GELU_TANH_NO_MUL,

495 + MoEActivation.RELU2_NO_MUL,

496 + ]

497

498 @staticmethod

499 def _supports_parallel_config(moe_parallel_config: FusedMoEParallelConfig) -> bool:

493 - raise NotImplementedError(

494 - "TritonWNA16Experts is not yet used by an Oracle. "

495 - "This method should not be called."

500 + # Why?

501 + return not (

502 + moe_parallel_config.use_fi_nvl_two_sided_kernels

503 + or moe_parallel_config.use_fi_nvl_one_sided_kernels

504 )

505

506 def apply(

@@ -515,7 +523,9 @@ class TritonWNA16Experts(TritonExperts):

523 ):

524 # Check constraints.

525 if self.quant_config.use_int4_w4a16:

518 - assert hidden_states.size(-1) // 2 == w1.size(2), "Hidden size mismatch"

526 + assert hidden_states.size(-1) // 2 == w1.size(2), (

527 + f"Hidden size mismatch {hidden_states.size(-1) // 2} == {w1.size(2)}"

528 + )

529 else:

530 assert hidden_states.size(-1) == w1.size(2), (

531 f"Hidden size mismatch {hidden_states.size(-1)} != {w1.size(2)}"

vllm/model_executor/layers/fused_moe/fused_moe.py

+6 -1

@@ -30,6 +30,7 @@ from vllm.model_executor.layers.fused_moe.utils import (

30 )

31 from vllm.platforms import current_platform

32 from vllm.triton_utils import tl, triton

33 +from vllm.utils.math_utils import next_power_of_2

34 from vllm.utils.torch_utils import direct_register_custom_op

36 logger = init_logger(__name__)

@@ -1243,7 +1244,11 @@ def get_default_config(

1244 bit = 4 if dtype == "int4_w4a16" else 8

1245 use_moe_wna16_cuda = should_moe_wna16_use_cuda(M * topk, block_shape[1], E, bit)

1246 if use_moe_wna16_cuda:

1246 - config = {"BLOCK_SIZE_M": min(16, M), "SPLIT_K": 1}

1247 + config = {

1248 + "BLOCK_SIZE_M": min(16, next_power_of_2(M)),

1249 + "GROUP_SIZE_M": 1,

1250 + "SPLIT_K": 1,

1251 + }

1252 elif M <= 20:

1253 config = {"BLOCK_SIZE_M": 16, "GROUP_SIZE_M": 1, "SPLIT_K": 1}

1254 elif M <= 40:

vllm/model_executor/layers/fused_moe/oracle/int_wna16.py

+91 -6

@@ -23,6 +23,9 @@ from vllm.model_executor.layers.fused_moe.experts.marlin_moe import (

23 MarlinExperts,

24 MarlinExpertsBase,

25 )

26 +from vllm.model_executor.layers.fused_moe.experts.triton_moe import (

27 + TritonWNA16Experts,

28 +)

29 from vllm.model_executor.layers.fused_moe.experts.trtllm_mxint4_moe import (

30 TrtLlmMxint4ExpertsMonolithic,

31 )

@@ -31,6 +34,7 @@ from vllm.model_executor.layers.quantization.utils.marlin_utils import (

34 marlin_act_int8_process_scales,

35 marlin_moe_permute_scales,

36 marlin_permute_bias,

37 + marlin_zero_points,

38 moe_awq_to_marlin_zero_points,

39 )

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

@@ -45,6 +49,7 @@ class WNA16MoEBackend(Enum):

49 MARLIN = "MARLIN"

50 BATCHED_MARLIN = "BATCHED_MARLIN"

51 FLASHINFER_TRTLLM = "FLASHINFER_TRTLLM"

52 + TRITON = "TRITON"

53 XPU = "XPU"

@@ -58,6 +63,8 @@ def backend_to_kernel_cls(

63 return [BatchedMarlinExperts]

64 elif backend == WNA16MoEBackend.FLASHINFER_TRTLLM:

65 return [TrtLlmMxint4ExpertsMonolithic]

66 + elif backend == WNA16MoEBackend.TRITON:

67 + return [TritonWNA16Experts]

68 elif backend == WNA16MoEBackend.XPU:

69 from vllm.model_executor.layers.fused_moe.experts.xpu_moe import (

70 XPUExpertsWNA16,

@@ -68,24 +75,38 @@ def backend_to_kernel_cls(

75 raise ValueError(f"Unknown WNA16 MoE backend: {backend.value}")

71 -def _get_priority_backends() -> list[WNA16MoEBackend]:

78 +def _get_priority_backends(

79 + may_have_zp: bool, may_have_bias: bool

80 +) -> list[WNA16MoEBackend]:

81 """

82 Get available backends in priority order based on platform and config.

83 """

84 if current_platform.is_xpu():

85 return [WNA16MoEBackend.XPU]

78 - _AVAILABLE_BACKENDS = [

79 - WNA16MoEBackend.FLASHINFER_TRTLLM,

87 + _AVAILABLE_BACKENDS = []

88 +

89 + if not may_have_zp and not may_have_bias:

90 + _AVAILABLE_BACKENDS.append(WNA16MoEBackend.FLASHINFER_TRTLLM)

91 +

92 + # Marlin supports ZP and bias

93 + _AVAILABLE_BACKENDS += [

94 WNA16MoEBackend.MARLIN,

95 WNA16MoEBackend.BATCHED_MARLIN,

96 ]

97 return _AVAILABLE_BACKENDS

99 + if not may_have_bias:

100 + _AVAILABLE_BACKENDS.append(WNA16MoEBackend.TRITON)

101 +

102 + return _AVAILABLE_BACKENDS

103 +

104

105 def select_wna16_moe_backend(

106 config: FusedMoEConfig,

107 weight_key: QuantKey,

108 + may_have_zp: bool,

109 + may_have_bias: bool,

110 ) -> tuple[WNA16MoEBackend, type[mk.FusedMoEExperts]]:

111 """Select the WNA16 MoE backend.

112

@@ -136,7 +157,7 @@ def select_wna16_moe_backend(

157 raise ValueError(_make_log_unsupported(backend, reason))

158

159 # Select kernels in order of backend.

139 - AVAILABLE_BACKENDS = _get_priority_backends()

160 + AVAILABLE_BACKENDS = _get_priority_backends(may_have_zp, may_have_bias)

161

162 for backend in AVAILABLE_BACKENDS:

163 activation_key = None # always BF16 activation for WNA16 MoE

@@ -218,6 +239,7 @@ def make_wna16_moe_kernel(

239 assert experts_cls in (

240 MarlinExperts,

241 BatchedMarlinExperts,

242 + TritonWNA16Experts,

243 TrtLlmMxint4ExpertsMonolithic,

244 XPUExpertsWNA16,

245 )

@@ -234,6 +256,7 @@ def make_wna16_moe_kernel(

256 assert prepare_finalize is not None

257

258 logger.info_once("Using %s", prepare_finalize.__class__.__name__, scope="local")

259 + logger.info_once("Using %s", experts_cls.__name__, scope="local")

260

261 extra_args: dict[str, Any] = {}

262 if issubclass(experts_cls, MarlinExpertsBase):

@@ -403,6 +426,8 @@ def _process_weights_marlin(

426 w2_input_global_scale: torch.Tensor | None = None

427 w13_bias_out: torch.Tensor | None = None

428 w2_bias_out: torch.Tensor | None = None

429 + w13_qzeros_out: torch.Tensor | None = None

430 + w2_qzeros_out: torch.Tensor | None = None

431

432 # --- FP8 weight / scale adjustment ---

433 if input_dtype == torch.float8_e4m3fn:

@@ -502,6 +527,24 @@ def _process_weights_marlin(

527 if w2_bias is not None:

528 w2_bias_out = marlin_permute_bias(w2_bias)

529

530 + if w13_qzeros is not None:

531 + w13_qzeros_out = marlin_zero_points(

532 + w13_qzeros,

533 + size_k=layer.intermediate_size_per_partition,

534 + size_n=w13_qzeros.shape[2],

535 + num_bits=num_bits,

536 + is_a_8bit=is_a_8bit,

537 + )

538 +

539 + if w2_qzeros is not None:

540 + w2_qzeros_out = marlin_zero_points(

541 + w2_qzeros,

542 + size_k=w2_qzeros.shape[1] * group_size_or_pack_factor,

543 + size_n=w2_qzeros.shape[2],

544 + num_bits=num_bits,

545 + is_a_8bit=is_a_8bit,

546 + )

547 +

548 return (

549 marlin_w13_qweight,

550 marlin_w2_qweight,

@@ -511,8 +554,8 @@ def _process_weights_marlin(

554 w2_g_idx,

555 w13_g_idx_sort_indices,

556 w2_g_idx_sort_indices,

514 - w13_qzeros,

515 - w2_qzeros,

557 + w13_qzeros_out,

558 + w2_qzeros_out,

559 w13_input_global_scale,

560 w2_input_global_scale,

561 w13_bias_out,

@@ -780,6 +823,9 @@ def convert_to_wna16_moe_kernel_format(

823 from vllm.model_executor.layers.quantization.awq_marlin import (

824 AWQMarlinConfig,

825 )

826 + from vllm.model_executor.layers.quantization.moe_wna16 import (

827 + MoeWNA16Config,

828 + )

829

830 if isinstance(quant_config, AWQMarlinConfig):

831 if w13_qzeros is None or w2_qzeros is None:

@@ -814,11 +860,17 @@ def convert_to_wna16_moe_kernel_format(

860 pack_factor = 32 // quant_config.num_bits

861 group_size = quant_config.group_size

862 actorder = quant_config.actorder

863 + elif isinstance(quant_config, MoeWNA16Config):

864 + num_bits = quant_config.weight_bits

865 + pack_factor = quant_config.bit8_pack_factor

866 + group_size = quant_config.group_size

867 + actorder = None

868 else:

869 raise TypeError(

870 "Marlin WNA16 MoE backend requires AutoGPTQConfig, AWQMarlinConfig or "

871 f"QuantizationArgs, got {type(quant_config).__name__}."

872 )

873 +

874 if w13_g_idx is None or w2_g_idx is None:

875 raise ValueError("GPTQ Marlin MoE requires g_idx tensors.")

876 return _process_weights_marlin(

@@ -850,6 +902,39 @@ def convert_to_wna16_moe_kernel_format(

902 w13_bias,

903 w2_bias,

904 )

905 + elif backend == WNA16MoEBackend.TRITON:

906 + # Convert from int32 to uint8 format for Triton kernel.

907 + # This changes the shape from (E, N, K // 8) to (E, N, K // 2) for int4,

908 + # which matches what the Triton kernel expects.

909 + w13_uint8 = w13.view(torch.uint8)

910 + w2_uint8 = w2.view(torch.uint8)

911 + return (

912 + w13_uint8,

913 + w2_uint8,

914 + w13_scale,

915 + w2_scale,

916 + None,

917 + None,

918 + None,

919 + None,

920 + w13_qzeros,

921 + w2_qzeros,

922 + None,

923 + None,

924 + w13_bias,

925 + w2_bias,

926 + )

927 + elif backend == WNA16MoEBackend.FLASHINFER_TRTLLM:

928 + return _process_weights_flashinfer(

929 + w13,

930 + w2,

931 + w13_scale,

932 + w2_scale,

933 + w13_g_idx,

934 + w2_g_idx,

935 + w13_bias,

936 + w2_bias,

937 + )

938 elif backend == WNA16MoEBackend.XPU:

939 assert quant_config is not None

940 (

vllm/model_executor/layers/quantization/auto_gptq.py

+26 -38

@@ -485,6 +485,8 @@ class AutoGPTQMoEMethod(FusedMoEMethodBase):

485 self.wna16_moe_backend, self.experts_cls = select_wna16_moe_backend(

486 moe,

487 weight_key,

488 + may_have_zp=True,

489 + may_have_bias=True,

490 )

491

492 def create_weights(

@@ -644,12 +646,22 @@ class AutoGPTQMoEMethod(FusedMoEMethodBase):

646 layer.workspace = marlin_make_workspace_new(device, 4)

647

648 def process_weights_after_loading(self, layer: RoutedExperts) -> None:

649 + def replace_or_register(name: str, val: torch.Tensor | None):

650 + if val is None:

651 + return

652 +

653 + if hasattr(layer, name):

654 + replace_parameter(layer, name, val)

655 + else:

656 + layer.register_parameter(

657 + name, torch.nn.Parameter(val, requires_grad=False)

658 + )

659 +

660 is_a_8bit = self.input_dtype is not None and self.input_dtype.itemsize == 1

661

649 - if is_a_8bit:

650 - assert self.quant_config.quant_type.size_bits == 8, (

651 - "W8A8-INT8 is not supported by marlin kernel."

652 - )

662 + assert not is_a_8bit or self.quant_config.quant_type.size_bits == 8, (

663 + "W8A8-INT8 is not supported by marlin kernel."

664 + )

665

666 (

667 w13,

@@ -660,8 +672,8 @@ class AutoGPTQMoEMethod(FusedMoEMethodBase):

672 w2_g_idx,

673 w13_g_idx_sort_indices,

674 w2_g_idx_sort_indices,

663 - _w13_qzeros,

664 - _w2_qzeros,

675 + w13_qzeros,

676 + w2_qzeros,

677 w13_input_global_scale,

678 w2_input_global_scale,

679 w13_bias,

@@ -679,6 +691,8 @@ class AutoGPTQMoEMethod(FusedMoEMethodBase):

691 w2_g_idx=layer.w2_g_idx,

692 w13_bias=getattr(layer, "w13_bias", None),

693 w2_bias=getattr(layer, "w2_bias", None),

694 + w13_qzeros=getattr(layer, "w13_qzeros", None),

695 + w2_qzeros=getattr(layer, "w2_qzeros", None),

696 )

697

698 replace_parameter(layer, "w13_qweight", w13)

@@ -689,38 +703,12 @@ class AutoGPTQMoEMethod(FusedMoEMethodBase):

703 replace_parameter(layer, "w2_g_idx", w2_g_idx)

704 replace_parameter(layer, "w13_g_idx_sort_indices", w13_g_idx_sort_indices)

705 replace_parameter(layer, "w2_g_idx_sort_indices", w2_g_idx_sort_indices)

692 - if w13_input_global_scale is not None:

693 - if hasattr(layer, "w13_input_global_scale"):

694 - replace_parameter(

695 - layer, "w13_input_global_scale", w13_input_global_scale

696 - )

697 - else:

698 - layer.register_parameter(

699 - "w13_input_global_scale",

700 - torch.nn.Parameter(w13_input_global_scale, requires_grad=False),

701 - )

702 - if w2_input_global_scale is not None:

703 - if hasattr(layer, "w2_input_global_scale"):

704 - replace_parameter(layer, "w2_input_global_scale", w2_input_global_scale)

705 - else:

706 - layer.register_parameter(

707 - "w2_input_global_scale",

708 - torch.nn.Parameter(w2_input_global_scale, requires_grad=False),

709 - )

710 - if w13_bias is not None:

711 - if hasattr(layer, "w13_bias"):

712 - replace_parameter(layer, "w13_bias", w13_bias)

713 - else:

714 - layer.register_parameter(

715 - "w13_bias", torch.nn.Parameter(w13_bias, requires_grad=False)

716 - )

717 - if w2_bias is not None:

718 - if hasattr(layer, "w2_bias"):

719 - replace_parameter(layer, "w2_bias", w2_bias)

720 - else:

721 - layer.register_parameter(

722 - "w2_bias", torch.nn.Parameter(w2_bias, requires_grad=False)

723 - )

706 + replace_or_register("w13_input_global_scale", w13_input_global_scale)

707 + replace_or_register("w2_input_global_scale", w2_input_global_scale)

708 + replace_or_register("w13_bias", w13_bias)

709 + replace_or_register("w2_bias", w2_bias)

710 + replace_or_register("w13_qzeros", w13_qzeros)

711 + replace_or_register("w2_qzeros", w2_qzeros)

712

713 self._setup_kernel(layer)

714

vllm/model_executor/layers/quantization/awq_marlin.py

@@ -524,6 +524,8 @@ class AWQMarlinMoEMethod(FusedMoEMethodBase):

524 self.wna16_moe_backend, self.experts_cls = select_wna16_moe_backend(

525 moe,

526 kInt4Static,

527 + may_have_zp=self.quant_config.zero_point,

528 + may_have_bias=True,

529 )

530

531 def create_weights(

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

+10 -7

@@ -88,7 +88,13 @@ class CompressedTensorsWNA16MarlinMoEMethod(CompressedTensorsMoEMethod):

88 self.wna16_backend, self.experts_cls = select_wna16_moe_backend(

89 config=self.moe,

90 weight_key=weight_key,

91 + may_have_zp=False,

92 + may_have_bias=False,

93 )

94 + self.is_marlin = self.wna16_backend in [

95 + WNA16MoEBackend.MARLIN,

96 + WNA16MoEBackend.BATCHED_MARLIN,

97 + ]

99 def get_weight_shape(

100 self,

@@ -114,7 +120,6 @@ class CompressedTensorsWNA16MarlinMoEMethod(CompressedTensorsMoEMethod):

120 "num_groups_w2 must be provided for weight scales"

121 )

122 w13_num_shards = 2 if self.moe.is_act_and_mul else 1

117 - is_flashinfer = self.wna16_backend == WNA16MoEBackend.FLASHINFER_TRTLLM

123 shape_map = {

124 "w13_weight": {

125 "Flashinfer": (

@@ -157,7 +162,7 @@ class CompressedTensorsWNA16MarlinMoEMethod(CompressedTensorsMoEMethod):

162 "Marlin": (num_experts, num_groups_w2, hidden_size),

163 },

164 }

160 - backend_key = "Flashinfer" if is_flashinfer else "Marlin"

165 + backend_key = "Marlin" if self.is_marlin else "Flashinfer"

166 return shape_map[weight_name][backend_key]

167

168 def create_weights(

@@ -174,9 +179,8 @@ class CompressedTensorsWNA16MarlinMoEMethod(CompressedTensorsMoEMethod):

179 # Will transpose the loaded weight along the

180 # intermediate and hidden dim sizes. Will

181 # shard for TP along the transposed dims

177 - is_transposed = self.wna16_backend != WNA16MoEBackend.FLASHINFER_TRTLLM

182 extra_weight_attrs.update(

179 - {"is_transposed": is_transposed, "quant_method": self.strategy}

183 + {"is_transposed": self.is_marlin, "quant_method": self.strategy}

184 )

185

186 w13_weight = torch.nn.Parameter(

@@ -324,7 +328,6 @@ class CompressedTensorsWNA16MarlinMoEMethod(CompressedTensorsMoEMethod):

328

329 def process_weights_after_loading(self, layer: torch.nn.Module) -> None:

330 # Process weights using the shared oracle infrastructure

327 - is_flashinfer = self.wna16_backend == WNA16MoEBackend.FLASHINFER_TRTLLM

331 (

332 w13_qweight,

333 w2_qweight,

@@ -360,7 +363,7 @@ class CompressedTensorsWNA16MarlinMoEMethod(CompressedTensorsMoEMethod):

363 replace_parameter(layer, "w2_weight_scale", w2_scales)

364

365 # Marlin-specific parameters (not needed for Flashinfer)

363 - if not is_flashinfer:

366 + if self.is_marlin:

367 replace_parameter(layer, "w13_weight_g_idx", w13_g_idx_processed)

368 replace_parameter(layer, "w2_weight_g_idx", w2_g_idx_processed)

369 replace_parameter(layer, "w13_g_idx_sort_indices", w13_g_idx_sort_indices)

@@ -392,7 +395,7 @@ class CompressedTensorsWNA16MarlinMoEMethod(CompressedTensorsMoEMethod):

395

396 # Add Marlin-specific arguments

397 marlin_args: dict[str, Any] = {}

395 - if not is_flashinfer:

398 + if self.is_marlin:

399 marlin_args = {

400 "w13_g_idx": layer.w13_weight_g_idx,

401 "w2_g_idx": layer.w2_weight_g_idx,

vllm/model_executor/layers/quantization/moe_wna16.py

+150 -24

@@ -13,11 +13,15 @@ from vllm.model_executor.layers.fused_moe import (

13 RoutedExperts,

14 SharedExperts,

15 )

16 -from vllm.model_executor.layers.fused_moe.activation import MoEActivation

16 from vllm.model_executor.layers.fused_moe.config import (

17 FusedMoEQuantConfig,

19 - int4_w4a16_moe_quant_config,

20 - int8_w8a16_moe_quant_config,

18 +)

19 +from vllm.model_executor.layers.fused_moe.oracle.int_wna16 import (

20 + WNA16MoEBackend,

21 + convert_to_wna16_moe_kernel_format,

22 + make_wna16_moe_kernel,

23 + make_wna16_moe_quant_config,

24 + select_wna16_moe_backend,

25 )

26 from vllm.model_executor.layers.fused_moe.unquantized_fused_moe_method import (

27 UnquantizedFusedMoEMethod,

@@ -31,7 +35,15 @@ from vllm.model_executor.layers.quantization.base_config import (

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

36 check_marlin_supports_layer,

37 )

34 -from vllm.model_executor.utils import set_weight_attrs

38 +from vllm.model_executor.layers.quantization.utils.quant_utils import (

39 + INT4_DTYPE,

40 + INT8_DTYPE,

41 + QuantKey,

42 + kInt4Static32GroupScale,

43 + kInt4StaticGroupScale,

44 + kInt8StaticGroupScale,

45 +)

46 +from vllm.model_executor.utils import replace_parameter, set_weight_attrs

47 from vllm.platforms import current_platform

@@ -216,6 +228,33 @@ class MoeWNA16Method(FusedMoEMethodBase):

228 super().__init__(moe)

229 self.quant_config = quant_config

230

231 + num_bits = self.quant_config.weight_bits

232 + group_size = self.quant_config.group_size

233 +

234 + if num_bits == 4:

235 + quant_type = INT4_DTYPE

236 + if group_size == 32:

237 + scale = kInt4Static32GroupScale

238 + else:

239 + scale = kInt4StaticGroupScale

240 + elif num_bits == 8:

241 + assert group_size == -1

242 + quant_type = INT8_DTYPE

243 + scale = kInt8StaticGroupScale

244 + else:

245 + raise ValueError("MoeWNA16Method only supports int4 and int8 now.")

246 +

247 + weight_key = QuantKey(quant_type, scale)

248 +

249 + # Select WNA16 MoE backend via oracle.

250 + # handle ZP?

251 + self.wna16_backend, self.experts_cls = select_wna16_moe_backend(

252 + config=self.moe,

253 + weight_key=weight_key,

254 + may_have_zp=self.quant_config.has_zp,

255 + may_have_bias=False,

256 + )

257 +

258 def create_weights(

259 self,

260 layer: RoutedExperts,

@@ -336,24 +375,89 @@ class MoeWNA16Method(FusedMoEMethodBase):

375 layer.register_parameter(key, param)

376 set_weight_attrs(param, extra_weight_attrs)

377

378 + def process_weights_after_loading(self, layer: torch.nn.Module) -> None:

379 + has_zp = self.quant_config.has_zp

380 + (

381 + w13_qweight,

382 + w2_qweight,

383 + w13_scales,

384 + w2_scales,

385 + w13_g_idx_processed,

386 + w2_g_idx_processed,

387 + w13_g_idx_sort_indices,

388 + w2_g_idx_sort_indices,

389 + w13_qzeros,

390 + w2_qzeros,

391 + w13_input_global_scale,

392 + w2_input_global_scale,

393 + _, # w13_bias

394 + _, # w2_bias

395 + ) = convert_to_wna16_moe_kernel_format(

396 + backend=self.wna16_backend,

397 + layer=layer,

398 + quant_config=self.quant_config,

399 + input_dtype=None,

400 + w13=layer.w13_qweight,

401 + w2=layer.w2_qweight,

402 + w13_scale=layer.w13_scales,

403 + w2_scale=layer.w2_scales,

404 + w13_g_idx=getattr(layer, "w13_g_idx", None),

405 + w2_g_idx=getattr(layer, "w2_g_idx", None),

406 + w13_qzeros=layer.w13_qzeros if has_zp else None,

407 + w2_qzeros=layer.w2_qzeros if has_zp else None,

408 + )

409 +

410 + # Replace common parameters

411 + replace_parameter(layer, "w13_qweight", w13_qweight)

412 + replace_parameter(layer, "w2_qweight", w2_qweight)

413 + replace_parameter(layer, "w13_scales", w13_scales)

414 + replace_parameter(layer, "w2_scales", w2_scales)

415 +

416 + if has_zp:

417 + assert w13_qzeros is not None and w2_qzeros is not None

418 + replace_parameter(layer, "w13_qzeros", w13_qzeros)

419 + replace_parameter(layer, "w2_qzeros", w2_qzeros)

420 +

421 + # Marlin-specific parameters (not needed for Flashinfer)

422 + if self.wna16_backend != WNA16MoEBackend.FLASHINFER_TRTLLM:

423 + replace_parameter(layer, "w13_g_idx", w13_g_idx_processed)

424 + replace_parameter(layer, "w2_g_idx", w2_g_idx_processed)

425 + replace_parameter(layer, "w13_g_idx_sort_indices", w13_g_idx_sort_indices)

426 + replace_parameter(layer, "w2_g_idx_sort_indices", w2_g_idx_sort_indices)

427 +

428 + # Register input global scales if present

429 + if w13_input_global_scale is not None:

430 + layer.register_parameter(

431 + "w13_input_global_scale",

432 + torch.nn.Parameter(w13_input_global_scale, requires_grad=False),

433 + )

434 + if w2_input_global_scale is not None:

435 + layer.register_parameter(

436 + "w2_input_global_scale",

437 + torch.nn.Parameter(w2_input_global_scale, requires_grad=False),

438 + )

439 +

440 + assert self.experts_cls is not None

441 + self.moe_quant_config = self.get_fused_moe_quant_config(layer)

442 + assert self.moe_quant_config is not None

443 + self.moe_kernel = make_wna16_moe_kernel(

444 + moe_quant_config=self.moe_quant_config,

445 + moe_config=self.moe,

446 + experts_cls=self.experts_cls,

447 + routing_tables=layer._expert_routing_tables(),

448 + )

449 +

450 def get_fused_moe_quant_config(

451 self, layer: RoutedExperts

452 ) -> FusedMoEQuantConfig | None:

342 - weight_bits = self.quant_config.weight_bits

453 has_zp = self.quant_config.has_zp

344 - assert weight_bits == 4 or weight_bits == 8

345 - config_builder = (

346 - int4_w4a16_moe_quant_config

347 - if weight_bits == 4

348 - else int8_w8a16_moe_quant_config

349 - )

350 -

351 - return config_builder(

454 + return make_wna16_moe_quant_config(

455 w1_scale=layer.w13_scales,

456 w2_scale=layer.w2_scales,

457 w1_zp=layer.w13_qzeros if has_zp else None,

458 w2_zp=layer.w2_qzeros if has_zp else None,

356 - block_shape=[0, layer.group_size],

459 + group_size=layer.group_size,

460 + num_bits=self.quant_config.weight_bits,

461 )

462

463 def apply(

@@ -365,22 +469,44 @@ class MoeWNA16Method(FusedMoEMethodBase):

469 shared_experts: SharedExperts | None,

470 shared_experts_input: torch.Tensor | None,

471 ) -> torch.Tensor:

368 - from vllm.model_executor.layers.fused_moe import fused_experts

369 -

370 - assert layer.activation == MoEActivation.SILU, (

371 - f"Only SiLU activation is supported, not {layer.activation}."

372 - )

373 -

374 - return fused_experts(

472 + assert not self.is_monolithic

473 + assert self.moe_kernel is not None

474 + return self.moe_kernel.apply(

475 x,

476 layer.w13_qweight,

477 layer.w2_qweight,

378 - topk_weights=topk_weights,

379 - topk_ids=topk_ids,

478 + topk_weights,

479 + topk_ids,

480 + activation=layer.activation,

481 + global_num_experts=layer.global_num_experts,

482 + expert_map=layer.expert_map,

483 apply_router_weight_on_input=layer.apply_router_weight_on_input,

484 + shared_experts=shared_experts,

485 + shared_experts_input=shared_experts_input,

486 + )

487 +

488 + def apply_monolithic(

489 + self,

490 + layer: RoutedExperts,

491 + x: torch.Tensor,

492 + router_logits: torch.Tensor,

493 + input_ids: torch.Tensor | None = None,

494 + ) -> torch.Tensor:

495 + assert self.is_monolithic

496 + assert self.moe_kernel is not None

497 + return self.moe_kernel.apply_monolithic(

498 + x,

499 + layer.w13_weight,

500 + layer.w2_weight,

501 + router_logits,

502 + activation=layer.activation,

503 global_num_experts=layer.global_num_experts,

504 expert_map=layer.expert_map,

383 - quant_config=self.moe_quant_config,

505 + apply_router_weight_on_input=layer.apply_router_weight_on_input,

506 + num_expert_group=layer.num_expert_group,

507 + topk_group=layer.topk_group,

508 + e_score_correction_bias=layer.e_score_correction_bias,

509 + routed_scaling_factor=layer.routed_scaling_factor,

510 )

511

512 @staticmethod

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

+17 -1

@@ -3,7 +3,7 @@

3 from collections.abc import Mapping

4 from copy import deepcopy

5 from types import MappingProxyType

6 -from typing import TYPE_CHECKING

6 +from typing import TYPE_CHECKING, Any

8 import regex as re

9 import torch

@@ -69,6 +69,20 @@ def get_dynamic_override(

69 return default_value

72 +def flatten_list(lst: list[Any]) -> list[Any]:

73 + output = []

74 +

75 + def _flatten(lst: list[Any]):

76 + for i in lst:

77 + if isinstance(i, list):

78 + _flatten(i)

79 + else:

80 + output.append(i)

81 +

82 + _flatten(lst)

83 + return output

84 +

85 +

86 def is_layer_gptq_quantized(

87 prefix: str,

88 quantized_layers: list[str],

@@ -83,6 +97,8 @@ def is_layer_gptq_quantized(

98 proj_name = prefix.split(".")[-1]

100 + quantized_layers = flatten_list(quantized_layers)

101 +

102 # Fused layers like gate_up_proj or qkv_proj will not be fused

103 # in the safetensors checkpoint. So, we convert the name

104 # from the fused version to unfused + check to make sure that