[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>
This commit is contained in:
JartX
+2
-1
@@ -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(
|
||||
|
||||
@@ -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 <cstdint>
|
||||
|
||||
#include <torch/all.h>
|
||||
#include <c10/cuda/CUDAGuard.h>
|
||||
#include <ATen/cuda/CUDAContext.h>
|
||||
|
||||
#include <hip/hip_runtime.h>
|
||||
#include <hip/hip_bf16.h>
|
||||
#include <hip/hip_fp16.h>
|
||||
|
||||
#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 <typename T>
|
||||
__forceinline__ __device__ T tzero();
|
||||
|
||||
template <>
|
||||
__forceinline__ __device__ half tzero<half>() {
|
||||
return __float2half_rn(0.0f);
|
||||
}
|
||||
|
||||
template <>
|
||||
__forceinline__ __device__ bf16_t tzero<bf16_t>() {
|
||||
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<unsigned long long*>(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<unsigned long long*>(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 <typename T>
|
||||
__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 <typename T, int BLOCK_SIZE_M>
|
||||
__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<T, half>::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<T>();
|
||||
}
|
||||
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<T>(sc_row, n, scales);
|
||||
if constexpr (std::is_same<T, half>::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<T, half>::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<const half*>(&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<const uint32_t*>(
|
||||
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<const uint32_t*>(&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<T, half>::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 <typename T, int BLOCK_SIZE_M>
|
||||
__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 <typename T, int BLOCK_SIZE_M>
|
||||
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<T, BLOCK_SIZE_M><<<grid, block, 0, stream>>>(
|
||||
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 <typename T>
|
||||
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<T, 1>(
|
||||
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<T, 2>(
|
||||
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<T, 4>(
|
||||
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<T, 8>(
|
||||
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<float>() : 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<std::remove_pointer_t<decltype(a_ptr)>>;
|
||||
vllm::moe_gptq_rdna3::dispatch_moe_gemm_q4<T>(
|
||||
a_ptr, c_ptr, (const uint32_t*)b_q_weight.data_ptr<int32_t>(), s_ptr,
|
||||
(const uint32_t*)b_qzeros.data_ptr<int32_t>(), topk_w_ptr,
|
||||
sorted_token_ids.data_ptr<int32_t>(), expert_ids.data_ptr<int32_t>(),
|
||||
num_tokens_post_padded.data_ptr<int32_t>(), 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());
|
||||
}
|
||||
}
|
||||
@@ -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,
|
||||
|
||||
@@ -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
|
||||
|
||||
@@ -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}"
|
||||
)
|
||||
@@ -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))
|
||||
@@ -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")
|
||||
|
||||
@@ -892,6 +892,7 @@ class FusedMoE(PluggableLayer):
|
||||
if quant_method_name in (
|
||||
"CompressedTensorsWNA16MarlinMoEMethod",
|
||||
"CompressedTensorsWNA16MoEMethod",
|
||||
"CompressedTensorsWNA16RDNA3MoEMethod",
|
||||
):
|
||||
if is_transposed:
|
||||
loaded_weight = loaded_weight.t().contiguous()
|
||||
|
||||
+14
-4
@@ -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
|
||||
|
||||
+261
@@ -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
|
||||
+48
@@ -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")
|
||||
Reference in New Issue
Block a user