[Misc] fix tests failure by using current_platform (#25825)

Signed-off-by: Juechen Liu <jueliu@meta.com>
This commit is contained in:
Juechen Liu 2025-09-28 21:18:57 -07:00 committed by GitHub
parent 0307428d65
commit a3ae45a38c
No known key found for this signature in database
GPG Key ID: B5690EEEBB952194

View File

@ -137,7 +137,7 @@ def triton_reshape_and_cache_flash(
# heuristics instead of autotuning
TILE_SIZE = min(2048, triton.next_power_of_2(n))
if torch.version.hip or torch.version.xpu:
if current_platform.is_rocm() or current_platform.is_xpu():
num_stages = 4
num_warps = 8
else: # cuda