[Bugfix] Fix cuda graph sizes when running with speculative decoding

[PatrykSaffer]

Purpose

When specifying max_seq_len and running with num_speculative_tokens > 1 not all cuda graph sizes are being captured.

Eg.
max_seq_len == 1, num_speculative_tokens==3
Only graph for batch_size==2 is being captured.

Test Plan

Manual test

Test Result

Tested manually

[chatgpt-codex-connector]

Codex usage limits have been reached for code reviews. Please check with the admins of this repo to increase the limits by adding credits.

[gemini-code-assist]

Code Review

The pull request addresses a bug where CUDA graph sizes were not correctly captured when speculative decoding was enabled, particularly when max_seq_len was small and num_speculative_tokens was greater than 1. The fix correctly incorporates the number of speculative tokens into the calculation of max_cudagraph_capture_size, ensuring that the system captures appropriate graph sizes for speculative decoding scenarios. This improves the correctness and efficiency of CUDA graph utilization under these conditions.

[mergify]

Hi @PatrykSaffer, the pre-commit checks have failed. Please run:

uv pip install pre-commit
pre-commit install
pre-commit run --all-files

Then, commit the changes and push to your branch.

For future commits, pre-commit will run automatically on changed files before each commit.

Tip

Is mypy or markdownlint failing?

mypy and markdownlint are run differently in CI. If the failure is related to either of these checks, please use the following commands to run them locally:

# For mypy (substitute "3.10" with the failing version if needed)
pre-commit run --hook-stage manual mypy-3.10
# For markdownlint
pre-commit run --hook-stage manual markdownlint