From 062b05ff3af405b8800f7c42fa3998cc03dd97eb Mon Sep 17 00:00:00 2001 From: JartX Date: Sat, 6 Jun 2026 22:30:39 +0200 Subject: [PATCH] [ROCm][Perf] Fused MoE W4A16 HIP kernel for AMD RDNA3 (gfx1100) (#44075) Signed-off-by: JartX Co-authored-by: TJian --- CMakeLists.txt | 3 +- csrc/rocm/moe_q_gemm_rdna3.cu | 639 ++++++++++++++++++ csrc/rocm/ops.h | 9 + csrc/rocm/torch_bindings.cpp | 9 + .../quantization/test_rdna3_compile_guards.py | 503 ++++++++++++++ .../quantization/test_rdna3_moe_w4a16.py | 367 ++++++++++ vllm/_custom_ops.py | 53 ++ vllm/model_executor/layers/fused_moe/layer.py | 1 + .../compressed_tensors_moe.py | 18 +- .../compressed_tensors_moe_wna16_rdna3.py | 261 +++++++ .../compressed_tensors_moe/rocm_moe.py | 48 ++ 11 files changed, 1906 insertions(+), 5 deletions(-) create mode 100644 csrc/rocm/moe_q_gemm_rdna3.cu create mode 100644 tests/kernels/quantization/test_rdna3_compile_guards.py create mode 100644 tests/kernels/quantization/test_rdna3_moe_w4a16.py create mode 100644 vllm/model_executor/layers/quantization/compressed_tensors/compressed_tensors_moe/compressed_tensors_moe_wna16_rdna3.py create mode 100644 vllm/model_executor/layers/quantization/compressed_tensors/compressed_tensors_moe/rocm_moe.py diff --git a/CMakeLists.txt b/CMakeLists.txt index 0b919e2314..529ba55c51 100644 --- a/CMakeLists.txt +++ b/CMakeLists.txt @@ -1310,7 +1310,8 @@ if(VLLM_GPU_LANG STREQUAL "HIP") set(VLLM_ROCM_HAS_GFX1100 ON) list(APPEND VLLM_ROCM_EXT_SRC "csrc/rocm/q_gemm_rdna3.cu" - "csrc/rocm/q_gemm_rdna3_wmma.cu") + "csrc/rocm/q_gemm_rdna3_wmma.cu" + "csrc/rocm/moe_q_gemm_rdna3.cu") endif() define_extension_target( diff --git a/csrc/rocm/moe_q_gemm_rdna3.cu b/csrc/rocm/moe_q_gemm_rdna3.cu new file mode 100644 index 0000000000..6c25ed7e4b --- /dev/null +++ b/csrc/rocm/moe_q_gemm_rdna3.cu @@ -0,0 +1,639 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright contributors to the vLLM project +// +// Fused MoE W4A16 GPTQ kernel for RDNA3 (gfx1100). +// +// Combines expert routing (sorted_token_ids / expert_ids) with the RDNA3 +// W4A16 dequant+dot from q_gemm_rdna3.cu into a single kernel launch. +// Each block processes BLOCK_SIZE_M tokens assigned to one expert, covering +// a tile of N output columns and K input positions. +// +// Weight format: same as the dense kernel — [E, K/8, N] uint32 shuffled, +// [E, groups, N] scales, [E, groups, N/8] packed zeros. +// +// Design: THREADS_X=256 (8 waves on wave32), BLOCK_KN_SIZE=256, each thread +// handles 4 N columns. Output via 64-bit packed CAS atomic-add directly to +// the pre-zeroed output tensor (no FP32 scratch buffer). + +#include + +#include +#include +#include + +#include +#include +#include + +#include "qdq_4_rdna3.cuh" + +#if defined(__HIPCC__) && defined(__gfx1100__) + #define __HIP__RDNA3__ +#endif + +namespace vllm { +namespace moe_gptq_rdna3 { + +#define BLOCK_KN_SIZE 256 +#define THREADS_X 256 + +#if defined(__HIP__RDNA3__) || !defined(__HIP_DEVICE_COMPILE__) + +using gptq_rdna3::bf162_t; +using gptq_rdna3::bf16_t; + +// --- Helpers (same as q_gemm_rdna3.cu) --- + +template +__forceinline__ __device__ T tzero(); + +template <> +__forceinline__ __device__ half tzero() { + return __float2half_rn(0.0f); +} + +template <> +__forceinline__ __device__ bf16_t tzero() { + return __float2bfloat16(0.0f); +} + +__forceinline__ __device__ float dot22_8_f(half2 (&dq)[4], const half* a_ptr) { + float result = 0.0f; + const half2* a2_ptr = (const half2*)a_ptr; + #pragma unroll + for (int i = 0; i < 4; i++) { + result = __builtin_amdgcn_fdot2(dq[i], *a2_ptr++, result, /*clamp=*/false); + } + return result; +} + +__forceinline__ __device__ float dot22_8_f(float (&dq)[8], + const bf16_t* a_ptr) { + float result = 0.0f; + #pragma unroll + for (int i = 0; i < 4; i++) { + uint32_t aw; + __builtin_memcpy(&aw, a_ptr + 2 * i, sizeof(uint32_t)); + float a_x = __uint_as_float((aw & 0xFFFFu) << 16); + float a_y = __uint_as_float(aw & 0xFFFF0000u); + result = __fmaf_rn(dq[2 * i + 0], a_x, result); + result = __fmaf_rn(dq[2 * i + 1], a_y, result); + } + return result; +} + +__forceinline__ __device__ void atomic_add_pk4_f16(half* addr, half2 v01, + half2 v23) { + unsigned long long* addr_u = reinterpret_cast(addr); + unsigned long long old = *addr_u; + while (true) { + union { + unsigned long long u; + half2 h2[2]; + } cur, sum; + cur.u = old; + sum.h2[0] = __hadd2(cur.h2[0], v01); + sum.h2[1] = __hadd2(cur.h2[1], v23); + unsigned long long prev = atomicCAS(addr_u, old, sum.u); + if (prev == old) break; + old = prev; + } +} + +__forceinline__ __device__ void atomic_add_pk4_bf16(bf16_t* addr, bf162_t v01, + bf162_t v23) { + unsigned long long* addr_u = reinterpret_cast(addr); + unsigned long long old = *addr_u; + while (true) { + union { + unsigned long long u; + bf162_t b2[2]; + } cur, sum; + cur.u = old; + sum.b2[0] = __hadd2(cur.b2[0], v01); + sum.b2[1] = __hadd2(cur.b2[1], v23); + unsigned long long prev = atomicCAS(addr_u, old, sum.u); + if (prev == old) break; + old = prev; + } +} + +__forceinline__ __device__ void load4_zeros(const uint32_t* qzeros_row, int n, + int (&zeros)[4]) { + int qcol = n / 8; + int shift = (n & 0x07) * 4; + uint32_t d = qzeros_row[qcol] >> shift; + zeros[0] = (int)(d & 0xF); + zeros[1] = (int)((d >> 4) & 0xF); + zeros[2] = (int)((d >> 8) & 0xF); + zeros[3] = (int)((d >> 12) & 0xF); +} + +template +__forceinline__ __device__ void load4_scales(const T* scales_row, int n, + T (&scales)[4]) { + scales[0] = scales_row[n + 0]; + scales[1] = scales_row[n + 1]; + scales[2] = scales_row[n + 2]; + scales[3] = scales_row[n + 3]; +} + +// --------------------------------------------------------------------------- +// Fused MoE kernel. +// --------------------------------------------------------------------------- + +template +__global__ void moe_gemm_q4_kernel_rdna3( + const T* __restrict__ a, // [size_m, size_k] or [M*topk, K] + T* __restrict__ c, // [M*topk, size_n] pre-zeroed + const uint32_t* __restrict__ b_q_weight, // [E, K/8, N] packed + const T* __restrict__ b_scales, // [E, groups, N] + const uint32_t* __restrict__ b_qzeros, // [E, groups, N/8] packed + const float* __restrict__ topk_weights, // [M*topk] or nullptr + const int32_t* __restrict__ sorted_token_ids, + const int32_t* __restrict__ expert_ids, + const int32_t* __restrict__ num_tokens_post_padded, + const int size_m, // total tokens (original M, or M*topk for w2) + const int size_n, // output features per expert + const int size_k, // input features + const int groups, // K / group_size + const int top_k, // routing top-k (1 for w2 pass) + // Per-expert strides (in elements, not bytes) + const int expert_weight_stride, // (K/8) * N + const int expert_scales_stride, // groups * N + const int expert_zeros_stride, // groups * (N/8) + const bool mul_topk_weight, + const int output_topk) { // >0: reduce output by token_id/output_topk + const int t = threadIdx.x; + const int token_block = blockIdx.x; + const int offset_n = blockIdx.y * BLOCK_KN_SIZE * 4; + const int offset_k = blockIdx.z * BLOCK_KN_SIZE; + const int end_k = min(offset_k + BLOCK_KN_SIZE, size_k); + const int n = offset_n + t * 4; + + // Early exit for padding blocks or invalid experts (expert_map = -1) + if (token_block * BLOCK_SIZE_M >= num_tokens_post_padded[0]) return; + + const int expert_id = expert_ids[token_block]; + if (expert_id == -1) return; + + // Expert-specific pointers + const uint32_t* expert_weights = + b_q_weight + (int64_t)expert_id * expert_weight_stride; + const T* expert_scales = b_scales + (int64_t)expert_id * expert_scales_stride; + const uint32_t* expert_qzeros = + b_qzeros + (int64_t)expert_id * expert_zeros_stride; + + // LDS for activations + constexpr int LDS_PAD = 8; + __shared__ T block_a[BLOCK_SIZE_M][BLOCK_KN_SIZE + LDS_PAD]; + + static_assert(BLOCK_KN_SIZE == THREADS_X, + "BLOCK_KN_SIZE must equal THREADS_X"); + + // For bf16 M=1, we can skip LDS and read A from global (same as dense). + // fp16 always needs LDS due to the dot22_8_f indexing pattern. + constexpr bool USE_LDS_A = (BLOCK_SIZE_M > 1) || std::is_same::value; + + const int offset_m_base = token_block * BLOCK_SIZE_M; + + if constexpr (USE_LDS_A) { + if (offset_k + t < end_k) { + #pragma unroll + for (int m = 0; m < BLOCK_SIZE_M; ++m) { + int32_t token_id = sorted_token_ids[offset_m_base + m]; + int token_row = token_id / top_k; + T av; + if (token_row < size_m) { + av = a[(int64_t)token_row * size_k + offset_k + t]; + } else { + av = tzero(); + } + block_a[m][t] = av; + } + } + __syncthreads(); + } + + if (n >= size_n) return; + + // Group bookkeeping + const int groupsize = size_k / groups; + int group = offset_k / groupsize; + int nextgroup = (group + 1) * groupsize; + + // Weight pointer for this expert + int qk = offset_k / 8; + const uint32_t* b_ptr = expert_weights + qk * size_n + n; + + // Per-column dequant constants (4 columns per thread) + half2 z1z16_h[4][2], y1y16_h[4][2]; + float z_b_f[4], y_b_f[4]; + + // GPTQv1: zero_offset = 1 + constexpr int zero_offset = 1; + + auto refresh_group = [&](int g) { + const uint32_t* qz_row = expert_qzeros + g * (size_n / 8); + const T* sc_row = expert_scales + g * size_n; + int zeros[4]; + T scales[4]; + load4_zeros(qz_row, n, zeros); + load4_scales(sc_row, n, scales); + if constexpr (std::is_same::value) { + #pragma unroll + for (int i = 0; i < 4; ++i) { + gptq_rdna3::prep_zero_scale_fp16((uint32_t)(zeros[i] + zero_offset), + scales[i], z1z16_h[i], y1y16_h[i]); + } + } else { + #pragma unroll + for (int i = 0; i < 4; ++i) { + gptq_rdna3::prep_zero_scale_bf16_f32((uint32_t)(zeros[i] + zero_offset), + scales[i], z_b_f[i], y_b_f[i]); + } + } + }; + + refresh_group(group); + + float block_c[BLOCK_SIZE_M][4]; + #pragma unroll + for (int m = 0; m < BLOCK_SIZE_M; ++m) { + #pragma unroll + for (int j = 0; j < 4; ++j) block_c[m][j] = 0.0f; + } + + // --- Main K-loop --- + int k = offset_k; + while (k < end_k) { + if (k == nextgroup) { + group++; + nextgroup += groupsize; + refresh_group(group); + } + + // Prefetch 4 weight words (128 bytes) + int4 b_w[4]; + #pragma unroll + for (int j = 0; j < 4; ++j) { + b_w[j] = *(const int4*)(b_ptr + j * size_n); + } + b_ptr += 4 * size_n; + + #pragma unroll + for (int j = 0; j < 4; ++j) { + const int a_off = (k - offset_k) + 8 * j; + + if constexpr (std::is_same::value) { + // fp16 path: dequant via bit-trick, dot via v_dot2_f32_f16 + half2 dq[4][4]; + gptq_rdna3::dequant_4bit_8_fp16((uint32_t)b_w[j].x, dq[0], z1z16_h[0], + y1y16_h[0]); + gptq_rdna3::dequant_4bit_8_fp16((uint32_t)b_w[j].y, dq[1], z1z16_h[1], + y1y16_h[1]); + gptq_rdna3::dequant_4bit_8_fp16((uint32_t)b_w[j].z, dq[2], z1z16_h[2], + y1y16_h[2]); + gptq_rdna3::dequant_4bit_8_fp16((uint32_t)b_w[j].w, dq[3], z1z16_h[3], + y1y16_h[3]); + + #pragma unroll + for (int m = 0; m < BLOCK_SIZE_M; ++m) { + const half* a_ptr = reinterpret_cast(&block_a[m][a_off]); + block_c[m][0] += dot22_8_f(dq[0], a_ptr); + block_c[m][1] += dot22_8_f(dq[1], a_ptr); + block_c[m][2] += dot22_8_f(dq[2], a_ptr); + block_c[m][3] += dot22_8_f(dq[3], a_ptr); + } + } else if constexpr (BLOCK_SIZE_M == 1) { + // bf16 M=1: v_dot2_f32_bf16 with InstCombine-defeating opacity + typedef short __attribute__((ext_vector_type(2))) bf16x2_t; + constexpr uint32_t BF16_MAGIC = 0x43004300u; + constexpr uint32_t BF16_ONES = 0x3F803F80u; + union pack4 { + float f[4]; + uint32_t u[4]; + }; + + uint32_t w[4]; + __builtin_memcpy(w, &b_w[j], sizeof(int4)); + + // Load activations — read from global (no LDS for bf16 M=1) + pack4 a_pack; + { + int32_t token_id = sorted_token_ids[offset_m_base]; + int token_row = token_id / top_k; + if (token_row < size_m) { + const uint32_t* a_words = reinterpret_cast( + a + (int64_t)token_row * size_k + offset_k + a_off); + a_pack.u[0] = a_words[0]; + a_pack.u[1] = a_words[1]; + a_pack.u[2] = a_words[2]; + a_pack.u[3] = a_words[3]; + } else { + a_pack.u[0] = 0; + a_pack.u[1] = 0; + a_pack.u[2] = 0; + a_pack.u[3] = 0; + } + } + + // sum_a for bias correction + float sum_a = 0.0f; + #pragma unroll + for (int b = 0; b < 4; ++b) { + sum_a = __builtin_amdgcn_fdot2_f32_bf16( + *((bf16x2_t*)(&a_pack.f[b])), *((const bf16x2_t*)&BF16_ONES), + sum_a, /*clamp=*/false); + } + + #pragma unroll 1 + for (int col = 0; col < 4; ++col) { + pack4 q_pack; + const uint32_t qa = w[col]; + q_pack.u[0] = ((qa >> 0) & 0x000F000Fu) | BF16_MAGIC; + q_pack.u[1] = ((qa >> 4) & 0x000F000Fu) | BF16_MAGIC; + q_pack.u[2] = ((qa >> 8) & 0x000F000Fu) | BF16_MAGIC; + q_pack.u[3] = ((qa >> 12) & 0x000F000Fu) | BF16_MAGIC; + + float partial = 0.0f; + #pragma unroll + for (int b = 0; b < 4; ++b) { + partial = __builtin_amdgcn_fdot2_f32_bf16( + *((bf16x2_t*)(&a_pack.f[b])), *((bf16x2_t*)(&q_pack.f[b])), + partial, /*clamp=*/false); + } + + block_c[0][col] = + __fmaf_rn(y_b_f[col], partial, + __fmaf_rn(z_b_f[col], sum_a, block_c[0][col])); + } + } else { + // bf16 M>1: v_dot2_f32_bf16 with LDS-staged activations + typedef short __attribute__((ext_vector_type(2))) bf16x2_t; + constexpr uint32_t BF16_MAGIC = 0x43004300u; + constexpr uint32_t BF16_ONES = 0x3F803F80u; + union pack4 { + float f[4]; + uint32_t u[4]; + }; + + uint32_t w[4]; + __builtin_memcpy(w, &b_w[j], sizeof(int4)); + + pack4 a_pack[BLOCK_SIZE_M]; + #pragma unroll + for (int m = 0; m < BLOCK_SIZE_M; ++m) { + const uint32_t* a_words = + reinterpret_cast(&block_a[m][a_off]); + a_pack[m].u[0] = a_words[0]; + a_pack[m].u[1] = a_words[1]; + a_pack[m].u[2] = a_words[2]; + a_pack[m].u[3] = a_words[3]; + } + + float sum_a[BLOCK_SIZE_M]; + #pragma unroll + for (int m = 0; m < BLOCK_SIZE_M; ++m) { + float s = 0.0f; + #pragma unroll + for (int b = 0; b < 4; ++b) { + s = __builtin_amdgcn_fdot2_f32_bf16(*((bf16x2_t*)(&a_pack[m].f[b])), + *((const bf16x2_t*)&BF16_ONES), + s, /*clamp=*/false); + } + sum_a[m] = s; + } + + #pragma unroll 1 + for (int col = 0; col < 4; ++col) { + pack4 q_pack; + const uint32_t qa = w[col]; + q_pack.u[0] = ((qa >> 0) & 0x000F000Fu) | BF16_MAGIC; + q_pack.u[1] = ((qa >> 4) & 0x000F000Fu) | BF16_MAGIC; + q_pack.u[2] = ((qa >> 8) & 0x000F000Fu) | BF16_MAGIC; + q_pack.u[3] = ((qa >> 12) & 0x000F000Fu) | BF16_MAGIC; + + #pragma unroll + for (int m = 0; m < BLOCK_SIZE_M; ++m) { + float partial = 0.0f; + #pragma unroll + for (int b = 0; b < 4; ++b) { + partial = __builtin_amdgcn_fdot2_f32_bf16( + *((bf16x2_t*)(&a_pack[m].f[b])), *((bf16x2_t*)(&q_pack.f[b])), + partial, /*clamp=*/false); + } + block_c[m][col] = + __fmaf_rn(y_b_f[col], partial, + __fmaf_rn(z_b_f[col], sum_a[m], block_c[m][col])); + } + } + } + } + k += 32; + } + + // --- Epilogue: apply topk_weight and atomic-add to output --- + #pragma unroll + for (int m = 0; m < BLOCK_SIZE_M; ++m) { + int32_t token_id = sorted_token_ids[offset_m_base + m]; + if (token_id / top_k >= size_m) continue; + + // Apply router weight + if (mul_topk_weight && topk_weights != nullptr) { + float tw = topk_weights[token_id]; + #pragma unroll + for (int j = 0; j < 4; ++j) block_c[m][j] *= tw; + } + + // output_topk > 0: reduce by mapping token_id back to original token + // (multiple experts write to the same row via atomics) + int64_t out_row = (output_topk > 0) ? (int64_t)(token_id / output_topk) + : (int64_t)token_id; + T* out = c + out_row * size_n + n; + if constexpr (std::is_same::value) { + half2 r01 = __halves2half2(__float2half_rn(block_c[m][0]), + __float2half_rn(block_c[m][1])); + half2 r23 = __halves2half2(__float2half_rn(block_c[m][2]), + __float2half_rn(block_c[m][3])); + atomic_add_pk4_f16(out, r01, r23); + } else { + bf162_t r01; + r01.x = __float2bfloat16(block_c[m][0]); + r01.y = __float2bfloat16(block_c[m][1]); + bf162_t r23; + r23.x = __float2bfloat16(block_c[m][2]); + r23.y = __float2bfloat16(block_c[m][3]); + atomic_add_pk4_bf16(out, r01, r23); + } + } +} + +#else // non-RDNA3: empty stub for symbol parity + +template +__global__ void moe_gemm_q4_kernel_rdna3( + const T*, T*, const uint32_t*, const T*, const uint32_t*, const float*, + const int32_t*, const int32_t*, const int32_t*, const int, const int, + const int, const int, const int, const int, const int, const int, + const bool, const int) {} + +#endif // __HIP__RDNA3__ || !__HIP_DEVICE_COMPILE__ + +// --------------------------------------------------------------------------- +// Launcher +// --------------------------------------------------------------------------- + +template +void launch_moe_gemm_q4( + const T* a, T* c, const uint32_t* b_q_weight, const T* b_scales, + const uint32_t* b_qzeros, const float* topk_weights, + const int32_t* sorted_token_ids, const int32_t* expert_ids, + const int32_t* num_tokens_post_padded, int num_token_blocks, int size_m, + int size_n, int size_k, int groups, int top_k, int expert_weight_stride, + int expert_scales_stride, int expert_zeros_stride, bool mul_topk_weight, + int output_topk, cudaStream_t stream) { + dim3 block(THREADS_X); + dim3 grid(num_token_blocks, + (size_n + BLOCK_KN_SIZE * 4 - 1) / (BLOCK_KN_SIZE * 4), + (size_k + BLOCK_KN_SIZE - 1) / BLOCK_KN_SIZE); + + moe_gemm_q4_kernel_rdna3<<>>( + a, c, b_q_weight, b_scales, b_qzeros, topk_weights, sorted_token_ids, + expert_ids, num_tokens_post_padded, size_m, size_n, size_k, groups, top_k, + expert_weight_stride, expert_scales_stride, expert_zeros_stride, + mul_topk_weight, output_topk); +} + +template +void dispatch_moe_gemm_q4( + const T* a, T* c, const uint32_t* b_q_weight, const T* b_scales, + const uint32_t* b_qzeros, const float* topk_weights, + const int32_t* sorted_token_ids, const int32_t* expert_ids, + const int32_t* num_tokens_post_padded, int num_token_blocks, int size_m, + int size_n, int size_k, int groups, int top_k, int block_size_m, + int expert_weight_stride, int expert_scales_stride, int expert_zeros_stride, + bool mul_topk_weight, int output_topk, cudaStream_t stream) { + // Dispatch to template instantiation based on block_size_m + switch (block_size_m) { + case 1: + launch_moe_gemm_q4( + a, c, b_q_weight, b_scales, b_qzeros, topk_weights, sorted_token_ids, + expert_ids, num_tokens_post_padded, num_token_blocks, size_m, size_n, + size_k, groups, top_k, expert_weight_stride, expert_scales_stride, + expert_zeros_stride, mul_topk_weight, output_topk, stream); + break; + case 2: + launch_moe_gemm_q4( + a, c, b_q_weight, b_scales, b_qzeros, topk_weights, sorted_token_ids, + expert_ids, num_tokens_post_padded, num_token_blocks, size_m, size_n, + size_k, groups, top_k, expert_weight_stride, expert_scales_stride, + expert_zeros_stride, mul_topk_weight, output_topk, stream); + break; + case 4: + launch_moe_gemm_q4( + a, c, b_q_weight, b_scales, b_qzeros, topk_weights, sorted_token_ids, + expert_ids, num_tokens_post_padded, num_token_blocks, size_m, size_n, + size_k, groups, top_k, expert_weight_stride, expert_scales_stride, + expert_zeros_stride, mul_topk_weight, output_topk, stream); + break; + case 8: + launch_moe_gemm_q4( + a, c, b_q_weight, b_scales, b_qzeros, topk_weights, sorted_token_ids, + expert_ids, num_tokens_post_padded, num_token_blocks, size_m, size_n, + size_k, groups, top_k, expert_weight_stride, expert_scales_stride, + expert_zeros_stride, mul_topk_weight, output_topk, stream); + break; + default: + TORCH_CHECK(false, + "moe_gptq_gemm_rdna3: block_size_m must be 1, 2, 4, or 8, " + "got ", + block_size_m); + } +} + +} // namespace moe_gptq_rdna3 +} // namespace vllm + +// --------------------------------------------------------------------------- +// Public entry point +// --------------------------------------------------------------------------- +// +// Inputs: +// a [M, K] or [M*top_k, K] half or bfloat16 +// c [M*top_k, N] same dtype (pre-zeroed!) +// b_q_weight [E, K/8, N] uint32 (shuffled) +// b_scales [E, groups, N] same dtype as a +// b_qzeros [E, groups, N/8] uint32 (packed 4-bit) +// topk_weights [M*top_k] or empty float32 +// sorted_token_ids [num_blocks * block_m] int32 +// expert_ids [num_blocks] int32 +// num_tokens_post_padded [1] int32 +// top_k int +// block_size_m int (1, 2, 4, or 8) +// mul_topk_weight bool + +void moe_gptq_gemm_rdna3(torch::Tensor a, torch::Tensor c, + torch::Tensor b_q_weight, torch::Tensor b_scales, + torch::Tensor b_qzeros, torch::Tensor topk_weights, + torch::Tensor sorted_token_ids, + torch::Tensor expert_ids, + torch::Tensor num_tokens_post_padded, int64_t top_k, + int64_t block_size_m, bool mul_topk_weight, + int64_t output_topk) { + TORCH_CHECK(a.is_cuda(), "a must be a CUDA/HIP tensor"); + TORCH_CHECK(c.is_cuda(), "c must be a CUDA/HIP tensor"); + TORCH_CHECK(b_q_weight.is_cuda(), "b_q_weight must be a CUDA/HIP tensor"); + TORCH_CHECK(a.dim() == 2, "a must be 2D"); + TORCH_CHECK(c.dim() == 2, "c must be 2D"); + TORCH_CHECK(b_q_weight.dim() == 3, "b_q_weight must be 3D [E, K/8, N]"); + TORCH_CHECK(b_scales.dim() == 3, "b_scales must be 3D [E, groups, N]"); + TORCH_CHECK(b_qzeros.dim() == 3, "b_qzeros must be 3D [E, groups, N/8]"); + TORCH_CHECK( + a.scalar_type() == torch::kHalf || a.scalar_type() == torch::kBFloat16, + "a must be half or bfloat16"); + TORCH_CHECK(a.scalar_type() == b_scales.scalar_type(), + "b_scales dtype must match a"); + + const at::cuda::OptionalCUDAGuard device_guard(device_of(a)); + auto stream = at::cuda::getCurrentCUDAStream(); + + int size_m = (int)a.size(0); + int size_k = (int)a.size(1); + int size_n = (int)b_q_weight.size(2); + int groups = (int)b_scales.size(1); + + // Per-expert strides + int expert_weight_stride = (int)(b_q_weight.size(1) * b_q_weight.size(2)); + int expert_scales_stride = (int)(b_scales.size(1) * b_scales.size(2)); + int expert_zeros_stride = (int)(b_qzeros.size(1) * b_qzeros.size(2)); + + int num_token_blocks = (int)(sorted_token_ids.size(0) / block_size_m); + + const float* topk_w_ptr = + (topk_weights.numel() > 0) ? topk_weights.data_ptr() : nullptr; + + // Manual dtype dispatch using HIP native types (c10::Half/BFloat16 don't + // implicitly convert to half/__hip_bfloat16 in device code). + using vllm::gptq_rdna3::bf16_t; + + auto dispatch = [&](auto* a_ptr, auto* c_ptr, const auto* s_ptr) { + using T = std::remove_const_t>; + vllm::moe_gptq_rdna3::dispatch_moe_gemm_q4( + a_ptr, c_ptr, (const uint32_t*)b_q_weight.data_ptr(), s_ptr, + (const uint32_t*)b_qzeros.data_ptr(), topk_w_ptr, + sorted_token_ids.data_ptr(), expert_ids.data_ptr(), + num_tokens_post_padded.data_ptr(), num_token_blocks, size_m, + size_n, size_k, groups, (int)top_k, (int)block_size_m, + expert_weight_stride, expert_scales_stride, expert_zeros_stride, + mul_topk_weight, (int)output_topk, stream); + }; + + if (a.scalar_type() == torch::kHalf) { + dispatch((const half*)a.data_ptr(), (half*)c.data_ptr(), + (const half*)b_scales.data_ptr()); + } else { + dispatch((const bf16_t*)a.data_ptr(), (bf16_t*)c.data_ptr(), + (const bf16_t*)b_scales.data_ptr()); + } +} diff --git a/csrc/rocm/ops.h b/csrc/rocm/ops.h index 73197d8a5e..549d50300d 100644 --- a/csrc/rocm/ops.h +++ b/csrc/rocm/ops.h @@ -27,6 +27,15 @@ torch::Tensor gptq_gemm_rdna3_wmma(torch::Tensor a, torch::Tensor b_q_weight, torch::Tensor b_scales, torch::Tensor b_g_idx, bool use_v2_format); +void moe_gptq_gemm_rdna3(torch::Tensor a, torch::Tensor c, + torch::Tensor b_q_weight, torch::Tensor b_scales, + torch::Tensor b_qzeros, torch::Tensor topk_weights, + torch::Tensor sorted_token_ids, + torch::Tensor expert_ids, + torch::Tensor num_tokens_post_padded, int64_t top_k, + int64_t block_size_m, bool mul_topk_weight, + int64_t output_topk); + void paged_attention( torch::Tensor& out, torch::Tensor& exp_sums, torch::Tensor& max_logits, torch::Tensor& tmp_out, torch::Tensor& query, torch::Tensor& key_cache, diff --git a/csrc/rocm/torch_bindings.cpp b/csrc/rocm/torch_bindings.cpp index 1e589598c7..03de6dcd15 100644 --- a/csrc/rocm/torch_bindings.cpp +++ b/csrc/rocm/torch_bindings.cpp @@ -50,6 +50,15 @@ TORCH_LIBRARY_EXPAND(TORCH_EXTENSION_NAME, rocm_ops) { "gptq_gemm_rdna3_wmma(Tensor a, Tensor b_q_weight, Tensor b_qzeros, " "Tensor b_scales, Tensor b_g_idx, bool use_v2_format) -> Tensor"); rocm_ops.impl("gptq_gemm_rdna3_wmma", torch::kCUDA, &gptq_gemm_rdna3_wmma); + + rocm_ops.def( + "moe_gptq_gemm_rdna3(Tensor a, Tensor! c, Tensor b_q_weight, " + "Tensor b_scales, Tensor b_qzeros, Tensor topk_weights, " + "Tensor sorted_token_ids, Tensor expert_ids, " + "Tensor num_tokens_post_padded, " + "int top_k, int block_size_m, bool mul_topk_weight, " + "int output_topk) -> ()"); + rocm_ops.impl("moe_gptq_gemm_rdna3", torch::kCUDA, &moe_gptq_gemm_rdna3); #endif // Custom attention op diff --git a/tests/kernels/quantization/test_rdna3_compile_guards.py b/tests/kernels/quantization/test_rdna3_compile_guards.py new file mode 100644 index 0000000000..5dfc021fb8 --- /dev/null +++ b/tests/kernels/quantization/test_rdna3_compile_guards.py @@ -0,0 +1,503 @@ +#!/usr/bin/env python3 +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""Compile-guard tests for the ROCm RDNA3 W4A16 kernels (dense + MoE). + +Verifies that the gfx1100 compilation and dispatch guards are hermetic: + - On gfx1100: all ops exist, dispatch selects RDNA3 kernels. + - On CDNA (gfx942/gfx950) or other non-gfx1100: ops must NOT exist, + dispatch must fall through to Triton/Marlin, and no RDNA3 code + path is reachable. + +The negative (non-gfx1100) tests verify at three layers: + 1. Compile-level: on non-gfx1100 hardware, the RDNA3 ops are absent + from the compiled _rocm_C extension — real binary verification. + 2. Static source analysis: parses CMakeLists.txt and torch_bindings.cpp + to verify that all RDNA3 .cu files and op registrations are inside + gfx1100-only guards. + 3. Runtime mock: patches on_gfx1100() to False and verifies that the + Python dispatch chain rejects the RDNA3 path. + +Run `pytest tests/kernels/quantization/test_rdna3_compile_guards.py`. +""" + +from pathlib import Path +from unittest.mock import patch + +import pytest +import regex as re +import torch + +import vllm +from vllm.platforms import current_platform + +if not current_platform.is_rocm(): + pytest.skip("RDNA3 compile-guard tests are ROCm-only", allow_module_level=True) + +from vllm.platforms.rocm import on_gfx1100 # noqa: E402 + +gfx1100_only = pytest.mark.skipif( + not on_gfx1100(), + reason="Requires gfx1100 hardware", +) + +not_gfx1100 = pytest.mark.skipif( + on_gfx1100(), + reason="This test verifies non-gfx1100 builds — skip on gfx1100", +) + +RDNA3_OPS = ["gptq_gemm_rdna3", "gptq_gemm_rdna3_wmma", "moe_gptq_gemm_rdna3"] +RDNA3_CU_FILES = [ + "q_gemm_rdna3.cu", + "q_gemm_rdna3_wmma.cu", + "moe_q_gemm_rdna3.cu", +] + + +def _find_repo_root() -> Path | None: + """Walk up from this file to find the repo root (has CMakeLists.txt).""" + for parent in [Path(__file__).resolve(), *Path(__file__).resolve().parents]: + if (parent / "CMakeLists.txt").exists() and (parent / "csrc").is_dir(): + return parent + return None + + +REPO_ROOT = _find_repo_root() + +# Directory of the *installed* vllm python package. The .py guard checks read +# from here so they verify the code that is actually imported at runtime — this +# works even on CI images that ship the wheel instead of the python source tree +# (where only csrc/ + CMakeLists.txt are checked out for building). +VLLM_PKG_DIR: Path | None = ( + Path(vllm.__file__).parent if getattr(vllm, "__file__", None) else None +) + +needs_source = pytest.mark.skipif( + REPO_ROOT is None, + reason="C/CMake source tree not available (installed package only)", +) + + +def _read_source_or_skip(*relparts: str) -> str: + """Read a C/CMake source file from the repo tree, or skip if absent. + + Used for csrc/ and CMakeLists.txt — these only exist in a source checkout, + not in the installed wheel. + """ + assert REPO_ROOT is not None # callers are gated by @needs_source + path = REPO_ROOT.joinpath(*relparts) + if not path.exists(): + pytest.skip(f"{path} not present in this source tree") + return path.read_text() + + +def _read_pkg_source_or_skip(*relparts: str) -> str: + """Read a python source file from the installed vllm package. + + Reflects the code actually loaded at runtime, so these guard checks run in + CI against the wheel — no source checkout required. Only skips for an + exotic install layout (namespace/zipimport) where __file__ is unavailable. + """ + if VLLM_PKG_DIR is None: + pytest.skip("vllm package directory not resolvable (zip/namespace?)") + assert VLLM_PKG_DIR is not None # narrow for mypy (skip above is NoReturn) + path = VLLM_PKG_DIR.joinpath(*relparts) + if not path.exists(): + pytest.skip(f"{path} not present in installed vllm package") + return path.read_text() + + +# ============================================================================ +# Part A: POSITIVE — on gfx1100, ops exist and dispatch works +# ============================================================================ + + +@gfx1100_only +@pytest.mark.parametrize("op_name", RDNA3_OPS) +def test_op_registered_on_gfx1100(op_name): + """On gfx1100, all RDNA3 ops must be registered in _rocm_C.""" + assert hasattr(torch.ops, "_rocm_C"), "_rocm_C module not loaded" + assert hasattr(torch.ops._rocm_C, op_name), ( + f"_rocm_C.{op_name} not registered — " + "check CMakeLists.txt VLLM_ROCM_HAS_GFX1100 " + "and torch_bindings.cpp #ifdef VLLM_ROCM_GFX1100" + ) + + +@gfx1100_only +def test_all_ops_present_or_all_absent(): + """The 3 RDNA3 ops are behind the same #ifdef — all present or all absent. + + Catches someone accidentally moving an op outside the guard. + """ + has_rocm_c = hasattr(torch.ops, "_rocm_C") + if not has_rocm_c: + pytest.skip("_rocm_C not loaded") + + present = {op: hasattr(torch.ops._rocm_C, op) for op in RDNA3_OPS} + values = set(present.values()) + assert len(values) == 1, ( + f"Guard inconsistency — some RDNA3 ops registered, others not: " + f"{present}. Check torch_bindings.cpp #ifdef VLLM_ROCM_GFX1100 block." + ) + + +# ============================================================================ +# Part B: NEGATIVE — compile-level verification on non-gfx1100 +# ============================================================================ + + +@not_gfx1100 +@pytest.mark.parametrize("op_name", RDNA3_OPS) +def test_op_absent_on_non_gfx1100(op_name): + """On non-gfx1100 (CDNA), RDNA3 ops must NOT exist in _rocm_C. + + This is the real compile-level check: the binary was built without + gfx1100 support, so the ops should not have been compiled or registered. + """ + if not hasattr(torch.ops, "_rocm_C"): + return + assert not hasattr(torch.ops._rocm_C, op_name), ( + f"_rocm_C.{op_name} is registered on non-gfx1100 hardware — " + "compile guard is broken: check CMakeLists.txt " + "VLLM_ROCM_HAS_GFX1100 and torch_bindings.cpp #ifdef" + ) + + +@not_gfx1100 +def test_rocm_moe_not_supported_on_non_gfx1100(): + """rocm_moe.is_supported() must return False on non-gfx1100 hardware.""" + from vllm.model_executor.layers.quantization.compressed_tensors.compressed_tensors_moe import ( # noqa: E501 + rocm_moe, + ) + + wq = type("WQ", (), {"num_bits": 4})() + assert rocm_moe.is_supported(wq) is False, ( + "rocm_moe.is_supported() returned True on non-gfx1100 — " + "dispatch guard is broken" + ) + + +@not_gfx1100 +def test_dense_kernel_rejects_on_non_gfx1100(): + """RDNA3W4A16LinearKernel.can_implement must reject on non-gfx1100.""" + from vllm.model_executor.kernels.linear.mixed_precision.MPLinearKernel import ( # noqa: E501 + MPLinearLayerConfig, + ) + from vllm.model_executor.kernels.linear.mixed_precision.rdna3_w4a16 import ( # noqa: E501 + RDNA3W4A16LinearKernel, + ) + from vllm.scalar_type import scalar_types + + config = MPLinearLayerConfig( + full_weight_shape=(1024, 256), + partition_weight_shape=(1024, 256), + weight_type=scalar_types.uint4b8, + act_type=torch.float16, + group_size=128, + zero_points=False, + has_g_idx=False, + ) + ok, reason = RDNA3W4A16LinearKernel.can_implement(config) + assert ok is False, f"RDNA3 dense kernel accepted on non-gfx1100: {reason}" + + +# ============================================================================ +# Part C: Static source analysis (build-level guards) +# ============================================================================ + + +@needs_source +class TestCMakeGuards: + """Verify CMakeLists.txt only compiles RDNA3 .cu files for gfx1100.""" + + @staticmethod + def _read_cmake(): + return _read_source_or_skip("CMakeLists.txt") + + def test_rdna3_cu_files_inside_gfx1100_conditional(self): + """All RDNA3 .cu files must be listed inside the + ``if(VLLM_GPU_ARCHES MATCHES "gfx1100")`` block, not unconditionally. + """ + cmake = self._read_cmake() + for cu_file in RDNA3_CU_FILES: + assert cu_file in cmake, f"{cu_file} not found in CMakeLists.txt" + + lines = cmake.splitlines() + in_gfx1100_block = False + for line in lines: + if 'VLLM_GPU_ARCHES MATCHES "gfx1100"' in line: + in_gfx1100_block = True + if in_gfx1100_block and "endif()" in line: + in_gfx1100_block = False + if cu_file in line: + assert in_gfx1100_block, ( + f"{cu_file} is listed OUTSIDE the gfx1100 " + f"conditional in CMakeLists.txt — CDNA builds " + f"would compile RDNA3 code. Line: {line.strip()}" + ) + + def test_compile_definition_only_for_gfx1100(self): + """VLLM_ROCM_GFX1100 compile definition must be conditional.""" + cmake = self._read_cmake() + lines = cmake.splitlines() + in_gfx1100_block = False + for line in lines: + if "VLLM_ROCM_HAS_GFX1100)" in line: + in_gfx1100_block = True + if in_gfx1100_block and "endif()" in line: + in_gfx1100_block = False + if "VLLM_ROCM_GFX1100" in line and "target_compile_definitions" in line: + assert in_gfx1100_block, ( + "VLLM_ROCM_GFX1100 compile definition is set outside " + "the VLLM_ROCM_HAS_GFX1100 conditional — CDNA builds " + f"would define it. Line: {line.strip()}" + ) + + +@needs_source +class TestTorchBindingsGuards: + """Verify torch_bindings.cpp gates all RDNA3 ops behind #ifdef.""" + + @staticmethod + def _read_bindings(): + return _read_source_or_skip("csrc", "rocm", "torch_bindings.cpp") + + def test_all_rdna3_ops_inside_ifdef(self): + """Every rdna3 op def/impl must be between #ifdef VLLM_ROCM_GFX1100 + and #endif. If any is outside, a CDNA build would try to register + the op and link a symbol that doesn't exist. + """ + src = self._read_bindings() + lines = src.splitlines() + + inside_guard = False + rdna3_lines_outside = [] + + for i, line in enumerate(lines, 1): + if "#ifdef VLLM_ROCM_GFX1100" in line: + inside_guard = True + elif line.strip() == "#endif" and inside_guard: + inside_guard = False + + if ( + "rdna3" in line.lower() + and not line.strip().startswith("//") + and not inside_guard + ): + rdna3_lines_outside.append((i, line.strip())) + + assert not rdna3_lines_outside, ( + "RDNA3 op references found OUTSIDE #ifdef VLLM_ROCM_GFX1100 " + "in torch_bindings.cpp — these would break CDNA builds:\n" + + "\n".join(f" L{n}: {s}" for n, s in rdna3_lines_outside) + ) + + def test_no_unconditional_rdna3_includes(self): + """No #include of RDNA3-specific headers outside the guard.""" + src = self._read_bindings() + lines = src.splitlines() + + inside_guard = False + for i, line in enumerate(lines, 1): + if "#ifdef VLLM_ROCM_GFX1100" in line: + inside_guard = True + elif line.strip() == "#endif" and inside_guard: + inside_guard = False + + if "#include" in line and "rdna3" in line.lower(): + assert inside_guard, ( + f"L{i}: RDNA3 include outside gfx1100 guard: {line.strip()}" + ) + + +class TestCustomOpsGuards: + """Verify _custom_ops.py gates register_fake behind hasattr checks.""" + + @staticmethod + def _read_custom_ops(): + return _read_pkg_source_or_skip("_custom_ops.py") + + def test_register_fake_guarded_by_hasattr(self): + """Every register_fake for an RDNA3 op must be preceded by a hasattr + check — otherwise it would crash on import on CDNA where the ops + don't exist. + """ + src = self._read_custom_ops() + for op in RDNA3_OPS: + pattern = rf'register_fake\(\s*"_rocm_C::{op}"\s*\)' + match = re.search(pattern, src) + if match is None: + continue + preceding = src[: match.start()] + last_hasattr = preceding.rfind(f'hasattr(torch.ops._rocm_C, "{op}")') + assert last_hasattr != -1, ( + f'register_fake("_rocm_C::{op}") is not preceded by a ' + f"hasattr check — would crash on CDNA import" + ) + gap = preceding[last_hasattr:].count("\n") + assert gap <= 5, ( + f"hasattr guard for {op} is {gap} lines before " + f"register_fake — suspiciously far; verify it's the " + f"actual guard and not a coincidence" + ) + + def test_no_toplevel_rocm_c_import(self): + """No top-level ``from vllm._rocm_C import`` — would crash on CDNA.""" + src = self._read_custom_ops() + for line in src.splitlines(): + stripped = line.strip() + if stripped.startswith("#") or stripped.startswith("//"): + continue + assert "from vllm._rocm_C import" not in stripped, ( + f"Top-level import of _rocm_C in _custom_ops.py would " + f"crash on CDNA: {stripped}" + ) + + +# ============================================================================ +# Part D: Runtime mock (simulate CDNA on gfx1100 hardware) +# ============================================================================ + + +class _FakeWeightQuant: + """Minimal stand-in for a weight quantization config.""" + + def __init__(self, num_bits): + self.num_bits = num_bits + + +class TestMoEDispatchMocked: + """Mock on_gfx1100() to False and verify RDNA3 MoE is unreachable.""" + + def test_is_supported_false_when_mocked_cdna(self): + """rocm_moe.is_supported() must return False when not on gfx1100.""" + from vllm.model_executor.layers.quantization.compressed_tensors.compressed_tensors_moe import ( # noqa: E501 + rocm_moe, + ) + + with patch("vllm.platforms.rocm.on_gfx1100", return_value=False): + assert rocm_moe.is_supported(_FakeWeightQuant(num_bits=4)) is False + + @pytest.mark.parametrize("num_bits", [2, 3, 8, 16]) + def test_is_supported_rejects_non_w4(self, num_bits): + """is_supported() rejects non-4-bit even before checking arch.""" + from vllm.model_executor.layers.quantization.compressed_tensors.compressed_tensors_moe import ( # noqa: E501 + rocm_moe, + ) + + assert rocm_moe.is_supported(_FakeWeightQuant(num_bits=num_bits)) is False + + def test_is_supported_false_when_op_missing(self): + """is_supported() returns False when the C++ op doesn't exist.""" + from vllm.model_executor.layers.quantization.compressed_tensors.compressed_tensors_moe import ( # noqa: E501 + rocm_moe, + ) + + fake_rocm_c = type("FakeRocmC", (), {"gptq_gemm_rdna3": None})() + with patch.object(torch, "ops", create=True) as mock_ops: + mock_ops._rocm_C = fake_rocm_c + assert rocm_moe.is_supported(_FakeWeightQuant(num_bits=4)) is False + + def test_is_supported_false_when_rocm_c_absent(self): + """is_supported() returns False when _rocm_C doesn't exist at all.""" + from vllm.model_executor.layers.quantization.compressed_tensors.compressed_tensors_moe import ( # noqa: E501 + rocm_moe, + ) + + fake_ops = type("FakeOps", (), {})() + with patch.object(torch, "ops", fake_ops): + assert rocm_moe.is_supported(_FakeWeightQuant(num_bits=4)) is False + + +class TestDenseKernelSelectionMocked: + """Mock on_gfx1100() and verify dense RDNA3 kernel is not selected.""" + + @gfx1100_only + def test_can_implement_rejects_when_mocked_cdna(self): + """RDNA3W4A16LinearKernel.can_implement must reject on mocked CDNA.""" + from vllm.model_executor.kernels.linear.mixed_precision.MPLinearKernel import ( # noqa: E501 + MPLinearLayerConfig, + ) + from vllm.model_executor.kernels.linear.mixed_precision.rdna3_w4a16 import ( # noqa: E501 + RDNA3W4A16LinearKernel, + ) + from vllm.scalar_type import scalar_types + + config = MPLinearLayerConfig( + full_weight_shape=(1024, 256), + partition_weight_shape=(1024, 256), + weight_type=scalar_types.uint4b8, + act_type=torch.float16, + group_size=128, + zero_points=False, + has_g_idx=False, + ) + ok, _ = RDNA3W4A16LinearKernel.can_implement(config) + assert ok is True + + with ( + patch("vllm.platforms.rocm.on_gfx1100", return_value=False), + patch("vllm.platforms.rocm._ON_GFX1100", False), + ): + ok, reason = RDNA3W4A16LinearKernel.can_implement(config) + assert ok is False, f"RDNA3 kernel accepted on simulated CDNA: {reason}" + + @gfx1100_only + def test_chooser_skips_rdna3_when_mocked_cdna(self): + """choose_mp_linear_kernel must NOT return RDNA3 on mocked CDNA.""" + from vllm.model_executor.kernels.linear import ( + choose_mp_linear_kernel, + ) + from vllm.model_executor.kernels.linear.mixed_precision.MPLinearKernel import ( # noqa: E501 + MPLinearLayerConfig, + ) + from vllm.scalar_type import scalar_types + + config = MPLinearLayerConfig( + full_weight_shape=(1024, 256), + partition_weight_shape=(1024, 256), + weight_type=scalar_types.uint4b8, + act_type=torch.float16, + group_size=128, + zero_points=False, + has_g_idx=False, + ) + with ( + patch("vllm.platforms.rocm.on_gfx1100", return_value=False), + patch("vllm.platforms.rocm._ON_GFX1100", False), + ): + chosen = choose_mp_linear_kernel(config) + assert chosen.__name__ != "RDNA3W4A16LinearKernel", ( + "RDNA3 kernel was selected on simulated CDNA — " + "choose_mp_linear_kernel guard is broken" + ) + + +class TestCompressedTensorsMoEDispatchGuard: + """Verify compressed_tensors_moe.py only enters rocm_moe under is_rocm().""" + + def test_rocm_guard_in_dispatch_source(self): + """The rocm_moe import and call must be inside an is_rocm() check.""" + src = _read_pkg_source_or_skip( + "model_executor", + "layers", + "quantization", + "compressed_tensors", + "compressed_tensors_moe", + "compressed_tensors_moe.py", + ) + lines = src.splitlines() + + for i, line in enumerate(lines, 1): + stripped = line.strip() + if "rocm_moe" in stripped and not stripped.startswith("#"): + found_guard = False + for j in range(i - 1, max(0, i - 15), -1): + if "is_rocm()" in lines[j - 1]: + found_guard = True + break + assert found_guard, ( + f"L{i}: rocm_moe reference not protected by " + f"is_rocm() guard: {stripped}" + ) diff --git a/tests/kernels/quantization/test_rdna3_moe_w4a16.py b/tests/kernels/quantization/test_rdna3_moe_w4a16.py new file mode 100644 index 0000000000..4248251635 --- /dev/null +++ b/tests/kernels/quantization/test_rdna3_moe_w4a16.py @@ -0,0 +1,367 @@ +#!/usr/bin/env python3 +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""Correctness tests for the ROCm RDNA3 fused MoE W4A16 HIP kernel (gfx1100). + +Tests ``moe_gptq_gemm_rdna3`` against the dense ``gptq_gemm_rdna3`` as +reference: builds RDNA3-format weights (shuffled int32, synthesized qzeros), +runs the fused MoE kernel, and compares per-expert results. + +Model parameters taken from: + - cyankiwi/Qwen3-30B-A3B-Instruct-2507-AWQ-4bit + (hidden=2048, inter=768, E=128, top_k=8, G=32) + - Qwen3.6-35B-A3B-GPTQ-W4A16-G32 + (hidden=2048, inter=512, E=256, top_k=8, G=32) + +Run `pytest tests/kernels/quantization/test_rdna3_moe_w4a16.py`. +""" + +import pytest +import torch + +from vllm.platforms import current_platform + +if not current_platform.is_rocm(): + pytest.skip("RDNA3 MoE W4A16 kernel is ROCm-only", allow_module_level=True) + +from vllm import _custom_ops as ops # noqa: E402 +from vllm.model_executor.layers.fused_moe.activation import ( # noqa: E402 + MoEActivation, + apply_moe_activation, +) +from vllm.model_executor.layers.fused_moe.moe_align_block_size import ( # noqa: E402 + moe_align_block_size, +) +from vllm.model_executor.layers.quantization.utils.quant_utils import ( # noqa: E402 + pack_quantized_values_into_int32, +) +from vllm.platforms.rocm import on_gfx1100 # noqa: E402 +from vllm.scalar_type import scalar_types # noqa: E402 + +device = "cuda" + +gfx1100_only = pytest.mark.skipif( + not ( + on_gfx1100() + and hasattr(torch.ops, "_rocm_C") + and hasattr(torch.ops._rocm_C, "moe_gptq_gemm_rdna3") + ), + reason="Requires gfx1100 with moe_gptq_gemm_rdna3 op", +) + +# Model configurations: real K/N/top_k/group_size dims, E capped at 16 to +# fit in test GPU memory (full E=128/256 would need >20GB for weights alone). +# Kernel behavior is E-independent (per-expert tiling), so E=16 is sufficient. +MODEL_CONFIGS = [ + # cyankiwi/Qwen3-30B-A3B-Instruct-2507-AWQ-4bit dims (E capped) + pytest.param(16, 2048, 768, 8, 32, id="Qwen3-30B-A3B"), + # Qwen3.6-35B-A3B-GPTQ-W4A16-G32 dims (E capped) + pytest.param(16, 2048, 512, 8, 32, id="Qwen3.6-35B-A3B"), +] + +# Token counts: decode (1), small batch (4), medium (16), prefill (64) +NUM_TOKENS = [1, 4, 16, 64, 256, 512] + + +def _make_packed_weights(E, K, N): + """Create random 4-bit packed weights [E, K/8, N] int32 + shuffle.""" + w = torch.randint(0, 16, (E, K, N), dtype=torch.int32, device=device) + packed = torch.zeros(E, K // 8, N, dtype=torch.int32, device=device) + for i in range(8): + packed |= (w[:, i::8, :] & 0xF) << (i * 4) + g_idx = torch.empty(0, dtype=torch.int32, device=device) + for e in range(E): + we = packed[e].contiguous() + ops.gptq_shuffle(we, g_idx, 4) + packed[e] = we + return packed + + +def _make_scales(E, groups, N, dtype): + return torch.rand(E, groups, N, dtype=dtype, device=device) * 0.1 + + +def _make_qzeros(E, groups, N): + zeros = torch.full( + (groups, N), + scalar_types.uint4b8.bias - 1, + dtype=torch.int32, + device=device, + ) + qz = pack_quantized_values_into_int32( + zeros, + scalar_types.uint4b8, + packed_dim=1, + ) + return qz.unsqueeze(0).expand(E, -1, -1).contiguous() + + +@gfx1100_only +@pytest.mark.parametrize("E, K, N_inter, top_k, group_size", MODEL_CONFIGS) +@pytest.mark.parametrize("M", NUM_TOKENS) +@pytest.mark.parametrize("dtype", [torch.float16, torch.bfloat16]) +@pytest.mark.parametrize("block_size_m", [1, 4]) +def test_fused_moe_w1_matches_dense( + E, K, N_inter, top_k, group_size, M, dtype, block_size_m +): + """w1 GEMM via fused kernel matches per-expert dense kernel.""" + N_gate_up = N_inter * 2 + groups = K // group_size + + torch.manual_seed(42) + x = torch.randn(M, K, dtype=dtype, device=device) + w13 = _make_packed_weights(E, K, N_gate_up) + w13_s = _make_scales(E, groups, N_gate_up, dtype) + w13_z = _make_qzeros(E, groups, N_gate_up) + g_idx = torch.empty(0, dtype=torch.int32, device=device) + + topk_ids = torch.randint(0, E, (M, top_k), device=device, dtype=torch.int32) + si, ei, ntp = moe_align_block_size(topk_ids, block_size_m, E) + + # Fused kernel + fused_out = torch.zeros(M * top_k, N_gate_up, dtype=dtype, device=device) + ops.moe_gptq_gemm_rdna3( + x, + fused_out, + w13, + w13_s, + w13_z, + torch.empty(0, device=device), + si, + ei, + ntp, + top_k, + block_size_m, + False, + 0, + ) + + # Per-expert dense reference + ref_out = torch.zeros(M * top_k, N_gate_up, dtype=dtype, device=device) + for m in range(M): + for k in range(top_k): + e = topk_ids[m, k].item() + flat = m * top_k + k + ref = ops.gptq_gemm_rdna3( + x[m : m + 1], + w13[e], + w13_z[e], + w13_s[e], + g_idx, + False, + ) + ref_out[flat] = ref.squeeze() + + # Split-K atomics can cause minor fp16/bf16 rounding differences + # at large K (e.g. K=2048 → 8 K-blocks). Use allclose, not equal. + atol = 0.5 if dtype == torch.bfloat16 else 0.1 + assert torch.allclose(fused_out, ref_out, atol=atol, rtol=0.01), ( + f"max diff: {(fused_out - ref_out).abs().max().item()}" + ) + + +@gfx1100_only +@pytest.mark.parametrize("E, K, N_inter, top_k, group_size", MODEL_CONFIGS) +@pytest.mark.parametrize("M", NUM_TOKENS) +@pytest.mark.parametrize("dtype", [torch.float16, torch.bfloat16]) +def test_fused_moe_output_topk_reduces(E, K, N_inter, top_k, group_size, M, dtype): + """output_topk fuses moe_sum: multiple experts write to same output row.""" + groups = K // group_size + + torch.manual_seed(123) + x = torch.randn(M * top_k, K, dtype=dtype, device=device) + w = _make_packed_weights(E, K, N_inter) + ws = _make_scales(E, groups, N_inter, dtype) + wz = _make_qzeros(E, groups, N_inter) + + topk_ids = torch.randint(0, E, (M, top_k), device=device, dtype=torch.int32) + topk_w = torch.softmax( + torch.randn(M, top_k, device=device), + dim=-1, + ).float() + + si, ei, ntp = moe_align_block_size(topk_ids, 1, E) + + # Without output_topk: write to [M*top_k, N] then moe_sum + flat_out = torch.zeros(M * top_k, N_inter, dtype=dtype, device=device) + ops.moe_gptq_gemm_rdna3( + x, + flat_out, + w, + ws, + wz, + topk_w.view(-1), + si, + ei, + ntp, + 1, + 1, + True, + 0, + ) + ref = torch.zeros(M, N_inter, dtype=dtype, device=device) + ops.moe_sum(flat_out.view(M, top_k, N_inter), ref) + + # With output_topk: write directly to [M, N] + fused = torch.zeros(M, N_inter, dtype=dtype, device=device) + ops.moe_gptq_gemm_rdna3( + x, + fused, + w, + ws, + wz, + topk_w.view(-1), + si, + ei, + ntp, + 1, + 1, + True, + top_k, + ) + + atol = 1.0 if dtype == torch.bfloat16 else 0.1 + assert torch.allclose(fused, ref, atol=atol, rtol=0.01), ( + f"max diff: {(fused - ref).abs().max().item()}" + ) + + +@gfx1100_only +@pytest.mark.parametrize("E, K, N_inter, top_k, group_size", MODEL_CONFIGS) +@pytest.mark.parametrize("M", NUM_TOKENS) +@pytest.mark.parametrize("dtype", [torch.float16, torch.bfloat16]) +def test_full_moe_e2e(E, K, N_inter, top_k, group_size, M, dtype): + """Full MoE forward: w1 + silu_and_mul + w2 with output_topk reduce.""" + N_gate_up = N_inter * 2 + hidden = K + + torch.manual_seed(7) + x = torch.randn(M, K, dtype=dtype, device=device) + w13 = _make_packed_weights(E, K, N_gate_up) + w13_s = _make_scales(E, K // group_size, N_gate_up, dtype) + w13_z = _make_qzeros(E, K // group_size, N_gate_up) + w2 = _make_packed_weights(E, N_inter, hidden) + w2_s = _make_scales(E, N_inter // group_size, hidden, dtype) + w2_z = _make_qzeros(E, N_inter // group_size, hidden) + g_idx = torch.empty(0, dtype=torch.int32, device=device) + + topk_ids = torch.randint(0, E, (M, top_k), device=device, dtype=torch.int32) + topk_w = torch.softmax( + torch.randn(M, top_k, device=device), + dim=-1, + ).float() + + si, ei, ntp = moe_align_block_size(topk_ids, 1, E) + + # Fused path (what apply() does) + w1_out = torch.zeros(M * top_k, N_gate_up, dtype=dtype, device=device) + ops.moe_gptq_gemm_rdna3( + x, + w1_out, + w13, + w13_s, + w13_z, + torch.empty(0, device=device), + si, + ei, + ntp, + top_k, + 1, + False, + 0, + ) + act_out = torch.empty(M * top_k, N_inter, dtype=dtype, device=device) + apply_moe_activation(MoEActivation.SILU, act_out, w1_out) + fused = torch.zeros(M, hidden, dtype=dtype, device=device) + ops.moe_gptq_gemm_rdna3( + act_out, + fused, + w2, + w2_s, + w2_z, + topk_w.view(-1), + si, + ei, + ntp, + 1, + 1, + True, + top_k, + ) + + # Per-expert reference + ref = torch.zeros(M, hidden, dtype=dtype, device=device) + for m_idx in range(M): + for k_idx in range(top_k): + e = topk_ids[m_idx, k_idx].item() + w = topk_w[m_idx, k_idx].item() + r1 = ops.gptq_gemm_rdna3( + x[m_idx : m_idx + 1], + w13[e], + w13_z[e], + w13_s[e], + g_idx, + False, + ) + a = torch.empty(1, N_inter, dtype=dtype, device=device) + apply_moe_activation(MoEActivation.SILU, a, r1) + r2 = ops.gptq_gemm_rdna3( + a, + w2[e], + w2_z[e], + w2_s[e], + g_idx, + False, + ) + ref[m_idx] += r2.squeeze() * w + + # E2E chains w1 + activation + w2 + topk_w + output_topk reduce. + # Each step accumulates rounding error (split-K atomics, topk_w + # multiply order). Use relative L2 norm like the dense kernel test. + diff_l2 = torch.norm(fused.float() - ref.float()) + ref_l2 = torch.norm(ref.float()) + rel_l2 = (diff_l2 / ref_l2).item() if ref_l2 > 0 else 0.0 + threshold = 0.05 if dtype == torch.float16 else 0.10 + assert rel_l2 < threshold, ( + f"rel L2 = {rel_l2:.4f} (threshold {threshold}), " + f"max abs diff: {(fused - ref).abs().max().item()}" + ) + + +@gfx1100_only +def test_expert_id_minus_one(): + """Kernel handles expert_id == -1 (expert parallelism) without crash.""" + # Qwen3-30B-A3B dims (E capped for memory) + E, K, N = 16, 2048, 768 + groups = K // 32 + + w = _make_packed_weights(E, K, N) + ws = _make_scales(E, groups, N, torch.bfloat16) + wz = _make_qzeros(E, groups, N) + x = torch.randn(1, K, dtype=torch.bfloat16, device=device) + + # Manually create sorted_token_ids/expert_ids with -1 + sorted_ids = torch.tensor([0], dtype=torch.int32, device=device) + expert_ids = torch.tensor([-1], dtype=torch.int32, device=device) + ntp = torch.tensor([1], dtype=torch.int32, device=device) + + out = torch.zeros(1, N, dtype=torch.bfloat16, device=device) + ops.moe_gptq_gemm_rdna3( + x, + out, + w, + ws, + wz, + torch.empty(0, device=device), + sorted_ids, + expert_ids, + ntp, + 1, + 1, + False, + 0, + ) + current_platform.synchronize() + + # Output should remain zero (expert skipped) + assert torch.equal(out, torch.zeros_like(out)) diff --git a/vllm/_custom_ops.py b/vllm/_custom_ops.py index bd8a19b6d2..7eaa10b141 100644 --- a/vllm/_custom_ops.py +++ b/vllm/_custom_ops.py @@ -695,6 +695,59 @@ if hasattr(torch.ops, "_rocm_C") and hasattr(torch.ops._rocm_C, "gptq_gemm_rdna3 ) +def moe_gptq_gemm_rdna3( + a: torch.Tensor, + c: torch.Tensor, + b_q_weight: torch.Tensor, + b_scales: torch.Tensor, + b_qzeros: torch.Tensor, + topk_weights: torch.Tensor, + sorted_token_ids: torch.Tensor, + expert_ids: torch.Tensor, + num_tokens_post_padded: torch.Tensor, + top_k: int, + block_size_m: int, + mul_topk_weight: bool, + output_topk: int = 0, +) -> None: + torch.ops._rocm_C.moe_gptq_gemm_rdna3( + a, + c, + b_q_weight, + b_scales, + b_qzeros, + topk_weights, + sorted_token_ids, + expert_ids, + num_tokens_post_padded, + top_k, + block_size_m, + mul_topk_weight, + output_topk, + ) + + +if hasattr(torch.ops, "_rocm_C") and hasattr(torch.ops._rocm_C, "moe_gptq_gemm_rdna3"): + + @register_fake("_rocm_C::moe_gptq_gemm_rdna3") + def _moe_gptq_gemm_rdna3_fake( + a: torch.Tensor, + c: torch.Tensor, + b_q_weight: torch.Tensor, + b_scales: torch.Tensor, + b_qzeros: torch.Tensor, + topk_weights: torch.Tensor, + sorted_token_ids: torch.Tensor, + expert_ids: torch.Tensor, + num_tokens_post_padded: torch.Tensor, + top_k: int, + block_size_m: int, + mul_topk_weight: bool, + output_topk: int = 0, + ) -> None: + return + + if hasattr(torch.ops._C, "allspark_w8a16_gemm"): @register_fake("_C::allspark_w8a16_gemm") diff --git a/vllm/model_executor/layers/fused_moe/layer.py b/vllm/model_executor/layers/fused_moe/layer.py index d1d6a4f03c..0abe83a635 100644 --- a/vllm/model_executor/layers/fused_moe/layer.py +++ b/vllm/model_executor/layers/fused_moe/layer.py @@ -892,6 +892,7 @@ class FusedMoE(PluggableLayer): if quant_method_name in ( "CompressedTensorsWNA16MarlinMoEMethod", "CompressedTensorsWNA16MoEMethod", + "CompressedTensorsWNA16RDNA3MoEMethod", ): if is_transposed: loaded_weight = loaded_weight.t().contiguous() diff --git a/vllm/model_executor/layers/quantization/compressed_tensors/compressed_tensors_moe/compressed_tensors_moe.py b/vllm/model_executor/layers/quantization/compressed_tensors/compressed_tensors_moe/compressed_tensors_moe.py index a2e9416219..1e33bfe6fc 100644 --- a/vllm/model_executor/layers/quantization/compressed_tensors/compressed_tensors_moe/compressed_tensors_moe.py +++ b/vllm/model_executor/layers/quantization/compressed_tensors/compressed_tensors_moe/compressed_tensors_moe.py @@ -98,10 +98,6 @@ class CompressedTensorsMoEMethod(FusedMoEMethodBase): not check_moe_marlin_supports_layer(layer, group_size) or current_platform.is_rocm() ): - from .compressed_tensors_moe_wna16 import ( - CompressedTensorsWNA16MoEMethod, - ) - if ( weight_quant.strategy == QuantizationStrategy.GROUP and weight_quant.actorder @@ -110,6 +106,20 @@ class CompressedTensorsMoEMethod(FusedMoEMethodBase): raise ValueError( "WNA16MoE is not supported with actorder=group/dynamic." ) + + # Native ROCm HIP kernels (RDNA3, etc.) + if current_platform.is_rocm(): + from . import rocm_moe + + if rocm_moe.is_supported(weight_quant): + return rocm_moe.make_method( + weight_quant, input_quant, layer.moe_config + ) + + from .compressed_tensors_moe_wna16 import ( + CompressedTensorsWNA16MoEMethod, + ) + logger.info_once("Using CompressedTensorsWNA16MoEMethod") return CompressedTensorsWNA16MoEMethod( weight_quant, input_quant, layer.moe_config diff --git a/vllm/model_executor/layers/quantization/compressed_tensors/compressed_tensors_moe/compressed_tensors_moe_wna16_rdna3.py b/vllm/model_executor/layers/quantization/compressed_tensors/compressed_tensors_moe/compressed_tensors_moe_wna16_rdna3.py new file mode 100644 index 0000000000..25cc7b115f --- /dev/null +++ b/vllm/model_executor/layers/quantization/compressed_tensors/compressed_tensors_moe/compressed_tensors_moe_wna16_rdna3.py @@ -0,0 +1,261 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""CompressedTensors MoE W4A16 using the fused RDNA3 (gfx1100) HIP kernel. + +Uses ``moe_gptq_gemm_rdna3`` — a single HIP kernel launch per GEMM that +handles expert routing + W4A16 dequant + dot product with atomic output. + +Weight format (per expert, same as dense RDNA3 W4A16): + - Packed int32 ``[E, K/8, N]`` with exllama shuffle + - Scales ``[E, groups, N]`` in activation dtype + - Zero points ``[E, groups, N/8]`` packed int32 (synthesized) +""" + +import torch + +from vllm import _custom_ops as ops +from vllm.logger import init_logger +from vllm.model_executor.layers.fused_moe import ( + RoutedExperts, + SharedExperts, +) +from vllm.model_executor.layers.fused_moe.activation import ( + MoEActivation, + apply_moe_activation, +) +from vllm.model_executor.layers.fused_moe.moe_align_block_size import ( + moe_align_block_size, +) +from vllm.model_executor.layers.quantization.compressed_tensors.compressed_tensors_moe.compressed_tensors_moe_wna16 import ( # noqa: E501 + CompressedTensorsWNA16MoEMethod, +) +from vllm.model_executor.layers.quantization.utils.quant_utils import ( + pack_quantized_values_into_int32, +) +from vllm.scalar_type import scalar_types + +logger = init_logger(__name__) + + +def _synthesize_qzeros( + groups: int, out_features: int, device: torch.device +) -> torch.Tensor: + """Create packed zero-point tensor for symmetric quant. + + GPTQv1 +1 quirk: kernel adds 1 to stored zeros, so encode + (bias - 1) = 7 for uint4b8 (bias=8). + """ + zeros = torch.full( + (groups, out_features), + scalar_types.uint4b8.bias - 1, + dtype=torch.int32, + device=device, + ) + return pack_quantized_values_into_int32(zeros, scalar_types.uint4b8, packed_dim=1) + + +class CompressedTensorsWNA16RDNA3MoEMethod(CompressedTensorsWNA16MoEMethod): + """W4A16 MoE using the fused RDNA3 HIP kernel (moe_gptq_gemm_rdna3). + + Weights are in RDNA3 format (shuffled int32 [E, K/8, N]), + NOT Triton format (transposed uint8). apply() dispatches through + the fused HIP kernel directly. + """ + + def process_weights_after_loading(self, layer: torch.nn.Module) -> None: + device = layer.w13_weight_packed.device + num_experts = layer.w13_weight_packed.shape[0] + empty_g_idx = torch.empty(0, dtype=torch.int32, device=device) + + # Shuffle weights in-place per expert (exllama nibble interleave) + for e in range(num_experts): + w13_e = layer.w13_weight_packed.data[e].contiguous() + ops.gptq_shuffle(w13_e, empty_g_idx, 4) + layer.w13_weight_packed.data[e] = w13_e + w2_e = layer.w2_weight_packed.data[e].contiguous() + ops.gptq_shuffle(w2_e, empty_g_idx, 4) + layer.w2_weight_packed.data[e] = w2_e + + # Keep scales as [E, groups, N] in activation dtype + act_dtype = layer.w13_weight_scale.dtype + layer.w13_weight_scale = torch.nn.Parameter( + layer.w13_weight_scale.to(dtype=act_dtype).contiguous(), + requires_grad=False, + ) + layer.w2_weight_scale = torch.nn.Parameter( + layer.w2_weight_scale.to(dtype=act_dtype).contiguous(), + requires_grad=False, + ) + + # Synthesize packed zero points: [E, groups, N/8] int32 + w13_groups = (layer.w13_weight_packed.shape[1] * 8) // self.group_size + w13_N = layer.w13_weight_packed.shape[2] + w2_groups = (layer.w2_weight_packed.shape[1] * 8) // self.group_size + w2_N = layer.w2_weight_packed.shape[2] + + w13_qz = _synthesize_qzeros(w13_groups, w13_N, device) + w2_qz = _synthesize_qzeros(w2_groups, w2_N, device) + layer.w13_qzeros = torch.nn.Parameter( + w13_qz.unsqueeze(0).expand(num_experts, -1, -1).contiguous(), + requires_grad=False, + ) + layer.w2_qzeros = torch.nn.Parameter( + w2_qz.unsqueeze(0).expand(num_experts, -1, -1).contiguous(), + requires_grad=False, + ) + + # Pre-allocate reusable buffers for decode (sizes based on top_k=8) + N_gate_up = w13_N + hidden_size = w2_N + intermediate = N_gate_up // 2 # gated activation + # Max tokens we expect in decode; prefill will re-allocate if needed + max_decode_tokens = 16 + top_k = 8 # conservative default + buf_size = max_decode_tokens * top_k + layer.rdna3_w1_buf = torch.zeros( + buf_size, N_gate_up, dtype=act_dtype, device=device + ) + layer.rdna3_act_buf = torch.empty( + buf_size, intermediate, dtype=act_dtype, device=device + ) + layer.rdna3_out_buf = torch.zeros( + max_decode_tokens, hidden_size, dtype=act_dtype, device=device + ) + layer.rdna3_empty_tw = torch.empty(0, device=device) + + def apply( + self, + layer: RoutedExperts, + x: torch.Tensor, + topk_weights: torch.Tensor, + topk_ids: torch.Tensor, + shared_experts: SharedExperts | None, + shared_experts_input: torch.Tensor | None, + ) -> torch.Tensor: + activation = ( + layer.activation + if isinstance(layer.activation, MoEActivation) + else MoEActivation.from_str(layer.activation) + ) + return _rdna3_fused_moe( + x, + topk_weights, + topk_ids, + layer=layer, + activation=activation, + apply_router_weight_on_input=(layer.apply_router_weight_on_input), + global_num_experts=layer.global_num_experts, + expert_map=layer.expert_map, + ) + + +def _rdna3_fused_moe( + hidden_states: torch.Tensor, + topk_weights: torch.Tensor, + topk_ids: torch.Tensor, + layer: RoutedExperts, + activation: MoEActivation, + apply_router_weight_on_input: bool, + global_num_experts: int, + expert_map: torch.Tensor | None, +) -> torch.Tensor: + """Fused MoE forward using the RDNA3 W4A16 HIP kernel. + + Optimizations vs naive dispatch: + - BLOCK_SIZE_M=1 for decode (no padding waste, bf16 fast path) + - Pre-allocated buffers (no torch.zeros per call) + - Inline token sorting for small M (skip moe_align_block_size) + - moe_sum fused into output accumulation + """ + num_tokens = hidden_states.shape[0] + top_k = topk_ids.shape[1] + total_tokens = num_tokens * top_k + N_gate_up = layer.w13_weight_packed.shape[2] + hidden_size = layer.w2_weight_packed.shape[2] + dtype = hidden_states.dtype + device = hidden_states.device + + intermediate_size = N_gate_up // 2 if activation.is_gated else N_gate_up + + if global_num_experts <= 0: + global_num_experts = layer.w13_weight_packed.shape[0] + + # BLOCK_SIZE_M=1 for decode (small M), 4 for prefill + block_size_m = 1 if num_tokens <= 4 else 4 + + # --- Token routing --- + sorted_token_ids, expert_ids, num_tokens_post_padded = moe_align_block_size( + topk_ids, + block_size_m, + global_num_experts, + expert_map, + ) + + # --- Reuse pre-allocated buffers when possible --- + if total_tokens <= layer.rdna3_w1_buf.shape[0]: + w1_out = layer.rdna3_w1_buf[:total_tokens] + w1_out.zero_() + act_out = layer.rdna3_act_buf[:total_tokens] + else: + w1_out = torch.zeros( + total_tokens, + N_gate_up, + dtype=dtype, + device=device, + ) + act_out = torch.empty( + total_tokens, + intermediate_size, + dtype=dtype, + device=device, + ) + + # --- topk weights (pre-cast to float32 for kernel) --- + topk_w_float = topk_weights.view(-1).float() + empty_tw = layer.rdna3_empty_tw + + # --- w1 GEMM: [M, K] -> [M*top_k, N_gate_up] --- + ops.moe_gptq_gemm_rdna3( + hidden_states, + w1_out, + layer.w13_weight_packed, + layer.w13_weight_scale, + layer.w13_qzeros, + topk_w_float if apply_router_weight_on_input else empty_tw, + sorted_token_ids, + expert_ids, + num_tokens_post_padded, + top_k, + block_size_m, + apply_router_weight_on_input, + ) + + # --- Activation (silu_and_mul etc.) --- + apply_moe_activation(activation, act_out, w1_out) + + # --- w2 GEMM: [M*top_k, intermediate] -> [M, hidden] (fused reduce) --- + # output_topk=top_k: kernel writes to out[token_id / top_k] directly, + # fusing moe_sum into the atomic accumulation — saves one kernel launch + # and the w2_out intermediate buffer. + out = torch.zeros( + num_tokens, + hidden_size, + dtype=dtype, + device=device, + ) + ops.moe_gptq_gemm_rdna3( + act_out, + out, + layer.w2_weight_packed, + layer.w2_weight_scale, + layer.w2_qzeros, + topk_w_float if not apply_router_weight_on_input else empty_tw, + sorted_token_ids, + expert_ids, + num_tokens_post_padded, + 1, + block_size_m, + not apply_router_weight_on_input, + output_topk=top_k, + ) + return out diff --git a/vllm/model_executor/layers/quantization/compressed_tensors/compressed_tensors_moe/rocm_moe.py b/vllm/model_executor/layers/quantization/compressed_tensors/compressed_tensors_moe/rocm_moe.py new file mode 100644 index 0000000000..e72caa0796 --- /dev/null +++ b/vllm/model_executor/layers/quantization/compressed_tensors/compressed_tensors_moe/rocm_moe.py @@ -0,0 +1,48 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""ROCm MoE kernel dispatcher. + +Selects architecture-specific native HIP MoE kernels in priority order. +Falls back to the Triton WNA16 path when no native kernel is available. +""" + +import torch + +from vllm.logger import init_logger + +logger = init_logger(__name__) + + +def is_supported(weight_quant) -> bool: + """Check if a native ROCm MoE kernel is available for this config.""" + if weight_quant.num_bits != 4: + return False + + from vllm.platforms.rocm import on_gfx1100 + + # RDNA3 (gfx1100). Future: add RDNA4 (gfx12x), CDNA (gfx94x), etc. + return ( + on_gfx1100() + and hasattr(torch.ops, "_rocm_C") + and hasattr(torch.ops._rocm_C, "moe_gptq_gemm_rdna3") + ) + + +def make_method(weight_quant, input_quant, moe_config): + """Create the native ROCm MoE method. Call only after is_supported().""" + from vllm.platforms.rocm import on_gfx1100 + + if on_gfx1100(): + from .compressed_tensors_moe_wna16_rdna3 import ( + CompressedTensorsWNA16RDNA3MoEMethod, + ) + + logger.info_once( + "Using CompressedTensorsWNA16RDNA3MoEMethod (native RDNA3 HIP kernel)" + ) + return CompressedTensorsWNA16RDNA3MoEMethod( + weight_quant, input_quant, moe_config + ) + + # Future: RDNA4, CDNA, etc. + raise RuntimeError("is_supported() returned True but no kernel matched")