[XPU][Fix] Pack device-pointer tables as uint64 to avoid 64-bit address overflow (#35051)

Co-authored-by: roopaksrivastav <roopak.srivastava@intel.com>
Co-authored-by: Claude Opus 5 <noreply@anthropic.com>

test/registered/unit/layers/test_mamba_state_scatter_triton.py

@@ -2,10 +2,12 @@ from sglang.test.ci.ci_register import (
     register_amd_ci,
     register_cpu_ci,
     register_cuda_ci,
+    register_xpu_ci,
 )
 
 register_cuda_ci(est_time=7, stage="base-b", runner_config="1-gpu-small")
 register_amd_ci(est_time=7, suite="stage-b-test-1-gpu-small-amd-mi35x")
+register_xpu_ci(est_time=20, suite="stage-b-test-1-gpu-xpu")
 # The dst layout-contract tests run on CPU (no kernel launch).
 register_cpu_ci(est_time=6, suite="base-a-test-cpu")
 
@@ -16,6 +18,7 @@ import torch
 try:
     from sglang.kernels.ops.mamba.mamba_state_scatter_triton import (
         _require_entry_contiguous_dst,
+        fused_conv_window_scatter_multi,
         fused_conv_window_scatter_with_mask,
         fused_mamba_state_scatter_with_mask,
     )
@@ -23,6 +26,7 @@ try:
     _FUSED_IMPORT_ERROR = None
 except Exception as e:  # pragma: no cover
     _require_entry_contiguous_dst = None
+    fused_conv_window_scatter_multi = None
     fused_conv_window_scatter_with_mask = None
     fused_mamba_state_scatter_with_mask = None
     _FUSED_IMPORT_ERROR = e
@@ -31,6 +35,25 @@ from sglang.srt.mem_cache.layout.page_major import (
     build_page_major_mamba_views,
     mamba_entry_bytes,
 )
+from sglang.srt.utils import get_device
+from sglang.srt.utils.common import get_device_module
+from sglang.test.test_utils import CustomTestCase
+
+# Backends the multi-type scatter is verified against; extend as others gain Triton.
+TRITON_DEVICES = ("cuda", "xpu")
+
+
+def _triton_device():
+    # get_device() raises when the host has no accelerator; the layout-contract
+    # cases below still run on the CPU suite.
+    try:
+        device = get_device()
+    except RuntimeError:
+        return None
+    return device if device in TRITON_DEVICES else None
+
+
+DEVICE = _triton_device()
 
 
 def _ref_scatter(dst, src, dst_indices, src_indices, step_indices):
@@ -365,5 +388,78 @@ class TestMambaStateScatterEnvelopeDst(unittest.TestCase):
         torch.testing.assert_close(conv_views[0], expect_conv)
 
 
+@unittest.skipUnless(
+    DEVICE is not None,
+    f"multi-type conv scatter needs one of {TRITON_DEVICES}",
+)
+class TestFusedConvWindowScatterMulti(CustomTestCase):
+    """Multi-type conv scatter must place every conv type correctly; it reaches
+    them through a host-built device-pointer table (issue #35047)."""
+
+    def _run(self, num_types, n2):
+        # n2 sizes the optional second (interval-crossing track) index set.
+        torch.manual_seed(11)
+        dev = DEVICE
+        layers, slots, batch, steps, dim, km1 = 2, 16, 5, 3, 8, 3
+        pairs = []
+        for t in range(num_types):
+            elems = dim * (t + 1)
+            pairs.append(
+                (
+                    torch.randn(
+                        (layers, slots, elems, km1), dtype=torch.bfloat16, device=dev
+                    ),
+                    torch.randn(
+                        (layers, batch, steps, elems, km1),
+                        dtype=torch.bfloat16,
+                        device=dev,
+                    ),
+                )
+            )
+
+        g = torch.Generator(device=dev).manual_seed(5)
+        # One permutation split across the two sets: their dst slots must be
+        # disjoint, else the two writes to a shared slot race.
+        perm = torch.randperm(slots, device=dev, generator=g).to(torch.int64)
+        dst1, dst2 = perm[:batch], perm[batch : batch + n2]
+
+        def _steps(n):
+            step = torch.randint(
+                0, steps, (n,), device=dev, dtype=torch.int64, generator=g
+            )
+            step[0] = -1  # one rejected row per set must be skipped
+            return step
+
+        step1 = _steps(batch)
+        if n2:
+            step2 = _steps(n2)
+        else:
+            dst2, step2 = None, None
+
+        expect = [dst.clone() for dst, _ in pairs]
+        for exp, (_, src) in zip(expect, pairs):
+            for dsts, stps in ((dst1, step1), (dst2, step2)):
+                if stps is None:
+                    continue
+                valid = stps >= 0
+                # src row = the row's position within its own index set
+                rows = torch.arange(stps.numel(), device=dev)[valid]
+                exp[:, dsts[valid]] = src[:, rows, stps[valid]]
+
+        fused_conv_window_scatter_multi(pairs, dst1, step1, dst2, step2)
+        get_device_module().synchronize()
+        for i, (exp, (got, _)) in enumerate(zip(expect, pairs)):
+            torch.testing.assert_close(got, exp, msg=f"conv type {i} mismatch")
+
+    def test_single_type(self):
+        self._run(num_types=1, n2=0)
+
+    def test_multi_type(self):
+        self._run(num_types=3, n2=0)
+
+    def test_multi_type_with_track_set(self):
+        self._run(num_types=2, n2=4)
+
+
 if __name__ == "__main__":  # pragma: no cover
     unittest.main()