[Eagle2] Fix multiple concurrent request crashes (#2730)

python/sglang/srt/speculative/eagle_utils.py

@@ -245,9 +245,10 @@ class EAGLEDraftInput(SpecInfo):
             )  # (b, topk)
             topk_cs_index, topk_cs_p = topk_cs.indices, topk_cs.values
 
-            selected_input_index = (
-                topk_cs_index.flatten() // self.topk
-            )  # shape: (b * topk)
+            selected_input_index = topk_cs_index.flatten() // self.topk + torch.arange(
+                0, batch.batch_size() * self.topk, step=self.topk, device="cuda"
+            ).repeat_interleave(self.topk)
+
             batch.spec_info.hidden_states = batch.spec_info.hidden_states[
                 selected_input_index, :
             ]
@@ -336,6 +337,7 @@ class EAGLEDraftInput(SpecInfo):
             triton.next_power_of_2(self.spec_steps + 1),
         )
 
+        batch.seq_lens_sum = sum(batch.seq_lens)
         batch.input_ids = self.verified_id
         self.verified_id = new_verified_id
 
@@ -439,7 +441,14 @@ class EAGLEDraftInput(SpecInfo):
         return kv_indices, cum_kv_seq_len, qo_indptr, None
 
     def merge_batch(self, spec_info: EAGLEDraftInput):
-
+        if self.hidden_states is None:
+            self.hidden_states = spec_info.hidden_states
+            self.verified_id = spec_info.verified_id
+            self.sample_output = spec_info.sample_output
+            self.prev_mode = spec_info.prev_mode
+            return
+        if spec_info.hidden_states is None:
+            return
         self.hidden_states = torch.cat(
             [self.hidden_states, spec_info.hidden_states], axis=0
         )

python/sglang/srt/speculative/eagle_worker.py

@@ -169,6 +169,8 @@ class EAGLEWorker(TpModelWorker):
         if not isinstance(reqs, List):
             reqs = [reqs]
         for req in reqs:
+            if req.rid not in self.finish_extend_len:
+                continue
             req_len = (
                 len(req.origin_input_ids)
                 + len(req.output_ids)