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 7, 12:45 UTC
Add padding support to wvSplitK solution for skinny GEMMs (#33762)
Signed-off-by: Hashem Hashemi <hashem.hashemi@amd.com>
Hashem Hashemi committed
Feb 28, 2026 at 01:02 UTC
7600642eaead7454fd977dde3513682244109e7c
3 files changed
+289
-444
csrc/rocm/skinny_gemms.cu
+256
-402
@@ -304,8 +304,9 @@ __device__ inline unsigned int min__(uint32_t a, uint32_t b) {
304
template <typename scalar_t, int THRDS, int YTILE, int WvPrGrp, int A_CHUNK,
305
int UNRL, int N>
306
__global__ void __launch_bounds__(WvPrGrp* THRDS)
307
- wvSplitK_hf_sml_(const int K, const int M, const int Bx, const int By,
308
- const scalar_t* B, const scalar_t* __restrict__ A,
307
+ wvSplitK_hf_sml_(const int K, const int Kbp, const int Kap, const int M,
308
+ const int Bx, const int By, const scalar_t* B,
309
+ const scalar_t* __restrict__ A,
310
const scalar_t* __restrict__ BIAS, scalar_t* C,
311
const int _WvPrGrp, const int CuCount) {
312
constexpr int max_lds_len = LDS_SIZE / 2;
@@ -314,7 +315,6 @@ __global__ void __launch_bounds__(WvPrGrp* THRDS)
315
#else
316
constexpr bool use_mfma = false;
317
#endif
317
-
318
using scalar8 =
319
__attribute__((__vector_size__((A_CHUNK / 2) * sizeof(float)))) float;
320
using half4 =
@@ -346,13 +346,13 @@ __global__ void __launch_bounds__(WvPrGrp* THRDS)
346
// - Then the WG will move to another 8 K elements
347
// TODO: Logic below will only work when K is multiple of 8
348
//----------------------------------------------------
349
- for (uint32_t k = 0; k < min__(K * N, max_lds_len);
350
- k += THRDS * WvPrGrp * A_CHUNK) {
351
- uint32_t k_in = k + ((threadIdx.y * THRDS + threadIdx.x) * A_CHUNK);
352
-
353
- if (k_in >= min__(K * N, max_lds_len)) break;
354
-
355
- *((bigType*)(&s[k_in])) = *((bigType*)(&A[k_in]));
349
+ for (uint32_t k = (threadIdx.y * THRDS + threadIdx.x) * A_CHUNK;
350
+ k < min__(Kap * N, max_lds_len); k += THRDS * WvPrGrp * A_CHUNK) {
351
+ #if defined(__gfx950__)
352
+ __builtin_amdgcn_global_load_lds((int*)(&A[k]), (int*)(&s[k]), 16, 0, 0);
353
+ #else
354
+ *((bigType*)(&s[k])) = *((bigType*)(&A[k]));
355
+ #endif
356
}
357
__syncthreads();
358
@@ -360,9 +360,6 @@ __global__ void __launch_bounds__(WvPrGrp* THRDS)
360
361
uint32_t m = (blockIdx.x * _WvPrGrp + (threadIdx.y % _WvPrGrp)) * YTILE;
362
363
- float sum[N][YTILE];
364
- scalar8 sum4[N][YTILE];
365
-
363
//----------------------------------------------------
364
// Each wave works on a single column of weight matrix.
365
// There are 16 waves per WG, and hence, each WG is
@@ -386,44 +383,20 @@ __global__ void __launch_bounds__(WvPrGrp* THRDS)
383
// YTILE represents how many column of weight matrix
384
// are being worked on by each wave.
385
//----------------------------------------------------
389
- for (int i = 0; i < YTILE; i++)
390
- for (int n = 0; n < N; n++)
391
- if constexpr (!use_mfma)
392
- sum[n][i] = 0;
393
- else
394
- sum4[n][i] = {0, 0, 0, 0};
395
-
396
- bigType bigA[N][UNRL];
397
- bigType bigB[YTILE][UNRL];
398
- //----------------------------------------------------
399
- // Fetch weight matrix B in interleaved K-split!
400
- // - Each thread (lane) is fetching 8 elements (A_Chunk)
401
- // - Each wave will fetch 64*8=> 512 elements (1024B)
402
- // - YTILE represents the number of column being serviced
403
- // by wave
404
- // - Loop for fetching weight matrix (B) are unrolled
405
- //
406
- // Fetch activation matrix A from LDS
407
- // - Loop for fetching activation matrix (A) are unrolled
408
- //
409
- // Finally, do the matrix multiplication in an unrolled
410
- // fashion. This provides lot of food for compiler
411
- // scheduling.
412
- //
413
- // TODO: Logic below will only work when K is multiple of 8
414
- //----------------------------------------------------
415
- // for (uint32_t k1 = 0; k1 < K; k1 += THRDS * A_CHUNK * UNRL) {
386
+ float sum[N][YTILE] = {};
387
+ scalar8 sum4[N][YTILE] = {};
388
+
389
for (uint32_t k1 = 0; k1 < K; k1 += THRDS * A_CHUNK * UNRL) {
390
+ bigType bigA[N][UNRL] = {};
391
+ bigType bigB[YTILE][UNRL];
392
// Fetch the weight matrix from memory!
393
#pragma unroll
394
for (uint32_t k2 = 0; k2 < UNRL; k2++) {
395
uint32_t k = k1 + k2 * THRDS * A_CHUNK;
396
uint32_t k_ = k + threadIdx.x * A_CHUNK;
422
- if (k_ >= K) break;
423
-
424
- const scalar_t* B_ = &B[(m + 0) * K + k_];
397
+ const scalar_t* B_ = &B[min__(k_, K - A_CHUNK)];
398
for (int y = 0; y < YTILE; y++)
426
- bigB[y][k2].h8 = (loadnt((scalar8*)(&B_[y * K])));
399
+ bigB[y][k2].h8 = (loadnt((scalar8*)(&B_[min__(y + m, M - 1) * Kbp])));
400
}
401
402
// Fetch activation matrix from either just LDS or from both LDS / memory
@@ -432,33 +405,20 @@ __global__ void __launch_bounds__(WvPrGrp* THRDS)
405
uint32_t k = k1 + k2 * THRDS * A_CHUNK;
406
uint32_t k_ = k + threadIdx.x * A_CHUNK;
407
if (k_ >= K) break;
435
-
436
- // Fetch A activation matrix in interleaved fashion from LDS or memory
437
-
408
for (int n = 0; n < N; n++) {
439
- bigA[n][k2] = *((const bigType*)(&(s[k_ + K * n])));
409
+ bigA[n][k2] = *((const bigType*)(&(s[k_ + Kap * n])));
410
}
411
}
412
413
// Do the matrix multiplication in interleaved manner
444
- #pragma unroll
414
for (uint32_t k2 = 0; k2 < UNRL; k2++) {
446
- uint32_t k = k1 + k2 * THRDS * A_CHUNK;
447
- uint32_t k_ = k + threadIdx.x * A_CHUNK;
448
- if (k_ >= K) break;
449
- // Do the matrix multiplication of activation and weight matrix
450
- // - Remember the accumulation is happening for K-split of 64!
451
- #pragma unroll
415
for (uint32_t n = 0; n < N; n++) {
453
- #pragma unroll
416
for (int y = 0; y < YTILE; y++) {
417
if constexpr (!use_mfma)
456
- #pragma unroll
418
for (uint32_t b = 0; b < A_CHUNK / 2; b++) {
419
DOT2C(sum[n][y], bigA[n][k2].f[b], bigB[y][k2].f[b])
420
}
421
else
461
- #pragma unroll
422
for (uint32_t b = 0; b < A_CHUNK / 4; b++)
423
sum4[n][y] = __builtin_amdgcn_mfma_f32_4x4x4bf16_1k(
424
bigA[n][k2].h4[b], bigB[y][k2].h4[b], sum4[n][y], 0, 0, 0);
@@ -466,46 +426,44 @@ __global__ void __launch_bounds__(WvPrGrp* THRDS)
426
}
427
}
428
}
469
-
429
+ __builtin_amdgcn_sched_barrier(0);
430
//----------------------------------------------------
431
// Final reduction step using shuffle
432
//----------------------------------------------------
433
if constexpr (!use_mfma) {
434
for (int n = 0; n < N; n++) {
435
for (int y = 0; y < YTILE; y++) {
476
- asm("s_nop 0\n\tv_add_f32 %0, %2, %3 row_shr:8 bound_ctrl:0 "
477
- : "=v"(sum[n][y])
478
- : "0"(sum[n][y]), "v"(sum[n][y]), "v"(sum[n][y]));
479
- asm("s_nop 0\n\tv_add_f32 %0, %2, %3 row_shr:4 bound_ctrl:0 "
480
- : "=v"(sum[n][y])
481
- : "0"(sum[n][y]), "v"(sum[n][y]), "v"(sum[n][y]));
482
- asm("s_nop 0\n\tv_add_f32 %0, %2, %3 row_shr:2 bound_ctrl:0 "
483
- : "=v"(sum[n][y])
484
- : "0"(sum[n][y]), "v"(sum[n][y]), "v"(sum[n][y]));
485
- asm("s_nop 0\n\tv_add_f32 %0, %2, %3 wave_shr:1 bound_ctrl:0"
486
- : "=v"(sum[n][y])
487
- : "0"(sum[n][y]), "v"(sum[n][y]), "v"(sum[n][y]));
488
- asm("s_nop 0\n\tv_add_f32 %0, %2, %3 row_bcast:15 bound_ctrl:0"
489
- : "=v"(sum[n][y])
490
- : "0"(sum[n][y]), "v"(sum[n][y]), "v"(sum[n][y]));
491
- asm("s_nop 0\n\tv_add_f32 %0, %2, %3 row_bcast:31 bound_ctrl:0"
492
- : "=v"(sum[n][y])
493
- : "0"(sum[n][y]), "v"(sum[n][y]), "v"(sum[n][y]));
436
+ sum[n][y] += __builtin_amdgcn_mov_dpp(sum[n][y], 0x118, 0xf, 0xf,
437
+ 1); // row_shr8
438
+ sum[n][y] += __builtin_amdgcn_mov_dpp(sum[n][y], 0x114, 0xf, 0xf,
439
+ 1); // row_shr4
440
+ sum[n][y] += __builtin_amdgcn_mov_dpp(sum[n][y], 0x112, 0xf, 0xf,
441
+ 1); // row_shr2
442
+ sum[n][y] += __builtin_amdgcn_mov_dpp(sum[n][y], 0x111, 0xf, 0xf,
443
+ 1); // row_shr1
444
+ sum[n][y] += __builtin_amdgcn_mov_dpp(sum[n][y], 0x142, 0xf, 0xf,
445
+ 1); // ROW_BCAST15
446
+ sum[n][y] += __builtin_amdgcn_mov_dpp(sum[n][y], 0x143, 0xf, 0xf,
447
+ 1); // ROW_BCAST31
448
}
449
}
450
451
if (threadIdx.x == 63) {
452
+ scalar_t biases[N][YTILE] = {};
453
+ if (BIAS)
454
+ for (int n = 0; n < N; n++) {
455
+ for (int y = 0; y < YTILE; y++) {
456
+ biases[n][y] = BIAS[(m + y) % Bx + (n % By) * Bx];
457
+ }
458
+ }
459
for (int n = 0; n < N; n++) {
499
- for (int i = 0; i < YTILE; i++) {
460
+ for (int y = 0; y < YTILE; y++) {
461
if constexpr (std::is_same_v<scalar_t, half>) {
501
- if (BIAS)
502
- sum[n][i] += __half2float(BIAS[(m + i) % Bx + (n % By) * M]);
462
+ sum[n][y] += __half2float(biases[n][y]);
463
} else if constexpr (std::is_same_v<scalar_t, __hip_bfloat16>) {
504
- if (BIAS)
505
- sum[n][i] +=
506
- __bfloat162float(BIAS[(m + i) % Bx + (n % By) * M]);
464
+ sum[n][y] += __bfloat162float(biases[n][y]);
465
}
508
- C[m + i + n * M] = __float2s<scalar_t>(sum[n][i]);
466
+ C[m + y + n * M] = __float2s<scalar_t>(sum[n][y]);
467
}
468
}
469
}
@@ -514,45 +472,43 @@ __global__ void __launch_bounds__(WvPrGrp* THRDS)
472
for (int n = 0; n < N; n++) {
473
#pragma unroll
474
for (int y = 0; y < YTILE; y++) {
517
- // float accm1 = 0;
518
- // for (int i=0; i<64; i++)
519
- // accm1 += __shfl(sum4[n][y][i%4], i);
475
+ /*float accm1 = 0;
476
+ for (int i=0; i<64; i++)
477
+ accm1 += __shfl(sum4[n][y][i%4], i);
478
+ sum4[n][y][0] = accm1;*/
479
float accm = sum4[n][y][0];
521
- asm("s_nop 0\n\tv_add_f32 %0, %2, %3 row_shl:1 bound_ctrl:0 "
522
- : "=v"(accm)
523
- : "0"(accm), "v"(sum4[n][y][1]), "v"(accm));
524
- asm("s_nop 0\n\tv_add_f32 %0, %2, %3 row_shl:2 bound_ctrl:0 "
525
- : "=v"(accm)
526
- : "0"(accm), "v"(sum4[n][y][2]), "v"(accm));
527
- asm("s_nop 0\n\tv_add_f32 %0, %2, %3 row_shl:3 bound_ctrl:0 "
528
- : "=v"(accm)
529
- : "0"(accm), "v"(sum4[n][y][3]), "v"(accm));
530
- asm("s_nop 0\n\tv_add_f32 %0, %2, %3 row_shl:4 bound_ctrl:0 "
531
- : "=v"(accm)
532
- : "0"(accm), "v"(accm), "v"(accm));
533
- asm("s_nop 0\n\tv_add_f32 %0, %2, %3 row_shl:8 bound_ctrl:0 "
534
- : "=v"(accm)
535
- : "0"(accm), "v"(accm), "v"(accm));
536
- asm("s_nop 0\n\tv_mov_b32 %0, %2 row_shr:15 bound_ctrl:0 "
537
- : "=v"(accm)
538
- : "0"(accm), "v"(accm));
539
- asm("s_nop 0\n\tv_add_f32 %0, %2, %3 row_bcast:15 bound_ctrl:0"
540
- : "=v"(accm)
541
- : "0"(accm), "v"(accm), "v"(accm));
542
- asm("s_nop 0\n\tv_add_f32 %0, %2, %3 row_bcast:31 bound_ctrl:0"
543
- : "=v"(accm)
544
- : "0"(accm), "v"(accm), "v"(accm));
480
+ accm += __builtin_amdgcn_mov_dpp(sum4[n][y][1], 0x101, 0xf, 0xf,
481
+ 1); // row_shl1
482
+ accm += __builtin_amdgcn_mov_dpp(sum4[n][y][2], 0x102, 0xf, 0xf,
483
+ 1); // row_shl2
484
+ accm += __builtin_amdgcn_mov_dpp(sum4[n][y][3], 0x103, 0xf, 0xf,
485
+ 1); // row_shl3
486
+ accm += __builtin_amdgcn_mov_dpp(accm, 0x104, 0xf, 0xf,
487
+ 1); // row_shl4
488
+ accm += __builtin_amdgcn_mov_dpp(accm, 0x108, 0xf, 0xf,
489
+ 1); // row_shl8
490
+ accm = __builtin_amdgcn_mov_dpp(accm, 0x11f, 0xf, 0xf,
491
+ 1); // row_shr15
492
+ accm += __builtin_amdgcn_mov_dpp(accm, 0x142, 0xf, 0xf,
493
+ 1); // ROW_BCAST15
494
+ accm += __builtin_amdgcn_mov_dpp(accm, 0x143, 0xf, 0xf,
495
+ 1); // ROW_BCAST31
496
497
sum4[n][y][0] = accm;
498
}
499
}
500
if (threadIdx.x == 63) {
501
+ scalar_t biases[N][YTILE] = {};
502
+ if (BIAS)
503
+ for (int n = 0; n < N; n++) {
504
+ for (int y = 0; y < YTILE; y++) {
505
+ biases[n][y] = BIAS[(m + y) % Bx + (n % By) * Bx];
506
+ }
507
+ }
508
for (int n = 0; n < N; n++) {
551
- for (int i = 0; i < YTILE; i++) {
552
- if (BIAS)
553
- sum4[n][i][0] +=
554
- __bfloat162float(BIAS[(m + i) % Bx + (n % By) * M]);
555
- C[m + i + n * M] = __float2bfloat16(sum4[n][i][0]);
509
+ for (int y = 0; y < YTILE; y++) {
510
+ sum4[n][y][0] += __bfloat162float(biases[n][y]);
511
+ C[m + y + n * M] = __float2bfloat16(sum4[n][y][0]);
512
}
513
}
514
}
@@ -563,8 +519,9 @@ __global__ void __launch_bounds__(WvPrGrp* THRDS)
519
#else // !defined(__HIP__GFX9__) TODO: Add NAVI support
520
template <typename scalar_t, int THRDS, int YTILE, int WvPrGrp, int A_CHUNK,
521
int UNRL, int N>
566
-__global__ void wvSplitK_hf_sml_(const int K, const int M, const int Bx,
567
- const int By, const scalar_t* B,
522
+__global__ void wvSplitK_hf_sml_(const int K, const int Kbp, const int Kap,
523
+ const int M, const int Bx, const int By,
524
+ const scalar_t* B,
525
const scalar_t* __restrict__ A,
526
const scalar_t* __restrict__ BIAS, scalar_t* C,
527
const int _WvPrGrp, const int CuCount) {
@@ -577,8 +534,9 @@ __global__ void wvSplitK_hf_sml_(const int K, const int M, const int Bx,
534
template <typename scalar_t, int THRDS, int YTILE, int WvPrGrp, int A_CHUNK,
535
int UNRL, int N>
536
__global__ void __launch_bounds__(WvPrGrp* THRDS)
580
- wvSplitK_hf_(const int K, const int M, const int Bx, const int By,
581
- const scalar_t* B, const scalar_t* __restrict__ A,
537
+ wvSplitK_hf_(const int K, const int Kbp, const int Kap, const int M,
538
+ const int Bx, const int By, const scalar_t* B,
539
+ const scalar_t* __restrict__ A,
540
const scalar_t* __restrict__ BIAS, scalar_t* C,
541
const int _WvPrGrp, const int CuCount) {
542
constexpr int max_lds_len = LDS_SIZE / 2;
@@ -601,13 +559,6 @@ __global__ void __launch_bounds__(WvPrGrp* THRDS)
559
scalar8 h8;
560
};
561
604
- //----------------------------------------------------
605
- // Reserving 64 KB of LDS to have 1 WG / CU
606
- // Goal is to bring the activation matrix A to the LDS
607
- // and use it across the lifetime of the work group
608
- // TODO: When activation matrix is larger than 64 KB
609
- // then this is not going to work!
610
- //----------------------------------------------------
562
__shared__ scalar_t s[max_lds_len];
563
564
//----------------------------------------------------
@@ -618,12 +569,6 @@ __global__ void __launch_bounds__(WvPrGrp* THRDS)
569
commitColumn[i] = 1;
570
}
571
621
- //----------------------------------------------------
622
- // Indexing function into the column of weight matrix B
623
- // Algorithm does 64 lane k-splitting / wave and uses
624
- // WG ID and Thread ID to find the index.
625
- //----------------------------------------------------
626
- // int _WvPrGrp = mindiv(N, CuCount * YTILE, WvPrGrp);
572
uint32_t m = (blockIdx.x * _WvPrGrp + threadIdx.y) * YTILE;
573
574
// Check whether there will be fragmentation!
@@ -636,91 +581,34 @@ __global__ void __launch_bounds__(WvPrGrp* THRDS)
581
m = startColumn;
582
}
583
639
- //----------------------------------------------------
640
- // Fetch the activation matrix to LDS
641
- // Loop iteration:
642
- // - Each thread (lane) is fetching 8 elements (A_Chunk)
643
- // - Each wave will fetch 64*8=> 512 elements
644
- // - Each WG will fetch 512 * 16 => 8K elements
645
- // - Then the WG will move to another 8 K elements
646
- // TODO: Logic below will only work when K is multiple of 8
647
- //----------------------------------------------------
648
- for (uint32_t k = 0; k < min__(K * N, max_lds_len);
649
- k += THRDS * WvPrGrp * A_CHUNK) {
650
- uint32_t k_in = k + ((threadIdx.y * THRDS + threadIdx.x) * A_CHUNK);
651
-
652
- if (k_in >= min__(K * N, max_lds_len)) break;
653
-
654
- *((bigType*)(&s[k_in])) = *((bigType*)(&A[k_in]));
584
+ for (uint32_t k = (threadIdx.y * THRDS + threadIdx.x) * A_CHUNK;
585
+ k < min__(Kap * N, max_lds_len); k += THRDS * WvPrGrp * A_CHUNK) {
586
+ #if defined(__gfx950__)
587
+ __builtin_amdgcn_global_load_lds((int*)(&A[k]), (int*)(&s[k]), 16, 0, 0);
588
+ #else
589
+ *((bigType*)(&s[k])) = *((bigType*)(&A[k]));
590
+ #endif
591
}
592
593
__syncthreads();
594
595
if (threadIdx.y >= _WvPrGrp) return;
596
661
- float sum[N][YTILE];
662
- scalar8 sum4[N][YTILE];
663
-
664
- //----------------------------------------------------
665
- // Each wave works on a single column of weight matrix.
666
- // There are 16 waves per WG, and hence, each WG is
667
- // working on 16 columns of weight matrix. Moreover,
668
- // we tile in column direction by YTILE, so when YTILE=1
669
- // the above math is right, however, when YTILE=2 then
670
- // each wave will be working on 2 columns and WG will
671
- // be working on 32 columns.
672
- //
673
- // Top level loop that makes WGs persistent!
674
- // - WGs iterates across columns of weight matrix
675
- // - Each wave within WG works on a given column(s)
676
- // - After completing first set of columns, WGs start
677
- // working on the next set of available columns
678
- //----------------------------------------------------
597
while (m < M) {
680
- //----------------------------------------------------
681
- // 'sum' accumulates the matrix A x B computation
682
- // split across 64 lanes.
683
- //
684
- // YTILE represents how many column of weight matrix
685
- // are being worked on by each wave.
686
- //----------------------------------------------------
687
- for (int i = 0; i < YTILE; i++)
688
- for (int n = 0; n < N; n++)
689
- if constexpr (!use_mfma)
690
- sum[n][i] = 0;
691
- else
692
- sum4[n][i] = {0, 0, 0, 0};
693
-
694
- bigType bigA[N][UNRL];
695
- bigType bigB[YTILE][UNRL];
696
- //----------------------------------------------------
697
- // Fetch weight matrix B in interleaved K-split!
698
- // - Each thread (lane) is fetching 8 elements (A_Chunk)
699
- // - Each wave will fetch 64*8=> 512 elements (1024B)
700
- // - YTILE represents the number of column being serviced
701
- // by wave
702
- // - Loop for fetching weight matrix (B) are unrolled
703
- //
704
- // Fetch activation matrix A from LDS
705
- // - Loop for fetching activation matrix (A) are unrolled
706
- //
707
- // Finally, do the matrix multiplication in an unrolled
708
- // fashion. This provides lot of food for compiler
709
- // scheduling.
710
- //
711
- // TODO: Logic below will only work when K is multiple of 8
712
- //----------------------------------------------------
598
+ float sum[N][YTILE] = {};
599
+ scalar8 sum4[N][YTILE] = {};
600
+
601
for (uint32_t k1 = 0; k1 < K; k1 += THRDS * A_CHUNK * UNRL) {
602
+ bigType bigA[N][UNRL] = {};
603
+ bigType bigB[YTILE][UNRL];
604
// Fetch the weight matrix from memory!
605
#pragma unroll
606
for (uint32_t k2 = 0; k2 < UNRL; k2++) {
607
uint32_t k = k1 + k2 * THRDS * A_CHUNK;
608
uint32_t k_ = k + threadIdx.x * A_CHUNK;
719
- if (k_ >= K) break;
720
-
721
- const scalar_t* B_ = &B[(m + 0) * K + k_];
722
- for (int b = 0; b < YTILE; b++)
723
- bigB[b][k2].h8 = (loadnt((scalar8*)(&B_[b * K])));
609
+ const scalar_t* B_ = &B[min__(k_, K - A_CHUNK)];
610
+ for (int y = 0; y < YTILE; y++)
611
+ bigB[y][k2].h8 = (loadnt((scalar8*)(&B_[min__(y + m, M - 1) * Kbp])));
612
}
613
614
// Fetch activation matrix from either just LDS or from both LDS / memory
@@ -729,36 +617,23 @@ __global__ void __launch_bounds__(WvPrGrp* THRDS)
617
uint32_t k = k1 + k2 * THRDS * A_CHUNK;
618
uint32_t k_ = k + threadIdx.x * A_CHUNK;
619
if (k_ >= K) break;
732
-
733
- // Fetch A activation matrix in interleaved fashion from LDS or memory
734
-
620
for (int n = 0; n < N; n++) {
736
- if (k_ + K * n < max_lds_len)
737
- bigA[n][k2] = *((const bigType*)(&(s[k_ + K * n])));
621
+ if (k_ + Kap * n < max_lds_len)
622
+ bigA[n][k2] = *((const bigType*)(&(s[k_ + Kap * n])));
623
else
739
- bigA[n][k2] = *((const bigType*)(&(A[k_ + K * n])));
624
+ bigA[n][k2] = *((const bigType*)(&(A[k_ + Kap * n])));
625
}
626
}
627
628
// Do the matrix multiplication in interleaved manner
744
- #pragma unroll
629
for (uint32_t n = 0; n < N; n++) {
746
- #pragma unroll
630
for (uint32_t k2 = 0; k2 < UNRL; k2++) {
748
- uint32_t k = k1 + k2 * THRDS * A_CHUNK;
749
- uint32_t k_ = k + threadIdx.x * A_CHUNK;
750
- if (k_ >= K) break;
751
- // Do the matrix multiplication of activation and weight matrix
752
- // - Remember the accumulation is happening for K-split of 64!
753
- #pragma unroll
631
for (int y = 0; y < YTILE; y++) {
632
if constexpr (!use_mfma)
756
- #pragma unroll
633
for (uint32_t b = 0; b < A_CHUNK / 2; b++) {
634
DOT2C(sum[n][y], bigA[n][k2].f[b], bigB[y][k2].f[b])
635
}
636
else
761
- #pragma unroll
637
for (uint32_t b = 0; b < A_CHUNK / 4; b++)
638
sum4[n][y] = __builtin_amdgcn_mfma_f32_4x4x4bf16_1k(
639
bigA[n][k2].h4[b], bigB[y][k2].h4[b], sum4[n][y], 0, 0, 0);
@@ -773,40 +648,38 @@ __global__ void __launch_bounds__(WvPrGrp* THRDS)
648
if constexpr (!use_mfma) {
649
for (int n = 0; n < N; n++) {
650
for (int y = 0; y < YTILE; y++) {
776
- asm("s_nop 0\n\tv_add_f32 %0, %2, %3 row_shr:8 bound_ctrl:0 "
777
- : "=v"(sum[n][y])
778
- : "0"(sum[n][y]), "v"(sum[n][y]), "v"(sum[n][y]));
779
- asm("s_nop 0\n\tv_add_f32 %0, %2, %3 row_shr:4 bound_ctrl:0 "
780
- : "=v"(sum[n][y])
781
- : "0"(sum[n][y]), "v"(sum[n][y]), "v"(sum[n][y]));
782
- asm("s_nop 0\n\tv_add_f32 %0, %2, %3 row_shr:2 bound_ctrl:0 "
783
- : "=v"(sum[n][y])
784
- : "0"(sum[n][y]), "v"(sum[n][y]), "v"(sum[n][y]));
785
- asm("s_nop 0\n\tv_add_f32 %0, %2, %3 wave_shr:1 bound_ctrl:0"
786
- : "=v"(sum[n][y])
787
- : "0"(sum[n][y]), "v"(sum[n][y]), "v"(sum[n][y]));
788
- asm("s_nop 0\n\tv_add_f32 %0, %2, %3 row_bcast:15 bound_ctrl:0"
789
- : "=v"(sum[n][y])
790
- : "0"(sum[n][y]), "v"(sum[n][y]), "v"(sum[n][y]));
791
- asm("s_nop 0\n\tv_add_f32 %0, %2, %3 row_bcast:31 bound_ctrl:0"
792
- : "=v"(sum[n][y])
793
- : "0"(sum[n][y]), "v"(sum[n][y]), "v"(sum[n][y]));
651
+ sum[n][y] += __builtin_amdgcn_mov_dpp(sum[n][y], 0x118, 0xf, 0xf,
652
+ 1); // row_shr8
653
+ sum[n][y] += __builtin_amdgcn_mov_dpp(sum[n][y], 0x114, 0xf, 0xf,
654
+ 1); // row_shr4
655
+ sum[n][y] += __builtin_amdgcn_mov_dpp(sum[n][y], 0x112, 0xf, 0xf,
656
+ 1); // row_shr2
657
+ sum[n][y] += __builtin_amdgcn_mov_dpp(sum[n][y], 0x111, 0xf, 0xf,
658
+ 1); // row_shr1
659
+ sum[n][y] += __builtin_amdgcn_mov_dpp(sum[n][y], 0x142, 0xf, 0xf,
660
+ 1); // ROW_BCAST15
661
+ sum[n][y] += __builtin_amdgcn_mov_dpp(sum[n][y], 0x143, 0xf, 0xf,
662
+ 1); // ROW_BCAST31
663
}
664
}
665
666
if (threadIdx.x == 63) {
667
+ scalar_t biases[N][YTILE] = {};
668
+ if (BIAS)
669
+ for (int n = 0; n < N; n++) {
670
+ for (int y = 0; y < YTILE; y++) {
671
+ biases[n][y] = BIAS[(m + y) % Bx + (n % By) * Bx];
672
+ }
673
+ }
674
for (int n = 0; n < N; n++) {
799
- for (int i = 0; i < YTILE; i++) {
800
- if (commitColumn[i]) {
675
+ for (int y = 0; y < YTILE; y++) {
676
+ if (commitColumn[y]) {
677
if constexpr (std::is_same_v<scalar_t, half>) {
802
- if (BIAS)
803
- sum[n][i] += __half2float(BIAS[(m + i) % Bx + (n % By) * M]);
678
+ sum[n][y] += __half2float(biases[n][y]);
679
} else if constexpr (std::is_same_v<scalar_t, __hip_bfloat16>) {
805
- if (BIAS)
806
- sum[n][i] +=
807
- __bfloat162float(BIAS[(m + i) % Bx + (n % By) * M]);
680
+ sum[n][y] += __bfloat162float(biases[n][y]);
681
}
809
- C[m + i + n * M] = __float2s<scalar_t>(sum[n][i]);
682
+ C[m + y + n * M] = __float2s<scalar_t>(sum[n][y]);
683
}
684
}
685
}
@@ -819,44 +692,39 @@ __global__ void __launch_bounds__(WvPrGrp* THRDS)
692
// float accm1 = 0;
693
// for (int i=0; i<64; i++)
694
// accm1 += __shfl(sum4[n][y][i%4], i);
822
-
695
float accm = sum4[n][y][0];
824
- asm("s_nop 0\n\tv_add_f32 %0, %2, %3 row_shl:1 bound_ctrl:0 "
825
- : "=v"(accm)
826
- : "0"(accm), "v"(sum4[n][y][1]), "v"(accm));
827
- asm("s_nop 0\n\tv_add_f32 %0, %2, %3 row_shl:2 bound_ctrl:0 "
828
- : "=v"(accm)
829
- : "0"(accm), "v"(sum4[n][y][2]), "v"(accm));
830
- asm("s_nop 0\n\tv_add_f32 %0, %2, %3 row_shl:3 bound_ctrl:0 "
831
- : "=v"(accm)
832
- : "0"(accm), "v"(sum4[n][y][3]), "v"(accm));
833
- asm("s_nop 0\n\tv_add_f32 %0, %2, %3 row_shl:4 bound_ctrl:0 "
834
- : "=v"(accm)
835
- : "0"(accm), "v"(accm), "v"(accm));
836
- asm("s_nop 0\n\tv_add_f32 %0, %2, %3 row_shl:8 bound_ctrl:0 "
837
- : "=v"(accm)
838
- : "0"(accm), "v"(accm), "v"(accm));
839
- asm("s_nop 0\n\tv_mov_b32 %0, %2 row_shr:15 bound_ctrl:0 "
840
- : "=v"(accm)
841
- : "0"(accm), "v"(accm));
842
- asm("s_nop 0\n\tv_add_f32 %0, %2, %3 row_bcast:15 bound_ctrl:0"
843
- : "=v"(accm)
844
- : "0"(accm), "v"(accm), "v"(accm));
845
- asm("s_nop 0\n\tv_add_f32 %0, %2, %3 row_bcast:31 bound_ctrl:0"
846
- : "=v"(accm)
847
- : "0"(accm), "v"(accm), "v"(accm));
848
-
696
+ accm += __builtin_amdgcn_mov_dpp(sum4[n][y][1], 0x101, 0xf, 0xf,
697
+ 1); // row_shl1
698
+ accm += __builtin_amdgcn_mov_dpp(sum4[n][y][2], 0x102, 0xf, 0xf,
699
+ 1); // row_shl2
700
+ accm += __builtin_amdgcn_mov_dpp(sum4[n][y][3], 0x103, 0xf, 0xf,
701
+ 1); // row_shl3
702
+ accm += __builtin_amdgcn_mov_dpp(accm, 0x104, 0xf, 0xf,
703
+ 1); // row_shl4
704
+ accm += __builtin_amdgcn_mov_dpp(accm, 0x108, 0xf, 0xf,
705
+ 1); // row_shl8
706
+ accm = __builtin_amdgcn_mov_dpp(accm, 0x11f, 0xf, 0xf,
707
+ 1); // row_shr15
708
+ accm += __builtin_amdgcn_mov_dpp(accm, 0x142, 0xf, 0xf,
709
+ 1); // ROW_BCAST15
710
+ accm += __builtin_amdgcn_mov_dpp(accm, 0x143, 0xf, 0xf,
711
+ 1); // ROW_BCAST31
712
sum4[n][y][0] = accm;
713
}
714
}
715
if (threadIdx.x == 63) {
716
+ scalar_t biases[N][YTILE] = {};
717
+ if (BIAS)
718
+ for (int n = 0; n < N; n++) {
719
+ for (int y = 0; y < YTILE; y++) {
720
+ biases[n][y] = BIAS[(m + y) % Bx + (n % By) * Bx];
721
+ }
722
+ }
723
for (int n = 0; n < N; n++) {
854
- for (int i = 0; i < YTILE; i++) {
855
- if (commitColumn[i]) {
856
- if (BIAS)
857
- sum4[n][i][0] +=
858
- __bfloat162float(BIAS[(m + i) % Bx + (n % By) * M]);
859
- C[m + i + n * M] = __float2bfloat16(sum4[n][i][0]);
724
+ for (int y = 0; y < YTILE; y++) {
725
+ if (commitColumn[y]) {
726
+ sum4[n][y][0] += __bfloat162float(biases[n][y]);
727
+ C[m + y + n * M] = __float2bfloat16(sum4[n][y][0]);
728
}
729
}
730
}
@@ -880,9 +748,9 @@ __global__ void __launch_bounds__(WvPrGrp* THRDS)
748
#else // !defined(__HIP__GFX9__) TODO: Add NAVI support
749
template <typename scalar_t, int THRDS, int YTILE, int WvPrGrp, int A_CHUNK,
750
int UNRL, int N>
883
-__global__ void wvSplitK_hf_(const int K, const int M, const int Bx,
884
- const int By, const scalar_t* B,
885
- const scalar_t* __restrict__ A,
751
+__global__ void wvSplitK_hf_(const int K, const int Kbp, const int Kap,
752
+ const int M, const int Bx, const int By,
753
+ const scalar_t* B, const scalar_t* __restrict__ A,
754
const scalar_t* __restrict__ BIAS, scalar_t* C,
755
const int _WvPrGrp, const int CuCount) {
756
UNREACHABLE_CODE
@@ -894,8 +762,9 @@ __global__ void wvSplitK_hf_(const int K, const int M, const int Bx,
762
template <typename scalar_t, int THRDS, int YTILE, int WvPrGrp, int A_CHUNK,
763
int UNRL, int N>
764
__global__ void __launch_bounds__(WvPrGrp* THRDS)
897
- wvSplitK_hf_big_(const int K, const int M, const int Bx, const int By,
898
- const scalar_t* B, const scalar_t* __restrict__ A,
765
+ wvSplitK_hf_big_(const int K, const int Kbp, const int Kap, const int M,
766
+ const int Bx, const int By, const scalar_t* B,
767
+ const scalar_t* __restrict__ A,
768
const scalar_t* __restrict__ BIAS, scalar_t* C,
769
const int _WvPrGrp, const int CuCount) {
770
constexpr int max_lds_len = LDS_SIZE / 2;
@@ -966,13 +835,13 @@ __global__ void __launch_bounds__(WvPrGrp* THRDS)
835
//----------------------------------------------------
836
#define PCML
837
#ifndef PCML
969
- for (uint32_t k = 0; k < min__(K * N, max_lds_len);
970
- k += THRDS * WvPrGrp * A_CHUNK) {
971
- uint32_t k_in = k + ((threadIdx.y * THRDS + threadIdx.x) * A_CHUNK);
972
-
973
- if (k_in >= min__(K * N, max_lds_len)) break;
974
-
975
- *((bigType*)(&s[k_in])) = *((bigType*)(&A[k_in]));
838
+ for (uint32_t k = (threadIdx.y * THRDS + threadIdx.x) * A_CHUNK;
839
+ k < min__(Kap * N, max_lds_len); k += THRDS * WvPrGrp * A_CHUNK) {
840
+ #if defined(__gfx950__)
841
+ __builtin_amdgcn_global_load_lds((int*)(&A[k]), (int*)(&s[k]), 16, 0, 0);
842
+ #else
843
+ *((bigType*)(&s[k])) = *((bigType*)(&A[k]));
844
+ #endif
845
}
846
__syncthreads();
847
#endif
@@ -987,10 +856,7 @@ __global__ void __launch_bounds__(WvPrGrp* THRDS)
856
? kFit
857
: (kFit - kFit % TUC); // round up to multiple of TUC
858
// if (kFit == 0) kFit = TUC;
990
- kFit = min__(kFit, K);
991
-
992
- float sum[N][YTILE];
993
- scalar8 sum4[N][YTILE];
859
+ kFit = min__(kFit, Kap);
860
861
//----------------------------------------------------
862
// Each wave works on a single column of weight matrix.
@@ -1021,15 +887,9 @@ __global__ void __launch_bounds__(WvPrGrp* THRDS)
887
// YTILE represents how many column of weight matrix
888
// are being worked on by each wave.
889
//----------------------------------------------------
1024
- for (int i = 0; i < YTILE; i++)
1025
- for (int n = 0; n < N; n++)
1026
- if constexpr (!use_mfma)
1027
- sum[n][i] = 0;
1028
- else
1029
- sum4[n][i] = {0, 0, 0, 0};
1030
-
1031
- bigType bigA[N][UNRL];
1032
- bigType bigB[YTILE][UNRL];
890
+ float sum[N][YTILE] = {};
891
+ scalar8 sum4[N][YTILE] = {};
892
+
893
//----------------------------------------------------
894
// Fetch weight matrix B in interleaved K-split!
895
// - Each thread (lane) is fetching 8 elements (A_Chunk)
@@ -1048,18 +908,26 @@ __global__ void __launch_bounds__(WvPrGrp* THRDS)
908
// TODO: Logic below will only work when K is multiple of 8
909
//----------------------------------------------------
910
for (uint32_t k1 = 0; k1 < K; k1 += THRDS * A_CHUNK * UNRL) {
911
+ bigType bigA[N][UNRL] = {};
912
+ bigType bigB[YTILE][UNRL];
913
+
914
#ifdef PCML
915
if ((k1 == 0) || (k1 == kBase + kFit)) { // load next chunk of A[] to LDS
916
if (k1 != 0) kBase += kFit;
917
__syncthreads();
918
for (uint32_t k = 0; k < kFit; k += THRDS * _WvPrGrp * A_CHUNK) {
919
uint32_t kOff = k + ((threadIdx.y * THRDS + threadIdx.x) * A_CHUNK);
1057
- if (kBase + kOff >= K) break;
920
+ if (kBase + kOff >= Kap) break;
921
if (kOff >= kFit) break;
922
for (uint32_t n = 0; n < N; n++) {
1060
- uint32_t k_in = kBase + n * K + kOff;
923
+ uint32_t k_in = kBase + n * Kap + kOff;
924
uint32_t k_ot = n * kFit + kOff;
925
+ #if defined(__gfx950__)
926
+ __builtin_amdgcn_global_load_lds((int*)(&A[k_in]), (int*)(&s[k_ot]),
927
+ 16, 0, 0);
928
+ #else
929
*((bigType*)(&s[k_ot])) = *((bigType*)(&A[k_in]));
930
+ #endif
931
}
932
}
933
__syncthreads();
@@ -1072,11 +940,9 @@ __global__ void __launch_bounds__(WvPrGrp* THRDS)
940
for (uint32_t k2 = 0; k2 < UNRL; k2++) {
941
uint32_t k = k1 + k2 * THRDS * A_CHUNK;
942
uint32_t k_ = k + threadIdx.x * A_CHUNK;
1075
- if (k_ >= K) break;
1076
-
1077
- const scalar_t* B_ = &B[(m + 0) * K + k_];
1078
- for (int b = 0; b < YTILE; b++)
1079
- bigB[b][k2].h8 = (loadnt((scalar8*)(&B_[b * K])));
943
+ const scalar_t* B_ = &B[min__(k_, K - A_CHUNK)];
944
+ for (int y = 0; y < YTILE; y++)
945
+ bigB[y][k2].h8 = (loadnt((scalar8*)(&B_[min__(y + m, M - 1) * Kbp])));
946
}
947
948
// Fetch activation matrix from either just LDS or from both LDS / memory
@@ -1085,17 +951,14 @@ __global__ void __launch_bounds__(WvPrGrp* THRDS)
951
uint32_t k = k1 + k2 * THRDS * A_CHUNK;
952
uint32_t k_ = k + threadIdx.x * A_CHUNK;
953
if (k_ >= K) break;
1088
-
1089
- // Fetch A activation matrix in interleaved fashion from LDS or memory
1090
-
954
for (int n = 0; n < N; n++) {
955
#ifdef PCML
956
bigA[n][k2] = *((const bigType*)(&(s[k_ - kBase + kFit * n])));
957
#else
1095
- if (k_ + K * n < 32 * 1024)
1096
- bigA[n][k2] = *((const bigType*)(&(s[k_ + K * n])));
958
+ if (k_ + Kap * n < max_lds_len)
959
+ bigA[n][k2] = *((const bigType*)(&(s[k_ + Kap * n])));
960
else
1098
- bigA[n][k2] = *((const bigType*)(&(A[k_ + K * n])));
961
+ bigA[n][k2] = *((const bigType*)(&(A[k_ + Kap * n])));
962
#endif
963
}
964
}
@@ -1103,22 +966,13 @@ __global__ void __launch_bounds__(WvPrGrp* THRDS)
966
// Do the matrix multiplication in interleaved manner
967
#pragma unroll
968
for (uint32_t k2 = 0; k2 < UNRL; k2++) {
1106
- uint32_t k = k1 + k2 * THRDS * A_CHUNK;
1107
- uint32_t k_ = k + threadIdx.x * A_CHUNK;
1108
- if (k_ >= K) break;
1109
- #pragma unroll
969
for (uint32_t n = 0; n < N; n++) {
1111
- // Do the matrix multiplication of activation and weight matrix
1112
- // - Remember the accumulation is happening for K-split of 64!
1113
- #pragma unroll
970
for (int y = 0; y < YTILE; y++) {
971
if constexpr (!use_mfma)
1116
- #pragma unroll
972
for (uint32_t b = 0; b < A_CHUNK / 2; b++) {
973
DOT2C(sum[n][y], bigA[n][k2].f[b], bigB[y][k2].f[b])
974
}
975
else
1121
- #pragma unroll
976
for (uint32_t b = 0; b < A_CHUNK / 4; b++)
977
sum4[n][y] = __builtin_amdgcn_mfma_f32_4x4x4bf16_1k(
978
bigA[n][k2].h4[b], bigB[y][k2].h4[b], sum4[n][y], 0, 0, 0);
@@ -1141,40 +995,38 @@ __global__ void __launch_bounds__(WvPrGrp* THRDS)
995
if constexpr (!use_mfma) {
996
for (int n = 0; n < N; n++) {
997
for (int y = 0; y < YTILE; y++) {
1144
- asm("s_nop 0\n\tv_add_f32 %0, %2, %3 row_shr:8 bound_ctrl:0 "
1145
- : "=v"(sum[n][y])
1146
- : "0"(sum[n][y]), "v"(sum[n][y]), "v"(sum[n][y]));
1147
- asm("s_nop 0\n\tv_add_f32 %0, %2, %3 row_shr:4 bound_ctrl:0 "
1148
- : "=v"(sum[n][y])
1149
- : "0"(sum[n][y]), "v"(sum[n][y]), "v"(sum[n][y]));
1150
- asm("s_nop 0\n\tv_add_f32 %0, %2, %3 row_shr:2 bound_ctrl:0 "
1151
- : "=v"(sum[n][y])
1152
- : "0"(sum[n][y]), "v"(sum[n][y]), "v"(sum[n][y]));
1153
- asm("s_nop 0\n\tv_add_f32 %0, %2, %3 wave_shr:1 bound_ctrl:0"
1154
- : "=v"(sum[n][y])
1155
- : "0"(sum[n][y]), "v"(sum[n][y]), "v"(sum[n][y]));
1156
- asm("s_nop 0\n\tv_add_f32 %0, %2, %3 row_bcast:15 bound_ctrl:0"
1157
- : "=v"(sum[n][y])
1158
- : "0"(sum[n][y]), "v"(sum[n][y]), "v"(sum[n][y]));
1159
- asm("s_nop 0\n\tv_add_f32 %0, %2, %3 row_bcast:31 bound_ctrl:0"
1160
- : "=v"(sum[n][y])
1161
- : "0"(sum[n][y]), "v"(sum[n][y]), "v"(sum[n][y]));
998
+ sum[n][y] += __builtin_amdgcn_mov_dpp(sum[n][y], 0x118, 0xf, 0xf,
999
+ 1); // row_shr8
1000
+ sum[n][y] += __builtin_amdgcn_mov_dpp(sum[n][y], 0x114, 0xf, 0xf,
1001
+ 1); // row_shr4
1002
+ sum[n][y] += __builtin_amdgcn_mov_dpp(sum[n][y], 0x112, 0xf, 0xf,
1003
+ 1); // row_shr2
1004
+ sum[n][y] += __builtin_amdgcn_mov_dpp(sum[n][y], 0x111, 0xf, 0xf,
1005
+ 1); // row_shr1
1006
+ sum[n][y] += __builtin_amdgcn_mov_dpp(sum[n][y], 0x142, 0xf, 0xf,
1007
+ 1); // ROW_BCAST15
1008
+ sum[n][y] += __builtin_amdgcn_mov_dpp(sum[n][y], 0x143, 0xf, 0xf,
1009
+ 1); // ROW_BCAST31
1010
}
1011
}
1012
1013
if (threadIdx.x == 63) {
1014
+ scalar_t biases[N][YTILE] = {};
1015
+ if (BIAS)
1016
+ for (int n = 0; n < N; n++) {
1017
+ for (int y = 0; y < YTILE; y++) {
1018
+ biases[n][y] = BIAS[(m + y) % Bx + (n % By) * Bx];
1019
+ }
1020
+ }
1021
for (int n = 0; n < N; n++) {
1167
- for (int i = 0; i < YTILE; i++) {
1168
- if (commitColumn[i]) {
1022
+ for (int y = 0; y < YTILE; y++) {
1023
+ if (commitColumn[y]) {
1024
if constexpr (std::is_same_v<scalar_t, half>) {
1170
- if (BIAS)
1171
- sum[n][i] += __half2float(BIAS[(m + i) % Bx + (n % By) * M]);
1025
+ sum[n][y] += __half2float(biases[n][y]);
1026
} else if constexpr (std::is_same_v<scalar_t, __hip_bfloat16>) {
1173
- if (BIAS)
1174
- sum[n][i] +=
1175
- __bfloat162float(BIAS[(m + i) % Bx + (n % By) * M]);
1027
+ sum[n][y] += __bfloat162float(biases[n][y]);
1028
}
1177
- C[m + i + n * M] = __float2s<scalar_t>(sum[n][i]);
1029
+ C[m + y + n * M] = __float2s<scalar_t>(sum[n][y]);
1030
}
1031
}
1032
}
@@ -1185,42 +1037,38 @@ __global__ void __launch_bounds__(WvPrGrp* THRDS)
1037
#pragma unroll
1038
for (int y = 0; y < YTILE; y++) {
1039
float accm = sum4[n][y][0];
1188
- asm("s_nop 0\n\tv_add_f32 %0, %2, %3 row_shl:1 bound_ctrl:0 "
1189
- : "=v"(accm)
1190
- : "0"(accm), "v"(sum4[n][y][1]), "v"(accm));
1191
- asm("s_nop 0\n\tv_add_f32 %0, %2, %3 row_shl:2 bound_ctrl:0 "
1192
- : "=v"(accm)
1193
- : "0"(accm), "v"(sum4[n][y][2]), "v"(accm));
1194
- asm("s_nop 0\n\tv_add_f32 %0, %2, %3 row_shl:3 bound_ctrl:0 "
1195
- : "=v"(accm)
1196
- : "0"(accm), "v"(sum4[n][y][3]), "v"(accm));
1197
- asm("s_nop 0\n\tv_add_f32 %0, %2, %3 row_shl:4 bound_ctrl:0 "
1198
- : "=v"(accm)
1199
- : "0"(accm), "v"(accm), "v"(accm));
1200
- asm("s_nop 0\n\tv_add_f32 %0, %2, %3 row_shl:8 bound_ctrl:0 "
1201
- : "=v"(accm)
1202
- : "0"(accm), "v"(accm), "v"(accm));
1203
- asm("s_nop 0\n\tv_mov_b32 %0, %2 row_shr:15 bound_ctrl:0 "
1204
- : "=v"(accm)
1205
- : "0"(accm), "v"(accm));
1206
- asm("s_nop 0\n\tv_add_f32 %0, %2, %3 row_bcast:15 bound_ctrl:0"
1207
- : "=v"(accm)
1208
- : "0"(accm), "v"(accm), "v"(accm));
1209
- asm("s_nop 0\n\tv_add_f32 %0, %2, %3 row_bcast:31 bound_ctrl:0"
1210
- : "=v"(accm)
1211
- : "0"(accm), "v"(accm), "v"(accm));
1212
-
1040
+ accm += __builtin_amdgcn_mov_dpp(sum4[n][y][1], 0x101, 0xf, 0xf,
1041
+ 1); // row_shl1
1042
+ accm += __builtin_amdgcn_mov_dpp(sum4[n][y][2], 0x102, 0xf, 0xf,
1043
+ 1); // row_shl2
1044
+ accm += __builtin_amdgcn_mov_dpp(sum4[n][y][3], 0x103, 0xf, 0xf,
1045
+ 1); // row_shl3
1046
+ accm += __builtin_amdgcn_mov_dpp(accm, 0x104, 0xf, 0xf,
1047
+ 1); // row_shl4
1048
+ accm += __builtin_amdgcn_mov_dpp(accm, 0x108, 0xf, 0xf,
1049
+ 1); // row_shl8
1050
+ accm = __builtin_amdgcn_mov_dpp(accm, 0x11f, 0xf, 0xf,
1051
+ 1); // row_shr15
1052
+ accm += __builtin_amdgcn_mov_dpp(accm, 0x142, 0xf, 0xf,
1053
+ 1); // ROW_BCAST15
1054
+ accm += __builtin_amdgcn_mov_dpp(accm, 0x143, 0xf, 0xf,
1055
+ 1); // ROW_BCAST31
1056
sum4[n][y][0] = accm;
1057
}
1058
}
1059
if (threadIdx.x == 63) {
1060
+ scalar_t biases[N][YTILE] = {};
1061
+ if (BIAS)
1062
+ for (int n = 0; n < N; n++) {
1063
+ for (int y = 0; y < YTILE; y++) {
1064
+ biases[n][y] = BIAS[(m + y) % Bx + (n % By) * Bx];
1065
+ }
1066
+ }
1067
for (int n = 0; n < N; n++) {
1218
- for (int i = 0; i < YTILE; i++) {
1219
- if (commitColumn[i]) {
1220
- if (BIAS)
1221
- sum4[n][i][0] +=
1222
- __bfloat162float(BIAS[(m + i) % Bx + (n % By) * M]);
1223
- C[m + i + n * M] = __float2bfloat16(sum4[n][i][0]);
1068
+ for (int y = 0; y < YTILE; y++) {
1069
+ if (commitColumn[y]) {
1070
+ sum4[n][y][0] += __bfloat162float(biases[n][y]);
1071
+ C[m + y + n * M] = __float2bfloat16(sum4[n][y][0]);
1072
}
1073
}
1074
}
@@ -1244,8 +1092,9 @@ __global__ void __launch_bounds__(WvPrGrp* THRDS)
1092
#else // !defined(__HIP__GFX9__) TODO: Add NAVI support
1093
template <typename scalar_t, int THRDS, int YTILE, int WvPrGrp, int A_CHUNK,
1094
int UNRL, int N>
1247
-__global__ void wvSplitK_hf_big_(const int K, const int M, const int Bx,
1248
- const int By, const scalar_t* B,
1095
+__global__ void wvSplitK_hf_big_(const int K, const int Kbp, const int Kap,
1096
+ const int M, const int Bx, const int By,
1097
+ const scalar_t* B,
1098
const scalar_t* __restrict__ A,
1099
const scalar_t* __restrict__ BIAS, scalar_t* C,
1100
const int _WvPrGrp, const int CuCount) {
@@ -1272,6 +1121,8 @@ torch::Tensor wvSplitK(const at::Tensor& in_a, const at::Tensor& in_b,
1121
auto M_in = in_a.size(0);
1122
auto K_in = in_a.size(1);
1123
auto N_in = in_b.size(0);
1124
+ auto Kap_in = in_a.stride(0);
1125
+ auto Kbp_in = in_b.stride(0);
1126
auto Bx_in =
1127
(in_bias.has_value() && in_bias->numel() > 0)
1128
? (in_bias->sizes().size() == 2) ? in_bias->size(1) : in_bias->size(0)
@@ -1296,27 +1147,30 @@ torch::Tensor wvSplitK(const at::Tensor& in_a, const at::Tensor& in_b,
1147
const cudaStream_t stream = at::cuda::getCurrentCUDAStream();
1148
const int max_lds_len = get_lds_size() / 2;
1149
1299
-#define WVSPLITK(_YTILE, _UNRL, _N) \
1300
- { \
1301
- dim3 block(64, 16); \
1302
- int __wvPrGrp = mindiv(M_in, CuCount * _YTILE, 16); \
1303
- if ((K_in * N_in <= max_lds_len) && (M_in % _YTILE == 0)) \
1304
- wvSplitK_hf_sml_<fptype, 64, _YTILE, 16, 8, _UNRL, _N> \
1305
- <<<grid, block, 0, stream>>>(K_in, M_in, Bx_in, By_in, af4, bf4, \
1306
- biasf4, c, __wvPrGrp, CuCount); \
1307
- else if (K_in * N_in <= max_lds_len * 1.2) \
1308
- wvSplitK_hf_<fptype, 64, _YTILE, 16, 8, _UNRL, _N> \
1309
- <<<grid, block, 0, stream>>>(K_in, M_in, Bx_in, By_in, af4, bf4, \
1310
- biasf4, c, __wvPrGrp, CuCount); \
1311
- else \
1312
- wvSplitK_hf_big_<fptype, 64, _YTILE, 16, 8, _UNRL, _N> \
1313
- <<<grid, block, 0, stream>>>(K_in, M_in, Bx_in, By_in, af4, bf4, \
1314
- biasf4, c, __wvPrGrp, CuCount); \
1150
+#define WVSPLITK(_YTILE, _UNRL, _N) \
1151
+ { \
1152
+ dim3 block(64, 16); \
1153
+ int __wvPrGrp = mindiv(M_in, CuCount * _YTILE, 16); \
1154
+ if ((Kbp_in * N_in <= max_lds_len) && (M_in % _YTILE == 0)) \
1155
+ wvSplitK_hf_sml_<fptype, 64, _YTILE, 16, 8, _UNRL, _N> \
1156
+ <<<grid, block, 0, stream>>>(K_in, Kap_in, Kbp_in, M_in, Bx_in, \
1157
+ By_in, af4, bf4, biasf4, c, __wvPrGrp, \
1158
+ CuCount); \
1159
+ else if (Kbp_in * N_in <= max_lds_len * 1.2) \
1160
+ wvSplitK_hf_<fptype, 64, _YTILE, 16, 8, _UNRL, _N> \
1161
+ <<<grid, block, 0, stream>>>(K_in, Kap_in, Kbp_in, M_in, Bx_in, \
1162
+ By_in, af4, bf4, biasf4, c, __wvPrGrp, \
1163
+ CuCount); \
1164
+ else \
1165
+ wvSplitK_hf_big_<fptype, 64, _YTILE, 16, 8, _UNRL, _N> \
1166
+ <<<grid, block, 0, stream>>>(K_in, Kap_in, Kbp_in, M_in, Bx_in, \
1167
+ By_in, af4, bf4, biasf4, c, __wvPrGrp, \
1168
+ CuCount); \
1169
}
1170
1171
#define WVSPLIT_TILE(_sYT, __N) \
1172
{ \
1319
- bool fit_lds = (K_in * N_in <= max_lds_len); \
1173
+ bool fit_lds = (Kbp_in * N_in <= max_lds_len); \
1174
if (_sYT <= 1) \
1175
WVSPLITK(1, 4, __N) \
1176
else if ((__N == 1) || (!fit_lds) || (_sYT <= 4 * 2)) \
tests/kernels/quantization/test_rocm_skinny_gemms.py
+33
-41
@@ -30,15 +30,22 @@ NKM_FACTORS_LLMM1 = [
30
31
NKM_FACTORS_WVSPLITK = [
32
# Different batch sizes with key dimensions
33
- (1, 16, 16),
33
+ (1, 32, 16),
34
(1, 64, 64),
35
(2, 256, 256),
36
(3, 1024, 1024),
37
(4, 4096, 4096),
38
+ (4, 4096, 4096 + 1),
39
+ (4, 4096 + 16, 4096),
40
+ (4, 4096 + 16, 4096 + 1),
41
# Extended K values
42
(1, 9216, 512),
43
(2, 10240, 1024),
44
(4, 16384, 8192),
45
+ (4, 16384 * 2, 8192),
46
+ (4, 16384 * 2, 8192 + 1),
47
+ (4, 16384 * 2 + 16, 8192),
48
+ (4, 16384 * 2 + 16, 8192 + 1),
49
# Minimum M constraint validation (m >= 8)
50
(1, 64, 8),
51
(2, 128, 8),
@@ -180,59 +187,44 @@ def test_rocm_llmm1_kernel(n, k, m, dtype, rows_per_block, seed):
187
torch.testing.assert_close(out, ref_out, atol=1e-8, rtol=1e-2)
188
189
190
+@pytest.mark.parametrize("xnorm", [False, True])
191
@pytest.mark.parametrize("n,k,m", NKM_FACTORS_WVSPLITK)
192
@pytest.mark.parametrize("dtype", DTYPES)
193
@pytest.mark.parametrize("seed", SEEDS)
194
@pytest.mark.skipif(not current_platform.is_rocm(), reason="only test for rocm")
187
-def test_rocm_wvsplitk_kernel(n, k, m, dtype, seed):
188
- torch.manual_seed(seed)
189
- cu_count = num_compute_units()
190
-
191
- A = torch.rand(n, k, dtype=dtype, device="cuda") - 0.5
192
- B = torch.rand(m, k, dtype=dtype, device="cuda") - 0.5
193
-
194
- ref_out = torch.nn.functional.linear(A, B)
195
- out = ops.wvSplitK(B, A.view(-1, A.size(-1)), cu_count)
196
-
197
- torch.testing.assert_close(out, ref_out, atol=1e-8, rtol=1e-2)
198
-
199
-
200
-@pytest.mark.parametrize("n,k,m", NKM_FACTORS_WVSPLITK)
201
-@pytest.mark.parametrize("dtype", DTYPES)
202
-@pytest.mark.parametrize("seed", SEEDS)
203
-@pytest.mark.skipif(not current_platform.is_rocm(), reason="only test for rocm")
204
-def test_rocm_wvsplitk_bias1D_kernel(n, k, m, dtype, seed):
195
+@pytest.mark.parametrize("bias_mode", BIAS_MODES)
196
+@pytest.mark.parametrize("padded_a", [False, True])
197
+@pytest.mark.parametrize("padded_b", [False, True])
198
+def test_rocm_wvsplitk_kernel(
199
+ xnorm, n, k, m, dtype, seed, bias_mode, padded_a, padded_b
200
+):
201
torch.manual_seed(seed)
202
cu_count = num_compute_units()
203
208
- xavier = math.sqrt(2 / k) # normalize to avoid large output-bias deltas
209
- A = (torch.rand(n, k, dtype=dtype, device="cuda") - 0.5) * xavier
210
- B = (torch.rand(m, k, dtype=dtype, device="cuda") - 0.5) * xavier
211
- BIAS = torch.rand(m, dtype=dtype, device="cuda") - 0.5
212
-
213
- ref_out = torch.nn.functional.linear(A, B, BIAS)
214
- out = ops.wvSplitK(B, A.view(-1, A.size(-1)), cu_count, BIAS)
215
-
216
- torch.testing.assert_close(out, ref_out, atol=1e-8, rtol=1e-2)
217
-
204
+ xavier = (
205
+ math.sqrt(2 / k) if xnorm else 1
206
+ ) # normalize to avoid large output-bias deltas
207
+ A = (torch.rand(n, k, dtype=dtype, device="cuda") * 2 - 1) * xavier
208
+ B = (torch.rand(m, k, dtype=dtype, device="cuda") * 2 - 1) * xavier
209
219
-@pytest.mark.parametrize("n,k,m", NKM_FACTORS_WVSPLITK)
220
-@pytest.mark.parametrize("dtype", DTYPES)
221
-@pytest.mark.parametrize("seed", SEEDS)
222
-@pytest.mark.skipif(not current_platform.is_rocm(), reason="only test for rocm")
223
-def test_rocm_wvsplitk_bias2D_kernel(n, k, m, dtype, seed):
224
- torch.manual_seed(seed)
225
- cu_count = num_compute_units()
210
+ BIAS = None
211
+ if bias_mode == 1:
212
+ BIAS = torch.rand(m, dtype=dtype, device="cuda") * 2 - 1
213
+ elif bias_mode == 2:
214
+ BIAS = torch.rand(n, m, dtype=dtype, device="cuda") * 2 - 1
215
227
- xavier = math.sqrt(2 / k) # normalize to avoid large output-bias deltas
228
- A = (torch.rand(n, k, dtype=dtype, device="cuda") - 0.5) * xavier
229
- B = (torch.rand(m, k, dtype=dtype, device="cuda") - 0.5) * xavier
230
- BIAS = torch.rand(n, m, dtype=dtype, device="cuda") - 0.5
216
+ if padded_a:
217
+ A = pad_fp8(A)
218
+ if padded_b:
219
+ B = pad_fp8(B)
220
221
ref_out = torch.nn.functional.linear(A, B, BIAS)
222
out = ops.wvSplitK(B, A.view(-1, A.size(-1)), cu_count, BIAS)
223
235
- torch.testing.assert_close(out, ref_out, atol=1e-8, rtol=1e-2)
224
+ if xnorm:
225
+ assert torch.allclose(out, ref_out, atol=1e-3, rtol=1e-8)
226
+ else:
227
+ assert torch.allclose(out, ref_out, atol=1e-3, rtol=1e-2)
228
229
230
@pytest.mark.parametrize("xnorm", [False, True])
vllm/model_executor/layers/utils.py
-1
@@ -191,7 +191,6 @@ def rocm_unquantized_gemm_impl(
191
and on_gfx9()
192
and x.dtype in [torch.float16, torch.bfloat16]
193
and k % 8 == 0
194
- and x.is_contiguous()
194
)
195
196
if use_skinny is not True: