@tduncan / Nvidia-Optimization-Stack-VLLM Mirror · read-only / commits / de98252f49
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 14, 05:30 UTC

Add GPT-OSS model code and config [1/N] (#22327)

Signed-off-by: Woosuk Kwon <woosuk.kwon@berkeley.edu>

Woosuk Kwon committed Aug 5, 2025 at 23:26 UTC de98252f497b8cde5b9f18a8dac53302f5c72db7
4 files changed +503
tests/models/registry.py
+1
@@ -197,6 +197,7 @@ _TEXT_GENERATION_EXAMPLE_MODELS = {
197 {"6b": "EleutherAI/gpt-j-6b"}),
198 "GPTNeoXForCausalLM": _HfExamplesInfo("EleutherAI/pythia-70m",
199 {"1b": "EleutherAI/pythia-1.4b"}),
200 + "GptOssForCausalLM": _HfExamplesInfo("openai/gpt-oss-20b"),
201 "GraniteForCausalLM": _HfExamplesInfo("ibm/PowerLM-3b"),
202 "GraniteMoeForCausalLM": _HfExamplesInfo("ibm/PowerMoE-3b"),
203 "GraniteMoeHybridForCausalLM": _HfExamplesInfo("ibm-granite/granite-4.0-tiny-preview"), # noqa: E501
vllm/model_executor/models/config.py
+29
@@ -247,6 +247,34 @@ class GraniteMoeHybridModelConfig(VerifyAndUpdateConfig):
247 config.max_model_len)
248
249
250 +class GptOssConfig(VerifyAndUpdateConfig):
251 +
252 + @staticmethod
253 + def verify_and_update_config(vllm_config: "VllmConfig") -> None:
254 + decoding_config = vllm_config.decoding_config
255 + if decoding_config.reasoning_backend == "":
256 + decoding_config.reasoning_backend = "openai"
257 +
258 + # Increase the max capture size from 512 to 1024 for performance.
259 + # NOTE(woosuk): This will increase the number of CUDA graphs
260 + # from 67 to 83.
261 + scheduler_config = vllm_config.scheduler_config
262 + if len(scheduler_config.cuda_graph_sizes) == 1:
263 + max_capture_size = scheduler_config.cuda_graph_sizes[0]
264 + # FIXME(woosuk): When using full cuda graph with FA3, the max
265 + # supported size is 992.
266 + if max_capture_size < 1024:
267 + cuda_graph_sizes = [1, 2, 4]
268 + # Step size 8 for small batch sizes
269 + cuda_graph_sizes += [i for i in range(8, 256, 8)]
270 + # Step size 16 for larger batch sizes
271 + cuda_graph_sizes += [i for i in range(256, 1025, 16)]
272 + scheduler_config.cuda_graph_sizes = cuda_graph_sizes
273 + logger.info(
274 + "Overriding max cuda graph capture size to "
275 + "%d for performance.", 1024)
276 +
277 +
278 class HybridAttentionMambaModelConfig(VerifyAndUpdateConfig):
279
280 @classmethod
@@ -345,4 +373,5 @@ MODELS_CONFIG_MAP: dict[str, type[VerifyAndUpdateConfig]] = {
373 "JinaVLForRanking": JinaVLForSequenceClassificationConfig,
374 "JambaForSequenceClassification": JambaForSequenceClassificationConfig,
375 "GraniteMoeHybridForCausalLM": GraniteMoeHybridModelConfig,
376 + "GptOssForCausalLM": GptOssConfig,
377 }
vllm/model_executor/models/gpt_oss.py new
+472
@@ -0,0 +1,472 @@
1 +# SPDX-License-Identifier: Apache-2.0
2 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 +from collections.abc import Iterable
4 +from typing import Optional
5 +
6 +import torch
7 +import torch.distributed as dist
8 +from torch import nn
9 +from transformers import GptOssConfig
10 +
11 +from vllm import envs
12 +from vllm.attention import Attention, AttentionType
13 +from vllm.compilation.decorators import support_torch_compile
14 +from vllm.config import CacheConfig, VllmConfig
15 +from vllm.distributed import (get_ep_group, 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
19 +from vllm.model_executor.layers.linear import (QKVParallelLinear,
20 + RowParallelLinear)
21 +from vllm.model_executor.layers.logits_processor import LogitsProcessor
22 +from vllm.model_executor.layers.quantization import QuantizationConfig
23 +from vllm.model_executor.layers.rotary_embedding import get_rope
24 +from vllm.model_executor.layers.vocab_parallel_embedding import (
25 + ParallelLMHead, VocabParallelEmbedding)
26 +from vllm.model_executor.model_loader.weight_utils import default_weight_loader
27 +from vllm.model_executor.sampling_metadata import SamplingMetadata
28 +from vllm.sequence import IntermediateTensors
29 +from vllm.utils import cdiv
30 +
31 +from .utils import extract_layer_index, maybe_prefix
32 +
33 +
34 +class OAIAttention(nn.Module):
35 +
36 + def __init__(
37 + self,
38 + config: GptOssConfig,
39 + quant_config: Optional[QuantizationConfig] = None,
40 + cache_config: Optional[CacheConfig] = None,
41 + prefix: str = "",
42 + ):
43 + super().__init__()
44 + self.layer_idx = extract_layer_index(prefix)
45 + self.head_dim = config.head_dim
46 + self.num_attention_heads = config.num_attention_heads
47 + self.num_key_value_heads = config.num_key_value_heads
48 + self.hidden_size = config.hidden_size
49 +
50 + self.rotary_emb = get_rope(
51 + self.head_dim,
52 + rotary_dim=self.head_dim,
53 + max_position=config.max_position_embeddings,
54 + base=config.rope_theta,
55 + dtype=torch.float32,
56 + rope_scaling={
57 + "rope_type":
58 + "yarn",
59 + "factor":
60 + config.rope_scaling["factor"],
61 + "original_max_position_embeddings":
62 + config.rope_scaling["original_max_position_embeddings"],
63 + "beta_fast":
64 + config.rope_ntk_beta,
65 + "beta_slow":
66 + config.rope_ntk_alpha,
67 + },
68 + is_neox_style=True,
69 + )
70 +
71 + tp_size = get_tensor_model_parallel_world_size()
72 +
73 + attention_sink_dtype = (
74 + torch.float32 if envs.VLLM_USE_TRTLLM_CONTEXT_ATTENTION
75 + or envs.VLLM_USE_TRTLLM_DECODE_ATTENTION else torch.bfloat16)
76 + self.sinks = torch.nn.Parameter(
77 + torch.empty(config.num_attention_heads // tp_size,
78 + dtype=attention_sink_dtype,
79 + requires_grad=False))
80 +
81 + self.norm = RMSNorm(config.hidden_size, eps=1e-5)
82 +
83 + self.q_size = self.num_attention_heads * self.head_dim // tp_size
84 + self.kv_size = self.num_key_value_heads * self.head_dim // tp_size
85 + self.scaling = self.head_dim**-0.5
86 + self.rope_theta = config.rope_theta
87 +
88 + self.qkv = QKVParallelLinear(
89 + hidden_size=self.hidden_size,
90 + head_size=self.head_dim,
91 + total_num_heads=self.num_attention_heads,
92 + total_num_kv_heads=self.num_key_value_heads,
93 + quant_config=quant_config,
94 + prefix=f"{prefix}.qkv_proj",
95 + )
96 +
97 + self.o_proj = RowParallelLinear(
98 + input_size=self.num_attention_heads * self.head_dim,
99 + output_size=self.hidden_size,
100 + quant_config=quant_config,
101 + prefix=f"{prefix}.o_proj",
102 + )
103 +
104 + self.num_local_attention_heads = config.num_attention_heads // tp_size
105 + self.num_local_key_value_heads = config.num_key_value_heads // tp_size
106 +
107 + # Only apply sliding window to every other layer
108 + sliding_window = (config.sliding_window if self.layer_idx %
109 + 2 == 0 else None)
110 + self.attn = Attention(
111 + self.num_local_attention_heads,
112 + self.head_dim,
113 + self.scaling,
114 + num_kv_heads=self.num_local_key_value_heads,
115 + cache_config=cache_config,
116 + quant_config=quant_config,
117 + per_layer_sliding_window=sliding_window,
118 + attn_type=AttentionType.DECODER,
119 + prefix=f"{prefix}.attn",
120 + sinks=self.sinks,
121 + )
122 +
123 + def forward(self, hidden_states: torch.Tensor,
124 + positions: torch.Tensor) -> torch.Tensor:
125 + t = self.norm(hidden_states)
126 +
127 + qkv, _ = self.qkv(t)
128 + q, k, v = qkv.split([self.q_size, self.kv_size, self.kv_size], dim=-1)
129 + q, k = self.rotary_emb(positions, q, k)
130 + v = v.contiguous()
131 + attn_output = self.attn(q, k, v)
132 + output, _ = self.o_proj(attn_output)
133 +
134 + return output + hidden_states
135 +
136 +
137 +class MLPBlock(torch.nn.Module):
138 +
139 + def __init__(
140 + self,
141 + config: GptOssConfig,
142 + layer_idx: int,
143 + quant_config: QuantizationConfig,
144 + prefix: str = "",
145 + ):
146 + super().__init__()
147 + self.layer_idx = layer_idx
148 + self.num_experts = config.num_local_experts
149 + self.experts_per_token = config.num_experts_per_tok
150 + self.world_size = dist.get_world_size() if dist.is_initialized() else 1
151 + self.norm = RMSNorm(config.hidden_size, eps=1e-5)
152 + self.router = torch.nn.Linear(config.hidden_size,
153 + config.num_local_experts,
154 + dtype=torch.bfloat16)
155 + assert config.intermediate_size % self.world_size == 0
156 + self.experts = FusedMoE(num_experts=config.num_local_experts,
157 + top_k=config.num_experts_per_token,
158 + hidden_size=config.hidden_size,
159 + intermediate_size=config.intermediate_size,
160 + reduce_results=True,
161 + renormalize=True,
162 + quant_config=quant_config,
163 + prefix=f"{prefix}.experts",
164 + apply_router_weight_on_input=False)
165 +
166 + def forward(self, x: torch.Tensor) -> torch.Tensor:
167 + t = self.norm(x)
168 + g = self.router(t)
169 + t = self.experts(hidden_states=t, router_logits=g)
170 + return x + t
171 +
172 +
173 +class TransformerBlock(torch.nn.Module):
174 +
175 + def __init__(
176 + self,
177 + config: GptOssConfig,
178 + quant_config: QuantizationConfig,
179 + prefix: str = "",
180 + ):
181 + super().__init__()
182 + self.layer_idx = extract_layer_index(prefix)
183 + self.attn = OAIAttention(config, prefix=f"{prefix}.attn")
184 + self.mlp = MLPBlock(config,
185 + self.layer_idx,
186 + quant_config=quant_config,
187 + prefix=f"{prefix}.mlp")
188 +
189 + def forward(self, hidden_states: torch.Tensor,
190 + positions: torch.Tensor) -> torch.Tensor:
191 + attn_output = self.attn(hidden_states, positions)
192 + output = self.mlp(attn_output)
193 + return output
194 +
195 +
196 +@support_torch_compile
197 +class GptOssModel(nn.Module):
198 +
199 + def __init__(
200 + self,
201 + *,
202 + vllm_config: VllmConfig,
203 + prefix: str = "",
204 + ):
205 + super().__init__()
206 + self.config = vllm_config.model_config.hf_config
207 + self.quant_config = vllm_config.quant_config
208 + self.config.hidden_size = self.config.hidden_size
209 + self.embedding = VocabParallelEmbedding(
210 + self.config.vocab_size,
211 + self.config.hidden_size,
212 + )
213 + self.layers = torch.nn.ModuleList([
214 + TransformerBlock(
215 + self.config,
216 + quant_config=self.quant_config,
217 + prefix=maybe_prefix(prefix, f"block.{layer_idx}"),
218 + ) for layer_idx in range(self.config.num_hidden_layers)
219 + ])
220 + self.norm = RMSNorm(self.config.hidden_size, eps=1e-5)
221 +
222 + def forward(self, input_ids: torch.Tensor,
223 + positions: torch.Tensor) -> torch.Tensor:
224 + x = self.embedding(input_ids)
225 + for layer in self.layers:
226 + x = layer(x, positions)
227 + x = self.norm(x)
228 + return x
229 +
230 +
231 +class GptOssForCausalLM(nn.Module):
232 +
233 + def __init__(
234 + self,
235 + vllm_config: VllmConfig,
236 + prefix: str = "",
237 + ):
238 + super().__init__()
239 + self.vllm_config = vllm_config
240 + self.model_config = vllm_config.model_config.hf_config
241 + self.model = GptOssModel(
242 + vllm_config=vllm_config,
243 + prefix=maybe_prefix(prefix, "model"),
244 + )
245 + self.lm_head = ParallelLMHead(
246 + self.model_config.vocab_size,
247 + self.model_config.hidden_size,
248 + )
249 + self.logits_processor = LogitsProcessor(self.model_config.vocab_size)
250 +
251 + def forward(self,
252 + input_ids: torch.Tensor,
253 + positions: torch.Tensor,
254 + intermediate_tensors: Optional[IntermediateTensors] = None,
255 + inputs_embeds: Optional[torch.Tensor] = None) -> torch.Tensor:
256 + assert intermediate_tensors is None
257 + assert inputs_embeds is None
258 + return self.model(input_ids, positions)
259 +
260 + def compute_logits(self, hidden_states: torch.Tensor,
261 + sampling_metadata: SamplingMetadata) -> torch.Tensor:
262 + logits = self.logits_processor(self.lm_head, hidden_states,
263 + sampling_metadata)
264 + return logits
265 +
266 + def load_weights(self, weights: Iterable[tuple[str,
267 + torch.Tensor]]) -> set[str]:
268 + rename_mapping = {
269 + "self_attn": "attn",
270 + "input_layernorm.weight": "attn.norm.weight",
271 + "post_attention_layernorm.weight": "mlp.norm.weight",
272 + "embed_tokens": "embedding",
273 + }
274 +
275 + def maybe_rename(name: str) -> str:
276 + for remap_name, new_name in rename_mapping.items():
277 + if remap_name in name:
278 + return name.replace(remap_name, new_name)
279 + return name
280 +
281 + params_dict = dict(self.named_parameters())
282 + loaded_params: set[str] = set()
283 + mxfp4_block = 32
284 +
285 + tp_rank = get_tensor_model_parallel_rank()
286 + tp_size = get_tensor_model_parallel_world_size()
287 + intermediate_size = self.model_config.intermediate_size
288 + intermediate_size_block = intermediate_size // mxfp4_block
289 + per_rank_intermediate_size_block = cdiv(intermediate_size_block,
290 + tp_size)
291 + per_rank_intermediate_size = (per_rank_intermediate_size_block *
292 + mxfp4_block)
293 +
294 + # Calculate common slicing bounds for current rank
295 + tp_rank_start = tp_rank * per_rank_intermediate_size
296 + tp_rank_end = min((tp_rank + 1) * per_rank_intermediate_size,
297 + intermediate_size)
298 +
299 + # Attention heads per rank
300 + heads_per_rank = self.model_config.num_attention_heads // tp_size
301 + head_start = tp_rank * heads_per_rank
302 +
303 + use_ep = self.vllm_config.parallel_config.enable_expert_parallel
304 + ep_size = get_ep_group().world_size
305 + ep_rank = get_ep_group().rank
306 + num_experts = self.model_config.num_local_experts
307 + experts_per_rank = num_experts // ep_size
308 + ep_rank_start = ep_rank * experts_per_rank
309 + ep_rank_end = (ep_rank + 1) * experts_per_rank
310 +
311 + for name, weight in weights:
312 + # FIXME(woosuk): Remove this after testing.
313 + weight = weight.cuda()
314 +
315 + if "gate_up_proj_blocks" in name:
316 + # Handle MLP gate and up projection weights
317 + new_name = name.replace("gate_up_proj_blocks", "w13_weight")
318 +
319 + # flat weight from (E, 2 * N, block_size, entry_per_block)
320 + # to (E, 2 * N, -1), shouldn't trigger copy for contiguous
321 + weight = weight.view(num_experts, 2 * intermediate_size,
322 + -1).contiguous()
323 +
324 + # Extract gate and up projection parts
325 + # since the weight is shuffled, we can slice directly
326 + if use_ep:
327 + narrow_weight = weight[ep_rank_start:ep_rank_end, ...]
328 + else:
329 + narrow_weight = weight[:,
330 + 2 * tp_rank_start:2 * tp_rank_end,
331 + ...]
332 +
333 + param = params_dict[new_name]
334 + weight_loader = getattr(param, "weight_loader",
335 + default_weight_loader)
336 + weight_loader(param,
337 + narrow_weight,
338 + weight_name=new_name,
339 + shard_id=None,
340 + expert_id=None)
341 + loaded_params.add(new_name)
342 +
343 + elif "down_proj_blocks" in name:
344 + # Handle MLP down projection weights
345 + new_name = name.replace("down_proj_blocks", "w2_weight")
346 + # same flatten here, but since 2 mx4 value are packed in 1
347 + # uint8, divide by 2
348 + weight = weight.view(num_experts, -1,
349 + intermediate_size // 2).contiguous()
350 + if use_ep:
351 + narrow_weight = weight[ep_rank_start:ep_rank_end, ...]
352 + else:
353 + narrow_weight = weight[...,
354 + tp_rank_start // 2:tp_rank_end // 2]
355 +
356 + param = params_dict[new_name]
357 + weight_loader = getattr(param, "weight_loader",
358 + default_weight_loader)
359 + weight_loader(param,
360 + narrow_weight,
361 + weight_name=new_name,
362 + shard_id=None,
363 + expert_id=None)
364 + loaded_params.add(new_name)
365 +
366 + elif "gate_up_proj_scales" in name:
367 + # Handle MLP gate and up projection weights scale
368 + new_name = name.replace("gate_up_proj_scales",
369 + "w13_weight_scale")
370 + if use_ep:
371 + narrow_weight = weight[ep_rank_start:ep_rank_end, ...]
372 + else:
373 + narrow_weight = weight[:,
374 + 2 * tp_rank_start:2 * tp_rank_end,
375 + ...]
376 +
377 + param = params_dict[new_name]
378 + weight_loader = getattr(param, "weight_loader",
379 + default_weight_loader)
380 + weight_loader(param,
381 + narrow_weight,
382 + weight_name=new_name,
383 + shard_id=None,
384 + expert_id=None)
385 + loaded_params.add(new_name)
386 +
387 + elif "down_proj_scales" in name:
388 + # Handle MLP down projection weights
389 + new_name = name.replace("down_proj_scales", "w2_weight_scale")
390 + if use_ep:
391 + narrow_weight = weight[ep_rank_start:ep_rank_end, ...]
392 + else:
393 + narrow_weight = weight[..., tp_rank_start //
394 + mxfp4_block:tp_rank_end //
395 + mxfp4_block]
396 +
397 + param = params_dict[new_name]
398 + weight_loader = getattr(param, "weight_loader",
399 + default_weight_loader)
400 + weight_loader(param,
401 + narrow_weight,
402 + weight_name=new_name,
403 + shard_id=None,
404 + expert_id=None)
405 + loaded_params.add(new_name)
406 + elif "gate_up_proj_bias" in name:
407 + # Handle MLP gate and up projection biases
408 + new_name = name.replace("gate_up_proj_bias", "w13_bias")
409 +
410 + # Extract gate and up projection bias parts
411 + if use_ep:
412 + narrow_weight = weight[ep_rank_start:ep_rank_end, ...]
413 + else:
414 + narrow_weight = weight[:,
415 + 2 * tp_rank_start:2 * tp_rank_end]
416 +
417 + param = params_dict[new_name]
418 + weight_loader = getattr(param, "weight_loader",
419 + default_weight_loader)
420 + weight_loader(param,
421 + narrow_weight,
422 + weight_name=new_name,
423 + shard_id=None,
424 + expert_id=None)
425 + loaded_params.add(new_name)
426 +
427 + elif "down_proj_bias" in name:
428 + # Handle MLP down projection bias
429 + new_name = name.replace("down_proj_bias", "w2_bias")
430 + param = params_dict[new_name]
431 + weight_loader = getattr(param, "weight_loader",
432 + default_weight_loader)
433 + if use_ep:
434 + weight = weight[ep_rank_start:ep_rank_end, ...]
435 + else:
436 + # (only load on rank 0 to avoid duplication)
437 + if tp_rank != 0:
438 + weight.zero_()
439 + weight_loader(param,
440 + weight,
441 + weight_name=new_name,
442 + shard_id=None,
443 + expert_id=None)
444 + loaded_params.add(new_name)
445 + elif "sinks" in name:
446 + # Handle attention sinks (distributed across ranks)
447 + name = name.replace("self_attn", "attn")
448 + param = params_dict[name]
449 + narrow_weight = weight.narrow(0, head_start, heads_per_rank)
450 + param.data.copy_(narrow_weight)
451 + loaded_params.add(name)
452 + elif "q_proj" in name or "k_proj" in name or "v_proj" in name:
453 + shard_id = ("q" if "q_proj" in name else
454 + "k" if "k_proj" in name else "v")
455 + name = name.replace("self_attn", "attn")
456 + param_name = name.replace(f"{shard_id}_proj", "qkv")
457 + param = params_dict[param_name]
458 + weight_loader = param.weight_loader
459 + weight_loader(param, weight, loaded_shard_id=shard_id)
460 + loaded_params.add(param_name)
461 + else:
462 + # Handle all other weights with potential renaming
463 + renamed_name = maybe_rename(name)
464 + if renamed_name not in params_dict:
465 + continue
466 + param = params_dict[renamed_name]
467 + weight_loader = getattr(param, "weight_loader",
468 + default_weight_loader)
469 + weight_loader(param, weight)
470 + loaded_params.add(renamed_name)
471 +
472 + return loaded_params
vllm/model_executor/models/registry.py
+1
@@ -74,6 +74,7 @@ _TEXT_GENERATION_MODELS = {
74 "GlmForCausalLM": ("glm", "GlmForCausalLM"),
75 "Glm4ForCausalLM": ("glm4", "Glm4ForCausalLM"),
76 "Glm4MoeForCausalLM": ("glm4_moe", "Glm4MoeForCausalLM"),
77 + "GptOssForCausalLM": ("gpt_oss", "GptOssForCausalLM"),
78 "GPT2LMHeadModel": ("gpt2", "GPT2LMHeadModel"),
79 "GPTBigCodeForCausalLM": ("gpt_bigcode", "GPTBigCodeForCausalLM"),
80 "GPTJForCausalLM": ("gpt_j", "GPTJForCausalLM"),