From bd922bc39b09748d5f0f1c355bb8dcf41c4cc5e1 Mon Sep 17 00:00:00 2001 From: Bingxu Chen Date: Thu, 10 Sep 2026 15:02:00 +0800 Subject: [PATCH] [AMD][CI] Fix UMBP test buffer after MoRI upgrade (#38672) --- test/registered/unit/mem_cache/test_umbp_store.py | 6 ++++-- 1 file changed, 4 insertions(+), 2 deletions(-) diff --git a/test/registered/unit/mem_cache/test_umbp_store.py b/test/registered/unit/mem_cache/test_umbp_store.py index cede83f4e..271b78a59 100755 --- a/test/registered/unit/mem_cache/test_umbp_store.py +++ b/test/registered/unit/mem_cache/test_umbp_store.py @@ -12,6 +12,7 @@ from typing import Optional from unittest.mock import MagicMock, patch import mori.umbp # noqa: F401 +import torch from sglang.test.ci.ci_register import register_amd_ci @@ -43,8 +44,9 @@ class MockHostKVCache: total_bytes = num_pages * 2 * element_size # K+V for each page self._buffer = (ctypes.c_char * total_bytes)() self._buffer_ptr = ctypes.addressof(self._buffer) - self.kv_buffer = MagicMock() - self.kv_buffer.data_ptr.return_value = self._buffer_ptr + self.kv_buffer = torch.frombuffer( + self._buffer, dtype=torch.uint8, count=total_bytes + ) def get_page_buffer_meta(self, indices): """Return (ptr_list, element_size_list) for MHA page_first layout.