[v1] Support allowed_token_ids in v1 Sampler (#13210)

Signed-off-by: Lu Fang <lufang@fb.com>
2025-02-21 22:13:05 -08:00
parent 8aca27fa11
commit bb78fb318e
7 changed files with 168 additions and 19 deletions
--- a/tests/v1/sample/test_rejection_sampler.py
+++ b/tests/v1/sample/test_rejection_sampler.py
@@ -43,6 +43,7 @@ def create_sampling_metadata(spec_tokens: List[List[int]]) -> SamplingMetadata:
        output_token_ids=[],
        min_tokens={},
        logit_bias=[None] * batch_size,
+        allowed_token_ids_mask=None,
    )