[BugFix] Fix spec decode + structured outputs + preemption edge case (#30916)

Signed-off-by: Nick Hill <nhill@redhat.com>
This commit is contained in:
Nick Hill 2025-12-18 12:59:55 -08:00 committed by GitHub
parent 634a14bd7d
commit b0b77c4655
No known key found for this signature in database
GPG Key ID: B5690EEEBB952194

View File

@ -3393,9 +3393,13 @@ class GPUModelRunner(
return async_output
def take_draft_token_ids(self) -> DraftTokenIds | None:
if self._draft_token_ids is None:
if not self.num_spec_tokens:
return None
req_ids = self.input_batch.req_ids
if self._draft_token_ids is None:
return DraftTokenIds(req_ids, [[] for _ in req_ids])
if isinstance(self._draft_token_ids, torch.Tensor):
draft_token_ids = self._draft_token_ids.tolist()
else: