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

python/sglang/srt/managers/cache_controller.py

@@ -1002,18 +1002,12 @@ class HiCacheController:
 
         storage_query_count = 0
         hash_value = []
+        page_hashes = self.get_hash_str(
+            tokens_to_fetch, last_hash, page_size=self.page_size
+        )
 
-        for start in range(
-            0, len(tokens_to_fetch), self.page_size * 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)
+        for start in range(0, len(page_hashes), STORAGE_BATCH_SIZE):
+            batch_hashes = page_hashes[start : start + STORAGE_BATCH_SIZE]
             extra_info = HiCacheStorageExtraInfo(prefix_keys=prefix_keys)
             hit_page_num = self.storage_backend.batch_exists(batch_hashes, extra_info)
             hash_value.extend(batch_hashes[:hit_page_num])

python/sglang/srt/mem_cache/cpp_utils/hash_binding.cpp

@@ -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());
+}

python/sglang/srt/mem_cache/cpp_utils/native_hash.py

@@ -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
+    )

python/sglang/srt/mem_cache/hybrid_cache/hybrid_cache_controller.py

@@ -585,13 +585,9 @@ class HybridCacheController(BaseHiCacheController):
         return operation.id
 
     def _storage_hit_query(self, operation) -> tuple[list[str], int]:
-        last_hash = operation.last_hash
-        hash_value = []
-        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)
+        hash_value = self.get_hash_str(
+            operation.token_ids, operation.last_hash, page_size=self.page_size
+        )
 
         extra_info = HiCacheStorageExtraInfo(
             prefix_keys=operation.prefix_keys.copy() if operation.prefix_keys else None

python/sglang/srt/mem_cache/radix_cache.py

@@ -21,7 +21,6 @@ limitations under the License.
 The radix tree data structure for managing the KV cache.
 """
 
-import hashlib
 import heapq
 import logging
 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.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:
     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:
         """SHA256 for logical units [start, end); bigram mode feeds overlapping (t_i, t_{i+1}) byte pairs."""
-        hasher = hashlib.sha256()
-        if prior_hash:
-            hasher.update(bytes.fromhex(prior_hash))
-        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()
+        hash_value = get_hash_str(self[start:end], prior_hash)
+        assert isinstance(hash_value, str)
+        return hash_value
 
 
 class TreeNode:

python/sglang/srt/mem_cache/utils.py

@@ -13,10 +13,10 @@
 # ==============================================================================
 """Common utilities."""
 
-import hashlib
 from typing import Any, Callable, List, Optional, Tuple
 
 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 (
     EvictionStrategy,
     FIFOStrategy,
@@ -103,22 +103,13 @@ def maybe_init_custom_mem_pool(
         return False, None, None
 
 
-def get_hash_str(token_ids: List[int], prior_hash: Optional[str] = None) -> str:
-    hasher = hashlib.sha256()
-
-    if prior_hash:
-        hasher.update(bytes.fromhex(prior_hash))
-
-    for t in token_ids:
-        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 get_hash_str(
+    token_ids: List[int],
+    prior_hash: Optional[str] = None,
+    page_size: Optional[int] = None,
+) -> str | List[str]:
+    prior_digest = bytes.fromhex(prior_hash) if prior_hash else None
+    return get_native_hash(token_ids, prior_digest, page_size)
 
 
 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]:
     """Compute SHA256-based hash values for position-aware KV block IDs."""
-    hash_values = []
-
     parent_hash = 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:
             parent_hash = node.parent.hash_value[-1]
 
-    logical_len = len(node.key)
-    for start in range(0, logical_len, page_size):
-        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
+    hash_values = get_hash_str(node.key, parent_hash, page_size=page_size)
+    assert isinstance(hash_values, list)
     return hash_values