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>

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)

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)