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 8, 11:15 UTC
[Bugfix][CPU] Fix llama4 inference on CPU (#34321)
Signed-off-by: jiang1.li <jiang1.li@intel.com>
Li, Jiang committed
Feb 11, 2026 at 19:07 UTC
05339a7b207e2f32b56c29398c18d577c74cef3b
6 files changed
+60
-18
.gitignore
+3
@@ -238,3 +238,6 @@ ep_kernels_workspace/
238
vllm/grpc/vllm_engine_pb2.py
239
vllm/grpc/vllm_engine_pb2_grpc.py
240
vllm/grpc/vllm_engine_pb2.pyi
241
+
242
+# Ignore generated cpu headers
243
+csrc/cpu/cpu_attn_dispatch_generated.h
csrc/cpu/cpu_fused_moe.cpp
+10
-3
@@ -147,7 +147,7 @@ void fused_moe_impl(scalar_t* __restrict__ output, scalar_t* __restrict__ input,
147
const int32_t token_num, const int32_t expert_num,
148
const int32_t topk_num, const int32_t input_size_13,
149
const int32_t output_size_13, const int32_t input_size_2,
150
- const int32_t output_size_2) {
150
+ const int32_t output_size_2, const bool skip_weighted) {
151
using scalar_vec_t = typename cpu_utils::VecTypeTrait<scalar_t>::vec_t;
152
constexpr int32_t gemm_n_tile_size = gemm_t::NSize;
153
constexpr int32_t gemm_m_tile_size = gemm_t::MaxMSize;
@@ -582,6 +582,11 @@ void fused_moe_impl(scalar_t* __restrict__ output, scalar_t* __restrict__ input,
582
scalar_t* __restrict__ curr_output_buffer =
583
output + token_id * output_size_2;
584
585
+ if (skip_weighted) {
586
+ // Only for topk_num == 1
587
+ *curr_weight = 1.0f;
588
+ }
589
+
590
if (topk_num > 1) {
591
{
592
int32_t w2_output_idx = curr_expand_token_id_index_buffer[0];
@@ -699,7 +704,7 @@ void cpu_fused_moe(
704
const std::optional<torch::Tensor>& w2_bias, // [expert_num, output_size_2]
705
const torch::Tensor& topk_weights, // [token_num, k], float32
706
const torch::Tensor& topk_id, // [token_num, k], int32
702
- const std::string& act, const std::string& isa) {
707
+ const bool skip_weighted, const std::string& act, const std::string& isa) {
708
const int32_t token_num = input.size(0);
709
const int32_t input_size_13 = input.size(1);
710
const int64_t input_stride = input.stride(0);
@@ -711,6 +716,8 @@ void cpu_fused_moe(
716
const int32_t topk_num = topk_id.size(1);
717
const FusedMOEAct act_type = get_act_type(act);
718
cpu_utils::ISA isa_type = cpu_utils::get_isa(isa);
719
+ TORCH_CHECK(!skip_weighted || topk_num == 1,
720
+ "skip_weighted is only supported for topk=1 on CPU");
721
722
VLLM_DISPATCH_FLOATING_TYPES(w13.scalar_type(), "cpu_fused_moe", [&]() {
723
CPU_ISA_DISPATCH_IMPL(isa_type, [&]() {
@@ -721,7 +728,7 @@ void cpu_fused_moe(
728
w2_bias.has_value() ? w2_bias->data_ptr<scalar_t>() : nullptr,
729
topk_weights.data_ptr<float>(), topk_id.data_ptr<int32_t>(), act_type,
730
token_num, expert_num, topk_num, input_size_13, output_size_13,
724
- input_size_2, output_size_2);
731
+ input_size_2, output_size_2, skip_weighted);
732
});
733
});
734
}
csrc/cpu/torch_bindings.cpp
+3
-2
@@ -119,8 +119,8 @@ void cpu_fused_moe(torch::Tensor& output, const torch::Tensor& input,
119
const std::optional<torch::Tensor>& w13_bias,
120
const std::optional<torch::Tensor>& w2_bias,
121
const torch::Tensor& topk_weights,
122
- const torch::Tensor& topk_id, const std::string& act,
123
- const std::string& isa);
122
+ const torch::Tensor& topk_id, const bool skip_weighted,
123
+ const std::string& act, const std::string& isa);
124
125
TORCH_LIBRARY_EXPAND(TORCH_EXTENSION_NAME, ops) {
126
// vLLM custom ops
@@ -320,6 +320,7 @@ TORCH_LIBRARY_EXPAND(TORCH_EXTENSION_NAME, ops) {
320
ops.def(
321
"cpu_fused_moe(Tensor(a0!) output, Tensor input, Tensor w13, Tensor w2, "
322
"Tensor? w13_bias, Tensor? w2_bias, Tensor topk_weights, Tensor topk_id, "
323
+ "bool skip_weighted, "
324
"str act, str isa) -> ()");
325
ops.impl("cpu_fused_moe", torch::kCPU, &cpu_fused_moe);
326
#endif
vllm/_custom_ops.py
+2
@@ -3078,6 +3078,7 @@ def cpu_fused_moe(
3078
topk_ids: torch.Tensor,
3079
act: str,
3080
isa: str,
3081
+ skip_weighted: bool = False,
3082
) -> torch.Tensor:
3083
output = torch.empty_like(input)
3084
torch.ops._C.cpu_fused_moe(
@@ -3089,6 +3090,7 @@ def cpu_fused_moe(
3090
w2_bias,
3091
topk_weights,
3092
topk_ids,
3093
+ skip_weighted,
3094
act,
3095
isa,
3096
)
vllm/model_executor/layers/fused_moe/cpu_fused_moe.py
+28
-8
@@ -238,7 +238,6 @@ class CPUFusedMOE:
238
activation: str = "silu",
239
) -> torch.Tensor:
240
assert activation in _CPU_MOE_ACT_FN, f"{activation} is not supported."
241
- assert not apply_router_weight_on_input
241
242
topk_weights, topk_ids = select_experts(
243
hidden_states=x,
@@ -261,6 +260,7 @@ class CPUFusedMOE:
260
topk_ids,
261
activation,
262
global_num_experts,
263
+ apply_router_weight_on_input,
264
)
265
266
def check_grouped_gemm(
@@ -355,7 +355,14 @@ class CPUFusedMOE:
355
topk_ids: torch.Tensor,
356
activation: str,
357
global_num_experts: int = -1,
358
+ skip_weighted: bool = False,
359
) -> torch.Tensor:
360
+ if skip_weighted:
361
+ assert topk_ids.size(1) == 1, (
362
+ "apply_router_weight_on_input is only implemented for topk=1"
363
+ )
364
+ input.mul_(topk_weights.to(input.dtype))
365
+
366
output = cpu_fused_moe(
367
input,
368
layer.w13_weight,
@@ -366,6 +373,7 @@ class CPUFusedMOE:
373
topk_ids,
374
activation,
375
self.isa,
376
+ skip_weighted,
377
)
378
return output
379
@@ -377,7 +385,14 @@ class CPUFusedMOE:
385
topk_ids: torch.Tensor,
386
activation: str,
387
global_num_experts: int = -1,
388
+ skip_weighted: bool = False,
389
) -> torch.Tensor:
390
+ if skip_weighted:
391
+ assert topk_ids.size(1) == 1, (
392
+ "apply_router_weight_on_input is only implemented for topk=1"
393
+ )
394
+ input.mul_(topk_weights.to(input.dtype))
395
+
396
output = torch.empty_like(input)
397
layer_id = id(layer)
398
torch.ops.vllm.cpu_fused_moe_torch(
@@ -388,6 +403,7 @@ class CPUFusedMOE:
403
topk_ids,
404
activation,
405
global_num_experts,
406
+ skip_weighted,
407
)
408
409
return output
@@ -401,6 +417,7 @@ def cpu_fused_moe_torch(
417
topk_ids: torch.Tensor,
418
activation: str,
419
global_num_experts: int = -1,
420
+ skip_weighted: bool = False,
421
) -> None:
422
layer = _CPU_MOE_LAYER_CACHE[layer_id]()
423
@@ -434,13 +451,16 @@ def cpu_fused_moe_torch(
451
new_x = torch.empty_like(outs)
452
453
new_x[idxs] = outs
437
- final_out = (
438
- new_x.view(*topk_ids.shape, -1)
439
- .type(topk_weights.dtype)
440
- .mul_(topk_weights.unsqueeze(dim=-1))
441
- .sum(dim=1)
442
- .type(new_x.dtype)
443
- )
454
+ if skip_weighted:
455
+ final_out = new_x
456
+ else:
457
+ final_out = (
458
+ new_x.view(*topk_ids.shape, -1)
459
+ .type(topk_weights.dtype)
460
+ .mul_(topk_weights.unsqueeze(dim=-1))
461
+ .sum(dim=1)
462
+ .type(new_x.dtype)
463
+ )
464
output.copy_(final_out)
465
466
vllm/v1/worker/cpu_worker.py
+14
-5
@@ -160,12 +160,21 @@ class CPUWorker(Worker):
160
x for x in logical_cpu_list if x.numa_node == selected_numa_node
161
]
162
else:
163
- assert len(logical_cpu_list) >= self.parallel_config.world_size
164
- logical_cpu_list = sorted(logical_cpu_list, key=lambda x: x.numa_node)
165
- sim_cpu_num_per_node = (
166
- len(logical_cpu_list) // self.parallel_config.world_size
163
+ # This is a bit tricky because the internal DP size
164
+ # is always 1 for non-MoE models
165
+ world_size_across_dp = (
166
+ self.parallel_config.world_size
167
+ * self.parallel_config._api_process_count
168
)
168
- start_idx = self.local_rank * sim_cpu_num_per_node
169
+ assert len(logical_cpu_list) >= world_size_across_dp
170
+ logical_cpu_list = sorted(logical_cpu_list, key=lambda x: x.numa_node)
171
+ sim_cpu_num_per_node = len(logical_cpu_list) // world_size_across_dp
172
+ assert self.parallel_config.data_parallel_rank_local is not None
173
+ start_idx = (
174
+ self.local_rank
175
+ + self.parallel_config.world_size
176
+ * self.parallel_config.data_parallel_rank_local
177
+ ) * sim_cpu_num_per_node
178
logical_cpu_list = logical_cpu_list[
179
start_idx : (start_idx + sim_cpu_num_per_node)
180
]