Fix

Signed-off-by: DarkLight1337 <[email protected]>
vllm-project · Dec 7, 2024 · 5c4fc15 · 5c4fc15
1 parent 68fc916
commit 5c4fc15
Showing 1 changed file with 8 additions and 1 deletion.
diff --git a/tests/distributed/test_pipeline_parallel.py b/tests/distributed/test_pipeline_parallel.py
@@ -28,6 +28,7 @@ class ParallelSetup(NamedTuple):
     chunked_prefill: bool
     speculative_model: Optional[str] = None
     num_speculative_tokens: Optional[int] = None
+    ngram_prompt_lookup_max: Optional[int] = None
 
 
 class PPTestOptions(NamedTuple):
@@ -84,7 +85,8 @@ def detailed(
                               eager_mode=False,
                               chunked_prefill=False,
                               speculative_model="[ngram]",
-                              num_speculative_tokens=5),
+                              num_speculative_tokens=5,
+                              ngram_prompt_lookup_max=3),
             ],
             distributed_backends=["mp", "ray"],
             task=task,
@@ -262,6 +264,7 @@ def _compare_tp(
         chunked_prefill,
         speculative_model,
         num_speculative_tokens,
+        ngram_prompt_lookup_max,
     ) = parallel_setup
     (
         multi_node_only,
@@ -308,6 +311,10 @@ def _compare_tp(
         common_args.extend(
             ["--num-speculative-tokens",
              str(num_speculative_tokens)])
+    if ngram_prompt_lookup_max:
+        common_args.extend(
+            ["--ngram-prompt-lookup-max",
+             str(ngram_prompt_lookup_max)])
 
     if (distributed_backend == "ray" and tp_size == 2 and pp_size == 2
             and chunked_prefill):