perf: expose FlexAttention loader override

README.md

@@ -139,6 +139,22 @@ fig.show()
 ```
 ![output of plot_multivariate function visualizing context and forecast of multivariate input](/resources/multivariate-prediction.png)
 
+### Optional FlexAttention for large multivariate batches
+
+The default dense attention backend avoids compilation overhead and is suitable for small batches. For large
+CUDA batches containing many independent multivariate series, opt into the block-sparse FlexAttention backend
+when loading the checkpoint:
+
+```python
+model = load_model("NX-AI/TiRex-2", device="cuda", use_flex_attention=True)
+forecasts = model.forecast(timeseries, prediction_length=64, batch_size=64)
+```
+
+FlexAttention compiles its CUDA kernel on first use and can be slower for small batches, so benchmark both the
+backend and `batch_size` on the target GPU. Leave `use_flex_attention` unset to preserve the checkpoint setting
+and package default, or pass `False` to force dense attention. The two CUDA kernels are numerically close but
+not bit-identical; re-evaluate forecast metrics when changing backends in reproducible benchmarks.
+
 
 
 ### Benchmarking

src/tirex2/base.py

@@ -51,6 +51,7 @@ def load_model(
     device: str = "cuda",
     *,
     hf_kwargs: dict[str, Any] | None = None,
+    use_flex_attention: bool | None = None,
 ) -> ForecastModel:
     """Load an inference-ready :class:`TiRex2` from a checkpoint directory or HF repo.
 
@@ -66,6 +67,12 @@ def load_model(
     hf_kwargs : dict, optional
         Extra keyword arguments forwarded to ``snapshot_download`` for Hugging
         Face paths, e.g. ``{"revision": "main", "local_files_only": True}``.
+    use_flex_attention : bool, optional
+        Override every variate mixer's checkpoint setting. ``True`` enables
+        block-sparse FlexAttention, which can reduce the cost of large grouped
+        multivariate batches on CUDA but adds first-call compilation overhead.
+        ``False`` forces dense attention. Leave as ``None`` to preserve the
+        checkpoint configuration and package defaults.
 
     Returns
     -------
@@ -89,6 +96,9 @@ def load_model(
         config: dict[str, Any] = yaml.safe_load(f)
 
     config["device"] = device
+    if use_flex_attention is not None:
+        for template in config["stack_config"]["templates"].values():
+            template["variate_mixer"]["use_flex_attention"] = use_flex_attention
     model = TiRex2(**config)
 
     checkpoint = torch.load(weights_file, map_location="cpu", weights_only=True)

test/test_tirex2_instantiation.py

@@ -84,6 +84,34 @@ def test_load_model_device_overrides_checkpoint_config(tmp_path, small_model_kwa
     assert {block.config.time_mixer.device for block in loaded.model.stack} == {"cpu"}
 
 
+@pytest.mark.parametrize(
+    "checkpoint_value, override",
+    [(False, True), (True, False)],
+    ids=["enable", "disable"],
+)
+def test_load_model_can_override_flex_attention(
+    tmp_path,
+    small_model_kwargs,
+    require_flex_attention,
+    checkpoint_value,
+    override,
+):
+    config = small_model_kwargs("cpu")
+    model = TiRex2(**config)
+    for template in config["stack_config"]["templates"].values():
+        template["variate_mixer"]["use_flex_attention"] = checkpoint_value
+
+    with (tmp_path / CONFIG_FILENAME).open("w") as f:
+        yaml.safe_dump(config, f)
+    torch.save(model.state_dict(), tmp_path / CKPT_FILENAME)
+
+    loaded = load_model(str(tmp_path), device="cpu", use_flex_attention=override)
+
+    assert all(block.config.variate_mixer.use_flex_attention is override for block in loaded.model.stack)
+    assert all(block.variate_mixer.use_flex_attention is override for block in loaded.model.stack)
+    assert all(block.variate_mixer.attn.use_flex_attention is override for block in loaded.model.stack)
+
+
 def test_tirex2_init_can_opt_into_matmul_precision(small_model_kwargs, monkeypatch):
     config = small_model_kwargs("cpu")
     precision_calls = []