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 12, 06:00 UTC
[CPU]Parallelize over tokens in int4 moe (#29600)
Signed-off-by: Zhang Xiangze <Xiangze.Zhang@arm.com>
Zhang Xiangze committed
Dec 2, 2025 at 14:21 UTC
13ea39bc09cf4c102ba4ad308df379dc5abc3ba4
1 file changed
+6
-6
csrc/moe/dynamic_4bit_int_moe_cpu.cpp
+6
-6
@@ -93,16 +93,16 @@ torch::Tensor dynamic_4bit_int_moe_cpu(
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) {
96
+ at::parallel_for(0, offsets[E], 0, [&](int64_t idx_begin, int64_t idx_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
+ for (int64_t e = 0; e < E; ++e) {
99
+ int64_t start = std::max(offsets[e], idx_begin);
100
+ int64_t end = std::min(offsets[e + 1], idx_end);
101
+ int64_t te = end - start;
102
+ if (te <= 0) {
103
continue;
104
}
105
104
- const int64_t start = offsets[e];
105
-
106
auto x_e = X_all.narrow(/*dim=*/0, /*start=*/start, /*length=*/te);
107
108
auto w13_e = w13_packed.select(/*dim=*/0, e);