diff --git a/python/sglang/srt/layers/attention/dsv4/indexer.py b/python/sglang/srt/layers/attention/dsv4/indexer.py index 93f026a5c..0f92d65d4 100644 --- a/python/sglang/srt/layers/attention/dsv4/indexer.py +++ b/python/sglang/srt/layers/attention/dsv4/indexer.py @@ -32,7 +32,7 @@ from sglang.srt.model_executor.runner_backend_utils.tc_piecewise_cuda_graph impo ) from sglang.srt.runtime_context import get_parallel from sglang.srt.state_capturer.indexer_topk import get_global_indexer_capturer -from sglang.srt.utils import add_prefix, is_cuda, is_hip +from sglang.srt.utils import add_prefix, is_cuda, is_hip, is_xpu from sglang.srt.utils.common import is_sm120_supported if TYPE_CHECKING: @@ -668,6 +668,12 @@ class C4IndexerBackendMixin: fn = fp8_paged_mqa_logits_torch_sm120 else: fn = fp8_paged_mqa_logits_torch + elif is_xpu(): + from sgl_kernel import fp8_paged_mqa_logits_triton + + # TODO: switch from triton to SYCL when OOM is resolved + + fn = fp8_paged_mqa_logits_triton else: from deep_gemm import fp8_paged_mqa_logits as fn diff --git a/python/sglang/srt/layers/attention/dsv4/metadata.py b/python/sglang/srt/layers/attention/dsv4/metadata.py index d26d1ccf5..c5335b0e0 100644 --- a/python/sglang/srt/layers/attention/dsv4/metadata.py +++ b/python/sglang/srt/layers/attention/dsv4/metadata.py @@ -7,7 +7,7 @@ from typing import TYPE_CHECKING, Any, List, Optional import torch from sglang.srt.environ import envs -from sglang.srt.utils import is_hip +from sglang.srt.utils import is_hip, is_xpu if TYPE_CHECKING: pass @@ -122,6 +122,7 @@ class PagedIndexerMetadata: def __post_init__(self): if ( envs.SGLANG_FP8_PAGED_MQA_LOGITS_TORCH.get() + or is_xpu() or envs.SGLANG_OPT_USE_AITER_INDEXER.get() ): self.deep_gemm_metadata = None