@tduncan / Nvidia-Optimization-Stack-VLLM Mirror · read-only / commits / 188b7f9b8c
Fork 0

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, 13:30 UTC

[Performance][ROCm] Add skinny gemms for unquantized linear on ROCm (#15830)

Signed-off-by: charlifu <charlifu@amd.com> Co-authored-by: Tyler Michael Smith <tysmith@redhat.com>

Charlie Fu committed Apr 21, 2025 at 22:46 UTC 188b7f9b8c30b09afb507beca480a1ee7d41c154
12 files changed +1957 -95
CMakeLists.txt
+1
@@ -678,6 +678,7 @@ if(VLLM_GPU_LANG STREQUAL "HIP")
678 #
679 set(VLLM_ROCM_EXT_SRC
680 "csrc/rocm/torch_bindings.cpp"
681 + "csrc/rocm/skinny_gemms.cu"
682 "csrc/rocm/attention.cu")
683
684 define_gpu_extension_target(
csrc/rocm/ops.h
+9
@@ -2,6 +2,15 @@
2
3 #include <torch/all.h>
4
5 +torch::Tensor LLMM1(at::Tensor& in_a, at::Tensor& in_b,
6 + const int64_t rows_per_block);
7 +
8 +torch::Tensor wvSplitK(at::Tensor& in_a, at::Tensor& in_b,
9 + const int64_t CuCount);
10 +
11 +void wvSplitKQ(at::Tensor& in_a, at::Tensor& in_b, at::Tensor& out_c,
12 + at::Tensor& scale_a, at::Tensor& scale_b, const int64_t CuCount);
13 +
14 void paged_attention(torch::Tensor& out, torch::Tensor& exp_sums,
15 torch::Tensor& max_logits, torch::Tensor& tmp_out,
16 torch::Tensor& query, torch::Tensor& key_cache,
csrc/rocm/skinny_gemms.cu new
+1600
@@ -0,0 +1,1600 @@
1 +#include <torch/all.h>
2 +#include <ATen/cuda/CUDAContext.h>
3 +#include <c10/cuda/CUDAGuard.h>
4 +
5 +#include <cuda_runtime.h>
6 +#include <cuda_fp16.h>
7 +#include <cuda_bf16.h>
8 +
9 +#include <stdexcept>
10 +#include <algorithm>
11 +
12 +#include "cuda_compat.h"
13 +#include "dispatch_utils.h"
14 +#include "quantization/fp8/common.cuh"
15 +
16 +#if defined(__HIPCC__) && (defined(__gfx90a__) || defined(__gfx942__))
17 + #define __HIP__MI300_MI250__
18 +#endif
19 +
20 +#if defined(__HIPCC__) && defined(__gfx942__)
21 + #define __HIP__MI300__
22 +#endif
23 +
24 +#if defined(NDEBUG)
25 + #undef NDEBUG
26 + #include <assert.h>
27 + #define UNREACHABLE_CODE assert(false);
28 + #define NDEBUG
29 +#else
30 + #define UNREACHABLE_CODE assert(false);
31 +#endif
32 +
33 +template <typename T>
34 +struct scalar {};
35 +
36 +template <typename T>
37 +struct scalar2 {};
38 +
39 +template <typename T>
40 +__device__ __forceinline__ float2 __s22float2(T v);
41 +
42 +template <typename T>
43 +__device__ __forceinline__ T __float2s(float v);
44 +
45 +template <typename T>
46 +__device__ __forceinline__ T __float22s2_rn(float2 v);
47 +
48 +// Definitions and cvt functions for fp16
49 +template <>
50 +struct scalar<c10::Half> {
51 + using type = half;
52 +};
53 +
54 +template <>
55 +struct scalar2<c10::Half> {
56 + using type = __half2;
57 +};
58 +
59 +template <>
60 +__device__ __forceinline__ half __float2s(float v) {
61 + return __float2half(v);
62 +}
63 +
64 +template <>
65 +__device__ __forceinline__ float2 __s22float2(__half2 v) {
66 + return __half22float2(v);
67 +}
68 +
69 +template <>
70 +__device__ __forceinline__ __half2 __float22s2_rn(float2 v) {
71 + return __float22half2_rn(v);
72 +}
73 +
74 +// Definitions and cvt functions for bf16
75 +template <>
76 +struct scalar<c10::BFloat16> {
77 + using type = __hip_bfloat16;
78 +};
79 +
80 +template <>
81 +struct scalar2<c10::BFloat16> {
82 + using type = __hip_bfloat162;
83 +};
84 +
85 +template <>
86 +__device__ __forceinline__ __hip_bfloat16 __float2s(float v) {
87 + return __float2bfloat16(v);
88 +}
89 +
90 +template <>
91 +__device__ __forceinline__ float2 __s22float2(__hip_bfloat162 v) {
92 + return __bfloat1622float2(v);
93 +}
94 +
95 +template <>
96 +__device__ __forceinline__ __hip_bfloat162 __float22s2_rn(float2 v) {
97 + return __float22bfloat162_rn(v);
98 +}
99 +
100 +template <typename T>
101 +__device__ __forceinline__ T loadnt(T* addr) {
102 + return __builtin_nontemporal_load(addr);
103 +}
104 +
105 +__device__ __forceinline__ float4 load_ntmprl(const float4* addr) {
106 + auto addr_alias = reinterpret_cast<const float*>(addr);
107 + auto dat0 = loadnt(addr_alias);
108 + auto dat1 = loadnt(addr_alias + 1);
109 + auto dat2 = loadnt(addr_alias + 2);
110 + auto dat3 = loadnt(addr_alias + 3);
111 + return make_float4(dat0, dat1, dat2, dat3);
112 +}
113 +
114 +// TBlock fetches entire rows of A, and entire col of B (K dimension); assume
115 +// N=1 for time being grid is M/A_NUM_ROWS blocks
116 +template <typename scalar_t, int NUM_A_ROWS_PER_BLOCK>
117 +__global__ void LLGemm1_kernel(const scalar_t* in_a, const scalar_t* in_b,
118 + scalar_t* out_c, const int K) {
119 + using scalar2_t = typename scalar2<scalar_t>::type;
120 + auto af4 = reinterpret_cast<const float4*>(in_a);
121 + auto bf4 = reinterpret_cast<const scalar2_t*>(in_b);
122 + auto c = reinterpret_cast<scalar2_t*>(out_c);
123 + __shared__ float red_smem[NUM_A_ROWS_PER_BLOCK][WARP_SIZE];
124 + const int row_addr = blockIdx.x * NUM_A_ROWS_PER_BLOCK * K / 8;
125 + const int threadid = threadIdx.x;
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;
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];
134 + scalar2_t acch2;
135 + scalar2_t oval;
136 +
137 + // As we later use warp shuffle operations, we may have more threads in the
138 + // block than the actual available data, hence the if guard here.
139 + if (threadid * 8 < K) {
140 +#pragma unroll
141 + for (int i = 0; i < NUM_A_ROWS_PER_BLOCK; i++) {
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 + }
146 +
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 +
152 + scalar2_t Af2;
153 + scalar2_t Bf2;
154 + float2 S;
155 +
156 + auto Ah2ptr = reinterpret_cast<scalar2_t*>(&rowA_elem4);
157 + scalar2_t* ah2lptr;
158 +
159 +#pragma unroll
160 + for (int i = 0; i < NUM_A_ROWS_PER_BLOCK; i++) {
161 + // Multiply-add on 8 scalar_t.
162 + ah2lptr = Ah2ptr + i * 4;
163 + Af2 = *(ah2lptr);
164 + acch2 = __hmul2(Af2, colB_elem4x);
165 + Af2 = *(ah2lptr + 1);
166 + acch2 = __hfma2(Af2, colB_elem4y, acch2);
167 + Af2 = *(ah2lptr + 2);
168 + acch2 = __hfma2(Af2, colB_elem4z, acch2);
169 + Af2 = *(ah2lptr + 3);
170 + acch2 = __hfma2(Af2, colB_elem4w, acch2);
171 + S = __s22float2(acch2);
172 +
173 + // See comment above concerning the if guard.
174 + acc[i] = (threadid * 8 < K ? S.x + S.y : 0.f);
175 + }
176 +
177 +// all reduce across warp.
178 +#pragma unroll
179 + for (int mask = WARP_SIZE / 2; mask >= 1; mask /= 2) {
180 +#pragma unroll
181 + for (int i = 0; i < NUM_A_ROWS_PER_BLOCK; i++) {
182 + acc[i] += __shfl_xor(acc[i], mask);
183 + }
184 + }
185 +
186 + // Warp leaders store the data to shared memory.
187 + if (lane < NUM_A_ROWS_PER_BLOCK) {
188 + red_smem[lane][warp] = acc[lane];
189 + }
190 +
191 + // Make sure the data is in shared memory.
192 + __syncthreads();
193 +
194 + if (qwarpid < NUM_A_ROWS_PER_BLOCK) {
195 + acc[qwarpid] = qthreadid < num_warps ? red_smem[qwarpid][qthreadid] : 0.f;
196 + for (int mask = num_warps / 2; mask >= 1; mask /= 2) {
197 + acc[qwarpid] += __shfl_xor(acc[qwarpid], mask);
198 + }
199 + float oval2 = __shfl_xor(acc[qwarpid], num_warps);
200 +
201 + if (lane % (num_warps * 2) == 0) {
202 + oval = __float22s2_rn<scalar2_t>(make_float2(acc[qwarpid], oval2));
203 + c[blockIdx.x * NUM_A_ROWS_PER_BLOCK / 2 + qwarpid / 2] = oval;
204 + }
205 + }
206 +}
207 +
208 +torch::Tensor LLMM1(at::Tensor& in_a, at::Tensor& in_b,
209 + const int64_t rows_per_block) {
210 + auto M = in_a.size(0);
211 + auto K = in_a.size(1);
212 + auto N = in_b.size(0);
213 +
214 + TORCH_CHECK(N == 1, "Row number of activation tensor must be 1.");
215 + TORCH_CHECK(in_a.dtype() == in_b.dtype());
216 + TORCH_CHECK(in_b.dtype() == torch::kFloat16 ||
217 + in_b.dtype() == torch::kBFloat16);
218 +
219 + auto out_c = torch::empty(
220 + {N, M}, torch::TensorOptions().dtype(in_b.dtype()).device(in_b.device()));
221 +
222 + // NUM_TREADS need to be a multiple of WARP_SIZE, as we are using warp shuffle
223 + // operations.
224 + 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);
228 +
229 + int NUM_BLOCKS = M / rows_per_block;
230 +
231 + const at::cuda::OptionalCUDAGuard device_guard(device_of(in_b));
232 + const cudaStream_t stream = at::cuda::getCurrentCUDAStream();
233 +
234 + // call the kernel function...
235 + AT_DISPATCH_REDUCED_FLOATING_TYPES(in_b.scalar_type(), "LLGemm1", [&] {
236 + auto a_ptr = in_a.data_ptr<scalar_t>();
237 + auto b_ptr = in_b.data_ptr<scalar_t>();
238 + auto c_ptr = out_c.data_ptr<scalar_t>();
239 + if (rows_per_block == 2) {
240 + LLGemm1_kernel<scalar_t, 2>
241 + <<<NUM_BLOCKS, NUM_THREADS, 0, stream>>>(a_ptr, b_ptr, c_ptr, K);
242 + } else if (rows_per_block == 4) {
243 + LLGemm1_kernel<scalar_t, 4>
244 + <<<NUM_BLOCKS, NUM_THREADS, 0, stream>>>(a_ptr, b_ptr, c_ptr, K);
245 + } else if (rows_per_block == 8) {
246 + LLGemm1_kernel<scalar_t, 8>
247 + <<<NUM_BLOCKS, NUM_THREADS, 0, stream>>>(a_ptr, b_ptr, c_ptr, K);
248 + } else if (rows_per_block == 16) {
249 + LLGemm1_kernel<scalar_t, 16>
250 + <<<NUM_BLOCKS, NUM_THREADS, 0, stream>>>(a_ptr, b_ptr, c_ptr, K);
251 + } else {
252 + NUM_BLOCKS = M / 4;
253 + LLGemm1_kernel<scalar_t, 4>
254 + <<<NUM_BLOCKS, NUM_THREADS, 0, stream>>>(a_ptr, b_ptr, c_ptr, K);
255 + }
256 + });
257 +
258 + return out_c;
259 +}
260 +
261 +#define DOT2C(V0, V2, V3) \
262 + if constexpr (std::is_same_v<scalar_t, half>) { \
263 + asm("v_dot2c_f32_f16 %0, %2, %3" : "=v"(V0) : "0"(V0), "v"(V2), "v"(V3)); \
264 + } else if constexpr (std::is_same_v<scalar_t, __hip_bfloat16>) { \
265 + float2 s = __bfloat1622float2(*((__hip_bfloat162*)(&(V2)))) * \
266 + __bfloat1622float2(*((__hip_bfloat162*)(&(V3)))); \
267 + V0 += (s.x + s.y); \
268 + }
269 +
270 +#if defined(__HIP__MI300_MI250__) // TODO: Add NAVI support
271 +// This version targets cases where A[] fits LDS capacity
272 +template <typename scalar_t, int THRDS, int YTILE, int WvPrGrp, int A_CHUNK,
273 + int UNRL, int N>
274 +__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 + using scalar8 =
279 + __attribute__((__vector_size__((A_CHUNK / 2) * sizeof(float)))) float;
280 + union bigType {
281 + scalar_t h[A_CHUNK];
282 + float f[A_CHUNK / 2];
283 + float2 f2[A_CHUNK / 4];
284 + double d[A_CHUNK / 4];
285 + scalar8 h8;
286 + };
287 +
288 + //----------------------------------------------------
289 + // Reserving 64 KB of LDS to have 1 WG / CU
290 + // Goal is to bring the activation matrix A to the LDS
291 + // and use it across the lifetime of the work group
292 + // TODO: When activation matrix is larger than 64 KB
293 + // then this is not goint to work!
294 + //----------------------------------------------------
295 + __shared__ scalar_t s[1024 * 32];
296 +
297 + //----------------------------------------------------
298 + // Fetch the activation matrix to LDS
299 + // Loop iteration:
300 + // - Each thread (lane) is fetching 8 elements (A_Chunk)
301 + // - Each wave will fetch 64*8=> 512 elements
302 + // - Each WG will fetch 512 * 16 => 8K elements
303 + // - Then the WG will move to another 8 K elements
304 + // TODO: Logic below will only work when K is multiple of 8
305 + //----------------------------------------------------
306 + for (uint32_t k = 0; k < min(K * N, 32 * 1024);
307 + k += THRDS * WvPrGrp * A_CHUNK) {
308 + uint32_t k_in = k + ((threadIdx.y * THRDS + threadIdx.x) * A_CHUNK);
309 +
310 + if (k_in >= min(K * N, 32 * 1024)) break;
311 +
312 + *((bigType*)(&s[k_in])) = *((bigType*)(&A[k_in]));
313 + }
314 + __syncthreads();
315 +
316 + if (threadIdx.y >= _WvPrGrp) return;
317 +
318 + uint32_t m = (blockIdx.x * _WvPrGrp + (threadIdx.y % _WvPrGrp)) * YTILE;
319 +
320 + float sum[N][YTILE];
321 +
322 + //----------------------------------------------------
323 + // Each wave works on a single column of weight matrix.
324 + // There are 16 waves per WG, and hence, each WG is
325 + // working on 16 columns of weight matrix. Moreover,
326 + // we tile in column direction by YTILE, so when YTILE=1
327 + // the above math is right, however, when YTILE=2 then
328 + // each wave will be working on 2 columns and WG will
329 + // be working on 32 columns.
330 + //
331 + // Top level loop that makes WGs persistent!
332 + // - WGs iterates across columns of weight matrix
333 + // - Each wave within WG works on a given column(s)
334 + // - After completing first set of columns, WGs start
335 + // working on the next set of available columns
336 + //----------------------------------------------------
337 + while (m < M) {
338 + //----------------------------------------------------
339 + // 'sum' accumulates the matrix A x B computation
340 + // split across 64 lanes.
341 + //
342 + // YTILE represents how many column of weight matrix
343 + // are being worked on by each wave.
344 + //----------------------------------------------------
345 + for (int i = 0; i < YTILE; i++)
346 + for (int n = 0; n < N; n++) sum[n][i] = 0;
347 +
348 + bigType bigA[N][UNRL];
349 + bigType bigB[YTILE][UNRL];
350 + //----------------------------------------------------
351 + // Fetch weight matrix B in interleaved K-split!
352 + // - Each thread (lane) is fetching 8 elements (A_Chunk)
353 + // - Each wave will fetch 64*8=> 512 elements (1024B)
354 + // - YTILE represents the number of column being serviced
355 + // by wave
356 + // - Loop for fetching weight matrix (B) are unrolled
357 + //
358 + // Fetch activation matrix A from LDS
359 + // - Loop for fetching activation matrix (A) are unrolled
360 + //
361 + // Finally, do the matrix multiplication in an unrolled
362 + // fashion. This provides lot of food for compiler
363 + // scheduling.
364 + //
365 + // TODO: Logic below will only work when K is multiple of 8
366 + //----------------------------------------------------
367 + // for (uint32_t k1 = 0; k1 < K; k1 += THRDS * A_CHUNK * UNRL) {
368 + for (uint32_t k1 = 0; k1 < K; k1 += THRDS * A_CHUNK * UNRL) {
369 + // Fetch the weight matrix from memory!
370 + #pragma unroll
371 + for (uint32_t k2 = 0; k2 < UNRL; k2++) {
372 + uint32_t k = k1 + k2 * THRDS * A_CHUNK;
373 + uint32_t k_ = k + threadIdx.x * A_CHUNK;
374 + if (k_ >= K) break;
375 +
376 + 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])));
395 + }
396 +
397 + // Fetch activation matrix from either just LDS or from both LDS / memory
398 + #pragma unroll
399 + for (uint32_t k2 = 0; k2 < UNRL; k2++) {
400 + uint32_t k = k1 + k2 * THRDS * A_CHUNK;
401 + uint32_t k_ = k + threadIdx.x * A_CHUNK;
402 + if (k_ >= K) break;
403 +
404 + // Fetch A activation matrix in interleaved fashion from LDS or memory
405 +
406 + for (int n = 0; n < N; n++) {
407 + bigA[n][k2] = *((const bigType*)(&(s[k_ + K * n])));
408 + }
409 + }
410 +
411 + // Do the matrix multiplication in interleaved manner
412 + #pragma unroll
413 + for (uint32_t k2 = 0; k2 < UNRL; k2++) {
414 + uint32_t k = k1 + k2 * THRDS * A_CHUNK;
415 + uint32_t k_ = k + threadIdx.x * A_CHUNK;
416 + if (k_ >= K) break;
417 + // Do the matrix multiplication of activation and weight matrix
418 + // - Remember the accumulation is happening for K-split of 64!
419 + #pragma unroll
420 + for (uint32_t n = 0; n < N; n++) {
421 + #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 + }
448 + }
449 + }
450 + }
451 + }
452 +
453 + //----------------------------------------------------
454 + // Final reduction step using shuffle
455 + //----------------------------------------------------
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]));
476 + }
477 + }
478 + if (threadIdx.x == 63) {
479 + 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]);
483 + }
484 + }
485 + }
486 +
487 + m += CuCount * _WvPrGrp * YTILE;
488 + }
489 +}
490 +#else // !defined(__HIP__MI300_MI250__) TODO: Add NAVI support
491 +template <typename scalar_t, int THRDS, int YTILE, int WvPrGrp, int A_CHUNK,
492 + int UNRL, int N>
493 +__global__ void wvSplitK_hf_sml_(const int K, const int M, const scalar_t* B,
494 + const scalar_t* __restrict__ A, scalar_t* C,
495 + const int _WvPrGrp, const int CuCount) {
496 + UNREACHABLE_CODE
497 +}
498 +#endif // defined(__HIP__MI300_MI250__) TODO: Add NAVI support
499 +
500 +#if defined(__HIP__MI300_MI250__) // TODO: Add NAVI support
501 +// This version targets cases where A[] marginally exceeds LDS capacity
502 +template <typename scalar_t, int THRDS, int YTILE, int WvPrGrp, int A_CHUNK,
503 + int UNRL, int N>
504 +__global__ void __launch_bounds__(WvPrGrp* THRDS)
505 + wvSplitK_hf_(const int K, const int M, const scalar_t* B,
506 + const scalar_t* __restrict__ A, scalar_t* C,
507 + const int _WvPrGrp, const int CuCount) {
508 + using scalar8 =
509 + __attribute__((__vector_size__((A_CHUNK / 2) * sizeof(float)))) float;
510 + union bigType {
511 + scalar_t h[A_CHUNK];
512 + float f[A_CHUNK / 2];
513 + float2 f2[A_CHUNK / 4];
514 + double d[A_CHUNK / 4];
515 + scalar8 h8;
516 + };
517 +
518 + //----------------------------------------------------
519 + // Reserving 64 KB of LDS to have 1 WG / CU
520 + // Goal is to bring the activation matrix A to the LDS
521 + // and use it across the lifetime of the work group
522 + // TODO: When activation matrix is larger than 64 KB
523 + // then this is not goint to work!
524 + //----------------------------------------------------
525 + __shared__ scalar_t s[1024 * 32];
526 +
527 + //----------------------------------------------------
528 + // Computation of columns that need to be committed to memory!
529 + //----------------------------------------------------
530 + uint32_t commitColumn[YTILE];
531 + for (uint32_t i = 0; i < YTILE; i++) {
532 + commitColumn[i] = 1;
533 + }
534 +
535 + //----------------------------------------------------
536 + // Indexing function into the column of weight matrix B
537 + // Algorithm does 64 lane k-splitting / wave and uses
538 + // WG ID and Thread ID to find the index.
539 + //----------------------------------------------------
540 + // int _WvPrGrp = mindiv(N, CuCount * YTILE, WvPrGrp);
541 + uint32_t m = (blockIdx.x * _WvPrGrp + threadIdx.y) * YTILE;
542 +
543 + // Check whether there will be fragmenation!
544 + // This will happen only for the last wave!
545 + if (m < M && (m + YTILE) >= M) {
546 + uint32_t startColumn = M - YTILE;
547 + for (uint32_t i = 0; i < (m - startColumn); i++) {
548 + commitColumn[i] = 0;
549 + }
550 + m = startColumn;
551 + }
552 +
553 + //----------------------------------------------------
554 + // Fetch the activation matrix to LDS
555 + // Loop iteration:
556 + // - Each thread (lane) is fetching 8 elements (A_Chunk)
557 + // - Each wave will fetch 64*8=> 512 elements
558 + // - Each WG will fetch 512 * 16 => 8K elements
559 + // - Then the WG will move to another 8 K elements
560 + // TODO: Logic below will only work when K is multiple of 8
561 + //----------------------------------------------------
562 + for (uint32_t k = 0; k < min(K * N, 32 * 1024);
563 + k += THRDS * WvPrGrp * A_CHUNK) {
564 + uint32_t k_in = k + ((threadIdx.y * THRDS + threadIdx.x) * A_CHUNK);
565 +
566 + if (k_in >= min(K * N, 32 * 1024)) break;
567 +
568 + *((bigType*)(&s[k_in])) = *((bigType*)(&A[k_in]));
569 + }
570 +
571 + __syncthreads();
572 +
573 + if (threadIdx.y >= _WvPrGrp) return;
574 +
575 + float sum[N][YTILE];
576 +
577 + //----------------------------------------------------
578 + // Each wave works on a single column of weight matrix.
579 + // There are 16 waves per WG, and hence, each WG is
580 + // working on 16 columns of weight matrix. Moreover,
581 + // we tile in column direction by YTILE, so when YTILE=1
582 + // the above math is right, however, when YTILE=2 then
583 + // each wave will be working on 2 columns and WG will
584 + // be working on 32 columns.
585 + //
586 + // Top level loop that makes WGs persistent!
587 + // - WGs iterates across columns of weight matrix
588 + // - Each wave within WG works on a given column(s)
589 + // - After completing first set of columns, WGs start
590 + // working on the next set of available columns
591 + //----------------------------------------------------
592 + while (m < M) {
593 + //----------------------------------------------------
594 + // 'sum' accumulates the matrix A x B computation
595 + // split across 64 lanes.
596 + //
597 + // YTILE represents how many column of weight matrix
598 + // are being worked on by each wave.
599 + //----------------------------------------------------
600 + for (int i = 0; i < YTILE; i++)
601 + for (int n = 0; n < N; n++) sum[n][i] = 0;
602 +
603 + bigType bigA[N][UNRL];
604 + bigType bigB[YTILE][UNRL];
605 + //----------------------------------------------------
606 + // Fetch weight matrix B in interleaved K-split!
607 + // - Each thread (lane) is fetching 8 elements (A_Chunk)
608 + // - Each wave will fetch 64*8=> 512 elements (1024B)
609 + // - YTILE represents the number of column being serviced
610 + // by wave
611 + // - Loop for fetching weight matrix (B) are unrolled
612 + //
613 + // Fetch activation matrix A from LDS
614 + // - Loop for fetching activation matrix (A) are unrolled
615 + //
616 + // Finally, do the matrix multiplication in an unrolled
617 + // fashion. This provides lot of food for compiler
618 + // scheduling.
619 + //
620 + // TODO: Logic below will only work when K is multiple of 8
621 + //----------------------------------------------------
622 + for (uint32_t k1 = 0; k1 < K; k1 += THRDS * A_CHUNK * UNRL) {
623 + // Fetch the weight matrix from memory!
624 + #pragma unroll
625 + for (uint32_t k2 = 0; k2 < UNRL; k2++) {
626 + uint32_t k = k1 + k2 * THRDS * A_CHUNK;
627 + uint32_t k_ = k + threadIdx.x * A_CHUNK;
628 + if (k_ >= K) break;
629 +
630 + 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])));
649 + }
650 +
651 + // Fetch activation matrix from either just LDS or from both LDS / memory
652 + #pragma unroll
653 + for (uint32_t k2 = 0; k2 < UNRL; k2++) {
654 + uint32_t k = k1 + k2 * THRDS * A_CHUNK;
655 + uint32_t k_ = k + threadIdx.x * A_CHUNK;
656 + if (k_ >= K) break;
657 +
658 + // Fetch A activation matrix in interleaved fashion from LDS or memory
659 +
660 + for (int n = 0; n < N; n++) {
661 + if (k_ + K * n < 32 * 1024)
662 + bigA[n][k2] = *((const bigType*)(&(s[k_ + K * n])));
663 + else
664 + bigA[n][k2] = *((const bigType*)(&(A[k_ + K * n])));
665 + }
666 + }
667 +
668 + // Do the matrix multiplication in interleaved manner
669 + #pragma unroll
670 + for (uint32_t n = 0; n < N; n++) {
671 + #pragma unroll
672 + for (uint32_t k2 = 0; k2 < UNRL; k2++) {
673 + uint32_t k = k1 + k2 * THRDS * A_CHUNK;
674 + uint32_t k_ = k + threadIdx.x * A_CHUNK;
675 + if (k_ >= K) break;
676 + // Do the matrix multiplication of activation and weight matrix
677 + // - Remember the accumulation is happening for K-split of 64!
678 + #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 + }
705 + }
706 + }
707 + }
708 + }
709 +
710 + //----------------------------------------------------
711 + // Final reduction step using shuffle
712 + //----------------------------------------------------
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]));
733 + }
734 + }
735 +
736 + if (threadIdx.x == 63) {
737 + 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]);
741 + }
742 + }
743 + }
744 +
745 + m += CuCount * _WvPrGrp * YTILE;
746 +
747 + // Check whether there will be fragmenation!
748 + // This will happen only for the last wave!
749 + if (m < M && (m + YTILE) >= M) {
750 + uint32_t startColumn = M - YTILE;
751 + for (uint32_t i = 0; i < (m - startColumn); i++) {
752 + commitColumn[i] = 0;
753 + }
754 + m = startColumn;
755 + }
756 + }
757 +}
758 +
759 +#else // !defined(__HIP__MI300_MI250__) TODO: Add NAVI support
760 +template <typename scalar_t, int THRDS, int YTILE, int WvPrGrp, int A_CHUNK,
761 + int UNRL, int N>
762 +__global__ void wvSplitK_hf_(const int K, const int M, const scalar_t* B,
763 + const scalar_t* __restrict__ A, scalar_t* C,
764 + const int _WvPrGrp, const int CuCount) {
765 + UNREACHABLE_CODE
766 +}
767 +#endif // defined(__HIP__MI300_MI250__) TODO: Add NAVI support
768 +
769 +#if defined(__HIP__MI300_MI250__) // TODO: Add NAVI support
770 +// This version targets big A[] cases, where it is much larger than LDS capacity
771 +template <typename scalar_t, int THRDS, int YTILE, int WvPrGrp, int A_CHUNK,
772 + int UNRL, int N>
773 +__global__ void __launch_bounds__(WvPrGrp* THRDS)
774 + wvSplitK_hf_big_(const int K, const int M, const scalar_t* B,
775 + const scalar_t* __restrict__ A, scalar_t* C,
776 + const int _WvPrGrp, const int CuCount) {
777 + using scalar8 =
778 + __attribute__((__vector_size__((A_CHUNK / 2) * sizeof(float)))) float;
779 +
780 + union bigType {
781 + scalar_t h[A_CHUNK];
782 + float f[A_CHUNK / 2];
783 + float2 f2[A_CHUNK / 4];
784 + double d[A_CHUNK / 4];
785 + scalar8 h8;
786 + };
787 +
788 + //----------------------------------------------------
789 + // Reserving 64 KB of LDS to have 1 WG / CU
790 + // Goal is to bring the activation matrix A to the LDS
791 + // and use it across the lifetime of the work group
792 + // TODO: When activation matrix is larger than 64 KB
793 + // then this is not goint to work!
794 + //----------------------------------------------------
795 + __shared__ scalar_t s[1024 * 32];
796 +
797 + //----------------------------------------------------
798 + // Computation of columns that need to be committed to memory!
799 + //----------------------------------------------------
800 + uint32_t commitColumn[YTILE];
801 + for (uint32_t i = 0; i < YTILE; i++) {
802 + commitColumn[i] = 1;
803 + }
804 +
805 + // int _WvPrGrp = mindiv(N, CuCount * YTILE, WvPrGrp);
806 + if (threadIdx.y >= _WvPrGrp) return;
807 +
808 + //----------------------------------------------------
809 + // Indexing function into the column of weight matrix B
810 + // Algorithm does 64 lane k-splitting / wave and uses
811 + // WG ID and Thread ID to find the index.
812 + //----------------------------------------------------
813 + uint32_t m = (blockIdx.x * _WvPrGrp + threadIdx.y) * YTILE;
814 +
815 + // Check whether there will be fragmenation!
816 + // This will happen only for the last wave!
817 + if (m < M && (m + YTILE) >= M) {
818 + uint32_t startColumn = M - YTILE;
819 + for (uint32_t i = 0; i < (m - startColumn); i++) {
820 + commitColumn[i] = 0;
821 + }
822 + m = startColumn;
823 + }
824 +
825 + //----------------------------------------------------
826 + // Fetch the activation matrix to LDS
827 + // Loop iteration:
828 + // - Each thread (lane) is fetching 8 elements (A_Chunk)
829 + // - Each wave will fetch 64*8=> 512 elements
830 + // - Each WG will fetch 512 * 16 => 8K elements
831 + // - Then the WG will move to another 8 K elements
832 + // TODO: Logic below will only work when K is multiple of 8
833 + //----------------------------------------------------
834 + #define PCML
835 + #ifndef PCML
836 + for (uint32_t k = 0; k < min(K * N, 32 * 1024);
837 + k += THRDS * WvPrGrp * A_CHUNK) {
838 + uint32_t k_in = k + ((threadIdx.y * THRDS + threadIdx.x) * A_CHUNK);
839 +
840 + if (k_in >= min(K * N, 32 * 1024)) break;
841 +
842 + *((bigType*)(&s[k_in])) = *((bigType*)(&A[k_in]));
843 + }
844 + __syncthreads();
845 + #endif
846 +
847 + #define TUC (THRDS * UNRL * A_CHUNK)
848 + uint32_t kBase = 0;
849 + // find biggest k size that fits in LDS
850 + uint32_t kFit = (32 * 1024) / N;
851 + // kFit = (kFit%TWC==0) ? kFit : (kFit-kFit%TWC+TWC); //round up to multiple
852 + // of TUC
853 + kFit = (kFit % TUC == 0)
854 + ? kFit
855 + : (kFit - kFit % TUC); // round up to multiple of TUC
856 + // if (kFit == 0) kFit = TUC;
857 + kFit = min(kFit, K);
858 +
859 + float sum[N][YTILE];
860 +
861 + //----------------------------------------------------
862 + // Each wave works on a single column of weight matrix.
863 + // There are 16 waves per WG, and hence, each WG is
864 + // working on 16 columns of weight matrix. Moreover,
865 + // we tile in column direction by YTILE, so when YTILE=1
866 + // the above math is right, however, when YTILE=2 then
867 + // each wave will be working on 2 columns and WG will
868 + // be working on 32 columns.
869 + //
870 + // Top level loop that makes WGs persistent!
871 + // - WGs iterates across columns of weight matrix
872 + // - Each wave within WG works on a given column(s)
873 + // - After completing first set of columns, WGs start
874 + // working on the next set of available columns
875 + //----------------------------------------------------
876 + #ifdef PCML
877 + int YW = (YTILE * _WvPrGrp);
878 + uint32_t Mrndp = (M % YW == 0) ? M : (M - M % YW + YW);
879 + while (m < Mrndp) {
880 + #else
881 + while (m < M) {
882 + #endif
883 + //----------------------------------------------------
884 + // 'sum' accumulates the matrix A x B computation
885 + // split across 64 lanes.
886 + //
887 + // YTILE represents how many column of weight matrix
888 + // are being worked on by each wave.
889 + //----------------------------------------------------
890 + for (int i = 0; i < YTILE; i++)
891 + for (int n = 0; n < N; n++) sum[n][i] = 0;
892 +
893 + bigType bigA[N][UNRL];
894 + bigType bigB[YTILE][UNRL];
895 + //----------------------------------------------------
896 + // Fetch weight matrix B in interleaved K-split!
897 + // - Each thread (lane) is fetching 8 elements (A_Chunk)
898 + // - Each wave will fetch 64*8=> 512 elements (1024B)
899 + // - YTILE represents the number of column being serviced
900 + // by wave
901 + // - Loop for fetching weight matrix (B) are unrolled
902 + //
903 + // Fetch activation matrix A from LDS
904 + // - Loop for fetching activation matrix (A) are unrolled
905 + //
906 + // Finally, do the matrix multiplication in an unrolled
907 + // fashion. This provides lot of food for compiler
908 + // scheduling.
909 + //
910 + // TODO: Logic below will only work when K is multiple of 8
911 + //----------------------------------------------------
912 + for (uint32_t k1 = 0; k1 < K; k1 += THRDS * A_CHUNK * UNRL) {
913 + #ifdef PCML
914 + if ((k1 == 0) || (k1 == kBase + kFit)) { // load next chunk of A[] to LDS
915 + if (k1 != 0) kBase += kFit;
916 + __syncthreads();
917 + for (uint32_t k = 0; k < kFit; k += THRDS * _WvPrGrp * A_CHUNK) {
918 + uint32_t kOff = k + ((threadIdx.y * THRDS + threadIdx.x) * A_CHUNK);
919 + if (kBase + kOff >= K) break;
920 + if (kOff >= kFit) break;
921 + for (uint32_t n = 0; n < N; n++) {
922 + uint32_t k_in = kBase + n * K + kOff;
923 + uint32_t k_ot = n * kFit + kOff;
924 + *((bigType*)(&s[k_ot])) = *((bigType*)(&A[k_in]));
925 + }
926 + }
927 + __syncthreads();
928 + }
929 + if (m >= M) continue;
930 + #endif
931 +
932 + // Fetch the weight matrix from memory!
933 + #pragma unroll
934 + for (uint32_t k2 = 0; k2 < UNRL; k2++) {
935 + uint32_t k = k1 + k2 * THRDS * A_CHUNK;
936 + uint32_t k_ = k + threadIdx.x * A_CHUNK;
937 + if (k_ >= K) break;
938 +
939 + 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])));
958 + }
959 +
960 + // Fetch activation matrix from either just LDS or from both LDS / memory
961 + #pragma unroll
962 + for (uint32_t k2 = 0; k2 < UNRL; k2++) {
963 + uint32_t k = k1 + k2 * THRDS * A_CHUNK;
964 + uint32_t k_ = k + threadIdx.x * A_CHUNK;
965 + if (k_ >= K) break;
966 +
967 + // Fetch A activation matrix in interleaved fashion from LDS or memory
968 +
969 + for (int n = 0; n < N; n++) {
970 + #ifdef PCML
971 + bigA[n][k2] = *((const bigType*)(&(s[k_ - kBase + kFit * n])));
972 + #else
973 + if (k_ + K * n < 32 * 1024)
974 + bigA[n][k2] = *((const bigType*)(&(s[k_ + K * n])));
975 + else
976 + bigA[n][k2] = *((const bigType*)(&(A[k_ + K * n])));
977 + #endif
978 + }
979 + }
980 +
981 + // Do the matrix multiplication in interleaved manner
982 + #pragma unroll
983 + for (uint32_t k2 = 0; k2 < UNRL; k2++) {
984 + uint32_t k = k1 + k2 * THRDS * A_CHUNK;
985 + uint32_t k_ = k + threadIdx.x * A_CHUNK;
986 + if (k_ >= K) break;
987 + #pragma unroll
988 + for (uint32_t n = 0; n < N; n++) {
989 + // Do the matrix multiplication of activation and weight matrix
990 + // - Remember the accumulation is happening for K-split of 64!
991 + #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 + }
1018 + }
1019 + }
1020 + }
1021 + }
1022 +
1023 + #ifdef PCML
1024 + if (m >= M) {
1025 + m += CuCount * _WvPrGrp * YTILE;
1026 + kBase = 0;
1027 + continue;
1028 + }
1029 + #endif
1030 +
1031 + //----------------------------------------------------
1032 + // Final reduction step using shuffle
1033 + //----------------------------------------------------
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]));
1054 + }
1055 + }
1056 +
1057 + if (threadIdx.x == 63) {
1058 + 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]);
1062 + }
1063 + }
1064 + }
1065 +
1066 + m += CuCount * _WvPrGrp * YTILE;
1067 + kBase = 0;
1068 +
1069 + // Check whether there will be fragmenation!
1070 + // This will happen only for the last wave!
1071 + if (m < M && (m + YTILE) >= M) {
1072 + uint32_t startColumn = M - YTILE;
1073 + for (uint32_t i = 0; i < (m - startColumn); i++) {
1074 + commitColumn[i] = 0;
1075 + }
1076 + m = startColumn;
1077 + }
1078 + }
1079 +}
1080 +#else // !defined(__HIP__MI300_MI250__) TODO: Add NAVI support
1081 +template <typename scalar_t, int THRDS, int YTILE, int WvPrGrp, int A_CHUNK,
1082 + int UNRL, int N>
1083 +__global__ void wvSplitK_hf_big_(const int K, const int M, const scalar_t* B,
1084 + const scalar_t* __restrict__ A, scalar_t* C,
1085 + const int _WvPrGrp, const int CuCount) {
1086 + UNREACHABLE_CODE
1087 +}
1088 +#endif // defined(__HIP__MI300_MI250__) TODO: Add NAVI support
1089 +
1090 +int mindiv(int N, int div1, int div2) {
1091 + int nPrRnd = div1 * div2;
1092 + int rnds0 = N / nPrRnd;
1093 + nPrRnd -= div1 * 3;
1094 + int rnds3 = N / nPrRnd;
1095 + nPrRnd -= div1;
1096 + int rnds4 = N / nPrRnd;
1097 + nPrRnd -= div1;
1098 + int rnds5 = N / nPrRnd;
1099 + nPrRnd -= div1;
1100 + int rnds6 = N / nPrRnd;
1101 + nPrRnd -= div1;
1102 + int rnds7 = N / nPrRnd;
1103 + nPrRnd -= div1;
1104 + int rnds8 = N / nPrRnd;
1105 + nPrRnd -= div1;
1106 + int rnds9 = N / nPrRnd;
1107 + nPrRnd -= div1;
1108 + int rtn = div2;
1109 + if (rnds0 == rnds3) rtn = div2 - 3;
1110 + if (rnds0 == rnds4) rtn = div2 - 4;
1111 + if (rnds0 == rnds5) rtn = div2 - 5;
1112 + if (rnds0 == rnds6) rtn = div2 - 6;
1113 + if (rnds0 == rnds7) rtn = div2 - 7;
1114 + if (rnds0 == rnds8) rtn = div2 - 8;
1115 + if (rnds0 == rnds9) rtn = div2 - 9;
1116 + return rtn;
1117 +}
1118 +
1119 +torch::Tensor wvSplitK(at::Tensor& in_a, at::Tensor& in_b,
1120 + const int64_t CuCount) {
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 +
1125 + TORCH_CHECK(in_a.dtype() == in_b.dtype());
1126 + TORCH_CHECK(K_in % 8 == 0, "k % 8 == 0");
1127 + TORCH_CHECK(in_a.dtype() == torch::kFloat16 ||
1128 + in_a.dtype() == torch::kBFloat16);
1129 +
1130 + auto out_c = torch::empty(
1131 + {N_in, M_in},
1132 + torch::TensorOptions().dtype(in_b.dtype()).device(in_b.device()));
1133 +
1134 + dim3 grid(CuCount);
1135 +
1136 + const at::cuda::OptionalCUDAGuard device_guard(device_of(in_a));
1137 + const cudaStream_t stream = at::cuda::getCurrentCUDAStream();
1138 +
1139 +#define WVSPLITK(_WvPrGrp, _YTILEs, _YTILEm, _YTILEb, _UNRLs, _UNRLm, _UNRLb, \
1140 + _N) \
1141 + { \
1142 + dim3 block(64, _WvPrGrp); \
1143 + if ((K_in * N_in <= 32 * 1024) && (M_in % _YTILEs == 0)) { \
1144 + int __wvPrGrp = mindiv(M_in, CuCount * _YTILEs, _WvPrGrp); \
1145 + wvSplitK_hf_sml_<fptype, 64, _YTILEs, _WvPrGrp, 8, _UNRLs, _N> \
1146 + <<<grid, block, 0, stream>>>(K_in, M_in, af4, bf4, c, __wvPrGrp, \
1147 + CuCount); \
1148 + } else if (K_in * N_in <= 32 * 1024 * 1.2) { \
1149 + int __wvPrGrp = mindiv(M_in, CuCount * _YTILEm, _WvPrGrp); \
1150 + wvSplitK_hf_<fptype, 64, _YTILEm, _WvPrGrp, 8, _UNRLm, _N> \
1151 + <<<grid, block, 0, stream>>>(K_in, M_in, af4, bf4, c, __wvPrGrp, \
1152 + CuCount); \
1153 + } else { \
1154 + int __wvPrGrp = mindiv(M_in, CuCount * _YTILEb, _WvPrGrp); \
1155 + wvSplitK_hf_big_<fptype, 64, _YTILEb, _WvPrGrp, 8, _UNRLb, _N> \
1156 + <<<grid, block, 0, stream>>>(K_in, M_in, af4, bf4, c, __wvPrGrp, \
1157 + CuCount); \
1158 + } \
1159 + }
1160 +
1161 + AT_DISPATCH_REDUCED_FLOATING_TYPES(in_b.scalar_type(), "wvSplitK", [&] {
1162 + using fptype = typename scalar<scalar_t>::type;
1163 + fptype* af4 = reinterpret_cast<fptype*>(in_a.data_ptr());
1164 + const fptype* bf4 = reinterpret_cast<const fptype*>(in_b.data_ptr());
1165 + fptype* c = reinterpret_cast<fptype*>(out_c.data_ptr());
1166 + switch (N_in) {
1167 + case 1:
1168 + WVSPLITK(16, 2, 2, 2, 2, 2, 2, 1)
1169 + break;
1170 + case 2:
1171 + WVSPLITK(16, 2, 2, 2, 2, 2, 2, 2)
1172 + break;
1173 + case 3:
1174 + WVSPLITK(16, 4, 7, 7, 1, 1, 1, 3)
1175 + break;
1176 + case 4:
1177 + WVSPLITK(16, 4, 7, 7, 1, 1, 1, 4)
1178 + break;
1179 + default:
1180 + throw std::runtime_error(
1181 + "Unsupported N value: " + std::to_string(M_in) + "," +
1182 + std::to_string(K_in) + "," + std::to_string(N_in));
1183 + }
1184 + });
1185 + return out_c;
1186 +}
1187 +
1188 +#if defined(__HIP__MI300__) // TODO: Add NAVI support
1189 +template <typename scalar_t, typename fp8_t, int THRDS, int YTILE, int WvPrGrp,
1190 + int A_CHUNK, int UNRL, int N>
1191 +__global__ void __launch_bounds__(WvPrGrp* THRDS)
1192 + wvSplitKQ_hf_sml_(const int K, const int Kp, const int M, const fp8_t* B,
1193 + const fp8_t* __restrict__ A, scalar_t* C,
1194 + const float* __restrict__ s_A,
1195 + const float* __restrict__ s_B, const int _WvPrGrp,
1196 + const int CuCount) {
1197 + using scalar8 =
1198 + __attribute__((__vector_size__((A_CHUNK / 4) * sizeof(float)))) float;
1199 + using intx2 = __attribute__((__vector_size__(2 * sizeof(int)))) int;
1200 + using intx4 = __attribute__((__vector_size__(4 * sizeof(int)))) int;
1201 + union bigType {
1202 + char f8[A_CHUNK];
1203 + char2 c2[A_CHUNK / 2];
1204 + scalar_t h[A_CHUNK / 2];
1205 + float f[A_CHUNK / 4];
1206 + int i[A_CHUNK / 4];
1207 + long l[A_CHUNK / 8];
1208 + intx4 l2[A_CHUNK / 16];
1209 + scalar8 h8;
1210 + };
1211 +
1212 + __shared__ fp8_t s[1024 * 64];
1213 +
1214 + for (uint32_t k = (threadIdx.y * THRDS + threadIdx.x) * A_CHUNK;
1215 + k < min(K * N, 64 * 1024); k += THRDS * WvPrGrp * A_CHUNK) {
1216 + *((bigType*)(&s[k])) = *((bigType*)(&A[k]));
1217 + }
1218 + __syncthreads();
1219 +
1220 + if (threadIdx.y >= _WvPrGrp) return;
1221 +
1222 + uint32_t m = (blockIdx.x * _WvPrGrp + (threadIdx.y % _WvPrGrp)) * YTILE;
1223 +
1224 + using floatx16 = __attribute__((__vector_size__(16 * sizeof(float)))) float;
1225 + floatx16 sum[N][YTILE];
1226 + float sA = *s_A;
1227 + float sB = *s_B;
1228 +
1229 + while (m < M) {
1230 + for (int i = 0; i < YTILE; i++)
1231 + for (int n = 0; n < N; n++) sum[n][i] = {0.f};
1232 +
1233 + bigType bigA[N][UNRL];
1234 + bigType bigB[YTILE][UNRL];
1235 +
1236 + for (uint32_t k1 = 0; k1 < K; k1 += THRDS * A_CHUNK * UNRL) {
1237 + #pragma unroll
1238 + for (uint32_t k2 = 0; k2 < UNRL; k2++) {
1239 + #pragma unroll
1240 + for (uint32_t n = 0; n < N; ++n) bigA[n][k2].h8 = {0.f};
1241 + #pragma unroll
1242 + for (uint32_t y = 0; y < YTILE; ++y) bigB[y][k2].h8 = {0.f};
1243 + }
1244 +
1245 + // Fetch the weight matrix from memory!
1246 + #pragma unroll
1247 + for (uint32_t k2 = 0; k2 < UNRL; k2++) {
1248 + uint32_t k = k1 + k2 * THRDS * A_CHUNK;
1249 + uint32_t k_ = k + threadIdx.x * A_CHUNK;
1250 + if (k_ >= K) break;
1251 +
1252 + const fp8_t* B_ = &B[(m + 0) * Kp + k_];
1253 + #pragma unroll
1254 + for (uint32_t y = 0; y < YTILE; ++y) {
1255 + bigB[y][k2].h8 = (loadnt((scalar8*)(&B_[y * Kp])));
1256 + }
1257 + }
1258 +
1259 + // Fetch activation matrix from either just LDS or from both LDS / memory
1260 + #pragma unroll
1261 + for (uint32_t k2 = 0; k2 < UNRL; k2++) {
1262 + uint32_t k = k1 + k2 * THRDS * A_CHUNK;
1263 + uint32_t k_ = k + threadIdx.x * A_CHUNK;
1264 + if (k_ >= K) break;
1265 + for (int n = 0; n < N; n++) {
1266 + bigA[n][k2] = *((const bigType*)(&(s[k_ + K * n])));
1267 + }
1268 + }
1269 +
1270 + // Do the matrix multiplication in interleaved manner
1271 + #pragma unroll
1272 + for (uint32_t k2 = 0; k2 < UNRL; k2++) {
1273 + uint32_t k = k1 + k2 * THRDS * A_CHUNK;
1274 + if (k >= K) break;
1275 +
1276 + for (uint32_t n = 0; n < N; n++) {
1277 + for (int i = 0; i < A_CHUNK; i += 8) {
1278 + for (int y = 0; y < YTILE; ++y) {
1279 + sum[n][y] = __builtin_amdgcn_mfma_f32_32x32x16_fp8_fp8(
1280 + bigA[n][k2].l[i / 8], bigB[y][k2].l[i / 8], sum[n][y], 0, 0,
1281 + 0);
1282 + }
1283 + }
1284 + }
1285 + }
1286 + }
1287 +
1288 + // Final reduction
1289 + for (int n = 0; n < N; n++) {
1290 + for (int y = 0; y < YTILE; y++) {
1291 + float accm0 = sum[n][y][0];
1292 + float accm16 = sum[n][y][8];
1293 + asm("v_add_f32 %0, %2, %3 row_shl:1 bound_ctrl:0 "
1294 + : "=v"(accm0)
1295 + : "0"(accm0), "v"(sum[n][y][1]), "v"(accm0));
1296 + asm("v_add_f32 %0, %2, %3 row_shl:1 bound_ctrl:0 "
1297 + : "=v"(accm16)
1298 + : "0"(accm16), "v"(sum[n][y][9]), "v"(accm16));
1299 + asm("v_add_f32 %0, %2, %3 row_shl:2 bound_ctrl:0 "
1300 + : "=v"(accm0)
1301 + : "0"(accm0), "v"(sum[n][y][2]), "v"(accm0));
1302 + asm("v_add_f32 %0, %2, %3 row_shl:2 bound_ctrl:0 "
1303 + : "=v"(accm16)
1304 + : "0"(accm16), "v"(sum[n][y][10]), "v"(accm16));
1305 + asm("v_add_f32 %0, %2, %3 row_shl:3 bound_ctrl:0 "
1306 + : "=v"(accm0)
1307 + : "0"(accm0), "v"(sum[n][y][3]), "v"(accm0));
1308 + asm("v_add_f32 %0, %2, %3 row_shl:3 bound_ctrl:0 "
1309 + : "=v"(accm16)
1310 + : "0"(accm16), "v"(sum[n][y][11]), "v"(accm16));
1311 + asm("v_add_f32 %0, %2, %3 row_shl:8 bound_ctrl:0 "
1312 + : "=v"(accm0)
1313 + : "0"(accm0), "v"(sum[n][y][4]), "v"(accm0));
1314 + asm("v_add_f32 %0, %2, %3 row_shl:8 bound_ctrl:0 "
1315 + : "=v"(accm16)
1316 + : "0"(accm16), "v"(sum[n][y][12]), "v"(accm16));
1317 + asm("v_add_f32 %0, %2, %3 row_shl:9 bound_ctrl:0 "
1318 + : "=v"(accm0)
1319 + : "0"(accm0), "v"(sum[n][y][5]), "v"(accm0));
1320 + asm("v_add_f32 %0, %2, %3 row_shl:9 bound_ctrl:0 "
1321 + : "=v"(accm16)
1322 + : "0"(accm16), "v"(sum[n][y][13]), "v"(accm16));
1323 + asm("v_add_f32 %0, %2, %3 row_shl:10 bound_ctrl:0 "
1324 + : "=v"(accm0)
1325 + : "0"(accm0), "v"(sum[n][y][6]), "v"(accm0));
1326 + asm("v_add_f32 %0, %2, %3 row_shl:10 bound_ctrl:0 "
1327 + : "=v"(accm16)
1328 + : "0"(accm16), "v"(sum[n][y][14]), "v"(accm16));
1329 + asm("v_add_f32 %0, %2, %3 row_shl:11 bound_ctrl:0 "
1330 + : "=v"(accm0)
1331 + : "0"(accm0), "v"(sum[n][y][7]), "v"(accm0));
1332 + asm("v_add_f32 %0, %2, %3 row_shl:11 bound_ctrl:0 "
1333 + : "=v"(accm16)
1334 + : "0"(accm16), "v"(sum[n][y][15]), "v"(accm16));
1335 + accm0 += __shfl(accm0, 36);
1336 + accm16 += __shfl(accm16, 52);
1337 + sum[n][y][0] = accm0 + __shfl(accm16, 16);
1338 + }
1339 + }
1340 +
1341 + if (threadIdx.x == 0) {
1342 + for (int n = 0; n < N; n++) {
1343 + for (int y = 0; y < YTILE; y++) {
1344 + C[m + y + n * M] = __float2s<scalar_t>(sum[n][y][0] * sA * sB);
1345 + }
1346 + }
1347 + }
1348 +
1349 + m += CuCount * _WvPrGrp * YTILE;
1350 + }
1351 +}
1352 +#else // !defined(__HIP__MI300__) TODO: Add NAVI support
1353 +template <typename scalar_t, typename fp8_t, int THRDS, int YTILE, int WvPrGrp,
1354 + int A_CHUNK, int UNRL, int N>
1355 +__global__ void wvSplitKQ_hf_sml_(const int K, const int Kp, const int M,
1356 + const fp8_t* B, const fp8_t* __restrict__ A,
1357 + scalar_t* C, const float* __restrict__ s_A,
1358 + const float* __restrict__ s_B,
1359 + const int _WvPrGrp, const int CuCount) {
1360 + UNREACHABLE_CODE
1361 +}
1362 +#endif // defined(__HIP__MI300__) TODO: Add NAVI support
1363 +
1364 +#if defined(__HIP__MI300__) // TODO: Add NAVI support
1365 +template <typename scalar_t, typename fp8_t, int THRDS, int YTILE, int WvPrGrp,
1366 + int A_CHUNK, int UNRL, int N>
1367 +__global__ void __launch_bounds__(WvPrGrp* THRDS)
1368 + wvSplitKQ_hf_(const int K, const int Kp, const int M, const fp8_t* B,
1369 + const fp8_t* __restrict__ A, scalar_t* C,
1370 + const float* __restrict__ s_A, const float* __restrict__ s_B,
1371 + const int _WvPrGrp, const int CuCount) {
1372 + using scalar8 =
1373 + __attribute__((__vector_size__((A_CHUNK / 4) * sizeof(float)))) float;
1374 + using intx2 = __attribute__((__vector_size__(2 * sizeof(int)))) int;
1375 + using intx4 = __attribute__((__vector_size__(4 * sizeof(int)))) int;
1376 + union bigType {
1377 + char f8[A_CHUNK];
1378 + char2 c2[A_CHUNK / 2];
1379 + scalar_t h[A_CHUNK / 2];
1380 + float f[A_CHUNK / 4];
1381 + int i[A_CHUNK / 4];
1382 + long l[A_CHUNK / 8];
1383 + intx4 l2[A_CHUNK / 16];
1384 + scalar8 h8;
1385 + };
1386 +
1387 + __shared__ fp8_t s[1024 * 64];
1388 +
1389 + for (uint32_t k = (threadIdx.y * THRDS + threadIdx.x) * A_CHUNK;
1390 + k < min(K * N, 64 * 1024); k += THRDS * WvPrGrp * A_CHUNK) {
1391 + *((bigType*)(&s[k])) = *((bigType*)(&A[k]));
1392 + }
1393 + __syncthreads();
1394 +
1395 + if (threadIdx.y >= _WvPrGrp) return;
1396 +
1397 + uint32_t m = (blockIdx.x * _WvPrGrp + (threadIdx.y % _WvPrGrp)) * YTILE;
1398 +
1399 + using floatx16 = __attribute__((__vector_size__(16 * sizeof(float)))) float;
1400 + floatx16 sum[N][YTILE];
1401 + float sA = *s_A;
1402 + float sB = *s_B;
1403 +
1404 + while (m < M) {
1405 + for (int i = 0; i < YTILE; i++)
1406 + for (int n = 0; n < N; n++) sum[n][i] = {0};
1407 +
1408 + bigType bigA[N][UNRL];
1409 + bigType bigB[YTILE][UNRL];
1410 +
1411 + for (uint32_t k1 = 0; k1 < K; k1 += THRDS * A_CHUNK * UNRL) {
1412 + // Fetch the weight matrix from memory!
1413 + #pragma unroll
1414 + for (uint32_t k2 = 0; k2 < UNRL; k2++) {
1415 + uint32_t k = k1 + k2 * THRDS * A_CHUNK;
1416 + uint32_t k_ = k + threadIdx.x * A_CHUNK;
1417 + if (k_ >= K) break;
1418 +
1419 + const fp8_t* B_ = &B[(m + 0) * Kp + k_];
1420 + for (int y = 0; y < YTILE; ++y) {
1421 + if (y + m >= M) break; // To avoid mem access fault.
1422 + bigB[y][k2].h8 = (loadnt((scalar8*)(&B_[y * Kp])));
1423 + }
1424 + }
1425 +
1426 + // Fetch activation matrix from either just LDS or from both LDS / memory
1427 + #pragma unroll
1428 + for (uint32_t k2 = 0; k2 < UNRL; k2++) {
1429 + uint32_t k = k1 + k2 * THRDS * A_CHUNK;
1430 + uint32_t k_ = k + threadIdx.x * A_CHUNK;
1431 + if (k_ >= K) break;
1432 + for (int n = 0; n < N; n++) {
1433 + if (k_ + K * n < 64 * 1024)
1434 + bigA[n][k2] = *((const bigType*)(&(s[k_ + K * n])));
1435 + else
1436 + bigA[n][k2] = *((const bigType*)(&(A[k_ + K * n])));
1437 + }
1438 + }
1439 +
1440 + // Do the matrix multiplication in interleaved manner
1441 + #pragma unroll
1442 + for (uint32_t k2 = 0; k2 < UNRL; k2++) {
1443 + uint32_t k = k1 + k2 * THRDS * A_CHUNK;
1444 + uint32_t k_ = k + threadIdx.x * A_CHUNK;
1445 + if (k_ >= K) break;
1446 +
1447 + for (uint32_t n = 0; n < N; n++) {
1448 + for (int i = 0; i < A_CHUNK; i += 8) {
1449 + for (int y = 0; y < YTILE; ++y) {
1450 + sum[n][y] = __builtin_amdgcn_mfma_f32_32x32x16_fp8_fp8(
1451 + bigA[n][k2].l[i / 8], bigB[y][k2].l[i / 8], sum[n][y], 0, 0,
1452 + 0);
1453 + }
1454 + }
1455 + }
1456 + }
1457 + }
1458 +
1459 + // Final reduction
1460 + for (int n = 0; n < N; n++) {
1461 + for (int y = 0; y < YTILE; y++) {
1462 + float accm0 = sum[n][y][0];
1463 + float accm16 = sum[n][y][8];
1464 + asm("v_add_f32 %0, %2, %3 row_shl:1 bound_ctrl:0 "
1465 + : "=v"(accm0)
1466 + : "0"(accm0), "v"(sum[n][y][1]), "v"(accm0));
1467 + asm("v_add_f32 %0, %2, %3 row_shl:1 bound_ctrl:0 "
1468 + : "=v"(accm16)
1469 + : "0"(accm16), "v"(sum[n][y][9]), "v"(accm16));
1470 + asm("v_add_f32 %0, %2, %3 row_shl:2 bound_ctrl:0 "
1471 + : "=v"(accm0)
1472 + : "0"(accm0), "v"(sum[n][y][2]), "v"(accm0));
1473 + asm("v_add_f32 %0, %2, %3 row_shl:2 bound_ctrl:0 "
1474 + : "=v"(accm16)
1475 + : "0"(accm16), "v"(sum[n][y][10]), "v"(accm16));
1476 + asm("v_add_f32 %0, %2, %3 row_shl:3 bound_ctrl:0 "
1477 + : "=v"(accm0)
1478 + : "0"(accm0), "v"(sum[n][y][3]), "v"(accm0));
1479 + asm("v_add_f32 %0, %2, %3 row_shl:3 bound_ctrl:0 "
1480 + : "=v"(accm16)
1481 + : "0"(accm16), "v"(sum[n][y][11]), "v"(accm16));
1482 + asm("v_add_f32 %0, %2, %3 row_shl:8 bound_ctrl:0 "
1483 + : "=v"(accm0)
1484 + : "0"(accm0), "v"(sum[n][y][4]), "v"(accm0));
1485 + asm("v_add_f32 %0, %2, %3 row_shl:8 bound_ctrl:0 "
1486 + : "=v"(accm16)
1487 + : "0"(accm16), "v"(sum[n][y][12]), "v"(accm16));
1488 + asm("v_add_f32 %0, %2, %3 row_shl:9 bound_ctrl:0 "
1489 + : "=v"(accm0)
1490 + : "0"(accm0), "v"(sum[n][y][5]), "v"(accm0));
1491 + asm("v_add_f32 %0, %2, %3 row_shl:9 bound_ctrl:0 "
1492 + : "=v"(accm16)
1493 + : "0"(accm16), "v"(sum[n][y][13]), "v"(accm16));
1494 + asm("v_add_f32 %0, %2, %3 row_shl:10 bound_ctrl:0 "
1495 + : "=v"(accm0)
1496 + : "0"(accm0), "v"(sum[n][y][6]), "v"(accm0));
1497 + asm("v_add_f32 %0, %2, %3 row_shl:10 bound_ctrl:0 "
1498 + : "=v"(accm16)
1499 + : "0"(accm16), "v"(sum[n][y][14]), "v"(accm16));
1500 + asm("v_add_f32 %0, %2, %3 row_shl:11 bound_ctrl:0 "
1501 + : "=v"(accm0)
1502 + : "0"(accm0), "v"(sum[n][y][7]), "v"(accm0));
1503 + asm("v_add_f32 %0, %2, %3 row_shl:11 bound_ctrl:0 "
1504 + : "=v"(accm16)
1505 + : "0"(accm16), "v"(sum[n][y][15]), "v"(accm16));
1506 + accm0 += __shfl(accm0, 36);
1507 + accm16 += __shfl(accm16, 52);
1508 + sum[n][y][0] = accm0 + __shfl(accm16, 16);
1509 + }
1510 + }
1511 +
1512 + if (threadIdx.x == 0) {
1513 + for (int n = 0; n < N; n++) {
1514 + for (int y = 0; y < YTILE; y++) {
1515 + if (y + m >= M) break; // To avoid mem access fault.
1516 + C[m + y + n * M] = __float2s<scalar_t>(sum[n][y][0] * sA * sB);
1517 + }
1518 + }
1519 + }
1520 +
1521 + m += CuCount * _WvPrGrp * YTILE;
1522 + }
1523 +}
1524 +#else // !defined(__HIP__MI300__) TODO: Add NAVI support
1525 +template <typename scalar_t, typename fp8_t, int THRDS, int YTILE, int WvPrGrp,
1526 + int A_CHUNK, int UNRL, int N>
1527 +__global__ void wvSplitKQ_hf_(const int K, const int Kp, const int M,
1528 + const fp8_t* B, const fp8_t* __restrict__ A,
1529 + scalar_t* C, const float* __restrict__ s_A,
1530 + const float* __restrict__ s_B, const int _WvPrGrp,
1531 + const int CuCount) {
1532 + UNREACHABLE_CODE
1533 +}
1534 +#endif // defined(__HIP__MI300__) TODO: Add NAVI support
1535 +
1536 +void wvSplitKQ(at::Tensor& in_a, at::Tensor& in_b, at::Tensor& out_c,
1537 + at::Tensor& scale_a, at::Tensor& scale_b,
1538 + const int64_t CuCount) {
1539 + static c10::ScalarType kFp8Type = is_fp8_ocp()
1540 + ? c10::ScalarType::Float8_e4m3fn
1541 + : c10::ScalarType::Float8_e4m3fnuz;
1542 + auto M_in = in_a.size(0);
1543 + auto K_in = in_a.size(1);
1544 + auto N_in = in_b.size(0);
1545 + auto Kp_in = in_a.stride(0);
1546 + TORCH_CHECK(K_in % 16 == 0, "k % 16 == 0");
1547 + TORCH_CHECK(in_a.dtype() == in_b.dtype() && in_a.dtype() == kFp8Type);
1548 + TORCH_CHECK(out_c.dtype() == torch::kFloat16 ||
1549 + out_c.dtype() == torch::kBFloat16);
1550 +
1551 + dim3 grid(CuCount);
1552 + const at::cuda::OptionalCUDAGuard device_guard(device_of(in_a));
1553 + const cudaStream_t stream = at::cuda::getCurrentCUDAStream();
1554 +
1555 +#define WVSPLITKQ(_WvPrGrp, _YTILEs, _YTILEm, _YTILEb, _UNRLs, _UNRLm, _UNRLb, \
1556 + _N) \
1557 + { \
1558 + dim3 block(64, _WvPrGrp); \
1559 + if ((K_in * N_in <= 64 * 1024) && (M_in % _YTILEs == 0)) { \
1560 + int __wvPrGrp = mindiv(M_in, CuCount * _YTILEs, _WvPrGrp); \
1561 + wvSplitKQ_hf_sml_<fptype, fp8_t, 64, _YTILEs, _WvPrGrp, 16, _UNRLs, _N> \
1562 + <<<grid, block, 0, stream>>>(K_in, Kp_in, M_in, a_ptr, b_ptr, c_ptr, \
1563 + s_a, s_b, __wvPrGrp, CuCount); \
1564 + } else { \
1565 + int __wvPrGrp = mindiv(M_in, CuCount * _YTILEm, _WvPrGrp); \
1566 + wvSplitKQ_hf_<fptype, fp8_t, 64, _YTILEm, _WvPrGrp, 16, _UNRLm, _N> \
1567 + <<<grid, block, 0, stream>>>(K_in, Kp_in, M_in, a_ptr, b_ptr, c_ptr, \
1568 + s_a, s_b, __wvPrGrp, CuCount); \
1569 + } \
1570 + }
1571 +
1572 + AT_DISPATCH_REDUCED_FLOATING_TYPES(out_c.scalar_type(), "wvSplitKQ", [&] {
1573 + using fptype = typename scalar<scalar_t>::type;
1574 + auto c_ptr = reinterpret_cast<fptype*>(out_c.data_ptr());
1575 + auto s_a = scale_a.data_ptr<float>();
1576 + auto s_b = scale_b.data_ptr<float>();
1577 + VLLM_DISPATCH_FP8_TYPES(in_a.scalar_type(), "wvSplitKQ", [&] {
1578 + auto a_ptr = in_a.data_ptr<fp8_t>();
1579 + auto b_ptr = in_b.data_ptr<fp8_t>();
1580 + switch (N_in) {
1581 + case 1:
1582 + WVSPLITKQ(16, 2, 2, 2, 2, 2, 2, 1)
1583 + break;
1584 + case 2:
1585 + WVSPLITKQ(16, 2, 2, 2, 2, 2, 2, 2)
1586 + break;
1587 + case 3:
1588 + WVSPLITKQ(16, 4, 7, 7, 1, 1, 1, 3)
1589 + break;
1590 + case 4:
1591 + WVSPLITKQ(16, 4, 7, 7, 1, 1, 1, 4)
1592 + break;
1593 + default:
1594 + throw std::runtime_error(
1595 + "Unsupported N value: " + std::to_string(M_in) + "," +
1596 + std::to_string(K_in) + "," + std::to_string(N_in));
1597 + }
1598 + });
1599 + });
1600 +}
\ No newline at end of file
csrc/rocm/torch_bindings.cpp
+18
@@ -14,6 +14,24 @@
14 TORCH_LIBRARY_EXPAND(TORCH_EXTENSION_NAME, rocm_ops) {
15 // vLLM custom ops for rocm
16
17 + // Custom gemm op for matrix-vector multiplication
18 + rocm_ops.def(
19 + "LLMM1(Tensor in_a, Tensor in_b, int rows_per_block) -> "
20 + "Tensor");
21 + rocm_ops.impl("LLMM1", torch::kCUDA, &LLMM1);
22 +
23 + // Custom gemm op for skinny matrix-matrix multiplication
24 + rocm_ops.def(
25 + "wvSplitK(Tensor in_a, Tensor in_b, int CuCount) -> "
26 + "Tensor");
27 + rocm_ops.impl("wvSplitK", torch::kCUDA, &wvSplitK);
28 +
29 + // wvSplitK for fp8
30 + rocm_ops.def(
31 + "wvSplitKQ(Tensor in_a, Tensor in_b, Tensor! out_c, Tensor scale_a, "
32 + " Tensor scale_b, int CuCount) -> ()");
33 + rocm_ops.impl("wvSplitKQ", torch::kCUDA, &wvSplitKQ);
34 +
35 // Custom attention op
36 // Compute the attention between an input query and the cached
37 // keys/values using PagedAttention.
tests/kernels/test_rocm_skinny_gemms.py new
+80
@@ -0,0 +1,80 @@
1 +# SPDX-License-Identifier: Apache-2.0
2 +import pytest
3 +import torch
4 +
5 +import vllm._custom_ops as ops
6 +from tests.kernels.quant_utils import ref_dynamic_per_tensor_fp8_quant
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
12 +N = [1, 2, 3, 4]
13 +SEEDS = [0]
14 +
15 +
16 +@pytest.mark.parametrize("n", [1]) # only test for batch size 1
17 +@pytest.mark.parametrize("k", K)
18 +@pytest.mark.parametrize("m", M)
19 +@pytest.mark.parametrize("dtype", DTYPES)
20 +@pytest.mark.parametrize("rows_per_block", [2, 4, 8, 16])
21 +@pytest.mark.parametrize("seed", SEEDS)
22 +@pytest.mark.skipif(not current_platform.is_rocm(),
23 + reason="only test for rocm")
24 +@torch.inference_mode()
25 +def test_rocm_llmm1_kernel(n, k, m, dtype, rows_per_block, seed):
26 + torch.manual_seed(seed)
27 + A = torch.rand(n, k, dtype=dtype, device="cuda")
28 + B = torch.rand(m, k, dtype=dtype, device="cuda")
29 +
30 + ref_out = torch.matmul(A, B.t())
31 + out = ops.LLMM1(B, A, rows_per_block)
32 +
33 + assert torch.allclose(out, ref_out, rtol=0.01)
34 +
35 +
36 +@pytest.mark.parametrize("n", N) # only test for batch size <= 4
37 +@pytest.mark.parametrize("k", K + [9216, 10240, 16384])
38 +@pytest.mark.parametrize("m", [8] + M) # m >= 8
39 +@pytest.mark.parametrize("dtype", DTYPES)
40 +@pytest.mark.parametrize("seed", SEEDS)
41 +@pytest.mark.skipif(not current_platform.is_rocm(),
42 + reason="only test for rocm")
43 +def test_rocm_wvsplitk_kernel(n, k, m, dtype, seed):
44 + torch.manual_seed(seed)
45 + cu_count = current_platform.get_cu_count()
46 +
47 + A = torch.rand(n, k, dtype=dtype, device="cuda")
48 + B = torch.rand(m, k, dtype=dtype, device="cuda")
49 +
50 + ref_out = torch.matmul(A, B.t())
51 + out = ops.wvSplitK(B, A, cu_count)
52 +
53 + assert torch.allclose(out, ref_out, rtol=0.01)
54 +
55 +
56 +@pytest.mark.parametrize("n", N) # only test for batch size <= 4
57 +@pytest.mark.parametrize("k", K[1:] + [14336, 24576, 32768]) # k % 16 == 0
58 +@pytest.mark.parametrize("m", M + [28672]) # m >= 16
59 +@pytest.mark.parametrize("dtype", DTYPES)
60 +@pytest.mark.parametrize("seed", SEEDS)
61 +@pytest.mark.skipif(not current_platform.is_rocm(),
62 + reason="only test for rocm")
63 +def test_rocm_wvsplitk_fp8_kernel(n, k, m, dtype, seed):
64 + torch.manual_seed(seed)
65 +
66 + A = torch.rand(n, k, device="cuda")
67 + B = torch.rand(m, k, device="cuda")
68 +
69 + A, scale_a = ref_dynamic_per_tensor_fp8_quant(A)
70 + B, scale_b = ref_dynamic_per_tensor_fp8_quant(B)
71 +
72 + ref_out = torch._scaled_mm(A,
73 + B.t(),
74 + out_dtype=dtype,
75 + scale_a=scale_a,
76 + scale_b=scale_b)
77 + out = ops.wvSplitKQ(B, A, dtype, scale_a, scale_b,
78 + current_platform.get_cu_count())
79 +
80 + assert torch.allclose(out, ref_out, rtol=0.01)
vllm/_custom_ops.py
+20
@@ -1196,6 +1196,26 @@ def selective_scan_fwd(u: torch.Tensor, delta: torch.Tensor, A: torch.Tensor,
1196 ssm_states, pad_slot_id)
1197
1198
1199 +# ROCm skinny gemms
1200 +def LLMM1(a: torch.Tensor, b: torch.Tensor,
1201 + rows_per_block: int) -> torch.Tensor:
1202 + return torch.ops._rocm_C.LLMM1(a, b, rows_per_block)
1203 +
1204 +
1205 +def wvSplitK(a: torch.Tensor, b: torch.Tensor, cu_count: int) -> torch.Tensor:
1206 + return torch.ops._rocm_C.wvSplitK(a, b, cu_count)
1207 +
1208 +
1209 +def wvSplitKQ(a: torch.Tensor, b: torch.Tensor, out_dtype: torch.dtype,
1210 + scale_a: torch.Tensor, scale_b: torch.Tensor,
1211 + cu_count: int) -> torch.Tensor:
1212 + out = torch.empty((b.shape[0], a.shape[0]),
1213 + dtype=out_dtype,
1214 + device=b.device)
1215 + torch.ops._rocm_C.wvSplitKQ(a, b, out, scale_a, scale_b, cu_count)
1216 + return out
1217 +
1218 +
1219 # moe
1220 def moe_sum(input: torch.Tensor, output: torch.Tensor):
1221 torch.ops._moe_C.moe_sum(input, output)
vllm/envs.py
+6
@@ -78,6 +78,7 @@ if TYPE_CHECKING:
78 VLLM_ROCM_USE_AITER_LINEAR: bool = True
79 VLLM_ROCM_USE_AITER_MOE: bool = True
80 VLLM_ROCM_USE_AITER_RMSNORM: bool = True
81 + VLLM_ROCM_USE_SKINNY_GEMM: bool = True
82 VLLM_ROCM_FP8_PADDING: bool = True
83 VLLM_ROCM_MOE_PADDING: bool = True
84 VLLM_ROCM_CUSTOM_PAGED_ATTN: bool = True
@@ -550,6 +551,11 @@ environment_variables: dict[str, Callable[[], Any]] = {
551 lambda: (os.getenv("VLLM_ROCM_USE_AITER_RMSNORM", "True").lower() in
552 ("true", "1")),
553
554 + # use rocm skinny gemms
555 + "VLLM_ROCM_USE_SKINNY_GEMM":
556 + lambda: (os.getenv("VLLM_ROCM_USE_SKINNY_GEMM", "True").lower() in
557 + ("true", "1")),
558 +
559 # Pad the fp8 weights to 256 bytes for ROCm
560 "VLLM_ROCM_FP8_PADDING":
561 lambda: bool(int(os.getenv("VLLM_ROCM_FP8_PADDING", "1"))),
vllm/model_executor/layers/linear.py
+2 -2
@@ -6,7 +6,6 @@ from typing import Any, Literal, Optional, Union
6
7 import torch
8 import torch.nn as nn
9 -import torch.nn.functional as F
9 from torch.nn.parameter import Parameter, UninitializedParameter
10
11 from vllm.distributed import (divide, get_tensor_model_parallel_rank,
@@ -17,6 +16,7 @@ from vllm.distributed import (divide, get_tensor_model_parallel_rank,
16 from vllm.logger import init_logger
17 from vllm.model_executor.layers.quantization.base_config import (
18 QuantizationConfig, QuantizeMethodBase)
19 +from vllm.model_executor.layers.utils import dispatch_unquantized_gemm
20 # yapf: disable
21 from vllm.model_executor.parameter import (BasevLLMParameter,
22 BlockQuantScaleParameter,
@@ -188,7 +188,7 @@ class UnquantizedLinearMethod(LinearMethodBase):
188 x: torch.Tensor,
189 bias: Optional[torch.Tensor] = None) -> torch.Tensor:
190
191 - return F.linear(x, layer.weight, bias)
191 + return dispatch_unquantized_gemm()(x, layer.weight, bias)
192
193
194 class LinearBase(torch.nn.Module):
vllm/model_executor/layers/quantization/utils/w8a8_utils.py
+173 -92
@@ -1,10 +1,11 @@
1 # SPDX-License-Identifier: Apache-2.0
2
3 -from typing import List, Optional, Tuple, Union
3 +from typing import Callable, List, Optional, Tuple, Union
4
5 import torch
6
7 from vllm import _custom_ops as ops
8 +from vllm import envs
9 from vllm.config import CompilationLevel, get_current_vllm_config
10 from vllm.platforms import current_platform
11
@@ -17,6 +18,7 @@ TORCH_DEVICE_IDENTITY = None
18 # The condition is determined once as the operations
19 # are time consuming.
20 USE_ROWWISE_TORCH_SCALED_MM = (current_platform.is_rocm()
21 + and torch.__version__[0:3] >= "2.7"
22 and current_platform.has_device_capability(94))
23
24
@@ -131,6 +133,159 @@ def maybe_create_device_identity():
133 TORCH_DEVICE_IDENTITY = torch.ones(1, dtype=torch.float32)
134
135
136 +def cutlass_w8a8_scaled_mm(*, qinput: torch.Tensor, weight: torch.Tensor,
137 + out_dtype: torch.dtype, scale_a: torch.Tensor,
138 + scale_b: torch.Tensor, bias: torch.Tensor,
139 + output_shape: List, **kwargs) -> torch.Tensor:
140 +
141 + # Fused GEMM_DQ
142 + output = ops.cutlass_scaled_mm(qinput,
143 + weight,
144 + out_dtype=out_dtype,
145 + scale_a=scale_a,
146 + scale_b=scale_b,
147 + bias=bias)
148 + return output.view(*output_shape)
149 +
150 +
151 +def rocm_per_tensor_w8a8_scaled_mm(*, qinput: torch.Tensor,
152 + weight: torch.Tensor,
153 + out_dtype: torch.dtype,
154 + scale_a: torch.Tensor,
155 + scale_b: torch.Tensor, bias: torch.Tensor,
156 + input_2d: torch.Tensor,
157 + output_shape: List) -> torch.Tensor:
158 + if envs.VLLM_ROCM_USE_SKINNY_GEMM and qinput.shape[
159 + 0] == 1 and qinput.shape[1] % 16 == 0:
160 + output = ops.wvSplitKQ(weight.t(), qinput, out_dtype, scale_a, scale_b,
161 + current_platform.get_cu_count())
162 + else:
163 + output = torch._scaled_mm(qinput,
164 + weight,
165 + out_dtype=out_dtype,
166 + scale_a=scale_a,
167 + scale_b=scale_b,
168 + bias=bias)
169 +
170 + return torch.narrow(output, 0, 0, input_2d.shape[0]).view(*output_shape)
171 +
172 +
173 +def torch_per_tensor_w8a8_scaled_mm(*, qinput: torch.Tensor,
174 + weight: torch.Tensor,
175 + out_dtype: torch.dtype,
176 + scale_a: torch.Tensor,
177 + scale_b: torch.Tensor, bias: torch.Tensor,
178 + input_2d: torch.Tensor,
179 + output_shape: List) -> torch.Tensor:
180 + output = torch._scaled_mm(qinput,
181 + weight,
182 + out_dtype=out_dtype,
183 + scale_a=scale_a,
184 + scale_b=scale_b,
185 + bias=bias)
186 + # A fix for discrepancy in scaled_mm which returns tuple
187 + # for torch < 2.5 and a single value in torch >= 2.5
188 + if type(output) is tuple and len(output) == 2:
189 + output = output[0]
190 +
191 + return torch.narrow(output, 0, 0, input_2d.shape[0]).view(*output_shape)
192 +
193 +
194 +def torch_per_token_w8a8_scaled_mm(*, qinput: torch.Tensor,
195 + weight: torch.Tensor,
196 + out_dtype: torch.dtype,
197 + scale_a: torch.Tensor,
198 + scale_b: torch.Tensor, bias: torch.Tensor,
199 + input_2d: torch.Tensor,
200 + output_shape: List) -> torch.Tensor:
201 + # Note: Callers of this function should check USE_ROWWISE_TORCH_SCALED_MM
202 + # when using it.
203 + # For now it has only been validated on ROCm platform.
204 + # fp8 rowwise scaling in torch._scaled_mm is introduced in
205 + # https://github.com/pytorch/pytorch/pull/144432 using
206 + # hipBLASLt and ROCm 6.3, which only exists in torch 2.7 and above.
207 + #
208 + # For CUDA platform please validate if the torch._scaled_mm supports
209 + # rowwise scaled GEMM before using it
210 +
211 + # Fused GEMM_DQ Rowwise GEMM
212 + output = torch._scaled_mm(qinput,
213 + weight,
214 + out_dtype=out_dtype,
215 + scale_a=scale_a,
216 + scale_b=scale_b.t(),
217 + bias=bias)
218 +
219 + output = torch.narrow(output, 0, 0, input_2d.shape[0])
220 + output = output.view(*output_shape)
221 + return output
222 +
223 +
224 +def torch_channelwise_w8a8_scaled_mm(*, qinput: torch.Tensor,
225 + weight: torch.Tensor,
226 + out_dtype: torch.dtype,
227 + scale_a: torch.Tensor,
228 + scale_b: torch.Tensor, bias: torch.Tensor,
229 + input_2d: torch.Tensor,
230 + output_shape: List,
231 + **kwargs) -> torch.Tensor:
232 + # Use unfused DQ due to limitations with scaled_mm
233 +
234 + # Symmetric quantized GEMM by definition computes the following:
235 + # C = (s_x * X) (s_w * W) + bias
236 + # This is equivalent to dequantizing the weights and activations
237 + # before applying a GEMM.
238 + #
239 + # In order to compute quantized operands, a quantized kernel
240 + # will rewrite the above like so:
241 + # C = s_w * s_x * (X * W) + bias
242 + #
243 + # For the scaled_mm fallback case, we break this down, since it
244 + # does not support s_w being a vector.
245 +
246 + # GEMM
247 + # This computes C = (X * W).
248 + # Output in fp32 to allow subsequent ops to happen in-place
249 + output = torch._scaled_mm(qinput,
250 + weight,
251 + scale_a=TORCH_DEVICE_IDENTITY,
252 + scale_b=TORCH_DEVICE_IDENTITY,
253 + out_dtype=torch.float32)
254 + # A fix for discrepancy in scaled_mm which returns tuple
255 + # for torch < 2.5 and a single value in torch >= 2.5
256 + if type(output) is tuple and len(output) == 2:
257 + output = output[0]
258 + # Unpad (undo num_token_padding)
259 + output = torch.narrow(output, 0, 0, input_2d.shape[0])
260 + x_scale = torch.narrow(scale_a, 0, 0, input_2d.shape[0])
261 +
262 + # DQ
263 + # C = sw * sx * (X * W) + bias
264 + output = output * x_scale * scale_b.t()
265 + if bias is not None:
266 + output = output + bias
267 + return output.to(out_dtype).view(*output_shape)
268 +
269 +
270 +def dispatch_w8a8_scaled_mm(
271 + cutlass_fp8_supported: bool, per_tensor_weights: bool,
272 + per_tensor_activations: bool, use_per_token_if_dynamic: Optional[bool]
273 +) -> Callable[..., torch.Tensor]:
274 +
275 + if cutlass_fp8_supported:
276 + return cutlass_w8a8_scaled_mm
277 + if per_tensor_weights and per_tensor_activations:
278 + if current_platform.is_rocm():
279 + return rocm_per_tensor_w8a8_scaled_mm
280 + return torch_per_tensor_w8a8_scaled_mm
281 + # torch.scaled_mm supports per tensor weights + activations only
282 + # so fallback to naive if per channel or per token
283 + if (use_per_token_if_dynamic and not per_tensor_weights
284 + and not per_tensor_activations and USE_ROWWISE_TORCH_SCALED_MM):
285 + return torch_per_token_w8a8_scaled_mm
286 + return torch_channelwise_w8a8_scaled_mm
287 +
288 +
289 # TODO(luka): follow similar pattern for marlin and block-fp8-linear
290 # https://github.com/vllm-project/vllm/issues/14397
291 class Fp8LinearOp:
@@ -156,7 +311,8 @@ class Fp8LinearOp:
311 if pad_output is None:
312 config = get_current_vllm_config().compilation_config
313 pad_output = config.level < CompilationLevel.PIECEWISE
159 - self.output_padding = 17 if pad_output else None
314 + self.output_padding = 17 if (
315 + pad_output and not current_platform.is_rocm()) else None
316
317 def apply(
318 self,
@@ -195,18 +351,6 @@ class Fp8LinearOp:
351 input_scale,
352 scale_ub=input_scale_ub,
353 use_per_token_if_dynamic=use_per_token_if_dynamic)
198 -
199 - # Fused GEMM_DQ
200 - output = ops.cutlass_scaled_mm(qinput,
201 - weight,
202 - out_dtype=out_dtype,
203 - scale_a=x_scale,
204 - scale_b=weight_scale,
205 - bias=bias)
206 - return output.view(*output_shape)
207 -
208 - # torch.scaled_mm supports per tensor weights + activations only
209 - # so fallback to naive if per channel or per token
354 else:
355 if input.dtype != current_platform.fp8_dtype():
356 # Maybe apply padding to output, see comment in __init__
@@ -218,84 +362,21 @@ class Fp8LinearOp:
362 else:
363 qinput, x_scale = input_2d, input_scale
364
221 - per_tensor_weights = (weight_scale.numel() == 1)
222 - per_tensor_activations = (x_scale.numel() == 1)
223 -
224 - if per_tensor_weights and per_tensor_activations:
225 - # Fused GEMM_DQ
226 - output = torch._scaled_mm(qinput,
227 - weight,
228 - out_dtype=out_dtype,
229 - scale_a=x_scale,
230 - scale_b=weight_scale,
231 - bias=bias)
232 - # A fix for discrepancy in scaled_mm which returns tuple
233 - # for torch < 2.5 and a single value in torch >= 2.5
234 - if type(output) is tuple and len(output) == 2:
235 - output = output[0]
236 -
237 - return torch.narrow(output, 0, 0,
238 - input_2d.shape[0]).view(*output_shape)
239 -
240 - elif (use_per_token_if_dynamic and not per_tensor_weights
241 - and not per_tensor_activations
242 - and USE_ROWWISE_TORCH_SCALED_MM):
243 - # For now validated on ROCm platform
244 - # fp8 rowwise scaling in torch._scaled_mm is introduced in
245 - # https://github.com/pytorch/pytorch/pull/144432 using hipBLASLt
246 - # and ROCm 6.3, which only exists in torch 2.7 and above.
247 - # For CUDA platform please validate if the
248 - # torch._scaled_mm support rowwise scaled GEMM
249 - # Fused GEMM_DQ Rowwise GEMM
250 - output = torch._scaled_mm(qinput,
251 - weight,
252 - out_dtype=out_dtype,
253 - scale_a=x_scale,
254 - scale_b=weight_scale.t(),
255 - bias=bias)
256 -
257 - output = torch.narrow(output, 0, 0, input_2d.shape[0])
258 - output = output.view(*output_shape)
259 - return output
260 -
261 - else:
262 - # Fallback for channelwise case, where we use unfused DQ
263 - # due to limitations with scaled_mm
264 -
265 - # Symmetric quantized GEMM by definition computes the following:
266 - # C = (s_x * X) (s_w * W) + bias
267 - # This is equivalent to dequantizing the weights and activations
268 - # before applying a GEMM.
269 - #
270 - # In order to compute quantized operands, a quantized kernel
271 - # will rewrite the above like so:
272 - # C = s_w * s_x * (X * W) + bias
273 - #
274 - # For the scaled_mm fallback case, we break this down, since it
275 - # does not support s_w being a vector.
276 -
277 - # GEMM
278 - # This computes C = (X * W).
279 - # Output in fp32 to allow subsequent ops to happen in-place
280 - output = torch._scaled_mm(qinput,
281 - weight,
282 - scale_a=TORCH_DEVICE_IDENTITY,
283 - scale_b=TORCH_DEVICE_IDENTITY,
284 - out_dtype=torch.float32)
285 - # A fix for discrepancy in scaled_mm which returns tuple
286 - # for torch < 2.5 and a single value in torch >= 2.5
287 - if type(output) is tuple and len(output) == 2:
288 - output = output[0]
289 - # Unpad (undo num_token_padding)
290 - output = torch.narrow(output, 0, 0, input_2d.shape[0])
291 - x_scale = torch.narrow(x_scale, 0, 0, input_2d.shape[0])
292 -
293 - # DQ
294 - # C = sw * sx * (X * W) + bias
295 - output = output * x_scale * weight_scale.t()
296 - if bias is not None:
297 - output = output + bias
298 - return output.to(dtype=input.dtype).view(*output_shape)
365 + per_tensor_weights = (weight_scale.numel() == 1)
366 + per_tensor_activations = (x_scale.numel() == 1)
367 +
368 + w8a8_scaled_mm_func = dispatch_w8a8_scaled_mm(
369 + self.cutlass_fp8_supported, per_tensor_weights,
370 + per_tensor_activations, use_per_token_if_dynamic)
371 +
372 + return w8a8_scaled_mm_func(qinput=qinput,
373 + weight=weight,
374 + out_dtype=input.dtype,
375 + scale_a=x_scale,
376 + scale_b=weight_scale,
377 + bias=bias,
378 + input_2d=input_2d,
379 + output_shape=output_shape)
380
381
382 def normalize_e4m3fn_to_e4m3fnuz(
vllm/model_executor/layers/utils.py
+36 -1
@@ -1,9 +1,13 @@
1 # SPDX-License-Identifier: Apache-2.0
2 """Utility methods for model layers."""
3 -from typing import Tuple
3 +from typing import Callable, Optional, Tuple
4
5 import torch
6
7 +from vllm import _custom_ops as ops
8 +from vllm import envs
9 +from vllm.platforms import current_platform
10 +
11
12 def get_token_bin_counts_and_mask(
13 tokens: torch.Tensor,
@@ -61,3 +65,34 @@ def apply_penalties(logits: torch.Tensor, prompt_tokens_tensor: torch.Tensor,
65 logits -= frequency_penalties.unsqueeze(dim=1) * output_bin_counts
66 logits -= presence_penalties.unsqueeze(dim=1) * output_mask
67 return logits
68 +
69 +
70 +def rocm_unquantized_gemm(x: torch.Tensor,
71 + weight: torch.Tensor,
72 + bias: Optional[torch.Tensor] = None):
73 + k = weight.shape[1]
74 + use_skinny = (envs.VLLM_ROCM_USE_SKINNY_GEMM and \
75 + x.dtype in [torch.float16, torch.bfloat16] \
76 + and k % 8 == 0 and bias is None)
77 +
78 + if use_skinny is not True:
79 + return torch.nn.functional.linear(x, weight, bias)
80 +
81 + x_view = x.view(-1, x.size(-1))
82 + n = x_view.shape[0]
83 + m = weight.shape[0]
84 + cu_count = current_platform.get_cu_count()
85 +
86 + if m > 8 and n < 4:
87 + out = ops.wvSplitK(weight, x_view, cu_count)
88 + return out.view(*x.shape[:-1], weight.shape[0])
89 + elif m % 4 == 0 and n == 1 and k <= 8192:
90 + out = ops.LLMM1(weight, x_view, out, 4)
91 + return out.view(*x.shape[:-1], weight.shape[0])
92 + return torch.nn.functional.linear(x, weight, bias)
93 +
94 +
95 +def dispatch_unquantized_gemm() -> Callable[..., torch.Tensor]:
96 + if current_platform.is_rocm():
97 + return rocm_unquantized_gemm
98 + return torch.nn.functional.linear
vllm/platforms/interface.py
+7
@@ -413,6 +413,13 @@ class Platform:
413 self.device_name, key)
414 return None
415
416 + @classmethod
417 + def get_cu_count(cls, device_id: int = 0) -> int:
418 + """
419 + Returns the total number of compute units (CU) on single GPU.
420 + """
421 + raise NotImplementedError
422 +
423
424 class UnspecifiedPlatform(Platform):
425 _enum = PlatformEnum.UNSPECIFIED
vllm/platforms/rocm.py
+5
@@ -312,3 +312,8 @@ class RocmPlatform(Platform):
312 gcn_arch = torch.cuda.get_device_properties(0).gcnArchName
313 supported_archs = ['gfx94']
314 return any(gfx in gcn_arch for gfx in supported_archs)
315 +
316 + @classmethod
317 + def get_cu_count(cls, device_id: int = 0) -> int:
318 + return torch.cuda.get_device_properties(
319 + device_id).multi_processor_count
\ No newline at end of file