gemm_w4a4.cuh 46.7 KB
Newer Older
muyangli's avatar
muyangli committed
1
2
3
#pragma once

#include "gemm_base.cuh"
sxtyzhangzk's avatar
sxtyzhangzk committed
4
#include "lora.cuh"
5
// #include "gemm_w4a4_block.cuh"
muyangli's avatar
muyangli committed
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23

namespace nunchaku::kernels {

template<typename Config>
class GEMM_W4A4;

#ifndef __INTELLISENSE__
template<typename Config>
class GEMM_W4A4 : public GEMMBase<Config> {
#else
template<>
class GEMM_W4A4<GEMMConfig_W4A4_FP16> : public GEMMBase<GEMMConfig_W4A4_FP16> {
    using Config = GEMMConfig_W4A4_FP16;
#endif

public:
    IMPORT_GEMM_BASE(Config);

24
25
26
27
28
29
30
31
32
33
public:
    // micro-scales for FP4 MMA
    // each uint32_t is a 4*32 matrix of scales (for MMA of 64*32)

#if defined(__CUDA_ARCH__) && __CUDA_ARCH__ >= 1200
    static constexpr bool FP4_AVAILABLE = true;
#else
    static constexpr bool FP4_AVAILABLE = false;
#endif

Muyang Li's avatar
Muyang Li committed
34
    __device__ __forceinline__ static void trap_no_fp4() {
35
36
37
38
39
40
41
42
43
44
45
        if (blockIdx.x == 0 && blockIdx.y == 0 && threadIdx.x == 0) {
            printf("FP4 is not available on this device\n");
        }
        __syncthreads();
        __nanosleep(1000000);
        __trap();
    }

    static_assert(WARP_N % 32 == 0);
    static_assert(WARP_M % 32 == 0);

Muyang Li's avatar
Muyang Li committed
46
47
    static constexpr int WMSCALES_PACK_SIZE   = clamp(WARP_N / 32, 1, 4);
    static constexpr int WMSCALES_NUM_PACKS   = ceilDiv(WARP_N / 32, WMSCALES_PACK_SIZE);
48
49
    static constexpr int WMSCALES_VALID_LANES = WARP_SIZE;

Muyang Li's avatar
Muyang Li committed
50
51
    static constexpr int AMSCALES_PACK_SIZE   = clamp(WARP_M / 32, 1, 4);
    static constexpr int AMSCALES_NUM_PACKS   = ceilDiv(WARP_M / 32, AMSCALES_PACK_SIZE);
52
53
54
55
56
57
58
59
60
61
62
63
    static constexpr int AMSCALES_VALID_LANES = WARP_SIZE;

    struct packed_wmscale_t {
        uint32_t data[WMSCALES_PACK_SIZE];
    };
    struct packed_amscale_t {
        uint32_t data[AMSCALES_PACK_SIZE];
    };
    using amscale_warp = std::array<packed_amscale_t, AMSCALES_NUM_PACKS>;
    using wmscale_warp = std::array<packed_wmscale_t, WMSCALES_NUM_PACKS>;

    // amscales: [M / BLOCK_M, K / group size, NUM_WARPS, AMSCALES_NUM_PACKS, WARP_SIZE] of packed_amscale_t
Muyang Li's avatar
Muyang Li committed
64
65
    __device__ __forceinline__ static void
    load_amscale(const packed_amscale_t *ptr, int group, amscale_warp &out, bool pred) {
66
67
        const int laneId = threadIdx.x % WARP_SIZE;
        const int warpId = threadIdx.x / WARP_SIZE;
Muyang Li's avatar
Muyang Li committed
68
#pragma unroll
69
        for (int i = 0; i < AMSCALES_NUM_PACKS; i++) {
Muyang Li's avatar
Muyang Li committed
70
71
72
            out[i] = load_pred(&ptr[(group * NUM_WARPS + warpId) * AMSCALES_NUM_PACKS * AMSCALES_VALID_LANES +
                                    i * AMSCALES_VALID_LANES + laneId],
                               pred);
73
74
75
        }
    }
    // wmscales: [N / BLOCK_N, 1, K / group size, WMSCALES_NUM_PACKS, WMSCALES_VALID_LANES] of packed_wmscale_t
Muyang Li's avatar
Muyang Li committed
76
77
    __device__ __forceinline__ static void
    load_wmscale(const packed_wmscale_t *ptr, int group, wmscale_warp &out, bool pred) {
78
        const int laneId = threadIdx.x % WARP_SIZE;
Muyang Li's avatar
Muyang Li committed
79
#pragma unroll
80
81
82
83
84
        for (int i = 0; i < WMSCALES_NUM_PACKS; i++) {
            out[i] = load_pred(&ptr[(group * WMSCALES_NUM_PACKS + i) * WMSCALES_VALID_LANES + laneId], pred);
        }
    }

Muyang Li's avatar
Muyang Li committed
85
86
    __device__ __forceinline__ static void quantize_w4a4_fp4_from_fpsum_warp(
        const packed_fpsum_t (&fpsum)[INSN_K / INSN_N], packed_act_t &output, uint32_t &output_scale, int ida) {
87
88
89
90

        constexpr int NUM_GROUPS = 4;
        static_assert(NUM_GROUPS == INSN_K / INSN_N);

Muyang Li's avatar
Muyang Li committed
91
        constexpr float QVALUE_MAX       = 6.0f;
92
        constexpr float RECPI_QVALUE_MAX = 1 / QVALUE_MAX;
Muyang Li's avatar
Muyang Li committed
93
        constexpr float MSCALE_MAX       = 448.0f;
94
95
96
97
98
99

        const int laneId = threadIdx.x % WARP_SIZE;

        // 0 for row 0-7; 1 for row 8-15
        // each half2_t represents a 8*8 matrix
        half2_t input[2][INSN_K / INSN_N * 2];
Muyang Li's avatar
Muyang Li committed
100
#pragma unroll
101
102
103
104
105
106
        for (int i = 0; i < INSN_K / INSN_N; i++) {
            input[0][i * 2 + 0] = fpsum[i].data[0];
            input[0][i * 2 + 1] = fpsum[i].data[2];
            input[1][i * 2 + 0] = fpsum[i].data[1];
            input[1][i * 2 + 1] = fpsum[i].data[3];
        }
Muyang Li's avatar
Muyang Li committed
107

108
109
110
111
112
113
114
        auto maxabs = [](half2_t val) ALWAYSINLINE {
            val = __habs2(val);
            return __hmax(val.x, val.y);
        };

        // each half_t represents maxvalue in a 8*16 matrix
        half_t maxvalue[2][NUM_GROUPS];
Muyang Li's avatar
Muyang Li committed
115
#pragma unroll
116
117
118
119
        for (int i = 0; i < NUM_GROUPS; i++) {
            maxvalue[0][i] = __hmax(maxabs(input[0][i * 2]), maxabs(input[0][i * 2 + 1]));
            maxvalue[1][i] = __hmax(maxabs(input[1][i * 2]), maxabs(input[1][i * 2 + 1]));
        }
Muyang Li's avatar
Muyang Li committed
120
#pragma unroll
121
        for (int mask = 2; mask > 0; mask /= 2) {
Muyang Li's avatar
Muyang Li committed
122
#pragma unroll
123
124
125
126
127
128
129
130
131
            for (int i = 0; i < NUM_GROUPS; i++) {
                maxvalue[0][i] = __hmax(maxvalue[0][i], __shfl_xor_sync(~0, maxvalue[0][i], mask));
                maxvalue[1][i] = __hmax(maxvalue[1][i], __shfl_xor_sync(~0, maxvalue[1][i], mask));
            }
        }
        // lane 0,1,2,3 / 4,5,6,7 / ...  should have identical maxvalue now

        float scale[2][NUM_GROUPS];
        float rscale[2][NUM_GROUPS];
Muyang Li's avatar
Muyang Li committed
132
#pragma unroll
133
        for (int i = 0; i < NUM_GROUPS; i++) {
Muyang Li's avatar
Muyang Li committed
134
135
            scale[0][i] = fminf(float(maxvalue[0][i]) * RECPI_QVALUE_MAX, MSCALE_MAX);
            scale[1][i] = fminf(float(maxvalue[1][i]) * RECPI_QVALUE_MAX, MSCALE_MAX);
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
            // TODO: check whether (1 / scale) or (1 / fp8scale) is better
            rscale[0][i] = cuda_frcp(scale[0][i]);
            rscale[1][i] = cuda_frcp(scale[1][i]);
        }
        uint32_t fp8scale[2];
        fp8scale[0] = quantize_float4_fp8(make_float4(scale[0][0], scale[0][1], scale[0][2], scale[0][3]));
        fp8scale[1] = quantize_float4_fp8(make_float4(scale[1][0], scale[1][1], scale[1][2], scale[1][3]));

        /**
         * output_scale pack format: (ida=0)
         * lane 0 => row 0 if ida==0
         * lane 1 => row 8 if ida==0
         * lane 2 => row 0 if ida==1
         * lane 3 => row 8 if ida==1
         * ...
         * lane i => quad (i/4) => row (i/4+8*(i%2)) if (i%4/2==ida) => srclane i, index i%2
         */
        if (laneId % 4 / 2 == ida) {
            output_scale = (laneId % 2 == 0) ? fp8scale[0] : fp8scale[1];
        }
Muyang Li's avatar
Muyang Li committed
156

157
        uint32_t qpacks[2][INSN_K / INSN_M * 2];
Muyang Li's avatar
Muyang Li committed
158
#pragma unroll
159
        for (int i = 0; i < INSN_K / INSN_M * 2; i++) {
Muyang Li's avatar
Muyang Li committed
160
#pragma unroll
161
            for (int j = 0; j < 2; j++) {
Muyang Li's avatar
Muyang Li committed
162
                float2 fval  = half22float2(input[j][i]) * make_float2(rscale[j][i / 2], rscale[j][i / 2]);
163
164
165
166
                qpacks[j][i] = quantize_float2_fp4(fval) << (laneId % 4 * 8);
            }
        }

Muyang Li's avatar
Muyang Li committed
167
#pragma unroll
168
        for (int mask = 1; mask <= 2; mask *= 2) {
Muyang Li's avatar
Muyang Li committed
169
#pragma unroll
170
            for (int i = 0; i < INSN_K / INSN_M * 2; i++) {
Muyang Li's avatar
Muyang Li committed
171
#pragma unroll
172
173
174
175
176
177
                for (int j = 0; j < 2; j++) {
                    qpacks[j][i] |= __shfl_xor_sync(~0, qpacks[j][i], mask);
                }
            }
        }
        // lane 0,1,2,3 / 4,5,6,7 / ...  should have identical qpacks now
Muyang Li's avatar
Muyang Li committed
178
#pragma unroll
179
180
181
182
183
184
185
186
187
188
189
190
        for (int i = 0; i < 4; i++) {
            if (laneId % 4 == i) {
                output.x = qpacks[0][0 + i];
                output.y = qpacks[1][0 + i];
                output.z = qpacks[0][4 + i];
                output.w = qpacks[1][4 + i];
            }
        }
    }

    // m16n16k64 MMA
    // ida, idb in {0, 1}
Muyang Li's avatar
Muyang Li committed
191
192
193
194
195
196
197
    __device__ __forceinline__ static packed_f32psum_t mma_fp4(packed_act_t act,
                                                               packed_wgt_t wgt,
                                                               packed_f32psum_t psum,
                                                               uint32_t amscale,
                                                               uint32_t wmscale,
                                                               int ida,
                                                               int idb) {
198
        packed_f32psum_t out;
Muyang Li's avatar
Muyang Li committed
199
        asm volatile(
200
201
202
203
204
205
            "mma.sync.aligned.m16n8k64.row.col.kind::mxf4nvf4.block_scale.scale_vec::4X.f32.e2m1.e2m1.f32.ue4m3 "
            "{%0, %1, %2, %3}, "
            "{%4, %5, %6, %7}, "
            "{%8, %9}, "
            "{%10, %11, %12, %13}, "
            "{%14}, {%15, %16}, "
Muyang Li's avatar
Muyang Li committed
206
207
208
209
210
211
212
213
214
215
216
217
218
219
220
221
222
223
224
            "{%17}, {%18, %19};"
            : "=f"(out.data[0]), "=f"(out.data[1]), "=f"(out.data[2]), "=f"(out.data[3])
            : "r"(act.x),
              "r"(act.y),
              "r"(act.z),
              "r"(act.w),
              "r"(wgt.x),
              "r"(wgt.y),
              "f"(psum.data[0]),
              "f"(psum.data[1]),
              "f"(psum.data[2]),
              "f"(psum.data[3]),
              "r"(amscale),
              "n"(0),
              "h"((short)ida),
              "r"(wmscale),
              "n"(0),
              "h"((short)(idb * 2)));
        asm volatile(
225
226
227
228
229
230
            "mma.sync.aligned.m16n8k64.row.col.kind::mxf4nvf4.block_scale.scale_vec::4X.f32.e2m1.e2m1.f32.ue4m3 "
            "{%0, %1, %2, %3}, "
            "{%4, %5, %6, %7}, "
            "{%8, %9}, "
            "{%10, %11, %12, %13}, "
            "{%14}, {%15, %16}, "
Muyang Li's avatar
Muyang Li committed
231
232
233
234
235
236
237
238
239
240
241
242
243
244
245
246
247
248
            "{%17}, {%18, %19};"
            : "=f"(out.data[4]), "=f"(out.data[5]), "=f"(out.data[6]), "=f"(out.data[7])
            : "r"(act.x),
              "r"(act.y),
              "r"(act.z),
              "r"(act.w),
              "r"(wgt.z),
              "r"(wgt.w),
              "f"(psum.data[4]),
              "f"(psum.data[5]),
              "f"(psum.data[6]),
              "f"(psum.data[7]),
              "r"(amscale),
              "n"(0),
              "h"((short)ida),
              "r"(wmscale),
              "n"(0),
              "h"((short)(idb * 2 + 1)));
249
250
251
        return out;
    }

Muyang Li's avatar
Muyang Li committed
252
253
    __device__ __forceinline__ static void
    compute_fp4(act_warp A, wgt_warp W, amscale_warp amscale, wmscale_warp wmscale, f32psum_warp &psum) {
254
255
256
        const int laneId = threadIdx.x % WARP_SIZE;
        const int warpId = threadIdx.x / WARP_SIZE;

Muyang Li's avatar
Muyang Li committed
257
#pragma unroll
258
        for (int j = 0; j < WARP_N_TILES; j++) {
Muyang Li's avatar
Muyang Li committed
259
#pragma unroll
260
            for (int i = 0; i < WARP_M_TILES; i++) {
Muyang Li's avatar
Muyang Li committed
261
262
263
264
265
266
267
268
                psum[i * WARP_N_TILES + j] =
                    mma_fp4(A[i],
                            W[j],
                            psum[i * WARP_N_TILES + j],
                            amscale[i / 2 / AMSCALES_PACK_SIZE].data[i / 2 % AMSCALES_PACK_SIZE],
                            wmscale[j / 2 / WMSCALES_PACK_SIZE].data[j / 2 % WMSCALES_PACK_SIZE],
                            i % 2,
                            j % 2);
269
270
271
272
273
            }
        }
    }

    template<typename Epilogue, bool USE_ALPHA>
Muyang Li's avatar
Muyang Li committed
274
275
276
277
278
279
280
281
282
283
284
    __device__ __forceinline__ static void gemm_w4a4_fp4_block(const BlockInfo binfo,
                                                               const packed_act_t *act,
                                                               const packed_wgt_t *wgt,
                                                               const packed_amscale_t *ascales,
                                                               const packed_wmscale_t *wscales,
                                                               float alpha, // per-tensor scale of weight
                                                               int M,
                                                               int N,
                                                               int K,
                                                               const Epilogue::Arguments &epilogueArgs,
                                                               bool alwaysfalse) {
285
286
287
288
289
        constexpr int NUM_STAGES = 2;

        const int laneId = threadIdx.x % WARP_SIZE;
        const int warpId = threadIdx.x / WARP_SIZE;

Muyang Li's avatar
Muyang Li committed
290
291
292
293
294
        act_warp A[NUM_STAGES];           // 8 * 2
        wgt_warp W[NUM_STAGES];           // 32 * 2
        amscale_warp amscale[NUM_STAGES]; // 1 * 2
        wmscale_warp wmscale[NUM_STAGES]; // 4 * 2
        f32psum_warp fpsum;               // 128
295
296
297
298
299
300
301
302

        for (int k = 0; k < NUM_STAGES - 1; k++) {
            load_act(act, k, K, A[k], true);
            load_wgt(wgt, k, K, W[k], true);
            load_amscale(ascales, k, amscale[k], true);
            load_wmscale(wscales, k, wmscale[k], true);
        }

Muyang Li's avatar
Muyang Li committed
303
#pragma unroll
304
        for (auto &pack : fpsum) {
Muyang Li's avatar
Muyang Li committed
305
#pragma unroll
306
307
308
309
            for (int i = 0; i < 8; i++) {
                pack.data[i] = 0;
            }
        }
Muyang Li's avatar
Muyang Li committed
310

311
312
313
        int dummy = 0;

        for (int k1 = 0; k1 < K / WARP_K; k1 += NUM_STAGES) {
Muyang Li's avatar
Muyang Li committed
314
#pragma unroll
315
316
            for (int k2 = 0; k2 < NUM_STAGES; k2++) {
                int nextk = k1 + k2 + NUM_STAGES - 1;
Muyang Li's avatar
Muyang Li committed
317
                int idx   = (k2 + NUM_STAGES - 1) % NUM_STAGES;
318
319
320
321
322
323
324
325
326
327
328
329
330
331
332
333
334
335
336
337
338
339
340
341
                bool pred = nextk < K / WARP_K;
                load_act(act, nextk, K, A[idx], pred);
                load_wgt(wgt, nextk, K, W[idx], pred);
                load_amscale(ascales, nextk, amscale[idx], pred);
                load_wmscale(wscales, nextk, wmscale[idx], pred);

                // __syncthreads();
                // if (alwaysfalse) {
                //     dummy = clock();
                // }

                compute_fp4(A[k2], W[k2], amscale[k2], wmscale[k2], fpsum);

                if (alwaysfalse) {
                    dummy = clock();
                }

                // asm volatile ("membar.cta;");
            }
        }

        unused_var(dummy, alwaysfalse);

        if constexpr (USE_ALPHA) {
Muyang Li's avatar
Muyang Li committed
342
#pragma unroll
343
            for (auto &pack : fpsum) {
Muyang Li's avatar
Muyang Li committed
344
#pragma unroll
345
346
347
348
349
350
351
352
353
354
355
356
357
358
359
                for (int i = 0; i < 8; i++) {
                    pack.data[i] *= alpha;
                }
            }
        }

        auto f16psum = packed_fp32_to_fp16(fpsum);

        CHECK_NAN(f16psum, "f16psum");

        Epilogue()(binfo, f16psum, M, N, K, epilogueArgs);
    }

    template<typename Epilogue, bool USE_ALPHA>
    struct gemm_w4a4_fp4_kernel {
360
        static constexpr int MIN_ARCH = 1200;
Muyang Li's avatar
Muyang Li committed
361
362
363
364
365
366
367
368
369
370
371
        __device__ void operator()(const packed_act_t *act,
                                   const packed_wgt_t *wgt,
                                   const packed_amscale_t *ascales,
                                   const packed_wmscale_t *wscales,
                                   float alpha,
                                   int M,
                                   int N,
                                   int K,
                                   Epilogue::Arguments epilogueArgs,
                                   bool swapBlockXY,
                                   bool alwaysfalse) {
372
            BlockInfo binfo = {
Muyang Li's avatar
Muyang Li committed
373
374
                .bm         = (int)blockIdx.x,
                .bn         = (int)blockIdx.y,
375
376
377
378
379
380
381
382
383
384
385
386
387
388
389
390
391
392
393
394
                .numBlocksM = (int)gridDim.x,
                .numBlocksN = (int)gridDim.y,
            };

            if (swapBlockXY) {
                std::swap(binfo.bm, binfo.bn);
                std::swap(binfo.numBlocksM, binfo.numBlocksN);
            }

            const int bm = binfo.bm;
            const int bn = binfo.bn;

            if constexpr (FP4_AVAILABLE) {
                gemm_w4a4_fp4_block<Epilogue, USE_ALPHA>(
                    binfo,
                    act + bm * (K / WARP_K) * NUM_WARPS * WARP_M_TILES * WARP_SIZE,
                    wgt + bn * (K / WARP_K) * WARP_N_TILES * WARP_SIZE,
                    ascales + bm * (K / WARP_K) * NUM_WARPS * AMSCALES_NUM_PACKS * AMSCALES_VALID_LANES,
                    wscales + bn * (K / WARP_K) * WMSCALES_NUM_PACKS * WMSCALES_VALID_LANES,
                    alpha,
Muyang Li's avatar
Muyang Li committed
395
396
397
                    M,
                    N,
                    K,
398
                    epilogueArgs,
Muyang Li's avatar
Muyang Li committed
399
                    alwaysfalse);
400
401
402
403
404
405
406

            } else {
                trap_no_fp4();
            }
        }
    };

muyangli's avatar
muyangli committed
407
408
public:
    template<bool ACT_UNSIGNED>
Muyang Li's avatar
Muyang Li committed
409
    __device__ __forceinline__ static packed_psum_t mma(packed_act_t act, packed_wgt_t wgt) {
muyangli's avatar
muyangli committed
410
411
        packed_psum_t psum;

Muyang Li's avatar
Muyang Li committed
412
413
414
415
        uint4 out1 = mma_m16n8kx_s32common<mma_helper::s4u4<ACT_UNSIGNED>, mma_helper::s4>(
            act, uint2(wgt.x, wgt.y), uint4(0, 0, 0, 0));
        uint4 out2 = mma_m16n8kx_s32common<mma_helper::s4u4<ACT_UNSIGNED>, mma_helper::s4>(
            act, uint2(wgt.z, wgt.w), uint4(0, 0, 0, 0));
416
417
418
419
420
421
422
423
        psum.data[0] = out1.x;
        psum.data[1] = out1.y;
        psum.data[2] = out1.z;
        psum.data[3] = out1.w;
        psum.data[4] = out2.x;
        psum.data[5] = out2.y;
        psum.data[6] = out2.z;
        psum.data[7] = out2.w;
Muyang Li's avatar
Muyang Li committed
424

muyangli's avatar
muyangli committed
425
426
427
428
429
        return psum;
    }

    // template<bool si>
    template<bool use_unsigned>
Muyang Li's avatar
Muyang Li committed
430
431
432
    __device__ __forceinline__ static void quantize_w4a4_from_fpsum_warp(const packed_fpsum_t (&fpsum)[INSN_K / INSN_N],
                                                                         packed_act_t &output,
                                                                         half_t *output_scale) {
muyangli's avatar
muyangli committed
433
434
        const int laneId = threadIdx.x % WARP_SIZE;

Muyang Li's avatar
Muyang Li committed
435
436
437
        constexpr float QVALUE_MAX_SIGNED         = 7.0f;
        constexpr float QVALUE_MAX_UNSIGNED       = 15.0f;
        constexpr float RECPI_QVALUE_MAX_SIGNED   = 1 / QVALUE_MAX_SIGNED;
muyangli's avatar
muyangli committed
438
439
        constexpr float RECPI_QVALUE_MAX_UNSIGNED = 1 / QVALUE_MAX_UNSIGNED;

Muyang Li's avatar
Muyang Li committed
440
        constexpr float QVALUE_MAX       = use_unsigned ? QVALUE_MAX_UNSIGNED : QVALUE_MAX_SIGNED;
muyangli's avatar
muyangli committed
441
442
443
444
445
446
        constexpr float RECPI_QVALUE_MAX = use_unsigned ? RECPI_QVALUE_MAX_UNSIGNED : RECPI_QVALUE_MAX_SIGNED;
        // constexpr int QUANTIZE_BITMASK = 0xf;

        // 0 for row 0-7; 1 for row 8-15
        half2_t input[2][INSN_K / INSN_N * 2];

Muyang Li's avatar
Muyang Li committed
447
#pragma unroll
muyangli's avatar
muyangli committed
448
449
450
451
452
453
454
455
456
457
        for (int i = 0; i < INSN_K / INSN_N; i++) {
            input[0][i * 2 + 0] = fpsum[i].data[0];
            input[0][i * 2 + 1] = fpsum[i].data[2];
            input[1][i * 2 + 0] = fpsum[i].data[1];
            input[1][i * 2 + 1] = fpsum[i].data[3];
        }

        half_t maxvalue[2];
        maxvalue[0] = 0;
        maxvalue[1] = 0;
Muyang Li's avatar
Muyang Li committed
458
#pragma unroll
muyangli's avatar
muyangli committed
459
460
461
        for (int i = 0; i < INSN_K / INSN_M * 2; i++) {
            half2_t abs0 = __habs2(input[0][i]);
            half2_t abs1 = __habs2(input[1][i]);
Muyang Li's avatar
Muyang Li committed
462
463
            maxvalue[0]  = __hmax(maxvalue[0], __hmax(abs0.x, abs0.y));
            maxvalue[1]  = __hmax(maxvalue[1], __hmax(abs1.x, abs1.y));
muyangli's avatar
muyangli committed
464
        }
Muyang Li's avatar
Muyang Li committed
465
#pragma unroll
muyangli's avatar
muyangli committed
466
467
468
469
470
471
472
473
474
475
476
477
478
        for (int mask = 2; mask > 0; mask /= 2) {
            maxvalue[0] = __hmax(maxvalue[0], __shfl_xor_sync(~0, maxvalue[0], mask));
            maxvalue[1] = __hmax(maxvalue[1], __shfl_xor_sync(~0, maxvalue[1], mask));
        }
        maxvalue[0] = __shfl_sync(~0, maxvalue[0], laneId / 4 * 4);
        maxvalue[1] = __shfl_sync(~0, maxvalue[1], laneId / 4 * 4);

        float scale[2];
        // scale[0] = float(maxvalue[0]) / QVALUE_MAX;
        // scale[1] = float(maxvalue[1]) / QVALUE_MAX;
        scale[0] = float(maxvalue[0]) * RECPI_QVALUE_MAX;
        scale[1] = float(maxvalue[1]) * RECPI_QVALUE_MAX;
        if (laneId % 4 == 0) {
Muyang Li's avatar
Muyang Li committed
479
            output_scale[laneId / 4]     = half_t(scale[0]);
muyangli's avatar
muyangli committed
480
481
482
483
484
485
486
487
488
489
            output_scale[laneId / 4 + 8] = half_t(scale[1]);
        }

        float rscale[2];
        // rscale[0] = QVALUE_MAX / float(maxvalue[0]);
        // rscale[1] = QVALUE_MAX / float(maxvalue[1]);
        rscale[0] = cuda_frcp(scale[0]);
        rscale[1] = cuda_frcp(scale[1]);

        uint32_t qpacks[2][INSN_K / INSN_M * 2];
Muyang Li's avatar
Muyang Li committed
490
#pragma unroll
muyangli's avatar
muyangli committed
491
        for (int i = 0; i < INSN_K / INSN_M * 2; i++) {
Muyang Li's avatar
Muyang Li committed
492
#pragma unroll
muyangli's avatar
muyangli committed
493
494
495
            for (int j = 0; j < 2; j++) {
                // half2_t hval = __hmul2(input[j][i], half2_t(rscale[j], rscale[j]));
                // float2 fval = half22float2(hval);
Muyang Li's avatar
Muyang Li committed
496
                float2 fval  = half22float2(input[j][i]) * make_float2(rscale[j], rscale[j]);
muyangli's avatar
muyangli committed
497
498
499
                qpacks[j][i] = quantize_float2<4, use_unsigned>(fval) << (laneId % 4 * 8);
            }
        }
Muyang Li's avatar
Muyang Li committed
500

muyangli's avatar
muyangli committed
501
        // 2 * 8 * 2 = 32 instructions => 256 cycles
Muyang Li's avatar
Muyang Li committed
502
#pragma unroll
muyangli's avatar
muyangli committed
503
        for (int mask = 1; mask <= 2; mask *= 2) {
Muyang Li's avatar
Muyang Li committed
504
#pragma unroll
muyangli's avatar
muyangli committed
505
            for (int i = 0; i < INSN_K / INSN_M * 2; i++) {
Muyang Li's avatar
Muyang Li committed
506
#pragma unroll
muyangli's avatar
muyangli committed
507
508
509
510
511
512
513
                for (int j = 0; j < 2; j++) {
                    qpacks[j][i] |= __shfl_xor_sync(~0, qpacks[j][i], mask);
                }
            }
        }
        // lane 0,1,2,3 / 4,5,6,7 / ...  should have identical qpacks now

Muyang Li's avatar
Muyang Li committed
514
#pragma unroll
muyangli's avatar
muyangli committed
515
516
517
518
519
520
521
522
523
524
525
526
527
528
529
530
        for (int i = 0; i < 4; i++) {
            if (laneId % 4 == i) {
                output.x = qpacks[0][0 + i];
                output.y = qpacks[1][0 + i];
                output.z = qpacks[0][4 + i];
                output.w = qpacks[1][4 + i];
            }
        }
    }

    /**
     * each warp quantizes a INSN_M * INSN_K (16 * 64) matrix
     * input is per-warp (in global memory)
     * output is per-thread (in regs)
     * output_scale is per-warp (in shared memory)
     * shmem must be at least INSN_M * INSN_K * sizeof(element) (16 * 64 * 0.5 = 512 Bytes)
Muyang Li's avatar
Muyang Li committed
531
532
     * default to quantize activation, if quantize weight, input should be column-majored and output should be
     * transposed ({x, y, z, w} = {x, z, y, w})
muyangli's avatar
muyangli committed
533
     */
Muyang Li's avatar
Muyang Li committed
534
535
    __device__ __forceinline__ static void
    quantize_w4a4_warp(const half_t *input, int stride, packed_act_t &output, half_t *output_scale, void *shmem) {
muyangli's avatar
muyangli committed
536
537
538
        const int laneId = threadIdx.x % WARP_SIZE;

        constexpr int QUANTIZE_BITWIDTH = 4;
Muyang Li's avatar
Muyang Li committed
539
        constexpr int QVALUE_MAX        = 7; // 4 bit => [-8, 7]
muyangli's avatar
muyangli committed
540
541
542

        // 1 lane = 1 pack
        // 1 warp = 32 lanes = 32 packs = 1 packwarp
Muyang Li's avatar
Muyang Li committed
543
544
545
546
547
        // a pack is {a0, ..., a7} in figure
        // https://docs.nvidia.com/cuda/parallel-thread-execution/index.html?highlight=ex2#mma-16864-a PACK_SIZE * 4 =
        // INSN_K / 2
        constexpr int PACK_SIZE             = INSN_K / 8; // = 8 for 4bit
        constexpr int NUM_PACKS_PER_ROW     = INSN_K / PACK_SIZE;
muyangli's avatar
muyangli committed
548
        constexpr int NUM_ROWS_PER_PACKWARP = PACK_SIZE * WARP_SIZE / INSN_K;
Muyang Li's avatar
Muyang Li committed
549
550
        constexpr int NUM_PACKWARPS         = INSN_M / NUM_ROWS_PER_PACKWARP;
        using packed_input                  = std::array<half_t, PACK_SIZE>;
muyangli's avatar
muyangli committed
551
552
553
554

        packed_input packs[NUM_PACKWARPS];

        // load
Muyang Li's avatar
Muyang Li committed
555
#pragma unroll
muyangli's avatar
muyangli committed
556
557
558
        for (int i = 0; i < NUM_PACKWARPS; i++) {
            int rowId = i * NUM_ROWS_PER_PACKWARP + laneId / NUM_PACKS_PER_ROW;
            int colId = laneId % NUM_PACKS_PER_ROW * PACK_SIZE;
Muyang Li's avatar
Muyang Li committed
559
            packs[i]  = load(reinterpret_cast<const packed_input *>(input + rowId * stride + colId));
muyangli's avatar
muyangli committed
560
561
562
563
        }

        // find max
        half_t maxvalue[NUM_PACKWARPS];
Muyang Li's avatar
Muyang Li committed
564
#pragma unroll
muyangli's avatar
muyangli committed
565
        for (int i = 0; i < NUM_PACKWARPS; i++) {
Muyang Li's avatar
Muyang Li committed
566
567
            maxvalue[i] = __habs(packs[i][0]);
#pragma unroll
muyangli's avatar
muyangli committed
568
            for (int j = 1; j < PACK_SIZE; j++) {
Muyang Li's avatar
Muyang Li committed
569

muyangli's avatar
muyangli committed
570
571
572
573
574
                maxvalue[i] = __hmax(maxvalue[i], __habs(packs[i][j]));
            }
        }

        // warp reduce (max)
Muyang Li's avatar
Muyang Li committed
575
#pragma unroll
muyangli's avatar
muyangli committed
576
        for (int mask = NUM_PACKS_PER_ROW / 2; mask > 0; mask /= 2) {
Muyang Li's avatar
Muyang Li committed
577
#pragma unroll
muyangli's avatar
muyangli committed
578
579
580
581
582
583
            for (int i = 0; i < NUM_PACKWARPS; i++) {
                maxvalue[i] = __hmax(maxvalue[i], __shfl_xor_sync(~0, maxvalue[i], mask));
            }
        }

        // broadcast (max)
Muyang Li's avatar
Muyang Li committed
584
#pragma unroll
muyangli's avatar
muyangli committed
585
586
587
588
589
590
        for (int i = 0; i < NUM_PACKWARPS; i++) {
            maxvalue[i] = __shfl_sync(~0, maxvalue[i], laneId / NUM_PACKS_PER_ROW * NUM_PACKS_PER_ROW);
        }

        // quantize
        using matrix_t = uint32_t[INSN_M][NUM_PACKS_PER_ROW];
Muyang Li's avatar
Muyang Li committed
591
592
        matrix_t &mat  = *reinterpret_cast<matrix_t *>(shmem);
#pragma unroll
muyangli's avatar
muyangli committed
593
594
595
596
597
598
599
600
        for (int i = 0; i < NUM_PACKWARPS; i++) {
            half_t scale  = maxvalue[i] / half_t(QVALUE_MAX);
            half_t rscale = half_t(QVALUE_MAX) / maxvalue[i];
            if (laneId % NUM_PACKS_PER_ROW == 0) {
                output_scale[i * NUM_ROWS_PER_PACKWARP + laneId / NUM_PACKS_PER_ROW] = scale;
            }

            uint32_t qpack = 0;
Muyang Li's avatar
Muyang Li committed
601
602
603
604
605
606
607
// #pragma unroll
//         for (int j = 0; j < PACK_SIZE; j++) {
//             int intvalue = __half2int_rn(packs[i][j] / scale);
//             intvalue = clamp(intvalue, -QVALUE_MAX, QVALUE_MAX);
//             qpack |= (intvalue & QUANTIZE_BITMASK) << (QUANTIZE_BITWIDTH * j);
//         }
#pragma unroll
muyangli's avatar
muyangli committed
608
609
610
611
612
613
614
            for (int j = 0; j < PACK_SIZE; j += 2) {
                half2_t hval = __hmul2(half2_t(rscale, rscale), half2_t(packs[i][j], packs[i][j + 1]));
                qpack |= quantize_float2<QUANTIZE_BITWIDTH, false>(half22float2(hval)) << (j * QUANTIZE_BITWIDTH);
            }
            mat[i * NUM_ROWS_PER_PACKWARP + laneId / NUM_PACKS_PER_ROW][laneId % NUM_PACKS_PER_ROW] = qpack;
        }
        __syncwarp();
Muyang Li's avatar
Muyang Li committed
615

muyangli's avatar
muyangli committed
616
617
618
619
620
621
622
623
624
625
        // convert to imma format
        int row = laneId % 16;
        int col = laneId / 16 * 4;
        ldmatrix(&mat[row][col], output);

        __syncwarp();
    }

    // each thread block (1 warp) quantize WARP_M * WARP_K tile (32 * 64)
    struct quantize_w4a4_act_kernel {
626
        static constexpr int MIN_ARCH = std::is_same_v<half_t, __nv_bfloat16> ? 800 : 750;
Muyang Li's avatar
Muyang Li committed
627
        __device__ void operator()(const half_t *input, packed_act_t *output, packed_ascale_t *oscales, int K) {
muyangli's avatar
muyangli committed
628
629
            const int laneId = threadIdx.x % WARP_SIZE;

Muyang Li's avatar
Muyang Li committed
630
631
            const int bm     = blockIdx.x / (BLOCK_M / WARP_M);
            const int bk     = blockIdx.y;
muyangli's avatar
muyangli committed
632
633
634
635
636
637
638
639
640
641
642
643
            const int warpId = blockIdx.x % (BLOCK_M / WARP_M);

            const int row = blockIdx.x * WARP_M;
            const int col = blockIdx.y * WARP_K;

            __shared__ alignas(128) half_t oscale_shmem[WARP_M];
            __shared__ alignas(128) uint8_t tmp_shmem[INSN_M * INSN_K / 2];

            for (int tileId = 0; tileId < WARP_M_TILES; tileId++) {
                packed_act_t tmpout;

                quantize_w4a4_warp(
Muyang Li's avatar
Muyang Li committed
644
                    input + (row + tileId * INSN_M) * K + col, K, tmpout, oscale_shmem + tileId * INSN_M, tmp_shmem);
muyangli's avatar
muyangli committed
645

Muyang Li's avatar
Muyang Li committed
646
647
648
                store(&output[(((bm * K / WARP_K + bk) * NUM_WARPS + warpId) * WARP_M_TILES + tileId) * WARP_SIZE +
                              laneId],
                      tmpout);
muyangli's avatar
muyangli committed
649
650
651
            }

            // if (threadIdx.x == 0) {
Muyang Li's avatar
Muyang Li committed
652
653
            //     printf("Block (%d, %d) => offset = %d\n", blockIdx.x, blockIdx.y, (bm * K / WARP_K + bk) * NUM_WARPS
            //     + warpId);
muyangli's avatar
muyangli committed
654
            // }
Muyang Li's avatar
Muyang Li committed
655
656
657
            pack_ascales(
                oscale_shmem,
                &oscales[((bm * K / WARP_K + bk) * NUM_WARPS + warpId) * ASCALES_NUM_PACKS * ASCALES_VALID_LANES]);
muyangli's avatar
muyangli committed
658
659
        }
    };
Muyang Li's avatar
Muyang Li committed
660

muyangli's avatar
muyangli committed
661
662
    // each thread block (1 warp) quantize WARP_N * WARP_K tile (128 * 64)
    struct quantize_w4a4_wgt_kernel {
663
        static constexpr int MIN_ARCH = std::is_same_v<half_t, __nv_bfloat16> ? 800 : 750;
Muyang Li's avatar
Muyang Li committed
664
        __device__ void operator()(const half_t *input, packed_wgt_t *output, packed_wscale_t *oscales, int K) {
muyangli's avatar
muyangli committed
665
666
667
668
669
670
671
672
673
674
675
676
677
678
679
            const int laneId = threadIdx.x % WARP_SIZE;

            const int bn = blockIdx.x / (BLOCK_N / WARP_N);
            const int bk = blockIdx.y;

            const int col = blockIdx.x * WARP_N;
            const int row = blockIdx.y * WARP_K;

            __shared__ alignas(128) half_t oscale_shmem[WARP_N];
            __shared__ alignas(128) uint8_t tmp_shmem[INSN_M * INSN_K / 2];

            for (int tileId = 0; tileId < WARP_N_TILES; tileId++) {
                packed_wgt_t tmpout;

                quantize_w4a4_warp(
Muyang Li's avatar
Muyang Li committed
680
                    input + (col + tileId * INSN_N) * K + row, K, tmpout, oscale_shmem + tileId * INSN_N, tmp_shmem);
muyangli's avatar
muyangli committed
681
682
683
684
685
686
687
688
689
690

                std::swap(tmpout.y, tmpout.z);

                store(&output[((bn * K / WARP_K + bk) * WARP_N_TILES + tileId) * WARP_SIZE + laneId], tmpout);
            }

            pack_wscales(oscale_shmem, &oscales[(bn * K / WARP_K + bk) * WSCALES_NUM_PACKS * WSCALES_VALID_LANES]);
        }
    };

691
    struct i2f_sm80 {
Muyang Li's avatar
Muyang Li committed
692
        __device__ __forceinline__ static float2 int2float2(int x, int y) {
693
694
695
            return make_float2(int2float_fast(x), int2float_fast(y));
        }

Muyang Li's avatar
Muyang Li committed
696
        __device__ __forceinline__ static half2_t int2half2(int x, int y) {
697
698
699
700
701
            return float22half2<half2_t>(int2float2(x, y));
        }
    };

    struct i2f_sm75 {
Muyang Li's avatar
Muyang Li committed
702
        __device__ __forceinline__ static float2 int2float2(int x, int y) {
703
704
705
            return make_float2(int2float_fast(x), int2float_fast(y));
        }

Muyang Li's avatar
Muyang Li committed
706
        __device__ __forceinline__ static half2_t int2half2(int x, int y) {
707
708
709
710
711
            return half2(__int2half_rn(x), __int2half_rn(y));
        }
    };

    struct i2f_sm75_fast {
Muyang Li's avatar
Muyang Li committed
712
        __device__ __forceinline__ static float2 int2float2(int x, int y) {
713
714
715
            return make_float2(int2float_fast(x), int2float_fast(y));
        }

Muyang Li's avatar
Muyang Li committed
716
        __device__ __forceinline__ static half2_t int2half2(int x, int y) {
717
718
719
720
            return int2half2_fast_512(x, y);
        }
    };

muyangli's avatar
muyangli committed
721
    template<bool ACT_UNSIGNED, typename T>
Muyang Li's avatar
Muyang Li committed
722
723
724
    __device__ __forceinline__ static void
    compute(act_warp A, wgt_warp W, ascale_warp ascale, wscale_warp wscale, T &fpsum) {
#if defined(__CUDA_ARCH__) && __CUDA_ARCH__ == 800
725
        using int2half2 = i2f_sm80;
Muyang Li's avatar
Muyang Li committed
726
727
728
729
730
731
732
733
734
#elif defined(__CUDA_ARCH__) && __CUDA_ARCH__ == 750
        using int2half2 = std::conditional_t<Config::FASTER_I2F, i2f_sm75_fast, i2f_sm75>;
        ;
#else
    using int2half2 = Base::i2f_normal;
#endif

        Base::template apply_scales<int2half2>(
            [&](int i, int j) { return mma<ACT_UNSIGNED>(A[i], W[j]); }, ascale, wscale, fpsum);
muyangli's avatar
muyangli committed
735
736
    }

Muyang Li's avatar
Muyang Li committed
737
    __device__ __forceinline__ static void checkNan(fpsum_warp fpsum, const char *info = "") {
muyangli's avatar
muyangli committed
738
739
740
741
742
743
744
745
#if ENABLE_NAN_CHECK
        const int laneId = threadIdx.x % WARP_SIZE;
        const int warpId = threadIdx.x / WARP_SIZE;

        for (int i = 0; i < fpsum.size(); i++) {
            for (int j = 0; j < 4; j++) {
                bool abnormal = !isfinite((float)fpsum[i].data[j].x) || !isfinite((float)fpsum[i].data[j].y);
                if (abnormal) {
Muyang Li's avatar
Muyang Li committed
746
747
748
749
750
751
752
753
754
755
756
                    printf("abnormal value detected at block.x=%d block.y=%d warpId=%d laneId=%d fpsum_warp (%s) i=%d "
                           "j=%d data.x=%f data.y=%f\n",
                           blockIdx.x,
                           blockIdx.y,
                           warpId,
                           laneId,
                           info,
                           i,
                           j,
                           (float)fpsum[i].data[j].x,
                           (float)fpsum[i].data[j].y);
muyangli's avatar
muyangli committed
757
758
759
760
761
762
763
                    __trap();
                }
            }
        }
#endif
    }

Muyang Li's avatar
Muyang Li committed
764
    __device__ __forceinline__ static void checkNan(packed_f32psum_t fpsum, const char *info = "") {
muyangli's avatar
muyangli committed
765
766
767
768
769
770
771
#if ENABLE_NAN_CHECK
        const int laneId = threadIdx.x % WARP_SIZE;
        const int warpId = threadIdx.x / WARP_SIZE;

        for (int j = 0; j < 8; j++) {
            bool abnormal = !isfinite(fpsum.data[j]);
            if (abnormal) {
Muyang Li's avatar
Muyang Li committed
772
773
774
775
776
777
                printf(
                    "abnormal value detected at bm=%d bn=%d warpId=%d laneId=%d packed_f32psum_t (%s) j=%d data=%f\n",
                    blockIdx.x,
                    blockIdx.y,
                    warpId,
                    laneId,
muyangli's avatar
muyangli committed
778
779
                    info,
                    j,
Muyang Li's avatar
Muyang Li committed
780
                    fpsum.data[j]);
muyangli's avatar
muyangli committed
781
782
783
784
785
786
                __trap();
            }
        }
#endif
    }

Muyang Li's avatar
Muyang Li committed
787
    __device__ __forceinline__ static void checkNan(packed_fpsum_t fpsum, const char *info = "") {
muyangli's avatar
muyangli committed
788
789
790
791
792
793
794
#if ENABLE_NAN_CHECK
        const int laneId = threadIdx.x % WARP_SIZE;
        const int warpId = threadIdx.x / WARP_SIZE;

        for (int j = 0; j < 4; j++) {
            bool abnormal = !isfinite((float)fpsum.data[j].x) || !isfinite((float)fpsum.data[j].y);
            if (abnormal) {
Muyang Li's avatar
Muyang Li committed
795
796
797
798
799
800
801
802
803
804
                printf("abnormal value detected at bm=%d bn=%d warpId=%d laneId=%d packed_fpsum_t (%s) j=%d data.x=%f "
                       "data.y=%f\n",
                       blockIdx.x,
                       blockIdx.y,
                       warpId,
                       laneId,
                       info,
                       j,
                       (float)fpsum.data[j].x,
                       (float)fpsum.data[j].y);
muyangli's avatar
muyangli committed
805
806
807
808
809
810
                __trap();
            }
        }
#endif
    }

Muyang Li's avatar
Muyang Li committed
811
    __device__ __forceinline__ static void checkNan(float data, const char *info = "") {
muyangli's avatar
muyangli committed
812
813
814
815
816
817
#if ENABLE_NAN_CHECK
        const int laneId = threadIdx.x % WARP_SIZE;
        const int warpId = threadIdx.x / WARP_SIZE;

        bool abnormal = !isfinite(data);
        if (abnormal) {
Muyang Li's avatar
Muyang Li committed
818
819
820
821
822
823
824
            printf("abnormal value detected at bm=%d bn=%d warpId=%d laneId=%d packed_fpsum_t (%s) data=%f\n",
                   blockIdx.x,
                   blockIdx.y,
                   warpId,
                   laneId,
                   info,
                   data);
muyangli's avatar
muyangli committed
825
826
827
828
829
830
            __trap();
        }
#endif
    }

    // out: [M / BLOCK_M, N / BLOCK_N, NUM_WARPS, 1, NUM_M_TILES, NUM_N_TILES, WARP_SIZE] of fpsum_warp
831
    template<typename Epilogue, bool ACT_UNSIGNED, bool USE_FP32_ACCUM>
Muyang Li's avatar
Muyang Li committed
832
833
834
835
836
837
838
839
840
841
842
843
    __device__ __forceinline__ static void gemm_w4a4_block(const BlockInfo binfo,
                                                           const packed_act_t *act,
                                                           const packed_wgt_t *wgt,
                                                           const packed_ascale_t *ascales,
                                                           const packed_wscale_t *wscales,
                                                           // const packed_wscale_t *bias_ptr,
                                                           // half_t *out,
                                                           int M,
                                                           int N,
                                                           int K,
                                                           const Epilogue::Arguments &epilogueArgs,
                                                           bool alwaysfalse) {
muyangli's avatar
muyangli committed
844
845
846
847
848
        constexpr int NUM_STAGES = 2;

        const int laneId = threadIdx.x % WARP_SIZE;
        const int warpId = threadIdx.x / WARP_SIZE;

849
850
851
852
#if 0
        fpsum_warp fpsum;
        GEMM_W4A4_Block<Config>()(act, wgt, ascales, wscales, K, fpsum, alwaysfalse);
#else
Muyang Li's avatar
Muyang Li committed
853
854
855
856
857
        act_warp A[NUM_STAGES];                                             // 8
        wgt_warp W[NUM_STAGES];                                             // 32
        ascale_warp ascale[NUM_STAGES];                                     // 1
        wscale_warp wscale[NUM_STAGES];                                     // 2
        std::conditional_t<USE_FP32_ACCUM, f32psum_warp, fpsum_warp> fpsum; // 64
muyangli's avatar
muyangli committed
858
859
860
861
862
863
864
865
866
867
868
869
870

        // load_wscale<true>(wscales, wscale[0], true);
        // load_wscale<false>(wscales, wscale[1], true);
        // load_wscale<false>(wscales, wscale[2], true);

        for (int k = 0; k < NUM_STAGES - 1; k++) {
            load_act(act, k, K, A[k], true);
            load_wgt(wgt, k, K, W[k], true);
            load_ascale(ascales, k, M, ascale[k], true);
            load_wscale(wscales, k, N, wscale[k], true);
        }

        for (auto &pack : fpsum) {
871
872
873
874
875
876
877
878
879
            if constexpr (USE_FP32_ACCUM) {
                for (int i = 0; i < 8; i++) {
                    pack.data[i] = 0;
                }
            } else {
                for (int i = 0; i < 4; i++) {
                    pack.data[i].x = 0;
                    pack.data[i].y = 0;
                }
muyangli's avatar
muyangli committed
880
881
            }
        }
Muyang Li's avatar
Muyang Li committed
882

muyangli's avatar
muyangli committed
883
884
885
        int dummy = 0;

        for (int k1 = 0; k1 < K / WARP_K; k1 += NUM_STAGES) {
Muyang Li's avatar
Muyang Li committed
886
#pragma unroll
muyangli's avatar
muyangli committed
887
888
            for (int k2 = 0; k2 < NUM_STAGES; k2++) {
                int nextk = k1 + k2 + NUM_STAGES - 1;
Muyang Li's avatar
Muyang Li committed
889
                int idx   = (k2 + NUM_STAGES - 1) % NUM_STAGES;
muyangli's avatar
muyangli committed
890
891
892
893
894
895
896
897
898
899
900
901
902
903
                bool pred = nextk < K / WARP_K;
                load_act(act, nextk, K, A[idx], pred);
                load_wgt(wgt, nextk, K, W[idx], pred);
                load_ascale(ascales, nextk, M, ascale[idx], pred);
                load_wscale(wscales, nextk, N, wscale[idx], pred);
                // load_wscale<false>(wscales, wscale[idx], pred);

                // __syncthreads();
                // if (alwaysfalse) {
                //     dummy = clock();
                // }

                compute<ACT_UNSIGNED>(A[k2], W[k2], ascale[k2], wscale[k2], fpsum);

Muyang Li's avatar
Muyang Li committed
904
                // #if defined(__CUDA_ARCH__) && __CUDA_ARCH__ >= 800
muyangli's avatar
muyangli committed
905
906
907
                if (alwaysfalse) {
                    dummy = clock();
                }
Muyang Li's avatar
Muyang Li committed
908
                // #endif
muyangli's avatar
muyangli committed
909
910
911
912
913
914
915

                // asm volatile ("membar.cta;");
            }
        }

        unused_var(dummy, alwaysfalse);

916
917
#endif

918
919
920
921
922
923
        fpsum_warp f16psum;
        if constexpr (USE_FP32_ACCUM) {
            f16psum = packed_fp32_to_fp16(fpsum);
        } else {
            f16psum = fpsum;
        }
muyangli's avatar
muyangli committed
924
925
926
927
928
929

        CHECK_NAN(f16psum, "f16psum");

        Epilogue()(binfo, f16psum, M, N, K, epilogueArgs);
    }

930
    template<bool FUSE_GELU, bool USE_UNSIGNED, bool USE_FP4>
muyangli's avatar
muyangli committed
931
    struct EpilogueQuantize {
932
933
        using oscales_t = typename std::conditional_t<USE_FP4, packed_amscale_t, packed_ascale_t>;

muyangli's avatar
muyangli committed
934
935
        struct Arguments {
            packed_act_t *qout;
936
            oscales_t *oscales;
muyangli's avatar
muyangli committed
937
938
939
940
941

            half_t shift_value;
            const packed_wscale_t *smooth_factor;
        };

Muyang Li's avatar
Muyang Li committed
942
        static constexpr int NUM_PACKS  = INSN_K / INSN_N;
muyangli's avatar
muyangli committed
943
944
        static constexpr int NUM_GROUPS = WARP_N_TILES / NUM_PACKS;

Muyang Li's avatar
Muyang Li committed
945
946
947
948
949
950
951
952
        __device__ __forceinline__ void apply_quantize(fpsum_warp fpsum,
                                                       int M,
                                                       int N,
                                                       int K,
                                                       packed_act_t *qout,
                                                       oscales_t *oscales,
                                                       half_t shift_value,
                                                       const packed_wscale_t *smooth_factor) {
muyangli's avatar
muyangli committed
953
954
955
956
957
958
959
960
            const int laneId = threadIdx.x % WARP_SIZE;
            const int warpId = threadIdx.x / WARP_SIZE;

            __shared__ half_t oscale_shmem[NUM_WARPS][WARP_M];

            wscale_warp smooth;
            load_wscale(smooth_factor, 0, N, smooth, true);

Muyang Li's avatar
Muyang Li committed
961
#pragma unroll
muyangli's avatar
muyangli committed
962
            for (int group = 0; group < NUM_GROUPS; group++) {
963
964
                amscale_warp omscale;

Muyang Li's avatar
Muyang Li committed
965
#pragma unroll
muyangli's avatar
muyangli committed
966
967
968
                for (int i = 0; i < WARP_M_TILES; i++) {
                    packed_fpsum_t tmp[NUM_PACKS];

Muyang Li's avatar
Muyang Li committed
969
#pragma unroll
muyangli's avatar
muyangli committed
970
971
972
                    for (int j = 0; j < NUM_PACKS; j++) {
                        half2_t ws1 = broadcast_wscale(smooth, (group * NUM_PACKS + j) * 4, laneId);
                        half2_t ws2 = broadcast_wscale(smooth, (group * NUM_PACKS + j) * 4 + 2, laneId);
Muyang Li's avatar
Muyang Li committed
973
#pragma unroll
muyangli's avatar
muyangli committed
974
                        for (int k = 0; k < 4; k++) {
Muyang Li's avatar
Muyang Li committed
975
                            half2_t src  = fpsum[i * WARP_N_TILES + group * NUM_PACKS + j].data[k];
muyangli's avatar
muyangli committed
976
977
978
979
980
981
982
983
984
985
986
987
988
989
990
991
992
993
994
995
996
                            half2_t &dst = tmp[j].data[k];

                            // dst.x = gelu(src.x);
                            // dst.y = gelu(src.y);
                            if constexpr (FUSE_GELU) {
                                dst = gelu_half2(src);
                            } else {
                                dst = src;
                            }

                            dst += half2_t(shift_value, shift_value);
                            // dst = src;
                        }

                        tmp[j].data[0] = h2div(tmp[j].data[0], ws1);
                        tmp[j].data[1] = h2div(tmp[j].data[1], ws1);
                        tmp[j].data[2] = h2div(tmp[j].data[2], ws2);
                        tmp[j].data[3] = h2div(tmp[j].data[3], ws2);
                    }

                    packed_act_t qresult;
997
                    if constexpr (USE_FP4) {
Muyang Li's avatar
Muyang Li committed
998
999
                        quantize_w4a4_fp4_from_fpsum_warp(
                            tmp, qresult, omscale[i / 2 / AMSCALES_PACK_SIZE].data[i / 2 % AMSCALES_PACK_SIZE], i % 2);
1000
1001
1002
                    } else {
                        quantize_w4a4_from_fpsum_warp<USE_UNSIGNED>(tmp, qresult, &oscale_shmem[warpId][i * INSN_M]);
                    }
muyangli's avatar
muyangli committed
1003
1004
1005
                    store(&qout[((group * NUM_WARPS + warpId) * WARP_M_TILES + i) * WARP_SIZE + laneId], qresult);
                }

1006
                if constexpr (USE_FP4) {
Muyang Li's avatar
Muyang Li committed
1007
#pragma unroll
1008
                    for (int k = 0; k < AMSCALES_NUM_PACKS; k++) {
Muyang Li's avatar
Muyang Li committed
1009
1010
1011
                        store(&oscales[((group * NUM_WARPS + warpId) * AMSCALES_NUM_PACKS + k) * AMSCALES_VALID_LANES +
                                       laneId],
                              omscale[k]);
1012
1013
1014
1015
                    }
                }
                if constexpr (!USE_FP4) {
                    __syncwarp();
Muyang Li's avatar
Muyang Li committed
1016
1017
                    pack_ascales(&oscale_shmem[warpId][0],
                                 &oscales[(group * NUM_WARPS + warpId) * ASCALES_NUM_PACKS * ASCALES_VALID_LANES]);
1018
1019
                    __syncwarp();
                }
muyangli's avatar
muyangli committed
1020
1021
1022
            }
        }

Muyang Li's avatar
Muyang Li committed
1023
1024
        __device__ __forceinline__ void
        operator()(const BlockInfo binfo, fpsum_warp fpsum, int M, int N, int K, const Arguments &args) {
muyangli's avatar
muyangli committed
1025
1026
1027
            const int bm = binfo.bm;
            const int bn = binfo.bn;

1028
            if constexpr (!USE_FP4 || FP4_AVAILABLE) {
Muyang Li's avatar
Muyang Li committed
1029
1030
1031
1032
1033
1034
1035
1036
1037
1038
                apply_quantize(fpsum,
                               M,
                               N,
                               K,
                               args.qout + (bm * N / WARP_K + bn * NUM_GROUPS) * NUM_WARPS * WARP_M_TILES * WARP_SIZE,
                               args.oscales + (bm * N / WARP_K + bn * NUM_GROUPS) * NUM_WARPS *
                                                  (USE_FP4 ? AMSCALES_NUM_PACKS * AMSCALES_VALID_LANES
                                                           : ASCALES_NUM_PACKS * ASCALES_VALID_LANES),
                               args.shift_value,
                               args.smooth_factor + bn * WSCALES_NUM_PACKS * WSCALES_VALID_LANES);
1039
1040
1041
            } else {
                trap_no_fp4();
            }
muyangli's avatar
muyangli committed
1042
1043
1044
1045
1046
1047
        }
    };
    // using EpilogueQuantizeFuseGelu = EpilogueQuantize<true>;

    template<typename Epilogue, bool ACT_UNSIGNED>
    struct gemm_w4a4_kernel {
1048
1049
1050
        static constexpr int MIN_ARCH = std::is_same_v<half_t, __nv_bfloat16> ? 800 : 750;
        static constexpr int MAX_ARCH = Config::FASTER_I2F ? 750 : INT_MAX; // FASTER_I2F is only needed on sm_75

Muyang Li's avatar
Muyang Li committed
1051
1052
1053
1054
1055
1056
1057
1058
1059
1060
        __device__ void operator()(const packed_act_t *act,
                                   const packed_wgt_t *wgt,
                                   const packed_ascale_t *ascales,
                                   const packed_wscale_t *wscales,
                                   int M,
                                   int N,
                                   int K,
                                   Epilogue::Arguments epilogueArgs,
                                   bool swapBlockXY,
                                   bool alwaysfalse) {
muyangli's avatar
muyangli committed
1061
1062
1063
            // printf("Device sizeof(args) = %d", (int)sizeof(epilogueArgs));

            BlockInfo binfo = {
Muyang Li's avatar
Muyang Li committed
1064
1065
                .bm         = (int)blockIdx.x,
                .bn         = (int)blockIdx.y,
muyangli's avatar
muyangli committed
1066
1067
1068
1069
1070
1071
1072
1073
1074
1075
1076
1077
1078
1079
                .numBlocksM = (int)gridDim.x,
                .numBlocksN = (int)gridDim.y,
            };

            if (swapBlockXY) {
                std::swap(binfo.bm, binfo.bn);
                std::swap(binfo.numBlocksM, binfo.numBlocksN);
            }

            const int bm = binfo.bm;
            const int bn = binfo.bn;

            // bool fusequant = !out;

1080
            gemm_w4a4_block<Epilogue, ACT_UNSIGNED, false>(
muyangli's avatar
muyangli committed
1081
1082
1083
1084
1085
1086
1087
1088
                binfo,
                act + bm * (K / WARP_K) * NUM_WARPS * WARP_M_TILES * WARP_SIZE,
                wgt + bn * (K / WARP_K) * WARP_N_TILES * WARP_SIZE,
                ascales + bm * (K / WARP_K) * NUM_WARPS * ASCALES_NUM_PACKS * ASCALES_VALID_LANES,
                wscales + bn * (K / WARP_K) * WSCALES_NUM_PACKS * WSCALES_VALID_LANES,
                // bias ? bias + bn * WSCALES_NUM_PACKS * WSCALES_VALID_LANES : nullptr,
                // out + (bm * BLOCK_M * N) + bn * BLOCK_N,
                // out + (bm * N / BLOCK_N + bn) * NUM_WARPS * WARP_M_TILES * WARP_N_TILES * WARP_SIZE,
Muyang Li's avatar
Muyang Li committed
1089
1090
1091
                M,
                N,
                K,
muyangli's avatar
muyangli committed
1092
                epilogueArgs,
Muyang Li's avatar
Muyang Li committed
1093
                alwaysfalse);
muyangli's avatar
muyangli committed
1094
1095
        }
    };
1096

sxtyzhangzk's avatar
sxtyzhangzk committed
1097
1098
1099
1100
    template<bool fuse_glu, bool use_fp4>
    struct quantize_w4a4_fuse_lora_kernel {
        using oscales_t = typename std::conditional_t<use_fp4, packed_amscale_t, packed_ascale_t>;

1101
        static constexpr int MIN_ARCH = std::is_same_v<half_t, __nv_bfloat16> ? 800 : 750;
Muyang Li's avatar
Muyang Li committed
1102
1103
        static constexpr size_t SHMEM_PER_WARP =
            ceilDiv<size_t>(Base::template load_act_to_fpsum<fuse_glu>::SHMEM_SIZE, 128) * 128;
1104
1105
1106
1107
        static constexpr size_t SHMEM_SIZE = SHMEM_PER_WARP * NUM_WARPS;

        struct Arguments {
            const half_t *input;
sxtyzhangzk's avatar
sxtyzhangzk committed
1108
1109
1110
1111
1112
1113
1114
            const packed_wscale_t *smooth_factor;
            packed_act_t *output;
            oscales_t *oscales;
            const packed_fpsum_t *lora_wgt_down;
            float *lora_act;

            int lora_rank;
1115
1116

            // aligned to BLOCK_M and BLOCK_N
Muyang Li's avatar
Muyang Li committed
1117
            int M, N; // N should be the actual K in the next GEMM (needs /2 if fuse_glu)
sxtyzhangzk's avatar
sxtyzhangzk committed
1118
            // the actual M and N   (no need to /2 if fuse_glu)
1119
1120
            int actualM, actualN;

sxtyzhangzk's avatar
sxtyzhangzk committed
1121
            bool alwaysfalse;
1122
1123
        };

Muyang Li's avatar
Muyang Li committed
1124
        __device__ __forceinline__ void operator()(Arguments args) {
1125
            const BlockInfo binfo = {
Muyang Li's avatar
Muyang Li committed
1126
1127
                .bm         = (int)blockIdx.x,
                .bn         = (int)blockIdx.y,
1128
1129
1130
1131
                .numBlocksM = (int)gridDim.x,
                .numBlocksN = (int)gridDim.y,
            };

Muyang Li's avatar
Muyang Li committed
1132
1133
            const int bm     = binfo.bm;
            const int bn     = binfo.bn;
1134
1135
1136
            const int warpId = threadIdx.x / WARP_SIZE;

            const int m_offset = bm * BLOCK_M + warpId * WARP_M;
sxtyzhangzk's avatar
sxtyzhangzk committed
1137
            const int n_offset = bn * BLOCK_N * (fuse_glu ? 2 : 1);
1138
1139
1140
1141
1142

            extern __shared__ uint8_t shmem[];

            fpsum_warp fpsum;

Muyang Li's avatar
Muyang Li committed
1143
1144
1145
1146
1147
1148
1149
            Base::template load_act_to_fpsum<fuse_glu>()(args.input + m_offset * args.actualN + n_offset,
                                                         args.actualN,
                                                         args.actualM - m_offset,
                                                         args.actualN - n_offset,
                                                         fpsum,
                                                         shmem + warpId * SHMEM_PER_WARP
                                                         // args.smooth_factor ? args.smooth_factor + n_offset : nullptr
1150
1151
            );

sxtyzhangzk's avatar
sxtyzhangzk committed
1152
1153
            CHECK_NAN(fpsum, "fpsum");
            // for (int i = 0; i < 16; i++) {
Muyang Li's avatar
Muyang Li committed
1154
            //     printf("bm=%d bn=%d warp=%d lane=%d fpsum[%d][0:1]=%f %f\n",
sxtyzhangzk's avatar
sxtyzhangzk committed
1155
1156
1157
1158
1159
1160
            //         bm, bn, warpId, threadIdx.x % WARP_SIZE, i,
            //         (float)fpsum[i].data[0].x, (float)fpsum[i].data[0].y);
            // }

            using EpilogueLoraDown = typename Lora<Config>::EpilogueLoraDown;

Muyang Li's avatar
Muyang Li committed
1161
1162
1163
1164
1165
1166
1167
1168
1169
1170
1171
1172
1173
1174
1175
1176
1177
1178
1179
1180
1181
1182
            EpilogueLoraDown()(binfo,
                               fpsum,
                               args.M,
                               args.N,
                               0,
                               typename EpilogueLoraDown::Arguments{
                                   .lora_wgt_down = args.lora_wgt_down,
                                   .lora_act      = args.lora_act,
                                   .rank          = args.lora_rank,
                                   .alwaysfalse   = args.alwaysfalse,
                               });

            EpilogueQuantize<false, false, use_fp4>()(
                binfo,
                fpsum,
                args.M,
                args.N,
                0,
                typename EpilogueQuantize<false, false, use_fp4>::Arguments{.qout          = args.output,
                                                                            .oscales       = args.oscales,
                                                                            .shift_value   = 0,
                                                                            .smooth_factor = args.smooth_factor});
1183
1184
        }
    };
muyangli's avatar
muyangli committed
1185
1186
};

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