@tduncan / Nvidia-Optimization-Stack-VLLM Mirror · read-only / commits / 9bd7231106
Fork 0

A high-throughput and memory-efficient inference and serving engine for LLMs

Mirrored from https://github.com/tduncan-lum/Nvidia-Optimization-Stack---VLLM.git

Synced Nov 5, 16:00 UTC

Revert "[Kernel] Add gpt-oss Router GEMM kernel (#37205)" (#38778)

Signed-off-by: Xin Yang <xyangx@amazon.com>

Xin Yang committed Apr 1, 2026 at 22:02 UTC 9bd72311068919b8f3430278d47859cf312039fa
13 files changed +12 -875
CMakeLists.txt
-1
@@ -1030,7 +1030,6 @@ if(VLLM_GPU_LANG STREQUAL "CUDA")
1030 list(APPEND VLLM_MOE_EXT_SRC
1031 "csrc/moe/moe_wna16.cu"
1032 "csrc/moe/grouped_topk_kernels.cu"
1033 - "csrc/moe/gpt_oss_router_gemm.cu"
1033 "csrc/moe/router_gemm.cu")
1034 endif()
1035
benchmarks/kernels/benchmark_router_gemm.py deleted
-134
@@ -1,134 +0,0 @@
1 -# SPDX-License-Identifier: Apache-2.0
2 -# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 -
4 -import torch
5 -import torch.nn.functional as F
6 -
7 -from vllm import _custom_ops as ops
8 -from vllm.platforms import current_platform
9 -from vllm.transformers_utils.config import get_config
10 -from vllm.triton_utils import triton
11 -from vllm.utils.argparse_utils import FlexibleArgumentParser
12 -
13 -# Dimensions supported by the DSV3 specialized kernel
14 -DSV3_SUPPORTED_NUM_EXPERTS = [256, 384]
15 -DSV3_SUPPORTED_HIDDEN_SIZES = [7168]
16 -
17 -# Dimensions supported by the gpt-oss specialized kernel
18 -GPT_OSS_SUPPORTED_NUM_EXPERTS = [32, 128]
19 -GPT_OSS_SUPPORTED_HIDDEN_SIZES = [2880]
20 -
21 -
22 -def get_batch_size_range(max_batch_size):
23 - return [2**x for x in range(14) if 2**x <= max_batch_size]
24 -
25 -
26 -def get_model_params(config):
27 - if config.architectures[0] in (
28 - "DeepseekV2ForCausalLM",
29 - "DeepseekV3ForCausalLM",
30 - "DeepseekV32ForCausalLM",
31 - ):
32 - num_experts = config.n_routed_experts
33 - hidden_size = config.hidden_size
34 - elif config.architectures[0] in ("GptOssForCausalLM",):
35 - num_experts = config.num_local_experts
36 - hidden_size = config.hidden_size
37 - else:
38 - raise ValueError(f"Unsupported architecture: {config.architectures}")
39 - return num_experts, hidden_size
40 -
41 -
42 -def get_benchmark(model, max_batch_size, trust_remote_code):
43 - @triton.testing.perf_report(
44 - triton.testing.Benchmark(
45 - x_names=["batch_size"],
46 - x_vals=get_batch_size_range(max_batch_size),
47 - x_log=False,
48 - line_arg="provider",
49 - line_vals=[
50 - "torch",
51 - "vllm",
52 - ],
53 - line_names=["PyTorch", "vLLM"],
54 - styles=([("blue", "-"), ("red", "-")]),
55 - ylabel="TFLOPs",
56 - plot_name=f"{model} router gemm throughput",
57 - args={},
58 - )
59 - )
60 - def benchmark(batch_size, provider):
61 - config = get_config(model=model, trust_remote_code=trust_remote_code)
62 - num_experts, hidden_size = get_model_params(config)
63 -
64 - mat_a = torch.randn(
65 - (batch_size, hidden_size), dtype=torch.bfloat16, device="cuda"
66 - ).contiguous()
67 - mat_b = torch.randn(
68 - (num_experts, hidden_size), dtype=torch.bfloat16, device="cuda"
69 - ).contiguous()
70 - bias = torch.randn(
71 - num_experts, dtype=torch.bfloat16, device="cuda"
72 - ).contiguous()
73 -
74 - is_hopper_or_blackwell = current_platform.is_device_capability(
75 - 90
76 - ) or current_platform.is_device_capability_family(100)
77 - allow_dsv3_router_gemm = (
78 - is_hopper_or_blackwell
79 - and num_experts in DSV3_SUPPORTED_NUM_EXPERTS
80 - and hidden_size in DSV3_SUPPORTED_HIDDEN_SIZES
81 - )
82 - allow_gpt_oss_router_gemm = (
83 - is_hopper_or_blackwell
84 - and num_experts in GPT_OSS_SUPPORTED_NUM_EXPERTS
85 - and hidden_size in GPT_OSS_SUPPORTED_HIDDEN_SIZES
86 - )
87 -
88 - has_bias = False
89 - if allow_gpt_oss_router_gemm:
90 - has_bias = True
91 -
92 - quantiles = [0.5, 0.2, 0.8]
93 -
94 - if provider == "torch":
95 -
96 - def runner():
97 - if has_bias:
98 - F.linear(mat_a, mat_b, bias)
99 - else:
100 - F.linear(mat_a, mat_b)
101 - elif provider == "vllm":
102 -
103 - def runner():
104 - if allow_dsv3_router_gemm:
105 - ops.dsv3_router_gemm(mat_a, mat_b, torch.bfloat16)
106 - elif allow_gpt_oss_router_gemm:
107 - ops.gpt_oss_router_gemm(mat_a, mat_b, bias)
108 - else:
109 - raise ValueError("Unsupported router gemm")
110 -
111 - ms, min_ms, max_ms = triton.testing.do_bench_cudagraph(
112 - runner, quantiles=quantiles
113 - )
114 -
115 - def tflops(t_ms):
116 - flops = 2 * batch_size * hidden_size * num_experts
117 - return flops / (t_ms * 1e-3) / 1e12
118 -
119 - return tflops(ms), tflops(max_ms), tflops(min_ms)
120 -
121 - return benchmark
122 -
123 -
124 -if __name__ == "__main__":
125 - parser = FlexibleArgumentParser()
126 - parser.add_argument("--model", type=str, default="openai/gpt-oss-20b")
127 - parser.add_argument("--max-batch-size", default=16, type=int)
128 - parser.add_argument("--trust-remote-code", action="store_true")
129 - args = parser.parse_args()
130 -
131 - # Get the benchmark function
132 - benchmark = get_benchmark(args.model, args.max_batch_size, args.trust_remote_code)
133 - # Run performance benchmark
134 - benchmark.run(print_data=True)
csrc/moe/gpt_oss_router_gemm.cu deleted
-144
@@ -1,144 +0,0 @@
1 -/*
2 - * Adapted from
3 - * https://github.com/NVIDIA/TensorRT-LLM/blob/v1.3.0rc7/cpp/tensorrt_llm/kernels/tinygemm2/tinygemm2_cuda.cu
4 - * Copyright (c) 2025, The vLLM team.
5 - * SPDX-FileCopyrightText: Copyright (c) 2025, NVIDIA CORPORATION.
6 - * All rights reserved. SPDX-License-Identifier: Apache-2.0
7 - *
8 - * Licensed under the Apache License, Version 2.0 (the "License");
9 - * you may not use this file except in compliance with the License.
10 - * You may obtain a copy of the License at
11 - *
12 - * http://www.apache.org/licenses/LICENSE-2.0
13 - *
14 - * Unless required by applicable law or agreed to in writing, software
15 - * distributed under the License is distributed on an "AS IS" BASIS,
16 - * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
17 - * See the License for the specific language governing permissions and
18 - * limitations under the License.
19 - */
20 -
21 -#include <ATen/cuda/CUDAContext.h>
22 -#include <c10/cuda/CUDAStream.h>
23 -#include <cuda.h>
24 -#include <cuda_runtime.h>
25 -#include <torch/all.h>
26 -#include "gpt_oss_router_gemm.cuh"
27 -
28 -void launch_gpt_oss_router_gemm(__nv_bfloat16* gA, __nv_bfloat16* gB,
29 - __nv_bfloat16* gC, __nv_bfloat16* bias,
30 - int batch_size, int output_features,
31 - int input_features, cudaStream_t stream) {
32 - static int const WARP_TILE_M = 16;
33 - static int const TILE_M = WARP_TILE_M;
34 - static int const TILE_N = 8;
35 - static int const TILE_K = 64;
36 - static int const STAGES = 16;
37 - static int const STAGE_UNROLL = 4;
38 - static bool const PROFILE = false;
39 -
40 - CUtensorMap weight_map{};
41 - CUtensorMap activation_map{};
42 -
43 - constexpr uint32_t rank = 2;
44 - uint64_t size[rank] = {(uint64_t)input_features, (uint64_t)output_features};
45 - uint64_t stride[rank - 1] = {input_features * sizeof(__nv_bfloat16)};
46 - uint32_t box_size[rank] = {TILE_K, TILE_M};
47 - uint32_t elem_stride[rank] = {1, 1};
48 -
49 - CUresult res = cuTensorMapEncodeTiled(
50 - &weight_map, CUtensorMapDataType::CU_TENSOR_MAP_DATA_TYPE_BFLOAT16, rank,
51 - gB, size, stride, box_size, elem_stride,
52 - CUtensorMapInterleave::CU_TENSOR_MAP_INTERLEAVE_NONE,
53 - CUtensorMapSwizzle::CU_TENSOR_MAP_SWIZZLE_128B,
54 - CUtensorMapL2promotion::CU_TENSOR_MAP_L2_PROMOTION_NONE,
55 - CUtensorMapFloatOOBfill::CU_TENSOR_MAP_FLOAT_OOB_FILL_NONE);
56 - TORCH_CHECK(res == CUDA_SUCCESS,
57 - "cuTensorMapEncodeTiled failed for weight_map, error code=",
58 - static_cast<int>(res));
59 -
60 - size[1] = batch_size;
61 - box_size[1] = TILE_N;
62 -
63 - res = cuTensorMapEncodeTiled(
64 - &activation_map, CUtensorMapDataType::CU_TENSOR_MAP_DATA_TYPE_BFLOAT16,
65 - rank, gA, size, stride, box_size, elem_stride,
66 - CUtensorMapInterleave::CU_TENSOR_MAP_INTERLEAVE_NONE,
67 - CUtensorMapSwizzle::CU_TENSOR_MAP_SWIZZLE_128B,
68 - CUtensorMapL2promotion::CU_TENSOR_MAP_L2_PROMOTION_NONE,
69 - CUtensorMapFloatOOBfill::CU_TENSOR_MAP_FLOAT_OOB_FILL_NONE);
70 - TORCH_CHECK(res == CUDA_SUCCESS,
71 - "cuTensorMapEncodeTiled failed for activation_map, error code=",
72 - static_cast<int>(res));
73 -
74 - int smem_size = STAGES * STAGE_UNROLL *
75 - (TILE_M * TILE_K * sizeof(__nv_bfloat16) +
76 - TILE_N * TILE_K * sizeof(__nv_bfloat16));
77 -
78 - gpuErrChk(cudaFuncSetAttribute(
79 - gpt_oss_router_gemm_kernel<WARP_TILE_M, TILE_M, TILE_N, TILE_K, STAGES,
80 - STAGE_UNROLL, PROFILE>,
81 - cudaFuncAttributeMaxDynamicSharedMemorySize, smem_size));
82 -
83 - int tiles_m = (output_features + TILE_M - 1) / TILE_M;
84 - int tiles_n = (batch_size + TILE_N - 1) / TILE_N;
85 -
86 - dim3 grid(tiles_m, tiles_n);
87 - dim3 block(384);
88 -
89 - cudaLaunchConfig_t config;
90 - cudaLaunchAttribute attrs[1];
91 - config.gridDim = grid;
92 - config.blockDim = block;
93 - config.dynamicSmemBytes = smem_size;
94 - config.stream = stream;
95 - config.attrs = attrs;
96 - attrs[0].id = cudaLaunchAttributeProgrammaticStreamSerialization;
97 - attrs[0].val.programmaticStreamSerializationAllowed = 1;
98 - config.numAttrs = 1;
99 -
100 - cudaLaunchKernelEx(
101 - &config,
102 - &gpt_oss_router_gemm_kernel<WARP_TILE_M, TILE_M, TILE_N, TILE_K, STAGES,
103 - STAGE_UNROLL, PROFILE>,
104 - gC, gA, gB, bias, output_features, batch_size, input_features, weight_map,
105 - activation_map, nullptr);
106 -}
107 -
108 -void gpt_oss_router_gemm_cuda_forward(torch::Tensor& output,
109 - torch::Tensor input, torch::Tensor weight,
110 - torch::Tensor bias) {
111 - auto const batch_size = input.size(0);
112 - auto const input_dim = input.size(1);
113 - auto const output_dim = weight.size(0);
114 -
115 - auto stream = at::cuda::getCurrentCUDAStream();
116 -
117 - if (input.scalar_type() == at::ScalarType::BFloat16) {
118 - launch_gpt_oss_router_gemm((__nv_bfloat16*)input.data_ptr(),
119 - (__nv_bfloat16*)weight.data_ptr(),
120 - (__nv_bfloat16*)output.mutable_data_ptr(),
121 - (__nv_bfloat16*)bias.data_ptr(), batch_size,
122 - output_dim, input_dim, stream);
123 - } else {
124 - throw std::invalid_argument("Unsupported dtype, only supports bfloat16");
125 - }
126 -}
127 -
128 -void gpt_oss_router_gemm(torch::Tensor& output, torch::Tensor input,
129 - torch::Tensor weight, torch::Tensor bias) {
130 - TORCH_CHECK(input.dim() == 2, "input must be 2D");
131 - TORCH_CHECK(weight.dim() == 2, "weight must be 2D");
132 - TORCH_CHECK(bias.dim() == 1, "bias must be 1D");
133 - TORCH_CHECK(input.sizes()[1] == weight.sizes()[1],
134 - "input.size(1) must match weight.size(1)");
135 - TORCH_CHECK(weight.sizes()[0] == bias.sizes()[0],
136 - "weight.size(0) must match bias.size(0)");
137 - TORCH_CHECK(input.scalar_type() == at::ScalarType::BFloat16,
138 - "input tensor must be bfloat16");
139 - TORCH_CHECK(weight.scalar_type() == at::ScalarType::BFloat16,
140 - "weight tensor must be bfloat16");
141 - TORCH_CHECK(bias.scalar_type() == at::ScalarType::BFloat16,
142 - "bias tensor must be bfloat16");
143 - gpt_oss_router_gemm_cuda_forward(output, input, weight, bias);
144 -}
csrc/moe/gpt_oss_router_gemm.cuh deleted
-447
@@ -1,447 +0,0 @@
1 -/*
2 - * Adapted from
3 - * https://github.com/NVIDIA/TensorRT-LLM/blob/v1.3.0rc7/cpp/tensorrt_llm/kernels/tinygemm2/tinygemm2_kernel.cuh
4 - * Copyright (c) 2025, The vLLM team.
5 - * SPDX-FileCopyrightText: Copyright (c) 2025, NVIDIA CORPORATION.
6 - * All rights reserved. SPDX-License-Identifier: Apache-2.0
7 - *
8 - * Licensed under the Apache License, Version 2.0 (the "License");
9 - * you may not use this file except in compliance with the License.
10 - * You may obtain a copy of the License at
11 - *
12 - * http://www.apache.org/licenses/LICENSE-2.0
13 - *
14 - * Unless required by applicable law or agreed to in writing, software
15 - * distributed under the License is distributed on an "AS IS" BASIS,
16 - * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
17 - * See the License for the specific language governing permissions and
18 - * limitations under the License.
19 - */
20 -
21 -#include "cuda_bf16.h"
22 -#include <stdint.h>
23 -#include <stdio.h>
24 -#include <vector>
25 -
26 -#include "cuda_pipeline.h"
27 -#include <cuda.h>
28 -#include <cuda/barrier>
29 -#include <cuda/std/utility>
30 -#include <cuda_runtime.h>
31 -
32 -using barrier = cuda::barrier<cuda::thread_scope_block>;
33 -namespace cde = cuda::device::experimental;
34 -namespace ptx = cuda::ptx;
35 -
36 -#define gpuErrChk(ans) \
37 - { \
38 - gpuAssert((ans), __FILE__, __LINE__); \
39 - }
40 -
41 -inline void gpuAssert(cudaError_t code, char const* file, int line,
42 - bool abort = true) {
43 - if (code != cudaSuccess) {
44 - fprintf(stderr, "GPUassert: %s %s %d\n", cudaGetErrorString(code), file,
45 - line);
46 - if (abort) {
47 - throw std::runtime_error(cudaGetErrorString(code));
48 - }
49 - }
50 -}
51 -
52 -#if defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 900)
53 -__device__ uint64_t gclock64() {
54 - unsigned long long int rv;
55 - asm volatile("mov.u64 %0, %%globaltimer;" : "=l"(rv));
56 - return rv;
57 -}
58 -
59 -__device__ void ldmatrix(__nv_bfloat16 rv[2], uint32_t smem_ptr) {
60 - int dst;
61 - asm volatile("ldmatrix.sync.aligned.x1.m8n8.shared.b16 {%0}, [%1];\n"
62 - : "=r"(dst)
63 - : "r"(smem_ptr));
64 - int* rvi = reinterpret_cast<int*>(&rv[0]);
65 - rvi[0] = dst;
66 -}
67 -
68 -__device__ void ldmatrix2(__nv_bfloat16 rv[4], uint32_t smem_ptr) {
69 - int x, y;
70 - asm volatile("ldmatrix.sync.aligned.x2.m8n8.shared.b16 {%0, %1}, [%2];\n"
71 - : "=r"(x), "=r"(y)
72 - : "r"(smem_ptr));
73 -
74 - int* rvi = reinterpret_cast<int*>(&rv[0]);
75 - rvi[0] = x;
76 - rvi[1] = y;
77 -}
78 -
79 -__device__ void ldmatrix4(__nv_bfloat16 rv[8], uint32_t smem_ptr) {
80 - int x, y, z, w;
81 - asm volatile(
82 - "ldmatrix.sync.aligned.x4.m8n8.shared.b16 {%0, %1, %2, %3}, [%4];"
83 - : "=r"(x), "=r"(y), "=r"(z), "=r"(w)
84 - : "r"(smem_ptr));
85 - int* rvi = reinterpret_cast<int*>(&rv[0]);
86 - rvi[0] = x;
87 - rvi[1] = y;
88 - rvi[2] = z;
89 - rvi[3] = w;
90 -}
91 -
92 -__device__ void HMMA_1688(float d[4], __nv_bfloat16 a[4], __nv_bfloat16 b[2],
93 - float c[4]) {
94 - uint32_t const* A = reinterpret_cast<uint32_t const*>(&a[0]);
95 - uint32_t const* B = reinterpret_cast<uint32_t const*>(&b[0]);
96 - float const* C = reinterpret_cast<float const*>(&c[0]);
97 - float* D = reinterpret_cast<float*>(&d[0]);
98 -
99 - asm volatile(
100 - "mma.sync.aligned.m16n8k8.row.col.f32.bf16.bf16.f32 "
101 - "{%0,%1,%2,%3}, {%4,%5}, {%6}, {%7,%8,%9,%10};\n"
102 - : "=f"(D[0]), "=f"(D[1]), "=f"(D[2]), "=f"(D[3])
103 - : "r"(A[0]), "r"(A[1]), "r"(B[0]), "f"(C[0]), "f"(C[1]), "f"(C[2]),
104 - "f"(C[3]));
105 -}
106 -
107 -__device__ void HMMA_16816(float d[4], __nv_bfloat16 a[8], __nv_bfloat16 b[4],
108 - float c[4]) {
109 - uint32_t const* A = reinterpret_cast<uint32_t const*>(&a[0]);
110 - uint32_t const* B = reinterpret_cast<uint32_t const*>(&b[0]);
111 - float const* C = reinterpret_cast<float const*>(&c[0]);
112 - float* D = reinterpret_cast<float*>(&d[0]);
113 -
114 - asm volatile(
115 - "mma.sync.aligned.m16n8k16.row.col.f32.bf16.bf16.f32 "
116 - "{%0,%1,%2,%3}, {%4,%5,%6,%7}, {%8,%9}, {%10,%11,%12,%13};\n"
117 - : "=f"(D[0]), "=f"(D[1]), "=f"(D[2]), "=f"(D[3])
118 - : "r"(A[0]), "r"(A[1]), "r"(A[2]), "r"(A[3]), "r"(B[0]), "r"(B[1]),
119 - "f"(C[0]), "f"(C[1]), "f"(C[2]), "f"(C[3]));
120 -}
121 -
122 -__device__ void bar_wait(uint32_t bar_ptr, int phase) {
123 - asm volatile(
124 - "{\n"
125 - ".reg .pred P1;\n"
126 - "LAB_WAIT:\n"
127 - "mbarrier.try_wait.parity.shared::cta.b64 P1, [%0], %1;\n"
128 - "@P1 bra.uni DONE;\n"
129 - "bra.uni LAB_WAIT;\n"
130 - "DONE:\n"
131 - "}\n" ::"r"(bar_ptr),
132 - "r"(phase));
133 -}
134 -
135 -__device__ bool bar_try_wait(uint32_t bar_ptr, int phase) {
136 - uint32_t success;
137 - #ifdef INTERNAL
138 - asm volatile(".pragma \"set knob DontInsertYield\";\n" : : : "memory");
139 - #endif
140 - asm volatile(
141 - "{\n\t"
142 - ".reg .pred P1; \n\t"
143 - "mbarrier.try_wait.parity.shared::cta.b64 P1, [%1], %2; \n\t"
144 - "selp.b32 %0, 1, 0, P1; \n\t"
145 - "}"
146 - : "=r"(success)
147 - : "r"(bar_ptr), "r"(phase));
148 - return success;
149 -}
150 -
151 -__device__ uint32_t elect_one_sync() {
152 - uint32_t pred = 0;
153 - uint32_t laneid = 0;
154 - asm volatile(
155 - "{\n"
156 - ".reg .b32 %%rx;\n"
157 - ".reg .pred %%px;\n"
158 - " elect.sync %%rx|%%px, %2;\n"
159 - "@%%px mov.s32 %1, 1;\n"
160 - " mov.s32 %0, %%rx;\n"
161 - "}\n"
162 - : "+r"(laneid), "+r"(pred)
163 - : "r"(0xFFFFFFFF));
164 - return pred;
165 -}
166 -#endif
167 -
168 -struct Profile {
169 - uint64_t start;
170 - uint64_t weight_load_start;
171 - uint64_t act_load_start;
172 - uint64_t compute_start;
173 - uint64_t complete;
174 -};
175 -
176 -template <int WARP_TILE_M, int TILE_M, int TILE_N, int TILE_K, int STAGES,
177 - int STAGE_UNROLL, bool PROFILE>
178 -__global__ __launch_bounds__(384, 1) void gpt_oss_router_gemm_kernel(
179 - __nv_bfloat16* output, __nv_bfloat16* weights, __nv_bfloat16* activations,
180 - __nv_bfloat16* bias, int M, int N, int K,
181 - const __grid_constant__ CUtensorMap weight_map,
182 - const __grid_constant__ CUtensorMap activation_map,
183 - Profile* profile = nullptr) {
184 -#if defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 900)
185 -
186 - if (PROFILE && threadIdx.x == 0 && blockIdx.y == 0)
187 - profile[blockIdx.x].start = gclock64();
188 -
189 - extern __shared__ __align__(128) char smem[];
190 -
191 - __nv_bfloat16* sh_weights = (__nv_bfloat16*)&smem[0];
192 - __nv_bfloat16* sh_activations =
193 - (__nv_bfloat16*)&smem[STAGES * STAGE_UNROLL * TILE_M * TILE_K *
194 - sizeof(__nv_bfloat16)];
195 -
196 - #pragma nv_diag_suppress static_var_with_dynamic_init
197 - __shared__ barrier bar_wt_ready[STAGES];
198 - __shared__ barrier bar_act_ready[STAGES];
199 - __shared__ barrier bar_data_consumed[STAGES];
200 -
201 - __shared__ float4 reduction_buffer[128];
202 -
203 - __shared__ nv_bfloat16 sh_bias[TILE_M];
204 -
205 - if (threadIdx.x == 0) {
206 - for (int i = 0; i < STAGES; i++) {
207 - init(&bar_wt_ready[i], 1);
208 - init(&bar_act_ready[i], 1);
209 - init(&bar_data_consumed[i], 32);
210 - }
211 - ptx::fence_proxy_async(ptx::space_shared);
212 - asm volatile("prefetch.tensormap [%0];"
213 - :
214 - : "l"(reinterpret_cast<uint64_t>(&weight_map))
215 - : "memory");
216 - asm volatile("prefetch.tensormap [%0];"
217 - :
218 - : "l"(reinterpret_cast<uint64_t>(&activation_map))
219 - : "memory");
220 - }
221 - __syncthreads();
222 -
223 - int warp_id = threadIdx.x / 32;
224 - int lane_id = threadIdx.x % 32;
225 -
226 - int phase = 0;
227 -
228 - int mib = blockIdx.x * TILE_M;
229 - int ni = blockIdx.y * TILE_N;
230 -
231 - float accum[4];
232 - for (int i = 0; i < 4; i++) accum[i] = 0.f;
233 -
234 - int const K_LOOPS_DMA =
235 - (K + 4 * TILE_K * STAGE_UNROLL - 1) / (4 * (TILE_K * STAGE_UNROLL));
236 - int const K_LOOPS_COMPUTE = K_LOOPS_DMA;
237 -
238 - // Data loading thread
239 - if (warp_id >= 4 && elect_one_sync()) {
240 - int stage = warp_id % 4;
241 -
242 - bool weight_warp = warp_id < 8;
243 - if (!weight_warp) {
244 - cudaGridDependencySynchronize();
245 - cudaTriggerProgrammaticLaunchCompletion();
246 - }
247 -
248 - for (int ki = 0; ki < K_LOOPS_DMA; ki++) {
249 - int k = (ki * 4 + (warp_id % 4)) * TILE_K * STAGE_UNROLL;
250 -
251 - uint64_t desc_ptr_wt = reinterpret_cast<uint64_t>(&weight_map);
252 - uint64_t desc_ptr_act = reinterpret_cast<uint64_t>(&activation_map);
253 -
254 - uint32_t bar_ptr_wt = __cvta_generic_to_shared(&bar_wt_ready[stage]);
255 - uint32_t bar_ptr_act = __cvta_generic_to_shared(&bar_act_ready[stage]);
256 - int bytes_wt = TILE_M * TILE_K * sizeof(__nv_bfloat16);
257 - int bytes_act = TILE_N * TILE_K * sizeof(__nv_bfloat16);
258 -
259 - bar_wait(__cvta_generic_to_shared(&bar_data_consumed[stage]), phase ^ 1);
260 -
261 - if (weight_warp)
262 - asm volatile("mbarrier.arrive.expect_tx.shared.b64 _, [%0], %1;"
263 - :
264 - : "r"(bar_ptr_wt), "r"(STAGE_UNROLL * bytes_wt));
265 - if (!weight_warp)
266 - asm volatile("mbarrier.arrive.expect_tx.shared.b64 _, [%0], %1;"
267 - :
268 - : "r"(bar_ptr_act), "r"(STAGE_UNROLL * bytes_act));
269 -
270 - if (PROFILE && blockIdx.y == 0 && ki == 0 && weight_warp)
271 - profile[blockIdx.x].weight_load_start = gclock64();
272 - if (PROFILE && blockIdx.y == 0 && ki == 0 && !weight_warp)
273 - profile[blockIdx.x].act_load_start = gclock64();
274 -
275 - for (int i = 0; i < STAGE_UNROLL; i++) {
276 - uint32_t smem_ptr_wt = __cvta_generic_to_shared(
277 - &sh_weights[(stage * STAGE_UNROLL + i) * TILE_M * TILE_K]);
278 - uint32_t crd0 = k + i * TILE_K;
279 - uint32_t crd1 = mib;
280 - if (weight_warp)
281 - asm volatile(
282 - "cp.async.bulk.tensor.2d.shared::cta.global.mbarrier::complete_"
283 - "tx::bytes [%0], [%1, {%3,%4}], "
284 - "[%2];"
285 - :
286 - : "r"(smem_ptr_wt), "l"(desc_ptr_wt), "r"(bar_ptr_wt), "r"(crd0),
287 - "r"(crd1)
288 - : "memory");
289 -
290 - uint32_t smem_ptr_act = __cvta_generic_to_shared(
291 - &sh_activations[(stage * STAGE_UNROLL + i) * TILE_N * TILE_K]);
292 - crd0 = k + i * TILE_K;
293 - crd1 = ni;
294 - if (!weight_warp)
295 - asm volatile(
296 - "cp.async.bulk.tensor.2d.shared::cta.global.mbarrier::complete_"
297 - "tx::bytes [%0], [%1, {%3,%4}], "
298 - "[%2];"
299 - :
300 - : "r"(smem_ptr_act), "l"(desc_ptr_act), "r"(bar_ptr_act),
301 - "r"(crd0), "r"(crd1)
302 - : "memory");
303 - }
304 -
305 - stage += 4;
306 - if (stage >= STAGES) {
307 - stage = warp_id % 4;
308 - phase ^= 1;
309 - }
310 - }
311 - // Wait for pending loads to be consumed before exiting, to avoid race
312 - for (int i = 0; i < (STAGES / 4) - 1; i++) {
313 - bar_wait(__cvta_generic_to_shared(&bar_data_consumed[stage]), phase ^ 1);
314 - stage += 4;
315 - if (stage >= STAGES) {
316 - stage = warp_id % 4;
317 - phase ^= 1;
318 - }
319 - }
320 - }
321 - // Compute threads
322 - else if (warp_id < 4) {
323 - // Sneak the bias load into the compute warps since they're just waiting for
324 - // stuff anyway
325 - if (threadIdx.x < TILE_M) sh_bias[threadIdx.x] = bias[mib + threadIdx.x];
326 -
327 - int stage = warp_id;
328 -
329 - int phase = 0;
330 - int lane_id_div8 = lane_id / 8;
331 - int lane_id_mod8 = lane_id % 8;
332 -
333 - int lane_row_offset_wt = (lane_id_div8 % 2) ? 8 : 0;
334 - int lane_col_offset_wt = (lane_id_div8 / 2) ? 1 : 0;
335 -
336 - int row_wt = lane_id_mod8 + lane_row_offset_wt;
337 - int row_act = lane_id_mod8;
338 -
339 - int row_offset_wt = (reinterpret_cast<uintptr_t>(sh_weights) / 128) % 8;
340 - int row_offset_act = row_offset_wt;
341 -
342 - uint32_t bar_ptr_wt = __cvta_generic_to_shared(&bar_wt_ready[stage]);
343 - uint32_t bar_ptr_act = __cvta_generic_to_shared(&bar_act_ready[stage]);
344 -
345 - bool weight_ready = bar_try_wait(bar_ptr_wt, phase);
346 - bool act_ready = bar_try_wait(bar_ptr_act, phase);
347 -
348 - #pragma unroll 2
349 - for (int ki = 0; ki < K_LOOPS_COMPUTE; ki++) {
350 - int next_stage = stage + 4;
351 - int next_phase = phase;
352 - if (next_stage >= STAGES) {
353 - next_stage = warp_id;
354 - next_phase ^= 1;
355 - }
356 -
357 - while (!weight_ready || !act_ready) {
358 - weight_ready = bar_try_wait(bar_ptr_wt, phase);
359 - act_ready = bar_try_wait(bar_ptr_act, phase);
360 - }
361 -
362 - if (PROFILE && blockIdx.y == 0 && threadIdx.x == 0 && ki == 0)
363 - profile[blockIdx.x].compute_start = gclock64();
364 -
365 - if (ki + 1 < K_LOOPS_COMPUTE) {
366 - weight_ready = bar_try_wait(
367 - __cvta_generic_to_shared(&bar_wt_ready[next_stage]), next_phase);
368 - act_ready = bar_try_wait(
369 - __cvta_generic_to_shared(&bar_act_ready[next_stage]), next_phase);
370 - }
371 -
372 - #pragma unroll
373 - for (int su = 0; su < STAGE_UNROLL; su++) {
374 - __nv_bfloat16* ptr_weights =
375 - &sh_weights[(stage * STAGE_UNROLL + su) * TILE_M * TILE_K];
376 - __nv_bfloat16* ptr_act =
377 - &sh_activations[(stage * STAGE_UNROLL + su) * TILE_N * TILE_K];
378 -
379 - #pragma unroll
380 - for (int kii = 0; kii < TILE_K / 16; kii++) {
381 - __nv_bfloat16 a[8];
382 - __nv_bfloat16 b[4];
383 -
384 - int col = 2 * kii + lane_col_offset_wt;
385 - int col_sw = ((row_wt + row_offset_wt) % 8) ^ col;
386 -
387 - ldmatrix4(a, __cvta_generic_to_shared(
388 - &ptr_weights[row_wt * TILE_K + col_sw * 8]));
389 -
390 - col = 2 * kii + lane_id_div8;
391 - col_sw = ((row_act + row_offset_act) % 8) ^ col;
392 -
393 - ldmatrix2(b, __cvta_generic_to_shared(
394 - &ptr_act[row_act * TILE_K + 8 * col_sw]));
395 -
396 - HMMA_16816(accum, a, b, accum);
397 - }
398 - }
399 -
400 - uint32_t bar_c = __cvta_generic_to_shared(&bar_data_consumed[stage]);
401 - asm volatile("mbarrier.arrive.shared::cta.b64 _, [%0];" : : "r"(bar_c));
402 -
403 - stage = next_stage;
404 - phase = next_phase;
405 - }
406 -
407 - float4 accum4;
408 - accum4.x = accum[0];
409 - accum4.y = accum[1];
410 - accum4.z = accum[2];
411 - accum4.w = accum[3];
412 - reduction_buffer[threadIdx.x] = accum4;
413 -
414 - __syncthreads();
415 -
416 - if (warp_id == 0) {
417 - int mi = mib + warp_id * WARP_TILE_M;
418 - int tm = mi + lane_id / 4;
419 - int tn = ni + 2 * (lane_id % 4);
420 -
421 - float4 accum1 = reduction_buffer[32 + threadIdx.x];
422 - float4 accum2 = reduction_buffer[64 + threadIdx.x];
423 - float4 accum3 = reduction_buffer[96 + threadIdx.x];
424 -
425 - accum[0] = accum[0] + accum1.x + accum2.x + accum3.x;
426 - accum[1] = accum[1] + accum1.y + accum2.y + accum3.y;
427 - accum[2] = accum[2] + accum1.z + accum2.z + accum3.z;
428 - accum[3] = accum[3] + accum1.w + accum2.w + accum3.w;
429 -
430 - float bias_lo = __bfloat162float(sh_bias[tm - mib]);
431 - float bias_hi = __bfloat162float(sh_bias[tm + 8 - mib]);
432 -
433 - if (tn < N && tm < M)
434 - output[tn * M + tm] = __float2bfloat16(accum[0] + bias_lo);
435 - if (tn + 1 < N && tm < M)
436 - output[(tn + 1) * M + tm] = __float2bfloat16(accum[1] + bias_lo);
437 - if (tn < N && tm + 8 < M)
438 - output[tn * M + tm + 8] = __float2bfloat16(accum[2] + bias_hi);
439 - if (tn + 1 < N && tm + 8 < M)
440 - output[(tn + 1) * M + tm + 8] = __float2bfloat16(accum[3] + bias_hi);
441 -
442 - if (PROFILE && blockIdx.y == 0 && threadIdx.x == 0)
443 - profile[blockIdx.x].complete = gclock64();
444 - }
445 - }
446 -#endif // end if defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 900)
447 -}
csrc/moe/moe_ops.h
-4
@@ -70,8 +70,4 @@ torch::Tensor router_gemm_bf16_fp32(torch::Tensor const& input,
70 // Supports num_tokens in [1, 16], num_experts in {256, 384}, hidden_dim = 7168
71 void dsv3_router_gemm(torch::Tensor& output, const torch::Tensor& mat_a,
72 const torch::Tensor& mat_b);
73 -
74 -// gpt-oss optimized router GEMM kernel for SM90+
75 -void gpt_oss_router_gemm(torch::Tensor& output, torch::Tensor input,
76 - torch::Tensor weight, torch::Tensor bias);
73 #endif
csrc/moe/torch_bindings.cpp
-6
@@ -132,12 +132,6 @@ TORCH_LIBRARY_EXPAND(TORCH_EXTENSION_NAME, m) {
132 // DeepSeek V3 optimized router GEMM for SM90+
133 m.def("dsv3_router_gemm(Tensor! output, Tensor mat_a, Tensor mat_b) -> ()");
134 // conditionally compiled so impl registration is in source file
135 -
136 - // gpt-oss optimized router GEMM kernel for SM90+
137 - m.def(
138 - "gpt_oss_router_gemm(Tensor! output, Tensor input, Tensor weights, "
139 - "Tensor bias) -> ()");
140 - m.impl("gpt_oss_router_gemm", torch::kCUDA, &gpt_oss_router_gemm);
135 #endif
136 }
137
tests/kernels/moe/test_router_gemm.py deleted
-37
@@ -1,37 +0,0 @@
1 -# SPDX-License-Identifier: Apache-2.0
2 -# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 -"""Tests for optimized router GEMM kernel
4 -
5 -Run `pytest tests/kernels/moe/test_router_gemm.py`.
6 -"""
7 -
8 -import pytest
9 -import torch
10 -
11 -import vllm._custom_ops as ops
12 -from vllm.platforms import current_platform
13 -from vllm.utils.torch_utils import set_random_seed
14 -
15 -
16 -@pytest.mark.skipif(
17 - not (
18 - current_platform.is_cuda()
19 - and (
20 - current_platform.is_device_capability(90)
21 - or current_platform.is_device_capability_family(100)
22 - )
23 - ),
24 - reason="This test only runs on Hopper or Blackwell GPUs.",
25 -)
26 -@pytest.mark.parametrize("batch_size", [1, 2, 4, 8])
27 -@pytest.mark.parametrize("input_dim", [360, 720, 1440, 2880])
28 -@pytest.mark.parametrize("output_dim", [32, 64, 128])
29 -def test_gpt_oss_router_gemm(batch_size, input_dim, output_dim):
30 - set_random_seed(0)
31 - x = torch.randn(batch_size, input_dim, device="cuda", dtype=torch.bfloat16)
32 - weight = torch.randn(output_dim, input_dim, device="cuda", dtype=torch.bfloat16)
33 - bias = torch.randn(output_dim, device="cuda", dtype=torch.bfloat16)
34 -
35 - output = ops.gpt_oss_router_gemm(x, weight, bias)
36 - output_ref = torch.nn.functional.linear(x, weight, bias)
37 - torch.testing.assert_close(output, output_ref, atol=1e-2, rtol=1e-2)
vllm/_custom_ops.py
-13
@@ -2327,19 +2327,6 @@ def dsv3_router_gemm(
2327 return output
2328
2329
2330 -def gpt_oss_router_gemm(
2331 - hidden_states: torch.Tensor, weight: torch.Tensor, bias: torch.Tensor
2332 -) -> torch.Tensor:
2333 - output = torch.empty(
2334 - hidden_states.shape[0],
2335 - weight.shape[0],
2336 - device=hidden_states.device,
2337 - dtype=hidden_states.dtype,
2338 - )
2339 - torch.ops._moe_C.gpt_oss_router_gemm(output, hidden_states, weight, bias)
2340 - return output
2341 -
2342 -
2330 def topk_softmax(
2331 topk_weights: torch.Tensor,
2332 topk_ids: torch.Tensor,
vllm/lora/layers/__init__.py
-2
@@ -13,7 +13,6 @@ from vllm.lora.layers.column_parallel_linear import (
13 QKVParallelLinearWithShardedLoRA,
14 )
15 from vllm.lora.layers.fused_moe import FusedMoE3DWithLoRA, FusedMoEWithLoRA
16 -from vllm.lora.layers.gate_linear import GateLinearWithLoRA
16 from vllm.lora.layers.logits_processor import LogitsProcessorWithLoRA
17 from vllm.lora.layers.replicated_linear import ReplicatedLinearWithLoRA
18 from vllm.lora.layers.row_parallel_linear import (
@@ -39,7 +38,6 @@ __all__ = [
38 "RowParallelLinearWithLoRA",
39 "RowParallelLinearWithShardedLoRA",
40 "ReplicatedLinearWithLoRA",
42 - "GateLinearWithLoRA",
41 "LoRAMapping",
42 "LoRAMappingType",
43 "FusedMoEWithLoRA",
vllm/lora/layers/gate_linear.py deleted
-30
@@ -1,30 +0,0 @@
1 -# SPDX-License-Identifier: Apache-2.0
2 -# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 -
4 -import torch.nn as nn
5 -from transformers import PretrainedConfig
6 -
7 -from vllm.config.lora import LoRAConfig
8 -from vllm.model_executor.custom_op import maybe_get_oot_by_class
9 -from vllm.model_executor.layers.fused_moe.router.gate_linear import GateLinear
10 -
11 -from .replicated_linear import ReplicatedLinearWithLoRA
12 -
13 -
14 -class GateLinearWithLoRA(ReplicatedLinearWithLoRA):
15 - def __init__(self, base_layer: GateLinear) -> None:
16 - super().__init__(
17 - base_layer,
18 - )
19 -
20 - # GateLinearWithLoRA should always be replaced, regardless of the fully
21 - # sharded LoRAs setting, because it is, by definition, copied per GPU.
22 - @classmethod
23 - def can_replace_layer(
24 - cls,
25 - source_layer: nn.Module,
26 - lora_config: LoRAConfig,
27 - packed_modules_list: list,
28 - model_config: PretrainedConfig | None = None,
29 - ) -> bool:
30 - return type(source_layer) is maybe_get_oot_by_class(GateLinear)
vllm/lora/utils.py
-2
@@ -21,7 +21,6 @@ from vllm.lora.layers import (
21 ColumnParallelLinearWithShardedLoRA,
22 FusedMoE3DWithLoRA,
23 FusedMoEWithLoRA,
24 - GateLinearWithLoRA,
24 LogitsProcessorWithLoRA,
25 MergedColumnParallelLinearVariableSliceWithLoRA,
26 MergedColumnParallelLinearWithLoRA,
@@ -82,7 +81,6 @@ _all_lora_classes: set[type[BaseLayerWithLoRA]] = {
81 MergedQKVParallelLinearWithLoRA,
82 RowParallelLinearWithLoRA,
83 ReplicatedLinearWithLoRA,
85 - GateLinearWithLoRA,
84 LogitsProcessorWithLoRA,
85 ColumnParallelLinearWithShardedLoRA,
86 QKVParallelLinearWithShardedLoRA,
vllm/model_executor/layers/fused_moe/router/gate_linear.py
+6 -52
@@ -3,11 +3,9 @@
3 import torch
4 from torch.nn.parameter import Parameter
5
6 -import vllm._custom_ops as ops
6 from vllm.model_executor.custom_op import PluggableLayer
7 from vllm.model_executor.layers.linear import ReplicatedLinear
8 from vllm.platforms import current_platform
10 -from vllm.utils.torch_utils import direct_register_custom_op
9
10
11 @PluggableLayer.register("gate_linear")
@@ -15,9 +13,8 @@ class GateLinear(ReplicatedLinear):
13 """MoE gate linear layer with three-tier GEMM dispatch:
14
15 1. DSV3 specialized kernel (SM90+, batch<=16, supported dims)
18 - 2. gpt-oss specialized kernel (SM90+, batch<=128, supported dims)
19 - 3. cuBLAS bf16×bf16→fp32 (SM90+ + bf16 + fp32 out_dtype)
20 - 4. F.linear via ReplicatedLinear (ultimate fallback)
16 + 2. cuBLAS bf16×bf16→fp32 (SM90+ + bf16 + fp32 out_dtype)
17 + 3. F.linear via ReplicatedLinear (ultimate fallback)
18
19 The ``out_dtype`` attribute is mutable and can be set after init
20 (e.g. when the required dtype depends on the expert quantization
@@ -28,10 +25,6 @@ class GateLinear(ReplicatedLinear):
25 DSV3_SUPPORTED_NUM_EXPERTS = [256, 384]
26 DSV3_SUPPORTED_HIDDEN_SIZES = [7168]
27
31 - # Dimensions supported by the gpt-oss specialized kernel
32 - GPT_OSS_SUPPORTED_NUM_EXPERTS = [32, 128]
33 - GPT_OSS_SUPPORTED_HIDDEN_SIZES = [2880]
34 -
28 def __init__(
29 self,
30 input_size: int,
@@ -72,15 +65,6 @@ class GateLinear(ReplicatedLinear):
65 and input_size in self.DSV3_SUPPORTED_HIDDEN_SIZES
66 )
67
75 - # gpt-oss specialized kernel eligibility (SM90+, exact dims)
76 - self.allow_gpt_oss_router_gemm = (
77 - self.weight.dtype == torch.bfloat16
78 - and current_platform.is_cuda()
79 - and is_hopper_or_blackwell
80 - and output_size in self.GPT_OSS_SUPPORTED_NUM_EXPERTS
81 - and input_size in self.GPT_OSS_SUPPORTED_HIDDEN_SIZES
82 - )
83 -
68 # cuBLAS bf16→fp32 eligibility
69 self.allow_cublas_router_gemm = (
70 self.allow_specialized_router_gemm
@@ -108,6 +92,8 @@ class GateLinear(ReplicatedLinear):
92 def forward(
93 self, x: torch.Tensor
94 ) -> torch.Tensor | tuple[torch.Tensor, Parameter | None]:
95 + import vllm._custom_ops as ops
96 +
97 # Tier 1: DSV3 specialized kernel
98 if self.allow_dsv3_router_gemm and x.shape[0] <= 16:
99 output = ops.dsv3_router_gemm(
@@ -117,47 +103,15 @@ class GateLinear(ReplicatedLinear):
103 )
104 return output, None
105
120 - # Tier 2: gpt-oss specialized kernel
121 - if self.allow_gpt_oss_router_gemm:
122 - output = torch.ops.vllm.gpt_oss_router_gemm(x, self.weight, self.bias)
123 - return output, None
124 -
125 - # Tier 3: cuBLAS bf16→fp32
106 + # Tier 2: cuBLAS bf16→fp32
107 if self.allow_cublas_router_gemm and x.dtype == torch.bfloat16:
108 output = ops.router_gemm_bf16_fp32(x, self.weight)
109 return output, None
110
130 - # Tier 4: F.linear (ReplicatedLinear)
111 + # Tier 3: F.linear (ReplicatedLinear)
112 if self.out_dtype is not None and x.dtype != self.weight.dtype:
113 x = x.to(self.weight.dtype)
114 output, output_bias = super().forward(x)
115 if self.out_dtype is not None and output.dtype != self.out_dtype:
116 output = output.to(self.out_dtype)
117 return output, output_bias
137 -
138 -
139 -def gpt_oss_router_gemm_impl(
140 - x: torch.Tensor, weight: torch.Tensor, bias: torch.Tensor
141 -) -> torch.Tensor:
142 - """
143 - Dynamically run min-latency gemm if num_tokens <= 128.
144 - This must be wrapped in a custom op because our torch.compile integration
145 - does not support runtime dispatching on num_tokens.
146 - """
147 - if x.shape[0] <= 128:
148 - return ops.gpt_oss_router_gemm(x, weight, bias)
149 - else:
150 - return torch.nn.functional.linear(x, weight, bias)
151 -
152 -
153 -def gpt_oss_router_gemm_fake(
154 - x: torch.Tensor, weight: torch.Tensor, bias: torch.Tensor
155 -) -> torch.Tensor:
156 - return x.new_empty((x.shape[0], weight.shape[0]))
157 -
158 -
159 -direct_register_custom_op(
160 - op_name="gpt_oss_router_gemm",
161 - op_func=gpt_oss_router_gemm_impl,
162 - fake_impl=gpt_oss_router_gemm_fake,
163 -)
vllm/model_executor/models/gpt_oss.py
+6 -3
@@ -20,11 +20,12 @@ from vllm.distributed import (
20 tensor_model_parallel_all_gather,
21 )
22 from vllm.model_executor.layers.attention import Attention
23 -from vllm.model_executor.layers.fused_moe import FusedMoE, GateLinear
23 +from vllm.model_executor.layers.fused_moe import FusedMoE
24 from vllm.model_executor.layers.fused_moe.config import FusedMoEParallelConfig
25 from vllm.model_executor.layers.layernorm import RMSNorm
26 from vllm.model_executor.layers.linear import (
27 QKVParallelLinear,
28 + ReplicatedLinear,
29 RowParallelLinear,
30 )
31 from vllm.model_executor.layers.logits_processor import LogitsProcessor
@@ -174,11 +175,13 @@ class MLPBlock(torch.nn.Module):
175 self.hidden_size = config.hidden_size
176 self.experts_per_token = config.num_experts_per_tok
177 self.world_size = dist.get_world_size() if dist.is_initialized() else 1
177 - self.router = GateLinear(
178 + self.router = ReplicatedLinear(
179 config.hidden_size,
180 config.num_local_experts,
181 bias=True,
182 + quant_config=None,
183 prefix=f"{prefix}.router",
184 + return_bias=False,
185 )
186 assert config.intermediate_size % self.world_size == 0
187 self.experts = FusedMoE(
@@ -206,7 +209,7 @@ class MLPBlock(torch.nn.Module):
209 self, x[:, : self.hidden_size], self.router.weight, self.router.bias
210 )
211 else:
209 - g, _ = self.router(x)
212 + g = self.router(x)
213 x = self.experts(hidden_states=x, router_logits=g)[:, : self.hidden_size]
214
215 if self.is_sequence_parallel: