FP32 router GEMV optimization (#48335)
Signed-off-by: peiyuanz <peiyuanz@inferact.ai> Signed-off-by: Jee Jee Li <jeejeelee@inferact.ai> Co-authored-by: peiyuanz <peiyuanz@inferact.ai> Co-authored-by: Claude <noreply@anthropic.com> Co-authored-by: zhouzhou <zhouzhou@zhouzhoudeMacBook-Pro.local>
This commit is contained in:
Jee Jee Li
co-authored by
peiyuanz
Claude
zhouzhou
peiyuanz
Claude
zhouzhou
parent
1bd8f80a64
commit
19069bcbd5
@@ -1,13 +1,18 @@
|
||||
// SPDX-License-Identifier: Apache-2.0
|
||||
// SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
||||
//
|
||||
// Router GEMM: activation(T) x weight(fp32) -> fp32, H=3072, E=256, M<=32.
|
||||
// Router GEMM: activation(T) x weight(fp32) -> fp32, M<=32, for the
|
||||
// supported (E, H) pairs listed at the bottom of this file.
|
||||
// Supports bf16 or fp32 activation; weight is always fp32.
|
||||
// Adapted from dsv3_router_gemm_float_out.cu.
|
||||
// (E=256, H=6144) bf16 uses a B300-tuned wide-block geometry; see
|
||||
// invokeFp32RouterGemm.
|
||||
|
||||
#include <cuda_bf16.h>
|
||||
#include <cuda_runtime.h>
|
||||
|
||||
#include <type_traits>
|
||||
|
||||
// ---------------------------------------------------------------------------
|
||||
// Load helpers
|
||||
// ---------------------------------------------------------------------------
|
||||
@@ -73,94 +78,113 @@ __device__ __forceinline__ void load_activation<__nv_bfloat16, 8>(
|
||||
// InputT : type of activation (float or __nv_bfloat16)
|
||||
// Weight is always fp32; output is always fp32.
|
||||
// VPT = 16 / sizeof(InputT): 4 for fp32, 8 for bf16
|
||||
template <typename InputT, int kBlockSize, int kNumTokens, int kNumExperts,
|
||||
int kHiddenDim>
|
||||
__global__ __launch_bounds__(128, 1) void fp32_router_gemm_kernel(
|
||||
float* out, InputT const* mat_a, float const* mat_b) {
|
||||
// Each block computes kEPB expert columns; wider blocks / kEPB > 1 are
|
||||
// selected per (shape, M) in invokeFp32RouterGemm (B300-tuned, see below).
|
||||
// kTGroups > 1 splits the tokens across groups of kBlockSize threads within
|
||||
// the block: all groups scan the same weight K-slices (group 0 misses to
|
||||
// DRAM, later groups hit L1) so weight traffic stays 1x, while per-thread
|
||||
// accumulator registers drop by kTGroups (at M=16 the 32 fp32 accumulators
|
||||
// push the kernel to 128 regs/thread and 1 block/SM).
|
||||
template <typename InputT, int kBlockSize, int kNumTokens, int kEPB,
|
||||
int kNumExperts, int kHiddenDim, int kTGroups = 1>
|
||||
__global__ __launch_bounds__(
|
||||
kBlockSize* kTGroups, 1) void fp32_router_gemm_kernel(float* out,
|
||||
InputT const* mat_a,
|
||||
float const* mat_b) {
|
||||
constexpr int VPT = 16 / sizeof(InputT);
|
||||
constexpr int k_elems_per_k_iteration = VPT * kBlockSize;
|
||||
constexpr int k_iterations = kHiddenDim / k_elems_per_k_iteration;
|
||||
static_assert(kHiddenDim % k_elems_per_k_iteration == 0);
|
||||
static_assert(kNumTokens % kTGroups == 0);
|
||||
constexpr int kWarpSize = 32;
|
||||
constexpr int kNumWarps = kBlockSize / kWarpSize;
|
||||
constexpr int kNumWarps = kBlockSize / kWarpSize; // per token group
|
||||
constexpr int kMG = kNumTokens / kTGroups; // tokens per group
|
||||
|
||||
int const n_idx = blockIdx.x;
|
||||
int const tid = threadIdx.x;
|
||||
int const e_base = blockIdx.x * kEPB;
|
||||
int const tid = threadIdx.x % kBlockSize;
|
||||
int const m0 = (threadIdx.x / kBlockSize) * kMG;
|
||||
int const warpId = tid / kWarpSize;
|
||||
int const laneId = tid % kWarpSize;
|
||||
|
||||
float acc[kNumTokens] = {};
|
||||
__shared__ float sm_reduction[kNumTokens][kNumWarps];
|
||||
|
||||
float const* b_col = mat_b + n_idx * kHiddenDim;
|
||||
|
||||
int k_bases[k_iterations];
|
||||
#pragma unroll
|
||||
for (int ki = 0; ki < k_iterations; ki++) {
|
||||
k_bases[ki] = ki * k_elems_per_k_iteration + tid * VPT;
|
||||
}
|
||||
float acc[kMG][kEPB] = {};
|
||||
__shared__ float sm_reduction[kNumTokens][kEPB][kNumWarps];
|
||||
|
||||
#if defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 900)
|
||||
cudaGridDependencySynchronize();
|
||||
// Fire the PDL trigger right after our own wait instead of at kernel end:
|
||||
// a gridsync-ing consumer is unaffected (its wait always targets full grid
|
||||
// completion), while a consumer that reads none of our outputs (e.g. the
|
||||
// NVFP4 activation quant, which reads the same hidden_states) can launch
|
||||
// now and fully overlap this kernel's body.
|
||||
cudaTriggerProgrammaticLaunchCompletion();
|
||||
#endif
|
||||
|
||||
#pragma unroll
|
||||
for (int ki = 0; ki < k_iterations; ki++) {
|
||||
int const k_base = k_bases[ki];
|
||||
int const k_base = ki * k_elems_per_k_iteration + tid * VPT;
|
||||
|
||||
float b_float[VPT];
|
||||
load_weight<VPT>(b_col + k_base, b_float);
|
||||
float b_float[kEPB][VPT];
|
||||
#pragma unroll
|
||||
for (int e = 0; e < kEPB; e++) {
|
||||
load_weight<VPT>(mat_b + (e_base + e) * kHiddenDim + k_base, b_float[e]);
|
||||
}
|
||||
|
||||
#pragma unroll
|
||||
for (int m_idx = 0; m_idx < kNumTokens; m_idx++) {
|
||||
for (int m_idx = 0; m_idx < kMG; m_idx++) {
|
||||
float a_float[VPT];
|
||||
load_activation<InputT, VPT>(mat_a + m_idx * kHiddenDim + k_base,
|
||||
a_float);
|
||||
load_activation<InputT, VPT>(
|
||||
mat_a + (size_t)(m0 + m_idx) * kHiddenDim + k_base, a_float);
|
||||
#pragma unroll
|
||||
for (int k = 0; k < VPT; k++) {
|
||||
acc[m_idx] += a_float[k] * b_float[k];
|
||||
for (int e = 0; e < kEPB; e++) {
|
||||
#pragma unroll
|
||||
for (int k = 0; k < VPT; k++) {
|
||||
acc[m_idx][e] += a_float[k] * b_float[e][k];
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// Warp-level butterfly reduction
|
||||
#pragma unroll
|
||||
for (int m = 0; m < kNumTokens; m++) {
|
||||
float sum = acc[m];
|
||||
sum += __shfl_xor_sync(0xffffffff, sum, 16);
|
||||
sum += __shfl_xor_sync(0xffffffff, sum, 8);
|
||||
sum += __shfl_xor_sync(0xffffffff, sum, 4);
|
||||
sum += __shfl_xor_sync(0xffffffff, sum, 2);
|
||||
sum += __shfl_xor_sync(0xffffffff, sum, 1);
|
||||
if (laneId == 0) sm_reduction[m][warpId] = sum;
|
||||
for (int m = 0; m < kMG; m++) {
|
||||
#pragma unroll
|
||||
for (int e = 0; e < kEPB; e++) {
|
||||
float sum = acc[m][e];
|
||||
sum += __shfl_xor_sync(0xffffffff, sum, 16);
|
||||
sum += __shfl_xor_sync(0xffffffff, sum, 8);
|
||||
sum += __shfl_xor_sync(0xffffffff, sum, 4);
|
||||
sum += __shfl_xor_sync(0xffffffff, sum, 2);
|
||||
sum += __shfl_xor_sync(0xffffffff, sum, 1);
|
||||
if (laneId == 0) sm_reduction[m0 + m][e][warpId] = sum;
|
||||
}
|
||||
}
|
||||
|
||||
__syncthreads();
|
||||
|
||||
if (tid == 0) {
|
||||
// Parallel finalize: one thread per (m, e) output.
|
||||
for (int idx = threadIdx.x; idx < kNumTokens * kEPB;
|
||||
idx += kBlockSize * kTGroups) {
|
||||
int const m = idx / kEPB;
|
||||
int const e = idx % kEPB;
|
||||
float final_sum = 0.0f;
|
||||
#pragma unroll
|
||||
for (int m = 0; m < kNumTokens; m++) {
|
||||
float final_sum = 0.0f;
|
||||
#pragma unroll
|
||||
for (int w = 0; w < kNumWarps; w++) final_sum += sm_reduction[m][w];
|
||||
out[m * kNumExperts + n_idx] = final_sum;
|
||||
}
|
||||
for (int w = 0; w < kNumWarps; w++) final_sum += sm_reduction[m][e][w];
|
||||
out[m * kNumExperts + e_base + e] = final_sum;
|
||||
}
|
||||
|
||||
#if defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 900)
|
||||
cudaTriggerProgrammaticLaunchCompletion();
|
||||
#endif
|
||||
}
|
||||
|
||||
// ---------------------------------------------------------------------------
|
||||
// Launcher
|
||||
// ---------------------------------------------------------------------------
|
||||
|
||||
template <typename InputT, int kNumTokens, int kNumExperts, int kHiddenDim>
|
||||
void invokeFp32RouterGemm(float* output, InputT const* mat_a,
|
||||
float const* mat_b, cudaStream_t stream) {
|
||||
constexpr int kBlockSize = 128;
|
||||
template <typename InputT, int kBlockSize, int kEPB, int kNumTokens,
|
||||
int kNumExperts, int kHiddenDim, int kTGroups = 1>
|
||||
static void launchFp32RouterGemm(float* output, InputT const* mat_a,
|
||||
float const* mat_b, cudaStream_t stream) {
|
||||
static_assert(kNumExperts % kEPB == 0);
|
||||
cudaLaunchConfig_t config;
|
||||
config.gridDim = kNumExperts;
|
||||
config.blockDim = kBlockSize;
|
||||
config.gridDim = kNumExperts / kEPB;
|
||||
config.blockDim = kBlockSize * kTGroups;
|
||||
config.dynamicSmemBytes = 0;
|
||||
config.stream = stream;
|
||||
cudaLaunchAttribute attrs[1];
|
||||
@@ -168,15 +192,112 @@ void invokeFp32RouterGemm(float* output, InputT const* mat_a,
|
||||
attrs[0].val.programmaticStreamSerializationAllowed = 1;
|
||||
config.numAttrs = 1;
|
||||
config.attrs = attrs;
|
||||
cudaLaunchKernelEx(&config,
|
||||
fp32_router_gemm_kernel<InputT, kBlockSize, kNumTokens,
|
||||
kNumExperts, kHiddenDim>,
|
||||
output, mat_a, mat_b);
|
||||
cudaLaunchKernelEx(
|
||||
&config,
|
||||
fp32_router_gemm_kernel<InputT, kBlockSize, kNumTokens, kEPB, kNumExperts,
|
||||
kHiddenDim, kTGroups>,
|
||||
output, mat_a, mat_b);
|
||||
}
|
||||
|
||||
static bool isBlackwellFamily() {
|
||||
static int sm = []() {
|
||||
int dev = 0, major = 0, minor = 0;
|
||||
cudaGetDevice(&dev);
|
||||
cudaDeviceGetAttribute(&major, cudaDevAttrComputeCapabilityMajor, dev);
|
||||
cudaDeviceGetAttribute(&minor, cudaDevAttrComputeCapabilityMinor, dev);
|
||||
return major * 10 + minor;
|
||||
}();
|
||||
return sm >= 100;
|
||||
}
|
||||
|
||||
template <typename InputT, int kNumTokens, int kNumExperts, int kHiddenDim>
|
||||
void invokeFp32RouterGemm(float* output, InputT const* mat_a,
|
||||
float const* mat_b, cudaStream_t stream) {
|
||||
// Geometry tuned on B300 per supported shape, bf16 activation, under a
|
||||
// production-fidelity harness (CUDA-graph replay, per-layer cold weights).
|
||||
// GLM-5.2 (E=256, H=6144):
|
||||
// M <= 4 : BS=768, EPB=1 (2.7us vs cast+cuBLAS 8.1us at M=1)
|
||||
// M in [5, 15]
|
||||
// or odd : BS=384, EPB=2 (crossover vs BS=768 measured in (4, 8))
|
||||
// M >= 16, even : BS=192, EPB=2, 2 token groups (M=16 4.79us vs 5.04,
|
||||
// M=24 5.71 vs 6.38, M=32 6.79 vs 7.72; M=12 loses at
|
||||
// 0.97x, so the boundary is 16).
|
||||
// Only enabled on the Blackwell family where it was validated; Hopper and
|
||||
// other shapes / fp32 activation keep the legacy geometry.
|
||||
if constexpr (std::is_same_v<InputT, __nv_bfloat16> && kNumExperts == 256 &&
|
||||
kHiddenDim == 6144) {
|
||||
if (!isBlackwellFamily()) {
|
||||
launchFp32RouterGemm<InputT, 128, 1, kNumTokens, kNumExperts, kHiddenDim>(
|
||||
output, mat_a, mat_b, stream);
|
||||
return;
|
||||
}
|
||||
if constexpr (kNumTokens <= 4) {
|
||||
launchFp32RouterGemm<InputT, 768, 1, kNumTokens, kNumExperts, kHiddenDim>(
|
||||
output, mat_a, mat_b, stream);
|
||||
} else if constexpr (kNumTokens >= 16 && kNumTokens % 2 == 0) {
|
||||
launchFp32RouterGemm<InputT, 192, 2, kNumTokens, kNumExperts, kHiddenDim,
|
||||
2>(output, mat_a, mat_b, stream);
|
||||
} else {
|
||||
launchFp32RouterGemm<InputT, 384, 2, kNumTokens, kNumExperts, kHiddenDim>(
|
||||
output, mat_a, mat_b, stream);
|
||||
}
|
||||
} else if constexpr (std::is_same_v<InputT, __nv_bfloat16> &&
|
||||
kNumExperts == 128 && kHiddenDim == 6144) {
|
||||
// MiniMax-M3. Legacy 128/1 only fills 128 blocks and pays the same
|
||||
// accumulator register cliffs; B300 sweep:
|
||||
// even M in [6, 10] : BS=384, EPB=1, 2 token groups (1.26-1.43x)
|
||||
// even M >= 12 : BS=192, EPB=1, 2 token groups (1.59-1.66x at
|
||||
// M >= 18; re-measured on B300+B200: 192 also wins
|
||||
// M=12/14 by 5-11%% on both, ties 384 at 16)
|
||||
// M <= 5 / odd : BS=384, EPB=1 (1.03-1.19x)
|
||||
if (!isBlackwellFamily()) {
|
||||
launchFp32RouterGemm<InputT, 128, 1, kNumTokens, kNumExperts, kHiddenDim>(
|
||||
output, mat_a, mat_b, stream);
|
||||
return;
|
||||
}
|
||||
if constexpr (kNumTokens >= 12 && kNumTokens % 2 == 0) {
|
||||
launchFp32RouterGemm<InputT, 192, 1, kNumTokens, kNumExperts, kHiddenDim,
|
||||
2>(output, mat_a, mat_b, stream);
|
||||
} else if constexpr (kNumTokens >= 6 && kNumTokens % 2 == 0) {
|
||||
launchFp32RouterGemm<InputT, 384, 1, kNumTokens, kNumExperts, kHiddenDim,
|
||||
2>(output, mat_a, mat_b, stream);
|
||||
} else {
|
||||
launchFp32RouterGemm<InputT, 384, 1, kNumTokens, kNumExperts, kHiddenDim>(
|
||||
output, mat_a, mat_b, stream);
|
||||
}
|
||||
} else if constexpr (std::is_same_v<InputT, __nv_bfloat16> &&
|
||||
kNumExperts == 256 && kHiddenDim == 3072) {
|
||||
// MiniMax-M2/M2.5. The 3.1MB weight is latency-floor bound at small M
|
||||
// (legacy already optimal); token groups win only at even M >= 8
|
||||
// (1.05-1.17x). EPB crossover measured between 12 and 16.
|
||||
if (!isBlackwellFamily()) {
|
||||
launchFp32RouterGemm<InputT, 128, 1, kNumTokens, kNumExperts, kHiddenDim>(
|
||||
output, mat_a, mat_b, stream);
|
||||
return;
|
||||
}
|
||||
if constexpr (kNumTokens >= 14 && kNumTokens % 2 == 0) {
|
||||
// M=14 originally measured 0.91x and stayed on legacy; two fresh
|
||||
// sweeps (B300 dev1 + B200) both put 192/2/tg2 ahead by 3.5-4%%.
|
||||
launchFp32RouterGemm<InputT, 192, 2, kNumTokens, kNumExperts, kHiddenDim,
|
||||
2>(output, mat_a, mat_b, stream);
|
||||
} else if constexpr (kNumTokens >= 8 && kNumTokens <= 12 &&
|
||||
kNumTokens % 2 == 0) {
|
||||
launchFp32RouterGemm<InputT, 192, 1, kNumTokens, kNumExperts, kHiddenDim,
|
||||
2>(output, mat_a, mat_b, stream);
|
||||
} else {
|
||||
launchFp32RouterGemm<InputT, 128, 1, kNumTokens, kNumExperts, kHiddenDim>(
|
||||
output, mat_a, mat_b, stream);
|
||||
}
|
||||
} else {
|
||||
launchFp32RouterGemm<InputT, 128, 1, kNumTokens, kNumExperts, kHiddenDim>(
|
||||
output, mat_a, mat_b, stream);
|
||||
}
|
||||
}
|
||||
|
||||
// ---------------------------------------------------------------------------
|
||||
// Explicit instantiations: M=1..32, for both input types, for the supported
|
||||
// (E, H) pairs: (256, 3072) [MiniMax-M2/M2.5] and (128, 6144) [MiniMax-M3].
|
||||
// (E, H) pairs: (256, 3072) [MiniMax-M2/M2.5], (128, 6144) [MiniMax-M3]
|
||||
// and (256, 6144) [GLM-5.2].
|
||||
// ---------------------------------------------------------------------------
|
||||
|
||||
#define INSTANTIATE(T, M, E, H) \
|
||||
@@ -221,6 +342,8 @@ INSTANTIATE_ALL(float, 256, 3072)
|
||||
INSTANTIATE_ALL(__nv_bfloat16, 256, 3072)
|
||||
INSTANTIATE_ALL(float, 128, 6144)
|
||||
INSTANTIATE_ALL(__nv_bfloat16, 128, 6144)
|
||||
INSTANTIATE_ALL(float, 256, 6144)
|
||||
INSTANTIATE_ALL(__nv_bfloat16, 256, 6144)
|
||||
|
||||
#undef INSTANTIATE_ALL
|
||||
#undef INSTANTIATE
|
||||
|
||||
@@ -25,10 +25,12 @@ inline int getSMVersion() {
|
||||
static constexpr int FP32_MAX_TOKENS = 32;
|
||||
|
||||
// Supported (hidden_dim, num_experts) pairs (must match the instantiations in
|
||||
// fp32_router_gemm.cu): (3072, 256) for MiniMax-M2/M2.5, (6144, 128) for M3.
|
||||
// fp32_router_gemm.cu): (3072, 256) for MiniMax-M2/M2.5, (6144, 128) for M3,
|
||||
// (6144, 256) for GLM-5.2.
|
||||
static inline bool fp32_router_gemm_supported(int hidden_dim, int num_experts) {
|
||||
return (hidden_dim == 3072 && num_experts == 256) ||
|
||||
(hidden_dim == 6144 && num_experts == 128);
|
||||
(hidden_dim == 6144 && num_experts == 128) ||
|
||||
(hidden_dim == 6144 && num_experts == 256);
|
||||
}
|
||||
|
||||
// Forward declarations — 4 template params must match fp32_router_gemm.cu
|
||||
@@ -77,6 +79,9 @@ void dispatchFp32RouterGemm(int num_experts, int hidden_dim, int num_tokens,
|
||||
} else if (num_experts == 128 && hidden_dim == 6144) {
|
||||
Fp32LoopUnroller<InputT, 128, 6144, 1, FP32_MAX_TOKENS>::unroll(
|
||||
num_tokens, output, mat_a, mat_b, stream);
|
||||
} else if (num_experts == 256 && hidden_dim == 6144) {
|
||||
Fp32LoopUnroller<InputT, 256, 6144, 1, FP32_MAX_TOKENS>::unroll(
|
||||
num_tokens, output, mat_a, mat_b, stream);
|
||||
} else {
|
||||
throw std::invalid_argument(
|
||||
"fp32_router_gemm: unsupported (hidden_dim, num_experts) pair");
|
||||
@@ -111,7 +116,7 @@ void fp32_router_gemm(
|
||||
STD_TORCH_CHECK(
|
||||
fp32_router_gemm_supported(hidden_dim, num_experts),
|
||||
"fp32_router_gemm: supported (hidden_dim, num_experts) pairs are "
|
||||
"(3072, 256) and (6144, 128)");
|
||||
"(3072, 256), (6144, 128) and (6144, 256)");
|
||||
STD_TORCH_CHECK(num_tokens <= FP32_MAX_TOKENS,
|
||||
"fp32_router_gemm: num_tokens must be in [0, 32]");
|
||||
STD_TORCH_CHECK(
|
||||
|
||||
@@ -3,9 +3,13 @@
|
||||
"""Tests for fp32_router_gemm kernel: activation×weight→fp32.
|
||||
|
||||
Supported (hidden_size, num_experts) pairs:
|
||||
(3072, 256) -> MiniMax-M2/M2.5, (6144, 128) -> MiniMax-M3
|
||||
(3072, 256) -> MiniMax-M2/M2.5, (6144, 128) -> MiniMax-M3,
|
||||
(6144, 256) -> GLM-5.2
|
||||
|
||||
Correctness baseline: torch.matmul in float64.
|
||||
Correctness baseline: F.linear in float32. Every M in [1, 32] is covered so
|
||||
all tuned geometries (wide-block, experts-per-block, token-group; boundaries
|
||||
at M=4/5, odd/even, M=15/16) are exercised on Blackwell, and the legacy
|
||||
128/1 geometry everywhere else.
|
||||
"""
|
||||
|
||||
import pytest
|
||||
@@ -14,7 +18,8 @@ import torch
|
||||
from vllm._custom_ops import fp32_router_gemm
|
||||
|
||||
# (hidden_size, num_experts)
|
||||
SHAPES = [(3072, 256), (6144, 128)]
|
||||
SHAPES = [(3072, 256), (6144, 128), (6144, 256)]
|
||||
ALL_M = list(range(1, 33))
|
||||
# Absolute tolerance for fp32 kernel vs float64 reference
|
||||
ATOL_FP32 = 2e-4
|
||||
ATOL_BF16 = 2e-2 # bf16 activation has lower precision
|
||||
@@ -34,7 +39,7 @@ def _ref(mat_a: torch.Tensor, mat_b: torch.Tensor) -> torch.Tensor:
|
||||
|
||||
|
||||
@pytest.mark.parametrize("hidden_dim,num_experts", SHAPES)
|
||||
@pytest.mark.parametrize("num_tokens", [1, 2, 4, 8, 16, 32])
|
||||
@pytest.mark.parametrize("num_tokens", ALL_M)
|
||||
def test_fp32_activation(num_tokens: int, hidden_dim: int, num_experts: int):
|
||||
"""fp32 activation → fp32 output should match reference closely."""
|
||||
_requires_sm90()
|
||||
@@ -52,7 +57,7 @@ def test_fp32_activation(num_tokens: int, hidden_dim: int, num_experts: int):
|
||||
|
||||
|
||||
@pytest.mark.parametrize("hidden_dim,num_experts", SHAPES)
|
||||
@pytest.mark.parametrize("num_tokens", [1, 2, 4, 8, 16, 32])
|
||||
@pytest.mark.parametrize("num_tokens", ALL_M)
|
||||
def test_bf16_activation(num_tokens: int, hidden_dim: int, num_experts: int):
|
||||
"""bf16 activation → fp32 output should match reference within bf16 error."""
|
||||
_requires_sm90()
|
||||
@@ -82,3 +87,87 @@ def test_output_shape_and_dtype(hidden_dim: int, num_experts: int):
|
||||
assert out.shape == (4, num_experts)
|
||||
assert out.dtype == torch.float32
|
||||
assert out.device.type == "cuda"
|
||||
|
||||
|
||||
@pytest.mark.parametrize("hidden_dim,num_experts", SHAPES)
|
||||
@pytest.mark.parametrize("num_tokens", [1, 2, 4, 8, 16, 24, 32])
|
||||
def test_topk_routing_consistency(num_tokens: int, hidden_dim: int, num_experts: int):
|
||||
"""The gate feeds top-k expert selection: the kernel's top-8 must match
|
||||
an fp64 reference's top-8 per token (ties tolerated). This is the
|
||||
business-level correctness of the router — numeric error only matters
|
||||
if it flips the argsort."""
|
||||
_requires_sm90()
|
||||
top_k = 8
|
||||
device = torch.device("cuda")
|
||||
for seed in range(5):
|
||||
torch.manual_seed(1000 + seed)
|
||||
mat_a = torch.randn(num_tokens, hidden_dim, dtype=torch.bfloat16, device=device)
|
||||
mat_b = torch.randn(num_experts, hidden_dim, dtype=torch.float32, device=device)
|
||||
out = fp32_router_gemm(mat_a, mat_b)
|
||||
ref = mat_a.double() @ mat_b.double().t()
|
||||
kernel_idx = out.topk(top_k, dim=-1).indices
|
||||
ref_vals, ref_idx = ref.topk(top_k, dim=-1)
|
||||
for t in range(num_tokens):
|
||||
got = set(kernel_idx[t].tolist())
|
||||
want = set(ref_idx[t].tolist())
|
||||
if got == want:
|
||||
continue
|
||||
# Tolerate genuine near-ties around the k-th value only.
|
||||
kth = ref_vals[t, -1].item()
|
||||
for e in got.symmetric_difference(want):
|
||||
gap = abs(ref[t, e].item() - kth)
|
||||
assert gap < 1e-3, (
|
||||
f"top-{top_k} mismatch beyond tie tolerance: token {t}, "
|
||||
f"expert {e}, gap {gap:.3e}"
|
||||
)
|
||||
|
||||
|
||||
def test_zero_tokens_returns_empty():
|
||||
"""M=0 is a graceful no-op returning an empty [0, E] tensor."""
|
||||
_requires_sm90()
|
||||
device = torch.device("cuda")
|
||||
hidden_dim, num_experts = SHAPES[0]
|
||||
mat_a = torch.empty(0, hidden_dim, dtype=torch.float32, device=device)
|
||||
mat_b = torch.randn(num_experts, hidden_dim, dtype=torch.float32, device=device)
|
||||
out = fp32_router_gemm(mat_a, mat_b)
|
||||
assert out.shape == (0, num_experts)
|
||||
assert out.dtype == torch.float32
|
||||
|
||||
|
||||
def test_rejects_invalid_inputs():
|
||||
"""The entry must fail loudly, never compute silently wrong results."""
|
||||
_requires_sm90()
|
||||
device = torch.device("cuda")
|
||||
hidden_dim, num_experts = SHAPES[0]
|
||||
mat_b = torch.randn(num_experts, hidden_dim, dtype=torch.float32, device=device)
|
||||
|
||||
# num_tokens > 32 (beyond the instantiated range)
|
||||
with pytest.raises(Exception, match="num_tokens"):
|
||||
fp32_router_gemm(
|
||||
torch.randn(33, hidden_dim, dtype=torch.float32, device=device), mat_b
|
||||
)
|
||||
|
||||
# unsupported (hidden_dim, num_experts) pair
|
||||
with pytest.raises(Exception, match="supported"):
|
||||
fp32_router_gemm(
|
||||
torch.randn(4, 1024, dtype=torch.float32, device=device),
|
||||
torch.randn(64, 1024, dtype=torch.float32, device=device),
|
||||
)
|
||||
|
||||
# non-contiguous activation (a column-slice view)
|
||||
wide = torch.randn(4, hidden_dim * 2, dtype=torch.float32, device=device)
|
||||
with pytest.raises(Exception, match="contiguous"):
|
||||
fp32_router_gemm(wide[:, :hidden_dim], mat_b)
|
||||
|
||||
# wrong weight dtype (bf16 weight is not a supported layout)
|
||||
with pytest.raises(Exception, match="float32"):
|
||||
fp32_router_gemm(
|
||||
torch.randn(4, hidden_dim, dtype=torch.float32, device=device),
|
||||
mat_b.to(torch.bfloat16),
|
||||
)
|
||||
|
||||
# fp16 activation (only fp32 / bf16 are accepted)
|
||||
with pytest.raises(Exception, match="float32 or bfloat16"):
|
||||
fp32_router_gemm(
|
||||
torch.randn(4, hidden_dim, dtype=torch.float16, device=device), mat_b
|
||||
)
|
||||
|
||||
@@ -15,8 +15,8 @@ class GateLinear(ReplicatedLinear):
|
||||
"""MoE gate linear layer with multi-tier GEMM dispatch:
|
||||
|
||||
1. DSV3 specialized kernel (SM90+, M<=16, H=7168 E=256/384, H=6144 E=256)
|
||||
2. fp32 specialized kernel (SM90+, bf16/fp32 in, fp32 out,
|
||||
M<=32, H=3072, E=256)
|
||||
2. fp32 specialized kernel (SM90+, bf16/fp32 in, fp32 out, M<=32,
|
||||
(H, E) in {(3072, 256), (6144, 128), (6144, 256)})
|
||||
3. cuBLAS bf16×bf16→fp32 (SM90+ + bf16 weight + fp32 out_dtype)
|
||||
4. F.linear via ReplicatedLinear (ultimate fallback)
|
||||
|
||||
@@ -36,7 +36,7 @@ class GateLinear(ReplicatedLinear):
|
||||
|
||||
# (hidden_size, num_experts) pairs with an instantiated fp32 kernel:
|
||||
# (3072, 256) -> MiniMax-M2/M2.5, (6144, 128) -> MiniMax-M3
|
||||
FP32_SUPPORTED_SHAPES = {(3072, 256), (6144, 128)}
|
||||
FP32_SUPPORTED_SHAPES = {(3072, 256), (6144, 128), (6144, 256)}
|
||||
FP32_MAX_TOKENS = 32
|
||||
|
||||
def __init__(
|
||||
|
||||
Reference in New Issue
Block a user