diff --git a/csrc/libtorch_stable/fp32_router_gemm.cu b/csrc/libtorch_stable/fp32_router_gemm.cu index 64393fad61..3207082d57 100644 --- a/csrc/libtorch_stable/fp32_router_gemm.cu +++ b/csrc/libtorch_stable/fp32_router_gemm.cu @@ -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 #include +#include + // --------------------------------------------------------------------------- // 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 -__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 +__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(b_col + k_base, b_float); + float b_float[kEPB][VPT]; +#pragma unroll + for (int e = 0; e < kEPB; e++) { + load_weight(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(mat_a + m_idx * kHiddenDim + k_base, - a_float); + load_activation( + 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 -void invokeFp32RouterGemm(float* output, InputT const* mat_a, - float const* mat_b, cudaStream_t stream) { - constexpr int kBlockSize = 128; +template +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, - output, mat_a, mat_b); + cudaLaunchKernelEx( + &config, + fp32_router_gemm_kernel, + 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 +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 && kNumExperts == 256 && + kHiddenDim == 6144) { + if (!isBlackwellFamily()) { + launchFp32RouterGemm( + output, mat_a, mat_b, stream); + return; + } + if constexpr (kNumTokens <= 4) { + launchFp32RouterGemm( + output, mat_a, mat_b, stream); + } else if constexpr (kNumTokens >= 16 && kNumTokens % 2 == 0) { + launchFp32RouterGemm(output, mat_a, mat_b, stream); + } else { + launchFp32RouterGemm( + output, mat_a, mat_b, stream); + } + } else if constexpr (std::is_same_v && + 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( + output, mat_a, mat_b, stream); + return; + } + if constexpr (kNumTokens >= 12 && kNumTokens % 2 == 0) { + launchFp32RouterGemm(output, mat_a, mat_b, stream); + } else if constexpr (kNumTokens >= 6 && kNumTokens % 2 == 0) { + launchFp32RouterGemm(output, mat_a, mat_b, stream); + } else { + launchFp32RouterGemm( + output, mat_a, mat_b, stream); + } + } else if constexpr (std::is_same_v && + 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( + 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(output, mat_a, mat_b, stream); + } else if constexpr (kNumTokens >= 8 && kNumTokens <= 12 && + kNumTokens % 2 == 0) { + launchFp32RouterGemm(output, mat_a, mat_b, stream); + } else { + launchFp32RouterGemm( + output, mat_a, mat_b, stream); + } + } else { + launchFp32RouterGemm( + 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 diff --git a/csrc/libtorch_stable/fp32_router_gemm_entry.cu b/csrc/libtorch_stable/fp32_router_gemm_entry.cu index fc09193c86..8643e33cbe 100644 --- a/csrc/libtorch_stable/fp32_router_gemm_entry.cu +++ b/csrc/libtorch_stable/fp32_router_gemm_entry.cu @@ -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::unroll( num_tokens, output, mat_a, mat_b, stream); + } else if (num_experts == 256 && hidden_dim == 6144) { + Fp32LoopUnroller::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( diff --git a/tests/kernels/test_fp32_router_gemm.py b/tests/kernels/test_fp32_router_gemm.py index 0673a438c5..571dfa077d 100644 --- a/tests/kernels/test_fp32_router_gemm.py +++ b/tests/kernels/test_fp32_router_gemm.py @@ -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 + ) diff --git a/vllm/model_executor/layers/fused_moe/router/gate_linear.py b/vllm/model_executor/layers/fused_moe/router/gate_linear.py index b11464b19d..37e7018a86 100644 --- a/vllm/model_executor/layers/fused_moe/router/gate_linear.py +++ b/vllm/model_executor/layers/fused_moe/router/gate_linear.py @@ -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__(