Upload folder using huggingface_hub

README.md

@@ -1,13 +1,31 @@
 ---
-title: Timebraid
-emoji: 📈
-colorFrom: gray
-colorTo: pink
+title: TimeBraid
+emoji: 🧵
+colorFrom: indigo
+colorTo: gray
 sdk: gradio
 sdk_version: 6.29.1
-python_version: '3.12'
 app_file: app.py
-pinned: false
+short_description: Time series + language model for analysis and forecasting
+python_version: "3.12"
+startup_duration_timeout: 30m
 ---
 
-Check out the configuration reference at https://huggingface.co/docs/hub/spaces-config-reference
+# TimeBraid
+
+Interactive demo of **TimeBraid-2.5B** — a unified time-series + language model that
+interleaves a Qwen3-1.7B language backbone with TimesFM 2.5 200M time-series experts
+through Mixture-of-Transformers layers.
+
+Paste numeric series (one series per line, values comma- or space-separated) and ask a
+question, or set a forecast horizon to get a numeric point forecast plotted against the
+observed history.
+
+- Model: [XinyueWangg/TimeBraid-2.5B](https://huggingface.co/XinyueWangg/TimeBraid-2.5B) (Apache-2.0)
+- Code: [CharonWangg/TimeBraid](https://github.com/CharonWangg/TimeBraid)
+- The `timebraid/` package in this Space is the authors' inference package, shipped
+  verbatim so `AutoModelForCausalLM` resolves the custom architecture with
+  `trust_remote_code=False`.
+
+Runs on ZeroGPU. Example inputs are drawn from the authors'
+`examples/inference_tasks.ipynb` notebook (Apache-2.0).

app.py

@@ -0,0 +1,331 @@
+"""TimeBraid: unified time series + language demo on ZeroGPU."""
+
+import os
+
+os.environ.setdefault("PYTORCH_CUDA_ALLOC_CONF", "expandable_segments:True")
+
+import spaces  # MUST come before any torch / CUDA-touching import
+import time
+
+import gradio as gr
+import matplotlib
+
+matplotlib.use("Agg")
+import matplotlib.pyplot as plt
+import torch
+from matplotlib.ticker import MaxNLocator
+from transformers import AutoModelForCausalLM, AutoProcessor
+
+from timebraid import TimeBraidProcessor  # noqa: F401 - registers HF AutoClasses
+
+MODEL_ID = "XinyueWangg/TimeBraid-2.5B"
+
+# Chart conventions from the authors' notebook (examples/inference_tasks.ipynb).
+_INPUT_COLOR = "#0072B2"
+_FORECAST_COLOR = "#D55E00"
+_BOUNDARY_COLOR = "#555555"
+_SERIES_COLORS = ("#0072B2", "#009E73", "#CC79A7", "#E69F00")
+
+processor = AutoProcessor.from_pretrained(MODEL_ID, trust_remote_code=False)
+model = (
+    AutoModelForCausalLM.from_pretrained(
+        MODEL_ID,
+        dtype=torch.bfloat16,
+        attn_implementation="flash_attention_2",
+        trust_remote_code=False,
+        use_safetensors=True,
+    )
+    .eval()
+    .to("cuda")
+)
+
+
+def parse_series(text: str) -> list[list[float]]:
+    """Parse one series per line; values comma- or whitespace-separated."""
+    if text is None:
+        return []
+    series: list[list[float]] = []
+    for line in text.splitlines():
+        line = line.strip().rstrip(",")
+        if not line:
+            continue
+        tokens = line.replace(",", " ").split()
+        values = [float(token) for token in tokens]
+        if not values:
+            raise ValueError("A series line parsed to zero values.")
+        series.append(values)
+    return series
+
+
+def render_plot(
+    history: list[list[float]],
+    forecast: list[float] | None,
+    target_index: int | None,
+) -> matplotlib.figure.Figure:
+    """Draw observed history and, when present, the forecast continuation."""
+    n = len(history)
+    fig, axes = plt.subplots(
+        n, 1, figsize=(9.5, 2.9 * n), constrained_layout=True, squeeze=False
+    )
+    for index, (ax, values) in enumerate(zip(axes[:, 0], history)):
+        color = _INPUT_COLOR if forecast is not None else _SERIES_COLORS[index % 4]
+        x = list(range(len(values)))
+        ax.plot(
+            x,
+            values,
+            color=color,
+            linewidth=2.0,
+            marker="o" if len(values) <= 64 else None,
+            markersize=4,
+            label="Observed history" if forecast is not None else f"Series {index + 1}",
+        )
+        title = f"INPUT TIMESERIES · Series {index + 1}"
+        if forecast is not None and target_index == index:
+            future_x = list(range(len(values), len(values) + len(forecast)))
+            boundary = len(values) - 0.5
+            ax.axvspan(boundary, future_x[-1] + 0.5, color=_FORECAST_COLOR, alpha=0.06)
+            ax.plot(
+                (len(values) - 1, *future_x),
+                (values[-1], *forecast),
+                color=_FORECAST_COLOR,
+                linewidth=2.0,
+                linestyle="--",
+                label="Point forecast",
+            )
+            ax.scatter(future_x, forecast, color=_FORECAST_COLOR, s=24, zorder=3)
+            ax.axvline(
+                boundary,
+                color=_BOUNDARY_COLOR,
+                linewidth=1.2,
+                linestyle=":",
+                label="Forecast boundary",
+            )
+            ax.text(
+                boundary,
+                1.01,
+                "forecast starts",
+                transform=ax.get_xaxis_transform(),
+                ha="center",
+                va="bottom",
+                fontsize=8,
+                color=_BOUNDARY_COLOR,
+            )
+            title = (
+                f"INPUT HISTORY + OUTPUT FORECAST · Series {index + 1} (forecast target)"
+            )
+        ax.set_title(title, loc="left", fontsize=11, fontweight="bold")
+        ax.set_xlabel("Observation / forecast index" if forecast is not None else "Observation index")
+        ax.set_ylabel("Value (raw scale)")
+        ax.xaxis.set_major_locator(MaxNLocator(integer=True))
+        ax.grid(axis="y", color="#D9DCE1", linewidth=0.8, alpha=0.7)
+        ax.set_axisbelow(True)
+        ax.spines["top"].set_visible(False)
+        ax.spines["right"].set_visible(False)
+        ax.legend(loc="upper left", frameon=False, ncol=3, fontsize=9)
+    return fig
+
+
+@spaces.GPU(duration=120)
+def run(
+    prompt: str,
+    series_text: str,
+    horizon: float,
+    target_series_index: float,
+    max_new_tokens: float,
+    progress=gr.Progress(track_tqdm=True),
+):
+    """Analyze or forecast a time series with TimeBraid.
+
+    Args:
+        prompt: Question or instruction about the series (or a plain text question).
+        series_text: Numeric series, one per line; values comma- or space-separated.
+        horizon: Number of future points to forecast; 0 = no forecast (understanding only).
+        target_series_index: 1-based index of the series to forecast (multi-series only).
+        max_new_tokens: Maximum generated text tokens.
+    Returns:
+        (answer text, forecast values text, plot, elapsed seconds text).
+    """
+    started = time.perf_counter()
+    try:
+        series = parse_series(series_text)
+    except ValueError as exc:
+        raise gr.Error(f"Could not parse the time series: {exc}")
+
+    horizon_i = int(horizon) if horizon and horizon > 0 else None
+    target_i: int | None = None
+    if horizon_i is not None:
+        if not series:
+            raise gr.Error("A forecast horizon needs at least one input series.")
+        if len(series) > 1:
+            if not target_series_index or target_series_index < 1:
+                raise gr.Error(
+                    "With multiple series, pick which one to forecast (Target series)."
+                )
+            if int(target_series_index) > len(series):
+                raise gr.Error(
+                    f"Target series is out of range: only {len(series)} series given."
+                )
+            target_i = int(target_series_index) - 1
+
+    messages = [{"role": "user", "content": prompt}]
+    inputs = processor.apply_chat_template(
+        messages,
+        timeseries=series if series else None,
+        horizon=horizon_i,
+        target_series_index=target_i,
+        add_generation_prompt=True,
+        tokenize=True,
+        return_dict=True,
+        return_tensors="pt",
+    ).to(model.device)
+    with torch.inference_mode():
+        outputs = model.generate(
+            **inputs,
+            max_new_tokens=int(max_new_tokens),
+            do_sample=False,
+            num_beams=1,
+            num_return_sequences=1,
+        )
+    result = processor.post_process_generation(outputs, model_inputs=inputs)
+    elapsed = time.perf_counter() - started
+
+    answer = (result.get("content") or "").strip()
+    if not answer and horizon_i is None:
+        answer = "(empty response)"
+    ts = result.get("timeseries")
+    forecast_values = None
+    forecast_text = "No numeric forecast (understanding route)."
+    if ts is not None and ts.get("values"):
+        forecast_values = [float(v) for v in ts["values"]]
+        forecast_text = ", ".join(f"{v:.4g}" for v in forecast_values)
+
+    plot = render_plot(series, forecast_values, result.get("target_series_index"))
+    stats = (
+        f"route: {'forecast' if horizon_i is not None else 'understanding'} · "
+        f"prompt tokens: {result.get('prompt_tokens')} · "
+        f"completion tokens: {result.get('completion_tokens')} · "
+        f"finish: {result.get('finish_reason')} · {elapsed:.1f}s on GPU"
+    )
+    return answer, forecast_text, plot, stats
+
+
+CSS = """
+#col-container { max-width: 1150px; margin: 0 auto; }
+.dark .gradio-container { color: var(--body-text-color); }
+"""
+
+with gr.Blocks(theme=gr.themes.Citrus(), css=CSS, title="TimeBraid") as demo:
+    with gr.Column(elem_id="col-container"):
+        gr.Markdown(
+            """
+            # 🧵 TimeBraid: Unifying Time Series and Language
+            Ask questions about numeric series, or request a point forecast.
+            One series per line, values comma- or space-separated. Leave the series
+            box empty for a pure text question. Set **Forecast horizon** to 0 for
+            analysis-only, or a positive number to get numeric forecast values.
+            **Model:** [XinyueWangg/TimeBraid-2.5B](https://huggingface.co/XinyueWangg/TimeBraid-2.5B) ·
+            **Code:** [GitHub](https://github.com/CharonWangg/TimeBraid) (Apache-2.0)
+            """
+        )
+        with gr.Row():
+            with gr.Column(scale=1):
+                prompt = gr.Textbox(
+                    label="Prompt",
+                    placeholder="e.g. Describe the dominant trend and turning points, or forecast the next 8 values.",
+                    lines=3,
+                )
+                series_box = gr.Textbox(
+                    label="Time series (one series per line)",
+                    placeholder="101.2, 101.8, 102.1, 102.5, 103.0, 103.4\n200, 208, 215, ...",
+                    lines=5,
+                )
+                horizon = gr.Number(
+                    label="Forecast horizon (0 = no forecast)",
+                    value=0,
+                    precision=0,
+                    minimum=0,
+                )
+                run_btn = gr.Button("Run", variant="primary")
+                with gr.Accordion("Advanced settings", open=False):
+                    target_idx = gr.Number(
+                        label="Target series (1-based; multi-series forecasts)",
+                        value=1,
+                        precision=0,
+                        minimum=1,
+                    )
+                    max_new_tokens = gr.Slider(
+                        label="Max new text tokens",
+                        minimum=64,
+                        maximum=1024,
+                        value=256,
+                        step=32,
+                    )
+            with gr.Column(scale=1):
+                answer_out = gr.Textbox(label="Answer", lines=8, show_copy_button=True)
+                forecast_out = gr.Textbox(
+                    label="Forecast values (raw scale)", show_copy_button=True
+                )
+                plot_out = gr.Plot(label="Series and forecast")
+                stats_out = gr.Markdown()
+        gr.Examples(
+            examples=[
+                [
+                    "Forecast the next 8 values from the observed seasonal pattern.",
+                    "100, 102, 105, 107, 103, 101, 99, 100, 103, 106, 108, 104, 102, 100, 101, 104",
+                    8,
+                    1,
+                    256,
+                ],
+                [
+                    "A promotion starts at the first forecast step and is expected to lift demand above the recent seasonal baseline. Forecast the next 6 values.",
+                    "200, 208, 215, 205, 198, 210, 218, 207, 201, 212, 220, 209",
+                    6,
+                    1,
+                    256,
+                ],
+                [
+                    "Describe the dominant trend, major turning points, and whether the series becomes more volatile.",
+                    "2.0, 2.1, 2.2, 2.4, 2.8, 3.1, 3.0, 2.9, 3.4, 3.8, 4.1, 4.3",
+                    0,
+                    1,
+                    256,
+                ],
+                [
+                    "Does this series contain an isolated upward spike? Say only whether it occurs near the beginning, middle, or end; do not give a numeric index. Compare it with the neighboring level.",
+                    "20.1, 20.0, 20.2, 20.1, 20.3, 34.8, 20.2, 20.1, 20.0, 20.2",
+                    0,
+                    1,
+                    256,
+                ],
+                [
+                    "Series 1 is website visits and Series 2 is conversion rate. Compare their trends, volatility, and co-movement.",
+                    "100, 104, 107, 111, 116, 120, 125, 129\n3.2, 3.1, 3.3, 3.4, 3.8, 3.7, 4.0, 4.2",
+                    0,
+                    1,
+                    256,
+                ],
+                [
+                    "Explain in one sentence what a moving average reveals about a time series.",
+                    "",
+                    0,
+                    1,
+                    256,
+                ],
+            ],
+            inputs=[prompt, series_box, horizon, target_idx, max_new_tokens],
+            outputs=[answer_out, forecast_out, plot_out, stats_out],
+            fn=run,
+            cache_examples=True,
+            cache_mode="lazy",
+            label="Examples (from the authors' inference notebook)",
+        )
+    run_btn.click(
+        run,
+        inputs=[prompt, series_box, horizon, target_idx, max_new_tokens],
+        outputs=[answer_out, forecast_out, plot_out, stats_out],
+        api_name="run",
+        concurrency_limit=1,
+    )
+
+if __name__ == "__main__":
+    demo.launch(mcp_server=True)

requirements.txt

@@ -0,0 +1,6 @@
+transformers==4.57.6
+accelerate==1.11.0
+numpy==2.1.3
+safetensors==0.5.3
+matplotlib
+https://huggingface.co/datasets/multimodalart/zerogpu-blackwell-wheels/resolve/main/wheels/pt211-cu130-cp312/flash_attn-2.8.3-cp312-cp312-linux_x86_64.whl

timebraid/__init__.py

@@ -0,0 +1,69 @@
+"""Public TimeBraid package with lazy Hugging Face runtime exports."""
+
+__all__ = [
+    # Runtime types
+    "GenerationRoute",
+    "TimeBraid",
+    "TimeBraidBatchFeature",
+    "TimeBraidConfig",
+    "TimeBraidGenerateOutput",
+    "TimeBraidLossBreakdown",
+    "TimeBraidOutput",
+    "TimeBraidProcessor",
+    # Registration
+    "register_timebraid_auto_classes",
+    # Loading and validation
+    "LoadedTimeBraid",
+    "ValidatedTimeBraidCheckpointSource",
+    "load_timebraid_checkpoint",
+    "load_timebraid_tokenizer",
+    "validate_timebraid_canonical_artifact",
+    "validate_timebraid_checkpoint_source",
+    "validate_timebraid_model",
+]
+
+
+def __getattr__(name: str):
+    if name not in __all__:
+        raise AttributeError(f"module {__name__!r} has no attribute {name!r}")
+    from .model import (
+        LoadedTimeBraid,
+        TimeBraid,
+        TimeBraidConfig,
+        TimeBraidGenerateOutput,
+        TimeBraidLossBreakdown,
+        TimeBraidOutput,
+        ValidatedTimeBraidCheckpointSource,
+        load_timebraid_checkpoint,
+        load_timebraid_tokenizer,
+        validate_timebraid_canonical_artifact,
+        validate_timebraid_checkpoint_source,
+        validate_timebraid_model,
+    )
+    from .model.mot.generation_route import GenerationRoute
+    from .model.timebraid import register_timebraid_auto_classes
+    from .processing_timebraid import TimeBraidBatchFeature, TimeBraidProcessor
+
+    # Importing any public name registers the Hugging Face auto classes, so
+    # `AutoModelForCausalLM.from_pretrained` resolves without remote code.
+    # `register_timebraid_auto_classes` is exported as well, because relying on
+    # an import side effect forces callers to suppress an unused-import warning.
+    register_timebraid_auto_classes()
+    return {
+        "GenerationRoute": GenerationRoute,
+        "LoadedTimeBraid": LoadedTimeBraid,
+        "TimeBraid": TimeBraid,
+        "TimeBraidBatchFeature": TimeBraidBatchFeature,
+        "TimeBraidConfig": TimeBraidConfig,
+        "TimeBraidGenerateOutput": TimeBraidGenerateOutput,
+        "TimeBraidLossBreakdown": TimeBraidLossBreakdown,
+        "TimeBraidOutput": TimeBraidOutput,
+        "TimeBraidProcessor": TimeBraidProcessor,
+        "ValidatedTimeBraidCheckpointSource": ValidatedTimeBraidCheckpointSource,
+        "load_timebraid_checkpoint": load_timebraid_checkpoint,
+        "load_timebraid_tokenizer": load_timebraid_tokenizer,
+        "register_timebraid_auto_classes": register_timebraid_auto_classes,
+        "validate_timebraid_canonical_artifact": validate_timebraid_canonical_artifact,
+        "validate_timebraid_checkpoint_source": validate_timebraid_checkpoint_source,
+        "validate_timebraid_model": validate_timebraid_model,
+    }[name]

timebraid/_vendor/__init__.py

@@ -0,0 +1 @@
+"""Third-party runtime code vendored for reproducible TimeBraid inference."""

timebraid/_vendor/timesfm/__init__.py

@@ -0,0 +1,20 @@
+# Modified by TimeBraid contributors for the vendored PyTorch-only runtime.
+# Copyright 2025 Google LLC
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+"""Vendored TimesFM module surface used by TimeBraid."""
+
+from .timesfm_2p5.timesfm_2p5_torch import TimesFM_2p5_200M_torch_module
+
+__all__ = ["TimesFM_2p5_200M_torch_module"]

timebraid/_vendor/timesfm/configs.py

@@ -0,0 +1,64 @@
+# Modified by TimeBraid contributors for the vendored PyTorch-only runtime.
+# Copyright 2025 Google LLC
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+"""Abstract configs for TimesFM layers."""
+
+import dataclasses
+from typing import Literal
+
+
+@dataclasses.dataclass(frozen=True)
+class ResidualBlockConfig:
+  """Framework-agnostic config for a residual block."""
+
+  input_dims: int
+  hidden_dims: int
+  output_dims: int
+  use_bias: bool
+  activation: Literal["relu", "swish", "none"]
+
+
+@dataclasses.dataclass(frozen=True)
+class RandomFourierFeaturesConfig:
+  """Framework-agnostic config for random fourier features."""
+
+  input_dims: int
+  output_dims: int
+  projection_stddev: float
+  use_bias: bool
+
+
+@dataclasses.dataclass(frozen=True)
+class TransformerConfig:
+  """Framework-agnostic config for a transformer."""
+
+  model_dims: int
+  hidden_dims: int
+  num_heads: int
+  attention_norm: Literal["rms"]
+  feedforward_norm: Literal["rms"]
+  qk_norm: Literal["rms", "none"]
+  use_bias: bool
+  use_rotary_position_embeddings: bool
+  ff_activation: Literal["relu", "swish", "none"]
+  fuse_qkv: bool
+
+
+@dataclasses.dataclass(frozen=True)
+class StackedTransformersConfig:
+  """Framework-agnostic config for a stacked transformers."""
+
+  num_layers: int
+  transformer: TransformerConfig

timebraid/_vendor/timesfm/timesfm_2p5/timesfm_2p5_base.py

@@ -0,0 +1,74 @@
+# Modified by TimeBraid contributors for the vendored PyTorch-only runtime.
+# Copyright 2025 Google LLC
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+"""TimesFM 2p5 base implementation."""
+
+import dataclasses
+
+from .. import configs
+
+ResidualBlockConfig = configs.ResidualBlockConfig
+StackedTransformersConfig = configs.StackedTransformersConfig
+TransformerConfig = configs.TransformerConfig
+
+
+@dataclasses.dataclass(frozen=True)
+class TimesFM_2p5_200M_Definition:
+  """Framework-agnostic config of TimesFM 2.5."""
+
+  context_limit = 16384
+  input_patch_len: int = 32
+  output_patch_len: int = 128
+  output_quantile_len: int = 1024
+  quantiles: list[float] = dataclasses.field(
+    default_factory=lambda: [0.1, 0.2, 0.3, 0.4, 0.5, 0.6, 0.7, 0.8, 0.9]
+  )
+  decode_index: int = 5
+  tokenizer: ResidualBlockConfig = ResidualBlockConfig(
+    input_dims=64,
+    hidden_dims=1280,
+    output_dims=1280,
+    use_bias=True,
+    activation="swish",
+  )
+  stacked_transformers: StackedTransformersConfig = StackedTransformersConfig(
+    num_layers=20,
+    transformer=TransformerConfig(
+      model_dims=1280,
+      hidden_dims=1280,
+      num_heads=16,
+      attention_norm="rms",
+      feedforward_norm="rms",
+      qk_norm="rms",
+      use_bias=False,
+      use_rotary_position_embeddings=True,
+      ff_activation="swish",
+      fuse_qkv=True,
+    ),
+  )
+  output_projection_point: ResidualBlockConfig = ResidualBlockConfig(
+    input_dims=1280,
+    hidden_dims=1280,
+    output_dims=1280,
+    use_bias=False,
+    activation="swish",
+  )
+  output_projection_quantiles: ResidualBlockConfig = ResidualBlockConfig(
+    input_dims=1280,
+    hidden_dims=1280,
+    output_dims=10240,
+    use_bias=False,
+    activation="swish",
+  )

timebraid/_vendor/timesfm/timesfm_2p5/timesfm_2p5_torch.py

@@ -0,0 +1,257 @@
+# Modified by TimeBraid contributors for the vendored PyTorch-only runtime.
+# Copyright 2025 Google LLC
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+"""TimesFM models."""
+
+from typing import Sequence
+
+import numpy as np
+import torch
+from safetensors.torch import load_file
+from torch import nn
+
+from ..torch import dense, transformer, util
+from . import timesfm_2p5_base
+
+revin = util.revin
+
+
+class TimesFM_2p5_200M_torch_module(nn.Module):
+  """TimesFM 2.5 with 200M parameters."""
+
+  config = timesfm_2p5_base.TimesFM_2p5_200M_Definition()
+
+  def __init__(self):
+    super().__init__()
+
+    # Names constants.
+    self.p = self.config.input_patch_len  # 32
+    self.o = self.config.output_patch_len  # 128
+    self.os = self.config.output_quantile_len  # 1024
+    self.m = self.o // self.p  # 4
+    self.x = self.config.stacked_transformers.num_layers  # 20
+    self.h = self.config.stacked_transformers.transformer.num_heads  # 16
+    self.md = self.config.stacked_transformers.transformer.model_dims  # 1280
+    self.hd = self.md // self.h  # 80
+    self.q = len(self.config.quantiles) + 1  # 10
+    self.aridx = self.config.decode_index  # 5
+
+    # Layers.
+    self.tokenizer = dense.ResidualBlock(self.config.tokenizer)
+    self.stacked_xf = nn.ModuleList(
+      [
+        transformer.Transformer(self.config.stacked_transformers.transformer)
+        for _ in range(self.x)
+      ]
+    )
+    self.output_projection_point = dense.ResidualBlock(
+      self.config.output_projection_point
+    )
+    self.output_projection_quantiles = dense.ResidualBlock(
+      self.config.output_projection_quantiles
+    )
+
+    # Device.
+    if torch.cuda.is_available():
+      self.device = torch.device("cuda:0")
+      self.device_count = torch.cuda.device_count()
+    else:
+      self.device = torch.device("cpu")
+      self.device_count = 1
+
+  def load_checkpoint(self, path: str, **kwargs):
+    """Loads a PyTorch TimesFM model from a checkpoint."""
+    tensors = load_file(path)
+    self.load_state_dict(tensors, strict=True)
+    self.to(self.device)
+    torch_compile = True
+    if "torch_compile" in kwargs:
+      torch_compile = kwargs["torch_compile"]
+    if torch_compile:
+      print("Compiling model...")
+      self = torch.compile(self)
+
+    self.eval()
+
+  def forward(
+    self,
+    inputs: torch.Tensor,
+    masks: torch.Tensor,
+    decode_caches: list[util.DecodeCache] | None = None,
+  ):
+    tokenizer_inputs = torch.cat([inputs, masks.to(inputs.dtype)], dim=-1)
+    input_embeddings = self.tokenizer(tokenizer_inputs)
+
+    if decode_caches is None:
+      decode_caches = [None] * self.x
+
+    output_embeddings = input_embeddings
+    new_decode_caches = []
+    for i, layer in enumerate(self.stacked_xf):
+      output_embeddings, new_cache = layer(
+        output_embeddings, masks[..., -1], decode_caches[i]
+      )
+      new_decode_caches.append(new_cache)
+    output_ts = self.output_projection_point(output_embeddings)
+    output_quantile_spread = self.output_projection_quantiles(output_embeddings)
+
+    return (
+      input_embeddings,
+      output_embeddings,
+      output_ts,
+      output_quantile_spread,
+    ), new_decode_caches
+
+  def decode(self, horizon: int, inputs, masks):
+    """Decodes the time series."""
+
+    with torch.no_grad():
+      batch_size, context = inputs.shape[0], inputs.shape[1]
+      num_decode_steps = (horizon - 1) // self.o
+      num_input_patches = context // self.p
+      decode_cache_size = num_input_patches + num_decode_steps * self.m
+
+      # Prefill
+      patched_inputs = torch.reshape(inputs, (batch_size, -1, self.p))
+      patched_masks = torch.reshape(masks, (batch_size, -1, self.p))
+
+      # running stats
+      n = torch.zeros(batch_size, device=inputs.device)
+      mu = torch.zeros(batch_size, device=inputs.device)
+      sigma = torch.zeros(batch_size, device=inputs.device)
+      patch_mu = []
+      patch_sigma = []
+      for i in range(num_input_patches):
+        (n, mu, sigma), _ = util.update_running_stats(
+          n, mu, sigma, patched_inputs[:, i], patched_masks[:, i]
+        )
+        patch_mu.append(mu)
+        patch_sigma.append(sigma)
+      last_n, last_mu, last_sigma = n, mu, sigma
+      context_mu = torch.stack(patch_mu, dim=1)
+      context_sigma = torch.stack(patch_sigma, dim=1)
+
+      decode_caches = [
+        util.DecodeCache(
+          next_index=torch.zeros(batch_size, dtype=torch.int32, device=inputs.device),
+          num_masked=torch.zeros(batch_size, dtype=torch.int32, device=inputs.device),
+          key=torch.zeros(
+            batch_size,
+            decode_cache_size,
+            self.h,
+            self.hd,
+            device=inputs.device,
+          ),
+          value=torch.zeros(
+            batch_size,
+            decode_cache_size,
+            self.h,
+            self.hd,
+            device=inputs.device,
+          ),
+        )
+        for _ in range(self.x)
+      ]
+
+      normed_inputs = revin(patched_inputs, context_mu, context_sigma, reverse=False)
+      normed_inputs = torch.where(patched_masks, 0.0, normed_inputs)
+      (_, _, normed_outputs, normed_quantile_spread), decode_caches = self(
+        normed_inputs, patched_masks, decode_caches
+      )
+      renormed_outputs = torch.reshape(
+        revin(normed_outputs, context_mu, context_sigma, reverse=True),
+        (batch_size, -1, self.o, self.q),
+      )
+      renormed_quantile_spread = torch.reshape(
+        revin(normed_quantile_spread, context_mu, context_sigma, reverse=True),
+        (batch_size, -1, self.os, self.q),
+      )[:, -1, ...]
+
+      # Autogressive decode
+      ar_outputs = []
+      last_renormed_output = renormed_outputs[:, -1, :, self.aridx]
+
+      for _ in range(num_decode_steps):
+        new_patched_input = torch.reshape(
+          last_renormed_output, (batch_size, self.m, self.p)
+        )
+        new_mask = torch.zeros_like(new_patched_input, dtype=torch.bool)
+
+        n, mu, sigma = last_n, last_mu, last_sigma
+        new_mus, new_sigmas = [], []
+        for i in range(self.m):
+          (n, mu, sigma), _ = util.update_running_stats(
+            n, mu, sigma, new_patched_input[:, i], new_mask[:, i]
+          )
+          new_mus.append(mu)
+          new_sigmas.append(sigma)
+        last_n, last_mu, last_sigma = n, mu, sigma
+        new_mu = torch.stack(new_mus, dim=1)
+        new_sigma = torch.stack(new_sigmas, dim=1)
+
+        new_normed_input = revin(new_patched_input, new_mu, new_sigma, reverse=False)
+        (_, _, new_normed_output, _), decode_caches = self(
+          new_normed_input, new_mask, decode_caches
+        )
+
+        new_renormed_output = torch.reshape(
+          revin(new_normed_output, new_mu, new_sigma, reverse=True),
+          (batch_size, self.m, self.o, self.q),
+        )
+        ar_outputs.append(new_renormed_output[:, -1, ...])
+        last_renormed_output = new_renormed_output[:, -1, :, self.aridx]
+
+      if num_decode_steps > 0:
+        ar_renormed_outputs = torch.stack(ar_outputs, dim=1)
+      else:
+        ar_renormed_outputs = None
+
+    return renormed_outputs, renormed_quantile_spread, ar_renormed_outputs
+
+  def forecast_naive(
+    self, horizon: int, inputs: Sequence[np.ndarray]
+  ) -> list[np.ndarray]:
+    """Forecasts the time series.
+
+    This is a naive implementation for debugging purposes. No forecasting
+    flags are used here. Forecasting quality can be subpar.
+
+    Args:
+      horizon: The number of time points to forecast.
+      inputs: A sequence of numpy arrays, each representing a time series to
+        query forecast for.
+
+    Returns:
+      A list of numpy arrays of forecasts.
+    """
+    outputs = []
+    for each_input in inputs:
+      input_t = torch.tensor(each_input, dtype=torch.float32)
+      mask = torch.zeros_like(input_t, dtype=torch.bool)
+      len_front_mask = self.p - (len(each_input) % self.p)
+      if len_front_mask < self.p:
+        input_t = torch.cat(
+          [torch.zeros(len_front_mask, dtype=torch.float32), input_t], dim=0
+        )
+        mask = torch.cat([torch.ones(len_front_mask, dtype=torch.bool), mask], dim=0)
+      input_t = input_t[None, ...]
+      mask = mask[None, ...]
+      t_pf, _, t_ar = self.decode(horizon, input_t, mask)
+      to_concat = [t_pf[:, -1, ...]]
+      if t_ar is not None:
+        to_concat.append(t_ar.reshape(1, -1, self.q))
+      torch_forecast = torch.cat(to_concat, dim=1)[..., :horizon]
+      torch_forecast = torch_forecast.squeeze(0)
+      outputs.append(torch_forecast.detach().cpu().numpy())
+    return outputs

timebraid/_vendor/timesfm/torch/__init__.py

@@ -0,0 +1,14 @@
+# Modified by TimeBraid contributors for the vendored PyTorch-only runtime.
+# Copyright 2025 Google LLC
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.

timebraid/_vendor/timesfm/torch/dense.py

@@ -0,0 +1,95 @@
+# Modified by TimeBraid contributors for the vendored PyTorch-only runtime.
+# Copyright 2025 Google LLC
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+"""Dense layers for TimesFM."""
+
+import torch
+from torch import nn
+
+from .. import configs
+
+
+class ResidualBlock(nn.Module):
+  """Residual block with two linear layers and a linear residual connection."""
+
+  def __init__(self, config: configs.ResidualBlockConfig):
+    super().__init__()
+    self.config = config
+    self.hidden_layer = nn.Linear(
+        in_features=config.input_dims,
+        out_features=config.hidden_dims,
+        bias=config.use_bias,
+    )
+    self.output_layer = nn.Linear(
+        in_features=config.hidden_dims,
+        out_features=config.output_dims,
+        bias=config.use_bias,
+    )
+    self.residual_layer = nn.Linear(
+        in_features=config.input_dims,
+        out_features=config.output_dims,
+        bias=config.use_bias,
+    )
+    if config.activation == "relu":
+      self.activation = nn.ReLU()
+    elif config.activation == "swish":
+      self.activation = nn.SiLU()
+    elif config.activation == "none":
+      self.activation = nn.Identity()
+    else:
+      raise ValueError(f"Activation: {config.activation} not supported.")
+
+  def forward(self, x: torch.Tensor) -> torch.Tensor:
+    return self.output_layer(
+        self.activation(self.hidden_layer(x))
+    ) + self.residual_layer(x)
+
+
+class RandomFourierFeatures(nn.Module):
+  """Random Fourier features layer."""
+
+  def __init__(self, config: configs.RandomFourierFeaturesConfig):
+    super().__init__()
+    self.config = config
+
+    if config.output_dims % 4 != 0:
+      raise ValueError(
+          f"Output dims must be a multiple of 4: {config.output_dims} % 4 != 0."
+      )
+    num_projected_features = config.output_dims // 4
+
+    self.phase_shifts = nn.Parameter(torch.zeros(2, num_projected_features))
+    self.projection_layer = nn.Linear(
+        in_features=config.input_dims,
+        out_features=num_projected_features,
+        bias=config.use_bias,
+    )
+    self.residual_layer = nn.Linear(
+        in_features=config.input_dims,
+        out_features=config.output_dims,
+        bias=config.use_bias,
+    )
+
+  def forward(self, x: torch.Tensor) -> torch.Tensor:
+    projected = self.projection_layer(x)
+    cos_features = torch.cos(projected)
+    sin_features = torch.sin(projected)
+    sq_wave_1 = torch.sign(torch.sin(projected + self.phase_shifts[0, :]))
+    sq_wave_2 = torch.sign(torch.sin(projected + self.phase_shifts[1, :]))
+    fourier_features = torch.cat(
+        [cos_features, sin_features, sq_wave_1, sq_wave_2], dim=-1
+    )
+    residual = self.residual_layer(x)
+    return fourier_features + residual

timebraid/_vendor/timesfm/torch/normalization.py

@@ -0,0 +1,40 @@
+# Modified by TimeBraid contributors for the vendored PyTorch-only runtime.
+# Copyright 2025 Google LLC
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+"""Normalization layers for TimesFM."""
+
+import torch
+from torch import nn
+
+
+class RMSNorm(nn.Module):
+  """RMS normalization."""
+
+  def __init__(
+      self,
+      num_features: int,
+      *,
+      epsilon: float = 1e-6,
+  ):
+    super().__init__()
+    self.scale = nn.Parameter(torch.zeros(num_features))
+    self.num_features = num_features
+    self.epsilon = epsilon
+
+  def forward(self, inputs: torch.Tensor) -> torch.Tensor:
+    var = torch.mean(torch.square(inputs), dim=-1, keepdim=True)
+    normed_inputs = inputs * torch.rsqrt(var + self.epsilon)
+    normed_inputs = normed_inputs * self.scale
+    return normed_inputs

timebraid/_vendor/timesfm/torch/transformer.py

@@ -0,0 +1,380 @@
+# Modified by TimeBraid contributors for the vendored PyTorch-only runtime.
+# Copyright 2025 Google LLC
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+"""Transformer layers for TimesFM."""
+
+import math
+from typing import Callable
+
+import torch
+import torch.nn.functional as F
+from torch import nn
+
+from .. import configs
+from . import normalization, util
+
+LayerNorm = nn.LayerNorm
+RMSNorm = normalization.RMSNorm
+DecodeCache = util.DecodeCache
+
+
+def make_attn_mask(
+  query_length: int,
+  num_all_masked_kv: torch.Tensor,
+  query_index_offset: torch.Tensor | None = None,
+  kv_length: int = 0,
+) -> torch.Tensor:
+  """Makes attention mask."""
+  if kv_length == 0:
+    kv_length = query_length
+
+  q_index = torch.arange(query_length, device=num_all_masked_kv.device)[
+    None, None, :, None
+  ]
+  if query_index_offset is not None:
+    q_index = q_index + query_index_offset[:, None, None, None]
+  kv_index = torch.arange(kv_length, device=num_all_masked_kv.device)[
+    None, None, None, :
+  ]
+  return torch.logical_and(
+    q_index >= kv_index,
+    kv_index >= num_all_masked_kv[:, None, None, None],
+  )
+
+
+class RotaryPositionalEmbedding(nn.Module):
+  """Rotary positional embedding."""
+
+  def __init__(
+    self,
+    embedding_dims: int,
+    min_timescale: float = 1.0,
+    max_timescale: float = 10000.0,
+  ):
+    super().__init__()
+    self.embedding_dims = embedding_dims
+    self.min_timescale = min_timescale
+    self.max_timescale = max_timescale
+
+  def forward(
+    self,
+    inputs: torch.Tensor,
+    position: torch.Tensor | None = None,
+  ):
+    """Generates a JTensor of sinusoids with different frequencies."""
+    if self.embedding_dims != inputs.shape[-1]:
+      raise ValueError(
+        "The embedding dims of the rotary position embedding"
+        "must match the hidden dimension of the inputs."
+      )
+    half_embedding_dim = self.embedding_dims // 2
+
+    # Keep RoPE numerics stable by building the sinusoid in fp32, but cast the
+    # final trigonometric factors back to the input dtype before mixing them
+    # with the activations. Without this, bf16 TimesFM activations get promoted
+    # to fp32 in the rotary path while V stays bf16, and fused SDPA then fails
+    # because Q/K/V no longer share a dtype.
+    rope_dtype = torch.float32
+    fraction = (
+      2
+      * torch.arange(0, half_embedding_dim, device=inputs.device, dtype=rope_dtype)
+      / self.embedding_dims
+    )
+    timescale = (
+      self.min_timescale * (self.max_timescale / self.min_timescale) ** fraction
+    ).to(device=inputs.device, dtype=rope_dtype)
+    if position is None:
+      seq_length = inputs.shape[1]
+      position = torch.arange(seq_length, dtype=rope_dtype, device=inputs.device)[
+        None, :
+      ]
+    else:
+      position = position.to(device=inputs.device, dtype=rope_dtype)
+
+    if len(inputs.shape) == 4:
+      position = position[..., None, None]
+      timescale = timescale[None, None, None, :]
+    elif len(inputs.shape) == 3:
+      position = position[..., None]
+      timescale = timescale[None, None, :]
+    else:
+      raise ValueError("Inputs must be of rank 3 or 4.")
+
+    sinusoid_inp = position / timescale
+    sin = torch.sin(sinusoid_inp).to(dtype=inputs.dtype)
+    cos = torch.cos(sinusoid_inp).to(dtype=inputs.dtype)
+    first_half, second_half = torch.chunk(inputs, 2, dim=-1)
+    first_part = first_half * cos - second_half * sin
+    second_part = second_half * cos + first_half * sin
+    return torch.cat([first_part, second_part], dim=-1)
+
+
+def _dot_product_attention(
+  query,
+  key,
+  value,
+  mask=None,
+):
+  """Computes dot-product attention given query, key, and value."""
+  attn_weights = torch.einsum("...qhd,...khd->...hqk", query, key)
+  if mask is not None:
+    attn_weights = torch.where(
+      mask, attn_weights, -torch.finfo(attn_weights.dtype).max / 2
+    )
+
+  attn_weights = F.softmax(attn_weights, dim=-1)
+
+  return torch.einsum("...hqk,...khd->...qhd", attn_weights, value)
+
+
+def _torch_dot_product_attention(query, key, value, mask=None):
+  """
+  Performs the exact same (unscaled) attention as the above function,
+  but using the fast and fused F.scaled_dot_product_attention kernel.
+  """
+
+  # 1. Permute inputs from (B, L, H, D) to the expected (B, H, L, D)
+  query = query.permute(0, 2, 1, 3)
+  key = key.permute(0, 2, 1, 3)
+  value = value.permute(0, 2, 1, 3)
+
+  # 2. Call the fused attention kernel
+  #    - Pass the mask to `attn_mask`.
+  #    - Set `scale=1.0` to disable the default 1/sqrt(d_k) scaling.
+  output = F.scaled_dot_product_attention(query, key, value, attn_mask=mask, scale=1.0)
+
+  # 3. Permute the output back to the original (B, L, H, D) layout
+  output = output.permute(0, 2, 1, 3)
+
+  return output
+
+
+class PerDimScale(nn.Module):
+  """Per-dimension scaling."""
+
+  def __init__(self, num_dims: int):
+    super().__init__()
+    self.num_dims = num_dims
+    self.per_dim_scale = nn.Parameter(torch.zeros(num_dims))
+
+  def forward(self, x: torch.Tensor) -> torch.Tensor:
+    scale_factor = (
+      1.442695041 / math.sqrt(self.num_dims) * F.softplus(self.per_dim_scale)
+    )
+    return x * scale_factor
+
+
+class MultiHeadAttention(nn.Module):
+  """Multi-head attention."""
+
+  def __init__(
+    self,
+    num_heads: int,
+    in_features: int,
+    *,
+    use_per_dim_scale: bool = True,
+    use_rotary_position_embeddings: bool = True,
+    use_bias: bool = False,
+    attention_fn: Callable[..., torch.Tensor] = _torch_dot_product_attention,
+    qk_norm: str = "rms",
+    fuse_qkv: bool = False,
+  ):
+    super().__init__()
+    self.num_heads = num_heads
+    self.in_features = in_features
+    self.head_dim = in_features // num_heads
+    self.use_bias = use_bias
+    self.attention_fn = attention_fn
+    self.qk_norm = qk_norm
+    self.fuse_qkv = fuse_qkv
+
+    if self.in_features % self.num_heads != 0:
+      raise ValueError(
+        f"Memory dimension ({self.in_features}) must be divisible by "
+        f"'num_heads' heads ({self.num_heads})."
+      )
+
+    if self.fuse_qkv:
+      self.qkv_proj = nn.Linear(self.in_features, 3 * self.in_features, bias=use_bias)
+    else:
+      self.query = nn.Linear(self.in_features, self.in_features, bias=use_bias)
+      self.key = nn.Linear(self.in_features, self.in_features, bias=use_bias)
+      self.value = nn.Linear(self.in_features, self.in_features, bias=use_bias)
+    self.out = nn.Linear(self.in_features, self.in_features, bias=use_bias)
+
+    if self.qk_norm == "rms":
+      self.query_ln = RMSNorm(self.head_dim)
+      self.key_ln = RMSNorm(self.head_dim)
+    else:
+      self.query_ln = nn.Identity()
+      self.key_ln = nn.Identity()
+
+    self.use_rotary_position_embeddings = use_rotary_position_embeddings
+    if self.use_rotary_position_embeddings:
+      self.rotary_position_embedding = RotaryPositionalEmbedding(
+        embedding_dims=self.head_dim,
+      )
+
+    self.use_per_dim_scale = use_per_dim_scale
+    if use_per_dim_scale:
+      self.per_dim_scale = PerDimScale(num_dims=self.head_dim)
+
+  def forward(
+    self,
+    inputs_q: torch.Tensor,
+    *,
+    decode_cache: DecodeCache | None = None,
+    patch_mask: torch.Tensor | None = None,
+  ) -> tuple[torch.Tensor, DecodeCache | None]:
+    b, n_patches, _ = inputs_q.shape
+    if patch_mask is None:
+      patch_mask = torch.zeros(b, n_patches, dtype=torch.bool, device=inputs_q.device)
+
+    if self.fuse_qkv:
+      qkv = self.qkv_proj(inputs_q)
+      query, key, value = torch.chunk(qkv, 3, dim=-1)
+      query = query.view(b, n_patches, self.num_heads, self.head_dim)
+      key = key.view(b, n_patches, self.num_heads, self.head_dim)
+      value = value.view(b, n_patches, self.num_heads, self.head_dim)
+    else:
+      query = self.query(inputs_q).view(b, n_patches, self.num_heads, self.head_dim)
+      key = self.key(inputs_q).view(b, n_patches, self.num_heads, self.head_dim)
+      value = self.value(inputs_q).view(b, n_patches, self.num_heads, self.head_dim)
+
+    if decode_cache is None:
+      num_masked = torch.sum(patch_mask.to(torch.int32), dim=-1)
+      next_index = torch.zeros_like(num_masked, dtype=torch.int32)
+    else:
+      num_masked = (
+        torch.sum(patch_mask.to(torch.int32), dim=-1) + decode_cache.num_masked
+      )
+      next_index = decode_cache.next_index.clone()
+
+    if self.use_rotary_position_embeddings:
+      position = (
+        torch.arange(n_patches, device=inputs_q.device)[None, :]
+        + next_index[:, None]
+        - num_masked[:, None]
+      )
+      query = self.rotary_position_embedding(query, position)
+      key = self.rotary_position_embedding(key, position)
+
+    query = self.query_ln(query)
+    key = self.key_ln(key)
+
+    if self.use_per_dim_scale:
+      query = self.per_dim_scale(query)
+
+    if decode_cache is not None:
+      _, decode_cache_size, _, _ = decode_cache.value.shape
+
+      start = decode_cache.next_index[0]
+      end = start + n_patches
+
+      # Perform a single, vectorized slice assignment for the entire batch.
+      # This is vastly more efficient than a Python for-loop.
+
+      decode_cache.key[:, start:end] = key
+      decode_cache.value[:, start:end] = value
+
+      key = decode_cache.key
+      value = decode_cache.value
+      decode_cache.next_index += n_patches
+      decode_cache.num_masked = num_masked
+      attn_mask = make_attn_mask(
+        query_length=n_patches,
+        num_all_masked_kv=num_masked,
+        query_index_offset=next_index,
+        kv_length=decode_cache_size,
+      )
+    else:
+      attn_mask = make_attn_mask(query_length=n_patches, num_all_masked_kv=num_masked)
+
+    x = self.attention_fn(
+      query,
+      key,
+      value,
+      mask=attn_mask,
+    )
+
+    x = x.reshape(b, n_patches, self.in_features)
+    out = self.out(x)
+    return out, decode_cache
+
+
+class Transformer(nn.Module):
+  """Classic Transformer used in TimesFM."""
+
+  def __init__(self, config: configs.TransformerConfig):
+    super().__init__()
+    self.config = config
+
+    if config.attention_norm == "rms":
+      self.pre_attn_ln = RMSNorm(num_features=config.model_dims)
+      self.post_attn_ln = RMSNorm(num_features=config.model_dims)
+    else:
+      raise ValueError(f"Layer norm: {config.attention_norm} not supported.")
+
+    self.attn = MultiHeadAttention(
+      num_heads=config.num_heads,
+      in_features=config.model_dims,
+      use_per_dim_scale=True,
+      use_rotary_position_embeddings=config.use_rotary_position_embeddings,
+      qk_norm=config.qk_norm,
+      fuse_qkv=config.fuse_qkv,
+    )
+
+    if config.feedforward_norm == "rms":
+      self.pre_ff_ln = RMSNorm(num_features=config.model_dims)
+      self.post_ff_ln = RMSNorm(num_features=config.model_dims)
+    else:
+      raise ValueError(f"Layer norm: {config.feedforward_norm} not supported.")
+
+    self.ff0 = nn.Linear(
+      in_features=config.model_dims,
+      out_features=config.hidden_dims,
+      bias=config.use_bias,
+    )
+    self.ff1 = nn.Linear(
+      in_features=config.hidden_dims,
+      out_features=config.model_dims,
+      bias=config.use_bias,
+    )
+    if config.ff_activation == "relu":
+      self.activation = nn.ReLU()
+    elif config.ff_activation == "swish":
+      self.activation = nn.SiLU()
+    elif config.ff_activation == "none":
+      self.activation = nn.Identity()
+    else:
+      raise ValueError(f"Activation: {config.ff_activation} not supported.")
+
+  def forward(
+    self,
+    input_embeddings: torch.Tensor,
+    patch_mask: torch.Tensor,
+    decode_cache: DecodeCache | None = None,
+  ) -> tuple[torch.Tensor, DecodeCache | None]:
+    attn_output, decode_cache = self.attn(
+      inputs_q=self.pre_attn_ln(input_embeddings),
+      decode_cache=decode_cache,
+      patch_mask=patch_mask,
+    )
+    attn_output = self.post_attn_ln(attn_output) + input_embeddings
+    output_embeddings = (
+      self.post_ff_ln(self.ff1(self.activation(self.ff0(self.pre_ff_ln(attn_output)))))
+      + attn_output
+    )
+    return output_embeddings, decode_cache

timebraid/_vendor/timesfm/torch/util.py

@@ -0,0 +1,95 @@
+# Modified by TimeBraid contributors for the vendored PyTorch-only runtime.
+# Copyright 2025 Google LLC
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+"""PyTorch utility functions for TimesFM layers."""
+
+import dataclasses
+import torch
+
+_TOLERANCE = 1e-6
+
+
+@dataclasses.dataclass(frozen=False)
+class DecodeCache:
+  """Cache for decoding."""
+
+  next_index: torch.Tensor
+  num_masked: torch.Tensor
+  key: torch.Tensor
+  value: torch.Tensor
+
+
+def update_running_stats(
+    n: torch.Tensor,
+    mu: torch.Tensor,
+    sigma: torch.Tensor,
+    x: torch.Tensor,
+    mask: torch.Tensor,
+) -> tuple[
+    tuple[torch.Tensor, torch.Tensor, torch.Tensor],
+    tuple[torch.Tensor, torch.Tensor, torch.Tensor],
+]:
+  """Updates the running stats."""
+  is_legit = torch.logical_not(mask)
+  inc_n = torch.sum(is_legit.to(x.dtype), dim=-1)
+
+  inc_mu_numerator = torch.sum(x * is_legit, dim=-1)
+  inc_n_safe = torch.where(inc_n == 0, 1.0, inc_n)
+  inc_mu = inc_mu_numerator / inc_n_safe
+  inc_mu = torch.where(inc_n == 0, 0.0, inc_mu)
+
+  inc_var_numerator = torch.sum(
+      ((x - inc_mu.unsqueeze(-1)) ** 2) * is_legit, dim=-1
+  )
+  inc_var = inc_var_numerator / inc_n_safe
+  inc_var = torch.where(inc_n == 0, 0.0, inc_var)
+  inc_sigma = torch.sqrt(inc_var)
+
+  new_n = n + inc_n
+  new_n_safe = torch.where(new_n == 0, 1.0, new_n)
+
+  new_mu = (n * mu + inc_mu * inc_n) / new_n_safe
+  new_mu = torch.where(new_n == 0, 0.0, new_mu)
+
+  term1 = n * sigma.pow(2)
+  term2 = inc_n * inc_sigma.pow(2)
+  term3 = n * (mu - new_mu).pow(2)
+  term4 = inc_n * (inc_mu - new_mu).pow(2)
+
+  new_var = (term1 + term2 + term3 + term4) / new_n_safe
+  new_var = torch.where(new_n == 0, 0.0, new_var)
+  new_sigma = torch.sqrt(torch.clamp(new_var, min=0.0))
+
+  return (w := (new_n, new_mu, new_sigma), w)
+
+
+def revin(
+    x: torch.Tensor,
+    mu: torch.Tensor,
+    sigma: torch.Tensor,
+    reverse: bool = False,
+):
+  """Reversible instance normalization."""
+  if len(mu.shape) == len(x.shape) - 1:
+    mu = mu[..., None]
+    sigma = sigma[..., None]
+  elif len(mu.shape) == len(x.shape) - 2:
+    mu = mu[..., None, None]
+    sigma = sigma[..., None, None]
+
+  if reverse:
+    return x * sigma + mu
+  else:
+    return (x - mu) / torch.where(sigma < _TOLERANCE, 1.0, sigma)

timebraid/data/__init__.py

@@ -0,0 +1 @@
+"""Inference tensorization internals for the TimeBraid processor."""

timebraid/data/mot_utils.py

@@ -0,0 +1,612 @@
+# Copyright 2026 the TimeBraid contributors.
+"""Inference-only validation and tensorization for TimeBraid time-series spans."""
+
+from __future__ import annotations
+
+import math
+from dataclasses import dataclass
+from typing import Any
+
+import torch
+
+from ..model.mot.structures import ROLE_CONTEXT, ROLE_OBSERVED, ROLE_TARGET
+
+# Derived from the model-side constants rather than restated. This module
+# writes `role_id` and the MoT runtime reads it back, so two independent
+# literals would route spans to the wrong expert the moment either side
+# renumbered — with correctly shaped tensors and no error.
+ROLE_TO_ID = {
+    "observed": ROLE_OBSERVED,
+    "context": ROLE_CONTEXT,
+    "target": ROLE_TARGET,
+}
+_GENERIC_OPEN_TAG = "<ts>"
+_GENERIC_CLOSE_TAG = "</ts>"
+_FLOAT32_MAX = float(torch.finfo(torch.float32).max)
+
+
+@dataclass(frozen=True, slots=True)
+class MoTNormalizedSpan:
+    """
+    One validated MoT span payload entry.
+
+    `loss_start` and `loss_roi_mask` remain part of the model payload ABI even
+    though the inference processor constructs conditioning-only spans.
+    """
+
+    length: int
+    role_id: int
+    segment_id: int
+    values: list[float]
+    loss_start: int = 0
+    loss_roi_mask: list[float] | None = None
+    text_start_token_idx: int | None = None
+    text_end_token_idx: int | None = None
+
+
+def _normalize_special_token_entry(token: Any) -> str:
+    """
+    Normalize tokenizer special-token declarations to their textual surface.
+
+    Tokenizers expose special tokens through a mix of:
+    - raw strings
+    - dicts with `content`
+    - AddedToken-like objects with `.content`
+
+    The processor uses this normalization to derive the tokenizer's exact TS
+    delimiter inventory.
+    """
+    if isinstance(token, str):
+        return token
+    if isinstance(token, dict):
+        content = token.get("content")
+        if isinstance(content, str):
+            return content
+        raise RuntimeError(
+            f"Tokenizer special-token dict is missing string `content`: {token!r}"
+        )
+
+    content = getattr(token, "content", None)
+    if isinstance(content, str):
+        return content
+
+    raise RuntimeError(
+        "Unsupported tokenizer special-token entry while resolving MoT TS delimiters: "
+        f"{type(token)} {token!r}"
+    )
+
+
+def _format_ts_tag(length: int | None, *, is_close: bool) -> str:
+    if length is not None:
+        raise RuntimeError(
+            f"Length-tagged TS delimiters are no longer supported: {length}."
+        )
+    return _GENERIC_CLOSE_TAG if is_close else _GENERIC_OPEN_TAG
+
+
+def _parse_ts_tag(token: str) -> tuple[int | None, bool]:
+    """
+    Parse one textual TS delimiter into `(tagged_len, is_close)`.
+
+    Canonical generic tags are `<ts>` and `</ts>`.
+    """
+    if token == _GENERIC_OPEN_TAG:
+        return None, False
+    if token == _GENERIC_CLOSE_TAG:
+        return None, True
+
+    raise RuntimeError(f"Unsupported TS tag syntax: {token}")
+
+
+def resolve_ts_delimiter_id_maps(
+    tokenizer: Any,
+) -> tuple[dict[int, int | None], dict[int, int | None]]:
+    """
+    Resolve tokenizer-visible TS delimiter ids to their optional tagged length.
+
+    The only supported tokenizer-visible TS delimiters are `<ts>` and `</ts>`.
+    """
+    candidate_sources: list[tuple[str, Any]] = []
+    for attr_name in (
+        "additional_special_tokens",
+        "extra_special_tokens",
+        "all_special_tokens",
+    ):
+        declared_tokens = getattr(tokenizer, attr_name, None)
+        if declared_tokens is not None:
+            candidate_sources.append((f"tokenizer.{attr_name}", declared_tokens))
+
+    added_tokens_decoder = getattr(tokenizer, "added_tokens_decoder", None)
+    if isinstance(added_tokens_decoder, dict) and len(added_tokens_decoder) > 0:
+        candidate_sources.append(
+            ("tokenizer.added_tokens_decoder", list(added_tokens_decoder.values()))
+        )
+
+    open_id_to_len: dict[int, int | None] = {}
+    close_id_to_len: dict[int, int | None] = {}
+    for source_name, declared_tokens in candidate_sources:
+        if isinstance(declared_tokens, dict):
+            iterable = declared_tokens.values()
+        elif isinstance(declared_tokens, (list, tuple, set)):
+            iterable = declared_tokens
+        else:
+            raise RuntimeError(
+                f"Unsupported tokenizer special-token collection while resolving MoT TS delimiters: "
+                f"{source_name} -> {type(declared_tokens)}"
+            )
+
+        for token in iterable:
+            token_str = _normalize_special_token_entry(token)
+            try:
+                tagged_len, is_close = _parse_ts_tag(token_str)
+            except RuntimeError:
+                continue
+
+            token_id = int(tokenizer.convert_tokens_to_ids(token_str))
+            target_map = close_id_to_len if is_close else open_id_to_len
+            previous = target_map.get(token_id)
+            if previous is not None or token_id in target_map:
+                if previous != tagged_len:
+                    raise RuntimeError(
+                        "Tokenizer resolves one TS delimiter id to multiple tagged lengths: "
+                        f"token={token_str!r}, token_id={token_id}, existing_len={previous}, new_len={tagged_len}"
+                    )
+                continue
+            target_map[token_id] = tagged_len
+
+    if not open_id_to_len or not close_id_to_len:
+        raise RuntimeError(
+            "MoT expected tokenizer-visible TS delimiters, but failed to resolve a complete id inventory. "
+            f"open_ids={sorted(open_id_to_len)}, close_ids={sorted(close_id_to_len)}."
+        )
+    return open_id_to_len, close_id_to_len
+
+
+def find_tokenized_ts_spans(
+    token_ids: list[int],
+    *,
+    open_id_to_len: dict[int, int | None],
+    close_id_to_len: dict[int, int | None],
+) -> list[tuple[int, int, int | None]]:
+    """
+    Parse one tokenized text row into ordered `(start_idx, end_idx, tagged_len)` spans.
+
+    This mirrors the strict delimiter parser the runtime historically performed
+    inside every forward pass. Moving it into preprocessing lets us cache
+    the result as batch metadata without weakening any structural checks.
+    """
+    spans: list[tuple[int, int, int | None]] = []
+    stack: list[tuple[int | None, int]] = []
+    for idx, raw_tid in enumerate(token_ids):
+        tid = int(raw_tid)
+        if tid in open_id_to_len:
+            if stack:
+                raise RuntimeError("Nested TS spans are not supported.")
+            stack.append((open_id_to_len[tid], idx))
+            continue
+
+        if tid not in close_id_to_len:
+            continue
+
+        if not stack:
+            raise RuntimeError(
+                f"Found closing TS token without opening token at position {idx}."
+            )
+
+        tagged_len, start_idx = stack.pop()
+        close_len = close_id_to_len[tid]
+        if close_len != tagged_len:
+            raise RuntimeError(
+                "Mismatched TS delimiters: "
+                f"open={_format_ts_tag(tagged_len, is_close=False)}, "
+                f"close={_format_ts_tag(close_len, is_close=True)}"
+            )
+        spans.append((start_idx, idx, tagged_len))
+
+    if stack:
+        raise RuntimeError("Unclosed TS tags found in sequence.")
+    spans.sort(key=lambda item: item[0])
+    return spans
+
+
+def _require_json_int(value: Any, *, owner: str, field_name: str) -> int:
+    """Require an actual JSON integer for schema fields that define supervision boundaries."""
+    if type(value) is not int:
+        raise RuntimeError(
+            f"{owner} `{field_name}` must be int, got {value!r} ({type(value)})."
+        )
+    return int(value)
+
+
+def normalize_timeseries_spans(
+    spans: Any,
+    *,
+    sample_idx: int,
+    require_segment_id: bool,
+    default_segment_id: int,
+    max_spans_per_sample: int,
+) -> list[MoTNormalizedSpan]:
+    """
+    Validate and normalize one sample's `timeseries` payload.
+
+    Expected schema for each span:
+    {
+      "len": int > 0,
+      "role": "observed" | "context" | "target",
+      "values": list[float] with exact `len`,
+      "segment_id": int > 0    # optional when require_segment_id=False
+      "loss_start": int in [0, len]   # required; len means no TS supervision
+      "loss_roi_mask": list[float]    # optional; raw-point aligned soft ROI in [0, 1]
+    }
+    """
+    if spans is None:
+        return []
+    if not isinstance(spans, list):
+        raise RuntimeError(
+            f"Sample[{sample_idx}] `timeseries` must be list, got {type(spans)}."
+        )
+    if len(spans) > max_spans_per_sample:
+        raise RuntimeError(
+            f"Sample[{sample_idx}] spans exceed max limit: "
+            f"{len(spans)} > {max_spans_per_sample}."
+        )
+
+    normalized: list[MoTNormalizedSpan] = []
+    for span_idx, span in enumerate(spans):
+        if not isinstance(span, dict):
+            raise RuntimeError(
+                f"Sample[{sample_idx}] span[{span_idx}] must be dict, got {type(span)}."
+            )
+
+        for key in ("len", "role", "values", "loss_start"):
+            if key not in span:
+                raise RuntimeError(
+                    f"Sample[{sample_idx}] span[{span_idx}] missing key `{key}`."
+                )
+
+        owner = f"Sample[{sample_idx}] span[{span_idx}]"
+        length = _require_json_int(span["len"], owner=owner, field_name="len")
+        if length <= 0:
+            raise RuntimeError(
+                f"Sample[{sample_idx}] span[{span_idx}] invalid len={length}."
+            )
+
+        role = str(span["role"])
+        if role not in ROLE_TO_ID:
+            raise RuntimeError(
+                f"Sample[{sample_idx}] span[{span_idx}] invalid role={role}."
+            )
+        role_id = ROLE_TO_ID[role]
+
+        values = span["values"]
+        if not isinstance(values, list):
+            raise RuntimeError(
+                f"Sample[{sample_idx}] span[{span_idx}] values must be list, got {type(values)}."
+            )
+        if len(values) != length:
+            raise RuntimeError(
+                f"Sample[{sample_idx}] span[{span_idx}] len mismatch: "
+                f"len={length}, values={len(values)}."
+            )
+
+        normalized_values: list[float] = []
+        for value_idx, value in enumerate(values):
+            if isinstance(value, bool):
+                raise RuntimeError(
+                    f"Sample[{sample_idx}] span[{span_idx}] value[{value_idx}] is bool."
+                )
+            if isinstance(value, (int, float)):
+                casted = float(value)
+            else:
+                raise RuntimeError(
+                    f"Sample[{sample_idx}] span[{span_idx}] value[{value_idx}] invalid type={type(value)}."
+                )
+            if not math.isfinite(casted):
+                raise RuntimeError(
+                    f"Sample[{sample_idx}] span[{span_idx}] value[{value_idx}] is not finite: {casted}."
+                )
+            if abs(casted) > _FLOAT32_MAX:
+                raise RuntimeError(
+                    f"Sample[{sample_idx}] span[{span_idx}] value[{value_idx}] exceeds float32 range: {casted}."
+                )
+            normalized_values.append(casted)
+
+        if "segment_id" in span and span["segment_id"] is not None:
+            segment_id = _require_json_int(
+                span["segment_id"],
+                owner=owner,
+                field_name="segment_id",
+            )
+        elif require_segment_id:
+            raise RuntimeError(
+                f"Sample[{sample_idx}] span[{span_idx}] missing `segment_id`."
+            )
+        else:
+            segment_id = int(default_segment_id)
+
+        if segment_id <= 0:
+            raise RuntimeError(
+                f"Sample[{sample_idx}] span[{span_idx}] invalid segment_id={segment_id}."
+            )
+
+        loss_start = _require_json_int(
+            span["loss_start"],
+            owner=owner,
+            field_name="loss_start",
+        )
+        if loss_start < 0 or loss_start > length:
+            raise RuntimeError(
+                f"Sample[{sample_idx}] span[{span_idx}] invalid loss_start={loss_start} for len={length}."
+            )
+
+        loss_roi_mask = None
+        if "loss_roi_mask" in span and span["loss_roi_mask"] is not None:
+            raw_loss_roi_mask = span["loss_roi_mask"]
+            if not isinstance(raw_loss_roi_mask, list):
+                raise RuntimeError(
+                    f"Sample[{sample_idx}] span[{span_idx}] loss_roi_mask must be list, "
+                    f"got {type(raw_loss_roi_mask)}."
+                )
+            if len(raw_loss_roi_mask) > 0:
+                if len(raw_loss_roi_mask) != length:
+                    raise RuntimeError(
+                        f"Sample[{sample_idx}] span[{span_idx}] loss_roi_mask length mismatch: "
+                        f"len={length}, mask={len(raw_loss_roi_mask)}."
+                    )
+                normalized_roi_mask: list[float] = []
+                for mask_idx, mask_value in enumerate(raw_loss_roi_mask):
+                    if isinstance(mask_value, bool) or not isinstance(
+                        mask_value, (int, float)
+                    ):
+                        raise RuntimeError(
+                            f"Sample[{sample_idx}] span[{span_idx}] loss_roi_mask[{mask_idx}] "
+                            f"invalid type={type(mask_value)}."
+                        )
+                    casted_mask = float(mask_value)
+                    if (
+                        not math.isfinite(casted_mask)
+                        or casted_mask < 0.0
+                        or casted_mask > 1.0
+                    ):
+                        raise RuntimeError(
+                            f"Sample[{sample_idx}] span[{span_idx}] loss_roi_mask[{mask_idx}] "
+                            f"must be finite and in [0, 1], got {casted_mask}."
+                        )
+                    normalized_roi_mask.append(casted_mask)
+                loss_roi_mask = normalized_roi_mask
+
+        text_start_token_idx = None
+        if "text_start_token_idx" in span and span["text_start_token_idx"] is not None:
+            text_start_token_idx = _require_json_int(
+                span["text_start_token_idx"],
+                owner=owner,
+                field_name="text_start_token_idx",
+            )
+            if text_start_token_idx < 0:
+                raise RuntimeError(
+                    f"{owner} text_start_token_idx must be >= 0, got {text_start_token_idx}."
+                )
+        text_end_token_idx = None
+        if "text_end_token_idx" in span and span["text_end_token_idx"] is not None:
+            text_end_token_idx = _require_json_int(
+                span["text_end_token_idx"],
+                owner=owner,
+                field_name="text_end_token_idx",
+            )
+            if text_end_token_idx < 0:
+                raise RuntimeError(
+                    f"{owner} text_end_token_idx must be >= 0, got {text_end_token_idx}."
+                )
+        if (text_start_token_idx is None) != (text_end_token_idx is None):
+            raise RuntimeError(
+                f"{owner} text_start_token_idx/text_end_token_idx must be provided together."
+            )
+        if (
+            text_start_token_idx is not None
+            and text_end_token_idx < text_start_token_idx
+        ):
+            raise RuntimeError(
+                f"{owner} text token bounds are inverted: start={text_start_token_idx}, end={text_end_token_idx}."
+            )
+
+        normalized.append(
+            MoTNormalizedSpan(
+                length=length,
+                role_id=role_id,
+                segment_id=segment_id,
+                values=normalized_values,
+                loss_start=loss_start,
+                loss_roi_mask=loss_roi_mask,
+                text_start_token_idx=text_start_token_idx,
+                text_end_token_idx=text_end_token_idx,
+            )
+        )
+    return normalized
+
+
+def build_timeseries_tensors(
+    spans_by_sample: list[list[MoTNormalizedSpan]],
+) -> dict[str, torch.Tensor]:
+    batch_size = len(spans_by_sample)
+    max_spans = max((len(spans) for spans in spans_by_sample), default=0)
+    max_len = max(
+        (span.length for spans in spans_by_sample for span in spans), default=0
+    )
+
+    ts_values = torch.zeros((batch_size, max_spans, max_len), dtype=torch.float32)
+    ts_lengths = torch.zeros((batch_size, max_spans), dtype=torch.int32)
+    ts_loss_start_idxs = torch.zeros((batch_size, max_spans), dtype=torch.int32)
+    ts_loss_roi_masks = torch.zeros(
+        (batch_size, max_spans, max_len), dtype=torch.float32
+    )
+    ts_roles = torch.full((batch_size, max_spans), -1, dtype=torch.int32)
+    ts_segment_ids = torch.full((batch_size, max_spans), -1, dtype=torch.int32)
+    ts_span_mask = torch.zeros((batch_size, max_spans), dtype=torch.bool)
+
+    for sample_idx, spans in enumerate(spans_by_sample):
+        for span_idx, span in enumerate(spans):
+            ts_values[sample_idx, span_idx, : span.length] = torch.tensor(
+                span.values, dtype=torch.float32
+            )
+            ts_lengths[sample_idx, span_idx] = int(span.length)
+            ts_loss_start_idxs[sample_idx, span_idx] = int(span.loss_start)
+            if span.loss_roi_mask is not None:
+                ts_loss_roi_masks[sample_idx, span_idx, : span.length] = torch.tensor(
+                    span.loss_roi_mask, dtype=torch.float32
+                )
+            ts_roles[sample_idx, span_idx] = int(span.role_id)
+            ts_segment_ids[sample_idx, span_idx] = int(span.segment_id)
+            ts_span_mask[sample_idx, span_idx] = True
+
+    return {
+        "ts_values": ts_values,
+        "ts_lengths": ts_lengths,
+        "ts_loss_start_idxs": ts_loss_start_idxs,
+        "ts_loss_roi_masks": ts_loss_roi_masks,
+        "ts_roles": ts_roles,
+        "ts_segment_ids": ts_segment_ids,
+        "ts_span_mask": ts_span_mask,
+    }
+
+
+def build_timeseries_text_token_tensors(
+    *,
+    input_ids_by_sample: list[list[int]],
+    spans_by_sample: list[list[MoTNormalizedSpan]],
+    max_spans_per_sample: int,
+    open_id_to_len: dict[int, int | None],
+    close_id_to_len: dict[int, int | None],
+) -> dict[str, torch.Tensor]:
+    """
+    Build per-span text-token boundary tensors aligned with the TS payload slots.
+
+    These tensors are intentionally batch-static:
+    - they depend only on tokenizer-visible text and already-normalized
+      `timeseries` payload ordering
+    - they do *not* depend on model hidden states or current TS layer depth
+
+    The runtime consumes these exact token boundaries directly instead of
+    rescanning every input row for `<ts> ... </ts>` on every forward pass.
+    """
+    if len(input_ids_by_sample) != len(spans_by_sample):
+        raise RuntimeError(
+            "MoT text-span tensor builder expects input_ids and spans to share batch size, "
+            f"got input_rows={len(input_ids_by_sample)}, span_rows={len(spans_by_sample)}."
+        )
+
+    batch_size = len(spans_by_sample)
+    max_spans = max((len(spans) for spans in spans_by_sample), default=0)
+    if max_spans > max_spans_per_sample:
+        raise RuntimeError(
+            f"Too many TS spans in one sample for processor metadata: {max_spans} > {max_spans_per_sample}."
+        )
+
+    ts_text_start_token_idxs = torch.full(
+        (batch_size, max_spans), -1, dtype=torch.int32
+    )
+    ts_text_end_token_idxs = torch.full((batch_size, max_spans), -1, dtype=torch.int32)
+
+    for sample_idx, (token_ids, normalized_spans) in enumerate(
+        zip(input_ids_by_sample, spans_by_sample)
+    ):
+        if all(
+            span.text_start_token_idx is not None
+            and span.text_end_token_idx is not None
+            for span in normalized_spans
+        ):
+            for span_idx, span in enumerate(normalized_spans):
+                start_idx = int(span.text_start_token_idx)
+                end_idx = int(span.text_end_token_idx)
+                open_target = end_idx < 0 and int(span.role_id) == ROLE_TO_ID["target"]
+                if (
+                    start_idx < 0
+                    or (not open_target and end_idx < start_idx)
+                    or (not open_target and end_idx >= len(token_ids))
+                    or (open_target and start_idx >= len(token_ids))
+                ):
+                    raise RuntimeError(
+                        "Precomputed MoT text bounds are invalid inside the processor: "
+                        f"sample={sample_idx}, span={span_idx}, start={start_idx}, end={end_idx}, "
+                        f"tokens={len(token_ids)}."
+                    )
+                ts_text_start_token_idxs[sample_idx, span_idx] = start_idx
+                ts_text_end_token_idxs[sample_idx, span_idx] = end_idx
+            continue
+
+        if any(
+            span.text_start_token_idx is not None or span.text_end_token_idx is not None
+            for span in normalized_spans
+        ):
+            raise RuntimeError(
+                f"Sample[{sample_idx}] mixes precomputed and tokenizer-derived TS text bounds."
+            )
+        parsed_spans = find_tokenized_ts_spans(
+            [int(token_id) for token_id in token_ids],
+            open_id_to_len=open_id_to_len,
+            close_id_to_len=close_id_to_len,
+        )
+        if len(parsed_spans) != len(normalized_spans):
+            raise RuntimeError(
+                "Parsed text spans and normalized TS payload spans disagree inside the processor: "
+                f"sample={sample_idx}, parsed={len(parsed_spans)}, payload={len(normalized_spans)}."
+            )
+
+        for span_idx, (span, (start_idx, end_idx, tagged_len)) in enumerate(
+            zip(normalized_spans, parsed_spans)
+        ):
+            if tagged_len is not None and int(span.length) != int(tagged_len):
+                raise RuntimeError(
+                    "TS payload length mismatch with text delimiter tags in processor input: "
+                    f"sample={sample_idx}, span={span_idx}, payload_len={span.length}, tag_len={tagged_len}."
+                )
+            ts_text_start_token_idxs[sample_idx, span_idx] = int(start_idx)
+            ts_text_end_token_idxs[sample_idx, span_idx] = int(end_idx)
+
+    return {
+        "ts_text_start_token_idxs": ts_text_start_token_idxs,
+        "ts_text_end_token_idxs": ts_text_end_token_idxs,
+    }
+
+
+def build_mot_batch_timeseries_tensors(
+    *,
+    tokenizer: Any,
+    input_ids_by_sample: list[list[int]],
+    spans_by_sample: list[list[MoTNormalizedSpan]],
+    max_spans_per_sample: int,
+) -> dict[str, torch.Tensor]:
+    """
+    Build the full processor-side TS tensor bundle for one inference request.
+
+    The model consumes two aligned views:
+    - dense span/value tensors for the bridge path
+    - text-token boundary tensors for exact `<ts> ... </ts>` alignment
+    """
+
+    needs_tokenized_delimiter_parse = any(
+        any(
+            span.text_start_token_idx is None or span.text_end_token_idx is None
+            for span in spans
+        )
+        for spans in spans_by_sample
+    )
+    if needs_tokenized_delimiter_parse:
+        open_id_to_len, close_id_to_len = resolve_ts_delimiter_id_maps(tokenizer)
+    else:
+        open_id_to_len, close_id_to_len = {}, {}
+    batch = build_timeseries_tensors(spans_by_sample)
+    text_bounds = build_timeseries_text_token_tensors(
+        input_ids_by_sample=input_ids_by_sample,
+        spans_by_sample=spans_by_sample,
+        max_spans_per_sample=max_spans_per_sample,
+        open_id_to_len=open_id_to_len,
+        close_id_to_len=close_id_to_len,
+    )
+    batch.update(text_bounds)
+    return batch
+
+
+__all__ = [
+    "build_mot_batch_timeseries_tensors",
+    "normalize_timeseries_spans",
+]

timebraid/model/__init__.py

@@ -0,0 +1,35 @@
+"""TimeBraid model core."""
+
+from .loading import (
+    LoadedTimeBraid,
+    ValidatedTimeBraidCheckpointSource,
+    load_timebraid_checkpoint,
+    load_timebraid_tokenizer,
+    validate_timebraid_canonical_artifact,
+    validate_timebraid_checkpoint_source,
+    validate_timebraid_model,
+)
+from .mot import TimeBraidGenerateOutput
+from .timebraid import (
+    TimeBraid,
+    TimeBraidConfig,
+    TimeBraidLossBreakdown,
+    TimeBraidOutput,
+    register_timebraid_auto_classes,
+)
+
+__all__ = [
+    "TimeBraid",
+    "TimeBraidConfig",
+    "TimeBraidGenerateOutput",
+    "TimeBraidLossBreakdown",
+    "TimeBraidOutput",
+    "LoadedTimeBraid",
+    "ValidatedTimeBraidCheckpointSource",
+    "register_timebraid_auto_classes",
+    "load_timebraid_checkpoint",
+    "load_timebraid_tokenizer",
+    "validate_timebraid_canonical_artifact",
+    "validate_timebraid_model",
+    "validate_timebraid_checkpoint_source",
+]

timebraid/model/loading.py

@@ -0,0 +1,1038 @@
+"""Strict loading for local or Hub TimeBraid checkpoints."""
+
+from __future__ import annotations
+
+import json
+import math
+from dataclasses import dataclass
+from pathlib import Path
+from typing import Any, Literal
+
+import torch
+from huggingface_hub import snapshot_download
+from safetensors import safe_open
+from transformers import (
+    AutoConfig,
+    AutoProcessor,
+    AutoTokenizer,
+    PreTrainedModel,
+    PreTrainedTokenizerBase,
+)
+
+from .mot.config_contract import collect_mot_hf_config_contract
+from .timebraid import (
+    TimeBraid,
+    TimeBraidConfig,
+    _validate_finite_model_tensors,
+    register_timebraid_auto_classes,
+)
+
+_MODEL_INDEX_NAME = "model.safetensors.index.json"
+_SINGLE_MODEL_NAME = "model.safetensors"
+# Keep structural preflight bounded independently of untrusted checkpoint
+# metadata; supported runtime payloads use a much smaller configured span cap.
+_MAX_SPANS_PER_SAMPLE = 2048
+
+
+@dataclass(frozen=True, slots=True)
+class LoadedTimeBraid:
+    """A frozen TimeBraid model, its tokenizer, and explicit load provenance."""
+
+    model: PreTrainedModel
+    tokenizer: PreTrainedTokenizerBase
+    source: str
+    weight_dtype: str
+    compute_dtype: str
+    device_map: str | dict[str, str] | None
+    physical_tensor_dtypes: tuple[tuple[str, str], ...] = ()
+
+
+@dataclass(frozen=True, slots=True)
+class ValidatedTimeBraidCheckpointSource:
+    """CPU-only structural preflight for one complete TimeBraid checkpoint."""
+
+    source: str
+    config: TimeBraidConfig
+    tokenizer: PreTrainedTokenizerBase
+    tensor_count: int
+    physical_tensor_dtypes: tuple[tuple[str, str], ...]
+    mot_ts_open_token_id: int
+    mot_ts_close_token_id: int
+
+
+def _loading_info_key(item: Any) -> str:
+    if isinstance(item, (tuple, list)) and item:
+        return str(item[0])
+    return str(item)
+
+
+def _normalize_source(source: str | Path, *, field_name: str) -> str:
+    if not isinstance(source, (str, Path)):
+        raise TypeError(
+            f"{field_name} must be a local directory path or Hub ID, got {type(source).__name__}."
+        )
+    normalized = str(source).strip()
+    if not normalized:
+        raise ValueError(f"{field_name} must be non-empty.")
+    return normalized
+
+
+def _resolve_hf_source(
+    source: str,
+    *,
+    revision: str | None,
+    token: str | None,
+    cache_dir: str | None,
+    local_files_only: bool,
+) -> str:
+    """Resolve an existing directory or one standard Hugging Face Hub repo."""
+    path = Path(source).expanduser()
+    if path.is_dir():
+        return str(path.resolve(strict=True))
+    if path.exists():
+        raise RuntimeError(f"source must be a directory or Hub ID, got {source!r}.")
+    return snapshot_download(
+        repo_id=source,
+        revision=revision,
+        token=token,
+        cache_dir=cache_dir,
+        local_files_only=local_files_only,
+    )
+
+
+def _strict_json_object(path: Path, *, owner: str) -> dict[str, Any]:
+    """Read one strict JSON object without duplicate or non-finite values."""
+
+    def reject_duplicates(pairs: list[tuple[str, Any]]) -> dict[str, Any]:
+        result: dict[str, Any] = {}
+        for key, value in pairs:
+            if key in result:
+                raise RuntimeError(f"{owner} contains duplicate key {key!r}: {path}")
+            result[key] = value
+        return result
+
+    def reject_constant(value: str) -> None:
+        raise RuntimeError(f"{owner} contains non-finite JSON number {value!r}: {path}")
+
+    def parse_finite_float(value: str) -> float:
+        parsed = float(value)
+        if not math.isfinite(parsed):
+            raise RuntimeError(
+                f"{owner} contains overflowing JSON number {value!r}: {path}"
+            )
+        return parsed
+
+    value = json.loads(
+        path.read_text(encoding="utf-8"),
+        object_pairs_hook=reject_duplicates,
+        parse_constant=reject_constant,
+        parse_float=parse_finite_float,
+    )
+    if not isinstance(value, dict):
+        raise RuntimeError(f"{owner} must contain one JSON object: {path}")
+    return value
+
+
+def _safe_top_level_name(value: object, *, owner: str) -> str:
+    if (
+        not isinstance(value, str)
+        or not value
+        or Path(value).name != value
+        or "/" in value
+        or "\\" in value
+        or value in {".", ".."}
+    ):
+        raise RuntimeError(f"{owner} must be a top-level filename, got {value!r}.")
+    return value
+
+
+def _resolved_source_hf_kwargs(*, trust_remote_code: bool) -> dict[str, Any]:
+    """Build reader kwargs after a Hub repo has resolved to one local snapshot."""
+    if trust_remote_code is not False:
+        raise ValueError(
+            "TimeBraid loading requires trust_remote_code=False; install and import "
+            "the TimeBraid package to register its Hugging Face classes."
+        )
+    return {
+        "local_files_only": True,
+        "trust_remote_code": False,
+    }
+
+
+def _reject_auto_map(value: Any, *, owner: str) -> None:
+    """Reject repository-provided Python entrypoints anywhere in HF metadata.
+
+    This is deliberately stricter than `trust_remote_code=False`: it refuses an
+    artifact that merely *carries* remote-code entrypoints, rather than only
+    declining to run them. That makes `load_timebraid_checkpoint` a loader for
+    canonical artifacts, which have no `auto_map`.
+
+    The published Hub artifact is not such an artifact — it ships `auto_map` on
+    purpose, so users who have not installed this package can load it with
+    `trust_remote_code=True`. Those users, and installed-package users alike,
+    go through `AutoModelForCausalLM.from_pretrained`, which works either way.
+    """
+    pending = [value]
+    while pending:
+        current = pending.pop()
+        if isinstance(current, dict):
+            if "auto_map" in current:
+                raise RuntimeError(
+                    f"{owner} must not contain auto_map; install and import the "
+                    "TimeBraid package instead."
+                )
+            pending.extend(current.values())
+        elif isinstance(current, list):
+            pending.extend(current)
+
+
+def _local_safetensors_tensor_dtypes(root: str) -> dict[str, str]:
+    directory = Path(root)
+    index_path = directory / _MODEL_INDEX_NAME
+    single_path = directory / _SINGLE_MODEL_NAME
+    unsafe_weight_files = sorted(
+        path.name
+        for path in directory.iterdir()
+        if path.name == "pytorch_model.bin"
+        or path.name == "pytorch_model.bin.index.json"
+        or (path.name.startswith("pytorch_model-") and path.name.endswith(".bin"))
+    )
+    if unsafe_weight_files:
+        raise RuntimeError(
+            f"Strict loading rejects pickle-backed weights: {unsafe_weight_files}."
+        )
+    selectors = (index_path.is_file(), single_path.is_file())
+    if sum(selectors) != 1:
+        raise RuntimeError(
+            "Resolved source must contain exactly one safetensors selector: "
+            "model.safetensors or model.safetensors.index.json."
+        )
+
+    declared: dict[str, str] | None = None
+    if selectors[0]:
+        index = _strict_json_object(index_path, owner="Model safetensors index")
+        weight_map = index.get("weight_map")
+        if not isinstance(weight_map, dict) or not weight_map:
+            raise RuntimeError(
+                "Model safetensors index must contain a non-empty weight_map."
+            )
+        declared = {}
+        for tensor_name, shard_value in weight_map.items():
+            if not isinstance(tensor_name, str) or not tensor_name:
+                raise RuntimeError(
+                    "Model safetensors index tensor names must be non-empty strings."
+                )
+            shard_name = _safe_top_level_name(
+                shard_value, owner=f"Shard for {tensor_name!r}"
+            )
+            if not shard_name.endswith(".safetensors"):
+                raise RuntimeError(
+                    f"Model shard must use safetensors, got {shard_name!r}."
+                )
+            declared[tensor_name] = shard_name
+        shard_names = sorted(set(declared.values()))
+    else:
+        shard_names = [_SINGLE_MODEL_NAME]
+
+    observed: dict[str, tuple[str, str]] = {}
+    for shard_name in shard_names:
+        shard_path = directory / shard_name
+        if not shard_path.is_file():
+            raise RuntimeError(
+                f"Model shard must resolve to a regular file: {shard_path}"
+            )
+        try:
+            with safe_open(shard_path, framework="pt", device="cpu") as reader:
+                shard_keys = list(reader.keys())
+                for tensor_name in shard_keys:
+                    tensor_slice = reader.get_slice(tensor_name)
+                    if tensor_name in observed:
+                        raise RuntimeError(
+                            f"Safetensors tensor appears in multiple shards: {tensor_name!r}."
+                        )
+                    observed[tensor_name] = (
+                        shard_name,
+                        str(tensor_slice.get_dtype()),
+                    )
+        except RuntimeError:
+            raise
+        except Exception as exc:
+            raise RuntimeError(
+                f"Could not validate safetensors shard: {shard_path}"
+            ) from exc
+    if not observed:
+        raise RuntimeError("Model safetensors payload contains no tensors.")
+    observed_placement = {
+        tensor_name: shard_name
+        for tensor_name, (shard_name, _dtype) in observed.items()
+    }
+    if declared is not None and observed_placement != declared:
+        missing = sorted(set(declared) - set(observed))
+        extra = sorted(set(observed) - set(declared))
+        misplaced = {
+            name: {"declared": declared[name], "observed": observed_placement[name]}
+            for name in sorted(set(declared) & set(observed))
+            if declared[name] != observed_placement[name]
+        }
+        raise RuntimeError(
+            "Safetensors index does not match shard tensor placement: "
+            f"missing={missing}, extra={extra}, misplaced={misplaced}."
+        )
+    observed_shards = {
+        path.name for path in directory.iterdir() if path.name.endswith(".safetensors")
+    }
+    if observed_shards != set(shard_names):
+        raise RuntimeError(
+            "Model source contains undeclared safetensors shards: "
+            f"declared={shard_names!r}, observed={sorted(observed_shards)!r}."
+        )
+    return {
+        tensor_name: dtype for tensor_name, (_shard_name, dtype) in observed.items()
+    }
+
+
+def _reject_noncanonical_tied_aliases(
+    config_payload: dict[str, Any], tensor_names: set[str]
+) -> None:
+    model_type = config_payload.get("model_type")
+    if model_type == TimeBraidConfig.model_type:
+        llm_config = config_payload.get("llm_config")
+        if not isinstance(llm_config, dict):
+            raise RuntimeError(
+                "TimeBraid config must contain a nested llm_config object."
+            )
+        tie_word_embeddings = llm_config.get("tie_word_embeddings", False)
+        lm_head_alias = "llm.lm_head.weight"
+    else:
+        tie_word_embeddings = config_payload.get("tie_word_embeddings", False)
+        lm_head_alias = "lm_head.weight"
+    if type(tie_word_embeddings) is not bool:
+        raise RuntimeError(
+            f"Model config tie_word_embeddings must be bool, got {tie_word_embeddings!r}."
+        )
+    forbidden = set()
+    if tie_word_embeddings and lm_head_alias in tensor_names:
+        forbidden.add(lm_head_alias)
+    if forbidden:
+        raise RuntimeError(
+            "Model safetensors contains noncanonical tied-weight aliases that HF may silently ignore: "
+            f"{sorted(forbidden)!r}."
+        )
+
+
+def _validate_local_model_payload(root: str) -> dict[str, str]:
+    directory = Path(root)
+    for forbidden_name in ("adapter_config.json", "additional_chat_templates"):
+        forbidden_path = directory / forbidden_name
+        try:
+            forbidden_path.lstat()
+        except FileNotFoundError:
+            continue
+        raise RuntimeError(
+            "Strict full-checkpoint loading rejects loader redirection or nested "
+            f"tokenizer templates: {forbidden_path}"
+        )
+
+    config_path = directory / "config.json"
+    if not config_path.is_file():
+        raise RuntimeError(
+            f"Model config must resolve to a regular file: {config_path}"
+        )
+    config_payload = _strict_json_object(config_path, owner="Model config")
+    _reject_auto_map(config_payload, owner="Model config")
+    tensor_dtypes = _local_safetensors_tensor_dtypes(root)
+    _reject_noncanonical_tied_aliases(config_payload, set(tensor_dtypes))
+    return tensor_dtypes
+
+
+def _validate_complete_checkpoint_load(
+    model: TimeBraid, loading_info: dict[str, Any]
+) -> None:
+    """Reject every tensor mismatch except live tied-weight aliases omitted by HF saves."""
+    tied_targets = set(getattr(model, "_tied_weights_keys", {}))
+    missing = sorted(
+        key
+        for item in (loading_info.get("missing_keys") or [])
+        if (key := _loading_info_key(item)) not in tied_targets
+    )
+    unexpected = sorted(
+        _loading_info_key(item) for item in (loading_info.get("unexpected_keys") or [])
+    )
+    mismatched = sorted(
+        str(item) for item in (loading_info.get("mismatched_keys") or [])
+    )
+    error_messages = [str(item) for item in (loading_info.get("error_msgs") or [])]
+    if missing or unexpected or mismatched or error_messages:
+        raise RuntimeError(
+            "TimeBraid checkpoint did not load exactly: "
+            f"missing={missing}, unexpected={unexpected}, mismatched={mismatched}, "
+            f"errors={error_messages}"
+        )
+
+
+def _validate_tokenizer_protocol(
+    tokenizer: PreTrainedTokenizerBase,
+    *,
+    token_id_upper_bound: int,
+) -> dict[str, int]:
+    if len(tokenizer) > token_id_upper_bound:
+        raise RuntimeError(
+            "TimeBraid tokenizer vocabulary exceeds the checkpoint embedding table: "
+            f"tokenizer={len(tokenizer)}, embeddings={token_id_upper_bound}."
+        )
+    delimiter_ids: dict[str, int] = {}
+    for delimiter in ("<ts>", "</ts>"):
+        token_ids = tokenizer(delimiter, add_special_tokens=False)["input_ids"]
+        if not isinstance(token_ids, list) or len(token_ids) != 1:
+            raise RuntimeError(
+                f"TimeBraid tokenizer must encode {delimiter!r} as exactly one token, got {token_ids!r}."
+            )
+        token_id = int(token_ids[0])
+        converted_id = tokenizer.convert_tokens_to_ids(delimiter)
+        if converted_id is None or int(converted_id) != token_id:
+            raise RuntimeError(
+                f"TimeBraid tokenizer has inconsistent ID resolution for {delimiter!r}: "
+                f"encode={token_id}, convert={converted_id!r}."
+            )
+        if tokenizer.unk_token_id is not None and token_id == int(
+            tokenizer.unk_token_id
+        ):
+            raise RuntimeError(
+                f"TimeBraid tokenizer resolves {delimiter!r} to unk_token_id={token_id}."
+            )
+        if token_id < 0 or token_id >= token_id_upper_bound:
+            raise RuntimeError(
+                f"TimeBraid tokenizer encodes {delimiter!r} outside the embedding table: {token_id}."
+            )
+        delimiter_ids[delimiter] = token_id
+    if delimiter_ids["<ts>"] == delimiter_ids["</ts>"]:
+        raise RuntimeError(
+            "TimeBraid tokenizer maps <ts> and </ts> to the same token ID: "
+            f"{delimiter_ids['<ts>']}."
+        )
+    return delimiter_ids
+
+
+def _load_exact_tokenizer(
+    source: str,
+    *,
+    use_fast_tokenizer: bool,
+    hf_kwargs: dict[str, Any],
+) -> PreTrainedTokenizerBase:
+    tokenizer_config_path = Path(source) / "tokenizer_config.json"
+    tokenizer_config = _strict_json_object(
+        tokenizer_config_path,
+        owner="Tokenizer config",
+    )
+    _reject_auto_map(tokenizer_config, owner="Tokenizer config")
+    extra_special_tokens_kwargs: dict[str, Any] = {}
+    if "extra_special_tokens" in tokenizer_config:
+        extra_special_tokens = tokenizer_config["extra_special_tokens"]
+        if isinstance(extra_special_tokens, dict):
+            pass
+        elif isinstance(extra_special_tokens, list) and not extra_special_tokens:
+            # Older Qwen checkpoints serialized the empty model-specific token
+            # mapping as []. Transformers 4.57 calls .keys() on this value; an
+            # explicit {} preserves the same empty set without changing bytes.
+            extra_special_tokens_kwargs["extra_special_tokens"] = {}
+        elif isinstance(extra_special_tokens, list):
+            raise RuntimeError(
+                "tokenizer_config.json extra_special_tokens must be an object or []; "
+                f"got list[{len(extra_special_tokens)}]."
+            )
+        else:
+            raise RuntimeError(
+                "tokenizer_config.json extra_special_tokens must be an object or []; "
+                f"got {type(extra_special_tokens).__name__}."
+            )
+    tokenizer = AutoTokenizer.from_pretrained(
+        source,
+        use_fast=use_fast_tokenizer,
+        padding_side="right",
+        # Preserve the serialized Qwen regex exactly; an automatic Mistral
+        # rewrite changes tokenizer identity and invalidates tokenized caches.
+        fix_mistral_regex=False,
+        **extra_special_tokens_kwargs,
+        **hf_kwargs,
+    )
+    if tokenizer.padding_side != "right":
+        raise RuntimeError(
+            f"TimeBraid tokenizer must use right padding, got {tokenizer.padding_side!r}."
+        )
+    return tokenizer
+
+
+def _validate_timebraid_config_surface(
+    config: TimeBraidConfig,
+) -> tuple[int, dict[str, Any]]:
+    if config.model_type != TimeBraidConfig.model_type:
+        raise RuntimeError(
+            "Strict TimeBraid checkpoint loading requires model_type='timebraid', "
+            f"got {config.model_type!r}."
+        )
+    mot_contract = collect_mot_hf_config_contract(config)
+    if str(getattr(config.llm_config, "model_type", "")) != "qwen3":
+        raise RuntimeError(
+            "TimeBraid currently requires a nested Qwen3 backbone config, got "
+            f"{getattr(config.llm_config, 'model_type', None)!r}."
+        )
+    config_vocab_size = getattr(config.llm_config, "vocab_size", 0)
+    if type(config_vocab_size) is not int or config_vocab_size <= 0:
+        raise RuntimeError(
+            f"TimeBraid config has invalid vocab_size={config_vocab_size}."
+        )
+    return config_vocab_size, mot_contract
+
+
+def _validate_config_tokenizer_delimiters(
+    *, mot_contract: dict[str, Any], tokenizer_ids: dict[str, int]
+) -> None:
+    configured = (
+        mot_contract["mot_ts_open_token_id"],
+        mot_contract["mot_ts_close_token_id"],
+    )
+    observed = (tokenizer_ids["<ts>"], tokenizer_ids["</ts>"])
+    if configured != observed:
+        raise RuntimeError(
+            "TimeBraid config TS delimiter IDs disagree with the tokenizer: "
+            f"config={configured}, tokenizer={observed}."
+        )
+
+
+def _load_timebraid_config(
+    source: str,
+    *,
+    hf_kwargs: dict[str, Any],
+) -> TimeBraidConfig:
+    """Load one current TimeBraid config without mutating serialized fields."""
+    AutoConfig.register(TimeBraidConfig.model_type, TimeBraidConfig, exist_ok=True)
+    config = AutoConfig.from_pretrained(source, **hf_kwargs)
+    if not isinstance(config, TimeBraidConfig):
+        raise RuntimeError(
+            "TimeBraid checkpoint must resolve to TimeBraidConfig, got "
+            f"{type(config).__name__}."
+        )
+    return config
+
+
+def validate_timebraid_checkpoint_source(
+    source: str | Path,
+    *,
+    revision: str | None = None,
+    token: str | None = None,
+    cache_dir: str | None = None,
+    local_files_only: bool = False,
+    trust_remote_code: bool = False,
+    use_fast_tokenizer: bool = True,
+) -> ValidatedTimeBraidCheckpointSource:
+    """Validate config, tokenizer, and safetensors without loading model tensors."""
+    hf_kwargs = _resolved_source_hf_kwargs(trust_remote_code=trust_remote_code)
+    source_str = _resolve_hf_source(
+        _normalize_source(source, field_name="source"),
+        revision=revision,
+        token=token,
+        cache_dir=cache_dir,
+        local_files_only=local_files_only,
+    )
+    tensor_dtypes = _validate_local_model_payload(source_str)
+    tokenizer = _load_exact_tokenizer(
+        source_str,
+        use_fast_tokenizer=use_fast_tokenizer,
+        hf_kwargs=hf_kwargs,
+    )
+    config = _load_timebraid_config(
+        source_str,
+        hf_kwargs=hf_kwargs,
+    )
+    config_vocab_size, mot_contract = _validate_timebraid_config_surface(config)
+    token_ids = _validate_tokenizer_protocol(
+        tokenizer,
+        token_id_upper_bound=config_vocab_size,
+    )
+    _validate_config_tokenizer_delimiters(
+        mot_contract=mot_contract,
+        tokenizer_ids=token_ids,
+    )
+
+    return ValidatedTimeBraidCheckpointSource(
+        source=source_str,
+        config=config,
+        tokenizer=tokenizer,
+        tensor_count=len(tensor_dtypes),
+        physical_tensor_dtypes=tuple(sorted(tensor_dtypes.items())),
+        mot_ts_open_token_id=token_ids["<ts>"],
+        mot_ts_close_token_id=token_ids["</ts>"],
+    )
+
+
+def load_timebraid_tokenizer(
+    source: str | Path,
+    *,
+    revision: str | None = None,
+    token: str | None = None,
+    cache_dir: str | None = None,
+    local_files_only: bool = False,
+    trust_remote_code: bool = False,
+    use_fast_tokenizer: bool = True,
+) -> PreTrainedTokenizerBase:
+    """Load the exact tokenizer from one local directory or Hub repo."""
+    hf_kwargs = _resolved_source_hf_kwargs(trust_remote_code=trust_remote_code)
+    source_str = _resolve_hf_source(
+        _normalize_source(source, field_name="source"),
+        revision=revision,
+        token=token,
+        cache_dir=cache_dir,
+        local_files_only=local_files_only,
+    )
+    tokenizer = _load_exact_tokenizer(
+        source_str,
+        use_fast_tokenizer=use_fast_tokenizer,
+        hf_kwargs=hf_kwargs,
+    )
+    # This source-only preflight fails before model allocation. The model loader
+    # repeats it against the actual embedding rows.
+    _validate_tokenizer_protocol(
+        tokenizer,
+        token_id_upper_bound=len(tokenizer),
+    )
+    return tokenizer
+
+
+def _validate_tied_weights(model: TimeBraid) -> None:
+    for target_name, source_name in getattr(model, "_tied_weights_keys", {}).items():
+        try:
+            target = model.get_parameter(target_name)
+            source = model.get_parameter(source_name)
+        except AttributeError as exc:
+            raise RuntimeError(
+                f"TimeBraid tied-weight path is missing: {target_name!r} -> {source_name!r}."
+            ) from exc
+        if target.data_ptr() != source.data_ptr():
+            raise RuntimeError(
+                f"TimeBraid tied weights do not share storage: {target_name!r} -> {source_name!r}."
+            )
+
+
+def _validate_no_meta_tensors(model: PreTrainedModel, *, owner: str) -> None:
+    meta_tensors = sorted(
+        name
+        for name, tensor in tuple(model.named_parameters())
+        + tuple(model.named_buffers())
+        if tensor.is_meta
+    )
+    if meta_tensors:
+        raise RuntimeError(
+            f"{owner} left meta tensors after loading: {meta_tensors[:50]}."
+        )
+
+
+def _validate_model(
+    model: TimeBraid,
+    tokenizer: PreTrainedTokenizerBase,
+) -> None:
+    input_embeddings = model.get_input_embeddings()
+    if input_embeddings is None or not hasattr(input_embeddings, "weight"):
+        raise RuntimeError("Loaded TimeBraid model has no input embedding weight.")
+    token_ids = _validate_tokenizer_protocol(
+        tokenizer,
+        token_id_upper_bound=int(input_embeddings.weight.shape[0]),
+    )
+    _validate_tied_weights(model)
+    _validate_no_meta_tensors(model, owner="TimeBraid checkpoint")
+    configured = (
+        getattr(model, "ts_open_token_id", None),
+        getattr(model, "ts_close_token_id", None),
+    )
+    observed = (token_ids["<ts>"], token_ids["</ts>"])
+    if configured != observed:
+        raise RuntimeError(
+            "TimeBraid TS delimiter IDs disagree with the tokenizer: "
+            f"model={configured}, tokenizer={observed}."
+        )
+
+
+def validate_timebraid_model(
+    model: TimeBraid,
+    tokenizer: PreTrainedTokenizerBase,
+) -> None:
+    """Validate the loaded TimeBraid tensor, tokenizer, and component surface."""
+    _validate_model(model, tokenizer)
+
+
+def _resolve_dtype(
+    value: torch.dtype | str | None, *, field_name: str
+) -> torch.dtype | Literal["auto"]:
+    if value is None or value == "auto":
+        return "auto"
+    if isinstance(value, torch.dtype):
+        return value
+    aliases = {
+        "bf16": torch.bfloat16,
+        "bfloat16": torch.bfloat16,
+        "fp16": torch.float16,
+        "float16": torch.float16,
+        "fp32": torch.float32,
+        "float32": torch.float32,
+    }
+    resolved = aliases.get(str(value).strip().lower())
+    if resolved is None:
+        raise ValueError(
+            f"{field_name} must be auto, bf16, fp16, or fp32, got {value!r}."
+        )
+    return resolved
+
+
+def _resolve_compute_dtype(
+    value: torch.dtype | str | None,
+    *,
+    model_dtype: torch.dtype | str | None,
+    accelerator_is_cuda: bool,
+) -> torch.dtype:
+    resolved = _resolve_dtype(value, field_name="compute_dtype")
+    if resolved != "auto":
+        return resolved
+    if not accelerator_is_cuda:
+        return torch.float32
+    normalized_model_dtype = (
+        str(model_dtype).removeprefix("torch.") if model_dtype is not None else None
+    )
+    if torch.cuda.is_bf16_supported() and normalized_model_dtype in {None, "bfloat16"}:
+        return torch.bfloat16
+    return torch.float16
+
+
+def _dtype_name(value: torch.dtype | Literal["auto"]) -> str:
+    return value if value == "auto" else str(value).removeprefix("torch.")
+
+
+def _resolve_device_map(
+    device: str | torch.device | None,
+    *,
+    low_cpu_mem_usage: bool,
+) -> tuple[str | dict[str, str] | None, bool]:
+    device_is_auto = isinstance(device, str) and device.strip().lower() == "auto"
+    if device_is_auto:
+        raise ValueError(
+            "device='auto' is unsupported because TimeBraid requires one explicit model device."
+        )
+    try:
+        requested_device = None if device is None else torch.device(device)
+    except (RuntimeError, TypeError) as exc:
+        raise ValueError(
+            f"device must be cpu, cuda, or cuda:N, got {device!r}."
+        ) from exc
+    accelerator_is_cuda = bool(
+        requested_device is not None and requested_device.type == "cuda"
+    )
+    if device is None:
+        return None, accelerator_is_cuda
+    if not low_cpu_mem_usage:
+        raise ValueError("device loading requires low_cpu_mem_usage=true.")
+    if requested_device is None:
+        raise RuntimeError("Explicit device resolution produced no device.")
+    if requested_device.type not in {"cpu", "cuda"}:
+        raise ValueError(f"device must be cpu or cuda, got {requested_device.type!r}.")
+    if requested_device.type == "cuda" and not torch.cuda.is_available():
+        raise RuntimeError(
+            f"CUDA device requested but CUDA is unavailable: {requested_device}."
+        )
+    return {"": str(requested_device)}, accelerator_is_cuda
+
+
+def _model_config_dtype(config: Any) -> torch.dtype | str | None:
+    return getattr(config, "dtype", None)
+
+
+def load_timebraid_checkpoint(
+    source: str | Path,
+    *,
+    compute_dtype: torch.dtype | str | None = "auto",
+    weight_dtype: torch.dtype | str | None = "auto",
+    device: str | torch.device | None = None,
+    attention_implementation: Literal["eager", "sdpa", "flash_attention_2"]
+    | None = None,
+    revision: str | None = None,
+    token: str | None = None,
+    cache_dir: str | None = None,
+    local_files_only: bool = False,
+    trust_remote_code: bool = False,
+    use_fast_tokenizer: bool = True,
+    low_cpu_mem_usage: bool = True,
+    max_spans_per_sample: int = 1024,
+) -> LoadedTimeBraid:
+    """Load one local or Hub safetensors-only TimeBraid checkpoint."""
+    _resolved_source_hf_kwargs(trust_remote_code=trust_remote_code)
+    resolved_device_map, accelerator_is_cuda = _resolve_device_map(
+        device,
+        low_cpu_mem_usage=low_cpu_mem_usage,
+    )
+    validated = validate_timebraid_checkpoint_source(
+        source,
+        revision=revision,
+        token=token,
+        cache_dir=cache_dir,
+        local_files_only=local_files_only,
+        trust_remote_code=trust_remote_code,
+        use_fast_tokenizer=use_fast_tokenizer,
+    )
+    source_str = validated.source
+    tokenizer = validated.tokenizer
+    config = validated.config
+    hf_kwargs = _resolved_source_hf_kwargs(trust_remote_code=trust_remote_code)
+    if attention_implementation is not None:
+        config.llm_config._attn_implementation = attention_implementation
+
+    if type(max_spans_per_sample) is not int:
+        raise ValueError(
+            f"max_spans_per_sample must be an exact integer, got {max_spans_per_sample!r}."
+        )
+    if not 1 <= max_spans_per_sample <= _MAX_SPANS_PER_SAMPLE:
+        raise ValueError(
+            "max_spans_per_sample must stay in the maintained runtime range "
+            f"[1, {_MAX_SPANS_PER_SAMPLE}], got {max_spans_per_sample}."
+        )
+    resolved_weight_dtype = _resolve_dtype(weight_dtype, field_name="weight_dtype")
+    resolved_compute_dtype = _resolve_compute_dtype(
+        compute_dtype,
+        model_dtype=_model_config_dtype(config.llm_config),
+        accelerator_is_cuda=accelerator_is_cuda,
+    )
+    runtime_options: dict[str, Any] = {
+        "mot_max_spans_per_sample": max_spans_per_sample,
+    }
+    # Transformers also forwards unknown model kwargs to GenerationConfig.
+    # Keep this custom value JSON-safe while the runtime accepts the same
+    # canonical dtype spelling.
+    runtime_options["mot_compute_dtype"] = _dtype_name(resolved_compute_dtype)
+
+    model_kwargs: dict[str, Any] = {
+        **hf_kwargs,
+        "config": config,
+        "low_cpu_mem_usage": low_cpu_mem_usage,
+        "output_loading_info": True,
+        "runtime_options": runtime_options,
+        "dtype": resolved_weight_dtype,
+        "use_safetensors": True,
+    }
+    if resolved_device_map is not None:
+        model_kwargs["device_map"] = resolved_device_map
+
+    model, loading_info = TimeBraid.from_pretrained(source_str, **model_kwargs)
+    _validate_complete_checkpoint_load(model, loading_info)
+    _validate_model(model, tokenizer)
+    model.requires_grad_(False)
+    model.eval()
+    return LoadedTimeBraid(
+        model=model,
+        tokenizer=tokenizer,
+        source=source_str,
+        weight_dtype=_dtype_name(resolved_weight_dtype),
+        compute_dtype=_dtype_name(resolved_compute_dtype),
+        device_map=resolved_device_map,
+        physical_tensor_dtypes=validated.physical_tensor_dtypes,
+    )
+
+
+def _safetensors_dtype_code(dtype: torch.dtype) -> str:
+    codes = {
+        torch.bool: "BOOL",
+        torch.uint8: "U8",
+        torch.int8: "I8",
+        torch.int16: "I16",
+        torch.int32: "I32",
+        torch.int64: "I64",
+        torch.float16: "F16",
+        torch.bfloat16: "BF16",
+        torch.float32: "F32",
+        torch.float64: "F64",
+    }
+    code = codes.get(dtype)
+    if code is None:
+        raise RuntimeError(f"TimeBraid release tensor has unsupported dtype {dtype}.")
+    return code
+
+
+def _validate_timebraid_processor_artifact(
+    source: str,
+    *,
+    use_fast_tokenizer: bool,
+    trust_remote_code: bool,
+) -> None:
+    """Require one self-contained HF processor package at a resolved source."""
+    hf_kwargs = _resolved_source_hf_kwargs(trust_remote_code=trust_remote_code)
+    root = Path(source)
+    processor_config_path = root / "processor_config.json"
+    if not processor_config_path.is_file():
+        raise RuntimeError(
+            "TimeBraid release artifact must contain processor_config.json: "
+            f"{processor_config_path}"
+        )
+    processor_config = _strict_json_object(
+        processor_config_path,
+        owner="Processor config",
+    )
+    tokenizer_config = _strict_json_object(
+        root / "tokenizer_config.json",
+        owner="Tokenizer config",
+    )
+    _reject_auto_map(processor_config, owner="Processor config")
+    _reject_auto_map(tokenizer_config, owner="Tokenizer config")
+    if processor_config.get("processor_class") != "TimeBraidProcessor":
+        raise RuntimeError(
+            "processor_config.json processor_class must be 'TimeBraidProcessor', got "
+            f"{processor_config.get('processor_class')!r}."
+        )
+    max_spans_per_sample = processor_config.get("max_spans_per_sample")
+    if (
+        type(max_spans_per_sample) is not int
+        or not 1 <= max_spans_per_sample <= _MAX_SPANS_PER_SAMPLE
+    ):
+        raise RuntimeError(
+            "processor_config.json max_spans_per_sample must be an integer in "
+            f"[1, {_MAX_SPANS_PER_SAMPLE}], got {max_spans_per_sample!r}."
+        )
+    normalization_epsilon = processor_config.get("normalization_epsilon")
+    if (
+        isinstance(normalization_epsilon, bool)
+        or not isinstance(normalization_epsilon, (int, float))
+        or not math.isfinite(float(normalization_epsilon))
+        or float(normalization_epsilon) <= 0.0
+    ):
+        raise RuntimeError(
+            "processor_config.json normalization_epsilon must be finite and positive, "
+            f"got {normalization_epsilon!r}."
+        )
+    if tokenizer_config.get("processor_class") != "TimeBraidProcessor":
+        raise RuntimeError(
+            "tokenizer_config.json processor_class must be 'TimeBraidProcessor', got "
+            f"{tokenizer_config.get('processor_class')!r}."
+        )
+    extra_special_tokens = tokenizer_config.get("extra_special_tokens")
+    if not isinstance(extra_special_tokens, dict):
+        raise RuntimeError(
+            "tokenizer_config.json extra_special_tokens must be an object in a release "
+            f"artifact, got {type(extra_special_tokens).__name__}."
+        )
+    if tokenizer_config.get("fix_mistral_regex") is not False:
+        raise RuntimeError(
+            "tokenizer_config.json fix_mistral_regex must be false to preserve the "
+            "trained tokenizer identity."
+        )
+
+    # Registration is process-local. The artifact itself stays free of remote
+    # auto_map code, while this gate exercises the same bare AutoProcessor call
+    # available to an installed TimeBraid package user.
+    register_timebraid_auto_classes()
+    from ..processing_timebraid import TimeBraidProcessor
+
+    processor = AutoProcessor.from_pretrained(
+        source,
+        use_fast=use_fast_tokenizer,
+        **hf_kwargs,
+    )
+    if not isinstance(processor, TimeBraidProcessor):
+        raise RuntimeError(
+            f"AutoProcessor must load TimeBraidProcessor, got {type(processor).__name__}."
+        )
+    if processor.max_spans_per_sample != max_spans_per_sample:
+        raise RuntimeError(
+            "AutoProcessor max_spans_per_sample disagrees with processor_config.json: "
+            f"{processor.max_spans_per_sample} != {max_spans_per_sample}."
+        )
+    if processor.normalization_epsilon != float(normalization_epsilon):
+        raise RuntimeError(
+            "AutoProcessor normalization_epsilon disagrees with processor_config.json: "
+            f"{processor.normalization_epsilon} != {normalization_epsilon}."
+        )
+
+
+def validate_timebraid_canonical_artifact(
+    source: str | Path,
+    *,
+    revision: str | None = None,
+    token: str | None = None,
+    cache_dir: str | None = None,
+    local_files_only: bool = False,
+    trust_remote_code: bool = False,
+    use_fast_tokenizer: bool = True,
+    low_cpu_mem_usage: bool = True,
+) -> LoadedTimeBraid:
+    """Run the one-time exact-load and finite-state release check."""
+    _resolved_source_hf_kwargs(trust_remote_code=trust_remote_code)
+    resolved_source = _resolve_hf_source(
+        _normalize_source(source, field_name="source"),
+        revision=revision,
+        token=token,
+        cache_dir=cache_dir,
+        local_files_only=local_files_only,
+    )
+    _validate_timebraid_processor_artifact(
+        resolved_source,
+        use_fast_tokenizer=use_fast_tokenizer,
+        trust_remote_code=trust_remote_code,
+    )
+    loaded = load_timebraid_checkpoint(
+        resolved_source,
+        compute_dtype="auto",
+        weight_dtype="auto",
+        device=None,
+        local_files_only=True,
+        trust_remote_code=trust_remote_code,
+        use_fast_tokenizer=use_fast_tokenizer,
+        low_cpu_mem_usage=low_cpu_mem_usage,
+    )
+    nested_config = getattr(getattr(loaded.model, "config", None), "llm_config", None)
+    declared_dtype = str(getattr(nested_config, "dtype", None)).removeprefix("torch.")
+    if declared_dtype != "float32":
+        raise RuntimeError(
+            "TimeBraid release artifact must declare llm_config.dtype='float32', "
+            f"got {declared_dtype!r}."
+        )
+    non_f32_tensors = {
+        tensor_name: physical_dtype
+        for tensor_name, physical_dtype in loaded.physical_tensor_dtypes
+        if physical_dtype != "F32"
+    }
+    if non_f32_tensors:
+        sample = dict(list(sorted(non_f32_tensors.items()))[:20])
+        raise RuntimeError(
+            "TimeBraid release safetensors must physically contain only F32 tensors; "
+            f"non_f32={sample}."
+        )
+    state = loaded.model.state_dict()
+    dtype_mismatches = {}
+    for tensor_name, physical_dtype in loaded.physical_tensor_dtypes:
+        tensor = state.get(tensor_name)
+        if tensor is None:
+            raise RuntimeError(
+                f"TimeBraid release tensor is absent after exact load: {tensor_name!r}."
+            )
+        loaded_dtype = _safetensors_dtype_code(tensor.dtype)
+        if physical_dtype != loaded_dtype:
+            dtype_mismatches[tensor_name] = {
+                "physical": physical_dtype,
+                "loaded": loaded_dtype,
+            }
+    if dtype_mismatches:
+        sample = dict(list(sorted(dtype_mismatches.items()))[:20])
+        raise RuntimeError(
+            "TimeBraid release safetensors must physically match their loaded dtypes; "
+            f"mismatches={sample}. Validate the final F32 artifact without "
+            "load-time dtype conversion."
+        )
+    _validate_finite_model_tensors(
+        loaded.model,
+        owner="TimeBraid release artifact",
+    )
+    return loaded
+
+
+__all__ = [
+    "LoadedTimeBraid",
+    "ValidatedTimeBraidCheckpointSource",
+    "load_timebraid_checkpoint",
+    "load_timebraid_tokenizer",
+    "validate_timebraid_canonical_artifact",
+    "validate_timebraid_model",
+    "validate_timebraid_checkpoint_source",
+]

timebraid/model/mot/__init__.py

@@ -0,0 +1,5 @@
+"""Public output type for native TimeBraid generation."""
+
+from .model import TimeBraidGenerateOutput
+
+__all__ = ["TimeBraidGenerateOutput"]

timebraid/model/mot/attention.py

@@ -0,0 +1,447 @@
+"""Mixed language and time-series attention for TimeBraid."""
+
+from __future__ import annotations
+
+from typing import Optional
+
+import torch
+import torch.nn as nn
+
+
+class _GlobalCausalJointSelfAttention(nn.Module):
+    """Projection-free packed-attention kernel shared by native and residual paths."""
+
+    def __init__(
+        self,
+        q_heads: int,
+        q_head_dim: int,
+        *,
+        lang_hidden_size: Optional[int] = None,
+    ):
+        super().__init__()
+        self.q_heads = q_heads
+        self.q_head_dim = q_head_dim
+        self.q_attention_hidden_size = self.q_heads * self.q_head_dim
+        # Residual fusion projects language, generation-TS and understanding-TS
+        # into one shared Qwen head geometry before packing, so its softmax scale
+        # is tied to `q_head_dim`. Native language attention keeps Qwen's
+        # `self_attn.scaling`; native TimesFM replay uses its own head dimension
+        # for the fixed attention-kernel scale.
+        self.scale = float(self.q_head_dim) ** -0.5
+        self.lang_hidden_size = int(lang_hidden_size or self.q_attention_hidden_size)
+        if self.lang_hidden_size <= 0:
+            raise ValueError(
+                f"`lang_hidden_size` must be positive, got {self.lang_hidden_size}."
+            )
+        self._flash_attn_varlen_func = None
+        self._fa_peft_integration_check = None
+
+    def _get_flash_attn_varlen_func(self):
+        if self._flash_attn_varlen_func is not None:
+            return self._flash_attn_varlen_func
+        try:
+            from flash_attn import flash_attn_varlen_func
+        except Exception as exc:
+            raise RuntimeError(
+                "MoT mixed attention is FA2-only but `flash_attn` could not be imported. "
+                "Install flash-attn in the active TimeBraid Python environment."
+            ) from exc
+        self._flash_attn_varlen_func = flash_attn_varlen_func
+        return self._flash_attn_varlen_func
+
+    def _get_fa_peft_integration_check(self):
+        if self._fa_peft_integration_check is not None:
+            return self._fa_peft_integration_check
+        try:
+            from transformers.modeling_flash_attention_utils import (
+                fa_peft_integration_check,
+            )
+        except Exception as exc:
+            raise RuntimeError(
+                "Failed to import transformers.modeling_flash_attention_utils.fa_peft_integration_check. "
+                "This runtime patch follows HF FA2 dtype reconciliation logic and requires this symbol."
+            ) from exc
+        self._fa_peft_integration_check = fa_peft_integration_check
+        return self._fa_peft_integration_check
+
+    def _infer_fa2_target_dtype_like_hf(self, q: torch.Tensor) -> Optional[torch.dtype]:
+        """Mirror HF FA2 target-dtype inference for packed-attention calls."""
+        target_dtype: Optional[torch.dtype] = None
+        if q.dtype == torch.float32:
+            if torch.is_autocast_enabled():
+                target_dtype = torch.get_autocast_gpu_dtype()
+            elif hasattr(self, "config") and hasattr(
+                self.config, "_pre_quantization_dtype"
+            ):
+                target_dtype = self.config._pre_quantization_dtype
+            else:
+                first_linear = next(
+                    (layer for layer in self.modules() if isinstance(layer, nn.Linear)),
+                    None,
+                )
+                if first_linear is None:
+                    raise RuntimeError(
+                        "Cannot infer FA2 target dtype because no nn.Linear module was found in bridge layer."
+                    )
+                target_dtype = first_linear.weight.dtype
+        return target_dtype
+
+    def _reconcile_fa2_qkv(
+        self,
+        *,
+        q: torch.Tensor,
+        k: torch.Tensor,
+        v: torch.Tensor,
+        mode_label: str,
+    ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
+        target_dtype = self._infer_fa2_target_dtype_like_hf(q)
+        fa_peft_integration_check = self._get_fa_peft_integration_check()
+        q_fa, k_fa, v_fa = fa_peft_integration_check(q, k, v, target_dtype)
+        if q_fa.dtype not in {torch.float16, torch.bfloat16}:
+            raise RuntimeError(
+                f"MoT mixed FA2 {mode_label} requires fp16/bf16 Q/K/V after HF-style dtype reconciliation, "
+                f"but got dtype={q_fa.dtype}, original_dtype={q.dtype}, "
+                f"target_dtype={target_dtype}, autocast_enabled={torch.is_autocast_enabled()}."
+            )
+        return q_fa, k_fa, v_fa
+
+    def _reconcile_fast_qkv(
+        self,
+        *,
+        q: torch.Tensor,
+        k: torch.Tensor,
+        v: torch.Tensor,
+        mode_label: str,
+    ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
+        """Normalize Q/K/V for the shared FA2 packed-attention kernel."""
+        return self._reconcile_fa2_qkv(q=q, k=k, v=v, mode_label=mode_label)
+
+    def _cu_seqlens(self, *, seqlens: list[int], device: torch.device) -> torch.Tensor:
+        cu = [0]
+        for seq_len in seqlens:
+            cu.append(cu[-1] + int(seq_len))
+        return torch.tensor(cu, device=device, dtype=torch.int32)
+
+    def _seqlens_from_cu_seqlens(self, *, cu_seqlens: torch.Tensor) -> list[int]:
+        if cu_seqlens.ndim != 1 or cu_seqlens.shape[0] == 0:
+            raise RuntimeError(
+                "Expected a non-empty 1D cu_seqlens tensor, got "
+                f"shape={tuple(cu_seqlens.shape)}."
+            )
+        cu_values = cu_seqlens.detach().to(device="cpu", dtype=torch.long).tolist()
+        return [
+            int(end) - int(start) for start, end in zip(cu_values[:-1], cu_values[1:])
+        ]
+
+    def _run_varlen_flash_attention(
+        self,
+        *,
+        q_sorted: torch.Tensor,
+        k_sorted: torch.Tensor,
+        v_sorted: torch.Tensor,
+        q_positions_sorted: torch.Tensor,
+        k_positions_sorted: torch.Tensor,
+        softmax_scale: float,
+        dropout_p: float,
+        mode_label: str,
+        q_seqlens: Optional[list[int]] = None,
+        k_seqlens: Optional[list[int]] = None,
+        cu_seqlens_q: Optional[torch.Tensor] = None,
+        cu_seqlens_k: Optional[torch.Tensor] = None,
+        max_seqlen_q: Optional[int] = None,
+        max_seqlen_k: Optional[int] = None,
+    ) -> torch.Tensor:
+        if q_sorted.ndim != 3 or k_sorted.ndim != 3 or v_sorted.ndim != 3:
+            raise RuntimeError(
+                "Fast MoT attention expects packed [N,H,D] Q/K/V tensors, got "
+                f"q={tuple(q_sorted.shape)}, k={tuple(k_sorted.shape)}, v={tuple(v_sorted.shape)}."
+            )
+        if k_sorted.shape != v_sorted.shape:
+            raise RuntimeError(
+                "Fast MoT attention expects aligned K/V shapes, got "
+                f"k={tuple(k_sorted.shape)}, v={tuple(v_sorted.shape)}."
+            )
+        if (
+            q_positions_sorted.ndim != 1
+            or q_positions_sorted.shape[0] != q_sorted.shape[0]
+        ):
+            raise RuntimeError(
+                "Fast MoT attention expects 1D query positions aligned to Q, got "
+                f"positions={tuple(q_positions_sorted.shape)}, q_tokens={q_sorted.shape[0]}."
+            )
+        if (
+            k_positions_sorted.ndim != 1
+            or k_positions_sorted.shape[0] != k_sorted.shape[0]
+        ):
+            raise RuntimeError(
+                "Fast MoT attention expects 1D key positions aligned to K, got "
+                f"positions={tuple(k_positions_sorted.shape)}, k_tokens={k_sorted.shape[0]}."
+            )
+
+        if q_seqlens is None:
+            if cu_seqlens_q is None:
+                raise RuntimeError(
+                    "Fast MoT attention requires either q_seqlens or cu_seqlens_q."
+                )
+            q_seqlens = self._seqlens_from_cu_seqlens(cu_seqlens=cu_seqlens_q)
+        if k_seqlens is None:
+            if cu_seqlens_k is None:
+                raise RuntimeError(
+                    "Fast MoT attention requires either k_seqlens or cu_seqlens_k."
+                )
+            k_seqlens = self._seqlens_from_cu_seqlens(cu_seqlens=cu_seqlens_k)
+        if max_seqlen_q is None:
+            max_seqlen_q = max(q_seqlens) if q_seqlens else 0
+        if max_seqlen_k is None:
+            max_seqlen_k = max(k_seqlens) if k_seqlens else 0
+
+        if q_sorted.device.type != "cuda":
+            raise RuntimeError(
+                f"MoT FA2 {mode_label} requires CUDA device, got {q_sorted.device.type}."
+            )
+        q_fast, k_fast, v_fast = self._reconcile_fast_qkv(
+            q=q_sorted,
+            k=k_sorted,
+            v=v_sorted,
+            mode_label=mode_label,
+        )
+        if cu_seqlens_q is None:
+            cu_seqlens_q = self._cu_seqlens(seqlens=q_seqlens, device=q_sorted.device)
+        if cu_seqlens_k is None:
+            cu_seqlens_k = self._cu_seqlens(seqlens=k_seqlens, device=q_sorted.device)
+        return self._get_flash_attn_varlen_func()(
+            q_fast,
+            k_fast,
+            v_fast,
+            cu_seqlens_q,
+            cu_seqlens_k,
+            max_seqlen_q,
+            max_seqlen_k,
+            dropout_p=dropout_p,
+            softmax_scale=softmax_scale,
+            causal=True,
+        )
+
+    def run_packed_self_attention(
+        self,
+        *,
+        q_sorted: torch.Tensor,
+        k_sorted: torch.Tensor,
+        v_sorted: torch.Tensor,
+        cu_seqlens: torch.Tensor,
+        max_seqlen: int,
+        rope_positions_sorted: torch.Tensor,
+        softmax_scale: float,
+        dropout_p: float,
+        mode_label: str,
+    ) -> torch.Tensor:
+        """Shared FA2 packed-attention entrypoint for native and residual callers."""
+        return self._run_varlen_flash_attention(
+            q_sorted=q_sorted,
+            k_sorted=k_sorted,
+            v_sorted=v_sorted,
+            q_positions_sorted=rope_positions_sorted,
+            k_positions_sorted=rope_positions_sorted,
+            softmax_scale=softmax_scale,
+            dropout_p=dropout_p,
+            mode_label=mode_label,
+            cu_seqlens_q=cu_seqlens,
+            cu_seqlens_k=cu_seqlens,
+            max_seqlen_q=max_seqlen,
+            max_seqlen_k=max_seqlen,
+        )
+
+
+class ResidualStreamProjection(nn.Module):
+    """Project one stream into and out of the shared residual-attention heads."""
+
+    def __init__(
+        self,
+        *,
+        input_size: int,
+        attention_heads: int,
+        head_dim: int,
+    ) -> None:
+        super().__init__()
+        if input_size <= 0:
+            raise ValueError(f"input_size must be positive, got {input_size}.")
+        if attention_heads <= 0:
+            raise ValueError(
+                f"attention_heads must be positive, got {attention_heads}."
+            )
+        if head_dim <= 0:
+            raise ValueError(f"head_dim must be positive, got {head_dim}.")
+
+        self.input_size = int(input_size)
+        self.attention_heads = int(attention_heads)
+        self.head_dim = int(head_dim)
+        attention_size = self.attention_heads * self.head_dim
+
+        # Keep registration order identical across all streams. Q/K use the
+        # Qwen-style post-projection head norm; V remains the raw projection.
+        #
+        # The epsilon is fixed at 1e-6 on purpose rather than mirrored from the
+        # host LM's `rms_norm_eps`. These norms carry trained weights, so the
+        # value they were trained under is part of this module's topology, the
+        # same as its head dimension; reading it from whatever backbone happens
+        # to be loaded would change a trained module's numerics. It matches
+        # Qwen3's default, so the two agree in practice today. Persisting it in
+        # the MoT restore contract is the proper fix and is deferred to the
+        # change that can update the training tree and the published configs
+        # together.
+        self.input_norm = nn.RMSNorm(
+            self.input_size, eps=1.0e-6, elementwise_affine=True
+        )
+        self.q_proj = nn.Linear(self.input_size, attention_size, bias=False)
+        self.q_norm = nn.RMSNorm(self.head_dim, eps=1.0e-6, elementwise_affine=True)
+        self.k_proj = nn.Linear(self.input_size, attention_size, bias=False)
+        self.k_norm = nn.RMSNorm(self.head_dim, eps=1.0e-6, elementwise_affine=True)
+        self.v_proj = nn.Linear(self.input_size, attention_size, bias=False)
+        self.out_proj = nn.Linear(attention_size, self.input_size, bias=False)
+        nn.init.zeros_(self.out_proj.weight)
+
+    @staticmethod
+    def _apply_head_norm(states: torch.Tensor, norm: nn.Module) -> torch.Tensor:
+        normalized = norm(states.to(dtype=norm.weight.dtype))
+        return normalized.to(dtype=states.dtype)
+
+    def project_qkv(
+        self, hidden: torch.Tensor
+    ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
+        if hidden.ndim != 2 or int(hidden.shape[-1]) != self.input_size:
+            raise RuntimeError(
+                "Residual stream projection expects [N,H], got "
+                f"shape={tuple(hidden.shape)}, expected_hidden={self.input_size}."
+            )
+        if int(hidden.shape[0]) == 0:
+            empty = hidden.new_empty((0, self.attention_heads, self.head_dim))
+            return empty, empty, empty
+
+        normalized = self.input_norm(hidden.to(dtype=self.input_norm.weight.dtype)).to(
+            dtype=self.q_proj.weight.dtype
+        )
+        token_count = int(normalized.shape[0])
+        head_shape = (token_count, self.attention_heads, self.head_dim)
+        q = self.q_proj(normalized).view(head_shape)
+        k = self.k_proj(normalized).view(head_shape)
+        v = self.v_proj(normalized).view(head_shape)
+        q = self._apply_head_norm(q, self.q_norm)
+        k = self._apply_head_norm(k, self.k_norm)
+        return q.contiguous(), k.contiguous(), v.contiguous()
+
+    def project_output(
+        self, out_heads: torch.Tensor, *, output_dtype: torch.dtype
+    ) -> torch.Tensor:
+        expected_head_shape = (self.attention_heads, self.head_dim)
+        if out_heads.ndim != 3 or tuple(out_heads.shape[1:]) != expected_head_shape:
+            raise RuntimeError(
+                "Residual stream output expects [N,H,D], got "
+                f"shape={tuple(out_heads.shape)}, expected_heads={expected_head_shape}."
+            )
+        flat = out_heads.reshape(out_heads.shape[0], -1)
+        delta = self.out_proj(flat.to(dtype=self.out_proj.weight.dtype))
+        return delta.to(dtype=output_dtype)
+
+
+class GlobalResidualAttention(_GlobalCausalJointSelfAttention):
+    """One joint residual-attention block over language and paired TS streams."""
+
+    def __init__(
+        self,
+        *,
+        language_hidden_size: int,
+        generation_hidden_size: Optional[int],
+        understanding_hidden_size: Optional[int],
+        attention_heads: int,
+        head_dim: int,
+    ) -> None:
+        if generation_hidden_size is None and understanding_hidden_size is None:
+            raise ValueError(
+                "GlobalResidualAttention requires generation or understanding."
+            )
+        super().__init__(
+            q_heads=attention_heads,
+            q_head_dim=head_dim,
+            lang_hidden_size=language_hidden_size,
+        )
+
+        # Module registration order is part of the checkpoint schema.
+        self.language = ResidualStreamProjection(
+            input_size=language_hidden_size,
+            attention_heads=attention_heads,
+            head_dim=head_dim,
+        )
+        self.generation = (
+            ResidualStreamProjection(
+                input_size=generation_hidden_size,
+                attention_heads=attention_heads,
+                head_dim=head_dim,
+            )
+            if generation_hidden_size is not None
+            else None
+        )
+        self.understanding = (
+            ResidualStreamProjection(
+                input_size=understanding_hidden_size,
+                attention_heads=attention_heads,
+                head_dim=head_dim,
+            )
+            if understanding_hidden_size is not None
+            else None
+        )
+        self.residual_attention_mode = True
+
+    @staticmethod
+    def _require_stream(
+        stream: Optional[ResidualStreamProjection], *, name: str
+    ) -> ResidualStreamProjection:
+        if stream is None:
+            raise RuntimeError(
+                f"GlobalResidualAttention has no {name} stream on this layer."
+            )
+        return stream
+
+    def project_language_qkv(
+        self, hidden: torch.Tensor
+    ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
+        return self.language.project_qkv(hidden)
+
+    def project_generation_qkv(
+        self, hidden: torch.Tensor
+    ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
+        return self._require_stream(self.generation, name="generation").project_qkv(
+            hidden
+        )
+
+    def project_understanding_qkv(
+        self, hidden: torch.Tensor
+    ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
+        return self._require_stream(
+            self.understanding, name="understanding"
+        ).project_qkv(hidden)
+
+    def project_language_delta(
+        self, out_heads: torch.Tensor, *, output_dtype: torch.dtype
+    ) -> torch.Tensor:
+        return self.language.project_output(out_heads, output_dtype=output_dtype)
+
+    def project_generation_delta(
+        self, out_heads: torch.Tensor, *, output_dtype: torch.dtype
+    ) -> torch.Tensor:
+        return self._require_stream(self.generation, name="generation").project_output(
+            out_heads, output_dtype=output_dtype
+        )
+
+    def project_understanding_delta(
+        self, out_heads: torch.Tensor, *, output_dtype: torch.dtype
+    ) -> torch.Tensor:
+        return self._require_stream(
+            self.understanding, name="understanding"
+        ).project_output(out_heads, output_dtype=output_dtype)
+
+
+__all__ = [
+    "GlobalResidualAttention",
+    "ResidualStreamProjection",
+    "_GlobalCausalJointSelfAttention",
+]

timebraid/model/mot/bridge_runtime.py

@@ -0,0 +1,471 @@
+"""Initialize TimeBraid's registered TimesFM and global-attention components."""
+
+from __future__ import annotations
+
+import copy
+import math
+from typing import List, Optional
+
+import torch
+import torch.nn as nn
+
+from .attention import GlobalResidualAttention, _GlobalCausalJointSelfAttention
+from .pairing import (
+    _build_pair_map,
+    _resolve_pairable_q_layer_indices,
+)
+from .runtime_state import MoTRuntimeOptions
+from .structures import build_timebraid_layer_plan
+
+
+def _build_upsampled_tsfm_layer_plan(source_depth: int, target_depth: int) -> List[int]:
+    """Build a concrete source-layer copy plan for an aligned TS stack."""
+    if source_depth <= 0:
+        raise ValueError(f"source_depth must be positive, got {source_depth}")
+    if target_depth <= 0:
+        raise ValueError(f"target_depth must be positive, got {target_depth}")
+    if target_depth < source_depth:
+        raise ValueError(
+            "TS layer upsampling expects target_depth >= source_depth, got "
+            f"source_depth={source_depth}, target_depth={target_depth}"
+        )
+    if target_depth == source_depth:
+        return list(range(source_depth))
+    if source_depth == 1:
+        return [0] * target_depth
+
+    source_last = source_depth - 1
+    target_last = target_depth - 1
+    plan: List[int] = []
+    for runtime_idx in range(target_depth):
+        source_idx = round(runtime_idx * source_last / target_last)
+        plan.append(int(source_idx))
+    return plan
+
+
+def _build_dsfp_tsfm_layer_plan(
+    source_depth: int, target_depth: int
+) -> tuple[List[int], List[int]]:
+    """Return the copied source layer and original-layer positions for DSFP."""
+    if source_depth <= 0:
+        raise ValueError(f"source_depth must be positive, got {source_depth}")
+    if target_depth < source_depth:
+        raise ValueError(
+            "DSFP expects target_depth >= source_depth, got "
+            f"source_depth={source_depth}, target_depth={target_depth}."
+        )
+
+    source_pair_map = _build_pair_map(
+        num_q_layers=target_depth,
+        num_t_layers=source_depth,
+        pairing_mode="interleaved",
+    )
+    source_positions = sorted(source_pair_map)
+    copy_plan = _build_upsampled_tsfm_layer_plan(
+        source_depth=source_depth, target_depth=target_depth
+    )
+    for runtime_idx, source_idx in source_pair_map.items():
+        if copy_plan[runtime_idx] != source_idx:
+            raise RuntimeError(
+                "DSFP source placement disagrees with the deterministic copy plan: "
+                f"runtime_idx={runtime_idx}, source_idx={source_idx}, "
+                f"copy_source_idx={copy_plan[runtime_idx]}."
+            )
+    return copy_plan, source_positions
+
+
+def _build_dsfp_tsfm_layers(
+    source_layers: List[nn.Module],
+    *,
+    target_depth: int,
+) -> nn.ModuleList:
+    """Keep source blocks once and fill every gap with a direct full-block copy."""
+    copy_plan, source_positions = _build_dsfp_tsfm_layer_plan(
+        source_depth=len(source_layers),
+        target_depth=target_depth,
+    )
+    source_position_set = set(source_positions)
+    layers: List[nn.Module] = []
+    for runtime_idx, source_idx in enumerate(copy_plan):
+        if runtime_idx in source_position_set:
+            layer = source_layers[source_idx]
+        else:
+            layer = copy.deepcopy(source_layers[source_idx])
+        layers.append(layer)
+
+    runtime_layers = nn.ModuleList(layers)
+    for param in runtime_layers.parameters():
+        param.requires_grad = False
+    return runtime_layers
+
+
+def initialize_timebraid_components(
+    owner: nn.Module,
+    llm_config,
+    runtime_options=None,
+) -> None:
+    """Register TimeBraid's canonical components directly on its model owner."""
+    if not isinstance(owner, nn.Module):
+        raise TypeError(
+            f"owner must be an initialized nn.Module, got {type(owner).__name__}."
+        )
+    if not hasattr(owner, "_modules"):
+        raise RuntimeError(
+            "owner nn.Module.__init__ must run before component initialization."
+        )
+    self = owner
+    config = llm_config
+    options = MoTRuntimeOptions(runtime_options)
+    # Host (Qwen) architecture info.
+    self.hidden_size = int(config.hidden_size)
+    self.num_heads = int(config.num_attention_heads)
+    self.num_q_layers = int(config.num_hidden_layers)
+    # Some Qwen configs expose only a subset of pairable full-attention layers.
+    self.q_pairable_layer_indices = _resolve_pairable_q_layer_indices(
+        config, self.num_q_layers
+    )
+    self.q_head_dim = int(
+        getattr(config, "head_dim", self.hidden_size // self.num_heads)
+    )
+    self.num_kv_heads = int(getattr(config, "num_key_value_heads", self.num_heads))
+    if self.num_heads % self.num_kv_heads != 0:
+        raise ValueError(
+            f"Qwen heads must be divisible by kv heads, got {self.num_heads}/{self.num_kv_heads}"
+        )
+    self.num_kv_groups = self.num_heads // self.num_kv_heads
+
+    self.pairing_mode = options.value("mot_pairing_mode").strip().lower()
+    if self.pairing_mode not in {"interleaved", "dsfp"}:
+        raise ValueError(
+            "`mot_pairing_mode` must be 'interleaved' or 'dsfp', got "
+            f"{self.pairing_mode!r}."
+        )
+    self.source_num_t_layers = options.value("mot_t_layers")
+    self.ts_gradient_checkpointing = options.value("mot_ts_gradient_checkpointing")
+    self.ts_gradient_checkpointing_use_reentrant = options.value(
+        "mot_ts_gradient_checkpointing_use_reentrant"
+    )
+    self.patch_size = options.value("mot_patch_size")
+    if self.patch_size <= 0:
+        raise ValueError(f"`mot_patch_size` must be positive, got {self.patch_size}.")
+    self.dsfp = self.pairing_mode == "dsfp"
+    # Optional understanding separation keeps the existing generation TS
+    # path intact, but adds a second TS tower for user/context-side spans.
+    # The depth selects the source TimesFM prefix. Both towers use the same
+    # pairing mode so one topology field owns every Qwen/TimesFM schedule.
+    self.understanding_pair_depth = options.value("mot_understanding_pair_depth")
+    if self.understanding_pair_depth < 0:
+        raise ValueError(
+            "`mot_understanding_pair_depth` must be >= 0, got "
+            f"{self.understanding_pair_depth}"
+        )
+    self.understanding_separation = self.understanding_pair_depth > 0
+    if self.understanding_separation:
+        if self.understanding_pair_depth > self.source_num_t_layers:
+            raise ValueError(
+                "`mot_understanding_pair_depth` cannot exceed `mot_t_layers`, got "
+                f"understanding_pair_depth={self.understanding_pair_depth}, "
+                f"source_t_layers={self.source_num_t_layers}"
+            )
+        if self.understanding_pair_depth > self.num_q_layers:
+            raise ValueError(
+                "Understanding TS pairing expects the Qwen depth to be at least the selected "
+                "source prefix depth, got "
+                f"understanding_pair_depth={self.understanding_pair_depth}, q_layers={self.num_q_layers}"
+            )
+    self.num_t_layers = int(self.source_num_t_layers)
+    if self.dsfp:
+        if len(self.q_pairable_layer_indices) != self.num_q_layers:
+            raise ValueError(
+                "`mot_pairing_mode=dsfp` requires every Qwen layer to be pairable, but "
+                "some layers are excluded by `config.layer_types`: "
+                f"pairable={len(self.q_pairable_layer_indices)}, total={self.num_q_layers}."
+            )
+        if self.source_num_t_layers != 20:
+            raise ValueError(
+                "`mot_pairing_mode=dsfp` requires the standard 20-layer TimesFM source stack, got "
+                f"mot_t_layers={self.source_num_t_layers}."
+            )
+        if self.source_num_t_layers > self.num_q_layers:
+            raise ValueError(
+                "`mot_pairing_mode=dsfp` expects the source TimesFM stack to fit inside Qwen depth, got "
+                f"source_t_layers={self.source_num_t_layers}, q_layers={self.num_q_layers}."
+            )
+        self.num_t_layers = int(self.num_q_layers)
+    self.num_understanding_t_layers = (
+        self.num_q_layers
+        if self.understanding_separation and self.dsfp
+        else self.understanding_pair_depth
+    )
+    # Packed rows can contain many short TS spans from multiple original samples.
+    self.max_spans_per_sample = options.value("mot_max_spans_per_sample")
+    if self.max_spans_per_sample <= 0:
+        raise ValueError(
+            f"max_spans_per_sample must be positive, got {self.max_spans_per_sample}."
+        )
+    # TimesFM structural defaults (2.5-200M).
+    self.tsfm_num_heads = options.value("mot_tsfm_num_heads")
+    self.tsfm_head_dim = options.value("mot_tsfm_head_dim")
+    self.tsfm_hidden_size = options.value("mot_tsfm_hidden_size")
+    if self.tsfm_hidden_size != self.tsfm_num_heads * self.tsfm_head_dim:
+        raise ValueError(
+            "TimesFM hidden/head mismatch: "
+            f"hidden={self.tsfm_hidden_size}, heads={self.tsfm_num_heads}, head_dim={self.tsfm_head_dim}"
+        )
+    self.understanding_tsfm_num_heads = self.tsfm_num_heads
+    self.understanding_tsfm_head_dim = self.tsfm_head_dim
+    self.understanding_tsfm_hidden_size = self.tsfm_hidden_size
+    # Mixed/Qwen positions are configurable at the span boundary.
+    # `patch_slot` reserves one mixed slot per realized TS patch;
+    # `span_slot` is the historical coarse lattice where one TS span
+    # advances later text by one slot and all patches share anchor+1.
+    # Native TS rotary state resets to local 0..N-1 inside each span either
+    # way. The restore contract always supplies this key, so the default
+    # only covers hand-built option maps and must match the trained lattice.
+    self.mixed_position_mode = options.value("mot_mixed_position_mode").strip().lower()
+    if self.mixed_position_mode not in {"span_slot", "patch_slot"}:
+        raise ValueError(
+            "`mot_mixed_position_mode` must be one of {'span_slot', 'patch_slot'}, got "
+            f"{self.mixed_position_mode!r}."
+        )
+    self.pair_map = _build_pair_map(
+        num_q_layers=self.num_q_layers,
+        num_t_layers=self.num_t_layers,
+        pairing_mode=self.pairing_mode,
+        q_layer_indices=self.q_pairable_layer_indices,
+    )
+    if self.understanding_separation:
+        self.understanding_pair_map = _build_pair_map(
+            num_q_layers=self.num_q_layers,
+            num_t_layers=self.num_understanding_t_layers,
+            pairing_mode=self.pairing_mode,
+            q_layer_indices=self.q_pairable_layer_indices,
+        )
+    else:
+        self.understanding_pair_map = {}
+    self.layer_plan = build_timebraid_layer_plan(
+        num_llm_layers=self.num_q_layers,
+        generation_tsfm_by_llm_layer=self.pair_map,
+        understanding_tsfm_by_llm_layer=self.understanding_pair_map,
+    )
+    # The delimiter ABI is persisted in config; no tokenizer object is
+    # required by the numeric runtime.
+    self._timesfm_util = None
+    self.ts_open_token_id = options.value("mot_ts_open_token_id")
+    self.ts_close_token_id = options.value("mot_ts_close_token_id")
+    if self.ts_open_token_id < 0 or self.ts_close_token_id < 0:
+        raise ValueError("TimeBraid runtime requires persisted TS delimiter token IDs.")
+    if self.ts_open_token_id == self.ts_close_token_id:
+        raise ValueError(
+            "TimeBraid runtime TS delimiter token IDs must be distinct, got "
+            f"{self.ts_open_token_id}."
+        )
+    # Complete TimeBraid checkpoints restore the registered TimesFM towers through
+    # HF loading; TimeBraid inference never resolves a second expert checkpoint.
+    self._timesfm_target_dtype = options.value("mot_compute_dtype")
+    self.ts_loss_weight = options.value("mot_ts_loss_weight")
+    if self.ts_loss_weight < 0:
+        raise ValueError(f"`mot_ts_loss_weight` must be >=0, got {self.ts_loss_weight}")
+    self.lm_loss_weight = options.value("mot_lm_loss_weight")
+    if self.lm_loss_weight < 0:
+        raise ValueError(f"`mot_lm_loss_weight` must be >=0, got {self.lm_loss_weight}")
+    self.ts_roi_mse_alpha = options.value("mot_ts_roi_mse_alpha")
+    if self.ts_roi_mse_alpha < 0:
+        raise ValueError(
+            f"`mot_ts_roi_mse_alpha` must be >= 0, got {self.ts_roi_mse_alpha}"
+        )
+    self.ts_understanding_loss_weight = options.value(
+        "mot_ts_understanding_loss_weight"
+    )
+    self.has_understanding_head = options.value("has_understanding_head")
+    if (
+        not math.isfinite(self.ts_understanding_loss_weight)
+        or self.ts_understanding_loss_weight < 0
+    ):
+        raise ValueError(
+            "`mot_ts_understanding_loss_weight` must be null or a finite value >= 0, got "
+            f"{self.ts_understanding_loss_weight}."
+        )
+    if self.has_understanding_head and not self.understanding_separation:
+        raise ValueError(
+            "has_understanding_head requires an understanding TimesFM tower."
+        )
+    if self.ts_understanding_loss_weight > 0.0 and not self.has_understanding_head:
+        raise ValueError(
+            "ts_understanding_loss_weight > 0 requires has_understanding_head=true."
+        )
+    self.tsfm_output_patch_len = 0
+    self.tsfm_q_channels = 0
+    self.tsfm_decode_index = 0
+    self.tsfm_quantile_taus: List[float] = []
+    self.generation_tsfm: Optional[nn.Module] = None
+    self.understanding_tsfm: Optional[nn.Module] = None
+    _init_timesfm_eager(self)
+    if self.generation_tsfm is not None:
+        native_patch_size = int(self.generation_tsfm.p)
+        if native_patch_size <= 0:
+            raise RuntimeError(
+                f"TimesFM must expose a positive input patch size, got p={native_patch_size}."
+            )
+        if self.patch_size != native_patch_size:
+            raise RuntimeError(
+                "`mot_patch_size` must match the loaded TimesFM input patch size, got "
+                f"mot_patch_size={self.patch_size}, timesfm.p={native_patch_size}."
+            )
+        if self.understanding_tsfm is not None:
+            understanding_patch_size = int(self.understanding_tsfm.p)
+            if understanding_patch_size != native_patch_size:
+                raise RuntimeError(
+                    "Generation and understanding TimesFM input patch sizes must match, got "
+                    f"generation={native_patch_size}, understanding={understanding_patch_size}."
+                )
+    self.global_residual_attention = nn.ModuleDict(
+        {
+            str(step.llm_layer): GlobalResidualAttention(
+                language_hidden_size=self.hidden_size,
+                generation_hidden_size=(
+                    self.tsfm_hidden_size
+                    if step.generation_tsfm_layer is not None
+                    else None
+                ),
+                understanding_hidden_size=(
+                    self.understanding_tsfm_hidden_size
+                    if step.understanding_tsfm_layer is not None
+                    else None
+                ),
+                attention_heads=self.num_heads,
+                head_dim=self.q_head_dim,
+            )
+            for step in self.layer_plan
+            if step.generation_tsfm_layer is not None
+            or step.understanding_tsfm_layer is not None
+        }
+    )
+    self.packed_attention = _GlobalCausalJointSelfAttention(
+        q_heads=self.num_heads,
+        q_head_dim=self.q_head_dim,
+        lang_hidden_size=self.hidden_size,
+    )
+    self.understanding_head: Optional[nn.Module] = None
+    if self.has_understanding_head:
+        if self.generation_tsfm is None:
+            raise RuntimeError(
+                "TS understanding reconstruction head requires initialized TimesFM."
+            )
+        # `has_understanding_head` persists this fixed-shape topology in the HF
+        # config. The TSFM-hidden MLP reconstructs one native TimesFM input patch;
+        # `ts_understanding_loss_weight` separately weights its fixed 2:8
+        # time-domain/FFT objective in the routed-loss helper.
+        self.understanding_head = nn.Sequential(
+            nn.Linear(int(self.tsfm_hidden_size), int(self.tsfm_hidden_size)),
+            nn.GELU(),
+            nn.Linear(int(self.tsfm_hidden_size), native_patch_size),
+        )
+        if not any(param.is_meta for param in self.understanding_head.parameters()):
+            self.understanding_head.to(
+                device=torch.device("cpu"), dtype=self._timesfm_target_dtype
+            )
+
+
+def _build_timesfm2p5_model():
+    from ..._vendor.timesfm.timesfm_2p5.timesfm_2p5_torch import (
+        TimesFM_2p5_200M_torch_module,
+    )
+    from ..._vendor.timesfm.torch import util as timesfm_torch_util
+
+    # TimeBraid checkpoints already contain this module's complete state. Construct
+    # the registered tower directly; the removed upstream Hub wrapper would
+    # allocate a second class-level 200M model and expose an unrelated download
+    # path before Transformers restores the embedded weights.
+    timesfm_model = TimesFM_2p5_200M_torch_module()
+
+    for param in timesfm_model.parameters():
+        param.requires_grad = False
+    timesfm_model.eval()
+    if not any(param.is_meta for param in timesfm_model.parameters()):
+        timesfm_model.to("cpu")
+    return timesfm_model, timesfm_torch_util
+
+
+def _init_timesfm_eager(self) -> None:
+    # Build TimesFM immediately so the complete registered component surface exists.
+    if self.generation_tsfm is not None:
+        return
+
+    self.generation_tsfm, self._timesfm_util = _build_timesfm2p5_model()
+
+    # Validate expected source structure before optional TS-depth upsampling
+    # replaces the registered TimesFM layer stack.
+    if len(self.generation_tsfm.stacked_xf) < self.source_num_t_layers:
+        raise RuntimeError(
+            "TimesFM has fewer source layers than requested by the active MoT runtime: "
+            f"got {len(self.generation_tsfm.stacked_xf)}, required {self.source_num_t_layers}"
+        )
+    source_layers = list(self.generation_tsfm.stacked_xf[: self.source_num_t_layers])
+    if self.understanding_separation:
+        # Deep-copy the complete expert first, then take the understanding
+        # prefix from that independent object graph. DSFP deliberately keeps
+        # its source-position modules rather than copying them, so passing
+        # generation modules into the expansion helper would otherwise make
+        # the two registered towers share parameters and tensor storage.
+        understanding_model = copy.deepcopy(self.generation_tsfm)
+        understanding_source_layers = list(
+            understanding_model.stacked_xf[: self.understanding_pair_depth]
+        )
+        if self.dsfp:
+            understanding_layers = _build_dsfp_tsfm_layers(
+                understanding_source_layers,
+                target_depth=self.num_q_layers,
+            )
+        else:
+            understanding_layers = nn.ModuleList(understanding_source_layers)
+        for param in understanding_model.parameters():
+            param.requires_grad = False
+        for param in understanding_layers.parameters():
+            param.requires_grad = False
+        understanding_layers.eval()
+        if not any(param.is_meta for param in understanding_layers.parameters()):
+            understanding_layers.to(
+                device=torch.device("cpu"), dtype=self._timesfm_target_dtype
+            )
+        understanding_model.stacked_xf = understanding_layers
+        understanding_model.x = self.num_understanding_t_layers
+        understanding_model.eval()
+        if not any(param.is_meta for param in understanding_model.parameters()):
+            understanding_model.to("cpu")
+        self.understanding_tsfm = understanding_model
+    if self.dsfp:
+        dsfp_layers = _build_dsfp_tsfm_layers(
+            source_layers,
+            target_depth=self.num_q_layers,
+        )
+        dsfp_layers.eval()
+        if not any(param.is_meta for param in dsfp_layers.parameters()):
+            dsfp_layers.to(device=torch.device("cpu"), dtype=self._timesfm_target_dtype)
+        self.generation_tsfm.stacked_xf = dsfp_layers
+        self.generation_tsfm.x = self.num_t_layers
+
+    self.tsfm_output_patch_len = int(self.generation_tsfm.o)
+    self.tsfm_q_channels = int(self.generation_tsfm.q)
+    self.tsfm_decode_index = int(self.generation_tsfm.aridx)
+    quantiles = list(self.generation_tsfm.config.quantiles)
+    if len(quantiles) + 1 != self.tsfm_q_channels:
+        raise RuntimeError(
+            f"TimesFM quantile channel mismatch: q_channels={self.tsfm_q_channels}, quantiles={len(quantiles)}"
+        )
+    if not (0 <= self.tsfm_decode_index < self.tsfm_q_channels):
+        raise RuntimeError(
+            f"TimesFM decode index out of range: decode_idx={self.tsfm_decode_index}, q={self.tsfm_q_channels}"
+        )
+    # Build channel->tau mapping, reserving decode channel for point forecast.
+    quantile_iter = iter(quantiles)
+    channel_taus: List[float] = []
+    for channel_idx in range(self.tsfm_q_channels):
+        if channel_idx == self.tsfm_decode_index:
+            channel_taus.append(-1.0)
+        else:
+            channel_taus.append(float(next(quantile_iter)))
+    self.tsfm_quantile_taus = channel_taus
+
+
+__all__ = ["initialize_timebraid_components"]

timebraid/model/mot/bridge_runtime_forecast.py

@@ -0,0 +1,1317 @@
+"""Forecast materialization and loss helpers for TimeBraid."""
+
+from __future__ import annotations
+
+import math
+from typing import TYPE_CHECKING, Any, Dict, List, NamedTuple, Optional
+
+import torch
+
+from .structures import (
+    ROLE_TARGET,
+    TS_ROUTE_GENERATION,
+    TS_ROUTE_UNDERSTANDING,
+    build_mot_forecast_loss_accounting,
+)
+
+if TYPE_CHECKING:
+    from .bridge_runtime_mot import _MoTForecastTarget, _MoTRuntime
+
+_FP32_SINH_LIMIT = math.asinh(torch.finfo(torch.float32).max)
+_MOT_TS_GLOBAL_LOSS_DETACHED_CAP = 2.0
+_MOT_TS_ROI_LOSS_DETACHED_CAP = 2.0
+
+
+class _FlatForecastLossParts(NamedTuple):
+    point: torch.Tensor
+    quantile: torch.Tensor
+    global_point: torch.Tensor
+    global_quantile: torch.Tensor
+    roi_point: torch.Tensor
+    roi_weighted_point: torch.Tensor
+    global_uncapped: torch.Tensor
+    global_cap_scale: torch.Tensor
+    roi_uncapped: torch.Tensor
+    roi_weighted_uncapped: torch.Tensor
+    roi_cap_scale: torch.Tensor
+    roi_supervised_points: torch.Tensor
+    roi_supervised_targets: torch.Tensor
+
+
+def _finite_stats(tensor: torch.Tensor) -> str:
+    """Return compact finite-value stats for hard forecast diagnostics."""
+    detached = tensor.detach()
+    finite = detached[torch.isfinite(detached)]
+    if finite.numel() == 0:
+        return f"shape={tuple(detached.shape)}, dtype={detached.dtype}, finite=0"
+    finite_f32 = finite.to(dtype=torch.float32)
+    return (
+        f"shape={tuple(detached.shape)}, dtype={detached.dtype}, "
+        f"finite={int(finite.numel())}/{int(detached.numel())}, "
+        f"min={float(finite_f32.min().item()):.6g}, "
+        f"max={float(finite_f32.max().item()):.6g}, "
+        f"absmax={float(finite_f32.abs().max().item()):.6g}"
+    )
+
+
+def _first_bad_index(mask: torch.Tensor) -> tuple[int, ...]:
+    bad = torch.nonzero(mask, as_tuple=False)
+    if int(bad.numel()) == 0:
+        return ()
+    return tuple(int(value) for value in bad[0].detach().cpu().tolist())
+
+
+def _target_debug_record(target: Any) -> str:
+    return (
+        f"sample={int(target.sample_idx)}, slot={int(target.slot_idx)}, "
+        f"role={int(target.role_id)}, segment={int(target.segment_id)}, "
+        f"ts_range=[{int(target.ts_start)},{int(target.ts_end)}), "
+        f"loss_start={int(target.loss_start_idx)}, "
+        f"generation_start={int(target.generation_start_index)}"
+    )
+
+
+def _require_finite_tensor(
+    tensor: torch.Tensor, *, label: str, targets: List[Any]
+) -> None:
+    if bool(torch.all(torch.isfinite(tensor))):
+        return
+    bad_index = _first_bad_index(torch.logical_not(torch.isfinite(tensor)))
+    target_text = ""
+    if bad_index:
+        target_axis = int(bad_index[0]) if len(bad_index) >= 1 else -1
+        if 0 <= target_axis < len(targets):
+            target_text = f", target=({_target_debug_record(targets[target_axis])})"
+    raise RuntimeError(
+        f"Forecast tensor became non-finite at {label}: "
+        f"bad_index={bad_index}, {_finite_stats(tensor)}{target_text}."
+    )
+
+
+def _require_sinh_safe(
+    tensor: torch.Tensor,
+    *,
+    targets: List[Any],
+    backend_label: str,
+    target_axis: int = 0,
+) -> None:
+    """Raise loudly when normalized-space head outputs would overflow fp32 sinh."""
+    unsafe = tensor.detach().to(dtype=torch.float32).abs() > _FP32_SINH_LIMIT
+    if not bool(torch.any(unsafe)):
+        return
+    bad_index = _first_bad_index(unsafe)
+    target_text = ""
+    if bad_index and 0 <= int(target_axis) < len(bad_index):
+        bad_target = int(bad_index[int(target_axis)])
+        if 0 <= bad_target < len(targets):
+            target_text = f", target=({_target_debug_record(targets[bad_target])})"
+    raise RuntimeError(
+        f"{backend_label} head output would overflow fp32 sinh during forecast "
+        f"inverse transform: bad_index={bad_index}, "
+        f"limit={_FP32_SINH_LIMIT:.6g}, "
+        f"value={float(tensor[bad_index].detach().to(torch.float32).item()):.6g}, "
+        f"{_finite_stats(tensor)}{target_text}."
+    )
+
+
+def _point_forecast_loss_values(
+    pred_point: torch.Tensor, target: torch.Tensor
+) -> torch.Tensor:
+    """Compute elementwise MSE on the TimesFM decode channel."""
+    pred_point = pred_point.to(dtype=torch.float32)
+    target = target.to(device=pred_point.device, dtype=torch.float32)
+    return (pred_point - target).square()
+
+
+def _detached_per_item_loss_cap_scale(
+    losses: torch.Tensor, cap_value: float
+) -> torch.Tensor:
+    """Return detached per-item cap scales for already-normalized record losses."""
+    if losses.ndim != 1:
+        raise RuntimeError(
+            f"Per-item loss cap expects a rank-1 tensor, got {tuple(losses.shape)}."
+        )
+    losses = losses.to(dtype=torch.float32)
+    if float(cap_value) <= 0.0:
+        return torch.ones_like(losses)
+    return torch.clamp(
+        losses.new_tensor(float(cap_value)) / losses.detach().clamp_min(1e-12), max=1.0
+    )
+
+
+def _apply_timesfm_value_transform(values: torch.Tensor) -> torch.Tensor:
+    """Apply the fixed post-scale asinh transform before TimesFM consumes values."""
+    return torch.asinh(values)
+
+
+def _inverse_timesfm_value_transform(
+    values: torch.Tensor,
+    *,
+    targets: Optional[List[Any]] = None,
+    backend_label: str = "TS value transform",
+    target_axis: int = 0,
+) -> torch.Tensor:
+    """Invert the fixed asinh transform and fail before fp32 sinh overflow."""
+    _require_sinh_safe(
+        values,
+        targets=[] if targets is None else targets,
+        backend_label=backend_label,
+        target_axis=target_axis,
+    )
+    return values.to(dtype=torch.float32).sinh()
+
+
+def _transform_forecast_targets_for_loss(
+    targets: torch.Tensor,
+    *,
+    owner_mu: torch.Tensor,
+    owner_sigma: torch.Tensor,
+) -> torch.Tensor:
+    """Move raw forecast targets into the same scaled value frame as tower predictions."""
+    scaled = (
+        targets.to(dtype=torch.float32) - owner_mu.unsqueeze(1)
+    ) / owner_sigma.unsqueeze(1).clamp_min(1e-6)
+    return _apply_timesfm_value_transform(scaled)
+
+
+def _global_count_tensor(
+    global_count: Optional[object],
+    *,
+    reference: torch.Tensor,
+    label: str,
+) -> Optional[torch.Tensor]:
+    if global_count is None:
+        return None
+    count = torch.as_tensor(
+        global_count, device=reference.device, dtype=reference.dtype
+    )
+    if int(count.numel()) != 1:
+        raise RuntimeError(
+            f"{label} denominator must be scalar, got shape={tuple(count.shape)}."
+        )
+    count = count.reshape(())
+    if bool(count <= 0):
+        raise RuntimeError(
+            f"{label} denominator must be positive, got {float(count.detach().cpu().item())}."
+        )
+    return count
+
+
+def _reduce_per_record_loss(
+    per_record: torch.Tensor,
+    *,
+    global_count: Optional[object],
+    label: str,
+) -> torch.Tensor:
+    if per_record.ndim != 1:
+        raise RuntimeError(
+            f"{label} reduction expects a rank-1 per-record tensor, got {tuple(per_record.shape)}."
+        )
+    if int(per_record.numel()) == 0:
+        raise RuntimeError(f"{label} reduction requires at least one local record.")
+    per_record = per_record.to(dtype=torch.float32)
+    count = _global_count_tensor(global_count, reference=per_record, label=label)
+    if count is None:
+        return per_record.mean()
+    return per_record.sum() / count
+
+
+def _global_count_for_horizon(
+    counts_by_horizon: Optional[object], horizon: int
+) -> Optional[object]:
+    if counts_by_horizon is None:
+        return None
+    if isinstance(counts_by_horizon, dict):
+        if int(horizon) not in counts_by_horizon:
+            raise RuntimeError(
+                "TS global denominator mapping is missing a forecast horizon: "
+                f"horizon={int(horizon)}, available={sorted(int(key) for key in counts_by_horizon.keys())}."
+            )
+        return counts_by_horizon[int(horizon)]
+    return counts_by_horizon
+
+
+def _roi_loss_from_value_losses(
+    *,
+    value_losses: torch.Tensor,
+    roi_records: torch.Tensor,
+    keep_mask: torch.Tensor,
+    record_target_idx: torch.Tensor,
+    target_count: int,
+    global_target_count: Optional[object] = None,
+) -> tuple[torch.Tensor, torch.Tensor]:
+    """Reduce a per-value loss over ROI points, normalized inside each target span."""
+    if value_losses.ndim != 2:
+        raise RuntimeError(
+            f"ROI loss expects per-value losses [record,horizon], got {tuple(value_losses.shape)}."
+        )
+    if roi_records.shape != value_losses.shape or keep_mask.shape != value_losses.shape:
+        raise RuntimeError(
+            "ROI loss expects mask/value shapes to match: "
+            f"loss={tuple(value_losses.shape)}, roi={tuple(roi_records.shape)}, keep={tuple(keep_mask.shape)}."
+        )
+
+    value_losses = value_losses.to(dtype=torch.float32)
+
+    # ROI is a sparse data-side reweighting over already-supervised points.  We
+    # normalize within each target before averaging target spans so `alpha`
+    # means "add one ROI objective" rather than "scale with the number of ROI
+    # points or the forecast horizon".
+    roi_mask_float = (roi_records * keep_mask).to(dtype=torch.float32)
+    roi_den_by_target = torch.zeros(
+        (target_count,), device=value_losses.device, dtype=torch.float32
+    )
+    roi_num_by_target = torch.zeros(
+        (target_count,), device=value_losses.device, dtype=torch.float32
+    )
+    roi_den_by_target.scatter_add_(0, record_target_idx, roi_mask_float.sum(dim=1))
+    if not bool(torch.any(roi_den_by_target > 0.0)):
+        zero = value_losses.new_zeros(())
+        return zero, value_losses.new_zeros((0,))
+
+    roi_num_by_target.scatter_add_(
+        0,
+        record_target_idx,
+        (value_losses.to(dtype=torch.float32) * roi_mask_float).sum(dim=1),
+    )
+    active_roi_targets = roi_den_by_target > 0.0
+    roi_per_target = (
+        roi_num_by_target[active_roi_targets] / roi_den_by_target[active_roi_targets]
+    )
+    reduced_roi_loss = _reduce_per_record_loss(
+        roi_per_target,
+        global_count=global_target_count,
+        label="ROI target",
+    )
+    return reduced_roi_loss, roi_per_target
+
+
+def _roi_supervision_counts(
+    *,
+    roi_records: torch.Tensor,
+    keep_mask: torch.Tensor,
+    record_target_idx: torch.Tensor,
+    target_count: int,
+) -> tuple[torch.Tensor, torch.Tensor]:
+    """Count ROI supervision after intersecting data masks with active loss positions."""
+    if roi_records.ndim != 2:
+        raise RuntimeError(
+            f"ROI supervision count expects [record,horizon] ROI masks, got {tuple(roi_records.shape)}."
+        )
+    if keep_mask.shape != roi_records.shape:
+        raise RuntimeError(
+            "ROI supervision count expects keep/ROI shapes to match: "
+            f"roi={tuple(roi_records.shape)}, keep={tuple(keep_mask.shape)}."
+        )
+    if record_target_idx.ndim != 1 or int(record_target_idx.shape[0]) != int(
+        roi_records.shape[0]
+    ):
+        raise RuntimeError(
+            "ROI supervision count expects one target index per record, got "
+            f"targets={tuple(record_target_idx.shape)}, roi_records={tuple(roi_records.shape)}."
+        )
+    if target_count < 0:
+        raise RuntimeError(
+            f"ROI supervision count requires non-negative target_count, got {target_count}."
+        )
+
+    roi_keep = torch.logical_and(roi_records > 0.0, keep_mask.to(dtype=torch.bool))
+    roi_points = roi_keep.to(dtype=torch.float32).sum()
+    if target_count == 0:
+        return roi_points, roi_records.new_zeros((), dtype=torch.float32)
+    roi_den_by_target = torch.zeros(
+        (target_count,), device=roi_records.device, dtype=torch.float32
+    )
+    roi_den_by_target.scatter_add_(
+        0, record_target_idx, roi_keep.to(dtype=torch.float32).sum(dim=1)
+    )
+    roi_targets = (roi_den_by_target > 0.0).to(dtype=torch.float32).sum()
+    return roi_points, roi_targets
+
+
+def _roi_value_losses_from_predictions(
+    *,
+    pred_records: torch.Tensor,
+    gt_records: torch.Tensor,
+    decode_index: int,
+) -> torch.Tensor:
+    """Build per-value ROI MSE on the TimesFM decode channel."""
+    pred_records = pred_records.to(dtype=torch.float32)
+    gt_records = gt_records.to(device=pred_records.device, dtype=torch.float32)
+    decode_index = int(decode_index)
+    if decode_index < 0 or decode_index >= int(pred_records.shape[-1]):
+        raise RuntimeError(
+            "ROI point loss decode index is outside the forecast channel axis: "
+            f"decode_index={decode_index}, q_channels={int(pred_records.shape[-1])}."
+        )
+    return _point_forecast_loss_values(
+        pred_records[:, :, decode_index],
+        gt_records,
+    )
+
+
+def _require_known_output_space(output_space: str) -> str:
+    output_space = str(output_space).strip().lower()
+    if output_space not in {"real", "normalized"}:
+        raise RuntimeError(
+            f"Forecast output space must be 'real' or 'normalized', got {output_space!r}."
+        )
+    return output_space
+
+
+def _predict_mot_span_quantiles_batched(
+    self,
+    *,
+    mot_runtime: "_MoTRuntime",
+    targets: List["_MoTForecastTarget"],
+    output_space: str = "real",
+) -> tuple[torch.Tensor, List[int]]:
+    """Run one numeric-head projection per non-empty TS-route sub-batch.
+
+    `output_space="real"` applies the full inverse transform (sinh, then RevIN
+    reverse) for generation/eval; `"normalized"` returns the tower-space head
+    output for normalized-space training losses.
+    """
+    output_space = _require_known_output_space(output_space)
+    if self.generation_tsfm is None or self._timesfm_util is None:
+        raise RuntimeError("TimesFM model/util is not initialized.")
+
+    output_patch_len = int(self.tsfm_output_patch_len)
+    q_channels = int(self.tsfm_q_channels)
+    if not targets:
+        return mot_runtime.ts_hidden.new_empty((0, 0, output_patch_len, q_channels)), []
+
+    # Group targets by explicit TS route so each uses its matching registered
+    # numeric head, then restore caller order for downstream loss or decoding.
+    predictions_by_target: list[Optional[torch.Tensor]] = [None] * len(targets)
+    mu_by_target: list[Optional[torch.Tensor]] = [None] * len(targets)
+    sigma_by_target: list[Optional[torch.Tensor]] = [None] * len(targets)
+    patch_counts = [0] * len(targets)
+    grouped_indices: dict[int, list[int]] = {
+        int(TS_ROUTE_GENERATION): [],
+        int(TS_ROUTE_UNDERSTANDING): [],
+    }
+    for target_idx, target in enumerate(targets):
+        route_id = int(target.ts_route_id)
+        if route_id not in grouped_indices:
+            raise RuntimeError(
+                f"Packed forecast head cannot read from TS route id {route_id}."
+            )
+        grouped_indices[route_id].append(target_idx)
+
+    for route_id, target_indices in grouped_indices.items():
+        if not target_indices:
+            continue
+        route_targets = [targets[index] for index in target_indices]
+        batched_hidden, context_mu, context_sigma, route_patch_counts = (
+            _collect_batched_forecast_head_inputs(
+                self,
+                mot_runtime=mot_runtime,
+                targets=route_targets,
+            )
+        )
+        route_timesfm = (
+            self.generation_tsfm
+            if route_id == int(TS_ROUTE_GENERATION)
+            else self.understanding_tsfm
+        )
+        if route_timesfm is None:
+            route_name = (
+                "generation"
+                if route_id == int(TS_ROUTE_GENERATION)
+                else "understanding"
+            )
+            raise RuntimeError(
+                f"{route_name.capitalize()} routed forecast target requires its TimesFM expert."
+            )
+        out_proj = route_timesfm.output_projection_point
+        out_dtype = next(out_proj.parameters()).dtype
+        route_outputs = out_proj(batched_hidden.to(dtype=out_dtype)).reshape(
+            len(route_targets),
+            -1,
+            output_patch_len,
+            q_channels,
+        )
+        if int(route_outputs.shape[1]) != max(route_patch_counts):
+            raise RuntimeError(
+                "Batched TimesFM forecast head produced an unexpected patch axis: "
+                f"route={route_id}, output_patches={route_outputs.shape[1]}, "
+                f"expected={max(route_patch_counts)}."
+            )
+        for route_row, target_idx in enumerate(target_indices):
+            patch_count = int(route_patch_counts[route_row])
+            predictions_by_target[target_idx] = route_outputs[route_row, :patch_count]
+            mu_by_target[target_idx] = context_mu[route_row, :patch_count]
+            sigma_by_target[target_idx] = context_sigma[route_row, :patch_count]
+            patch_counts[target_idx] = patch_count
+
+    if any(prediction is None for prediction in predictions_by_target):
+        raise RuntimeError("Batched TimesFM forecast head left an unprojected target.")
+    normed_batched = torch.nn.utils.rnn.pad_sequence(
+        [prediction for prediction in predictions_by_target if prediction is not None],
+        batch_first=True,
+    )
+    if output_space == "normalized":
+        return normed_batched, patch_counts
+    else:
+        # The TimesFM tower consumes asinh(RevIN(x)), so sinh must run before
+        # the RevIN reverse transform.
+        normed_outputs = normed_batched.reshape(
+            len(targets), -1, output_patch_len * q_channels
+        )
+        real_outputs = _inverse_timesfm_value_transform(
+            normed_outputs,
+            targets=targets,
+            backend_label="TimesFM",
+        )
+        _require_finite_tensor(
+            real_outputs,
+            label="inverse_timesfm_value_transform(timesfm_head)",
+            targets=targets,
+        )
+        context_mu = torch.nn.utils.rnn.pad_sequence(
+            [value for value in mu_by_target if value is not None],
+            batch_first=True,
+        )
+        context_sigma = torch.nn.utils.rnn.pad_sequence(
+            [value for value in sigma_by_target if value is not None],
+            batch_first=True,
+            padding_value=1.0,
+        )
+        head_outputs = self._timesfm_util.revin(
+            real_outputs,
+            context_mu.to(device=real_outputs.device, dtype=real_outputs.dtype),
+            context_sigma.to(device=real_outputs.device, dtype=real_outputs.dtype),
+            reverse=True,
+        )
+    forecast_batched = head_outputs.reshape(
+        len(targets),
+        -1,
+        output_patch_len,
+        q_channels,
+    )
+    if int(forecast_batched.shape[1]) != max(patch_counts):
+        raise RuntimeError(
+            "Batched TimesFM forecast head produced an unexpected patch axis: "
+            f"output_patches={forecast_batched.shape[1]}, expected={max(patch_counts)}"
+        )
+    return forecast_batched, patch_counts
+
+
+def _collect_batched_forecast_head_inputs(
+    self,
+    *,
+    mot_runtime: "_MoTRuntime",
+    targets: List["_MoTForecastTarget"],
+) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, List[int]]:
+    """Pad TS hidden states and RevIN stats for a batched forecast head call."""
+    hidden_slices: List[torch.Tensor] = []
+    context_mu_parts: List[torch.Tensor] = []
+    context_sigma_parts: List[torch.Tensor] = []
+    patch_counts: List[int] = []
+    for target_idx, target in enumerate(targets):
+        if target.ts_end < target.ts_start:
+            raise RuntimeError(
+                "Packed forecast target has negative TS slice: "
+                f"target_idx={target_idx}, start={target.ts_start}, end={target.ts_end}"
+            )
+
+        patch_count = int(target.ts_end) - int(target.ts_start)
+        if patch_count <= 0:
+            raise RuntimeError(
+                "Batched packed forecast head requires every supervised target to have at least "
+                f"one TS patch, got target_idx={target_idx}, patches={patch_count}."
+            )
+        if patch_count != int(target.patch_valid_lengths.shape[0]):
+            raise RuntimeError(
+                "Packed TimesFM patch count mismatch before batched forecast head: "
+                f"target_idx={target_idx}, hidden_patches={patch_count}, "
+                f"valid_len_patches={target.patch_valid_lengths.shape[0]}"
+            )
+
+        context_mu = target.context_mu.reshape(-1)
+        context_sigma = target.context_sigma.reshape(-1)
+        if (
+            int(context_mu.shape[0]) != patch_count
+            or int(context_sigma.shape[0]) != patch_count
+        ):
+            raise RuntimeError(
+                "Packed TimesFM RevIN stats must align one value per target patch: "
+                f"target_idx={target_idx}, patches={patch_count}, "
+                f"mu={tuple(target.context_mu.shape)}, sigma={tuple(target.context_sigma.shape)}"
+            )
+
+        route_id = int(target.ts_route_id)
+        if route_id == int(TS_ROUTE_UNDERSTANDING):
+            source_hidden = mot_runtime.understanding_hidden
+            if source_hidden is None:
+                raise RuntimeError(
+                    "Understanding routed forecast target requires `mot_runtime.understanding_hidden`."
+                )
+        elif route_id == int(TS_ROUTE_GENERATION):
+            source_hidden = mot_runtime.ts_hidden
+        else:
+            raise RuntimeError(
+                f"Packed forecast head cannot read from TS route id {route_id}."
+            )
+
+        hidden_slices.append(source_hidden[target.ts_start : target.ts_end])
+        context_mu_parts.append(context_mu.to(device=source_hidden.device))
+        context_sigma_parts.append(context_sigma.to(device=source_hidden.device))
+        patch_counts.append(patch_count)
+
+    batched_hidden = torch.nn.utils.rnn.pad_sequence(hidden_slices, batch_first=True)
+    context_mu = torch.nn.utils.rnn.pad_sequence(context_mu_parts, batch_first=True)
+    context_sigma = torch.nn.utils.rnn.pad_sequence(
+        context_sigma_parts,
+        batch_first=True,
+        padding_value=1.0,
+    )
+    return batched_hidden, context_mu, context_sigma, patch_counts
+
+
+def _compute_flat_forecast_losses_from_predictions(
+    self,
+    *,
+    pred_quantiles_batched: torch.Tensor,
+    targets: List[Any],
+    patch_counts: List[int],
+    mot_ts_num_records_in_batch: Optional[object] = None,
+    mot_ts_num_roi_targets_in_batch: Optional[object] = None,
+    loss_space: str = "real",
+) -> Optional[_FlatForecastLossParts]:
+    """
+    Reduce all patch forecasts in one masked tensor while preserving patch-level weighting.
+
+    The previous path launched point/quantile loss kernels once per patch. This
+    helper builds one flat `[num_patch_forecasts, horizon]` mask and performs the
+    loss reductions in batched tensor ops. Each record is owned by patch i and
+    aligns it to the raw suffix immediately after that patch; detached per-record
+    caps may rescale its final objective contribution.
+
+    Predictions must be in the tower-normalized space (producer called with
+    `output_space="normalized"`); targets are scored against
+    `asinh((gt - mu) / sigma)` built from each record's owner-patch causal stats.
+    `loss_space` accepts only `"normalized"`; the `"real"` default exists so a
+    caller that forgets to state the space fails loudly instead of silently
+    picking a value frame.
+    """
+    loss_space = _require_known_output_space(loss_space)
+    if loss_space != "normalized":
+        raise RuntimeError("Routed TS losses must use loss_space='normalized'.")
+    if pred_quantiles_batched.ndim != 4:
+        raise RuntimeError(
+            "Flat forecast loss expects predictions shaped [target, patch, horizon, channel], got "
+            f"{tuple(pred_quantiles_batched.shape)}"
+        )
+    if int(pred_quantiles_batched.shape[0]) != len(targets):
+        raise RuntimeError(
+            "Flat forecast loss target count mismatch: "
+            f"pred_targets={pred_quantiles_batched.shape[0]}, targets={len(targets)}"
+        )
+    if len(patch_counts) != len(targets):
+        raise RuntimeError(
+            "Flat forecast loss patch-count metadata mismatch: "
+            f"patch_counts={len(patch_counts)}, targets={len(targets)}"
+        )
+
+    device = pred_quantiles_batched.device
+    dtype = torch.float32
+    max_pred_patches = int(pred_quantiles_batched.shape[1])
+    forecast_horizon = int(pred_quantiles_batched.shape[2])
+    q_channels = int(pred_quantiles_batched.shape[3])
+
+    values_by_target: List[torch.Tensor] = []
+    roi_masks_by_target: List[torch.Tensor] = []
+    patch_length_parts: List[torch.Tensor] = []
+    context_mu_parts: List[torch.Tensor] = []
+    context_sigma_parts: List[torch.Tensor] = []
+    value_lengths: List[int] = []
+    loss_start_idxs: List[int] = []
+
+    for target_idx, target in enumerate(targets):
+        pred_patch_count = int(patch_counts[target_idx])
+        if pred_patch_count < 0 or pred_patch_count > max_pred_patches:
+            raise RuntimeError(
+                "Flat forecast loss received invalid per-target patch count: "
+                f"target_idx={target_idx}, patches={pred_patch_count}, max_patches={max_pred_patches}"
+            )
+        context_mu = target.context_mu.reshape(-1).to(
+            device=device, dtype=torch.float32
+        )
+        context_sigma = target.context_sigma.reshape(-1).to(
+            device=device, dtype=torch.float32
+        )
+        if (
+            int(context_mu.shape[0]) != pred_patch_count
+            or int(context_sigma.shape[0]) != pred_patch_count
+        ):
+            raise RuntimeError(
+                "Normalized-space forecast loss stats must align one value per owner patch: "
+                f"target_idx={target_idx}, patches={pred_patch_count}, "
+                f"mu={tuple(target.context_mu.shape)}, sigma={tuple(target.context_sigma.shape)}."
+            )
+        context_mu_parts.append(context_mu)
+        context_sigma_parts.append(context_sigma)
+
+        if target.values.ndim != 1:
+            raise RuntimeError(
+                "Flat forecast loss expects raw TS values to be one-dimensional, got "
+                f"target_idx={target_idx}, values_shape={tuple(target.values.shape)}"
+            )
+        values_len = int(target.values.shape[0])
+        if pred_patch_count != int(target.patch_valid_lengths.shape[0]):
+            raise RuntimeError(
+                "Flat forecast loss patch count mismatch with valid lengths: "
+                f"target_idx={target_idx}, pred_patches={pred_patch_count}, "
+                f"valid_len_patches={target.patch_valid_lengths.shape[0]}"
+            )
+        if pred_patch_count == 0:
+            raise RuntimeError(
+                "Routed TS supervision requires at least one real TS patch per span."
+            )
+
+        # Raw targets must stay FP32 until after causal normalization. Casting
+        # large-offset series to BF16 here can erase the complete future delta.
+        values_by_target.append(target.values.to(device=device, dtype=torch.float32))
+        target_roi_mask = target.loss_roi_mask
+        if target_roi_mask is None:
+            target_roi_mask = torch.zeros_like(target.values, dtype=torch.float32)
+        target_roi_mask = target_roi_mask.to(device=device, dtype=torch.float32)
+        if target_roi_mask.ndim != 1 or int(target_roi_mask.shape[0]) != values_len:
+            raise RuntimeError(
+                "Flat forecast loss expects loss_roi_mask to align with raw TS values, got "
+                f"target_idx={target_idx}, mask_shape={tuple(target_roi_mask.shape)}, values_len={values_len}."
+            )
+        if bool(
+            torch.any(torch.logical_or(target_roi_mask < 0.0, target_roi_mask > 1.0))
+        ):
+            raise RuntimeError(
+                "Flat forecast loss expects loss_roi_mask values in [0, 1], got "
+                f"target_idx={target_idx}."
+            )
+        roi_masks_by_target.append(target_roi_mask)
+        patch_length_parts.append(
+            target.patch_valid_lengths.to(device=device, dtype=torch.long)
+        )
+        value_lengths.append(values_len)
+        loss_start_idx = int(target.loss_start_idx)
+        loss_start_idxs.append(loss_start_idx)
+
+    if not values_by_target:
+        return None
+
+    values_padded = torch.nn.utils.rnn.pad_sequence(values_by_target, batch_first=True)
+    roi_masks_padded = torch.nn.utils.rnn.pad_sequence(
+        roi_masks_by_target, batch_first=True
+    )
+    patch_lengths = torch.nn.utils.rnn.pad_sequence(
+        patch_length_parts, batch_first=True
+    )
+    context_mu_padded = torch.nn.utils.rnn.pad_sequence(
+        context_mu_parts, batch_first=True
+    )
+    context_sigma_padded = torch.nn.utils.rnn.pad_sequence(
+        context_sigma_parts, batch_first=True, padding_value=1.0
+    )
+    max_values_len = int(values_padded.shape[1])
+    if max_values_len <= 0:
+        raise RuntimeError(
+            "Flat forecast loss built supervision records without any raw target values."
+        )
+    if int(patch_lengths.shape[1]) != max_pred_patches:
+        raise RuntimeError(
+            "Flat forecast loss patch-length padding axis does not match predictions: "
+            f"patch_lengths={patch_lengths.shape[1]}, pred_patches={max_pred_patches}"
+        )
+
+    patch_counts_tensor = torch.tensor(patch_counts, device=device, dtype=torch.long)
+    values_lengths_tensor = torch.tensor(value_lengths, device=device, dtype=torch.long)
+    loss_start_idxs_tensor = torch.tensor(
+        loss_start_idxs, device=device, dtype=torch.long
+    )
+    prefix_ends = torch.cumsum(patch_lengths, dim=1)
+    last_patch_indices = torch.clamp(patch_counts_tensor - 1, min=0)
+    covered_lengths = prefix_ends.gather(1, last_patch_indices[:, None]).squeeze(1)
+    patch_axis = torch.arange(
+        max_pred_patches, device=device, dtype=torch.long
+    ).unsqueeze(0)
+    successor_patch_mask = patch_axis < (patch_counts_tensor - 1).unsqueeze(1)
+    bad_contract = torch.logical_or(patch_counts_tensor <= 0, patch_lengths[:, 0] <= 0)
+    bad_contract = torch.logical_or(
+        bad_contract, covered_lengths != values_lengths_tensor
+    )
+    if max_pred_patches > 1:
+        bad_successors = torch.any(
+            torch.logical_and(patch_lengths[:, 1:] <= 0, successor_patch_mask[:, :-1]),
+            dim=1,
+        )
+        bad_contract = torch.logical_or(bad_contract, bad_successors)
+    if bool(torch.any(bad_contract)):
+        bad_idx = int(torch.nonzero(bad_contract, as_tuple=False)[0].item())
+        raise RuntimeError(
+            "Packed routed TS supervision received an invalid patch/value contract: "
+            f"target_idx={bad_idx}, patches={int(patch_counts[bad_idx])}, "
+            f"covered_raw_points={int(covered_lengths[bad_idx].item())}, "
+            f"raw_values={int(values_lengths_tensor[bad_idx].item())}."
+        )
+
+    decode_index = int(self.tsfm_decode_index)
+    if decode_index < 0 or decode_index >= q_channels:
+        raise RuntimeError(
+            "TS forecast decode index is outside the forecast channel axis: "
+            f"decode_index={decode_index}, q_channels={q_channels}"
+        )
+    accounting = build_mot_forecast_loss_accounting(
+        patch_lengths=patch_lengths,
+        patch_counts=patch_counts_tensor,
+        value_lengths=values_lengths_tensor,
+        loss_start_idxs=loss_start_idxs_tensor,
+        forecast_horizon=forecast_horizon,
+    )
+    if accounting is None:
+        return None
+
+    keep_mask = accounting.keep_mask
+    record_target_idx = accounting.record_target_idx
+    record_patch_idx = accounting.record_patch_idx
+    raw_indices = accounting.raw_indices
+    valid_counts = accounting.valid_counts.to(dtype=dtype)
+    if bool(torch.any(valid_counts <= 0)):
+        raise RuntimeError(
+            "Flat forecast loss built a record with no supervised points."
+        )
+    if bool(torch.any(torch.logical_and(keep_mask, raw_indices >= max_values_len))):
+        raise RuntimeError(
+            "Flat forecast loss built raw target indices outside the padded target width."
+        )
+
+    # Head projections may run in BF16, but every target transform and loss
+    # reduction is FP32. The cast remains differentiable back to the head.
+    pred_records = pred_quantiles_batched[record_target_idx, record_patch_idx].to(
+        dtype=torch.float32
+    )
+    safe_raw_indices = torch.clamp(raw_indices, max=max_values_len - 1)
+    gt_records = values_padded[record_target_idx.unsqueeze(1), safe_raw_indices]
+    roi_records = roi_masks_padded[record_target_idx.unsqueeze(1), safe_raw_indices]
+    # Each record was predicted in its owner patch's causal-stats frame, so
+    # the targets move into the same configured post-scale value frame.
+    owner_mu = context_mu_padded[record_target_idx, record_patch_idx]
+    owner_sigma = context_sigma_padded[record_target_idx, record_patch_idx].clamp_min(
+        1e-6
+    )
+    gt_records = _transform_forecast_targets_for_loss(
+        gt_records,
+        owner_mu=owner_mu,
+        owner_sigma=owner_sigma,
+    )
+    mask_float = keep_mask.to(dtype=dtype)
+
+    zero = pred_records.new_zeros(())
+    roi_supervised_points, roi_supervised_targets = _roi_supervision_counts(
+        roi_records=roi_records,
+        keep_mask=keep_mask,
+        record_target_idx=record_target_idx,
+        target_count=len(targets),
+    )
+    global_point = zero
+    total_point = zero
+    roi_point_loss = global_point.new_zeros(())
+    roi_weighted_point_loss = global_point.new_zeros(())
+    global_quantile = zero
+    total_quantile = zero
+    global_uncapped = zero
+    global_cap_scale = zero.new_ones(())
+    roi_uncapped = zero
+    roi_weighted_uncapped = zero
+    roi_cap_scale = zero.new_ones(())
+    record_count = int(record_target_idx.numel())
+    point_per_record = pred_records.new_zeros((record_count,))
+    quantile_per_record = pred_records.new_zeros((record_count,))
+
+    taus = torch.tensor(self.tsfm_quantile_taus, device=device, dtype=dtype)
+    if int(taus.numel()) != q_channels:
+        raise RuntimeError(
+            "TS expert quantile tau count must match the forecast channel axis: "
+            f"taus={int(taus.numel())}, q_channels={q_channels}"
+        )
+    quantile_mask = taus >= 0
+    if torch.any(quantile_mask):
+        active_taus = taus[quantile_mask]
+        pred_q = pred_records[:, :, quantile_mask]
+        err = gt_records[:, :, None] - pred_q
+        quantile_by_value = torch.maximum(
+            active_taus.view(1, 1, -1) * err,
+            (active_taus.view(1, 1, -1) - 1.0) * err,
+        )
+        quantile_denominator = valid_counts * int(active_taus.shape[0])
+        quantile_per_record = (quantile_by_value * mask_float[:, :, None]).sum(
+            dim=(1, 2)
+        ) / quantile_denominator
+
+    roi_alpha = float(self.ts_roi_mse_alpha)
+    point_by_value = _point_forecast_loss_values(
+        pred_records[:, :, decode_index],
+        gt_records,
+    )
+    point_per_record = (point_by_value * mask_float).sum(dim=1) / valid_counts
+    global_point = _reduce_per_record_loss(
+        point_per_record,
+        global_count=mot_ts_num_records_in_batch,
+        label="TS point record",
+    )
+
+    # Point and quantile are the same TS objective family. Keep their internal
+    # contribution unweighted here; the public knob is the outer
+    # `mot_ts_loss_weight` applied to the complete `ts_aux_loss`.
+    global_per_record_uncapped = point_per_record + quantile_per_record
+    global_uncapped = _reduce_per_record_loss(
+        global_per_record_uncapped,
+        global_count=mot_ts_num_records_in_batch,
+        label="TS global record",
+    )
+    global_record_cap_scales = _detached_per_item_loss_cap_scale(
+        global_per_record_uncapped,
+        _MOT_TS_GLOBAL_LOSS_DETACHED_CAP,
+    )
+    global_cap_scale = global_record_cap_scales.mean()
+    global_point = _reduce_per_record_loss(
+        point_per_record * global_record_cap_scales,
+        global_count=mot_ts_num_records_in_batch,
+        label="TS capped point record",
+    )
+    global_quantile = _reduce_per_record_loss(
+        quantile_per_record * global_record_cap_scales,
+        global_count=mot_ts_num_records_in_batch,
+        label="TS capped quantile record",
+    )
+    total_point = global_point
+    total_quantile = global_quantile
+
+    if roi_alpha > 0.0:
+        roi_value_losses = _roi_value_losses_from_predictions(
+            pred_records=pred_records,
+            gt_records=gt_records,
+            decode_index=decode_index,
+        )
+        reduced_roi_loss, roi_per_target = _roi_loss_from_value_losses(
+            value_losses=roi_value_losses,
+            roi_records=roi_records,
+            keep_mask=mask_float,
+            record_target_idx=record_target_idx,
+            target_count=len(targets),
+            global_target_count=mot_ts_num_roi_targets_in_batch,
+        )
+        roi_uncapped = reduced_roi_loss
+        if int(roi_per_target.numel()) > 0:
+            roi_target_cap_scales = _detached_per_item_loss_cap_scale(
+                roi_per_target,
+                _MOT_TS_ROI_LOSS_DETACHED_CAP,
+            )
+            roi_loss_for_objective = _reduce_per_record_loss(
+                roi_per_target * roi_target_cap_scales,
+                global_count=mot_ts_num_roi_targets_in_batch,
+                label="capped ROI target",
+            )
+            roi_cap_scale = roi_target_cap_scales.mean()
+        else:
+            roi_loss_for_objective = roi_uncapped
+        roi_weighted_uncapped = total_quantile.new_tensor(roi_alpha) * roi_uncapped
+        roi_point_loss = roi_loss_for_objective
+        roi_weighted_point_loss = total_point.new_tensor(roi_alpha) * roi_point_loss
+        total_point = total_point + roi_weighted_point_loss
+
+    return _FlatForecastLossParts(
+        point=total_point,
+        quantile=total_quantile,
+        global_point=global_point,
+        global_quantile=global_quantile,
+        roi_point=roi_point_loss,
+        roi_weighted_point=roi_weighted_point_loss,
+        global_uncapped=global_uncapped,
+        global_cap_scale=global_cap_scale,
+        roi_uncapped=roi_uncapped,
+        roi_weighted_uncapped=roi_weighted_uncapped,
+        roi_cap_scale=roi_cap_scale,
+        roi_supervised_points=roi_supervised_points,
+        roi_supervised_targets=roi_supervised_targets,
+    )
+
+
+def _predict_mot_span_quantiles(
+    self,
+    *,
+    mot_runtime: "_MoTRuntime",
+    target: "_MoTForecastTarget",
+) -> torch.Tensor:
+    """Run the native numeric head on one packed TS target slice (real space)."""
+    pred_batched, patch_counts = _predict_mot_span_quantiles_batched(
+        self,
+        mot_runtime=mot_runtime,
+        targets=[target],
+        output_space="real",
+    )
+    if len(patch_counts) != 1:
+        raise RuntimeError(
+            f"Single-target forecast expected one patch count, got {patch_counts}."
+        )
+    return pred_batched[0, : int(patch_counts[0])]
+
+
+def _finalize_forecast_routed_losses(
+    self,
+    *,
+    main_losses: _FlatForecastLossParts,
+    extra_losses: Optional[Dict[str, torch.Tensor]] = None,
+) -> Dict[str, torch.Tensor]:
+    """Assemble the routed TS loss dictionary."""
+    total_point = main_losses.point
+    total_quantile = main_losses.quantile
+    total_ts_aux = total_point + total_quantile
+    total_roi_loss = main_losses.roi_weighted_point
+    routed_losses: Dict[str, torch.Tensor] = {
+        "ts_aux_loss": total_ts_aux,
+        "ts_aux_point_loss": total_point,
+        "ts_aux_quantile_loss": total_quantile,
+        "ts_aux_global_point_loss": main_losses.global_point,
+        "ts_aux_global_quantile_loss": main_losses.global_quantile,
+        "ts_aux_global_uncapped_loss": main_losses.global_uncapped,
+        "ts_aux_global_cap_scale": main_losses.global_cap_scale,
+        "ts_aux_roi_loss": total_roi_loss,
+        "ts_aux_roi_point_loss": main_losses.roi_point,
+        "ts_aux_roi_weighted_point_loss": main_losses.roi_weighted_point,
+        "ts_aux_roi_uncapped_loss": main_losses.roi_uncapped,
+        "ts_aux_roi_weighted_uncapped_loss": main_losses.roi_weighted_uncapped,
+        "ts_aux_roi_cap_scale": main_losses.roi_cap_scale,
+        "ts_aux_roi_supervised_points": main_losses.roi_supervised_points,
+        "ts_aux_roi_supervised_targets": main_losses.roi_supervised_targets,
+        "ts_aux_roi_alpha": total_point.new_tensor(float(self.ts_roi_mse_alpha)),
+    }
+
+    if extra_losses:
+        weighted_loss = extra_losses.get("ts_aux_understanding_weighted_loss")
+        if isinstance(weighted_loss, torch.Tensor):
+            total_ts_aux = total_ts_aux + weighted_loss
+            routed_losses["ts_aux_loss"] = total_ts_aux
+        routed_losses.update(extra_losses)
+
+    return routed_losses
+
+
+def _zero_flat_forecast_loss_parts(reference: torch.Tensor) -> _FlatForecastLossParts:
+    zero = reference.new_zeros((), dtype=torch.float32)
+    one = reference.new_ones((), dtype=torch.float32)
+    return _FlatForecastLossParts(
+        point=zero,
+        quantile=zero,
+        global_point=zero,
+        global_quantile=zero,
+        roi_point=zero,
+        roi_weighted_point=zero,
+        global_uncapped=zero,
+        global_cap_scale=one,
+        roi_uncapped=zero,
+        roi_weighted_uncapped=zero,
+        roi_cap_scale=one,
+        roi_supervised_points=zero,
+        roi_supervised_targets=zero,
+    )
+
+
+def _compute_understanding_head_losses(
+    self,
+    *,
+    mot_runtime: "_MoTRuntime",
+    targets: List[Any],
+    mot_ts_understanding_num_points: Optional[object] = None,
+    mot_ts_understanding_num_fft_elements: Optional[object] = None,
+) -> Dict[str, torch.Tensor]:
+    """Reconstruct observed/context TS patches from understanding hidden states."""
+    weight = float(self.ts_understanding_loss_weight)
+    understanding_head = self.understanding_head
+    understanding_hidden = mot_runtime.understanding_hidden
+    reference = understanding_hidden
+    if reference is None and targets:
+        reference = targets[0].values
+    if reference is None:
+        reference = torch.zeros((), dtype=torch.float32)
+    zero = reference.new_zeros((), dtype=torch.float32)
+    count_zero = reference.new_zeros((), dtype=torch.float32)
+    disabled = weight <= 0.0
+    if (
+        disabled
+        or understanding_hidden is None
+        or int(understanding_hidden.shape[0]) == 0
+    ):
+        return {
+            "ts_aux_understanding_loss": zero,
+            "ts_aux_understanding_weighted_loss": zero,
+            "ts_aux_understanding_time_loss": zero,
+            "ts_aux_understanding_fft_loss": zero,
+            "ts_aux_understanding_patches": count_zero,
+            "ts_aux_understanding_points": count_zero,
+            "ts_aux_understanding_weight": count_zero,
+        }
+    if understanding_head is None:
+        raise RuntimeError(
+            "ts_understanding_loss_weight > 0 requires understanding_head."
+        )
+
+    hidden_parts: List[torch.Tensor] = []
+    value_parts: List[torch.Tensor] = []
+    mask_parts: List[torch.Tensor] = []
+    for target in targets:
+        if int(target.role_id) == int(ROLE_TARGET):
+            continue
+        if int(target.ts_route_id) != int(TS_ROUTE_UNDERSTANDING):
+            continue
+        ts_start = int(target.ts_start)
+        ts_end = int(target.ts_end)
+        if ts_end <= ts_start:
+            continue
+        if not (0 <= ts_start < ts_end <= int(understanding_hidden.shape[0])):
+            raise RuntimeError(
+                "TS understanding reconstruction target points outside packed understanding hidden states: "
+                f"sample={int(target.sample_idx)}, "
+                f"slot={int(target.slot_idx)}, "
+                f"ts_range=[{ts_start},{ts_end}), hidden_count={int(understanding_hidden.shape[0])}."
+            )
+        reconstruction_values = target.reconstruction_values
+        reconstruction_masks = target.reconstruction_masks
+        if reconstruction_values is None or reconstruction_masks is None:
+            raise RuntimeError(
+                "TS understanding reconstruction requires patch-aligned values and masks."
+            )
+        if reconstruction_values.ndim != 3 or int(reconstruction_values.shape[0]) != 1:
+            raise RuntimeError(
+                "TS understanding reconstruction values must be shaped [1,N_patch,patch], got "
+                f"{tuple(reconstruction_values.shape)}."
+            )
+        if reconstruction_masks.shape != reconstruction_values.shape:
+            raise RuntimeError(
+                "TS understanding reconstruction masks must match reconstruction values, got "
+                f"values={tuple(reconstruction_values.shape)}, masks={tuple(reconstruction_masks.shape)}."
+            )
+        patch_count = ts_end - ts_start
+        if int(reconstruction_values.shape[1]) != patch_count:
+            raise RuntimeError(
+                "TS understanding reconstruction patch count must match hidden span length, got "
+                f"values={int(reconstruction_values.shape[1])}, hidden={patch_count}."
+            )
+        valid_mask = torch.logical_not(
+            reconstruction_masks[0].to(
+                device=understanding_hidden.device, dtype=torch.bool
+            )
+        )
+        if not bool(torch.any(valid_mask)):
+            continue
+        hidden_parts.append(understanding_hidden[ts_start:ts_end])
+        value_parts.append(
+            reconstruction_values[0].to(
+                device=understanding_hidden.device, dtype=torch.float32
+            )
+        )
+        mask_parts.append(valid_mask)
+
+    if not hidden_parts:
+        return {
+            "ts_aux_understanding_loss": zero,
+            "ts_aux_understanding_weighted_loss": zero,
+            "ts_aux_understanding_time_loss": zero,
+            "ts_aux_understanding_fft_loss": zero,
+            "ts_aux_understanding_patches": count_zero,
+            "ts_aux_understanding_points": count_zero,
+            "ts_aux_understanding_weight": reference.new_tensor(
+                float(weight), dtype=torch.float32
+            ),
+        }
+
+    head_inputs = torch.cat(hidden_parts, dim=0)
+    target_values = torch.cat(value_parts, dim=0)
+    valid_mask = torch.cat(mask_parts, dim=0)
+    pred_values = understanding_head(head_inputs)
+    if pred_values.ndim != 2 or tuple(pred_values.shape) != tuple(target_values.shape):
+        raise RuntimeError(
+            "TS understanding head output must match patch reconstruction targets, got "
+            f"pred={tuple(pred_values.shape)}, target={tuple(target_values.shape)}."
+        )
+
+    pred_f32 = pred_values.to(dtype=torch.float32)
+    target_f32 = target_values.to(dtype=torch.float32)
+    valid_mask_f32 = valid_mask.to(dtype=torch.float32)
+    local_valid_points = valid_mask_f32.sum()
+    global_valid_points = _global_count_tensor(
+        mot_ts_understanding_num_points,
+        reference=pred_f32,
+        label="TS understanding valid-point",
+    )
+    time_denominator = (
+        local_valid_points if global_valid_points is None else global_valid_points
+    )
+    time_loss = (
+        (pred_f32 - target_f32).square() * valid_mask_f32
+    ).sum() / time_denominator
+
+    pred_fft = torch.fft.rfft(pred_f32 * valid_mask_f32, dim=-1, norm="ortho")
+    target_fft = torch.fft.rfft(target_f32 * valid_mask_f32, dim=-1, norm="ortho")
+    fft_squared_error = (pred_fft - target_fft).abs().square()
+    global_fft_elements = _global_count_tensor(
+        mot_ts_understanding_num_fft_elements,
+        reference=pred_f32,
+        label="TS understanding FFT-element",
+    )
+    fft_denominator = (
+        fft_squared_error.new_tensor(float(fft_squared_error.numel()))
+        if global_fft_elements is None
+        else global_fft_elements
+    )
+    fft_loss = fft_squared_error.sum() / fft_denominator
+    loss = 0.2 * time_loss + 0.8 * fft_loss
+    weighted_loss = loss.new_tensor(weight) * loss
+    return {
+        "ts_aux_understanding_loss": loss,
+        "ts_aux_understanding_weighted_loss": weighted_loss,
+        "ts_aux_understanding_time_loss": time_loss,
+        "ts_aux_understanding_fft_loss": fft_loss,
+        "ts_aux_understanding_patches": pred_values.new_tensor(
+            float(int(pred_values.shape[0])), dtype=torch.float32
+        ),
+        "ts_aux_understanding_points": pred_values.new_tensor(
+            float(valid_mask.sum().item()), dtype=torch.float32
+        ),
+        "ts_aux_understanding_weight": pred_values.new_tensor(
+            float(weight), dtype=torch.float32
+        ),
+    }
+
+
+def _sum_flat_loss_parts(parts: List[_FlatForecastLossParts]) -> _FlatForecastLossParts:
+    if not parts:
+        raise RuntimeError("Cannot sum empty TS loss parts.")
+    if len(parts) == 1:
+        return parts[0]
+    values = []
+    for field_idx in range(len(_FlatForecastLossParts._fields)):
+        values.append(sum((part[field_idx] for part in parts[1:]), parts[0][field_idx]))
+    return _FlatForecastLossParts(*values)
+
+
+def compute_routed_mot_losses(
+    self,
+    mot_runtime: "_MoTRuntime",
+    *,
+    mot_ts_num_records_by_horizon: Optional[object] = None,
+    mot_ts_num_roi_targets_by_horizon: Optional[object] = None,
+    mot_ts_understanding_num_points: Optional[object] = None,
+    mot_ts_understanding_num_fft_elements: Optional[object] = None,
+) -> Optional[Dict[str, torch.Tensor]]:
+    """Route TS losses by target type, directly from the packed runtime.
+
+    Contract:
+    - next-patch generation records share one TS record/ROI reducer and the
+      exact GA/DDP denominator bundle
+    - `<ts> -> first patch` and `last patch -> </ts>` boundary routes are removed
+    """
+    if self.generation_tsfm is None:
+        return None
+
+    understanding_head_losses = _compute_understanding_head_losses(
+        self,
+        mot_runtime=mot_runtime,
+        targets=list(mot_runtime.forecast_targets),
+        mot_ts_understanding_num_points=mot_ts_understanding_num_points,
+        mot_ts_understanding_num_fft_elements=mot_ts_understanding_num_fft_elements,
+    )
+    supervised_targets_all = [
+        target for target in mot_runtime.forecast_targets if target.supervise
+    ]
+    if not supervised_targets_all:
+        understanding_weighted = understanding_head_losses.get(
+            "ts_aux_understanding_weighted_loss"
+        )
+        if isinstance(understanding_weighted, torch.Tensor) and bool(
+            torch.any(understanding_weighted != 0)
+        ):
+            return _finalize_forecast_routed_losses(
+                self,
+                main_losses=_zero_flat_forecast_loss_parts(understanding_weighted),
+                extra_losses=understanding_head_losses,
+            )
+        return None
+    # Keep conditioning spans out of next-patch loss even for hand-built
+    # runtimes; target spans must use the generation route checked below.
+    generation_next_patch_targets = [
+        target
+        for target in supervised_targets_all
+        if int(target.role_id) == int(ROLE_TARGET)
+        and int(target.ts_route_id) == int(TS_ROUTE_GENERATION)
+    ]
+    bad_next_patch_target_routes = [
+        (
+            int(target.sample_idx),
+            int(target.slot_idx),
+            int(target.ts_route_id),
+        )
+        for target in supervised_targets_all
+        if int(target.role_id) == int(ROLE_TARGET)
+        and int(target.ts_route_id) != int(TS_ROUTE_GENERATION)
+    ]
+    if bad_next_patch_target_routes:
+        raise RuntimeError(
+            "assistant forecasting target spans must route through the generation TS expert; "
+            f"bad_sample_slot_route={bad_next_patch_target_routes}."
+        )
+
+    loss_parts: List[_FlatForecastLossParts] = []
+    extra_losses: Dict[str, torch.Tensor] = dict(understanding_head_losses)
+    routed_targets = generation_next_patch_targets
+    if routed_targets:
+        pred_quantiles_batched, patch_counts = _predict_mot_span_quantiles_batched(
+            self,
+            mot_runtime=mot_runtime,
+            targets=routed_targets,
+            output_space="normalized",
+        )
+        routed_losses = _compute_flat_forecast_losses_from_predictions(
+            self,
+            pred_quantiles_batched=pred_quantiles_batched,
+            targets=routed_targets,
+            patch_counts=patch_counts,
+            mot_ts_num_records_in_batch=_global_count_for_horizon(
+                mot_ts_num_records_by_horizon,
+                int(self.tsfm_output_patch_len),
+            ),
+            mot_ts_num_roi_targets_in_batch=_global_count_for_horizon(
+                mot_ts_num_roi_targets_by_horizon,
+                int(self.tsfm_output_patch_len),
+            ),
+            loss_space="normalized",
+        )
+        if routed_losses is not None:
+            loss_parts.append(routed_losses)
+
+    if not loss_parts:
+        understanding_weighted = extra_losses.get("ts_aux_understanding_weighted_loss")
+        if isinstance(understanding_weighted, torch.Tensor) and bool(
+            torch.any(understanding_weighted != 0)
+        ):
+            return _finalize_forecast_routed_losses(
+                self,
+                main_losses=_zero_flat_forecast_loss_parts(understanding_weighted),
+                extra_losses=extra_losses,
+            )
+        return None
+    flat_losses = _sum_flat_loss_parts(loss_parts)
+
+    return _finalize_forecast_routed_losses(
+        self,
+        main_losses=flat_losses,
+        extra_losses=extra_losses,
+    )
+
+
+__all__ = [
+    "compute_routed_mot_losses",
+    "_compute_flat_forecast_losses_from_predictions",
+    "_compute_understanding_head_losses",
+]

timebraid/model/mot/config_contract.py

@@ -0,0 +1,186 @@
+"""TimeBraid Hugging Face config restore contract."""
+
+from __future__ import annotations
+
+from collections.abc import Mapping
+from dataclasses import dataclass
+from typing import Any
+
+
+class _Required:
+    """Sentinel for an option the restore contract always supplies."""
+
+    def __repr__(self) -> str:  # pragma: no cover - debugging aid
+        return "REQUIRED"
+
+
+REQUIRED = _Required()
+
+
+@dataclass(frozen=True, slots=True)
+class MoTOption:
+    """One MoT construction knob, declared once for every boundary it crosses.
+
+    The same `name` is used in the checkpoint's `config.json`, in the runtime
+    option mapping, and in the error messages — so grepping one spelling finds
+    every place the knob exists.
+    """
+
+    name: str
+    kind: str  # boolean | integer | number | string | torch_dtype
+    default: Any  # REQUIRED when the restore contract always supplies it
+    persisted: bool  # does it live in the checkpoint's config.json
+
+    @property
+    def is_required(self) -> bool:
+        return self.default is REQUIRED
+
+
+# The single declaration of the MoT option surface.
+#
+# Persisted options are REQUIRED rather than defaulted on purpose. The restore
+# contract always supplies them, so a default could only ever apply to a
+# hand-built option map — and for something like the position lattice that
+# means plausible-but-wrong forecasts instead of an error. Defaults belong
+# only to options no checkpoint carries.
+MOT_OPTIONS: tuple[MoTOption, ...] = (
+    # --- persisted: module topology and tokenizer ABI ---
+    MoTOption("mot_pairing_mode", "string", REQUIRED, persisted=True),
+    MoTOption("mot_t_layers", "integer", REQUIRED, persisted=True),
+    MoTOption("mot_understanding_pair_depth", "integer", REQUIRED, persisted=True),
+    # The one contract field without the `mot_` prefix. Renaming it would break
+    # every checkpoint the training tree has written, so it is left alone until
+    # both trees move together.
+    MoTOption("has_understanding_head", "boolean", REQUIRED, persisted=True),
+    MoTOption("mot_mixed_position_mode", "string", REQUIRED, persisted=True),
+    MoTOption("mot_tsfm_num_heads", "integer", REQUIRED, persisted=True),
+    MoTOption("mot_tsfm_head_dim", "integer", REQUIRED, persisted=True),
+    MoTOption("mot_tsfm_hidden_size", "integer", REQUIRED, persisted=True),
+    MoTOption("mot_patch_size", "integer", REQUIRED, persisted=True),
+    MoTOption("mot_ts_open_token_id", "integer", REQUIRED, persisted=True),
+    MoTOption("mot_ts_close_token_id", "integer", REQUIRED, persisted=True),
+    # --- not persisted: supplied per load, or not at all ---
+    MoTOption("mot_compute_dtype", "torch_dtype", "float32", persisted=False),
+    MoTOption("mot_max_spans_per_sample", "integer", 64, persisted=False),
+    MoTOption("mot_ts_gradient_checkpointing", "boolean", False, persisted=False),
+    MoTOption(
+        "mot_ts_gradient_checkpointing_use_reentrant",
+        "boolean",
+        False,
+        persisted=False,
+    ),
+    # No loader supplies these four; they exist because the loss path is still
+    # reachable from the public `runtime_options` argument. See the release
+    # audit's "training machinery on the inference path".
+    MoTOption("mot_ts_loss_weight", "number", 1.0, persisted=False),
+    MoTOption("mot_lm_loss_weight", "number", 1.0, persisted=False),
+    MoTOption("mot_ts_roi_mse_alpha", "number", 0.0, persisted=False),
+    MoTOption("mot_ts_understanding_loss_weight", "number", 0.0, persisted=False),
+)
+
+MOT_OPTION_BY_NAME: dict[str, MoTOption] = {
+    option.name: option for option in MOT_OPTIONS
+}
+
+# Derived, not restated: the persisted subset is a property of the registry.
+MOT_HF_CONFIG_CONTRACT_FIELDS = tuple(
+    option.name for option in MOT_OPTIONS if option.persisted
+)
+
+
+def collect_mot_hf_config_contract(config: Any) -> dict[str, object]:
+    """Collect the complete MoT HF restore contract from a config object."""
+    is_mapping = isinstance(config, Mapping)
+
+    def _has(field: str) -> bool:
+        return field in config if is_mapping else hasattr(config, field)
+
+    def _get(field: str) -> object:
+        return config.get(field) if is_mapping else getattr(config, field, None)
+
+    missing = [field for field in MOT_HF_CONFIG_CONTRACT_FIELDS if not _has(field)]
+    if missing:
+        raise RuntimeError(f"MoT HF config restore metadata is missing {missing}.")
+
+    contract: dict[str, object] = {}
+    for field in MOT_HF_CONFIG_CONTRACT_FIELDS:
+        contract[field] = _get(field)
+    return validate_mot_hf_config_contract(contract)
+
+
+def validate_mot_hf_config_contract(
+    contract: dict[str, object],
+) -> dict[str, object]:
+    """Validate and canonicalize every persisted MoT restore field."""
+    missing = [
+        field for field in MOT_HF_CONFIG_CONTRACT_FIELDS if field not in contract
+    ]
+    if missing:
+        raise RuntimeError(f"MoT HF config restore metadata is missing {missing}.")
+
+    validated = dict(contract)
+    allowed_strings = {
+        "mot_pairing_mode": {"interleaved", "dsfp"},
+        "mot_mixed_position_mode": {"span_slot", "patch_slot"},
+    }
+    for field, allowed in allowed_strings.items():
+        value = validated[field]
+        if not isinstance(value, str) or value not in allowed:
+            raise RuntimeError(
+                f"MoT HF config `{field}` must be one of {sorted(allowed)!r}, got {value!r}."
+            )
+
+    positive_integer_fields = (
+        "mot_t_layers",
+        "mot_tsfm_num_heads",
+        "mot_tsfm_head_dim",
+        "mot_tsfm_hidden_size",
+        "mot_patch_size",
+    )
+    for field in positive_integer_fields:
+        value = validated[field]
+        if type(value) is not int or value <= 0:
+            raise RuntimeError(
+                f"MoT HF config `{field}` must be a positive integer, got {value!r}."
+            )
+    pair_depth = validated["mot_understanding_pair_depth"]
+    if type(pair_depth) is not int or not 0 <= pair_depth <= validated["mot_t_layers"]:
+        raise RuntimeError(
+            "MoT HF config `mot_understanding_pair_depth` must be an integer in "
+            f"[0, mot_t_layers], got {pair_depth!r}."
+        )
+    has_understanding_head = validated["has_understanding_head"]
+    if type(has_understanding_head) is not bool:
+        raise RuntimeError(
+            "MoT HF config `has_understanding_head` must be a bool, got "
+            f"{has_understanding_head!r}."
+        )
+    if has_understanding_head and pair_depth == 0:
+        raise RuntimeError(
+            "MoT HF config cannot enable `has_understanding_head` without an "
+            "understanding TimesFM tower."
+        )
+    expected_hidden_size = (
+        validated["mot_tsfm_num_heads"] * validated["mot_tsfm_head_dim"]
+    )
+    if validated["mot_tsfm_hidden_size"] != expected_hidden_size:
+        raise RuntimeError(
+            "MoT HF config TSFM geometry is inconsistent: "
+            f"hidden_size={validated['mot_tsfm_hidden_size']}, "
+            f"num_heads*head_dim={expected_hidden_size}."
+        )
+
+    delimiter_ids = []
+    for field in ("mot_ts_open_token_id", "mot_ts_close_token_id"):
+        value = validated[field]
+        if type(value) is not int or value < 0:
+            raise RuntimeError(
+                f"MoT HF config `{field}` must be a non-negative integer, got {value!r}."
+            )
+        delimiter_ids.append(value)
+    if delimiter_ids[0] == delimiter_ids[1]:
+        raise RuntimeError(
+            "MoT HF config TS delimiter token IDs must be distinct, got "
+            f"{delimiter_ids[0]}."
+        )
+    return validated

timebraid/model/mot/generation_route.py

@@ -0,0 +1,223 @@
+"""Explicit routing decision for TimeBraid generation.
+
+`TimeBraid.generate` accepts ordinary Hugging Face keyword arguments and has to
+decide, from the inputs alone, whether a request wants plain text from the host
+language model or the mixed text/time-series scheduler. That decision used to
+live inline in a five-condition lattice with no name and no way to report
+itself, so a request that fell through to plain text produced text with no
+diagnostic. Everything needed to make that decision is gathered here instead:
+the route is a named value, the reason is a sentence, and both are reachable
+from a test without constructing a model.
+"""
+
+from __future__ import annotations
+
+from collections.abc import Mapping, Sequence
+from dataclasses import dataclass
+from enum import Enum
+from typing import Optional
+
+import torch
+
+from .structures import TIMEBRAID_REQUIRED_TS_PAYLOAD_FIELDS, TimeBraidPayload
+
+
+class FinishReason(str, Enum):
+    """Why the mixed scheduler stopped producing tokens.
+
+    These two spellings used to be bare string literals at five construction
+    sites, with the processor raising on anything it did not recognise — so a
+    third reason would have broken decoding at a distance rather than at the
+    call site that introduced it.
+    """
+
+    TEXT_BUDGET = "text_budget"
+    EOS_OR_PROTOCOL_STOP = "eos_or_protocol_stop"
+
+    # Enum overrides str.__str__, so without this `str(member)` and any
+    # f-string yield 'ClassName.MEMBER' instead of the value — which is
+    # exactly what a JSON payload or a log line would then carry.
+    __str__ = str.__str__
+
+
+class GenerationRoute(str, Enum):
+    """Which generation implementation a request resolves to.
+
+    Inherits from `str` so the value survives JSON serialization and compares
+    equal to its own name in rollout records and logs.
+    """
+
+    PLAIN_TEXT = "plain_text"
+    MIXED_TIMESERIES = "mixed_timeseries"
+
+    # Enum overrides str.__str__, so without this `str(member)` and any
+    # f-string yield 'ClassName.MEMBER' instead of the value — which is
+    # exactly what a JSON payload or a log line would then carry.
+    __str__ = str.__str__
+
+
+@dataclass(frozen=True, slots=True)
+class RouteDecision:
+    """A resolved route together with the reason it was chosen."""
+
+    route: GenerationRoute
+    reason: str
+
+    @property
+    def is_mixed(self) -> bool:
+        return self.route is GenerationRoute.MIXED_TIMESERIES
+
+
+def _contains_ts_delimiter(
+    route_input: object, *, ts_open_token_id: int, ts_close_token_id: int
+) -> bool:
+    """Scan a token tensor for either TS delimiter id."""
+    if not isinstance(route_input, torch.Tensor):
+        return False
+    return bool(
+        torch.logical_or(
+            route_input.eq(int(ts_open_token_id)),
+            route_input.eq(int(ts_close_token_id)),
+        )
+        .any()
+        .item()
+    )
+
+
+def resolve_generation_route(
+    *,
+    payload: TimeBraidPayload,
+    route_input: object,
+    target_horizons: Optional[Sequence[int]],
+    target_total_lengths: Optional[Sequence[int]],
+    ts_open_token_id: int,
+    ts_close_token_id: int,
+) -> RouteDecision:
+    """Decide which generation route a request takes.
+
+    Three independent signals request the mixed scheduler: a tensor-valued TS
+    payload, a `<ts>`/`</ts>` delimiter in the prompt, or explicit
+    `target_horizons`. Absent all three, the request is plain text.
+
+    One tolerance is preserved from the original lattice: an explicit all-zero
+    `target_horizons` on a text-only batch still produces text, because
+    external evaluation harnesses pass `target_horizons=0` to mean "forbid TS
+    rollout" rather than "run the TS scheduler with nothing to forecast".
+
+    `target_total_lengths` is deliberately asymmetric against that tolerance:
+    all-zero lengths are an explicit request for the scheduler's empty-target
+    validation, so they defeat it. An absent field is *not* all-zero.
+
+    This is only consulted by the `generate` compatibility shim, which is the
+    only entry point that has to infer intent from loose keyword arguments.
+    `generate_mixed` states its intent in its signature and validates its
+    inputs directly.
+    """
+    has_payload = payload.has_runtime_inputs()
+
+    horizons_given = target_horizons is not None
+    horizons_all_zero = horizons_given and all(
+        horizon == 0 for horizon in target_horizons
+    )
+    lengths_all_zero = target_total_lengths is not None and all(
+        total == 0 for total in target_total_lengths
+    )
+
+    # A positive horizon, or an explicit empty-target request, goes straight to
+    # the scheduler; the delimiter scan below is then pure overhead.
+    horizons_force_mixed = horizons_given and (
+        not horizons_all_zero or lengths_all_zero
+    )
+
+    has_delimiter = False
+    if not has_payload and not horizons_force_mixed:
+        has_delimiter = _contains_ts_delimiter(
+            route_input,
+            ts_open_token_id=ts_open_token_id,
+            ts_close_token_id=ts_close_token_id,
+        )
+
+    horizons_request_mixed = horizons_given
+    if (
+        horizons_request_mixed
+        and horizons_all_zero
+        and not lengths_all_zero
+        and not has_payload
+        and isinstance(route_input, torch.Tensor)
+        and not has_delimiter
+    ):
+        horizons_request_mixed = False
+
+    if has_payload:
+        return RouteDecision(
+            GenerationRoute.MIXED_TIMESERIES,
+            "a tensor-valued TS payload was supplied",
+        )
+    if has_delimiter:
+        return RouteDecision(
+            GenerationRoute.MIXED_TIMESERIES,
+            "the prompt contains a TS delimiter token",
+        )
+    if horizons_request_mixed:
+        return RouteDecision(
+            GenerationRoute.MIXED_TIMESERIES,
+            "explicit target_horizons were supplied",
+        )
+
+    absent = []
+    if not has_payload:
+        absent.append("no TS payload")
+    if not has_delimiter:
+        absent.append("no TS delimiter token in the prompt")
+    if not horizons_given:
+        absent.append("no target_horizons")
+    elif horizons_all_zero:
+        absent.append("target_horizons are all zero")
+    return RouteDecision(
+        GenerationRoute.PLAIN_TEXT,
+        "plain text because " + ", ".join(absent),
+    )
+
+
+def require_mixed_generation_inputs(inputs: Mapping[str, object]) -> None:
+    """Reject a mapping that cannot drive mixed generation.
+
+    Validated structurally rather than by type so a test can pass a plain dict
+    and so a caller that assembles its own batch is not forced through the
+    processor. The error names the missing fields and where they come from,
+    because their absence used to route silently to plain text.
+    """
+    if not isinstance(inputs, Mapping):
+        raise TypeError(
+            "generate_mixed requires the mapping returned by "
+            "TimeBraidProcessor.apply_chat_template (or an equivalent mapping), "
+            f"got {type(inputs).__name__}."
+        )
+
+    if "input_ids" not in inputs or inputs["input_ids"] is None:
+        raise ValueError(
+            "generate_mixed requires `input_ids`; "
+            "TimeBraidProcessor.apply_chat_template supplies it."
+        )
+
+    missing = [
+        name
+        for name in TIMEBRAID_REQUIRED_TS_PAYLOAD_FIELDS
+        if inputs.get(name) is None
+    ]
+    if missing:
+        raise ValueError(
+            "generate_mixed requires a complete TS payload; missing "
+            f"{missing}. Pass `timeseries=` to "
+            "TimeBraidProcessor.apply_chat_template, which emits every TS "
+            "payload field, or call `generate` for plain text."
+        )
+
+
+__all__ = [
+    "FinishReason",
+    "GenerationRoute",
+    "RouteDecision",
+    "require_mixed_generation_inputs",
+    "resolve_generation_route",
+]

timebraid/model/mot/kv_cache.py

@@ -0,0 +1,1614 @@
+"""Cache state and validation for TimeBraid generation."""
+
+from __future__ import annotations
+
+import hashlib
+from dataclasses import dataclass, fields, is_dataclass
+from typing import Mapping, Sequence
+
+import torch
+from transformers.cache_utils import DynamicCache, DynamicLayer, StaticCache
+
+_RequestId = str | int
+_SUPPORTED_CACHE_IMPLEMENTATIONS = {None, "dynamic", "dynamic_full"}
+
+# One incremental KV stream registry key: (kind, row, layer, span).
+# kind="llm"/"residual" streams live for the whole request row (span=None,
+# layer is the Qwen layer index). kind="ts" streams live for one generated
+# span (span is the target slot, layer is the TimesFM layer index).
+StreamKey = tuple[str, int, int, int | None]
+_STREAM_KINDS = ("llm", "residual", "ts")
+_STREAM_LIFECYCLE = {"llm": "row", "residual": "row", "ts": "span"}
+
+
+def _validate_stream_key(key: object) -> StreamKey:
+    if not isinstance(key, tuple) or len(key) != 4:
+        raise TypeError(
+            f"Stream key must be a (kind, row, layer, span) tuple, got {key!r}."
+        )
+    kind, row, layer, span = key
+    if kind not in _STREAM_KINDS:
+        raise ValueError(f"Stream kind must be one of {_STREAM_KINDS}, got {kind!r}.")
+    if isinstance(row, bool) or not isinstance(row, int) or row < 0:
+        raise ValueError(f"Stream row must be a non-negative int, got {row!r}.")
+    if isinstance(layer, bool) or not isinstance(layer, int) or layer < 0:
+        raise ValueError(f"Stream layer must be a non-negative int, got {layer!r}.")
+    if _STREAM_LIFECYCLE[kind] == "row":
+        if span is not None:
+            raise ValueError(
+                f"Row-lifetime {kind!r} streams must use span=None, got {span!r}."
+            )
+    elif isinstance(span, bool) or not isinstance(span, int) or span < 0:
+        raise ValueError(
+            f"Span-lifetime {kind!r} streams need a non-negative span slot, got {span!r}."
+        )
+    return key
+
+
+def _validate_request_id(request_id: object) -> _RequestId:
+    if isinstance(request_id, bool) or not isinstance(request_id, (str, int)):
+        raise TypeError(
+            f"request_id must be str or int, got {type(request_id).__name__}."
+        )
+    if isinstance(request_id, str) and not request_id:
+        raise ValueError("request_id must not be empty.")
+    return request_id
+
+
+def _normalize_request_ids(request_ids: Sequence[_RequestId]) -> tuple[_RequestId, ...]:
+    normalized = tuple(_validate_request_id(request_id) for request_id in request_ids)
+    if not normalized:
+        raise ValueError("MoTDynamicCache requires at least one request_id.")
+    if len(set(normalized)) != len(normalized):
+        raise ValueError(f"Initial request_ids must be unique, got {normalized}.")
+    return normalized
+
+
+def _update_payload_hash(digest: "hashlib._Hash", value: object) -> None:
+    if isinstance(value, torch.Tensor):
+        try:
+            version = int(value._version)
+        except RuntimeError as exc:
+            raise RuntimeError(
+                "Payload tensors must be created outside torch.inference_mode so in-place mutation is observable."
+            ) from exc
+        tensor_state = (
+            "tensor",
+            str(value.dtype),
+            str(value.device),
+            tuple(value.shape),
+            tuple(value.stride()),
+            int(value.storage_offset()),
+            int(value.untyped_storage().data_ptr()) if value.numel() > 0 else 0,
+            version,
+        )
+        digest.update(repr(tensor_state).encode("utf-8"))
+        return
+
+    if value is None or isinstance(value, (bool, int, float, str)):
+        digest.update(repr((type(value).__name__, value)).encode("utf-8"))
+        return
+
+    if isinstance(value, Mapping):
+        digest.update(b"mapping{")
+        for key in sorted(value):
+            if not isinstance(key, str):
+                raise TypeError(
+                    f"Payload mapping keys must be strings, got {type(key).__name__}."
+                )
+            digest.update(key.encode("utf-8"))
+            _update_payload_hash(digest, value[key])
+        digest.update(b"}")
+        return
+
+    if isinstance(value, (list, tuple)):
+        digest.update(f"{type(value).__name__}[".encode("utf-8"))
+        for item in value:
+            _update_payload_hash(digest, item)
+        digest.update(b"]")
+        return
+
+    if is_dataclass(value) and not isinstance(value, type):
+        digest.update(f"dataclass:{type(value).__qualname__}{{".encode("utf-8"))
+        for field in fields(value):
+            digest.update(field.name.encode("utf-8"))
+            try:
+                _update_payload_hash(digest, getattr(value, field.name))
+            except RuntimeError as exc:
+                raise RuntimeError(
+                    f"Payload field `{field.name}` is not fingerprintable: {exc}"
+                ) from exc
+        digest.update(b"}")
+        return
+
+    raise TypeError(f"Unsupported payload fingerprint value: {type(value).__name__}.")
+
+
+def fingerprint_payload_state(payload: object) -> str:
+    """Fingerprint tensor identity/version metadata without copying payload values."""
+    digest = hashlib.sha256()
+    _update_payload_hash(digest, payload)
+    return digest.hexdigest()
+
+
+def validate_mot_cache_inputs(
+    *,
+    use_cache: bool | None,
+    past_key_values: object | None,
+    cache_implementation: str | None = None,
+) -> None:
+    """Validate the supported MoT cache surface before any cache can mutate."""
+    if cache_implementation not in _SUPPORTED_CACHE_IMPLEMENTATIONS:
+        if cache_implementation == "static" or isinstance(past_key_values, StaticCache):
+            raise RuntimeError(
+                "MoT KV cache does not support StaticCache: HF static generation replaces the required 2D "
+                "layout mask with a compiled 4D mask. Use the dynamic cache implementation."
+            )
+        raise RuntimeError(
+            "MoT KV cache supports cache_implementation in {None, 'dynamic', 'dynamic_full'}, got "
+            f"{cache_implementation!r}."
+        )
+
+    if isinstance(past_key_values, StaticCache):
+        raise RuntimeError(
+            "MoT KV cache does not support StaticCache; use MoTDynamicCache."
+        )
+    if past_key_values is not None and use_cache is not True:
+        raise RuntimeError(
+            "past_key_values requires use_cache=True because Qwen mutates a supplied cache even when use_cache=False."
+        )
+    if use_cache is False and cache_implementation is not None:
+        raise RuntimeError("cache_implementation must be unset when use_cache=False.")
+    if past_key_values is not None and not isinstance(past_key_values, MoTDynamicCache):
+        raise TypeError(
+            "MoT continuation requires MoTDynamicCache with runtime sidecars, got "
+            f"{type(past_key_values).__name__}."
+        )
+
+
+def _require_bool_mask(
+    mask: torch.Tensor, *, name: str, shape: tuple[int, ...]
+) -> torch.Tensor:
+    if not isinstance(mask, torch.Tensor):
+        raise TypeError(f"{name} must be a tensor, got {type(mask).__name__}.")
+    if mask.dtype is not torch.bool:
+        raise TypeError(f"{name} must have dtype torch.bool, got {mask.dtype}.")
+    if tuple(mask.shape) != shape:
+        raise ValueError(f"{name} must have shape {shape}, got {tuple(mask.shape)}.")
+    return mask
+
+
+def _require_positions(
+    positions: torch.Tensor,
+    *,
+    name: str,
+    shape: tuple[int, ...],
+    device: torch.device,
+) -> torch.Tensor:
+    if not isinstance(positions, torch.Tensor):
+        raise TypeError(f"{name} must be a tensor, got {type(positions).__name__}.")
+    if positions.dtype is not torch.long:
+        raise TypeError(f"{name} must have dtype torch.long, got {positions.dtype}.")
+    if tuple(positions.shape) != shape:
+        raise ValueError(
+            f"{name} must have shape {shape}, got {tuple(positions.shape)}."
+        )
+    if positions.device != device:
+        raise ValueError(f"{name} must be on {device}, got {positions.device}.")
+    return positions
+
+
+def _expanded_positions(
+    positions: torch.Tensor,
+    *,
+    batch_size: int,
+    token_count: int,
+    name: str,
+    device: torch.device,
+) -> torch.Tensor:
+    if not isinstance(positions, torch.Tensor):
+        raise TypeError(f"{name} must be a tensor, got {type(positions).__name__}.")
+    if positions.dtype is not torch.long:
+        raise TypeError(f"{name} must have dtype torch.long, got {positions.dtype}.")
+    if positions.device != device:
+        raise ValueError(f"{name} must be on {device}, got {positions.device}.")
+    if tuple(positions.shape) == (token_count,):
+        return positions.view(1, token_count).expand(batch_size, token_count)
+    if tuple(positions.shape) == (batch_size, token_count):
+        return positions
+    raise ValueError(
+        f"{name} must have shape {(token_count,)} or {(batch_size, token_count)}, got {tuple(positions.shape)}."
+    )
+
+
+def _geometric_capacity(current: int, required: int, initial: int) -> int:
+    if required <= current:
+        return current
+    capacity = max(1, current, initial)
+    while capacity < required:
+        capacity *= 2
+    return capacity
+
+
+@dataclass(frozen=True)
+class MoTSidecarView:
+    """Logical view over one named per-layer sidecar stream."""
+
+    keys: torch.Tensor
+    values: torch.Tensor
+    valid_mask: torch.Tensor
+    physical_positions: torch.Tensor
+    logical_positions: torch.Tensor
+    rope_positions: torch.Tensor
+    capacity: int
+
+
+class _PositionBuffer:
+    def __init__(self, *, batch_size: int, initial_capacity: int) -> None:
+        self.batch_size = batch_size
+        self.initial_capacity = initial_capacity
+        self.length = 0
+        self.capacity = 0
+        self.valid_mask: torch.Tensor | None = None
+        self.logical_positions: torch.Tensor | None = None
+        self.rope_positions: torch.Tensor | None = None
+
+    @property
+    def device(self) -> torch.device | None:
+        return None if self.valid_mask is None else self.valid_mask.device
+
+    def _ensure_capacity(self, required: int, *, device: torch.device) -> None:
+        target = _geometric_capacity(self.capacity, required, self.initial_capacity)
+        if target == self.capacity:
+            if self.device != device:
+                raise ValueError(
+                    f"Position metadata is on {self.device}, got append on {device}."
+                )
+            return
+
+        valid = torch.zeros((self.batch_size, target), dtype=torch.bool, device=device)
+        logical = torch.full(
+            (self.batch_size, target), -1, dtype=torch.long, device=device
+        )
+        rope = torch.full(
+            (self.batch_size, target), -1, dtype=torch.long, device=device
+        )
+        if self.length > 0:
+            if self.device != device:
+                raise ValueError(
+                    f"Position metadata is on {self.device}, got append on {device}."
+                )
+            valid[:, : self.length].copy_(self.valid_mask[:, : self.length])
+            logical[:, : self.length].copy_(self.logical_positions[:, : self.length])
+            rope[:, : self.length].copy_(self.rope_positions[:, : self.length])
+        self.valid_mask = valid
+        self.logical_positions = logical
+        self.rope_positions = rope
+        self.capacity = target
+
+    def append(
+        self,
+        *,
+        valid_mask: torch.Tensor,
+        logical_positions: torch.Tensor,
+        rope_positions: torch.Tensor,
+    ) -> None:
+        token_count = int(valid_mask.shape[1])
+        required = self.length + token_count
+        self._ensure_capacity(required, device=valid_mask.device)
+        start, end = self.length, required
+        self.valid_mask[:, start:end].copy_(valid_mask)
+        self.logical_positions[:, start:end].copy_(logical_positions)
+        self.rope_positions[:, start:end].copy_(rope_positions)
+        invalid = ~valid_mask
+        self.logical_positions[:, start:end].masked_fill_(invalid, -1)
+        self.rope_positions[:, start:end].masked_fill_(invalid, -1)
+        self.length = required
+
+    def crop(self, target_length: int) -> None:
+        self.length = min(self.length, target_length)
+        if self.capacity == 0:
+            return
+        self.valid_mask[:, self.length :].zero_()
+        self.logical_positions[:, self.length :].fill_(-1)
+        self.rope_positions[:, self.length :].fill_(-1)
+
+    def reset(self) -> None:
+        self.length = 0
+        if self.capacity == 0:
+            return
+        self.valid_mask.zero_()
+        self.logical_positions.fill_(-1)
+        self.rope_positions.fill_(-1)
+
+    def index_select(self, indices: torch.LongTensor) -> None:
+        if self.capacity == 0:
+            self.batch_size = int(indices.numel())
+            return
+        local_indices = indices.to(device=self.valid_mask.device)
+        self.valid_mask = self.valid_mask.index_select(0, local_indices)
+        self.logical_positions = self.logical_positions.index_select(0, local_indices)
+        self.rope_positions = self.rope_positions.index_select(0, local_indices)
+        self.batch_size = int(indices.numel())
+
+    def repeat_interleave(self, repeats: int) -> None:
+        if self.capacity > 0:
+            self.valid_mask = self.valid_mask.repeat_interleave(repeats, dim=0)
+            self.logical_positions = self.logical_positions.repeat_interleave(
+                repeats, dim=0
+            )
+            self.rope_positions = self.rope_positions.repeat_interleave(repeats, dim=0)
+        self.batch_size *= repeats
+
+
+class _SidecarBuffer:
+    # Sidecar K/V uses packed-runtime order [B, capacity, H, D]. The native HF
+    # cache remains compact Qwen order [B, H_kv, T, D].
+    def __init__(self, *, batch_size: int, initial_capacity: int) -> None:
+        self.batch_size = batch_size
+        self.initial_capacity = initial_capacity
+        self.length = 0
+        self.capacity = 0
+        self.keys: torch.Tensor | None = None
+        self.values: torch.Tensor | None = None
+        self.valid_mask: torch.Tensor | None = None
+        self.physical_positions: torch.Tensor | None = None
+        self.logical_positions: torch.Tensor | None = None
+        self.rope_positions: torch.Tensor | None = None
+
+    def _ensure_capacity(self, required: int, *, reference: torch.Tensor) -> None:
+        target = _geometric_capacity(self.capacity, required, self.initial_capacity)
+        if target == self.capacity:
+            if (
+                self.keys.device != reference.device
+                or self.keys.dtype != reference.dtype
+            ):
+                raise ValueError(
+                    "Sidecar append dtype/device changed: "
+                    f"expected {self.keys.dtype} on {self.keys.device}, got {reference.dtype} on {reference.device}."
+                )
+            return
+
+        _, _, num_heads, head_dim = reference.shape
+        keys = torch.zeros(
+            (self.batch_size, target, num_heads, head_dim),
+            dtype=reference.dtype,
+            device=reference.device,
+        )
+        values = torch.zeros_like(keys)
+        valid = torch.zeros(
+            (self.batch_size, target), dtype=torch.bool, device=reference.device
+        )
+        physical = torch.full(
+            (self.batch_size, target), -1, dtype=torch.long, device=reference.device
+        )
+        logical = torch.full_like(physical, -1)
+        rope = torch.full_like(physical, -1)
+        if self.length > 0:
+            if (
+                self.keys.device != reference.device
+                or self.keys.dtype != reference.dtype
+            ):
+                raise ValueError(
+                    "Sidecar append dtype/device changed: "
+                    f"expected {self.keys.dtype} on {self.keys.device}, got {reference.dtype} on {reference.device}."
+                )
+            if tuple(self.keys.shape[2:]) != (num_heads, head_dim):
+                raise ValueError(
+                    "Sidecar head shape changed: "
+                    f"expected {tuple(self.keys.shape[2:])}, got {(num_heads, head_dim)}."
+                )
+            keys[:, : self.length].copy_(self.keys[:, : self.length])
+            values[:, : self.length].copy_(self.values[:, : self.length])
+            valid[:, : self.length].copy_(self.valid_mask[:, : self.length])
+            physical[:, : self.length].copy_(self.physical_positions[:, : self.length])
+            logical[:, : self.length].copy_(self.logical_positions[:, : self.length])
+            rope[:, : self.length].copy_(self.rope_positions[:, : self.length])
+        self.keys = keys
+        self.values = values
+        self.valid_mask = valid
+        self.physical_positions = physical
+        self.logical_positions = logical
+        self.rope_positions = rope
+        self.capacity = target
+
+    def append(
+        self,
+        *,
+        keys: torch.Tensor,
+        values: torch.Tensor,
+        valid_mask: torch.Tensor,
+        physical_positions: torch.Tensor,
+        logical_positions: torch.Tensor,
+        rope_positions: torch.Tensor,
+    ) -> None:
+        required = self.length + int(keys.shape[1])
+        self._ensure_capacity(required, reference=keys)
+        if tuple(self.keys.shape[2:]) != tuple(keys.shape[2:]):
+            raise ValueError(
+                f"Sidecar head shape changed: expected {tuple(self.keys.shape[2:])}, got {tuple(keys.shape[2:])}."
+            )
+        start, end = self.length, required
+        self.keys[:, start:end].copy_(keys)
+        self.values[:, start:end].copy_(values)
+        self.valid_mask[:, start:end].copy_(valid_mask)
+        self.physical_positions[:, start:end].copy_(physical_positions)
+        self.logical_positions[:, start:end].copy_(logical_positions)
+        self.rope_positions[:, start:end].copy_(rope_positions)
+        invalid = ~valid_mask
+        expanded_invalid = invalid[:, :, None, None]
+        self.keys[:, start:end].masked_fill_(expanded_invalid, 0)
+        self.values[:, start:end].masked_fill_(expanded_invalid, 0)
+        self.physical_positions[:, start:end].masked_fill_(invalid, -1)
+        self.logical_positions[:, start:end].masked_fill_(invalid, -1)
+        self.rope_positions[:, start:end].masked_fill_(invalid, -1)
+        self.length = required
+
+    def view(self) -> MoTSidecarView:
+        if self.capacity == 0:
+            raise RuntimeError("Sidecar buffer is not initialized.")
+        return MoTSidecarView(
+            keys=self.keys[:, : self.length],
+            values=self.values[:, : self.length],
+            valid_mask=self.valid_mask[:, : self.length],
+            physical_positions=self.physical_positions[:, : self.length],
+            logical_positions=self.logical_positions[:, : self.length],
+            rope_positions=self.rope_positions[:, : self.length],
+            capacity=self.capacity,
+        )
+
+    def crop_physical(self, target_length: int) -> None:
+        if self.capacity == 0:
+            return
+        remove = self.valid_mask[:, : self.length] & self.physical_positions[
+            :, : self.length
+        ].ge(target_length)
+        if torch.any(remove):
+            expanded_remove = remove[:, :, None, None]
+            self.keys[:, : self.length].masked_fill_(expanded_remove, 0)
+            self.values[:, : self.length].masked_fill_(expanded_remove, 0)
+            self.valid_mask[:, : self.length].masked_fill_(remove, False)
+            self.physical_positions[:, : self.length].masked_fill_(remove, -1)
+            self.logical_positions[:, : self.length].masked_fill_(remove, -1)
+            self.rope_positions[:, : self.length].masked_fill_(remove, -1)
+        while self.length > 0 and not bool(
+            torch.any(self.valid_mask[:, self.length - 1]).item()
+        ):
+            self.length -= 1
+
+    def truncate(self, target_length: int) -> None:
+        if target_length < 0 or target_length > self.length:
+            raise ValueError(
+                f"Sidecar target_length must be in [0, {self.length}], got {target_length}."
+            )
+        if target_length == self.length:
+            return
+        self.keys[:, target_length : self.length].zero_()
+        self.values[:, target_length : self.length].zero_()
+        self.valid_mask[:, target_length : self.length].zero_()
+        self.physical_positions[:, target_length : self.length].fill_(-1)
+        self.logical_positions[:, target_length : self.length].fill_(-1)
+        self.rope_positions[:, target_length : self.length].fill_(-1)
+        self.length = target_length
+
+    def reset(self) -> None:
+        self.length = 0
+        if self.capacity == 0:
+            return
+        self.keys.zero_()
+        self.values.zero_()
+        self.valid_mask.zero_()
+        self.physical_positions.fill_(-1)
+        self.logical_positions.fill_(-1)
+        self.rope_positions.fill_(-1)
+
+    def index_select(self, indices: torch.LongTensor) -> None:
+        if self.capacity == 0:
+            self.batch_size = int(indices.numel())
+            return
+        local_indices = indices.to(device=self.keys.device)
+        self.keys = self.keys.index_select(0, local_indices)
+        self.values = self.values.index_select(0, local_indices)
+        self.valid_mask = self.valid_mask.index_select(0, local_indices)
+        self.physical_positions = self.physical_positions.index_select(0, local_indices)
+        self.logical_positions = self.logical_positions.index_select(0, local_indices)
+        self.rope_positions = self.rope_positions.index_select(0, local_indices)
+        self.batch_size = int(indices.numel())
+
+    def repeat_interleave(self, repeats: int) -> None:
+        if self.capacity > 0:
+            self.keys = self.keys.repeat_interleave(repeats, dim=0)
+            self.values = self.values.repeat_interleave(repeats, dim=0)
+            self.valid_mask = self.valid_mask.repeat_interleave(repeats, dim=0)
+            self.physical_positions = self.physical_positions.repeat_interleave(
+                repeats, dim=0
+            )
+            self.logical_positions = self.logical_positions.repeat_interleave(
+                repeats, dim=0
+            )
+            self.rope_positions = self.rope_positions.repeat_interleave(repeats, dim=0)
+        self.batch_size *= repeats
+
+
+class _GeometricDynamicLayer(DynamicLayer):
+    """HF-compatible dynamic layer whose appends reuse geometrically grown storage."""
+
+    def __init__(self, *, initial_capacity: int) -> None:
+        super().__init__()
+        self._initial_capacity = int(initial_capacity)
+        self._length = 0
+        self._capacity = 0
+        self._key_storage: torch.Tensor | None = None
+        self._value_storage: torch.Tensor | None = None
+
+    @property
+    def capacity(self) -> int:
+        return self._capacity
+
+    def lazy_initialization(self, key_states: torch.Tensor) -> None:
+        self.dtype = key_states.dtype
+        self.device = key_states.device
+        self.keys = key_states.new_empty(
+            (*key_states.shape[:-2], 0, key_states.shape[-1])
+        )
+        self.values = self.keys.clone()
+        self.is_initialized = True
+
+    def _refresh_views(self) -> None:
+        if self._key_storage is None or self._value_storage is None:
+            return
+        self.keys = self._key_storage[..., : self._length, :]
+        self.values = self._value_storage[..., : self._length, :]
+
+    def _ensure_capacity(self, key_states: torch.Tensor, required: int) -> None:
+        if required <= self._capacity:
+            return
+        if self._capacity == 0:
+            capacity = max(
+                self._initial_capacity,
+                required + max(64, required // 8),
+            )
+        else:
+            capacity = max(required, self._capacity * 2)
+        storage_shape = (*key_states.shape[:-2], capacity, key_states.shape[-1])
+        key_storage = key_states.new_empty(storage_shape)
+        value_storage = key_states.new_empty(storage_shape)
+        if self._length > 0:
+            key_storage[..., : self._length, :].copy_(
+                self._key_storage[..., : self._length, :]
+            )
+            value_storage[..., : self._length, :].copy_(
+                self._value_storage[..., : self._length, :]
+            )
+        self._key_storage = key_storage
+        self._value_storage = value_storage
+        self._capacity = capacity
+
+    def update(
+        self,
+        key_states: torch.Tensor,
+        value_states: torch.Tensor,
+        cache_kwargs: Mapping[str, object] | None = None,
+    ) -> tuple[torch.Tensor, torch.Tensor]:
+        del cache_kwargs
+        if not self.is_initialized:
+            self.lazy_initialization(key_states)
+        if tuple(key_states.shape) != tuple(value_states.shape):
+            raise ValueError(
+                f"Native layer K/V shape mismatch: keys={tuple(key_states.shape)}, values={tuple(value_states.shape)}."
+            )
+        if self._length > 0 and (
+            key_states.dtype != self._key_storage.dtype
+            or key_states.device != self._key_storage.device
+            or tuple(key_states.shape[:-2]) != tuple(self._key_storage.shape[:-2])
+            or int(key_states.shape[-1]) != int(self._key_storage.shape[-1])
+        ):
+            raise ValueError(
+                "Native layer append dtype, device, batch, head, and head_dim must stay fixed."
+            )
+        token_count = int(key_states.shape[-2])
+        required = self._length + token_count
+        self._ensure_capacity(key_states, required)
+        self._key_storage[..., self._length : required, :].copy_(key_states)
+        self._value_storage[..., self._length : required, :].copy_(value_states)
+        self._length = required
+        self._refresh_views()
+        return self.keys, self.values
+
+    def get_seq_length(self) -> int:
+        return self._length
+
+    def crop(self, max_length: int) -> None:
+        if max_length < 0:
+            max_length = self._length - abs(max_length)
+        self._length = max(0, min(self._length, int(max_length)))
+        self._refresh_views()
+
+    def reset(self) -> None:
+        if self._length > 0:
+            self._key_storage[..., : self._length, :].zero_()
+            self._value_storage[..., : self._length, :].zero_()
+        self._length = 0
+        self._refresh_views()
+
+    def _select_batch(self, indices: torch.LongTensor) -> None:
+        if self._capacity == 0:
+            return
+        local_indices = indices.to(device=self._key_storage.device)
+        self._key_storage = self._key_storage.index_select(0, local_indices)
+        self._value_storage = self._value_storage.index_select(0, local_indices)
+        self._refresh_views()
+
+    def reorder_cache(self, beam_idx: torch.LongTensor) -> None:
+        self._select_batch(beam_idx)
+
+    def batch_select_indices(self, indices: torch.Tensor) -> None:
+        self._select_batch(indices.to(dtype=torch.long))
+
+    def batch_repeat_interleave(self, repeats: int) -> None:
+        if self._capacity > 0:
+            self._key_storage = self._key_storage.repeat_interleave(repeats, dim=0)
+            self._value_storage = self._value_storage.repeat_interleave(repeats, dim=0)
+            self._refresh_views()
+
+    def offload(self) -> None:
+        if self._capacity > 0:
+            self._key_storage = self._key_storage.to("cpu", non_blocking=True)
+            self._value_storage = self._value_storage.to("cpu", non_blocking=True)
+            self._refresh_views()
+
+    def prefetch(self) -> None:
+        if self._capacity > 0 and self._key_storage.device != self.device:
+            self._key_storage = self._key_storage.to(self.device, non_blocking=True)
+            self._value_storage = self._value_storage.to(self.device, non_blocking=True)
+            self._refresh_views()
+
+
+class MoTDynamicCache(DynamicCache):
+    """HF dynamic cache plus batch-stable MoT runtime sidecars."""
+
+    is_mot_generation_cache = True
+
+    def __init__(
+        self,
+        *,
+        request_ids: Sequence[_RequestId],
+        config=None,
+        initial_sidecar_capacity: int = 16,
+        offloading: bool = False,
+        offload_only_non_sliding: bool = False,
+    ) -> None:
+        if initial_sidecar_capacity <= 0:
+            raise ValueError(
+                f"initial_sidecar_capacity must be positive, got {initial_sidecar_capacity}."
+            )
+        if offloading or offload_only_non_sliding:
+            raise RuntimeError(
+                "MoTDynamicCache does not support offloading because runtime sidecars must remain colocated."
+            )
+        super().__init__(
+            config=config,
+            offloading=False,
+            offload_only_non_sliding=False,
+        )
+        if any(bool(getattr(layer, "is_sliding", False)) for layer in self.layers):
+            raise RuntimeError(
+                "MoTDynamicCache does not support sliding-window decoder layers."
+            )
+        self.layers = [
+            _GeometricDynamicLayer(initial_capacity=initial_sidecar_capacity)
+            for _layer in self.layers
+        ]
+        # Config-less caches are still supported: `update` appends the same
+        # geometric layer type instead of DynamicCache's concatenating default.
+        self.layer_class_to_replicate = None
+        self._request_ids = _normalize_request_ids(request_ids)
+        self._batch_size = len(self._request_ids)
+        self._initial_sidecar_capacity = int(initial_sidecar_capacity)
+        self._position_buffer = _PositionBuffer(
+            batch_size=self._batch_size,
+            initial_capacity=self._initial_sidecar_capacity,
+        )
+        self._sidecars: dict[str, dict[int, _SidecarBuffer]] = {}
+        self._physical_length = 0
+        self._row_logical_next_positions = torch.zeros(
+            self._batch_size, dtype=torch.long
+        )
+        self._row_rope_next_positions = torch.zeros(self._batch_size, dtype=torch.long)
+        self._payload_fingerprints: tuple[str, ...] | None = None
+        self._payload_epochs = torch.zeros(self._batch_size, dtype=torch.long)
+        self._config = (
+            config.get_text_config(decoder=True) if config is not None else None
+        )
+        # Cached mixed generation owns these request-local values. Declaring the
+        # complete schema here keeps cache lifetime and reset behavior explicit.
+        # `_streams` is the single registry for every incremental KV stream,
+        # keyed by (kind, row, layer, span):
+        #   kind="llm"/"residual": layer is the Qwen layer index, span is None,
+        #     and the stream lives for the whole request row.
+        #   kind="ts": layer is the TimesFM layer index, span is the target
+        #     slot, and the stream lives for one generated span.
+        # `None` means no incremental runtime streams are currently bound.
+        self._streams: dict[StreamKey, dict[str, object]] | None = None
+        self._mot_active_segment_ids: tuple[int, ...] = ()
+        self._mot_identity_runtime: bool = False
+        # Non-KV per-span decode state (hidden, patch_count, target, ...).
+        self._span_states: dict[tuple[int, int], dict[str, object]] = {}
+        self._mot_pending_text_logits: torch.Tensor | None = None
+        self._mot_stats: list[dict[str, int]] = []
+        self._mot_pad_token_id: int = 0
+
+    @classmethod
+    def from_legacy_cache(
+        cls,
+        past_key_values: tuple[tuple[torch.Tensor, torch.Tensor], ...],
+    ) -> "MoTDynamicCache":
+        del past_key_values
+        raise RuntimeError(
+            "MoT cannot recover runtime, position, and payload sidecars from a legacy tuple cache."
+        )
+
+    def to_legacy_cache(self) -> tuple[tuple[torch.Tensor, torch.Tensor], ...]:
+        raise RuntimeError(
+            "MoT cache cannot be converted to a legacy tuple without losing required sidecars."
+        )
+
+    def offload(self, layer_idx: int, only_non_sliding: bool = True) -> None:
+        del layer_idx, only_non_sliding
+        raise RuntimeError(
+            "MoTDynamicCache does not support offloading required MoT sidecars."
+        )
+
+    def prefetch(self, layer_idx: int, only_non_sliding: bool = True) -> None:
+        del layer_idx, only_non_sliding
+        raise RuntimeError(
+            "MoTDynamicCache does not support offloaded sidecar prefetch."
+        )
+
+    @property
+    def request_ids(self) -> tuple[_RequestId, ...]:
+        return self._request_ids
+
+    @property
+    def batch_size(self) -> int:
+        return self._batch_size
+
+    @property
+    def physical_length(self) -> int:
+        return self._physical_length
+
+    @property
+    def row_logical_next_positions(self) -> torch.LongTensor:
+        return self._row_logical_next_positions
+
+    @property
+    def row_rope_next_positions(self) -> torch.LongTensor:
+        return self._row_rope_next_positions
+
+    @property
+    def payload_fingerprints(self) -> tuple[str, ...] | None:
+        return self._payload_fingerprints
+
+    @property
+    def payload_epochs(self) -> tuple[int, ...]:
+        return tuple(int(epoch) for epoch in self._payload_epochs.tolist())
+
+    @property
+    def has_streams(self) -> bool:
+        return self._streams is not None
+
+    def bind_streams(self, streams: Mapping[StreamKey, dict]) -> None:
+        """Bind the complete prefill stream registry exactly once."""
+        if self._streams is not None:
+            raise RuntimeError(
+                "MoT stream registry is already bound; reset the cache first."
+            )
+        validated: dict[StreamKey, dict[str, object]] = {}
+        for key, entry in streams.items():
+            _validate_stream_key(key)
+            if not isinstance(entry, dict):
+                raise TypeError(
+                    f"Stream entry for {key!r} must be a dict, got {type(entry).__name__}."
+                )
+            if int(key[1]) >= self._batch_size:
+                raise ValueError(
+                    f"Stream key {key!r} row exceeds cache batch size {self._batch_size}."
+                )
+            validated[key] = entry
+        self._streams = validated
+
+    def add_stream(self, key: StreamKey, entry: dict) -> None:
+        """Register one lazily created stream (e.g. a new TS target layer)."""
+        if self._streams is None:
+            raise RuntimeError("Cannot add a stream before the registry is bound.")
+        _validate_stream_key(key)
+        if key in self._streams:
+            raise RuntimeError(f"Stream {key!r} is already registered.")
+        if not isinstance(entry, dict):
+            raise TypeError(
+                f"Stream entry for {key!r} must be a dict, got {type(entry).__name__}."
+            )
+        self._streams[key] = entry
+
+    def get_stream(self, key: StreamKey) -> dict | None:
+        """Fetch one stream; None means the key is legitimately absent (e.g. unpaired layer)."""
+        if self._streams is None:
+            raise RuntimeError("MoT stream registry is not bound.")
+        return self._streams.get(_validate_stream_key(key))
+
+    def require_stream(self, key: StreamKey) -> dict:
+        stream = self.get_stream(key)
+        if stream is None:
+            raise RuntimeError(f"MoT stream registry has no entry for {key!r}.")
+        return stream
+
+    def iter_streams(
+        self,
+        *,
+        kind: str | None = None,
+        row: int | None = None,
+        span: int | None = None,
+    ):
+        """Yield (key, entry) pairs; the single enumeration point for reset/rollback/stats."""
+        if self._streams is None:
+            return
+        for key, entry in self._streams.items():
+            if kind is not None and key[0] != kind:
+                continue
+            if row is not None and key[1] != row:
+                continue
+            if span is not None and key[3] != span:
+                continue
+            yield key, entry
+
+    def _validate_native_kv(
+        self,
+        key_states: torch.Tensor,
+        value_states: torch.Tensor,
+        layer_idx: int,
+        cache_kwargs: Mapping[str, object] | None,
+    ) -> torch.LongTensor:
+        if not isinstance(key_states, torch.Tensor) or not isinstance(
+            value_states, torch.Tensor
+        ):
+            raise TypeError("Native cache key_states and value_states must be tensors.")
+        if key_states.ndim != 4 or value_states.ndim != 4:
+            raise ValueError(
+                "Native Qwen K/V must have shape [B,H_kv,T,D], got "
+                f"keys={tuple(key_states.shape)}, values={tuple(value_states.shape)}."
+            )
+        if tuple(key_states.shape) != tuple(value_states.shape):
+            raise ValueError(
+                f"Native Qwen K/V shape mismatch: keys={tuple(key_states.shape)}, values={tuple(value_states.shape)}."
+            )
+        if (
+            key_states.dtype != value_states.dtype
+            or key_states.device != value_states.device
+        ):
+            raise ValueError(
+                "Native Qwen K/V dtype/device mismatch: "
+                f"keys={key_states.dtype}/{key_states.device}, values={value_states.dtype}/{value_states.device}."
+            )
+        if not torch.is_floating_point(key_states):
+            raise TypeError(
+                f"Native Qwen K/V must be floating point, got {key_states.dtype}."
+            )
+        if int(key_states.shape[0]) != self._batch_size:
+            raise ValueError(
+                f"Native Qwen K/V batch must match request_ids: expected {self._batch_size}, got {key_states.shape[0]}."
+            )
+        if int(key_states.shape[2]) <= 0:
+            raise ValueError("Native Qwen K/V update must contain at least one token.")
+        if layer_idx < 0 or (
+            self._config is not None
+            and layer_idx >= int(self._config.num_hidden_layers)
+        ):
+            raise ValueError(f"Native cache layer_idx out of range: {layer_idx}.")
+        if self._config is not None:
+            expected_heads = int(self._config.num_key_value_heads)
+            expected_head_dim = int(
+                getattr(
+                    self._config,
+                    "head_dim",
+                    self._config.hidden_size // self._config.num_attention_heads,
+                )
+            )
+            if tuple(key_states.shape[1::2]) != (expected_heads, expected_head_dim):
+                raise ValueError(
+                    "Native cache must keep compact Qwen KV heads: "
+                    f"expected H_kv/D={(expected_heads, expected_head_dim)}, "
+                    f"got {(int(key_states.shape[1]), int(key_states.shape[3]))}."
+                )
+
+        previous_length = int(self.get_seq_length(layer_idx))
+        token_count = int(key_states.shape[2])
+        raw_cache_position = (
+            None if cache_kwargs is None else cache_kwargs.get("cache_position")
+        )
+        if raw_cache_position is None:
+            cache_position = torch.arange(
+                previous_length,
+                previous_length + token_count,
+                device=key_states.device,
+                dtype=torch.long,
+            )
+        else:
+            if not isinstance(raw_cache_position, torch.Tensor):
+                raise TypeError(
+                    "cache_position must be a tensor for native cache update, got "
+                    f"{type(raw_cache_position).__name__}."
+                )
+            cache_position = _require_positions(
+                raw_cache_position,
+                name="cache_position",
+                shape=(token_count,),
+                device=key_states.device,
+            )
+        expected = torch.arange(
+            previous_length,
+            previous_length + token_count,
+            device=key_states.device,
+            dtype=torch.long,
+        )
+        if not torch.equal(cache_position, expected):
+            raise ValueError(
+                "Dynamic cache_position must be a contiguous append: "
+                f"expected {expected.tolist()}, got {cache_position.tolist()}."
+            )
+        return cache_position
+
+    def update(
+        self,
+        key_states: torch.Tensor,
+        value_states: torch.Tensor,
+        layer_idx: int,
+        cache_kwargs: Mapping[str, object] | None = None,
+    ) -> tuple[torch.Tensor, torch.Tensor]:
+        cache_position = self._validate_native_kv(
+            key_states, value_states, layer_idx, cache_kwargs
+        )
+        while len(self.layers) <= layer_idx:
+            self.layers.append(
+                _GeometricDynamicLayer(initial_capacity=self._initial_sidecar_capacity)
+            )
+        keys, values = super().update(
+            key_states, value_states, layer_idx, dict(cache_kwargs or {})
+        )
+        resulting_length = int(self.get_seq_length(layer_idx))
+        expected_length = int(cache_position[-1].item()) + 1
+        if resulting_length != expected_length:
+            raise RuntimeError(
+                f"Native cache length drifted at layer {layer_idx}: expected {expected_length}, got {resulting_length}."
+            )
+        if layer_idx == 0:
+            if self._physical_length not in (
+                int(cache_position[0].item()),
+                expected_length,
+            ):
+                raise RuntimeError(
+                    "Native layer 0 update disagrees with physical_length: "
+                    f"physical_length={self._physical_length}, cache_position={cache_position.tolist()}."
+                )
+            self._physical_length = expected_length
+        elif self._physical_length != expected_length:
+            raise RuntimeError(
+                "Native layers must update with the same physical timeline: "
+                f"layer={layer_idx}, layer_length={expected_length}, physical_length={self._physical_length}."
+            )
+        return keys, values
+
+    def record_position_metadata(
+        self,
+        *,
+        cache_position: torch.LongTensor,
+        valid_mask: torch.Tensor,
+        logical_positions: torch.LongTensor,
+        rope_positions: torch.LongTensor,
+        active_rows: torch.Tensor | None = None,
+    ) -> None:
+        if not isinstance(cache_position, torch.Tensor) or cache_position.ndim != 1:
+            raise ValueError(
+                "cache_position metadata must be a 1D tensor, got "
+                f"{None if not isinstance(cache_position, torch.Tensor) else tuple(cache_position.shape)}."
+            )
+        token_count = int(cache_position.numel())
+        if token_count <= 0:
+            raise ValueError(
+                "Position metadata must contain at least one physical token."
+            )
+        valid_mask = _require_bool_mask(
+            valid_mask,
+            name="valid_mask",
+            shape=(self._batch_size, token_count),
+        )
+        if (
+            cache_position.device != valid_mask.device
+            or cache_position.dtype is not torch.long
+        ):
+            raise ValueError(
+                "cache_position must be torch.long on the metadata device, got "
+                f"{cache_position.dtype} on {cache_position.device}, metadata on {valid_mask.device}."
+            )
+        logical_positions = _require_positions(
+            logical_positions,
+            name="logical_positions",
+            shape=(self._batch_size, token_count),
+            device=valid_mask.device,
+        )
+        rope_positions = _require_positions(
+            rope_positions,
+            name="rope_positions",
+            shape=(self._batch_size, token_count),
+            device=valid_mask.device,
+        )
+        expected_start = self._position_buffer.length
+        expected = torch.arange(
+            expected_start,
+            expected_start + token_count,
+            device=cache_position.device,
+            dtype=torch.long,
+        )
+        if not torch.equal(cache_position, expected):
+            raise ValueError(
+                "Position metadata must append contiguously: "
+                f"expected {expected.tolist()}, got {cache_position.tolist()}."
+            )
+        if int(cache_position[-1].item()) >= self._physical_length:
+            raise RuntimeError(
+                "Position metadata cannot run ahead of native K/V: "
+                f"last_position={int(cache_position[-1].item())}, physical_length={self._physical_length}."
+            )
+        if active_rows is not None:
+            active_rows = _require_bool_mask(
+                active_rows,
+                name="active_rows",
+                shape=(self._batch_size,),
+            )
+            if active_rows.device != valid_mask.device:
+                raise ValueError(
+                    f"active_rows must be on {valid_mask.device}, got {active_rows.device}."
+                )
+            valid_mask = valid_mask & active_rows[:, None]
+        if torch.any(logical_positions[valid_mask] < 0):
+            raise ValueError("Valid logical_positions must be non-negative.")
+        if torch.any(rope_positions[valid_mask] < 0):
+            raise ValueError("Valid rope_positions must be non-negative.")
+
+        self._position_buffer.append(
+            valid_mask=valid_mask,
+            logical_positions=logical_positions,
+            rope_positions=rope_positions,
+        )
+        if self._row_logical_next_positions.device != valid_mask.device:
+            if expected_start != 0:
+                raise RuntimeError(
+                    "Row logical position metadata changed device after prefill."
+                )
+            self._row_logical_next_positions = self._row_logical_next_positions.to(
+                valid_mask.device
+            )
+            self._row_rope_next_positions = self._row_rope_next_positions.to(
+                valid_mask.device
+            )
+        for row_idx in range(self._batch_size):
+            row_valid = valid_mask[row_idx]
+            if bool(torch.any(row_valid).item()):
+                last_idx = int(torch.nonzero(row_valid, as_tuple=False)[-1].item())
+                self._row_logical_next_positions[row_idx] = (
+                    logical_positions[row_idx, last_idx] + 1
+                )
+                self._row_rope_next_positions[row_idx] = (
+                    rope_positions[row_idx, last_idx] + 1
+                )
+
+    def _write_sidecar(
+        self,
+        *,
+        name: str,
+        layer_idx: int,
+        keys: torch.Tensor,
+        values: torch.Tensor,
+        valid_mask: torch.Tensor,
+        physical_positions: torch.LongTensor,
+        logical_positions: torch.LongTensor,
+        rope_positions: torch.LongTensor,
+        active_rows: torch.Tensor | None = None,
+        replace_start: int | None = None,
+    ) -> MoTSidecarView:
+        if not isinstance(name, str) or not name:
+            raise ValueError(f"Sidecar name must be a non-empty string, got {name!r}.")
+        if layer_idx < 0:
+            raise ValueError(
+                f"Sidecar layer_idx must be non-negative, got {layer_idx}."
+            )
+        if not isinstance(keys, torch.Tensor) or not isinstance(values, torch.Tensor):
+            raise TypeError("Sidecar keys and values must be tensors.")
+        if keys.ndim != 4 or values.ndim != 4:
+            raise ValueError(
+                "Sidecar K/V must have shape [B,T,H,D], got "
+                f"keys={tuple(keys.shape)}, values={tuple(values.shape)}."
+            )
+        if tuple(keys.shape) != tuple(values.shape):
+            raise ValueError(
+                f"Sidecar K/V shape mismatch: keys={tuple(keys.shape)}, values={tuple(values.shape)}."
+            )
+        if keys.dtype != values.dtype or keys.device != values.device:
+            raise ValueError(
+                "Sidecar K/V dtype/device mismatch: "
+                f"keys={keys.dtype}/{keys.device}, values={values.dtype}/{values.device}."
+            )
+        if not torch.is_floating_point(keys):
+            raise TypeError(f"Sidecar K/V must be floating point, got {keys.dtype}.")
+        batch_size, token_count = int(keys.shape[0]), int(keys.shape[1])
+        if batch_size != self._batch_size:
+            raise ValueError(
+                f"Sidecar batch must be {self._batch_size}, got {batch_size}."
+            )
+        if token_count <= 0 or int(keys.shape[2]) <= 0 or int(keys.shape[3]) <= 0:
+            raise ValueError(
+                f"Sidecar K/V dimensions must be positive, got {tuple(keys.shape)}."
+            )
+        valid_mask = _require_bool_mask(
+            valid_mask,
+            name="valid_mask",
+            shape=(self._batch_size, token_count),
+        )
+        if valid_mask.device != keys.device:
+            raise ValueError(
+                f"valid_mask must be on {keys.device}, got {valid_mask.device}."
+            )
+        if active_rows is not None:
+            active_rows = _require_bool_mask(
+                active_rows,
+                name="active_rows",
+                shape=(self._batch_size,),
+            )
+            if active_rows.device != keys.device:
+                raise ValueError(
+                    f"active_rows must be on {keys.device}, got {active_rows.device}."
+                )
+            valid_mask = valid_mask & active_rows[:, None]
+        physical_positions = _expanded_positions(
+            physical_positions,
+            batch_size=self._batch_size,
+            token_count=token_count,
+            name="physical_positions",
+            device=keys.device,
+        )
+        logical_positions = _require_positions(
+            logical_positions,
+            name="logical_positions",
+            shape=(self._batch_size, token_count),
+            device=keys.device,
+        )
+        rope_positions = _require_positions(
+            rope_positions,
+            name="rope_positions",
+            shape=(self._batch_size, token_count),
+            device=keys.device,
+        )
+        for positions, position_name in (
+            (physical_positions, "physical_positions"),
+            (logical_positions, "logical_positions"),
+            (rope_positions, "rope_positions"),
+        ):
+            if torch.any(positions[valid_mask] < 0):
+                raise ValueError(f"Valid {position_name} must be non-negative.")
+        if torch.any(physical_positions[valid_mask] >= self._physical_length):
+            raise ValueError(
+                "Sidecar physical_positions must refer to existing native cache columns: "
+                f"physical_length={self._physical_length}."
+            )
+
+        layer_buffers = self._sidecars.setdefault(name, {})
+        buffer = layer_buffers.get(layer_idx)
+        if buffer is None:
+            if replace_start is not None:
+                raise KeyError(
+                    f"Cannot replace absent MoT sidecar {name!r} at layer {layer_idx}."
+                )
+            buffer = _SidecarBuffer(
+                batch_size=self._batch_size,
+                initial_capacity=self._initial_sidecar_capacity,
+            )
+            layer_buffers[layer_idx] = buffer
+        if replace_start is not None:
+            if isinstance(replace_start, bool) or not isinstance(replace_start, int):
+                raise TypeError(
+                    f"replace_start must be an integer, got {type(replace_start).__name__}."
+                )
+            if buffer.capacity > 0 and (
+                buffer.keys.dtype != keys.dtype
+                or buffer.keys.device != keys.device
+                or tuple(buffer.keys.shape[2:]) != tuple(keys.shape[2:])
+            ):
+                raise ValueError(
+                    "Replacement sidecar dtype/device/head shape must match the existing stream: "
+                    f"existing={buffer.keys.dtype}/{buffer.keys.device}/{tuple(buffer.keys.shape[2:])}, "
+                    f"replacement={keys.dtype}/{keys.device}/{tuple(keys.shape[2:])}."
+                )
+            buffer.truncate(replace_start)
+        buffer.append(
+            keys=keys,
+            values=values,
+            valid_mask=valid_mask,
+            physical_positions=physical_positions,
+            logical_positions=logical_positions,
+            rope_positions=rope_positions,
+        )
+        return buffer.view()
+
+    def append_sidecar(
+        self,
+        *,
+        name: str,
+        layer_idx: int,
+        keys: torch.Tensor,
+        values: torch.Tensor,
+        valid_mask: torch.Tensor,
+        physical_positions: torch.LongTensor,
+        logical_positions: torch.LongTensor,
+        rope_positions: torch.LongTensor,
+        active_rows: torch.Tensor | None = None,
+    ) -> MoTSidecarView:
+        return self._write_sidecar(
+            name=name,
+            layer_idx=layer_idx,
+            keys=keys,
+            values=values,
+            valid_mask=valid_mask,
+            physical_positions=physical_positions,
+            logical_positions=logical_positions,
+            rope_positions=rope_positions,
+            active_rows=active_rows,
+        )
+
+    def rollback_sidecar(
+        self, name: str, layer_idx: int, *, target_length: int
+    ) -> MoTSidecarView:
+        """Drop one mutable sidecar tail without touching native or sibling streams."""
+        try:
+            buffer = self._sidecars[name][layer_idx]
+        except KeyError as exc:
+            raise KeyError(
+                f"No MoT sidecar named {name!r} at layer {layer_idx}."
+            ) from exc
+        if isinstance(target_length, bool) or not isinstance(target_length, int):
+            raise TypeError(
+                f"target_length must be an integer, got {type(target_length).__name__}."
+            )
+        buffer.truncate(target_length)
+        return buffer.view()
+
+    def replace_sidecar_tail(
+        self,
+        *,
+        name: str,
+        layer_idx: int,
+        start: int,
+        keys: torch.Tensor,
+        values: torch.Tensor,
+        valid_mask: torch.Tensor,
+        physical_positions: torch.LongTensor,
+        logical_positions: torch.LongTensor,
+        rope_positions: torch.LongTensor,
+        active_rows: torch.Tensor | None = None,
+    ) -> MoTSidecarView:
+        """Validate a replacement, then atomically discard and rewrite one stream tail."""
+        return self._write_sidecar(
+            name=name,
+            layer_idx=layer_idx,
+            keys=keys,
+            values=values,
+            valid_mask=valid_mask,
+            physical_positions=physical_positions,
+            logical_positions=logical_positions,
+            rope_positions=rope_positions,
+            active_rows=active_rows,
+            replace_start=start,
+        )
+
+    def get_sidecar(self, name: str, layer_idx: int) -> MoTSidecarView:
+        try:
+            return self._sidecars[name][layer_idx].view()
+        except KeyError as exc:
+            raise KeyError(
+                f"No MoT sidecar named {name!r} at layer {layer_idx}."
+            ) from exc
+
+    def bind_payload_fingerprints(self, fingerprints: Sequence[str]) -> None:
+        normalized = tuple(fingerprints)
+        if len(normalized) != self._batch_size or any(
+            not isinstance(value, str) or not value for value in normalized
+        ):
+            raise ValueError(
+                f"Payload fingerprints must contain {self._batch_size} non-empty strings, got {normalized}."
+            )
+        if self._payload_fingerprints is None:
+            self._payload_fingerprints = normalized
+            return
+        self.validate_payload_fingerprints(normalized)
+
+    def validate_payload_fingerprints(self, fingerprints: Sequence[str]) -> None:
+        if self._payload_fingerprints is None:
+            raise RuntimeError("Payload fingerprints are not bound to this cache.")
+        candidate = tuple(fingerprints)
+        if len(candidate) != self._batch_size:
+            raise ValueError(
+                f"Expected {self._batch_size} payload fingerprints, got {len(candidate)}."
+            )
+        changed = [
+            self._request_ids[row_idx]
+            for row_idx, (cached, current) in enumerate(
+                zip(self._payload_fingerprints, candidate, strict=True)
+            )
+            if cached != current
+        ]
+        if changed:
+            raise RuntimeError(
+                "MoT payload changed while KV cache was live; start a new payload epoch for request_ids="
+                f"{changed}."
+            )
+
+    def start_payload_epoch(self, fingerprints: Sequence[str]) -> None:
+        if self._streams is not None:
+            raise RuntimeError(
+                "External payload epochs cannot preserve bound MoT runtime streams; start a fresh prefill cache."
+            )
+        if self._payload_fingerprints is None:
+            raise RuntimeError(
+                "Cannot advance payload epoch before binding initial fingerprints."
+            )
+        candidate = tuple(fingerprints)
+        if len(candidate) != self._batch_size or any(
+            not isinstance(value, str) or not value for value in candidate
+        ):
+            raise ValueError(
+                f"Expected {self._batch_size} non-empty payload fingerprints, got {candidate}."
+            )
+        changed = torch.tensor(
+            [
+                cached != current
+                for cached, current in zip(
+                    self._payload_fingerprints, candidate, strict=True
+                )
+            ],
+            dtype=torch.bool,
+        )
+        if not bool(torch.any(changed).item()):
+            raise ValueError(
+                "start_payload_epoch requires at least one changed payload fingerprint."
+            )
+        next_epochs = self._payload_epochs + changed.to(dtype=torch.long)
+        self._clear_content(clear_payload=False)
+        self._payload_fingerprints = candidate
+        self._payload_epochs = next_epochs
+
+    def commit_payload_mutation(
+        self,
+        fingerprints: Sequence[str],
+        *,
+        changed_rows: Sequence[int],
+    ) -> None:
+        """Record a payload mutation only after its TS and fusion streams were synchronized."""
+        if self._payload_fingerprints is None:
+            raise RuntimeError(
+                "Cannot commit payload mutation before binding initial fingerprints."
+            )
+        candidate = tuple(fingerprints)
+        if len(candidate) != self._batch_size or any(
+            not isinstance(value, str) or not value for value in candidate
+        ):
+            raise ValueError(
+                f"Expected {self._batch_size} non-empty payload fingerprints, got {candidate}."
+            )
+        normalized_rows = tuple(int(row_idx) for row_idx in changed_rows)
+        if len(set(normalized_rows)) != len(normalized_rows):
+            raise ValueError(f"changed_rows must be unique, got {normalized_rows}.")
+        if any(
+            row_idx < 0 or row_idx >= self._batch_size for row_idx in normalized_rows
+        ):
+            raise IndexError(
+                f"changed_rows must stay in [0,{self._batch_size}), got {normalized_rows}."
+            )
+        observed_changed = {
+            row_idx
+            for row_idx, (cached, current) in enumerate(
+                zip(self._payload_fingerprints, candidate, strict=True)
+            )
+            if cached != current
+        }
+        expected_changed = set(normalized_rows)
+        if observed_changed != expected_changed:
+            raise RuntimeError(
+                "Controlled payload mutation rows do not match fingerprint changes: "
+                f"expected={sorted(expected_changed)}, observed={sorted(observed_changed)}."
+            )
+        if not observed_changed:
+            raise ValueError(
+                "commit_payload_mutation requires at least one changed payload fingerprint."
+            )
+        changed_mask = torch.tensor(
+            [row_idx in observed_changed for row_idx in range(self._batch_size)],
+            dtype=torch.long,
+        )
+        self._payload_epochs = self._payload_epochs + changed_mask
+        self._payload_fingerprints = candidate
+
+    def _reject_bound_batch_transform(self, operation: str) -> None:
+        if self._streams is not None:
+            raise RuntimeError(
+                f"MoT cache {operation} is unavailable while request-local runtime streams are bound."
+            )
+
+    def _validated_batch_indices(
+        self, indices: torch.LongTensor, *, name: str
+    ) -> torch.LongTensor:
+        if (
+            not isinstance(indices, torch.Tensor)
+            or indices.dtype is not torch.long
+            or indices.ndim != 1
+        ):
+            raise TypeError(f"{name} must be a 1D torch.long tensor.")
+        if indices.numel() <= 0:
+            raise ValueError(f"{name} must select at least one row.")
+        indices_cpu = indices.detach().to(device="cpu")
+        if (
+            int(indices_cpu.min().item()) < 0
+            or int(indices_cpu.max().item()) >= self._batch_size
+        ):
+            raise IndexError(
+                f"{name} values must be in [0, {self._batch_size}), got {indices_cpu.tolist()}."
+            )
+        return indices_cpu
+
+    def _select_metadata(self, indices: torch.LongTensor) -> None:
+        selected = indices.tolist()
+        self._request_ids = tuple(self._request_ids[index] for index in selected)
+        self._payload_epochs = self._payload_epochs.index_select(0, indices)
+        if self._payload_fingerprints is not None:
+            self._payload_fingerprints = tuple(
+                self._payload_fingerprints[index] for index in selected
+            )
+        row_indices = indices.to(device=self._row_logical_next_positions.device)
+        self._row_logical_next_positions = (
+            self._row_logical_next_positions.index_select(0, row_indices)
+        )
+        self._row_rope_next_positions = self._row_rope_next_positions.index_select(
+            0, row_indices
+        )
+        self._position_buffer.index_select(indices)
+        for layer_buffers in self._sidecars.values():
+            for buffer in layer_buffers.values():
+                buffer.index_select(indices)
+        self._batch_size = int(indices.numel())
+
+    def reorder_cache(self, beam_idx: torch.LongTensor) -> None:
+        self._reject_bound_batch_transform("reorder")
+        indices = self._validated_batch_indices(beam_idx, name="beam_idx")
+        super().reorder_cache(indices)
+        self._select_metadata(indices)
+
+    def batch_select_indices(self, indices: torch.Tensor) -> None:
+        self._reject_bound_batch_transform("batch selection")
+        indices_cpu = self._validated_batch_indices(indices, name="indices")
+        super().batch_select_indices(indices_cpu)
+        self._select_metadata(indices_cpu)
+
+    def batch_repeat_interleave(self, repeats: int) -> None:
+        self._reject_bound_batch_transform("batch repeat")
+        if isinstance(repeats, bool) or not isinstance(repeats, int) or repeats <= 0:
+            raise ValueError(f"repeats must be a positive integer, got {repeats!r}.")
+        super().batch_repeat_interleave(repeats)
+        self._request_ids = tuple(
+            request_id for request_id in self._request_ids for _ in range(repeats)
+        )
+        self._payload_epochs = self._payload_epochs.repeat_interleave(repeats)
+        if self._payload_fingerprints is not None:
+            self._payload_fingerprints = tuple(
+                fingerprint
+                for fingerprint in self._payload_fingerprints
+                for _ in range(repeats)
+            )
+        self._row_logical_next_positions = (
+            self._row_logical_next_positions.repeat_interleave(repeats)
+        )
+        self._row_rope_next_positions = self._row_rope_next_positions.repeat_interleave(
+            repeats
+        )
+        self._position_buffer.repeat_interleave(repeats)
+        for layer_buffers in self._sidecars.values():
+            for buffer in layer_buffers.values():
+                buffer.repeat_interleave(repeats)
+        self._batch_size *= repeats
+
+    def _target_crop_length(self, max_length: int) -> int:
+        if isinstance(max_length, bool) or not isinstance(max_length, int):
+            raise TypeError(
+                f"max_length must be an integer, got {type(max_length).__name__}."
+            )
+        target = (
+            self._physical_length - abs(max_length) if max_length < 0 else max_length
+        )
+        return max(0, min(self._physical_length, target))
+
+    def crop(self, max_length: int) -> None:
+        self._reject_bound_batch_transform("crop")
+        target = self._target_crop_length(max_length)
+        super().crop(target)
+        self._physical_length = target
+        self._position_buffer.crop(target)
+        for layer_buffers in self._sidecars.values():
+            for buffer in layer_buffers.values():
+                buffer.crop_physical(target)
+        self._recompute_row_next_positions()
+
+    def _recompute_row_next_positions(self) -> None:
+        self._row_logical_next_positions.zero_()
+        self._row_rope_next_positions.zero_()
+        if self._position_buffer.length == 0:
+            return
+        valid = self._position_buffer.valid_mask[:, : self._position_buffer.length]
+        logical = self._position_buffer.logical_positions[
+            :, : self._position_buffer.length
+        ]
+        rope = self._position_buffer.rope_positions[:, : self._position_buffer.length]
+        for row_idx in range(self._batch_size):
+            row_valid = valid[row_idx]
+            if bool(torch.any(row_valid).item()):
+                last_idx = int(torch.nonzero(row_valid, as_tuple=False)[-1].item())
+                self._row_logical_next_positions[row_idx] = (
+                    logical[row_idx, last_idx] + 1
+                )
+                self._row_rope_next_positions[row_idx] = rope[row_idx, last_idx] + 1
+
+    def _clear_content(self, *, clear_payload: bool) -> None:
+        super().crop(0)
+        self._physical_length = 0
+        self._position_buffer.reset()
+        self._row_logical_next_positions.zero_()
+        self._row_rope_next_positions.zero_()
+        for layer_buffers in self._sidecars.values():
+            for buffer in layer_buffers.values():
+                buffer.reset()
+        if clear_payload:
+            self._payload_fingerprints = None
+            self._payload_epochs.zero_()
+
+    def reset(self) -> None:
+        self._clear_content(clear_payload=True)
+        self._streams = None
+        self._mot_active_segment_ids = ()
+        self._mot_identity_runtime = False
+        self._span_states = {}
+        self._mot_pending_text_logits = None
+        self._mot_stats = []
+        self._mot_pad_token_id = 0
+
+    def validate_consistency(self, *, require_position_metadata: bool = True) -> None:
+        for layer_idx, layer in enumerate(self.layers):
+            if (
+                layer.is_initialized
+                and int(layer.get_seq_length()) != self._physical_length
+            ):
+                raise RuntimeError(
+                    "Native cache layer length mismatch: "
+                    f"layer={layer_idx}, layer_length={int(layer.get_seq_length())}, "
+                    f"physical_length={self._physical_length}."
+                )
+        if (
+            require_position_metadata
+            and self._position_buffer.length != self._physical_length
+        ):
+            raise RuntimeError(
+                "Position metadata length mismatch: "
+                f"metadata={self._position_buffer.length}, physical_length={self._physical_length}."
+            )
+
+
+__all__ = [
+    "MoTDynamicCache",
+    "MoTSidecarView",
+    "fingerprint_payload_state",
+    "validate_mot_cache_inputs",
+]

timebraid/model/mot/pairing.py

@@ -0,0 +1,97 @@
+"""Layer-pairing helpers for the TimeBraid model core."""
+
+from __future__ import annotations
+
+from typing import Any, Dict, List, Optional
+
+_QWEN_PAIRABLE_LAYER_TYPE = "full_attention"
+
+
+def _build_pair_map(
+    num_q_layers: int,
+    num_t_layers: int,
+    pairing_mode: str,
+    q_layer_indices: Optional[List[int]] = None,
+) -> Dict[int, int]:
+    """Build the Q-layer -> TimesFM-layer mapping for one pairing mode."""
+    if num_t_layers <= 0:
+        raise ValueError(f"num_t_layers must be positive, got {num_t_layers}")
+    if q_layer_indices is None:
+        available_q_layers = list(range(num_q_layers))
+    else:
+        available_q_layers = [int(idx) for idx in q_layer_indices]
+    if not available_q_layers:
+        raise ValueError("No pairable Q layers are available for MoT bridge mapping.")
+    if available_q_layers != sorted(available_q_layers):
+        raise ValueError(
+            f"q_layer_indices must be sorted ascending, got {available_q_layers}"
+        )
+    if available_q_layers[0] < 0 or available_q_layers[-1] >= num_q_layers:
+        raise ValueError(
+            "q_layer_indices must stay inside [0, num_q_layers). "
+            f"got indices={available_q_layers}, q_layers={num_q_layers}"
+        )
+
+    if len(available_q_layers) > len(set(available_q_layers)):
+        raise ValueError(f"q_layer_indices must be unique, got {available_q_layers}")
+    if num_t_layers > num_q_layers and q_layer_indices is None:
+        raise ValueError(
+            f"TimesFM layer count exceeds Qwen layers: t_layers={num_t_layers}, q_layers={num_q_layers}"
+        )
+
+    if pairing_mode == "interleaved":
+        if len(available_q_layers) < num_t_layers:
+            raise ValueError(
+                "`interleaved` requires at least as many pairable Qwen layers as active TimesFM layers, "
+                f"got pairable_q_layers={len(available_q_layers)}, t_layers={num_t_layers}."
+            )
+        if num_t_layers == 1:
+            return {available_q_layers[0]: 0}
+        q_last = len(available_q_layers) - 1
+        t_last = num_t_layers - 1
+        selected_q_layers = [
+            available_q_layers[round(t_idx * q_last / t_last)]
+            for t_idx in range(num_t_layers)
+        ]
+        return {q_idx: t_idx for t_idx, q_idx in enumerate(selected_q_layers)}
+
+    if pairing_mode == "dsfp":
+        if len(available_q_layers) != num_q_layers:
+            raise ValueError(
+                "`dsfp` requires every Qwen layer to be pairable, got "
+                f"pairable_q_layers={len(available_q_layers)}, q_layers={num_q_layers}."
+            )
+        if num_t_layers != num_q_layers:
+            raise ValueError(
+                "`dsfp` maps the expanded TimesFM stack one-to-one onto Qwen depth, got "
+                f"t_layers={num_t_layers}, q_layers={num_q_layers}."
+            )
+        return {q_idx: t_idx for t_idx, q_idx in enumerate(available_q_layers)}
+
+    raise ValueError(
+        f"Unsupported MoT pairing mode: {pairing_mode!r}. Expected 'interleaved' or 'dsfp'."
+    )
+
+
+def _resolve_pairable_q_layer_indices(config: Any, num_q_layers: int) -> List[int]:
+    """Resolve which Qwen decoder layers can participate in the TimeBraid layer plan."""
+    layer_types = getattr(config, "layer_types", None)
+    if layer_types is None:
+        return list(range(num_q_layers))
+    if len(layer_types) != num_q_layers:
+        raise ValueError(
+            "config.layer_types length mismatch with num_hidden_layers: "
+            f"len(layer_types)={len(layer_types)}, num_hidden_layers={num_q_layers}"
+        )
+
+    pairable = [
+        idx
+        for idx, layer_type in enumerate(layer_types)
+        if layer_type == _QWEN_PAIRABLE_LAYER_TYPE
+    ]
+    if not pairable:
+        raise ValueError(
+            f"MoT requires at least one `{_QWEN_PAIRABLE_LAYER_TYPE}` layer, but none were found in "
+            "config.layer_types."
+        )
+    return pairable

timebraid/model/mot/runtime_state.py

@@ -0,0 +1,130 @@
+"""Immutable per-model MoT construction options."""
+
+from __future__ import annotations
+
+import math
+from types import MappingProxyType
+from typing import Any, Mapping
+
+import torch
+
+from .config_contract import MOT_OPTION_BY_NAME
+
+
+class MoTRuntimeOptions:
+    """Read-only option view used while initializing one TimeBraid model."""
+
+    __slots__ = ("_values",)
+
+    def __init__(self, values: Mapping[str, Any] | None = None) -> None:
+        if values is None:
+            values = {}
+        if not isinstance(values, Mapping):
+            raise TypeError(
+                f"MoT runtime options must be a mapping, got {type(values).__name__}."
+            )
+        # A read-only copy prevents both caller mutation and cross-model state
+        # leakage. Model construction can therefore overlap safely in threads;
+        # there is no process-global configuration transaction to race on.
+        self._values = MappingProxyType(dict(values))
+
+    def _value(self, name: str, default: Any) -> Any:
+        value = self._values.get(name, default)
+        return default if value is None else value
+
+    def value(self, name: str) -> Any:
+        """Read one option, taking its type and default from the registry.
+
+        Call sites used to restate both at every read, which is how the same
+        knob ended up with disagreeing defaults in different files. Here the
+        registry answers once.
+        """
+        option = MOT_OPTION_BY_NAME.get(name)
+        if option is None:
+            raise KeyError(
+                f"Unknown MoT runtime option {name!r}; declare it in MOT_OPTIONS."
+            )
+
+        present = self._values.get(name)
+        if option.is_required and present is None:
+            raise RuntimeError(
+                f"MoT runtime option {name!r} is supplied by the checkpoint restore "
+                "contract and has no default. A hand-built option map must set it "
+                "explicitly; substituting a value here would silently change the "
+                "trained topology."
+            )
+
+        # For a required option the value is present, so passing it as its own
+        # default simply routes it through the type check below.
+        default = present if option.is_required else option.default
+        if option.kind == "boolean":
+            return self.boolean(name, default)
+        if option.kind == "integer":
+            return self.integer(name, default)
+        if option.kind == "number":
+            return self.number(name, default)
+        if option.kind == "string":
+            return self.string(name, default)
+        if option.kind == "torch_dtype":
+            return self.torch_dtype(name, default)
+        raise KeyError(f"MoT option {name!r} has an unknown kind {option.kind!r}.")
+
+    def boolean(self, name: str, default: bool = False) -> bool:
+        value = self._values.get(name, default)
+        if type(value) is not bool:
+            raise TypeError(
+                f"MoT runtime option {name!r} must be a boolean, got {value!r}."
+            )
+        return value
+
+    def integer(self, name: str, default: int) -> int:
+        value = self._value(name, default)
+        if type(value) is not int:
+            raise TypeError(
+                f"MoT runtime option {name!r} must be an integer, got {value!r}."
+            )
+        return value
+
+    def string(self, name: str, default: str) -> str:
+        value = self._value(name, default)
+        if not isinstance(value, str):
+            raise TypeError(
+                f"MoT runtime option {name!r} must be a string, got {value!r}."
+            )
+        return value
+
+    def number(self, name: str, default: float) -> float:
+        value = self._value(name, default)
+        if isinstance(value, bool) or not isinstance(value, (int, float)):
+            raise TypeError(
+                f"MoT runtime option {name!r} must be numeric, got {value!r}."
+            )
+        result = float(value)
+        if not math.isfinite(result):
+            raise ValueError(
+                f"MoT runtime option {name!r} must be finite, got {value!r}."
+            )
+        return result
+
+    def torch_dtype(self, name: str, default: torch.dtype) -> torch.dtype:
+        value = self._value(name, default)
+
+        if isinstance(value, torch.dtype):
+            return value
+
+        if not isinstance(value, str):
+            raise TypeError(
+                f"MoT runtime option {name!r} must resolve to a torch dtype, "
+                f"got {value!r}."
+            )
+        text = value.strip().lower()
+        if text in {"torch.bfloat16", "bfloat16", "bf16"}:
+            return torch.bfloat16
+        if text in {"torch.float16", "float16", "fp16", "half"}:
+            return torch.float16
+        if text in {"torch.float32", "float32", "fp32", "float"}:
+            return torch.float32
+
+        raise ValueError(
+            f"MoT runtime option {name!r} must resolve to a torch dtype, got {value!r}."
+        )

timebraid/model/mot/structures.py

@@ -0,0 +1,270 @@
+"""Shared structural types and tag helpers for TimeBraid."""
+
+from __future__ import annotations
+
+from collections.abc import Mapping
+from dataclasses import dataclass
+from typing import NamedTuple, Optional
+
+import torch
+
+IGNORE_INDEX = -100
+
+
+def _extract_text_position_ids(
+    position_ids: Optional[torch.LongTensor],
+) -> Optional[torch.LongTensor]:
+    """Normalize position ids to MoT's 2D text timeline."""
+    if position_ids is None:
+        return None
+    if position_ids.ndim == 2:
+        return position_ids
+    if position_ids.ndim == 3 and position_ids.shape[0] in {3, 4}:
+        return position_ids[0]
+    raise RuntimeError(f"Unsupported position_ids shape: {tuple(position_ids.shape)}")
+
+
+ROLE_OBSERVED = 0
+ROLE_CONTEXT = 1
+ROLE_TARGET = 2
+TS_ROUTE_GENERATION = 0
+TS_ROUTE_UNDERSTANDING = 2
+
+
+TIMEBRAID_TS_PAYLOAD_FIELDS = (
+    "ts_values",
+    "ts_lengths",
+    "ts_loss_start_idxs",
+    "ts_loss_roi_masks",
+    "ts_roles",
+    "ts_segment_ids",
+    "ts_span_mask",
+    "ts_text_start_token_idxs",
+    "ts_text_end_token_idxs",
+)
+TIMEBRAID_REQUIRED_TS_PAYLOAD_FIELDS = tuple(
+    name for name in TIMEBRAID_TS_PAYLOAD_FIELDS if name != "ts_loss_roi_masks"
+)
+
+
+@dataclass
+class TimeBraidPayload:
+    """Curve tensors carried together through TimeBraid runtime compilation and rollout."""
+
+    ts_values: Optional[torch.Tensor] = None
+    ts_lengths: Optional[torch.Tensor] = None
+    ts_loss_start_idxs: Optional[torch.Tensor] = None
+    ts_loss_roi_masks: Optional[torch.Tensor] = None
+    ts_roles: Optional[torch.Tensor] = None
+    ts_segment_ids: Optional[torch.Tensor] = None
+    ts_span_mask: Optional[torch.Tensor] = None
+    ts_text_start_token_idxs: Optional[torch.Tensor] = None
+    ts_text_end_token_idxs: Optional[torch.Tensor] = None
+
+    @classmethod
+    def pop_from_kwargs(cls, kwargs: dict[str, object]) -> "TimeBraidPayload":
+        return cls(
+            **{name: kwargs.pop(name, None) for name in TIMEBRAID_TS_PAYLOAD_FIELDS}
+        )
+
+    def inject_generation_inputs(self, model_inputs: dict[str, object]) -> None:
+        """Flatten the payload only at the public Hugging Face model-kwargs boundary."""
+        for name in TIMEBRAID_TS_PAYLOAD_FIELDS:
+            value = getattr(self, name)
+            if value is not None:
+                model_inputs[name] = value
+
+    def has_runtime_inputs(self) -> bool:
+        return any(
+            getattr(self, name) is not None for name in TIMEBRAID_TS_PAYLOAD_FIELDS
+        )
+
+    def runtime_kwargs(self) -> dict[str, object]:
+        """Flatten the payload for public model calls that retain the HF tensor ABI."""
+        return {name: getattr(self, name) for name in TIMEBRAID_TS_PAYLOAD_FIELDS}
+
+
+@dataclass(frozen=True)
+class TimeBraidLayerStep:
+    """One LLM layer and its optional generation/understanding TimesFM layers."""
+
+    llm_layer: int
+    generation_tsfm_layer: Optional[int] = None
+    understanding_tsfm_layer: Optional[int] = None
+
+
+def build_timebraid_layer_plan(
+    *,
+    num_llm_layers: int,
+    generation_tsfm_by_llm_layer: Mapping[int, int],
+    understanding_tsfm_by_llm_layer: Mapping[int, int],
+) -> tuple[TimeBraidLayerStep, ...]:
+    """Materialize validated per-layer execution steps from paired-layer maps."""
+    if type(num_llm_layers) is not int or num_llm_layers <= 0:
+        raise ValueError(
+            f"num_llm_layers must be a positive integer, got {num_llm_layers!r}."
+        )
+
+    def _validate_tsfm_map(mapping: Mapping[int, int], *, role: str) -> dict[int, int]:
+        if not isinstance(mapping, Mapping):
+            raise TypeError(
+                f"{role}_tsfm_by_llm_layer must be a mapping, got "
+                f"{type(mapping).__name__}."
+            )
+        normalized: dict[int, int] = {}
+        for llm_layer, tsfm_layer in mapping.items():
+            if type(llm_layer) is not int or not 0 <= llm_layer < num_llm_layers:
+                raise ValueError(
+                    f"{role} LLM layer must be an integer in [0, {num_llm_layers}), "
+                    f"got {llm_layer!r}."
+                )
+            if type(tsfm_layer) is not int or tsfm_layer < 0:
+                raise ValueError(
+                    f"{role} TimesFM layer must be a non-negative integer, got "
+                    f"{tsfm_layer!r}."
+                )
+            normalized[llm_layer] = tsfm_layer
+
+        observed_tsfm_layers = [
+            normalized[llm_layer] for llm_layer in sorted(normalized)
+        ]
+        expected_tsfm_layers = list(range(len(normalized)))
+        if observed_tsfm_layers != expected_tsfm_layers:
+            raise ValueError(
+                f"{role} TimesFM layers must advance in contiguous 0..N-1 order, "
+                f"got {observed_tsfm_layers}."
+            )
+        return normalized
+
+    generation = _validate_tsfm_map(generation_tsfm_by_llm_layer, role="generation")
+    if not generation:
+        raise ValueError("generation_tsfm_by_llm_layer must not be empty.")
+    understanding = _validate_tsfm_map(
+        understanding_tsfm_by_llm_layer, role="understanding"
+    )
+    return tuple(
+        TimeBraidLayerStep(
+            llm_layer=llm_layer,
+            generation_tsfm_layer=generation.get(llm_layer),
+            understanding_tsfm_layer=understanding.get(llm_layer),
+        )
+        for llm_layer in range(num_llm_layers)
+    )
+
+
+class MoTForecastLossAccounting(NamedTuple):
+    record_target_idx: torch.Tensor
+    record_patch_idx: torch.Tensor
+    target_starts: torch.Tensor
+    target_lens: torch.Tensor
+    raw_indices: torch.Tensor
+    keep_mask: torch.Tensor
+    valid_counts: torch.Tensor
+
+
+def build_mot_forecast_loss_accounting(
+    *,
+    patch_lengths: torch.Tensor,
+    patch_counts: torch.Tensor,
+    value_lengths: torch.Tensor,
+    loss_start_idxs: torch.Tensor,
+    forecast_horizon: int,
+) -> Optional[MoTForecastLossAccounting]:
+    """
+    Build the canonical TS forecast record mask.
+
+    `loss_start` is both a point-level supervision boundary and an owner-patch
+    eligibility boundary. A historical owner patch whose raw end is before
+    `loss_start` must not receive loss just because a long horizon overlaps the
+    future suffix.
+    """
+    if patch_lengths.ndim != 2:
+        raise RuntimeError(
+            f"TS loss accounting expects patch_lengths [target,patch], got {tuple(patch_lengths.shape)}."
+        )
+    target_count, max_patches = int(patch_lengths.shape[0]), int(patch_lengths.shape[1])
+    for name, tensor in (
+        ("patch_counts", patch_counts),
+        ("value_lengths", value_lengths),
+        ("loss_start_idxs", loss_start_idxs),
+    ):
+        if tensor.ndim != 1 or int(tensor.shape[0]) != target_count:
+            raise RuntimeError(
+                f"TS loss accounting expects {name} [{target_count}], got {tuple(tensor.shape)}."
+            )
+    forecast_horizon = int(forecast_horizon)
+    if forecast_horizon <= 0:
+        raise RuntimeError(
+            f"TS loss accounting requires positive forecast_horizon, got {forecast_horizon}."
+        )
+
+    device = patch_lengths.device
+    prefix_ends = torch.cumsum(patch_lengths, dim=1)
+    patch_axis = torch.arange(max_patches, device=device, dtype=torch.long).unsqueeze(0)
+    successor_patch_mask = patch_axis < (patch_counts - 1).unsqueeze(1)
+
+    # Each record is owned by a patch and predicts the raw suffix starting right
+    # after that patch's last raw point.
+    target_starts_by_patch = prefix_ends
+    next_available_lengths = value_lengths[:, None] - prefix_ends
+    forecast_horizon_tensor = torch.full_like(prefix_ends, forecast_horizon)
+    target_lens_by_patch = torch.minimum(
+        forecast_horizon_tensor, next_available_lengths
+    )
+    keep_starts_by_patch = torch.clamp(
+        loss_start_idxs[:, None] - target_starts_by_patch, min=0
+    )
+
+    # An owner patch before the history/future boundary is context only. The
+    # boundary owner itself is the first record allowed to predict supervised
+    # future points.
+    owner_reaches_loss_start = prefix_ends >= loss_start_idxs[:, None]
+    base_record_mask = torch.logical_and(successor_patch_mask, owner_reaches_loss_start)
+    record_mask = torch.logical_and(
+        base_record_mask,
+        torch.logical_and(
+            target_lens_by_patch > 0, keep_starts_by_patch < target_lens_by_patch
+        ),
+    )
+    record_target_idx, record_patch_idx = torch.nonzero(record_mask, as_tuple=True)
+    if int(record_target_idx.numel()) == 0:
+        return None
+
+    target_lens = target_lens_by_patch[record_target_idx, record_patch_idx]
+    target_starts = target_starts_by_patch[record_target_idx, record_patch_idx]
+    horizon_idx = torch.arange(forecast_horizon, device=device, dtype=torch.long)
+    raw_indices = target_starts.unsqueeze(1) + horizon_idx.unsqueeze(0)
+    keep_mask = torch.logical_and(
+        horizon_idx.unsqueeze(0) < target_lens.unsqueeze(1),
+        raw_indices >= loss_start_idxs[record_target_idx].unsqueeze(1),
+    )
+    nonempty_record_mask = torch.any(keep_mask, dim=1)
+    if not bool(torch.any(nonempty_record_mask)):
+        return None
+
+    keep_mask = keep_mask[nonempty_record_mask]
+    target_lens = target_lens[nonempty_record_mask]
+    target_starts = target_starts[nonempty_record_mask]
+    record_target_idx = record_target_idx[nonempty_record_mask]
+    record_patch_idx = record_patch_idx[nonempty_record_mask]
+    raw_indices = raw_indices[nonempty_record_mask]
+    valid_counts = keep_mask.to(dtype=torch.long).sum(dim=1)
+    if bool(torch.any(valid_counts <= 0)):
+        raise RuntimeError(
+            "TS loss accounting built a record with no supervised points."
+        )
+    return MoTForecastLossAccounting(
+        record_target_idx=record_target_idx,
+        record_patch_idx=record_patch_idx,
+        target_starts=target_starts,
+        target_lens=target_lens,
+        raw_indices=raw_indices,
+        keep_mask=keep_mask,
+        valid_counts=valid_counts,
+    )
+
+
+@dataclass
+class _Span:
+    start_token_idx: int
+    end_token_idx: int

timebraid/model/timebraid.py

@@ -0,0 +1,1021 @@
+"""HF-native TimeBraid model with directly registered language and TS components."""
+
+from __future__ import annotations
+
+from collections.abc import Iterable, Mapping, Sequence
+from dataclasses import dataclass, fields
+from typing import TYPE_CHECKING, Any, Optional, Union
+
+import torch
+import torch.nn as nn
+from transformers import (
+    AutoConfig,
+    AutoModelForCausalLM,
+    PretrainedConfig,
+    PreTrainedModel,
+)
+from transformers.cache_utils import Cache
+from transformers.modeling_outputs import CausalLMOutputWithPast
+from transformers.modeling_utils import get_parameter_dtype
+from transformers.utils import ModelOutput
+
+from .mot.bridge_runtime import initialize_timebraid_components
+from .mot.bridge_runtime_mot import sync_gradient_checkpointing_from_decoder_layers
+from .mot.config_contract import (
+    MOT_HF_CONFIG_CONTRACT_FIELDS,
+    collect_mot_hf_config_contract,
+)
+from .mot.generation_route import (
+    GenerationRoute,
+    require_mixed_generation_inputs,
+    resolve_generation_route,
+)
+from .mot.kv_cache import MoTDynamicCache
+from .mot.model import (
+    NormalizedMoTGenerationControls,
+    TimeBraidGenerateOutput,
+    TimeBraidPayload,
+    merge_timebraid_losses,
+    normalize_mot_generation_controls,
+    run_timebraid_decoder,
+    run_timebraid_generate,
+    validate_mot_generation_input_ids,
+)
+from .mot.structures import IGNORE_INDEX, TIMEBRAID_TS_PAYLOAD_FIELDS
+
+if TYPE_CHECKING:  # pragma: no cover - import cycle only matters at runtime
+    from ..processing_timebraid import TimeBraidBatchFeature
+
+_FINITE_SCAN_CHUNK_ELEMENTS = 8 * 1024 * 1024
+_NON_INFERENCE_CONFIG_FIELDS = (
+    "mot_timesfm_model_name_or_path",
+    "mot_ts_roi_mse_alpha",
+    "mot_ts_loss_weight",
+    "mot_lm_loss_weight",
+    "mot_ts_understanding_loss_weight",
+)
+_FORBIDDEN_CONFIG_FIELDS = _NON_INFERENCE_CONFIG_FIELDS
+_MOT_GENERATION_CONTROL_FIELDS = (
+    "mot_target_horizons",
+    "mot_forecast_head_len",
+    "mot_target_total_lengths",
+    "mot_target_history_span_idxs",
+    "mot_return_forecast_quantiles",
+)
+
+
+@dataclass
+class TimeBraidLossBreakdown(ModelOutput):
+    """Fixed scalar loss accounting returned beside the differentiable loss."""
+
+    loss: Optional[torch.Tensor] = None
+    lm_loss: Optional[torch.Tensor] = None
+    ts_aux_loss: Optional[torch.Tensor] = None
+    ts_aux_point_loss: Optional[torch.Tensor] = None
+    ts_aux_quantile_loss: Optional[torch.Tensor] = None
+    ts_aux_global_point_loss: Optional[torch.Tensor] = None
+    ts_aux_global_quantile_loss: Optional[torch.Tensor] = None
+    ts_aux_global_uncapped_loss: Optional[torch.Tensor] = None
+    ts_aux_global_cap_scale: Optional[torch.Tensor] = None
+    ts_aux_roi_loss: Optional[torch.Tensor] = None
+    ts_aux_roi_point_loss: Optional[torch.Tensor] = None
+    ts_aux_roi_weighted_point_loss: Optional[torch.Tensor] = None
+    ts_aux_roi_uncapped_loss: Optional[torch.Tensor] = None
+    ts_aux_roi_weighted_uncapped_loss: Optional[torch.Tensor] = None
+    ts_aux_roi_cap_scale: Optional[torch.Tensor] = None
+    ts_aux_roi_supervised_points: Optional[torch.Tensor] = None
+    ts_aux_roi_supervised_targets: Optional[torch.Tensor] = None
+    ts_aux_roi_alpha: Optional[torch.Tensor] = None
+    ts_aux_understanding_loss: Optional[torch.Tensor] = None
+    ts_aux_understanding_weighted_loss: Optional[torch.Tensor] = None
+    ts_aux_understanding_time_loss: Optional[torch.Tensor] = None
+    ts_aux_understanding_fft_loss: Optional[torch.Tensor] = None
+    ts_aux_understanding_patches: Optional[torch.Tensor] = None
+    ts_aux_understanding_points: Optional[torch.Tensor] = None
+    ts_aux_understanding_weight: Optional[torch.Tensor] = None
+
+
+@dataclass
+class TimeBraidOutput(CausalLMOutputWithPast):
+    """Stable causal-LM output for every dictionary-returning TimeBraid forward."""
+
+    token_accuracy: Optional[torch.Tensor] = None
+    predicted_tokens: Optional[torch.Tensor] = None
+    loss_breakdown: Optional[TimeBraidLossBreakdown] = None
+
+
+def _present_forbidden_config_fields(config: object) -> list[str]:
+    if isinstance(config, Mapping):
+        return sorted(field for field in _FORBIDDEN_CONFIG_FIELDS if field in config)
+    return sorted(field for field in _FORBIDDEN_CONFIG_FIELDS if hasattr(config, field))
+
+
+def _validate_finite_named_tensors(
+    named_tensors: Iterable[tuple[str, torch.Tensor]],
+    *,
+    owner: str,
+) -> None:
+    for name, tensor in named_tensors:
+        if not tensor.is_floating_point() and not tensor.is_complex():
+            continue
+        flat = tensor.detach().contiguous().view(-1)
+        for start in range(0, int(flat.numel()), _FINITE_SCAN_CHUNK_ELEMENTS):
+            chunk = flat[start : start + _FINITE_SCAN_CHUNK_ELEMENTS]
+            finite_mask = torch.isfinite(chunk)
+            if bool(torch.all(finite_mask)):
+                continue
+            chunk_index = int(
+                torch.nonzero(torch.logical_not(finite_mask), as_tuple=False)[0]
+                .detach()
+                .cpu()
+                .item()
+            )
+            raise RuntimeError(
+                f"{owner} contains a non-finite tensor value: "
+                f"name={name!r}, flat_index={start + chunk_index}, dtype={tensor.dtype}."
+            )
+
+
+def _validate_finite_model_tensors(model: nn.Module, *, owner: str) -> None:
+    """Reject non-finite floating or complex model state before use."""
+    _validate_finite_named_tensors(
+        (*model.named_parameters(), *model.named_buffers()),
+        owner=owner,
+    )
+
+
+def _runtime_options_from_timebraid_config(config: "TimeBraidConfig") -> dict[str, Any]:
+    """Build direct-component initialization options from checkpoint metadata.
+
+    A contract field's name *is* its runtime option name, so this copies
+    rather than translates. The hand-maintained rename table it replaced was
+    what allowed one knob to carry two spellings and drift apart.
+    """
+    contract = collect_mot_hf_config_contract(config)
+    options: dict[str, Any] = {
+        name: contract[name] for name in MOT_HF_CONFIG_CONTRACT_FIELDS
+    }
+    # Never restored from a checkpoint: inference always runs with gradient
+    # checkpointing off.
+    options["mot_ts_gradient_checkpointing"] = False
+    options["mot_ts_gradient_checkpointing_use_reentrant"] = False
+    compute_dtype = getattr(config, "dtype", None) or getattr(
+        config.llm_config, "dtype", None
+    )
+    if compute_dtype is not None:
+        options["mot_compute_dtype"] = compute_dtype
+    return options
+
+
+_LLM_MIRROR_FIELDS = (
+    "vocab_size",
+    "hidden_size",
+    "intermediate_size",
+    "num_hidden_layers",
+    "num_attention_heads",
+    "num_key_value_heads",
+    "head_dim",
+    "max_position_embeddings",
+    "tie_word_embeddings",
+    "dtype",
+    "bos_token_id",
+    "eos_token_id",
+    "pad_token_id",
+    "use_cache",
+    "sliding_window",
+    "layer_types",
+    "_attn_implementation",
+    "_attn_implementation_internal",
+)
+_LLM_SERIALIZED_MIRROR_FIELDS = tuple(
+    field for field in _LLM_MIRROR_FIELDS if not field.startswith("_")
+)
+
+
+class TimeBraidConfig(PretrainedConfig):
+    r"""Composite config for TimeBraid's language and time-series components."""
+
+    model_type = "timebraid"
+    is_composition = True
+    has_no_defaults_at_init = True
+    sub_configs = {"llm_config": AutoConfig}
+
+    def __init__(
+        self,
+        llm_config: Optional[PretrainedConfig | dict[str, Any]] = None,
+        **kwargs: Any,
+    ) -> None:
+        if llm_config is None:
+            raise ValueError("TimeBraidConfig requires `llm_config`.")
+        forbidden_fields = _present_forbidden_config_fields(kwargs)
+        if forbidden_fields:
+            raise ValueError(
+                "TimeBraidConfig does not restore training-only or external "
+                f"TimesFM fields: {forbidden_fields}."
+            )
+        nested_forbidden_fields = _present_forbidden_config_fields(llm_config)
+        if nested_forbidden_fields:
+            raise ValueError(
+                "TimeBraidConfig llm_config does not restore training-only or "
+                f"external TimesFM fields: {nested_forbidden_fields}."
+            )
+
+        if isinstance(llm_config, PretrainedConfig):
+            resolved_llm_config = llm_config
+        else:
+            llm_config_dict = dict(llm_config)
+            llm_model_type = llm_config_dict.pop("model_type", None)
+            if llm_model_type is None:
+                raise ValueError(
+                    "TimeBraidConfig `llm_config` must include `model_type`."
+                )
+            resolved_llm_config = AutoConfig.for_model(
+                llm_model_type, **llm_config_dict
+            )
+
+        outer_mirrors = sorted(
+            field for field in _LLM_SERIALIZED_MIRROR_FIELDS if field in kwargs
+        )
+        if outer_mirrors:
+            raise ValueError(
+                "TimeBraidConfig stores LLM fields only in llm_config; remove outer fields "
+                f"{outer_mirrors}."
+            )
+        kwargs.setdefault("architectures", ["TimeBraid"])
+        # PretrainedConfig applies dtype and attention overrides during its
+        # initializer. Register the nested config first so those composite
+        # controls reach Qwen instead of remaining only on the outer wrapper.
+        self.llm_config = resolved_llm_config
+        super().__init__(**kwargs)
+        self._sync_llm_surface()
+
+    @classmethod
+    def from_llm_config(
+        cls, llm_config: PretrainedConfig, **kwargs: Any
+    ) -> "TimeBraidConfig":
+        return cls(llm_config=llm_config, **kwargs)
+
+    def _sync_llm_surface(self) -> None:
+        """Derive the non-serialized wrapper surface generic HF helpers consume."""
+        for field in _LLM_MIRROR_FIELDS:
+            if hasattr(self.llm_config, field):
+                setattr(self, field, getattr(self.llm_config, field))
+
+    def to_dict(self) -> dict[str, Any]:
+        """Serialize only inference restore state, never training objectives."""
+        payload = super().to_dict()
+        forbidden_fields = _present_forbidden_config_fields(payload)
+        if forbidden_fields:
+            raise RuntimeError(
+                "TimeBraidConfig contains forbidden training-only or external "
+                f"TimesFM fields: {forbidden_fields}."
+            )
+        for field in _LLM_SERIALIZED_MIRROR_FIELDS:
+            payload.pop(field, None)
+        nested = payload.get("llm_config")
+        if isinstance(nested, dict):
+            nested_forbidden_fields = _present_forbidden_config_fields(nested)
+            if nested_forbidden_fields:
+                raise RuntimeError(
+                    "TimeBraidConfig llm_config contains forbidden training-only, "
+                    f"or external TimesFM fields: {nested_forbidden_fields}."
+                )
+        return payload
+
+
+def _labels_with_language_supervision(
+    labels: Optional[torch.LongTensor],
+) -> Optional[torch.LongTensor]:
+    if isinstance(labels, torch.Tensor) and (
+        labels.numel() == 0 or not bool(torch.any(labels.ne(IGNORE_INDEX)).item())
+    ):
+        return None
+    return labels
+
+
+def _require_binary_plain_attention_mask(attention_mask: object) -> None:
+    """Reject packed segment ids before delegating a plain call to native Qwen."""
+    # Rank-4 prepared masks remain part of Qwen's public ABI. TimeBraid packing
+    # uses positive segment ids only in the rank-2 mask produced by its
+    # processor/collator, so the plain route must accept exactly binary values
+    # there instead of letting Qwen collapse every nonzero id to one segment.
+    if not isinstance(attention_mask, torch.Tensor) or attention_mask.ndim != 2:
+        return
+    binary_values = torch.logical_or(attention_mask.eq(0), attention_mask.eq(1))
+    if not bool(torch.all(binary_values).item()):
+        raise RuntimeError(
+            "Plain Qwen fallback requires a binary 2D attention_mask with values "
+            "in {0, 1}; packed segment ids require the TimeBraid processor/collator "
+            f"payload, got dtype={attention_mask.dtype}, shape={tuple(attention_mask.shape)}."
+        )
+
+
+_TIMEBRAID_KWARG_PREFIXES = ("mot_", "ts_")
+_KNOWN_TIMEBRAID_KWARGS = frozenset(
+    TIMEBRAID_TS_PAYLOAD_FIELDS + _MOT_GENERATION_CONTROL_FIELDS
+)
+
+
+def _reject_unknown_timebraid_kwargs(generation_kwargs: Mapping[str, object]) -> None:
+    """Reject misspelled TimeBraid arguments instead of forwarding them.
+
+    `mot_` and `ts_` are TimeBraid's own argument namespace, so an unrecognized
+    name there is a mistake we can name. Unprefixed names belong to Hugging
+    Face and must keep flowing through, or this package would break whenever
+    transformers adds a generation argument.
+
+    Without this, a typo such as `mot_target_horizon` was silently forwarded,
+    the request routed to plain text, and the caller received prose with no
+    forecast and no error.
+    """
+    unknown = sorted(
+        name
+        for name in generation_kwargs
+        if name.startswith(_TIMEBRAID_KWARG_PREFIXES)
+        and name not in _KNOWN_TIMEBRAID_KWARGS
+    )
+    if unknown:
+        raise TypeError(
+            f"Unknown TimeBraid generation argument(s) {unknown}. Supported "
+            f"names are {sorted(_KNOWN_TIMEBRAID_KWARGS)}."
+        )
+
+
+def _fixed_loss_breakdown(
+    breakdown: Optional[Mapping[str, torch.Tensor]],
+) -> Optional[TimeBraidLossBreakdown]:
+    if breakdown is None:
+        return None
+    reference = next(
+        (value for value in breakdown.values() if isinstance(value, torch.Tensor)),
+        None,
+    )
+    if reference is None:
+        raise RuntimeError("TimeBraid loss accounting returned no tensor values.")
+    zero = reference.detach().new_zeros(())
+    values: dict[str, torch.Tensor] = {}
+    for output_field in fields(TimeBraidLossBreakdown):
+        value = breakdown.get(output_field.name)
+        if value is None:
+            value = zero.clone()
+        if not isinstance(value, torch.Tensor) or value.numel() != 1:
+            raise RuntimeError(
+                "TimeBraid loss accounting values must be scalar tensors: "
+                f"field={output_field.name!r}, value={value!r}."
+            )
+        values[output_field.name] = value.detach().reshape(())
+    return TimeBraidLossBreakdown(**values)
+
+
+def _as_timebraid_output(
+    outputs: object,
+    breakdown: Optional[Mapping[str, torch.Tensor]],
+) -> TimeBraidOutput:
+    return TimeBraidOutput(
+        loss=getattr(outputs, "loss", None),
+        logits=getattr(outputs, "logits", None),
+        past_key_values=getattr(outputs, "past_key_values", None),
+        hidden_states=getattr(outputs, "hidden_states", None),
+        attentions=getattr(outputs, "attentions", None),
+        token_accuracy=getattr(outputs, "token_accuracy", None),
+        predicted_tokens=getattr(outputs, "predicted_tokens", None),
+        loss_breakdown=_fixed_loss_breakdown(breakdown),
+    )
+
+
+def _build_lm_output(
+    *,
+    llm: PreTrainedModel,
+    hidden_states: torch.Tensor,
+    past_key_values: object,
+    decoder_hidden_states: object,
+    decoder_attentions: object,
+    labels: Optional[torch.LongTensor],
+    labels_for_loss: Optional[torch.LongTensor],
+    shift_labels: Optional[torch.LongTensor],
+    logits_to_keep: Union[int, torch.Tensor],
+    skip_logits: Optional[bool],
+    loss_kwargs: Mapping[str, object],
+) -> TimeBraidOutput:
+    """Project decoder hidden states and own the complete language-loss contract."""
+    if skip_logits is not None and type(skip_logits) is not bool:
+        raise TypeError(
+            f"TimeBraid forward skip_logits must be a bool or None, got {skip_logits!r}."
+        )
+
+    # `skip_logits=True` only ever meant "the fused Liger loss will consume the
+    # hidden states instead of logits". That path was unreachable — the flag
+    # enabling it was hardcoded False with no setter — so it is gone, and the
+    # request is refused rather than silently ignored.
+    if skip_logits:
+        raise RuntimeError(
+            "skip_logits=True is not supported: TimeBraid always projects logits "
+            "through the language head."
+        )
+
+    slice_indices = (
+        slice(-logits_to_keep, None)
+        if isinstance(logits_to_keep, int)
+        else logits_to_keep
+    )
+    kept_hidden_states = hidden_states[:, slice_indices, :]
+    has_effective_supervision = labels_for_loss is not None or shift_labels is not None
+    loss = None
+    # Only the removed fused-loss path produced these; the language head does
+    # not, so they stay None and keep TimeBraidOutput's shape stable.
+    token_accuracy = None
+    predicted_tokens = None
+    logits = llm.lm_head(kept_hidden_states)
+    if has_effective_supervision:
+        loss = llm.loss_function(
+            logits=logits,
+            labels=labels_for_loss,
+            shift_labels=shift_labels,
+            vocab_size=llm.config.vocab_size,
+            **loss_kwargs,
+        )
+
+    return TimeBraidOutput(
+        loss=loss,
+        logits=logits,
+        past_key_values=past_key_values,
+        hidden_states=decoder_hidden_states,
+        attentions=decoder_attentions,
+        token_accuracy=token_accuracy,
+        predicted_tokens=predicted_tokens,
+    )
+
+
+class TimeBraid(PreTrainedModel):
+    r"""Top-level model with direct language, time-series, and attention components."""
+
+    # This is the complete trainable/inference component graph. Keeping the
+    # ownership list on the public model makes optional topology readable
+    # without inspecting the numeric execution helpers.
+    llm: PreTrainedModel
+    generation_tsfm: nn.Module
+    understanding_tsfm: Optional[nn.Module]
+    global_residual_attention: nn.ModuleDict
+    packed_attention: nn.Module
+    understanding_head: Optional[nn.Module]
+
+    config_class = TimeBraidConfig
+    base_model_prefix = "llm"
+    main_input_name = "input_ids"
+    _supports_flash_attn_2 = True
+    _supports_sdpa = True
+    supports_gradient_checkpointing = True
+    _tied_weights_keys = {"llm.lm_head.weight": "llm.model.embed_tokens.weight"}
+
+    def __init__(
+        self,
+        config: TimeBraidConfig,
+        llm: Optional[PreTrainedModel] = None,
+        runtime_options: Optional[dict[str, Any]] = None,
+    ) -> None:
+        if not isinstance(config, TimeBraidConfig):
+            raise ValueError(
+                f"TimeBraid requires TimeBraidConfig, got {type(config).__name__}."
+            )
+        super().__init__(config)
+        self.llm = (
+            llm
+            if llm is not None
+            else AutoModelForCausalLM.from_config(config.llm_config)
+        )
+        self.config.llm_config = self.llm.config
+        self.config._sync_llm_surface()
+        self.generation_config = self.llm.generation_config
+
+        resolved_runtime_options = _runtime_options_from_timebraid_config(config)
+        if runtime_options:
+            resolved_runtime_options.update(runtime_options)
+        initialize_timebraid_components(
+            self,
+            self.llm.model.config,
+            resolved_runtime_options,
+        )
+        # Optional training integration may enable the fused hidden-to-loss tail
+        # after construction; that tail does not replace either Qwen forward method.
+        self.generation_tsfm.eval()
+        if self.understanding_tsfm is not None:
+            self.understanding_tsfm.eval()
+        self._sync_hf_loading_contracts()
+
+    def _sync_hf_loading_contracts(self) -> None:
+        """Synchronize tied-weight and parallelism metadata without generic post-init."""
+        self._tied_weights_keys = {
+            "llm.lm_head.weight": "llm.model.embed_tokens.weight"
+        }
+        self.all_tied_weights_keys = dict(self._tied_weights_keys)
+        tied_keys = getattr(self.llm, "all_tied_weights_keys", None)
+        if tied_keys:
+            self.all_tied_weights_keys.update(
+                {f"llm.{key}": f"llm.{value}" for key, value in tied_keys.items()}
+            )
+        # Transformers expects composite wrappers to expose concrete empty
+        # parallelism plans during distributed `from_pretrained()` allocator
+        # warmup, so mirror the child plans here.
+        for plan_attr in ("_tp_plan", "_ep_plan", "_pp_plan"):
+            merged_plan = {}
+            child_plan = getattr(self.llm, plan_attr, None)
+            if child_plan:
+                merged_plan.update(
+                    {f"llm.{key}": value for key, value in child_plan.copy().items()}
+                )
+            setattr(self, plan_attr, merged_plan)
+
+    def train(self, mode: bool = True) -> "TimeBraid":
+        super().train(mode)
+        # TimesFM towers stay in eval mode during full-parameter training. Module
+        # mode controls stochastic behavior independently of `requires_grad`.
+        self.generation_tsfm.eval()
+        if self.understanding_tsfm is not None:
+            self.understanding_tsfm.eval()
+        return self
+
+    def gradient_checkpointing_enable(
+        self, gradient_checkpointing_kwargs: Optional[dict[str, Any]] = None
+    ) -> None:
+        super().gradient_checkpointing_enable(
+            gradient_checkpointing_kwargs=gradient_checkpointing_kwargs
+        )
+        sync_gradient_checkpointing_from_decoder_layers(self, self.llm.model.layers)
+
+    def gradient_checkpointing_disable(self) -> None:
+        super().gradient_checkpointing_disable()
+        sync_gradient_checkpointing_from_decoder_layers(self, self.llm.model.layers)
+
+    def forward(
+        self,
+        input_ids: Optional[torch.LongTensor] = None,
+        attention_mask: Optional[torch.Tensor] = None,
+        position_ids: Optional[torch.LongTensor] = None,
+        past_key_values: Optional[Cache] = None,
+        inputs_embeds: Optional[torch.FloatTensor] = None,
+        labels: Optional[torch.LongTensor] = None,
+        use_cache: Optional[bool] = None,
+        cache_position: Optional[torch.LongTensor] = None,
+        logits_to_keep: Union[int, torch.Tensor] = 0,
+        output_attentions: Optional[bool] = None,
+        output_hidden_states: Optional[bool] = None,
+        return_dict: Optional[bool] = None,
+        skip_logits: Optional[bool] = None,
+        shift_labels: Optional[torch.LongTensor] = None,
+        ts_values: object = None,
+        ts_lengths: object = None,
+        ts_loss_start_idxs: object = None,
+        ts_loss_roi_masks: object = None,
+        ts_roles: object = None,
+        ts_segment_ids: object = None,
+        ts_span_mask: object = None,
+        ts_text_start_token_idxs: object = None,
+        ts_text_end_token_idxs: object = None,
+        active_rows: Optional[torch.Tensor] = None,
+        num_items_in_batch: Optional[object] = None,
+        mot_ts_num_records_by_horizon: Optional[object] = None,
+        mot_ts_num_roi_targets_by_horizon: Optional[object] = None,
+        mot_ts_understanding_num_points: Optional[object] = None,
+        mot_ts_understanding_num_fft_elements: Optional[object] = None,
+        **kwargs: Any,
+    ) -> TimeBraidOutput | tuple[Any, ...]:
+        # Generation-only controls can accompany processor batches, but never
+        # participate in the causal forward computation.
+        for key in _MOT_GENERATION_CONTROL_FIELDS:
+            kwargs.pop(key, None)
+
+        payload = TimeBraidPayload(
+            ts_values=ts_values,
+            ts_lengths=ts_lengths,
+            ts_loss_start_idxs=ts_loss_start_idxs,
+            ts_loss_roi_masks=ts_loss_roi_masks,
+            ts_roles=ts_roles,
+            ts_segment_ids=ts_segment_ids,
+            ts_span_mask=ts_span_mask,
+            ts_text_start_token_idxs=ts_text_start_token_idxs,
+            ts_text_end_token_idxs=ts_text_end_token_idxs,
+        )
+        use_timebraid_decoder = (
+            isinstance(past_key_values, MoTDynamicCache) or payload.has_runtime_inputs()
+        )
+        if not use_timebraid_decoder and isinstance(input_ids, torch.Tensor):
+            use_timebraid_decoder = bool(
+                torch.logical_or(
+                    input_ids.eq(int(self.ts_open_token_id)),
+                    input_ids.eq(int(self.ts_close_token_id)),
+                )
+                .any()
+                .item()
+            )
+        if active_rows is not None and not isinstance(past_key_values, MoTDynamicCache):
+            raise RuntimeError("active_rows requires an existing MoTDynamicCache.")
+
+        if return_dict is None:
+            resolved_return_dict = bool(self.llm.config.use_return_dict)
+        elif type(return_dict) is bool:
+            resolved_return_dict = return_dict
+        else:
+            raise TypeError(
+                f"TimeBraid forward return_dict must be a bool or None, got {return_dict!r}."
+            )
+        labels_for_loss = _labels_with_language_supervision(labels)
+        if not use_timebraid_decoder:
+            _require_binary_plain_attention_mask(attention_mask)
+            if skip_logits is not None and type(skip_logits) is not bool:
+                raise TypeError(
+                    "TimeBraid forward skip_logits must be a bool or None, got "
+                    f"{skip_logits!r}."
+                )
+            use_outer_lm_head = shift_labels is not None and labels_for_loss is None
+            if skip_logits and not use_outer_lm_head:
+                raise ValueError(
+                    "skip_logits=True requires fused TimeBraid language loss."
+                )
+            plain_loss_kwargs = dict(kwargs)
+            if num_items_in_batch is not None:
+                plain_loss_kwargs["num_items_in_batch"] = num_items_in_batch
+            plain_forward_kwargs = dict(plain_loss_kwargs)
+            if shift_labels is not None:
+                plain_forward_kwargs["shift_labels"] = shift_labels
+            if use_outer_lm_head:
+                # Run Qwen's ordinary decoder, then let TimeBraid consume its
+                # hidden states. This outer tail exists for explicit
+                # shift-label-only loss, where Qwen cannot build the labels.
+                decoder_outputs = self.llm.model(
+                    input_ids=input_ids,
+                    attention_mask=attention_mask,
+                    position_ids=position_ids,
+                    past_key_values=past_key_values,
+                    inputs_embeds=inputs_embeds,
+                    use_cache=use_cache,
+                    cache_position=cache_position,
+                    output_attentions=output_attentions,
+                    output_hidden_states=output_hidden_states,
+                    return_dict=True,
+                    **kwargs,
+                )
+                outputs = _build_lm_output(
+                    llm=self.llm,
+                    hidden_states=decoder_outputs.last_hidden_state,
+                    past_key_values=decoder_outputs.past_key_values,
+                    decoder_hidden_states=decoder_outputs.hidden_states,
+                    decoder_attentions=decoder_outputs.attentions,
+                    labels=labels,
+                    labels_for_loss=labels_for_loss,
+                    shift_labels=shift_labels,
+                    logits_to_keep=logits_to_keep,
+                    skip_logits=skip_logits,
+                    loss_kwargs=plain_loss_kwargs,
+                )
+            else:
+                outputs = self.llm(
+                    input_ids=input_ids,
+                    attention_mask=attention_mask,
+                    position_ids=position_ids,
+                    past_key_values=past_key_values,
+                    inputs_embeds=inputs_embeds,
+                    labels=labels_for_loss,
+                    use_cache=use_cache,
+                    cache_position=cache_position,
+                    logits_to_keep=logits_to_keep,
+                    output_attentions=output_attentions,
+                    output_hidden_states=output_hidden_states,
+                    return_dict=True,
+                    **plain_forward_kwargs,
+                )
+            outputs, breakdown = merge_timebraid_losses(
+                model=self,
+                mot_runtime=None,
+                outputs=outputs,
+            )
+            timebraid_output = _as_timebraid_output(outputs, breakdown)
+            if resolved_return_dict:
+                return timebraid_output
+            # Preserve Qwen's tuple prefix and ordering. Logits keeps its slot
+            # when fused execution leaves it as None, followed by optional decoder
+            # outputs; TimeBraid's token statistics append afterward.
+            tuple_output = (timebraid_output.logits,)
+            tuple_output += tuple(
+                value
+                for value in (
+                    timebraid_output.past_key_values,
+                    timebraid_output.hidden_states,
+                    timebraid_output.attentions,
+                )
+                if value is not None
+            )
+            if timebraid_output.loss is not None:
+                tuple_output = (timebraid_output.loss,) + tuple_output
+            if timebraid_output.token_accuracy is not None:
+                tuple_output += (timebraid_output.token_accuracy,)
+            if timebraid_output.predicted_tokens is not None:
+                tuple_output += (timebraid_output.predicted_tokens,)
+            return tuple_output
+
+        if not resolved_return_dict:
+            raise RuntimeError("TS-routed TimeBraid forward requires return_dict=True.")
+        cache_implementation = kwargs.pop("cache_implementation", None)
+        if use_cache is None:
+            configured_use_cache = self.llm.model.config.use_cache
+            if type(configured_use_cache) is not bool:
+                raise TypeError(
+                    "Qwen config use_cache must be a bool, got "
+                    f"{configured_use_cache!r}."
+                )
+            resolved_use_cache = configured_use_cache and not self.llm.model.training
+        elif type(use_cache) is bool:
+            resolved_use_cache = use_cache
+        else:
+            raise TypeError(
+                f"TimeBraid forward use_cache must be a bool or None, got {use_cache!r}."
+            )
+        decoder_result = run_timebraid_decoder(
+            owner=self,
+            text_model=self.llm.model,
+            payload=payload,
+            input_ids=input_ids,
+            attention_mask=attention_mask,
+            position_ids=position_ids,
+            past_key_values=past_key_values,
+            inputs_embeds=inputs_embeds,
+            use_cache=resolved_use_cache,
+            cache_position=cache_position,
+            output_attentions=output_attentions,
+            output_hidden_states=output_hidden_states,
+            return_dict=True,
+            active_rows=active_rows,
+            cache_implementation=cache_implementation,
+            collect_incremental_kv_cache=resolved_use_cache,
+        )
+
+        loss_kwargs = dict(kwargs)
+        if num_items_in_batch is not None:
+            loss_kwargs["num_items_in_batch"] = num_items_in_batch
+        outputs = _build_lm_output(
+            llm=self.llm,
+            hidden_states=decoder_result.hidden_states,
+            past_key_values=decoder_result.past_key_values,
+            decoder_hidden_states=None,
+            decoder_attentions=None,
+            labels=labels,
+            labels_for_loss=labels_for_loss,
+            shift_labels=shift_labels,
+            logits_to_keep=logits_to_keep,
+            skip_logits=skip_logits,
+            loss_kwargs=loss_kwargs,
+        )
+        outputs, breakdown = merge_timebraid_losses(
+            model=self,
+            mot_runtime=decoder_result.mot_runtime,
+            outputs=outputs,
+            mot_ts_num_records_by_horizon=mot_ts_num_records_by_horizon,
+            mot_ts_num_roi_targets_by_horizon=mot_ts_num_roi_targets_by_horizon,
+            mot_ts_understanding_num_points=mot_ts_understanding_num_points,
+            mot_ts_understanding_num_fft_elements=(
+                mot_ts_understanding_num_fft_elements
+            ),
+        )
+        return _as_timebraid_output(outputs, breakdown)
+
+    def save_pretrained(self, *args: Any, **kwargs: Any) -> None:
+        # Match Transformers' own dtype authority, but update both composite
+        # config surfaces before the base implementation serializes either one.
+        if self.config.llm_config is not self.llm.config:
+            raise RuntimeError(
+                "TimeBraid save requires config.llm_config to be the live LLM config."
+            )
+        # The checkpoint is self-contained, so machine-local source paths are
+        # neither restore semantics nor useful public provenance.
+        self.config._name_or_path = ""
+        self.config.llm_config._name_or_path = ""
+        self.config.llm_config.dtype = get_parameter_dtype(self)
+        self.config._sync_llm_surface()
+        return super().save_pretrained(*args, **kwargs)
+
+    def prepare_inputs_for_generation(
+        self, *args: Any, **kwargs: Any
+    ) -> dict[str, Any]:
+        return self.llm.prepare_inputs_for_generation(*args, **kwargs)
+
+    def generate_mixed(
+        self,
+        inputs: "TimeBraidBatchFeature",
+        *,
+        max_new_tokens: Optional[int] = None,
+        max_length: Optional[int] = None,
+        mot_target_total_lengths: Optional[Sequence[int]] = None,
+        mot_forecast_head_len: Optional[int] = None,
+        mot_return_forecast_quantiles: bool = False,
+        eos_token_id: Optional[Union[int, Sequence[int]]] = None,
+        pad_token_id: Optional[int] = None,
+        use_cache: Optional[bool] = None,
+        cache_implementation: Optional[str] = None,
+    ) -> TimeBraidGenerateOutput:
+        """Generate mixed text and time series, or raise.
+
+        `inputs` is the mapping returned by
+        `TimeBraidProcessor.apply_chat_template`, which carries `input_ids`,
+        `attention_mask`, the complete TS payload, and the target horizons.
+        Everything the processor cannot know is a keyword argument here.
+
+        This is the explicit counterpart to `generate`. It states its whole
+        contract in its signature, so a misspelled argument is a `TypeError`
+        rather than a silent downgrade, and an argument this scheduler cannot
+        honour — `temperature`, `logits_processor`, anything sampling-related —
+        is simply not a parameter. Unlike `generate`, it never falls back to
+        plain text: if the request does not resolve to mixed generation it
+        raises and says which signal was missing.
+
+        Supply exactly one of `max_new_tokens` or `max_length`.
+
+        TimeBraid's own arguments keep the `mot_` prefix they carry everywhere
+        else — in `config.json`, in the processor's output, and in `generate`'s
+        keyword arguments — so one name means one thing at every boundary.
+        """
+        require_mixed_generation_inputs(inputs)
+        input_ids = validate_mot_generation_input_ids(inputs["input_ids"])
+        batch_size = int(input_ids.shape[0])
+
+        payload = TimeBraidPayload(
+            **{name: inputs.get(name) for name in TIMEBRAID_TS_PAYLOAD_FIELDS}
+        )
+        controls = NormalizedMoTGenerationControls.normalize(
+            batch_size=batch_size,
+            # Supplied by the processor, which owns the prompt surface.
+            target_horizons=inputs.get("mot_target_horizons"),
+            target_history_span_idxs=inputs.get("mot_target_history_span_idxs"),
+            # Supplied by the caller: the processor cannot know these.
+            target_total_lengths=mot_target_total_lengths,
+            forecast_head_len=mot_forecast_head_len,
+            return_forecast_quantiles=mot_return_forecast_quantiles,
+        )
+
+        # No route resolution here: `require_mixed_generation_inputs` above has
+        # already established the one signal that matters, a complete TS
+        # payload. Inferring intent is the compatibility shim's problem.
+        if controls.target_horizons is None:
+            raise RuntimeError(
+                "generate_mixed requires `mot_target_horizons` in `inputs`. Pass "
+                "`horizon=` to TimeBraidProcessor.apply_chat_template, which "
+                "supplies it; use 0 to forbid TS rollout."
+            )
+
+        generation_kwargs: dict[str, Any] = {}
+        for name, value in (
+            ("max_new_tokens", max_new_tokens),
+            ("max_length", max_length),
+            ("eos_token_id", eos_token_id),
+            ("pad_token_id", pad_token_id),
+            ("use_cache", use_cache),
+            ("cache_implementation", cache_implementation),
+        ):
+            if value is not None:
+                generation_kwargs[name] = value
+
+        return run_timebraid_generate(
+            model=self,
+            text_model=self.llm.model,
+            payload=payload,
+            input_ids=input_ids,
+            attention_mask=inputs.get("attention_mask"),
+            controls=controls,
+            kwargs=generation_kwargs,
+        )
+
+    def generate(
+        self, *args: Any, **kwargs: Any
+    ) -> TimeBraidGenerateOutput | torch.Tensor | ModelOutput:
+        """Hugging Face-compatible dispatch over `**kwargs`.
+
+        Kept so `model.generate(**processor_output, max_new_tokens=...)` keeps
+        working, and so plain-text requests still reach the host language
+        model. It resolves a route and delegates; the contract for mixed
+        generation lives on `generate_mixed`.
+        """
+        generation_kwargs = dict(kwargs)
+        payload = TimeBraidPayload.pop_from_kwargs(generation_kwargs)
+        generation_controls = {
+            key: generation_kwargs.pop(key)
+            for key in _MOT_GENERATION_CONTROL_FIELDS
+            if key in generation_kwargs
+        }
+        _reject_unknown_timebraid_kwargs(generation_kwargs)
+
+        positional_input = args[0] if args else None
+        keyword_input = generation_kwargs.get("input_ids")
+        hf_input = generation_kwargs.get("inputs")
+        route_input = (
+            positional_input
+            if positional_input is not None
+            else (keyword_input if keyword_input is not None else hf_input)
+        )
+        if isinstance(route_input, torch.Tensor) and route_input.ndim == 2:
+            control_batch_size = int(route_input.shape[0])
+        elif isinstance(payload.ts_values, torch.Tensor) and payload.ts_values.ndim > 0:
+            control_batch_size = int(payload.ts_values.shape[0])
+        else:
+            control_batch_size = 1
+        controls = normalize_mot_generation_controls(
+            generation_controls,
+            batch_size=control_batch_size,
+        )
+
+        if isinstance(generation_kwargs.get("past_key_values"), MoTDynamicCache):
+            raise RuntimeError(
+                "MoT native mixed generate accepts a fresh prompt and owns its full "
+                "phase scheduler; continue a public forward cache through "
+                "model.forward instead."
+            )
+
+        decision = resolve_generation_route(
+            payload=payload,
+            route_input=route_input,
+            target_horizons=controls.target_horizons,
+            target_total_lengths=controls.target_total_lengths,
+            ts_open_token_id=int(self.ts_open_token_id),
+            ts_close_token_id=int(self.ts_close_token_id),
+        )
+        if decision.route is GenerationRoute.PLAIN_TEXT:
+            _require_binary_plain_attention_mask(
+                generation_kwargs.get("attention_mask")
+            )
+            return self.llm.generate(*args, **generation_kwargs)
+        if controls.target_horizons is None:
+            raise RuntimeError(
+                "Mixed generation requires explicit `mot_target_horizons`. Pass "
+                "`horizon=` to TimeBraidProcessor.apply_chat_template, which "
+                "supplies it; use 0 to forbid TS rollout, or a positive integer "
+                "to allow one generated TS span."
+            )
+
+        if len(args) > 1:
+            raise TypeError(
+                "TS generation accepts at most one positional input tensor."
+            )
+        provided_inputs = [
+            value
+            for value in (positional_input, keyword_input, hf_input)
+            if value is not None
+        ]
+        if len(provided_inputs) != 1:
+            raise ValueError(
+                "TS generation requires exactly one of positional inputs, "
+                "input_ids, or inputs."
+            )
+        input_ids = validate_mot_generation_input_ids(provided_inputs[0])
+        generation_kwargs.pop("input_ids", None)
+        generation_kwargs.pop("inputs", None)
+        attention_mask = generation_kwargs.pop("attention_mask", None)
+        return run_timebraid_generate(
+            model=self,
+            text_model=self.llm.model,
+            payload=payload,
+            input_ids=input_ids,
+            attention_mask=attention_mask,
+            controls=controls,
+            kwargs=generation_kwargs,
+        )
+
+    def get_input_embeddings(self) -> nn.Module:
+        return self.llm.get_input_embeddings()
+
+    def set_input_embeddings(self, value: nn.Module) -> None:
+        self.llm.set_input_embeddings(value)
+
+    def get_output_embeddings(self) -> Optional[nn.Module]:
+        return self.llm.get_output_embeddings()
+
+    def set_output_embeddings(self, value: nn.Module) -> None:
+        self.llm.set_output_embeddings(value)
+
+    def tie_weights(self) -> None:
+        super().tie_weights()
+        self._sync_hf_loading_contracts()
+
+    def can_generate(self) -> bool:
+        llm = self._modules.get("llm") if hasattr(self, "_modules") else None
+        return True if llm is None else llm.can_generate()
+
+
+def register_timebraid_auto_classes() -> None:
+    """Register TimeBraid with HF auto classes in the active Python process."""
+    from transformers import AutoProcessor
+
+    from ..processing_timebraid import TimeBraidProcessor
+
+    AutoConfig.register(TimeBraidConfig.model_type, TimeBraidConfig, exist_ok=True)
+    AutoProcessor.register(TimeBraidConfig, TimeBraidProcessor, exist_ok=True)
+    # Deliberately not registered with plain `AutoModel`: that class means
+    # "base model without a task head", and TimeBraid is a causal LM. The
+    # published auto_map must not advertise it either — see AUTO_MAP in
+    # scripts/build-huggingface-timebraid-model.py.
+    AutoModelForCausalLM.register(TimeBraidConfig, TimeBraid, exist_ok=True)
+
+
+__all__ = [
+    "TimeBraid",
+    "TimeBraidConfig",
+    "TimeBraidLossBreakdown",
+    "TimeBraidOutput",
+    "register_timebraid_auto_classes",
+]

timebraid/processing_timebraid.py

@@ -0,0 +1,726 @@
+"""Hugging Face processor for TimeBraid text and time-series generation."""
+
+from __future__ import annotations
+
+import math
+from collections.abc import Mapping, Sequence
+from dataclasses import dataclass
+from typing import Any
+
+import torch
+from transformers import ProcessorMixin
+from transformers.feature_extraction_utils import BatchFeature
+
+from .data.mot_utils import (
+    ROLE_TO_ID,
+    build_mot_batch_timeseries_tensors,
+    normalize_timeseries_spans,
+)
+from .model.mot.generation_route import FinishReason
+
+_ASSISTANT_PREFIX = "<|im_start|>assistant\n<think>\n\n</think>\n\n"
+_ASSISTANT_SUFFIX = "<|im_end|>"
+_NO_THINK_SUFFIX = "/no_think"
+_STAT_SIGNIFICANT_DIGITS = 6
+_STAT_DECIMAL_LOWER_BOUND = 1.0e-4
+_STAT_DECIMAL_UPPER_BOUND = 1.0e6
+_QWEN_STOP_TOKENS = ("<|im_end|>", "<|endoftext|>")
+
+
+def _coerce_finite_float(value: float, *, name: str) -> float:
+    coerced = float(value)
+    if not math.isfinite(coerced):
+        raise ValueError(f"{name} must be finite, got {value!r}")
+    return coerced
+
+
+def _normalize_scientific_notation(text: str) -> str:
+    mantissa, exponent = text.split("e")
+    mantissa = mantissa.rstrip("0").rstrip(".")
+    if mantissa in {"", "-0"}:
+        return "0"
+    sign = ""
+    if exponent.startswith("-"):
+        sign = "-"
+    digits = exponent.lstrip("+-0") or "0"
+    return f"{mantissa}e{sign}{digits}"
+
+
+def _strip_fixed_decimal(text: str) -> str:
+    stripped = text.rstrip("0").rstrip(".") if "." in text else text
+    return "0" if stripped in {"", "-0"} else stripped
+
+
+def format_stat(
+    value: float,
+    significant_digits: int = _STAT_SIGNIFICANT_DIGITS,
+    decimal_lower_bound: float = _STAT_DECIMAL_LOWER_BOUND,
+    decimal_upper_bound: float = _STAT_DECIMAL_UPPER_BOUND,
+) -> str:
+    """Format one prompt-visible normalization statistic."""
+    if significant_digits <= 0:
+        raise ValueError(
+            f"significant_digits must be positive, got {significant_digits}"
+        )
+    if decimal_lower_bound <= 0:
+        raise ValueError(
+            f"decimal_lower_bound must be positive, got {decimal_lower_bound}"
+        )
+    if decimal_upper_bound <= decimal_lower_bound:
+        raise ValueError(
+            "decimal_upper_bound must be greater than decimal_lower_bound, got "
+            f"{decimal_upper_bound} <= {decimal_lower_bound}"
+        )
+    coerced = _coerce_finite_float(value, name="value")
+    if coerced == 0.0:
+        return "0"
+    magnitude = abs(coerced)
+    if decimal_lower_bound <= magnitude < decimal_upper_bound:
+        exponent = math.floor(math.log10(magnitude))
+        decimals = max(0, significant_digits - 1 - exponent)
+        rendered = _strip_fixed_decimal(f"{coerced:.{decimals}f}")
+    else:
+        rendered = _normalize_scientific_notation(
+            f"{coerced:.{significant_digits - 1}e}"
+        )
+    parsed = float(rendered)
+    if rendered in {"", "-0"} or (coerced != 0.0 and parsed == 0.0):
+        raise ValueError(
+            "Formatted normalization statistic lost its nonzero magnitude: "
+            f"value={coerced!r}, rendered={rendered!r}."
+        )
+    relative_error = abs(parsed - coerced) / abs(coerced)
+    if relative_error > 10.0 ** (1 - significant_digits):
+        raise ValueError(
+            "Formatted normalization statistic lost too much precision: "
+            f"value={coerced!r}, rendered={rendered!r}."
+        )
+    return rendered
+
+
+def compute_zscore_stats(
+    values: Sequence[float], eps: float = 1.0e-6
+) -> tuple[float, float]:
+    """Return population mean and epsilon-floored population standard deviation."""
+    if not values:
+        raise ValueError("Time-series values must be non-empty.")
+    mean = sum(values) / float(len(values))
+    variance = sum((value - mean) ** 2 for value in values) / float(len(values))
+    std = max(variance**0.5, float(eps))
+    if not math.isfinite(mean) or not math.isfinite(std):
+        raise ValueError("Time-series z-score statistics must be finite.")
+    return float(mean), float(std)
+
+
+def zscore_with_stats(values: Sequence[float], mean: float, std: float) -> list[float]:
+    """Normalize finite values with explicit statistics."""
+    if not math.isfinite(std) or std <= 0.0:
+        raise ValueError(f"std must be finite and positive, got {std!r}.")
+    normalized = [(float(value) - mean) / std for value in values]
+    if any(not math.isfinite(value) for value in normalized):
+        raise ValueError("Time-series z-score values must be finite.")
+    return normalized
+
+
+def build_inline_named_zscore_span_reference(
+    *, length_tag: int, mean: float, std: float, precision: int | None = None
+) -> str:
+    """Render the canonical stats block and generic TS delimiter pair."""
+    del precision
+    if type(length_tag) is not int or length_tag <= 0:
+        raise ValueError(f"length_tag must be a positive integer, got {length_tag!r}.")
+    return (
+        f"<stats>len={length_tag}, mean={format_stat(mean)}, "
+        f"std={format_stat(std)}</stats> <ts></ts>"
+    )
+
+
+def format_qwen_chat_turns(turns: Sequence[tuple[str, str]]) -> str:
+    """Render the canonical Qwen transcript used by TimeBraid."""
+    if not turns:
+        raise ValueError("Chat turns must be non-empty.")
+    rendered: list[str] = []
+    for index, (role, content) in enumerate(turns):
+        if role not in {"system", "user", "assistant"}:
+            raise ValueError(f"Unsupported chat role {role!r} at index {index}.")
+        allow_empty_scaffold = role == "assistant" and index == len(turns) - 1
+        if not content and not allow_empty_scaffold:
+            raise ValueError(f"Chat content must be non-empty at index {index}.")
+        if role == "assistant":
+            rendered.append(
+                f"<|im_start|>assistant\n<think>\n\n</think>\n\n{content}<|im_end|>\n"
+            )
+        else:
+            rendered.append(f"<|im_start|>{role}\n{content}<|im_end|>\n")
+    return "".join(rendered)
+
+
+def _qwen_stop_token_ids(tokenizer: Any) -> list[int]:
+    """Resolve only the canonical single-token Qwen generation stops."""
+    resolved: list[int] = []
+    eos_token_id = getattr(tokenizer, "eos_token_id", None)
+    if type(eos_token_id) is int and eos_token_id >= 0:
+        resolved.append(eos_token_id)
+    for token in _QWEN_STOP_TOKENS:
+        token_ids = tokenizer(token, add_special_tokens=False)["input_ids"]
+        if not isinstance(token_ids, list) or len(token_ids) != 1:
+            continue
+        token_id = int(token_ids[0])
+        if tokenizer.unk_token_id is not None and token_id == int(
+            tokenizer.unk_token_id
+        ):
+            continue
+        if token_id not in resolved:
+            resolved.append(token_id)
+    return resolved
+
+
+@dataclass(frozen=True, slots=True)
+class TimeBraidRequestContext:
+    """What `post_process_generation` needs that the tensors do not carry.
+
+    Typed rather than a bare dict because it travels beside the batch: a
+    serving worker that hands the tensors to another process needs to know
+    exactly what else must go with them, and a 9-key dict documented nowhere
+    could not tell it.
+    """
+
+    prompt_text: str
+    horizon: int | None
+    target_series_index: int
+    normalization: Any
+    prompt_width: int
+    prompt_tokens: int
+    num_timeseries_spans: int
+    eos_token_ids: Sequence[int]
+    pad_token_id: int | None
+
+
+class TimeBraidBatchFeature(BatchFeature):
+    """Tensor inputs plus the non-mapping context needed for decoding.
+
+    The context rides as an attribute rather than as batch data because it is
+    not a tensor. `BatchFeature.to()` returns self, so the documented
+    processor -> generate -> post_process flow keeps it; `{**inputs}` and any
+    cross-process hand-off do not. Pass `context=` explicitly there.
+    """
+
+    def __init__(
+        self, data: dict[str, Any], *, postprocess_context: TimeBraidRequestContext
+    ):
+        if not isinstance(postprocess_context, TimeBraidRequestContext):
+            raise TypeError(
+                "postprocess_context must be a TimeBraidRequestContext, got "
+                f"{type(postprocess_context).__name__}."
+            )
+        super().__init__(data=data)
+        self.postprocess_context = postprocess_context
+
+
+class TimeBraidProcessor(ProcessorMixin):
+    """Prepare one TimeBraid completion and decode its mixed generation output."""
+
+    attributes = ["tokenizer"]
+    tokenizer_class = "AutoTokenizer"
+
+    def __init__(
+        self,
+        tokenizer,
+        max_spans_per_sample: int = 64,
+        normalization_epsilon: float = 1.0e-6,
+    ) -> None:
+        if type(max_spans_per_sample) is not int or max_spans_per_sample <= 0:
+            raise ValueError(
+                "max_spans_per_sample must be a positive integer, got "
+                f"{max_spans_per_sample!r}."
+            )
+        if (
+            isinstance(normalization_epsilon, bool)
+            or not isinstance(normalization_epsilon, (int, float))
+            or not math.isfinite(float(normalization_epsilon))
+            or float(normalization_epsilon) <= 0.0
+        ):
+            raise ValueError(
+                "normalization_epsilon must be finite and positive, got "
+                f"{normalization_epsilon!r}."
+            )
+        self.max_spans_per_sample = max_spans_per_sample
+        self.normalization_epsilon = float(normalization_epsilon)
+        super().__init__(tokenizer)
+        self._validate_delimiters()
+
+    @property
+    def model_input_names(self) -> list[str]:
+        return [
+            "input_ids",
+            "attention_mask",
+            "ts_values",
+            "ts_lengths",
+            "ts_loss_start_idxs",
+            "ts_loss_roi_masks",
+            "ts_roles",
+            "ts_segment_ids",
+            "ts_span_mask",
+            "ts_text_start_token_idxs",
+            "ts_text_end_token_idxs",
+            "mot_target_horizons",
+            "mot_target_history_span_idxs",
+        ]
+
+    def _validate_delimiters(self) -> None:
+        resolved: list[int] = []
+        for delimiter in ("<ts>", "</ts>"):
+            token_ids = self.tokenizer(delimiter, add_special_tokens=False)["input_ids"]
+            if not isinstance(token_ids, list) or len(token_ids) != 1:
+                raise ValueError(
+                    f"TimeBraid tokenizer must encode {delimiter!r} as one token, got {token_ids!r}."
+                )
+            token_id = int(token_ids[0])
+            if self.tokenizer.unk_token_id is not None and token_id == int(
+                self.tokenizer.unk_token_id
+            ):
+                raise ValueError(
+                    f"TimeBraid tokenizer resolves {delimiter!r} to unk_token_id={token_id}."
+                )
+            resolved.append(token_id)
+        if resolved[0] == resolved[1]:
+            raise ValueError("TimeBraid TS delimiters must use distinct token IDs.")
+
+    def __call__(
+        self,
+        *,
+        messages: Sequence[Mapping[str, str]],
+        timeseries: Sequence[Sequence[float]] | None = None,
+        horizon: int | None = None,
+        target_series_index: int | None = None,
+        return_tensors: str = "pt",
+    ) -> TimeBraidBatchFeature:
+        return self.apply_chat_template(
+            messages,
+            timeseries=timeseries,
+            horizon=horizon,
+            target_series_index=target_series_index,
+            add_generation_prompt=True,
+            tokenize=True,
+            return_dict=True,
+            return_tensors=return_tensors,
+        )
+
+    def apply_chat_template(
+        self,
+        messages: Sequence[Mapping[str, str]],
+        *,
+        timeseries: Sequence[Sequence[float]] | None = None,
+        horizon: int | None = None,
+        target_series_index: int | None = None,
+        add_generation_prompt: bool = True,
+        tokenize: bool = True,
+        return_dict: bool = True,
+        return_tensors: str = "pt",
+    ) -> TimeBraidBatchFeature:
+        """Prepare one request; a positive horizon starts numeric forecasting."""
+        if add_generation_prompt is not True:
+            raise ValueError(
+                "add_generation_prompt must be true for TimeBraid inference."
+            )
+        if tokenize is not True:
+            raise ValueError("tokenize must be true for TimeBraid inference.")
+        if return_dict is not True:
+            raise ValueError("return_dict must be true for TimeBraid inference.")
+        if return_tensors != "pt":
+            raise ValueError(f"return_tensors must be 'pt', got {return_tensors!r}.")
+        normalized_messages = self._normalize_messages(messages)
+        raw_series = self._normalize_raw_timeseries(
+            [] if timeseries is None else timeseries
+        )
+        if horizon is not None and (type(horizon) is not int or horizon <= 0):
+            raise ValueError(f"horizon must be a positive integer, got {horizon!r}.")
+        resolved_target_series_index: int | None = None
+        if horizon is None:
+            if target_series_index is not None:
+                raise ValueError("target_series_index requires a forecast horizon.")
+        else:
+            if not raw_series:
+                raise ValueError("horizon requires at least one input time series.")
+            if target_series_index is None:
+                if len(raw_series) > 1:
+                    raise ValueError(
+                        "target_series_index is required when forecasting from "
+                        "multiple input time series."
+                    )
+                resolved_target_series_index = 0
+            elif type(target_series_index) is not int:
+                raise ValueError(
+                    "target_series_index must be an integer, got "
+                    f"{target_series_index!r}."
+                )
+            elif not 0 <= target_series_index < len(raw_series):
+                raise ValueError(
+                    "target_series_index must select an input time series, got "
+                    f"{target_series_index!r} for {len(raw_series)} series."
+                )
+            else:
+                resolved_target_series_index = target_series_index
+
+        spans: list[dict[str, Any]] = []
+        normalization: list[dict[str, float | str]] = []
+        references: list[str] = []
+        for values in raw_series:
+            try:
+                mean, std = compute_zscore_stats(values, self.normalization_epsilon)
+            except (OverflowError, ValueError) as exc:
+                raise ValueError(
+                    "Time-series values cannot be z-score normalized."
+                ) from exc
+            normalized_values = zscore_with_stats(values, mean=mean, std=std)
+            role = "context" if horizon is not None else "observed"
+            spans.append(
+                {
+                    "len": len(normalized_values),
+                    "role": role,
+                    "values": normalized_values,
+                    "loss_start": len(normalized_values),
+                }
+            )
+            references.append(
+                build_inline_named_zscore_span_reference(
+                    length_tag=len(values), mean=mean, std=std
+                )
+            )
+            normalization.append(
+                {
+                    "method": "history_population_zscore",
+                    "mean": mean,
+                    "std": std,
+                    "epsilon": self.normalization_epsilon,
+                }
+            )
+
+        if references:
+            references_text = "\n".join(
+                f"Series {index + 1}: {reference}"
+                for index, reference in enumerate(references)
+            )
+            target_text = ""
+            if horizon is not None and len(references) > 1:
+                target_text = (
+                    "\nForecast target: "
+                    f"Series {int(resolved_target_series_index) + 1}."
+                )
+            normalized_messages[-1] = (
+                "user",
+                normalized_messages[-1][1]
+                + "\n\nTime series inputs:\n"
+                + references_text
+                + target_text,
+            )
+        last_role, last_content = normalized_messages[-1]
+        normalized_messages[-1] = (
+            last_role,
+            last_content.rstrip() + "\n" + _NO_THINK_SUFFIX,
+        )
+        transcript = format_qwen_chat_turns([*normalized_messages, ("assistant", "")])
+        if not transcript.endswith(_ASSISTANT_PREFIX + _ASSISTANT_SUFFIX + "\n"):
+            raise RuntimeError("TimeBraid assistant scaffold rendering drifted.")
+        prompt = transcript[: -len(_ASSISTANT_SUFFIX + "\n")]
+
+        previous_padding_side = getattr(self.tokenizer, "padding_side", None)
+        self.tokenizer.padding_side = "left"
+        try:
+            tokenized = self.tokenizer(
+                [prompt],
+                add_special_tokens=False,
+                padding=True,
+                return_tensors="pt",
+            )
+        finally:
+            if previous_padding_side is not None:
+                self.tokenizer.padding_side = previous_padding_side
+        input_ids = tokenized["input_ids"]
+        attention_mask = tokenized["attention_mask"]
+        data: dict[str, Any] = {
+            "input_ids": input_ids,
+            "attention_mask": attention_mask,
+        }
+        normalized_spans = normalize_timeseries_spans(
+            spans,
+            sample_idx=0,
+            require_segment_id=False,
+            default_segment_id=1,
+            max_spans_per_sample=self.max_spans_per_sample,
+        )
+        if normalized_spans:
+            data.update(
+                build_mot_batch_timeseries_tensors(
+                    tokenizer=self.tokenizer,
+                    input_ids_by_sample=input_ids.tolist(),
+                    spans_by_sample=[normalized_spans],
+                    max_spans_per_sample=self.max_spans_per_sample,
+                )
+            )
+            data["mot_target_horizons"] = torch.tensor(
+                [0 if horizon is None else horizon], dtype=torch.long
+            )
+            if horizon is not None:
+                data["mot_target_history_span_idxs"] = torch.tensor(
+                    [int(resolved_target_series_index)], dtype=torch.long
+                )
+                if len(normalized_spans) >= self.max_spans_per_sample:
+                    raise ValueError(
+                        "Forecasting requires one additional output span; provide "
+                        f"fewer than {self.max_spans_per_sample} input time series."
+                    )
+                # Forecasting is an explicit output request, so use the existing
+                # open-target protocol instead of asking the LM to emit <ts>.
+                # All input series and text remain visible; only the selected
+                # series supplies the numeric target's history and scale.
+                self._open_forecast_target(data, int(resolved_target_series_index))
+                prompt += "<ts>"
+                input_ids = data["input_ids"]
+                attention_mask = data["attention_mask"]
+        return TimeBraidBatchFeature(
+            data,
+            postprocess_context=TimeBraidRequestContext(
+                prompt_text=prompt,
+                horizon=horizon,
+                target_series_index=resolved_target_series_index,
+                normalization=normalization,
+                prompt_width=int(input_ids.shape[1]),
+                prompt_tokens=int(attention_mask[0].sum().item()),
+                num_timeseries_spans=len(normalized_spans)
+                + (1 if horizon is not None else 0),
+                eos_token_ids=_qwen_stop_token_ids(self.tokenizer),
+                pad_token_id=self.tokenizer.pad_token_id,
+            ),
+        )
+
+    def _open_forecast_target(
+        self, data: dict[str, Any], target_series_index: int
+    ) -> None:
+        """Seed a separate assistant target from the selected observed history."""
+        input_ids = data["input_ids"]
+        open_position = int(input_ids.shape[1])
+        open_id = self.tokenizer.convert_tokens_to_ids("<ts>")
+        data["input_ids"] = torch.cat(
+            [input_ids, input_ids.new_tensor([[open_id]])], dim=1
+        )
+        mask = data["attention_mask"]
+        data["attention_mask"] = torch.cat([mask, mask.new_ones((1, 1))], dim=1)
+        for field, value in list(data.items()):
+            if field.startswith("ts_"):
+                data[field] = torch.cat(
+                    [value, value[:, target_series_index : target_series_index + 1]],
+                    dim=1,
+                )
+        data["ts_roles"][0, -1] = ROLE_TO_ID["target"]
+        data["ts_text_start_token_idxs"][0, -1] = open_position
+        data["ts_text_end_token_idxs"][0, -1] = -1
+
+    def _normalize_messages(
+        self, messages: Sequence[Mapping[str, str]]
+    ) -> list[tuple[str, str]]:
+        if not isinstance(messages, Sequence) or isinstance(messages, (str, bytes)):
+            raise ValueError("messages must be a non-empty sequence.")
+        normalized: list[tuple[str, str]] = []
+        for index, message in enumerate(messages):
+            if not isinstance(message, Mapping):
+                raise ValueError(f"messages[{index}] must be an object.")
+            role = message.get("role")
+            content = message.get("content")
+            if role not in {"system", "user", "assistant"}:
+                raise ValueError(f"messages[{index}] has unsupported role {role!r}.")
+            if not isinstance(content, str) or not content.strip():
+                raise ValueError(f"messages[{index}].content must be non-empty text.")
+            try:
+                content.encode("utf-8", errors="strict")
+            except UnicodeEncodeError as exc:
+                raise ValueError(
+                    f"messages[{index}].content must be valid UTF-8 text."
+                ) from exc
+            normalized.append((str(role), content.strip()))
+        if not normalized or normalized[-1][0] != "user":
+            raise ValueError("messages must end with a user turn.")
+        return normalized
+
+    def _normalize_raw_timeseries(
+        self, timeseries: Sequence[Sequence[float]]
+    ) -> list[list[float]]:
+        if not isinstance(timeseries, Sequence) or isinstance(timeseries, (str, bytes)):
+            raise ValueError("timeseries must be a sequence of numeric sequences.")
+        normalized: list[list[float]] = []
+        for series_index, series in enumerate(timeseries):
+            if not isinstance(series, Sequence) or isinstance(series, (str, bytes)):
+                raise ValueError(
+                    f"timeseries[{series_index}] must be a numeric sequence."
+                )
+            values: list[float] = []
+            for value_index, value in enumerate(series):
+                if isinstance(value, bool) or not isinstance(value, (int, float)):
+                    raise ValueError(
+                        f"timeseries[{series_index}][{value_index}] must be finite."
+                    )
+                try:
+                    casted = float(value)
+                except (OverflowError, ValueError) as exc:
+                    raise ValueError(
+                        f"timeseries[{series_index}][{value_index}] must be finite."
+                    ) from exc
+                if not math.isfinite(casted):
+                    raise ValueError(
+                        f"timeseries[{series_index}][{value_index}] must be finite."
+                    )
+                values.append(casted)
+            if not values:
+                raise ValueError(f"timeseries[{series_index}] must be non-empty.")
+            normalized.append(values)
+        return normalized
+
+    def post_process_generation(
+        self,
+        output: Any,
+        *,
+        model_inputs: TimeBraidBatchFeature | None = None,
+        context: TimeBraidRequestContext | None = None,
+    ) -> dict[str, Any]:
+        """Decode one mixed generation result.
+
+        Supply either `model_inputs` — the object `apply_chat_template`
+        returned, which carries its own context — or `context` directly. The
+        second form exists because the context is an attribute rather than
+        batch data, so `{**inputs}` and any cross-process hand-off drop it.
+        """
+        if model_inputs is not None and context is not None:
+            raise TypeError("Pass either model_inputs or context, not both.")
+        if context is None:
+            if model_inputs is None:
+                raise TypeError(
+                    "post_process_generation needs the request context: pass "
+                    "`model_inputs=` with the object returned by "
+                    "apply_chat_template, or `context=` with its "
+                    "`.postprocess_context` if the batch was rebuilt or moved "
+                    "between processes."
+                )
+            if not isinstance(model_inputs, TimeBraidBatchFeature):
+                raise TypeError("model_inputs must be returned by TimeBraidProcessor.")
+            context = model_inputs.postprocess_context
+        if not isinstance(context, TimeBraidRequestContext):
+            raise TypeError(
+                "context must be a TimeBraidRequestContext, got "
+                f"{type(context).__name__}."
+            )
+        generated_ids = getattr(output, "sequences", output)
+        if not isinstance(generated_ids, torch.Tensor) or generated_ids.ndim != 2:
+            raise TypeError("generation output must expose sequences shaped [1, L].")
+        if int(generated_ids.shape[0]) != 1:
+            raise ValueError("TimeBraid postprocessing supports one request at a time.")
+        prompt_width = int(context.prompt_width)
+        suffix_ids = [
+            int(token_id)
+            for token_id in generated_ids[0, prompt_width:].detach().cpu().tolist()
+        ]
+        completion_tokens = len(suffix_ids)
+        eos_hit = False
+        eos_token_ids = set(context.eos_token_ids)
+        for index, token_id in enumerate(suffix_ids):
+            if token_id in eos_token_ids:
+                completion_tokens = index + 1
+                suffix_ids = suffix_ids[:index]
+                eos_hit = True
+                break
+        if not eos_hit:
+            pad_token_id = context.pad_token_id
+            while (
+                suffix_ids
+                and pad_token_id is not None
+                and suffix_ids[-1] == int(pad_token_id)
+            ):
+                suffix_ids.pop()
+            completion_tokens = len(suffix_ids)
+        decoded = self.tokenizer.decode(suffix_ids, skip_special_tokens=False)
+        text = decoded.split(_ASSISTANT_SUFFIX, 1)[0].strip()
+        if "</think>" in text:
+            text = text.split("</think>", 1)[1].strip()
+        text = text.replace("<ts>", "").replace("</ts>", "").strip()
+
+        horizon = context.horizon
+        generated_ts = getattr(output, "generated_ts_values", None)
+        normalized_values: list[float] = []
+        if generated_ts is not None:
+            if generated_ts == []:
+                generated_ts = [[]]
+            if not isinstance(generated_ts, list) or len(generated_ts) != 1:
+                raise ValueError(
+                    "Generated time-series values are not request-aligned."
+                )
+            normalized_values = [float(value) for value in generated_ts[0]]
+        if horizon is None and normalized_values:
+            raise ValueError("A horizon-free completion emitted time-series values.")
+        if horizon is not None and len(normalized_values) != int(horizon):
+            raise ValueError(
+                "Generated time-series horizon mismatch: "
+                f"expected={horizon}, got={len(normalized_values)}."
+            )
+        if any(not math.isfinite(value) for value in normalized_values):
+            raise ValueError("Generated time-series values must be finite.")
+
+        rollout_records = getattr(output, "rollout_records", None)
+        record = None
+        if rollout_records is not None:
+            if rollout_records == []:
+                rollout_records = None
+        if rollout_records is not None:
+            if not isinstance(rollout_records, list) or len(rollout_records) != 1:
+                raise ValueError("Generation rollout records are not request-aligned.")
+            record = rollout_records[0]
+            if not isinstance(record, Mapping):
+                raise TypeError("Generation rollout record must be an object.")
+        timeseries_result = None
+        if normalized_values:
+            normalization = context.normalization
+            target_series_index = context.target_series_index
+            if type(
+                target_series_index
+            ) is not int or not 0 <= target_series_index < len(normalization):
+                raise ValueError(
+                    "Numeric generation requires a valid target normalization record."
+                )
+            mean = float(normalization[target_series_index]["mean"])
+            std = float(normalization[target_series_index]["std"])
+            values = [value * std + mean for value in normalized_values]
+            if any(not math.isfinite(value) for value in values):
+                raise ValueError("Generated time-series denormalization overflowed.")
+            timeseries_result = {
+                "values": values,
+                "normalized_values": normalized_values,
+                "target_series_index": target_series_index,
+            }
+
+        decode_impl = "hf_generate"
+        finish_reason = "stop" if eos_hit else "length"
+        if record is not None:
+            decode_impl = str(record.get("decode_impl") or "")
+            scheduler_finish = str(record.get("finish_reason") or "")
+            if scheduler_finish == FinishReason.EOS_OR_PROTOCOL_STOP:
+                finish_reason = "stop"
+            elif scheduler_finish == FinishReason.TEXT_BUDGET:
+                finish_reason = "length"
+            else:
+                raise ValueError(
+                    f"Unknown TimeBraid scheduler finish reason {scheduler_finish!r}."
+                )
+        return {
+            "content": text,
+            "timeseries": timeseries_result,
+            "target_series_index": context.target_series_index,
+            "normalization": context.normalization or None,
+            "finish_reason": finish_reason,
+            "prompt_tokens": int(context.prompt_tokens),
+            "completion_tokens": completion_tokens,
+            "decode_impl": decode_impl,
+        }
+
+
+__all__ = [
+    "TimeBraidBatchFeature",
+    "TimeBraidProcessor",
+]