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 7, 14:00 UTC
[Bugfix][Model] Fix gpt-oss batch invariance (#35404)
Signed-off-by: Jakub Zakrzewski <jzakrzewski@nvidia.com>
Jakub Zakrzewski committed
Feb 27, 2026 at 21:41 UTC
1f3dbd95fd13849e974f1f31ff36b3e91d7768bc
2 files changed
+14
-8
vllm/model_executor/layers/linear.py
+1
-6
@@ -28,7 +28,6 @@ from vllm.model_executor.layers.quantization.base_config import (
28
)
29
from vllm.model_executor.layers.utils import (
30
dispatch_unquantized_gemm,
31
- is_layer_moe_router_gate,
31
)
32
from vllm.model_executor.parameter import (
33
BasevLLMParameter,
@@ -257,11 +256,7 @@ class UnquantizedLinearMethod(LinearMethodBase):
256
x: torch.Tensor,
257
bias: torch.Tensor | None = None,
258
) -> torch.Tensor:
260
- if (
261
- vllm_is_batch_invariant()
262
- and current_platform.is_cuda_alike()
263
- and is_layer_moe_router_gate(getattr(layer, "prefix", ""))
264
- ):
259
+ if vllm_is_batch_invariant() and current_platform.is_cuda_alike():
260
return linear_batch_invariant(x, layer.weight, bias)
261
return dispatch_unquantized_gemm()(layer, x, layer.weight, bias)
262
vllm/model_executor/models/gpt_oss.py
+13
-2
@@ -23,7 +23,11 @@ from vllm.model_executor.layers.attention import Attention
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 QKVParallelLinear, RowParallelLinear
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
32
from vllm.model_executor.layers.quantization import QuantizationConfig
33
from vllm.model_executor.layers.quantization.utils.ocp_mx_utils import OCP_MX_BLOCK_SIZE
@@ -165,7 +169,14 @@ class MLPBlock(torch.nn.Module):
169
self.hidden_size = config.hidden_size
170
self.experts_per_token = config.num_experts_per_tok
171
self.world_size = dist.get_world_size() if dist.is_initialized() else 1
168
- self.router = torch.nn.Linear(config.hidden_size, config.num_local_experts)
172
+ self.router = ReplicatedLinear(
173
+ config.hidden_size,
174
+ config.num_local_experts,
175
+ bias=True,
176
+ quant_config=None,
177
+ prefix=f"{prefix}.router",
178
+ return_bias=False,
179
+ )
180
assert config.intermediate_size % self.world_size == 0
181
self.experts = FusedMoE(
182
num_experts=config.num_local_experts,