diff --git a/python/sglang/srt/environ.py b/python/sglang/srt/environ.py index 2e26548de..3a7f612fa 100644 --- a/python/sglang/srt/environ.py +++ b/python/sglang/srt/environ.py @@ -215,6 +215,7 @@ class Envs: SGLANG_CP_SHARED_KV_MLA_PREFETCH_MIN_PREFIX_PAGES = EnvInt(-1) SGLANG_CP_DRAFT_SHARED_KV = EnvBool(False) SGLANG_CP_DRAFT_SHARED_KV_DEBUG = EnvBool(False) + SGLANG_DISABLE_TAI_BIGRAM = EnvBool(False) SGLANG_TEST_REQUEST_TIME_STATS = EnvBool(False) SGLANG_DISABLE_TP_MEMORY_INBALANCE_CHECK = EnvBool(False) SGLANG_SIMULATE_ACC_LEN = EnvFloat(-1) diff --git a/python/sglang/srt/mem_cache/utils.py b/python/sglang/srt/mem_cache/utils.py index 3aec3dd89..7b9e8c87f 100644 --- a/python/sglang/srt/mem_cache/utils.py +++ b/python/sglang/srt/mem_cache/utils.py @@ -13,6 +13,7 @@ # ============================================================================== """Common utilities.""" +import logging from typing import Any, List, Optional, Tuple import torch @@ -21,6 +22,8 @@ import triton.language as tl from sglang.srt.environ import envs +logger = logging.getLogger(__name__) + @triton.jit def set_mla_kv_buffer_kernel( @@ -264,7 +267,7 @@ def maybe_init_custom_mem_pool( return False, None, None -def convert_to_bigram_key(tokens: List[int]) -> List[Tuple[int, int]]: +def _python_convert_to_bigram_key(tokens: List[int]) -> List[Tuple[int, int]]: # EAGLE uses bigram keys in the radix tree since draft sequence is the one-token-shifted version of target # [1, 2, 3, 4] -> [(1,2), (2,3), (3,4)] if len(tokens) and isinstance(tokens[0], tuple): @@ -272,3 +275,28 @@ def convert_to_bigram_key(tokens: List[int]) -> List[Tuple[int, int]]: if len(tokens) < 2: return [] return [(tokens[i], tokens[i + 1]) for i in range(len(tokens) - 1)] + + +try: + from tai_kernel.radix import convert_to_bigram_key as _tai_bigram_impl +except Exception: + _tai_bigram_impl = None + + +def convert_to_bigram_key(tokens: List[int]) -> List[Tuple[int, int]]: + global _tai_bigram_impl + if _tai_bigram_impl is not None and not envs.SGLANG_DISABLE_TAI_BIGRAM.get(): + try: + return _tai_bigram_impl(tokens) + except Exception as e: + # JIT compile or runtime failures (missing compiler, stale cache, + # CPython ABI mismatch) must degrade to the Python path rather + # than crash a serving loop. Disable after one failure so we + # don't keep paying the broken-extension cost on every call. + logger.warning( + "tai-kernel convert_to_bigram_key failed, falling back to " + "Python implementation for the rest of this process: %s", + e, + ) + _tai_bigram_impl = None + return _python_convert_to_bigram_key(tokens) diff --git a/test/registered/unit/mem_cache/test_convert_to_bigram_key.py b/test/registered/unit/mem_cache/test_convert_to_bigram_key.py new file mode 100644 index 000000000..d34c9f28c --- /dev/null +++ b/test/registered/unit/mem_cache/test_convert_to_bigram_key.py @@ -0,0 +1,87 @@ +"""Unit tests for sglang.srt.mem_cache.utils.convert_to_bigram_key. + +Covers both the pure-Python fallback and the tai-kernel pybind11 path +(`tai_kernel.radix.convert_to_bigram_key`) when the extension is importable. + +Run: + python -m pytest test/registered/unit/mem_cache/test_convert_to_bigram_key.py -v +""" + +import sys +import unittest +import unittest.mock + +for _mod in ("sgl_kernel", "sgl_kernel.kvcacheio"): + if _mod not in sys.modules: + sys.modules[_mod] = unittest.mock.MagicMock() + +from sglang.srt.environ import envs +from sglang.srt.mem_cache import utils as mem_cache_utils +from sglang.srt.mem_cache.utils import ( + _python_convert_to_bigram_key, + convert_to_bigram_key, +) +from sglang.test.ci.ci_register import register_amd_ci, register_cuda_ci + +register_cuda_ci(est_time=2, suite="stage-b-test-1-gpu-small") +register_amd_ci(est_time=2, suite="stage-b-test-1-gpu-small-amd") + + +class TestConvertToBigramKey(unittest.TestCase): + """The dispatcher must agree with the pure-Python reference on every input, + regardless of whether the tai-kernel extension is loaded.""" + + cases = [ + [], + [42], + [1, 2], + [1, 2, 3, 4, 5], + list(range(1024)), + # Already-bigram input should be returned as-is (no re-pairing). + [(1, 2), (2, 3), (3, 4)], + ] + + def _assert_equivalent(self, tokens): + expected = _python_convert_to_bigram_key(list(tokens)) + got = convert_to_bigram_key(list(tokens)) + self.assertEqual(got, expected) + self.assertIsInstance(got, list) + for item in got: + self.assertIsInstance(item, tuple) + self.assertEqual(len(item), 2) + + def test_python_fallback(self): + # Force the fallback regardless of whether tai_kernel is available. + with unittest.mock.patch.object( + envs.SGLANG_DISABLE_TAI_BIGRAM, "get", return_value=True + ): + for tokens in self.cases: + with self.subTest(tokens=tokens): + self._assert_equivalent(tokens) + + def test_tai_kernel_path(self): + if mem_cache_utils._tai_bigram_impl is None: + self.skipTest("tai_kernel.radix not importable in this environment") + with unittest.mock.patch.object( + envs.SGLANG_DISABLE_TAI_BIGRAM, "get", return_value=False + ): + for tokens in self.cases: + with self.subTest(tokens=tokens): + self._assert_equivalent(tokens) + + def test_distinct_tuple_identity(self): + # Bigrams that share token values must still be distinct objects so + # downstream radix dicts don't alias keys. + for force_fallback in (True, False): + if not force_fallback and mem_cache_utils._tai_bigram_impl is None: + continue + with unittest.mock.patch.object( + envs.SGLANG_DISABLE_TAI_BIGRAM, "get", return_value=force_fallback + ): + got = convert_to_bigram_key([7, 7, 7]) + self.assertEqual(got, [(7, 7), (7, 7)]) + self.assertIsNot(got[0], got[1]) + + +if __name__ == "__main__": + unittest.main()