[1/n][Chunked Prefill] Refactor input query shapes (#3236)

2024-03-21 06:46:05 +09:00
parent 426ec4ec67
commit 6e435de766
18 changed files with 579 additions and 263 deletions
--- a/tests/spec_decode/test_multi_step_worker.py
+++ b/tests/spec_decode/test_multi_step_worker.py
@@ -92,8 +92,8 @@ def test_same_output_for_single_step():
        num_gpu_blocks,
        seed,
    )
-    multi_step_worker.model_runner = worker.model_runner
-    multi_step_worker.cache_engine = worker.cache_engine
+    # multi_step_worker.model_runner = worker.model_runner
+    # multi_step_worker.cache_engine = worker.cache_engine

    num_steps = 1