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_channel_first_fwd_gpu_no_bias
def causal_conv1d_channel_first_fwd_gpu_no_bias[x_dtype: DType, weight_dtype: DType, output_dtype: DType, kNThreads: Int, kWidth: Int, kNElts: Int, x_LT: TensorLayout, weight_LT: TensorLayout, output_LT: TensorLayout](batch: Int, dim: Int, seqlen: Int, width: Int, x: TileTensor[x_dtype, x_LT, MutUntrackedOrigin], weight: TileTensor[weight_dtype, weight_LT, MutUntrackedOrigin], output: TileTensor[output_dtype, output_LT, MutUntrackedOrigin], x_batch_stride: UInt32, x_c_stride: UInt32, x_l_stride: UInt32, weight_c_stride: UInt32, weight_width_stride: UInt32, out_batch_stride: UInt32, out_c_stride: UInt32, out_l_stride: UInt32, silu_activation: Int8)
Optimized causal conv1d implementation for channel first data layout using SIMD operations (no bias).
Key optimizations:
- SIMD vectorization for input/output operations
- Efficient memory access patterns with coalesced loads
- Vectorized weight loading and computation
- Optimized activation function with SIMD operations
- Better thread utilization and memory bandwidth usage
Grid: (ceildiv(seqlen, kNThreads * kNElts), dim, batch) Block: kNThreads
Parameters:
- βx_dtype (
DType): Element type of the input tensorx. - βweight_dtype (
DType): Element type of the weight tensorweight. - βoutput_dtype (
DType): Element type of the output tensoroutput. - βkNThreads (
Int): Number of threads per block used to process the sequence dimension. - βkWidth (
Int): Compile-time convolution kernel width; must match the runtimewidthargument. - βkNElts (
Int): Number of sequence elements each thread processes, used for SIMD vectorization and ILP. - βx_LT (
TensorLayout): TensorLayout of the input tensorx. - βweight_LT (
TensorLayout): TensorLayout of the weight tensorweight. - βoutput_LT (
TensorLayout): TensorLayout of the output tensoroutput.
Args:
- βbatch (
Int): Batch size. - βdim (
Int): Number of channels. - βseqlen (
Int): Sequence length. - βwidth (
Int): Kernel width (must match kWidth compile-time parameter). - βx (
TileTensor[x_dtype, x_LT, MutUntrackedOrigin]): Input tensor of shape (B, C, L). - βweight (
TileTensor[weight_dtype, weight_LT, MutUntrackedOrigin]): Weight tensor of shape (C, W). - βoutput (
TileTensor[output_dtype, output_LT, MutUntrackedOrigin]): Output tensor of shape (B, C, L). - βx_batch_stride (
UInt32): Stride for the batch dimension of the input tensor. - βx_c_stride (
UInt32): Stride for the channel dimension of the input tensor. - βx_l_stride (
UInt32): Stride for the sequence length dimension of the input tensor. - βweight_c_stride (
UInt32): Stride for the channel dimension of the weight tensor. - βweight_width_stride (
UInt32): Stride for the width dimension of the weight tensor. - βout_batch_stride (
UInt32): Stride for the batch dimension of the output tensor. - βout_c_stride (
UInt32): Stride for the channel dimension of the output tensor. - βout_l_stride (
UInt32): Stride for the sequence length dimension of the output tensor. - βsilu_activation (
Int8): Whether to apply SiLU activation (Int8: 0 or 1).
Was this page helpful?
Thank you! We'll create more content like this.
Thank you for helping us improve!