Skip to content

Commit

Permalink
Fix failing test
Browse files Browse the repository at this point in the history
Signed-off-by: Sourashis Roy <[email protected]>
  • Loading branch information
sroy745 authored and andoorve committed Nov 25, 2024
1 parent e875afa commit a8d2210
Showing 1 changed file with 2 additions and 1 deletion.
3 changes: 2 additions & 1 deletion tests/spec_decode/test_spec_decode_worker.py
Original file line number Diff line number Diff line change
Expand Up @@ -867,7 +867,8 @@ def test_chunked_prefill_flow(k: int, batch_size: int, batch_composition: str):
target_group_metadata_list = prefill + decodes
execute_model_req = ExecuteModelRequest(
seq_group_metadata_list=target_group_metadata_list,
num_lookahead_slots=k)
# For prefill only batches we expect num_lookahead_slots = 0.
num_lookahead_slots=k if n_decodes > 0 else 0)

target_token_ids = torch.randint(low=0,
high=vocab_size,
Expand Down

0 comments on commit a8d2210

Please sign in to comment.