@tduncan / Nvidia-Optimization-Stack-VLLM Mirror · read-only / commits / 87985077a4

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 6, 19:15 UTC

[Speculative Decoding] Add `norm_before_fc` for gpt-oss draft models (#36545)

Signed-off-by: Shubhra Pandit <shubhra.pandit@gmail.com> Co-authored-by: Benjamin Chislett <chislett.ben@gmail.com> Co-authored-by: Benjamin Chislett <bchislett@nvidia.com>

Shubhra Pandit committed Mar 12, 2026 at 19:03 UTC 87985077a45b721355efd7c406384254910f963f

1 file changed +13

vllm/model_executor/models/llama_eagle3.py

+13

@@ -150,6 +150,7 @@ class LlamaModel(nn.Module):

150 self.use_aux_hidden_state = eagle_config["use_aux_hidden_state"]

151 else:

152 self.use_aux_hidden_state = True

153 + self.norm_before_fc = getattr(self.config, "norm_before_fc", False)

154

155 current_vllm_config = get_current_vllm_config()

156

@@ -175,6 +176,13 @@ class LlamaModel(nn.Module):

176 fc_input_size = self.config.target_hidden_size * 3

177 else:

178 fc_input_size = self.config.hidden_size * 3

179 + if self.norm_before_fc:

180 + self.input_norm = RMSNorm(

181 + fc_input_size,

182 + eps=self.config.rms_norm_eps,

183 + )

184 + else:

185 + self.input_norm = None

186 self.fc = ReplicatedLinear(

187 input_size=fc_input_size,

188 output_size=self.config.hidden_size,

@@ -357,6 +365,9 @@ class Eagle3LlamaForCausalLM(LlamaForCausalLM):

365 if not self.model.use_aux_hidden_state:

366 return hidden_states

367 # combine multiple auxiliary hidden states returned by eagle3

368 +

369 + if self.model.norm_before_fc:

370 + hidden_states = self.model.input_norm(hidden_states)

371 return self.model.fc(hidden_states)

372

373 def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]):

@@ -403,6 +414,8 @@ class Eagle3LlamaForCausalLM(LlamaForCausalLM):

414 skip_substrs.append("embed_tokens")

415 if not self.model.use_aux_hidden_state:

416 skip_substrs.append("fc.")

417 + if not self.model.norm_before_fc:

418 + skip_substrs.append("input_norm.")

419 loader = AutoWeightsLoader(

420 self,

421 skip_prefixes=None,