Skip to content

[HiCache] fix unit test with changed new APIs #13498

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
stmatengss merged 3 commits into from
Nov 21, 2025
Merged

[HiCache] fix unit test with changed new APIs #13498

Changes from all commits
Commits
Show all changes
3 commits
Select commit Hold shift + click to select a range
465e7aa
[HiCache] fix unit test with changed new APIs
stmatengss Nov 18, 2025
24ca76f
Merge branch 'main' into fix_mooncake_store_test_bug
stmatengss Nov 18, 2025
8cf0c66
Merge branch 'main' into fix_mooncake_store_test_bug
stmatengss Nov 19, 2025
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
36 changes: 32 additions & 4 deletions python/sglang/srt/mem_cache/storage/mooncake_store/test_mooncake_store.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -29,6 +29,29 @@ def generate_batch_query_keys(kv_num: int, config: HiCacheStorageConfig):
return set_keys, get_keys, exist_keys


def create_mock_host_kv_cache(buffer_size, dtype=torch.float32):
"""Create a mock HostKVCache-like object for testing."""
buffer = torch.randn(buffer_size, dtype=dtype)

class MockHostKVCache:
def __init__(self, buffer):
self.kv_buffer = buffer
self.layout = "page_first"
self.page_size = 1 # Simple page size for testing

def get_page_buffer_meta(self, indices):
"""Mock implementation of get_page_buffer_meta."""
ptr_list = []
element_size_list = []
for idx in indices:
# Create mock pointers and sizes for each page
ptr_list.append(idx * self.page_size * self.kv_buffer.element_size())
element_size_list.append(self.page_size * self.kv_buffer.element_size())
return ptr_list, element_size_list
Comment on lines +42 to +50
Copy link
Contributor

@gemini-code-assist gemini-code-assist bot Nov 18, 2025

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

medium

The mock implementation of get_page_buffer_meta calculates byte offsets but returns them as if they were pointers. The C++ backend that consumes these values expects actual memory pointers. This will likely lead to a crash if this code path is ever exercised in tests. To make the mock more correct and robust, you should add the buffer's base data pointer to the calculated offset.

Suggested change
def get_page_buffer_meta(self, indices):
"""Mock implementation of get_page_buffer_meta."""
ptr_list = []
element_size_list = []
for idx in indices:
# Create mock pointers and sizes for each page
ptr_list.append(idx * self.page_size * self.kv_buffer.element_size())
element_size_list.append(self.page_size * self.kv_buffer.element_size())
return ptr_list, element_size_list
def get_page_buffer_meta(self, indices):
"""Mock implementation of get_page_buffer_meta."""
ptr_list = []
element_size_list = []
base_ptr = self.kv_buffer.data_ptr()
element_size = self.kv_buffer.element_size()
for idx in indices:
# Create mock pointers and sizes for each page
ptr_list.append(base_ptr + idx * self.page_size * element_size)
element_size_list.append(self.page_size * element_size)
return ptr_list, element_size_list


return MockHostKVCache(buffer), buffer


def test_single_operation():
"""Test the set API with a single key-value pair."""
print("=" * 100)
Expand All @@ -37,8 +60,11 @@ def test_single_operation():
buffer_size = 1024 * 1024 * 16 # 16MB
value_elements = 1024
store = MooncakeStore()
buffer = torch.randn(buffer_size, dtype=torch.float32)
store.register_buffer(buffer)
mock_host_kv_cache, buffer = create_mock_host_kv_cache(buffer_size)

# Register the memory pool host - this is the proper workflow
store.register_mem_pool_host(mock_host_kv_cache)

value_size = value_elements * buffer.element_size()

key = str(uuid.uuid4())
Expand Down Expand Up @@ -75,8 +101,10 @@ def test_batch_operation(config: HiCacheStorageConfig):
value_elements = 256
kv_num = 13
store = MooncakeStore(config)
buffer = torch.randn(buffer_size, dtype=torch.float32)
store.register_buffer(buffer)
mock_host_kv_cache, buffer = create_mock_host_kv_cache(buffer_size)

store.register_mem_pool_host(mock_host_kv_cache)

value_size = value_elements * buffer.element_size()

set_keys, get_keys, exist_keys = generate_batch_query_keys(kv_num, config)
Expand Down
Loading