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 13, 20:45 UTC
[CPU]Improve dynamic 4bit moe performance (#27240)
Signed-off-by: Zhang Xiangze <Xiangze.Zhang@arm.com>
xiangze-arm committed
Nov 4, 2025 at 14:33 UTC
f32cbc9a0c456966df300076a3a9f2889151b024
1 file changed
+12
-21
csrc/moe/dynamic_4bit_int_moe_cpu.cpp
+12
-21
@@ -87,30 +87,23 @@ torch::Tensor dynamic_4bit_int_moe_cpu(
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);
90
+ auto X_all = x_c.index_select(/*dim=*/0, expert_tokens);
91
+ if (apply_router_weight_on_input) {
92
+ X_all = X_all.mul(expert_gates.unsqueeze(1));
93
+ }
94
+ auto Y_all = at::empty({offsets[E], H}, x_c.options());
95
96
at::parallel_for(0, E, 1, [&](int64_t e_begin, int64_t e_end) {
97
+ c10::InferenceMode guard;
98
for (int64_t e = e_begin; e < e_end; ++e) {
99
const int64_t te = counts[e];
100
if (te == 0) {
98
- y_list[e] = at::empty({0, H}, x_c.options());
101
continue;
102
}
103
104
const int64_t start = offsets[e];
105
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
- }
106
+ auto x_e = X_all.narrow(/*dim=*/0, /*start=*/start, /*length=*/te);
107
108
auto w13_e = w13_packed.select(/*dim=*/0, e);
109
auto w2_e = w2_packed.select(/*dim=*/0, e);
@@ -137,17 +130,15 @@ torch::Tensor dynamic_4bit_int_moe_cpu(
130
// W2
131
auto y = mm(act, w2_e, g_eff_2, /*in_features=*/I, /*out_features=*/H);
132
140
- if (!apply_router_weight_on_input) {
141
- y = y.mul(gates_e.unsqueeze(1));
142
- }
143
-
133
// Store per-expert result
145
- y_list[e] = y;
134
+ Y_all.narrow(/*dim=*/0, /*start=*/start, /*length=*/te).copy_(y);
135
}
136
});
137
149
- // Concatenate all expert outputs to match expert_tokens order
150
- auto Y_all = at::cat(y_list, /*dim=*/0);
138
+ if (!apply_router_weight_on_input) {
139
+ Y_all = Y_all.mul(expert_gates.unsqueeze(1));
140
+ }
141
+
142
auto out = at::zeros({T, H}, x.options());
143
out =
144
at::index_add(out, /*dim=*/0, /*index=*/expert_tokens, /*source=*/Y_all);