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

[Bugfix] Fix gpt_oss packed_modules_mapping (#28536)

Signed-off-by: Jee Jee Li <pandaleefree@gmail.com>

Jee Jee Li committed Nov 12, 2025 at 21:02 UTC a9d18b51078d29f732e2811e91e8c002b1d1d0c4
1 file changed +5 -5
vllm/model_executor/models/gpt_oss.py
+5 -5
@@ -92,7 +92,7 @@ class OAIAttention(nn.Module):
92 self.scaling = self.head_dim**-0.5
93 self.rope_theta = config.rope_theta
94
95 - self.qkv = QKVParallelLinear(
95 + self.qkv_proj = QKVParallelLinear(
96 hidden_size=self.hidden_size,
97 head_size=self.head_dim,
98 total_num_heads=self.num_attention_heads,
@@ -129,7 +129,7 @@ class OAIAttention(nn.Module):
129 def forward(
130 self, hidden_states: torch.Tensor, positions: torch.Tensor
131 ) -> torch.Tensor:
132 - qkv, _ = self.qkv(hidden_states)
132 + qkv, _ = self.qkv_proj(hidden_states)
133 q, k, v = qkv.split([self.q_size, self.kv_size, self.kv_size], dim=-1)
134 q, k = self.rotary_emb(positions, q, k)
135 v = v.contiguous()
@@ -606,9 +606,9 @@ class GptOssModel(nn.Module):
606 def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]:
607 stacked_params_mapping = [
608 # (param_name, shard_name, shard_id)
609 - (".qkv", ".q_proj", "q"),
610 - (".qkv", ".k_proj", "k"),
611 - (".qkv", ".v_proj", "v"),
609 + (".qkv_proj", ".q_proj", "q"),
610 + (".qkv_proj", ".k_proj", "k"),
611 + (".qkv_proj", ".v_proj", "v"),
612 ]
613
614 tp_rank = get_tensor_model_parallel_rank()