chenchenghao/sglang
1
0
Fork 0
Code Issues Pull Requests Actions 11 Packages Projects Releases Wiki Activity

Refactor eagle bigram key matching (#13714)

Browse Source
This commit is contained in:
Liangsheng Yin
2025-11-22 20:40:42 +08:00
committed by GitHub
parent a92afb00c6
commit ac43822634
7 changed files with 89 additions and 181 deletions
Download Patch File Download Diff File Expand all files Collapse all files

6
python/sglang/srt/managers/schedule_batch.py
View File

@@ -562,8 +562,8 @@ class Req:
self.host_hit_length = 0
# The node to lock until for swa radix tree lock ref
self.swa_uuid_for_lock: Optional[int] = None
# The prefix length of the last prefix matching
self.last_matched_prefix_len: int = 0
# The prefix length that is inserted into the tree cache
self.cache_protected_len: int = 0
# Whether or not if it is chunked. It increments whenever
# it is chunked, and decrement whenever chunked request is
@@ -775,7 +775,7 @@ class Req:
else {}
),
)
self.last_matched_prefix_len = len(self.prefix_indices)
self.cache_protected_len = len(self.prefix_indices)
self.extend_input_len = len(self.fill_ids) - len(self.prefix_indices)
# Based on https://github.com/vllm-project/vllm/blob/7a64d24aad69e4d2548aa0bf528d9fe63428ab01/vllm/transformers_utils/detokenizer.py#L194-L313

2
python/sglang/srt/managers/schedule_policy.py
View File

@@ -603,7 +603,7 @@ class PrefillAdder:
req.prefix_indices = torch.cat([req.prefix_indices, new_indices])
req.extend_input_len = len(req.fill_ids) - len(req.prefix_indices)
prefix_len = len(req.prefix_indices)
req.last_matched_prefix_len = prefix_len
req.cache_protected_len = prefix_len
input_tokens = self.ceil_paged_tokens(req.extend_input_len)

4
python/sglang/srt/mem_cache/hiradix_cache.py
View File

@@ -685,7 +685,7 @@ class HiRadixCache(RadixCache):
def match_prefix(self, key: RadixKey, **kwargs):
empty_value = torch.empty((0,), dtype=torch.int64, device=self.device)
key.token_ids = self.key_convert_fn(key.token_ids)
key, _ = self.maybe_bigram_convert(key)
if self.disable or len(key) == 0:
return MatchResult(
device_indices=empty_value,
@@ -857,7 +857,7 @@ class HiRadixCache(RadixCache):
):
if priority is None:
priority = 0
key.token_ids = self.key_convert_fn(key.token_ids)
key, value = self.maybe_bigram_convert(key, value)
if len(key) == 0:
return 0

170
python/sglang/srt/mem_cache/radix_cache.py
View File

@@ -1,5 +1,7 @@
from __future__ import annotations
from sglang.srt.mem_cache.utils import convert_to_bigram_key
"""
Copyright 2023-2024 SGLang Team
Licensed under the Apache License, Version 2.0 (the "License");
@@ -51,12 +53,18 @@ if TYPE_CHECKING:
class RadixKey:
def __init__(self, token_ids: List[int], extra_key: Optional[str] = None):
def __init__(
self,
token_ids: List[int],
extra_key: Optional[str] = None,
is_bigram: bool = False,
):
# token ids sequence
self.token_ids = token_ids
# extra key (e.g. lora_id, cache_salt)
self.extra_key = extra_key
# is bigram key
self.is_bigram = is_bigram
def __len__(self) -> int:
return len(self.token_ids)
@@ -178,16 +186,6 @@ def get_child_key(key: RadixKey, page_size: int = 1):
return (key.extra_key, plain_key)
def _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) < 2:
return []
if isinstance(tokens[0], tuple):
return tokens
return [(tokens[i], tokens[i + 1]) for i in range(len(tokens) - 1)]
class RadixCache(BasePrefixCache):
def __init__(
self,
@@ -225,11 +223,6 @@ class RadixCache(BasePrefixCache):
self.key_match_fn = partial(_key_match_paged, page_size=page_size)
self.get_child_key_fn = partial(get_child_key, page_size=page_size)
if is_eagle:
self.key_convert_fn = _convert_to_bigram_key
else:
self.key_convert_fn = lambda key: key
if eviction_policy.lower() == "lru":
self.eviction_strategy: EvictionStrategy = LRUStrategy()
elif eviction_policy.lower() == "lfu":
@@ -261,6 +254,16 @@ class RadixCache(BasePrefixCache):
self.protected_size_ = 0
self._record_all_cleared_event()
def maybe_bigram_convert(
self, key: RadixKey, value: Optional[torch.Tensor] = None
) -> Tuple[RadixKey, Optional[torch.Tensor]]:
if self.is_eagle and not key.is_bigram:
key.token_ids = convert_to_bigram_key(key.token_ids)
if value is not None:
value = value[: len(key)]
return key, value
def match_prefix(self, key: RadixKey, **kwargs) -> MatchResult:
"""Find the longest cached prefix of ``key`` in the radix tree.
@@ -299,7 +302,7 @@ class RadixCache(BasePrefixCache):
to expose a precise boundary; this structural refinement improves
subsequent match efficiency and does not duplicate data.
"""
key.token_ids = self.key_convert_fn(key.token_ids)
key, _ = self.maybe_bigram_convert(key)
def empty_match_result():
return MatchResult(
@@ -337,78 +340,60 @@ class RadixCache(BasePrefixCache):
if self.disable:
return 0
key.token_ids = self.key_convert_fn(key.token_ids)
if value is None:
value = torch.tensor(key.token_ids, dtype=torch.int64)
if self.is_eagle:
# Make sure the value len equal to the EAGLE bigram key len
value = value[: len(key)]
key, value = self.maybe_bigram_convert(key, value)
return self._insert_helper(self.root_node, key, value, priority)
def _page_align_keys(self, key: list) -> list:
if self.page_size == 1:
return key
page_aligned_len = len(key) // self.page_size * self.page_size
return key[:page_aligned_len]
def cache_finished_req(self, req: Req, is_insert: bool = True):
"""Cache request when it finishes."""
# In deterministic mode, disable finished request insertion to radix cache
if self.disable_finished_insert:
is_insert = False
committed_kv_len = req.pop_committed_kv_cache()
kv_committed_len = req.pop_committed_kv_cache()
if self.disable:
kv_indices = self.req_to_token_pool.req_to_token[
req.req_pool_idx, :committed_kv_len
req.req_pool_idx, :kv_committed_len
]
self.token_to_kv_pool_allocator.free(kv_indices)
self.req_to_token_pool.free(req.req_pool_idx)
return
token_ids = (req.origin_input_ids + req.output_ids)[:committed_kv_len]
# For EAGLE radix cache, we will convert the key to bigram key, e.g. [1,2,3,4] -> [(1,2), (2,3), (3,4)], the length will -1. ((len([(1,2), (2,3), (3,4)]) = len([1,2,3,4]) - 1))
# So for the corresponding kv length should also -1. Then we get the actual_kv_len, and use it to do later calculation and slicing.
actual_kv_len = committed_kv_len - 1 if self.is_eagle else committed_kv_len
token_ids = (req.origin_input_ids + req.output_ids)[:kv_committed_len]
kv_indices = self.req_to_token_pool.req_to_token[
req.req_pool_idx, :committed_kv_len
req.req_pool_idx, : len(token_ids)
]
if self.page_size != 1:
page_aligned_len = actual_kv_len // self.page_size * self.page_size
page_aligned_kv_indices = kv_indices[:page_aligned_len].to(
dtype=torch.int64, copy=True
)
else:
page_aligned_len = actual_kv_len
page_aligned_kv_indices = kv_indices.to(dtype=torch.int64, copy=True)
page_aligned_token_len = (
page_aligned_len + 1 if self.is_eagle else page_aligned_len
)
old_prefix_len = len(req.prefix_indices)
if self.is_eagle and old_prefix_len > req.last_matched_prefix_len:
# In EAGLE chunked prefill case, the prefix_indices included one unmatched token (kv_indices[actual_kv_len:])
# Here we -1 to make sure the kv of the unmatched token can be freed correctly to avoid memory leak
old_prefix_len -= 1
# Maybe convert to bigram keys for EAGLE
keys = convert_to_bigram_key(req.fill_ids) if self.is_eagle else req.fill_ids
keys = self._page_align_keys(keys)
values = kv_indices[: len(keys)].to(dtype=torch.int64, copy=True)
radix_key = RadixKey(keys, req.extra_key, is_bigram=self.is_eagle)
# Radix Cache takes one ref in memory pool
if is_insert:
priority = getattr(req, "priority", 0) or 0
new_prefix_len = self.insert(
RadixKey(token_ids[:page_aligned_token_len], req.extra_key),
page_aligned_kv_indices,
priority=priority,
)
new_prefix_len = self.insert(radix_key, values, priority=priority)
# Free the duplicates that were already in the tree
self.token_to_kv_pool_allocator.free(
kv_indices[old_prefix_len:new_prefix_len]
kv_indices[req.cache_protected_len : new_prefix_len]
)
else:
self.token_to_kv_pool_allocator.free(
kv_indices[old_prefix_len:page_aligned_len]
kv_indices[req.cache_protected_len : len(keys)]
)
# free the unaligned tail
self.token_to_kv_pool_allocator.free(kv_indices[page_aligned_len:])
self.token_to_kv_pool_allocator.free(kv_indices[len(keys) :])
# Remove req slot release the cache lock
self.req_to_token_pool.free(req.req_pool_idx)
@@ -420,77 +405,56 @@ class RadixCache(BasePrefixCache):
return
token_ids = req.fill_ids
all_token_len = len(token_ids)
# For EAGLE radix cache, we will convert the key to bigram key, e.g. [1,2,3,4] -> [(1,2), (2,3), (3,4)], the length will -1. ((len([(1,2), (2,3), (3,4)]) = len([1,2,3,4]) - 1))
# So for the corresponding kv length should also -1. Then we get the actual_kv_len, and use it to do later calculation and slicing.
actual_kv_len = all_token_len - 1 if self.is_eagle else all_token_len
kv_indices = self.req_to_token_pool.req_to_token[
req.req_pool_idx, :all_token_len
req.req_pool_idx, : len(token_ids)
]
if self.page_size != 1:
page_aligned_len = actual_kv_len // self.page_size * self.page_size
page_aligned_kv_indices = kv_indices[:page_aligned_len].to(
dtype=torch.int64, copy=True
)
else:
page_aligned_len = actual_kv_len
page_aligned_kv_indices = kv_indices.to(dtype=torch.int64, copy=True)
# For EAGLE, the page_aligned_len is for the bigram key, the normal key len should +1
page_aligned_token_len = (
page_aligned_len + 1 if self.is_eagle else page_aligned_len
)
page_aligned_token_ids = token_ids[:page_aligned_token_len]
old_prefix_len = len(req.prefix_indices)
if self.is_eagle and old_prefix_len > req.last_matched_prefix_len:
# In EAGLE chunked prefill case, the prefix_indices included one unmatched token (kv_indices[actual_kv_len:])
# Here we -1 to make sure the kv of the unmatched token can be freed correctly to avoid memory leak
old_prefix_len -= 1
# Maybe convert to bigram keys for EAGLE
keys = convert_to_bigram_key(req.fill_ids) if self.is_eagle else req.fill_ids
keys = self._page_align_keys(keys)
values = kv_indices[: len(keys)].to(dtype=torch.int64, copy=True)
radix_key = RadixKey(keys, req.extra_key, is_bigram=self.is_eagle)
# Radix Cache takes one ref in memory pool
priority = getattr(req, "priority", 0) or 0
new_prefix_len = self.insert(
RadixKey(page_aligned_token_ids, req.extra_key),
page_aligned_kv_indices,
radix_key,
values,
chunked=chunked,
priority=priority,
priority=getattr(req, "priority", 0) or 0,
)
self.token_to_kv_pool_allocator.free(
kv_indices[req.cache_protected_len : new_prefix_len]
)
self.token_to_kv_pool_allocator.free(kv_indices[old_prefix_len:new_prefix_len])
# The prefix indices could be updated, reuse it
new_indices, new_last_node, _, _ = self.match_prefix(
RadixKey(token_ids=page_aligned_token_ids, extra_key=req.extra_key)
)
new_indices, new_last_node, _, _ = self.match_prefix(radix_key)
assert len(new_indices) == len(keys), f"{len(new_indices)=}, {len(keys)=}"
self.req_to_token_pool.write(
(req.req_pool_idx, slice(old_prefix_len, len(new_indices))),
new_indices[old_prefix_len:],
(req.req_pool_idx, slice(req.cache_protected_len, len(new_indices))),
new_indices[req.cache_protected_len :],
)
# The last_matched_prefix_len is not always equal to len(req.prefix_indices)
# The cache_protected_len is not always equal to len(req.prefix_indices)
# since for page_size > 1, the partial part is added to req.prefix_indices, but that part of kv indices is not added to the tree.
# It should be freed in the next cache_unfinished_req and final cache_finished_req to avoid memory leak.
# So we introduce this `last_matched_prefix_len` field to make sure the partial part can be freed correctly.
req.last_matched_prefix_len = len(new_indices)
# So we introduce this `cache_protected_len` field to make sure the partial part can be freed correctly.
req.cache_protected_len = len(new_indices)
self.dec_lock_ref(req.last_node)
self.inc_lock_ref(new_last_node)
# `req.prefix_indices` will be used in `PrefillAdder::add_chunked_req` later
if self.page_size != 1:
# Handle partial page, the partial part should be freed in the next cache_unfinished_req and final cache_finished_req.
# - page_size != 1: there is a partial page at the end, keep the full kv_indices
# - eagle case: bigram keys will only cache len - 1 kv indices
if len(new_indices) < len(kv_indices):
req.prefix_indices = torch.cat(
[new_indices, kv_indices[len(new_indices) :]]
)
else:
if self.is_eagle:
# Attach the kv index of the last token for EAGLE, it can be used in chunked prefill
req.prefix_indices = torch.cat(
[new_indices, kv_indices[actual_kv_len:]]
)
else:
req.prefix_indices = new_indices
req.prefix_indices = new_indices
req.last_node = new_last_node
def pretty_print(self):

10
python/sglang/srt/mem_cache/swa_radix_cache.py
View File

@@ -33,11 +33,11 @@ from sglang.srt.mem_cache.base_prefix_cache import BasePrefixCache, MatchResult
from sglang.srt.mem_cache.memory_pool import ReqToTokenPool
from sglang.srt.mem_cache.radix_cache import (
RadixKey,
_convert_to_bigram_key,
_key_match_page_size1,
_key_match_paged,
get_child_key,
)
from sglang.srt.mem_cache.utils import convert_to_bigram_key
if TYPE_CHECKING:
from sglang.srt.managers.schedule_batch import Req
@@ -359,7 +359,7 @@ class SWARadixCache(BasePrefixCache):
self.get_child_key_fn = partial(get_child_key, page_size=page_size)
if is_eagle:
self.key_convert_fn = _convert_to_bigram_key
self.key_convert_fn = convert_to_bigram_key
else:
self.key_convert_fn = lambda key: key
@@ -474,7 +474,7 @@ class SWARadixCache(BasePrefixCache):
)
old_prefix_len = len(req.prefix_indices)
if self.is_eagle and old_prefix_len > req.last_matched_prefix_len:
if self.is_eagle and old_prefix_len > req.cache_protected_len:
# In EAGLE chunked prefill case, the prefix_indices included one unmatched token (kv_indices[actual_kv_len:])
# Here we -1 to make sure the kv of the unmatched token can be freed correctly to avoid memory leak
old_prefix_len -= 1
@@ -536,7 +536,7 @@ class SWARadixCache(BasePrefixCache):
page_aligned_token_ids = token_ids[:page_aligned_token_len]
old_prefix_len = len(req.prefix_indices)
if self.is_eagle and old_prefix_len > req.last_matched_prefix_len:
if self.is_eagle and old_prefix_len > req.cache_protected_len:
# In EAGLE chunked prefill case, the prefix_indices included one unmatched token (kv_indices[actual_kv_len:])
# Here we -1 to make sure the kv of the unmatched token can be freed correctly to avoid memory leak
old_prefix_len -= 1
@@ -562,7 +562,7 @@ class SWARadixCache(BasePrefixCache):
new_indices[old_prefix_len:],
)
req.last_matched_prefix_len = len(new_indices)
req.cache_protected_len = len(new_indices)
self.dec_lock_ref(req.last_node, req.swa_uuid_for_lock)
swa_uuid_for_lock = self.inc_lock_ref(new_last_node)

12
python/sglang/srt/mem_cache/utils.py
View File

@@ -13,7 +13,7 @@
# ==============================================================================
"""Common utilities."""
from typing import Any, Optional, Tuple
from typing import Any, List, Optional, Tuple
import torch
import triton
@@ -241,3 +241,13 @@ def maybe_init_custom_mem_pool(
return init_mooncake_custom_mem_pool(device)
else:
return False, None, None
def 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):
return tokens
if len(tokens) < 2:
return []
return [(tokens[i], tokens[i + 1]) for i in range(len(tokens) - 1)]