debug

yeyu-nvidia · yeyu-nvidia · commit b0edb38c7750 · 2025-09-26T14:18:01.000-07:00
Signed-off-by: Ye Yu &lt;yeyu@nvidia.com&gt;
diff --git a/modelopt/torch/speculative/plugins/megatron_eagle.py b/modelopt/torch/speculative/plugins/megatron_eagle.py
@@ -945,7 +945,8 @@ def _eagle_forward(
         )
 
         # Update inference_context.sequence_len_offset after each call of eagle_module
-        inference_context.sequence_len_offset += eagle_inputs["input_ids"].shape[1]
+        if inference_context is not None:
+            inference_context.sequence_len_offset += eagle_inputs["input_ids"].shape[1]
 
         if hasattr(self.eagle_module, "eagle_output_layer"):
             eagle_logits, _ = self.eagle_module.eagle_output_layer(eagle_hidden_states)

Original file line number	Diff line number	Diff line change
`@@ -945,7 +945,8 @@ def _eagle_forward(`
`945`	`945`	`)`
`946`	`946`
`947`	`947`	`# Update inference_context.sequence_len_offset after each call of eagle_module`
`948`		`- inference_context.sequence_len_offset += eagle_inputs["input_ids"].shape[1]`
	`948`	`+ if inference_context is not None:`
	`949`	`+ inference_context.sequence_len_offset += eagle_inputs["input_ids"].shape[1]`
`949`	`950`
`950`	`951`	`if hasattr(self.eagle_module, "eagle_output_layer"):`
`951`	`952`	`eagle_logits, _ = self.eagle_module.eagle_output_layer(eagle_hidden_states)`