diff --git a/tests/quantization/test_auto_gptq.py b/tests/quantization/test_auto_gptq.py index b733ee486c18..9eefd88e915e 100644 --- a/tests/quantization/test_auto_gptq.py +++ b/tests/quantization/test_auto_gptq.py @@ -5,13 +5,17 @@ Run `pytest tests/quantization/test_auto_gptq.py -v -s`. """ +from types import SimpleNamespace + import pytest import torch from tests.quantization.utils import is_quant_method_supported +from vllm.model_executor.layers.fused_moe import RoutedExperts from vllm.model_executor.layers.quantization.auto_gptq import ( AutoGPTQConfig, AutoGPTQLinearMethod, + AutoGPTQMoEMethod, ) PROMPT = "On the surface of Mars, we found" @@ -54,3 +58,78 @@ def check_model(model): def test_auto_gptq_config_get_name(): """Test that AutoGPTQConfig.get_name() returns 'auto_gptq'.""" assert AutoGPTQConfig.get_name() == "auto_gptq" + + +def test_auto_gptq_moe_creates_zero_initialized_expert_biases(): + method = object.__new__(AutoGPTQMoEMethod) + method.quant_config = AutoGPTQConfig(4, 128, False, True, False, {}, {}) + method.input_dtype = None + method.experts_cls = None + layer = torch.nn.Module() + + method.create_weights( + layer=layer, + num_experts=2, + hidden_size=8, + intermediate_size_per_partition=4, + params_dtype=torch.float16, + intermediate_size_full=4, + weight_loader=lambda *args, **kwargs: None, + ) + + assert layer.w13_bias.shape == (2, 8) + assert layer.w2_bias.shape == (2, 8) + assert torch.count_nonzero(layer.w13_bias) == 0 + assert torch.count_nonzero(layer.w2_bias) == 0 + + +def test_routed_experts_loads_per_expert_biases(): + class Loader: + quant_config = None + quant_method = object() + moe_config = SimpleNamespace( + is_act_and_mul=True, + tp_rank=0, + moe_parallel_config=SimpleNamespace(tp_size=1), + ) + _get_hidden_dim = staticmethod(RoutedExperts._get_hidden_dim) + _narrow_expert_data_for_padding = staticmethod( + RoutedExperts._narrow_expert_data_for_padding + ) + _load_w13 = RoutedExperts._load_w13 + _loaded_expert_biases = set() + + @staticmethod + def _map_global_expert_id_to_local_expert_id(expert_id): + return expert_id + + loader = Loader() + w13_bias = torch.nn.Parameter(torch.zeros(1, 8), requires_grad=False) + w2_bias = torch.nn.Parameter(torch.zeros(1, 4), requires_grad=False) + + for shard_id, loaded in ( + ("w1", torch.tensor([1.0, 2.0, 3.0, 4.0])), + ("w3", torch.tensor([5.0, 6.0, 7.0, 8.0])), + ): + assert RoutedExperts.weight_loader( + loader, + w13_bias, + loaded, + weight_name="model.layers.0.mlp.experts.w13_bias", + shard_id=shard_id, + expert_id=0, + return_success=True, + ) + + assert RoutedExperts.weight_loader( + loader, + w2_bias, + torch.tensor([9.0, 10.0, 11.0, 12.0]), + weight_name="model.layers.0.mlp.experts.w2_bias", + shard_id="w2", + expert_id=0, + return_success=True, + ) + assert torch.equal(w13_bias, torch.arange(1, 9, dtype=torch.float32).reshape(1, 8)) + assert torch.equal(w2_bias, torch.arange(9, 13, dtype=torch.float32).reshape(1, 4)) + assert loader._loaded_expert_biases == {"w13_bias", "w2_bias"} diff --git a/tests/quantization/test_moe_wna16.py b/tests/quantization/test_moe_wna16.py index c4b0ab5a8464..16c68fd1c79f 100644 --- a/tests/quantization/test_moe_wna16.py +++ b/tests/quantization/test_moe_wna16.py @@ -5,45 +5,194 @@ import pytest import torch +from compressed_tensors.quantization import ( + ActivationOrdering, + QuantizationArgs, + QuantizationStrategy, + QuantizationType, +) -from vllm.model_executor.layers.fused_moe.activation import MoEActivation -from vllm.model_executor.layers.quantization.moe_wna16 import MoeWNA16Method -from vllm.platforms import current_platform +from vllm.model_executor.layers.fused_moe.oracle.int_wna16 import ( + WNA16MoEBackend, + _backend_incompatibility_reason, + _convert_moe_wna16_humming_tensors, + convert_to_wna16_moe_kernel_format, + map_wna16_backend, +) +from vllm.model_executor.layers.quantization import moe_wna16 +from vllm.model_executor.layers.quantization.auto_awq import AutoAWQConfig +from vllm.model_executor.layers.quantization.auto_gptq import AutoGPTQConfig +from vllm.model_executor.layers.quantization.moe_wna16 import ( + MoeWNA16Config, + MoeWNA16Method, +) -@pytest.mark.skipif(not current_platform.is_cuda(), reason="Only test on CUDA") -def test_moe_wna16_apply_passes_layer_activation(monkeypatch): - captured_kwargs = {} +def test_map_wna16_backend_supports_triton(): + assert map_wna16_backend("triton") == WNA16MoEBackend.TRITON - def fake_fused_experts(*args, **kwargs): - captured_kwargs.update(kwargs) - return torch.empty(1, 2) - monkeypatch.setattr( - "vllm.model_executor.layers.fused_moe.fused_experts", - fake_fused_experts, +@pytest.mark.parametrize( + ("backend", "quant_config", "may_have_zp", "may_have_bias", "expected"), + [ + ( + WNA16MoEBackend.TRITON, + AutoAWQConfig(4, 128, True, False), + True, + False, + "AutoAWQ weight layout", + ), + ( + WNA16MoEBackend.TRITON, + AutoGPTQConfig(4, 128, True, True, False, {}, {}), + False, + False, + "activation ordering", + ), + ( + WNA16MoEBackend.TRITON, + QuantizationArgs( + num_bits=4, + type=QuantizationType.INT, + strategy=QuantizationStrategy.GROUP, + symmetric=True, + dynamic=False, + group_size=128, + actorder=ActivationOrdering.GROUP, + ), + False, + False, + "activation ordering", + ), + ( + WNA16MoEBackend.TRITON, + AutoGPTQConfig(4, 128, False, True, False, {}, {}), + False, + True, + "bias", + ), + ( + WNA16MoEBackend.MARLIN, + MoeWNA16Config( + linear_quant_method="gptq", + weight_bits=4, + group_size=128, + has_zp=False, + lm_head_quantized=False, + modules_to_not_convert=None, + full_config={}, + ), + False, + False, + "MoeWNA16 checkpoint layout", + ), + ], +) +def test_wna16_oracle_rejects_incompatible_quant_structures( + backend, quant_config, may_have_zp, may_have_bias, expected +): + reason = _backend_incompatibility_reason( + backend=backend, + quant_config=quant_config, + may_have_zp=may_have_zp, + may_have_bias=may_have_bias, + ) + + assert reason is not None + assert expected in reason + + +def test_compressed_tensors_weights_are_transposed_for_triton(): + quant_config = QuantizationArgs( + num_bits=4, + type=QuantizationType.INT, + strategy=QuantizationStrategy.GROUP, + symmetric=True, + dynamic=False, + group_size=32, ) + w13 = torch.arange(16, dtype=torch.int32).reshape(1, 2, 8) + w2 = torch.arange(12, dtype=torch.int32).reshape(1, 2, 6) + w13_scale = torch.arange(32, dtype=torch.float16).reshape(1, 4, 8) + w2_scale = torch.arange(18, dtype=torch.float16).reshape(1, 3, 6) + converted = convert_to_wna16_moe_kernel_format( + backend=WNA16MoEBackend.TRITON, + layer=torch.nn.Module(), + quant_config=quant_config, + input_dtype=None, + w13=w13, + w2=w2, + w13_scale=w13_scale, + w2_scale=w2_scale, + ) + + assert converted is not None + assert torch.equal(converted[0], w13.transpose(1, 2).contiguous().view(torch.uint8)) + assert torch.equal(converted[1], w2.transpose(1, 2).contiguous().view(torch.uint8)) + assert torch.equal(converted[2], w13_scale.transpose(1, 2).contiguous()) + assert torch.equal(converted[3], w2_scale.transpose(1, 2).contiguous()) + + +def test_moe_wna16_setup_forwards_selected_backend(monkeypatch): method = object.__new__(MoeWNA16Method) - method.moe = SimpleNamespace(disable_inplace=False) - method.moe_quant_config = object() - layer = SimpleNamespace( - w13_qweight=torch.empty(1, 2), - w2_qweight=torch.empty(1, 2), - activation=MoEActivation.GELU_TANH, - apply_router_weight_on_input=False, - global_num_experts=1, - expert_map=None, + method.experts_cls = object + method.wna16_backend = WNA16MoEBackend.HUMMING + method.moe = object() + quant_config = object() + method.get_fused_moe_quant_config = lambda layer: quant_config + layer = SimpleNamespace(_expert_routing_tables=lambda: (None, None, None)) + captured = {} + kernel = object() + + def fake_make_wna16_moe_kernel(**kwargs): + captured.update(kwargs) + return kernel + + monkeypatch.setattr(moe_wna16, "make_wna16_moe_kernel", fake_make_wna16_moe_kernel) + + method._setup_kernel(layer) + + assert method.moe_kernel is kernel + assert captured["backend"] == WNA16MoEBackend.HUMMING + assert captured["layer"] is layer + + +def test_moe_wna16_humming_adapter_repacks_uint8_tensors(): + qweight = torch.arange(32, dtype=torch.uint8).reshape(1, 4, 8) + scales = torch.arange(16, dtype=torch.float16).reshape(1, 4, 4) + qzeros = torch.arange(16, dtype=torch.uint8).reshape(1, 8, 2) + + converted = _convert_moe_wna16_humming_tensors( + {"qweight": qweight, "scales": scales, "qzeros": qzeros}, + has_zero_point=True, ) - output = method.apply( - layer, - x=torch.empty(1, 2), - topk_weights=torch.empty(1, 1), - topk_ids=torch.empty(1, 1, dtype=torch.int32), - shared_experts=None, - shared_experts_input=None, + assert torch.equal(converted["weight"], qweight.view(torch.int32)) + assert converted["weight"].shape == (1, 4, 2) + assert torch.equal(converted["weight_scale"], scales) + expected_qzeros = ( + qzeros.transpose(-1, -2) + .contiguous() + .view(torch.int32) + .transpose(-1, -2) + .contiguous() + ) + assert torch.equal(converted["zero_point"], expected_qzeros) + assert converted["zero_point"].shape == (1, 2, 2) + + +def test_moe_wna16_uses_humming_quant_config(monkeypatch): + from vllm.model_executor.layers.quantization.utils import humming_utils + + method = object.__new__(MoeWNA16Method) + method.wna16_backend = WNA16MoEBackend.HUMMING + layer = object() + quant_config = object() + monkeypatch.setattr( + humming_utils, + "get_humming_moe_quant_config", + lambda actual_layer: quant_config if actual_layer is layer else None, ) - assert output.shape == (1, 2) - assert captured_kwargs["activation"] is MoEActivation.GELU_TANH + assert method.get_fused_moe_quant_config(layer) is quant_config diff --git a/vllm/model_executor/layers/fused_moe/experts/triton_moe.py b/vllm/model_executor/layers/fused_moe/experts/triton_moe.py index 3196667b3f7a..8610e6b1a70b 100644 --- a/vllm/model_executor/layers/fused_moe/experts/triton_moe.py +++ b/vllm/model_executor/layers/fused_moe/experts/triton_moe.py @@ -43,7 +43,10 @@ kFp8Static128BlockSym, kFp8StaticChannelSym, kFp8StaticTensorSym, + kInt4Static, + kInt4Static32, kInt8DynamicTokenSym, + kInt8Static, kInt8StaticChannelSym, ) from vllm.platforms import current_platform @@ -542,40 +545,45 @@ def moe_sum(self, input: torch.Tensor, output: torch.Tensor) -> None: class TritonWNA16Experts(TritonExperts): @staticmethod def _supports_current_device() -> bool: - raise NotImplementedError( - "TritonWNA16Experts is not yet used by an Oracle. " - "This method should not be called." - ) + return current_platform.is_cuda_alike() or current_platform.is_xpu() @staticmethod def _supports_no_act_and_mul() -> bool: - raise NotImplementedError( - "TritonWNA16Experts is not yet used by an Oracle. " - "This method should not be called." - ) + return True @staticmethod def _supports_quant_scheme( weight_key: QuantKey | None, activation_key: QuantKey | None, ) -> bool: - raise NotImplementedError( - "TritonWNA16Experts is not yet used by an Oracle. " - "This method should not be called." - ) + SUPPORTED_W = [ + kInt4Static, + kInt8Static, + kInt4Static32, + # other group sizes? + ] + return weight_key in SUPPORTED_W @staticmethod def _supports_activation(activation: MoEActivation) -> bool: - raise NotImplementedError( - "TritonWNA16Experts is not yet used by an Oracle. " - "This method should not be called." - ) + return activation in [ + MoEActivation.SILU, + MoEActivation.GELU, + MoEActivation.GELU_TANH, + MoEActivation.SWIGLUOAI, + MoEActivation.SWIGLUSTEP, + MoEActivation.SILU_NO_MUL, + MoEActivation.GELU_NO_MUL, + MoEActivation.GELU_TANH_NO_MUL, + MoEActivation.RELU2_NO_MUL, + ] @staticmethod def _supports_parallel_config(moe_parallel_config: FusedMoEParallelConfig) -> bool: - raise NotImplementedError( - "TritonWNA16Experts is not yet used by an Oracle. " - "This method should not be called." + # Why? + return not ( + moe_parallel_config.use_fi_nvl_two_sided_kernels + or moe_parallel_config.use_fi_nvl_one_sided_kernels ) def apply( @@ -598,7 +606,9 @@ def apply( ): # Check constraints. if self.quant_config.use_int4_w4a16: - assert hidden_states.size(-1) // 2 == w1.size(2), "Hidden size mismatch" + assert hidden_states.size(-1) // 2 == w1.size(2), ( + f"Hidden size mismatch {hidden_states.size(-1) // 2} == {w1.size(2)}" + ) else: assert hidden_states.size(-1) == w1.size(2), ( f"Hidden size mismatch {hidden_states.size(-1)} != {w1.size(2)}" diff --git a/vllm/model_executor/layers/fused_moe/fused_moe.py b/vllm/model_executor/layers/fused_moe/fused_moe.py index 62cf12ea24e2..0d357fcbf456 100644 --- a/vllm/model_executor/layers/fused_moe/fused_moe.py +++ b/vllm/model_executor/layers/fused_moe/fused_moe.py @@ -31,6 +31,7 @@ ) from vllm.platforms import current_platform from vllm.triton_utils import tl, triton +from vllm.utils.math_utils import next_power_of_2 from vllm.utils.platform_utils import get_device_name_as_file_name from vllm.utils.torch_utils import direct_register_custom_op @@ -1298,7 +1299,11 @@ def get_default_config( bit = 4 if dtype == "int4_w4a16" else 8 use_moe_wna16_cuda = should_moe_wna16_use_cuda(M * topk, block_shape[1], E, bit) if use_moe_wna16_cuda: - config = {"BLOCK_SIZE_M": min(16, M), "SPLIT_K": 1} + config = { + "BLOCK_SIZE_M": min(16, next_power_of_2(M)), + "GROUP_SIZE_M": 1, + "SPLIT_K": 1, + } elif M <= 20: config = {"BLOCK_SIZE_M": 16, "GROUP_SIZE_M": 1, "SPLIT_K": 1} elif M <= 40: diff --git a/vllm/model_executor/layers/fused_moe/oracle/int_wna16.py b/vllm/model_executor/layers/fused_moe/oracle/int_wna16.py index a80013f65baa..9d9d3aa76569 100644 --- a/vllm/model_executor/layers/fused_moe/oracle/int_wna16.py +++ b/vllm/model_executor/layers/fused_moe/oracle/int_wna16.py @@ -24,6 +24,9 @@ MarlinExperts, MarlinExpertsBase, ) +from vllm.model_executor.layers.fused_moe.experts.triton_moe import ( + TritonWNA16Experts, +) from vllm.model_executor.layers.fused_moe.experts.trtllm_mxint4_moe import ( TrtLlmMxint4ExpertsMonolithic, ) @@ -50,6 +53,7 @@ class WNA16MoEBackend(Enum): HUMMING = "HUMMING" CPU = "CPU" FLASHINFER_TRTLLM = "FLASHINFER_TRTLLM" + TRITON = "TRITON" XPU = "XPU" EMULATION = "EMULATION" @@ -76,6 +80,8 @@ def backend_to_kernel_cls( return [BatchedMarlinExperts] elif backend == WNA16MoEBackend.FLASHINFER_TRTLLM: return [TrtLlmMxint4ExpertsMonolithic] + elif backend == WNA16MoEBackend.TRITON: + return [TritonWNA16Experts] elif backend == WNA16MoEBackend.XPU: from vllm.model_executor.layers.fused_moe.experts.xpu_moe import ( XPUExpertsWNA16, @@ -107,19 +113,56 @@ def _get_priority_backends() -> list[WNA16MoEBackend]: if current_platform.is_xpu(): return [WNA16MoEBackend.XPU] - _AVAILABLE_BACKENDS = [ + return [ WNA16MoEBackend.FLASHINFER_TRTLLM, WNA16MoEBackend.MARLIN, WNA16MoEBackend.BATCHED_MARLIN, + WNA16MoEBackend.TRITON, WNA16MoEBackend.HUMMING, WNA16MoEBackend.EMULATION, ] - return _AVAILABLE_BACKENDS + + +def _backend_incompatibility_reason( + backend: WNA16MoEBackend, + quant_config: QuantizationConfig | QuantizationArgs, + may_have_zp: bool, + may_have_bias: bool, +) -> str | None: + if backend == WNA16MoEBackend.FLASHINFER_TRTLLM and (may_have_zp or may_have_bias): + return "zero points and bias are not supported" + + from vllm.model_executor.layers.quantization.auto_awq import AutoAWQConfig + from vllm.model_executor.layers.quantization.auto_gptq import AutoGPTQConfig + from vllm.model_executor.layers.quantization.moe_wna16 import MoeWNA16Config + + if backend == WNA16MoEBackend.TRITON: + if may_have_bias: + return "expert bias is not supported" + if isinstance(quant_config, AutoAWQConfig): + return "the AutoAWQ weight layout is not supported" + if isinstance(quant_config, AutoGPTQConfig) and quant_config.desc_act: + return "GPTQ activation ordering is not supported" + if ( + isinstance(quant_config, QuantizationArgs) + and quant_config.actorder == "group" + ): + return "group activation ordering is not supported" + + if isinstance(quant_config, MoeWNA16Config) and backend in ( + WNA16MoEBackend.MARLIN, + WNA16MoEBackend.BATCHED_MARLIN, + WNA16MoEBackend.EMULATION, + ): + return "the MoeWNA16 checkpoint layout is not supported" + + return None def map_wna16_backend(runner_backend: MoEBackend) -> WNA16MoEBackend: """Map user's MoEBackend to WNA16MoEBackend.""" mapping = { + "triton": WNA16MoEBackend.TRITON, "marlin": WNA16MoEBackend.MARLIN, "humming": WNA16MoEBackend.HUMMING, "flashinfer_trtllm": WNA16MoEBackend.FLASHINFER_TRTLLM, @@ -136,6 +179,9 @@ def map_wna16_backend(runner_backend: MoEBackend) -> WNA16MoEBackend: def select_wna16_moe_backend( config: FusedMoEConfig, weight_key: QuantKey, + quant_config: QuantizationConfig | QuantizationArgs, + may_have_zp: bool, + may_have_bias: bool, ) -> tuple[WNA16MoEBackend, type[mk.FusedMoEExperts]]: """Select the WNA16 MoE backend. @@ -143,6 +189,9 @@ def select_wna16_moe_backend( config: the shared ``FusedMoEConfig`` for this layer. weight_key: The QuantKey describing the weight quantization. Must have int4 or int8 type. + quant_config: Quantization structure and checkpoint format description. + may_have_zp: Whether the integration can provide weight zero points. + may_have_bias: Whether the integration can provide expert bias. Returns: A tuple of (``WNA16MoEBackend``, experts class or ``None``). @@ -189,6 +238,11 @@ def _return_or_raise( runner_backend = config.moe_backend if runner_backend != "auto": requested_backend = map_wna16_backend(runner_backend) + reason = _backend_incompatibility_reason( + requested_backend, quant_config, may_have_zp, may_have_bias + ) + if reason is not None: + raise ValueError(_make_log_unsupported(requested_backend, reason)) return _return_or_raise( requested_backend, config, weight_key, None, activation_format ) @@ -197,6 +251,12 @@ def _return_or_raise( AVAILABLE_BACKENDS = _get_priority_backends() for backend in AVAILABLE_BACKENDS: + reason = _backend_incompatibility_reason( + backend, quant_config, may_have_zp, may_have_bias + ) + if reason is not None: + logger.debug_once(_make_log_unsupported(backend, reason), scope="local") + continue activation_key = None # always BF16 activation for WNA16 MoE for k_cls in backend_to_kernel_cls(backend): supported, reason = k_cls.is_supported_config( @@ -294,6 +354,7 @@ def make_wna16_moe_kernel( allowed_experts: tuple[type[mk.FusedMoEExperts], ...] = ( MarlinExperts, BatchedMarlinExperts, + TritonWNA16Experts, TrtLlmMxint4ExpertsMonolithic, XPUExpertsWNA16, CPUExpertsInt4, @@ -315,6 +376,7 @@ def make_wna16_moe_kernel( assert prepare_finalize is not None logger.info_once("Using %s", prepare_finalize.__class__.__name__, scope="local") + logger.info_once("Using %s", experts_cls.__name__, scope="local") extra_args: dict[str, Any] = {} if backend == WNA16MoEBackend.HUMMING: @@ -329,35 +391,17 @@ def make_wna16_moe_kernel( "is_k_full": is_k_full, } - if experts_cls is XPUExpertsWNA16: - assert ( - prepare_finalize.activation_format == mk.FusedMoEActivationFormat.Standard - ), ( - "XPUExpertsWNA16 only supports the Standard activation format; " - "xpu_fused_moe(is_int4=True) does not implement BatchedExperts." - ) - experts: mk.FusedMoEExperts = XPUExpertsWNA16( - moe_config=moe_config, - quant_config=moe_quant_config, - ) - elif ( - prepare_finalize.activation_format == mk.FusedMoEActivationFormat.BatchedExperts - ): + if prepare_finalize.activation_format == mk.FusedMoEActivationFormat.BatchedExperts: max_num_tokens = prepare_finalize.max_num_tokens_per_rank() assert max_num_tokens is not None - experts = experts_cls( - max_num_tokens=max_num_tokens, - num_dispatchers=prepare_finalize.num_dispatchers(), - moe_config=moe_config, - quant_config=moe_quant_config, - **extra_args, - ) - else: - experts = experts_cls( - moe_config=moe_config, - quant_config=moe_quant_config, - **extra_args, - ) + extra_args["max_num_tokens"] = max_num_tokens + extra_args["num_dispatchers"] = prepare_finalize.num_dispatchers() + + experts = experts_cls( + moe_config=moe_config, + quant_config=moe_quant_config, + **extra_args, + ) return mk.FusedMoEKernel( prepare_finalize, @@ -1034,11 +1078,68 @@ def _humming_wna16_weight_schema( "sym": quant_config.is_sym, } raise TypeError( - "Humming WNA16 MoE requires AutoAWQConfig or AutoGPTQConfig, " + "Humming WNA16 checkpoint schema requires AutoAWQConfig or " + "AutoGPTQConfig, " f"got {type(quant_config).__name__}." ) +def _convert_moe_wna16_humming_tensors( + tensors: dict[str, torch.Tensor], has_zero_point: bool +) -> dict[str, torch.Tensor]: + """Convert MoeWNA16's N-first uint8 packing to Humming's int32 packing.""" + if sys.byteorder != "little": + raise NotImplementedError( + "MoeWNA16 to Humming conversion requires a little-endian host." + ) + + output = { + "weight": tensors["qweight"].contiguous().view(torch.int32), + "weight_scale": tensors["scales"], + } + if has_zero_point: + qzeros = tensors["qzeros"] + output["zero_point"] = ( + qzeros.transpose(-1, -2) + .contiguous() + .view(torch.int32) + .transpose(-1, -2) + .contiguous() + ) + return output + + +class _MoeWNA16HummingWeightSchema: + """Adapter from MoeWNA16's generic packed layout to Humming's layout.""" + + def __init__(self, bits: int, group_size: int, has_zero_point: bool) -> None: + self.bits = bits + self.group_size = group_size + self.has_zero_point = has_zero_point + + def convert_humming( + self, + tensors: dict[str, torch.Tensor], + shape_n_stacks: list[int], + shape_k_stacks: list[int], + param_dtype: torch.dtype, + num_experts: int | None = None, + ) -> tuple[Any, dict[str, torch.Tensor]]: + del shape_n_stacks, shape_k_stacks, num_experts + from vllm.utils.humming import HummingWeightSchema, dtypes + + output = _convert_moe_wna16_humming_tensors( + tensors, has_zero_point=self.has_zero_point + ) + output["weight_scale"] = output["weight_scale"].to(param_dtype) + schema = HummingWeightSchema( + b_dtype=dtypes.DataType.from_str(f"uint{self.bits}"), + weight_scale_group_size=self.group_size, + has_zero_point=self.has_zero_point, + ) + return schema, output + + def _unpack_and_dequant_int4_gptq( w_int32: torch.Tensor, scale: torch.Tensor, @@ -1323,13 +1424,27 @@ def convert_to_wna16_moe_kernel_format( input_dtype: optional activation dtype, usually should be 16 bit. """ if backend == WNA16MoEBackend.HUMMING: + from vllm.model_executor.layers.quantization.moe_wna16 import MoeWNA16Config from vllm.model_executor.layers.quantization.utils.humming_utils import ( convert_to_humming_moe_kernel_format, ) - convert_to_humming_moe_kernel_format( - layer, quant_config=_humming_wna16_weight_schema(quant_config) - ) + if isinstance(quant_config, MoeWNA16Config): + from vllm.utils.humming import HummingInputSchema + + convert_to_humming_moe_kernel_format( + layer, + weight_schema=_MoeWNA16HummingWeightSchema( + bits=quant_config.weight_bits, + group_size=layer.group_size, + has_zero_point=quant_config.has_zp, + ), + input_schema=HummingInputSchema(), + ) + else: + convert_to_humming_moe_kernel_format( + layer, quant_config=_humming_wna16_weight_schema(quant_config) + ) return None if backend in ( @@ -1344,16 +1459,32 @@ def convert_to_wna16_moe_kernel_format( ) if isinstance(quant_config, AutoAWQConfig): - if w13_qzeros is None or w2_qzeros is None: - raise ValueError("AWQ Marlin MoE requires zero-point tensors.") - - weight_bits = quant_config.weight_bits + num_bits = quant_config.weight_bits pack_factor = quant_config.pack_factor group_size = quant_config.group_size + elif isinstance(quant_config, AutoGPTQConfig): + num_bits = quant_config.quant_type.size_bits + pack_factor = quant_config.pack_factor + group_size = quant_config.group_size + actorder = "group" if quant_config.desc_act else None + elif isinstance(quant_config, QuantizationArgs): + num_bits = quant_config.num_bits + pack_factor = 32 // quant_config.num_bits + group_size = quant_config.group_size + actorder = quant_config.actorder + else: + raise TypeError( + "Marlin WNA16 MoE backend requires AutoGPTQConfig, AutoAWQConfig or " + f"QuantizationArgs, got {type(quant_config).__name__}." + ) + + if isinstance(quant_config, AutoAWQConfig): + if w13_qzeros is None or w2_qzeros is None: + raise ValueError("AWQ Marlin MoE requires zero-point tensors.") return _process_awq_weights_marlin( layer, - weight_bits, + num_bits, pack_factor, group_size, input_dtype, @@ -1366,41 +1497,28 @@ def convert_to_wna16_moe_kernel_format( w13_bias, w2_bias, ) - elif isinstance(quant_config, AutoGPTQConfig): - num_bits = quant_config.quant_type.size_bits - pack_factor = quant_config.pack_factor - group_size = quant_config.group_size - actorder = "group" if quant_config.desc_act else None - elif isinstance(quant_config, QuantizationArgs): - num_bits = quant_config.num_bits - pack_factor = 32 // quant_config.num_bits - group_size = quant_config.group_size - actorder = quant_config.actorder else: - raise TypeError( - "Marlin WNA16 MoE backend requires AutoAWQConfig, AutoGPTQConfig or " - f"QuantizationArgs, got {type(quant_config).__name__}." + if w13_g_idx is None or w2_g_idx is None: + raise ValueError("GPTQ Marlin MoE requires g_idx tensors.") + + return _process_weights_marlin( + layer, + input_dtype, + num_bits, + pack_factor, + group_size, + actorder, + w13, + w2, + w13_scale, + w2_scale, + w13_g_idx, + w2_g_idx, + w13_qzeros, + w2_qzeros, + w13_bias, + w2_bias, ) - if w13_g_idx is None or w2_g_idx is None: - raise ValueError("GPTQ Marlin MoE requires g_idx tensors.") - return _process_weights_marlin( - layer, - input_dtype, - num_bits, - pack_factor, - group_size, - actorder, - w13, - w2, - w13_scale, - w2_scale, - w13_g_idx, - w2_g_idx, - w13_qzeros, - w2_qzeros, - w13_bias, - w2_bias, - ) elif backend == WNA16MoEBackend.CPU: return _process_weights_cpu( quant_config, @@ -1482,5 +1600,47 @@ def convert_to_wna16_moe_kernel_format( w13_qzeros, w2_qzeros, ) + elif backend == WNA16MoEBackend.TRITON: + # Two possible input layouts depending on the quantization source: + # + # MoeWNA16 (uint8): (E, N_out, K // bit8_pack) — N-first + # → just view as uint8 (no-op) + # + # AutoGPTQ/compressed-tensors (int32, K-first): + # (E, K // pack32, N_out) + # → transpose to N-first, then view as uint8 to get + # (E, N_out, K // bit8_pack) [int32 = 4 bytes → 4 uint8s] + # Scales: (E, K // gs, N_out) → transpose → (E, N_out, K // gs) + from vllm.model_executor.layers.quantization.auto_gptq import ( + AutoGPTQConfig, + ) + + if isinstance(quant_config, (AutoGPTQConfig, QuantizationArgs)): + # These integrations build in K-first format even when the Triton + # backend is selected. Transpose to N-first first. + w13_uint8 = w13.transpose(1, 2).contiguous().view(torch.uint8) + w2_uint8 = w2.transpose(1, 2).contiguous().view(torch.uint8) + w13_scale = w13_scale.transpose(1, 2).contiguous() + w2_scale = w2_scale.transpose(1, 2).contiguous() + else: + # MoeWNA16 uses N-first uint8 weights and scales. + w13_uint8 = w13.view(torch.uint8) + w2_uint8 = w2.view(torch.uint8) + return ( + w13_uint8, + w2_uint8, + w13_scale, + w2_scale, + None, + None, + None, + None, + w13_qzeros, + w2_qzeros, + None, + None, + w13_bias, + w2_bias, + ) else: raise ValueError(f"Unsupported wna16 MoE backend: {backend.value}") diff --git a/vllm/model_executor/layers/fused_moe/routed_experts.py b/vllm/model_executor/layers/fused_moe/routed_experts.py index 80e4064ea6b6..8bf0123e57fe 100644 --- a/vllm/model_executor/layers/fused_moe/routed_experts.py +++ b/vllm/model_executor/layers/fused_moe/routed_experts.py @@ -114,6 +114,7 @@ def __init__( self.e_score_correction_bias = e_score_correction_bias self.apply_router_weight_on_input = apply_router_weight_on_input # End random parameters + self._loaded_expert_biases: set[str] = set() self.quant_method = self._get_quant_method( self.layer_name, @@ -691,6 +692,25 @@ def weight_loader( expert_data = param.data if full_load else param.data[expert_id] + if "bias" in weight_name: + self._loaded_expert_biases.add(weight_name.rsplit(".", 1)[-1]) + if shard_id == "w2": + expert_data = self._narrow_expert_data_for_padding( + expert_data, + loaded_weight, + hidden_dim=0, + ) + expert_data.copy_(loaded_weight) + else: + self._load_w13( + shard_id=shard_id, + shard_dim=0, + loaded_weight=loaded_weight, + expert_data=expert_data, + tp_rank=self.moe_config.tp_rank, + ) + return True if return_success else None + # Case input scale: input_scale loading is only supported for fp8 if "input_scale" in weight_name: # this is needed for compressed-tensors only diff --git a/vllm/model_executor/layers/quantization/auto_awq.py b/vllm/model_executor/layers/quantization/auto_awq.py index 1e49ec3387e4..a77fef05254e 100644 --- a/vllm/model_executor/layers/quantization/auto_awq.py +++ b/vllm/model_executor/layers/quantization/auto_awq.py @@ -560,6 +560,9 @@ def __init__( self.wna16_moe_backend, self.experts_cls = select_wna16_moe_backend( moe, kInt4Static, + quant_config=self.quant_config, + may_have_zp=self.quant_config.zero_point, + may_have_bias=True, ) def create_weights( diff --git a/vllm/model_executor/layers/quantization/auto_gptq.py b/vllm/model_executor/layers/quantization/auto_gptq.py index b056f5222af8..f94bbb25b3ca 100644 --- a/vllm/model_executor/layers/quantization/auto_gptq.py +++ b/vllm/model_executor/layers/quantization/auto_gptq.py @@ -489,6 +489,9 @@ def __init__( self.wna16_moe_backend, self.experts_cls = select_wna16_moe_backend( moe, weight_key, + quant_config=self.quant_config, + may_have_zp=True, + may_have_bias=True, ) def create_weights( @@ -577,25 +580,25 @@ def create_weights( set_weight_attrs(w2_scales, extra_weight_attrs) # don't shard the w2 scales when running act order set_weight_attrs(w2_scales, {"load_full_w2": self.quant_config.desc_act}) - # up_proj scales + # up_proj zero points w13_qzeros = torch.nn.Parameter( torch.empty( num_experts, scales_size13, 2 * intermediate_size_per_partition // self.quant_config.pack_factor, - dtype=params_dtype, + dtype=torch.int32, ), requires_grad=False, ) layer.register_parameter("w13_qzeros", w13_qzeros) set_weight_attrs(w13_qzeros, extra_weight_attrs) - # down_proj scales + # down_proj zero points w2_qzeros = torch.nn.Parameter( torch.empty( num_experts, scales_size2, hidden_size // self.quant_config.pack_factor, - dtype=params_dtype, + dtype=torch.int32, ), requires_grad=False, ) @@ -644,6 +647,26 @@ def create_weights( layer.register_parameter("w2_g_idx_sort_indices", w2_g_idx_sort_indices) set_weight_attrs(w2_g_idx_sort_indices, extra_weight_attrs) + # Some GPTQ checkpoints contain expert biases even when the model + # architecture does not declare them. Zero initialization keeps + # checkpoints without biases equivalent to the bias-free path. + w13_bias = torch.nn.Parameter( + torch.zeros( + num_experts, + 2 * intermediate_size_per_partition, + dtype=params_dtype, + ), + requires_grad=False, + ) + layer.register_parameter("w13_bias", w13_bias) + set_weight_attrs(w13_bias, extra_weight_attrs) + w2_bias = torch.nn.Parameter( + torch.zeros(num_experts, hidden_size, dtype=params_dtype), + requires_grad=False, + ) + layer.register_parameter("w2_bias", w2_bias) + set_weight_attrs(w2_bias, extra_weight_attrs) + if self.experts_cls is not None and issubclass( self.experts_cls, FusedMoEExpertsModular ): @@ -651,12 +674,31 @@ def create_weights( layer.workspace = marlin_make_workspace_new(device, 4) def process_weights_after_loading(self, layer: RoutedExperts) -> None: + def replace_or_register(name: str, val: torch.Tensor | None): + if val is None: + return + + if hasattr(layer, name): + replace_parameter(layer, name, val) + else: + layer.register_parameter( + name, torch.nn.Parameter(val, requires_grad=False) + ) + is_a_8bit = self.input_dtype is not None and self.input_dtype.itemsize == 1 - if is_a_8bit: - assert self.quant_config.quant_type.size_bits == 8, ( - "W8A8-INT8 is not supported by marlin kernel." - ) + assert not is_a_8bit or self.quant_config.quant_type.size_bits == 8, ( + "W8A8-INT8 is not supported by marlin kernel." + ) + + w13_bias = getattr(layer, "w13_bias", None) + if "w13_bias" not in layer._loaded_expert_biases: + layer.register_parameter("w13_bias", None) + w13_bias = None + w2_bias = getattr(layer, "w2_bias", None) + if "w2_bias" not in layer._loaded_expert_biases: + layer.register_parameter("w2_bias", None) + w2_bias = None converted = convert_to_wna16_moe_kernel_format( backend=self.wna16_moe_backend, @@ -669,8 +711,10 @@ def process_weights_after_loading(self, layer: RoutedExperts) -> None: w2_scale=layer.w2_scales, w13_g_idx=layer.w13_g_idx, w2_g_idx=layer.w2_g_idx, - w13_bias=getattr(layer, "w13_bias", None), - w2_bias=getattr(layer, "w2_bias", None), + w13_bias=w13_bias, + w2_bias=w2_bias, + w13_qzeros=getattr(layer, "w13_qzeros", None), + w2_qzeros=getattr(layer, "w2_qzeros", None), ) if converted is None: @@ -703,42 +747,12 @@ def process_weights_after_loading(self, layer: RoutedExperts) -> None: replace_parameter(layer, "w2_g_idx", w2_g_idx) replace_parameter(layer, "w13_g_idx_sort_indices", w13_g_idx_sort_indices) replace_parameter(layer, "w2_g_idx_sort_indices", w2_g_idx_sort_indices) - if w13_qzeros is not None: - replace_parameter(layer, "w13_qzeros", w13_qzeros) - if w2_qzeros is not None: - replace_parameter(layer, "w2_qzeros", w2_qzeros) - if w13_input_global_scale is not None: - if hasattr(layer, "w13_input_global_scale"): - replace_parameter( - layer, "w13_input_global_scale", w13_input_global_scale - ) - else: - layer.register_parameter( - "w13_input_global_scale", - torch.nn.Parameter(w13_input_global_scale, requires_grad=False), - ) - if w2_input_global_scale is not None: - if hasattr(layer, "w2_input_global_scale"): - replace_parameter(layer, "w2_input_global_scale", w2_input_global_scale) - else: - layer.register_parameter( - "w2_input_global_scale", - torch.nn.Parameter(w2_input_global_scale, requires_grad=False), - ) - if w13_bias is not None: - if hasattr(layer, "w13_bias"): - replace_parameter(layer, "w13_bias", w13_bias) - else: - layer.register_parameter( - "w13_bias", torch.nn.Parameter(w13_bias, requires_grad=False) - ) - if w2_bias is not None: - if hasattr(layer, "w2_bias"): - replace_parameter(layer, "w2_bias", w2_bias) - else: - layer.register_parameter( - "w2_bias", torch.nn.Parameter(w2_bias, requires_grad=False) - ) + replace_or_register("w13_input_global_scale", w13_input_global_scale) + replace_or_register("w2_input_global_scale", w2_input_global_scale) + replace_or_register("w13_bias", w13_bias) + replace_or_register("w2_bias", w2_bias) + replace_or_register("w13_qzeros", w13_qzeros) + replace_or_register("w2_qzeros", w2_qzeros) # The modular kernel reads w13_weight/w2_weight; marlin keeps *_qweight. layer.w13_weight = layer.w13_qweight diff --git a/vllm/model_executor/layers/quantization/compressed_tensors/compressed_tensors_moe/compressed_tensors_moe_wna16_marlin.py b/vllm/model_executor/layers/quantization/compressed_tensors/compressed_tensors_moe/compressed_tensors_moe_wna16_marlin.py index 2b3317d00f3c..516d94fae6e9 100644 --- a/vllm/model_executor/layers/quantization/compressed_tensors/compressed_tensors_moe/compressed_tensors_moe_wna16_marlin.py +++ b/vllm/model_executor/layers/quantization/compressed_tensors/compressed_tensors_moe/compressed_tensors_moe_wna16_marlin.py @@ -92,7 +92,15 @@ def __init__( self.wna16_backend, self.experts_cls = select_wna16_moe_backend( config=self.moe, weight_key=weight_key, + quant_config=self.weight_quant, + may_have_zp=not self.symmetric, + may_have_bias=False, ) + self.is_marlin = self.wna16_backend in [ + WNA16MoEBackend.MARLIN, + WNA16MoEBackend.BATCHED_MARLIN, + ] + self.is_transposed = self.wna16_backend != WNA16MoEBackend.FLASHINFER_TRTLLM def get_weight_shape( self, @@ -118,7 +126,6 @@ def get_weight_shape( "num_groups_w2 must be provided for weight scales/zero_points" ) w13_num_shards = 2 if self.moe.is_act_and_mul else 1 - is_flashinfer = self.wna16_backend == WNA16MoEBackend.FLASHINFER_TRTLLM shape_map = { "w13_weight": { "Flashinfer": ( @@ -177,7 +184,7 @@ def get_weight_shape( ), }, } - backend_key = "Flashinfer" if is_flashinfer else "Marlin" + backend_key = "Marlin" if self.is_transposed else "Flashinfer" return shape_map[weight_name][backend_key] @staticmethod @@ -217,9 +224,8 @@ def create_weights( # Will transpose the loaded weight along the # intermediate and hidden dim sizes. Will # shard for TP along the transposed dims - is_transposed = self.wna16_backend != WNA16MoEBackend.FLASHINFER_TRTLLM extra_weight_attrs.update( - {"is_transposed": is_transposed, "quant_method": self.strategy} + {"is_transposed": self.is_transposed, "quant_method": self.strategy} ) w13_weight = torch.nn.Parameter( @@ -415,7 +421,6 @@ def create_weights( def process_weights_after_loading(self, layer: torch.nn.Module) -> None: # Process weights using the shared oracle infrastructure - is_flashinfer = self.wna16_backend == WNA16MoEBackend.FLASHINFER_TRTLLM converted = convert_to_wna16_moe_kernel_format( backend=self.wna16_backend, layer=layer, @@ -430,6 +435,7 @@ def process_weights_after_loading(self, layer: torch.nn.Module) -> None: w13_qzeros=getattr(layer, "w13_weight_zero_point", None), w2_qzeros=getattr(layer, "w2_weight_zero_point", None), ) + if converted is None: # In-place backends (e.g. Humming) are not wired through this # marlin-only method; fail clearly rather than unpacking None. @@ -437,6 +443,7 @@ def process_weights_after_loading(self, layer: torch.nn.Module) -> None: f"{type(self).__name__} does not support the " f"{self.wna16_backend.value} MoE backend." ) + ( w13_qweight, w2_qweight, @@ -466,7 +473,7 @@ def process_weights_after_loading(self, layer: torch.nn.Module) -> None: replace_parameter(layer, "w2_weight_zero_point", w2_qzeros) # Marlin-specific parameters (not needed for Flashinfer) - if not is_flashinfer: + if self.is_marlin: if w13_g_idx_processed is not None: replace_parameter(layer, "w13_weight_g_idx", w13_g_idx_processed) if w2_g_idx_processed is not None: @@ -510,7 +517,7 @@ def process_weights_after_loading(self, layer: torch.nn.Module) -> None: # Add Marlin-specific arguments marlin_args: dict[str, Any] = {} - if not is_flashinfer: + if self.is_marlin: marlin_args = { "w13_g_idx": layer.w13_weight_g_idx, "w2_g_idx": layer.w2_weight_g_idx, diff --git a/vllm/model_executor/layers/quantization/moe_wna16.py b/vllm/model_executor/layers/quantization/moe_wna16.py index 23e175fd6249..76f2d740bc73 100644 --- a/vllm/model_executor/layers/quantization/moe_wna16.py +++ b/vllm/model_executor/layers/quantization/moe_wna16.py @@ -15,8 +15,13 @@ ) from vllm.model_executor.layers.fused_moe.config import ( FusedMoEQuantConfig, - int4_w4a16_moe_quant_config, - int8_w8a16_moe_quant_config, +) +from vllm.model_executor.layers.fused_moe.oracle.int_wna16 import ( + WNA16MoEBackend, + convert_to_wna16_moe_kernel_format, + make_wna16_moe_kernel, + make_wna16_moe_quant_config, + select_wna16_moe_backend, ) from vllm.model_executor.layers.fused_moe.unquantized_fused_moe_method import ( UnquantizedFusedMoEMethod, @@ -27,7 +32,15 @@ QuantizationConfig, QuantizeMethodBase, ) -from vllm.model_executor.utils import set_weight_attrs +from vllm.model_executor.layers.quantization.utils.quant_utils import ( + INT4_DTYPE, + INT8_DTYPE, + QuantKey, + kInt4Static32GroupScale, + kInt4StaticGroupScale, + kInt8StaticGroupScale, +) +from vllm.model_executor.utils import replace_parameter, set_weight_attrs from vllm.platforms import current_platform @@ -199,6 +212,34 @@ def __init__(self, quant_config: MoeWNA16Config, moe: "FusedMoEConfig") -> None: super().__init__(moe) self.quant_config = quant_config + num_bits = self.quant_config.weight_bits + group_size = self.quant_config.group_size + + if num_bits == 4: + quant_type = INT4_DTYPE + if group_size == 32: + scale = kInt4Static32GroupScale + else: + scale = kInt4StaticGroupScale + elif num_bits == 8: + assert group_size == -1 + quant_type = INT8_DTYPE + scale = kInt8StaticGroupScale + else: + raise ValueError("MoeWNA16Method only supports int4 and int8 now.") + + weight_key = QuantKey(quant_type, scale) + + # Select WNA16 MoE backend via oracle. + # handle ZP? + self.wna16_backend, self.experts_cls = select_wna16_moe_backend( + config=self.moe, + weight_key=weight_key, + quant_config=self.quant_config, + may_have_zp=self.quant_config.has_zp, + may_have_bias=False, + ) + def create_weights( self, layer: RoutedExperts, @@ -322,23 +363,104 @@ def create_weights( def get_fused_moe_quant_config( self, layer: RoutedExperts ) -> FusedMoEQuantConfig | None: - weight_bits = self.quant_config.weight_bits - has_zp = self.quant_config.has_zp - assert weight_bits == 4 or weight_bits == 8 - config_builder = ( - int4_w4a16_moe_quant_config - if weight_bits == 4 - else int8_w8a16_moe_quant_config - ) + if self.wna16_backend == WNA16MoEBackend.HUMMING: + from vllm.model_executor.layers.quantization.utils.humming_utils import ( + get_humming_moe_quant_config, + ) + + return get_humming_moe_quant_config(layer) - return config_builder( + has_zp = self.quant_config.has_zp + return make_wna16_moe_quant_config( w1_scale=layer.w13_scales, w2_scale=layer.w2_scales, w1_zp=layer.w13_qzeros if has_zp else None, w2_zp=layer.w2_qzeros if has_zp else None, - block_shape=[0, layer.group_size], + group_size=layer.group_size, + num_bits=self.quant_config.weight_bits, ) + def _setup_kernel(self, layer: RoutedExperts): + assert self.experts_cls is not None + self.moe_quant_config = self.get_fused_moe_quant_config(layer) + assert self.moe_quant_config is not None + self.moe_kernel = make_wna16_moe_kernel( + moe_quant_config=self.moe_quant_config, + moe_config=self.moe, + experts_cls=self.experts_cls, + backend=self.wna16_backend, + layer=layer, + routing_tables=layer._expert_routing_tables(), + ) + + def process_weights_after_loading(self, layer: RoutedExperts) -> None: + has_zp = self.quant_config.has_zp + converted = convert_to_wna16_moe_kernel_format( + backend=self.wna16_backend, + layer=layer, + quant_config=self.quant_config, + input_dtype=None, + w13=layer.w13_qweight, + w2=layer.w2_qweight, + w13_scale=layer.w13_scales, + w2_scale=layer.w2_scales, + w13_g_idx=None, + w2_g_idx=None, + w13_qzeros=layer.w13_qzeros if has_zp else None, + w2_qzeros=layer.w2_qzeros if has_zp else None, + ) + + if converted is None: + # Backend rewrote the layer's params in place (e.g. Humming). + self._setup_kernel(layer) + return + + ( + w13_qweight, + w2_qweight, + w13_scales, + w2_scales, + _, + _, + _, + _, + w13_qzeros, + w2_qzeros, + w13_input_global_scale, + w2_input_global_scale, + _, # w13_bias + _, # w2_bias + ) = converted + + # Replace common parameters + replace_parameter(layer, "w13_qweight", w13_qweight) + replace_parameter(layer, "w2_qweight", w2_qweight) + replace_parameter(layer, "w13_scales", w13_scales) + replace_parameter(layer, "w2_scales", w2_scales) + layer.w13_weight = layer.w13_qweight + layer.w2_weight = layer.w2_qweight + + if has_zp: + assert w13_qzeros is not None and w2_qzeros is not None + replace_parameter(layer, "w13_qzeros", w13_qzeros) + replace_parameter(layer, "w2_qzeros", w2_qzeros) + + # Marlin-specific parameters (not needed for Flashinfer) + if self.wna16_backend != WNA16MoEBackend.FLASHINFER_TRTLLM: + # Register input global scales if present + if w13_input_global_scale is not None: + layer.register_parameter( + "w13_input_global_scale", + torch.nn.Parameter(w13_input_global_scale, requires_grad=False), + ) + if w2_input_global_scale is not None: + layer.register_parameter( + "w2_input_global_scale", + torch.nn.Parameter(w2_input_global_scale, requires_grad=False), + ) + + self._setup_kernel(layer) + def apply( self, layer: RoutedExperts, @@ -348,19 +470,44 @@ def apply( shared_experts: SharedExperts | None, shared_experts_input: torch.Tensor | None, ) -> torch.Tensor: - from vllm.model_executor.layers.fused_moe import fused_experts - - return fused_experts( + assert not self.is_monolithic + assert self.moe_kernel is not None + return self.moe_kernel.apply( x, - layer.w13_qweight, - layer.w2_qweight, + layer.w13_weight, + layer.w2_weight, topk_weights=topk_weights, topk_ids=topk_ids, activation=layer.activation, + global_num_experts=layer.global_num_experts, + expert_map=layer.expert_map, apply_router_weight_on_input=layer.apply_router_weight_on_input, + shared_experts=shared_experts, + shared_experts_input=shared_experts_input, + ) + + def apply_monolithic( + self, + layer: RoutedExperts, + x: torch.Tensor, + router_logits: torch.Tensor, + input_ids: torch.Tensor | None = None, + ) -> torch.Tensor: + assert self.is_monolithic + assert self.moe_kernel is not None + return self.moe_kernel.apply_monolithic( + x, + layer.w13_weight, + layer.w2_weight, + router_logits, + activation=layer.activation, global_num_experts=layer.global_num_experts, expert_map=layer.expert_map, - quant_config=self.moe_quant_config, + apply_router_weight_on_input=layer.apply_router_weight_on_input, + num_expert_group=layer.num_expert_group, + topk_group=layer.topk_group, + e_score_correction_bias=layer.e_score_correction_bias, + routed_scaling_factor=layer.routed_scaling_factor, ) @staticmethod diff --git a/vllm/model_executor/layers/quantization/utils/gptq_utils.py b/vllm/model_executor/layers/quantization/utils/gptq_utils.py index 691d80b0b747..c73c42f4ff02 100644 --- a/vllm/model_executor/layers/quantization/utils/gptq_utils.py +++ b/vllm/model_executor/layers/quantization/utils/gptq_utils.py @@ -3,7 +3,7 @@ from collections.abc import Mapping from copy import deepcopy from types import MappingProxyType -from typing import TYPE_CHECKING +from typing import TYPE_CHECKING, Any import regex as re import torch @@ -69,6 +69,20 @@ def get_dynamic_override( return default_value +def flatten_list(lst: list[Any]) -> list[Any]: + output = [] + + def _flatten(lst: list[Any]): + for i in lst: + if isinstance(i, list): + _flatten(i) + else: + output.append(i) + + _flatten(lst) + return output + + def is_layer_gptq_quantized( prefix: str, quantized_layers: list[str], @@ -83,6 +97,8 @@ def is_layer_gptq_quantized( proj_name = prefix.split(".")[-1] + quantized_layers = flatten_list(quantized_layers) + # Fused layers like gate_up_proj or qkv_proj will not be fused # in the safetensors checkpoint. So, we convert the name # from the fused version to unfused + check to make sure that