From 1fc24bdb97c74730bf73ba269165d96438eeaeb9 Mon Sep 17 00:00:00 2001 From: JartX Date: Wed, 24 Jun 2026 23:17:15 +0200 Subject: [PATCH 01/36] [ROCm][gfx1100] Add native HIP MXFP4 dense weight-only GEMM for RDNA3 MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit RDNA3 has no native FP4 matmul, and vLLM's dense MXFP4 path falls back to Marlin (CUDA-only) or Quark emulation (dequant + torch.linear, which loses the 4-bit VRAM saving). This adds a fused weight-only MXFP4 kernel for gfx1100/1101/1102 written as native HIP (hand-written WMMA + scalar, no Triton or external library), dequantizing E2M1 weights on the fly into the v_wmma_f32_16x16x16_{f16,bf16} pipeline (activations stay 16-bit, W4A16). E2M1 -> {bf16,fp16} bits is a near-direct field copy, and the E8M0 block scale (power of two, group 32) folds in as an exponent add rather than a multiply. Compressed-tensors [N, K/2] weights reinterpret as little-endian int32 directly into the [K/8, N] layout the kernel reads, so the repack is a view + transpose. One op (mxfp4_gemm_rdna3) dispatches internally: - M <= 8 (decode): scalar GEMV (templated M_COUNT 1/2/4/8, split-K atomic) - M > 8 (prefill/batch): WMMA tile ladder 16x16 -> 128x64 A Rdna3MxFp4LinearKernel wires it into CompressedTensorsW4A4Mxfp4 on gfx11xx, replacing the Marlin/emulation fallback. Tests: tests/kernels/quantization/test_mxfp4_rdna3.py validates the GEMM against a torch dequant reference (bf16 + fp16, M=1..128) — all pass on a 7900 XTX. End-to-end on a Qwen3-8B checkpoint quantized to MXFP4: 1.48x faster decode and 2.7x smaller weights vs the bf16 baseline, coherent output. Signed-off-by: JartX --- CMakeLists.txt | 1 + csrc/rocm/mxfp4_gemm_rdna3.cu | 1003 +++++++++++++++++ csrc/rocm/ops.h | 3 + csrc/rocm/qdq_mxfp4_rdna3.cuh | 121 ++ csrc/rocm/torch_bindings.cpp | 5 + .../kernels/quantization/test_mxfp4_rdna3.py | 68 ++ .../model_executor/kernels/linear/__init__.py | 6 + .../kernels/linear/mxfp4/rocm.py | 89 ++ 8 files changed, 1296 insertions(+) create mode 100644 csrc/rocm/mxfp4_gemm_rdna3.cu create mode 100644 csrc/rocm/qdq_mxfp4_rdna3.cuh create mode 100644 tests/kernels/quantization/test_mxfp4_rdna3.py create mode 100644 vllm/model_executor/kernels/linear/mxfp4/rocm.py diff --git a/CMakeLists.txt b/CMakeLists.txt index 36b8e66f2c64..d20cd6e8a995 100644 --- a/CMakeLists.txt +++ b/CMakeLists.txt @@ -1352,6 +1352,7 @@ if(VLLM_GPU_LANG STREQUAL "HIP") list(APPEND VLLM_ROCM_EXT_SRC "csrc/rocm/q_gemm_rdna3.cu" "csrc/rocm/q_gemm_rdna3_wmma.cu" + "csrc/rocm/mxfp4_gemm_rdna3.cu" "csrc/rocm/moe_q_gemm_rdna3.cu") endif() diff --git a/csrc/rocm/mxfp4_gemm_rdna3.cu b/csrc/rocm/mxfp4_gemm_rdna3.cu new file mode 100644 index 000000000000..72de607410d1 --- /dev/null +++ b/csrc/rocm/mxfp4_gemm_rdna3.cu @@ -0,0 +1,1003 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright contributors to the vLLM project +// +// MXFP4 W4A16 GEMM for RDNA3 (gfx1100/1101/1102), WMMA path. Forked from the +// GPTQ WMMA kernel (q_gemm_rdna3_wmma.cu): identical f16/bf16 tiling, only the +// B-tile dequant differs (E2M1 + E8M0 group-32, no zeros). Weights are repacked +// Python-side into the [K/8, N] uint32 layout the GPTQ kernel reads. + +#include + +#include +#include +#include + +#include +#include +#include + +#include "qdq_mxfp4_rdna3.cuh" + +#if defined(__HIPCC__) && defined(__gfx1100__) + #define __HIP__RDNA3__ +#endif + +namespace vllm { +namespace mxfp4_rdna3_wmma { + +using bf16_t = __hip_bfloat16; +using mxfp4_rdna3::dequant_mxfp4_8_bf16; +using mxfp4_rdna3::dequant_mxfp4_8_f32; +using mxfp4_rdna3::dequant_mxfp4_8_fp16; +using mxfp4_rdna3::mxfp4_e8m0_bias; + +// K-split heuristic: more blocks along gridDim.z → more resident waves. Mirror +// the GPTQ kernel's K-only split (good enough for the v1 path). +__host__ __device__ static inline int compute_mxfp4_k_split(int size_k) { + if (size_k >= 4096) return 4; + if (size_k >= 2048) return 2; + return 1; +} + +#if defined(__HIP__RDNA3__) || !defined(__HIP_DEVICE_COMPILE__) + +// Native AMDGPU vector types expected by the WMMA built-ins. +using v16fp16 = _Float16 __attribute__((ext_vector_type(16))); +using v16bf16 = __bf16 __attribute__((ext_vector_type(16))); +using v8fp32 = float __attribute__((ext_vector_type(8))); + +__device__ __forceinline__ v8fp32 wmma_mma(v16fp16 a, v16fp16 b, v8fp32 c) { + return __builtin_amdgcn_wmma_f32_16x16x16_f16_w32(a, b, c); +} +__device__ __forceinline__ v8fp32 wmma_mma(v16bf16 a, v16bf16 b, v8fp32 c) { + return __builtin_amdgcn_wmma_f32_16x16x16_bf16_w32(a, b, c); +} + +template +struct WmmaNative; +template <> +struct WmmaNative { + using elem = _Float16; + using v16 = v16fp16; +}; +template <> +struct WmmaNative { + using elem = __bf16; + using v16 = v16bf16; +}; + +template +__device__ __forceinline__ TO bitcast_elem(FROM x) { + static_assert(sizeof(FROM) == sizeof(TO), + "bitcast_elem requires equal-sized types"); + TO r; + __builtin_memcpy(&r, &x, sizeof(TO)); + return r; +} + +// Packed atomic add for the K-split (gridDim.z > 1) epilogue: CAS-32 retry on +// a uint32 word covering two adjacent fp16/bf16 output cells. +__device__ __forceinline__ void atomic_add_pk_f16(half2* addr, half2 val) { + uint32_t* a32 = reinterpret_cast(addr); + uint32_t old = *a32, assumed; + do { + assumed = old; + half2 cur = bitcast_elem(assumed); + half2 sum = __hadd2(cur, val); + old = atomicCAS(a32, assumed, bitcast_elem(sum)); + } while (assumed != old); +} +__device__ __forceinline__ void atomic_add_pk_bf16(__hip_bfloat162* addr, + __hip_bfloat162 val) { + uint32_t* a32 = reinterpret_cast(addr); + uint32_t old = *a32, assumed; + do { + assumed = old; + __hip_bfloat162 cur = bitcast_elem(assumed); + __hip_bfloat162 sum = __hadd2(cur, val); + old = atomicCAS(a32, assumed, bitcast_elem<__hip_bfloat162, uint32_t>(sum)); + } while (assumed != old); +} + +#endif // helpers guard + +#if defined(__HIP__RDNA3__) || !defined(__HIP_DEVICE_COMPILE__) + +// v1: 16M x 16N tile, 1 wave, full K. Fragment layout matches the GPTQ v1 +// kernel; only the B-tile dequant differs. Handles all M (fallback for M<32). +template +__global__ void gemm_mxfp4_wmma_kernel_16x16_1w( + const T* __restrict__ a, const uint32_t* __restrict__ b_q, + const uint8_t* __restrict__ b_scales_e8m0, T* __restrict__ c, + const int size_m, const int size_n, const int size_k, const int groups) { + using E = typename WmmaNative::elem; + using V16 = typename WmmaNative::v16; + + const int m_tile = blockIdx.y * 16; + const int n_tile = blockIdx.x * 16; + if (m_tile >= size_m || n_tile >= size_n) return; + + const int lane = threadIdx.x; // 0..31 + const int lane_lo = lane & 15; // row index within fragment + const int lane_hi = lane >> 4; // 0 or 1 + + v8fp32 c_acc = {0, 0, 0, 0, 0, 0, 0, 0}; + + const int groupsize = size_k / groups; // == 32 for MXFP4 + constexpr uint32_t mant_bits = std::is_same::value ? 10u : 7u; + + const int k_per_split = size_k / gridDim.z; + const int k_start = blockIdx.z * k_per_split; + const int k_end = k_start + k_per_split; + + __shared__ T b_lds[16][16]; + + for (int k_tile = k_start; k_tile < k_end; k_tile += 16) { + // ---- Dequant 16x16 B tile into LDS ---- + // 32 lanes = 16 N cols × 2 K-octets per col. + const int my_n = lane_lo; + const int my_k_octet = lane_hi; // 0 → K[0..7], 1 → K[8..15] + const int actual_n = n_tile + my_n; + + if (actual_n < size_n) { + const int qk_row = (k_tile / 8) + my_k_octet; + const uint32_t qa = b_q[qk_row * size_n + actual_n]; + + // group_size == 32, k_tile steps by 16 → both octets of a tile share the + // same group, and group = k_tile / groupsize is exact. + const int g = k_tile / groupsize; + const uint32_t s8 = b_scales_e8m0[g * size_n + actual_n]; + const int32_t bias_u16 = mxfp4_e8m0_bias(s8, mant_bits); + + const int k_base = my_k_octet * 8; + + if constexpr (std::is_same::value) { + half dq[8]; + dequant_mxfp4_8_fp16(qa, bias_u16, dq); + #pragma unroll + for (int i = 0; i < 8; i++) b_lds[k_base + i][my_n] = dq[i]; + } else { + bf16_t dq[8]; + dequant_mxfp4_8_bf16(qa, bias_u16, dq); + #pragma unroll + for (int i = 0; i < 8; i++) b_lds[k_base + i][my_n] = dq[i]; + } + } + + // Single-wave block: no __syncthreads needed. + V16 a_frag, b_frag; + const int m_row = m_tile + lane_lo; + + if (m_row < size_m) { + const T* a_row = a + m_row * size_k; + static_assert(sizeof(a_frag) == 32, "V16 must be 32 bytes (16 × 2)"); + __builtin_memcpy(&a_frag, a_row + k_tile, sizeof(a_frag)); + } else { + #pragma unroll + for (int i = 0; i < 16; i++) a_frag[i] = (E)0; + } + + #pragma unroll + for (int i = 0; i < 16; i++) { + b_frag[i] = bitcast_elem(b_lds[i][lane_lo]); + } + + c_acc = wmma_mma(a_frag, b_frag, c_acc); + } + + // ---- Store C ---- (lane = N column, slot i → row m = 2*i + lane_hi) + if (gridDim.z > 1) { + const bool is_even_lane = (lane_lo & 1) == 0; + const int out_n_pair = n_tile + lane_lo; + #pragma unroll + for (int i = 0; i < 8; i++) { + float other_f = __shfl_xor(c_acc[i], 1); + if (!is_even_lane) continue; + const int out_m = m_tile + 2 * i + lane_hi; + if (out_m >= size_m || out_n_pair >= size_n) continue; + T* dst = c + out_m * size_n + out_n_pair; + if constexpr (std::is_same::value) { + half2 packed = + __halves2half2(__float2half_rn(c_acc[i]), __float2half_rn(other_f)); + atomic_add_pk_f16(reinterpret_cast(dst), packed); + } else { + __hip_bfloat162 packed; + packed.x = __float2bfloat16(c_acc[i]); + packed.y = __float2bfloat16(other_f); + atomic_add_pk_bf16(reinterpret_cast<__hip_bfloat162*>(dst), packed); + } + } + } else { + const int out_n = n_tile + lane_lo; + if (out_n < size_n) { + #pragma unroll + for (int i = 0; i < 8; i++) { + const int out_m = m_tile + 2 * i + lane_hi; + if (out_m < size_m) { + T* dst = c + out_m * size_n + out_n; + if constexpr (std::is_same::value) { + *dst = __float2half_rn(c_acc[i]); + } else { + *dst = __float2bfloat16(c_acc[i]); + } + } + } + } + } +} + +#else // non-RDNA3 device pass: empty kernel for symbol parity. +template +__global__ void gemm_mxfp4_wmma_kernel_16x16_1w(const T*, const uint32_t*, + const uint8_t*, T*, const int, + const int, const int, + const int) {} +#endif + +template +void launch_gemm_mxfp4_wmma_16x16_1w(const T* a, const uint32_t* b_q, + const uint8_t* b_scales_e8m0, T* c, + int size_m, int size_n, int size_k, + int groups, cudaStream_t stream) { + int k_split = compute_mxfp4_k_split(size_k); + // Each K-segment must be a whole number of 16-wide tiles, else the k_tile + // loop would drop the tail of K. Fall back to no split otherwise. + while (k_split > 1 && (size_k % k_split != 0 || (size_k / k_split) % 16 != 0)) { + k_split /= 2; + } + dim3 block(32); + dim3 grid((size_n + 15) / 16, (size_m + 15) / 16, k_split); + gemm_mxfp4_wmma_kernel_16x16_1w<<>>( + a, b_q, b_scales_e8m0, c, size_m, size_n, size_k, groups); +} + +// =========================================================================== +// Multi-wave kernels for the prefill regime (large M). Same WMMA tiling / +// double-buffering / K-split as the GPTQ kernels; only the B-tile dequant +// differs (E2M1 + cached E8M0 exponent bias, no zeros). group_size == 32. +// =========================================================================== +#if defined(__HIP__RDNA3__) || !defined(__HIP_DEVICE_COMPILE__) + +// ---- 32x16_2w: 2 waves, 32M x 16N, double-buffered LDS (32 <= M < 64) ---- +template +__global__ void gemm_mxfp4_wmma_kernel_32x16_2w( + const T* __restrict__ a, const uint32_t* __restrict__ b_q, + const uint8_t* __restrict__ b_scales_e8m0, T* __restrict__ c, + const int size_m, const int size_n, const int size_k, const int groups) { + using E = typename WmmaNative::elem; + using V16 = typename WmmaNative::v16; + + const int m_tile = blockIdx.y * 32; + const int n_tile = blockIdx.x * 16; + if (m_tile >= size_m || n_tile >= size_n) return; + + const int tid = threadIdx.x; + const int wave_id = tid >> 5; + const int lane = tid & 31; + const int lane_lo = lane & 15; + const int lane_hi = lane >> 4; + + v8fp32 c_acc = {0, 0, 0, 0, 0, 0, 0, 0}; + const int groupsize = size_k / groups; + constexpr uint32_t mant_bits = std::is_same::value ? 10u : 7u; + + const int k_per_split = size_k / gridDim.z; + const int k_start = blockIdx.z * k_per_split; + const int k_end = k_start + k_per_split; + + __shared__ T b_lds[2][16][16]; + + auto dequant_into = [&](int buf, int k_tile) { + if (wave_id != 0) return; + const int my_n = lane_lo; + const int actual_n = n_tile + my_n; + if (actual_n >= size_n) return; + const int qk_row = (k_tile / 8) + lane_hi; + const uint32_t qa = b_q[qk_row * size_n + actual_n]; + const int g = k_tile / groupsize; + const uint32_t s8 = b_scales_e8m0[g * size_n + actual_n]; + const int32_t bias = mxfp4_e8m0_bias(s8, mant_bits); + const int k_base = lane_hi * 8; + if constexpr (std::is_same::value) { + half dq[8]; + dequant_mxfp4_8_fp16(qa, bias, dq); + #pragma unroll + for (int i = 0; i < 8; i++) b_lds[buf][k_base + i][my_n] = dq[i]; + } else { + bf16_t dq[8]; + dequant_mxfp4_8_bf16(qa, bias, dq); + #pragma unroll + for (int i = 0; i < 8; i++) b_lds[buf][k_base + i][my_n] = dq[i]; + } + }; + + dequant_into(0, k_start); + __syncthreads(); + + int cur_buf = 0; + for (int k_tile = k_start; k_tile < k_end; k_tile += 16) { + const int next_buf = 1 - cur_buf; + if (k_tile + 16 < k_end) dequant_into(next_buf, k_tile + 16); + + const int m_row = m_tile + wave_id * 16 + lane_lo; + V16 a_frag, b_frag; + if (m_row < size_m) { + static_assert(sizeof(a_frag) == 32, "V16 must be 32 bytes"); + __builtin_memcpy(&a_frag, a + m_row * size_k + k_tile, sizeof(a_frag)); + } else { + #pragma unroll + for (int i = 0; i < 16; i++) a_frag[i] = (E)0; + } + #pragma unroll + for (int i = 0; i < 16; i++) + b_frag[i] = bitcast_elem(b_lds[cur_buf][i][lane_lo]); + + c_acc = wmma_mma(a_frag, b_frag, c_acc); + __syncthreads(); + cur_buf = next_buf; + } + + const int m_tile_wave = m_tile + wave_id * 16; + if (gridDim.z > 1) { + const bool is_even_lane = (lane_lo & 1) == 0; + const int out_n_pair = n_tile + lane_lo; + #pragma unroll + for (int i = 0; i < 8; i++) { + float other_f = __shfl_xor(c_acc[i], 1); + if (!is_even_lane) continue; + const int out_m = m_tile_wave + 2 * i + lane_hi; + if (out_m >= size_m || out_n_pair >= size_n) continue; + T* dst = c + out_m * size_n + out_n_pair; + if constexpr (std::is_same::value) { + half2 packed = + __halves2half2(__float2half_rn(c_acc[i]), __float2half_rn(other_f)); + atomic_add_pk_f16(reinterpret_cast(dst), packed); + } else { + __hip_bfloat162 packed; + packed.x = __float2bfloat16(c_acc[i]); + packed.y = __float2bfloat16(other_f); + atomic_add_pk_bf16(reinterpret_cast<__hip_bfloat162*>(dst), packed); + } + } + } else { + const int out_n = n_tile + lane_lo; + if (out_n < size_n) { + #pragma unroll + for (int i = 0; i < 8; i++) { + const int out_m = m_tile_wave + 2 * i + lane_hi; + if (out_m < size_m) { + T* dst = c + out_m * size_n + out_n; + if constexpr (std::is_same::value) + *dst = __float2half_rn(c_acc[i]); + else + *dst = __float2bfloat16(c_acc[i]); + } + } + } + } +} + +// ---- 64x64_4w: 4 waves, 64M x 64N, K=16/iter (64 <= M < 128) ---- +template +__global__ void gemm_mxfp4_wmma_kernel_64x64_4w( + const T* __restrict__ a, const uint32_t* __restrict__ b_q, + const uint8_t* __restrict__ b_scales_e8m0, T* __restrict__ c, + const int size_m, const int size_n, const int size_k, const int groups) { + using E = typename WmmaNative::elem; + using V16 = typename WmmaNative::v16; + + const int m_tile = blockIdx.y * 64; + const int n_tile = blockIdx.x * 64; + if (m_tile >= size_m || n_tile >= size_n) return; + + const int tid = threadIdx.x; + const int wave_id = tid >> 5; + const int lane = tid & 31; + const int lane_lo = lane & 15; + const int lane_hi = lane >> 4; + + v8fp32 c_acc0 = {0, 0, 0, 0, 0, 0, 0, 0}; + v8fp32 c_acc1 = {0, 0, 0, 0, 0, 0, 0, 0}; + v8fp32 c_acc2 = {0, 0, 0, 0, 0, 0, 0, 0}; + v8fp32 c_acc3 = {0, 0, 0, 0, 0, 0, 0, 0}; + + const int groupsize = size_k / groups; + constexpr uint32_t mant_bits = std::is_same::value ? 10u : 7u; + + const int k_per_split = size_k / gridDim.z; + const int k_start = blockIdx.z * k_per_split; + const int k_end = k_start + k_per_split; + + __shared__ T b_lds[2][16][64]; + + auto dequant_into = [&](int buf, int k_tile) { + const int my_n_in_tile = wave_id * 16 + lane_lo; + const int actual_n = n_tile + my_n_in_tile; + if (actual_n >= size_n) return; + const int qk_row = (k_tile / 8) + lane_hi; + const uint32_t qa = b_q[qk_row * size_n + actual_n]; + const int g = k_tile / groupsize; + const uint32_t s8 = b_scales_e8m0[g * size_n + actual_n]; + const int32_t bias = mxfp4_e8m0_bias(s8, mant_bits); + const int k_base = lane_hi * 8; + if constexpr (std::is_same::value) { + half dq[8]; + dequant_mxfp4_8_fp16(qa, bias, dq); + #pragma unroll + for (int i = 0; i < 8; i++) b_lds[buf][k_base + i][my_n_in_tile] = dq[i]; + } else { + bf16_t dq[8]; + dequant_mxfp4_8_bf16(qa, bias, dq); + #pragma unroll + for (int i = 0; i < 8; i++) b_lds[buf][k_base + i][my_n_in_tile] = dq[i]; + } + }; + + dequant_into(0, k_start); + __syncthreads(); + + int cur_buf = 0; + for (int k_tile = k_start; k_tile < k_end; k_tile += 16) { + const int next_buf = 1 - cur_buf; + if (k_tile + 16 < k_end) dequant_into(next_buf, k_tile + 16); + + const int m_row = m_tile + wave_id * 16 + lane_lo; + V16 a_frag, b_frag0, b_frag1, b_frag2, b_frag3; + if (m_row < size_m) { + static_assert(sizeof(a_frag) == 32, "V16 must be 32 bytes"); + __builtin_memcpy(&a_frag, a + m_row * size_k + k_tile, sizeof(a_frag)); + } else { + #pragma unroll + for (int i = 0; i < 16; i++) a_frag[i] = (E)0; + } + #pragma unroll + for (int i = 0; i < 16; i++) { + b_frag0[i] = bitcast_elem(b_lds[cur_buf][i][lane_lo + 0]); + b_frag1[i] = bitcast_elem(b_lds[cur_buf][i][lane_lo + 16]); + b_frag2[i] = bitcast_elem(b_lds[cur_buf][i][lane_lo + 32]); + b_frag3[i] = bitcast_elem(b_lds[cur_buf][i][lane_lo + 48]); + } + c_acc0 = wmma_mma(a_frag, b_frag0, c_acc0); + c_acc1 = wmma_mma(a_frag, b_frag1, c_acc1); + c_acc2 = wmma_mma(a_frag, b_frag2, c_acc2); + c_acc3 = wmma_mma(a_frag, b_frag3, c_acc3); + __syncthreads(); + cur_buf = next_buf; + } + + const int m_tile_wave = m_tile + wave_id * 16; + auto store_acc = [&](const v8fp32& acc, int n_base) { + if (gridDim.z > 1) { + const bool is_even_lane = (lane_lo & 1) == 0; + const int out_n_pair = n_base + lane_lo; + #pragma unroll + for (int i = 0; i < 8; i++) { + float other_f = __shfl_xor(acc[i], 1); + if (!is_even_lane) continue; + const int out_m = m_tile_wave + 2 * i + lane_hi; + if (out_m >= size_m || out_n_pair >= size_n) continue; + T* dst = c + out_m * size_n + out_n_pair; + if constexpr (std::is_same::value) { + half2 packed = + __halves2half2(__float2half_rn(acc[i]), __float2half_rn(other_f)); + atomic_add_pk_f16(reinterpret_cast(dst), packed); + } else { + __hip_bfloat162 packed; + packed.x = __float2bfloat16(acc[i]); + packed.y = __float2bfloat16(other_f); + atomic_add_pk_bf16(reinterpret_cast<__hip_bfloat162*>(dst), packed); + } + } + } else { + const int out_n = n_base + lane_lo; + if (out_n >= size_n) return; + #pragma unroll + for (int i = 0; i < 8; i++) { + const int out_m = m_tile_wave + 2 * i + lane_hi; + if (out_m < size_m) { + T* dst = c + out_m * size_n + out_n; + if constexpr (std::is_same::value) + *dst = __float2half_rn(acc[i]); + else + *dst = __float2bfloat16(acc[i]); + } + } + } + }; + store_acc(c_acc0, n_tile + 0); + store_acc(c_acc1, n_tile + 16); + store_acc(c_acc2, n_tile + 32); + store_acc(c_acc3, n_tile + 48); +} + +// ---- 128x64_k32: 8 waves, 128M x 64N, K=32/iter, cached group bias (M >= 128) ---- +template +__global__ void gemm_mxfp4_wmma_kernel_128x64_k32( + const T* __restrict__ a, const uint32_t* __restrict__ b_q, + const uint8_t* __restrict__ b_scales_e8m0, T* __restrict__ c, + const int size_m, const int size_n, const int size_k, const int groups) { + using E = typename WmmaNative::elem; + using V16 = typename WmmaNative::v16; + + const int m_tile = blockIdx.y * 128; + const int n_tile = blockIdx.x * 64; + if (m_tile >= size_m || n_tile >= size_n) return; + + const int tid = threadIdx.x; + const int wave_id = tid >> 5; + const int lane = tid & 31; + const int lane_lo = lane & 15; + const int lane_hi = lane >> 4; + + v8fp32 c_acc0 = {0, 0, 0, 0, 0, 0, 0, 0}; + v8fp32 c_acc1 = {0, 0, 0, 0, 0, 0, 0, 0}; + v8fp32 c_acc2 = {0, 0, 0, 0, 0, 0, 0, 0}; + v8fp32 c_acc3 = {0, 0, 0, 0, 0, 0, 0, 0}; + + const int groupsize = size_k / groups; + constexpr uint32_t mant_bits = std::is_same::value ? 10u : 7u; + + const int k_per_split = size_k / gridDim.z; + const int k_start = blockIdx.z * k_per_split; + const int k_end = k_start + k_per_split; + + __shared__ T b_lds[2][64][34]; // +2 pad to break LDS bank conflicts + + const int dq_wave4 = wave_id & 3; + const int dq_k_half = (wave_id >= 4) ? 1 : 0; + const int dq_n = dq_wave4 * 16 + lane_lo; + const int dq_an = n_tile + dq_n; + const bool dq_ok = (dq_an < size_n); + const int dq_oct = lane_hi + dq_k_half * 2; + const int dq_kb = dq_oct * 8; + int cached_g = -1; + int32_t cached_bias = 0; + + auto dequant_into = [&](int buf, int k_tile) __attribute__((always_inline)) { + if (!dq_ok) return; + const int g = k_tile / groupsize; + if (g != cached_g) { + cached_g = g; + const uint32_t s8 = b_scales_e8m0[g * size_n + dq_an]; + cached_bias = mxfp4_e8m0_bias(s8, mant_bits); + } + const int qk_row = (k_tile / 8) + dq_oct; + const uint32_t qa = b_q[qk_row * size_n + dq_an]; + if constexpr (std::is_same::value) { + half dq[8]; + dequant_mxfp4_8_fp16(qa, cached_bias, dq); + #pragma unroll + for (int i = 0; i < 8; i++) b_lds[buf][dq_n][dq_kb + i] = dq[i]; + } else { + bf16_t dq[8]; + dequant_mxfp4_8_bf16(qa, cached_bias, dq); + #pragma unroll + for (int i = 0; i < 8; i++) b_lds[buf][dq_n][dq_kb + i] = dq[i]; + } + }; + + dequant_into(0, k_start); + __syncthreads(); + + int cur_buf = 0; + const int m_row = m_tile + wave_id * 16 + lane_lo; + const T* a_row_ptr = (m_row < size_m) ? (a + m_row * size_k) : nullptr; + + for (int k_tile = k_start; k_tile < k_end; k_tile += 32) { + const int next_buf = 1 - cur_buf; + if (k_tile + 32 < k_end) dequant_into(next_buf, k_tile + 32); + + V16 a_frag_lo, a_frag_hi; + V16 b_frag0, b_frag1, b_frag2, b_frag3; + if (a_row_ptr) { + __builtin_memcpy(&a_frag_lo, a_row_ptr + k_tile, sizeof(V16)); + __builtin_memcpy(&a_frag_hi, a_row_ptr + k_tile + 16, sizeof(V16)); + } else { + #pragma unroll + for (int i = 0; i < 16; i++) a_frag_lo[i] = (E)0; + #pragma unroll + for (int i = 0; i < 16; i++) a_frag_hi[i] = (E)0; + } + + static_assert(sizeof(V16) == 32, "V16 must be 32 bytes for memcpy"); + __builtin_memcpy(&b_frag0, &b_lds[cur_buf][lane_lo + 0][0], 32); + __builtin_memcpy(&b_frag1, &b_lds[cur_buf][lane_lo + 16][0], 32); + __builtin_memcpy(&b_frag2, &b_lds[cur_buf][lane_lo + 32][0], 32); + __builtin_memcpy(&b_frag3, &b_lds[cur_buf][lane_lo + 48][0], 32); + c_acc0 = wmma_mma(a_frag_lo, b_frag0, c_acc0); + c_acc1 = wmma_mma(a_frag_lo, b_frag1, c_acc1); + c_acc2 = wmma_mma(a_frag_lo, b_frag2, c_acc2); + c_acc3 = wmma_mma(a_frag_lo, b_frag3, c_acc3); + + __builtin_memcpy(&b_frag0, &b_lds[cur_buf][lane_lo + 0][16], 32); + __builtin_memcpy(&b_frag1, &b_lds[cur_buf][lane_lo + 16][16], 32); + __builtin_memcpy(&b_frag2, &b_lds[cur_buf][lane_lo + 32][16], 32); + __builtin_memcpy(&b_frag3, &b_lds[cur_buf][lane_lo + 48][16], 32); + c_acc0 = wmma_mma(a_frag_hi, b_frag0, c_acc0); + c_acc1 = wmma_mma(a_frag_hi, b_frag1, c_acc1); + c_acc2 = wmma_mma(a_frag_hi, b_frag2, c_acc2); + c_acc3 = wmma_mma(a_frag_hi, b_frag3, c_acc3); + + __syncthreads(); + cur_buf = next_buf; + } + + const int m_tile_wave = m_tile + wave_id * 16; + auto store_acc = [&](const v8fp32& acc, int n_base) { + if (gridDim.z > 1) { + const bool is_even_lane = (lane_lo & 1) == 0; + const int out_n_pair = n_base + lane_lo; + #pragma unroll + for (int i = 0; i < 8; i++) { + float other_f = __shfl_xor(acc[i], 1); + if (!is_even_lane) continue; + const int out_m = m_tile_wave + 2 * i + lane_hi; + if (out_m >= size_m || out_n_pair >= size_n) continue; + T* dst = c + out_m * size_n + out_n_pair; + if constexpr (std::is_same::value) { + half2 packed = + __halves2half2(__float2half_rn(acc[i]), __float2half_rn(other_f)); + atomic_add_pk_f16(reinterpret_cast(dst), packed); + } else { + __hip_bfloat162 packed; + packed.x = __float2bfloat16(acc[i]); + packed.y = __float2bfloat16(other_f); + atomic_add_pk_bf16(reinterpret_cast<__hip_bfloat162*>(dst), packed); + } + } + } else { + const int out_n = n_base + lane_lo; + if (out_n >= size_n) return; + #pragma unroll + for (int i = 0; i < 8; i++) { + const int out_m = m_tile_wave + 2 * i + lane_hi; + if (out_m < size_m) { + T* dst = c + out_m * size_n + out_n; + if constexpr (std::is_same::value) + *dst = __float2half_rn(acc[i]); + else + *dst = __float2bfloat16(acc[i]); + } + } + } + }; + store_acc(c_acc0, n_tile + 0); + store_acc(c_acc1, n_tile + 16); + store_acc(c_acc2, n_tile + 32); + store_acc(c_acc3, n_tile + 48); +} + +#else // non-RDNA3 device pass: empty kernels for symbol parity. +template +__global__ void gemm_mxfp4_wmma_kernel_32x16_2w(const T*, const uint32_t*, + const uint8_t*, T*, const int, + const int, const int, + const int) {} +template +__global__ void gemm_mxfp4_wmma_kernel_64x64_4w(const T*, const uint32_t*, + const uint8_t*, T*, const int, + const int, const int, + const int) {} +template +__global__ void gemm_mxfp4_wmma_kernel_128x64_k32(const T*, const uint32_t*, + const uint8_t*, T*, const int, + const int, const int, + const int) {} +#endif + +// K-split that keeps each segment a whole multiple of `step` 16-wide tiles. +__host__ static inline int mxfp4_k_split_aligned(int size_k, int step) { + int s = compute_mxfp4_k_split(size_k); + while (s > 1 && (size_k % s != 0 || (size_k / s) % step != 0)) s /= 2; + return s; +} + +// Dispatcher: pick the widest tile that fits M, falling back to v1 for M < 32. +template +void launch_gemm_mxfp4_wmma(const T* a, const uint32_t* b_q, + const uint8_t* b_scales_e8m0, T* c, int size_m, + int size_n, int size_k, int groups, + cudaStream_t stream) { + if (size_m >= 128 && size_n >= 64) { + const int k_split = mxfp4_k_split_aligned(size_k, 32); + dim3 block(256); + dim3 grid((size_n + 63) / 64, (size_m + 127) / 128, k_split); + gemm_mxfp4_wmma_kernel_128x64_k32<<>>( + a, b_q, b_scales_e8m0, c, size_m, size_n, size_k, groups); + return; + } + if (size_m >= 64 && size_n >= 64) { + const int k_split = mxfp4_k_split_aligned(size_k, 16); + dim3 block(128); + dim3 grid((size_n + 63) / 64, (size_m + 63) / 64, k_split); + gemm_mxfp4_wmma_kernel_64x64_4w<<>>( + a, b_q, b_scales_e8m0, c, size_m, size_n, size_k, groups); + return; + } + if (size_m >= 32) { + const int k_split = mxfp4_k_split_aligned(size_k, 16); + dim3 block(64); + dim3 grid((size_n + 15) / 16, (size_m + 31) / 32, k_split); + gemm_mxfp4_wmma_kernel_32x16_2w<<>>( + a, b_q, b_scales_e8m0, c, size_m, size_n, size_k, groups); + return; + } + launch_gemm_mxfp4_wmma_16x16_1w(a, b_q, b_scales_e8m0, c, size_m, size_n, + size_k, groups, stream); +} + +// Scalar GEMV for decode (small M): WMMA wastes 15/16 of its compute at M=1, +// so this reads each weight column once and does M scalar dots (bandwidth-bound, +// the 4-bit read wins). 256 threads, 4 N cols each; gridDim.z splits K (atomic). +#define MXFP4_BLOCK_KN 256 +#define MXFP4_THREADS 256 + +#if defined(__HIP__RDNA3__) || !defined(__HIP_DEVICE_COMPILE__) + +__device__ __forceinline__ float mxfp4_to_f(half v) { return __half2float(v); } +__device__ __forceinline__ float mxfp4_to_f(bf16_t v) { + return __bfloat162float(v); +} + +__device__ __forceinline__ void mxfp4_atomic_add_pk4(half* addr, half2 v01, + half2 v23) { + unsigned long long* a = reinterpret_cast(addr); + unsigned long long old = *a; + 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(a, old, sum.u); + if (prev == old) break; + old = prev; + } +} +__device__ __forceinline__ void mxfp4_atomic_add_pk4(bf16_t* addr, + __hip_bfloat162 v01, + __hip_bfloat162 v23) { + unsigned long long* a = reinterpret_cast(addr); + unsigned long long old = *a; + while (true) { + union { + unsigned long long u; + __hip_bfloat162 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(a, old, sum.u); + if (prev == old) break; + old = prev; + } +} + +template +__global__ void gemm_mxfp4_scalar_rdna3( + const T* __restrict__ a, const uint32_t* __restrict__ b_q, + const uint8_t* __restrict__ b_scales_e8m0, T* __restrict__ c, + const int size_m, const int size_n, const int size_k, const int groups) { + const int t = threadIdx.x; + const int offset_n = blockIdx.x * MXFP4_BLOCK_KN * 4; + const int offset_m = blockIdx.y * M_COUNT; + const int offset_k = blockIdx.z * MXFP4_BLOCK_KN; + const int end_k = min(offset_k + MXFP4_BLOCK_KN, size_k); + const int n = offset_n + t * 4; + + constexpr int LDS_PAD = 8; + __shared__ T block_a[M_COUNT][MXFP4_BLOCK_KN + LDS_PAD]; + if (offset_k + t < end_k) { + #pragma unroll + for (int m = 0; m < M_COUNT; ++m) { + block_a[m][t] = (offset_m + m < size_m) + ? a[(offset_m + m) * size_k + offset_k + t] + : (T)0; + } + } + __syncthreads(); + if (n >= size_n) return; + + const int groupsize = size_k / groups; + int group = offset_k / groupsize; + int nextgroup = (group + 1) * groupsize; + // dequant_mxfp4_8_f32 always decodes via bf16 bits (mant 7), independent of + // the activation dtype T, so the E8M0 bias must use mant_bits=7. + constexpr uint32_t mant_bits = 7u; + + int qk = offset_k / 8; + const uint32_t* b_ptr = b_q + qk * size_n + n; + + int32_t bias[4]; + auto refresh = [&](int g) { + #pragma unroll + for (int col = 0; col < 4; ++col) + bias[col] = mxfp4_e8m0_bias(b_scales_e8m0[g * size_n + n + col], mant_bits); + }; + refresh(group); + + float block_c[M_COUNT][4]; + #pragma unroll + for (int m = 0; m < M_COUNT; ++m) + #pragma unroll + for (int col = 0; col < 4; ++col) block_c[m][col] = 0.0f; + + int k = offset_k; + while (k < end_k) { + if (k == nextgroup) { + group++; + nextgroup += groupsize; + refresh(group); + } + 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; + uint32_t wcol[4]; + __builtin_memcpy(wcol, &b_w[j], sizeof(int4)); + #pragma unroll + for (int col = 0; col < 4; ++col) { + float dq[8]; + dequant_mxfp4_8_f32(wcol[col], bias[col], dq); + #pragma unroll + for (int m = 0; m < M_COUNT; ++m) { + const T* ap = &block_a[m][a_off]; + float s = 0.0f; + #pragma unroll + for (int i = 0; i < 8; ++i) s += dq[i] * mxfp4_to_f(ap[i]); + block_c[m][col] += s; + } + } + } + k += 32; + } + + #pragma unroll + for (int m = 0; m < M_COUNT; ++m) { + if (offset_m + m >= size_m) continue; + T* out = c + (offset_m + m) * size_n + n; + if constexpr (std::is_same::value) { + half2 r01 = __halves2half2(__float2half_rn(block_c[m][0]), + __float2half_rn(block_c[m][1])); + half2 r23 = __halves2half2(__float2half_rn(block_c[m][2]), + __float2half_rn(block_c[m][3])); + mxfp4_atomic_add_pk4(out, r01, r23); + } else { + __hip_bfloat162 r01, r23; + r01.x = __float2bfloat16(block_c[m][0]); + r01.y = __float2bfloat16(block_c[m][1]); + r23.x = __float2bfloat16(block_c[m][2]); + r23.y = __float2bfloat16(block_c[m][3]); + mxfp4_atomic_add_pk4(out, r01, r23); + } + } +} + +#else // non-RDNA3 stub +template +__global__ void gemm_mxfp4_scalar_rdna3(const T*, const uint32_t*, + const uint8_t*, T*, const int, + const int, const int, const int) {} +#endif + +template +void launch_mxfp4_scalar(const T* a, const uint32_t* b_q, + const uint8_t* b_scales_e8m0, T* c, int size_m, + int size_n, int size_k, int groups, + cudaStream_t stream) { + dim3 block(MXFP4_BLOCK_KN); + dim3 grid((size_n + MXFP4_BLOCK_KN * 4 - 1) / (MXFP4_BLOCK_KN * 4), + (size_m + M_COUNT - 1) / M_COUNT, + (size_k + MXFP4_BLOCK_KN - 1) / MXFP4_BLOCK_KN); + gemm_mxfp4_scalar_rdna3<<>>( + a, b_q, b_scales_e8m0, c, size_m, size_n, size_k, groups); +} + +template +void launch_gemm_mxfp4_scalar(const T* a, const uint32_t* b_q, + const uint8_t* b_scales_e8m0, T* c, int size_m, + int size_n, int size_k, int groups, + cudaStream_t stream) { + if (size_m == 1) + launch_mxfp4_scalar(a, b_q, b_scales_e8m0, c, size_m, size_n, size_k, + groups, stream); + else if (size_m <= 3) + launch_mxfp4_scalar(a, b_q, b_scales_e8m0, c, size_m, size_n, size_k, + groups, stream); + else if (size_m <= 7) + launch_mxfp4_scalar(a, b_q, b_scales_e8m0, c, size_m, size_n, size_k, + groups, stream); + else + launch_mxfp4_scalar(a, b_q, b_scales_e8m0, c, size_m, size_n, size_k, + groups, stream); +} + +} // namespace mxfp4_rdna3_wmma +} // namespace vllm + +// --------------------------------------------------------------------------- +// Public entry point. +// --------------------------------------------------------------------------- +// +// Inputs: +// a [M, K] half or bfloat16 +// b_q_weight [K/8, N] uint32 (E2M1, 8 sequential K nibbles per word, +// repacked from compressed-tensors [N, K/2]) +// b_scales_e8m0 [K/32, N] uint8 (E8M0 block scale, group_size = 32) +// +// Output: +// c [M, N] same dtype as a +// +// Requirements: N % 16 == 0, K % 32 == 0. +torch::Tensor mxfp4_gemm_rdna3(torch::Tensor a, torch::Tensor b_q_weight, + torch::Tensor b_scales_e8m0) { + TORCH_CHECK(a.is_cuda(), "a must be a CUDA/HIP tensor"); + TORCH_CHECK(b_q_weight.is_cuda(), "b_q_weight must be a CUDA/HIP tensor"); + TORCH_CHECK(b_scales_e8m0.is_cuda(), "b_scales_e8m0 must be a CUDA/HIP tensor"); + TORCH_CHECK(a.dim() == 2, "a must be 2D [M, K]"); + TORCH_CHECK(b_q_weight.dim() == 2, "b_q_weight must be 2D [K/8, N]"); + TORCH_CHECK( + a.scalar_type() == torch::kHalf || a.scalar_type() == torch::kBFloat16, + "a must be half or bfloat16"); + TORCH_CHECK(b_q_weight.scalar_type() == torch::kUInt32 || + b_q_weight.scalar_type() == torch::kInt32, + "b_q_weight must be (u)int32"); + TORCH_CHECK(b_scales_e8m0.scalar_type() == torch::kUInt8, + "b_scales_e8m0 must be uint8 (E8M0)"); + + 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(1); + int groups = (int)b_scales_e8m0.size(0); + + TORCH_CHECK(b_q_weight.size(0) * 8 == size_k, + "b_q_weight first dim must be K/8"); + TORCH_CHECK(b_scales_e8m0.size(1) == size_n, "b_scales_e8m0 last dim must be N"); + TORCH_CHECK(groups * 32 == size_k, "E8M0 group count must be K/32"); + TORCH_CHECK(size_n % 16 == 0, "WMMA path requires N % 16 == 0"); + TORCH_CHECK(size_k % 32 == 0, "MXFP4 path requires K % 32 == 0"); + + auto opts = torch::TensorOptions().dtype(a.dtype()).device(a.device()); + // Zero-init: K-split (gridDim.z > 1) accumulates atomically into c. + at::Tensor c = torch::zeros({size_m, size_n}, opts); + + // Small M (decode) -> scalar GEMV; larger M (prefill/batch) -> WMMA. + const bool use_scalar = size_m <= 8; + if (a.scalar_type() == torch::kHalf) { + namespace mx = vllm::mxfp4_rdna3_wmma; + auto* ap = (const half*)a.data_ptr(); + auto* bq = (const uint32_t*)b_q_weight.data_ptr(); + auto* bs = (const uint8_t*)b_scales_e8m0.data_ptr(); + auto* cp = (half*)c.data_ptr(); + if (use_scalar) + mx::launch_gemm_mxfp4_scalar(ap, bq, bs, cp, size_m, size_n, size_k, + groups, stream); + else + mx::launch_gemm_mxfp4_wmma(ap, bq, bs, cp, size_m, size_n, size_k, + groups, stream); + } else { + using bf16_t = vllm::mxfp4_rdna3_wmma::bf16_t; + namespace mx = vllm::mxfp4_rdna3_wmma; + auto* ap = (const bf16_t*)a.data_ptr(); + auto* bq = (const uint32_t*)b_q_weight.data_ptr(); + auto* bs = (const uint8_t*)b_scales_e8m0.data_ptr(); + auto* cp = (bf16_t*)c.data_ptr(); + if (use_scalar) + mx::launch_gemm_mxfp4_scalar(ap, bq, bs, cp, size_m, size_n, + size_k, groups, stream); + else + mx::launch_gemm_mxfp4_wmma(ap, bq, bs, cp, size_m, size_n, size_k, + groups, stream); + } + + return c; +} diff --git a/csrc/rocm/ops.h b/csrc/rocm/ops.h index 549d50300d6f..b8fdd9b8a136 100644 --- a/csrc/rocm/ops.h +++ b/csrc/rocm/ops.h @@ -27,6 +27,9 @@ 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); +torch::Tensor mxfp4_gemm_rdna3(torch::Tensor a, torch::Tensor b_q_weight, + torch::Tensor b_scales_e8m0); + 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, diff --git a/csrc/rocm/qdq_mxfp4_rdna3.cuh b/csrc/rocm/qdq_mxfp4_rdna3.cuh new file mode 100644 index 000000000000..b41712f5a9bd --- /dev/null +++ b/csrc/rocm/qdq_mxfp4_rdna3.cuh @@ -0,0 +1,121 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright contributors to the vLLM project +// +// MXFP4 dequant for RDNA3: E2M1 weight (mags {0,.5,1,1.5,2,3,4,6}) + E8M0 +// block scale (2^(s8-127), group of 32), no zero-point. E2M1->{bf16,fp16} +// is a field copy (only 0 and the 0.5 subnormal are special); the E8M0 +// scale folds in as an exponent add `bits += (s8-127)< + +// Device-only HIP dependencies are gated so the pure-integer decode below can +// be #included and unit-tested by a plain host compiler. +#if defined(__HIPCC__) +#include +#include +#define MXFP4_HD __host__ __device__ __forceinline__ +#define MXFP4_D __device__ __forceinline__ +#else +#define MXFP4_HD inline +#endif + +namespace vllm { +namespace mxfp4_rdna3 { + +#if defined(__HIPCC__) +using bf16_t = __hip_bfloat16; +using bf162_t = __hip_bfloat162; +#endif + +// Pure-integer E2M1 -> float bits of the unscaled magnitude (no HIP types). +// bf16: bias 127, mant top bit 6. fp16: bias 15, mant top bit 9. +MXFP4_HD uint16_t mxfp4_e2m1_bits(uint32_t nib, uint32_t exp_bias, + uint32_t mant_pos) { + const uint32_t sign = (nib & 0x8u) << 12; + const uint32_t em = nib & 0x7u; // exp(2) | mant(1) + const uint32_t e = em >> 1; + const uint32_t m = em & 1u; + // Normal: E = e-1+exp_bias, M = m at top. The 0.5 subnormal (e==0,m==1) + // reuses that exponent with M forced to 0; the zero code clears to 0x0000. + const uint32_t out_exp = (e + exp_bias - 1u) << (mant_pos + 1u); + const uint32_t out_mant = (e != 0u ? m : 0u) << mant_pos; + return (uint16_t)(em == 0u ? sign : (sign | out_exp | out_mant)); +} + +MXFP4_HD uint16_t mxfp4_e2m1_to_bf16_bits(uint32_t nib) { + return mxfp4_e2m1_bits(nib, /*exp_bias=*/127u, /*mant_pos=*/6u); +} + +MXFP4_HD uint16_t mxfp4_e2m1_to_fp16_bits(uint32_t nib) { + return mxfp4_e2m1_bits(nib, /*exp_bias=*/15u, /*mant_pos=*/9u); +} + +// Apply the E8M0 scale as an exponent add (zero stays zero). +MXFP4_HD uint16_t mxfp4_apply_e8m0_bits(uint16_t bits, int32_t bias_u16) { + return (bits & 0x7FFFu) == 0u ? bits + : (uint16_t)((int32_t)bits + bias_u16); +} + +// E8M0 byte -> additive exponent bias. mant_bits = 7 (bf16) or 10 (fp16). +MXFP4_HD int32_t mxfp4_e8m0_bias(uint32_t s8, uint32_t mant_bits) { + return ((int32_t)s8 - 127) << mant_bits; +} + +#if defined(__HIPCC__) +MXFP4_D bf16_t mxfp4_nib_to_bf16(uint32_t nib, int32_t bias_u16) { + union { + uint16_t u; + bf16_t b; + } o; + o.u = mxfp4_apply_e8m0_bits(mxfp4_e2m1_to_bf16_bits(nib), bias_u16); + return o.b; +} + +MXFP4_D half mxfp4_nib_to_fp16(uint32_t nib, int32_t bias_u16) { + union { + uint16_t u; + half h; + } o; + o.u = mxfp4_apply_e8m0_bits(mxfp4_e2m1_to_fp16_bits(nib), bias_u16); + return o.h; +} + +// Unpack 8 codes from a packed uint32 (nibble i = K position k0+i, low nibble +// = even K). All 8 share one E8M0 block since K is tiled in multiples of 32. +MXFP4_D void dequant_mxfp4_8_bf16(uint32_t qa, int32_t bias_u16, + bf16_t (&dq)[8]) { +#pragma unroll + for (int i = 0; i < 8; i++) { + dq[i] = mxfp4_nib_to_bf16((qa >> (4 * i)) & 0xFu, bias_u16); + } +} + +MXFP4_D void dequant_mxfp4_8_fp16(uint32_t qa, int32_t bias_u16, + half (&dq)[8]) { +#pragma unroll + for (int i = 0; i < 8; i++) { + dq[i] = mxfp4_nib_to_fp16((qa >> (4 * i)) & 0xFu, bias_u16); + } +} + +// fp32 output (bf16 bits widened by a free <<16) for a scalar dot path. +MXFP4_D void dequant_mxfp4_8_f32(uint32_t qa, int32_t bias_u16, + float (&dq)[8]) { +#pragma unroll + for (int i = 0; i < 8; i++) { + uint16_t b = mxfp4_apply_e8m0_bits( + mxfp4_e2m1_to_bf16_bits((qa >> (4 * i)) & 0xFu), bias_u16); + dq[i] = __uint_as_float((uint32_t)b << 16); + } +} + +#endif // defined(__HIPCC__) + +} // namespace mxfp4_rdna3 +} // namespace vllm + +#endif // _qdq_mxfp4_rdna3_cuh diff --git a/csrc/rocm/torch_bindings.cpp b/csrc/rocm/torch_bindings.cpp index 03de6dcd1576..76983023b27e 100644 --- a/csrc/rocm/torch_bindings.cpp +++ b/csrc/rocm/torch_bindings.cpp @@ -51,6 +51,11 @@ TORCH_LIBRARY_EXPAND(TORCH_EXTENSION_NAME, rocm_ops) { "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( + "mxfp4_gemm_rdna3(Tensor a, Tensor b_q_weight, " + "Tensor b_scales_e8m0) -> Tensor"); + rocm_ops.impl("mxfp4_gemm_rdna3", torch::kCUDA, &mxfp4_gemm_rdna3); + rocm_ops.def( "moe_gptq_gemm_rdna3(Tensor a, Tensor! c, Tensor b_q_weight, " "Tensor b_scales, Tensor b_qzeros, Tensor topk_weights, " diff --git a/tests/kernels/quantization/test_mxfp4_rdna3.py b/tests/kernels/quantization/test_mxfp4_rdna3.py new file mode 100644 index 000000000000..764e0b3f8e0f --- /dev/null +++ b/tests/kernels/quantization/test_mxfp4_rdna3.py @@ -0,0 +1,68 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""Correctness test for the RDNA3 (gfx1100) MXFP4 weight-only WMMA GEMM. + +Run inside the ROCm container after building _rocm_C: + .venv/bin/python -m pytest tests/kernels/quantization/test_mxfp4_rdna3_wmma.py -v +""" + +import pytest +import torch + +from vllm.platforms import current_platform + +E2M1_MAG = torch.tensor( + [0.0, 0.5, 1.0, 1.5, 2.0, 3.0, 4.0, 6.0], dtype=torch.float32 +) + +pytestmark = pytest.mark.skipif( + not current_platform.is_rocm(), reason="RDNA3 MXFP4 WMMA kernel is ROCm-only" +) + + +def _dequant_ref(w_packed: torch.Tensor, scale_e8m0: torch.Tensor) -> torch.Tensor: + """[N, K/2] uint8 (E2M1) + [N, K/32] uint8 (E8M0) -> [N, K] fp32.""" + N, k_half = w_packed.shape + K = k_half * 2 + lo = (w_packed & 0xF).to(torch.long) # even K + hi = (w_packed >> 4).to(torch.long) # odd K + codes = torch.stack([lo, hi], dim=-1).reshape(N, K) # interleave even/odd + sign = torch.where((codes & 0x8) != 0, -1.0, 1.0) + mag = E2M1_MAG.to(w_packed.device)[codes & 0x7] + scale = torch.pow(2.0, scale_e8m0.float() - 127.0) # [N, K/32] + scale = scale.repeat_interleave(32, dim=1) # [N, K] + return sign * mag * scale + + +def _repack(w_packed, scale_e8m0): + b_q = w_packed.contiguous().view(torch.int32).t().contiguous() # [K/8, N] + b_scale = scale_e8m0.t().contiguous() # [K/32, N] + return b_q, b_scale + + +@pytest.mark.parametrize("dtype", [torch.bfloat16, torch.float16]) +@pytest.mark.parametrize("M", [1, 16, 64, 128]) +@pytest.mark.parametrize("N,K", [(64, 256), (256, 512), (128, 4096)]) +def test_mxfp4_wmma_matches_dequant_reference(dtype, M, N, K): + if not hasattr(torch.ops._rocm_C, "mxfp4_gemm_rdna3"): + pytest.skip("_rocm_C.mxfp4_gemm_rdna3 not built") + + torch.manual_seed(0) + dev = "cuda" + # Random E2M1 codes (full 0..255 byte = two 4-bit codes) and mild E8M0 + # scales centered on 127 (== 2^0) to stay well inside fp16/bf16 range. + w_packed = torch.randint(0, 256, (N, K // 2), dtype=torch.uint8, device=dev) + scale_e8m0 = torch.randint( + 124, 131, (N, K // 32), dtype=torch.uint8, device=dev + ) + x = torch.randn(M, K, dtype=dtype, device=dev) * 0.1 + + w_deq = _dequant_ref(w_packed, scale_e8m0) # [N, K] fp32 + ref = (x.float() @ w_deq.t()).to(dtype) + + b_q, b_scale = _repack(w_packed, scale_e8m0) + out = torch.ops._rocm_C.mxfp4_gemm_rdna3(x, b_q, b_scale) + + assert out.shape == (M, N) + # bf16/fp16 accumulate-in-fp32; tolerate dtype rounding of the reduction. + torch.testing.assert_close(out.float(), ref.float(), atol=2e-2, rtol=2e-2) diff --git a/vllm/model_executor/kernels/linear/__init__.py b/vllm/model_executor/kernels/linear/__init__.py index 4ac8d49cd58e..319a0e86a3c5 100644 --- a/vllm/model_executor/kernels/linear/__init__.py +++ b/vllm/model_executor/kernels/linear/__init__.py @@ -77,6 +77,9 @@ from vllm.model_executor.kernels.linear.mxfp4.marlin import ( MarlinMxFp4LinearKernel, ) +from vllm.model_executor.kernels.linear.mxfp4.rocm import ( + Rdna3MxFp4LinearKernel, +) from vllm.model_executor.kernels.linear.mxfp4.xpu import ( XPUMxFp4LinearKernel, ) @@ -431,6 +434,9 @@ def _filter_kernels_by_backend( PlatformEnum.XPU: [ XPUMxFp4LinearKernel, ], + PlatformEnum.ROCM: [ + Rdna3MxFp4LinearKernel, + ], } # TODO make all kernels inherit from MMLinearKernel diff --git a/vllm/model_executor/kernels/linear/mxfp4/rocm.py b/vllm/model_executor/kernels/linear/mxfp4/rocm.py new file mode 100644 index 000000000000..6fc55dbabc7c --- /dev/null +++ b/vllm/model_executor/kernels/linear/mxfp4/rocm.py @@ -0,0 +1,89 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +import torch +from torch.nn.parameter import Parameter + +from vllm.platforms import current_platform + +from .base import MxFp4LinearKernel, MxFp4LinearLayerConfig + +_MXFP4_GROUP_SIZE = 32 + +# Fake/meta impl so the op is traceable under torch.compile + CUDA graphs. +if hasattr(torch.ops, "_rocm_C") and hasattr( + torch.ops._rocm_C, "mxfp4_gemm_rdna3" +): + try: + + @torch.library.register_fake("_rocm_C::mxfp4_gemm_rdna3") + def _mxfp4_gemm_rdna3_fake(a, b_q_weight, b_scales_e8m0): + return a.new_empty((a.shape[0], b_q_weight.shape[1])) + except RuntimeError: + pass # already registered + + +class Rdna3MxFp4LinearKernel(MxFp4LinearKernel): + """Weight-only MXFP4 GEMM for RDNA3 (gfx11xx) via the native HIP kernel + ``_rocm_C.mxfp4_gemm_rdna3`` (scalar GEMV for decode, WMMA for prefill). + RDNA3 has no native FP4 matmul, so E2M1 weights are dequantized on the fly + into the f16/bf16 pipeline (W4A16). + """ + + @classmethod + def is_supported( + cls, compute_capability: int | None = None + ) -> tuple[bool, str | None]: + if not current_platform.is_rocm(): + return False, "not running on ROCm" + from vllm.platforms.rocm import on_gfx11 + + if not on_gfx11(): + return False, "RDNA3 (gfx11xx) WMMA required" + if not hasattr(torch.ops._rocm_C, "mxfp4_gemm_rdna3"): + return False, "_rocm_C.mxfp4_gemm_rdna3 op not built" + return True, None + + @classmethod + def can_implement(cls, config: MxFp4LinearLayerConfig) -> tuple[bool, str | None]: + return True, None + + def process_weights_after_loading(self, layer: torch.nn.Module) -> None: + # layer.weight: [N, K/2] uint8, two E2M1 codes per byte (low nibble at + # the even K index). layer.weight_scale: [N, K/32] uint8 E8M0. + w = layer.weight.data + scale = layer.weight_scale.data + N, k_half = w.shape + K = k_half * 2 + + if N % 16 != 0 or K % _MXFP4_GROUP_SIZE != 0: + raise ValueError( + f"RDNA3 MXFP4 WMMA kernel needs N%16==0 and K%32==0, got " + f"N={N}, K={K}" + ) + + # [N, K/2] uint8 reinterpreted little-endian as int32 is already the + # [K/8, N] word layout the kernel reads, so repack = view + transpose. + b_q = w.contiguous().view(torch.int32).t().contiguous() # [K/8, N] + b_scale = scale.t().contiguous() # E8M0 [K/32, N] + + layer.weight = Parameter(b_q, requires_grad=False) + layer.weight_scale = Parameter(b_scale, requires_grad=False) + + def apply_weights( + self, + layer: torch.nn.Module, + x: torch.Tensor, + bias: torch.Tensor | None = None, + ) -> torch.Tensor: + out_shape = x.shape[:-1] + (layer.output_size_per_partition,) + x_2d = x.reshape(-1, x.shape[-1]) + if not x_2d.is_contiguous(): + x_2d = x_2d.contiguous() + + out = torch.ops._rocm_C.mxfp4_gemm_rdna3( + x_2d, layer.weight, layer.weight_scale + ) + if bias is not None: + out = out + bias + return out.view(out_shape) From 1f205e4533ca0830b7ba643e68a58024c95f7cd9 Mon Sep 17 00:00:00 2001 From: JartX Date: Thu, 25 Jun 2026 01:18:49 +0200 Subject: [PATCH 02/36] [ROCm][gfx1100] Add native HIP MXFP4 fused MoE for RDNA3 MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Compressed-tensors MXFP4 MoE models (e.g. MiniMax-MXFP4, Kimi-MXFP4) on RDNA3 fall back to the Triton-unfused expert path (AITER CK MXFP4 MoE is gfx950-only). This adds a native HIP fused-MoE kernel for gfx1100/1101/1102 that combines expert routing (sorted_token_ids / expert_ids) with the E2M1/E8M0 dequant + dot in one launch per GEMM, mirroring the existing moe_gptq_gemm_rdna3 W4A16 path but for MXFP4 (no zero-point; E8M0 block scale [E, K/32, N] uint8). The op moe_mxfp4_gemm_rdna3 is templated on BLOCK_SIZE_M (1/2/4/8) and reuses qdq_mxfp4_rdna3.cuh. CompressedTensorsW4A4Mxfp4RDNA3MoEMethod repacks experts to [E, K/8, N] (memory-lean per-expert pass so a multi-GB MoE fits next to the loaded weights) and runs gate_up -> SwiGLU -> down via RDNA3FusedMoEMixin, shared with the W4A16 path. Selected via rocm_moe_rdna on gfx1100. Tests: validated on a 7900 XTX vs a per-expert torch reference — gate_up, decode (block_size_m=1), the fused down-proj reduction, and a full gate_up -> SwiGLU -> down forward all match (rel < 6e-3, bf16+fp16). Signed-off-by: JartX --- CMakeLists.txt | 3 +- csrc/rocm/moe_mxfp4_gemm_rdna3.cu | 358 ++++++++++++++++++ csrc/rocm/ops.h | 9 + csrc/rocm/torch_bindings.cpp | 8 + vllm/_custom_ops.py | 50 +++ .../compressed_tensors_moe.py | 5 + ...compressed_tensors_moe_w4a4_mxfp4_rdna3.py | 53 +++ .../compressed_tensors_moe_wna16_rdna3.py | 160 +------- .../rdna3_moe_common.py | 154 ++++++++ .../compressed_tensors_moe/rocm_moe_rdna.py | 23 ++ 10 files changed, 677 insertions(+), 146 deletions(-) create mode 100644 csrc/rocm/moe_mxfp4_gemm_rdna3.cu create mode 100644 vllm/model_executor/layers/quantization/compressed_tensors/compressed_tensors_moe/compressed_tensors_moe_w4a4_mxfp4_rdna3.py create mode 100644 vllm/model_executor/layers/quantization/compressed_tensors/compressed_tensors_moe/rdna3_moe_common.py diff --git a/CMakeLists.txt b/CMakeLists.txt index d20cd6e8a995..39a161f1bb44 100644 --- a/CMakeLists.txt +++ b/CMakeLists.txt @@ -1353,7 +1353,8 @@ if(VLLM_GPU_LANG STREQUAL "HIP") "csrc/rocm/q_gemm_rdna3.cu" "csrc/rocm/q_gemm_rdna3_wmma.cu" "csrc/rocm/mxfp4_gemm_rdna3.cu" - "csrc/rocm/moe_q_gemm_rdna3.cu") + "csrc/rocm/moe_q_gemm_rdna3.cu" + "csrc/rocm/moe_mxfp4_gemm_rdna3.cu") endif() define_extension_target( diff --git a/csrc/rocm/moe_mxfp4_gemm_rdna3.cu b/csrc/rocm/moe_mxfp4_gemm_rdna3.cu new file mode 100644 index 000000000000..a90ccb014954 --- /dev/null +++ b/csrc/rocm/moe_mxfp4_gemm_rdna3.cu @@ -0,0 +1,358 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright contributors to the vLLM project +// +// MXFP4 fused MoE GEMM for RDNA3 (gfx1100/1101/1102). Fork of +// moe_q_gemm_rdna3.cu: expert routing (sorted_token_ids / expert_ids) + the +// MXFP4 scalar dequant+dot, in one kernel. No zero-point; E8M0 block scale is +// [E, K/32, N] uint8. Per-expert weights are [E, K/8, N] uint32 (E2M1). + +#include + +#include +#include +#include + +#include +#include +#include + +#include "qdq_mxfp4_rdna3.cuh" + +#if defined(__HIPCC__) && defined(__gfx1100__) + #define __HIP__RDNA3__ +#endif + +#define MOE_MXFP4_BLOCK_KN 256 +#define MOE_MXFP4_THREADS 256 + +namespace vllm { +namespace moe_mxfp4_rdna3 { + +using bf16_t = __hip_bfloat16; +using bf162_t = __hip_bfloat162; +using mxfp4_rdna3::dequant_mxfp4_8_f32; +using mxfp4_rdna3::mxfp4_e8m0_bias; + +#if defined(__HIP__RDNA3__) || !defined(__HIP_DEVICE_COMPILE__) + +template +__device__ __forceinline__ T tzero(); +template <> +__device__ __forceinline__ half tzero() { + return __float2half_rn(0.0f); +} +template <> +__device__ __forceinline__ bf16_t tzero() { + return __float2bfloat16(0.0f); +} + +__device__ __forceinline__ float mxfp4_to_f(half v) { return __half2float(v); } +__device__ __forceinline__ float mxfp4_to_f(bf16_t v) { + return __bfloat162float(v); +} + +__device__ __forceinline__ void atomic_add_pk4(half* addr, half2 v01, + half2 v23) { + unsigned long long* a = reinterpret_cast(addr); + unsigned long long old = *a; + 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(a, old, sum.u); + if (prev == old) break; + old = prev; + } +} +__device__ __forceinline__ void atomic_add_pk4(bf16_t* addr, bf162_t v01, + bf162_t v23) { + unsigned long long* a = reinterpret_cast(addr); + unsigned long long old = *a; + 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(a, old, sum.u); + if (prev == old) break; + old = prev; + } +} + +template +__global__ void moe_gemm_mxfp4_kernel_rdna3( + const T* __restrict__ a, T* __restrict__ c, + const uint32_t* __restrict__ b_q_weight, + const uint8_t* __restrict__ b_scales_e8m0, + const float* __restrict__ topk_weights, + 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, + const int size_n, const int size_k, const int groups, const int top_k, + const int expert_weight_stride, const int expert_scales_stride, + const bool mul_topk_weight, const int output_topk) { + const int t = threadIdx.x; + const int token_block = blockIdx.x; + const int offset_n = blockIdx.y * MOE_MXFP4_BLOCK_KN * 4; + const int offset_k = blockIdx.z * MOE_MXFP4_BLOCK_KN; + const int end_k = min(offset_k + MOE_MXFP4_BLOCK_KN, size_k); + const int n = offset_n + t * 4; + + 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; + + const uint32_t* expert_weights = + b_q_weight + (int64_t)expert_id * expert_weight_stride; + const uint8_t* expert_scales = + b_scales_e8m0 + (int64_t)expert_id * expert_scales_stride; + + constexpr int LDS_PAD = 8; + __shared__ T block_a[BLOCK_SIZE_M][MOE_MXFP4_BLOCK_KN + LDS_PAD]; + static_assert(MOE_MXFP4_BLOCK_KN == MOE_MXFP4_THREADS, + "BLOCK_KN must equal THREADS"); + const int offset_m_base = token_block * BLOCK_SIZE_M; + + // Stage activations into LDS via the routing table. + 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; + block_a[m][t] = (token_row < size_m) + ? a[(int64_t)token_row * size_k + offset_k + t] + : tzero(); + } + } + __syncthreads(); + if (n >= size_n) return; + + const int groupsize = size_k / groups; // == 32 for MXFP4 + int group = offset_k / groupsize; + int nextgroup = (group + 1) * groupsize; + // dequant_mxfp4_8_f32 decodes via bf16 bits (mant 7) for both dtypes. + constexpr uint32_t mant_bits = 7u; + + int qk = offset_k / 8; + const uint32_t* b_ptr = expert_weights + qk * size_n + n; + + int32_t bias[4]; + auto refresh = [&](int g) { + const uint8_t* sc = expert_scales + (int64_t)g * size_n; + #pragma unroll + for (int col = 0; col < 4; ++col) + bias[col] = mxfp4_e8m0_bias(sc[n + col], mant_bits); + }; + refresh(group); + + float block_c[BLOCK_SIZE_M][4]; + #pragma unroll + for (int m = 0; m < BLOCK_SIZE_M; ++m) + #pragma unroll + for (int col = 0; col < 4; ++col) block_c[m][col] = 0.0f; + + int k = offset_k; + while (k < end_k) { + if (k == nextgroup) { + group++; + nextgroup += groupsize; + refresh(group); + } + 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; + uint32_t wcol[4]; + __builtin_memcpy(wcol, &b_w[j], sizeof(int4)); + #pragma unroll + for (int col = 0; col < 4; ++col) { + float dq[8]; + dequant_mxfp4_8_f32(wcol[col], bias[col], dq); + #pragma unroll + for (int m = 0; m < BLOCK_SIZE_M; ++m) { + const T* ap = &block_a[m][a_off]; + float s = 0.0f; + #pragma unroll + for (int i = 0; i < 8; ++i) s += dq[i] * mxfp4_to_f(ap[i]); + block_c[m][col] += s; + } + } + } + k += 32; + } + + // Epilogue: topk weight + atomic-add (with optional output_topk reduction). + #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; + 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; + } + int64_t out_row = (output_topk > 0) ? (int64_t)(token_id / output_topk) + : (int64_t)token_id; + T* out = c + out_row * size_n + n; + if constexpr (std::is_same::value) { + half2 r01 = __halves2half2(__float2half_rn(block_c[m][0]), + __float2half_rn(block_c[m][1])); + half2 r23 = __halves2half2(__float2half_rn(block_c[m][2]), + __float2half_rn(block_c[m][3])); + atomic_add_pk4(out, r01, r23); + } else { + bf162_t r01, r23; + r01.x = __float2bfloat16(block_c[m][0]); + r01.y = __float2bfloat16(block_c[m][1]); + r23.x = __float2bfloat16(block_c[m][2]); + r23.y = __float2bfloat16(block_c[m][3]); + atomic_add_pk4(out, r01, r23); + } + } +} + +#else // non-RDNA3: empty stub for symbol parity +template +__global__ void moe_gemm_mxfp4_kernel_rdna3( + const T*, T*, const uint32_t*, const uint8_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 bool, + const int) {} +#endif + +template +void launch_moe_gemm_mxfp4(const T* a, T* c, const uint32_t* b_q_weight, + const uint8_t* b_scales_e8m0, + 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, + bool mul_topk_weight, int output_topk, + cudaStream_t stream) { + dim3 block(MOE_MXFP4_THREADS); + dim3 grid(num_token_blocks, + (size_n + MOE_MXFP4_BLOCK_KN * 4 - 1) / (MOE_MXFP4_BLOCK_KN * 4), + (size_k + MOE_MXFP4_BLOCK_KN - 1) / MOE_MXFP4_BLOCK_KN); + moe_gemm_mxfp4_kernel_rdna3<<>>( + a, c, b_q_weight, b_scales_e8m0, 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, mul_topk_weight, output_topk); +} + +template +void dispatch_moe_gemm_mxfp4( + const T* a, T* c, const uint32_t* b_q_weight, const uint8_t* b_scales_e8m0, + 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, bool mul_topk_weight, int output_topk, + cudaStream_t stream) { + auto L = [&](auto bsm) { + launch_moe_gemm_mxfp4( + a, c, b_q_weight, b_scales_e8m0, 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, + mul_topk_weight, output_topk, stream); + }; + switch (block_size_m) { + case 1: + L(std::integral_constant{}); + break; + case 2: + L(std::integral_constant{}); + break; + case 4: + L(std::integral_constant{}); + break; + case 8: + L(std::integral_constant{}); + break; + default: + TORCH_CHECK(false, "moe_mxfp4_gemm_rdna3: block_size_m must be 1/2/4/8"); + } +} + +} // namespace moe_mxfp4_rdna3 +} // namespace vllm + +// --------------------------------------------------------------------------- +// Public entry point. +// a [M, K] or [M*top_k, K] half/bf16 +// c [M*top_k, N] or reduced same dtype (pre-zeroed!) +// b_q_weight [E, K/8, N] uint32 (E2M1) +// b_scales_e8m0 [E, K/32, N] uint8 (E8M0) +// 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 +// --------------------------------------------------------------------------- +void moe_mxfp4_gemm_rdna3(torch::Tensor a, torch::Tensor c, + torch::Tensor b_q_weight, torch::Tensor b_scales_e8m0, + 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() && c.is_cuda() && b_q_weight.is_cuda(), + "tensors must be CUDA/HIP"); + TORCH_CHECK(a.dim() == 2 && c.dim() == 2, "a and c must be 2D"); + TORCH_CHECK(b_q_weight.dim() == 3, "b_q_weight must be [E, K/8, N]"); + TORCH_CHECK(b_scales_e8m0.dim() == 3, "b_scales_e8m0 must be [E, K/32, N]"); + TORCH_CHECK( + a.scalar_type() == torch::kHalf || a.scalar_type() == torch::kBFloat16, + "a must be half or bfloat16"); + TORCH_CHECK(b_scales_e8m0.scalar_type() == torch::kUInt8, + "b_scales_e8m0 must be uint8 (E8M0)"); + + 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_e8m0.size(1); // K/32 + int expert_weight_stride = (int)(b_q_weight.size(1) * b_q_weight.size(2)); + int expert_scales_stride = + (int)(b_scales_e8m0.size(1) * b_scales_e8m0.size(2)); + int num_token_blocks = (int)(sorted_token_ids.size(0) / block_size_m); + + const float* topk_w_ptr = + (topk_weights.numel() > 0) ? topk_weights.data_ptr() : nullptr; + + using bf16_t = vllm::moe_mxfp4_rdna3::bf16_t; + const uint32_t* bq = (const uint32_t*)b_q_weight.data_ptr(); + const uint8_t* bs = (const uint8_t*)b_scales_e8m0.data_ptr(); + const int32_t* sti = sorted_token_ids.data_ptr(); + const int32_t* eid = expert_ids.data_ptr(); + const int32_t* ntp = num_tokens_post_padded.data_ptr(); + + if (a.scalar_type() == torch::kHalf) { + vllm::moe_mxfp4_rdna3::dispatch_moe_gemm_mxfp4( + (const half*)a.data_ptr(), (half*)c.data_ptr(), bq, bs, topk_w_ptr, sti, + eid, ntp, num_token_blocks, size_m, size_n, size_k, groups, (int)top_k, + (int)block_size_m, expert_weight_stride, expert_scales_stride, + mul_topk_weight, (int)output_topk, stream); + } else { + vllm::moe_mxfp4_rdna3::dispatch_moe_gemm_mxfp4( + (const bf16_t*)a.data_ptr(), (bf16_t*)c.data_ptr(), bq, bs, topk_w_ptr, + sti, eid, ntp, num_token_blocks, size_m, size_n, size_k, groups, + (int)top_k, (int)block_size_m, expert_weight_stride, + expert_scales_stride, mul_topk_weight, (int)output_topk, stream); + } +} diff --git a/csrc/rocm/ops.h b/csrc/rocm/ops.h index b8fdd9b8a136..1256785f53e7 100644 --- a/csrc/rocm/ops.h +++ b/csrc/rocm/ops.h @@ -39,6 +39,15 @@ void moe_gptq_gemm_rdna3(torch::Tensor a, torch::Tensor c, int64_t block_size_m, bool mul_topk_weight, int64_t output_topk); +void moe_mxfp4_gemm_rdna3(torch::Tensor a, torch::Tensor c, + torch::Tensor b_q_weight, torch::Tensor b_scales_e8m0, + torch::Tensor topk_weights, + torch::Tensor sorted_token_ids, + torch::Tensor expert_ids, + torch::Tensor num_tokens_post_padded, int64_t top_k, + int64_t block_size_m, bool mul_topk_weight, + int64_t output_topk); + void paged_attention( torch::Tensor& out, torch::Tensor& exp_sums, torch::Tensor& max_logits, torch::Tensor& tmp_out, torch::Tensor& query, torch::Tensor& key_cache, diff --git a/csrc/rocm/torch_bindings.cpp b/csrc/rocm/torch_bindings.cpp index 76983023b27e..129822362bf1 100644 --- a/csrc/rocm/torch_bindings.cpp +++ b/csrc/rocm/torch_bindings.cpp @@ -64,6 +64,14 @@ TORCH_LIBRARY_EXPAND(TORCH_EXTENSION_NAME, rocm_ops) { "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); + + rocm_ops.def( + "moe_mxfp4_gemm_rdna3(Tensor a, Tensor! c, Tensor b_q_weight, " + "Tensor b_scales_e8m0, 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_mxfp4_gemm_rdna3", torch::kCUDA, &moe_mxfp4_gemm_rdna3); #endif // Custom attention op diff --git a/vllm/_custom_ops.py b/vllm/_custom_ops.py index 16e0df0df648..4b9a16d48e96 100644 --- a/vllm/_custom_ops.py +++ b/vllm/_custom_ops.py @@ -754,6 +754,56 @@ def _moe_gptq_gemm_rdna3_fake( return +def moe_mxfp4_gemm_rdna3( + a: torch.Tensor, + c: torch.Tensor, + b_q_weight: torch.Tensor, + b_scales_e8m0: 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_mxfp4_gemm_rdna3( + a, + c, + b_q_weight, + b_scales_e8m0, + 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_mxfp4_gemm_rdna3"): + + @register_fake("_rocm_C::moe_mxfp4_gemm_rdna3") + def _moe_mxfp4_gemm_rdna3_fake( + a: torch.Tensor, + c: torch.Tensor, + b_q_weight: torch.Tensor, + b_scales_e8m0: torch.Tensor, + topk_weights: torch.Tensor, + sorted_token_ids: torch.Tensor, + expert_ids: torch.Tensor, + num_tokens_post_padded: torch.Tensor, + top_k: int, + block_size_m: int, + mul_topk_weight: bool, + output_topk: int = 0, + ) -> None: + return + + if hasattr(torch.ops._C, "allspark_w8a16_gemm"): @register_fake("_C::allspark_w8a16_gemm") diff --git a/vllm/model_executor/layers/quantization/compressed_tensors/compressed_tensors_moe/compressed_tensors_moe.py b/vllm/model_executor/layers/quantization/compressed_tensors/compressed_tensors_moe/compressed_tensors_moe.py index 00221485233f..ead1f233aab7 100644 --- a/vllm/model_executor/layers/quantization/compressed_tensors/compressed_tensors_moe/compressed_tensors_moe.py +++ b/vllm/model_executor/layers/quantization/compressed_tensors/compressed_tensors_moe/compressed_tensors_moe.py @@ -64,6 +64,11 @@ def get_moe_method( format = scheme_dict.get("format") if quant_config._is_mxfp4(weight_quant): + from . import rocm_moe_rdna + + if rocm_moe_rdna.is_supported_mxfp4(weight_quant): + return rocm_moe_rdna.make_mxfp4_method(layer.moe_config) + from .compressed_tensors_moe_w4a4_mxfp4 import ( CompressedTensorsW4A4Mxfp4MoEMethod, ) diff --git a/vllm/model_executor/layers/quantization/compressed_tensors/compressed_tensors_moe/compressed_tensors_moe_w4a4_mxfp4_rdna3.py b/vllm/model_executor/layers/quantization/compressed_tensors/compressed_tensors_moe/compressed_tensors_moe_w4a4_mxfp4_rdna3.py new file mode 100644 index 000000000000..21b54d8e0e66 --- /dev/null +++ b/vllm/model_executor/layers/quantization/compressed_tensors/compressed_tensors_moe/compressed_tensors_moe_w4a4_mxfp4_rdna3.py @@ -0,0 +1,53 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""CompressedTensors MXFP4 MoE using the fused RDNA3 (gfx1100) HIP kernel. + +Uses ``moe_mxfp4_gemm_rdna3`` — one HIP launch per GEMM that fuses expert +routing + E2M1/E8M0 dequant + dot product with atomic output. No zero-point; +E8M0 block scale is ``[E, K/32, N]`` uint8. The two-GEMM forward is shared +with the W4A16 path via RDNA3FusedMoEMixin. +""" + +import torch + +from vllm import _custom_ops as ops +from vllm.model_executor.layers.quantization.compressed_tensors.compressed_tensors_moe.compressed_tensors_moe_w4a4_mxfp4 import ( # noqa: E501 + CompressedTensorsW4A4Mxfp4MoEMethod, +) + +from .rdna3_moe_common import RDNA3FusedMoEMixin, repack_experts + + +class CompressedTensorsW4A4Mxfp4RDNA3MoEMethod( + RDNA3FusedMoEMixin, CompressedTensorsW4A4Mxfp4MoEMethod +): + """MXFP4 MoE via the native RDNA3 HIP kernel (moe_mxfp4_gemm_rdna3).""" + + def process_weights_after_loading(self, layer: torch.nn.Module) -> None: + device = layer.w13_weight_packed.device + b_q13, b_s13 = repack_experts( + layer.w13_weight_packed.data, layer.w13_weight_scale.data, False + ) + del layer.w13_weight_packed, layer.w13_weight_scale + layer.w13_weight_packed = torch.nn.Parameter(b_q13, requires_grad=False) + layer.w13_weight_scale = torch.nn.Parameter(b_s13, requires_grad=False) + + b_q2, b_s2 = repack_experts( + layer.w2_weight_packed.data, layer.w2_weight_scale.data, False + ) + del layer.w2_weight_packed, layer.w2_weight_scale + layer.w2_weight_packed = torch.nn.Parameter(b_q2, requires_grad=False) + layer.w2_weight_scale = torch.nn.Parameter(b_s2, requires_grad=False) + layer.rdna3_empty_tw = torch.empty(0, device=device) + + def _gemm_w13(self, layer, a, c, tw, sti, eid, ntp, top_k, block_size_m, mul_tw): + ops.moe_mxfp4_gemm_rdna3( + a, c, layer.w13_weight_packed, layer.w13_weight_scale, tw, sti, eid, + ntp, top_k, block_size_m, mul_tw, 0, + ) + + def _gemm_w2(self, layer, a, c, tw, sti, eid, ntp, block_size_m, mul_tw, output_topk): + ops.moe_mxfp4_gemm_rdna3( + a, c, layer.w2_weight_packed, layer.w2_weight_scale, tw, sti, eid, + ntp, 1, block_size_m, mul_tw, output_topk, + ) diff --git a/vllm/model_executor/layers/quantization/compressed_tensors/compressed_tensors_moe/compressed_tensors_moe_wna16_rdna3.py b/vllm/model_executor/layers/quantization/compressed_tensors/compressed_tensors_moe/compressed_tensors_moe_wna16_rdna3.py index 25cc7b115f02..d28b103a30e2 100644 --- a/vllm/model_executor/layers/quantization/compressed_tensors/compressed_tensors_moe/compressed_tensors_moe_wna16_rdna3.py +++ b/vllm/model_executor/layers/quantization/compressed_tensors/compressed_tensors_moe/compressed_tensors_moe_wna16_rdna3.py @@ -15,17 +15,6 @@ 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, ) @@ -34,6 +23,8 @@ ) from vllm.scalar_type import scalar_types +from .rdna3_moe_common import RDNA3FusedMoEMixin + logger = init_logger(__name__) @@ -54,7 +45,9 @@ def _synthesize_qzeros( return pack_quantized_values_into_int32(zeros, scalar_types.uint4b8, packed_dim=1) -class CompressedTensorsWNA16RDNA3MoEMethod(CompressedTensorsWNA16MoEMethod): +class CompressedTensorsWNA16RDNA3MoEMethod( + RDNA3FusedMoEMixin, CompressedTensorsWNA16MoEMethod +): """W4A16 MoE using the fused RDNA3 HIP kernel (moe_gptq_gemm_rdna3). Weights are in RDNA3 format (shuffled int32 [E, K/8, N]), @@ -123,139 +116,16 @@ def process_weights_after_loading(self, layer: torch.nn.Module) -> None: ) 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, + def _gemm_w13(self, layer, a, c, tw, sti, eid, ntp, top_k, block_size_m, mul_tw): + ops.moe_gptq_gemm_rdna3( + a, c, layer.w13_weight_packed, layer.w13_weight_scale, + layer.w13_qzeros, tw, sti, eid, ntp, top_k, block_size_m, mul_tw, ) - 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 + def _gemm_w2(self, layer, a, c, tw, sti, eid, ntp, block_size_m, mul_tw, output_topk): + ops.moe_gptq_gemm_rdna3( + a, c, layer.w2_weight_packed, layer.w2_weight_scale, + layer.w2_qzeros, tw, sti, eid, ntp, 1, block_size_m, mul_tw, + output_topk, + ) diff --git a/vllm/model_executor/layers/quantization/compressed_tensors/compressed_tensors_moe/rdna3_moe_common.py b/vllm/model_executor/layers/quantization/compressed_tensors/compressed_tensors_moe/rdna3_moe_common.py new file mode 100644 index 000000000000..a604dbdfe139 --- /dev/null +++ b/vllm/model_executor/layers/quantization/compressed_tensors/compressed_tensors_moe/rdna3_moe_common.py @@ -0,0 +1,154 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""Shared orchestration for the native RDNA3 (gfx1100) fused-MoE methods. + +The W4A16 (INT4) and MXFP4 RDNA3 paths run the same two-GEMM forward +(gate_up -> SwiGLU -> down, with the down-proj output_topk reduction fusing +moe_sum); only the per-GEMM HIP op differs. That common flow lives here; each +method supplies ``_gemm_w13`` / ``_gemm_w2``. +""" + +import torch + +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, +) + + +def repack_experts(packed: torch.Tensor, scale: torch.Tensor, deinterleave: bool): + """[E, N, K/2] uint8 + [E, N, K/32] uint8 -> [E, K/8, N] int32 + [E, K/32, N]. + + Per-expert into pre-allocated outputs (peak ~2x ONE expert tensor, not the + whole stack) so a multi-GB MoE repack fits alongside the loaded weights. + Four consecutive packed bytes little-endian = a uint32 of 8 E2M1 codes + along K — exactly the [K/8, N] word layout the kernel reads. ``deinterleave`` + splits gate/up rows of w13 ([g0,u0,...] -> [g..., u...]) for GPT-OSS. + """ + E, N, kh = packed.shape + dev = packed.device + b_q = torch.empty(E, kh // 4, N, dtype=torch.int32, device=dev) + b_scale = torch.empty(E, scale.shape[2], N, dtype=torch.uint8, device=dev) + for e in range(E): + pe = packed[e] + se = scale[e] + if deinterleave: + pe = torch.cat([pe[::2], pe[1::2]], dim=0) + se = torch.cat([se[::2], se[1::2]], dim=0) + b_q[e] = pe.contiguous().view(torch.int32).t() + b_scale[e] = se.t() + return b_q, b_scale + + +class RDNA3FusedMoEMixin: + """Mixin: shared apply() + two-GEMM orchestration. Subclasses implement + ``_gemm_w13`` and ``_gemm_w2`` (the op-specific HIP launch).""" + + def _gemm_w13(self, layer, a, c, tw, sti, eid, ntp, top_k, block_size_m, mul_tw): + raise NotImplementedError + + def _gemm_w2(self, layer, a, c, tw, sti, eid, ntp, block_size_m, mul_tw, output_topk): + raise NotImplementedError + + def _rdna3_run( + self, + layer, + hidden_states, + topk_weights, + topk_ids, + activation, + apply_router_weight_on_input, + global_num_experts, + expert_map, + ): + 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 = 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 if num_tokens <= 4 else 4 + + sorted_token_ids, expert_ids, num_tokens_post_padded = moe_align_block_size( + topk_ids, block_size_m, global_num_experts, expert_map + ) + + # Use pre-allocated decode buffers when present and big enough. + w1_buf = getattr(layer, "rdna3_w1_buf", None) + if w1_buf is not None and total_tokens <= w1_buf.shape[0]: + w1_out = 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, dtype=dtype, device=device + ) + + topk_w_float = topk_weights.view(-1).float() + empty_tw = getattr(layer, "rdna3_empty_tw", None) + if empty_tw is None: + empty_tw = torch.empty(0, device=device) + + self._gemm_w13( + layer, + hidden_states, + w1_out, + 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, + ) + apply_moe_activation(activation, act_out, w1_out) + + out = torch.zeros(num_tokens, hidden_size, dtype=dtype, device=device) + self._gemm_w2( + layer, + act_out, + out, + topk_w_float if not apply_router_weight_on_input else empty_tw, + sorted_token_ids, + expert_ids, + num_tokens_post_padded, + block_size_m, + not apply_router_weight_on_input, + top_k, + ) + return out + + 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 self._rdna3_run( + layer, + x, + topk_weights, + topk_ids, + activation, + layer.apply_router_weight_on_input, + layer.global_num_experts, + layer.expert_map, + ) diff --git a/vllm/model_executor/layers/quantization/compressed_tensors/compressed_tensors_moe/rocm_moe_rdna.py b/vllm/model_executor/layers/quantization/compressed_tensors/compressed_tensors_moe/rocm_moe_rdna.py index e72caa0796a5..0a8056e34e6b 100644 --- a/vllm/model_executor/layers/quantization/compressed_tensors/compressed_tensors_moe/rocm_moe_rdna.py +++ b/vllm/model_executor/layers/quantization/compressed_tensors/compressed_tensors_moe/rocm_moe_rdna.py @@ -28,6 +28,29 @@ def is_supported(weight_quant) -> bool: ) +def is_supported_mxfp4(weight_quant) -> bool: + """Check if the native RDNA3 MXFP4 MoE kernel is available.""" + from vllm.platforms.rocm import on_gfx1100 + + return ( + on_gfx1100() + and hasattr(torch.ops, "_rocm_C") + and hasattr(torch.ops._rocm_C, "moe_mxfp4_gemm_rdna3") + ) + + +def make_mxfp4_method(moe_config): + """Create the native RDNA3 MXFP4 MoE method. Call after is_supported_mxfp4.""" + from .compressed_tensors_moe_w4a4_mxfp4_rdna3 import ( + CompressedTensorsW4A4Mxfp4RDNA3MoEMethod, + ) + + logger.info_once( + "Using CompressedTensorsW4A4Mxfp4RDNA3MoEMethod (native RDNA3 HIP kernel)" + ) + return CompressedTensorsW4A4Mxfp4RDNA3MoEMethod(moe_config) + + 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 From b3446d7c31697210e77f06eee62bdda6d83417b8 Mon Sep 17 00:00:00 2001 From: JartX Date: Thu, 25 Jun 2026 01:18:49 +0200 Subject: [PATCH 03/36] [ROCm][gfx1100] Add native HIP MXFP4 MoE for GPT-OSS on RDNA3 GPT-OSS is native-mxfp4 (GptOssMxfp4MoEMethod via the modular MoE runner, not compressed-tensors) and on gfx1100 falls back to Triton-unfused. This routes it to the native moe_mxfp4_gemm_rdna3 HIP kernel, handling the two GPT-OSS specifics: per-expert biases (w13_bias, w2_bias) and the SwiGLU-OAI clamped activation (alpha=1.702, beta=1.0, limit=7.0). w13 ships gate/up interleaved and is de-interleaved at load; biases and the topk-weighted moe_sum run in Python so the kernel is reused unchanged. The modular runner calls apply() through forward_modular once the method reports supports_internal_mk, non-monolithic, and a null fused-moe quant config. Selected from Mxfp4Config.get_quant_method on gfx1100. Tests: validated on an RX 7900 XTX end to end on openai/gpt-oss-20b (coherent generation, ~76 tok/s decode) and unit-tested vs an interleaved swigluoai per-expert reference (rel < 5e-3, bf16). Signed-off-by: JartX --- .../quantization/gpt_oss_mxfp4_rdna3.py | 129 ++++++++++++++++++ .../layers/quantization/mxfp4.py | 20 +++ 2 files changed, 149 insertions(+) create mode 100644 vllm/model_executor/layers/quantization/gpt_oss_mxfp4_rdna3.py diff --git a/vllm/model_executor/layers/quantization/gpt_oss_mxfp4_rdna3.py b/vllm/model_executor/layers/quantization/gpt_oss_mxfp4_rdna3.py new file mode 100644 index 000000000000..c15af665d338 --- /dev/null +++ b/vllm/model_executor/layers/quantization/gpt_oss_mxfp4_rdna3.py @@ -0,0 +1,129 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""GPT-OSS MXFP4 MoE on RDNA3 (gfx1100) via the native HIP kernel. + +GPT-OSS is native-mxfp4 (GptOssMxfp4MoEMethod), not compressed-tensors, and +differs from the plain MXFP4 MoE in two ways handled here: + - per-expert biases (w13_bias, w2_bias) + - the SwiGLU-OAI clamped activation (alpha=1.702, beta=1.0, limit=7.0) + +w13 ships gate/up *interleaved*; we de-interleave at load so the gate_up +output is contiguous [gate || up]. Biases and the moe_sum reduction run in +Python (output_topk=0), so the moe_mxfp4_gemm_rdna3 kernel is reused as-is. +""" + +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.moe_align_block_size import ( + moe_align_block_size, +) +from vllm.model_executor.layers.quantization.compressed_tensors.compressed_tensors_moe.rdna3_moe_common import ( # noqa: E501 + repack_experts, +) +from vllm.model_executor.layers.quantization.mxfp4 import GptOssMxfp4MoEMethod + +logger = init_logger(__name__) + +_ALPHA, _BETA, _LIMIT = 1.702, 1.0, 7.0 + + +class GptOssMxfp4RDNA3MoEMethod(GptOssMxfp4MoEMethod): + """GPT-OSS MXFP4 MoE via the native RDNA3 HIP kernel (moe_mxfp4_gemm_rdna3). + + Bypasses the modular MoE kernel: the runner calls apply() directly through + forward_modular. We report supports_internal_mk (skip MK construction), + not monolithic (route through forward_modular -> apply), and a null + fused-moe quant config (apply does the dequant itself). + """ + + @property + def supports_internal_mk(self) -> bool: + return True + + @property + def is_monolithic(self) -> bool: + return False + + def get_fused_moe_quant_config(self, layer): + return None + + def process_weights_after_loading(self, layer: RoutedExperts) -> None: + device = layer.w13_weight.device + # w13: de-interleave gate/up + repack, then free the source. + b_q13, b_s13 = repack_experts( + layer.w13_weight.data, layer.w13_weight_scale.data, deinterleave=True + ) + b13 = torch.cat( + [layer.w13_bias.data[:, ::2], layer.w13_bias.data[:, 1::2]], dim=1 + ).contiguous() + del layer.w13_weight, layer.w13_weight_scale + layer.w13_weight_packed = torch.nn.Parameter(b_q13, requires_grad=False) + layer.w13_weight_scale = torch.nn.Parameter(b_s13, requires_grad=False) + layer.w13_bias = torch.nn.Parameter(b13, requires_grad=False) + + # w2: plain repack (no de-interleave), then free the source. + b_q2, b_s2 = repack_experts( + layer.w2_weight.data, layer.w2_weight_scale.data, deinterleave=False + ) + del layer.w2_weight, layer.w2_weight_scale + layer.w2_weight_packed = torch.nn.Parameter(b_q2, requires_grad=False) + layer.w2_weight_scale = torch.nn.Parameter(b_s2, requires_grad=False) + layer.rdna3_empty_tw = torch.empty(0, device=device) + + @staticmethod + def _swiglu_oai(x: torch.Tensor) -> torch.Tensor: + # gate||up contiguous (w13 de-interleaved at load). SwiGLU-OAI: + # out = clamp(gate) * sigmoid(alpha*gate) * (clamp(up) + beta) + d = x.shape[-1] // 2 + gate = torch.clamp(x[..., :d], max=_LIMIT) + up = torch.clamp(x[..., d:], min=-_LIMIT, max=_LIMIT) + return gate * torch.sigmoid(_ALPHA * gate) * (up + _BETA) + + 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: + M = x.shape[0] + top_k = topk_ids.shape[1] + total = M * top_k + N_gate_up = layer.w13_weight_packed.shape[2] + hidden = layer.w2_weight_packed.shape[2] + dtype = x.dtype + device = x.device + + gne = layer.global_num_experts + if gne <= 0: + gne = layer.w13_weight_packed.shape[0] + block_size_m = 1 if M <= 4 else 4 + sti, eid, ntp = moe_align_block_size( + topk_ids, block_size_m, gne, layer.expert_map + ) + flat_experts = topk_ids.reshape(-1).long() + empty_tw = layer.rdna3_empty_tw + + # gate_up GEMM (no reduction) + per-expert bias + w1 = torch.zeros(total, N_gate_up, dtype=dtype, device=device) + ops.moe_mxfp4_gemm_rdna3( + x, w1, layer.w13_weight_packed, layer.w13_weight_scale, empty_tw, + sti, eid, ntp, top_k, block_size_m, False, 0, + ) + w1 = w1 + layer.w13_bias[flat_experts].to(dtype) + act = self._swiglu_oai(w1) + + # down GEMM (no reduction) + per-expert bias, then weighted moe_sum + w2o = torch.zeros(total, hidden, dtype=dtype, device=device) + ops.moe_mxfp4_gemm_rdna3( + act, w2o, layer.w2_weight_packed, layer.w2_weight_scale, empty_tw, + sti, eid, ntp, 1, block_size_m, False, 0, + ) + w2o = w2o + layer.w2_bias[flat_experts].to(dtype) + w2o = w2o * topk_weights.reshape(-1, 1).to(dtype) + return w2o.view(M, top_k, hidden).sum(dim=1) diff --git a/vllm/model_executor/layers/quantization/mxfp4.py b/vllm/model_executor/layers/quantization/mxfp4.py index 5ef5fd40d5eb..c6e16b59b175 100644 --- a/vllm/model_executor/layers/quantization/mxfp4.py +++ b/vllm/model_executor/layers/quantization/mxfp4.py @@ -86,6 +86,26 @@ def get_quant_method( ) return UnquantizedLinearMethod() elif isinstance(layer, RoutedExperts): + import torch + + from vllm.platforms import current_platform + + if ( + current_platform.is_rocm() + and hasattr(torch.ops, "_rocm_C") + and hasattr(torch.ops._rocm_C, "moe_mxfp4_gemm_rdna3") + ): + from vllm.platforms.rocm import on_gfx1100 + + if on_gfx1100(): + from vllm.model_executor.layers.quantization.gpt_oss_mxfp4_rdna3 import ( # noqa: E501 + GptOssMxfp4RDNA3MoEMethod, + ) + + logger.info_once( + "Using GptOssMxfp4RDNA3MoEMethod (native RDNA3 HIP kernel)" + ) + return GptOssMxfp4RDNA3MoEMethod(layer.moe_config) return GptOssMxfp4MoEMethod(layer.moe_config) elif isinstance(layer, Attention): logger.debug_once( From 1d28923b07194e921146b3e99dc549b94bdacb28 Mon Sep 17 00:00:00 2001 From: JartX Date: Thu, 25 Jun 2026 03:30:20 +0200 Subject: [PATCH 04/36] [ROCm][gfx1100] MXFP4 dense: route batched decode (M>=2) through WMMA The op sent M<=8 to the scalar GEMV, which is O(M) per weight read and degrades hard with batch size (M=8 took ~205us). The 16x16 WMMA tile is ~flat up to M=16 and far faster, so only single-token decode (M=1) needs the scalar path. Route M>=2 to WMMA: M=8 drops 205us -> ~44us (4.6x), and M=2..16 now beat bf16 ~1.5x on a 7900 XTX (they were losing 0.3-0.85x before). M=1 stays scalar (best there). Tested on RX 7900 XTX: correctness rel < 0.01 vs torch dequant ref (bf16+fp16, M=2..8, multiple shapes). Signed-off-by: JartX --- csrc/rocm/mxfp4_gemm_rdna3.cu | 7 +++++-- 1 file changed, 5 insertions(+), 2 deletions(-) diff --git a/csrc/rocm/mxfp4_gemm_rdna3.cu b/csrc/rocm/mxfp4_gemm_rdna3.cu index 72de607410d1..316c60cf64a7 100644 --- a/csrc/rocm/mxfp4_gemm_rdna3.cu +++ b/csrc/rocm/mxfp4_gemm_rdna3.cu @@ -970,8 +970,11 @@ torch::Tensor mxfp4_gemm_rdna3(torch::Tensor a, torch::Tensor b_q_weight, // Zero-init: K-split (gridDim.z > 1) accumulates atomically into c. at::Tensor c = torch::zeros({size_m, size_n}, opts); - // Small M (decode) -> scalar GEMV; larger M (prefill/batch) -> WMMA. - const bool use_scalar = size_m <= 8; + // M=1 (single-token decode) -> scalar GEMV; M>=2 (batched decode / prefill) + // -> WMMA. The scalar path is O(M) per weight read, so it degrades fast with + // batch; the WMMA 16x16 tile is flat ~constant up to M=16 and far faster + // (e.g. M=8: 205us scalar -> ~49us WMMA on a 7900 XTX). + const bool use_scalar = size_m == 1; if (a.scalar_type() == torch::kHalf) { namespace mx = vllm::mxfp4_rdna3_wmma; auto* ap = (const half*)a.data_ptr(); From 2631a7189e3c88fe85eae988f4f5de3a759cfbdf Mon Sep 17 00:00:00 2001 From: JartX Date: Thu, 25 Jun 2026 03:39:46 +0200 Subject: [PATCH 05/36] [ROCm][gfx1100] MXFP4: scope kernel comments to gfx1100 The kernels are gated on __gfx1100__ (compile) and on_gfx1100() (dispatch) and are only validated on gfx1100; drop the aspirational gfx1101/1102 mention from the file headers. Signed-off-by: JartX --- csrc/rocm/moe_mxfp4_gemm_rdna3.cu | 2 +- csrc/rocm/mxfp4_gemm_rdna3.cu | 2 +- 2 files changed, 2 insertions(+), 2 deletions(-) diff --git a/csrc/rocm/moe_mxfp4_gemm_rdna3.cu b/csrc/rocm/moe_mxfp4_gemm_rdna3.cu index a90ccb014954..109ce6414544 100644 --- a/csrc/rocm/moe_mxfp4_gemm_rdna3.cu +++ b/csrc/rocm/moe_mxfp4_gemm_rdna3.cu @@ -1,7 +1,7 @@ // SPDX-License-Identifier: Apache-2.0 // SPDX-FileCopyrightText: Copyright contributors to the vLLM project // -// MXFP4 fused MoE GEMM for RDNA3 (gfx1100/1101/1102). Fork of +// MXFP4 fused MoE GEMM for RDNA3 (gfx1100). Fork of // moe_q_gemm_rdna3.cu: expert routing (sorted_token_ids / expert_ids) + the // MXFP4 scalar dequant+dot, in one kernel. No zero-point; E8M0 block scale is // [E, K/32, N] uint8. Per-expert weights are [E, K/8, N] uint32 (E2M1). diff --git a/csrc/rocm/mxfp4_gemm_rdna3.cu b/csrc/rocm/mxfp4_gemm_rdna3.cu index 316c60cf64a7..76379ee932a0 100644 --- a/csrc/rocm/mxfp4_gemm_rdna3.cu +++ b/csrc/rocm/mxfp4_gemm_rdna3.cu @@ -1,7 +1,7 @@ // SPDX-License-Identifier: Apache-2.0 // SPDX-FileCopyrightText: Copyright contributors to the vLLM project // -// MXFP4 W4A16 GEMM for RDNA3 (gfx1100/1101/1102), WMMA path. Forked from the +// MXFP4 W4A16 GEMM for RDNA3 (gfx1100), WMMA path. Forked from the // GPTQ WMMA kernel (q_gemm_rdna3_wmma.cu): identical f16/bf16 tiling, only the // B-tile dequant differs (E2M1 + E8M0 group-32, no zeros). Weights are repacked // Python-side into the [K/8, N] uint32 layout the GPTQ kernel reads. From 1f5f284e179a592af28046bc53e6a4019eb02f9c Mon Sep 17 00:00:00 2001 From: JartX Date: Thu, 25 Jun 2026 12:49:41 +0200 Subject: [PATCH 06/36] [ROCm][gfx1100] MXFP4 MoE: add WMMA path for batched decode / prefill MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit The fused MoE kernel was scalar-only (per-tile dequant amortized over BLOCK_SIZE_M tokens with scalar FMAs, O(M) per weight read), so at high decode concurrency it stopped scaling and only tied the Triton-unfused path. This adds a 16x16x16 WMMA tile (one expert, single wave) — the dense mxfp4_gemm tiling plus the MoE expert gather (A rows via sorted_token_ids, weights via expert_ids) and a scatter epilogue with paired-atomic accumulation (supports the output_topk reduction). Selected via block_size_m == 16; the Python dispatch picks it once there are enough tokens to fill the tile. Both the compressed-tensors and GPT-OSS RDNA3 MoE methods inherit it (same op). Tested on RX 7900 XTX: WMMA output matches the scalar kernel (rel < 2e-4) and a torch dequant reference (rel < 4e-3), bf16+fp16, gate_up / +topk-weight / down+reduce. End to end on openai/gpt-oss-20b (vllm bench serve, --ignore-eos, 512/256) vs the Triton-unfused fallback, output tok/s: max_concurrency Triton HIP speedup 1 8.2 46.6 5.7x 8 88.7 109.8 1.24x 32 192.0 376.0 1.96x (at concurrency 32, 188 -> 376 tok/s and TPOT 227 -> 117 ms vs the scalar path.) Signed-off-by: JartX --- csrc/rocm/moe_mxfp4_gemm_rdna3.cu | 209 +++++++++++++++++- .../rdna3_moe_common.py | 4 +- .../quantization/gpt_oss_mxfp4_rdna3.py | 2 +- 3 files changed, 212 insertions(+), 3 deletions(-) diff --git a/csrc/rocm/moe_mxfp4_gemm_rdna3.cu b/csrc/rocm/moe_mxfp4_gemm_rdna3.cu index 109ce6414544..0c867939dbef 100644 --- a/csrc/rocm/moe_mxfp4_gemm_rdna3.cu +++ b/csrc/rocm/moe_mxfp4_gemm_rdna3.cu @@ -86,6 +86,182 @@ __device__ __forceinline__ void atomic_add_pk4(bf16_t* addr, bf162_t v01, } } +// --------------------------------------------------------------------------- +// WMMA path for the batched/prefill regime (BLOCK_SIZE_M == 16). The scalar +// kernel above amortizes the per-tile dequant over BLOCK_SIZE_M tokens with +// scalar FMAs (O(M) per weight read); at high decode concurrency that loses to +// a matrix-engine GEMM. This reuses the 16x16x16 WMMA tiling from the dense +// mxfp4_gemm_rdna3 kernel, with the MoE expert gather (A rows via +// sorted_token_ids, weights via expert_ids) and a scatter epilogue. +// --------------------------------------------------------------------------- +using v16fp16 = _Float16 __attribute__((ext_vector_type(16))); +using v16bf16 = __bf16 __attribute__((ext_vector_type(16))); +using v8fp32 = float __attribute__((ext_vector_type(8))); + +__device__ __forceinline__ v8fp32 wmma_mma(v16fp16 a, v16fp16 b, v8fp32 c) { + return __builtin_amdgcn_wmma_f32_16x16x16_f16_w32(a, b, c); +} +__device__ __forceinline__ v8fp32 wmma_mma(v16bf16 a, v16bf16 b, v8fp32 c) { + return __builtin_amdgcn_wmma_f32_16x16x16_bf16_w32(a, b, c); +} + +template +struct WmmaNative; +template <> +struct WmmaNative { + using elem = _Float16; + using v16 = v16fp16; +}; +template <> +struct WmmaNative { + using elem = __bf16; + using v16 = v16bf16; +}; + +using mxfp4_rdna3::dequant_mxfp4_8_bf16; +using mxfp4_rdna3::dequant_mxfp4_8_fp16; + +template +__device__ __forceinline__ TO wmoe_bitcast(FROM x) { + static_assert(sizeof(FROM) == sizeof(TO), "size"); + TO r; + __builtin_memcpy(&r, &x, sizeof(TO)); + return r; +} + +// Atomic-add a pair of adjacent fp16/bf16 cells (one uint32 word), CAS retry. +__device__ __forceinline__ void atomic_add_pk2(half* addr, half2 v) { + uint32_t* a = reinterpret_cast(addr); + uint32_t old = *a, prev; + do { + prev = old; + half2 sum = __hadd2(wmoe_bitcast(prev), v); + old = atomicCAS(a, prev, wmoe_bitcast(sum)); + } while (prev != old); +} +__device__ __forceinline__ void atomic_add_pk2(bf16_t* addr, bf162_t v) { + uint32_t* a = reinterpret_cast(addr); + uint32_t old = *a, prev; + do { + prev = old; + bf162_t sum = __hadd2(wmoe_bitcast(prev), v); + old = atomicCAS(a, prev, wmoe_bitcast(sum)); + } while (prev != old); +} + +// 16 tokens x 16 N tile, one expert, single wave (32 threads), full K. +template +__global__ void moe_gemm_mxfp4_wmma_kernel_rdna3( + const T* __restrict__ a, T* __restrict__ c, + const uint32_t* __restrict__ b_q_weight, + const uint8_t* __restrict__ b_scales_e8m0, + const float* __restrict__ topk_weights, + 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, + const int size_n, const int size_k, const int groups, const int top_k, + const int expert_weight_stride, const int expert_scales_stride, + const bool mul_topk_weight, const int output_topk) { + using E = typename WmmaNative::elem; + using V16 = typename WmmaNative::v16; + + const int token_block = blockIdx.x; + const int n_tile = blockIdx.y * 16; + if (token_block * 16 >= num_tokens_post_padded[0]) return; + const int expert_id = expert_ids[token_block]; + if (expert_id == -1 || n_tile >= size_n) return; + + const int lane = threadIdx.x; // 0..31 + const int lane_lo = lane & 15; + const int lane_hi = lane >> 4; + + const uint32_t* expert_weights = + b_q_weight + (int64_t)expert_id * expert_weight_stride; + const uint8_t* expert_scales = + b_scales_e8m0 + (int64_t)expert_id * expert_scales_stride; + + // Gather: lane_lo selects token slot; its source A row is token_id / top_k. + const int my_tok_id = sorted_token_ids[token_block * 16 + lane_lo]; + const int my_tok_row = my_tok_id / top_k; + const T* a_row = + (my_tok_row < size_m) ? (a + (int64_t)my_tok_row * size_k) : nullptr; + + constexpr uint32_t mant_bits = std::is_same::value ? 10u : 7u; + const int groupsize = size_k / groups; // == 32 + + v8fp32 c_acc = {0, 0, 0, 0, 0, 0, 0, 0}; + __shared__ T b_lds[16][16]; + const int actual_n = n_tile + lane_lo; // this lane's N column for B-dequant + + for (int k_tile = 0; k_tile < size_k; k_tile += 16) { + if (actual_n < size_n) { + const int my_k_octet = lane_hi; // 0 -> K[0..7], 1 -> K[8..15] + const uint32_t qa = + expert_weights[((k_tile / 8) + my_k_octet) * size_n + actual_n]; + const int32_t bias = mxfp4_e8m0_bias( + expert_scales[(k_tile / groupsize) * size_n + actual_n], mant_bits); + const int k_base = my_k_octet * 8; + if constexpr (std::is_same::value) { + half dq[8]; + dequant_mxfp4_8_fp16(qa, bias, dq); + #pragma unroll + for (int i = 0; i < 8; i++) b_lds[k_base + i][lane_lo] = dq[i]; + } else { + bf16_t dq[8]; + dequant_mxfp4_8_bf16(qa, bias, dq); + #pragma unroll + for (int i = 0; i < 8; i++) b_lds[k_base + i][lane_lo] = dq[i]; + } + } + // Single wave: lanes that wrote b_lds are the same wave that reads it; the + // WMMA reads all 16 rows, so a wave barrier is needed for cross-lane reads. + __syncthreads(); + + V16 a_frag, b_frag; + if (a_row) { + __builtin_memcpy(&a_frag, a_row + k_tile, sizeof(V16)); + } else { + #pragma unroll + for (int i = 0; i < 16; i++) a_frag[i] = (E)0; + } + #pragma unroll + for (int i = 0; i < 16; i++) b_frag[i] = wmoe_bitcast(b_lds[i][lane_lo]); + c_acc = wmma_mma(a_frag, b_frag, c_acc); + __syncthreads(); + } + + // Scatter epilogue: c_acc[i] -> token slot m = 2*i + lane_hi, col = lane_lo. + // Pair adjacent columns (lane_lo, lane_lo^1) into one uint32 atomic add. + const bool is_even_lane = (lane_lo & 1) == 0; + const int out_n = n_tile + lane_lo; + #pragma unroll + for (int i = 0; i < 8; i++) { + float other = __shfl_xor(c_acc[i], 1); + if (!is_even_lane) continue; + const int m_slot = 2 * i + lane_hi; + const int tok_id = sorted_token_ids[token_block * 16 + m_slot]; + if (tok_id / top_k >= size_m || out_n >= size_n) continue; + float v0 = c_acc[i], v1 = other; + if (mul_topk_weight && topk_weights != nullptr) { + const float tw = topk_weights[tok_id]; + v0 *= tw; + v1 *= tw; + } + const int64_t out_row = + (output_topk > 0) ? (int64_t)(tok_id / output_topk) : (int64_t)tok_id; + T* dst = c + out_row * size_n + out_n; + if constexpr (std::is_same::value) { + atomic_add_pk2(dst, __halves2half2(__float2half_rn(v0), + __float2half_rn(v1))); + } else { + bf162_t p; + p.x = __float2bfloat16(v0); + p.y = __float2bfloat16(v1); + atomic_add_pk2(dst, p); + } + } +} + template __global__ void moe_gemm_mxfp4_kernel_rdna3( const T* __restrict__ a, T* __restrict__ c, @@ -229,8 +405,32 @@ __global__ void moe_gemm_mxfp4_kernel_rdna3( const int32_t*, const int32_t*, const int32_t*, const int, const int, const int, const int, const int, const int, const int, const bool, const int) {} +template +__global__ void moe_gemm_mxfp4_wmma_kernel_rdna3( + const T*, T*, const uint32_t*, const uint8_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 bool, + const int) {} #endif +// WMMA launcher: 16x16 tiles, single wave, no K-split (enough blocks at the +// batch sizes that select this path). +template +void launch_moe_gemm_mxfp4_wmma( + const T* a, T* c, const uint32_t* b_q_weight, const uint8_t* b_scales_e8m0, + 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, + bool mul_topk_weight, int output_topk, cudaStream_t stream) { + dim3 block(32); + dim3 grid(num_token_blocks, (size_n + 15) / 16, 1); + moe_gemm_mxfp4_wmma_kernel_rdna3<<>>( + a, c, b_q_weight, b_scales_e8m0, 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, mul_topk_weight, output_topk); +} + template void launch_moe_gemm_mxfp4(const T* a, T* c, const uint32_t* b_q_weight, const uint8_t* b_scales_e8m0, @@ -282,8 +482,15 @@ void dispatch_moe_gemm_mxfp4( case 8: L(std::integral_constant{}); break; + case 16: + launch_moe_gemm_mxfp4_wmma( + a, c, b_q_weight, b_scales_e8m0, 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, + mul_topk_weight, output_topk, stream); + break; default: - TORCH_CHECK(false, "moe_mxfp4_gemm_rdna3: block_size_m must be 1/2/4/8"); + TORCH_CHECK(false, "moe_mxfp4_gemm_rdna3: block_size_m must be 1/2/4/8/16"); } } diff --git a/vllm/model_executor/layers/quantization/compressed_tensors/compressed_tensors_moe/rdna3_moe_common.py b/vllm/model_executor/layers/quantization/compressed_tensors/compressed_tensors_moe/rdna3_moe_common.py index a604dbdfe139..b26d02e8b0dc 100644 --- a/vllm/model_executor/layers/quantization/compressed_tensors/compressed_tensors_moe/rdna3_moe_common.py +++ b/vllm/model_executor/layers/quantization/compressed_tensors/compressed_tensors_moe/rdna3_moe_common.py @@ -76,7 +76,9 @@ def _rdna3_run( if global_num_experts <= 0: global_num_experts = layer.w13_weight_packed.shape[0] - block_size_m = 1 if num_tokens <= 4 else 4 + # Scalar tiles (1/4) for low batch; the 16-wide WMMA tile wins once + # there are enough tokens to fill it (high decode concurrency / prefill). + block_size_m = 1 if num_tokens <= 4 else (4 if num_tokens < 16 else 16) sorted_token_ids, expert_ids, num_tokens_post_padded = moe_align_block_size( topk_ids, block_size_m, global_num_experts, expert_map diff --git a/vllm/model_executor/layers/quantization/gpt_oss_mxfp4_rdna3.py b/vllm/model_executor/layers/quantization/gpt_oss_mxfp4_rdna3.py index c15af665d338..eb1a9fa44a47 100644 --- a/vllm/model_executor/layers/quantization/gpt_oss_mxfp4_rdna3.py +++ b/vllm/model_executor/layers/quantization/gpt_oss_mxfp4_rdna3.py @@ -102,7 +102,7 @@ def apply( gne = layer.global_num_experts if gne <= 0: gne = layer.w13_weight_packed.shape[0] - block_size_m = 1 if M <= 4 else 4 + block_size_m = 1 if M <= 4 else (4 if M < 16 else 16) sti, eid, ntp = moe_align_block_size( topk_ids, block_size_m, gne, layer.expert_map ) From 2b034a468c5f9c7f8d091572d2332a8201bf6b2e Mon Sep 17 00:00:00 2001 From: JartX Date: Thu, 25 Jun 2026 14:12:24 +0200 Subject: [PATCH 07/36] [ROCm][gfx1100] MXFP4 MoE: occupancy-aware block_size_m Share one select_block_size_m helper between the compressed-tensors and GPT-OSS RDNA3 MoE dispatch instead of the inlined per-method rule. Sub-16 batches leave the WMMA-16 tile mostly padding, so size them by expected occupancy (avg tokens/expert): bsm=1 at very low occupancy beats bsm=4 (+9% e2e at concurrency 8 on Qwen3.6-35B-A3B). Batches >=16 keep the WMMA tile, which clustered routing fills well even on 256-expert models and which also covers prefill. Signed-off-by: JartX --- .../rdna3_moe_common.py | 22 ++++++++++++++++--- .../quantization/gpt_oss_mxfp4_rdna3.py | 3 ++- 2 files changed, 21 insertions(+), 4 deletions(-) diff --git a/vllm/model_executor/layers/quantization/compressed_tensors/compressed_tensors_moe/rdna3_moe_common.py b/vllm/model_executor/layers/quantization/compressed_tensors/compressed_tensors_moe/rdna3_moe_common.py index b26d02e8b0dc..b86add850b65 100644 --- a/vllm/model_executor/layers/quantization/compressed_tensors/compressed_tensors_moe/rdna3_moe_common.py +++ b/vllm/model_executor/layers/quantization/compressed_tensors/compressed_tensors_moe/rdna3_moe_common.py @@ -20,6 +20,24 @@ ) +def select_block_size_m(num_tokens: int, top_k: int, num_experts: int) -> int: + """Pick the fused-MoE GEMM tile for the current batch. + + The WMMA-16 tile wins once a step has >=16 rows: real (clustered) routing + fills it well even on high-expert models, and this is also where prefill + lives (validated e2e: WMMA is fastest at concurrency 32 even on Qwen's 256 + experts). For sub-16 batches the 16-tile is mostly padding, so we fall back + to a scalar tile sized to the expected occupancy avg tokens/expert = + num_tokens*top_k/E (bsm=1 beats bsm=4 at low concurrency, +9% e2e at C=8). + """ + if num_experts <= 0 or top_k <= 0: + return 1 + if num_tokens >= 16: + return 16 + tpe = num_tokens * top_k / num_experts + return 4 if tpe >= 1.5 else 1 + + def repack_experts(packed: torch.Tensor, scale: torch.Tensor, deinterleave: bool): """[E, N, K/2] uint8 + [E, N, K/32] uint8 -> [E, K/8, N] int32 + [E, K/32, N]. @@ -76,9 +94,7 @@ def _rdna3_run( if global_num_experts <= 0: global_num_experts = layer.w13_weight_packed.shape[0] - # Scalar tiles (1/4) for low batch; the 16-wide WMMA tile wins once - # there are enough tokens to fill it (high decode concurrency / prefill). - block_size_m = 1 if num_tokens <= 4 else (4 if num_tokens < 16 else 16) + block_size_m = select_block_size_m(num_tokens, top_k, global_num_experts) sorted_token_ids, expert_ids, num_tokens_post_padded = moe_align_block_size( topk_ids, block_size_m, global_num_experts, expert_map diff --git a/vllm/model_executor/layers/quantization/gpt_oss_mxfp4_rdna3.py b/vllm/model_executor/layers/quantization/gpt_oss_mxfp4_rdna3.py index eb1a9fa44a47..437b8b4f356d 100644 --- a/vllm/model_executor/layers/quantization/gpt_oss_mxfp4_rdna3.py +++ b/vllm/model_executor/layers/quantization/gpt_oss_mxfp4_rdna3.py @@ -22,6 +22,7 @@ ) from vllm.model_executor.layers.quantization.compressed_tensors.compressed_tensors_moe.rdna3_moe_common import ( # noqa: E501 repack_experts, + select_block_size_m, ) from vllm.model_executor.layers.quantization.mxfp4 import GptOssMxfp4MoEMethod @@ -102,7 +103,7 @@ def apply( gne = layer.global_num_experts if gne <= 0: gne = layer.w13_weight_packed.shape[0] - block_size_m = 1 if M <= 4 else (4 if M < 16 else 16) + block_size_m = select_block_size_m(M, top_k, gne) sti, eid, ntp = moe_align_block_size( topk_ids, block_size_m, gne, layer.expert_map ) From 307d76afd625a4b7c0923f414b9763cf88f2cafd Mon Sep 17 00:00:00 2001 From: JartX Date: Thu, 25 Jun 2026 16:00:00 +0200 Subject: [PATCH 08/36] [ROCm][gfx1100] MXFP4 dense: LUT dequant for the scalar decode GEMV The M=1 scalar GEMV is compute-bound on gfx1100 (profiling a dense 27B decode: the MXFP4 GEMV is 43% of GPU time at ~47% of peak bandwidth; an isolated read-only kernel runs 5-7x faster than the GEMV, so the per-nibble arithmetic E2M1 decode is the limiter, not memory). Replace it with a 16-entry signed-magnitude LUT in LDS (one lookup per nibble) and fold the E8M0 block scale in once per group as 2^(s8-127) = uint_as_float(s8<<23). Exact fp32 mags make this more precise than the bf16-bits decode it replaces. Measured 1.5-2.6x on the real op at Qwen3.5-27B MLP shapes (gate/up, down, fused). Signed-off-by: JartX --- csrc/rocm/mxfp4_gemm_rdna3.cu | 26 +++++++++++++++++--------- 1 file changed, 17 insertions(+), 9 deletions(-) diff --git a/csrc/rocm/mxfp4_gemm_rdna3.cu b/csrc/rocm/mxfp4_gemm_rdna3.cu index 76379ee932a0..370427583d52 100644 --- a/csrc/rocm/mxfp4_gemm_rdna3.cu +++ b/csrc/rocm/mxfp4_gemm_rdna3.cu @@ -789,6 +789,14 @@ __global__ void gemm_mxfp4_scalar_rdna3( constexpr int LDS_PAD = 8; __shared__ T block_a[M_COUNT][MXFP4_BLOCK_KN + LDS_PAD]; + // E2M1 signed-magnitude LUT (exact fp32): one lookup replaces the per-nibble + // bit decode; the E8M0 scale folds in once per group as a multiply. Decode is + // compute-bound on gfx1100, so this is ~2.8x over the arithmetic decode. + __shared__ float w_lut[16]; + if (t < 16) { + const float mag[8] = {0.f, .5f, 1.f, 1.5f, 2.f, 3.f, 4.f, 6.f}; + w_lut[t] = (t & 0x8 ? -1.f : 1.f) * mag[t & 0x7]; + } if (offset_k + t < end_k) { #pragma unroll for (int m = 0; m < M_COUNT; ++m) { @@ -803,18 +811,18 @@ __global__ void gemm_mxfp4_scalar_rdna3( const int groupsize = size_k / groups; int group = offset_k / groupsize; int nextgroup = (group + 1) * groupsize; - // dequant_mxfp4_8_f32 always decodes via bf16 bits (mant 7), independent of - // the activation dtype T, so the E8M0 bias must use mant_bits=7. - constexpr uint32_t mant_bits = 7u; int qk = offset_k / 8; const uint32_t* b_ptr = b_q + qk * size_n + n; - int32_t bias[4]; + // Per-group E8M0 scale as a pure bit construct: 2^(s8-127) is the fp32 whose + // exponent field is s8 (IEEE bias 127), i.e. uint_as_float(s8 << 23). + float gscale[4]; auto refresh = [&](int g) { #pragma unroll for (int col = 0; col < 4; ++col) - bias[col] = mxfp4_e8m0_bias(b_scales_e8m0[g * size_n + n + col], mant_bits); + gscale[col] = __uint_as_float((uint32_t)b_scales_e8m0[g * size_n + n + col] + << 23); }; refresh(group); @@ -843,15 +851,15 @@ __global__ void gemm_mxfp4_scalar_rdna3( __builtin_memcpy(wcol, &b_w[j], sizeof(int4)); #pragma unroll for (int col = 0; col < 4; ++col) { - float dq[8]; - dequant_mxfp4_8_f32(wcol[col], bias[col], dq); + const uint32_t qa = wcol[col]; #pragma unroll for (int m = 0; m < M_COUNT; ++m) { const T* ap = &block_a[m][a_off]; float s = 0.0f; #pragma unroll - for (int i = 0; i < 8; ++i) s += dq[i] * mxfp4_to_f(ap[i]); - block_c[m][col] += s; + for (int i = 0; i < 8; ++i) + s += w_lut[(qa >> (4 * i)) & 0xFu] * mxfp4_to_f(ap[i]); + block_c[m][col] += gscale[col] * s; } } } From 5c71b949407c2d1d0d50d829018132d4503dc9c6 Mon Sep 17 00:00:00 2001 From: JartX Date: Thu, 25 Jun 2026 16:04:25 +0200 Subject: [PATCH 09/36] [ROCm][gfx1100] MXFP4 MoE: LUT dequant for the scalar expert GEMV Same compute-bound decode fix as the dense kernel, applied to the fused-MoE scalar path (the WMMA tile is unchanged). 16-entry signed-magnitude LUT in LDS + E8M0 scale as 2^(s8-127)=uint_as_float(s8<<23) per group, replacing the per-nibble arithmetic decode. ~2.3-2.5x on the scalar expert GEMVs at decode batch sizes (Qwen3.6-35B-A3B and GPT-OSS shapes); benefits both the compressed-tensors and GPT-OSS MoE paths. Correctness rel<3e-3 vs reference. Signed-off-by: JartX --- csrc/rocm/moe_mxfp4_gemm_rdna3.cu | 25 +++++++++++++++++-------- 1 file changed, 17 insertions(+), 8 deletions(-) diff --git a/csrc/rocm/moe_mxfp4_gemm_rdna3.cu b/csrc/rocm/moe_mxfp4_gemm_rdna3.cu index 0c867939dbef..b89f6278b272 100644 --- a/csrc/rocm/moe_mxfp4_gemm_rdna3.cu +++ b/csrc/rocm/moe_mxfp4_gemm_rdna3.cu @@ -296,6 +296,15 @@ __global__ void moe_gemm_mxfp4_kernel_rdna3( "BLOCK_KN must equal THREADS"); const int offset_m_base = token_block * BLOCK_SIZE_M; + // E2M1 signed-magnitude LUT (exact fp32): one lookup replaces the per-nibble + // bit decode; the E8M0 scale folds in once per group as a multiply. The + // decode is compute-bound on gfx1100 (~2.8x over the arithmetic decode). + __shared__ float w_lut[16]; + if (t < 16) { + const float mag[8] = {0.f, .5f, 1.f, 1.5f, 2.f, 3.f, 4.f, 6.f}; + w_lut[t] = (t & 0x8 ? -1.f : 1.f) * mag[t & 0x7]; + } + // Stage activations into LDS via the routing table. if (offset_k + t < end_k) { #pragma unroll @@ -313,18 +322,18 @@ __global__ void moe_gemm_mxfp4_kernel_rdna3( const int groupsize = size_k / groups; // == 32 for MXFP4 int group = offset_k / groupsize; int nextgroup = (group + 1) * groupsize; - // dequant_mxfp4_8_f32 decodes via bf16 bits (mant 7) for both dtypes. - constexpr uint32_t mant_bits = 7u; int qk = offset_k / 8; const uint32_t* b_ptr = expert_weights + qk * size_n + n; - int32_t bias[4]; + // Per-group E8M0 scale as a pure bit construct: 2^(s8-127) is the fp32 whose + // exponent field is s8 (IEEE bias 127), i.e. uint_as_float(s8 << 23). + float gscale[4]; auto refresh = [&](int g) { const uint8_t* sc = expert_scales + (int64_t)g * size_n; #pragma unroll for (int col = 0; col < 4; ++col) - bias[col] = mxfp4_e8m0_bias(sc[n + col], mant_bits); + gscale[col] = __uint_as_float((uint32_t)sc[n + col] << 23); }; refresh(group); @@ -353,15 +362,15 @@ __global__ void moe_gemm_mxfp4_kernel_rdna3( __builtin_memcpy(wcol, &b_w[j], sizeof(int4)); #pragma unroll for (int col = 0; col < 4; ++col) { - float dq[8]; - dequant_mxfp4_8_f32(wcol[col], bias[col], dq); + const uint32_t qa = wcol[col]; #pragma unroll for (int m = 0; m < BLOCK_SIZE_M; ++m) { const T* ap = &block_a[m][a_off]; float s = 0.0f; #pragma unroll - for (int i = 0; i < 8; ++i) s += dq[i] * mxfp4_to_f(ap[i]); - block_c[m][col] += s; + for (int i = 0; i < 8; ++i) + s += w_lut[(qa >> (4 * i)) & 0xFu] * mxfp4_to_f(ap[i]); + block_c[m][col] += gscale[col] * s; } } } From 3f7ef40288a519800c3dacb466ce2b589866949f Mon Sep 17 00:00:00 2001 From: JartX Date: Thu, 25 Jun 2026 18:22:18 +0200 Subject: [PATCH 10/36] [ROCm][gfx1100] MXFP4: LUT dequant for WMMA path + raise scalar dispatch Extends the LUT decode to the WMMA tiles (dense 16x16/32x16/64x64 + the fused MoE tile) via a shared dequant_mxfp4_8_lut helper: per nibble it looks up the E2M1 magnitude bits and applies only the E8M0 exponent add, skipping the field construction. Also raises the dense scalar dispatch to size_m<=8 (was ==1): the LUT made the scalar GEMV beat the WMMA 16-tile up to ~M=8 (1.3-1.6x), and the tile wastes most rows there anyway. Measured e2e (vllm bench serve, vs the scalar-only-LUT build): GPT-OSS C=32 +25% (307->383 tok/s, WMMA-LUT), Qwen3.5-27B dense C=8 +9% (scalar dispatch), C=32 +6%; Qwen3.6-35B-A3B C=16/32 +4-5%. WMMA correctness rel<2.3e-3. Signed-off-by: JartX --- csrc/rocm/moe_mxfp4_gemm_rdna3.cu | 21 ++++++----- csrc/rocm/mxfp4_gemm_rdna3.cu | 62 +++++++++++++++---------------- csrc/rocm/qdq_mxfp4_rdna3.cuh | 13 +++++++ 3 files changed, 55 insertions(+), 41 deletions(-) diff --git a/csrc/rocm/moe_mxfp4_gemm_rdna3.cu b/csrc/rocm/moe_mxfp4_gemm_rdna3.cu index b89f6278b272..f2effee3988e 100644 --- a/csrc/rocm/moe_mxfp4_gemm_rdna3.cu +++ b/csrc/rocm/moe_mxfp4_gemm_rdna3.cu @@ -120,6 +120,9 @@ struct WmmaNative { using mxfp4_rdna3::dequant_mxfp4_8_bf16; using mxfp4_rdna3::dequant_mxfp4_8_fp16; +using mxfp4_rdna3::dequant_mxfp4_8_lut; +using mxfp4_rdna3::mxfp4_e2m1_to_bf16_bits; +using mxfp4_rdna3::mxfp4_e2m1_to_fp16_bits; template __device__ __forceinline__ TO wmoe_bitcast(FROM x) { @@ -191,6 +194,11 @@ __global__ void moe_gemm_mxfp4_wmma_kernel_rdna3( v8fp32 c_acc = {0, 0, 0, 0, 0, 0, 0, 0}; __shared__ T b_lds[16][16]; + __shared__ uint16_t w_lut[16]; // E2M1 magnitude bits for T, filled once + if (lane < 16) + w_lut[lane] = std::is_same::value ? mxfp4_e2m1_to_fp16_bits(lane) + : mxfp4_e2m1_to_bf16_bits(lane); + __syncthreads(); const int actual_n = n_tile + lane_lo; // this lane's N column for B-dequant for (int k_tile = 0; k_tile < size_k; k_tile += 16) { @@ -201,17 +209,10 @@ __global__ void moe_gemm_mxfp4_wmma_kernel_rdna3( const int32_t bias = mxfp4_e8m0_bias( expert_scales[(k_tile / groupsize) * size_n + actual_n], mant_bits); const int k_base = my_k_octet * 8; - if constexpr (std::is_same::value) { - half dq[8]; - dequant_mxfp4_8_fp16(qa, bias, dq); - #pragma unroll - for (int i = 0; i < 8; i++) b_lds[k_base + i][lane_lo] = dq[i]; - } else { - bf16_t dq[8]; - dequant_mxfp4_8_bf16(qa, bias, dq); + T dq[8]; + dequant_mxfp4_8_lut(qa, bias, w_lut, dq); #pragma unroll - for (int i = 0; i < 8; i++) b_lds[k_base + i][lane_lo] = dq[i]; - } + for (int i = 0; i < 8; i++) b_lds[k_base + i][lane_lo] = dq[i]; } // Single wave: lanes that wrote b_lds are the same wave that reads it; the // WMMA reads all 16 rows, so a wave barrier is needed for cross-lane reads. diff --git a/csrc/rocm/mxfp4_gemm_rdna3.cu b/csrc/rocm/mxfp4_gemm_rdna3.cu index 370427583d52..688946c2534c 100644 --- a/csrc/rocm/mxfp4_gemm_rdna3.cu +++ b/csrc/rocm/mxfp4_gemm_rdna3.cu @@ -29,6 +29,9 @@ using bf16_t = __hip_bfloat16; using mxfp4_rdna3::dequant_mxfp4_8_bf16; using mxfp4_rdna3::dequant_mxfp4_8_f32; using mxfp4_rdna3::dequant_mxfp4_8_fp16; +using mxfp4_rdna3::dequant_mxfp4_8_lut; +using mxfp4_rdna3::mxfp4_e2m1_to_bf16_bits; +using mxfp4_rdna3::mxfp4_e2m1_to_fp16_bits; using mxfp4_rdna3::mxfp4_e8m0_bias; // K-split heuristic: more blocks along gridDim.z → more resident waves. Mirror @@ -131,6 +134,11 @@ __global__ void gemm_mxfp4_wmma_kernel_16x16_1w( const int k_end = k_start + k_per_split; __shared__ T b_lds[16][16]; + __shared__ uint16_t w_lut[16]; // E2M1 magnitude bits for T, filled once + if (lane < 16) + w_lut[lane] = std::is_same::value ? mxfp4_e2m1_to_fp16_bits(lane) + : mxfp4_e2m1_to_bf16_bits(lane); + __syncthreads(); for (int k_tile = k_start; k_tile < k_end; k_tile += 16) { // ---- Dequant 16x16 B tile into LDS ---- @@ -151,17 +159,10 @@ __global__ void gemm_mxfp4_wmma_kernel_16x16_1w( const int k_base = my_k_octet * 8; - if constexpr (std::is_same::value) { - half dq[8]; - dequant_mxfp4_8_fp16(qa, bias_u16, dq); + T dq[8]; + dequant_mxfp4_8_lut(qa, bias_u16, w_lut, dq); #pragma unroll - for (int i = 0; i < 8; i++) b_lds[k_base + i][my_n] = dq[i]; - } else { - bf16_t dq[8]; - dequant_mxfp4_8_bf16(qa, bias_u16, dq); - #pragma unroll - for (int i = 0; i < 8; i++) b_lds[k_base + i][my_n] = dq[i]; - } + for (int i = 0; i < 8; i++) b_lds[k_base + i][my_n] = dq[i]; } // Single-wave block: no __syncthreads needed. @@ -286,6 +287,11 @@ __global__ void gemm_mxfp4_wmma_kernel_32x16_2w( const int k_end = k_start + k_per_split; __shared__ T b_lds[2][16][16]; + __shared__ uint16_t w_lut[16]; // E2M1 magnitude bits for T, filled once + if (tid < 16) + w_lut[tid] = std::is_same::value ? mxfp4_e2m1_to_fp16_bits(tid) + : mxfp4_e2m1_to_bf16_bits(tid); + __syncthreads(); auto dequant_into = [&](int buf, int k_tile) { if (wave_id != 0) return; @@ -298,17 +304,10 @@ __global__ void gemm_mxfp4_wmma_kernel_32x16_2w( const uint32_t s8 = b_scales_e8m0[g * size_n + actual_n]; const int32_t bias = mxfp4_e8m0_bias(s8, mant_bits); const int k_base = lane_hi * 8; - if constexpr (std::is_same::value) { - half dq[8]; - dequant_mxfp4_8_fp16(qa, bias, dq); + T dq[8]; + dequant_mxfp4_8_lut(qa, bias, w_lut, dq); #pragma unroll - for (int i = 0; i < 8; i++) b_lds[buf][k_base + i][my_n] = dq[i]; - } else { - bf16_t dq[8]; - dequant_mxfp4_8_bf16(qa, bias, dq); - #pragma unroll - for (int i = 0; i < 8; i++) b_lds[buf][k_base + i][my_n] = dq[i]; - } + for (int i = 0; i < 8; i++) b_lds[buf][k_base + i][my_n] = dq[i]; }; dequant_into(0, k_start); @@ -409,6 +408,12 @@ __global__ void gemm_mxfp4_wmma_kernel_64x64_4w( const int k_end = k_start + k_per_split; __shared__ T b_lds[2][16][64]; + __shared__ uint16_t w_lut[16]; // E2M1 magnitude bits for T, filled once + if (threadIdx.x < 16) + w_lut[threadIdx.x] = std::is_same::value + ? mxfp4_e2m1_to_fp16_bits(threadIdx.x) + : mxfp4_e2m1_to_bf16_bits(threadIdx.x); + __syncthreads(); auto dequant_into = [&](int buf, int k_tile) { const int my_n_in_tile = wave_id * 16 + lane_lo; @@ -420,17 +425,10 @@ __global__ void gemm_mxfp4_wmma_kernel_64x64_4w( const uint32_t s8 = b_scales_e8m0[g * size_n + actual_n]; const int32_t bias = mxfp4_e8m0_bias(s8, mant_bits); const int k_base = lane_hi * 8; - if constexpr (std::is_same::value) { - half dq[8]; - dequant_mxfp4_8_fp16(qa, bias, dq); + T dq[8]; + dequant_mxfp4_8_lut(qa, bias, w_lut, dq); #pragma unroll - for (int i = 0; i < 8; i++) b_lds[buf][k_base + i][my_n_in_tile] = dq[i]; - } else { - bf16_t dq[8]; - dequant_mxfp4_8_bf16(qa, bias, dq); - #pragma unroll - for (int i = 0; i < 8; i++) b_lds[buf][k_base + i][my_n_in_tile] = dq[i]; - } + for (int i = 0; i < 8; i++) b_lds[buf][k_base + i][my_n_in_tile] = dq[i]; }; dequant_into(0, k_start); @@ -982,7 +980,9 @@ torch::Tensor mxfp4_gemm_rdna3(torch::Tensor a, torch::Tensor b_q_weight, // -> WMMA. The scalar path is O(M) per weight read, so it degrades fast with // batch; the WMMA 16x16 tile is flat ~constant up to M=16 and far faster // (e.g. M=8: 205us scalar -> ~49us WMMA on a 7900 XTX). - const bool use_scalar = size_m == 1; + // The LUT decode made the scalar GEMV fast enough to beat the WMMA 16-tile up + // to M~8 (1.3-1.6x at M=4/8 on gfx1100); above that the matrix engine wins. + const bool use_scalar = size_m <= 8; if (a.scalar_type() == torch::kHalf) { namespace mx = vllm::mxfp4_rdna3_wmma; auto* ap = (const half*)a.data_ptr(); diff --git a/csrc/rocm/qdq_mxfp4_rdna3.cuh b/csrc/rocm/qdq_mxfp4_rdna3.cuh index b41712f5a9bd..4e2d47676081 100644 --- a/csrc/rocm/qdq_mxfp4_rdna3.cuh +++ b/csrc/rocm/qdq_mxfp4_rdna3.cuh @@ -102,6 +102,19 @@ MXFP4_D void dequant_mxfp4_8_fp16(uint32_t qa, int32_t bias_u16, } } +// LUT decode: `lut` holds the 16 E2M1 magnitude bit-patterns for T (filled once +// from mxfp4_e2m1_to_{bf16,fp16}_bits), so per nibble we skip the arithmetic +// field construction and just apply the E8M0 exponent add. T = bf16_t or half. +template +MXFP4_D void dequant_mxfp4_8_lut(uint32_t qa, int32_t bias_u16, + const uint16_t* lut, T (&dq)[8]) { +#pragma unroll + for (int i = 0; i < 8; i++) { + uint16_t b = mxfp4_apply_e8m0_bits(lut[(qa >> (4 * i)) & 0xFu], bias_u16); + __builtin_memcpy(&dq[i], &b, sizeof(T)); + } +} + // fp32 output (bf16 bits widened by a free <<16) for a scalar dot path. MXFP4_D void dequant_mxfp4_8_f32(uint32_t qa, int32_t bias_u16, float (&dq)[8]) { From 23722a3c6dfaaaf5c05f1da248e1616424cf2b5e Mon Sep 17 00:00:00 2001 From: JartX Date: Thu, 25 Jun 2026 18:51:20 +0200 Subject: [PATCH 11/36] [ROCm][gfx1100] MXFP4: move GPT-OSS gfx1100 selection out of generic Mxfp4Config The gfx1100 GPT-OSS MoE method selection was inlined in the generic Mxfp4Config.get_quant_method. Move it behind a maybe_rdna3_gpt_oss_method() helper in the gfx1100-scoped module so the generic config only carries a one-line delegation (mirrors rocm_moe_rdna for the compressed-tensors path). No behavior change; selection log/path verified on gfx1100. Signed-off-by: JartX --- .../quantization/gpt_oss_mxfp4_rdna3.py | 23 ++++++++++++++++ .../layers/quantization/mxfp4.py | 26 +++++-------------- 2 files changed, 29 insertions(+), 20 deletions(-) diff --git a/vllm/model_executor/layers/quantization/gpt_oss_mxfp4_rdna3.py b/vllm/model_executor/layers/quantization/gpt_oss_mxfp4_rdna3.py index 437b8b4f356d..1413ae96884d 100644 --- a/vllm/model_executor/layers/quantization/gpt_oss_mxfp4_rdna3.py +++ b/vllm/model_executor/layers/quantization/gpt_oss_mxfp4_rdna3.py @@ -31,6 +31,29 @@ _ALPHA, _BETA, _LIMIT = 1.702, 1.0, 7.0 +def maybe_rdna3_gpt_oss_method(layer): + """Return the native RDNA3 GPT-OSS MXFP4 MoE method on gfx1100, else None. + + Keeps the gfx1100 selection out of the generic ``Mxfp4Config`` so the only + footprint there is a one-line delegation (mirrors ``rocm_moe_rdna`` for the + compressed-tensors path). + """ + from vllm.platforms import current_platform + + if not ( + current_platform.is_rocm() + and hasattr(torch.ops, "_rocm_C") + and hasattr(torch.ops._rocm_C, "moe_mxfp4_gemm_rdna3") + ): + return None + from vllm.platforms.rocm import on_gfx1100 + + if not on_gfx1100(): + return None + logger.info_once("Using GptOssMxfp4RDNA3MoEMethod (native RDNA3 HIP kernel)") + return GptOssMxfp4RDNA3MoEMethod(layer.moe_config) + + class GptOssMxfp4RDNA3MoEMethod(GptOssMxfp4MoEMethod): """GPT-OSS MXFP4 MoE via the native RDNA3 HIP kernel (moe_mxfp4_gemm_rdna3). diff --git a/vllm/model_executor/layers/quantization/mxfp4.py b/vllm/model_executor/layers/quantization/mxfp4.py index c6e16b59b175..8174a860d221 100644 --- a/vllm/model_executor/layers/quantization/mxfp4.py +++ b/vllm/model_executor/layers/quantization/mxfp4.py @@ -86,27 +86,13 @@ def get_quant_method( ) return UnquantizedLinearMethod() elif isinstance(layer, RoutedExperts): - import torch - - from vllm.platforms import current_platform + from vllm.model_executor.layers.quantization.gpt_oss_mxfp4_rdna3 import ( + maybe_rdna3_gpt_oss_method, + ) - if ( - current_platform.is_rocm() - and hasattr(torch.ops, "_rocm_C") - and hasattr(torch.ops._rocm_C, "moe_mxfp4_gemm_rdna3") - ): - from vllm.platforms.rocm import on_gfx1100 - - if on_gfx1100(): - from vllm.model_executor.layers.quantization.gpt_oss_mxfp4_rdna3 import ( # noqa: E501 - GptOssMxfp4RDNA3MoEMethod, - ) - - logger.info_once( - "Using GptOssMxfp4RDNA3MoEMethod (native RDNA3 HIP kernel)" - ) - return GptOssMxfp4RDNA3MoEMethod(layer.moe_config) - return GptOssMxfp4MoEMethod(layer.moe_config) + return maybe_rdna3_gpt_oss_method(layer) or GptOssMxfp4MoEMethod( + layer.moe_config + ) elif isinstance(layer, Attention): logger.debug_once( "MXFP4 attention layer is not implemented. " From b16e783548488c30075d3584d90d66a08f662540 Mon Sep 17 00:00:00 2001 From: JartX Date: Thu, 25 Jun 2026 18:56:16 +0200 Subject: [PATCH 12/36] [ROCm][gfx1100] MXFP4: apply pre-commit formatting (ruff + clang-format) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit No functional change — ruff-format line reflow on the Python paths and clang-format on the HIP kernels/headers to satisfy the lint hooks. Signed-off-by: JartX --- csrc/rocm/moe_mxfp4_gemm_rdna3.cu | 44 +++++++++---------- csrc/rocm/mxfp4_gemm_rdna3.cu | 23 ++++++---- csrc/rocm/qdq_mxfp4_rdna3.cuh | 21 +++++---- .../kernels/quantization/test_mxfp4_rdna3.py | 8 +--- .../kernels/linear/mxfp4/rocm.py | 11 ++--- ...compressed_tensors_moe_w4a4_mxfp4_rdna3.py | 32 +++++++++++--- .../compressed_tensors_moe_wna16_rdna3.py | 33 +++++++++++--- .../rdna3_moe_common.py | 4 +- .../quantization/gpt_oss_mxfp4_rdna3.py | 28 ++++++++++-- 9 files changed, 130 insertions(+), 74 deletions(-) diff --git a/csrc/rocm/moe_mxfp4_gemm_rdna3.cu b/csrc/rocm/moe_mxfp4_gemm_rdna3.cu index f2effee3988e..8cfb4261462b 100644 --- a/csrc/rocm/moe_mxfp4_gemm_rdna3.cu +++ b/csrc/rocm/moe_mxfp4_gemm_rdna3.cu @@ -174,7 +174,7 @@ __global__ void moe_gemm_mxfp4_wmma_kernel_rdna3( const int expert_id = expert_ids[token_block]; if (expert_id == -1 || n_tile >= size_n) return; - const int lane = threadIdx.x; // 0..31 + const int lane = threadIdx.x; // 0..31 const int lane_lo = lane & 15; const int lane_hi = lane >> 4; @@ -226,7 +226,8 @@ __global__ void moe_gemm_mxfp4_wmma_kernel_rdna3( for (int i = 0; i < 16; i++) a_frag[i] = (E)0; } #pragma unroll - for (int i = 0; i < 16; i++) b_frag[i] = wmoe_bitcast(b_lds[i][lane_lo]); + for (int i = 0; i < 16; i++) + b_frag[i] = wmoe_bitcast(b_lds[i][lane_lo]); c_acc = wmma_mma(a_frag, b_frag, c_acc); __syncthreads(); } @@ -252,8 +253,8 @@ __global__ void moe_gemm_mxfp4_wmma_kernel_rdna3( (output_topk > 0) ? (int64_t)(tok_id / output_topk) : (int64_t)tok_id; T* dst = c + out_row * size_n + out_n; if constexpr (std::is_same::value) { - atomic_add_pk2(dst, __halves2half2(__float2half_rn(v0), - __float2half_rn(v1))); + atomic_add_pk2(dst, + __halves2half2(__float2half_rn(v0), __float2half_rn(v1))); } else { bf162_t p; p.x = __float2bfloat16(v0); @@ -411,16 +412,14 @@ __global__ void moe_gemm_mxfp4_kernel_rdna3( #else // non-RDNA3: empty stub for symbol parity template __global__ void moe_gemm_mxfp4_kernel_rdna3( - const T*, T*, const uint32_t*, const uint8_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 bool, - const int) {} + const T*, T*, const uint32_t*, const uint8_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 bool, const int) {} template __global__ void moe_gemm_mxfp4_wmma_kernel_rdna3( - const T*, T*, const uint32_t*, const uint8_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 bool, - const int) {} + const T*, T*, const uint32_t*, const uint8_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 bool, const int) {} #endif // WMMA launcher: 16x16 tiles, single wave, no K-split (enough blocks at the @@ -442,17 +441,13 @@ void launch_moe_gemm_mxfp4_wmma( } template -void launch_moe_gemm_mxfp4(const T* a, T* c, const uint32_t* b_q_weight, - const uint8_t* b_scales_e8m0, - 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, - bool mul_topk_weight, int output_topk, - cudaStream_t stream) { +void launch_moe_gemm_mxfp4( + const T* a, T* c, const uint32_t* b_q_weight, const uint8_t* b_scales_e8m0, + 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, + bool mul_topk_weight, int output_topk, cudaStream_t stream) { dim3 block(MOE_MXFP4_THREADS); dim3 grid(num_token_blocks, (size_n + MOE_MXFP4_BLOCK_KN * 4 - 1) / (MOE_MXFP4_BLOCK_KN * 4), @@ -500,7 +495,8 @@ void dispatch_moe_gemm_mxfp4( mul_topk_weight, output_topk, stream); break; default: - TORCH_CHECK(false, "moe_mxfp4_gemm_rdna3: block_size_m must be 1/2/4/8/16"); + TORCH_CHECK(false, + "moe_mxfp4_gemm_rdna3: block_size_m must be 1/2/4/8/16"); } } diff --git a/csrc/rocm/mxfp4_gemm_rdna3.cu b/csrc/rocm/mxfp4_gemm_rdna3.cu index 688946c2534c..6d9048479995 100644 --- a/csrc/rocm/mxfp4_gemm_rdna3.cu +++ b/csrc/rocm/mxfp4_gemm_rdna3.cu @@ -126,7 +126,7 @@ __global__ void gemm_mxfp4_wmma_kernel_16x16_1w( v8fp32 c_acc = {0, 0, 0, 0, 0, 0, 0, 0}; - const int groupsize = size_k / groups; // == 32 for MXFP4 + const int groupsize = size_k / groups; // == 32 for MXFP4 constexpr uint32_t mant_bits = std::is_same::value ? 10u : 7u; const int k_per_split = size_k / gridDim.z; @@ -243,7 +243,8 @@ void launch_gemm_mxfp4_wmma_16x16_1w(const T* a, const uint32_t* b_q, int k_split = compute_mxfp4_k_split(size_k); // Each K-segment must be a whole number of 16-wide tiles, else the k_tile // loop would drop the tail of K. Fall back to no split otherwise. - while (k_split > 1 && (size_k % k_split != 0 || (size_k / k_split) % 16 != 0)) { + while (k_split > 1 && + (size_k % k_split != 0 || (size_k / k_split) % 16 != 0)) { k_split /= 2; } dim3 block(32); @@ -508,7 +509,8 @@ __global__ void gemm_mxfp4_wmma_kernel_64x64_4w( store_acc(c_acc3, n_tile + 48); } -// ---- 128x64_k32: 8 waves, 128M x 64N, K=32/iter, cached group bias (M >= 128) ---- +// ---- 128x64_k32: 8 waves, 128M x 64N, K=32/iter, cached group bias (M >= 128) +// ---- template __global__ void gemm_mxfp4_wmma_kernel_128x64_k32( const T* __restrict__ a, const uint32_t* __restrict__ b_q, @@ -725,8 +727,9 @@ void launch_gemm_mxfp4_wmma(const T* a, const uint32_t* b_q, } // Scalar GEMV for decode (small M): WMMA wastes 15/16 of its compute at M=1, -// so this reads each weight column once and does M scalar dots (bandwidth-bound, -// the 4-bit read wins). 256 threads, 4 N cols each; gridDim.z splits K (atomic). +// so this reads each weight column once and does M scalar dots +// (bandwidth-bound, the 4-bit read wins). 256 threads, 4 N cols each; gridDim.z +// splits K (atomic). #define MXFP4_BLOCK_KN 256 #define MXFP4_THREADS 256 @@ -819,8 +822,8 @@ __global__ void gemm_mxfp4_scalar_rdna3( auto refresh = [&](int g) { #pragma unroll for (int col = 0; col < 4; ++col) - gscale[col] = __uint_as_float((uint32_t)b_scales_e8m0[g * size_n + n + col] - << 23); + gscale[col] = + __uint_as_float((uint32_t)b_scales_e8m0[g * size_n + n + col] << 23); }; refresh(group); @@ -945,7 +948,8 @@ torch::Tensor mxfp4_gemm_rdna3(torch::Tensor a, torch::Tensor b_q_weight, torch::Tensor b_scales_e8m0) { TORCH_CHECK(a.is_cuda(), "a must be a CUDA/HIP tensor"); TORCH_CHECK(b_q_weight.is_cuda(), "b_q_weight must be a CUDA/HIP tensor"); - TORCH_CHECK(b_scales_e8m0.is_cuda(), "b_scales_e8m0 must be a CUDA/HIP tensor"); + TORCH_CHECK(b_scales_e8m0.is_cuda(), + "b_scales_e8m0 must be a CUDA/HIP tensor"); TORCH_CHECK(a.dim() == 2, "a must be 2D [M, K]"); TORCH_CHECK(b_q_weight.dim() == 2, "b_q_weight must be 2D [K/8, N]"); TORCH_CHECK( @@ -967,7 +971,8 @@ torch::Tensor mxfp4_gemm_rdna3(torch::Tensor a, torch::Tensor b_q_weight, TORCH_CHECK(b_q_weight.size(0) * 8 == size_k, "b_q_weight first dim must be K/8"); - TORCH_CHECK(b_scales_e8m0.size(1) == size_n, "b_scales_e8m0 last dim must be N"); + TORCH_CHECK(b_scales_e8m0.size(1) == size_n, + "b_scales_e8m0 last dim must be N"); TORCH_CHECK(groups * 32 == size_k, "E8M0 group count must be K/32"); TORCH_CHECK(size_n % 16 == 0, "WMMA path requires N % 16 == 0"); TORCH_CHECK(size_k % 32 == 0, "MXFP4 path requires K % 32 == 0"); diff --git a/csrc/rocm/qdq_mxfp4_rdna3.cuh b/csrc/rocm/qdq_mxfp4_rdna3.cuh index 4e2d47676081..8595631ee3cf 100644 --- a/csrc/rocm/qdq_mxfp4_rdna3.cuh +++ b/csrc/rocm/qdq_mxfp4_rdna3.cuh @@ -15,12 +15,12 @@ // Device-only HIP dependencies are gated so the pure-integer decode below can // be #included and unit-tested by a plain host compiler. #if defined(__HIPCC__) -#include -#include -#define MXFP4_HD __host__ __device__ __forceinline__ -#define MXFP4_D __device__ __forceinline__ + #include + #include + #define MXFP4_HD __host__ __device__ __forceinline__ + #define MXFP4_D __device__ __forceinline__ #else -#define MXFP4_HD inline + #define MXFP4_HD inline #endif namespace vllm { @@ -56,8 +56,7 @@ MXFP4_HD uint16_t mxfp4_e2m1_to_fp16_bits(uint32_t nib) { // Apply the E8M0 scale as an exponent add (zero stays zero). MXFP4_HD uint16_t mxfp4_apply_e8m0_bits(uint16_t bits, int32_t bias_u16) { - return (bits & 0x7FFFu) == 0u ? bits - : (uint16_t)((int32_t)bits + bias_u16); + return (bits & 0x7FFFu) == 0u ? bits : (uint16_t)((int32_t)bits + bias_u16); } // E8M0 byte -> additive exponent bias. mant_bits = 7 (bf16) or 10 (fp16). @@ -88,7 +87,7 @@ MXFP4_D half mxfp4_nib_to_fp16(uint32_t nib, int32_t bias_u16) { // = even K). All 8 share one E8M0 block since K is tiled in multiples of 32. MXFP4_D void dequant_mxfp4_8_bf16(uint32_t qa, int32_t bias_u16, bf16_t (&dq)[8]) { -#pragma unroll + #pragma unroll for (int i = 0; i < 8; i++) { dq[i] = mxfp4_nib_to_bf16((qa >> (4 * i)) & 0xFu, bias_u16); } @@ -96,7 +95,7 @@ MXFP4_D void dequant_mxfp4_8_bf16(uint32_t qa, int32_t bias_u16, MXFP4_D void dequant_mxfp4_8_fp16(uint32_t qa, int32_t bias_u16, half (&dq)[8]) { -#pragma unroll + #pragma unroll for (int i = 0; i < 8; i++) { dq[i] = mxfp4_nib_to_fp16((qa >> (4 * i)) & 0xFu, bias_u16); } @@ -108,7 +107,7 @@ MXFP4_D void dequant_mxfp4_8_fp16(uint32_t qa, int32_t bias_u16, template MXFP4_D void dequant_mxfp4_8_lut(uint32_t qa, int32_t bias_u16, const uint16_t* lut, T (&dq)[8]) { -#pragma unroll + #pragma unroll for (int i = 0; i < 8; i++) { uint16_t b = mxfp4_apply_e8m0_bits(lut[(qa >> (4 * i)) & 0xFu], bias_u16); __builtin_memcpy(&dq[i], &b, sizeof(T)); @@ -118,7 +117,7 @@ MXFP4_D void dequant_mxfp4_8_lut(uint32_t qa, int32_t bias_u16, // fp32 output (bf16 bits widened by a free <<16) for a scalar dot path. MXFP4_D void dequant_mxfp4_8_f32(uint32_t qa, int32_t bias_u16, float (&dq)[8]) { -#pragma unroll + #pragma unroll for (int i = 0; i < 8; i++) { uint16_t b = mxfp4_apply_e8m0_bits( mxfp4_e2m1_to_bf16_bits((qa >> (4 * i)) & 0xFu), bias_u16); diff --git a/tests/kernels/quantization/test_mxfp4_rdna3.py b/tests/kernels/quantization/test_mxfp4_rdna3.py index 764e0b3f8e0f..acfc8e6d94e7 100644 --- a/tests/kernels/quantization/test_mxfp4_rdna3.py +++ b/tests/kernels/quantization/test_mxfp4_rdna3.py @@ -11,9 +11,7 @@ from vllm.platforms import current_platform -E2M1_MAG = torch.tensor( - [0.0, 0.5, 1.0, 1.5, 2.0, 3.0, 4.0, 6.0], dtype=torch.float32 -) +E2M1_MAG = torch.tensor([0.0, 0.5, 1.0, 1.5, 2.0, 3.0, 4.0, 6.0], dtype=torch.float32) pytestmark = pytest.mark.skipif( not current_platform.is_rocm(), reason="RDNA3 MXFP4 WMMA kernel is ROCm-only" @@ -52,9 +50,7 @@ def test_mxfp4_wmma_matches_dequant_reference(dtype, M, N, K): # Random E2M1 codes (full 0..255 byte = two 4-bit codes) and mild E8M0 # scales centered on 127 (== 2^0) to stay well inside fp16/bf16 range. w_packed = torch.randint(0, 256, (N, K // 2), dtype=torch.uint8, device=dev) - scale_e8m0 = torch.randint( - 124, 131, (N, K // 32), dtype=torch.uint8, device=dev - ) + scale_e8m0 = torch.randint(124, 131, (N, K // 32), dtype=torch.uint8, device=dev) x = torch.randn(M, K, dtype=dtype, device=dev) * 0.1 w_deq = _dequant_ref(w_packed, scale_e8m0) # [N, K] fp32 diff --git a/vllm/model_executor/kernels/linear/mxfp4/rocm.py b/vllm/model_executor/kernels/linear/mxfp4/rocm.py index 6fc55dbabc7c..c234bd275e41 100644 --- a/vllm/model_executor/kernels/linear/mxfp4/rocm.py +++ b/vllm/model_executor/kernels/linear/mxfp4/rocm.py @@ -11,9 +11,7 @@ _MXFP4_GROUP_SIZE = 32 # Fake/meta impl so the op is traceable under torch.compile + CUDA graphs. -if hasattr(torch.ops, "_rocm_C") and hasattr( - torch.ops._rocm_C, "mxfp4_gemm_rdna3" -): +if hasattr(torch.ops, "_rocm_C") and hasattr(torch.ops._rocm_C, "mxfp4_gemm_rdna3"): try: @torch.library.register_fake("_rocm_C::mxfp4_gemm_rdna3") @@ -58,8 +56,7 @@ def process_weights_after_loading(self, layer: torch.nn.Module) -> None: if N % 16 != 0 or K % _MXFP4_GROUP_SIZE != 0: raise ValueError( - f"RDNA3 MXFP4 WMMA kernel needs N%16==0 and K%32==0, got " - f"N={N}, K={K}" + f"RDNA3 MXFP4 WMMA kernel needs N%16==0 and K%32==0, got N={N}, K={K}" ) # [N, K/2] uint8 reinterpreted little-endian as int32 is already the @@ -81,9 +78,7 @@ def apply_weights( if not x_2d.is_contiguous(): x_2d = x_2d.contiguous() - out = torch.ops._rocm_C.mxfp4_gemm_rdna3( - x_2d, layer.weight, layer.weight_scale - ) + out = torch.ops._rocm_C.mxfp4_gemm_rdna3(x_2d, layer.weight, layer.weight_scale) if bias is not None: out = out + bias return out.view(out_shape) diff --git a/vllm/model_executor/layers/quantization/compressed_tensors/compressed_tensors_moe/compressed_tensors_moe_w4a4_mxfp4_rdna3.py b/vllm/model_executor/layers/quantization/compressed_tensors/compressed_tensors_moe/compressed_tensors_moe_w4a4_mxfp4_rdna3.py index 21b54d8e0e66..dfb5df04a123 100644 --- a/vllm/model_executor/layers/quantization/compressed_tensors/compressed_tensors_moe/compressed_tensors_moe_w4a4_mxfp4_rdna3.py +++ b/vllm/model_executor/layers/quantization/compressed_tensors/compressed_tensors_moe/compressed_tensors_moe_w4a4_mxfp4_rdna3.py @@ -42,12 +42,34 @@ def process_weights_after_loading(self, layer: torch.nn.Module) -> None: def _gemm_w13(self, layer, a, c, tw, sti, eid, ntp, top_k, block_size_m, mul_tw): ops.moe_mxfp4_gemm_rdna3( - a, c, layer.w13_weight_packed, layer.w13_weight_scale, tw, sti, eid, - ntp, top_k, block_size_m, mul_tw, 0, + a, + c, + layer.w13_weight_packed, + layer.w13_weight_scale, + tw, + sti, + eid, + ntp, + top_k, + block_size_m, + mul_tw, + 0, ) - def _gemm_w2(self, layer, a, c, tw, sti, eid, ntp, block_size_m, mul_tw, output_topk): + def _gemm_w2( + self, layer, a, c, tw, sti, eid, ntp, block_size_m, mul_tw, output_topk + ): ops.moe_mxfp4_gemm_rdna3( - a, c, layer.w2_weight_packed, layer.w2_weight_scale, tw, sti, eid, - ntp, 1, block_size_m, mul_tw, output_topk, + a, + c, + layer.w2_weight_packed, + layer.w2_weight_scale, + tw, + sti, + eid, + ntp, + 1, + block_size_m, + mul_tw, + output_topk, ) diff --git a/vllm/model_executor/layers/quantization/compressed_tensors/compressed_tensors_moe/compressed_tensors_moe_wna16_rdna3.py b/vllm/model_executor/layers/quantization/compressed_tensors/compressed_tensors_moe/compressed_tensors_moe_wna16_rdna3.py index d28b103a30e2..9f6e44507228 100644 --- a/vllm/model_executor/layers/quantization/compressed_tensors/compressed_tensors_moe/compressed_tensors_moe_wna16_rdna3.py +++ b/vllm/model_executor/layers/quantization/compressed_tensors/compressed_tensors_moe/compressed_tensors_moe_wna16_rdna3.py @@ -116,16 +116,37 @@ def process_weights_after_loading(self, layer: torch.nn.Module) -> None: ) layer.rdna3_empty_tw = torch.empty(0, device=device) - def _gemm_w13(self, layer, a, c, tw, sti, eid, ntp, top_k, block_size_m, mul_tw): ops.moe_gptq_gemm_rdna3( - a, c, layer.w13_weight_packed, layer.w13_weight_scale, - layer.w13_qzeros, tw, sti, eid, ntp, top_k, block_size_m, mul_tw, + a, + c, + layer.w13_weight_packed, + layer.w13_weight_scale, + layer.w13_qzeros, + tw, + sti, + eid, + ntp, + top_k, + block_size_m, + mul_tw, ) - def _gemm_w2(self, layer, a, c, tw, sti, eid, ntp, block_size_m, mul_tw, output_topk): + def _gemm_w2( + self, layer, a, c, tw, sti, eid, ntp, block_size_m, mul_tw, output_topk + ): ops.moe_gptq_gemm_rdna3( - a, c, layer.w2_weight_packed, layer.w2_weight_scale, - layer.w2_qzeros, tw, sti, eid, ntp, 1, block_size_m, mul_tw, + a, + c, + layer.w2_weight_packed, + layer.w2_weight_scale, + layer.w2_qzeros, + tw, + sti, + eid, + ntp, + 1, + block_size_m, + mul_tw, output_topk, ) diff --git a/vllm/model_executor/layers/quantization/compressed_tensors/compressed_tensors_moe/rdna3_moe_common.py b/vllm/model_executor/layers/quantization/compressed_tensors/compressed_tensors_moe/rdna3_moe_common.py index b86add850b65..f7d098d84b55 100644 --- a/vllm/model_executor/layers/quantization/compressed_tensors/compressed_tensors_moe/rdna3_moe_common.py +++ b/vllm/model_executor/layers/quantization/compressed_tensors/compressed_tensors_moe/rdna3_moe_common.py @@ -69,7 +69,9 @@ class RDNA3FusedMoEMixin: def _gemm_w13(self, layer, a, c, tw, sti, eid, ntp, top_k, block_size_m, mul_tw): raise NotImplementedError - def _gemm_w2(self, layer, a, c, tw, sti, eid, ntp, block_size_m, mul_tw, output_topk): + def _gemm_w2( + self, layer, a, c, tw, sti, eid, ntp, block_size_m, mul_tw, output_topk + ): raise NotImplementedError def _rdna3_run( diff --git a/vllm/model_executor/layers/quantization/gpt_oss_mxfp4_rdna3.py b/vllm/model_executor/layers/quantization/gpt_oss_mxfp4_rdna3.py index 1413ae96884d..59fef5817850 100644 --- a/vllm/model_executor/layers/quantization/gpt_oss_mxfp4_rdna3.py +++ b/vllm/model_executor/layers/quantization/gpt_oss_mxfp4_rdna3.py @@ -136,8 +136,18 @@ def apply( # gate_up GEMM (no reduction) + per-expert bias w1 = torch.zeros(total, N_gate_up, dtype=dtype, device=device) ops.moe_mxfp4_gemm_rdna3( - x, w1, layer.w13_weight_packed, layer.w13_weight_scale, empty_tw, - sti, eid, ntp, top_k, block_size_m, False, 0, + x, + w1, + layer.w13_weight_packed, + layer.w13_weight_scale, + empty_tw, + sti, + eid, + ntp, + top_k, + block_size_m, + False, + 0, ) w1 = w1 + layer.w13_bias[flat_experts].to(dtype) act = self._swiglu_oai(w1) @@ -145,8 +155,18 @@ def apply( # down GEMM (no reduction) + per-expert bias, then weighted moe_sum w2o = torch.zeros(total, hidden, dtype=dtype, device=device) ops.moe_mxfp4_gemm_rdna3( - act, w2o, layer.w2_weight_packed, layer.w2_weight_scale, empty_tw, - sti, eid, ntp, 1, block_size_m, False, 0, + act, + w2o, + layer.w2_weight_packed, + layer.w2_weight_scale, + empty_tw, + sti, + eid, + ntp, + 1, + block_size_m, + False, + 0, ) w2o = w2o + layer.w2_bias[flat_experts].to(dtype) w2o = w2o * topk_weights.reshape(-1, 1).to(dtype) From 6b0a024b293577f4482856c876e02d18e7085df7 Mon Sep 17 00:00:00 2001 From: JartX Date: Thu, 25 Jun 2026 19:05:45 +0200 Subject: [PATCH 13/36] [ROCm][gfx1100] MXFP4 MoE: rename tpe -> tokens_per_expert (typos hook) The typos pre-commit hook flags 'tpe' as a misspelling of 'type'. Spell it out; clearer anyway. No functional change. Signed-off-by: JartX --- .../compressed_tensors_moe/rdna3_moe_common.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/vllm/model_executor/layers/quantization/compressed_tensors/compressed_tensors_moe/rdna3_moe_common.py b/vllm/model_executor/layers/quantization/compressed_tensors/compressed_tensors_moe/rdna3_moe_common.py index f7d098d84b55..2f532cf56345 100644 --- a/vllm/model_executor/layers/quantization/compressed_tensors/compressed_tensors_moe/rdna3_moe_common.py +++ b/vllm/model_executor/layers/quantization/compressed_tensors/compressed_tensors_moe/rdna3_moe_common.py @@ -34,8 +34,8 @@ def select_block_size_m(num_tokens: int, top_k: int, num_experts: int) -> int: return 1 if num_tokens >= 16: return 16 - tpe = num_tokens * top_k / num_experts - return 4 if tpe >= 1.5 else 1 + tokens_per_expert = num_tokens * top_k / num_experts + return 4 if tokens_per_expert >= 1.5 else 1 def repack_experts(packed: torch.Tensor, scale: torch.Tensor, deinterleave: bool): From faf7db9224ae70d1c96d3515e946c5d799acccef Mon Sep 17 00:00:00 2001 From: JartX Date: Thu, 25 Jun 2026 23:02:53 +0200 Subject: [PATCH 14/36] [ROCm][gfx1100] MXFP4 dense: gate linear kernel on gfx1100, not gfx11xx MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Rdna3MxFp4LinearKernel.is_supported() used on_gfx11() (substring-matches every gfx11xx: gfx1101/1102/1150/...), but the HIP kernel is compiled only under __gfx1100__ — other gfx11xx device passes get empty stub kernels. On a multi-arch wheel that includes gfx1100, a non-gfx1100 RDNA3 card would select this kernel, run the stub, and return the zero-initialized output: silent all-zero linear outputs, no crash. Gate on on_gfx1100() to match the kernel's compile gate and every other RDNA3 selector in this series. Reported in review (Claude). Signed-off-by: JartX --- vllm/model_executor/kernels/linear/mxfp4/rocm.py | 6 +++--- 1 file changed, 3 insertions(+), 3 deletions(-) diff --git a/vllm/model_executor/kernels/linear/mxfp4/rocm.py b/vllm/model_executor/kernels/linear/mxfp4/rocm.py index c234bd275e41..200b8f92d81d 100644 --- a/vllm/model_executor/kernels/linear/mxfp4/rocm.py +++ b/vllm/model_executor/kernels/linear/mxfp4/rocm.py @@ -34,10 +34,10 @@ def is_supported( ) -> tuple[bool, str | None]: if not current_platform.is_rocm(): return False, "not running on ROCm" - from vllm.platforms.rocm import on_gfx11 + from vllm.platforms.rocm import on_gfx1100 - if not on_gfx11(): - return False, "RDNA3 (gfx11xx) WMMA required" + if not on_gfx1100(): + return False, "RDNA3 gfx1100 WMMA required" if not hasattr(torch.ops._rocm_C, "mxfp4_gemm_rdna3"): return False, "_rocm_C.mxfp4_gemm_rdna3 op not built" return True, None From df10623adcb96a63e51d56778cb38e29f441ca34 Mon Sep 17 00:00:00 2001 From: JartX Date: Thu, 25 Jun 2026 23:11:34 +0200 Subject: [PATCH 15/36] [ROCm][gfx1100] MXFP4: don't route W4A16 MoE through the WMMA-16 tile The shared RDNA3FusedMoEMixin is inherited by both the MXFP4 MoE method and the existing W4A16 (moe_gptq_gemm_rdna3) method. select_block_size_m returns 16 for >=16 tokens, but moe_gptq_gemm_rdna3 only has tiles {1,2,4,8} and TORCH_CHECKs the default arm, so every W4A16 prefill / batched-decode >=16 crashed. Make tile selection an overridable mixin method (_select_block_size_m) and override it in the W4A16 method to its proven 1/4 cap; the MXFP4 paths keep the 16-tile. Reported in review (Claude). Verified the W4A16 override caps at {1,4} for all batch sizes while the MXFP4 default still returns 16. Signed-off-by: JartX --- .../compressed_tensors_moe_wna16_rdna3.py | 6 ++++++ .../compressed_tensors_moe/rdna3_moe_common.py | 12 +++++++++++- 2 files changed, 17 insertions(+), 1 deletion(-) diff --git a/vllm/model_executor/layers/quantization/compressed_tensors/compressed_tensors_moe/compressed_tensors_moe_wna16_rdna3.py b/vllm/model_executor/layers/quantization/compressed_tensors/compressed_tensors_moe/compressed_tensors_moe_wna16_rdna3.py index 9f6e44507228..a24800df5f6f 100644 --- a/vllm/model_executor/layers/quantization/compressed_tensors/compressed_tensors_moe/compressed_tensors_moe_wna16_rdna3.py +++ b/vllm/model_executor/layers/quantization/compressed_tensors/compressed_tensors_moe/compressed_tensors_moe_wna16_rdna3.py @@ -116,6 +116,12 @@ def process_weights_after_loading(self, layer: torch.nn.Module) -> None: ) layer.rdna3_empty_tw = torch.empty(0, device=device) + def _select_block_size_m(self, num_tokens, top_k, num_experts): + # moe_gptq_gemm_rdna3 only has tiles {1,2,4,8} (no WMMA-16 path); keep + # the proven W4A16 cap so prefill / batched decode don't hit the + # kernel's TORCH_CHECK default arm. + return 1 if num_tokens <= 4 else 4 + def _gemm_w13(self, layer, a, c, tw, sti, eid, ntp, top_k, block_size_m, mul_tw): ops.moe_gptq_gemm_rdna3( a, diff --git a/vllm/model_executor/layers/quantization/compressed_tensors/compressed_tensors_moe/rdna3_moe_common.py b/vllm/model_executor/layers/quantization/compressed_tensors/compressed_tensors_moe/rdna3_moe_common.py index 2f532cf56345..177a6128ec22 100644 --- a/vllm/model_executor/layers/quantization/compressed_tensors/compressed_tensors_moe/rdna3_moe_common.py +++ b/vllm/model_executor/layers/quantization/compressed_tensors/compressed_tensors_moe/rdna3_moe_common.py @@ -66,6 +66,14 @@ class RDNA3FusedMoEMixin: """Mixin: shared apply() + two-GEMM orchestration. Subclasses implement ``_gemm_w13`` and ``_gemm_w2`` (the op-specific HIP launch).""" + def _select_block_size_m( + self, num_tokens: int, top_k: int, num_experts: int + ) -> int: + """Tile size for the GEMM op. The default targets the MXFP4 kernel + (tiles 1/2/4/8/16); subclasses whose kernel supports a smaller set + (e.g. the W4A16 ``moe_gptq_gemm_rdna3``: 1/2/4/8) must override.""" + return select_block_size_m(num_tokens, top_k, num_experts) + def _gemm_w13(self, layer, a, c, tw, sti, eid, ntp, top_k, block_size_m, mul_tw): raise NotImplementedError @@ -96,7 +104,9 @@ def _rdna3_run( if global_num_experts <= 0: global_num_experts = layer.w13_weight_packed.shape[0] - block_size_m = select_block_size_m(num_tokens, top_k, global_num_experts) + block_size_m = self._select_block_size_m( + num_tokens, top_k, global_num_experts + ) sorted_token_ids, expert_ids, num_tokens_post_padded = moe_align_block_size( topk_ids, block_size_m, global_num_experts, expert_map From bfe3c81b062f36b95e2bbb0bb3a7889a5efc17cb Mon Sep 17 00:00:00 2001 From: JartX Date: Thu, 25 Jun 2026 23:17:09 +0200 Subject: [PATCH 16/36] [ROCm][gfx1100] MXFP4: test the W4A16 MoE tile-cap regression Pure-Python dispatch tests (no GPU) for the shared RDNA3 MoE block-size selection: the W4A16 method must never select a tile outside the moe_gptq_gemm_rdna3 set {1,2,4,8}, while the MXFP4 default may return 16. Guards the regression where the shared mixin routed W4A16 >=16-token batches through the WMMA-16 tile and crashed the kernel. Verified the test fails (112 cases) when the W4A16 override is removed and passes (418) with it. Signed-off-by: JartX --- .../quantization/test_mxfp4_rdna3_dispatch.py | 74 +++++++++++++++++++ 1 file changed, 74 insertions(+) create mode 100644 tests/kernels/quantization/test_mxfp4_rdna3_dispatch.py diff --git a/tests/kernels/quantization/test_mxfp4_rdna3_dispatch.py b/tests/kernels/quantization/test_mxfp4_rdna3_dispatch.py new file mode 100644 index 000000000000..9cbdd873de9e --- /dev/null +++ b/tests/kernels/quantization/test_mxfp4_rdna3_dispatch.py @@ -0,0 +1,74 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""Block-size dispatch tests for the RDNA3 (gfx1100) fused-MoE methods. + +``RDNA3FusedMoEMixin`` is shared by the MXFP4 MoE method and the existing W4A16 +(``moe_gptq_gemm_rdna3``) method. The MXFP4 kernel has a WMMA-16 tile; the W4A16 +kernel only instantiates tiles {1,2,4,8} and TORCH_CHECKs anything else. So the +W4A16 method must override tile selection to stay inside its set — otherwise any +>=16-token prefill / batched decode hard-crashes. These are pure-Python checks +(no GPU), so the regression is caught on any CI host. +""" + +import pytest + +from vllm.model_executor.layers.quantization.compressed_tensors.compressed_tensors_moe.rdna3_moe_common import ( # noqa: E501 + RDNA3FusedMoEMixin, + select_block_size_m, +) + +# Tiles each HIP kernel actually instantiates (the rest hit a TORCH_CHECK). +GPTQ_TILES = {1, 2, 4, 8} # moe_gptq_gemm_rdna3 (W4A16) +MXFP4_TILES = {1, 2, 4, 8, 16} # moe_mxfp4_gemm_rdna3 + +BATCH_SIZES = [1, 2, 3, 4, 8, 15, 16, 17, 32, 64, 256, 512, 4096] +TOP_KS = [1, 2, 4, 8] +NUM_EXPERTS = [8, 32, 128, 256] + + +def _w4a16_cls(): + from vllm.model_executor.layers.quantization.compressed_tensors.compressed_tensors_moe.compressed_tensors_moe_wna16_rdna3 import ( # noqa: E501 + CompressedTensorsWNA16RDNA3MoEMethod, + ) + + return CompressedTensorsWNA16RDNA3MoEMethod + + +@pytest.mark.parametrize("num_tokens", BATCH_SIZES) +@pytest.mark.parametrize("top_k", TOP_KS) +@pytest.mark.parametrize("num_experts", NUM_EXPERTS) +def test_mxfp4_tile_in_supported_set(num_tokens, top_k, num_experts): + bsm = select_block_size_m(num_tokens, top_k, num_experts) + assert bsm in MXFP4_TILES + # the WMMA-16 tile is only picked once a step can fill it + if num_tokens < 16: + assert bsm != 16 + + +@pytest.mark.parametrize("num_tokens", BATCH_SIZES) +@pytest.mark.parametrize("top_k", TOP_KS) +@pytest.mark.parametrize("num_experts", NUM_EXPERTS) +def test_w4a16_never_selects_unsupported_tile(num_tokens, top_k, num_experts): + """Regression guard: the W4A16 method must never hand moe_gptq_gemm_rdna3 a + tile outside {1,2,4,8} (>=16-token batches used to select 16 and crash).""" + # _select_block_size_m is stateless, so call it without building a method. + bsm = _w4a16_cls()._select_block_size_m(None, num_tokens, top_k, num_experts) + assert bsm in GPTQ_TILES, ( + f"W4A16 selected block_size_m={bsm} for num_tokens={num_tokens}; " + f"moe_gptq_gemm_rdna3 only supports {sorted(GPTQ_TILES)}" + ) + + +def test_w4a16_override_differs_from_mixin_default(): + """The mixin default returns the WMMA-16 tile for a 32-token batch; the + W4A16 override must replace it with a kernel-supported tile.""" + assert RDNA3FusedMoEMixin._select_block_size_m(None, 32, 8, 256) == 16 + w4a16 = _w4a16_cls()._select_block_size_m(None, 32, 8, 256) + assert w4a16 != 16 + assert w4a16 in GPTQ_TILES + + +def test_degenerate_inputs_return_safe_tile(): + # num_experts/top_k <= 0 must not divide-by-zero; fall back to the 1-tile. + assert select_block_size_m(32, 0, 256) == 1 + assert select_block_size_m(32, 8, 0) == 1 From 7fd46d311b8f36626cc47d1bc17268c52745bbce Mon Sep 17 00:00:00 2001 From: JartX Date: Thu, 25 Jun 2026 23:19:37 +0200 Subject: [PATCH 17/36] [ROCm][gfx1100] MXFP4: pre-commit fixups (ruff-format + typos) ruff-format collapse in rdna3_moe_common and reword 'TORCH_CHECKs' -> 'rejects' in the dispatch test so the typos hook stops flagging it. No functional change. Signed-off-by: JartX --- tests/kernels/quantization/test_mxfp4_rdna3_dispatch.py | 2 +- .../compressed_tensors_moe/rdna3_moe_common.py | 4 +--- 2 files changed, 2 insertions(+), 4 deletions(-) diff --git a/tests/kernels/quantization/test_mxfp4_rdna3_dispatch.py b/tests/kernels/quantization/test_mxfp4_rdna3_dispatch.py index 9cbdd873de9e..0c745fba2d2b 100644 --- a/tests/kernels/quantization/test_mxfp4_rdna3_dispatch.py +++ b/tests/kernels/quantization/test_mxfp4_rdna3_dispatch.py @@ -4,7 +4,7 @@ ``RDNA3FusedMoEMixin`` is shared by the MXFP4 MoE method and the existing W4A16 (``moe_gptq_gemm_rdna3``) method. The MXFP4 kernel has a WMMA-16 tile; the W4A16 -kernel only instantiates tiles {1,2,4,8} and TORCH_CHECKs anything else. So the +kernel only instantiates tiles {1,2,4,8} and rejects anything else. So the W4A16 method must override tile selection to stay inside its set — otherwise any >=16-token prefill / batched decode hard-crashes. These are pure-Python checks (no GPU), so the regression is caught on any CI host. diff --git a/vllm/model_executor/layers/quantization/compressed_tensors/compressed_tensors_moe/rdna3_moe_common.py b/vllm/model_executor/layers/quantization/compressed_tensors/compressed_tensors_moe/rdna3_moe_common.py index 177a6128ec22..0db818230664 100644 --- a/vllm/model_executor/layers/quantization/compressed_tensors/compressed_tensors_moe/rdna3_moe_common.py +++ b/vllm/model_executor/layers/quantization/compressed_tensors/compressed_tensors_moe/rdna3_moe_common.py @@ -104,9 +104,7 @@ def _rdna3_run( if global_num_experts <= 0: global_num_experts = layer.w13_weight_packed.shape[0] - block_size_m = self._select_block_size_m( - num_tokens, top_k, global_num_experts - ) + block_size_m = self._select_block_size_m(num_tokens, top_k, global_num_experts) sorted_token_ids, expert_ids, num_tokens_post_padded = moe_align_block_size( topk_ids, block_size_m, global_num_experts, expert_map From d43f1434cc28762afc18da51e7b4b1a30bbb0154 Mon Sep 17 00:00:00 2001 From: JartX Date: Fri, 26 Jun 2026 03:14:34 +0200 Subject: [PATCH 18/36] [ROCm][gfx1100] MXFP4 MoE: register an oracle backend instead of bespoke methods MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Per maintainer review: add the RDNA3 MXFP4 MoE path as a proper Mxfp4MoeBackend.RDNA3_MXFP4 + a modular RDNA3Mxfp4Experts (FusedMoEExpertsModular) and let the oracle dispatch to it, instead of introducing/overriding quantization methods. - New experts/rdna3_mxfp4_moe.py: drives moe_mxfp4_gemm_rdna3; Standard activation format (own moe_align routing), gfx1100-gated via _supports_current_device; one apply() covers compressed-tensors (SiLU, no bias) and GPT-OSS (per-expert bias + clamped SwiGLU-OAI) by reading bias and gemm1_alpha/beta/clamp from the quant_config; allocates its own scratch. - oracle/mxfp4.py: enum value, backend_to_kernel_cls, gfx1100-first priority, W4A16 quant-config branch, and the [K/8,N] repack in convert_gpt_oss_weight_to_mxfp4_moe_kernel_format (w13 gate/up de-interleave). - compressed_tensors_moe_w4a4_mxfp4: pick the backend on gfx1100 + repack in process_weights (no de-interleave). - Drop the bespoke GptOssMxfp4RDNA3MoEMethod, the get_quant_method hook, the rocm_moe_rdna MXFP4 selector, and CompressedTensorsW4A4Mxfp4RDNA3MoEMethod. (The W4A16 INT4 MoE path keeps its existing method — different quant.) Validated e2e on gfx1100: GPT-OSS-20B and Qwen3.6-35B-A3B-MXFP4 both serve via 'RDNA3_MXFP4' with correct output (17*23=391; 25*4=100; Tokyo; Rayleigh). Signed-off-by: JartX --- .../fused_moe/experts/rdna3_mxfp4_moe.py | 232 ++++++++++++++++++ .../layers/fused_moe/oracle/mxfp4.py | 33 ++- .../compressed_tensors_moe.py | 5 - .../compressed_tensors_moe_w4a4_mxfp4.py | 21 ++ ...compressed_tensors_moe_w4a4_mxfp4_rdna3.py | 75 ------ .../compressed_tensors_moe/rocm_moe_rdna.py | 23 -- .../quantization/gpt_oss_mxfp4_rdna3.py | 173 ------------- .../layers/quantization/mxfp4.py | 8 +- 8 files changed, 286 insertions(+), 284 deletions(-) create mode 100644 vllm/model_executor/layers/fused_moe/experts/rdna3_mxfp4_moe.py delete mode 100644 vllm/model_executor/layers/quantization/compressed_tensors/compressed_tensors_moe/compressed_tensors_moe_w4a4_mxfp4_rdna3.py delete mode 100644 vllm/model_executor/layers/quantization/gpt_oss_mxfp4_rdna3.py diff --git a/vllm/model_executor/layers/fused_moe/experts/rdna3_mxfp4_moe.py b/vllm/model_executor/layers/fused_moe/experts/rdna3_mxfp4_moe.py new file mode 100644 index 000000000000..065cf98c0a43 --- /dev/null +++ b/vllm/model_executor/layers/fused_moe/experts/rdna3_mxfp4_moe.py @@ -0,0 +1,232 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""Native RDNA3 (gfx1100) MXFP4 fused-MoE experts backend. + +A modular ``FusedMoEExperts`` that drives the ``moe_mxfp4_gemm_rdna3`` HIP kernel +(E2M1 weights + E8M0 group-32 scale, no zero point). Standard activation format: +``apply`` receives the full hidden states + topk ids/weights and does its own +expert routing (moe_align_block_size), the two GEMMs, the activation, the +per-expert bias and the weighted reduction — so it covers both the +compressed-tensors MXFP4 MoE (no bias, SiLU) and GPT-OSS native-mxfp4 (per-expert +bias + clamped SwiGLU-OAI) by reading everything from ``quant_config``. + +Registered as ``Mxfp4MoeBackend.RDNA3_MXFP4`` in the oracle and gated on gfx1100. +""" + +import torch + +import vllm.model_executor.layers.fused_moe.modular_kernel as mk +from vllm import _custom_ops as ops +from vllm.model_executor.layers.fused_moe.activation import ( + MoEActivation, + apply_moe_activation, +) +from vllm.model_executor.layers.fused_moe.config import ( + FusedMoEParallelConfig, +) +from vllm.model_executor.layers.fused_moe.moe_align_block_size import ( + moe_align_block_size, +) +from vllm.model_executor.layers.fused_moe.topk_weight_and_reduce import ( + TopKWeightAndReduceNoOP, +) +from vllm.model_executor.layers.quantization.utils.quant_utils import ( + QuantKey, + kMxfp4Static, +) + + +def _select_block_size_m(num_tokens: int, top_k: int, num_experts: int) -> int: + """Pick the GEMM tile from expected occupancy (avg tokens/expert). + + >=16-row steps fill the WMMA-16 tile (clustered routing fills it even on + 256-expert models; also where prefill lives); sub-16 batches use a scalar + tile sized to occupancy (bsm=1 beats bsm=4 at low concurrency). + """ + if num_experts <= 0 or top_k <= 0: + return 1 + if num_tokens >= 16: + return 16 + tokens_per_expert = num_tokens * top_k / num_experts + return 4 if tokens_per_expert >= 1.5 else 1 + + +def _rocm_op_available() -> bool: + return hasattr(torch.ops, "_rocm_C") and hasattr( + torch.ops._rocm_C, "moe_mxfp4_gemm_rdna3" + ) + + +def repack_experts_rdna3( + packed: torch.Tensor, scale: torch.Tensor, deinterleave: bool +) -> tuple[torch.Tensor, torch.Tensor]: + """[E, N, K/2] u8 + [E, N, K/32] u8 -> [E, K/8, N] int32 + [E, K/32, N] u8. + + The kernel reads weights as [E, K/8, N] uint32 (4 consecutive packed bytes + little-endian = a uint32 of 8 E2M1 codes along K). Per-expert into + pre-allocated outputs so a multi-GB repack fits next to the loaded weights. + ``deinterleave`` splits GPT-OSS's gate/up rows ([g0,u0,...] -> [g...,u...]). + """ + E, N, kh = packed.shape + dev = packed.device + b_q = torch.empty(E, kh // 4, N, dtype=torch.int32, device=dev) + b_scale = torch.empty(E, scale.shape[2], N, dtype=torch.uint8, device=dev) + for e in range(E): + pe = packed[e] + se = scale[e] + if deinterleave: + pe = torch.cat([pe[::2], pe[1::2]], dim=0) + se = torch.cat([se[::2], se[1::2]], dim=0) + b_q[e] = pe.contiguous().view(torch.int32).t() + b_scale[e] = se.t() + return b_q, b_scale + + +class RDNA3Mxfp4Experts(mk.FusedMoEExpertsModular): + """MXFP4 MoE experts on the native RDNA3 HIP kernel (gfx1100 only).""" + + @staticmethod + def activation_format() -> mk.FusedMoEActivationFormat: + return mk.FusedMoEActivationFormat.Standard + + @staticmethod + def _supports_current_device() -> bool: + from vllm.platforms import current_platform + + if not current_platform.is_rocm() or not _rocm_op_available(): + return False + from vllm.platforms.rocm import on_gfx1100 + + return on_gfx1100() + + @staticmethod + def _supports_no_act_and_mul() -> bool: + return False # gated MLP only (gate_up -> act -> down) + + @staticmethod + def _supports_quant_scheme( + weight_key: QuantKey | None, + activation_key: QuantKey | None, + ) -> bool: + # MXFP4 weight-only (W4A16); no activation quantization. + return (weight_key, activation_key) == (kMxfp4Static, None) + + @staticmethod + def _supports_activation(activation: MoEActivation) -> bool: + return activation in (MoEActivation.SILU, MoEActivation.SWIGLUOAI) + + @staticmethod + def _supports_parallel_config(moe_parallel_config: FusedMoEParallelConfig) -> bool: + return True + + def finalize_weight_and_reduce_impl(self) -> mk.TopKWeightAndReduce: + # The kernel/epilogue applies the topk weights and reduces, so the + # prepare/finalize step must not reduce again. + return TopKWeightAndReduceNoOP() + + def workspace_shapes( + self, + M: int, + N: int, + K: int, + topk: int, + global_num_experts: int, + local_num_experts: int, + expert_tokens_meta: mk.ExpertTokensMetadata | None, + activation: MoEActivation, + ) -> tuple[tuple[int, ...], tuple[int, ...], tuple[int, ...]]: + # apply() allocates its own scratch (shapes derived from the weights), + # so request no framework workspaces; output is the reduced [M, K]. + return ((0,), (0,), (M, K)) + + def apply( + self, + output: torch.Tensor, + hidden_states: torch.Tensor, + w1: torch.Tensor, + w2: torch.Tensor, + topk_weights: torch.Tensor, + topk_ids: torch.Tensor, + activation: MoEActivation, + global_num_experts: int, + expert_map: torch.Tensor | None, + a1q_scale: torch.Tensor | None, + a2_scale: torch.Tensor | None, + workspace13: torch.Tensor, + workspace2: torch.Tensor, + expert_tokens_meta: mk.ExpertTokensMetadata | None, + apply_router_weight_on_input: bool, + ) -> None: + qc = self.quant_config + M = hidden_states.shape[0] + top_k = topk_ids.shape[1] + total = M * top_k + n_gate_up = w1.shape[2] # [E, K/8, 2*intermediate] + intermediate = n_gate_up // 2 + hidden = w2.shape[2] + dtype = hidden_states.dtype + device = hidden_states.device + + gne = global_num_experts if global_num_experts > 0 else w1.shape[0] + block_size_m = _select_block_size_m(M, top_k, gne) + sti, eid, ntp = moe_align_block_size(topk_ids, block_size_m, gne, expert_map) + + topk_w_f32 = topk_weights.reshape(-1).float() + empty_tw = torch.empty(0, device=device) + w1_bias = qc.w1_bias + w2_bias = qc.w2_bias + flat_experts = topk_ids.reshape(-1).long() if ( + w1_bias is not None or w2_bias is not None + ) else None + + # gate_up GEMM (no reduction); router weight optionally folded in here. + w1_out = torch.zeros(total, n_gate_up, dtype=dtype, device=device) + ops.moe_mxfp4_gemm_rdna3( + hidden_states, w1_out, w1, qc.w1_scale, + topk_w_f32 if apply_router_weight_on_input else empty_tw, + sti, eid, ntp, top_k, block_size_m, apply_router_weight_on_input, 0, + ) + if w1_bias is not None: + w1_out = w1_out + w1_bias[flat_experts].to(dtype) + + act_out = torch.empty(total, intermediate, dtype=dtype, device=device) + if activation == MoEActivation.SWIGLUOAI: + self._swiglu_oai(w1_out, act_out, qc) + else: + apply_moe_activation(activation, act_out, w1_out) + + # down GEMM; if there is a per-expert bias we must add it before the + # weighted reduction, so reduce in Python; otherwise fuse the reduction + # (output_topk) and the topk-weight multiply into the kernel epilogue. + if w2_bias is not None: + w2_out = torch.zeros(total, hidden, dtype=dtype, device=device) + ops.moe_mxfp4_gemm_rdna3( + act_out, w2_out, w2, qc.w2_scale, empty_tw, + sti, eid, ntp, 1, block_size_m, False, 0, + ) + w2_out = w2_out + w2_bias[flat_experts].to(dtype) + if not apply_router_weight_on_input: + w2_out = w2_out * topk_weights.reshape(-1, 1).to(dtype) + output.copy_(w2_out.view(M, top_k, hidden).sum(dim=1)) + else: + output.zero_() + ops.moe_mxfp4_gemm_rdna3( + act_out, output, w2, qc.w2_scale, + empty_tw if apply_router_weight_on_input else topk_w_f32, + sti, eid, ntp, top_k, block_size_m, + not apply_router_weight_on_input, top_k, + ) + + @staticmethod + def _swiglu_oai(x: torch.Tensor, out: torch.Tensor, qc) -> None: + # gate||up contiguous; OAI clamped SwiGLU with quant-config params. + alpha = qc.gemm1_alpha if qc.gemm1_alpha is not None else 1.702 + beta = qc.gemm1_beta if qc.gemm1_beta is not None else 1.0 + limit = qc.gemm1_clamp_limit + d = x.shape[-1] // 2 + gate = x[..., :d] + up = x[..., d:] + if limit is not None: + gate = torch.clamp(gate, max=limit) + up = torch.clamp(up, min=-limit, max=limit) + out.copy_(gate * torch.sigmoid(alpha * gate) * (up + beta)) diff --git a/vllm/model_executor/layers/fused_moe/oracle/mxfp4.py b/vllm/model_executor/layers/fused_moe/oracle/mxfp4.py index e6e9d17925ae..bd7fa6b3c1f7 100644 --- a/vllm/model_executor/layers/fused_moe/oracle/mxfp4.py +++ b/vllm/model_executor/layers/fused_moe/oracle/mxfp4.py @@ -110,6 +110,8 @@ class Mxfp4MoeBackend(Enum): AITER = "AITER_MXFP4_BF16" AITER_MXFP4_FP8 = "AITER_MXFP4_FP8" # W4A8: triton kernel AITER_MXFP4_MXFP4 = "AITER_MXFP4_MXFP4" # W4A4: CK kernel + # ROCm RDNA3 (gfx1100) native HIP kernel + RDNA3_MXFP4 = "RDNA3_MXFP4" # Triton TRITON = "TRITON" TRITON_UNFUSED = "TRITON_UNFUSED" @@ -191,6 +193,13 @@ def backend_to_kernel_cls( return [UnfusedOAITritonExperts] + elif backend == Mxfp4MoeBackend.RDNA3_MXFP4: + from vllm.model_executor.layers.fused_moe.experts.rdna3_mxfp4_moe import ( + RDNA3Mxfp4Experts, + ) + + return [RDNA3Mxfp4Experts] + elif backend == Mxfp4MoeBackend.HUMMING: from vllm.model_executor.layers.fused_moe.experts.fused_humming_moe import ( BatchedHummingGroupedExperts, @@ -305,6 +314,8 @@ def _get_priority_backends_for_gpt_oss() -> list[Mxfp4MoeBackend]: """Available backends in priority order, BF16-act variant before activation-quantized variant within each vendor family.""" _AVAILABLE_BACKENDS = [ + # gfx1100 native HIP kernel; is_supported_config gates it off elsewhere. + Mxfp4MoeBackend.RDNA3_MXFP4, Mxfp4MoeBackend.FLASHINFER_TRTLLM_MXFP4_BF16, Mxfp4MoeBackend.FLASHINFER_TRTLLM_MXFP4_MXFP8, Mxfp4MoeBackend.AITER_MXFP4_BF16, @@ -330,7 +341,7 @@ def _get_priority_backends() -> list[Mxfp4MoeBackend]: backend-level ``is_supported_config`` check filters by device capability). """ if current_platform.is_rocm(): - return [Mxfp4MoeBackend.AITER_MXFP4_BF16] + return [Mxfp4MoeBackend.RDNA3_MXFP4, Mxfp4MoeBackend.AITER_MXFP4_BF16] if current_platform.is_xpu(): return [Mxfp4MoeBackend.XPU] _AVAILABLE_BACKENDS = [ @@ -742,6 +753,25 @@ def convert_gpt_oss_weight_to_mxfp4_moe_kernel_format( w2_bias, ) + elif mxfp4_backend == Mxfp4MoeBackend.RDNA3_MXFP4: + from vllm.model_executor.layers.fused_moe.experts.rdna3_mxfp4_moe import ( + repack_experts_rdna3, + ) + + # [E, K/8, N] uint32 + [E, K/32, N] uint8; w13 gate/up de-interleaved so + # the gate_up output is contiguous [gate || up]. + b_q13, b_s13 = repack_experts_rdna3( + w13_weight.data, w13_weight_scale.data, deinterleave=True + ) + b_q2, b_s2 = repack_experts_rdna3( + w2_weight.data, w2_weight_scale.data, deinterleave=False + ) + if w13_bias is not None: + w13_bias = torch.cat( + [w13_bias.data[:, ::2], w13_bias.data[:, 1::2]], dim=1 + ).contiguous() + return b_q13, b_q2, b_s13, b_s2, w13_bias, w2_bias + elif mxfp4_backend in TRTLLM_BACKENDS: assert _cache_permute_indices is not None from flashinfer.fp4_quantization import nvfp4_block_scale_interleave @@ -1634,6 +1664,7 @@ def make_mxfp4_moe_quant_config( Mxfp4MoeBackend.FLASHINFER_TRTLLM_MXFP4_BF16, Mxfp4MoeBackend.FLASHINFER_CUTLASS_MXFP4_BF16, Mxfp4MoeBackend.AITER_MXFP4_BF16, + Mxfp4MoeBackend.RDNA3_MXFP4, Mxfp4MoeBackend.CPU, ): return mxfp4_w4a16_moe_quant_config( diff --git a/vllm/model_executor/layers/quantization/compressed_tensors/compressed_tensors_moe/compressed_tensors_moe.py b/vllm/model_executor/layers/quantization/compressed_tensors/compressed_tensors_moe/compressed_tensors_moe.py index ead1f233aab7..00221485233f 100644 --- a/vllm/model_executor/layers/quantization/compressed_tensors/compressed_tensors_moe/compressed_tensors_moe.py +++ b/vllm/model_executor/layers/quantization/compressed_tensors/compressed_tensors_moe/compressed_tensors_moe.py @@ -64,11 +64,6 @@ def get_moe_method( format = scheme_dict.get("format") if quant_config._is_mxfp4(weight_quant): - from . import rocm_moe_rdna - - if rocm_moe_rdna.is_supported_mxfp4(weight_quant): - return rocm_moe_rdna.make_mxfp4_method(layer.moe_config) - from .compressed_tensors_moe_w4a4_mxfp4 import ( CompressedTensorsW4A4Mxfp4MoEMethod, ) diff --git a/vllm/model_executor/layers/quantization/compressed_tensors/compressed_tensors_moe/compressed_tensors_moe_w4a4_mxfp4.py b/vllm/model_executor/layers/quantization/compressed_tensors/compressed_tensors_moe/compressed_tensors_moe_w4a4_mxfp4.py index 906b0727b182..1ec4b26be429 100644 --- a/vllm/model_executor/layers/quantization/compressed_tensors/compressed_tensors_moe/compressed_tensors_moe_w4a4_mxfp4.py +++ b/vllm/model_executor/layers/quantization/compressed_tensors/compressed_tensors_moe/compressed_tensors_moe_w4a4_mxfp4.py @@ -21,6 +21,10 @@ from vllm.model_executor.layers.fused_moe.experts.marlin_moe import ( MarlinExperts, ) +from vllm.model_executor.layers.fused_moe.experts.rdna3_mxfp4_moe import ( + RDNA3Mxfp4Experts, + repack_experts_rdna3, +) from vllm.model_executor.layers.fused_moe.experts.xpu_moe import ( XPUExpertsMxFp4, ) @@ -56,6 +60,10 @@ def __init__(self, moe): self.mxfp4_backend = Mxfp4MoeBackend.XPU self.experts_cls = XPUExpertsMxFp4 logger.info_once("Using XPUExpertsMxFp4 for MXFP4 MoE on XPU platform") + elif RDNA3Mxfp4Experts._supports_current_device(): + self.mxfp4_backend = Mxfp4MoeBackend.RDNA3_MXFP4 + self.experts_cls = RDNA3Mxfp4Experts + logger.info_once("Using RDNA3Mxfp4Experts (native gfx1100 HIP kernel)") else: logger.info_once("Using MarlinExperts for MXFP4 MoE") self.experts_cls = MarlinExperts @@ -189,6 +197,19 @@ def process_weights_after_loading(self, layer: RoutedExperts) -> None: ) elif current_platform.is_xpu(): pass + elif self.mxfp4_backend == Mxfp4MoeBackend.RDNA3_MXFP4: + # Repack to the [E, K/8, N] uint32 + [E, K/32, N] layout the HIP + # kernel reads (no gate/up de-interleave for compressed-tensors). + b_q13, b_s13 = repack_experts_rdna3( + layer.w13_weight.data, layer.w13_weight_scale.data, False + ) + b_q2, b_s2 = repack_experts_rdna3( + layer.w2_weight.data, layer.w2_weight_scale.data, False + ) + layer.w13_weight = torch.nn.Parameter(b_q13, requires_grad=False) + layer.w13_weight_scale = torch.nn.Parameter(b_s13, requires_grad=False) + layer.w2_weight = torch.nn.Parameter(b_q2, requires_grad=False) + layer.w2_weight_scale = torch.nn.Parameter(b_s2, requires_grad=False) else: logger.warning_once( "Your GPU does not have native support for FP4 computation " diff --git a/vllm/model_executor/layers/quantization/compressed_tensors/compressed_tensors_moe/compressed_tensors_moe_w4a4_mxfp4_rdna3.py b/vllm/model_executor/layers/quantization/compressed_tensors/compressed_tensors_moe/compressed_tensors_moe_w4a4_mxfp4_rdna3.py deleted file mode 100644 index dfb5df04a123..000000000000 --- a/vllm/model_executor/layers/quantization/compressed_tensors/compressed_tensors_moe/compressed_tensors_moe_w4a4_mxfp4_rdna3.py +++ /dev/null @@ -1,75 +0,0 @@ -# SPDX-License-Identifier: Apache-2.0 -# SPDX-FileCopyrightText: Copyright contributors to the vLLM project -"""CompressedTensors MXFP4 MoE using the fused RDNA3 (gfx1100) HIP kernel. - -Uses ``moe_mxfp4_gemm_rdna3`` — one HIP launch per GEMM that fuses expert -routing + E2M1/E8M0 dequant + dot product with atomic output. No zero-point; -E8M0 block scale is ``[E, K/32, N]`` uint8. The two-GEMM forward is shared -with the W4A16 path via RDNA3FusedMoEMixin. -""" - -import torch - -from vllm import _custom_ops as ops -from vllm.model_executor.layers.quantization.compressed_tensors.compressed_tensors_moe.compressed_tensors_moe_w4a4_mxfp4 import ( # noqa: E501 - CompressedTensorsW4A4Mxfp4MoEMethod, -) - -from .rdna3_moe_common import RDNA3FusedMoEMixin, repack_experts - - -class CompressedTensorsW4A4Mxfp4RDNA3MoEMethod( - RDNA3FusedMoEMixin, CompressedTensorsW4A4Mxfp4MoEMethod -): - """MXFP4 MoE via the native RDNA3 HIP kernel (moe_mxfp4_gemm_rdna3).""" - - def process_weights_after_loading(self, layer: torch.nn.Module) -> None: - device = layer.w13_weight_packed.device - b_q13, b_s13 = repack_experts( - layer.w13_weight_packed.data, layer.w13_weight_scale.data, False - ) - del layer.w13_weight_packed, layer.w13_weight_scale - layer.w13_weight_packed = torch.nn.Parameter(b_q13, requires_grad=False) - layer.w13_weight_scale = torch.nn.Parameter(b_s13, requires_grad=False) - - b_q2, b_s2 = repack_experts( - layer.w2_weight_packed.data, layer.w2_weight_scale.data, False - ) - del layer.w2_weight_packed, layer.w2_weight_scale - layer.w2_weight_packed = torch.nn.Parameter(b_q2, requires_grad=False) - layer.w2_weight_scale = torch.nn.Parameter(b_s2, requires_grad=False) - layer.rdna3_empty_tw = torch.empty(0, device=device) - - def _gemm_w13(self, layer, a, c, tw, sti, eid, ntp, top_k, block_size_m, mul_tw): - ops.moe_mxfp4_gemm_rdna3( - a, - c, - layer.w13_weight_packed, - layer.w13_weight_scale, - tw, - sti, - eid, - ntp, - top_k, - block_size_m, - mul_tw, - 0, - ) - - def _gemm_w2( - self, layer, a, c, tw, sti, eid, ntp, block_size_m, mul_tw, output_topk - ): - ops.moe_mxfp4_gemm_rdna3( - a, - c, - layer.w2_weight_packed, - layer.w2_weight_scale, - tw, - sti, - eid, - ntp, - 1, - block_size_m, - mul_tw, - output_topk, - ) diff --git a/vllm/model_executor/layers/quantization/compressed_tensors/compressed_tensors_moe/rocm_moe_rdna.py b/vllm/model_executor/layers/quantization/compressed_tensors/compressed_tensors_moe/rocm_moe_rdna.py index 0a8056e34e6b..e72caa0796a5 100644 --- a/vllm/model_executor/layers/quantization/compressed_tensors/compressed_tensors_moe/rocm_moe_rdna.py +++ b/vllm/model_executor/layers/quantization/compressed_tensors/compressed_tensors_moe/rocm_moe_rdna.py @@ -28,29 +28,6 @@ def is_supported(weight_quant) -> bool: ) -def is_supported_mxfp4(weight_quant) -> bool: - """Check if the native RDNA3 MXFP4 MoE kernel is available.""" - from vllm.platforms.rocm import on_gfx1100 - - return ( - on_gfx1100() - and hasattr(torch.ops, "_rocm_C") - and hasattr(torch.ops._rocm_C, "moe_mxfp4_gemm_rdna3") - ) - - -def make_mxfp4_method(moe_config): - """Create the native RDNA3 MXFP4 MoE method. Call after is_supported_mxfp4.""" - from .compressed_tensors_moe_w4a4_mxfp4_rdna3 import ( - CompressedTensorsW4A4Mxfp4RDNA3MoEMethod, - ) - - logger.info_once( - "Using CompressedTensorsW4A4Mxfp4RDNA3MoEMethod (native RDNA3 HIP kernel)" - ) - return CompressedTensorsW4A4Mxfp4RDNA3MoEMethod(moe_config) - - 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 diff --git a/vllm/model_executor/layers/quantization/gpt_oss_mxfp4_rdna3.py b/vllm/model_executor/layers/quantization/gpt_oss_mxfp4_rdna3.py deleted file mode 100644 index 59fef5817850..000000000000 --- a/vllm/model_executor/layers/quantization/gpt_oss_mxfp4_rdna3.py +++ /dev/null @@ -1,173 +0,0 @@ -# SPDX-License-Identifier: Apache-2.0 -# SPDX-FileCopyrightText: Copyright contributors to the vLLM project -"""GPT-OSS MXFP4 MoE on RDNA3 (gfx1100) via the native HIP kernel. - -GPT-OSS is native-mxfp4 (GptOssMxfp4MoEMethod), not compressed-tensors, and -differs from the plain MXFP4 MoE in two ways handled here: - - per-expert biases (w13_bias, w2_bias) - - the SwiGLU-OAI clamped activation (alpha=1.702, beta=1.0, limit=7.0) - -w13 ships gate/up *interleaved*; we de-interleave at load so the gate_up -output is contiguous [gate || up]. Biases and the moe_sum reduction run in -Python (output_topk=0), so the moe_mxfp4_gemm_rdna3 kernel is reused as-is. -""" - -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.moe_align_block_size import ( - moe_align_block_size, -) -from vllm.model_executor.layers.quantization.compressed_tensors.compressed_tensors_moe.rdna3_moe_common import ( # noqa: E501 - repack_experts, - select_block_size_m, -) -from vllm.model_executor.layers.quantization.mxfp4 import GptOssMxfp4MoEMethod - -logger = init_logger(__name__) - -_ALPHA, _BETA, _LIMIT = 1.702, 1.0, 7.0 - - -def maybe_rdna3_gpt_oss_method(layer): - """Return the native RDNA3 GPT-OSS MXFP4 MoE method on gfx1100, else None. - - Keeps the gfx1100 selection out of the generic ``Mxfp4Config`` so the only - footprint there is a one-line delegation (mirrors ``rocm_moe_rdna`` for the - compressed-tensors path). - """ - from vllm.platforms import current_platform - - if not ( - current_platform.is_rocm() - and hasattr(torch.ops, "_rocm_C") - and hasattr(torch.ops._rocm_C, "moe_mxfp4_gemm_rdna3") - ): - return None - from vllm.platforms.rocm import on_gfx1100 - - if not on_gfx1100(): - return None - logger.info_once("Using GptOssMxfp4RDNA3MoEMethod (native RDNA3 HIP kernel)") - return GptOssMxfp4RDNA3MoEMethod(layer.moe_config) - - -class GptOssMxfp4RDNA3MoEMethod(GptOssMxfp4MoEMethod): - """GPT-OSS MXFP4 MoE via the native RDNA3 HIP kernel (moe_mxfp4_gemm_rdna3). - - Bypasses the modular MoE kernel: the runner calls apply() directly through - forward_modular. We report supports_internal_mk (skip MK construction), - not monolithic (route through forward_modular -> apply), and a null - fused-moe quant config (apply does the dequant itself). - """ - - @property - def supports_internal_mk(self) -> bool: - return True - - @property - def is_monolithic(self) -> bool: - return False - - def get_fused_moe_quant_config(self, layer): - return None - - def process_weights_after_loading(self, layer: RoutedExperts) -> None: - device = layer.w13_weight.device - # w13: de-interleave gate/up + repack, then free the source. - b_q13, b_s13 = repack_experts( - layer.w13_weight.data, layer.w13_weight_scale.data, deinterleave=True - ) - b13 = torch.cat( - [layer.w13_bias.data[:, ::2], layer.w13_bias.data[:, 1::2]], dim=1 - ).contiguous() - del layer.w13_weight, layer.w13_weight_scale - layer.w13_weight_packed = torch.nn.Parameter(b_q13, requires_grad=False) - layer.w13_weight_scale = torch.nn.Parameter(b_s13, requires_grad=False) - layer.w13_bias = torch.nn.Parameter(b13, requires_grad=False) - - # w2: plain repack (no de-interleave), then free the source. - b_q2, b_s2 = repack_experts( - layer.w2_weight.data, layer.w2_weight_scale.data, deinterleave=False - ) - del layer.w2_weight, layer.w2_weight_scale - layer.w2_weight_packed = torch.nn.Parameter(b_q2, requires_grad=False) - layer.w2_weight_scale = torch.nn.Parameter(b_s2, requires_grad=False) - layer.rdna3_empty_tw = torch.empty(0, device=device) - - @staticmethod - def _swiglu_oai(x: torch.Tensor) -> torch.Tensor: - # gate||up contiguous (w13 de-interleaved at load). SwiGLU-OAI: - # out = clamp(gate) * sigmoid(alpha*gate) * (clamp(up) + beta) - d = x.shape[-1] // 2 - gate = torch.clamp(x[..., :d], max=_LIMIT) - up = torch.clamp(x[..., d:], min=-_LIMIT, max=_LIMIT) - return gate * torch.sigmoid(_ALPHA * gate) * (up + _BETA) - - 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: - M = x.shape[0] - top_k = topk_ids.shape[1] - total = M * top_k - N_gate_up = layer.w13_weight_packed.shape[2] - hidden = layer.w2_weight_packed.shape[2] - dtype = x.dtype - device = x.device - - gne = layer.global_num_experts - if gne <= 0: - gne = layer.w13_weight_packed.shape[0] - block_size_m = select_block_size_m(M, top_k, gne) - sti, eid, ntp = moe_align_block_size( - topk_ids, block_size_m, gne, layer.expert_map - ) - flat_experts = topk_ids.reshape(-1).long() - empty_tw = layer.rdna3_empty_tw - - # gate_up GEMM (no reduction) + per-expert bias - w1 = torch.zeros(total, N_gate_up, dtype=dtype, device=device) - ops.moe_mxfp4_gemm_rdna3( - x, - w1, - layer.w13_weight_packed, - layer.w13_weight_scale, - empty_tw, - sti, - eid, - ntp, - top_k, - block_size_m, - False, - 0, - ) - w1 = w1 + layer.w13_bias[flat_experts].to(dtype) - act = self._swiglu_oai(w1) - - # down GEMM (no reduction) + per-expert bias, then weighted moe_sum - w2o = torch.zeros(total, hidden, dtype=dtype, device=device) - ops.moe_mxfp4_gemm_rdna3( - act, - w2o, - layer.w2_weight_packed, - layer.w2_weight_scale, - empty_tw, - sti, - eid, - ntp, - 1, - block_size_m, - False, - 0, - ) - w2o = w2o + layer.w2_bias[flat_experts].to(dtype) - w2o = w2o * topk_weights.reshape(-1, 1).to(dtype) - return w2o.view(M, top_k, hidden).sum(dim=1) diff --git a/vllm/model_executor/layers/quantization/mxfp4.py b/vllm/model_executor/layers/quantization/mxfp4.py index 8174a860d221..5ef5fd40d5eb 100644 --- a/vllm/model_executor/layers/quantization/mxfp4.py +++ b/vllm/model_executor/layers/quantization/mxfp4.py @@ -86,13 +86,7 @@ def get_quant_method( ) return UnquantizedLinearMethod() elif isinstance(layer, RoutedExperts): - from vllm.model_executor.layers.quantization.gpt_oss_mxfp4_rdna3 import ( - maybe_rdna3_gpt_oss_method, - ) - - return maybe_rdna3_gpt_oss_method(layer) or GptOssMxfp4MoEMethod( - layer.moe_config - ) + return GptOssMxfp4MoEMethod(layer.moe_config) elif isinstance(layer, Attention): logger.debug_once( "MXFP4 attention layer is not implemented. " From c8db136bbcb5e66cd58108a0696fc338870cb90f Mon Sep 17 00:00:00 2001 From: JartX Date: Fri, 26 Jun 2026 03:23:35 +0200 Subject: [PATCH 19/36] [ROCm][gfx1100] MXFP4 MoE: leave the W4A16 INT4 path untouched The oracle-backend refactor no longer needs the shared RDNA3FusedMoEMixin, so revert the pre-existing W4A16 (moe_gptq_gemm_rdna3) path to its main version instead of carrying it through a now-single-user mixin: restore compressed_tensors_moe_wna16_rdna3.py to main (its own inline block_size_m cap) and drop rdna3_moe_common.py. This PR now only touches the MXFP4 MoE paths. Repoint the dispatch test to the MXFP4 experts tile selector. Signed-off-by: JartX --- .../quantization/test_mxfp4_rdna3_dispatch.py | 70 ++----- .../compressed_tensors_moe_wna16_rdna3.py | 187 ++++++++++++++---- .../rdna3_moe_common.py | 182 ----------------- 3 files changed, 165 insertions(+), 274 deletions(-) delete mode 100644 vllm/model_executor/layers/quantization/compressed_tensors/compressed_tensors_moe/rdna3_moe_common.py diff --git a/tests/kernels/quantization/test_mxfp4_rdna3_dispatch.py b/tests/kernels/quantization/test_mxfp4_rdna3_dispatch.py index 0c745fba2d2b..2326f0ee80e9 100644 --- a/tests/kernels/quantization/test_mxfp4_rdna3_dispatch.py +++ b/tests/kernels/quantization/test_mxfp4_rdna3_dispatch.py @@ -1,74 +1,44 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project -"""Block-size dispatch tests for the RDNA3 (gfx1100) fused-MoE methods. +"""Block-size dispatch tests for the RDNA3 (gfx1100) MXFP4 MoE experts backend. -``RDNA3FusedMoEMixin`` is shared by the MXFP4 MoE method and the existing W4A16 -(``moe_gptq_gemm_rdna3``) method. The MXFP4 kernel has a WMMA-16 tile; the W4A16 -kernel only instantiates tiles {1,2,4,8} and rejects anything else. So the -W4A16 method must override tile selection to stay inside its set — otherwise any ->=16-token prefill / batched decode hard-crashes. These are pure-Python checks -(no GPU), so the regression is caught on any CI host. +``RDNA3Mxfp4Experts`` (the ``Mxfp4MoeBackend.RDNA3_MXFP4`` kernel) picks a GEMM +tile from expected occupancy: the WMMA-16 tile once a step has >=16 rows, else a +scalar tile (1/4). The ``moe_mxfp4_gemm_rdna3`` kernel only instantiates tiles +{1,2,4,8,16}, so the selector must never return anything else. Pure-Python (no +GPU), so this runs on any CI host. """ import pytest -from vllm.model_executor.layers.quantization.compressed_tensors.compressed_tensors_moe.rdna3_moe_common import ( # noqa: E501 - RDNA3FusedMoEMixin, - select_block_size_m, +from vllm.model_executor.layers.fused_moe.experts.rdna3_mxfp4_moe import ( + _select_block_size_m, ) -# Tiles each HIP kernel actually instantiates (the rest hit a TORCH_CHECK). -GPTQ_TILES = {1, 2, 4, 8} # moe_gptq_gemm_rdna3 (W4A16) -MXFP4_TILES = {1, 2, 4, 8, 16} # moe_mxfp4_gemm_rdna3 - +SUPPORTED_TILES = {1, 2, 4, 8, 16} # moe_mxfp4_gemm_rdna3 instantiations BATCH_SIZES = [1, 2, 3, 4, 8, 15, 16, 17, 32, 64, 256, 512, 4096] TOP_KS = [1, 2, 4, 8] NUM_EXPERTS = [8, 32, 128, 256] -def _w4a16_cls(): - from vllm.model_executor.layers.quantization.compressed_tensors.compressed_tensors_moe.compressed_tensors_moe_wna16_rdna3 import ( # noqa: E501 - CompressedTensorsWNA16RDNA3MoEMethod, - ) - - return CompressedTensorsWNA16RDNA3MoEMethod - - @pytest.mark.parametrize("num_tokens", BATCH_SIZES) @pytest.mark.parametrize("top_k", TOP_KS) @pytest.mark.parametrize("num_experts", NUM_EXPERTS) -def test_mxfp4_tile_in_supported_set(num_tokens, top_k, num_experts): - bsm = select_block_size_m(num_tokens, top_k, num_experts) - assert bsm in MXFP4_TILES - # the WMMA-16 tile is only picked once a step can fill it +def test_tile_in_supported_set(num_tokens, top_k, num_experts): + bsm = _select_block_size_m(num_tokens, top_k, num_experts) + assert bsm in SUPPORTED_TILES + # the WMMA-16 tile is only chosen once a step can fill it if num_tokens < 16: assert bsm != 16 -@pytest.mark.parametrize("num_tokens", BATCH_SIZES) -@pytest.mark.parametrize("top_k", TOP_KS) -@pytest.mark.parametrize("num_experts", NUM_EXPERTS) -def test_w4a16_never_selects_unsupported_tile(num_tokens, top_k, num_experts): - """Regression guard: the W4A16 method must never hand moe_gptq_gemm_rdna3 a - tile outside {1,2,4,8} (>=16-token batches used to select 16 and crash).""" - # _select_block_size_m is stateless, so call it without building a method. - bsm = _w4a16_cls()._select_block_size_m(None, num_tokens, top_k, num_experts) - assert bsm in GPTQ_TILES, ( - f"W4A16 selected block_size_m={bsm} for num_tokens={num_tokens}; " - f"moe_gptq_gemm_rdna3 only supports {sorted(GPTQ_TILES)}" - ) - - -def test_w4a16_override_differs_from_mixin_default(): - """The mixin default returns the WMMA-16 tile for a 32-token batch; the - W4A16 override must replace it with a kernel-supported tile.""" - assert RDNA3FusedMoEMixin._select_block_size_m(None, 32, 8, 256) == 16 - w4a16 = _w4a16_cls()._select_block_size_m(None, 32, 8, 256) - assert w4a16 != 16 - assert w4a16 in GPTQ_TILES +def test_large_batches_use_wmma_tile(): + # >=16-row steps (prefill / high-concurrency decode) take the WMMA tile. + assert _select_block_size_m(16, 8, 256) == 16 + assert _select_block_size_m(512, 8, 256) == 16 def test_degenerate_inputs_return_safe_tile(): - # num_experts/top_k <= 0 must not divide-by-zero; fall back to the 1-tile. - assert select_block_size_m(32, 0, 256) == 1 - assert select_block_size_m(32, 8, 0) == 1 + # top_k/num_experts <= 0 must not divide-by-zero; fall back to the 1-tile. + assert _select_block_size_m(32, 0, 256) == 1 + assert _select_block_size_m(32, 8, 0) == 1 diff --git a/vllm/model_executor/layers/quantization/compressed_tensors/compressed_tensors_moe/compressed_tensors_moe_wna16_rdna3.py b/vllm/model_executor/layers/quantization/compressed_tensors/compressed_tensors_moe/compressed_tensors_moe_wna16_rdna3.py index a24800df5f6f..25cc7b115f02 100644 --- a/vllm/model_executor/layers/quantization/compressed_tensors/compressed_tensors_moe/compressed_tensors_moe_wna16_rdna3.py +++ b/vllm/model_executor/layers/quantization/compressed_tensors/compressed_tensors_moe/compressed_tensors_moe_wna16_rdna3.py @@ -15,6 +15,17 @@ 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, ) @@ -23,8 +34,6 @@ ) from vllm.scalar_type import scalar_types -from .rdna3_moe_common import RDNA3FusedMoEMixin - logger = init_logger(__name__) @@ -45,9 +54,7 @@ def _synthesize_qzeros( return pack_quantized_values_into_int32(zeros, scalar_types.uint4b8, packed_dim=1) -class CompressedTensorsWNA16RDNA3MoEMethod( - RDNA3FusedMoEMixin, CompressedTensorsWNA16MoEMethod -): +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]), @@ -116,43 +123,139 @@ def process_weights_after_loading(self, layer: torch.nn.Module) -> None: ) layer.rdna3_empty_tw = torch.empty(0, device=device) - def _select_block_size_m(self, num_tokens, top_k, num_experts): - # moe_gptq_gemm_rdna3 only has tiles {1,2,4,8} (no WMMA-16 path); keep - # the proven W4A16 cap so prefill / batched decode don't hit the - # kernel's TORCH_CHECK default arm. - return 1 if num_tokens <= 4 else 4 - - def _gemm_w13(self, layer, a, c, tw, sti, eid, ntp, top_k, block_size_m, mul_tw): - ops.moe_gptq_gemm_rdna3( - a, - c, - layer.w13_weight_packed, - layer.w13_weight_scale, - layer.w13_qzeros, - tw, - sti, - eid, - ntp, - top_k, - block_size_m, - mul_tw, + 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 _gemm_w2( - self, layer, a, c, tw, sti, eid, ntp, block_size_m, mul_tw, output_topk - ): - ops.moe_gptq_gemm_rdna3( - a, - c, - layer.w2_weight_packed, - layer.w2_weight_scale, - layer.w2_qzeros, - tw, - sti, - eid, - ntp, - 1, - block_size_m, - mul_tw, - output_topk, + +def _rdna3_fused_moe( + hidden_states: torch.Tensor, + topk_weights: torch.Tensor, + topk_ids: torch.Tensor, + layer: RoutedExperts, + activation: MoEActivation, + apply_router_weight_on_input: bool, + global_num_experts: int, + expert_map: torch.Tensor | None, +) -> torch.Tensor: + """Fused MoE forward using the RDNA3 W4A16 HIP kernel. + + Optimizations vs naive dispatch: + - BLOCK_SIZE_M=1 for decode (no padding waste, bf16 fast path) + - Pre-allocated buffers (no torch.zeros per call) + - Inline token sorting for small M (skip moe_align_block_size) + - moe_sum fused into output accumulation + """ + num_tokens = hidden_states.shape[0] + top_k = topk_ids.shape[1] + total_tokens = num_tokens * top_k + N_gate_up = layer.w13_weight_packed.shape[2] + hidden_size = layer.w2_weight_packed.shape[2] + dtype = hidden_states.dtype + device = hidden_states.device + + intermediate_size = N_gate_up // 2 if activation.is_gated else N_gate_up + + if global_num_experts <= 0: + global_num_experts = layer.w13_weight_packed.shape[0] + + # BLOCK_SIZE_M=1 for decode (small M), 4 for prefill + block_size_m = 1 if num_tokens <= 4 else 4 + + # --- Token routing --- + sorted_token_ids, expert_ids, num_tokens_post_padded = moe_align_block_size( + topk_ids, + block_size_m, + global_num_experts, + expert_map, + ) + + # --- Reuse pre-allocated buffers when possible --- + if total_tokens <= layer.rdna3_w1_buf.shape[0]: + w1_out = layer.rdna3_w1_buf[:total_tokens] + w1_out.zero_() + act_out = layer.rdna3_act_buf[:total_tokens] + else: + w1_out = torch.zeros( + total_tokens, + N_gate_up, + dtype=dtype, + device=device, ) + act_out = torch.empty( + total_tokens, + intermediate_size, + dtype=dtype, + device=device, + ) + + # --- topk weights (pre-cast to float32 for kernel) --- + topk_w_float = topk_weights.view(-1).float() + empty_tw = layer.rdna3_empty_tw + + # --- w1 GEMM: [M, K] -> [M*top_k, N_gate_up] --- + ops.moe_gptq_gemm_rdna3( + hidden_states, + w1_out, + layer.w13_weight_packed, + layer.w13_weight_scale, + layer.w13_qzeros, + topk_w_float if apply_router_weight_on_input else empty_tw, + sorted_token_ids, + expert_ids, + num_tokens_post_padded, + top_k, + block_size_m, + apply_router_weight_on_input, + ) + + # --- Activation (silu_and_mul etc.) --- + apply_moe_activation(activation, act_out, w1_out) + + # --- w2 GEMM: [M*top_k, intermediate] -> [M, hidden] (fused reduce) --- + # output_topk=top_k: kernel writes to out[token_id / top_k] directly, + # fusing moe_sum into the atomic accumulation — saves one kernel launch + # and the w2_out intermediate buffer. + out = torch.zeros( + num_tokens, + hidden_size, + dtype=dtype, + device=device, + ) + ops.moe_gptq_gemm_rdna3( + act_out, + out, + layer.w2_weight_packed, + layer.w2_weight_scale, + layer.w2_qzeros, + topk_w_float if not apply_router_weight_on_input else empty_tw, + sorted_token_ids, + expert_ids, + num_tokens_post_padded, + 1, + block_size_m, + not apply_router_weight_on_input, + output_topk=top_k, + ) + return out diff --git a/vllm/model_executor/layers/quantization/compressed_tensors/compressed_tensors_moe/rdna3_moe_common.py b/vllm/model_executor/layers/quantization/compressed_tensors/compressed_tensors_moe/rdna3_moe_common.py deleted file mode 100644 index 0db818230664..000000000000 --- a/vllm/model_executor/layers/quantization/compressed_tensors/compressed_tensors_moe/rdna3_moe_common.py +++ /dev/null @@ -1,182 +0,0 @@ -# SPDX-License-Identifier: Apache-2.0 -# SPDX-FileCopyrightText: Copyright contributors to the vLLM project -"""Shared orchestration for the native RDNA3 (gfx1100) fused-MoE methods. - -The W4A16 (INT4) and MXFP4 RDNA3 paths run the same two-GEMM forward -(gate_up -> SwiGLU -> down, with the down-proj output_topk reduction fusing -moe_sum); only the per-GEMM HIP op differs. That common flow lives here; each -method supplies ``_gemm_w13`` / ``_gemm_w2``. -""" - -import torch - -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, -) - - -def select_block_size_m(num_tokens: int, top_k: int, num_experts: int) -> int: - """Pick the fused-MoE GEMM tile for the current batch. - - The WMMA-16 tile wins once a step has >=16 rows: real (clustered) routing - fills it well even on high-expert models, and this is also where prefill - lives (validated e2e: WMMA is fastest at concurrency 32 even on Qwen's 256 - experts). For sub-16 batches the 16-tile is mostly padding, so we fall back - to a scalar tile sized to the expected occupancy avg tokens/expert = - num_tokens*top_k/E (bsm=1 beats bsm=4 at low concurrency, +9% e2e at C=8). - """ - if num_experts <= 0 or top_k <= 0: - return 1 - if num_tokens >= 16: - return 16 - tokens_per_expert = num_tokens * top_k / num_experts - return 4 if tokens_per_expert >= 1.5 else 1 - - -def repack_experts(packed: torch.Tensor, scale: torch.Tensor, deinterleave: bool): - """[E, N, K/2] uint8 + [E, N, K/32] uint8 -> [E, K/8, N] int32 + [E, K/32, N]. - - Per-expert into pre-allocated outputs (peak ~2x ONE expert tensor, not the - whole stack) so a multi-GB MoE repack fits alongside the loaded weights. - Four consecutive packed bytes little-endian = a uint32 of 8 E2M1 codes - along K — exactly the [K/8, N] word layout the kernel reads. ``deinterleave`` - splits gate/up rows of w13 ([g0,u0,...] -> [g..., u...]) for GPT-OSS. - """ - E, N, kh = packed.shape - dev = packed.device - b_q = torch.empty(E, kh // 4, N, dtype=torch.int32, device=dev) - b_scale = torch.empty(E, scale.shape[2], N, dtype=torch.uint8, device=dev) - for e in range(E): - pe = packed[e] - se = scale[e] - if deinterleave: - pe = torch.cat([pe[::2], pe[1::2]], dim=0) - se = torch.cat([se[::2], se[1::2]], dim=0) - b_q[e] = pe.contiguous().view(torch.int32).t() - b_scale[e] = se.t() - return b_q, b_scale - - -class RDNA3FusedMoEMixin: - """Mixin: shared apply() + two-GEMM orchestration. Subclasses implement - ``_gemm_w13`` and ``_gemm_w2`` (the op-specific HIP launch).""" - - def _select_block_size_m( - self, num_tokens: int, top_k: int, num_experts: int - ) -> int: - """Tile size for the GEMM op. The default targets the MXFP4 kernel - (tiles 1/2/4/8/16); subclasses whose kernel supports a smaller set - (e.g. the W4A16 ``moe_gptq_gemm_rdna3``: 1/2/4/8) must override.""" - return select_block_size_m(num_tokens, top_k, num_experts) - - def _gemm_w13(self, layer, a, c, tw, sti, eid, ntp, top_k, block_size_m, mul_tw): - raise NotImplementedError - - def _gemm_w2( - self, layer, a, c, tw, sti, eid, ntp, block_size_m, mul_tw, output_topk - ): - raise NotImplementedError - - def _rdna3_run( - self, - layer, - hidden_states, - topk_weights, - topk_ids, - activation, - apply_router_weight_on_input, - global_num_experts, - expert_map, - ): - 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 = 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 = self._select_block_size_m(num_tokens, top_k, global_num_experts) - - sorted_token_ids, expert_ids, num_tokens_post_padded = moe_align_block_size( - topk_ids, block_size_m, global_num_experts, expert_map - ) - - # Use pre-allocated decode buffers when present and big enough. - w1_buf = getattr(layer, "rdna3_w1_buf", None) - if w1_buf is not None and total_tokens <= w1_buf.shape[0]: - w1_out = 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, dtype=dtype, device=device - ) - - topk_w_float = topk_weights.view(-1).float() - empty_tw = getattr(layer, "rdna3_empty_tw", None) - if empty_tw is None: - empty_tw = torch.empty(0, device=device) - - self._gemm_w13( - layer, - hidden_states, - w1_out, - 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, - ) - apply_moe_activation(activation, act_out, w1_out) - - out = torch.zeros(num_tokens, hidden_size, dtype=dtype, device=device) - self._gemm_w2( - layer, - act_out, - out, - topk_w_float if not apply_router_weight_on_input else empty_tw, - sorted_token_ids, - expert_ids, - num_tokens_post_padded, - block_size_m, - not apply_router_weight_on_input, - top_k, - ) - return out - - 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 self._rdna3_run( - layer, - x, - topk_weights, - topk_ids, - activation, - layer.apply_router_weight_on_input, - layer.global_num_experts, - layer.expert_map, - ) From 95f26261a9dcda9c16df553ebac1f97306ac67c9 Mon Sep 17 00:00:00 2001 From: JartX Date: Fri, 26 Jun 2026 03:26:12 +0200 Subject: [PATCH 20/36] [ROCm][gfx1100] MXFP4 MoE: ruff-format the experts backend Apply ruff-format (one-arg-per-line on the op calls) so the pre-commit format hook passes in CI. No functional change. Signed-off-by: JartX --- .../fused_moe/experts/rdna3_mxfp4_moe.py | 49 +++++++++++++++---- 1 file changed, 39 insertions(+), 10 deletions(-) diff --git a/vllm/model_executor/layers/fused_moe/experts/rdna3_mxfp4_moe.py b/vllm/model_executor/layers/fused_moe/experts/rdna3_mxfp4_moe.py index 065cf98c0a43..e723751e1e02 100644 --- a/vllm/model_executor/layers/fused_moe/experts/rdna3_mxfp4_moe.py +++ b/vllm/model_executor/layers/fused_moe/experts/rdna3_mxfp4_moe.py @@ -175,16 +175,27 @@ def apply( empty_tw = torch.empty(0, device=device) w1_bias = qc.w1_bias w2_bias = qc.w2_bias - flat_experts = topk_ids.reshape(-1).long() if ( - w1_bias is not None or w2_bias is not None - ) else None + flat_experts = ( + topk_ids.reshape(-1).long() + if (w1_bias is not None or w2_bias is not None) + else None + ) # gate_up GEMM (no reduction); router weight optionally folded in here. w1_out = torch.zeros(total, n_gate_up, dtype=dtype, device=device) ops.moe_mxfp4_gemm_rdna3( - hidden_states, w1_out, w1, qc.w1_scale, + hidden_states, + w1_out, + w1, + qc.w1_scale, topk_w_f32 if apply_router_weight_on_input else empty_tw, - sti, eid, ntp, top_k, block_size_m, apply_router_weight_on_input, 0, + sti, + eid, + ntp, + top_k, + block_size_m, + apply_router_weight_on_input, + 0, ) if w1_bias is not None: w1_out = w1_out + w1_bias[flat_experts].to(dtype) @@ -201,8 +212,18 @@ def apply( if w2_bias is not None: w2_out = torch.zeros(total, hidden, dtype=dtype, device=device) ops.moe_mxfp4_gemm_rdna3( - act_out, w2_out, w2, qc.w2_scale, empty_tw, - sti, eid, ntp, 1, block_size_m, False, 0, + act_out, + w2_out, + w2, + qc.w2_scale, + empty_tw, + sti, + eid, + ntp, + 1, + block_size_m, + False, + 0, ) w2_out = w2_out + w2_bias[flat_experts].to(dtype) if not apply_router_weight_on_input: @@ -211,10 +232,18 @@ def apply( else: output.zero_() ops.moe_mxfp4_gemm_rdna3( - act_out, output, w2, qc.w2_scale, + act_out, + output, + w2, + qc.w2_scale, empty_tw if apply_router_weight_on_input else topk_w_f32, - sti, eid, ntp, top_k, block_size_m, - not apply_router_weight_on_input, top_k, + sti, + eid, + ntp, + top_k, + block_size_m, + not apply_router_weight_on_input, + top_k, ) @staticmethod From 2afb7c087370a8ae7465eed848ba445046e766b1 Mon Sep 17 00:00:00 2001 From: JartX Date: Fri, 26 Jun 2026 13:14:19 +0200 Subject: [PATCH 21/36] [ROCm][gfx1100] MXFP4 MoE: fix tensor-parallel reduce + route Quark/W4A4 to the native kernel RDNA3Mxfp4Experts produced wrong output under tensor parallelism: its w2 GEMM fuses the top-k reduction into the kernel epilogue (output_topk), but each TP rank only holds an intermediate shard, so its down-proj is a partial that the layer all-reduces afterwards (correct at TP1, garbage at TP2). Use the unreduced-rows + Python top-k reduce path (already taken for a per-expert bias) whenever tp_world_size > 1; TP1 keeps the fused fast path. Also route activation-quantized MXFP4 MoE (W4A4/W4A8, e.g. AMD Quark w_mxfp4_a_mxfp4) to the gfx1100 kernel: ROCm has no native FP4 so these degrade to weight-only, and the oracle ROCm fallback now prefers RDNA3_MXFP4 over the Triton-unfused path. w13 gate/up de-interleave is gated on model_type=="gpt_oss" instead of hard-coded, so separate gate_proj/up_proj checkpoints repack correctly. Supported and validated on gfx1100 (RX 7900 XTX): AMD Quark amd/Qwen3.5-35B-A3B-MXFP4 (W4A4) and compressed-tensors MXFP4 MoE on TP2, GPT-OSS-20B on TP1 and TP2, and compressed-tensors MXFP4 dense -- all on the native HIP kernel. Signed-off-by: JartX --- .../fused_moe/experts/rdna3_mxfp4_moe.py | 16 +++++++---- .../layers/fused_moe/oracle/mxfp4.py | 28 +++++++++++++++---- 2 files changed, 33 insertions(+), 11 deletions(-) diff --git a/vllm/model_executor/layers/fused_moe/experts/rdna3_mxfp4_moe.py b/vllm/model_executor/layers/fused_moe/experts/rdna3_mxfp4_moe.py index e723751e1e02..d84c8b3d027a 100644 --- a/vllm/model_executor/layers/fused_moe/experts/rdna3_mxfp4_moe.py +++ b/vllm/model_executor/layers/fused_moe/experts/rdna3_mxfp4_moe.py @@ -17,6 +17,7 @@ import vllm.model_executor.layers.fused_moe.modular_kernel as mk from vllm import _custom_ops as ops +from vllm.distributed import get_tensor_model_parallel_world_size from vllm.model_executor.layers.fused_moe.activation import ( MoEActivation, apply_moe_activation, @@ -206,10 +207,14 @@ def apply( else: apply_moe_activation(activation, act_out, w1_out) - # down GEMM; if there is a per-expert bias we must add it before the - # weighted reduction, so reduce in Python; otherwise fuse the reduction - # (output_topk) and the topk-weight multiply into the kernel epilogue. - if w2_bias is not None: + # down GEMM. The kernel can fuse the top-k reduction + weight multiply + # into its epilogue (output_topk), but that path produces wrong results + # under tensor parallelism (verified: correct at TP1, garbage at TP2 -- + # the per-rank partial down-proj is later all-reduced by the layer). + # A per-expert bias must also be added before the weighted reduction. + # In both cases write unreduced rows and reduce in Python instead. + unfused = w2_bias is not None or get_tensor_model_parallel_world_size() > 1 + if unfused: w2_out = torch.zeros(total, hidden, dtype=dtype, device=device) ops.moe_mxfp4_gemm_rdna3( act_out, @@ -225,7 +230,8 @@ def apply( False, 0, ) - w2_out = w2_out + w2_bias[flat_experts].to(dtype) + if w2_bias is not None: + w2_out = w2_out + w2_bias[flat_experts].to(dtype) if not apply_router_weight_on_input: w2_out = w2_out * topk_weights.reshape(-1, 1).to(dtype) output.copy_(w2_out.view(M, top_k, hidden).sum(dim=1)) diff --git a/vllm/model_executor/layers/fused_moe/oracle/mxfp4.py b/vllm/model_executor/layers/fused_moe/oracle/mxfp4.py index bd7fa6b3c1f7..87c36c301247 100644 --- a/vllm/model_executor/layers/fused_moe/oracle/mxfp4.py +++ b/vllm/model_executor/layers/fused_moe/oracle/mxfp4.py @@ -549,10 +549,21 @@ def select_mxfp4_moe_backend( ) if current_platform.is_rocm(): - backend = Mxfp4MoeBackend.TRITON_UNFUSED + # ROCm has no native FP4 compute, so activation-quantized MXFP4 + # (W4A4/W4A8) degrades to weight-only. On gfx1100 the native HIP + # kernel handles that faster than the Triton-unfused fallback. + from vllm.model_executor.layers.fused_moe.experts.rdna3_mxfp4_moe import ( + RDNA3Mxfp4Experts, + ) + + backend = ( + Mxfp4MoeBackend.RDNA3_MXFP4 + if RDNA3Mxfp4Experts._supports_current_device() + else Mxfp4MoeBackend.TRITON_UNFUSED + ) logger.info_once(_make_log_backend(backend)) return _return_or_raise( - Mxfp4MoeBackend.TRITON_UNFUSED, + backend, config, kMxfp4Static, None, @@ -754,19 +765,24 @@ def convert_gpt_oss_weight_to_mxfp4_moe_kernel_format( ) elif mxfp4_backend == Mxfp4MoeBackend.RDNA3_MXFP4: + # [E, K/8, N] uint32 + [E, K/32, N] uint8. GPT-OSS stores gate/up + # interleaved in a fused tensor (de-interleave so the gate_up output is + # contiguous [gate || up]); standard checkpoints with separate + # gate_proj/up_proj are already [gate_all, up_all] after _load_w13. + from vllm.config import get_current_vllm_config from vllm.model_executor.layers.fused_moe.experts.rdna3_mxfp4_moe import ( repack_experts_rdna3, ) - # [E, K/8, N] uint32 + [E, K/32, N] uint8; w13 gate/up de-interleaved so - # the gate_up output is contiguous [gate || up]. + model_type = get_current_vllm_config().model_config.hf_config.model_type + deint = model_type == "gpt_oss" b_q13, b_s13 = repack_experts_rdna3( - w13_weight.data, w13_weight_scale.data, deinterleave=True + w13_weight.data, w13_weight_scale.data, deinterleave=deint ) b_q2, b_s2 = repack_experts_rdna3( w2_weight.data, w2_weight_scale.data, deinterleave=False ) - if w13_bias is not None: + if deint and w13_bias is not None: w13_bias = torch.cat( [w13_bias.data[:, ::2], w13_bias.data[:, 1::2]], dim=1 ).contiguous() From 2495c0676b4c5b9ff58f65326b0a08d210e20ff1 Mon Sep 17 00:00:00 2001 From: JartX Date: Fri, 26 Jun 2026 13:26:34 +0200 Subject: [PATCH 22/36] [ROCm][gfx1100] Quark MXFP4: run weight-only dense linears on the native RDNA3 kernel QuarkOCP_MX linears fall back to dequant-to-bf16 emulation on gfx1100 (no native FP4). For statically-quantized weight-only MXFP4 their packed layout ([N, K/2] uint8 + [N, K/32] uint8 E8M0) is exactly what Rdna3MxFp4LinearKernel consumes, so route them through the native HIP GEMM instead, mirroring the compressed-tensors MXFP4 path. Activations stay in bf16 (W4A16). The kernel is chosen per-layer and only when N%16==0 and K%32==0; everything else (mxfp6, dynamic-quant, non-conforming shapes) keeps the existing emulation path. Verified on gfx1100 that the kernel path matches the dequant reference on real Quark MXFP4 weights (cosine 0.99999). Signed-off-by: JartX --- .../quark/schemes/quark_ocp_mx.py | 33 +++++++++++++++++++ 1 file changed, 33 insertions(+) diff --git a/vllm/model_executor/layers/quantization/quark/schemes/quark_ocp_mx.py b/vllm/model_executor/layers/quantization/quark/schemes/quark_ocp_mx.py index 70a7e81cc455..05c8e97f82df 100644 --- a/vllm/model_executor/layers/quantization/quark/schemes/quark_ocp_mx.py +++ b/vllm/model_executor/layers/quantization/quark/schemes/quark_ocp_mx.py @@ -207,6 +207,21 @@ def __init__( self.input_dtype != "mxfp4" or self.weight_dtype != "mxfp4" ) + # gfx1100 has no native FP4, so statically-quantized weight-only MXFP4 + # linears run on the native RDNA3 HIP GEMM instead of dequant-to-bf16 + # emulation (activations stay in bf16 == W4A16). The kernel is chosen + # per-layer in process_weights_after_loading once shapes are known. + self.rdna3_kernel = None + if self.weight_dtype == "mxfp4" and not self.dynamic_mxfp4_quant: + from vllm.model_executor.kernels.linear.mxfp4.rocm import ( + Rdna3MxFp4LinearKernel, + ) + + if Rdna3MxFp4LinearKernel.is_supported()[0]: + from vllm.model_executor.kernels.linear import MxFp4LinearLayerConfig + + self.rdna3_kernel = Rdna3MxFp4LinearKernel(MxFp4LinearLayerConfig()) + self.rocm_use_aiter_fp4_asm_gemm = ( rocm_aiter_ops.is_asm_fp4_gemm_dynamic_quant_enabled() ) @@ -265,9 +280,24 @@ def process_dynamic_mxfp4_weights_after_loading( layer.weight_scale = torch.nn.Parameter(w_s.T.contiguous(), requires_grad=False) layer.weight = torch.nn.Parameter(w_q, requires_grad=False) + def _rdna3_dims_ok(self, layer: torch.nn.Module) -> bool: + # The HIP kernel needs N%16==0 and K%32==0; layer.weight is [N, K/2] u8. + n, k_half = layer.weight.shape + return n % 16 == 0 and (k_half * 2) % 32 == 0 + def process_weights_after_loading(self, layer: torch.nn.Module) -> None: layer.weight = torch.nn.Parameter(layer.weight.data, requires_grad=False) + if self.rdna3_kernel is not None and self._rdna3_dims_ok(layer): + layer.output_size_per_partition = layer.weight.shape[0] + layer.weight_scale = torch.nn.Parameter( + layer.weight_scale.data, requires_grad=False + ) + self.rdna3_kernel.process_weights_after_loading(layer) + layer._rdna3 = True + return + layer._rdna3 = False + if self.emulate: if self.dynamic_mxfp4_quant: self.process_dynamic_mxfp4_weights_after_loading(layer) @@ -363,6 +393,9 @@ def apply_weights( x: torch.Tensor, bias: torch.Tensor | None = None, ) -> torch.Tensor: + if getattr(layer, "_rdna3", False): + assert self.rdna3_kernel is not None + return self.rdna3_kernel.apply_weights(layer, x, bias) if self.emulate: dq_w = self.dequant_func(layer.weight, layer.weight_scale, x.dtype) qdq_x = self.quant_dequant_func(x) From 09ebe428552707004a2bdac72fa7525498c20873 Mon Sep 17 00:00:00 2001 From: JartX Date: Fri, 26 Jun 2026 22:13:20 +0200 Subject: [PATCH 23/36] [ROCm][gfx1100] MXFP4 MoE: select RDNA3 via the priority list, not the fallback Addressing review feedback: - Move RDNA3_MXFP4 above TRITON in the priority list (drop the extra comment), and revert the deepseek-v4 _get_priority_backends ROCm branch. - Keep the ROCm fallback a plain TRITON_UNFUSED catch. Model the W4A4 weight-only downgrade where it belongs instead: QuarkOCP_MX_MoEMethod requests activation_key=None when the platform has no native FP4 (supports_mx()), so a Quark w_mxfp4_a_mxfp4 MoE is selected through the normal priority list (RDNA3 weight-only) rather than via the fallback. - QuarkOCP_MX dense scheme: use a per-layer self.rdna3_kernel (None when the shapes are unsupported) instead of a layer._rdna3 attribute. Verified on 2x RX 7900 XTX: Quark W4A4 and GPT-OSS W4A16 both select RDNA3_MXFP4 through the priority loop and produce correct output. Signed-off-by: JartX --- .../layers/fused_moe/oracle/mxfp4.py | 20 ++++--------------- .../layers/quantization/quark/quark_moe.py | 8 ++++++-- .../quark/schemes/quark_ocp_mx.py | 7 +++---- 3 files changed, 13 insertions(+), 22 deletions(-) diff --git a/vllm/model_executor/layers/fused_moe/oracle/mxfp4.py b/vllm/model_executor/layers/fused_moe/oracle/mxfp4.py index 87c36c301247..e4adcaeae33b 100644 --- a/vllm/model_executor/layers/fused_moe/oracle/mxfp4.py +++ b/vllm/model_executor/layers/fused_moe/oracle/mxfp4.py @@ -314,13 +314,12 @@ def _get_priority_backends_for_gpt_oss() -> list[Mxfp4MoeBackend]: """Available backends in priority order, BF16-act variant before activation-quantized variant within each vendor family.""" _AVAILABLE_BACKENDS = [ - # gfx1100 native HIP kernel; is_supported_config gates it off elsewhere. - Mxfp4MoeBackend.RDNA3_MXFP4, Mxfp4MoeBackend.FLASHINFER_TRTLLM_MXFP4_BF16, Mxfp4MoeBackend.FLASHINFER_TRTLLM_MXFP4_MXFP8, Mxfp4MoeBackend.AITER_MXFP4_BF16, Mxfp4MoeBackend.AITER_MXFP4_FP8, Mxfp4MoeBackend.AITER_MXFP4_MXFP4, + Mxfp4MoeBackend.RDNA3_MXFP4, Mxfp4MoeBackend.TRITON, Mxfp4MoeBackend.FLASHINFER_CUTLASS_MXFP4_BF16, Mxfp4MoeBackend.FLASHINFER_CUTLASS_MXFP4_MXFP8, @@ -341,7 +340,7 @@ def _get_priority_backends() -> list[Mxfp4MoeBackend]: backend-level ``is_supported_config`` check filters by device capability). """ if current_platform.is_rocm(): - return [Mxfp4MoeBackend.RDNA3_MXFP4, Mxfp4MoeBackend.AITER_MXFP4_BF16] + return [Mxfp4MoeBackend.AITER_MXFP4_BF16] if current_platform.is_xpu(): return [Mxfp4MoeBackend.XPU] _AVAILABLE_BACKENDS = [ @@ -549,21 +548,10 @@ def select_mxfp4_moe_backend( ) if current_platform.is_rocm(): - # ROCm has no native FP4 compute, so activation-quantized MXFP4 - # (W4A4/W4A8) degrades to weight-only. On gfx1100 the native HIP - # kernel handles that faster than the Triton-unfused fallback. - from vllm.model_executor.layers.fused_moe.experts.rdna3_mxfp4_moe import ( - RDNA3Mxfp4Experts, - ) - - backend = ( - Mxfp4MoeBackend.RDNA3_MXFP4 - if RDNA3Mxfp4Experts._supports_current_device() - else Mxfp4MoeBackend.TRITON_UNFUSED - ) + backend = Mxfp4MoeBackend.TRITON_UNFUSED logger.info_once(_make_log_backend(backend)) return _return_or_raise( - backend, + Mxfp4MoeBackend.TRITON_UNFUSED, config, kMxfp4Static, None, diff --git a/vllm/model_executor/layers/quantization/quark/quark_moe.py b/vllm/model_executor/layers/quantization/quark/quark_moe.py index 70b1e25959e1..010d8c8f7f55 100644 --- a/vllm/model_executor/layers/quantization/quark/quark_moe.py +++ b/vllm/model_executor/layers/quantization/quark/quark_moe.py @@ -1013,9 +1013,13 @@ def __init__( moe, activation_key=kFp8StaticTensorSym ) elif self.ocp_mx_scheme == "w_mxfp4_a_mxfp4": - # W4A4: MXFP4 weights + MXFP4 activations + # W4A4: MXFP4 weights + MXFP4 activations. Platforms without native + # FP4 (e.g. gfx1100) compute this weight-only with bf16 activations, + # so request a weight-only backend instead of the mxfp4-activation + # variant. + act_key = kMxfp4Dynamic if current_platform.supports_mx() else None self.mxfp4_backend, self.experts_cls = select_mxfp4_moe_backend( - moe, activation_key=kMxfp4Dynamic + moe, activation_key=act_key ) # Validation for unsupported schemes diff --git a/vllm/model_executor/layers/quantization/quark/schemes/quark_ocp_mx.py b/vllm/model_executor/layers/quantization/quark/schemes/quark_ocp_mx.py index 05c8e97f82df..727b8d3df248 100644 --- a/vllm/model_executor/layers/quantization/quark/schemes/quark_ocp_mx.py +++ b/vllm/model_executor/layers/quantization/quark/schemes/quark_ocp_mx.py @@ -294,9 +294,9 @@ def process_weights_after_loading(self, layer: torch.nn.Module) -> None: layer.weight_scale.data, requires_grad=False ) self.rdna3_kernel.process_weights_after_loading(layer) - layer._rdna3 = True return - layer._rdna3 = False + # Shapes not supported by the HIP kernel; fall back to emulation. + self.rdna3_kernel = None if self.emulate: if self.dynamic_mxfp4_quant: @@ -393,8 +393,7 @@ def apply_weights( x: torch.Tensor, bias: torch.Tensor | None = None, ) -> torch.Tensor: - if getattr(layer, "_rdna3", False): - assert self.rdna3_kernel is not None + if self.rdna3_kernel is not None: return self.rdna3_kernel.apply_weights(layer, x, bias) if self.emulate: dq_w = self.dequant_func(layer.weight, layer.weight_scale, x.dtype) From 1608b9960aa9cc4c0e5984291a9ddca5f8cb8599 Mon Sep 17 00:00:00 2001 From: JartX Date: Fri, 26 Jun 2026 22:34:11 +0200 Subject: [PATCH 24/36] [ROCm][gfx1100] MXFP4 MoE: trim verbose comments (no functional change) Comment-only line trim of the previous commits per review feedback: shorten the W4A4 weight-only downgrade, the Quark dense-kernel setup, the TP-reduce note and the gate/up de-interleave note to the essential rationale. No code changes. Signed-off-by: JartX --- .../layers/fused_moe/experts/rdna3_mxfp4_moe.py | 10 ++++------ vllm/model_executor/layers/fused_moe/oracle/mxfp4.py | 6 ++---- .../layers/quantization/quark/quark_moe.py | 5 +---- .../layers/quantization/quark/schemes/quark_ocp_mx.py | 6 ++---- 4 files changed, 9 insertions(+), 18 deletions(-) diff --git a/vllm/model_executor/layers/fused_moe/experts/rdna3_mxfp4_moe.py b/vllm/model_executor/layers/fused_moe/experts/rdna3_mxfp4_moe.py index d84c8b3d027a..da757bf2c2e1 100644 --- a/vllm/model_executor/layers/fused_moe/experts/rdna3_mxfp4_moe.py +++ b/vllm/model_executor/layers/fused_moe/experts/rdna3_mxfp4_moe.py @@ -207,12 +207,10 @@ def apply( else: apply_moe_activation(activation, act_out, w1_out) - # down GEMM. The kernel can fuse the top-k reduction + weight multiply - # into its epilogue (output_topk), but that path produces wrong results - # under tensor parallelism (verified: correct at TP1, garbage at TP2 -- - # the per-rank partial down-proj is later all-reduced by the layer). - # A per-expert bias must also be added before the weighted reduction. - # In both cases write unreduced rows and reduce in Python instead. + # down GEMM. The kernel's fused output_topk reduce is wrong under TP + # (each rank's down-proj is a partial the layer all-reduces afterwards); + # a per-expert bias likewise needs adding before the reduction. In both + # cases write unreduced rows and reduce in Python instead. unfused = w2_bias is not None or get_tensor_model_parallel_world_size() > 1 if unfused: w2_out = torch.zeros(total, hidden, dtype=dtype, device=device) diff --git a/vllm/model_executor/layers/fused_moe/oracle/mxfp4.py b/vllm/model_executor/layers/fused_moe/oracle/mxfp4.py index e4adcaeae33b..975f657851e0 100644 --- a/vllm/model_executor/layers/fused_moe/oracle/mxfp4.py +++ b/vllm/model_executor/layers/fused_moe/oracle/mxfp4.py @@ -753,10 +753,8 @@ def convert_gpt_oss_weight_to_mxfp4_moe_kernel_format( ) elif mxfp4_backend == Mxfp4MoeBackend.RDNA3_MXFP4: - # [E, K/8, N] uint32 + [E, K/32, N] uint8. GPT-OSS stores gate/up - # interleaved in a fused tensor (de-interleave so the gate_up output is - # contiguous [gate || up]); standard checkpoints with separate - # gate_proj/up_proj are already [gate_all, up_all] after _load_w13. + # Only GPT-OSS stores gate/up interleaved in a fused tensor; separate + # gate_proj/up_proj checkpoints are already contiguous after _load_w13. from vllm.config import get_current_vllm_config from vllm.model_executor.layers.fused_moe.experts.rdna3_mxfp4_moe import ( repack_experts_rdna3, diff --git a/vllm/model_executor/layers/quantization/quark/quark_moe.py b/vllm/model_executor/layers/quantization/quark/quark_moe.py index 010d8c8f7f55..b8fd0428ecec 100644 --- a/vllm/model_executor/layers/quantization/quark/quark_moe.py +++ b/vllm/model_executor/layers/quantization/quark/quark_moe.py @@ -1013,10 +1013,7 @@ def __init__( moe, activation_key=kFp8StaticTensorSym ) elif self.ocp_mx_scheme == "w_mxfp4_a_mxfp4": - # W4A4: MXFP4 weights + MXFP4 activations. Platforms without native - # FP4 (e.g. gfx1100) compute this weight-only with bf16 activations, - # so request a weight-only backend instead of the mxfp4-activation - # variant. + # W4A4, or weight-only when the platform has no native FP4 (gfx1100). act_key = kMxfp4Dynamic if current_platform.supports_mx() else None self.mxfp4_backend, self.experts_cls = select_mxfp4_moe_backend( moe, activation_key=act_key diff --git a/vllm/model_executor/layers/quantization/quark/schemes/quark_ocp_mx.py b/vllm/model_executor/layers/quantization/quark/schemes/quark_ocp_mx.py index 727b8d3df248..6c728fb445d8 100644 --- a/vllm/model_executor/layers/quantization/quark/schemes/quark_ocp_mx.py +++ b/vllm/model_executor/layers/quantization/quark/schemes/quark_ocp_mx.py @@ -207,10 +207,8 @@ def __init__( self.input_dtype != "mxfp4" or self.weight_dtype != "mxfp4" ) - # gfx1100 has no native FP4, so statically-quantized weight-only MXFP4 - # linears run on the native RDNA3 HIP GEMM instead of dequant-to-bf16 - # emulation (activations stay in bf16 == W4A16). The kernel is chosen - # per-layer in process_weights_after_loading once shapes are known. + # On gfx1100 run weight-only MXFP4 linears on the native RDNA3 GEMM + # instead of dequant emulation; chosen per-layer once shapes are known. self.rdna3_kernel = None if self.weight_dtype == "mxfp4" and not self.dynamic_mxfp4_quant: from vllm.model_executor.kernels.linear.mxfp4.rocm import ( From daadfeebb9085342cb6f5a3a7dc4225c862cc127 Mon Sep 17 00:00:00 2001 From: JartX Date: Thu, 13 Aug 2026 14:30:47 +0200 Subject: [PATCH 25/36] Register the RDNA3 MXFP4 linear kernel in the ROCm priority list MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit The merge left ``_POSSIBLE_MXFP4_KERNELS`` with two ``PlatformEnum.ROCM`` entries — upstream's [AiterMxfp4LinearKernel, EmulationMxfp4LinearKernel] and ours [Rdna3MxFp4LinearKernel]. In a dict literal the last one wins, so ROCm would have lost the aiter and emulation backends entirely. Merge them into one list with the RDNA3 kernel first: its ``is_supported`` gates on gfx1100, so every other ROCm target falls through to aiter / emulation as before. Signed-off-by: JartX --- vllm/model_executor/kernels/linear/__init__.py | 4 +--- 1 file changed, 1 insertion(+), 3 deletions(-) diff --git a/vllm/model_executor/kernels/linear/__init__.py b/vllm/model_executor/kernels/linear/__init__.py index 7fe07febfb4d..c56a6e6fd0f5 100644 --- a/vllm/model_executor/kernels/linear/__init__.py +++ b/vllm/model_executor/kernels/linear/__init__.py @@ -535,15 +535,13 @@ def _resolve_backend_kernels( EmulationMxfp4LinearKernel, ], PlatformEnum.ROCM: [ + Rdna3MxFp4LinearKernel, AiterMxfp4LinearKernel, EmulationMxfp4LinearKernel, ], PlatformEnum.XPU: [ XPUMxFp4LinearKernel, ], - PlatformEnum.ROCM: [ - Rdna3MxFp4LinearKernel, - ], } # TODO make all kernels inherit from MMLinearKernel From 6e5392e570735c793038540e6c87f97c2a3b159e Mon Sep 17 00:00:00 2001 From: JartX Date: Thu, 13 Aug 2026 14:34:08 +0200 Subject: [PATCH 26/36] Make the RDNA3 MXFP4 linear kernel follow the weight-only contract Upstream's kernel-selection tests split MXFP4 linear kernels in two: true-W4A4 kernels (FlashInfer / XPU / aiter) that quantize activations themselves and must reject an unset activation key, and weight-only (A16) kernels (Marlin / Humming) that accept None or MXFP4-dynamic and log a warning saying the requested activation quantization is ignored. gfx1100 has no native FP4 matmul, so the RDNA3 kernel is weight-only: adopt the Marlin/Humming can_implement verbatim and add it to the _WEIGHT_ONLY_KERNELS list so the existing tests cover it. This replaces the branch's old `act_key = ... if supports_mx() else None` decision for the dense path, which upstream has since settled the same way. Signed-off-by: JartX --- .../quantization/test_mxfp4_kernel_selection.py | 7 ++++++- vllm/model_executor/kernels/linear/mxfp4/rocm.py | 12 ++++++++++++ 2 files changed, 18 insertions(+), 1 deletion(-) diff --git a/tests/kernels/quantization/test_mxfp4_kernel_selection.py b/tests/kernels/quantization/test_mxfp4_kernel_selection.py index 28ce615ce655..6ef117a15dd2 100644 --- a/tests/kernels/quantization/test_mxfp4_kernel_selection.py +++ b/tests/kernels/quantization/test_mxfp4_kernel_selection.py @@ -18,6 +18,7 @@ MarlinMxFp4LinearKernel, MxFp4LinearKernel, MxFp4LinearLayerConfig, + Rdna3MxFp4LinearKernel, XPUMxFp4LinearKernel, init_mxfp4_linear_kernel, register_linear_kernel, @@ -44,7 +45,11 @@ # Weight-only (A16) kernels: they never quantize activations. They still accept # MXFP4 activation keys as an intentional compatibility fallback. -_WEIGHT_ONLY_KERNELS = [MarlinMxFp4LinearKernel, HummingMxFp4LinearKernel] +_WEIGHT_ONLY_KERNELS = [ + MarlinMxFp4LinearKernel, + HummingMxFp4LinearKernel, + Rdna3MxFp4LinearKernel, +] def test_can_implement_is_abstract(): diff --git a/vllm/model_executor/kernels/linear/mxfp4/rocm.py b/vllm/model_executor/kernels/linear/mxfp4/rocm.py index 200b8f92d81d..f238e07147b4 100644 --- a/vllm/model_executor/kernels/linear/mxfp4/rocm.py +++ b/vllm/model_executor/kernels/linear/mxfp4/rocm.py @@ -4,10 +4,14 @@ import torch from torch.nn.parameter import Parameter +from vllm.logger import init_logger +from vllm.model_executor.layers.quantization.utils.quant_utils import kMxfp4Dynamic from vllm.platforms import current_platform from .base import MxFp4LinearKernel, MxFp4LinearLayerConfig +logger = init_logger(__name__) + _MXFP4_GROUP_SIZE = 32 # Fake/meta impl so the op is traceable under torch.compile + CUDA graphs. @@ -44,6 +48,14 @@ def is_supported( @classmethod def can_implement(cls, config: MxFp4LinearLayerConfig) -> tuple[bool, str | None]: + if config.activation_quant_key not in (None, kMxfp4Dynamic): + return False, "only supports MXFP4 dynamic or unquantized activations" + if config.activation_quant_key is not None: + logger.warning_once( + "Rdna3MxFp4LinearKernel is a weight-only (A16) kernel; " + "the requested activation quantization (%s) is ignored.", + config.activation_quant_key, + ) return True, None def process_weights_after_loading(self, layer: torch.nn.Module) -> None: From 7105001f294d548b677b6306cb35181642febd08 Mon Sep 17 00:00:00 2001 From: JartX Date: Thu, 13 Aug 2026 19:14:04 +0200 Subject: [PATCH 27/36] Drop the redundant MXFP4 op-availability guards Both the linear kernel and the MoE experts backend checked that torch.ops._rocm_C exposes their op before declaring support. CMake builds mxfp4_gemm_rdna3 / moe_mxfp4_gemm_rdna3 for every gfx11xx target, so on gfx1100 the op is always there and the check never fires; on any other target the on_gfx1100() guard has already returned False. Addresses two review comments asking when the op would not be built. Signed-off-by: JartX --- vllm/model_executor/kernels/linear/mxfp4/rocm.py | 2 -- .../layers/fused_moe/experts/rdna3_mxfp4_moe.py | 8 +------- 2 files changed, 1 insertion(+), 9 deletions(-) diff --git a/vllm/model_executor/kernels/linear/mxfp4/rocm.py b/vllm/model_executor/kernels/linear/mxfp4/rocm.py index f238e07147b4..f9cc954cb787 100644 --- a/vllm/model_executor/kernels/linear/mxfp4/rocm.py +++ b/vllm/model_executor/kernels/linear/mxfp4/rocm.py @@ -42,8 +42,6 @@ def is_supported( if not on_gfx1100(): return False, "RDNA3 gfx1100 WMMA required" - if not hasattr(torch.ops._rocm_C, "mxfp4_gemm_rdna3"): - return False, "_rocm_C.mxfp4_gemm_rdna3 op not built" return True, None @classmethod diff --git a/vllm/model_executor/layers/fused_moe/experts/rdna3_mxfp4_moe.py b/vllm/model_executor/layers/fused_moe/experts/rdna3_mxfp4_moe.py index da757bf2c2e1..fe1c0d4955c4 100644 --- a/vllm/model_executor/layers/fused_moe/experts/rdna3_mxfp4_moe.py +++ b/vllm/model_executor/layers/fused_moe/experts/rdna3_mxfp4_moe.py @@ -52,12 +52,6 @@ def _select_block_size_m(num_tokens: int, top_k: int, num_experts: int) -> int: return 4 if tokens_per_expert >= 1.5 else 1 -def _rocm_op_available() -> bool: - return hasattr(torch.ops, "_rocm_C") and hasattr( - torch.ops._rocm_C, "moe_mxfp4_gemm_rdna3" - ) - - def repack_experts_rdna3( packed: torch.Tensor, scale: torch.Tensor, deinterleave: bool ) -> tuple[torch.Tensor, torch.Tensor]: @@ -94,7 +88,7 @@ def activation_format() -> mk.FusedMoEActivationFormat: def _supports_current_device() -> bool: from vllm.platforms import current_platform - if not current_platform.is_rocm() or not _rocm_op_available(): + if not current_platform.is_rocm(): return False from vllm.platforms.rocm import on_gfx1100 From f92711adca2d6e6aee067916926fc5031f940138 Mon Sep 17 00:00:00 2001 From: JartX Date: Thu, 13 Aug 2026 19:14:11 +0200 Subject: [PATCH 28/36] Take the MXFP4 MoE gate/up layout as a parameter, not from the config The RDNA3 branch of convert_gpt_oss_weight_to_mxfp4_moe_kernel_format reached into get_current_vllm_config() to read model_type and decide whether w13 needs de-interleaving. Despite its name that converter is also called for non-GPT-OSS checkpoints (quark_moe.py), which is why the branch needed to distinguish them at all. Make the layout an explicit gate_up_interleaved argument instead. The default is True, so GPT-OSS callers are unchanged; QuarkOCP_MX_MoEMethod passes its own self.model_type, which upstream already computes. Signed-off-by: JartX --- .../layers/fused_moe/oracle/mxfp4.py | 18 ++++++++++-------- .../layers/quantization/quark/quark_moe.py | 1 + 2 files changed, 11 insertions(+), 8 deletions(-) diff --git a/vllm/model_executor/layers/fused_moe/oracle/mxfp4.py b/vllm/model_executor/layers/fused_moe/oracle/mxfp4.py index 7902492d90f2..79475c0ac235 100644 --- a/vllm/model_executor/layers/fused_moe/oracle/mxfp4.py +++ b/vllm/model_executor/layers/fused_moe/oracle/mxfp4.py @@ -716,6 +716,7 @@ def convert_gpt_oss_weight_to_mxfp4_moe_kernel_format( w13_input_scale: torch.Tensor | None = None, w2_input_scale: torch.Tensor | None = None, _cache_permute_indices: dict[torch.Size, torch.Tensor] | None = None, + gate_up_interleaved: bool = True, ) -> tuple[ torch.Tensor, torch.Tensor, @@ -724,7 +725,13 @@ def convert_gpt_oss_weight_to_mxfp4_moe_kernel_format( torch.Tensor | None, torch.Tensor | None, ]: - """Convert loaded weights into backend-specific kernel format.""" + """Convert loaded weights into backend-specific kernel format. + + ``gate_up_interleaved`` describes the loaded ``w13`` layout: GPT-OSS fuses + gate/up row-interleaved ([g0, u0, g1, u1, ...]), while checkpoints with + separate gate_proj/up_proj tensors are already contiguous after ``_load_w13``. + Backends that only ever see GPT-OSS weights ignore it. + """ if mxfp4_backend == Mxfp4MoeBackend.DEEPGEMM_MXFP4: w13_weight_scale, w2_weight_scale = _pack_deepgemm_mxfp4_scales( @@ -784,22 +791,17 @@ def convert_gpt_oss_weight_to_mxfp4_moe_kernel_format( ) elif mxfp4_backend == Mxfp4MoeBackend.RDNA3_MXFP4: - # Only GPT-OSS stores gate/up interleaved in a fused tensor; separate - # gate_proj/up_proj checkpoints are already contiguous after _load_w13. - from vllm.config import get_current_vllm_config from vllm.model_executor.layers.fused_moe.experts.rdna3_mxfp4_moe import ( repack_experts_rdna3, ) - model_type = get_current_vllm_config().model_config.hf_config.model_type - deint = model_type == "gpt_oss" b_q13, b_s13 = repack_experts_rdna3( - w13_weight.data, w13_weight_scale.data, deinterleave=deint + w13_weight.data, w13_weight_scale.data, deinterleave=gate_up_interleaved ) b_q2, b_s2 = repack_experts_rdna3( w2_weight.data, w2_weight_scale.data, deinterleave=False ) - if deint and w13_bias is not None: + if gate_up_interleaved and w13_bias is not None: w13_bias = torch.cat( [w13_bias.data[:, ::2], w13_bias.data[:, 1::2]], dim=1 ).contiguous() diff --git a/vllm/model_executor/layers/quantization/quark/quark_moe.py b/vllm/model_executor/layers/quantization/quark/quark_moe.py index 6c2ca44ed045..609ebe2445b3 100644 --- a/vllm/model_executor/layers/quantization/quark/quark_moe.py +++ b/vllm/model_executor/layers/quantization/quark/quark_moe.py @@ -1299,6 +1299,7 @@ def _setup_kernel(self, layer: RoutedExperts): w2_bias=w2_bias, w13_input_scale=layer.w13_input_scale, w2_input_scale=layer.w2_input_scale, + gate_up_interleaved=self.model_type == "gpt_oss", ) ) From 47495e26c1ed477b8b77f44a57c76f1c5e199f60 Mon Sep 17 00:00:00 2001 From: JartX Date: Thu, 13 Aug 2026 19:14:20 +0200 Subject: [PATCH 29/36] Route the compressed-tensors MXFP4 MoE repack through the oracle CompressedTensorsW4A4Mxfp4MoEMethod repacked the RDNA3 weights inline with its own repack_experts_rdna3 calls, duplicating what the oracle already does for the quark and GPT-OSS checkpoint formats. Call convert_gpt_oss_weight_to_mxfp4_moe_kernel_format with gate_up_interleaved=False instead (compressed-tensors keeps gate_proj and up_proj separate, so w13 is contiguous after _load_w13), leaving the repack logic in one place. Verified bit-identical to the inline path on synthetic tensors: both produce the same w13/w2 int32 and the same uint8 scales, biases untouched. Addresses the review comment asking this to use the shared conversion function. The backend *selection* in the same file still bypasses select_mxfp4_moe_backend; that needs CutlassExpertsMxfp4 registering in the oracle first and is left to the follow-up PR. Signed-off-by: JartX --- .../compressed_tensors_moe_w4a4_mxfp4.py | 29 +++++++++++-------- 1 file changed, 17 insertions(+), 12 deletions(-) diff --git a/vllm/model_executor/layers/quantization/compressed_tensors/compressed_tensors_moe/compressed_tensors_moe_w4a4_mxfp4.py b/vllm/model_executor/layers/quantization/compressed_tensors/compressed_tensors_moe/compressed_tensors_moe_w4a4_mxfp4.py index 120fabc48eaa..1c9808845daf 100644 --- a/vllm/model_executor/layers/quantization/compressed_tensors/compressed_tensors_moe/compressed_tensors_moe_w4a4_mxfp4.py +++ b/vllm/model_executor/layers/quantization/compressed_tensors/compressed_tensors_moe/compressed_tensors_moe_w4a4_mxfp4.py @@ -23,13 +23,13 @@ ) from vllm.model_executor.layers.fused_moe.experts.rdna3_mxfp4_moe import ( RDNA3Mxfp4Experts, - repack_experts_rdna3, ) from vllm.model_executor.layers.fused_moe.experts.xpu_moe import ( XPUExpertsMxFp4, ) from vllm.model_executor.layers.fused_moe.oracle.mxfp4 import ( Mxfp4MoeBackend, + convert_gpt_oss_weight_to_mxfp4_moe_kernel_format, make_mxfp4_moe_kernel, make_mxfp4_moe_quant_config, ) @@ -199,18 +199,23 @@ def process_weights_after_loading(self, layer: RoutedExperts) -> None: elif current_platform.is_xpu(): pass elif self.mxfp4_backend == Mxfp4MoeBackend.RDNA3_MXFP4: - # Repack to the [E, K/8, N] uint32 + [E, K/32, N] layout the HIP - # kernel reads (no gate/up de-interleave for compressed-tensors). - b_q13, b_s13 = repack_experts_rdna3( - layer.w13_weight.data, layer.w13_weight_scale.data, False + # compressed-tensors keeps gate_proj/up_proj separate, so w13 is + # already contiguous after _load_w13. + w13, w2, w13_scale, w2_scale, _, _ = ( + convert_gpt_oss_weight_to_mxfp4_moe_kernel_format( + mxfp4_backend=self.mxfp4_backend, + layer=layer, + w13_weight=layer.w13_weight, + w2_weight=layer.w2_weight, + w13_weight_scale=layer.w13_weight_scale, + w2_weight_scale=layer.w2_weight_scale, + gate_up_interleaved=False, + ) ) - b_q2, b_s2 = repack_experts_rdna3( - layer.w2_weight.data, layer.w2_weight_scale.data, False - ) - layer.w13_weight = torch.nn.Parameter(b_q13, requires_grad=False) - layer.w13_weight_scale = torch.nn.Parameter(b_s13, requires_grad=False) - layer.w2_weight = torch.nn.Parameter(b_q2, requires_grad=False) - layer.w2_weight_scale = torch.nn.Parameter(b_s2, requires_grad=False) + layer.w13_weight = torch.nn.Parameter(w13, requires_grad=False) + layer.w13_weight_scale = torch.nn.Parameter(w13_scale, requires_grad=False) + layer.w2_weight = torch.nn.Parameter(w2, requires_grad=False) + layer.w2_weight_scale = torch.nn.Parameter(w2_scale, requires_grad=False) else: logger.warning_once( "Your GPU does not have native support for FP4 computation " From fd7833ed64f0d3796131ef2ba9bcf14fb6b28ed9 Mon Sep 17 00:00:00 2001 From: JartX Date: Sun, 6 Sep 2026 09:12:14 +0200 Subject: [PATCH 30/36] [ROCm][MXFP4] Make the RDNA3 MXFP4 experts follow the modular-kernel contract Four things the class was doing its own way: - moe_problem_size() was inherited, so N was read off a packed dimension (K // 8). Nothing consumes N today, so it was harmless by luck; override it like MarlinExperts does and report the real N/K. - workspace_shapes() asked for no workspaces and apply() allocated three tensors per call. Take the framework workspaces instead: the activation output and both GEMM outputs come out of workspace2, since the modular kernel provisions `output` out of workspace13 and the down GEMM reads the activation while writing `output`. The gate/up buffer is dead after the activation, so the unfused down GEMM reuses that region. - the clamped SwiGLU-OAI was reimplemented in Python. The repack de-interleaves gate/up, which is exactly what MoEActivation SWIGLUOAI_UNINTERLEAVE describes, so go through self.activation() and let the fused kernel read alpha/beta/clamp from the activation config. Same math (gate clamped above, up clamped both ways, silu(alpha*gate) * (up + beta)), computed in fp32 instead of the activation dtype. - the per-expert bias indexed w1_bias/w2_bias with global expert ids. Under expert parallelism those are not local indices: the lookup wrapped around or went out of range, and rows routed to another rank -- which the GEMM skips -- still got a bias added. Map through expert_map and zero the invalid rows. Also reads tp_size from the parallel config instead of calling into the distributed state on every forward. Signed-off-by: JartX --- .../fused_moe/experts/rdna3_mxfp4_moe.py | 109 +++++++++++------- 1 file changed, 68 insertions(+), 41 deletions(-) diff --git a/vllm/model_executor/layers/fused_moe/experts/rdna3_mxfp4_moe.py b/vllm/model_executor/layers/fused_moe/experts/rdna3_mxfp4_moe.py index fe1c0d4955c4..a6a7d7028699 100644 --- a/vllm/model_executor/layers/fused_moe/experts/rdna3_mxfp4_moe.py +++ b/vllm/model_executor/layers/fused_moe/experts/rdna3_mxfp4_moe.py @@ -17,10 +17,8 @@ import vllm.model_executor.layers.fused_moe.modular_kernel as mk from vllm import _custom_ops as ops -from vllm.distributed import get_tensor_model_parallel_world_size from vllm.model_executor.layers.fused_moe.activation import ( MoEActivation, - apply_moe_activation, ) from vllm.model_executor.layers.fused_moe.config import ( FusedMoEParallelConfig, @@ -119,6 +117,21 @@ def finalize_weight_and_reduce_impl(self) -> mk.TopKWeightAndReduce: # prepare/finalize step must not reduce again. return TopKWeightAndReduceNoOP() + def moe_problem_size( + self, + a1: torch.Tensor, + w1: torch.Tensor, + w2: torch.Tensor, + topk_ids: torch.Tensor, + ) -> tuple[int, int, int, int, int]: + # Weights are packed along the reduction dim -- w1 is + # [E, K // 8, 2 * N] and w2 is [E, N // 8, K] -- so neither N nor K can + # be read off a trailing dimension the way the base implementation does. + assert w1.dim() == 3 and w2.dim() == 3 + assert a1.dim() == 2 + assert topk_ids.size(0) == a1.size(0), f"{topk_ids.size(0)} != {a1.size(0)}" + return w1.size(0), a1.size(0), w1.size(2) // 2, a1.size(-1), topk_ids.size(1) + def workspace_shapes( self, M: int, @@ -130,9 +143,34 @@ def workspace_shapes( expert_tokens_meta: mk.ExpertTokensMetadata | None, activation: MoEActivation, ) -> tuple[tuple[int, ...], tuple[int, ...], tuple[int, ...]]: - # apply() allocates its own scratch (shapes derived from the weights), - # so request no framework workspaces; output is the reduced [M, K]. - return ((0,), (0,), (M, K)) + # The modular kernel provisions `output` out of workspace13, so every + # buffer a GEMM reads while another writes `output` has to come out of + # workspace2: the activation output stays live across the down GEMM. + # The gate/up GEMM output is dead once the activation has run, so the + # unfused down GEMM reuses that region. + scratch = M * topk * (N + max(2 * N, K)) + return ((M, K), (scratch,), (M, K)) + + def _expert_bias_rows( + self, + bias: torch.Tensor, + topk_ids: torch.Tensor, + expert_map: torch.Tensor | None, + dtype: torch.dtype, + ) -> torch.Tensor: + """Per-row expert bias, [M * top_k, N]. + + ``topk_ids`` are global expert ids; the bias is stored per *local* + expert, so under EP they have to go through ``expert_map`` first. Rows + routed to another rank map to -1: the GEMM skips them, so their bias + must be zero rather than a wrapped-around lookup. + """ + ids = topk_ids if expert_map is None else expert_map[topk_ids] + ids = ids.reshape(-1) + rows = bias[ids.clamp_min(0).long()].to(dtype) + if expert_map is not None: + rows = rows * (ids >= 0).unsqueeze(1).to(dtype) + return rows def apply( self, @@ -153,31 +191,30 @@ def apply( apply_router_weight_on_input: bool, ) -> None: qc = self.quant_config - M = hidden_states.shape[0] - top_k = topk_ids.shape[1] + E, M, N, K, top_k = self.moe_problem_size(hidden_states, w1, w2, topk_ids) + n_gate_up = 2 * N total = M * top_k - n_gate_up = w1.shape[2] # [E, K/8, 2*intermediate] - intermediate = n_gate_up // 2 - hidden = w2.shape[2] dtype = hidden_states.dtype device = hidden_states.device - gne = global_num_experts if global_num_experts > 0 else w1.shape[0] + gne = global_num_experts if global_num_experts > 0 else E block_size_m = _select_block_size_m(M, top_k, gne) sti, eid, ntp = moe_align_block_size(topk_ids, block_size_m, gne, expert_map) + # workspace2 holds the activation output (live across the down GEMM) + # followed by the GEMM output region, which the two GEMMs share. + scratch = workspace2.view(-1) + act_out = scratch[: total * N].view(total, N) + gemm_out = scratch[total * N :] + w1_out = gemm_out[: total * n_gate_up].view(total, n_gate_up) + topk_w_f32 = topk_weights.reshape(-1).float() empty_tw = torch.empty(0, device=device) w1_bias = qc.w1_bias w2_bias = qc.w2_bias - flat_experts = ( - topk_ids.reshape(-1).long() - if (w1_bias is not None or w2_bias is not None) - else None - ) # gate_up GEMM (no reduction); router weight optionally folded in here. - w1_out = torch.zeros(total, n_gate_up, dtype=dtype, device=device) + w1_out.zero_() ops.moe_mxfp4_gemm_rdna3( hidden_states, w1_out, @@ -193,21 +230,23 @@ def apply( 0, ) if w1_bias is not None: - w1_out = w1_out + w1_bias[flat_experts].to(dtype) + w1_out.add_(self._expert_bias_rows(w1_bias, topk_ids, expert_map, dtype)) - act_out = torch.empty(total, intermediate, dtype=dtype, device=device) + # The repack de-interleaves GPT-OSS's gate/up rows, so the contiguous + # SwiGLU-OAI variant is the matching primitive; both read alpha/beta/ + # clamp from the activation config. if activation == MoEActivation.SWIGLUOAI: - self._swiglu_oai(w1_out, act_out, qc) - else: - apply_moe_activation(activation, act_out, w1_out) + activation = MoEActivation.SWIGLUOAI_UNINTERLEAVE + self.activation(activation, act_out, w1_out) # down GEMM. The kernel's fused output_topk reduce is wrong under TP # (each rank's down-proj is a partial the layer all-reduces afterwards); # a per-expert bias likewise needs adding before the reduction. In both # cases write unreduced rows and reduce in Python instead. - unfused = w2_bias is not None or get_tensor_model_parallel_world_size() > 1 + unfused = w2_bias is not None or self.moe_config.moe_parallel_config.tp_size > 1 if unfused: - w2_out = torch.zeros(total, hidden, dtype=dtype, device=device) + w2_out = gemm_out[: total * K].view(total, K) + w2_out.zero_() ops.moe_mxfp4_gemm_rdna3( act_out, w2_out, @@ -223,10 +262,12 @@ def apply( 0, ) if w2_bias is not None: - w2_out = w2_out + w2_bias[flat_experts].to(dtype) + w2_out.add_( + self._expert_bias_rows(w2_bias, topk_ids, expert_map, dtype) + ) if not apply_router_weight_on_input: - w2_out = w2_out * topk_weights.reshape(-1, 1).to(dtype) - output.copy_(w2_out.view(M, top_k, hidden).sum(dim=1)) + w2_out.mul_(topk_weights.reshape(-1, 1).to(dtype)) + torch.sum(w2_out.view(M, top_k, K), dim=1, out=output) else: output.zero_() ops.moe_mxfp4_gemm_rdna3( @@ -243,17 +284,3 @@ def apply( not apply_router_weight_on_input, top_k, ) - - @staticmethod - def _swiglu_oai(x: torch.Tensor, out: torch.Tensor, qc) -> None: - # gate||up contiguous; OAI clamped SwiGLU with quant-config params. - alpha = qc.gemm1_alpha if qc.gemm1_alpha is not None else 1.702 - beta = qc.gemm1_beta if qc.gemm1_beta is not None else 1.0 - limit = qc.gemm1_clamp_limit - d = x.shape[-1] // 2 - gate = x[..., :d] - up = x[..., d:] - if limit is not None: - gate = torch.clamp(gate, max=limit) - up = torch.clamp(up, min=-limit, max=limit) - out.copy_(gate * torch.sigmoid(alpha * gate) * (up + beta)) From 10bb691ec18d82ce50147fadae42e0404cc6930d Mon Sep 17 00:00:00 2001 From: JartX Date: Sun, 6 Sep 2026 10:07:26 +0200 Subject: [PATCH 31/36] [ROCm][MXFP4] Bound the RDNA3 MXFP4 GEMM test by the rounding it can't avoid The dense test compared against a fixed atol/rtol of 2e-2. That is below the floor for the case it generates: full-range E2M1 codes with E8M0 exponents spread over 2^-3..2^3, so at K=4096 the accumulated magnitude sum|x||w| is ~2e3 while the sums themselves cancel down to ~1e2. Rounding the *exact* result to bf16 already misses by ~0.5 there, so 8 of the 24 cases failed on gfx1100 (up to 1.5 absolute) with a kernel that is fine. Measured against an fp64 reference, the kernel lands at 0.06-0.08 of eps * sum|x||w| across dtypes and shapes -- roughly twice the unavoidable output-rounding floor, and it tracks the input magnitudes exactly as a rounding error should. So bound it by that quantity instead, at 0.5 * eps (about 7x slack), which still catches anything off by more than a percent of the accumulated magnitude. 24/24 pass on a 7900 XTX. Signed-off-by: JartX --- .../kernels/quantization/test_mxfp4_rdna3.py | 19 ++++++++++++++++--- 1 file changed, 16 insertions(+), 3 deletions(-) diff --git a/tests/kernels/quantization/test_mxfp4_rdna3.py b/tests/kernels/quantization/test_mxfp4_rdna3.py index acfc8e6d94e7..956133774b30 100644 --- a/tests/kernels/quantization/test_mxfp4_rdna3.py +++ b/tests/kernels/quantization/test_mxfp4_rdna3.py @@ -54,11 +54,24 @@ def test_mxfp4_wmma_matches_dequant_reference(dtype, M, N, K): x = torch.randn(M, K, dtype=dtype, device=dev) * 0.1 w_deq = _dequant_ref(w_packed, scale_e8m0) # [N, K] fp32 - ref = (x.float() @ w_deq.t()).to(dtype) + ref = (x.double() @ w_deq.double().t()).float() b_q, b_scale = _repack(w_packed, scale_e8m0) out = torch.ops._rocm_C.mxfp4_gemm_rdna3(x, b_q, b_scale) assert out.shape == (M, N) - # bf16/fp16 accumulate-in-fp32; tolerate dtype rounding of the reduction. - torch.testing.assert_close(out.float(), ref.float(), atol=2e-2, rtol=2e-2) + # A K-term dot product rounds with the *accumulated* magnitude, not with the + # result: these are full-range E2M1 codes, so at K=4096 sum|x||w| is ~2e3 + # while the sums themselves cancel down to ~1e2. Rounding the exact result + # to bf16 already misses by ~0.5 there, so a fixed atol cannot work. Bound + # the kernel by eps * sum|x||w| instead; measured slack is ~7x (the kernel + # lands at 0.06-0.08 of this bound across dtypes and shapes, roughly twice + # the unavoidable output-rounding floor). + accum = x.abs().float() @ w_deq.abs().t() + tol = 0.5 * torch.finfo(dtype).eps * accum + 1e-3 + err = (out.float() - ref).abs() + worst = int(err.argmax()) + assert bool((err <= tol).all()), ( + f"error {err.flatten()[worst]:.4f} exceeds the rounding bound " + f"{tol.flatten()[worst]:.4f} at flat index {worst}" + ) From 224c825953e4496ffc4d2187ddecc5af29b8ac4b Mon Sep 17 00:00:00 2001 From: JartX Date: Tue, 6 Oct 2026 11:50:37 +0200 Subject: [PATCH 32/36] [ROCm][MXFP4] RDNA3 MoE: align sizes to the MXFP4 block, not 256 RDNA3_MXFP4 fell into the generic ROCm branch of mxfp4_round_up_hidden_size_and_intermediate_size, which rounds to 256. The RDNA3 kernels bound-check N and K, so only the MXFP4 block (32) constrains them. With amd/Qwen3-Coder-Next-MXFP4 at TP4 the 256 round-up padded the per-rank intermediate size 128 -> 256, doubling the experts (20 GiB per GPU) and leaving no room for the KV cache. Tested on 4x RX 7900 XTX (gfx1100): - amd/Qwen3-Coder-Next-MXFP4 TP4: starts, GSM8K 87.0% (500 q). - openai/gpt-oss-20b TP1 (2880 no longer padded to 3072): GSM8K 86.5% (200 q), 442 tok/s, same as before. Signed-off-by: JartX --- vllm/model_executor/layers/fused_moe/oracle/mxfp4.py | 8 +++++++- 1 file changed, 7 insertions(+), 1 deletion(-) diff --git a/vllm/model_executor/layers/fused_moe/oracle/mxfp4.py b/vllm/model_executor/layers/fused_moe/oracle/mxfp4.py index da3cbee339f3..91273f74db3e 100644 --- a/vllm/model_executor/layers/fused_moe/oracle/mxfp4.py +++ b/vllm/model_executor/layers/fused_moe/oracle/mxfp4.py @@ -786,7 +786,13 @@ def mxfp4_round_up_hidden_size_and_intermediate_size( # b12x plans for the exact model dimensions. B12xExperts validates the # required MXFP4 block alignment before selecting the backend. return hidden_size, intermediate_size - if backend == Mxfp4MoeBackend.EMULATION: + if backend == Mxfp4MoeBackend.RDNA3_MXFP4: + # The RDNA3 kernels bound-check N and K, so only the MXFP4 block (K) and + # the 4-column vector loads (N) constrain them. The generic ROCm 256 + # round-up below would double the experts at TP4 (512 / 4 = 128). + intermediate_size = round_up(intermediate_size, OCP_MX_BLOCK_SIZE) + hidden_size = round_up(hidden_size, OCP_MX_BLOCK_SIZE) + elif backend == Mxfp4MoeBackend.EMULATION: # Emulation has no kernel tile; it only needs OCP MX block alignment so the # per-block scale buffers (`dim // OCP_MX_BLOCK_SIZE`) aren't floor-truncated # by a non-block-aligned TP/DP shard (e.g. 2880 // 4 = 720). From ff32da4c847545161725920295999f2ec1f4b3f3 Mon Sep 17 00:00:00 2001 From: JartX Date: Tue, 6 Oct 2026 12:08:02 +0200 Subject: [PATCH 33/36] [ROCm][MXFP4] Drop the blank line after the converter docstring (ruff) Signed-off-by: JartX --- vllm/model_executor/layers/fused_moe/oracle/mxfp4.py | 1 - 1 file changed, 1 deletion(-) diff --git a/vllm/model_executor/layers/fused_moe/oracle/mxfp4.py b/vllm/model_executor/layers/fused_moe/oracle/mxfp4.py index 91273f74db3e..8d450b5eac67 100644 --- a/vllm/model_executor/layers/fused_moe/oracle/mxfp4.py +++ b/vllm/model_executor/layers/fused_moe/oracle/mxfp4.py @@ -880,7 +880,6 @@ def convert_gpt_oss_weight_to_mxfp4_moe_kernel_format( separate gate_proj/up_proj tensors are already contiguous after ``_load_w13``. Backends that only ever see GPT-OSS weights ignore it. """ - if mxfp4_backend == Mxfp4MoeBackend.DEEPGEMM_MXFP4: w13_weight_scale, w2_weight_scale = _pack_deepgemm_mxfp4_scales( w13_weight, From e060ed0ab92a0bdbee2d74f9e1454f61d8153f97 Mon Sep 17 00:00:00 2001 From: JartX Date: Fri, 23 Oct 2026 00:24:20 +0200 Subject: [PATCH 34/36] [ROCm][MXFP4] Address review: fix fused down-GEMM rows, follow the oracle - RDNA3Mxfp4Experts: the fused (output_topk) down GEMM passed top_k as the input row divisor, but its input is the [M * top_k, N] activation buffer, so every token-expert slot read the wrong row. Pass 1, as the W4A16 backend does. Only MoE without bias at TP1 took this path. - SWIGLUOAI falls back to the constants swigluoai_and_mul hardcodes when the quant path does not forward them (Quark + gpt_oss.py). - Drop the gate_up_interleaved flag from the oracle converter; de-interleave when layer.activation is SWIGLUOAI. - Revert the quark_moe.py activation-key change: the checkpoint's activation scheme is honoured again and the oracle's BF16 fallback (with its warning) reaches RDNA3. --moe-backend rdna3 now maps to RDNA3_MXFP4 to opt in. - Compressed-tensors only picks RDNA3 for supported activations. - Move the dense op wrapper and its fake to _custom_ops.py; rename the linear kernel module to mxfp4/rdna3.py (*_hip* is gitignored for hipify output). - Host-side checks: same device, contiguity, N % 4, block_size_m > 0 and expert_ids length. - Tests: end-to-end RDNA3Mxfp4Experts.apply against a dequantized reference (fused and unfused paths, SiLU and SwiGLU-OAI), oracle selection on gfx1100, gfx1100-only skip. Signed-off-by: JartX --- csrc/rocm/moe_mxfp4_gemm_rdna3.cu | 13 ++ csrc/rocm/mxfp4_gemm_rdna3.cu | 6 + .../kernels/quantization/test_mxfp4_rdna3.py | 129 +++++++++++++++++- .../quantization/test_mxfp4_rdna3_dispatch.py | 108 +++++++++++++++ vllm/_custom_ops.py | 21 +++ .../model_executor/kernels/linear/__init__.py | 2 +- .../linear/mxfp4/{rocm.py => rdna3.py} | 15 +- .../fused_moe/experts/rdna3_mxfp4_moe.py | 45 +++++- .../layers/fused_moe/oracle/mxfp4.py | 15 +- .../compressed_tensors_moe_w4a4_mxfp4.py | 7 +- .../layers/quantization/quark/quark_moe.py | 6 +- 11 files changed, 324 insertions(+), 43 deletions(-) rename vllm/model_executor/kernels/linear/mxfp4/{rocm.py => rdna3.py} (84%) diff --git a/csrc/rocm/moe_mxfp4_gemm_rdna3.cu b/csrc/rocm/moe_mxfp4_gemm_rdna3.cu index 8cfb4261462b..ccf5bc79f226 100644 --- a/csrc/rocm/moe_mxfp4_gemm_rdna3.cu +++ b/csrc/rocm/moe_mxfp4_gemm_rdna3.cu @@ -532,6 +532,17 @@ void moe_mxfp4_gemm_rdna3(torch::Tensor a, torch::Tensor c, "a must be half or bfloat16"); TORCH_CHECK(b_scales_e8m0.scalar_type() == torch::kUInt8, "b_scales_e8m0 must be uint8 (E8M0)"); + for (const auto& t : {c, b_q_weight, b_scales_e8m0, topk_weights, + sorted_token_ids, expert_ids, num_tokens_post_padded}) { + TORCH_CHECK(t.numel() == 0 || t.device() == a.device(), + "all tensors must be on the same device"); + } + TORCH_CHECK(a.is_contiguous() && c.is_contiguous() && + b_q_weight.is_contiguous() && b_scales_e8m0.is_contiguous(), + "a, c and the weights must be contiguous"); + TORCH_CHECK(b_q_weight.size(2) % 4 == 0, + "N must be a multiple of 4 (vector loads, packed atomics)"); + TORCH_CHECK(block_size_m > 0, "block_size_m must be positive"); const at::cuda::OptionalCUDAGuard device_guard(device_of(a)); auto stream = at::cuda::getCurrentCUDAStream(); @@ -544,6 +555,8 @@ void moe_mxfp4_gemm_rdna3(torch::Tensor a, torch::Tensor c, int expert_scales_stride = (int)(b_scales_e8m0.size(1) * b_scales_e8m0.size(2)); int num_token_blocks = (int)(sorted_token_ids.size(0) / block_size_m); + TORCH_CHECK(expert_ids.size(0) >= num_token_blocks, + "expert_ids must have one entry per token block"); const float* topk_w_ptr = (topk_weights.numel() > 0) ? topk_weights.data_ptr() : nullptr; diff --git a/csrc/rocm/mxfp4_gemm_rdna3.cu b/csrc/rocm/mxfp4_gemm_rdna3.cu index 6d9048479995..2ac65b3630dc 100644 --- a/csrc/rocm/mxfp4_gemm_rdna3.cu +++ b/csrc/rocm/mxfp4_gemm_rdna3.cu @@ -960,6 +960,12 @@ torch::Tensor mxfp4_gemm_rdna3(torch::Tensor a, torch::Tensor b_q_weight, "b_q_weight must be (u)int32"); TORCH_CHECK(b_scales_e8m0.scalar_type() == torch::kUInt8, "b_scales_e8m0 must be uint8 (E8M0)"); + TORCH_CHECK( + b_q_weight.device() == a.device() && b_scales_e8m0.device() == a.device(), + "all tensors must be on the same device"); + TORCH_CHECK(a.is_contiguous() && b_q_weight.is_contiguous() && + b_scales_e8m0.is_contiguous(), + "all tensors must be contiguous"); const at::cuda::OptionalCUDAGuard device_guard(device_of(a)); auto stream = at::cuda::getCurrentCUDAStream(); diff --git a/tests/kernels/quantization/test_mxfp4_rdna3.py b/tests/kernels/quantization/test_mxfp4_rdna3.py index 956133774b30..dda0247c782c 100644 --- a/tests/kernels/quantization/test_mxfp4_rdna3.py +++ b/tests/kernels/quantization/test_mxfp4_rdna3.py @@ -1,20 +1,38 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project -"""Correctness test for the RDNA3 (gfx1100) MXFP4 weight-only WMMA GEMM. +"""Correctness tests for the RDNA3 (gfx1100) MXFP4 weight-only kernels: the +dense GEMM and the ``RDNA3Mxfp4Experts`` fused-MoE backend. Run inside the ROCm container after building _rocm_C: - .venv/bin/python -m pytest tests/kernels/quantization/test_mxfp4_rdna3_wmma.py -v + .venv/bin/python -m pytest tests/kernels/quantization/test_mxfp4_rdna3.py -v """ import pytest import torch +import torch.nn.functional as F from vllm.platforms import current_platform +if not current_platform.is_rocm(): + pytest.skip("RDNA3 MXFP4 kernels are ROCm-only", allow_module_level=True) + +from tests.kernels.moe.utils import make_dummy_moe_config # noqa: E402 +from vllm.model_executor.layers.fused_moe.activation import ( # noqa: E402 + MoEActivation, +) +from vllm.model_executor.layers.fused_moe.config import ( # noqa: E402 + mxfp4_w4a16_moe_quant_config, +) +from vllm.model_executor.layers.fused_moe.experts.rdna3_mxfp4_moe import ( # noqa: E402 + RDNA3Mxfp4Experts, + repack_experts_rdna3, +) +from vllm.platforms.rocm import on_gfx1100 # noqa: E402 + E2M1_MAG = torch.tensor([0.0, 0.5, 1.0, 1.5, 2.0, 3.0, 4.0, 6.0], dtype=torch.float32) pytestmark = pytest.mark.skipif( - not current_platform.is_rocm(), reason="RDNA3 MXFP4 WMMA kernel is ROCm-only" + not on_gfx1100(), reason="RDNA3 MXFP4 kernels require gfx1100" ) @@ -42,9 +60,6 @@ def _repack(w_packed, scale_e8m0): @pytest.mark.parametrize("M", [1, 16, 64, 128]) @pytest.mark.parametrize("N,K", [(64, 256), (256, 512), (128, 4096)]) def test_mxfp4_wmma_matches_dequant_reference(dtype, M, N, K): - if not hasattr(torch.ops._rocm_C, "mxfp4_gemm_rdna3"): - pytest.skip("_rocm_C.mxfp4_gemm_rdna3 not built") - torch.manual_seed(0) dev = "cuda" # Random E2M1 codes (full 0..255 byte = two 4-bit codes) and mild E8M0 @@ -75,3 +90,105 @@ def test_mxfp4_wmma_matches_dequant_reference(dtype, M, N, K): f"error {err.flatten()[worst]:.4f} exceeds the rounding bound " f"{tol.flatten()[worst]:.4f} at flat index {worst}" ) + + +def _swigluoai_interleaved(h: torch.Tensor) -> torch.Tensor: + gate = h[..., ::2].clamp(max=7.0) + up = h[..., 1::2].clamp(-7.0, 7.0) + return (up + 1.0) * gate * torch.sigmoid(1.702 * gate) + + +@pytest.mark.parametrize("M", [1, 3, 16, 64]) +@pytest.mark.parametrize( + "activation, with_bias", + [ + # No bias: the down GEMM reduces over top_k in its epilogue. + (MoEActivation.SILU, False), + # GPT-OSS: interleaved gate/up, per-expert bias, constants not passed + # in the quant config (as on the Quark path); reduced in Python. + (MoEActivation.SWIGLUOAI, True), + ], +) +def test_rdna3_mxfp4_experts_match_dequant_reference(M, activation, with_bias): + """``RDNA3Mxfp4Experts.apply`` against a dequantize-then-matmul MoE.""" + E, K, N, top_k = 8, 256, 128, 4 + dtype = torch.bfloat16 + dev = "cuda" + torch.manual_seed(0) + + def rand_mxfp4(rows, cols): + packed = torch.randint(0, 256, (E, rows, cols // 2), dtype=torch.uint8) + scale = torch.randint(122, 126, (E, rows, cols // 32), dtype=torch.uint8) + return packed.to(dev), scale.to(dev) + + w13, w13_s = rand_mxfp4(2 * N, K) + w2, w2_s = rand_mxfp4(K, N) + w13_ref = torch.stack([_dequant_ref(w13[e], w13_s[e]) for e in range(E)]) + w2_ref = torch.stack([_dequant_ref(w2[e], w2_s[e]) for e in range(E)]) + b13 = torch.randn(E, 2 * N, device=dev) if with_bias else None + b2 = torch.randn(E, K, device=dev) if with_bias else None + + x = torch.randn(M, K, device=dev, dtype=dtype) + topk_weights, topk_ids = torch.randn(M, E, device=dev).softmax(-1).topk(top_k) + topk_ids = topk_ids.to(torch.int32) + + ref = torch.zeros(M, K, device=dev) + for m in range(M): + for j in range(top_k): + e = topk_ids[m, j] + h = x[m].float() @ w13_ref[e].t() + if with_bias: + h = h + b13[e] + if activation == MoEActivation.SWIGLUOAI: + a = _swigluoai_interleaved(h) + else: + a = F.silu(h[:N]) * h[N:] + y = a.to(dtype).float() @ w2_ref[e].t() + if with_bias: + y = y + b2[e] + ref[m] += topk_weights[m, j] * y + + deinterleave = activation == MoEActivation.SWIGLUOAI + b_q13, b_s13 = repack_experts_rdna3(w13, w13_s, deinterleave=deinterleave) + b_q2, b_s2 = repack_experts_rdna3(w2, w2_s, deinterleave=False) + if deinterleave: + b13 = torch.cat([b13[:, ::2], b13[:, 1::2]], dim=1) + quant_config = mxfp4_w4a16_moe_quant_config( + w1_scale=b_s13, + w2_scale=b_s2, + w1_bias=None if b13 is None else b13.to(dtype), + w2_bias=None if b2 is None else b2.to(dtype), + ) + moe_config = make_dummy_moe_config( + num_experts=E, + experts_per_token=top_k, + hidden_dim=K, + intermediate_size=N, + in_dtype=dtype, + activation=activation, + ) + experts = RDNA3Mxfp4Experts(moe_config, quant_config) + ws13, ws2, out_shape = experts.workspace_shapes( + M, N, K, top_k, E, E, None, activation + ) + out = torch.empty(out_shape, device=dev, dtype=dtype) + experts.apply( + out, + x, + b_q13, + b_q2, + topk_weights, + topk_ids, + activation, + E, + None, + None, + None, + torch.empty(ws13, device=dev, dtype=dtype), + torch.empty(ws2, device=dev, dtype=dtype), + None, + False, + ) + + rel = (out.float() - ref).norm() / ref.norm() + assert rel < 2e-2, f"relative L2 error {rel:.4f}" diff --git a/tests/kernels/quantization/test_mxfp4_rdna3_dispatch.py b/tests/kernels/quantization/test_mxfp4_rdna3_dispatch.py index 2326f0ee80e9..65ce73ad6fed 100644 --- a/tests/kernels/quantization/test_mxfp4_rdna3_dispatch.py +++ b/tests/kernels/quantization/test_mxfp4_rdna3_dispatch.py @@ -7,13 +7,36 @@ scalar tile (1/4). The ``moe_mxfp4_gemm_rdna3`` kernel only instantiates tiles {1,2,4,8,16}, so the selector must never return anything else. Pure-Python (no GPU), so this runs on any CI host. + +The oracle tests at the bottom run on real gfx1100 hardware: which backend a +checkpoint that asks for quantized activations lands on. """ +from unittest.mock import patch + import pytest +import torch +from vllm.config.kernel import MoEBackend +from vllm.model_executor.layers.fused_moe.activation import MoEActivation +from vllm.model_executor.layers.fused_moe.config import ( + FusedMoEConfig, + FusedMoEParallelConfig, + RoutingMethodType, +) from vllm.model_executor.layers.fused_moe.experts.rdna3_mxfp4_moe import ( _select_block_size_m, ) +from vllm.model_executor.layers.fused_moe.oracle.mxfp4 import ( + Mxfp4MoeBackend, + select_mxfp4_moe_backend, +) +from vllm.model_executor.layers.quantization.utils.quant_utils import ( + QuantKey, + kFp8StaticTensorSym, + kMxfp4Dynamic, +) +from vllm.platforms import current_platform SUPPORTED_TILES = {1, 2, 4, 8, 16} # moe_mxfp4_gemm_rdna3 instantiations BATCH_SIZES = [1, 2, 3, 4, 8, 15, 16, 17, 32, 64, 256, 512, 4096] @@ -42,3 +65,88 @@ def test_degenerate_inputs_return_safe_tile(): # top_k/num_experts <= 0 must not divide-by-zero; fall back to the 1-tile. assert _select_block_size_m(32, 0, 256) == 1 assert _select_block_size_m(32, 8, 0) == 1 + + +ON_GFX1100 = False +if current_platform.is_rocm(): + from vllm.platforms.rocm import on_gfx1100 + + ON_GFX1100 = on_gfx1100() + +# W4A16 checkpoints, Quark W4A4 (e.g. amd/Qwen3-Coder-Next-MXFP4) and Quark W4A8 +# (e.g. amd/gpt-oss-120b-w-mxfp4-a-fp8). gfx1100 can run none of the quantized +# activations natively. +ACTIVATION_KEYS = [None, kMxfp4Dynamic, kFp8StaticTensorSym] + + +def _make_moe_config(moe_backend: MoEBackend) -> FusedMoEConfig: + return FusedMoEConfig( + num_experts=8, + experts_per_token=2, + hidden_dim=256, + intermediate_size=256, + num_local_experts=8, + num_logical_experts=8, + moe_parallel_config=FusedMoEParallelConfig.make_no_parallel(), + activation=MoEActivation.SILU, + in_dtype=torch.bfloat16, + device="cuda", + routing_method=RoutingMethodType.Renormalize, + moe_backend=moe_backend, + ) + + +@pytest.fixture +def no_user_override(): + with patch( + "vllm.model_executor.layers.fused_moe.oracle.mxfp4.get_current_vllm_config" + ) as mock_get_config: + mock_get_config.return_value.model_config.quantization_config = None + yield + + +def _with_quark(installed: bool): + return patch( + "vllm.model_executor.layers.fused_moe.experts.ocp_mx_emulation_moe.has_quark", + return_value=installed, + ) + + +@pytest.mark.skipif(not ON_GFX1100, reason="Requires gfx1100") +@pytest.mark.parametrize("activation_key", ACTIVATION_KEYS) +def test_explicit_rdna3_backend(no_user_override, activation_key: QuantKey | None): + """--moe-backend rdna3 selects the native kernel whatever the checkpoint's + activation scheme, running activations in BF16.""" + with _with_quark(True): + backend, _ = select_mxfp4_moe_backend( + _make_moe_config("rdna3"), activation_key=activation_key + ) + assert backend == Mxfp4MoeBackend.RDNA3_MXFP4 + + +@pytest.mark.skipif(not ON_GFX1100, reason="Requires gfx1100") +def test_auto_w4a16_selects_rdna3(no_user_override): + with _with_quark(True): + backend, _ = select_mxfp4_moe_backend(_make_moe_config("auto")) + assert backend == Mxfp4MoeBackend.RDNA3_MXFP4 + + +@pytest.mark.skipif(not ON_GFX1100, reason="Requires gfx1100") +@pytest.mark.parametrize("activation_key", [kMxfp4Dynamic, kFp8StaticTensorSym]) +@pytest.mark.parametrize( + "quark_installed, expected", + [ + # Emulation honours the checkpoint's activation quantization, so it + # wins under auto; the native kernel is the BF16 fallback. + (True, Mxfp4MoeBackend.EMULATION), + (False, Mxfp4MoeBackend.RDNA3_MXFP4), + ], +) +def test_auto_quantized_activations( + no_user_override, activation_key, quark_installed, expected +): + with _with_quark(quark_installed): + backend, _ = select_mxfp4_moe_backend( + _make_moe_config("auto"), activation_key=activation_key + ) + assert backend == expected diff --git a/vllm/_custom_ops.py b/vllm/_custom_ops.py index 77ed7560dd59..d26f39da5279 100644 --- a/vllm/_custom_ops.py +++ b/vllm/_custom_ops.py @@ -713,6 +713,27 @@ def _gptq_gemm_rdna3_wmma_fake( ) +def mxfp4_gemm_rdna3( + a: torch.Tensor, + b_q_weight: torch.Tensor, + b_scales_e8m0: torch.Tensor, +) -> torch.Tensor: + return torch.ops._rocm_C.mxfp4_gemm_rdna3(a, b_q_weight, b_scales_e8m0) + + +if hasattr(torch.ops, "_rocm_C") and hasattr(torch.ops._rocm_C, "mxfp4_gemm_rdna3"): + + @register_fake("_rocm_C::mxfp4_gemm_rdna3") + def _mxfp4_gemm_rdna3_fake( + a: torch.Tensor, + b_q_weight: torch.Tensor, + b_scales_e8m0: torch.Tensor, + ) -> torch.Tensor: + return torch.empty( + (a.size(0), b_q_weight.size(1)), dtype=a.dtype, device=a.device + ) + + def moe_gptq_gemm_rdna3( a: torch.Tensor, c: torch.Tensor, diff --git a/vllm/model_executor/kernels/linear/__init__.py b/vllm/model_executor/kernels/linear/__init__.py index c745447fae90..1cfcce25a041 100644 --- a/vllm/model_executor/kernels/linear/__init__.py +++ b/vllm/model_executor/kernels/linear/__init__.py @@ -91,7 +91,7 @@ from vllm.model_executor.kernels.linear.mxfp4.pytorch import ( TorchMxFp4LinearKernel, ) -from vllm.model_executor.kernels.linear.mxfp4.rocm import ( +from vllm.model_executor.kernels.linear.mxfp4.rdna3 import ( Rdna3MxFp4LinearKernel, ) from vllm.model_executor.kernels.linear.mxfp4.xpu import ( diff --git a/vllm/model_executor/kernels/linear/mxfp4/rocm.py b/vllm/model_executor/kernels/linear/mxfp4/rdna3.py similarity index 84% rename from vllm/model_executor/kernels/linear/mxfp4/rocm.py rename to vllm/model_executor/kernels/linear/mxfp4/rdna3.py index f9cc954cb787..7a0f5895183c 100644 --- a/vllm/model_executor/kernels/linear/mxfp4/rocm.py +++ b/vllm/model_executor/kernels/linear/mxfp4/rdna3.py @@ -4,6 +4,7 @@ import torch from torch.nn.parameter import Parameter +from vllm import _custom_ops as ops from vllm.logger import init_logger from vllm.model_executor.layers.quantization.utils.quant_utils import kMxfp4Dynamic from vllm.platforms import current_platform @@ -14,19 +15,9 @@ _MXFP4_GROUP_SIZE = 32 -# Fake/meta impl so the op is traceable under torch.compile + CUDA graphs. -if hasattr(torch.ops, "_rocm_C") and hasattr(torch.ops._rocm_C, "mxfp4_gemm_rdna3"): - try: - - @torch.library.register_fake("_rocm_C::mxfp4_gemm_rdna3") - def _mxfp4_gemm_rdna3_fake(a, b_q_weight, b_scales_e8m0): - return a.new_empty((a.shape[0], b_q_weight.shape[1])) - except RuntimeError: - pass # already registered - class Rdna3MxFp4LinearKernel(MxFp4LinearKernel): - """Weight-only MXFP4 GEMM for RDNA3 (gfx11xx) via the native HIP kernel + """Weight-only MXFP4 GEMM for RDNA3 (gfx1100) via the native HIP kernel ``_rocm_C.mxfp4_gemm_rdna3`` (scalar GEMV for decode, WMMA for prefill). RDNA3 has no native FP4 matmul, so E2M1 weights are dequantized on the fly into the f16/bf16 pipeline (W4A16). @@ -88,7 +79,7 @@ def apply_weights( if not x_2d.is_contiguous(): x_2d = x_2d.contiguous() - out = torch.ops._rocm_C.mxfp4_gemm_rdna3(x_2d, layer.weight, layer.weight_scale) + out = ops.mxfp4_gemm_rdna3(x_2d, layer.weight, layer.weight_scale) if bias is not None: out = out + bias return out.view(out_shape) diff --git a/vllm/model_executor/layers/fused_moe/experts/rdna3_mxfp4_moe.py b/vllm/model_executor/layers/fused_moe/experts/rdna3_mxfp4_moe.py index a6a7d7028699..2065f3dbb8ab 100644 --- a/vllm/model_executor/layers/fused_moe/experts/rdna3_mxfp4_moe.py +++ b/vllm/model_executor/layers/fused_moe/experts/rdna3_mxfp4_moe.py @@ -13,15 +13,20 @@ Registered as ``Mxfp4MoeBackend.RDNA3_MXFP4`` in the oracle and gated on gfx1100. """ +import dataclasses + import torch import vllm.model_executor.layers.fused_moe.modular_kernel as mk from vllm import _custom_ops as ops from vllm.model_executor.layers.fused_moe.activation import ( MoEActivation, + apply_moe_activation, ) from vllm.model_executor.layers.fused_moe.config import ( + FusedMoEConfig, FusedMoEParallelConfig, + FusedMoEQuantConfig, ) from vllm.model_executor.layers.fused_moe.moe_align_block_size import ( moe_align_block_size, @@ -75,9 +80,34 @@ def repack_experts_rdna3( return b_q, b_scale +def _first_set(*values: float | None) -> float: + return next(v for v in values if v is not None) + + class RDNA3Mxfp4Experts(mk.FusedMoEExpertsModular): """MXFP4 MoE experts on the native RDNA3 HIP kernel (gfx1100 only).""" + def __init__( + self, + moe_config: FusedMoEConfig, + quant_config: FusedMoEQuantConfig, + max_num_tokens: int | None = None, + num_dispatchers: int | None = None, + ): + super().__init__(moe_config, quant_config, max_num_tokens, num_dispatchers) + # The repack de-interleaves GPT-OSS's gate/up rows, so SWIGLUOAI runs as + # the packed SWIGLUOAI_UNINTERLEAVE. Not every quant path forwards the + # SwiGLU constants (gpt_oss.py does not), so fall back to the values + # swigluoai_and_mul hardcodes. + self.swigluoai_config = dataclasses.replace( + self.activation_config, + clamp_limit=_first_set( + quant_config.gemm1_clamp_limit, moe_config.swiglu_limit, 7.0 + ), + alpha=_first_set(quant_config.gemm1_alpha, moe_config.swiglu_alpha, 1.702), + beta=_first_set(quant_config.gemm1_beta, moe_config.swiglu_beta, 1.0), + ) + @staticmethod def activation_format() -> mk.FusedMoEActivationFormat: return mk.FusedMoEActivationFormat.Standard @@ -232,12 +262,15 @@ def apply( if w1_bias is not None: w1_out.add_(self._expert_bias_rows(w1_bias, topk_ids, expert_map, dtype)) - # The repack de-interleaves GPT-OSS's gate/up rows, so the contiguous - # SwiGLU-OAI variant is the matching primitive; both read alpha/beta/ - # clamp from the activation config. if activation == MoEActivation.SWIGLUOAI: - activation = MoEActivation.SWIGLUOAI_UNINTERLEAVE - self.activation(activation, act_out, w1_out) + apply_moe_activation( + MoEActivation.SWIGLUOAI_UNINTERLEAVE, + act_out, + w1_out, + activation_config=self.swigluoai_config, + ) + else: + self.activation(activation, act_out, w1_out) # down GEMM. The kernel's fused output_topk reduce is wrong under TP # (each rank's down-proj is a partial the layer all-reduces afterwards); @@ -279,7 +312,7 @@ def apply( sti, eid, ntp, - top_k, + 1, block_size_m, not apply_router_weight_on_input, top_k, diff --git a/vllm/model_executor/layers/fused_moe/oracle/mxfp4.py b/vllm/model_executor/layers/fused_moe/oracle/mxfp4.py index a010a58558ad..83707bb31dee 100644 --- a/vllm/model_executor/layers/fused_moe/oracle/mxfp4.py +++ b/vllm/model_executor/layers/fused_moe/oracle/mxfp4.py @@ -398,6 +398,7 @@ def map_mxfp4_backend(runner_backend: MoEBackend) -> list[Mxfp4MoeBackend]: "xpu": [Mxfp4MoeBackend.XPU], "cpu": [Mxfp4MoeBackend.CPU], "emulation": [Mxfp4MoeBackend.EMULATION], + "rdna3": [Mxfp4MoeBackend.RDNA3_MXFP4], } if backends := mapping.get(runner_backend): return backends @@ -925,7 +926,6 @@ def convert_gpt_oss_weight_to_mxfp4_moe_kernel_format( w13_input_scale: torch.Tensor | None = None, w2_input_scale: torch.Tensor | None = None, _cache_permute_indices: dict[torch.Size, torch.Tensor] | None = None, - gate_up_interleaved: bool = True, ) -> tuple[ torch.Tensor, torch.Tensor, @@ -934,13 +934,7 @@ def convert_gpt_oss_weight_to_mxfp4_moe_kernel_format( torch.Tensor | None, torch.Tensor | None, ]: - """Convert loaded weights into backend-specific kernel format. - - ``gate_up_interleaved`` describes the loaded ``w13`` layout: GPT-OSS fuses - gate/up row-interleaved ([g0, u0, g1, u1, ...]), while checkpoints with - separate gate_proj/up_proj tensors are already contiguous after ``_load_w13``. - Backends that only ever see GPT-OSS weights ignore it. - """ + """Convert loaded weights into backend-specific kernel format.""" if mxfp4_backend == Mxfp4MoeBackend.DEEPGEMM_MXFP4: w13_weight_scale, w2_weight_scale = _pack_deepgemm_mxfp4_scales( w13_weight, @@ -1003,13 +997,14 @@ def convert_gpt_oss_weight_to_mxfp4_moe_kernel_format( repack_experts_rdna3, ) + deinterleave = layer.activation == MoEActivation.SWIGLUOAI b_q13, b_s13 = repack_experts_rdna3( - w13_weight.data, w13_weight_scale.data, deinterleave=gate_up_interleaved + w13_weight.data, w13_weight_scale.data, deinterleave=deinterleave ) b_q2, b_s2 = repack_experts_rdna3( w2_weight.data, w2_weight_scale.data, deinterleave=False ) - if gate_up_interleaved and w13_bias is not None: + if deinterleave and w13_bias is not None: w13_bias = torch.cat( [w13_bias.data[:, ::2], w13_bias.data[:, 1::2]], dim=1 ).contiguous() diff --git a/vllm/model_executor/layers/quantization/compressed_tensors/compressed_tensors_moe/compressed_tensors_moe_w4a4_mxfp4.py b/vllm/model_executor/layers/quantization/compressed_tensors/compressed_tensors_moe/compressed_tensors_moe_w4a4_mxfp4.py index b7dbdf295387..a9cf3c8e1ca7 100644 --- a/vllm/model_executor/layers/quantization/compressed_tensors/compressed_tensors_moe/compressed_tensors_moe_w4a4_mxfp4.py +++ b/vllm/model_executor/layers/quantization/compressed_tensors/compressed_tensors_moe/compressed_tensors_moe_w4a4_mxfp4.py @@ -77,7 +77,9 @@ def __init__(self, moe): self.mxfp4_backend = Mxfp4MoeBackend.XPU self.experts_cls = XPUExpertsMxFp4 logger.info_once("Using XPUExpertsMxFp4 for MXFP4 MoE on XPU platform") - elif RDNA3Mxfp4Experts._supports_current_device(): + elif RDNA3Mxfp4Experts._supports_current_device() and ( + RDNA3Mxfp4Experts._supports_activation(moe.activation) + ): self.mxfp4_backend = Mxfp4MoeBackend.RDNA3_MXFP4 self.experts_cls = RDNA3Mxfp4Experts logger.info_once("Using RDNA3Mxfp4Experts (native gfx1100 HIP kernel)") @@ -231,8 +233,6 @@ def process_weights_after_loading(self, layer: RoutedExperts) -> None: elif self.mxfp4_backend in B12X_BACKENDS or current_platform.is_xpu(): pass elif self.mxfp4_backend == Mxfp4MoeBackend.RDNA3_MXFP4: - # compressed-tensors keeps gate_proj/up_proj separate, so w13 is - # already contiguous after _load_w13. w13, w2, w13_scale, w2_scale, _, _ = ( convert_gpt_oss_weight_to_mxfp4_moe_kernel_format( mxfp4_backend=self.mxfp4_backend, @@ -241,7 +241,6 @@ def process_weights_after_loading(self, layer: RoutedExperts) -> None: w2_weight=layer.w2_weight, w13_weight_scale=layer.w13_weight_scale, w2_weight_scale=layer.w2_weight_scale, - gate_up_interleaved=False, ) ) layer.w13_weight = torch.nn.Parameter(w13, requires_grad=False) diff --git a/vllm/model_executor/layers/quantization/quark/quark_moe.py b/vllm/model_executor/layers/quantization/quark/quark_moe.py index 49fbebdc9cd2..a17f7b28da87 100644 --- a/vllm/model_executor/layers/quantization/quark/quark_moe.py +++ b/vllm/model_executor/layers/quantization/quark/quark_moe.py @@ -1635,10 +1635,9 @@ def __init__( moe, activation_key=kFp8StaticTensorSym ) elif self.ocp_mx_scheme == "w_mxfp4_a_mxfp4": - # W4A4, or weight-only when the platform has no native FP4 (gfx1100). - act_key = kMxfp4Dynamic if current_platform.supports_mx() else None + # W4A4: MXFP4 weights + MXFP4 activations self.mxfp4_backend, self.experts_cls = select_mxfp4_moe_backend( - moe, activation_key=act_key + moe, activation_key=kMxfp4Dynamic ) # Validation for unsupported schemes @@ -1841,7 +1840,6 @@ def _setup_kernel(self, layer: RoutedExperts): w2_bias=w2_bias, w13_input_scale=layer.w13_input_scale, w2_input_scale=layer.w2_input_scale, - gate_up_interleaved=self.model_type == "gpt_oss", ) ) From 69ceda1cef8ea5bca4d9b9320bf7612e98f456da Mon Sep 17 00:00:00 2001 From: JartX Date: Fri, 23 Oct 2026 00:46:22 +0200 Subject: [PATCH 35/36] [ROCm][MXFP4] Make the experts test biases non-optional (mypy) Signed-off-by: JartX --- tests/kernels/quantization/test_mxfp4_rdna3.py | 17 +++++++---------- 1 file changed, 7 insertions(+), 10 deletions(-) diff --git a/tests/kernels/quantization/test_mxfp4_rdna3.py b/tests/kernels/quantization/test_mxfp4_rdna3.py index dda0247c782c..bb03c6d95c2e 100644 --- a/tests/kernels/quantization/test_mxfp4_rdna3.py +++ b/tests/kernels/quantization/test_mxfp4_rdna3.py @@ -125,8 +125,9 @@ def rand_mxfp4(rows, cols): w2, w2_s = rand_mxfp4(K, N) w13_ref = torch.stack([_dequant_ref(w13[e], w13_s[e]) for e in range(E)]) w2_ref = torch.stack([_dequant_ref(w2[e], w2_s[e]) for e in range(E)]) - b13 = torch.randn(E, 2 * N, device=dev) if with_bias else None - b2 = torch.randn(E, K, device=dev) if with_bias else None + bias_init = torch.randn if with_bias else torch.zeros + b13 = bias_init(E, 2 * N, device=dev) + b2 = bias_init(E, K, device=dev) x = torch.randn(M, K, device=dev, dtype=dtype) topk_weights, topk_ids = torch.randn(M, E, device=dev).softmax(-1).topk(top_k) @@ -136,16 +137,12 @@ def rand_mxfp4(rows, cols): for m in range(M): for j in range(top_k): e = topk_ids[m, j] - h = x[m].float() @ w13_ref[e].t() - if with_bias: - h = h + b13[e] + h = x[m].float() @ w13_ref[e].t() + b13[e] if activation == MoEActivation.SWIGLUOAI: a = _swigluoai_interleaved(h) else: a = F.silu(h[:N]) * h[N:] - y = a.to(dtype).float() @ w2_ref[e].t() - if with_bias: - y = y + b2[e] + y = a.to(dtype).float() @ w2_ref[e].t() + b2[e] ref[m] += topk_weights[m, j] * y deinterleave = activation == MoEActivation.SWIGLUOAI @@ -156,8 +153,8 @@ def rand_mxfp4(rows, cols): quant_config = mxfp4_w4a16_moe_quant_config( w1_scale=b_s13, w2_scale=b_s2, - w1_bias=None if b13 is None else b13.to(dtype), - w2_bias=None if b2 is None else b2.to(dtype), + w1_bias=b13.to(dtype) if with_bias else None, + w2_bias=b2.to(dtype) if with_bias else None, ) moe_config = make_dummy_moe_config( num_experts=E, From 91f53ed7d60ca0e48cc74e25ba21ca39f5459502 Mon Sep 17 00:00:00 2001 From: JartX Date: Fri, 23 Oct 2026 01:13:51 +0200 Subject: [PATCH 36/36] [ROCm][MXFP4] Correct the down-GEMM reduction comment, document the RDNA3 path The fused top-k reduce is not known to be wrong under TP: the earlier TP2 corruption was the top_k row-divisor bug fixed in e060ed0ab9. Keep TP>1 on the unfused path until it is validated, and say so. Document the gfx1100 native MXFP4 path and the --moe-backend rdna3 opt-in in the Quark guide. Signed-off-by: JartX --- docs/features/quantization/quark.md | 2 ++ .../layers/fused_moe/experts/rdna3_mxfp4_moe.py | 7 +++---- 2 files changed, 5 insertions(+), 4 deletions(-) diff --git a/docs/features/quantization/quark.md b/docs/features/quantization/quark.md index e5cfb8653bd6..fb6bbce32b36 100644 --- a/docs/features/quantization/quark.md +++ b/docs/features/quantization/quark.md @@ -270,6 +270,8 @@ vllm serve fxmarty/qwen1.5_moe_a2.7b_chat_w_fp4_a_fp6_e2m3 --tensor-parallel-siz A simulation of the matrix multiplication execution in MXFP4/MXFP6 can be run on devices that do not support OCP MX operations natively (e.g. AMD Instinct MI325, MI300 and MI250), dequantizing weights from FP4/FP6 to half precision on the fly, using a fused kernel. This is useful e.g. to evaluate FP4/FP6 models using vLLM, or alternatively to benefit from the ~2.5-4x memory savings (compared to float16 and bfloat16). +On Radeon RX 7900 series GPUs (gfx1100), MXFP4 weights also run on native HIP kernels with BF16/FP16 activations. Dense layers use them by default. For MoE layers of checkpoints that request activation quantization (e.g. `w_mxfp4_a_mxfp4`, `w_mxfp4_a_fp8`), pass `--moe-backend rdna3` to use them instead of the simulation; the activation quantization is then not applied. + To generate offline models quantized using MXFP4 data type, the easiest approach is to use AMD Quark's [quantization script](https://quark.docs.amd.com/latest/pytorch/example_quark_torch_llm_ptq.html), as an example: ```bash diff --git a/vllm/model_executor/layers/fused_moe/experts/rdna3_mxfp4_moe.py b/vllm/model_executor/layers/fused_moe/experts/rdna3_mxfp4_moe.py index 2065f3dbb8ab..60116e2c17b7 100644 --- a/vllm/model_executor/layers/fused_moe/experts/rdna3_mxfp4_moe.py +++ b/vllm/model_executor/layers/fused_moe/experts/rdna3_mxfp4_moe.py @@ -272,10 +272,9 @@ def apply( else: self.activation(activation, act_out, w1_out) - # down GEMM. The kernel's fused output_topk reduce is wrong under TP - # (each rank's down-proj is a partial the layer all-reduces afterwards); - # a per-expert bias likewise needs adding before the reduction. In both - # cases write unreduced rows and reduce in Python instead. + # down GEMM. A per-expert bias has to be added before the top-k + # reduction, so that case writes unreduced rows and reduces here. TP>1 + # takes the same path until the fused reduce is validated there. unfused = w2_bias is not None or self.moe_config.moe_parallel_config.tp_size > 1 if unfused: w2_out = gemm_out[: total * K].view(total, K)