Skip to main content

Module gather

Module gather 

Source
Expand description

GPU-accelerated gather / index_select along dim=0.

Gathers rows from a 2D source tensor using an index array: output[i, :] = src[indices[i], :]

Used for MoE scale factor gathering by expert index.

Statics§

GATHER_SHADER_SOURCE
MSL source for the gather kernel (embedded at compile time).

Functions§

dispatch_gather_f32
Dispatch a gather / index_select operation on the GPU.
register
Register gather shader source with the given kernel registry.