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
[Model] Add PP support and VLM backbone compatability for GPT-OSS (#23680)
Signed-off-by: Isotr0py <mozf@mail2.sysu.edu.cn>
Isotr0py committed
Aug 28, 2025 at 16:03 UTC
c5d004aaaf3b2106d33974c673bec0568c18f762
2 files changed
+87
-34
docs/models/supported_models.md
+1
-1
@@ -358,7 +358,7 @@ th {
358
| `GPTBigCodeForCausalLM` | StarCoder, SantaCoder, WizardCoder | `bigcode/starcoder`, `bigcode/gpt_bigcode-santacoder`, `WizardLM/WizardCoder-15B-V1.0`, etc. | ✅︎ | ✅︎ | ✅︎ |
359
| `GPTJForCausalLM` | GPT-J | `EleutherAI/gpt-j-6b`, `nomic-ai/gpt4all-j`, etc. | | ✅︎ | ✅︎ |
360
| `GPTNeoXForCausalLM` | GPT-NeoX, Pythia, OpenAssistant, Dolly V2, StableLM | `EleutherAI/gpt-neox-20b`, `EleutherAI/pythia-12b`, `OpenAssistant/oasst-sft-4-pythia-12b-epoch-3.5`, `databricks/dolly-v2-12b`, `stabilityai/stablelm-tuned-alpha-7b`, etc. | | ✅︎ | ✅︎ |
361
-| `GptOssForCausalLM` | GPT-OSS | `openai/gpt-oss-120b`, `openai/gpt-oss-20b` | | | ✅︎ |
361
+| `GptOssForCausalLM` | GPT-OSS | `openai/gpt-oss-120b`, `openai/gpt-oss-20b` | | ✅︎ | ✅︎ |
362
| `GraniteForCausalLM` | Granite 3.0, Granite 3.1, PowerLM | `ibm-granite/granite-3.0-2b-base`, `ibm-granite/granite-3.1-8b-instruct`, `ibm/PowerLM-3b`, etc. | ✅︎ | ✅︎ | ✅︎ |
363
| `GraniteMoeForCausalLM` | Granite 3.0 MoE, PowerMoE | `ibm-granite/granite-3.0-1b-a400m-base`, `ibm-granite/granite-3.0-3b-a800m-instruct`, `ibm/PowerMoE-3b`, etc. | ✅︎ | ✅︎ | ✅︎ |
364
| `GraniteMoeHybridForCausalLM` | Granite 4.0 MoE Hybrid | `ibm-granite/granite-4.0-tiny-preview`, etc. | ✅︎ | ✅︎ | ✅︎ |
vllm/model_executor/models/gpt_oss.py
+86
-33
@@ -11,7 +11,8 @@ from transformers import GptOssConfig
11
from vllm.attention import Attention, AttentionType
12
from vllm.compilation.decorators import support_torch_compile
13
from vllm.config import CacheConfig, VllmConfig
14
-from vllm.distributed import (get_ep_group, get_tensor_model_parallel_rank,
14
+from vllm.distributed import (get_ep_group, get_pp_group,
15
+ get_tensor_model_parallel_rank,
16
get_tensor_model_parallel_world_size)
17
from vllm.model_executor.layers.fused_moe import FusedMoE
18
from vllm.model_executor.layers.layernorm import RMSNorm
@@ -27,7 +28,10 @@ from vllm.model_executor.sampling_metadata import SamplingMetadata
28
from vllm.sequence import IntermediateTensors
29
from vllm.utils import cdiv
30
31
+from .interfaces import SupportsPP
32
from .utils import (AutoWeightsLoader, WeightsMapper, extract_layer_index,
33
+ is_pp_missing_parameter,
34
+ make_empty_intermediate_tensors_factory, make_layers,
35
maybe_prefix)
36
37
@@ -75,8 +79,6 @@ class OAIAttention(nn.Module):
79
dtype=torch.bfloat16,
80
requires_grad=False))
81
78
- self.norm = RMSNorm(config.hidden_size, eps=1e-5)
79
-
82
self.q_size = self.num_attention_heads * self.head_dim // tp_size
83
self.kv_size = self.num_key_value_heads * self.head_dim // tp_size
84
self.scaling = self.head_dim**-0.5
@@ -119,16 +121,13 @@ class OAIAttention(nn.Module):
121
122
def forward(self, hidden_states: torch.Tensor,
123
positions: torch.Tensor) -> torch.Tensor:
122
- t = self.norm(hidden_states)
123
-
124
- qkv, _ = self.qkv(t)
124
+ qkv, _ = self.qkv(hidden_states)
125
q, k, v = qkv.split([self.q_size, self.kv_size, self.kv_size], dim=-1)
126
q, k = self.rotary_emb(positions, q, k)
127
v = v.contiguous()
128
attn_output = self.attn(q, k, v)
129
output, _ = self.o_proj(attn_output)
130
-
131
- return output + hidden_states
130
+ return output
131
132
133
class MLPBlock(torch.nn.Module):
@@ -145,7 +144,6 @@ class MLPBlock(torch.nn.Module):
144
self.num_experts = config.num_local_experts
145
self.experts_per_token = config.num_experts_per_tok
146
self.world_size = dist.get_world_size() if dist.is_initialized() else 1
148
- self.norm = RMSNorm(config.hidden_size, eps=1e-5)
147
self.router = torch.nn.Linear(config.hidden_size,
148
config.num_local_experts,
149
dtype=torch.bfloat16)
@@ -163,10 +161,9 @@ class MLPBlock(torch.nn.Module):
161
activation="swigluoai")
162
163
def forward(self, x: torch.Tensor) -> torch.Tensor:
166
- t = self.norm(x)
167
- g = self.router(t)
168
- t = self.experts(hidden_states=t, router_logits=g)
169
- return x + t
164
+ g = self.router(x)
165
+ x = self.experts(hidden_states=x, router_logits=g)
166
+ return x
167
168
169
class TransformerBlock(torch.nn.Module):
@@ -187,12 +184,28 @@ class TransformerBlock(torch.nn.Module):
184
self.layer_idx,
185
quant_config=quant_config,
186
prefix=f"{prefix}.mlp")
187
+ self.input_layernorm = RMSNorm(config.hidden_size, eps=1e-5)
188
+ self.post_attention_layernorm = RMSNorm(config.hidden_size, eps=1e-5)
189
191
- def forward(self, hidden_states: torch.Tensor,
192
- positions: torch.Tensor) -> torch.Tensor:
193
- attn_output = self.attn(hidden_states, positions)
194
- output = self.mlp(attn_output)
195
- return output
190
+ def forward(
191
+ self,
192
+ hidden_states: torch.Tensor,
193
+ positions: torch.Tensor,
194
+ residual: Optional[torch.Tensor],
195
+ ) -> torch.Tensor:
196
+ # Self Attention
197
+ if residual is None:
198
+ residual = hidden_states
199
+ hidden_states = self.input_layernorm(hidden_states)
200
+ else:
201
+ hidden_states, residual = self.input_layernorm(
202
+ hidden_states, residual)
203
+ hidden_states = self.attn(hidden_states, positions)
204
+ # Fully Connected
205
+ hidden_states, residual = self.post_attention_layernorm(
206
+ hidden_states, residual)
207
+ output = self.mlp(hidden_states)
208
+ return output, residual
209
210
211
@support_torch_compile
@@ -214,22 +227,52 @@ class GptOssModel(nn.Module):
227
self.config.vocab_size,
228
self.config.hidden_size,
229
)
217
- self.layers = torch.nn.ModuleList([
218
- TransformerBlock(
230
+ self.start_layer, self.end_layer, self.layers = make_layers(
231
+ self.config.num_hidden_layers,
232
+ lambda prefix: TransformerBlock(
233
self.config,
234
cache_config=self.cache_config,
235
quant_config=self.quant_config,
222
- prefix=maybe_prefix(prefix, f"block.{layer_idx}"),
223
- ) for layer_idx in range(self.config.num_hidden_layers)
224
- ])
236
+ prefix=prefix,
237
+ ),
238
+ prefix=f"{prefix}.layers",
239
+ )
240
self.norm = RMSNorm(self.config.hidden_size, eps=1e-5)
241
+ self.make_empty_intermediate_tensors = (
242
+ make_empty_intermediate_tensors_factory(
243
+ ["hidden_states", "residual"], self.config.hidden_size))
244
227
- def forward(self, input_ids: torch.Tensor,
228
- positions: torch.Tensor) -> torch.Tensor:
229
- x = self.embedding(input_ids)
230
- for layer in self.layers:
231
- x = layer(x, positions)
232
- x = self.norm(x)
245
+ def get_input_embeddings(self, input_ids: torch.Tensor) -> torch.Tensor:
246
+ return self.embedding(input_ids)
247
+
248
+ def forward(
249
+ self,
250
+ input_ids: torch.Tensor,
251
+ positions: torch.Tensor,
252
+ intermediate_tensors: Optional[IntermediateTensors] = None,
253
+ inputs_embeds: Optional[torch.Tensor] = None,
254
+ ) -> torch.Tensor:
255
+ if get_pp_group().is_first_rank:
256
+ if inputs_embeds is not None:
257
+ x = inputs_embeds
258
+ else:
259
+ x = self.get_input_embeddings(input_ids)
260
+
261
+ residual = None
262
+ else:
263
+ assert intermediate_tensors is not None
264
+ x = intermediate_tensors["hidden_states"]
265
+ residual = intermediate_tensors["residual"]
266
+
267
+ for i in range(self.start_layer, self.end_layer):
268
+ layer = self.layers[i]
269
+ x, residual = layer(x, positions, residual)
270
+ if not get_pp_group().is_last_rank:
271
+ return IntermediateTensors({
272
+ "hidden_states": x,
273
+ "residual": residual
274
+ })
275
+ x, _ = self.norm(x, residual)
276
return x
277
278
def _load_weights_mxfp4(
@@ -264,6 +307,10 @@ class GptOssModel(nn.Module):
307
intermediate_size)
308
309
for name, weight in weights:
310
+ # Skip layers on other devices.
311
+ if is_pp_missing_parameter(name, self):
312
+ continue
313
+
314
# FIXME(woosuk): Remove this after testing.
315
weight = weight.cuda()
316
@@ -445,6 +492,10 @@ class GptOssModel(nn.Module):
492
intermediate_size)
493
494
for name, weight in weights:
495
+ # Skip layers on other devices.
496
+ if is_pp_missing_parameter(name, self):
497
+ continue
498
+
499
if ".w13_weight" in name:
500
# Handle MLP gate and up projection weights
501
# Extract gate and up projection parts
@@ -562,18 +613,15 @@ class GptOssModel(nn.Module):
613
weights, stacked_params_mapping)
614
615
565
-class GptOssForCausalLM(nn.Module):
616
+class GptOssForCausalLM(nn.Module, SupportsPP):
617
packed_modules_mapping = {"qkv": ["q_proj", "k_proj", "v_proj"]}
618
619
hf_to_vllm_mapper = WeightsMapper(
620
orig_to_new_substr={
621
".self_attn.": ".attn.",
571
- ".post_attention_layernorm.": ".mlp.norm.",
622
},
623
orig_to_new_suffix={
624
".embed_tokens.weight": ".embedding.weight",
575
- ".input_layernorm.weight": ".attn.norm.weight",
576
- ".post_attention_layernorm.weight": ".mlp.norm.weight",
625
626
# MoE MXFP4 weights
627
".gate_up_proj_blocks": ".w13_weight",
@@ -609,6 +657,11 @@ class GptOssForCausalLM(nn.Module):
657
self.config.hidden_size,
658
)
659
self.logits_processor = LogitsProcessor(self.config.vocab_size)
660
+ self.make_empty_intermediate_tensors = (
661
+ self.model.make_empty_intermediate_tensors)
662
+
663
+ def get_input_embeddings(self, input_ids: torch.Tensor) -> torch.Tensor:
664
+ return self.model.get_input_embeddings(input_ids)
665
666
def forward(self,
667
input_ids: torch.Tensor,