A high-throughput and memory-efficient inference and serving engine for LLMs
Mirrored from https://github.com/tduncan-lum/Nvidia-Optimization-Stack---VLLM.git
Synced Nov 15, 11:00 UTC
[Kernel][Hardware][AMD] Bf16 mfma opt for ROCm skinny GEMMs (#17071)
Signed-off-by: Hashem Hashemi <hashem.hashemi@amd.com> Signed-off-by: charlifu <charlifu@amd.com> Co-authored-by: charlifu <charlifu@amd.com>
Hashem Hashemi committed
May 7, 2025 at 22:34 UTC
5a499e70d5a8b8d89b524e570e313ca93d8692cd
4 files changed
+316
-228
csrc/rocm/skinny_gemms.cu
+313
-226
@@ -126,8 +126,8 @@ __global__ void LLGemm1_kernel(const scalar_t* in_a, const scalar_t* in_b,
126
const int warp = threadIdx.x / WARP_SIZE;
127
const int lane = threadIdx.x % WARP_SIZE;
128
const int num_warps = blockDim.x / WARP_SIZE;
129
- const int qwarpid = threadid / num_warps;
130
- const int qthreadid = threadid % num_warps;
129
+ const int qwarpid = threadid / 16;
130
+ const int qthreadid = threadid % 16;
131
float4 rowA_elem4[NUM_A_ROWS_PER_BLOCK];
132
scalar2_t colB_elem4x, colB_elem4y, colB_elem4z, colB_elem4w;
133
float acc[NUM_A_ROWS_PER_BLOCK];
@@ -142,15 +142,13 @@ __global__ void LLGemm1_kernel(const scalar_t* in_a, const scalar_t* in_b,
142
// rowA_elem4[i] holds 8 * half numbers seen as a single float4.
143
rowA_elem4[i] = load_ntmprl(&af4[row_addr + threadid + K / 8 * i]);
144
}
145
+ colB_elem4x = bf4[threadid * 4 + 0];
146
+ colB_elem4y = bf4[threadid * 4 + 1];
147
+ colB_elem4z = bf4[threadid * 4 + 2];
148
+ colB_elem4w = bf4[threadid * 4 + 3];
149
}
150
147
- colB_elem4x = bf4[threadid * 4 + 0];
148
- colB_elem4y = bf4[threadid * 4 + 1];
149
- colB_elem4z = bf4[threadid * 4 + 2];
150
- colB_elem4w = bf4[threadid * 4 + 3];
151
-
151
scalar2_t Af2;
153
- [[maybe_unused]] scalar2_t Bf2;
152
float2 S;
153
154
auto Ah2ptr = reinterpret_cast<scalar2_t*>(&rowA_elem4);
@@ -193,12 +191,13 @@ __global__ void LLGemm1_kernel(const scalar_t* in_a, const scalar_t* in_b,
191
192
if (qwarpid < NUM_A_ROWS_PER_BLOCK) {
193
acc[qwarpid] = qthreadid < num_warps ? red_smem[qwarpid][qthreadid] : 0.f;
196
- for (int mask = num_warps / 2; mask >= 1; mask /= 2) {
194
+#pragma unroll
195
+ for (int mask = 16 / 2; mask >= 1; mask /= 2) {
196
acc[qwarpid] += __shfl_xor(acc[qwarpid], mask);
197
}
199
- float oval2 = __shfl_xor(acc[qwarpid], num_warps);
198
+ float oval2 = __shfl_xor(acc[qwarpid], 16);
199
201
- if (lane % (num_warps * 2) == 0) {
200
+ if (lane % 32 == 0) {
201
oval = __float22s2_rn<scalar2_t>(make_float2(acc[qwarpid], oval2));
202
c[blockIdx.x * NUM_A_ROWS_PER_BLOCK / 2 + qwarpid / 2] = oval;
203
}
@@ -222,9 +221,10 @@ torch::Tensor LLMM1(at::Tensor& in_a, at::Tensor& in_b,
221
// NUM_TREADS need to be a multiple of WARP_SIZE, as we are using warp shuffle
222
// operations.
223
const int NUM_THREADS =
225
- K * 2 / 16 % WARP_SIZE == 0
226
- ? K * 2 / 16
227
- : K * 2 / 16 + (WARP_SIZE - K * 2 / 16 % WARP_SIZE);
224
+ max(rows_per_block * 16,
225
+ K * 2 / 16 % WARP_SIZE == 0
226
+ ? K * 2 / 16
227
+ : K * 2 / 16 + (WARP_SIZE - K * 2 / 16 % WARP_SIZE));
228
229
int NUM_BLOCKS = M / rows_per_block;
230
@@ -275,13 +275,22 @@ __global__ void __launch_bounds__(WvPrGrp* THRDS)
275
wvSplitK_hf_sml_(const int K, const int M, const scalar_t* B,
276
const scalar_t* __restrict__ A, scalar_t* C,
277
const int _WvPrGrp, const int CuCount) {
278
+ #if defined(__HIP__MI300__)
279
+ constexpr bool use_mfma = (std::is_same_v<scalar_t, __hip_bfloat16>);
280
+ #else
281
+ constexpr bool use_mfma = false;
282
+ #endif
283
+
284
using scalar8 =
285
__attribute__((__vector_size__((A_CHUNK / 2) * sizeof(float)))) float;
286
+ using half4 =
287
+ __attribute__((__vector_size__((A_CHUNK / 2) * sizeof(__bf16)))) __bf16;
288
union bigType {
289
scalar_t h[A_CHUNK];
290
float f[A_CHUNK / 2];
291
float2 f2[A_CHUNK / 4];
292
double d[A_CHUNK / 4];
293
+ half4 h4[A_CHUNK / 4];
294
scalar8 h8;
295
};
296
@@ -318,6 +327,7 @@ __global__ void __launch_bounds__(WvPrGrp* THRDS)
327
uint32_t m = (blockIdx.x * _WvPrGrp + (threadIdx.y % _WvPrGrp)) * YTILE;
328
329
float sum[N][YTILE];
330
+ scalar8 sum4[N][YTILE];
331
332
//----------------------------------------------------
333
// Each wave works on a single column of weight matrix.
@@ -343,7 +353,11 @@ __global__ void __launch_bounds__(WvPrGrp* THRDS)
353
// are being worked on by each wave.
354
//----------------------------------------------------
355
for (int i = 0; i < YTILE; i++)
346
- for (int n = 0; n < N; n++) sum[n][i] = 0;
356
+ for (int n = 0; n < N; n++)
357
+ if constexpr (!use_mfma)
358
+ sum[n][i] = 0;
359
+ else
360
+ sum4[n][i] = {0, 0, 0, 0};
361
362
bigType bigA[N][UNRL];
363
bigType bigB[YTILE][UNRL];
@@ -374,24 +388,8 @@ __global__ void __launch_bounds__(WvPrGrp* THRDS)
388
if (k_ >= K) break;
389
390
const scalar_t* B_ = &B[(m + 0) * K + k_];
377
- bigB[0][k2].h8 = (loadnt((scalar8*)(&B_[0 * K])));
378
- //----------------------------------------------------
379
- // The following code with YTILE > 1 has to be deleted
380
- //----------------------------------------------------
381
- if constexpr (YTILE >= 2)
382
- bigB[1][k2].h8 = (loadnt((scalar8*)(&B_[1 * K])));
383
- if constexpr (YTILE >= 3)
384
- bigB[2][k2].h8 = (loadnt((scalar8*)(&B_[2 * K])));
385
- if constexpr (YTILE >= 4)
386
- bigB[3][k2].h8 = (loadnt((scalar8*)(&B_[3 * K])));
387
- if constexpr (YTILE >= 5)
388
- bigB[4][k2].h8 = (loadnt((scalar8*)(&B_[4 * K])));
389
- if constexpr (YTILE >= 6)
390
- bigB[5][k2].h8 = (loadnt((scalar8*)(&B_[5 * K])));
391
- if constexpr (YTILE >= 7)
392
- bigB[6][k2].h8 = (loadnt((scalar8*)(&B_[6 * K])));
393
- if constexpr (YTILE >= 8)
394
- bigB[7][k2].h8 = (loadnt((scalar8*)(&B_[7 * K])));
391
+ for (int y = 0; y < YTILE; y++)
392
+ bigB[y][k2].h8 = (loadnt((scalar8*)(&B_[y * K])));
393
}
394
395
// Fetch activation matrix from either just LDS or from both LDS / memory
@@ -419,32 +417,17 @@ __global__ void __launch_bounds__(WvPrGrp* THRDS)
417
#pragma unroll
418
for (uint32_t n = 0; n < N; n++) {
419
#pragma unroll
422
- for (uint32_t b = 0; b < A_CHUNK / 2; b++) {
423
- DOT2C(sum[n][0], bigA[n][k2].f[b], bigB[0][k2].f[b])
424
- //----------------------------------------------------
425
- // The following code with YTILE > 1
426
- //----------------------------------------------------
427
- if constexpr (YTILE >= 2) {
428
- DOT2C(sum[n][1], bigA[n][k2].f[b], bigB[1][k2].f[b]);
429
- }
430
- if constexpr (YTILE >= 3) {
431
- DOT2C(sum[n][2], bigA[n][k2].f[b], bigB[2][k2].f[b]);
432
- }
433
- if constexpr (YTILE >= 4) {
434
- DOT2C(sum[n][3], bigA[n][k2].f[b], bigB[3][k2].f[b]);
435
- }
436
- if constexpr (YTILE >= 5) {
437
- DOT2C(sum[n][4], bigA[n][k2].f[b], bigB[4][k2].f[b]);
438
- }
439
- if constexpr (YTILE >= 6) {
440
- DOT2C(sum[n][5], bigA[n][k2].f[b], bigB[5][k2].f[b]);
441
- }
442
- if constexpr (YTILE >= 7) {
443
- DOT2C(sum[n][6], bigA[n][k2].f[b], bigB[6][k2].f[b]);
444
- }
445
- if constexpr (YTILE >= 8) {
446
- DOT2C(sum[n][7], bigA[n][k2].f[b], bigB[7][k2].f[b]);
447
- }
420
+ for (int y = 0; y < YTILE; y++) {
421
+ if constexpr (!use_mfma)
422
+ #pragma unroll
423
+ for (uint32_t b = 0; b < A_CHUNK / 2; b++) {
424
+ DOT2C(sum[n][y], bigA[n][k2].f[b], bigB[y][k2].f[b])
425
+ }
426
+ else
427
+ #pragma unroll
428
+ for (uint32_t b = 0; b < A_CHUNK / 4; b++)
429
+ sum4[n][y] = __builtin_amdgcn_mfma_f32_4x4x4bf16_1k(
430
+ bigA[n][k2].h4[b], bigB[y][k2].h4[b], sum4[n][y], 0, 0, 0);
431
}
432
}
433
}
@@ -453,37 +436,84 @@ __global__ void __launch_bounds__(WvPrGrp* THRDS)
436
//----------------------------------------------------
437
// Final reduction step using shuffle
438
//----------------------------------------------------
456
- for (int n = 0; n < N; n++) {
457
- for (int y = 0; y < YTILE; y++) {
458
- asm("s_nop 0\n\tv_add_f32 %0, %2, %3 row_shr:8 bound_ctrl:0 "
459
- : "=v"(sum[n][y])
460
- : "0"(sum[n][y]), "v"(sum[n][y]), "v"(sum[n][y]));
461
- asm("s_nop 0\n\tv_add_f32 %0, %2, %3 row_shr:4 bound_ctrl:0 "
462
- : "=v"(sum[n][y])
463
- : "0"(sum[n][y]), "v"(sum[n][y]), "v"(sum[n][y]));
464
- asm("s_nop 0\n\tv_add_f32 %0, %2, %3 row_shr:2 bound_ctrl:0 "
465
- : "=v"(sum[n][y])
466
- : "0"(sum[n][y]), "v"(sum[n][y]), "v"(sum[n][y]));
467
- asm("s_nop 0\n\tv_add_f32 %0, %2, %3 wave_shr:1 bound_ctrl:0"
468
- : "=v"(sum[n][y])
469
- : "0"(sum[n][y]), "v"(sum[n][y]), "v"(sum[n][y]));
470
- asm("s_nop 0\n\tv_add_f32 %0, %2, %3 row_bcast:15 bound_ctrl:0"
471
- : "=v"(sum[n][y])
472
- : "0"(sum[n][y]), "v"(sum[n][y]), "v"(sum[n][y]));
473
- asm("s_nop 0\n\tv_add_f32 %0, %2, %3 row_bcast:31 bound_ctrl:0"
474
- : "=v"(sum[n][y])
475
- : "0"(sum[n][y]), "v"(sum[n][y]), "v"(sum[n][y]));
439
+ if constexpr (!use_mfma) {
440
+ for (int n = 0; n < N; n++) {
441
+ for (int y = 0; y < YTILE; y++) {
442
+ asm("s_nop 0\n\tv_add_f32 %0, %2, %3 row_shr:8 bound_ctrl:0 "
443
+ : "=v"(sum[n][y])
444
+ : "0"(sum[n][y]), "v"(sum[n][y]), "v"(sum[n][y]));
445
+ asm("s_nop 0\n\tv_add_f32 %0, %2, %3 row_shr:4 bound_ctrl:0 "
446
+ : "=v"(sum[n][y])
447
+ : "0"(sum[n][y]), "v"(sum[n][y]), "v"(sum[n][y]));
448
+ asm("s_nop 0\n\tv_add_f32 %0, %2, %3 row_shr:2 bound_ctrl:0 "
449
+ : "=v"(sum[n][y])
450
+ : "0"(sum[n][y]), "v"(sum[n][y]), "v"(sum[n][y]));
451
+ asm("s_nop 0\n\tv_add_f32 %0, %2, %3 wave_shr:1 bound_ctrl:0"
452
+ : "=v"(sum[n][y])
453
+ : "0"(sum[n][y]), "v"(sum[n][y]), "v"(sum[n][y]));
454
+ asm("s_nop 0\n\tv_add_f32 %0, %2, %3 row_bcast:15 bound_ctrl:0"
455
+ : "=v"(sum[n][y])
456
+ : "0"(sum[n][y]), "v"(sum[n][y]), "v"(sum[n][y]));
457
+ asm("s_nop 0\n\tv_add_f32 %0, %2, %3 row_bcast:31 bound_ctrl:0"
458
+ : "=v"(sum[n][y])
459
+ : "0"(sum[n][y]), "v"(sum[n][y]), "v"(sum[n][y]));
460
+ }
461
}
477
- }
478
- if (threadIdx.x == 63) {
462
+
463
+ if (threadIdx.x == 63) {
464
+ for (int n = 0; n < N; n++) {
465
+ for (int i = 0; i < YTILE; i++) {
466
+ // if (commitColumn[i]) C[m + i + n * M] = __float2half(sum[n][i]);
467
+ C[m + i + n * M] = __float2s<scalar_t>(sum[n][i]);
468
+ }
469
+ }
470
+ }
471
+ } else {
472
+ #pragma unroll
473
for (int n = 0; n < N; n++) {
480
- for (int i = 0; i < YTILE; i++) {
481
- // if (commitColumn[i]) C[m + i + n * M] = __float2half(sum[n][i]);
482
- C[m + i + n * M] = __float2s<scalar_t>(sum[n][i]);
474
+ #pragma unroll
475
+ for (int y = 0; y < YTILE; y++) {
476
+ // float accm1 = 0;
477
+ // for (int i=0; i<64; i++)
478
+ // accm1 += __shfl(sum4[n][y][i%4], i);
479
+ float accm = sum4[n][y][0];
480
+ asm("s_nop 0\n\tv_add_f32 %0, %2, %3 row_shl:1 bound_ctrl:0 "
481
+ : "=v"(accm)
482
+ : "0"(accm), "v"(sum4[n][y][1]), "v"(accm));
483
+ asm("s_nop 0\n\tv_add_f32 %0, %2, %3 row_shl:2 bound_ctrl:0 "
484
+ : "=v"(accm)
485
+ : "0"(accm), "v"(sum4[n][y][2]), "v"(accm));
486
+ asm("s_nop 0\n\tv_add_f32 %0, %2, %3 row_shl:3 bound_ctrl:0 "
487
+ : "=v"(accm)
488
+ : "0"(accm), "v"(sum4[n][y][3]), "v"(accm));
489
+ asm("s_nop 0\n\tv_add_f32 %0, %2, %3 row_shl:4 bound_ctrl:0 "
490
+ : "=v"(accm)
491
+ : "0"(accm), "v"(accm), "v"(accm));
492
+ asm("s_nop 0\n\tv_add_f32 %0, %2, %3 row_shl:8 bound_ctrl:0 "
493
+ : "=v"(accm)
494
+ : "0"(accm), "v"(accm), "v"(accm));
495
+ asm("s_nop 0\n\tv_mov_b32 %0, %2 row_shr:15 bound_ctrl:0 "
496
+ : "=v"(accm)
497
+ : "0"(accm), "v"(accm));
498
+ asm("s_nop 0\n\tv_add_f32 %0, %2, %3 row_bcast:15 bound_ctrl:0"
499
+ : "=v"(accm)
500
+ : "0"(accm), "v"(accm), "v"(accm));
501
+ asm("s_nop 0\n\tv_add_f32 %0, %2, %3 row_bcast:31 bound_ctrl:0"
502
+ : "=v"(accm)
503
+ : "0"(accm), "v"(accm), "v"(accm));
504
+
505
+ sum4[n][y][0] = accm;
506
+ }
507
+ }
508
+ if (threadIdx.x == 63) {
509
+ for (int n = 0; n < N; n++) {
510
+ for (int i = 0; i < YTILE; i++) {
511
+ // if (commitColumn[i]) C[n + i + m * N] = __float2half(sum[n][i]);
512
+ C[m + i + n * M] = __float2bfloat16(sum4[n][i][0]);
513
+ }
514
}
515
}
516
}
486
-
517
m += CuCount * _WvPrGrp * YTILE;
518
}
519
}
@@ -505,13 +535,22 @@ __global__ void __launch_bounds__(WvPrGrp* THRDS)
535
wvSplitK_hf_(const int K, const int M, const scalar_t* B,
536
const scalar_t* __restrict__ A, scalar_t* C,
537
const int _WvPrGrp, const int CuCount) {
538
+ #if defined(__HIP__MI300__)
539
+ constexpr bool use_mfma = (std::is_same_v<scalar_t, __hip_bfloat16>);
540
+ #else
541
+ constexpr bool use_mfma = false;
542
+ #endif
543
+
544
using scalar8 =
545
__attribute__((__vector_size__((A_CHUNK / 2) * sizeof(float)))) float;
546
+ using half4 =
547
+ __attribute__((__vector_size__((A_CHUNK / 2) * sizeof(__bf16)))) __bf16;
548
union bigType {
549
scalar_t h[A_CHUNK];
550
float f[A_CHUNK / 2];
551
float2 f2[A_CHUNK / 4];
552
double d[A_CHUNK / 4];
553
+ half4 h4[A_CHUNK / 4];
554
scalar8 h8;
555
};
556
@@ -573,6 +612,7 @@ __global__ void __launch_bounds__(WvPrGrp* THRDS)
612
if (threadIdx.y >= _WvPrGrp) return;
613
614
float sum[N][YTILE];
615
+ scalar8 sum4[N][YTILE];
616
617
//----------------------------------------------------
618
// Each wave works on a single column of weight matrix.
@@ -598,7 +638,11 @@ __global__ void __launch_bounds__(WvPrGrp* THRDS)
638
// are being worked on by each wave.
639
//----------------------------------------------------
640
for (int i = 0; i < YTILE; i++)
601
- for (int n = 0; n < N; n++) sum[n][i] = 0;
641
+ for (int n = 0; n < N; n++)
642
+ if constexpr (!use_mfma)
643
+ sum[n][i] = 0;
644
+ else
645
+ sum4[n][i] = {0, 0, 0, 0};
646
647
bigType bigA[N][UNRL];
648
bigType bigB[YTILE][UNRL];
@@ -628,24 +672,8 @@ __global__ void __launch_bounds__(WvPrGrp* THRDS)
672
if (k_ >= K) break;
673
674
const scalar_t* B_ = &B[(m + 0) * K + k_];
631
- bigB[0][k2].h8 = (loadnt((scalar8*)(&B_[0 * K])));
632
- //----------------------------------------------------
633
- // The following code with YTILE > 1 has to be deleted
634
- //----------------------------------------------------
635
- if constexpr (YTILE >= 2)
636
- bigB[1][k2].h8 = (loadnt((scalar8*)(&B_[1 * K])));
637
- if constexpr (YTILE >= 3)
638
- bigB[2][k2].h8 = (loadnt((scalar8*)(&B_[2 * K])));
639
- if constexpr (YTILE >= 4)
640
- bigB[3][k2].h8 = (loadnt((scalar8*)(&B_[3 * K])));
641
- if constexpr (YTILE >= 5)
642
- bigB[4][k2].h8 = (loadnt((scalar8*)(&B_[4 * K])));
643
- if constexpr (YTILE >= 6)
644
- bigB[5][k2].h8 = (loadnt((scalar8*)(&B_[5 * K])));
645
- if constexpr (YTILE >= 7)
646
- bigB[6][k2].h8 = (loadnt((scalar8*)(&B_[6 * K])));
647
- if constexpr (YTILE >= 8)
648
- bigB[7][k2].h8 = (loadnt((scalar8*)(&B_[7 * K])));
675
+ for (int b = 0; b < YTILE; b++)
676
+ bigB[b][k2].h8 = (loadnt((scalar8*)(&B_[b * K])));
677
}
678
679
// Fetch activation matrix from either just LDS or from both LDS / memory
@@ -676,32 +704,17 @@ __global__ void __launch_bounds__(WvPrGrp* THRDS)
704
// Do the matrix multiplication of activation and weight matrix
705
// - Remember the accumulation is happening for K-split of 64!
706
#pragma unroll
679
- for (uint32_t b = 0; b < A_CHUNK / 2; b++) {
680
- DOT2C(sum[n][0], bigA[n][k2].f[b], bigB[0][k2].f[b]);
681
- //----------------------------------------------------
682
- // The following code with YTILE > 1
683
- //----------------------------------------------------
684
- if constexpr (YTILE >= 2) {
685
- DOT2C(sum[n][1], bigA[n][k2].f[b], bigB[1][k2].f[b]);
686
- }
687
- if constexpr (YTILE >= 3) {
688
- DOT2C(sum[n][2], bigA[n][k2].f[b], bigB[2][k2].f[b]);
689
- }
690
- if constexpr (YTILE >= 4) {
691
- DOT2C(sum[n][3], bigA[n][k2].f[b], bigB[3][k2].f[b]);
692
- }
693
- if constexpr (YTILE >= 5) {
694
- DOT2C(sum[n][4], bigA[n][k2].f[b], bigB[4][k2].f[b]);
695
- }
696
- if constexpr (YTILE >= 6) {
697
- DOT2C(sum[n][5], bigA[n][k2].f[b], bigB[5][k2].f[b]);
698
- }
699
- if constexpr (YTILE >= 7) {
700
- DOT2C(sum[n][6], bigA[n][k2].f[b], bigB[6][k2].f[b]);
701
- }
702
- if constexpr (YTILE >= 8) {
703
- DOT2C(sum[n][7], bigA[n][k2].f[b], bigB[7][k2].f[b]);
704
- }
707
+ for (int y = 0; y < YTILE; y++) {
708
+ if constexpr (!use_mfma)
709
+ #pragma unroll
710
+ for (uint32_t b = 0; b < A_CHUNK / 2; b++) {
711
+ DOT2C(sum[n][y], bigA[n][k2].f[b], bigB[y][k2].f[b])
712
+ }
713
+ else
714
+ #pragma unroll
715
+ for (uint32_t b = 0; b < A_CHUNK / 4; b++)
716
+ sum4[n][y] = __builtin_amdgcn_mfma_f32_4x4x4bf16_1k(
717
+ bigA[n][k2].h4[b], bigB[y][k2].h4[b], sum4[n][y], 0, 0, 0);
718
}
719
}
720
}
@@ -710,34 +723,82 @@ __global__ void __launch_bounds__(WvPrGrp* THRDS)
723
//----------------------------------------------------
724
// Final reduction step using shuffle
725
//----------------------------------------------------
713
- for (int n = 0; n < N; n++) {
714
- for (int y = 0; y < YTILE; y++) {
715
- asm("s_nop 0\n\tv_add_f32 %0, %2, %3 row_shr:8 bound_ctrl:0 "
716
- : "=v"(sum[n][y])
717
- : "0"(sum[n][y]), "v"(sum[n][y]), "v"(sum[n][y]));
718
- asm("s_nop 0\n\tv_add_f32 %0, %2, %3 row_shr:4 bound_ctrl:0 "
719
- : "=v"(sum[n][y])
720
- : "0"(sum[n][y]), "v"(sum[n][y]), "v"(sum[n][y]));
721
- asm("s_nop 0\n\tv_add_f32 %0, %2, %3 row_shr:2 bound_ctrl:0 "
722
- : "=v"(sum[n][y])
723
- : "0"(sum[n][y]), "v"(sum[n][y]), "v"(sum[n][y]));
724
- asm("s_nop 0\n\tv_add_f32 %0, %2, %3 wave_shr:1 bound_ctrl:0"
725
- : "=v"(sum[n][y])
726
- : "0"(sum[n][y]), "v"(sum[n][y]), "v"(sum[n][y]));
727
- asm("s_nop 0\n\tv_add_f32 %0, %2, %3 row_bcast:15 bound_ctrl:0"
728
- : "=v"(sum[n][y])
729
- : "0"(sum[n][y]), "v"(sum[n][y]), "v"(sum[n][y]));
730
- asm("s_nop 0\n\tv_add_f32 %0, %2, %3 row_bcast:31 bound_ctrl:0"
731
- : "=v"(sum[n][y])
732
- : "0"(sum[n][y]), "v"(sum[n][y]), "v"(sum[n][y]));
726
+ if constexpr (!use_mfma) {
727
+ for (int n = 0; n < N; n++) {
728
+ for (int y = 0; y < YTILE; y++) {
729
+ asm("s_nop 0\n\tv_add_f32 %0, %2, %3 row_shr:8 bound_ctrl:0 "
730
+ : "=v"(sum[n][y])
731
+ : "0"(sum[n][y]), "v"(sum[n][y]), "v"(sum[n][y]));
732
+ asm("s_nop 0\n\tv_add_f32 %0, %2, %3 row_shr:4 bound_ctrl:0 "
733
+ : "=v"(sum[n][y])
734
+ : "0"(sum[n][y]), "v"(sum[n][y]), "v"(sum[n][y]));
735
+ asm("s_nop 0\n\tv_add_f32 %0, %2, %3 row_shr:2 bound_ctrl:0 "
736
+ : "=v"(sum[n][y])
737
+ : "0"(sum[n][y]), "v"(sum[n][y]), "v"(sum[n][y]));
738
+ asm("s_nop 0\n\tv_add_f32 %0, %2, %3 wave_shr:1 bound_ctrl:0"
739
+ : "=v"(sum[n][y])
740
+ : "0"(sum[n][y]), "v"(sum[n][y]), "v"(sum[n][y]));
741
+ asm("s_nop 0\n\tv_add_f32 %0, %2, %3 row_bcast:15 bound_ctrl:0"
742
+ : "=v"(sum[n][y])
743
+ : "0"(sum[n][y]), "v"(sum[n][y]), "v"(sum[n][y]));
744
+ asm("s_nop 0\n\tv_add_f32 %0, %2, %3 row_bcast:31 bound_ctrl:0"
745
+ : "=v"(sum[n][y])
746
+ : "0"(sum[n][y]), "v"(sum[n][y]), "v"(sum[n][y]));
747
+ }
748
}
734
- }
749
736
- if (threadIdx.x == 63) {
750
+ if (threadIdx.x == 63) {
751
+ for (int n = 0; n < N; n++) {
752
+ for (int i = 0; i < YTILE; i++) {
753
+ if (commitColumn[i])
754
+ C[m + i + n * M] = __float2s<scalar_t>(sum[n][i]);
755
+ }
756
+ }
757
+ }
758
+ } else {
759
+ #pragma unroll
760
for (int n = 0; n < N; n++) {
738
- for (int i = 0; i < YTILE; i++) {
739
- if (commitColumn[i])
740
- C[m + i + n * M] = __float2s<scalar_t>(sum[n][i]);
761
+ #pragma unroll
762
+ for (int y = 0; y < YTILE; y++) {
763
+ // float accm1 = 0;
764
+ // for (int i=0; i<64; i++)
765
+ // accm1 += __shfl(sum4[n][y][i%4], i);
766
+
767
+ float accm = sum4[n][y][0];
768
+ asm("s_nop 0\n\tv_add_f32 %0, %2, %3 row_shl:1 bound_ctrl:0 "
769
+ : "=v"(accm)
770
+ : "0"(accm), "v"(sum4[n][y][1]), "v"(accm));
771
+ asm("s_nop 0\n\tv_add_f32 %0, %2, %3 row_shl:2 bound_ctrl:0 "
772
+ : "=v"(accm)
773
+ : "0"(accm), "v"(sum4[n][y][2]), "v"(accm));
774
+ asm("s_nop 0\n\tv_add_f32 %0, %2, %3 row_shl:3 bound_ctrl:0 "
775
+ : "=v"(accm)
776
+ : "0"(accm), "v"(sum4[n][y][3]), "v"(accm));
777
+ asm("s_nop 0\n\tv_add_f32 %0, %2, %3 row_shl:4 bound_ctrl:0 "
778
+ : "=v"(accm)
779
+ : "0"(accm), "v"(accm), "v"(accm));
780
+ asm("s_nop 0\n\tv_add_f32 %0, %2, %3 row_shl:8 bound_ctrl:0 "
781
+ : "=v"(accm)
782
+ : "0"(accm), "v"(accm), "v"(accm));
783
+ asm("s_nop 0\n\tv_mov_b32 %0, %2 row_shr:15 bound_ctrl:0 "
784
+ : "=v"(accm)
785
+ : "0"(accm), "v"(accm));
786
+ asm("s_nop 0\n\tv_add_f32 %0, %2, %3 row_bcast:15 bound_ctrl:0"
787
+ : "=v"(accm)
788
+ : "0"(accm), "v"(accm), "v"(accm));
789
+ asm("s_nop 0\n\tv_add_f32 %0, %2, %3 row_bcast:31 bound_ctrl:0"
790
+ : "=v"(accm)
791
+ : "0"(accm), "v"(accm), "v"(accm));
792
+
793
+ sum4[n][y][0] = accm;
794
+ }
795
+ }
796
+ if (threadIdx.x == 63) {
797
+ for (int n = 0; n < N; n++) {
798
+ for (int i = 0; i < YTILE; i++) {
799
+ // if (commitColumn[i]) C[n + i + m * N] = __float2half(sum[n][i]);
800
+ C[m + i + n * M] = __float2bfloat16(sum4[n][i][0]);
801
+ }
802
}
803
}
804
}
@@ -774,14 +835,22 @@ __global__ void __launch_bounds__(WvPrGrp* THRDS)
835
wvSplitK_hf_big_(const int K, const int M, const scalar_t* B,
836
const scalar_t* __restrict__ A, scalar_t* C,
837
const int _WvPrGrp, const int CuCount) {
838
+ #if defined(__HIP__MI300__)
839
+ constexpr bool use_mfma = (std::is_same_v<scalar_t, __hip_bfloat16>);
840
+ #else
841
+ constexpr bool use_mfma = false;
842
+ #endif
843
+
844
using scalar8 =
845
__attribute__((__vector_size__((A_CHUNK / 2) * sizeof(float)))) float;
779
-
846
+ using half4 =
847
+ __attribute__((__vector_size__((A_CHUNK / 2) * sizeof(__bf16)))) __bf16;
848
union bigType {
849
scalar_t h[A_CHUNK];
850
float f[A_CHUNK / 2];
851
float2 f2[A_CHUNK / 4];
852
double d[A_CHUNK / 4];
853
+ half4 h4[A_CHUNK / 4];
854
scalar8 h8;
855
};
856
@@ -857,6 +926,7 @@ __global__ void __launch_bounds__(WvPrGrp* THRDS)
926
kFit = min(kFit, K);
927
928
float sum[N][YTILE];
929
+ scalar8 sum4[N][YTILE];
930
931
//----------------------------------------------------
932
// Each wave works on a single column of weight matrix.
@@ -888,7 +958,11 @@ __global__ void __launch_bounds__(WvPrGrp* THRDS)
958
// are being worked on by each wave.
959
//----------------------------------------------------
960
for (int i = 0; i < YTILE; i++)
891
- for (int n = 0; n < N; n++) sum[n][i] = 0;
961
+ for (int n = 0; n < N; n++)
962
+ if constexpr (!use_mfma)
963
+ sum[n][i] = 0;
964
+ else
965
+ sum4[n][i] = {0, 0, 0, 0};
966
967
bigType bigA[N][UNRL];
968
bigType bigB[YTILE][UNRL];
@@ -937,24 +1011,8 @@ __global__ void __launch_bounds__(WvPrGrp* THRDS)
1011
if (k_ >= K) break;
1012
1013
const scalar_t* B_ = &B[(m + 0) * K + k_];
940
- bigB[0][k2].h8 = (loadnt((scalar8*)(&B_[0 * K])));
941
- //----------------------------------------------------
942
- // The following code with YTILE > 1 has to be deleted
943
- //----------------------------------------------------
944
- if constexpr (YTILE >= 2)
945
- bigB[1][k2].h8 = (loadnt((scalar8*)(&B_[1 * K])));
946
- if constexpr (YTILE >= 3)
947
- bigB[2][k2].h8 = (loadnt((scalar8*)(&B_[2 * K])));
948
- if constexpr (YTILE >= 4)
949
- bigB[3][k2].h8 = (loadnt((scalar8*)(&B_[3 * K])));
950
- if constexpr (YTILE >= 5)
951
- bigB[4][k2].h8 = (loadnt((scalar8*)(&B_[4 * K])));
952
- if constexpr (YTILE >= 6)
953
- bigB[5][k2].h8 = (loadnt((scalar8*)(&B_[5 * K])));
954
- if constexpr (YTILE >= 7)
955
- bigB[6][k2].h8 = (loadnt((scalar8*)(&B_[6 * K])));
956
- if constexpr (YTILE >= 8)
957
- bigB[7][k2].h8 = (loadnt((scalar8*)(&B_[7 * K])));
1014
+ for (int b = 0; b < YTILE; b++)
1015
+ bigB[b][k2].h8 = (loadnt((scalar8*)(&B_[b * K])));
1016
}
1017
1018
// Fetch activation matrix from either just LDS or from both LDS / memory
@@ -989,32 +1047,17 @@ __global__ void __launch_bounds__(WvPrGrp* THRDS)
1047
// Do the matrix multiplication of activation and weight matrix
1048
// - Remember the accumulation is happening for K-split of 64!
1049
#pragma unroll
992
- for (uint32_t b = 0; b < A_CHUNK / 2; b++) {
993
- DOT2C(sum[n][0], bigA[n][k2].f[b], bigB[0][k2].f[b]);
994
- //----------------------------------------------------
995
- // The following code with YTILE > 1
996
- //----------------------------------------------------
997
- if constexpr (YTILE >= 2) {
998
- DOT2C(sum[n][1], bigA[n][k2].f[b], bigB[1][k2].f[b]);
999
- }
1000
- if constexpr (YTILE >= 3) {
1001
- DOT2C(sum[n][2], bigA[n][k2].f[b], bigB[2][k2].f[b]);
1002
- }
1003
- if constexpr (YTILE >= 4) {
1004
- DOT2C(sum[n][3], bigA[n][k2].f[b], bigB[3][k2].f[b]);
1005
- }
1006
- if constexpr (YTILE >= 5) {
1007
- DOT2C(sum[n][4], bigA[n][k2].f[b], bigB[4][k2].f[b]);
1008
- }
1009
- if constexpr (YTILE >= 6) {
1010
- DOT2C(sum[n][5], bigA[n][k2].f[b], bigB[5][k2].f[b]);
1011
- }
1012
- if constexpr (YTILE >= 7) {
1013
- DOT2C(sum[n][6], bigA[n][k2].f[b], bigB[6][k2].f[b]);
1014
- }
1015
- if constexpr (YTILE >= 8) {
1016
- DOT2C(sum[n][7], bigA[n][k2].f[b], bigB[7][k2].f[b]);
1017
- }
1050
+ for (int y = 0; y < YTILE; y++) {
1051
+ if constexpr (!use_mfma)
1052
+ #pragma unroll
1053
+ for (uint32_t b = 0; b < A_CHUNK / 2; b++) {
1054
+ DOT2C(sum[n][y], bigA[n][k2].f[b], bigB[y][k2].f[b])
1055
+ }
1056
+ else
1057
+ #pragma unroll
1058
+ for (uint32_t b = 0; b < A_CHUNK / 4; b++)
1059
+ sum4[n][y] = __builtin_amdgcn_mfma_f32_4x4x4bf16_1k(
1060
+ bigA[n][k2].h4[b], bigB[y][k2].h4[b], sum4[n][y], 0, 0, 0);
1061
}
1062
}
1063
}
@@ -1031,34 +1074,78 @@ __global__ void __launch_bounds__(WvPrGrp* THRDS)
1074
//----------------------------------------------------
1075
// Final reduction step using shuffle
1076
//----------------------------------------------------
1034
- for (int n = 0; n < N; n++) {
1035
- for (int y = 0; y < YTILE; y++) {
1036
- asm("s_nop 0\n\tv_add_f32 %0, %2, %3 row_shr:8 bound_ctrl:0 "
1037
- : "=v"(sum[n][y])
1038
- : "0"(sum[n][y]), "v"(sum[n][y]), "v"(sum[n][y]));
1039
- asm("s_nop 0\n\tv_add_f32 %0, %2, %3 row_shr:4 bound_ctrl:0 "
1040
- : "=v"(sum[n][y])
1041
- : "0"(sum[n][y]), "v"(sum[n][y]), "v"(sum[n][y]));
1042
- asm("s_nop 0\n\tv_add_f32 %0, %2, %3 row_shr:2 bound_ctrl:0 "
1043
- : "=v"(sum[n][y])
1044
- : "0"(sum[n][y]), "v"(sum[n][y]), "v"(sum[n][y]));
1045
- asm("s_nop 0\n\tv_add_f32 %0, %2, %3 wave_shr:1 bound_ctrl:0"
1046
- : "=v"(sum[n][y])
1047
- : "0"(sum[n][y]), "v"(sum[n][y]), "v"(sum[n][y]));
1048
- asm("s_nop 0\n\tv_add_f32 %0, %2, %3 row_bcast:15 bound_ctrl:0"
1049
- : "=v"(sum[n][y])
1050
- : "0"(sum[n][y]), "v"(sum[n][y]), "v"(sum[n][y]));
1051
- asm("s_nop 0\n\tv_add_f32 %0, %2, %3 row_bcast:31 bound_ctrl:0"
1052
- : "=v"(sum[n][y])
1053
- : "0"(sum[n][y]), "v"(sum[n][y]), "v"(sum[n][y]));
1077
+ if constexpr (!use_mfma) {
1078
+ for (int n = 0; n < N; n++) {
1079
+ for (int y = 0; y < YTILE; y++) {
1080
+ asm("s_nop 0\n\tv_add_f32 %0, %2, %3 row_shr:8 bound_ctrl:0 "
1081
+ : "=v"(sum[n][y])
1082
+ : "0"(sum[n][y]), "v"(sum[n][y]), "v"(sum[n][y]));
1083
+ asm("s_nop 0\n\tv_add_f32 %0, %2, %3 row_shr:4 bound_ctrl:0 "
1084
+ : "=v"(sum[n][y])
1085
+ : "0"(sum[n][y]), "v"(sum[n][y]), "v"(sum[n][y]));
1086
+ asm("s_nop 0\n\tv_add_f32 %0, %2, %3 row_shr:2 bound_ctrl:0 "
1087
+ : "=v"(sum[n][y])
1088
+ : "0"(sum[n][y]), "v"(sum[n][y]), "v"(sum[n][y]));
1089
+ asm("s_nop 0\n\tv_add_f32 %0, %2, %3 wave_shr:1 bound_ctrl:0"
1090
+ : "=v"(sum[n][y])
1091
+ : "0"(sum[n][y]), "v"(sum[n][y]), "v"(sum[n][y]));
1092
+ asm("s_nop 0\n\tv_add_f32 %0, %2, %3 row_bcast:15 bound_ctrl:0"
1093
+ : "=v"(sum[n][y])
1094
+ : "0"(sum[n][y]), "v"(sum[n][y]), "v"(sum[n][y]));
1095
+ asm("s_nop 0\n\tv_add_f32 %0, %2, %3 row_bcast:31 bound_ctrl:0"
1096
+ : "=v"(sum[n][y])
1097
+ : "0"(sum[n][y]), "v"(sum[n][y]), "v"(sum[n][y]));
1098
+ }
1099
}
1055
- }
1100
1057
- if (threadIdx.x == 63) {
1101
+ if (threadIdx.x == 63) {
1102
+ for (int n = 0; n < N; n++) {
1103
+ for (int i = 0; i < YTILE; i++) {
1104
+ if (commitColumn[i])
1105
+ C[m + i + n * M] = __float2s<scalar_t>(sum[n][i]);
1106
+ }
1107
+ }
1108
+ }
1109
+ } else {
1110
+ #pragma unroll
1111
for (int n = 0; n < N; n++) {
1059
- for (int i = 0; i < YTILE; i++) {
1060
- if (commitColumn[i])
1061
- C[m + i + n * M] = __float2s<scalar_t>(sum[n][i]);
1112
+ #pragma unroll
1113
+ for (int y = 0; y < YTILE; y++) {
1114
+ float accm = sum4[n][y][0];
1115
+ asm("s_nop 0\n\tv_add_f32 %0, %2, %3 row_shl:1 bound_ctrl:0 "
1116
+ : "=v"(accm)
1117
+ : "0"(accm), "v"(sum4[n][y][1]), "v"(accm));
1118
+ asm("s_nop 0\n\tv_add_f32 %0, %2, %3 row_shl:2 bound_ctrl:0 "
1119
+ : "=v"(accm)
1120
+ : "0"(accm), "v"(sum4[n][y][2]), "v"(accm));
1121
+ asm("s_nop 0\n\tv_add_f32 %0, %2, %3 row_shl:3 bound_ctrl:0 "
1122
+ : "=v"(accm)
1123
+ : "0"(accm), "v"(sum4[n][y][3]), "v"(accm));
1124
+ asm("s_nop 0\n\tv_add_f32 %0, %2, %3 row_shl:4 bound_ctrl:0 "
1125
+ : "=v"(accm)
1126
+ : "0"(accm), "v"(accm), "v"(accm));
1127
+ asm("s_nop 0\n\tv_add_f32 %0, %2, %3 row_shl:8 bound_ctrl:0 "
1128
+ : "=v"(accm)
1129
+ : "0"(accm), "v"(accm), "v"(accm));
1130
+ asm("s_nop 0\n\tv_mov_b32 %0, %2 row_shr:15 bound_ctrl:0 "
1131
+ : "=v"(accm)
1132
+ : "0"(accm), "v"(accm));
1133
+ asm("s_nop 0\n\tv_add_f32 %0, %2, %3 row_bcast:15 bound_ctrl:0"
1134
+ : "=v"(accm)
1135
+ : "0"(accm), "v"(accm), "v"(accm));
1136
+ asm("s_nop 0\n\tv_add_f32 %0, %2, %3 row_bcast:31 bound_ctrl:0"
1137
+ : "=v"(accm)
1138
+ : "0"(accm), "v"(accm), "v"(accm));
1139
+
1140
+ sum4[n][y][0] = accm;
1141
+ }
1142
+ }
1143
+ if (threadIdx.x == 63) {
1144
+ for (int n = 0; n < N; n++) {
1145
+ for (int i = 0; i < YTILE; i++) {
1146
+ // if (commitColumn[i]) C[n + i + m * N] = __float2half(sum[n][i]);
1147
+ C[m + i + n * M] = __float2bfloat16(sum4[n][i][0]);
1148
+ }
1149
}
1150
}
1151
}
tests/kernels/quantization/test_rocm_skinny_gemms.py
+1
-1
@@ -8,7 +8,7 @@ from vllm.platforms import current_platform
8
9
DTYPES = [torch.bfloat16, torch.float16]
10
M = [16, 32, 64, 128, 256, 512, 1024, 4096, 8192]
11
-K = [8, 16, 32, 64, 128, 256, 512, 1024, 2048, 4096, 8192] # k % 8 == 0
11
+K = [8, 16, 32, 64, 128, 256, 512, 1024, 2048, 4096, 6144, 8192] # k % 8 == 0
12
N = [1, 2, 3, 4]
13
SEEDS = [0]
14
vllm/model_executor/layers/utils.py
+1
-1
@@ -84,7 +84,7 @@ def rocm_unquantized_gemm(x: torch.Tensor,
84
m = weight.shape[0]
85
cu_count = current_platform.get_cu_count()
86
87
- if m > 8 and 0 < n < 4:
87
+ if m > 8 and 0 < n <= 4:
88
out = ops.wvSplitK(weight, x_view, cu_count)
89
return out.view(*x.shape[:-1], weight.shape[0])
90
elif m % 4 == 0 and n == 1 and k <= 8192:
vllm/platforms/rocm.py
+1
@@ -104,6 +104,7 @@ def device_id_to_physical_device_id(device_id: int) -> int:
104
return device_id
105
106
107
+@cache
108
def on_mi250_mi300() -> bool:
109
GPU_ARCH = torch.cuda.get_device_properties("cuda").gcnArchName
110
return any(arch in GPU_ARCH for arch in ["gfx90a", "gfx942"])