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

[Core][Optimization] change python dict to pytorch tensor (#4607)

Browse Source
This commit is contained in:
youkaichao
2024-05-06 21:30:27 -07:00
committed by GitHub
parent a98187cf72
commit 63575bc2e1
19 changed files with 77 additions and 81 deletions
Download Patch File Download Diff File Expand all files Collapse all files

8
tests/core/test_scheduler.py
View File

@@ -568,7 +568,7 @@ def test_decode_schedule_preempted():
# Both should be preempted, not swapped.
assert output.blocks_to_swap_out == {}
# Nothing is copied.
assert output.blocks_to_copy == {}
assert output.blocks_to_copy == []
def test_decode_swap_beam_search():
@@ -618,7 +618,7 @@ def test_decode_swap_beam_search():
# Both should be preempted, not swapped.
assert output.blocks_to_swap_out == expected_swap_mapping
# Nothing is copied.
assert output.blocks_to_copy == {}
assert output.blocks_to_copy == []
def test_schedule_decode_blocks_to_copy_update():
@@ -650,7 +650,7 @@ def test_schedule_decode_blocks_to_copy_update():
assert output.blocks_to_swap_out == {}
# Since append_slot returns the source -> dist mapping, it should
# applied.
assert output.blocks_to_copy == {2: [3]}
assert output.blocks_to_copy == [(2, 3)]
def test_schedule_swapped_simple():
@@ -853,7 +853,7 @@ def test_schedule_swapped_blocks_to_copy():
assert len(remaining_swapped) == 0
assert len(output.decode_seq_groups) == 1
assert len(output.prefill_seq_groups) == 0
assert output.blocks_to_copy == {2: [3]}
assert output.blocks_to_copy == [(2, 3)]
def test_scheduling_budget():