[Bugfix][CPU][MoE] Fix out-of-bounds write and segfault when router weights are fp32 (#56168)
Signed-off-by: Farzad Abdolhosseini <farzad@elastix.ai> Co-authored-by: Claude <noreply@anthropic.com> Co-authored-by: Li, Jiang <jiang1.li@intel.com>
Farzad Abdolhosseini committed
Sep 24, 2026 at 01:15 UTC
8b06a343eca7646ca7f7042c57b9ce953fef1892
2 files changed
+35
-3
tests/model_executor/test_cpu_unquantized_gemm_dispatch.py
+23
@@ -7,6 +7,7 @@ import torch
7
8
from vllm.model_executor.layers import utils
9
from vllm.platforms import current_platform
10
+from vllm.utils.torch_utils import set_default_torch_dtype
11
12
13
@pytest.fixture(scope="module")
@@ -89,3 +90,25 @@ def test_dispatch_cpu_unquantized_gemm_logs_zentorch_dispatch(monkeypatch):
90
expected_prepacked,
91
)
92
]
93
+
94
+
95
+@pytest.mark.usefixtures("_mock_zentorch_linear_unary")
96
+@pytest.mark.parametrize("weight_dtype", [torch.bfloat16, torch.float16, torch.float32])
97
+def test_dispatch_cpu_unquantized_gemm_remove_weight_keeps_dtype(
98
+ monkeypatch, weight_dtype
99
+):
100
+ monkeypatch.setattr(current_platform, "is_zen_cpu", lambda: True)
101
+
102
+ layer = torch.nn.Linear(16, 8, bias=False, dtype=weight_dtype)
103
+ loading_dtype = (
104
+ torch.float32 if weight_dtype is not torch.float32 else torch.bfloat16
105
+ )
106
+ with set_default_torch_dtype(loading_dtype):
107
+ utils.dispatch_cpu_unquantized_gemm(layer, remove_weight=True)
108
+
109
+ assert layer.weight.numel() == 0
110
+ assert layer.weight.dtype is weight_dtype
111
+
112
+ x = torch.randn(4, 16, dtype=weight_dtype)
113
+ output = layer.cpu_linear(x, layer.weight, None)
114
+ assert not output.isnan().any()
vllm/model_executor/layers/utils.py
+12
-3
@@ -513,7 +513,10 @@ def dispatch_cpu_unquantized_gemm(
513
)
514
)
515
if remove_weight:
516
- layer.weight = torch.nn.Parameter(torch.empty(0), requires_grad=False)
516
+ layer.weight = torch.nn.Parameter(
517
+ torch.empty(0, dtype=dtype, device=layer.weight.device),
518
+ requires_grad=False,
519
+ )
520
logger.debug_once(
521
"CPU unquantized GEMM dispatch: using zentorch_linear_unary (prepacked=%s)",
522
is_prepacked,
@@ -539,7 +542,10 @@ def dispatch_cpu_unquantized_gemm(
542
bias,
543
)
544
if remove_weight:
542
- layer.weight = torch.nn.Parameter(torch.empty(0), requires_grad=False)
545
+ layer.weight = torch.nn.Parameter(
546
+ torch.empty(0, dtype=dtype, device=layer.weight.device),
547
+ requires_grad=False,
548
+ )
549
logger.debug_once(
550
"CPU unquantized GEMM dispatch: using sgl-kernel weight_packed_linear"
551
)
@@ -554,7 +560,10 @@ def dispatch_cpu_unquantized_gemm(
560
handler = ops.create_onednn_mm(origin_weight.t(), 32)
561
layer.cpu_linear = lambda x, weight, bias: ops.onednn_mm(handler, x, bias)
562
if remove_weight:
557
- layer.weight = torch.nn.Parameter(torch.empty(0), requires_grad=False)
563
+ layer.weight = torch.nn.Parameter(
564
+ torch.empty(0, dtype=dtype, device=layer.weight.device),
565
+ requires_grad=False,
566
+ )
567
logger.debug_once("CPU unquantized GEMM dispatch: using oneDNN onednn_mm")
568
return
569
except RuntimeError as e: