feat(xpu_graph): batch sizes, per batch size routing, and Inductor kernels

.claude/skills/torch-tweak-assistant/SKILL.md

@@ -71,6 +71,7 @@ Do NOT activate this skill when:
 
 ### XPU Graph Backend
 - **Best for**: replaying a module that was recorded once during tuning
+- **Batch sizes**: recordings for 1, 2, 4, and every multiple of 8 up to `max_recorded_batch_size` (256); other sizes are padded
 - **Rejection**: failed, empty, or incorrect recordings are dropped for that module and input
 
 ### OpenVINO Backend

ChangeLog.md

@@ -11,6 +11,12 @@ SPDX-License-Identifier: Apache-2.0
 - fix: a tuned module that a pipeline holds under several paths is saved once and restored as one module.
 - feat: `save` warns about tuned modules it does not save, e.g. a module the pipeline holds in a plain list.
 - fix: a passthrough module without parameters keeps its inputs on their device instead of moving them to CPU.
+- feat (breaking): a module without a strategy is tuned with `tt.default_strategy(device)` instead of `FirstWinsStrategy`.
+  Tuning measures every backend and keeps the fastest, so it takes longer.
+- feat: `tt.default_strategy(device)` returns `HighestThroughputStrategy` over every backend that runs on the device.
+- feat (breaking): `FirstWinsStrategy` and `HighestThroughputStrategy` without `backends` use `default_backends(device)` for the tuning device.
+  `HighestThroughputStrategy()` now also measures OpenVINO, and neither strategy tries XPU graph on CPU.
+  Tuning on a device other than XPU or CPU without `backends` raises `ValueError`.
 - feat (breaking): `inspect` measures module depth in executed modules.
   Containers such as `nn.ModuleList` and modules that never run no longer add a level.
   `min_depth` is renamed to `depth`.

docs/api/backends/batch_size_router_backend.md

@@ -0,0 +1,13 @@
+<!--
+Copyright (c) 2026 Intel Corporation
+SPDX-License-Identifier: Apache-2.0
+-->
+
+# Batch Size Router Backend API
+
+## BatchSizeRouterBackend
+
+::: torch_tweak.torch.backend.BatchSizeRouterBackend
+    options:
+      show_root_heading: true
+      heading_level: 3

docs/guides/backends/xpu_graph_backend.md

@@ -10,11 +10,57 @@ Recording is not repeated for each inference call.
 
 ## Overview
 
-- **One recording**: warmup and capture run in `build`, while the module is being tuned
+- **Recording**: warmup and capture run in `build`, while the module is being tuned
 - **Replay**: inference copies new tensor values into the recorded inputs and calls `replay`
 - **Rejection**: a failed recording, an empty recording, or outputs that disagree with eager drop this backend for that module and input
-- **Static shapes**: a batch axis or another dynamic axis rejects the backend. Replay cannot change tensor shapes
-- **Fallback**: an inference call whose inputs differ from the recorded call in shape, dtype, device, or a non-tensor value runs the original module
+- **Dynamic axes**: an axis that changes without following the batch size, such as sequence length, rejects the backend
+- **Fallback**: an inference call whose inputs differ from the recordings in shape, dtype, device, or a non-tensor value runs the original module
+
+## Batch Sizes
+
+A graph spec with static shapes records one graph.
+A graph spec with a batch axis records graphs for batch sizes 1, 2, 4, and then every multiple of 8, up to the largest batch size seen during tuning.
+`max_recorded_batch_size` caps that size. The default is 256, which records 35 graphs.
+
+A call pads its batch with zeros up to the smallest recorded size that holds it, replays that graph, and returns only the real rows.
+A batch larger than the largest recorded size replays in chunks of that size.
+Padding costs compute. A batch of 3 runs the graph recorded for 4, and a batch of 17 runs the graph recorded for 24.
+
+Padding is only correct when every row of the batch is computed on its own.
+Tuning checks that: it runs eager once on the largest recorded batch size.
+Every graph then replays a prefix of that batch, full and partly padded with zeros and with copies of a real row, and is compared with the matching rows of eager.
+The backend is rejected when:
+
+- an output tensor has no batch axis, for example `x.sum()`
+- one row of the output depends on another row, for example `x - x.mean(dim=0)`
+- the module rearranges the rows of its batch, for example `torch.cat([y, y])`
+
+All graphs of one module share one memory pool and one input buffer sized for the largest recorded graph.
+Peak memory follows the largest recorded graph.
+Tuning time grows with the number of graphs: each one is warmed up, recorded, and replayed by the check.
+An out-of-memory error while recording a size skips that size and keeps the smaller graphs.
+
+```python
+from torch_tweak.torch.backend import XPUGraphBackend, XPUGraphBackendConfig
+
+backend = XPUGraphBackend(XPUGraphBackendConfig(max_recorded_batch_size=64))
+```
+
+## Inductor Kernels
+
+`compile_with_inductor=True` records the kernels Torch Inductor generates instead of the eager kernels.
+Replay then runs the fused Inductor kernels without launching each one from Python.
+On CUDA, `torch.compile(mode="reduce-overhead")` does the same with CUDA graphs. That mode has no effect on XPU.
+
+Compilation runs during the warmup before each recording.
+Inductor compiles about three times for all recorded sizes: once for the largest size, once with a dynamic batch axis, and once for batch size 1.
+Replay is still checked against the eager module.
+
+```python
+from torch_tweak.torch.backend import XPUGraphBackend, XPUGraphBackendConfig
+
+backend = XPUGraphBackend(XPUGraphBackendConfig(compile_with_inductor=True))
+```
 
 ## Quick Start
 
@@ -25,8 +71,11 @@ from torch_tweak.torch.tune_strategy import HighestThroughputStrategy
 strategy = HighestThroughputStrategy(backends=[XPUGraphBackend()])
 ```
 
-The backend is first in the default `FirstWinsStrategy` list, before Torch Inductor.
-It is also in the default `HighestThroughputStrategy` list.
+On XPU the backend is first in `default_backends(device)`, before Torch Inductor.
+`FirstWinsStrategy`, `HighestThroughputStrategy`, and `tt.default_strategy(device)` use that list when no backends are given.
+On CPU the list leaves this backend out.
+An XPU graph often wins at small batch sizes and loses at large ones.
+`HighestThroughputStrategy` then sends small batches to it and large batches to the faster backend.
 
 PyTorch records a native SYCL graph when it was built with a SYCL compiler of 2026.1 or newer, except on Intel Data Center GPU Max Series.
 An older build records the legacy SYCL command graph.

docs/guides/tune_strategies/tune_strategies.md

@@ -18,6 +18,21 @@ Torch Tweak provides three built-in strategies:
 - **FirstWinsStrategy**: Tries backends in order, uses the first that succeeds
 - **HighestThroughputStrategy**: Profiles all backends, selects the fastest
 
+## Default Strategy
+
+A module without a strategy is tuned with `tt.default_strategy(device)`.
+It is `HighestThroughputStrategy` over `default_backends(device)`, every backend that runs on the device and keeps the results of the model:
+
+| Device | Backends |
+|---|---|
+| XPU | XPU graph, Torch Inductor, OpenVINO, Torch Eager |
+| CPU | Torch Inductor, OpenVINO, Torch Eager |
+
+TorchAO is left out because quantization changes the results.
+Every module gets its own strategy instance, built for the device it is tuned on.
+`FirstWinsStrategy` and `HighestThroughputStrategy` without a backend list use the same `default_backends(device)`.
+Without a strategy, tuning on any other device raises `ValueError`.
+
 ## Why Backends Can Fail
 
 Not every backend can successfully tune every model. Each backend relies on a different compilation or export technology, and each has its own limitations:
@@ -82,7 +97,7 @@ tt.tune(model, input_data)
 
 ## FirstWinsStrategy
 
-Tries backends in priority order and returns the first one that successfully builds and validates. If a backend fails, the strategy moves on to the next candidate instead of aborting. If all backends fail, the error is caught and the original model is used as-is. The default order is XPU graph, Torch Inductor, OpenVINO, then eager. Put a backend earlier only when you want that one to win as soon as it builds.
+Tries backends in priority order and returns the first one that successfully builds and validates. If a backend fails, the strategy moves on to the next candidate instead of aborting. If all backends fail, the error is caught and the original model is used as-is. Without a backend list it tries `default_backends(device)`: XPU graph, Torch Inductor, OpenVINO, then eager, with XPU graph skipped on CPU. Put a backend earlier only when you want that one to win as soon as it builds.
 
 ### Usage
 
@@ -181,6 +196,23 @@ tt.tune(model, input_data)
 3. Selects backend with highest throughput
 4. Returns best performing backend
 
+### Different Winners per Batch Size
+
+With a batch axis, every backend is profiled on every power of two from 1 up to the max batch size.
+The best backend is the one with the highest throughput at any batch size, usually the largest one.
+Another backend takes a range of batch sizes when it is at least 10% faster there.
+The strategy then returns a `BatchSizeRouterBackend` that sends each call to the backend tuned for its batch size.
+A batch size between two profiled ones goes to the backend of the larger one.
+
+For example, an XPU graph launches a small batch faster and Torch Inductor computes a large batch faster:
+
+```
+BatchSizeRouterBackend(1-16: XPUGraphBackend(), 17+: TorchInductorBackend())
+```
+
+Every routed backend stays in memory.
+A checkpoint stores each of them.
+
 ### When to Use
 
 ✅ **Good for**:

docs/learn/backends.md

@@ -35,7 +35,8 @@ Guide: [Torch Inductor Backend](../guides/backends/torch_inductor_backend.md).
 
 ## XPU Graph Backend
 
-The XPU graph backend records the module once while tuning and replays that recording for later inference.
+The XPU graph backend records the module while tuning and replays that recording for later inference.
+A module with a batch axis gets recordings for batch sizes 1, 2, 4, and every multiple of 8, and other batch sizes are padded.
 PyTorch chooses native or legacy recording. This backend uses whichever one that build records.
 A recording that fails, or that disagrees with eager, is rejected for that module and input.
 

docs/learn/core_functionalities.md

@@ -118,7 +118,7 @@ strategy = FirstWinsStrategy(
 )
 ```
 
-That order is also the default when you construct `FirstWinsStrategy` with no backend list. The first backend that builds and validates is kept.
+`FirstWinsStrategy` with no backend list tries `default_backends(device)` in that order. On CPU it skips XPU graph. The first backend that builds and validates is kept.
 
 ## OneBackendStrategy
 
@@ -132,7 +132,8 @@ strategy = OneBackendStrategy(backend=OpenVINOBackend())
 
 ## HighestThroughputStrategy
 
-Profiles all compatible backends and selects the fastest. Use this when maximum throughput matters and you can afford longer tuning time. `HighestThroughputStrategy()` with no backend list compares Torch Inductor, XPU graph, and eager.
+Profiles all compatible backends and selects the fastest. Use this when maximum throughput matters and you can afford longer tuning time. `HighestThroughputStrategy()` with no backend list compares `default_backends(device)`: XPU graph, Torch Inductor, OpenVINO, and eager on XPU, and the same without XPU graph on CPU.
+When different backends win at different batch sizes, it returns a `BatchSizeRouterBackend` that sends each call to the backend tuned for its batch size.
 
 ```python
 from torch_tweak.torch.tune_strategy import HighestThroughputStrategy

examples/Minimal/demo.py

@@ -23,7 +23,6 @@ def main() -> None:
     model = nn.Sequential(nn.Linear(32, 32), nn.ReLU(), nn.Linear(32, 8)).eval().to(device)
     sample = torch.randn(32, device=device)
 
-    # One batch size keeps every tensor shape static, so the XPU graph can be recorded.
     strategy = tt.HighestThroughputStrategy(
         backends=[
             TorchInductorBackend(),

properdocs.yml

@@ -117,6 +117,7 @@ nav:
                     -   Torch Inductor Backend: api/backends/torch_inductor_backend.md
                     -   XPU Graph Backend: api/backends/xpu_graph_backend.md
                     -   Torch Eager Backend: api/backends/torch_eager_backend.md
+                    -   Batch Size Router Backend: api/backends/batch_size_router_backend.md
             -   Tuning Strategies API:
                     -   FirstWinsStrategy: api/tune_strategies/first_wins_strategy.md
                     -   OneBackendStrategy: api/tune_strategies/one_backend_strategy.md

tests/unit/torch/backend/test_batch_size_router_backend.py

@@ -0,0 +1,151 @@
+# Copyright (c) 2026 Intel Corporation
+# SPDX-License-Identifier: Apache-2.0
+
+import time
+from pathlib import Path
+from typing import Any
+
+import pytest
+import torch
+
+from tests.toy_models.torch_models import ToyTorchModel
+from torch_tweak.torch.backend import Backend, BatchSizeRouterBackend
+from torch_tweak.torch.backend.backend import BackendState
+from torch_tweak.torch.task.profiling import NumStepsMeasuringStopStrategy
+from torch_tweak.torch.task.profiling.profiling_stop_strategy import AllSamplesProfilingStopStrategy
+from torch_tweak.torch.tune_strategy.highest_throughput_strategy import HighestThroughputStrategy
+
+
+class TimedBackend(Backend):
+    """Sleeps per_call plus per_row times the batch size, then runs the module."""
+
+    requires_original_module = True
+
+    def __init__(self, label: str = "", per_call: float = 0.0, per_row: float = 0.0):
+        super().__init__()
+        self.label, self.per_call, self.per_row = label, per_call, per_row
+        self.calls: list[int] = []
+        self._module: Any = None
+
+    def key(self) -> str:
+        return f"TimedBackend_{self.label}"
+
+    def describe(self) -> str:
+        return f"TimedBackend({self.label})"
+
+    def _build(self, module: Any, graph_spec: Any, data: list[Any], cache_dir: Path) -> Backend:
+        self._module = module
+        return self
+
+    def _activate(self):
+        pass
+
+    def _deactivate(self):
+        pass
+
+    def _deploy(self):
+        pass
+
+    def _infer(self, *args, **kwargs):
+        batch_size = args[0].shape[0]
+        self.calls.append(batch_size)
+        time.sleep(self.per_call + self.per_row * batch_size)
+        return self._module(*args, **kwargs)
+
+    def to_dict(self):
+        return {"type": "TimedBackend", "label": self.label}
+
+    @classmethod
+    def from_dict(cls, module, state_dict: dict):
+        backend = cls(label=state_dict["label"])
+        backend._module = module
+        backend.state = BackendState.CHECKPOINT_LOADED
+        return backend
+
+
+def built(backend: TimedBackend, module, torch_device, tmp_path) -> TimedBackend:
+    sample = ((torch.rand(1, 4, device=torch_device),), {})
+    backend.build(module, None, [sample], torch_device, tmp_path)
+    return backend
+
+
+@pytest.fixture
+def model(torch_device):
+    return ToyTorchModel().eval().to(torch_device)
+
+
+@pytest.fixture
+def graph_spec(model, torch_device):
+    return model.graph_spec(batch_sizes=[1, 16], device=torch_device)
+
+
+def test_routes_send_each_batch_size_to_its_backend(model, graph_spec, torch_device, tmp_path):
+    small = built(TimedBackend("small"), model, torch_device, tmp_path)
+    large = built(TimedBackend("large"), model, torch_device, tmp_path)
+    router = BatchSizeRouterBackend([(4, small), (None, large)], graph_spec.input_spec)
+
+    assert router.is_active
+    for batch_size in (1, 4, 5, 64):
+        router.infer(model.samples([batch_size], device=torch_device)[0][0][0])
+    assert small.calls == [1, 4]
+    assert large.calls == [5, 64]
+    assert router.describe() == "BatchSizeRouterBackend(1-4: TimedBackend(small), 5+: TimedBackend(large))"
+    assert router.requires_original_module
+
+
+def test_last_route_must_take_every_larger_batch_size(model, graph_spec, torch_device, tmp_path):
+    backend = built(TimedBackend(), model, torch_device, tmp_path)
+    with pytest.raises(ValueError, match="every larger batch size"):
+        BatchSizeRouterBackend([(4, backend)], graph_spec.input_spec)
+
+
+def test_serialization(model, graph_spec, torch_device, tmp_path):
+    small = built(TimedBackend("small"), model, torch_device, tmp_path)
+    large = built(TimedBackend("large"), model, torch_device, tmp_path)
+    router = BatchSizeRouterBackend([(4, small), (None, large)], graph_spec.input_spec)
+
+    loaded = BatchSizeRouterBackend.from_dict(model, router.to_dict())
+
+    assert loaded.state == BackendState.CHECKPOINT_LOADED
+    assert loaded.describe() == router.describe()
+    loaded.deploy(torch_device)
+    sample = model.samples([2], device=torch_device)[0][0][0]
+    torch.testing.assert_close(loaded.infer(sample), model(sample))
+    assert loaded.backend_for(2).calls == [2]
+
+
+def tune(backends, model, graph_spec, torch_device, tmp_path) -> Backend:
+    strategy = HighestThroughputStrategy(
+        backends=backends,
+        measurement_stop_strategy=NumStepsMeasuringStopStrategy(num_steps=3),
+        profiling_stop_strategy=AllSamplesProfilingStopStrategy(),
+    ).enable_find_max_batch_size(False)
+    strategy.enable_correctness_check(False)
+    sample = ((model.sample(device=torch_device).unsqueeze(0),), {})
+    return strategy.tune(model, "test", graph_spec, [sample], torch_device, tmp_path)
+
+
+def test_strategy_routes_small_batches_to_the_backend_faster_there(model, graph_spec, torch_device, tmp_path):
+    # Per call: per_row 1 ms per row against a fixed 6 ms. Rows are faster up to batch size 4.
+    rows = TimedBackend("rows", per_row=1e-3)
+    fixed = TimedBackend("fixed", per_call=6e-3)
+
+    tuned = tune([rows, fixed], model, graph_spec, torch_device, tmp_path)
+
+    assert isinstance(tuned, BatchSizeRouterBackend)
+    assert [(largest, backend.describe()) for largest, backend in tuned.routes] == [
+        (4, "TimedBackend(rows)"),
+        (None, "TimedBackend(fixed)"),
+    ]
+
+
+def test_strategy_keeps_one_backend_when_it_wins_every_batch_size(model, graph_spec, torch_device, tmp_path):
+    tuned = tune(
+        [TimedBackend("slow", per_call=5e-3), TimedBackend("fast", per_call=1e-3)],
+        model,
+        graph_spec,
+        torch_device,
+        tmp_path,
+    )
+
+    assert tuned.describe() == "TimedBackend(fast)"