[ROCm][CI] Align spec decode logprob test prefill settings - #41335
Conversation
Signed-off-by: Andreas Karatzas <akaratza@amd.com>
![claude[bot]](https://avatars.githubusercontent.com/in/1236702?s=60&v=4){kind=small}
![gemini-code-assist[bot]](https://avatars.githubusercontent.com/in/956858?s=60&v=4){kind=small}
There was a problem hiding this comment.
Code Review
This pull request refactors the test_spec_decode_logprobs test to use a shared configuration dictionary for both the reference and speculative LLM instances. By enforcing consistent chunked prefill settings and batch sizes, the changes ensure that logprob comparisons remain deterministic and isolated from batch geometry variations on ROCm hardware. I have no feedback to provide as there were no review comments.
Bortlesboat
left a comment
Bortlesboat
left a comment
There was a problem hiding this comment.
Looks right. Worth flagging that this also narrows what the test exercises — the prior setup had ref_llm running in default scheduling and spec_llm in chunked prefill, so the comparison was implicitly mixing "spec decode correctness" with "logprobs are stable across scheduling shifts." The cross-scheduling part is what was producing the ROCm noise in the first place. Forcing both runs into the same enable_chunked_prefill=True / max_num_batched_tokens=32 isolates it to spec decode only, which is the test's actual job. The shared llm_kwargs factoring also makes any future drift between the two configs impossible by construction.
…ect#41335) Signed-off-by: Andreas Karatzas <akaratza@amd.com>
…ect#41335) Signed-off-by: Andreas Karatzas <akaratza@amd.com> Co-authored-by: hongbolv <33214277+hongbolv@users.noreply.github.com>
…ect#41335) Signed-off-by: Andreas Karatzas <akaratza@amd.com> Signed-off-by: Ifta Khairul Alam Adil <ikaadil007@gmail.com>
…ect#41335) Signed-off-by: Andreas Karatzas <akaratza@amd.com>
…ect#41335) Signed-off-by: Andreas Karatzas <akaratza@amd.com>
…ect#41335) Signed-off-by: Andreas Karatzas <akaratza@amd.com>
…ect#41335) Signed-off-by: Andreas Karatzas <akaratza@amd.com> Signed-off-by: Matt Van Horn <455140+mvanhorn@users.noreply.github.com>
…ect#41335) Signed-off-by: Andreas Karatzas <akaratza@amd.com>
3 participants
This PR fixes a ROCm-only failure in
test_spec_decode_logprobsfor the ngram speculative decoding cases. The test compares logprobs from a regular base-model run against logprobs from a speculative decoding run. For that comparison to isolate speculative decoding, both LLM instances need to use the same execution settings. Before this change, the spec-decode LLM forced chunked prefill withenable_chunked_prefill=Trueandmax_num_batched_tokens=32, while the reference LLM did not.On ROCm, floating-point reductions can vary with batch geometry and scheduling. That means the two runs could produce slightly different logprobs even when speculative decoding itself was behaving correctly. The observed failure was a logprob difference of about
0.235, which exceeded the test tolerance.This change factors the shared LLM options into
llm_kwargsand passes them to both the reference and spec-decode LLMs. The only intended difference between the two runs is now the speculative decoding configuration, so the test better serves its original purpose.cc @kenroche