skinny_gemms.cu · tduncan/Nvidia-Optimization-Stack-VLLM

@tduncan / Nvidia-Optimization-Stack-VLLM

Code Commits Issues Pulls Cloud Agent CI/CD

main

cu 2,544 lines 97.6 KB

Raw

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	#include <map>
12	#include <mutex>
13	#include <utility>
14
15	#include "../cuda_compat.h"
16	#include "dispatch_utils.h"
17	#include "quantization/w8a8/fp8/common.cuh"
18	#include "core/batch_invariant.hpp"
19
20	// Number of streams the pre-allocated split-K pool covers. ~7.5 MiB per slot.
21	static constexpr int64_t kWvSlots = 8;
22	// Slots held back for capturing streams, which cannot allocate one themselves.
23	static constexpr int64_t kWvCaptureSlots = 8;
24
25	// TODO(rasmith): The kernels in this file are susceptible to integer overflow
26	// issues, do not take strides, and are unable to handle PyTorch tensors that
27	// return is_contiguous() as False (the tensors may actually be contiguous
28	// in memory).
29	//
30	// However, it may be possible to fix these kernels to handle both issues.
31
32	#if defined(__HIPCC__) && \
33	(defined(__gfx90a__) \|\| defined(__gfx942__) \|\| defined(__gfx950__))
34	#define __HIP__GFX9__
35	#endif
36
37	// Combined RDNA macro (gfx11 + gfx12) - both use 32-wide wavefronts
38	#if defined(__GFX11__) \|\| defined(__GFX12__)
39	#define __HIP__GFX1X__
40	#endif
41
42	#if defined(__HIPCC__) && (defined(__gfx942__) \|\| defined(__gfx950__))
43	#define __HIP__MI3XX__
44	#endif
45
46	#if defined(__gfx950__)
47	#define LDS_SIZE 160 * 1024
48	#else
49	#define LDS_SIZE 64 * 1024
50	#endif
51
52	int get_lds_size() {
53	static const int result = [] {
54	const auto* dprops = at::cuda::getCurrentDeviceProperties();
55	const std::string device_arch = dprops->gcnArchName;
56	return device_arch.find("gfx95") == std::string::npos ? 64 * 1024
57	: 160 * 1024;
58	}();
59	return result;
60	}
61
62	bool on_gfx1x() {
63	static const bool result = [] {
64	const auto* dprops = at::cuda::getCurrentDeviceProperties();
65	const std::string device_arch = dprops->gcnArchName;
66	return device_arch.find("gfx11") != std::string::npos \|\|
67	device_arch.find("gfx12") != std::string::npos;
68	}();
69	return result;
70	}
71
72	bool on_gfx12() {
73	static const bool result = [] {
74	const auto* dprops = at::cuda::getCurrentDeviceProperties();
75	const std::string device_arch = dprops->gcnArchName;
76	return device_arch.find("gfx12") != std::string::npos;
77	}();
78	return result;
79	}
80
81	bool on_gfx1151() {
82	static const bool result = [] {
83	const auto* dprops = at::cuda::getCurrentDeviceProperties();
84	const std::string device_arch = dprops->gcnArchName;
85	return device_arch.find("gfx1151") != std::string::npos;
86	}();
87	return result;
88	}
89
90	#if defined(NDEBUG)
91	#undef NDEBUG
92	#include <assert.h>
93	#define UNREACHABLE_CODE assert(false);
94	#define NDEBUG
95	#else
96	#define UNREACHABLE_CODE assert(false);
97	#endif
98
99	template <typename T>
100	struct scalar {};
101
102	template <typename T>
103	struct scalar2 {};
104
105	template <typename T>
106	__device__ __forceinline__ float2 __s22float2(T v);
107
108	template <typename T>
109	__device__ __forceinline__ T __float2s(float v);
110
111	template <typename T>
112	__device__ __forceinline__ T __float22s2_rn(float2 v);
113
114	// Definitions and cvt functions for fp16
115	template <>
116	struct scalar<c10::Half> {
117	using type = half;
118	};
119
120	template <>
121	struct scalar2<c10::Half> {
122	using type = __half2;
123	};
124
125	template <>
126	__device__ __forceinline__ half __float2s(float v) {
127	return __float2half(v);
128	}
129
130	template <>
131	__device__ __forceinline__ float2 __s22float2(__half2 v) {
132	return __half22float2(v);
133	}
134
135	template <>
136	__device__ __forceinline__ __half2 __float22s2_rn(float2 v) {
137	return __float22half2_rn(v);
138	}
139
140	// Definitions and cvt functions for bf16
141	template <>
142	struct scalar<c10::BFloat16> {
143	using type = __hip_bfloat16;
144	};
145
146	template <>
147	struct scalar2<c10::BFloat16> {
148	using type = __hip_bfloat162;
149	};
150
151	template <>
152	__device__ __forceinline__ __hip_bfloat16 __float2s(float v) {
153	return __float2bfloat16(v);
154	}
155
156	template <>
157	__device__ __forceinline__ float2 __s22float2(__hip_bfloat162 v) {
158	return __bfloat1622float2(v);
159	}
160
161	template <>
162	__device__ __forceinline__ __hip_bfloat162 __float22s2_rn(float2 v) {
163	return __float22bfloat162_rn(v);
164	}
165
166	template <typename T>
167	__device__ __forceinline__ T loadnt(T* addr) {
168	return __builtin_nontemporal_load(addr);
169	}
170
171	__device__ __forceinline__ float4 load_ntmprl(const float4* addr) {
172	auto addr_alias = reinterpret_cast<const float*>(addr);
173	auto dat0 = loadnt(addr_alias);
174	auto dat1 = loadnt(addr_alias + 1);
175	auto dat2 = loadnt(addr_alias + 2);
176	auto dat3 = loadnt(addr_alias + 3);
177	return make_float4(dat0, dat1, dat2, dat3);
178	}
179
180	// TBlock fetches entire rows of A, and entire col of B (K dimension); assume
181	// N=1 for time being grid is M/A_NUM_ROWS blocks
182	template <typename scalar_t, int NUM_A_ROWS_PER_BLOCK>
183	__global__ void LLGemm1_kernel(const scalar_t* in_a, const scalar_t* in_b,
184	scalar_t* out_c, const int K) {
185	using scalar2_t = typename scalar2<scalar_t>::type;
186	auto af4 = reinterpret_cast<const float4*>(in_a);
187	auto bf4 = reinterpret_cast<const scalar2_t*>(in_b);
188	auto c = reinterpret_cast<scalar2_t*>(out_c);
189	__shared__ float red_smem[NUM_A_ROWS_PER_BLOCK][WARP_SIZE];
190	const int row_addr = blockIdx.x * NUM_A_ROWS_PER_BLOCK * K / 8;
191	const int threadid = threadIdx.x;
192	const int warp = threadIdx.x / WARP_SIZE;
193	const int lane = threadIdx.x % WARP_SIZE;
194	const int num_warps = blockDim.x / WARP_SIZE;
195	const int qwarpid = threadid / 16;
196	const int qthreadid = threadid % 16;
197	float4 rowA_elem4[NUM_A_ROWS_PER_BLOCK];
198	scalar2_t colB_elem4x, colB_elem4y, colB_elem4z, colB_elem4w;
199	float acc[NUM_A_ROWS_PER_BLOCK];
200	scalar2_t acch2;
201	scalar2_t oval;
202
203	// As we later use warp shuffle operations, we may have more threads in the
204	// block than the actual available data, hence the if guard here.
205	if (threadid * 8 < K) {
206	#pragma unroll
207	for (int i = 0; i < NUM_A_ROWS_PER_BLOCK; i++) {
208	// rowA_elem4[i] holds 8 * half numbers seen as a single float4.
209	rowA_elem4[i] = load_ntmprl(&af4[row_addr + threadid + K / 8 * i]);
210	}
211	colB_elem4x = bf4[threadid * 4 + 0];
212	colB_elem4y = bf4[threadid * 4 + 1];
213	colB_elem4z = bf4[threadid * 4 + 2];
214	colB_elem4w = bf4[threadid * 4 + 3];
215	}
216
217	scalar2_t Af2;
218	float2 S;
219
220	auto Ah2ptr = reinterpret_cast<scalar2_t*>(&rowA_elem4);
221	scalar2_t* ah2lptr;
222
223	#pragma unroll
224	for (int i = 0; i < NUM_A_ROWS_PER_BLOCK; i++) {
225	// Multiply-add on 8 scalar_t.
226	ah2lptr = Ah2ptr + i * 4;
227	Af2 = *(ah2lptr);
228	acch2 = __hmul2(Af2, colB_elem4x);
229	Af2 = *(ah2lptr + 1);
230	acch2 = __hfma2(Af2, colB_elem4y, acch2);
231	Af2 = *(ah2lptr + 2);
232	acch2 = __hfma2(Af2, colB_elem4z, acch2);
233	Af2 = *(ah2lptr + 3);
234	acch2 = __hfma2(Af2, colB_elem4w, acch2);
235	S = __s22float2(acch2);
236
237	// See comment above concerning the if guard.
238	acc[i] = (threadid * 8 < K ? S.x + S.y : 0.f);
239	}
240
241	// all reduce across warp.
242	#pragma unroll
243	for (int mask = WARP_SIZE / 2; mask >= 1; mask /= 2) {
244	#pragma unroll
245	for (int i = 0; i < NUM_A_ROWS_PER_BLOCK; i++) {
246	acc[i] += __shfl_xor(acc[i], mask);
247	}
248	}
249
250	// Warp leaders store the data to shared memory.
251	if (lane < NUM_A_ROWS_PER_BLOCK) {
252	red_smem[lane][warp] = acc[lane];
253	}
254
255	// Make sure the data is in shared memory.
256	__syncthreads();
257
258	if (qwarpid < NUM_A_ROWS_PER_BLOCK) {
259	acc[qwarpid] = qthreadid < num_warps ? red_smem[qwarpid][qthreadid] : 0.f;
260	#pragma unroll
261	for (int mask = 16 / 2; mask >= 1; mask /= 2) {
262	acc[qwarpid] += __shfl_xor(acc[qwarpid], mask);
263	}
264	float oval2 = __shfl_xor(acc[qwarpid], 16);
265
266	if (lane % 32 == 0) {
267	oval = __float22s2_rn<scalar2_t>(make_float2(acc[qwarpid], oval2));
268	c[blockIdx.x * NUM_A_ROWS_PER_BLOCK / 2 + qwarpid / 2] = oval;
269	}
270	}
271	}
272
273	torch::Tensor LLMM1(at::Tensor& in_a, at::Tensor& in_b,
274	const int64_t rows_per_block) {
275	auto M = in_a.size(0);
276	auto K = in_a.size(1);
277	auto N = in_b.size(0);
278
279	TORCH_CHECK(N == 1, "Row number of activation tensor must be 1.");
280	TORCH_CHECK(in_a.dtype() == in_b.dtype());
281	TORCH_CHECK(in_b.dtype() == torch::kFloat16 \|\|
282	in_b.dtype() == torch::kBFloat16);
283
284	auto out_c = torch::empty(
285	{N, M}, torch::TensorOptions().dtype(in_b.dtype()).device(in_b.device()));
286
287	// NUM_TREADS need to be a multiple of WARP_SIZE, as we are using warp shuffle
288	// operations.
289	const int NUM_THREADS =
290	max(rows_per_block * 16,
291	K * 2 / 16 % WARP_SIZE == 0
292	? K * 2 / 16
293	: K * 2 / 16 + (WARP_SIZE - K * 2 / 16 % WARP_SIZE));
294
295	int NUM_BLOCKS = M / rows_per_block;
296
297	const at::cuda::OptionalCUDAGuard device_guard(device_of(in_b));
298	const cudaStream_t stream = at::cuda::getCurrentCUDAStream();
299
300	// call the kernel function...
301	AT_DISPATCH_REDUCED_FLOATING_TYPES(in_b.scalar_type(), "LLGemm1", [&] {
302	auto a_ptr = in_a.data_ptr<scalar_t>();
303	auto b_ptr = in_b.data_ptr<scalar_t>();
304	auto c_ptr = out_c.data_ptr<scalar_t>();
305	if (rows_per_block == 2) {
306	LLGemm1_kernel<scalar_t, 2>
307	<<<NUM_BLOCKS, NUM_THREADS, 0, stream>>>(a_ptr, b_ptr, c_ptr, K);
308	} else if (rows_per_block == 4) {
309	LLGemm1_kernel<scalar_t, 4>
310	<<<NUM_BLOCKS, NUM_THREADS, 0, stream>>>(a_ptr, b_ptr, c_ptr, K);
311	} else if (rows_per_block == 8) {
312	LLGemm1_kernel<scalar_t, 8>
313	<<<NUM_BLOCKS, NUM_THREADS, 0, stream>>>(a_ptr, b_ptr, c_ptr, K);
314	} else if (rows_per_block == 16) {
315	LLGemm1_kernel<scalar_t, 16>
316	<<<NUM_BLOCKS, NUM_THREADS, 0, stream>>>(a_ptr, b_ptr, c_ptr, K);
317	} else {
318	NUM_BLOCKS = M / 4;
319	LLGemm1_kernel<scalar_t, 4>
320	<<<NUM_BLOCKS, NUM_THREADS, 0, stream>>>(a_ptr, b_ptr, c_ptr, K);
321	}
322	});
323
324	return out_c;
325	}
326
327	#if defined(__HIP__GFX9__) && !defined(__HIP__GFX1X__)
328	#define DOT2C(V0, V2, V3) \
329	if constexpr (std::is_same_v<scalar_t, half>) { \
330	asm("v_dot2c_f32_f16 %0, %2, %3" \
331	: "=v"(V0) \
332	: "0"(V0), "v"(V2), "v"(V3)); \
333	} else if constexpr (std::is_same_v<scalar_t, __hip_bfloat16>) { \
334	float2 s = __bfloat1622float2(((__hip_bfloat162)(&(V2)))) * \
335	__bfloat1622float2(((__hip_bfloat162)(&(V3)))); \
336	V0 += (s.x + s.y); \
337	}
338	#elif defined(__HIP__GFX1X__)
339	// gfx1x: v_dot2_f32_f16 (VOP3-P, dot10-insts, available on gfx11+gfx12)
340	#define DOT2C(V0, V2, V3) \
341	if constexpr (std::is_same_v<scalar_t, half>) { \
342	asm("v_dot2_f32_f16 %0, %1, %2, %0" : "+v"(V0) : "v"(V2), "v"(V3)); \
343	} else if constexpr (std::is_same_v<scalar_t, __hip_bfloat16>) { \
344	float2 s = __bfloat1622float2(((__hip_bfloat162)(&(V2)))) * \
345	__bfloat1622float2(((__hip_bfloat162)(&(V3)))); \
346	V0 += (s.x + s.y); \
347	}
348	#endif
349
350	// To avoid LLVM silently upcasting to double
351	__device__ inline unsigned int min__(uint32_t a, uint32_t b) {
352	return min(a, b);
353	}
354
355	#if defined(__HIP__GFX9__) \|\| defined(__HIP__GFX1X__)
356	// This version targets cases where A[] fits LDS capacity
357	template <typename scalar_t, int THRDS, int YTILE, int WvPrGrp, int A_CHUNK,
358	int UNRL, int N>
359	__global__ void __launch_bounds__(WvPrGrp* THRDS)
360	wvSplitK_hf_sml_(const int K, const int Kbp, const int Kap, const int M,
361	const int Bx, const int By, const scalar_t* B,
362	const scalar_t* __restrict__ A,
363	const scalar_t* __restrict__ BIAS, scalar_t* C,
364	const int _WvPrGrp, const int CuCount) {
365	constexpr int max_lds_len = LDS_SIZE / 2;
366	#if defined(__HIP__MI3XX__)
367	constexpr bool use_mfma = (std::is_same_v<scalar_t, __hip_bfloat16>);
368	#else
369	constexpr bool use_mfma = false;
370	#endif
371	using scalar8 =
372	__attribute__((__vector_size__((A_CHUNK / 2) * sizeof(float)))) float;
373	using half4 =
374	__attribute__((__vector_size__((A_CHUNK / 2) * sizeof(__bf16)))) __bf16;
375	union bigType {
376	scalar_t h[A_CHUNK];
377	float f[A_CHUNK / 2];
378	float2 f2[A_CHUNK / 4];
379	double d[A_CHUNK / 4];
380	half4 h4[A_CHUNK / 4];
381	scalar8 h8;
382	};
383
384	//----------------------------------------------------
385	// Reserving 64/160 KB of LDS to have 1 WG / CU
386	// Goal is to bring the activation matrix A to the LDS
387	// and use it across the lifetime of the work group
388	// TODO: When activation matrix is larger than 64 KB
389	// then this is not going to work!
390	//----------------------------------------------------
391	__shared__ scalar_t s[max_lds_len];
392
393	//----------------------------------------------------
394	// Fetch the activation matrix to LDS
395	// Loop iteration:
396	// - Each thread (lane) is fetching 8 elements (A_Chunk)
397	// - Each wave will fetch 64*8=> 512 elements
398	// - Each WG will fetch 512 * 16 => 8K elements
399	// - Then the WG will move to another 8 K elements
400	// TODO: Logic below will only work when K is multiple of 8
401	//----------------------------------------------------
402	for (uint32_t k = (threadIdx.y * THRDS + threadIdx.x) * A_CHUNK;
403	k < min__(Kap * N, max_lds_len); k += THRDS * WvPrGrp * A_CHUNK) {
404	#if defined(__gfx950__)
405	__builtin_amdgcn_global_load_lds((int)(&A[k]), (int)(&s[k]), 16, 0, 0);
406	#else
407	((bigType)(&s[k])) = ((bigType)(&A[k]));
408	#endif
409	}
410	__syncthreads();
411
412	if (threadIdx.y >= _WvPrGrp) return;
413
414	uint32_t m = (blockIdx.x * _WvPrGrp + (threadIdx.y % _WvPrGrp)) * YTILE;
415
416	//----------------------------------------------------
417	// Each wave works on a single column of weight matrix.
418	// There are 16 waves per WG, and hence, each WG is
419	// working on 16 columns of weight matrix. Moreover,
420	// we tile in column direction by YTILE, so when YTILE=1
421	// the above math is right, however, when YTILE=2 then
422	// each wave will be working on 2 columns and WG will
423	// be working on 32 columns.
424	//
425	// Top level loop that makes WGs persistent!
426	// - WGs iterates across columns of weight matrix
427	// - Each wave within WG works on a given column(s)
428	// - After completing first set of columns, WGs start
429	// working on the next set of available columns
430	//----------------------------------------------------
431	while (m < M) {
432	//----------------------------------------------------
433	// 'sum' accumulates the matrix A x B computation
434	// split across 64 lanes.
435	//
436	// YTILE represents how many column of weight matrix
437	// are being worked on by each wave.
438	//----------------------------------------------------
439	float sum[N][YTILE] = {};
440	scalar8 sum4[N][YTILE] = {};
441
442	for (uint32_t k1 = 0; k1 < K; k1 += THRDS * A_CHUNK * UNRL) {
443	bigType bigA[N][UNRL] = {};
444	bigType bigB[YTILE][UNRL];
445	// Fetch the weight matrix from memory!
446	#pragma unroll
447	for (uint32_t k2 = 0; k2 < UNRL; k2++) {
448	uint32_t k = k1 + k2 * THRDS * A_CHUNK;
449	uint32_t k_ = k + threadIdx.x * A_CHUNK;
450	const scalar_t* B_ = &B[min__(k_, K - A_CHUNK)];
451	for (int y = 0; y < YTILE; y++)
452	bigB[y][k2].h8 = (loadnt((scalar8)(&B_[min__(y + m, M - 1) Kbp])));
453	}
454
455	// Fetch activation matrix from either just LDS or from both LDS / memory
456	#pragma unroll
457	for (uint32_t k2 = 0; k2 < UNRL; k2++) {
458	uint32_t k = k1 + k2 * THRDS * A_CHUNK;
459	uint32_t k_ = k + threadIdx.x * A_CHUNK;
460	if (k_ >= K) break;
461	for (int n = 0; n < N; n++) {
462	bigA[n][k2] = ((const bigType)(&(s[k_ + Kap * n])));
463	}
464	}
465
466	// Do the matrix multiplication in interleaved manner
467	for (uint32_t k2 = 0; k2 < UNRL; k2++) {
468	for (uint32_t n = 0; n < N; n++) {
469	for (int y = 0; y < YTILE; y++) {
470	if constexpr (!use_mfma)
471	for (uint32_t b = 0; b < A_CHUNK / 2; b++) {
472	DOT2C(sum[n][y], bigA[n][k2].f[b], bigB[y][k2].f[b])
473	}
474	else
475	for (uint32_t b = 0; b < A_CHUNK / 4; b++)
476	sum4[n][y] = __builtin_amdgcn_mfma_f32_4x4x4bf16_1k(
477	bigA[n][k2].h4[b], bigB[y][k2].h4[b], sum4[n][y], 0, 0, 0);
478	}
479	}
480	}
481	}
482	__builtin_amdgcn_sched_barrier(0);
483	//----------------------------------------------------
484	// Final reduction step using shuffle
485	//----------------------------------------------------
486	if constexpr (!use_mfma) {
487	for (int n = 0; n < N; n++) {
488	for (int y = 0; y < YTILE; y++) {
489	sum[n][y] += __builtin_amdgcn_mov_dpp(sum[n][y], 0x118, 0xf, 0xf,
490	1); // row_shr8
491	sum[n][y] += __builtin_amdgcn_mov_dpp(sum[n][y], 0x114, 0xf, 0xf,
492	1); // row_shr4
493	sum[n][y] += __builtin_amdgcn_mov_dpp(sum[n][y], 0x112, 0xf, 0xf,
494	1); // row_shr2
495	sum[n][y] += __builtin_amdgcn_mov_dpp(sum[n][y], 0x111, 0xf, 0xf,
496	1); // row_shr1
497	#if defined(__HIP__GFX9__)
498	sum[n][y] += __builtin_amdgcn_mov_dpp(sum[n][y], 0x142, 0xf, 0xf,
499	1); // ROW_BCAST15
500	sum[n][y] += __builtin_amdgcn_mov_dpp(sum[n][y], 0x143, 0xf, 0xf,
501	1); // ROW_BCAST31
502	#else
503	sum[n][y] += __shfl_xor(sum[n][y], 16);
504	#endif
505	}
506	}
507
508	if (threadIdx.x == (THRDS - 1)) {
509	scalar_t biases[N][YTILE] = {};
510	if (BIAS)
511	for (int n = 0; n < N; n++) {
512	for (int y = 0; y < YTILE; y++) {
513	biases[n][y] = BIAS[(m + y) % Bx + (n % By) * Bx];
514	}
515	}
516	for (int n = 0; n < N; n++) {
517	for (int y = 0; y < YTILE; y++) {
518	if constexpr (std::is_same_v<scalar_t, half>) {
519	sum[n][y] += __half2float(biases[n][y]);
520	} else if constexpr (std::is_same_v<scalar_t, __hip_bfloat16>) {
521	sum[n][y] += __bfloat162float(biases[n][y]);
522	}
523	C[m + y + n * M] = __float2s<scalar_t>(sum[n][y]);
524	}
525	}
526	}
527	} else {
528	#ifdef __HIP__GFX9__
529	#pragma unroll
530	for (int n = 0; n < N; n++) {
531	#pragma unroll
532	for (int y = 0; y < YTILE; y++) {
533	/*float accm1 = 0;
534	for (int i=0; i<64; i++)
535	accm1 += __shfl(sum4[n][y][i%4], i);
536	sum4[n][y][0] = accm1;*/
537	float accm = sum4[n][y][0];
538	accm += __builtin_amdgcn_mov_dpp(sum4[n][y][1], 0x101, 0xf, 0xf,
539	1); // row_shl1
540	accm += __builtin_amdgcn_mov_dpp(sum4[n][y][2], 0x102, 0xf, 0xf,
541	1); // row_shl2
542	accm += __builtin_amdgcn_mov_dpp(sum4[n][y][3], 0x103, 0xf, 0xf,
543	1); // row_shl3
544	accm += __builtin_amdgcn_mov_dpp(accm, 0x104, 0xf, 0xf,
545	1); // row_shl4
546	accm += __builtin_amdgcn_mov_dpp(accm, 0x108, 0xf, 0xf,
547	1); // row_shl8
548	accm = __builtin_amdgcn_mov_dpp(accm, 0x11f, 0xf, 0xf,
549	1); // row_shr15
550	accm += __builtin_amdgcn_mov_dpp(accm, 0x142, 0xf, 0xf,
551	1); // ROW_BCAST15
552	accm += __builtin_amdgcn_mov_dpp(accm, 0x143, 0xf, 0xf,
553	1); // ROW_BCAST31
554
555	sum4[n][y][0] = accm;
556	}
557	}
558	if (threadIdx.x == (THRDS - 1)) {
559	scalar_t biases[N][YTILE] = {};
560	if (BIAS)
561	for (int n = 0; n < N; n++) {
562	for (int y = 0; y < YTILE; y++) {
563	biases[n][y] = BIAS[(m + y) % Bx + (n % By) * Bx];
564	}
565	}
566	for (int n = 0; n < N; n++) {
567	for (int y = 0; y < YTILE; y++) {
568	sum4[n][y][0] += __bfloat162float(biases[n][y]);
569	C[m + y + n * M] = __float2bfloat16(sum4[n][y][0]);
570	}
571	}
572	}
573	#endif // __HIP__GFX9__ (MFMA path)
574	}
575	m += CuCount * _WvPrGrp * YTILE;
576	}
577	}
578	#else
579	template <typename scalar_t, int THRDS, int YTILE, int WvPrGrp, int A_CHUNK,
580	int UNRL, int N>
581	__global__ void wvSplitK_hf_sml_(const int K, const int Kbp, const int Kap,
582	const int M, const int Bx, const int By,
583	const scalar_t* B,
584	const scalar_t* __restrict__ A,
585	const scalar_t* __restrict__ BIAS, scalar_t* C,
586	const int _WvPrGrp, const int CuCount) {
587	UNREACHABLE_CODE
588	}
589	#endif
590
591	#if defined(__HIP__GFX9__) \|\| defined(__HIP__GFX1X__)
592	// This version targets cases where A[] marginally exceeds LDS capacity
593	template <typename scalar_t, int THRDS, int YTILE, int WvPrGrp, int A_CHUNK,
594	int UNRL, int N>
595	__global__ void __launch_bounds__(WvPrGrp* THRDS)
596	wvSplitK_hf_(const int K, const int Kbp, const int Kap, const int M,
597	const int Bx, const int By, const scalar_t* B,
598	const scalar_t* __restrict__ A,
599	const scalar_t* __restrict__ BIAS, scalar_t* C,
600	const int _WvPrGrp, const int CuCount) {
601	constexpr int max_lds_len = LDS_SIZE / 2;
602	#if defined(__HIP__MI3XX__)
603	constexpr bool use_mfma = (std::is_same_v<scalar_t, __hip_bfloat16>);
604	#else
605	constexpr bool use_mfma = false;
606	#endif
607
608	using scalar8 =
609	__attribute__((__vector_size__((A_CHUNK / 2) * sizeof(float)))) float;
610	using half4 =
611	__attribute__((__vector_size__((A_CHUNK / 2) * sizeof(__bf16)))) __bf16;
612	union bigType {
613	scalar_t h[A_CHUNK];
614	float f[A_CHUNK / 2];
615	float2 f2[A_CHUNK / 4];
616	double d[A_CHUNK / 4];
617	half4 h4[A_CHUNK / 4];
618	scalar8 h8;
619	};
620
621	__shared__ scalar_t s[max_lds_len];
622
623	//----------------------------------------------------
624	// Computation of columns that need to be committed to memory!
625	//----------------------------------------------------
626	uint32_t commitColumn[YTILE];
627	for (uint32_t i = 0; i < YTILE; i++) {
628	commitColumn[i] = 1;
629	}
630
631	uint32_t m = (blockIdx.x * _WvPrGrp + threadIdx.y) * YTILE;
632
633	// Check whether there will be fragmentation!
634	// This will happen only for the last wave!
635	if (m < M && (m + YTILE) >= M) {
636	uint32_t startColumn = M - YTILE;
637	for (uint32_t i = 0; i < (m - startColumn); i++) {
638	commitColumn[i] = 0;
639	}
640	m = startColumn;
641	}
642
643	for (uint32_t k = (threadIdx.y * THRDS + threadIdx.x) * A_CHUNK;
644	k < min__(Kap * N, max_lds_len); k += THRDS * WvPrGrp * A_CHUNK) {
645	#if defined(__gfx950__)
646	__builtin_amdgcn_global_load_lds((int)(&A[k]), (int)(&s[k]), 16, 0, 0);
647	#else
648	((bigType)(&s[k])) = ((bigType)(&A[k]));
649	#endif
650	}
651
652	__syncthreads();
653
654	if (threadIdx.y >= _WvPrGrp) return;
655
656	while (m < M) {
657	float sum[N][YTILE] = {};
658	scalar8 sum4[N][YTILE] = {};
659
660	for (uint32_t k1 = 0; k1 < K; k1 += THRDS * A_CHUNK * UNRL) {
661	bigType bigA[N][UNRL] = {};
662	bigType bigB[YTILE][UNRL];
663	// Fetch the weight matrix from memory!
664	#pragma unroll
665	for (uint32_t k2 = 0; k2 < UNRL; k2++) {
666	uint32_t k = k1 + k2 * THRDS * A_CHUNK;
667	uint32_t k_ = k + threadIdx.x * A_CHUNK;
668	const scalar_t* B_ = &B[min__(k_, K - A_CHUNK)];
669	for (int y = 0; y < YTILE; y++)
670	bigB[y][k2].h8 = (loadnt((scalar8)(&B_[min__(y + m, M - 1) Kbp])));
671	}
672
673	// Fetch activation matrix from either just LDS or from both LDS / memory
674	#pragma unroll
675	for (uint32_t k2 = 0; k2 < UNRL; k2++) {
676	uint32_t k = k1 + k2 * THRDS * A_CHUNK;
677	uint32_t k_ = k + threadIdx.x * A_CHUNK;
678	if (k_ >= K) break;
679	for (int n = 0; n < N; n++) {
680	if (k_ + Kap * n < max_lds_len)
681	bigA[n][k2] = ((const bigType)(&(s[k_ + Kap * n])));
682	else
683	bigA[n][k2] = ((const bigType)(&(A[k_ + Kap * n])));
684	}
685	}
686
687	// Do the matrix multiplication in interleaved manner
688	for (uint32_t n = 0; n < N; n++) {
689	for (uint32_t k2 = 0; k2 < UNRL; k2++) {
690	for (int y = 0; y < YTILE; y++) {
691	if constexpr (!use_mfma)
692	for (uint32_t b = 0; b < A_CHUNK / 2; b++) {
693	DOT2C(sum[n][y], bigA[n][k2].f[b], bigB[y][k2].f[b])
694	}
695	else
696	for (uint32_t b = 0; b < A_CHUNK / 4; b++)
697	sum4[n][y] = __builtin_amdgcn_mfma_f32_4x4x4bf16_1k(
698	bigA[n][k2].h4[b], bigB[y][k2].h4[b], sum4[n][y], 0, 0, 0);
699	}
700	}
701	}
702	}
703
704	//----------------------------------------------------
705	// Final reduction step using shuffle
706	//----------------------------------------------------
707	if constexpr (!use_mfma) {
708	for (int n = 0; n < N; n++) {
709	for (int y = 0; y < YTILE; y++) {
710	sum[n][y] += __builtin_amdgcn_mov_dpp(sum[n][y], 0x118, 0xf, 0xf,
711	1); // row_shr8
712	sum[n][y] += __builtin_amdgcn_mov_dpp(sum[n][y], 0x114, 0xf, 0xf,
713	1); // row_shr4
714	sum[n][y] += __builtin_amdgcn_mov_dpp(sum[n][y], 0x112, 0xf, 0xf,
715	1); // row_shr2
716	sum[n][y] += __builtin_amdgcn_mov_dpp(sum[n][y], 0x111, 0xf, 0xf,
717	1); // row_shr1
718	#if defined(__HIP__GFX9__)
719	sum[n][y] += __builtin_amdgcn_mov_dpp(sum[n][y], 0x142, 0xf, 0xf,
720	1); // ROW_BCAST15
721	sum[n][y] += __builtin_amdgcn_mov_dpp(sum[n][y], 0x143, 0xf, 0xf,
722	1); // ROW_BCAST31
723	#else
724	sum[n][y] += __shfl_xor(sum[n][y], 16);
725	#endif
726	}
727	}
728
729	if (threadIdx.x == (THRDS - 1)) {
730	scalar_t biases[N][YTILE] = {};
731	if (BIAS)
732	for (int n = 0; n < N; n++) {
733	for (int y = 0; y < YTILE; y++) {
734	biases[n][y] = BIAS[(m + y) % Bx + (n % By) * Bx];
735	}
736	}
737	for (int n = 0; n < N; n++) {
738	for (int y = 0; y < YTILE; y++) {
739	if (commitColumn[y]) {
740	if constexpr (std::is_same_v<scalar_t, half>) {
741	sum[n][y] += __half2float(biases[n][y]);
742	} else if constexpr (std::is_same_v<scalar_t, __hip_bfloat16>) {
743	sum[n][y] += __bfloat162float(biases[n][y]);
744	}
745	C[m + y + n * M] = __float2s<scalar_t>(sum[n][y]);
746	}
747	}
748	}
749	}
750	} else {
751	#ifdef __HIP__GFX9__
752	#pragma unroll
753	for (int n = 0; n < N; n++) {
754	#pragma unroll
755	for (int y = 0; y < YTILE; y++) {
756	// float accm1 = 0;
757	// for (int i=0; i<64; i++)
758	// accm1 += __shfl(sum4[n][y][i%4], i);
759	float accm = sum4[n][y][0];
760	accm += __builtin_amdgcn_mov_dpp(sum4[n][y][1], 0x101, 0xf, 0xf,
761	1); // row_shl1
762	accm += __builtin_amdgcn_mov_dpp(sum4[n][y][2], 0x102, 0xf, 0xf,
763	1); // row_shl2
764	accm += __builtin_amdgcn_mov_dpp(sum4[n][y][3], 0x103, 0xf, 0xf,
765	1); // row_shl3
766	accm += __builtin_amdgcn_mov_dpp(accm, 0x104, 0xf, 0xf,
767	1); // row_shl4
768	accm += __builtin_amdgcn_mov_dpp(accm, 0x108, 0xf, 0xf,
769	1); // row_shl8
770	accm = __builtin_amdgcn_mov_dpp(accm, 0x11f, 0xf, 0xf,
771	1); // row_shr15
772	accm += __builtin_amdgcn_mov_dpp(accm, 0x142, 0xf, 0xf,
773	1); // ROW_BCAST15
774	accm += __builtin_amdgcn_mov_dpp(accm, 0x143, 0xf, 0xf,
775	1); // ROW_BCAST31
776	sum4[n][y][0] = accm;
777	}
778	}
779	if (threadIdx.x == (THRDS - 1)) {
780	scalar_t biases[N][YTILE] = {};
781	if (BIAS)
782	for (int n = 0; n < N; n++) {
783	for (int y = 0; y < YTILE; y++) {
784	biases[n][y] = BIAS[(m + y) % Bx + (n % By) * Bx];
785	}
786	}
787	for (int n = 0; n < N; n++) {
788	for (int y = 0; y < YTILE; y++) {
789	if (commitColumn[y]) {
790	sum4[n][y][0] += __bfloat162float(biases[n][y]);
791	C[m + y + n * M] = __float2bfloat16(sum4[n][y][0]);
792	}
793	}
794	}
795	}
796	#endif // __HIP__GFX9__ (MFMA path)
797	}
798
799	m += CuCount * _WvPrGrp * YTILE;
800
801	// Check whether there will be fragmentation!
802	// This will happen only for the last wave!
803	if (m < M && (m + YTILE) >= M) {
804	uint32_t startColumn = M - YTILE;
805	for (uint32_t i = 0; i < (m - startColumn); i++) {
806	commitColumn[i] = 0;
807	}
808	m = startColumn;
809	}
810	}
811	}
812
813	#else
814	template <typename scalar_t, int THRDS, int YTILE, int WvPrGrp, int A_CHUNK,
815	int UNRL, int N>
816	__global__ void wvSplitK_hf_(const int K, const int Kbp, const int Kap,
817	const int M, const int Bx, const int By,
818	const scalar_t* B, const scalar_t* __restrict__ A,
819	const scalar_t* __restrict__ BIAS, scalar_t* C,
820	const int _WvPrGrp, const int CuCount) {
821	UNREACHABLE_CODE
822	}
823	#endif
824
825	#if defined(__HIP__GFX9__) \|\| defined(__HIP__GFX1X__)
826	// This version targets big A[] cases, where it is much larger than LDS capacity
827	template <typename scalar_t, int THRDS, int YTILE, int WvPrGrp, int A_CHUNK,
828	int UNRL, int N>
829	__global__ void __launch_bounds__(WvPrGrp* THRDS)
830	wvSplitK_hf_big_(const int K, const int Kbp, const int Kap, const int M,
831	const int Bx, const int By, const scalar_t* B,
832	const scalar_t* __restrict__ A,
833	const scalar_t* __restrict__ BIAS, scalar_t* C,
834	const int _WvPrGrp, const int CuCount) {
835	constexpr int max_lds_len = LDS_SIZE / 2;
836	#if defined(__HIP__MI3XX__)
837	constexpr bool use_mfma = (std::is_same_v<scalar_t, __hip_bfloat16>);
838	#else
839	constexpr bool use_mfma = false;
840	#endif
841
842	using scalar8 =
843	__attribute__((__vector_size__((A_CHUNK / 2) * sizeof(float)))) float;
844	using half4 =
845	__attribute__((__vector_size__((A_CHUNK / 2) * sizeof(__bf16)))) __bf16;
846	union bigType {
847	scalar_t h[A_CHUNK];
848	float f[A_CHUNK / 2];
849	float2 f2[A_CHUNK / 4];
850	double d[A_CHUNK / 4];
851	half4 h4[A_CHUNK / 4];
852	scalar8 h8;
853	};
854
855	//----------------------------------------------------
856	// Reserving 64/160 KB of LDS to have 1 WG / CU
857	// Goal is to bring the activation matrix A to the LDS
858	// and use it across the lifetime of the work group
859	// TODO: When activation matrix is larger than 64 KB
860	// then this is not going to work!
861	//----------------------------------------------------
862	__shared__ scalar_t s[max_lds_len];
863
864	//----------------------------------------------------
865	// Computation of columns that need to be committed to memory!
866	//----------------------------------------------------
867	uint32_t commitColumn[YTILE];
868	for (uint32_t i = 0; i < YTILE; i++) {
869	commitColumn[i] = 1;
870	}
871
872	// int _WvPrGrp = mindiv(N, CuCount * YTILE, WvPrGrp);
873	if (threadIdx.y >= _WvPrGrp) return;
874
875	//----------------------------------------------------
876	// Indexing function into the column of weight matrix B
877	// Algorithm does 64 lane k-splitting / wave and uses
878	// WG ID and Thread ID to find the index.
879	//----------------------------------------------------
880	uint32_t m = (blockIdx.x * _WvPrGrp + threadIdx.y) * YTILE;
881
882	// Check whether there will be fragmentation!
883	// This will happen only for the last wave!
884	if (m < M && (m + YTILE) >= M) {
885	uint32_t startColumn = M - YTILE;
886	for (uint32_t i = 0; i < (m - startColumn); i++) {
887	commitColumn[i] = 0;
888	}
889	m = startColumn;
890	}
891
892	//----------------------------------------------------
893	// Fetch the activation matrix to LDS
894	// Loop iteration:
895	// - Each thread (lane) is fetching 8 elements (A_Chunk)
896	// - Each wave will fetch 64*8=> 512 elements
897	// - Each WG will fetch 512 * 16 => 8K elements
898	// - Then the WG will move to another 8 K elements
899	// TODO: Logic below will only work when K is multiple of 8
900	//----------------------------------------------------
901	#define PCML
902	#ifndef PCML
903	for (uint32_t k = (threadIdx.y * THRDS + threadIdx.x) * A_CHUNK;
904	k < min__(Kap * N, max_lds_len); k += THRDS * WvPrGrp * A_CHUNK) {
905	#if defined(__gfx950__)
906	__builtin_amdgcn_global_load_lds((int)(&A[k]), (int)(&s[k]), 16, 0, 0);
907	#else
908	((bigType)(&s[k])) = ((bigType)(&A[k]));
909	#endif
910	}
911	__syncthreads();
912	#endif
913
914	#define TUC (THRDS * UNRL * A_CHUNK)
915	uint32_t kBase = 0;
916	// find biggest k size that fits in LDS
917	uint32_t kFit = (max_lds_len) / N;
918	// kFit = (kFit%TWC==0) ? kFit : (kFit-kFit%TWC+TWC); //round up to multiple
919	// of TUC
920	kFit = (kFit % TUC == 0)
921	? kFit
922	: (kFit - kFit % TUC); // round up to multiple of TUC
923	// if (kFit == 0) kFit = TUC;
924	kFit = min__(kFit, Kap);
925
926	//----------------------------------------------------
927	// Each wave works on a single column of weight matrix.
928	// There are 16 waves per WG, and hence, each WG is
929	// working on 16 columns of weight matrix. Moreover,
930	// we tile in column direction by YTILE, so when YTILE=1
931	// the above math is right, however, when YTILE=2 then
932	// each wave will be working on 2 columns and WG will
933	// be working on 32 columns.
934	//
935	// Top level loop that makes WGs persistent!
936	// - WGs iterates across columns of weight matrix
937	// - Each wave within WG works on a given column(s)
938	// - After completing first set of columns, WGs start
939	// working on the next set of available columns
940	//----------------------------------------------------
941	#ifdef PCML
942	int YW = (YTILE * _WvPrGrp);
943	uint32_t Mrndp = (M % YW == 0) ? M : (M - M % YW + YW);
944	while (m < Mrndp) {
945	#else
946	while (m < M) {
947	#endif
948	//----------------------------------------------------
949	// 'sum' accumulates the matrix A x B computation
950	// split across 64 lanes.
951	//
952	// YTILE represents how many column of weight matrix
953	// are being worked on by each wave.
954	//----------------------------------------------------
955	float sum[N][YTILE] = {};
956	scalar8 sum4[N][YTILE] = {};
957
958	//----------------------------------------------------
959	// Fetch weight matrix B in interleaved K-split!
960	// - Each thread (lane) is fetching 8 elements (A_Chunk)
961	// - Each wave will fetch 64*8=> 512 elements (1024B)
962	// - YTILE represents the number of column being serviced
963	// by wave
964	// - Loop for fetching weight matrix (B) are unrolled
965	//
966	// Fetch activation matrix A from LDS
967	// - Loop for fetching activation matrix (A) are unrolled
968	//
969	// Finally, do the matrix multiplication in an unrolled
970	// fashion. This provides lot of food for compiler
971	// scheduling.
972	//
973	// TODO: Logic below will only work when K is multiple of 8
974	//----------------------------------------------------
975	for (uint32_t k1 = 0; k1 < K; k1 += THRDS * A_CHUNK * UNRL) {
976	bigType bigA[N][UNRL] = {};
977	bigType bigB[YTILE][UNRL];
978
979	#ifdef PCML
980	if ((k1 == 0) \|\| (k1 == kBase + kFit)) { // load next chunk of A[] to LDS
981	if (k1 != 0) kBase += kFit;
982	__syncthreads();
983	for (uint32_t k = 0; k < kFit; k += THRDS * _WvPrGrp * A_CHUNK) {
984	uint32_t kOff = k + ((threadIdx.y * THRDS + threadIdx.x) * A_CHUNK);
985	if (kBase + kOff >= Kap) break;
986	if (kOff >= kFit) break;
987	for (uint32_t n = 0; n < N; n++) {
988	uint32_t k_in = kBase + n * Kap + kOff;
989	uint32_t k_ot = n * kFit + kOff;
990	#if defined(__gfx950__)
991	__builtin_amdgcn_global_load_lds((int)(&A[k_in]), (int)(&s[k_ot]),
992	16, 0, 0);
993	#else
994	((bigType)(&s[k_ot])) = ((bigType)(&A[k_in]));
995	#endif
996	}
997	}
998	__syncthreads();
999	}
1000	if (m >= M) continue;
1001	#endif
1002
1003	// Fetch the weight matrix from memory!
1004	#pragma unroll
1005	for (uint32_t k2 = 0; k2 < UNRL; k2++) {
1006	uint32_t k = k1 + k2 * THRDS * A_CHUNK;
1007	uint32_t k_ = k + threadIdx.x * A_CHUNK;
1008	const scalar_t* B_ = &B[min__(k_, K - A_CHUNK)];
1009	for (int y = 0; y < YTILE; y++)
1010	bigB[y][k2].h8 = (loadnt((scalar8)(&B_[min__(y + m, M - 1) Kbp])));
1011	}
1012
1013	// Fetch activation matrix from either just LDS or from both LDS / memory
1014	#pragma unroll
1015	for (uint32_t k2 = 0; k2 < UNRL; k2++) {
1016	uint32_t k = k1 + k2 * THRDS * A_CHUNK;
1017	uint32_t k_ = k + threadIdx.x * A_CHUNK;
1018	if (k_ >= K) break;
1019	for (int n = 0; n < N; n++) {
1020	#ifdef PCML
1021	bigA[n][k2] = ((const bigType)(&(s[k_ - kBase + kFit * n])));
1022	#else
1023	if (k_ + Kap * n < max_lds_len)
1024	bigA[n][k2] = ((const bigType)(&(s[k_ + Kap * n])));
1025	else
1026	bigA[n][k2] = ((const bigType)(&(A[k_ + Kap * n])));
1027	#endif
1028	}
1029	}
1030
1031	// Do the matrix multiplication in interleaved manner
1032	#pragma unroll
1033	for (uint32_t k2 = 0; k2 < UNRL; k2++) {
1034	for (uint32_t n = 0; n < N; n++) {
1035	for (int y = 0; y < YTILE; y++) {
1036	if constexpr (!use_mfma)
1037	for (uint32_t b = 0; b < A_CHUNK / 2; b++) {
1038	DOT2C(sum[n][y], bigA[n][k2].f[b], bigB[y][k2].f[b])
1039	}
1040	else
1041	for (uint32_t b = 0; b < A_CHUNK / 4; b++)
1042	sum4[n][y] = __builtin_amdgcn_mfma_f32_4x4x4bf16_1k(
1043	bigA[n][k2].h4[b], bigB[y][k2].h4[b], sum4[n][y], 0, 0, 0);
1044	}
1045	}
1046	}
1047	}
1048
1049	#ifdef PCML
1050	if (m >= M) {
1051	m += CuCount * _WvPrGrp * YTILE;
1052	kBase = 0;
1053	continue;
1054	}
1055	#endif
1056
1057	//----------------------------------------------------
1058	// Final reduction step using shuffle
1059	//----------------------------------------------------
1060	if constexpr (!use_mfma) {
1061	for (int n = 0; n < N; n++) {
1062	for (int y = 0; y < YTILE; y++) {
1063	sum[n][y] += __builtin_amdgcn_mov_dpp(sum[n][y], 0x118, 0xf, 0xf,
1064	1); // row_shr8
1065	sum[n][y] += __builtin_amdgcn_mov_dpp(sum[n][y], 0x114, 0xf, 0xf,
1066	1); // row_shr4
1067	sum[n][y] += __builtin_amdgcn_mov_dpp(sum[n][y], 0x112, 0xf, 0xf,
1068	1); // row_shr2
1069	sum[n][y] += __builtin_amdgcn_mov_dpp(sum[n][y], 0x111, 0xf, 0xf,
1070	1); // row_shr1
1071	#if defined(__HIP__GFX9__)
1072	sum[n][y] += __builtin_amdgcn_mov_dpp(sum[n][y], 0x142, 0xf, 0xf,
1073	1); // ROW_BCAST15
1074	sum[n][y] += __builtin_amdgcn_mov_dpp(sum[n][y], 0x143, 0xf, 0xf,
1075	1); // ROW_BCAST31
1076	#else
1077	sum[n][y] += __shfl_xor(sum[n][y], 16);
1078	#endif
1079	}
1080	}
1081
1082	if (threadIdx.x == (THRDS - 1)) {
1083	scalar_t biases[N][YTILE] = {};
1084	if (BIAS)
1085	for (int n = 0; n < N; n++) {
1086	for (int y = 0; y < YTILE; y++) {
1087	biases[n][y] = BIAS[(m + y) % Bx + (n % By) * Bx];
1088	}
1089	}
1090	for (int n = 0; n < N; n++) {
1091	for (int y = 0; y < YTILE; y++) {
1092	if (commitColumn[y]) {
1093	if constexpr (std::is_same_v<scalar_t, half>) {
1094	sum[n][y] += __half2float(biases[n][y]);
1095	} else if constexpr (std::is_same_v<scalar_t, __hip_bfloat16>) {
1096	sum[n][y] += __bfloat162float(biases[n][y]);
1097	}
1098	C[m + y + n * M] = __float2s<scalar_t>(sum[n][y]);
1099	}
1100	}
1101	}
1102	}
1103	} else {
1104	#ifdef __HIP__GFX9__
1105	#pragma unroll
1106	for (int n = 0; n < N; n++) {
1107	#pragma unroll
1108	for (int y = 0; y < YTILE; y++) {
1109	float accm = sum4[n][y][0];
1110	accm += __builtin_amdgcn_mov_dpp(sum4[n][y][1], 0x101, 0xf, 0xf,
1111	1); // row_shl1
1112	accm += __builtin_amdgcn_mov_dpp(sum4[n][y][2], 0x102, 0xf, 0xf,
1113	1); // row_shl2
1114	accm += __builtin_amdgcn_mov_dpp(sum4[n][y][3], 0x103, 0xf, 0xf,
1115	1); // row_shl3
1116	accm += __builtin_amdgcn_mov_dpp(accm, 0x104, 0xf, 0xf,
1117	1); // row_shl4
1118	accm += __builtin_amdgcn_mov_dpp(accm, 0x108, 0xf, 0xf,
1119	1); // row_shl8
1120	accm = __builtin_amdgcn_mov_dpp(accm, 0x11f, 0xf, 0xf,
1121	1); // row_shr15
1122	accm += __builtin_amdgcn_mov_dpp(accm, 0x142, 0xf, 0xf,
1123	1); // ROW_BCAST15
1124	accm += __builtin_amdgcn_mov_dpp(accm, 0x143, 0xf, 0xf,
1125	1); // ROW_BCAST31
1126	sum4[n][y][0] = accm;
1127	}
1128	}
1129	if (threadIdx.x == (THRDS - 1)) {
1130	scalar_t biases[N][YTILE] = {};
1131	if (BIAS)
1132	for (int n = 0; n < N; n++) {
1133	for (int y = 0; y < YTILE; y++) {
1134	biases[n][y] = BIAS[(m + y) % Bx + (n % By) * Bx];
1135	}
1136	}
1137	for (int n = 0; n < N; n++) {
1138	for (int y = 0; y < YTILE; y++) {
1139	if (commitColumn[y]) {
1140	sum4[n][y][0] += __bfloat162float(biases[n][y]);
1141	C[m + y + n * M] = __float2bfloat16(sum4[n][y][0]);
1142	}
1143	}
1144	}
1145	}
1146	#endif // __HIP__GFX9__ (MFMA path)
1147	}
1148
1149	m += CuCount * _WvPrGrp * YTILE;
1150	kBase = 0;
1151
1152	// Check whether there will be fragmentation!
1153	// This will happen only for the last wave!
1154	if (m < M && (m + YTILE) >= M) {
1155	uint32_t startColumn = M - YTILE;
1156	for (uint32_t i = 0; i < (m - startColumn); i++) {
1157	commitColumn[i] = 0;
1158	}
1159	m = startColumn;
1160	}
1161	}
1162	}
1163	#else
1164	template <typename scalar_t, int THRDS, int YTILE, int WvPrGrp, int A_CHUNK,
1165	int UNRL, int N>
1166	__global__ void wvSplitK_hf_big_(const int K, const int Kbp, const int Kap,
1167	const int M, const int Bx, const int By,
1168	const scalar_t* B,
1169	const scalar_t* __restrict__ A,
1170	const scalar_t* __restrict__ BIAS, scalar_t* C,
1171	const int _WvPrGrp, const int CuCount) {
1172	UNREACHABLE_CODE
1173	}
1174	#endif
1175
1176	// Find the min val of div2 that doesn't increase N/(div1*div2)
1177	int mindiv(int N, int div1, int div2) {
1178	int nPrRnd = div1 * div2;
1179	int rnds[13];
1180	for (int i = 0; i < 13; i++) {
1181	rnds[i] = (N + nPrRnd - 1) / nPrRnd;
1182	nPrRnd -= div1;
1183	}
1184	for (int i = 12; i >= 0; i--)
1185	if (rnds[0] == rnds[i]) return (div2 - i);
1186	return 0;
1187	}
1188
1189	torch::Tensor wvSplitK(const at::Tensor& in_a, const at::Tensor& in_b,
1190	const std::optional<at::Tensor>& in_bias,
1191	const int64_t CuCount) {
1192	auto M_in = in_a.size(0);
1193	auto K_in = in_a.size(1);
1194	auto N_in = in_b.size(0);
1195	auto Kap_in = in_a.stride(0);
1196	auto Kbp_in = in_b.stride(0);
1197	auto Bx_in =
1198	(in_bias.has_value() && in_bias->numel() > 0)
1199	? (in_bias->sizes().size() == 2) ? in_bias->size(1) : in_bias->size(0)
1200	: 1;
1201	auto By_in = (in_bias.has_value() && in_bias->numel() > 0 &&
1202	in_bias->sizes().size() == 2)
1203	? in_bias->size(0)
1204	: 1;
1205
1206	TORCH_CHECK(in_a.dtype() == in_b.dtype());
1207	TORCH_CHECK(K_in % 8 == 0, "k % 8 == 0");
1208	TORCH_CHECK(in_a.dtype() == torch::kFloat16 \|\|
1209	in_a.dtype() == torch::kBFloat16);
1210
1211	auto out_c = torch::empty(
1212	{N_in, M_in},
1213	torch::TensorOptions().dtype(in_b.dtype()).device(in_b.device()));
1214
1215	dim3 grid(CuCount);
1216
1217	const at::cuda::OptionalCUDAGuard device_guard(device_of(in_a));
1218	const cudaStream_t stream = at::cuda::getCurrentCUDAStream();
1219	const int max_lds_len = get_lds_size() / 2;
1220
1221	#define WVSPLITK_CFG(_THRDS, _WVPRGRP, _YTILE, _UNRL, _N) \
1222	{ \
1223	dim3 block(_THRDS, _WVPRGRP); \
1224	int __wvPrGrp = mindiv(M_in, CuCount * _YTILE, _WVPRGRP); \
1225	if ((Kbp_in * N_in <= max_lds_len) && (M_in % _YTILE == 0)) \
1226	wvSplitK_hf_sml_<fptype, _THRDS, _YTILE, _WVPRGRP, 8, _UNRL, _N> \
1227	<<<grid, block, 0, stream>>>(K_in, Kap_in, Kbp_in, M_in, Bx_in, \
1228	By_in, af4, bf4, biasf4, c, __wvPrGrp, \
1229	CuCount); \
1230	else if (Kbp_in * N_in <= max_lds_len * 1.2) \
1231	wvSplitK_hf_<fptype, _THRDS, _YTILE, _WVPRGRP, 8, _UNRL, _N> \
1232	<<<grid, block, 0, stream>>>(K_in, Kap_in, Kbp_in, M_in, Bx_in, \
1233	By_in, af4, bf4, biasf4, c, __wvPrGrp, \
1234	CuCount); \
1235	else \
1236	wvSplitK_hf_big_<fptype, _THRDS, _YTILE, _WVPRGRP, 8, _UNRL, _N> \
1237	<<<grid, block, 0, stream>>>(K_in, Kap_in, Kbp_in, M_in, Bx_in, \
1238	By_in, af4, bf4, biasf4, c, __wvPrGrp, \
1239	CuCount); \
1240	}
1241
1242	#define WVSPLIT_TILE_CFG(_THRDS, _WVPRGRP, _sYT, __N) \
1243	{ \
1244	bool fit_lds = (Kbp_in * N_in <= max_lds_len); \
1245	if (_sYT <= 1) \
1246	WVSPLITK_CFG(_THRDS, _WVPRGRP, 1, 4, __N) \
1247	else if ((__N == 1) \|\| (!fit_lds) \|\| (_sYT <= 4 * 2)) \
1248	WVSPLITK_CFG(_THRDS, _WVPRGRP, 2, 2, __N) \
1249	else if (_sYT <= 4 * 3) \
1250	WVSPLITK_CFG(_THRDS, _WVPRGRP, 3, 2, __N) \
1251	else if (__N == 4) \
1252	WVSPLITK_CFG(_THRDS, _WVPRGRP, 4, 1, __N) \
1253	else \
1254	WVSPLITK_CFG(_THRDS, _WVPRGRP, 4, 2, __N) \
1255	}
1256
1257	// WVSPLITK_CFG arguments are: (THRDS, WVPRGRP, YTILE, UNRL, N).
1258	// THRDS = wavefront width (32 on GFX11/GFX12, 64 on GFX9)
1259	// WVPRGRP= waves per group (always 16)
1260	// YTILE = output rows per thread tile
1261	// UNRL = K-loop unroll factor
1262	// N = batch size (passed through from the switch in wvSplitK)
1263	#define WVSPLIT_TILE(_sYT, __N) \
1264	{ \
1265	if (on_gfx1151()) { \
1266	bool fit_lds = (Kbp_in * N_in <= max_lds_len); \
1267	if (_sYT <= 1) \
1268	WVSPLITK_CFG(/THRDS=/32, /WVPRGRP=/16, /YTILE=/1, /UNRL=/4, \
1269	__N) \
1270	else if ((K_in % 1024 == 512) && K_in >= 1536 && \
1271	(_sYT >= 40 \|\| K_in >= 4096)) \
1272	WVSPLITK_CFG(/THRDS=/32, /WVPRGRP=/16, /YTILE=/4, /UNRL=/1, \
1273	__N) \
1274	else if (K_in < 1024) \
1275	WVSPLITK_CFG(/THRDS=/32, /WVPRGRP=/16, /YTILE=/2, /UNRL=/4, \
1276	__N) \
1277	else if (K_in <= 2048 && (__N >= 2 \|\| _sYT <= 26)) \
1278	WVSPLITK_CFG(/THRDS=/32, /WVPRGRP=/16, /YTILE=/1, /UNRL=/4, \
1279	__N) \
1280	else if (__N >= 2 && !fit_lds) \
1281	WVSPLITK_CFG(/THRDS=/32, /WVPRGRP=/16, /YTILE=/1, /UNRL=/4, \
1282	__N) \
1283	else if (__N == 1) \
1284	WVSPLITK_CFG(/THRDS=/32, /WVPRGRP=/16, /YTILE=/1, /UNRL=/2, \
1285	__N) \
1286	else \
1287	WVSPLITK_CFG(/THRDS=/32, /WVPRGRP=/16, /YTILE=/1, /UNRL=/1, \
1288	__N) \
1289	} else if (on_gfx1x()) { /* gfx1100/gfx1150/GFX12, wave32 */ \
1290	WVSPLIT_TILE_CFG(/THRDS=/32, /WVPRGRP=/16, _sYT, __N) \
1291	} else { /* GFX9, wave64 */ \
1292	WVSPLIT_TILE_CFG(/THRDS=/64, /WVPRGRP=/16, _sYT, __N) \
1293	} \
1294	}
1295
1296	AT_DISPATCH_REDUCED_FLOATING_TYPES(in_b.scalar_type(), "wvSplitK", [&] {
1297	using fptype = typename scalar<scalar_t>::type;
1298	fptype* af4 = reinterpret_cast<fptype*>(in_a.data_ptr());
1299	const fptype* bf4 = reinterpret_cast<const fptype*>(in_b.data_ptr());
1300	const fptype* biasf4 =
1301	(in_bias.has_value() && in_bias->numel() > 0)
1302	? reinterpret_cast<const fptype*>(in_bias->data_ptr())
1303	: nullptr;
1304	fptype* c = reinterpret_cast<fptype*>(out_c.data_ptr());
1305
1306	// first shoot for biggest tile-size that keeps all simd busy,
1307	// then cut the active waves to balance their distribution...
1308	int sYT = (M_in + CuCount * 4 - 1) / (CuCount * 4);
1309
1310	switch (N_in) {
1311	case 1:
1312	WVSPLIT_TILE(sYT, 1)
1313	break;
1314	case 2:
1315	WVSPLIT_TILE(sYT, 2)
1316	break;
1317	case 3:
1318	WVSPLIT_TILE(sYT, 3)
1319	break;
1320	case 4:
1321	WVSPLIT_TILE(sYT, 4)
1322	break;
1323	case 5:
1324	WVSPLIT_TILE(sYT, 5)
1325	break;
1326	default:
1327	throw std::runtime_error(
1328	"Unsupported N value: " + std::to_string(M_in) + "," +
1329	std::to_string(K_in) + "," + std::to_string(N_in));
1330	}
1331	});
1332	return out_c;
1333	}
1334
1335	// This version targets cases skinny where CUs are not filled
1336	// Wave-SplitK is used with reduction done via atomics.
1337	#if defined(__gfx950__)
1338	#define WVSPLITKRC_1KPASS
1339	template <typename scalar_t, int THRDS, int YTILE, int WvPrGrp, int A_CHUNK,
1340	int UNRL, int N, int GrpsShrB, int CHUNKK, int DTRMNSTC>
1341	__global__ void __launch_bounds__(WvPrGrp* THRDS)
1342	__attribute__((amdgpu_waves_per_eu(1, 1)))
1343	wvSplitKrc_(const int actlN, const int K, const int Kap, const int M,
1344	const int Bx, const int By, const scalar_t* __restrict__ A,
1345	const scalar_t* __restrict__ B,
1346	const scalar_t* __restrict__ BIAS, float* glbl, int* cntr,
1347	scalar_t* C, const int CuCount) {
1348	constexpr int NTILE = 16;
1349	constexpr int APAD = 1;
1350	constexpr int ASTRD = 64;
1351	constexpr int BPAD = 1;
1352	constexpr int WVLDS_ = THRDS * A_CHUNK / CHUNKK;
1353	constexpr int WVLDS = ((WVLDS_ + A_CHUNK * BPAD)) * YTILE;
1354
1355	constexpr int max_lds_len = LDS_SIZE / 2;
1356
1357	using scalar16 =
1358	__attribute__((__vector_size__((A_CHUNK * 2) * sizeof(float)))) float;
1359	using scalar8 =
1360	__attribute__((__vector_size__((A_CHUNK / 2) * sizeof(float)))) float;
1361	using half4 =
1362	__attribute__((__vector_size__((A_CHUNK / 2) * sizeof(__bf16)))) __bf16;
1363	union bigType {
1364	scalar_t h[A_CHUNK];
1365	float f[A_CHUNK / 2];
1366	unsigned int i[A_CHUNK / 2];
1367	float2 f2[A_CHUNK / 4];
1368	unsigned long l[A_CHUNK / 4];
1369	double d[A_CHUNK / 4];
1370	half4 h4[A_CHUNK / 4];
1371	scalar8 h8;
1372	};
1373	using big4 = __attribute__((__vector_size__(4 * sizeof(bigType)))) __bf16;
1374
1375	// The main-loop staging areas are separated from the split-K readback by a
1376	// __syncthreads(), so the readback may reuse all of LDS.
1377	__shared__ union {
1378	struct {
1379	scalar_t stg[WvPrGrp * WVLDS / GrpsShrB];
1380	scalar_t s[max_lds_len - WvPrGrp * WVLDS / GrpsShrB];
1381	} mainloop;
1382	scalar_t all[max_lds_len];
1383	} lds;
1384	auto& stg = lds.mainloop.stg;
1385	auto& s = lds.mainloop.s;
1386	unsigned int* myStg = (unsigned int)(&stg[WVLDS (threadIdx.y / GrpsShrB)]);
1387
1388	#ifndef WVSPLITKRC_1KPASS
1389	constexpr int TUC_ = (THRDS * UNRL * A_CHUNK);
1390	// find biggest k size that fits padded into LDS
1391	constexpr uint32_t kFit__ = (max_lds_len - WvPrGrp * WVLDS / GrpsShrB) / N;
1392	constexpr uint32_t kFit_ = (kFit__ * ASTRD) / (APAD + ASTRD);
1393	uint32_t kFit = kFit_ - (kFit_ % TUC_);
1394	uint32_t kfitsPerRdc = (K + kFit - 1) / kFit;
1395
1396	// find best k split to fill the CUs
1397	if (((K + kfitsPerRdc * kFit - 1) / (kfitsPerRdc * kFit)) * numCuWithFullK <=
1398	CuCount)
1399	while (true) {
1400	while (kFit > TUC_) {
1401	uint32_t kFit_ = kFit - TUC_;
1402	if (((K + (kfitsPerRdc * kFit_ - 1)) / (kfitsPerRdc * kFit_)) *
1403	numCuWithFullK >
1404	CuCount)
1405	break;
1406	kFit = kFit_;
1407	}
1408	if (((K + ((kfitsPerRdc - 1) * kFit - 1)) / ((kfitsPerRdc - 1) * kFit)) *
1409	numCuWithFullK <=
1410	CuCount)
1411	kfitsPerRdc--;
1412	else
1413	break;
1414	}
1415	#else
1416	int constexpr kFit = 512 / CHUNKK;
1417	int constexpr kfitsPerRdc = 1;
1418	#endif
1419
1420	bool doRdc = true; // Assuming (kfitsPerRdc * kFit < K) is always true
1421	uint32_t numCuWithFullK =
1422	((M + (WvPrGrp * YTILE / GrpsShrB) - 1) / (WvPrGrp * YTILE / GrpsShrB));
1423	uint32_t Mmod = numCuWithFullK * (WvPrGrp * YTILE / GrpsShrB);
1424
1425	// given above k-split, find this wave's position
1426	uint32_t kFitPdd = kFit * CHUNKK + ((kFit * CHUNKK) / ASTRD) * APAD;
1427	uint32_t m0 = (blockIdx.x * WvPrGrp / GrpsShrB) * YTILE;
1428	uint32_t m1 = ((threadIdx.y % WvPrGrp) / GrpsShrB) * YTILE;
1429	uint32_t m = (m0 + m1) % Mmod;
1430	const uint32_t k_str = (m0 / Mmod) * kFit * kfitsPerRdc;
1431	uint32_t k_end = (m0 / Mmod + 1) * kFit * kfitsPerRdc;
1432	const uint32_t k_rnd = (K + kFit * kfitsPerRdc - 1) / (kFit * kfitsPerRdc);
1433
1434	scalar8 sum4[N / NTILE / GrpsShrB][1] = {0};
1435	bigType bigB_[YTILE / GrpsShrB / CHUNKK][UNRL];
1436	const uint32_t bLoader = (threadIdx.y % GrpsShrB);
1437	uint32_t kBase = 0;
1438	if (k_str >= K) return;
1439	if (m >= Mmod) return;
1440
1441	bool noreloada = false;
1442	constexpr bool FAST_UNSAFE_RDC_INIT = false;
1443
1444	#ifdef WVSPLITKRC_1KPASS
1445	// Early glbl init, B[] loading, if 1KPASS
1446	if constexpr (FAST_UNSAFE_RDC_INIT) {
1447	if (m + (threadIdx.x % 16) < M)
1448	if (doRdc)
1449	if (k_str == 0) {
1450	int mindx = m + (threadIdx.x % 16);
1451	int nindx_ = (0 + (threadIdx.x / 16) * 4) + 0 * NTILE +
1452	(N / GrpsShrB) * (threadIdx.y % GrpsShrB);
1453	int adr_ = mindx + M * nindx_ / 4;
1454	__hip_atomic_store(&cntr[adr_], 0, __ATOMIC_RELAXED,
1455	__HIP_MEMORY_SCOPE_AGENT);
1456	for (uint32_t nt = 0; nt < N / NTILE / GrpsShrB; nt++) {
1457	for (uint32_t j = 0; j < 4; j++) {
1458	int nindx = (j + (threadIdx.x / 16) * 4) + nt * NTILE +
1459	(N / GrpsShrB) * (threadIdx.y % GrpsShrB);
1460	int adr = mindx + M * nindx;
1461	__hip_atomic_store(&glbl[adr], 0, __ATOMIC_RELAXED,
1462	__HIP_MEMORY_SCOPE_AGENT);
1463	}
1464	}
1465	}
1466	}
1467
1468	// Load first B[] chunk
1469	#pragma unroll
1470	for (uint32_t k2 = 0; k2 < UNRL; k2++) {
1471	uint32_t k = k_str + k2 * THRDS * A_CHUNK;
1472	uint32_t k_ = k + (threadIdx.x % (THRDS / CHUNKK)) * A_CHUNK;
1473	const scalar_t* B_ = &B[min__(k_, K - A_CHUNK)];
1474	#pragma unroll
1475	for (uint32_t y = 0; y < YTILE / GrpsShrB; y += CHUNKK)
1476	bigB_[y / CHUNKK][k2].h8 = (loadnt(
1477	(scalar8)(&B_[min__((y + threadIdx.x / (THRDS / CHUNKK)) GrpsShrB +
1478	bLoader + m,
1479	M - 1) *
1480	K])));
1481	}
1482	{
1483	#else
1484	while (m < Mmod) {
1485	#endif
1486
1487	#ifndef WVSPLITKRC_1KPASS
1488	if constexpr (FAST_UNSAFE_RDC_INIT) {
1489	if (m + (threadIdx.x % 16) < M)
1490	if (doRdc)
1491	if (k_str == 0) {
1492	int mindx = m + (threadIdx.x % 16);
1493	int nindx_ = (0 + (threadIdx.x / 16) * 4) + 0 * NTILE +
1494	(N / GrpsShrB) * (threadIdx.y % GrpsShrB);
1495	int adr_ = mindx + M * nindx_ / 4;
1496	__hip_atomic_store(&cntr[adr_], 0, __ATOMIC_RELAXED,
1497	__HIP_MEMORY_SCOPE_AGENT);
1498	for (uint32_t nt = 0; nt < N / NTILE / GrpsShrB; nt++) {
1499	for (uint32_t j = 0; j < 4; j++) {
1500	int nindx = (j + (threadIdx.x / 16) * 4) + nt * NTILE +
1501	(N / GrpsShrB) * (threadIdx.y % GrpsShrB);
1502	int adr = mindx + M * nindx;
1503	__hip_atomic_store(&glbl[adr], 0, __ATOMIC_RELAXED,
1504	__HIP_MEMORY_SCOPE_AGENT);
1505	}
1506	}
1507	}
1508	}
1509
1510	#endif
1511
1512	#ifndef WVSPLITKRC_1KPASS
1513	for (uint32_t k1 = k_str; k1 < k_end; k1 += THRDS * A_CHUNK * UNRL) {
1514	#else
1515	const uint32_t k1 = k_str;
1516	{
1517	#endif
1518	#ifndef WVSPLITKRC_1KPASS
1519	const bool reloada = (!noreloada) &&
1520	((k1 == k_str) \|\| (k1 == k_str + kBase + kFit)) &&
1521	(k1 < k_end);
1522	// load next chunk of A[] to LDS
1523	if (reloada) {
1524	if (k1 != k_str) kBase += kFit;
1525	__syncthreads();
1526	#else
1527	const bool reloada = (!noreloada) &&
1528	((k1 == k_str) \|\| (k1 == k_str + kBase + kFit)) &&
1529	(k1 < k_end);
1530	if (reloada) {
1531	#endif
1532	constexpr int sprdN = 4;
1533	const uint32_t thrd = threadIdx.x % (THRDS / CHUNKK);
1534
1535	#ifndef WVSPLITKRC_1KPASS
1536	#pragma unroll
1537	for (int k = 0; k < kFit;
1538	k += (THRDS * (WvPrGrp / sprdN) * A_CHUNK) / CHUNKK) {
1539	#else
1540	const unsigned int k = 0;
1541	{
1542	#endif
1543	unsigned int kOff = k + (thrd * A_CHUNK);
1544	unsigned int kOffcp = min__(K - A_CHUNK, k_str + kOff);
1545	for (unsigned int n = 0; n < N; n += CHUNKK * sprdN) {
1546	__builtin_amdgcn_global_load_lds(
1547	(int)(&A[min__(Kap actlN - A_CHUNK,
1548	kOffcp + Kap * (n / CHUNKK +
1549	(N / CHUNKK) * (threadIdx.x /
1550	(64 / CHUNKK)) +
1551	(threadIdx.y % sprdN)))]),
1552	(int*)(&s[(k +
1553	kFitPdd * ((n / CHUNKK) + (threadIdx.y % sprdN)))]),
1554	16, 0, 0);
1555	}
1556
1557	// Stage loaded B[] to LDS for MFMA swizzling...
1558	for (uint32_t k2 = 0; k2 < UNRL; k2++) {
1559	uint32_t k = k1 + k2 * THRDS * A_CHUNK;
1560	uint32_t k_ = k + (threadIdx.x % (THRDS / CHUNKK)) * A_CHUNK;
1561	const bool oob_k = (k_ >= K);
1562	for (uint32_t y = 0; y < YTILE / GrpsShrB; y += CHUNKK) {
1563	uint32_t idx =
1564	(threadIdx.x % (THRDS / CHUNKK)) * 4 +
1565	((y + threadIdx.x / (THRDS / CHUNKK)) * GrpsShrB + bLoader) *
1566	((THRDS / CHUNKK + BPAD) * 4);
1567	// zero out if oob
1568	((scalar8)&myStg[idx]) =
1569	(oob_k) // TODO: ever necessary (y*GrpsShrB+bLoader+m>=M) ?
1570	? 0
1571	: bigB_[y / CHUNKK][k2].h8;
1572	}
1573	}
1574	}
1575	}
1576	}
1577	#ifndef WVSPLITKRC_1KPASS
1578	// Fire load of next B[] chunk...
1579	if ((k1 + THRDS * A_CHUNK * UNRL < k_end) &&
1580	(k1 + THRDS * A_CHUNK * UNRL < K))
1581	#pragma unroll
1582	for (uint32_t k2 = 0; k2 < UNRL; k2++) {
1583	uint32_t k = k1 + THRDS * A_CHUNK * UNRL + k2 * THRDS * A_CHUNK;
1584	uint32_t k_ = k + threadIdx.x * A_CHUNK;
1585	const scalar_t* B_ = &B[min__(k_, K - A_CHUNK)];
1586	#pragma unroll
1587	for (uint32_t y = 0; y < YTILE / GrpsShrB; y += CHUNKK)
1588	bigB_[y / CHUNKK][k2].h8 = (loadnt(
1589	(scalar8)(&B_[min__((y + threadIdx.x / (THRDS / CHUNKK))
1590	GrpsShrB +
1591	bLoader + m,
1592	M - 1) *
1593	K])));
1594	}
1595	#endif
1596
1597	// B[] staging is cooperative across GrpsShrB, so sync here before reading
1598	// back. This wait is currently inserted by compiler, but not guaranteed.
1599	asm volatile("s_waitcnt 0");
1600	__syncthreads();
1601
1602	// read back B[] swizzled for MFMA...
1603	bigType bigB[YTILE / CHUNKK][UNRL];
1604	for (uint32_t k2 = 0; k2 < UNRL; k2++) {
1605	for (uint32_t y = 0; y < YTILE / CHUNKK; y++) {
1606	unsigned int idx =
1607	(threadIdx.x % YTILE) * ((THRDS / CHUNKK + BPAD) * 4) +
1608	(threadIdx.x / YTILE) * 4 + y * 16;
1609	bigB[y][k2].h8 = ((scalar8)&myStg[idx]);
1610	}
1611	}
1612
1613	// rReadback A[] swizzled for MFMA...
1614	bigType bigA[N / GrpsShrB / CHUNKK][UNRL];
1615	#pragma unroll
1616	for (uint32_t k2 = 0; k2 < UNRL; k2++) {
1617	uint32_t k = k1 + k2 * THRDS * A_CHUNK - kBase - k_str;
1618	#pragma unroll
1619	for (uint32_t nt = 0; nt < N / GrpsShrB; nt += NTILE)
1620	#pragma unroll
1621	for (uint32_t n = 0; n < NTILE / CHUNKK; n++) {
1622	uint32_t idxa =
1623	((nt + (N / GrpsShrB) * (threadIdx.y % GrpsShrB)) % (N / CHUNKK) +
1624	(threadIdx.x % NTILE)) *
1625	kFitPdd +
1626	((nt + (N / GrpsShrB) * (threadIdx.y % GrpsShrB)) /
1627	(N / CHUNKK)) *
1628	A_CHUNK * (64 / CHUNKK) +
1629	A_CHUNK * ((threadIdx.x / NTILE) + n * 4) + k;
1630	bigA[nt / CHUNKK + n][k2] = ((const bigType)(&(s[idxa])));
1631	}
1632	}
1633
1634	// Do the MFMAs
1635	#pragma unroll
1636	for (uint32_t k2 = 0; k2 < UNRL; k2++) {
1637	#pragma unroll
1638	for (uint32_t nt = 0; nt < N / NTILE / GrpsShrB; nt++) {
1639	#pragma unroll
1640	for (uint32_t j = 0; j < YTILE / CHUNKK; j++) {
1641	if constexpr (std::is_same_v<scalar_t, half>) {
1642	sum4[nt][0] = __builtin_amdgcn_mfma_f32_16x16x32_f16(
1643	bigA[nt * (YTILE / CHUNKK) + j][k2].h8, bigB[j][k2].h8,
1644	sum4[nt][0], 0, 0, 0);
1645	} else { // bf16
1646	sum4[nt][0] = __builtin_amdgcn_mfma_f32_16x16x32_bf16(
1647	bigA[nt * (YTILE / CHUNKK) + j][k2].h8, bigB[j][k2].h8,
1648	sum4[nt][0], 0, 0, 0);
1649	}
1650	}
1651	}
1652	}
1653	}
1654
1655	union flt4 {
1656	scalar8 s8;
1657	float2 f2[2];
1658	float4 f4;
1659	};
1660	if (m + (threadIdx.x % 16) < M) {
1661	int my_cntr;
1662	int mindx = m + (threadIdx.x % 16);
1663	int g_mindx = m * 4 + (threadIdx.x % 64); // coalesced atomic reduction
1664	scalar_t biases[N / NTILE / GrpsShrB][4] = {};
1665	// Atomic add the output, read biases
1666	for (uint32_t nt = 0; nt < N / NTILE / GrpsShrB; nt++) {
1667	int g_nindx =
1668	(nt * NTILE + (N / GrpsShrB) * (threadIdx.y % GrpsShrB)) / 4;
1669	int g_adr = g_mindx * 4 + 0 + M * g_nindx * 4;
1670	if (DTRMNSTC) {
1671	flt4 flt4_ = {.s8 = sum4[nt][0]};
1672	__hip_atomic_store((float2)&glbl[g_adr + M N * (m0 / Mmod)],
1673	flt4_.f2[0], __ATOMIC_RELAXED,
1674	__HIP_MEMORY_SCOPE_AGENT);
1675	__hip_atomic_store((float2)&glbl[g_adr + 2 + M N * (m0 / Mmod)],
1676	flt4_.f2[1], __ATOMIC_RELAXED,
1677	__HIP_MEMORY_SCOPE_AGENT);
1678	} else {
1679	for (uint32_t j = 0; j < 4; j++)
1680	atomicAdd((&glbl[g_adr + j]), sum4[nt][0][j]);
1681	}
1682	}
1683
1684	__atomic_signal_fence(__ATOMIC_SEQ_CST);
1685	asm volatile("s_waitcnt vmcnt(0)" ::: "memory");
1686	__atomic_signal_fence(__ATOMIC_SEQ_CST);
1687
1688	int nindx_ = (0 + (threadIdx.x / 16) * 4) + 0 * NTILE +
1689	(N / GrpsShrB) * (threadIdx.y % GrpsShrB);
1690	int adr_ = mindx + M * nindx_ / 4;
1691	my_cntr = atomicAdd(&cntr[adr_], 1);
1692
1693	// make sure LDS is free for write out staging
1694	if (DTRMNSTC) __syncthreads();
1695
1696	// Update the complete counter
1697	flt4 vals[N / NTILE / GrpsShrB] = {};
1698	// If we're the last k-shard, read back the value and convert...
1699	if (my_cntr + 1 == k_rnd) {
1700	cntr[adr_] = 0; // clear for next round
1701	if constexpr (DTRMNSTC) {
1702	constexpr uint32_t NT = N / NTILE / GrpsShrB;
1703	// The readback stages k-shard partials through LDS, which bounds how
1704	// many shards may be in flight at once.
1705	constexpr uint32_t LDS_F4 =
1706	max_lds_len * sizeof(scalar_t) / sizeof(float4);
1707	constexpr uint32_t KSBMAX = LDS_F4 / (THRDS * 4 * NT);
1708	static_assert(KSBMAX >= 1, "LDS cannot stage even one k-shard");
1709	const uint32_t KSB = (k_rnd < KSBMAX) ? k_rnd : KSBMAX;
1710
1711	// First batch is peeled; it is the only batch for most shapes.
1712	#pragma unroll
1713	for (uint32_t kb = 0; kb < KSB; kb++) {
1714	for (uint32_t nt = 0; nt < NT; nt++) {
1715	int g_nindx =
1716	(nt * NTILE + (N / GrpsShrB) * (threadIdx.y % GrpsShrB)) / 4;
1717	int g_adr = g_mindx * 4 + 0 + M * g_nindx * 4;
1718	__builtin_amdgcn_global_load_lds(
1719	(float4)(&glbl[g_adr + M N * kb]),
1720	&(((float4)lds.all)[(threadIdx.y THRDS) + kb * THRDS * 4 +
1721	nt * THRDS * 4 * KSB]),
1722	16, 0, 0);
1723	}
1724	}
1725	if (BIAS) // overlaps the DMAs above
1726	for (uint32_t nt = 0; nt < NT; nt++) {
1727	for (uint32_t j = 0; j < 4; j++) {
1728	int nindx = (j + (threadIdx.x / 16) * 4) + nt * NTILE +
1729	(N / GrpsShrB) * (threadIdx.y % GrpsShrB);
1730	biases[nt][j] = BIAS[(mindx % Bx) + (nindx % By) * Bx];
1731	}
1732	}
1733	asm volatile("s_waitcnt vmcnt(0)" ::: "memory");
1734	for (uint32_t kb = 0; kb < KSB; kb++) {
1735	for (uint32_t nt = 0; nt < NT; nt++) {
1736	vals[nt].f4 +=
1737	((float4)lds.all)[(threadIdx.x + threadIdx.y THRDS) +
1738	kb * THRDS * 4 + nt * THRDS * 4 * KSB];
1739	}
1740	}
1741
1742	for (uint32_t ks0 = KSB; ks0 < k_rnd; ks0 += KSB) {
1743	const uint32_t ksn = (k_rnd - ks0 < KSB) ? (k_rnd - ks0) : KSB;
1744	// prior reads must retire before the staging is overwritten
1745	asm volatile("s_waitcnt lgkmcnt(0)" ::: "memory");
1746	for (uint32_t kb = 0; kb < ksn; kb++) {
1747	for (uint32_t nt = 0; nt < NT; nt++) {
1748	int g_nindx =
1749	(nt * NTILE + (N / GrpsShrB) * (threadIdx.y % GrpsShrB)) / 4;
1750	int g_adr = g_mindx * 4 + 0 + M * g_nindx * 4;
1751	__builtin_amdgcn_global_load_lds(
1752	(float4)(&glbl[g_adr + M N * (ks0 + kb)]),
1753	&(((float4)lds.all)[(threadIdx.y THRDS) + kb * THRDS * 4 +
1754	nt * THRDS * 4 * KSB]),
1755	16, 0, 0);
1756	}
1757	}
1758	asm volatile("s_waitcnt vmcnt(0)" ::: "memory");
1759	for (uint32_t kb = 0; kb < ksn; kb++) {
1760	for (uint32_t nt = 0; nt < NT; nt++) {
1761	vals[nt].f4 +=
1762	((float4)lds.all)[(threadIdx.x + threadIdx.y THRDS) +
1763	kb * THRDS * 4 + nt * THRDS * 4 * KSB];
1764	}
1765	}
1766	}
1767	} else {
1768	for (uint32_t nt = 0; nt < N / NTILE / GrpsShrB; nt++) {
1769	int g_nindx =
1770	(nt * NTILE + (N / GrpsShrB) * (threadIdx.y % GrpsShrB)) / 4;
1771	int g_adr = g_mindx * 4 + 0 + M * g_nindx * 4;
1772	vals[nt].f4 = (float4)(&glbl[g_adr]);
1773	(float4)(&glbl[g_adr]) = {}; // clear out for next round
1774	}
1775	if (BIAS)
1776	for (uint32_t nt = 0; nt < N / NTILE / GrpsShrB; nt++) {
1777	for (uint32_t j = 0; j < 4; j++) {
1778	int nindx = (j + (threadIdx.x / 16) * 4) + nt * NTILE +
1779	(N / GrpsShrB) * (threadIdx.y % GrpsShrB);
1780	biases[nt][j] = BIAS[(mindx % Bx) + (nindx % By) * Bx];
1781	}
1782	}
1783	}
1784	__builtin_amdgcn_sched_barrier(0);
1785	for (uint32_t nt = 0; nt < N / NTILE / GrpsShrB; nt++) {
1786	for (uint32_t j = 0; j < 4; j++) {
1787	int nindx = (j + (threadIdx.x / 16) * 4) + nt * NTILE +
1788	(N / GrpsShrB) * (threadIdx.y % GrpsShrB);
1789	if (nindx < actlN) {
1790	int adr = mindx + M * nindx;
1791	if constexpr (std::is_same_v<scalar_t, __hip_bfloat16>) {
1792	vals[nt].s8[j] += __bfloat162float(biases[nt][j]);
1793	C[adr] = __float2bfloat16(vals[nt].s8[j]);
1794	} else {
1795	vals[nt].s8[j] += __half2float(biases[nt][j]);
1796	C[adr] = __float2half(vals[nt].s8[j]);
1797	}
1798	}
1799	}
1800	}
1801	}
1802
1803	#ifndef WVSPLITKRC_1KPASS
1804	m0 += CuCount * WvPrGrp * YTILE / GrpsShrB;
1805	m = (m0 + m1) % Mmod;
1806	k_str = (m0 / Mmod) * kFit * kfitsPerRdc;
1807	k_end = (m0 / Mmod + 1) * kFit * kfitsPerRdc;
1808	if (k_str >= K) break;
1809	kBase = 0;
1810	#endif
1811	}
1812	}
1813	#else
1814	template <typename scalar_t, int THRDS, int YTILE, int WvPrGrp, int A_CHUNK,
1815	int UNRL, int N, int GrpsShrB, int CHUNKK, int DTRMNSTC>
1816	__global__ void wvSplitKrc_(const int actlN, const int K, const int Kap,
1817	const int M, const int Bx, const int By,
1818	const scalar_t* B, const scalar_t* __restrict__ A,
1819	const scalar_t* __restrict__ BIAS, float* glbl,
1820	int* cntr, scalar_t* C,
1821	const int CuCount){UNREACHABLE_CODE}
1822	#endif // defined(__HIP__GFX9__) TODO: Add NAVI support
1823
1824	torch::Tensor wvSplitKrc(const at::Tensor& in_a, const at::Tensor& in_b,
1825	const std::optional<at::Tensor>& in_bias,
1826	const int64_t CuCount) {
1827	static constexpr bool _DTRMNSTC = true;
1828
1829	auto M_in = in_b.size(0);
1830	auto N_in = in_a.size(0);
1831	auto K_in = in_b.size(1);
1832	auto Kap_in = in_a.stride(0);
1833
1834	auto Bx_in =
1835	(in_bias.has_value() && in_bias->numel() > 0)
1836	? (in_bias->sizes().size() == 2) ? in_bias->size(1) : in_bias->size(0)
1837	: 1;
1838	auto By_in = (in_bias.has_value() && in_bias->numel() > 0 &&
1839	in_bias->sizes().size() == 2)
1840	? in_bias->size(0)
1841	: 1;
1842
1843	TORCH_CHECK(in_a.dtype() == in_b.dtype());
1844	TORCH_CHECK(K_in % 8 == 0, "k % 8 == 0");
1845	TORCH_CHECK(in_a.dtype() == torch::kFloat16 \|\|
1846	in_a.dtype() == torch::kBFloat16);
1847
1848	const at::cuda::OptionalCUDAGuard device_guard(device_of(in_a));
1849
1850	auto out_c = torch::empty(
1851	{N_in, M_in},
1852	torch::TensorOptions().dtype(in_a.dtype()).device(in_a.device()));
1853
1854	auto N_p2 = 1U << (32 - __builtin_clz(N_in - 1));
1855
1856	dim3 grid(CuCount);
1857
1858	const cudaStream_t stream = at::cuda::getCurrentCUDAStream();
1859	// const int max_lds_len = get_lds_size() / 2;
1860
1861	// With 64 Ms per CU (each of 4 SIMDs working on a 16x16 tile),
1862	// and each working on a 512-shard of K, how many CUs would we need?
1863	int rndup_cus = ((M_in + 64 - 1) / 64) * ((K_in + 512 - 1) / 512);
1864
1865	// How many of 4 waves in a group can work on same 16 Ms at same time? First
1866	// try to maximize this. This reduces the Ms each group works on, i.e.
1867	// increasing the number of CUs needed.
1868	int GrpsShrB = min(N_p2 / 16, 4);
1869
1870	// Given the above, how many CUs would we need?
1871	int CuNeeded = rndup_cus * GrpsShrB;
1872
1873	if (CuNeeded > CuCount) throw std::runtime_error("Invalid wvSplitKrc size");
1874
1875	// Can we increase SplitK by shrinking the K-shard to 256? That doubles both
1876	// the CUs needed and the k-shards to read back, and the readback is serial in
1877	// the shard count, so it pays only up to a crossover. The crossover is not a
1878	// constant -- it falls as M grows -- so one cap cannot sit on it for every
1879	// shape; this value minimizes total time over the admitted shapes. Retune it
1880	// if the readback cost changes, and note the directions are not symmetric:
1881	// capping too high costs about twice what capping too low does.
1882	// Mirrored by rocm_unquantized_gemm_impl() in
1883	// vllm/model_executor/layers/utils.py; both must pick the same chunkk or the
1884	// host-side fit check bounds the wrong k_rnd.
1885	constexpr int64_t CHUNKK2_MAX_SHARDS = 11;
1886	const int64_t shards_chunkk2 = (K_in + 256 - 1) / 256; // 256 == 512 / 2
1887	const int chunkk = (N_p2 != 16 && CuNeeded * 2 <= CuCount &&
1888	shards_chunkk2 <= CHUNKK2_MAX_SHARDS)
1889	? 2
1890	: 1;
1891
1892	// One fp32 partial per (M, N, k-shard) must fit the split-K workspace.
1893	const int64_t k_rnd = (K_in + 512 / chunkk - 1) / (512 / chunkk);
1894	constexpr int64_t glbl_cap = 128 * 1024 * (_DTRMNSTC ? 12 : 1);
1895	const int64_t glbl_needed =
1896	_DTRMNSTC ? (int64_t)N_p2 * M_in * k_rnd : (int64_t)N_p2 * M_in;
1897	TORCH_CHECK(glbl_needed <= glbl_cap,
1898	"wvSplitKrc: split-K workspace too small", " (need ", glbl_needed,
1899	" > ", glbl_cap, " floats) for N=", N_in, " K=", K_in,
1900	" M=", M_in);
1901
1902	// Split-K workspaces: one pre-partitioned pool per device.
1903	// These must be zeroed before first use, and the kernel re-zeros them after
1904	// each use.
1905	// To avoid unnecessary overhead, we do not allocate them zeroed
1906	// per-invocation. However, a simple static allocation does not work with
1907	// multi-streams, because concurrent streams can alias into the same
1908	// workspace. Solution: pre-allocate kWvSlots = 8 slots per device, plus
1909	// kWvCaptureSlots reserved for capturing streams, and zero the whole pool
1910	// once at creation. On first use, a stream takes one of the slots if any are
1911	// still available. On subsequent uses, it re-uses the slot. If no slots are
1912	// available, fall back to allocating a zeroed one per call, which is safe but
1913	// adds overhead.
1914	// A capturing stream must neither allocate nor enqueue anything -- either
1915	// would be captured into the graph -- so it is only handed a reserved slot.
1916	struct WvSplitKrcPool {
1917	torch::Tensor glbl, cntr;
1918	int64_t glbl_stride = 0, cntr_stride = 0;
1919	int next_slot = 0;
1920	int next_capture_slot = 0;
1921	std::map<cudaStream_t, int> slot_of;
1922	std::map<cudaStream_t, int> capture_slot_of;
1923	};
1924	// mutex not strictly necessary if caller is always single-threaded (Python)
1925	// but this preserves Torch thread_local semantics for the C++ API
1926	static std::mutex wv_pool_mu;
1927	static std::map<int, WvSplitKrcPool> wv_pools;
1928
1929	float* glbl;
1930	int* cntr;
1931	// Fallback space
1932	torch::Tensor spill_glbl, spill_cntr;
1933	{
1934	std::lock_guard<std::mutex> wv_lk(wv_pool_mu);
1935	const int wv_dev = static_cast<int>(in_a.device().index());
1936	WvSplitKrcPool& P = wv_pools[wv_dev];
1937	constexpr int64_t wv_gf = 128 * 1024 * (_DTRMNSTC ? 12 : 1);
1938	constexpr int64_t wv_ci = wv_gf / 4;
1939	constexpr int64_t wv_all_slots = kWvSlots + kWvCaptureSlots;
1940
1941	// A stream whose capture was already invalidated reports an error; treat
1942	// that as not capturing, since the capture is lost either way.
1943	cudaStreamCaptureStatus capture_status = cudaStreamCaptureStatusNone;
1944	unsigned long long capture_id = 0;
1945	if (cudaStreamGetCaptureInfo(stream, &capture_status, &capture_id) !=
1946	cudaSuccess) {
1947	capture_status = cudaStreamCaptureStatusNone;
1948	}
1949	const bool wv_capturing = capture_status != cudaStreamCaptureStatusNone;
1950
1951	if (!P.glbl.defined()) {
1952	TORCH_CHECK(!wv_capturing,
1953	"wvSplitKrc: the split-K workspace pool cannot be allocated "
1954	"during graph capture. Call "
1955	"warmup_rocm_skinny_gemm_workspaces() "
1956	"(vllm/model_executor/layers/utils.py) eagerly first.");
1957	P.glbl_stride = wv_gf;
1958	P.cntr_stride = wv_ci;
1959	P.glbl = torch::zeros(wv_gf * wv_all_slots, torch::TensorOptions()
1960	.dtype(torch::kFloat32)
1961	.device(in_a.device()))
1962	.detach();
1963	P.cntr =
1964	torch::zeros(
1965	wv_ci * wv_all_slots,
1966	torch::TensorOptions().dtype(torch::kInt).device(in_a.device()))
1967	.detach();
1968	// The mutex orders host access to the map, not device execution, so it
1969	// alone would let another stream be handed a slot the fill above is
1970	// still clearing. Synchronizing here, inside the lock, is what makes
1971	// every slot safe to use without any further per-slot zeroing.
1972	const cudaError_t wv_err = cudaStreamSynchronize(stream);
1973	TORCH_CHECK(wv_err == cudaSuccess,
1974	"wvSplitKrc: failed to zero the split-K workspace pool: ",
1975	cudaGetErrorString(wv_err));
1976	}
1977
1978	int wv_slot = -1;
1979	if (wv_capturing) {
1980	// One workspace per capturing stream, so graphs captured on the same
1981	// stream share it -- safe because at most one of them replays at a time
1982	// (vLLM captures one graph per padded batch size, and runs one per step).
1983	auto cap_it = P.capture_slot_of.find(stream);
1984	if (cap_it == P.capture_slot_of.end()) {
1985	TORCH_CHECK(P.next_capture_slot < kWvCaptureSlots,
1986	"wvSplitKrc: no free capture workspace slot; more than ",
1987	kWvCaptureSlots,
1988	" distinct streams have captured a graph containing "
1989	"wvSplitKrc. Raise kWvCaptureSlots in "
1990	"csrc/rocm/skinny_gemms.cu.");
1991	cap_it = P.capture_slot_of
1992	.emplace(stream, static_cast<int>(kWvSlots) +
1993	P.next_capture_slot++)
1994	.first;
1995	}
1996	wv_slot = cap_it->second;
1997	} else {
1998	auto slot_it = P.slot_of.find(stream);
1999	if (slot_it != P.slot_of.end()) {
2000	wv_slot = slot_it->second;
2001	} else if (P.next_slot < kWvSlots) {
2002	wv_slot = P.next_slot++;
2003	P.slot_of.emplace(stream, wv_slot);
2004	}
2005	}
2006
2007	if (wv_slot >= 0) {
2008	glbl = P.glbl.data_ptr<float>() + wv_slot * P.glbl_stride;
2009	cntr = P.cntr.data_ptr<int>() + wv_slot * P.cntr_stride;
2010	} else {
2011	spill_glbl = torch::zeros(
2012	wv_gf,
2013	torch::TensorOptions().dtype(torch::kFloat32).device(in_a.device()));
2014	spill_cntr = torch::zeros(
2015	wv_ci,
2016	torch::TensorOptions().dtype(torch::kInt).device(in_a.device()));
2017	glbl = spill_glbl.data_ptr<float>();
2018	cntr = spill_cntr.data_ptr<int>();
2019	}
2020	}
2021
2022	#define WVSPLITKrc(_N, _GrpsShrB, _CHUNKK) \
2023	{ \
2024	dim3 block(64, 4); \
2025	if constexpr (_DTRMNSTC) \
2026	wvSplitKrc_<fptype, 64, 16, 4, 8, 1, _N, _GrpsShrB, _CHUNKK, 1> \
2027	<<<grid, block, 0, stream>>>(N_in, K_in, Kap_in, M_in, Bx_in, By_in, \
2028	af4, bf4, biasf4, glbl, cntr, c, \
2029	CuCount); \
2030	else \
2031	wvSplitKrc_<fptype, 64, 16, 4, 8, 1, _N, _GrpsShrB, _CHUNKK, 0> \
2032	<<<grid, block, 0, stream>>>(N_in, K_in, Kap_in, M_in, Bx_in, By_in, \
2033	af4, bf4, biasf4, glbl, cntr, c, \
2034	CuCount); \
2035	}
2036
2037	AT_DISPATCH_REDUCED_FLOATING_TYPES(in_a.scalar_type(), "wvSplitKrc", [&] {
2038	using fptype = typename scalar<scalar_t>::type;
2039	const fptype* af4 = reinterpret_cast<const fptype*>(in_a.data_ptr());
2040	const fptype* bf4 = reinterpret_cast<const fptype*>(in_b.data_ptr());
2041	const fptype* biasf4 =
2042	(in_bias.has_value() && in_bias->numel() > 0)
2043	? reinterpret_cast<const fptype*>(in_bias->data_ptr())
2044	: nullptr;
2045	fptype* c = reinterpret_cast<fptype*>(out_c.data_ptr());
2046
2047	switch (N_p2) {
2048	case 16:
2049	WVSPLITKrc(16, 1, 1) break;
2050	case 32:
2051	if (chunkk == 2) WVSPLITKrc(32, 2, 2) else WVSPLITKrc(32, 2, 1) break;
2052	case 64:
2053	if (chunkk == 2) WVSPLITKrc(64, 4, 2) else WVSPLITKrc(64, 4, 1) break;
2054	case 128:
2055	if (chunkk == 2) WVSPLITKrc(128, 4, 2) else WVSPLITKrc(128, 4, 1) break;
2056	default:
2057	throw std::runtime_error(
2058	"Unsupported N value: " + std::to_string(M_in) + "," +
2059	std::to_string(K_in) + "," + std::to_string(N_in));
2060	}
2061	});
2062	return out_c;
2063	}
2064
2065	#if defined(__HIP__MI3XX__) \|\| defined(__GFX12__)
2066	template <typename scalar_t, typename fp8_t, int THRDS, int YTILE, int WvPrGrp,
2067	int A_CHUNK, int UNRL, int N>
2068	__global__ void __launch_bounds__(WvPrGrp* THRDS)
2069	wvSplitKQ_hf_sml_(const int K, const int Kap, const int Kbp, const int M,
2070	const int Bx, const int By, const fp8_t* B,
2071	const fp8_t* __restrict__ A,
2072	const scalar_t* __restrict__ BIAS, scalar_t* C,
2073	const float* __restrict__ s_A,
2074	const float* __restrict__ s_B, const int _WvPrGrp,
2075	const int CuCount) {
2076	constexpr int max_lds_len = LDS_SIZE;
2077	using scalar8 =
2078	__attribute__((__vector_size__((A_CHUNK / 4) * sizeof(float)))) float;
2079	using intx2 = __attribute__((__vector_size__(2 * sizeof(int)))) int;
2080	using intx4 = __attribute__((__vector_size__(4 * sizeof(int)))) int;
2081	union bigType {
2082	char f8[A_CHUNK];
2083	char2 c2[A_CHUNK / 2];
2084	scalar_t h[A_CHUNK / 2];
2085	float f[A_CHUNK / 4];
2086	int i[A_CHUNK / 4];
2087	long l[A_CHUNK / 8];
2088	intx4 l2[A_CHUNK / 16];
2089	scalar8 h8;
2090	};
2091
2092	__shared__ fp8_t s[max_lds_len];
2093
2094	for (uint32_t k = (threadIdx.y * THRDS + threadIdx.x) * A_CHUNK;
2095	k < min__(Kap * N, max_lds_len); k += THRDS * WvPrGrp * A_CHUNK) {
2096	#if defined(__gfx950__)
2097	__builtin_amdgcn_global_load_lds((int)(&A[k]), (int)(&s[k]), 16, 0, 0);
2098	#else
2099	((bigType)(&s[k])) = ((bigType)(&A[k]));
2100	#endif
2101	}
2102	asm volatile("s_waitcnt vmcnt(0)");
2103	__syncthreads();
2104
2105	if (threadIdx.y >= _WvPrGrp) return;
2106
2107	uint32_t m = (blockIdx.x * _WvPrGrp + (threadIdx.y % _WvPrGrp)) * YTILE;
2108
2109	float sA = *s_A;
2110	float sB = *s_B;
2111
2112	while (m < M) {
2113	#ifdef __GFX12__
2114	// gfx12: per-lane scalar accumulation via v_dot4_f32_fp8_fp8
2115	float sum[N][YTILE] = {};
2116	#else
2117	// gfx9: MFMA accumulation
2118	scalar8 sum[N][YTILE] = {};
2119	#endif
2120	for (uint32_t k1 = 0; k1 < K; k1 += THRDS * A_CHUNK * UNRL) {
2121	bigType bigA[N][UNRL] = {};
2122	bigType bigB[YTILE][UNRL];
2123
2124	// Fetch the weight matrix from memory!
2125	#pragma unroll
2126	for (uint32_t k2 = 0; k2 < UNRL; k2++) {
2127	uint32_t k = k1 + k2 * THRDS * A_CHUNK;
2128	uint32_t k_ = k + threadIdx.x * A_CHUNK;
2129	const fp8_t* B_ = &B[min__(k_, K - A_CHUNK)];
2130	#pragma unroll
2131	for (uint32_t y = 0; y < YTILE; ++y) {
2132	bigB[y][k2].h8 = (loadnt((scalar8)(&B_[min__(y + m, M - 1) Kbp])));
2133	}
2134	}
2135
2136	// Fetch activation matrix from either just LDS or from both LDS / memory
2137	#pragma unroll
2138	for (uint32_t k2 = 0; k2 < UNRL; k2++) {
2139	uint32_t k = k1 + k2 * THRDS * A_CHUNK;
2140	uint32_t k_ = k + threadIdx.x * A_CHUNK;
2141	if (k_ >= K) break;
2142	for (int n = 0; n < N; n++) {
2143	bigA[n][k2] = ((const bigType)(&(s[k_ + Kap * n])));
2144	}
2145	}
2146
2147	// Do the matrix multiplication in interleaved manner
2148	#pragma unroll
2149	for (uint32_t k2 = 0; k2 < UNRL; k2++) {
2150	for (uint32_t n = 0; n < N; n++) {
2151	#ifdef __GFX12__
2152	// gfx12: 4 x dot4 per A_CHUNK=16 bytes (4 FP8 per dot4)
2153	for (int y = 0; y < YTILE; ++y) {
2154	#pragma unroll
2155	for (int i = 0; i < A_CHUNK / 4; i++) {
2156	sum[n][y] = __builtin_amdgcn_dot4_f32_fp8_fp8(
2157	bigA[n][k2].i[i], bigB[y][k2].i[i], sum[n][y]);
2158	}
2159	}
2160	#else
2161	// gfx9: MFMA path
2162	for (int i = 0; i < A_CHUNK; i += 8) {
2163	for (int y = 0; y < YTILE; ++y) {
2164	sum[n][y] = __builtin_amdgcn_mfma_f32_16x16x32_fp8_fp8(
2165	bigA[n][k2].l[i / 8], bigB[y][k2].l[i / 8], sum[n][y], 0, 0,
2166	0);
2167	}
2168	}
2169	#endif
2170	}
2171	}
2172	}
2173
2174	// Final reduction
2175	#ifdef __GFX12__
2176	// gfx12 wave32: DPP row_shr within 16-lane rows + cross-row shuffle
2177	for (int n = 0; n < N; n++) {
2178	for (int y = 0; y < YTILE; y++) {
2179	asm("s_nop 0\n\tv_add_f32 %0, %2, %3 row_shr:8 bound_ctrl:0 "
2180	: "=v"(sum[n][y])
2181	: "0"(sum[n][y]), "v"(sum[n][y]), "v"(sum[n][y]));
2182	asm("s_nop 0\n\tv_add_f32 %0, %2, %3 row_shr:4 bound_ctrl:0 "
2183	: "=v"(sum[n][y])
2184	: "0"(sum[n][y]), "v"(sum[n][y]), "v"(sum[n][y]));
2185	asm("s_nop 0\n\tv_add_f32 %0, %2, %3 row_shr:2 bound_ctrl:0 "
2186	: "=v"(sum[n][y])
2187	: "0"(sum[n][y]), "v"(sum[n][y]), "v"(sum[n][y]));
2188	asm("s_nop 0\n\tv_add_f32 %0, %2, %3 row_shr:1 bound_ctrl:0 "
2189	: "=v"(sum[n][y])
2190	: "0"(sum[n][y]), "v"(sum[n][y]), "v"(sum[n][y]));
2191	sum[n][y] += __shfl_xor(sum[n][y], 16);
2192	}
2193	}
2194	#else
2195	// gfx9 MFMA reduction
2196	for (int n = 0; n < N; n++) {
2197	for (int y = 0; y < YTILE; y++) {
2198	float accm0 = sum[n][y][0];
2199	accm0 += __builtin_amdgcn_mov_dpp(sum[n][y][1], 0x101, 0xf, 0xf,
2200	1); // row_shl1
2201	accm0 += __builtin_amdgcn_mov_dpp(sum[n][y][2], 0x102, 0xf, 0xf,
2202	1); // row_shl2
2203	accm0 += __builtin_amdgcn_mov_dpp(sum[n][y][3], 0x103, 0xf, 0xf,
2204	1); // row_shl3
2205	accm0 += __shfl_down(accm0, 20);
2206	accm0 += __shfl_down(accm0, 40);
2207	sum[n][y][0] = accm0;
2208	}
2209	}
2210	#endif
2211
2212	const bool writeback_lane =
2213	#ifdef __GFX12__
2214	threadIdx.x == (THRDS - 1);
2215	#else
2216	threadIdx.x == 0;
2217	#endif
2218	if (writeback_lane) {
2219	scalar_t biases[N][YTILE] = {};
2220	if (BIAS)
2221	for (int n = 0; n < N; n++) {
2222	for (int y = 0; y < YTILE; y++) {
2223	biases[n][y] = BIAS[(m + y) % Bx + (n % By) * Bx];
2224	}
2225	}
2226	for (int n = 0; n < N; n++) {
2227	for (int y = 0; y < YTILE; y++) {
2228	if (y + m >= M) break; // To avoid mem access fault.
2229	#ifdef __GFX12__
2230	float result = sum[n][y] * sA * sB;
2231	#else
2232	float result = sum[n][y][0] * sA * sB;
2233	#endif
2234	if constexpr (std::is_same_v<scalar_t, half>) {
2235	result += __half2float(biases[n][y]);
2236	} else if constexpr (std::is_same_v<scalar_t, __hip_bfloat16>) {
2237	result += __bfloat162float(biases[n][y]);
2238	}
2239	C[m + y + n * M] = __float2s<scalar_t>(result);
2240	}
2241	}
2242	}
2243
2244	m += CuCount * _WvPrGrp * YTILE;
2245	}
2246	}
2247	#else // !defined(__HIP__MI3XX__) && !defined(__GFX12__)
2248	template <typename scalar_t, typename fp8_t, int THRDS, int YTILE, int WvPrGrp,
2249	int A_CHUNK, int UNRL, int N>
2250	__global__ void wvSplitKQ_hf_sml_(const int K, const int Kap, const int Kbp,
2251	const int M, const int Bx, const int By,
2252	const fp8_t* B, const fp8_t* __restrict__ A,
2253	const scalar_t* __restrict__ BIAS,
2254	scalar_t* C, const float* __restrict__ s_A,
2255	const float* __restrict__ s_B,
2256	const int _WvPrGrp, const int CuCount) {
2257	UNREACHABLE_CODE
2258	}
2259	#endif // defined(__HIP__MI3XX__) \|\| defined(__GFX12__)
2260
2261	#if defined(__HIP__MI3XX__) \|\| defined(__GFX12__)
2262	template <typename scalar_t, typename fp8_t, int THRDS, int YTILE, int WvPrGrp,
2263	int A_CHUNK, int UNRL, int N>
2264	__global__ void __launch_bounds__(WvPrGrp* THRDS)
2265	wvSplitKQ_hf_(const int K, const int Kap, const int Kbp, const int M,
2266	const int Bx, const int By, const fp8_t* B,
2267	const fp8_t* __restrict__ A,
2268	const scalar_t* __restrict__ BIAS, scalar_t* C,
2269	const float* __restrict__ s_A, const float* __restrict__ s_B,
2270	const int _WvPrGrp, const int CuCount) {
2271	constexpr int max_lds_len = LDS_SIZE;
2272	using scalar8 =
2273	__attribute__((__vector_size__((A_CHUNK / 4) * sizeof(float)))) float;
2274	using intx2 = __attribute__((__vector_size__(2 * sizeof(int)))) int;
2275	using intx4 = __attribute__((__vector_size__(4 * sizeof(int)))) int;
2276	union bigType {
2277	char f8[A_CHUNK];
2278	char2 c2[A_CHUNK / 2];
2279	scalar_t h[A_CHUNK / 2];
2280	float f[A_CHUNK / 4];
2281	int i[A_CHUNK / 4];
2282	long l[A_CHUNK / 8];
2283	intx4 l2[A_CHUNK / 16];
2284	scalar8 h8;
2285	};
2286
2287	__shared__ fp8_t s[max_lds_len];
2288
2289	for (uint32_t k = (threadIdx.y * THRDS + threadIdx.x) * A_CHUNK;
2290	k < min__(Kap * N, max_lds_len); k += THRDS * WvPrGrp * A_CHUNK) {
2291	#if defined(__gfx950__)
2292	__builtin_amdgcn_global_load_lds((int)(&A[k]), (int)(&s[k]), 16, 0, 0);
2293	#else
2294	((bigType)(&s[k])) = ((bigType)(&A[k]));
2295	#endif
2296	}
2297	asm volatile("s_waitcnt vmcnt(0)");
2298	__syncthreads();
2299
2300	if (threadIdx.y >= _WvPrGrp) return;
2301
2302	uint32_t m = (blockIdx.x * _WvPrGrp + (threadIdx.y % _WvPrGrp)) * YTILE;
2303
2304	float sA = *s_A;
2305	float sB = *s_B;
2306
2307	while (m < M) {
2308	#ifdef __GFX12__
2309	// gfx12: per-lane scalar accumulation via v_dot4_f32_fp8_fp8
2310	float sum[N][YTILE] = {};
2311	#else
2312	// gfx9: MFMA accumulation
2313	scalar8 sum[N][YTILE] = {};
2314	#endif
2315	for (uint32_t k1 = 0; k1 < K; k1 += THRDS * A_CHUNK * UNRL) {
2316	bigType bigA[N][UNRL] = {};
2317	bigType bigB[YTILE][UNRL];
2318
2319	// Fetch the weight matrix from memory!
2320	#pragma unroll
2321	for (uint32_t k2 = 0; k2 < UNRL; k2++) {
2322	uint32_t k = k1 + k2 * THRDS * A_CHUNK;
2323	uint32_t k_ = k + threadIdx.x * A_CHUNK;
2324	const fp8_t* B_ = &B[min__(k_, K - A_CHUNK)];
2325	for (int y = 0; y < YTILE; ++y) {
2326	bigB[y][k2].h8 = (loadnt((scalar8)(&B_[min__(y + m, M - 1) Kbp])));
2327	}
2328	}
2329
2330	// Fetch activation matrix from either just LDS or from both LDS / memory
2331	#pragma unroll
2332	for (uint32_t k2 = 0; k2 < UNRL; k2++) {
2333	uint32_t k = k1 + k2 * THRDS * A_CHUNK;
2334	uint32_t k_ = k + threadIdx.x * A_CHUNK;
2335	if (k_ >= K) break;
2336	for (int n = 0; n < N; n++) {
2337	if (k_ + Kap * n < max_lds_len)
2338	bigA[n][k2] = ((const bigType)(&(s[k_ + Kap * n])));
2339	else
2340	bigA[n][k2] = ((const bigType)(&(A[k_ + Kap * n])));
2341	}
2342	}
2343
2344	// Do the matrix multiplication in interleaved manner
2345	#pragma unroll
2346	for (uint32_t k2 = 0; k2 < UNRL; k2++) {
2347	for (uint32_t n = 0; n < N; n++) {
2348	#ifdef __GFX12__
2349	// gfx12: 4 x dot4 per A_CHUNK=16 bytes (4 FP8 per dot4)
2350	for (int y = 0; y < YTILE; ++y) {
2351	#pragma unroll
2352	for (int i = 0; i < A_CHUNK / 4; i++) {
2353	sum[n][y] = __builtin_amdgcn_dot4_f32_fp8_fp8(
2354	bigA[n][k2].i[i], bigB[y][k2].i[i], sum[n][y]);
2355	}
2356	}
2357	#else
2358	// gfx9: MFMA path
2359	for (int i = 0; i < A_CHUNK; i += 8) {
2360	for (int y = 0; y < YTILE; ++y) {
2361	sum[n][y] = __builtin_amdgcn_mfma_f32_16x16x32_fp8_fp8(
2362	bigA[n][k2].l[i / 8], bigB[y][k2].l[i / 8], sum[n][y], 0, 0,
2363	0);
2364	}
2365	}
2366	#endif
2367	}
2368	}
2369	}
2370
2371	// Final reduction
2372	#ifdef __GFX12__
2373	// gfx12 wave32: DPP row_shr within 16-lane rows + cross-row shuffle
2374	for (int n = 0; n < N; n++) {
2375	for (int y = 0; y < YTILE; y++) {
2376	asm("s_nop 0\n\tv_add_f32 %0, %2, %3 row_shr:8 bound_ctrl:0 "
2377	: "=v"(sum[n][y])
2378	: "0"(sum[n][y]), "v"(sum[n][y]), "v"(sum[n][y]));
2379	asm("s_nop 0\n\tv_add_f32 %0, %2, %3 row_shr:4 bound_ctrl:0 "
2380	: "=v"(sum[n][y])
2381	: "0"(sum[n][y]), "v"(sum[n][y]), "v"(sum[n][y]));
2382	asm("s_nop 0\n\tv_add_f32 %0, %2, %3 row_shr:2 bound_ctrl:0 "
2383	: "=v"(sum[n][y])
2384	: "0"(sum[n][y]), "v"(sum[n][y]), "v"(sum[n][y]));
2385	asm("s_nop 0\n\tv_add_f32 %0, %2, %3 row_shr:1 bound_ctrl:0 "
2386	: "=v"(sum[n][y])
2387	: "0"(sum[n][y]), "v"(sum[n][y]), "v"(sum[n][y]));
2388	sum[n][y] += __shfl_xor(sum[n][y], 16);
2389	}
2390	}
2391	#else
2392	// gfx9 MFMA reduction
2393	for (int n = 0; n < N; n++) {
2394	for (int y = 0; y < YTILE; y++) {
2395	float accm0 = sum[n][y][0];
2396	accm0 += __builtin_amdgcn_mov_dpp(sum[n][y][1], 0x101, 0xf, 0xf,
2397	1); // row_shl1
2398	accm0 += __builtin_amdgcn_mov_dpp(sum[n][y][2], 0x102, 0xf, 0xf,
2399	1); // row_shl2
2400	accm0 += __builtin_amdgcn_mov_dpp(sum[n][y][3], 0x103, 0xf, 0xf,
2401	1); // row_shl3
2402	accm0 += __shfl_down(accm0, 20);
2403	accm0 += __shfl_down(accm0, 40);
2404	sum[n][y][0] = accm0;
2405	}
2406	}
2407	#endif
2408
2409	const bool writeback_lane =
2410	#ifdef __GFX12__
2411	threadIdx.x == (THRDS - 1);
2412	#else
2413	threadIdx.x == 0;
2414	#endif
2415	if (writeback_lane) {
2416	scalar_t biases[N][YTILE] = {};
2417	if (BIAS)
2418	for (int n = 0; n < N; n++) {
2419	for (int y = 0; y < YTILE; y++) {
2420	biases[n][y] = BIAS[(m + y) % Bx + (n % By) * Bx];
2421	}
2422	}
2423	for (int n = 0; n < N; n++) {
2424	for (int y = 0; y < YTILE; y++) {
2425	if (y + m >= M) break; // To avoid mem access fault.
2426	#ifdef __GFX12__
2427	float result = sum[n][y] * sA * sB;
2428	#else
2429	float result = sum[n][y][0] * sA * sB;
2430	#endif
2431	if constexpr (std::is_same_v<scalar_t, half>) {
2432	result += __half2float(biases[n][y]);
2433	} else if constexpr (std::is_same_v<scalar_t, __hip_bfloat16>) {
2434	result += __bfloat162float(biases[n][y]);
2435	}
2436	C[m + y + n * M] = __float2s<scalar_t>(result);
2437	}
2438	}
2439	}
2440
2441	m += CuCount * _WvPrGrp * YTILE;
2442	}
2443	}
2444	#else // !defined(__HIP__MI3XX__) && !defined(__GFX12__)
2445	template <typename scalar_t, typename fp8_t, int THRDS, int YTILE, int WvPrGrp,
2446	int A_CHUNK, int UNRL, int N>
2447	__global__ void wvSplitKQ_hf_(const int K, const int Kap, const int Kbp,
2448	const int M, const int Bx, const int By,
2449	const fp8_t* B, const fp8_t* __restrict__ A,
2450	const scalar_t* __restrict__ BIAS, scalar_t* C,
2451	const float* __restrict__ s_A,
2452	const float* __restrict__ s_B, const int _WvPrGrp,
2453	const int CuCount) {
2454	UNREACHABLE_CODE
2455	}
2456	#endif // defined(__HIP__MI3XX__) \|\| defined(__GFX12__)
2457
2458	void wvSplitKQ(const at::Tensor& in_b, const at::Tensor& in_a,
2459	const std::optional<at::Tensor>& in_bias, at::Tensor& out_c,
2460	const at::Tensor& scale_a, const at::Tensor& scale_b,
2461	const int64_t CuCount) {
2462	static c10::ScalarType kFp8Type = is_fp8_ocp()
2463	? c10::ScalarType::Float8_e4m3fn
2464	: c10::ScalarType::Float8_e4m3fnuz;
2465	auto M_in = in_b.size(0);
2466	auto K_in = in_b.size(1);
2467	auto N_in = in_a.size(0);
2468	auto Kap_in = in_a.stride(0);
2469	auto Kbp_in = in_b.stride(0);
2470	auto Bx_in =
2471	(in_bias.has_value() && in_bias->numel() > 0)
2472	? (in_bias->sizes().size() == 2) ? in_bias->size(1) : in_bias->size(0)
2473	: 1;
2474	auto By_in = (in_bias.has_value() && in_bias->numel() > 0 &&
2475	in_bias->sizes().size() == 2)
2476	? in_bias->size(0)
2477	: 1;
2478
2479	TORCH_CHECK(K_in % 16 == 0, "k % 16 == 0");
2480	TORCH_CHECK(in_a.dtype() == in_b.dtype() && in_a.dtype() == kFp8Type);
2481	TORCH_CHECK(out_c.dtype() == torch::kFloat16 \|\|
2482	out_c.dtype() == torch::kBFloat16);
2483
2484	dim3 grid(CuCount);
2485	const at::cuda::OptionalCUDAGuard device_guard(device_of(in_a));
2486	const cudaStream_t stream = at::cuda::getCurrentCUDAStream();
2487	const int max_lds_len = get_lds_size();
2488
2489	#define WVSPLITKQ_IMPL(_THRDS, _WvPrGrp, _YTILEs, _YTILEm, _UNRLs, _UNRLm, _N) \
2490	{ \
2491	dim3 block(_THRDS, _WvPrGrp); \
2492	if ((Kap_in * N_in <= max_lds_len) && (M_in % _YTILEs == 0)) { \
2493	int __wvPrGrp = min(_WvPrGrp, mindiv(M_in, CuCount * _YTILEs, 16)); \
2494	wvSplitKQ_hf_sml_<fptype, fp8_t, _THRDS, _YTILEs, _WvPrGrp, 16, _UNRLs, \
2495	_N><<<grid, block, 0, stream>>>( \
2496	K_in, Kap_in, Kbp_in, M_in, Bx_in, By_in, b_ptr, a_ptr, bias_ptr, \
2497	c_ptr, s_a, s_b, __wvPrGrp, CuCount); \
2498	} else { \
2499	int __wvPrGrp = min(_WvPrGrp, mindiv(M_in, CuCount * _YTILEm, 16)); \
2500	wvSplitKQ_hf_<fptype, fp8_t, _THRDS, _YTILEm, _WvPrGrp, 16, _UNRLm, _N> \
2501	<<<grid, block, 0, stream>>>(K_in, Kap_in, Kbp_in, M_in, Bx_in, \
2502	By_in, b_ptr, a_ptr, bias_ptr, c_ptr, \
2503	s_a, s_b, __wvPrGrp, CuCount); \
2504	} \
2505	}
2506
2507	#define WVSPLITKQ(_WvPrGrp, _YTILEs, _YTILEm, _UNRLs, _UNRLm, _N) \
2508	if (on_gfx12()) \
2509	WVSPLITKQ_IMPL(32, _WvPrGrp, _YTILEs, _YTILEm, _UNRLs, _UNRLm, _N) \
2510	else \
2511	WVSPLITKQ_IMPL(64, _WvPrGrp, _YTILEs, _YTILEm, _UNRLs, _UNRLm, _N)
2512
2513	AT_DISPATCH_REDUCED_FLOATING_TYPES(out_c.scalar_type(), "wvSplitKQ", [&] {
2514	using fptype = typename scalar<scalar_t>::type;
2515	auto c_ptr = reinterpret_cast<fptype*>(out_c.data_ptr());
2516	auto s_a = scale_a.data_ptr<float>();
2517	auto s_b = scale_b.data_ptr<float>();
2518	VLLM_DISPATCH_FP8_TYPES(in_a.scalar_type(), "wvSplitKQ", [&] {
2519	auto a_ptr = in_a.data_ptr<fp8_t>();
2520	auto b_ptr = in_b.data_ptr<fp8_t>();
2521	auto bias_ptr = (in_bias.has_value() && in_bias->numel() > 0)
2522	? reinterpret_cast<fptype*>(in_bias->data_ptr())
2523	: nullptr;
2524	switch (N_in) {
2525	case 1:
2526	WVSPLITKQ(16, 2, 2, 2, 2, 1)
2527	break;
2528	case 2:
2529	WVSPLITKQ(16, 2, 2, 2, 2, 2)
2530	break;
2531	case 3:
2532	WVSPLITKQ(16, 2, 2, 1, 1, 3)
2533	break;
2534	case 4:
2535	WVSPLITKQ(16, 2, 2, 1, 1, 4)
2536	break;
2537	default:
2538	throw std::runtime_error(
2539	"Unsupported N value: " + std::to_string(M_in) + "," +
2540	std::to_string(K_in) + "," + std::to_string(N_in));
2541	}
2542	});
2543	});
2544	}