fix(perf): per-module benchmark works for ResNet (and similar HF models)

src/winml/modelkit/commands/perf.py

@@ -575,11 +575,21 @@ def _perf_modules(
 
     console.print(f"[dim]Found {len(module_configs)} {module_class} instances[/dim]")
 
-    # Instantiate parent with init weights (no pretrained download)
+    # Instantiate parent with init weights (no pretrained download).
+    # Submodule configs intentionally drop `loader.task`, so re-resolve the
+    # parent task from the model_id — the same path `generate_hf_build_config`
+    # used to compute module_path. Without this, models whose `architectures`
+    # field maps to a different task than `get_supported_tasks(model_type)[0]`
+    # instantiate the wrong parent class and `get_submodule()` raises
+    # AttributeError.
     model_type = module_configs[0].loader.model_type
     if not model_type:
         raise click.ClickException("module configs missing model_type")
-    parent_model = _instantiate_parent_model(model_type, task=task)
+
+    from ..loader import resolve_loader_config
+
+    parent_loader_cfg, _, _ = resolve_loader_config(model_id=hf_model, task=task)
+    parent_model = _instantiate_parent_model(model_type, task=parent_loader_cfg.task)
 
     all_results: list[dict[str, Any]] = []
     for i, cfg in enumerate(module_configs):

src/winml/modelkit/config/build.py

@@ -44,6 +44,7 @@
 
 import copy
 import hashlib
+import inspect
 import json
 import logging
 from dataclasses import dataclass, field
@@ -259,6 +260,9 @@ class SubmoduleInfo:
         output_shapes: Shape of each output tensor (e.g., [[1,16,64]])
         input_dtypes: Dtype of each input tensor (e.g., ["float32", "float32"])
         output_dtypes: Dtype of each output tensor (e.g., ["float32"])
+        input_names: Forward-arg names for each input (e.g., ["hidden_state"]
+            or ["pixel_values"]). Empty when hook capture didn't run; callers
+            then fall back to generic ``input_{i}`` names.
     """
 
     class_name: str
@@ -267,6 +271,7 @@ class SubmoduleInfo:
     output_shapes: list[list[int]]
     input_dtypes: list[str]
     output_dtypes: list[str]
+    input_names: list[str] = field(default_factory=list)
 
 
 # =============================================================================
@@ -814,10 +819,20 @@ def _build_submodule_config(
         - Inherited optim/compile from parent
         - Quant with task=None, model_name=None (RandomDataset fallback)
     """
-    # Build InputTensorSpec for EACH input tensor (not just the first)
+
+    # Build InputTensorSpec for EACH input tensor (not just the first).
+    # Use the submodule's actual forward-arg names so build_hf_model can
+    # call submodule(**kwargs) correctly — submodule forward args may be
+    # positional (e.g. `input`) or keyword (e.g. `hidden_state`). Fall back
+    # to generic input_{i} only when names were not discovered.
+    def _input_name(i: int) -> str:
+        if i < len(sub_info.input_names) and sub_info.input_names[i]:
+            return sub_info.input_names[i]
+        return f"input_{i}"
+
     input_tensors = [
         InputTensorSpec(
-            name=f"input_{i}",
+            name=_input_name(i),
             shape=tuple(shape),
             dtype=sub_info.input_dtypes[i] if i < len(sub_info.input_dtypes) else None,
         )
@@ -1081,12 +1096,14 @@ def _find_submodules_by_class(
     results = []
     for full_path, layer_info in torchinfo_modules:
         io_info = hook_data.get(full_path)
+        layer_input_names: list[str] = []
         if io_info and io_info.input_shapes:
             # Prefer hook-captured data (has complete multi-input info)
             layer_input_shapes = io_info.input_shapes
             layer_output_shapes = io_info.output_shapes
             layer_input_dtypes = io_info.input_dtypes
             layer_output_dtypes = io_info.output_dtypes
+            layer_input_names = io_info.input_names
         else:
             # Fall back to torchinfo data (single input only)
             layer_input_shapes = [layer_info.input_size] if layer_info.input_size else []
@@ -1102,6 +1119,16 @@ def _find_submodules_by_class(
             layer_input_dtypes = [param_dtype] * len(layer_input_shapes)
             layer_output_dtypes = [param_dtype] * len(layer_output_shapes)
 
+            # Without hook data, derive names from the forward signature so
+            # build_hf_model can invoke the submodule with the correct kwargs.
+            try:
+                sig = inspect.signature(layer_info.module.forward)
+                layer_input_names = [p.name for p in sig.parameters.values() if p.name != "self"][
+                    : len(layer_input_shapes)
+                ]
+            except (TypeError, ValueError):
+                layer_input_names = []
+
         results.append(
             SubmoduleInfo(
                 class_name=layer_info.class_name,
@@ -1110,6 +1137,7 @@ def _find_submodules_by_class(
                 output_shapes=layer_output_shapes,
                 input_dtypes=layer_input_dtypes,
                 output_dtypes=layer_output_dtypes,
+                input_names=layer_input_names,
             )
         )
 

tests/unit/commands/test_perf_module.py

@@ -101,6 +101,12 @@ def test_device_and_ep_forwarded_through_module_path(self, tmp_path: Path) -> No
         fake_session = MagicMock()
         fake_session.perf.side_effect = RuntimeError("test-skip-benchmark")
 
+        # _perf_modules calls resolve_loader_config(model_id=...) to recover the
+        # parent task (submodule configs strip it). Stub it so "fake/model" never
+        # hits the HF Hub.
+        fake_loader_cfg = MagicMock()
+        fake_loader_cfg.task = "fill-mask"
+
         with (
             patch(
                 "winml.modelkit.sysinfo.resolve_device",
@@ -110,6 +116,10 @@ def test_device_and_ep_forwarded_through_module_path(self, tmp_path: Path) -> No
                 "winml.modelkit.config.generate_hf_build_config",
                 return_value=[fake_cfg],
             ) as mock_gen,
+            patch(
+                "winml.modelkit.loader.resolve_loader_config",
+                return_value=(fake_loader_cfg, MagicMock(), MagicMock()),
+            ),
             patch(
                 "winml.modelkit.commands.build._instantiate_parent_model",
                 return_value=MagicMock(),

tests/unit/config/test_build.py

@@ -796,6 +796,117 @@ def test_empty_inputs(self, parent_config: WinMLBuildConfig) -> None:
         # Empty list is falsy, so input_tensors should be set to None
         assert result.export.input_tensors is None
 
+    def test_input_names_propagate(self, parent_config: WinMLBuildConfig) -> None:
+        """SubmoduleInfo.input_names propagate to InputTensorSpec.name."""
+        sub_info = SubmoduleInfo(
+            class_name="ResNetBottleNeckLayer",
+            module_path="encoder.stages.0.layers.0",
+            input_shapes=[[1, 64, 32, 32]],
+            output_shapes=[[1, 256, 32, 32]],
+            input_dtypes=["float32"],
+            output_dtypes=["float32"],
+            input_names=["hidden_state"],
+        )
+
+        result = _build_submodule_config(sub_info, parent_config)
+
+        assert result.export.input_tensors is not None
+        assert result.export.input_tensors[0].name == "hidden_state"
+
+    def test_input_names_fallback(self, parent_config: WinMLBuildConfig) -> None:
+        """Missing, short, or empty-string input_names fall back to input_{i}."""
+        # Case 1: empty input_names → input_0
+        sub_info_empty = SubmoduleInfo(
+            class_name="Conv2d",
+            module_path="encoder.conv",
+            input_shapes=[[1, 64, 32, 32]],
+            output_shapes=[[1, 128, 16, 16]],
+            input_dtypes=["float32"],
+            output_dtypes=["float32"],
+            input_names=[],
+        )
+        result_empty = _build_submodule_config(sub_info_empty, parent_config)
+        assert result_empty.export.input_tensors is not None
+        assert result_empty.export.input_tensors[0].name == "input_0"
+
+        # Case 2: input_names shorter than input_shapes → known name then fallback
+        sub_info_short = SubmoduleInfo(
+            class_name="CrossAttention",
+            module_path="decoder.cross_attn",
+            input_shapes=[[1, 16, 64], [1, 16, 64]],
+            output_shapes=[[1, 16, 64]],
+            input_dtypes=["float32", "float32"],
+            output_dtypes=["float32"],
+            input_names=["hidden_state"],
+        )
+        result_short = _build_submodule_config(sub_info_short, parent_config)
+        assert result_short.export.input_tensors is not None
+        assert result_short.export.input_tensors[0].name == "hidden_state"
+        assert result_short.export.input_tensors[1].name == "input_1"
+
+        # Case 3: empty-string entry → input_{i}
+        sub_info_blank = SubmoduleInfo(
+            class_name="Conv2d",
+            module_path="encoder.conv",
+            input_shapes=[[1, 64, 32, 32]],
+            output_shapes=[[1, 128, 16, 16]],
+            input_dtypes=["float32"],
+            output_dtypes=["float32"],
+            input_names=[""],
+        )
+        result_blank = _build_submodule_config(sub_info_blank, parent_config)
+        assert result_blank.export.input_tensors is not None
+        assert result_blank.export.input_tensors[0].name == "input_0"
+
+
+# =============================================================================
+# TestFindSubmodulesByClass - exercises the signature-fallback branch
+# =============================================================================
+
+
+class TestFindSubmodulesByClass:
+    """Tests for _find_submodules_by_class signature-fallback branch."""
+
+    def test_signature_fallback_when_hook_data_empty(self) -> None:
+        """Empty hook_data triggers inspect.signature fallback for input_names."""
+        import torch
+        from torch import nn
+
+        from winml.modelkit.config.build import _find_submodules_by_class
+
+        class SignatureFallbackSubmodule(nn.Module):
+            def __init__(self) -> None:
+                super().__init__()
+                self.linear = nn.Linear(8, 16)
+
+            def forward(self, hidden_state: torch.Tensor) -> torch.Tensor:
+                return self.linear(hidden_state)
+
+        class SignatureFallbackWrapper(nn.Module):
+            def __init__(self) -> None:
+                super().__init__()
+                self.layer = SignatureFallbackSubmodule()
+
+            def forward(self, hidden_state: torch.Tensor) -> torch.Tensor:
+                return self.layer(hidden_state)
+
+        model = SignatureFallbackWrapper()
+
+        # Force the fallback path by short-circuiting hook capture.
+        with patch(
+            "winml.modelkit.inspect.module_io_capture.capture_module_io",
+            return_value={},
+        ):
+            results = _find_submodules_by_class(
+                model,
+                "SignatureFallbackSubmodule",
+                input_shapes=[(1, 8)],
+                input_dtypes=["float32"],
+            )
+
+        assert len(results) == 1
+        assert results[0].input_names == ["hidden_state"]
+
 
 # =============================================================================
 # TestConfigCliOverride - CLI tests for --config flag