[CPU][BugFix] Multiple fixes to w4a8_int8 CPU MoE path (#46739)
Signed-off-by: Fadi Arafeh <fadi.arafeh@arm.com>
Fadi Arafeh committed
Jul 6, 2026 at 06:39 UTC
f1073c050c1c0149430738fd1c7cc078d4883900
10 files changed
+192
-190
.buildkite/scripts/hardware_ci/run-cpu-test-arm.sh
+2
-1
@@ -38,7 +38,8 @@ function cpu_tests() {
38
pytest -x -v -s tests/kernels/attention/test_cpu_attn.py
39
pytest -x -v -s tests/kernels/core/test_cpu_activation.py
40
pytest -x -v -s tests/kernels/moe/test_cpu_fused_moe.py
41
- pytest -x -v -s tests/kernels/mamba/cpu/test_cpu_gdn_ops.py"
41
+ pytest -x -v -s tests/kernels/mamba/cpu/test_cpu_gdn_ops.py
42
+ pytest -x -v -s tests/kernels/moe/test_cpu_int4_moe.py"
43
44
# skip tests requiring model downloads if HF_TOKEN is not set
45
# due to rate-limits
csrc/cpu/torch_bindings.cpp
+2
-1
@@ -278,7 +278,8 @@ TORCH_LIBRARY_EXPAND(TORCH_EXTENSION_NAME, ops) {
278
ops.def(
279
"dynamic_4bit_int_moe("
280
"Tensor x, Tensor topk_ids, Tensor topk_weights,"
281
- "Tensor w13_packed, Tensor w2_packed, int H, int I, int I2,"
281
+ "Tensor w13_packed, Tensor w2_packed,"
282
+ "int hidden_size, int intermediate_size,"
283
"int group_size, bool apply_router_weight_on_input, int activation_kind"
284
") -> Tensor");
285
csrc/moe/dynamic_4bit_int_moe_cpu.cpp
+54
-29
@@ -29,25 +29,37 @@ enum ActivationKind : int64_t {
29
30
torch::Tensor dynamic_4bit_int_moe_cpu(
31
torch::Tensor x, torch::Tensor topk_ids, torch::Tensor topk_weights,
32
- torch::Tensor w13_packed, torch::Tensor w2_packed, int64_t H, int64_t I,
33
- int64_t I2, int64_t group_size, bool apply_router_weight_on_input,
34
- int64_t activation_kind) {
32
+ torch::Tensor w13_packed, torch::Tensor w2_packed, int64_t hidden_size,
33
+ int64_t intermediate_size, int64_t group_size,
34
+ bool apply_router_weight_on_input, int64_t activation_kind) {
35
TORCH_CHECK(x.dim() == 2, "x must be 2D");
36
TORCH_CHECK(topk_ids.dim() == 2 && topk_weights.dim() == 2,
37
"topk tensors must be [T, K]");
38
TORCH_CHECK(
39
w13_packed.size(0) == w2_packed.size(0),
40
"w13_packed and w2_packed must have same number of experts in dim 0");
41
- TORCH_CHECK(I2 == 2 * I, "I2 must equal 2*I");
41
42
const int64_t T = x.size(0);
43
const int64_t K = topk_ids.size(1);
44
const int64_t E = w13_packed.size(0);
45
const int64_t N = T * K;
46
+ const int64_t w13_out_features = 2 * intermediate_size;
47
48
auto x_c = x.contiguous();
49
+ // _dyn_quant_matmul_4bit kernel natively supports these pre-quant activation
50
+ // dtypes:
51
+ // - fp32: with channelwise and groupwise
52
+ // - bf16: with channelwise -> upcast to fp32 for groupwise
53
+ // - fp16: not supported -> upcast to fp32 for groupwise & channelwise
54
+ const auto output_dtype = x_c.scalar_type();
55
+ const bool should_cast_input =
56
+ ((group_size != -1) && output_dtype == at::kBFloat16) ||
57
+ output_dtype == at::kHalf;
58
+ if (should_cast_input) {
59
+ x_c = x_c.to(at::kFloat);
60
+ }
61
auto ids_c = topk_ids.contiguous();
50
- auto gates_c = topk_weights.to(at::kFloat).contiguous();
62
+ auto gates_c = topk_weights.to(x_c.scalar_type()).contiguous();
63
64
// bucketing tokens -> experts
65
c10::SmallVector<int64_t, 64> counts(
@@ -63,35 +75,42 @@ torch::Tensor dynamic_4bit_int_moe_cpu(
75
c10::SmallVector<int64_t, 65> offsets(E + 1, 0); // ( E +1 )
76
for (int64_t e = 0; e < E; ++e) offsets[e + 1] = offsets[e] + counts[e];
77
78
+ // expert_tokens = [tokens indices for expert 0, ...]
79
+ // expert_gates = [router weights for tokens assigned to expert 0, ...]
80
auto expert_tokens = at::empty({offsets[E]}, ids_c.options());
81
auto expert_gates = at::empty({offsets[E]}, gates_c.options());
82
{
83
c10::SmallVector<int64_t, 64> cursor(E, 0);
70
- const auto* ids_ptr = ids_c.data_ptr<int64_t>();
71
- const auto* gts_ptr = gates_c.data_ptr<float>();
72
- auto* tok_ptr = expert_tokens.data_ptr<int64_t>();
73
- auto* gate_ptr = expert_gates.data_ptr<float>();
74
-
75
- for (int64_t t = 0; t < T; ++t) {
76
- const int64_t base = t * K;
77
- for (int64_t k = 0; k < K; ++k) {
78
- const int64_t idx = base + k;
79
- const int64_t e = ids_ptr[idx];
80
- const int64_t p = offsets[e] + (cursor[e]++);
81
- tok_ptr[p] = t;
82
- gate_ptr[p] = gts_ptr[idx];
83
- }
84
- }
84
+ AT_DISPATCH_FLOATING_TYPES_AND2(
85
+ at::ScalarType::BFloat16, at::ScalarType::Half, gates_c.scalar_type(),
86
+ "bucket_expert_tokens_and_gates", [&] {
87
+ const auto* ids_ptr = ids_c.data_ptr<int64_t>();
88
+ const auto* gts_ptr = gates_c.data_ptr<scalar_t>();
89
+ auto* tok_ptr = expert_tokens.data_ptr<int64_t>();
90
+ auto* gate_ptr = expert_gates.data_ptr<scalar_t>();
91
+
92
+ for (int64_t t = 0; t < T; ++t) {
93
+ const int64_t base = t * K;
94
+ for (int64_t k = 0; k < K; ++k) {
95
+ const int64_t idx = base + k;
96
+ const int64_t e = ids_ptr[idx];
97
+ const int64_t p = offsets[e] + (cursor[e]++);
98
+ tok_ptr[p] = t;
99
+ gate_ptr[p] = gts_ptr[idx];
100
+ }
101
+ }
102
+ });
103
}
104
87
- const int64_t g_eff_13 = (group_size != -1) ? group_size : H;
88
- const int64_t g_eff_2 = (group_size != -1) ? group_size : I;
105
+ const int64_t g_eff_13 = (group_size != -1) ? group_size : hidden_size;
106
+ const int64_t g_eff_2 = (group_size != -1) ? group_size : intermediate_size;
107
108
+ // X_all [num_tokens * K, hidden_size]
109
auto X_all = x_c.index_select(/*dim=*/0, expert_tokens);
110
if (apply_router_weight_on_input) {
111
X_all = X_all.mul(expert_gates.unsqueeze(1));
112
}
94
- auto Y_all = at::empty({offsets[E], H}, x_c.options());
113
+ auto Y_all = at::empty({offsets[E], hidden_size}, x_c.options());
114
115
at::parallel_for(0, offsets[E], 0, [&](int64_t idx_begin, int64_t idx_end) {
116
c10::InferenceMode guard;
@@ -109,11 +128,13 @@ torch::Tensor dynamic_4bit_int_moe_cpu(
128
auto w2_e = w2_packed.select(/*dim=*/0, e);
129
130
// W13
112
- auto y13 =
113
- mm(x_e, w13_e, g_eff_13, /*in_features=*/H, /*out_features=*/I2);
131
+ auto y13 = mm(x_e, w13_e, g_eff_13, /*in_features=*/hidden_size,
132
+ /*out_features=*/w13_out_features);
133
115
- auto g_part = y13.narrow(/*dim=*/1, /*start=*/0, /*length=*/I);
116
- auto u_part = y13.narrow(/*dim=*/1, /*start=*/I, /*length=*/I);
134
+ auto g_part =
135
+ y13.narrow(/*dim=*/1, /*start=*/0, /*length=*/intermediate_size);
136
+ auto u_part = y13.narrow(/*dim=*/1, /*start=*/intermediate_size,
137
+ /*length=*/intermediate_size);
138
139
torch::Tensor act;
140
if (activation_kind == ActivationKind::SwiGLUOAI) { // SwiGLUOAI
@@ -128,7 +149,8 @@ torch::Tensor dynamic_4bit_int_moe_cpu(
149
}
150
151
// W2
131
- auto y = mm(act, w2_e, g_eff_2, /*in_features=*/I, /*out_features=*/H);
152
+ auto y = mm(act, w2_e, g_eff_2, /*in_features=*/intermediate_size,
153
+ /*out_features=*/hidden_size);
154
155
// Store per-expert result
156
Y_all.narrow(/*dim=*/0, /*start=*/start, /*length=*/te).copy_(y);
@@ -138,8 +160,11 @@ torch::Tensor dynamic_4bit_int_moe_cpu(
160
if (!apply_router_weight_on_input) {
161
Y_all = Y_all.mul(expert_gates.unsqueeze(1));
162
}
163
+ if (Y_all.scalar_type() != output_dtype) {
164
+ Y_all = Y_all.to(output_dtype);
165
+ }
166
142
- auto out = at::zeros({T, H}, x.options());
167
+ auto out = at::zeros({T, hidden_size}, x.options());
168
out =
169
at::index_add(out, /*dim=*/0, /*index=*/expert_tokens, /*source=*/Y_all);
170
csrc/ops.h
+3
-3
@@ -53,9 +53,9 @@ void dynamic_scaled_int8_quant(torch::Tensor& out, torch::Tensor const& input,
53
54
torch::Tensor dynamic_4bit_int_moe_cpu(
55
torch::Tensor x, torch::Tensor topk_ids, torch::Tensor topk_weights,
56
- torch::Tensor w13_packed, torch::Tensor w2_packed, int64_t H, int64_t I,
57
- int64_t I2, int64_t group_size, bool apply_router_weight_on_input,
58
- int64_t activation_kind);
56
+ torch::Tensor w13_packed, torch::Tensor w2_packed, int64_t hidden_size,
57
+ int64_t intermediate_size, int64_t group_size,
58
+ bool apply_router_weight_on_input, int64_t activation_kind);
59
60
using fptr_t = int64_t;
61
#ifdef USE_ROCM
tests/kernels/moe/test_cpu_int4_moe.py
+44
-114
@@ -8,19 +8,21 @@ import pytest
8
import torch
9
import torch.nn.functional as F
10
11
-from vllm.platforms import current_platform
11
+from vllm.model_executor.layers.fused_moe.activation import MoEActivation
12
+from vllm.model_executor.layers.fused_moe.experts.cpu_int4_moe import (
13
+ CPUExpertsInt4,
14
+)
15
+from vllm.model_executor.layers.fused_moe.oracle.w4a8_int8 import (
16
+ convert_to_w4a8_int8_moe_format,
17
+)
18
+from vllm.platforms import CpuArchEnum, current_platform
19
from vllm.utils.torch_utils import set_random_seed
20
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)
21
+if (
22
+ not current_platform.is_cpu()
23
+ or current_platform.get_cpu_architecture() != CpuArchEnum.ARM
24
+):
25
+ pytest.skip("skipping Arm CPU-only tests", allow_module_level=True)
26
27
28
# Tolerance for INT4 W4A8
@@ -34,49 +36,6 @@ def _silu_and_mul(x: torch.Tensor) -> torch.Tensor:
36
return F.silu(x[..., :d]) * x[..., d:]
37
38
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
-
39
def _make_int4_moe_weights(
40
E: int,
41
N: int,
@@ -124,59 +83,29 @@ def _make_int4_moe_weights(
83
w13_bias = torch.randn(E, 2 * N, dtype=torch.float32) * 0.01
84
w2_bias = torch.randn(E, K, dtype=torch.float32) * 0.01
85
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()
86
+ w13_packed, w2_packed, *_ = convert_to_w4a8_int8_moe_format(
87
+ w13_weight=w13_int4,
88
+ w2_weight=w2_int4,
89
+ w13_weight_scale=w13_scales,
90
+ w2_weight_scale=w2_scales,
91
+ group_size=group_size,
92
+ w13_bias=w13_bias if has_bias else None,
93
+ w2_bias=w2_bias if has_bias else None,
94
+ )
95
+
96
+ if group_size == -1:
97
+ w13_scale = w13_scales.float()
98
+ w2_scale = w2_scales.float()
99
+ else:
100
+ w13_scale = w13_scales.float().repeat_interleave(group_size, dim=-1)
101
+ w2_scale = w2_scales.float().repeat_interleave(group_size, dim=-1)
102
+
103
+ w13_ref = w13_int4.float() * w13_scale
104
+ w2_ref = w2_int4.float() * w2_scale
105
+ if has_bias and w13_bias is not None:
106
+ w13_ref = w13_ref + w13_bias.float().unsqueeze(-1)
107
+ if has_bias and w2_bias is not None:
108
+ w2_ref = w2_ref + w2_bias.float().unsqueeze(-1)
109
110
return w13_packed, w2_packed, w13_ref, w2_ref, w13_bias, w2_bias
111
@@ -233,17 +162,20 @@ MoE_CONFIGS = [
162
(768, 2048, 16, 4, 64),
163
]
164
SEEDS = [0, 42]
165
+ACTIVATION_DTYPES = [torch.float32, torch.bfloat16, torch.float16]
166
167
168
@pytest.mark.parametrize("M", NUM_TOKENS)
169
@pytest.mark.parametrize("N,K,E,topk,group_size", MoE_CONFIGS)
170
@pytest.mark.parametrize("seed", SEEDS)
241
-def test_cpu_int4_moe_kernel(M, N, K, E, topk, group_size, seed):
171
+@pytest.mark.parametrize("activation_dtype", ACTIVATION_DTYPES)
172
+def test_cpu_int4_moe_kernel(M, N, K, E, topk, group_size, seed, activation_dtype):
173
"""Test dynamic_4bit_int_moe kernel against dequantized torch reference."""
174
set_random_seed(seed)
175
+ activation = MoEActivation.SILU
176
177
# Generate input activations
246
- a = torch.randn(M, K, dtype=torch.bfloat16) / (K**0.5)
178
+ a = torch.randn(M, K, dtype=activation_dtype) / (K**0.5)
179
180
# Generate INT4 weights
181
w13_packed, w2_packed, w13_ref, w2_ref, w13_bias, w2_bias = _make_int4_moe_weights(
@@ -266,8 +198,6 @@ def test_cpu_int4_moe_kernel(M, N, K, E, topk, group_size, seed):
198
)
199
200
# Test dynamic_4bit_int_moe kernel
269
- # Activation kind: 1 = SwiGLU_Ug (SiLU(u)*g) for OAI-style
270
- activation_kind = 1
201
apply_router_weight_on_input = False
202
203
out = torch.ops._C.dynamic_4bit_int_moe(
@@ -278,14 +208,14 @@ def test_cpu_int4_moe_kernel(M, N, K, E, topk, group_size, seed):
208
w2_packed,
209
K, # H (hidden_size / w2_out_features)
210
N, # I (intermediate_size / w2_in_features)
281
- 2 * N, # I2 (2*intermediate_size / w13_out_features)
211
group_size,
212
apply_router_weight_on_input,
284
- activation_kind,
213
+ CPUExpertsInt4._activation_kind(activation),
214
)
215
216
+ assert out.dtype == activation_dtype
217
torch.testing.assert_close(
288
- ref_out.bfloat16(),
218
+ ref_out,
219
out,
220
atol=INT4_W4A8_ATOL,
221
rtol=INT4_W4A8_RTOL,
vllm/_custom_ops.py
+18
@@ -2974,6 +2974,24 @@ if hasattr(torch.ops._C, "fused_experts_cpu"):
2974
return torch.empty_like(hidden_states)
2975
2976
2977
+if hasattr(torch.ops._C, "dynamic_4bit_int_moe"):
2978
+
2979
+ @register_fake("_C::dynamic_4bit_int_moe")
2980
+ def dynamic_4bit_int_moe_fake(
2981
+ x: torch.Tensor,
2982
+ topk_ids: torch.Tensor,
2983
+ topk_weights: torch.Tensor,
2984
+ w13_packed: torch.Tensor,
2985
+ w2_packed: torch.Tensor,
2986
+ hidden_size: int,
2987
+ intermediate_size: int,
2988
+ group_size: int,
2989
+ apply_router_weight_on_input: bool,
2990
+ activation_kind: int,
2991
+ ) -> torch.Tensor:
2992
+ return x.new_empty((x.size(0), hidden_size))
2993
+
2994
+
2995
def fused_experts_cpu(
2996
hidden_states: torch.Tensor,
2997
w1: torch.Tensor,
vllm/model_executor/layers/fused_moe/experts/cpu_int4_moe.py
+54
-16
@@ -14,12 +14,13 @@ from vllm.model_executor.layers.fused_moe.config import (
14
)
15
from vllm.model_executor.layers.quantization.utils.quant_utils import (
16
QuantKey,
17
+ kInt4W4A8StaticChannelSym,
18
kInt4W4A8StaticGroup32Sym,
19
kInt4W4A8StaticGroup64Sym,
20
kInt4W4A8StaticGroup128Sym,
21
kInt4W4A8StaticGroupSym,
22
)
22
-from vllm.platforms import current_platform
23
+from vllm.platforms import CpuArchEnum, current_platform
24
25
26
class CPUExpertsInt4(mk.FusedMoEExpertsMonolithic):
@@ -48,6 +49,48 @@ class CPUExpertsInt4(mk.FusedMoEExpertsMonolithic):
49
def activation_format() -> mk.FusedMoEActivationFormat:
50
return mk.FusedMoEActivationFormat.Standard
51
52
+ @staticmethod
53
+ def is_supported_config(
54
+ cls: type[mk.FusedMoEExperts],
55
+ moe_config: FusedMoEConfig,
56
+ weight_key: QuantKey | None,
57
+ activation_key: QuantKey | None,
58
+ activation_format: mk.FusedMoEActivationFormat,
59
+ ) -> tuple[bool, str | None]:
60
+ if (
61
+ not current_platform.is_cpu()
62
+ or current_platform.get_cpu_architecture() != CpuArchEnum.ARM
63
+ ):
64
+ return False, "kernel only supports Arm CPU"
65
+
66
+ if moe_config.in_dtype not in (
67
+ torch.float32,
68
+ torch.bfloat16,
69
+ torch.float16,
70
+ ):
71
+ return (
72
+ False,
73
+ f"kernel does not support {moe_config.in_dtype} input/output dtype",
74
+ )
75
+
76
+ try:
77
+ _ = torch.ops.aten._dyn_quant_matmul_4bit
78
+ _ = torch.ops.aten._dyn_quant_pack_4bit_weight
79
+ except AttributeError:
80
+ return (
81
+ False,
82
+ f"PyTorch {torch.__version__} does not support "
83
+ "_dyn_quant_* 4bit ops. Install a newer version",
84
+ )
85
+
86
+ return mk.FusedMoEExperts.is_supported_config(
87
+ cls,
88
+ moe_config,
89
+ weight_key,
90
+ activation_key,
91
+ activation_format,
92
+ )
93
+
94
@staticmethod
95
def _supports_current_device() -> bool:
96
return current_platform.is_cpu()
@@ -86,8 +129,9 @@ class CPUExpertsInt4(mk.FusedMoEExpertsMonolithic):
129
Can be channel-wise or group-wise quantization
130
- Activations: dynamic per-token 8-bit integer quantization
131
"""
89
- # group size must be multiple of 32
132
+ # channelwise or groupwise with group size being a multiple of 32
133
SUPPORTED_W_A = [
134
+ (kInt4W4A8StaticChannelSym, None),
135
(kInt4W4A8StaticGroup128Sym, None),
136
(kInt4W4A8StaticGroup64Sym, None),
137
(kInt4W4A8StaticGroup32Sym, None),
@@ -120,7 +164,8 @@ class CPUExpertsInt4(mk.FusedMoEExpertsMonolithic):
164
"""Expert parallelism not yet supported."""
165
return False
166
123
- def _activation_kind(self, activation: MoEActivation) -> int:
167
+ @staticmethod
168
+ def _activation_kind(activation: MoEActivation) -> int:
169
"""Convert MoEActivation to kernel activation kind integer.
170
171
Returns:
@@ -200,29 +245,22 @@ class CPUExpertsInt4(mk.FusedMoEExpertsMonolithic):
245
e_score_correction_bias=e_score_correction_bias,
246
)
247
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
248
+ hidden_size = self.moe_config.hidden_dim
249
+ intermediate_size = self.moe_config.intermediate_size_per_partition
250
assert self.quant_config.block_shape is not None
211
- if self.quant_config.is_per_act_token:
251
+ # C++ kernel expects an int: -1 for channelwise, and group size for groupwise
252
+ if self.quant_config.block_shape == [-1, 1]:
253
group_size = -1
254
else:
255
group_size = self.quant_config.block_shape[1]
215
-
216
- # Call the dynamic 4-bit int MoE kernel
256
return torch.ops._C.dynamic_4bit_int_moe(
257
hidden_states,
258
topk_ids.to(torch.long),
259
topk_weights,
260
w1, # w13_weight_packed
261
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)
262
+ hidden_size,
263
+ intermediate_size,
264
group_size,
265
apply_router_weight_on_input,
266
self._activation_kind(activation),
vllm/model_executor/layers/fused_moe/oracle/w4a8_int8.py
+2
-2
@@ -270,7 +270,7 @@ def convert_to_w4a8_int8_moe_format(
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)
273
+ w13_out_features = w13_weight.shape[1] # 2 * intermediate_size
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
@@ -286,7 +286,7 @@ def convert_to_w4a8_int8_moe_format(
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,
289
+ w13_out_features,
290
group_size,
291
)
292
)
vllm/model_executor/layers/quantization/compressed_tensors/compressed_tensors_moe/compressed_tensors_moe_w4a8_int8.py
+7
-24
@@ -5,7 +5,6 @@
5
import torch
6
from compressed_tensors.quantization import (
7
QuantizationArgs,
8
- QuantizationStrategy,
8
)
9
10
from vllm.logger import init_logger
@@ -32,7 +31,6 @@ from vllm.model_executor.layers.quantization.utils.quant_utils import (
31
ScaleDesc,
32
)
33
from vllm.model_executor.utils import replace_parameter, set_weight_attrs
35
-from vllm.platforms import CpuArchEnum, current_platform
34
35
logger = init_logger(__name__)
36
@@ -64,32 +62,17 @@ class CompressedTensorsW4A8Int8MoEMethod(CompressedTensorsMoEMethod):
62
weight_quant.group_size if (weight_quant.group_size is not None) else -1
63
)
64
67
- # Validate scheme: weights=W4 (channel or group),
68
- # activations=dynamic TOKEN (A8)
69
-
70
- # Must be dynamic per-token activations
71
- if (
72
- input_quant.strategy != QuantizationStrategy.TOKEN
73
- or not input_quant.dynamic
65
+ # make sure group size is valid
66
+ if self.group_size != -1 and (
67
+ moe.hidden_dim % self.group_size != 0
68
+ or moe.intermediate_size_per_partition % self.group_size != 0
69
):
70
raise ValueError(
76
- "W4A8-int MoE needs dynamic per-token activation quantization."
71
+ f"Group size ({self.group_size}) must evenly divide both "
72
+ f"hidden size ({moe.hidden_dim}) and intermediate size per "
73
+ f"partition ({moe.intermediate_size_per_partition})."
74
)
75
79
- if weight_quant.num_bits != 4:
80
- raise ValueError("This method only supports 4-bit weights (num_bits=4).")
81
-
82
- # Arm: check _dyn ops availability
83
- if current_platform.get_cpu_architecture() == CpuArchEnum.ARM:
84
- try:
85
- _ = torch.ops.aten._dyn_quant_matmul_4bit
86
- _ = torch.ops.aten._dyn_quant_pack_4bit_weight
87
- except AttributeError as err:
88
- raise RuntimeError(
89
- f"""PyTorch {torch.__version__} lacks _dyn_quant_* 4bit ops;
90
- install a newer build."""
91
- ) from err
92
-
76
# Construct QuantKey for weights from QuantizationArgs
77
# W4A8 INT4: 4-bit weights (stored as int8), static quantization
78
if self.group_size == -1:
vllm/model_executor/layers/quantization/utils/quant_utils.py
+6
@@ -216,6 +216,12 @@ kInt4W4A8StaticGroupSym = QuantKey(
216
torch.int8, kInt4W4A8StaticGroupScale, symmetric=True
217
)
218
219
+kInt4W4A8StaticChannelSym = QuantKey(
220
+ torch.int8,
221
+ ScaleDesc(torch.float32, True, GroupShape.PER_CHANNEL),
222
+ symmetric=True,
223
+)
224
+
225
226
def create_fp8_quant_key(
227
static: bool,