flush_cache() 在权重更新中的作用及其对 Draft 和 Target Worker 的影响

您询问的是在 SchedulerUpdateWeightsMixin.update_weights_from_tensor() 中调用的 flush_cache() 做了什么，以及是否会同时刷新 draft 和 target worker 的缓存。

flush_cache() 的具体操作

flush_cache() 方法会清空调度器中所有与 KV cache 相关的状态： 1

关键操作包括：

1. 重置 RadixCache：self.tree_cache.reset() 清空所有缓存的前缀树节点
2. 清空请求到 Token 的映射池：self.req_to_token_pool.clear()
3. 清空 Token 到 KV 池的分配器：self.token_to_kv_pool_allocator.clear()
4. 清空 Draft Worker 缓存：如果存在 self.draft_worker，调用 self.draft_worker.clear_cache_pool()
5. 重置统计指标：清零生成 token 数、投机推理统计等
6. 释放 GPU 内存：调用 torch.cuda.empty_cache()

是否同时刷新 Draft 和 Target Worker？

答案：是的，会同时刷新两者的缓存，但方式不同。

1. Draft Worker 的显式清理

如果启用了投机推理（self.draft_worker 存在），flush_cache() 会显式调用 draft worker 的清理方法： 2

在 EAGLE Worker 的实现中，clear_cache_pool() 是一个空方法，因为 draft 和 target worker 共享同一个内存池分配器： 3

2. Target Worker 的隐式清理

Target worker 的 KV cache 通过共享的内存池被清理。关键点在于：

• Draft 和 target worker 共享 req_to_token_pool 和 token_to_kv_pool_allocator
• 当调度器调用 self.req_to_token_pool.clear() 和 self.token_to_kv_pool_allocator.clear() 时，这些共享的分配器被清空
• 因此，target worker 的 KV cache 索引也被清空 4

3. 共享内存架构的关键

Draft 和 target worker 在初始化时获取相同的内存池引用： 5

这意味着：

• 分配器是共享的：req_to_token_pool 和 token_to_kv_pool_allocator 在两个 worker 之间共享
• 物理 KV cache 是独立的：每个 worker 有自己的 token_to_kv_pool（实际存储 KV 数据的 buffer）
• 清空分配器 = 清空两者的索引：当调度器清空共享的分配器时，两个 worker 的 KV cache 索引都被清空

完整的清理流程

Notes

关键设计要点：

1. 统一清理：通过共享内存分配器，一次 flush_cache() 调用就能清空 draft 和 target 两个 worker 的 KV cache 索引

2. 避免重复清理：Draft worker 的 clear_cache_pool() 是空方法，因为清理工作已经通过共享分配器完成

3. 物理内存独立：虽然索引被清空，但每个 worker 的物理 KV buffer 是独立的，这允许它们维护不同的 KV cache 内容（在投机推理的不同阶段）

4. 原子性保证：flush_cache() 在权重更新成功后才执行，并且使用 assert 确保清理成功，保证了系统状态的一致性

5. 分布式同步：在 flush_cache() 之后，代码调用 torch.distributed.barrier() 确保所有 TP ranks 同步完成清理 6

这种设计确保了权重更新后，所有旧的 KV cache（无论是 draft 还是 target 的）都被清空，避免了使用旧权重计算的 KV 与新权重模型混用的问题。

Wiki pages you might want to explore:

• Overview (sgl-project/sglang)

Citations

File: python/sglang/srt/managers/scheduler.py (L2274-2304)

    def flush_cache(self):
        """Flush the memory pool and cache."""
        if self._is_no_request():
            self.cur_batch = None
            self.last_batch = None
            self.tree_cache.reset()
            if self.grammar_backend:
                self.grammar_backend.reset()
            self.req_to_token_pool.clear()
            self.token_to_kv_pool_allocator.clear()

            if self.draft_worker:
                self.draft_worker.clear_cache_pool()

            self.num_generated_tokens = 0
            self.forward_ct_decode = 0
            self.spec_num_accepted_tokens = 0
            self.spec_num_forward_ct = 0
            self.spec_total_num_accepted_tokens = 0
            self.spec_total_num_forward_ct = 0
            torch.cuda.empty_cache()
            logger.info("Cache flushed successfully!")
            if_success = True
        else:
            logging.warning(
                f"Cache not flushed because there are pending requests. "
                f"#queue-req: {len(self.waiting_queue)}, "
                f"#running-req: {len(self.running_batch.reqs)}"
            )
            if_success = False
        return if_success

File: python/sglang/srt/speculative/eagle_worker_v2.py (L104-108)

        # Share the allocator with a target worker.
        # Draft and target worker own their own KV cache pools.
        self.req_to_token_pool, self.token_to_kv_pool_allocator = (
            target_worker.get_memory_pool()
        )

File: python/sglang/srt/speculative/eagle_worker_v2.py (L564-566)

        self.req_to_token_pool, self.token_to_kv_pool_allocator = (
            target_worker.get_memory_pool()
        )

File: python/sglang/srt/speculative/eagle_worker_v2.py (L591-593)

    def clear_cache_pool(self):
        # allocator and kv cache pool are shared with target worker, which are cleared in scheduler
        pass

File: python/sglang/srt/managers/scheduler_update_weights_mixin.py (L88-88)

        torch.distributed.barrier(group=self.tp_cpu_group)