[Bugfix][XPU] Fix Mamba state pointer overflow (#48109)
Signed-off-by: Oxygen56 <jiangth99@163.com>
Willow Lopez committed
Aug 19, 2026 at 10:45 UTC
f936a267f9a774b7ac09f987226d8e5a591139da
2 files changed
+98
-2
tests/v1/worker/test_mamba_utils.py
+89
@@ -18,6 +18,7 @@ from vllm.v1.kv_cache_interface import KVCacheConfig, KVCacheGroupSpec, MambaSpe
18
from vllm.v1.worker.mamba_utils import (
19
MambaCopyBuffers,
20
MambaSpecDecodeGPUContext,
21
+ _reinterpret_u64_as_i64,
22
batch_memcpy,
23
collect_mamba_copy_meta,
24
do_mamba_copy_block,
@@ -189,6 +190,94 @@ class _MockCpuGpuBuffer:
190
return self.gpu[:n].copy_(self.cpu[:n], non_blocking=True)
191
192
193
+class _FakeDataPtrTensor:
194
+ """Tensor wrapper that exposes a controlled data_ptr for metadata tests."""
195
+
196
+ def __init__(self, tensor: torch.Tensor, data_ptr: int):
197
+ self._tensor = tensor
198
+ self._data_ptr = data_ptr
199
+ self.shape = tensor.shape
200
+
201
+ def data_ptr(self) -> int:
202
+ return self._data_ptr
203
+
204
+ def dim(self) -> int:
205
+ return self._tensor.dim()
206
+
207
+ def stride(self, *args):
208
+ return self._tensor.stride(*args)
209
+
210
+ def numel(self) -> int:
211
+ return self._tensor.numel()
212
+
213
+ def element_size(self) -> int:
214
+ return self._tensor.element_size()
215
+
216
+ def size(self, *args):
217
+ return self._tensor.size(*args)
218
+
219
+ def __getitem__(self, item):
220
+ return self._tensor[item]
221
+
222
+
223
+def test_reinterpret_u64_as_i64_preserves_pointer_bits():
224
+ ptrs = [
225
+ 0,
226
+ 1,
227
+ (1 << 63) - 1,
228
+ 1 << 63,
229
+ (1 << 63) + 1234,
230
+ (1 << 64) - 1,
231
+ ]
232
+ ptr_tensor = torch.zeros(len(ptrs), dtype=torch.int64)
233
+
234
+ for idx, ptr in enumerate(ptrs):
235
+ ptr_tensor[idx] = _reinterpret_u64_as_i64(ptr)
236
+
237
+ assert ptr_tensor.numpy().view(np.uint64).tolist() == ptrs
238
+
239
+
240
+def test_gpu_context_reinterprets_high_data_ptrs_for_int64_metadata():
241
+ cfg = _TestConfig(num_layers=1)
242
+ device = torch.device("cpu")
243
+ kv_cache_config = _make_kv_cache_config(cfg, ["layer_0"])
244
+ gpu_ctx = _make_gpu_ctx(cfg, kv_cache_config, device)
245
+ conv_ptr = 1 << 63
246
+ temporal_ptr = (1 << 64) - 8
247
+ block_table_ptr = (1 << 63) + 42
248
+
249
+ conv_state = _FakeDataPtrTensor(
250
+ torch.empty(
251
+ cfg.num_blocks,
252
+ cfg.conv_width,
253
+ cfg.conv_inner_dim,
254
+ dtype=cfg.dtype,
255
+ ),
256
+ conv_ptr,
257
+ )
258
+ temporal_state = _FakeDataPtrTensor(
259
+ torch.empty(cfg.num_blocks, cfg.temporal_state_dim, dtype=cfg.dtype),
260
+ temporal_ptr,
261
+ )
262
+ block_table = _FakeDataPtrTensor(
263
+ torch.empty(1, 4, dtype=torch.int32),
264
+ block_table_ptr,
265
+ )
266
+ forward_context = {"layer_0": _make_mock_attention(conv_state, temporal_state)}
267
+
268
+ gpu_ctx.initialize_from_forward_context(
269
+ kv_cache_config, forward_context, _COPY_FUNCS, [block_table]
270
+ )
271
+
272
+ assert gpu_ctx.state_base_addrs.tolist() == [
273
+ _reinterpret_u64_as_i64(conv_ptr),
274
+ _reinterpret_u64_as_i64(temporal_ptr),
275
+ ]
276
+ assert gpu_ctx.block_table_ptrs.tolist() == [
277
+ _reinterpret_u64_as_i64(block_table_ptr)
278
+ ]
279
+
280
+
281
def _make_postprocess_scheduler_output(
282
req_ids: list[str],
283
num_scheduled_tokens: dict[str, int],
vllm/v1/worker/mamba_utils.py
+9
-2
@@ -114,6 +114,11 @@ def _memcpy_u64_tiled(
114
tl.store(dst_u8 + i + offsets, data, mask=mask)
115
116
117
+def _reinterpret_u64_as_i64(value: int) -> int:
118
+ """Preserve a uint64 pointer bit pattern in a torch.int64 tensor."""
119
+ return value if value < (1 << 63) else value - (1 << 64)
120
+
121
+
122
@triton.jit
123
def _copy_mamba_state_block(
124
state_idx,
@@ -836,7 +841,9 @@ class MambaSpecDecodeGPUContext:
841
for state_type_idx, copy_func in enumerate(mamba_state_copy_funcs):
842
state = kv_caches[state_type_idx]
843
# Base address
839
- self.state_base_addrs[idx] = state.data_ptr()
844
+ self.state_base_addrs[idx] = _reinterpret_u64_as_i64(
845
+ state.data_ptr()
846
+ )
847
848
# Block stride (bytes between consecutive blocks)
849
# state shape: [num_blocks, ...], stride(0) = elements per block
@@ -923,7 +930,7 @@ class MambaSpecDecodeGPUContext:
930
)
931
self.block_table_stride_req = int(next(iter(strides)))
932
for i, bt in enumerate(block_tables):
926
- self.block_table_ptrs[i] = bt.data_ptr()
933
+ self.block_table_ptrs[i] = _reinterpret_u64_as_i64(bt.data_ptr())
934
935
self.is_initialized = True
936