@tduncan / Nvidia-Optimization-Stack-VLLM Mirror · read-only / commits / b1169d7be8
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 6, 12:00 UTC

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

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

Xin Yang committed Mar 18, 2026 at 08:15 UTC b1169d7be8add20ab1db4bc93c2b5c6336ef9754
13 files changed +875 -13
CMakeLists.txt
+1
@@ -999,6 +999,7 @@ if(VLLM_GPU_LANG STREQUAL "CUDA")
999 list(APPEND VLLM_MOE_EXT_SRC
1000 "csrc/moe/moe_wna16.cu"
1001 "csrc/moe/grouped_topk_kernels.cu"
1002 + "csrc/moe/gpt_oss_router_gemm.cu"
1003 "csrc/moe/router_gemm.cu")
1004 endif()
1005
benchmarks/kernels/benchmark_router_gemm.py new
+134
@@ -0,0 +1,134 @@
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 new
+144
@@ -0,0 +1,144 @@
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 new
+447
@@ -0,0 +1,447 @@
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,4 +70,8 @@ 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);
77 #endif
csrc/moe/torch_bindings.cpp
+6
@@ -132,6 +132,12 @@ 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);
141 #endif
142 }
143
tests/kernels/moe/test_router_gemm.py new
+37
@@ -0,0 +1,37 @@
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
@@ -2362,6 +2362,19 @@ def dsv3_router_gemm(
2362 return output
2363
2364
2365 +def gpt_oss_router_gemm(
2366 + hidden_states: torch.Tensor, weight: torch.Tensor, bias: torch.Tensor
2367 +) -> torch.Tensor:
2368 + output = torch.empty(
2369 + hidden_states.shape[0],
2370 + weight.shape[0],
2371 + device=hidden_states.device,
2372 + dtype=hidden_states.dtype,
2373 + )
2374 + torch.ops._moe_C.gpt_oss_router_gemm(output, hidden_states, weight, bias)
2375 + return output
2376 +
2377 +
2378 def topk_softmax(
2379 topk_weights: torch.Tensor,
2380 topk_ids: torch.Tensor,
vllm/lora/layers/__init__.py
+2
@@ -13,6 +13,7 @@ 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
17 from vllm.lora.layers.logits_processor import LogitsProcessorWithLoRA
18 from vllm.lora.layers.replicated_linear import ReplicatedLinearWithLoRA
19 from vllm.lora.layers.row_parallel_linear import (
@@ -38,6 +39,7 @@ __all__ = [
39 "RowParallelLinearWithLoRA",
40 "RowParallelLinearWithShardedLoRA",
41 "ReplicatedLinearWithLoRA",
42 + "GateLinearWithLoRA",
43 "LoRAMapping",
44 "LoRAMappingType",
45 "FusedMoEWithLoRA",
vllm/lora/layers/gate_linear.py new
+30
@@ -0,0 +1,30 @@
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,6 +21,7 @@ from vllm.lora.layers import (
21 ColumnParallelLinearWithShardedLoRA,
22 FusedMoE3DWithLoRA,
23 FusedMoEWithLoRA,
24 + GateLinearWithLoRA,
25 LogitsProcessorWithLoRA,
26 MergedColumnParallelLinearVariableSliceWithLoRA,
27 MergedColumnParallelLinearWithLoRA,
@@ -81,6 +82,7 @@ _all_lora_classes: set[type[BaseLayerWithLoRA]] = {
82 MergedQKVParallelLinearWithLoRA,
83 RowParallelLinearWithLoRA,
84 ReplicatedLinearWithLoRA,
85 + GateLinearWithLoRA,
86 LogitsProcessorWithLoRA,
87 ColumnParallelLinearWithShardedLoRA,
88 QKVParallelLinearWithShardedLoRA,
vllm/model_executor/layers/fused_moe/router/gate_linear.py
+52 -6
@@ -3,9 +3,11 @@
3 import torch
4 from torch.nn.parameter import Parameter
5
6 +import vllm._custom_ops as ops
7 from vllm.model_executor.custom_op import PluggableLayer
8 from vllm.model_executor.layers.linear import ReplicatedLinear
9 from vllm.platforms import current_platform
10 +from vllm.utils.torch_utils import direct_register_custom_op
11
12
13 @PluggableLayer.register("gate_linear")
@@ -13,8 +15,9 @@ class GateLinear(ReplicatedLinear):
15 """MoE gate linear layer with three-tier GEMM dispatch:
16
17 1. DSV3 specialized kernel (SM90+, batch<=16, supported dims)
16 - 2. cuBLAS bf16×bf16→fp32 (SM90+ + bf16 + fp32 out_dtype)
17 - 3. F.linear via ReplicatedLinear (ultimate fallback)
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)
21
22 The ``out_dtype`` attribute is mutable and can be set after init
23 (e.g. when the required dtype depends on the expert quantization
@@ -25,6 +28,10 @@ class GateLinear(ReplicatedLinear):
28 DSV3_SUPPORTED_NUM_EXPERTS = [256, 384]
29 DSV3_SUPPORTED_HIDDEN_SIZES = [7168]
30
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 +
35 def __init__(
36 self,
37 input_size: int,
@@ -65,6 +72,15 @@ class GateLinear(ReplicatedLinear):
72 and input_size in self.DSV3_SUPPORTED_HIDDEN_SIZES
73 )
74
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 +
84 # cuBLAS bf16→fp32 eligibility
85 self.allow_cublas_router_gemm = (
86 self.allow_specialized_router_gemm
@@ -92,8 +108,6 @@ class GateLinear(ReplicatedLinear):
108 def forward(
109 self, x: torch.Tensor
110 ) -> torch.Tensor | tuple[torch.Tensor, Parameter | None]:
95 - import vllm._custom_ops as ops
96 -
111 # Tier 1: DSV3 specialized kernel
112 if self.allow_dsv3_router_gemm and x.shape[0] <= 16:
113 output = ops.dsv3_router_gemm(
@@ -103,15 +117,47 @@ class GateLinear(ReplicatedLinear):
117 )
118 return output, None
119
106 - # Tier 2: cuBLAS bf16→fp32
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
126 if self.allow_cublas_router_gemm and x.dtype == torch.bfloat16:
127 output = ops.router_gemm_bf16_fp32(x, self.weight)
128 return output, None
129
111 - # Tier 3: F.linear (ReplicatedLinear)
130 + # Tier 4: F.linear (ReplicatedLinear)
131 if self.out_dtype is not None and x.dtype != self.weight.dtype:
132 x = x.to(self.weight.dtype)
133 output, output_bias = super().forward(x)
134 if self.out_dtype is not None and output.dtype != self.out_dtype:
135 output = output.to(self.out_dtype)
136 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
+3 -7
@@ -20,12 +20,11 @@ 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
23 +from vllm.model_executor.layers.fused_moe import FusedMoE, GateLinear
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,
28 RowParallelLinear,
29 )
30 from vllm.model_executor.layers.logits_processor import LogitsProcessor
@@ -175,13 +174,11 @@ class MLPBlock(torch.nn.Module):
174 self.hidden_size = config.hidden_size
175 self.experts_per_token = config.num_experts_per_tok
176 self.world_size = dist.get_world_size() if dist.is_initialized() else 1
178 - self.router = ReplicatedLinear(
177 + self.router = GateLinear(
178 config.hidden_size,
179 config.num_local_experts,
180 bias=True,
182 - quant_config=None,
181 prefix=f"{prefix}.router",
184 - return_bias=False,
182 )
183 assert config.intermediate_size % self.world_size == 0
184 self.experts = FusedMoE(
@@ -209,7 +206,7 @@ class MLPBlock(torch.nn.Module):
206 self, x[:, : self.hidden_size], self.router.weight, self.router.bias
207 )
208 else:
212 - g = self.router(x)
209 + g, _ = self.router(x)
210 x = self.experts(hidden_states=x, router_logits=g)[:, : self.hidden_size]
211
212 if self.is_sequence_parallel:
@@ -273,7 +270,6 @@ class GptOssModel(nn.Module, EagleModelMixin):
270 self.config = vllm_config.model_config.hf_config
271 self.quant_config = vllm_config.quant_config
272 self.parallel_config = vllm_config.parallel_config
276 - self.config.hidden_size = self.config.hidden_size
273 self.embedding = VocabParallelEmbedding(
274 self.config.vocab_size,
275 self.config.hidden_size,