[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>

python/sglang/kernels/ops/mamba/mamba_state_scatter_triton.py

@@ -10,6 +10,8 @@ import torch
 import triton
 import triton.language as tl
 
+from sglang.kernels.ops.memory.ptr_table import make_ptr_table
+
 
 def _require_entry_contiguous_dst(
     dst: torch.Tensor, entry_start_dim: int, fn_name: str
@@ -578,7 +580,7 @@ def _conv_multi_build_meta(pairs, block_size: int):
             ]
         )
         block_start += triton.cdiv(elem, block_size)
-    meta = torch.tensor(rows, dtype=torch.int64, device=pairs[0][0].device)
+    meta = make_ptr_table(rows, device=pairs[0][0].device)
     return meta, block_start
 
 

python/sglang/kernels/ops/memory/ptr_table.py

@@ -0,0 +1,22 @@
+"""Device-pointer tables for kernels that address several tensors per launch."""
+
+from __future__ import annotations
+
+from typing import Sequence, Union
+
+import torch
+
+
+def make_ptr_table(
+    rows: Union[Sequence[int], Sequence[Sequence[int]]],
+    device: Union[torch.device, str],
+) -> torch.Tensor:
+    """Pack ``data_ptr()`` values -- flat, or 2-D with companion columns such as
+    strides -- into an ``int64`` table a kernel bitcasts back to pointers.
+
+    Built unsigned because XPU USM addresses set the top bit, which
+    ``dtype=torch.int64`` rejects while unpacking through ``long long``;
+    ``view`` moves no bits, so kernels keep their signed element type.
+    Values must be in ``[0, 2**64)``.
+    """
+    return torch.tensor(rows, dtype=torch.uint64, device=device).view(torch.int64)

python/sglang/srt/mem_cache/mamba_slot_fused.py

@@ -24,11 +24,13 @@ import torch
 import triton
 import triton.language as tl
 
+from sglang.kernels.ops.memory.ptr_table import make_ptr_table
+
 _BLOCK = 1024
 
 
 class ConvSlotDescriptor(NamedTuple):
-    ptr: torch.Tensor  # [T] int64 base byte-addresses
+    ptr: torch.Tensor  # [T] int64-viewed base byte-addresses (make_ptr_table)
     feat: torch.Tensor  # [T] int64 per-slot feature length (elements)
     layer_stride: torch.Tensor  # [T] int64 element stride between layers
     slot_stride: torch.Tensor  # [T] int64 element stride between slots
@@ -116,7 +118,7 @@ def build_conv_slot_descriptor(tensors: List[torch.Tensor]) -> ConvSlotDescripto
         max_feat = max(max_feat, t[0, 0].numel())
     to_i64 = lambda xs: torch.tensor(xs, dtype=torch.int64, device=device)
     return ConvSlotDescriptor(
-        ptr=to_i64(ptr),
+        ptr=make_ptr_table(ptr, device=device),
         feat=to_i64(feat),
         layer_stride=to_i64(layer_stride),
         slot_stride=to_i64(slot_stride),

python/sglang/srt/mem_cache/memory_pool.py

@@ -489,7 +489,7 @@ class MambaPool:
                 *physical_conv_shape,
             ),
             dtype=conv_dtype,
-            device="cuda",
+            device=self.device,
         )
         physical_conv_strides = phys.stride()[2:]
         window_stride = physical_conv_strides[window_axis]
@@ -771,7 +771,7 @@ class MambaPool:
                             temporal_state_shape[2],
                         ),
                         dtype=ssm_dtype,
-                        device="cuda",
+                        device=device,
                     )
                 # Cache intermediate conv windows (last K-1 inputs) per draft token
                 # during target verify.
@@ -839,7 +839,7 @@ class MambaPool:
                                 conv_shape[1],
                             ),
                             dtype=conv_dtype,
-                            device="cuda",
+                            device=device,
                         )
                         for conv_shape in dense_conv_shapes
                     ]