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

Mojo function

causal_conv1d_varlen_fwd_seqparallel_gpu

def causal_conv1d_varlen_fwd_seqparallel_gpu[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, WIDTH: Int, BLOCK_DIM: Int, TILE_SEQ: Int, x_LT: TensorLayout, weight_LT: TensorLayout, bias_LT: TensorLayout, query_start_loc_LT: TensorLayout, cache_indices_LT: TensorLayout, has_initial_state_LT: TensorLayout, conv_states_LT: TensorLayout, output_LT: TensorLayout](dim_dev: Int32, total_seqlen_dev: Int32, batch_dev: Int32, x: TileTensor[x_dtype, x_LT, MutUntrackedOrigin], weight: TileTensor[weight_dtype, weight_LT, MutUntrackedOrigin], bias: TileTensor[bias_dtype, bias_LT, MutUntrackedOrigin], query_start_loc: TileTensor[cu_seqlens_dtype, query_start_loc_LT, MutUntrackedOrigin], cache_indices: TileTensor[cache_indices_dtype, cache_indices_LT, MutUntrackedOrigin], has_initial_state: TileTensor[has_initial_state_dtype, has_initial_state_LT, MutUntrackedOrigin], conv_states: TileTensor[conv_states_dtype, conv_states_LT, MutUntrackedOrigin], output: TileTensor[output_dtype, output_LT, MutUntrackedOrigin], 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: Int8, pad_slot_id: Int32, has_cache_indices: Int8, has_initial_state_flag: Int8, has_conv_states: Int8, has_bias: Int8)

GPU kernel for causal conv1d forward with variable length sequences, sequence-parallel prefill variant (NVIDIA B200/sm_100, generic elsewhere).

Grid: (batch, ceildiv(dim, BLOCK_DIM), num_tiles_ub) Block: (BLOCK_DIM, 1)

Each grid-(x,y,z) block handles one (sequence, channel-tile, seq-tile). The z-dimension tiles the sequence into TILE_SEQ-sized chunks so a long prefill sequence is spread across many blocks instead of walking the whole sequence serially in one thread (see causal_conv1d_varlen_fwd_gpu, which stays byte-identical and is reused verbatim for decode). Depthwise conv (WIDTH<=4) has no cross-position recurrence: the causal gather reads GLOBAL read-only x across tile boundaries, so tiles need no shared memory or cross-block synchronization. conv_states is written exactly once, by the tail tile of each sequence (tile_end == seqlen).

Dispatched from CausalConv1DVarlenFwd.launch_gpu only when total_seqlen > batch (i.e. at least one multi-token prefill segment is present); pure decode (total_seqlen == batch) keeps using the serial kernel above unmodified. Per Kernels/claude_kb patterns/kv-buffer-pipeline-style host-vs-device tiling notes: num_tiles_ub is a safe host-side upper bound (ceildiv(total_seqlen, TILE_SEQ) + batch) that avoids a host-side max-reduction over ragged per-sequence lengths; blocks whose z-index exceeds a given sequence's actual tile count early-return.

Note: silu_activation and flag parameters are Int8 (0 or 1) instead of Bool for DevicePassable compatibility on GPU.