fix: correct register_xpu_ci args in DSA indexer XPU test (#24959)

python/sglang/kernels/ops/attention/dsa/index_buf_accessor.py

@@ -6,9 +6,10 @@
 
 from sglang.kernels.ops.quantization.fp8_kernel import is_fp8_fnuz
 from sglang.srt.layers.attention.dsa.utils import aiter_can_use_preshuffle_paged_mqa
-from sglang.srt.utils import get_bool_env_var, is_hip
+from sglang.srt.utils import get_bool_env_var, is_hip, is_xpu
 
 _is_hip = is_hip()
+_is_xpu = is_xpu()
 _is_fp8_fnuz = is_fp8_fnuz()
 _use_aiter = get_bool_env_var("SGLANG_USE_AITER") and _is_hip
 # aiter cp_gather kernel with preshuffle=True is only valid when the indexer
@@ -308,6 +309,11 @@ def _set_k_and_s_triton(
             assert (
                 page_size % 16 == 0
             ), f"HIP preshuffle requires page_size to be a multiple of 16, got {page_size}"
+    elif _is_xpu:
+        assert page_size in (
+            64,
+            128,
+        ), f"XPU DSA requires page_size 64 or 128, got {page_size}"
     else:
         assert page_size == 64
 

python/sglang/srt/hardware_backend/xpu/kernels/dsa/act_quant.py

@@ -0,0 +1,49 @@
+"""Torch-native per-group FP8 activation quantization for XPU."""
+
+from typing import Optional, Tuple
+
+import torch
+
+
+def act_quant(
+    x: torch.Tensor, block_size: int = 128, scale_fmt: Optional[str] = None
+) -> Tuple[torch.Tensor, torch.Tensor]:
+    """Per-group FP8 activation quantization using PyTorch ops.
+
+    For each group of `block_size` columns, computes the abs-max, derives a
+    per-group scale, and quantizes to float8_e4m3fn.
+
+    Args:
+        x: Input tensor (contiguous, last dim divisible by block_size).
+        block_size: Number of columns per quantization group.
+        scale_fmt: If not None, round scales to nearest power of 2.
+
+    Returns:
+        (y_fp8, scale) where y_fp8 has dtype float8_e4m3fn and scale is float32.
+    """
+    assert x.is_contiguous(), "Input tensor must be contiguous"
+    N = x.size(-1)
+    assert (
+        N % block_size == 0
+    ), f"Last dim must be divisible by block_size ({block_size})"
+
+    FP8_MAX = 448.0
+    x_flat = x.view(-1, N).float()
+    M = x_flat.size(0)
+    n_groups = N // block_size
+
+    # Reshape to (M, n_groups, block_size) for per-group quantization
+    x_grouped = x_flat.view(M, n_groups, block_size)
+    amax = x_grouped.abs().amax(dim=-1).clamp(min=1e-4)  # (M, n_groups)
+
+    if scale_fmt is not None:
+        # Round scale to power of 2
+        scale = torch.exp2(torch.log2(amax / FP8_MAX).ceil())
+    else:
+        scale = amax / FP8_MAX
+
+    # Quantize
+    y = (x_grouped / scale.unsqueeze(-1)).clamp(-FP8_MAX, FP8_MAX)
+    y = y.view(M, N).to(torch.float8_e4m3fn).view(*x.shape[:-1], N)
+    scale = scale.view(*x.shape[:-1], n_groups)
+    return y, scale

python/sglang/srt/layers/attention/dsa/dsa_indexer.py

@@ -98,6 +98,10 @@
     except ImportError as e:
         deep_gemm = e
 
+if _is_xpu:
+    from sgl_kernel import fp8_mqa_logits as sgl_fp8_mqa_logits
+    from sgl_kernel import fp8_paged_mqa_logits as sgl_fp8_paged_mqa_logits
+
 if _use_aiter:
     from aiter.ops.cache import indexer_k_quant_and_cache
 
@@ -337,7 +341,6 @@ def topk_transform(
 
 
 def rotate_activation(x: torch.Tensor) -> torch.Tensor:
-    # from sgl_kernel import hadamard_transform
     if _is_hip:
         from fast_hadamard_transform import hadamard_transform
     elif _is_xpu:
@@ -894,6 +897,11 @@ def _get_topk_paged(
                 assert (
                     page_size == 1
                 ), f"HIP legacy DSA path requires page_size == 1, got {page_size}"
+        elif _is_xpu:
+            assert page_size in (
+                64,
+                128,
+            ), f"XPU DSA only supports page_size 64 or 128, got {page_size}"
         else:
             assert page_size == 64, "only support page size 64"
         # NOTE(dark): this support extend/decode/decode+graph
@@ -985,6 +993,17 @@ def _get_topk_paged(
                 preshuffle=_use_aiter_preshuffle,
                 kv_block_size=block_kv,
             )
+        elif _is_xpu:
+            logits = sgl_fp8_paged_mqa_logits(
+                q_fp8[:q_offset],
+                kv_cache_fp8,
+                weights[:q_offset],
+                seqlens_32_2d,
+                block_tables,
+                None,
+                max_seq_len,
+                clean_logits=False,
+            )
         elif use_cute_dsl:
             logits = cutedsl_paged_mqa_logits(
                 q_fp8,
@@ -1052,7 +1071,10 @@ def _get_mqa_logits_budget_bytes(self, device_index: int) -> int:
         if cached_budget is not None:
             return cached_budget
 
-        total_mem = torch.cuda.get_device_properties(device_index).total_memory
+        if _is_xpu:
+            total_mem = torch.xpu.get_device_properties(device_index).total_memory
+        else:
+            total_mem = torch.cuda.get_device_properties(device_index).total_memory
 
         total_mem_budget = int(total_mem * self._MQA_LOGITS_TOTAL_MEM_FRACTION)
         mem_fraction_static = get_server_args().mem_fraction_static
@@ -1073,10 +1095,15 @@ def _get_mqa_logits_budget_bytes(self, device_index: int) -> int:
             return static_budget
 
         # Match the original free-memory guard: logits_bytes * 2 > free_mem.
-        # torch.cuda.mem_get_info synchronizes the host, so cache the result,
-        # capped by the workload-independent serving-memory headroom.
-        free_mem, _ = torch.cuda.mem_get_info(device_index)
-        budget_bytes = min(int(free_mem * free_mem_fraction), static_budget)
+        # Synchronizes the host; cache the result capped by serving-memory headroom.
+        if _is_xpu:
+            # On XPU, use total_mem budget as the free-memory estimate;
+            # dynamic free-memory query is not supported the same way as CUDA.
+            # TODO Use torch.xpu.mem_get_info() when available (planned end of 2026).
+            budget_bytes = static_budget
+        else:
+            free_mem, _ = torch.cuda.mem_get_info(device_index)
+            budget_bytes = min(int(free_mem * free_mem_fraction), static_budget)
 
         budget_bytes = max(1, budget_bytes)
         self._mqa_logits_budget_bytes[device_index] = budget_bytes
@@ -1125,13 +1152,19 @@ def _get_topk_ragged(
                     page_size == 1
                 ), f"HIP legacy DSA path requires page_size == 1, got {page_size}"
         else:
-            assert page_size == 64, "only support page size 64"
+            if _is_xpu:
+                assert page_size in (
+                    64,
+                    128,
+                ), f"XPU DSA requires page_size 64 or 128, got {page_size}"
+            else:
+                assert page_size == 64, "only support page size 64"
 
-        assert len(weights.shape) == 3
-        assert (
-            forward_batch.seq_lens_cpu is not None
-            and forward_batch.extend_seq_lens_cpu is not None
-        )
+            assert len(weights.shape) == 3
+            assert (
+                forward_batch.seq_lens_cpu is not None
+                and forward_batch.extend_seq_lens_cpu is not None
+            )
         weights = weights.squeeze(-1)
 
         if _is_hip and not _use_aiter_preshuffle:
@@ -1206,6 +1239,15 @@ def _get_topk_ragged(
                         ke,
                         clean_logits=False,
                     )
+                elif _is_xpu:
+                    logits = sgl_fp8_mqa_logits(
+                        q_fp8[:q_offset],
+                        kv_fp8,
+                        weights[:q_offset],
+                        ks,
+                        ke,
+                        clean_logits=False,
+                    )
                 else:
                     q_padded, w_padded, _ = self._pad_heads_for_deep_gemm(
                         q_fp8[:q_offset], weights[:q_offset]
@@ -1262,6 +1304,15 @@ def _get_topk_ragged(
                         ke[start:end],
                         clean_logits=False,
                     )
+                elif _is_xpu:
+                    logits_chunk = sgl_fp8_mqa_logits(
+                        q_fp8[start:end],
+                        kv_fp8,
+                        weights[start:end],
+                        ks[start:end],
+                        ke[start:end],
+                        clean_logits=False,
+                    )
                 else:
                     q_padded, w_padded, _ = self._pad_heads_for_deep_gemm(
                         q_fp8[start:end], weights[start:end]
@@ -1406,7 +1457,13 @@ def _get_topk_ragged_with_cp(
             assert isinstance(get_token_to_kv_pool(), DSATokenToKVPool)
 
         page_size = get_token_to_kv_pool().page_size
-        assert page_size == 64, "only support page size 64"
+        if _is_xpu:
+            assert page_size in (
+                64,
+                128,
+            ), f"XPU DSA requires page_size 64 or 128, got {page_size}"
+        else:
+            assert page_size == 64, "only support page size 64"
         assert len(weights.shape) == 3
         weights = weights.squeeze(-1)
         k_fp8_list = []
@@ -1558,7 +1615,13 @@ def forward_indexer(
             from sglang.kernels.ops.attention.dsa.tilelang_kernel import fp8_index
 
         page_size = get_token_to_kv_pool().page_size
-        assert page_size == 64, "only support page size 64"
+        if _is_xpu:
+            assert page_size in (
+                64,
+                128,
+            ), f"XPU DSA requires page_size 64 or 128, got {page_size}"
+        else:
+            assert page_size == 64, "only support page size 64"
 
         assert len(weights.shape) == 3
         weights = weights.squeeze(-1)
@@ -1735,6 +1798,10 @@ def forward_cuda(
     ) -> Optional[torch.Tensor]:
         if _is_hip:
             from sglang.kernels.ops.attention.dsa.tilelang_kernel import act_quant
+        elif _is_xpu:
+            from sglang.srt.hardware_backend.xpu.kernels.dsa.act_quant import (
+                act_quant,
+            )
         elif not _is_npu:
             from sglang.kernels.ops.attention.dsa.triton_kernel import act_quant
 
@@ -1968,7 +2035,7 @@ def forward_cuda(
             else:
                 weights = self._get_logits_head_gate(x_for_gate, q_scale)
 
-        if _is_cuda or _is_hip:
+        if _is_cuda or _is_hip or _is_xpu:
             # In piecewise/breakable CUDA graph, any access to seq_lens_cpu
             # creates a Dynamo shape guard. These graph modes never have empty
             # batches.