DS4Server/metal/mtplx_qwen.metal

// MTPLX bodies: Copyright 2026 MTPLX. SPDX-License-Identifier: Apache-2.0
// Dependency bodies carry their own license notice below.
// Generated by tools/mtplx-kernel-source.py; do not edit bodies.
// MTPLX revision: e652d55e2652137a4abcf1312357abbf3eb9d692
// Only entry-point ABI and template instantiations are adapted.
#include <metal_stdlib>
using namespace metal;

// Source: mtplx/kernels/hyper_connection.py::_SOURCE
// File SHA256: 561b635a4bb36108b617ce2a5e0eda2b1868a8b2242d00112e1858bbf1fe89ea
// Body SHA256: 84beb52939e8965c8038a049bef36b7a10f9d368779883c214770122cff2da86
template <typename T, int HAS_INJECT>
kernel void kernel_qwen_mtplx_hyper_read(
    device const T* x [[buffer(0)]],
    device const T* gamma [[buffer(1)]],
    device const T* wd [[buffer(2)]],
    device const T* wu [[buffer(3)]],
    device const T* wi [[buffer(4)]],
    device T* mixed [[buffer(5)]],
    device T* inject [[buffer(6)]],
    uint3 thread_position_in_grid [[thread_position_in_grid]],
    uint3 thread_position_in_threadgroup [[thread_position_in_threadgroup]]) {
    constexpr int HC = 4;
    constexpr int D = 2560;
    constexpr int HCD = HC * D;
    constexpr int R = 320;
    constexpr float EPS = 1e-6f;

    const uint row = thread_position_in_grid.y;
    const uint tid = thread_position_in_threadgroup.x;
    const uint sg = tid / 32;
    const uint lane = tid % 32;

    threadgroup float gsum[32];
    threadgroup float nrm_scale[HC];
    threadgroup T normed[HCD];
    threadgroup float mixv[R];

    device const T* xr = x + (size_t)row * HCD;

    // ---- 1) per-group sum of squares (group g owns threads [g*256,(g+1)*256))
    const int g = tid >> 8;
    const int t_in = tid & 255;
    float ss = 0.0f;
    for (int k = 0; k < 10; ++k) {
        const int i = g * D + t_in + k * 256;
        const float v = (float)xr[i];
        ss += v * v;
    }
    ss = simd_sum(ss);
    if (lane == 0) gsum[sg] = ss;
    threadgroup_barrier(mem_flags::mem_threadgroup);
    if (tid < HC) {
        float tot = 0.0f;
        for (int j = 0; j < 8; ++j) tot += gsum[tid * 8 + j];
        nrm_scale[tid] = metal::rsqrt(tot / (float)D + EPS);
    }
    threadgroup_barrier(mem_flags::mem_threadgroup);

    // ---- 2) normed = x * scale_g * gamma (bf16 boundary like the module)
    for (int k = 0; k < 10; ++k) {
        const int i = g * D + t_in + k * 256;
        const float nv = (float)((T)((float)xr[i] * nrm_scale[g]));
        normed[i] = (T)(nv * (float)gamma[i]);
    }
    threadgroup_barrier(mem_flags::mem_threadgroup);

    // ---- 3) down proj rows (320 = 32 simdgroups x 10), /HC, silu
    for (int rr = 0; rr < 10; ++rr) {
        const int orow = sg * 10 + rr;
        float acc = 0.0f;
        device const T* wrow = wd + (size_t)orow * HCD;
        for (int i = lane; i < HCD; i += 32) {
            acc += (float)normed[i] * (float)wrow[i];
        }
        acc = simd_sum(acc);
        if (lane == 0) {
            float t0 = (float)((T)acc);        // Linear out cast
            t0 = (float)((T)(t0 * 0.25f));     // /hc_count
            mixv[orow] = (float)((T)(t0 / (1.0f + metal::exp(-t0))));  // silu
        }
    }
    threadgroup_barrier(mem_flags::mem_threadgroup);

    // ---- 4) inject = 2*sigmoid((Wi @ normed)/HC) — BEFORE the mix, which
    // repurposes the `normed` staging below.
    if (HAS_INJECT && sg < HC) {
        float acc = 0.0f;
        device const T* irow = wi + (size_t)sg * HCD;
        for (int i = lane; i < HCD; i += 32) {
            acc += (float)normed[i] * (float)irow[i];
        }
        acc = simd_sum(acc);
        if (lane == 0) {
            float t1 = (float)((T)acc);
            t1 = (float)((T)(t1 * 0.25f));
            inject[(size_t)row * HC + sg] = (T)(2.0f / (1.0f + metal::exp(-t1)));
        }
    }
    threadgroup_barrier(mem_flags::mem_threadgroup);

    // ---- 5a) up proj + sigmoid, gated product written back over `normed`
    // (one simdgroup per output row: lanes stride R — coalesced wu reads)
    for (int pos = (int)sg; pos < HCD; pos += 32) {
        device const T* urow = wu + (size_t)pos * R;
        float a = 0.0f;
        for (int j = lane; j < R; j += 32) a += mixv[j] * (float)urow[j];
        a = simd_sum(a);
        if (lane == 0) {
            a = (float)((T)a);                                   // Linear out
            float w = 1.0f / (1.0f + metal::exp(-a));            // sigmoid
            w = (float)((T)w);
            normed[pos] = (T)(w * (float)normed[pos]);           // bf16 mul
        }
    }
    threadgroup_barrier(mem_flags::mem_threadgroup);

    // ---- 5b) mean over the HC groups
    for (int d0 = tid; d0 < D; d0 += 1024) {
        const float acc_mix = (float)normed[d0] + (float)normed[D + d0]
            + (float)normed[2 * D + d0] + (float)normed[3 * D + d0];
        mixed[(size_t)row * D + d0] = (T)(acc_mix * 0.25f);
    }
}
typedef decltype(kernel_qwen_mtplx_hyper_read<bfloat, 0>) kernel_qwen_mtplx_hyper_read_mix_bf16_type;
template [[host_name("kernel_qwen_mtplx_hyper_read_mix_bf16")]]
kernel kernel_qwen_mtplx_hyper_read_mix_bf16_type kernel_qwen_mtplx_hyper_read<bfloat, 0>;
typedef decltype(kernel_qwen_mtplx_hyper_read<bfloat, 1>) kernel_qwen_mtplx_hyper_read_inj_bf16_type;
template [[host_name("kernel_qwen_mtplx_hyper_read_inj_bf16")]]
kernel kernel_qwen_mtplx_hyper_read_inj_bf16_type kernel_qwen_mtplx_hyper_read<bfloat, 1>;

// Source: mtplx/kernels/hyper_connection_v3.py::_SRC_R1
// File SHA256: bc4cc7c59403845c31881f7526dade0f80a7704303c6c5273a794d605cf4e82a
// Body SHA256: 0df4ae0ff5df0352cb1ba2ec67277763ebce77a4425c3004767c19c10895a9cb
template <typename T>
kernel void kernel_qwen_mtplx_hyper_v3_r1(
    device const T* x [[buffer(0)]],
    device const T* wn [[buffer(1)]],
    device const uint32_t* qw [[buffer(2)]],
    device const T* qs [[buffer(3)]],
    device const T* qb [[buffer(4)]],
    device T* mix_out [[buffer(5)]],
    device T* inject_out [[buffer(6)]],
    device float* rms_out [[buffer(7)]],
    uint3 thread_position_in_threadgroup [[thread_position_in_threadgroup]],
    uint3 threadgroup_position_in_grid [[threadgroup_position_in_grid]]) {
    constexpr int K = 10240;
    constexpr int GROUP = 2560;          // hc group width
    constexpr int QGS = 64;              // quant group size
    constexpr int NGROUPS = K / QGS;     // 160 quant groups per row
    constexpr int N_DOWN = 320;
    constexpr int N_TOTAL = 324;

    const uint tid = thread_position_in_threadgroup.x;
    const uint sg = tid / 32;
    const uint lane = tid % 32;
    const uint n = threadgroup_position_in_grid.x * 32 + sg;

    threadgroup float tg_sums[4];
    threadgroup float tg_partial[32];

    // Cooperative x^2 group sums (all 1024 threads, 10 elems each).
    float part[4] = {0.0f, 0.0f, 0.0f, 0.0f};
    for (int i = tid; i < K; i += 1024) {
        const float xv = (float)x[i];
        part[i / GROUP] += xv * xv;
    }
    for (int g = 0; g < 4; ++g) {
        float v = simd_sum(part[g]);
        if (lane == 0) tg_partial[sg] = v;
        threadgroup_barrier(mem_flags::mem_threadgroup);
        if (sg == 0) {
            float acc = (lane < 32) ? tg_partial[lane] : 0.0f;
            acc = simd_sum(acc);
            if (lane == 0) tg_sums[g] = acc;
        }
        threadgroup_barrier(mem_flags::mem_threadgroup);
    }
    float rms[4];
    for (int g = 0; g < 4; ++g) {
        rms[g] = metal::rsqrt(tg_sums[g] / (float)GROUP + 1e-6f);
    }
    if (threadgroup_position_in_grid.x == 0 && tid < 4) {
        rms_out[tid] = rms[tid];
    }
    if (n >= (uint)N_TOTAL) return;

    const device uint32_t* wrow = qw + (size_t)n * (K / 4);
    const device T* srow = qs + (size_t)n * NGROUPS;
    const device T* brow = qb + (size_t)n * NGROUPS;
    float acc = 0.0f;
    for (int g = lane; g < NGROUPS; g += 32) {
        const float s = (float)srow[g];
        const float b = (float)brow[g];
        const device uint32_t* wg = wrow + g * (QGS / 4);
        const int base = g * QGS;
        float qacc = 0.0f;
        float nsum = 0.0f;
        for (int wi = 0; wi < QGS / 4; ++wi) {
            const uint32_t word = wg[wi];
            const int i0 = base + wi * 4;
            for (int by = 0; by < 4; ++by) {
                const int i = i0 + by;
                const float nx = (float)x[i] * (float)wn[i] * rms[i / GROUP];
                qacc += (float)((word >> (8 * by)) & 0xFF) * nx;
                nsum += nx;
            }
        }
        acc += s * qacc + b * nsum;
    }
    acc = simd_sum(acc);
    if (lane == 0) {
        if (n < (uint)N_DOWN) {
            const float v = acc * 0.25f;               // / hc_count
            mix_out[n] = (T)(v / (1.0f + metal::exp(-v)));   // silu
        } else {
            const float v = acc * 0.25f;
            inject_out[n - N_DOWN] = (T)(2.0f / (1.0f + metal::exp(-v)));
        }
    }
}
typedef decltype(kernel_qwen_mtplx_hyper_v3_r1<bfloat>) kernel_qwen_mtplx_hyper_v3_r1_bf16_type;
template [[host_name("kernel_qwen_mtplx_hyper_v3_r1_bf16")]]
kernel kernel_qwen_mtplx_hyper_v3_r1_bf16_type kernel_qwen_mtplx_hyper_v3_r1<bfloat>;

// Source: mtplx/kernels/hyper_connection_v3.py::_SRC_R2
// File SHA256: bc4cc7c59403845c31881f7526dade0f80a7704303c6c5273a794d605cf4e82a
// Body SHA256: f5ac42120a2c20ebf825bdc290b3b9c8afbfa6ad4a5d18454d9eff542174a74f
template <typename T>
kernel void kernel_qwen_mtplx_hyper_v3_r2(
    device const T* x [[buffer(0)]],
    device const T* wn [[buffer(1)]],
    device const T* mixv [[buffer(2)]],
    device const float* rms_in [[buffer(3)]],
    device const uint32_t* qw [[buffer(4)]],
    device const T* qs [[buffer(5)]],
    device const T* qb [[buffer(6)]],
    device T* y [[buffer(7)]],
    uint3 thread_position_in_threadgroup [[thread_position_in_threadgroup]],
    uint3 threadgroup_position_in_grid [[threadgroup_position_in_grid]]) {
    constexpr int KUP = 320;
    constexpr int GROUP = 2560;
    constexpr int QGS = 64;
    constexpr int NG = KUP / QGS;        // 5 quant groups per up row

    const uint tid = thread_position_in_threadgroup.x;
    const uint sg = tid / 32;
    const uint lane = tid % 32;
    const uint d = threadgroup_position_in_grid.x * 32 + sg;
    if (d >= (uint)GROUP) return;

    float acc_mix = 0.0f;
    for (int h = 0; h < 4; ++h) {
        const uint row = (uint)h * GROUP + d;
        const device uint32_t* wrow = qw + (size_t)row * (KUP / 4);
        const device T* srow = qs + (size_t)row * NG;
        const device T* brow = qb + (size_t)row * NG;
        float dot = 0.0f;
        for (int g = 0; g < NG; ++g) {
            const float s = (float)srow[g];
            const float b = (float)brow[g];
            const device uint32_t* wg = wrow + g * (QGS / 4);
            const int base = g * QGS;
            float qacc = 0.0f;
            float msum = 0.0f;
            for (int wi = (int)lane; wi < QGS / 4; wi += 32) {
                const uint32_t word = wg[wi];
                const int i0 = base + wi * 4;
                for (int by = 0; by < 4; ++by) {
                    const float mv = (float)mixv[i0 + by];
                    qacc += (float)((word >> (8 * by)) & 0xFF) * mv;
                    msum += mv;
                }
            }
            dot += s * qacc + b * msum;
        }
        dot = simd_sum(dot);
        const float m = 1.0f / (1.0f + metal::exp(-dot));   // sigmoid
        const int i = h * GROUP + (int)d;
        const float normed = (float)x[i] * (float)wn[i] * rms_in[h];
        acc_mix += m * normed;
    }
    if (lane == 0) {
        y[d] = (T)(acc_mix * 0.25f);
    }
}
typedef decltype(kernel_qwen_mtplx_hyper_v3_r2<bfloat>) kernel_qwen_mtplx_hyper_v3_r2_bf16_type;
template [[host_name("kernel_qwen_mtplx_hyper_v3_r2_bf16")]]
kernel kernel_qwen_mtplx_hyper_v3_r2_bf16_type kernel_qwen_mtplx_hyper_v3_r2<bfloat>;

// Source: mtplx/kernels/gdn_conv_norm.py::_SRC
// File SHA256: b414389736c2f7c166acc7846d4d0b4e021524d75886b7062214306adeab85eb
// Body SHA256: f11cee0261d66aa1ddac3370d747c22a5eb943f1e0cd306764c0a5d280179e09
template <typename T>
kernel void kernel_qwen_mtplx_gdn_conv_norm(
    device const T* xnew [[buffer(0)]],
    device const T* state [[buffer(1)]],
    device const T* cw [[buffer(2)]],
    device T* q_out [[buffer(3)]],
    device T* k_out [[buffer(4)]],
    device T* v_out [[buffer(5)]],
    device T* state_out [[buffer(6)]],
    uint3 thread_position_in_threadgroup [[thread_position_in_threadgroup]],
    uint3 threadgroup_position_in_grid [[threadgroup_position_in_grid]]) {
    constexpr int C = 10240;                  // conv channels
    constexpr int QK = 2048;                  // q width == k width
    constexpr int HD = 128;                   // head dim
    constexpr float INV_SCALE = 0.08838834764831845f;   // 128^-0.5

    const uint tid = thread_position_in_threadgroup.x;
    const uint sg = tid / 32;
    const uint lane = tid % 32;
    const uint c = threadgroup_position_in_grid.x * 1024 + tid;
    if (c >= (uint)C) return;

    threadgroup float tg_vals[1024];
    threadgroup float tg_partial[32];

    // depthwise conv over [state0, state1, state2, new] + silu
    const float acc =
        (float)cw[c * 4 + 0] * (float)state[0 * C + c] +
        (float)cw[c * 4 + 1] * (float)state[1 * C + c] +
        (float)cw[c * 4 + 2] * (float)state[2 * C + c] +
        (float)cw[c * 4 + 3] * (float)xnew[c];
    const float sv = acc / (1.0f + metal::exp(-acc));

    // rolled conv state
    state_out[0 * C + c] = state[1 * C + c];
    state_out[1 * C + c] = state[2 * C + c];
    state_out[2 * C + c] = xnew[c];

    if (c >= (uint)(2 * QK)) {
        v_out[c - 2 * QK] = (T)sv;
        return;
    }

    // q/k: per-head l2norm. This TG holds 8 aligned heads of 128 channels;
    // 4 consecutive simdgroups own one head.
    tg_vals[tid] = sv;
    float part = simd_sum(sv * sv);
    if (lane == 0) tg_partial[sg] = part;
    threadgroup_barrier(mem_flags::mem_threadgroup);
    const uint head_sg0 = (sg / 4) * 4;
    const float ssum = tg_partial[head_sg0] + tg_partial[head_sg0 + 1]
                     + tg_partial[head_sg0 + 2] + tg_partial[head_sg0 + 3];
    const float inv = metal::rsqrt(ssum + 1e-6f);
    const float normed = tg_vals[tid] * inv;
    if (c < (uint)QK) {
        q_out[c] = (T)(normed * INV_SCALE);
    } else {
        k_out[c - QK] = (T)normed;
    }
}
typedef decltype(kernel_qwen_mtplx_gdn_conv_norm<bfloat>) kernel_qwen_mtplx_gdn_conv_norm_bf16_type;
template [[host_name("kernel_qwen_mtplx_gdn_conv_norm_bf16")]]
kernel kernel_qwen_mtplx_gdn_conv_norm_bf16_type kernel_qwen_mtplx_gdn_conv_norm<bfloat>;

// Source: mtplx/kernels/gdn_conv_norm.py::_SRC_ROWS
// File SHA256: b414389736c2f7c166acc7846d4d0b4e021524d75886b7062214306adeab85eb
// Body SHA256: 291df3d66fccd0feff041313a61e3a8f942068e8ef3557da05a813fe587304f0
template <typename T, int S>
kernel void kernel_qwen_mtplx_gdn_conv_norm_rows(
    device const T* xnew [[buffer(0)]],
    device const T* state [[buffer(1)]],
    device const T* cw [[buffer(2)]],
    device T* q_out [[buffer(3)]],
    device T* k_out [[buffer(4)]],
    device T* v_out [[buffer(5)]],
    device T* state_out [[buffer(6)]],
    uint3 thread_position_in_threadgroup [[thread_position_in_threadgroup]],
    uint3 threadgroup_position_in_grid [[threadgroup_position_in_grid]]) {
    constexpr int C = 10240;                  // conv channels
    constexpr int QK = 2048;                  // q width == k width
    constexpr float INV_SCALE = 0.08838834764831845f;   // 128^-0.5

    const uint tid = thread_position_in_threadgroup.x;
    const uint sg = tid / 32;
    const uint lane = tid % 32;
    const uint c = threadgroup_position_in_grid.x * 1024 + tid;
    if (c >= (uint)C) return;

    threadgroup float tg_vals[1024];
    threadgroup float tg_partial[32];

    const float w0 = (float)cw[c * 4 + 0];
    const float w1 = (float)cw[c * 4 + 1];
    const float w2 = (float)cw[c * 4 + 2];
    const float w3 = (float)cw[c * 4 + 3];

    // stream(t): t<3 -> conv state row t, else xnew row t-3
    #define STREAM(t) ((t) < 3 ? (float)state[(t) * C + c] : (float)xnew[((t) - 3) * C + c])

    const bool is_v = (c >= (uint)(2 * QK));

    for (int s = 0; s < S; ++s) {
        const float acc = w0 * STREAM(s) + w1 * STREAM(s + 1)
                        + w2 * STREAM(s + 2) + w3 * STREAM(s + 3);
        const float sv = acc / (1.0f + metal::exp(-acc));
        if (is_v) {
            v_out[s * (C - 2 * QK) + (c - 2 * QK)] = (T)sv;
            continue;
        }
        tg_vals[tid] = sv;
        float part = simd_sum(sv * sv);
        if (lane == 0) tg_partial[sg] = part;
        threadgroup_barrier(mem_flags::mem_threadgroup);
        const uint head_sg0 = (sg / 4) * 4;
        const float ssum = tg_partial[head_sg0] + tg_partial[head_sg0 + 1]
                         + tg_partial[head_sg0 + 2] + tg_partial[head_sg0 + 3];
        const float inv = metal::rsqrt(ssum + 1e-6f);
        const float normed = tg_vals[tid] * inv;
        if (c < (uint)QK) {
            q_out[s * QK + c] = (T)(normed * INV_SCALE);
        } else {
            k_out[s * QK + (c - QK)] = (T)normed;
        }
        threadgroup_barrier(mem_flags::mem_threadgroup);
    }

    // rolled conv state: last 3 rows of the stream
    state_out[0 * C + c] = (T)STREAM(S + 0);
    state_out[1 * C + c] = (T)STREAM(S + 1);
    state_out[2 * C + c] = (T)STREAM(S + 2);
    #undef STREAM
}
typedef decltype(kernel_qwen_mtplx_gdn_conv_norm_rows<bfloat, 2>) kernel_qwen_mtplx_gdn_conv_norm_rows_s2_bf16_type;
template [[host_name("kernel_qwen_mtplx_gdn_conv_norm_rows_s2_bf16")]]
kernel kernel_qwen_mtplx_gdn_conv_norm_rows_s2_bf16_type kernel_qwen_mtplx_gdn_conv_norm_rows<bfloat, 2>;
typedef decltype(kernel_qwen_mtplx_gdn_conv_norm_rows<bfloat, 3>) kernel_qwen_mtplx_gdn_conv_norm_rows_s3_bf16_type;
template [[host_name("kernel_qwen_mtplx_gdn_conv_norm_rows_s3_bf16")]]
kernel kernel_qwen_mtplx_gdn_conv_norm_rows_s3_bf16_type kernel_qwen_mtplx_gdn_conv_norm_rows<bfloat, 3>;
typedef decltype(kernel_qwen_mtplx_gdn_conv_norm_rows<bfloat, 4>) kernel_qwen_mtplx_gdn_conv_norm_rows_s4_bf16_type;
template [[host_name("kernel_qwen_mtplx_gdn_conv_norm_rows_s4_bf16")]]
kernel kernel_qwen_mtplx_gdn_conv_norm_rows_s4_bf16_type kernel_qwen_mtplx_gdn_conv_norm_rows<bfloat, 4>;
typedef decltype(kernel_qwen_mtplx_gdn_conv_norm_rows<bfloat, 5>) kernel_qwen_mtplx_gdn_conv_norm_rows_s5_bf16_type;
template [[host_name("kernel_qwen_mtplx_gdn_conv_norm_rows_s5_bf16")]]
kernel kernel_qwen_mtplx_gdn_conv_norm_rows_s5_bf16_type kernel_qwen_mtplx_gdn_conv_norm_rows<bfloat, 5>;
typedef decltype(kernel_qwen_mtplx_gdn_conv_norm_rows<bfloat, 6>) kernel_qwen_mtplx_gdn_conv_norm_rows_s6_bf16_type;
template [[host_name("kernel_qwen_mtplx_gdn_conv_norm_rows_s6_bf16")]]
kernel kernel_qwen_mtplx_gdn_conv_norm_rows_s6_bf16_type kernel_qwen_mtplx_gdn_conv_norm_rows<bfloat, 6>;

// Source: mtplx/kernels/gdn_step_fused.py::_SRC
// File SHA256: d194412dbea6c3e5758a3a345d5a15fc127ea525f71e9b23f6fe44e19a1c0ad4
// Body SHA256: 433a6eef305d6be98f2816d6947ea66e93b0193c2f42717adf36ba1e5f5b0771
template <typename T, typename StT>
kernel void kernel_qwen_mtplx_gdn_step_fused(
    device const T* xnew [[buffer(0)]],
    device const T* z_row [[buffer(1)]],
    device const T* a_row [[buffer(2)]],
    device const T* b_row [[buffer(3)]],
    device const T* state [[buffer(4)]],
    device const T* cw [[buffer(5)]],
    device const T* A_log [[buffer(6)]],
    device const T* dt_bias [[buffer(7)]],
    device const StT* dstate [[buffer(8)]],
    device const T* norm_w [[buffer(9)]],
    device T* y [[buffer(10)]],
    device T* state_out [[buffer(11)]],
    device StT* dstate_out [[buffer(12)]],
    uint3 thread_position_in_threadgroup [[thread_position_in_threadgroup]],
    uint3 threadgroup_position_in_grid [[threadgroup_position_in_grid]]) {
    constexpr int C = 10240;                 // conv channels
    constexpr int QK = 2048;                 // q width == k width
    constexpr int HD = 128;                  // head dim (Dk == Dv)
    constexpr int HV_PER_HK = 3;             // 48 v heads / 16 k heads
    constexpr float INV_SCALE = 0.08838834764831845f;  // 128^-0.5

    const uint tid = thread_position_in_threadgroup.x;
    const uint sg = tid / 32;
    const uint lane = tid % 32;
    const uint hv = threadgroup_position_in_grid.z;
    const uint hk = hv / HV_PER_HK;

    threadgroup float qn[HD];
    threadgroup float kn[HD];
    threadgroup float vraw[HD];
    threadgroup float outv[HD];
    threadgroup float part[8];

    // ---- phase 1: conv + silu over this head's 384 channels ----
    for (uint i = tid; i < 384; i += 256) {
      uint g;
      if (i < HD) {
        g = hk * HD + i;                       // q section
      } else if (i < 2 * HD) {
        g = QK + hk * HD + (i - HD);           // k section
      } else {
        g = 2 * QK + hv * HD + (i - 2 * HD);   // v section
      }
      const float acc =
          (float)cw[g * 4 + 0] * (float)state[0 * C + g] +
          (float)cw[g * 4 + 1] * (float)state[1 * C + g] +
          (float)cw[g * 4 + 2] * (float)state[2 * C + g] +
          (float)cw[g * 4 + 3] * (float)xnew[g];
      const float sv = acc / (1.0f + metal::exp(-acc));

      // rolled conv state: v channels owned uniquely; q/k written by the
      // hv%3==0 sibling so each channel has exactly one writer.
      if (i >= 2 * HD || (hv % HV_PER_HK) == 0) {
        state_out[0 * C + g] = state[1 * C + g];
        state_out[1 * C + g] = state[2 * C + g];
        state_out[2 * C + g] = xnew[g];
      }

      if (i < HD) {
        qn[i] = sv;
      } else if (i < 2 * HD) {
        kn[i - HD] = sv;
      } else {
        // stock path rounds v to T at the conv-kernel boundary
        vraw[i - 2 * HD] = (float)(T)sv;
      }
    }
    threadgroup_barrier(mem_flags::mem_threadgroup);

    // ---- phase 1b: per-head l2norm (fp32, eps on sum of squares) ----
    {
      float p = 0.0f;
      if (tid < HD) {
        p = qn[tid] * qn[tid];
      } else if (tid < 2 * HD) {
        p = kn[tid - HD] * kn[tid - HD];
      }
      p = simd_sum(p);
      if (lane == 0) part[sg] = p;
      threadgroup_barrier(mem_flags::mem_threadgroup);
      const float qsum = part[0] + part[1] + part[2] + part[3];
      const float ksum = part[4] + part[5] + part[6] + part[7];
      const float qinv = metal::rsqrt(qsum + 1e-6f) * INV_SCALE;
      const float kinv = metal::rsqrt(ksum + 1e-6f);
      threadgroup_barrier(mem_flags::mem_threadgroup);
      if (tid < HD) {
        qn[tid] = (float)(T)(qn[tid] * qinv);
      } else if (tid < 2 * HD) {
        kn[tid - HD] = (float)(T)(kn[tid - HD] * kinv);
      }
      threadgroup_barrier(mem_flags::mem_threadgroup);
    }

    // ---- g / beta (redundant per thread; matches compute_g + sigmoid) ----
    const float a_v = (float)a_row[hv] + (float)dt_bias[hv];
    // softplus = logaddexp(x, 0) = max(x,0) + log1p(exp(-|x|))
    const float sp = metal::max(a_v, 0.0f)
                   + metal::log(1.0f + metal::exp(-metal::abs(a_v)));
    const float g_dec = metal::exp(-metal::exp((float)A_log[hv]) * sp);
    const float b_v = (float)b_row[hv];
    const float beta = 1.0f / (1.0f + metal::exp(-b_v));

    // ---- phase 2: delta recurrence, one simdgroup per row, stride 8 ----
    for (uint dv = sg; dv < HD; dv += 8) {
      const uint srow = (hv * HD + dv) * HD;
      float s0 = (float)dstate[srow + 4 * lane + 0] * g_dec;
      float s1 = (float)dstate[srow + 4 * lane + 1] * g_dec;
      float s2 = (float)dstate[srow + 4 * lane + 2] * g_dec;
      float s3 = (float)dstate[srow + 4 * lane + 3] * g_dec;
      const float k0 = kn[4 * lane + 0];
      const float k1 = kn[4 * lane + 1];
      const float k2 = kn[4 * lane + 2];
      const float k3 = kn[4 * lane + 3];
      float kv_mem = s0 * k0 + s1 * k1 + s2 * k2 + s3 * k3;
      kv_mem = simd_sum(kv_mem);
      const float delta = (vraw[dv] - kv_mem) * beta;
      s0 += k0 * delta;
      s1 += k1 * delta;
      s2 += k2 * delta;
      s3 += k3 * delta;
      float out = s0 * qn[4 * lane + 0] + s1 * qn[4 * lane + 1]
                + s2 * qn[4 * lane + 2] + s3 * qn[4 * lane + 3];
      out = simd_sum(out);
      dstate_out[srow + 4 * lane + 0] = (StT)s0;
      dstate_out[srow + 4 * lane + 1] = (StT)s1;
      dstate_out[srow + 4 * lane + 2] = (StT)s2;
      dstate_out[srow + 4 * lane + 3] = (StT)s3;
      if (lane == 0) {
        // the stock kernel emits InT here; keep that rounding
        outv[dv] = (float)(T)out;
      }
    }
    threadgroup_barrier(mem_flags::mem_threadgroup);

    // ---- phase 3: SigmoidRMSNormGated epilogue ----
    {
      float e = 0.0f;
      if (tid < HD) e = outv[tid] * outv[tid];
      e = simd_sum(e);
      if (lane == 0 && sg < 4) part[sg] = e;
      threadgroup_barrier(mem_flags::mem_threadgroup);
      const float mean = (part[0] + part[1] + part[2] + part[3]) / (float)HD;
      const float rinv = metal::rsqrt(mean + 1e-6f);
      if (tid < HD) {
        const float normed = (float)(T)(outv[tid] * rinv * (float)norm_w[tid]);
        const float zf = (float)z_row[hv * HD + tid];
        const float gate = 1.0f / (1.0f + metal::exp(-zf));
        y[hv * HD + tid] = (T)(gate * normed);
      }
    }
}
typedef decltype(kernel_qwen_mtplx_gdn_step_fused<bfloat, float>) kernel_qwen_mtplx_gdn_step_fused_f32_state_bf16_type;
template [[host_name("kernel_qwen_mtplx_gdn_step_fused_f32_state_bf16")]]
kernel kernel_qwen_mtplx_gdn_step_fused_f32_state_bf16_type kernel_qwen_mtplx_gdn_step_fused<bfloat, float>;

// Source: mtplx/kernels/gdn_out_fused.py::_SRC
// File SHA256: f8698820407b5b2c469b8393c20cb2e7841c2fd8e3eef004234d1cd6b38bacdc
// Body SHA256: 21043fd910132a1cd9fd7d99185c85b4b69cc8b24f3747e0ff108d6af2185043
template <typename T, int GS_C>
kernel void kernel_qwen_mtplx_gdn_out_fused(
    device const T* x [[buffer(0)]],
    device const bfloat* z [[buffer(1)]],
    device const bfloat* wn [[buffer(2)]],
    device const uint32_t* qw [[buffer(3)]],
    device const bfloat* qs [[buffer(4)]],
    device const bfloat* qb [[buffer(5)]],
    device T* y [[buffer(6)]],
    uint3 thread_position_in_threadgroup [[thread_position_in_threadgroup]],
    uint3 threadgroup_position_in_grid [[threadgroup_position_in_grid]]) {
    constexpr int NH = 48;                    // value heads
    constexpr int HD = 128;                   // head dim
    constexpr int K = NH * HD;                // 6144
    constexpr int GS = GS_C;                  // quant group size
    constexpr int NGROUPS = K / GS;
    constexpr int WPG = GS / 8;               // u32 words per group
    constexpr int DMODEL = 2560;

    const uint tid = thread_position_in_threadgroup.x;
    const uint sg = tid / 32;
    const uint lane = tid % 32;

    threadgroup float tg_rms[NH];
    threadgroup T tg_v[K];
    threadgroup float tg_part[32];

    // Phase 1a: per-head sum of squares. Head h owned by simdgroup h%32
    // (simdgroups 0..15 own two heads).
    for (uint h = sg; h < (uint)NH; h += 32) {
        float part = 0.0f;
        for (uint d = lane; d < (uint)HD; d += 32) {
            const float xv = (float)x[h * HD + d];
            part += xv * xv;
        }
        part = simd_sum(part);
        if (lane == 0) {
            tg_rms[h] = metal::rsqrt(part / (float)HD + 1e-6f);
        }
    }
    threadgroup_barrier(mem_flags::mem_threadgroup);

    // Phase 1b: gated-normed values, bf16-cast through T.
    for (uint i = tid; i < (uint)K; i += 1024) {
        const uint h = i / HD;
        const uint d = i % HD;
        const float normed = (float)x[i] * tg_rms[h] * (float)wn[d];
        const float g = 1.0f / (1.0f + metal::exp(-(float)z[i]));
        tg_v[i] = (T)(g * normed);
    }
    threadgroup_barrier(mem_flags::mem_threadgroup);

    // Phase 2: one simdgroup per output row.
    const uint row = threadgroup_position_in_grid.x * 32 + sg;
    if (row >= (uint)DMODEL) return;
    const device uint32_t* wrow = qw + (size_t)row * (K / 8);
    // scale/bias buffers keep their own dtype (bf16 in shipped packs) —
    // never retype them through T (T follows x, which is fp32 when the
    // delta kernel hands over float state rows)
    const device auto* srow = qs + (size_t)row * NGROUPS;
    const device auto* brow = qb + (size_t)row * NGROUPS;
    float acc = 0.0f;
    for (int g = lane; g < NGROUPS; g += 32) {
        const float s = (float)srow[g];
        const float b = (float)brow[g];
        const device uint32_t* wg = wrow + g * WPG;
        const int base = g * GS;
        float qacc = 0.0f;
        float vsum = 0.0f;
        for (int wi = 0; wi < WPG; ++wi) {
            const uint32_t word = wg[wi];
            const int i0 = base + wi * 8;
            for (int nib = 0; nib < 8; ++nib) {
                const float vv = (float)tg_v[i0 + nib];
                qacc += (float)((word >> (4 * nib)) & 0xF) * vv;
                vsum += vv;
            }
        }
        acc += s * qacc + b * vsum;
    }
    acc = simd_sum(acc);
    if (lane == 0) {
        y[row] = (T)acc;
    }
}
typedef decltype(kernel_qwen_mtplx_gdn_out_fused<bfloat, 32>) kernel_qwen_mtplx_gdn_out_fused_gs32_bf16_type;
template [[host_name("kernel_qwen_mtplx_gdn_out_fused_gs32_bf16")]]
kernel kernel_qwen_mtplx_gdn_out_fused_gs32_bf16_type kernel_qwen_mtplx_gdn_out_fused<bfloat, 32>;
typedef decltype(kernel_qwen_mtplx_gdn_out_fused<float, 32>) kernel_qwen_mtplx_gdn_out_fused_gs32_f32_type;
template [[host_name("kernel_qwen_mtplx_gdn_out_fused_gs32_f32")]]
kernel kernel_qwen_mtplx_gdn_out_fused_gs32_f32_type kernel_qwen_mtplx_gdn_out_fused<float, 32>;
typedef decltype(kernel_qwen_mtplx_gdn_out_fused<bfloat, 64>) kernel_qwen_mtplx_gdn_out_fused_gs64_bf16_type;
template [[host_name("kernel_qwen_mtplx_gdn_out_fused_gs64_bf16")]]
kernel kernel_qwen_mtplx_gdn_out_fused_gs64_bf16_type kernel_qwen_mtplx_gdn_out_fused<bfloat, 64>;
typedef decltype(kernel_qwen_mtplx_gdn_out_fused<float, 64>) kernel_qwen_mtplx_gdn_out_fused_gs64_f32_type;
template [[host_name("kernel_qwen_mtplx_gdn_out_fused_gs64_f32")]]
kernel kernel_qwen_mtplx_gdn_out_fused_gs64_f32_type kernel_qwen_mtplx_gdn_out_fused<float, 64>;

// Source: mtplx/kernels/moe_glu_decode.py::_SRC_A
// File SHA256: ed8b307cd8acd2b111237a57f6adc711d455fac1feb7800df72821c4061b8912
// Body SHA256: 85d5d607e7d647a17124013d98b9f4b3882b8f024b5131bab2d20c7379a2474b
template <typename T, int n_inter_c, int topk_c, int GS_GU>
kernel void kernel_qwen_mtplx_moe_glu_h(
    device const T* x [[buffer(0)]],
    device const uint32_t* gw [[buffer(1)]],
    device const T* gs [[buffer(2)]],
    device const T* gb [[buffer(3)]],
    device const uint32_t* experts [[buffer(4)]],
    device T* h [[buffer(5)]],
    uint3 thread_position_in_threadgroup [[thread_position_in_threadgroup]],
    uint3 threadgroup_position_in_grid [[threadgroup_position_in_grid]]) {
    constexpr int K = 2560;
    constexpr int GS = GS_GU;                 // quant group size (32/64 forges)
    constexpr int NGROUPS = K / GS;
    constexpr int WPG = GS / 8;               // uint32 words per group (8 nibbles/word)

    const uint tid = thread_position_in_threadgroup.x;
    const uint sg = tid / 32;
    const uint lane = tid % 32;
    const uint gsid = threadgroup_position_in_grid.x * 32 + sg;  // global simd id
    const int n_inter = n_inter_c;
    const int topk = topk_c;
    if (gsid >= (uint)(topk * n_inter)) return;
    const int e_slot = gsid / n_inter;
    const int j = gsid % n_inter;
    const uint e = experts[e_slot];

    float acc_g = 0.0f;
    float acc_u = 0.0f;
    for (int hh = 0; hh < 2; ++hh) {
        const int row = hh * n_inter + j;
        const device uint32_t* wrow = gw + ((size_t)e * 2 * n_inter + row) * (K / 8);
        const device T* srow = gs + ((size_t)e * 2 * n_inter + row) * NGROUPS;
        const device T* brow = gb + ((size_t)e * 2 * n_inter + row) * NGROUPS;
        float acc = 0.0f;
        for (int g = lane; g < NGROUPS; g += 32) {
            const float s = (float)srow[g];
            const float b = (float)brow[g];
            const device uint32_t* wg = wrow + g * WPG;
            const device T* xg = x + g * GS;
            float qacc = 0.0f;
            float xsum = 0.0f;
            for (int wi = 0; wi < WPG; ++wi) {
                uint32_t word = wg[wi];
                const device T* xv = xg + wi * 8;
                for (int nib = 0; nib < 8; ++nib) {
                    const float xf = (float)xv[nib];
                    qacc += (float)((word >> (4 * nib)) & 0xF) * xf;
                    xsum += xf;
                }
            }
            acc += s * qacc + b * xsum;
        }
        acc = simd_sum(acc);
        if (hh == 0) acc_g = acc; else acc_u = acc;
    }
    if (lane == 0) {
        const float gv = (float)((T)acc_g);
        const float uv = (float)((T)acc_u);
        const float sw = gv / (1.0f + metal::exp(-gv));   // silu
        h[(size_t)e_slot * n_inter + j] = (T)((float)((T)sw) * uv);
    }
}
typedef decltype(kernel_qwen_mtplx_moe_glu_h<bfloat, 640, 10, 32>) kernel_qwen_mtplx_moe_glu_h_g32_bf16_type;
template [[host_name("kernel_qwen_mtplx_moe_glu_h_g32_bf16")]]
kernel kernel_qwen_mtplx_moe_glu_h_g32_bf16_type kernel_qwen_mtplx_moe_glu_h<bfloat, 640, 10, 32>;
typedef decltype(kernel_qwen_mtplx_moe_glu_h<bfloat, 640, 10, 64>) kernel_qwen_mtplx_moe_glu_h_g64_bf16_type;
template [[host_name("kernel_qwen_mtplx_moe_glu_h_g64_bf16")]]
kernel kernel_qwen_mtplx_moe_glu_h_g64_bf16_type kernel_qwen_mtplx_moe_glu_h<bfloat, 640, 10, 64>;

// Source: mtplx/kernels/moe_glu_decode.py::_SRC_B
// File SHA256: ed8b307cd8acd2b111237a57f6adc711d455fac1feb7800df72821c4061b8912
// Body SHA256: 698055e6a7703cffe318dbcab2e04e3aaff4ffb6a056be60957e1ca3815e746c
template <typename T, int dmodel_c, int n_inter_c, int topk_c, int GS_DN>
kernel void kernel_qwen_mtplx_moe_down_y(
    device const T* h [[buffer(0)]],
    device const uint32_t* dw [[buffer(1)]],
    device const T* ds [[buffer(2)]],
    device const T* db [[buffer(3)]],
    device const uint32_t* experts [[buffer(4)]],
    device const float* rw [[buffer(5)]],
    device T* y [[buffer(6)]],
    uint3 thread_position_in_threadgroup [[thread_position_in_threadgroup]],
    uint3 threadgroup_position_in_grid [[threadgroup_position_in_grid]]) {
    constexpr int GS = GS_DN;                 // quant group size (32/64 forges)

    const uint tid = thread_position_in_threadgroup.x;
    const uint sg = tid / 32;
    const uint lane = tid % 32;
    const uint d = threadgroup_position_in_grid.x * 32 + sg;
    const int dmodel = dmodel_c;
    const int n_inter = n_inter_c;
    const int topk = topk_c;
    if (d >= (uint)dmodel) return;

    const int ngroups = n_inter / GS;
    float acc = 0.0f;
    for (int e_slot = 0; e_slot < topk; ++e_slot) {
        const uint e = experts[e_slot];
        const device uint32_t* wrow = dw + ((size_t)e * dmodel + d) * (n_inter / 8);
        const device T* srow = ds + ((size_t)e * dmodel + d) * ngroups;
        const device T* brow = db + ((size_t)e * dmodel + d) * ngroups;
        const device T* he = h + (size_t)e_slot * n_inter;
        float dot = 0.0f;
        for (int k = (int)lane; k < n_inter; k += 32) {
            const uint32_t word = wrow[k / 8];
            const float q = (float)((word >> (4 * (k % 8))) & 0xF);
            const int g = k / GS;
            dot = metal::fma((float)srow[g] * q + (float)brow[g], (float)he[k], dot);
        }
        dot = simd_sum(dot);
        acc += (float)rw[e_slot] * (float)((T)dot);
    }
    if (lane == 0) {
        y[d] = (T)acc;
    }
}
typedef decltype(kernel_qwen_mtplx_moe_down_y<bfloat, 2560, 640, 10, 32>) kernel_qwen_mtplx_moe_down_y_g32_bf16_type;
template [[host_name("kernel_qwen_mtplx_moe_down_y_g32_bf16")]]
kernel kernel_qwen_mtplx_moe_down_y_g32_bf16_type kernel_qwen_mtplx_moe_down_y<bfloat, 2560, 640, 10, 32>;
typedef decltype(kernel_qwen_mtplx_moe_down_y<bfloat, 2560, 640, 10, 64>) kernel_qwen_mtplx_moe_down_y_g64_bf16_type;
template [[host_name("kernel_qwen_mtplx_moe_down_y_g64_bf16")]]
kernel kernel_qwen_mtplx_moe_down_y_g64_bf16_type kernel_qwen_mtplx_moe_down_y<bfloat, 2560, 640, 10, 64>;

// Source: mtplx/kernels/moe_glu_decode.py::_SRC_A_M
// File SHA256: ed8b307cd8acd2b111237a57f6adc711d455fac1feb7800df72821c4061b8912
// Body SHA256: 01cd8217814d3417c6daeff9103f9585422b5417467abf4eafaec2f5f4ee34c2
template <typename T, int n_inter_c, int topk_c, int m_rows_c, int GS_GU>
kernel void kernel_qwen_mtplx_moe_glu_h_m(
    device const T* x [[buffer(0)]],
    device const uint32_t* gw [[buffer(1)]],
    device const T* gs [[buffer(2)]],
    device const T* gb [[buffer(3)]],
    device const uint32_t* experts [[buffer(4)]],
    device T* h [[buffer(5)]],
    uint3 thread_position_in_threadgroup [[thread_position_in_threadgroup]],
    uint3 threadgroup_position_in_grid [[threadgroup_position_in_grid]]) {
    constexpr int K = 2560;
    constexpr int GS = GS_GU;                 // quant group size (32/64 forges)
    constexpr int NGROUPS = K / GS;
    constexpr int WPG = GS / 8;               // uint32 words per group (8 nibbles/word)

    const uint tid = thread_position_in_threadgroup.x;
    const uint sg = tid / 32;
    const uint lane = tid % 32;
    const uint gsid = threadgroup_position_in_grid.x * 32 + sg;  // global simd id
    const int n_inter = n_inter_c;
    const int topk = topk_c;
    const int m_rows = m_rows_c;
    if (gsid >= (uint)(m_rows * topk * n_inter)) return;
    const int m = gsid / (topk * n_inter);
    const int rem = gsid % (topk * n_inter);
    const int e_slot = rem / n_inter;
    const int j = rem % n_inter;
    const uint e = experts[m * topk + e_slot];
    const device T* xm = x + (size_t)m * K;

    float acc_g = 0.0f;
    float acc_u = 0.0f;
    for (int hh = 0; hh < 2; ++hh) {
        const int row = hh * n_inter + j;
        const device uint32_t* wrow = gw + ((size_t)e * 2 * n_inter + row) * (K / 8);
        const device T* srow = gs + ((size_t)e * 2 * n_inter + row) * NGROUPS;
        const device T* brow = gb + ((size_t)e * 2 * n_inter + row) * NGROUPS;
        float acc = 0.0f;
        for (int g = lane; g < NGROUPS; g += 32) {
            const float s = (float)srow[g];
            const float b = (float)brow[g];
            const device uint32_t* wg = wrow + g * WPG;
            const device T* xg = xm + g * GS;
            float qacc = 0.0f;
            float xsum = 0.0f;
            for (int wi = 0; wi < WPG; ++wi) {
                uint32_t word = wg[wi];
                const device T* xv = xg + wi * 8;
                for (int nib = 0; nib < 8; ++nib) {
                    const float xf = (float)xv[nib];
                    qacc += (float)((word >> (4 * nib)) & 0xF) * xf;
                    xsum += xf;
                }
            }
            acc += s * qacc + b * xsum;
        }
        acc = simd_sum(acc);
        if (hh == 0) acc_g = acc; else acc_u = acc;
    }
    if (lane == 0) {
        const float gv = (float)((T)acc_g);
        const float uv = (float)((T)acc_u);
        const float sw = gv / (1.0f + metal::exp(-gv));   // silu
        h[((size_t)m * topk + e_slot) * n_inter + j] = (T)((float)((T)sw) * uv);
    }
}
typedef decltype(kernel_qwen_mtplx_moe_glu_h_m<bfloat, 640, 10, 2, 32>) kernel_qwen_mtplx_moe_glu_h_m_m2_g32_bf16_type;
template [[host_name("kernel_qwen_mtplx_moe_glu_h_m_m2_g32_bf16")]]
kernel kernel_qwen_mtplx_moe_glu_h_m_m2_g32_bf16_type kernel_qwen_mtplx_moe_glu_h_m<bfloat, 640, 10, 2, 32>;
typedef decltype(kernel_qwen_mtplx_moe_glu_h_m<bfloat, 640, 10, 2, 64>) kernel_qwen_mtplx_moe_glu_h_m_m2_g64_bf16_type;
template [[host_name("kernel_qwen_mtplx_moe_glu_h_m_m2_g64_bf16")]]
kernel kernel_qwen_mtplx_moe_glu_h_m_m2_g64_bf16_type kernel_qwen_mtplx_moe_glu_h_m<bfloat, 640, 10, 2, 64>;
typedef decltype(kernel_qwen_mtplx_moe_glu_h_m<bfloat, 640, 10, 3, 32>) kernel_qwen_mtplx_moe_glu_h_m_m3_g32_bf16_type;
template [[host_name("kernel_qwen_mtplx_moe_glu_h_m_m3_g32_bf16")]]
kernel kernel_qwen_mtplx_moe_glu_h_m_m3_g32_bf16_type kernel_qwen_mtplx_moe_glu_h_m<bfloat, 640, 10, 3, 32>;
typedef decltype(kernel_qwen_mtplx_moe_glu_h_m<bfloat, 640, 10, 3, 64>) kernel_qwen_mtplx_moe_glu_h_m_m3_g64_bf16_type;
template [[host_name("kernel_qwen_mtplx_moe_glu_h_m_m3_g64_bf16")]]
kernel kernel_qwen_mtplx_moe_glu_h_m_m3_g64_bf16_type kernel_qwen_mtplx_moe_glu_h_m<bfloat, 640, 10, 3, 64>;
typedef decltype(kernel_qwen_mtplx_moe_glu_h_m<bfloat, 640, 10, 4, 32>) kernel_qwen_mtplx_moe_glu_h_m_m4_g32_bf16_type;
template [[host_name("kernel_qwen_mtplx_moe_glu_h_m_m4_g32_bf16")]]
kernel kernel_qwen_mtplx_moe_glu_h_m_m4_g32_bf16_type kernel_qwen_mtplx_moe_glu_h_m<bfloat, 640, 10, 4, 32>;
typedef decltype(kernel_qwen_mtplx_moe_glu_h_m<bfloat, 640, 10, 4, 64>) kernel_qwen_mtplx_moe_glu_h_m_m4_g64_bf16_type;
template [[host_name("kernel_qwen_mtplx_moe_glu_h_m_m4_g64_bf16")]]
kernel kernel_qwen_mtplx_moe_glu_h_m_m4_g64_bf16_type kernel_qwen_mtplx_moe_glu_h_m<bfloat, 640, 10, 4, 64>;

// Source: mtplx/kernels/moe_glu_decode.py::_SRC_B_M
// File SHA256: ed8b307cd8acd2b111237a57f6adc711d455fac1feb7800df72821c4061b8912
// Body SHA256: fd09231330f1b53b21b7dc2245c652b339834b60ea866afc700784658d1ed08f
template <typename T, int dmodel_c, int n_inter_c, int topk_c, int m_rows_c, int GS_DN>
kernel void kernel_qwen_mtplx_moe_down_y_m(
    device const T* h [[buffer(0)]],
    device const uint32_t* dw [[buffer(1)]],
    device const T* ds [[buffer(2)]],
    device const T* db [[buffer(3)]],
    device const uint32_t* experts [[buffer(4)]],
    device const float* rw [[buffer(5)]],
    device T* y [[buffer(6)]],
    uint3 thread_position_in_threadgroup [[thread_position_in_threadgroup]],
    uint3 threadgroup_position_in_grid [[threadgroup_position_in_grid]]) {
    constexpr int GS = GS_DN;                 // quant group size (32/64 forges)

    const uint tid = thread_position_in_threadgroup.x;
    const uint sg = tid / 32;
    const uint lane = tid % 32;
    const uint gid = threadgroup_position_in_grid.x * 32 + sg;
    const int dmodel = dmodel_c;
    const int n_inter = n_inter_c;
    const int topk = topk_c;
    const int m_rows = m_rows_c;
    if (gid >= (uint)(m_rows * dmodel)) return;
    const int m = gid / dmodel;
    const uint d = gid % dmodel;

    const int ngroups = n_inter / GS;
    float acc = 0.0f;
    for (int e_slot = 0; e_slot < topk; ++e_slot) {
        const uint e = experts[m * topk + e_slot];
        const device uint32_t* wrow = dw + ((size_t)e * dmodel + d) * (n_inter / 8);
        const device T* srow = ds + ((size_t)e * dmodel + d) * ngroups;
        const device T* brow = db + ((size_t)e * dmodel + d) * ngroups;
        const device T* he = h + ((size_t)m * topk + e_slot) * n_inter;
        float dot = 0.0f;
        for (int k = (int)lane; k < n_inter; k += 32) {
            const uint32_t word = wrow[k / 8];
            const float q = (float)((word >> (4 * (k % 8))) & 0xF);
            const int g = k / GS;
            dot = metal::fma((float)srow[g] * q + (float)brow[g], (float)he[k], dot);
        }
        dot = simd_sum(dot);
        acc += (float)rw[m * topk + e_slot] * (float)((T)dot);
    }
    if (lane == 0) {
        y[(size_t)m * dmodel + d] = (T)acc;
    }
}
typedef decltype(kernel_qwen_mtplx_moe_down_y_m<bfloat, 2560, 640, 10, 2, 32>) kernel_qwen_mtplx_moe_down_y_m_m2_g32_bf16_type;
template [[host_name("kernel_qwen_mtplx_moe_down_y_m_m2_g32_bf16")]]
kernel kernel_qwen_mtplx_moe_down_y_m_m2_g32_bf16_type kernel_qwen_mtplx_moe_down_y_m<bfloat, 2560, 640, 10, 2, 32>;
typedef decltype(kernel_qwen_mtplx_moe_down_y_m<bfloat, 2560, 640, 10, 2, 64>) kernel_qwen_mtplx_moe_down_y_m_m2_g64_bf16_type;
template [[host_name("kernel_qwen_mtplx_moe_down_y_m_m2_g64_bf16")]]
kernel kernel_qwen_mtplx_moe_down_y_m_m2_g64_bf16_type kernel_qwen_mtplx_moe_down_y_m<bfloat, 2560, 640, 10, 2, 64>;
typedef decltype(kernel_qwen_mtplx_moe_down_y_m<bfloat, 2560, 640, 10, 3, 32>) kernel_qwen_mtplx_moe_down_y_m_m3_g32_bf16_type;
template [[host_name("kernel_qwen_mtplx_moe_down_y_m_m3_g32_bf16")]]
kernel kernel_qwen_mtplx_moe_down_y_m_m3_g32_bf16_type kernel_qwen_mtplx_moe_down_y_m<bfloat, 2560, 640, 10, 3, 32>;
typedef decltype(kernel_qwen_mtplx_moe_down_y_m<bfloat, 2560, 640, 10, 3, 64>) kernel_qwen_mtplx_moe_down_y_m_m3_g64_bf16_type;
template [[host_name("kernel_qwen_mtplx_moe_down_y_m_m3_g64_bf16")]]
kernel kernel_qwen_mtplx_moe_down_y_m_m3_g64_bf16_type kernel_qwen_mtplx_moe_down_y_m<bfloat, 2560, 640, 10, 3, 64>;
typedef decltype(kernel_qwen_mtplx_moe_down_y_m<bfloat, 2560, 640, 10, 4, 32>) kernel_qwen_mtplx_moe_down_y_m_m4_g32_bf16_type;
template [[host_name("kernel_qwen_mtplx_moe_down_y_m_m4_g32_bf16")]]
kernel kernel_qwen_mtplx_moe_down_y_m_m4_g32_bf16_type kernel_qwen_mtplx_moe_down_y_m<bfloat, 2560, 640, 10, 4, 32>;
typedef decltype(kernel_qwen_mtplx_moe_down_y_m<bfloat, 2560, 640, 10, 4, 64>) kernel_qwen_mtplx_moe_down_y_m_m4_g64_bf16_type;
template [[host_name("kernel_qwen_mtplx_moe_down_y_m_m4_g64_bf16")]]
kernel kernel_qwen_mtplx_moe_down_y_m_m4_g64_bf16_type kernel_qwen_mtplx_moe_down_y_m<bfloat, 2560, 640, 10, 4, 64>;

// mlx-lm 0.31.3 dependency kernel: Copyright 2023 Apple Inc.
// SPDX-License-Identifier: MIT; see MLX-LM-LICENSE.txt.

// Source: mlx_lm/models/gated_delta.py::_make_gated_delta_kernel
// File SHA256: 79c8376a51c694b03e54d2f996ced6ea6c8c42868b8571529f97334db165a3e1
// Body SHA256: 4f2c168c5c59a5f1086db6792b650c5e6dc6915d08174804d8aed8c16e8482a9
template <typename InT, typename StT, int Dk, int Dv, int Hk, int Hv>
kernel void kernel_qwen_mtplx_gated_delta(
    device const InT* q [[buffer(0)]],
    device const InT* k [[buffer(1)]],
    device const InT* v [[buffer(2)]],
    device const float* g [[buffer(3)]],
    device const InT* beta [[buffer(4)]],
    device const StT* state_in [[buffer(5)]],
    constant const int32_t& T [[buffer(6)]],
    device InT* y [[buffer(7)]],
    device StT* state_out [[buffer(8)]],
    uint thread_index_in_simdgroup [[thread_index_in_simdgroup]],
    uint3 thread_position_in_grid [[thread_position_in_grid]],
    uint3 thread_position_in_threadgroup [[thread_position_in_threadgroup]]) {
        auto n = thread_position_in_grid.z;
        auto b_idx = n / Hv;
        auto hv_idx = n % Hv;
        auto hk_idx = hv_idx / (Hv / Hk);
        constexpr int n_per_t = Dk / 32;

        // q, k: [B, T, Hk, Dk]
        auto q_ = q + b_idx * T * Hk * Dk + hk_idx * Dk;
        auto k_ = k + b_idx * T * Hk * Dk + hk_idx * Dk;

        // v, y: [B, T, Hv, Dv]
        auto v_ = v + b_idx * T * Hv * Dv + hv_idx * Dv;
        y += b_idx * T * Hv * Dv + hv_idx * Dv;

        auto dk_idx = thread_position_in_threadgroup.x;
        auto dv_idx = thread_position_in_grid.y;

        // state_in, state_out: [B, Hv, Dv, Dk]
        auto i_state = state_in + (n * Dv + dv_idx) * Dk;
        auto o_state = state_out + (n * Dv + dv_idx) * Dk;

        float state[n_per_t];
        for (int i = 0; i < n_per_t; ++i) {
          auto s_idx = n_per_t * dk_idx + i;
          state[i] = static_cast<float>(i_state[s_idx]);
        }

        // g: [B, T, Hv]
        auto g_ = g + b_idx * T * Hv;
        auto beta_ = beta + b_idx * T * Hv;

        for (int t = 0; t < T; ++t) {
          if (true) {
            float kv_mem = 0.0f;
            for (int i = 0; i < n_per_t; ++i) {
              auto s_idx = n_per_t * dk_idx + i;
              state[i] = state[i] * g_[hv_idx];
              kv_mem += state[i] * k_[s_idx];
            }
            kv_mem = simd_sum(kv_mem);

            auto delta = (v_[dv_idx] - kv_mem) * beta_[hv_idx];

            float out = 0.0f;
            for (int i = 0; i < n_per_t; ++i) {
              auto s_idx = n_per_t * dk_idx + i;
              state[i] = state[i] + k_[s_idx] * delta;
              out += state[i] * q_[s_idx];
            }
            out = simd_sum(out);
            if (thread_index_in_simdgroup == 0) {
              y[dv_idx] = static_cast<InT>(out);
            }
          } else {
            y[dv_idx] = static_cast<InT>(0);
          }
          // Increment data pointers to next time step
          q_ += Hk * Dk;
          k_ += Hk * Dk;
          v_ += Hv * Dv;
          y += Hv * Dv;
          g_ += Hv;
          beta_ += Hv;
        }
        for (int i = 0; i < n_per_t; ++i) {
          auto s_idx = n_per_t * dk_idx + i;
          o_state[s_idx] = static_cast<StT>(state[i]);
        }
    }
typedef decltype(kernel_qwen_mtplx_gated_delta<bfloat, float, 128, 128, 16, 48>) kernel_qwen_mtplx_gated_delta_bf16_type;
template [[host_name("kernel_qwen_mtplx_gated_delta_bf16")]]
kernel kernel_qwen_mtplx_gated_delta_bf16_type kernel_qwen_mtplx_gated_delta<bfloat, float, 128, 128, 16, 48>;

// Source: mlx_lm/models/gated_delta.py::_make_gated_delta_kernel_masked
// File SHA256: 79c8376a51c694b03e54d2f996ced6ea6c8c42868b8571529f97334db165a3e1
// Body SHA256: e7b6c63d60e9bc2b71f3a705c76d4375562a376af9730b1c8ef74d1836eb69f4
template <typename InT, typename StT, int Dk, int Dv, int Hk, int Hv>
kernel void kernel_qwen_mtplx_gated_delta_masked(
    device const InT* q [[buffer(0)]],
    device const InT* k [[buffer(1)]],
    device const InT* v [[buffer(2)]],
    device const float* g [[buffer(3)]],
    device const InT* beta [[buffer(4)]],
    device const StT* state_in [[buffer(5)]],
    constant const int32_t& T [[buffer(6)]],
    device const bool* mask [[buffer(7)]],
    device InT* y [[buffer(8)]],
    device StT* state_out [[buffer(9)]],
    uint thread_index_in_simdgroup [[thread_index_in_simdgroup]],
    uint3 thread_position_in_grid [[thread_position_in_grid]],
    uint3 thread_position_in_threadgroup [[thread_position_in_threadgroup]]) {
        auto n = thread_position_in_grid.z;
        auto b_idx = n / Hv;
        auto hv_idx = n % Hv;
        auto hk_idx = hv_idx / (Hv / Hk);
        constexpr int n_per_t = Dk / 32;

        // q, k: [B, T, Hk, Dk]
        auto q_ = q + b_idx * T * Hk * Dk + hk_idx * Dk;
        auto k_ = k + b_idx * T * Hk * Dk + hk_idx * Dk;

        // v, y: [B, T, Hv, Dv]
        auto v_ = v + b_idx * T * Hv * Dv + hv_idx * Dv;
        y += b_idx * T * Hv * Dv + hv_idx * Dv;

        auto dk_idx = thread_position_in_threadgroup.x;
        auto dv_idx = thread_position_in_grid.y;

        // state_in, state_out: [B, Hv, Dv, Dk]
        auto i_state = state_in + (n * Dv + dv_idx) * Dk;
        auto o_state = state_out + (n * Dv + dv_idx) * Dk;

        float state[n_per_t];
        for (int i = 0; i < n_per_t; ++i) {
          auto s_idx = n_per_t * dk_idx + i;
          state[i] = static_cast<float>(i_state[s_idx]);
        }

        // g: [B, T, Hv]
        auto g_ = g + b_idx * T * Hv;
        auto beta_ = beta + b_idx * T * Hv;

        for (int t = 0; t < T; ++t) {
          if (mask[b_idx * T + t]) {
            float kv_mem = 0.0f;
            for (int i = 0; i < n_per_t; ++i) {
              auto s_idx = n_per_t * dk_idx + i;
              state[i] = state[i] * g_[hv_idx];
              kv_mem += state[i] * k_[s_idx];
            }
            kv_mem = simd_sum(kv_mem);

            auto delta = (v_[dv_idx] - kv_mem) * beta_[hv_idx];

            float out = 0.0f;
            for (int i = 0; i < n_per_t; ++i) {
              auto s_idx = n_per_t * dk_idx + i;
              state[i] = state[i] + k_[s_idx] * delta;
              out += state[i] * q_[s_idx];
            }
            out = simd_sum(out);
            if (thread_index_in_simdgroup == 0) {
              y[dv_idx] = static_cast<InT>(out);
            }
          } else {
            y[dv_idx] = static_cast<InT>(0);
          }
          // Increment data pointers to next time step
          q_ += Hk * Dk;
          k_ += Hk * Dk;
          v_ += Hv * Dv;
          y += Hv * Dv;
          g_ += Hv;
          beta_ += Hv;
        }
        for (int i = 0; i < n_per_t; ++i) {
          auto s_idx = n_per_t * dk_idx + i;
          o_state[s_idx] = static_cast<StT>(state[i]);
        }
    }
typedef decltype(kernel_qwen_mtplx_gated_delta_masked<bfloat, float, 128, 128, 16, 48>) kernel_qwen_mtplx_gated_delta_masked_bf16_type;
template [[host_name("kernel_qwen_mtplx_gated_delta_masked_bf16")]]
kernel kernel_qwen_mtplx_gated_delta_masked_bf16_type kernel_qwen_mtplx_gated_delta_masked<bfloat, float, 128, 128, 16, 48>;

// MTPLX runtime JIT shaders: Copyright Apple Inc. SPDX-License-Identifier: MIT
// See MLX-LM-LICENSE.txt; only includes are resolved and instantiations supplied.
// Runtime unit: indexing.h; file SHA256: e820b8ee2b5132a97122780c12433ebb5100d8078d31e211d0429400a11415bb
// Runtime unit SHA256: 6b177621bd445d19f9b45bfa1f0930e23cf0fe952aa5477cb03789253bea2bf0
// BEGIN RUNTIME UNIT
// Copyright © 2023-2024 Apple Inc.



template <typename IdxT, int NIDX>
struct Indices {
  const array<const device IdxT*, NIDX> buffers;
  const constant int* shapes;
  const constant int64_t* strides;
  const constant bool* row_contiguous;
  const int ndim;
};

template <typename IdxT>
METAL_FUNC size_t offset_neg_idx(IdxT idx, int size) {
  if (is_unsigned_v<IdxT>) {
    return idx;
  } else {
    return (idx < 0) ? idx + size : idx;
  }
}
// END RUNTIME UNIT
// Runtime unit: gather_front.h; file SHA256: 64aacebf6576dfcd389383564fa1214bc87f2a091dd33cc64c598c5367ecab96
// Runtime unit SHA256: 8fefc704f2a7a374684dd7e1793056aab16f0ab75aaa141f965669f2ad857257
// BEGIN RUNTIME UNIT
// Copyright © 2025 Apple Inc.



template <typename T, typename IdxT, typename LocT, int N>
[[kernel]] void gather_front(
    const device T* src,
    const device IdxT* indices,
    device T* out,
    const constant int64_t& stride,
    const constant int& size,
    uint2 index [[thread_position_in_grid]],
    uint2 grid_dim [[threads_per_grid]]) {
  auto idx = offset_neg_idx(indices[index.y], size);
  LocT src_idx = static_cast<LocT>(stride) * idx;
  LocT out_idx = static_cast<LocT>(stride) * index.y;

  int s_idx = N * index.x;
  for (int i = 0; i < N && s_idx < stride; ++i, ++s_idx) {
    out[out_idx + s_idx] = src[src_idx + s_idx];
  }
}
// END RUNTIME UNIT
template [[host_name("kernel_qwen_mtplx_gather_front_bf16_int_1")]]
[[kernel]] decltype(gather_front<bfloat, uint32_t, int, 1>) gather_front<bfloat, uint32_t, int, 1>;
template [[host_name("kernel_qwen_mtplx_gather_front_bf16_int64_t_1")]]
[[kernel]] decltype(gather_front<bfloat, uint32_t, int64_t, 1>) gather_front<bfloat, uint32_t, int64_t, 1>;
template [[host_name("kernel_qwen_mtplx_gather_front_bf16_int_2")]]
[[kernel]] decltype(gather_front<bfloat, uint32_t, int, 2>) gather_front<bfloat, uint32_t, int, 2>;
template [[host_name("kernel_qwen_mtplx_gather_front_bf16_int64_t_2")]]
[[kernel]] decltype(gather_front<bfloat, uint32_t, int64_t, 2>) gather_front<bfloat, uint32_t, int64_t, 2>;
template [[host_name("kernel_qwen_mtplx_gather_front_f16_int_1")]]
[[kernel]] decltype(gather_front<half, uint32_t, int, 1>) gather_front<half, uint32_t, int, 1>;
template [[host_name("kernel_qwen_mtplx_gather_front_f16_int64_t_1")]]
[[kernel]] decltype(gather_front<half, uint32_t, int64_t, 1>) gather_front<half, uint32_t, int64_t, 1>;
template [[host_name("kernel_qwen_mtplx_gather_front_f16_int_2")]]
[[kernel]] decltype(gather_front<half, uint32_t, int, 2>) gather_front<half, uint32_t, int, 2>;
template [[host_name("kernel_qwen_mtplx_gather_front_f16_int64_t_2")]]
[[kernel]] decltype(gather_front<half, uint32_t, int64_t, 2>) gather_front<half, uint32_t, int64_t, 2>;
template [[host_name("kernel_qwen_mtplx_gather_front_u32_int_1")]]
[[kernel]] decltype(gather_front<uint32_t, uint32_t, int, 1>) gather_front<uint32_t, uint32_t, int, 1>;
template [[host_name("kernel_qwen_mtplx_gather_front_u32_int64_t_1")]]
[[kernel]] decltype(gather_front<uint32_t, uint32_t, int64_t, 1>) gather_front<uint32_t, uint32_t, int64_t, 1>;
template [[host_name("kernel_qwen_mtplx_gather_front_f32_int_1")]]
[[kernel]] decltype(gather_front<float, uint32_t, int, 1>) gather_front<float, uint32_t, int, 1>;
template [[host_name("kernel_qwen_mtplx_gather_front_f32_int64_t_1")]]
[[kernel]] decltype(gather_front<float, uint32_t, int64_t, 1>) gather_front<float, uint32_t, int64_t, 1>;
template [[host_name("kernel_qwen_mtplx_gather_front_bf16_idxi64_int_1")]]
[[kernel]] decltype(gather_front<bfloat, int64_t, int, 1>) gather_front<bfloat, int64_t, int, 1>;
template [[host_name("kernel_qwen_mtplx_gather_front_bf16_idxi64_int64_t_1")]]
[[kernel]] decltype(gather_front<bfloat, int64_t, int64_t, 1>) gather_front<bfloat, int64_t, int64_t, 1>;
template [[host_name("kernel_qwen_mtplx_gather_front_bf16_idxi64_int_2")]]
[[kernel]] decltype(gather_front<bfloat, int64_t, int, 2>) gather_front<bfloat, int64_t, int, 2>;
template [[host_name("kernel_qwen_mtplx_gather_front_bf16_idxi64_int64_t_2")]]
[[kernel]] decltype(gather_front<bfloat, int64_t, int64_t, 2>) gather_front<bfloat, int64_t, int64_t, 2>;
template [[host_name("kernel_qwen_mtplx_gather_front_f16_idxi64_int_1")]]
[[kernel]] decltype(gather_front<half, int64_t, int, 1>) gather_front<half, int64_t, int, 1>;
template [[host_name("kernel_qwen_mtplx_gather_front_f16_idxi64_int64_t_1")]]
[[kernel]] decltype(gather_front<half, int64_t, int64_t, 1>) gather_front<half, int64_t, int64_t, 1>;
template [[host_name("kernel_qwen_mtplx_gather_front_f16_idxi64_int_2")]]
[[kernel]] decltype(gather_front<half, int64_t, int, 2>) gather_front<half, int64_t, int, 2>;
template [[host_name("kernel_qwen_mtplx_gather_front_f16_idxi64_int64_t_2")]]
[[kernel]] decltype(gather_front<half, int64_t, int64_t, 2>) gather_front<half, int64_t, int64_t, 2>;
template [[host_name("kernel_qwen_mtplx_gather_front_u32_idxi64_int_1")]]
[[kernel]] decltype(gather_front<uint32_t, int64_t, int, 1>) gather_front<uint32_t, int64_t, int, 1>;
template [[host_name("kernel_qwen_mtplx_gather_front_u32_idxi64_int64_t_1")]]
[[kernel]] decltype(gather_front<uint32_t, int64_t, int64_t, 1>) gather_front<uint32_t, int64_t, int64_t, 1>;
template [[host_name("kernel_qwen_mtplx_gather_front_f32_idxi64_int_1")]]
[[kernel]] decltype(gather_front<float, int64_t, int, 1>) gather_front<float, int64_t, int, 1>;
template [[host_name("kernel_qwen_mtplx_gather_front_f32_idxi64_int64_t_1")]]
[[kernel]] decltype(gather_front<float, int64_t, int64_t, 1>) gather_front<float, int64_t, int64_t, 1>;
template [[host_name("kernel_qwen_mtplx_gather_front_bf16_idxi32_int_1")]]
[[kernel]] decltype(gather_front<bfloat, int32_t, int, 1>) gather_front<bfloat, int32_t, int, 1>;
template [[host_name("kernel_qwen_mtplx_gather_front_bf16_idxi32_int64_t_1")]]
[[kernel]] decltype(gather_front<bfloat, int32_t, int64_t, 1>) gather_front<bfloat, int32_t, int64_t, 1>;
template [[host_name("kernel_qwen_mtplx_gather_front_bf16_idxi32_int_2")]]
[[kernel]] decltype(gather_front<bfloat, int32_t, int, 2>) gather_front<bfloat, int32_t, int, 2>;
template [[host_name("kernel_qwen_mtplx_gather_front_bf16_idxi32_int64_t_2")]]
[[kernel]] decltype(gather_front<bfloat, int32_t, int64_t, 2>) gather_front<bfloat, int32_t, int64_t, 2>;
template [[host_name("kernel_qwen_mtplx_gather_front_f16_idxi32_int_1")]]
[[kernel]] decltype(gather_front<half, int32_t, int, 1>) gather_front<half, int32_t, int, 1>;
template [[host_name("kernel_qwen_mtplx_gather_front_f16_idxi32_int64_t_1")]]
[[kernel]] decltype(gather_front<half, int32_t, int64_t, 1>) gather_front<half, int32_t, int64_t, 1>;
template [[host_name("kernel_qwen_mtplx_gather_front_f16_idxi32_int_2")]]
[[kernel]] decltype(gather_front<half, int32_t, int, 2>) gather_front<half, int32_t, int, 2>;
template [[host_name("kernel_qwen_mtplx_gather_front_f16_idxi32_int64_t_2")]]
[[kernel]] decltype(gather_front<half, int32_t, int64_t, 2>) gather_front<half, int32_t, int64_t, 2>;
template [[host_name("kernel_qwen_mtplx_gather_front_u32_idxi32_int_1")]]
[[kernel]] decltype(gather_front<uint32_t, int32_t, int, 1>) gather_front<uint32_t, int32_t, int, 1>;
template [[host_name("kernel_qwen_mtplx_gather_front_u32_idxi32_int64_t_1")]]
[[kernel]] decltype(gather_front<uint32_t, int32_t, int64_t, 1>) gather_front<uint32_t, int32_t, int64_t, 1>;
template [[host_name("kernel_qwen_mtplx_gather_front_f32_idxi32_int_1")]]
[[kernel]] decltype(gather_front<float, int32_t, int, 1>) gather_front<float, int32_t, int, 1>;
template [[host_name("kernel_qwen_mtplx_gather_front_f32_idxi32_int64_t_1")]]
[[kernel]] decltype(gather_front<float, int32_t, int64_t, 1>) gather_front<float, int32_t, int64_t, 1>;

// Actual runtime-generated SiLU, captured by tools/mtplx-jit-reference.py.
// Copyright Apple Inc. SPDX-License-Identifier: MIT. Only host aliases change.
// Runtime unit: bf16.h; file SHA256: abd87446a310b77ac530ef52a324feae5cb285d03ec9613e3a88ebb71410fdcb
// Runtime unit SHA256: 16104811a899f47193dbafe95ea825d839423d3fe0c0bf9897015aaabf4cf44c
// BEGIN RUNTIME UNIT
// Copyright © 2023 Apple Inc.



using namespace metal;

typedef bfloat bfloat16_t;
inline uint16_t bfloat16_to_uint16(const bfloat16_t x) {
  return as_type<uint16_t>(x);
}

inline bfloat16_t uint16_to_bfloat16(const uint16_t x) {
  return as_type<bfloat16_t>(x);
}
// END RUNTIME UNIT
// Runtime unit: bf16_math.h; file SHA256: 1f374f8380f756eb89acf6a847741cb8fecbe642945e159fb6208d804cc06496
// Runtime unit SHA256: 5ac7275c794d3c7c66a8dda727587b0299788436e6f6a49cf1b5f599608da106
// BEGIN RUNTIME UNIT
// Copyright © 2023 Apple Inc.


///////////////////////////////////////////////////////////////////////////////
// Metal math for bfloat16
///////////////////////////////////////////////////////////////////////////////

/*

Following the Metal Shading Language Specification (Metal 3.1)

"bfloat is an extended itypeing point type that only allows implicit conversion
 to a type of greater itypeing point rank. While bfloat can be implicitly
 converted to itype, it cannot be implicitly converted to half, and neither
 itype nor half can be implicitly converted to bfloat."

Further, as far as I can tell, the stdlib math/simd functions are not defined
for bfloat and calling with an argument of type bfloat will result in that
argument getting implicitly converted to itype which then returns an output
that is (likely) a itype which cannot be implicitly converted into a bfloat

This leads to situations where
bfloat a = 5.0bf;
bfloat b = metal::abs(a); // this will throw an error since abs return itype
bfloat c = static_cast<bfloat>(metal::abs(a)); // this is fine

For the moment, I will be adding overloaded instantiations of the math
functions to accordingly automatically handle the casting

*/

#define instantiate_metal_math_funcs(itype, otype, ctype, mfast)               \
                                                                               \
  METAL_FUNC otype abs(itype x) {                                              \
    return static_cast<otype>(__metal_fabs(static_cast<ctype>(x), mfast));     \
  }                                                                            \
  METAL_FUNC otype acos(itype x) {                                             \
    return static_cast<otype>(__metal_acos(static_cast<ctype>(x), mfast));     \
  }                                                                            \
  METAL_FUNC otype acosh(itype x) {                                            \
    return static_cast<otype>(__metal_acosh(static_cast<ctype>(x), mfast));    \
  }                                                                            \
  METAL_FUNC otype asin(itype x) {                                             \
    return static_cast<otype>(__metal_asin(static_cast<ctype>(x), mfast));     \
  }                                                                            \
  METAL_FUNC otype asinh(itype x) {                                            \
    return static_cast<otype>(__metal_asinh(static_cast<ctype>(x), mfast));    \
  }                                                                            \
  METAL_FUNC otype atan(itype y_over_x) {                                      \
    return static_cast<otype>(                                                 \
        __metal_atan(static_cast<ctype>(y_over_x), mfast));                    \
  }                                                                            \
  METAL_FUNC otype atan2(itype y, itype x) {                                   \
    return static_cast<otype>(                                                 \
        __metal_atan2(static_cast<ctype>(y), static_cast<ctype>(x), mfast));   \
  }                                                                            \
  METAL_FUNC otype atanh(itype x) {                                            \
    return static_cast<otype>(__metal_atanh(static_cast<ctype>(x), mfast));    \
  }                                                                            \
  METAL_FUNC otype ceil(itype x) {                                             \
    return static_cast<otype>(__metal_ceil(static_cast<ctype>(x), mfast));     \
  }                                                                            \
  METAL_FUNC otype cos(itype x) {                                              \
    return static_cast<otype>(__metal_cos(static_cast<ctype>(x), mfast));      \
  }                                                                            \
  METAL_FUNC otype cosh(itype x) {                                             \
    return static_cast<otype>(__metal_cosh(static_cast<ctype>(x), mfast));     \
  }                                                                            \
  METAL_FUNC otype cospi(itype x) {                                            \
    return static_cast<otype>(__metal_cospi(static_cast<ctype>(x), mfast));    \
  }                                                                            \
  METAL_FUNC otype divide(itype x, itype y) {                                  \
    return static_cast<otype>(                                                 \
        __metal_divide(static_cast<ctype>(x), static_cast<ctype>(y), mfast));  \
  }                                                                            \
  METAL_FUNC otype exp(itype x) {                                              \
    return static_cast<otype>(__metal_exp(static_cast<ctype>(x), mfast));      \
  }                                                                            \
  METAL_FUNC otype exp10(itype x) {                                            \
    return static_cast<otype>(__metal_exp10(static_cast<ctype>(x), mfast));    \
  }                                                                            \
  METAL_FUNC otype exp2(itype x) {                                             \
    return static_cast<otype>(__metal_exp2(static_cast<ctype>(x), mfast));     \
  }                                                                            \
  METAL_FUNC otype fabs(itype x) {                                             \
    return static_cast<otype>(__metal_fabs(static_cast<ctype>(x), mfast));     \
  }                                                                            \
  METAL_FUNC otype fdim(itype x, itype y) {                                    \
    ctype t = static_cast<ctype>(x - y);                                       \
    return static_cast<otype>(select(t, ctype(0), t < ctype(0) || x == y));    \
  }                                                                            \
  METAL_FUNC otype floor(itype x) {                                            \
    return static_cast<otype>(__metal_floor(static_cast<ctype>(x), mfast));    \
  }                                                                            \
  METAL_FUNC otype fma(itype x, itype y, itype z) {                            \
    return static_cast<otype>(__metal_fma(                                     \
        static_cast<ctype>(x), static_cast<ctype>(y), static_cast<ctype>(z))); \
  }                                                                            \
  METAL_FUNC otype fmax(itype x, itype y) {                                    \
    return static_cast<otype>(                                                 \
        __metal_fmax(static_cast<ctype>(x), static_cast<ctype>(y), mfast));    \
  }                                                                            \
  METAL_FUNC otype fmax3(itype x, itype y, itype z) {                          \
    return static_cast<otype>(__metal_fmax3(                                   \
        static_cast<ctype>(x),                                                 \
        static_cast<ctype>(y),                                                 \
        static_cast<ctype>(z),                                                 \
        mfast));                                                               \
  }                                                                            \
  METAL_FUNC otype fmedian3(itype x, itype y, itype z) {                       \
    return static_cast<otype>(__metal_fmedian3(                                \
        static_cast<ctype>(x),                                                 \
        static_cast<ctype>(y),                                                 \
        static_cast<ctype>(z),                                                 \
        mfast));                                                               \
  }                                                                            \
  METAL_FUNC otype fmin(itype x, itype y) {                                    \
    return static_cast<otype>(                                                 \
        __metal_fmin(static_cast<ctype>(x), static_cast<ctype>(y), mfast));    \
  }                                                                            \
  METAL_FUNC otype fmin3(itype x, itype y, itype z) {                          \
    return static_cast<otype>(__metal_fmin3(                                   \
        static_cast<ctype>(x),                                                 \
        static_cast<ctype>(y),                                                 \
        static_cast<ctype>(z),                                                 \
        mfast));                                                               \
  }                                                                            \
  METAL_FUNC otype fmod(itype x, itype y) {                                    \
    return static_cast<otype>(                                                 \
        __metal_fmod(static_cast<ctype>(x), static_cast<ctype>(y), mfast));    \
  }                                                                            \
  METAL_FUNC otype fract(itype x) {                                            \
    return static_cast<otype>(__metal_fract(static_cast<ctype>(x), mfast));    \
  }                                                                            \
  METAL_FUNC otype frexp(itype x, thread int& exp) {                           \
    return static_cast<otype>(__metal_frexp(static_cast<ctype>(x), &exp));     \
  }                                                                            \
  METAL_FUNC otype ldexp(itype x, int k) {                                     \
    return static_cast<otype>(__metal_ldexp(static_cast<ctype>(x), k, mfast)); \
  }                                                                            \
  METAL_FUNC otype log(itype x) {                                              \
    return static_cast<otype>(__metal_log(static_cast<ctype>(x), mfast));      \
  }                                                                            \
  METAL_FUNC otype log10(itype x) {                                            \
    return static_cast<otype>(__metal_log10(static_cast<ctype>(x), mfast));    \
  }                                                                            \
  METAL_FUNC otype log2(itype x) {                                             \
    return static_cast<otype>(__metal_log2(static_cast<ctype>(x), mfast));     \
  }                                                                            \
  METAL_FUNC otype max(itype x, itype y) {                                     \
    return static_cast<otype>(                                                 \
        __metal_fmax(static_cast<ctype>(x), static_cast<ctype>(y), mfast));    \
  }                                                                            \
  METAL_FUNC otype max3(itype x, itype y, itype z) {                           \
    return static_cast<otype>(__metal_fmax3(                                   \
        static_cast<ctype>(x),                                                 \
        static_cast<ctype>(y),                                                 \
        static_cast<ctype>(z),                                                 \
        mfast));                                                               \
  }                                                                            \
  METAL_FUNC otype median3(itype x, itype y, itype z) {                        \
    return static_cast<otype>(__metal_fmedian3(                                \
        static_cast<ctype>(x),                                                 \
        static_cast<ctype>(y),                                                 \
        static_cast<ctype>(z),                                                 \
        mfast));                                                               \
  }                                                                            \
  METAL_FUNC otype min(itype x, itype y) {                                     \
    return static_cast<otype>(                                                 \
        __metal_fmin(static_cast<ctype>(x), static_cast<ctype>(y), mfast));    \
  }                                                                            \
  METAL_FUNC otype min3(itype x, itype y, itype z) {                           \
    return static_cast<otype>(__metal_fmin3(                                   \
        static_cast<ctype>(x),                                                 \
        static_cast<ctype>(y),                                                 \
        static_cast<ctype>(z),                                                 \
        mfast));                                                               \
  }                                                                            \
  METAL_FUNC otype nextafter(itype x, itype y) {                               \
    return static_cast<otype>(                                                 \
        __metal_nextafter(static_cast<ctype>(x), static_cast<ctype>(y)));      \
  }                                                                            \
  METAL_FUNC otype pow(itype x, itype y) {                                     \
    return static_cast<otype>(                                                 \
        __metal_pow(static_cast<ctype>(x), static_cast<ctype>(y), mfast));     \
  }                                                                            \
  METAL_FUNC otype powr(itype x, itype y) {                                    \
    return static_cast<otype>(                                                 \
        __metal_powr(static_cast<ctype>(x), static_cast<ctype>(y), mfast));    \
  }                                                                            \
  METAL_FUNC otype rint(itype x) {                                             \
    return static_cast<otype>(__metal_rint(static_cast<ctype>(x), mfast));     \
  }                                                                            \
  METAL_FUNC otype round(itype x) {                                            \
    return static_cast<otype>(__metal_round(static_cast<ctype>(x), mfast));    \
  }                                                                            \
  METAL_FUNC otype rsqrt(itype x) {                                            \
    return static_cast<otype>(__metal_rsqrt(static_cast<ctype>(x), mfast));    \
  }                                                                            \
  METAL_FUNC otype sin(itype x) {                                              \
    return static_cast<otype>(__metal_sin(static_cast<ctype>(x), mfast));      \
  }                                                                            \
  METAL_FUNC otype sinh(itype x) {                                             \
    return static_cast<otype>(__metal_sinh(static_cast<ctype>(x), mfast));     \
  }                                                                            \
  METAL_FUNC otype sinpi(itype x) {                                            \
    return static_cast<otype>(__metal_sinpi(static_cast<ctype>(x), mfast));    \
  }                                                                            \
  METAL_FUNC otype sqrt(itype x) {                                             \
    return static_cast<otype>(__metal_sqrt(static_cast<ctype>(x), mfast));     \
  }                                                                            \
  METAL_FUNC otype tan(itype x) {                                              \
    return static_cast<otype>(__metal_tan(static_cast<ctype>(x), mfast));      \
  }                                                                            \
  METAL_FUNC otype tanh(itype x) {                                             \
    return static_cast<otype>(__metal_tanh(static_cast<ctype>(x), mfast));     \
  }                                                                            \
  METAL_FUNC otype tanpi(itype x) {                                            \
    return static_cast<otype>(__metal_tanpi(static_cast<ctype>(x), mfast));    \
  }                                                                            \
  METAL_FUNC otype trunc(itype x) {                                            \
    return static_cast<otype>(__metal_trunc(static_cast<ctype>(x), mfast));    \
  }

namespace metal {

instantiate_metal_math_funcs(
    bfloat16_t,
    bfloat16_t,
    float,
    __METAL_MAYBE_FAST_MATH__);

namespace fast {

instantiate_metal_math_funcs(
    bfloat16_t,
    bfloat16_t,
    float,
    __METAL_FAST_MATH__);

} // namespace fast

namespace precise {

instantiate_metal_math_funcs(
    bfloat16_t,
    bfloat16_t,
    float,
    __METAL_PRECISE_MATH__);

} // namespace precise

} // namespace metal

///////////////////////////////////////////////////////////////////////////////
// Metal simd for bfloat16
///////////////////////////////////////////////////////////////////////////////

#define instantiate_metal_simd_comm_funcs(                                   \
    itype, otype, ctype, itype_to_ctype, ctype_to_otype)                     \
                                                                             \
  METAL_FUNC otype simd_broadcast(itype data, ushort broadcast_lane_id) {    \
    return ctype_to_otype(                                                   \
        __metal_simd_broadcast(itype_to_ctype(data), broadcast_lane_id));    \
  }                                                                          \
                                                                             \
  METAL_FUNC otype simd_shuffle(itype data, ushort simd_lane_id) {           \
    return ctype_to_otype(                                                   \
        __metal_simd_shuffle(itype_to_ctype(data), simd_lane_id));           \
  }                                                                          \
                                                                             \
  METAL_FUNC otype simd_shuffle_and_fill_down(                               \
      itype data, itype filling_data, ushort delta, ushort modulo) {         \
    return ctype_to_otype(__metal_simd_shuffle_and_fill_down(                \
        itype_to_ctype(data), itype_to_ctype(filling_data), delta, modulo)); \
  }                                                                          \
                                                                             \
  METAL_FUNC otype simd_shuffle_and_fill_down(                               \
      itype data, itype filling_data, ushort delta) {                        \
    return ctype_to_otype(__metal_simd_shuffle_and_fill_down(                \
        itype_to_ctype(data),                                                \
        itype_to_ctype(filling_data),                                        \
        delta,                                                               \
        __metal_get_simdgroup_size(ushort())));                              \
  }                                                                          \
                                                                             \
  METAL_FUNC otype simd_shuffle_and_fill_up(                                 \
      itype data, itype filling_data, ushort delta, ushort modulo) {         \
    return ctype_to_otype(__metal_simd_shuffle_and_fill_up(                  \
        itype_to_ctype(data), itype_to_ctype(filling_data), delta, modulo)); \
  }                                                                          \
                                                                             \
  METAL_FUNC otype simd_shuffle_and_fill_up(                                 \
      itype data, itype filling_data, ushort delta) {                        \
    return ctype_to_otype(__metal_simd_shuffle_and_fill_up(                  \
        itype_to_ctype(data),                                                \
        itype_to_ctype(filling_data),                                        \
        delta,                                                               \
        __metal_get_simdgroup_size(ushort())));                              \
  }                                                                          \
                                                                             \
  METAL_FUNC otype simd_shuffle_down(itype data, ushort delta) {             \
    return ctype_to_otype(                                                   \
        __metal_simd_shuffle_down(itype_to_ctype(data), delta));             \
  }                                                                          \
                                                                             \
  METAL_FUNC otype simd_shuffle_rotate_down(itype data, ushort delta) {      \
    return ctype_to_otype(                                                   \
        __metal_simd_shuffle_rotate_down(itype_to_ctype(data), delta));      \
  }                                                                          \
                                                                             \
  METAL_FUNC otype simd_shuffle_rotate_up(itype data, ushort delta) {        \
    return ctype_to_otype(                                                   \
        __metal_simd_shuffle_rotate_up(itype_to_ctype(data), delta));        \
  }                                                                          \
                                                                             \
  METAL_FUNC otype simd_shuffle_up(itype data, ushort delta) {               \
    return ctype_to_otype(                                                   \
        __metal_simd_shuffle_up(itype_to_ctype(data), delta));               \
  }                                                                          \
                                                                             \
  METAL_FUNC otype simd_shuffle_xor(itype data, ushort mask) {               \
    return ctype_to_otype(                                                   \
        __metal_simd_shuffle_xor(itype_to_ctype(data), mask));               \
  }

#define instantiate_metal_simd_reduction_funcs(itype, otype, ctype)            \
                                                                               \
  METAL_FUNC otype simd_max(itype data) {                                      \
    return static_cast<otype>(__metal_simd_max(static_cast<ctype>(data)));     \
  }                                                                            \
                                                                               \
  METAL_FUNC otype simd_min(itype data) {                                      \
    return static_cast<otype>(__metal_simd_min(static_cast<ctype>(data)));     \
  }                                                                            \
                                                                               \
  METAL_FUNC otype simd_prefix_exclusive_product(itype data) {                 \
    return static_cast<otype>(                                                 \
        __metal_simd_prefix_exclusive_product(static_cast<ctype>(data)));      \
  }                                                                            \
                                                                               \
  METAL_FUNC otype simd_prefix_exclusive_sum(itype data) {                     \
    return static_cast<otype>(                                                 \
        __metal_simd_prefix_exclusive_sum(static_cast<ctype>(data)));          \
  }                                                                            \
                                                                               \
  METAL_FUNC otype simd_prefix_inclusive_product(itype data) {                 \
    return static_cast<otype>(                                                 \
        __metal_simd_prefix_inclusive_product(static_cast<ctype>(data)));      \
  }                                                                            \
                                                                               \
  METAL_FUNC otype simd_prefix_inclusive_sum(itype data) {                     \
    return static_cast<otype>(                                                 \
        __metal_simd_prefix_inclusive_sum(static_cast<ctype>(data)));          \
  }                                                                            \
                                                                               \
  METAL_FUNC otype simd_product(itype data) {                                  \
    return static_cast<otype>(__metal_simd_product(static_cast<ctype>(data))); \
  }                                                                            \
                                                                               \
  METAL_FUNC otype simd_sum(itype data) {                                      \
    return static_cast<otype>(__metal_simd_sum(static_cast<ctype>(data)));     \
  }                                                                            \
                                                                               \
  METAL_FUNC otype simd_xor(itype data) {                                      \
    return static_cast<otype>(__metal_simd_xor(static_cast<ctype>(data)));     \
  }

namespace metal {

instantiate_metal_simd_comm_funcs(
    bfloat16_t,
    bfloat16_t,
    uint16_t,
    bfloat16_to_uint16,
    uint16_to_bfloat16);
instantiate_metal_simd_reduction_funcs(bfloat16_t, bfloat16_t, float);

} // namespace metal
// END RUNTIME UNIT
// Runtime unit: utils.h; file SHA256: 5e1568e9edde9d05dbf86f68fa0d6c6240f2c32b973c7c6a76166b9c0d91543d
// Runtime unit SHA256: 4d161c065a8f4eb7ae4088d48e7e9fe07f19b1492eebd78585bfb67727b52961
// BEGIN RUNTIME UNIT
template <typename IdxT = int64_t>
METAL_FUNC IdxT elem_to_loc_1(uint elem, constant const int64_t& stride) {
  return elem * IdxT(stride);
}

template <typename IdxT = int64_t>
METAL_FUNC IdxT elem_to_loc_2(uint2 elem, constant const int64_t strides[2]) {
  return elem.x * IdxT(strides[1]) + elem.y * IdxT(strides[0]);
}

template <typename IdxT = int64_t>
METAL_FUNC IdxT elem_to_loc_3(uint3 elem, constant const int64_t strides[3]) {
  return elem.x * IdxT(strides[2]) + elem.y * IdxT(strides[1]) +
      elem.z * IdxT(strides[0]);
}

template <typename U, typename T>
inline U cast_to(T val) {
  return static_cast<U>(val);
}

// END RUNTIME UNIT
// Runtime unit: unary_ops.h; file SHA256: 0a5492b65ae39ecb6d8b04e64ea5007e0a4ff60d8d9428559bfbfd03387ece2a
// Runtime unit SHA256: 21d84493d7e2ee35a0ee7de563fc67c4773f7a46b09eded543c5a4db429906cd
// BEGIN RUNTIME UNIT
struct Sigmoid {
  template <typename T>
  T operator()(T x) thread {
    auto y = 1 / (1 + metal::exp(metal::abs(x)));
    return (x < 0) ? y : 1 - y;
  }
};

// END RUNTIME UNIT
// Runtime unit: binary_ops.h; file SHA256: 2dd13c2496f5d6f0856e4ca99db2ceb5c6d7dc0c344b9ba8e41ef3b7d7ed8a97
// Runtime unit SHA256: 8db3e4fac1958175bcf399aac30c935de5150c68cd3b864b1d55896bbc4bef9a
// BEGIN RUNTIME UNIT
struct Multiply {
  template <typename T>
  T operator()(T x, T y) thread {
    return x * y;
  }
};

// END RUNTIME UNIT
// Runtime unit: captured-silu-jit; file SHA256: ed0ccd2cbbbcdacbb3e1c1fe8eb7f6b4d104a00197c97a43c93904fc0124d0af
// Runtime unit SHA256: 6f57e91b38fbad364609a3faa5f2aca2e90bfb015e67921d045ed826601f9d21
// BEGIN RUNTIME UNIT
[[host_name("kernel_qwen_mtplx_silu_bf16_contiguous")]]
[[kernel]] void BV2ISigmoidACV2IBroadcastABDV2IBroadcastBAEV2OMultiplyCD_V_V2_11160318154034397263_contiguous(
    device const bfloat16_t* A [[buffer(0)]],
    device bfloat16_t* B [[buffer(1)]],
    constant const uint& size [[buffer(2)]],
    uint3 pos [[thread_position_in_grid]],
    uint3 grid [[threads_per_grid]]) {
  constexpr int N_ = 1;
  uint index = N_ * pos.x;
  bfloat16_t tmp_A = A[index];
  bfloat16_t tmp_C = Sigmoid()(tmp_A);
  bfloat16_t tmp_D = cast_to<bfloat16_t>(tmp_A);
  bfloat16_t tmp_E = cast_to<bfloat16_t>(tmp_C);
  bfloat16_t tmp_B = Multiply()(tmp_D, tmp_E);
  B[index] = tmp_B;
}
[[host_name("kernel_qwen_mtplx_silu_bf16_contiguous_n")]]
[[kernel]] void BV2ISigmoidACV2IBroadcastABDV2IBroadcastBAEV2OMultiplyCD_V_V2_11160318154034397263_contiguous_n(
    device const bfloat16_t* A [[buffer(0)]],
    device bfloat16_t* B [[buffer(1)]],
    constant const uint& size [[buffer(2)]],
    uint3 pos [[thread_position_in_grid]],
    uint3 grid [[threads_per_grid]]) {
  constexpr int N_ = 4;
  uint index = N_ * pos.x;
  for (int i = 0; i < N_ && index < size; ++i) {
  bfloat16_t tmp_A = A[index];
  bfloat16_t tmp_C = Sigmoid()(tmp_A);
  bfloat16_t tmp_D = cast_to<bfloat16_t>(tmp_A);
  bfloat16_t tmp_E = cast_to<bfloat16_t>(tmp_C);
  bfloat16_t tmp_B = Multiply()(tmp_D, tmp_E);
  B[index] = tmp_B;
  index++;
  }
}
[[host_name("kernel_qwen_mtplx_silu_bf16_contiguous_large")]]
[[kernel]] void BV2ISigmoidACV2IBroadcastABDV2IBroadcastBAEV2OMultiplyCD_V_V2_11160318154034397263_contiguous_large(
    device const bfloat16_t* A [[buffer(0)]],
    device bfloat16_t* B [[buffer(1)]],
    constant const int64_t& size [[buffer(2)]],
    uint3 pos [[thread_position_in_grid]],
    uint3 grid [[threads_per_grid]]) {
  constexpr int N_ = 4;
  int64_t index = N_ * (pos.x + grid.x * int64_t(pos.y));
  for (int i = 0; i < N_ && index < size; ++i) {
  bfloat16_t tmp_A = A[index];
  bfloat16_t tmp_C = Sigmoid()(tmp_A);
  bfloat16_t tmp_D = cast_to<bfloat16_t>(tmp_A);
  bfloat16_t tmp_E = cast_to<bfloat16_t>(tmp_C);
  bfloat16_t tmp_B = Multiply()(tmp_D, tmp_E);
  B[index] = tmp_B;
  index++;
  }
}
[[host_name("kernel_qwen_mtplx_silu_bf16_strided_1")]]
[[kernel]] void BV2ISigmoidACV2IBroadcastABDV2IBroadcastBAEV2OMultiplyCD_V_V2_11160318154034397263_strided_1(
    device const bfloat16_t* A [[buffer(0)]],
    constant const int64_t* in_strides [[buffer(1)]],
    device bfloat16_t* B [[buffer(2)]],
    constant const int* output_shape [[buffer(3)]],
    uint3 pos [[thread_position_in_grid]],
    uint3 grid [[threads_per_grid]]) {
  constexpr int N_ = 1;
  uint index = pos.x + grid.x * (pos.y + uint(grid.y) * pos.z);
  uint index_A = elem_to_loc_1<uint>(pos.x, in_strides[0]);
  bfloat16_t tmp_A = A[index_A];
  bfloat16_t tmp_C = Sigmoid()(tmp_A);
  bfloat16_t tmp_D = cast_to<bfloat16_t>(tmp_A);
  bfloat16_t tmp_E = cast_to<bfloat16_t>(tmp_C);
  bfloat16_t tmp_B = Multiply()(tmp_D, tmp_E);
  B[index] = tmp_B;
}
[[host_name("kernel_qwen_mtplx_silu_bf16_strided_1_large")]]
[[kernel]] void BV2ISigmoidACV2IBroadcastABDV2IBroadcastBAEV2OMultiplyCD_V_V2_11160318154034397263_strided_1_large(
    device const bfloat16_t* A [[buffer(0)]],
    constant const int64_t* in_strides [[buffer(1)]],
    device bfloat16_t* B [[buffer(2)]],
    constant const int* output_shape [[buffer(3)]],
    uint3 pos [[thread_position_in_grid]],
    uint3 grid [[threads_per_grid]]) {
  constexpr int N_ = 1;
  int64_t index = pos.x + grid.x * (pos.y + int64_t(grid.y) * pos.z);
  int64_t index_A = elem_to_loc_1<int64_t>(pos.x, in_strides[0]);
  bfloat16_t tmp_A = A[index_A];
  bfloat16_t tmp_C = Sigmoid()(tmp_A);
  bfloat16_t tmp_D = cast_to<bfloat16_t>(tmp_A);
  bfloat16_t tmp_E = cast_to<bfloat16_t>(tmp_C);
  bfloat16_t tmp_B = Multiply()(tmp_D, tmp_E);
  B[index] = tmp_B;
}
[[host_name("kernel_qwen_mtplx_silu_bf16_strided_2")]]
[[kernel]] void BV2ISigmoidACV2IBroadcastABDV2IBroadcastBAEV2OMultiplyCD_V_V2_11160318154034397263_strided_2(
    device const bfloat16_t* A [[buffer(0)]],
    constant const int64_t* in_strides [[buffer(1)]],
    device bfloat16_t* B [[buffer(2)]],
    constant const int* output_shape [[buffer(3)]],
    uint3 pos [[thread_position_in_grid]],
    uint3 grid [[threads_per_grid]]) {
  constexpr int N_ = 1;
  uint index = pos.x + grid.x * (pos.y + uint(grid.y) * pos.z);
  uint index_A = elem_to_loc_2<uint>({pos.x, pos.y}, in_strides + 0);
  bfloat16_t tmp_A = A[index_A];
  bfloat16_t tmp_C = Sigmoid()(tmp_A);
  bfloat16_t tmp_D = cast_to<bfloat16_t>(tmp_A);
  bfloat16_t tmp_E = cast_to<bfloat16_t>(tmp_C);
  bfloat16_t tmp_B = Multiply()(tmp_D, tmp_E);
  B[index] = tmp_B;
}
[[host_name("kernel_qwen_mtplx_silu_bf16_strided_2_large")]]
[[kernel]] void BV2ISigmoidACV2IBroadcastABDV2IBroadcastBAEV2OMultiplyCD_V_V2_11160318154034397263_strided_2_large(
    device const bfloat16_t* A [[buffer(0)]],
    constant const int64_t* in_strides [[buffer(1)]],
    device bfloat16_t* B [[buffer(2)]],
    constant const int* output_shape [[buffer(3)]],
    uint3 pos [[thread_position_in_grid]],
    uint3 grid [[threads_per_grid]]) {
  constexpr int N_ = 1;
  int64_t index = pos.x + grid.x * (pos.y + int64_t(grid.y) * pos.z);
  int64_t index_A = elem_to_loc_2<int64_t>({pos.x, pos.y}, in_strides + 0);
  bfloat16_t tmp_A = A[index_A];
  bfloat16_t tmp_C = Sigmoid()(tmp_A);
  bfloat16_t tmp_D = cast_to<bfloat16_t>(tmp_A);
  bfloat16_t tmp_E = cast_to<bfloat16_t>(tmp_C);
  bfloat16_t tmp_B = Multiply()(tmp_D, tmp_E);
  B[index] = tmp_B;
}
[[host_name("kernel_qwen_mtplx_silu_bf16_strided_3")]]
[[kernel]] void BV2ISigmoidACV2IBroadcastABDV2IBroadcastBAEV2OMultiplyCD_V_V2_11160318154034397263_strided_3(
    device const bfloat16_t* A [[buffer(0)]],
    constant const int64_t* in_strides [[buffer(1)]],
    device bfloat16_t* B [[buffer(2)]],
    constant const int* output_shape [[buffer(3)]],
    uint3 pos [[thread_position_in_grid]],
    uint3 grid [[threads_per_grid]]) {
  constexpr int N_ = 1;
  uint index = pos.x + grid.x * (pos.y + uint(grid.y) * pos.z);
  uint index_A = elem_to_loc_3<uint>(pos, in_strides + 0);
  bfloat16_t tmp_A = A[index_A];
  bfloat16_t tmp_C = Sigmoid()(tmp_A);
  bfloat16_t tmp_D = cast_to<bfloat16_t>(tmp_A);
  bfloat16_t tmp_E = cast_to<bfloat16_t>(tmp_C);
  bfloat16_t tmp_B = Multiply()(tmp_D, tmp_E);
  B[index] = tmp_B;
}
[[host_name("kernel_qwen_mtplx_silu_bf16_strided_3_large")]]
[[kernel]] void BV2ISigmoidACV2IBroadcastABDV2IBroadcastBAEV2OMultiplyCD_V_V2_11160318154034397263_strided_3_large(
    device const bfloat16_t* A [[buffer(0)]],
    constant const int64_t* in_strides [[buffer(1)]],
    device bfloat16_t* B [[buffer(2)]],
    constant const int* output_shape [[buffer(3)]],
    uint3 pos [[thread_position_in_grid]],
    uint3 grid [[threads_per_grid]]) {
  constexpr int N_ = 1;
  int64_t index = pos.x + grid.x * (pos.y + int64_t(grid.y) * pos.z);
  int64_t index_A = elem_to_loc_3<int64_t>(pos, in_strides + 0);
  bfloat16_t tmp_A = A[index_A];
  bfloat16_t tmp_C = Sigmoid()(tmp_A);
  bfloat16_t tmp_D = cast_to<bfloat16_t>(tmp_A);
  bfloat16_t tmp_E = cast_to<bfloat16_t>(tmp_C);
  bfloat16_t tmp_B = Multiply()(tmp_D, tmp_E);
  B[index] = tmp_B;
}
[[host_name("kernel_qwen_mtplx_silu_bf16_strided_4")]]
[[kernel]] void BV2ISigmoidACV2IBroadcastABDV2IBroadcastBAEV2OMultiplyCD_V_V2_11160318154034397263_strided_4(
    device const bfloat16_t* A [[buffer(0)]],
    constant const int64_t* in_strides [[buffer(1)]],
    device bfloat16_t* B [[buffer(2)]],
    constant const int* output_shape [[buffer(3)]],
    uint3 pos [[thread_position_in_grid]],
    uint3 grid [[threads_per_grid]]) {
  constexpr int N_ = 2;
  int xshape = output_shape[3];
  uint index = N_ * pos.x + xshape * (pos.y + uint(grid.y) * pos.z);
  uint index_A = N_ * pos.x * uint(in_strides[3]) + pos.y * uint(in_strides[2]);
  uint zpos = pos.z;
  for (int d = 1; d >= 0; --d) {
    uint l = zpos % output_shape[d];
    index_A += l * uint(in_strides[0 + d]);
    zpos /= output_shape[d];
  }
  for (int i = 0; i < N_ && (int(N_ * pos.x) + i) < xshape; ++i) {
  bfloat16_t tmp_A = A[index_A];
  bfloat16_t tmp_C = Sigmoid()(tmp_A);
  bfloat16_t tmp_D = cast_to<bfloat16_t>(tmp_A);
  bfloat16_t tmp_E = cast_to<bfloat16_t>(tmp_C);
  bfloat16_t tmp_B = Multiply()(tmp_D, tmp_E);
  B[index] = tmp_B;
  index_A += in_strides[3];
  index++;
  }
}
[[host_name("kernel_qwen_mtplx_silu_bf16_strided_4_large")]]
[[kernel]] void BV2ISigmoidACV2IBroadcastABDV2IBroadcastBAEV2OMultiplyCD_V_V2_11160318154034397263_strided_4_large(
    device const bfloat16_t* A [[buffer(0)]],
    constant const int64_t* in_strides [[buffer(1)]],
    device bfloat16_t* B [[buffer(2)]],
    constant const int* output_shape [[buffer(3)]],
    uint3 pos [[thread_position_in_grid]],
    uint3 grid [[threads_per_grid]]) {
  constexpr int N_ = 4;
  int xshape = output_shape[3];
  int64_t index = N_ * pos.x + xshape * (pos.y + int64_t(grid.y) * pos.z);
  int64_t index_A = N_ * pos.x * int64_t(in_strides[3]) + pos.y * int64_t(in_strides[2]);
  uint zpos = pos.z;
  for (int d = 1; d >= 0; --d) {
    uint l = zpos % output_shape[d];
    index_A += l * int64_t(in_strides[0 + d]);
    zpos /= output_shape[d];
  }
  for (int i = 0; i < N_ && (int(N_ * pos.x) + i) < xshape; ++i) {
  bfloat16_t tmp_A = A[index_A];
  bfloat16_t tmp_C = Sigmoid()(tmp_A);
  bfloat16_t tmp_D = cast_to<bfloat16_t>(tmp_A);
  bfloat16_t tmp_E = cast_to<bfloat16_t>(tmp_C);
  bfloat16_t tmp_B = Multiply()(tmp_D, tmp_E);
  B[index] = tmp_B;
  index_A += in_strides[3];
  index++;
  }
}
[[host_name("kernel_qwen_mtplx_silu_bf16_strided_5")]]
[[kernel]] void BV2ISigmoidACV2IBroadcastABDV2IBroadcastBAEV2OMultiplyCD_V_V2_11160318154034397263_strided_5(
    device const bfloat16_t* A [[buffer(0)]],
    constant const int64_t* in_strides [[buffer(1)]],
    device bfloat16_t* B [[buffer(2)]],
    constant const int* output_shape [[buffer(3)]],
    uint3 pos [[thread_position_in_grid]],
    uint3 grid [[threads_per_grid]]) {
  constexpr int N_ = 2;
  int xshape = output_shape[4];
  uint index = N_ * pos.x + xshape * (pos.y + uint(grid.y) * pos.z);
  uint index_A = N_ * pos.x * uint(in_strides[4]) + pos.y * uint(in_strides[3]);
  uint zpos = pos.z;
  for (int d = 2; d >= 0; --d) {
    uint l = zpos % output_shape[d];
    index_A += l * uint(in_strides[0 + d]);
    zpos /= output_shape[d];
  }
  for (int i = 0; i < N_ && (int(N_ * pos.x) + i) < xshape; ++i) {
  bfloat16_t tmp_A = A[index_A];
  bfloat16_t tmp_C = Sigmoid()(tmp_A);
  bfloat16_t tmp_D = cast_to<bfloat16_t>(tmp_A);
  bfloat16_t tmp_E = cast_to<bfloat16_t>(tmp_C);
  bfloat16_t tmp_B = Multiply()(tmp_D, tmp_E);
  B[index] = tmp_B;
  index_A += in_strides[4];
  index++;
  }
}
[[host_name("kernel_qwen_mtplx_silu_bf16_strided_5_large")]]
[[kernel]] void BV2ISigmoidACV2IBroadcastABDV2IBroadcastBAEV2OMultiplyCD_V_V2_11160318154034397263_strided_5_large(
    device const bfloat16_t* A [[buffer(0)]],
    constant const int64_t* in_strides [[buffer(1)]],
    device bfloat16_t* B [[buffer(2)]],
    constant const int* output_shape [[buffer(3)]],
    uint3 pos [[thread_position_in_grid]],
    uint3 grid [[threads_per_grid]]) {
  constexpr int N_ = 4;
  int xshape = output_shape[4];
  int64_t index = N_ * pos.x + xshape * (pos.y + int64_t(grid.y) * pos.z);
  int64_t index_A = N_ * pos.x * int64_t(in_strides[4]) + pos.y * int64_t(in_strides[3]);
  uint zpos = pos.z;
  for (int d = 2; d >= 0; --d) {
    uint l = zpos % output_shape[d];
    index_A += l * int64_t(in_strides[0 + d]);
    zpos /= output_shape[d];
  }
  for (int i = 0; i < N_ && (int(N_ * pos.x) + i) < xshape; ++i) {
  bfloat16_t tmp_A = A[index_A];
  bfloat16_t tmp_C = Sigmoid()(tmp_A);
  bfloat16_t tmp_D = cast_to<bfloat16_t>(tmp_A);
  bfloat16_t tmp_E = cast_to<bfloat16_t>(tmp_C);
  bfloat16_t tmp_B = Multiply()(tmp_D, tmp_E);
  B[index] = tmp_B;
  index_A += in_strides[4];
  index++;
  }
}
[[host_name("kernel_qwen_mtplx_silu_bf16_strided_6")]]
[[kernel]] void BV2ISigmoidACV2IBroadcastABDV2IBroadcastBAEV2OMultiplyCD_V_V2_11160318154034397263_strided_6(
    device const bfloat16_t* A [[buffer(0)]],
    constant const int64_t* in_strides [[buffer(1)]],
    device bfloat16_t* B [[buffer(2)]],
    constant const int* output_shape [[buffer(3)]],
    uint3 pos [[thread_position_in_grid]],
    uint3 grid [[threads_per_grid]]) {
  constexpr int N_ = 2;
  int xshape = output_shape[5];
  uint index = N_ * pos.x + xshape * (pos.y + uint(grid.y) * pos.z);
  uint index_A = N_ * pos.x * uint(in_strides[5]) + pos.y * uint(in_strides[4]);
  uint zpos = pos.z;
  for (int d = 3; d >= 0; --d) {
    uint l = zpos % output_shape[d];
    index_A += l * uint(in_strides[0 + d]);
    zpos /= output_shape[d];
  }
  for (int i = 0; i < N_ && (int(N_ * pos.x) + i) < xshape; ++i) {
  bfloat16_t tmp_A = A[index_A];
  bfloat16_t tmp_C = Sigmoid()(tmp_A);
  bfloat16_t tmp_D = cast_to<bfloat16_t>(tmp_A);
  bfloat16_t tmp_E = cast_to<bfloat16_t>(tmp_C);
  bfloat16_t tmp_B = Multiply()(tmp_D, tmp_E);
  B[index] = tmp_B;
  index_A += in_strides[5];
  index++;
  }
}
[[host_name("kernel_qwen_mtplx_silu_bf16_strided_6_large")]]
[[kernel]] void BV2ISigmoidACV2IBroadcastABDV2IBroadcastBAEV2OMultiplyCD_V_V2_11160318154034397263_strided_6_large(
    device const bfloat16_t* A [[buffer(0)]],
    constant const int64_t* in_strides [[buffer(1)]],
    device bfloat16_t* B [[buffer(2)]],
    constant const int* output_shape [[buffer(3)]],
    uint3 pos [[thread_position_in_grid]],
    uint3 grid [[threads_per_grid]]) {
  constexpr int N_ = 4;
  int xshape = output_shape[5];
  int64_t index = N_ * pos.x + xshape * (pos.y + int64_t(grid.y) * pos.z);
  int64_t index_A = N_ * pos.x * int64_t(in_strides[5]) + pos.y * int64_t(in_strides[4]);
  uint zpos = pos.z;
  for (int d = 3; d >= 0; --d) {
    uint l = zpos % output_shape[d];
    index_A += l * int64_t(in_strides[0 + d]);
    zpos /= output_shape[d];
  }
  for (int i = 0; i < N_ && (int(N_ * pos.x) + i) < xshape; ++i) {
  bfloat16_t tmp_A = A[index_A];
  bfloat16_t tmp_C = Sigmoid()(tmp_A);
  bfloat16_t tmp_D = cast_to<bfloat16_t>(tmp_A);
  bfloat16_t tmp_E = cast_to<bfloat16_t>(tmp_C);
  bfloat16_t tmp_B = Multiply()(tmp_D, tmp_E);
  B[index] = tmp_B;
  index_A += in_strides[5];
  index++;
  }
}
[[host_name("kernel_qwen_mtplx_silu_bf16_strided_7")]]
[[kernel]] void BV2ISigmoidACV2IBroadcastABDV2IBroadcastBAEV2OMultiplyCD_V_V2_11160318154034397263_strided_7(
    device const bfloat16_t* A [[buffer(0)]],
    constant const int64_t* in_strides [[buffer(1)]],
    device bfloat16_t* B [[buffer(2)]],
    constant const int* output_shape [[buffer(3)]],
    uint3 pos [[thread_position_in_grid]],
    uint3 grid [[threads_per_grid]]) {
  constexpr int N_ = 2;
  int xshape = output_shape[6];
  uint index = N_ * pos.x + xshape * (pos.y + uint(grid.y) * pos.z);
  uint index_A = N_ * pos.x * uint(in_strides[6]) + pos.y * uint(in_strides[5]);
  uint zpos = pos.z;
  for (int d = 4; d >= 0; --d) {
    uint l = zpos % output_shape[d];
    index_A += l * uint(in_strides[0 + d]);
    zpos /= output_shape[d];
  }
  for (int i = 0; i < N_ && (int(N_ * pos.x) + i) < xshape; ++i) {
  bfloat16_t tmp_A = A[index_A];
  bfloat16_t tmp_C = Sigmoid()(tmp_A);
  bfloat16_t tmp_D = cast_to<bfloat16_t>(tmp_A);
  bfloat16_t tmp_E = cast_to<bfloat16_t>(tmp_C);
  bfloat16_t tmp_B = Multiply()(tmp_D, tmp_E);
  B[index] = tmp_B;
  index_A += in_strides[6];
  index++;
  }
}
[[host_name("kernel_qwen_mtplx_silu_bf16_strided_7_large")]]
[[kernel]] void BV2ISigmoidACV2IBroadcastABDV2IBroadcastBAEV2OMultiplyCD_V_V2_11160318154034397263_strided_7_large(
    device const bfloat16_t* A [[buffer(0)]],
    constant const int64_t* in_strides [[buffer(1)]],
    device bfloat16_t* B [[buffer(2)]],
    constant const int* output_shape [[buffer(3)]],
    uint3 pos [[thread_position_in_grid]],
    uint3 grid [[threads_per_grid]]) {
  constexpr int N_ = 4;
  int xshape = output_shape[6];
  int64_t index = N_ * pos.x + xshape * (pos.y + int64_t(grid.y) * pos.z);
  int64_t index_A = N_ * pos.x * int64_t(in_strides[6]) + pos.y * int64_t(in_strides[5]);
  uint zpos = pos.z;
  for (int d = 4; d >= 0; --d) {
    uint l = zpos % output_shape[d];
    index_A += l * int64_t(in_strides[0 + d]);
    zpos /= output_shape[d];
  }
  for (int i = 0; i < N_ && (int(N_ * pos.x) + i) < xshape; ++i) {
  bfloat16_t tmp_A = A[index_A];
  bfloat16_t tmp_C = Sigmoid()(tmp_A);
  bfloat16_t tmp_D = cast_to<bfloat16_t>(tmp_A);
  bfloat16_t tmp_E = cast_to<bfloat16_t>(tmp_C);
  bfloat16_t tmp_B = Multiply()(tmp_D, tmp_E);
  B[index] = tmp_B;
  index_A += in_strides[6];
  index++;
  }
}
[[host_name("kernel_qwen_mtplx_silu_bf16_strided_dynamic")]]
[[kernel]] void BV2ISigmoidACV2IBroadcastABDV2IBroadcastBAEV2OMultiplyCD_V_V2_11160318154034397263_strided_dynamic(
    device const bfloat16_t* A [[buffer(0)]],
    constant const int64_t* in_strides [[buffer(1)]],
    device bfloat16_t* B [[buffer(2)]],
    constant const int* output_shape [[buffer(3)]],
    constant const int& ndim [[buffer(4)]],
    uint3 pos [[thread_position_in_grid]],
    uint3 grid [[threads_per_grid]]) {
  constexpr int N_ = 2;
  int xshape = output_shape[ndim - 1];
  uint index = N_ * pos.x + xshape * (pos.y + uint(grid.y) * pos.z);
  uint index_A = N_ * pos.x * uint(in_strides[ndim * 0 + ndim - 1]) + pos.y * uint(in_strides[ndim * 0 + ndim - 2]);
  uint zpos = pos.z;
  for (int d = ndim - 3; d >= 0; --d) {
    uint l = zpos % output_shape[d];
    index_A += l * uint(in_strides[0 * ndim + d]);
    zpos /= output_shape[d];
  }
  for (int i = 0; i < N_ && (int(N_ * pos.x) + i) < xshape; ++i) {
  bfloat16_t tmp_A = A[index_A];
  bfloat16_t tmp_C = Sigmoid()(tmp_A);
  bfloat16_t tmp_D = cast_to<bfloat16_t>(tmp_A);
  bfloat16_t tmp_E = cast_to<bfloat16_t>(tmp_C);
  bfloat16_t tmp_B = Multiply()(tmp_D, tmp_E);
  B[index] = tmp_B;
  index_A += in_strides[0 * ndim + ndim - 1];
  index++;
  }
}
[[host_name("kernel_qwen_mtplx_silu_bf16_strided_dynamic_large")]]
[[kernel]] void BV2ISigmoidACV2IBroadcastABDV2IBroadcastBAEV2OMultiplyCD_V_V2_11160318154034397263_strided_dynamic_large(
    device const bfloat16_t* A [[buffer(0)]],
    constant const int64_t* in_strides [[buffer(1)]],
    device bfloat16_t* B [[buffer(2)]],
    constant const int* output_shape [[buffer(3)]],
    constant const int& ndim [[buffer(4)]],
    uint3 pos [[thread_position_in_grid]],
    uint3 grid [[threads_per_grid]]) {
  constexpr int N_ = 4;
  int xshape = output_shape[ndim - 1];
  int64_t index = N_ * pos.x + xshape * (pos.y + int64_t(grid.y) * pos.z);
  int64_t index_A = N_ * pos.x * int64_t(in_strides[ndim * 0 + ndim - 1]) + pos.y * int64_t(in_strides[ndim * 0 + ndim - 2]);
  uint zpos = pos.z;
  for (int d = ndim - 3; d >= 0; --d) {
    uint l = zpos % output_shape[d];
    index_A += l * int64_t(in_strides[0 * ndim + d]);
    zpos /= output_shape[d];
  }
  for (int i = 0; i < N_ && (int(N_ * pos.x) + i) < xshape; ++i) {
  bfloat16_t tmp_A = A[index_A];
  bfloat16_t tmp_C = Sigmoid()(tmp_A);
  bfloat16_t tmp_D = cast_to<bfloat16_t>(tmp_A);
  bfloat16_t tmp_E = cast_to<bfloat16_t>(tmp_C);
  bfloat16_t tmp_B = Multiply()(tmp_D, tmp_E);
  B[index] = tmp_B;
  index_A += in_strides[0 * ndim + ndim - 1];
  index++;
  }
}
// END RUNTIME UNIT

// Runtime GatherAxis shader: Copyright Apple Inc. SPDX-License-Identifier: MIT
// Runtime unit: utils.h; file SHA256: 5e1568e9edde9d05dbf86f68fa0d6c6240f2c32b973c7c6a76166b9c0d91543d
// Runtime unit SHA256: b37b61719692317574e960fd7586b5051e6a8c60749705f5bef18f9fc156090a
// BEGIN RUNTIME UNIT
template <typename IdxT = int64_t>
METAL_FUNC IdxT elem_to_loc(
    IdxT elem,
    constant const int* shape,
    constant const int64_t* strides,
    int ndim) {
  IdxT loc = 0;
  for (int i = ndim - 1; i >= 0 && elem > 0; --i) {
    loc += (elem % shape[i]) * IdxT(strides[i]);
    elem /= shape[i];
  }
  return loc;
}

// END RUNTIME UNIT
// Runtime unit: indexing/gather_axis.h; file SHA256: e1a745391ff4990f3f1ad75c5687c3b102dcdc4833d8fbbac38e10f54af29af4
// Runtime unit SHA256: 5d5723cce37465e7b10b6c4df615214e4c67317a6c810e7dcf9ec319036cbd25
// BEGIN RUNTIME UNIT
// Copyright © 2025 Apple Inc.


template <typename T, typename IdxT, typename LocT, bool SrcC, bool IdxC>
[[kernel]] void gather_axis(
    const device T* src [[buffer(0)]],
    const device IdxT* indices [[buffer(1)]],
    device T* out [[buffer(2)]],
    const constant int* shape [[buffer(3)]],
    const constant int64_t* src_strides [[buffer(4)]],
    const constant int64_t* idx_strides [[buffer(5)]],
    const constant size_t& ndim [[buffer(6)]],
    const constant int& axis [[buffer(7)]],
    const constant int& axis_size [[buffer(8)]],
    const constant size_t& src_ax_stride [[buffer(9)]],
    const constant size_t& idx_ax_stride [[buffer(10)]],
    uint3 index [[thread_position_in_grid]],
    uint3 grid_dim [[threads_per_grid]]) {
  LocT elem_idx = index.z * static_cast<LocT>(grid_dim.x);
  LocT out_idx = elem_idx * grid_dim.y + index.x;

  LocT idx_loc = index.y * static_cast<LocT>(idx_ax_stride);
  if (IdxC) {
    idx_loc += out_idx;
  } else {
    idx_loc += elem_to_loc<LocT>(elem_idx + index.x, shape, idx_strides, ndim);
  }

  auto idx_val = indices[idx_loc];
  if (is_signed_v<IdxT>) {
    idx_val = (idx_val < 0) ? idx_val + axis_size : idx_val;
  }

  LocT src_idx = idx_val * static_cast<LocT>(src_ax_stride);
  if (SrcC) {
    src_idx += elem_idx * axis_size + index.x;
  } else {
    src_idx += elem_to_loc<LocT>(elem_idx + index.x, shape, src_strides, ndim);
  }

  out_idx += index.y * static_cast<LocT>(grid_dim.x);
  out[out_idx] = src[src_idx];
}
// END RUNTIME UNIT
template [[host_name("kernel_qwen_mtplx_gather_axis_bf16_int_00")]]
[[kernel]] decltype(gather_axis<bfloat16_t, uint32_t, int, false, false>) gather_axis<bfloat16_t, uint32_t, int, false, false>;
template [[host_name("kernel_qwen_mtplx_gather_axis_bf16_int_01")]]
[[kernel]] decltype(gather_axis<bfloat16_t, uint32_t, int, false, true>) gather_axis<bfloat16_t, uint32_t, int, false, true>;
template [[host_name("kernel_qwen_mtplx_gather_axis_bf16_int_10")]]
[[kernel]] decltype(gather_axis<bfloat16_t, uint32_t, int, true, false>) gather_axis<bfloat16_t, uint32_t, int, true, false>;
template [[host_name("kernel_qwen_mtplx_gather_axis_bf16_int_11")]]
[[kernel]] decltype(gather_axis<bfloat16_t, uint32_t, int, true, true>) gather_axis<bfloat16_t, uint32_t, int, true, true>;
template [[host_name("kernel_qwen_mtplx_gather_axis_bf16_int64_t_00")]]
[[kernel]] decltype(gather_axis<bfloat16_t, uint32_t, int64_t, false, false>) gather_axis<bfloat16_t, uint32_t, int64_t, false, false>;
template [[host_name("kernel_qwen_mtplx_gather_axis_bf16_int64_t_01")]]
[[kernel]] decltype(gather_axis<bfloat16_t, uint32_t, int64_t, false, true>) gather_axis<bfloat16_t, uint32_t, int64_t, false, true>;
template [[host_name("kernel_qwen_mtplx_gather_axis_bf16_int64_t_10")]]
[[kernel]] decltype(gather_axis<bfloat16_t, uint32_t, int64_t, true, false>) gather_axis<bfloat16_t, uint32_t, int64_t, true, false>;
template [[host_name("kernel_qwen_mtplx_gather_axis_bf16_int64_t_11")]]
[[kernel]] decltype(gather_axis<bfloat16_t, uint32_t, int64_t, true, true>) gather_axis<bfloat16_t, uint32_t, int64_t, true, true>;
template [[host_name("kernel_qwen_mtplx_gather_axis_bf16_idxi32_int_00")]]
[[kernel]] decltype(gather_axis<bfloat16_t, int32_t, int, false, false>) gather_axis<bfloat16_t, int32_t, int, false, false>;
template [[host_name("kernel_qwen_mtplx_gather_axis_bf16_idxi32_int_01")]]
[[kernel]] decltype(gather_axis<bfloat16_t, int32_t, int, false, true>) gather_axis<bfloat16_t, int32_t, int, false, true>;
template [[host_name("kernel_qwen_mtplx_gather_axis_bf16_idxi32_int_10")]]
[[kernel]] decltype(gather_axis<bfloat16_t, int32_t, int, true, false>) gather_axis<bfloat16_t, int32_t, int, true, false>;
template [[host_name("kernel_qwen_mtplx_gather_axis_bf16_idxi32_int_11")]]
[[kernel]] decltype(gather_axis<bfloat16_t, int32_t, int, true, true>) gather_axis<bfloat16_t, int32_t, int, true, true>;
template [[host_name("kernel_qwen_mtplx_gather_axis_bf16_idxi32_int64_t_00")]]
[[kernel]] decltype(gather_axis<bfloat16_t, int32_t, int64_t, false, false>) gather_axis<bfloat16_t, int32_t, int64_t, false, false>;
template [[host_name("kernel_qwen_mtplx_gather_axis_bf16_idxi32_int64_t_01")]]
[[kernel]] decltype(gather_axis<bfloat16_t, int32_t, int64_t, false, true>) gather_axis<bfloat16_t, int32_t, int64_t, false, true>;
template [[host_name("kernel_qwen_mtplx_gather_axis_bf16_idxi32_int64_t_10")]]
[[kernel]] decltype(gather_axis<bfloat16_t, int32_t, int64_t, true, false>) gather_axis<bfloat16_t, int32_t, int64_t, true, false>;
template [[host_name("kernel_qwen_mtplx_gather_axis_bf16_idxi32_int64_t_11")]]
[[kernel]] decltype(gather_axis<bfloat16_t, int32_t, int64_t, true, true>) gather_axis<bfloat16_t, int32_t, int64_t, true, true>;
template [[host_name("kernel_qwen_mtplx_gather_axis_f16_int_00")]]
[[kernel]] decltype(gather_axis<half, uint32_t, int, false, false>) gather_axis<half, uint32_t, int, false, false>;
template [[host_name("kernel_qwen_mtplx_gather_axis_f16_int_01")]]
[[kernel]] decltype(gather_axis<half, uint32_t, int, false, true>) gather_axis<half, uint32_t, int, false, true>;
template [[host_name("kernel_qwen_mtplx_gather_axis_f16_int_10")]]
[[kernel]] decltype(gather_axis<half, uint32_t, int, true, false>) gather_axis<half, uint32_t, int, true, false>;
template [[host_name("kernel_qwen_mtplx_gather_axis_f16_int_11")]]
[[kernel]] decltype(gather_axis<half, uint32_t, int, true, true>) gather_axis<half, uint32_t, int, true, true>;
template [[host_name("kernel_qwen_mtplx_gather_axis_f16_int64_t_00")]]
[[kernel]] decltype(gather_axis<half, uint32_t, int64_t, false, false>) gather_axis<half, uint32_t, int64_t, false, false>;
template [[host_name("kernel_qwen_mtplx_gather_axis_f16_int64_t_01")]]
[[kernel]] decltype(gather_axis<half, uint32_t, int64_t, false, true>) gather_axis<half, uint32_t, int64_t, false, true>;
template [[host_name("kernel_qwen_mtplx_gather_axis_f16_int64_t_10")]]
[[kernel]] decltype(gather_axis<half, uint32_t, int64_t, true, false>) gather_axis<half, uint32_t, int64_t, true, false>;
template [[host_name("kernel_qwen_mtplx_gather_axis_f16_int64_t_11")]]
[[kernel]] decltype(gather_axis<half, uint32_t, int64_t, true, true>) gather_axis<half, uint32_t, int64_t, true, true>;
template [[host_name("kernel_qwen_mtplx_gather_axis_f16_idxi32_int_00")]]
[[kernel]] decltype(gather_axis<half, int32_t, int, false, false>) gather_axis<half, int32_t, int, false, false>;
template [[host_name("kernel_qwen_mtplx_gather_axis_f16_idxi32_int_01")]]
[[kernel]] decltype(gather_axis<half, int32_t, int, false, true>) gather_axis<half, int32_t, int, false, true>;
template [[host_name("kernel_qwen_mtplx_gather_axis_f16_idxi32_int_10")]]
[[kernel]] decltype(gather_axis<half, int32_t, int, true, false>) gather_axis<half, int32_t, int, true, false>;
template [[host_name("kernel_qwen_mtplx_gather_axis_f16_idxi32_int_11")]]
[[kernel]] decltype(gather_axis<half, int32_t, int, true, true>) gather_axis<half, int32_t, int, true, true>;
template [[host_name("kernel_qwen_mtplx_gather_axis_f16_idxi32_int64_t_00")]]
[[kernel]] decltype(gather_axis<half, int32_t, int64_t, false, false>) gather_axis<half, int32_t, int64_t, false, false>;
template [[host_name("kernel_qwen_mtplx_gather_axis_f16_idxi32_int64_t_01")]]
[[kernel]] decltype(gather_axis<half, int32_t, int64_t, false, true>) gather_axis<half, int32_t, int64_t, false, true>;
template [[host_name("kernel_qwen_mtplx_gather_axis_f16_idxi32_int64_t_10")]]
[[kernel]] decltype(gather_axis<half, int32_t, int64_t, true, false>) gather_axis<half, int32_t, int64_t, true, false>;
template [[host_name("kernel_qwen_mtplx_gather_axis_f16_idxi32_int64_t_11")]]
[[kernel]] decltype(gather_axis<half, int32_t, int64_t, true, true>) gather_axis<half, int32_t, int64_t, true, true>;
template [[host_name("kernel_qwen_mtplx_gather_axis_f32_int_00")]]
[[kernel]] decltype(gather_axis<float, uint32_t, int, false, false>) gather_axis<float, uint32_t, int, false, false>;
template [[host_name("kernel_qwen_mtplx_gather_axis_f32_int_01")]]
[[kernel]] decltype(gather_axis<float, uint32_t, int, false, true>) gather_axis<float, uint32_t, int, false, true>;
template [[host_name("kernel_qwen_mtplx_gather_axis_f32_int_10")]]
[[kernel]] decltype(gather_axis<float, uint32_t, int, true, false>) gather_axis<float, uint32_t, int, true, false>;
template [[host_name("kernel_qwen_mtplx_gather_axis_f32_int_11")]]
[[kernel]] decltype(gather_axis<float, uint32_t, int, true, true>) gather_axis<float, uint32_t, int, true, true>;
template [[host_name("kernel_qwen_mtplx_gather_axis_f32_int64_t_00")]]
[[kernel]] decltype(gather_axis<float, uint32_t, int64_t, false, false>) gather_axis<float, uint32_t, int64_t, false, false>;
template [[host_name("kernel_qwen_mtplx_gather_axis_f32_int64_t_01")]]
[[kernel]] decltype(gather_axis<float, uint32_t, int64_t, false, true>) gather_axis<float, uint32_t, int64_t, false, true>;
template [[host_name("kernel_qwen_mtplx_gather_axis_f32_int64_t_10")]]
[[kernel]] decltype(gather_axis<float, uint32_t, int64_t, true, false>) gather_axis<float, uint32_t, int64_t, true, false>;
template [[host_name("kernel_qwen_mtplx_gather_axis_f32_int64_t_11")]]
[[kernel]] decltype(gather_axis<float, uint32_t, int64_t, true, true>) gather_axis<float, uint32_t, int64_t, true, true>;
template [[host_name("kernel_qwen_mtplx_gather_axis_f32_idxi32_int_00")]]
[[kernel]] decltype(gather_axis<float, int32_t, int, false, false>) gather_axis<float, int32_t, int, false, false>;
template [[host_name("kernel_qwen_mtplx_gather_axis_f32_idxi32_int_01")]]
[[kernel]] decltype(gather_axis<float, int32_t, int, false, true>) gather_axis<float, int32_t, int, false, true>;
template [[host_name("kernel_qwen_mtplx_gather_axis_f32_idxi32_int_10")]]
[[kernel]] decltype(gather_axis<float, int32_t, int, true, false>) gather_axis<float, int32_t, int, true, false>;
template [[host_name("kernel_qwen_mtplx_gather_axis_f32_idxi32_int_11")]]
[[kernel]] decltype(gather_axis<float, int32_t, int, true, true>) gather_axis<float, int32_t, int, true, true>;
template [[host_name("kernel_qwen_mtplx_gather_axis_f32_idxi32_int64_t_00")]]
[[kernel]] decltype(gather_axis<float, int32_t, int64_t, false, false>) gather_axis<float, int32_t, int64_t, false, false>;
template [[host_name("kernel_qwen_mtplx_gather_axis_f32_idxi32_int64_t_01")]]
[[kernel]] decltype(gather_axis<float, int32_t, int64_t, false, true>) gather_axis<float, int32_t, int64_t, false, true>;
template [[host_name("kernel_qwen_mtplx_gather_axis_f32_idxi32_int64_t_10")]]
[[kernel]] decltype(gather_axis<float, int32_t, int64_t, true, false>) gather_axis<float, int32_t, int64_t, true, false>;
template [[host_name("kernel_qwen_mtplx_gather_axis_f32_idxi32_int64_t_11")]]
[[kernel]] decltype(gather_axis<float, int32_t, int64_t, true, true>) gather_axis<float, int32_t, int64_t, true, true>;
template [[host_name("kernel_qwen_mtplx_gather_axis_u32_int_00")]]
[[kernel]] decltype(gather_axis<uint32_t, uint32_t, int, false, false>) gather_axis<uint32_t, uint32_t, int, false, false>;
template [[host_name("kernel_qwen_mtplx_gather_axis_u32_int_01")]]
[[kernel]] decltype(gather_axis<uint32_t, uint32_t, int, false, true>) gather_axis<uint32_t, uint32_t, int, false, true>;
template [[host_name("kernel_qwen_mtplx_gather_axis_u32_int_10")]]
[[kernel]] decltype(gather_axis<uint32_t, uint32_t, int, true, false>) gather_axis<uint32_t, uint32_t, int, true, false>;
template [[host_name("kernel_qwen_mtplx_gather_axis_u32_int_11")]]
[[kernel]] decltype(gather_axis<uint32_t, uint32_t, int, true, true>) gather_axis<uint32_t, uint32_t, int, true, true>;
template [[host_name("kernel_qwen_mtplx_gather_axis_u32_int64_t_00")]]
[[kernel]] decltype(gather_axis<uint32_t, uint32_t, int64_t, false, false>) gather_axis<uint32_t, uint32_t, int64_t, false, false>;
template [[host_name("kernel_qwen_mtplx_gather_axis_u32_int64_t_01")]]
[[kernel]] decltype(gather_axis<uint32_t, uint32_t, int64_t, false, true>) gather_axis<uint32_t, uint32_t, int64_t, false, true>;
template [[host_name("kernel_qwen_mtplx_gather_axis_u32_int64_t_10")]]
[[kernel]] decltype(gather_axis<uint32_t, uint32_t, int64_t, true, false>) gather_axis<uint32_t, uint32_t, int64_t, true, false>;
template [[host_name("kernel_qwen_mtplx_gather_axis_u32_int64_t_11")]]
[[kernel]] decltype(gather_axis<uint32_t, uint32_t, int64_t, true, true>) gather_axis<uint32_t, uint32_t, int64_t, true, true>;
template [[host_name("kernel_qwen_mtplx_gather_axis_u32_idxi32_int_00")]]
[[kernel]] decltype(gather_axis<uint32_t, int32_t, int, false, false>) gather_axis<uint32_t, int32_t, int, false, false>;
template [[host_name("kernel_qwen_mtplx_gather_axis_u32_idxi32_int_01")]]
[[kernel]] decltype(gather_axis<uint32_t, int32_t, int, false, true>) gather_axis<uint32_t, int32_t, int, false, true>;
template [[host_name("kernel_qwen_mtplx_gather_axis_u32_idxi32_int_10")]]
[[kernel]] decltype(gather_axis<uint32_t, int32_t, int, true, false>) gather_axis<uint32_t, int32_t, int, true, false>;
template [[host_name("kernel_qwen_mtplx_gather_axis_u32_idxi32_int_11")]]
[[kernel]] decltype(gather_axis<uint32_t, int32_t, int, true, true>) gather_axis<uint32_t, int32_t, int, true, true>;
template [[host_name("kernel_qwen_mtplx_gather_axis_u32_idxi32_int64_t_00")]]
[[kernel]] decltype(gather_axis<uint32_t, int32_t, int64_t, false, false>) gather_axis<uint32_t, int32_t, int64_t, false, false>;
template [[host_name("kernel_qwen_mtplx_gather_axis_u32_idxi32_int64_t_01")]]
[[kernel]] decltype(gather_axis<uint32_t, int32_t, int64_t, false, true>) gather_axis<uint32_t, int32_t, int64_t, false, true>;
template [[host_name("kernel_qwen_mtplx_gather_axis_u32_idxi32_int64_t_10")]]
[[kernel]] decltype(gather_axis<uint32_t, int32_t, int64_t, true, false>) gather_axis<uint32_t, int32_t, int64_t, true, false>;
template [[host_name("kernel_qwen_mtplx_gather_axis_u32_idxi32_int64_t_11")]]
[[kernel]] decltype(gather_axis<uint32_t, int32_t, int64_t, true, true>) gather_axis<uint32_t, int32_t, int64_t, true, true>;
template [[host_name("kernel_qwen_mtplx_gather_axis_bool_idxi64_int_00")]]
[[kernel]] decltype(gather_axis<bool, int64_t, int, false, false>) gather_axis<bool, int64_t, int, false, false>;
template [[host_name("kernel_qwen_mtplx_gather_axis_bool_idxi64_int_01")]]
[[kernel]] decltype(gather_axis<bool, int64_t, int, false, true>) gather_axis<bool, int64_t, int, false, true>;
template [[host_name("kernel_qwen_mtplx_gather_axis_bool_idxi64_int_10")]]
[[kernel]] decltype(gather_axis<bool, int64_t, int, true, false>) gather_axis<bool, int64_t, int, true, false>;
template [[host_name("kernel_qwen_mtplx_gather_axis_bool_idxi64_int_11")]]
[[kernel]] decltype(gather_axis<bool, int64_t, int, true, true>) gather_axis<bool, int64_t, int, true, true>;
template [[host_name("kernel_qwen_mtplx_gather_axis_bool_idxi64_int64_t_00")]]
[[kernel]] decltype(gather_axis<bool, int64_t, int64_t, false, false>) gather_axis<bool, int64_t, int64_t, false, false>;
template [[host_name("kernel_qwen_mtplx_gather_axis_bool_idxi64_int64_t_01")]]
[[kernel]] decltype(gather_axis<bool, int64_t, int64_t, false, true>) gather_axis<bool, int64_t, int64_t, false, true>;
template [[host_name("kernel_qwen_mtplx_gather_axis_bool_idxi64_int64_t_10")]]
[[kernel]] decltype(gather_axis<bool, int64_t, int64_t, true, false>) gather_axis<bool, int64_t, int64_t, true, false>;
template [[host_name("kernel_qwen_mtplx_gather_axis_bool_idxi64_int64_t_11")]]
[[kernel]] decltype(gather_axis<bool, int64_t, int64_t, true, true>) gather_axis<bool, int64_t, int64_t, true, true>;
template [[host_name("kernel_qwen_mtplx_gather_axis_i64_idxi64_int_00")]]
[[kernel]] decltype(gather_axis<int64_t, int64_t, int, false, false>) gather_axis<int64_t, int64_t, int, false, false>;
template [[host_name("kernel_qwen_mtplx_gather_axis_i64_idxi64_int_01")]]
[[kernel]] decltype(gather_axis<int64_t, int64_t, int, false, true>) gather_axis<int64_t, int64_t, int, false, true>;
template [[host_name("kernel_qwen_mtplx_gather_axis_i64_idxi64_int_10")]]
[[kernel]] decltype(gather_axis<int64_t, int64_t, int, true, false>) gather_axis<int64_t, int64_t, int, true, false>;
template [[host_name("kernel_qwen_mtplx_gather_axis_i64_idxi64_int_11")]]
[[kernel]] decltype(gather_axis<int64_t, int64_t, int, true, true>) gather_axis<int64_t, int64_t, int, true, true>;
template [[host_name("kernel_qwen_mtplx_gather_axis_i64_idxi64_int64_t_00")]]
[[kernel]] decltype(gather_axis<int64_t, int64_t, int64_t, false, false>) gather_axis<int64_t, int64_t, int64_t, false, false>;
template [[host_name("kernel_qwen_mtplx_gather_axis_i64_idxi64_int64_t_01")]]
[[kernel]] decltype(gather_axis<int64_t, int64_t, int64_t, false, true>) gather_axis<int64_t, int64_t, int64_t, false, true>;
template [[host_name("kernel_qwen_mtplx_gather_axis_i64_idxi64_int64_t_10")]]
[[kernel]] decltype(gather_axis<int64_t, int64_t, int64_t, true, false>) gather_axis<int64_t, int64_t, int64_t, true, false>;
template [[host_name("kernel_qwen_mtplx_gather_axis_i64_idxi64_int64_t_11")]]
[[kernel]] decltype(gather_axis<int64_t, int64_t, int64_t, true, true>) gather_axis<int64_t, int64_t, int64_t, true, true>;
// Runtime unit: gather.h; file SHA256: 3b2f5b21cd2e71427c9641368a457840f5cb131e202095cd1352f983ce7541ce
// Runtime unit SHA256: a883c916496d227fb6bc2c9bf1586e56a4331d097f63d3e42faa69786d4e2992
// BEGIN RUNTIME UNIT
// Copyright © 2024 Apple Inc.



template <typename T, typename IdxT, int NIDX, int IDX_NDIM, typename LocT>
METAL_FUNC void gather_impl(
    const device T* src [[buffer(0)]],
    device T* out [[buffer(1)]],
    const constant int* src_shape [[buffer(2)]],
    const constant int64_t* src_strides [[buffer(3)]],
    const constant size_t& src_ndim [[buffer(4)]],
    const constant int* slice_sizes [[buffer(5)]],
    const constant int* axes [[buffer(6)]],
    const thread Indices<IdxT, NIDX>& indices,
    uint3 index [[thread_position_in_grid]],
    uint3 grid_dim [[threads_per_grid]]) {
  LocT src_idx = 0;
  for (int i = 0; i < NIDX; ++i) {
    LocT idx_loc;
    if (IDX_NDIM == 0) {
      idx_loc = 0;
    } else if (IDX_NDIM == 1) {
      idx_loc = index.x * static_cast<LocT>(indices.strides[indices.ndim * i]);
    } else {
      idx_loc = index.x * static_cast<LocT>(indices.strides[indices.ndim * i]);
      idx_loc += indices.row_contiguous[i]
          ? index.y
          : elem_to_loc<LocT>(
                index.y,
                &indices.shapes[indices.ndim * i + 1],
                &indices.strides[indices.ndim * i + 1],
                indices.ndim - 1);
    }
    auto ax = axes[i];
    auto idx_val = offset_neg_idx(indices.buffers[i][idx_loc], src_shape[ax]);
    src_idx += static_cast<LocT>(idx_val) * static_cast<LocT>(src_strides[ax]);
  }

  auto src_offset =
      elem_to_loc<LocT>(index.z, slice_sizes, src_strides, src_ndim);

  LocT out_idx = index.z;
  if (IDX_NDIM == 1) {
    out_idx += static_cast<LocT>(grid_dim.z) * index.x;
  } else if (IDX_NDIM >= 2) {
    out_idx += grid_dim.z * (index.x * static_cast<LocT>(grid_dim.y) + index.y);
  }
  out[out_idx] = src[src_offset + src_idx];
}
// END RUNTIME UNIT
// Runtime unit: gather JIT wrappers; file SHA256: 1b38dbdf3120eca3e5266ffd6d69691591c9432c0b5acece4c5ab2de7a9e802f
// Runtime unit SHA256: 8afe2e20ea04b1e2f3504bc579d99e8766382ad22775179a03ff30393bb837c9
// BEGIN RUNTIME UNIT

[[kernel]] void kernel_qwen_mtplx_gather_bf16_idxi64_1_0_int(
    const device bfloat* src [[buffer(0)]],
    device bfloat* out [[buffer(1)]],
    const constant int* src_shape [[buffer(2)]],
    const constant int64_t* src_strides [[buffer(3)]],
    const constant size_t& src_ndim [[buffer(4)]],
    const constant int* slice_sizes [[buffer(5)]],
    const constant int* axes [[buffer(6)]],
    const constant int* idx_shapes [[buffer(7)]],
    const constant int64_t* idx_strides [[buffer(8)]],
    const constant bool* idx_contigs [[buffer(9)]],
    const constant int& idx_ndim [[buffer(10)]],
    const device int64_t *idx0 [[buffer(20)]],
    uint3 index [[thread_position_in_grid]],
    uint3 grid_dim [[threads_per_grid]]) {
  Indices<int64_t, 1> idxs{
    { idx0 }, idx_shapes, idx_strides, idx_contigs, idx_ndim};

  return gather_impl<bfloat, int64_t, 1, 0, int>(
      src,
      out,
      src_shape,
      src_strides,
      src_ndim,
      slice_sizes,
      axes,
      idxs,
      index,
      grid_dim);
}

[[kernel]] void kernel_qwen_mtplx_gather_bf16_idxi64_1_0_int64_t(
    const device bfloat* src [[buffer(0)]],
    device bfloat* out [[buffer(1)]],
    const constant int* src_shape [[buffer(2)]],
    const constant int64_t* src_strides [[buffer(3)]],
    const constant size_t& src_ndim [[buffer(4)]],
    const constant int* slice_sizes [[buffer(5)]],
    const constant int* axes [[buffer(6)]],
    const constant int* idx_shapes [[buffer(7)]],
    const constant int64_t* idx_strides [[buffer(8)]],
    const constant bool* idx_contigs [[buffer(9)]],
    const constant int& idx_ndim [[buffer(10)]],
    const device int64_t *idx0 [[buffer(20)]],
    uint3 index [[thread_position_in_grid]],
    uint3 grid_dim [[threads_per_grid]]) {
  Indices<int64_t, 1> idxs{
    { idx0 }, idx_shapes, idx_strides, idx_contigs, idx_ndim};

  return gather_impl<bfloat, int64_t, 1, 0, int64_t>(
      src,
      out,
      src_shape,
      src_strides,
      src_ndim,
      slice_sizes,
      axes,
      idxs,
      index,
      grid_dim);
}

[[kernel]] void kernel_qwen_mtplx_gather_bf16_idxi64_1_1_int(
    const device bfloat* src [[buffer(0)]],
    device bfloat* out [[buffer(1)]],
    const constant int* src_shape [[buffer(2)]],
    const constant int64_t* src_strides [[buffer(3)]],
    const constant size_t& src_ndim [[buffer(4)]],
    const constant int* slice_sizes [[buffer(5)]],
    const constant int* axes [[buffer(6)]],
    const constant int* idx_shapes [[buffer(7)]],
    const constant int64_t* idx_strides [[buffer(8)]],
    const constant bool* idx_contigs [[buffer(9)]],
    const constant int& idx_ndim [[buffer(10)]],
    const device int64_t *idx0 [[buffer(20)]],
    uint3 index [[thread_position_in_grid]],
    uint3 grid_dim [[threads_per_grid]]) {
  Indices<int64_t, 1> idxs{
    { idx0 }, idx_shapes, idx_strides, idx_contigs, idx_ndim};

  return gather_impl<bfloat, int64_t, 1, 1, int>(
      src,
      out,
      src_shape,
      src_strides,
      src_ndim,
      slice_sizes,
      axes,
      idxs,
      index,
      grid_dim);
}

[[kernel]] void kernel_qwen_mtplx_gather_bf16_idxi64_1_1_int64_t(
    const device bfloat* src [[buffer(0)]],
    device bfloat* out [[buffer(1)]],
    const constant int* src_shape [[buffer(2)]],
    const constant int64_t* src_strides [[buffer(3)]],
    const constant size_t& src_ndim [[buffer(4)]],
    const constant int* slice_sizes [[buffer(5)]],
    const constant int* axes [[buffer(6)]],
    const constant int* idx_shapes [[buffer(7)]],
    const constant int64_t* idx_strides [[buffer(8)]],
    const constant bool* idx_contigs [[buffer(9)]],
    const constant int& idx_ndim [[buffer(10)]],
    const device int64_t *idx0 [[buffer(20)]],
    uint3 index [[thread_position_in_grid]],
    uint3 grid_dim [[threads_per_grid]]) {
  Indices<int64_t, 1> idxs{
    { idx0 }, idx_shapes, idx_strides, idx_contigs, idx_ndim};

  return gather_impl<bfloat, int64_t, 1, 1, int64_t>(
      src,
      out,
      src_shape,
      src_strides,
      src_ndim,
      slice_sizes,
      axes,
      idxs,
      index,
      grid_dim);
}

[[kernel]] void kernel_qwen_mtplx_gather_bf16_idxi64_1_2_int(
    const device bfloat* src [[buffer(0)]],
    device bfloat* out [[buffer(1)]],
    const constant int* src_shape [[buffer(2)]],
    const constant int64_t* src_strides [[buffer(3)]],
    const constant size_t& src_ndim [[buffer(4)]],
    const constant int* slice_sizes [[buffer(5)]],
    const constant int* axes [[buffer(6)]],
    const constant int* idx_shapes [[buffer(7)]],
    const constant int64_t* idx_strides [[buffer(8)]],
    const constant bool* idx_contigs [[buffer(9)]],
    const constant int& idx_ndim [[buffer(10)]],
    const device int64_t *idx0 [[buffer(20)]],
    uint3 index [[thread_position_in_grid]],
    uint3 grid_dim [[threads_per_grid]]) {
  Indices<int64_t, 1> idxs{
    { idx0 }, idx_shapes, idx_strides, idx_contigs, idx_ndim};

  return gather_impl<bfloat, int64_t, 1, 2, int>(
      src,
      out,
      src_shape,
      src_strides,
      src_ndim,
      slice_sizes,
      axes,
      idxs,
      index,
      grid_dim);
}

[[kernel]] void kernel_qwen_mtplx_gather_bf16_idxi64_1_2_int64_t(
    const device bfloat* src [[buffer(0)]],
    device bfloat* out [[buffer(1)]],
    const constant int* src_shape [[buffer(2)]],
    const constant int64_t* src_strides [[buffer(3)]],
    const constant size_t& src_ndim [[buffer(4)]],
    const constant int* slice_sizes [[buffer(5)]],
    const constant int* axes [[buffer(6)]],
    const constant int* idx_shapes [[buffer(7)]],
    const constant int64_t* idx_strides [[buffer(8)]],
    const constant bool* idx_contigs [[buffer(9)]],
    const constant int& idx_ndim [[buffer(10)]],
    const device int64_t *idx0 [[buffer(20)]],
    uint3 index [[thread_position_in_grid]],
    uint3 grid_dim [[threads_per_grid]]) {
  Indices<int64_t, 1> idxs{
    { idx0 }, idx_shapes, idx_strides, idx_contigs, idx_ndim};

  return gather_impl<bfloat, int64_t, 1, 2, int64_t>(
      src,
      out,
      src_shape,
      src_strides,
      src_ndim,
      slice_sizes,
      axes,
      idxs,
      index,
      grid_dim);
}

[[kernel]] void kernel_qwen_mtplx_gather_bf16_idxi64_1_3_int(
    const device bfloat* src [[buffer(0)]],
    device bfloat* out [[buffer(1)]],
    const constant int* src_shape [[buffer(2)]],
    const constant int64_t* src_strides [[buffer(3)]],
    const constant size_t& src_ndim [[buffer(4)]],
    const constant int* slice_sizes [[buffer(5)]],
    const constant int* axes [[buffer(6)]],
    const constant int* idx_shapes [[buffer(7)]],
    const constant int64_t* idx_strides [[buffer(8)]],
    const constant bool* idx_contigs [[buffer(9)]],
    const constant int& idx_ndim [[buffer(10)]],
    const device int64_t *idx0 [[buffer(20)]],
    uint3 index [[thread_position_in_grid]],
    uint3 grid_dim [[threads_per_grid]]) {
  Indices<int64_t, 1> idxs{
    { idx0 }, idx_shapes, idx_strides, idx_contigs, idx_ndim};

  return gather_impl<bfloat, int64_t, 1, 3, int>(
      src,
      out,
      src_shape,
      src_strides,
      src_ndim,
      slice_sizes,
      axes,
      idxs,
      index,
      grid_dim);
}

[[kernel]] void kernel_qwen_mtplx_gather_bf16_idxi64_1_3_int64_t(
    const device bfloat* src [[buffer(0)]],
    device bfloat* out [[buffer(1)]],
    const constant int* src_shape [[buffer(2)]],
    const constant int64_t* src_strides [[buffer(3)]],
    const constant size_t& src_ndim [[buffer(4)]],
    const constant int* slice_sizes [[buffer(5)]],
    const constant int* axes [[buffer(6)]],
    const constant int* idx_shapes [[buffer(7)]],
    const constant int64_t* idx_strides [[buffer(8)]],
    const constant bool* idx_contigs [[buffer(9)]],
    const constant int& idx_ndim [[buffer(10)]],
    const device int64_t *idx0 [[buffer(20)]],
    uint3 index [[thread_position_in_grid]],
    uint3 grid_dim [[threads_per_grid]]) {
  Indices<int64_t, 1> idxs{
    { idx0 }, idx_shapes, idx_strides, idx_contigs, idx_ndim};

  return gather_impl<bfloat, int64_t, 1, 3, int64_t>(
      src,
      out,
      src_shape,
      src_strides,
      src_ndim,
      slice_sizes,
      axes,
      idxs,
      index,
      grid_dim);
}

[[kernel]] void kernel_qwen_mtplx_gather_bf16_idxi64_1_4_int(
    const device bfloat* src [[buffer(0)]],
    device bfloat* out [[buffer(1)]],
    const constant int* src_shape [[buffer(2)]],
    const constant int64_t* src_strides [[buffer(3)]],
    const constant size_t& src_ndim [[buffer(4)]],
    const constant int* slice_sizes [[buffer(5)]],
    const constant int* axes [[buffer(6)]],
    const constant int* idx_shapes [[buffer(7)]],
    const constant int64_t* idx_strides [[buffer(8)]],
    const constant bool* idx_contigs [[buffer(9)]],
    const constant int& idx_ndim [[buffer(10)]],
    const device int64_t *idx0 [[buffer(20)]],
    uint3 index [[thread_position_in_grid]],
    uint3 grid_dim [[threads_per_grid]]) {
  Indices<int64_t, 1> idxs{
    { idx0 }, idx_shapes, idx_strides, idx_contigs, idx_ndim};

  return gather_impl<bfloat, int64_t, 1, 4, int>(
      src,
      out,
      src_shape,
      src_strides,
      src_ndim,
      slice_sizes,
      axes,
      idxs,
      index,
      grid_dim);
}

[[kernel]] void kernel_qwen_mtplx_gather_bf16_idxi64_1_4_int64_t(
    const device bfloat* src [[buffer(0)]],
    device bfloat* out [[buffer(1)]],
    const constant int* src_shape [[buffer(2)]],
    const constant int64_t* src_strides [[buffer(3)]],
    const constant size_t& src_ndim [[buffer(4)]],
    const constant int* slice_sizes [[buffer(5)]],
    const constant int* axes [[buffer(6)]],
    const constant int* idx_shapes [[buffer(7)]],
    const constant int64_t* idx_strides [[buffer(8)]],
    const constant bool* idx_contigs [[buffer(9)]],
    const constant int& idx_ndim [[buffer(10)]],
    const device int64_t *idx0 [[buffer(20)]],
    uint3 index [[thread_position_in_grid]],
    uint3 grid_dim [[threads_per_grid]]) {
  Indices<int64_t, 1> idxs{
    { idx0 }, idx_shapes, idx_strides, idx_contigs, idx_ndim};

  return gather_impl<bfloat, int64_t, 1, 4, int64_t>(
      src,
      out,
      src_shape,
      src_strides,
      src_ndim,
      slice_sizes,
      axes,
      idxs,
      index,
      grid_dim);
}

[[kernel]] void kernel_qwen_mtplx_gather_f16_idxi64_1_0_int(
    const device half* src [[buffer(0)]],
    device half* out [[buffer(1)]],
    const constant int* src_shape [[buffer(2)]],
    const constant int64_t* src_strides [[buffer(3)]],
    const constant size_t& src_ndim [[buffer(4)]],
    const constant int* slice_sizes [[buffer(5)]],
    const constant int* axes [[buffer(6)]],
    const constant int* idx_shapes [[buffer(7)]],
    const constant int64_t* idx_strides [[buffer(8)]],
    const constant bool* idx_contigs [[buffer(9)]],
    const constant int& idx_ndim [[buffer(10)]],
    const device int64_t *idx0 [[buffer(20)]],
    uint3 index [[thread_position_in_grid]],
    uint3 grid_dim [[threads_per_grid]]) {
  Indices<int64_t, 1> idxs{
    { idx0 }, idx_shapes, idx_strides, idx_contigs, idx_ndim};

  return gather_impl<half, int64_t, 1, 0, int>(
      src,
      out,
      src_shape,
      src_strides,
      src_ndim,
      slice_sizes,
      axes,
      idxs,
      index,
      grid_dim);
}

[[kernel]] void kernel_qwen_mtplx_gather_f16_idxi64_1_0_int64_t(
    const device half* src [[buffer(0)]],
    device half* out [[buffer(1)]],
    const constant int* src_shape [[buffer(2)]],
    const constant int64_t* src_strides [[buffer(3)]],
    const constant size_t& src_ndim [[buffer(4)]],
    const constant int* slice_sizes [[buffer(5)]],
    const constant int* axes [[buffer(6)]],
    const constant int* idx_shapes [[buffer(7)]],
    const constant int64_t* idx_strides [[buffer(8)]],
    const constant bool* idx_contigs [[buffer(9)]],
    const constant int& idx_ndim [[buffer(10)]],
    const device int64_t *idx0 [[buffer(20)]],
    uint3 index [[thread_position_in_grid]],
    uint3 grid_dim [[threads_per_grid]]) {
  Indices<int64_t, 1> idxs{
    { idx0 }, idx_shapes, idx_strides, idx_contigs, idx_ndim};

  return gather_impl<half, int64_t, 1, 0, int64_t>(
      src,
      out,
      src_shape,
      src_strides,
      src_ndim,
      slice_sizes,
      axes,
      idxs,
      index,
      grid_dim);
}

[[kernel]] void kernel_qwen_mtplx_gather_f16_idxi64_1_1_int(
    const device half* src [[buffer(0)]],
    device half* out [[buffer(1)]],
    const constant int* src_shape [[buffer(2)]],
    const constant int64_t* src_strides [[buffer(3)]],
    const constant size_t& src_ndim [[buffer(4)]],
    const constant int* slice_sizes [[buffer(5)]],
    const constant int* axes [[buffer(6)]],
    const constant int* idx_shapes [[buffer(7)]],
    const constant int64_t* idx_strides [[buffer(8)]],
    const constant bool* idx_contigs [[buffer(9)]],
    const constant int& idx_ndim [[buffer(10)]],
    const device int64_t *idx0 [[buffer(20)]],
    uint3 index [[thread_position_in_grid]],
    uint3 grid_dim [[threads_per_grid]]) {
  Indices<int64_t, 1> idxs{
    { idx0 }, idx_shapes, idx_strides, idx_contigs, idx_ndim};

  return gather_impl<half, int64_t, 1, 1, int>(
      src,
      out,
      src_shape,
      src_strides,
      src_ndim,
      slice_sizes,
      axes,
      idxs,
      index,
      grid_dim);
}

[[kernel]] void kernel_qwen_mtplx_gather_f16_idxi64_1_1_int64_t(
    const device half* src [[buffer(0)]],
    device half* out [[buffer(1)]],
    const constant int* src_shape [[buffer(2)]],
    const constant int64_t* src_strides [[buffer(3)]],
    const constant size_t& src_ndim [[buffer(4)]],
    const constant int* slice_sizes [[buffer(5)]],
    const constant int* axes [[buffer(6)]],
    const constant int* idx_shapes [[buffer(7)]],
    const constant int64_t* idx_strides [[buffer(8)]],
    const constant bool* idx_contigs [[buffer(9)]],
    const constant int& idx_ndim [[buffer(10)]],
    const device int64_t *idx0 [[buffer(20)]],
    uint3 index [[thread_position_in_grid]],
    uint3 grid_dim [[threads_per_grid]]) {
  Indices<int64_t, 1> idxs{
    { idx0 }, idx_shapes, idx_strides, idx_contigs, idx_ndim};

  return gather_impl<half, int64_t, 1, 1, int64_t>(
      src,
      out,
      src_shape,
      src_strides,
      src_ndim,
      slice_sizes,
      axes,
      idxs,
      index,
      grid_dim);
}

[[kernel]] void kernel_qwen_mtplx_gather_f16_idxi64_1_2_int(
    const device half* src [[buffer(0)]],
    device half* out [[buffer(1)]],
    const constant int* src_shape [[buffer(2)]],
    const constant int64_t* src_strides [[buffer(3)]],
    const constant size_t& src_ndim [[buffer(4)]],
    const constant int* slice_sizes [[buffer(5)]],
    const constant int* axes [[buffer(6)]],
    const constant int* idx_shapes [[buffer(7)]],
    const constant int64_t* idx_strides [[buffer(8)]],
    const constant bool* idx_contigs [[buffer(9)]],
    const constant int& idx_ndim [[buffer(10)]],
    const device int64_t *idx0 [[buffer(20)]],
    uint3 index [[thread_position_in_grid]],
    uint3 grid_dim [[threads_per_grid]]) {
  Indices<int64_t, 1> idxs{
    { idx0 }, idx_shapes, idx_strides, idx_contigs, idx_ndim};

  return gather_impl<half, int64_t, 1, 2, int>(
      src,
      out,
      src_shape,
      src_strides,
      src_ndim,
      slice_sizes,
      axes,
      idxs,
      index,
      grid_dim);
}

[[kernel]] void kernel_qwen_mtplx_gather_f16_idxi64_1_2_int64_t(
    const device half* src [[buffer(0)]],
    device half* out [[buffer(1)]],
    const constant int* src_shape [[buffer(2)]],
    const constant int64_t* src_strides [[buffer(3)]],
    const constant size_t& src_ndim [[buffer(4)]],
    const constant int* slice_sizes [[buffer(5)]],
    const constant int* axes [[buffer(6)]],
    const constant int* idx_shapes [[buffer(7)]],
    const constant int64_t* idx_strides [[buffer(8)]],
    const constant bool* idx_contigs [[buffer(9)]],
    const constant int& idx_ndim [[buffer(10)]],
    const device int64_t *idx0 [[buffer(20)]],
    uint3 index [[thread_position_in_grid]],
    uint3 grid_dim [[threads_per_grid]]) {
  Indices<int64_t, 1> idxs{
    { idx0 }, idx_shapes, idx_strides, idx_contigs, idx_ndim};

  return gather_impl<half, int64_t, 1, 2, int64_t>(
      src,
      out,
      src_shape,
      src_strides,
      src_ndim,
      slice_sizes,
      axes,
      idxs,
      index,
      grid_dim);
}

[[kernel]] void kernel_qwen_mtplx_gather_f16_idxi64_1_3_int(
    const device half* src [[buffer(0)]],
    device half* out [[buffer(1)]],
    const constant int* src_shape [[buffer(2)]],
    const constant int64_t* src_strides [[buffer(3)]],
    const constant size_t& src_ndim [[buffer(4)]],
    const constant int* slice_sizes [[buffer(5)]],
    const constant int* axes [[buffer(6)]],
    const constant int* idx_shapes [[buffer(7)]],
    const constant int64_t* idx_strides [[buffer(8)]],
    const constant bool* idx_contigs [[buffer(9)]],
    const constant int& idx_ndim [[buffer(10)]],
    const device int64_t *idx0 [[buffer(20)]],
    uint3 index [[thread_position_in_grid]],
    uint3 grid_dim [[threads_per_grid]]) {
  Indices<int64_t, 1> idxs{
    { idx0 }, idx_shapes, idx_strides, idx_contigs, idx_ndim};

  return gather_impl<half, int64_t, 1, 3, int>(
      src,
      out,
      src_shape,
      src_strides,
      src_ndim,
      slice_sizes,
      axes,
      idxs,
      index,
      grid_dim);
}

[[kernel]] void kernel_qwen_mtplx_gather_f16_idxi64_1_3_int64_t(
    const device half* src [[buffer(0)]],
    device half* out [[buffer(1)]],
    const constant int* src_shape [[buffer(2)]],
    const constant int64_t* src_strides [[buffer(3)]],
    const constant size_t& src_ndim [[buffer(4)]],
    const constant int* slice_sizes [[buffer(5)]],
    const constant int* axes [[buffer(6)]],
    const constant int* idx_shapes [[buffer(7)]],
    const constant int64_t* idx_strides [[buffer(8)]],
    const constant bool* idx_contigs [[buffer(9)]],
    const constant int& idx_ndim [[buffer(10)]],
    const device int64_t *idx0 [[buffer(20)]],
    uint3 index [[thread_position_in_grid]],
    uint3 grid_dim [[threads_per_grid]]) {
  Indices<int64_t, 1> idxs{
    { idx0 }, idx_shapes, idx_strides, idx_contigs, idx_ndim};

  return gather_impl<half, int64_t, 1, 3, int64_t>(
      src,
      out,
      src_shape,
      src_strides,
      src_ndim,
      slice_sizes,
      axes,
      idxs,
      index,
      grid_dim);
}

[[kernel]] void kernel_qwen_mtplx_gather_f16_idxi64_1_4_int(
    const device half* src [[buffer(0)]],
    device half* out [[buffer(1)]],
    const constant int* src_shape [[buffer(2)]],
    const constant int64_t* src_strides [[buffer(3)]],
    const constant size_t& src_ndim [[buffer(4)]],
    const constant int* slice_sizes [[buffer(5)]],
    const constant int* axes [[buffer(6)]],
    const constant int* idx_shapes [[buffer(7)]],
    const constant int64_t* idx_strides [[buffer(8)]],
    const constant bool* idx_contigs [[buffer(9)]],
    const constant int& idx_ndim [[buffer(10)]],
    const device int64_t *idx0 [[buffer(20)]],
    uint3 index [[thread_position_in_grid]],
    uint3 grid_dim [[threads_per_grid]]) {
  Indices<int64_t, 1> idxs{
    { idx0 }, idx_shapes, idx_strides, idx_contigs, idx_ndim};

  return gather_impl<half, int64_t, 1, 4, int>(
      src,
      out,
      src_shape,
      src_strides,
      src_ndim,
      slice_sizes,
      axes,
      idxs,
      index,
      grid_dim);
}

[[kernel]] void kernel_qwen_mtplx_gather_f16_idxi64_1_4_int64_t(
    const device half* src [[buffer(0)]],
    device half* out [[buffer(1)]],
    const constant int* src_shape [[buffer(2)]],
    const constant int64_t* src_strides [[buffer(3)]],
    const constant size_t& src_ndim [[buffer(4)]],
    const constant int* slice_sizes [[buffer(5)]],
    const constant int* axes [[buffer(6)]],
    const constant int* idx_shapes [[buffer(7)]],
    const constant int64_t* idx_strides [[buffer(8)]],
    const constant bool* idx_contigs [[buffer(9)]],
    const constant int& idx_ndim [[buffer(10)]],
    const device int64_t *idx0 [[buffer(20)]],
    uint3 index [[thread_position_in_grid]],
    uint3 grid_dim [[threads_per_grid]]) {
  Indices<int64_t, 1> idxs{
    { idx0 }, idx_shapes, idx_strides, idx_contigs, idx_ndim};

  return gather_impl<half, int64_t, 1, 4, int64_t>(
      src,
      out,
      src_shape,
      src_strides,
      src_ndim,
      slice_sizes,
      axes,
      idxs,
      index,
      grid_dim);
}

[[kernel]] void kernel_qwen_mtplx_gather_u32_idxi64_1_0_int(
    const device uint32_t* src [[buffer(0)]],
    device uint32_t* out [[buffer(1)]],
    const constant int* src_shape [[buffer(2)]],
    const constant int64_t* src_strides [[buffer(3)]],
    const constant size_t& src_ndim [[buffer(4)]],
    const constant int* slice_sizes [[buffer(5)]],
    const constant int* axes [[buffer(6)]],
    const constant int* idx_shapes [[buffer(7)]],
    const constant int64_t* idx_strides [[buffer(8)]],
    const constant bool* idx_contigs [[buffer(9)]],
    const constant int& idx_ndim [[buffer(10)]],
    const device int64_t *idx0 [[buffer(20)]],
    uint3 index [[thread_position_in_grid]],
    uint3 grid_dim [[threads_per_grid]]) {
  Indices<int64_t, 1> idxs{
    { idx0 }, idx_shapes, idx_strides, idx_contigs, idx_ndim};

  return gather_impl<uint32_t, int64_t, 1, 0, int>(
      src,
      out,
      src_shape,
      src_strides,
      src_ndim,
      slice_sizes,
      axes,
      idxs,
      index,
      grid_dim);
}

[[kernel]] void kernel_qwen_mtplx_gather_u32_idxi64_1_0_int64_t(
    const device uint32_t* src [[buffer(0)]],
    device uint32_t* out [[buffer(1)]],
    const constant int* src_shape [[buffer(2)]],
    const constant int64_t* src_strides [[buffer(3)]],
    const constant size_t& src_ndim [[buffer(4)]],
    const constant int* slice_sizes [[buffer(5)]],
    const constant int* axes [[buffer(6)]],
    const constant int* idx_shapes [[buffer(7)]],
    const constant int64_t* idx_strides [[buffer(8)]],
    const constant bool* idx_contigs [[buffer(9)]],
    const constant int& idx_ndim [[buffer(10)]],
    const device int64_t *idx0 [[buffer(20)]],
    uint3 index [[thread_position_in_grid]],
    uint3 grid_dim [[threads_per_grid]]) {
  Indices<int64_t, 1> idxs{
    { idx0 }, idx_shapes, idx_strides, idx_contigs, idx_ndim};

  return gather_impl<uint32_t, int64_t, 1, 0, int64_t>(
      src,
      out,
      src_shape,
      src_strides,
      src_ndim,
      slice_sizes,
      axes,
      idxs,
      index,
      grid_dim);
}

[[kernel]] void kernel_qwen_mtplx_gather_u32_idxi64_1_1_int(
    const device uint32_t* src [[buffer(0)]],
    device uint32_t* out [[buffer(1)]],
    const constant int* src_shape [[buffer(2)]],
    const constant int64_t* src_strides [[buffer(3)]],
    const constant size_t& src_ndim [[buffer(4)]],
    const constant int* slice_sizes [[buffer(5)]],
    const constant int* axes [[buffer(6)]],
    const constant int* idx_shapes [[buffer(7)]],
    const constant int64_t* idx_strides [[buffer(8)]],
    const constant bool* idx_contigs [[buffer(9)]],
    const constant int& idx_ndim [[buffer(10)]],
    const device int64_t *idx0 [[buffer(20)]],
    uint3 index [[thread_position_in_grid]],
    uint3 grid_dim [[threads_per_grid]]) {
  Indices<int64_t, 1> idxs{
    { idx0 }, idx_shapes, idx_strides, idx_contigs, idx_ndim};

  return gather_impl<uint32_t, int64_t, 1, 1, int>(
      src,
      out,
      src_shape,
      src_strides,
      src_ndim,
      slice_sizes,
      axes,
      idxs,
      index,
      grid_dim);
}

[[kernel]] void kernel_qwen_mtplx_gather_u32_idxi64_1_1_int64_t(
    const device uint32_t* src [[buffer(0)]],
    device uint32_t* out [[buffer(1)]],
    const constant int* src_shape [[buffer(2)]],
    const constant int64_t* src_strides [[buffer(3)]],
    const constant size_t& src_ndim [[buffer(4)]],
    const constant int* slice_sizes [[buffer(5)]],
    const constant int* axes [[buffer(6)]],
    const constant int* idx_shapes [[buffer(7)]],
    const constant int64_t* idx_strides [[buffer(8)]],
    const constant bool* idx_contigs [[buffer(9)]],
    const constant int& idx_ndim [[buffer(10)]],
    const device int64_t *idx0 [[buffer(20)]],
    uint3 index [[thread_position_in_grid]],
    uint3 grid_dim [[threads_per_grid]]) {
  Indices<int64_t, 1> idxs{
    { idx0 }, idx_shapes, idx_strides, idx_contigs, idx_ndim};

  return gather_impl<uint32_t, int64_t, 1, 1, int64_t>(
      src,
      out,
      src_shape,
      src_strides,
      src_ndim,
      slice_sizes,
      axes,
      idxs,
      index,
      grid_dim);
}

[[kernel]] void kernel_qwen_mtplx_gather_u32_idxi64_1_2_int(
    const device uint32_t* src [[buffer(0)]],
    device uint32_t* out [[buffer(1)]],
    const constant int* src_shape [[buffer(2)]],
    const constant int64_t* src_strides [[buffer(3)]],
    const constant size_t& src_ndim [[buffer(4)]],
    const constant int* slice_sizes [[buffer(5)]],
    const constant int* axes [[buffer(6)]],
    const constant int* idx_shapes [[buffer(7)]],
    const constant int64_t* idx_strides [[buffer(8)]],
    const constant bool* idx_contigs [[buffer(9)]],
    const constant int& idx_ndim [[buffer(10)]],
    const device int64_t *idx0 [[buffer(20)]],
    uint3 index [[thread_position_in_grid]],
    uint3 grid_dim [[threads_per_grid]]) {
  Indices<int64_t, 1> idxs{
    { idx0 }, idx_shapes, idx_strides, idx_contigs, idx_ndim};

  return gather_impl<uint32_t, int64_t, 1, 2, int>(
      src,
      out,
      src_shape,
      src_strides,
      src_ndim,
      slice_sizes,
      axes,
      idxs,
      index,
      grid_dim);
}

[[kernel]] void kernel_qwen_mtplx_gather_u32_idxi64_1_2_int64_t(
    const device uint32_t* src [[buffer(0)]],
    device uint32_t* out [[buffer(1)]],
    const constant int* src_shape [[buffer(2)]],
    const constant int64_t* src_strides [[buffer(3)]],
    const constant size_t& src_ndim [[buffer(4)]],
    const constant int* slice_sizes [[buffer(5)]],
    const constant int* axes [[buffer(6)]],
    const constant int* idx_shapes [[buffer(7)]],
    const constant int64_t* idx_strides [[buffer(8)]],
    const constant bool* idx_contigs [[buffer(9)]],
    const constant int& idx_ndim [[buffer(10)]],
    const device int64_t *idx0 [[buffer(20)]],
    uint3 index [[thread_position_in_grid]],
    uint3 grid_dim [[threads_per_grid]]) {
  Indices<int64_t, 1> idxs{
    { idx0 }, idx_shapes, idx_strides, idx_contigs, idx_ndim};

  return gather_impl<uint32_t, int64_t, 1, 2, int64_t>(
      src,
      out,
      src_shape,
      src_strides,
      src_ndim,
      slice_sizes,
      axes,
      idxs,
      index,
      grid_dim);
}

[[kernel]] void kernel_qwen_mtplx_gather_u32_idxi64_1_3_int(
    const device uint32_t* src [[buffer(0)]],
    device uint32_t* out [[buffer(1)]],
    const constant int* src_shape [[buffer(2)]],
    const constant int64_t* src_strides [[buffer(3)]],
    const constant size_t& src_ndim [[buffer(4)]],
    const constant int* slice_sizes [[buffer(5)]],
    const constant int* axes [[buffer(6)]],
    const constant int* idx_shapes [[buffer(7)]],
    const constant int64_t* idx_strides [[buffer(8)]],
    const constant bool* idx_contigs [[buffer(9)]],
    const constant int& idx_ndim [[buffer(10)]],
    const device int64_t *idx0 [[buffer(20)]],
    uint3 index [[thread_position_in_grid]],
    uint3 grid_dim [[threads_per_grid]]) {
  Indices<int64_t, 1> idxs{
    { idx0 }, idx_shapes, idx_strides, idx_contigs, idx_ndim};

  return gather_impl<uint32_t, int64_t, 1, 3, int>(
      src,
      out,
      src_shape,
      src_strides,
      src_ndim,
      slice_sizes,
      axes,
      idxs,
      index,
      grid_dim);
}

[[kernel]] void kernel_qwen_mtplx_gather_u32_idxi64_1_3_int64_t(
    const device uint32_t* src [[buffer(0)]],
    device uint32_t* out [[buffer(1)]],
    const constant int* src_shape [[buffer(2)]],
    const constant int64_t* src_strides [[buffer(3)]],
    const constant size_t& src_ndim [[buffer(4)]],
    const constant int* slice_sizes [[buffer(5)]],
    const constant int* axes [[buffer(6)]],
    const constant int* idx_shapes [[buffer(7)]],
    const constant int64_t* idx_strides [[buffer(8)]],
    const constant bool* idx_contigs [[buffer(9)]],
    const constant int& idx_ndim [[buffer(10)]],
    const device int64_t *idx0 [[buffer(20)]],
    uint3 index [[thread_position_in_grid]],
    uint3 grid_dim [[threads_per_grid]]) {
  Indices<int64_t, 1> idxs{
    { idx0 }, idx_shapes, idx_strides, idx_contigs, idx_ndim};

  return gather_impl<uint32_t, int64_t, 1, 3, int64_t>(
      src,
      out,
      src_shape,
      src_strides,
      src_ndim,
      slice_sizes,
      axes,
      idxs,
      index,
      grid_dim);
}

[[kernel]] void kernel_qwen_mtplx_gather_u32_idxi64_1_4_int(
    const device uint32_t* src [[buffer(0)]],
    device uint32_t* out [[buffer(1)]],
    const constant int* src_shape [[buffer(2)]],
    const constant int64_t* src_strides [[buffer(3)]],
    const constant size_t& src_ndim [[buffer(4)]],
    const constant int* slice_sizes [[buffer(5)]],
    const constant int* axes [[buffer(6)]],
    const constant int* idx_shapes [[buffer(7)]],
    const constant int64_t* idx_strides [[buffer(8)]],
    const constant bool* idx_contigs [[buffer(9)]],
    const constant int& idx_ndim [[buffer(10)]],
    const device int64_t *idx0 [[buffer(20)]],
    uint3 index [[thread_position_in_grid]],
    uint3 grid_dim [[threads_per_grid]]) {
  Indices<int64_t, 1> idxs{
    { idx0 }, idx_shapes, idx_strides, idx_contigs, idx_ndim};

  return gather_impl<uint32_t, int64_t, 1, 4, int>(
      src,
      out,
      src_shape,
      src_strides,
      src_ndim,
      slice_sizes,
      axes,
      idxs,
      index,
      grid_dim);
}

[[kernel]] void kernel_qwen_mtplx_gather_u32_idxi64_1_4_int64_t(
    const device uint32_t* src [[buffer(0)]],
    device uint32_t* out [[buffer(1)]],
    const constant int* src_shape [[buffer(2)]],
    const constant int64_t* src_strides [[buffer(3)]],
    const constant size_t& src_ndim [[buffer(4)]],
    const constant int* slice_sizes [[buffer(5)]],
    const constant int* axes [[buffer(6)]],
    const constant int* idx_shapes [[buffer(7)]],
    const constant int64_t* idx_strides [[buffer(8)]],
    const constant bool* idx_contigs [[buffer(9)]],
    const constant int& idx_ndim [[buffer(10)]],
    const device int64_t *idx0 [[buffer(20)]],
    uint3 index [[thread_position_in_grid]],
    uint3 grid_dim [[threads_per_grid]]) {
  Indices<int64_t, 1> idxs{
    { idx0 }, idx_shapes, idx_strides, idx_contigs, idx_ndim};

  return gather_impl<uint32_t, int64_t, 1, 4, int64_t>(
      src,
      out,
      src_shape,
      src_strides,
      src_ndim,
      slice_sizes,
      axes,
      idxs,
      index,
      grid_dim);
}

[[kernel]] void kernel_qwen_mtplx_gather_f32_idxi64_1_0_int(
    const device float* src [[buffer(0)]],
    device float* out [[buffer(1)]],
    const constant int* src_shape [[buffer(2)]],
    const constant int64_t* src_strides [[buffer(3)]],
    const constant size_t& src_ndim [[buffer(4)]],
    const constant int* slice_sizes [[buffer(5)]],
    const constant int* axes [[buffer(6)]],
    const constant int* idx_shapes [[buffer(7)]],
    const constant int64_t* idx_strides [[buffer(8)]],
    const constant bool* idx_contigs [[buffer(9)]],
    const constant int& idx_ndim [[buffer(10)]],
    const device int64_t *idx0 [[buffer(20)]],
    uint3 index [[thread_position_in_grid]],
    uint3 grid_dim [[threads_per_grid]]) {
  Indices<int64_t, 1> idxs{
    { idx0 }, idx_shapes, idx_strides, idx_contigs, idx_ndim};

  return gather_impl<float, int64_t, 1, 0, int>(
      src,
      out,
      src_shape,
      src_strides,
      src_ndim,
      slice_sizes,
      axes,
      idxs,
      index,
      grid_dim);
}

[[kernel]] void kernel_qwen_mtplx_gather_f32_idxi64_1_0_int64_t(
    const device float* src [[buffer(0)]],
    device float* out [[buffer(1)]],
    const constant int* src_shape [[buffer(2)]],
    const constant int64_t* src_strides [[buffer(3)]],
    const constant size_t& src_ndim [[buffer(4)]],
    const constant int* slice_sizes [[buffer(5)]],
    const constant int* axes [[buffer(6)]],
    const constant int* idx_shapes [[buffer(7)]],
    const constant int64_t* idx_strides [[buffer(8)]],
    const constant bool* idx_contigs [[buffer(9)]],
    const constant int& idx_ndim [[buffer(10)]],
    const device int64_t *idx0 [[buffer(20)]],
    uint3 index [[thread_position_in_grid]],
    uint3 grid_dim [[threads_per_grid]]) {
  Indices<int64_t, 1> idxs{
    { idx0 }, idx_shapes, idx_strides, idx_contigs, idx_ndim};

  return gather_impl<float, int64_t, 1, 0, int64_t>(
      src,
      out,
      src_shape,
      src_strides,
      src_ndim,
      slice_sizes,
      axes,
      idxs,
      index,
      grid_dim);
}

[[kernel]] void kernel_qwen_mtplx_gather_f32_idxi64_1_1_int(
    const device float* src [[buffer(0)]],
    device float* out [[buffer(1)]],
    const constant int* src_shape [[buffer(2)]],
    const constant int64_t* src_strides [[buffer(3)]],
    const constant size_t& src_ndim [[buffer(4)]],
    const constant int* slice_sizes [[buffer(5)]],
    const constant int* axes [[buffer(6)]],
    const constant int* idx_shapes [[buffer(7)]],
    const constant int64_t* idx_strides [[buffer(8)]],
    const constant bool* idx_contigs [[buffer(9)]],
    const constant int& idx_ndim [[buffer(10)]],
    const device int64_t *idx0 [[buffer(20)]],
    uint3 index [[thread_position_in_grid]],
    uint3 grid_dim [[threads_per_grid]]) {
  Indices<int64_t, 1> idxs{
    { idx0 }, idx_shapes, idx_strides, idx_contigs, idx_ndim};

  return gather_impl<float, int64_t, 1, 1, int>(
      src,
      out,
      src_shape,
      src_strides,
      src_ndim,
      slice_sizes,
      axes,
      idxs,
      index,
      grid_dim);
}

[[kernel]] void kernel_qwen_mtplx_gather_f32_idxi64_1_1_int64_t(
    const device float* src [[buffer(0)]],
    device float* out [[buffer(1)]],
    const constant int* src_shape [[buffer(2)]],
    const constant int64_t* src_strides [[buffer(3)]],
    const constant size_t& src_ndim [[buffer(4)]],
    const constant int* slice_sizes [[buffer(5)]],
    const constant int* axes [[buffer(6)]],
    const constant int* idx_shapes [[buffer(7)]],
    const constant int64_t* idx_strides [[buffer(8)]],
    const constant bool* idx_contigs [[buffer(9)]],
    const constant int& idx_ndim [[buffer(10)]],
    const device int64_t *idx0 [[buffer(20)]],
    uint3 index [[thread_position_in_grid]],
    uint3 grid_dim [[threads_per_grid]]) {
  Indices<int64_t, 1> idxs{
    { idx0 }, idx_shapes, idx_strides, idx_contigs, idx_ndim};

  return gather_impl<float, int64_t, 1, 1, int64_t>(
      src,
      out,
      src_shape,
      src_strides,
      src_ndim,
      slice_sizes,
      axes,
      idxs,
      index,
      grid_dim);
}

[[kernel]] void kernel_qwen_mtplx_gather_f32_idxi64_1_2_int(
    const device float* src [[buffer(0)]],
    device float* out [[buffer(1)]],
    const constant int* src_shape [[buffer(2)]],
    const constant int64_t* src_strides [[buffer(3)]],
    const constant size_t& src_ndim [[buffer(4)]],
    const constant int* slice_sizes [[buffer(5)]],
    const constant int* axes [[buffer(6)]],
    const constant int* idx_shapes [[buffer(7)]],
    const constant int64_t* idx_strides [[buffer(8)]],
    const constant bool* idx_contigs [[buffer(9)]],
    const constant int& idx_ndim [[buffer(10)]],
    const device int64_t *idx0 [[buffer(20)]],
    uint3 index [[thread_position_in_grid]],
    uint3 grid_dim [[threads_per_grid]]) {
  Indices<int64_t, 1> idxs{
    { idx0 }, idx_shapes, idx_strides, idx_contigs, idx_ndim};

  return gather_impl<float, int64_t, 1, 2, int>(
      src,
      out,
      src_shape,
      src_strides,
      src_ndim,
      slice_sizes,
      axes,
      idxs,
      index,
      grid_dim);
}

[[kernel]] void kernel_qwen_mtplx_gather_f32_idxi64_1_2_int64_t(
    const device float* src [[buffer(0)]],
    device float* out [[buffer(1)]],
    const constant int* src_shape [[buffer(2)]],
    const constant int64_t* src_strides [[buffer(3)]],
    const constant size_t& src_ndim [[buffer(4)]],
    const constant int* slice_sizes [[buffer(5)]],
    const constant int* axes [[buffer(6)]],
    const constant int* idx_shapes [[buffer(7)]],
    const constant int64_t* idx_strides [[buffer(8)]],
    const constant bool* idx_contigs [[buffer(9)]],
    const constant int& idx_ndim [[buffer(10)]],
    const device int64_t *idx0 [[buffer(20)]],
    uint3 index [[thread_position_in_grid]],
    uint3 grid_dim [[threads_per_grid]]) {
  Indices<int64_t, 1> idxs{
    { idx0 }, idx_shapes, idx_strides, idx_contigs, idx_ndim};

  return gather_impl<float, int64_t, 1, 2, int64_t>(
      src,
      out,
      src_shape,
      src_strides,
      src_ndim,
      slice_sizes,
      axes,
      idxs,
      index,
      grid_dim);
}

[[kernel]] void kernel_qwen_mtplx_gather_f32_idxi64_1_3_int(
    const device float* src [[buffer(0)]],
    device float* out [[buffer(1)]],
    const constant int* src_shape [[buffer(2)]],
    const constant int64_t* src_strides [[buffer(3)]],
    const constant size_t& src_ndim [[buffer(4)]],
    const constant int* slice_sizes [[buffer(5)]],
    const constant int* axes [[buffer(6)]],
    const constant int* idx_shapes [[buffer(7)]],
    const constant int64_t* idx_strides [[buffer(8)]],
    const constant bool* idx_contigs [[buffer(9)]],
    const constant int& idx_ndim [[buffer(10)]],
    const device int64_t *idx0 [[buffer(20)]],
    uint3 index [[thread_position_in_grid]],
    uint3 grid_dim [[threads_per_grid]]) {
  Indices<int64_t, 1> idxs{
    { idx0 }, idx_shapes, idx_strides, idx_contigs, idx_ndim};

  return gather_impl<float, int64_t, 1, 3, int>(
      src,
      out,
      src_shape,
      src_strides,
      src_ndim,
      slice_sizes,
      axes,
      idxs,
      index,
      grid_dim);
}

[[kernel]] void kernel_qwen_mtplx_gather_f32_idxi64_1_3_int64_t(
    const device float* src [[buffer(0)]],
    device float* out [[buffer(1)]],
    const constant int* src_shape [[buffer(2)]],
    const constant int64_t* src_strides [[buffer(3)]],
    const constant size_t& src_ndim [[buffer(4)]],
    const constant int* slice_sizes [[buffer(5)]],
    const constant int* axes [[buffer(6)]],
    const constant int* idx_shapes [[buffer(7)]],
    const constant int64_t* idx_strides [[buffer(8)]],
    const constant bool* idx_contigs [[buffer(9)]],
    const constant int& idx_ndim [[buffer(10)]],
    const device int64_t *idx0 [[buffer(20)]],
    uint3 index [[thread_position_in_grid]],
    uint3 grid_dim [[threads_per_grid]]) {
  Indices<int64_t, 1> idxs{
    { idx0 }, idx_shapes, idx_strides, idx_contigs, idx_ndim};

  return gather_impl<float, int64_t, 1, 3, int64_t>(
      src,
      out,
      src_shape,
      src_strides,
      src_ndim,
      slice_sizes,
      axes,
      idxs,
      index,
      grid_dim);
}

[[kernel]] void kernel_qwen_mtplx_gather_f32_idxi64_1_4_int(
    const device float* src [[buffer(0)]],
    device float* out [[buffer(1)]],
    const constant int* src_shape [[buffer(2)]],
    const constant int64_t* src_strides [[buffer(3)]],
    const constant size_t& src_ndim [[buffer(4)]],
    const constant int* slice_sizes [[buffer(5)]],
    const constant int* axes [[buffer(6)]],
    const constant int* idx_shapes [[buffer(7)]],
    const constant int64_t* idx_strides [[buffer(8)]],
    const constant bool* idx_contigs [[buffer(9)]],
    const constant int& idx_ndim [[buffer(10)]],
    const device int64_t *idx0 [[buffer(20)]],
    uint3 index [[thread_position_in_grid]],
    uint3 grid_dim [[threads_per_grid]]) {
  Indices<int64_t, 1> idxs{
    { idx0 }, idx_shapes, idx_strides, idx_contigs, idx_ndim};

  return gather_impl<float, int64_t, 1, 4, int>(
      src,
      out,
      src_shape,
      src_strides,
      src_ndim,
      slice_sizes,
      axes,
      idxs,
      index,
      grid_dim);
}

[[kernel]] void kernel_qwen_mtplx_gather_f32_idxi64_1_4_int64_t(
    const device float* src [[buffer(0)]],
    device float* out [[buffer(1)]],
    const constant int* src_shape [[buffer(2)]],
    const constant int64_t* src_strides [[buffer(3)]],
    const constant size_t& src_ndim [[buffer(4)]],
    const constant int* slice_sizes [[buffer(5)]],
    const constant int* axes [[buffer(6)]],
    const constant int* idx_shapes [[buffer(7)]],
    const constant int64_t* idx_strides [[buffer(8)]],
    const constant bool* idx_contigs [[buffer(9)]],
    const constant int& idx_ndim [[buffer(10)]],
    const device int64_t *idx0 [[buffer(20)]],
    uint3 index [[thread_position_in_grid]],
    uint3 grid_dim [[threads_per_grid]]) {
  Indices<int64_t, 1> idxs{
    { idx0 }, idx_shapes, idx_strides, idx_contigs, idx_ndim};

  return gather_impl<float, int64_t, 1, 4, int64_t>(
      src,
      out,
      src_shape,
      src_strides,
      src_ndim,
      slice_sizes,
      axes,
      idxs,
      index,
      grid_dim);
}

[[kernel]] void kernel_qwen_mtplx_gather_bf16_idxi32_1_0_int(
    const device bfloat* src [[buffer(0)]],
    device bfloat* out [[buffer(1)]],
    const constant int* src_shape [[buffer(2)]],
    const constant int64_t* src_strides [[buffer(3)]],
    const constant size_t& src_ndim [[buffer(4)]],
    const constant int* slice_sizes [[buffer(5)]],
    const constant int* axes [[buffer(6)]],
    const constant int* idx_shapes [[buffer(7)]],
    const constant int64_t* idx_strides [[buffer(8)]],
    const constant bool* idx_contigs [[buffer(9)]],
    const constant int& idx_ndim [[buffer(10)]],
    const device int32_t *idx0 [[buffer(20)]],
    uint3 index [[thread_position_in_grid]],
    uint3 grid_dim [[threads_per_grid]]) {
  Indices<int32_t, 1> idxs{
    { idx0 }, idx_shapes, idx_strides, idx_contigs, idx_ndim};

  return gather_impl<bfloat, int32_t, 1, 0, int>(
      src,
      out,
      src_shape,
      src_strides,
      src_ndim,
      slice_sizes,
      axes,
      idxs,
      index,
      grid_dim);
}

[[kernel]] void kernel_qwen_mtplx_gather_bf16_idxi32_1_0_int64_t(
    const device bfloat* src [[buffer(0)]],
    device bfloat* out [[buffer(1)]],
    const constant int* src_shape [[buffer(2)]],
    const constant int64_t* src_strides [[buffer(3)]],
    const constant size_t& src_ndim [[buffer(4)]],
    const constant int* slice_sizes [[buffer(5)]],
    const constant int* axes [[buffer(6)]],
    const constant int* idx_shapes [[buffer(7)]],
    const constant int64_t* idx_strides [[buffer(8)]],
    const constant bool* idx_contigs [[buffer(9)]],
    const constant int& idx_ndim [[buffer(10)]],
    const device int32_t *idx0 [[buffer(20)]],
    uint3 index [[thread_position_in_grid]],
    uint3 grid_dim [[threads_per_grid]]) {
  Indices<int32_t, 1> idxs{
    { idx0 }, idx_shapes, idx_strides, idx_contigs, idx_ndim};

  return gather_impl<bfloat, int32_t, 1, 0, int64_t>(
      src,
      out,
      src_shape,
      src_strides,
      src_ndim,
      slice_sizes,
      axes,
      idxs,
      index,
      grid_dim);
}

[[kernel]] void kernel_qwen_mtplx_gather_bf16_idxi32_1_1_int(
    const device bfloat* src [[buffer(0)]],
    device bfloat* out [[buffer(1)]],
    const constant int* src_shape [[buffer(2)]],
    const constant int64_t* src_strides [[buffer(3)]],
    const constant size_t& src_ndim [[buffer(4)]],
    const constant int* slice_sizes [[buffer(5)]],
    const constant int* axes [[buffer(6)]],
    const constant int* idx_shapes [[buffer(7)]],
    const constant int64_t* idx_strides [[buffer(8)]],
    const constant bool* idx_contigs [[buffer(9)]],
    const constant int& idx_ndim [[buffer(10)]],
    const device int32_t *idx0 [[buffer(20)]],
    uint3 index [[thread_position_in_grid]],
    uint3 grid_dim [[threads_per_grid]]) {
  Indices<int32_t, 1> idxs{
    { idx0 }, idx_shapes, idx_strides, idx_contigs, idx_ndim};

  return gather_impl<bfloat, int32_t, 1, 1, int>(
      src,
      out,
      src_shape,
      src_strides,
      src_ndim,
      slice_sizes,
      axes,
      idxs,
      index,
      grid_dim);
}

[[kernel]] void kernel_qwen_mtplx_gather_bf16_idxi32_1_1_int64_t(
    const device bfloat* src [[buffer(0)]],
    device bfloat* out [[buffer(1)]],
    const constant int* src_shape [[buffer(2)]],
    const constant int64_t* src_strides [[buffer(3)]],
    const constant size_t& src_ndim [[buffer(4)]],
    const constant int* slice_sizes [[buffer(5)]],
    const constant int* axes [[buffer(6)]],
    const constant int* idx_shapes [[buffer(7)]],
    const constant int64_t* idx_strides [[buffer(8)]],
    const constant bool* idx_contigs [[buffer(9)]],
    const constant int& idx_ndim [[buffer(10)]],
    const device int32_t *idx0 [[buffer(20)]],
    uint3 index [[thread_position_in_grid]],
    uint3 grid_dim [[threads_per_grid]]) {
  Indices<int32_t, 1> idxs{
    { idx0 }, idx_shapes, idx_strides, idx_contigs, idx_ndim};

  return gather_impl<bfloat, int32_t, 1, 1, int64_t>(
      src,
      out,
      src_shape,
      src_strides,
      src_ndim,
      slice_sizes,
      axes,
      idxs,
      index,
      grid_dim);
}

[[kernel]] void kernel_qwen_mtplx_gather_bf16_idxi32_1_2_int(
    const device bfloat* src [[buffer(0)]],
    device bfloat* out [[buffer(1)]],
    const constant int* src_shape [[buffer(2)]],
    const constant int64_t* src_strides [[buffer(3)]],
    const constant size_t& src_ndim [[buffer(4)]],
    const constant int* slice_sizes [[buffer(5)]],
    const constant int* axes [[buffer(6)]],
    const constant int* idx_shapes [[buffer(7)]],
    const constant int64_t* idx_strides [[buffer(8)]],
    const constant bool* idx_contigs [[buffer(9)]],
    const constant int& idx_ndim [[buffer(10)]],
    const device int32_t *idx0 [[buffer(20)]],
    uint3 index [[thread_position_in_grid]],
    uint3 grid_dim [[threads_per_grid]]) {
  Indices<int32_t, 1> idxs{
    { idx0 }, idx_shapes, idx_strides, idx_contigs, idx_ndim};

  return gather_impl<bfloat, int32_t, 1, 2, int>(
      src,
      out,
      src_shape,
      src_strides,
      src_ndim,
      slice_sizes,
      axes,
      idxs,
      index,
      grid_dim);
}

[[kernel]] void kernel_qwen_mtplx_gather_bf16_idxi32_1_2_int64_t(
    const device bfloat* src [[buffer(0)]],
    device bfloat* out [[buffer(1)]],
    const constant int* src_shape [[buffer(2)]],
    const constant int64_t* src_strides [[buffer(3)]],
    const constant size_t& src_ndim [[buffer(4)]],
    const constant int* slice_sizes [[buffer(5)]],
    const constant int* axes [[buffer(6)]],
    const constant int* idx_shapes [[buffer(7)]],
    const constant int64_t* idx_strides [[buffer(8)]],
    const constant bool* idx_contigs [[buffer(9)]],
    const constant int& idx_ndim [[buffer(10)]],
    const device int32_t *idx0 [[buffer(20)]],
    uint3 index [[thread_position_in_grid]],
    uint3 grid_dim [[threads_per_grid]]) {
  Indices<int32_t, 1> idxs{
    { idx0 }, idx_shapes, idx_strides, idx_contigs, idx_ndim};

  return gather_impl<bfloat, int32_t, 1, 2, int64_t>(
      src,
      out,
      src_shape,
      src_strides,
      src_ndim,
      slice_sizes,
      axes,
      idxs,
      index,
      grid_dim);
}

[[kernel]] void kernel_qwen_mtplx_gather_bf16_idxi32_1_3_int(
    const device bfloat* src [[buffer(0)]],
    device bfloat* out [[buffer(1)]],
    const constant int* src_shape [[buffer(2)]],
    const constant int64_t* src_strides [[buffer(3)]],
    const constant size_t& src_ndim [[buffer(4)]],
    const constant int* slice_sizes [[buffer(5)]],
    const constant int* axes [[buffer(6)]],
    const constant int* idx_shapes [[buffer(7)]],
    const constant int64_t* idx_strides [[buffer(8)]],
    const constant bool* idx_contigs [[buffer(9)]],
    const constant int& idx_ndim [[buffer(10)]],
    const device int32_t *idx0 [[buffer(20)]],
    uint3 index [[thread_position_in_grid]],
    uint3 grid_dim [[threads_per_grid]]) {
  Indices<int32_t, 1> idxs{
    { idx0 }, idx_shapes, idx_strides, idx_contigs, idx_ndim};

  return gather_impl<bfloat, int32_t, 1, 3, int>(
      src,
      out,
      src_shape,
      src_strides,
      src_ndim,
      slice_sizes,
      axes,
      idxs,
      index,
      grid_dim);
}

[[kernel]] void kernel_qwen_mtplx_gather_bf16_idxi32_1_3_int64_t(
    const device bfloat* src [[buffer(0)]],
    device bfloat* out [[buffer(1)]],
    const constant int* src_shape [[buffer(2)]],
    const constant int64_t* src_strides [[buffer(3)]],
    const constant size_t& src_ndim [[buffer(4)]],
    const constant int* slice_sizes [[buffer(5)]],
    const constant int* axes [[buffer(6)]],
    const constant int* idx_shapes [[buffer(7)]],
    const constant int64_t* idx_strides [[buffer(8)]],
    const constant bool* idx_contigs [[buffer(9)]],
    const constant int& idx_ndim [[buffer(10)]],
    const device int32_t *idx0 [[buffer(20)]],
    uint3 index [[thread_position_in_grid]],
    uint3 grid_dim [[threads_per_grid]]) {
  Indices<int32_t, 1> idxs{
    { idx0 }, idx_shapes, idx_strides, idx_contigs, idx_ndim};

  return gather_impl<bfloat, int32_t, 1, 3, int64_t>(
      src,
      out,
      src_shape,
      src_strides,
      src_ndim,
      slice_sizes,
      axes,
      idxs,
      index,
      grid_dim);
}

[[kernel]] void kernel_qwen_mtplx_gather_bf16_idxi32_1_4_int(
    const device bfloat* src [[buffer(0)]],
    device bfloat* out [[buffer(1)]],
    const constant int* src_shape [[buffer(2)]],
    const constant int64_t* src_strides [[buffer(3)]],
    const constant size_t& src_ndim [[buffer(4)]],
    const constant int* slice_sizes [[buffer(5)]],
    const constant int* axes [[buffer(6)]],
    const constant int* idx_shapes [[buffer(7)]],
    const constant int64_t* idx_strides [[buffer(8)]],
    const constant bool* idx_contigs [[buffer(9)]],
    const constant int& idx_ndim [[buffer(10)]],
    const device int32_t *idx0 [[buffer(20)]],
    uint3 index [[thread_position_in_grid]],
    uint3 grid_dim [[threads_per_grid]]) {
  Indices<int32_t, 1> idxs{
    { idx0 }, idx_shapes, idx_strides, idx_contigs, idx_ndim};

  return gather_impl<bfloat, int32_t, 1, 4, int>(
      src,
      out,
      src_shape,
      src_strides,
      src_ndim,
      slice_sizes,
      axes,
      idxs,
      index,
      grid_dim);
}

[[kernel]] void kernel_qwen_mtplx_gather_bf16_idxi32_1_4_int64_t(
    const device bfloat* src [[buffer(0)]],
    device bfloat* out [[buffer(1)]],
    const constant int* src_shape [[buffer(2)]],
    const constant int64_t* src_strides [[buffer(3)]],
    const constant size_t& src_ndim [[buffer(4)]],
    const constant int* slice_sizes [[buffer(5)]],
    const constant int* axes [[buffer(6)]],
    const constant int* idx_shapes [[buffer(7)]],
    const constant int64_t* idx_strides [[buffer(8)]],
    const constant bool* idx_contigs [[buffer(9)]],
    const constant int& idx_ndim [[buffer(10)]],
    const device int32_t *idx0 [[buffer(20)]],
    uint3 index [[thread_position_in_grid]],
    uint3 grid_dim [[threads_per_grid]]) {
  Indices<int32_t, 1> idxs{
    { idx0 }, idx_shapes, idx_strides, idx_contigs, idx_ndim};

  return gather_impl<bfloat, int32_t, 1, 4, int64_t>(
      src,
      out,
      src_shape,
      src_strides,
      src_ndim,
      slice_sizes,
      axes,
      idxs,
      index,
      grid_dim);
}

[[kernel]] void kernel_qwen_mtplx_gather_f16_idxi32_1_0_int(
    const device half* src [[buffer(0)]],
    device half* out [[buffer(1)]],
    const constant int* src_shape [[buffer(2)]],
    const constant int64_t* src_strides [[buffer(3)]],
    const constant size_t& src_ndim [[buffer(4)]],
    const constant int* slice_sizes [[buffer(5)]],
    const constant int* axes [[buffer(6)]],
    const constant int* idx_shapes [[buffer(7)]],
    const constant int64_t* idx_strides [[buffer(8)]],
    const constant bool* idx_contigs [[buffer(9)]],
    const constant int& idx_ndim [[buffer(10)]],
    const device int32_t *idx0 [[buffer(20)]],
    uint3 index [[thread_position_in_grid]],
    uint3 grid_dim [[threads_per_grid]]) {
  Indices<int32_t, 1> idxs{
    { idx0 }, idx_shapes, idx_strides, idx_contigs, idx_ndim};

  return gather_impl<half, int32_t, 1, 0, int>(
      src,
      out,
      src_shape,
      src_strides,
      src_ndim,
      slice_sizes,
      axes,
      idxs,
      index,
      grid_dim);
}

[[kernel]] void kernel_qwen_mtplx_gather_f16_idxi32_1_0_int64_t(
    const device half* src [[buffer(0)]],
    device half* out [[buffer(1)]],
    const constant int* src_shape [[buffer(2)]],
    const constant int64_t* src_strides [[buffer(3)]],
    const constant size_t& src_ndim [[buffer(4)]],
    const constant int* slice_sizes [[buffer(5)]],
    const constant int* axes [[buffer(6)]],
    const constant int* idx_shapes [[buffer(7)]],
    const constant int64_t* idx_strides [[buffer(8)]],
    const constant bool* idx_contigs [[buffer(9)]],
    const constant int& idx_ndim [[buffer(10)]],
    const device int32_t *idx0 [[buffer(20)]],
    uint3 index [[thread_position_in_grid]],
    uint3 grid_dim [[threads_per_grid]]) {
  Indices<int32_t, 1> idxs{
    { idx0 }, idx_shapes, idx_strides, idx_contigs, idx_ndim};

  return gather_impl<half, int32_t, 1, 0, int64_t>(
      src,
      out,
      src_shape,
      src_strides,
      src_ndim,
      slice_sizes,
      axes,
      idxs,
      index,
      grid_dim);
}

[[kernel]] void kernel_qwen_mtplx_gather_f16_idxi32_1_1_int(
    const device half* src [[buffer(0)]],
    device half* out [[buffer(1)]],
    const constant int* src_shape [[buffer(2)]],
    const constant int64_t* src_strides [[buffer(3)]],
    const constant size_t& src_ndim [[buffer(4)]],
    const constant int* slice_sizes [[buffer(5)]],
    const constant int* axes [[buffer(6)]],
    const constant int* idx_shapes [[buffer(7)]],
    const constant int64_t* idx_strides [[buffer(8)]],
    const constant bool* idx_contigs [[buffer(9)]],
    const constant int& idx_ndim [[buffer(10)]],
    const device int32_t *idx0 [[buffer(20)]],
    uint3 index [[thread_position_in_grid]],
    uint3 grid_dim [[threads_per_grid]]) {
  Indices<int32_t, 1> idxs{
    { idx0 }, idx_shapes, idx_strides, idx_contigs, idx_ndim};

  return gather_impl<half, int32_t, 1, 1, int>(
      src,
      out,
      src_shape,
      src_strides,
      src_ndim,
      slice_sizes,
      axes,
      idxs,
      index,
      grid_dim);
}

[[kernel]] void kernel_qwen_mtplx_gather_f16_idxi32_1_1_int64_t(
    const device half* src [[buffer(0)]],
    device half* out [[buffer(1)]],
    const constant int* src_shape [[buffer(2)]],
    const constant int64_t* src_strides [[buffer(3)]],
    const constant size_t& src_ndim [[buffer(4)]],
    const constant int* slice_sizes [[buffer(5)]],
    const constant int* axes [[buffer(6)]],
    const constant int* idx_shapes [[buffer(7)]],
    const constant int64_t* idx_strides [[buffer(8)]],
    const constant bool* idx_contigs [[buffer(9)]],
    const constant int& idx_ndim [[buffer(10)]],
    const device int32_t *idx0 [[buffer(20)]],
    uint3 index [[thread_position_in_grid]],
    uint3 grid_dim [[threads_per_grid]]) {
  Indices<int32_t, 1> idxs{
    { idx0 }, idx_shapes, idx_strides, idx_contigs, idx_ndim};

  return gather_impl<half, int32_t, 1, 1, int64_t>(
      src,
      out,
      src_shape,
      src_strides,
      src_ndim,
      slice_sizes,
      axes,
      idxs,
      index,
      grid_dim);
}

[[kernel]] void kernel_qwen_mtplx_gather_f16_idxi32_1_2_int(
    const device half* src [[buffer(0)]],
    device half* out [[buffer(1)]],
    const constant int* src_shape [[buffer(2)]],
    const constant int64_t* src_strides [[buffer(3)]],
    const constant size_t& src_ndim [[buffer(4)]],
    const constant int* slice_sizes [[buffer(5)]],
    const constant int* axes [[buffer(6)]],
    const constant int* idx_shapes [[buffer(7)]],
    const constant int64_t* idx_strides [[buffer(8)]],
    const constant bool* idx_contigs [[buffer(9)]],
    const constant int& idx_ndim [[buffer(10)]],
    const device int32_t *idx0 [[buffer(20)]],
    uint3 index [[thread_position_in_grid]],
    uint3 grid_dim [[threads_per_grid]]) {
  Indices<int32_t, 1> idxs{
    { idx0 }, idx_shapes, idx_strides, idx_contigs, idx_ndim};

  return gather_impl<half, int32_t, 1, 2, int>(
      src,
      out,
      src_shape,
      src_strides,
      src_ndim,
      slice_sizes,
      axes,
      idxs,
      index,
      grid_dim);
}

[[kernel]] void kernel_qwen_mtplx_gather_f16_idxi32_1_2_int64_t(
    const device half* src [[buffer(0)]],
    device half* out [[buffer(1)]],
    const constant int* src_shape [[buffer(2)]],
    const constant int64_t* src_strides [[buffer(3)]],
    const constant size_t& src_ndim [[buffer(4)]],
    const constant int* slice_sizes [[buffer(5)]],
    const constant int* axes [[buffer(6)]],
    const constant int* idx_shapes [[buffer(7)]],
    const constant int64_t* idx_strides [[buffer(8)]],
    const constant bool* idx_contigs [[buffer(9)]],
    const constant int& idx_ndim [[buffer(10)]],
    const device int32_t *idx0 [[buffer(20)]],
    uint3 index [[thread_position_in_grid]],
    uint3 grid_dim [[threads_per_grid]]) {
  Indices<int32_t, 1> idxs{
    { idx0 }, idx_shapes, idx_strides, idx_contigs, idx_ndim};

  return gather_impl<half, int32_t, 1, 2, int64_t>(
      src,
      out,
      src_shape,
      src_strides,
      src_ndim,
      slice_sizes,
      axes,
      idxs,
      index,
      grid_dim);
}

[[kernel]] void kernel_qwen_mtplx_gather_f16_idxi32_1_3_int(
    const device half* src [[buffer(0)]],
    device half* out [[buffer(1)]],
    const constant int* src_shape [[buffer(2)]],
    const constant int64_t* src_strides [[buffer(3)]],
    const constant size_t& src_ndim [[buffer(4)]],
    const constant int* slice_sizes [[buffer(5)]],
    const constant int* axes [[buffer(6)]],
    const constant int* idx_shapes [[buffer(7)]],
    const constant int64_t* idx_strides [[buffer(8)]],
    const constant bool* idx_contigs [[buffer(9)]],
    const constant int& idx_ndim [[buffer(10)]],
    const device int32_t *idx0 [[buffer(20)]],
    uint3 index [[thread_position_in_grid]],
    uint3 grid_dim [[threads_per_grid]]) {
  Indices<int32_t, 1> idxs{
    { idx0 }, idx_shapes, idx_strides, idx_contigs, idx_ndim};

  return gather_impl<half, int32_t, 1, 3, int>(
      src,
      out,
      src_shape,
      src_strides,
      src_ndim,
      slice_sizes,
      axes,
      idxs,
      index,
      grid_dim);
}

[[kernel]] void kernel_qwen_mtplx_gather_f16_idxi32_1_3_int64_t(
    const device half* src [[buffer(0)]],
    device half* out [[buffer(1)]],
    const constant int* src_shape [[buffer(2)]],
    const constant int64_t* src_strides [[buffer(3)]],
    const constant size_t& src_ndim [[buffer(4)]],
    const constant int* slice_sizes [[buffer(5)]],
    const constant int* axes [[buffer(6)]],
    const constant int* idx_shapes [[buffer(7)]],
    const constant int64_t* idx_strides [[buffer(8)]],
    const constant bool* idx_contigs [[buffer(9)]],
    const constant int& idx_ndim [[buffer(10)]],
    const device int32_t *idx0 [[buffer(20)]],
    uint3 index [[thread_position_in_grid]],
    uint3 grid_dim [[threads_per_grid]]) {
  Indices<int32_t, 1> idxs{
    { idx0 }, idx_shapes, idx_strides, idx_contigs, idx_ndim};

  return gather_impl<half, int32_t, 1, 3, int64_t>(
      src,
      out,
      src_shape,
      src_strides,
      src_ndim,
      slice_sizes,
      axes,
      idxs,
      index,
      grid_dim);
}

[[kernel]] void kernel_qwen_mtplx_gather_f16_idxi32_1_4_int(
    const device half* src [[buffer(0)]],
    device half* out [[buffer(1)]],
    const constant int* src_shape [[buffer(2)]],
    const constant int64_t* src_strides [[buffer(3)]],
    const constant size_t& src_ndim [[buffer(4)]],
    const constant int* slice_sizes [[buffer(5)]],
    const constant int* axes [[buffer(6)]],
    const constant int* idx_shapes [[buffer(7)]],
    const constant int64_t* idx_strides [[buffer(8)]],
    const constant bool* idx_contigs [[buffer(9)]],
    const constant int& idx_ndim [[buffer(10)]],
    const device int32_t *idx0 [[buffer(20)]],
    uint3 index [[thread_position_in_grid]],
    uint3 grid_dim [[threads_per_grid]]) {
  Indices<int32_t, 1> idxs{
    { idx0 }, idx_shapes, idx_strides, idx_contigs, idx_ndim};

  return gather_impl<half, int32_t, 1, 4, int>(
      src,
      out,
      src_shape,
      src_strides,
      src_ndim,
      slice_sizes,
      axes,
      idxs,
      index,
      grid_dim);
}

[[kernel]] void kernel_qwen_mtplx_gather_f16_idxi32_1_4_int64_t(
    const device half* src [[buffer(0)]],
    device half* out [[buffer(1)]],
    const constant int* src_shape [[buffer(2)]],
    const constant int64_t* src_strides [[buffer(3)]],
    const constant size_t& src_ndim [[buffer(4)]],
    const constant int* slice_sizes [[buffer(5)]],
    const constant int* axes [[buffer(6)]],
    const constant int* idx_shapes [[buffer(7)]],
    const constant int64_t* idx_strides [[buffer(8)]],
    const constant bool* idx_contigs [[buffer(9)]],
    const constant int& idx_ndim [[buffer(10)]],
    const device int32_t *idx0 [[buffer(20)]],
    uint3 index [[thread_position_in_grid]],
    uint3 grid_dim [[threads_per_grid]]) {
  Indices<int32_t, 1> idxs{
    { idx0 }, idx_shapes, idx_strides, idx_contigs, idx_ndim};

  return gather_impl<half, int32_t, 1, 4, int64_t>(
      src,
      out,
      src_shape,
      src_strides,
      src_ndim,
      slice_sizes,
      axes,
      idxs,
      index,
      grid_dim);
}

[[kernel]] void kernel_qwen_mtplx_gather_u32_idxi32_1_0_int(
    const device uint32_t* src [[buffer(0)]],
    device uint32_t* out [[buffer(1)]],
    const constant int* src_shape [[buffer(2)]],
    const constant int64_t* src_strides [[buffer(3)]],
    const constant size_t& src_ndim [[buffer(4)]],
    const constant int* slice_sizes [[buffer(5)]],
    const constant int* axes [[buffer(6)]],
    const constant int* idx_shapes [[buffer(7)]],
    const constant int64_t* idx_strides [[buffer(8)]],
    const constant bool* idx_contigs [[buffer(9)]],
    const constant int& idx_ndim [[buffer(10)]],
    const device int32_t *idx0 [[buffer(20)]],
    uint3 index [[thread_position_in_grid]],
    uint3 grid_dim [[threads_per_grid]]) {
  Indices<int32_t, 1> idxs{
    { idx0 }, idx_shapes, idx_strides, idx_contigs, idx_ndim};

  return gather_impl<uint32_t, int32_t, 1, 0, int>(
      src,
      out,
      src_shape,
      src_strides,
      src_ndim,
      slice_sizes,
      axes,
      idxs,
      index,
      grid_dim);
}

[[kernel]] void kernel_qwen_mtplx_gather_u32_idxi32_1_0_int64_t(
    const device uint32_t* src [[buffer(0)]],
    device uint32_t* out [[buffer(1)]],
    const constant int* src_shape [[buffer(2)]],
    const constant int64_t* src_strides [[buffer(3)]],
    const constant size_t& src_ndim [[buffer(4)]],
    const constant int* slice_sizes [[buffer(5)]],
    const constant int* axes [[buffer(6)]],
    const constant int* idx_shapes [[buffer(7)]],
    const constant int64_t* idx_strides [[buffer(8)]],
    const constant bool* idx_contigs [[buffer(9)]],
    const constant int& idx_ndim [[buffer(10)]],
    const device int32_t *idx0 [[buffer(20)]],
    uint3 index [[thread_position_in_grid]],
    uint3 grid_dim [[threads_per_grid]]) {
  Indices<int32_t, 1> idxs{
    { idx0 }, idx_shapes, idx_strides, idx_contigs, idx_ndim};

  return gather_impl<uint32_t, int32_t, 1, 0, int64_t>(
      src,
      out,
      src_shape,
      src_strides,
      src_ndim,
      slice_sizes,
      axes,
      idxs,
      index,
      grid_dim);
}

[[kernel]] void kernel_qwen_mtplx_gather_u32_idxi32_1_1_int(
    const device uint32_t* src [[buffer(0)]],
    device uint32_t* out [[buffer(1)]],
    const constant int* src_shape [[buffer(2)]],
    const constant int64_t* src_strides [[buffer(3)]],
    const constant size_t& src_ndim [[buffer(4)]],
    const constant int* slice_sizes [[buffer(5)]],
    const constant int* axes [[buffer(6)]],
    const constant int* idx_shapes [[buffer(7)]],
    const constant int64_t* idx_strides [[buffer(8)]],
    const constant bool* idx_contigs [[buffer(9)]],
    const constant int& idx_ndim [[buffer(10)]],
    const device int32_t *idx0 [[buffer(20)]],
    uint3 index [[thread_position_in_grid]],
    uint3 grid_dim [[threads_per_grid]]) {
  Indices<int32_t, 1> idxs{
    { idx0 }, idx_shapes, idx_strides, idx_contigs, idx_ndim};

  return gather_impl<uint32_t, int32_t, 1, 1, int>(
      src,
      out,
      src_shape,
      src_strides,
      src_ndim,
      slice_sizes,
      axes,
      idxs,
      index,
      grid_dim);
}

[[kernel]] void kernel_qwen_mtplx_gather_u32_idxi32_1_1_int64_t(
    const device uint32_t* src [[buffer(0)]],
    device uint32_t* out [[buffer(1)]],
    const constant int* src_shape [[buffer(2)]],
    const constant int64_t* src_strides [[buffer(3)]],
    const constant size_t& src_ndim [[buffer(4)]],
    const constant int* slice_sizes [[buffer(5)]],
    const constant int* axes [[buffer(6)]],
    const constant int* idx_shapes [[buffer(7)]],
    const constant int64_t* idx_strides [[buffer(8)]],
    const constant bool* idx_contigs [[buffer(9)]],
    const constant int& idx_ndim [[buffer(10)]],
    const device int32_t *idx0 [[buffer(20)]],
    uint3 index [[thread_position_in_grid]],
    uint3 grid_dim [[threads_per_grid]]) {
  Indices<int32_t, 1> idxs{
    { idx0 }, idx_shapes, idx_strides, idx_contigs, idx_ndim};

  return gather_impl<uint32_t, int32_t, 1, 1, int64_t>(
      src,
      out,
      src_shape,
      src_strides,
      src_ndim,
      slice_sizes,
      axes,
      idxs,
      index,
      grid_dim);
}

[[kernel]] void kernel_qwen_mtplx_gather_u32_idxi32_1_2_int(
    const device uint32_t* src [[buffer(0)]],
    device uint32_t* out [[buffer(1)]],
    const constant int* src_shape [[buffer(2)]],
    const constant int64_t* src_strides [[buffer(3)]],
    const constant size_t& src_ndim [[buffer(4)]],
    const constant int* slice_sizes [[buffer(5)]],
    const constant int* axes [[buffer(6)]],
    const constant int* idx_shapes [[buffer(7)]],
    const constant int64_t* idx_strides [[buffer(8)]],
    const constant bool* idx_contigs [[buffer(9)]],
    const constant int& idx_ndim [[buffer(10)]],
    const device int32_t *idx0 [[buffer(20)]],
    uint3 index [[thread_position_in_grid]],
    uint3 grid_dim [[threads_per_grid]]) {
  Indices<int32_t, 1> idxs{
    { idx0 }, idx_shapes, idx_strides, idx_contigs, idx_ndim};

  return gather_impl<uint32_t, int32_t, 1, 2, int>(
      src,
      out,
      src_shape,
      src_strides,
      src_ndim,
      slice_sizes,
      axes,
      idxs,
      index,
      grid_dim);
}

[[kernel]] void kernel_qwen_mtplx_gather_u32_idxi32_1_2_int64_t(
    const device uint32_t* src [[buffer(0)]],
    device uint32_t* out [[buffer(1)]],
    const constant int* src_shape [[buffer(2)]],
    const constant int64_t* src_strides [[buffer(3)]],
    const constant size_t& src_ndim [[buffer(4)]],
    const constant int* slice_sizes [[buffer(5)]],
    const constant int* axes [[buffer(6)]],
    const constant int* idx_shapes [[buffer(7)]],
    const constant int64_t* idx_strides [[buffer(8)]],
    const constant bool* idx_contigs [[buffer(9)]],
    const constant int& idx_ndim [[buffer(10)]],
    const device int32_t *idx0 [[buffer(20)]],
    uint3 index [[thread_position_in_grid]],
    uint3 grid_dim [[threads_per_grid]]) {
  Indices<int32_t, 1> idxs{
    { idx0 }, idx_shapes, idx_strides, idx_contigs, idx_ndim};

  return gather_impl<uint32_t, int32_t, 1, 2, int64_t>(
      src,
      out,
      src_shape,
      src_strides,
      src_ndim,
      slice_sizes,
      axes,
      idxs,
      index,
      grid_dim);
}

[[kernel]] void kernel_qwen_mtplx_gather_u32_idxi32_1_3_int(
    const device uint32_t* src [[buffer(0)]],
    device uint32_t* out [[buffer(1)]],
    const constant int* src_shape [[buffer(2)]],
    const constant int64_t* src_strides [[buffer(3)]],
    const constant size_t& src_ndim [[buffer(4)]],
    const constant int* slice_sizes [[buffer(5)]],
    const constant int* axes [[buffer(6)]],
    const constant int* idx_shapes [[buffer(7)]],
    const constant int64_t* idx_strides [[buffer(8)]],
    const constant bool* idx_contigs [[buffer(9)]],
    const constant int& idx_ndim [[buffer(10)]],
    const device int32_t *idx0 [[buffer(20)]],
    uint3 index [[thread_position_in_grid]],
    uint3 grid_dim [[threads_per_grid]]) {
  Indices<int32_t, 1> idxs{
    { idx0 }, idx_shapes, idx_strides, idx_contigs, idx_ndim};

  return gather_impl<uint32_t, int32_t, 1, 3, int>(
      src,
      out,
      src_shape,
      src_strides,
      src_ndim,
      slice_sizes,
      axes,
      idxs,
      index,
      grid_dim);
}

[[kernel]] void kernel_qwen_mtplx_gather_u32_idxi32_1_3_int64_t(
    const device uint32_t* src [[buffer(0)]],
    device uint32_t* out [[buffer(1)]],
    const constant int* src_shape [[buffer(2)]],
    const constant int64_t* src_strides [[buffer(3)]],
    const constant size_t& src_ndim [[buffer(4)]],
    const constant int* slice_sizes [[buffer(5)]],
    const constant int* axes [[buffer(6)]],
    const constant int* idx_shapes [[buffer(7)]],
    const constant int64_t* idx_strides [[buffer(8)]],
    const constant bool* idx_contigs [[buffer(9)]],
    const constant int& idx_ndim [[buffer(10)]],
    const device int32_t *idx0 [[buffer(20)]],
    uint3 index [[thread_position_in_grid]],
    uint3 grid_dim [[threads_per_grid]]) {
  Indices<int32_t, 1> idxs{
    { idx0 }, idx_shapes, idx_strides, idx_contigs, idx_ndim};

  return gather_impl<uint32_t, int32_t, 1, 3, int64_t>(
      src,
      out,
      src_shape,
      src_strides,
      src_ndim,
      slice_sizes,
      axes,
      idxs,
      index,
      grid_dim);
}

[[kernel]] void kernel_qwen_mtplx_gather_u32_idxi32_1_4_int(
    const device uint32_t* src [[buffer(0)]],
    device uint32_t* out [[buffer(1)]],
    const constant int* src_shape [[buffer(2)]],
    const constant int64_t* src_strides [[buffer(3)]],
    const constant size_t& src_ndim [[buffer(4)]],
    const constant int* slice_sizes [[buffer(5)]],
    const constant int* axes [[buffer(6)]],
    const constant int* idx_shapes [[buffer(7)]],
    const constant int64_t* idx_strides [[buffer(8)]],
    const constant bool* idx_contigs [[buffer(9)]],
    const constant int& idx_ndim [[buffer(10)]],
    const device int32_t *idx0 [[buffer(20)]],
    uint3 index [[thread_position_in_grid]],
    uint3 grid_dim [[threads_per_grid]]) {
  Indices<int32_t, 1> idxs{
    { idx0 }, idx_shapes, idx_strides, idx_contigs, idx_ndim};

  return gather_impl<uint32_t, int32_t, 1, 4, int>(
      src,
      out,
      src_shape,
      src_strides,
      src_ndim,
      slice_sizes,
      axes,
      idxs,
      index,
      grid_dim);
}

[[kernel]] void kernel_qwen_mtplx_gather_u32_idxi32_1_4_int64_t(
    const device uint32_t* src [[buffer(0)]],
    device uint32_t* out [[buffer(1)]],
    const constant int* src_shape [[buffer(2)]],
    const constant int64_t* src_strides [[buffer(3)]],
    const constant size_t& src_ndim [[buffer(4)]],
    const constant int* slice_sizes [[buffer(5)]],
    const constant int* axes [[buffer(6)]],
    const constant int* idx_shapes [[buffer(7)]],
    const constant int64_t* idx_strides [[buffer(8)]],
    const constant bool* idx_contigs [[buffer(9)]],
    const constant int& idx_ndim [[buffer(10)]],
    const device int32_t *idx0 [[buffer(20)]],
    uint3 index [[thread_position_in_grid]],
    uint3 grid_dim [[threads_per_grid]]) {
  Indices<int32_t, 1> idxs{
    { idx0 }, idx_shapes, idx_strides, idx_contigs, idx_ndim};

  return gather_impl<uint32_t, int32_t, 1, 4, int64_t>(
      src,
      out,
      src_shape,
      src_strides,
      src_ndim,
      slice_sizes,
      axes,
      idxs,
      index,
      grid_dim);
}

[[kernel]] void kernel_qwen_mtplx_gather_f32_idxi32_1_0_int(
    const device float* src [[buffer(0)]],
    device float* out [[buffer(1)]],
    const constant int* src_shape [[buffer(2)]],
    const constant int64_t* src_strides [[buffer(3)]],
    const constant size_t& src_ndim [[buffer(4)]],
    const constant int* slice_sizes [[buffer(5)]],
    const constant int* axes [[buffer(6)]],
    const constant int* idx_shapes [[buffer(7)]],
    const constant int64_t* idx_strides [[buffer(8)]],
    const constant bool* idx_contigs [[buffer(9)]],
    const constant int& idx_ndim [[buffer(10)]],
    const device int32_t *idx0 [[buffer(20)]],
    uint3 index [[thread_position_in_grid]],
    uint3 grid_dim [[threads_per_grid]]) {
  Indices<int32_t, 1> idxs{
    { idx0 }, idx_shapes, idx_strides, idx_contigs, idx_ndim};

  return gather_impl<float, int32_t, 1, 0, int>(
      src,
      out,
      src_shape,
      src_strides,
      src_ndim,
      slice_sizes,
      axes,
      idxs,
      index,
      grid_dim);
}

[[kernel]] void kernel_qwen_mtplx_gather_f32_idxi32_1_0_int64_t(
    const device float* src [[buffer(0)]],
    device float* out [[buffer(1)]],
    const constant int* src_shape [[buffer(2)]],
    const constant int64_t* src_strides [[buffer(3)]],
    const constant size_t& src_ndim [[buffer(4)]],
    const constant int* slice_sizes [[buffer(5)]],
    const constant int* axes [[buffer(6)]],
    const constant int* idx_shapes [[buffer(7)]],
    const constant int64_t* idx_strides [[buffer(8)]],
    const constant bool* idx_contigs [[buffer(9)]],
    const constant int& idx_ndim [[buffer(10)]],
    const device int32_t *idx0 [[buffer(20)]],
    uint3 index [[thread_position_in_grid]],
    uint3 grid_dim [[threads_per_grid]]) {
  Indices<int32_t, 1> idxs{
    { idx0 }, idx_shapes, idx_strides, idx_contigs, idx_ndim};

  return gather_impl<float, int32_t, 1, 0, int64_t>(
      src,
      out,
      src_shape,
      src_strides,
      src_ndim,
      slice_sizes,
      axes,
      idxs,
      index,
      grid_dim);
}

[[kernel]] void kernel_qwen_mtplx_gather_f32_idxi32_1_1_int(
    const device float* src [[buffer(0)]],
    device float* out [[buffer(1)]],
    const constant int* src_shape [[buffer(2)]],
    const constant int64_t* src_strides [[buffer(3)]],
    const constant size_t& src_ndim [[buffer(4)]],
    const constant int* slice_sizes [[buffer(5)]],
    const constant int* axes [[buffer(6)]],
    const constant int* idx_shapes [[buffer(7)]],
    const constant int64_t* idx_strides [[buffer(8)]],
    const constant bool* idx_contigs [[buffer(9)]],
    const constant int& idx_ndim [[buffer(10)]],
    const device int32_t *idx0 [[buffer(20)]],
    uint3 index [[thread_position_in_grid]],
    uint3 grid_dim [[threads_per_grid]]) {
  Indices<int32_t, 1> idxs{
    { idx0 }, idx_shapes, idx_strides, idx_contigs, idx_ndim};

  return gather_impl<float, int32_t, 1, 1, int>(
      src,
      out,
      src_shape,
      src_strides,
      src_ndim,
      slice_sizes,
      axes,
      idxs,
      index,
      grid_dim);
}

[[kernel]] void kernel_qwen_mtplx_gather_f32_idxi32_1_1_int64_t(
    const device float* src [[buffer(0)]],
    device float* out [[buffer(1)]],
    const constant int* src_shape [[buffer(2)]],
    const constant int64_t* src_strides [[buffer(3)]],
    const constant size_t& src_ndim [[buffer(4)]],
    const constant int* slice_sizes [[buffer(5)]],
    const constant int* axes [[buffer(6)]],
    const constant int* idx_shapes [[buffer(7)]],
    const constant int64_t* idx_strides [[buffer(8)]],
    const constant bool* idx_contigs [[buffer(9)]],
    const constant int& idx_ndim [[buffer(10)]],
    const device int32_t *idx0 [[buffer(20)]],
    uint3 index [[thread_position_in_grid]],
    uint3 grid_dim [[threads_per_grid]]) {
  Indices<int32_t, 1> idxs{
    { idx0 }, idx_shapes, idx_strides, idx_contigs, idx_ndim};

  return gather_impl<float, int32_t, 1, 1, int64_t>(
      src,
      out,
      src_shape,
      src_strides,
      src_ndim,
      slice_sizes,
      axes,
      idxs,
      index,
      grid_dim);
}

[[kernel]] void kernel_qwen_mtplx_gather_f32_idxi32_1_2_int(
    const device float* src [[buffer(0)]],
    device float* out [[buffer(1)]],
    const constant int* src_shape [[buffer(2)]],
    const constant int64_t* src_strides [[buffer(3)]],
    const constant size_t& src_ndim [[buffer(4)]],
    const constant int* slice_sizes [[buffer(5)]],
    const constant int* axes [[buffer(6)]],
    const constant int* idx_shapes [[buffer(7)]],
    const constant int64_t* idx_strides [[buffer(8)]],
    const constant bool* idx_contigs [[buffer(9)]],
    const constant int& idx_ndim [[buffer(10)]],
    const device int32_t *idx0 [[buffer(20)]],
    uint3 index [[thread_position_in_grid]],
    uint3 grid_dim [[threads_per_grid]]) {
  Indices<int32_t, 1> idxs{
    { idx0 }, idx_shapes, idx_strides, idx_contigs, idx_ndim};

  return gather_impl<float, int32_t, 1, 2, int>(
      src,
      out,
      src_shape,
      src_strides,
      src_ndim,
      slice_sizes,
      axes,
      idxs,
      index,
      grid_dim);
}

[[kernel]] void kernel_qwen_mtplx_gather_f32_idxi32_1_2_int64_t(
    const device float* src [[buffer(0)]],
    device float* out [[buffer(1)]],
    const constant int* src_shape [[buffer(2)]],
    const constant int64_t* src_strides [[buffer(3)]],
    const constant size_t& src_ndim [[buffer(4)]],
    const constant int* slice_sizes [[buffer(5)]],
    const constant int* axes [[buffer(6)]],
    const constant int* idx_shapes [[buffer(7)]],
    const constant int64_t* idx_strides [[buffer(8)]],
    const constant bool* idx_contigs [[buffer(9)]],
    const constant int& idx_ndim [[buffer(10)]],
    const device int32_t *idx0 [[buffer(20)]],
    uint3 index [[thread_position_in_grid]],
    uint3 grid_dim [[threads_per_grid]]) {
  Indices<int32_t, 1> idxs{
    { idx0 }, idx_shapes, idx_strides, idx_contigs, idx_ndim};

  return gather_impl<float, int32_t, 1, 2, int64_t>(
      src,
      out,
      src_shape,
      src_strides,
      src_ndim,
      slice_sizes,
      axes,
      idxs,
      index,
      grid_dim);
}

[[kernel]] void kernel_qwen_mtplx_gather_f32_idxi32_1_3_int(
    const device float* src [[buffer(0)]],
    device float* out [[buffer(1)]],
    const constant int* src_shape [[buffer(2)]],
    const constant int64_t* src_strides [[buffer(3)]],
    const constant size_t& src_ndim [[buffer(4)]],
    const constant int* slice_sizes [[buffer(5)]],
    const constant int* axes [[buffer(6)]],
    const constant int* idx_shapes [[buffer(7)]],
    const constant int64_t* idx_strides [[buffer(8)]],
    const constant bool* idx_contigs [[buffer(9)]],
    const constant int& idx_ndim [[buffer(10)]],
    const device int32_t *idx0 [[buffer(20)]],
    uint3 index [[thread_position_in_grid]],
    uint3 grid_dim [[threads_per_grid]]) {
  Indices<int32_t, 1> idxs{
    { idx0 }, idx_shapes, idx_strides, idx_contigs, idx_ndim};

  return gather_impl<float, int32_t, 1, 3, int>(
      src,
      out,
      src_shape,
      src_strides,
      src_ndim,
      slice_sizes,
      axes,
      idxs,
      index,
      grid_dim);
}

[[kernel]] void kernel_qwen_mtplx_gather_f32_idxi32_1_3_int64_t(
    const device float* src [[buffer(0)]],
    device float* out [[buffer(1)]],
    const constant int* src_shape [[buffer(2)]],
    const constant int64_t* src_strides [[buffer(3)]],
    const constant size_t& src_ndim [[buffer(4)]],
    const constant int* slice_sizes [[buffer(5)]],
    const constant int* axes [[buffer(6)]],
    const constant int* idx_shapes [[buffer(7)]],
    const constant int64_t* idx_strides [[buffer(8)]],
    const constant bool* idx_contigs [[buffer(9)]],
    const constant int& idx_ndim [[buffer(10)]],
    const device int32_t *idx0 [[buffer(20)]],
    uint3 index [[thread_position_in_grid]],
    uint3 grid_dim [[threads_per_grid]]) {
  Indices<int32_t, 1> idxs{
    { idx0 }, idx_shapes, idx_strides, idx_contigs, idx_ndim};

  return gather_impl<float, int32_t, 1, 3, int64_t>(
      src,
      out,
      src_shape,
      src_strides,
      src_ndim,
      slice_sizes,
      axes,
      idxs,
      index,
      grid_dim);
}

[[kernel]] void kernel_qwen_mtplx_gather_f32_idxi32_1_4_int(
    const device float* src [[buffer(0)]],
    device float* out [[buffer(1)]],
    const constant int* src_shape [[buffer(2)]],
    const constant int64_t* src_strides [[buffer(3)]],
    const constant size_t& src_ndim [[buffer(4)]],
    const constant int* slice_sizes [[buffer(5)]],
    const constant int* axes [[buffer(6)]],
    const constant int* idx_shapes [[buffer(7)]],
    const constant int64_t* idx_strides [[buffer(8)]],
    const constant bool* idx_contigs [[buffer(9)]],
    const constant int& idx_ndim [[buffer(10)]],
    const device int32_t *idx0 [[buffer(20)]],
    uint3 index [[thread_position_in_grid]],
    uint3 grid_dim [[threads_per_grid]]) {
  Indices<int32_t, 1> idxs{
    { idx0 }, idx_shapes, idx_strides, idx_contigs, idx_ndim};

  return gather_impl<float, int32_t, 1, 4, int>(
      src,
      out,
      src_shape,
      src_strides,
      src_ndim,
      slice_sizes,
      axes,
      idxs,
      index,
      grid_dim);
}

[[kernel]] void kernel_qwen_mtplx_gather_f32_idxi32_1_4_int64_t(
    const device float* src [[buffer(0)]],
    device float* out [[buffer(1)]],
    const constant int* src_shape [[buffer(2)]],
    const constant int64_t* src_strides [[buffer(3)]],
    const constant size_t& src_ndim [[buffer(4)]],
    const constant int* slice_sizes [[buffer(5)]],
    const constant int* axes [[buffer(6)]],
    const constant int* idx_shapes [[buffer(7)]],
    const constant int64_t* idx_strides [[buffer(8)]],
    const constant bool* idx_contigs [[buffer(9)]],
    const constant int& idx_ndim [[buffer(10)]],
    const device int32_t *idx0 [[buffer(20)]],
    uint3 index [[thread_position_in_grid]],
    uint3 grid_dim [[threads_per_grid]]) {
  Indices<int32_t, 1> idxs{
    { idx0 }, idx_shapes, idx_strides, idx_contigs, idx_ndim};

  return gather_impl<float, int32_t, 1, 4, int64_t>(
      src,
      out,
      src_shape,
      src_strides,
      src_ndim,
      slice_sizes,
      axes,
      idxs,
      index,
      grid_dim);
}

[[kernel]] void kernel_qwen_mtplx_gather_bf16_idxu32_1_0_int(
    const device bfloat* src [[buffer(0)]],
    device bfloat* out [[buffer(1)]],
    const constant int* src_shape [[buffer(2)]],
    const constant int64_t* src_strides [[buffer(3)]],
    const constant size_t& src_ndim [[buffer(4)]],
    const constant int* slice_sizes [[buffer(5)]],
    const constant int* axes [[buffer(6)]],
    const constant int* idx_shapes [[buffer(7)]],
    const constant int64_t* idx_strides [[buffer(8)]],
    const constant bool* idx_contigs [[buffer(9)]],
    const constant int& idx_ndim [[buffer(10)]],
    const device uint32_t *idx0 [[buffer(20)]],
    uint3 index [[thread_position_in_grid]],
    uint3 grid_dim [[threads_per_grid]]) {
  Indices<uint32_t, 1> idxs{
    { idx0 }, idx_shapes, idx_strides, idx_contigs, idx_ndim};

  return gather_impl<bfloat, uint32_t, 1, 0, int>(
      src,
      out,
      src_shape,
      src_strides,
      src_ndim,
      slice_sizes,
      axes,
      idxs,
      index,
      grid_dim);
}

[[kernel]] void kernel_qwen_mtplx_gather_bf16_idxu32_1_0_int64_t(
    const device bfloat* src [[buffer(0)]],
    device bfloat* out [[buffer(1)]],
    const constant int* src_shape [[buffer(2)]],
    const constant int64_t* src_strides [[buffer(3)]],
    const constant size_t& src_ndim [[buffer(4)]],
    const constant int* slice_sizes [[buffer(5)]],
    const constant int* axes [[buffer(6)]],
    const constant int* idx_shapes [[buffer(7)]],
    const constant int64_t* idx_strides [[buffer(8)]],
    const constant bool* idx_contigs [[buffer(9)]],
    const constant int& idx_ndim [[buffer(10)]],
    const device uint32_t *idx0 [[buffer(20)]],
    uint3 index [[thread_position_in_grid]],
    uint3 grid_dim [[threads_per_grid]]) {
  Indices<uint32_t, 1> idxs{
    { idx0 }, idx_shapes, idx_strides, idx_contigs, idx_ndim};

  return gather_impl<bfloat, uint32_t, 1, 0, int64_t>(
      src,
      out,
      src_shape,
      src_strides,
      src_ndim,
      slice_sizes,
      axes,
      idxs,
      index,
      grid_dim);
}

[[kernel]] void kernel_qwen_mtplx_gather_bf16_idxu32_1_1_int(
    const device bfloat* src [[buffer(0)]],
    device bfloat* out [[buffer(1)]],
    const constant int* src_shape [[buffer(2)]],
    const constant int64_t* src_strides [[buffer(3)]],
    const constant size_t& src_ndim [[buffer(4)]],
    const constant int* slice_sizes [[buffer(5)]],
    const constant int* axes [[buffer(6)]],
    const constant int* idx_shapes [[buffer(7)]],
    const constant int64_t* idx_strides [[buffer(8)]],
    const constant bool* idx_contigs [[buffer(9)]],
    const constant int& idx_ndim [[buffer(10)]],
    const device uint32_t *idx0 [[buffer(20)]],
    uint3 index [[thread_position_in_grid]],
    uint3 grid_dim [[threads_per_grid]]) {
  Indices<uint32_t, 1> idxs{
    { idx0 }, idx_shapes, idx_strides, idx_contigs, idx_ndim};

  return gather_impl<bfloat, uint32_t, 1, 1, int>(
      src,
      out,
      src_shape,
      src_strides,
      src_ndim,
      slice_sizes,
      axes,
      idxs,
      index,
      grid_dim);
}

[[kernel]] void kernel_qwen_mtplx_gather_bf16_idxu32_1_1_int64_t(
    const device bfloat* src [[buffer(0)]],
    device bfloat* out [[buffer(1)]],
    const constant int* src_shape [[buffer(2)]],
    const constant int64_t* src_strides [[buffer(3)]],
    const constant size_t& src_ndim [[buffer(4)]],
    const constant int* slice_sizes [[buffer(5)]],
    const constant int* axes [[buffer(6)]],
    const constant int* idx_shapes [[buffer(7)]],
    const constant int64_t* idx_strides [[buffer(8)]],
    const constant bool* idx_contigs [[buffer(9)]],
    const constant int& idx_ndim [[buffer(10)]],
    const device uint32_t *idx0 [[buffer(20)]],
    uint3 index [[thread_position_in_grid]],
    uint3 grid_dim [[threads_per_grid]]) {
  Indices<uint32_t, 1> idxs{
    { idx0 }, idx_shapes, idx_strides, idx_contigs, idx_ndim};

  return gather_impl<bfloat, uint32_t, 1, 1, int64_t>(
      src,
      out,
      src_shape,
      src_strides,
      src_ndim,
      slice_sizes,
      axes,
      idxs,
      index,
      grid_dim);
}

[[kernel]] void kernel_qwen_mtplx_gather_bf16_idxu32_1_2_int(
    const device bfloat* src [[buffer(0)]],
    device bfloat* out [[buffer(1)]],
    const constant int* src_shape [[buffer(2)]],
    const constant int64_t* src_strides [[buffer(3)]],
    const constant size_t& src_ndim [[buffer(4)]],
    const constant int* slice_sizes [[buffer(5)]],
    const constant int* axes [[buffer(6)]],
    const constant int* idx_shapes [[buffer(7)]],
    const constant int64_t* idx_strides [[buffer(8)]],
    const constant bool* idx_contigs [[buffer(9)]],
    const constant int& idx_ndim [[buffer(10)]],
    const device uint32_t *idx0 [[buffer(20)]],
    uint3 index [[thread_position_in_grid]],
    uint3 grid_dim [[threads_per_grid]]) {
  Indices<uint32_t, 1> idxs{
    { idx0 }, idx_shapes, idx_strides, idx_contigs, idx_ndim};

  return gather_impl<bfloat, uint32_t, 1, 2, int>(
      src,
      out,
      src_shape,
      src_strides,
      src_ndim,
      slice_sizes,
      axes,
      idxs,
      index,
      grid_dim);
}

[[kernel]] void kernel_qwen_mtplx_gather_bf16_idxu32_1_2_int64_t(
    const device bfloat* src [[buffer(0)]],
    device bfloat* out [[buffer(1)]],
    const constant int* src_shape [[buffer(2)]],
    const constant int64_t* src_strides [[buffer(3)]],
    const constant size_t& src_ndim [[buffer(4)]],
    const constant int* slice_sizes [[buffer(5)]],
    const constant int* axes [[buffer(6)]],
    const constant int* idx_shapes [[buffer(7)]],
    const constant int64_t* idx_strides [[buffer(8)]],
    const constant bool* idx_contigs [[buffer(9)]],
    const constant int& idx_ndim [[buffer(10)]],
    const device uint32_t *idx0 [[buffer(20)]],
    uint3 index [[thread_position_in_grid]],
    uint3 grid_dim [[threads_per_grid]]) {
  Indices<uint32_t, 1> idxs{
    { idx0 }, idx_shapes, idx_strides, idx_contigs, idx_ndim};

  return gather_impl<bfloat, uint32_t, 1, 2, int64_t>(
      src,
      out,
      src_shape,
      src_strides,
      src_ndim,
      slice_sizes,
      axes,
      idxs,
      index,
      grid_dim);
}

[[kernel]] void kernel_qwen_mtplx_gather_bf16_idxu32_1_3_int(
    const device bfloat* src [[buffer(0)]],
    device bfloat* out [[buffer(1)]],
    const constant int* src_shape [[buffer(2)]],
    const constant int64_t* src_strides [[buffer(3)]],
    const constant size_t& src_ndim [[buffer(4)]],
    const constant int* slice_sizes [[buffer(5)]],
    const constant int* axes [[buffer(6)]],
    const constant int* idx_shapes [[buffer(7)]],
    const constant int64_t* idx_strides [[buffer(8)]],
    const constant bool* idx_contigs [[buffer(9)]],
    const constant int& idx_ndim [[buffer(10)]],
    const device uint32_t *idx0 [[buffer(20)]],
    uint3 index [[thread_position_in_grid]],
    uint3 grid_dim [[threads_per_grid]]) {
  Indices<uint32_t, 1> idxs{
    { idx0 }, idx_shapes, idx_strides, idx_contigs, idx_ndim};

  return gather_impl<bfloat, uint32_t, 1, 3, int>(
      src,
      out,
      src_shape,
      src_strides,
      src_ndim,
      slice_sizes,
      axes,
      idxs,
      index,
      grid_dim);
}

[[kernel]] void kernel_qwen_mtplx_gather_bf16_idxu32_1_3_int64_t(
    const device bfloat* src [[buffer(0)]],
    device bfloat* out [[buffer(1)]],
    const constant int* src_shape [[buffer(2)]],
    const constant int64_t* src_strides [[buffer(3)]],
    const constant size_t& src_ndim [[buffer(4)]],
    const constant int* slice_sizes [[buffer(5)]],
    const constant int* axes [[buffer(6)]],
    const constant int* idx_shapes [[buffer(7)]],
    const constant int64_t* idx_strides [[buffer(8)]],
    const constant bool* idx_contigs [[buffer(9)]],
    const constant int& idx_ndim [[buffer(10)]],
    const device uint32_t *idx0 [[buffer(20)]],
    uint3 index [[thread_position_in_grid]],
    uint3 grid_dim [[threads_per_grid]]) {
  Indices<uint32_t, 1> idxs{
    { idx0 }, idx_shapes, idx_strides, idx_contigs, idx_ndim};

  return gather_impl<bfloat, uint32_t, 1, 3, int64_t>(
      src,
      out,
      src_shape,
      src_strides,
      src_ndim,
      slice_sizes,
      axes,
      idxs,
      index,
      grid_dim);
}

[[kernel]] void kernel_qwen_mtplx_gather_bf16_idxu32_1_4_int(
    const device bfloat* src [[buffer(0)]],
    device bfloat* out [[buffer(1)]],
    const constant int* src_shape [[buffer(2)]],
    const constant int64_t* src_strides [[buffer(3)]],
    const constant size_t& src_ndim [[buffer(4)]],
    const constant int* slice_sizes [[buffer(5)]],
    const constant int* axes [[buffer(6)]],
    const constant int* idx_shapes [[buffer(7)]],
    const constant int64_t* idx_strides [[buffer(8)]],
    const constant bool* idx_contigs [[buffer(9)]],
    const constant int& idx_ndim [[buffer(10)]],
    const device uint32_t *idx0 [[buffer(20)]],
    uint3 index [[thread_position_in_grid]],
    uint3 grid_dim [[threads_per_grid]]) {
  Indices<uint32_t, 1> idxs{
    { idx0 }, idx_shapes, idx_strides, idx_contigs, idx_ndim};

  return gather_impl<bfloat, uint32_t, 1, 4, int>(
      src,
      out,
      src_shape,
      src_strides,
      src_ndim,
      slice_sizes,
      axes,
      idxs,
      index,
      grid_dim);
}

[[kernel]] void kernel_qwen_mtplx_gather_bf16_idxu32_1_4_int64_t(
    const device bfloat* src [[buffer(0)]],
    device bfloat* out [[buffer(1)]],
    const constant int* src_shape [[buffer(2)]],
    const constant int64_t* src_strides [[buffer(3)]],
    const constant size_t& src_ndim [[buffer(4)]],
    const constant int* slice_sizes [[buffer(5)]],
    const constant int* axes [[buffer(6)]],
    const constant int* idx_shapes [[buffer(7)]],
    const constant int64_t* idx_strides [[buffer(8)]],
    const constant bool* idx_contigs [[buffer(9)]],
    const constant int& idx_ndim [[buffer(10)]],
    const device uint32_t *idx0 [[buffer(20)]],
    uint3 index [[thread_position_in_grid]],
    uint3 grid_dim [[threads_per_grid]]) {
  Indices<uint32_t, 1> idxs{
    { idx0 }, idx_shapes, idx_strides, idx_contigs, idx_ndim};

  return gather_impl<bfloat, uint32_t, 1, 4, int64_t>(
      src,
      out,
      src_shape,
      src_strides,
      src_ndim,
      slice_sizes,
      axes,
      idxs,
      index,
      grid_dim);
}

[[kernel]] void kernel_qwen_mtplx_gather_f16_idxu32_1_0_int(
    const device half* src [[buffer(0)]],
    device half* out [[buffer(1)]],
    const constant int* src_shape [[buffer(2)]],
    const constant int64_t* src_strides [[buffer(3)]],
    const constant size_t& src_ndim [[buffer(4)]],
    const constant int* slice_sizes [[buffer(5)]],
    const constant int* axes [[buffer(6)]],
    const constant int* idx_shapes [[buffer(7)]],
    const constant int64_t* idx_strides [[buffer(8)]],
    const constant bool* idx_contigs [[buffer(9)]],
    const constant int& idx_ndim [[buffer(10)]],
    const device uint32_t *idx0 [[buffer(20)]],
    uint3 index [[thread_position_in_grid]],
    uint3 grid_dim [[threads_per_grid]]) {
  Indices<uint32_t, 1> idxs{
    { idx0 }, idx_shapes, idx_strides, idx_contigs, idx_ndim};

  return gather_impl<half, uint32_t, 1, 0, int>(
      src,
      out,
      src_shape,
      src_strides,
      src_ndim,
      slice_sizes,
      axes,
      idxs,
      index,
      grid_dim);
}

[[kernel]] void kernel_qwen_mtplx_gather_f16_idxu32_1_0_int64_t(
    const device half* src [[buffer(0)]],
    device half* out [[buffer(1)]],
    const constant int* src_shape [[buffer(2)]],
    const constant int64_t* src_strides [[buffer(3)]],
    const constant size_t& src_ndim [[buffer(4)]],
    const constant int* slice_sizes [[buffer(5)]],
    const constant int* axes [[buffer(6)]],
    const constant int* idx_shapes [[buffer(7)]],
    const constant int64_t* idx_strides [[buffer(8)]],
    const constant bool* idx_contigs [[buffer(9)]],
    const constant int& idx_ndim [[buffer(10)]],
    const device uint32_t *idx0 [[buffer(20)]],
    uint3 index [[thread_position_in_grid]],
    uint3 grid_dim [[threads_per_grid]]) {
  Indices<uint32_t, 1> idxs{
    { idx0 }, idx_shapes, idx_strides, idx_contigs, idx_ndim};

  return gather_impl<half, uint32_t, 1, 0, int64_t>(
      src,
      out,
      src_shape,
      src_strides,
      src_ndim,
      slice_sizes,
      axes,
      idxs,
      index,
      grid_dim);
}

[[kernel]] void kernel_qwen_mtplx_gather_f16_idxu32_1_1_int(
    const device half* src [[buffer(0)]],
    device half* out [[buffer(1)]],
    const constant int* src_shape [[buffer(2)]],
    const constant int64_t* src_strides [[buffer(3)]],
    const constant size_t& src_ndim [[buffer(4)]],
    const constant int* slice_sizes [[buffer(5)]],
    const constant int* axes [[buffer(6)]],
    const constant int* idx_shapes [[buffer(7)]],
    const constant int64_t* idx_strides [[buffer(8)]],
    const constant bool* idx_contigs [[buffer(9)]],
    const constant int& idx_ndim [[buffer(10)]],
    const device uint32_t *idx0 [[buffer(20)]],
    uint3 index [[thread_position_in_grid]],
    uint3 grid_dim [[threads_per_grid]]) {
  Indices<uint32_t, 1> idxs{
    { idx0 }, idx_shapes, idx_strides, idx_contigs, idx_ndim};

  return gather_impl<half, uint32_t, 1, 1, int>(
      src,
      out,
      src_shape,
      src_strides,
      src_ndim,
      slice_sizes,
      axes,
      idxs,
      index,
      grid_dim);
}

[[kernel]] void kernel_qwen_mtplx_gather_f16_idxu32_1_1_int64_t(
    const device half* src [[buffer(0)]],
    device half* out [[buffer(1)]],
    const constant int* src_shape [[buffer(2)]],
    const constant int64_t* src_strides [[buffer(3)]],
    const constant size_t& src_ndim [[buffer(4)]],
    const constant int* slice_sizes [[buffer(5)]],
    const constant int* axes [[buffer(6)]],
    const constant int* idx_shapes [[buffer(7)]],
    const constant int64_t* idx_strides [[buffer(8)]],
    const constant bool* idx_contigs [[buffer(9)]],
    const constant int& idx_ndim [[buffer(10)]],
    const device uint32_t *idx0 [[buffer(20)]],
    uint3 index [[thread_position_in_grid]],
    uint3 grid_dim [[threads_per_grid]]) {
  Indices<uint32_t, 1> idxs{
    { idx0 }, idx_shapes, idx_strides, idx_contigs, idx_ndim};

  return gather_impl<half, uint32_t, 1, 1, int64_t>(
      src,
      out,
      src_shape,
      src_strides,
      src_ndim,
      slice_sizes,
      axes,
      idxs,
      index,
      grid_dim);
}

[[kernel]] void kernel_qwen_mtplx_gather_f16_idxu32_1_2_int(
    const device half* src [[buffer(0)]],
    device half* out [[buffer(1)]],
    const constant int* src_shape [[buffer(2)]],
    const constant int64_t* src_strides [[buffer(3)]],
    const constant size_t& src_ndim [[buffer(4)]],
    const constant int* slice_sizes [[buffer(5)]],
    const constant int* axes [[buffer(6)]],
    const constant int* idx_shapes [[buffer(7)]],
    const constant int64_t* idx_strides [[buffer(8)]],
    const constant bool* idx_contigs [[buffer(9)]],
    const constant int& idx_ndim [[buffer(10)]],
    const device uint32_t *idx0 [[buffer(20)]],
    uint3 index [[thread_position_in_grid]],
    uint3 grid_dim [[threads_per_grid]]) {
  Indices<uint32_t, 1> idxs{
    { idx0 }, idx_shapes, idx_strides, idx_contigs, idx_ndim};

  return gather_impl<half, uint32_t, 1, 2, int>(
      src,
      out,
      src_shape,
      src_strides,
      src_ndim,
      slice_sizes,
      axes,
      idxs,
      index,
      grid_dim);
}

[[kernel]] void kernel_qwen_mtplx_gather_f16_idxu32_1_2_int64_t(
    const device half* src [[buffer(0)]],
    device half* out [[buffer(1)]],
    const constant int* src_shape [[buffer(2)]],
    const constant int64_t* src_strides [[buffer(3)]],
    const constant size_t& src_ndim [[buffer(4)]],
    const constant int* slice_sizes [[buffer(5)]],
    const constant int* axes [[buffer(6)]],
    const constant int* idx_shapes [[buffer(7)]],
    const constant int64_t* idx_strides [[buffer(8)]],
    const constant bool* idx_contigs [[buffer(9)]],
    const constant int& idx_ndim [[buffer(10)]],
    const device uint32_t *idx0 [[buffer(20)]],
    uint3 index [[thread_position_in_grid]],
    uint3 grid_dim [[threads_per_grid]]) {
  Indices<uint32_t, 1> idxs{
    { idx0 }, idx_shapes, idx_strides, idx_contigs, idx_ndim};

  return gather_impl<half, uint32_t, 1, 2, int64_t>(
      src,
      out,
      src_shape,
      src_strides,
      src_ndim,
      slice_sizes,
      axes,
      idxs,
      index,
      grid_dim);
}

[[kernel]] void kernel_qwen_mtplx_gather_f16_idxu32_1_3_int(
    const device half* src [[buffer(0)]],
    device half* out [[buffer(1)]],
    const constant int* src_shape [[buffer(2)]],
    const constant int64_t* src_strides [[buffer(3)]],
    const constant size_t& src_ndim [[buffer(4)]],
    const constant int* slice_sizes [[buffer(5)]],
    const constant int* axes [[buffer(6)]],
    const constant int* idx_shapes [[buffer(7)]],
    const constant int64_t* idx_strides [[buffer(8)]],
    const constant bool* idx_contigs [[buffer(9)]],
    const constant int& idx_ndim [[buffer(10)]],
    const device uint32_t *idx0 [[buffer(20)]],
    uint3 index [[thread_position_in_grid]],
    uint3 grid_dim [[threads_per_grid]]) {
  Indices<uint32_t, 1> idxs{
    { idx0 }, idx_shapes, idx_strides, idx_contigs, idx_ndim};

  return gather_impl<half, uint32_t, 1, 3, int>(
      src,
      out,
      src_shape,
      src_strides,
      src_ndim,
      slice_sizes,
      axes,
      idxs,
      index,
      grid_dim);
}

[[kernel]] void kernel_qwen_mtplx_gather_f16_idxu32_1_3_int64_t(
    const device half* src [[buffer(0)]],
    device half* out [[buffer(1)]],
    const constant int* src_shape [[buffer(2)]],
    const constant int64_t* src_strides [[buffer(3)]],
    const constant size_t& src_ndim [[buffer(4)]],
    const constant int* slice_sizes [[buffer(5)]],
    const constant int* axes [[buffer(6)]],
    const constant int* idx_shapes [[buffer(7)]],
    const constant int64_t* idx_strides [[buffer(8)]],
    const constant bool* idx_contigs [[buffer(9)]],
    const constant int& idx_ndim [[buffer(10)]],
    const device uint32_t *idx0 [[buffer(20)]],
    uint3 index [[thread_position_in_grid]],
    uint3 grid_dim [[threads_per_grid]]) {
  Indices<uint32_t, 1> idxs{
    { idx0 }, idx_shapes, idx_strides, idx_contigs, idx_ndim};

  return gather_impl<half, uint32_t, 1, 3, int64_t>(
      src,
      out,
      src_shape,
      src_strides,
      src_ndim,
      slice_sizes,
      axes,
      idxs,
      index,
      grid_dim);
}

[[kernel]] void kernel_qwen_mtplx_gather_f16_idxu32_1_4_int(
    const device half* src [[buffer(0)]],
    device half* out [[buffer(1)]],
    const constant int* src_shape [[buffer(2)]],
    const constant int64_t* src_strides [[buffer(3)]],
    const constant size_t& src_ndim [[buffer(4)]],
    const constant int* slice_sizes [[buffer(5)]],
    const constant int* axes [[buffer(6)]],
    const constant int* idx_shapes [[buffer(7)]],
    const constant int64_t* idx_strides [[buffer(8)]],
    const constant bool* idx_contigs [[buffer(9)]],
    const constant int& idx_ndim [[buffer(10)]],
    const device uint32_t *idx0 [[buffer(20)]],
    uint3 index [[thread_position_in_grid]],
    uint3 grid_dim [[threads_per_grid]]) {
  Indices<uint32_t, 1> idxs{
    { idx0 }, idx_shapes, idx_strides, idx_contigs, idx_ndim};

  return gather_impl<half, uint32_t, 1, 4, int>(
      src,
      out,
      src_shape,
      src_strides,
      src_ndim,
      slice_sizes,
      axes,
      idxs,
      index,
      grid_dim);
}

[[kernel]] void kernel_qwen_mtplx_gather_f16_idxu32_1_4_int64_t(
    const device half* src [[buffer(0)]],
    device half* out [[buffer(1)]],
    const constant int* src_shape [[buffer(2)]],
    const constant int64_t* src_strides [[buffer(3)]],
    const constant size_t& src_ndim [[buffer(4)]],
    const constant int* slice_sizes [[buffer(5)]],
    const constant int* axes [[buffer(6)]],
    const constant int* idx_shapes [[buffer(7)]],
    const constant int64_t* idx_strides [[buffer(8)]],
    const constant bool* idx_contigs [[buffer(9)]],
    const constant int& idx_ndim [[buffer(10)]],
    const device uint32_t *idx0 [[buffer(20)]],
    uint3 index [[thread_position_in_grid]],
    uint3 grid_dim [[threads_per_grid]]) {
  Indices<uint32_t, 1> idxs{
    { idx0 }, idx_shapes, idx_strides, idx_contigs, idx_ndim};

  return gather_impl<half, uint32_t, 1, 4, int64_t>(
      src,
      out,
      src_shape,
      src_strides,
      src_ndim,
      slice_sizes,
      axes,
      idxs,
      index,
      grid_dim);
}

[[kernel]] void kernel_qwen_mtplx_gather_u32_idxu32_1_0_int(
    const device uint32_t* src [[buffer(0)]],
    device uint32_t* out [[buffer(1)]],
    const constant int* src_shape [[buffer(2)]],
    const constant int64_t* src_strides [[buffer(3)]],
    const constant size_t& src_ndim [[buffer(4)]],
    const constant int* slice_sizes [[buffer(5)]],
    const constant int* axes [[buffer(6)]],
    const constant int* idx_shapes [[buffer(7)]],
    const constant int64_t* idx_strides [[buffer(8)]],
    const constant bool* idx_contigs [[buffer(9)]],
    const constant int& idx_ndim [[buffer(10)]],
    const device uint32_t *idx0 [[buffer(20)]],
    uint3 index [[thread_position_in_grid]],
    uint3 grid_dim [[threads_per_grid]]) {
  Indices<uint32_t, 1> idxs{
    { idx0 }, idx_shapes, idx_strides, idx_contigs, idx_ndim};

  return gather_impl<uint32_t, uint32_t, 1, 0, int>(
      src,
      out,
      src_shape,
      src_strides,
      src_ndim,
      slice_sizes,
      axes,
      idxs,
      index,
      grid_dim);
}

[[kernel]] void kernel_qwen_mtplx_gather_u32_idxu32_1_0_int64_t(
    const device uint32_t* src [[buffer(0)]],
    device uint32_t* out [[buffer(1)]],
    const constant int* src_shape [[buffer(2)]],
    const constant int64_t* src_strides [[buffer(3)]],
    const constant size_t& src_ndim [[buffer(4)]],
    const constant int* slice_sizes [[buffer(5)]],
    const constant int* axes [[buffer(6)]],
    const constant int* idx_shapes [[buffer(7)]],
    const constant int64_t* idx_strides [[buffer(8)]],
    const constant bool* idx_contigs [[buffer(9)]],
    const constant int& idx_ndim [[buffer(10)]],
    const device uint32_t *idx0 [[buffer(20)]],
    uint3 index [[thread_position_in_grid]],
    uint3 grid_dim [[threads_per_grid]]) {
  Indices<uint32_t, 1> idxs{
    { idx0 }, idx_shapes, idx_strides, idx_contigs, idx_ndim};

  return gather_impl<uint32_t, uint32_t, 1, 0, int64_t>(
      src,
      out,
      src_shape,
      src_strides,
      src_ndim,
      slice_sizes,
      axes,
      idxs,
      index,
      grid_dim);
}

[[kernel]] void kernel_qwen_mtplx_gather_u32_idxu32_1_1_int(
    const device uint32_t* src [[buffer(0)]],
    device uint32_t* out [[buffer(1)]],
    const constant int* src_shape [[buffer(2)]],
    const constant int64_t* src_strides [[buffer(3)]],
    const constant size_t& src_ndim [[buffer(4)]],
    const constant int* slice_sizes [[buffer(5)]],
    const constant int* axes [[buffer(6)]],
    const constant int* idx_shapes [[buffer(7)]],
    const constant int64_t* idx_strides [[buffer(8)]],
    const constant bool* idx_contigs [[buffer(9)]],
    const constant int& idx_ndim [[buffer(10)]],
    const device uint32_t *idx0 [[buffer(20)]],
    uint3 index [[thread_position_in_grid]],
    uint3 grid_dim [[threads_per_grid]]) {
  Indices<uint32_t, 1> idxs{
    { idx0 }, idx_shapes, idx_strides, idx_contigs, idx_ndim};

  return gather_impl<uint32_t, uint32_t, 1, 1, int>(
      src,
      out,
      src_shape,
      src_strides,
      src_ndim,
      slice_sizes,
      axes,
      idxs,
      index,
      grid_dim);
}

[[kernel]] void kernel_qwen_mtplx_gather_u32_idxu32_1_1_int64_t(
    const device uint32_t* src [[buffer(0)]],
    device uint32_t* out [[buffer(1)]],
    const constant int* src_shape [[buffer(2)]],
    const constant int64_t* src_strides [[buffer(3)]],
    const constant size_t& src_ndim [[buffer(4)]],
    const constant int* slice_sizes [[buffer(5)]],
    const constant int* axes [[buffer(6)]],
    const constant int* idx_shapes [[buffer(7)]],
    const constant int64_t* idx_strides [[buffer(8)]],
    const constant bool* idx_contigs [[buffer(9)]],
    const constant int& idx_ndim [[buffer(10)]],
    const device uint32_t *idx0 [[buffer(20)]],
    uint3 index [[thread_position_in_grid]],
    uint3 grid_dim [[threads_per_grid]]) {
  Indices<uint32_t, 1> idxs{
    { idx0 }, idx_shapes, idx_strides, idx_contigs, idx_ndim};

  return gather_impl<uint32_t, uint32_t, 1, 1, int64_t>(
      src,
      out,
      src_shape,
      src_strides,
      src_ndim,
      slice_sizes,
      axes,
      idxs,
      index,
      grid_dim);
}

[[kernel]] void kernel_qwen_mtplx_gather_u32_idxu32_1_2_int(
    const device uint32_t* src [[buffer(0)]],
    device uint32_t* out [[buffer(1)]],
    const constant int* src_shape [[buffer(2)]],
    const constant int64_t* src_strides [[buffer(3)]],
    const constant size_t& src_ndim [[buffer(4)]],
    const constant int* slice_sizes [[buffer(5)]],
    const constant int* axes [[buffer(6)]],
    const constant int* idx_shapes [[buffer(7)]],
    const constant int64_t* idx_strides [[buffer(8)]],
    const constant bool* idx_contigs [[buffer(9)]],
    const constant int& idx_ndim [[buffer(10)]],
    const device uint32_t *idx0 [[buffer(20)]],
    uint3 index [[thread_position_in_grid]],
    uint3 grid_dim [[threads_per_grid]]) {
  Indices<uint32_t, 1> idxs{
    { idx0 }, idx_shapes, idx_strides, idx_contigs, idx_ndim};

  return gather_impl<uint32_t, uint32_t, 1, 2, int>(
      src,
      out,
      src_shape,
      src_strides,
      src_ndim,
      slice_sizes,
      axes,
      idxs,
      index,
      grid_dim);
}

[[kernel]] void kernel_qwen_mtplx_gather_u32_idxu32_1_2_int64_t(
    const device uint32_t* src [[buffer(0)]],
    device uint32_t* out [[buffer(1)]],
    const constant int* src_shape [[buffer(2)]],
    const constant int64_t* src_strides [[buffer(3)]],
    const constant size_t& src_ndim [[buffer(4)]],
    const constant int* slice_sizes [[buffer(5)]],
    const constant int* axes [[buffer(6)]],
    const constant int* idx_shapes [[buffer(7)]],
    const constant int64_t* idx_strides [[buffer(8)]],
    const constant bool* idx_contigs [[buffer(9)]],
    const constant int& idx_ndim [[buffer(10)]],
    const device uint32_t *idx0 [[buffer(20)]],
    uint3 index [[thread_position_in_grid]],
    uint3 grid_dim [[threads_per_grid]]) {
  Indices<uint32_t, 1> idxs{
    { idx0 }, idx_shapes, idx_strides, idx_contigs, idx_ndim};

  return gather_impl<uint32_t, uint32_t, 1, 2, int64_t>(
      src,
      out,
      src_shape,
      src_strides,
      src_ndim,
      slice_sizes,
      axes,
      idxs,
      index,
      grid_dim);
}

[[kernel]] void kernel_qwen_mtplx_gather_u32_idxu32_1_3_int(
    const device uint32_t* src [[buffer(0)]],
    device uint32_t* out [[buffer(1)]],
    const constant int* src_shape [[buffer(2)]],
    const constant int64_t* src_strides [[buffer(3)]],
    const constant size_t& src_ndim [[buffer(4)]],
    const constant int* slice_sizes [[buffer(5)]],
    const constant int* axes [[buffer(6)]],
    const constant int* idx_shapes [[buffer(7)]],
    const constant int64_t* idx_strides [[buffer(8)]],
    const constant bool* idx_contigs [[buffer(9)]],
    const constant int& idx_ndim [[buffer(10)]],
    const device uint32_t *idx0 [[buffer(20)]],
    uint3 index [[thread_position_in_grid]],
    uint3 grid_dim [[threads_per_grid]]) {
  Indices<uint32_t, 1> idxs{
    { idx0 }, idx_shapes, idx_strides, idx_contigs, idx_ndim};

  return gather_impl<uint32_t, uint32_t, 1, 3, int>(
      src,
      out,
      src_shape,
      src_strides,
      src_ndim,
      slice_sizes,
      axes,
      idxs,
      index,
      grid_dim);
}

[[kernel]] void kernel_qwen_mtplx_gather_u32_idxu32_1_3_int64_t(
    const device uint32_t* src [[buffer(0)]],
    device uint32_t* out [[buffer(1)]],
    const constant int* src_shape [[buffer(2)]],
    const constant int64_t* src_strides [[buffer(3)]],
    const constant size_t& src_ndim [[buffer(4)]],
    const constant int* slice_sizes [[buffer(5)]],
    const constant int* axes [[buffer(6)]],
    const constant int* idx_shapes [[buffer(7)]],
    const constant int64_t* idx_strides [[buffer(8)]],
    const constant bool* idx_contigs [[buffer(9)]],
    const constant int& idx_ndim [[buffer(10)]],
    const device uint32_t *idx0 [[buffer(20)]],
    uint3 index [[thread_position_in_grid]],
    uint3 grid_dim [[threads_per_grid]]) {
  Indices<uint32_t, 1> idxs{
    { idx0 }, idx_shapes, idx_strides, idx_contigs, idx_ndim};

  return gather_impl<uint32_t, uint32_t, 1, 3, int64_t>(
      src,
      out,
      src_shape,
      src_strides,
      src_ndim,
      slice_sizes,
      axes,
      idxs,
      index,
      grid_dim);
}

[[kernel]] void kernel_qwen_mtplx_gather_u32_idxu32_1_4_int(
    const device uint32_t* src [[buffer(0)]],
    device uint32_t* out [[buffer(1)]],
    const constant int* src_shape [[buffer(2)]],
    const constant int64_t* src_strides [[buffer(3)]],
    const constant size_t& src_ndim [[buffer(4)]],
    const constant int* slice_sizes [[buffer(5)]],
    const constant int* axes [[buffer(6)]],
    const constant int* idx_shapes [[buffer(7)]],
    const constant int64_t* idx_strides [[buffer(8)]],
    const constant bool* idx_contigs [[buffer(9)]],
    const constant int& idx_ndim [[buffer(10)]],
    const device uint32_t *idx0 [[buffer(20)]],
    uint3 index [[thread_position_in_grid]],
    uint3 grid_dim [[threads_per_grid]]) {
  Indices<uint32_t, 1> idxs{
    { idx0 }, idx_shapes, idx_strides, idx_contigs, idx_ndim};

  return gather_impl<uint32_t, uint32_t, 1, 4, int>(
      src,
      out,
      src_shape,
      src_strides,
      src_ndim,
      slice_sizes,
      axes,
      idxs,
      index,
      grid_dim);
}

[[kernel]] void kernel_qwen_mtplx_gather_u32_idxu32_1_4_int64_t(
    const device uint32_t* src [[buffer(0)]],
    device uint32_t* out [[buffer(1)]],
    const constant int* src_shape [[buffer(2)]],
    const constant int64_t* src_strides [[buffer(3)]],
    const constant size_t& src_ndim [[buffer(4)]],
    const constant int* slice_sizes [[buffer(5)]],
    const constant int* axes [[buffer(6)]],
    const constant int* idx_shapes [[buffer(7)]],
    const constant int64_t* idx_strides [[buffer(8)]],
    const constant bool* idx_contigs [[buffer(9)]],
    const constant int& idx_ndim [[buffer(10)]],
    const device uint32_t *idx0 [[buffer(20)]],
    uint3 index [[thread_position_in_grid]],
    uint3 grid_dim [[threads_per_grid]]) {
  Indices<uint32_t, 1> idxs{
    { idx0 }, idx_shapes, idx_strides, idx_contigs, idx_ndim};

  return gather_impl<uint32_t, uint32_t, 1, 4, int64_t>(
      src,
      out,
      src_shape,
      src_strides,
      src_ndim,
      slice_sizes,
      axes,
      idxs,
      index,
      grid_dim);
}

[[kernel]] void kernel_qwen_mtplx_gather_f32_idxu32_1_0_int(
    const device float* src [[buffer(0)]],
    device float* out [[buffer(1)]],
    const constant int* src_shape [[buffer(2)]],
    const constant int64_t* src_strides [[buffer(3)]],
    const constant size_t& src_ndim [[buffer(4)]],
    const constant int* slice_sizes [[buffer(5)]],
    const constant int* axes [[buffer(6)]],
    const constant int* idx_shapes [[buffer(7)]],
    const constant int64_t* idx_strides [[buffer(8)]],
    const constant bool* idx_contigs [[buffer(9)]],
    const constant int& idx_ndim [[buffer(10)]],
    const device uint32_t *idx0 [[buffer(20)]],
    uint3 index [[thread_position_in_grid]],
    uint3 grid_dim [[threads_per_grid]]) {
  Indices<uint32_t, 1> idxs{
    { idx0 }, idx_shapes, idx_strides, idx_contigs, idx_ndim};

  return gather_impl<float, uint32_t, 1, 0, int>(
      src,
      out,
      src_shape,
      src_strides,
      src_ndim,
      slice_sizes,
      axes,
      idxs,
      index,
      grid_dim);
}

[[kernel]] void kernel_qwen_mtplx_gather_f32_idxu32_1_0_int64_t(
    const device float* src [[buffer(0)]],
    device float* out [[buffer(1)]],
    const constant int* src_shape [[buffer(2)]],
    const constant int64_t* src_strides [[buffer(3)]],
    const constant size_t& src_ndim [[buffer(4)]],
    const constant int* slice_sizes [[buffer(5)]],
    const constant int* axes [[buffer(6)]],
    const constant int* idx_shapes [[buffer(7)]],
    const constant int64_t* idx_strides [[buffer(8)]],
    const constant bool* idx_contigs [[buffer(9)]],
    const constant int& idx_ndim [[buffer(10)]],
    const device uint32_t *idx0 [[buffer(20)]],
    uint3 index [[thread_position_in_grid]],
    uint3 grid_dim [[threads_per_grid]]) {
  Indices<uint32_t, 1> idxs{
    { idx0 }, idx_shapes, idx_strides, idx_contigs, idx_ndim};

  return gather_impl<float, uint32_t, 1, 0, int64_t>(
      src,
      out,
      src_shape,
      src_strides,
      src_ndim,
      slice_sizes,
      axes,
      idxs,
      index,
      grid_dim);
}

[[kernel]] void kernel_qwen_mtplx_gather_f32_idxu32_1_1_int(
    const device float* src [[buffer(0)]],
    device float* out [[buffer(1)]],
    const constant int* src_shape [[buffer(2)]],
    const constant int64_t* src_strides [[buffer(3)]],
    const constant size_t& src_ndim [[buffer(4)]],
    const constant int* slice_sizes [[buffer(5)]],
    const constant int* axes [[buffer(6)]],
    const constant int* idx_shapes [[buffer(7)]],
    const constant int64_t* idx_strides [[buffer(8)]],
    const constant bool* idx_contigs [[buffer(9)]],
    const constant int& idx_ndim [[buffer(10)]],
    const device uint32_t *idx0 [[buffer(20)]],
    uint3 index [[thread_position_in_grid]],
    uint3 grid_dim [[threads_per_grid]]) {
  Indices<uint32_t, 1> idxs{
    { idx0 }, idx_shapes, idx_strides, idx_contigs, idx_ndim};

  return gather_impl<float, uint32_t, 1, 1, int>(
      src,
      out,
      src_shape,
      src_strides,
      src_ndim,
      slice_sizes,
      axes,
      idxs,
      index,
      grid_dim);
}

[[kernel]] void kernel_qwen_mtplx_gather_f32_idxu32_1_1_int64_t(
    const device float* src [[buffer(0)]],
    device float* out [[buffer(1)]],
    const constant int* src_shape [[buffer(2)]],
    const constant int64_t* src_strides [[buffer(3)]],
    const constant size_t& src_ndim [[buffer(4)]],
    const constant int* slice_sizes [[buffer(5)]],
    const constant int* axes [[buffer(6)]],
    const constant int* idx_shapes [[buffer(7)]],
    const constant int64_t* idx_strides [[buffer(8)]],
    const constant bool* idx_contigs [[buffer(9)]],
    const constant int& idx_ndim [[buffer(10)]],
    const device uint32_t *idx0 [[buffer(20)]],
    uint3 index [[thread_position_in_grid]],
    uint3 grid_dim [[threads_per_grid]]) {
  Indices<uint32_t, 1> idxs{
    { idx0 }, idx_shapes, idx_strides, idx_contigs, idx_ndim};

  return gather_impl<float, uint32_t, 1, 1, int64_t>(
      src,
      out,
      src_shape,
      src_strides,
      src_ndim,
      slice_sizes,
      axes,
      idxs,
      index,
      grid_dim);
}

[[kernel]] void kernel_qwen_mtplx_gather_f32_idxu32_1_2_int(
    const device float* src [[buffer(0)]],
    device float* out [[buffer(1)]],
    const constant int* src_shape [[buffer(2)]],
    const constant int64_t* src_strides [[buffer(3)]],
    const constant size_t& src_ndim [[buffer(4)]],
    const constant int* slice_sizes [[buffer(5)]],
    const constant int* axes [[buffer(6)]],
    const constant int* idx_shapes [[buffer(7)]],
    const constant int64_t* idx_strides [[buffer(8)]],
    const constant bool* idx_contigs [[buffer(9)]],
    const constant int& idx_ndim [[buffer(10)]],
    const device uint32_t *idx0 [[buffer(20)]],
    uint3 index [[thread_position_in_grid]],
    uint3 grid_dim [[threads_per_grid]]) {
  Indices<uint32_t, 1> idxs{
    { idx0 }, idx_shapes, idx_strides, idx_contigs, idx_ndim};

  return gather_impl<float, uint32_t, 1, 2, int>(
      src,
      out,
      src_shape,
      src_strides,
      src_ndim,
      slice_sizes,
      axes,
      idxs,
      index,
      grid_dim);
}

[[kernel]] void kernel_qwen_mtplx_gather_f32_idxu32_1_2_int64_t(
    const device float* src [[buffer(0)]],
    device float* out [[buffer(1)]],
    const constant int* src_shape [[buffer(2)]],
    const constant int64_t* src_strides [[buffer(3)]],
    const constant size_t& src_ndim [[buffer(4)]],
    const constant int* slice_sizes [[buffer(5)]],
    const constant int* axes [[buffer(6)]],
    const constant int* idx_shapes [[buffer(7)]],
    const constant int64_t* idx_strides [[buffer(8)]],
    const constant bool* idx_contigs [[buffer(9)]],
    const constant int& idx_ndim [[buffer(10)]],
    const device uint32_t *idx0 [[buffer(20)]],
    uint3 index [[thread_position_in_grid]],
    uint3 grid_dim [[threads_per_grid]]) {
  Indices<uint32_t, 1> idxs{
    { idx0 }, idx_shapes, idx_strides, idx_contigs, idx_ndim};

  return gather_impl<float, uint32_t, 1, 2, int64_t>(
      src,
      out,
      src_shape,
      src_strides,
      src_ndim,
      slice_sizes,
      axes,
      idxs,
      index,
      grid_dim);
}

[[kernel]] void kernel_qwen_mtplx_gather_f32_idxu32_1_3_int(
    const device float* src [[buffer(0)]],
    device float* out [[buffer(1)]],
    const constant int* src_shape [[buffer(2)]],
    const constant int64_t* src_strides [[buffer(3)]],
    const constant size_t& src_ndim [[buffer(4)]],
    const constant int* slice_sizes [[buffer(5)]],
    const constant int* axes [[buffer(6)]],
    const constant int* idx_shapes [[buffer(7)]],
    const constant int64_t* idx_strides [[buffer(8)]],
    const constant bool* idx_contigs [[buffer(9)]],
    const constant int& idx_ndim [[buffer(10)]],
    const device uint32_t *idx0 [[buffer(20)]],
    uint3 index [[thread_position_in_grid]],
    uint3 grid_dim [[threads_per_grid]]) {
  Indices<uint32_t, 1> idxs{
    { idx0 }, idx_shapes, idx_strides, idx_contigs, idx_ndim};

  return gather_impl<float, uint32_t, 1, 3, int>(
      src,
      out,
      src_shape,
      src_strides,
      src_ndim,
      slice_sizes,
      axes,
      idxs,
      index,
      grid_dim);
}

[[kernel]] void kernel_qwen_mtplx_gather_f32_idxu32_1_3_int64_t(
    const device float* src [[buffer(0)]],
    device float* out [[buffer(1)]],
    const constant int* src_shape [[buffer(2)]],
    const constant int64_t* src_strides [[buffer(3)]],
    const constant size_t& src_ndim [[buffer(4)]],
    const constant int* slice_sizes [[buffer(5)]],
    const constant int* axes [[buffer(6)]],
    const constant int* idx_shapes [[buffer(7)]],
    const constant int64_t* idx_strides [[buffer(8)]],
    const constant bool* idx_contigs [[buffer(9)]],
    const constant int& idx_ndim [[buffer(10)]],
    const device uint32_t *idx0 [[buffer(20)]],
    uint3 index [[thread_position_in_grid]],
    uint3 grid_dim [[threads_per_grid]]) {
  Indices<uint32_t, 1> idxs{
    { idx0 }, idx_shapes, idx_strides, idx_contigs, idx_ndim};

  return gather_impl<float, uint32_t, 1, 3, int64_t>(
      src,
      out,
      src_shape,
      src_strides,
      src_ndim,
      slice_sizes,
      axes,
      idxs,
      index,
      grid_dim);
}

[[kernel]] void kernel_qwen_mtplx_gather_f32_idxu32_1_4_int(
    const device float* src [[buffer(0)]],
    device float* out [[buffer(1)]],
    const constant int* src_shape [[buffer(2)]],
    const constant int64_t* src_strides [[buffer(3)]],
    const constant size_t& src_ndim [[buffer(4)]],
    const constant int* slice_sizes [[buffer(5)]],
    const constant int* axes [[buffer(6)]],
    const constant int* idx_shapes [[buffer(7)]],
    const constant int64_t* idx_strides [[buffer(8)]],
    const constant bool* idx_contigs [[buffer(9)]],
    const constant int& idx_ndim [[buffer(10)]],
    const device uint32_t *idx0 [[buffer(20)]],
    uint3 index [[thread_position_in_grid]],
    uint3 grid_dim [[threads_per_grid]]) {
  Indices<uint32_t, 1> idxs{
    { idx0 }, idx_shapes, idx_strides, idx_contigs, idx_ndim};

  return gather_impl<float, uint32_t, 1, 4, int>(
      src,
      out,
      src_shape,
      src_strides,
      src_ndim,
      slice_sizes,
      axes,
      idxs,
      index,
      grid_dim);
}

[[kernel]] void kernel_qwen_mtplx_gather_f32_idxu32_1_4_int64_t(
    const device float* src [[buffer(0)]],
    device float* out [[buffer(1)]],
    const constant int* src_shape [[buffer(2)]],
    const constant int64_t* src_strides [[buffer(3)]],
    const constant size_t& src_ndim [[buffer(4)]],
    const constant int* slice_sizes [[buffer(5)]],
    const constant int* axes [[buffer(6)]],
    const constant int* idx_shapes [[buffer(7)]],
    const constant int64_t* idx_strides [[buffer(8)]],
    const constant bool* idx_contigs [[buffer(9)]],
    const constant int& idx_ndim [[buffer(10)]],
    const device uint32_t *idx0 [[buffer(20)]],
    uint3 index [[thread_position_in_grid]],
    uint3 grid_dim [[threads_per_grid]]) {
  Indices<uint32_t, 1> idxs{
    { idx0 }, idx_shapes, idx_strides, idx_contigs, idx_ndim};

  return gather_impl<float, uint32_t, 1, 4, int64_t>(
      src,
      out,
      src_shape,
      src_strides,
      src_ndim,
      slice_sizes,
      axes,
      idxs,
      index,
      grid_dim);
}
// END RUNTIME UNIT

// Runtime ScatterAxis: Copyright Apple Inc. SPDX-License-Identifier: MIT
#include <metal_atomic>
namespace mtplx_eager_scatter {
// Runtime unit: complex.h; file SHA256: 16e8a815b2cbdb6070e0824e64fe33fccb6e918f1b84ea5c792bd89d33e57bf1
// Runtime unit SHA256: 0393eb41ea618f7b09343672f598d8c4c3e4b98eb8536466fb53c8d286fbf1f7
// BEGIN RUNTIME UNIT
// Copyright © 2023 Apple Inc.




using namespace metal;

template <typename T>
struct complex_t;

template <typename T>
static constexpr constant bool is_complex_v = false;

template <typename T>
static constexpr constant bool is_complex_v<complex_t<T>> = true;

// Metal accepts explicit bfloat casts that is_convertible_v reports as false.
template <typename From, typename To>
static constexpr constant bool is_lane_convertible_v =
    is_convertible_v<From, To> ||
    (is_same_v<To, bfloat16_t> && is_convertible_v<From, float>) ||
    (is_same_v<From, bfloat16_t> && is_convertible_v<float, To>);

template <typename T>
struct complex_t {
  using value_type = T;

  T real;
  T imag;

  // Constructors
  constexpr complex_t(T real, T imag) thread : real(real), imag(imag) {};
  constexpr complex_t() thread : real(0), imag(0) {};
  constexpr complex_t() threadgroup : real(0), imag(0) {};

  // Conversions from scalar types
  template <
      typename U,
      typename = typename enable_if<
          !is_complex_v<U> && is_lane_convertible_v<U, T>>::type>
  constexpr complex_t(U x) thread : real(static_cast<T>(x)),
                                    imag(static_cast<T>(0)) {}

  template <
      typename U,
      typename = typename enable_if<
          !is_complex_v<U> && is_lane_convertible_v<U, T>>::type>
  constexpr complex_t(U x) threadgroup : real(static_cast<T>(x)),
                                         imag(static_cast<T>(0)) {}

  template <
      typename U,
      typename = typename enable_if<
          !is_complex_v<U> && is_lane_convertible_v<U, T>>::type>
  constexpr complex_t(U x) device : real(static_cast<T>(x)),
                                    imag(static_cast<T>(0)) {}

  template <
      typename U,
      typename = typename enable_if<
          !is_complex_v<U> && is_lane_convertible_v<U, T>>::type>
  constexpr complex_t(U x) constant : real(static_cast<T>(x)),
                                      imag(static_cast<T>(0)) {}

  // Conversions between complex types
  template <
      typename U,
      typename = typename enable_if<
          !is_same_v<U, T> && is_lane_convertible_v<U, T>>::type>
  constexpr complex_t(complex_t<U> x) thread : real(static_cast<T>(x.real)),
                                               imag(static_cast<T>(x.imag)) {}

  template <
      typename U,
      typename = typename enable_if<
          !is_same_v<U, T> && is_lane_convertible_v<U, T>>::type>
  constexpr complex_t(complex_t<U> x) threadgroup
      : real(static_cast<T>(x.real)),
        imag(static_cast<T>(x.imag)) {}

  template <
      typename U,
      typename = typename enable_if<
          !is_same_v<U, T> && is_lane_convertible_v<U, T>>::type>
  constexpr complex_t(complex_t<U> x) device : real(static_cast<T>(x.real)),
                                               imag(static_cast<T>(x.imag)) {}

  template <
      typename U,
      typename = typename enable_if<
          !is_same_v<U, T> && is_lane_convertible_v<U, T>>::type>
  constexpr complex_t(complex_t<U> x) constant : real(static_cast<T>(x.real)),
                                                 imag(static_cast<T>(x.imag)) {}

  // Conversions to and from two-lane vectors (the FFT lane representation)
  constexpr complex_t(vec<T, 2> v) thread : real(v.x), imag(v.y) {};
  constexpr complex_t(vec<T, 2> v) threadgroup : real(v.x), imag(v.y) {};
  constexpr complex_t(vec<T, 2> v) device : real(v.x), imag(v.y) {};
  constexpr complex_t(vec<T, 2> v) constant : real(v.x), imag(v.y) {};

  constexpr operator vec<T, 2>() const thread {
    return vec<T, 2>(real, imag);
  }

  constexpr operator vec<T, 2>() const threadgroup {
    return vec<T, 2>(real, imag);
  }

  constexpr operator vec<T, 2>() const device {
    return vec<T, 2>(real, imag);
  }

  constexpr operator vec<T, 2>() const constant {
    return vec<T, 2>(real, imag);
  }

  // Conversions to scalar types
  template <
      typename U,
      typename = typename enable_if<
          !is_complex_v<U> && is_lane_convertible_v<T, U>>::type>
  constexpr operator U() const thread {
    return static_cast<U>(real);
  }

  template <
      typename U,
      typename = typename enable_if<
          !is_complex_v<U> && is_lane_convertible_v<T, U>>::type>
  constexpr operator U() const threadgroup {
    return static_cast<U>(real);
  }

  template <
      typename U,
      typename = typename enable_if<
          !is_complex_v<U> && is_lane_convertible_v<T, U>>::type>
  constexpr operator U() const device {
    return static_cast<U>(real);
  }

  template <
      typename U,
      typename = typename enable_if<
          !is_complex_v<U> && is_lane_convertible_v<T, U>>::type>
  constexpr operator U() const constant {
    return static_cast<U>(real);
  }
};

using complex32_t = complex_t<half>;
using complex64_t = complex_t<float>;

static_assert(sizeof(complex32_t) == 2 * sizeof(half));
static_assert(sizeof(complex64_t) == 2 * sizeof(float));
static_assert(sizeof(complex_t<bfloat16_t>) == 2 * sizeof(bfloat16_t));

template <typename T>
constexpr complex_t<T> operator-(complex_t<T> x) {
  return {-x.real, -x.imag};
}

template <typename T>
constexpr bool operator>=(complex_t<T> a, complex_t<T> b) {
  return (a.real > b.real) || (a.real == b.real && a.imag >= b.imag);
}

template <typename T>
constexpr bool operator>(complex_t<T> a, complex_t<T> b) {
  return (a.real > b.real) || (a.real == b.real && a.imag > b.imag);
}

template <typename T>
constexpr bool operator<=(complex_t<T> a, complex_t<T> b) {
  return operator>=(b, a);
}

template <typename T>
constexpr bool operator<(complex_t<T> a, complex_t<T> b) {
  return operator>(b, a);
}

template <typename T>
constexpr bool operator==(complex_t<T> a, complex_t<T> b) {
  return a.real == b.real && a.imag == b.imag;
}

template <typename T>
constexpr complex_t<T> operator+(complex_t<T> a, complex_t<T> b) {
  return {a.real + b.real, a.imag + b.imag};
}

template <typename T>
constexpr thread complex_t<T>& operator+=(
    thread complex_t<T>& a,
    complex_t<T> b) {
  a.real += b.real;
  a.imag += b.imag;
  return a;
}

template <typename T>
constexpr threadgroup complex_t<T>& operator+=(
    threadgroup complex_t<T>& a,
    complex_t<T> b) {
  a.real += b.real;
  a.imag += b.imag;
  return a;
}

template <typename T>
constexpr device complex_t<T>& operator+=(
    device complex_t<T>& a,
    complex_t<T> b) {
  a.real += b.real;
  a.imag += b.imag;
  return a;
}

template <
    typename T,
    typename U,
    enable_if_t<!is_complex_v<U> && is_lane_convertible_v<U, T>, bool> = true>
constexpr complex_t<T> operator+(U a, complex_t<T> b) {
  return {static_cast<T>(a) + b.real, b.imag};
}

template <
    typename T,
    typename U,
    enable_if_t<!is_complex_v<U> && is_lane_convertible_v<U, T>, bool> = true>
constexpr complex_t<T> operator+(complex_t<T> a, U b) {
  return {a.real + static_cast<T>(b), a.imag};
}

template <typename T>
constexpr complex_t<T> operator-(complex_t<T> a, complex_t<T> b) {
  return {a.real - b.real, a.imag - b.imag};
}

template <
    typename T,
    typename U,
    enable_if_t<!is_complex_v<U> && is_lane_convertible_v<U, T>, bool> = true>
constexpr complex_t<T> operator-(U a, complex_t<T> b) {
  return {static_cast<T>(a) - b.real, -b.imag};
}

template <
    typename T,
    typename U,
    enable_if_t<!is_complex_v<U> && is_lane_convertible_v<U, T>, bool> = true>
constexpr complex_t<T> operator-(complex_t<T> a, U b) {
  return {a.real - static_cast<T>(b), a.imag};
}

template <typename T>
constexpr complex_t<T> operator*(complex_t<T> a, complex_t<T> b) {
  return {a.real * b.real - a.imag * b.imag, a.real * b.imag + a.imag * b.real};
}

template <typename T>
constexpr complex_t<T> operator/(complex_t<T> a, complex_t<T> b) {
  auto denom = b.real * b.real + b.imag * b.imag;
  auto x = a.real * b.real + a.imag * b.imag;
  auto y = a.imag * b.real - a.real * b.imag;
  return {x / denom, y / denom};
}

template <
    typename T,
    typename U,
    enable_if_t<!is_complex_v<U> && is_lane_convertible_v<U, T>, bool> = true>
constexpr complex_t<T> operator/(U a, complex_t<T> b) {
  auto scalar = static_cast<T>(a);
  auto denom = b.real * b.real + b.imag * b.imag;
  auto x = scalar * b.real;
  auto y = -scalar * b.imag;
  return {x / denom, y / denom};
}

template <typename T>
constexpr complex_t<T> operator%(complex_t<T> a, complex_t<T> b) {
  auto real = a.real - (b.real * static_cast<int64_t>(a.real / b.real));
  auto imag = a.imag - (b.imag * static_cast<int64_t>(a.imag / b.imag));
  if (real != 0 && (real < 0 != b.real < 0)) {
    real += b.real;
  }
  if (imag != 0 && (imag < 0 != b.imag < 0)) {
    imag += b.imag;
  }
  return {real, imag};
}

static_assert(
    (complex_t<half>{1.0h, 2.0h} * complex_t<half>{3.0h, 4.0h}).real == -5.0h);
static_assert(
    (complex_t<bfloat16_t>{bfloat16_t(1.0f), bfloat16_t(2.0f)} *
     complex_t<bfloat16_t>{bfloat16_t(3.0f), bfloat16_t(4.0f)})
        .real == bfloat16_t(-5.0f));
// END RUNTIME UNIT
// Runtime unit: atomic.h; file SHA256: 4c35ea2798a2335502865247aee878149fc9ada0d7e84c05d771baef0c7fcc60
// Runtime unit SHA256: 7c800c32a72badbf76980e746956577d168dd396e002f6cc2eaa83e90b913457
// BEGIN RUNTIME UNIT
// Copyright © 2023 Apple Inc.



using namespace metal;

///////////////////////////////////////////////////////////////////////////////
// Atomic utils
///////////////////////////////////////////////////////////////////////////////

#pragma METAL internals : enable
template <typename T>
constexpr constant bool is_metal_atomic = _disjunction<
    is_same<T, int>,
    is_same<T, uint>,
    is_same<T, ulong>,
    is_same<T, float>>::value;

#pragma METAL internals : disable

template <typename T, typename = void>
struct mlx_atomic {
  atomic<uint> val;
};

template <typename T>
struct mlx_atomic<T, enable_if_t<is_metal_atomic<T>>> {
  atomic<T> val;
};

///////////////////////////////////////////////////////////////////////////////
// Native metal atomics
///////////////////////////////////////////////////////////////////////////////

template <typename T, enable_if_t<is_metal_atomic<T>, bool> = true>
METAL_FUNC T
mlx_atomic_load_explicit(device mlx_atomic<T>* object, size_t offset) {
  return atomic_load_explicit(&(object[offset].val), memory_order_relaxed);
}

template <typename T, enable_if_t<is_metal_atomic<T>, bool> = true>
METAL_FUNC void
mlx_atomic_store_explicit(device mlx_atomic<T>* object, T val, size_t offset) {
  atomic_store_explicit(&(object[offset].val), val, memory_order_relaxed);
}

template <typename T, enable_if_t<is_metal_atomic<T>, bool> = true>
METAL_FUNC void mlx_atomic_fetch_and_explicit(
    device mlx_atomic<T>* object,
    T val,
    size_t offset) {
  atomic_fetch_and_explicit(&(object[offset].val), val, memory_order_relaxed);
}

template <typename T, enable_if_t<is_metal_atomic<T>, bool> = true>
METAL_FUNC void mlx_atomic_fetch_or_explicit(
    device mlx_atomic<T>* object,
    T val,
    size_t offset) {
  atomic_fetch_or_explicit(&(object[offset].val), val, memory_order_relaxed);
}

template <typename T, enable_if_t<is_metal_atomic<T>, bool> = true>
METAL_FUNC void mlx_atomic_fetch_min_explicit(
    device mlx_atomic<T>* object,
    T val,
    size_t offset) {
  atomic_fetch_min_explicit(&(object[offset].val), val, memory_order_relaxed);
}

template <typename T, enable_if_t<is_metal_atomic<T>, bool> = true>
METAL_FUNC void mlx_atomic_fetch_max_explicit(
    device mlx_atomic<T>* object,
    T val,
    size_t offset) {
  atomic_fetch_max_explicit(&(object[offset].val), val, memory_order_relaxed);
}

template <typename T, enable_if_t<is_metal_atomic<T>, bool> = true>
METAL_FUNC void mlx_atomic_fetch_add_explicit(
    device mlx_atomic<T>* object,
    T val,
    size_t offset) {
  atomic_fetch_add_explicit(&(object[offset].val), val, memory_order_relaxed);
}

template <typename T, enable_if_t<is_metal_atomic<T>, bool> = true>
METAL_FUNC void mlx_atomic_fetch_mul_explicit(
    device mlx_atomic<T>* object,
    T val,
    size_t offset) {
  T expected = mlx_atomic_load_explicit(object, offset);
  while (!mlx_atomic_compare_exchange_weak_explicit(
      object, &expected, val * expected, offset)) {
    // Workaround: Metal's atomic_compare_exchange_weak_explicit<float> does
    // not perform bitwise comparison as required by the C++ atomics spec.
    // The compiler lowers the success check to `fcmp fast ueq` under
    // no-nans-fp-math, which evaluates to false when either operand is NaN -
    // even when the bit patterns are identical. With NaN in memory the CAS
    // can never succeed, so the loop spins. Bail out instead: memory is
    // already NaN and that is the correct reduction result regardless of
    // this thread's update.
    if constexpr (metal::is_floating_point_v<T>) {
      if (isnan(expected)) {
        break;
      }
    }
  }
}

template <typename T, enable_if_t<is_metal_atomic<T>, bool> = true>
METAL_FUNC bool mlx_atomic_compare_exchange_weak_explicit(
    device mlx_atomic<T>* object,
    thread T* expected,
    T val,
    size_t offset) {
  return atomic_compare_exchange_weak_explicit(
      &(object[offset].val),
      expected,
      val,
      memory_order_relaxed,
      memory_order_relaxed);
}

// Specialization for float since it does not atomic_fetch_min_explicit
template <>
METAL_FUNC void mlx_atomic_fetch_min_explicit<float>(
    device mlx_atomic<float>* object,
    float val,
    size_t offset) {
  float expected = mlx_atomic_load_explicit(object, offset);
  while (val < expected) {
    if (mlx_atomic_compare_exchange_weak_explicit(
            object, &expected, val, offset)) {
      return;
    }
  }
}

// Specialization for float since it does not atomic_fetch_max_explicit
template <>
METAL_FUNC void mlx_atomic_fetch_max_explicit<float>(
    device mlx_atomic<float>* object,
    float val,
    size_t offset) {
  float expected = mlx_atomic_load_explicit(object, offset);
  while (val > expected) {
    if (mlx_atomic_compare_exchange_weak_explicit(
            object, &expected, val, offset)) {
      return;
    }
  }
}

///////////////////////////////////////////////////////////////////////////////
// Custom atomics
///////////////////////////////////////////////////////////////////////////////

namespace {

template <typename T>
constexpr constant uint packing_size = sizeof(uint) / sizeof(T);

template <typename T>
union uint_or_packed {
  T val[packing_size<T>];
  uint bits;
};

template <typename T, typename Op>
struct mlx_atomic_update_helper {
  uint operator()(uint_or_packed<T> init, T update, size_t elem_offset) thread {
    Op op;
    init.val[elem_offset] = op(update, init.val[elem_offset]);
    return init.bits;
  }
};

template <typename T, typename Op>
METAL_FUNC void mlx_atomic_update_and_store(
    device mlx_atomic<T>* object,
    T update,
    size_t offset) {
  size_t pack_offset = offset / packing_size<T>;
  size_t elem_offset = offset % packing_size<T>;

  mlx_atomic_update_helper<T, Op> helper;
  uint_or_packed<T> expected;
  expected.bits =
      atomic_load_explicit(&(object[pack_offset].val), memory_order_relaxed);

  while (Op::condition(update, expected.val[elem_offset]) &&
         !mlx_atomic_compare_exchange_weak_explicit(
             object,
             &(expected.bits),
             helper(expected, update, elem_offset),
             pack_offset)) {
  }
}

template <typename T>
struct __None {
  static bool condition(T a, T b) {
#pragma unused(a)
#pragma unused(b)
    return true;
  }

  T operator()(T a, T b) thread {
#pragma unused(b)
    return a;
  }
};

template <typename T>
struct __Add {
  static bool condition(T a, T b) {
#pragma unused(a)
#pragma unused(b)
    return true;
  }

  T operator()(T a, T b) thread {
    return a + b;
  }
};

template <typename T>
struct __Mul {
  static bool condition(T a, T b) {
#pragma unused(a)
    return b != 0;
  }

  T operator()(T a, T b) thread {
    return a * b;
  }
};

template <typename T>
struct __Max {
  static bool condition(T a, T b) {
    return a > b;
  }

  T operator()(T a, T b) thread {
    return max(a, b);
  }
};

template <typename T>
struct __Min {
  static bool condition(T a, T b) {
    return a < b;
  }

  T operator()(T a, T b) thread {
    return min(a, b);
  }
};

} // namespace

template <typename T, enable_if_t<!is_metal_atomic<T>, bool> = true>
METAL_FUNC T
mlx_atomic_load_explicit(device mlx_atomic<T>* object, size_t offset) {
  size_t pack_offset = offset / sizeof(T);
  size_t elem_offset = offset % sizeof(T);
  uint_or_packed<T> packed_val;
  packed_val.bits =
      atomic_load_explicit(&(object[pack_offset].val), memory_order_relaxed);
  return packed_val.val[elem_offset];
}

template <typename T, enable_if_t<!is_metal_atomic<T>, bool> = true>
METAL_FUNC void
mlx_atomic_store_explicit(device mlx_atomic<T>* object, T val, size_t offset) {
  mlx_atomic_update_and_store<T, __None<T>>(object, val, offset);
}

template <typename T, enable_if_t<!is_metal_atomic<T>, bool> = true>
METAL_FUNC void mlx_atomic_fetch_and_explicit(
    device mlx_atomic<T>* object,
    T val,
    size_t offset) {
  size_t pack_offset = offset / packing_size<T>;
  size_t elem_offset = offset % packing_size<T>;
  uint_or_packed<T> identity;
  identity.bits = __UINT32_MAX__;
  identity.val[elem_offset] = val;

  atomic_fetch_and_explicit(
      &(object[pack_offset].val), identity.bits, memory_order_relaxed);
}

template <typename T, enable_if_t<!is_metal_atomic<T>, bool> = true>
METAL_FUNC void mlx_atomic_fetch_or_explicit(
    device mlx_atomic<T>* object,
    T val,
    size_t offset) {
  size_t pack_offset = offset / packing_size<T>;
  size_t elem_offset = offset % packing_size<T>;
  uint_or_packed<T> identity;
  identity.bits = 0;
  identity.val[elem_offset] = val;

  atomic_fetch_or_explicit(
      &(object[pack_offset].val), identity.bits, memory_order_relaxed);
}

template <typename T, enable_if_t<!is_metal_atomic<T>, bool> = true>
METAL_FUNC void mlx_atomic_fetch_min_explicit(
    device mlx_atomic<T>* object,
    T val,
    size_t offset) {
  mlx_atomic_update_and_store<T, __Min<T>>(object, val, offset);
}

template <typename T, enable_if_t<!is_metal_atomic<T>, bool> = true>
METAL_FUNC void mlx_atomic_fetch_max_explicit(
    device mlx_atomic<T>* object,
    T val,
    size_t offset) {
  mlx_atomic_update_and_store<T, __Max<T>>(object, val, offset);
}

template <typename T, enable_if_t<!is_metal_atomic<T>, bool> = true>
METAL_FUNC void mlx_atomic_fetch_add_explicit(
    device mlx_atomic<T>* object,
    T val,
    size_t offset) {
  mlx_atomic_update_and_store<T, __Add<T>>(object, val, offset);
}

template <typename T, enable_if_t<!is_metal_atomic<T>, bool> = true>
METAL_FUNC void mlx_atomic_fetch_mul_explicit(
    device mlx_atomic<T>* object,
    T val,
    size_t offset) {
  mlx_atomic_update_and_store<T, __Mul<T>>(object, val, offset);
}

template <typename T, enable_if_t<!is_metal_atomic<T>, bool> = true>
METAL_FUNC bool mlx_atomic_compare_exchange_weak_explicit(
    device mlx_atomic<T>* object,
    thread uint* expected,
    uint val,
    size_t offset) {
  return atomic_compare_exchange_weak_explicit(
      &(object[offset].val),
      expected,
      val,
      memory_order_relaxed,
      memory_order_relaxed);
}
// END RUNTIME UNIT
// Runtime unit: reduction/ops.h; file SHA256: 78d06730fc9564a73944e7f1fe3897d25c8789b28a939bf418e1968db311da41
// Runtime unit SHA256: e73aa8c073780f3d0fa97f431c350d48fcdfb13d60b02910b3d342d38ecf96b7
// BEGIN RUNTIME UNIT
#define DEFINE_SIMD_REDUCE()                                             \
  template <typename T, metal::enable_if_t<sizeof(T) < 8, bool> = true>  \
  T simd_reduce(T val) thread {                                          \
    return simd_reduce_impl(val);                                        \
  }                                                                      \
                                                                         \
  template <typename T, metal::enable_if_t<sizeof(T) == 8, bool> = true> \
  T simd_reduce(T val) thread {                                          \
    for (short i = simd_size / 2; i > 0; i /= 2) {                       \
      val = operator()(val, simd_shuffle_down(val, i));                  \
    }                                                                    \
    return val;                                                          \
  }

static constant constexpr const uint8_t simd_size = 32;

union bool4_or_uint {
  bool4 b;
  unsigned int i;
};

struct None {
  template <typename T>
  void atomic_update(device mlx_atomic<T>* out, T val, size_t offset = 0)
      thread {
    mlx_atomic_store_explicit(out, val, offset);
  }
};

template <typename U = bool>
struct And {
  DEFINE_SIMD_REDUCE()

  bool simd_reduce_impl(bool val) thread {
    return simd_all(val);
  }

  static constexpr constant bool init = true;

  void atomic_update(
      device mlx_atomic<unsigned int>* out,
      bool val,
      int elem_idx,
      size_t offset = 0) thread {
    if (!val) {
      bool4_or_uint update;
      update.b = {true, true, true, true};
      update.b[elem_idx] = false;
      mlx_atomic_fetch_and_explicit(out, update.i, offset);
    }
  }

  void atomic_update(device mlx_atomic<bool>* out, bool val, size_t offset = 0)
      thread {
    if (!val) {
      mlx_atomic_store_explicit(out, val, offset);
    }
  }

  // Non atomic update
  void update(device bool* out, bool val) thread {
    *out &= val;
  }

  // Operator
  bool operator()(bool a, bool b) thread {
    return a && b;
  }
};

template <typename U = bool>
struct Or {
  DEFINE_SIMD_REDUCE()

  bool simd_reduce_impl(bool val) thread {
    return simd_any(val);
  }

  static constexpr constant bool init = false;

  void atomic_update(
      device mlx_atomic<unsigned int>* out,
      bool val,
      int elem_idx,
      size_t offset = 0) thread {
    if (val) {
      bool4_or_uint update;
      update.b = {false, false, false, false};
      update.b[elem_idx] = true;
      mlx_atomic_fetch_or_explicit(out, update.i, offset);
    }
  }

  void atomic_update(device mlx_atomic<bool>* out, bool val, size_t offset = 0)
      thread {
    if (val) {
      mlx_atomic_store_explicit(out, val, offset);
    }
  }

  // Non atomic update
  void update(device bool* out, bool val) thread {
    *out |= val;
  }

  // Operator
  bool operator()(bool a, bool b) thread {
    return a || b;
  }
};

template <typename U>
struct Sum {
  DEFINE_SIMD_REDUCE()

  template <typename T>
  T simd_reduce_impl(T val) thread {
    return simd_sum(val);
  }

  static constexpr constant U init = U(0);

  template <typename T>
  void atomic_update(device mlx_atomic<T>* out, T val, size_t offset = 0)
      thread {
    mlx_atomic_fetch_add_explicit(out, val, offset);
  }

  void atomic_update(
      device mlx_atomic<complex64_t>* out,
      complex64_t val,
      size_t offset = 0) thread {
    auto out_lanes = reinterpret_cast<device mlx_atomic<float>*>(out);
    mlx_atomic_fetch_add_explicit(out_lanes, val.real, 2 * offset);
    mlx_atomic_fetch_add_explicit(out_lanes, val.imag, 2 * offset + 1);
  }

  // Operator
  U operator()(U a, U b) thread {
    return a + b;
  }
};

// END RUNTIME UNIT
// Runtime unit: indexing/indexing.h; file SHA256: e820b8ee2b5132a97122780c12433ebb5100d8078d31e211d0429400a11415bb
// Runtime unit SHA256: 6b177621bd445d19f9b45bfa1f0930e23cf0fe952aa5477cb03789253bea2bf0
// BEGIN RUNTIME UNIT
// Copyright © 2023-2024 Apple Inc.



template <typename IdxT, int NIDX>
struct Indices {
  const array<const device IdxT*, NIDX> buffers;
  const constant int* shapes;
  const constant int64_t* strides;
  const constant bool* row_contiguous;
  const int ndim;
};

template <typename IdxT>
METAL_FUNC size_t offset_neg_idx(IdxT idx, int size) {
  if (is_unsigned_v<IdxT>) {
    return idx;
  } else {
    return (idx < 0) ? idx + size : idx;
  }
}
// END RUNTIME UNIT
// Runtime unit: indexing/scatter.h; file SHA256: fa799c286378c59fbb3aeb973e74bf471861e3356ce2817968fe5d6872c3d9ad
// Runtime unit SHA256: c2ee76965834378eb9bf72db83fcca20e9b37c6c5c60a81de424e2325a02f7c3
// BEGIN RUNTIME UNIT
// Copyright © 2024 Apple Inc.



template <
    typename T,
    typename IdxT,
    typename Op,
    int NIDX,
    bool UPD_ROW_CONTIG,
    int NWORK,
    typename LocT>
METAL_FUNC void scatter_impl(
    const device T* updates,
    device mlx_atomic<T>* out,
    const constant int* upd_shape,
    const constant int64_t* upd_strides,
    const constant size_t& upd_ndim,
    const constant size_t& upd_size,
    const constant int* out_shape,
    const constant int64_t* out_strides,
    const constant size_t& out_ndim,
    const constant int* axes,
    const constant size_t& idx_size,
    const thread Indices<IdxT, NIDX>& indices,
    uint2 gid [[thread_position_in_grid]]) {
  Op op;

  auto ind_idx = gid.y * NWORK;
  LocT out_offset = 0;
  if (upd_size > 1) {
    out_offset = elem_to_loc<LocT>(
        gid.x, upd_shape + indices.ndim, out_strides, out_ndim);
  }

  for (int j = 0; j < NWORK && ind_idx < idx_size; ++j, ind_idx++) {
    LocT out_idx = out_offset;
    for (int i = 0; i < NIDX; ++i) {
      auto idx_loc = indices.row_contiguous[i]
          ? ind_idx
          : elem_to_loc<LocT>(
                ind_idx,
                &indices.shapes[indices.ndim * i],
                &indices.strides[indices.ndim * i],
                indices.ndim);
      auto ax = axes[i];
      auto idx_val = offset_neg_idx(indices.buffers[i][idx_loc], out_shape[ax]);
      out_idx +=
          static_cast<LocT>(idx_val) * static_cast<LocT>(out_strides[ax]);
    }
    auto upd_idx = ind_idx * static_cast<LocT>(upd_size) + gid.x;
    if constexpr (!UPD_ROW_CONTIG) {
      upd_idx = elem_to_loc<LocT>(upd_idx, upd_shape, upd_strides, upd_ndim);
    }
    op.atomic_update(out, updates[upd_idx], out_idx);
  }
}

template <
    typename T,
    typename IdxT,
    typename Op,
    bool OUT_ROW_CONTIG,
    bool UPD_ROW_CONTIG,
    bool UPD_SCALAR,
    int NWORK,
    int NDIM>
[[kernel]] void slice_update_op_impl(
    const device T* updates [[buffer(0)]],
    device T* out [[buffer(1)]],
    const constant int* update_shape [[buffer(2)]],
    const constant int64_t* update_strides [[buffer(3)]],
    const constant int& update_ndim [[buffer(4)]],
    const constant int64_t& update_size [[buffer(5)]],
    const constant int64_t* output_strides [[buffer(6)]],
    const constant int64_t& output_offset [[buffer(7)]],
    uint3 gid [[thread_position_in_grid]],
    uint3 gsize [[threads_per_grid]]) {
  Op op;

  IdxT idx =
      (IdxT(gid.z) * IdxT(gsize.y) + IdxT(gid.y)) * IdxT(gsize.x) * NWORK +
      IdxT(gid.x) * NWORK;
  IdxT out_idx;
  IdxT update_idx;

  if constexpr (OUT_ROW_CONTIG) {
    out_idx = idx;
  } else if constexpr (NDIM == 1) {
    out_idx = NWORK * gid.x * output_strides[0];
  } else if constexpr (NDIM == 2) {
    out_idx = gid.y * output_strides[0] + NWORK * gid.x * output_strides[1];
  } else if constexpr (NDIM == 3) {
    out_idx = gid.z * output_strides[0] + gid.y * output_strides[1] +
        NWORK * gid.x * output_strides[2];
  } else {
    out_idx = elem_to_loc<IdxT>(idx, update_shape, output_strides, update_ndim);
  }

  if constexpr (UPD_SCALAR) {
    update_idx = 0;
  } else if constexpr (UPD_ROW_CONTIG) {
    update_idx = idx;
  } else if constexpr (NDIM == 1) {
    update_idx = NWORK * gid.x * update_strides[0];
  } else if constexpr (NDIM == 2) {
    update_idx = gid.y * update_strides[0] + NWORK * gid.x * update_strides[1];
  } else if constexpr (NDIM == 3) {
    update_idx = gid.z * update_strides[0] + gid.y * update_strides[1] +
        NWORK * gid.x * update_strides[2];
  } else {
    update_idx =
        elem_to_loc<IdxT>(idx, update_shape, update_strides, update_ndim);
  }

  out += output_offset;

  if constexpr (OUT_ROW_CONTIG && (UPD_ROW_CONTIG || UPD_SCALAR)) {
    for (int j = 0; j < NWORK; j++) {
      out[out_idx] = op(out[out_idx], updates[update_idx]);
      out_idx++;
      if constexpr (!UPD_SCALAR) {
        update_idx++;
      }
    }
  } else {
    auto out_stride = output_strides[update_ndim - 1];
    auto update_stride = update_strides[update_ndim - 1];
    for (int j = 0; j < NWORK; j++) {
      out[out_idx] = op(out[out_idx], updates[update_idx]);
      out_idx += out_stride;
      if constexpr (!UPD_SCALAR) {
        update_idx += update_stride;
      }
    }
  }
}
// END RUNTIME UNIT
// Runtime unit: indexing/scatter_axis.h; file SHA256: 43eabd0216101f8e32f5cdd19ce40b7f954564be27fad98a5e0fe345e7b94ce5
// Runtime unit SHA256: 67c5db62df1160220db9342cebdb1a9b0ce96ca2e2d7c25af0a4ba4954c6419e
// BEGIN RUNTIME UNIT
// Copyright © 2025 Apple Inc.


template <
    typename T,
    typename IdxT,
    typename LocT,
    typename Op,
    bool UpdC,
    bool IdxC>
[[kernel]] void scatter_axis(
    const device T* upd [[buffer(0)]],
    const device IdxT* indices [[buffer(1)]],
    device mlx_atomic<T>* out [[buffer(2)]],
    const constant int* shape [[buffer(3)]],
    const constant int64_t* upd_strides [[buffer(4)]],
    const constant int64_t* idx_strides [[buffer(5)]],
    const constant size_t& ndim [[buffer(6)]],
    const constant int& axis [[buffer(7)]],
    const constant int& out_axis_size [[buffer(8)]],
    const constant size_t& upd_ax_stride [[buffer(9)]],
    const constant size_t& idx_ax_stride [[buffer(10)]],
    uint3 index [[thread_position_in_grid]],
    uint3 grid_dim [[threads_per_grid]]) {
  Op op;

  LocT elem_idx = index.z * static_cast<LocT>(grid_dim.x);

  LocT idx_loc = index.y * static_cast<LocT>(idx_ax_stride);
  if (IdxC) {
    idx_loc += elem_idx * grid_dim.y + index.x;
  } else {
    idx_loc += elem_to_loc<LocT>(elem_idx + index.x, shape, idx_strides, ndim);
  }

  auto idx_val = indices[idx_loc];
  if (is_signed_v<IdxT>) {
    idx_val = (idx_val < 0) ? idx_val + out_axis_size : idx_val;
  }

  LocT upd_idx = index.y * static_cast<LocT>(upd_ax_stride);
  if (UpdC) {
    upd_idx += elem_idx * grid_dim.y + index.x;
  } else {
    upd_idx += elem_to_loc<LocT>(elem_idx + index.x, shape, upd_strides, ndim);
  }

  LocT out_idx = elem_idx * static_cast<LocT>(out_axis_size) +
      idx_val * grid_dim.x + index.x;
  op.atomic_update(out, upd[upd_idx], out_idx);
}
// END RUNTIME UNIT
template [[host_name("kernel_qwen_mtplx_scatter_axis_bool_idxi64_int_00")]]
[[kernel]] decltype(scatter_axis<bool, int64_t, int, None, false, false>) scatter_axis<bool, int64_t, int, None, false, false>;
template [[host_name("kernel_qwen_mtplx_scatter_axis_bool_idxi64_int_01")]]
[[kernel]] decltype(scatter_axis<bool, int64_t, int, None, false, true>) scatter_axis<bool, int64_t, int, None, false, true>;
template [[host_name("kernel_qwen_mtplx_scatter_axis_bool_idxi64_int_10")]]
[[kernel]] decltype(scatter_axis<bool, int64_t, int, None, true, false>) scatter_axis<bool, int64_t, int, None, true, false>;
template [[host_name("kernel_qwen_mtplx_scatter_axis_bool_idxi64_int_11")]]
[[kernel]] decltype(scatter_axis<bool, int64_t, int, None, true, true>) scatter_axis<bool, int64_t, int, None, true, true>;
template [[host_name("kernel_qwen_mtplx_scatter_axis_bool_idxi64_int64_t_00")]]
[[kernel]] decltype(scatter_axis<bool, int64_t, int64_t, None, false, false>) scatter_axis<bool, int64_t, int64_t, None, false, false>;
template [[host_name("kernel_qwen_mtplx_scatter_axis_bool_idxi64_int64_t_01")]]
[[kernel]] decltype(scatter_axis<bool, int64_t, int64_t, None, false, true>) scatter_axis<bool, int64_t, int64_t, None, false, true>;
template [[host_name("kernel_qwen_mtplx_scatter_axis_bool_idxi64_int64_t_10")]]
[[kernel]] decltype(scatter_axis<bool, int64_t, int64_t, None, true, false>) scatter_axis<bool, int64_t, int64_t, None, true, false>;
template [[host_name("kernel_qwen_mtplx_scatter_axis_bool_idxi64_int64_t_11")]]
[[kernel]] decltype(scatter_axis<bool, int64_t, int64_t, None, true, true>) scatter_axis<bool, int64_t, int64_t, None, true, true>;
// Runtime unit: scatter JIT wrappers; file SHA256: 1b38dbdf3120eca3e5266ffd6d69691591c9432c0b5acece4c5ab2de7a9e802f
// Runtime unit SHA256: 1206fc4eda346eccb7778882e03cfbde8608f2195d32a660dc14a49077a637ad
// BEGIN RUNTIME UNIT

[[host_name("kernel_qwen_mtplx_scatter_bf16_idxi64_sum_1_updc_false_nwork1_int")]] [[kernel]] void kernel_qwen_mtplx_scatter_bf16_idxi64_sum_1_updc_false_nwork1_int(
    const device bfloat* updates [[buffer(1)]],
    device mlx_atomic<bfloat>* out [[buffer(2)]],
    const constant int* upd_shape [[buffer(3)]],
    const constant int64_t* upd_strides [[buffer(4)]],
    const constant size_t& upd_ndim [[buffer(5)]],
    const constant size_t& upd_size [[buffer(6)]],
    const constant int* out_shape [[buffer(7)]],
    const constant int64_t* out_strides [[buffer(8)]],
    const constant size_t& out_ndim [[buffer(9)]],
    const constant int* axes [[buffer(10)]],
    const constant int* idx_shapes [[buffer(11)]],
    const constant int64_t* idx_strides [[buffer(12)]],
    const constant bool* idx_contigs [[buffer(13)]],
    const constant int& idx_ndim [[buffer(14)]],
    const constant size_t& idx_size [[buffer(15)]],
    const device int64_t *idx0 [[buffer(20)]],
    uint2 gid [[thread_position_in_grid]]) {
  Indices<int64_t, 1> idxs{ { idx0 }, idx_shapes, idx_strides, idx_contigs, idx_ndim};

  return scatter_impl<bfloat, int64_t, Sum<bfloat>, 1, false, 1, int>(
      updates,
      out,
      upd_shape,
      upd_strides,
      upd_ndim,
      upd_size,
      out_shape,
      out_strides,
      out_ndim,
      axes,
      idx_size,
      idxs,
      gid);
}

[[host_name("kernel_qwen_mtplx_scatter_bf16_idxi64_sum_1_updc_false_nwork1_int64_t")]] [[kernel]] void kernel_qwen_mtplx_scatter_bf16_idxi64_sum_1_updc_false_nwork1_int64_t(
    const device bfloat* updates [[buffer(1)]],
    device mlx_atomic<bfloat>* out [[buffer(2)]],
    const constant int* upd_shape [[buffer(3)]],
    const constant int64_t* upd_strides [[buffer(4)]],
    const constant size_t& upd_ndim [[buffer(5)]],
    const constant size_t& upd_size [[buffer(6)]],
    const constant int* out_shape [[buffer(7)]],
    const constant int64_t* out_strides [[buffer(8)]],
    const constant size_t& out_ndim [[buffer(9)]],
    const constant int* axes [[buffer(10)]],
    const constant int* idx_shapes [[buffer(11)]],
    const constant int64_t* idx_strides [[buffer(12)]],
    const constant bool* idx_contigs [[buffer(13)]],
    const constant int& idx_ndim [[buffer(14)]],
    const constant size_t& idx_size [[buffer(15)]],
    const device int64_t *idx0 [[buffer(20)]],
    uint2 gid [[thread_position_in_grid]]) {
  Indices<int64_t, 1> idxs{ { idx0 }, idx_shapes, idx_strides, idx_contigs, idx_ndim};

  return scatter_impl<bfloat, int64_t, Sum<bfloat>, 1, false, 1, int64_t>(
      updates,
      out,
      upd_shape,
      upd_strides,
      upd_ndim,
      upd_size,
      out_shape,
      out_strides,
      out_ndim,
      axes,
      idx_size,
      idxs,
      gid);
}

[[host_name("kernel_qwen_mtplx_scatter_bf16_idxi64_sum_1_updc_true_nwork1_int")]] [[kernel]] void kernel_qwen_mtplx_scatter_bf16_idxi64_sum_1_updc_true_nwork1_int(
    const device bfloat* updates [[buffer(1)]],
    device mlx_atomic<bfloat>* out [[buffer(2)]],
    const constant int* upd_shape [[buffer(3)]],
    const constant int64_t* upd_strides [[buffer(4)]],
    const constant size_t& upd_ndim [[buffer(5)]],
    const constant size_t& upd_size [[buffer(6)]],
    const constant int* out_shape [[buffer(7)]],
    const constant int64_t* out_strides [[buffer(8)]],
    const constant size_t& out_ndim [[buffer(9)]],
    const constant int* axes [[buffer(10)]],
    const constant int* idx_shapes [[buffer(11)]],
    const constant int64_t* idx_strides [[buffer(12)]],
    const constant bool* idx_contigs [[buffer(13)]],
    const constant int& idx_ndim [[buffer(14)]],
    const constant size_t& idx_size [[buffer(15)]],
    const device int64_t *idx0 [[buffer(20)]],
    uint2 gid [[thread_position_in_grid]]) {
  Indices<int64_t, 1> idxs{ { idx0 }, idx_shapes, idx_strides, idx_contigs, idx_ndim};

  return scatter_impl<bfloat, int64_t, Sum<bfloat>, 1, true, 1, int>(
      updates,
      out,
      upd_shape,
      upd_strides,
      upd_ndim,
      upd_size,
      out_shape,
      out_strides,
      out_ndim,
      axes,
      idx_size,
      idxs,
      gid);
}

[[host_name("kernel_qwen_mtplx_scatter_bf16_idxi64_sum_1_updc_true_nwork1_int64_t")]] [[kernel]] void kernel_qwen_mtplx_scatter_bf16_idxi64_sum_1_updc_true_nwork1_int64_t(
    const device bfloat* updates [[buffer(1)]],
    device mlx_atomic<bfloat>* out [[buffer(2)]],
    const constant int* upd_shape [[buffer(3)]],
    const constant int64_t* upd_strides [[buffer(4)]],
    const constant size_t& upd_ndim [[buffer(5)]],
    const constant size_t& upd_size [[buffer(6)]],
    const constant int* out_shape [[buffer(7)]],
    const constant int64_t* out_strides [[buffer(8)]],
    const constant size_t& out_ndim [[buffer(9)]],
    const constant int* axes [[buffer(10)]],
    const constant int* idx_shapes [[buffer(11)]],
    const constant int64_t* idx_strides [[buffer(12)]],
    const constant bool* idx_contigs [[buffer(13)]],
    const constant int& idx_ndim [[buffer(14)]],
    const constant size_t& idx_size [[buffer(15)]],
    const device int64_t *idx0 [[buffer(20)]],
    uint2 gid [[thread_position_in_grid]]) {
  Indices<int64_t, 1> idxs{ { idx0 }, idx_shapes, idx_strides, idx_contigs, idx_ndim};

  return scatter_impl<bfloat, int64_t, Sum<bfloat>, 1, true, 1, int64_t>(
      updates,
      out,
      upd_shape,
      upd_strides,
      upd_ndim,
      upd_size,
      out_shape,
      out_strides,
      out_ndim,
      axes,
      idx_size,
      idxs,
      gid);
}

[[host_name("kernel_qwen_mtplx_scatter_f16_idxi64_sum_1_updc_false_nwork1_int")]] [[kernel]] void kernel_qwen_mtplx_scatter_f16_idxi64_sum_1_updc_false_nwork1_int(
    const device half* updates [[buffer(1)]],
    device mlx_atomic<half>* out [[buffer(2)]],
    const constant int* upd_shape [[buffer(3)]],
    const constant int64_t* upd_strides [[buffer(4)]],
    const constant size_t& upd_ndim [[buffer(5)]],
    const constant size_t& upd_size [[buffer(6)]],
    const constant int* out_shape [[buffer(7)]],
    const constant int64_t* out_strides [[buffer(8)]],
    const constant size_t& out_ndim [[buffer(9)]],
    const constant int* axes [[buffer(10)]],
    const constant int* idx_shapes [[buffer(11)]],
    const constant int64_t* idx_strides [[buffer(12)]],
    const constant bool* idx_contigs [[buffer(13)]],
    const constant int& idx_ndim [[buffer(14)]],
    const constant size_t& idx_size [[buffer(15)]],
    const device int64_t *idx0 [[buffer(20)]],
    uint2 gid [[thread_position_in_grid]]) {
  Indices<int64_t, 1> idxs{ { idx0 }, idx_shapes, idx_strides, idx_contigs, idx_ndim};

  return scatter_impl<half, int64_t, Sum<half>, 1, false, 1, int>(
      updates,
      out,
      upd_shape,
      upd_strides,
      upd_ndim,
      upd_size,
      out_shape,
      out_strides,
      out_ndim,
      axes,
      idx_size,
      idxs,
      gid);
}

[[host_name("kernel_qwen_mtplx_scatter_f16_idxi64_sum_1_updc_false_nwork1_int64_t")]] [[kernel]] void kernel_qwen_mtplx_scatter_f16_idxi64_sum_1_updc_false_nwork1_int64_t(
    const device half* updates [[buffer(1)]],
    device mlx_atomic<half>* out [[buffer(2)]],
    const constant int* upd_shape [[buffer(3)]],
    const constant int64_t* upd_strides [[buffer(4)]],
    const constant size_t& upd_ndim [[buffer(5)]],
    const constant size_t& upd_size [[buffer(6)]],
    const constant int* out_shape [[buffer(7)]],
    const constant int64_t* out_strides [[buffer(8)]],
    const constant size_t& out_ndim [[buffer(9)]],
    const constant int* axes [[buffer(10)]],
    const constant int* idx_shapes [[buffer(11)]],
    const constant int64_t* idx_strides [[buffer(12)]],
    const constant bool* idx_contigs [[buffer(13)]],
    const constant int& idx_ndim [[buffer(14)]],
    const constant size_t& idx_size [[buffer(15)]],
    const device int64_t *idx0 [[buffer(20)]],
    uint2 gid [[thread_position_in_grid]]) {
  Indices<int64_t, 1> idxs{ { idx0 }, idx_shapes, idx_strides, idx_contigs, idx_ndim};

  return scatter_impl<half, int64_t, Sum<half>, 1, false, 1, int64_t>(
      updates,
      out,
      upd_shape,
      upd_strides,
      upd_ndim,
      upd_size,
      out_shape,
      out_strides,
      out_ndim,
      axes,
      idx_size,
      idxs,
      gid);
}

[[host_name("kernel_qwen_mtplx_scatter_f16_idxi64_sum_1_updc_true_nwork1_int")]] [[kernel]] void kernel_qwen_mtplx_scatter_f16_idxi64_sum_1_updc_true_nwork1_int(
    const device half* updates [[buffer(1)]],
    device mlx_atomic<half>* out [[buffer(2)]],
    const constant int* upd_shape [[buffer(3)]],
    const constant int64_t* upd_strides [[buffer(4)]],
    const constant size_t& upd_ndim [[buffer(5)]],
    const constant size_t& upd_size [[buffer(6)]],
    const constant int* out_shape [[buffer(7)]],
    const constant int64_t* out_strides [[buffer(8)]],
    const constant size_t& out_ndim [[buffer(9)]],
    const constant int* axes [[buffer(10)]],
    const constant int* idx_shapes [[buffer(11)]],
    const constant int64_t* idx_strides [[buffer(12)]],
    const constant bool* idx_contigs [[buffer(13)]],
    const constant int& idx_ndim [[buffer(14)]],
    const constant size_t& idx_size [[buffer(15)]],
    const device int64_t *idx0 [[buffer(20)]],
    uint2 gid [[thread_position_in_grid]]) {
  Indices<int64_t, 1> idxs{ { idx0 }, idx_shapes, idx_strides, idx_contigs, idx_ndim};

  return scatter_impl<half, int64_t, Sum<half>, 1, true, 1, int>(
      updates,
      out,
      upd_shape,
      upd_strides,
      upd_ndim,
      upd_size,
      out_shape,
      out_strides,
      out_ndim,
      axes,
      idx_size,
      idxs,
      gid);
}

[[host_name("kernel_qwen_mtplx_scatter_f16_idxi64_sum_1_updc_true_nwork1_int64_t")]] [[kernel]] void kernel_qwen_mtplx_scatter_f16_idxi64_sum_1_updc_true_nwork1_int64_t(
    const device half* updates [[buffer(1)]],
    device mlx_atomic<half>* out [[buffer(2)]],
    const constant int* upd_shape [[buffer(3)]],
    const constant int64_t* upd_strides [[buffer(4)]],
    const constant size_t& upd_ndim [[buffer(5)]],
    const constant size_t& upd_size [[buffer(6)]],
    const constant int* out_shape [[buffer(7)]],
    const constant int64_t* out_strides [[buffer(8)]],
    const constant size_t& out_ndim [[buffer(9)]],
    const constant int* axes [[buffer(10)]],
    const constant int* idx_shapes [[buffer(11)]],
    const constant int64_t* idx_strides [[buffer(12)]],
    const constant bool* idx_contigs [[buffer(13)]],
    const constant int& idx_ndim [[buffer(14)]],
    const constant size_t& idx_size [[buffer(15)]],
    const device int64_t *idx0 [[buffer(20)]],
    uint2 gid [[thread_position_in_grid]]) {
  Indices<int64_t, 1> idxs{ { idx0 }, idx_shapes, idx_strides, idx_contigs, idx_ndim};

  return scatter_impl<half, int64_t, Sum<half>, 1, true, 1, int64_t>(
      updates,
      out,
      upd_shape,
      upd_strides,
      upd_ndim,
      upd_size,
      out_shape,
      out_strides,
      out_ndim,
      axes,
      idx_size,
      idxs,
      gid);
}

[[host_name("kernel_qwen_mtplx_scatter_f32_idxi64_sum_1_updc_false_nwork1_int")]] [[kernel]] void kernel_qwen_mtplx_scatter_f32_idxi64_sum_1_updc_false_nwork1_int(
    const device float* updates [[buffer(1)]],
    device mlx_atomic<float>* out [[buffer(2)]],
    const constant int* upd_shape [[buffer(3)]],
    const constant int64_t* upd_strides [[buffer(4)]],
    const constant size_t& upd_ndim [[buffer(5)]],
    const constant size_t& upd_size [[buffer(6)]],
    const constant int* out_shape [[buffer(7)]],
    const constant int64_t* out_strides [[buffer(8)]],
    const constant size_t& out_ndim [[buffer(9)]],
    const constant int* axes [[buffer(10)]],
    const constant int* idx_shapes [[buffer(11)]],
    const constant int64_t* idx_strides [[buffer(12)]],
    const constant bool* idx_contigs [[buffer(13)]],
    const constant int& idx_ndim [[buffer(14)]],
    const constant size_t& idx_size [[buffer(15)]],
    const device int64_t *idx0 [[buffer(20)]],
    uint2 gid [[thread_position_in_grid]]) {
  Indices<int64_t, 1> idxs{ { idx0 }, idx_shapes, idx_strides, idx_contigs, idx_ndim};

  return scatter_impl<float, int64_t, Sum<float>, 1, false, 1, int>(
      updates,
      out,
      upd_shape,
      upd_strides,
      upd_ndim,
      upd_size,
      out_shape,
      out_strides,
      out_ndim,
      axes,
      idx_size,
      idxs,
      gid);
}

[[host_name("kernel_qwen_mtplx_scatter_f32_idxi64_sum_1_updc_false_nwork1_int64_t")]] [[kernel]] void kernel_qwen_mtplx_scatter_f32_idxi64_sum_1_updc_false_nwork1_int64_t(
    const device float* updates [[buffer(1)]],
    device mlx_atomic<float>* out [[buffer(2)]],
    const constant int* upd_shape [[buffer(3)]],
    const constant int64_t* upd_strides [[buffer(4)]],
    const constant size_t& upd_ndim [[buffer(5)]],
    const constant size_t& upd_size [[buffer(6)]],
    const constant int* out_shape [[buffer(7)]],
    const constant int64_t* out_strides [[buffer(8)]],
    const constant size_t& out_ndim [[buffer(9)]],
    const constant int* axes [[buffer(10)]],
    const constant int* idx_shapes [[buffer(11)]],
    const constant int64_t* idx_strides [[buffer(12)]],
    const constant bool* idx_contigs [[buffer(13)]],
    const constant int& idx_ndim [[buffer(14)]],
    const constant size_t& idx_size [[buffer(15)]],
    const device int64_t *idx0 [[buffer(20)]],
    uint2 gid [[thread_position_in_grid]]) {
  Indices<int64_t, 1> idxs{ { idx0 }, idx_shapes, idx_strides, idx_contigs, idx_ndim};

  return scatter_impl<float, int64_t, Sum<float>, 1, false, 1, int64_t>(
      updates,
      out,
      upd_shape,
      upd_strides,
      upd_ndim,
      upd_size,
      out_shape,
      out_strides,
      out_ndim,
      axes,
      idx_size,
      idxs,
      gid);
}

[[host_name("kernel_qwen_mtplx_scatter_f32_idxi64_sum_1_updc_true_nwork1_int")]] [[kernel]] void kernel_qwen_mtplx_scatter_f32_idxi64_sum_1_updc_true_nwork1_int(
    const device float* updates [[buffer(1)]],
    device mlx_atomic<float>* out [[buffer(2)]],
    const constant int* upd_shape [[buffer(3)]],
    const constant int64_t* upd_strides [[buffer(4)]],
    const constant size_t& upd_ndim [[buffer(5)]],
    const constant size_t& upd_size [[buffer(6)]],
    const constant int* out_shape [[buffer(7)]],
    const constant int64_t* out_strides [[buffer(8)]],
    const constant size_t& out_ndim [[buffer(9)]],
    const constant int* axes [[buffer(10)]],
    const constant int* idx_shapes [[buffer(11)]],
    const constant int64_t* idx_strides [[buffer(12)]],
    const constant bool* idx_contigs [[buffer(13)]],
    const constant int& idx_ndim [[buffer(14)]],
    const constant size_t& idx_size [[buffer(15)]],
    const device int64_t *idx0 [[buffer(20)]],
    uint2 gid [[thread_position_in_grid]]) {
  Indices<int64_t, 1> idxs{ { idx0 }, idx_shapes, idx_strides, idx_contigs, idx_ndim};

  return scatter_impl<float, int64_t, Sum<float>, 1, true, 1, int>(
      updates,
      out,
      upd_shape,
      upd_strides,
      upd_ndim,
      upd_size,
      out_shape,
      out_strides,
      out_ndim,
      axes,
      idx_size,
      idxs,
      gid);
}

[[host_name("kernel_qwen_mtplx_scatter_f32_idxi64_sum_1_updc_true_nwork1_int64_t")]] [[kernel]] void kernel_qwen_mtplx_scatter_f32_idxi64_sum_1_updc_true_nwork1_int64_t(
    const device float* updates [[buffer(1)]],
    device mlx_atomic<float>* out [[buffer(2)]],
    const constant int* upd_shape [[buffer(3)]],
    const constant int64_t* upd_strides [[buffer(4)]],
    const constant size_t& upd_ndim [[buffer(5)]],
    const constant size_t& upd_size [[buffer(6)]],
    const constant int* out_shape [[buffer(7)]],
    const constant int64_t* out_strides [[buffer(8)]],
    const constant size_t& out_ndim [[buffer(9)]],
    const constant int* axes [[buffer(10)]],
    const constant int* idx_shapes [[buffer(11)]],
    const constant int64_t* idx_strides [[buffer(12)]],
    const constant bool* idx_contigs [[buffer(13)]],
    const constant int& idx_ndim [[buffer(14)]],
    const constant size_t& idx_size [[buffer(15)]],
    const device int64_t *idx0 [[buffer(20)]],
    uint2 gid [[thread_position_in_grid]]) {
  Indices<int64_t, 1> idxs{ { idx0 }, idx_shapes, idx_strides, idx_contigs, idx_ndim};

  return scatter_impl<float, int64_t, Sum<float>, 1, true, 1, int64_t>(
      updates,
      out,
      upd_shape,
      upd_strides,
      upd_ndim,
      upd_size,
      out_shape,
      out_strides,
      out_ndim,
      axes,
      idx_size,
      idxs,
      gid);
}
// END RUNTIME UNIT
} // namespace mtplx_eager_scatter

// Actual runtime-generated SwiGLU: Copyright Apple Inc.; SPDX-License-Identifier: MIT
// Runtime unit: captured-swiglu-jit; file SHA256: 8cb2a51a7d9f0cb91a8a7669ef844c3964dc469ec9e4436a719e47f025b8e781
// Runtime unit SHA256: 58a220047d41020fd7f714c16ab50fcb3955523aa5a036857ff718abdd84b7dc
// BEGIN RUNTIME UNIT
[[host_name("kernel_qwen_mtplx_swiglu_bf16_contiguous")]]
[[kernel]] void CV2ISigmoidADV2IBroadcastACEV2IBroadcastCAFV2IMultiplyDEGV2IBroadcastFBHV2IBroadcastBFIV2OMultiplyGH_VV_V2V2_11160318154034397263_contiguous(
    device const bfloat16_t* A [[buffer(0)]],
    device const bfloat16_t* B [[buffer(1)]],
    device bfloat16_t* C [[buffer(2)]],
    constant const uint& size [[buffer(3)]],
    uint3 pos [[thread_position_in_grid]],
    uint3 grid [[threads_per_grid]]) {
  constexpr int N_ = 1;
  uint index = N_ * pos.x;
  bfloat16_t tmp_A = A[index];
  bfloat16_t tmp_B = B[index];
  bfloat16_t tmp_D = Sigmoid()(tmp_A);
  bfloat16_t tmp_E = cast_to<bfloat16_t>(tmp_A);
  bfloat16_t tmp_F = cast_to<bfloat16_t>(tmp_D);
  bfloat16_t tmp_G = Multiply()(tmp_E, tmp_F);
  bfloat16_t tmp_H = cast_to<bfloat16_t>(tmp_G);
  bfloat16_t tmp_I = cast_to<bfloat16_t>(tmp_B);
  bfloat16_t tmp_C = Multiply()(tmp_H, tmp_I);
  C[index] = tmp_C;
}
[[host_name("kernel_qwen_mtplx_swiglu_bf16_contiguous_n")]]
[[kernel]] void CV2ISigmoidADV2IBroadcastACEV2IBroadcastCAFV2IMultiplyDEGV2IBroadcastFBHV2IBroadcastBFIV2OMultiplyGH_VV_V2V2_11160318154034397263_contiguous_n(
    device const bfloat16_t* A [[buffer(0)]],
    device const bfloat16_t* B [[buffer(1)]],
    device bfloat16_t* C [[buffer(2)]],
    constant const uint& size [[buffer(3)]],
    uint3 pos [[thread_position_in_grid]],
    uint3 grid [[threads_per_grid]]) {
  constexpr int N_ = 4;
  uint index = N_ * pos.x;
  for (int i = 0; i < N_ && index < size; ++i) {
  bfloat16_t tmp_A = A[index];
  bfloat16_t tmp_B = B[index];
  bfloat16_t tmp_D = Sigmoid()(tmp_A);
  bfloat16_t tmp_E = cast_to<bfloat16_t>(tmp_A);
  bfloat16_t tmp_F = cast_to<bfloat16_t>(tmp_D);
  bfloat16_t tmp_G = Multiply()(tmp_E, tmp_F);
  bfloat16_t tmp_H = cast_to<bfloat16_t>(tmp_G);
  bfloat16_t tmp_I = cast_to<bfloat16_t>(tmp_B);
  bfloat16_t tmp_C = Multiply()(tmp_H, tmp_I);
  C[index] = tmp_C;
  index++;
  }
}
[[host_name("kernel_qwen_mtplx_swiglu_bf16_contiguous_large")]]
[[kernel]] void CV2ISigmoidADV2IBroadcastACEV2IBroadcastCAFV2IMultiplyDEGV2IBroadcastFBHV2IBroadcastBFIV2OMultiplyGH_VV_V2V2_11160318154034397263_contiguous_large(
    device const bfloat16_t* A [[buffer(0)]],
    device const bfloat16_t* B [[buffer(1)]],
    device bfloat16_t* C [[buffer(2)]],
    constant const int64_t& size [[buffer(3)]],
    uint3 pos [[thread_position_in_grid]],
    uint3 grid [[threads_per_grid]]) {
  constexpr int N_ = 4;
  int64_t index = N_ * (pos.x + grid.x * int64_t(pos.y));
  for (int i = 0; i < N_ && index < size; ++i) {
  bfloat16_t tmp_A = A[index];
  bfloat16_t tmp_B = B[index];
  bfloat16_t tmp_D = Sigmoid()(tmp_A);
  bfloat16_t tmp_E = cast_to<bfloat16_t>(tmp_A);
  bfloat16_t tmp_F = cast_to<bfloat16_t>(tmp_D);
  bfloat16_t tmp_G = Multiply()(tmp_E, tmp_F);
  bfloat16_t tmp_H = cast_to<bfloat16_t>(tmp_G);
  bfloat16_t tmp_I = cast_to<bfloat16_t>(tmp_B);
  bfloat16_t tmp_C = Multiply()(tmp_H, tmp_I);
  C[index] = tmp_C;
  index++;
  }
}
[[host_name("kernel_qwen_mtplx_swiglu_bf16_strided_1")]]
[[kernel]] void CV2ISigmoidADV2IBroadcastACEV2IBroadcastCAFV2IMultiplyDEGV2IBroadcastFBHV2IBroadcastBFIV2OMultiplyGH_VV_V2V2_11160318154034397263_strided_1(
    device const bfloat16_t* A [[buffer(0)]],
    device const bfloat16_t* B [[buffer(1)]],
    constant const int64_t* in_strides [[buffer(2)]],
    device bfloat16_t* C [[buffer(3)]],
    constant const int* output_shape [[buffer(4)]],
    uint3 pos [[thread_position_in_grid]],
    uint3 grid [[threads_per_grid]]) {
  constexpr int N_ = 1;
  uint index = pos.x + grid.x * (pos.y + uint(grid.y) * pos.z);
  uint index_A = elem_to_loc_1<uint>(pos.x, in_strides[0]);
  uint index_B = elem_to_loc_1<uint>(pos.x, in_strides[1]);
  bfloat16_t tmp_A = A[index_A];
  bfloat16_t tmp_B = B[index_B];
  bfloat16_t tmp_D = Sigmoid()(tmp_A);
  bfloat16_t tmp_E = cast_to<bfloat16_t>(tmp_A);
  bfloat16_t tmp_F = cast_to<bfloat16_t>(tmp_D);
  bfloat16_t tmp_G = Multiply()(tmp_E, tmp_F);
  bfloat16_t tmp_H = cast_to<bfloat16_t>(tmp_G);
  bfloat16_t tmp_I = cast_to<bfloat16_t>(tmp_B);
  bfloat16_t tmp_C = Multiply()(tmp_H, tmp_I);
  C[index] = tmp_C;
}
[[host_name("kernel_qwen_mtplx_swiglu_bf16_strided_1_large")]]
[[kernel]] void CV2ISigmoidADV2IBroadcastACEV2IBroadcastCAFV2IMultiplyDEGV2IBroadcastFBHV2IBroadcastBFIV2OMultiplyGH_VV_V2V2_11160318154034397263_strided_1_large(
    device const bfloat16_t* A [[buffer(0)]],
    device const bfloat16_t* B [[buffer(1)]],
    constant const int64_t* in_strides [[buffer(2)]],
    device bfloat16_t* C [[buffer(3)]],
    constant const int* output_shape [[buffer(4)]],
    uint3 pos [[thread_position_in_grid]],
    uint3 grid [[threads_per_grid]]) {
  constexpr int N_ = 1;
  int64_t index = pos.x + grid.x * (pos.y + int64_t(grid.y) * pos.z);
  int64_t index_A = elem_to_loc_1<int64_t>(pos.x, in_strides[0]);
  int64_t index_B = elem_to_loc_1<int64_t>(pos.x, in_strides[1]);
  bfloat16_t tmp_A = A[index_A];
  bfloat16_t tmp_B = B[index_B];
  bfloat16_t tmp_D = Sigmoid()(tmp_A);
  bfloat16_t tmp_E = cast_to<bfloat16_t>(tmp_A);
  bfloat16_t tmp_F = cast_to<bfloat16_t>(tmp_D);
  bfloat16_t tmp_G = Multiply()(tmp_E, tmp_F);
  bfloat16_t tmp_H = cast_to<bfloat16_t>(tmp_G);
  bfloat16_t tmp_I = cast_to<bfloat16_t>(tmp_B);
  bfloat16_t tmp_C = Multiply()(tmp_H, tmp_I);
  C[index] = tmp_C;
}
[[host_name("kernel_qwen_mtplx_swiglu_bf16_strided_2")]]
[[kernel]] void CV2ISigmoidADV2IBroadcastACEV2IBroadcastCAFV2IMultiplyDEGV2IBroadcastFBHV2IBroadcastBFIV2OMultiplyGH_VV_V2V2_11160318154034397263_strided_2(
    device const bfloat16_t* A [[buffer(0)]],
    device const bfloat16_t* B [[buffer(1)]],
    constant const int64_t* in_strides [[buffer(2)]],
    device bfloat16_t* C [[buffer(3)]],
    constant const int* output_shape [[buffer(4)]],
    uint3 pos [[thread_position_in_grid]],
    uint3 grid [[threads_per_grid]]) {
  constexpr int N_ = 1;
  uint index = pos.x + grid.x * (pos.y + uint(grid.y) * pos.z);
  uint index_A = elem_to_loc_2<uint>({pos.x, pos.y}, in_strides + 0);
  uint index_B = elem_to_loc_2<uint>({pos.x, pos.y}, in_strides + 2);
  bfloat16_t tmp_A = A[index_A];
  bfloat16_t tmp_B = B[index_B];
  bfloat16_t tmp_D = Sigmoid()(tmp_A);
  bfloat16_t tmp_E = cast_to<bfloat16_t>(tmp_A);
  bfloat16_t tmp_F = cast_to<bfloat16_t>(tmp_D);
  bfloat16_t tmp_G = Multiply()(tmp_E, tmp_F);
  bfloat16_t tmp_H = cast_to<bfloat16_t>(tmp_G);
  bfloat16_t tmp_I = cast_to<bfloat16_t>(tmp_B);
  bfloat16_t tmp_C = Multiply()(tmp_H, tmp_I);
  C[index] = tmp_C;
}
[[host_name("kernel_qwen_mtplx_swiglu_bf16_strided_2_large")]]
[[kernel]] void CV2ISigmoidADV2IBroadcastACEV2IBroadcastCAFV2IMultiplyDEGV2IBroadcastFBHV2IBroadcastBFIV2OMultiplyGH_VV_V2V2_11160318154034397263_strided_2_large(
    device const bfloat16_t* A [[buffer(0)]],
    device const bfloat16_t* B [[buffer(1)]],
    constant const int64_t* in_strides [[buffer(2)]],
    device bfloat16_t* C [[buffer(3)]],
    constant const int* output_shape [[buffer(4)]],
    uint3 pos [[thread_position_in_grid]],
    uint3 grid [[threads_per_grid]]) {
  constexpr int N_ = 1;
  int64_t index = pos.x + grid.x * (pos.y + int64_t(grid.y) * pos.z);
  int64_t index_A = elem_to_loc_2<int64_t>({pos.x, pos.y}, in_strides + 0);
  int64_t index_B = elem_to_loc_2<int64_t>({pos.x, pos.y}, in_strides + 2);
  bfloat16_t tmp_A = A[index_A];
  bfloat16_t tmp_B = B[index_B];
  bfloat16_t tmp_D = Sigmoid()(tmp_A);
  bfloat16_t tmp_E = cast_to<bfloat16_t>(tmp_A);
  bfloat16_t tmp_F = cast_to<bfloat16_t>(tmp_D);
  bfloat16_t tmp_G = Multiply()(tmp_E, tmp_F);
  bfloat16_t tmp_H = cast_to<bfloat16_t>(tmp_G);
  bfloat16_t tmp_I = cast_to<bfloat16_t>(tmp_B);
  bfloat16_t tmp_C = Multiply()(tmp_H, tmp_I);
  C[index] = tmp_C;
}
[[host_name("kernel_qwen_mtplx_swiglu_bf16_strided_3")]]
[[kernel]] void CV2ISigmoidADV2IBroadcastACEV2IBroadcastCAFV2IMultiplyDEGV2IBroadcastFBHV2IBroadcastBFIV2OMultiplyGH_VV_V2V2_11160318154034397263_strided_3(
    device const bfloat16_t* A [[buffer(0)]],
    device const bfloat16_t* B [[buffer(1)]],
    constant const int64_t* in_strides [[buffer(2)]],
    device bfloat16_t* C [[buffer(3)]],
    constant const int* output_shape [[buffer(4)]],
    uint3 pos [[thread_position_in_grid]],
    uint3 grid [[threads_per_grid]]) {
  constexpr int N_ = 1;
  uint index = pos.x + grid.x * (pos.y + uint(grid.y) * pos.z);
  uint index_A = elem_to_loc_3<uint>(pos, in_strides + 0);
  uint index_B = elem_to_loc_3<uint>(pos, in_strides + 3);
  bfloat16_t tmp_A = A[index_A];
  bfloat16_t tmp_B = B[index_B];
  bfloat16_t tmp_D = Sigmoid()(tmp_A);
  bfloat16_t tmp_E = cast_to<bfloat16_t>(tmp_A);
  bfloat16_t tmp_F = cast_to<bfloat16_t>(tmp_D);
  bfloat16_t tmp_G = Multiply()(tmp_E, tmp_F);
  bfloat16_t tmp_H = cast_to<bfloat16_t>(tmp_G);
  bfloat16_t tmp_I = cast_to<bfloat16_t>(tmp_B);
  bfloat16_t tmp_C = Multiply()(tmp_H, tmp_I);
  C[index] = tmp_C;
}
[[host_name("kernel_qwen_mtplx_swiglu_bf16_strided_3_large")]]
[[kernel]] void CV2ISigmoidADV2IBroadcastACEV2IBroadcastCAFV2IMultiplyDEGV2IBroadcastFBHV2IBroadcastBFIV2OMultiplyGH_VV_V2V2_11160318154034397263_strided_3_large(
    device const bfloat16_t* A [[buffer(0)]],
    device const bfloat16_t* B [[buffer(1)]],
    constant const int64_t* in_strides [[buffer(2)]],
    device bfloat16_t* C [[buffer(3)]],
    constant const int* output_shape [[buffer(4)]],
    uint3 pos [[thread_position_in_grid]],
    uint3 grid [[threads_per_grid]]) {
  constexpr int N_ = 1;
  int64_t index = pos.x + grid.x * (pos.y + int64_t(grid.y) * pos.z);
  int64_t index_A = elem_to_loc_3<int64_t>(pos, in_strides + 0);
  int64_t index_B = elem_to_loc_3<int64_t>(pos, in_strides + 3);
  bfloat16_t tmp_A = A[index_A];
  bfloat16_t tmp_B = B[index_B];
  bfloat16_t tmp_D = Sigmoid()(tmp_A);
  bfloat16_t tmp_E = cast_to<bfloat16_t>(tmp_A);
  bfloat16_t tmp_F = cast_to<bfloat16_t>(tmp_D);
  bfloat16_t tmp_G = Multiply()(tmp_E, tmp_F);
  bfloat16_t tmp_H = cast_to<bfloat16_t>(tmp_G);
  bfloat16_t tmp_I = cast_to<bfloat16_t>(tmp_B);
  bfloat16_t tmp_C = Multiply()(tmp_H, tmp_I);
  C[index] = tmp_C;
}
[[host_name("kernel_qwen_mtplx_swiglu_bf16_strided_4")]]
[[kernel]] void CV2ISigmoidADV2IBroadcastACEV2IBroadcastCAFV2IMultiplyDEGV2IBroadcastFBHV2IBroadcastBFIV2OMultiplyGH_VV_V2V2_11160318154034397263_strided_4(
    device const bfloat16_t* A [[buffer(0)]],
    device const bfloat16_t* B [[buffer(1)]],
    constant const int64_t* in_strides [[buffer(2)]],
    device bfloat16_t* C [[buffer(3)]],
    constant const int* output_shape [[buffer(4)]],
    uint3 pos [[thread_position_in_grid]],
    uint3 grid [[threads_per_grid]]) {
  constexpr int N_ = 2;
  int xshape = output_shape[3];
  uint index = N_ * pos.x + xshape * (pos.y + uint(grid.y) * pos.z);
  uint index_A = N_ * pos.x * uint(in_strides[3]) + pos.y * uint(in_strides[2]);
  uint index_B = N_ * pos.x * uint(in_strides[7]) + pos.y * uint(in_strides[6]);
  uint zpos = pos.z;
  for (int d = 1; d >= 0; --d) {
    uint l = zpos % output_shape[d];
    index_A += l * uint(in_strides[0 + d]);
    index_B += l * uint(in_strides[4 + d]);
    zpos /= output_shape[d];
  }
  for (int i = 0; i < N_ && (int(N_ * pos.x) + i) < xshape; ++i) {
  bfloat16_t tmp_A = A[index_A];
  bfloat16_t tmp_B = B[index_B];
  bfloat16_t tmp_D = Sigmoid()(tmp_A);
  bfloat16_t tmp_E = cast_to<bfloat16_t>(tmp_A);
  bfloat16_t tmp_F = cast_to<bfloat16_t>(tmp_D);
  bfloat16_t tmp_G = Multiply()(tmp_E, tmp_F);
  bfloat16_t tmp_H = cast_to<bfloat16_t>(tmp_G);
  bfloat16_t tmp_I = cast_to<bfloat16_t>(tmp_B);
  bfloat16_t tmp_C = Multiply()(tmp_H, tmp_I);
  C[index] = tmp_C;
  index_A += in_strides[3];
  index_B += in_strides[7];
  index++;
  }
}
[[host_name("kernel_qwen_mtplx_swiglu_bf16_strided_4_large")]]
[[kernel]] void CV2ISigmoidADV2IBroadcastACEV2IBroadcastCAFV2IMultiplyDEGV2IBroadcastFBHV2IBroadcastBFIV2OMultiplyGH_VV_V2V2_11160318154034397263_strided_4_large(
    device const bfloat16_t* A [[buffer(0)]],
    device const bfloat16_t* B [[buffer(1)]],
    constant const int64_t* in_strides [[buffer(2)]],
    device bfloat16_t* C [[buffer(3)]],
    constant const int* output_shape [[buffer(4)]],
    uint3 pos [[thread_position_in_grid]],
    uint3 grid [[threads_per_grid]]) {
  constexpr int N_ = 4;
  int xshape = output_shape[3];
  int64_t index = N_ * pos.x + xshape * (pos.y + int64_t(grid.y) * pos.z);
  int64_t index_A = N_ * pos.x * int64_t(in_strides[3]) + pos.y * int64_t(in_strides[2]);
  int64_t index_B = N_ * pos.x * int64_t(in_strides[7]) + pos.y * int64_t(in_strides[6]);
  uint zpos = pos.z;
  for (int d = 1; d >= 0; --d) {
    uint l = zpos % output_shape[d];
    index_A += l * int64_t(in_strides[0 + d]);
    index_B += l * int64_t(in_strides[4 + d]);
    zpos /= output_shape[d];
  }
  for (int i = 0; i < N_ && (int(N_ * pos.x) + i) < xshape; ++i) {
  bfloat16_t tmp_A = A[index_A];
  bfloat16_t tmp_B = B[index_B];
  bfloat16_t tmp_D = Sigmoid()(tmp_A);
  bfloat16_t tmp_E = cast_to<bfloat16_t>(tmp_A);
  bfloat16_t tmp_F = cast_to<bfloat16_t>(tmp_D);
  bfloat16_t tmp_G = Multiply()(tmp_E, tmp_F);
  bfloat16_t tmp_H = cast_to<bfloat16_t>(tmp_G);
  bfloat16_t tmp_I = cast_to<bfloat16_t>(tmp_B);
  bfloat16_t tmp_C = Multiply()(tmp_H, tmp_I);
  C[index] = tmp_C;
  index_A += in_strides[3];
  index_B += in_strides[7];
  index++;
  }
}
[[host_name("kernel_qwen_mtplx_swiglu_bf16_strided_5")]]
[[kernel]] void CV2ISigmoidADV2IBroadcastACEV2IBroadcastCAFV2IMultiplyDEGV2IBroadcastFBHV2IBroadcastBFIV2OMultiplyGH_VV_V2V2_11160318154034397263_strided_5(
    device const bfloat16_t* A [[buffer(0)]],
    device const bfloat16_t* B [[buffer(1)]],
    constant const int64_t* in_strides [[buffer(2)]],
    device bfloat16_t* C [[buffer(3)]],
    constant const int* output_shape [[buffer(4)]],
    uint3 pos [[thread_position_in_grid]],
    uint3 grid [[threads_per_grid]]) {
  constexpr int N_ = 2;
  int xshape = output_shape[4];
  uint index = N_ * pos.x + xshape * (pos.y + uint(grid.y) * pos.z);
  uint index_A = N_ * pos.x * uint(in_strides[4]) + pos.y * uint(in_strides[3]);
  uint index_B = N_ * pos.x * uint(in_strides[9]) + pos.y * uint(in_strides[8]);
  uint zpos = pos.z;
  for (int d = 2; d >= 0; --d) {
    uint l = zpos % output_shape[d];
    index_A += l * uint(in_strides[0 + d]);
    index_B += l * uint(in_strides[5 + d]);
    zpos /= output_shape[d];
  }
  for (int i = 0; i < N_ && (int(N_ * pos.x) + i) < xshape; ++i) {
  bfloat16_t tmp_A = A[index_A];
  bfloat16_t tmp_B = B[index_B];
  bfloat16_t tmp_D = Sigmoid()(tmp_A);
  bfloat16_t tmp_E = cast_to<bfloat16_t>(tmp_A);
  bfloat16_t tmp_F = cast_to<bfloat16_t>(tmp_D);
  bfloat16_t tmp_G = Multiply()(tmp_E, tmp_F);
  bfloat16_t tmp_H = cast_to<bfloat16_t>(tmp_G);
  bfloat16_t tmp_I = cast_to<bfloat16_t>(tmp_B);
  bfloat16_t tmp_C = Multiply()(tmp_H, tmp_I);
  C[index] = tmp_C;
  index_A += in_strides[4];
  index_B += in_strides[9];
  index++;
  }
}
[[host_name("kernel_qwen_mtplx_swiglu_bf16_strided_5_large")]]
[[kernel]] void CV2ISigmoidADV2IBroadcastACEV2IBroadcastCAFV2IMultiplyDEGV2IBroadcastFBHV2IBroadcastBFIV2OMultiplyGH_VV_V2V2_11160318154034397263_strided_5_large(
    device const bfloat16_t* A [[buffer(0)]],
    device const bfloat16_t* B [[buffer(1)]],
    constant const int64_t* in_strides [[buffer(2)]],
    device bfloat16_t* C [[buffer(3)]],
    constant const int* output_shape [[buffer(4)]],
    uint3 pos [[thread_position_in_grid]],
    uint3 grid [[threads_per_grid]]) {
  constexpr int N_ = 4;
  int xshape = output_shape[4];
  int64_t index = N_ * pos.x + xshape * (pos.y + int64_t(grid.y) * pos.z);
  int64_t index_A = N_ * pos.x * int64_t(in_strides[4]) + pos.y * int64_t(in_strides[3]);
  int64_t index_B = N_ * pos.x * int64_t(in_strides[9]) + pos.y * int64_t(in_strides[8]);
  uint zpos = pos.z;
  for (int d = 2; d >= 0; --d) {
    uint l = zpos % output_shape[d];
    index_A += l * int64_t(in_strides[0 + d]);
    index_B += l * int64_t(in_strides[5 + d]);
    zpos /= output_shape[d];
  }
  for (int i = 0; i < N_ && (int(N_ * pos.x) + i) < xshape; ++i) {
  bfloat16_t tmp_A = A[index_A];
  bfloat16_t tmp_B = B[index_B];
  bfloat16_t tmp_D = Sigmoid()(tmp_A);
  bfloat16_t tmp_E = cast_to<bfloat16_t>(tmp_A);
  bfloat16_t tmp_F = cast_to<bfloat16_t>(tmp_D);
  bfloat16_t tmp_G = Multiply()(tmp_E, tmp_F);
  bfloat16_t tmp_H = cast_to<bfloat16_t>(tmp_G);
  bfloat16_t tmp_I = cast_to<bfloat16_t>(tmp_B);
  bfloat16_t tmp_C = Multiply()(tmp_H, tmp_I);
  C[index] = tmp_C;
  index_A += in_strides[4];
  index_B += in_strides[9];
  index++;
  }
}
[[host_name("kernel_qwen_mtplx_swiglu_bf16_strided_6")]]
[[kernel]] void CV2ISigmoidADV2IBroadcastACEV2IBroadcastCAFV2IMultiplyDEGV2IBroadcastFBHV2IBroadcastBFIV2OMultiplyGH_VV_V2V2_11160318154034397263_strided_6(
    device const bfloat16_t* A [[buffer(0)]],
    device const bfloat16_t* B [[buffer(1)]],
    constant const int64_t* in_strides [[buffer(2)]],
    device bfloat16_t* C [[buffer(3)]],
    constant const int* output_shape [[buffer(4)]],
    uint3 pos [[thread_position_in_grid]],
    uint3 grid [[threads_per_grid]]) {
  constexpr int N_ = 2;
  int xshape = output_shape[5];
  uint index = N_ * pos.x + xshape * (pos.y + uint(grid.y) * pos.z);
  uint index_A = N_ * pos.x * uint(in_strides[5]) + pos.y * uint(in_strides[4]);
  uint index_B = N_ * pos.x * uint(in_strides[11]) + pos.y * uint(in_strides[10]);
  uint zpos = pos.z;
  for (int d = 3; d >= 0; --d) {
    uint l = zpos % output_shape[d];
    index_A += l * uint(in_strides[0 + d]);
    index_B += l * uint(in_strides[6 + d]);
    zpos /= output_shape[d];
  }
  for (int i = 0; i < N_ && (int(N_ * pos.x) + i) < xshape; ++i) {
  bfloat16_t tmp_A = A[index_A];
  bfloat16_t tmp_B = B[index_B];
  bfloat16_t tmp_D = Sigmoid()(tmp_A);
  bfloat16_t tmp_E = cast_to<bfloat16_t>(tmp_A);
  bfloat16_t tmp_F = cast_to<bfloat16_t>(tmp_D);
  bfloat16_t tmp_G = Multiply()(tmp_E, tmp_F);
  bfloat16_t tmp_H = cast_to<bfloat16_t>(tmp_G);
  bfloat16_t tmp_I = cast_to<bfloat16_t>(tmp_B);
  bfloat16_t tmp_C = Multiply()(tmp_H, tmp_I);
  C[index] = tmp_C;
  index_A += in_strides[5];
  index_B += in_strides[11];
  index++;
  }
}
[[host_name("kernel_qwen_mtplx_swiglu_bf16_strided_6_large")]]
[[kernel]] void CV2ISigmoidADV2IBroadcastACEV2IBroadcastCAFV2IMultiplyDEGV2IBroadcastFBHV2IBroadcastBFIV2OMultiplyGH_VV_V2V2_11160318154034397263_strided_6_large(
    device const bfloat16_t* A [[buffer(0)]],
    device const bfloat16_t* B [[buffer(1)]],
    constant const int64_t* in_strides [[buffer(2)]],
    device bfloat16_t* C [[buffer(3)]],
    constant const int* output_shape [[buffer(4)]],
    uint3 pos [[thread_position_in_grid]],
    uint3 grid [[threads_per_grid]]) {
  constexpr int N_ = 4;
  int xshape = output_shape[5];
  int64_t index = N_ * pos.x + xshape * (pos.y + int64_t(grid.y) * pos.z);
  int64_t index_A = N_ * pos.x * int64_t(in_strides[5]) + pos.y * int64_t(in_strides[4]);
  int64_t index_B = N_ * pos.x * int64_t(in_strides[11]) + pos.y * int64_t(in_strides[10]);
  uint zpos = pos.z;
  for (int d = 3; d >= 0; --d) {
    uint l = zpos % output_shape[d];
    index_A += l * int64_t(in_strides[0 + d]);
    index_B += l * int64_t(in_strides[6 + d]);
    zpos /= output_shape[d];
  }
  for (int i = 0; i < N_ && (int(N_ * pos.x) + i) < xshape; ++i) {
  bfloat16_t tmp_A = A[index_A];
  bfloat16_t tmp_B = B[index_B];
  bfloat16_t tmp_D = Sigmoid()(tmp_A);
  bfloat16_t tmp_E = cast_to<bfloat16_t>(tmp_A);
  bfloat16_t tmp_F = cast_to<bfloat16_t>(tmp_D);
  bfloat16_t tmp_G = Multiply()(tmp_E, tmp_F);
  bfloat16_t tmp_H = cast_to<bfloat16_t>(tmp_G);
  bfloat16_t tmp_I = cast_to<bfloat16_t>(tmp_B);
  bfloat16_t tmp_C = Multiply()(tmp_H, tmp_I);
  C[index] = tmp_C;
  index_A += in_strides[5];
  index_B += in_strides[11];
  index++;
  }
}
[[host_name("kernel_qwen_mtplx_swiglu_bf16_strided_7")]]
[[kernel]] void CV2ISigmoidADV2IBroadcastACEV2IBroadcastCAFV2IMultiplyDEGV2IBroadcastFBHV2IBroadcastBFIV2OMultiplyGH_VV_V2V2_11160318154034397263_strided_7(
    device const bfloat16_t* A [[buffer(0)]],
    device const bfloat16_t* B [[buffer(1)]],
    constant const int64_t* in_strides [[buffer(2)]],
    device bfloat16_t* C [[buffer(3)]],
    constant const int* output_shape [[buffer(4)]],
    uint3 pos [[thread_position_in_grid]],
    uint3 grid [[threads_per_grid]]) {
  constexpr int N_ = 2;
  int xshape = output_shape[6];
  uint index = N_ * pos.x + xshape * (pos.y + uint(grid.y) * pos.z);
  uint index_A = N_ * pos.x * uint(in_strides[6]) + pos.y * uint(in_strides[5]);
  uint index_B = N_ * pos.x * uint(in_strides[13]) + pos.y * uint(in_strides[12]);
  uint zpos = pos.z;
  for (int d = 4; d >= 0; --d) {
    uint l = zpos % output_shape[d];
    index_A += l * uint(in_strides[0 + d]);
    index_B += l * uint(in_strides[7 + d]);
    zpos /= output_shape[d];
  }
  for (int i = 0; i < N_ && (int(N_ * pos.x) + i) < xshape; ++i) {
  bfloat16_t tmp_A = A[index_A];
  bfloat16_t tmp_B = B[index_B];
  bfloat16_t tmp_D = Sigmoid()(tmp_A);
  bfloat16_t tmp_E = cast_to<bfloat16_t>(tmp_A);
  bfloat16_t tmp_F = cast_to<bfloat16_t>(tmp_D);
  bfloat16_t tmp_G = Multiply()(tmp_E, tmp_F);
  bfloat16_t tmp_H = cast_to<bfloat16_t>(tmp_G);
  bfloat16_t tmp_I = cast_to<bfloat16_t>(tmp_B);
  bfloat16_t tmp_C = Multiply()(tmp_H, tmp_I);
  C[index] = tmp_C;
  index_A += in_strides[6];
  index_B += in_strides[13];
  index++;
  }
}
[[host_name("kernel_qwen_mtplx_swiglu_bf16_strided_7_large")]]
[[kernel]] void CV2ISigmoidADV2IBroadcastACEV2IBroadcastCAFV2IMultiplyDEGV2IBroadcastFBHV2IBroadcastBFIV2OMultiplyGH_VV_V2V2_11160318154034397263_strided_7_large(
    device const bfloat16_t* A [[buffer(0)]],
    device const bfloat16_t* B [[buffer(1)]],
    constant const int64_t* in_strides [[buffer(2)]],
    device bfloat16_t* C [[buffer(3)]],
    constant const int* output_shape [[buffer(4)]],
    uint3 pos [[thread_position_in_grid]],
    uint3 grid [[threads_per_grid]]) {
  constexpr int N_ = 4;
  int xshape = output_shape[6];
  int64_t index = N_ * pos.x + xshape * (pos.y + int64_t(grid.y) * pos.z);
  int64_t index_A = N_ * pos.x * int64_t(in_strides[6]) + pos.y * int64_t(in_strides[5]);
  int64_t index_B = N_ * pos.x * int64_t(in_strides[13]) + pos.y * int64_t(in_strides[12]);
  uint zpos = pos.z;
  for (int d = 4; d >= 0; --d) {
    uint l = zpos % output_shape[d];
    index_A += l * int64_t(in_strides[0 + d]);
    index_B += l * int64_t(in_strides[7 + d]);
    zpos /= output_shape[d];
  }
  for (int i = 0; i < N_ && (int(N_ * pos.x) + i) < xshape; ++i) {
  bfloat16_t tmp_A = A[index_A];
  bfloat16_t tmp_B = B[index_B];
  bfloat16_t tmp_D = Sigmoid()(tmp_A);
  bfloat16_t tmp_E = cast_to<bfloat16_t>(tmp_A);
  bfloat16_t tmp_F = cast_to<bfloat16_t>(tmp_D);
  bfloat16_t tmp_G = Multiply()(tmp_E, tmp_F);
  bfloat16_t tmp_H = cast_to<bfloat16_t>(tmp_G);
  bfloat16_t tmp_I = cast_to<bfloat16_t>(tmp_B);
  bfloat16_t tmp_C = Multiply()(tmp_H, tmp_I);
  C[index] = tmp_C;
  index_A += in_strides[6];
  index_B += in_strides[13];
  index++;
  }
}
[[host_name("kernel_qwen_mtplx_swiglu_bf16_strided_dynamic")]]
[[kernel]] void CV2ISigmoidADV2IBroadcastACEV2IBroadcastCAFV2IMultiplyDEGV2IBroadcastFBHV2IBroadcastBFIV2OMultiplyGH_VV_V2V2_11160318154034397263_strided_dynamic(
    device const bfloat16_t* A [[buffer(0)]],
    device const bfloat16_t* B [[buffer(1)]],
    constant const int64_t* in_strides [[buffer(2)]],
    device bfloat16_t* C [[buffer(3)]],
    constant const int* output_shape [[buffer(4)]],
    constant const int& ndim [[buffer(5)]],
    uint3 pos [[thread_position_in_grid]],
    uint3 grid [[threads_per_grid]]) {
  constexpr int N_ = 2;
  int xshape = output_shape[ndim - 1];
  uint index = N_ * pos.x + xshape * (pos.y + uint(grid.y) * pos.z);
  uint index_A = N_ * pos.x * uint(in_strides[ndim * 0 + ndim - 1]) + pos.y * uint(in_strides[ndim * 0 + ndim - 2]);
  uint index_B = N_ * pos.x * uint(in_strides[ndim * 1 + ndim - 1]) + pos.y * uint(in_strides[ndim * 1 + ndim - 2]);
  uint zpos = pos.z;
  for (int d = ndim - 3; d >= 0; --d) {
    uint l = zpos % output_shape[d];
    index_A += l * uint(in_strides[0 * ndim + d]);
    index_B += l * uint(in_strides[1 * ndim + d]);
    zpos /= output_shape[d];
  }
  for (int i = 0; i < N_ && (int(N_ * pos.x) + i) < xshape; ++i) {
  bfloat16_t tmp_A = A[index_A];
  bfloat16_t tmp_B = B[index_B];
  bfloat16_t tmp_D = Sigmoid()(tmp_A);
  bfloat16_t tmp_E = cast_to<bfloat16_t>(tmp_A);
  bfloat16_t tmp_F = cast_to<bfloat16_t>(tmp_D);
  bfloat16_t tmp_G = Multiply()(tmp_E, tmp_F);
  bfloat16_t tmp_H = cast_to<bfloat16_t>(tmp_G);
  bfloat16_t tmp_I = cast_to<bfloat16_t>(tmp_B);
  bfloat16_t tmp_C = Multiply()(tmp_H, tmp_I);
  C[index] = tmp_C;
  index_A += in_strides[0 * ndim + ndim - 1];
  index_B += in_strides[1 * ndim + ndim - 1];
  index++;
  }
}
[[host_name("kernel_qwen_mtplx_swiglu_bf16_strided_dynamic_large")]]
[[kernel]] void CV2ISigmoidADV2IBroadcastACEV2IBroadcastCAFV2IMultiplyDEGV2IBroadcastFBHV2IBroadcastBFIV2OMultiplyGH_VV_V2V2_11160318154034397263_strided_dynamic_large(
    device const bfloat16_t* A [[buffer(0)]],
    device const bfloat16_t* B [[buffer(1)]],
    constant const int64_t* in_strides [[buffer(2)]],
    device bfloat16_t* C [[buffer(3)]],
    constant const int* output_shape [[buffer(4)]],
    constant const int& ndim [[buffer(5)]],
    uint3 pos [[thread_position_in_grid]],
    uint3 grid [[threads_per_grid]]) {
  constexpr int N_ = 4;
  int xshape = output_shape[ndim - 1];
  int64_t index = N_ * pos.x + xshape * (pos.y + int64_t(grid.y) * pos.z);
  int64_t index_A = N_ * pos.x * int64_t(in_strides[ndim * 0 + ndim - 1]) + pos.y * int64_t(in_strides[ndim * 0 + ndim - 2]);
  int64_t index_B = N_ * pos.x * int64_t(in_strides[ndim * 1 + ndim - 1]) + pos.y * int64_t(in_strides[ndim * 1 + ndim - 2]);
  uint zpos = pos.z;
  for (int d = ndim - 3; d >= 0; --d) {
    uint l = zpos % output_shape[d];
    index_A += l * int64_t(in_strides[0 * ndim + d]);
    index_B += l * int64_t(in_strides[1 * ndim + d]);
    zpos /= output_shape[d];
  }
  for (int i = 0; i < N_ && (int(N_ * pos.x) + i) < xshape; ++i) {
  bfloat16_t tmp_A = A[index_A];
  bfloat16_t tmp_B = B[index_B];
  bfloat16_t tmp_D = Sigmoid()(tmp_A);
  bfloat16_t tmp_E = cast_to<bfloat16_t>(tmp_A);
  bfloat16_t tmp_F = cast_to<bfloat16_t>(tmp_D);
  bfloat16_t tmp_G = Multiply()(tmp_E, tmp_F);
  bfloat16_t tmp_H = cast_to<bfloat16_t>(tmp_G);
  bfloat16_t tmp_I = cast_to<bfloat16_t>(tmp_B);
  bfloat16_t tmp_C = Multiply()(tmp_H, tmp_I);
  C[index] = tmp_C;
  index_A += in_strides[0 * ndim + ndim - 1];
  index_B += in_strides[1 * ndim + ndim - 1];
  index++;
  }
}
// END RUNTIME UNIT

// Original compute_g JIT dependencies; Apple MIT except cexpf.h (Apache-2.0).
// Runtime unit: complex.h; file SHA256: 16e8a815b2cbdb6070e0824e64fe33fccb6e918f1b84ea5c792bd89d33e57bf1
// Runtime unit SHA256: 0393eb41ea618f7b09343672f598d8c4c3e4b98eb8536466fb53c8d286fbf1f7
// BEGIN RUNTIME UNIT
// Copyright © 2023 Apple Inc.




using namespace metal;

template <typename T>
struct complex_t;

template <typename T>
static constexpr constant bool is_complex_v = false;

template <typename T>
static constexpr constant bool is_complex_v<complex_t<T>> = true;

// Metal accepts explicit bfloat casts that is_convertible_v reports as false.
template <typename From, typename To>
static constexpr constant bool is_lane_convertible_v =
    is_convertible_v<From, To> ||
    (is_same_v<To, bfloat16_t> && is_convertible_v<From, float>) ||
    (is_same_v<From, bfloat16_t> && is_convertible_v<float, To>);

template <typename T>
struct complex_t {
  using value_type = T;

  T real;
  T imag;

  // Constructors
  constexpr complex_t(T real, T imag) thread : real(real), imag(imag) {};
  constexpr complex_t() thread : real(0), imag(0) {};
  constexpr complex_t() threadgroup : real(0), imag(0) {};

  // Conversions from scalar types
  template <
      typename U,
      typename = typename enable_if<
          !is_complex_v<U> && is_lane_convertible_v<U, T>>::type>
  constexpr complex_t(U x) thread : real(static_cast<T>(x)),
                                    imag(static_cast<T>(0)) {}

  template <
      typename U,
      typename = typename enable_if<
          !is_complex_v<U> && is_lane_convertible_v<U, T>>::type>
  constexpr complex_t(U x) threadgroup : real(static_cast<T>(x)),
                                         imag(static_cast<T>(0)) {}

  template <
      typename U,
      typename = typename enable_if<
          !is_complex_v<U> && is_lane_convertible_v<U, T>>::type>
  constexpr complex_t(U x) device : real(static_cast<T>(x)),
                                    imag(static_cast<T>(0)) {}

  template <
      typename U,
      typename = typename enable_if<
          !is_complex_v<U> && is_lane_convertible_v<U, T>>::type>
  constexpr complex_t(U x) constant : real(static_cast<T>(x)),
                                      imag(static_cast<T>(0)) {}

  // Conversions between complex types
  template <
      typename U,
      typename = typename enable_if<
          !is_same_v<U, T> && is_lane_convertible_v<U, T>>::type>
  constexpr complex_t(complex_t<U> x) thread : real(static_cast<T>(x.real)),
                                               imag(static_cast<T>(x.imag)) {}

  template <
      typename U,
      typename = typename enable_if<
          !is_same_v<U, T> && is_lane_convertible_v<U, T>>::type>
  constexpr complex_t(complex_t<U> x) threadgroup
      : real(static_cast<T>(x.real)),
        imag(static_cast<T>(x.imag)) {}

  template <
      typename U,
      typename = typename enable_if<
          !is_same_v<U, T> && is_lane_convertible_v<U, T>>::type>
  constexpr complex_t(complex_t<U> x) device : real(static_cast<T>(x.real)),
                                               imag(static_cast<T>(x.imag)) {}

  template <
      typename U,
      typename = typename enable_if<
          !is_same_v<U, T> && is_lane_convertible_v<U, T>>::type>
  constexpr complex_t(complex_t<U> x) constant : real(static_cast<T>(x.real)),
                                                 imag(static_cast<T>(x.imag)) {}

  // Conversions to and from two-lane vectors (the FFT lane representation)
  constexpr complex_t(vec<T, 2> v) thread : real(v.x), imag(v.y) {};
  constexpr complex_t(vec<T, 2> v) threadgroup : real(v.x), imag(v.y) {};
  constexpr complex_t(vec<T, 2> v) device : real(v.x), imag(v.y) {};
  constexpr complex_t(vec<T, 2> v) constant : real(v.x), imag(v.y) {};

  constexpr operator vec<T, 2>() const thread {
    return vec<T, 2>(real, imag);
  }

  constexpr operator vec<T, 2>() const threadgroup {
    return vec<T, 2>(real, imag);
  }

  constexpr operator vec<T, 2>() const device {
    return vec<T, 2>(real, imag);
  }

  constexpr operator vec<T, 2>() const constant {
    return vec<T, 2>(real, imag);
  }

  // Conversions to scalar types
  template <
      typename U,
      typename = typename enable_if<
          !is_complex_v<U> && is_lane_convertible_v<T, U>>::type>
  constexpr operator U() const thread {
    return static_cast<U>(real);
  }

  template <
      typename U,
      typename = typename enable_if<
          !is_complex_v<U> && is_lane_convertible_v<T, U>>::type>
  constexpr operator U() const threadgroup {
    return static_cast<U>(real);
  }

  template <
      typename U,
      typename = typename enable_if<
          !is_complex_v<U> && is_lane_convertible_v<T, U>>::type>
  constexpr operator U() const device {
    return static_cast<U>(real);
  }

  template <
      typename U,
      typename = typename enable_if<
          !is_complex_v<U> && is_lane_convertible_v<T, U>>::type>
  constexpr operator U() const constant {
    return static_cast<U>(real);
  }
};

using complex32_t = complex_t<half>;
using complex64_t = complex_t<float>;

static_assert(sizeof(complex32_t) == 2 * sizeof(half));
static_assert(sizeof(complex64_t) == 2 * sizeof(float));
static_assert(sizeof(complex_t<bfloat16_t>) == 2 * sizeof(bfloat16_t));

template <typename T>
constexpr complex_t<T> operator-(complex_t<T> x) {
  return {-x.real, -x.imag};
}

template <typename T>
constexpr bool operator>=(complex_t<T> a, complex_t<T> b) {
  return (a.real > b.real) || (a.real == b.real && a.imag >= b.imag);
}

template <typename T>
constexpr bool operator>(complex_t<T> a, complex_t<T> b) {
  return (a.real > b.real) || (a.real == b.real && a.imag > b.imag);
}

template <typename T>
constexpr bool operator<=(complex_t<T> a, complex_t<T> b) {
  return operator>=(b, a);
}

template <typename T>
constexpr bool operator<(complex_t<T> a, complex_t<T> b) {
  return operator>(b, a);
}

template <typename T>
constexpr bool operator==(complex_t<T> a, complex_t<T> b) {
  return a.real == b.real && a.imag == b.imag;
}

template <typename T>
constexpr complex_t<T> operator+(complex_t<T> a, complex_t<T> b) {
  return {a.real + b.real, a.imag + b.imag};
}

template <typename T>
constexpr thread complex_t<T>& operator+=(
    thread complex_t<T>& a,
    complex_t<T> b) {
  a.real += b.real;
  a.imag += b.imag;
  return a;
}

template <typename T>
constexpr threadgroup complex_t<T>& operator+=(
    threadgroup complex_t<T>& a,
    complex_t<T> b) {
  a.real += b.real;
  a.imag += b.imag;
  return a;
}

template <typename T>
constexpr device complex_t<T>& operator+=(
    device complex_t<T>& a,
    complex_t<T> b) {
  a.real += b.real;
  a.imag += b.imag;
  return a;
}

template <
    typename T,
    typename U,
    enable_if_t<!is_complex_v<U> && is_lane_convertible_v<U, T>, bool> = true>
constexpr complex_t<T> operator+(U a, complex_t<T> b) {
  return {static_cast<T>(a) + b.real, b.imag};
}

template <
    typename T,
    typename U,
    enable_if_t<!is_complex_v<U> && is_lane_convertible_v<U, T>, bool> = true>
constexpr complex_t<T> operator+(complex_t<T> a, U b) {
  return {a.real + static_cast<T>(b), a.imag};
}

template <typename T>
constexpr complex_t<T> operator-(complex_t<T> a, complex_t<T> b) {
  return {a.real - b.real, a.imag - b.imag};
}

template <
    typename T,
    typename U,
    enable_if_t<!is_complex_v<U> && is_lane_convertible_v<U, T>, bool> = true>
constexpr complex_t<T> operator-(U a, complex_t<T> b) {
  return {static_cast<T>(a) - b.real, -b.imag};
}

template <
    typename T,
    typename U,
    enable_if_t<!is_complex_v<U> && is_lane_convertible_v<U, T>, bool> = true>
constexpr complex_t<T> operator-(complex_t<T> a, U b) {
  return {a.real - static_cast<T>(b), a.imag};
}

template <typename T>
constexpr complex_t<T> operator*(complex_t<T> a, complex_t<T> b) {
  return {a.real * b.real - a.imag * b.imag, a.real * b.imag + a.imag * b.real};
}

template <typename T>
constexpr complex_t<T> operator/(complex_t<T> a, complex_t<T> b) {
  auto denom = b.real * b.real + b.imag * b.imag;
  auto x = a.real * b.real + a.imag * b.imag;
  auto y = a.imag * b.real - a.real * b.imag;
  return {x / denom, y / denom};
}

template <
    typename T,
    typename U,
    enable_if_t<!is_complex_v<U> && is_lane_convertible_v<U, T>, bool> = true>
constexpr complex_t<T> operator/(U a, complex_t<T> b) {
  auto scalar = static_cast<T>(a);
  auto denom = b.real * b.real + b.imag * b.imag;
  auto x = scalar * b.real;
  auto y = -scalar * b.imag;
  return {x / denom, y / denom};
}

template <typename T>
constexpr complex_t<T> operator%(complex_t<T> a, complex_t<T> b) {
  auto real = a.real - (b.real * static_cast<int64_t>(a.real / b.real));
  auto imag = a.imag - (b.imag * static_cast<int64_t>(a.imag / b.imag));
  if (real != 0 && (real < 0 != b.real < 0)) {
    real += b.real;
  }
  if (imag != 0 && (imag < 0 != b.imag < 0)) {
    imag += b.imag;
  }
  return {real, imag};
}

static_assert(
    (complex_t<half>{1.0h, 2.0h} * complex_t<half>{3.0h, 4.0h}).real == -5.0h);
static_assert(
    (complex_t<bfloat16_t>{bfloat16_t(1.0f), bfloat16_t(2.0f)} *
     complex_t<bfloat16_t>{bfloat16_t(3.0f), bfloat16_t(4.0f)})
        .real == bfloat16_t(-5.0f));
// END RUNTIME UNIT
// Runtime unit: cexpf.h; file SHA256: 88b6e15a52a5800d98d9bc6da840ca5cf70bf572fda136409580c1f17b1e0aab
// Runtime unit SHA256: 772a0d51398a8a5cc0d8776e32e3b4a7e80323e58e80dd81da426c03c3bb2ad4
// BEGIN RUNTIME UNIT
// Copyright © 2025 Apple Inc.
// Copyright © 2008-2013 NVIDIA Corporation
// Copyright © 2013 Filipe RNC Maia
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
//     http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.
//
// Forked from
// https://github.com/NVIDIA/cccl/blob/main/thrust/thrust/detail/complex/cexpf.h

// TODO: We should use thrust::exp but the thrust header in old CUDA versions
// can not be used in JIT.



using ieee_float_shape_type = union {
  float value;
  uint32_t word;
};

inline void get_float_word(thread uint32_t& i, float d) {
  ieee_float_shape_type gf_u;
  gf_u.value = (d);
  (i) = gf_u.word;
}

inline void get_float_word(thread int32_t& i, float d) {
  ieee_float_shape_type gf_u;
  gf_u.value = (d);
  (i) = gf_u.word;
}

inline void set_float_word(thread float& d, uint32_t i) {
  ieee_float_shape_type sf_u;
  sf_u.word = (i);
  (d) = sf_u.value;
}

inline float frexp_expf(float x, thread int* expt) {
  const uint32_t k = 235;
  const float kln2 = 162.88958740F;

  float exp_x;
  uint32_t hx;

  exp_x = metal::exp(x - kln2);
  get_float_word(hx, exp_x);
  *expt = (hx >> 23) - (0x7f + 127) + k;
  set_float_word(exp_x, (hx & 0x7fffff) | ((0x7f + 127) << 23));
  return exp_x;
}

inline complex64_t ldexp_cexpf(complex64_t z, int expt) {
  float x, y, exp_x, scale1, scale2;
  int ex_expt, half_expt;

  x = z.real;
  y = z.imag;
  exp_x = frexp_expf(x, &ex_expt);
  expt += ex_expt;

  half_expt = expt / 2;
  set_float_word(scale1, (0x7f + half_expt) << 23);
  half_expt = expt - half_expt;
  set_float_word(scale2, (0x7f + half_expt) << 23);

  return complex64_t{
      metal::cos(y) * exp_x * scale1 * scale2,
      metal::sin(y) * exp_x * scale1 * scale2};
}

inline complex64_t cexpf(const thread complex64_t& z) {
  float x, y, exp_x;
  uint32_t hx, hy;

  const uint32_t exp_ovfl = 0x42b17218, cexp_ovfl = 0x43400074;

  x = z.real;
  y = z.imag;

  get_float_word(hy, y);
  hy &= 0x7fffffff;

  /* cexp(x + I 0) = exp(x) + I 0 */
  if (hy == 0) {
    return complex64_t{metal::exp(x), y};
  }
  get_float_word(hx, x);
  /* cexp(0 + I y) = cos(y) + I sin(y) */
  if ((hx & 0x7fffffff) == 0) {
    return complex64_t{metal::cos(y), metal::sin(y)};
  }
  if (hy >= 0x7f800000) {
    if ((hx & 0x7fffffff) != 0x7f800000) {
      /* cexp(finite|NaN +- I Inf|NaN) = NaN + I NaN */
      return complex64_t{y - y, y - y};
    } else if (hx & 0x80000000) {
      /* cexp(-Inf +- I Inf|NaN) = 0 + I 0 */
      return complex64_t{0.0, 0.0};
    } else {
      /* cexp(+Inf +- I Inf|NaN) = Inf + I NaN */
      return complex64_t{x, y - y};
    }
  }

  if (hx >= exp_ovfl && hx <= cexp_ovfl) {
    /*
     * x is between 88.7 and 192, so we must scale to avoid
     * overflow in expf(x).
     */
    return ldexp_cexpf(z, 0);
  } else {
    /*
     * Cases covered here:
     *  -  x < exp_ovfl and exp(x) won't overflow (common case)
     *  -  x > cexp_ovfl, so exp(x) * s overflows for all s > 0
     *  -  x = +-Inf (generated by exp())
     *  -  x = NaN (spurious inexact exception from y)
     */
    exp_x = metal::exp(x);
    return complex64_t{exp_x * metal::cos(y), exp_x * metal::sin(y)};
  }
}
// END RUNTIME UNIT
// Runtime unit: utils.h; file SHA256: 5e1568e9edde9d05dbf86f68fa0d6c6240f2c32b973c7c6a76166b9c0d91543d
// Runtime unit SHA256: f2f4265df9688cd7ade514d56e7692eff22844a8ae768d0e861813c9590f3eda
// BEGIN RUNTIME UNIT
template <typename U>
struct Limits {
  static const constant U max = metal::numeric_limits<U>::max();
  static const constant U min = metal::numeric_limits<U>::min();
  static const constant U finite_max = metal::numeric_limits<U>::max();
  static const constant U finite_min = metal::numeric_limits<U>::min();
};

#define instantiate_default_limit(type)                                      \
  template <>                                                                \
  struct Limits<type> {                                                      \
    static constexpr constant type max = metal::numeric_limits<type>::max(); \
    static constexpr constant type min = metal::numeric_limits<type>::min(); \
    static constexpr constant type finite_max =                              \
        metal::numeric_limits<type>::max();                                  \
    static constexpr constant type finite_min =                              \
        metal::numeric_limits<type>::min();                                  \
  };

instantiate_default_limit(uint8_t);
instantiate_default_limit(uint16_t);
instantiate_default_limit(uint32_t);
instantiate_default_limit(uint64_t);
instantiate_default_limit(int8_t);
instantiate_default_limit(int16_t);
instantiate_default_limit(int32_t);
instantiate_default_limit(int64_t);

#define instantiate_float_limit(type)             \
  template <>                                     \
  struct Limits<type> {                           \
    static constexpr constant type max =          \
        metal::numeric_limits<type>::infinity();  \
    static constexpr constant type min =          \
        -metal::numeric_limits<type>::infinity(); \
    static constexpr constant type finite_max =   \
        metal::numeric_limits<type>::max();       \
    static constexpr constant type finite_min =   \
        -metal::numeric_limits<type>::max();      \
  };

instantiate_float_limit(half);
instantiate_float_limit(float);
instantiate_float_limit(bfloat16_t);

template <>
struct Limits<bool> {
  static constexpr constant bool max = true;
  static constexpr constant bool min = false;
};

template <typename T>
struct Limits<complex_t<T>> {
  inline static constexpr constant complex_t<T> max = complex_t<T>(
      metal::numeric_limits<T>::infinity(),
      metal::numeric_limits<T>::infinity());
  inline static constexpr constant complex_t<T> min = complex_t<T>(
      -metal::numeric_limits<T>::infinity(),
      -metal::numeric_limits<T>::infinity());
};

inline float log1p(float x) {
  float xp1 = 1.0f + x;
  if (xp1 == Limits<float>::max) {
    return Limits<float>::max;
  }
  if (xp1 == 1.0f) {
    return x;
  }

  return x * (metal::log(xp1) / (xp1 - 1.0f));
}

inline bfloat16_t log1p(bfloat16_t x) {
  float xp1 = 1.0f + static_cast<float>(x);
  if (xp1 == Limits<float>::max) {
    return Limits<bfloat16_t>::max;
  }
  if (xp1 == 1.0f) {
    return x;
  }

  return bfloat16_t(x * (metal::log(xp1) / (xp1 - 1.0f)));
}

inline complex64_t log1p(complex64_t in) {
  float x = in.real;
  float y = in.imag;
  float zabs = metal::precise::sqrt(x * x + y * y);
  float theta = metal::atan2(y, x + 1);
  if (zabs < 0.5f) {
    float r = x * (2 + x) + y * y;
    if (r == 0) { // handle underflow
      return {x, theta};
    }
    return {0.5f * log1p(r), theta};
  } else {
    auto z0 = metal::sqrt((x + 1) * (x + 1) + y * y);
    return {metal::log(z0), theta};
  }
}

// END RUNTIME UNIT
// Runtime unit: unary_ops.h; file SHA256: 0a5492b65ae39ecb6d8b04e64ea5007e0a4ff60d8d9428559bfbfd03387ece2a
// Runtime unit SHA256: 167c0343e9f1257ae104ea8edfe17878e0721e1ab2b1d142ec6eeb5d010344ca
// BEGIN RUNTIME UNIT
struct Exp {
  template <typename T>
  T operator()(T x) thread {
    return metal::precise::exp(x);
  };
  complex64_t operator()(complex64_t x) thread {
    return cexpf(x);
  }
};

struct Negative {
  template <typename T>
  T operator()(T x) thread {
    return -x;
  };
};

// END RUNTIME UNIT
// Runtime unit: binary_ops.h; file SHA256: 2dd13c2496f5d6f0856e4ca99db2ceb5c6d7dc0c344b9ba8e41ef3b7d7ed8a97
// Runtime unit SHA256: bfbc13f52cc73ee117f99d3167d8ac53c33a61ecef1ab811f6cc32266d44cfd4
// BEGIN RUNTIME UNIT
struct Add {
  template <typename T>
  T operator()(T x, T y) thread {
    return x + y;
  }
};

struct LogAddExp {
  template <typename T>
  T operator()(T x, T y) thread {
    if (metal::isnan(x) || metal::isnan(y)) {
      return metal::numeric_limits<T>::quiet_NaN();
    }
    constexpr T inf = metal::numeric_limits<T>::infinity();
    T maxval = metal::max(x, y);
    T minval = metal::min(x, y);
    return (minval == -inf || maxval == inf)
        ? maxval
        : (maxval + log1p(metal::exp(minval - maxval)));
  };

  complex64_t operator()(complex64_t x, complex64_t y) thread {
    if (metal::isnan(x.real) || metal::isnan(x.imag) || metal::isnan(y.real) ||
        metal::isnan(y.imag)) {
      return metal::numeric_limits<float>::quiet_NaN();
    }
    constexpr float inf = metal::numeric_limits<float>::infinity();
    complex64_t maxval = x > y ? x : y;
    complex64_t minval = x < y ? x : y;
    if (minval.real == -inf || maxval.real == inf)
      return maxval;
    float m = metal::exp(minval.real - maxval.real);
    complex64_t dexp{
        m * metal::cos(minval.imag - maxval.imag),
        m * metal::sin(minval.imag - maxval.imag),
    };
    return maxval + log1p(dexp);
  }
};

// END RUNTIME UNIT
// Runtime unit: captured-compute_g-jit; file SHA256: 701e2f54b7cb8bf97616f83256657e6c5e8cc8b46f4b65c559ccc030ab111dbf
// Runtime unit SHA256: 670ec574f1274bea8709b585f5b1383aa3f511d72472465f559c5c0c482bbf1f
// BEGIN RUNTIME UNIT
[[host_name("kernel_qwen_mtplx_compute_g_bf16_contiguous")]]
[[kernel]] void Ef4IAsTypeAFf4IExpEGf4INegativeFHV2IBroadcastBCIV2IBroadcastCBJV2IAddHIKV2IBroadcastJDLV2IBroadcastDJMV2ILogAddExpKLNf4IAsTypeMOf4IBroadcastGNPf4IBroadcastNGQf4IMultiplyOPRf4OExpQ_VVVC_V2V2V2_10408321403207385874_contiguous(
    device const bfloat16_t* A [[buffer(0)]],
    device const bfloat16_t* B [[buffer(1)]],
    device const bfloat16_t* C [[buffer(2)]],
    device float* D [[buffer(3)]],
    constant const uint& size [[buffer(4)]],
    uint3 pos [[thread_position_in_grid]],
    uint3 grid [[threads_per_grid]]) {
  constexpr int N_ = 1;
  uint index = N_ * pos.x;
  bfloat16_t tmp_A = A[index];
  bfloat16_t tmp_B = B[index];
  bfloat16_t tmp_C = C[index];
  auto tmp_E = static_cast<bfloat16_t>(0);
  float tmp_F = cast_to<float>(tmp_A);
  float tmp_G = Exp()(tmp_F);
  float tmp_H = Negative()(tmp_G);
  bfloat16_t tmp_I = cast_to<bfloat16_t>(tmp_B);
  bfloat16_t tmp_J = cast_to<bfloat16_t>(tmp_C);
  bfloat16_t tmp_K = Add()(tmp_I, tmp_J);
  bfloat16_t tmp_L = cast_to<bfloat16_t>(tmp_K);
  bfloat16_t tmp_M = cast_to<bfloat16_t>(tmp_E);
  bfloat16_t tmp_N = LogAddExp()(tmp_L, tmp_M);
  float tmp_O = cast_to<float>(tmp_N);
  float tmp_P = cast_to<float>(tmp_H);
  float tmp_Q = cast_to<float>(tmp_O);
  float tmp_R = Multiply()(tmp_P, tmp_Q);
  float tmp_D = Exp()(tmp_R);
  D[index] = tmp_D;
}
[[host_name("kernel_qwen_mtplx_compute_g_bf16_contiguous_n")]]
[[kernel]] void Ef4IAsTypeAFf4IExpEGf4INegativeFHV2IBroadcastBCIV2IBroadcastCBJV2IAddHIKV2IBroadcastJDLV2IBroadcastDJMV2ILogAddExpKLNf4IAsTypeMOf4IBroadcastGNPf4IBroadcastNGQf4IMultiplyOPRf4OExpQ_VVVC_V2V2V2_10408321403207385874_contiguous_n(
    device const bfloat16_t* A [[buffer(0)]],
    device const bfloat16_t* B [[buffer(1)]],
    device const bfloat16_t* C [[buffer(2)]],
    device float* D [[buffer(3)]],
    constant const uint& size [[buffer(4)]],
    uint3 pos [[thread_position_in_grid]],
    uint3 grid [[threads_per_grid]]) {
  constexpr int N_ = 2;
  uint index = N_ * pos.x;
  for (int i = 0; i < N_ && index < size; ++i) {
  bfloat16_t tmp_A = A[index];
  bfloat16_t tmp_B = B[index];
  bfloat16_t tmp_C = C[index];
  auto tmp_E = static_cast<bfloat16_t>(0);
  float tmp_F = cast_to<float>(tmp_A);
  float tmp_G = Exp()(tmp_F);
  float tmp_H = Negative()(tmp_G);
  bfloat16_t tmp_I = cast_to<bfloat16_t>(tmp_B);
  bfloat16_t tmp_J = cast_to<bfloat16_t>(tmp_C);
  bfloat16_t tmp_K = Add()(tmp_I, tmp_J);
  bfloat16_t tmp_L = cast_to<bfloat16_t>(tmp_K);
  bfloat16_t tmp_M = cast_to<bfloat16_t>(tmp_E);
  bfloat16_t tmp_N = LogAddExp()(tmp_L, tmp_M);
  float tmp_O = cast_to<float>(tmp_N);
  float tmp_P = cast_to<float>(tmp_H);
  float tmp_Q = cast_to<float>(tmp_O);
  float tmp_R = Multiply()(tmp_P, tmp_Q);
  float tmp_D = Exp()(tmp_R);
  D[index] = tmp_D;
  index++;
  }
}
[[host_name("kernel_qwen_mtplx_compute_g_bf16_contiguous_large")]]
[[kernel]] void Ef4IAsTypeAFf4IExpEGf4INegativeFHV2IBroadcastBCIV2IBroadcastCBJV2IAddHIKV2IBroadcastJDLV2IBroadcastDJMV2ILogAddExpKLNf4IAsTypeMOf4IBroadcastGNPf4IBroadcastNGQf4IMultiplyOPRf4OExpQ_VVVC_V2V2V2_10408321403207385874_contiguous_large(
    device const bfloat16_t* A [[buffer(0)]],
    device const bfloat16_t* B [[buffer(1)]],
    device const bfloat16_t* C [[buffer(2)]],
    device float* D [[buffer(3)]],
    constant const int64_t& size [[buffer(4)]],
    uint3 pos [[thread_position_in_grid]],
    uint3 grid [[threads_per_grid]]) {
  constexpr int N_ = 2;
  int64_t index = N_ * (pos.x + grid.x * int64_t(pos.y));
  for (int i = 0; i < N_ && index < size; ++i) {
  bfloat16_t tmp_A = A[index];
  bfloat16_t tmp_B = B[index];
  bfloat16_t tmp_C = C[index];
  auto tmp_E = static_cast<bfloat16_t>(0);
  float tmp_F = cast_to<float>(tmp_A);
  float tmp_G = Exp()(tmp_F);
  float tmp_H = Negative()(tmp_G);
  bfloat16_t tmp_I = cast_to<bfloat16_t>(tmp_B);
  bfloat16_t tmp_J = cast_to<bfloat16_t>(tmp_C);
  bfloat16_t tmp_K = Add()(tmp_I, tmp_J);
  bfloat16_t tmp_L = cast_to<bfloat16_t>(tmp_K);
  bfloat16_t tmp_M = cast_to<bfloat16_t>(tmp_E);
  bfloat16_t tmp_N = LogAddExp()(tmp_L, tmp_M);
  float tmp_O = cast_to<float>(tmp_N);
  float tmp_P = cast_to<float>(tmp_H);
  float tmp_Q = cast_to<float>(tmp_O);
  float tmp_R = Multiply()(tmp_P, tmp_Q);
  float tmp_D = Exp()(tmp_R);
  D[index] = tmp_D;
  index++;
  }
}
[[host_name("kernel_qwen_mtplx_compute_g_bf16_strided_1")]]
[[kernel]] void Ef4IAsTypeAFf4IExpEGf4INegativeFHV2IBroadcastBCIV2IBroadcastCBJV2IAddHIKV2IBroadcastJDLV2IBroadcastDJMV2ILogAddExpKLNf4IAsTypeMOf4IBroadcastGNPf4IBroadcastNGQf4IMultiplyOPRf4OExpQ_VVVC_V2V2V2_10408321403207385874_strided_1(
    device const bfloat16_t* A [[buffer(0)]],
    device const bfloat16_t* B [[buffer(1)]],
    device const bfloat16_t* C [[buffer(2)]],
    constant const int64_t* in_strides [[buffer(3)]],
    device float* D [[buffer(4)]],
    constant const int* output_shape [[buffer(5)]],
    uint3 pos [[thread_position_in_grid]],
    uint3 grid [[threads_per_grid]]) {
  constexpr int N_ = 1;
  uint index = pos.x + grid.x * (pos.y + uint(grid.y) * pos.z);
  auto tmp_E = static_cast<bfloat16_t>(0);
  uint index_A = elem_to_loc_1<uint>(pos.x, in_strides[0]);
  uint index_B = elem_to_loc_1<uint>(pos.x, in_strides[1]);
  uint index_C = elem_to_loc_1<uint>(pos.x, in_strides[2]);
  bfloat16_t tmp_A = A[index_A];
  bfloat16_t tmp_B = B[index_B];
  bfloat16_t tmp_C = C[index_C];
  float tmp_F = cast_to<float>(tmp_A);
  float tmp_G = Exp()(tmp_F);
  float tmp_H = Negative()(tmp_G);
  bfloat16_t tmp_I = cast_to<bfloat16_t>(tmp_B);
  bfloat16_t tmp_J = cast_to<bfloat16_t>(tmp_C);
  bfloat16_t tmp_K = Add()(tmp_I, tmp_J);
  bfloat16_t tmp_L = cast_to<bfloat16_t>(tmp_K);
  bfloat16_t tmp_M = cast_to<bfloat16_t>(tmp_E);
  bfloat16_t tmp_N = LogAddExp()(tmp_L, tmp_M);
  float tmp_O = cast_to<float>(tmp_N);
  float tmp_P = cast_to<float>(tmp_H);
  float tmp_Q = cast_to<float>(tmp_O);
  float tmp_R = Multiply()(tmp_P, tmp_Q);
  float tmp_D = Exp()(tmp_R);
  D[index] = tmp_D;
}
[[host_name("kernel_qwen_mtplx_compute_g_bf16_strided_1_large")]]
[[kernel]] void Ef4IAsTypeAFf4IExpEGf4INegativeFHV2IBroadcastBCIV2IBroadcastCBJV2IAddHIKV2IBroadcastJDLV2IBroadcastDJMV2ILogAddExpKLNf4IAsTypeMOf4IBroadcastGNPf4IBroadcastNGQf4IMultiplyOPRf4OExpQ_VVVC_V2V2V2_10408321403207385874_strided_1_large(
    device const bfloat16_t* A [[buffer(0)]],
    device const bfloat16_t* B [[buffer(1)]],
    device const bfloat16_t* C [[buffer(2)]],
    constant const int64_t* in_strides [[buffer(3)]],
    device float* D [[buffer(4)]],
    constant const int* output_shape [[buffer(5)]],
    uint3 pos [[thread_position_in_grid]],
    uint3 grid [[threads_per_grid]]) {
  constexpr int N_ = 1;
  int64_t index = pos.x + grid.x * (pos.y + int64_t(grid.y) * pos.z);
  auto tmp_E = static_cast<bfloat16_t>(0);
  int64_t index_A = elem_to_loc_1<int64_t>(pos.x, in_strides[0]);
  int64_t index_B = elem_to_loc_1<int64_t>(pos.x, in_strides[1]);
  int64_t index_C = elem_to_loc_1<int64_t>(pos.x, in_strides[2]);
  bfloat16_t tmp_A = A[index_A];
  bfloat16_t tmp_B = B[index_B];
  bfloat16_t tmp_C = C[index_C];
  float tmp_F = cast_to<float>(tmp_A);
  float tmp_G = Exp()(tmp_F);
  float tmp_H = Negative()(tmp_G);
  bfloat16_t tmp_I = cast_to<bfloat16_t>(tmp_B);
  bfloat16_t tmp_J = cast_to<bfloat16_t>(tmp_C);
  bfloat16_t tmp_K = Add()(tmp_I, tmp_J);
  bfloat16_t tmp_L = cast_to<bfloat16_t>(tmp_K);
  bfloat16_t tmp_M = cast_to<bfloat16_t>(tmp_E);
  bfloat16_t tmp_N = LogAddExp()(tmp_L, tmp_M);
  float tmp_O = cast_to<float>(tmp_N);
  float tmp_P = cast_to<float>(tmp_H);
  float tmp_Q = cast_to<float>(tmp_O);
  float tmp_R = Multiply()(tmp_P, tmp_Q);
  float tmp_D = Exp()(tmp_R);
  D[index] = tmp_D;
}
[[host_name("kernel_qwen_mtplx_compute_g_bf16_strided_2")]]
[[kernel]] void Ef4IAsTypeAFf4IExpEGf4INegativeFHV2IBroadcastBCIV2IBroadcastCBJV2IAddHIKV2IBroadcastJDLV2IBroadcastDJMV2ILogAddExpKLNf4IAsTypeMOf4IBroadcastGNPf4IBroadcastNGQf4IMultiplyOPRf4OExpQ_VVVC_V2V2V2_10408321403207385874_strided_2(
    device const bfloat16_t* A [[buffer(0)]],
    device const bfloat16_t* B [[buffer(1)]],
    device const bfloat16_t* C [[buffer(2)]],
    constant const int64_t* in_strides [[buffer(3)]],
    device float* D [[buffer(4)]],
    constant const int* output_shape [[buffer(5)]],
    uint3 pos [[thread_position_in_grid]],
    uint3 grid [[threads_per_grid]]) {
  constexpr int N_ = 1;
  uint index = pos.x + grid.x * (pos.y + uint(grid.y) * pos.z);
  auto tmp_E = static_cast<bfloat16_t>(0);
  uint index_A = elem_to_loc_2<uint>({pos.x, pos.y}, in_strides + 0);
  uint index_B = elem_to_loc_2<uint>({pos.x, pos.y}, in_strides + 2);
  uint index_C = elem_to_loc_2<uint>({pos.x, pos.y}, in_strides + 4);
  bfloat16_t tmp_A = A[index_A];
  bfloat16_t tmp_B = B[index_B];
  bfloat16_t tmp_C = C[index_C];
  float tmp_F = cast_to<float>(tmp_A);
  float tmp_G = Exp()(tmp_F);
  float tmp_H = Negative()(tmp_G);
  bfloat16_t tmp_I = cast_to<bfloat16_t>(tmp_B);
  bfloat16_t tmp_J = cast_to<bfloat16_t>(tmp_C);
  bfloat16_t tmp_K = Add()(tmp_I, tmp_J);
  bfloat16_t tmp_L = cast_to<bfloat16_t>(tmp_K);
  bfloat16_t tmp_M = cast_to<bfloat16_t>(tmp_E);
  bfloat16_t tmp_N = LogAddExp()(tmp_L, tmp_M);
  float tmp_O = cast_to<float>(tmp_N);
  float tmp_P = cast_to<float>(tmp_H);
  float tmp_Q = cast_to<float>(tmp_O);
  float tmp_R = Multiply()(tmp_P, tmp_Q);
  float tmp_D = Exp()(tmp_R);
  D[index] = tmp_D;
}
[[host_name("kernel_qwen_mtplx_compute_g_bf16_strided_2_large")]]
[[kernel]] void Ef4IAsTypeAFf4IExpEGf4INegativeFHV2IBroadcastBCIV2IBroadcastCBJV2IAddHIKV2IBroadcastJDLV2IBroadcastDJMV2ILogAddExpKLNf4IAsTypeMOf4IBroadcastGNPf4IBroadcastNGQf4IMultiplyOPRf4OExpQ_VVVC_V2V2V2_10408321403207385874_strided_2_large(
    device const bfloat16_t* A [[buffer(0)]],
    device const bfloat16_t* B [[buffer(1)]],
    device const bfloat16_t* C [[buffer(2)]],
    constant const int64_t* in_strides [[buffer(3)]],
    device float* D [[buffer(4)]],
    constant const int* output_shape [[buffer(5)]],
    uint3 pos [[thread_position_in_grid]],
    uint3 grid [[threads_per_grid]]) {
  constexpr int N_ = 1;
  int64_t index = pos.x + grid.x * (pos.y + int64_t(grid.y) * pos.z);
  auto tmp_E = static_cast<bfloat16_t>(0);
  int64_t index_A = elem_to_loc_2<int64_t>({pos.x, pos.y}, in_strides + 0);
  int64_t index_B = elem_to_loc_2<int64_t>({pos.x, pos.y}, in_strides + 2);
  int64_t index_C = elem_to_loc_2<int64_t>({pos.x, pos.y}, in_strides + 4);
  bfloat16_t tmp_A = A[index_A];
  bfloat16_t tmp_B = B[index_B];
  bfloat16_t tmp_C = C[index_C];
  float tmp_F = cast_to<float>(tmp_A);
  float tmp_G = Exp()(tmp_F);
  float tmp_H = Negative()(tmp_G);
  bfloat16_t tmp_I = cast_to<bfloat16_t>(tmp_B);
  bfloat16_t tmp_J = cast_to<bfloat16_t>(tmp_C);
  bfloat16_t tmp_K = Add()(tmp_I, tmp_J);
  bfloat16_t tmp_L = cast_to<bfloat16_t>(tmp_K);
  bfloat16_t tmp_M = cast_to<bfloat16_t>(tmp_E);
  bfloat16_t tmp_N = LogAddExp()(tmp_L, tmp_M);
  float tmp_O = cast_to<float>(tmp_N);
  float tmp_P = cast_to<float>(tmp_H);
  float tmp_Q = cast_to<float>(tmp_O);
  float tmp_R = Multiply()(tmp_P, tmp_Q);
  float tmp_D = Exp()(tmp_R);
  D[index] = tmp_D;
}
[[host_name("kernel_qwen_mtplx_compute_g_bf16_strided_3")]]
[[kernel]] void Ef4IAsTypeAFf4IExpEGf4INegativeFHV2IBroadcastBCIV2IBroadcastCBJV2IAddHIKV2IBroadcastJDLV2IBroadcastDJMV2ILogAddExpKLNf4IAsTypeMOf4IBroadcastGNPf4IBroadcastNGQf4IMultiplyOPRf4OExpQ_VVVC_V2V2V2_10408321403207385874_strided_3(
    device const bfloat16_t* A [[buffer(0)]],
    device const bfloat16_t* B [[buffer(1)]],
    device const bfloat16_t* C [[buffer(2)]],
    constant const int64_t* in_strides [[buffer(3)]],
    device float* D [[buffer(4)]],
    constant const int* output_shape [[buffer(5)]],
    uint3 pos [[thread_position_in_grid]],
    uint3 grid [[threads_per_grid]]) {
  constexpr int N_ = 1;
  uint index = pos.x + grid.x * (pos.y + uint(grid.y) * pos.z);
  auto tmp_E = static_cast<bfloat16_t>(0);
  uint index_A = elem_to_loc_3<uint>(pos, in_strides + 0);
  uint index_B = elem_to_loc_3<uint>(pos, in_strides + 3);
  uint index_C = elem_to_loc_3<uint>(pos, in_strides + 6);
  bfloat16_t tmp_A = A[index_A];
  bfloat16_t tmp_B = B[index_B];
  bfloat16_t tmp_C = C[index_C];
  float tmp_F = cast_to<float>(tmp_A);
  float tmp_G = Exp()(tmp_F);
  float tmp_H = Negative()(tmp_G);
  bfloat16_t tmp_I = cast_to<bfloat16_t>(tmp_B);
  bfloat16_t tmp_J = cast_to<bfloat16_t>(tmp_C);
  bfloat16_t tmp_K = Add()(tmp_I, tmp_J);
  bfloat16_t tmp_L = cast_to<bfloat16_t>(tmp_K);
  bfloat16_t tmp_M = cast_to<bfloat16_t>(tmp_E);
  bfloat16_t tmp_N = LogAddExp()(tmp_L, tmp_M);
  float tmp_O = cast_to<float>(tmp_N);
  float tmp_P = cast_to<float>(tmp_H);
  float tmp_Q = cast_to<float>(tmp_O);
  float tmp_R = Multiply()(tmp_P, tmp_Q);
  float tmp_D = Exp()(tmp_R);
  D[index] = tmp_D;
}
[[host_name("kernel_qwen_mtplx_compute_g_bf16_strided_3_large")]]
[[kernel]] void Ef4IAsTypeAFf4IExpEGf4INegativeFHV2IBroadcastBCIV2IBroadcastCBJV2IAddHIKV2IBroadcastJDLV2IBroadcastDJMV2ILogAddExpKLNf4IAsTypeMOf4IBroadcastGNPf4IBroadcastNGQf4IMultiplyOPRf4OExpQ_VVVC_V2V2V2_10408321403207385874_strided_3_large(
    device const bfloat16_t* A [[buffer(0)]],
    device const bfloat16_t* B [[buffer(1)]],
    device const bfloat16_t* C [[buffer(2)]],
    constant const int64_t* in_strides [[buffer(3)]],
    device float* D [[buffer(4)]],
    constant const int* output_shape [[buffer(5)]],
    uint3 pos [[thread_position_in_grid]],
    uint3 grid [[threads_per_grid]]) {
  constexpr int N_ = 1;
  int64_t index = pos.x + grid.x * (pos.y + int64_t(grid.y) * pos.z);
  auto tmp_E = static_cast<bfloat16_t>(0);
  int64_t index_A = elem_to_loc_3<int64_t>(pos, in_strides + 0);
  int64_t index_B = elem_to_loc_3<int64_t>(pos, in_strides + 3);
  int64_t index_C = elem_to_loc_3<int64_t>(pos, in_strides + 6);
  bfloat16_t tmp_A = A[index_A];
  bfloat16_t tmp_B = B[index_B];
  bfloat16_t tmp_C = C[index_C];
  float tmp_F = cast_to<float>(tmp_A);
  float tmp_G = Exp()(tmp_F);
  float tmp_H = Negative()(tmp_G);
  bfloat16_t tmp_I = cast_to<bfloat16_t>(tmp_B);
  bfloat16_t tmp_J = cast_to<bfloat16_t>(tmp_C);
  bfloat16_t tmp_K = Add()(tmp_I, tmp_J);
  bfloat16_t tmp_L = cast_to<bfloat16_t>(tmp_K);
  bfloat16_t tmp_M = cast_to<bfloat16_t>(tmp_E);
  bfloat16_t tmp_N = LogAddExp()(tmp_L, tmp_M);
  float tmp_O = cast_to<float>(tmp_N);
  float tmp_P = cast_to<float>(tmp_H);
  float tmp_Q = cast_to<float>(tmp_O);
  float tmp_R = Multiply()(tmp_P, tmp_Q);
  float tmp_D = Exp()(tmp_R);
  D[index] = tmp_D;
}
[[host_name("kernel_qwen_mtplx_compute_g_bf16_strided_4")]]
[[kernel]] void Ef4IAsTypeAFf4IExpEGf4INegativeFHV2IBroadcastBCIV2IBroadcastCBJV2IAddHIKV2IBroadcastJDLV2IBroadcastDJMV2ILogAddExpKLNf4IAsTypeMOf4IBroadcastGNPf4IBroadcastNGQf4IMultiplyOPRf4OExpQ_VVVC_V2V2V2_10408321403207385874_strided_4(
    device const bfloat16_t* A [[buffer(0)]],
    device const bfloat16_t* B [[buffer(1)]],
    device const bfloat16_t* C [[buffer(2)]],
    constant const int64_t* in_strides [[buffer(3)]],
    device float* D [[buffer(4)]],
    constant const int* output_shape [[buffer(5)]],
    uint3 pos [[thread_position_in_grid]],
    uint3 grid [[threads_per_grid]]) {
  constexpr int N_ = 2;
  int xshape = output_shape[3];
  uint index = N_ * pos.x + xshape * (pos.y + uint(grid.y) * pos.z);
  auto tmp_E = static_cast<bfloat16_t>(0);
  uint index_A = N_ * pos.x * uint(in_strides[3]) + pos.y * uint(in_strides[2]);
  uint index_B = N_ * pos.x * uint(in_strides[7]) + pos.y * uint(in_strides[6]);
  uint index_C = N_ * pos.x * uint(in_strides[11]) + pos.y * uint(in_strides[10]);
  uint zpos = pos.z;
  for (int d = 1; d >= 0; --d) {
    uint l = zpos % output_shape[d];
    index_A += l * uint(in_strides[0 + d]);
    index_B += l * uint(in_strides[4 + d]);
    index_C += l * uint(in_strides[8 + d]);
    zpos /= output_shape[d];
  }
  for (int i = 0; i < N_ && (int(N_ * pos.x) + i) < xshape; ++i) {
  bfloat16_t tmp_A = A[index_A];
  bfloat16_t tmp_B = B[index_B];
  bfloat16_t tmp_C = C[index_C];
  float tmp_F = cast_to<float>(tmp_A);
  float tmp_G = Exp()(tmp_F);
  float tmp_H = Negative()(tmp_G);
  bfloat16_t tmp_I = cast_to<bfloat16_t>(tmp_B);
  bfloat16_t tmp_J = cast_to<bfloat16_t>(tmp_C);
  bfloat16_t tmp_K = Add()(tmp_I, tmp_J);
  bfloat16_t tmp_L = cast_to<bfloat16_t>(tmp_K);
  bfloat16_t tmp_M = cast_to<bfloat16_t>(tmp_E);
  bfloat16_t tmp_N = LogAddExp()(tmp_L, tmp_M);
  float tmp_O = cast_to<float>(tmp_N);
  float tmp_P = cast_to<float>(tmp_H);
  float tmp_Q = cast_to<float>(tmp_O);
  float tmp_R = Multiply()(tmp_P, tmp_Q);
  float tmp_D = Exp()(tmp_R);
  D[index] = tmp_D;
  index_A += in_strides[3];
  index_B += in_strides[7];
  index_C += in_strides[11];
  index++;
  }
}
[[host_name("kernel_qwen_mtplx_compute_g_bf16_strided_4_large")]]
[[kernel]] void Ef4IAsTypeAFf4IExpEGf4INegativeFHV2IBroadcastBCIV2IBroadcastCBJV2IAddHIKV2IBroadcastJDLV2IBroadcastDJMV2ILogAddExpKLNf4IAsTypeMOf4IBroadcastGNPf4IBroadcastNGQf4IMultiplyOPRf4OExpQ_VVVC_V2V2V2_10408321403207385874_strided_4_large(
    device const bfloat16_t* A [[buffer(0)]],
    device const bfloat16_t* B [[buffer(1)]],
    device const bfloat16_t* C [[buffer(2)]],
    constant const int64_t* in_strides [[buffer(3)]],
    device float* D [[buffer(4)]],
    constant const int* output_shape [[buffer(5)]],
    uint3 pos [[thread_position_in_grid]],
    uint3 grid [[threads_per_grid]]) {
  constexpr int N_ = 4;
  int xshape = output_shape[3];
  int64_t index = N_ * pos.x + xshape * (pos.y + int64_t(grid.y) * pos.z);
  auto tmp_E = static_cast<bfloat16_t>(0);
  int64_t index_A = N_ * pos.x * int64_t(in_strides[3]) + pos.y * int64_t(in_strides[2]);
  int64_t index_B = N_ * pos.x * int64_t(in_strides[7]) + pos.y * int64_t(in_strides[6]);
  int64_t index_C = N_ * pos.x * int64_t(in_strides[11]) + pos.y * int64_t(in_strides[10]);
  uint zpos = pos.z;
  for (int d = 1; d >= 0; --d) {
    uint l = zpos % output_shape[d];
    index_A += l * int64_t(in_strides[0 + d]);
    index_B += l * int64_t(in_strides[4 + d]);
    index_C += l * int64_t(in_strides[8 + d]);
    zpos /= output_shape[d];
  }
  for (int i = 0; i < N_ && (int(N_ * pos.x) + i) < xshape; ++i) {
  bfloat16_t tmp_A = A[index_A];
  bfloat16_t tmp_B = B[index_B];
  bfloat16_t tmp_C = C[index_C];
  float tmp_F = cast_to<float>(tmp_A);
  float tmp_G = Exp()(tmp_F);
  float tmp_H = Negative()(tmp_G);
  bfloat16_t tmp_I = cast_to<bfloat16_t>(tmp_B);
  bfloat16_t tmp_J = cast_to<bfloat16_t>(tmp_C);
  bfloat16_t tmp_K = Add()(tmp_I, tmp_J);
  bfloat16_t tmp_L = cast_to<bfloat16_t>(tmp_K);
  bfloat16_t tmp_M = cast_to<bfloat16_t>(tmp_E);
  bfloat16_t tmp_N = LogAddExp()(tmp_L, tmp_M);
  float tmp_O = cast_to<float>(tmp_N);
  float tmp_P = cast_to<float>(tmp_H);
  float tmp_Q = cast_to<float>(tmp_O);
  float tmp_R = Multiply()(tmp_P, tmp_Q);
  float tmp_D = Exp()(tmp_R);
  D[index] = tmp_D;
  index_A += in_strides[3];
  index_B += in_strides[7];
  index_C += in_strides[11];
  index++;
  }
}
[[host_name("kernel_qwen_mtplx_compute_g_bf16_strided_5")]]
[[kernel]] void Ef4IAsTypeAFf4IExpEGf4INegativeFHV2IBroadcastBCIV2IBroadcastCBJV2IAddHIKV2IBroadcastJDLV2IBroadcastDJMV2ILogAddExpKLNf4IAsTypeMOf4IBroadcastGNPf4IBroadcastNGQf4IMultiplyOPRf4OExpQ_VVVC_V2V2V2_10408321403207385874_strided_5(
    device const bfloat16_t* A [[buffer(0)]],
    device const bfloat16_t* B [[buffer(1)]],
    device const bfloat16_t* C [[buffer(2)]],
    constant const int64_t* in_strides [[buffer(3)]],
    device float* D [[buffer(4)]],
    constant const int* output_shape [[buffer(5)]],
    uint3 pos [[thread_position_in_grid]],
    uint3 grid [[threads_per_grid]]) {
  constexpr int N_ = 2;
  int xshape = output_shape[4];
  uint index = N_ * pos.x + xshape * (pos.y + uint(grid.y) * pos.z);
  auto tmp_E = static_cast<bfloat16_t>(0);
  uint index_A = N_ * pos.x * uint(in_strides[4]) + pos.y * uint(in_strides[3]);
  uint index_B = N_ * pos.x * uint(in_strides[9]) + pos.y * uint(in_strides[8]);
  uint index_C = N_ * pos.x * uint(in_strides[14]) + pos.y * uint(in_strides[13]);
  uint zpos = pos.z;
  for (int d = 2; d >= 0; --d) {
    uint l = zpos % output_shape[d];
    index_A += l * uint(in_strides[0 + d]);
    index_B += l * uint(in_strides[5 + d]);
    index_C += l * uint(in_strides[10 + d]);
    zpos /= output_shape[d];
  }
  for (int i = 0; i < N_ && (int(N_ * pos.x) + i) < xshape; ++i) {
  bfloat16_t tmp_A = A[index_A];
  bfloat16_t tmp_B = B[index_B];
  bfloat16_t tmp_C = C[index_C];
  float tmp_F = cast_to<float>(tmp_A);
  float tmp_G = Exp()(tmp_F);
  float tmp_H = Negative()(tmp_G);
  bfloat16_t tmp_I = cast_to<bfloat16_t>(tmp_B);
  bfloat16_t tmp_J = cast_to<bfloat16_t>(tmp_C);
  bfloat16_t tmp_K = Add()(tmp_I, tmp_J);
  bfloat16_t tmp_L = cast_to<bfloat16_t>(tmp_K);
  bfloat16_t tmp_M = cast_to<bfloat16_t>(tmp_E);
  bfloat16_t tmp_N = LogAddExp()(tmp_L, tmp_M);
  float tmp_O = cast_to<float>(tmp_N);
  float tmp_P = cast_to<float>(tmp_H);
  float tmp_Q = cast_to<float>(tmp_O);
  float tmp_R = Multiply()(tmp_P, tmp_Q);
  float tmp_D = Exp()(tmp_R);
  D[index] = tmp_D;
  index_A += in_strides[4];
  index_B += in_strides[9];
  index_C += in_strides[14];
  index++;
  }
}
[[host_name("kernel_qwen_mtplx_compute_g_bf16_strided_5_large")]]
[[kernel]] void Ef4IAsTypeAFf4IExpEGf4INegativeFHV2IBroadcastBCIV2IBroadcastCBJV2IAddHIKV2IBroadcastJDLV2IBroadcastDJMV2ILogAddExpKLNf4IAsTypeMOf4IBroadcastGNPf4IBroadcastNGQf4IMultiplyOPRf4OExpQ_VVVC_V2V2V2_10408321403207385874_strided_5_large(
    device const bfloat16_t* A [[buffer(0)]],
    device const bfloat16_t* B [[buffer(1)]],
    device const bfloat16_t* C [[buffer(2)]],
    constant const int64_t* in_strides [[buffer(3)]],
    device float* D [[buffer(4)]],
    constant const int* output_shape [[buffer(5)]],
    uint3 pos [[thread_position_in_grid]],
    uint3 grid [[threads_per_grid]]) {
  constexpr int N_ = 4;
  int xshape = output_shape[4];
  int64_t index = N_ * pos.x + xshape * (pos.y + int64_t(grid.y) * pos.z);
  auto tmp_E = static_cast<bfloat16_t>(0);
  int64_t index_A = N_ * pos.x * int64_t(in_strides[4]) + pos.y * int64_t(in_strides[3]);
  int64_t index_B = N_ * pos.x * int64_t(in_strides[9]) + pos.y * int64_t(in_strides[8]);
  int64_t index_C = N_ * pos.x * int64_t(in_strides[14]) + pos.y * int64_t(in_strides[13]);
  uint zpos = pos.z;
  for (int d = 2; d >= 0; --d) {
    uint l = zpos % output_shape[d];
    index_A += l * int64_t(in_strides[0 + d]);
    index_B += l * int64_t(in_strides[5 + d]);
    index_C += l * int64_t(in_strides[10 + d]);
    zpos /= output_shape[d];
  }
  for (int i = 0; i < N_ && (int(N_ * pos.x) + i) < xshape; ++i) {
  bfloat16_t tmp_A = A[index_A];
  bfloat16_t tmp_B = B[index_B];
  bfloat16_t tmp_C = C[index_C];
  float tmp_F = cast_to<float>(tmp_A);
  float tmp_G = Exp()(tmp_F);
  float tmp_H = Negative()(tmp_G);
  bfloat16_t tmp_I = cast_to<bfloat16_t>(tmp_B);
  bfloat16_t tmp_J = cast_to<bfloat16_t>(tmp_C);
  bfloat16_t tmp_K = Add()(tmp_I, tmp_J);
  bfloat16_t tmp_L = cast_to<bfloat16_t>(tmp_K);
  bfloat16_t tmp_M = cast_to<bfloat16_t>(tmp_E);
  bfloat16_t tmp_N = LogAddExp()(tmp_L, tmp_M);
  float tmp_O = cast_to<float>(tmp_N);
  float tmp_P = cast_to<float>(tmp_H);
  float tmp_Q = cast_to<float>(tmp_O);
  float tmp_R = Multiply()(tmp_P, tmp_Q);
  float tmp_D = Exp()(tmp_R);
  D[index] = tmp_D;
  index_A += in_strides[4];
  index_B += in_strides[9];
  index_C += in_strides[14];
  index++;
  }
}
[[host_name("kernel_qwen_mtplx_compute_g_bf16_strided_6")]]
[[kernel]] void Ef4IAsTypeAFf4IExpEGf4INegativeFHV2IBroadcastBCIV2IBroadcastCBJV2IAddHIKV2IBroadcastJDLV2IBroadcastDJMV2ILogAddExpKLNf4IAsTypeMOf4IBroadcastGNPf4IBroadcastNGQf4IMultiplyOPRf4OExpQ_VVVC_V2V2V2_10408321403207385874_strided_6(
    device const bfloat16_t* A [[buffer(0)]],
    device const bfloat16_t* B [[buffer(1)]],
    device const bfloat16_t* C [[buffer(2)]],
    constant const int64_t* in_strides [[buffer(3)]],
    device float* D [[buffer(4)]],
    constant const int* output_shape [[buffer(5)]],
    uint3 pos [[thread_position_in_grid]],
    uint3 grid [[threads_per_grid]]) {
  constexpr int N_ = 2;
  int xshape = output_shape[5];
  uint index = N_ * pos.x + xshape * (pos.y + uint(grid.y) * pos.z);
  auto tmp_E = static_cast<bfloat16_t>(0);
  uint index_A = N_ * pos.x * uint(in_strides[5]) + pos.y * uint(in_strides[4]);
  uint index_B = N_ * pos.x * uint(in_strides[11]) + pos.y * uint(in_strides[10]);
  uint index_C = N_ * pos.x * uint(in_strides[17]) + pos.y * uint(in_strides[16]);
  uint zpos = pos.z;
  for (int d = 3; d >= 0; --d) {
    uint l = zpos % output_shape[d];
    index_A += l * uint(in_strides[0 + d]);
    index_B += l * uint(in_strides[6 + d]);
    index_C += l * uint(in_strides[12 + d]);
    zpos /= output_shape[d];
  }
  for (int i = 0; i < N_ && (int(N_ * pos.x) + i) < xshape; ++i) {
  bfloat16_t tmp_A = A[index_A];
  bfloat16_t tmp_B = B[index_B];
  bfloat16_t tmp_C = C[index_C];
  float tmp_F = cast_to<float>(tmp_A);
  float tmp_G = Exp()(tmp_F);
  float tmp_H = Negative()(tmp_G);
  bfloat16_t tmp_I = cast_to<bfloat16_t>(tmp_B);
  bfloat16_t tmp_J = cast_to<bfloat16_t>(tmp_C);
  bfloat16_t tmp_K = Add()(tmp_I, tmp_J);
  bfloat16_t tmp_L = cast_to<bfloat16_t>(tmp_K);
  bfloat16_t tmp_M = cast_to<bfloat16_t>(tmp_E);
  bfloat16_t tmp_N = LogAddExp()(tmp_L, tmp_M);
  float tmp_O = cast_to<float>(tmp_N);
  float tmp_P = cast_to<float>(tmp_H);
  float tmp_Q = cast_to<float>(tmp_O);
  float tmp_R = Multiply()(tmp_P, tmp_Q);
  float tmp_D = Exp()(tmp_R);
  D[index] = tmp_D;
  index_A += in_strides[5];
  index_B += in_strides[11];
  index_C += in_strides[17];
  index++;
  }
}
[[host_name("kernel_qwen_mtplx_compute_g_bf16_strided_6_large")]]
[[kernel]] void Ef4IAsTypeAFf4IExpEGf4INegativeFHV2IBroadcastBCIV2IBroadcastCBJV2IAddHIKV2IBroadcastJDLV2IBroadcastDJMV2ILogAddExpKLNf4IAsTypeMOf4IBroadcastGNPf4IBroadcastNGQf4IMultiplyOPRf4OExpQ_VVVC_V2V2V2_10408321403207385874_strided_6_large(
    device const bfloat16_t* A [[buffer(0)]],
    device const bfloat16_t* B [[buffer(1)]],
    device const bfloat16_t* C [[buffer(2)]],
    constant const int64_t* in_strides [[buffer(3)]],
    device float* D [[buffer(4)]],
    constant const int* output_shape [[buffer(5)]],
    uint3 pos [[thread_position_in_grid]],
    uint3 grid [[threads_per_grid]]) {
  constexpr int N_ = 4;
  int xshape = output_shape[5];
  int64_t index = N_ * pos.x + xshape * (pos.y + int64_t(grid.y) * pos.z);
  auto tmp_E = static_cast<bfloat16_t>(0);
  int64_t index_A = N_ * pos.x * int64_t(in_strides[5]) + pos.y * int64_t(in_strides[4]);
  int64_t index_B = N_ * pos.x * int64_t(in_strides[11]) + pos.y * int64_t(in_strides[10]);
  int64_t index_C = N_ * pos.x * int64_t(in_strides[17]) + pos.y * int64_t(in_strides[16]);
  uint zpos = pos.z;
  for (int d = 3; d >= 0; --d) {
    uint l = zpos % output_shape[d];
    index_A += l * int64_t(in_strides[0 + d]);
    index_B += l * int64_t(in_strides[6 + d]);
    index_C += l * int64_t(in_strides[12 + d]);
    zpos /= output_shape[d];
  }
  for (int i = 0; i < N_ && (int(N_ * pos.x) + i) < xshape; ++i) {
  bfloat16_t tmp_A = A[index_A];
  bfloat16_t tmp_B = B[index_B];
  bfloat16_t tmp_C = C[index_C];
  float tmp_F = cast_to<float>(tmp_A);
  float tmp_G = Exp()(tmp_F);
  float tmp_H = Negative()(tmp_G);
  bfloat16_t tmp_I = cast_to<bfloat16_t>(tmp_B);
  bfloat16_t tmp_J = cast_to<bfloat16_t>(tmp_C);
  bfloat16_t tmp_K = Add()(tmp_I, tmp_J);
  bfloat16_t tmp_L = cast_to<bfloat16_t>(tmp_K);
  bfloat16_t tmp_M = cast_to<bfloat16_t>(tmp_E);
  bfloat16_t tmp_N = LogAddExp()(tmp_L, tmp_M);
  float tmp_O = cast_to<float>(tmp_N);
  float tmp_P = cast_to<float>(tmp_H);
  float tmp_Q = cast_to<float>(tmp_O);
  float tmp_R = Multiply()(tmp_P, tmp_Q);
  float tmp_D = Exp()(tmp_R);
  D[index] = tmp_D;
  index_A += in_strides[5];
  index_B += in_strides[11];
  index_C += in_strides[17];
  index++;
  }
}
[[host_name("kernel_qwen_mtplx_compute_g_bf16_strided_7")]]
[[kernel]] void Ef4IAsTypeAFf4IExpEGf4INegativeFHV2IBroadcastBCIV2IBroadcastCBJV2IAddHIKV2IBroadcastJDLV2IBroadcastDJMV2ILogAddExpKLNf4IAsTypeMOf4IBroadcastGNPf4IBroadcastNGQf4IMultiplyOPRf4OExpQ_VVVC_V2V2V2_10408321403207385874_strided_7(
    device const bfloat16_t* A [[buffer(0)]],
    device const bfloat16_t* B [[buffer(1)]],
    device const bfloat16_t* C [[buffer(2)]],
    constant const int64_t* in_strides [[buffer(3)]],
    device float* D [[buffer(4)]],
    constant const int* output_shape [[buffer(5)]],
    uint3 pos [[thread_position_in_grid]],
    uint3 grid [[threads_per_grid]]) {
  constexpr int N_ = 2;
  int xshape = output_shape[6];
  uint index = N_ * pos.x + xshape * (pos.y + uint(grid.y) * pos.z);
  auto tmp_E = static_cast<bfloat16_t>(0);
  uint index_A = N_ * pos.x * uint(in_strides[6]) + pos.y * uint(in_strides[5]);
  uint index_B = N_ * pos.x * uint(in_strides[13]) + pos.y * uint(in_strides[12]);
  uint index_C = N_ * pos.x * uint(in_strides[20]) + pos.y * uint(in_strides[19]);
  uint zpos = pos.z;
  for (int d = 4; d >= 0; --d) {
    uint l = zpos % output_shape[d];
    index_A += l * uint(in_strides[0 + d]);
    index_B += l * uint(in_strides[7 + d]);
    index_C += l * uint(in_strides[14 + d]);
    zpos /= output_shape[d];
  }
  for (int i = 0; i < N_ && (int(N_ * pos.x) + i) < xshape; ++i) {
  bfloat16_t tmp_A = A[index_A];
  bfloat16_t tmp_B = B[index_B];
  bfloat16_t tmp_C = C[index_C];
  float tmp_F = cast_to<float>(tmp_A);
  float tmp_G = Exp()(tmp_F);
  float tmp_H = Negative()(tmp_G);
  bfloat16_t tmp_I = cast_to<bfloat16_t>(tmp_B);
  bfloat16_t tmp_J = cast_to<bfloat16_t>(tmp_C);
  bfloat16_t tmp_K = Add()(tmp_I, tmp_J);
  bfloat16_t tmp_L = cast_to<bfloat16_t>(tmp_K);
  bfloat16_t tmp_M = cast_to<bfloat16_t>(tmp_E);
  bfloat16_t tmp_N = LogAddExp()(tmp_L, tmp_M);
  float tmp_O = cast_to<float>(tmp_N);
  float tmp_P = cast_to<float>(tmp_H);
  float tmp_Q = cast_to<float>(tmp_O);
  float tmp_R = Multiply()(tmp_P, tmp_Q);
  float tmp_D = Exp()(tmp_R);
  D[index] = tmp_D;
  index_A += in_strides[6];
  index_B += in_strides[13];
  index_C += in_strides[20];
  index++;
  }
}
[[host_name("kernel_qwen_mtplx_compute_g_bf16_strided_7_large")]]
[[kernel]] void Ef4IAsTypeAFf4IExpEGf4INegativeFHV2IBroadcastBCIV2IBroadcastCBJV2IAddHIKV2IBroadcastJDLV2IBroadcastDJMV2ILogAddExpKLNf4IAsTypeMOf4IBroadcastGNPf4IBroadcastNGQf4IMultiplyOPRf4OExpQ_VVVC_V2V2V2_10408321403207385874_strided_7_large(
    device const bfloat16_t* A [[buffer(0)]],
    device const bfloat16_t* B [[buffer(1)]],
    device const bfloat16_t* C [[buffer(2)]],
    constant const int64_t* in_strides [[buffer(3)]],
    device float* D [[buffer(4)]],
    constant const int* output_shape [[buffer(5)]],
    uint3 pos [[thread_position_in_grid]],
    uint3 grid [[threads_per_grid]]) {
  constexpr int N_ = 4;
  int xshape = output_shape[6];
  int64_t index = N_ * pos.x + xshape * (pos.y + int64_t(grid.y) * pos.z);
  auto tmp_E = static_cast<bfloat16_t>(0);
  int64_t index_A = N_ * pos.x * int64_t(in_strides[6]) + pos.y * int64_t(in_strides[5]);
  int64_t index_B = N_ * pos.x * int64_t(in_strides[13]) + pos.y * int64_t(in_strides[12]);
  int64_t index_C = N_ * pos.x * int64_t(in_strides[20]) + pos.y * int64_t(in_strides[19]);
  uint zpos = pos.z;
  for (int d = 4; d >= 0; --d) {
    uint l = zpos % output_shape[d];
    index_A += l * int64_t(in_strides[0 + d]);
    index_B += l * int64_t(in_strides[7 + d]);
    index_C += l * int64_t(in_strides[14 + d]);
    zpos /= output_shape[d];
  }
  for (int i = 0; i < N_ && (int(N_ * pos.x) + i) < xshape; ++i) {
  bfloat16_t tmp_A = A[index_A];
  bfloat16_t tmp_B = B[index_B];
  bfloat16_t tmp_C = C[index_C];
  float tmp_F = cast_to<float>(tmp_A);
  float tmp_G = Exp()(tmp_F);
  float tmp_H = Negative()(tmp_G);
  bfloat16_t tmp_I = cast_to<bfloat16_t>(tmp_B);
  bfloat16_t tmp_J = cast_to<bfloat16_t>(tmp_C);
  bfloat16_t tmp_K = Add()(tmp_I, tmp_J);
  bfloat16_t tmp_L = cast_to<bfloat16_t>(tmp_K);
  bfloat16_t tmp_M = cast_to<bfloat16_t>(tmp_E);
  bfloat16_t tmp_N = LogAddExp()(tmp_L, tmp_M);
  float tmp_O = cast_to<float>(tmp_N);
  float tmp_P = cast_to<float>(tmp_H);
  float tmp_Q = cast_to<float>(tmp_O);
  float tmp_R = Multiply()(tmp_P, tmp_Q);
  float tmp_D = Exp()(tmp_R);
  D[index] = tmp_D;
  index_A += in_strides[6];
  index_B += in_strides[13];
  index_C += in_strides[20];
  index++;
  }
}
[[host_name("kernel_qwen_mtplx_compute_g_bf16_strided_dynamic")]]
[[kernel]] void Ef4IAsTypeAFf4IExpEGf4INegativeFHV2IBroadcastBCIV2IBroadcastCBJV2IAddHIKV2IBroadcastJDLV2IBroadcastDJMV2ILogAddExpKLNf4IAsTypeMOf4IBroadcastGNPf4IBroadcastNGQf4IMultiplyOPRf4OExpQ_VVVC_V2V2V2_10408321403207385874_strided_dynamic(
    device const bfloat16_t* A [[buffer(0)]],
    device const bfloat16_t* B [[buffer(1)]],
    device const bfloat16_t* C [[buffer(2)]],
    constant const int64_t* in_strides [[buffer(3)]],
    device float* D [[buffer(4)]],
    constant const int* output_shape [[buffer(5)]],
    constant const int& ndim [[buffer(6)]],
    uint3 pos [[thread_position_in_grid]],
    uint3 grid [[threads_per_grid]]) {
  constexpr int N_ = 2;
  int xshape = output_shape[ndim - 1];
  uint index = N_ * pos.x + xshape * (pos.y + uint(grid.y) * pos.z);
  auto tmp_E = static_cast<bfloat16_t>(0);
  uint index_A = N_ * pos.x * uint(in_strides[ndim * 0 + ndim - 1]) + pos.y * uint(in_strides[ndim * 0 + ndim - 2]);
  uint index_B = N_ * pos.x * uint(in_strides[ndim * 1 + ndim - 1]) + pos.y * uint(in_strides[ndim * 1 + ndim - 2]);
  uint index_C = N_ * pos.x * uint(in_strides[ndim * 2 + ndim - 1]) + pos.y * uint(in_strides[ndim * 2 + ndim - 2]);
  uint zpos = pos.z;
  for (int d = ndim - 3; d >= 0; --d) {
    uint l = zpos % output_shape[d];
    index_A += l * uint(in_strides[0 * ndim + d]);
    index_B += l * uint(in_strides[1 * ndim + d]);
    index_C += l * uint(in_strides[2 * ndim + d]);
    zpos /= output_shape[d];
  }
  for (int i = 0; i < N_ && (int(N_ * pos.x) + i) < xshape; ++i) {
  bfloat16_t tmp_A = A[index_A];
  bfloat16_t tmp_B = B[index_B];
  bfloat16_t tmp_C = C[index_C];
  float tmp_F = cast_to<float>(tmp_A);
  float tmp_G = Exp()(tmp_F);
  float tmp_H = Negative()(tmp_G);
  bfloat16_t tmp_I = cast_to<bfloat16_t>(tmp_B);
  bfloat16_t tmp_J = cast_to<bfloat16_t>(tmp_C);
  bfloat16_t tmp_K = Add()(tmp_I, tmp_J);
  bfloat16_t tmp_L = cast_to<bfloat16_t>(tmp_K);
  bfloat16_t tmp_M = cast_to<bfloat16_t>(tmp_E);
  bfloat16_t tmp_N = LogAddExp()(tmp_L, tmp_M);
  float tmp_O = cast_to<float>(tmp_N);
  float tmp_P = cast_to<float>(tmp_H);
  float tmp_Q = cast_to<float>(tmp_O);
  float tmp_R = Multiply()(tmp_P, tmp_Q);
  float tmp_D = Exp()(tmp_R);
  D[index] = tmp_D;
  index_A += in_strides[0 * ndim + ndim - 1];
  index_B += in_strides[1 * ndim + ndim - 1];
  index_C += in_strides[2 * ndim + ndim - 1];
  index++;
  }
}
[[host_name("kernel_qwen_mtplx_compute_g_bf16_strided_dynamic_large")]]
[[kernel]] void Ef4IAsTypeAFf4IExpEGf4INegativeFHV2IBroadcastBCIV2IBroadcastCBJV2IAddHIKV2IBroadcastJDLV2IBroadcastDJMV2ILogAddExpKLNf4IAsTypeMOf4IBroadcastGNPf4IBroadcastNGQf4IMultiplyOPRf4OExpQ_VVVC_V2V2V2_10408321403207385874_strided_dynamic_large(
    device const bfloat16_t* A [[buffer(0)]],
    device const bfloat16_t* B [[buffer(1)]],
    device const bfloat16_t* C [[buffer(2)]],
    constant const int64_t* in_strides [[buffer(3)]],
    device float* D [[buffer(4)]],
    constant const int* output_shape [[buffer(5)]],
    constant const int& ndim [[buffer(6)]],
    uint3 pos [[thread_position_in_grid]],
    uint3 grid [[threads_per_grid]]) {
  constexpr int N_ = 4;
  int xshape = output_shape[ndim - 1];
  int64_t index = N_ * pos.x + xshape * (pos.y + int64_t(grid.y) * pos.z);
  auto tmp_E = static_cast<bfloat16_t>(0);
  int64_t index_A = N_ * pos.x * int64_t(in_strides[ndim * 0 + ndim - 1]) + pos.y * int64_t(in_strides[ndim * 0 + ndim - 2]);
  int64_t index_B = N_ * pos.x * int64_t(in_strides[ndim * 1 + ndim - 1]) + pos.y * int64_t(in_strides[ndim * 1 + ndim - 2]);
  int64_t index_C = N_ * pos.x * int64_t(in_strides[ndim * 2 + ndim - 1]) + pos.y * int64_t(in_strides[ndim * 2 + ndim - 2]);
  uint zpos = pos.z;
  for (int d = ndim - 3; d >= 0; --d) {
    uint l = zpos % output_shape[d];
    index_A += l * int64_t(in_strides[0 * ndim + d]);
    index_B += l * int64_t(in_strides[1 * ndim + d]);
    index_C += l * int64_t(in_strides[2 * ndim + d]);
    zpos /= output_shape[d];
  }
  for (int i = 0; i < N_ && (int(N_ * pos.x) + i) < xshape; ++i) {
  bfloat16_t tmp_A = A[index_A];
  bfloat16_t tmp_B = B[index_B];
  bfloat16_t tmp_C = C[index_C];
  float tmp_F = cast_to<float>(tmp_A);
  float tmp_G = Exp()(tmp_F);
  float tmp_H = Negative()(tmp_G);
  bfloat16_t tmp_I = cast_to<bfloat16_t>(tmp_B);
  bfloat16_t tmp_J = cast_to<bfloat16_t>(tmp_C);
  bfloat16_t tmp_K = Add()(tmp_I, tmp_J);
  bfloat16_t tmp_L = cast_to<bfloat16_t>(tmp_K);
  bfloat16_t tmp_M = cast_to<bfloat16_t>(tmp_E);
  bfloat16_t tmp_N = LogAddExp()(tmp_L, tmp_M);
  float tmp_O = cast_to<float>(tmp_N);
  float tmp_P = cast_to<float>(tmp_H);
  float tmp_Q = cast_to<float>(tmp_O);
  float tmp_R = Multiply()(tmp_P, tmp_Q);
  float tmp_D = Exp()(tmp_R);
  D[index] = tmp_D;
  index_A += in_strides[0 * ndim + ndim - 1];
  index_B += in_strides[1 * ndim + ndim - 1];
  index_C += in_strides[2 * ndim + ndim - 1];
  index++;
  }
}
// END RUNTIME UNIT

namespace mtplx_qsa_prepare_q {

        using namespace metal;
        constant constexpr uint HEADS = 4;
        constant constexpr uint HEAD_DIM = 128;
        constant constexpr uint ROTARY_DIM = 64;
        constant constexpr uint HALF_ROTARY = 32;
        constant constexpr float RMS_EPS = 1e-06f;
        constant constexpr float ROPE_ATTENTION_SCALE = 1.0f;
    
// Source: mtplx/kernels/qsa_indexer_prepare.py::_prepare_queries_kernel
// File SHA256: a77f6ca5ae805e729519c4629ae88b455a6dbf473a457a6e1c8219174eb59091
// Body SHA256: 67329208d2a7906094e742871793de678d2272f669fd928d929943979be8b741
template <typename T>
kernel void kernel_qwen_mtplx_qsa_prepare_q(
    device const T* raw_q [[buffer(0)]],
    constant const int64_t* raw_q_strides [[buffer(1)]],
    device const T* norm_weight [[buffer(2)]],
    constant const int64_t* norm_weight_strides [[buffer(3)]],
    device const float* inv_freq [[buffer(4)]],
    constant const int64_t* inv_freq_strides [[buffer(5)]],
    device const int* pos_start [[buffer(6)]],
    device T* prepared_q [[buffer(7)]],
    uint3 threadgroup_position_in_grid [[threadgroup_position_in_grid]],
    uint thread_index_in_simdgroup [[thread_index_in_simdgroup]]) {
        const uint row_head = threadgroup_position_in_grid.x;
        const uint lane = thread_index_in_simdgroup;
        const uint row = row_head / HEADS;
        const uint head = row_head - row * HEADS;
        const size_t src_base =
            (size_t)row * raw_q_strides[1] +
            (size_t)head * raw_q_strides[2];
        const size_t out_base =
            ((size_t)row * HEADS + head) * HEAD_DIM;

        // MLX rms_single_row uses 32 lanes and four contiguous values per
        // lane for every axis up through 128.  Reproduce that reduction and
        // its cast-before-weight ordering exactly.
        float square_sum = 0.0f;
        const uint lane_base = lane * 4u;
        for (uint i = 0; i < 4u; ++i) {
            const uint dim = lane_base + i;
            if (dim < HEAD_DIM) {
                const float value = float(raw_q[
                    src_base + (size_t)dim * raw_q_strides[3]]);
                square_sum += value * value;
            }
        }
        square_sum = simd_sum(square_sum);
        const float inverse_rms = metal::precise::rsqrt(
            square_sum / float(HEAD_DIM) + RMS_EPS);
        const float position = float(pos_start[0] + int(row));

        // Non-interleaved/rotate-half partial RoPE, matching
        // _rope_cos_sin + _apply_partial_rope.  mx.cos/mx.sin use the Metal
        // precise variants (unlike mx.fast.rope, which intentionally uses
        // fast trig), so this kernel does too.
        for (uint pair = lane; pair < HALF_ROTARY; pair += 32u) {
            const size_t first_at =
                src_base + (size_t)pair * raw_q_strides[3];
            const size_t second_at = src_base +
                (size_t)(pair + HALF_ROTARY) * raw_q_strides[3];
            const T first_norm = norm_weight[
                (size_t)pair * norm_weight_strides[0]] *
                static_cast<T>(float(raw_q[first_at]) * inverse_rms);
            const T second_norm = norm_weight[
                (size_t)(pair + HALF_ROTARY) * norm_weight_strides[0]] *
                static_cast<T>(float(raw_q[second_at]) * inverse_rms);
            const float theta = position * float(inv_freq[
                (size_t)pair * inv_freq_strides[0]]);
            const float cosine =
                metal::precise::cos(theta) * ROPE_ATTENTION_SCALE;
            const float sine =
                metal::precise::sin(theta) * ROPE_ATTENTION_SCALE;
            const float first = float(first_norm);
            const float second = float(second_norm);
            // Keep the products as distinct fp32 operations.  The stock
            // _apply_partial_rope graph rounds both multiplies before its
            // add/subtract; allowing Metal to contract this expression into
            // an FMA changes a handful of bf16 cutoff values at large
            // positions (and can therefore perturb an exact top-k set).
            const float first_cosine = first * cosine;
            const float second_sine = second * sine;
            const float second_cosine = second * cosine;
            const float first_sine = first * sine;
            prepared_q[out_base + pair] =
                static_cast<T>(first_cosine - second_sine);
            prepared_q[out_base + pair + HALF_ROTARY] =
                static_cast<T>(second_cosine + first_sine);
        }

        for (uint dim = ROTARY_DIM + lane; dim < HEAD_DIM; dim += 32u) {
            const T normalized = norm_weight[
                (size_t)dim * norm_weight_strides[0]] *
                static_cast<T>(float(raw_q[
                    src_base + (size_t)dim * raw_q_strides[3]]) * inverse_rms);
            prepared_q[out_base + dim] = normalized;
        }
    }
typedef decltype(kernel_qwen_mtplx_qsa_prepare_q<bfloat>) kernel_qwen_mtplx_qsa_prepare_q_bf16_type;
template [[host_name("kernel_qwen_mtplx_qsa_prepare_q_bf16")]]
kernel kernel_qwen_mtplx_qsa_prepare_q_bf16_type kernel_qwen_mtplx_qsa_prepare_q<bfloat>;
} // namespace mtplx_qsa_prepare_q

namespace mtplx_qsa_pool_k {

        using namespace metal;
        constant constexpr uint HEAD_DIM = 128;
        constant constexpr uint ROTARY_DIM = 64;
        constant constexpr uint HALF_ROTARY = 32;
        constant constexpr uint RATIO = 4;
        constant constexpr float RMS_EPS = 1e-06f;
        constant constexpr float ROPE_ATTENTION_SCALE = 1.0f;
    
// Source: mtplx/kernels/qsa_indexer_prepare.py::_pool_keys_kernel
// File SHA256: a77f6ca5ae805e729519c4629ae88b455a6dbf473a457a6e1c8219174eb59091
// Body SHA256: 269cdfebf5e4fb08470e1f4b40c0f7568e9001b3ccb003e7162f2ba477e0c8cd
template <typename T>
kernel void kernel_qwen_mtplx_qsa_pool_k(
    device const T* raw_keys [[buffer(0)]],
    constant const int64_t* raw_keys_strides [[buffer(1)]],
    device const T* norm_weight [[buffer(2)]],
    constant const int64_t* norm_weight_strides [[buffer(3)]],
    device const float* inv_freq [[buffer(4)]],
    constant const int64_t* inv_freq_strides [[buffer(5)]],
    device const int* block_start [[buffer(6)]],
    device T* pooled [[buffer(7)]],
    uint3 threadgroup_position_in_grid [[threadgroup_position_in_grid]],
    uint thread_index_in_simdgroup [[thread_index_in_simdgroup]]) {
        const uint block = threadgroup_position_in_grid.x;
        const uint lane = thread_index_in_simdgroup;
        const uint lane_base = lane * 4u;
        threadgroup float rounded_means[HEAD_DIM];

        // The stock path reduces raw keys in float32, divides by the block
        // width, then casts back to the raw-key dtype before RMSNorm.
        float square_sum = 0.0f;
        for (uint i = 0; i < 4u; ++i) {
            const uint dim = lane_base + i;
            if (dim < HEAD_DIM) {
                float sum = 0.0f;
                for (uint within = 0; within < RATIO; ++within) {
                    const uint token = block * RATIO + within;
                    sum += float(raw_keys[
                        (size_t)token * raw_keys_strides[1] +
                        (size_t)dim * raw_keys_strides[2]]);
                }
                const T rounded = static_cast<T>(sum / float(RATIO));
                const float mean_value = float(rounded);
                rounded_means[dim] = mean_value;
                square_sum += mean_value * mean_value;
            }
        }
        square_sum = simd_sum(square_sum);
        const float inverse_rms = metal::precise::rsqrt(
            square_sum / float(HEAD_DIM) + RMS_EPS);
        threadgroup_barrier(mem_flags::mem_threadgroup);

        const size_t out_base = (size_t)block * HEAD_DIM;
        const float position = float(
            (block_start[0] + int(block)) * int(RATIO));
        for (uint pair = lane; pair < HALF_ROTARY; pair += 32u) {
            const T first_norm = norm_weight[
                (size_t)pair * norm_weight_strides[0]] *
                static_cast<T>(rounded_means[pair] * inverse_rms);
            const T second_norm = norm_weight[
                (size_t)(pair + HALF_ROTARY) * norm_weight_strides[0]] *
                static_cast<T>(
                    rounded_means[pair + HALF_ROTARY] * inverse_rms);
            const float theta = position * float(inv_freq[
                (size_t)pair * inv_freq_strides[0]]);
            const float cosine =
                metal::precise::cos(theta) * ROPE_ATTENTION_SCALE;
            const float sine =
                metal::precise::sin(theta) * ROPE_ATTENTION_SCALE;
            const float first = float(first_norm);
            const float second = float(second_norm);
            const float first_cosine = first * cosine;
            const float second_sine = second * sine;
            const float second_cosine = second * cosine;
            const float first_sine = first * sine;
            pooled[out_base + pair] =
                static_cast<T>(first_cosine - second_sine);
            pooled[out_base + pair + HALF_ROTARY] =
                static_cast<T>(second_cosine + first_sine);
        }
        for (uint dim = ROTARY_DIM + lane; dim < HEAD_DIM; dim += 32u) {
            pooled[out_base + dim] = norm_weight[
                (size_t)dim * norm_weight_strides[0]] *
                static_cast<T>(rounded_means[dim] * inverse_rms);
        }
    }
typedef decltype(kernel_qwen_mtplx_qsa_pool_k<bfloat>) kernel_qwen_mtplx_qsa_pool_k_bf16_type;
template [[host_name("kernel_qwen_mtplx_qsa_pool_k_bf16")]]
kernel kernel_qwen_mtplx_qsa_pool_k_bf16_type kernel_qwen_mtplx_qsa_pool_k<bfloat>;
} // namespace mtplx_qsa_pool_k
#include <metal_stdlib>

namespace mtplx_qsa_flash_skip {

using namespace metal;

// Source: mtplx/kernels/qsa_flash_skip.py::_SRC
// File SHA256: 27d49eb93f0de82013fdc0b271e77076b19d1972450e1bbc44cdd92984669103
// Body SHA256: eca43ef00a564482937984a1c9f8d05d0f7a432a1ed19a6cbbb5b4475f7b246a
template <typename T, int GQA>
kernel void kernel_qwen_mtplx_qsa_flash_skip(
    device const T* q [[buffer(0)]],
    device const T* k [[buffer(1)]],
    device const T* v [[buffer(2)]],
    device const int* blocks [[buffer(3)]],
    device const int* params [[buffer(4)]],
    device const float* scale [[buffer(5)]],
    device T* out [[buffer(6)]],
    uint3 threadgroup_position_in_grid [[threadgroup_position_in_grid]],
    uint3 thread_position_in_threadgroup [[thread_position_in_threadgroup]]) {
    constexpr int HD = 256;
    constexpr int BLK = 4;                    // indexer compress ratio

    const uint hq = threadgroup_position_in_grid.x;
    const uint tid = thread_position_in_threadgroup.x;
    const uint sg = tid / 32;
    const uint lane = tid % 32;

    const int t_total = params[0];            // valid kv length T
    const int tail_start = params[1];         // first tail token
    const int nsel = params[2];               // selected block count
    const int cap = params[3];                // kv allocation stride (dim2)

    threadgroup float q_s[HD];
    if (tid < (uint)HD) q_s[tid] = (float)q[hq * HD + tid] * scale[0];
    threadgroup_barrier(mem_flags::mem_threadgroup);

    const uint hkv = hq / GQA;
    device const T* kh = k + (size_t)hkv * (size_t)cap * HD;
    device const T* vh = v + (size_t)hkv * (size_t)cap * HD;

    float m = -3.0e38f;
    float l = 0.0f;
    float o[8] = {0.0f, 0.0f, 0.0f, 0.0f, 0.0f, 0.0f, 0.0f, 0.0f};
    const uint dbase = lane * 8;

    // selected complete blocks, striped across the 8 simdgroups
    for (int bi = (int)sg; bi < nsel; bi += 8) {
        const int t0 = blocks[bi] * BLK;
        for (int j = 0; j < BLK; ++j) {
            const size_t t = (size_t)(t0 + j);
            float part = 0.0f;
            device const T* kr = kh + t * HD + dbase;
            for (int d = 0; d < 8; ++d) part += q_s[dbase + d] * (float)kr[d];
            const float s = simd_sum(part);
            const float mn = metal::max(m, s);
            const float c = metal::exp(m - mn);
            const float p = metal::exp(s - mn);
            l = l * c + p;
            device const T* vr = vh + t * HD + dbase;
            for (int d = 0; d < 8; ++d) o[d] = o[d] * c + p * (float)vr[d];
            m = mn;
        }
    }
    // visible tail keys, same striping
    for (int t = tail_start + (int)sg; t < t_total; t += 8) {
        float part = 0.0f;
        device const T* kr = kh + (size_t)t * HD + dbase;
        for (int d = 0; d < 8; ++d) part += q_s[dbase + d] * (float)kr[d];
        const float s = simd_sum(part);
        const float mn = metal::max(m, s);
        const float c = metal::exp(m - mn);
        const float p = metal::exp(s - mn);
        l = l * c + p;
        device const T* vr = vh + (size_t)t * HD + dbase;
        for (int d = 0; d < 8; ++d) o[d] = o[d] * c + p * (float)vr[d];
        m = mn;
    }

    // two-level merge across the 8 simdgroups
    threadgroup float m_tg[8];
    threadgroup float l_tg[8];
    threadgroup float o_tg[8 * HD];
    if (lane == 0) { m_tg[sg] = m; l_tg[sg] = l; }
    for (int d = 0; d < 8; ++d) o_tg[sg * HD + dbase + d] = o[d];
    threadgroup_barrier(mem_flags::mem_threadgroup);

    // each of the 256 threads owns one output dim
    float M = m_tg[0];
    for (int s2 = 1; s2 < 8; ++s2) M = metal::max(M, m_tg[s2]);
    float L = 0.0f;
    float acc = 0.0f;
    for (int s2 = 0; s2 < 8; ++s2) {
        const float w = metal::exp(m_tg[s2] - M);
        L += l_tg[s2] * w;
        acc += o_tg[s2 * HD + tid] * w;
    }
    out[hq * HD + tid] = (T)(acc / L);
}
typedef decltype(kernel_qwen_mtplx_qsa_flash_skip<bfloat, 1>) kernel_qwen_mtplx_qsa_flash_skip_g1_bf16_type;
template [[host_name("kernel_qwen_mtplx_qsa_flash_skip_g1_bf16")]]
kernel kernel_qwen_mtplx_qsa_flash_skip_g1_bf16_type kernel_qwen_mtplx_qsa_flash_skip<bfloat, 1>;
typedef decltype(kernel_qwen_mtplx_qsa_flash_skip<half, 1>) kernel_qwen_mtplx_qsa_flash_skip_g1_f16_type;
template [[host_name("kernel_qwen_mtplx_qsa_flash_skip_g1_f16")]]
kernel kernel_qwen_mtplx_qsa_flash_skip_g1_f16_type kernel_qwen_mtplx_qsa_flash_skip<half, 1>;
typedef decltype(kernel_qwen_mtplx_qsa_flash_skip<bfloat, 12>) kernel_qwen_mtplx_qsa_flash_skip_g12_bf16_type;
template [[host_name("kernel_qwen_mtplx_qsa_flash_skip_g12_bf16")]]
kernel kernel_qwen_mtplx_qsa_flash_skip_g12_bf16_type kernel_qwen_mtplx_qsa_flash_skip<bfloat, 12>;
typedef decltype(kernel_qwen_mtplx_qsa_flash_skip<half, 12>) kernel_qwen_mtplx_qsa_flash_skip_g12_f16_type;
template [[host_name("kernel_qwen_mtplx_qsa_flash_skip_g12_f16")]]
kernel kernel_qwen_mtplx_qsa_flash_skip_g12_f16_type kernel_qwen_mtplx_qsa_flash_skip<half, 12>;
} // namespace mtplx_qsa_flash_skip
#include <MetalPerformancePrimitives/MetalPerformancePrimitives.h>

namespace mtplx_qsa_prefill_flash {

using namespace metal;

constant constexpr short QSA_ELEMS_PER_FRAG = 8;
constant constexpr short QSA_ELEM_COLS = 4;
constant constexpr short QSA_ELEM_ROWS_JUMP = 8;

// NAX fragment lane -> (column, row) mapping used by sdpa_nax_tile.py and the
// proven M5/G17G MetalPerformancePrimitives 16x32x16 descriptor.
inline short2 qsa_nax_coord(ushort lane) {
    short quad = short(lane >> 2);
    short row = ((quad & 4) | ((short(lane) >> 1) & 3));
    short col = ((quad & 2) | (short(lane) & 1)) * 4;
    return short2{col, row};
}

// Source: mtplx/kernels/qsa_prefill_flash.py::_SOURCE
// File SHA256: 5ca852ea441810a47b374a1c94a49a636055dab31f7f11c4ce0118f7f70b388d
// Body SHA256: 26ac0e2c2cb8c45a4eb20b18ff7542ddae9c5db8a6940dce7c04c25126a5306e
template <typename T>
kernel void kernel_qwen_mtplx_qsa_prefill_flash(
    device const T* q [[buffer(0)]],
    constant const int64_t* q_strides [[buffer(1)]],
    device const T* k [[buffer(2)]],
    constant const int64_t* k_strides [[buffer(3)]],
    device const T* v [[buffer(4)]],
    constant const int64_t* v_strides [[buffer(5)]],
    device const int* block_ids [[buffer(6)]],
    constant const int64_t* block_ids_strides [[buffer(7)]],
    device const bool* block_valid [[buffer(8)]],
    constant const int64_t* block_valid_strides [[buffer(9)]],
    device const int* params [[buffer(10)]],
    device const float* scale [[buffer(11)]],
    device T* out [[buffer(12)]],
    uint3 threadgroup_position_in_grid [[threadgroup_position_in_grid]],
    uint thread_index_in_simdgroup [[thread_index_in_simdgroup]]) {
    constexpr int HEAD_DIM = 256;
    constexpr int KV_HEADS = 2;
    constexpr int GQA = 12;
    constexpr int BLOCK_TOKENS = 4;
    constexpr int TOP_K_BLOCKS = 512;
    constexpr int M_ROWS = 16;
    constexpr int TILE_BLOCKS = 8;
    constexpr int TOKENS_PER_TILE = TILE_BLOCKS * BLOCK_TOKENS;
    constexpr int MAX_SELECTED_TILES = TOP_K_BLOCKS / TILE_BLOCKS;
    constexpr int D_FRAGS = HEAD_DIM / 16;
    constexpr int OUT_GROUPS = HEAD_DIM / 32;

    // MLX Steel's wide-load idiom: a byte aggregate aligned only to one T.
    // Unlike vec<T,4>, this remains defined for unit-stride views whose base
    // begins at an odd fp16/bf16 element offset.
    struct alignas(sizeof(T)) QSAReadVector8 {
        uchar bytes[sizeof(T) * 8];
    };

    const ushort lane = ushort(thread_index_in_simdgroup);
    const int work = int(threadgroup_position_in_grid.x);
    const int row = work / KV_HEADS;
    const int kv_head = work - row * KV_HEADS;
    const int pos_start = int(params[0]);
    const int total_tokens = int(params[1]);
    const int query_pos = pos_start + row;
    const int complete_blocks = (query_pos + 1) / BLOCK_TOKENS;
    const int tail_start = complete_blocks * BLOCK_TOKENS;
    const int tail_count = query_pos + 1 - tail_start;
    // Find the highest selection slot that can actually contribute.  The
    // production selector emits a chronological valid prefix, but deriving
    // this from the slots themselves preserves the public kernel's existing
    // semantics for arbitrary validity holes as well.  Encoding slot+1 lets
    // zero mean "no selected tile" in the simd-wide max reduction.
    uint local_active_slots = 0u;
    for (uint block_slot = uint(lane);
         block_slot < uint(TOP_K_BLOCKS);
         block_slot += 32u) {
        const size_t id_at = size_t(row) * block_ids_strides[0] +
            size_t(block_slot) * block_ids_strides[1];
        const size_t valid_at = size_t(row) * block_valid_strides[0] +
            size_t(block_slot) * block_valid_strides[1];
        const int block_id = int(block_ids[id_at]);
        if (bool(block_valid[valid_at]) &&
            block_id >= 0 && block_id < complete_blocks) {
            local_active_slots = metal::max(local_active_slots, block_slot + 1u);
        }
    }
    const uint active_blocks = simd_max(local_active_slots);
    const int active_selected_tiles =
        (int(active_blocks) + TILE_BLOCKS - 1) / TILE_BLOCKS;
    const int active_tiles = active_selected_tiles + (tail_count > 0 ? 1 : 0);
    const short2 sc = qsa_nax_coord(lane);

    constexpr auto desc = mpp::tensor_ops::matmul2d_descriptor(
        16, 32, 16, false, true, true,
        mpp::tensor_ops::matmul2d_descriptor::mode::multiply_accumulate);
    mpp::tensor_ops::matmul2d<desc, metal::execution_simdgroup> mm;
    auto ct_a = mm.get_left_input_cooperative_tensor<T, T, float>();
    auto ct_b = mm.get_right_input_cooperative_tensor<T, T, float>();
    auto ct_c =
        mm.get_destination_cooperative_tensor<decltype(ct_a), decltype(ct_b), float>();

    // Exactly 32 * 256 * sizeof(T) == 16 KiB. The allocation first holds K
    // row-major and is then overwritten with V-transpose for the PV multiply.
    threadgroup T tg_tile[TOKENS_PER_TILE * HEAD_DIM];

    // Each lane carries two padded M rows. Four lanes share a logical row;
    // xor(1),xor(8) reduce the four score fragments for that row.
    float row_max[2] = {-1.0e38f, -1.0e38f};
    float row_sum[2] = {0.0f, 0.0f};
    float out_frag[OUT_GROUPS][2][QSA_ELEMS_PER_FRAG];
    for (int group = 0; group < OUT_GROUPS; ++group) {
        for (short row_part = 0; row_part < 2; ++row_part) {
            for (short elem = 0; elem < QSA_ELEMS_PER_FRAG; ++elem) {
                out_frag[group][row_part][elem] = 0.0f;
            }
        }
    }

    for (int tile_index = 0; tile_index < active_tiles; ++tile_index) {
        const bool tail_tile =
            tail_count > 0 && tile_index == active_selected_tiles;
        const int block_base = tile_index * TILE_BLOCKS;

        // Gather K into tg_tile[token,dim]. Metadata is deliberately recomputed
        // by the one simdgroup so the kernel owns only one 16 KiB TG allocation.
        for (int token_slot = 0; token_slot < TOKENS_PER_TILE; ++token_slot) {
            int token = 0;
            bool token_valid = false;
            if (tail_tile) {
                token = tail_start + token_slot;
                token_valid = token_slot < tail_count && token < total_tokens;
            } else {
                const int local_block = token_slot / BLOCK_TOKENS;
                const int within = token_slot - local_block * BLOCK_TOKENS;
                const int block_slot = block_base + local_block;
                const size_t id_at = size_t(row) * block_ids_strides[0] +
                    size_t(block_slot) * block_ids_strides[1];
                const size_t valid_at = size_t(row) * block_valid_strides[0] +
                    size_t(block_slot) * block_valid_strides[1];
                const int block_id = int(block_ids[id_at]);
                token_valid = bool(block_valid[valid_at]) &&
                    block_id >= 0 && block_id < complete_blocks;
                token = token_valid ? block_id * BLOCK_TOKENS + within : 0;
            }
            if (k_strides[3] == 1) {
                // Each lane exposes one compiler-visible 16-byte copy. Across
                // the simdgroup the full 256-wide cache row is contiguous.
                const int dim0 = int(lane) * 8;
                const int destination = token_slot * HEAD_DIM + dim0;
                if (token_valid) {
                    const size_t k_at = size_t(kv_head) * k_strides[1] +
                        size_t(token) * k_strides[2] + size_t(dim0);
                    *reinterpret_cast<threadgroup QSAReadVector8*>(
                        &tg_tile[destination]) =
                        *reinterpret_cast<const device QSAReadVector8*>(&k[k_at]);
                } else {
                    for (short elem = 0; elem < 8; ++elem) {
                        tg_tile[destination + int(elem)] = T(0);
                    }
                }
            } else {
                // Fail-correct for unusual cache views whose feature axis is
                // not contiguous; the production backing takes the wide-copy lane.
                for (int dim = int(lane); dim < HEAD_DIM; dim += 32) {
                    const size_t k_at = size_t(kv_head) * k_strides[1] +
                        size_t(token) * k_strides[2] +
                        size_t(dim) * k_strides[3];
                    tg_tile[token_slot * HEAD_DIM + dim] =
                        token_valid ? k[k_at] : T(0);
                }
            }
        }
        simdgroup_barrier(mem_flags::mem_threadgroup);

        // QK: padded [16,256] x gathered [32,256]^T -> [16,32].
        for (short elem = 0; elem < 2 * QSA_ELEMS_PER_FRAG; ++elem) {
            ct_c[elem] = 0.0f;
        }
        for (int frag = 0; frag < D_FRAGS; ++frag) {
            for (short row_part = 0; row_part < 2; ++row_part) {
                const int m_row = int(sc.y) + int(row_part) * QSA_ELEM_ROWS_JUMP;
                if (m_row < GQA) {
                    const int q_head = kv_head * GQA + m_row;
                    const size_t q_base = size_t(q_head) * q_strides[1] +
                        size_t(row) * q_strides[2] +
                        size_t(frag * 16 + int(sc.x)) * q_strides[3];
                    for (short col = 0; col < QSA_ELEM_COLS; ++col) {
                        ct_a[row_part * QSA_ELEM_COLS + col] =
                            q[q_base + size_t(col) * q_strides[3]];
                    }
                } else {
                    for (short col = 0; col < QSA_ELEM_COLS; ++col) {
                        ct_a[row_part * QSA_ELEM_COLS + col] = T(0);
                    }
                }
            }
            for (short n_half = 0; n_half < 2; ++n_half) {
                for (short row_half = 0; row_half < 2; ++row_half) {
                    const int token_slot = int(n_half) * 16 + int(sc.y) +
                        int(row_half) * QSA_ELEM_ROWS_JUMP;
                    const int tile_base = token_slot * HEAD_DIM +
                        frag * 16 + int(sc.x);
                    for (short col = 0; col < QSA_ELEM_COLS; ++col) {
                        ct_b[n_half * QSA_ELEMS_PER_FRAG +
                             row_half * QSA_ELEM_COLS + col] =
                            tg_tile[tile_base + int(col)];
                    }
                }
            }
            mm.run(ct_a, ct_b, ct_c);
        }

        // Mask, scale, and exponentiate the score fragment in registers.
        float probabilities[2][QSA_ELEMS_PER_FRAG];
        float correction[2];
        for (short row_part = 0; row_part < 2; ++row_part) {
            const int m_row = int(sc.y) + int(row_part) * QSA_ELEM_ROWS_JUMP;
            const bool live_row = m_row < GQA;
            float tile_max = -1.0e38f;
            for (short n_half = 0; n_half < 2; ++n_half) {
                for (short col = 0; col < QSA_ELEM_COLS; ++col) {
                    const int token_slot = int(n_half) * 16 + int(sc.x) + int(col);
                    bool token_valid = false;
                    if (tail_tile) {
                        const int token = tail_start + token_slot;
                        token_valid = token_slot < tail_count && token < total_tokens;
                    } else {
                        const int local_block = token_slot / BLOCK_TOKENS;
                        const int block_slot = block_base + local_block;
                        const size_t id_at = size_t(row) * block_ids_strides[0] +
                            size_t(block_slot) * block_ids_strides[1];
                        const size_t valid_at = size_t(row) * block_valid_strides[0] +
                            size_t(block_slot) * block_valid_strides[1];
                        const int block_id = int(block_ids[id_at]);
                        token_valid = bool(block_valid[valid_at]) &&
                            block_id >= 0 && block_id < complete_blocks;
                    }
                    const short at = row_part * QSA_ELEM_COLS + col;
                    const float score = live_row && token_valid
                        ? ct_c[n_half * QSA_ELEMS_PER_FRAG + at] * scale[0]
                        : -1.0e38f;
                    probabilities[n_half][at] = score;
                    tile_max = metal::max(tile_max, score);
                }
            }
            tile_max = metal::max(tile_max, simd_shuffle_xor(tile_max, ushort(1)));
            tile_max = metal::max(tile_max, simd_shuffle_xor(tile_max, ushort(8)));
            const float new_max = metal::max(row_max[row_part], tile_max);
            correction[row_part] = metal::exp(row_max[row_part] - new_max);
            float tile_sum = 0.0f;
            for (short n_half = 0; n_half < 2; ++n_half) {
                for (short col = 0; col < QSA_ELEM_COLS; ++col) {
                    const short at = row_part * QSA_ELEM_COLS + col;
                    const float score = probabilities[n_half][at];
                    const float probability = score > -1.0e37f
                        ? metal::exp(score - new_max)
                        : 0.0f;
                    probabilities[n_half][at] = probability;
                    tile_sum += probability;
                }
            }
            tile_sum += simd_shuffle_xor(tile_sum, ushort(1));
            tile_sum += simd_shuffle_xor(tile_sum, ushort(8));
            row_max[row_part] = new_max;
            row_sum[row_part] = row_sum[row_part] * correction[row_part] + tile_sum;
        }

        for (short row_part = 0; row_part < 2; ++row_part) {
            const float factor = correction[row_part];
            for (int group = 0; group < OUT_GROUPS; ++group) {
                for (short dim_half = 0; dim_half < 2; ++dim_half) {
                    for (short col = 0; col < QSA_ELEM_COLS; ++col) {
                        out_frag[group][dim_half]
                                [row_part * QSA_ELEM_COLS + col] *= factor;
                    }
                }
            }
        }

        // QK is finished: overwrite the same 16 KiB with V^T[dim,token].
        simdgroup_barrier(mem_flags::mem_threadgroup);
        for (int token_slot = 0; token_slot < TOKENS_PER_TILE; ++token_slot) {
            int token = 0;
            bool token_valid = false;
            if (tail_tile) {
                token = tail_start + token_slot;
                token_valid = token_slot < tail_count && token < total_tokens;
            } else {
                const int local_block = token_slot / BLOCK_TOKENS;
                const int within = token_slot - local_block * BLOCK_TOKENS;
                const int block_slot = block_base + local_block;
                const size_t id_at = size_t(row) * block_ids_strides[0] +
                    size_t(block_slot) * block_ids_strides[1];
                const size_t valid_at = size_t(row) * block_valid_strides[0] +
                    size_t(block_slot) * block_valid_strides[1];
                const int block_id = int(block_ids[id_at]);
                token_valid = bool(block_valid[valid_at]) &&
                    block_id >= 0 && block_id < complete_blocks;
                token = token_valid ? block_id * BLOCK_TOKENS + within : 0;
            }
            if (v_strides[3] == 1) {
                const int dim0 = int(lane) * 8;
                thread T v_lane[8];
                if (token_valid) {
                    const size_t v_at = size_t(kv_head) * v_strides[1] +
                        size_t(token) * v_strides[2] + size_t(dim0);
                    *reinterpret_cast<thread QSAReadVector8*>(&v_lane[0]) =
                        *reinterpret_cast<const device QSAReadVector8*>(&v[v_at]);
                } else {
                    for (short elem = 0; elem < 8; ++elem) {
                        v_lane[elem] = T(0);
                    }
                }
                for (short elem = 0; elem < 8; ++elem) {
                    tg_tile[(dim0 + int(elem)) * TOKENS_PER_TILE + token_slot] =
                        v_lane[elem];
                }
            } else {
                for (int dim = int(lane); dim < HEAD_DIM; dim += 32) {
                    const size_t v_at = size_t(kv_head) * v_strides[1] +
                        size_t(token) * v_strides[2] +
                        size_t(dim) * v_strides[3];
                    tg_tile[dim * TOKENS_PER_TILE + token_slot] =
                        token_valid ? v[v_at] : T(0);
                }
            }
        }
        simdgroup_barrier(mem_flags::mem_threadgroup);

        // PV: [16,32] x [256,32]^T -> [16,256]. MPP consumes probabilities
        // in T, while accumulation and all persistent online state remain fp32.
        for (int group = 0; group < OUT_GROUPS; ++group) {
            for (short elem = 0; elem < QSA_ELEMS_PER_FRAG; ++elem) {
                ct_c[elem] = out_frag[group][0][elem];
                ct_c[QSA_ELEMS_PER_FRAG + elem] = out_frag[group][1][elem];
            }
            for (short token_half = 0; token_half < 2; ++token_half) {
                for (short row_part = 0; row_part < 2; ++row_part) {
                    for (short col = 0; col < QSA_ELEM_COLS; ++col) {
                        ct_a[row_part * QSA_ELEM_COLS + col] = T(
                            probabilities[token_half][row_part * QSA_ELEM_COLS + col]);
                    }
                }
                for (short dim_half = 0; dim_half < 2; ++dim_half) {
                    for (short row_half = 0; row_half < 2; ++row_half) {
                        const int dim = group * 32 + int(dim_half) * 16 +
                            int(sc.y) + int(row_half) * QSA_ELEM_ROWS_JUMP;
                        const int tile_base = dim * TOKENS_PER_TILE +
                            int(token_half) * 16 + int(sc.x);
                        for (short col = 0; col < QSA_ELEM_COLS; ++col) {
                            ct_b[dim_half * QSA_ELEMS_PER_FRAG +
                                 row_half * QSA_ELEM_COLS + col] =
                                tg_tile[tile_base + int(col)];
                        }
                    }
                }
                mm.run(ct_a, ct_b, ct_c);
            }
            for (short elem = 0; elem < QSA_ELEMS_PER_FRAG; ++elem) {
                out_frag[group][0][elem] = ct_c[elem];
                out_frag[group][1][elem] = ct_c[QSA_ELEMS_PER_FRAG + elem];
            }
        }
        simdgroup_barrier(mem_flags::mem_threadgroup);
    }

    // Normalize and store the twelve live rows into contiguous [1,24,S,256].
    for (short row_part = 0; row_part < 2; ++row_part) {
        const int m_row = int(sc.y) + int(row_part) * QSA_ELEM_ROWS_JUMP;
        if (m_row >= GQA) continue;
        const int q_head = kv_head * GQA + m_row;
        const float inv_sum = row_sum[row_part] > 0.0f ? 1.0f / row_sum[row_part] : 0.0f;
        const size_t out_base =
            (size_t(q_head) * size_t(params[2]) + size_t(row)) * HEAD_DIM;
        for (int group = 0; group < OUT_GROUPS; ++group) {
            for (short dim_half = 0; dim_half < 2; ++dim_half) {
                for (short col = 0; col < QSA_ELEM_COLS; ++col) {
                    const int dim = group * 32 + int(dim_half) * 16 +
                        int(sc.x) + int(col);
                    out[out_base + size_t(dim)] = T(
                        out_frag[group][dim_half][row_part * QSA_ELEM_COLS + col] *
                        inv_sum);
                }
            }
        }
    }
}
typedef decltype(kernel_qwen_mtplx_qsa_prefill_flash<bfloat>) kernel_qwen_mtplx_qsa_prefill_flash_bf16_type;
template [[host_name("kernel_qwen_mtplx_qsa_prefill_flash_bf16")]]
kernel kernel_qwen_mtplx_qsa_prefill_flash_bf16_type kernel_qwen_mtplx_qsa_prefill_flash<bfloat>;
typedef decltype(kernel_qwen_mtplx_qsa_prefill_flash<half>) kernel_qwen_mtplx_qsa_prefill_flash_f16_type;
template [[host_name("kernel_qwen_mtplx_qsa_prefill_flash_f16")]]
kernel kernel_qwen_mtplx_qsa_prefill_flash_f16_type kernel_qwen_mtplx_qsa_prefill_flash<half>;
} // namespace mtplx_qsa_prefill_flash