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 14, 05:30 UTC
[fix]: remove data type hardcoding from gptoss model implementation (#23807)
Signed-off-by: Nikhil Gupta <nikhil.gupta2@arm.com>
Nikhil Gupta committed
Sep 18, 2025 at 19:15 UTC
064cac7bb7251862a841d8057d83581350edf837
1 file changed
+1
-3
vllm/model_executor/models/gpt_oss.py
+1
-3
@@ -76,7 +76,6 @@ class OAIAttention(nn.Module):
76
77
self.sinks = torch.nn.Parameter(
78
torch.empty(config.num_attention_heads // tp_size,
79
- dtype=torch.bfloat16,
79
requires_grad=False))
80
81
self.q_size = self.num_attention_heads * self.head_dim // tp_size
@@ -145,8 +144,7 @@ class MLPBlock(torch.nn.Module):
144
self.experts_per_token = config.num_experts_per_tok
145
self.world_size = dist.get_world_size() if dist.is_initialized() else 1
146
self.router = torch.nn.Linear(config.hidden_size,
148
- config.num_local_experts,
149
- dtype=torch.bfloat16)
147
+ config.num_local_experts)
148
assert config.intermediate_size % self.world_size == 0
149
self.experts = FusedMoE(num_experts=config.num_local_experts,
150
top_k=config.num_experts_per_tok,