| 1 | // SPDX-License-Identifier: Apache-2.0 |
| 2 | // SPDX-FileCopyrightText: Copyright contributors to the vLLM project |
| 3 | // |
| 4 | // Fused MoE W4A16 GPTQ kernel for RDNA3 (gfx1100). |
| 5 | // |
| 6 | // Combines expert routing (sorted_token_ids / expert_ids) with the RDNA3 |
| 7 | // W4A16 dequant+dot from q_gemm_rdna3.cu into a single kernel launch. |
| 8 | // Each block processes BLOCK_SIZE_M tokens assigned to one expert, covering |
| 9 | // a tile of N output columns and K input positions. |
| 10 | // |
| 11 | // Weight format: same as the dense kernel — [E, K/8, N] uint32 shuffled, |
| 12 | // [E, groups, N] scales, [E, groups, N/8] packed zeros. |
| 13 | // |
| 14 | // Design: THREADS_X=256 (8 waves on wave32), BLOCK_KN_SIZE=256, each thread |
| 15 | // handles 4 N columns. Output via 64-bit packed CAS atomic-add directly to |
| 16 | // the pre-zeroed output tensor (no FP32 scratch buffer). |
| 17 | |
| 18 | #include <cstdint> |
| 19 | |
| 20 | #include <torch/all.h> |
| 21 | #include <c10/cuda/CUDAGuard.h> |
| 22 | #include <ATen/cuda/CUDAContext.h> |
| 23 | |
| 24 | #include <hip/hip_runtime.h> |
| 25 | #include <hip/hip_bf16.h> |
| 26 | #include <hip/hip_fp16.h> |
| 27 | |
| 28 | #include "qdq_4_rdna3.cuh" |
| 29 | |
| 30 | #if defined(__HIPCC__) && defined(__gfx1100__) |
| 31 | #define __HIP__RDNA3__ |
| 32 | #endif |
| 33 | |
| 34 | namespace vllm { |
| 35 | namespace moe_gptq_rdna3 { |
| 36 | |
| 37 | #define BLOCK_KN_SIZE 256 |
| 38 | #define THREADS_X 256 |
| 39 | |
| 40 | #if defined(__HIP__RDNA3__) || !defined(__HIP_DEVICE_COMPILE__) |
| 41 | |
| 42 | using gptq_rdna3::bf162_t; |
| 43 | using gptq_rdna3::bf16_t; |
| 44 | |
| 45 | // --- Helpers (same as q_gemm_rdna3.cu) --- |
| 46 | |
| 47 | template <typename T> |
| 48 | __forceinline__ __device__ T tzero(); |
| 49 | |
| 50 | template <> |
| 51 | __forceinline__ __device__ half tzero<half>() { |
| 52 | return __float2half_rn(0.0f); |
| 53 | } |
| 54 | |
| 55 | template <> |
| 56 | __forceinline__ __device__ bf16_t tzero<bf16_t>() { |
| 57 | return __float2bfloat16(0.0f); |
| 58 | } |
| 59 | |
| 60 | __forceinline__ __device__ float dot22_8_f(half2 (&dq)[4], const half* a_ptr) { |
| 61 | float result = 0.0f; |
| 62 | const half2* a2_ptr = (const half2*)a_ptr; |
| 63 | #pragma unroll |
| 64 | for (int i = 0; i < 4; i++) { |
| 65 | result = __builtin_amdgcn_fdot2(dq[i], *a2_ptr++, result, /*clamp=*/false); |
| 66 | } |
| 67 | return result; |
| 68 | } |
| 69 | |
| 70 | __forceinline__ __device__ void atomic_add_pk4_f16(half* addr, half2 v01, |
| 71 | half2 v23) { |
| 72 | unsigned long long* addr_u = reinterpret_cast<unsigned long long*>(addr); |
| 73 | unsigned long long old = *addr_u; |
| 74 | while (true) { |
| 75 | union { |
| 76 | unsigned long long u; |
| 77 | half2 h2[2]; |
| 78 | } cur, sum; |
| 79 | cur.u = old; |
| 80 | sum.h2[0] = __hadd2(cur.h2[0], v01); |
| 81 | sum.h2[1] = __hadd2(cur.h2[1], v23); |
| 82 | unsigned long long prev = atomicCAS(addr_u, old, sum.u); |
| 83 | if (prev == old) break; |
| 84 | old = prev; |
| 85 | } |
| 86 | } |
| 87 | |
| 88 | __forceinline__ __device__ void atomic_add_pk4_bf16(bf16_t* addr, bf162_t v01, |
| 89 | bf162_t v23) { |
| 90 | unsigned long long* addr_u = reinterpret_cast<unsigned long long*>(addr); |
| 91 | unsigned long long old = *addr_u; |
| 92 | while (true) { |
| 93 | union { |
| 94 | unsigned long long u; |
| 95 | bf162_t b2[2]; |
| 96 | } cur, sum; |
| 97 | cur.u = old; |
| 98 | sum.b2[0] = __hadd2(cur.b2[0], v01); |
| 99 | sum.b2[1] = __hadd2(cur.b2[1], v23); |
| 100 | unsigned long long prev = atomicCAS(addr_u, old, sum.u); |
| 101 | if (prev == old) break; |
| 102 | old = prev; |
| 103 | } |
| 104 | } |
| 105 | |
| 106 | __forceinline__ __device__ void load4_zeros(const uint32_t* qzeros_row, int n, |
| 107 | int (&zeros)[4]) { |
| 108 | int qcol = n / 8; |
| 109 | int shift = (n & 0x07) * 4; |
| 110 | uint32_t d = qzeros_row[qcol] >> shift; |
| 111 | zeros[0] = (int)(d & 0xF); |
| 112 | zeros[1] = (int)((d >> 4) & 0xF); |
| 113 | zeros[2] = (int)((d >> 8) & 0xF); |
| 114 | zeros[3] = (int)((d >> 12) & 0xF); |
| 115 | } |
| 116 | |
| 117 | template <typename T> |
| 118 | __forceinline__ __device__ void load4_scales(const T* scales_row, int n, |
| 119 | T (&scales)[4]) { |
| 120 | scales[0] = scales_row[n + 0]; |
| 121 | scales[1] = scales_row[n + 1]; |
| 122 | scales[2] = scales_row[n + 2]; |
| 123 | scales[3] = scales_row[n + 3]; |
| 124 | } |
| 125 | |
| 126 | // --------------------------------------------------------------------------- |
| 127 | // Fused MoE kernel. |
| 128 | // --------------------------------------------------------------------------- |
| 129 | |
| 130 | template <typename T, int BLOCK_SIZE_M> |
| 131 | __global__ void moe_gemm_q4_kernel_rdna3( |
| 132 | const T* __restrict__ a, // [size_m, size_k] or [M*topk, K] |
| 133 | T* __restrict__ c, // [M*topk, size_n] pre-zeroed |
| 134 | const uint32_t* __restrict__ b_q_weight, // [E, K/8, N] packed |
| 135 | const T* __restrict__ b_scales, // [E, groups, N] |
| 136 | const uint32_t* __restrict__ b_qzeros, // [E, groups, N/8] packed |
| 137 | const float* __restrict__ topk_weights, // [M*topk] or nullptr |
| 138 | const int32_t* __restrict__ sorted_token_ids, |
| 139 | const int32_t* __restrict__ expert_ids, |
| 140 | const int32_t* __restrict__ num_tokens_post_padded, |
| 141 | const int size_m, // total tokens (original M, or M*topk for w2) |
| 142 | const int size_n, // output features per expert |
| 143 | const int size_k, // input features |
| 144 | const int groups, // K / group_size |
| 145 | const int top_k, // routing top-k (1 for w2 pass) |
| 146 | // Per-expert strides (in elements, not bytes) |
| 147 | const int expert_weight_stride, // (K/8) * N |
| 148 | const int expert_scales_stride, // groups * N |
| 149 | const int expert_zeros_stride, // groups * (N/8) |
| 150 | const bool mul_topk_weight, |
| 151 | const int output_topk) { // >0: reduce output by token_id/output_topk |
| 152 | const int t = threadIdx.x; |
| 153 | const int token_block = blockIdx.x; |
| 154 | const int offset_n = blockIdx.y * BLOCK_KN_SIZE * 4; |
| 155 | const int offset_k = blockIdx.z * BLOCK_KN_SIZE; |
| 156 | const int end_k = min(offset_k + BLOCK_KN_SIZE, size_k); |
| 157 | const int n = offset_n + t * 4; |
| 158 | |
| 159 | // Early exit for padding blocks or invalid experts (expert_map = -1) |
| 160 | if (token_block * BLOCK_SIZE_M >= num_tokens_post_padded[0]) return; |
| 161 | |
| 162 | const int expert_id = expert_ids[token_block]; |
| 163 | if (expert_id == -1) return; |
| 164 | |
| 165 | // Expert-specific pointers |
| 166 | const uint32_t* expert_weights = |
| 167 | b_q_weight + (int64_t)expert_id * expert_weight_stride; |
| 168 | const T* expert_scales = b_scales + (int64_t)expert_id * expert_scales_stride; |
| 169 | const uint32_t* expert_qzeros = |
| 170 | b_qzeros + (int64_t)expert_id * expert_zeros_stride; |
| 171 | |
| 172 | // LDS for activations |
| 173 | constexpr int LDS_PAD = 8; |
| 174 | __shared__ T block_a[BLOCK_SIZE_M][BLOCK_KN_SIZE + LDS_PAD]; |
| 175 | |
| 176 | static_assert(BLOCK_KN_SIZE == THREADS_X, |
| 177 | "BLOCK_KN_SIZE must equal THREADS_X"); |
| 178 | |
| 179 | // For bf16 M=1, we can skip LDS and read A from global (same as dense). |
| 180 | // fp16 always needs LDS due to the dot22_8_f indexing pattern. |
| 181 | constexpr bool USE_LDS_A = (BLOCK_SIZE_M > 1) || std::is_same<T, half>::value; |
| 182 | |
| 183 | const int offset_m_base = token_block * BLOCK_SIZE_M; |
| 184 | |
| 185 | if constexpr (USE_LDS_A) { |
| 186 | if (offset_k + t < end_k) { |
| 187 | #pragma unroll |
| 188 | for (int m = 0; m < BLOCK_SIZE_M; ++m) { |
| 189 | int32_t token_id = sorted_token_ids[offset_m_base + m]; |
| 190 | int token_row = token_id / top_k; |
| 191 | T av; |
| 192 | if (token_row < size_m) { |
| 193 | av = a[(int64_t)token_row * size_k + offset_k + t]; |
| 194 | } else { |
| 195 | av = tzero<T>(); |
| 196 | } |
| 197 | block_a[m][t] = av; |
| 198 | } |
| 199 | } |
| 200 | __syncthreads(); |
| 201 | } |
| 202 | |
| 203 | if (n >= size_n) return; |
| 204 | |
| 205 | // Group bookkeeping |
| 206 | const int groupsize = size_k / groups; |
| 207 | int group = offset_k / groupsize; |
| 208 | int nextgroup = (group + 1) * groupsize; |
| 209 | |
| 210 | // Weight pointer for this expert |
| 211 | int qk = offset_k / 8; |
| 212 | const uint32_t* b_ptr = expert_weights + qk * size_n + n; |
| 213 | |
| 214 | // Per-column dequant constants (4 columns per thread) |
| 215 | half2 z1z16_h[4][2], y1y16_h[4][2]; |
| 216 | float z_b_f[4], y_b_f[4]; |
| 217 | |
| 218 | // GPTQv1: zero_offset = 1 |
| 219 | constexpr int zero_offset = 1; |
| 220 | |
| 221 | auto refresh_group = [&](int g) { |
| 222 | const uint32_t* qz_row = expert_qzeros + g * (size_n / 8); |
| 223 | const T* sc_row = expert_scales + g * size_n; |
| 224 | int zeros[4]; |
| 225 | T scales[4]; |
| 226 | load4_zeros(qz_row, n, zeros); |
| 227 | load4_scales<T>(sc_row, n, scales); |
| 228 | if constexpr (std::is_same<T, half>::value) { |
| 229 | #pragma unroll |
| 230 | for (int i = 0; i < 4; ++i) { |
| 231 | gptq_rdna3::prep_zero_scale_fp16((uint32_t)(zeros[i] + zero_offset), |
| 232 | scales[i], z1z16_h[i], y1y16_h[i]); |
| 233 | } |
| 234 | } else { |
| 235 | #pragma unroll |
| 236 | for (int i = 0; i < 4; ++i) { |
| 237 | gptq_rdna3::prep_zero_scale_bf16_f32((uint32_t)(zeros[i] + zero_offset), |
| 238 | scales[i], z_b_f[i], y_b_f[i]); |
| 239 | } |
| 240 | } |
| 241 | }; |
| 242 | |
| 243 | refresh_group(group); |
| 244 | |
| 245 | float block_c[BLOCK_SIZE_M][4]; |
| 246 | #pragma unroll |
| 247 | for (int m = 0; m < BLOCK_SIZE_M; ++m) { |
| 248 | #pragma unroll |
| 249 | for (int j = 0; j < 4; ++j) block_c[m][j] = 0.0f; |
| 250 | } |
| 251 | |
| 252 | // --- Main K-loop --- |
| 253 | int k = offset_k; |
| 254 | while (k < end_k) { |
| 255 | if (k == nextgroup) { |
| 256 | group++; |
| 257 | nextgroup += groupsize; |
| 258 | refresh_group(group); |
| 259 | } |
| 260 | |
| 261 | // Prefetch 4 weight words (128 bytes) |
| 262 | int4 b_w[4]; |
| 263 | #pragma unroll |
| 264 | for (int j = 0; j < 4; ++j) { |
| 265 | b_w[j] = *(const int4*)(b_ptr + j * size_n); |
| 266 | } |
| 267 | b_ptr += 4 * size_n; |
| 268 | |
| 269 | #pragma unroll |
| 270 | for (int j = 0; j < 4; ++j) { |
| 271 | const int a_off = (k - offset_k) + 8 * j; |
| 272 | |
| 273 | if constexpr (std::is_same<T, half>::value) { |
| 274 | // fp16 path: dequant via bit-trick, dot via v_dot2_f32_f16 |
| 275 | half2 dq[4][4]; |
| 276 | gptq_rdna3::dequant_4bit_8_fp16((uint32_t)b_w[j].x, dq[0], z1z16_h[0], |
| 277 | y1y16_h[0]); |
| 278 | gptq_rdna3::dequant_4bit_8_fp16((uint32_t)b_w[j].y, dq[1], z1z16_h[1], |
| 279 | y1y16_h[1]); |
| 280 | gptq_rdna3::dequant_4bit_8_fp16((uint32_t)b_w[j].z, dq[2], z1z16_h[2], |
| 281 | y1y16_h[2]); |
| 282 | gptq_rdna3::dequant_4bit_8_fp16((uint32_t)b_w[j].w, dq[3], z1z16_h[3], |
| 283 | y1y16_h[3]); |
| 284 | |
| 285 | #pragma unroll |
| 286 | for (int m = 0; m < BLOCK_SIZE_M; ++m) { |
| 287 | const half* a_ptr = reinterpret_cast<const half*>(&block_a[m][a_off]); |
| 288 | block_c[m][0] += dot22_8_f(dq[0], a_ptr); |
| 289 | block_c[m][1] += dot22_8_f(dq[1], a_ptr); |
| 290 | block_c[m][2] += dot22_8_f(dq[2], a_ptr); |
| 291 | block_c[m][3] += dot22_8_f(dq[3], a_ptr); |
| 292 | } |
| 293 | } else if constexpr (BLOCK_SIZE_M == 1) { |
| 294 | // bf16 M=1: v_dot2_f32_bf16 with InstCombine-defeating opacity |
| 295 | typedef short __attribute__((ext_vector_type(2))) bf16x2_t; |
| 296 | constexpr uint32_t BF16_MAGIC = 0x43004300u; |
| 297 | constexpr uint32_t BF16_ONES = 0x3F803F80u; |
| 298 | union pack4 { |
| 299 | float f[4]; |
| 300 | uint32_t u[4]; |
| 301 | }; |
| 302 | |
| 303 | uint32_t w[4]; |
| 304 | __builtin_memcpy(w, &b_w[j], sizeof(int4)); |
| 305 | |
| 306 | // Load activations — read from global (no LDS for bf16 M=1) |
| 307 | pack4 a_pack; |
| 308 | { |
| 309 | int32_t token_id = sorted_token_ids[offset_m_base]; |
| 310 | int token_row = token_id / top_k; |
| 311 | if (token_row < size_m) { |
| 312 | const uint32_t* a_words = reinterpret_cast<const uint32_t*>( |
| 313 | a + (int64_t)token_row * size_k + offset_k + a_off); |
| 314 | a_pack.u[0] = a_words[0]; |
| 315 | a_pack.u[1] = a_words[1]; |
| 316 | a_pack.u[2] = a_words[2]; |
| 317 | a_pack.u[3] = a_words[3]; |
| 318 | } else { |
| 319 | a_pack.u[0] = 0; |
| 320 | a_pack.u[1] = 0; |
| 321 | a_pack.u[2] = 0; |
| 322 | a_pack.u[3] = 0; |
| 323 | } |
| 324 | } |
| 325 | |
| 326 | // sum_a for bias correction |
| 327 | float sum_a = 0.0f; |
| 328 | #pragma unroll |
| 329 | for (int b = 0; b < 4; ++b) { |
| 330 | sum_a = __builtin_amdgcn_fdot2_f32_bf16( |
| 331 | *((bf16x2_t*)(&a_pack.f[b])), *((const bf16x2_t*)&BF16_ONES), |
| 332 | sum_a, /*clamp=*/false); |
| 333 | } |
| 334 | |
| 335 | #pragma unroll 1 |
| 336 | for (int col = 0; col < 4; ++col) { |
| 337 | pack4 q_pack; |
| 338 | const uint32_t qa = w[col]; |
| 339 | q_pack.u[0] = ((qa >> 0) & 0x000F000Fu) | BF16_MAGIC; |
| 340 | q_pack.u[1] = ((qa >> 4) & 0x000F000Fu) | BF16_MAGIC; |
| 341 | q_pack.u[2] = ((qa >> 8) & 0x000F000Fu) | BF16_MAGIC; |
| 342 | q_pack.u[3] = ((qa >> 12) & 0x000F000Fu) | BF16_MAGIC; |
| 343 | |
| 344 | float partial = 0.0f; |
| 345 | #pragma unroll |
| 346 | for (int b = 0; b < 4; ++b) { |
| 347 | partial = __builtin_amdgcn_fdot2_f32_bf16( |
| 348 | *((bf16x2_t*)(&a_pack.f[b])), *((bf16x2_t*)(&q_pack.f[b])), |
| 349 | partial, /*clamp=*/false); |
| 350 | } |
| 351 | |
| 352 | block_c[0][col] = |
| 353 | __fmaf_rn(y_b_f[col], partial, |
| 354 | __fmaf_rn(z_b_f[col], sum_a, block_c[0][col])); |
| 355 | } |
| 356 | } else { |
| 357 | // bf16 M>1: v_dot2_f32_bf16 with LDS-staged activations |
| 358 | typedef short __attribute__((ext_vector_type(2))) bf16x2_t; |
| 359 | constexpr uint32_t BF16_MAGIC = 0x43004300u; |
| 360 | constexpr uint32_t BF16_ONES = 0x3F803F80u; |
| 361 | union pack4 { |
| 362 | float f[4]; |
| 363 | uint32_t u[4]; |
| 364 | }; |
| 365 | |
| 366 | uint32_t w[4]; |
| 367 | __builtin_memcpy(w, &b_w[j], sizeof(int4)); |
| 368 | |
| 369 | pack4 a_pack[BLOCK_SIZE_M]; |
| 370 | #pragma unroll |
| 371 | for (int m = 0; m < BLOCK_SIZE_M; ++m) { |
| 372 | const uint32_t* a_words = |
| 373 | reinterpret_cast<const uint32_t*>(&block_a[m][a_off]); |
| 374 | a_pack[m].u[0] = a_words[0]; |
| 375 | a_pack[m].u[1] = a_words[1]; |
| 376 | a_pack[m].u[2] = a_words[2]; |
| 377 | a_pack[m].u[3] = a_words[3]; |
| 378 | } |
| 379 | |
| 380 | float sum_a[BLOCK_SIZE_M]; |
| 381 | #pragma unroll |
| 382 | for (int m = 0; m < BLOCK_SIZE_M; ++m) { |
| 383 | float s = 0.0f; |
| 384 | #pragma unroll |
| 385 | for (int b = 0; b < 4; ++b) { |
| 386 | s = __builtin_amdgcn_fdot2_f32_bf16(*((bf16x2_t*)(&a_pack[m].f[b])), |
| 387 | *((const bf16x2_t*)&BF16_ONES), |
| 388 | s, /*clamp=*/false); |
| 389 | } |
| 390 | sum_a[m] = s; |
| 391 | } |
| 392 | |
| 393 | #pragma unroll 1 |
| 394 | for (int col = 0; col < 4; ++col) { |
| 395 | pack4 q_pack; |
| 396 | const uint32_t qa = w[col]; |
| 397 | q_pack.u[0] = ((qa >> 0) & 0x000F000Fu) | BF16_MAGIC; |
| 398 | q_pack.u[1] = ((qa >> 4) & 0x000F000Fu) | BF16_MAGIC; |
| 399 | q_pack.u[2] = ((qa >> 8) & 0x000F000Fu) | BF16_MAGIC; |
| 400 | q_pack.u[3] = ((qa >> 12) & 0x000F000Fu) | BF16_MAGIC; |
| 401 | |
| 402 | #pragma unroll |
| 403 | for (int m = 0; m < BLOCK_SIZE_M; ++m) { |
| 404 | float partial = 0.0f; |
| 405 | #pragma unroll |
| 406 | for (int b = 0; b < 4; ++b) { |
| 407 | partial = __builtin_amdgcn_fdot2_f32_bf16( |
| 408 | *((bf16x2_t*)(&a_pack[m].f[b])), *((bf16x2_t*)(&q_pack.f[b])), |
| 409 | partial, /*clamp=*/false); |
| 410 | } |
| 411 | block_c[m][col] = |
| 412 | __fmaf_rn(y_b_f[col], partial, |
| 413 | __fmaf_rn(z_b_f[col], sum_a[m], block_c[m][col])); |
| 414 | } |
| 415 | } |
| 416 | } |
| 417 | } |
| 418 | k += 32; |
| 419 | } |
| 420 | |
| 421 | // --- Epilogue: apply topk_weight and atomic-add to output --- |
| 422 | #pragma unroll |
| 423 | for (int m = 0; m < BLOCK_SIZE_M; ++m) { |
| 424 | int32_t token_id = sorted_token_ids[offset_m_base + m]; |
| 425 | if (token_id / top_k >= size_m) continue; |
| 426 | |
| 427 | // Apply router weight |
| 428 | if (mul_topk_weight && topk_weights != nullptr) { |
| 429 | float tw = topk_weights[token_id]; |
| 430 | #pragma unroll |
| 431 | for (int j = 0; j < 4; ++j) block_c[m][j] *= tw; |
| 432 | } |
| 433 | |
| 434 | // output_topk > 0: reduce by mapping token_id back to original token |
| 435 | // (multiple experts write to the same row via atomics) |
| 436 | int64_t out_row = (output_topk > 0) ? (int64_t)(token_id / output_topk) |
| 437 | : (int64_t)token_id; |
| 438 | T* out = c + out_row * size_n + n; |
| 439 | if constexpr (std::is_same<T, half>::value) { |
| 440 | half2 r01 = __halves2half2(__float2half_rn(block_c[m][0]), |
| 441 | __float2half_rn(block_c[m][1])); |
| 442 | half2 r23 = __halves2half2(__float2half_rn(block_c[m][2]), |
| 443 | __float2half_rn(block_c[m][3])); |
| 444 | atomic_add_pk4_f16(out, r01, r23); |
| 445 | } else { |
| 446 | bf162_t r01; |
| 447 | r01.x = __float2bfloat16(block_c[m][0]); |
| 448 | r01.y = __float2bfloat16(block_c[m][1]); |
| 449 | bf162_t r23; |
| 450 | r23.x = __float2bfloat16(block_c[m][2]); |
| 451 | r23.y = __float2bfloat16(block_c[m][3]); |
| 452 | atomic_add_pk4_bf16(out, r01, r23); |
| 453 | } |
| 454 | } |
| 455 | } |
| 456 | |
| 457 | #else // non-RDNA3: empty stub for symbol parity |
| 458 | |
| 459 | template <typename T, int BLOCK_SIZE_M> |
| 460 | __global__ void moe_gemm_q4_kernel_rdna3( |
| 461 | const T*, T*, const uint32_t*, const T*, const uint32_t*, const float*, |
| 462 | const int32_t*, const int32_t*, const int32_t*, const int, const int, |
| 463 | const int, const int, const int, const int, const int, const int, |
| 464 | const bool, const int) {} |
| 465 | |
| 466 | #endif // __HIP__RDNA3__ || !__HIP_DEVICE_COMPILE__ |
| 467 | |
| 468 | // --------------------------------------------------------------------------- |
| 469 | // Launcher |
| 470 | // --------------------------------------------------------------------------- |
| 471 | |
| 472 | template <typename T, int BLOCK_SIZE_M> |
| 473 | void launch_moe_gemm_q4( |
| 474 | const T* a, T* c, const uint32_t* b_q_weight, const T* b_scales, |
| 475 | const uint32_t* b_qzeros, const float* topk_weights, |
| 476 | const int32_t* sorted_token_ids, const int32_t* expert_ids, |
| 477 | const int32_t* num_tokens_post_padded, int num_token_blocks, int size_m, |
| 478 | int size_n, int size_k, int groups, int top_k, int expert_weight_stride, |
| 479 | int expert_scales_stride, int expert_zeros_stride, bool mul_topk_weight, |
| 480 | int output_topk, cudaStream_t stream) { |
| 481 | dim3 block(THREADS_X); |
| 482 | dim3 grid(num_token_blocks, |
| 483 | (size_n + BLOCK_KN_SIZE * 4 - 1) / (BLOCK_KN_SIZE * 4), |
| 484 | (size_k + BLOCK_KN_SIZE - 1) / BLOCK_KN_SIZE); |
| 485 | |
| 486 | moe_gemm_q4_kernel_rdna3<T, BLOCK_SIZE_M><<<grid, block, 0, stream>>>( |
| 487 | a, c, b_q_weight, b_scales, b_qzeros, topk_weights, sorted_token_ids, |
| 488 | expert_ids, num_tokens_post_padded, size_m, size_n, size_k, groups, top_k, |
| 489 | expert_weight_stride, expert_scales_stride, expert_zeros_stride, |
| 490 | mul_topk_weight, output_topk); |
| 491 | } |
| 492 | |
| 493 | template <typename T> |
| 494 | void dispatch_moe_gemm_q4( |
| 495 | const T* a, T* c, const uint32_t* b_q_weight, const T* b_scales, |
| 496 | const uint32_t* b_qzeros, const float* topk_weights, |
| 497 | const int32_t* sorted_token_ids, const int32_t* expert_ids, |
| 498 | const int32_t* num_tokens_post_padded, int num_token_blocks, int size_m, |
| 499 | int size_n, int size_k, int groups, int top_k, int block_size_m, |
| 500 | int expert_weight_stride, int expert_scales_stride, int expert_zeros_stride, |
| 501 | bool mul_topk_weight, int output_topk, cudaStream_t stream) { |
| 502 | // Dispatch to template instantiation based on block_size_m |
| 503 | switch (block_size_m) { |
| 504 | case 1: |
| 505 | launch_moe_gemm_q4<T, 1>( |
| 506 | a, c, b_q_weight, b_scales, b_qzeros, topk_weights, sorted_token_ids, |
| 507 | expert_ids, num_tokens_post_padded, num_token_blocks, size_m, size_n, |
| 508 | size_k, groups, top_k, expert_weight_stride, expert_scales_stride, |
| 509 | expert_zeros_stride, mul_topk_weight, output_topk, stream); |
| 510 | break; |
| 511 | case 2: |
| 512 | launch_moe_gemm_q4<T, 2>( |
| 513 | a, c, b_q_weight, b_scales, b_qzeros, topk_weights, sorted_token_ids, |
| 514 | expert_ids, num_tokens_post_padded, num_token_blocks, size_m, size_n, |
| 515 | size_k, groups, top_k, expert_weight_stride, expert_scales_stride, |
| 516 | expert_zeros_stride, mul_topk_weight, output_topk, stream); |
| 517 | break; |
| 518 | case 4: |
| 519 | launch_moe_gemm_q4<T, 4>( |
| 520 | a, c, b_q_weight, b_scales, b_qzeros, topk_weights, sorted_token_ids, |
| 521 | expert_ids, num_tokens_post_padded, num_token_blocks, size_m, size_n, |
| 522 | size_k, groups, top_k, expert_weight_stride, expert_scales_stride, |
| 523 | expert_zeros_stride, mul_topk_weight, output_topk, stream); |
| 524 | break; |
| 525 | case 8: |
| 526 | launch_moe_gemm_q4<T, 8>( |
| 527 | a, c, b_q_weight, b_scales, b_qzeros, topk_weights, sorted_token_ids, |
| 528 | expert_ids, num_tokens_post_padded, num_token_blocks, size_m, size_n, |
| 529 | size_k, groups, top_k, expert_weight_stride, expert_scales_stride, |
| 530 | expert_zeros_stride, mul_topk_weight, output_topk, stream); |
| 531 | break; |
| 532 | default: |
| 533 | TORCH_CHECK(false, |
| 534 | "moe_gptq_gemm_rdna3: block_size_m must be 1, 2, 4, or 8, " |
| 535 | "got ", |
| 536 | block_size_m); |
| 537 | } |
| 538 | } |
| 539 | |
| 540 | } // namespace moe_gptq_rdna3 |
| 541 | } // namespace vllm |
| 542 | |
| 543 | // --------------------------------------------------------------------------- |
| 544 | // Public entry point |
| 545 | // --------------------------------------------------------------------------- |
| 546 | // |
| 547 | // Inputs: |
| 548 | // a [M, K] or [M*top_k, K] half or bfloat16 |
| 549 | // c [M*top_k, N] same dtype (pre-zeroed!) |
| 550 | // b_q_weight [E, K/8, N] uint32 (shuffled) |
| 551 | // b_scales [E, groups, N] same dtype as a |
| 552 | // b_qzeros [E, groups, N/8] uint32 (packed 4-bit) |
| 553 | // topk_weights [M*top_k] or empty float32 |
| 554 | // sorted_token_ids [num_blocks * block_m] int32 |
| 555 | // expert_ids [num_blocks] int32 |
| 556 | // num_tokens_post_padded [1] int32 |
| 557 | // top_k int |
| 558 | // block_size_m int (1, 2, 4, or 8) |
| 559 | // mul_topk_weight bool |
| 560 | |
| 561 | void moe_gptq_gemm_rdna3(torch::Tensor a, torch::Tensor c, |
| 562 | torch::Tensor b_q_weight, torch::Tensor b_scales, |
| 563 | torch::Tensor b_qzeros, torch::Tensor topk_weights, |
| 564 | torch::Tensor sorted_token_ids, |
| 565 | torch::Tensor expert_ids, |
| 566 | torch::Tensor num_tokens_post_padded, int64_t top_k, |
| 567 | int64_t block_size_m, bool mul_topk_weight, |
| 568 | int64_t output_topk) { |
| 569 | TORCH_CHECK(a.is_cuda(), "a must be a CUDA/HIP tensor"); |
| 570 | TORCH_CHECK(c.is_cuda(), "c must be a CUDA/HIP tensor"); |
| 571 | TORCH_CHECK(b_q_weight.is_cuda(), "b_q_weight must be a CUDA/HIP tensor"); |
| 572 | TORCH_CHECK(a.dim() == 2, "a must be 2D"); |
| 573 | TORCH_CHECK(c.dim() == 2, "c must be 2D"); |
| 574 | TORCH_CHECK(b_q_weight.dim() == 3, "b_q_weight must be 3D [E, K/8, N]"); |
| 575 | TORCH_CHECK(b_scales.dim() == 3, "b_scales must be 3D [E, groups, N]"); |
| 576 | TORCH_CHECK(b_qzeros.dim() == 3, "b_qzeros must be 3D [E, groups, N/8]"); |
| 577 | TORCH_CHECK( |
| 578 | a.scalar_type() == torch::kHalf || a.scalar_type() == torch::kBFloat16, |
| 579 | "a must be half or bfloat16"); |
| 580 | TORCH_CHECK(a.scalar_type() == b_scales.scalar_type(), |
| 581 | "b_scales dtype must match a"); |
| 582 | |
| 583 | const at::cuda::OptionalCUDAGuard device_guard(device_of(a)); |
| 584 | auto stream = at::cuda::getCurrentCUDAStream(); |
| 585 | |
| 586 | int size_m = (int)a.size(0); |
| 587 | int size_k = (int)a.size(1); |
| 588 | int size_n = (int)b_q_weight.size(2); |
| 589 | int groups = (int)b_scales.size(1); |
| 590 | |
| 591 | // Per-expert strides |
| 592 | int expert_weight_stride = (int)(b_q_weight.size(1) * b_q_weight.size(2)); |
| 593 | int expert_scales_stride = (int)(b_scales.size(1) * b_scales.size(2)); |
| 594 | int expert_zeros_stride = (int)(b_qzeros.size(1) * b_qzeros.size(2)); |
| 595 | |
| 596 | int num_token_blocks = (int)(sorted_token_ids.size(0) / block_size_m); |
| 597 | |
| 598 | const float* topk_w_ptr = |
| 599 | (topk_weights.numel() > 0) ? topk_weights.data_ptr<float>() : nullptr; |
| 600 | |
| 601 | // Manual dtype dispatch using HIP native types (c10::Half/BFloat16 don't |
| 602 | // implicitly convert to half/__hip_bfloat16 in device code). |
| 603 | using vllm::gptq_rdna3::bf16_t; |
| 604 | |
| 605 | auto dispatch = [&](auto* a_ptr, auto* c_ptr, const auto* s_ptr) { |
| 606 | using T = std::remove_const_t<std::remove_pointer_t<decltype(a_ptr)>>; |
| 607 | vllm::moe_gptq_rdna3::dispatch_moe_gemm_q4<T>( |
| 608 | a_ptr, c_ptr, (const uint32_t*)b_q_weight.data_ptr<int32_t>(), s_ptr, |
| 609 | (const uint32_t*)b_qzeros.data_ptr<int32_t>(), topk_w_ptr, |
| 610 | sorted_token_ids.data_ptr<int32_t>(), expert_ids.data_ptr<int32_t>(), |
| 611 | num_tokens_post_padded.data_ptr<int32_t>(), num_token_blocks, size_m, |
| 612 | size_n, size_k, groups, (int)top_k, (int)block_size_m, |
| 613 | expert_weight_stride, expert_scales_stride, expert_zeros_stride, |
| 614 | mul_topk_weight, (int)output_topk, stream); |
| 615 | }; |
| 616 | |
| 617 | if (a.scalar_type() == torch::kHalf) { |
| 618 | dispatch((const half*)a.data_ptr(), (half*)c.data_ptr(), |
| 619 | (const half*)b_scales.data_ptr()); |
| 620 | } else { |
| 621 | dispatch((const bf16_t*)a.data_ptr(), (bf16_t*)c.data_ptr(), |
| 622 | (const bf16_t*)b_scales.data_ptr()); |
| 623 | } |
| 624 | } |