fix(export): accept traced 0-dim tensors as KV cache dimensions

src/winml/modelkit/models/winml/kv_cache.py

@@ -55,6 +55,17 @@
     from transformers.cache_utils import CacheLayerMixin
 
 
+def _as_static_dim(value: Any) -> Any:
+    """Normalize a (possibly traced) shape value to Python ``int``/``list[int]``.
+
+    ``Tensor.size(dim)`` returns a 0-dim tensor while TorchScript tracing, which
+    breaks callers that type-check for ``int``.
+    """
+    if isinstance(value, (list, tuple)):
+        return [int(item) for item in value]
+    return int(value)
+
+
 # =============================================================================
 # WinMLCache — common interface
 # =============================================================================
@@ -103,6 +114,32 @@ def set_trace_position(self, position: torch.Tensor) -> None:
         """Provide the position tensor when a model omits cache update kwargs."""
         self._trace_position = position
 
+    def early_initialization(
+        self,
+        batch_size: int,
+        num_heads: int | list[int],
+        head_dim: int | list[int],
+        dtype: torch.dtype,
+        device: torch.device,
+    ) -> None:
+        """Initialize all layers, accepting traced 0-dim tensors as dimensions.
+
+        Export wrappers derive ``batch_size``/``num_heads``/``head_dim`` from
+        the past-KV graph inputs via ``Tensor.size(dim)``. Under the
+        TorchScript tracer used by ``torch.onnx.export`` that call returns a
+        0-dim tensor instead of an ``int``, so the upstream implementation
+        skips its ``isinstance(..., int)`` broadcast and then calls ``len()``
+        on the value, which fails. These dimensions are static for an exported
+        graph, so normalize them back to Python ints before delegating.
+        """
+        super().early_initialization(
+            batch_size=_as_static_dim(batch_size),
+            num_heads=_as_static_dim(num_heads),
+            head_dim=_as_static_dim(head_dim),
+            dtype=dtype,
+            device=device,
+        )
+
     # ----- Interface for WinMLEncoderDecoderModel.forward -----
 
     @abstractmethod

tests/unit/export/test_io.py

@@ -1102,6 +1102,45 @@ def test_partial_chunk_left_padded(self) -> None:
         assert pos_ids[0].tolist() == [0, 0, 4, 5]
 
 
+# =============================================================================
+# WinMLCache.early_initialization — traced shape arguments
+# =============================================================================
+
+
+class TestWinMLCacheEarlyInitialization:
+    """``early_initialization`` accepts the 0-dim tensors emitted while tracing.
+
+    Export wrappers derive the cache dimensions from the past-KV graph inputs
+    with ``Tensor.size(dim)``. Under the TorchScript tracer used by
+    ``torch.onnx.export`` that returns a 0-dim tensor rather than an ``int``,
+    which upstream's ``isinstance(..., int)`` broadcast does not handle.
+    """
+
+    @pytest.mark.parametrize("cache_cls_name", ["WinMLStaticCache", "WinMLSlidingWindowCache"])
+    def test_zero_dim_tensor_dims_initialize_layers(self, cache_cls_name: str) -> None:
+        from transformers import PretrainedConfig
+
+        import winml.modelkit.models.winml as winml_models
+
+        cache_cls = getattr(winml_models, cache_cls_name)
+        config = PretrainedConfig(num_hidden_layers=2)
+        cache = cache_cls(config, max_cache_len=16)
+
+        cache.early_initialization(
+            batch_size=torch.tensor(1),
+            num_heads=torch.tensor(2),
+            head_dim=torch.tensor(8),
+            dtype=torch.float32,
+            device=torch.device("cpu"),
+        )
+
+        assert all(layer.is_initialized for layer in cache.layers)
+        keys = cache.layers[0].keys
+        assert keys.shape[0] == 1
+        assert keys.shape[1] == 2
+        assert keys.shape[3] == 8
+
+
 # =============================================================================
 # WinMLStaticCache.update — multi-dim index_put_ writes correct positions
 # =============================================================================