Is this page helpful?

core.ssm.causal_conv1d#

Causal convolution over contiguous context-parallel sequence shards, and the determinism guard the SSM mixers apply to the causal_conv1d backward.

Module Contents#

Functions#

_use_causal_conv1d_deterministic_mode

Whether causal_conv1d’s backward will take its deterministic reduction.

assert_causal_conv1d_deterministic

Refuse a deterministic run whose convolution cannot be bit-reproducible.

_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._use_causal_conv1d_deterministic_mode()#

Whether causal_conv1d’s backward will take its deterministic reduction.

Mirrors the kernel’s own use_deterministic_mode() (csrc/causal_conv1d.cpp): only a leading '1' or '0' decides, anything else falls through to torch.

core.ssm.causal_conv1d.assert_causal_conv1d_deterministic(deterministic_mode)#

Refuse a deterministic run whose convolution cannot be bit-reproducible.

The conv backward combines each weight-gradient element’s per-block partials with atomicAdd, which fixes no order; 1.6.0+ uses a per-block workspace and an ordered reduce instead. Worst on the channel-last layout GDP and Mamba’s fused path both feed the conv, where an element takes batch * ceil(seqlen / 128) contributions rather than the channels-first batch.

Call once at construction. Keyed on deterministic_mode, not torch’s global flag, which unrelated tests set and never restore.

core.ssm.causal_conv1d._exchange_initial_states(
x: torch.Tensor,
state_len: int,
cp_group: torch.distributed.ProcessGroup,
initial_state_mask: torch.Tensor | None = None,
) → 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. The optional mask removes tail tokens outside the first local packed sequence.

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,
global_seq_idx: torch.Tensor | None = None,
) → torch.Tensor#

Apply causal Conv1d to a contiguous context-parallel shard.

Parameters:

  • x – Input tensor of shape [B, T, D]. THD callers flatten packed tokens along T and use B=1.

  • 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.

  • global_seq_idx – Global per-token sequence IDs for packed THD input, replicated across CP ranks. IDs must be non-negative. Pass None for non-packed input. The convolution state resets at each sequence boundary.

Returns:

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

Raises:

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

  • ValueError – If global_seq_idx has an invalid shape, dtype, or device.