@tduncan / Nvidia-Optimization-Stack-VLLM Mirror · read-only / commits / f51bbc694d
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, 19:30 UTC

[MoE Refactor] W4a8 int8 oracle (#42789)

Signed-off-by: Bill Nell <bnell@redhat.com> Co-authored-by: Robert Shaw <114415538+robertgshaw2-redhat@users.noreply.github.com>

bnellnm committed May 26, 2026 at 11:15 UTC f51bbc694d9704daeefc0549229858b657bce20c
5 files changed +1016 -122
tests/kernels/moe/test_cpu_int4_moe.py new
+296
@@ -0,0 +1,296 @@
1 +# SPDX-License-Identifier: Apache-2.0
2 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 +"""Tests for CPU INT4 W4A8 dynamic quantized fused MoE kernel (CPUExpertsInt4)."""
4 +
5 +import sys
6 +
7 +import pytest
8 +import torch
9 +import torch.nn.functional as F
10 +
11 +from vllm.platforms import current_platform
12 +from vllm.utils.torch_utils import set_random_seed
13 +
14 +if not current_platform.is_cpu():
15 + pytest.skip("skipping CPU-only tests", allow_module_level=True)
16 +
17 +# Check if the dynamic_4bit_int_moe op is available
18 +if not hasattr(torch.ops._C, "dynamic_4bit_int_moe"):
19 + pytest.skip("dynamic_4bit_int_moe op not available", allow_module_level=True)
20 +
21 +# Check if KleidiAI ops are available
22 +if not hasattr(torch.ops.aten, "_dyn_quant_pack_4bit_weight"):
23 + pytest.skip("KleidiAI 4-bit ops not available", allow_module_level=True)
24 +
25 +
26 +# Tolerance for INT4 W4A8
27 +INT4_W4A8_ATOL = 2e-2
28 +INT4_W4A8_RTOL = 2e-2
29 +
30 +
31 +def _silu_and_mul(x: torch.Tensor) -> torch.Tensor:
32 + """SwiGLU activation: SiLU(gate) * up."""
33 + d = x.shape[-1] // 2
34 + return F.silu(x[..., :d]) * x[..., d:]
35 +
36 +
37 +def _pack_int4_weight_to_kleidi(
38 + int4_as_int8: torch.Tensor,
39 + scales: torch.Tensor,
40 + bias: torch.Tensor | None,
41 + group_size: int,
42 + in_features: int,
43 + out_features: int,
44 +) -> torch.Tensor:
45 + """Pack INT4 weights (stored as int8 in [-8,7]) to KleidiAI format.
46 +
47 + Args:
48 + int4_as_int8: [out, in] int8 tensor with values in [-8, 7]
49 + scales: [out, in//group_size] or [out, 1] for channel-wise
50 + bias: [out] optional bias
51 + group_size: Quantization group size (-1 for channel-wise)
52 + in_features: Input dimension
53 + out_features: Output dimension
54 +
55 + Returns:
56 + Packed weight tensor in KleidiAI format
57 + """
58 + # Shift to unsigned nibble [0, 15]
59 + tmp = int4_as_int8.add(8)
60 + # Pack pairs along input dimension
61 + uint8_nibbles = ((tmp[:, 1::2] << 4) | tmp[:, ::2]).to(torch.uint8)
62 +
63 + # Determine scale dtype based on group_size
64 + scale_dtype = torch.float32 if group_size == -1 else torch.bfloat16
65 + scales_typed = scales.to(scale_dtype)
66 + bias_typed = None if bias is None else bias.to(torch.float32)
67 +
68 + # Pack using KleidiAI op
69 + actual_group_size = in_features if group_size == -1 else group_size
70 + return torch.ops.aten._dyn_quant_pack_4bit_weight(
71 + uint8_nibbles,
72 + scales_typed,
73 + bias_typed,
74 + actual_group_size,
75 + in_features,
76 + out_features,
77 + )
78 +
79 +
80 +def _make_int4_moe_weights(
81 + E: int,
82 + N: int,
83 + K: int,
84 + group_size: int,
85 + has_bias: bool = False,
86 +) -> tuple[
87 + torch.Tensor,
88 + torch.Tensor,
89 + torch.Tensor,
90 + torch.Tensor,
91 + torch.Tensor | None,
92 + torch.Tensor | None,
93 +]:
94 + """Generate random INT4 MoE weights with random scales.
95 +
96 + Args:
97 + E: Number of experts
98 + N: Intermediate size
99 + K: Hidden size
100 + group_size: Quantization group size (-1 for channel-wise)
101 + has_bias: Whether to include bias
102 +
103 + Returns:
104 + (w13_packed, w2_packed, w13_ref, w2_ref, w13_bias, w2_bias)
105 + where *_ref are the dequantized float reference weights
106 + """
107 + # Generate INT4 weights as int8 values in [-8, 7]
108 + w13_int4 = torch.randint(-8, 8, (E, 2 * N, K), dtype=torch.int8)
109 + w2_int4 = torch.randint(-8, 8, (E, K, N), dtype=torch.int8)
110 +
111 + # Determine number of scale columns
112 + def _n_scale_cols(in_features: int) -> int:
113 + return 1 if group_size == -1 else (in_features // group_size)
114 +
115 + # Generate random scales
116 + scale_dtype = torch.float32 if group_size == -1 else torch.bfloat16
117 + w13_scales = torch.rand(E, 2 * N, _n_scale_cols(K), dtype=scale_dtype) * 0.01
118 + w2_scales = torch.rand(E, K, _n_scale_cols(N), dtype=scale_dtype) * 0.01
119 +
120 + # Generate biases if needed
121 + w13_bias = None
122 + w2_bias = None
123 + if has_bias:
124 + w13_bias = torch.randn(E, 2 * N, dtype=torch.float32) * 0.01
125 + w2_bias = torch.randn(E, K, dtype=torch.float32) * 0.01
126 +
127 + # Pack weights for each expert
128 + w13_packed_list = []
129 + w2_packed_list = []
130 +
131 + for e in range(E):
132 + w13_packed_list.append(
133 + _pack_int4_weight_to_kleidi(
134 + w13_int4[e],
135 + w13_scales[e],
136 + w13_bias[e] if (has_bias and w13_bias is not None) else None,
137 + group_size,
138 + K,
139 + 2 * N,
140 + )
141 + )
142 + w2_packed_list.append(
143 + _pack_int4_weight_to_kleidi(
144 + w2_int4[e],
145 + w2_scales[e],
146 + w2_bias[e] if (has_bias and w2_bias is not None) else None,
147 + group_size,
148 + N,
149 + K,
150 + )
151 + )
152 +
153 + w13_packed = torch.stack(w13_packed_list, dim=0)
154 + w2_packed = torch.stack(w2_packed_list, dim=0)
155 +
156 + # Create reference dequantized weights
157 + w13_ref = torch.zeros(E, 2 * N, K, dtype=torch.float32)
158 + w2_ref = torch.zeros(E, K, N, dtype=torch.float32)
159 +
160 + for e in range(E):
161 + # Dequantize w13
162 + for i in range(2 * N):
163 + for j in range(K):
164 + group_idx = 0 if group_size == -1 else (j // group_size)
165 + w13_ref[e, i, j] = (
166 + w13_int4[e, i, j].float() * w13_scales[e, i, group_idx].float()
167 + )
168 + if has_bias and w13_bias is not None:
169 + w13_ref[e, i, j] += w13_bias[e, i].float()
170 +
171 + # Dequantize w2
172 + for i in range(K):
173 + for j in range(N):
174 + group_idx = 0 if group_size == -1 else (j // group_size)
175 + w2_ref[e, i, j] = (
176 + w2_int4[e, i, j].float() * w2_scales[e, i, group_idx].float()
177 + )
178 + if has_bias and w2_bias is not None:
179 + w2_ref[e, i, j] += w2_bias[e, i].float()
180 +
181 + return w13_packed, w2_packed, w13_ref, w2_ref, w13_bias, w2_bias
182 +
183 +
184 +def ref_int4_moe(
185 + a: torch.Tensor,
186 + w13_ref: torch.Tensor,
187 + w2_ref: torch.Tensor,
188 + topk_weight: torch.Tensor,
189 + topk_ids: torch.Tensor,
190 +) -> torch.Tensor:
191 + """Reference INT4 W4A8 fused MoE using dequantized weights.
192 +
193 + Steps:
194 + 1. Use dequantized float weights
195 + 2. For each expert: matmul → SwiGLU → matmul
196 + 3. Weighted sum across top-k experts
197 + """
198 + B, D = a.shape
199 + topk = topk_ids.size(1)
200 +
201 + a_exp = a.view(B, -1, D).repeat(1, topk, 1).reshape(-1, D).float()
202 + out = torch.zeros(B * topk, w2_ref.shape[1], dtype=torch.float32)
203 +
204 + topk_weight_flat = topk_weight.view(-1)
205 + topk_ids_flat = topk_ids.view(-1)
206 +
207 + for i in range(w13_ref.shape[0]):
208 + mask = topk_ids_flat == i
209 + if mask.sum():
210 + # w13: [2N, K], input: [B, K] -> output: [B, 2N]
211 + gate_up = torch.matmul(a_exp[mask], w13_ref[i].transpose(0, 1))
212 + # SwiGLU activation
213 + hidden = _silu_and_mul(gate_up)
214 + # w2: [K, N], hidden: [B, N] -> output: [B, K]
215 + out[mask] = torch.matmul(hidden, w2_ref[i].transpose(0, 1))
216 +
217 + return (
218 + (out.view(B, -1, w2_ref.shape[1]) * topk_weight_flat.view(B, -1, 1))
219 + .sum(dim=1)
220 + .to(a.dtype)
221 + )
222 +
223 +
224 +NUM_TOKENS = [1, 2, 64, 128]
225 +# (intermediate_size N, hidden_size K, num_experts E, topk, group_size)
226 +MoE_CONFIGS = [
227 + (256, 512, 8, 2, 128),
228 + (256, 512, 8, 2, 64),
229 + (256, 512, 8, 2, -1), # channel-wise
230 + (512, 256, 8, 4, 128),
231 + (512, 512, 8, 2, 128),
232 + (768, 2048, 8, 2, 128),
233 + (768, 2048, 16, 4, 64),
234 +]
235 +SEEDS = [0, 42]
236 +
237 +
238 +@pytest.mark.parametrize("M", NUM_TOKENS)
239 +@pytest.mark.parametrize("N,K,E,topk,group_size", MoE_CONFIGS)
240 +@pytest.mark.parametrize("seed", SEEDS)
241 +def test_cpu_int4_moe_kernel(M, N, K, E, topk, group_size, seed):
242 + """Test dynamic_4bit_int_moe kernel against dequantized torch reference."""
243 + set_random_seed(seed)
244 +
245 + # Generate input activations
246 + a = torch.randn(M, K, dtype=torch.bfloat16) / (K**0.5)
247 +
248 + # Generate INT4 weights
249 + w13_packed, w2_packed, w13_ref, w2_ref, w13_bias, w2_bias = _make_int4_moe_weights(
250 + E, N, K, group_size, has_bias=False
251 + )
252 +
253 + # Generate router logits and topk
254 + score = torch.randn(M, E, dtype=torch.bfloat16)
255 + score = torch.softmax(score, dim=-1, dtype=torch.float32)
256 + topk_weight, topk_ids = torch.topk(score, topk)
257 + topk_ids = topk_ids.to(torch.long)
258 +
259 + # Reference output using dequantized weights
260 + ref_out = ref_int4_moe(
261 + a,
262 + w13_ref,
263 + w2_ref,
264 + topk_weight,
265 + topk_ids,
266 + )
267 +
268 + # Test dynamic_4bit_int_moe kernel
269 + # Activation kind: 1 = SwiGLU_Ug (SiLU(u)*g) for OAI-style
270 + activation_kind = 1
271 + apply_router_weight_on_input = False
272 +
273 + out = torch.ops._C.dynamic_4bit_int_moe(
274 + a,
275 + topk_ids,
276 + topk_weight,
277 + w13_packed,
278 + w2_packed,
279 + K, # H (hidden_size / w2_out_features)
280 + N, # I (intermediate_size / w2_in_features)
281 + 2 * N, # I2 (2*intermediate_size / w13_out_features)
282 + group_size,
283 + apply_router_weight_on_input,
284 + activation_kind,
285 + )
286 +
287 + torch.testing.assert_close(
288 + ref_out.bfloat16(),
289 + out,
290 + atol=INT4_W4A8_ATOL,
291 + rtol=INT4_W4A8_RTOL,
292 + )
293 +
294 +
295 +if __name__ == "__main__":
296 + sys.exit(pytest.main([__file__, "-v"]))
vllm/model_executor/layers/fused_moe/experts/cpu_int4_moe.py new
+229
@@ -0,0 +1,229 @@
1 +# SPDX-License-Identifier: Apache-2.0
2 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 +"""CPU INT4 W4A8 dynamic quantized fused MoE experts."""
4 +
5 +import torch
6 +
7 +import vllm.model_executor.layers.fused_moe.modular_kernel as mk
8 +from vllm.model_executor.layers.fused_moe.activation import MoEActivation
9 +from vllm.model_executor.layers.fused_moe.config import (
10 + FusedMoEConfig,
11 + FusedMoEParallelConfig,
12 + FusedMoEQuantConfig,
13 + RoutingMethodType,
14 +)
15 +from vllm.model_executor.layers.quantization.utils.quant_utils import (
16 + QuantKey,
17 + kInt4W4A8StaticGroup32Sym,
18 + kInt4W4A8StaticGroup64Sym,
19 + kInt4W4A8StaticGroup128Sym,
20 + kInt4W4A8StaticGroupSym,
21 +)
22 +from vllm.platforms import current_platform
23 +
24 +
25 +class CPUExpertsInt4(mk.FusedMoEExpertsMonolithic):
26 + """CPU INT4 W4A8 dynamic quantized monolithic MoE experts.
27 +
28 + Uses the dynamic_4bit_int_moe kernel for efficient 4-bit weight,
29 + 8-bit activation MoE inference on CPU.
30 + """
31 +
32 + def __init__(
33 + self,
34 + moe_config: FusedMoEConfig,
35 + quant_config: FusedMoEQuantConfig,
36 + ):
37 + super().__init__(
38 + moe_config,
39 + quant_config,
40 + )
41 +
42 + @property
43 + def expects_unquantized_inputs(self) -> bool:
44 + """Expects unquantized inputs (quantization happens in kernel)."""
45 + return True
46 +
47 + @staticmethod
48 + def activation_format() -> mk.FusedMoEActivationFormat:
49 + return mk.FusedMoEActivationFormat.Standard
50 +
51 + @staticmethod
52 + def _supports_current_device() -> bool:
53 + return current_platform.is_cpu()
54 +
55 + @staticmethod
56 + def _supports_no_act_and_mul() -> bool:
57 + """Does not support no_act_and_mul (requires SwiGLU or SiLU)."""
58 + return False
59 +
60 + @staticmethod
61 + def _supports_activation(activation: MoEActivation) -> bool:
62 + """Supports SiLU and SwiGLU variants."""
63 + return activation in (
64 + MoEActivation.SILU,
65 + MoEActivation.SWIGLUOAI,
66 + MoEActivation.SWIGLUSTEP,
67 + )
68 +
69 + @staticmethod
70 + def _supports_parallel_config(
71 + moe_parallel_config: FusedMoEParallelConfig,
72 + ) -> bool:
73 + """Currently does not support expert parallelism."""
74 + # Based on compressed_tensors implementation check
75 + return moe_parallel_config.ep_size == 1
76 +
77 + @staticmethod
78 + def _supports_quant_scheme(
79 + weight_key: QuantKey | None,
80 + activation_key: QuantKey | None,
81 + ) -> bool:
82 + """Supports INT4 weights with INT8 dynamic activations.
83 +
84 + This is W4A8 with:
85 + - Weights: 4-bit integer (stored as int8, packed to uint8 nibbles)
86 + Can be channel-wise or group-wise quantization
87 + - Activations: dynamic per-token 8-bit integer quantization
88 + """
89 + # group size must be multiple of 32
90 + SUPPORTED_W_A = [
91 + (kInt4W4A8StaticGroup128Sym, None),
92 + (kInt4W4A8StaticGroup64Sym, None),
93 + (kInt4W4A8StaticGroup32Sym, None),
94 + (kInt4W4A8StaticGroupSym, None),
95 + ]
96 + return (weight_key, activation_key) in SUPPORTED_W_A
97 +
98 + @staticmethod
99 + def _supports_routing_method(
100 + routing_method: RoutingMethodType,
101 + weight_key: QuantKey | None,
102 + activation_key: QuantKey | None,
103 + ) -> bool:
104 + """Supports standard routing methods."""
105 + return routing_method in [
106 + RoutingMethodType.Default,
107 + RoutingMethodType.Renormalize,
108 + RoutingMethodType.RenormalizeNaive,
109 + ]
110 +
111 + @staticmethod
112 + def _supports_router_logits_dtype(
113 + router_logits_dtype: torch.dtype | None,
114 + routing_method: RoutingMethodType,
115 + ) -> bool:
116 + """Accepts any router logits dtype."""
117 + return True
118 +
119 + def supports_expert_map(self) -> bool:
120 + """Expert parallelism not yet supported."""
121 + return False
122 +
123 + def _activation_kind(self, activation: MoEActivation) -> int:
124 + """Convert MoEActivation to kernel activation kind integer.
125 +
126 + Returns:
127 + 0 = SwiGLU_Gu (SiLU(g)*u)
128 + 1 = SwiGLU_Ug (SiLU(u)*g)
129 + 2 = SiLU
130 + """
131 + if activation == MoEActivation.SWIGLUSTEP:
132 + return 0
133 + if activation == MoEActivation.SWIGLUOAI:
134 + return 1
135 + if activation == MoEActivation.SILU:
136 + return 2
137 + raise ValueError(f"Unsupported activation '{activation}'")
138 +
139 + def apply(
140 + self,
141 + hidden_states: torch.Tensor,
142 + w1: torch.Tensor, # w13_weight_packed
143 + w2: torch.Tensor, # w2_weight_packed
144 + router_logits: torch.Tensor,
145 + activation: MoEActivation,
146 + global_num_experts: int,
147 + expert_map: torch.Tensor | None,
148 + a1q_scale: torch.Tensor | None,
149 + apply_router_weight_on_input: bool,
150 + # grouped topk + fused topk bias parameters
151 + num_expert_group: int | None = None,
152 + e_score_correction_bias: torch.Tensor | None = None,
153 + routed_scaling_factor: float | None = None,
154 + topk_group: int | None = None,
155 + ) -> torch.Tensor:
156 + """Apply the monolithic 4-bit INT MoE forward pass.
157 +
158 + Args:
159 + hidden_states: Input tensor [num_tokens, hidden_size]
160 + w1: Packed w13 weights (w1+w3 gated weights)
161 + w2: Packed w2 weights (down projection)
162 + router_logits: Router output logits [num_tokens, num_experts]
163 + activation: Activation function type
164 + global_num_experts: Total number of experts
165 + expert_map: Expert mapping for EP (not supported)
166 + a1q_scale: Activation quantization scale (not used, dynamic)
167 + apply_router_weight_on_input: Whether to apply routing on input
168 + num_expert_group: For grouped topk
169 + e_score_correction_bias: Bias for expert scores
170 + routed_scaling_factor: Scaling factor for routing
171 + topk_group: Group size for topk
172 +
173 + Returns:
174 + Output tensor after MoE computation
175 + """
176 + from vllm.model_executor.layers.fused_moe.cpu_fused_moe import (
177 + select_experts,
178 + )
179 +
180 + renormalize = self.moe_config.routing_method in (
181 + RoutingMethodType.Renormalize,
182 + RoutingMethodType.RenormalizeNaive,
183 + )
184 +
185 + # TODO(bnell): this could be factored into a CPURouter class and
186 + # turn this into a modular kernel
187 + # Perform topk selection
188 + topk_weights, topk_ids = select_experts(
189 + hidden_states=hidden_states,
190 + router_logits=router_logits,
191 + use_grouped_topk=num_expert_group is not None,
192 + top_k=self.moe_config.experts_per_token,
193 + renormalize=renormalize,
194 + topk_group=topk_group,
195 + num_expert_group=num_expert_group,
196 + scoring_func="softmax",
197 + routed_scaling_factor=(
198 + routed_scaling_factor if routed_scaling_factor is not None else 1.0
199 + ),
200 + e_score_correction_bias=e_score_correction_bias,
201 + )
202 +
203 + # Extract dimensions from weight tensors
204 + # w1 is w13_packed: [num_experts, packed_data...]
205 + # w2 is w2_packed: [num_experts, packed_data...]
206 + # These dimensions should be available from the layer
207 + # For now, we'll extract from moe_config
208 + K = self.moe_config.hidden_dim
209 + N = self.moe_config.intermediate_size_per_partition
210 + assert self.quant_config.block_shape is not None
211 + if self.quant_config.is_per_act_token:
212 + group_size = -1
213 + else:
214 + group_size = self.quant_config.block_shape[1]
215 +
216 + # Call the dynamic 4-bit int MoE kernel
217 + return torch.ops._C.dynamic_4bit_int_moe(
218 + hidden_states,
219 + topk_ids.to(torch.long),
220 + topk_weights,
221 + w1, # w13_weight_packed
222 + w2, # w2_weight_packed
223 + K, # hidden_size (w2_out_features)
224 + N, # intermediate_size (w2_in_features)
225 + N * 2, # 2*intermediate_size (w13_out_features)
226 + group_size,
227 + apply_router_weight_on_input,
228 + self._activation_kind(activation),
229 + )
vllm/model_executor/layers/fused_moe/oracle/w4a8_int8.py new
+361
@@ -0,0 +1,361 @@
1 +# SPDX-License-Identifier: Apache-2.0
2 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 +
4 +from enum import Enum
5 +
6 +import torch
7 +
8 +import vllm.model_executor.layers.fused_moe.modular_kernel as mk
9 +from vllm.config.kernel import MoEBackend
10 +from vllm.logger import init_logger
11 +from vllm.model_executor.layers.fused_moe.all2all_utils import (
12 + maybe_make_prepare_finalize,
13 +)
14 +from vllm.model_executor.layers.fused_moe.config import (
15 + FusedMoEConfig,
16 + FusedMoEQuantConfig,
17 + FusedMoEQuantDesc,
18 +)
19 +from vllm.model_executor.layers.quantization.utils.quant_utils import (
20 + GroupShape,
21 + QuantKey,
22 +)
23 +from vllm.platforms import current_platform
24 +
25 +logger = init_logger(__name__)
26 +
27 +
28 +class W4A8Int8MoeBackend(Enum):
29 + CPU_INT4 = "CPU_INT4"
30 +
31 +
32 +def _get_priority_backends(
33 + moe_config: FusedMoEConfig,
34 +) -> list[W4A8Int8MoeBackend]:
35 + """
36 + Get available backends in priority order based on platform and config.
37 +
38 + Currently only CPU INT4 backend is available for W4A8 INT8 MoE.
39 + """
40 + if current_platform.is_cpu():
41 + return [W4A8Int8MoeBackend.CPU_INT4]
42 + return []
43 +
44 +
45 +def backend_to_kernel_cls(
46 + backend: W4A8Int8MoeBackend,
47 +) -> list[type[mk.FusedMoEExperts]]:
48 + """Map W4A8Int8MoeBackend to kernel class."""
49 + if backend == W4A8Int8MoeBackend.CPU_INT4:
50 + from vllm.model_executor.layers.fused_moe.experts.cpu_int4_moe import (
51 + CPUExpertsInt4,
52 + )
53 +
54 + return [CPUExpertsInt4]
55 +
56 + else:
57 + raise ValueError(f"Unknown W4A8 Int8 MoE backend: {backend.value}")
58 +
59 +
60 +def map_w4a8_int8_backend(runner_backend: MoEBackend) -> W4A8Int8MoeBackend:
61 + """Map user's MoEBackend to W4A8Int8MoeBackend."""
62 + mapping = {
63 + "cpu": W4A8Int8MoeBackend.CPU_INT4,
64 + }
65 + if backend := mapping.get(runner_backend):
66 + return backend
67 + raise ValueError(
68 + f"moe_backend='{runner_backend}' is not supported for W4A8 Int8 MoE. "
69 + f"Expected one of {list(mapping.keys())}."
70 + )
71 +
72 +
73 +def select_w4a8_int8_moe_backend(
74 + config: FusedMoEConfig,
75 + weight_key: QuantKey | None,
76 + activation_key: QuantKey | None,
77 +) -> tuple[W4A8Int8MoeBackend, type[mk.FusedMoEExperts]]:
78 + """
79 + Select the primary W4A8 Int8 MoE backend.
80 +
81 + Args:
82 + config: MoE configuration
83 + weight_key: Weight quantization key (should be one of kInt4W4A8Static*)
84 + activation_key: Activation quantization key (currently unused for W4A8)
85 +
86 + Returns:
87 + Tuple of (backend, kernel_class)
88 + """
89 +
90 + AVAILABLE_BACKENDS = _get_priority_backends(config)
91 +
92 + if not AVAILABLE_BACKENDS:
93 + raise NotImplementedError("W4A8 Int8 MoE is only supported on CPU platforms")
94 +
95 + activation_format = (
96 + mk.FusedMoEActivationFormat.BatchedExperts
97 + if config.moe_parallel_config.use_batched_activation_format
98 + else mk.FusedMoEActivationFormat.Standard
99 + )
100 +
101 + def _make_log_backend(backend: W4A8Int8MoeBackend) -> str:
102 + available_backend_strs = [b.value for b in AVAILABLE_BACKENDS]
103 + return (
104 + f"Using {backend.value} W4A8 Int8 MoE backend out "
105 + f"of potential backends: {available_backend_strs}."
106 + )
107 +
108 + def _make_log_unsupported(backend: W4A8Int8MoeBackend, reason: str | None) -> str:
109 + if reason:
110 + return (
111 + f"W4A8 Int8 MoE backend {backend.value} does not support the "
112 + f"deployment configuration since {reason}."
113 + )
114 + else:
115 + return (
116 + f"W4A8 Int8 MoE backend '{backend.value}' does not support the "
117 + "deployment configuration."
118 + )
119 +
120 + def _return_or_raise(
121 + backend: W4A8Int8MoeBackend,
122 + ) -> tuple[W4A8Int8MoeBackend, type[mk.FusedMoEExperts]]:
123 + reason = None
124 + for k_cls in backend_to_kernel_cls(backend):
125 + supported, reason = k_cls.is_supported_config(
126 + k_cls, config, weight_key, activation_key, activation_format
127 + )
128 + if supported:
129 + logger.info_once(_make_log_backend(backend))
130 + return backend, k_cls
131 + raise ValueError(_make_log_unsupported(backend, reason))
132 +
133 + # Handle explicit moe_backend from user.
134 + runner_backend = config.moe_backend
135 + if runner_backend != "auto":
136 + requested_backend = map_w4a8_int8_backend(runner_backend)
137 + return _return_or_raise(requested_backend)
138 +
139 + # Select kernels in order of backend.
140 + for backend in AVAILABLE_BACKENDS:
141 + for k_cls in backend_to_kernel_cls(backend):
142 + supported, reason = k_cls.is_supported_config(
143 + k_cls,
144 + config,
145 + weight_key,
146 + activation_key,
147 + activation_format,
148 + )
149 + if supported:
150 + logger.info_once(_make_log_backend(backend))
151 + return backend, k_cls
152 + else:
153 + logger.debug_once(_make_log_unsupported(backend, reason))
154 +
155 + raise NotImplementedError(
156 + "No W4A8 Int8 MoE backend supports the deployment configuration."
157 + )
158 +
159 +
160 +def make_w4a8_int8_moe_quant_config(
161 + block_shape: tuple[int, int] | None = None,
162 +) -> FusedMoEQuantConfig:
163 + """
164 + Create FusedMoEQuantConfig for W4A8 Int8 MoE.
165 +
166 + Args:
167 + block_shape: Quantization block shape (row, col).
168 + For channel-wise: (-1, 1) or None
169 + For group-wise: (1, group_size)
170 +
171 + Returns:
172 + FusedMoEQuantConfig with appropriate settings for W4A8 Int8
173 + """
174 + # W4A8 Int8 uses static weight quantization, dynamic activation quantization
175 + # Weights are 4-bit (stored as int8, packed to uint8),
176 + # activations are dynamically quantized to 8-bit in kernel
177 +
178 + group_shape = GroupShape(*block_shape) if block_shape is not None else None
179 +
180 + return FusedMoEQuantConfig(
181 + # Activations: unquantized (FP/BF16), dynamically quantized in kernel
182 + _a1=FusedMoEQuantDesc(shape=group_shape),
183 + _a2=FusedMoEQuantDesc(shape=group_shape),
184 + # Weights: INT8 (4-bit values), pre-packed with scales
185 + # dtype=None means already quantized/packed
186 + _w1=FusedMoEQuantDesc(dtype=None, shape=group_shape),
187 + _w2=FusedMoEQuantDesc(dtype=None, shape=group_shape),
188 + )
189 +
190 +
191 +def pack_int4_weights_for_kleidi(
192 + int4_as_int8: torch.Tensor,
193 + scales: torch.Tensor,
194 + bias: torch.Tensor | None,
195 + in_features: int,
196 + out_features: int,
197 + group_size: int,
198 +) -> torch.Tensor:
199 + """
200 + Pack INT4 weights (stored as int8 in [-8,7]) to KleidiAI format.
201 +
202 + Args:
203 + int4_as_int8: [out, in] int8 tensor with values in [-8, 7]
204 + scales: [out, in//group_size] or [out, 1] for channel-wise
205 + bias: [out] optional bias
206 + in_features: Input dimension
207 + out_features: Output dimension
208 + group_size: Quantization group size (-1 for channel-wise)
209 +
210 + Returns:
211 + Packed weight tensor in KleidiAI format
212 + """
213 + # Shift to unsigned nibble [0, 15]
214 + tmp = int4_as_int8.add(8)
215 + # Pack pairs along input dimension
216 + uint8_nibbles = ((tmp[:, 1::2] << 4) | tmp[:, ::2]).to(torch.uint8)
217 +
218 + # Determine scale dtype based on group_size
219 + # KleidiAI groupwise kernels accept bfloat16 scales
220 + # KleidiAI channelwise kernels accept float32 scales
221 + scale_dtype = torch.float32 if group_size == -1 else torch.bfloat16
222 + scales_typed = scales.to(scale_dtype)
223 + bias_typed = None if bias is None else bias.to(torch.float32)
224 +
225 + # Pack using KleidiAI op
226 + actual_group_size = in_features if group_size == -1 else group_size
227 + return torch.ops.aten._dyn_quant_pack_4bit_weight(
228 + uint8_nibbles,
229 + scales_typed,
230 + bias_typed,
231 + actual_group_size,
232 + in_features,
233 + out_features,
234 + )
235 +
236 +
237 +def convert_to_w4a8_int8_moe_format(
238 + w13_weight: torch.Tensor,
239 + w2_weight: torch.Tensor,
240 + w13_weight_scale: torch.Tensor,
241 + w2_weight_scale: torch.Tensor,
242 + group_size: int,
243 + w13_bias: torch.Tensor | None = None,
244 + w2_bias: torch.Tensor | None = None,
245 +) -> tuple[
246 + torch.Tensor,
247 + torch.Tensor,
248 + torch.Tensor,
249 + torch.Tensor,
250 + torch.Tensor | None,
251 + torch.Tensor | None,
252 +]:
253 + """
254 + Pack INT4 MoE weights to KleidiAI format.
255 +
256 + This function packs the INT4 weights (stored as int8 values) into
257 + the format expected by the KleidiAI dynamic_4bit_int_moe kernel.
258 +
259 + Args:
260 + w13_weight: [E, 2*IN, H] int8 tensor (int4 values in [-8,7])
261 + w2_weight: [E, H, IN] int8 tensor (int4 values in [-8,7])
262 + w13_weight_scale: [E, 2*IN, H/g or 1] scale tensor
263 + w2_weight_scale: [E, H, IN/g or 1] scale tensor
264 + group_size: Quantization group size (-1 for channel-wise)
265 + w13_bias: Optional [E, 2*IN] bias tensor
266 + w2_bias: Optional [E, H] bias tensor
267 +
268 + Returns:
269 + Tuple of (w13_packed, w2_packed) tensors
270 + """
271 + # Derive dimensions from tensor shapes
272 + E = w13_weight.shape[0] # num_experts
273 + I2 = w13_weight.shape[1] # w13_out_features (2*IN)
274 + H = w13_weight.shape[2] # w13_in_features (hidden_size)
275 + IN = w2_weight.shape[2] # w2_in_features (intermediate_size)
276 + w2_out_features = w2_weight.shape[1] # Should equal H
277 +
278 + # Pack per expert
279 + w13_packed_list = []
280 + w2_packed_list = []
281 +
282 + for e in range(E):
283 + w13_packed_list.append(
284 + pack_int4_weights_for_kleidi(
285 + w13_weight[e], # [2I, H]
286 + w13_weight_scale[e], # [2I, H/g or 1]
287 + w13_bias[e] if w13_bias is not None else None, # [2I]
288 + H,
289 + I2,
290 + group_size,
291 + )
292 + )
293 + w2_packed_list.append(
294 + pack_int4_weights_for_kleidi(
295 + w2_weight[e], # [H, IN]
296 + w2_weight_scale[e], # [H, IN/g or 1]
297 + w2_bias[e] if w2_bias is not None else None, # [H]
298 + IN,
299 + w2_out_features, # in_features=IN, out_features=H
300 + group_size,
301 + )
302 + )
303 +
304 + # Stack all experts
305 + w13_packed = torch.stack(w13_packed_list, dim=0)
306 + w2_packed = torch.stack(w2_packed_list, dim=0)
307 + empty = torch.empty(0)
308 +
309 + return w13_packed, w2_packed, empty, empty, empty, empty
310 +
311 +
312 +def make_w4a8_int8_moe_kernel(
313 + moe_quant_config: FusedMoEQuantConfig,
314 + moe_config: FusedMoEConfig,
315 + experts_cls: type[mk.FusedMoEExperts],
316 + routing_tables: tuple[torch.Tensor, torch.Tensor, torch.Tensor] | None = None,
317 +) -> mk.FusedMoEKernel:
318 + """
319 + Create FusedMoEKernel for W4A8 Int8 MoE.
320 +
321 + Args:
322 + moe_quant_config: Quantization configuration
323 + moe_config: MoE configuration
324 + experts_cls: Expert kernel class (should be CPUExpertsInt4)
325 + routing_tables: Optional routing tables for expert parallelism
326 +
327 + Returns:
328 + Configured FusedMoEKernel instance
329 + """
330 + # Create Prepare/Finalize.
331 + prepare_finalize = maybe_make_prepare_finalize(
332 + moe=moe_config,
333 + quant_config=moe_quant_config,
334 + routing_tables=routing_tables,
335 + allow_new_interface=True,
336 + use_monolithic=issubclass(experts_cls, mk.FusedMoEExpertsMonolithic),
337 + )
338 + assert prepare_finalize is not None
339 +
340 + logger.info_once("Using %s", prepare_finalize.__class__.__name__)
341 +
342 + # Create Experts.
343 + # W4A8 Int8 currently only supports monolithic interface
344 + if not issubclass(experts_cls, mk.FusedMoEExpertsMonolithic):
345 + raise ValueError(
346 + f"W4A8 Int8 MoE only supports monolithic experts, "
347 + f"but got {experts_cls.__name__}"
348 + )
349 +
350 + experts = experts_cls(
351 + moe_config=moe_config,
352 + quant_config=moe_quant_config,
353 + )
354 +
355 + kernel = mk.FusedMoEKernel(
356 + prepare_finalize,
357 + experts,
358 + inplace=not moe_config.disable_inplace,
359 + )
360 +
361 + return kernel
vllm/model_executor/layers/quantization/compressed_tensors/compressed_tensors_moe/compressed_tensors_moe_w4a8_int8.py
+105 -122
@@ -11,16 +11,26 @@ from compressed_tensors.quantization import (
11 from vllm.logger import init_logger
12 from vllm.model_executor.layers.fused_moe import (
13 RoutedExperts,
14 + SharedExperts,
15 )
15 -from vllm.model_executor.layers.fused_moe.activation import MoEActivation
16 from vllm.model_executor.layers.fused_moe.config import (
17 FusedMoEConfig,
18 FusedMoEQuantConfig,
19 )
20 -from vllm.model_executor.layers.fused_moe.cpu_fused_moe import select_experts
20 +from vllm.model_executor.layers.fused_moe.oracle.w4a8_int8 import (
21 + convert_to_w4a8_int8_moe_format,
22 + make_w4a8_int8_moe_kernel,
23 + make_w4a8_int8_moe_quant_config,
24 + select_w4a8_int8_moe_backend,
25 +)
26 from vllm.model_executor.layers.quantization.compressed_tensors.compressed_tensors_moe import ( # noqa E501
27 CompressedTensorsMoEMethod,
28 )
29 +from vllm.model_executor.layers.quantization.utils.quant_utils import (
30 + GroupShape,
31 + QuantKey,
32 + ScaleDesc,
33 +)
34 from vllm.model_executor.utils import replace_parameter, set_weight_attrs
35 from vllm.platforms import CpuArchEnum, current_platform
36
@@ -48,6 +58,11 @@ class CompressedTensorsW4A8Int8MoEMethod(CompressedTensorsMoEMethod):
58 self.has_bias = self.moe.has_bias
59 self.weight_quant = weight_quant
60 self.input_quant = input_quant
61 + self.static_input_scales = False # always dynamic per token
62 + # Weight can be channel-wise (group_size=None) or group-wise
63 + self.group_size = (
64 + weight_quant.group_size if (weight_quant.group_size is not None) else -1
65 + )
66
67 # Validate scheme: weights=W4 (channel or group),
68 # activations=dynamic TOKEN (A8)
@@ -61,17 +76,9 @@ class CompressedTensorsW4A8Int8MoEMethod(CompressedTensorsMoEMethod):
76 "W4A8-int MoE needs dynamic per-token activation quantization."
77 )
78
64 - # Weight can be channel-wise (group_size=None) or group-wise
65 - self.group_size = (
66 - weight_quant.group_size if (weight_quant.group_size is not None) else -1
67 - )
79 if weight_quant.num_bits != 4:
80 raise ValueError("This method only supports 4-bit weights (num_bits=4).")
81
71 - # CPU only
72 - if not current_platform.is_cpu():
73 - raise ValueError("CompressedTensorsW4A8Int8MoEMethod is CPU-only.")
74 -
82 # Arm: check _dyn ops availability
83 if current_platform.get_cpu_architecture() == CpuArchEnum.ARM:
84 try:
@@ -82,7 +89,26 @@ class CompressedTensorsW4A8Int8MoEMethod(CompressedTensorsMoEMethod):
89 f"""PyTorch {torch.__version__} lacks _dyn_quant_* 4bit ops;
90 install a newer build."""
91 ) from err
85 - self.static_input_scales = False # always dynamic per token
92 +
93 + # Construct QuantKey for weights from QuantizationArgs
94 + # W4A8 INT4: 4-bit weights (stored as int8), static quantization
95 + if self.group_size == -1:
96 + # Channel-wise quantization
97 + group_shape = GroupShape(-1, 1)
98 + scale_dtype = torch.float32
99 + else:
100 + # Group-wise quantization
101 + group_shape = GroupShape(1, self.group_size)
102 + scale_dtype = torch.bfloat16
103 +
104 + weight_scale_desc = ScaleDesc(scale_dtype, static=True, group_shape=group_shape)
105 + weight_key = QuantKey(torch.int8, weight_scale_desc, symmetric=True)
106 +
107 + self.backend, self.experts_cls = select_w4a8_int8_moe_backend(
108 + moe,
109 + weight_key,
110 + activation_key=None, # unquantized inputs
111 + )
112
113 # ---- parameter creation ----
114 def create_weights(
@@ -182,72 +208,20 @@ class CompressedTensorsW4A8Int8MoEMethod(CompressedTensorsMoEMethod):
208
209 # post-load packing to dyn-4bit KleidiAI kernel's format
210 def process_weights_after_loading(self, layer: torch.nn.Module) -> None:
185 - E = layer.w13_weight.shape[0]
186 - H = layer.w13_in_features
187 - I2 = layer.w13_out_features
188 - IN = layer.w2_in_features
189 - g = layer.group_size
190 -
191 - def _pack_matrix(
192 - int4_as_int8_2d: torch.Tensor,
193 - scales_2d: torch.Tensor,
194 - bias_1d: torch.Tensor | None,
195 - in_features: int,
196 - out_features: int,
197 - ) -> torch.Tensor:
198 - # int4 values are stored as int8 in [-8,7].
199 - # Shift to unsigned nibble and pack pairs along input-dim.
200 - tmp = int4_as_int8_2d.add(8) # [out, in]
201 - uint8_nibbles = ((tmp[:, 1::2] << 4) | tmp[:, ::2]).to(
202 - torch.uint8
203 - ) # [out, in//2]
204 -
205 - # KleidiAI groupwise kernels accepts float32 scales
206 - # KleidiAI groupwise kernels accepts bfloat16 scales
207 - scale_dtype = torch.float32 if g == -1 else torch.bfloat16
208 - scales = scales_2d.to(scale_dtype)
209 - bias = None if bias_1d is None else bias_1d.to(torch.float32)
210 - return torch.ops.aten._dyn_quant_pack_4bit_weight(
211 - uint8_nibbles,
212 - scales,
213 - bias,
214 - g if g != -1 else in_features,
215 - in_features,
216 - out_features,
217 - )
218 -
219 - # Pack per expert
220 - w13_packed_list = []
221 - w2_packed_list = []
222 -
223 - has_w13_bias = hasattr(layer, "w13_bias") and layer.w13_bias is not None
224 - has_w2_bias = hasattr(layer, "w2_bias") and layer.w2_bias is not None
225 -
226 - for e in range(E):
227 - w13_packed_list.append(
228 - _pack_matrix(
229 - layer.w13_weight[e], # [2I, H]
230 - layer.w13_weight_scale[e], # [2I, H/g or 1]
231 - layer.w13_bias[e] if has_w13_bias else None, # [2I]
232 - H,
233 - I2,
234 - )
211 + # Use oracle to pack weights.
212 + w13_packed, w2_packed, w13_weight_scale, w2_weight_scale, w13_bias, w2_bias = (
213 + convert_to_w4a8_int8_moe_format(
214 + w13_weight=layer.w13_weight,
215 + w2_weight=layer.w2_weight,
216 + w13_weight_scale=layer.w13_weight_scale,
217 + w2_weight_scale=layer.w2_weight_scale,
218 + group_size=self.group_size,
219 + w13_bias=layer.w13_bias if self.has_bias else None,
220 + w2_bias=layer.w2_bias if self.has_bias else None,
221 )
236 - w2_packed_list.append(
237 - _pack_matrix(
238 - # w2 shape is [H, IN]; we need [out, in] == [H, IN].
239 - layer.w2_weight[e], # [H, IN]
240 - layer.w2_weight_scale[e], # [H, IN/g or 1]
241 - layer.w2_bias[e] if has_w2_bias else None, # [H]
242 - IN,
243 - layer.w2_out_features, # in_features=IN, out_features=H
244 - )
245 - )
246 -
247 - # each packed tensor has identical shape per expert; stack on dim 0
248 - w13_packed = torch.stack(w13_packed_list, dim=0)
249 - w2_packed = torch.stack(w2_packed_list, dim=0)
222 + )
223
224 + # Register packed weights as parameters
225 replace_parameter(
226 layer,
227 "w13_weight_packed",
@@ -259,7 +233,6 @@ class CompressedTensorsW4A8Int8MoEMethod(CompressedTensorsMoEMethod):
233 torch.nn.Parameter(w2_packed, requires_grad=False),
234 )
235
262 - # free raw tensors/scales/bias now that they're packed into the payload.
236 replace_parameter(
237 layer, "w13_weight", torch.nn.Parameter(torch.empty(0), requires_grad=False)
238 )
@@ -269,36 +242,46 @@ class CompressedTensorsW4A8Int8MoEMethod(CompressedTensorsMoEMethod):
242 replace_parameter(
243 layer,
244 "w13_weight_scale",
272 - torch.nn.Parameter(torch.empty(0), requires_grad=False),
245 + torch.nn.Parameter(w13_weight_scale, requires_grad=False),
246 )
247 replace_parameter(
248 layer,
249 "w2_weight_scale",
277 - torch.nn.Parameter(torch.empty(0), requires_grad=False),
250 + torch.nn.Parameter(w2_weight_scale, requires_grad=False),
251 )
279 - if has_w13_bias:
252 +
253 + if self.has_bias:
254 replace_parameter(
255 layer,
256 "w13_bias",
283 - torch.nn.Parameter(torch.empty(0), requires_grad=False),
257 + torch.nn.Parameter(w13_bias, requires_grad=False),
258 )
285 - if has_w2_bias:
259 + if self.has_bias:
260 replace_parameter(
261 layer,
262 "w2_bias",
289 - torch.nn.Parameter(torch.empty(0), requires_grad=False),
263 + torch.nn.Parameter(w2_bias, requires_grad=False),
264 )
265
266 + quant_config = self.get_fused_moe_quant_config(layer)
267 + assert quant_config is not None
268 + assert self.experts_cls is not None
269 + self.moe_kernel = make_w4a8_int8_moe_kernel(
270 + moe_quant_config=quant_config,
271 + moe_config=self.moe,
272 + experts_cls=self.experts_cls,
273 + routing_tables=layer._expert_routing_tables(),
274 + )
275 +
276 def get_fused_moe_quant_config(
277 self, layer: torch.nn.Module
278 ) -> FusedMoEQuantConfig | None:
295 - # CPU dynamic 4-bit MoE path does not use modular kernels or
296 - # fused_experts; quant config is not needed.
297 - return None
279 + # Determine block shape from group_size
280 + # group_size=-1 means channel-wise: (-1, 1)
281 + # group_size=N means group-wise: (1, N)
282 + block_shape = (-1, 1) if self.group_size == -1 else (1, self.group_size)
283
299 - @property
300 - def is_monolithic(self) -> bool:
301 - return True
284 + return make_w4a8_int8_moe_quant_config(block_shape=block_shape)
285
286 def apply_monolithic(
287 self,
@@ -307,43 +290,43 @@ class CompressedTensorsW4A8Int8MoEMethod(CompressedTensorsMoEMethod):
290 router_logits: torch.Tensor,
291 input_ids: torch.Tensor | None = None,
292 ) -> torch.Tensor:
310 - assert layer.activation in (
311 - MoEActivation.SILU,
312 - MoEActivation.SWIGLUOAI,
313 - MoEActivation.SWIGLUSTEP,
314 - ), "Only SiLU/SwiGLUGU/SwiGLUUG are supported."
315 - assert layer.expert_map is None, """expert_map/EP not implemented
316 - for CPU dyn-4bit MoE."""
317 -
318 - def _act_kind(s: MoEActivation) -> int:
319 - # 0 = SwiGLU_Gu (SiLU(g)*u), 1 = SwiGLU_Ug (SiLU(u)*g), 2 = SiLU
320 - if s == MoEActivation.SWIGLUSTEP:
321 - return 0
322 - if s == MoEActivation.SWIGLUOAI:
323 - return 1
324 - if s == MoEActivation.SILU:
325 - return 2
326 - raise ValueError(f"Unknown activation '{s}'")
327 -
328 - # Apply topk softmax on router output
329 - topk_weights, topk_ids = select_experts(
330 - hidden_states=x,
331 - router_logits=router_logits,
332 - top_k=layer.top_k,
333 - use_grouped_topk=layer.use_grouped_topk,
334 - renormalize=layer.renormalize,
293 + assert self.is_monolithic
294 + assert self.moe_kernel is not None
295 + return self.moe_kernel.apply_monolithic(
296 + x,
297 + layer.w13_weight_packed,
298 + layer.w2_weight_packed,
299 + router_logits,
300 + activation=layer.activation,
301 + global_num_experts=layer.global_num_experts,
302 + expert_map=layer.expert_map,
303 + apply_router_weight_on_input=layer.apply_router_weight_on_input,
304 + num_expert_group=layer.num_expert_group,
305 + topk_group=layer.topk_group,
306 + e_score_correction_bias=layer.e_score_correction_bias,
307 + routed_scaling_factor=layer.routed_scaling_factor,
308 )
309
337 - return torch.ops._C.dynamic_4bit_int_moe(
310 + def apply(
311 + self,
312 + layer: RoutedExperts,
313 + x: torch.Tensor,
314 + topk_weights: torch.Tensor,
315 + topk_ids: torch.Tensor,
316 + shared_experts: SharedExperts | None,
317 + shared_experts_input: torch.Tensor | None,
318 + ) -> torch.Tensor:
319 + assert self.moe_kernel is not None
320 + return self.moe_kernel.apply(
321 x,
339 - topk_ids.to(torch.long),
340 - topk_weights,
322 layer.w13_weight_packed,
323 layer.w2_weight_packed,
343 - layer.w2_out_features,
344 - layer.w2_in_features,
345 - layer.w13_out_features,
346 - layer.group_size,
347 - layer.apply_router_weight_on_input,
348 - int(_act_kind(layer.activation)),
324 + topk_weights,
325 + topk_ids,
326 + activation=layer.activation,
327 + global_num_experts=layer.global_num_experts,
328 + expert_map=layer.expert_map,
329 + apply_router_weight_on_input=layer.apply_router_weight_on_input,
330 + shared_experts=shared_experts,
331 + shared_experts_input=shared_experts_input,
332 )
vllm/model_executor/layers/quantization/utils/quant_utils.py
+25
@@ -181,6 +181,31 @@ kInt8Static = QuantKey(INT8_DTYPE, scale=kInt8StaticGroupScale, symmetric=True)
181 kInt8StaticChannelSym = QuantKey(torch.int8, kStaticChannelScale, symmetric=True)
182 kInt8DynamicTokenSym = QuantKey(torch.int8, kDynamicTokenScale, symmetric=True)
183
184 +# INT4 W4A8 quantization keys
185 +
186 +# For group-wise quantization (e.g., group_size=128)
187 +# Note: group_size will be specified at runtime, this is a generic group scale
188 +kInt4W4A8StaticGroupScale128 = ScaleDesc(torch.bfloat16, True, GroupShape(1, 128))
189 +kInt4W4A8StaticGroup128Sym = QuantKey(
190 + torch.int8, kInt4W4A8StaticGroupScale128, symmetric=True
191 +)
192 +
193 +kInt4W4A8StaticGroupScale64 = ScaleDesc(torch.bfloat16, True, GroupShape(1, 64))
194 +kInt4W4A8StaticGroup64Sym = QuantKey(
195 + torch.int8, kInt4W4A8StaticGroupScale64, symmetric=True
196 +)
197 +
198 +kInt4W4A8StaticGroupScale32 = ScaleDesc(torch.bfloat16, True, GroupShape(1, 32))
199 +kInt4W4A8StaticGroup32Sym = QuantKey(
200 + torch.int8, kInt4W4A8StaticGroupScale32, symmetric=True
201 +)
202 +
203 +# Generic group-wise with flexible group size (per-token groups)
204 +kInt4W4A8StaticGroupScale = ScaleDesc(torch.bfloat16, True, GroupShape(1, -1))
205 +kInt4W4A8StaticGroupSym = QuantKey(
206 + torch.int8, kInt4W4A8StaticGroupScale, symmetric=True
207 +)
208 +
209
210 def create_fp8_quant_key(
211 static: bool,