core.ssm.causal_conv1d

Causal convolution over contiguous context-parallel sequence shards.

Module Contents

Functions

_exchange_initial_states	Exchange the preceding rank’s tail as the local convolution state.
causal_conv1d_cp	Apply causal Conv1d to a contiguous context-parallel shard.

API

core.ssm.causal_conv1d._exchange_initial_states(

x: torch.Tensor,

state_len: int,

cp_group: torch.distributed.ProcessGroup,

) → torch.Tensor | None

Exchange the preceding rank’s tail as the local convolution state.

All ranks participate in a differentiable ring exchange. Rank 0 zeros the wrapped tail to preserve the global causal boundary.

core.ssm.causal_conv1d.causal_conv1d_cp(

x: torch.Tensor,

weight: torch.Tensor,

bias: torch.Tensor | None,

activation: str | None,

cp_group: torch.distributed.ProcessGroup,

) → torch.Tensor

Apply causal Conv1d to a contiguous context-parallel shard.

Parameters:

• x – Input tensor of shape [B, T, D].

• weight – Depthwise weights of shape [D, W].

• bias – Optional channel-wise bias.

• activation – Optional activation passed to causal_conv1d_fn.

• cp_group – Context-parallel process group ordered by sequence shard.

Returns:

Output tensor of shape [B, T, D].

Raises:

ImportError – If the optional causal-conv1d dependency is unavailable.