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 15, 19:45 UTC
[Bugfix][IPEX] Add `VLLM_CPU_MOE_PREPACK` to allow disabling MoE prepack when CPU does not support it (#14681)
Signed-off-by: Thien Tran <gau.nernst@yahoo.com.sg>
Thien Tran committed
Mar 14, 2025 at 11:43 UTC
95d680b8620b4f1a087a2857142a05c15c45aceb
3 files changed
+10
-1
docs/source/getting_started/installation/cpu.md
+1
@@ -195,6 +195,7 @@ vLLM CPU backend supports the following vLLM features:
195
196
- `VLLM_CPU_KVCACHE_SPACE`: specify the KV Cache size (e.g, `VLLM_CPU_KVCACHE_SPACE=40` means 40 GB space for KV cache), larger setting will allow vLLM running more requests in parallel. This parameter should be set based on the hardware configuration and memory management pattern of users.
197
- `VLLM_CPU_OMP_THREADS_BIND`: specify the CPU cores dedicated to the OpenMP threads. For example, `VLLM_CPU_OMP_THREADS_BIND=0-31` means there will be 32 OpenMP threads bound on 0-31 CPU cores. `VLLM_CPU_OMP_THREADS_BIND=0-31|32-63` means there will be 2 tensor parallel processes, 32 OpenMP threads of rank0 are bound on 0-31 CPU cores, and the OpenMP threads of rank1 are bound on 32-63 CPU cores.
198
+- `VLLM_CPU_MOE_PREPACK`: whether to use prepack for MoE layer. This will be passed to `ipex.llm.modules.GatedMLPMOE`. Default is `1` (True). On unsupported CPUs, you might need to set this to `0` (False).
199
200
## Performance tips
201
vllm/envs.py
+7
@@ -40,6 +40,7 @@ if TYPE_CHECKING:
40
VLLM_PP_LAYER_PARTITION: Optional[str] = None
41
VLLM_CPU_KVCACHE_SPACE: int = 0
42
VLLM_CPU_OMP_THREADS_BIND: str = ""
43
+ VLLM_CPU_MOE_PREPACK: bool = True
44
VLLM_OPENVINO_DEVICE: str = "CPU"
45
VLLM_OPENVINO_KVCACHE_SPACE: int = 0
46
VLLM_OPENVINO_CPU_KV_CACHE_PRECISION: Optional[str] = None
@@ -349,6 +350,12 @@ environment_variables: dict[str, Callable[[], Any]] = {
350
"VLLM_CPU_OMP_THREADS_BIND":
351
lambda: os.getenv("VLLM_CPU_OMP_THREADS_BIND", "all"),
352
353
+ # (CPU backend only) whether to use prepack for MoE layer. This will be
354
+ # passed to ipex.llm.modules.GatedMLPMOE. On unsupported CPUs, you might
355
+ # need to set this to "0" (False).
356
+ "VLLM_CPU_MOE_PREPACK":
357
+ lambda: bool(int(os.getenv("VLLM_CPU_MOE_PREPACK", "1"))),
358
+
359
# OpenVINO device selection
360
# default is CPU
361
"VLLM_OPENVINO_DEVICE":
vllm/model_executor/layers/fused_moe/layer.py
+2
-1
@@ -7,6 +7,7 @@ from typing import Callable, List, Optional, Tuple
7
import torch
8
from torch.nn.parameter import UninitializedParameter
9
10
+from vllm import envs
11
from vllm.config import get_current_vllm_config
12
from vllm.distributed import (get_dp_group, get_tensor_model_parallel_rank,
13
get_tensor_model_parallel_world_size,
@@ -104,7 +105,7 @@ class UnquantizedFusedMoEMethod(FusedMoEMethodBase, CustomOp):
105
layer.ipex_fusion = ipex.llm.modules.GatedMLPMOE(
106
layer.w13_weight,
107
layer.w2_weight,
107
- use_prepack=True,
108
+ use_prepack=envs.VLLM_CPU_MOE_PREPACK,
109
)
110
else:
111
raise NotImplementedError("CPU MOE only supports x86 arch.")