Skip to content

Commit 6b533d5

Browse files
fix: DRY sampler never applies during decoding (stale token history views)
Co-Authored-By: Claude Opus 5 <noreply@anthropic.com>
1 parent 851494a commit 6b533d5

1 file changed

Lines changed: 7 additions & 10 deletions

File tree

aphrodite/v1/worker/gpu_input_batch.py

Lines changed: 7 additions & 10 deletions
Original file line numberDiff line numberDiff line change
@@ -1261,9 +1261,13 @@ def _make_sampling_metadata(self) -> SamplingMetadata:
12611261
output_token_ids = cast(list[list[int]], self.req_output_token_ids) if needs_output_token_ids else []
12621262
output_token_ids_tensor = None
12631263
token_history_ids, token_history_lens = None, None
1264-
token_history_ids_cpu, token_history_lens_cpu = (
1265-
self._get_token_history_cpu_views(num_reqs) if not self.no_dry else (None, None)
1266-
)
1264+
# The DRY CPU kernel cannot be fed from token_ids_cpu here: cached views
1265+
# freeze at metadata-build width and hide every generated token, while
1266+
# live full-width views expose async-scheduling placeholder ids (-1) at
1267+
# sampling time. Passing None makes Sampler.apply_dry fall back to the
1268+
# persistent state, which update_dry_state keeps in sync with the real
1269+
# sampled ids each step.
1270+
token_history_ids_cpu, token_history_lens_cpu = (None, None)
12671271

12681272
allowed_token_ids_mask: torch.Tensor | None = None
12691273
if not self.no_allowed_token_ids:
@@ -1414,13 +1418,6 @@ def _make_output_token_ids_tensor(self, num_reqs: int) -> torch.Tensor:
14141418
)
14151419
return output_token_ids_cpu_tensor.to(device=self.device, non_blocking=True)
14161420

1417-
def _get_token_history_cpu_views(self, num_reqs: int) -> tuple[torch.Tensor, torch.Tensor]:
1418-
max_history_len = int(self.num_tokens_no_spec[:num_reqs].max()) if num_reqs else 0
1419-
return (
1420-
self.token_ids_cpu_tensor[:num_reqs, :max_history_len],
1421-
self.num_tokens_no_spec_cpu_tensor[:num_reqs],
1422-
)
1423-
14241421
def _make_token_history_tensors(self, num_reqs: int) -> tuple[torch.Tensor, torch.Tensor]:
14251422
history_lens = torch.as_tensor(
14261423
self.num_tokens_no_spec[:num_reqs],

0 commit comments

Comments
 (0)