Add Kumo-Forecast to SDM’s Kumo time-series models.

[agautam478]

Replacement for #908, submitted from the public fork. Preserves the forecasting implementation and review fixes from that PR, including the latest merge of main.

Ports tensor-native Kumo-Forecast from NVIDIA/Kumo-TS into SDM’s Kumo time-series package.

Includes patch extraction and embedding, explicit-state reversible instance normalization, cross-variate attention, a fixed-horizon forecast head, and a PyTorch-native T5 encoder without a Transformers runtime dependency.

Provides strict loading for the base and cross-channel checkpoints, sdm.models.KumoForecasting with forward and fit/predict, and a recipe matching the published standardizer and restoring predictions to original units.

Uses nvidia/Kumo-Forecast with the existing pinned checkpoint revision. The SDM public API, model architecture, checkpoint filenames, and normalization constants are unchanged by the rebranding.

[copy-pr-bot]

This pull request requires additional validation before any workflows can run on NVIDIA's runners.

Pull request vetters can view their responsibilities here.

Contributors can view more details about this message here.

[coderabbitai]

Navigate logical layers of code changes, visualize relationships, and explore their blast radius.

🧰 Additional context used

📚 Code guidelines (1)

AGENTS.md — auto-discovered

No actionable comments were generated in the recent review. 🎉

ℹ️ Recent review info

⚙️ Run configuration

• Configuration used: Repository: NVIDIA/structured-data-models/.coderabbit.yaml
• Review profile: QUIET
• Plan: Enterprise
• Run ID: 62be30c6-ef2a-4d01-98a1-2a0b35ab45a7

📥 Commits

Reviewing files that changed from the base of the PR and between 98f6128 and 2dde259.

📒 Files selected for processing (22)

• docs/source/api/models.rst
• examples/README.md
• examples/kumo/timeseries/forecasting.py
• sdm/models/__init__.py
• sdm/models/kumo/__init__.py
• sdm/models/kumo/timeseries/__init__.py
• sdm/models/kumo/timeseries/forecasting/__init__.py
• sdm/models/kumo/timeseries/forecasting/attention.py
• sdm/models/kumo/timeseries/forecasting/ckpt.py
• sdm/models/kumo/timeseries/forecasting/encoder.py
• sdm/models/kumo/timeseries/forecasting/head.py
• sdm/models/kumo/timeseries/forecasting/model.py
• sdm/models/kumo/timeseries/forecasting/normalization.py
• sdm/models/kumo/timeseries/forecasting/patch.py
• sdm/models/kumo/timeseries/forecasting/recipe.py
• sdm/nn/attention.py
• test/models/kumo/timeseries/forecasting/conftest.py
• test/models/kumo/timeseries/forecasting/test_encoder.py
• test/models/kumo/timeseries/forecasting/test_layers.py
• test/models/kumo/timeseries/forecasting/test_model.py
• test/models/kumo/timeseries/forecasting/test_recipe.py
• test/nn/test_attention.py

Included review availability: This review used your included allowance. Your plan provides up to 12 included reviews per hour; 11 remain after this review.

---

📝 Summary

Summary by CodeRabbit

• New Features
  • Added Kumo forecasting with pretrained weights, supporting forecasts up to 72 steps from historical time-series data.
  • Forecasts can be generated directly or through a cached prediction workflow, with results returned in the target columns’ original scale.
  • Attention now supports configurable dropout and additive attention bias.
• Documentation
  • Added a runnable forecasting example and listed the model’s parameter counts and Apache-2.0 licenses.

Walkthrough

Adds KumoForecasting for numerical time-series forecasts. The model uses up to 512 history steps and supports forecast horizons of up to 72 steps. The change also adds its neural-network components, checkpoint handling, package exports, tests, and a runnable example.

Changes

Kumo forecasting

Layer / File(s)	Summary
Attention and T5 encoder sdm/nn/attention.py, sdm/models/kumo/timeseries/forecasting/encoder.py, sdm/models/kumo/timeseries/forecasting/ckpt.py, test/nn/test_attention.py, test/models/kumo/timeseries/forecasting/test_encoder.py	SDPA gains additive attention bias and configurable dropout. The T5 encoder adds shared relative-position bias, masked-key handling, and gated feed-forward layers. Checkpoint keys are remapped to SDM parameter names. Tests cover encoder outputs, masking, gradients, attention bias, and dropout.
Time-series processing and prediction components sdm/models/kumo/timeseries/forecasting/{normalization,patch,attention,head,recipe}.py, test/models/kumo/timeseries/forecasting/{conftest,test_layers,test_recipe}.py	Adds reversible instance normalization, patching and patch embeddings, cross-channel attention, a forecasting head, and a recipe that uses fixed statistics. Tests cover component behavior, shapes, masking, and recipe transformations.
Model API and checkpoint integration sdm/models/kumo/timeseries/forecasting/model.py, sdm/models/{__init__.py,kumo/__init__.py,kumo/timeseries/__init__.py,kumo/timeseries/forecasting/__init__.py}, test/models/kumo/timeseries/forecasting/test_model.py	Adds KumoForecasting, including input and horizon validation, cached prediction replay, backbone assembly, and pretrained or supplied checkpoint loading. The model is exported through the package hierarchy. Tests cover prediction paths, input cases, checkpoint loading, and device construction.
Model documentation and runnable example docs/source/api/models.rst, examples/README.md, examples/kumo/timeseries/forecasting.py	Adds the model to the API overview and examples index. The example runs direct and cached prediction on synthetic history and checks the outputs.

Priority: ➖ Normal

Estimated code review effort: 4 (Complex) | ~45 minutes

Sequence Diagram(s)

sequenceDiagram
  participant Caller
  participant KumoForecasting
  participant RevIN
  participant PatchEmbedding
  participant T5Encoder
  participant CrossChannelAttention
  participant ForecastingHead
  Caller->>KumoForecasting: Submit context and forecast horizon
  KumoForecasting->>RevIN: Normalize time-series inputs
  KumoForecasting->>PatchEmbedding: Create patch embeddings
  KumoForecasting->>T5Encoder: Encode patch embeddings
  KumoForecasting->>CrossChannelAttention: Mix variates when enabled
  KumoForecasting->>ForecastingHead: Project encoded patches to forecast
  KumoForecasting->>RevIN: Restore forecast scale
  KumoForecasting-->>Caller: Return forecast and target columns

Loading

Merge Risk: ⚪ Minimal · up to 2dde2

This adds a new forecasting model with tests, docs and an example. No concrete merge-blocking risk was identified in the supplied changes.

🚥 Pre-merge checks | ✅ 4 | ❌ 1

❌ Failed checks (1 warning)

Check name	Status	Explanation	Resolution
Docstring Coverage	⚠️ Warning	Docstring coverage is 23.53% which is insufficient. The required threshold is 80.00%. Docstring coverage is scoped to functions touched by this diff. Analyzed 68 functions across 20 files. (2 skipped:…	Write docstrings for the functions missing them to satisfy the coverage threshold.

✅ Passed checks (4 passed)

Check name	Status	Explanation
Title check	✅ Passed	The title clearly describes the main change: adding Kumo-Forecast to SDM’s Kumo time-series models.
Description check	✅ Passed	The description directly explains the forecasting implementation, model API, checkpoints, and preprocessing included in the changeset.
Linked Issues check	✅ Passed	Check skipped because no linked issues were found for this pull request.
Out of Scope Changes check	✅ Passed	Check skipped because no linked issues were found for this pull request.

Full details: Docstring Coverage

Explanation

Docstring coverage is 23.53% which is insufficient. The required threshold is 80.00%. Docstring coverage is scoped to functions touched by this diff. Analyzed 68 functions across 20 files. (2 skipped: 2 unsupported.)

• Fix all pre-merge checks with AI

✨ Finishing Touches

🧪 Generate unit tests (beta)

• Create a new PR

• Autopilot · Keep fixing CodeRabbit findings and required CI, and resolving merge conflicts

---

_{Comment @coderabbitai help to get the list of available commands.}