include all state updates

Signed-off-by: Leo Tian <[email protected]>

Author: leo-cf-tian

vllm/v1/spec_decode/eagle.py

@@ -171,12 +171,6 @@ def propose(
                                            hidden_states, attn_metadata,
                                            batch_size)
 
-            # Increment the sequence lengths.
-            attn_metadata.max_seq_len += 1
-            # Consider max model length.
-            attn_metadata.max_seq_len = min(attn_metadata.max_seq_len,
-                                            self.max_model_len)
-
             # copy inputs to buffer for cudagraph
             # Run the model.
             with set_forward_context(attn_metadata,
@@ -231,6 +225,12 @@ def advance_speculative_state(self, draft_token_ids: torch.Tensor,
             BLOCK_SIZE=1024,
             PADDING_SLOT_ID=PADDING_SLOT_ID)
 
+        # Increment the sequence lengths.
+        attn_metadata.max_seq_len += 1
+        # Consider max model length.
+        attn_metadata.max_seq_len = min(attn_metadata.max_seq_len,
+                                        self.max_model_len)
+
     @staticmethod
     def prepare_inputs(
         # [batch_size + 1]