From 21ffc11de4c8752308aa497cebe4efd576f01dc8 Mon Sep 17 00:00:00 2001
From: rusty1s <matthias.fey@tu-dortmund.de>
Date: Fri, 16 Oct 2026 11:40:41 +0000
Subject: [PATCH 1/5] update

---
 sdm/models/kumo/timeseries/forecast/block.py  |  83 +++++++
 .../kumo/timeseries/forecast/encoder.py       | 209 ------------------
 sdm/models/kumo/timeseries/forecast/icl.py    |  87 ++++++++
 sdm/models/kumo/timeseries/forecast/model.py  |   6 +-
 .../kumo/timeseries/forecast/test_encoder.py  | 168 --------------
 .../kumo/timeseries/forecast/test_icl.py      |  37 ++++
 6 files changed, 211 insertions(+), 379 deletions(-)
 create mode 100644 sdm/models/kumo/timeseries/forecast/block.py
 delete mode 100644 sdm/models/kumo/timeseries/forecast/encoder.py
 create mode 100644 sdm/models/kumo/timeseries/forecast/icl.py
 delete mode 100644 test/models/kumo/timeseries/forecast/test_encoder.py
 create mode 100644 test/models/kumo/timeseries/forecast/test_icl.py

diff --git a/sdm/models/kumo/timeseries/forecast/block.py b/sdm/models/kumo/timeseries/forecast/block.py
new file mode 100644
index 000000000..e2bae299b
--- /dev/null
+++ b/sdm/models/kumo/timeseries/forecast/block.py
@@ -0,0 +1,83 @@
+# SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+
+# ruff: noqa: D101, D102
+
+import math
+from typing import Any
+
+import torch
+import torch.nn.functional as F
+from torch import Tensor
+from torch.nn import Linear
+
+from sdm.nn import RMSNorm, TransformerBlock
+
+
+class MLP(torch.nn.Module):
+    def __init__(
+        self,
+        channels: int,
+        hidden_channels: int,
+        device: torch.device | str | None = None,
+        dtype: torch.dtype | None = None,
+    ) -> None:
+        super().__init__()
+        factory_kwargs: dict[str, Any] = {"device": device, "dtype": dtype}
+        self.norm = RMSNorm(channels, eps=1e-6, **factory_kwargs)
+        self.wi_0 = Linear(channels, hidden_channels, False, **factory_kwargs)
+        self.wi_1 = Linear(channels, hidden_channels, False, **factory_kwargs)
+        self.wo = Linear(hidden_channels, channels, False, **factory_kwargs)
+
+    def forward(self, x: Tensor) -> Tensor:
+        x = self.norm(x)
+        x = F.gelu(self.wi_0(x), approximate="tanh") * self.wi_1(x)
+        return self.wo(x)
+
+
+class KumoForecastTransformerBlock(TransformerBlock):
+    def __init__(
+        self,
+        channels: int,
+        hidden_channels: int,
+        num_heads: int,
+        device: torch.device | str | None = None,
+        dtype: torch.dtype | None = None,
+    ) -> None:
+        factory_kwargs: dict[str, Any] = {"device": device, "dtype": dtype}
+        norm = RMSNorm(channels, eps=1e-6, **factory_kwargs)
+        super().__init__(
+            channels=channels,
+            num_query_heads=num_heads,
+            mlp=MLP(channels, hidden_channels, **factory_kwargs),
+            query_norm=norm,
+            key_value_norm=norm,
+            scale=1.0,
+            bias=False,
+            **factory_kwargs,
+        )
+        self.hidden_channels = hidden_channels
+        self.attn.out_lin.reset_parameters()
+
+    def peak_bytes_per_example(
+        self,
+        element_size: int,
+        query_length: int,
+        key_value_length: int | None = None,
+    ) -> int:
+        r""":meta private:"""  # noqa: D415
+        length = max(query_length, key_value_length or 0)
+        return math.ceil(11.25 * length * element_size * self.attn.q_dim)
+
+
+if __name__ == "__main__":
+    from sdm.testing.memory import benchmark_transformer_block_memory_peak
+
+    benchmark_transformer_block_memory_peak(
+        block=lambda channels, num_heads: KumoForecastTransformerBlock(
+            channels=channels,
+            hidden_channels=2816,
+            num_heads=num_heads,
+        ),
+        channels_and_heads=[(1024, 16)],
+    )
diff --git a/sdm/models/kumo/timeseries/forecast/encoder.py b/sdm/models/kumo/timeseries/forecast/encoder.py
deleted file mode 100644
index 505269d7c..000000000
--- a/sdm/models/kumo/timeseries/forecast/encoder.py
+++ /dev/null
@@ -1,209 +0,0 @@
-# SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved.
-# SPDX-License-Identifier: Apache-2.0
-
-import math
-from typing import Any
-
-import torch
-import torch.nn.functional as F
-from torch import Tensor
-from torch.nn import Dropout, Embedding, Linear, ModuleList
-
-from sdm.nn import SDPA, RMSNorm
-
-
-class _T5Block(torch.nn.Module):
-    def __init__(
-        self,
-        channels: int,
-        hidden_channels: int,
-        num_heads: int,
-        dropout: float,
-        eps: float,
-        device: torch.device | str | None,
-        dtype: torch.dtype | None,
-    ) -> None:
-        super().__init__()
-        factory_kwargs: dict[str, Any] = {"device": device, "dtype": dtype}
-        self.num_heads = num_heads
-        self.attention_norm = RMSNorm(channels, eps=eps, **factory_kwargs)
-        self.ffn_norm = RMSNorm(channels, eps=eps, **factory_kwargs)
-        self.sdpa = SDPA(
-            num_query_heads=num_heads,
-            scale=1.0,
-            dropout=dropout,
-        )
-        self.q = Linear(
-            in_features=channels,
-            out_features=channels,
-            bias=False,
-            **factory_kwargs,
-        )
-        self.k = Linear(
-            in_features=channels,
-            out_features=channels,
-            bias=False,
-            **factory_kwargs,
-        )
-        self.v = Linear(
-            in_features=channels,
-            out_features=channels,
-            bias=False,
-            **factory_kwargs,
-        )
-        self.o = Linear(
-            in_features=channels,
-            out_features=channels,
-            bias=False,
-            **factory_kwargs,
-        )
-        self.wi_0 = Linear(
-            in_features=channels,
-            out_features=hidden_channels,
-            bias=False,
-            **factory_kwargs,
-        )
-        self.wi_1 = Linear(
-            in_features=channels,
-            out_features=hidden_channels,
-            bias=False,
-            **factory_kwargs,
-        )
-        self.wo = Linear(
-            in_features=hidden_channels,
-            out_features=channels,
-            bias=False,
-            **factory_kwargs,
-        )
-        self.dropout = Dropout(dropout)
-
-    def forward(self, x: Tensor, bias: Tensor) -> Tensor:
-        normalized = self.attention_norm(x)
-        shape = (*x.shape[:-1], self.num_heads, -1)
-        query = self.q(normalized).view(shape)
-        key = self.k(normalized).view(shape)
-        value = self.v(normalized).view(shape)
-        # T5 uses unscaled dot products and shares position bias across layers.
-        attended = self.sdpa(
-            query=query,
-            key=key,
-            value=value,
-            attn_mask=bias.to(query.dtype),
-        )
-        attended = attended.flatten(-2)
-        x = x + self.dropout(self.o(attended))
-        normalized = self.ffn_norm(x)
-        hidden = F.gelu(self.wi_0(normalized), approximate="tanh")
-        hidden = hidden * self.wi_1(normalized)
-        return x + self.dropout(self.wo(self.dropout(hidden)))
-
-
-class T5Encoder(torch.nn.Module):
-    r"""Bidirectional encoder from the
-    `"Exploring the Limits of Transfer Learning with a Unified Text-to-Text
-    Transformer" <https://arxiv.org/abs/1910.10683>`_ paper.
-
-    Accepts continuous patch embeddings; no vocabulary embedding or decoder is
-    needed. Relative position bias is learned once and shared by all layers.
-
-    Args:
-        channels: Input and output embedding width.
-        hidden_channels: Feed-forward hidden width.
-        num_layers: Number of encoder blocks.
-        num_heads: Number of attention heads. Must divide ``channels``.
-        num_buckets: Number of bidirectional relative-position buckets.
-            Must be at least ``4``.
-        max_distance: Distance at which relative-position buckets saturate.
-            Must exceed ``num_buckets // 4``.
-        dropout: Attention, residual, and feed-forward dropout probability.
-        eps: RMS normalization epsilon.
-        device: The device.
-        dtype: The dtype.
-    """  # noqa: D205
-
-    def __init__(
-        self,
-        channels: int = 1024,
-        hidden_channels: int = 2816,
-        num_layers: int = 24,
-        num_heads: int = 16,
-        num_buckets: int = 32,
-        max_distance: int = 128,
-        dropout: float = 0.1,
-        eps: float = 1e-6,
-        device: torch.device | str | None = None,
-        dtype: torch.dtype | None = None,
-    ) -> None:
-        super().__init__()
-        if num_buckets < 4 or max_distance <= num_buckets // 4:
-            raise ValueError(
-                "Expected num_buckets >= 4 and max_distance > num_buckets // 4"
-            )
-        factory_kwargs: dict[str, Any] = {"device": device, "dtype": dtype}
-        self.num_buckets = num_buckets
-        self.max_distance = max_distance
-        self.relative_attention_bias = Embedding(
-            num_embeddings=num_buckets,
-            embedding_dim=num_heads,
-            **factory_kwargs,
-        )
-        self.layers = ModuleList(
-            [
-                _T5Block(
-                    channels=channels,
-                    hidden_channels=hidden_channels,
-                    num_heads=num_heads,
-                    dropout=dropout,
-                    eps=eps,
-                    **factory_kwargs,
-                )
-                for _ in range(num_layers)
-            ]
-        )
-        self.norm = RMSNorm(channels, eps=eps, **factory_kwargs)
-        self.dropout = Dropout(dropout)
-
-    def _position_bias(self, length: int, device: torch.device) -> Tensor:
-        position = torch.arange(length, device=device)
-        distance = position.unsqueeze(0) - position.unsqueeze(1)
-        half = self.num_buckets // 2
-        exact = half // 2
-        magnitude = distance.abs()
-        logarithmic = (
-            exact
-            + (
-                (magnitude.float().clamp_min(exact) / exact).log()
-                * ((half - exact) / math.log(self.max_distance / exact))
-            ).long()
-        )
-        bucket = torch.where(
-            magnitude < exact, magnitude, logarithmic.clamp_max(half - 1)
-        )
-        bucket = bucket + (distance > 0).long() * half
-        return self.relative_attention_bias(bucket).permute(2, 0, 1)
-
-    def forward(self, x: Tensor, mask: Tensor | None = None) -> Tensor:
-        """Encode continuous patch embeddings.
-
-        Args:
-            x: Embeddings with shape ``[..., P, C]``, where ``P`` is the
-                number of patches and ``C`` is the embedding width.
-            mask: Observed-patch mask with shape ``[..., P]``. Batch
-                dimensions may broadcast. False patches cannot be attended
-                to as keys.
-
-        Returns:
-            Encoded patches with shape ``[..., P, C]``.
-        """
-        bias = self._position_bias(x.size(-2), x.device)
-        if mask is not None:
-            padding = x.new_zeros((*mask.shape[:-1], 1, 1, mask.size(-1)))
-            padding = padding.masked_fill(
-                ~mask.bool().unsqueeze(-2).unsqueeze(-2),
-                -torch.inf,
-            )
-            bias = bias + padding
-        x = self.dropout(x)
-        for layer in self.layers:
-            x = layer(x, bias)
-        return self.dropout(self.norm(x))
diff --git a/sdm/models/kumo/timeseries/forecast/icl.py b/sdm/models/kumo/timeseries/forecast/icl.py
new file mode 100644
index 000000000..d40c70679
--- /dev/null
+++ b/sdm/models/kumo/timeseries/forecast/icl.py
@@ -0,0 +1,87 @@
+# SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+
+# ruff: noqa: D101, D102
+
+import math
+from typing import Any, Literal
+
+import torch
+from torch import Tensor
+from torch.nn import Embedding, ModuleList
+
+from sdm.models.kumo.timeseries.forecast.block import (
+    KumoForecastTransformerBlock,
+)
+from sdm.nn import RMSNorm
+
+
+class ICLBlock(torch.nn.Module):
+    def __init__(
+        self,
+        channels: int,
+        hidden_channels: int,
+        num_layers: int,
+        num_heads: int,
+        num_buckets: int,
+        max_distance: int,
+        device: torch.device | str | None = None,
+        dtype: torch.dtype | None = None,
+    ) -> None:
+        super().__init__()
+        factory_kwargs: dict[str, Any] = {"device": device, "dtype": dtype}
+
+        assert num_buckets >= 4
+        assert max_distance > num_buckets // 4
+
+        self.num_buckets = num_buckets
+        self.max_distance = max_distance
+        self.relative_attention_bias = Embedding(
+            num_embeddings=num_buckets,
+            embedding_dim=num_heads,
+            **factory_kwargs,
+        )
+        self.layers = ModuleList(
+            [
+                KumoForecastTransformerBlock(
+                    channels=channels,
+                    hidden_channels=hidden_channels,
+                    num_heads=num_heads,
+                    **factory_kwargs,
+                )
+                for _ in range(num_layers)
+            ]
+        )
+        self.norm = RMSNorm(channels, eps=1e-6, **factory_kwargs)
+
+    def get_position_bias(self, length: int, device: torch.device) -> Tensor:
+        num_buckets = self.num_buckets // 2  # Buckets per direction.
+        max_exact = num_buckets // 2
+        scale = num_buckets - max_exact
+        scale /= math.log(self.max_distance / max_exact)
+
+        pos = torch.arange(length, device=device)
+        offset = pos.unsqueeze(0) - pos.unsqueeze(1)  # [P, P]
+        dist = offset.abs()
+        log_dist = (dist.float().clamp(min=max_exact) / max_exact).log()
+        bucket = max_exact + (log_dist * scale).long()
+        bucket = bucket.clamp_max(num_buckets - 1)
+        bucket = torch.where(dist < max_exact, dist, bucket)
+        bucket = bucket + (offset > 0) * num_buckets
+        return self.relative_attention_bias(bucket).permute(2, 0, 1)
+
+    def forward(
+        self,
+        x: Tensor,  # [..., N, D]
+        *,
+        batch_size_limit: int | Literal["auto"] | None = "auto",
+    ) -> Tensor:  # [..., N, D]
+        bias = self.get_position_bias(length=x.size(-2), device=x.device)
+        for layer in self.layers:
+            x = layer(
+                query=x,
+                attn_mask=bias,
+                batch_size_limit=batch_size_limit,
+                out=None if torch.is_grad_enabled() else x,
+            )
+        return self.norm(x)
diff --git a/sdm/models/kumo/timeseries/forecast/model.py b/sdm/models/kumo/timeseries/forecast/model.py
index 79059c605..45a365473 100644
--- a/sdm/models/kumo/timeseries/forecast/model.py
+++ b/sdm/models/kumo/timeseries/forecast/model.py
@@ -16,7 +16,7 @@
 from sdm.models import ICLModel
 from sdm.models._huggingface import download_checkpoint
 from sdm.models.kumo.timeseries.forecast.attention import CrossChannelAttention
-from sdm.models.kumo.timeseries.forecast.encoder import T5Encoder
+from sdm.models.kumo.timeseries.forecast.icl import ICLBlock
 from sdm.models.kumo.timeseries.forecast.recipe import default_recipe
 from sdm.nn import RevIN, SinusoidalEmbedding
 from sdm.tensor.table import TableSchema
@@ -246,11 +246,13 @@ def __init__(
         torch.nn.init.orthogonal_(self.lin.weight, gain=1.41)
         self.pe = SinusoidalEmbedding()
 
-        self.encoder = T5Encoder(
+        self.encoder = ICLBlock(
             channels=channels,
             hidden_channels=hidden_channels,
             num_layers=num_layers,
             num_heads=num_heads,
+            num_buckets=32,
+            max_distance=128,
             **factory_kwargs,
         )
         self.cross_channel_attn: CrossChannelAttention | None = None
diff --git a/test/models/kumo/timeseries/forecast/test_encoder.py b/test/models/kumo/timeseries/forecast/test_encoder.py
deleted file mode 100644
index 599b4c103..000000000
--- a/test/models/kumo/timeseries/forecast/test_encoder.py
+++ /dev/null
@@ -1,168 +0,0 @@
-# SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved.
-# SPDX-License-Identifier: Apache-2.0
-
-import pytest
-import torch
-
-from sdm.models.kumo.timeseries.forecast.encoder import T5Encoder
-from sdm.testing import withCUDA
-
-
-@withCUDA
-@pytest.mark.parametrize("length", [9, 140])
-@pytest.mark.parametrize("dtype", [torch.float32, torch.bfloat16])
-@pytest.mark.parametrize(("num_buckets", "max_distance"), [(4, 2), (32, 128)])
-def test_encoder_matches_t5(
-    device: torch.device,
-    length: int,
-    dtype: torch.dtype,
-    num_buckets: int,
-    max_distance: int,
-) -> None:
-    transformers = pytest.importorskip("transformers")
-    reference = transformers.T5EncoderModel(
-        transformers.T5Config(
-            vocab_size=16,
-            d_model=32,
-            d_ff=48,
-            num_layers=2,
-            num_heads=4,
-            d_kv=8,
-            dropout_rate=0.0,
-            feed_forward_proj="gated-gelu",
-            relative_attention_num_buckets=num_buckets,
-            relative_attention_max_distance=max_distance,
-        )
-    ).eval()
-    torch.nn.Module.to(reference, device=device, dtype=dtype)
-    model = T5Encoder(
-        channels=32,
-        hidden_channels=48,
-        num_layers=2,
-        num_heads=4,
-        dropout=0.0,
-        dtype=dtype,
-        device=device,
-        num_buckets=num_buckets,
-        max_distance=max_distance,
-    ).eval()
-    model.relative_attention_bias.load_state_dict(
-        reference.encoder.block[0]
-        .layer[0]
-        .SelfAttention.relative_attention_bias.state_dict()
-    )
-    model.norm.load_state_dict(reference.encoder.final_layer_norm.state_dict())
-    for i, block in enumerate(reference.encoder.block):
-        model.get_submodule(f"layers.{i}.attention_norm").load_state_dict(
-            block.layer[0].layer_norm.state_dict()
-        )
-        model.get_submodule(f"layers.{i}.ffn_norm").load_state_dict(
-            block.layer[1].layer_norm.state_dict()
-        )
-        for name in ["q", "k", "v", "o"]:
-            model.get_submodule(f"layers.{i}.{name}").load_state_dict(
-                getattr(block.layer[0].SelfAttention, name).state_dict()
-            )
-        for name in ["wi_0", "wi_1", "wo"]:
-            model.get_submodule(f"layers.{i}.{name}").load_state_dict(
-                getattr(block.layer[1].DenseReluDense, name).state_dict()
-            )
-    x = torch.randn(2, length, 32, dtype=dtype, device=device)
-    mask = torch.ones(2, length, dtype=torch.bool, device=device)
-    mask[0, :3] = False
-    mask[1, -5:] = False
-
-    with torch.inference_mode():
-        expected = reference(
-            inputs_embeds=x, attention_mask=mask
-        ).last_hidden_state
-        actual = model(x, mask)
-    assert actual.dtype == dtype
-    if dtype == torch.bfloat16:
-        # Fused SDPA/GELU avoid intermediate bf16 rounding in the reference.
-        # Bound aggregate numerical error rather than near-zero elements.
-        relative_error = (
-            actual.float() - expected.float()
-        ).norm() / expected.float().norm()
-        assert relative_error < 0.02
-    else:
-        torch.testing.assert_close(actual, expected, atol=2e-5, rtol=2e-5)
-
-
-def test_masked_keys_cannot_change_observed_outputs() -> None:
-    model = T5Encoder(
-        channels=16,
-        hidden_channels=24,
-        num_layers=2,
-        num_heads=2,
-    ).eval()
-    x = torch.randn(2, 12, 16)
-    mask = torch.ones(2, 12, dtype=torch.bool)
-    mask[:, :4] = False
-    changed = x.clone()
-    changed[:, :4] += 100
-    with torch.inference_mode():
-        torch.testing.assert_close(
-            model(x, mask)[:, 4:], model(changed, mask)[:, 4:]
-        )
-
-
-def test_encoder_backward() -> None:
-    model = T5Encoder(
-        channels=16,
-        hidden_channels=24,
-        num_layers=2,
-        num_heads=2,
-        dropout=0.0,
-    )
-    x = torch.randn(2, 12, 16, requires_grad=True)
-    model(x).square().mean().backward()
-    assert x.grad is not None
-    assert x.grad.isfinite().all()
-    assert x.grad.abs().sum() > 0
-
-
-@pytest.mark.parametrize("dtype", [torch.float32, torch.bfloat16])
-def test_all_masked_keys_cannot_change_other_outputs(
-    dtype: torch.dtype,
-) -> None:
-    model = T5Encoder(
-        channels=16,
-        hidden_channels=24,
-        num_layers=2,
-        num_heads=2,
-        dropout=0.0,
-        dtype=dtype,
-    )
-    x = torch.randn(2, 12, 16, dtype=dtype, requires_grad=True)
-    mask = torch.zeros(2, 12, dtype=torch.bool)
-    changed = x.detach().clone()
-    changed[:, -1] += torch.arange(16, dtype=dtype)
-
-    actual = model(x, mask)
-    with torch.no_grad():
-        other = model(changed, mask)
-    assert actual.isfinite().all()
-    # Residual/FFN paths may change the final patch, but attention must not
-    # carry it to any other patch when every key is masked.
-    torch.testing.assert_close(actual[:, :-1], other[:, :-1], atol=0, rtol=0)
-    actual.float().square().sum().backward()
-    assert x.grad is not None
-    assert x.grad.isfinite().all()
-
-
-@pytest.mark.parametrize(
-    ("num_buckets", "max_distance"), [(2, 128), (32, 8), (32, 7)]
-)
-def test_invalid_relative_position_configuration(
-    num_buckets: int, max_distance: int
-) -> None:
-    with pytest.raises(ValueError, match=r"num_buckets.*max_distance"):
-        T5Encoder(
-            channels=16,
-            hidden_channels=24,
-            num_layers=1,
-            num_heads=2,
-            num_buckets=num_buckets,
-            max_distance=max_distance,
-        )
diff --git a/test/models/kumo/timeseries/forecast/test_icl.py b/test/models/kumo/timeseries/forecast/test_icl.py
new file mode 100644
index 000000000..13dbae5f3
--- /dev/null
+++ b/test/models/kumo/timeseries/forecast/test_icl.py
@@ -0,0 +1,37 @@
+# SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+
+import pytest
+import torch
+
+from sdm.models.kumo.timeseries.forecast.icl import ICLBlock
+from sdm.testing import withCUDA
+
+
+@withCUDA
+@pytest.mark.parametrize("length", [9, 140])
+@pytest.mark.parametrize(("num_buckets", "max_distance"), [(4, 2), (32, 128)])
+def test_icl_block(
+    device: torch.device,
+    length: int,
+    num_buckets: int,
+    max_distance: int,
+) -> None:
+    model = ICLBlock(
+        channels=32,
+        hidden_channels=48,
+        num_layers=2,
+        num_heads=4,
+        num_buckets=num_buckets,
+        max_distance=max_distance,
+        device=device,
+    )
+
+    x = torch.randn(2, length, 32, device=device)
+
+    with torch.no_grad():
+        out = model(x)
+
+    assert out.size() == (2, length, 32)
+    assert out.dtype == x.dtype
+    assert out.device == device

From f5ed2c150cd2338ef891195a9f3d852e2ada6a4e Mon Sep 17 00:00:00 2001
From: rusty1s <matthias.fey@tu-dortmund.de>
Date: Fri, 16 Oct 2026 11:57:59 +0000
Subject: [PATCH 2/5] update

---
 .../kumo/timeseries/forecast/attention.py     | 61 -------------------
 sdm/models/kumo/timeseries/forecast/cross.py  | 48 +++++++++++++++
 sdm/models/kumo/timeseries/forecast/model.py  | 14 ++---
 .../timeseries/forecast/test_attention.py     | 48 ---------------
 .../kumo/timeseries/forecast/test_cross.py    | 31 ++++++++++
 .../kumo/timeseries/forecast/test_model.py    | 17 ++++--
 6 files changed, 98 insertions(+), 121 deletions(-)
 delete mode 100644 sdm/models/kumo/timeseries/forecast/attention.py
 create mode 100644 sdm/models/kumo/timeseries/forecast/cross.py
 delete mode 100644 test/models/kumo/timeseries/forecast/test_attention.py
 create mode 100644 test/models/kumo/timeseries/forecast/test_cross.py

diff --git a/sdm/models/kumo/timeseries/forecast/attention.py b/sdm/models/kumo/timeseries/forecast/attention.py
deleted file mode 100644
index 6e8bc8d55..000000000
--- a/sdm/models/kumo/timeseries/forecast/attention.py
+++ /dev/null
@@ -1,61 +0,0 @@
-# SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved.
-# SPDX-License-Identifier: Apache-2.0
-
-from typing import Any
-
-import torch
-from torch import Tensor
-
-from sdm.nn import Attention
-
-
-class CrossChannelAttention(torch.nn.Module):
-    """Apply self-attention across variates at every patch position.
-
-    Args:
-        channels: Number of embedding channels.
-        num_heads: Number of attention heads.
-        dropout: Dropout probability.
-        device: The device.
-        dtype: The dtype.
-    """
-
-    def __init__(
-        self,
-        channels: int,
-        num_heads: int = 8,
-        dropout: float = 0.1,
-        device: torch.device | str | None = None,
-        dtype: torch.dtype | None = None,
-    ) -> None:
-        super().__init__()
-        factory_kwargs: dict[str, Any] = {"device": device, "dtype": dtype}
-        self.attn = Attention(
-            channels=channels,
-            num_query_heads=num_heads,
-            dropout=dropout,
-            **factory_kwargs,
-        )
-        # Preserve the upstream MultiheadAttention initialization.
-        torch.nn.init.xavier_uniform_(self.attn.qkv_lin.weight)
-        torch.nn.init.zeros_(self.attn.qkv_lin.bias)
-        self.attn.out_lin.reset_parameters()
-        torch.nn.init.zeros_(self.attn.out_lin.bias)
-        self.norm = torch.nn.LayerNorm(channels, **factory_kwargs)
-        self.dropout = torch.nn.Dropout(dropout)
-
-    def forward(self, x: Tensor) -> Tensor:
-        """Mix variates independently at every patch position.
-
-        Args:
-            x: Patch embeddings with shape ``[..., V, P, C]``, where ``V``
-                is the number of variates, ``P`` is the number of patches,
-                and ``C`` is the embedding width.
-
-        Returns:
-            Patch embeddings with shape ``[..., V, P, C]``.
-        """
-        x = x.transpose(-3, -2)  # [..., P, V, C]
-        attended = self.attn(x)
-        out = self.norm(x + self.dropout(attended))
-        return out.transpose(-3, -2)
diff --git a/sdm/models/kumo/timeseries/forecast/cross.py b/sdm/models/kumo/timeseries/forecast/cross.py
new file mode 100644
index 000000000..9a3509f5f
--- /dev/null
+++ b/sdm/models/kumo/timeseries/forecast/cross.py
@@ -0,0 +1,48 @@
+# SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+
+# ruff: noqa: D101, D102
+
+from typing import Any, Literal
+
+import torch
+from torch import Tensor
+
+from sdm.nn import TransformerBlock
+
+
+class CrossChannelAttention(torch.nn.Module):
+    def __init__(
+        self,
+        channels: int,
+        num_heads: int,
+        device: torch.device | str | None = None,
+        dtype: torch.dtype | None = None,
+    ) -> None:
+        super().__init__()
+        factory_kwargs: dict[str, Any] = {"device": device, "dtype": dtype}
+        self.block = TransformerBlock(
+            channels=channels,
+            num_query_heads=num_heads,
+            mlp=None,
+            **factory_kwargs,
+        )
+        torch.nn.init.xavier_uniform_(self.block.attn.qkv_lin.weight)
+        torch.nn.init.zeros_(self.block.attn.qkv_lin.bias)
+        self.block.attn.out_lin.reset_parameters()
+        torch.nn.init.zeros_(self.block.attn.out_lin.bias)
+        self.norm = torch.nn.LayerNorm(channels, **factory_kwargs)
+
+    def forward(
+        self,
+        x: Tensor,  # [..., C, N, D]
+        *,
+        batch_size_limit: int | Literal["auto"] | None = "auto",
+    ) -> Tensor:  # [..., C, N, D]
+        x = x.transpose(-3, -2)  # [..., P, V, C]
+        x = self.block(
+            query=x,
+            batch_size_limit=batch_size_limit,
+            out=None if torch.is_grad_enabled() else x,
+        ).transpose(-3, -2)
+        return self.norm(x)
diff --git a/sdm/models/kumo/timeseries/forecast/model.py b/sdm/models/kumo/timeseries/forecast/model.py
index 45a365473..84889c10d 100644
--- a/sdm/models/kumo/timeseries/forecast/model.py
+++ b/sdm/models/kumo/timeseries/forecast/model.py
@@ -15,7 +15,7 @@
 from sdm.cache import Cache
 from sdm.models import ICLModel
 from sdm.models._huggingface import download_checkpoint
-from sdm.models.kumo.timeseries.forecast.attention import CrossChannelAttention
+from sdm.models.kumo.timeseries.forecast.cross import CrossChannelAttention
 from sdm.models.kumo.timeseries.forecast.icl import ICLBlock
 from sdm.models.kumo.timeseries.forecast.recipe import default_recipe
 from sdm.nn import RevIN, SinusoidalEmbedding
@@ -246,7 +246,7 @@ def __init__(
         torch.nn.init.orthogonal_(self.lin.weight, gain=1.41)
         self.pe = SinusoidalEmbedding()
 
-        self.encoder = ICLBlock(
+        self.icl = ICLBlock(
             channels=channels,
             hidden_channels=hidden_channels,
             num_layers=num_layers,
@@ -255,9 +255,9 @@ def __init__(
             max_distance=128,
             **factory_kwargs,
         )
-        self.cross_channel_attn: CrossChannelAttention | None = None
+        self.cross_channel: CrossChannelAttention | None = None
         if cross_channel:
-            self.cross_channel_attn = CrossChannelAttention(
+            self.cross_channel = CrossChannelAttention(
                 channels=channels,
                 num_heads=cross_channel_heads,
                 **factory_kwargs,
@@ -280,9 +280,9 @@ def forward(
         x = x.unflatten(-1, (-1, self.patch_size))  # [..., C, N, P]
         x = self.pe(self.lin(x))  # [..., C, N, D]
 
-        x = self.encoder(x)
-        if self.cross_channel_attn is not None:
-            x = self.cross_channel_attn(x)
+        x = self.icl(x)  # [..., C, N, D]
+        if self.cross_channel is not None:
+            x = self.cross_channel(x)  # [..., C, N, D]
 
         x = self.head(x.flatten(-2))  # [..., C, prediction_size]
         return self.revin.inverse(x, mean, scale)
diff --git a/test/models/kumo/timeseries/forecast/test_attention.py b/test/models/kumo/timeseries/forecast/test_attention.py
deleted file mode 100644
index 3620de281..000000000
--- a/test/models/kumo/timeseries/forecast/test_attention.py
+++ /dev/null
@@ -1,48 +0,0 @@
-# SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved.
-# SPDX-License-Identifier: Apache-2.0
-
-import pytest
-import torch
-
-from sdm.models.kumo.timeseries.forecast.attention import CrossChannelAttention
-from sdm.testing import withCUDA
-
-
-@withCUDA
-@pytest.mark.parametrize("batch_shape", [(), (2,), (2, 3)])
-def test_cross_channel_attention(
-    device: torch.device,
-    batch_shape: tuple[int, ...],
-) -> None:
-    reference = torch.nn.MultiheadAttention(
-        embed_dim=8, num_heads=2, batch_first=True, device=device
-    ).eval()
-    norm = torch.nn.LayerNorm(8, device=device)
-    model = CrossChannelAttention(
-        channels=8, num_heads=2, device=device
-    ).eval()
-    with torch.no_grad():
-        model.attn.qkv_lin.weight.copy_(reference.in_proj_weight)
-        model.attn.qkv_lin.bias.copy_(reference.in_proj_bias)
-        model.attn.out_lin.load_state_dict(reference.out_proj.state_dict())
-        model.norm.load_state_dict(norm.state_dict())
-    x = torch.randn(*batch_shape, 3, 4, 8, device=device)
-    flattened = x.transpose(-3, -2).reshape(-1, 3, 8)
-    with torch.inference_mode():
-        attended, _ = reference(flattened, flattened, flattened)
-        expected = norm(flattened + attended)
-        expected = expected.reshape(*batch_shape, 4, 3, 8).transpose(-3, -2)
-        actual = model(x)
-        torch.testing.assert_close(actual, expected, atol=1e-6, rtol=1e-5)
-        permutation = torch.tensor([2, 0, 1], device=device)
-        torch.testing.assert_close(
-            model(x[..., permutation, :, :]),
-            actual[..., permutation, :, :],
-            atol=1e-6,
-            rtol=1e-5,
-        )
-        changed = x.clone()
-        changed[..., 0, :] += 10
-        torch.testing.assert_close(
-            model(changed)[..., 1:, :], actual[..., 1:, :]
-        )
diff --git a/test/models/kumo/timeseries/forecast/test_cross.py b/test/models/kumo/timeseries/forecast/test_cross.py
new file mode 100644
index 000000000..75866969f
--- /dev/null
+++ b/test/models/kumo/timeseries/forecast/test_cross.py
@@ -0,0 +1,31 @@
+# SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+
+
+import pytest
+import torch
+
+from sdm.models.kumo.timeseries.forecast.cross import CrossChannelAttention
+from sdm.testing import withCUDA
+
+
+@withCUDA
+@pytest.mark.parametrize("batch_shape", [(), (2,), (2, 3)])
+def test_cross_channel_attention(
+    batch_shape: tuple[int, ...],
+    device: torch.device,
+) -> None:
+    model = CrossChannelAttention(
+        channels=8,
+        num_heads=2,
+        device=device,
+    )
+
+    x = torch.randn(*batch_shape, 3, 4, 8, device=device)
+
+    with torch.no_grad():
+        out = model(x)
+
+    assert out.size() == (*batch_shape, 3, 4, 8)
+    assert out.dtype == x.dtype
+    assert out.device == device
diff --git a/test/models/kumo/timeseries/forecast/test_model.py b/test/models/kumo/timeseries/forecast/test_model.py
index 967afae2a..434a885cb 100644
--- a/test/models/kumo/timeseries/forecast/test_model.py
+++ b/test/models/kumo/timeseries/forecast/test_model.py
@@ -9,10 +9,16 @@
 from sdm import TableTensor
 from sdm.models import KumoForecast
 from sdm.models.kumo.timeseries.forecast import model as forecasting
+from sdm.testing import withCUDA
 
 
+@withCUDA
 @pytest.mark.parametrize("cross_channel", [False, True])
-def test_forward(monkeypatch: pytest.MonkeyPatch, cross_channel: bool) -> None:
+def test_forward(
+    monkeypatch: pytest.MonkeyPatch,
+    cross_channel: bool,
+    device: torch.device,
+) -> None:
     monkeypatch.setattr(
         forecasting,
         "_KumoForecast",
@@ -25,23 +31,24 @@ def test_forward(monkeypatch: pytest.MonkeyPatch, cross_channel: bool) -> None:
             cross_channel_heads=4,
         ),
     )
-    model = KumoForecast(cross_channel, pretrained=False)
+    model = KumoForecast(cross_channel, pretrained=False, device=device)
 
     # Past and past-and-future covariates:
     x_context = TableTensor.from_tensor(
-        torch.randn(512, 2),
+        torch.randn(512, 2, device=device),
         columns=["x1", "x2"],
     )
-    x_query = TableTensor.from_tensor(torch.randn(3, 0))
+    x_query = TableTensor.from_tensor(torch.randn(3, 0, device=device))
 
     # Target variates:
     y_context = TableTensor.from_tensor(
-        torch.randn(512, 3),
+        torch.randn(512, 3, device=device),
         columns=["y1", "y2", "y3"],
     )
 
     out = model(x_context, y_context, x_query)
     assert out.size() == (3, 3)
+    assert out.device == device
     assert out.numerical.isfinite().all()
     assert "y1" in out.columns["numerical"]
     assert "y2" in out.columns["numerical"]

From 2edbb96029ee4967d74b3873192236c86dcdf5ad Mon Sep 17 00:00:00 2001
From: rusty1s <matthias.fey@tu-dortmund.de>
Date: Fri, 16 Oct 2026 12:12:03 +0000
Subject: [PATCH 3/5] update

---
 sdm/models/kumo/timeseries/forecast/ckpt.py  | 59 ++++++++++++++++++++
 sdm/models/kumo/timeseries/forecast/model.py |  4 +-
 2 files changed, 62 insertions(+), 1 deletion(-)
 create mode 100644 sdm/models/kumo/timeseries/forecast/ckpt.py

diff --git a/sdm/models/kumo/timeseries/forecast/ckpt.py b/sdm/models/kumo/timeseries/forecast/ckpt.py
new file mode 100644
index 000000000..cda65f060
--- /dev/null
+++ b/sdm/models/kumo/timeseries/forecast/ckpt.py
@@ -0,0 +1,59 @@
+# SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+
+import torch
+from torch import Tensor
+
+
+def remap_ckpt(ckpt: dict[str, Tensor]) -> dict[str, Tensor]:  # noqa: D103
+    out = {
+        "lin.weight": ckpt["patch_embedding.value_embedding.weight"],
+        "icl.relative_attention_bias.weight": ckpt[
+            "encoder.block.0.layer.0.SelfAttention.relative_attention_bias.weight"
+        ],
+        "icl.norm.weight": ckpt["encoder.final_layer_norm.weight"],
+        "head.weight": ckpt["head.linear.weight"],
+        "head.bias": ckpt["head.linear.bias"],
+    }
+    layers = sorted(
+        {
+            key.split(".")[2]
+            for key in ckpt
+            if key.startswith("encoder.block.")
+        },
+        key=int,
+    )
+    for layer in layers:
+        source = f"encoder.block.{layer}.layer."
+        target = f"icl.layers.{layer}."
+        norm = ckpt[source + "0.layer_norm.weight"]
+        out[target + "query_norm.weight"] = norm
+        out[target + "key_value_norm.weight"] = norm
+        out[target + "attn.qkv_lin.weight"] = torch.cat(
+            [
+                ckpt[source + f"0.SelfAttention.{name}.weight"]
+                for name in "qkv"
+            ],
+            dim=0,
+        )
+        out[target + "attn.out_lin.weight"] = ckpt[
+            source + "0.SelfAttention.o.weight"
+        ]
+        out[target + "mlp.norm.weight"] = ckpt[source + "1.layer_norm.weight"]
+        for name in ("wi_0", "wi_1", "wo"):
+            out[target + f"mlp.{name}.weight"] = ckpt[
+                source + f"1.DenseReluDense.{name}.weight"
+            ]
+
+    if "cross_channel_attn.attn.in_proj_weight" in ckpt:
+        for old, new in (
+            ("attn.in_proj_weight", "block.attn.qkv_lin.weight"),
+            ("attn.in_proj_bias", "block.attn.qkv_lin.bias"),
+            ("attn.out_proj.weight", "block.attn.out_lin.weight"),
+            ("attn.out_proj.bias", "block.attn.out_lin.bias"),
+            ("norm.weight", "norm.weight"),
+            ("norm.bias", "norm.bias"),
+        ):
+            out["cross_channel." + new] = ckpt["cross_channel_attn." + old]
+
+    return out
diff --git a/sdm/models/kumo/timeseries/forecast/model.py b/sdm/models/kumo/timeseries/forecast/model.py
index 84889c10d..7f6b8e075 100644
--- a/sdm/models/kumo/timeseries/forecast/model.py
+++ b/sdm/models/kumo/timeseries/forecast/model.py
@@ -15,6 +15,7 @@
 from sdm.cache import Cache
 from sdm.models import ICLModel
 from sdm.models._huggingface import download_checkpoint
+from sdm.models.kumo.timeseries.forecast.ckpt import remap_ckpt
 from sdm.models.kumo.timeseries.forecast.cross import CrossChannelAttention
 from sdm.models.kumo.timeseries.forecast.icl import ICLBlock
 from sdm.models.kumo.timeseries.forecast.recipe import default_recipe
@@ -100,7 +101,8 @@ def _load_from_pretrained(
             else "moment_head_512_6hr.pt",
             revision="abff20a58834638b28227ff4ab934f26206e4b09",
         )
-        torch.load(checkpoint, map_location=device, weights_only=True)
+        ckpt = torch.load(checkpoint, map_location=device, weights_only=True)
+        self.model.load_state_dict(remap_ckpt(ckpt), assign=True)
 
         return self
 

From 315a035dbe257e29e4c02256acf0a7aa0a84a4ca Mon Sep 17 00:00:00 2001
From: rusty1s <matthias.fey@tu-dortmund.de>
Date: Fri, 16 Oct 2026 12:33:34 +0000
Subject: [PATCH 4/5] update

---
 sdm/models/kumo/timeseries/forecast/model.py | 11 +++++++++++
 1 file changed, 11 insertions(+)

diff --git a/sdm/models/kumo/timeseries/forecast/model.py b/sdm/models/kumo/timeseries/forecast/model.py
index 7f6b8e075..60a9ada38 100644
--- a/sdm/models/kumo/timeseries/forecast/model.py
+++ b/sdm/models/kumo/timeseries/forecast/model.py
@@ -198,6 +198,17 @@ def _forward(
             y_schema = y_context.schema
             x = y_context.numerical
             if x_context_only is not None:
+                if (
+                    x_context_only.size(-1) > 0
+                    and self.model.cross_channel is None
+                ):
+                    warn_once(
+                        key="kumo-forecast-historical-covariate",
+                        message=(
+                            "Historical covariates have no effect when "
+                            "'cross_channel' is disabled."
+                        ),
+                    )
                 x = torch.cat([x, x_context_only.numerical], dim=-1)
 
             x = x[..., -self.model.context_size :, :].transpose(-1, -2)

From d5bd45b45977a05decd64d696377684c64d9d73d Mon Sep 17 00:00:00 2001
From: rusty1s <matthias.fey@tu-dortmund.de>
Date: Fri, 16 Oct 2026 12:48:37 +0000
Subject: [PATCH 5/5] update

---
 sdm/models/kumo/timeseries/forecast/model.py | 7 ++-----
 1 file changed, 2 insertions(+), 5 deletions(-)

diff --git a/sdm/models/kumo/timeseries/forecast/model.py b/sdm/models/kumo/timeseries/forecast/model.py
index 60a9ada38..b312ef67c 100644
--- a/sdm/models/kumo/timeseries/forecast/model.py
+++ b/sdm/models/kumo/timeseries/forecast/model.py
@@ -197,11 +197,8 @@ def _forward(
             assert y_context is not None
             y_schema = y_context.schema
             x = y_context.numerical
-            if x_context_only is not None:
-                if (
-                    x_context_only.size(-1) > 0
-                    and self.model.cross_channel is None
-                ):
+            if x_context_only is not None and x_context_only.size(-1) > 0:
+                if self.model.cross_channel is None:
                     warn_once(
                         key="kumo-forecast-historical-covariate",
                         message=(