# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. # SPDX-License-Identifier: MIT # # Permission is hereby granted, free of charge, to any person obtaining a # copy of this software and associated documentation files (the "Software"), # to deal in the Software without restriction, including without limitation # the rights to use, copy, modify, merge, publish, distribute, sublicense, # and/or sell copies of the Software, and to permit persons to whom the # Software is furnished to do so, subject to the following conditions: # # The above copyright notice and this permission notice shall be included in # all copies or substantial portions of the Software. # # THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR # IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, # FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL # THE AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER # LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING # FROM, OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER # DEALINGS IN THE SOFTWARE. # Copyright (c) Kyutai, all rights reserved. # This source code is licensed under the license found in the # LICENSE file in the root directory of this source tree. # Copyright (c) Meta Platforms, Inc. and affiliates. # All rights reserved. # # This source code is licensed under the license found in the # LICENSE file in the root directory of this source tree. """ Streaming module API that should be implemented by all Streaming components, """ import abc from contextlib import contextmanager from dataclasses import dataclass, fields, is_dataclass import itertools import math import json from typing import List, Union, Protocol, TypeVar, Generic, Any, Optional import torch from safetensors.torch import save_file, load_file class Resetable(Protocol): def reset(self) -> None: pass State = TypeVar("State", bound=Resetable) StreamingStateDict = dict[str, Union[torch.Tensor, int, float, str, bool, None]] def is_dataclass_instance(obj): """Check if obj is an instance of a dataclass (not the class itself). Parameters ---------- obj : Any Object to check. Returns ------- bool True if obj is an instance of a dataclass, False otherwise. """ return is_dataclass(obj) and not isinstance(obj, type) def _restore_streaming_state_pt(value: torch.Tensor, name: str, state_dict: dict[str, torch.Tensor], ): """Restore the streaming state from the given pt_state dict Parameters ---------- value : torch.Tensor Specific streaming state tensor that needs to be set. name : str Name of the tensor in the state dict. state_dict: StreamingStateDict Flattened state dict containing the values to set. """ if name in state_dict: value.copy_(state_dict[name].to(value.device)) state_dict.pop(name) else: raise KeyError(f"Expected to find a streaming state for {name}.") def _set_streaming_state_inplace(streaming_state: State, state_dict: StreamingStateDict, prefix: str, device: torch.device, ): """Set the streaming state in-place from the given `state_dict` dict. Parameters ---------- streaming_state : State Specific streaming state object that needs to be set. state_dict: StreamingStateDict Flattened state dict containing the values to set. prefix : str Prefix to add to each key when looking up in `state_dict`. device : torch.device Device to move tensors to if needed. """ if isinstance(streaming_state, torch.Tensor): _restore_streaming_state_pt(streaming_state, prefix, state_dict) elif is_dataclass_instance(streaming_state): _restore_streaming_state_from_keys(streaming_state, state_dict, prefix, [field.name for field in fields(streaming_state)], device) elif hasattr(streaming_state, "asdict"): _restore_streaming_state_from_keys(streaming_state, state_dict, prefix, list(streaming_state.asdict().keys()), device) else: raise TypeError(f"Unsupported type {type(streaming_state)} for streaming state with prefix {prefix}.") def _restore_streaming_state_from_keys(streaming_state: State, state_dict: StreamingStateDict, prefix: str, keys: List[str], device: torch.device, ): """Restores the streaming state from the given `state_dict` dict looking up fields by adding `prefix` to each key in `keys` to look up values. `torch.Tensor` are copied to `device` if no `torch.Tensor` is already present otherwise, the data is copied to the device of the existing `torch.Tensor`. Parameters ---------- streaming_state : State Specific streaming state object that needs to be set. state_dict: StreamingStateDict Flattened state dict containing the values to set. prefix : str Prefix to add to each key when looking up in `state_dict`. keys : List[str] List of keys to look up in `state_dict`. device : torch.device Device to move tensors to if needed. """ for key in keys: full_key = f"{prefix}.{key}" existing_value = getattr(streaming_state, key) if isinstance(existing_value, torch.Tensor): _restore_streaming_state_pt(existing_value, full_key, state_dict) elif isinstance(existing_value, (int, float, str, bool, type(None))): if full_key in state_dict: restored_value = state_dict[full_key] if isinstance(restored_value, torch.Tensor): restored_value = restored_value.to(device) setattr(streaming_state, key, restored_value) state_dict.pop(full_key) else: raise RuntimeError(f"Key {full_key} not found in state_dict.") else: _set_streaming_state_inplace(existing_value, state_dict, full_key, device) def safe_asdict(dataclass_obj): """ safe_asdict(dataclass_obj) Converts a dataclass object to a dict, skipping empty nested dataclasses without requiring values to be pickleable. Parameters ---------- dataclass_obj : Any Dataclass object to convert. Returns ------- dict Dictionary representation of the dataclass object. """ out = {} for field in fields(dataclass_obj): value = getattr(dataclass_obj, field.name) if is_dataclass_instance(value): subvalue = safe_asdict(value) if len(subvalue) > 0: out[field.name] = subvalue else: out[field.name] = value return out def _flatten_streaming_state(state_dict: dict[str, torch.Tensor], state_dict_metadata: dict[str, Union[int, float, str, None]], state: dict[str, State], prefix: str, ): """ _flatten_streaming_state(state_dict, state_dict_metadata, state, prefix) Helper function for recursively flattening the streaming state into a dict of tensors and a dict of metadata (non-tensor values). Parameters ---------- state_dict : dict[str, torch.Tensor] Dictionary to store the flattened tensor states. state_dict_metadata : dict[str, Union[int, float, str, None]] Dictionary to store the flattened non-tensor states. state : dict[str, State] The streaming state to flatten. prefix : str Prefix to add to each key in the flattened state. """ for key, value in state.items(): if isinstance(value, torch.Tensor): state_dict[f"{prefix}{key}"] = value.contiguous() elif is_dataclass_instance(value): _flatten_streaming_state(state_dict, state_dict_metadata, safe_asdict(value), prefix=f"{prefix}{key}.") elif isinstance(value, dict): _flatten_streaming_state(state_dict, state_dict_metadata, value, prefix=f"{prefix}{key}.") elif isinstance(value, (str, int, float, bool, type(None))): state_dict_metadata[f"{prefix}{key}"] = value elif hasattr(value, "asdict"): _flatten_streaming_state(state_dict, state_dict_metadata, value.asdict(), prefix=f"{prefix}{key}.") else: raise TypeError(f"Unsupported type {type(value)} for key {key} (prefix={prefix}) in streaming state.") def load_streaming_state(path: str, metadata_path: str, device: Union[str, int] = 'cpu', ) -> StreamingStateDict: """ load_streaming_state(path, metadata_path) Loads a streaming state from a safetensors file and its associated metadata json file. Parameters ---------- str : path Path to the safetensors file. str : metadata_path Path to the metadata json file. device : Union[str, int], optional Device to load the tensors onto, by default 'cpu'. Returns ------- dict The loaded streaming state flattened as a dictionary. """ state_dict = load_file(path, device=device) with open(metadata_path, "rt", encoding="utf-8") as fin: state_dict_metadata = json.load(fin) state_dict.update(state_dict_metadata) return state_dict class StreamingModule(abc.ABC, torch.nn.Module, Generic[State]): """Common API for streaming components. Each streaming component has a streaming state, which is just a dict[str, Tensor]. By convention, the first dim of each tensor must be the batch size. Don't use dots in the key names, as this would clash with submodules (like in state_dict). If `self._is_streaming` is True, the component should use and remember the proper state inside `self._streaming_state`. To set a streaming component in streaming state, use with module.streaming(): ... This will automatically reset the streaming state when exiting the context manager. This also automatically propagates to all streaming children module. Some module might also implement the `StreamingModule.flush` method, although this one is trickier, as all parents module must be StreamingModule and implement it as well for it to work properly. See `StreamingSequential` after. """ def __init__(self) -> None: super().__init__() self._streaming_state: State | None = None self._streaming_propagate: bool = True @property def is_streaming(self): return self._streaming_state is not None def set_streaming_propagate(self, streaming_propagate: bool): self._streaming_propagate = streaming_propagate def _apply_named_streaming(self, fn: Any): def _handle_module(prefix: str, module: torch.nn.Module, recurse: bool = True): propagate = True if isinstance(module, StreamingModule): if module._streaming_propagate: fn(prefix, module) else: propagate = False if not recurse: return if propagate: for name, child in module.named_children(): _handle_module(prefix + "." + name, child) _handle_module("", self, recurse=False) for name, child in self.named_children(): _handle_module(name, child) def _start_streaming(self, batch_size: int): def _start_streaming(name: str, module: StreamingModule): module._streaming_state = module._init_streaming_state(batch_size) self._apply_named_streaming(_start_streaming) def _stop_streaming(self): def _stop_streaming(name: str, module: StreamingModule): module._streaming_state = None self._apply_named_streaming(_stop_streaming) @abc.abstractmethod def _init_streaming_state(self, batch_size: int) -> State: ... def streaming_forever(self, batch_size: int): self._start_streaming(batch_size) @contextmanager def streaming(self, batch_size: int): """Context manager to enter streaming mode. Reset streaming state on exit.""" self._start_streaming(batch_size) try: yield finally: self._stop_streaming() def reset_streaming(self): """Reset the streaming state.""" def _reset(name: str, module: StreamingModule): state = module._streaming_state if state is None: raise ValueError( f"Trying to reset streaming, but {name} wasn't streaming." ) state.reset() self._apply_named_streaming(_reset) def get_streaming_state(self) -> dict[str, Any]: """Return the complete streaming state, including that of sub-modules.""" state: dict[str, Any] = {} def _add(name: str, module: StreamingModule): state[name] = module._streaming_state self._apply_named_streaming(_add) return state def save_streaming_state(self, save_path: str, metadata_save_path: str, extra_state_dict: Optional[dict[str, torch.Tensor]] = None, ): """Save the streaming state, including that of sub-modules, to the given paths. Parameters ---------- save_path : str Path to save the streaming state tensors (safetensors format). metadata_save_path : str Path to save the streaming state metadata (json format). extra_state_dict : Optional[dict[str, torch.Tensor]], optional Extra state dict to include in the saved streaming state tensors, by default None. """ state_dict = {} if extra_state_dict is not None: state_dict.update(extra_state_dict) state_dict_metadata = {} state = self.get_streaming_state() _flatten_streaming_state(state_dict, state_dict_metadata, state, prefix="") save_file(state_dict, save_path) with open(metadata_save_path, "wt", encoding="utf-8") as fout: json.dump(state_dict_metadata, fout) def set_streaming_state_inplace(self, state: StreamingStateDict): """ Set the streaming state in-place, including that of sub-modules using a flattened-state dict. """ device = next(self.parameters()).device def _set(name: str, module: StreamingModule): _set_streaming_state_inplace(module._streaming_state, state, prefix=name, device=device) self._apply_named_streaming(_set) if state: raise RuntimeError(f"Some states were not consumed: {list(state.keys())}") def set_streaming_state(self, state: dict[str, Any]): """Set the streaming state, including that of sub-modules.""" state = dict(state) def _set(name: str, module: StreamingModule): if name in state: module._streaming_state = state[name] state.pop(name) else: raise RuntimeError(f"Expected to find a streaming state for {name}.") self._apply_named_streaming(_set) if state: raise RuntimeError(f"Some states were not consumed: {list(state.keys())}") @dataclass class _NullState: pass def reset(self) -> None: pass class StreamingContainer(StreamingModule[_NullState]): def _init_streaming_state(self, batch_size: int) -> _NullState: return _NullState() @dataclass class _StreamingAddState: previous_x: torch.Tensor | None = None previous_y: torch.Tensor | None = None def reset(self): self.previous_x = None self.previous_y = None class StreamingAdd(StreamingModule[_StreamingAddState]): def _init_streaming_state(self, batch_size: int) -> _StreamingAddState: return _StreamingAddState() def forward(self, x: torch.Tensor, y: torch.Tensor): if self._streaming_state is None: return x + y else: prev_x = self._streaming_state.previous_x prev_y = self._streaming_state.previous_y if prev_x is not None: x = torch.cat([prev_x, x], dim=-1) if prev_y is not None: y = torch.cat([prev_y, y], dim=-1) m_l = min(x.shape[-1], y.shape[-1]) self._streaming_state.previous_x = x[..., m_l:] self._streaming_state.previous_y = y[..., m_l:] return x[..., :m_l] + y[..., :m_l] @dataclass class _StreamingConvState: previous: torch.Tensor | None = None def reset(self): self.previous = None class RawStreamingConv1d(torch.nn.Conv1d, StreamingModule[_StreamingConvState]): def __init__(self, *args, **kwargs): super().__init__(*args, **kwargs) assert self.padding[0] == 0, "Padding should be handled outside." assert ( self.stride[0] <= self.kernel_size[0] ), "stride must be less than kernel_size." def _init_streaming_state(self, batch_size: int) -> _StreamingConvState: return _StreamingConvState() def forward(self, input: torch.Tensor) -> torch.Tensor: stride = self.stride[0] # Effective kernel size accounting for dilation. kernel = (self.kernel_size[0] - 1) * self.dilation[0] + 1 if self._streaming_state is None: return super().forward(input) else: # Due to the potential overlap, we might have some cache of the previous time steps. previous = self._streaming_state.previous if previous is not None: input = torch.cat([previous, input], dim=-1) B, C, T = input.shape # We now compute the number of full convolution frames, i.e. the frames # that are ready to be computed. num_frames = max(0, int(math.floor((T - kernel) / stride) + 1)) offset = num_frames * stride # We will compute `num_frames` outputs, and we are advancing by `stride` # for each of the frame, so we know the data before `stride * num_frames` # will never be used again. self._streaming_state.previous = input[..., offset:] if num_frames > 0: input_length = (num_frames - 1) * stride + kernel out = super().forward(input[..., :input_length]) else: # Not enough data as this point to output some new frames. out = torch.empty( B, self.out_channels, 0, device=input.device, dtype=input.dtype ) return out @dataclass class _StreamingConvTrState: partial: torch.Tensor | None = None def reset(self): self.partial = None class RawStreamingConvTranspose1d( torch.nn.ConvTranspose1d, StreamingModule[_StreamingConvTrState] ): def __init__(self, *args, **kwargs): super().__init__(*args, **kwargs) assert self.padding[0] == 0, "Padding should be handled outside." assert self.dilation[0] == 1, "No dilation for now" assert ( self.stride[0] <= self.kernel_size[0] ), "stride must be less than kernel_size." assert self.output_padding[0] == 0, "Output padding not supported." def _init_streaming_state(self, batch_size: int) -> _StreamingConvTrState: return _StreamingConvTrState() def forward(self, x: torch.Tensor) -> torch.Tensor: # type: ignore B, C, T = x.shape stride = self.stride[0] kernel = self.kernel_size[0] if self._streaming_state is None: return super().forward(x) else: if T == 0: return torch.empty( B, self.out_channels, 0, device=x.device, dtype=x.dtype ) out = super().forward(x) OT = out.shape[-1] partial = self._streaming_state.partial if partial is not None: # Due to the potential overlap, the rightmost output of the conv transpose is not # ready to be output, as it will receive contributions from the next input frames. # Here we recover those `partial` output frames. We know that the first time step # of the `partial` tensor corresponds to the first time step of `out` as anything # coming before the first time step of `out` would have been already flushed. PT = partial.shape[-1] if self.bias is not None: out[..., :PT] += partial - self.bias[:, None] else: out[..., :PT] += partial # The input is T, the output is S * (T - 1) + K. # The offset of the left of the next frame will be S * T # so everything between 0 and S * T is ready to be output, and we need # to keep in the internal state everything beyond that, i.e. S (T - 1) + K - S T = K - S invalid_steps = kernel - stride partial = out[..., OT - invalid_steps :] out = out[..., : OT - invalid_steps] self._streaming_state.partial = partial return out def test(): torch.manual_seed(1234) device = "cpu" if torch.cuda.is_available(): # Avoid the cuda optimizations that would take place on single precision # floats for convolutions. torch.backends.cudnn.enabled = True torch.backends.cudnn.benchmark = False torch.backends.cudnn.deterministic = True torch.backends.cuda.matmul.allow_tf32 = False torch.backends.cudnn.allow_tf32 = False device = "cuda:0" kernel_sizes = [1, 3, 4, 8, 15, 16] strides = [1, 2, 3, 4, 5, 6, 7, 8, 9] chin = 6 chout = 12 for kernel, stride in itertools.product(kernel_sizes, strides): if stride > kernel: continue conv = RawStreamingConv1d(chin, chout, kernel, stride).to(device) convtr = RawStreamingConvTranspose1d(chout, chin, kernel, stride).to(device) for length in [4, 8, 32, 54, 65, 128, 1043]: print(f"ksize {kernel} strides {stride} len {length}") if length < kernel: continue batch_size = 3 x = torch.randn(batch_size, chin, length).to(device) y = conv(x) z = convtr(y) for chunk_size in [1, 3, 5, 8]: ys = [] zs = [] with conv.streaming(batch_size), convtr.streaming(batch_size): for offset in range(0, length, chunk_size): chunk = x[..., offset : offset + chunk_size] ys.append(conv(chunk)) zs.append(convtr(ys[-1])) y_stream = torch.cat(ys, dim=-1) z_stream = torch.cat(zs, dim=-1) y = y[..., : y_stream.shape[-1]] z = z[..., : z_stream.shape[-1]] assert y.shape == y_stream.shape, (y.shape, y_stream.shape) delta = (y_stream - y).norm() / y.norm() assert delta <= 1e-6, delta num_frames = int((length - kernel) / stride) + 1 assert num_frames == y_stream.shape[-1] assert z.shape == z_stream.shape, (z.shape, z_stream.shape) delta = (z_stream - z).norm() / z.norm() assert delta <= 1e-6, (delta, (z_stream - z).abs().mean(dim=(0, 1))) if __name__ == "__main__": with torch.no_grad(): test()