eCity-TS-313M

eCity-TS-313M is a multivariate time series forecasting model for water utility scenarios. It is built on a decoder-only, patch-based Transformer whose attention alternates between the time axis (causal) and the variate axis (fully connected), and whose quantile head outputs uncertainty intervals.

This document covers inference only: how to load the artifacts with ONNX Runtime and TensorRT, feed data, and read results. Given a history window of up to 4096 steps × 15 variables, a single forward pass predicts the next 96 steps at 9 quantiles, supporting interval forecasting and scheduling decisions for supply volume, distribution network flow, pressure, and tank level.

1. Deliverables

Artifact	Path	Notes
ONNX graph	artifacts/ecity_ts_313m_15x4096.onnx	Input/output signature below
ONNX external weights	artifacts/ecity_ts_313m_15x4096.onnx.data	~1.3 GB, must sit next to the .onnx file
TensorRT engine	artifacts/ecity_ts_313m_15x4096_fp32.engine	~1216 MB, bound to the TensorRT / CUDA / GPU used at build time
Reference sample	example/ecity_ts_313m_15x4096.npz	Fixed input (batch=4) plus FP32 reference output; used by both demo scripts for regression checks
Benchmark snapshot	example/ecity_ts_313m_15x4096.json	One-shot TensorRT batch=4 latency / peak-memory snapshot taken at release time

Input/Output Contract

Name	Shape	Type	Notes
target	[B, 15, 4096]	float32	Historical observations; unobserved points can hold any value and are masked by target_mask
target_mask	[B, 15, 4096]	bool	True means the point was observed
series_ids	[B, 15]	int64	Variate identity; keep it stable across batches, usually 0..14
quantiles	[9, B, 15, 96]	float32	quantiles 0.1–0.9 × batch × variates × forecast steps; the median is index 4

• B is the only dynamic dimension; both ONNX and TensorRT use an optimization profile of min / opt / max = 1 / 4 / 8.
• Quantiles are laid out along axis 0 in the ONNX graph; TensorRT internally rewrites TopK onto the last dimension and transposes back, so the external shape is unchanged.

2. Environment Setup

# 1. Create a dedicated Python environment (3.12 recommended)
conda create -n ecity-ts-313m python=3.12 -y
conda activate ecity-ts-313m

# 2. Install PyTorch matching CUDA 12.8 (used by the TensorRT path for memory and streams)
python -m pip install torch --index-url https://download.pytorch.org/whl/cu128

# 3. Install runtime dependencies (numpy, onnxruntime-gpu, tensorrt, torch)
python -m pip install -r requirements.txt

• Requires a CUDA 12.x driver and cuDNN 9 (cuDNN comes with the torch cu128 wheel). demo_onnx.py calls ort.preload_dlls() so the CUDA EP does not fall back to CPU when libcudnn.so.9 is missing from the library search path.
• The ONNX path only needs numpy + onnxruntime-gpu; the TensorRT path additionally needs tensorrt + torch (demo_trt.py uses PyTorch for device memory and stream management, and never loads model weights).
• Environment self-check:

python -c "import numpy, torch, onnxruntime, tensorrt; print(numpy.__version__, torch.__version__, onnxruntime.__version__, tensorrt.__version__)"

• Verified versions: Python 3.12, numpy 1.26.4, torch 2.9.1+cu128, onnxruntime-gpu 1.26.0, TensorRT 10.6.0, 8× NVIDIA H200.
• The engine is tightly coupled to its build environment; changing the GPU model, CUDA, or TensorRT version requires a rebuild.

3. Download

All files are hosted in the HuggingFace model repository eCityAI-Sup/eCity-TS-313M (also mirrored on ModelScope as eCityAI/eCity-TS-313M). The download reproduces the repository layout one-to-one (artifacts/, example/, and the two demo scripts), so the demos run out of the box:

pip install -U "huggingface_hub[cli]"
huggingface-cli download eCityAI-Sup/eCity-TS-313M --local-dir ./eCity-TS-313M
cd eCity-TS-313M

Keep the .onnx and .onnx.data files together in artifacts/. The TensorRT engine is hardware-bound; if your GPU / CUDA / TensorRT combination differs from the release environment (see section 2), use the ONNX artifact instead.

4. ONNX Runtime Inference

import numpy as np
import onnxruntime as ort

# The pip-installed cuDNN is not on the system library path; preload it or the CUDA EP silently falls back to CPU.
ort.preload_dlls()

session = ort.InferenceSession(
    "artifacts/ecity_ts_313m_15x4096.onnx",
    # The CUDA EP enables TF32 (10-bit mantissa) by default, which inflates the deviation to ~1e-1; disable it.
    providers=[("CUDAExecutionProvider", {"use_tf32": "0"}), "CPUExecutionProvider"],
)

target = np.asarray(x, dtype=np.float32)[None]              # [1, 15, 4096]
target_mask = np.isfinite(x)[None]                          # [1, 15, 4096]
series_ids = np.arange(15, dtype=np.int64)[None]            # [1, 15]

quantiles = session.run(
    ["quantiles"],
    {"target": target, "target_mask": target_mask, "series_ids": series_ids},
)[0]
# quantiles.shape == (9, 1, 15, 96)

• x is a [15, 4096] history matrix; multiple sites can be stacked into [B, 15, 4096] for a single pass.
• See demo_onnx.py for a runnable example (it loads example/ecity_ts_313m_15x4096.npz and compares against the reference output).
• The session only falls back to CPU when the CUDA provider is unavailable; results are identical but much slower. Check the effective provider with session.get_providers().
• Read results directly: point = quantiles[4], interval quantiles[0] (P10) / quantiles[8] (P90).

5. TensorRT Inference

The repository ships EngineRunner in demo_trt.py, which uses PyTorch CUDA tensors directly as device buffers with no host/device copies:

import torch
from demo_trt import EngineRunner

runner = EngineRunner("artifacts/ecity_ts_313m_15x4096_fp32.engine")

target = torch.as_tensor(x, dtype=torch.float32, device="cuda")[None]      # [1, 15, 4096]
target_mask = torch.isfinite(target).bool()                                # [1, 15, 4096]
series_ids = torch.arange(15, dtype=torch.int64, device="cuda")[None]      # [1, 15]

quantiles = runner.infer(target, target_mask, series_ids)   # GPU tensor, (9, 1, 15, 96)
torch.cuda.synchronize()                                     # required before reading or timing

Key points when using the TensorRT Python API directly:

1. Deserialize with trt.Runtime(logger).deserialize_cuda_engine(engine_bytes), then call create_execution_context().
2. For each request, call set_input_shape(name, shape) first (it returns False outside the 1/4/8 profile) and then set_tensor_address(name, data_ptr).
3. Inputs must be contiguous and on CUDA; allocate the output tensor on the same device using get_tensor_shape("quantiles").
4. Submit with execute_async_v3(torch.cuda.current_stream().cuda_stream) so it shares the stream with upstream PyTorch ops.
5. An execution context keeps the shape and address state of the current call; do not share it across threads (keep one context per worker).

6. Dataset Performance

Lower is better on all three metrics. On the Xi'an Xingfu Lindai project dataset (Building Division), eCity-TS-313M cuts MAE by ~83% versus TimesFM and ~56% versus Moirai, and RMSE by ~79% and ~49% respectively.

Model	MAE	MSE	RMSE
TimesFM	229.5565	504803.1905	371.7470
Moirai	87.3544	23699.9194	153.0700
eCity-TS-313M	38.0379	6158.1171	78.4775

7. Measured Baseline (H200, 15×4096, fp32)

Backend	Metric	Result
ONNX Runtime (CUDA EP, TF32 disabled)	max_abs / mean_abs / p99_abs	3.91e-4 / 4.73e-6 / 6.48e-5
TensorRT	max_abs / mean_abs / p99_abs	3.47e-4 / 4.57e-6 / 6.29e-5
TensorRT, batch=4	P50 / P95 / P99 / throughput	142.0 ms / 143.6 ms / 163.2 ms / 28.4 samples/s

Errors are measured against the FP32 reference output; both backends stay far below the 1e-2 acceptable for production. example/ecity_ts_313m_15x4096.json is a separate one-shot release snapshot (latency P50/P95 and peak memory at batch 4); small run-to-run differences against the table above are expected. This network is dominated by normalization, element-wise ops, and long-sequence reductions, so kernel launch and memory traffic dominate. The main benefit of TensorRT over a general-purpose runtime is removing Python/framework overhead and delivering more stable memory usage rather than an order-of-magnitude latency drop.

8. Runtime Constraints

• Batch range: 1 ≤ B ≤ 8; larger batches require rebuilding the engine or changing the profile.
• Fixed shapes: V=15, T=4096, H=96 are baked into the graph; other lengths cannot be fed directly.
• fp32 only: with a 4096 context, fp16 overflows to inf in transformer.layers.0.ffn.fc2 and the following RMSNorm turns entirely NaN.
• Engine is machine-bound: a TensorRT engine is tied to the GPU architecture, CUDA, and TensorRT version, so a new environment requires a rebuild.
• Thread safety: a single execution context is not meant for concurrent calls; create one context per thread for concurrency.
• Missing-value semantics: the current artifact was exported with has_missing_values=false, i.e. every point in the history window is treated as observed. If you need to model gaps and missing data, use an artifact exported with has_missing_values=true.

9. FAQ

Q: How should the forecast be used? quantiles[4] is the median point forecast, while quantiles[0] / quantiles[8] are the P10 / P90 interval, usable for reported values, alarm thresholds, and scheduling margin.

Q: Can ONNX run without a GPU? Yes, drop the CUDA provider to use CPU. Results are identical, but a 15×4096 input is not suitable for online inference on CPU.

Q: Must series_ids be contiguous? No, but they must map one-to-one onto the physical meaning of each variable and stay stable across requests; otherwise the model treats different variables as the same identity.

Q: Why is there an .onnx.data file next to the .onnx file? The weights are ~1.3 GB, which exceeds the single-file protobuf limit, so external data is used. Both files must stay in the same directory and be moved together.

Q: Can multiple sites share one batch? Yes, stack their target / target_mask along dimension 0 (same variate order and history window) and predict the whole batch in one pass.

Citation

The eCity-TS-313M weights come from the open-source time series foundation model Toto 2.0. If you use this model in your work, please cite the base paper:

@misc{khwaja2026toto20timeseries,
      title={Toto 2.0: Time Series Forecasting Enters the Scaling Era},
      author={Emaad Khwaja and Chris Lettieri and Gerald Woo and Eden Belouadah and Marc Cenac and Guillaume Jarry and Enguerrand Paquin and Xunyi Zhao and Viktoriya Zhukov and Othmane Abou-Amal and Chenghao Liu and Ameet Talwalkar and David Asker},
      year={2026},
      eprint={2605.20119},
      archivePrefix={arXiv},
      primaryClass={cs.LG},
      url={https://arxiv.org/abs/2605.20119},
}

Acknowledgement

eCity-TS-313M is built on the open-source time series foundation model Toto 2.0; we thank its authors for releasing the model weights and inference code. This repository only covers inference for water utility scenarios, the ONNX / TensorRT deployment adaptation, and numerical validation; it contains no training code or training data.

License

This repository is released under the Apache-2.0 License; the weights inherit the Apache-2.0 license of the base model. Architecture hyperparameters and the quantile layout are embedded in the ONNX graph itself, so no separate config file is shipped.