[ModelRunner V2] Speculative Decoding NGram GPU Implementations

.buildkite/test_areas/model_runner_v2.yaml

@@ -132,13 +132,15 @@ steps:
   - tests/v1/spec_decode/test_max_len.py
   - tests/v1/spec_decode/test_rejection_sampler_utils.py
   - tests/v1/spec_decode/test_synthetic_rejection_sampler_utils.py
+  - tests/v1/spec_decode/test_ngram_gpu.py
   - tests/v1/e2e/spec_decode/
   commands:
   - set -x
   - export VLLM_USE_V2_MODEL_RUNNER=1
-  - pytest -v -s v1/spec_decode/test_max_len.py -k "eagle or mtp"
+  - pytest -v -s v1/spec_decode/test_max_len.py -k "eagle or mtp or ngram_gpu"
   - pytest -v -s v1/spec_decode/test_rejection_sampler_utils.py
   - pytest -v -s v1/spec_decode/test_synthetic_rejection_sampler_utils.py
+  - pytest -v -s v1/spec_decode/test_ngram_gpu.py
   - pytest -v -s v1/e2e/spec_decode/eagle/
   - pytest -v -s v1/e2e/spec_decode/speculators/
   - >-

tests/test_config.py

@@ -293,6 +293,38 @@ def test_resolve_cudagraph_mode_adjusts_spec_decode_sizes_only_for_v1(
     assert compilation_config.cudagraph_capture_sizes == expected_capture_sizes
 
 
+@pytest.mark.parametrize(
+    ("requested", "expected"),
+    [
+        # Only a mode without a separate decode routine has to change: varlen
+        # decode batches would otherwise replay on a mixed full graph.
+        (CUDAGraphMode.FULL, CUDAGraphMode.FULL_AND_PIECEWISE),
+        (CUDAGraphMode.FULL_AND_PIECEWISE, CUDAGraphMode.FULL_AND_PIECEWISE),
+        (CUDAGraphMode.FULL_DECODE_ONLY, CUDAGraphMode.FULL_DECODE_ONLY),
+        (CUDAGraphMode.PIECEWISE, CUDAGraphMode.PIECEWISE),
+        (CUDAGraphMode.NONE, CUDAGraphMode.NONE),
+    ],
+)
+def test_resolve_cudagraph_mode_varlen_decode(requested, expected):
+    """varlen_decode requires a separate decode routine for full cudagraphs."""
+    compilation_config = CompilationConfig(
+        cudagraph_mode=requested,
+        cudagraph_capture_sizes=[1, 2, 4, 8],
+    )
+    compilation_config.max_cudagraph_capture_size = 8
+    compilation_config.post_init_cudagraph_sizes()
+
+    cudagraph_mode = compilation_config.resolve_cudagraph_mode_and_sizes(
+        AttentionCGSupport.ALWAYS,
+        "FakeAttentionBackend",
+        use_v2_model_runner=True,
+        varlen_decode=True,
+    )
+
+    assert cudagraph_mode == expected
+    assert compilation_config.cudagraph_mode == expected
+
+
 @pytest.mark.parametrize(
     ("model_config", "expected"),
     [

tests/v1/spec_decode/test_max_len.py

@@ -36,6 +36,33 @@ def test_ngram_max_len(num_speculative_tokens: int, vllm_runner):
         runner.llm.generate(_PROMPTS, sampling_params)
 
 
+@pytest.mark.parametrize("num_speculative_tokens", [1, 3, 10])
+@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 (where "ngram" and "ngram_gpu" both
+    resolve to the GPU implementation) correctly
+    handles the ``max_model_len`` boundary across various speculative-token
+    counts.
+    """
+    with vllm_runner(
+        "facebook/opt-125m",
+        trust_remote_code=False,
+        max_model_len=100,
+        enable_chunked_prefill=None,
+        enforce_eager=True,  # For faster initialization.
+        speculative_config={
+            "method": method,
+            "prompt_lookup_max": 5,
+            "prompt_lookup_min": 3,
+            "num_speculative_tokens": num_speculative_tokens,
+        },
+    ) as runner:
+        sampling_params = SamplingParams(max_tokens=100, ignore_eos=True)
+        runner.llm.generate(_PROMPTS, sampling_params)
+
+
 @pytest.mark.parametrize("num_speculative_tokens", [1, 3, 10])
 @pytest.mark.parametrize("attn_backend", get_attn_backend_list_based_on_platform())
 def test_eagle_max_len(