[Core] Factor out common code in SequenceData and Sequence (#8675)

This commit is contained in:
Cyrus Leung
2024-09-21 10:30:39 +08:00
committed by GitHub
parent d4bf085ad0
commit 0455c46ed4
8 changed files with 64 additions and 97 deletions

View File

@@ -1,10 +1,7 @@
from array import array
import pytest
from vllm.model_executor.layers.sampler import SamplerOutput
from vllm.sequence import (VLLM_TOKEN_ID_ARRAY_TYPE,
CompletionSequenceGroupOutput, SequenceData,
from vllm.sequence import (CompletionSequenceGroupOutput, SequenceData,
SequenceOutput)
from .core.utils import create_dummy_prompt
@@ -58,7 +55,7 @@ def test_sampler_output_eq(sample_outputs):
def test_sequence_data_prefill():
seq_data = SequenceData(array(VLLM_TOKEN_ID_ARRAY_TYPE, [1, 2, 3, 4]))
seq_data = SequenceData.from_seqs([1, 2, 3, 4])
assert seq_data.get_num_uncomputed_tokens() == 4
assert seq_data.get_num_computed_tokens() == 0
# advance by 2