[Kernel][XPU] Adjust kernel unit tests for XPU (#45140)
Signed-off-by: Dobrzyniewicz, Agata <agata.dobrzyniewicz@intel.com> Co-authored-by: Kunshang Ji <kunshang.ji@intel.com>
Agata Dobrzyniewicz committed
Jun 30, 2026 at 11:57 UTC
8e9d70fdd5842a3bc0ef2a45e5413be1f4ffd2d6
1 file changed
+5
-5
tests/kernels/mamba/test_mamba_ssm.py
+5
-5
@@ -347,7 +347,7 @@ def test_selective_state_update(dim, dstate, has_z, itype):
347
rtol, atol = (3e-4, 1e-3) if itype == torch.float32 else (5e-3, 1e-2)
348
if itype == torch.bfloat16:
349
rtol, atol = 1e-2, 5e-2
350
- if torch.version.hip:
350
+ if current_platform.is_rocm() or current_platform.is_xpu():
351
atol *= 2
352
# set seed
353
set_random_seed(0)
@@ -437,7 +437,7 @@ def test_selective_state_update_varlen(dim, dstate, has_z, itype, max_seq_len):
437
rtol, atol = (3e-4, 1e-3) if itype == torch.float32 else (5e-3, 1e-2)
438
if itype == torch.bfloat16:
439
rtol, atol = 5e-2, 1.5e-1
440
- if torch.version.hip:
440
+ if current_platform.is_rocm() or current_platform.is_xpu():
441
atol *= 2
442
# set seed
443
set_random_seed(0)
@@ -700,7 +700,7 @@ def test_selective_state_update_with_batch_indices(
700
rtol, atol = (3e-4, 1e-3) if itype == torch.float32 else (5e-3, 1e-2)
701
if itype == torch.bfloat16:
702
rtol, atol = 1e-1, 1e-1
703
- if torch.version.hip:
703
+ if current_platform.is_rocm() or current_platform.is_xpu():
704
atol *= 2
705
# set seed
706
torch.random.manual_seed(0)
@@ -865,7 +865,7 @@ def test_selective_state_update_with_num_accepted_tokens(
865
rtol, atol = (3e-4, 1e-3) if itype == torch.float32 else (5e-3, 1e-2)
866
if itype == torch.bfloat16:
867
rtol, atol = 5e-2, 1.5e-1
868
- if torch.version.hip:
868
+ if current_platform.is_rocm() or current_platform.is_xpu():
869
atol *= 2
870
871
set_random_seed(0)
@@ -991,7 +991,7 @@ def test_selective_state_update_varlen_with_num_accepted(
991
rtol, atol = (3e-4, 1e-3) if itype == torch.float32 else (5e-3, 1e-2)
992
if itype == torch.bfloat16:
993
rtol, atol = 5e-2, 1.5e-1
994
- if torch.version.hip:
994
+ if current_platform.is_rocm() or current_platform.is_xpu():
995
atol *= 2
996
997
set_random_seed(0)