diff --git a/CMakeLists.txt b/CMakeLists.txt index bb0f51b43e..ed51927362 100644 --- a/CMakeLists.txt +++ b/CMakeLists.txt @@ -1417,25 +1417,6 @@ if(VLLM_GPU_LANG STREQUAL "CUDA") message(STATUS "Not building Marlin MOE kernels as no compatible archs found" " in CUDA target architectures") endif() - - # DeepSeek V3 router GEMM kernel requires SM90+ and CUDA >= 12.0. - # (fp32_router_gemm has been migrated to _C_stable_libtorch above.) - cuda_archs_sm90plus(SM90PLUS_ROUTER_GEMM_ARCHS "${CUDA_ARCHS}") - if(${CMAKE_CUDA_COMPILER_VERSION} VERSION_GREATER_EQUAL 12.0 AND SM90PLUS_ROUTER_GEMM_ARCHS) - set(DSV3_ROUTER_GEMM_SRC - "csrc/libtorch_stable/moe/dsv3_router_gemm_entry.cu" - "csrc/libtorch_stable/moe/dsv3_router_gemm_float_out.cu" - "csrc/libtorch_stable/moe/dsv3_router_gemm_bf16_out.cu") - set_gencode_flags_for_srcs( - SRCS "${DSV3_ROUTER_GEMM_SRC}" - CUDA_ARCHS "${SM90PLUS_ROUTER_GEMM_ARCHS}") - list(APPEND VLLM_MOE_EXT_SRC "${DSV3_ROUTER_GEMM_SRC}") - - message(STATUS "Building DSV3 router GEMM kernels for archs: ${SM90PLUS_ROUTER_GEMM_ARCHS}") - else() - message(STATUS "Not building DSV3 router GEMM kernels as no compatible archs found" - " (requires SM90+ and CUDA >= 12.0)") - endif() endif() message(STATUS "Enabling MoE C_stable extension.") diff --git a/benchmarks/kernels/benchmark_router_gemm.py b/benchmarks/kernels/benchmark_router_gemm.py index ba46a7fd78..f4550da9df 100644 --- a/benchmarks/kernels/benchmark_router_gemm.py +++ b/benchmarks/kernels/benchmark_router_gemm.py @@ -10,10 +10,6 @@ from vllm.transformers_utils.config import get_config from vllm.triton_utils import triton from vllm.utils.argparse_utils import FlexibleArgumentParser -# Dimensions supported by the DSV3 specialized kernel -DSV3_SUPPORTED_NUM_EXPERTS = [256, 384] -DSV3_SUPPORTED_HIDDEN_SIZES = [7168] - # Dimensions supported by the gpt-oss specialized kernel GPT_OSS_SUPPORTED_NUM_EXPERTS = [32, 128] GPT_OSS_SUPPORTED_HIDDEN_SIZES = [2880] @@ -71,11 +67,6 @@ def get_benchmark(model, max_batch_size, trust_remote_code): is_hopper_or_blackwell = current_platform.is_device_capability( 90 ) or current_platform.is_device_capability_family(100) - allow_dsv3_router_gemm = ( - is_hopper_or_blackwell - and num_experts in DSV3_SUPPORTED_NUM_EXPERTS - and hidden_size in DSV3_SUPPORTED_HIDDEN_SIZES - ) allow_gpt_oss_router_gemm = ( is_hopper_or_blackwell and num_experts in GPT_OSS_SUPPORTED_NUM_EXPERTS @@ -116,9 +107,7 @@ def get_benchmark(model, max_batch_size, trust_remote_code): elif provider == "vllm": def runner(): - if allow_dsv3_router_gemm: - ops.dsv3_router_gemm(mat_a, mat_b, torch.bfloat16) - elif allow_fp32_router_gemm: + if allow_fp32_router_gemm: ops.fp32_router_gemm(mat_a, mat_b) elif allow_gpt_oss_router_gemm: ops.gpt_oss_router_gemm(mat_a, mat_b, bias) diff --git a/csrc/libtorch_stable/fp32_router_gemm.cu b/csrc/libtorch_stable/fp32_router_gemm.cu index 3207082d57..9719adb513 100644 --- a/csrc/libtorch_stable/fp32_router_gemm.cu +++ b/csrc/libtorch_stable/fp32_router_gemm.cu @@ -4,9 +4,6 @@ // 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 @@ -215,34 +212,10 @@ 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 && + if constexpr (std::is_same_v && kNumExperts == 128 && 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) @@ -296,8 +269,7 @@ void invokeFp32RouterGemm(float* output, InputT const* mat_a, // --------------------------------------------------------------------------- // Explicit instantiations: M=1..32, for both input types, for the supported -// (E, H) pairs: (256, 3072) [MiniMax-M2/M2.5], (128, 6144) [MiniMax-M3] -// and (256, 6144) [GLM-5.2]. +// (E, H) pairs: (256, 3072) [MiniMax-M2/M2.5] and (128, 6144) [MiniMax-M3]. // --------------------------------------------------------------------------- #define INSTANTIATE(T, M, E, H) \ @@ -342,8 +314,6 @@ 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 8643e33cbe..fc09193c86 100644 --- a/csrc/libtorch_stable/fp32_router_gemm_entry.cu +++ b/csrc/libtorch_stable/fp32_router_gemm_entry.cu @@ -25,12 +25,10 @@ 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, -// (6144, 256) for GLM-5.2. +// fp32_router_gemm.cu): (3072, 256) for MiniMax-M2/M2.5, (6144, 128) for M3. 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 == 256); + (hidden_dim == 6144 && num_experts == 128); } // Forward declarations — 4 template params must match fp32_router_gemm.cu @@ -79,9 +77,6 @@ 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"); @@ -116,7 +111,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), (6144, 128) and (6144, 256)"); + "(3072, 256) and (6144, 128)"); STD_TORCH_CHECK(num_tokens <= FP32_MAX_TOKENS, "fp32_router_gemm: num_tokens must be in [0, 32]"); STD_TORCH_CHECK( diff --git a/csrc/libtorch_stable/moe/dsv3_router_gemm_bf16_out.cu b/csrc/libtorch_stable/moe/dsv3_router_gemm_bf16_out.cu deleted file mode 100644 index 8695d1e808..0000000000 --- a/csrc/libtorch_stable/moe/dsv3_router_gemm_bf16_out.cu +++ /dev/null @@ -1,337 +0,0 @@ -/* - * Adapted from SGLang's sgl-kernel implementation, which was adapted from - * https://github.com/NVIDIA/TensorRT-LLM/blob/main/cpp/tensorrt_llm/kernels/dsv3MinLatencyKernels/dsv3RouterGemm.cu - * https://github.com/NVIDIA/TensorRT-LLM/blob/main/cpp/tensorrt_llm/thop/dsv3RouterGemmOp.cpp - * - * Copyright (c) 2019-2023, NVIDIA CORPORATION. All rights reserved. - * - * Licensed under the Apache License, Version 2.0 (the "License"); - * you may not use this file except in compliance with the License. - * You may obtain a copy of the License at - * - * http://www.apache.org/licenses/LICENSE-2.0 - * - * Unless required by applicable law or agreed to in writing, software - * distributed under the License is distributed on an "AS IS" BASIS, - * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. - * See the License for the specific language governing permissions and - * limitations under the License. - */ - -#include - -#include -#include - -// Custom FMA implementation using PTX assembly instructions -__device__ __forceinline__ void fma(float2& d, float2 const& a, float2 const& b, - float2 const& c) { - asm volatile("fma.rn.f32x2 %0, %1, %2, %3;\n" - : "=l"(reinterpret_cast(d)) - : "l"(reinterpret_cast(a)), - "l"(reinterpret_cast(b)), - "l"(reinterpret_cast(c))); -} - -// Convert 8 bfloat16 values from a uint4 to float array - optimized conversion -template -__device__ __forceinline__ void bf16_uint4_to_float8(uint4 const& vec, - float* dst) { - __nv_bfloat16* bf16_ptr = - reinterpret_cast<__nv_bfloat16*>(const_cast(&vec)); - -#pragma unroll - for (int i = 0; i < VPT; i++) { - dst[i] = __bfloat162float(bf16_ptr[i]); - } -} - -template -__global__ __launch_bounds__(128, 1) void router_gemm_kernel_bf16_output( - __nv_bfloat16* out, T const* mat_a, T const* mat_b) { - // Each block handles one expert column - int const n_idx = blockIdx.x; - int const tid = threadIdx.x; - constexpr int kWarpSize = 32; - constexpr int kNumWarps = kBlockSize / kWarpSize; - // Constants for this kernel - constexpr int k_elems_per_k_iteration = VPT * kBlockSize; - constexpr int k_iterations = - kHiddenDim / k_elems_per_k_iteration; // Total K iterations - - // Initialize accumulators for all M rows - float acc[kNumTokens] = {}; - - // Shared memory for warp-level reduction - __shared__ float sm_reduction[kNumTokens][kNumWarps]; // kNumWarps - - // B matrix is in column-major order, so we can directly load a column for the - // n_idx expert - T const* b_col = mat_b + n_idx * kHiddenDim; - - // Pre-compute k_base values for each iteration to help compiler optimize - 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; - } - -#if (defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 900)) - cudaGridDependencySynchronize(); -#endif - - // Process the GEMM in chunks - for (int ki = 0; ki < k_iterations; ki++) { - int const k_base = k_bases[ki]; - - // Load B matrix values using vector load (8 bf16 values) - uint4 b_vec = *reinterpret_cast(b_col + k_base); - - // Convert B values to float - float b_float[VPT]; - bf16_uint4_to_float8(b_vec, b_float); - -// Process each token -#pragma unroll - for (int m_idx = 0; m_idx < kNumTokens; m_idx++) { - // Load both rows of A matrix using vector loads - uint4 a_vec = *reinterpret_cast( - mat_a + (m_idx * kHiddenDim) + k_base); - - // Convert A values to float - float a_float[VPT]; - bf16_uint4_to_float8(a_vec, a_float); - -// Process elements in this chunk -#pragma unroll - for (int k = 0; k < VPT; k++) { - float a = a_float[k]; - float b = b_float[k]; - acc[m_idx] += a * b; - } - } - } - - // Perform warp-level reduction - int const warpSize = 32; - int const warpId = tid / warpSize; - int const laneId = tid % warpSize; - - // Register for warp-level reduction results - float warp_result[kNumTokens]; - -#pragma unroll - for (int m_idx = 0; m_idx < kNumTokens; m_idx++) { - warp_result[m_idx] = acc[m_idx]; - } - -// Perform warp-level reduction using optimized butterfly pattern -#pragma unroll - for (int m = 0; m < kNumTokens; m++) { - float sum = warp_result[m]; - - // Butterfly reduction pattern - 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); - - // Only the first thread in each warp stores to shared memory - if (laneId == 0) { - sm_reduction[m][warpId] = sum; - } - } - - __syncthreads(); - - // Final reduction across warps (only first thread) - if (tid == 0) { -#pragma unroll - for (int m = 0; m < kNumTokens; m++) { - float final_sum = 0.0f; - -// Sum across the kNumWarps -#pragma unroll - for (int w = 0; w < kNumWarps; w++) { - final_sum += sm_reduction[m][w]; - } - - // Write final result - out[m * kNumExperts + n_idx] = __float2bfloat16(final_sum); - } - } -#if (defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 900)) - cudaTriggerProgrammaticLaunchCompletion(); -#endif -} - -template -void invokeRouterGemmBf16Output(__nv_bfloat16* output, T const* mat_a, - T const* mat_b, cudaStream_t stream) { - constexpr int VPT = 16 / sizeof(T); - constexpr int kBlockSize = 128; - cudaLaunchConfig_t config; - config.gridDim = kNumExperts; - config.blockDim = kBlockSize; - config.dynamicSmemBytes = 0; - config.stream = stream; - cudaLaunchAttribute attrs[1]; - attrs[0].id = cudaLaunchAttributeProgrammaticStreamSerialization; - attrs[0].val.programmaticStreamSerializationAllowed = 1; - config.numAttrs = 1; - config.attrs = attrs; - cudaLaunchKernelEx( - &config, - router_gemm_kernel_bf16_output, - output, mat_a, mat_b); -} - -// Template instantiations for DEFAULT_NUM_EXPERTS experts -template void invokeRouterGemmBf16Output<__nv_bfloat16, 1, 256, 7168>( - __nv_bfloat16*, __nv_bfloat16 const*, __nv_bfloat16 const*, cudaStream_t); - -template void invokeRouterGemmBf16Output<__nv_bfloat16, 2, 256, 7168>( - __nv_bfloat16*, __nv_bfloat16 const*, __nv_bfloat16 const*, cudaStream_t); - -template void invokeRouterGemmBf16Output<__nv_bfloat16, 3, 256, 7168>( - __nv_bfloat16*, __nv_bfloat16 const*, __nv_bfloat16 const*, cudaStream_t); - -template void invokeRouterGemmBf16Output<__nv_bfloat16, 4, 256, 7168>( - __nv_bfloat16*, __nv_bfloat16 const*, __nv_bfloat16 const*, cudaStream_t); - -template void invokeRouterGemmBf16Output<__nv_bfloat16, 5, 256, 7168>( - __nv_bfloat16*, __nv_bfloat16 const*, __nv_bfloat16 const*, cudaStream_t); - -template void invokeRouterGemmBf16Output<__nv_bfloat16, 6, 256, 7168>( - __nv_bfloat16*, __nv_bfloat16 const*, __nv_bfloat16 const*, cudaStream_t); - -template void invokeRouterGemmBf16Output<__nv_bfloat16, 7, 256, 7168>( - __nv_bfloat16*, __nv_bfloat16 const*, __nv_bfloat16 const*, cudaStream_t); - -template void invokeRouterGemmBf16Output<__nv_bfloat16, 8, 256, 7168>( - __nv_bfloat16*, __nv_bfloat16 const*, __nv_bfloat16 const*, cudaStream_t); - -template void invokeRouterGemmBf16Output<__nv_bfloat16, 9, 256, 7168>( - __nv_bfloat16*, __nv_bfloat16 const*, __nv_bfloat16 const*, cudaStream_t); - -template void invokeRouterGemmBf16Output<__nv_bfloat16, 10, 256, 7168>( - __nv_bfloat16*, __nv_bfloat16 const*, __nv_bfloat16 const*, cudaStream_t); - -template void invokeRouterGemmBf16Output<__nv_bfloat16, 11, 256, 7168>( - __nv_bfloat16*, __nv_bfloat16 const*, __nv_bfloat16 const*, cudaStream_t); - -template void invokeRouterGemmBf16Output<__nv_bfloat16, 12, 256, 7168>( - __nv_bfloat16*, __nv_bfloat16 const*, __nv_bfloat16 const*, cudaStream_t); - -template void invokeRouterGemmBf16Output<__nv_bfloat16, 13, 256, 7168>( - __nv_bfloat16*, __nv_bfloat16 const*, __nv_bfloat16 const*, cudaStream_t); - -template void invokeRouterGemmBf16Output<__nv_bfloat16, 14, 256, 7168>( - __nv_bfloat16*, __nv_bfloat16 const*, __nv_bfloat16 const*, cudaStream_t); - -template void invokeRouterGemmBf16Output<__nv_bfloat16, 15, 256, 7168>( - __nv_bfloat16*, __nv_bfloat16 const*, __nv_bfloat16 const*, cudaStream_t); - -template void invokeRouterGemmBf16Output<__nv_bfloat16, 16, 256, 7168>( - __nv_bfloat16*, __nv_bfloat16 const*, __nv_bfloat16 const*, cudaStream_t); - -// Template instantiations for KIMI_K2_NUM_EXPERTS experts -template void invokeRouterGemmBf16Output<__nv_bfloat16, 1, 384, 7168>( - __nv_bfloat16*, __nv_bfloat16 const*, __nv_bfloat16 const*, cudaStream_t); - -template void invokeRouterGemmBf16Output<__nv_bfloat16, 2, 384, 7168>( - __nv_bfloat16*, __nv_bfloat16 const*, __nv_bfloat16 const*, cudaStream_t); - -template void invokeRouterGemmBf16Output<__nv_bfloat16, 3, 384, 7168>( - __nv_bfloat16*, __nv_bfloat16 const*, __nv_bfloat16 const*, cudaStream_t); - -template void invokeRouterGemmBf16Output<__nv_bfloat16, 4, 384, 7168>( - __nv_bfloat16*, __nv_bfloat16 const*, __nv_bfloat16 const*, cudaStream_t); - -template void invokeRouterGemmBf16Output<__nv_bfloat16, 5, 384, 7168>( - __nv_bfloat16*, __nv_bfloat16 const*, __nv_bfloat16 const*, cudaStream_t); - -template void invokeRouterGemmBf16Output<__nv_bfloat16, 6, 384, 7168>( - __nv_bfloat16*, __nv_bfloat16 const*, __nv_bfloat16 const*, cudaStream_t); - -template void invokeRouterGemmBf16Output<__nv_bfloat16, 7, 384, 7168>( - __nv_bfloat16*, __nv_bfloat16 const*, __nv_bfloat16 const*, cudaStream_t); - -template void invokeRouterGemmBf16Output<__nv_bfloat16, 8, 384, 7168>( - __nv_bfloat16*, __nv_bfloat16 const*, __nv_bfloat16 const*, cudaStream_t); - -template void invokeRouterGemmBf16Output<__nv_bfloat16, 9, 384, 7168>( - __nv_bfloat16*, __nv_bfloat16 const*, __nv_bfloat16 const*, cudaStream_t); - -template void invokeRouterGemmBf16Output<__nv_bfloat16, 10, 384, 7168>( - __nv_bfloat16*, __nv_bfloat16 const*, __nv_bfloat16 const*, cudaStream_t); - -template void invokeRouterGemmBf16Output<__nv_bfloat16, 11, 384, 7168>( - __nv_bfloat16*, __nv_bfloat16 const*, __nv_bfloat16 const*, cudaStream_t); - -template void invokeRouterGemmBf16Output<__nv_bfloat16, 12, 384, 7168>( - __nv_bfloat16*, __nv_bfloat16 const*, __nv_bfloat16 const*, cudaStream_t); - -template void invokeRouterGemmBf16Output<__nv_bfloat16, 13, 384, 7168>( - __nv_bfloat16*, __nv_bfloat16 const*, __nv_bfloat16 const*, cudaStream_t); - -template void invokeRouterGemmBf16Output<__nv_bfloat16, 14, 384, 7168>( - __nv_bfloat16*, __nv_bfloat16 const*, __nv_bfloat16 const*, cudaStream_t); - -template void invokeRouterGemmBf16Output<__nv_bfloat16, 15, 384, 7168>( - __nv_bfloat16*, __nv_bfloat16 const*, __nv_bfloat16 const*, cudaStream_t); - -template void invokeRouterGemmBf16Output<__nv_bfloat16, 16, 384, 7168>( - __nv_bfloat16*, __nv_bfloat16 const*, __nv_bfloat16 const*, cudaStream_t); - -// Template instantiations for GLM-5 (DEFAULT_NUM_EXPERTS, hidden_dim=6144) -template void invokeRouterGemmBf16Output<__nv_bfloat16, 1, 256, 6144>( - __nv_bfloat16*, __nv_bfloat16 const*, __nv_bfloat16 const*, cudaStream_t); - -template void invokeRouterGemmBf16Output<__nv_bfloat16, 2, 256, 6144>( - __nv_bfloat16*, __nv_bfloat16 const*, __nv_bfloat16 const*, cudaStream_t); - -template void invokeRouterGemmBf16Output<__nv_bfloat16, 3, 256, 6144>( - __nv_bfloat16*, __nv_bfloat16 const*, __nv_bfloat16 const*, cudaStream_t); - -template void invokeRouterGemmBf16Output<__nv_bfloat16, 4, 256, 6144>( - __nv_bfloat16*, __nv_bfloat16 const*, __nv_bfloat16 const*, cudaStream_t); - -template void invokeRouterGemmBf16Output<__nv_bfloat16, 5, 256, 6144>( - __nv_bfloat16*, __nv_bfloat16 const*, __nv_bfloat16 const*, cudaStream_t); - -template void invokeRouterGemmBf16Output<__nv_bfloat16, 6, 256, 6144>( - __nv_bfloat16*, __nv_bfloat16 const*, __nv_bfloat16 const*, cudaStream_t); - -template void invokeRouterGemmBf16Output<__nv_bfloat16, 7, 256, 6144>( - __nv_bfloat16*, __nv_bfloat16 const*, __nv_bfloat16 const*, cudaStream_t); - -template void invokeRouterGemmBf16Output<__nv_bfloat16, 8, 256, 6144>( - __nv_bfloat16*, __nv_bfloat16 const*, __nv_bfloat16 const*, cudaStream_t); - -template void invokeRouterGemmBf16Output<__nv_bfloat16, 9, 256, 6144>( - __nv_bfloat16*, __nv_bfloat16 const*, __nv_bfloat16 const*, cudaStream_t); - -template void invokeRouterGemmBf16Output<__nv_bfloat16, 10, 256, 6144>( - __nv_bfloat16*, __nv_bfloat16 const*, __nv_bfloat16 const*, cudaStream_t); - -template void invokeRouterGemmBf16Output<__nv_bfloat16, 11, 256, 6144>( - __nv_bfloat16*, __nv_bfloat16 const*, __nv_bfloat16 const*, cudaStream_t); - -template void invokeRouterGemmBf16Output<__nv_bfloat16, 12, 256, 6144>( - __nv_bfloat16*, __nv_bfloat16 const*, __nv_bfloat16 const*, cudaStream_t); - -template void invokeRouterGemmBf16Output<__nv_bfloat16, 13, 256, 6144>( - __nv_bfloat16*, __nv_bfloat16 const*, __nv_bfloat16 const*, cudaStream_t); - -template void invokeRouterGemmBf16Output<__nv_bfloat16, 14, 256, 6144>( - __nv_bfloat16*, __nv_bfloat16 const*, __nv_bfloat16 const*, cudaStream_t); - -template void invokeRouterGemmBf16Output<__nv_bfloat16, 15, 256, 6144>( - __nv_bfloat16*, __nv_bfloat16 const*, __nv_bfloat16 const*, cudaStream_t); - -template void invokeRouterGemmBf16Output<__nv_bfloat16, 16, 256, 6144>( - __nv_bfloat16*, __nv_bfloat16 const*, __nv_bfloat16 const*, cudaStream_t); diff --git a/csrc/libtorch_stable/moe/dsv3_router_gemm_entry.cu b/csrc/libtorch_stable/moe/dsv3_router_gemm_entry.cu deleted file mode 100644 index c5ebaf11f1..0000000000 --- a/csrc/libtorch_stable/moe/dsv3_router_gemm_entry.cu +++ /dev/null @@ -1,210 +0,0 @@ -/* - * Adapted from SGLang's sgl-kernel implementation, which was adapted from - * https://github.com/NVIDIA/TensorRT-LLM/blob/main/cpp/tensorrt_llm/kernels/dsv3MinLatencyKernels/dsv3RouterGemm.cu - * https://github.com/NVIDIA/TensorRT-LLM/blob/main/cpp/tensorrt_llm/thop/dsv3RouterGemmOp.cpp - * - * Copyright (c) 2019-2023, NVIDIA CORPORATION. All rights reserved. - * - * Licensed under the Apache License, Version 2.0 (the "License"); - * you may not use this file except in compliance with the License. - * You may obtain a copy of the License at - * - * http://www.apache.org/licenses/LICENSE-2.0 - * - * Unless required by applicable law or agreed to in writing, software - * distributed under the License is distributed on an "AS IS" BASIS, - * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. - * See the License for the specific language governing permissions and - * limitations under the License. - */ - -#include -#include -#include - -#include "libtorch_stable/torch_utils.h" - -#include -#include - -#include - -namespace { - -inline int getSMVersion() { - auto* props = get_device_prop(); - return props->major * 10 + props->minor; -} - -} // namespace - -static constexpr int DEFAULT_NUM_EXPERTS = 256; -static constexpr int KIMI_K2_NUM_EXPERTS = 384; -static constexpr int DEFAULT_HIDDEN_DIM = 7168; -static constexpr int GLM_5_HIDDEN_DIM = 6144; - -template -void invokeRouterGemmFloatOutput(float* output, T const* mat_a, T const* mat_b, - cudaStream_t stream); - -template -void invokeRouterGemmBf16Output(__nv_bfloat16* output, T const* mat_a, - T const* mat_b, cudaStream_t stream); - -template -struct LoopUnroller { - static void unroll_float_output(int num_tokens, float* output, - __nv_bfloat16 const* input, - __nv_bfloat16 const* weights, - cudaStream_t stream) { - if (num_tokens == kBegin) { - invokeRouterGemmFloatOutput<__nv_bfloat16, kBegin, kNumExperts, - kHiddenDim>(output, input, weights, stream); - } else { - LoopUnroller::unroll_float_output(num_tokens, output, input, - weights, stream); - } - } - - static void unroll_bf16_output(int num_tokens, __nv_bfloat16* output, - __nv_bfloat16 const* input, - __nv_bfloat16 const* weights, - cudaStream_t stream) { - if (num_tokens == kBegin) { - invokeRouterGemmBf16Output<__nv_bfloat16, kBegin, kNumExperts, - kHiddenDim>(output, input, weights, stream); - } else { - LoopUnroller::unroll_bf16_output(num_tokens, output, input, - weights, stream); - } - } -}; - -template -struct LoopUnroller { - static void unroll_float_output(int num_tokens, float* output, - __nv_bfloat16 const* input, - __nv_bfloat16 const* weights, - cudaStream_t stream) { - if (num_tokens == kEnd) { - invokeRouterGemmFloatOutput<__nv_bfloat16, kEnd, kNumExperts, kHiddenDim>( - output, input, weights, stream); - } else { - throw std::invalid_argument("Invalid num_tokens, only supports 1 to 16"); - } - } - - static void unroll_bf16_output(int num_tokens, __nv_bfloat16* output, - __nv_bfloat16 const* input, - __nv_bfloat16 const* weights, - cudaStream_t stream) { - if (num_tokens == kEnd) { - invokeRouterGemmBf16Output<__nv_bfloat16, kEnd, kNumExperts, kHiddenDim>( - output, input, weights, stream); - } else { - throw std::invalid_argument("Invalid num_tokens, only supports 1 to 16"); - } - } -}; - -void dsv3_router_gemm( - torch::stable::Tensor& output, // [num_tokens, num_experts] - torch::stable::Tensor const& mat_a, // [num_tokens, hidden_dim] - torch::stable::Tensor const& mat_b // [num_experts, hidden_dim] -) { - STD_TORCH_CHECK(output.dim() == 2 && mat_a.dim() == 2 && mat_b.dim() == 2); - - const int num_tokens = mat_a.size(0); - const int num_experts = mat_b.size(0); - const int hidden_dim = mat_a.size(1); - - STD_TORCH_CHECK(mat_a.size(1) == mat_b.size(1), - "mat_a and mat_b must have the same hidden_dim"); - STD_TORCH_CHECK( - hidden_dim == DEFAULT_HIDDEN_DIM || hidden_dim == GLM_5_HIDDEN_DIM, - "Expected hidden_dim=", DEFAULT_HIDDEN_DIM, - " or hidden_dim=", GLM_5_HIDDEN_DIM, ", but got hidden_dim=", hidden_dim); - STD_TORCH_CHECK( - num_experts == DEFAULT_NUM_EXPERTS || num_experts == KIMI_K2_NUM_EXPERTS, - "Expected num_experts=", DEFAULT_NUM_EXPERTS, - " or num_experts=", KIMI_K2_NUM_EXPERTS, - ", but got num_experts=", num_experts); - // KIMI_K2_NUM_EXPERTS is only instantiated for the default hidden_dim. - STD_TORCH_CHECK( - hidden_dim == DEFAULT_HIDDEN_DIM || num_experts == DEFAULT_NUM_EXPERTS, - "hidden_dim=", GLM_5_HIDDEN_DIM, - " only supports num_experts=", DEFAULT_NUM_EXPERTS, - ", but got num_experts=", num_experts); - STD_TORCH_CHECK(num_tokens >= 1 && num_tokens <= 16, - "currently num_tokens must be less than or equal to 16 for " - "router_gemm"); - STD_TORCH_CHECK( - mat_a.scalar_type() == torch::headeronly::ScalarType::BFloat16, - "mat_a must be bf16"); - STD_TORCH_CHECK( - mat_b.scalar_type() == torch::headeronly::ScalarType::BFloat16, - "mat_b must be bf16"); - STD_TORCH_CHECK( - output.scalar_type() == torch::headeronly::ScalarType::Float || - output.scalar_type() == torch::headeronly::ScalarType::BFloat16, - "output must be float32 or bf16"); - - const torch::stable::accelerator::DeviceGuard device_guard( - mat_a.get_device_index()); - const int sm = getSMVersion(); - STD_TORCH_CHECK(sm >= 90, "required CUDA ARCH >= SM_90"); - - const cudaStream_t stream = get_current_cuda_stream(mat_a.get_device_index()); - - __nv_bfloat16 const* a_ptr = - reinterpret_cast<__nv_bfloat16 const*>(mat_a.data_ptr()); - __nv_bfloat16 const* b_ptr = - reinterpret_cast<__nv_bfloat16 const*>(mat_b.data_ptr()); - - if (output.scalar_type() == torch::headeronly::ScalarType::Float) { - float* out_ptr = reinterpret_cast(output.mutable_data_ptr()); - if (hidden_dim == DEFAULT_HIDDEN_DIM) { - if (num_experts == DEFAULT_NUM_EXPERTS) { - LoopUnroller<1, 16, DEFAULT_NUM_EXPERTS, - DEFAULT_HIDDEN_DIM>::unroll_float_output(num_tokens, - out_ptr, a_ptr, - b_ptr, stream); - } else { - LoopUnroller<1, 16, KIMI_K2_NUM_EXPERTS, - DEFAULT_HIDDEN_DIM>::unroll_float_output(num_tokens, - out_ptr, a_ptr, - b_ptr, stream); - } - } else { // GLM_5_HIDDEN_DIM - LoopUnroller<1, 16, DEFAULT_NUM_EXPERTS, - GLM_5_HIDDEN_DIM>::unroll_float_output(num_tokens, out_ptr, - a_ptr, b_ptr, stream); - } - } else if (output.scalar_type() == torch::headeronly::ScalarType::BFloat16) { - __nv_bfloat16* out_ptr = - reinterpret_cast<__nv_bfloat16*>(output.mutable_data_ptr()); - if (hidden_dim == DEFAULT_HIDDEN_DIM) { - if (num_experts == DEFAULT_NUM_EXPERTS) { - LoopUnroller<1, 16, DEFAULT_NUM_EXPERTS, - DEFAULT_HIDDEN_DIM>::unroll_bf16_output(num_tokens, - out_ptr, a_ptr, - b_ptr, stream); - } else { - LoopUnroller<1, 16, KIMI_K2_NUM_EXPERTS, - DEFAULT_HIDDEN_DIM>::unroll_bf16_output(num_tokens, - out_ptr, a_ptr, - b_ptr, stream); - } - } else { // GLM_5_HIDDEN_DIM - LoopUnroller<1, 16, DEFAULT_NUM_EXPERTS, - GLM_5_HIDDEN_DIM>::unroll_bf16_output(num_tokens, out_ptr, - a_ptr, b_ptr, stream); - } - } -} - -STABLE_TORCH_LIBRARY_IMPL(_moe_C, CUDA, m) { - m.impl("dsv3_router_gemm", TORCH_BOX(&dsv3_router_gemm)); -} diff --git a/csrc/libtorch_stable/moe/dsv3_router_gemm_float_out.cu b/csrc/libtorch_stable/moe/dsv3_router_gemm_float_out.cu deleted file mode 100644 index 58a2b44ae2..0000000000 --- a/csrc/libtorch_stable/moe/dsv3_router_gemm_float_out.cu +++ /dev/null @@ -1,337 +0,0 @@ -/* - * Adapted from SGLang's sgl-kernel implementation, which was adapted from - * https://github.com/NVIDIA/TensorRT-LLM/blob/main/cpp/tensorrt_llm/kernels/dsv3MinLatencyKernels/dsv3RouterGemm.cu - * https://github.com/NVIDIA/TensorRT-LLM/blob/main/cpp/tensorrt_llm/thop/dsv3RouterGemmOp.cpp - * - * Copyright (c) 2019-2023, NVIDIA CORPORATION. All rights reserved. - * - * Licensed under the Apache License, Version 2.0 (the "License"); - * you may not use this file except in compliance with the License. - * You may obtain a copy of the License at - * - * http://www.apache.org/licenses/LICENSE-2.0 - * - * Unless required by applicable law or agreed to in writing, software - * distributed under the License is distributed on an "AS IS" BASIS, - * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. - * See the License for the specific language governing permissions and - * limitations under the License. - */ - -#include - -#include -#include - -// Custom FMA implementation using PTX assembly instructions -__device__ __forceinline__ void fma(float2& d, float2 const& a, float2 const& b, - float2 const& c) { - asm volatile("fma.rn.f32x2 %0, %1, %2, %3;\n" - : "=l"(reinterpret_cast(d)) - : "l"(reinterpret_cast(a)), - "l"(reinterpret_cast(b)), - "l"(reinterpret_cast(c))); -} - -// Convert 8 bfloat16 values from a uint4 to float array - optimized conversion -template -__device__ __forceinline__ void bf16_uint4_to_float8(uint4 const& vec, - float* dst) { - __nv_bfloat16* bf16_ptr = - reinterpret_cast<__nv_bfloat16*>(const_cast(&vec)); - -#pragma unroll - for (int i = 0; i < VPT; i++) { - dst[i] = __bfloat162float(bf16_ptr[i]); - } -} - -template -__global__ __launch_bounds__(128, 1) void router_gemm_kernel_float_output( - float* out, T const* mat_a, T const* mat_b) { - // Each block handles one expert column - int const n_idx = blockIdx.x; - int const tid = threadIdx.x; - constexpr int kWarpSize = 32; - constexpr int kNumWarps = kBlockSize / kWarpSize; - // Constants for this kernel - constexpr int k_elems_per_k_iteration = VPT * kBlockSize; - constexpr int k_iterations = - kHiddenDim / k_elems_per_k_iteration; // Total K iterations - - // Initialize accumulators for all M rows - float acc[kNumTokens] = {}; - - // Shared memory for warp-level reduction - __shared__ float sm_reduction[kNumTokens][kNumWarps]; // kNumWarps - - // B matrix is in column-major order, so we can directly load a column for the - // n_idx expert - T const* b_col = mat_b + n_idx * kHiddenDim; - - // Pre-compute k_base values for each iteration to help compiler optimize - 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; - } - -#if (defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 900)) - cudaGridDependencySynchronize(); -#endif - - // Process the GEMM in chunks - for (int ki = 0; ki < k_iterations; ki++) { - int const k_base = k_bases[ki]; - - // Load B matrix values using vector load (8 bf16 values) - uint4 b_vec = *reinterpret_cast(b_col + k_base); - - // Convert B values to float - float b_float[VPT]; - bf16_uint4_to_float8(b_vec, b_float); - -// Process each token -#pragma unroll - for (int m_idx = 0; m_idx < kNumTokens; m_idx++) { - // Load both rows of A matrix using vector loads - uint4 a_vec = *reinterpret_cast( - mat_a + (m_idx * kHiddenDim) + k_base); - - // Convert A values to float - float a_float[VPT]; - bf16_uint4_to_float8(a_vec, a_float); - -// Process elements in this chunk -#pragma unroll - for (int k = 0; k < VPT; k++) { - float a = a_float[k]; - float b = b_float[k]; - acc[m_idx] += a * b; - } - } - } - - // Perform warp-level reduction - int const warpSize = 32; - int const warpId = tid / warpSize; - int const laneId = tid % warpSize; - - // Register for warp-level reduction results - float warp_result[kNumTokens]; - -#pragma unroll - for (int m_idx = 0; m_idx < kNumTokens; m_idx++) { - warp_result[m_idx] = acc[m_idx]; - } - -// Perform warp-level reduction using optimized butterfly pattern -#pragma unroll - for (int m = 0; m < kNumTokens; m++) { - float sum = warp_result[m]; - - // Butterfly reduction pattern - 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); - - // Only the first thread in each warp stores to shared memory - if (laneId == 0) { - sm_reduction[m][warpId] = sum; - } - } - - __syncthreads(); - - // Final reduction across warps (only first thread) - if (tid == 0) { -#pragma unroll - for (int m = 0; m < kNumTokens; m++) { - float final_sum = 0.0f; - -// Sum across the kNumWarps -#pragma unroll - for (int w = 0; w < kNumWarps; w++) { - final_sum += sm_reduction[m][w]; - } - - // Write final result - out[m * kNumExperts + n_idx] = final_sum; - } - } -#if (defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 900)) - cudaTriggerProgrammaticLaunchCompletion(); -#endif -} - -template -void invokeRouterGemmFloatOutput(float* output, T const* mat_a, T const* mat_b, - cudaStream_t stream) { - constexpr int VPT = 16 / sizeof(T); - constexpr int kBlockSize = 128; - cudaLaunchConfig_t config; - config.gridDim = kNumExperts; - config.blockDim = kBlockSize; - config.dynamicSmemBytes = 0; - config.stream = stream; - cudaLaunchAttribute attrs[1]; - attrs[0].id = cudaLaunchAttributeProgrammaticStreamSerialization; - attrs[0].val.programmaticStreamSerializationAllowed = 1; - config.numAttrs = 1; - config.attrs = attrs; - cudaLaunchKernelEx( - &config, - router_gemm_kernel_float_output, - output, mat_a, mat_b); -} - -// Template instantiations for DEFAULT_NUM_EXPERTS experts -template void invokeRouterGemmFloatOutput<__nv_bfloat16, 1, 256, 7168>( - float*, __nv_bfloat16 const*, __nv_bfloat16 const*, cudaStream_t); - -template void invokeRouterGemmFloatOutput<__nv_bfloat16, 2, 256, 7168>( - float*, __nv_bfloat16 const*, __nv_bfloat16 const*, cudaStream_t); - -template void invokeRouterGemmFloatOutput<__nv_bfloat16, 3, 256, 7168>( - float*, __nv_bfloat16 const*, __nv_bfloat16 const*, cudaStream_t); - -template void invokeRouterGemmFloatOutput<__nv_bfloat16, 4, 256, 7168>( - float*, __nv_bfloat16 const*, __nv_bfloat16 const*, cudaStream_t); - -template void invokeRouterGemmFloatOutput<__nv_bfloat16, 5, 256, 7168>( - float*, __nv_bfloat16 const*, __nv_bfloat16 const*, cudaStream_t); - -template void invokeRouterGemmFloatOutput<__nv_bfloat16, 6, 256, 7168>( - float*, __nv_bfloat16 const*, __nv_bfloat16 const*, cudaStream_t); - -template void invokeRouterGemmFloatOutput<__nv_bfloat16, 7, 256, 7168>( - float*, __nv_bfloat16 const*, __nv_bfloat16 const*, cudaStream_t); - -template void invokeRouterGemmFloatOutput<__nv_bfloat16, 8, 256, 7168>( - float*, __nv_bfloat16 const*, __nv_bfloat16 const*, cudaStream_t); - -template void invokeRouterGemmFloatOutput<__nv_bfloat16, 9, 256, 7168>( - float*, __nv_bfloat16 const*, __nv_bfloat16 const*, cudaStream_t); - -template void invokeRouterGemmFloatOutput<__nv_bfloat16, 10, 256, 7168>( - float*, __nv_bfloat16 const*, __nv_bfloat16 const*, cudaStream_t); - -template void invokeRouterGemmFloatOutput<__nv_bfloat16, 11, 256, 7168>( - float*, __nv_bfloat16 const*, __nv_bfloat16 const*, cudaStream_t); - -template void invokeRouterGemmFloatOutput<__nv_bfloat16, 12, 256, 7168>( - float*, __nv_bfloat16 const*, __nv_bfloat16 const*, cudaStream_t); - -template void invokeRouterGemmFloatOutput<__nv_bfloat16, 13, 256, 7168>( - float*, __nv_bfloat16 const*, __nv_bfloat16 const*, cudaStream_t); - -template void invokeRouterGemmFloatOutput<__nv_bfloat16, 14, 256, 7168>( - float*, __nv_bfloat16 const*, __nv_bfloat16 const*, cudaStream_t); - -template void invokeRouterGemmFloatOutput<__nv_bfloat16, 15, 256, 7168>( - float*, __nv_bfloat16 const*, __nv_bfloat16 const*, cudaStream_t); - -template void invokeRouterGemmFloatOutput<__nv_bfloat16, 16, 256, 7168>( - float*, __nv_bfloat16 const*, __nv_bfloat16 const*, cudaStream_t); - -// Template instantiations for KIMI_K2_NUM_EXPERTS experts -template void invokeRouterGemmFloatOutput<__nv_bfloat16, 1, 384, 7168>( - float*, __nv_bfloat16 const*, __nv_bfloat16 const*, cudaStream_t); - -template void invokeRouterGemmFloatOutput<__nv_bfloat16, 2, 384, 7168>( - float*, __nv_bfloat16 const*, __nv_bfloat16 const*, cudaStream_t); - -template void invokeRouterGemmFloatOutput<__nv_bfloat16, 3, 384, 7168>( - float*, __nv_bfloat16 const*, __nv_bfloat16 const*, cudaStream_t); - -template void invokeRouterGemmFloatOutput<__nv_bfloat16, 4, 384, 7168>( - float*, __nv_bfloat16 const*, __nv_bfloat16 const*, cudaStream_t); - -template void invokeRouterGemmFloatOutput<__nv_bfloat16, 5, 384, 7168>( - float*, __nv_bfloat16 const*, __nv_bfloat16 const*, cudaStream_t); - -template void invokeRouterGemmFloatOutput<__nv_bfloat16, 6, 384, 7168>( - float*, __nv_bfloat16 const*, __nv_bfloat16 const*, cudaStream_t); - -template void invokeRouterGemmFloatOutput<__nv_bfloat16, 7, 384, 7168>( - float*, __nv_bfloat16 const*, __nv_bfloat16 const*, cudaStream_t); - -template void invokeRouterGemmFloatOutput<__nv_bfloat16, 8, 384, 7168>( - float*, __nv_bfloat16 const*, __nv_bfloat16 const*, cudaStream_t); - -template void invokeRouterGemmFloatOutput<__nv_bfloat16, 9, 384, 7168>( - float*, __nv_bfloat16 const*, __nv_bfloat16 const*, cudaStream_t); - -template void invokeRouterGemmFloatOutput<__nv_bfloat16, 10, 384, 7168>( - float*, __nv_bfloat16 const*, __nv_bfloat16 const*, cudaStream_t); - -template void invokeRouterGemmFloatOutput<__nv_bfloat16, 11, 384, 7168>( - float*, __nv_bfloat16 const*, __nv_bfloat16 const*, cudaStream_t); - -template void invokeRouterGemmFloatOutput<__nv_bfloat16, 12, 384, 7168>( - float*, __nv_bfloat16 const*, __nv_bfloat16 const*, cudaStream_t); - -template void invokeRouterGemmFloatOutput<__nv_bfloat16, 13, 384, 7168>( - float*, __nv_bfloat16 const*, __nv_bfloat16 const*, cudaStream_t); - -template void invokeRouterGemmFloatOutput<__nv_bfloat16, 14, 384, 7168>( - float*, __nv_bfloat16 const*, __nv_bfloat16 const*, cudaStream_t); - -template void invokeRouterGemmFloatOutput<__nv_bfloat16, 15, 384, 7168>( - float*, __nv_bfloat16 const*, __nv_bfloat16 const*, cudaStream_t); - -template void invokeRouterGemmFloatOutput<__nv_bfloat16, 16, 384, 7168>( - float*, __nv_bfloat16 const*, __nv_bfloat16 const*, cudaStream_t); - -// Template instantiations for GLM-5 (DEFAULT_NUM_EXPERTS, hidden_dim=6144) -template void invokeRouterGemmFloatOutput<__nv_bfloat16, 1, 256, 6144>( - float*, __nv_bfloat16 const*, __nv_bfloat16 const*, cudaStream_t); - -template void invokeRouterGemmFloatOutput<__nv_bfloat16, 2, 256, 6144>( - float*, __nv_bfloat16 const*, __nv_bfloat16 const*, cudaStream_t); - -template void invokeRouterGemmFloatOutput<__nv_bfloat16, 3, 256, 6144>( - float*, __nv_bfloat16 const*, __nv_bfloat16 const*, cudaStream_t); - -template void invokeRouterGemmFloatOutput<__nv_bfloat16, 4, 256, 6144>( - float*, __nv_bfloat16 const*, __nv_bfloat16 const*, cudaStream_t); - -template void invokeRouterGemmFloatOutput<__nv_bfloat16, 5, 256, 6144>( - float*, __nv_bfloat16 const*, __nv_bfloat16 const*, cudaStream_t); - -template void invokeRouterGemmFloatOutput<__nv_bfloat16, 6, 256, 6144>( - float*, __nv_bfloat16 const*, __nv_bfloat16 const*, cudaStream_t); - -template void invokeRouterGemmFloatOutput<__nv_bfloat16, 7, 256, 6144>( - float*, __nv_bfloat16 const*, __nv_bfloat16 const*, cudaStream_t); - -template void invokeRouterGemmFloatOutput<__nv_bfloat16, 8, 256, 6144>( - float*, __nv_bfloat16 const*, __nv_bfloat16 const*, cudaStream_t); - -template void invokeRouterGemmFloatOutput<__nv_bfloat16, 9, 256, 6144>( - float*, __nv_bfloat16 const*, __nv_bfloat16 const*, cudaStream_t); - -template void invokeRouterGemmFloatOutput<__nv_bfloat16, 10, 256, 6144>( - float*, __nv_bfloat16 const*, __nv_bfloat16 const*, cudaStream_t); - -template void invokeRouterGemmFloatOutput<__nv_bfloat16, 11, 256, 6144>( - float*, __nv_bfloat16 const*, __nv_bfloat16 const*, cudaStream_t); - -template void invokeRouterGemmFloatOutput<__nv_bfloat16, 12, 256, 6144>( - float*, __nv_bfloat16 const*, __nv_bfloat16 const*, cudaStream_t); - -template void invokeRouterGemmFloatOutput<__nv_bfloat16, 13, 256, 6144>( - float*, __nv_bfloat16 const*, __nv_bfloat16 const*, cudaStream_t); - -template void invokeRouterGemmFloatOutput<__nv_bfloat16, 14, 256, 6144>( - float*, __nv_bfloat16 const*, __nv_bfloat16 const*, cudaStream_t); - -template void invokeRouterGemmFloatOutput<__nv_bfloat16, 15, 256, 6144>( - float*, __nv_bfloat16 const*, __nv_bfloat16 const*, cudaStream_t); - -template void invokeRouterGemmFloatOutput<__nv_bfloat16, 16, 256, 6144>( - float*, __nv_bfloat16 const*, __nv_bfloat16 const*, cudaStream_t); diff --git a/csrc/libtorch_stable/moe/moe_ops.h b/csrc/libtorch_stable/moe/moe_ops.h index 03525a05f9..7c3ca10371 100644 --- a/csrc/libtorch_stable/moe/moe_ops.h +++ b/csrc/libtorch_stable/moe/moe_ops.h @@ -79,15 +79,3 @@ int64_t moe_permute_sort_workspace_size(int64_t num_expanded_rows, void shuffle_rows(const torch::stable::Tensor& input_tensor, const torch::stable::Tensor& dst2src_map, torch::stable::Tensor& output_tensor); - -#ifndef USE_ROCM -// DeepSeek V3 optimized router GEMM kernel for SM90+ -// Computes output = mat_a @ mat_b.T where: -// mat_a: [num_tokens, hidden_dim] in bf16 -// mat_b: [num_experts, hidden_dim] in bf16 -// output: [num_tokens, num_experts] in bf16 or fp32 -// Supports num_tokens in [1, 16], num_experts in {256, 384}, hidden_dim = 7168 -void dsv3_router_gemm(torch::stable::Tensor& output, - const torch::stable::Tensor& mat_a, - const torch::stable::Tensor& mat_b); -#endif diff --git a/csrc/libtorch_stable/moe/torch_bindings.cpp b/csrc/libtorch_stable/moe/torch_bindings.cpp index c6467f1c1f..a321043848 100644 --- a/csrc/libtorch_stable/moe/torch_bindings.cpp +++ b/csrc/libtorch_stable/moe/torch_bindings.cpp @@ -126,10 +126,6 @@ STABLE_TORCH_LIBRARY_FRAGMENT(_moe_C, m) { "topk_group, int topk, bool renormalize, float " "routed_scaling_factor, Tensor bias, int scoring_func) -> (Tensor, " "Tensor)"); - - // DeepSeek V3 optimized router GEMM for SM90+ - m.def("dsv3_router_gemm(Tensor! output, Tensor mat_a, Tensor mat_b) -> ()"); - // conditionally compiled so impl registration is in source file #endif } diff --git a/tests/kernels/test_fp32_router_gemm.py b/tests/kernels/test_fp32_router_gemm.py index 6384cae7cd..d330abba53 100644 --- a/tests/kernels/test_fp32_router_gemm.py +++ b/tests/kernels/test_fp32_router_gemm.py @@ -3,8 +3,7 @@ """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, - (6144, 256) -> GLM-5.2 + (3072, 256) -> MiniMax-M2/M2.5, (6144, 128) -> MiniMax-M3 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 @@ -19,7 +18,7 @@ from vllm._custom_ops import fp32_router_gemm from vllm.platforms import current_platform # (hidden_size, num_experts) -SHAPES = [(3072, 256), (6144, 128), (6144, 256)] +SHAPES = [(3072, 256), (6144, 128)] ALL_M = list(range(1, 33)) # Absolute tolerance for fp32 kernel vs float64 reference ATOL_FP32 = 2e-4 diff --git a/vllm/_custom_ops.py b/vllm/_custom_ops.py index 4d8c1e40a5..dba31ef3e2 100644 --- a/vllm/_custom_ops.py +++ b/vllm/_custom_ops.py @@ -2384,21 +2384,6 @@ def moe_wna16_gemm( ) -def dsv3_router_gemm( - hidden_states: torch.Tensor, - router_weight: torch.Tensor, - output_dtype: torch.dtype, -) -> torch.Tensor: - output = torch.empty( - hidden_states.shape[0], - router_weight.shape[0], - device=hidden_states.device, - dtype=output_dtype, - ) - torch.ops._moe_C.dsv3_router_gemm(output, hidden_states, router_weight) - return output - - def fp32_router_gemm( hidden_states: torch.Tensor, router_weight: torch.Tensor, 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 308d8520f5..0e4a402599 100644 --- a/vllm/model_executor/layers/fused_moe/router/gate_linear.py +++ b/vllm/model_executor/layers/fused_moe/router/gate_linear.py @@ -20,30 +20,20 @@ class GateLinear(ReplicatedLinear): 1. cuteDSL ll_bf16_gemm (SM90+, M<=16, bf16 in, fp32 out, K divisible by 8) - 2. DSV3 specialized kernel (SM90+, M<=16, H=7168 E=256/384, H=6144 E=256) - 3. fp32 specialized kernel (SM90+, bf16/fp32 in, fp32 out, M<=32, - (H, E) in {(3072, 256), (6144, 128), (6144, 256)}) - 4. experimental bf16x3 CuteDSL kernel (opt-in, SM100, bf16 in, fp32 weight) - 5. cuBLAS bf16×bf16→fp32 (SM90+ + bf16 weight + fp32 out_dtype) - 6. F.linear via ReplicatedLinear (ultimate fallback) + 2. fp32 specialized kernel (SM90+, bf16/fp32 in, fp32 out, M<=32, + (H, E) in {(3072, 256), (6144, 128)}) + 3. experimental bf16x3 CuteDSL kernel (opt-in, SM100, bf16 in, fp32 weight) + 4. cuBLAS bf16×bf16→fp32 (SM90+ + bf16 weight + fp32 out_dtype) + 5. F.linear via ReplicatedLinear (ultimate fallback) The ``out_dtype`` attribute is mutable and can be set after init (e.g. when the required dtype depends on the expert quantization method which is only known later). """ - # Dimensions supported by the DSV3 specialized kernel. - # Valid (hidden_size, num_experts) combinations: - # (7168, 256) -> DeepSeek-V3, (7168, 384) -> Kimi-K2, - # (6144, 256) -> GLM-5 - DSV3_SUPPORTED_NUM_EXPERTS = [256, 384] - DSV3_SUPPORTED_HIDDEN_SIZES = [7168, 6144] - # num_experts=384 is only instantiated for hidden_size=7168. - DSV3_UNSUPPORTED_SHAPES = {(6144, 384)} - # (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), (6144, 256)} + FP32_SUPPORTED_SHAPES = {(3072, 256), (6144, 128)} FP32_MAX_TOKENS = 32 def __init__( @@ -77,18 +67,7 @@ class GateLinear(ReplicatedLinear): ) self.out_dtype = out_dtype - # DSV3 specialized kernel eligibility (SM90+, exact dims) self.allow_specialized_router_gemm = can_use_specialized_kernels - self.allow_dsv3_router_gemm = ( - self.allow_specialized_router_gemm - and self.weight.dtype == torch.bfloat16 - and output_size in self.DSV3_SUPPORTED_NUM_EXPERTS - and input_size in self.DSV3_SUPPORTED_HIDDEN_SIZES - and (input_size, output_size) not in self.DSV3_UNSUPPORTED_SHAPES - ) - # See https://github.com/vllm-project/vllm/pull/44217 - # for more details. - self._dsv3_max_batch = 16 if is_hopper else 8 # fp32 specialized kernel eligibility (SM90+, exact dims, fp32 weight) vllm_config = get_current_vllm_config_or_none() @@ -187,16 +166,7 @@ class GateLinear(ReplicatedLinear): output = ll_bf16_gemm(x, self.weight) return output, None - # Tier 2: DSV3 specialized kernel (fallback for when cuteDSL unavailable) - if self.allow_dsv3_router_gemm and x.shape[0] <= self._dsv3_max_batch: - output = ops.dsv3_router_gemm( - hidden_states=x, - router_weight=self.weight, - output_dtype=self.out_dtype, - ) - return output, None - - # Tier 3: fp32 specialized kernel (H=3072, E=256, M<=32) + # Tier 2: fp32 specialized kernel (H=3072, E=256, M<=32) # Dispatch is wrapped in a custom op so that torch.compile/CUDA-graph # capture does not freeze the runtime num_tokens branch. if self.allow_fp32_router_gemm and x.dtype in ( @@ -208,7 +178,7 @@ class GateLinear(ReplicatedLinear): ) return output, None - # Tier 4: experimental bf16x3 CuteDSL kernel for fp32 router weights + # Tier 3: experimental bf16x3 CuteDSL kernel for fp32 router weights if self.allow_bf16x3_router_gemm and x.dtype == torch.bfloat16: from vllm.model_executor.layers.fused_moe.router.bf16x3_router_gemm_cutedsl import ( # noqa: E501 bf16x3_router_gemm, @@ -217,12 +187,12 @@ class GateLinear(ReplicatedLinear): output = bf16x3_router_gemm(x, self.weight) return output, None - # Tier 5: cuBLAS bf16→fp32 + # Tier 4: cuBLAS bf16→fp32 if self.allow_cublas_router_gemm and x.dtype == torch.bfloat16: output = torch.mm(x, self.weight.T, out_dtype=torch.float32) return output, None - # Tier 6: F.linear (ReplicatedLinear) + # Tier 5: F.linear (ReplicatedLinear) if self.out_dtype is not None and x.dtype != self.weight.dtype: x = x.to(self.weight.dtype) output, output_bias = super().forward(x)