[ROCm][CI] Fix spec decode profile assertion and logprob test determinism

[AndreasKaratzas]

Fixes two issues blocking spec decode logprob tests on ROCm:

1. Profile run assertion failure (gpu_model_runner.py): The _dummy_run method asserted num_tokens <= self.scheduler_config.max_num_batched_tokens, but with speculative decoding max_num_tokens (which accounts for verification tokens) can exceed max_num_batched_tokens. Updated the assertion to use self.max_num_tokens, consistent with the rest of the runner.

2. Non-deterministic logprob comparison (test_logprobs.py): The ref LLM and spec-decode LLM used different batch sizes, which on ROCm triggers non-associative floating-point reduction differences in attention/GEMM kernels. These numerical divergences were misattributed to spec decode incorrectness. Added ROCM_DETERMINISM_KWARGS (max_num_seqs=1) applied to both LLM instances on ROCm only, pinning identical execution paths. No behavioral change on other platforms.

Test Plan

pytest -s -v tests/v1/sample/test_logprobs.py

[dosubot]

Related Documentation

Checked 0 published document(s) in 1 knowledge base(s). No updates required.

^{How did I do? Any feedback?}  

[mergify]

Hi @AndreasKaratzas, 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

[gemini-code-assist]

Code Review

The pull request addresses two issues related to speculative decoding on ROCm: an assertion failure in gpu_model_runner.py and non-deterministic logprob comparisons in test_logprobs.py. The changes correctly update the assertion to use self.max_num_tokens and introduce ROCM_DETERMINISM_KWARGS to ensure deterministic execution for logprob tests on ROCm. The changes are well-explained and directly address the identified problems.

[gemini-code-assist]

The ROCM_DETERMINISM_KWARGS dictionary currently only sets max_num_seqs=1. The PR description mentions enforce_eager and async_scheduling=False as part of the determinism kwargs. These should also be included in the dictionary to fully align with the described fix and ensure consistent execution paths on ROCm.

Suggested change

	ROCM_DETERMINISM_KWARGS: dict = (
	dict(
	max_num_seqs=1,
	)
	if current_platform.is_rocm()
	else {}
	)
	ROCM_DETERMINISM_KWARGS: dict = (
	dict(
	enforce_eager=True,
	async_scheduling=False,
	max_num_seqs=1,
	)
	if current_platform.is_rocm()
	else {}
	)

[AndreasKaratzas]

I've updated the description already, apparently those args were unnecessary.

[AndreasKaratzas]

This PR depends on:

• [BugFix] Refactor add max_num_tokens_per_forward_pass to account for drafting #35038

That's why pre-commit is failing.

EDIT: This is no longer true. I have reverted the change in vllm/v1/worker/gpu_model_runner.py in order to merge this PR quickly.

[AndreasKaratzas]

cc @LucasWilkinson

[AndreasKaratzas]

I'm going to revert the change in vllm/v1/worker/gpu_model_runner.py since this PR is critical for AMD CI.

[tjtanaa]

LGTM