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

spatial_merge_kernel

def spatial_merge_kernel[dtype: DType, InputLayoutType: TensorLayout, input_origin: ImmOrigin, OutputLayoutType: TensorLayout, output_origin: MutOrigin, GridThwLayoutType: TensorLayout, grid_thw_origin: ImmOrigin](output: TileTensor[dtype, OutputLayoutType, output_origin], input: TileTensor[dtype, InputLayoutType, input_origin], grid_thw: TileTensor[DType.int64, GridThwLayoutType, grid_thw_origin], batch_size: Int, hidden_size: Int, merge_size: Int)

Spatial merge kernel.

Grid: 1D over all output patches (one block per output patch). Threads: loop over channels (hidden_size x merge_size^2).

Parameters:

dtype (DType): Element type of the input and output tensors.
InputLayoutType (TensorLayout): Compile-time TensorLayout of the input tensor.
input_origin (ImmOrigin): Immutable origin of the input tensor.
OutputLayoutType (TensorLayout): Compile-time TensorLayout of the output tensor.
output_origin (MutOrigin): Mutable origin of the output tensor.
GridThwLayoutType (TensorLayout): Compile-time TensorLayout of the grid_thw tensor.
grid_thw_origin (ImmOrigin): Immutable origin of the grid_thw tensor.

Args:

output (TileTensor[dtype, OutputLayoutType, output_origin]): Output tensor.
input (TileTensor[dtype, InputLayoutType, input_origin]): Input tensor.
grid_thw (TileTensor[DType.int64, GridThwLayoutType, grid_thw_origin]): Grid dimensions tensor (B, 3) containing [t, h, w] for each item.
batch_size (Int): Number of items in batch.
hidden_size (Int): Hidden dimension size.
merge_size (Int): Size of spatial merge blocks.