`core.ssm.ops.common.causal_conv1d_varlen`#

Triton varlen depthwise causal 1D convolution with per-sequence initial states and fused SiLU.

Supports packed variable-length sequences where causal_conv1d_fn cannot accept both seq_idx and initial_states simultaneously.

Module Contents#

Functions#

`_causal_conv1d_varlen_kernel`	Depthwise causal conv1d over packed varlen sequences with initial states and SiLU.
`causal_conv1d_varlen_fn`	Depthwise causal 1D convolution over packed variable-length sequences.
`_causal_conv1d_varlen_simple`	Simple PyTorch implementation of varlen causal conv1d with initial states and SiLU.
`_causal_conv1d_carry_states_kernel`	Per-request conv state after appending a slice, carrying prior history.
`causal_conv1d_varlen_carry_states`	Per-request conv state after appending a slice, carrying prior history.

API#

core.ssm.ops.common.causal_conv1d_varlen._causal_conv1d_varlen_kernel( x_ptr, weight_ptr, bias_ptr, seq_idx_ptr, seq_start_ptr, initial_states_ptr, out_ptr, total_tokens, conv_dim: triton.language.constexpr, initial_states_stride_req, initial_states_stride_dim, WIDTH: triton.language.constexpr, BLOCK_T: triton.language.constexpr, BLOCK_C: triton.language.constexpr, HAS_INITIAL_STATES: triton.language.constexpr, HAS_BIAS: triton.language.constexpr, )#

Depthwise causal conv1d over packed varlen sequences with initial states and SiLU.

Fully vectorized over BLOCK_T tokens x BLOCK_C channels per thread block.

core.ssm.ops.common.causal_conv1d_varlen.causal_conv1d_varlen_fn( x: torch.Tensor, weight: torch.Tensor, bias: torch.Tensor, cu_seqlens: torch.Tensor, initial_states: torch.Tensor = None, activation: str = 'silu', precomputed_seq_idx: torch.Tensor = None, precomputed_seq_start: torch.Tensor = None, ) → torch.Tensor#

Depthwise causal 1D convolution over packed variable-length sequences.

Supports both cu_seqlens (sequence boundaries) and initial_states simultaneously, unlike causal_conv1d_fn which requires mutual exclusivity between seq_idx and initial_states.

Parameters:

x – Input tensor of shape (total_tokens, conv_dim), channels-last packed.
weight – Convolution weights of shape (conv_dim, d_conv).
bias – Bias of shape (conv_dim,), or None for a bias-free convolution (Gated Delta Product layers default to conv_bias=False).
cu_seqlens – Cumulative sequence lengths of shape (num_requests + 1,), int32.
initial_states – Per-request initial conv states of shape (num_requests, conv_dim, d_conv - 1). If None, uses zeros.
activation – Activation function, must be “silu”.
precomputed_seq_idx – Precomputed per-token request ID of shape (total_tokens,). If provided, skips repeat_interleave (CUDA graph compatible). Padding tokens should use 0 as sentinel.
precomputed_seq_start – Precomputed per-token request start position of shape (total_tokens,). Must be provided together with precomputed_seq_idx.

Returns:

Output tensor of shape (total_tokens, conv_dim).

core.ssm.ops.common.causal_conv1d_varlen._causal_conv1d_varlen_simple( x: torch.Tensor, weight: torch.Tensor, bias: torch.Tensor, cu_seqlens: torch.Tensor, initial_states: torch.Tensor, out: torch.Tensor, ) → None#

Simple PyTorch implementation of varlen causal conv1d with initial states and SiLU.

This is a reference implementation for testing. Processes each request and token sequentially.

core.ssm.ops.common.causal_conv1d_varlen._causal_conv1d_carry_states_kernel( x_ptr, cu_seqlens_ptr, previous_states_ptr, out_ptr, total_tokens, conv_dim, x_stride_token, x_stride_dim, prev_stride_req, prev_stride_dim, out_stride_req, out_stride_dim, D_CONV: triton.language.constexpr, BLOCK_C: triton.language.constexpr, )#

Per-request conv state after appending a slice, carrying prior history.

One program handles one request and a block of channels. The column-by-column rule is in causal_conv1d_varlen_carry_states; here it is a static_range over D_CONV, so the index arithmetic lands in registers and needs no precomputed plan.

core.ssm.ops.common.causal_conv1d_varlen.causal_conv1d_varlen_carry_states( x: torch.Tensor, cu_seqlens: torch.Tensor, previous_states: torch.Tensor, ) → torch.Tensor#

Per-request conv state after appending a slice, carrying prior history.

A request’s conv state holds the last d_conv tokens it has seen. It exists so that the next call can convolve across the boundary: the causal convolution at token t reads tokens t - d_conv + 1 .. t, so continuing a sequence requires the tokens that came before the current input.

When a prompt is processed one slice at a time, the new state is generally a mix of both sources – the tail of the slice, and whatever the request was already carrying. With d_conv = 4, a request whose incoming state is [a b c d] and whose slice this call is [e f]:

incoming state   a b c d
slice                    e f
new state            c d e f

Two columns survive from the incoming state, two come from the slice. The same request handed a slice of [e f g h] or longer takes all four columns from the slice, and the incoming state does not contribute at all:

incoming state   a b c d
slice                    e f g h
new state                e f g h

So the incoming state matters exactly when a slice is shorter than d_conv. That is not an edge case to design around – a prompt of any length can end in a short slice, since the split points are chosen by a token budget that knows nothing about d_conv. Reading only the slice in the first example would produce [0 0 e f], and the next call would convolve e and f against zeros instead of against c and d.

Zero-length (padding) requests take every column from the incoming state and so reproduce it unchanged.

The launch geometry is fixed by the padded request count, conv_dim and d_conv, and no value crosses to the host, so this is safe to capture in a CUDA graph.