gemm_base.cuh 41.5 KB
Newer Older
fengzch-das's avatar
fengzch-das committed
1
#include "hip/hip_runtime.h"
muyangli's avatar
muyangli committed
2
3
4
5
6
7
8
#pragma once

#include "common.h"

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

#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 {

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

Muyang Li's avatar
Muyang Li committed
35
36
    static constexpr int BLOCK_M   = 256;
    static constexpr int BLOCK_N   = 128;
muyangli's avatar
muyangli committed
37
38
39
40
41
42
43
    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;

44
45
46
47
    // 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
48
49
    using half_t  = typename std::conditional_t<bf16, __hip_bfloat16, half>;
    using half2_t = typename std::conditional_t<bf16, __hip_bfloat162, half2>;
muyangli's avatar
muyangli committed
50
51
};

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

class GEMMConfig_W8A8 {
public:
Muyang Li's avatar
Muyang Li committed
58
59
    static constexpr int BLOCK_M   = 256;
    static constexpr int BLOCK_N   = 128;
muyangli's avatar
muyangli committed
60
61
62
63
64
65
66
67
68
69
70
    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
71
72
    using half_t  = __hip_bfloat16;
    using half2_t = __hip_bfloat162;
muyangli's avatar
muyangli committed
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
99
#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
100
     *
muyangli's avatar
muyangli committed
101
102
103
104
105
     * 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
106
     *
muyangli's avatar
muyangli committed
107
108
109
110
111
112
113
     *  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
114
     *
muyangli's avatar
muyangli committed
115
116
117
118
119
     * 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
120
     *
muyangli's avatar
muyangli committed
121
122
123
     *  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
124
125
126
127
     *
     * {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
128
129
     * 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
130
131
132
     * */
    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
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
    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
149
150
151
     *
     * {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
152
     */
Muyang Li's avatar
Muyang Li committed
153
154
    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
155
156
157
158
    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
159

muyangli's avatar
muyangli committed
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
188
    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
189
190
191
192
193
194
    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
195
196
197
198
199
200
201
202
203
    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
204
205
    __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
206
        static_assert(std::is_same_v<half_t, half> || std::is_same_v<half_t, __hip_bfloat16>);
muyangli's avatar
muyangli committed
207

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

        uint4 out1 = mma_m16n8k16_f32f16f16f32<is_bf16>(
Muyang Li's avatar
Muyang Li committed
211
212
            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
213
            kernels::bit_cast<uint4>(float4(psum.data[0], psum.data[1], psum.data[2], psum.data[3])));
214
        uint4 out2 = mma_m16n8k16_f32f16f16f32<is_bf16>(
Muyang Li's avatar
Muyang Li committed
215
216
            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
217
218
219
220
221
222
223
224
225
            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
226

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

Muyang Li's avatar
Muyang Li committed
230
    __device__ __forceinline__ static packed_fpsum_t packed_fp32_to_fp16(packed_f32psum_t input) {
muyangli's avatar
muyangli committed
231
232
233
234
235
236
237
        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
238
    __device__ __forceinline__ static packed_f32psum_t packed_fp16_to_fp32(packed_fpsum_t input) {
muyangli's avatar
muyangli committed
239
240
        packed_f32psum_t results;
        for (int i = 0; i < 4; i++) {
Muyang Li's avatar
Muyang Li committed
241
242
            float2 tmp              = half22float2(input.data[i]);
            results.data[i * 2]     = tmp.x;
muyangli's avatar
muyangli committed
243
244
245
246
247
            results.data[i * 2 + 1] = tmp.y;
        }
        return results;
    }

Muyang Li's avatar
Muyang Li committed
248
    __device__ __forceinline__ static fpsum_warp packed_fp32_to_fp16(f32psum_warp input) {
muyangli's avatar
muyangli committed
249
        fpsum_warp results;
Muyang Li's avatar
Muyang Li committed
250
#pragma unroll
muyangli's avatar
muyangli committed
251
252
253
254
255
256
        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
257
    __device__ __forceinline__ static f32psum_warp packed_fp16_to_fp32(fpsum_warp input) {
sxtyzhangzk's avatar
sxtyzhangzk committed
258
        f32psum_warp results;
Muyang Li's avatar
Muyang Li committed
259
#pragma unroll
sxtyzhangzk's avatar
sxtyzhangzk committed
260
261
262
263
264
265
        for (int i = 0; i < results.size(); i++) {
            results[i] = packed_fp16_to_fp32(input[i]);
        }
        return results;
    }

muyangli's avatar
muyangli committed
266
    // 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
267
    __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
268
269
        int laneId = threadIdx.x % WARP_SIZE;
        int warpId = threadIdx.x / WARP_SIZE;
Muyang Li's avatar
Muyang Li committed
270
#pragma unroll
muyangli's avatar
muyangli committed
271
        for (int i = 0; i < WARP_M_TILES; i++) {
Muyang Li's avatar
Muyang Li committed
272
273
274
            // 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);
275
            //}
muyangli's avatar
muyangli committed
276
277
278
279
        }
    }

    // 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
280
    __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
281
        int laneId = threadIdx.x % WARP_SIZE;
Muyang Li's avatar
Muyang Li committed
282

muyangli's avatar
muyangli committed
283
284
285
        // 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
286
#pragma unroll
muyangli's avatar
muyangli committed
287
        for (int i = 0; i < WARP_N_TILES; i++) {
Muyang Li's avatar
Muyang Li committed
288
289
290
291
292
            // 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;
293
            //}
muyangli's avatar
muyangli committed
294
295
296
        }
    }

Muyang Li's avatar
Muyang Li committed
297
298
299
300
    // 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
301
302
        int laneId = threadIdx.x % WARP_SIZE;
        int warpId = threadIdx.x / WARP_SIZE;
Muyang Li's avatar
Muyang Li committed
303
#pragma unroll
muyangli's avatar
muyangli committed
304
        for (int i = 0; i < ASCALES_NUM_PACKS; i++) {
305
            // if (pred && laneId < ASCALES_VALID_LANES) {
Muyang Li's avatar
Muyang Li committed
306
307
308
309
310
            // 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
311

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

Muyang Li's avatar
Muyang Li committed
316
317
318
319
    // 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
320
321
322
323
324
325
        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
326
327
        // // 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
328

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

            // 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]);
339
            // }
muyangli's avatar
muyangli committed
340
341
342
343
        }
    }

    // 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
344
345
346
    __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
347
348
349
350
        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
351
352
353
    __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
354
355
356
357
        const int elementIdx = k % ASCALES_PACK_SIZE / 2;
        return __shfl_sync(~0, block[packIdx].data[elementIdx], srcLane);
    }

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

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

    template<typename i2f = i2f_normal, typename F>
Muyang Li's avatar
Muyang Li committed
369
370
    __device__ __forceinline__ static void
    apply_scales(F &&getpsum, ascale_warp ascale, wscale_warp wscale, fpsum_warp &fpsum) {
muyangli's avatar
muyangli committed
371
372
373
374
375
376
377
378
        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
379
380
            asx[i]     = half2_t(as.x, as.x);
            asy[i]     = half2_t(as.y, as.y);
muyangli's avatar
muyangli committed
381
382
383
384
385
386
387
388
389
390
391
392
393
        }

        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
394
395
396

                //     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
397
398
                // }

399
400
401
402
403
                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
404
                // if (threadIdx.x == 3 && j == 1 && i == 0) {
Muyang Li's avatar
Muyang Li committed
405
406
                //     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
407
408
409
410
411
                // }
            }
        }
    }

412
    template<typename i2f = i2f_normal, typename F>
Muyang Li's avatar
Muyang Li committed
413
414
    __device__ __forceinline__ static void
    apply_scales(F &&getpsum, ascale_warp ascale, wscale_warp wscale, f32psum_warp &fpsum) {
muyangli's avatar
muyangli committed
415
416
417
418
419
420
421
422
        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
423
424
            asx[i]     = half22float2(half2_t(as.x, as.x));
            asy[i]     = half22float2(half2_t(as.y, as.y));
muyangli's avatar
muyangli committed
425
426
427
428
429
430
431
432
433
434
435
436
437
438
439
440
        }

        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);

441
442
443
444
                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
445
446
447
448
449
450
451
452
            }
        }
    }

    /**
     * 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
453
    __device__ __forceinline__ static void pack_ascales(const half_t *input, packed_ascale_t *output) {
muyangli's avatar
muyangli committed
454
        const int laneId = threadIdx.x % WARP_SIZE;
Muyang Li's avatar
Muyang Li committed
455
#pragma unroll
muyangli's avatar
muyangli committed
456
457
458
        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
459
#pragma unroll
muyangli's avatar
muyangli committed
460
                for (int i = 0; i < ASCALES_PACK_SIZE; i += 2) {
Muyang Li's avatar
Muyang Li committed
461
462
463
464
                    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
465
466
467
468
469
470
471
472
473
474
                }
                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
475
    __device__ __forceinline__ static void pack_wscales(const half_t *input, packed_wscale_t *output) {
muyangli's avatar
muyangli committed
476
        const int laneId = threadIdx.x % WARP_SIZE;
Muyang Li's avatar
Muyang Li committed
477
#pragma unroll
muyangli's avatar
muyangli committed
478
479
480
        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
481
#pragma unroll
muyangli's avatar
muyangli committed
482
                for (int i = 0; i < WSCALES_PACK_SIZE; i += 2) {
Muyang Li's avatar
Muyang Li committed
483
484
485
                    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
486
487
488
489
490
491
492
493
494
495
496
                }
                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
497
498
        static constexpr int PACK_SIZE  = WARP_N / WARP_SIZE;
        using pack_t                    = std::array<half_t, PACK_SIZE>;
muyangli's avatar
muyangli committed
499
500

        // F (int rowId, pack_t &pack)
Muyang Li's avatar
Muyang Li committed
501
502
503
        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
504
505
            const int laneId = threadIdx.x % WARP_SIZE;

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

            // 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
523
#pragma unroll
muyangli's avatar
muyangli committed
524
            for (int i = 0; i < WARP_M_TILES; i++) {
Muyang Li's avatar
Muyang Li committed
525
#pragma unroll
muyangli's avatar
muyangli committed
526
                for (int j = 0; j < WARP_N_TILES; j++) {
Muyang Li's avatar
Muyang Li committed
527
528
529
                    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
530
531
532
533
534
                    *reinterpret_cast<half2_t *>(&mat[row][col + 0]) = fsum.data[0];
                    *reinterpret_cast<half2_t *>(&mat[row][col + 8]) = fsum.data[2];
                }
                __syncwarp();

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
552
553
                }
                __syncwarp();

Muyang Li's avatar
Muyang Li committed
554
#pragma unroll
muyangli's avatar
muyangli committed
555
                for (int j = 0; j < WARP_N_TILES; j++) {
Muyang Li's avatar
Muyang Li committed
556
557
558
                    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
559
560
561
562
563
                    *reinterpret_cast<half2_t *>(&mat[row][col + 0]) = fsum.data[1];
                    *reinterpret_cast<half2_t *>(&mat[row][col + 8]) = fsum.data[3];
                }
                __syncwarp();

Muyang Li's avatar
Muyang Li committed
564
#pragma unroll
muyangli's avatar
muyangli committed
565
566
567
568
569
570
571
572
573
574
                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;
575
                    // if (pred) {
Muyang Li's avatar
Muyang Li committed
576
577
578
579
                    store_pred(
                        reinterpret_cast<pack_t *>(&output[(i * INSN_M + 8 + row) * stride + laneId * PACK_SIZE]),
                        pack,
                        pred);
580
                    // }
muyangli's avatar
muyangli committed
581
582
583
584
585
586
587
588
589
                }
                __syncwarp();
            }
            // if (enableReduce) {
            //     store<true>(reinterpret_cast<pack_t *>(&reduce_result[laneId * PACK_SIZE]), reduce_tmp);
            // }
        }
    };

sxtyzhangzk's avatar
sxtyzhangzk committed
590
591
592
593
    // 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
594
        using matrix_t                     = half_t[INSN_M][WARP_N + 8];
sxtyzhangzk's avatar
sxtyzhangzk committed
595
596
        static constexpr size_t SHMEM_SIZE = sizeof(matrix_t);

Muyang Li's avatar
Muyang Li committed
597
598
        __device__ __forceinline__ void
        operator()(const half_t *input, int stride, int maxRows, int maxCols, fpsum_warp &out, void *shmem) {
sxtyzhangzk's avatar
sxtyzhangzk committed
599
600
601
602
603
            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
604
605
            using packed_input      = std::array<half_t, PACK_SIZE>;
            using packed_raw_input  = std::array<half2_t, PACK_SIZE>;
sxtyzhangzk's avatar
sxtyzhangzk committed
606

Muyang Li's avatar
Muyang Li committed
607
#pragma unroll
sxtyzhangzk's avatar
sxtyzhangzk committed
608
            for (int m = 0; m < WARP_M_TILES; m++) {
Muyang Li's avatar
Muyang Li committed
609
#pragma unroll
sxtyzhangzk's avatar
sxtyzhangzk committed
610
611
612
613
614
615
616
617
                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
618
619
                            raw = load(reinterpret_cast<const packed_raw_input *>(input + (m * INSN_M + row) * stride +
                                                                                  laneId * PACK_SIZE * 2));
sxtyzhangzk's avatar
sxtyzhangzk committed
620
                        }
Muyang Li's avatar
Muyang Li committed
621
#pragma unroll
sxtyzhangzk's avatar
sxtyzhangzk committed
622
623
624
625
626
627
628
                        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
629
630
                            pack = load(reinterpret_cast<const packed_input *>(input + (m * INSN_M + row) * stride +
                                                                               laneId * PACK_SIZE));
sxtyzhangzk's avatar
sxtyzhangzk committed
631
632
633
634
635
636
637
638
639
640
641
642
643
644
645
646
647
648
                        }
                    }
                    store<true>(reinterpret_cast<packed_input *>(&mat[row][laneId * PACK_SIZE]), pack);
                }
                __syncwarp();

                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;
                }
                __syncwarp();
            }
        }
    };

muyangli's avatar
muyangli committed
649
    template<typename F>
Muyang Li's avatar
Muyang Li committed
650
    __device__ __forceinline__ static fpsum_warp apply_act(fpsum_warp fpsum, F func) {
muyangli's avatar
muyangli committed
651
        fpsum_warp result;
Muyang Li's avatar
Muyang Li committed
652
#pragma unroll
muyangli's avatar
muyangli committed
653
        for (int i = 0; i < WARP_M_TILES; i++) {
Muyang Li's avatar
Muyang Li committed
654
#pragma unroll
muyangli's avatar
muyangli committed
655
            for (int j = 0; j < WARP_N_TILES; j++) {
Muyang Li's avatar
Muyang Li committed
656
#pragma unroll
muyangli's avatar
muyangli committed
657
658
                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
659
660
661
                    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
662
663
664
665
666
667
668
669
670
671
672
673
                }
            }
        }
        return result;
    }

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

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

muyangli's avatar
muyangli committed
678
679
680
681
682
            __shared__ alignas(128) uint8_t shmem[NUM_WARPS][ceilDiv(unpack_fpsum::SHMEM_SIZE, 128) * 128];

            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
683
684
685
686
687
688
689
690
691
692
693
694
695
696
697
            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
698
699
700
701
702
        }
    };

    struct EpilogueNop {
        // workaround for layout mismatch between host and device code
Muyang Li's avatar
Muyang Li committed
703
704
705
        struct Arguments {
            size_t unused;
        };
muyangli's avatar
muyangli committed
706

Muyang Li's avatar
Muyang Li committed
707
708
        __device__ __forceinline__ void
        operator()(const BlockInfo binfo, fpsum_warp fpsum, int M, int N, int K, const Arguments &args) {}
muyangli's avatar
muyangli committed
709
710
    };

711
    template<bool USE_BIAS = true, bool USE_SCALE = false>
muyangli's avatar
muyangli committed
712
713
    struct EpilogueBias {
        struct Arguments {
Muyang Li's avatar
Muyang Li committed
714
            const packed_wscale_t *bias; // [N / BLOCK_N, WSCALES_NUM_PACKS, WSCALES_VALID_LANES] of packed_wscale_t
715
            const packed_wscale_t *scale;
muyangli's avatar
muyangli committed
716
717
        };

Muyang Li's avatar
Muyang Li committed
718
719
        __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
720
721
722
            const int laneId = threadIdx.x % WARP_SIZE;

            // if (laneId == 0) {
Muyang Li's avatar
Muyang Li committed
723
724
            //     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
725
726
            // }

727
728
729
730
731
732
733
            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
734
735

            for (int j = 0; j < WARP_N_TILES; j++) {
736
737
738
739
740
741
742
743
744
745
                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
746
747
748
749

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

750
751
752
753
754
755
756
757
758
759
760
761
762
763
764
765
                    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
766
767
768
769
                }
            }
        }

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

    struct EpilogueSilu {
Muyang Li's avatar
Muyang Li committed
785
786
787
        struct Arguments {
            size_t unused;
        };
muyangli's avatar
muyangli committed
788

Muyang Li's avatar
Muyang Li committed
789
790
        __device__ __forceinline__ void
        operator()(const BlockInfo binfo, fpsum_warp &fpsum, int M, int N, int K, const Arguments &args) {
muyangli's avatar
muyangli committed
791
792
793
794
            fpsum = apply_act(fpsum, [](half_t x) { return silu(x); });
        }
    };

Muyang Li's avatar
Muyang Li committed
795
    template<typename... Epilogues>
muyangli's avatar
muyangli committed
796
797
798
    struct EpilogueCombination {
        using Arguments = std::tuple<typename Epilogues::Arguments...>;

Muyang Li's avatar
Muyang Li committed
799
800
        __device__ __forceinline__ void
        operator()(const BlockInfo binfo, fpsum_warp &fpsum, int M, int N, int K, const Arguments &args) {
muyangli's avatar
muyangli committed
801
            // this function makes intellisense crashes :(
Muyang Li's avatar
Muyang Li committed
802
803
804
#if __INTELLISENSE__
            __trap(); // should not happen when actually compiling
#else
muyangli's avatar
muyangli committed
805
806
807
808
            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
809
810
811
812
            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
813
814
815
816
        }
    };
};

Muyang Li's avatar
Muyang Li committed
817
818
819
820
821
822
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
866
867
868
869
870
871
872
#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>                                                                            \
    using EpilogueBias = typename Base::EpilogueBias<USE_BIAS, USE_SCALE>;                                             \
    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
873
    using Base::apply_act;
muyangli's avatar
muyangli committed
874

875
876
template<typename kernel>
constexpr int min_arch() {
Muyang Li's avatar
Muyang Li committed
877
    if constexpr (requires { kernel::MIN_ARCH; }) {
878
879
880
881
882
883
884
        return kernel::MIN_ARCH;
    } else {
        return 0;
    }
}
template<typename kernel>
constexpr int max_arch() {
Muyang Li's avatar
Muyang Li committed
885
    if constexpr (requires { kernel::MAX_ARCH; }) {
886
887
888
889
890
        return kernel::MAX_ARCH;
    } else {
        return INT_MAX;
    }
}
muyangli's avatar
muyangli committed
891

Muyang Li's avatar
Muyang Li committed
892
893
template<typename kernel, typename... T>
__global__ static void invoke_kernel(T... args) {
fengzch-das's avatar
fengzch-das committed
894
895
#ifdef __DTK_ARCH__
    if constexpr (__DTK_ARCH__ >= min_arch<kernel>() && __DTK_ARCH__ <= max_arch<kernel>()) {
896
897
898
899
900
901
        kernel()(args...);
    } else {
        trap_unsupported_arch();
    }
#else
    // ???
muyangli's avatar
muyangli committed
902
    kernel()(args...);
903
#endif
muyangli's avatar
muyangli committed
904
905
906
}

template<typename T>
Muyang Li's avatar
Muyang Li committed
907
__global__ static void test_sizeof_device() {
muyangli's avatar
muyangli committed
908
909
910
911
912
913
914
915
916
917
918
919
    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
920
921
   hipLaunchKernelGGL(( test_sizeof_device<T>), dim3(1), dim3(1), 0, 0, );
    checkCUDA(hipDeviceSynchronize());
muyangli's avatar
muyangli committed
922
923
}

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