gemm_base.cuh 41.7 KB
Newer Older
muyangli's avatar
muyangli committed
1
2
3
4
5
6
7
#pragma once

#include "common.h"

#include "../utils.cuh"
#include "../dispatch_utils.h"
#include "gemm_utils.cuh"
sxtyzhangzk's avatar
sxtyzhangzk committed
8
#include "mma_earlycuda.cuh"
muyangli's avatar
muyangli committed
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28

#pragma nv_diag_suppress 177

#ifdef _MSC_VER
#define ALWAYSINLINE [[msvc::forceinline]]
#else
#define ALWAYSINLINE __attribute__((always_inline))
#endif

// #define ENABLE_NAN_CHECK 1
#if ENABLE_NAN_CHECK
#define STRINGIZE(x) STRINGIZE2(x)
#define STRINGIZE2(x) #x
#define CHECK_NAN(data, name) checkNan(data, name " at " STRINGIZE(__LINE__))
#else
#define CHECK_NAN(data, name)
#endif

namespace nunchaku::kernels {

29
template<bool bf16, bool faster_i2f = false>
muyangli's avatar
muyangli committed
30
31
32
33
class GEMMConfig_W4A4 {
public:
    // BE CAREFUL: weights need to be repacked when the tiling size changes

Muyang Li's avatar
Muyang Li committed
34
35
    static constexpr int BLOCK_M   = 256;
    static constexpr int BLOCK_N   = 128;
muyangli's avatar
muyangli committed
36
37
38
39
40
41
42
    static constexpr int WARP_SIZE = 32;
    static constexpr int NUM_WARPS = 8;

    static constexpr int INSN_M = 16;
    static constexpr int INSN_N = 16;
    static constexpr int INSN_K = 64;

43
44
45
46
    // faster i2f conversion on sm_75
    // may generate incorrect results in certain circumstances
    static constexpr bool FASTER_I2F = faster_i2f;

fengzch-das's avatar
fengzch-das committed
47
48
    using half_t  = typename std::conditional_t<bf16, __nv_bfloat16, half>;
    using half2_t = typename std::conditional_t<bf16, __nv_bfloat162, half2>;
muyangli's avatar
muyangli committed
49
50
};

Muyang Li's avatar
Muyang Li committed
51
52
using GEMMConfig_W4A4_FP16           = GEMMConfig_W4A4<false>;
using GEMMConfig_W4A4_BF16           = GEMMConfig_W4A4<true>;
53
using GEMMConfig_W4A4_FP16_FasterI2F = GEMMConfig_W4A4<false, true>;
muyangli's avatar
muyangli committed
54
55
56

class GEMMConfig_W8A8 {
public:
Muyang Li's avatar
Muyang Li committed
57
58
    static constexpr int BLOCK_M   = 256;
    static constexpr int BLOCK_N   = 128;
muyangli's avatar
muyangli committed
59
60
61
62
63
64
65
66
67
68
69
    static constexpr int WARP_SIZE = 32;
    static constexpr int NUM_WARPS = 8;

    static constexpr int INSN_M = 16;
    static constexpr int INSN_N = 16;
    static constexpr int INSN_K = 32;

#if 0
    using half_t  = half;
    using half2_t = half2;
#else
fengzch-das's avatar
fengzch-das committed
70
71
    using half_t  = __nv_bfloat16;
    using half2_t = __nv_bfloat162;
muyangli's avatar
muyangli committed
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
#endif
};

template<class Config>
class GEMMBase : public Config {
public:
    using Config::BLOCK_M;
    using Config::BLOCK_N;
    using Config::WARP_SIZE;
    using Config::NUM_WARPS;
    using Config::INSN_M;
    using Config::INSN_N;
    using Config::INSN_K;

    using typename Config::half_t;
    using typename Config::half2_t;

    static constexpr int WARP_M = BLOCK_M / NUM_WARPS;
    static constexpr int WARP_N = BLOCK_N;
    static constexpr int WARP_K = INSN_K;

    static constexpr int WARP_M_TILES = WARP_M / INSN_M;
    static constexpr int WARP_N_TILES = WARP_N / INSN_N;
    static constexpr int WARP_K_TILES = WARP_K / INSN_K;

    /**
     * refer to https://docs.nvidia.com/cuda/parallel-thread-execution/index.html#mma-16864-c
Muyang Li's avatar
Muyang Li committed
99
     *
muyangli's avatar
muyangli committed
100
101
102
103
104
     * wscales store order: (pack = 4)
     *  0   1   8   9   <-- load by lane 0, broadcast to lane {0, 4, 8, ..., 28} (8x)
     *  2   3   10  11  <-- load by lane 1, broadcast to lane {1, 5, 9, ..., 29} (8x)
     *  4   5   12  13  <-- load by lane 2, broadcast to lane {2, 6, 10, ..., 30} (8x)
     *  6   7   14  15  <-- load by lane 3, broadcast to lane {3, 7, 11, ..., 31} (8x)
Muyang Li's avatar
Muyang Li committed
105
     *
muyangli's avatar
muyangli committed
106
107
108
109
110
111
112
     *  16  17  24  25  <-- load by lane 4, broadcast to lane {0, 4, 8, ..., 28} (8x)
     *  ...
     *  22  23  30  31  <-- load by lane 7, broadcast to lane {3, 7, 11, ..., 31} (8x)
     *  ... ...
     *  112 113 120 121 <-- load by lane 28, broadcast to lane {0, 4, 8, ..., 28} (8x)
     *  ...
     *  118 119 126 127 <-- load by lane 31, broadcast to lane {3, 7, 11, ..., 31} (8x)
Muyang Li's avatar
Muyang Li committed
113
     *
muyangli's avatar
muyangli committed
114
115
116
117
118
     * wscales store order: (pack = 8)
     *  0   1   8   9   16  17  24  25  <-- load by lane 0, broadcast to lane {0, 4, 8, ..., 28} (8x)
     *  2   3   10  11  18  19  26  27  <-- load by lane 1, broadcast to lane {1, 5, 9, ..., 29} (8x)
     *  4   5   12  13  20  21  28  29  <-- load by lane 2, broadcast to lane {2, 6, 10, ..., 30} (8x)
     *  6   7   14  15  22  23  30  31  <-- load by lane 3, broadcast to lane {3, 7, 11, ..., 31} (8x)
Muyang Li's avatar
Muyang Li committed
119
     *
muyangli's avatar
muyangli committed
120
121
122
     *  224 225 232 233 240 241 248 249 <-- load by lane 28, broadcast to lane {0, 4, 8, ..., 28} (8x)
     *  ...
     *  230 231 238 239 246 247 254 255 <-- load by lane 31, broadcast to lane {3, 7, 11, ..., 31} (8x)
Muyang Li's avatar
Muyang Li committed
123
124
125
126
     *
     * {k}-th wscale used by lane {i} => {k // (WSCALES_PACK_SIZE * WARP_SIZE)}-th pack, in lane {4*(k //
     * WSCALES_PACK_SIZE) + i % 4}, element {k % WSCALES_PACK_SIZE}
     *
muyangli's avatar
muyangli committed
127
128
     * max pack size set to 8 since max load size is 16 bytes / lane
     * min pack size set to 2 since shuffle granularity is 32b 2*half
Muyang Li's avatar
Muyang Li committed
129
130
131
     * */
    static constexpr int WSCALES_PACK_SIZE   = clamp(WARP_N / WARP_SIZE, 4 / sizeof(half), 16 / sizeof(half));
    static constexpr int WSCALES_NUM_PACKS   = ceilDiv(WARP_N, (WSCALES_PACK_SIZE * WARP_SIZE));
muyangli's avatar
muyangli committed
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
    static constexpr int WSCALES_VALID_LANES = std::min(WARP_SIZE, WARP_N / WSCALES_PACK_SIZE);

    /**
     * ascales store order: (pack = 2)
     *  0   8   <-- load by lane 0, broadcast to lane {0, 1, 2, 3} (4x)
     *  1   9   <-- load by lane 1, broadcast to lane {4, 5, 6, 7} (4x)
     *  2   10
     *  ...
     *  6   14
     *  7   15  <-- load by lane 7, broadcast to lane {28, 29, 30, 31} (4x)
     *  ... ...
     *  48  56  <-- load by lane 24, broadcast to lane {0, 1, 2, 3} (4x)
     *  49  57
     *  ...
     *  54  62
     *  55  63  <-- load by lane 31, broadcast to lane {28, 29, 30, 31}  (4x)
Muyang Li's avatar
Muyang Li committed
148
149
150
     *
     * {k}-th wscale used by lane {i} => {k // (ASCALES_PACK_SIZE * WARP_SIZE)}-th pack, in lane {8*(k //
     * ASCALES_PACK_SIZE) + i // 4}, element {k % ASCALES_PACK_SIZE}
muyangli's avatar
muyangli committed
151
     */
Muyang Li's avatar
Muyang Li committed
152
153
    static constexpr int ASCALES_PACK_SIZE   = clamp(WARP_M / WARP_SIZE, 4 / sizeof(half), 16 / sizeof(half));
    static constexpr int ASCALES_NUM_PACKS   = ceilDiv(WARP_M, (ASCALES_PACK_SIZE * WARP_SIZE));
muyangli's avatar
muyangli committed
154
155
156
157
    static constexpr int ASCALES_VALID_LANES = std::min(WARP_SIZE, WARP_M / ASCALES_PACK_SIZE);

    using packed_act_t = uint4;
    using packed_wgt_t = uint4;
Muyang Li's avatar
Muyang Li committed
158

muyangli's avatar
muyangli committed
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
    struct alignas(32) packed_psum_t {
        int data[8];
    };
    struct alignas(16) packed_fpsum_t {
        half2_t data[4];
    };
    struct alignas(8) packed_gated_fpsum_t {
        half_t data[4];
    };
    // 16 * 16 matrix
    struct alignas(32) packed_f32psum_t {
        float data[8];

        static constexpr packed_f32psum_t zeros() {
            packed_f32psum_t result;
            for (int i = 0; i < 8; i++) {
                result.data[i] = 0;
            }
            return result;
        };
    };

    struct packed_wscale_t {
        half2_t data[WSCALES_PACK_SIZE / 2];
    };
    struct packed_ascale_t {
        half2_t data[ASCALES_PACK_SIZE / 2];
    };

Muyang Li's avatar
Muyang Li committed
188
189
190
191
192
193
    using act_warp         = std::array<packed_act_t, WARP_M_TILES>;
    using wgt_warp         = std::array<packed_wgt_t, WARP_N_TILES>;
    using ascale_warp      = std::array<packed_ascale_t, ASCALES_NUM_PACKS>;
    using wscale_warp      = std::array<packed_wscale_t, WSCALES_NUM_PACKS>;
    using fpsum_warp       = std::array<packed_fpsum_t, WARP_M_TILES * WARP_N_TILES>;
    using f32psum_warp     = std::array<packed_f32psum_t, WARP_M_TILES * WARP_N_TILES>;
muyangli's avatar
muyangli committed
194
195
196
197
198
199
200
201
202
    using gated_fpsum_warp = std::array<packed_gated_fpsum_t, WARP_M_TILES * WARP_N_TILES>;

    struct BlockInfo {
        int bm;
        int bn;
        int numBlocksM;
        int numBlocksN;
    };

Muyang Li's avatar
Muyang Li committed
203
204
    __device__ __forceinline__ static packed_f32psum_t
    mma_f16xf16_f32(packed_fpsum_t a, packed_fpsum_t b, packed_f32psum_t psum) {
fengzch-das's avatar
fengzch-das committed
205
        static_assert(std::is_same_v<half_t, half> || std::is_same_v<half_t, __nv_bfloat16>);
muyangli's avatar
muyangli committed
206

fengzch-das's avatar
fengzch-das committed
207
        static constexpr bool is_bf16 = std::is_same_v<half_t, __nv_bfloat16>;
208
209

        uint4 out1 = mma_m16n8k16_f32f16f16f32<is_bf16>(
Muyang Li's avatar
Muyang Li committed
210
211
            kernels::bit_cast<uint4>(a),
            kernels::bit_cast<uint2>(std::array<half2_t, 2>(b.data[0], b.data[1])),
sxtyzhangzk's avatar
sxtyzhangzk committed
212
            kernels::bit_cast<uint4>(float4(psum.data[0], psum.data[1], psum.data[2], psum.data[3])));
213
        uint4 out2 = mma_m16n8k16_f32f16f16f32<is_bf16>(
Muyang Li's avatar
Muyang Li committed
214
215
            kernels::bit_cast<uint4>(a),
            kernels::bit_cast<uint2>(std::array<half2_t, 2>(b.data[2], b.data[3])),
sxtyzhangzk's avatar
sxtyzhangzk committed
216
217
218
219
220
221
222
223
224
            kernels::bit_cast<uint4>(float4(psum.data[4], psum.data[5], psum.data[6], psum.data[7])));
        psum.data[0] = kernels::bit_cast<float>(out1.x);
        psum.data[1] = kernels::bit_cast<float>(out1.y);
        psum.data[2] = kernels::bit_cast<float>(out1.z);
        psum.data[3] = kernels::bit_cast<float>(out1.w);
        psum.data[4] = kernels::bit_cast<float>(out2.x);
        psum.data[5] = kernels::bit_cast<float>(out2.y);
        psum.data[6] = kernels::bit_cast<float>(out2.z);
        psum.data[7] = kernels::bit_cast<float>(out2.w);
Muyang Li's avatar
Muyang Li committed
225

muyangli's avatar
muyangli committed
226
227
228
        return psum;
    }

Muyang Li's avatar
Muyang Li committed
229
    __device__ __forceinline__ static packed_fpsum_t packed_fp32_to_fp16(packed_f32psum_t input) {
muyangli's avatar
muyangli committed
230
231
232
233
234
235
236
        packed_fpsum_t results;
        for (int i = 0; i < 4; i++) {
            results.data[i] = float22half2<half2_t>(float2(input.data[i * 2], input.data[i * 2 + 1]));
        }
        return results;
    }

Muyang Li's avatar
Muyang Li committed
237
    __device__ __forceinline__ static packed_f32psum_t packed_fp16_to_fp32(packed_fpsum_t input) {
muyangli's avatar
muyangli committed
238
239
        packed_f32psum_t results;
        for (int i = 0; i < 4; i++) {
Muyang Li's avatar
Muyang Li committed
240
241
            float2 tmp              = half22float2(input.data[i]);
            results.data[i * 2]     = tmp.x;
muyangli's avatar
muyangli committed
242
243
244
245
246
            results.data[i * 2 + 1] = tmp.y;
        }
        return results;
    }

Muyang Li's avatar
Muyang Li committed
247
    __device__ __forceinline__ static fpsum_warp packed_fp32_to_fp16(f32psum_warp input) {
muyangli's avatar
muyangli committed
248
        fpsum_warp results;
Muyang Li's avatar
Muyang Li committed
249
#pragma unroll
muyangli's avatar
muyangli committed
250
251
252
253
254
255
        for (int i = 0; i < results.size(); i++) {
            results[i] = packed_fp32_to_fp16(input[i]);
        }
        return results;
    }

Muyang Li's avatar
Muyang Li committed
256
    __device__ __forceinline__ static f32psum_warp packed_fp16_to_fp32(fpsum_warp input) {
sxtyzhangzk's avatar
sxtyzhangzk committed
257
        f32psum_warp results;
Muyang Li's avatar
Muyang Li committed
258
#pragma unroll
sxtyzhangzk's avatar
sxtyzhangzk committed
259
260
261
262
263
264
        for (int i = 0; i < results.size(); i++) {
            results[i] = packed_fp16_to_fp32(input[i]);
        }
        return results;
    }

muyangli's avatar
muyangli committed
265
    // activation: row major, [M / BLOCK_M, K / WARP_K, NUM_WARPS, WARP_M_TILES, WARP_SIZE] of packed_act_t
Muyang Li's avatar
Muyang Li committed
266
    __device__ __forceinline__ static void load_act(const packed_act_t *act, int k, int K, act_warp &out, bool pred) {
muyangli's avatar
muyangli committed
267
268
        int laneId = threadIdx.x % WARP_SIZE;
        int warpId = threadIdx.x / WARP_SIZE;
Muyang Li's avatar
Muyang Li committed
269
#pragma unroll
muyangli's avatar
muyangli committed
270
        for (int i = 0; i < WARP_M_TILES; i++) {
Muyang Li's avatar
Muyang Li committed
271
272
273
            // if (pred) {
            //  out[i] = load(&act[((warpId * WARP_M_TILES + i) * K / WARP_K + k) * WARP_SIZE + laneId]);
            out[i] = load_pred(&act[((k * NUM_WARPS + warpId) * WARP_M_TILES + i) * WARP_SIZE + laneId], pred);
274
            //}
muyangli's avatar
muyangli committed
275
276
277
278
        }
    }

    // weight: column major: [N / BLOCK_N, 1, K / WARP_K, WARP_N_TILES, WARP_SIZE] of packed_wgt_t
Muyang Li's avatar
Muyang Li committed
279
    __device__ __forceinline__ static void load_wgt(const packed_wgt_t *wgt, int k, int K, wgt_warp &out, bool pred) {
muyangli's avatar
muyangli committed
280
        int laneId = threadIdx.x % WARP_SIZE;
Muyang Li's avatar
Muyang Li committed
281

muyangli's avatar
muyangli committed
282
283
284
        // const packed_wgt_t *ptr = &wgt[(0 * K / WARP_K + k) * WARP_SIZE + laneId];
        const packed_wgt_t *ptr = &wgt[(0 + k * WARP_N_TILES) * WARP_SIZE + laneId];
        // int offset = K / WARP_K * WARP_SIZE;
Muyang Li's avatar
Muyang Li committed
285
#pragma unroll
muyangli's avatar
muyangli committed
286
        for (int i = 0; i < WARP_N_TILES; i++) {
Muyang Li's avatar
Muyang Li committed
287
288
289
290
291
            // if (pred) {
            //  out[i] = load(&wgt[(i * K / WARP_K + k) * WARP_SIZE + laneId]);
            //  out[i] = load(&wgt[(i + k * WARP_N_TILES) * WARP_SIZE + laneId]);
            out[i] = load_pred(&ptr[i * WARP_SIZE], pred);
            // ptr += offset;
292
            //}
muyangli's avatar
muyangli committed
293
294
295
        }
    }

Muyang Li's avatar
Muyang Li committed
296
297
298
299
    // ascales: row major [M / BLOCK_M, K / group size, NUM_WARPS, ASCALES_NUM_PACKS, ASCALES_VALID_LANES] of
    // packed_ascale_t
    __device__ __forceinline__ static void
    load_ascale(const packed_ascale_t *ascales, int group, int M, ascale_warp &out, bool pred) {
muyangli's avatar
muyangli committed
300
301
        int laneId = threadIdx.x % WARP_SIZE;
        int warpId = threadIdx.x / WARP_SIZE;
Muyang Li's avatar
Muyang Li committed
302
#pragma unroll
muyangli's avatar
muyangli committed
303
        for (int i = 0; i < ASCALES_NUM_PACKS; i++) {
304
            // if (pred && laneId < ASCALES_VALID_LANES) {
Muyang Li's avatar
Muyang Li committed
305
306
307
308
309
            // out[i] = ascales[(group * M / WARP_M + warpId) * ASCALES_VALID_LANES * ASCALES_NUM_PACKS + i *
            // ASCALES_VALID_LANES + laneId];
            out[i] = load_pred(&ascales[(group * NUM_WARPS + warpId) * ASCALES_NUM_PACKS * ASCALES_VALID_LANES +
                                        i * ASCALES_VALID_LANES + laneId],
                               pred && laneId < ASCALES_VALID_LANES);
muyangli's avatar
muyangli committed
310

311
            // }
muyangli's avatar
muyangli committed
312
313
314
        }
    }

Muyang Li's avatar
Muyang Li committed
315
316
317
318
    // wscales: column major [N / BLOCK_N, K / group size, 1, WSCALES_NUM_PACKS, WSCALES_VALID_LANES] of packed_wscale_t
    // </del>
    __device__ __forceinline__ static void
    load_wscale(const packed_wscale_t *wscales, int group, int N, wscale_warp &out, bool pred) {
muyangli's avatar
muyangli committed
319
320
321
322
323
324
        int laneId = threadIdx.x % WARP_SIZE;

        // static_assert(WSCALES_NUM_PACKS * WSCALES_VALID_LANES == 32);
        // static_assert(sizeof(packed_wscale_t) == 8);

        // const packed_wscale_t *ptr = &wscales[(group * WSCALES_NUM_PACKS + 0) * WSCALES_VALID_LANES + laneId];
Muyang Li's avatar
Muyang Li committed
325
326
        // // const packed_wscale_t *ptr = (const packed_wscale_t *)((const char *)wscales) + ((group *
        // WSCALES_NUM_PACKS + 0) * WSCALES_VALID_LANES + laneId) * sizeof(packed_wscale_t);
muyangli's avatar
muyangli committed
327

Muyang Li's avatar
Muyang Li committed
328
#pragma unroll
muyangli's avatar
muyangli committed
329
        for (int i = 0; i < WSCALES_NUM_PACKS; i++) {
330
            // if (pred && laneId < WSCALES_VALID_LANES) {
Muyang Li's avatar
Muyang Li committed
331
332
333
334
335
336
337

            // out[i] = wscales[group * N / WARP_N * WSCALES_VALID_LANES * WSCALES_NUM_PACKS + i * WSCALES_VALID_LANES +
            // laneId]; out[i] = load(&wscales[group * N / WARP_N * WSCALES_VALID_LANES * WSCALES_NUM_PACKS + i *
            // WSCALES_VALID_LANES + laneId]);
            out[i] = load_pred(&wscales[(group * WSCALES_NUM_PACKS + i) * WSCALES_VALID_LANES + laneId],
                               pred && laneId < WSCALES_VALID_LANES);
            // out[i] = load(&ptr[i * WSCALES_VALID_LANES]);
338
            // }
muyangli's avatar
muyangli committed
339
340
341
342
        }
    }

    // get {k}-th and {k+1}-th wscale from the block, k must be multiples of 2, k must be uniform across all lanes
Muyang Li's avatar
Muyang Li committed
343
344
345
    __device__ __forceinline__ static half2_t broadcast_wscale(wscale_warp block, int k, int laneId) {
        const int packIdx    = k / (WSCALES_PACK_SIZE * WARP_SIZE);
        const int srcLane    = 4 * (k / WSCALES_PACK_SIZE) + laneId % 4;
muyangli's avatar
muyangli committed
346
347
348
349
        const int elementIdx = k % WSCALES_PACK_SIZE / 2;
        return __shfl_sync(~0, block[packIdx].data[elementIdx], srcLane);
    }
    // get {k}-th and {k+1}-th ascale from the block, k must be multiples of 2, k must be uniform across all lanes
Muyang Li's avatar
Muyang Li committed
350
351
352
    __device__ __forceinline__ static half2_t broadcast_ascale(ascale_warp block, int k, int laneId) {
        const int packIdx    = k / (ASCALES_PACK_SIZE * WARP_SIZE);
        const int srcLane    = 8 * (k / ASCALES_PACK_SIZE) + laneId / 4;
muyangli's avatar
muyangli committed
353
354
355
356
        const int elementIdx = k % ASCALES_PACK_SIZE / 2;
        return __shfl_sync(~0, block[packIdx].data[elementIdx], srcLane);
    }

357
    struct i2f_normal {
Muyang Li's avatar
Muyang Li committed
358
        __device__ __forceinline__ static float2 int2float2(int x, int y) {
359
360
361
            return make_float2(__int2float_rn(x), __int2float_rn(y));
        }

Muyang Li's avatar
Muyang Li committed
362
        __device__ __forceinline__ static half2_t int2half2(int x, int y) {
363
364
365
366
367
            return float22half2<half2_t>(int2float2(x, y));
        }
    };

    template<typename i2f = i2f_normal, typename F>
Muyang Li's avatar
Muyang Li committed
368
369
    __device__ __forceinline__ static void
    apply_scales(F &&getpsum, ascale_warp ascale, wscale_warp wscale, fpsum_warp &fpsum) {
muyangli's avatar
muyangli committed
370
371
372
373
374
375
376
377
        const int laneId = threadIdx.x % WARP_SIZE;
        const int warpId = threadIdx.x / WARP_SIZE;

        half2_t asx[WARP_M_TILES];
        half2_t asy[WARP_M_TILES];

        for (int i = 0; i < WARP_M_TILES; i++) {
            half2_t as = broadcast_ascale(ascale, i * 2, laneId);
Muyang Li's avatar
Muyang Li committed
378
379
            asx[i]     = half2_t(as.x, as.x);
            asy[i]     = half2_t(as.y, as.y);
muyangli's avatar
muyangli committed
380
381
382
383
384
385
386
387
388
389
390
391
392
        }

        for (int j = 0; j < WARP_N_TILES; j++) {
            half2_t ws1 = broadcast_wscale(wscale, j * 4, laneId);
            half2_t ws2 = broadcast_wscale(wscale, j * 4 + 2, laneId);

            for (int i = 0; i < WARP_M_TILES; i++) {
                auto &fsum = fpsum[i * WARP_N_TILES + j];

                packed_psum_t psum = getpsum(i, j);

                // constexpr int target = 0;
                // if (threadIdx.x == 3 && j == 1 && i == 0) {
Muyang Li's avatar
Muyang Li committed
393
394
395

                //     printf("before ws2 = %f %f fsum.data[%d] = %f %f\n", (float)ws2.x, (float)ws2.y, target,
                //     (float)fsum.data[target].x, (float)fsum.data[target].y);
muyangli's avatar
muyangli committed
396
397
                // }

398
399
400
401
402
                fsum.data[0] = __hfma2(i2f::int2half2(psum.data[0], psum.data[1]), __hmul2(asx[i], ws1), fsum.data[0]);
                fsum.data[1] = __hfma2(i2f::int2half2(psum.data[2], psum.data[3]), __hmul2(asy[i], ws1), fsum.data[1]);
                fsum.data[2] = __hfma2(i2f::int2half2(psum.data[4], psum.data[5]), __hmul2(asx[i], ws2), fsum.data[2]);
                fsum.data[3] = __hfma2(i2f::int2half2(psum.data[6], psum.data[7]), __hmul2(asy[i], ws2), fsum.data[3]);

muyangli's avatar
muyangli committed
403
                // if (threadIdx.x == 3 && j == 1 && i == 0) {
Muyang Li's avatar
Muyang Li committed
404
405
                //     printf("before ws2 = %f %f fsum.data[%d] = %f %f\n", (float)ws2.x, (float)ws2.y, target,
                //     (float)fsum.data[target].x, (float)fsum.data[target].y);
muyangli's avatar
muyangli committed
406
407
408
409
410
                // }
            }
        }
    }

411
    template<typename i2f = i2f_normal, typename F>
Muyang Li's avatar
Muyang Li committed
412
413
    __device__ __forceinline__ static void
    apply_scales(F &&getpsum, ascale_warp ascale, wscale_warp wscale, f32psum_warp &fpsum) {
muyangli's avatar
muyangli committed
414
415
416
417
418
419
420
421
        const int laneId = threadIdx.x % WARP_SIZE;
        const int warpId = threadIdx.x / WARP_SIZE;

        float2 asx[WARP_M_TILES];
        float2 asy[WARP_M_TILES];

        for (int i = 0; i < WARP_M_TILES; i++) {
            half2_t as = broadcast_ascale(ascale, i * 2, laneId);
Muyang Li's avatar
Muyang Li committed
422
423
            asx[i]     = half22float2(half2_t(as.x, as.x));
            asy[i]     = half22float2(half2_t(as.y, as.y));
muyangli's avatar
muyangli committed
424
425
426
427
428
429
430
431
432
433
434
435
436
437
438
439
        }

        auto fma2 = [](float2 a, float2 b, float &cx, float &cy) ALWAYSINLINE {
            cx += a.x * b.x;
            cy += a.y * b.y;
        };

        for (int j = 0; j < WARP_N_TILES; j++) {
            float2 ws1 = half22float2(broadcast_wscale(wscale, j * 4, laneId));
            float2 ws2 = half22float2(broadcast_wscale(wscale, j * 4 + 2, laneId));

            for (int i = 0; i < WARP_M_TILES; i++) {
                auto &fsum = fpsum[i * WARP_N_TILES + j];

                packed_psum_t psum = getpsum(i, j);

440
441
442
443
                fma2(i2f::int2float2(psum.data[0], psum.data[1]), asx[i] * ws1, fsum.data[0], fsum.data[1]);
                fma2(i2f::int2float2(psum.data[2], psum.data[3]), asy[i] * ws1, fsum.data[2], fsum.data[3]);
                fma2(i2f::int2float2(psum.data[4], psum.data[5]), asx[i] * ws2, fsum.data[4], fsum.data[5]);
                fma2(i2f::int2float2(psum.data[6], psum.data[7]), asy[i] * ws2, fsum.data[6], fsum.data[7]);
muyangli's avatar
muyangli committed
444
445
446
447
448
449
450
451
            }
        }
    }

    /**
     * input: WARP_M of half (in shared memory, per-warp)
     * output: [..., ASCALES_NUM_PACKS, ASCALES_VALID_LANES] in global memory, per-warp
     */
Muyang Li's avatar
Muyang Li committed
452
    __device__ __forceinline__ static void pack_ascales(const half_t *input, packed_ascale_t *output) {
muyangli's avatar
muyangli committed
453
        const int laneId = threadIdx.x % WARP_SIZE;
Muyang Li's avatar
Muyang Li committed
454
#pragma unroll
muyangli's avatar
muyangli committed
455
456
457
        for (int j = 0; j < ASCALES_NUM_PACKS; j++) {
            if (laneId < ASCALES_VALID_LANES) {
                packed_ascale_t tmp;
Muyang Li's avatar
Muyang Li committed
458
#pragma unroll
muyangli's avatar
muyangli committed
459
                for (int i = 0; i < ASCALES_PACK_SIZE; i += 2) {
Muyang Li's avatar
Muyang Li committed
460
461
462
463
                    tmp.data[i / 2].x = input[j * ASCALES_PACK_SIZE * WARP_SIZE + laneId / 8 * 8 * ASCALES_PACK_SIZE +
                                              laneId % 8 + i * 8];
                    tmp.data[i / 2].y = input[j * ASCALES_PACK_SIZE * WARP_SIZE + laneId / 8 * 8 * ASCALES_PACK_SIZE +
                                              laneId % 8 + (i + 1) * 8];
muyangli's avatar
muyangli committed
464
465
466
467
468
469
470
471
472
473
                }
                output[j * ASCALES_VALID_LANES + laneId] = tmp;
            }
        }
    }

    /**
     * input: WARP_N of half (in shared memory, per-warp)
     * output: [..., WSCALES_NUM_PACKS, WSCALES_VALID_LANES] in global memory, per-warp
     */
Muyang Li's avatar
Muyang Li committed
474
    __device__ __forceinline__ static void pack_wscales(const half_t *input, packed_wscale_t *output) {
muyangli's avatar
muyangli committed
475
        const int laneId = threadIdx.x % WARP_SIZE;
Muyang Li's avatar
Muyang Li committed
476
#pragma unroll
muyangli's avatar
muyangli committed
477
478
479
        for (int j = 0; j < WSCALES_NUM_PACKS; j++) {
            if (laneId < WSCALES_VALID_LANES) {
                packed_wscale_t tmp;
Muyang Li's avatar
Muyang Li committed
480
#pragma unroll
muyangli's avatar
muyangli committed
481
                for (int i = 0; i < WSCALES_PACK_SIZE; i += 2) {
Muyang Li's avatar
Muyang Li committed
482
483
484
                    tmp.data[i / 2] = *reinterpret_cast<const half2_t *>(
                        &input[j * WSCALES_PACK_SIZE * WARP_SIZE + laneId / 4 * 4 * WSCALES_PACK_SIZE + laneId % 4 * 2 +
                               i * 4]);
muyangli's avatar
muyangli committed
485
486
487
488
489
490
491
492
493
494
495
                }
                store(&output[j * WSCALES_VALID_LANES + laneId], tmp);
            }
        }
    }

    struct unpack_fpsum {
        // +8 to prevent bank conflicts
        using matrix_t = half_t[8][WARP_N + 8];

        static constexpr int SHMEM_SIZE = sizeof(matrix_t);
Muyang Li's avatar
Muyang Li committed
496
497
        static constexpr int PACK_SIZE  = WARP_N / WARP_SIZE;
        using pack_t                    = std::array<half_t, PACK_SIZE>;
muyangli's avatar
muyangli committed
498
499

        // F (int rowId, pack_t &pack)
Muyang Li's avatar
Muyang Li committed
500
501
502
        template<typename... F>
        __device__ __forceinline__ void operator()(
            fpsum_warp fpsum, half_t *output, int stride, int maxRows, int maxCols, void *shmem, F &&...plugins) {
muyangli's avatar
muyangli committed
503
504
            const int laneId = threadIdx.x % WARP_SIZE;

Muyang Li's avatar
Muyang Li committed
505
            matrix_t &mat = *reinterpret_cast<matrix_t *>(shmem);
muyangli's avatar
muyangli committed
506
507
508
509
510
511
512
513
514
515
516
517
518
519
520
521

            // pack_t reduce_tmp;
            // constexpr bool enableReduce = !std::is_void_v<FuncReduce>;

            // if constexpr (enableReduce) {
            //     reduce_tmp.fill(reduce_initval);
            //     // reduce_tmp = load<true>(reinterpret_cast<pack_t *>(&reduce_result[laneId * PACK_SIZE]));
            // }
            // auto doReduce = [&reduce_tmp](pack_t pack) {
            //     if constexpr (enableReduce) {
            //         for (int i = 0; i < PACK_SIZE; i++) {
            //             reduce_tmp[i] = FuncReduce()(reduce_tmp[i], pack[i]);
            //         }
            //     }
            // };

Muyang Li's avatar
Muyang Li committed
522
#pragma unroll
muyangli's avatar
muyangli committed
523
            for (int i = 0; i < WARP_M_TILES; i++) {
Muyang Li's avatar
Muyang Li committed
524
#pragma unroll
muyangli's avatar
muyangli committed
525
                for (int j = 0; j < WARP_N_TILES; j++) {
Muyang Li's avatar
Muyang Li committed
526
527
528
                    packed_fpsum_t &fsum                             = fpsum[i * WARP_N_TILES + j];
                    int row                                          = laneId / 4;
                    int col                                          = laneId % 4 * 2 + j * INSN_N;
muyangli's avatar
muyangli committed
529
530
531
                    *reinterpret_cast<half2_t *>(&mat[row][col + 0]) = fsum.data[0];
                    *reinterpret_cast<half2_t *>(&mat[row][col + 8]) = fsum.data[2];
                }
limm's avatar
limm committed
532
533
                // __syncwarp();
		__builtin_amdgcn_wave_barrier();
muyangli's avatar
muyangli committed
534

Muyang Li's avatar
Muyang Li committed
535
#pragma unroll
muyangli's avatar
muyangli committed
536
537
538
539
540
541
542
543
544
545
                for (int row = 0; row < 8; row++) {
                    pack_t pack = *reinterpret_cast<pack_t *>(&mat[row][laneId * PACK_SIZE]);

                    // if constexpr (enableReduce) {
                    //     doReduce(pack);
                    // }

                    (plugins(i * INSN_M + row, pack), ...);

                    bool pred = i * INSN_M + row < maxRows && laneId * PACK_SIZE < maxCols;
546
                    // if (pred) {
Muyang Li's avatar
Muyang Li committed
547
548
549
                    store_pred(reinterpret_cast<pack_t *>(&output[(i * INSN_M + row) * stride + laneId * PACK_SIZE]),
                               pack,
                               pred);
550
                    // }
muyangli's avatar
muyangli committed
551
                }
limm's avatar
limm committed
552
553
                // __syncwarp();
		__builtin_amdgcn_wave_barrier();
muyangli's avatar
muyangli committed
554

Muyang Li's avatar
Muyang Li committed
555
#pragma unroll
muyangli's avatar
muyangli committed
556
                for (int j = 0; j < WARP_N_TILES; j++) {
Muyang Li's avatar
Muyang Li committed
557
558
559
                    packed_fpsum_t &fsum                             = fpsum[i * WARP_N_TILES + j];
                    int row                                          = laneId / 4;
                    int col                                          = laneId % 4 * 2 + j * INSN_N;
muyangli's avatar
muyangli committed
560
561
562
                    *reinterpret_cast<half2_t *>(&mat[row][col + 0]) = fsum.data[1];
                    *reinterpret_cast<half2_t *>(&mat[row][col + 8]) = fsum.data[3];
                }
limm's avatar
limm committed
563
564
                // __syncwarp();
		__builtin_amdgcn_wave_barrier();
muyangli's avatar
muyangli committed
565

Muyang Li's avatar
Muyang Li committed
566
#pragma unroll
muyangli's avatar
muyangli committed
567
568
569
570
571
572
573
574
575
576
                for (int row = 0; row < 8; row++) {
                    pack_t pack = *reinterpret_cast<pack_t *>(&mat[row][laneId * PACK_SIZE]);

                    // if constexpr (enableReduce) {
                    //     doReduce(pack);
                    // }

                    (plugins(i * INSN_M + 8 + row, pack), ...);

                    bool pred = i * INSN_M + 8 + row < maxRows && laneId * PACK_SIZE < maxCols;
577
                    // if (pred) {
Muyang Li's avatar
Muyang Li committed
578
579
580
581
                    store_pred(
                        reinterpret_cast<pack_t *>(&output[(i * INSN_M + 8 + row) * stride + laneId * PACK_SIZE]),
                        pack,
                        pred);
582
                    // }
muyangli's avatar
muyangli committed
583
                }
limm's avatar
limm committed
584
585
                // __syncwarp();
		__builtin_amdgcn_wave_barrier();
muyangli's avatar
muyangli committed
586
587
588
589
590
591
592
            }
            // if (enableReduce) {
            //     store<true>(reinterpret_cast<pack_t *>(&reduce_result[laneId * PACK_SIZE]), reduce_tmp);
            // }
        }
    };

sxtyzhangzk's avatar
sxtyzhangzk committed
593
594
595
596
    // loads act of [WARP_M, WARP_N] and stores to fpsum_warp
    // [WARP_M, WARP_N * 2] when fuse_glu
    template<bool fuse_glu>
    struct load_act_to_fpsum {
Muyang Li's avatar
Muyang Li committed
597
        using matrix_t                     = half_t[INSN_M][WARP_N + 8];
sxtyzhangzk's avatar
sxtyzhangzk committed
598
599
        static constexpr size_t SHMEM_SIZE = sizeof(matrix_t);

Muyang Li's avatar
Muyang Li committed
600
601
        __device__ __forceinline__ void
        operator()(const half_t *input, int stride, int maxRows, int maxCols, fpsum_warp &out, void *shmem) {
sxtyzhangzk's avatar
sxtyzhangzk committed
602
603
604
605
606
            const int laneId = threadIdx.x % WARP_SIZE;

            matrix_t &mat = *reinterpret_cast<matrix_t *>(shmem);

            constexpr int PACK_SIZE = WARP_N / WARP_SIZE;
Muyang Li's avatar
Muyang Li committed
607
608
            using packed_input      = std::array<half_t, PACK_SIZE>;
            using packed_raw_input  = std::array<half2_t, PACK_SIZE>;
sxtyzhangzk's avatar
sxtyzhangzk committed
609

Muyang Li's avatar
Muyang Li committed
610
#pragma unroll
sxtyzhangzk's avatar
sxtyzhangzk committed
611
            for (int m = 0; m < WARP_M_TILES; m++) {
Muyang Li's avatar
Muyang Li committed
612
#pragma unroll
sxtyzhangzk's avatar
sxtyzhangzk committed
613
614
615
616
617
618
619
620
                for (int row = 0; row < INSN_M; row++) {
                    packed_input pack;
                    // TODO: numCols not multiples of PACK_SIZE
                    if constexpr (fuse_glu) {
                        packed_raw_input raw;
                        raw.fill(half2_t(0, 0));
                        bool pred = (m * INSN_M + row) < maxRows && laneId * PACK_SIZE * 2 < maxCols;
                        if (pred) {
Muyang Li's avatar
Muyang Li committed
621
622
                            raw = load(reinterpret_cast<const packed_raw_input *>(input + (m * INSN_M + row) * stride +
                                                                                  laneId * PACK_SIZE * 2));
sxtyzhangzk's avatar
sxtyzhangzk committed
623
                        }
Muyang Li's avatar
Muyang Li committed
624
#pragma unroll
sxtyzhangzk's avatar
sxtyzhangzk committed
625
626
627
628
629
630
631
                        for (int j = 0; j < PACK_SIZE; j++) {
                            pack[j] = raw[j].x * silu(raw[j].y);
                        }
                    } else {
                        pack.fill(half_t(0));
                        bool pred = (m * INSN_M + row) < maxRows && laneId * PACK_SIZE < maxCols;
                        if (pred) {
Muyang Li's avatar
Muyang Li committed
632
633
                            pack = load(reinterpret_cast<const packed_input *>(input + (m * INSN_M + row) * stride +
                                                                               laneId * PACK_SIZE));
sxtyzhangzk's avatar
sxtyzhangzk committed
634
635
636
637
                        }
                    }
                    store<true>(reinterpret_cast<packed_input *>(&mat[row][laneId * PACK_SIZE]), pack);
                }
limm's avatar
limm committed
638
639
                // __syncwarp();
		__builtin_amdgcn_wave_barrier();
sxtyzhangzk's avatar
sxtyzhangzk committed
640
641
642
643
644
645
646
647

                for (int n = 0; n < WARP_N_TILES; n++) {
                    const int row = laneId % 16;
                    const int col = n * INSN_N + laneId / 16 * 8;
                    uint4 tmp;
                    ldmatrix(&mat[row][col], tmp);
                    *reinterpret_cast<uint4 *>(&out[m * WARP_N_TILES + n]) = tmp;
                }
limm's avatar
limm committed
648
649
                // __syncwarp();
		__builtin_amdgcn_wave_barrier();
sxtyzhangzk's avatar
sxtyzhangzk committed
650
651
652
653
            }
        }
    };

muyangli's avatar
muyangli committed
654
    template<typename F>
Muyang Li's avatar
Muyang Li committed
655
    __device__ __forceinline__ static fpsum_warp apply_act(fpsum_warp fpsum, F func) {
muyangli's avatar
muyangli committed
656
        fpsum_warp result;
Muyang Li's avatar
Muyang Li committed
657
#pragma unroll
muyangli's avatar
muyangli committed
658
        for (int i = 0; i < WARP_M_TILES; i++) {
Muyang Li's avatar
Muyang Li committed
659
#pragma unroll
muyangli's avatar
muyangli committed
660
            for (int j = 0; j < WARP_N_TILES; j++) {
Muyang Li's avatar
Muyang Li committed
661
#pragma unroll
muyangli's avatar
muyangli committed
662
663
                for (int k = 0; k < 4; k++) {
                    half2_t &dst = result[i * WARP_N_TILES + j].data[k];
Muyang Li's avatar
Muyang Li committed
664
665
666
                    half2_t src  = fpsum[i * WARP_N_TILES + j].data[k];
                    dst.x        = func(src.x);
                    dst.y        = func(src.y);
muyangli's avatar
muyangli committed
667
668
669
670
671
672
673
674
675
676
677
678
                }
            }
        }
        return result;
    }

    struct EpilogueDefault {
        struct Arguments {
            half_t *out;
            int actualM, actualN;
        };

Muyang Li's avatar
Muyang Li committed
679
680
        __device__ __forceinline__ void
        operator()(const BlockInfo binfo, fpsum_warp fpsum, int M, int N, int K, const Arguments &args) {
muyangli's avatar
muyangli committed
681
            const int warpId = threadIdx.x / WARP_SIZE;
Muyang Li's avatar
Muyang Li committed
682

limm's avatar
limm committed
683
684
            // __shared__ alignas(128) uint8_t shmem[NUM_WARPS][ceilDiv(unpack_fpsum::SHMEM_SIZE, 128) * 128];
	    __shared__ uint8_t shmem[NUM_WARPS][ceilDiv(unpack_fpsum::SHMEM_SIZE, 128) * 128] __attribute__((aligned(128)));
muyangli's avatar
muyangli committed
685
686
687
688

            const int m_offset = binfo.bm * BLOCK_M + warpId * WARP_M;
            const int n_offset = binfo.bn * BLOCK_N;

Muyang Li's avatar
Muyang Li committed
689
690
691
692
693
694
695
696
697
698
699
700
701
702
703
            unpack_fpsum()(fpsum,
                           args.out + m_offset * args.actualN + n_offset,
                           args.actualN,
                           args.actualM - m_offset,
                           args.actualN - n_offset,
                           shmem[warpId],
                           [&](int rowId, unpack_fpsum::pack_t &pack) ALWAYSINLINE {
                               if constexpr (std::is_same_v<half_t, half>) {
#pragma unroll
                                   for (int i = 0; i < pack.size(); i++) {
                                       pack[i] = __hmin(pack[i], (half)65504);
                                       pack[i] = __hmax(pack[i], (half)-65504);
                                   }
                               }
                           });
muyangli's avatar
muyangli committed
704
705
706
707
708
        }
    };

    struct EpilogueNop {
        // workaround for layout mismatch between host and device code
Muyang Li's avatar
Muyang Li committed
709
710
711
        struct Arguments {
            size_t unused;
        };
muyangli's avatar
muyangli committed
712

Muyang Li's avatar
Muyang Li committed
713
714
        __device__ __forceinline__ void
        operator()(const BlockInfo binfo, fpsum_warp fpsum, int M, int N, int K, const Arguments &args) {}
muyangli's avatar
muyangli committed
715
716
    };

717
    template<bool USE_BIAS = true, bool USE_SCALE = false>
muyangli's avatar
muyangli committed
718
719
    struct EpilogueBias {
        struct Arguments {
Muyang Li's avatar
Muyang Li committed
720
            const packed_wscale_t *bias; // [N / BLOCK_N, WSCALES_NUM_PACKS, WSCALES_VALID_LANES] of packed_wscale_t
721
            const packed_wscale_t *scale;
muyangli's avatar
muyangli committed
722
723
        };

Muyang Li's avatar
Muyang Li committed
724
725
        __device__ __forceinline__ void
        apply_bias(fpsum_warp &fpsum, int M, int N, int K, const packed_wscale_t *bias, const packed_wscale_t *scale) {
muyangli's avatar
muyangli committed
726
727
728
            const int laneId = threadIdx.x % WARP_SIZE;

            // if (laneId == 0) {
Muyang Li's avatar
Muyang Li committed
729
730
            //     printf("block.x=%d block.y=%d warpId=%d bias=%p\n", blockIdx.x, blockIdx.y, threadIdx.x / WARP_SIZE,
            //     bias);
muyangli's avatar
muyangli committed
731
732
            // }

733
734
735
736
737
738
739
            wscale_warp b, s;
            if constexpr (USE_BIAS) {
                load_wscale(bias, 0, N, b, true);
            }
            if constexpr (USE_SCALE) {
                load_wscale(scale, 0, N, s, true);
            }
muyangli's avatar
muyangli committed
740
741

            for (int j = 0; j < WARP_N_TILES; j++) {
742
743
744
745
746
747
748
749
750
751
                half2_t b1, b2;
                half2_t s1, s2;
                if constexpr (USE_BIAS) {
                    b1 = broadcast_wscale(b, j * 4, laneId);
                    b2 = broadcast_wscale(b, j * 4 + 2, laneId);
                }
                if constexpr (USE_SCALE) {
                    s1 = broadcast_wscale(s, j * 4, laneId);
                    s2 = broadcast_wscale(s, j * 4 + 2, laneId);
                }
muyangli's avatar
muyangli committed
752
753
754
755

                for (int i = 0; i < WARP_M_TILES; i++) {
                    auto &fsum = fpsum[i * WARP_N_TILES + j];

756
757
758
759
760
761
762
763
764
765
766
767
768
769
770
771
                    if constexpr (USE_SCALE && USE_BIAS) {
                        fsum.data[0] = __hfma2(fsum.data[0], s1, b1);
                        fsum.data[1] = __hfma2(fsum.data[1], s1, b1);
                        fsum.data[2] = __hfma2(fsum.data[2], s2, b2);
                        fsum.data[3] = __hfma2(fsum.data[3], s2, b2);
                    } else if constexpr (USE_SCALE) {
                        fsum.data[0] = __hmul2(fsum.data[0], s1);
                        fsum.data[1] = __hmul2(fsum.data[1], s1);
                        fsum.data[2] = __hmul2(fsum.data[2], s2);
                        fsum.data[3] = __hmul2(fsum.data[3], s2);
                    } else if constexpr (USE_BIAS) {
                        fsum.data[0] = __hadd2(fsum.data[0], b1);
                        fsum.data[1] = __hadd2(fsum.data[1], b1);
                        fsum.data[2] = __hadd2(fsum.data[2], b2);
                        fsum.data[3] = __hadd2(fsum.data[3], b2);
                    }
muyangli's avatar
muyangli committed
772
773
774
775
                }
            }
        }

Muyang Li's avatar
Muyang Li committed
776
777
        __device__ __forceinline__ void
        operator()(const BlockInfo binfo, fpsum_warp &fpsum, int M, int N, int K, const Arguments &args) {
muyangli's avatar
muyangli committed
778
            const int bn = binfo.bn;
779
            if constexpr (USE_BIAS || USE_SCALE) {
Muyang Li's avatar
Muyang Li committed
780
781
782
783
784
785
                apply_bias(fpsum,
                           M,
                           N,
                           K,
                           args.bias + bn * WSCALES_NUM_PACKS * WSCALES_VALID_LANES,
                           args.scale + bn * WSCALES_NUM_PACKS * WSCALES_VALID_LANES);
786
            }
muyangli's avatar
muyangli committed
787
788
789
790
        }
    };

    struct EpilogueSilu {
Muyang Li's avatar
Muyang Li committed
791
792
793
        struct Arguments {
            size_t unused;
        };
muyangli's avatar
muyangli committed
794

Muyang Li's avatar
Muyang Li committed
795
796
        __device__ __forceinline__ void
        operator()(const BlockInfo binfo, fpsum_warp &fpsum, int M, int N, int K, const Arguments &args) {
muyangli's avatar
muyangli committed
797
798
799
800
            fpsum = apply_act(fpsum, [](half_t x) { return silu(x); });
        }
    };

Muyang Li's avatar
Muyang Li committed
801
    template<typename... Epilogues>
muyangli's avatar
muyangli committed
802
803
804
    struct EpilogueCombination {
        using Arguments = std::tuple<typename Epilogues::Arguments...>;

Muyang Li's avatar
Muyang Li committed
805
806
        __device__ __forceinline__ void
        operator()(const BlockInfo binfo, fpsum_warp &fpsum, int M, int N, int K, const Arguments &args) {
muyangli's avatar
muyangli committed
807
            // this function makes intellisense crashes :(
Muyang Li's avatar
Muyang Li committed
808
809
810
#if __INTELLISENSE__
            __trap(); // should not happen when actually compiling
#else
muyangli's avatar
muyangli committed
811
812
813
814
            std::tuple<Epilogues...> epilogues;
            auto run = [&]<size_t idx>() {
                std::get<idx>(epilogues).operator()(binfo, fpsum, M, N, K, std::get<idx>(args));
            };
Muyang Li's avatar
Muyang Li committed
815
816
817
818
            auto foreach = [&]<size_t... Is>(std::index_sequence<Is...>) { (run.template operator()<Is>(), ...); };
            foreach (std::make_index_sequence<sizeof...(Epilogues)>())
                ;
#endif
muyangli's avatar
muyangli committed
819
820
821
822
        }
    };
};

Muyang Li's avatar
Muyang Li committed
823
824
825
826
827
828
829
830
831
832
833
834
835
836
837
838
839
840
841
842
843
844
845
846
847
848
849
850
851
852
853
854
855
856
857
858
859
860
861
862
863
864
865
#define IMPORT_GEMM_BASE(config)                                                                                       \
    using Base = GEMMBase<config>;                                                                                     \
    using Base::BLOCK_M;                                                                                               \
    using Base::BLOCK_N;                                                                                               \
    using Base::WARP_SIZE;                                                                                             \
    using Base::NUM_WARPS;                                                                                             \
    using Base::INSN_M;                                                                                                \
    using Base::INSN_N;                                                                                                \
    using Base::INSN_K;                                                                                                \
    using typename Base::half_t;                                                                                       \
    using typename Base::half2_t;                                                                                      \
    using Base::WARP_M;                                                                                                \
    using Base::WARP_N;                                                                                                \
    using Base::WARP_K;                                                                                                \
    using Base::WARP_M_TILES;                                                                                          \
    using Base::WARP_N_TILES;                                                                                          \
    using Base::WARP_K_TILES;                                                                                          \
    using Base::WSCALES_PACK_SIZE;                                                                                     \
    using Base::WSCALES_NUM_PACKS;                                                                                     \
    using Base::WSCALES_VALID_LANES;                                                                                   \
    using Base::ASCALES_PACK_SIZE;                                                                                     \
    using Base::ASCALES_NUM_PACKS;                                                                                     \
    using Base::ASCALES_VALID_LANES;                                                                                   \
    using typename Base::packed_act_t;                                                                                 \
    using typename Base::packed_wgt_t;                                                                                 \
    using typename Base::packed_psum_t;                                                                                \
    using typename Base::packed_fpsum_t;                                                                               \
    using typename Base::packed_gated_fpsum_t;                                                                         \
    using typename Base::packed_f32psum_t;                                                                             \
    using typename Base::packed_wscale_t;                                                                              \
    using typename Base::packed_ascale_t;                                                                              \
    using typename Base::act_warp;                                                                                     \
    using typename Base::wgt_warp;                                                                                     \
    using typename Base::ascale_warp;                                                                                  \
    using typename Base::wscale_warp;                                                                                  \
    using typename Base::fpsum_warp;                                                                                   \
    using typename Base::f32psum_warp;                                                                                 \
    using typename Base::gated_fpsum_warp;                                                                             \
    using typename Base::BlockInfo;                                                                                    \
    using typename Base::unpack_fpsum;                                                                                 \
    using typename Base::EpilogueDefault;                                                                              \
    using typename Base::EpilogueNop;                                                                                  \
    template<bool USE_BIAS, bool USE_SCALE>                                                                            \
limm's avatar
limm committed
866
    using EpilogueBias = typename Base::template EpilogueBias<USE_BIAS, USE_SCALE>;                                             \
Muyang Li's avatar
Muyang Li committed
867
868
869
870
871
872
873
874
875
876
877
878
    using Base::mma_f16xf16_f32;                                                                                       \
    using Base::packed_fp32_to_fp16;                                                                                   \
    using Base::packed_fp16_to_fp32;                                                                                   \
    using Base::load_act;                                                                                              \
    using Base::load_wgt;                                                                                              \
    using Base::load_ascale;                                                                                           \
    using Base::load_wscale;                                                                                           \
    using Base::broadcast_wscale;                                                                                      \
    using Base::broadcast_ascale;                                                                                      \
    using Base::apply_scales;                                                                                          \
    using Base::pack_ascales;                                                                                          \
    using Base::pack_wscales;                                                                                          \
sxtyzhangzk's avatar
sxtyzhangzk committed
879
    using Base::apply_act;
muyangli's avatar
muyangli committed
880

881
882
template<typename kernel>
constexpr int min_arch() {
Muyang Li's avatar
Muyang Li committed
883
    if constexpr (requires { kernel::MIN_ARCH; }) {
884
885
886
887
888
889
890
        return kernel::MIN_ARCH;
    } else {
        return 0;
    }
}
template<typename kernel>
constexpr int max_arch() {
Muyang Li's avatar
Muyang Li committed
891
    if constexpr (requires { kernel::MAX_ARCH; }) {
892
893
894
895
896
        return kernel::MAX_ARCH;
    } else {
        return INT_MAX;
    }
}
muyangli's avatar
muyangli committed
897

Muyang Li's avatar
Muyang Li committed
898
899
template<typename kernel, typename... T>
__global__ static void invoke_kernel(T... args) {
fengzch-das's avatar
fengzch-das committed
900
901
#ifdef __CUDA_ARCH__
    if constexpr (__CUDA_ARCH__ >= min_arch<kernel>() && __CUDA_ARCH__ <= max_arch<kernel>()) {
902
903
904
905
906
907
        kernel()(args...);
    } else {
        trap_unsupported_arch();
    }
#else
    // ???
muyangli's avatar
muyangli committed
908
    kernel()(args...);
909
#endif
muyangli's avatar
muyangli committed
910
911
912
}

template<typename T>
Muyang Li's avatar
Muyang Li committed
913
__global__ static void test_sizeof_device() {
muyangli's avatar
muyangli committed
914
915
916
917
918
919
920
921
922
923
924
925
    printf("sizeof on device = %d\n", (int)sizeof(T));
}

template<typename T>
static void test_sizeof_host() {
    printf("sizeof on host = %d\n", (int)sizeof(T));
}

template<typename T>
static void test_sizeof() {
    printf("typeid = %s\n", typeid(T).name());
    test_sizeof_host<T>();
fengzch-das's avatar
fengzch-das committed
926
927
    test_sizeof_device<T><<<1, 1>>>();
    checkCUDA(cudaDeviceSynchronize());
muyangli's avatar
muyangli committed
928
929
}

Muyang Li's avatar
Muyang Li committed
930
}; // namespace nunchaku::kernels