Dispatch EAGLE radix bigram builder to tai-kernel
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>
This commit is contained in:
@@ -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)
|
||||
|
||||
Reference in New Issue
Block a user