[CPU][Recipes] Detect model head constraints for automatic tensor parallel selection in vLLM Recipes Tool (#60170)

Signed-off-by: louie-tsai <louie.tsai@intel.com>

tools/recipes/RUNTIME_TUNING.md

@@ -79,7 +79,7 @@ on a validation environment that differs from the user's target deployment.
 
 | Runtime parameter | Why it may need deployment-time refinement | Main decision input | Current policy |
 | --- | --- | --- | --- |
-| `tensor-parallel-size` | The effective CPU/NUMA topology available to a container or pod can differ from the system used to validate the recipe. | Hardware topology | Use the largest power-of-two TP value that does not exceed the effective NUMA-node count. |
+| `tensor-parallel-size` | Available topology and model head counts constrain TP. | Hardware topology + model configuration | Use the largest head-compatible power-of-two TP within the effective NUMA-node count, for DP=1 and PP=1. |
 | `gpu-memory-utilization` | Available memory can differ by machine size, container limits, and other memory use. The vLLM option name is also used by the CPU backend. | Hardware memory + recipe baseline | Calculate a conservative fraction from the most constrained NUMA node. |
 | `max-num-seqs` | The useful scheduler concurrency depends on the number of requests expected to be active at the same time. | Workload concurrency | Set `max-num-seqs` to `--concurrency` when supplied. |
 | `max-num-batched-tokens` | Each scheduler iteration must share its token budget between active decodes and incoming prefills. | Input/output token shape, concurrency, and optional QPS/TPOT | Calculate decode budget + expected prefill demand, with vLLM scheduler constraints as floors. |
@@ -99,6 +99,35 @@ For example, 2 effective NUMA nodes produce TP=2, while 6 nodes currently
 produce TP=4. This is a topology-based starting point and avoids automatically
 selecting unusual non-power-of-two TP sizes before they are validated.
 
+The policy then reads the model configuration through vLLM's config loader,
+without loading weights. It honors the model path/ID, revision, code revision,
+`trust-remote-code`, config format and JSON `hf-overrides`. No model-specific
+numeric head-count table is maintained.
+
+Each candidate must divide the attention-head count. KV heads must divide
+across TP ranks when KV heads >= TP; otherwise TP must be a multiple of the KV
+head count so the heads can be replicated. An incompatible candidate is halved
+until it passes every detected constraint. Nested text configuration, global KV
+heads, Whisper encoder/decoder heads and hybrid linear-attention heads are
+checked. Vision head divisibility is checked conservatively unless
+`mm-encoder-tp-mode: data` is configured.
+
+For Phi-4-reasoning on four NUMA nodes, the detected 40 attention heads and 10
+KV heads reduce the initial TP=4 suggestion to TP=2. The printed tuning notes
+show the detected counts and the reduction; `config.yml` receives TP=2.
+
+If configuration loading or head extraction fails, the tool keeps the recipe
+TP (vLLM's default TP=1 when absent) and prints why detection failed. That
+fallback does not certify an existing recipe TP as compatible. Plain conversion
+without hardware detection does not fetch the model configuration.
+
+Automatic TP selection currently requires recipe DP=1 and PP=1; otherwise the
+recipe's parallelism is retained. Automatic DP tuning remains disabled pending
+CPU DP binding and capacity validation. No unused NUMA nodes are converted into
+DP replicas. These checks establish head compatibility, not memory fit or
+complete model/backend support. Explicit `vllm serve` CLI overrides still take
+precedence over the generated YAML and are validated by vLLM at startup.
+
 ### `gpu-memory-utilization`
 
 For every effective NUMA memory node, the policy calculates:

tools/recipes/model_constraints.py

@@ -0,0 +1,111 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+
+"""Read model head constraints without loading model weights."""
+
+from __future__ import annotations
+
+import json
+from dataclasses import dataclass
+from typing import Any
+
+
+@dataclass(frozen=True)
+class HeadConstraint:
+    name: str
+    attention_heads: int
+    kv_heads: int
+
+    def supports(self, tp: int) -> bool:
+        if self.attention_heads % tp:
+            return False
+        if self.kv_heads >= tp:
+            return self.kv_heads % tp == 0
+        return tp % self.kv_heads == 0
+
+
+def _positive_int(value: object, name: str) -> int:
+    if isinstance(value, bool) or not isinstance(value, int) or value < 1:
+        raise ValueError(f"Cannot determine a positive integer for {name}: {value!r}")
+    return value
+
+
+def extract_head_constraints(
+    hf_config: Any, *, include_vision: bool = True
+) -> tuple[HeadConstraint, ...]:
+    """Extract decoder, global-attention, hybrid and encoder head constraints."""
+    constraints = []
+    if getattr(hf_config, "model_type", None) == "whisper":
+        for component in ("encoder", "decoder"):
+            heads = _positive_int(
+                getattr(hf_config, f"{component}_attention_heads", None), component
+            )
+            constraints.append(HeadConstraint(component, heads, heads))
+        return tuple(constraints)
+
+    text = hf_config.get_text_config()
+    heads = _positive_int(getattr(text, "num_attention_heads", None), "attention heads")
+    kv = getattr(text, "num_key_value_heads", None)
+    if kv is None:
+        kv = heads
+    kv = _positive_int(kv, "KV heads")
+    constraints.append(HeadConstraint("text", heads, kv))
+
+    global_kv = getattr(text, "num_global_key_value_heads", None)
+    if global_kv is not None:
+        constraints.append(
+            HeadConstraint(
+                "global attention", heads, _positive_int(global_kv, "global KV")
+            )
+        )
+
+    # Hybrid layers have separate head counts; use conservative divisibility
+    # rather than assuming they support standard attention's KV replication.
+    for field in ("linear_num_key_heads", "linear_num_value_heads"):
+        count = getattr(text, field, None)
+        if count is not None:
+            count = _positive_int(count, field)
+            constraints.append(HeadConstraint(field, count, count))
+
+    vision = getattr(hf_config, "vision_config", None)
+    if include_vision and vision is not None:
+        vision_heads = getattr(vision, "num_attention_heads", None)
+        if vision_heads is None:
+            vision_heads = getattr(vision, "num_heads", None)
+        vision_heads = _positive_int(vision_heads, "vision heads")
+        constraints.append(HeadConstraint("vision", vision_heads, vision_heads))
+    return tuple(constraints)
+
+
+def load_head_constraints(config: dict[str, Any]) -> tuple[HeadConstraint, ...]:
+    """Load the effective serving configuration through vLLM's config loader."""
+    from vllm.transformers_utils.config import get_config
+
+    hf_overrides = config.get("hf-overrides")
+    if isinstance(hf_overrides, str):
+        hf_overrides = json.loads(hf_overrides)
+    if hf_overrides is not None and not isinstance(hf_overrides, dict):
+        raise ValueError("Head detection requires hf-overrides to be a JSON object")
+    hf_config = get_config(
+        config["model"],
+        trust_remote_code=config.get("trust-remote-code", False),
+        revision=config.get("revision"),
+        code_revision=config.get("code-revision"),
+        config_format=config.get("config-format", "auto"),
+        hf_overrides_kw=hf_overrides,
+    )
+    return extract_head_constraints(
+        hf_config, include_vision=config.get("mm-encoder-tp-mode") != "data"
+    )
+
+
+def suggest_tensor_parallel_size(
+    candidate: int, constraints: tuple[HeadConstraint, ...]
+) -> int:
+    """Reduce an automatic power-of-two candidate until all head checks pass."""
+    _positive_int(candidate, "TP candidate")
+    if candidate & (candidate - 1) or not constraints:
+        raise ValueError("Expected a power-of-two TP candidate and model constraints")
+    while not all(constraint.supports(candidate) for constraint in constraints):
+        candidate //= 2
+    return candidate

tools/recipes/runtime_tuning.py

@@ -19,6 +19,7 @@ from math import ceil
 from typing import Any
 
 from hardware_detection import HardwareInfo
+from model_constraints import load_head_constraints, suggest_tensor_parallel_size
 
 from vllm.config.scheduler import SchedulerConfig
 
@@ -119,19 +120,50 @@ def _resolve_tensor_parallel_size(
     workload: WorkloadHints,
     result: TuningResult,
 ) -> None:
-    del config, workload
+    del workload
     if hardware is None or hardware.numa_node_count < 1:
         return
 
+    if any(
+        config.get(key, 1) != 1
+        for key in ("data-parallel-size", "pipeline-parallel-size")
+    ):
+        result.notes.append(
+            "tensor-parallel-size: kept recipe value; automatic TP selection "
+            "currently requires DP=1 and PP=1"
+        )
+        return
+
     # Draft policy: use the largest power-of-two TP that does not exceed the
     # effective NUMA-node count. This deliberately avoids hard-coding unusual
     # non-power-of-two topologies into the converter.
     candidate = 1 << (hardware.numa_node_count.bit_length() - 1)
+    try:
+        constraints = load_head_constraints(config)
+    except Exception as exc:
+        result.notes.append(
+            "tensor-parallel-size: kept recipe value (vLLM default TP=1 if absent); "
+            f"model head detection failed: {exc}"
+        )
+        return
+
+    for constraint in constraints:
+        result.notes.append(
+            f"detected {constraint.name}: "
+            f"attention_heads={constraint.attention_heads}, "
+            f"kv_heads={constraint.kv_heads}"
+        )
+    selected = suggest_tensor_parallel_size(candidate, constraints)
+    if selected != candidate:
+        result.notes.append(
+            f"tensor-parallel-size: reduced NUMA-derived TP={candidate} to "
+            f"TP={selected} for model head partitioning/replication"
+        )
     _record_override(
         result,
         "tensor-parallel-size",
-        candidate,
-        "derived from the effective NUMA-node count",
+        selected,
+        "derived from effective NUMA nodes and detected model head constraints",
     )
 
 
@@ -256,8 +288,8 @@ def _resolve_data_parallel_size(
         return
 
     result.notes.append(
-        "data-parallel-size: target QPS supplied; kept recipe value because "
-        "per-replica capacity is not known yet"
+        "data-parallel-size: automatic DP tuning is disabled; kept recipe value "
+        "because CPU DP auto-binding and per-replica capacity need validation"
     )