[ROCm][CI] Fix order-dependent failure in test_flash_attn_accepts_handled_fp8_variants (MI355)

[stefankoncarevic]

Purpose

tests/kernels/attention/test_attention_selector.py::test_flash_attn_accepts_handled_fp8_variants[fp8|fp8_e4m3]
fails on the Kernels (B200-MI355) CI group when the file is run in full, while
passing when the two cases are run in isolation:

assert FlashAttentionBackend.supports_kv_cache_dtype('fp8_e4m3')
E   AssertionError: assert False
E    +  where False = supports_kv_cache_dtype('fp8_e4m3')

This is a test-ordering bug, not a production regression.

Root cause

The test verifies that FlashAttentionBackend accepts fp8/fp8_e4m3 on
platforms that support them. On ROCm these dtypes are genuinely unsupported
(get_flash_attn_version() returns None), so the test simulates XPU by
patching is_xpu to return True.

The accept/reject decision is made in
vllm/v1/attention/backends/fa_utils.py::flash_attn_supports_kv_cache_dtype:

if current_platform.is_xpu():
    return True

Because a Python function resolves globals from its own defining module,
this reads fa_utils.current_platform — regardless of who calls it.
FlashAttentionBackend.supports_kv_cache_dtype simply delegates to that helper
and never reads is_xpu itself.

The original test patched is_xpu on flash_attn.current_platform (a different
module's binding). current_platform is a lazily-resolved singleton exposed via
vllm.platforms.__getattr__, and each module captures its own reference at
import time (from vllm.platforms import current_platform). Earlier tests in the
file run with patch("vllm.platforms.current_platform", RocmPlatform()) and
trigger backend imports inside that context, which can leave fa_utils and
flash_attn bound to different platform objects depending on import order. When
they diverge, patching flash_attn's binding has no effect on the object
fa_utils actually checks, so is_xpu() stays False and the assertion fails.

In isolation the two bindings are the same object, so the original patch happens
to work — hence the isolated-pass / full-file-fail behavior.

Why it started failing now

Introduced by #46080 ("[Hardware][AMD][CI] Fix Kernels Attention test groups"),
which (a) routed the full test_attention_selector.py into the
Kernels (B200-MI355) group, and (b) made test_non_causal_backend_selection
run on ROCm (previously skipped when CudaPlatform is None) under
patch("vllm.platforms.current_platform", RocmPlatform()). That exposed the
pre-existing fragile patch target added in #42685.

Fix

Patch is_xpu on the binding the decision actually uses (fa_utils), which is
order-robust:

import vllm.v1.attention.backends.fa_utils as fa_utils_mod
from vllm.v1.attention.backends.flash_attn import FlashAttentionBackend

# The fp8 decision is made in fa_utils, using its own current_platform
# binding, so patch is_xpu there (not on flash_attn's) to stay robust to
# import order across earlier tests that patch vllm.platforms.current_platform.
monkeypatch.setattr(fa_utils_mod.current_platform, "is_xpu", lambda: True)
assert FlashAttentionBackend.supports_kv_cache_dtype(kv_cache_dtype)

This is a test-only change; no kernel or runtime code is modified. The
monkeypatch fixture is function-scoped and auto-reverts, so no state leaks to
other tests or groups.

Test Plan

On a gfx950 (MI355) host inside the ROCm test image:

# full file (as CI runs it)
pytest -v -s tests/kernels/attention/test_attention_selector.py
# and the two cases in isolation
pytest -v tests/kernels/attention/test_attention_selector.py::test_flash_attn_accepts_handled_fp8_variants

Test Result

Before: full-file run fails the two fp8/fp8_e4m3 cases (isolated run passes).

After:

# full file
23 passed, 7 skipped in ~8s
# isolated
2 passed in ~1s

Notes

• Root fix alternative (not taken): import flash_attn/fa_utils at the top of
  the test module before any patch("vllm.platforms.current_platform", ...) so
  all bindings point at the real singleton. That addresses the broader
  contamination class but is more invasive and relies on import-ordering
  guarantees; the targeted patch above is the minimal correct change.

---

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.

[claude]

Claude Code Review

This pull request is from a fork — automated review is disabled. A repository maintainer can comment @claude review to run a one-time review.

[AndreasKaratzas]

LGTM