diff --git a/families/qwen/tests/runtime_receipt.py b/families/qwen/tests/runtime_receipt.py index b3fd45c5c6..102c711e28 100644 --- a/families/qwen/tests/runtime_receipt.py +++ b/families/qwen/tests/runtime_receipt.py @@ -1,13 +1,16 @@ # SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. # SPDX-License-Identifier: Apache-2.0 -"""Qwen native-KV runtime receipt checks used by family E2E.""" +"""Qwen prompt and native-KV runtime receipt checks used by family E2E.""" from __future__ import annotations import re -_PREFILL = re.compile(r"^\[trtmc\.prefill\] tokens=(\d+) launches=(\d+) max_chunk=(\d+)$") +_PREFILL = re.compile( + r"^(?:\[[^,\]]+,\d+\]:)?\s*" + r"\[trtmc\.prefill\] tokens=(\d+) launches=(\d+) max_chunk=(\d+)$" +) _RUNTIME_ERROR = re.compile( r"\[trt\]\s+ERROR:|IExecutionContext::enqueueV3:\s+Error Code|" r"Internal Error:|Cuda Runtime|illegal memory access", @@ -24,6 +27,16 @@ def prefill_observations(stderr: str) -> tuple[tuple[int, int, int], ...]: return tuple(values) +def assert_prompt_token_count(payload: dict, prompt_tokens: int) -> None: + observations = prefill_observations(str(payload["runtime_stderr"])) + assert observations, "native execution did not report its prompt token count" + # Tensor-parallel ranks each report the complete prompt, not disjoint pieces. + for tokens, _, _ in observations: + assert tokens == prompt_tokens, ( + f"native prompt has {tokens} tokens; reference has {prompt_tokens}" + ) + + def assert_native_kv_receipt(payload: dict, case: dict, prompt_tokens: int) -> None: stderr = str(payload["runtime_stderr"]) expected_rows = int(case["expected_kv_cache_rows"]) diff --git a/families/qwen/tests/test_e2e.py b/families/qwen/tests/test_e2e.py index 58d95f9a31..53840eee8d 100644 --- a/families/qwen/tests/test_e2e.py +++ b/families/qwen/tests/test_e2e.py @@ -278,7 +278,7 @@ def _run_native( return payload -def _raw_prompt_token_count(model_dir: Path, manifest: dict, prompt: str) -> int: +def _prompt_token_count(model_dir: Path, manifest: dict, case: dict, prompt: str) -> int: from transformers import AutoTokenizer tokenizer = AutoTokenizer.from_pretrained( @@ -286,7 +286,7 @@ def _raw_prompt_token_count(model_dir: Path, manifest: dict, prompt: str) -> int local_files_only=True, trust_remote_code=bool(manifest.get("trust_remote_code", False)), ) - return len(tokenizer.encode(prompt, add_special_tokens=False)) + return int(_render_prompt(tokenizer, prompt, case)["input_ids"].shape[-1]) @cache @@ -765,7 +765,10 @@ def test_e2e(case_name: str, request, tmp_path: Path) -> None: ) if manifest["task"] == "embedding": return _embedding_e2e(manifest, case, model_dir, runtime_root, torch, tmp_path) + from families.qwen.tests.runtime_receipt import assert_prompt_token_count + prompt = _prompt(case) + prompt_tokens = _prompt_token_count(model_dir, manifest, case, prompt) record_evidence("inputs", {"prompt": prompt}) bundle = tmp_path / manifest["bundle"] @@ -784,6 +787,8 @@ def test_e2e(case_name: str, request, tmp_path: Path) -> None: tmp_path, ) record_evidence("native", payload) + with evidence_stage("compare"): + assert_prompt_token_count(payload, prompt_tokens) if case_name in _LOGIT_ORACLES: with evidence_stage("native"): payload["logits_trace"] = _native_logits_trace(bundle, prompt, case, tp_size, tmp_path) @@ -802,6 +807,7 @@ def test_e2e(case_name: str, request, tmp_path: Path) -> None: ) record_evidence("native", repeated) with evidence_stage("compare"): + assert_prompt_token_count(repeated, prompt_tokens) assert repeated["token_ids"] == payload["token_ids"] with evidence_stage("compare"): assert repeated["text"] == payload["text"] @@ -817,7 +823,6 @@ def test_e2e(case_name: str, request, tmp_path: Path) -> None: from families.qwen.tests.runtime_receipt import assert_native_kv_receipt assert "expected_prompt_tokens" in case - prompt_tokens = _raw_prompt_token_count(model_dir, manifest, prompt) with evidence_stage("compare"): assert_native_kv_receipt(payload, case, prompt_tokens) record_evidence("reference", {"mode": "contract_only", "oracle": "family runtime contract"}) diff --git a/families/qwen/tests/test_runtime_receipt.py b/families/qwen/tests/test_runtime_receipt.py index b9e9c19b10..130ec1ed54 100644 --- a/families/qwen/tests/test_runtime_receipt.py +++ b/families/qwen/tests/test_runtime_receipt.py @@ -5,7 +5,7 @@ import pytest -from families.qwen.tests.runtime_receipt import assert_native_kv_receipt +from families.qwen.tests.runtime_receipt import assert_native_kv_receipt, assert_prompt_token_count def _case() -> dict: @@ -60,3 +60,34 @@ def test_long_regression_keeps_exact_prompt_and_observed_token_gates() -> None: assert_native_kv_receipt(_payload(stderr), case, 65) with pytest.raises(AssertionError): assert_native_kv_receipt(_payload(stderr), case, 64) + + +@pytest.mark.parametrize("tokens", [20, 24, 144]) +def test_prompt_count_matches_without_kv_expectations(tokens: int) -> None: + assert_prompt_token_count( + _payload(f"[trtmc.prefill] tokens={tokens} launches=1 max_chunk={tokens}"), tokens + ) + + +def test_prompt_count_rejects_extra_thinking_prefix() -> None: + with pytest.raises(AssertionError, match="native prompt has 24 tokens; reference has 20"): + assert_prompt_token_count( + _payload("[trtmc.prefill] tokens=24 launches=1 max_chunk=24"), 20 + ) + + +def test_prompt_count_requires_native_receipt() -> None: + with pytest.raises(AssertionError, match="did not report"): + assert_prompt_token_count(_payload(""), 20) + + +@pytest.mark.parametrize("tagged", [False, True]) +def test_prompt_count_checks_each_tensor_parallel_rank(tagged: bool) -> None: + receipt = "[trtmc.prefill] tokens=20 launches=1 max_chunk=20" + rank0 = "[1,0]:" if tagged else "" + rank1 = "[1,1]:" if tagged else "" + assert_prompt_token_count(_payload(f"{rank0}{receipt}\n{rank1}{receipt}"), 20) + with pytest.raises(AssertionError, match="native prompt has 24 tokens"): + assert_prompt_token_count( + _payload(f"{rank0}{receipt}\n{rank1}[trtmc.prefill] tokens=24 launches=1 max_chunk=24"), 20 + )