[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>
This commit is contained in:
Agata Dobrzyniewicz
2026-06-30 09:57:27 +00:00
committed by GitHub
co-authored by Kunshang Ji
parent 364ee36af1
commit 8e9d70fdd5
+5 -5
View File
@@ -347,7 +347,7 @@ def test_selective_state_update(dim, dstate, has_z, itype):
rtol, atol = (3e-4, 1e-3) if itype == torch.float32 else (5e-3, 1e-2)
if itype == torch.bfloat16:
rtol, atol = 1e-2, 5e-2
if torch.version.hip:
if current_platform.is_rocm() or current_platform.is_xpu():
atol *= 2
# set seed
set_random_seed(0)
@@ -437,7 +437,7 @@ def test_selective_state_update_varlen(dim, dstate, has_z, itype, max_seq_len):
rtol, atol = (3e-4, 1e-3) if itype == torch.float32 else (5e-3, 1e-2)
if itype == torch.bfloat16:
rtol, atol = 5e-2, 1.5e-1
if torch.version.hip:
if current_platform.is_rocm() or current_platform.is_xpu():
atol *= 2
# set seed
set_random_seed(0)
@@ -700,7 +700,7 @@ def test_selective_state_update_with_batch_indices(
rtol, atol = (3e-4, 1e-3) if itype == torch.float32 else (5e-3, 1e-2)
if itype == torch.bfloat16:
rtol, atol = 1e-1, 1e-1
if torch.version.hip:
if current_platform.is_rocm() or current_platform.is_xpu():
atol *= 2
# set seed
torch.random.manual_seed(0)
@@ -865,7 +865,7 @@ def test_selective_state_update_with_num_accepted_tokens(
rtol, atol = (3e-4, 1e-3) if itype == torch.float32 else (5e-3, 1e-2)
if itype == torch.bfloat16:
rtol, atol = 5e-2, 1.5e-1
if torch.version.hip:
if current_platform.is_rocm() or current_platform.is_xpu():
atol *= 2
set_random_seed(0)
@@ -991,7 +991,7 @@ def test_selective_state_update_varlen_with_num_accepted(
rtol, atol = (3e-4, 1e-3) if itype == torch.float32 else (5e-3, 1e-2)
if itype == torch.bfloat16:
rtol, atol = 5e-2, 1.5e-1
if torch.version.hip:
if current_platform.is_rocm() or current_platform.is_xpu():
atol *= 2
set_random_seed(0)