@tduncan / Nvidia-Optimization-Stack-VLLM Mirror · read-only / commits / 062b05ff3a
Fork 0

A high-throughput and memory-efficient inference and serving engine for LLMs

Mirrored from https://github.com/tduncan-lum/Nvidia-Optimization-Stack---VLLM.git

Synced Nov 1, 18:45 UTC

[ROCm][Perf] Fused MoE W4A16 HIP kernel for AMD RDNA3 (gfx1100) (#44075)

Signed-off-by: JartX <sagformas@epdcenter.es> Co-authored-by: TJian <tunjian.tan@embeddedllm.com>

JartX committed Jun 6, 2026 at 22:30 UTC 062b05ff3af405b8800f7c42fa3998cc03dd97eb
11 files changed +1906 -5
CMakeLists.txt
+2 -1
@@ -1310,7 +1310,8 @@ if(VLLM_GPU_LANG STREQUAL "HIP")
1310 set(VLLM_ROCM_HAS_GFX1100 ON)
1311 list(APPEND VLLM_ROCM_EXT_SRC
1312 "csrc/rocm/q_gemm_rdna3.cu"
1313 - "csrc/rocm/q_gemm_rdna3_wmma.cu")
1313 + "csrc/rocm/q_gemm_rdna3_wmma.cu"
1314 + "csrc/rocm/moe_q_gemm_rdna3.cu")
1315 endif()
1316
1317 define_extension_target(
csrc/rocm/moe_q_gemm_rdna3.cu new
+639
@@ -0,0 +1,639 @@
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__ float dot22_8_f(float (&dq)[8],
71 + const bf16_t* a_ptr) {
72 + float result = 0.0f;
73 + #pragma unroll
74 + for (int i = 0; i < 4; i++) {
75 + uint32_t aw;
76 + __builtin_memcpy(&aw, a_ptr + 2 * i, sizeof(uint32_t));
77 + float a_x = __uint_as_float((aw & 0xFFFFu) << 16);
78 + float a_y = __uint_as_float(aw & 0xFFFF0000u);
79 + result = __fmaf_rn(dq[2 * i + 0], a_x, result);
80 + result = __fmaf_rn(dq[2 * i + 1], a_y, result);
81 + }
82 + return result;
83 +}
84 +
85 +__forceinline__ __device__ void atomic_add_pk4_f16(half* addr, half2 v01,
86 + half2 v23) {
87 + unsigned long long* addr_u = reinterpret_cast<unsigned long long*>(addr);
88 + unsigned long long old = *addr_u;
89 + while (true) {
90 + union {
91 + unsigned long long u;
92 + half2 h2[2];
93 + } cur, sum;
94 + cur.u = old;
95 + sum.h2[0] = __hadd2(cur.h2[0], v01);
96 + sum.h2[1] = __hadd2(cur.h2[1], v23);
97 + unsigned long long prev = atomicCAS(addr_u, old, sum.u);
98 + if (prev == old) break;
99 + old = prev;
100 + }
101 +}
102 +
103 +__forceinline__ __device__ void atomic_add_pk4_bf16(bf16_t* addr, bf162_t v01,
104 + bf162_t v23) {
105 + unsigned long long* addr_u = reinterpret_cast<unsigned long long*>(addr);
106 + unsigned long long old = *addr_u;
107 + while (true) {
108 + union {
109 + unsigned long long u;
110 + bf162_t b2[2];
111 + } cur, sum;
112 + cur.u = old;
113 + sum.b2[0] = __hadd2(cur.b2[0], v01);
114 + sum.b2[1] = __hadd2(cur.b2[1], v23);
115 + unsigned long long prev = atomicCAS(addr_u, old, sum.u);
116 + if (prev == old) break;
117 + old = prev;
118 + }
119 +}
120 +
121 +__forceinline__ __device__ void load4_zeros(const uint32_t* qzeros_row, int n,
122 + int (&zeros)[4]) {
123 + int qcol = n / 8;
124 + int shift = (n & 0x07) * 4;
125 + uint32_t d = qzeros_row[qcol] >> shift;
126 + zeros[0] = (int)(d & 0xF);
127 + zeros[1] = (int)((d >> 4) & 0xF);
128 + zeros[2] = (int)((d >> 8) & 0xF);
129 + zeros[3] = (int)((d >> 12) & 0xF);
130 +}
131 +
132 +template <typename T>
133 +__forceinline__ __device__ void load4_scales(const T* scales_row, int n,
134 + T (&scales)[4]) {
135 + scales[0] = scales_row[n + 0];
136 + scales[1] = scales_row[n + 1];
137 + scales[2] = scales_row[n + 2];
138 + scales[3] = scales_row[n + 3];
139 +}
140 +
141 +// ---------------------------------------------------------------------------
142 +// Fused MoE kernel.
143 +// ---------------------------------------------------------------------------
144 +
145 +template <typename T, int BLOCK_SIZE_M>
146 +__global__ void moe_gemm_q4_kernel_rdna3(
147 + const T* __restrict__ a, // [size_m, size_k] or [M*topk, K]
148 + T* __restrict__ c, // [M*topk, size_n] pre-zeroed
149 + const uint32_t* __restrict__ b_q_weight, // [E, K/8, N] packed
150 + const T* __restrict__ b_scales, // [E, groups, N]
151 + const uint32_t* __restrict__ b_qzeros, // [E, groups, N/8] packed
152 + const float* __restrict__ topk_weights, // [M*topk] or nullptr
153 + const int32_t* __restrict__ sorted_token_ids,
154 + const int32_t* __restrict__ expert_ids,
155 + const int32_t* __restrict__ num_tokens_post_padded,
156 + const int size_m, // total tokens (original M, or M*topk for w2)
157 + const int size_n, // output features per expert
158 + const int size_k, // input features
159 + const int groups, // K / group_size
160 + const int top_k, // routing top-k (1 for w2 pass)
161 + // Per-expert strides (in elements, not bytes)
162 + const int expert_weight_stride, // (K/8) * N
163 + const int expert_scales_stride, // groups * N
164 + const int expert_zeros_stride, // groups * (N/8)
165 + const bool mul_topk_weight,
166 + const int output_topk) { // >0: reduce output by token_id/output_topk
167 + const int t = threadIdx.x;
168 + const int token_block = blockIdx.x;
169 + const int offset_n = blockIdx.y * BLOCK_KN_SIZE * 4;
170 + const int offset_k = blockIdx.z * BLOCK_KN_SIZE;
171 + const int end_k = min(offset_k + BLOCK_KN_SIZE, size_k);
172 + const int n = offset_n + t * 4;
173 +
174 + // Early exit for padding blocks or invalid experts (expert_map = -1)
175 + if (token_block * BLOCK_SIZE_M >= num_tokens_post_padded[0]) return;
176 +
177 + const int expert_id = expert_ids[token_block];
178 + if (expert_id == -1) return;
179 +
180 + // Expert-specific pointers
181 + const uint32_t* expert_weights =
182 + b_q_weight + (int64_t)expert_id * expert_weight_stride;
183 + const T* expert_scales = b_scales + (int64_t)expert_id * expert_scales_stride;
184 + const uint32_t* expert_qzeros =
185 + b_qzeros + (int64_t)expert_id * expert_zeros_stride;
186 +
187 + // LDS for activations
188 + constexpr int LDS_PAD = 8;
189 + __shared__ T block_a[BLOCK_SIZE_M][BLOCK_KN_SIZE + LDS_PAD];
190 +
191 + static_assert(BLOCK_KN_SIZE == THREADS_X,
192 + "BLOCK_KN_SIZE must equal THREADS_X");
193 +
194 + // For bf16 M=1, we can skip LDS and read A from global (same as dense).
195 + // fp16 always needs LDS due to the dot22_8_f indexing pattern.
196 + constexpr bool USE_LDS_A = (BLOCK_SIZE_M > 1) || std::is_same<T, half>::value;
197 +
198 + const int offset_m_base = token_block * BLOCK_SIZE_M;
199 +
200 + if constexpr (USE_LDS_A) {
201 + if (offset_k + t < end_k) {
202 + #pragma unroll
203 + for (int m = 0; m < BLOCK_SIZE_M; ++m) {
204 + int32_t token_id = sorted_token_ids[offset_m_base + m];
205 + int token_row = token_id / top_k;
206 + T av;
207 + if (token_row < size_m) {
208 + av = a[(int64_t)token_row * size_k + offset_k + t];
209 + } else {
210 + av = tzero<T>();
211 + }
212 + block_a[m][t] = av;
213 + }
214 + }
215 + __syncthreads();
216 + }
217 +
218 + if (n >= size_n) return;
219 +
220 + // Group bookkeeping
221 + const int groupsize = size_k / groups;
222 + int group = offset_k / groupsize;
223 + int nextgroup = (group + 1) * groupsize;
224 +
225 + // Weight pointer for this expert
226 + int qk = offset_k / 8;
227 + const uint32_t* b_ptr = expert_weights + qk * size_n + n;
228 +
229 + // Per-column dequant constants (4 columns per thread)
230 + half2 z1z16_h[4][2], y1y16_h[4][2];
231 + float z_b_f[4], y_b_f[4];
232 +
233 + // GPTQv1: zero_offset = 1
234 + constexpr int zero_offset = 1;
235 +
236 + auto refresh_group = [&](int g) {
237 + const uint32_t* qz_row = expert_qzeros + g * (size_n / 8);
238 + const T* sc_row = expert_scales + g * size_n;
239 + int zeros[4];
240 + T scales[4];
241 + load4_zeros(qz_row, n, zeros);
242 + load4_scales<T>(sc_row, n, scales);
243 + if constexpr (std::is_same<T, half>::value) {
244 + #pragma unroll
245 + for (int i = 0; i < 4; ++i) {
246 + gptq_rdna3::prep_zero_scale_fp16((uint32_t)(zeros[i] + zero_offset),
247 + scales[i], z1z16_h[i], y1y16_h[i]);
248 + }
249 + } else {
250 + #pragma unroll
251 + for (int i = 0; i < 4; ++i) {
252 + gptq_rdna3::prep_zero_scale_bf16_f32((uint32_t)(zeros[i] + zero_offset),
253 + scales[i], z_b_f[i], y_b_f[i]);
254 + }
255 + }
256 + };
257 +
258 + refresh_group(group);
259 +
260 + float block_c[BLOCK_SIZE_M][4];
261 + #pragma unroll
262 + for (int m = 0; m < BLOCK_SIZE_M; ++m) {
263 + #pragma unroll
264 + for (int j = 0; j < 4; ++j) block_c[m][j] = 0.0f;
265 + }
266 +
267 + // --- Main K-loop ---
268 + int k = offset_k;
269 + while (k < end_k) {
270 + if (k == nextgroup) {
271 + group++;
272 + nextgroup += groupsize;
273 + refresh_group(group);
274 + }
275 +
276 + // Prefetch 4 weight words (128 bytes)
277 + int4 b_w[4];
278 + #pragma unroll
279 + for (int j = 0; j < 4; ++j) {
280 + b_w[j] = *(const int4*)(b_ptr + j * size_n);
281 + }
282 + b_ptr += 4 * size_n;
283 +
284 + #pragma unroll
285 + for (int j = 0; j < 4; ++j) {
286 + const int a_off = (k - offset_k) + 8 * j;
287 +
288 + if constexpr (std::is_same<T, half>::value) {
289 + // fp16 path: dequant via bit-trick, dot via v_dot2_f32_f16
290 + half2 dq[4][4];
291 + gptq_rdna3::dequant_4bit_8_fp16((uint32_t)b_w[j].x, dq[0], z1z16_h[0],
292 + y1y16_h[0]);
293 + gptq_rdna3::dequant_4bit_8_fp16((uint32_t)b_w[j].y, dq[1], z1z16_h[1],
294 + y1y16_h[1]);
295 + gptq_rdna3::dequant_4bit_8_fp16((uint32_t)b_w[j].z, dq[2], z1z16_h[2],
296 + y1y16_h[2]);
297 + gptq_rdna3::dequant_4bit_8_fp16((uint32_t)b_w[j].w, dq[3], z1z16_h[3],
298 + y1y16_h[3]);
299 +
300 + #pragma unroll
301 + for (int m = 0; m < BLOCK_SIZE_M; ++m) {
302 + const half* a_ptr = reinterpret_cast<const half*>(&block_a[m][a_off]);
303 + block_c[m][0] += dot22_8_f(dq[0], a_ptr);
304 + block_c[m][1] += dot22_8_f(dq[1], a_ptr);
305 + block_c[m][2] += dot22_8_f(dq[2], a_ptr);
306 + block_c[m][3] += dot22_8_f(dq[3], a_ptr);
307 + }
308 + } else if constexpr (BLOCK_SIZE_M == 1) {
309 + // bf16 M=1: v_dot2_f32_bf16 with InstCombine-defeating opacity
310 + typedef short __attribute__((ext_vector_type(2))) bf16x2_t;
311 + constexpr uint32_t BF16_MAGIC = 0x43004300u;
312 + constexpr uint32_t BF16_ONES = 0x3F803F80u;
313 + union pack4 {
314 + float f[4];
315 + uint32_t u[4];
316 + };
317 +
318 + uint32_t w[4];
319 + __builtin_memcpy(w, &b_w[j], sizeof(int4));
320 +
321 + // Load activations — read from global (no LDS for bf16 M=1)
322 + pack4 a_pack;
323 + {
324 + int32_t token_id = sorted_token_ids[offset_m_base];
325 + int token_row = token_id / top_k;
326 + if (token_row < size_m) {
327 + const uint32_t* a_words = reinterpret_cast<const uint32_t*>(
328 + a + (int64_t)token_row * size_k + offset_k + a_off);
329 + a_pack.u[0] = a_words[0];
330 + a_pack.u[1] = a_words[1];
331 + a_pack.u[2] = a_words[2];
332 + a_pack.u[3] = a_words[3];
333 + } else {
334 + a_pack.u[0] = 0;
335 + a_pack.u[1] = 0;
336 + a_pack.u[2] = 0;
337 + a_pack.u[3] = 0;
338 + }
339 + }
340 +
341 + // sum_a for bias correction
342 + float sum_a = 0.0f;
343 + #pragma unroll
344 + for (int b = 0; b < 4; ++b) {
345 + sum_a = __builtin_amdgcn_fdot2_f32_bf16(
346 + *((bf16x2_t*)(&a_pack.f[b])), *((const bf16x2_t*)&BF16_ONES),
347 + sum_a, /*clamp=*/false);
348 + }
349 +
350 + #pragma unroll 1
351 + for (int col = 0; col < 4; ++col) {
352 + pack4 q_pack;
353 + const uint32_t qa = w[col];
354 + q_pack.u[0] = ((qa >> 0) & 0x000F000Fu) | BF16_MAGIC;
355 + q_pack.u[1] = ((qa >> 4) & 0x000F000Fu) | BF16_MAGIC;
356 + q_pack.u[2] = ((qa >> 8) & 0x000F000Fu) | BF16_MAGIC;
357 + q_pack.u[3] = ((qa >> 12) & 0x000F000Fu) | BF16_MAGIC;
358 +
359 + float partial = 0.0f;
360 + #pragma unroll
361 + for (int b = 0; b < 4; ++b) {
362 + partial = __builtin_amdgcn_fdot2_f32_bf16(
363 + *((bf16x2_t*)(&a_pack.f[b])), *((bf16x2_t*)(&q_pack.f[b])),
364 + partial, /*clamp=*/false);
365 + }
366 +
367 + block_c[0][col] =
368 + __fmaf_rn(y_b_f[col], partial,
369 + __fmaf_rn(z_b_f[col], sum_a, block_c[0][col]));
370 + }
371 + } else {
372 + // bf16 M>1: v_dot2_f32_bf16 with LDS-staged activations
373 + typedef short __attribute__((ext_vector_type(2))) bf16x2_t;
374 + constexpr uint32_t BF16_MAGIC = 0x43004300u;
375 + constexpr uint32_t BF16_ONES = 0x3F803F80u;
376 + union pack4 {
377 + float f[4];
378 + uint32_t u[4];
379 + };
380 +
381 + uint32_t w[4];
382 + __builtin_memcpy(w, &b_w[j], sizeof(int4));
383 +
384 + pack4 a_pack[BLOCK_SIZE_M];
385 + #pragma unroll
386 + for (int m = 0; m < BLOCK_SIZE_M; ++m) {
387 + const uint32_t* a_words =
388 + reinterpret_cast<const uint32_t*>(&block_a[m][a_off]);
389 + a_pack[m].u[0] = a_words[0];
390 + a_pack[m].u[1] = a_words[1];
391 + a_pack[m].u[2] = a_words[2];
392 + a_pack[m].u[3] = a_words[3];
393 + }
394 +
395 + float sum_a[BLOCK_SIZE_M];
396 + #pragma unroll
397 + for (int m = 0; m < BLOCK_SIZE_M; ++m) {
398 + float s = 0.0f;
399 + #pragma unroll
400 + for (int b = 0; b < 4; ++b) {
401 + s = __builtin_amdgcn_fdot2_f32_bf16(*((bf16x2_t*)(&a_pack[m].f[b])),
402 + *((const bf16x2_t*)&BF16_ONES),
403 + s, /*clamp=*/false);
404 + }
405 + sum_a[m] = s;
406 + }
407 +
408 + #pragma unroll 1
409 + for (int col = 0; col < 4; ++col) {
410 + pack4 q_pack;
411 + const uint32_t qa = w[col];
412 + q_pack.u[0] = ((qa >> 0) & 0x000F000Fu) | BF16_MAGIC;
413 + q_pack.u[1] = ((qa >> 4) & 0x000F000Fu) | BF16_MAGIC;
414 + q_pack.u[2] = ((qa >> 8) & 0x000F000Fu) | BF16_MAGIC;
415 + q_pack.u[3] = ((qa >> 12) & 0x000F000Fu) | BF16_MAGIC;
416 +
417 + #pragma unroll
418 + for (int m = 0; m < BLOCK_SIZE_M; ++m) {
419 + float partial = 0.0f;
420 + #pragma unroll
421 + for (int b = 0; b < 4; ++b) {
422 + partial = __builtin_amdgcn_fdot2_f32_bf16(
423 + *((bf16x2_t*)(&a_pack[m].f[b])), *((bf16x2_t*)(&q_pack.f[b])),
424 + partial, /*clamp=*/false);
425 + }
426 + block_c[m][col] =
427 + __fmaf_rn(y_b_f[col], partial,
428 + __fmaf_rn(z_b_f[col], sum_a[m], block_c[m][col]));
429 + }
430 + }
431 + }
432 + }
433 + k += 32;
434 + }
435 +
436 + // --- Epilogue: apply topk_weight and atomic-add to output ---
437 + #pragma unroll
438 + for (int m = 0; m < BLOCK_SIZE_M; ++m) {
439 + int32_t token_id = sorted_token_ids[offset_m_base + m];
440 + if (token_id / top_k >= size_m) continue;
441 +
442 + // Apply router weight
443 + if (mul_topk_weight && topk_weights != nullptr) {
444 + float tw = topk_weights[token_id];
445 + #pragma unroll
446 + for (int j = 0; j < 4; ++j) block_c[m][j] *= tw;
447 + }
448 +
449 + // output_topk > 0: reduce by mapping token_id back to original token
450 + // (multiple experts write to the same row via atomics)
451 + int64_t out_row = (output_topk > 0) ? (int64_t)(token_id / output_topk)
452 + : (int64_t)token_id;
453 + T* out = c + out_row * size_n + n;
454 + if constexpr (std::is_same<T, half>::value) {
455 + half2 r01 = __halves2half2(__float2half_rn(block_c[m][0]),
456 + __float2half_rn(block_c[m][1]));
457 + half2 r23 = __halves2half2(__float2half_rn(block_c[m][2]),
458 + __float2half_rn(block_c[m][3]));
459 + atomic_add_pk4_f16(out, r01, r23);
460 + } else {
461 + bf162_t r01;
462 + r01.x = __float2bfloat16(block_c[m][0]);
463 + r01.y = __float2bfloat16(block_c[m][1]);
464 + bf162_t r23;
465 + r23.x = __float2bfloat16(block_c[m][2]);
466 + r23.y = __float2bfloat16(block_c[m][3]);
467 + atomic_add_pk4_bf16(out, r01, r23);
468 + }
469 + }
470 +}
471 +
472 +#else // non-RDNA3: empty stub for symbol parity
473 +
474 +template <typename T, int BLOCK_SIZE_M>
475 +__global__ void moe_gemm_q4_kernel_rdna3(
476 + const T*, T*, const uint32_t*, const T*, const uint32_t*, const float*,
477 + const int32_t*, const int32_t*, const int32_t*, const int, const int,
478 + const int, const int, const int, const int, const int, const int,
479 + const bool, const int) {}
480 +
481 +#endif // __HIP__RDNA3__ || !__HIP_DEVICE_COMPILE__
482 +
483 +// ---------------------------------------------------------------------------
484 +// Launcher
485 +// ---------------------------------------------------------------------------
486 +
487 +template <typename T, int BLOCK_SIZE_M>
488 +void launch_moe_gemm_q4(
489 + const T* a, T* c, const uint32_t* b_q_weight, const T* b_scales,
490 + const uint32_t* b_qzeros, const float* topk_weights,
491 + const int32_t* sorted_token_ids, const int32_t* expert_ids,
492 + const int32_t* num_tokens_post_padded, int num_token_blocks, int size_m,
493 + int size_n, int size_k, int groups, int top_k, int expert_weight_stride,
494 + int expert_scales_stride, int expert_zeros_stride, bool mul_topk_weight,
495 + int output_topk, cudaStream_t stream) {
496 + dim3 block(THREADS_X);
497 + dim3 grid(num_token_blocks,
498 + (size_n + BLOCK_KN_SIZE * 4 - 1) / (BLOCK_KN_SIZE * 4),
499 + (size_k + BLOCK_KN_SIZE - 1) / BLOCK_KN_SIZE);
500 +
501 + moe_gemm_q4_kernel_rdna3<T, BLOCK_SIZE_M><<<grid, block, 0, stream>>>(
502 + a, c, b_q_weight, b_scales, b_qzeros, topk_weights, sorted_token_ids,
503 + expert_ids, num_tokens_post_padded, size_m, size_n, size_k, groups, top_k,
504 + expert_weight_stride, expert_scales_stride, expert_zeros_stride,
505 + mul_topk_weight, output_topk);
506 +}
507 +
508 +template <typename T>
509 +void dispatch_moe_gemm_q4(
510 + const T* a, T* c, const uint32_t* b_q_weight, const T* b_scales,
511 + const uint32_t* b_qzeros, const float* topk_weights,
512 + const int32_t* sorted_token_ids, const int32_t* expert_ids,
513 + const int32_t* num_tokens_post_padded, int num_token_blocks, int size_m,
514 + int size_n, int size_k, int groups, int top_k, int block_size_m,
515 + int expert_weight_stride, int expert_scales_stride, int expert_zeros_stride,
516 + bool mul_topk_weight, int output_topk, cudaStream_t stream) {
517 + // Dispatch to template instantiation based on block_size_m
518 + switch (block_size_m) {
519 + case 1:
520 + launch_moe_gemm_q4<T, 1>(
521 + a, c, b_q_weight, b_scales, b_qzeros, topk_weights, sorted_token_ids,
522 + expert_ids, num_tokens_post_padded, num_token_blocks, size_m, size_n,
523 + size_k, groups, top_k, expert_weight_stride, expert_scales_stride,
524 + expert_zeros_stride, mul_topk_weight, output_topk, stream);
525 + break;
526 + case 2:
527 + launch_moe_gemm_q4<T, 2>(
528 + a, c, b_q_weight, b_scales, b_qzeros, topk_weights, sorted_token_ids,
529 + expert_ids, num_tokens_post_padded, num_token_blocks, size_m, size_n,
530 + size_k, groups, top_k, expert_weight_stride, expert_scales_stride,
531 + expert_zeros_stride, mul_topk_weight, output_topk, stream);
532 + break;
533 + case 4:
534 + launch_moe_gemm_q4<T, 4>(
535 + a, c, b_q_weight, b_scales, b_qzeros, topk_weights, sorted_token_ids,
536 + expert_ids, num_tokens_post_padded, num_token_blocks, size_m, size_n,
537 + size_k, groups, top_k, expert_weight_stride, expert_scales_stride,
538 + expert_zeros_stride, mul_topk_weight, output_topk, stream);
539 + break;
540 + case 8:
541 + launch_moe_gemm_q4<T, 8>(
542 + a, c, b_q_weight, b_scales, b_qzeros, topk_weights, sorted_token_ids,
543 + expert_ids, num_tokens_post_padded, num_token_blocks, size_m, size_n,
544 + size_k, groups, top_k, expert_weight_stride, expert_scales_stride,
545 + expert_zeros_stride, mul_topk_weight, output_topk, stream);
546 + break;
547 + default:
548 + TORCH_CHECK(false,
549 + "moe_gptq_gemm_rdna3: block_size_m must be 1, 2, 4, or 8, "
550 + "got ",
551 + block_size_m);
552 + }
553 +}
554 +
555 +} // namespace moe_gptq_rdna3
556 +} // namespace vllm
557 +
558 +// ---------------------------------------------------------------------------
559 +// Public entry point
560 +// ---------------------------------------------------------------------------
561 +//
562 +// Inputs:
563 +// a [M, K] or [M*top_k, K] half or bfloat16
564 +// c [M*top_k, N] same dtype (pre-zeroed!)
565 +// b_q_weight [E, K/8, N] uint32 (shuffled)
566 +// b_scales [E, groups, N] same dtype as a
567 +// b_qzeros [E, groups, N/8] uint32 (packed 4-bit)
568 +// topk_weights [M*top_k] or empty float32
569 +// sorted_token_ids [num_blocks * block_m] int32
570 +// expert_ids [num_blocks] int32
571 +// num_tokens_post_padded [1] int32
572 +// top_k int
573 +// block_size_m int (1, 2, 4, or 8)
574 +// mul_topk_weight bool
575 +
576 +void moe_gptq_gemm_rdna3(torch::Tensor a, torch::Tensor c,
577 + torch::Tensor b_q_weight, torch::Tensor b_scales,
578 + torch::Tensor b_qzeros, torch::Tensor topk_weights,
579 + torch::Tensor sorted_token_ids,
580 + torch::Tensor expert_ids,
581 + torch::Tensor num_tokens_post_padded, int64_t top_k,
582 + int64_t block_size_m, bool mul_topk_weight,
583 + int64_t output_topk) {
584 + TORCH_CHECK(a.is_cuda(), "a must be a CUDA/HIP tensor");
585 + TORCH_CHECK(c.is_cuda(), "c must be a CUDA/HIP tensor");
586 + TORCH_CHECK(b_q_weight.is_cuda(), "b_q_weight must be a CUDA/HIP tensor");
587 + TORCH_CHECK(a.dim() == 2, "a must be 2D");
588 + TORCH_CHECK(c.dim() == 2, "c must be 2D");
589 + TORCH_CHECK(b_q_weight.dim() == 3, "b_q_weight must be 3D [E, K/8, N]");
590 + TORCH_CHECK(b_scales.dim() == 3, "b_scales must be 3D [E, groups, N]");
591 + TORCH_CHECK(b_qzeros.dim() == 3, "b_qzeros must be 3D [E, groups, N/8]");
592 + TORCH_CHECK(
593 + a.scalar_type() == torch::kHalf || a.scalar_type() == torch::kBFloat16,
594 + "a must be half or bfloat16");
595 + TORCH_CHECK(a.scalar_type() == b_scales.scalar_type(),
596 + "b_scales dtype must match a");
597 +
598 + const at::cuda::OptionalCUDAGuard device_guard(device_of(a));
599 + auto stream = at::cuda::getCurrentCUDAStream();
600 +
601 + int size_m = (int)a.size(0);
602 + int size_k = (int)a.size(1);
603 + int size_n = (int)b_q_weight.size(2);
604 + int groups = (int)b_scales.size(1);
605 +
606 + // Per-expert strides
607 + int expert_weight_stride = (int)(b_q_weight.size(1) * b_q_weight.size(2));
608 + int expert_scales_stride = (int)(b_scales.size(1) * b_scales.size(2));
609 + int expert_zeros_stride = (int)(b_qzeros.size(1) * b_qzeros.size(2));
610 +
611 + int num_token_blocks = (int)(sorted_token_ids.size(0) / block_size_m);
612 +
613 + const float* topk_w_ptr =
614 + (topk_weights.numel() > 0) ? topk_weights.data_ptr<float>() : nullptr;
615 +
616 + // Manual dtype dispatch using HIP native types (c10::Half/BFloat16 don't
617 + // implicitly convert to half/__hip_bfloat16 in device code).
618 + using vllm::gptq_rdna3::bf16_t;
619 +
620 + auto dispatch = [&](auto* a_ptr, auto* c_ptr, const auto* s_ptr) {
621 + using T = std::remove_const_t<std::remove_pointer_t<decltype(a_ptr)>>;
622 + vllm::moe_gptq_rdna3::dispatch_moe_gemm_q4<T>(
623 + a_ptr, c_ptr, (const uint32_t*)b_q_weight.data_ptr<int32_t>(), s_ptr,
624 + (const uint32_t*)b_qzeros.data_ptr<int32_t>(), topk_w_ptr,
625 + sorted_token_ids.data_ptr<int32_t>(), expert_ids.data_ptr<int32_t>(),
626 + num_tokens_post_padded.data_ptr<int32_t>(), num_token_blocks, size_m,
627 + size_n, size_k, groups, (int)top_k, (int)block_size_m,
628 + expert_weight_stride, expert_scales_stride, expert_zeros_stride,
629 + mul_topk_weight, (int)output_topk, stream);
630 + };
631 +
632 + if (a.scalar_type() == torch::kHalf) {
633 + dispatch((const half*)a.data_ptr(), (half*)c.data_ptr(),
634 + (const half*)b_scales.data_ptr());
635 + } else {
636 + dispatch((const bf16_t*)a.data_ptr(), (bf16_t*)c.data_ptr(),
637 + (const bf16_t*)b_scales.data_ptr());
638 + }
639 +}
csrc/rocm/ops.h
+9
@@ -27,6 +27,15 @@ torch::Tensor gptq_gemm_rdna3_wmma(torch::Tensor a, torch::Tensor b_q_weight,
27 torch::Tensor b_scales,
28 torch::Tensor b_g_idx, bool use_v2_format);
29
30 +void moe_gptq_gemm_rdna3(torch::Tensor a, torch::Tensor c,
31 + torch::Tensor b_q_weight, torch::Tensor b_scales,
32 + torch::Tensor b_qzeros, torch::Tensor topk_weights,
33 + torch::Tensor sorted_token_ids,
34 + torch::Tensor expert_ids,
35 + torch::Tensor num_tokens_post_padded, int64_t top_k,
36 + int64_t block_size_m, bool mul_topk_weight,
37 + int64_t output_topk);
38 +
39 void paged_attention(
40 torch::Tensor& out, torch::Tensor& exp_sums, torch::Tensor& max_logits,
41 torch::Tensor& tmp_out, torch::Tensor& query, torch::Tensor& key_cache,
csrc/rocm/torch_bindings.cpp
+9
@@ -50,6 +50,15 @@ TORCH_LIBRARY_EXPAND(TORCH_EXTENSION_NAME, rocm_ops) {
50 "gptq_gemm_rdna3_wmma(Tensor a, Tensor b_q_weight, Tensor b_qzeros, "
51 "Tensor b_scales, Tensor b_g_idx, bool use_v2_format) -> Tensor");
52 rocm_ops.impl("gptq_gemm_rdna3_wmma", torch::kCUDA, &gptq_gemm_rdna3_wmma);
53 +
54 + rocm_ops.def(
55 + "moe_gptq_gemm_rdna3(Tensor a, Tensor! c, Tensor b_q_weight, "
56 + "Tensor b_scales, Tensor b_qzeros, Tensor topk_weights, "
57 + "Tensor sorted_token_ids, Tensor expert_ids, "
58 + "Tensor num_tokens_post_padded, "
59 + "int top_k, int block_size_m, bool mul_topk_weight, "
60 + "int output_topk) -> ()");
61 + rocm_ops.impl("moe_gptq_gemm_rdna3", torch::kCUDA, &moe_gptq_gemm_rdna3);
62 #endif
63
64 // Custom attention op
tests/kernels/quantization/test_rdna3_compile_guards.py new
+503
@@ -0,0 +1,503 @@
1 +#!/usr/bin/env python3
2 +# SPDX-License-Identifier: Apache-2.0
3 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
4 +"""Compile-guard tests for the ROCm RDNA3 W4A16 kernels (dense + MoE).
5 +
6 +Verifies that the gfx1100 compilation and dispatch guards are hermetic:
7 + - On gfx1100: all ops exist, dispatch selects RDNA3 kernels.
8 + - On CDNA (gfx942/gfx950) or other non-gfx1100: ops must NOT exist,
9 + dispatch must fall through to Triton/Marlin, and no RDNA3 code
10 + path is reachable.
11 +
12 +The negative (non-gfx1100) tests verify at three layers:
13 + 1. Compile-level: on non-gfx1100 hardware, the RDNA3 ops are absent
14 + from the compiled _rocm_C extension — real binary verification.
15 + 2. Static source analysis: parses CMakeLists.txt and torch_bindings.cpp
16 + to verify that all RDNA3 .cu files and op registrations are inside
17 + gfx1100-only guards.
18 + 3. Runtime mock: patches on_gfx1100() to False and verifies that the
19 + Python dispatch chain rejects the RDNA3 path.
20 +
21 +Run `pytest tests/kernels/quantization/test_rdna3_compile_guards.py`.
22 +"""
23 +
24 +from pathlib import Path
25 +from unittest.mock import patch
26 +
27 +import pytest
28 +import regex as re
29 +import torch
30 +
31 +import vllm
32 +from vllm.platforms import current_platform
33 +
34 +if not current_platform.is_rocm():
35 + pytest.skip("RDNA3 compile-guard tests are ROCm-only", allow_module_level=True)
36 +
37 +from vllm.platforms.rocm import on_gfx1100 # noqa: E402
38 +
39 +gfx1100_only = pytest.mark.skipif(
40 + not on_gfx1100(),
41 + reason="Requires gfx1100 hardware",
42 +)
43 +
44 +not_gfx1100 = pytest.mark.skipif(
45 + on_gfx1100(),
46 + reason="This test verifies non-gfx1100 builds — skip on gfx1100",
47 +)
48 +
49 +RDNA3_OPS = ["gptq_gemm_rdna3", "gptq_gemm_rdna3_wmma", "moe_gptq_gemm_rdna3"]
50 +RDNA3_CU_FILES = [
51 + "q_gemm_rdna3.cu",
52 + "q_gemm_rdna3_wmma.cu",
53 + "moe_q_gemm_rdna3.cu",
54 +]
55 +
56 +
57 +def _find_repo_root() -> Path | None:
58 + """Walk up from this file to find the repo root (has CMakeLists.txt)."""
59 + for parent in [Path(__file__).resolve(), *Path(__file__).resolve().parents]:
60 + if (parent / "CMakeLists.txt").exists() and (parent / "csrc").is_dir():
61 + return parent
62 + return None
63 +
64 +
65 +REPO_ROOT = _find_repo_root()
66 +
67 +# Directory of the *installed* vllm python package. The .py guard checks read
68 +# from here so they verify the code that is actually imported at runtime — this
69 +# works even on CI images that ship the wheel instead of the python source tree
70 +# (where only csrc/ + CMakeLists.txt are checked out for building).
71 +VLLM_PKG_DIR: Path | None = (
72 + Path(vllm.__file__).parent if getattr(vllm, "__file__", None) else None
73 +)
74 +
75 +needs_source = pytest.mark.skipif(
76 + REPO_ROOT is None,
77 + reason="C/CMake source tree not available (installed package only)",
78 +)
79 +
80 +
81 +def _read_source_or_skip(*relparts: str) -> str:
82 + """Read a C/CMake source file from the repo tree, or skip if absent.
83 +
84 + Used for csrc/ and CMakeLists.txt — these only exist in a source checkout,
85 + not in the installed wheel.
86 + """
87 + assert REPO_ROOT is not None # callers are gated by @needs_source
88 + path = REPO_ROOT.joinpath(*relparts)
89 + if not path.exists():
90 + pytest.skip(f"{path} not present in this source tree")
91 + return path.read_text()
92 +
93 +
94 +def _read_pkg_source_or_skip(*relparts: str) -> str:
95 + """Read a python source file from the installed vllm package.
96 +
97 + Reflects the code actually loaded at runtime, so these guard checks run in
98 + CI against the wheel — no source checkout required. Only skips for an
99 + exotic install layout (namespace/zipimport) where __file__ is unavailable.
100 + """
101 + if VLLM_PKG_DIR is None:
102 + pytest.skip("vllm package directory not resolvable (zip/namespace?)")
103 + assert VLLM_PKG_DIR is not None # narrow for mypy (skip above is NoReturn)
104 + path = VLLM_PKG_DIR.joinpath(*relparts)
105 + if not path.exists():
106 + pytest.skip(f"{path} not present in installed vllm package")
107 + return path.read_text()
108 +
109 +
110 +# ============================================================================
111 +# Part A: POSITIVE — on gfx1100, ops exist and dispatch works
112 +# ============================================================================
113 +
114 +
115 +@gfx1100_only
116 +@pytest.mark.parametrize("op_name", RDNA3_OPS)
117 +def test_op_registered_on_gfx1100(op_name):
118 + """On gfx1100, all RDNA3 ops must be registered in _rocm_C."""
119 + assert hasattr(torch.ops, "_rocm_C"), "_rocm_C module not loaded"
120 + assert hasattr(torch.ops._rocm_C, op_name), (
121 + f"_rocm_C.{op_name} not registered — "
122 + "check CMakeLists.txt VLLM_ROCM_HAS_GFX1100 "
123 + "and torch_bindings.cpp #ifdef VLLM_ROCM_GFX1100"
124 + )
125 +
126 +
127 +@gfx1100_only
128 +def test_all_ops_present_or_all_absent():
129 + """The 3 RDNA3 ops are behind the same #ifdef — all present or all absent.
130 +
131 + Catches someone accidentally moving an op outside the guard.
132 + """
133 + has_rocm_c = hasattr(torch.ops, "_rocm_C")
134 + if not has_rocm_c:
135 + pytest.skip("_rocm_C not loaded")
136 +
137 + present = {op: hasattr(torch.ops._rocm_C, op) for op in RDNA3_OPS}
138 + values = set(present.values())
139 + assert len(values) == 1, (
140 + f"Guard inconsistency — some RDNA3 ops registered, others not: "
141 + f"{present}. Check torch_bindings.cpp #ifdef VLLM_ROCM_GFX1100 block."
142 + )
143 +
144 +
145 +# ============================================================================
146 +# Part B: NEGATIVE — compile-level verification on non-gfx1100
147 +# ============================================================================
148 +
149 +
150 +@not_gfx1100
151 +@pytest.mark.parametrize("op_name", RDNA3_OPS)
152 +def test_op_absent_on_non_gfx1100(op_name):
153 + """On non-gfx1100 (CDNA), RDNA3 ops must NOT exist in _rocm_C.
154 +
155 + This is the real compile-level check: the binary was built without
156 + gfx1100 support, so the ops should not have been compiled or registered.
157 + """
158 + if not hasattr(torch.ops, "_rocm_C"):
159 + return
160 + assert not hasattr(torch.ops._rocm_C, op_name), (
161 + f"_rocm_C.{op_name} is registered on non-gfx1100 hardware — "
162 + "compile guard is broken: check CMakeLists.txt "
163 + "VLLM_ROCM_HAS_GFX1100 and torch_bindings.cpp #ifdef"
164 + )
165 +
166 +
167 +@not_gfx1100
168 +def test_rocm_moe_not_supported_on_non_gfx1100():
169 + """rocm_moe.is_supported() must return False on non-gfx1100 hardware."""
170 + from vllm.model_executor.layers.quantization.compressed_tensors.compressed_tensors_moe import ( # noqa: E501
171 + rocm_moe,
172 + )
173 +
174 + wq = type("WQ", (), {"num_bits": 4})()
175 + assert rocm_moe.is_supported(wq) is False, (
176 + "rocm_moe.is_supported() returned True on non-gfx1100 — "
177 + "dispatch guard is broken"
178 + )
179 +
180 +
181 +@not_gfx1100
182 +def test_dense_kernel_rejects_on_non_gfx1100():
183 + """RDNA3W4A16LinearKernel.can_implement must reject on non-gfx1100."""
184 + from vllm.model_executor.kernels.linear.mixed_precision.MPLinearKernel import ( # noqa: E501
185 + MPLinearLayerConfig,
186 + )
187 + from vllm.model_executor.kernels.linear.mixed_precision.rdna3_w4a16 import ( # noqa: E501
188 + RDNA3W4A16LinearKernel,
189 + )
190 + from vllm.scalar_type import scalar_types
191 +
192 + config = MPLinearLayerConfig(
193 + full_weight_shape=(1024, 256),
194 + partition_weight_shape=(1024, 256),
195 + weight_type=scalar_types.uint4b8,
196 + act_type=torch.float16,
197 + group_size=128,
198 + zero_points=False,
199 + has_g_idx=False,
200 + )
201 + ok, reason = RDNA3W4A16LinearKernel.can_implement(config)
202 + assert ok is False, f"RDNA3 dense kernel accepted on non-gfx1100: {reason}"
203 +
204 +
205 +# ============================================================================
206 +# Part C: Static source analysis (build-level guards)
207 +# ============================================================================
208 +
209 +
210 +@needs_source
211 +class TestCMakeGuards:
212 + """Verify CMakeLists.txt only compiles RDNA3 .cu files for gfx1100."""
213 +
214 + @staticmethod
215 + def _read_cmake():
216 + return _read_source_or_skip("CMakeLists.txt")
217 +
218 + def test_rdna3_cu_files_inside_gfx1100_conditional(self):
219 + """All RDNA3 .cu files must be listed inside the
220 + ``if(VLLM_GPU_ARCHES MATCHES "gfx1100")`` block, not unconditionally.
221 + """
222 + cmake = self._read_cmake()
223 + for cu_file in RDNA3_CU_FILES:
224 + assert cu_file in cmake, f"{cu_file} not found in CMakeLists.txt"
225 +
226 + lines = cmake.splitlines()
227 + in_gfx1100_block = False
228 + for line in lines:
229 + if 'VLLM_GPU_ARCHES MATCHES "gfx1100"' in line:
230 + in_gfx1100_block = True
231 + if in_gfx1100_block and "endif()" in line:
232 + in_gfx1100_block = False
233 + if cu_file in line:
234 + assert in_gfx1100_block, (
235 + f"{cu_file} is listed OUTSIDE the gfx1100 "
236 + f"conditional in CMakeLists.txt — CDNA builds "
237 + f"would compile RDNA3 code. Line: {line.strip()}"
238 + )
239 +
240 + def test_compile_definition_only_for_gfx1100(self):
241 + """VLLM_ROCM_GFX1100 compile definition must be conditional."""
242 + cmake = self._read_cmake()
243 + lines = cmake.splitlines()
244 + in_gfx1100_block = False
245 + for line in lines:
246 + if "VLLM_ROCM_HAS_GFX1100)" in line:
247 + in_gfx1100_block = True
248 + if in_gfx1100_block and "endif()" in line:
249 + in_gfx1100_block = False
250 + if "VLLM_ROCM_GFX1100" in line and "target_compile_definitions" in line:
251 + assert in_gfx1100_block, (
252 + "VLLM_ROCM_GFX1100 compile definition is set outside "
253 + "the VLLM_ROCM_HAS_GFX1100 conditional — CDNA builds "
254 + f"would define it. Line: {line.strip()}"
255 + )
256 +
257 +
258 +@needs_source
259 +class TestTorchBindingsGuards:
260 + """Verify torch_bindings.cpp gates all RDNA3 ops behind #ifdef."""
261 +
262 + @staticmethod
263 + def _read_bindings():
264 + return _read_source_or_skip("csrc", "rocm", "torch_bindings.cpp")
265 +
266 + def test_all_rdna3_ops_inside_ifdef(self):
267 + """Every rdna3 op def/impl must be between #ifdef VLLM_ROCM_GFX1100
268 + and #endif. If any is outside, a CDNA build would try to register
269 + the op and link a symbol that doesn't exist.
270 + """
271 + src = self._read_bindings()
272 + lines = src.splitlines()
273 +
274 + inside_guard = False
275 + rdna3_lines_outside = []
276 +
277 + for i, line in enumerate(lines, 1):
278 + if "#ifdef VLLM_ROCM_GFX1100" in line:
279 + inside_guard = True
280 + elif line.strip() == "#endif" and inside_guard:
281 + inside_guard = False
282 +
283 + if (
284 + "rdna3" in line.lower()
285 + and not line.strip().startswith("//")
286 + and not inside_guard
287 + ):
288 + rdna3_lines_outside.append((i, line.strip()))
289 +
290 + assert not rdna3_lines_outside, (
291 + "RDNA3 op references found OUTSIDE #ifdef VLLM_ROCM_GFX1100 "
292 + "in torch_bindings.cpp — these would break CDNA builds:\n"
293 + + "\n".join(f" L{n}: {s}" for n, s in rdna3_lines_outside)
294 + )
295 +
296 + def test_no_unconditional_rdna3_includes(self):
297 + """No #include of RDNA3-specific headers outside the guard."""
298 + src = self._read_bindings()
299 + lines = src.splitlines()
300 +
301 + inside_guard = False
302 + for i, line in enumerate(lines, 1):
303 + if "#ifdef VLLM_ROCM_GFX1100" in line:
304 + inside_guard = True
305 + elif line.strip() == "#endif" and inside_guard:
306 + inside_guard = False
307 +
308 + if "#include" in line and "rdna3" in line.lower():
309 + assert inside_guard, (
310 + f"L{i}: RDNA3 include outside gfx1100 guard: {line.strip()}"
311 + )
312 +
313 +
314 +class TestCustomOpsGuards:
315 + """Verify _custom_ops.py gates register_fake behind hasattr checks."""
316 +
317 + @staticmethod
318 + def _read_custom_ops():
319 + return _read_pkg_source_or_skip("_custom_ops.py")
320 +
321 + def test_register_fake_guarded_by_hasattr(self):
322 + """Every register_fake for an RDNA3 op must be preceded by a hasattr
323 + check — otherwise it would crash on import on CDNA where the ops
324 + don't exist.
325 + """
326 + src = self._read_custom_ops()
327 + for op in RDNA3_OPS:
328 + pattern = rf'register_fake\(\s*"_rocm_C::{op}"\s*\)'
329 + match = re.search(pattern, src)
330 + if match is None:
331 + continue
332 + preceding = src[: match.start()]
333 + last_hasattr = preceding.rfind(f'hasattr(torch.ops._rocm_C, "{op}")')
334 + assert last_hasattr != -1, (
335 + f'register_fake("_rocm_C::{op}") is not preceded by a '
336 + f"hasattr check — would crash on CDNA import"
337 + )
338 + gap = preceding[last_hasattr:].count("\n")
339 + assert gap <= 5, (
340 + f"hasattr guard for {op} is {gap} lines before "
341 + f"register_fake — suspiciously far; verify it's the "
342 + f"actual guard and not a coincidence"
343 + )
344 +
345 + def test_no_toplevel_rocm_c_import(self):
346 + """No top-level ``from vllm._rocm_C import`` — would crash on CDNA."""
347 + src = self._read_custom_ops()
348 + for line in src.splitlines():
349 + stripped = line.strip()
350 + if stripped.startswith("#") or stripped.startswith("//"):
351 + continue
352 + assert "from vllm._rocm_C import" not in stripped, (
353 + f"Top-level import of _rocm_C in _custom_ops.py would "
354 + f"crash on CDNA: {stripped}"
355 + )
356 +
357 +
358 +# ============================================================================
359 +# Part D: Runtime mock (simulate CDNA on gfx1100 hardware)
360 +# ============================================================================
361 +
362 +
363 +class _FakeWeightQuant:
364 + """Minimal stand-in for a weight quantization config."""
365 +
366 + def __init__(self, num_bits):
367 + self.num_bits = num_bits
368 +
369 +
370 +class TestMoEDispatchMocked:
371 + """Mock on_gfx1100() to False and verify RDNA3 MoE is unreachable."""
372 +
373 + def test_is_supported_false_when_mocked_cdna(self):
374 + """rocm_moe.is_supported() must return False when not on gfx1100."""
375 + from vllm.model_executor.layers.quantization.compressed_tensors.compressed_tensors_moe import ( # noqa: E501
376 + rocm_moe,
377 + )
378 +
379 + with patch("vllm.platforms.rocm.on_gfx1100", return_value=False):
380 + assert rocm_moe.is_supported(_FakeWeightQuant(num_bits=4)) is False
381 +
382 + @pytest.mark.parametrize("num_bits", [2, 3, 8, 16])
383 + def test_is_supported_rejects_non_w4(self, num_bits):
384 + """is_supported() rejects non-4-bit even before checking arch."""
385 + from vllm.model_executor.layers.quantization.compressed_tensors.compressed_tensors_moe import ( # noqa: E501
386 + rocm_moe,
387 + )
388 +
389 + assert rocm_moe.is_supported(_FakeWeightQuant(num_bits=num_bits)) is False
390 +
391 + def test_is_supported_false_when_op_missing(self):
392 + """is_supported() returns False when the C++ op doesn't exist."""
393 + from vllm.model_executor.layers.quantization.compressed_tensors.compressed_tensors_moe import ( # noqa: E501
394 + rocm_moe,
395 + )
396 +
397 + fake_rocm_c = type("FakeRocmC", (), {"gptq_gemm_rdna3": None})()
398 + with patch.object(torch, "ops", create=True) as mock_ops:
399 + mock_ops._rocm_C = fake_rocm_c
400 + assert rocm_moe.is_supported(_FakeWeightQuant(num_bits=4)) is False
401 +
402 + def test_is_supported_false_when_rocm_c_absent(self):
403 + """is_supported() returns False when _rocm_C doesn't exist at all."""
404 + from vllm.model_executor.layers.quantization.compressed_tensors.compressed_tensors_moe import ( # noqa: E501
405 + rocm_moe,
406 + )
407 +
408 + fake_ops = type("FakeOps", (), {})()
409 + with patch.object(torch, "ops", fake_ops):
410 + assert rocm_moe.is_supported(_FakeWeightQuant(num_bits=4)) is False
411 +
412 +
413 +class TestDenseKernelSelectionMocked:
414 + """Mock on_gfx1100() and verify dense RDNA3 kernel is not selected."""
415 +
416 + @gfx1100_only
417 + def test_can_implement_rejects_when_mocked_cdna(self):
418 + """RDNA3W4A16LinearKernel.can_implement must reject on mocked CDNA."""
419 + from vllm.model_executor.kernels.linear.mixed_precision.MPLinearKernel import ( # noqa: E501
420 + MPLinearLayerConfig,
421 + )
422 + from vllm.model_executor.kernels.linear.mixed_precision.rdna3_w4a16 import ( # noqa: E501
423 + RDNA3W4A16LinearKernel,
424 + )
425 + from vllm.scalar_type import scalar_types
426 +
427 + config = MPLinearLayerConfig(
428 + full_weight_shape=(1024, 256),
429 + partition_weight_shape=(1024, 256),
430 + weight_type=scalar_types.uint4b8,
431 + act_type=torch.float16,
432 + group_size=128,
433 + zero_points=False,
434 + has_g_idx=False,
435 + )
436 + ok, _ = RDNA3W4A16LinearKernel.can_implement(config)
437 + assert ok is True
438 +
439 + with (
440 + patch("vllm.platforms.rocm.on_gfx1100", return_value=False),
441 + patch("vllm.platforms.rocm._ON_GFX1100", False),
442 + ):
443 + ok, reason = RDNA3W4A16LinearKernel.can_implement(config)
444 + assert ok is False, f"RDNA3 kernel accepted on simulated CDNA: {reason}"
445 +
446 + @gfx1100_only
447 + def test_chooser_skips_rdna3_when_mocked_cdna(self):
448 + """choose_mp_linear_kernel must NOT return RDNA3 on mocked CDNA."""
449 + from vllm.model_executor.kernels.linear import (
450 + choose_mp_linear_kernel,
451 + )
452 + from vllm.model_executor.kernels.linear.mixed_precision.MPLinearKernel import ( # noqa: E501
453 + MPLinearLayerConfig,
454 + )
455 + from vllm.scalar_type import scalar_types
456 +
457 + config = MPLinearLayerConfig(
458 + full_weight_shape=(1024, 256),
459 + partition_weight_shape=(1024, 256),
460 + weight_type=scalar_types.uint4b8,
461 + act_type=torch.float16,
462 + group_size=128,
463 + zero_points=False,
464 + has_g_idx=False,
465 + )
466 + with (
467 + patch("vllm.platforms.rocm.on_gfx1100", return_value=False),
468 + patch("vllm.platforms.rocm._ON_GFX1100", False),
469 + ):
470 + chosen = choose_mp_linear_kernel(config)
471 + assert chosen.__name__ != "RDNA3W4A16LinearKernel", (
472 + "RDNA3 kernel was selected on simulated CDNA — "
473 + "choose_mp_linear_kernel guard is broken"
474 + )
475 +
476 +
477 +class TestCompressedTensorsMoEDispatchGuard:
478 + """Verify compressed_tensors_moe.py only enters rocm_moe under is_rocm()."""
479 +
480 + def test_rocm_guard_in_dispatch_source(self):
481 + """The rocm_moe import and call must be inside an is_rocm() check."""
482 + src = _read_pkg_source_or_skip(
483 + "model_executor",
484 + "layers",
485 + "quantization",
486 + "compressed_tensors",
487 + "compressed_tensors_moe",
488 + "compressed_tensors_moe.py",
489 + )
490 + lines = src.splitlines()
491 +
492 + for i, line in enumerate(lines, 1):
493 + stripped = line.strip()
494 + if "rocm_moe" in stripped and not stripped.startswith("#"):
495 + found_guard = False
496 + for j in range(i - 1, max(0, i - 15), -1):
497 + if "is_rocm()" in lines[j - 1]:
498 + found_guard = True
499 + break
500 + assert found_guard, (
501 + f"L{i}: rocm_moe reference not protected by "
502 + f"is_rocm() guard: {stripped}"
503 + )
tests/kernels/quantization/test_rdna3_moe_w4a16.py new
+367
@@ -0,0 +1,367 @@
1 +#!/usr/bin/env python3
2 +# SPDX-License-Identifier: Apache-2.0
3 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
4 +"""Correctness tests for the ROCm RDNA3 fused MoE W4A16 HIP kernel (gfx1100).
5 +
6 +Tests ``moe_gptq_gemm_rdna3`` against the dense ``gptq_gemm_rdna3`` as
7 +reference: builds RDNA3-format weights (shuffled int32, synthesized qzeros),
8 +runs the fused MoE kernel, and compares per-expert results.
9 +
10 +Model parameters taken from:
11 + - cyankiwi/Qwen3-30B-A3B-Instruct-2507-AWQ-4bit
12 + (hidden=2048, inter=768, E=128, top_k=8, G=32)
13 + - Qwen3.6-35B-A3B-GPTQ-W4A16-G32
14 + (hidden=2048, inter=512, E=256, top_k=8, G=32)
15 +
16 +Run `pytest tests/kernels/quantization/test_rdna3_moe_w4a16.py`.
17 +"""
18 +
19 +import pytest
20 +import torch
21 +
22 +from vllm.platforms import current_platform
23 +
24 +if not current_platform.is_rocm():
25 + pytest.skip("RDNA3 MoE W4A16 kernel is ROCm-only", allow_module_level=True)
26 +
27 +from vllm import _custom_ops as ops # noqa: E402
28 +from vllm.model_executor.layers.fused_moe.activation import ( # noqa: E402
29 + MoEActivation,
30 + apply_moe_activation,
31 +)
32 +from vllm.model_executor.layers.fused_moe.moe_align_block_size import ( # noqa: E402
33 + moe_align_block_size,
34 +)
35 +from vllm.model_executor.layers.quantization.utils.quant_utils import ( # noqa: E402
36 + pack_quantized_values_into_int32,
37 +)
38 +from vllm.platforms.rocm import on_gfx1100 # noqa: E402
39 +from vllm.scalar_type import scalar_types # noqa: E402
40 +
41 +device = "cuda"
42 +
43 +gfx1100_only = pytest.mark.skipif(
44 + not (
45 + on_gfx1100()
46 + and hasattr(torch.ops, "_rocm_C")
47 + and hasattr(torch.ops._rocm_C, "moe_gptq_gemm_rdna3")
48 + ),
49 + reason="Requires gfx1100 with moe_gptq_gemm_rdna3 op",
50 +)
51 +
52 +# Model configurations: real K/N/top_k/group_size dims, E capped at 16 to
53 +# fit in test GPU memory (full E=128/256 would need >20GB for weights alone).
54 +# Kernel behavior is E-independent (per-expert tiling), so E=16 is sufficient.
55 +MODEL_CONFIGS = [
56 + # cyankiwi/Qwen3-30B-A3B-Instruct-2507-AWQ-4bit dims (E capped)
57 + pytest.param(16, 2048, 768, 8, 32, id="Qwen3-30B-A3B"),
58 + # Qwen3.6-35B-A3B-GPTQ-W4A16-G32 dims (E capped)
59 + pytest.param(16, 2048, 512, 8, 32, id="Qwen3.6-35B-A3B"),
60 +]
61 +
62 +# Token counts: decode (1), small batch (4), medium (16), prefill (64)
63 +NUM_TOKENS = [1, 4, 16, 64, 256, 512]
64 +
65 +
66 +def _make_packed_weights(E, K, N):
67 + """Create random 4-bit packed weights [E, K/8, N] int32 + shuffle."""
68 + w = torch.randint(0, 16, (E, K, N), dtype=torch.int32, device=device)
69 + packed = torch.zeros(E, K // 8, N, dtype=torch.int32, device=device)
70 + for i in range(8):
71 + packed |= (w[:, i::8, :] & 0xF) << (i * 4)
72 + g_idx = torch.empty(0, dtype=torch.int32, device=device)
73 + for e in range(E):
74 + we = packed[e].contiguous()
75 + ops.gptq_shuffle(we, g_idx, 4)
76 + packed[e] = we
77 + return packed
78 +
79 +
80 +def _make_scales(E, groups, N, dtype):
81 + return torch.rand(E, groups, N, dtype=dtype, device=device) * 0.1
82 +
83 +
84 +def _make_qzeros(E, groups, N):
85 + zeros = torch.full(
86 + (groups, N),
87 + scalar_types.uint4b8.bias - 1,
88 + dtype=torch.int32,
89 + device=device,
90 + )
91 + qz = pack_quantized_values_into_int32(
92 + zeros,
93 + scalar_types.uint4b8,
94 + packed_dim=1,
95 + )
96 + return qz.unsqueeze(0).expand(E, -1, -1).contiguous()
97 +
98 +
99 +@gfx1100_only
100 +@pytest.mark.parametrize("E, K, N_inter, top_k, group_size", MODEL_CONFIGS)
101 +@pytest.mark.parametrize("M", NUM_TOKENS)
102 +@pytest.mark.parametrize("dtype", [torch.float16, torch.bfloat16])
103 +@pytest.mark.parametrize("block_size_m", [1, 4])
104 +def test_fused_moe_w1_matches_dense(
105 + E, K, N_inter, top_k, group_size, M, dtype, block_size_m
106 +):
107 + """w1 GEMM via fused kernel matches per-expert dense kernel."""
108 + N_gate_up = N_inter * 2
109 + groups = K // group_size
110 +
111 + torch.manual_seed(42)
112 + x = torch.randn(M, K, dtype=dtype, device=device)
113 + w13 = _make_packed_weights(E, K, N_gate_up)
114 + w13_s = _make_scales(E, groups, N_gate_up, dtype)
115 + w13_z = _make_qzeros(E, groups, N_gate_up)
116 + g_idx = torch.empty(0, dtype=torch.int32, device=device)
117 +
118 + topk_ids = torch.randint(0, E, (M, top_k), device=device, dtype=torch.int32)
119 + si, ei, ntp = moe_align_block_size(topk_ids, block_size_m, E)
120 +
121 + # Fused kernel
122 + fused_out = torch.zeros(M * top_k, N_gate_up, dtype=dtype, device=device)
123 + ops.moe_gptq_gemm_rdna3(
124 + x,
125 + fused_out,
126 + w13,
127 + w13_s,
128 + w13_z,
129 + torch.empty(0, device=device),
130 + si,
131 + ei,
132 + ntp,
133 + top_k,
134 + block_size_m,
135 + False,
136 + 0,
137 + )
138 +
139 + # Per-expert dense reference
140 + ref_out = torch.zeros(M * top_k, N_gate_up, dtype=dtype, device=device)
141 + for m in range(M):
142 + for k in range(top_k):
143 + e = topk_ids[m, k].item()
144 + flat = m * top_k + k
145 + ref = ops.gptq_gemm_rdna3(
146 + x[m : m + 1],
147 + w13[e],
148 + w13_z[e],
149 + w13_s[e],
150 + g_idx,
151 + False,
152 + )
153 + ref_out[flat] = ref.squeeze()
154 +
155 + # Split-K atomics can cause minor fp16/bf16 rounding differences
156 + # at large K (e.g. K=2048 → 8 K-blocks). Use allclose, not equal.
157 + atol = 0.5 if dtype == torch.bfloat16 else 0.1
158 + assert torch.allclose(fused_out, ref_out, atol=atol, rtol=0.01), (
159 + f"max diff: {(fused_out - ref_out).abs().max().item()}"
160 + )
161 +
162 +
163 +@gfx1100_only
164 +@pytest.mark.parametrize("E, K, N_inter, top_k, group_size", MODEL_CONFIGS)
165 +@pytest.mark.parametrize("M", NUM_TOKENS)
166 +@pytest.mark.parametrize("dtype", [torch.float16, torch.bfloat16])
167 +def test_fused_moe_output_topk_reduces(E, K, N_inter, top_k, group_size, M, dtype):
168 + """output_topk fuses moe_sum: multiple experts write to same output row."""
169 + groups = K // group_size
170 +
171 + torch.manual_seed(123)
172 + x = torch.randn(M * top_k, K, dtype=dtype, device=device)
173 + w = _make_packed_weights(E, K, N_inter)
174 + ws = _make_scales(E, groups, N_inter, dtype)
175 + wz = _make_qzeros(E, groups, N_inter)
176 +
177 + topk_ids = torch.randint(0, E, (M, top_k), device=device, dtype=torch.int32)
178 + topk_w = torch.softmax(
179 + torch.randn(M, top_k, device=device),
180 + dim=-1,
181 + ).float()
182 +
183 + si, ei, ntp = moe_align_block_size(topk_ids, 1, E)
184 +
185 + # Without output_topk: write to [M*top_k, N] then moe_sum
186 + flat_out = torch.zeros(M * top_k, N_inter, dtype=dtype, device=device)
187 + ops.moe_gptq_gemm_rdna3(
188 + x,
189 + flat_out,
190 + w,
191 + ws,
192 + wz,
193 + topk_w.view(-1),
194 + si,
195 + ei,
196 + ntp,
197 + 1,
198 + 1,
199 + True,
200 + 0,
201 + )
202 + ref = torch.zeros(M, N_inter, dtype=dtype, device=device)
203 + ops.moe_sum(flat_out.view(M, top_k, N_inter), ref)
204 +
205 + # With output_topk: write directly to [M, N]
206 + fused = torch.zeros(M, N_inter, dtype=dtype, device=device)
207 + ops.moe_gptq_gemm_rdna3(
208 + x,
209 + fused,
210 + w,
211 + ws,
212 + wz,
213 + topk_w.view(-1),
214 + si,
215 + ei,
216 + ntp,
217 + 1,
218 + 1,
219 + True,
220 + top_k,
221 + )
222 +
223 + atol = 1.0 if dtype == torch.bfloat16 else 0.1
224 + assert torch.allclose(fused, ref, atol=atol, rtol=0.01), (
225 + f"max diff: {(fused - ref).abs().max().item()}"
226 + )
227 +
228 +
229 +@gfx1100_only
230 +@pytest.mark.parametrize("E, K, N_inter, top_k, group_size", MODEL_CONFIGS)
231 +@pytest.mark.parametrize("M", NUM_TOKENS)
232 +@pytest.mark.parametrize("dtype", [torch.float16, torch.bfloat16])
233 +def test_full_moe_e2e(E, K, N_inter, top_k, group_size, M, dtype):
234 + """Full MoE forward: w1 + silu_and_mul + w2 with output_topk reduce."""
235 + N_gate_up = N_inter * 2
236 + hidden = K
237 +
238 + torch.manual_seed(7)
239 + x = torch.randn(M, K, dtype=dtype, device=device)
240 + w13 = _make_packed_weights(E, K, N_gate_up)
241 + w13_s = _make_scales(E, K // group_size, N_gate_up, dtype)
242 + w13_z = _make_qzeros(E, K // group_size, N_gate_up)
243 + w2 = _make_packed_weights(E, N_inter, hidden)
244 + w2_s = _make_scales(E, N_inter // group_size, hidden, dtype)
245 + w2_z = _make_qzeros(E, N_inter // group_size, hidden)
246 + g_idx = torch.empty(0, dtype=torch.int32, device=device)
247 +
248 + topk_ids = torch.randint(0, E, (M, top_k), device=device, dtype=torch.int32)
249 + topk_w = torch.softmax(
250 + torch.randn(M, top_k, device=device),
251 + dim=-1,
252 + ).float()
253 +
254 + si, ei, ntp = moe_align_block_size(topk_ids, 1, E)
255 +
256 + # Fused path (what apply() does)
257 + w1_out = torch.zeros(M * top_k, N_gate_up, dtype=dtype, device=device)
258 + ops.moe_gptq_gemm_rdna3(
259 + x,
260 + w1_out,
261 + w13,
262 + w13_s,
263 + w13_z,
264 + torch.empty(0, device=device),
265 + si,
266 + ei,
267 + ntp,
268 + top_k,
269 + 1,
270 + False,
271 + 0,
272 + )
273 + act_out = torch.empty(M * top_k, N_inter, dtype=dtype, device=device)
274 + apply_moe_activation(MoEActivation.SILU, act_out, w1_out)
275 + fused = torch.zeros(M, hidden, dtype=dtype, device=device)
276 + ops.moe_gptq_gemm_rdna3(
277 + act_out,
278 + fused,
279 + w2,
280 + w2_s,
281 + w2_z,
282 + topk_w.view(-1),
283 + si,
284 + ei,
285 + ntp,
286 + 1,
287 + 1,
288 + True,
289 + top_k,
290 + )
291 +
292 + # Per-expert reference
293 + ref = torch.zeros(M, hidden, dtype=dtype, device=device)
294 + for m_idx in range(M):
295 + for k_idx in range(top_k):
296 + e = topk_ids[m_idx, k_idx].item()
297 + w = topk_w[m_idx, k_idx].item()
298 + r1 = ops.gptq_gemm_rdna3(
299 + x[m_idx : m_idx + 1],
300 + w13[e],
301 + w13_z[e],
302 + w13_s[e],
303 + g_idx,
304 + False,
305 + )
306 + a = torch.empty(1, N_inter, dtype=dtype, device=device)
307 + apply_moe_activation(MoEActivation.SILU, a, r1)
308 + r2 = ops.gptq_gemm_rdna3(
309 + a,
310 + w2[e],
311 + w2_z[e],
312 + w2_s[e],
313 + g_idx,
314 + False,
315 + )
316 + ref[m_idx] += r2.squeeze() * w
317 +
318 + # E2E chains w1 + activation + w2 + topk_w + output_topk reduce.
319 + # Each step accumulates rounding error (split-K atomics, topk_w
320 + # multiply order). Use relative L2 norm like the dense kernel test.
321 + diff_l2 = torch.norm(fused.float() - ref.float())
322 + ref_l2 = torch.norm(ref.float())
323 + rel_l2 = (diff_l2 / ref_l2).item() if ref_l2 > 0 else 0.0
324 + threshold = 0.05 if dtype == torch.float16 else 0.10
325 + assert rel_l2 < threshold, (
326 + f"rel L2 = {rel_l2:.4f} (threshold {threshold}), "
327 + f"max abs diff: {(fused - ref).abs().max().item()}"
328 + )
329 +
330 +
331 +@gfx1100_only
332 +def test_expert_id_minus_one():
333 + """Kernel handles expert_id == -1 (expert parallelism) without crash."""
334 + # Qwen3-30B-A3B dims (E capped for memory)
335 + E, K, N = 16, 2048, 768
336 + groups = K // 32
337 +
338 + w = _make_packed_weights(E, K, N)
339 + ws = _make_scales(E, groups, N, torch.bfloat16)
340 + wz = _make_qzeros(E, groups, N)
341 + x = torch.randn(1, K, dtype=torch.bfloat16, device=device)
342 +
343 + # Manually create sorted_token_ids/expert_ids with -1
344 + sorted_ids = torch.tensor([0], dtype=torch.int32, device=device)
345 + expert_ids = torch.tensor([-1], dtype=torch.int32, device=device)
346 + ntp = torch.tensor([1], dtype=torch.int32, device=device)
347 +
348 + out = torch.zeros(1, N, dtype=torch.bfloat16, device=device)
349 + ops.moe_gptq_gemm_rdna3(
350 + x,
351 + out,
352 + w,
353 + ws,
354 + wz,
355 + torch.empty(0, device=device),
356 + sorted_ids,
357 + expert_ids,
358 + ntp,
359 + 1,
360 + 1,
361 + False,
362 + 0,
363 + )
364 + current_platform.synchronize()
365 +
366 + # Output should remain zero (expert skipped)
367 + assert torch.equal(out, torch.zeros_like(out))
vllm/_custom_ops.py
+53
@@ -695,6 +695,59 @@ if hasattr(torch.ops, "_rocm_C") and hasattr(torch.ops._rocm_C, "gptq_gemm_rdna3
695 )
696
697
698 +def moe_gptq_gemm_rdna3(
699 + a: torch.Tensor,
700 + c: torch.Tensor,
701 + b_q_weight: torch.Tensor,
702 + b_scales: torch.Tensor,
703 + b_qzeros: torch.Tensor,
704 + topk_weights: torch.Tensor,
705 + sorted_token_ids: torch.Tensor,
706 + expert_ids: torch.Tensor,
707 + num_tokens_post_padded: torch.Tensor,
708 + top_k: int,
709 + block_size_m: int,
710 + mul_topk_weight: bool,
711 + output_topk: int = 0,
712 +) -> None:
713 + torch.ops._rocm_C.moe_gptq_gemm_rdna3(
714 + a,
715 + c,
716 + b_q_weight,
717 + b_scales,
718 + b_qzeros,
719 + topk_weights,
720 + sorted_token_ids,
721 + expert_ids,
722 + num_tokens_post_padded,
723 + top_k,
724 + block_size_m,
725 + mul_topk_weight,
726 + output_topk,
727 + )
728 +
729 +
730 +if hasattr(torch.ops, "_rocm_C") and hasattr(torch.ops._rocm_C, "moe_gptq_gemm_rdna3"):
731 +
732 + @register_fake("_rocm_C::moe_gptq_gemm_rdna3")
733 + def _moe_gptq_gemm_rdna3_fake(
734 + a: torch.Tensor,
735 + c: torch.Tensor,
736 + b_q_weight: torch.Tensor,
737 + b_scales: torch.Tensor,
738 + b_qzeros: torch.Tensor,
739 + topk_weights: torch.Tensor,
740 + sorted_token_ids: torch.Tensor,
741 + expert_ids: torch.Tensor,
742 + num_tokens_post_padded: torch.Tensor,
743 + top_k: int,
744 + block_size_m: int,
745 + mul_topk_weight: bool,
746 + output_topk: int = 0,
747 + ) -> None:
748 + return
749 +
750 +
751 if hasattr(torch.ops._C, "allspark_w8a16_gemm"):
752
753 @register_fake("_C::allspark_w8a16_gemm")
vllm/model_executor/layers/fused_moe/layer.py
+1
@@ -892,6 +892,7 @@ class FusedMoE(PluggableLayer):
892 if quant_method_name in (
893 "CompressedTensorsWNA16MarlinMoEMethod",
894 "CompressedTensorsWNA16MoEMethod",
895 + "CompressedTensorsWNA16RDNA3MoEMethod",
896 ):
897 if is_transposed:
898 loaded_weight = loaded_weight.t().contiguous()
vllm/model_executor/layers/quantization/compressed_tensors/compressed_tensors_moe/compressed_tensors_moe.py
+14 -4
@@ -98,10 +98,6 @@ class CompressedTensorsMoEMethod(FusedMoEMethodBase):
98 not check_moe_marlin_supports_layer(layer, group_size)
99 or current_platform.is_rocm()
100 ):
101 - from .compressed_tensors_moe_wna16 import (
102 - CompressedTensorsWNA16MoEMethod,
103 - )
104 -
101 if (
102 weight_quant.strategy == QuantizationStrategy.GROUP
103 and weight_quant.actorder
@@ -110,6 +106,20 @@ class CompressedTensorsMoEMethod(FusedMoEMethodBase):
106 raise ValueError(
107 "WNA16MoE is not supported with actorder=group/dynamic."
108 )
109 +
110 + # Native ROCm HIP kernels (RDNA3, etc.)
111 + if current_platform.is_rocm():
112 + from . import rocm_moe
113 +
114 + if rocm_moe.is_supported(weight_quant):
115 + return rocm_moe.make_method(
116 + weight_quant, input_quant, layer.moe_config
117 + )
118 +
119 + from .compressed_tensors_moe_wna16 import (
120 + CompressedTensorsWNA16MoEMethod,
121 + )
122 +
123 logger.info_once("Using CompressedTensorsWNA16MoEMethod")
124 return CompressedTensorsWNA16MoEMethod(
125 weight_quant, input_quant, layer.moe_config
vllm/model_executor/layers/quantization/compressed_tensors/compressed_tensors_moe/compressed_tensors_moe_wna16_rdna3.py new
+261
@@ -0,0 +1,261 @@
1 +# SPDX-License-Identifier: Apache-2.0
2 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 +"""CompressedTensors MoE W4A16 using the fused RDNA3 (gfx1100) HIP kernel.
4 +
5 +Uses ``moe_gptq_gemm_rdna3`` — a single HIP kernel launch per GEMM that
6 +handles expert routing + W4A16 dequant + dot product with atomic output.
7 +
8 +Weight format (per expert, same as dense RDNA3 W4A16):
9 + - Packed int32 ``[E, K/8, N]`` with exllama shuffle
10 + - Scales ``[E, groups, N]`` in activation dtype
11 + - Zero points ``[E, groups, N/8]`` packed int32 (synthesized)
12 +"""
13 +
14 +import torch
15 +
16 +from vllm import _custom_ops as ops
17 +from vllm.logger import init_logger
18 +from vllm.model_executor.layers.fused_moe import (
19 + RoutedExperts,
20 + SharedExperts,
21 +)
22 +from vllm.model_executor.layers.fused_moe.activation import (
23 + MoEActivation,
24 + apply_moe_activation,
25 +)
26 +from vllm.model_executor.layers.fused_moe.moe_align_block_size import (
27 + moe_align_block_size,
28 +)
29 +from vllm.model_executor.layers.quantization.compressed_tensors.compressed_tensors_moe.compressed_tensors_moe_wna16 import ( # noqa: E501
30 + CompressedTensorsWNA16MoEMethod,
31 +)
32 +from vllm.model_executor.layers.quantization.utils.quant_utils import (
33 + pack_quantized_values_into_int32,
34 +)
35 +from vllm.scalar_type import scalar_types
36 +
37 +logger = init_logger(__name__)
38 +
39 +
40 +def _synthesize_qzeros(
41 + groups: int, out_features: int, device: torch.device
42 +) -> torch.Tensor:
43 + """Create packed zero-point tensor for symmetric quant.
44 +
45 + GPTQv1 +1 quirk: kernel adds 1 to stored zeros, so encode
46 + (bias - 1) = 7 for uint4b8 (bias=8).
47 + """
48 + zeros = torch.full(
49 + (groups, out_features),
50 + scalar_types.uint4b8.bias - 1,
51 + dtype=torch.int32,
52 + device=device,
53 + )
54 + return pack_quantized_values_into_int32(zeros, scalar_types.uint4b8, packed_dim=1)
55 +
56 +
57 +class CompressedTensorsWNA16RDNA3MoEMethod(CompressedTensorsWNA16MoEMethod):
58 + """W4A16 MoE using the fused RDNA3 HIP kernel (moe_gptq_gemm_rdna3).
59 +
60 + Weights are in RDNA3 format (shuffled int32 [E, K/8, N]),
61 + NOT Triton format (transposed uint8). apply() dispatches through
62 + the fused HIP kernel directly.
63 + """
64 +
65 + def process_weights_after_loading(self, layer: torch.nn.Module) -> None:
66 + device = layer.w13_weight_packed.device
67 + num_experts = layer.w13_weight_packed.shape[0]
68 + empty_g_idx = torch.empty(0, dtype=torch.int32, device=device)
69 +
70 + # Shuffle weights in-place per expert (exllama nibble interleave)
71 + for e in range(num_experts):
72 + w13_e = layer.w13_weight_packed.data[e].contiguous()
73 + ops.gptq_shuffle(w13_e, empty_g_idx, 4)
74 + layer.w13_weight_packed.data[e] = w13_e
75 + w2_e = layer.w2_weight_packed.data[e].contiguous()
76 + ops.gptq_shuffle(w2_e, empty_g_idx, 4)
77 + layer.w2_weight_packed.data[e] = w2_e
78 +
79 + # Keep scales as [E, groups, N] in activation dtype
80 + act_dtype = layer.w13_weight_scale.dtype
81 + layer.w13_weight_scale = torch.nn.Parameter(
82 + layer.w13_weight_scale.to(dtype=act_dtype).contiguous(),
83 + requires_grad=False,
84 + )
85 + layer.w2_weight_scale = torch.nn.Parameter(
86 + layer.w2_weight_scale.to(dtype=act_dtype).contiguous(),
87 + requires_grad=False,
88 + )
89 +
90 + # Synthesize packed zero points: [E, groups, N/8] int32
91 + w13_groups = (layer.w13_weight_packed.shape[1] * 8) // self.group_size
92 + w13_N = layer.w13_weight_packed.shape[2]
93 + w2_groups = (layer.w2_weight_packed.shape[1] * 8) // self.group_size
94 + w2_N = layer.w2_weight_packed.shape[2]
95 +
96 + w13_qz = _synthesize_qzeros(w13_groups, w13_N, device)
97 + w2_qz = _synthesize_qzeros(w2_groups, w2_N, device)
98 + layer.w13_qzeros = torch.nn.Parameter(
99 + w13_qz.unsqueeze(0).expand(num_experts, -1, -1).contiguous(),
100 + requires_grad=False,
101 + )
102 + layer.w2_qzeros = torch.nn.Parameter(
103 + w2_qz.unsqueeze(0).expand(num_experts, -1, -1).contiguous(),
104 + requires_grad=False,
105 + )
106 +
107 + # Pre-allocate reusable buffers for decode (sizes based on top_k=8)
108 + N_gate_up = w13_N
109 + hidden_size = w2_N
110 + intermediate = N_gate_up // 2 # gated activation
111 + # Max tokens we expect in decode; prefill will re-allocate if needed
112 + max_decode_tokens = 16
113 + top_k = 8 # conservative default
114 + buf_size = max_decode_tokens * top_k
115 + layer.rdna3_w1_buf = torch.zeros(
116 + buf_size, N_gate_up, dtype=act_dtype, device=device
117 + )
118 + layer.rdna3_act_buf = torch.empty(
119 + buf_size, intermediate, dtype=act_dtype, device=device
120 + )
121 + layer.rdna3_out_buf = torch.zeros(
122 + max_decode_tokens, hidden_size, dtype=act_dtype, device=device
123 + )
124 + layer.rdna3_empty_tw = torch.empty(0, device=device)
125 +
126 + def apply(
127 + self,
128 + layer: RoutedExperts,
129 + x: torch.Tensor,
130 + topk_weights: torch.Tensor,
131 + topk_ids: torch.Tensor,
132 + shared_experts: SharedExperts | None,
133 + shared_experts_input: torch.Tensor | None,
134 + ) -> torch.Tensor:
135 + activation = (
136 + layer.activation
137 + if isinstance(layer.activation, MoEActivation)
138 + else MoEActivation.from_str(layer.activation)
139 + )
140 + return _rdna3_fused_moe(
141 + x,
142 + topk_weights,
143 + topk_ids,
144 + layer=layer,
145 + activation=activation,
146 + apply_router_weight_on_input=(layer.apply_router_weight_on_input),
147 + global_num_experts=layer.global_num_experts,
148 + expert_map=layer.expert_map,
149 + )
150 +
151 +
152 +def _rdna3_fused_moe(
153 + hidden_states: torch.Tensor,
154 + topk_weights: torch.Tensor,
155 + topk_ids: torch.Tensor,
156 + layer: RoutedExperts,
157 + activation: MoEActivation,
158 + apply_router_weight_on_input: bool,
159 + global_num_experts: int,
160 + expert_map: torch.Tensor | None,
161 +) -> torch.Tensor:
162 + """Fused MoE forward using the RDNA3 W4A16 HIP kernel.
163 +
164 + Optimizations vs naive dispatch:
165 + - BLOCK_SIZE_M=1 for decode (no padding waste, bf16 fast path)
166 + - Pre-allocated buffers (no torch.zeros per call)
167 + - Inline token sorting for small M (skip moe_align_block_size)
168 + - moe_sum fused into output accumulation
169 + """
170 + num_tokens = hidden_states.shape[0]
171 + top_k = topk_ids.shape[1]
172 + total_tokens = num_tokens * top_k
173 + N_gate_up = layer.w13_weight_packed.shape[2]
174 + hidden_size = layer.w2_weight_packed.shape[2]
175 + dtype = hidden_states.dtype
176 + device = hidden_states.device
177 +
178 + intermediate_size = N_gate_up // 2 if activation.is_gated else N_gate_up
179 +
180 + if global_num_experts <= 0:
181 + global_num_experts = layer.w13_weight_packed.shape[0]
182 +
183 + # BLOCK_SIZE_M=1 for decode (small M), 4 for prefill
184 + block_size_m = 1 if num_tokens <= 4 else 4
185 +
186 + # --- Token routing ---
187 + sorted_token_ids, expert_ids, num_tokens_post_padded = moe_align_block_size(
188 + topk_ids,
189 + block_size_m,
190 + global_num_experts,
191 + expert_map,
192 + )
193 +
194 + # --- Reuse pre-allocated buffers when possible ---
195 + if total_tokens <= layer.rdna3_w1_buf.shape[0]:
196 + w1_out = layer.rdna3_w1_buf[:total_tokens]
197 + w1_out.zero_()
198 + act_out = layer.rdna3_act_buf[:total_tokens]
199 + else:
200 + w1_out = torch.zeros(
201 + total_tokens,
202 + N_gate_up,
203 + dtype=dtype,
204 + device=device,
205 + )
206 + act_out = torch.empty(
207 + total_tokens,
208 + intermediate_size,
209 + dtype=dtype,
210 + device=device,
211 + )
212 +
213 + # --- topk weights (pre-cast to float32 for kernel) ---
214 + topk_w_float = topk_weights.view(-1).float()
215 + empty_tw = layer.rdna3_empty_tw
216 +
217 + # --- w1 GEMM: [M, K] -> [M*top_k, N_gate_up] ---
218 + ops.moe_gptq_gemm_rdna3(
219 + hidden_states,
220 + w1_out,
221 + layer.w13_weight_packed,
222 + layer.w13_weight_scale,
223 + layer.w13_qzeros,
224 + topk_w_float if apply_router_weight_on_input else empty_tw,
225 + sorted_token_ids,
226 + expert_ids,
227 + num_tokens_post_padded,
228 + top_k,
229 + block_size_m,
230 + apply_router_weight_on_input,
231 + )
232 +
233 + # --- Activation (silu_and_mul etc.) ---
234 + apply_moe_activation(activation, act_out, w1_out)
235 +
236 + # --- w2 GEMM: [M*top_k, intermediate] -> [M, hidden] (fused reduce) ---
237 + # output_topk=top_k: kernel writes to out[token_id / top_k] directly,
238 + # fusing moe_sum into the atomic accumulation — saves one kernel launch
239 + # and the w2_out intermediate buffer.
240 + out = torch.zeros(
241 + num_tokens,
242 + hidden_size,
243 + dtype=dtype,
244 + device=device,
245 + )
246 + ops.moe_gptq_gemm_rdna3(
247 + act_out,
248 + out,
249 + layer.w2_weight_packed,
250 + layer.w2_weight_scale,
251 + layer.w2_qzeros,
252 + topk_w_float if not apply_router_weight_on_input else empty_tw,
253 + sorted_token_ids,
254 + expert_ids,
255 + num_tokens_post_padded,
256 + 1,
257 + block_size_m,
258 + not apply_router_weight_on_input,
259 + output_topk=top_k,
260 + )
261 + return out
vllm/model_executor/layers/quantization/compressed_tensors/compressed_tensors_moe/rocm_moe.py new
+48
@@ -0,0 +1,48 @@
1 +# SPDX-License-Identifier: Apache-2.0
2 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 +"""ROCm MoE kernel dispatcher.
4 +
5 +Selects architecture-specific native HIP MoE kernels in priority order.
6 +Falls back to the Triton WNA16 path when no native kernel is available.
7 +"""
8 +
9 +import torch
10 +
11 +from vllm.logger import init_logger
12 +
13 +logger = init_logger(__name__)
14 +
15 +
16 +def is_supported(weight_quant) -> bool:
17 + """Check if a native ROCm MoE kernel is available for this config."""
18 + if weight_quant.num_bits != 4:
19 + return False
20 +
21 + from vllm.platforms.rocm import on_gfx1100
22 +
23 + # RDNA3 (gfx1100). Future: add RDNA4 (gfx12x), CDNA (gfx94x), etc.
24 + return (
25 + on_gfx1100()
26 + and hasattr(torch.ops, "_rocm_C")
27 + and hasattr(torch.ops._rocm_C, "moe_gptq_gemm_rdna3")
28 + )
29 +
30 +
31 +def make_method(weight_quant, input_quant, moe_config):
32 + """Create the native ROCm MoE method. Call only after is_supported()."""
33 + from vllm.platforms.rocm import on_gfx1100
34 +
35 + if on_gfx1100():
36 + from .compressed_tensors_moe_wna16_rdna3 import (
37 + CompressedTensorsWNA16RDNA3MoEMethod,
38 + )
39 +
40 + logger.info_once(
41 + "Using CompressedTensorsWNA16RDNA3MoEMethod (native RDNA3 HIP kernel)"
42 + )
43 + return CompressedTensorsWNA16RDNA3MoEMethod(
44 + weight_quant, input_quant, moe_config
45 + )
46 +
47 + # Future: RDNA4, CDNA, etc.
48 + raise RuntimeError("is_supported() returned True but no kernel matched")