gemm_base.cuh 42.3 KB
Newer Older
muyangli's avatar
muyangli committed
1
#pragma once
fengzch's avatar
fengzch committed
2
#include <hip/amd_detail/amd_hip_bf16.h>
muyangli's avatar
muyangli committed
3
4
5
6
7
8

#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

#pragma nv_diag_suppress 177
fengzch's avatar
fengzch committed
12
#define __DTK_ARCH__ 1200
muyangli's avatar
muyangli committed
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
#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, __nv_bfloat16, half>;
    using half2_t = typename std::conditional_t<bf16, __nv_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  = __nv_bfloat16;
    using half2_t = __nv_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, __nv_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, __nv_bfloat16>;
209
210

        uint4 out1 = mma_m16n8k16_f32f16f16f32<is_bf16>(
Muyang Li's avatar
Muyang Li committed
211
            kernels::bit_cast<uint4>(a),
fengzch's avatar
fengzch committed
212
            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
            kernels::bit_cast<uint4>(a),
fengzch's avatar
fengzch committed
216
            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
        const int elementIdx = k % WSCALES_PACK_SIZE / 2;
fengzch's avatar
fengzch committed
348
349
350
351
352
353
354
355
        half2 temp;
        temp.x = float(block[packIdx].data[elementIdx].x);
        temp.y = float(block[packIdx].data[elementIdx].y);
        half2 res = __shfl(temp, srcLane);
        half2_t result;
        result.x = (float)res.x;
        result.y = (float)res.y;
        return result;
muyangli's avatar
muyangli committed
356
357
    }
    // 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
358
359
360
    __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
361
        const int elementIdx = k % ASCALES_PACK_SIZE / 2;
fengzch's avatar
fengzch committed
362
363
364
365
366
367
368
369
        half2 temp;
        temp.x = float(block[packIdx].data[elementIdx].x);
        temp.y = float(block[packIdx].data[elementIdx].y);
        half2 res = __shfl(temp, srcLane);
        half2_t result;
        result.x = (float)res.x;
        result.y = (float)res.y;
        return result;
muyangli's avatar
muyangli committed
370
371
    }

372
    struct i2f_normal {
Muyang Li's avatar
Muyang Li committed
373
        __device__ __forceinline__ static float2 int2float2(int x, int y) {
374
375
376
            return make_float2(__int2float_rn(x), __int2float_rn(y));
        }

Muyang Li's avatar
Muyang Li committed
377
        __device__ __forceinline__ static half2_t int2half2(int x, int y) {
378
379
380
381
382
            return float22half2<half2_t>(int2float2(x, y));
        }
    };

    template<typename i2f = i2f_normal, typename F>
Muyang Li's avatar
Muyang Li committed
383
384
    __device__ __forceinline__ static void
    apply_scales(F &&getpsum, ascale_warp ascale, wscale_warp wscale, fpsum_warp &fpsum) {
muyangli's avatar
muyangli committed
385
386
387
388
389
390
391
392
        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
393
394
            asx[i]     = half2_t(as.x, as.x);
            asy[i]     = half2_t(as.y, as.y);
muyangli's avatar
muyangli committed
395
396
397
398
399
400
401
402
403
404
405
406
407
        }

        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
408
409
410

                //     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
411
412
                // }

413
414
415
416
417
                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
418
                // if (threadIdx.x == 3 && j == 1 && i == 0) {
Muyang Li's avatar
Muyang Li committed
419
420
                //     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
421
422
423
424
425
                // }
            }
        }
    }

426
    template<typename i2f = i2f_normal, typename F>
Muyang Li's avatar
Muyang Li committed
427
428
    __device__ __forceinline__ static void
    apply_scales(F &&getpsum, ascale_warp ascale, wscale_warp wscale, f32psum_warp &fpsum) {
muyangli's avatar
muyangli committed
429
430
431
432
433
434
435
436
        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
437
438
            asx[i]     = half22float2(half2_t(as.x, as.x));
            asy[i]     = half22float2(half2_t(as.y, as.y));
muyangli's avatar
muyangli committed
439
440
441
442
443
444
445
446
447
448
449
450
451
452
453
454
        }

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

455
456
457
458
                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
459
460
461
462
463
464
465
466
            }
        }
    }

    /**
     * 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
467
    __device__ __forceinline__ static void pack_ascales(const half_t *input, packed_ascale_t *output) {
muyangli's avatar
muyangli committed
468
        const int laneId = threadIdx.x % WARP_SIZE;
Muyang Li's avatar
Muyang Li committed
469
#pragma unroll
muyangli's avatar
muyangli committed
470
471
472
        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
473
#pragma unroll
muyangli's avatar
muyangli committed
474
                for (int i = 0; i < ASCALES_PACK_SIZE; i += 2) {
Muyang Li's avatar
Muyang Li committed
475
476
477
478
                    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
479
480
481
482
483
484
485
486
487
488
                }
                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
489
    __device__ __forceinline__ static void pack_wscales(const half_t *input, packed_wscale_t *output) {
muyangli's avatar
muyangli committed
490
        const int laneId = threadIdx.x % WARP_SIZE;
Muyang Li's avatar
Muyang Li committed
491
#pragma unroll
muyangli's avatar
muyangli committed
492
493
494
        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
495
#pragma unroll
muyangli's avatar
muyangli committed
496
                for (int i = 0; i < WSCALES_PACK_SIZE; i += 2) {
Muyang Li's avatar
Muyang Li committed
497
498
499
                    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
500
501
502
503
504
505
506
507
508
509
510
                }
                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
511
512
        static constexpr int PACK_SIZE  = WARP_N / WARP_SIZE;
        using pack_t                    = std::array<half_t, PACK_SIZE>;
muyangli's avatar
muyangli committed
513
514

        // F (int rowId, pack_t &pack)
Muyang Li's avatar
Muyang Li committed
515
516
517
        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
518
519
            const int laneId = threadIdx.x % WARP_SIZE;

Muyang Li's avatar
Muyang Li committed
520
            matrix_t &mat = *reinterpret_cast<matrix_t *>(shmem);
muyangli's avatar
muyangli committed
521
522
523
524
525
526
527
528
529
530
531
532
533
534
535
536

            // 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
537
#pragma unroll
muyangli's avatar
muyangli committed
538
            for (int i = 0; i < WARP_M_TILES; i++) {
Muyang Li's avatar
Muyang Li committed
539
#pragma unroll
muyangli's avatar
muyangli committed
540
                for (int j = 0; j < WARP_N_TILES; j++) {
Muyang Li's avatar
Muyang Li committed
541
542
543
                    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
544
545
546
                    *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
547
548
                // __syncwarp();
		__builtin_amdgcn_wave_barrier();
muyangli's avatar
muyangli committed
549

Muyang Li's avatar
Muyang Li committed
550
#pragma unroll
muyangli's avatar
muyangli committed
551
552
553
554
555
556
557
558
559
560
                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;
561
                    // if (pred) {
Muyang Li's avatar
Muyang Li committed
562
563
564
                    store_pred(reinterpret_cast<pack_t *>(&output[(i * INSN_M + row) * stride + laneId * PACK_SIZE]),
                               pack,
                               pred);
565
                    // }
muyangli's avatar
muyangli committed
566
                }
limm's avatar
limm committed
567
568
                // __syncwarp();
		__builtin_amdgcn_wave_barrier();
muyangli's avatar
muyangli committed
569

Muyang Li's avatar
Muyang Li committed
570
#pragma unroll
muyangli's avatar
muyangli committed
571
                for (int j = 0; j < WARP_N_TILES; j++) {
Muyang Li's avatar
Muyang Li committed
572
573
574
                    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
575
576
577
                    *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
578
579
                // __syncwarp();
		__builtin_amdgcn_wave_barrier();
muyangli's avatar
muyangli committed
580

Muyang Li's avatar
Muyang Li committed
581
#pragma unroll
muyangli's avatar
muyangli committed
582
583
584
585
586
587
588
589
590
591
                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;
592
                    // if (pred) {
Muyang Li's avatar
Muyang Li committed
593
594
595
596
                    store_pred(
                        reinterpret_cast<pack_t *>(&output[(i * INSN_M + 8 + row) * stride + laneId * PACK_SIZE]),
                        pack,
                        pred);
597
                    // }
muyangli's avatar
muyangli committed
598
                }
limm's avatar
limm committed
599
600
                // __syncwarp();
		__builtin_amdgcn_wave_barrier();
muyangli's avatar
muyangli committed
601
602
603
604
605
606
607
            }
            // if (enableReduce) {
            //     store<true>(reinterpret_cast<pack_t *>(&reduce_result[laneId * PACK_SIZE]), reduce_tmp);
            // }
        }
    };

sxtyzhangzk's avatar
sxtyzhangzk committed
608
609
610
611
    // 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
612
        using matrix_t                     = half_t[INSN_M][WARP_N + 8];
sxtyzhangzk's avatar
sxtyzhangzk committed
613
614
        static constexpr size_t SHMEM_SIZE = sizeof(matrix_t);

Muyang Li's avatar
Muyang Li committed
615
616
        __device__ __forceinline__ void
        operator()(const half_t *input, int stride, int maxRows, int maxCols, fpsum_warp &out, void *shmem) {
sxtyzhangzk's avatar
sxtyzhangzk committed
617
618
619
620
621
            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
622
623
            using packed_input      = std::array<half_t, PACK_SIZE>;
            using packed_raw_input  = std::array<half2_t, PACK_SIZE>;
sxtyzhangzk's avatar
sxtyzhangzk committed
624

Muyang Li's avatar
Muyang Li committed
625
#pragma unroll
sxtyzhangzk's avatar
sxtyzhangzk committed
626
            for (int m = 0; m < WARP_M_TILES; m++) {
Muyang Li's avatar
Muyang Li committed
627
#pragma unroll
sxtyzhangzk's avatar
sxtyzhangzk committed
628
629
630
631
632
633
634
635
                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
636
637
                            raw = load(reinterpret_cast<const packed_raw_input *>(input + (m * INSN_M + row) * stride +
                                                                                  laneId * PACK_SIZE * 2));
sxtyzhangzk's avatar
sxtyzhangzk committed
638
                        }
Muyang Li's avatar
Muyang Li committed
639
#pragma unroll
sxtyzhangzk's avatar
sxtyzhangzk committed
640
641
642
643
644
645
646
                        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
647
648
                            pack = load(reinterpret_cast<const packed_input *>(input + (m * INSN_M + row) * stride +
                                                                               laneId * PACK_SIZE));
sxtyzhangzk's avatar
sxtyzhangzk committed
649
650
651
652
                        }
                    }
                    store<true>(reinterpret_cast<packed_input *>(&mat[row][laneId * PACK_SIZE]), pack);
                }
limm's avatar
limm committed
653
654
                // __syncwarp();
		__builtin_amdgcn_wave_barrier();
sxtyzhangzk's avatar
sxtyzhangzk committed
655
656
657
658
659
660
661
662

                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
663
664
                // __syncwarp();
		__builtin_amdgcn_wave_barrier();
sxtyzhangzk's avatar
sxtyzhangzk committed
665
666
667
668
            }
        }
    };

muyangli's avatar
muyangli committed
669
    template<typename F>
Muyang Li's avatar
Muyang Li committed
670
    __device__ __forceinline__ static fpsum_warp apply_act(fpsum_warp fpsum, F func) {
muyangli's avatar
muyangli committed
671
        fpsum_warp result;
Muyang Li's avatar
Muyang Li committed
672
#pragma unroll
muyangli's avatar
muyangli committed
673
        for (int i = 0; i < WARP_M_TILES; i++) {
Muyang Li's avatar
Muyang Li committed
674
#pragma unroll
muyangli's avatar
muyangli committed
675
            for (int j = 0; j < WARP_N_TILES; j++) {
Muyang Li's avatar
Muyang Li committed
676
#pragma unroll
muyangli's avatar
muyangli committed
677
678
                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
679
680
681
                    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
682
683
684
685
686
687
688
689
690
691
692
693
                }
            }
        }
        return result;
    }

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

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

limm's avatar
limm committed
698
699
            // __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
700
701
702
703

            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
704
705
706
707
708
709
710
711
712
713
714
715
716
717
718
            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
719
720
721
722
723
        }
    };

    struct EpilogueNop {
        // workaround for layout mismatch between host and device code
Muyang Li's avatar
Muyang Li committed
724
725
726
        struct Arguments {
            size_t unused;
        };
muyangli's avatar
muyangli committed
727

Muyang Li's avatar
Muyang Li committed
728
729
        __device__ __forceinline__ void
        operator()(const BlockInfo binfo, fpsum_warp fpsum, int M, int N, int K, const Arguments &args) {}
muyangli's avatar
muyangli committed
730
731
    };

732
    template<bool USE_BIAS = true, bool USE_SCALE = false>
muyangli's avatar
muyangli committed
733
734
    struct EpilogueBias {
        struct Arguments {
Muyang Li's avatar
Muyang Li committed
735
            const packed_wscale_t *bias; // [N / BLOCK_N, WSCALES_NUM_PACKS, WSCALES_VALID_LANES] of packed_wscale_t
736
            const packed_wscale_t *scale;
muyangli's avatar
muyangli committed
737
738
        };

Muyang Li's avatar
Muyang Li committed
739
740
        __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
741
742
743
            const int laneId = threadIdx.x % WARP_SIZE;

            // if (laneId == 0) {
Muyang Li's avatar
Muyang Li committed
744
745
            //     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
746
747
            // }

748
749
750
751
752
753
754
            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
755
756

            for (int j = 0; j < WARP_N_TILES; j++) {
757
758
759
760
761
762
763
764
765
766
                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
767
768
769
770

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

771
772
773
774
775
776
777
778
779
780
781
782
783
784
785
786
                    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
787
788
789
790
                }
            }
        }

Muyang Li's avatar
Muyang Li committed
791
792
        __device__ __forceinline__ void
        operator()(const BlockInfo binfo, fpsum_warp &fpsum, int M, int N, int K, const Arguments &args) {
muyangli's avatar
muyangli committed
793
            const int bn = binfo.bn;
794
            if constexpr (USE_BIAS || USE_SCALE) {
Muyang Li's avatar
Muyang Li committed
795
796
797
798
799
800
                apply_bias(fpsum,
                           M,
                           N,
                           K,
                           args.bias + bn * WSCALES_NUM_PACKS * WSCALES_VALID_LANES,
                           args.scale + bn * WSCALES_NUM_PACKS * WSCALES_VALID_LANES);
801
            }
muyangli's avatar
muyangli committed
802
803
804
805
        }
    };

    struct EpilogueSilu {
Muyang Li's avatar
Muyang Li committed
806
807
808
        struct Arguments {
            size_t unused;
        };
muyangli's avatar
muyangli committed
809

Muyang Li's avatar
Muyang Li committed
810
811
        __device__ __forceinline__ void
        operator()(const BlockInfo binfo, fpsum_warp &fpsum, int M, int N, int K, const Arguments &args) {
muyangli's avatar
muyangli committed
812
813
814
815
            fpsum = apply_act(fpsum, [](half_t x) { return silu(x); });
        }
    };

Muyang Li's avatar
Muyang Li committed
816
    template<typename... Epilogues>
muyangli's avatar
muyangli committed
817
818
819
    struct EpilogueCombination {
        using Arguments = std::tuple<typename Epilogues::Arguments...>;

Muyang Li's avatar
Muyang Li committed
820
821
        __device__ __forceinline__ void
        operator()(const BlockInfo binfo, fpsum_warp &fpsum, int M, int N, int K, const Arguments &args) {
muyangli's avatar
muyangli committed
822
            // this function makes intellisense crashes :(
Muyang Li's avatar
Muyang Li committed
823
824
825
#if __INTELLISENSE__
            __trap(); // should not happen when actually compiling
#else
muyangli's avatar
muyangli committed
826
827
828
829
            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
830
831
832
833
            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
834
835
836
837
        }
    };
};

Muyang Li's avatar
Muyang Li committed
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
873
874
875
876
877
878
879
880
#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
881
    using EpilogueBias = typename Base::template EpilogueBias<USE_BIAS, USE_SCALE>;                                             \
Muyang Li's avatar
Muyang Li committed
882
883
884
885
886
887
888
889
890
891
892
893
    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
894
    using Base::apply_act;
muyangli's avatar
muyangli committed
895

896
897
template<typename kernel>
constexpr int min_arch() {
Muyang Li's avatar
Muyang Li committed
898
    if constexpr (requires { kernel::MIN_ARCH; }) {
899
900
901
902
903
904
905
        return kernel::MIN_ARCH;
    } else {
        return 0;
    }
}
template<typename kernel>
constexpr int max_arch() {
Muyang Li's avatar
Muyang Li committed
906
    if constexpr (requires { kernel::MAX_ARCH; }) {
907
908
909
910
911
        return kernel::MAX_ARCH;
    } else {
        return INT_MAX;
    }
}
muyangli's avatar
muyangli committed
912

Muyang Li's avatar
Muyang Li committed
913
914
template<typename kernel, typename... T>
__global__ static void invoke_kernel(T... args) {
fengzch's avatar
fengzch committed
915
916
917
918
919
920
921
// #ifdef __CUDA_ARCH__
//     if constexpr (__CUDA_ARCH__ >= min_arch<kernel>() && __CUDA_ARCH__ <= max_arch<kernel>()) {
//         kernel()(args...);
//     } else {
//         trap_unsupported_arch();
//     }
// #else
922
    // ???
muyangli's avatar
muyangli committed
923
    kernel()(args...);
fengzch's avatar
fengzch committed
924
// #endif
muyangli's avatar
muyangli committed
925
926
927
}

template<typename T>
Muyang Li's avatar
Muyang Li committed
928
__global__ static void test_sizeof_device() {
muyangli's avatar
muyangli committed
929
930
931
932
933
934
935
936
937
938
939
940
    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
941
942
    test_sizeof_device<T><<<1, 1>>>();
    checkCUDA(cudaDeviceSynchronize());
muyangli's avatar
muyangli committed
943
944
}

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