fix unexcepted answer in EAGLE mode (#9252)

This commit is contained in:
zyksir
2025-08-17 08:45:36 +08:00
committed by GitHub
parent 94371dbbd6
commit 6a9d6ca33c
2 changed files with 33 additions and 5 deletions

View File

@@ -836,6 +836,21 @@ class EAGLEWorker(TpModelWorker):
assert isinstance(forward_batch.spec_info, EagleDraftInput)
assert forward_batch.spec_info is batch.spec_info
self.capture_for_decode(logits_output, forward_batch.spec_info)
has_finished, unfinished_req_index = False, []
for i, req in enumerate(batch.reqs):
if req.finished():
has_finished = True
else:
unfinished_req_index.append(i)
if has_finished:
unfinished_index_device = torch.tensor(
unfinished_req_index,
dtype=torch.int64,
device=batch.spec_info.topk_p.device,
)
batch.spec_info.filter_batch(
unfinished_index_device, has_been_filtered=False
)
def forward_draft_extend_after_decode(self, batch: ScheduleBatch):
assert isinstance(batch.spec_info, EagleDraftInput)