[Spec Decode] Resolve both ngram methods to the GPU implementation on the V2 runner

On the V2 model runner, method="ngram" and method="ngram_gpu" now both use
NgramGPUSpeculator (via a new SpeculativeConfig.use_ngram() helper); the V1
runner keeps its separate CPU and GPU proposers. Also scope the
torch.compile cache disable to the V1 ngram-gpu proposer: it exists for
V1's @support_torch_compile kernel, while the V2 implementation is pure
Triton and does not need it.

Co-authored-by: Claude <noreply@anthropic.com>

Signed-off-by: Nick Hill <nickhill123@gmail.com>

Author: njhill

tests/v1/spec_decode/test_max_len.py

@@ -37,10 +37,12 @@ def test_ngram_max_len(num_speculative_tokens: int, vllm_runner):
 
 
 @pytest.mark.parametrize("num_speculative_tokens", [1, 3, 10])
-def test_ngram_gpu_max_len(num_speculative_tokens: int, vllm_runner):
+@pytest.mark.parametrize("method", ["ngram", "ngram_gpu"])
+def test_ngram_gpu_max_len(method: str, num_speculative_tokens: int, vllm_runner):
     """V2 GPU n-gram counterpart of ``test_ngram_max_len``.
 
-    Verifies that the V2 model runner with ``method="ngram_gpu"`` correctly
+    Verifies that the V2 model runner (where "ngram" and "ngram_gpu" both
+    resolve to the GPU implementation) correctly
     handles the ``max_model_len`` boundary across various speculative-token
     counts.
     """
@@ -51,7 +53,7 @@ def test_ngram_gpu_max_len(num_speculative_tokens: int, vllm_runner):
         enable_chunked_prefill=None,
         enforce_eager=True,  # For faster initialization.
         speculative_config={
-            "method": "ngram_gpu",
+            "method": method,
             "prompt_lookup_max": 5,
             "prompt_lookup_min": 3,
             "num_speculative_tokens": num_speculative_tokens,

tests/v1/spec_decode/test_ngram_gpu.py

@@ -49,6 +49,7 @@ def _make_vllm_config(
     k: int,
     max_num_seqs: int = 8,
     max_model_len: int = 64,
+    method: str = "ngram_gpu",
 ) -> VllmConfig:
     model_config = ModelConfig(
         model="facebook/opt-125m",
@@ -60,7 +61,7 @@ def _make_vllm_config(
         max_model_len=max_model_len,
     )
     speculative_config = SpeculativeConfig(
-        method="ngram_gpu",
+        method=method,
         prompt_lookup_min=min_n,
         prompt_lookup_max=max_n,
         num_speculative_tokens=k,
@@ -360,14 +361,23 @@ def test_propose_requires_req_states():
         )
 
 
+@pytest.mark.parametrize("method", ["ngram", "ngram_gpu"])
+def test_both_ngram_methods_resolve_to_gpu_speculator(method: str):
+    """On the V2 runner, "ngram" and "ngram_gpu" use the same implementation."""
+    from vllm.v1.worker.gpu.spec_decode import init_speculator
+
+    cfg = _make_vllm_config(min_n=2, max_n=3, k=2, method=method)
+    spec = init_speculator(cfg, DEVICE)
+    assert isinstance(spec, NgramGPUSpeculator)
+
+
 def test_construction_validates_speculative_config():
     spec = _make_speculator(min_n=2, max_n=3, k=2)
     assert spec.min_n == 2
     assert spec.max_n == 3
     assert spec.num_speculative_steps == 2
-    # No-op hooks must not raise.
-    spec.load_model(target_model=None)
-    spec.set_attn()
+    # Inherited no-op hooks must not raise.
+    spec.init_cudagraph_manager(None)
     spec.capture()
 
 

vllm/compilation/backends.py

@@ -1081,10 +1081,13 @@ def __call__(self, graph: fx.GraphModule, example_inputs: Sequence[Any]) -> Any:
         # Honors opt-outs such as CompilationMode.NONE or VLLM_DISABLE_COMPILE_CACHE.
         disable_cache = not is_compile_cache_enabled(self.inductor_config)
 
-        # TODO(patchy): ngram gpu kernel will cause vllm torch compile cache errors.
+        # TODO(patchy): the V1 torch.compile ngram-gpu kernel causes vllm
+        # torch compile cache errors. The V2 implementation is pure Triton and
+        # does not need the cache disabled.
         is_ngram_gpu_enabled = (
             vllm_config.speculative_config is not None
             and vllm_config.speculative_config.use_ngram_gpu()
+            and not vllm_config.use_v2_model_runner
         )
         disable_cache = disable_cache or is_ngram_gpu_enabled
 

vllm/config/speculative.py

@@ -1498,6 +1498,9 @@ def uses_extract_hidden_states(self) -> bool:
     def use_ngram_gpu(self) -> bool:
         return self.method == "ngram_gpu"
 
+    def use_ngram(self) -> bool:
+        return self.method in ("ngram", "ngram_gpu")
+
     def use_multi_module_mtp(self) -> bool:
         if self.method != "mtp" or self.draft_model_config is None:
             return False

vllm/config/vllm.py

@@ -2395,15 +2395,15 @@ def _get_v2_model_runner_unsupported_features(self) -> list[str]:
             unsupported.append("pipeline parallelism with external_launcher")
 
         if speculative_config is not None:
-            # CPU ngram is not supported by the V2 model runner yet.
-            if speculative_config.method == "ngram":
-                unsupported.append("ngram speculative decoding")
-            elif speculative_config.method not in (
+            # Both ngram methods resolve to the same GPU implementation on
+            # the V2 model runner.
+            if speculative_config.method not in (
                 "eagle",
                 "eagle3",
                 "mtp",
                 "dflash",
                 "dspark",
+                "ngram",
                 "ngram_gpu",
             ):
                 unsupported.append(f"speculative method '{speculative_config.method}'")

vllm/v1/worker/gpu/model_runner.py

@@ -271,8 +271,7 @@ def __init__(self, vllm_config: VllmConfig, device: torch.device):
 
         # General request states.
         use_dense_all_token_ids = (
-            self.speculative_config is not None
-            and self.speculative_config.use_ngram_gpu()
+            self.speculative_config is not None and self.speculative_config.use_ngram()
         )
         self.req_states = RequestState(
             max_num_reqs=self.max_num_reqs,

vllm/v1/worker/gpu/spec_decode/__init__.py

@@ -42,7 +42,7 @@ def init_speculator(vllm_config: VllmConfig, device: torch.device):
         )
 
         return EagleSpeculator(vllm_config, device)
-    elif speculative_config.use_ngram_gpu():
+    elif speculative_config.use_ngram():
         from vllm.v1.worker.gpu.spec_decode.ngram.speculator import (
             NgramGPUSpeculator,
         )