perf: expose FlexAttention loader override

[danieleb1861]

This PR turns TiRex2's existing FlexAttention implementation into a testable public configuration option by proposing an optional use_flex_attention keyword argument to load_model, while preserving existing defaults.
Performance, compilation, batching, and numerical trade offs are documented with a micro-benchmark.

Context

During large scale zero-shot evaluation of independent multivariate forecasting windows on a personal dataset, the cost of dense grouped variate attention grows with the total number of packed variates, even though interactions between unrelated timeseries groups are masked.

TiRex2 already provides a block-sparse FlexAttention path capable of exploiting this group structure. However, the backend is currently not selectable through the public load_model interface. Enabling it requires modifying checkpoint configuration or mutating internal model attributes after loading.

To my personal knowledge, this makes an existing performance capability difficult to discover, configure, test, and reproduce from user code.

Proposed interface

The contribution exposes the existing backend through an additive keyword-only argument:

from tirex2 import load_model
 
model = load_model(
    "NX-AI/TiRex-2",
    device="cuda",
    use_flex_attention=True,
)

The behavior is explicit:

• None, the default, preserves the checkpoint setting and current package behavior;
• True enables FlexAttention in every variate mixer;
• False explicitly forces dense attention.

Existing callers therefore maintain the same behavior without code or checkpoint changes.

Rationale

FlexAttention is deliberately exposed as an opt-in backend rather than enabled globally.

The measurements show that backend selection depends on workload size, as already known by Authors:

• dense attention remains faster for small grouped batches;
• FlexAttention becomes beneficial as the number of independent multivariate groups increases.

However, very large batches can become memory-bound before attention throughput is fully utilized and FlexAttention introduces first-use CUDA compilation overhead.

Keeping the existing default avoids changing latency, compilation behavior, or numerical reproducibility for current users.

Performance characterization

Benchmarks were run with on the released TiRex2 checkpoint with the following HW and SW setup:

• PyTorch 2.9.1+cu128;
• NVIDIA GeForce RTX 4070 Laptop GPU;
• context length 20;
• prediction length 15;
• one target and 16 past covariates per time-series window;
• test-time differencing and sign-flip augmentation disabled.

Runs with other context/prediction length combinations showed larger FlexAttention gains; the tables below are representative of a single configuration rather than a best-case result.

Measurement methodology

Two complementary measurements were used.

Variate-attention microbenchmark

• The model dimensions matched the released checkpoint: embedding dimension 512 and four attention heads.
• Each configuration used two warm-up executions followed by five timed executions.
• Latency was measured with CUDA events.
• Dense and FlexAttention layers used identical weights and inputs.

First-use FlexAttention compilation was excluded from warmed latency and evaluated separately. The reported values of the latency is the median of the five timed executions.

Forecasting benchmark

• The released TiRex2 checkpoint was loaded on CUDA in float32.
• Inputs were generated once and reused across backend comparisons.
• Test-time differencing and sign-flip augmentation were disabled.
• Each backend and batch-size configuration was executed once for warm-up before the measured pass.
• Elapsed wall time was measured with time.perf_counter().
• torch.cuda.synchronize() was called before stopping the timer.
• Model loading and first-use kernel compilation were excluded.
• Input transformation, GPU transfers, model execution, output transformation, and forecast reconstruction were included.
• Throughput was calculated as number_of_windows / elapsed_seconds.
• Peak allocated CUDA memory was collected with torch.cuda.max_memory_allocated() after resetting peak statistics.

The measurements were collected on a single NVIDIA GeForce RTX 4070 Laptop GPU without fixed GPU clocks. They should therefore be interpreted as workload-specific results from the user point of view rather than hardware-independent performances.

End-to-end forecasting throughput

Internal batch	Dense attention	FlexAttention	Observation
8	37.74 windows/s	26.21 windows/s	Dense is faster
32	35.09 windows/s	37.00 windows/s	Small FlexAttention benefit
64	32.89 windows/s	37.61 windows/s	Best observed FlexAttention result
128	9.22 windows/s	10.10 windows/s	Both paths become memory-bound

At batch size 64, FlexAttention improved end-to-end throughput by approximately 14.3% on this hardware and workload.

Variate-attention layer

For 128 independent groups with 17 variates each, corresponding to 2,176 packed variate rows:

Backend	Median layer latency
Dense attention	149.59 ms
FlexAttention	47.37 ms

This is a 3.16× reduction in variate-attention latency for that isolated workload.

The smaller end-to-end improvement is expected because TiRex2's recurrent time-mixing blocks must still process every target and covariate. This PR therefore addresses the grouped variate-attention component rather than claiming a universal model-level speedup, that was not benchmarked.
However, improvements from the proposed implementation can nonetheless be meaningful from a user perspective, as seen above.

Numerical behavior

Dense attention and FlexAttention are numerically close but not bit-identical.

On a 64-window batch of a held-out navigation forecasting benchmark from my own work:

• maximum absolute median-forecast difference: 3.25e-5;
• mean absolute difference: 2.38e-6;
• dense MAPE: 0.2426895%;
• FlexAttention MAPE: 0.2426810%;
• MAPE difference: −8.57e-6 percentage points.

The documentation explicitly recommends re-evaluating forecast metrics when changing attention backends in reproducible experiments.

Compatibility

The change is intentionally small and additive:

• no model weights are modified;
• no checkpoint migration is required;
• no tensor shapes, existing default or state-dict keys change;
• CPU and small-batch users remain on the current path unless they explicitly opt in.

Validation

The finalized change was checked with:

• the full package test suite: 61 passed, 16 skipped;
• Ruff lint;
• package wheel build;
• explicit tests for enabling and disabling the loader override;
• end-to-end dense/FlexAttention forecast comparison;
• CUDA latency and memory measurements across multiple batch sizes with a demo model configuration.

Scope

This PR does not introduce a new sparse-attention algorithm. It provides a supported public interface for TiRex2's existing FlexAttention implementation and documents the workload conditions under which it may be useful.

It also does not attempt to remove the fundamental cost of processing multiple covariates through the recurrent time-mixing stack.

Thank you for open-sourcing TiRex-2!
I'm using it in my research on shipboard energy management. Happy to contribute further as that work progresses.

[github-actions]

All contributors have signed the CLA ✍️ ✅
_{Posted by the CLA Assistant Lite bot.}

[danieleb1861]

I have read the CLA Document and I hereby sign the CLA