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_cpu
causal_conv1d_channel_first_fwd_cpu[x_dtype: DType, weight_dtype: DType, output_dtype: DType, bias_dtype: DType](batch: Int, dim: Int, seqlen: Int, width: Int, x: TileTensor[x_dtype, address_space=x.address_space, linear_idx_type=x.linear_idx_type, element_size=x.element_size], weight: TileTensor[weight_dtype, address_space=weight.address_space, linear_idx_type=weight.linear_idx_type, element_size=weight.element_size], output: TileTensor[output_dtype, address_space=output.address_space, linear_idx_type=output.linear_idx_type, element_size=output.element_size], bias: TileTensor[bias_dtype, address_space=bias.address_space, linear_idx_type=bias.linear_idx_type, element_size=bias.element_size], 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, bias_stride: UInt32, silu_activation: Bool, ctx: Optional[DeviceContext] = None)
CPU implementation of causal conv1d for channel-first layout with bias.
Optimizations:
- Parallelization across batch*channel dimensions using sync_parallelize.
- Pre-loaded weights in registers to reduce memory access.
Args:
- βbatch (
Int): Batch size. - βdim (
Int): Number of channels. - βseqlen (
Int): Sequence length. - βwidth (
Int): Kernel width. - βx (
TileTensor[x_dtype, address_space=x.address_space, linear_idx_type=x.linear_idx_type, element_size=x.element_size]): Input tensor of shape (B, C, L). - βweight (
TileTensor[weight_dtype, address_space=weight.address_space, linear_idx_type=weight.linear_idx_type, element_size=weight.element_size]): Weight tensor of shape (C, W). - βoutput (
TileTensor[output_dtype, address_space=output.address_space, linear_idx_type=output.linear_idx_type, element_size=output.element_size]): Output tensor of shape (B, C, L). - βbias (
TileTensor[bias_dtype, address_space=bias.address_space, linear_idx_type=bias.linear_idx_type, element_size=bias.element_size]): Bias tensor of shape (C,). - β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. - βbias_stride (
UInt32): Stride for the bias tensor. - βsilu_activation (
Bool): Whether to apply SiLU activation. - βctx (
Optional[DeviceContext]): The context to execute the work on.
Was this page helpful?
Thank you! We'll create more content like this.
Thank you for helping us improve!