Port Qwen row-map SwiGLU expert gathers

metal/MLX-LICENSE.txt

@@ -0,0 +1,21 @@
+MIT License
+
+Copyright © 2023 Apple Inc.
+
+Permission is hereby granted, free of charge, to any person obtaining a copy
+of this software and associated documentation files (the "Software"), to deal
+in the Software without restriction, including without limitation the rights
+to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
+copies of the Software, and to permit persons to whom the Software is
+furnished to do so, subject to the following conditions:
+
+The above copyright notice and this permission notice shall be included in all
+copies or substantial portions of the Software.
+
+THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
+IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
+FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
+AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
+LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
+OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
+SOFTWARE.

metal/mtplx-moe-sorted-body.metal

@@ -0,0 +1,199 @@
+// 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));
+        }
+    }