[CI-Build] Fixed numeric issue of test_prefix_prefill on AMD

[him-rh-nm]

Resolves issue #28490.

Purpose

- relax the prefix prefill test tolerances on ROCm to account for the larger numerical error seen on MI300
- apply the same tolerance adjustment to the ALiBi variant of the test so both paths remain consistent

Test Plan

pytest -s -v 'tests/kernels/attention/test_prefix_prefill.py'

Test Result

---

Essential Elements of an Effective PR Description Checklist

• The purpose of the PR, such as "Fix some issue (link existing issues this PR will resolve)".
• The test plan, such as providing test command.
• The test results, such as pasting the results comparison before and after, or e2e results
• (Optional) The necessary documentation update, such as updating supported_models.md and examples for a new model.
• (Optional) Release notes update. If your change is user facing, please update the release notes draft in the Google Doc.

[github-actions]

👋 Hi! Thank you for contributing to the vLLM project.

💬 Join our developer Slack at https://slack.vllm.ai to discuss your PR in #pr-reviews, coordinate on features in #feat- channels, or join special interest groups in #sig- channels.

Just a reminder: PRs would not trigger full CI run by default. Instead, it would only run fastcheck CI which starts running only a small and essential subset of CI tests to quickly catch errors.

You ask your reviewers to trigger select CI tests on top of fastcheck CI.

Once the PR is approved and ready to go, your PR reviewer(s) can run CI to test the changes comprehensively before merging.

To run CI, PR reviewers can either: Add ready label to the PR or enable auto-merge.

If you have any questions, please reach out to us on Slack at https://slack.vllm.ai.

🚀

[gemini-code-assist]

Code Review

This pull request addresses a numerical precision issue on AMD hardware by relaxing the test tolerances for prefix prefill operations. The changes correctly introduce platform-specific logic to increase the absolute tolerance (atol) for assertions when running on ROCm, which is a standard approach for handling hardware-dependent numerical differences. My review identifies an area for improvement regarding code duplication. The logic for adjusting the tolerance is repeated in two separate test functions. I've suggested refactoring this into a shared helper function to enhance maintainability and prevent potential inconsistencies in the future.

[gemini-code-assist]

This logic block for adjusting tolerance on ROCm is duplicated from lines 317-323. This duplication, including the hardcoded tolerance values, can lead to maintenance issues, such as forgetting to update both locations if the tolerances need to be changed.

To improve maintainability, consider extracting this logic into a helper function. For example:

def _adjust_atol_for_rocm(atol: float, op: Callable, kv_cache_dtype: str) -> float:
    if current_platform.is_rocm() and op is chunked_prefill_paged_decode:
        rocm_atol = 5e-3 if "fp8" in kv_cache_dtype else 2e-3
        return max(atol, rocm_atol)
    return atol

Then you could simplify the code at both locations to a single line:
atol = _adjust_atol_for_rocm(atol, op, kv_cache_dtype)

[SageMoore]

This is a good suggestion

[SageMoore]

Do we need to relax the non-fp8 case as well?

[SageMoore]

This is a good suggestion

[mergify]

This pull request has merge conflicts that must be resolved before it can be
merged. Please rebase the PR, @him-rh-nm.

https://docs.github.com/en/pull-requests/collaborating-with-pull-requests/working-with-forks/syncing-a-fork