faster decoding related changes

README.md

@@ -57,8 +57,12 @@ pip install -e .
 ### Code Example
 
 ```python
+import torch
 import numpy as np
 import timesfm
+
+torch.set_float32_matmul_precision("high")
+
 model = TimesFM_2p5_200M_torch.from_pretrained("google/timesfm-2.5-200m-pytorch")
 
 model.compile(

src/timesfm/configs.py

@@ -94,6 +94,7 @@ class TransformerConfig:
   use_bias: bool
   use_rotary_position_embeddings: bool
   ff_activation: Literal["relu", "swish", "none"]
+  fuse_qkv: bool
 
 
 @dataclasses.dataclass(frozen=True)

src/timesfm/timesfm_2p5/timesfm_2p5_base.py

@@ -16,7 +16,9 @@
 
 import dataclasses
 from typing import Any, Callable
+
 import numpy as np
+
 from .. import configs
 
 ResidualBlockConfig = configs.ResidualBlockConfig
@@ -107,6 +109,57 @@ class TimesFM_2p5_200M_Definition:
       use_bias=False,
       use_rotary_position_embeddings=True,
       ff_activation="swish",
+      fuse_qkv=False,
+    ),
+  )
+  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",
+  )
+
+
+@dataclasses.dataclass(frozen=True)
+class TimesFM_2p5_new_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(
@@ -179,9 +232,7 @@ class TimesFM_2p5:
       idx += 1
       if idx == self.global_batch_size:
         idx = 0
-        point_forecast, quantile_forecast = self.compiled_decode(
+        point_forecast, quantile_forecast = self.compiled_decode(horizon, values, masks)
-            horizon, values, masks
-        )
         output_points.append(point_forecast)
         output_quantiles.append(quantile_forecast)
         values = []

src/timesfm/timesfm_2p5/timesfm_2p5_torch.py

@@ -22,7 +22,7 @@ from typing import Dict, Optional, Sequence, Union
 import numpy as np
 import torch
 from huggingface_hub import ModelHubMixin, hf_hub_download
-from safetensors.torch import load_file
+from safetensors.torch import load_file, save_file
 from torch import nn
 
 from .. import configs
@@ -78,7 +78,7 @@ class TimesFM_2p5_200M_torch_module(nn.Module):
   def load_checkpoint(self, path: str):
     """Loads a PyTorch TimesFM model from a checkpoint."""
     tensors = load_file(path)
-    self.load_state_dict(tensors)
+    self.load_state_dict(tensors, strict=True)
     self.to(self.device)
 
   def forward(
@@ -113,9 +113,6 @@ class TimesFM_2p5_200M_torch_module(nn.Module):
   def decode(self, horizon: int, inputs, masks):
     """Decodes the time series."""
 
-    inputs = inputs.to(self.device)
-    masks = masks.to(self.device)
-
     with torch.no_grad():
       batch_size, context = inputs.shape[0], inputs.shape[1]
       num_decode_steps = (horizon - 1) // self.o
@@ -257,10 +254,16 @@ class TimesFM_2p5_200M_torch_module(nn.Module):
     return outputs
 
 
+class TimesFM_2p5_200M_torch_new_module(TimesFM_2p5_200M_torch_module):
+  """TimesFM 2.5 with 200M parameters."""
+
+  config = timesfm_2p5_base.TimesFM_2p5_new_200M_Definition()
+
+
 class TimesFM_2p5_200M_torch(timesfm_2p5_base.TimesFM_2p5, ModelHubMixin):
   """PyTorch implementation of TimesFM 2.5 with 200M parameters."""
 
-  model: nn.Module = TimesFM_2p5_200M_torch_module()
+  model: nn.Module = TimesFM_2p5_200M_torch_new_module()
 
   @classmethod
   def _from_pretrained(
@@ -310,6 +313,17 @@ class TimesFM_2p5_200M_torch(timesfm_2p5_base.TimesFM_2p5, ModelHubMixin):
     instance.model.load_checkpoint(model_file_path)
     return instance
 
+  def _save_pretrained(self, save_directory: Union[str, Path]):
+    """
+    Saves the model's state dictionary to a safetensors file. This method
+    is called by the `save_pretrained` method from `ModelHubMixin`.
+    """
+    if not os.path.exists(save_directory):
+      os.makedirs(save_directory)
+
+    weights_path = os.path.join(save_directory, "model.safetensors")
+    save_file(self.model.state_dict(), weights_path)
+
   def compile(self, forecast_config: configs.ForecastConfig, **kwargs) -> None:
     """Attempts to compile the model for fast decoding.
 
@@ -320,8 +334,9 @@ class TimesFM_2p5_200M_torch(timesfm_2p5_base.TimesFM_2p5, ModelHubMixin):
       **kwargs: Additional keyword arguments to pass to model.compile().
     """
 
-    if kwargs.get("backend", None) is not None:
+    if forecast_config.torch_compile:
-      self.model.compile(**kwargs)
+      self.model = torch.compile(self.model)
+    self.model.eval()
     self.global_batch_size = (
       forecast_config.per_core_batch_size * self.model.device_count
     )
@@ -363,8 +378,10 @@ class TimesFM_2p5_200M_torch(timesfm_2p5_base.TimesFM_2p5, ModelHubMixin):
           f"Horizon must be less than the max horizon. {horizon} > {fc.max_horizon}."
         )
 
-      inputs = torch.Tensor(np.array(inputs)).to(self.model.device)
+      inputs = (
-      masks = torch.Tensor(np.array(masks)).to(self.model.device).to(torch.bool)
+        torch.from_numpy(np.array(inputs)).to(self.model.device).to(torch.float32)
+      )
+      masks = torch.from_numpy(np.array(masks)).to(self.model.device).to(torch.bool)
       batch_size = inputs.shape[0]
 
       if fc.infer_is_positive:

src/timesfm/torch/transformer.py

@@ -18,12 +18,11 @@ import math
 from typing import Callable
 
 import torch
-from torch import nn
 import torch.nn.functional as F
+from torch import nn
 
 from .. import configs
-from . import normalization
+from . import normalization, util
-from . import util
 
 LayerNorm = nn.LayerNorm
 RMSNorm = normalization.RMSNorm
@@ -86,14 +85,13 @@ class RotaryPositionalEmbedding(nn.Module):
       / self.embedding_dims
     )
     timescale = (
-        self.min_timescale
+      self.min_timescale * (self.max_timescale / self.min_timescale) ** fraction
-        * (self.max_timescale / self.min_timescale) ** fraction
     ).to(inputs.device)
     if position is None:
       seq_length = inputs.shape[1]
-      position = torch.arange(
+      position = torch.arange(seq_length, dtype=torch.float32, device=inputs.device)[
-          seq_length, dtype=torch.float32, device=inputs.device
+        None, :
-      )[None, :]
+      ]
 
     if len(inputs.shape) == 4:
       position = position[..., None, None]
@@ -159,6 +157,7 @@ class MultiHeadAttention(nn.Module):
     use_bias: bool = False,
     attention_fn: Callable[..., torch.Tensor] = _dot_product_attention,
     qk_norm: str = "rms",
+    fuse_qkv: bool = False,
   ):
     super().__init__()
     self.num_heads = num_heads
@@ -167,6 +166,7 @@ class MultiHeadAttention(nn.Module):
     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(
@@ -174,6 +174,9 @@ class MultiHeadAttention(nn.Module):
         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)
@@ -205,25 +208,25 @@ class MultiHeadAttention(nn.Module):
   ) -> tuple[torch.Tensor, DecodeCache | None]:
     b, n_patches, _ = inputs_q.shape
     if patch_mask is None:
-      patch_mask = torch.zeros(
+      patch_mask = torch.zeros(b, n_patches, dtype=torch.bool, device=inputs_q.device)
-          b, n_patches, dtype=torch.bool, device=inputs_q.device
-      )
 
-    query = self.query(inputs_q).view(
+    if self.fuse_qkv:
-        b, n_patches, self.num_heads, self.head_dim
+      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(
+      value = self.value(inputs_q).view(b, n_patches, self.num_heads, self.head_dim)
-        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)
+        torch.sum(patch_mask.to(torch.int32), dim=-1) + decode_cache.num_masked
-          + decode_cache.num_masked
       )
       next_index = decode_cache.next_index.clone()
 
@@ -244,13 +247,18 @@ class MultiHeadAttention(nn.Module):
 
     if decode_cache is not None:
       _, decode_cache_size, _, _ = decode_cache.value.shape
-      for i in range(b):
+
-        start = decode_cache.next_index[i]
+      start = decode_cache.next_index[0]
       end = start + n_patches
-        decode_cache.key[i, start:end] = key[i].clone()
+
-        decode_cache.value[i, start:end] = value[i].clone()
+      # Perform a single, vectorized slice assignment for the entire batch.
-      key = decode_cache.key.clone()
+      # This is vastly more efficient than a Python for-loop.
-      value = decode_cache.value.clone()
+
+      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(
@@ -260,9 +268,7 @@ class MultiHeadAttention(nn.Module):
         kv_length=decode_cache_size,
       )
     else:
-      attn_mask = make_attn_mask(
+      attn_mask = make_attn_mask(query_length=n_patches, num_all_masked_kv=num_masked)
-          query_length=n_patches, num_all_masked_kv=num_masked
-      )
 
     x = self.attention_fn(
       query,
@@ -270,6 +276,7 @@ class MultiHeadAttention(nn.Module):
       value,
       mask=attn_mask,
     )
+
     x = x.reshape(b, n_patches, self.in_features)
     out = self.out(x)
     return out, decode_cache
@@ -294,6 +301,7 @@ class Transformer(nn.Module):
       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":
@@ -334,9 +342,7 @@ class Transformer(nn.Module):
     )
     attn_output = self.post_attn_ln(attn_output) + input_embeddings
     output_embeddings = (
-        self.post_ff_ln(
+      self.post_ff_ln(self.ff1(self.activation(self.ff0(self.pre_ff_ln(attn_output)))))
-            self.ff1(self.activation(self.ff0(self.pre_ff_ln(attn_output))))
-        )
       + attn_output
     )
     return output_embeddings, decode_cache