leavelet
d655fad040
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>
Registered Tests
Tests under this directory are auto-discovered by run_suite.py via CI registration decorators.
Where Should I Put My New Test?
No server / engine launch required
| What you're testing | Directory | Requires |
|---|---|---|
| Component logic in isolation (cache, scheduler, config, parser, etc.) | unit/<module>/ |
CPU or GPU |
| CUDA kernel correctness | kernels/ |
GPU |
Server / engine launch required (E2E)
| What you're testing | Directory | Requires |
|---|---|---|
| Model inference correctness | models/, 4-gpu-models/, 8-gpu-models/ |
GPU |
| Feature-specific (OpenAI API, LoRA, speculative, distributed, VLM, etc.) | openai_server/, lora/, spec/, distributed/, ... |
GPU |
| Benchmarks (performance, accuracy, stress) | benchmark/ |
GPU |
| Platform-specific | amd/, ascend/ |
Vendor GPU |
See unit/README.md for unit test conventions.