Eagle speculative decoding part 2: Fix cuda graph + DP attention hanging (#2684)

Co-authored-by: yukavio <kavioyu@gmail.com>
2024-12-31 02:25:05 -08:00
parent 6c42fa229d
commit b0524c3789
7 changed files with 131 additions and 58 deletions
@@ -150,12 +150,18 @@ class TpModelWorker:
        self,
        model_worker_batch: ModelWorkerBatch,
        launch_done: Optional[threading.Event] = None,
+        skip_sample: bool = False,
    ):
        forward_batch = ForwardBatch.init_new(model_worker_batch, self.model_runner)
        logits_output = self.model_runner.forward(forward_batch)
        if launch_done:
            launch_done.set()
-        next_token_ids = self.model_runner.sample(logits_output, model_worker_batch)
+
+        if skip_sample:
+            next_token_ids = None
+        else:
+            next_token_ids = self.model_runner.sample(logits_output, model_worker_batch)
+
        return logits_output, next_token_ids

    def forward_batch_embedding(self, model_worker_batch: ModelWorkerBatch):