A high-throughput and memory-efficient inference and serving engine for LLMs
Mirrored from https://github.com/tduncan-lum/Nvidia-Optimization-Stack---VLLM.git
Synced Nov 14, 05:30 UTC
[fix]: add Arm 4bit fused moe support (#23809)
Signed-off-by: Nikhil Gupta <nikhil.gupta2@arm.com>
Nikhil Gupta committed
Sep 24, 2025 at 02:32 UTC
359d2930063038c78b9a54a1b0176fbabb1c204b
7 files changed
+488
-11
cmake/cpu_extension.cmake
+2
-1
@@ -258,7 +258,8 @@ set(VLLM_EXT_SRC
258
"csrc/cpu/layernorm.cpp"
259
"csrc/cpu/mla_decode.cpp"
260
"csrc/cpu/pos_encoding.cpp"
261
- "csrc/cpu/torch_bindings.cpp")
261
+ "csrc/cpu/torch_bindings.cpp"
262
+ "csrc/moe/dynamic_4bit_int_moe_cpu.cpp")
263
264
if (AVX512_FOUND AND NOT AVX512_DISABLED)
265
set(VLLM_EXT_SRC
csrc/cpu/torch_bindings.cpp
+10
@@ -88,8 +88,18 @@ TORCH_LIBRARY_EXPAND(TORCH_EXTENSION_NAME, ops) {
88
" int tp_rank, int blocksparse_local_blocks,"
89
" int blocksparse_vert_stride, int blocksparse_block_size,"
90
" int blocksparse_head_sliding_step) -> ()");
91
+
92
ops.impl("paged_attention_v1", torch::kCPU, &paged_attention_v1);
93
94
+ ops.def(
95
+ "dynamic_4bit_int_moe("
96
+ "Tensor x, Tensor topk_ids, Tensor topk_weights,"
97
+ "Tensor w13_packed, Tensor w2_packed, int H, int I, int I2,"
98
+ "int group_size, bool apply_router_weight_on_input, int activation_kind"
99
+ ") -> Tensor");
100
+
101
+ ops.impl("dynamic_4bit_int_moe", torch::kCPU, &dynamic_4bit_int_moe_cpu);
102
+
103
// PagedAttention V2.
104
ops.def(
105
"paged_attention_v2("
csrc/moe/dynamic_4bit_int_moe_cpu.cpp
new
+156
@@ -0,0 +1,156 @@
1
+#include <ATen/ATen.h>
2
+#include <ATen/Parallel.h>
3
+#include <torch/all.h>
4
+
5
+// _dyn_quant_matmul_4bit is only available on AArch64.
6
+#if defined(__aarch64__)
7
+ #include <ATen/ops/_dyn_quant_matmul_4bit.h>
8
+#endif
9
+
10
+inline torch::Tensor mm(const torch::Tensor& a, const torch::Tensor& packed_w,
11
+ int64_t group_size_eff, int64_t in_features,
12
+ int64_t out_features) {
13
+#if defined(__aarch64__)
14
+ return at::_ops::_dyn_quant_matmul_4bit::call(a, packed_w, group_size_eff,
15
+ in_features, out_features);
16
+#else
17
+ TORCH_CHECK(false,
18
+ "dynamic 4-bit int MoE path requires AArch64 (ARM64); "
19
+ "_dyn_quant_matmul_4bit is unavailable on this architecture");
20
+ return {};
21
+#endif
22
+}
23
+
24
+enum ActivationKind : int64_t {
25
+ SwiGLU_Gu = 0, // act = SiLU(g) * u
26
+ SwiGLUOAI = 1, // act = SiLU(u) * g
27
+ SiLU = 2 // SiLU
28
+};
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) {
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");
42
+
43
+ const int64_t T = x.size(0);
44
+ const int64_t K = topk_ids.size(1);
45
+ const int64_t E = w13_packed.size(0);
46
+ const int64_t N = T * K;
47
+
48
+ auto x_c = x.contiguous();
49
+ auto ids_c = topk_ids.contiguous();
50
+ auto gates_c = topk_weights.to(at::kFloat).contiguous();
51
+
52
+ // bucketing tokens -> experts
53
+ c10::SmallVector<int64_t, 64> counts(
54
+ E, 0); // Small vector uses stack allocation
55
+ {
56
+ const auto* ids_ptr = ids_c.data_ptr<int64_t>();
57
+ for (int64_t i = 0; i < N; ++i) {
58
+ const int64_t e_id = ids_ptr[i];
59
+ TORCH_CHECK(0 <= e_id && e_id < E, "expert id out of range");
60
+ counts[e_id]++;
61
+ }
62
+ }
63
+ c10::SmallVector<int64_t, 65> offsets(E + 1, 0); // ( E +1 )
64
+ for (int64_t e = 0; e < E; ++e) offsets[e + 1] = offsets[e] + counts[e];
65
+
66
+ auto expert_tokens = at::empty({offsets[E]}, ids_c.options());
67
+ auto expert_gates = at::empty({offsets[E]}, gates_c.options());
68
+ {
69
+ 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
+ }
85
+ }
86
+
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;
89
+
90
+ // Per-expert outputs filled in parallel
91
+ std::vector<torch::Tensor> y_list(E);
92
+ y_list.resize(E);
93
+
94
+ at::parallel_for(0, E, 1, [&](int64_t e_begin, int64_t e_end) {
95
+ for (int64_t e = e_begin; e < e_end; ++e) {
96
+ const int64_t te = counts[e];
97
+ if (te == 0) {
98
+ y_list[e] = at::empty({0, H}, x_c.options());
99
+ continue;
100
+ }
101
+
102
+ const int64_t start = offsets[e];
103
+
104
+ auto sel_tokens =
105
+ expert_tokens.narrow(/*dim=*/0, /*start=*/start, /*length=*/te);
106
+ auto gates_e =
107
+ expert_gates.narrow(/*dim=*/0, /*start=*/start, /*length=*/te);
108
+
109
+ auto x_e = x_c.index_select(/*dim=*/0, sel_tokens);
110
+
111
+ if (apply_router_weight_on_input) {
112
+ x_e = x_e.mul(gates_e.unsqueeze(1));
113
+ }
114
+
115
+ auto w13_e = w13_packed.select(/*dim=*/0, e);
116
+ auto w2_e = w2_packed.select(/*dim=*/0, e);
117
+
118
+ // W13
119
+ auto y13 =
120
+ mm(x_e, w13_e, g_eff_13, /*in_features=*/H, /*out_features=*/I2);
121
+
122
+ auto g_part = y13.narrow(/*dim=*/1, /*start=*/0, /*length=*/I);
123
+ auto u_part = y13.narrow(/*dim=*/1, /*start=*/I, /*length=*/I);
124
+
125
+ torch::Tensor act;
126
+ if (activation_kind == ActivationKind::SwiGLUOAI) { // SwiGLUOAI
127
+ constexpr double kAlpha = 1.702; // GPT-OSS default
128
+ constexpr double kLimit = 7.0; // GPT-OSS default
129
+ auto gate_c = at::clamp_max(g_part, kLimit);
130
+ auto up_c = at::clamp(u_part, -kLimit, kLimit);
131
+ auto glu = gate_c.mul(at::sigmoid(gate_c.mul(kAlpha)));
132
+ act = up_c.add(1.0).mul(glu);
133
+ } else { // SiLU , SwiGLU_GU, vLLM maps silu to SiluAndMul()
134
+ act = at::silu(g_part).mul(u_part);
135
+ }
136
+
137
+ // W2
138
+ auto y = mm(act, w2_e, g_eff_2, /*in_features=*/I, /*out_features=*/H);
139
+
140
+ if (!apply_router_weight_on_input) {
141
+ y = y.mul(gates_e.unsqueeze(1));
142
+ }
143
+
144
+ // Store per-expert result
145
+ y_list[e] = y;
146
+ }
147
+ });
148
+
149
+ // Concatenate all expert outputs to match expert_tokens order
150
+ auto Y_all = at::cat(y_list, /*dim=*/0);
151
+ auto out = at::zeros({T, H}, x.options());
152
+ out =
153
+ at::index_add(out, /*dim=*/0, /*index=*/expert_tokens, /*source=*/Y_all);
154
+
155
+ return out;
156
+}
csrc/ops.h
+6
@@ -328,6 +328,12 @@ void selective_scan_fwd(const torch::Tensor& u, const torch::Tensor& delta,
328
const std::optional<torch::Tensor>& has_initial_state,
329
const torch::Tensor& ssm_states, int64_t pad_slot_id);
330
331
+torch::Tensor dynamic_4bit_int_moe_cpu(
332
+ torch::Tensor x, torch::Tensor topk_ids, torch::Tensor topk_weights,
333
+ torch::Tensor w13_packed, torch::Tensor w2_packed, int64_t H, int64_t I,
334
+ int64_t I2, int64_t group_size, bool apply_router_weight_on_input,
335
+ int64_t activation_kind);
336
+
337
using fptr_t = int64_t;
338
fptr_t init_custom_ar(const std::vector<int64_t>& fake_ipc_ptrs,
339
torch::Tensor& rank_data, int64_t rank,
vllm/model_executor/layers/fused_moe/cpu_fused_moe.py
+9
-6
@@ -98,13 +98,16 @@ def select_experts(
98
e_score_correction_bias=e_score_correction_bias)
99
elif custom_routing_function is None:
100
assert scoring_func == "softmax"
101
- topk_weights = torch.nn.functional.softmax(router_logits,
102
- dim=1,
103
- dtype=torch.float32)
104
- topk_weights, topk_ids = torch.topk(topk_weights, top_k, dim=-1)
101
+ topk_logit_vals, topk_idx = torch.topk(router_logits,
102
+ k=top_k,
103
+ dim=-1,
104
+ sorted=False)
105
if renormalize:
106
- topk_weights /= topk_weights.sum(dim=-1, keepdim=True)
107
- return topk_weights, topk_ids.to(torch.int32)
106
+ topk_vals = torch.softmax(topk_logit_vals, dim=-1)
107
+ else:
108
+ logZ = torch.logsumexp(router_logits, dim=-1, keepdim=True)
109
+ topk_vals = (topk_logit_vals - logZ).exp()
110
+ return topk_vals.to(torch.float32), topk_idx.to(torch.int32)
111
else:
112
return custom_routing_function(hidden_states=hidden_states,
113
gating_output=router_logits,
vllm/model_executor/layers/fused_moe/layer.py
-2
@@ -69,8 +69,6 @@ else:
69
if is_rocm_aiter_moe_enabled():
70
from vllm.model_executor.layers.fused_moe.rocm_aiter_fused_moe import ( # noqa: E501
71
rocm_aiter_grouped_topk as grouped_topk)
72
-elif current_platform.is_cpu():
73
- pass
72
else:
73
from vllm.model_executor.layers.fused_moe.fused_moe import grouped_topk
74
if current_platform.is_tpu():
vllm/model_executor/layers/quantization/compressed_tensors/compressed_tensors_moe.py
+305
-2
@@ -22,6 +22,7 @@ from vllm.model_executor.layers.fused_moe.config import (
22
FusedMoEQuantConfig, fp8_w8a8_moe_quant_config,
23
int4_w4a16_moe_quant_config, int8_w8a8_moe_quant_config,
24
int8_w8a16_moe_quant_config, nvfp4_moe_quant_config)
25
+from vllm.model_executor.layers.fused_moe.cpu_fused_moe import select_experts
26
from vllm.model_executor.layers.fused_moe.flashinfer_cutlass_moe import (
27
is_valid_flashinfer_cutlass_fused_moe)
28
from vllm.model_executor.layers.quantization.compressed_tensors.schemes.compressed_tensors_wNa16 import ( # noqa
@@ -47,7 +48,7 @@ from vllm.model_executor.layers.quantization.utils.quant_utils import (
48
from vllm.model_executor.layers.quantization.utils.w8a8_utils import (
49
all_close_1d, normalize_e4m3fn_to_e4m3fnuz, per_tensor_dequantize)
50
from vllm.model_executor.utils import set_weight_attrs
50
-from vllm.platforms import current_platform
51
+from vllm.platforms import CpuArchEnum, current_platform
52
from vllm.scalar_type import scalar_types
53
from vllm.utils.deep_gemm import is_deep_gemm_e8m0_used
54
@@ -63,7 +64,7 @@ __all__ = [
64
"CompressedTensorsMoEMethod", "CompressedTensorsW8A8Fp8MoEMethod",
65
"CompressedTensorsW8A8Int8MoEMethod",
66
"CompressedTensorsWNA16MarlinMoEMethod", "CompressedTensorsWNA16MoEMethod",
66
- "CompressedTensorsW4A4MoeMethod"
67
+ "CompressedTensorsW4A4MoeMethod", "CompressedTensorsW4A8Int8MoEMethod"
68
]
69
70
@@ -139,6 +140,10 @@ class CompressedTensorsMoEMethod(FusedMoEMethodBase):
140
elif quant_config._is_dynamic_token_w8a8(weight_quant, input_quant):
141
return CompressedTensorsW8A8Int8MoEMethod(quant_config,
142
layer.moe_config)
143
+ elif quant_config._is_dynamic_token_w4a8_int(weight_quant,
144
+ input_quant):
145
+ return CompressedTensorsW4A8Int8MoEMethod(quant_config,
146
+ layer.moe_config)
147
else:
148
raise RuntimeError(
149
f"Unsupported FusedMoe scheme: {weight_quant}, {input_quant}")
@@ -1769,3 +1774,301 @@ class CompressedTensorsWNA16MoEMethod(CompressedTensorsMoEMethod):
1774
expert_map=expert_map,
1775
quant_config=self.moe_quant_config,
1776
)
1777
+
1778
+
1779
+class CompressedTensorsW4A8Int8MoEMethod(CompressedTensorsMoEMethod):
1780
+ """
1781
+ CPU-only MoE method using dynamic 4-bit matmul kernels on Arm Platform
1782
+ - Weights: int4 (stored as int8 values in [-8,7], packed to uint8 nibbles)
1783
+ - Scales: Fp32 for Channelwise , bf16 for groupwise quantization
1784
+ - Bias: Same data type as original weights
1785
+ - Activations: FP32/Bf16 dynamic per-token (A8 Int),
1786
+ quantized inside the kernel
1787
+ """
1788
+
1789
+ def __init__(
1790
+ self,
1791
+ quant_config: "CompressedTensorsConfig", # type: ignore # noqa E501
1792
+ moe: FusedMoEConfig):
1793
+ super().__init__(moe)
1794
+ self.has_bias = self.moe.has_bias
1795
+ self.quant_config = quant_config
1796
+
1797
+ # Validate scheme: weights=W4 (channel or group),
1798
+ # activations=dynamic TOKEN (A8)
1799
+ wq = self.quant_config.target_scheme_map["Linear"].get("weights")
1800
+ aq = self.quant_config.target_scheme_map["Linear"].get(
1801
+ "input_activations")
1802
+
1803
+ # Must be dynamic per-token activations
1804
+ if aq.strategy != QuantizationStrategy.TOKEN or not aq.dynamic:
1805
+ raise ValueError(
1806
+ "W4A8-int MoE needs dynamic per-token activation quantization."
1807
+ )
1808
+
1809
+ # Weight can be channel-wise (group_size=None) or group-wise
1810
+ self.group_size = wq.group_size if (wq.group_size is not None) else -1
1811
+ if wq.num_bits != 4:
1812
+ raise ValueError(
1813
+ "This method only supports 4-bit weights (num_bits=4).")
1814
+
1815
+ # CPU only
1816
+ if not current_platform.is_cpu():
1817
+ raise ValueError("CompressedTensorsW4A8Int8MoEMethod is CPU-only.")
1818
+
1819
+ # Arm: check _dyn ops availability
1820
+ if current_platform.get_cpu_architecture() == CpuArchEnum.ARM:
1821
+ try:
1822
+ _ = torch.ops.aten._dyn_quant_matmul_4bit
1823
+ _ = torch.ops.aten._dyn_quant_pack_4bit_weight
1824
+ except AttributeError as err:
1825
+ raise RuntimeError(
1826
+ f"""PyTorch {torch.__version__} lacks _dyn_quant_* 4bit ops;
1827
+ install a newer build.""") from err
1828
+ self.static_input_scales = False # always dynamic per token
1829
+
1830
+ # ---- parameter creation ----
1831
+ def create_weights(self, layer: torch.nn.Module, num_experts: int,
1832
+ hidden_size: int, intermediate_size_per_partition: int,
1833
+ params_dtype: torch.dtype, **extra_weight_attrs):
1834
+
1835
+ # Shapes per local rank (TP/EP):
1836
+ # w13: [E, 2*I_local, H] int8 (int4 values in [-8,7])
1837
+ # w2 : [E, H, I_local] int8
1838
+ # Scales:
1839
+ # channel-wise: group_size=-1 -> per-output-row, single scale per row
1840
+ # group-wise : group_size=g ->
1841
+ # per-output-row, (in_features/g) scales
1842
+
1843
+ E = num_experts
1844
+ H = hidden_size
1845
+ IN = intermediate_size_per_partition
1846
+ g = self.group_size
1847
+
1848
+ # Per-row scale columns
1849
+ def _n_scale_cols(in_features: int) -> int:
1850
+ return 1 if g == -1 else (in_features // g)
1851
+
1852
+ # Register unpacked int4-as-int8 weights the loader will fill.
1853
+ w13 = torch.nn.Parameter(torch.empty(E, 2 * IN, H, dtype=torch.int8),
1854
+ requires_grad=False)
1855
+ set_weight_attrs(w13, extra_weight_attrs)
1856
+ layer.register_parameter("w13_weight", w13)
1857
+
1858
+ w2 = torch.nn.Parameter(torch.empty(E, H, IN, dtype=torch.int8),
1859
+ requires_grad=False)
1860
+ set_weight_attrs(w2, extra_weight_attrs)
1861
+ layer.register_parameter("w2_weight", w2)
1862
+
1863
+ # Register scales
1864
+ # KleidiAI groupwise kernels accepts float32 scales
1865
+ # KleidiAI groupwise kernels accepts bfloat16 scales
1866
+ scale_dtype = torch.float32 if g == -1 else torch.bfloat16
1867
+
1868
+ w13_s = torch.nn.Parameter(torch.ones(E,
1869
+ 2 * IN,
1870
+ _n_scale_cols(H),
1871
+ dtype=scale_dtype),
1872
+ requires_grad=False)
1873
+ set_weight_attrs(
1874
+ w13_s, {
1875
+ "quant_method": "channel" if g == -1 else "group",
1876
+ **extra_weight_attrs
1877
+ })
1878
+ layer.register_parameter("w13_weight_scale", w13_s)
1879
+
1880
+ w2_s = torch.nn.Parameter(torch.ones(E,
1881
+ H,
1882
+ _n_scale_cols(IN),
1883
+ dtype=scale_dtype),
1884
+ requires_grad=False)
1885
+ set_weight_attrs(
1886
+ w2_s, {
1887
+ "quant_method": "channel" if g == -1 else "group",
1888
+ **extra_weight_attrs
1889
+ })
1890
+ layer.register_parameter("w2_weight_scale", w2_s)
1891
+
1892
+ if self.has_bias:
1893
+ w13_bias = torch.nn.Parameter(torch.zeros(E,
1894
+ 2 * IN,
1895
+ dtype=params_dtype),
1896
+ requires_grad=False)
1897
+ layer.register_parameter("w13_bias", w13_bias)
1898
+ set_weight_attrs(w13_bias, extra_weight_attrs)
1899
+
1900
+ w2_bias = torch.nn.Parameter(torch.zeros(num_experts,
1901
+ hidden_size,
1902
+ dtype=params_dtype),
1903
+ requires_grad=False)
1904
+ layer.register_parameter("w2_bias", w2_bias)
1905
+ set_weight_attrs(w2_bias, extra_weight_attrs)
1906
+
1907
+ # Placeholders for packed weights (will be replaced after packing)
1908
+ layer.register_parameter(
1909
+ "w13_weight_packed",
1910
+ torch.nn.Parameter(torch.empty(0), requires_grad=False))
1911
+ set_weight_attrs(layer.w13_weight_packed, extra_weight_attrs)
1912
+
1913
+ layer.register_parameter(
1914
+ "w2_weight_packed",
1915
+ torch.nn.Parameter(torch.empty(0), requires_grad=False))
1916
+ set_weight_attrs(layer.w2_weight_packed, extra_weight_attrs)
1917
+
1918
+ # dims for 4 bit fused matmuls
1919
+ layer.w13_in_features = H
1920
+ layer.w13_out_features = 2 * IN
1921
+ layer.w2_in_features = IN
1922
+ layer.w2_out_features = H
1923
+ layer.group_size = g
1924
+
1925
+ # post-load packing to dyn-4bit KleidiAI kernel's format
1926
+ def process_weights_after_loading(self, layer: torch.nn.Module) -> None:
1927
+ E = layer.w13_weight.shape[0]
1928
+ H = layer.w13_in_features
1929
+ I2 = layer.w13_out_features
1930
+ IN = layer.w2_in_features
1931
+ g = layer.group_size
1932
+
1933
+ def _pack_matrix(int4_as_int8_2d: torch.Tensor,
1934
+ scales_2d: torch.Tensor,
1935
+ bias_1d: Optional[torch.Tensor], in_features: int,
1936
+ out_features: int) -> torch.Tensor:
1937
+ # int4 values are stored as int8 in [-8,7].
1938
+ # Shift to unsigned nibble and pack pairs along input-dim.
1939
+ tmp = int4_as_int8_2d.add(8) # [out, in]
1940
+ uint8_nibbles = ((tmp[:, 1::2] << 4) | tmp[:, ::2]).to(
1941
+ torch.uint8) # [out, in//2]
1942
+
1943
+ # KleidiAI groupwise kernels accepts float32 scales
1944
+ # KleidiAI groupwise kernels accepts bfloat16 scales
1945
+ scale_dtype = torch.float32 if g == -1 else torch.bfloat16
1946
+ scales = scales_2d.to(scale_dtype)
1947
+ bias = None if bias_1d is None else bias_1d.to(torch.float32)
1948
+ return torch.ops.aten._dyn_quant_pack_4bit_weight(
1949
+ uint8_nibbles, scales, bias, g if g != -1 else in_features,
1950
+ in_features, out_features)
1951
+
1952
+ # Pack per expert
1953
+ w13_packed_list = []
1954
+ w2_packed_list = []
1955
+
1956
+ has_w13_bias = hasattr(layer,
1957
+ "w13_bias") and layer.w13_bias is not None
1958
+ has_w2_bias = hasattr(layer, "w2_bias") and layer.w2_bias is not None
1959
+
1960
+ for e in range(E):
1961
+ w13_packed_list.append(
1962
+ _pack_matrix(
1963
+ layer.w13_weight[e], # [2I, H]
1964
+ layer.w13_weight_scale[e], # [2I, H/g or 1]
1965
+ layer.w13_bias[e] if has_w13_bias else None, # [2I]
1966
+ H,
1967
+ I2))
1968
+ w2_packed_list.append(
1969
+ _pack_matrix(
1970
+ # w2 shape is [H, IN]; we need [out, in] == [H, IN].
1971
+ layer.w2_weight[e], # [H, IN]
1972
+ layer.w2_weight_scale[e], # [H, IN/g or 1]
1973
+ layer.w2_bias[e] if has_w2_bias else None, # [H]
1974
+ IN,
1975
+ layer.w2_out_features # in_features=IN, out_features=H
1976
+ ))
1977
+
1978
+ # each packed tensor has identical shape per expert; stack on dim 0
1979
+ w13_packed = torch.stack(w13_packed_list, dim=0)
1980
+ w2_packed = torch.stack(w2_packed_list, dim=0)
1981
+
1982
+ replace_parameter(layer, "w13_weight_packed",
1983
+ torch.nn.Parameter(w13_packed, requires_grad=False))
1984
+ replace_parameter(layer, "w2_weight_packed",
1985
+ torch.nn.Parameter(w2_packed, requires_grad=False))
1986
+
1987
+ # free raw tensors/scales/bias now that they're packed into the payload.
1988
+ replace_parameter(
1989
+ layer, "w13_weight",
1990
+ torch.nn.Parameter(torch.empty(0), requires_grad=False))
1991
+ replace_parameter(
1992
+ layer, "w2_weight",
1993
+ torch.nn.Parameter(torch.empty(0), requires_grad=False))
1994
+ replace_parameter(
1995
+ layer, "w13_weight_scale",
1996
+ torch.nn.Parameter(torch.empty(0), requires_grad=False))
1997
+ replace_parameter(
1998
+ layer, "w2_weight_scale",
1999
+ torch.nn.Parameter(torch.empty(0), requires_grad=False))
2000
+ if has_w13_bias:
2001
+ replace_parameter(
2002
+ layer, "w13_bias",
2003
+ torch.nn.Parameter(torch.empty(0), requires_grad=False))
2004
+ if has_w2_bias:
2005
+ replace_parameter(
2006
+ layer, "w2_bias",
2007
+ torch.nn.Parameter(torch.empty(0), requires_grad=False))
2008
+
2009
+ def get_fused_moe_quant_config(
2010
+ self, layer: torch.nn.Module) -> Optional[FusedMoEQuantConfig]:
2011
+ # CPU dynamic 4-bit MoE path does not use modular kernels or
2012
+ # fused_experts; quant config is not needed.
2013
+ return None
2014
+
2015
+ def apply(
2016
+ self,
2017
+ layer: torch.nn.Module,
2018
+ x: torch.Tensor,
2019
+ router_logits: torch.Tensor,
2020
+ top_k: int,
2021
+ renormalize: bool,
2022
+ use_grouped_topk: bool = False,
2023
+ topk_group: Optional[int] = None,
2024
+ num_expert_group: Optional[int] = None,
2025
+ global_num_experts: int = -1,
2026
+ expert_map: Optional[torch.Tensor] = None,
2027
+ custom_routing_function: Optional[Callable] = None,
2028
+ scoring_func: str = "softmax",
2029
+ routed_scaling_factor: float = 1.0,
2030
+ e_score_correction_bias: Optional[torch.Tensor] = None,
2031
+ apply_router_weight_on_input: bool = False,
2032
+ activation: str = "silu",
2033
+ enable_eplb: bool = False,
2034
+ expert_load_view: Optional[torch.Tensor] = None,
2035
+ logical_to_physical_map: Optional[torch.Tensor] = None,
2036
+ logical_replica_count: Optional[torch.Tensor] = None,
2037
+ ) -> torch.Tensor:
2038
+ assert not enable_eplb, "EPLB not supported for W4A8-int MoE yet."
2039
+ assert activation in (
2040
+ "silu", "swigluoai",
2041
+ "swiglu"), "Only SiLU/SwiGLUGU/SwiGLUUG are supported."
2042
+ assert expert_map is None, """expert_map/EP not implemented
2043
+ for CPU dyn-4bit MoE."""
2044
+
2045
+ def _act_kind(s: str) -> int:
2046
+ # 0 = SwiGLU_Gu (SiLU(g)*u), 1 = SwiGLU_Ug (SiLU(u)*g), 2 = SiLU
2047
+ if s == "swiglu":
2048
+ return 0
2049
+ if s == "swigluoai":
2050
+ return 1
2051
+ if s == "silu":
2052
+ return 2
2053
+ raise ValueError(f"Unknown activation '{s}'")
2054
+
2055
+ # Apply topk softmax on router output
2056
+ topk_weights, topk_ids = select_experts(
2057
+ hidden_states=x,
2058
+ router_logits=router_logits,
2059
+ use_grouped_topk=use_grouped_topk,
2060
+ top_k=top_k,
2061
+ renormalize=renormalize,
2062
+ topk_group=topk_group,
2063
+ num_expert_group=num_expert_group,
2064
+ custom_routing_function=custom_routing_function,
2065
+ scoring_func=scoring_func,
2066
+ routed_scaling_factor=routed_scaling_factor,
2067
+ e_score_correction_bias=e_score_correction_bias,
2068
+ )
2069
+
2070
+ return torch.ops._C.dynamic_4bit_int_moe(
2071
+ x, topk_ids.to(torch.long), topk_weights, layer.w13_weight_packed,
2072
+ layer.w2_weight_packed, layer.w2_out_features,
2073
+ layer.w2_in_features, layer.w13_out_features, layer.group_size,
2074
+ apply_router_weight_on_input, int(_act_kind(activation)))
\ No newline at end of file