biondizzle/vllm
1
0
Fork 0
Code Issues Pull Requests Actions 2 Packages Projects Releases Wiki Activity

[V1] Logits processor docs (#22919)

Browse Source 
Signed-off-by: Andrew Feldman <afeldman@redhat.com>
Signed-off-by: afeldman-nm <156691304+afeldman-nm@users.noreply.github.com>
Co-authored-by: Joseph Marinier <Joseph.Marinier@gmail.com>
This commit is contained in:
afeldman-nm
2025-09-17 14:53:12 -04:00
committed by GitHub
parent e3db5ebb66
commit 7ae9887542
7 changed files with 1065 additions and 16 deletions
Download Patch File Download Diff File Expand all files Collapse all files

10
examples/offline_inference/logits_processor/custom.py
View File

@@ -56,7 +56,6 @@ class DummyLogitsProcessor(LogitsProcessor):
self.req_info: dict[int, int] = {}
def is_argmax_invariant(self) -> bool:
"""Never impacts greedy sampling"""
return False
def update_state(self, batch_update: Optional[BatchUpdate]):
@@ -75,13 +74,12 @@ class DummyLogitsProcessor(LogitsProcessor):
return logits
# Save target values before modification
rows_list = list(self.req_info.keys())
cols = torch.tensor(
[self.req_info[i] for i in rows_list],
dtype=torch.long,
device=logits.device,
list(self.req_info.values()), dtype=torch.long, device=logits.device
)
rows = torch.tensor(
list(self.req_info.keys()), dtype=torch.long, device=logits.device
)
rows = torch.tensor(rows_list, dtype=torch.long, device=logits.device)
values_to_keep = logits[rows, cols].clone()
# Mask all but target tokens