SGLang 源码完全拆解（二十七）：调度策略——FCFS、LPM、DFS 权重与 HRRN，一张等待队列的排序账

1
if (running_batch.batch_is_full or len(self.waiting_queue) == 0) and self.chunked_req is None:
2
    return None, running_batch
3
...
4
# Get priority queue
5
self.policy.calc_priority(
6
    self.waiting_queue,
7
    running_batch,
8
    processed_tokens=self.processed_tokens_counter,
9
    bootstrap_queue=(
10
        self.disagg_prefill_bootstrap_queue.queue
11
        if self.disagg_prefill_bootstrap_queue is not None
12
        else []
13
    ),
14
)
15

16
adder = self.get_prefill_adder_class()(...)
17
if self.chunked_req is not None:
18
    self.chunked_req.init_next_round_input()
19
    adder.chunked_req_limit = self.policy.shortest_prefill_chunk_limit(...)
20
    self.chunked_req = adder.add_chunked_req(self.chunked_req)
21
...
22
for req in self.waiting_queue:
23
    ...
24
    req.init_next_round_input(self.tree_cache)
25
    res = adder.add_one_req(...)
26
    if res != AddReqResult.CONTINUE:
27
        ...
28
        break

1
policy = self._determine_active_policy(waiting_queue)
2

3
# Cache-agnostic policies sort without matching; match anyway so the load
4
# snapshot sees num_matched_prefix_tokens. Decode never prefills.
5
if (
6
    not isinstance(policy, CacheAwarePolicy)
7
    and self.tree_cache.supports_fast_match_prefix()
8
    and get_disagg().disaggregation_mode != "decode"
9
):
10
    for r in waiting_queue:
11
        match_kv_cache(r, self.tree_cache)
12

13
if self.policy == CacheAgnosticPolicy.FCFS:
14
    if self.enable_priority_scheduling:
15
        SchedulePolicy._sort_by_priority_and_fcfs(waiting_queue, self.priority_sign)
16
elif isinstance(policy, CacheAwarePolicy):
17
    ...

1
def _compute_prefix_matches(
2
    self, waiting_queue: List[Req], policy: CacheAwarePolicy
3
) -> Set[int]:
4
    temporary_deprioritized: Set[int] = set()
5
    self.waiting_queue_radix_tree.reset()
6

7
    for r in waiting_queue:
8
        prefix_ids = r.origin_input_ids + r.output_ids
9
        extra_key = r.extra_key
10
        cache_salt = r.cache_salt
11
        match_result = match_kv_cache(r, self.tree_cache, prefix_ids)
12

13
        # NOTE(sang): This logic is for in-batch prefix caching;
14
        # If there are more than 1 request that have small matching prefix from
15
        # existing cache, but all those requests share the same prefix, we prefer
16
        # to schedule only one of them so that we can increase the cache hit rate.
17
        if r.prefix_len <= IN_BATCH_PREFIX_CACHING_CHECK_THRESHOLD:
18
            match_result = self.waiting_queue_radix_tree.match_prefix(
19
                MatchPrefixParams(
20
                    key=RadixKey(
21
                        token_ids=prefix_ids,
22
                        extra_key=extra_key,
23
                        cache_salt=cache_salt,
24
                    )
25
                )
26
            )
27
            if envs.SGLANG_RADIX_FORCE_MISS.get():
28
                match_result = zero_match_result(
29
                    self.waiting_queue_radix_tree, match_result, extra_key=extra_key
30
                )
31
            if match_result.device_prefix_len >= IN_BATCH_PREFIX_CACHING_DEPRIORITIZE_THRESHOLD:
32
                temporary_deprioritized.add(r.rid)
33
            else:
34
                # Insert with a dummy key
35
                self.waiting_queue_radix_tree.insert(
36
                    InsertParams(
37
                        key=RadixKey(token_ids=prefix_ids, extra_key=extra_key, cache_salt=cache_salt),
38
                        value=torch.empty(len(prefix_ids), dtype=torch.bool),
39
                    )
40
                )
41
    return temporary_deprioritized

1
def _dfs_weight_order(root_node, node_handles, resolve_node_handle) -> list[int]:
2
    last_node_to_indices = {}
3
    for index, node_handle in enumerate(node_handles):
4
        node = resolve_node_handle(node_handle)
5
        last_node_to_indices.setdefault(node, []).append(index)
6

7
    node_to_weight = {node: len(indices) for node, indices in last_node_to_indices.items()}
8

9
    stack = [(root_node, False)]
10
    while stack:
11
        node, visited = stack.pop()
12
        if visited:
13
            weight = node_to_weight.get(node, 0)
14
            for child in node.children.values():
15
                weight += node_to_weight.get(child, 0)
16
            node_to_weight[node] = weight
17
            continue
18
        stack.append((node, True))
19
        for child in reversed(list(node.children.values())):
20
            stack.append((child, False))
21

22
    order = []
23
    stack = [(root_node, False)]
24
    while stack:
25
        node, visited = stack.pop()
26
        if visited:
27
            order.extend(last_node_to_indices.get(node, ()))
28
            continue
29
        children = list(node.children.values())
30
        children.sort(key=lambda child: -node_to_weight.get(child, 0))
31
        stack.append((node, True))
32
        for child in reversed(children):
33
            stack.append((child, False))
34
    return order

1
@staticmethod
2
def _sort_by_hrrn(
3
    waiting_queue: List[Req],
4
    temporary_deprioritized: Set[int],
5
    processed_tokens: int,
6
) -> None:
7
    def _key(r: Req):
8
        rid = r.rid
9
        if rid in temporary_deprioritized:
10
            return (float("inf"), rid)
11
        uncached = SchedulePolicy._uncached_len(r)
12
        if uncached <= 0:
13
            # No prefill work; drain immediately.
14
            return (-float("inf"), rid)
15
        waited_tokens = max(0, processed_tokens - r.arrival_processed_tokens)
16
        ratio_delta = waited_tokens / uncached
17
        return (-ratio_delta, rid)
18

19
    waiting_queue.sort(key=_key)

1
@staticmethod
2
def _shortest_prefill_work(r: Req) -> int:
3
    return max(
4
        1,
5
        len(r.origin_input_ids) + len(r.output_ids) - r.num_matched_prefix_tokens,
6
    )
7

8
@staticmethod
9
def _sort_by_shortest_prefill(
10
    waiting_queue: List[Req], temporary_deprioritized: Set[int]
11
) -> None:
12
    waiting_queue.sort(
13
        key=lambda r: (
14
            r.rid in temporary_deprioritized,
15
            SchedulePolicy._shortest_prefill_work(r),
16
            r.time_stats.wait_queue_entry_time,
17
        )
18
    )

1
def shortest_prefill_chunk_limit(
2
    self, chunked_req: Req, waiting_queue: List[Req], budget: int, page_size: int
3
) -> Optional[int]:
4
    """Cap the active prefill chunk to reserve tokens for shorter waiting requests."""
5
    if (
6
        self.policy != CacheAwarePolicy.SHORTEST_PREFILL_FIRST
7
        or budget < 2 * page_size
8
    ):
9
        return None
10
    remaining = len(chunked_req.full_untruncated_fill_ids) - chunked_req.prefix_len
11
    reserved = 0
12
    for req in waiting_queue:
13
        work = self._shortest_prefill_work(req)
14
        charge = _ceil_div(work, page_size) * page_size
15
        if work >= remaining or reserved + charge > budget - page_size:
16
            break
17
        reserved += charge
18
    if not reserved:
19
        return None
20
    # Page alignment keeps continuation boundaries allocator-compatible.
21
    return (budget - reserved) // page_size * page_size

1
@staticmethod
2
def _sort_by_routing_key(waiting_queue: List[Req], running_batch: ScheduleBatch) -> None:
3
    routing_key_counts = Counter(
4
        r.routing_key for r in running_batch.reqs if r.routing_key
5
    )
6
    if not routing_key_counts:
7
        return
8

9
    def sort_key(req: Req):
10
        key = req.routing_key
11
        if key and key in routing_key_counts:
12
            return (0, -count, key)
13
        else:
14
            return (1, 0, key or "")
15

16
    waiting_queue.sort(key=sort_key)

1
def _touch_waiting_prefixes(self, policy, waiting_queue, bootstrap_queue=[]):
2
    # Head last, so LRU evicts in reverse admission order; cache-aware policies
3
    # already touch while matching, and other eviction strategies ignore recency.
4
    bootstrap_max = envs.SGLANG_WAITING_PREFIX_REFRESH_BOOTSTRAP_MAX.get()
5
    if (
6
        isinstance(policy, CacheAwarePolicy)
7
        or not (waiting_queue or (bootstrap_queue and bootstrap_max > 0))
8
        or not envs.SGLANG_ENABLE_WAITING_PREFIX_REFRESH.get()
9
        or get_disagg().disaggregation_mode == "decode"
10
        or get_memory().radix_eviction_policy.lower() not in ("lru", "joint_lru")
11
        or self.tree_cache.supports_fast_match_prefix()
12
        or not self.tree_cache.supports_prefix_sharing()
13
    ):
14
        return
15
    refresh = (
16
        self._waiting_prefix_refresh_calls % WAITING_PREFIX_REFRESH_INTERVAL == 0
17
    )
18
    self._waiting_prefix_refresh_calls += 1
19
    if not refresh:
20
        return
21
    for r in reversed(bootstrap_queue[:bootstrap_max]):
22
        touch_waiting_prefix(r, self.tree_cache)
23
    for r in reversed(waiting_queue[:WAITING_QUEUE_PREFIX_MATCH_MAX]):
24
        touch_waiting_prefix(r, self.tree_cache)

1
class PrefillDelayerSinglePassExecutor:
2
    def __init__(self, prefill_delayer: PrefillDelayer, token_usage: float):
3
        self._prefill_delayer = prefill_delayer
4
        self._token_usage = token_usage
5
        self._result: Optional[_NegotiateOutput] = None
6
        self._attempted_prefill_bs = 0
7
    ...
8
    def negotiate_should_allow_prefill(
9
        self,
10
        local_prefillable: bool,
11
        running_batch: int = 0,
12
        max_prefill_bs: int = 0,
13
        max_running_requests: int = 0,
14
        waiting_queue_len: int = 0,
15
    ) -> bool:
16
        if local_prefillable:
17
            self._attempted_prefill_bs = max(
18
                self._attempted_prefill_bs,
19
                self._estimate_attempted_prefill_bs(
20
                    running_batch=running_batch,
21
                    max_running_requests=max_running_requests,
22
                    waiting_queue_len=waiting_queue_len,
23
                ),
24
            )
25
        if not self._called:
26
            self._result = self._prefill_delayer._negotiate_should_allow_prefill(
27
                local_prefillable=local_prefillable,
28
                token_usage=self._token_usage,
29
                running_batch=running_batch,
30
                max_prefill_bs=max_prefill_bs,
31
                max_running_requests=max_running_requests,
32
                waiting_queue_len=waiting_queue_len,
33
            )
34
        return self._result.output_allow

1
event_loop
2
  └─ get_next_batch_to_run
3
       ├─ 合并上一轮 prefill batch、过滤已完成请求
4
       ├─ chunked_req 暂摘
5
       ├─ _should_defer_prefill?  → 是：这一轮只跑 decode
6
       └─ get_new_batch_prefill
7
            ├─ 预取 grammar 队列、HiCache 状态
8
            ├─ PrefillDelayerSinglePassExecutor 协商（DP 下跨 rank）
9
            ├─ SchedulePolicy.calc_priority
10
            │    ├─ _validate_and_adjust_policy（构造时已决定）
11
            │    ├─ _determine_active_policy（队列 > 128 → FCFS）
12
            │    ├─ 缓存无关策略补 match_kv_cache（负载快照用）
13
            │    ├─ _compute_prefix_matches（同批去重 + 临时树）
14
            │    ├─ 按策略排序：FCFS / LPM / DFS / HRRN / SPF / LOF / routing-key
15
            │    └─ _touch_waiting_prefixes（无缓存策略、每 32 轮）
16
            ├─ shortest_prefill_chunk_limit（仅 SPF）
17
            ├─ add_chunked_req（续算上一轮长提示）
18
            ├─ PrefillAdder(...)
19
            └─ for req in waiting_queue:
20
                 ├─ add_one_req → NO_TOKEN/OTHER → break
21
                 └─ CONTINUE