DS4Server/src/engine/metal/qwen_mtplx/vision_kernels.rs

//! Vision-only dispatches from pinned MLX 0.32.2 normalization, conv and SDPA.
//! Rust supplies the same layouts and parameters to the existing Metal library.
use super::super::gpu::{Buffer, MetalConstant};
use super::array::{Array, Dtype, Layout};
use super::{bytes, dispatch_geometry, dispatch_specialized, ops, stream::Stream, tensor};

#[derive(Clone, Copy, PartialEq)]
pub(super) enum Kernel {
    LayerNorm,
    PatchConv,
    Attention,
    Linear,
}

pub(super) fn call(
    kernel: Kernel,
    inputs: &[Array],
    shape: &[i32],
    stream: Stream,
) -> Result<Array, String> {
    Array::make_operation_with_inputs(
        stream,
        inputs,
        Layout::new(shape, Dtype::BF16)?,
        ops::Operation::Vision(kernel),
    )
}

pub(super) fn evaluate(kernel: Kernel, inputs: &[Array], out: &Array) -> Result<(), String> {
    if kernel == Kernel::Linear {
        return super::dense::evaluate_bias(&inputs[0], &inputs[1], Some(&inputs[2]), out);
    }
    let nbytes = out.layout().nbytes() as u64;
    out.set_data(Buffer::mtplx_bytes(nbytes)?)?;
    let buffers = inputs.iter().map(Array::buffer).collect::<Vec<_>>();
    let ob = out.buffer();
    let mut bindings = inputs
        .iter()
        .zip(&buffers)
        .enumerate()
        .map(|(i, (a, b))| {
            tensor(i as u32, b)
                .at_byte_offset(a.offset())
                .input(a.data_size())
        })
        .collect::<Vec<_>>();
    bindings.push(tensor(inputs.len() as u32, &ob).output(out.data_size()));
    match kernel {
        Kernel::Linear => {
            super::dense::evaluate_bias(&inputs[0], &inputs[1], Some(&inputs[2]), out)
        }
        Kernel::LayerNorm => {
            let axis = out.layout().dim(-1)? as u32;
            let rows = out.layout().size() as u32 / axis;
            let threads = axis.div_ceil(8).div_ceil(32) * 32;
            let eps = 1e-6_f32;
            let stride = 1_u32;
            bindings.extend([
                bytes(4, &eps),
                bytes(5, &axis),
                bytes(6, &stride),
                bytes(7, &stride),
            ]);
            dispatch_geometry(
                "layer_normbfloat16",
                &bindings,
                &[],
                [rows * threads, 1, 1],
                [threads, 1, 1],
                true,
            )
        }
        Kernel::PatchConv => {
            // Inputs are contiguous [patches,2,16,16,16], weights
            // [1152,2,16,16,16]. Original Conv3d pads RGB channels to 16.
            #[repr(C)]
            struct Params {
                n: i32,
                c: i32,
                o: i32,
                input: [i32; 3],
                weight: [i32; 3],
                output: [i32; 3],
                stride: [i32; 3],
                pad: [i32; 3],
                kdil: [i32; 3],
                idil: [i32; 3],
                is: [i64; 5],
                ws: [i64; 5],
                os: [i64; 5],
                groups: i32,
                flip: u8,
                padding: [u8; 3],
            }
            let n = out.layout().dim(0)?;
            let params = Params {
                n,
                c: 16,
                o: 1152,
                input: [2, 16, 16],
                weight: [2, 16, 16],
                output: [1, 1, 1],
                stride: [2, 16, 16],
                pad: [0; 3],
                kdil: [1; 3],
                idil: [1; 3],
                is: [8192, 4096, 256, 16, 1],
                ws: [8192, 4096, 256, 16, 1],
                os: [1152, 1152, 1152, 1152, 1],
                groups: 1,
                flip: 0,
                padding: [0; 3],
            };
            let gemm = [n, 1152, 8192, 512, 16, 16, 16, -8160, 18, (n + 31) / 32, 0];
            bindings.extend([bytes(3, &params), bytes(4, &gemm)]);
            dispatch_specialized(
                "implicit_gemm_conv_3d_bfloat16_bm32_bn64_bk16_wm2_wn2_filter_s",
                &bindings,
                &[],
                [18, ((n + 31) / 32) as u32, 1],
                [32, 2, 2],
            )
        }
        Kernel::Attention => {
            #[repr(C)]
            struct Params {
                b: i32,
                h: i32,
                d: i32,
                ql: i32,
                kl: i32,
                gqa: i32,
                scale: f32,
                nq: i32,
                nk: i32,
                aq: i32,
                ak: i32,
                rq: i32,
                rk: i32,
                off: i32,
                strides: [[i64; 3]; 4],
            }
            let shape = out.layout().shape().to_vec();
            let n = shape[2];
            let mut strides = [[0; 3]; 4];
            for (s, a) in strides
                .iter_mut()
                .zip(inputs.iter().chain(std::iter::once(out)))
            {
                s.copy_from_slice(&a.layout().strides()[..3]);
            }
            let params = Params {
                b: 1,
                h: 16,
                d: 80,
                ql: n,
                kl: n,
                gqa: 1,
                scale: (72_f64.powf(-0.5)) as f32,
                nq: ((n + 31) / 32),
                nk: ((n + 31) / 32),
                aq: n / 32,
                ak: n / 32,
                rq: n % 32,
                rk: n % 32,
                off: 0,
                strides,
            };
            let constants = [
                (200, n % 32 == 0),
                (201, n % 32 == 0),
                (300, false),
                (301, false),
                (302, false),
            ]
            .map(|(index, v)| MetalConstant {
                index,
                value: u32::from(v),
                kind: 0,
            });
            bindings.push(bytes(4, &params));
            dispatch_specialized(
                "steel_attention_bfloat16_bq32_bk32_bd80_wm4_wn1_maskbfloat16",
                &bindings,
                &constants,
                [((n + 31) / 32) as u32, 16, 1],
                [32, 4, 1],
            )
        }
    }
}