| |
| |
| |
| |
| |
| |
|
|
| """Short convolution implementation for efficient causal convolutions.""" |
|
|
| import warnings |
|
|
| import torch |
| import torch.nn as nn |
| from einops import rearrange |
|
|
| try: |
| from causal_conv1d import causal_conv1d_fn as causal_conv1d_fn_cuda |
| from causal_conv1d import causal_conv1d_update as causal_conv1d_update_cuda |
| except ImportError: |
| causal_conv1d_fn_cuda = None |
| causal_conv1d_update_cuda = None |
|
|
|
|
| class ShortConvolution(nn.Conv1d): |
| """Short convolution layer for efficient causal convolution operations. |
| |
| This class implements a depthwise 1D convolution with causal padding, |
| designed for efficient sequence processing. It supports multiple backends (Triton/CUDA) |
| and optional activation functions. |
| |
| Args: |
| hidden_size (int): Number of input/output channels (must be equal for depthwise conv) |
| kernel_size (int): Size of the convolution kernel |
| bias (bool, optional): Whether to include learnable bias. Defaults to False. |
| activation (Optional[str], optional): Activation function ('silu' or 'swish'). Defaults to 'silu'. |
| backend (Optional[str], optional): Backend implementation ('triton' or 'cuda'). Defaults to 'triton'. |
| device (Optional[torch.device], optional): Device to place the layer on. Defaults to None. |
| dtype (Optional[torch.dtype], optional): Data type for layer parameters. Defaults to None. |
| **kwargs: Additional keyword arguments (deprecated 'use_fast_conv1d' supported for compatibility) |
| |
| Attributes: |
| hidden_size (int): Number of channels |
| activation (Optional[str]): Selected activation function |
| backend (str): Actual backend being used (may differ from input due to availability) |
| |
| Note: |
| - Uses depthwise convolution (groups=hidden_size) for efficiency |
| - Applies causal padding (kernel_size-1) to ensure no future information leakage |
| - Falls back to Triton backend if CUDA backend is unavailable |
| """ |
|
|
| def __init__( |
| self, |
| hidden_size: int, |
| kernel_size: int, |
| bias: bool = False, |
| activation: str | None = 'silu', |
| backend: str | None = 'triton', |
| device: torch.device | None = None, |
| dtype: torch.dtype | None = None, |
| **kwargs, |
| ): |
| super().__init__( |
| in_channels=hidden_size, |
| out_channels=hidden_size, |
| kernel_size=kernel_size, |
| groups=hidden_size, |
| bias=bias, |
| padding=kernel_size - 1, |
| device=device, |
| dtype=dtype, |
| ) |
|
|
| self.hidden_size = hidden_size |
| self.activation = None |
|
|
| if activation is not None: |
| assert activation in ['silu', 'swish'], f"Activation `{activation}` not supported yet." |
| self.activation = activation |
|
|
| if 'use_fast_conv1d' in kwargs: |
| warnings.warn( |
| "The `use_fast_conv1d` parameter is deprecated and will be ignored. " |
| "Please use the `backend` parameter instead.", |
| ) |
| import os |
| self.backend = os.environ.get('FLA_CONV_BACKEND', backend) |
| if backend not in ['cuda', 'triton']: |
| raise ValueError(f"Invalid backend: {backend}, must be one of ['cuda', 'triton']") |
| if backend == 'cuda': |
| if causal_conv1d_fn_cuda is None: |
| warnings.warn( |
| "The `backend` parameter is set to `cuda`, but `causal_conv1d_fn` is not available. " |
| "Switching to the Triton implementation instead. " |
| "Consider installing `causal_conv1d` to enable the CUDA backend.", |
| ) |
| self.backend = 'triton' |
|
|
| def extra_repr(self): |
| s = ('{in_channels}, {out_channels}, kernel_size={kernel_size}' |
| ', stride={stride}') |
| if self.padding != (0,) * len(self.padding): |
| s += ', padding={padding}' |
| if self.dilation != (1,) * len(self.dilation): |
| s += ', dilation={dilation}' |
| if self.output_padding != (0,) * len(self.output_padding): |
| s += ', output_padding={output_padding}' |
| if self.groups != 1: |
| s += ', groups={groups}' |
| if self.bias is None: |
| s += ', bias=False' |
| if self.padding_mode != 'zeros': |
| s += ', padding_mode={padding_mode}' |
| if self.activation is not None: |
| s += ', activation={activation}' |
| s += f', backend={self.backend}' |
| return s.format(**self.__dict__) |
|
|
| def forward( |
| self, |
| x: torch.Tensor, |
| residual: torch.Tensor | None = None, |
| mask: torch.Tensor | None = None, |
| cache: torch.Tensor | None = None, |
| output_final_state: bool = False, |
| cu_seqlens: torch.LongTensor | None = None, |
| chunk_indices: torch.LongTensor | None = None, |
| **kwargs, |
| ) -> tuple[torch.Tensor, torch.Tensor | None]: |
| """ |
| Args: |
| x (`torch.Tensor`): |
| Tensor of shape `[B, T, D]`. `B` must be 1 if `cu_seqlens` is provided. |
| residual (`Optional[torch.Tensor]`): |
| Residual tensor of shape `[B, T, D]`. Default: `None`. |
| mask (`Optional[torch.Tensor]`): |
| Attention mask dealing with padded positions. |
| cache (`Optional[torch.Tensor]`): |
| Previous cache tensor of shape `[N, D, W]`, where `W` is the kernel size. |
| If provided, the cache is updated **inplace**. |
| output_final_state (Optional[bool]): |
| Whether to output the final state of shape `[N, D, W]`. Default: `False`. |
| cu_seqlens (Optional[torch.LongTensor]): |
| Cumulative sequence lengths for each batch. Used for varlen. Default: `None`. |
| Shape: [B+1] |
| chunk_indices (Optional[torch.LongTensor]): |
| Chunk indices for variable-length sequences. Default: `None`. |
| |
| Returns: |
| Tensor of shape `[B, T, D]`. |
| """ |
| |
| from ...modules.conv.causal_conv1d import causal_conv1d |
|
|
| B, T, *_ = x.shape |
| N = B if cu_seqlens is None else len(cu_seqlens) - 1 |
| if mask is not None: |
| if cu_seqlens is not None: |
| raise ValueError("`mask` and `cu_seqlens` cannot be provided at the same time") |
| x = x.mul_(mask.unsqueeze(-1)) |
|
|
| |
| if B * T == N: |
| y, cache = self.step( |
| x=x, |
| residual=residual, |
| cache=cache, |
| output_final_state=output_final_state, |
| cu_seqlens=cu_seqlens, |
| ) |
| return y, cache |
|
|
| |
| |
| |
| |
| |
| if self.backend == 'cuda' and cache is not None: |
| warnings.warn( |
| "The CUDA backend does not support both `cu_seqlens` and `cache` being provided, " |
| "or both `cu_seqlens` and `output_final_state` being provided. " |
| "Switching to the Triton backend instead. ", |
| stacklevel=2, |
| ) |
| self.backend = 'triton' |
|
|
| return causal_conv1d( |
| x=x, |
| weight=rearrange(self.weight, "d 1 w -> d w"), |
| bias=self.bias, |
| residual=residual, |
| initial_state=cache, |
| output_final_state=output_final_state, |
| activation=self.activation, |
| backend=self.backend, |
| cu_seqlens=cu_seqlens, |
| chunk_indices=chunk_indices, |
| **kwargs, |
| ) |
|
|
| def step( |
| self, |
| x: torch.Tensor, |
| residual: torch.Tensor | None, |
| cache: torch.Tensor | None, |
| output_final_state: bool = False, |
| cu_seqlens: torch.LongTensor | None = None, |
| ) -> tuple[torch.Tensor, torch.Tensor | None]: |
| from ...modules.conv.triton.ops import causal_conv1d_update |
|
|
| B, _, D, W = *x.shape, self.kernel_size[0] |
| N = B if cu_seqlens is None else len(cu_seqlens) - 1 |
| |
| |
| |
| if cache is None: |
| cache = x.new_zeros(N, D, W) |
| |
| if self.backend == 'triton': |
| y, cache = causal_conv1d_update( |
| x=x, |
| cache=cache, |
| residual=residual, |
| weight=rearrange(self.weight, "d 1 w -> d w"), |
| bias=self.bias, |
| activation=self.activation, |
| ) |
| return y, (cache if output_final_state else None) |
|
|
| shape = x.shape |
| x = x.squeeze(0) if cu_seqlens is not None else x.squeeze(1) |
| |
| |
| |
| |
| y = causal_conv1d_update_cuda( |
| x=x, |
| conv_state=cache, |
| weight=rearrange(self.weight, "d 1 w -> d w"), |
| bias=self.bias, |
| activation=self.activation, |
| ) |
| y = y.view(shape) |
| if residual is not None: |
| y.add_(residual) |
| return y, (cache if output_final_state else None) |
|
|
| @property |
| def state_size(self) -> int: |
| return self.hidden_size * self.kernel_size |
|
|
