[Bugfix][SpecDecode] kv corruption with bonus tokens in spec decode (#9730)

Co-authored-by: LiuXiaoxuanPKU <lilyliupku@gmail.com>
This commit is contained in:
Sungjae Lee
2024-11-06 10:45:45 +09:00
committed by GitHub
parent 966e31697b
commit 0c63c34f72
4 changed files with 159 additions and 10 deletions

View File

@@ -68,12 +68,14 @@ def create_worker(cls: Callable[..., T],
seed: int,
is_driver_worker: bool = True,
enforce_eager: bool = True,
model_runner_cls: Optional[ModelRunner] = None) -> T:
model_runner_cls: Optional[ModelRunner] = None,
dtype: Optional[str] = "auto") -> T:
engine_args = EngineArgs(
model=model_name,
seed=seed,
block_size=block_size,
enforce_eager=enforce_eager,
dtype=dtype,
)
engine_config = engine_args.create_engine_config()