[CPU][Bugfix] Fix flaky ShortConv prefill test on ARM (uninitialized weights) (#47848)
Signed-off-by: Rahul Vishwakarma <Rahul.Vishwakarma2@ibm.com>
Rahul Vishwakarma committed
Jul 8, 2026 at 06:50 UTC
f7efab58ec27ab233c021d20f69973a7cf913c2f
1 file changed
+9
tests/kernels/mamba/test_cpu_short_conv.py
+9
@@ -60,6 +60,12 @@ def test_short_conv_forward_native_prefill(vllm_config):
60
layer = ShortConv(config=config, dim=dim, layer_idx=0, prefix=prefix)
61
62
layer.to("cpu")
63
+ # vLLM Linear layers allocate weights with torch.empty (uninitialized).
64
+ # On ARM these come back as zero-filled pages, so in_proj output is zero and
65
+ # the prefill state stays zero. Seed + init to make the test platform-safe.
66
+ torch.manual_seed(0)
67
+ for p in layer.parameters():
68
+ torch.nn.init.normal_(p)
69
dispatch_cpu_unquantized_gemm(layer.in_proj, remove_weight=False)
70
dispatch_cpu_unquantized_gemm(layer.out_proj, remove_weight=False)
71
@@ -117,6 +123,9 @@ def test_short_conv_forward_native_decode(vllm_config):
123
layer = ShortConv(config=config, dim=dim, layer_idx=0, prefix=prefix)
124
125
layer.to("cpu")
126
+ torch.manual_seed(0)
127
+ for p in layer.parameters():
128
+ torch.nn.init.normal_(p)
129
dispatch_cpu_unquantized_gemm(layer.in_proj, remove_weight=False)
130
dispatch_cpu_unquantized_gemm(layer.out_proj, remove_weight=False)
131