From e7233bc10e4a4ade4386b3dac1eb93b352846d33 Mon Sep 17 00:00:00 2001
From: Ard <88337265+aw471@users.noreply.github.com>
Date: Wed, 21 Oct 2026 09:43:34 -0700
Subject: [PATCH 1/6] Add `KumoAnomaly` reconstruction and scoring

Co-authored-by: Aditi Gautam <adgautam@nvidia.com>
---
 sdm/models/kumo/timeseries/anomaly/model.py   | 189 +++++++++++++++++-
 .../kumo/timeseries/anomaly/test_model.py     |  88 ++++++++
 2 files changed, 274 insertions(+), 3 deletions(-)
 create mode 100644 test/models/kumo/timeseries/anomaly/test_model.py

diff --git a/sdm/models/kumo/timeseries/anomaly/model.py b/sdm/models/kumo/timeseries/anomaly/model.py
index 46a3534a7..51178bc97 100644
--- a/sdm/models/kumo/timeseries/anomaly/model.py
+++ b/sdm/models/kumo/timeseries/anomaly/model.py
@@ -5,16 +5,25 @@
 
 from __future__ import annotations
 
+import copy
 from collections.abc import Sequence
-from typing import ClassVar
+from functools import partial
+from typing import Any, ClassVar, cast
 
 import torch
+import torch.nn.functional as F
 from torch import Tensor
 
 from sdm import Recipe, Stype, TableTensor
+from sdm._inference import inference_mode
+from sdm._warnings import warn_once
 from sdm.models import Callback, Model
 from sdm.models._huggingface import download_checkpoint
+from sdm.models.kumo.timeseries.anomaly.denoiser import DiffusionDenoiser
+from sdm.models.kumo.timeseries.anomaly.embedding import SideEmbedding
 from sdm.models.kumo.timeseries.anomaly.recipe import default_recipe
+from sdm.models.kumo.timeseries.anomaly.sampling import DPMSolver
+from sdm.models.kumo.timeseries.anomaly.schedule import DiffusionSchedule
 
 
 class KumoAnomaly(Model):
@@ -96,7 +105,7 @@ def forward(
         Args:
             x: The feature tensor with shape ``[..., R, D]`` with ``R`` rows
                 and ``D`` columns.
-            recipe: The custom recipe for pre- and post-processing.
+            recipe: The custom recipe for pre-processing.
             num_samples: The number of reconstruction samples per complementary
                 masking strategy. Reconstructions are combined using the
                 sample median.
@@ -111,13 +120,187 @@ def forward(
             absolute reconstruction error across preprocessed features at each
             time step.
         """
-        raise NotImplementedError
+        callbacks = () if callbacks is None else callbacks
+        requires_grad = self.training
+        requires_grad |= any(callback.requires_grad for callback in callbacks)
+
+        param = next(self.parameters())
+        if not isinstance(x, TableTensor):
+            x = TableTensor.from_tensor(x)
+        x = cast(TableTensor, x.to(param.device))
+
+        recipe = (
+            self.default_recipe() if recipe is None else copy.deepcopy(recipe)
+        )
+        with (
+            torch.amp.autocast(param.device.type, enabled=False),
+            inference_mode("no_grad" if requires_grad else "inference"),
+        ):
+            x = recipe.features.fit_transform(x, generator=generator)
+
+        for callback in callbacks:
+            x, _ = callback.on_query_preprocessing_end(self, x, None)
+
+        invalid = x.active_stypes - self.supported_feature_stypes - {Stype.id}
+        if len(invalid) > 0:
+            stypes = ", ".join(f"{str(stype)!r}" for stype in invalid)
+            warn_once(
+                key="model-unsupported-feature-stypes",
+                message=(
+                    f"{self.__class__.__name__!r} received unsupported "
+                    f"feature stypes {stypes}. Columns with unsupported "
+                    f"feature stypes will not be consumed by the model."
+                ),
+            )
+        if x.numerical.size(-1) > self.model.num_features:
+            raise ValueError(
+                f"{self.__class__.__name__!r} supports at most "
+                f"{self.model.num_features} features after pre-processing "
+                f"(got {x.numerical.size(-1)})"
+            )
+
+        with inference_mode("grad" if requires_grad else "inference"):
+            out = self.model(
+                x.numerical.to(param.dtype),
+                num_samples=num_samples,
+                num_steps=num_steps,
+                generator=generator,
+            )
+            out = TableTensor(
+                columns={Stype.numerical: ["score"]},
+                numerical=out.unsqueeze(-1),
+            )
+            for callback in callbacks:
+                out = callback.on_model_forward_end(self, out)
+        return cast(TableTensor, out.to(x.dtype))
 
 
 class _KumoAnomaly(torch.nn.Module):
     def __init__(
         self,
+        num_features: int = 40,
+        time_channels: int = 256,
+        feature_channels: int = 64,
+        channels: int = 128,
+        num_layers: int = 6,
+        num_heads: int = 8,
+        embedding_channels: int = 256,
+        num_steps: int = 500,
+        window_size: int = 100,
+        num_blocks: int = 4,
+        batch_size: int = 32,
         device: torch.device | str | None = None,
         dtype: torch.dtype | None = None,
     ) -> None:
         super().__init__()
+        factory_kwargs: dict[str, Any] = {"device": device, "dtype": dtype}
+
+        self.num_features = num_features
+        self.num_steps = num_steps
+        self.window_size = window_size
+        self.num_blocks = num_blocks
+        self.batch_size = batch_size
+
+        self.side_embedding = SideEmbedding(
+            num_features=num_features,
+            time_channels=time_channels,
+            feature_channels=feature_channels,
+            **factory_kwargs,
+        )
+        self.denoiser = DiffusionDenoiser(
+            channels=channels,
+            side_channels=time_channels + feature_channels,
+            num_steps=num_steps,
+            num_layers=num_layers,
+            num_heads=num_heads,
+            embedding_channels=embedding_channels,
+            **factory_kwargs,
+        )
+
+    def forward(
+        self,
+        x: Tensor,  # [..., R, D]
+        num_samples: int,
+        num_steps: int,
+        generator: torch.Generator | None,
+    ) -> Tensor:  # [..., R]
+
+        *size, length, features = x.size()
+        x = x.reshape(-1, length, features)
+        x = F.pad(x, (0, self.num_features - features))
+
+        # Pad the last window by repeating the last row:
+        padding = -length % self.window_size
+        x = torch.cat([x, x[:, -1:].expand(-1, padding, -1)], dim=1)
+        windows = x.reshape(-1, self.window_size, self.num_features)
+
+        schedule = DiffusionSchedule(
+            num_steps=self.num_steps,
+            beta_start=0.0001,
+            beta_end=0.01,
+            schedule="linear",
+            device="cpu",
+        )
+        sampler = DPMSolver(schedule.beta.float(), num_steps).to(x.device)
+
+        out = torch.cat(
+            [
+                self.reconstruct(batch, sampler, num_samples, generator)
+                for batch in windows.split(self.batch_size)
+            ]
+        )
+        out = out.reshape(x.size())[:, :length, :features]
+        out = (out - x[:, :length, :features]).abs().mean(dim=-1)
+        return out.reshape(*size, length)
+
+    def reconstruct(
+        self,
+        x: Tensor,  # [B, L, K]
+        sampler: DPMSolver,
+        num_samples: int,
+        generator: torch.Generator | None,
+    ) -> Tensor:  # [B, L, K]
+
+        x = x.transpose(1, 2)  # [B, K, L]
+        batch, _, length = x.size()
+        mean = x.mean(dim=(1, 2), keepdim=True)
+        std = x.std(dim=(1, 2), keepdim=True) + 1e-5
+        x = ((x - mean) / std).clamp(-10, 10)
+
+        time = torch.arange(length, device=x.device).expand(batch, -1)
+        blocks = time // (length // self.num_blocks)
+
+        out = x.new_zeros((num_samples, *x.size()))
+        for strategy in range(2):
+            mask = ((blocks % 2) == strategy).to(x.dtype)
+            mask = mask.unsqueeze(1).expand_as(x)  # [B, K, L]
+            predict_noise = partial(
+                self.predict_noise,
+                cond=mask * x,
+                target_mask=1 - mask,
+                side_info=self.side_embedding(time, mask),
+                strategy=torch.tensor([strategy], device=x.device),
+            )
+            for i in range(num_samples):
+                noise = torch.empty_like(x).normal_(generator=generator)
+                sample = sampler(noise, predict_noise) * std + mean
+                out[i] += (1 - mask) * sample
+
+        return out.median(dim=0).values.transpose(1, 2)
+
+    def predict_noise(
+        self,
+        x: Tensor,  # [B, K, L]
+        step: Tensor,  # [1]
+        *,
+        cond: Tensor,  # [B, K, L]
+        target_mask: Tensor,  # [B, K, L]
+        side_info: Tensor,  # [B, S, K, L]
+        strategy: Tensor,  # [1]
+    ) -> Tensor:  # [B, K, L]
+        return self.denoiser(
+            x=torch.stack([cond, target_mask * x], dim=1),
+            side_info=side_info,
+            step=step,
+            strategy=strategy,
+        )
diff --git a/test/models/kumo/timeseries/anomaly/test_model.py b/test/models/kumo/timeseries/anomaly/test_model.py
new file mode 100644
index 000000000..04a46f0ad
--- /dev/null
+++ b/test/models/kumo/timeseries/anomaly/test_model.py
@@ -0,0 +1,88 @@
+# SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+
+import pytest
+import torch
+
+from sdm import CategoricalTensor, Recipe, Stype, TableTensor
+from sdm.explain.gradient import CaptureInputs, EnableInputGradients
+from sdm.models import Callback, KumoAnomaly
+from sdm.testing import withCUDA
+
+
+@withCUDA
+def test_forward(device: torch.device) -> None:
+    model = KumoAnomaly(pretrained=False, device=device)
+    x = torch.randn(2, 130, 3, device=device)
+
+    def generator() -> torch.Generator:
+        return torch.Generator(device).manual_seed(0)
+
+    out = model(x, num_samples=1, num_steps=1, generator=generator())
+    assert out.size() == (2, 130, 1)
+    assert out.columns["numerical"] == ("score",)
+    assert out.device == device
+    assert out.numerical.isfinite().all()
+    assert (out.numerical >= 0).all()
+
+    again = model(x, num_samples=1, num_steps=1, generator=generator())
+    assert torch.equal(again.numerical, out.numerical)
+
+
+def test_wide_float64_forward() -> None:
+    model = KumoAnomaly(pretrained=False)
+    x = torch.randn(60, 50, dtype=torch.float64)
+
+    out = model(x, num_samples=1, num_steps=1)
+    assert out.size() == (60, 1)
+    assert out.numerical.dtype == torch.float64
+
+
+class DoubleScores(Callback):
+    def on_model_forward_end(
+        self,
+        model: torch.nn.Module,
+        out: TableTensor,
+    ) -> TableTensor:
+        return out.replace_blocks(numerical=2 * out.numerical.double())
+
+
+def test_callbacks() -> None:
+    model = KumoAnomaly(pretrained=False)
+    x = torch.randn(30, 2)
+    capture = CaptureInputs()
+    callbacks = [EnableInputGradients(), capture, DoubleScores()]
+
+    def generator() -> torch.Generator:
+        return torch.Generator().manual_seed(0)
+
+    out = model(x, num_samples=1, num_steps=1, generator=generator())
+    doubled = model(
+        x,
+        num_samples=1,
+        num_steps=1,
+        callbacks=callbacks,
+        generator=generator(),
+    )
+    assert doubled.numerical.dtype == torch.float32
+    torch.testing.assert_close(doubled.numerical, 2 * out.numerical)
+
+    doubled.numerical.sum().backward()
+    ((x, _),) = capture.inputs
+    assert x.numerical.grad is not None
+
+
+def test_invalid_inputs() -> None:
+    model = KumoAnomaly(pretrained=False)
+    x = TableTensor(
+        columns={Stype.numerical: ["a"], Stype.categorical: ["b"]},
+        numerical=torch.randn(10, 1),
+        categorical=CategoricalTensor.from_tensor(
+            torch.randint(0, 5, (10, 1))
+        ),
+    )
+
+    with pytest.warns(RuntimeWarning, match="unsupported feature stypes"):
+        model(x, num_samples=1, num_steps=1)
+    with pytest.raises(ValueError, match="at most 40 features"):
+        model(torch.randn(10, 41), recipe=Recipe())

From 7d3adf1169b00ea60dac3af946366fa6b9212286 Mon Sep 17 00:00:00 2001
From: Ard <88337265+aw471@users.noreply.github.com>
Date: Mon, 19 Oct 2026 11:08:29 -0700
Subject: [PATCH 2/6] Support post-norm and dropout in `TransformerBlock`

Add optional `attn_residual_norm` and `mlp_residual_norm` modules to `TransformerBlock`, and apply its dropout to the attention output before the residual addition. Build the Kumo-Anomaly `ResidualBlock` on a new `KumoAnomalyTransformerBlock` instead of `torch.nn.TransformerEncoderLayer`.
---
 sdm/models/kumo/timeseries/anomaly/block.py   | 71 ++++++++++++++-----
 sdm/nn/attention.py                           | 30 +++++++-
 .../kumo/timeseries/anomaly/test_block.py     | 51 ++++++++-----
 test/nn/test_attention.py                     | 58 +++++++++++++++
 test/nn/test_compile.py                       | 26 +++++++
 5 files changed, 199 insertions(+), 37 deletions(-)

diff --git a/sdm/models/kumo/timeseries/anomaly/block.py b/sdm/models/kumo/timeseries/anomaly/block.py
index 975e50468..4cf5f8d30 100644
--- a/sdm/models/kumo/timeseries/anomaly/block.py
+++ b/sdm/models/kumo/timeseries/anomaly/block.py
@@ -6,7 +6,53 @@
 
 import torch
 from torch import Tensor
-from torch.nn import Conv1d, LayerNorm, Linear, TransformerEncoderLayer
+from torch.nn import GELU, Conv1d, Dropout, LayerNorm, Linear, Sequential
+
+from sdm.nn import TransformerBlock
+
+
+class KumoAnomalyTransformerBlock(TransformerBlock):
+    """Post-norm transformer block of Kumo-Anomaly.
+
+    Args:
+        channels: Input and output width.
+        num_heads: Attention heads.
+        hidden_channels: Feedforward width.
+        dropout: Attention and feedforward dropout probability.
+        device: The device.
+        dtype: The dtype.
+    """
+
+    def __init__(
+        self,
+        channels: int,
+        num_heads: int,
+        hidden_channels: int,
+        dropout: float,
+        device: torch.device | str | None = None,
+        dtype: torch.dtype | None = None,
+    ) -> None:
+        factory_kwargs: dict[str, Any] = {"device": device, "dtype": dtype}
+        super().__init__(
+            channels=channels,
+            num_query_heads=num_heads,
+            mlp=Sequential(
+                Linear(channels, hidden_channels, **factory_kwargs),
+                GELU(),
+                Dropout(dropout),
+                Linear(hidden_channels, channels, **factory_kwargs),
+                Dropout(dropout),
+            ),
+            attn_residual_norm=LayerNorm(channels, **factory_kwargs),
+            mlp_residual_norm=LayerNorm(channels, **factory_kwargs),
+            dropout=dropout,
+            **factory_kwargs,
+        )
+        # Match the initialization of `torch.nn.MultiheadAttention`.
+        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)
 
 
 class ResidualBlock(torch.nn.Module):
@@ -84,25 +130,18 @@ def __init__(
             with torch.no_grad():
                 projection.weight.mul_(0.5)
 
-        # Preserve post-residual normalization and attention dropout, which
-        # SDM's TransformerBlock does not currently provide together.
-        # TODO: Switch to SDM's TransformerBlock once it supports both.
-        self.time_layer = TransformerEncoderLayer(
-            d_model=channels,
-            nhead=num_heads,
-            dim_feedforward=hidden_channels,
+        self.time_layer = KumoAnomalyTransformerBlock(
+            channels=channels,
+            num_heads=num_heads,
+            hidden_channels=hidden_channels,
             dropout=dropout,
-            activation="gelu",
-            batch_first=True,
             **factory_kwargs,
         )
-        self.feature_layer = TransformerEncoderLayer(
-            d_model=channels,
-            nhead=num_heads,
-            dim_feedforward=hidden_channels,
+        self.feature_layer = KumoAnomalyTransformerBlock(
+            channels=channels,
+            num_heads=num_heads,
+            hidden_channels=hidden_channels,
             dropout=dropout,
-            activation="gelu",
-            batch_first=True,
             **factory_kwargs,
         )
         self.norm_after_time = LayerNorm(channels, **factory_kwargs)
diff --git a/sdm/nn/attention.py b/sdm/nn/attention.py
index fa4721ff7..c6446e2e1 100644
--- a/sdm/nn/attention.py
+++ b/sdm/nn/attention.py
@@ -458,6 +458,10 @@ class TransformerBlock(torch.nn.Module):
             attention.
         post_attn_norm: Normalization applied to the attention output before
             its residual addition.
+        attn_residual_norm: Normalization applied after the attention residual
+            addition.
+        mlp_residual_norm: Normalization applied after the feedforward
+            residual addition. Requires ``mlp``.
         query_transform: Transformation applied to projected query heads before
             scaled dot-product attention.
         key_transform: Transformation applied to projected key heads before
@@ -468,7 +472,8 @@ class TransformerBlock(torch.nn.Module):
         scale: Scaling factor passed to
             :func:`torch.nn.functional.scaled_dot_product_attention`.
             ``None`` uses ``1 / sqrt(channels_per_head)``.
-        dropout: Attention dropout probability.
+        dropout: Dropout probability applied to attention weights and to the
+            attention output before its residual addition.
         bias: If set to ``False``, the module will not learn an additive bias.
         device: The device.
         dtype: The dtype.
@@ -483,6 +488,8 @@ def __init__(
         query_norm: torch.nn.Module | None = None,
         key_value_norm: torch.nn.Module | None = None,
         post_attn_norm: torch.nn.Module | None = None,
+        attn_residual_norm: torch.nn.Module | None = None,
+        mlp_residual_norm: torch.nn.Module | None = None,
         query_transform: torch.nn.Module | None = None,
         key_transform: torch.nn.Module | None = None,
         query_scaling: QueryScaling | None = None,
@@ -494,12 +501,18 @@ def __init__(
         dtype: torch.dtype | None = None,
     ) -> None:
         super().__init__()
+        if mlp is None and mlp_residual_norm is not None:
+            raise ValueError("'mlp_residual_norm' requires 'mlp'")
+
         factory_kwargs: dict[str, Any] = {"device": device, "dtype": dtype}
 
         self.mlp = mlp
         self.query_norm = query_norm
         self.key_value_norm = key_value_norm
         self.post_attn_norm = post_attn_norm
+        self.attn_residual_norm = attn_residual_norm
+        self.mlp_residual_norm = mlp_residual_norm
+        self.dropout = dropout
 
         self.attn = Attention(
             channels=channels,
@@ -811,6 +824,8 @@ def _forward(
 
         if self.post_attn_norm is not None:
             attn_out = self.post_attn_norm(attn_out)
+        if self.training and self.dropout > 0.0:
+            attn_out = F.dropout(attn_out, p=self.dropout)
 
         if (
             out is not None
@@ -818,13 +833,24 @@ def _forward(
             and not out.is_contiguous()
         ):
             tmp = (attn_out + query).to(out.dtype)
+            if self.attn_residual_norm is not None:
+                tmp = self.attn_residual_norm(tmp)
             if self.mlp is not None:
-                tmp += self.mlp(tmp)
+                tmp = (tmp + self.mlp(tmp)).to(out.dtype)
+                if self.mlp_residual_norm is not None:
+                    tmp = self.mlp_residual_norm(tmp)
             out.copy_(tmp)
         else:
             tmp = torch.add(attn_out, query, out=out)
+            if self.attn_residual_norm is not None:
+                tmp = self.attn_residual_norm(tmp)
             if self.mlp is not None:
                 tmp = torch.add(tmp, self.mlp(tmp), out=out)
+                if self.mlp_residual_norm is not None:
+                    tmp = self.mlp_residual_norm(tmp)
+            # Normalization allocates, so copy its result into `out`.
+            if out is not None and tmp is not out:
+                tmp = out.copy_(tmp)
             out = tmp
 
         return (out, kv) if return_key_value else out
diff --git a/test/models/kumo/timeseries/anomaly/test_block.py b/test/models/kumo/timeseries/anomaly/test_block.py
index 92347cd88..c2269441b 100644
--- a/test/models/kumo/timeseries/anomaly/test_block.py
+++ b/test/models/kumo/timeseries/anomaly/test_block.py
@@ -2,35 +2,48 @@
 # SPDX-License-Identifier: Apache-2.0
 
 import math
-from collections.abc import Iterator
 
 import pytest
 import torch
-import torch.nn.functional as F
 from torch import Tensor
 
-from sdm.models.kumo.timeseries.anomaly.block import ResidualBlock
+from sdm.models.kumo.timeseries.anomaly.block import (
+    KumoAnomalyTransformerBlock,
+    ResidualBlock,
+)
 from sdm.testing import withCUDA
 
 
-@pytest.fixture(autouse=True)
-def exact_gelu() -> Iterator[None]:
-    # At inference time on GPU, PyTorch's encoder layer switches to a faster
-    # kernel with an approximate GELU. Disable it to compare exact outputs.
-    enabled = torch.backends.mha.get_fastpath_enabled()
-    torch.backends.mha.set_fastpath_enabled(False)
-    yield
-    torch.backends.mha.set_fastpath_enabled(enabled)
-
-
 def reference_attention(
-    layer: torch.nn.TransformerEncoderLayer, x: Tensor
+    layer: KumoAnomalyTransformerBlock, x: Tensor
 ) -> Tensor:
-    """Evaluate post-norm attention without the encoder's fused fast path."""
-    attn, _ = layer.self_attn(x, x, x, need_weights=False)
-    x = layer.norm1(x + layer.dropout1(attn))
-    hidden = layer.linear2(layer.dropout(F.gelu(layer.linear1(x))))
-    return layer.norm2(x + layer.dropout2(hidden))
+    """Evaluate a sequence-first PyTorch encoder layer with equal weights."""
+    state = layer.state_dict()
+    reference = torch.nn.TransformerEncoderLayer(
+        d_model=x.size(-1),
+        nhead=layer.attn.num_query_heads,
+        dim_feedforward=state["mlp.0.weight"].size(0),
+        dropout=0.0,
+        activation="gelu",
+        device=x.device,
+        dtype=x.dtype,
+    )
+    names = {
+        "self_attn.in_proj_": "attn.qkv_lin.",
+        "self_attn.out_proj.": "attn.out_lin.",
+        "linear1.": "mlp.0.",
+        "linear2.": "mlp.3.",
+        "norm1.": "attn_residual_norm.",
+        "norm2.": "mlp_residual_norm.",
+    }
+    reference.load_state_dict(
+        {
+            target + key: state[source + key]
+            for target, source in names.items()
+            for key in ("weight", "bias")
+        }
+    )
+    return reference.eval()(x.transpose(0, 1)).transpose(0, 1)
 
 
 def reference_block(
diff --git a/test/nn/test_attention.py b/test/nn/test_attention.py
index 4e45506ef..3e7e14b4d 100644
--- a/test/nn/test_attention.py
+++ b/test/nn/test_attention.py
@@ -751,3 +751,61 @@ def test_transformer_block_kv_cache() -> None:
 
     torch.testing.assert_close(cache_out, direct_out)
     torch.testing.assert_close(cached_out, direct_out)
+
+
+@withCUDA
+def test_transformer_block_residual_norms(device: torch.device) -> None:
+    channels = 8
+    mlp = torch.nn.Linear(channels, channels, device=device)
+    attn_residual_norm = torch.nn.LayerNorm(channels, device=device)
+    mlp_residual_norm = torch.nn.LayerNorm(channels, device=device)
+    module = TransformerBlock(
+        channels=channels,
+        num_query_heads=2,
+        mlp=mlp,
+        attn_residual_norm=attn_residual_norm,
+        mlp_residual_norm=mlp_residual_norm,
+        device=device,
+    )
+    torch.nn.init.normal_(module.attn.out_lin.weight)
+    query = torch.randn(2, 4, 3, channels, device=device).transpose(1, 2)
+
+    expected = attn_residual_norm(query + module.attn(query))
+    expected = mlp_residual_norm(expected + mlp(expected))
+    torch.testing.assert_close(module(query), expected)
+
+    with torch.no_grad():
+        out = torch.empty(2, 3, 4, channels, device=device)
+        assert module(query, batch_size_limit=1, out=out) is out
+        torch.testing.assert_close(out, expected)
+
+        alias = query.clone()
+        assert module(alias, out=alias) is alias
+        torch.testing.assert_close(alias, expected)
+
+    with pytest.raises(ValueError, match="requires 'mlp'"):
+        TransformerBlock(
+            channels=channels,
+            num_query_heads=2,
+            mlp=None,
+            mlp_residual_norm=mlp_residual_norm,
+        )
+
+
+@withCUDA
+def test_transformer_block_dropout(device: torch.device) -> None:
+    module = TransformerBlock(
+        channels=4,
+        num_query_heads=1,
+        mlp=None,
+        dropout=0.5,
+        device=device,
+    )
+    torch.nn.init.eye_(module.attn.out_lin.weight)
+    # A single key receives the full attention weight. Dropping weights and
+    # outputs each scales the unit values by 0 or 2.
+    query = torch.zeros(64, 1, 4, device=device)
+    with torch.no_grad():
+        module.attn.qkv_lin.bias[8:].fill_(1.0)
+        assert set(module(query).unique().tolist()) == {0.0, 4.0}
+        torch.testing.assert_close(module.eval()(query), query + 1)
diff --git a/test/nn/test_compile.py b/test/nn/test_compile.py
index d61f15a43..8a0f292ae 100644
--- a/test/nn/test_compile.py
+++ b/test/nn/test_compile.py
@@ -280,6 +280,32 @@ def test_transformer_block_compile(
     torch.testing.assert_close(out, expected)
 
 
+@withCUDA
+def test_transformer_block_residual_norms_compile(
+    device: torch.device,
+) -> None:
+    channels = 8
+    module = TransformerBlock(
+        channels=channels,
+        num_query_heads=2,
+        mlp=torch.nn.Linear(channels, channels, device=device),
+        attn_residual_norm=torch.nn.LayerNorm(channels, device=device),
+        mlp_residual_norm=torch.nn.LayerNorm(channels, device=device),
+        device=device,
+    )
+    query = torch.randn(2, 3, channels, device=device)
+    expected = torch.empty(2, 3, channels, device=device, dtype=torch.bfloat16)
+    out = expected.new_empty(3, 2, channels).transpose(0, 1)
+
+    # CUDA autocast runs `LayerNorm` in float32, so residual sums must be
+    # rounded to the dtype of `out` as in eager mode.
+    with torch.no_grad(), torch.autocast(device.type, dtype=torch.bfloat16):
+        module(query=query, out=expected)
+        actual = fullgraph(module)(query=query, out=out)
+    assert actual is out
+    assert torch.equal(out, expected)
+
+
 @withCUDA
 @pytest.mark.parametrize("self_attn", [False, True])
 def test_induced_transformer_block_compile(

From 78f96b5488b7756e9b87c4ab0e7285c84fa1b7d2 Mon Sep 17 00:00:00 2001
From: Ard <88337265+aw471@users.noreply.github.com>
Date: Wed, 21 Oct 2026 09:54:47 -0700
Subject: [PATCH 3/6] Load the pretrained `KumoAnomaly` checkpoint

Co-authored-by: Aditi Gautam <adgautam@nvidia.com>
---
 sdm/models/kumo/timeseries/anomaly/ckpt.py  | 32 +++++++++++++++++++++
 sdm/models/kumo/timeseries/anomaly/model.py | 10 ++++---
 2 files changed, 38 insertions(+), 4 deletions(-)
 create mode 100644 sdm/models/kumo/timeseries/anomaly/ckpt.py

diff --git a/sdm/models/kumo/timeseries/anomaly/ckpt.py b/sdm/models/kumo/timeseries/anomaly/ckpt.py
new file mode 100644
index 000000000..3c6cc7a36
--- /dev/null
+++ b/sdm/models/kumo/timeseries/anomaly/ckpt.py
@@ -0,0 +1,32 @@
+# SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+
+from torch import Tensor
+
+
+def remap_ckpt(ckpt: dict[str, Tensor]) -> dict[str, Tensor]:  # noqa: D103
+    out = {}
+    for key, value in ckpt.items():
+        if key.startswith("diffmodel."):
+            key = "denoiser." + key.removeprefix("diffmodel.")
+            for old, new in (
+                ("layers.0.self_attn.in_proj_", "attn.qkv_lin."),
+                ("layers.0.self_attn.out_proj.", "attn.out_lin."),
+                ("layers.0.linear1.", "mlp.0."),
+                ("layers.0.linear2.", "mlp.3."),
+                ("layers.0.norm1.", "attn_residual_norm."),
+                ("layers.0.norm2.", "mlp_residual_norm."),
+            ):
+                key = key.replace(old, new)
+            if key in (
+                "denoiser.input_projection.weight",
+                "denoiser.output_projection1.weight",
+                "denoiser.output_projection2.weight",
+            ):
+                value = value.squeeze(-1)  # `Conv1d` to `Linear`.
+        elif key.startswith(
+            ("embed_layer.", "embed_norm.", "time_embed_norm.")
+        ):
+            key = "side_embedding." + key
+        out[key] = value
+    return out
diff --git a/sdm/models/kumo/timeseries/anomaly/model.py b/sdm/models/kumo/timeseries/anomaly/model.py
index 51178bc97..b1e88afe4 100644
--- a/sdm/models/kumo/timeseries/anomaly/model.py
+++ b/sdm/models/kumo/timeseries/anomaly/model.py
@@ -19,6 +19,7 @@
 from sdm._warnings import warn_once
 from sdm.models import Callback, Model
 from sdm.models._huggingface import download_checkpoint
+from sdm.models.kumo.timeseries.anomaly.ckpt import remap_ckpt
 from sdm.models.kumo.timeseries.anomaly.denoiser import DiffusionDenoiser
 from sdm.models.kumo.timeseries.anomaly.embedding import SideEmbedding
 from sdm.models.kumo.timeseries.anomaly.recipe import default_recipe
@@ -61,9 +62,9 @@ def __init__(
     ) -> None:
         super().__init__()
 
-        self.model = _KumoAnomaly(
-            device="meta" if pretrained else device,
-        )
+        # Build on the target device, since the checkpoint does not contain
+        # the non-persistent sinusoidal lookup of the step embedding:
+        self.model = _KumoAnomaly(device=device)
 
         if pretrained:
             self._load_from_pretrained(device=device)
@@ -81,7 +82,8 @@ def _load_from_pretrained(
             filename="final_model.pth",
             revision="226c5003b0582adc1cac433123f47324bfae39f2",
         )
-        torch.load(path, map_location=device, weights_only=True)  # TODO
+        ckpt = torch.load(path, map_location=device, weights_only=True)
+        self.model.load_state_dict(remap_ckpt(ckpt["model"]))
 
         return self
 

From c692650cd7f4f3ad038fd6678595f84d7d0eac33 Mon Sep 17 00:00:00 2001
From: rusty1s <matthias.fey@tu-dortmund.de>
Date: Thu, 22 Oct 2026 09:43:49 +0000
Subject: [PATCH 4/6] update

---
 sdm/models/kumo/timeseries/anomaly/block.py   | 71 +++++--------------
 sdm/models/kumo/timeseries/anomaly/ckpt.py    | 10 +--
 sdm/models/kumo/timeseries/anomaly/model.py   |  3 +-
 sdm/nn/attention.py                           | 30 +-------
 .../kumo/timeseries/anomaly/test_block.py     | 51 +++++--------
 .../kumo/timeseries/anomaly/test_model.py     | 47 ++++++++++++
 test/nn/test_attention.py                     | 58 ---------------
 test/nn/test_compile.py                       | 26 -------
 8 files changed, 87 insertions(+), 209 deletions(-)

diff --git a/sdm/models/kumo/timeseries/anomaly/block.py b/sdm/models/kumo/timeseries/anomaly/block.py
index 4cf5f8d30..975e50468 100644
--- a/sdm/models/kumo/timeseries/anomaly/block.py
+++ b/sdm/models/kumo/timeseries/anomaly/block.py
@@ -6,53 +6,7 @@
 
 import torch
 from torch import Tensor
-from torch.nn import GELU, Conv1d, Dropout, LayerNorm, Linear, Sequential
-
-from sdm.nn import TransformerBlock
-
-
-class KumoAnomalyTransformerBlock(TransformerBlock):
-    """Post-norm transformer block of Kumo-Anomaly.
-
-    Args:
-        channels: Input and output width.
-        num_heads: Attention heads.
-        hidden_channels: Feedforward width.
-        dropout: Attention and feedforward dropout probability.
-        device: The device.
-        dtype: The dtype.
-    """
-
-    def __init__(
-        self,
-        channels: int,
-        num_heads: int,
-        hidden_channels: int,
-        dropout: float,
-        device: torch.device | str | None = None,
-        dtype: torch.dtype | None = None,
-    ) -> None:
-        factory_kwargs: dict[str, Any] = {"device": device, "dtype": dtype}
-        super().__init__(
-            channels=channels,
-            num_query_heads=num_heads,
-            mlp=Sequential(
-                Linear(channels, hidden_channels, **factory_kwargs),
-                GELU(),
-                Dropout(dropout),
-                Linear(hidden_channels, channels, **factory_kwargs),
-                Dropout(dropout),
-            ),
-            attn_residual_norm=LayerNorm(channels, **factory_kwargs),
-            mlp_residual_norm=LayerNorm(channels, **factory_kwargs),
-            dropout=dropout,
-            **factory_kwargs,
-        )
-        # Match the initialization of `torch.nn.MultiheadAttention`.
-        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)
+from torch.nn import Conv1d, LayerNorm, Linear, TransformerEncoderLayer
 
 
 class ResidualBlock(torch.nn.Module):
@@ -130,18 +84,25 @@ def __init__(
             with torch.no_grad():
                 projection.weight.mul_(0.5)
 
-        self.time_layer = KumoAnomalyTransformerBlock(
-            channels=channels,
-            num_heads=num_heads,
-            hidden_channels=hidden_channels,
+        # Preserve post-residual normalization and attention dropout, which
+        # SDM's TransformerBlock does not currently provide together.
+        # TODO: Switch to SDM's TransformerBlock once it supports both.
+        self.time_layer = TransformerEncoderLayer(
+            d_model=channels,
+            nhead=num_heads,
+            dim_feedforward=hidden_channels,
             dropout=dropout,
+            activation="gelu",
+            batch_first=True,
             **factory_kwargs,
         )
-        self.feature_layer = KumoAnomalyTransformerBlock(
-            channels=channels,
-            num_heads=num_heads,
-            hidden_channels=hidden_channels,
+        self.feature_layer = TransformerEncoderLayer(
+            d_model=channels,
+            nhead=num_heads,
+            dim_feedforward=hidden_channels,
             dropout=dropout,
+            activation="gelu",
+            batch_first=True,
             **factory_kwargs,
         )
         self.norm_after_time = LayerNorm(channels, **factory_kwargs)
diff --git a/sdm/models/kumo/timeseries/anomaly/ckpt.py b/sdm/models/kumo/timeseries/anomaly/ckpt.py
index 3c6cc7a36..014c66648 100644
--- a/sdm/models/kumo/timeseries/anomaly/ckpt.py
+++ b/sdm/models/kumo/timeseries/anomaly/ckpt.py
@@ -9,15 +9,7 @@ def remap_ckpt(ckpt: dict[str, Tensor]) -> dict[str, Tensor]:  # noqa: D103
     for key, value in ckpt.items():
         if key.startswith("diffmodel."):
             key = "denoiser." + key.removeprefix("diffmodel.")
-            for old, new in (
-                ("layers.0.self_attn.in_proj_", "attn.qkv_lin."),
-                ("layers.0.self_attn.out_proj.", "attn.out_lin."),
-                ("layers.0.linear1.", "mlp.0."),
-                ("layers.0.linear2.", "mlp.3."),
-                ("layers.0.norm1.", "attn_residual_norm."),
-                ("layers.0.norm2.", "mlp_residual_norm."),
-            ):
-                key = key.replace(old, new)
+            key = key.replace(".layers.0.", ".")
             if key in (
                 "denoiser.input_projection.weight",
                 "denoiser.output_projection1.weight",
diff --git a/sdm/models/kumo/timeseries/anomaly/model.py b/sdm/models/kumo/timeseries/anomaly/model.py
index ffb94cadb..758849575 100644
--- a/sdm/models/kumo/timeseries/anomaly/model.py
+++ b/sdm/models/kumo/timeseries/anomaly/model.py
@@ -18,6 +18,7 @@
 from sdm._inference import inference_mode
 from sdm.models import Callback, Model
 from sdm.models._huggingface import download_checkpoint
+from sdm.models.kumo.timeseries.anomaly.ckpt import remap_ckpt
 from sdm.models.kumo.timeseries.anomaly.denoiser import DiffusionDenoiser
 from sdm.models.kumo.timeseries.anomaly.embedding import SideEmbedding
 from sdm.models.kumo.timeseries.anomaly.recipe import default_recipe
@@ -106,7 +107,7 @@ def forward(
         Args:
             x: The feature tensor with shape ``[..., R, D]`` with ``R`` rows
                 and ``D`` columns.
-            recipe: The custom recipe for pre-processing.
+            recipe: The custom recipe for pre- and post-processing.
             num_samples: The number of reconstruction samples per complementary
                 masking strategy. Reconstructions are combined using the
                 sample median.
diff --git a/sdm/nn/attention.py b/sdm/nn/attention.py
index c6446e2e1..fa4721ff7 100644
--- a/sdm/nn/attention.py
+++ b/sdm/nn/attention.py
@@ -458,10 +458,6 @@ class TransformerBlock(torch.nn.Module):
             attention.
         post_attn_norm: Normalization applied to the attention output before
             its residual addition.
-        attn_residual_norm: Normalization applied after the attention residual
-            addition.
-        mlp_residual_norm: Normalization applied after the feedforward
-            residual addition. Requires ``mlp``.
         query_transform: Transformation applied to projected query heads before
             scaled dot-product attention.
         key_transform: Transformation applied to projected key heads before
@@ -472,8 +468,7 @@ class TransformerBlock(torch.nn.Module):
         scale: Scaling factor passed to
             :func:`torch.nn.functional.scaled_dot_product_attention`.
             ``None`` uses ``1 / sqrt(channels_per_head)``.
-        dropout: Dropout probability applied to attention weights and to the
-            attention output before its residual addition.
+        dropout: Attention dropout probability.
         bias: If set to ``False``, the module will not learn an additive bias.
         device: The device.
         dtype: The dtype.
@@ -488,8 +483,6 @@ def __init__(
         query_norm: torch.nn.Module | None = None,
         key_value_norm: torch.nn.Module | None = None,
         post_attn_norm: torch.nn.Module | None = None,
-        attn_residual_norm: torch.nn.Module | None = None,
-        mlp_residual_norm: torch.nn.Module | None = None,
         query_transform: torch.nn.Module | None = None,
         key_transform: torch.nn.Module | None = None,
         query_scaling: QueryScaling | None = None,
@@ -501,18 +494,12 @@ def __init__(
         dtype: torch.dtype | None = None,
     ) -> None:
         super().__init__()
-        if mlp is None and mlp_residual_norm is not None:
-            raise ValueError("'mlp_residual_norm' requires 'mlp'")
-
         factory_kwargs: dict[str, Any] = {"device": device, "dtype": dtype}
 
         self.mlp = mlp
         self.query_norm = query_norm
         self.key_value_norm = key_value_norm
         self.post_attn_norm = post_attn_norm
-        self.attn_residual_norm = attn_residual_norm
-        self.mlp_residual_norm = mlp_residual_norm
-        self.dropout = dropout
 
         self.attn = Attention(
             channels=channels,
@@ -824,8 +811,6 @@ def _forward(
 
         if self.post_attn_norm is not None:
             attn_out = self.post_attn_norm(attn_out)
-        if self.training and self.dropout > 0.0:
-            attn_out = F.dropout(attn_out, p=self.dropout)
 
         if (
             out is not None
@@ -833,24 +818,13 @@ def _forward(
             and not out.is_contiguous()
         ):
             tmp = (attn_out + query).to(out.dtype)
-            if self.attn_residual_norm is not None:
-                tmp = self.attn_residual_norm(tmp)
             if self.mlp is not None:
-                tmp = (tmp + self.mlp(tmp)).to(out.dtype)
-                if self.mlp_residual_norm is not None:
-                    tmp = self.mlp_residual_norm(tmp)
+                tmp += self.mlp(tmp)
             out.copy_(tmp)
         else:
             tmp = torch.add(attn_out, query, out=out)
-            if self.attn_residual_norm is not None:
-                tmp = self.attn_residual_norm(tmp)
             if self.mlp is not None:
                 tmp = torch.add(tmp, self.mlp(tmp), out=out)
-                if self.mlp_residual_norm is not None:
-                    tmp = self.mlp_residual_norm(tmp)
-            # Normalization allocates, so copy its result into `out`.
-            if out is not None and tmp is not out:
-                tmp = out.copy_(tmp)
             out = tmp
 
         return (out, kv) if return_key_value else out
diff --git a/test/models/kumo/timeseries/anomaly/test_block.py b/test/models/kumo/timeseries/anomaly/test_block.py
index c2269441b..92347cd88 100644
--- a/test/models/kumo/timeseries/anomaly/test_block.py
+++ b/test/models/kumo/timeseries/anomaly/test_block.py
@@ -2,48 +2,35 @@
 # SPDX-License-Identifier: Apache-2.0
 
 import math
+from collections.abc import Iterator
 
 import pytest
 import torch
+import torch.nn.functional as F
 from torch import Tensor
 
-from sdm.models.kumo.timeseries.anomaly.block import (
-    KumoAnomalyTransformerBlock,
-    ResidualBlock,
-)
+from sdm.models.kumo.timeseries.anomaly.block import ResidualBlock
 from sdm.testing import withCUDA
 
 
+@pytest.fixture(autouse=True)
+def exact_gelu() -> Iterator[None]:
+    # At inference time on GPU, PyTorch's encoder layer switches to a faster
+    # kernel with an approximate GELU. Disable it to compare exact outputs.
+    enabled = torch.backends.mha.get_fastpath_enabled()
+    torch.backends.mha.set_fastpath_enabled(False)
+    yield
+    torch.backends.mha.set_fastpath_enabled(enabled)
+
+
 def reference_attention(
-    layer: KumoAnomalyTransformerBlock, x: Tensor
+    layer: torch.nn.TransformerEncoderLayer, x: Tensor
 ) -> Tensor:
-    """Evaluate a sequence-first PyTorch encoder layer with equal weights."""
-    state = layer.state_dict()
-    reference = torch.nn.TransformerEncoderLayer(
-        d_model=x.size(-1),
-        nhead=layer.attn.num_query_heads,
-        dim_feedforward=state["mlp.0.weight"].size(0),
-        dropout=0.0,
-        activation="gelu",
-        device=x.device,
-        dtype=x.dtype,
-    )
-    names = {
-        "self_attn.in_proj_": "attn.qkv_lin.",
-        "self_attn.out_proj.": "attn.out_lin.",
-        "linear1.": "mlp.0.",
-        "linear2.": "mlp.3.",
-        "norm1.": "attn_residual_norm.",
-        "norm2.": "mlp_residual_norm.",
-    }
-    reference.load_state_dict(
-        {
-            target + key: state[source + key]
-            for target, source in names.items()
-            for key in ("weight", "bias")
-        }
-    )
-    return reference.eval()(x.transpose(0, 1)).transpose(0, 1)
+    """Evaluate post-norm attention without the encoder's fused fast path."""
+    attn, _ = layer.self_attn(x, x, x, need_weights=False)
+    x = layer.norm1(x + layer.dropout1(attn))
+    hidden = layer.linear2(layer.dropout(F.gelu(layer.linear1(x))))
+    return layer.norm2(x + layer.dropout2(hidden))
 
 
 def reference_block(
diff --git a/test/models/kumo/timeseries/anomaly/test_model.py b/test/models/kumo/timeseries/anomaly/test_model.py
index 849dd2874..d4705e590 100644
--- a/test/models/kumo/timeseries/anomaly/test_model.py
+++ b/test/models/kumo/timeseries/anomaly/test_model.py
@@ -1,6 +1,9 @@
 # SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved.
 # SPDX-License-Identifier: Apache-2.0
 
+from pathlib import Path
+
+import pytest
 import torch
 
 from sdm.models import KumoAnomaly
@@ -18,3 +21,47 @@ def test_forward(device: torch.device) -> None:
     assert out.device == device
     assert out.numerical.isfinite().all()
     assert (out.numerical >= 0).all()
+
+
+def test_load_from_pretrained(
+    tmp_path: Path,
+    monkeypatch: pytest.MonkeyPatch,
+) -> None:
+    reference = KumoAnomaly(pretrained=False)
+    ckpt = {}
+    for key, value in reference.model.state_dict().items():
+        if key.startswith("denoiser."):
+            key = "diffmodel." + key.removeprefix("denoiser.")
+            key = key.replace(".time_layer.", ".time_layer.layers.0.")
+            key = key.replace(".feature_layer.", ".feature_layer.layers.0.")
+            if key in (
+                "diffmodel.input_projection.weight",
+                "diffmodel.output_projection1.weight",
+                "diffmodel.output_projection2.weight",
+            ):
+                value = value.unsqueeze(-1)
+        else:
+            key = key.removeprefix("side_embedding.")
+        ckpt[key] = value
+
+    path = tmp_path / "final_model.pth"
+    torch.save({"model": ckpt}, path)
+    monkeypatch.setattr(
+        "sdm.models.kumo.timeseries.anomaly.model.download_checkpoint",
+        lambda **kwargs: path,
+    )
+    model = KumoAnomaly(pretrained=True)
+    x = torch.randn(17, 3)
+    expected = reference(
+        x=x,
+        num_samples=1,
+        num_steps=2,
+        generator=torch.Generator().manual_seed(42),
+    )
+    out = model(
+        x=x,
+        num_samples=1,
+        num_steps=2,
+        generator=torch.Generator().manual_seed(42),
+    )
+    torch.testing.assert_close(out.numerical, expected.numerical)
diff --git a/test/nn/test_attention.py b/test/nn/test_attention.py
index 3e7e14b4d..4e45506ef 100644
--- a/test/nn/test_attention.py
+++ b/test/nn/test_attention.py
@@ -751,61 +751,3 @@ def test_transformer_block_kv_cache() -> None:
 
     torch.testing.assert_close(cache_out, direct_out)
     torch.testing.assert_close(cached_out, direct_out)
-
-
-@withCUDA
-def test_transformer_block_residual_norms(device: torch.device) -> None:
-    channels = 8
-    mlp = torch.nn.Linear(channels, channels, device=device)
-    attn_residual_norm = torch.nn.LayerNorm(channels, device=device)
-    mlp_residual_norm = torch.nn.LayerNorm(channels, device=device)
-    module = TransformerBlock(
-        channels=channels,
-        num_query_heads=2,
-        mlp=mlp,
-        attn_residual_norm=attn_residual_norm,
-        mlp_residual_norm=mlp_residual_norm,
-        device=device,
-    )
-    torch.nn.init.normal_(module.attn.out_lin.weight)
-    query = torch.randn(2, 4, 3, channels, device=device).transpose(1, 2)
-
-    expected = attn_residual_norm(query + module.attn(query))
-    expected = mlp_residual_norm(expected + mlp(expected))
-    torch.testing.assert_close(module(query), expected)
-
-    with torch.no_grad():
-        out = torch.empty(2, 3, 4, channels, device=device)
-        assert module(query, batch_size_limit=1, out=out) is out
-        torch.testing.assert_close(out, expected)
-
-        alias = query.clone()
-        assert module(alias, out=alias) is alias
-        torch.testing.assert_close(alias, expected)
-
-    with pytest.raises(ValueError, match="requires 'mlp'"):
-        TransformerBlock(
-            channels=channels,
-            num_query_heads=2,
-            mlp=None,
-            mlp_residual_norm=mlp_residual_norm,
-        )
-
-
-@withCUDA
-def test_transformer_block_dropout(device: torch.device) -> None:
-    module = TransformerBlock(
-        channels=4,
-        num_query_heads=1,
-        mlp=None,
-        dropout=0.5,
-        device=device,
-    )
-    torch.nn.init.eye_(module.attn.out_lin.weight)
-    # A single key receives the full attention weight. Dropping weights and
-    # outputs each scales the unit values by 0 or 2.
-    query = torch.zeros(64, 1, 4, device=device)
-    with torch.no_grad():
-        module.attn.qkv_lin.bias[8:].fill_(1.0)
-        assert set(module(query).unique().tolist()) == {0.0, 4.0}
-        torch.testing.assert_close(module.eval()(query), query + 1)
diff --git a/test/nn/test_compile.py b/test/nn/test_compile.py
index 8a0f292ae..d61f15a43 100644
--- a/test/nn/test_compile.py
+++ b/test/nn/test_compile.py
@@ -280,32 +280,6 @@ def test_transformer_block_compile(
     torch.testing.assert_close(out, expected)
 
 
-@withCUDA
-def test_transformer_block_residual_norms_compile(
-    device: torch.device,
-) -> None:
-    channels = 8
-    module = TransformerBlock(
-        channels=channels,
-        num_query_heads=2,
-        mlp=torch.nn.Linear(channels, channels, device=device),
-        attn_residual_norm=torch.nn.LayerNorm(channels, device=device),
-        mlp_residual_norm=torch.nn.LayerNorm(channels, device=device),
-        device=device,
-    )
-    query = torch.randn(2, 3, channels, device=device)
-    expected = torch.empty(2, 3, channels, device=device, dtype=torch.bfloat16)
-    out = expected.new_empty(3, 2, channels).transpose(0, 1)
-
-    # CUDA autocast runs `LayerNorm` in float32, so residual sums must be
-    # rounded to the dtype of `out` as in eager mode.
-    with torch.no_grad(), torch.autocast(device.type, dtype=torch.bfloat16):
-        module(query=query, out=expected)
-        actual = fullgraph(module)(query=query, out=out)
-    assert actual is out
-    assert torch.equal(out, expected)
-
-
 @withCUDA
 @pytest.mark.parametrize("self_attn", [False, True])
 def test_induced_transformer_block_compile(

From 01e54e7ee14ac85c65cff25c73f3bc52d73b4b34 Mon Sep 17 00:00:00 2001
From: rusty1s <matthias.fey@tu-dortmund.de>
Date: Thu, 22 Oct 2026 09:44:24 +0000
Subject: [PATCH 5/6] update

---
 sdm/models/kumo/timeseries/anomaly/model.py   |  2 -
 .../kumo/timeseries/anomaly/test_model.py     | 47 -------------------
 2 files changed, 49 deletions(-)

diff --git a/sdm/models/kumo/timeseries/anomaly/model.py b/sdm/models/kumo/timeseries/anomaly/model.py
index 758849575..e24c35a04 100644
--- a/sdm/models/kumo/timeseries/anomaly/model.py
+++ b/sdm/models/kumo/timeseries/anomaly/model.py
@@ -62,8 +62,6 @@ def __init__(
     ) -> None:
         super().__init__()
 
-        # Build on the target device, since the checkpoint does not contain
-        # the non-persistent sinusoidal lookup of the step embedding:
         self.model = _KumoAnomaly(device=device)
 
         if pretrained:
diff --git a/test/models/kumo/timeseries/anomaly/test_model.py b/test/models/kumo/timeseries/anomaly/test_model.py
index d4705e590..849dd2874 100644
--- a/test/models/kumo/timeseries/anomaly/test_model.py
+++ b/test/models/kumo/timeseries/anomaly/test_model.py
@@ -1,9 +1,6 @@
 # SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved.
 # SPDX-License-Identifier: Apache-2.0
 
-from pathlib import Path
-
-import pytest
 import torch
 
 from sdm.models import KumoAnomaly
@@ -21,47 +18,3 @@ def test_forward(device: torch.device) -> None:
     assert out.device == device
     assert out.numerical.isfinite().all()
     assert (out.numerical >= 0).all()
-
-
-def test_load_from_pretrained(
-    tmp_path: Path,
-    monkeypatch: pytest.MonkeyPatch,
-) -> None:
-    reference = KumoAnomaly(pretrained=False)
-    ckpt = {}
-    for key, value in reference.model.state_dict().items():
-        if key.startswith("denoiser."):
-            key = "diffmodel." + key.removeprefix("denoiser.")
-            key = key.replace(".time_layer.", ".time_layer.layers.0.")
-            key = key.replace(".feature_layer.", ".feature_layer.layers.0.")
-            if key in (
-                "diffmodel.input_projection.weight",
-                "diffmodel.output_projection1.weight",
-                "diffmodel.output_projection2.weight",
-            ):
-                value = value.unsqueeze(-1)
-        else:
-            key = key.removeprefix("side_embedding.")
-        ckpt[key] = value
-
-    path = tmp_path / "final_model.pth"
-    torch.save({"model": ckpt}, path)
-    monkeypatch.setattr(
-        "sdm.models.kumo.timeseries.anomaly.model.download_checkpoint",
-        lambda **kwargs: path,
-    )
-    model = KumoAnomaly(pretrained=True)
-    x = torch.randn(17, 3)
-    expected = reference(
-        x=x,
-        num_samples=1,
-        num_steps=2,
-        generator=torch.Generator().manual_seed(42),
-    )
-    out = model(
-        x=x,
-        num_samples=1,
-        num_steps=2,
-        generator=torch.Generator().manual_seed(42),
-    )
-    torch.testing.assert_close(out.numerical, expected.numerical)

From d9af792824c66bc747a3b75697091ed6a9f22fe3 Mon Sep 17 00:00:00 2001
From: rusty1s <matthias.fey@tu-dortmund.de>
Date: Thu, 22 Oct 2026 10:09:26 +0000
Subject: [PATCH 6/6] update

---
 sdm/models/kumo/timeseries/anomaly/recipe.py | 1 +
 1 file changed, 1 insertion(+)

diff --git a/sdm/models/kumo/timeseries/anomaly/recipe.py b/sdm/models/kumo/timeseries/anomaly/recipe.py
index 3f6425b6c..27f040df0 100644
--- a/sdm/models/kumo/timeseries/anomaly/recipe.py
+++ b/sdm/models/kumo/timeseries/anomaly/recipe.py
@@ -8,6 +8,7 @@ def default_recipe() -> sp.Recipe:  # noqa: D103
     return sp.Recipe(
         features=sp.StypeDispatch(
             numerical=[
+                sp.ImputeMean(),
                 sp.PCA(num_components=40),
                 sp.MinMaxScale(),
             ]