From f7efab58ec27ab233c021d20f69973a7cf913c2f Mon Sep 17 00:00:00 2001 From: Rahul Vishwakarma <168823860+rahulssv-ibm@users.noreply.github.com> Date: Wed, 8 Jul 2026 06:50:09 +0530 Subject: [PATCH] [CPU][Bugfix] Fix flaky ShortConv prefill test on ARM (uninitialized weights) (#47848) Signed-off-by: Rahul Vishwakarma --- tests/kernels/mamba/test_cpu_short_conv.py | 9 +++++++++ 1 file changed, 9 insertions(+) diff --git a/tests/kernels/mamba/test_cpu_short_conv.py b/tests/kernels/mamba/test_cpu_short_conv.py index cd32e0901a7..c8e85a45511 100644 --- a/tests/kernels/mamba/test_cpu_short_conv.py +++ b/tests/kernels/mamba/test_cpu_short_conv.py @@ -60,6 +60,12 @@ def test_short_conv_forward_native_prefill(vllm_config): layer = ShortConv(config=config, dim=dim, layer_idx=0, prefix=prefix) layer.to("cpu") + # vLLM Linear layers allocate weights with torch.empty (uninitialized). + # On ARM these come back as zero-filled pages, so in_proj output is zero and + # the prefill state stays zero. Seed + init to make the test platform-safe. + torch.manual_seed(0) + for p in layer.parameters(): + torch.nn.init.normal_(p) dispatch_cpu_unquantized_gemm(layer.in_proj, remove_weight=False) dispatch_cpu_unquantized_gemm(layer.out_proj, remove_weight=False) @@ -117,6 +123,9 @@ def test_short_conv_forward_native_decode(vllm_config): layer = ShortConv(config=config, dim=dim, layer_idx=0, prefix=prefix) layer.to("cpu") + torch.manual_seed(0) + for p in layer.parameters(): + torch.nn.init.normal_(p) dispatch_cpu_unquantized_gemm(layer.in_proj, remove_weight=False) dispatch_cpu_unquantized_gemm(layer.out_proj, remove_weight=False)