[Speculative decoding] CUDA graph support

[heeju-kim2]

We add a end-to-end test that verifies speculative decoding works with CUDA graphs. Specifically, we ensure that a model produces the same output when speculative decoding is enabled vs not, when cuda graphs are enabled.

[cadedaniel]

Thanks for the PR Heeju!

I will take a look today.

[mgoin]

These look like mistaken changes/invalid syntax - please restore this file

[cadedaniel]

Thanks again @heeju-kim2 for the PR.

So I ran spec decode locally with cuda graphs enabled, and I find that it works. Are you misconfiguring spec decode?

You can see what I saw by running the test below, which uses run_greedy_equality_correctness_test from test_correctness.py.

By the way, if this works for you, you can update this PR to simply include this test, remove the other changes, and we can merge it. Thanks for contributing to vLLM :)

@pytest.mark.parametrize(
    "common_llm_kwargs",
    [{  
        # Required for spec decode.
        "use_v2_block_manager": True,
    }]) 
@pytest.mark.parametrize("per_test_common_llm_kwargs", [
{
    # Identical models.
    "model": "JackFram/llama-68m",
    "speculative_model": "JackFram/llama-68m",
    "num_speculative_tokens": 5,
},
{
    # Distinct models.
    "model": "JackFram/llama-160m",
    "speculative_model": "JackFram/llama-68m",
    "num_speculative_tokens": 5,
}
])
@pytest.mark.parametrize("baseline_llm_kwargs", [{}])
@pytest.mark.parametrize("test_llm_kwargs", [{
    "enforce_eager": False,
}])
@pytest.mark.parametrize("batch_size", [1, 32])
@pytest.mark.parametrize("output_len", [128])
@pytest.mark.parametrize("seed", [1])
def test_spec_decode_cuda_graph(baseline_llm_generator, test_llm_generator, batch_size, output_len):
    run_greedy_equality_correctness_test(
        baseline_llm_generator,
        test_llm_generator,
        batch_size,
        max_output_len=output_len,
        force_output_len=True,
    )

[cadedaniel]

Thanks for the PR @heeju-kim2 !