For the complete documentation index, see llms.txt. Markdown versions of all pages are available by appending .md to any URL (e.g. /max/get-started.md).

Mojo function

causal_conv1d_varlen_fwd_cpu

def causal_conv1d_varlen_fwd_cpu[x_dtype: DType, weight_dtype: DType, bias_dtype: DType, output_dtype: DType, cu_seqlens_dtype: DType, cache_indices_dtype: DType, has_initial_state_dtype: DType, conv_states_dtype: DType](dim: Int, total_seqlen: Int, width: Int, batch: Int, x: TileTensor[x_dtype, Storage=x.Storage, address_space=x.address_space, linear_idx_type=x.linear_idx_type], weight: TileTensor[weight_dtype, Storage=weight.Storage, address_space=weight.address_space, linear_idx_type=weight.linear_idx_type], bias: TileTensor[bias_dtype, Storage=bias.Storage, address_space=bias.address_space, linear_idx_type=bias.linear_idx_type], query_start_loc: TileTensor[cu_seqlens_dtype, Storage=query_start_loc.Storage, address_space=query_start_loc.address_space, linear_idx_type=query_start_loc.linear_idx_type], cache_indices: TileTensor[cache_indices_dtype, Storage=cache_indices.Storage, address_space=cache_indices.address_space, linear_idx_type=cache_indices.linear_idx_type], has_initial_state: TileTensor[has_initial_state_dtype, Storage=has_initial_state.Storage, address_space=has_initial_state.address_space, linear_idx_type=has_initial_state.linear_idx_type], conv_states: TileTensor[conv_states_dtype, Storage=conv_states.Storage, address_space=conv_states.address_space, linear_idx_type=conv_states.linear_idx_type], output: TileTensor[output_dtype, Storage=output.Storage, address_space=output.address_space, linear_idx_type=output.linear_idx_type], x_dim_stride: UInt32, x_seqlen_stride: UInt32, weight_dim_stride: UInt32, weight_width_stride: UInt32, out_dim_stride: UInt32, out_seqlen_stride: UInt32, conv_states_batch_stride: UInt32, conv_states_dim_stride: UInt32, conv_states_width_stride: UInt32, silu_activation: Bool, pad_slot_id: Int32, has_cache_indices: Bool, has_initial_state_flag: Bool, has_conv_states: Bool, has_bias: Bool)

Forward pass for causal conv1d with variable length sequences.

Performs causal 1D convolution on variable length sequences that are concatenated together. Uses cumulative sequence lengths to identify sequence boundaries.

This is the CPU reference implementation for causal_conv1d_varlen_fwd.

Parameters:

x_dtype (DType): Data type of the input tensor.
weight_dtype (DType): Data type of the weight tensor.
bias_dtype (DType): Data type of the bias tensor.
output_dtype (DType): Data type of the output tensor.
cu_seqlens_dtype (DType): Data type of the cumulative sequence lengths.
cache_indices_dtype (DType): Data type of the cache indices.
has_initial_state_dtype (DType): Data type of the initial-state flags.
conv_states_dtype (DType): Data type of the convolution states.

Args:

dim (Int): Number of channels in the convolution.
total_seqlen (Int): Total length of the concatenated sequences.
width (Int): Convolution width (number of taps).
batch (Int): Number of sequences in the batch.
x (TileTensor[x_dtype, Storage=x.Storage, address_space=x.address_space, linear_idx_type=x.linear_idx_type]): Input tensor of shape (dim, total_seqlen).
weight (TileTensor[weight_dtype, Storage=weight.Storage, address_space=weight.address_space, linear_idx_type=weight.linear_idx_type]): Weight tensor of shape (dim, width).
bias (TileTensor[bias_dtype, Storage=bias.Storage, address_space=bias.address_space, linear_idx_type=bias.linear_idx_type]): Bias tensor of shape (dim,).
query_start_loc (TileTensor[cu_seqlens_dtype, Storage=query_start_loc.Storage, address_space=query_start_loc.address_space, linear_idx_type=query_start_loc.linear_idx_type]): Cumulative sequence lengths of shape (batch + 1,).
cache_indices (TileTensor[cache_indices_dtype, Storage=cache_indices.Storage, address_space=cache_indices.address_space, linear_idx_type=cache_indices.linear_idx_type]): Per-sequence indices into conv_states of shape (batch,).
has_initial_state (TileTensor[has_initial_state_dtype, Storage=has_initial_state.Storage, address_space=has_initial_state.address_space, linear_idx_type=has_initial_state.linear_idx_type]): Per-sequence flags of shape (batch,) indicating whether to use the initial state.
conv_states (TileTensor[conv_states_dtype, Storage=conv_states.Storage, address_space=conv_states.address_space, linear_idx_type=conv_states.linear_idx_type]): Convolution states of shape (..., dim, width - 1), updated in place.
output (TileTensor[output_dtype, Storage=output.Storage, address_space=output.address_space, linear_idx_type=output.linear_idx_type]): Output tensor of shape (dim, total_seqlen).
x_dim_stride (UInt32): Stride for the channel dimension in x.
x_seqlen_stride (UInt32): Stride for the sequence dimension in x.
weight_dim_stride (UInt32): Stride for the channel dimension in weight.
weight_width_stride (UInt32): Stride for the tap dimension in weight.
out_dim_stride (UInt32): Stride for the channel dimension in output.
out_seqlen_stride (UInt32): Stride for the sequence dimension in output.
conv_states_batch_stride (UInt32): Stride for the batch dimension in conv_states.
conv_states_dim_stride (UInt32): Stride for the channel dimension in conv_states.
conv_states_width_stride (UInt32): Stride for the tap dimension in conv_states.
silu_activation (Bool): Whether to apply the SiLU activation to the output.
pad_slot_id (Int32): Slot ID identifying padded entries to skip.
has_cache_indices (Bool): Whether to consult cache_indices for slot lookup.
has_initial_state_flag (Bool): Whether to consult has_initial_state per sequence.
has_conv_states (Bool): Whether conv_states is provided and should be updated.
has_bias (Bool): Whether to add bias to the convolution sum.