The pure-Python convert_to_bigram_key list comprehension runs on every cache_finished_req / cache_unfinished_req of every EAGLE radix variant (radix_cache, hiradix_cache, swa_radix_cache), with token-list lengths that scale with prompt + output. Scheduler profiles consistently flag it as the largest mem_cache-side CPU hotspot. This commit wires sglang.srt.mem_cache.utils.convert_to_bigram_key to tai_kernel.radix.convert_to_bigram_key when the extension is importable, falling back to the pure-Python implementation otherwise. The tai-kernel path uses a pybind11 module that calls the CPython C API directly (PyTuple_New + PyTuple_SET_ITEM with ref-stealing) rather than rebuilding the list comprehension's bytecode-level tuple allocations. Measured 1.4x at n=131k and up to 2.5x for n=1k on g0034 Python 3.11; allocator-bound at large n because CPython's 2-tuple freelist already amortises the construction. The int64-packing follow-up that bypasses tuple allocation entirely is parked as a separate work item. Runtime safety: - The dispatcher catches any first-call JIT compile / runtime failure, logs once, and falls through to the pure-Python path for the rest of the process — JIT failures must degrade rather than crash a serving loop. - SGLANG_DISABLE_TAI_BIGRAM forces the Python path for bisecting. Constraint: Output must be a real Python List[Tuple[int, int]] because downstream radix dicts use the tuples as hashable keys. Rejected: Pre-allocated tuple slab pool | CPython's per-interpreter 2-tuple freelist already serves this case, and we cannot recycle tuples that become radix-tree keys without changing the consumer. Rejected: int64-packed keys this round | requires changes to RadixKey, get_child_key_fn, key_match_fn, and EAGLE bigram detection; deserves its own plan. Confidence: high Scope-risk: low Directive: Keep _python_convert_to_bigram_key reachable; if the tai-kernel path is ever removed, the EAGLE radix cache must continue to work unchanged. Tested: tai-kernel side validated on the cluster (``python benchmark/radix/benchmark_convert_to_bigram_key.py --check`` prints byte-exact correctness then 1.4-2.5x speedup across sizes 128..131072 on g0034). The new ``test/registered/unit/mem_cache/test_convert_to_bigram_key.py`` exercises both dispatch paths via SGLANG_DISABLE_TAI_BIGRAM patching. Not-tested: End-to-end EAGLE serving accuracy + scheduler-time delta on a real workload. Co-Authored-By: Claude Opus 4.7 (1M context) <noreply@anthropic.com>
Unit Tests
Component-level tests that do not launch a server or load model weights. Tests can use CPU or GPU — the key criterion is no server process.
Quick Start
- Find the source file under
python/sglang/srt/. - Create the corresponding test here, mirroring the source tree:
srt/mem_cache/radix_cache.py → unit/mem_cache/test_radix_cache.py srt/sampling/sampling_params.py → unit/sampling/test_sampling_params.py - Register for CI at the top of the file (after imports, before test classes):
from sglang.test.ci.ci_register import register_cpu_ci register_cpu_ci(est_time=5, suite="stage-a-test-cpu") # or: register_cuda_ci(est_time=10, suite="stage-b-test-1-gpu-small") - Run locally:
pytest test/registered/unit/ -v # all unit tests pytest test/registered/unit/mem_cache/ -v # one module - Run with coverage:
# summary pytest test/registered/unit/ --cov --cov-config=.coveragerc -v # PR incremental check (require ≥60% on changed lines) pytest test/registered/unit/ --cov --cov-config=.coveragerc --cov-report=xml diff-cover coverage.xml --compare-branch=origin/main --fail-under=60
Example
"""Unit tests for <module> — no server, no model loading."""
from sglang.test.ci.ci_register import register_cpu_ci
register_cpu_ci(est_time=5, suite="stage-a-test-cpu")
import unittest
from sglang.srt.<module> import TargetClass
from sglang.test.test_utils import CustomTestCase
class TestTargetClass(CustomTestCase):
def test_basic_behavior(self):
obj = TargetClass(...)
self.assertEqual(obj.method(), expected)
if __name__ == "__main__":
unittest.main()
Rules
- No
popen_launch_server()orEngine(...). - No model weight loading.
- Use
CustomTestCase(fromsglang.test.test_utils, adds CI retry). - Use
unittest.mockfor dependencies that are expensive to construct.