remove hard code for block_size

This commit is contained in:
guodongxiaren
2025-12-30 01:53:55 +08:00
parent 2f21442653
commit 55c64e7fdf
2 changed files with 6 additions and 1 deletions
+1
View File
@@ -18,6 +18,7 @@ class LLMEngine:
config_fields = {field.name for field in fields(Config)} config_fields = {field.name for field in fields(Config)}
config_kwargs = {k: v for k, v in kwargs.items() if k in config_fields} config_kwargs = {k: v for k, v in kwargs.items() if k in config_fields}
config = Config(model, **config_kwargs) config = Config(model, **config_kwargs)
Sequence.set_block_size(config.kvcache_block_size)
self.ps = [] self.ps = []
self.events = [] self.events = []
ctx = mp.get_context("spawn") ctx = mp.get_context("spawn")
+5 -1
View File
@@ -12,9 +12,13 @@ class SequenceStatus(Enum):
class Sequence: class Sequence:
block_size = 256 block_size: int = 0 # invalid value, will be set by set_block_size
counter = count() counter = count()
@classmethod
def set_block_size(cls, block_size: int):
cls.block_size = block_size
def __init__(self, token_ids: list[int], sampling_params = SamplingParams()): def __init__(self, token_ids: list[int], sampling_params = SamplingParams()):
self.seq_id = next(Sequence.counter) self.seq_id = next(Sequence.counter)
self.status = SequenceStatus.WAITING self.status = SequenceStatus.WAITING