From 811769dac6907b76d222dffbaec3d03970b6c3df Mon Sep 17 00:00:00 2001 From: Robert Esclapez Garcia Date: Wed, 24 Jun 2026 06:53:28 -0700 Subject: [PATCH 1/3] [ROCm][MoE] rdna_moe_gemm W4A16 MoE prefill WMMA gemm1 for gfx11 (default-on) Adds a producer/consumer WMMA kernel (rdna_moe_gemm) for the W4A16 MoE prefill GEMM1 on AMD RDNA3 (gfx11), a faster alternative to the Triton fused_moe_kernel_gptq_awq. Enabled by default on gfx11 (VLLM_MOE_HIP=0 forces Triton, =1 forces on). On the rdna_moe_gemm path gemm1 (up/gate proj, top_k>1) runs the rdna_moe_gemm WMMA kernel and gemm2 (down proj, top_k=1) runs Triton, both at a single block_m=32 moe_align alignment. The per-routed-token activations are stored compact (flat-topk indexed -- the kernels write C[sorted_token_ids[slot]]), so gemm2 gathers them via the shared sorted_token_ids with no re-permute. gemm1 at block_m=32 is faster than at 16 on gfx1151 (the larger WMMA tile more than pays for the extra alignment padding), so there is a single alignment for both gemms. - csrc/rocm/moe_gemm_w4a16_wmma.cu: the WMMA kernel with K, N and the weight N-row stride as runtime args (one instantiation per tile family handles any compliant shape and any weight padding). Compiled into _rocm_C; the WMMA body is gfx11-only with a stub on other arches. Registered as torch.ops._rocm_C.moe_gemm_w4a16; an unsupported shape TORCH_CHECKs. - vllm/envs.py: VLLM_MOE_HIP tri-state (unset = default-on on gfx11). - vllm/model_executor/layers/fused_moe/moe_hip_w4a16.py: host-side shape predicate (prefill_uses_rdna_moe_gemm) plus a graph-safe vLLM custom op wrapper with a no-op fake, so the path is torch.compile-safe. - vllm/model_executor/layers/fused_moe/hybrid_w4a16_moe.py: apply() dispatch -- rdna_moe_gemm gemm1 when the shape is supported, else Triton; gemm2 always Triton. - tests/kernels/quantization/test_moe_gemm_w4a16.py: compiled op vs Triton reference across shapes and block sizes. AI assistance (Claude) was used. Co-authored-by: Claude Signed-off-by: Robert Esclapez Garcia --- CMakeLists.txt | 1 + csrc/rocm/moe_gemm_w4a16_wmma.cu | 389 ++++++++++++++++++ csrc/rocm/ops.h | 9 + csrc/rocm/torch_bindings.cpp | 10 + .../quantization/test_moe_gemm_w4a16.py | 136 ++++++ vllm/envs.py | 5 + .../layers/fused_moe/hybrid_w4a16_moe.py | 110 +++-- .../layers/fused_moe/moe_hip_w4a16.py | 123 ++++++ 8 files changed, 747 insertions(+), 36 deletions(-) create mode 100644 csrc/rocm/moe_gemm_w4a16_wmma.cu create mode 100644 tests/kernels/quantization/test_moe_gemm_w4a16.py create mode 100644 vllm/model_executor/layers/fused_moe/moe_hip_w4a16.py diff --git a/CMakeLists.txt b/CMakeLists.txt index c84157bf5455..f29a019d46dd 100644 --- a/CMakeLists.txt +++ b/CMakeLists.txt @@ -1300,6 +1300,7 @@ if(VLLM_GPU_LANG STREQUAL "HIP") "csrc/rocm/skinny_gemms_w8a8/instantiate_n3.cu" "csrc/rocm/skinny_gemms_w8a8/instantiate_n4.cu" "csrc/rocm/skinny_gemms_w8a8/instantiate_n5.cu" + "csrc/rocm/moe_gemm_w4a16_wmma.cu" "csrc/rocm/attention.cu") set(VLLM_ROCM_FLAGS ${VLLM_GPU_FLAGS}) diff --git a/csrc/rocm/moe_gemm_w4a16_wmma.cu b/csrc/rocm/moe_gemm_w4a16_wmma.cu new file mode 100644 index 000000000000..faaaf13714e7 --- /dev/null +++ b/csrc/rocm/moe_gemm_w4a16_wmma.cu @@ -0,0 +1,389 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright contributors to the vLLM project +// +// W4A16 MoE producer/consumer WMMA prefill GEMM for AMD RDNA3 (gfx11). +// Always compiled into _rocm_C; the WMMA body below is gfx11-only (empty stub +// on other device arches, so multi-arch/CDNA builds still link) and registered +// as torch.ops._rocm_C.moe_gemm_w4a16 in torch_bindings.cpp. Python gates calls +// on on_gfx11(). +// +// moe_gemm_w4a16(A, w_packed, w_scale, sorted_token_ids, expert_ids, C, +// n_valid_tokens, top_k, block_m, num_blocks) -> () +// computes C[slot] = A[token//top_k] @ dequant(W[expert]) in place with +// identical semantics/layout to invoke_fused_moe_kernel_hybrid_triton. Callers +// gate the shape (Python prefill_uses_rdna_moe_gemm), so an unsupported shape +// raises via TORCH_CHECK (skinny_gemms-style) rather than returning a fall-back +// status. Reads the weight N-row stride at runtime, so any padding (incl. the +// +128B gfx11x cache-cliff layout) is handled with no separate code path. +// +// The dequant helpers + moe_gemm1_kernel template are tuned for the +// Qwen3.6-35B-A3B prefill MoE shapes. + +#include +#include +#include +#include + +// ----------------------------- device code ----------------------------- +typedef __bf16 bf16x16 __attribute__((ext_vector_type(16))); +typedef float f32x8 __attribute__((ext_vector_type(8))); +// 4 uints = 8 bf16 = 128 bits: one global_load_b128 per lane for the A producer +// load (replaces the per-element 16-bit gather global_load_d16_b16/_hi_b16). +typedef unsigned int v4u __attribute__((ext_vector_type(4))); + +// ---- int4 -> bf16 dequant ("magic" trick) ---- +__device__ __forceinline__ unsigned int two_nibbles_to_bf16( + unsigned int packed) { + unsigned int out; + asm volatile("v_and_or_b32 %0, %1, %2, %3" + : "=v"(out) + : "v"(packed), "v"(0x000F000Fu), "v"(0x43004300u)); + return out; +} +__device__ __forceinline__ void unpack_8_int4(unsigned int* out4, + unsigned int packed) { + out4[0] = two_nibbles_to_bf16(packed); + packed >>= 4; + out4[1] = two_nibbles_to_bf16(packed); + packed >>= 4; + out4[2] = two_nibbles_to_bf16(packed); + packed >>= 4; + out4[3] = two_nibbles_to_bf16(packed); +} + +// The kernel body uses RDNA3 __builtin_amdgcn_wmma_* intrinsics, which only +// exist on gfx11 device passes. __GFX11__ is predefined by the compiler on a +// gfx11 device pass (the same macro skinny_gemms / attention use). Compile the +// real body for gfx11 (and on the host pass, which needs the full definition); +// emit an empty stub for any other device arch so the TU links in non-gfx11 / +// multi-arch builds. The op is never called off gfx11 (gated by on_gfx11()). +#if defined(__HIPCC__) && defined(__GFX11__) + #define MOE_W4A16_GFX11 1 +#endif + +// Template parameters. K, N and the weight N-row stride are RUNTIME args, so a +// single instantiation serves any shape meeting the tile divisibility +// (K % StepK == 0, N % TileN == 0) and any weight N-row padding; only the tile +// shape below is compile-time: +// GroupSize quant group size (scales are [E, N, K / GroupSize]). +// TopK MoE top_k; a power of two, so A's source row is +// token >> log2(TopK). +// TileN N columns one workgroup computes per N-tile +// (requires N % TileN == 0). +// BlockM rows per moe_align block (the alignment block_m). +// StepK K elements consumed per ring stage; a multiple of 16 +// (the WMMA K dim) that divides GroupSize so a stage +// stays within a single quant group. +// LdsPad per-row padding on the LDS tiles that breaks shared- +// memory bank conflicts on the WMMA loads (tuned). +// ProducerWaves waves that copy A + packed weights global -> LDS ring. +// ConsumerWaves waves that dequant int4 -> bf16 and WMMA-accumulate. +// RingBufferDepth number of LDS ring stages (software-pipeline depth). +// MinWorkgroupsPerCU __launch_bounds__ occupancy hint (min blocks per CU). +template +__global__ __launch_bounds__( + (ProducerWaves + ConsumerWaves) * 32, + MinWorkgroupsPerCU) void moe_gemm1_kernel(const __hip_bfloat16* __restrict__ A, // [tokens, K] acts + const int* __restrict__ w_packed, // [E, N, w_row_stride] + const __hip_bfloat16* __restrict__ w_scale, // [E, N, K/G] + const int* __restrict__ sorted_token_ids, // slot -> flat token + const int* __restrict__ expert_ids, // block -> expert + __hip_bfloat16* __restrict__ C, // [slots, N] output + int K, int N, int w_row_stride, + int n_valid_tokens, int n_slots, + int n_blocks) { +#if defined(MOE_W4A16_GFX11) || !defined(__HIP_DEVICE_COMPILE__) + const int quant_groups = K / GroupSize; + const int num_k_steps = K / StepK; + constexpr int M_TILES = BlockM / 16; + constexpr int N_PER_CWAVE = TileN / ConsumerWaves; + constexpr int N_TILES = N_PER_CWAVE / 16; + constexpr int lds_A_stride = StepK + LdsPad; + constexpr int lds_W_stride = TileN + LdsPad; + constexpr int producer_threads = ProducerWaves * 32; + constexpr int A_per_thread = (BlockM * StepK) / producer_threads; + constexpr int W_per_thread = ((StepK / 8) * TileN) / producer_threads; + constexpr int A_VEC = 8; // bf16 per A vector load (v4u -> global_load_b128) + + static_assert(StepK % 16 == 0, + "StepK must be a multiple of 16 (WMMA K-step)"); + static_assert(GroupSize % StepK == 0, + "StepK must divide the quant GroupSize"); + static_assert(StepK <= GroupSize, "assumes >=1 ring stage per quant group"); + static_assert(N_PER_CWAVE % 16 == 0, + "TileN/ConsumerWaves must be a mult of 16"); + static_assert(BlockM % 16 == 0, + "BlockM must be a multiple of 16 (WMMA M-tile)"); + static_assert((TopK & (TopK - 1)) == 0, + "TopK should be a power of two (divide -> shift)"); + static_assert((BlockM * StepK) % producer_threads == 0, + "A tile must distribute evenly over producer threads"); + static_assert(((StepK / 8) * TileN) % producer_threads == 0, + "weight tile must distribute evenly over producer threads"); + static_assert(StepK % A_VEC == 0, "StepK must be divisible by A_VEC"); + static_assert((BlockM * StepK) % (producer_threads * A_VEC) == 0, + "A tile must distribute evenly over producer_threads * A_VEC"); + + const int tid = threadIdx.x; + const int wave = tid >> 5; + const int lane = tid & 31; + const bool producer = (wave < ProducerWaves); + const int consumer_wave = wave - ProducerWaves; + const int wave_n_base = consumer_wave * N_PER_CWAVE; + + __shared__ alignas(16) __bf16 lds_A[RingBufferDepth][BlockM * lds_A_stride]; + __shared__ int lds_W[RingBufferDepth][(StepK / 8) * lds_W_stride]; + __shared__ int s_token_row[BlockM]; + __shared__ int s_token[BlockM]; + __shared__ unsigned char s_valid[BlockM]; + + bf16x16 ones; + { + unsigned int* words = (unsigned int*)&ones; + #pragma unroll + for (int i = 0; i < 8; i++) words[i] = 0x3F803F80u; + } + + #define PRODUCE(stage, k_base) \ + do { \ + v4u a_buf[A_per_thread / A_VEC]; \ + int w_buf[W_per_thread]; \ + _Pragma("unroll") for (int v = 0; v < A_per_thread / A_VEC; v++) { \ + int c = v * producer_threads + tid; \ + int row = (c * A_VEC) / StepK, col = (c * A_VEC) % StepK; \ + a_buf[v] = \ + *(const v4u*)&A[(size_t)s_token_row[row] * K + (k_base) + col]; \ + } \ + _Pragma("unroll") for (int i = 0; i < W_per_thread; i++) { \ + int flat = i * producer_threads + tid; \ + int ncol = flat % TileN, pk = flat / TileN; \ + w_buf[i] = \ + w_packed[w_expert_base + (long)(n_base + ncol) * w_row_stride + \ + ((k_base) >> 3) + pk]; \ + } \ + _Pragma("unroll") for (int v = 0; v < A_per_thread / A_VEC; v++) { \ + int c = v * producer_threads + tid; \ + int row = (c * A_VEC) / StepK, col = (c * A_VEC) % StepK; \ + if constexpr (lds_A_stride % 8 == 0) { \ + *(v4u*)&lds_A[stage][row * lds_A_stride + col] = a_buf[v]; \ + } else { \ + int* d = (int*)&lds_A[stage][row * lds_A_stride + col]; \ + int* s = (int*)&a_buf[v]; \ + _Pragma("unroll") for (int j = 0; j < 4; j++) d[j] = s[j]; \ + } \ + } \ + _Pragma("unroll") for (int i = 0; i < W_per_thread; i++) { \ + int flat = i * producer_threads + tid; \ + int ncol = flat % TileN, pk = flat / TileN; \ + lds_W[stage][pk * lds_W_stride + ncol] = w_buf[i]; \ + } \ + } while (0) + + #define CONSUME(stage, group) \ + do { \ + float scale[N_TILES]; \ + _Pragma("unroll") for (int nt = 0; nt < N_TILES; nt++) scale[nt] = \ + (float) \ + w_scale[scale_expert_base + \ + (long)(n_base + wave_n_base + nt * 16 + (lane & 15)) * \ + quant_groups + \ + (group)]; \ + f32x8 raw_acc[M_TILES * N_TILES]; \ + f32x8 sum_a[M_TILES]; \ + _Pragma("unroll") for (int i = 0; i < M_TILES * N_TILES; i++) \ + raw_acc[i] = (f32x8){0, 0, 0, 0, 0, 0, 0, 0}; \ + _Pragma("unroll") for (int i = 0; i < M_TILES; i++) sum_a[i] = \ + (f32x8){0, 0, 0, 0, 0, 0, 0, 0}; \ + _Pragma("unroll") for (int kk = 0; kk < StepK / 16; kk++) { \ + bf16x16 a_frag[M_TILES], b_frag[N_TILES]; \ + _Pragma("unroll") for (int mt = 0; mt < M_TILES; mt++) { \ + int row = mt * 16 + (lane & 15); \ + _Pragma("unroll") for (int el = 0; el < 16; el++) a_frag[mt][el] = \ + lds_A[stage][row * lds_A_stride + kk * 16 + el]; \ + } \ + _Pragma("unroll") for (int nt = 0; nt < N_TILES; nt++) { \ + int col = wave_n_base + nt * 16 + (lane & 15); \ + int w_lo = lds_W[stage][(kk * 2) * lds_W_stride + col]; \ + int w_hi = lds_W[stage][(kk * 2 + 1) * lds_W_stride + col]; \ + unsigned int* frag = (unsigned int*)&b_frag[nt]; \ + unpack_8_int4(frag, (unsigned int)w_lo); \ + unpack_8_int4(frag + 4, (unsigned int)w_hi); \ + } \ + _Pragma("unroll") for (int mt = 0; mt < M_TILES; mt++) { \ + sum_a[mt] = __builtin_amdgcn_wmma_f32_16x16x16_bf16_w32( \ + a_frag[mt], ones, sum_a[mt]); \ + _Pragma("unroll") for (int nt = 0; nt < N_TILES; nt++) \ + raw_acc[mt * N_TILES + nt] = \ + __builtin_amdgcn_wmma_f32_16x16x16_bf16_w32( \ + a_frag[mt], b_frag[nt], raw_acc[mt * N_TILES + nt]); \ + } \ + } \ + _Pragma("unroll") for (int mt = 0; mt < M_TILES; mt++) \ + _Pragma("unroll") for (int nt = 0; nt < N_TILES; nt++) { \ + float sc = scale[nt]; \ + int tile = mt * N_TILES + nt; \ + _Pragma("unroll") for (int el = 0; el < 8; el++) acc[tile][el] += \ + sc * (raw_acc[tile][el] - 136.0f * sum_a[mt][el]); \ + } \ + } while (0) + + const int block = blockIdx.x % n_blocks; + const int n_tile = blockIdx.x / n_blocks; + const int expert = expert_ids[block]; + // Padding blocks (over-launch up to the EM cap) carry expert_id == -1. + // Skip them exactly like Triton's `if off_experts == -1: return`, so the + // host can launch a sync-free upper-bound grid. + if (expert < 0) return; + const int n_base = n_tile * TileN; + const long w_expert_base = (long)expert * N * w_row_stride; + const long scale_expert_base = (long)expert * N * quant_groups; + + if (tid < BlockM) { + int slot = block * BlockM + tid; + int token = + (slot < n_slots) ? sorted_token_ids[slot] : (n_valid_tokens + 1); + s_valid[tid] = (token < n_valid_tokens) ? 1 : 0; + s_token_row[tid] = (token < n_valid_tokens) ? (token / TopK) : 0; + s_token[tid] = token; + } + __syncthreads(); + + f32x8 acc[M_TILES * N_TILES]; + #pragma unroll + for (int i = 0; i < M_TILES * N_TILES; i++) + acc[i] = (f32x8){0, 0, 0, 0, 0, 0, 0, 0}; + + if (producer) { + #pragma unroll + for (int s = 0; s < RingBufferDepth - 1; s++) + if (s < num_k_steps) PRODUCE(s, s * StepK); + } + __syncthreads(); + for (int step = 0; step < num_k_steps; step++) { + int prefetch = step + (RingBufferDepth - 1); + if (producer && prefetch < num_k_steps) + PRODUCE(prefetch % RingBufferDepth, prefetch * StepK); + if (!producer) CONSUME(step % RingBufferDepth, (step * StepK) / GroupSize); + __syncthreads(); + } + + if (!producer) { + int col_lane = lane & 15, row_parity = lane >> 4; + #pragma unroll + for (int mt = 0; mt < M_TILES; mt++) + #pragma unroll + for (int nt = 0; nt < N_TILES; nt++) { + int tile = mt * N_TILES + nt; + int out_col = n_base + wave_n_base + nt * 16 + col_lane; + #pragma unroll + for (int el = 0; el < 8; el++) { + int out_row = mt * 16 + 2 * el + row_parity; + if (out_row < BlockM && s_valid[out_row]) + C[(size_t)s_token[out_row] * N + out_col] = + (__hip_bfloat16)acc[tile][el]; + } + } + } + #undef PRODUCE + #undef CONSUME +#endif // MOE_W4A16_GFX11 || !__HIP_DEVICE_COMPILE__ +} +// --------------------------- end device code --------------------------- + +// ---- host launch dispatch ---- +// valid_blocks = ceil(ntpp / BlockM); but the op receives n_valid_tokens (= +// M*top_k) and the routing tensors. The grid uses the number of populated +// blocks = expert_ids.numel() (one expert id per BlockM block, the alignment +// block count), matching the standalone harness's valid_blocks. +template +static void launch_moe(const at::Tensor& A, const at::Tensor& w_packed, + const at::Tensor& w_scale, const at::Tensor& sti, + const at::Tensor& eid, at::Tensor& C, int K, int N, + int n_valid_tokens, int n_blocks) { + const int n_slots = (int)sti.size(0); + // Weight N-row stride is read from the tensor, so any padding (incl. the + // gfx11x cache-cliff +128B layout) is handled automatically. + const int w_row_stride = (int)w_packed.stride(1); + // n_blocks = sync-free launch upper bound (EM cap, like Triton). Padding + // blocks read expert_id == -1 and early-return via the in-kernel guard. + constexpr int THREADS = (ProducerWaves + ConsumerWaves) * 32; + dim3 block(THREADS); + dim3 grid((N / TileN) * n_blocks); + hipStream_t stream = c10::hip::getCurrentHIPStreamMasqueradingAsCUDA(); + moe_gemm1_kernel + <<>>( + reinterpret_cast(A.data_ptr()), + w_packed.data_ptr(), + reinterpret_cast(w_scale.data_ptr()), + sti.data_ptr(), eid.data_ptr(), + reinterpret_cast<__hip_bfloat16*>(C.data_ptr()), K, N, w_row_stride, + n_valid_tokens, n_slots, n_blocks); +} + +// Runs a tuned instantiation for the given (top_k, K, N) shape; the in-place C +// is the only output. Callers (the vLLM op + its Python predicate +// prefill_uses_rdna_moe_gemm) gate the shape first, so the TORCH_CHECK guards +// below are "can't happen" invariants -- mirroring skinny_gemms' switch/throw +// dispatch rather than returning a fall-back boolean. +void moe_gemm_w4a16(at::Tensor A, at::Tensor w_packed, at::Tensor w_scale, + at::Tensor sorted_token_ids, at::Tensor expert_ids, + at::Tensor C, int64_t n_valid_tokens, int64_t top_k, + int64_t block_m, int64_t num_blocks) { + TORCH_CHECK(A.scalar_type() == at::kBFloat16, "A must be bf16"); + TORCH_CHECK(C.scalar_type() == at::kBFloat16, "C must be bf16"); + TORCH_CHECK(w_packed.scalar_type() == at::kInt, "w_packed must be int32"); + TORCH_CHECK(w_scale.scalar_type() == at::kBFloat16, "w_scale must be bf16"); + // A, C, w_scale must be contiguous; w_packed only needs a contiguous + // innermost (K//8) dim -- its N-row stride is read at runtime, so any padding + // (incl. the gfx11x cache-cliff +128B layout) is handled automatically. + TORCH_CHECK(A.is_contiguous() && w_scale.is_contiguous() && C.is_contiguous(), + "A, C and w_scale must be contiguous"); + TORCH_CHECK(w_packed.stride(2) == 1, "w_packed K//8 dim must be contiguous"); + + const int K = (int)A.size(1); + const int N = (int)w_packed.size(1); + TORCH_CHECK((int)w_packed.size(2) == K / 8, "w_packed last dim must be K//8"); + // The kernel assumes the expert stride is N * row_stride. + const long w_row_stride = w_packed.stride(1); + TORCH_CHECK(w_packed.stride(0) == (long)N * w_row_stride, + "w_packed expert stride must equal N * row_stride"); + const int G = K / (int)w_scale.size(2); + const int nvt = (int)n_valid_tokens; + int nb = (int)num_blocks; + const int nb_cap = (int)expert_ids.size(0); + if (nb > nb_cap) nb = nb_cap; // never index past the expert_ids buffer + + // Constraint-based tile selection. Two tile families keyed by top_k (the MoE + // structure: gemm1 = routed up/gate proj, top_k>1; gemm2 = down proj, + // top_k==1). Each accepts ANY (K, N) meeting the tile's divisibility -- + // K % G == 0 (implies K % StepK == 0 since StepK | G) and N % TileN == 0 -- + // so the kernel is not pinned to the exact Qwen3.6 shapes. Only G==128 and + // block_m==32 are instantiated. + // launch_moe + TORCH_CHECK(G == 128 && (K % G) == 0 && block_m == 32, + "moe_gemm_w4a16 requires G==128, K%G==0, block_m==32 " + "(got G=", + G, " K=", K, " block_m=", block_m, ")"); + + switch (top_k) { + case 8: // gemm1 family: TileN=256, StepK=128. + TORCH_CHECK((N % 256) == 0, "moe_gemm_w4a16 gemm1 requires N % 256 == 0"); + launch_moe<128, 8, 256, 32, 128, 4, 2, 4, 2, 3>( + A, w_packed, w_scale, sorted_token_ids, expert_ids, C, K, N, nvt, nb); + break; + case 1: // gemm2 family: TileN=512, StepK=64. + TORCH_CHECK((N % 512) == 0, "moe_gemm_w4a16 gemm2 requires N % 512 == 0"); + launch_moe<128, 1, 512, 32, 64, 2, 4, 8, 3, 2>( + A, w_packed, w_scale, sorted_token_ids, expert_ids, C, K, N, nvt, nb); + break; + default: + TORCH_CHECK(false, "moe_gemm_w4a16 unsupported top_k=", top_k, + " (expected 1 or 8)"); + } +} diff --git a/csrc/rocm/ops.h b/csrc/rocm/ops.h index ebd0d84276bf..bc46a99f62aa 100644 --- a/csrc/rocm/ops.h +++ b/csrc/rocm/ops.h @@ -113,6 +113,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); +// W4A16 MoE prefill WMMA GEMM for gfx11 (defined in moe_gemm_w4a16_wmma.cu; +// real body is gfx11-only, stub elsewhere). Mutates C in place; callers gate +// the shape (Python prefill_uses_rdna_moe_gemm), so an unsupported shape raises +// via TORCH_CHECK. +void moe_gemm_w4a16(at::Tensor A, at::Tensor w_packed, at::Tensor w_scale, + at::Tensor sorted_token_ids, at::Tensor expert_ids, + at::Tensor C, int64_t n_valid_tokens, int64_t top_k, + int64_t block_m, int64_t num_blocks); + 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 1d832a5c0dc8..967301155ba0 100644 --- a/csrc/rocm/torch_bindings.cpp +++ b/csrc/rocm/torch_bindings.cpp @@ -155,6 +155,16 @@ TORCH_LIBRARY_EXPAND(TORCH_EXTENSION_NAME, rocm_ops) { rocm_ops.impl("gptq_gemm_rdna3_wmma", torch::kCUDA, &gptq_gemm_rdna3_wmma); #endif + // W4A16 MoE prefill WMMA GEMM for AMD RDNA3 (gfx11). Always registered; the + // kernel body is gfx11-only (stub elsewhere) and Python gates calls on + // on_gfx1x(). Mutates c in place; an unsupported shape raises via TORCH_CHECK + // (callers gate via the Python prefill_uses_rdna_moe_gemm predicate). + rocm_ops.def( + "moe_gemm_w4a16(Tensor a, Tensor w_packed, Tensor w_scale, " + "Tensor sorted_token_ids, Tensor expert_ids, Tensor! c, " + "int n_valid_tokens, int top_k, int block_m, int num_blocks) -> ()"); + rocm_ops.impl("moe_gemm_w4a16", torch::kCUDA, &moe_gemm_w4a16); + // Custom attention op // Compute the attention between an input query and the cached // keys/values using PagedAttention. diff --git a/tests/kernels/quantization/test_moe_gemm_w4a16.py b/tests/kernels/quantization/test_moe_gemm_w4a16.py new file mode 100644 index 000000000000..8e98f2b03cef --- /dev/null +++ b/tests/kernels/quantization/test_moe_gemm_w4a16.py @@ -0,0 +1,136 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""Correctness test for the gfx11 W4A16 MoE prefill WMMA GEMM. + +Validates the compiled op ``torch.ops._rocm_C.moe_gemm_w4a16`` (gated behind +``VLLM_MOE_HIP`` in production, called directly here) against the Triton +reference ``invoke_fused_moe_kernel_hybrid_triton`` on identical inputs, using +the vLLM MoE weight layout (ExLlama shuffle-packed INT4). + +Run ``pytest tests/kernels/quantization/test_moe_gemm_w4a16.py``. +""" + +import pytest +import torch + +# The op is compiled into _rocm_C on all ROCm builds but has a real body only on +# gfx11 (stub elsewhere). Skip cleanly off ROCm, off gfx11, or if the op is +# absent (older build). Probe op presence before importing rocm.py so CUDA skips +# without touching ROCm-only platform code. +try: + import vllm._rocm_C # noqa: F401 + + _have_op = hasattr(torch.ops._rocm_C, "moe_gemm_w4a16") +except Exception: + _have_op = False +if not _have_op: + pytest.skip("_rocm_C.moe_gemm_w4a16 not available", allow_module_level=True) + +from vllm.platforms.rocm import on_gfx11 + +if not on_gfx11(): + pytest.skip("requires gfx11 (RDNA3 WMMA)", allow_module_level=True) + +pytest.importorskip("triton") + +from vllm.model_executor.kernels.linear.mixed_precision.hybrid_w4a16 import ( + pack_int4_exllama_shuffle, +) +from vllm.model_executor.layers.fused_moe import moe_hip_w4a16 # noqa: E402 +from vllm.model_executor.layers.fused_moe.fused_moe import ( + invoke_fused_moe_kernel_hybrid_triton, +) +from vllm.model_executor.layers.fused_moe.moe_align_block_size import ( + moe_align_block_size, +) +from vllm.triton_utils import tl + +DEV = "cuda" + + +@pytest.mark.parametrize("block_m", [32]) +@pytest.mark.parametrize( + "K, N, TOPK", + [ + (2048, 1024, 8), # gemm1 family, tuned Qwen3.6 shape + (1024, 512, 8), # gemm1 family, different K/N -> exercises runtime K,N + (512, 2048, 1), # gemm2 family, tuned Qwen3.6 shape + ], +) +def test_moe_gemm_w4a16_matches_triton(K, N, TOPK, block_m, monkeypatch): + # The rdna_moe_gemm kernel is default-on under gfx11, so force it off here -- the + # reference below dispatches through invoke_fused_moe_kernel_hybrid_triton + # and must stay on the Triton path to be a real reference. + monkeypatch.setenv("VLLM_MOE_HIP", "0") + + E, G = 256, 128 + T = 994 + dt = torch.bfloat16 + torch.manual_seed(0) + + A = torch.randn(T, K, dtype=dt, device=DEV) * 0.1 + w_uint4 = torch.randint(0, 16, (E, N, K), dtype=torch.int32, device=DEV) + w_packed = torch.stack([pack_int4_exllama_shuffle(w_uint4[e]) for e in range(E)]) + w_scale = torch.randn(E, N, K // G, dtype=dt, device=DEV).abs() * 0.01 + + logits = torch.randn(T, E, device=DEV) + _, topk_ids = torch.topk(logits.softmax(-1), TOPK, dim=-1) + topk_ids = topk_ids.to(torch.int32) + + sorted_token_ids, expert_ids, num_tokens_post_padded = moe_align_block_size( + topk_ids, block_m, E, ignore_invalid_experts=True + ) + num_slots = sorted_token_ids.size(0) + + # Triton reference. VLLM_MOE_HIP=0 (set above) makes the in-kernel dispatch + # predicate False so this stays on the Triton path. + assert not moe_hip_w4a16.is_enabled(), ( + "VLLM_MOE_HIP must be forced off so the reference uses Triton" + ) + c_ref = torch.zeros(num_slots, N, dtype=dt, device=DEV) + invoke_fused_moe_kernel_hybrid_triton( + A=A, + B=w_packed, + C=c_ref, + B_scale=w_scale, + topk_weights=None, + sorted_token_ids=sorted_token_ids, + expert_ids=expert_ids, + num_tokens_post_padded=num_tokens_post_padded, + mul_routed_weight=False, + top_k=TOPK, + config=dict( + BLOCK_SIZE_M=block_m, + BLOCK_SIZE_N=64, + BLOCK_SIZE_K=128, + GROUP_SIZE_M=4, + num_warps=2, + num_stages=1, + ), + compute_type=tl.bfloat16, + group_size=G, + align_block_size_m=block_m, + ) + + # rdna_moe_gemm compiled op (called directly; expert_ids.numel() == num blocks). + # The op returns nothing and mutates c in place; an unsupported shape would + # TORCH_CHECK rather than return a status. + c = torch.zeros(num_slots, N, dtype=dt, device=DEV) + torch.ops._rocm_C.moe_gemm_w4a16( + A, + w_packed, + w_scale, + sorted_token_ids, + expert_ids, + c, + T * TOPK, + TOPK, + block_m, + expert_ids.numel(), + ) + torch.accelerator.synchronize() + + rel = ( + (c.float() - c_ref.float()).abs().sum() / (c_ref.float().abs().sum() + 1e-9) + ).item() + assert rel < 0.01, f"rel={rel:.4f} exceeds tolerance vs Triton reference" diff --git a/vllm/envs.py b/vllm/envs.py index 5950757e166b..18a401f14805 100755 --- a/vllm/envs.py +++ b/vllm/envs.py @@ -118,6 +118,7 @@ VLLM_MOE_AWQ_GEMV_HIP: bool = False VLLM_MOE_GPTQ_EXLLAMA: bool = False VLLM_MOE_HYBRID_W4A16: bool = False + VLLM_MOE_HIP: str | None = None VLLM_ROCM_USE_MOE_WNA16_CUDA_KERNEL: bool = False VLLM_ROCM_USE_AITER: bool = False VLLM_ROCM_USE_AITER_PAGED_ATTN: bool = False @@ -1117,6 +1118,10 @@ def _resolve_rust_frontend_path() -> str | None: "VLLM_MOE_GPTQ_EXLLAMA": lambda: ( os.getenv("VLLM_MOE_GPTQ_EXLLAMA", "true").lower() in ("true", "1") ), + # rdna_moe_gemm gfx11 W4A16 MoE prefill WMMA GEMM. Tri-state: unset = + # default-on wherever the kernel is built (gfx11), "1" forces on, "0" forces + # the Triton path. Read via moe_hip_w4a16.is_enabled(). + "VLLM_MOE_HIP": lambda: os.environ.get("VLLM_MOE_HIP", None), # Optional: enable external Oink custom ops (e.g., Blackwell RMSNorm). # Disabled by default. "VLLM_USE_OINK_OPS": lambda: ( diff --git a/vllm/model_executor/layers/fused_moe/hybrid_w4a16_moe.py b/vllm/model_executor/layers/fused_moe/hybrid_w4a16_moe.py index 962ceeae5605..25a4de2e3d02 100644 --- a/vllm/model_executor/layers/fused_moe/hybrid_w4a16_moe.py +++ b/vllm/model_executor/layers/fused_moe/hybrid_w4a16_moe.py @@ -141,12 +141,13 @@ def finalize_weight_and_reduce_impl(self) -> mk.TopKWeightAndReduce: SMALL_GROUP_SIZE_THRESHOLD = 64 def _select_block_size_m(self, num_tokens: int, topk: int, E: int) -> int: - """Select block size in the M dimension. + """Select block size in the M dimension (the moe_align alignment). Decode (num_tokens <= MAX_SKINNY_BATCH_SIZE): small block sizes compatible with the wvSplitK_int4 HIP kernel. - Prefill: TRITON_BLOCK_SIZE_M, or TRITON_BLOCK_SIZE_M_SMALL_GS - for small group_size. + Prefill: TRITON_BLOCK_SIZE_M (or TRITON_BLOCK_SIZE_M_SMALL_GS for small + group_size). The rdna_moe_gemm WMMA gemm1 and the Triton gemm2 share this one + alignment. """ if num_tokens > HybridW4A16MoEExperts.MAX_SKINNY_BATCH_SIZE: if self._group_size <= HybridW4A16MoEExperts.SMALL_GROUP_SIZE_THRESHOLD: @@ -157,6 +158,19 @@ def _select_block_size_m(self, num_tokens: int, topk: int, E: int) -> int: return 4 if avg >= 4 else 2 if avg >= 2 else 1 return 1 + def _prefill_uses_rdna_moe_gemm( + self, N: int, K: int, topk: int, activation: MoEActivation + ) -> bool: + """Whether the rdna_moe_gemm WMMA kernel can run gemm1 for this shape (gemm1 + K-by-N top_k; gemm2 down proj also checked since both must be supported). + Gates the gemm1 rdna_moe_gemm-vs-Triton dispatch in apply().""" + from vllm.model_executor.layers.fused_moe import moe_hip_w4a16 + + act_dim = self.adjust_N_for_activation(N, activation) + return moe_hip_w4a16.prefill_uses_rdna_moe_gemm( + K, N, act_dim, topk, self._group_size + ) + def moe_problem_size( self, a1: torch.Tensor, @@ -325,6 +339,20 @@ def apply( P = num_tokens * top_k_num use_triton = num_tokens > self.MAX_SKINNY_BATCH_SIZE + # rdna_moe_gemm prefill path: gemm1 runs the rdna_moe_gemm WMMA kernel + # and gemm2 runs Triton, both at the single block_size_m alignment (32). + # gemm1 at 32 is faster than at 16 (the better tile beats the extra + # padding), so there is no separate gemm2 alignment. + use_rdna_moe_gemm = use_triton and self._prefill_uses_rdna_moe_gemm( + N, K, top_k_num, activation + ) + # The rdna_moe_gemm kernel cannot fold routing weights on input; that flag is + # only set for top_k==1 layers, which the gemm1 family (top_k>1) never + # matches, so the two are mutually exclusive in practice. + assert not (use_rdna_moe_gemm and apply_router_weight_on_input), ( + "VLLM_MOE_HIP rdna_moe_gemm kernel does not support " + "apply_router_weight_on_input" + ) # ---- Route tokens to experts ---- scattered = block_size_m == 1 @@ -374,22 +402,17 @@ def apply( if w1.dtype == torch.int32 and hidden_states.dtype == torch.float16 else tl.bfloat16 ) - # Compute one config per GEMM: gemm1's contraction is K - # (hidden_dim), gemm2's is activation_out_dim. Pass the - # alignment block_size_m so _triton_config only emits - # BLOCK_SIZE_M values that divide it. + # Per-GEMM Triton config for the fallback (gemm1 contraction is K, + # gemm2's is activation_out_dim). M_routed = num_tokens * top_k_num config_gemm1 = self._triton_config(K, M_routed, block_size_m) config_gemm2 = self._triton_config( activation_out_dim, M_routed, block_size_m ) - # When a config picks a smaller BLOCK_SIZE_M than the - # alignment used, repeat-interleave expert_ids so the kernel - # reads a valid expert id for every (smaller) block. This - # is exact: kernel block 4i..4i+3 with BLOCK_M=32 covers - # the same 128-row range that one BLOCK_M=128 block did, so - # they all process the same expert. + # When a Triton config picks a smaller BLOCK_SIZE_M than the + # alignment, repeat-interleave expert_ids so the kernel reads a valid + # expert id for every (smaller) block. def _expert_ids_for(cfg: dict) -> torch.Tensor: kbm = cfg["BLOCK_SIZE_M"] if kbm == block_size_m: @@ -400,33 +423,48 @@ def _expert_ids_for(cfg: dict) -> torch.Tensor: ) return expert_ids.repeat_interleave(block_size_m // kbm) - # GEMM 1 (Triton prefill path) - # The kernel gathers from hidden_states via - # sorted_token_ids // top_k, so pass the original top_k. - invoke_fused_moe_kernel_hybrid_triton( - A=gemm1_in, - B=w1, - C=gemm1_out, - B_scale=self.w1_scale, - topk_weights=topk_weights if apply_router_weight_on_input else None, - sorted_token_ids=sorted_token_ids, - expert_ids=_expert_ids_for(config_gemm1), - num_tokens_post_padded=num_tokens_post_padded, - mul_routed_weight=apply_router_weight_on_input, - top_k=top_k_num, - config=config_gemm1, - compute_type=compute_type, - group_size=self._group_size, - align_block_size_m=block_size_m, - ) + # The rdna_moe_gemm WMMA kernel (VLLM_MOE_HIP, gfx11) replaces only the + # Triton GEMM1 (identical output layout); GEMM2 always runs on Triton. + # num_blocks == expert_ids.numel() (one id per align block); + # n_valid_tokens == A.rows * top_k, matching Triton. + + # GEMM 1 (gathers from hidden_states via sorted_token_ids // top_k). + if use_rdna_moe_gemm: + torch.ops.vllm.moe_gemm_w4a16( + gemm1_in, + w1, + self.w1_scale, + sorted_token_ids, + expert_ids, + gemm1_out, + num_tokens * top_k_num, + top_k_num, + block_size_m, + expert_ids.numel(), + ) + else: + invoke_fused_moe_kernel_hybrid_triton( + A=gemm1_in, + B=w1, + C=gemm1_out, + B_scale=self.w1_scale, + topk_weights=topk_weights if apply_router_weight_on_input else None, + sorted_token_ids=sorted_token_ids, + expert_ids=_expert_ids_for(config_gemm1), + num_tokens_post_padded=num_tokens_post_padded, + mul_routed_weight=apply_router_weight_on_input, + top_k=top_k_num, + config=config_gemm1, + compute_type=compute_type, + group_size=self._group_size, + align_block_size_m=block_size_m, + ) # Activation (only the [0:P] real rows; padding rows unread) apply_moe_activation(activation, act_out[:P], gemm1_out[:P]) - # GEMM 2 (Triton prefill path) - # act_out is in slot-space (GEMM1 wrote at sorted_token_ids - # positions). Pass top_k=1 so the kernel reads act_out[slot] - # directly. + # GEMM 2 (act_out is compact flat-topk order; top_k=1) -- always + # Triton, on the same alignment as gemm1. invoke_fused_moe_kernel_hybrid_triton( A=act_out, B=w2, diff --git a/vllm/model_executor/layers/fused_moe/moe_hip_w4a16.py b/vllm/model_executor/layers/fused_moe/moe_hip_w4a16.py new file mode 100644 index 000000000000..26d8b5ca6d38 --- /dev/null +++ b/vllm/model_executor/layers/fused_moe/moe_hip_w4a16.py @@ -0,0 +1,123 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""Dispatch for the rdna_moe_gemm W4A16 MoE prefill GEMM (gfx11 / RDNA3). + +Enabled by default on gfx11 builds (where the kernel is compiled); set +``VLLM_MOE_HIP=0`` to force the Triton path or ``VLLM_MOE_HIP=1`` to force-enable. +When active and the routed problem matches one of the tuned Qwen3.6-35B-A3B +prefill MoE shapes, the producer/consumer WMMA kernel replaces the Triton +``fused_moe_kernel_gptq_awq`` for that GEMM. Any other shape, arch, or an +explicit disable falls through to Triton unchanged, so the path is A/B-able. + +The kernel itself is compiled into ``_rocm_C`` +(``csrc/rocm/moe_gemm_w4a16_wmma.cu``; the WMMA body is gfx11-only, a stub +elsewhere) and exposed as ``torch.ops._rocm_C.moe_gemm_w4a16``. This module just +makes the dispatch decision host-side (env + shape predicate) and wraps the call +as a registered vLLM custom op with a no-op fake, so the path is graph-safe under +torch.compile. +""" + +from __future__ import annotations + +import torch + +import vllm.envs as envs +from vllm.platforms.rocm import on_gfx11 +from vllm.utils.torch_utils import direct_register_custom_op + +# The kernel takes K, N and the weight row stride at runtime, so it handles any +# shape meeting the tile constraints and any weight N-row padding. What stays +# fixed (compile-time instantiations) is the tile family keyed by top_k -> +# required N divisor (BN), the group size G, and block_m. Keep in sync with the +# host dispatch in csrc/rocm/moe_gemm_w4a16_wmma.cu. +_TILE_N_DIVISOR_BY_TOPK = {8: 256, 1: 512} # top_k -> BN +_SUPPORTED_G = 128 + + +def is_enabled() -> bool: + # Default-on on gfx11 (the only arch with a real kernel body; the op is a + # stub elsewhere). VLLM_MOE_HIP overrides: "1" forces on, "0" forces Triton. + val = envs.VLLM_MOE_HIP + if val is not None: + return val == "1" + return on_gfx11() + + +def _shape_supported(K: int, N: int, G: int, top_k: int) -> bool: + """Divisibility constraints a (K, N, G, top_k) GEMM must meet for a tuned + tile family (top_k picks BN; K % G == 0 also gives K % GK == 0 since GK|G).""" + bn = _TILE_N_DIVISOR_BY_TOPK.get(top_k) + return bn is not None and G == _SUPPORTED_G and K % G == 0 and N % bn == 0 + + +def prefill_uses_rdna_moe_gemm( + K_hidden: int, N_gemm1: int, act_dim: int, top_k: int, group_size: int +) -> bool: + """True iff BOTH MoE prefill GEMMs will run on the rdna_moe_gemm kernel for this + shape. Shape-only (no tensors), so workspace_shapes and apply() can agree on + block_m up front. gemm1 = (K_hidden, N_gemm1, top_k); gemm2 = down proj + (act_dim, K_hidden, top_k=1).""" + if not is_enabled() or not on_gfx11(): + return False + return _shape_supported(K_hidden, N_gemm1, group_size, top_k) and _shape_supported( + act_dim, K_hidden, group_size, 1 + ) + + +def _moe_gemm_w4a16_impl( + A: torch.Tensor, + B: torch.Tensor, + B_scale: torch.Tensor, + sorted_token_ids: torch.Tensor, + expert_ids: torch.Tensor, + C: torch.Tensor, + n_valid_tokens: int, + top_k: int, + block_size_m: int, + num_blocks: int, +) -> None: + """Run the rdna_moe_gemm WMMA W4A16 MoE GEMM, writing into C in place. + + Callers MUST gate on ``prefill_uses_rdna_moe_gemm`` first (shape-only, so the + block_m choice and this dispatch agree). The kernel validates the shape with + TORCH_CHECK, so a mismatch raises rather than silently falling back. + + num_blocks: sync-free launch upper bound (Triton's EM cdiv block_m); padding + blocks carry expert_id == -1 and early-return in the kernel. + """ + torch.ops._rocm_C.moe_gemm_w4a16( + A, + B, + B_scale, + sorted_token_ids, + expert_ids, + C, + int(n_valid_tokens), + int(top_k), + int(block_size_m), + int(num_blocks), + ) + + +def _moe_gemm_w4a16_fake( + A: torch.Tensor, + B: torch.Tensor, + B_scale: torch.Tensor, + sorted_token_ids: torch.Tensor, + expert_ids: torch.Tensor, + C: torch.Tensor, + n_valid_tokens: int, + top_k: int, + block_size_m: int, + num_blocks: int, +) -> None: + # C is pre-allocated and mutated in place; nothing to allocate. + return None + + +direct_register_custom_op( + op_name="moe_gemm_w4a16", + op_func=_moe_gemm_w4a16_impl, + mutates_args=["C"], + fake_impl=_moe_gemm_w4a16_fake, +) From f50fdea20a1f08f8e6220a00292bf0452fa0d2b7 Mon Sep 17 00:00:00 2001 From: Robert Esclapez Garcia Date: Wed, 24 Jun 2026 07:17:59 -0700 Subject: [PATCH 2/3] [ROCm][MoE] Add roofline profiling scope to the rdna_moe_gemm gemm1 Wrap the rdna_moe_gemm gemm1 launch (torch.ops.vllm.moe_gemm_w4a16) in apply() with an apply()-level record_function scope carrying the dims the roofline tool needs: M, N, K, E, top_k, the quant group size g, block_m, valid_blocks (= num_tokens_post_padded // block_m) and n_routed (M*top_k), plus per-expert tok_hist / vtok_hist histograms describing the routing skew that drives the block_m padding. g is the real self._group_size (the per-group scale bytes and dequant FLOPs scale with K/g), not an assumed 128. The scope is gated on VLLM_CUSTOM_SCOPES_FOR_PROFILING / VLLM_NVTX_SCOPES_FOR_PROFILING -- the valid_blocks/histogram reads force a device->host sync, taken only when profiling; production gets a nullcontext. AI assistance (Claude) was used. Co-authored-by: Claude Signed-off-by: Robert Esclapez Garcia --- .../layers/fused_moe/hybrid_w4a16_moe.py | 98 +++++++++++++++++-- 1 file changed, 88 insertions(+), 10 deletions(-) diff --git a/vllm/model_executor/layers/fused_moe/hybrid_w4a16_moe.py b/vllm/model_executor/layers/fused_moe/hybrid_w4a16_moe.py index 25a4de2e3d02..554f0b11884b 100644 --- a/vllm/model_executor/layers/fused_moe/hybrid_w4a16_moe.py +++ b/vllm/model_executor/layers/fused_moe/hybrid_w4a16_moe.py @@ -10,8 +10,11 @@ CUDA-graph compatible. """ +import contextlib + import torch +import vllm.envs as envs import vllm.model_executor.layers.fused_moe.modular_kernel as mk from vllm.model_executor.layers.fused_moe.activation import ( MoEActivation, @@ -40,6 +43,69 @@ from vllm.v1.utils import record_function_or_nullcontext +def _moe_gemm_w4a16_scope( + M: int, + N: int, + K: int, + E: int, + top_k: int, + group_size: int, + block_m: int, + num_tokens_post_padded: torch.Tensor | None, + n_routed: int, + topk_ids: torch.Tensor | None = None, +): + """record_function scope around an rdna_moe_gemm (moe_gemm_w4a16) launch, + surfacing the dims the roofline tool needs -- including the quantities no + tensor shape exposes: the quant group size g (= group_size; the per-group + scale bytes and dequant FLOPs scale with K/g, so it must be the real value, + not a 128 default), valid_blocks (non-padding blocks the kernel actually + runs = num_tokens_post_padded // block_m), n_routed (M*top_k useful rows), + and a tokens-per-expert histogram (tok_hist): the full distribution of how + many block_m blocks each expert fills, as counts indexed by blocks (index + 0 = inactive experts, index i = experts filling i blocks) -- the routing + skew that drives the block_m padding. + + These are only known at runtime (they depend on the routing), so reading them + forces a device->host sync. That sync only perturbs the CPU timeline, not the + GPU kernel duration the roofline measures, and we take it solely when + profiling scopes are enabled -- production (where the context is a + nullcontext) pays nothing. + """ + if not ( + envs.VLLM_CUSTOM_SCOPES_FOR_PROFILING or envs.VLLM_NVTX_SCOPES_FOR_PROFILING + ): + return contextlib.nullcontext() + valid_blocks = ( + int(num_tokens_post_padded.item()) // block_m + if num_tokens_post_padded is not None + else 0 + ) + hist_str = "" + if topk_ids is not None: + ids = topk_ids.flatten() + ids = ids[(ids >= 0) & (ids < E)] + counts = torch.bincount(ids, minlength=E) # tokens per expert + blocks = (counts + (block_m - 1)) // block_m # blocks filled (0=inactive) + # tok_hist: full distribution, index i = #experts filling i blocks. + hist = torch.bincount(blocks).tolist() # one sync + hist_str = " tok_hist=" + ",".join(str(int(h)) for h in hist) + # vtok_hist: valid tokens per block (block occupancy), index v = + # #blocks holding exactly v real tokens (1..block_m). Each active expert + # contributes (nblk-1) full blocks + one last block with last_valid; the + # gap to block_m is the padding the kernel computes but discards. + active = counts[counts > 0] + nblk = (active + (block_m - 1)) // block_m + last_valid = active - (nblk - 1) * block_m # 1..block_m + vhist = torch.bincount(last_valid, minlength=block_m + 1) + vhist[block_m] = vhist[block_m] + (nblk - 1).sum() # add the full blocks + hist_str += " vtok_hist=" + ",".join(str(int(h)) for h in vhist.tolist()) + return record_function_or_nullcontext( + f"moe_gemm_w4a16 {M}x{N}x{K} E={E} top_k={top_k} g={group_size} " + f"block_m={block_m} valid_blocks={valid_blocks} n_routed={n_routed}{hist_str}" + ) + + class HybridW4A16MoEExperts(mk.FusedMoEExpertsModular): """MoE experts using fused wvSplitK_int4_g kernel. @@ -430,18 +496,30 @@ def _expert_ids_for(cfg: dict) -> torch.Tensor: # GEMM 1 (gathers from hidden_states via sorted_token_ids // top_k). if use_rdna_moe_gemm: - torch.ops.vllm.moe_gemm_w4a16( - gemm1_in, - w1, - self.w1_scale, - sorted_token_ids, - expert_ids, - gemm1_out, - num_tokens * top_k_num, + with _moe_gemm_w4a16_scope( + num_tokens, + N, + K, + global_num_experts, top_k_num, + self._group_size, block_size_m, - expert_ids.numel(), - ) + num_tokens_post_padded, + num_tokens * top_k_num, + topk_ids, + ): + torch.ops.vllm.moe_gemm_w4a16( + gemm1_in, + w1, + self.w1_scale, + sorted_token_ids, + expert_ids, + gemm1_out, + num_tokens * top_k_num, + top_k_num, + block_size_m, + expert_ids.numel(), + ) else: invoke_fused_moe_kernel_hybrid_triton( A=gemm1_in, From b9faa78c2f4aaf55986f2ae0a26b9f81d9594243 Mon Sep 17 00:00:00 2001 From: Robert Esclapez Garcia Date: Thu, 25 Jun 2026 03:48:51 -0700 Subject: [PATCH 3/3] [ROCm][MoE] Add W4A16 MoE routing-distribution benchmark suite Adds benchmarks/kernels/moe_w4a16_bench/, a routing-distribution harness for the gfx11 W4A16 MoE prefill GEMM (rdna_moe_gemm): - gen_distributions.py: synthetic routing archetypes (balanced/uniform/ zipf1/zipf2/hot16) + a moe_align padding/fill space-map. synth_topk() is importable so the benchmark builds routing in-process. - bench_moe_gemm.py: times torch.ops._rocm_C.moe_gemm_w4a16 against the Triton reference (invoke_fused_moe_kernel_hybrid_triton) on identical inputs across the archetypes, for both prefill GEMMs (--gemm 1,2), with Triton as the correctness oracle. Co-authored-by: Claude Signed-off-by: Robert Esclapez Garcia --- .../kernels/moe_w4a16_bench/bench_moe_gemm.py | 208 ++++++++++++++++++ .../moe_w4a16_bench/gen_distributions.py | 133 +++++++++++ 2 files changed, 341 insertions(+) create mode 100644 benchmarks/kernels/moe_w4a16_bench/bench_moe_gemm.py create mode 100644 benchmarks/kernels/moe_w4a16_bench/gen_distributions.py diff --git a/benchmarks/kernels/moe_w4a16_bench/bench_moe_gemm.py b/benchmarks/kernels/moe_w4a16_bench/bench_moe_gemm.py new file mode 100644 index 000000000000..b74750591145 --- /dev/null +++ b/benchmarks/kernels/moe_w4a16_bench/bench_moe_gemm.py @@ -0,0 +1,208 @@ +#!/usr/bin/env python3 +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""rdna_moe_gemm-vs-Triton across routing archetypes, for both MoE prefill GEMMs. + +gemm1 (up/gate, top_k=8, K=hidden=2048 N=intermediate=1024) and gemm2 (down, +top_k=1, K=intermediate=512 N=hidden=2048) of the W4A16 MoE share a single +moe_align@32 layout. For each (archetype, T) from ``gen_distributions.synth_topk`` +this times ``torch.ops._rocm_C.moe_gemm_w4a16`` against the Triton reference +(``invoke_fused_moe_kernel_hybrid_triton``) on identical inputs, so the kernel's +win/loss can be correlated with the per-block fill of the routing. Triton is the +correctness oracle (``rel`` = kernel-vs-Triton). + +The W4A16 weights/scales are routing-independent, so they are built once per gemm +and reused; only the routing and A change per case. Requires gfx11 (RDNA3 WMMA); +off gfx11 the compiled op is a stub. + + python bench_moe_gemm.py # both gemms, all archetypes + python bench_moe_gemm.py --gemm 1 # gemm1 only + python bench_moe_gemm.py --arches zipf2,hot16 --ts 256,994 +""" + +import os +import sys + +import numpy as np +import torch + +# Import the sibling routing generator (same directory). +sys.path.insert(0, os.path.dirname(os.path.abspath(__file__))) + +from gen_distributions import Archetype, synth_topk # noqa: E402 + +from vllm.model_executor.kernels.linear.mixed_precision.hybrid_w4a16 import ( # noqa: E402 + pack_int4_exllama_shuffle, +) +from vllm.model_executor.layers.fused_moe.fused_moe import ( # noqa: E402 + invoke_fused_moe_kernel_hybrid_triton, +) +from vllm.model_executor.layers.fused_moe.moe_align_block_size import ( # noqa: E402 + moe_align_block_size, +) +from vllm.triton_utils import tl, triton # noqa: E402 +from vllm.utils.argparse_utils import FlexibleArgumentParser # noqa: E402 + +# Shared routing/quant params (must match gen_distributions: E=256, top_k=8). +E, ROUTING_TOPK, G = 256, 8, 128 +BLOCK_M = 32 # the only block_m the rdna_moe_gemm op supports +DEV, DT = "cuda", torch.bfloat16 + +# Per-gemm shape and tuned Triton reference config (block_m=32). ``op_top_k`` is +# the op's top_k arg; ``a_rows(T)`` is the number of activation rows the gemm +# consumes (gemm1 reads per-token, gemm2 reads per-routed-slot). +GEMM_SPECS = { + 1: dict(K=2048, N=1024, op_top_k=ROUTING_TOPK, tri_k=128, tri_gm=4, a_rows=1), + 2: dict(K=512, N=2048, op_top_k=1, tri_k=64, tri_gm=1, a_rows=ROUTING_TOPK), +} + + +def time_us(fn): + # do_bench flushes the L2 cache between iterations, so each launch starts + # cold -- it measures one isolated launch (the production pattern, where the + # same gemm does not run back-to-back) rather than rewarding cross-launch + # weight cache-reuse that real inference never sees. + return triton.testing.do_bench(fn, return_mode="median") * 1000.0 + + +def ramp_clocks(): + # This APU starts at low clocks; without a ramp the first do_bench call reads + # inflated and biases whichever op runs first. + w = torch.randn(4096, 4096, dtype=DT, device=DEV) + for _ in range(60): + w = w @ w.t() * 1e-4 + 0.1 + torch.accelerator.synchronize() + + +def run_gemm(gemm, arches, ts, seed): + spec = GEMM_SPECS[gemm] + K, N = spec["K"], spec["N"] + op_top_k = spec["op_top_k"] + triton_cfg = dict( + BLOCK_SIZE_M=BLOCK_M, + BLOCK_SIZE_N=64, + BLOCK_SIZE_K=spec["tri_k"], + GROUP_SIZE_M=spec["tri_gm"], + num_warps=2, + num_stages=1, + ) + + torch.manual_seed(seed) + rng = np.random.default_rng(seed) # routing rng (so --seed varies routing too) + print(f"\n### gemm{gemm}: building weights E={E} N={N} K={K} (packed int4)...") + w_uint4 = torch.randint(0, 16, (E, N, K), dtype=torch.int32, device=DEV) + w_packed = torch.stack([pack_int4_exllama_shuffle(w_uint4[e]) for e in range(E)]) + w_scale = torch.randn(E, N, K // G, dtype=DT, device=DEV).abs() * 0.01 + del w_uint4 + + hdr = ( + f"{'arch':9s} {'T':>5s} {'nvt':>6s} {'blocks':>6s} {'use%':>5s} | " + f"{'triton_us':>9s} {'kernel_us':>9s} {'speedup%':>8s} " + f"{'k_TFLOP/s':>9s} {'rel':>7s}" + ) + print(hdr) + print("-" * len(hdr)) + + rows = [] + for name in arches: + for T in ts: + topk_ids = ( + torch.from_numpy(synth_topk(name, T, rng=rng)).to(torch.int32).to(DEV) + ) + sti, eid, ntpp_t = moe_align_block_size( + topk_ids, BLOCK_M, E, ignore_invalid_experts=True + ) + ntpp = int(ntpp_t.item()) + num_slots = sti.size(0) + nvt = T * ROUTING_TOPK # useful routed rows (== op's n_valid_tokens) + # valid_blocks = num_tokens_post_padded // block_m; the kernel computes + # every row of these blocks, so padded_rows is the real GEMM work. + valid_blocks = ntpp // BLOCK_M + padded_rows = valid_blocks * BLOCK_M + use = 100.0 * nvt / padded_rows if padded_rows else 0.0 + A = torch.randn(spec["a_rows"] * T, K, dtype=DT, device=DEV) * 0.1 + + C_ref = torch.zeros(num_slots, N, dtype=DT, device=DEV) + + def run_triton(C=C_ref, sti=sti, eid=eid, ntpp_t=ntpp_t, A=A): + invoke_fused_moe_kernel_hybrid_triton( + A=A, + B=w_packed, + C=C, + B_scale=w_scale, + topk_weights=None, + sorted_token_ids=sti, + expert_ids=eid, + num_tokens_post_padded=ntpp_t, + mul_routed_weight=False, + top_k=op_top_k, + config=triton_cfg, + compute_type=tl.bfloat16, + group_size=G, + align_block_size_m=BLOCK_M, + ) + + run_triton() + torch.accelerator.synchronize() + C_ref = C_ref.clone() # freeze reference + + C_k = torch.zeros(num_slots, N, dtype=DT, device=DEV) + + def run_kernel(C=C_k, sti=sti, eid=eid, nb=valid_blocks, A=A, nvt=nvt): + torch.ops._rocm_C.moe_gemm_w4a16( + A, w_packed, w_scale, sti, eid, C, nvt, op_top_k, BLOCK_M, nb + ) + + run_kernel() + torch.accelerator.synchronize() + rel = ( + (C_k.float() - C_ref.float()).abs().sum() + / (C_ref.float().abs().sum() + 1e-9) + ).item() + + t_us = time_us(run_triton) + k_us = time_us(run_kernel) + spd = (t_us / k_us - 1) * 100.0 + # Real GEMM work = padded_rows (valid_blocks * block_m), not just nvt: + # the kernel computes the padded rows too. FLOPs = 2 * M * N * K. + k_tflops = 2 * padded_rows * N * K / (k_us * 1e6) + flag = "" if rel < 0.01 else " BADREL" + print( + f"{name.value:9s} {T:5d} {nvt:6d} {valid_blocks:6d} {use:5.0f} | " + f"{t_us:9.1f} {k_us:9.1f} {spd:+8.1f} {k_tflops:9.1f} " + f"{rel:7.4f}{flag}", + flush=True, + ) + rows.append((spd, name, T, use, t_us, k_us, rel)) + + print(f"\n==== gemm{gemm} kernel vs Triton, worst-first ====") + for spd, name, T, use, t_us, k_us, rel in sorted(rows): + print( + f" {spd:+6.1f}% {name.value:9s} T={T:<5d} use={use:3.0f}% " + f"(triton {t_us:.0f} / kernel {k_us:.0f} us) rel={rel:.4f}" + ) + + +def main(): + ap = FlexibleArgumentParser(description=__doc__) + ap.add_argument("--gemm", default="1,2", help="which gemms to bench: 1, 2, or 1,2") + ap.add_argument("--arches", default=",".join(a.value for a in Archetype)) + ap.add_argument("--ts", default="16,64,128,256,512,768,994,1024,1536,2048,4096") + ap.add_argument("--seed", type=int, default=0) + args = ap.parse_args() + + if not torch.cuda.is_available(): + raise SystemExit("HIP/CUDA not available") + if not hasattr(torch.ops._rocm_C, "moe_gemm_w4a16"): + raise SystemExit("torch.ops._rocm_C.moe_gemm_w4a16 not available (need gfx11)") + + gemms = [int(g) for g in args.gemm.split(",")] + arches = [Archetype(x) for x in args.arches.split(",")] + ts = [int(x) for x in args.ts.split(",")] + ramp_clocks() + for gemm in gemms: + run_gemm(gemm, arches, ts, args.seed) + + +if __name__ == "__main__": + main() diff --git a/benchmarks/kernels/moe_w4a16_bench/gen_distributions.py b/benchmarks/kernels/moe_w4a16_bench/gen_distributions.py new file mode 100644 index 000000000000..0fd7b65914f7 --- /dev/null +++ b/benchmarks/kernels/moe_w4a16_bench/gen_distributions.py @@ -0,0 +1,133 @@ +#!/usr/bin/env python3 +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""Synthetic MoE routing-distribution generator + padding/fill space-map. + +Covers the space of token->expert distributions and batch sizes that real +inputs can hit, so the kernel-vs-Triton benchmarks (``bench_archetype_map.py``, +``bench_gemm2.py``) can probe the win/loss surface keyed by the routing shape +and the number of valid tokens (num_valid = T*top_k). + +``synth_topk(name, T)`` is importable: the benchmarks call it to build routing +in-process, so there is no on-disk dataset step. Running this file as a script +prints a per-archetype padding/fill space-map. + +Archetypes (per-expert popularity over E experts): + - balanced: deterministic round-robin, every expert gets an equal share. + - uniform: near-balanced random (Poisson variance around the mean). + - zipf1: moderate power-law skew (1/r). + - zipf2: heavy power-law skew (1/r^2), a few hot experts. + - hot16: near-collapse, ~16 hot experts take almost all tokens. + +Block-fill model matches moe_align_block_size: an expert with c tokens occupies +ceil(c/BM) blocks; the first ceil-1 are full (BM rows), the last holds c%BM rows. +""" + +import argparse # plain argparse: this file is numpy-only, no torch/vllm import +from enum import Enum + +import numpy as np + +E, TOPK = 256, 8 +TS = [16, 64, 128, 256, 512, 768, 994, 1024, 1536, 2048, 4096] + + +class Archetype(str, Enum): + """Token->expert routing distributions, from flat to heavily skewed.""" + + BALANCED = "balanced" # deterministic round-robin, every expert equal + UNIFORM = "uniform" # near-balanced random (Poisson variance) + ZIPF1 = "zipf1" # moderate power-law skew (1/r) + ZIPF2 = "zipf2" # heavy power-law skew (1/r^2), a few hot experts + HOT16 = "hot16" # near-collapse, ~16 hot experts take almost all tokens + + +def popularity(arch): + """Per-expert popularity weights over the E experts (len E), or None for + the deterministic ``BALANCED`` archetype.""" + r = np.arange(1, E + 1) + if arch is Archetype.UNIFORM: + return np.ones(E) + if arch is Archetype.ZIPF1: + return 1.0 / r**1.0 + if arch is Archetype.ZIPF2: + return 1.0 / r**2.0 + if arch is Archetype.HOT16: + return np.where(r <= 16, 1.0, 1e-6) + if arch is Archetype.BALANCED: + return None + raise ValueError(arch) + + +def synth_topk(arch, T, rng=None): + """Per-token top_k distinct experts drawn from the archetype popularity. + + Returns an int32 ``[T, TOPK]`` array. ``rng`` lets callers fix the seed. + """ + if rng is None: + rng = np.random.default_rng(0) + if arch is Archetype.BALANCED: + # round-robin so every expert gets as equal a share as possible + flat = np.tile(np.arange(E), (T * TOPK + E - 1) // E)[: T * TOPK] + return flat.reshape(T, TOPK).astype(np.int32) + logp = np.log(popularity(arch) + 1e-12) + toks = np.empty((T, TOPK), dtype=np.int32) + for t in range(T): # gumbel-top-k = sample top_k w/o replacement ~ popularity + g = rng.gumbel(size=E) + toks[t] = np.argpartition(-(logp + g), TOPK)[:TOPK] + return toks + + +def counts(toks): + return np.bincount(toks.reshape(-1), minlength=E) + + +def pad_stats(cnt, BM): + """useful% = real rows / padded rows; sk8% = fraction of blocks with <=8 fill.""" + used = cnt[cnt > 0] + if used.size == 0: + return dict(useful=0.0, sk8=0.0) + blocks = -(-used // BM) # ceil + fills = [] + for c, b in zip(used, blocks): + fills += [BM] * (b - 1) + [int(c - (b - 1) * BM)] + fills = np.array(fills) + padr = int(blocks.sum()) * BM + return dict(useful=100 * int(cnt.sum()) / padr, sk8=100 * (fills <= 8).mean()) + + +def main(): + ap = argparse.ArgumentParser(description=__doc__) + ap.add_argument("--arches", default=",".join(a.value for a in Archetype)) + ap.add_argument("--ts", default=",".join(str(t) for t in TS)) + ap.add_argument("--seed", type=int, default=0) + args = ap.parse_args() + + arches = [Archetype(x) for x in args.arches.split(",")] + ts = [int(x) for x in args.ts.split(",")] + rng = np.random.default_rng(args.seed) + + hdr = ( + f"{'archetype':10s} {'T':>5s} {'nvalid':>7s} | " + f"{'mean':>5s} {'med':>4s} {'max':>5s} {'dead':>4s} | " + f"{'bm16 use%':>9s} {'sk8%':>5s} | {'bm32 use%':>9s} {'sk8%':>5s}" + ) + print(hdr) + print("-" * len(hdr)) + for name in arches: + for T in ts: + cnt = counts(synth_topk(name, T, rng=rng)) + u = cnt[cnt > 0] + s16, s32 = pad_stats(cnt, 16), pad_stats(cnt, 32) + print( + f"{name.value:10s} {T:5d} {T * TOPK:7d} | {u.mean():5.1f} " + f"{int(np.median(u)):4d} {int(u.max()):5d} " + f"{int((cnt == 0).sum()):4d} | " + f"{s16['useful']:8.1f} {s16['sk8']:5.0f} | " + f"{s32['useful']:8.1f} {s32['sk8']:5.0f}" + ) + print() + + +if __name__ == "__main__": + main()