From ad44adbf0d2362a664d57beeb9b7c1af409dfbfd Mon Sep 17 00:00:00 2001
From: mikemikimike <13286568797@163.com>
Date: Fri, 6 Nov 2026 14:44:03 +0800
Subject: [PATCH] fix: derive vLLM context limits from checkpoint metadata

---
 src/quantpreflight/planner.py |  4 ++--
 src/quantpreflight/rules.py   | 42 ++++++++++++++++++++++++++++++++---
 tests/test_planner.py         |  7 ++++++
 tests/test_rule_boundaries.py | 16 ++++++++++++-
 4 files changed, 63 insertions(+), 6 deletions(-)

diff --git a/src/quantpreflight/planner.py b/src/quantpreflight/planner.py
index 5b8e30f..ded9c50 100644
--- a/src/quantpreflight/planner.py
+++ b/src/quantpreflight/planner.py
@@ -8,7 +8,7 @@
 
 from .models import CheckpointMeta, Finding, ScanReport, SetupPlan, TargetSpec
 from .profiles import resolve_gpu_profile
-from .rules import _effective_text_config, _pos_int, analyze_checkpoint, estimate_per_gpu_vram_gb
+from .rules import _effective_text_config, analyze_checkpoint, context_limit, estimate_per_gpu_vram_gb
 
 TP_CHOICES = (1, 2, 4, 8)
 # They say a model is risky to load, not that it won't start
@@ -147,7 +147,7 @@ def plan_setup(
 ) -> SetupPlan:
     """Find the closest setup to `target` that clears every finding a setting can clear."""
     tcfg = _effective_text_config(meta.config)
-    limit = _pos_int(tcfg.get("max_position_embeddings"))
+    limit = context_limit(meta.config, meta.tokenizer_config)
     gpu = resolve_gpu_profile(target.gpu) if target.gpu else None
     vram = target.vram_gb_override if target.vram_gb_override is not None else (gpu.vram_gb if gpu else None)
     usable = vram * target.gpu_memory_utilization if vram else None
diff --git a/src/quantpreflight/rules.py b/src/quantpreflight/rules.py
index 4d32438..3d00ca0 100644
--- a/src/quantpreflight/rules.py
+++ b/src/quantpreflight/rules.py
@@ -134,6 +134,42 @@ def _effective_text_config(config: dict[str, Any]) -> dict[str, Any]:
     return config
 
 
+_CONTEXT_LENGTH_KEYS = ("max_position_embeddings", "n_positions", "max_seq_len", "seq_length")
+_ROPE_TYPES_WITHOUT_FACTOR = frozenset(
+    {"su", "longrope", "llama3", "yarn", "deepseek_yarn", "deepseek_llama_scaling"}
+)
+
+
+def context_limit(config: dict[str, Any], tokenizer_config: dict[str, Any] | None = None) -> int | None:
+    """Derive the serving context limit using vLLM-compatible config precedence."""
+    text_config = _effective_text_config(config)
+    limits: list[int] = []
+    for key in _CONTEXT_LENGTH_KEYS:
+        value = _pos_int(text_config.get(key))
+        if value is not None:
+            limits.append(value)
+    model_max = _pos_int(text_config.get("model_max_length"))
+    if model_max is not None:
+        limits.append(model_max)
+    if not limits:
+        return None
+    limit = min(limits)
+    rope = text_config.get("rope_scaling")
+    model_type = str(text_config.get("model_type", "")).lower()
+    architectures = text_config.get("architectures")
+    architecture_names = architectures if isinstance(architectures, list) else ()
+    is_gemma3 = model_type == "gemma3" or any("gemma3" in str(name).lower() for name in architecture_names)
+    if isinstance(rope, dict) and not is_gemma3:
+        rope_type = str(rope.get("rope_type", rope.get("type", "default"))).lower()
+        factor = rope.get("factor")
+        if rope_type not in _ROPE_TYPES_WITHOUT_FACTOR and isinstance(factor, (int, float)) and factor > 0:
+            limit = int(limit * factor)
+    tokenizer_max = _pos_int((tokenizer_config or {}).get("model_max_length"))
+    if tokenizer_max is not None:
+        limit = min(limit, tokenizer_max)
+    return limit
+
+
 def _to_int_set(val: Any) -> set[int]:
     if isinstance(val, int) and not isinstance(val, bool):
         return {val}
@@ -1347,8 +1383,8 @@ def _version_in_range(version: tuple[int, ...], item: dict[str, Any]) -> bool:
     return True
 
 
-def _check_qpf016_context_limit(tcfg: dict[str, Any], target: TargetSpec) -> list[Finding]:
-    limit = _pos_int(tcfg.get("max_position_embeddings"))
+def _check_qpf016_context_limit(meta: CheckpointMeta, target: TargetSpec) -> list[Finding]:
+    limit = context_limit(meta.config, meta.tokenizer_config)
     if target.engine != "vllm" or not target.context_len or not limit or target.context_len <= limit:
         return []
     return [
@@ -1460,7 +1496,7 @@ def analyze_checkpoint(
     findings.extend(_check_qpf013_safetensors_spans(meta))
     findings.extend(_check_qpf014_legacy_pickle_weights(meta))
     findings.extend(_check_qpf015_engine_advisories(tgt, engine_prof))
-    findings.extend(_check_qpf016_context_limit(tcfg, tgt))
+    findings.extend(_check_qpf016_context_limit(meta, tgt))
 
     est_vram, vram_findings = _check_qpf012_vram_budget(meta, tcfg, tgt, gpu_prof)
     findings.extend(vram_findings)
diff --git a/tests/test_planner.py b/tests/test_planner.py
index 3f12f77..4e6b06d 100644
--- a/tests/test_planner.py
+++ b/tests/test_planner.py
@@ -57,6 +57,13 @@ def test_a_model_that_fits_is_ready_as_asked() -> None:
     assert plan.command == "vllm serve org/model"
 
 
+def test_setup_uses_tokenizer_context_cap_and_rope_scaling() -> None:
+    config = {**LLAMA, "max_position_embeddings": 4096, "rope_scaling": {"type": "linear", "factor": 4}}
+    meta = _meta(config=config, tokenizer_config={"model_max_length": 12000})
+    plan = plan_setup(meta, TargetSpec(gpu="H100", context_len=20000))
+    assert plan.context_len == 12000
+
+
 def test_heads_that_do_not_split_move_to_the_closest_gpu_count_with_fewer_gpus_on_a_tie() -> None:
     plan = plan_setup(_meta(10), TargetSpec(gpu="H100", tp=3))
     assert plan.status == "changed" and plan.tp == 2
diff --git a/tests/test_rule_boundaries.py b/tests/test_rule_boundaries.py
index c906e72..5886911 100644
--- a/tests/test_rule_boundaries.py
+++ b/tests/test_rule_boundaries.py
@@ -9,7 +9,7 @@
 from quantpreflight.formatters import _fmt_bytes, format_json_report, format_sarif_report, format_text_report
 from quantpreflight.models import CheckpointMeta, QuantSpec, Severity, TargetSpec, TensorSpec, TokenizerMeta
 from quantpreflight.profiles import load_advisories
-from quantpreflight.rules import analyze_checkpoint, estimate_per_gpu_vram_gb
+from quantpreflight.rules import analyze_checkpoint, context_limit, estimate_per_gpu_vram_gb
 
 LLAMA = {"architectures": ["LlamaForCausalLM"], "model_type": "llama"}
 AWQ4 = QuantSpec(method="awq", bits=4.0)
@@ -752,6 +752,20 @@ def test_qpf016_names_a_setting_that_works() -> None:
     assert report.findings[0].remediation == "Set --context-len to 2048 or lower."
 
 
+@pytest.mark.parametrize(
+    ("config", "tokenizer", "expected"),
+    [
+        ({"max_position_embeddings": 8192, "n_positions": 4096}, {}, 4096),
+        ({"max_position_embeddings": 8192, "model_max_length": 2048}, {}, 2048),
+        ({"max_position_embeddings": 8192}, {"model_max_length": 4096}, 4096),
+        ({"max_position_embeddings": 4096, "rope_scaling": {"type": "linear", "factor": 4}}, {}, 16384),
+        ({"max_position_embeddings": 4096, "rope_scaling": {"type": "yarn", "factor": 4}}, {}, 4096),
+    ],
+)
+def test_context_limit_matches_vllm_precedence(config: dict[str, Any], tokenizer: dict[str, Any], expected: int) -> None:
+    assert context_limit(config, tokenizer) == expected
+
+
 # Corrections from scanning 500 public models