diff --git a/scripts/e2e_eval/run_eval.py b/scripts/e2e_eval/run_eval.py index 394cb75cb..5393eb5d9 100644 --- a/scripts/e2e_eval/run_eval.py +++ b/scripts/e2e_eval/run_eval.py @@ -102,6 +102,12 @@ # An explicit per-model precision (``ModelEntry.precision``) overrides this. This # expansion is NPU-only; off-NPU devices never build quantized variants. _NPU_FALLBACK_PRECISIONS: tuple[str, ...] = ("w8a8", "w8a16") +_FLOAT_PRECISIONS = frozenset({"auto", "fp16", "fp32"}) +_NAMED_QUANTIZED_PRECISIONS = frozenset({"int4", "int8", "int16"}) +_MIXED_PRECISION_RE = re.compile(r"^w(\d+)a(\d+)$") +_VALID_WEIGHT_BITS = frozenset({4, 8, 16}) +_VALID_ACTIVATION_BITS = frozenset({8, 16, 32}) +_QDQ_WEIGHT_BITS = frozenset({8, 16}) # EPs whose eval track keeps the model unquantized (the "fp" variant) # rather than running winml's QDQ pass on top. This is an eval-setup @@ -868,7 +874,8 @@ def _extract_onnx_path(build_proc: dict, hf_id: str, task: str | None) -> str | # Patterns used by winml build to report the artifact path markers = ("Final artifact:", "Existing artifact found:", "Artifact:") onnx_path = None - for line in (build_proc["stderr"] + build_proc["stdout"]).splitlines(): + text = _strip_ansi(build_proc["stderr"] + build_proc["stdout"]) + for line in text.splitlines(): for marker in markers: if marker in line: candidate = line.split(marker)[-1].strip() @@ -878,12 +885,44 @@ def _extract_onnx_path(build_proc: dict, hf_id: str, task: str | None) -> str | if onnx_path: break + if not onnx_path: + onnx_path = _extract_wrapped_onnx_path(text, markers) + if not onnx_path or not Path(onnx_path).exists(): onnx_path = _find_cached_model(hf_id, build_proc, task) return onnx_path +_ANSI_ESCAPE_RE = re.compile(r"\x1b\[[0-?]*[ -/]*[@-~]") + + +def _strip_ansi(text: str) -> str: + """Remove ANSI SGR/control escapes from captured Rich output.""" + return _ANSI_ESCAPE_RE.sub("", text) + + +def _extract_wrapped_onnx_path(text: str, markers: tuple[str, ...]) -> str | None: + """Extract an artifact path that Rich wrapped onto following lines.""" + stop_markers = ("Build config:", "Export ", "Optimize ", "Quantize ", "Compile ") + for marker in markers: + start = text.find(marker) + if start == -1: + continue + fragments: list[str] = [] + for raw_line in text[start + len(marker) :].splitlines(): + line = raw_line.strip() + if not line: + continue + if any(stop in line for stop in stop_markers): + break + fragments.append(line) + for candidate in (line, "".join(fragments), " ".join(fragments)): + if candidate and Path(candidate).exists(): + return candidate + return None + + def _extract_task_from_config(config_path: Path) -> str | None: """Read the task from a build config JSON file.""" try: @@ -2364,12 +2403,22 @@ def _is_quantized_precision(precision: str) -> bool: """True if a recipe precision implies quantization. Off-NPU devices skip quantized recipe variants (see :func:`_build_jobs`). - Delegates to winml's precision policy so the classification never drifts - from the CLI (``fp16`` -> False, ``w8a16``/``w8a8`` -> True). + Keep this local so the eval harness can classify recipe variants without + importing the full winml config package and its ONNX Runtime dependencies. """ - from winml.modelkit.config.precision import is_quantized_precision - - return is_quantized_precision(precision) + p = precision.lower() + if p in _FLOAT_PRECISIONS: + return False + if p in _NAMED_QUANTIZED_PRECISIONS: + return True + match = _MIXED_PRECISION_RE.match(p) + if not match: + return False + weight_bits, activation_bits = int(match.group(1)), int(match.group(2)) + if weight_bits not in _VALID_WEIGHT_BITS or activation_bits not in _VALID_ACTIVATION_BITS: + return False + # w8a32/w16a32 are invalid float-activation mixes; w4a32 is weight-only. + return not (activation_bits == 32 and weight_bits in _QDQ_WEIGHT_BITS) def _build_jobs( @@ -2422,8 +2471,7 @@ def _build_jobs( # explicit per-model precision (e.g. fp16) skips this and is honored # by the single-fallback branch below via _resolve_precision. jobs.extend( - EvalJob(entry, None, fallback_precision=prec) - for prec in _NPU_FALLBACK_PRECISIONS + EvalJob(entry, None, fallback_precision=prec) for prec in _NPU_FALLBACK_PRECISIONS ) else: jobs.append(EvalJob(entry, None)) diff --git a/src/winml/modelkit/models/hf/t5.py b/src/winml/modelkit/models/hf/t5.py index d4c3fdc06..48a1be8b9 100644 --- a/src/winml/modelkit/models/hf/t5.py +++ b/src/winml/modelkit/models/hf/t5.py @@ -39,7 +39,7 @@ from ...optim import WinMLOptimizationConfig from ..winml.composite_model import register_composite_model from ..winml.encoder_decoder import EncoderDecoderInputGenerator, WinMLEncoderDecoderModel -from ..winml.kv_cache import PastKeyValueInputGenerator, WinMLSlidingWindowCache +from ..winml.kv_cache import PastKeyValueInputGenerator, WinMLStaticCache if TYPE_CHECKING: @@ -87,31 +87,16 @@ def forward( class T5DecoderWrapper(nn.Module): - """Wraps T5ForConditionalGeneration with sliding-window KV cache I/O. + """Wraps T5ForConditionalGeneration with static KV cache I/O. Input: full buffer ``[batch, heads, max_decode, d_kv]`` per layer. Output: only the new token's KV ``[batch, heads, 1, d_kv]`` per layer. - Uses ``WinMLSlidingWindowCache`` (Slice+Concat eviction) wrapped in - ``EncoderDecoderCache`` (cross-attn empty → always recomputed from - ``encoder_hidden_states``). - - ``cache_position`` is intentionally NOT an ONNX input — it is pinned to - ``[max_cache_len - 1]`` (the rightmost buffer slot) inside ``forward`` and - traced as a Constant. For single-token generation with a sliding window, - the new token is always written to the rightmost slot, so this value is - invariant. Baking it in lets ONNX constant-fold the entire - ``compute_bias`` subgraph (``memory_position - context_position`` is - constant → learned-bias Gather becomes a fixed tensor) and collapses the - causal mask ``kv_idx <= q_idx`` (all-True since ``q_idx == W-1``). - - This couples the exported graph to sliding-window semantics at build - time. ``WinMLStaticCache`` cannot be used as the *inference* cache for - this ONNX — its buffer layout (left-aligned, index_copy_) does not match - the graph's internal Slice+Concat. Callers who want static-cache - semantics must subclass ``T5DecoderWrapper``, take ``cache_position`` as - an input again, and re-export. ``WinMLStaticCache`` itself remains - fully functional for that path. + Uses ``WinMLStaticCache`` wrapped in ``EncoderDecoderCache``. Transformers + 5.x derives T5 causal mask and relative position bias from + ``past_key_values.get_seq_length()`` instead of the now-deprecated + ``cache_position`` kwarg, so the wrapper stores the ONNX ``cache_position`` + input on the cache before calling the HF model. """ def __init__(self, model: nn.Module, num_layers: int) -> None: @@ -125,7 +110,7 @@ def __init__(self, model: nn.Module, num_layers: int) -> None: @classmethod def from_pretrained(cls, model_name_or_path: str, **kwargs: Any) -> T5DecoderWrapper: - """Load full T5, wrap with sliding-window cache.""" + """Load full T5, wrap with static cache.""" full_model = T5ForConditionalGeneration.from_pretrained(model_name_or_path, **kwargs) num_layers = full_model.config.num_layers wrapper = cls(full_model, num_layers) @@ -137,11 +122,11 @@ def get_export_args(self, inputs: dict[str, torch.Tensor]) -> tuple[torch.Tensor return tuple(inputs.values()) def forward(self, *args: torch.Tensor) -> tuple[torch.Tensor, ...]: - """Run decoder with sliding-window KV cache. + """Run decoder with static KV cache. Positional args (order matches OnnxConfig.inputs): decoder_input_ids, encoder_hidden_states, attention_mask, - decoder_attention_mask, + decoder_attention_mask, cache_position, past_0_key, past_0_value, past_1_key, past_1_value, ... Returns: @@ -152,15 +137,14 @@ def forward(self, *args: torch.Tensor) -> tuple[torch.Tensor, ...]: encoder_hidden_states = args[1] attention_mask = args[2] decoder_attention_mask = args[3] - kv_start = 4 + cache_position = args[4] + kv_start = 5 - # Build WinMLSlidingWindowCache from input KV tensors. - # update() does Slice+Concat (not index_copy_/ScatterElements) — evicting - # the N oldest entries and appending the N new ones at the right. The - # incoming key/value states are captured for direct ONNX output - # (avoiding a scatter→gather round-trip in the graph). + # Build WinMLStaticCache from input KV tensors. The incoming key/value + # states are captured for direct ONNX output, avoiding a scatter->gather + # round-trip in the graph. max_cache_len = args[kv_start].size(2) - self_attn_cache = WinMLSlidingWindowCache(self.config, max_cache_len=max_cache_len) + self_attn_cache = WinMLStaticCache(self.config, max_cache_len=max_cache_len) self_attn_cache.early_initialization( batch_size=decoder_input_ids.size(0), num_heads=args[kv_start].size(1), @@ -175,13 +159,9 @@ def forward(self, *args: torch.Tensor) -> tuple[torch.Tensor, ...]: layer.keys = args[kv_start + i * 2] layer.values = args[kv_start + i * 2 + 1] - # Sliding window + single-token gen: the query is always at the - # rightmost slot. Constructing this constant inside forward traces it - # as a Constant node — downstream compute_bias and causal-mask subgraphs - # then constant-fold through ONNX optimization. - cache_position = torch.tensor( - [max_cache_len - 1], dtype=torch.int64, device=decoder_input_ids.device - ) + # T5 in transformers 5.x asks the cache for the current sequence length + # before it calls update(); make the traced ONNX input answer that query. + self_attn_cache.set_trace_position(cache_position) # EncoderDecoderCache is structurally required: T5Attention routes # self-attention → self_attention_cache, cross-attention → cross_attention_cache. @@ -246,18 +226,12 @@ def outputs(self) -> dict[str, dict[int, str]]: # noqa: D102 @register_onnx_overwrite("t5", "text2text-generation", library_name="transformers") class T5DecoderIOConfig(OnnxConfig): # type: ignore[misc] # optimum base is untyped - """ONNX config for T5 decoder with sliding-window KV cache. + """ONNX config for T5 decoder with static KV cache. Inputs: decoder_input_ids, encoder_hidden_states, attention_mask, - decoder_attention_mask, past_{i}_key/value + decoder_attention_mask, cache_position, past_{i}_key/value Outputs: logits, present_{i}_key/value - ``cache_position`` is *not* an input: ``T5DecoderWrapper.forward`` pins it - to ``[max_cache_len - 1]`` (rightmost buffer slot) as a Constant in the - graph. This couples the exported model to sliding-window semantics at - build time; see ``T5DecoderWrapper`` docstring for the static-cache - re-export path if needed. - Input past KV: full buffer [batch, heads, max_decode, d_kv]. Output present KV: new token only [batch, heads, 1, d_kv]. """ @@ -287,6 +261,7 @@ def inputs(self) -> dict[str, dict[int, str]]: # noqa: D102 "encoder_hidden_states": {0: "batch_size"}, "attention_mask": {0: "batch_size"}, "decoder_attention_mask": {0: "batch_size"}, + "cache_position": {}, } num_layers = self._normalized_config.num_layers for i in range(num_layers): @@ -347,22 +322,8 @@ class WinMLT5Model(WinMLEncoderDecoderModel): @classmethod def get_cache_class(cls) -> type: - """T5 defaults to ``WinMLSlidingWindowCache`` (Slice+Concat; no ScatterElements). - - Correctness with T5's learned relative position bias hinges on a single - invariant: ``cache_position`` is always the query's *buffer index*, not - its absolute sequence position. ``get_query_cache_position`` on each - cache class supplies the right value — ``[step]`` for static, - ``[max_cache_len-1]`` for sliding. Under that convention, - ``T5Attention.compute_bias`` computes ``memory_position - context_position - = j - (W-1)`` which gives correct relative distances regardless of - overflow, and HF's ``create_causal_mask`` (``kv_idx <= q_idx``) allows - every buffer slot while the 2D decoder mask selects the filled region. - - ``WinMLStaticCache`` remains fully supported — subclass ``WinMLT5Model`` - and override this method to get index_copy_ semantics instead. - """ - return WinMLSlidingWindowCache + """T5 defaults to ``WinMLStaticCache`` for the exported decoder graph.""" + return WinMLStaticCache @property def generation_config(self) -> GenerationConfig: # noqa: D102 diff --git a/src/winml/modelkit/models/winml/kv_cache.py b/src/winml/modelkit/models/winml/kv_cache.py index 1bcfa1fe4..c39f1ef1b 100644 --- a/src/winml/modelkit/models/winml/kv_cache.py +++ b/src/winml/modelkit/models/winml/kv_cache.py @@ -294,6 +294,18 @@ def update( return k_out, v_out + def get_seq_length(self, layer_idx: int = 0) -> int | torch.Tensor: + """Return filled length, or the traced query position during export.""" + if layer_idx >= len(self.layers): + return 0 + trace_position = getattr(self, "_trace_position", None) + if trace_position is not None: + flattened = trace_position.reshape(-1) + if flattened.numel() == 0: + return 0 + return flattened[0] + return self.step + def build_decoder_mask(self, max_len: int, num_new_tokens: int = 1) -> torch.Tensor: """Left-aligned: first ``step + num_new_tokens`` positions are 1.""" import torch diff --git a/tests/e2e/test_build_e2e.py b/tests/e2e/test_build_e2e.py index 1223ff181..e618570ad 100644 --- a/tests/e2e/test_build_e2e.py +++ b/tests/e2e/test_build_e2e.py @@ -53,6 +53,22 @@ @pytest.fixture(autouse=True) def _mock_resolve_device(): """Mock hardware detection to avoid failures in test environments.""" + from winml.modelkit.session import EPDeviceTarget + from winml.modelkit.utils.constants import normalize_ep_name + + cpu_target = EPDeviceTarget(ep="CPUExecutionProvider", device="cpu") + + def resolve_without_hardware_detection(target: EPDeviceTarget) -> EPDeviceTarget: + ep = target.ep or "auto" + device = (target.device or "auto").lower() + if ep == "auto": + return EPDeviceTarget(ep=cpu_target.ep, device=cpu_target.device, source=target.source) + return EPDeviceTarget( + ep=normalize_ep_name(ep), + device=cpu_target.device if device == "auto" else device, + source=target.source, + ) + with ( patch( "winml.modelkit.session.auto_detect_device", @@ -62,10 +78,24 @@ def _mock_resolve_device(): "winml.modelkit.sysinfo.hardware.get_available_devices", return_value=["cpu"], ), + patch( + "winml.modelkit.session.resolve_device", + side_effect=resolve_without_hardware_detection, + ), ): yield +def test_mock_resolve_device_preserves_explicit_target(): + """The e2e hardware mock must not rewrite explicit EP/device requests.""" + from winml.modelkit.session import EPDeviceTarget, resolve_device + + target = resolve_device(EPDeviceTarget(ep="qnn", device="npu")) + + assert target.ep == "QNNExecutionProvider" + assert target.device == "npu" + + # --------------------------------------------------------------------------- # Helpers # --------------------------------------------------------------------------- diff --git a/tests/unit/e2e/test_build_e2e_fixture_contract.py b/tests/unit/e2e/test_build_e2e_fixture_contract.py new file mode 100644 index 000000000..c21f5ea77 --- /dev/null +++ b/tests/unit/e2e/test_build_e2e_fixture_contract.py @@ -0,0 +1,45 @@ +# ------------------------------------------------------------------------- +# Copyright (c) Microsoft Corporation. All rights reserved. +# Licensed under the MIT License. +# -------------------------------------------------------------------------- +"""Contract tests for e2e build-test hardware mocking. + +The target e2e module imports optional image dependencies at module import +time, so these tests inspect the fixture source without importing it. +""" + +from __future__ import annotations + +import ast +from pathlib import Path + + +SOURCE = Path(__file__).resolve().parents[2] / "e2e" / "test_build_e2e.py" + + +def _fixture_tree() -> ast.FunctionDef: + tree = ast.parse(SOURCE.read_text(encoding="utf-8")) + for node in tree.body: + if isinstance(node, ast.FunctionDef) and node.name == "_mock_resolve_device": + return node + raise AssertionError("_mock_resolve_device fixture not found") + + +def test_resolve_device_mock_uses_side_effect_not_constant_cpu() -> None: + fixture = _fixture_tree() + patch_calls = [ + node + for node in ast.walk(fixture) + if isinstance(node, ast.Call) + and isinstance(node.func, ast.Name) + and node.func.id == "patch" + and node.args + and isinstance(node.args[0], ast.Constant) + and node.args[0].value == "winml.modelkit.session.resolve_device" + ] + assert patch_calls + + keyword_names = {keyword.arg for keyword in patch_calls[0].keywords} + + assert "side_effect" in keyword_names + assert "return_value" not in keyword_names diff --git a/tests/unit/eval/test_run_eval_script.py b/tests/unit/eval/test_run_eval_script.py index 773dd0e65..d1e8f5ec5 100644 --- a/tests/unit/eval/test_run_eval_script.py +++ b/tests/unit/eval/test_run_eval_script.py @@ -791,6 +791,40 @@ def test_fail(self, run_eval): assert run_eval.accuracy_status({"winml_eval_status": "FAIL"}) == "FAIL" +class TestQuantizedPrecisionClassifier: + @pytest.mark.parametrize( + ("precision", "expected"), + [ + ("auto", False), + ("fp16", False), + ("fp32", False), + ("int4", True), + ("int8", True), + ("int16", True), + ("w4a8", True), + ("w4a16", True), + ("w4a32", True), + ("w8a8", True), + ("w8a16", True), + ("w8a32", False), + ("w16a8", True), + ("w16a16", True), + ("w16a32", False), + ("w8a4", False), + ("w4a4", False), + ("w2a8", False), + ("garbage", False), + ("wXaY", False), + ("", False), + ("W8A16", True), + ("INT8", True), + ("w08a16", True), + ], + ) + def test_matches_precision_policy_cases(self, run_eval, precision, expected): + assert run_eval._is_quantized_precision(precision) is expected + + class TestRecipeConfigHelpers: """Eval-section detection, meta-config pick, and trust-remote-code gate.""" @@ -871,9 +905,7 @@ def test_non_npu_keeps_only_non_quantized_variants(self, run_eval, tmp_path): def test_non_npu_recipe_only_quantized_falls_back(self, run_eval, tmp_path): # A recipe with no non-quantized variant leaves nothing to run off-NPU, # so the model builds a single winml-config fallback. - self._make_single_recipe( - tmp_path, "microsoft_resnet-50", "image-classification", ["w8a16"] - ) + self._make_single_recipe(tmp_path, "microsoft_resnet-50", "image-classification", ["w8a16"]) entry = _entry() jobs = run_eval._build_jobs([entry], tmp_path, "cpu") assert len(jobs) == 1 @@ -923,9 +955,7 @@ def test_npu_skip_quant_ep_drops_quantized_recipe_variants(self, run_eval, tmp_p def test_npu_skip_quant_ep_recipe_only_quantized_falls_back(self, run_eval, tmp_path): # Dropping every quantized variant leaves nothing to build, so the model # goes through the single unquantized winml-config fallback. - self._make_single_recipe( - tmp_path, "microsoft_resnet-50", "image-classification", ["w8a16"] - ) + self._make_single_recipe(tmp_path, "microsoft_resnet-50", "image-classification", ["w8a16"]) entry = _entry() jobs = run_eval._build_jobs([entry], tmp_path, "npu", ep="vitisai") assert len(jobs) == 1 @@ -956,6 +986,20 @@ def test_non_npu_recipes_disabled_single_fallback(self, run_eval, tmp_path): assert jobs[0].variant is None +class TestExtractOnnxPath: + """``_extract_onnx_path`` understands Rich-wrapped final artifact output.""" + + def test_final_artifact_path_may_start_on_next_line(self, run_eval, tmp_path): + artifact = tmp_path / "mask_abc_model.onnx" + artifact.write_text("onnx", encoding="utf-8") + proc = { + "stdout": f"Build complete\n Final artifact: \n{artifact}\n Build config: x\n", + "stderr": "", + } + + assert run_eval._extract_onnx_path(proc, "acme/model", "fill-mask") == str(artifact) + + class TestRunRecipeBuild: """``_run_recipe_build`` builds authored configs with ``winml build -c --use-cache``.""" diff --git a/tests/unit/models/t5/test_static_cache_contract.py b/tests/unit/models/t5/test_static_cache_contract.py new file mode 100644 index 000000000..40ec4b0cf --- /dev/null +++ b/tests/unit/models/t5/test_static_cache_contract.py @@ -0,0 +1,131 @@ +# ------------------------------------------------------------------------- +# Copyright (c) Microsoft Corporation. All rights reserved. +# Licensed under the MIT License. +# -------------------------------------------------------------------------- +"""Static-cache contract tests for the T5 encoder-decoder wrapper. + +These tests intentionally inspect source structure so they can run in minimal +environments where torch wheels are unavailable. The end-to-end export tests +exercise the same contract with the real libraries on supported hosts. +""" + +from __future__ import annotations + +import ast +from pathlib import Path + + +REPO_ROOT = Path(__file__).resolve().parents[4] +T5_SOURCE = REPO_ROOT / "src" / "winml" / "modelkit" / "models" / "hf" / "t5.py" +KV_CACHE_SOURCE = REPO_ROOT / "src" / "winml" / "modelkit" / "models" / "winml" / "kv_cache.py" + + +def _tree(path: Path) -> ast.Module: + return ast.parse(path.read_text(encoding="utf-8")) + + +def _class(tree: ast.Module, name: str) -> ast.ClassDef: + for node in tree.body: + if isinstance(node, ast.ClassDef) and node.name == name: + return node + raise AssertionError(f"class {name} not found") + + +def _method(cls: ast.ClassDef, name: str) -> ast.FunctionDef: + for node in cls.body: + if isinstance(node, ast.FunctionDef) and node.name == name: + return node + raise AssertionError(f"{cls.name}.{name} not found") + + +def test_t5_decoder_declares_cache_position_input() -> None: + inputs = _method(_class(_tree(T5_SOURCE), "T5DecoderIOConfig"), "inputs") + + string_constants = { + node.value + for node in ast.walk(inputs) + if isinstance(node, ast.Constant) and isinstance(node.value, str) + } + + assert "cache_position" in string_constants + + +def test_t5_model_defaults_to_static_cache() -> None: + get_cache_class = _method(_class(_tree(T5_SOURCE), "WinMLT5Model"), "get_cache_class") + + returns = [node.value for node in ast.walk(get_cache_class) if isinstance(node, ast.Return)] + + assert any(isinstance(value, ast.Name) and value.id == "WinMLStaticCache" for value in returns) + + +def test_t5_decoder_forward_uses_static_cache_and_threads_trace_position() -> None: + forward = _method(_class(_tree(T5_SOURCE), "T5DecoderWrapper"), "forward") + + names = {node.id for node in ast.walk(forward) if isinstance(node, ast.Name)} + called_attrs = { + node.func.attr + for node in ast.walk(forward) + if isinstance(node, ast.Call) and isinstance(node.func, ast.Attribute) + } + + assert "WinMLStaticCache" in names + assert "set_trace_position" in called_attrs + + +def test_t5_decoder_forward_reads_cache_position_before_past_kv() -> None: + forward = _method(_class(_tree(T5_SOURCE), "T5DecoderWrapper"), "forward") + assignments = { + target.id: value + for node in ast.walk(forward) + if isinstance(node, ast.Assign) and len(node.targets) == 1 + for target, value in [(node.targets[0], node.value)] + if isinstance(target, ast.Name) + } + + cache_position = assignments["cache_position"] + kv_start = assignments["kv_start"] + + assert isinstance(cache_position, ast.Subscript) + assert isinstance(cache_position.value, ast.Name) + assert cache_position.value.id == "args" + assert isinstance(cache_position.slice, ast.Constant) + assert cache_position.slice.value == 4 + assert isinstance(kv_start, ast.Constant) + assert kv_start.value == 5 + + +def test_static_cache_trace_position_drives_hf_seq_length_queries() -> None: + static_cache = _class(_tree(KV_CACHE_SOURCE), "WinMLStaticCache") + get_seq_length = _method(static_cache, "get_seq_length") + + attrs = {node.attr for node in ast.walk(get_seq_length) if isinstance(node, ast.Attribute)} + string_constants = { + node.value + for node in ast.walk(get_seq_length) + if isinstance(node, ast.Constant) and isinstance(node.value, str) + } + + assert "_trace_position" in attrs | string_constants + assert "step" in attrs + + +def test_cache_base_normalizes_traced_dimensions_before_early_initialization() -> None: + tree = _tree(KV_CACHE_SOURCE) + cache_base = _class(tree, "WinMLCache") + early_initialization = _method(cache_base, "early_initialization") + module_functions = {node.name for node in tree.body if isinstance(node, ast.FunctionDef)} + + normalized_args = { + keyword.arg + for call in ast.walk(early_initialization) + if isinstance(call, ast.Call) + and isinstance(call.func, ast.Attribute) + and call.func.attr == "early_initialization" + for keyword in call.keywords + if isinstance(keyword.value, ast.Call) + and isinstance(keyword.value.func, ast.Name) + and keyword.value.func.id == "_as_static_dim" + } + + assert "_as_static_dim" in module_functions + assert normalized_args == {"batch_size", "num_heads", "head_dim"}