// MTPLX 9882703, MIT; see MTPLX-LICENSE.txt.

    constexpr int BM = 64;
    constexpr int BN = 64;
    constexpr int BK = 64;
    constexpr int WN = 2;
    constexpr short SM = 32;
    constexpr short SN = 32;
    constexpr short SK = 32;
    constexpr short TM = SM / 16;
    constexpr short TN = SN / 16;
    constexpr short TK = SK / 16;
    constexpr int PACK = get_pack_factor<BITS, 8>();
    constexpr int PACK_BYTES = get_bytes_per_pack<BITS>();
    constexpr int BK_PAD = BK + 16 / sizeof(T);
    constexpr int ROW_BYTES = KD * PACK_BYTES / PACK;
    constexpr int ROW_GROUPS = KD / GS;
    constexpr int OUT_COLS = OUTC;
    // SWIGLU: the tile's first 32 weight rows are gate rows, the last 32 the
    // matching up rows, each half loaded by 64 threads from its own row range.
    using WeightLoader = metal::conditional_t<
        SWIGLU,
        QuantizedBlockLoader<T, BN / 2, BK, BK_PAD, 1, 64, GS, BITS>,
        QuantizedBlockLoader<T, BN, BK, BK_PAD, 1, 128, GS, BITS>>;

    threadgroup T w_tile[BN * BK_PAD];

    // The tile this threadgroup owns, and the one expert whose rows it holds:
    // the last expert whose first tile is at or before it.
    const int tile = int(threadgroup_position_in_grid.y);
    if (tile >= tile_start[ED]) {
        return;
    }
    int lo = 0;
    int hi = ED - 1;
    while (lo < hi) {
        const int mid = (lo + hi + 1) >> 1;
        if (tile_start[mid] <= tile) {
            lo = mid;
        } else {
            hi = mid - 1;
        }
    }
    const int expert = lo;
    const int row0 = row_start[expert] + (tile - tile_start[expert]) * BM;
    const int tile_rows = min(BM, row_start[expert + 1] - row0);
    const int col0 = int(threadgroup_position_in_grid.x) * BN;

    const ushort sg = ushort(simdgroup_index_in_threadgroup);
    const ushort lane = ushort(thread_index_in_simdgroup);
    const short tm = SM * short(sg / WN);
    const short tn = SN * short(sg % WN);
    // Rows of this tile that fall in this simdgroup's 32-row half (0 to 32).
    const short rows = short(clamp(tile_rows - int(tm), 0, int(SM)));

    ushort load_sg = sg;
    const device uint8_t* w_src = (const device uint8_t*)w;
    const device T* s_src = scales;
    const device T* b_src = biases;
    size_t w_row = size_t(expert) * ND + size_t(col0);
    threadgroup T* w_dst = w_tile;
    if (SWIGLU) {
        // Output columns [col0 / 2, col0 / 2 + 32): the matching gate rows
        // of w and up rows of wu (the same array at an offset when the
        // weight is a fused [gate | up]).
        const int up_half = int(load_sg / 2);
        const size_t c = size_t(col0 / 2);
        if (up_half) {
            w_src = (const device uint8_t*)wu;
            s_src = scales_u;
            b_src = biases_u;
            w_row = size_t(expert) * UD + size_t(UOFF) + c;
        } else {
            w_row = size_t(expert) * ND + c;
        }
        w_dst = w_tile + up_half * (BN / 2) * BK_PAD;
        load_sg = load_sg % 2;
    }
    thread WeightLoader loader(
        w_src + w_row * ROW_BYTES,
        s_src + w_row * ROW_GROUPS,
        b_src + w_row * ROW_GROUPS,
        KD,
        w_dst,
        load_sg,
        lane);

    NAXTile<float, TM, TN> acc;
    acc.clear();

    // A lane holds rows (y, y + 8) and four consecutive columns from x of every
    // 16x16 fragment (MLX's fragment coordinates).  Its four token rows are
    // fixed for the whole K loop; rows past the tile read as zero, as MLX's
    // bounded load does.
    const short2 coord = BaseNAXFrag::get_coord();
    const device T* lane_rows[TM][2];
    bool lane_live[TM][2];
    for (short i = 0; i < TM; i++) {
        for (short r = 0; r < 2; r++) {
            const short local = i * 16 + coord.y + r * 8;
            lane_live[i][r] = local < rows;
            const int token = lane_live[i][r] ? int(row_map[row0 + tm + local]) : 0;
            lane_rows[i][r] = x + size_t(token) * KD + coord.x;
        }
    }

    // The K loop, compiled twice: for a simdgroup whose 32 rows are all live
    // (no bounds) and for a partial or empty half.
    dispatch_bool(rows == SM, [&](auto full) {
        for (int k = 0; k < KD / BK; k++) {
            threadgroup_barrier(mem_flags::mem_threadgroup);
            loader.load_unsafe();
            threadgroup_barrier(mem_flags::mem_threadgroup);

            STEEL_PRAGMA_NO_UNROLL
            for (short kk = 0; kk < BK; kk += SK) {
                if (decltype(full)::value || rows > 0) {
                    NAXTile<T, TM, TK> a;
                    NAXTile<T, TN, TK> b;
                    volatile int keep_order;
                    const int col = k * BK + kk;
                    STEEL_PRAGMA_UNROLL
                    for (short i = 0; i < TM; i++) {
                        STEEL_PRAGMA_UNROLL
                        for (short j = 0; j < TK; j++) {
                            thread auto& frag = a.frag_at(i, j);
                            STEEL_PRAGMA_UNROLL
                            for (short r = 0; r < 2; r++) {
                                const device T* src = lane_rows[i][r] + col + j * 16;
                                STEEL_PRAGMA_UNROLL
                                for (short c = 0; c < 4; c++) {
                                    if constexpr (decltype(full)::value) {
                                        frag[r * 4 + c] = src[c];
                                    } else {
                                        frag[r * 4 + c] = lane_live[i][r] ? src[c] : T(0);
                                    }
                                }
                            }
                        }
                    }
                    b.template load<T, BK_PAD, 1>(w_tile + tn * BK_PAD + kk);
                    tile_matmad_nax(
                        acc,
                        a,
                        metal::bool_constant<false>{},
                        b,
                        metal::bool_constant<true>{});
                    (void)keep_order;
                }
            }
            loader.next();
        }
    });

    if constexpr (SWIGLU) {
        // The up simdgroup of each 32-row half hands its values, rounded to T
        // as the stock gather's output is, to the gate simdgroup of the same
        // rows through the (now free) weight tile, lane to lane.  The gate
        // simdgroup writes silu(gate) * up with MLX's own functors, rounding to
        // T after each op as the stock chain does.
        constexpr short VALS = TM * TN * 8;
        static_assert(2 * 32 * VALS <= BN * BK_PAD, "the hand-off must fit the weight tile");
        threadgroup T* handoff = w_tile + (sg / 2) * (32 * VALS) + lane * VALS;
        const thread float* vals = acc.elems();
        threadgroup_barrier(mem_flags::mem_threadgroup);
        if (tn != 0) {
            for (short e = 0; e < VALS; e++) {
                handoff[e] = static_cast<T>(vals[e]);
            }
        }
        threadgroup_barrier(mem_flags::mem_threadgroup);
        if (tn == 0 && rows > 0) {
            device T* out = y + size_t(row0 + tm) * OUT_COLS + size_t(col0 / 2);
            for (short i = 0; i < TM; i++) {
                for (short j = 0; j < TN; j++) {
                    for (short r = 0; r < 2; r++) {
                        const short row = i * 16 + coord.y + r * 8;
                        if (row >= rows) {
                            continue;
                        }
                        for (short c = 0; c < 4; c++) {
                            const short e = (i * TN + j) * 8 + r * 4 + c;
                            const T gate = static_cast<T>(vals[e]);
                            const T act = Multiply()(gate, Sigmoid()(gate));
                            out[size_t(row) * OUT_COLS + j * 16 + coord.x + c] =
                                Multiply()(act, handoff[e]);
                        }
                    }
                }
            }
        }
    } else if (rows > 0) {
        device T* out = y + size_t(row0 + tm) * ND + size_t(col0 + tn);
        if (rows == SM) {
            acc.store(out, ND);
        } else {
            acc.store_safe(out, ND, short2(SN, rows));
        }
    }