[CPU][Bugfix] Fix flaky ShortConv prefill test on ARM (uninitialized weights) (#47848)

Signed-off-by: Rahul Vishwakarma <Rahul.Vishwakarma2@ibm.com>
(cherry picked from commit f7efab58ec)

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)
 

tests/models/transformers/fusers/test_linear.py

@@ -258,7 +258,8 @@ def _apply_qkv_fuser_with_stubs(module: nn.Module, fuser: QKVFuser):
         merged.weight.copy_(torch.cat([q.weight, k.weight, v.weight], dim=0))
         if q.bias is not None:
             merged.bias.copy_(torch.cat([q.bias, k.bias, v.bias], dim=0))
-    merged.split_sizes = [q.out_features, k.out_features, v.out_features]
+    merged.output_sizes = [q.out_features, k.out_features, v.out_features]
+    merged.tp_size = 1
     setattr(module, fuser.merged_name, merged)
     for name in (fuser.q_name, fuser.k_name, fuser.v_name):
         delattr(module, name)
@@ -332,7 +333,8 @@ def test_detects_and_rewrites_qkv(attn_cls, kv_heads):
     # original semantics (branches, kwargs, attribute reads)
     code = fuser.fused_forward.__code__
     names = code.co_names
-    assert "qkv_proj" in names and "split_sizes" in names and "o_proj" in names
+    assert "qkv_proj" in names and "output_sizes" in names and "o_proj" in names
+    assert "tp_size" in names
     assert not {"q_proj", "k_proj", "v_proj"} & set(names)
     if attn_cls is FakeAttention:
         assert "update" in names  # the cache branch survives

tests/quantization/test_modelopt.py

@@ -133,6 +133,49 @@ def test_modelopt_mixed_precision_quantizes_parallel_lm_head():
     assert isinstance(method, ModelOptNvFp4LinearMethod)
 
 
+def test_modelopt_mixed_precision_infers_fused_gate_up_projection():
+    from vllm.model_executor.layers.linear import LinearBase
+
+    config = _mixed_precision_config(
+        {
+            "model.layers.0.mlp.gate_proj": {"quant_algo": "NVFP4"},
+            "model.layers.0.mlp.up_proj": {"quant_algo": "NVFP4"},
+        }
+    )
+
+    fake_layer = MagicMock(spec=LinearBase)
+    with patch(
+        "vllm.model_executor.layers.quantization.modelopt.init_nvfp4_linear_kernel"
+    ):
+        method = config.get_quant_method(fake_layer, "model.layers.0.mlp.gate_up_proj")
+
+    assert isinstance(method, ModelOptNvFp4LinearMethod)
+
+
+@pytest.mark.parametrize(
+    ("quantized_prefix", "missing_prefix"),
+    [
+        ("model.layers.0.mlp.gate_proj", "model.layers.0.mlp.down_proj"),
+        ("model.layers.0.self_attn.o_proj", "model.layers.0.self_attn.qkv_proj"),
+    ],
+)
+def test_modelopt_mixed_precision_does_not_infer_missing_sibling_linear(
+    quantized_prefix, missing_prefix
+):
+    from vllm.model_executor.layers.linear import LinearBase
+
+    config = _mixed_precision_config(
+        {
+            quantized_prefix: {"quant_algo": "NVFP4"},
+        }
+    )
+
+    fake_layer = MagicMock(spec=LinearBase)
+    method = config.get_quant_method(fake_layer, missing_prefix)
+
+    assert isinstance(method, UnquantizedLinearMethod)
+
+
 def test_vocab_parallel_embedding_weight_loader_accepts_scalar_scale():
     holder = Mock()
     scale = torch.nn.Parameter(torch.empty(1))

vllm/model_executor/layers/linear.py

@@ -1000,16 +1000,12 @@ class QKVParallelLinear(ColumnParallelLinear):
             self.num_kv_heads = divide(self.total_num_kv_heads, tp_size)
             self.num_kv_head_replicas = 1
         input_size = self.hidden_size
-        output_size = (
-            self.num_heads * self.head_size
-            + self.num_kv_heads * self.head_size
-            + self.num_kv_heads * self.v_head_size
-        ) * tp_size
         self.output_sizes = [
             self.num_heads * self.head_size * tp_size,  # q_proj
             self.num_kv_heads * self.head_size * tp_size,  # k_proj
             self.num_kv_heads * self.v_head_size * tp_size,  # v_proj
         ]
+        output_size = sum(self.output_sizes)
 
         super().__init__(
             input_size=input_size,

vllm/model_executor/layers/quantization/modelopt.py

@@ -2457,16 +2457,29 @@ class ModelOptMixedPrecisionConfig(ModelOptQuantConfigBase):
                 if key.startswith(parent_dot):
                     return info["quant_algo"].upper()
 
-        # 4. Parent-prefix fallback for fused projections (qkv_proj, gate_up_proj).
-        for candidate in self._quantized_layer_prefix_candidates(prefix):
-            parent_dot = candidate.rsplit(".", 1)[0] + "."
-            algos = {
-                info["quant_algo"].upper()
-                for key, info in self.quantized_layers.items()
-                if key.startswith(parent_dot) and "." not in key[len(parent_dot) :]
-            }
-            if len(algos) == 1:
-                return algos.pop()
+        # 4. Parent-prefix fallback for fused projections whose config lists
+        # shard names instead of vLLM's packed module name.
+        fused_projection_shards = {
+            "qkv_proj": ("q_proj", "k_proj", "v_proj"),
+            "gate_up_proj": ("gate_proj", "up_proj"),
+        }
+        shard_names = fused_projection_shards.get(proj_name)
+        if shard_names is not None:
+            for candidate in self._quantized_layer_prefix_candidates(prefix):
+                parent_dot = candidate.rsplit(".", 1)[0] + "."
+                shard_algos: set[str] = set()
+                for shard_name in shard_names:
+                    shard_prefix = f"{parent_dot}{shard_name}"
+                    if shard_prefix in self.quantized_layers:
+                        algo = self.quantized_layers[shard_prefix]["quant_algo"].upper()
+                        shard_algos.add(algo)
+                if len(shard_algos) == 1:
+                    return shard_algos.pop()
+                if len(shard_algos) > 1:
+                    raise ValueError(
+                        f"Mixed quant_algo within fused layer {prefix}: "
+                        f"{shard_algos}. All shards must use the same quantization."
+                    )
 
         return None
 

vllm/model_executor/models/transformers/fusers/qkv.py

@@ -127,15 +127,14 @@ class QKVFuser(StackedFuser):
         if len({id(block) for block, _ in blocks}) != 1:
             raise ValueError("projection calls are in different blocks")
 
-        # q(x), k(x), v(x) -> q, k, v = qkv(x).split(self.qkv.split_sizes, -1)
+        # q(x), k(x), v(x) -> q, k, v = qkv(x).split(qkv.output_sizes / qkv.tp_size, -1)
         names = {node.id for node in ast.walk(funcdef) if isinstance(node, ast.Name)}
         temps = [f"{name}_fused" for name in (self.q_name, self.k_name, self.v_name)]
         if names & set(temps):
             raise ValueError("fused temporaries would shadow existing names")
         merged = f"self.{self.merged_name}"
-        template = (
-            f"{', '.join(temps)} = {merged}(__arg__).split({merged}.split_sizes, -1)"
-        )
+        sections = f"[s // {merged}.tp_size for s in {merged}.output_sizes]"
+        template = f"{', '.join(temps)} = {merged}(__arg__).split({sections}, -1)"
         assign = ast.parse(template).body[0]
         arg = next(
             node
@@ -198,13 +197,6 @@ class QKVFuser(StackedFuser):
             self.merged_name,
             merged,
         )
-        # The rewritten forward splits the merged projection into this rank's
-        # shard sizes (see `update_forward`)
-        merged.split_sizes = [
-            merged.num_heads * merged.head_size,
-            merged.num_kv_heads * merged.head_size,
-            merged.num_kv_heads * merged.v_head_size,
-        ]
         setattr(module, self.merged_name, merged)
         # Drop the consumed submodules so their (meta) params are not expected.
         for name in (self.q_name, self.k_name, self.v_name):