[HiCache] Optimize HiCache hash generation with bulk token byte conversion (#28287)

This commit is contained in:
huangtingwei
2026-07-01 15:44:23 +08:00
committed by GitHub
parent df0dfbaa45
commit 5e1ccd9320
7 changed files with 675 additions and 180 deletions
+5 -11
View File
@@ -1002,18 +1002,12 @@ class HiCacheController:
storage_query_count = 0 storage_query_count = 0
hash_value = [] hash_value = []
page_hashes = self.get_hash_str(
tokens_to_fetch, last_hash, page_size=self.page_size
)
for start in range( for start in range(0, len(page_hashes), STORAGE_BATCH_SIZE):
0, len(tokens_to_fetch), self.page_size * STORAGE_BATCH_SIZE batch_hashes = page_hashes[start : start + STORAGE_BATCH_SIZE]
):
end = min(start + self.page_size * STORAGE_BATCH_SIZE, len(tokens_to_fetch))
batch_tokens = tokens_to_fetch[start:end]
batch_hashes = []
for i in range(0, len(batch_tokens), self.page_size):
last_hash = self.get_hash_str(
batch_tokens[i : i + self.page_size], last_hash
)
batch_hashes.append(last_hash)
extra_info = HiCacheStorageExtraInfo(prefix_keys=prefix_keys) extra_info = HiCacheStorageExtraInfo(prefix_keys=prefix_keys)
hit_page_num = self.storage_backend.batch_exists(batch_hashes, extra_info) hit_page_num = self.storage_backend.batch_exists(batch_hashes, extra_info)
hash_value.extend(batch_hashes[:hit_page_num]) hash_value.extend(batch_hashes[:hit_page_num])
@@ -0,0 +1,351 @@
#include <Python.h>
#if defined(__AVX2__)
#include <immintrin.h>
#endif
#include <openssl/sha.h>
#include <pybind11/pybind11.h>
#include <algorithm>
#include <array>
#include <cstdint>
#include <cstring>
#include <limits>
#include <stdexcept>
#include <string>
#include <type_traits>
#include <vector>
namespace py = pybind11;
namespace {
constexpr std::size_t kDigestLen = SHA256_DIGEST_LENGTH;
constexpr std::size_t kHexLen = SHA256_DIGEST_LENGTH * 2;
inline std::uint32_t checked_u32(std::uint64_t value) {
if (value > UINT32_MAX) {
throw std::out_of_range("token id does not fit in uint32");
}
return static_cast<std::uint32_t>(value);
}
inline void digest_to_hex_chars(const unsigned char *digest, char *out) {
static constexpr char kHex[] = "0123456789abcdef";
for (std::size_t i = 0; i < kDigestLen; ++i) {
const unsigned char byte = digest[i];
out[i * 2] = kHex[byte >> 4];
out[i * 2 + 1] = kHex[byte & 0x0f];
}
}
std::string digest_to_hex_string(const unsigned char *digest) {
std::string out(kHexLen, '\0');
digest_to_hex_chars(digest, out.data());
return out;
}
std::array<unsigned char, kDigestLen>
parse_prior_digest(py::object prior_digest_obj, bool *has_prior_digest) {
std::array<unsigned char, kDigestLen> prior_digest{};
*has_prior_digest = false;
if (!prior_digest_obj.is_none()) {
std::string prior = prior_digest_obj.cast<std::string>();
if (prior.size() != kDigestLen) {
throw std::invalid_argument("prior_digest must be exactly 32 bytes");
}
std::copy(prior.begin(), prior.end(), prior_digest.begin());
*has_prior_digest = true;
}
return prior_digest;
}
inline void hash_page(const unsigned char *data, std::size_t len,
bool &has_prior_digest,
std::array<unsigned char, kDigestLen> &prior_digest) {
SHA256_CTX ctx;
SHA256_Init(&ctx);
if (has_prior_digest) {
SHA256_Update(&ctx, prior_digest.data(), prior_digest.size());
}
if (len > 0) {
SHA256_Update(&ctx, data, len);
}
SHA256_Final(prior_digest.data(), &ctx);
has_prior_digest = true;
}
template <typename RawToken>
inline void fill_regular_page(const RawToken *raw, std::size_t start,
std::size_t count, std::uint32_t *out) {
for (std::size_t i = 0; i < count; ++i) {
out[i] = checked_u32(raw[start + i]);
}
}
template <>
inline void
fill_regular_page<std::uint32_t>(const std::uint32_t *raw, std::size_t start,
std::size_t count, std::uint32_t *out) {
std::copy(raw + start, raw + start + count, out);
}
template <>
inline void
fill_regular_page<std::uint64_t>(const std::uint64_t *raw, std::size_t start,
std::size_t count, std::uint32_t *out) {
#if defined(__AVX2__)
std::size_t i = 0;
for (; i + 4 <= count; i += 4) {
const __m256i v =
_mm256_loadu_si256(reinterpret_cast<const __m256i *>(raw + start + i));
const __m256i high = _mm256_srli_epi64(v, 32);
if (!_mm256_testz_si256(high, high)) {
throw std::out_of_range("token id does not fit in uint32");
}
const __m256i low_pairs = _mm256_shuffle_epi32(v, _MM_SHUFFLE(2, 0, 2, 0));
const __m128i lane0 = _mm256_castsi256_si128(low_pairs);
const __m128i lane1 = _mm256_extracti128_si256(low_pairs, 1);
_mm_storel_epi64(reinterpret_cast<__m128i *>(out + i), lane0);
_mm_storel_epi64(reinterpret_cast<__m128i *>(out + i + 2), lane1);
}
for (; i < count; ++i) {
out[i] = checked_u32(raw[start + i]);
}
#else
for (std::size_t i = 0; i < count; ++i) {
out[i] = checked_u32(raw[start + i]);
}
#endif
}
template <typename RawToken>
inline void fill_bigram_page(const RawToken *raw, std::size_t start,
std::size_t count, std::uint32_t *out) {
std::uint32_t prev = checked_u32(raw[start]);
for (std::size_t i = 0; i < count; ++i) {
const std::uint32_t next = checked_u32(raw[start + i + 1]);
out[i * 2] = prev;
out[i * 2 + 1] = next;
prev = next;
}
}
template <>
inline void
fill_bigram_page<std::uint32_t>(const std::uint32_t *raw, std::size_t start,
std::size_t count, std::uint32_t *out) {
std::uint32_t prev = raw[start];
for (std::size_t i = 0; i < count; ++i) {
const std::uint32_t next = raw[start + i + 1];
out[i * 2] = prev;
out[i * 2 + 1] = next;
prev = next;
}
}
template <>
inline void
fill_bigram_page<std::uint64_t>(const std::uint64_t *raw, std::size_t start,
std::size_t count, std::uint32_t *out) {
std::uint32_t prev = checked_u32(raw[start]);
for (std::size_t i = 0; i < count; ++i) {
const std::uint32_t next = checked_u32(raw[start + i + 1]);
out[i * 2] = prev;
out[i * 2 + 1] = next;
prev = next;
}
}
struct RawTokenBuffer {
py::buffer_info info;
std::size_t logical_len;
bool is_bigram;
};
RawTokenBuffer get_raw_token_buffer(const py::buffer &raw_tokens,
std::size_t logical_len,
std::size_t unit_width, bool is_bigram) {
py::buffer_info info = raw_tokens.request();
if (info.ndim != 1) {
throw std::invalid_argument("raw_tokens must be a one-dimensional buffer");
}
if (info.itemsize != 4 && info.itemsize != 8) {
throw std::invalid_argument("raw_tokens itemsize must be 4 or 8 bytes");
}
const std::size_t need_raw_tokens =
is_bigram && logical_len > 0 ? logical_len + 1 : logical_len * unit_width;
if (static_cast<std::size_t>(info.size) < need_raw_tokens) {
throw std::invalid_argument("raw_tokens is shorter than logical_len");
}
return RawTokenBuffer{std::move(info), logical_len, is_bigram};
}
template <typename RawToken>
void hash_pages_to_hex_blob(const RawToken *raw, std::size_t logical_len,
std::size_t page_size, std::size_t unit_width,
bool is_bigram, bool has_prior_digest,
std::array<unsigned char, kDigestLen> prior_digest,
std::string &hex_blob) {
if (page_size == 0) {
throw std::invalid_argument("page_size must be positive");
}
if (is_bigram) {
unit_width = 2;
}
const bool can_hash_raw_bytes =
std::is_same_v<RawToken, std::uint32_t> && !is_bigram;
std::vector<std::uint32_t> page_words;
if (!can_hash_raw_bytes) {
page_words.resize(page_size * unit_width);
}
for (std::size_t start = 0, page_idx = 0; start < logical_len;
start += page_size, ++page_idx) {
const std::size_t page_units = std::min(page_size, logical_len - start);
const std::size_t page_bytes =
page_units * unit_width * sizeof(std::uint32_t);
const unsigned char *bytes = nullptr;
if (can_hash_raw_bytes) {
bytes = reinterpret_cast<const unsigned char *>(raw + start * unit_width);
} else {
if (is_bigram) {
fill_bigram_page(raw, start, page_units, page_words.data());
} else {
fill_regular_page(raw, start * unit_width, page_units * unit_width,
page_words.data());
}
bytes = reinterpret_cast<const unsigned char *>(page_words.data());
}
hash_page(bytes, page_bytes, has_prior_digest, prior_digest);
digest_to_hex_chars(prior_digest.data(),
hex_blob.data() + page_idx * kHexLen);
}
}
template <typename RawToken>
std::string hash_all(const RawToken *raw, std::size_t logical_len,
std::size_t unit_width, bool is_bigram,
bool has_prior_digest,
std::array<unsigned char, kDigestLen> prior_digest) {
if (is_bigram) {
unit_width = 2;
}
const bool can_hash_raw_bytes =
std::is_same_v<RawToken, std::uint32_t> && !is_bigram;
const unsigned char *bytes = nullptr;
std::size_t num_bytes = logical_len * unit_width * sizeof(std::uint32_t);
std::vector<std::uint32_t> words;
if (can_hash_raw_bytes) {
bytes = reinterpret_cast<const unsigned char *>(raw);
} else {
words.resize(logical_len * unit_width);
if (logical_len > 0) {
if (is_bigram) {
fill_bigram_page(raw, 0, logical_len, words.data());
} else {
fill_regular_page(raw, 0, logical_len * unit_width, words.data());
}
}
bytes = reinterpret_cast<const unsigned char *>(words.data());
}
hash_page(bytes, num_bytes, has_prior_digest, prior_digest);
return digest_to_hex_string(prior_digest.data());
}
py::object hex_blob_to_pylist(const std::string &hex_blob) {
const std::size_t num_pages = hex_blob.size() / kHexLen;
if (num_pages >
static_cast<std::size_t>(std::numeric_limits<Py_ssize_t>::max())) {
throw std::overflow_error("too many hash pages");
}
PyObject *raw_list = PyList_New(static_cast<Py_ssize_t>(num_pages));
if (raw_list == nullptr) {
throw py::error_already_set();
}
py::object list = py::reinterpret_steal<py::object>(raw_list);
const char *data = hex_blob.data();
for (std::size_t i = 0; i < num_pages; ++i) {
PyObject *item = PyUnicode_FromStringAndSize(
data + i * kHexLen, static_cast<Py_ssize_t>(kHexLen));
if (item == nullptr) {
throw py::error_already_set();
}
PyList_SET_ITEM(raw_list, static_cast<Py_ssize_t>(i), item);
}
return list;
}
std::string hash_str(const py::buffer &raw_tokens, std::size_t logical_len,
std::size_t unit_width, bool is_bigram,
py::object prior_digest_obj) {
RawTokenBuffer buffer =
get_raw_token_buffer(raw_tokens, logical_len, unit_width, is_bigram);
bool has_prior_digest = false;
auto prior_digest = parse_prior_digest(prior_digest_obj, &has_prior_digest);
py::gil_scoped_release release;
if (buffer.info.itemsize == 4) {
return hash_all(static_cast<const std::uint32_t *>(buffer.info.ptr),
logical_len, unit_width, is_bigram, has_prior_digest,
prior_digest);
}
return hash_all(static_cast<const std::uint64_t *>(buffer.info.ptr),
logical_len, unit_width, is_bigram, has_prior_digest,
prior_digest);
}
py::object pages_hashes(const py::buffer &raw_tokens, std::size_t logical_len,
std::size_t page_size, std::size_t unit_width,
bool is_bigram, py::object prior_digest_obj) {
RawTokenBuffer buffer =
get_raw_token_buffer(raw_tokens, logical_len, unit_width, is_bigram);
const std::size_t num_pages =
page_size == 0 ? 0 : (logical_len + page_size - 1) / page_size;
bool has_prior_digest = false;
auto prior_digest = parse_prior_digest(prior_digest_obj, &has_prior_digest);
std::string hex_blob(num_pages * kHexLen, '\0');
{
py::gil_scoped_release release;
if (buffer.info.itemsize == 4) {
hash_pages_to_hex_blob(
static_cast<const std::uint32_t *>(buffer.info.ptr), logical_len,
page_size, unit_width, is_bigram, has_prior_digest, prior_digest,
hex_blob);
} else {
hash_pages_to_hex_blob(
static_cast<const std::uint64_t *>(buffer.info.ptr), logical_len,
page_size, unit_width, is_bigram, has_prior_digest, prior_digest,
hex_blob);
}
}
return hex_blob_to_pylist(hex_blob);
}
py::object get_hash(const py::buffer &raw_tokens, std::size_t logical_len,
std::size_t unit_width, bool is_bigram,
py::object prior_digest_obj, py::object page_size_obj) {
if (page_size_obj.is_none()) {
return py::cast(hash_str(raw_tokens, logical_len, unit_width, is_bigram,
prior_digest_obj));
}
return pages_hashes(raw_tokens, logical_len,
page_size_obj.cast<std::size_t>(), unit_width, is_bigram,
prior_digest_obj);
}
} // namespace
PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) {
m.def("get_hash", &get_hash, py::arg("raw_tokens"), py::arg("logical_len"),
py::arg("unit_width"), py::arg("is_bigram"), py::arg("prior_digest"),
py::arg("page_size") = py::none());
}
@@ -0,0 +1,86 @@
import os
import platform
import sys
from array import array
from functools import lru_cache
from typing import Any, Optional
def _cpu_supports_avx2() -> bool:
if platform.machine().lower() not in ("x86_64", "amd64"):
return False
try:
with open("/proc/cpuinfo", "r", encoding="utf-8", errors="ignore") as f:
return "avx2" in f.read().lower()
except OSError:
return False
@lru_cache(maxsize=1)
def _load_native_hash_module() -> Any:
if sys.byteorder != "little" or not sys.platform.startswith("linux"):
raise RuntimeError(
"HiCache native hash is only supported on little-endian Linux"
)
try:
from torch.utils.cpp_extension import load
abs_path = os.path.dirname(os.path.abspath(__file__))
extra_cflags = ["-O3", "-std=c++17", "-DNDEBUG"]
if _cpu_supports_avx2():
extra_cflags.append("-mavx2")
return load(
name="hicache_hash_cpp",
sources=[f"{abs_path}/hash_binding.cpp"],
extra_cflags=extra_cflags,
extra_ldflags=["-lcrypto"],
with_cuda=False,
verbose=False,
)
except Exception as exc:
raise RuntimeError("Failed to load HiCache native hash extension") from exc
def _native_hash_input(token_ids: Any) -> tuple[array, int, int, bool]:
raw_token_ids = getattr(token_ids, "raw_token_ids", None)
raw = (
raw_token_ids()
if raw_token_ids is not None
else getattr(token_ids, "token_ids", token_ids)
)
logical_len = len(token_ids)
is_bigram = getattr(token_ids, "is_bigram", False)
if isinstance(raw, array) and raw.typecode in ("I", "q", "Q", "L"):
if is_bigram and logical_len > 0 and len(raw) < logical_len + 1:
raise ValueError("bigram token buffer is shorter than logical length")
return raw, logical_len, 2 if is_bigram else 1, is_bigram
if is_bigram:
return array("I", raw[: logical_len + 1]), logical_len, 2, is_bigram
if logical_len == 0:
return array("I"), logical_len, 1, is_bigram
first_token = raw[0]
if isinstance(first_token, tuple):
unit_width = len(first_token)
return (
array("I", (elem for token in raw[:logical_len] for elem in token)),
logical_len,
unit_width,
is_bigram,
)
return array("I", raw[:logical_len]), logical_len, 1, is_bigram
def get_native_hash(
token_ids: Any, prior_digest: Optional[bytes], page_size: Optional[int] = None
) -> str | list[str]:
raw, logical_len, unit_width, is_bigram = _native_hash_input(token_ids)
return _load_native_hash_module().get_hash(
raw, logical_len, unit_width, is_bigram, prior_digest, page_size
)
@@ -585,13 +585,9 @@ class HybridCacheController(BaseHiCacheController):
return operation.id return operation.id
def _storage_hit_query(self, operation) -> tuple[list[str], int]: def _storage_hit_query(self, operation) -> tuple[list[str], int]:
last_hash = operation.last_hash hash_value = self.get_hash_str(
hash_value = [] operation.token_ids, operation.last_hash, page_size=self.page_size
for start in range(0, len(operation.token_ids), self.page_size): )
last_hash = self.get_hash_str(
operation.token_ids[start : start + self.page_size], last_hash
)
hash_value.append(last_hash)
extra_info = HiCacheStorageExtraInfo( extra_info = HiCacheStorageExtraInfo(
prefix_keys=operation.prefix_keys.copy() if operation.prefix_keys else None prefix_keys=operation.prefix_keys.copy() if operation.prefix_keys else None
+8 -14
View File
@@ -21,7 +21,6 @@ limitations under the License.
The radix tree data structure for managing the KV cache. The radix tree data structure for managing the KV cache.
""" """
import hashlib
import heapq import heapq
import logging import logging
import sys import sys
@@ -48,7 +47,11 @@ from sglang.srt.mem_cache.base_prefix_cache import (
) )
from sglang.srt.mem_cache.events import KVCacheEventMixin from sglang.srt.mem_cache.events import KVCacheEventMixin
from sglang.srt.mem_cache.session_radix_cache import SessionRadixCacheMixin from sglang.srt.mem_cache.session_radix_cache import SessionRadixCacheMixin
from sglang.srt.mem_cache.utils import get_eviction_strategy, split_node_hash_value from sglang.srt.mem_cache.utils import (
get_eviction_strategy,
get_hash_str,
split_node_hash_value,
)
if TYPE_CHECKING: if TYPE_CHECKING:
from sglang.srt.managers.schedule_batch import Req from sglang.srt.managers.schedule_batch import Req
@@ -206,18 +209,9 @@ class RadixKey:
def hash_page(self, start: int, end: int, prior_hash: Optional[str] = None) -> str: def hash_page(self, start: int, end: int, prior_hash: Optional[str] = None) -> str:
"""SHA256 for logical units [start, end); bigram mode feeds overlapping (t_i, t_{i+1}) byte pairs.""" """SHA256 for logical units [start, end); bigram mode feeds overlapping (t_i, t_{i+1}) byte pairs."""
hasher = hashlib.sha256() hash_value = get_hash_str(self[start:end], prior_hash)
if prior_hash: assert isinstance(hash_value, str)
hasher.update(bytes.fromhex(prior_hash)) return hash_value
t = self.token_ids
if self.is_bigram:
for j in range(start, end):
hasher.update(t[j].to_bytes(4, byteorder="little", signed=False))
hasher.update(t[j + 1].to_bytes(4, byteorder="little", signed=False))
else:
for j in range(start, end):
hasher.update(t[j].to_bytes(4, byteorder="little", signed=False))
return hasher.hexdigest()
class TreeNode: class TreeNode:
+10 -27
View File
@@ -13,10 +13,10 @@
# ============================================================================== # ==============================================================================
"""Common utilities.""" """Common utilities."""
import hashlib
from typing import Any, Callable, List, Optional, Tuple from typing import Any, Callable, List, Optional, Tuple
from sglang.srt.environ import envs from sglang.srt.environ import envs
from sglang.srt.mem_cache.cpp_utils.native_hash import get_native_hash
from sglang.srt.mem_cache.evict_policy import ( from sglang.srt.mem_cache.evict_policy import (
EvictionStrategy, EvictionStrategy,
FIFOStrategy, FIFOStrategy,
@@ -103,22 +103,13 @@ def maybe_init_custom_mem_pool(
return False, None, None return False, None, None
def get_hash_str(token_ids: List[int], prior_hash: Optional[str] = None) -> str: def get_hash_str(
hasher = hashlib.sha256() token_ids: List[int],
prior_hash: Optional[str] = None,
if prior_hash: page_size: Optional[int] = None,
hasher.update(bytes.fromhex(prior_hash)) ) -> str | List[str]:
prior_digest = bytes.fromhex(prior_hash) if prior_hash else None
for t in token_ids: return get_native_hash(token_ids, prior_digest, page_size)
if isinstance(t, tuple):
# EAGLE bigram mode: hash both elements to uniquely identify the bigram
for elem in t:
hasher.update(elem.to_bytes(4, byteorder="little", signed=False))
else:
# Regular mode: single integer token
hasher.update(t.to_bytes(4, byteorder="little", signed=False))
return hasher.hexdigest()
def hash_str_to_int64(hash_str: str) -> int: def hash_str_to_int64(hash_str: str) -> int:
@@ -134,21 +125,13 @@ def hash_str_to_int64(hash_str: str) -> int:
def compute_node_hash_values(node: Any, page_size: int) -> List[str]: def compute_node_hash_values(node: Any, page_size: int) -> List[str]:
"""Compute SHA256-based hash values for position-aware KV block IDs.""" """Compute SHA256-based hash values for position-aware KV block IDs."""
hash_values = []
parent_hash = None parent_hash = None
if node.parent is not None and node.parent.hash_value is not None: if node.parent is not None and node.parent.hash_value is not None:
if len(node.parent.key) > 0 and len(node.parent.hash_value) > 0: if len(node.parent.key) > 0 and len(node.parent.hash_value) > 0:
parent_hash = node.parent.hash_value[-1] parent_hash = node.parent.hash_value[-1]
logical_len = len(node.key) hash_values = get_hash_str(node.key, parent_hash, page_size=page_size)
for start in range(0, logical_len, page_size): assert isinstance(hash_values, list)
end = min(start + page_size, logical_len)
if end <= start:
continue
hash_val = node.key.hash_page(start, end, parent_hash)
hash_values.append(hash_val)
parent_hash = hash_val
return hash_values return hash_values
@@ -1,7 +1,10 @@
"""Unit tests for mem_cache/utils.py — no server, no model loading.""" """Unit tests for mem_cache/utils.py — no server, no model loading."""
import hashlib import hashlib
import sys
import types
import unittest import unittest
from array import array
from unittest.mock import MagicMock, patch from unittest.mock import MagicMock, patch
from sglang.srt.mem_cache.evict_policy import ( from sglang.srt.mem_cache.evict_policy import (
@@ -22,12 +25,117 @@ from sglang.srt.mem_cache.utils import (
split_node_hash_value, split_node_hash_value,
) )
from sglang.test.ci.ci_register import register_cpu_ci from sglang.test.ci.ci_register import register_cpu_ci
from sglang.test.test_utils import CustomTestCase
register_cpu_ci(est_time=8, suite="base-a-test-cpu") register_cpu_ci(est_time=8, suite="base-a-test-cpu")
class TestGetEvictionStrategy(CustomTestCase): def _legacy_get_hash_str(token_ids, prior_hash=None):
hasher = hashlib.sha256()
if prior_hash:
hasher.update(bytes.fromhex(prior_hash))
for t in token_ids:
if isinstance(t, tuple):
for elem in t:
hasher.update(elem.to_bytes(4, byteorder="little", signed=False))
else:
hasher.update(t.to_bytes(4, byteorder="little", signed=False))
return hasher.hexdigest()
def _legacy_page_hashes(key, page_size, prior_hash=None):
hashes = []
running_hash = prior_hash
for start in range(0, len(key), page_size):
running_hash = _legacy_get_hash_str(
key[start : start + page_size], running_hash
)
hashes.append(running_hash)
return hashes
class _HashKey:
def __init__(self, token_ids, is_bigram=False):
self.token_ids = token_ids
self.is_bigram = is_bigram
def __len__(self):
if self.is_bigram:
return max(0, len(self.token_ids) - 1)
return len(self.token_ids)
def __getitem__(self, index):
if isinstance(index, slice):
start = index.start or 0
stop = index.stop if index.stop is not None else len(self)
if self.is_bigram:
return _HashKey(self.token_ids[start : stop + 1], is_bigram=True)
return _HashKey(self.token_ids[start:stop])
if self.is_bigram:
return (self.token_ids[index], self.token_ids[index + 1])
return self.token_ids[index]
def raw_token_ids(self):
return self.token_ids
def hash_page(self, start, end, prior_hash=None):
return _legacy_get_hash_str(self[start:end], prior_hash)
def _single_hash_compatibility_cases():
prior_hash = _legacy_get_hash_str([7, 8, 9])
return [
("empty_list", [], None),
("plain_list", [1, 2, 3, 4, 5], None),
("array_q", _HashKey(array("q", range(1, 258))), None),
("array_i_with_prior", _HashKey(array("I", range(1, 258))), prior_hash),
(
"tuple_bigram_with_prior",
[(10, 20), (20, 30), (30, 40), (40, 50)],
prior_hash,
),
(
"eagle_bigram_with_prior",
_HashKey(
array("q", ((i * 2654435761) & 0x00FFFFFF for i in range(258))),
is_bigram=True,
),
prior_hash,
),
("empty_bigram", _HashKey(array("q"), is_bigram=True), None),
]
def _page_hash_compatibility_cases():
prior_hash = _legacy_get_hash_str([7, 8, 9])
return [
("empty_list", [], 8, None),
("empty_bigram", _HashKey(array("q"), is_bigram=True), 8, None),
("array_q_page_64", _HashKey(array("q", range(1, 258))), 64, None),
(
"array_i_page_64_with_prior",
_HashKey(array("I", range(1, 258))),
64,
prior_hash,
),
(
"eagle_bigram_page_64_with_prior",
_HashKey(
array("q", ((i * 2654435761) & 0x00FFFFFF for i in range(258))),
is_bigram=True,
),
64,
prior_hash,
),
(
"eagle_bigram_page_1_with_prior",
_HashKey(array("q", [11, 22, 33, 44, 55]), is_bigram=True),
1,
prior_hash,
),
]
class TestGetEvictionStrategy(unittest.TestCase):
def test_lru(self): def test_lru(self):
self.assertIsInstance(get_eviction_strategy("lru"), LRUStrategy) self.assertIsInstance(get_eviction_strategy("lru"), LRUStrategy)
@@ -73,7 +181,7 @@ class TestGetEvictionStrategy(CustomTestCase):
self.assertIsNot(s1, s2) self.assertIsNot(s1, s2)
class TestMaybeInitCustomMemPool(CustomTestCase): class TestMaybeInitCustomMemPool(unittest.TestCase):
@patch("sglang.srt.mem_cache.utils.envs.SGLANG_MOONCAKE_CUSTOM_MEM_POOL.get") @patch("sglang.srt.mem_cache.utils.envs.SGLANG_MOONCAKE_CUSTOM_MEM_POOL.get")
def test_disabled_by_default(self, mock_env_get): def test_disabled_by_default(self, mock_env_get):
mock_env_get.return_value = None mock_env_get.return_value = None
@@ -83,69 +191,76 @@ class TestMaybeInitCustomMemPool(CustomTestCase):
self.assertIsNone(pool_type) self.assertIsNone(pool_type)
@patch("sglang.srt.mem_cache.utils.envs.SGLANG_MOONCAKE_CUSTOM_MEM_POOL.get") @patch("sglang.srt.mem_cache.utils.envs.SGLANG_MOONCAKE_CUSTOM_MEM_POOL.get")
@patch("sglang.srt.disaggregation.mooncake.utils.init_mooncake_custom_mem_pool") def test_enabled_via_env(self, mock_env_get):
def test_enabled_via_env(self, mock_init, mock_env_get):
mock_env_get.return_value = "enabled" mock_env_get.return_value = "enabled"
mock_init = MagicMock()
mock_init.return_value = (True, "mock_pool_instance", "mooncake") mock_init.return_value = (True, "mock_pool_instance", "mooncake")
enabled, pool, pool_type = maybe_init_custom_mem_pool("cuda:0") mooncake_pkg = types.ModuleType("sglang.srt.disaggregation.mooncake")
mooncake_utils = types.ModuleType("sglang.srt.disaggregation.mooncake.utils")
mooncake_utils.init_mooncake_custom_mem_pool = mock_init
with patch.dict(
sys.modules,
{
"sglang.srt.disaggregation.mooncake": mooncake_pkg,
"sglang.srt.disaggregation.mooncake.utils": mooncake_utils,
},
):
enabled, pool, pool_type = maybe_init_custom_mem_pool("cuda:0")
self.assertTrue(enabled) self.assertTrue(enabled)
self.assertEqual(pool, "mock_pool_instance") self.assertEqual(pool, "mock_pool_instance")
self.assertEqual(pool_type, "mooncake") self.assertEqual(pool_type, "mooncake")
mock_init.assert_called_once_with("cuda:0") mock_init.assert_called_once_with("cuda:0")
class TestGetHashStr(CustomTestCase): class TestGetHashStr(unittest.TestCase):
def test_empty_list(self): def test_hash_str_matches_pre_optimization_per_token_loop(self):
result = get_hash_str([]) for name, tokens, prior_hash in _single_hash_compatibility_cases():
self.assertIsInstance(result, str) with self.subTest(name=name):
self.assertEqual(len(result), 64) self.assertEqual(
expected = hashlib.sha256().hexdigest() get_hash_str(tokens, prior_hash),
self.assertEqual(result, expected) _legacy_get_hash_str(tokens, prior_hash),
)
def test_page_hashes_match_pre_optimization_per_token_loop(self):
for name, tokens, page_size, prior_hash in _page_hash_compatibility_cases():
with self.subTest(name=name):
self.assertEqual(
get_hash_str(tokens, prior_hash, page_size=page_size),
_legacy_page_hashes(tokens, page_size, prior_hash),
)
def test_hash_properties(self):
self.assertEqual(get_hash_str([]), hashlib.sha256().hexdigest())
def test_different_sequence_different_hash(self):
h1 = get_hash_str([1, 2, 3]) h1 = get_hash_str([1, 2, 3])
h2 = get_hash_str([3, 2, 1]) h2 = get_hash_str([3, 2, 1])
self.assertNotEqual(h1, h2) self.assertNotEqual(h1, h2)
def test_different_values_different_hash(self): self.assertNotEqual(get_hash_str([100]), get_hash_str([200]))
h1 = get_hash_str([100]) self.assertEqual(get_hash_str([1, 2]), get_hash_str([(1, 2)]))
h2 = get_hash_str([200]) self.assertNotEqual(get_hash_str([(1, 2)]), get_hash_str([(2, 1)]))
self.assertNotEqual(h1, h2)
def test_bigram_vs_flat_token_equivalent(self):
h_flat = get_hash_str([1, 2])
h_bigram = get_hash_str([(1, 2)])
self.assertEqual(h_flat, h_bigram)
def test_bigram_order_matters(self):
h1 = get_hash_str([(1, 2)])
h2 = get_hash_str([(2, 1)])
self.assertNotEqual(h1, h2)
def test_prior_hash_chaining(self):
chained = get_hash_str([3, 4], prior_hash=get_hash_str([1, 2])) chained = get_hash_str([3, 4], prior_hash=get_hash_str([1, 2]))
# prior_hash must fold into the digest, so chaining differs from
# hashing [3, 4] alone...
self.assertNotEqual(chained, get_hash_str([3, 4])) self.assertNotEqual(chained, get_hash_str([3, 4]))
# ...and a different prior_hash must yield a different chained digest.
self.assertNotEqual( self.assertNotEqual(
chained, get_hash_str([3, 4], prior_hash=get_hash_str([9, 9])) chained, get_hash_str([3, 4], prior_hash=get_hash_str([9, 9]))
) )
def test_prior_hash_single_step(self):
step1 = get_hash_str([1])
step2 = get_hash_str([2], prior_hash=step1)
direct = get_hash_str([1, 2])
self.assertNotEqual(step2, direct)
def test_returns_64_char_hex(self):
for tokens in [[], [1], [1, 2, 3], [(1, 2)], [1, 2, 3, 4, 5]]: for tokens in [[], [1], [1, 2, 3], [(1, 2)], [1, 2, 3, 4, 5]]:
result = get_hash_str(tokens) with self.subTest(tokens=tokens):
self.assertRegex(result, r"^[0-9a-f]{64}$") self.assertRegex(get_hash_str(tokens), r"^[0-9a-f]{64}$")
def test_hash_key_hash_page_matches_get_hash_str(self):
key = _HashKey(array("q", [1, 2, 3, 4, 5, 6]), is_bigram=True)
prior_hash = get_hash_str([(9, 10)])
self.assertEqual(
key.hash_page(1, 4, prior_hash),
get_hash_str(key[1:4], prior_hash),
)
class TestHashStrToInt64(CustomTestCase): class TestHashStrToInt64(unittest.TestCase):
def test_zero_hash(self): def test_zero_hash(self):
result = hash_str_to_int64("0" * 64) result = hash_str_to_int64("0" * 64)
self.assertEqual(result, 0) self.assertEqual(result, 0)
@@ -178,96 +293,72 @@ class TestHashStrToInt64(CustomTestCase):
self.assertTrue(-(2**63) <= int64_val < 2**63) self.assertTrue(-(2**63) <= int64_val < 2**63)
class TestComputeNodeHashValues(CustomTestCase): class TestComputeNodeHashValues(unittest.TestCase):
def setUp(self): def _make_node(self, key, parent=None, parent_hash_values=None):
def mock_hash_page(start, end, parent_hash):
parts = [f"p{start}-{end}"]
if parent_hash is not None:
parts.append(parent_hash)
return "-".join(parts)
self.mock_hash_page = mock_hash_page
def _make_node(self, key_len, parent=None, parent_hash_values=None):
node = MagicMock() node = MagicMock()
node.key.__len__.return_value = key_len node.key = key
node.key.hash_page = self.mock_hash_page
node.parent = parent node.parent = parent
if parent is not None: if parent is not None:
parent.hash_value = parent_hash_values parent.hash_value = parent_hash_values
return node return node
def test_single_page_root(self): def test_root_node_hashes_match_legacy_page_hashes(self):
node = self._make_node(key_len=3) cases = [
result = compute_node_hash_values(node, page_size=16) ("single_page", _HashKey(array("q", [1, 2, 3])), 16),
self.assertEqual(len(result), 1) ("multiple_pages", _HashKey(array("q", range(1, 31))), 16),
self.assertIn("p0-3", result[0]) ("page_aligned_boundary", _HashKey(array("q", range(1, 33))), 8),
("shorter_than_page", _HashKey(array("q", [1, 2, 3, 4, 5])), 16),
]
def test_multiple_pages(self): for name, key, page_size in cases:
node = self._make_node(key_len=30) with self.subTest(name=name):
result = compute_node_hash_values(node, page_size=16) node = self._make_node(key)
self.assertEqual(len(result), 2) self.assertEqual(
self.assertIn("p0-16", result[0]) compute_node_hash_values(node, page_size=page_size),
self.assertIn("p16-30", result[1]) _legacy_page_hashes(key, page_size=page_size),
)
def test_page_aligned_boundary(self): def test_parent_hash_is_used_only_when_parent_has_nonempty_key_and_hash(self):
node = self._make_node(key_len=32)
result = compute_node_hash_values(node, page_size=8)
self.assertEqual(len(result), 4)
self.assertIn("p24-32", result[3])
def test_key_shorter_than_page_size(self):
node = self._make_node(key_len=5)
result = compute_node_hash_values(node, page_size=16)
self.assertEqual(result, ["p0-5"])
def test_chained_parent_hash(self):
parent = MagicMock() parent = MagicMock()
parent.key.__len__.return_value = 8 parent.key = _HashKey(array("q", range(1, 17)))
parent.hash_value = ["parent_hash_0", "parent_hash_1"] parent.hash_value = _legacy_page_hashes(parent.key, page_size=8)
parent.key.hash_page = self.mock_hash_page child_key = _HashKey(array("q", range(101, 109)))
cases = [
("valid_parent_hash", parent, parent.hash_value, parent.hash_value[-1]),
(
"empty_parent_key",
self._make_node(_HashKey(array("q"))),
[get_hash_str([1, 2, 3])],
None,
),
(
"empty_parent_hash",
self._make_node(_HashKey(array("q", range(1, 9)))),
[],
None,
),
(
"none_parent_hash",
self._make_node(_HashKey(array("q", range(1, 9)))),
None,
None,
),
]
child = self._make_node( for name, parent, parent_hash_values, expected_prior in cases:
key_len=16, parent=parent, parent_hash_values=parent.hash_value with self.subTest(name=name):
) child = self._make_node(
result = compute_node_hash_values(child, page_size=8) child_key, parent=parent, parent_hash_values=parent_hash_values
self.assertEqual(result, ["p0-8-parent_hash_1", "p8-16-p0-8-parent_hash_1"]) )
self.assertEqual(
def test_parent_with_empty_key(self): compute_node_hash_values(child, page_size=8),
parent = MagicMock() _legacy_page_hashes(
parent.key.__len__.return_value = 0 child_key, page_size=8, prior_hash=expected_prior
parent.hash_value = ["some_hash"] ),
parent.key.hash_page = self.mock_hash_page )
child = self._make_node(
key_len=8, parent=parent, parent_hash_values=parent.hash_value
)
result = compute_node_hash_values(child, page_size=8)
self.assertEqual(len(result), 1)
self.assertNotIn("some_hash", result[0])
def test_parent_without_hash_value(self):
parent = MagicMock()
parent.key.__len__.return_value = 8
parent.hash_value = []
parent.key.hash_page = self.mock_hash_page
child = self._make_node(key_len=8, parent=parent, parent_hash_values=[])
result = compute_node_hash_values(child, page_size=8)
self.assertEqual(result, ["p0-8"])
def test_parent_with_none_hash_value(self):
parent = MagicMock()
parent.key.__len__.return_value = 8
parent.hash_value = None
parent.key.hash_page = self.mock_hash_page
child = self._make_node(key_len=8, parent=parent, parent_hash_values=None)
result = compute_node_hash_values(child, page_size=8)
self.assertEqual(result, ["p0-8"])
class TestSplitNodeHashValue(CustomTestCase): class TestSplitNodeHashValue(unittest.TestCase):
def test_none_input_returns_none_tuple(self): def test_none_input_returns_none_tuple(self):
result = split_node_hash_value(None, 10, 4) result = split_node_hash_value(None, 10, 4)
self.assertEqual(result, (None, None)) self.assertEqual(result, (None, None))