[XPU] Enable Kimi K3 KDA kernel tests on XPU (#51809)
Signed-off-by: pmanczak <pawel.manczak@intel.com>
pmanczak committed
Aug 17, 2026 at 11:21 UTC
cc7cf71fc819df579664adc2e944438d464dbc30
2 files changed
+8
-2
tests/models/kimi_k3/test_kda.py
+7
-1
@@ -30,9 +30,15 @@ from vllm.models.kimi_k3.nvidia.ops.third_party.kda import (
30
fused_recurrent_kda_fwd,
31
fused_recurrent_kda_packed_decode,
32
)
33
+from vllm.platforms import current_platform
34
from vllm.third_party.flash_linear_attention.ops.l2norm import l2norm_fwd
35
35
-DEVICE = "cuda"
36
+DEVICE = current_platform.device_type
37
+
38
+pytestmark = pytest.mark.skipif(
39
+ not (current_platform.is_cuda_alike() or current_platform.is_xpu()),
40
+ reason="The KDA kernels require a CUDA-alike or XPU device.",
41
+)
42
43
# The AMD and NVIDIA copies of the KDA kernels are vendored separately and are
44
# free to diverge, so the shared-semantics tests below run against both.
vllm/model_executor/layers/mamba/ops/gather_initial_states.py
+1
-1
@@ -49,7 +49,7 @@ def gather_initial_states(
49
) -> torch.Tensor:
50
"""Gather dense state rows, replacing uninitialized rows with zeros."""
51
assert state.ndim >= 2
52
- assert state.is_cuda
52
+ assert state.is_cuda or state.is_xpu
53
assert indices.ndim == 1 and has_initial_state.ndim == 1
54
assert indices.shape == has_initial_state.shape
55
assert indices.device == state.device