深入 SGLang Overlap 调度：CPU-GPU Overlap、SBO 与 TBO 源码分析

┌─────────────────────────────────────────────────────┐
│  L1: CPU-GPU Overlap (Scheduler 级别)                │
│  ┌───────────────────────────────────────────────┐  │
│  │  L3: TBO (Forward 内部，child batch 级流水)    │  │
│  │  ┌─────────────────────────────────────────┐  │  │
│  │  │  L2: SBO (MoE layer 内部，stream 级别)   │  │  │
│  │  └─────────────────────────────────────────┘  │  │
│  └───────────────────────────────────────────────┘  │
└─────────────────────────────────────────────────────┘

┌─────────────────────────────────────────────────────┐
│  L1: CPU-GPU Overlap (Scheduler 级别)                │
│  ┌───────────────────────────────────────────────┐  │
│  │  L3: TBO (Forward 内部，child batch 级流水)    │  │
│  │  ┌─────────────────────────────────────────┐  │  │
│  │  │  L2: SBO (MoE layer 内部，stream 级别)   │  │  │
│  │  └─────────────────────────────────────────┘  │  │
│  └───────────────────────────────────────────────┘  │
└─────────────────────────────────────────────────────┘

disable_overlap_schedule: bool = False          # 禁用调度层 overlap（默认开启）
enable_two_batch_overlap: bool = False           # 启用 TBO
enable_single_batch_overlap: bool = False        # 启用 SBO
tbo_token_distribution_threshold: float = 0.48   # TBO 拆分平衡阈值（控制 two-chunk split）

disable_overlap_schedule: bool = False          # 禁用调度层 overlap（默认开启）
enable_two_batch_overlap: bool = False           # 启用 TBO
enable_single_batch_overlap: bool = False        # 启用 SBO
tbo_token_distribution_threshold: float = 0.48   # TBO 拆分平衡阈值（控制 two-chunk split）

迭代 N:
  CPU: schedule(batch N) + process_result(batch N-1)
  GPU:               forward(batch N)

迭代 N:
  CPU: schedule(batch N) + process_result(batch N-1)
  GPU:               forward(batch N)

def run_event_loop(self) -> None:
    # schedule_stream 用来承载调度阶段提交的异步 GPU 准备工作。
    self.schedule_stream = self.device_module.Stream(priority=0)
    if self.device == "cpu":
        # CPU 模式下没有真实的设备流同步，覆写成空操作即可。
        self.schedule_stream.synchronize = lambda: None
    # 整个 scheduler 事件循环都运行在 schedule_stream 上下文中。
    with self.device_module.StreamContext(self.schedule_stream):
        dispatch_event_loop(self)
 
def dispatch_event_loop(scheduler: Scheduler):
    # 按运行模式分发到不同事件循环；overlap 主路径在这里切入。
    if disaggregation_mode == DisaggregationMode.NULL:
        if scheduler.enable_pdmux:
            scheduler.event_loop_pdmux()
        elif server_args.pp_size > 1:
            scheduler.event_loop_pp()
        elif scheduler.enable_overlap:
            scheduler.event_loop_overlap()
        else:
            scheduler.event_loop_normal()

def run_event_loop(self) -> None:
    # schedule_stream 用来承载调度阶段提交的异步 GPU 准备工作。
    self.schedule_stream = self.device_module.Stream(priority=0)
    if self.device == "cpu":
        # CPU 模式下没有真实的设备流同步，覆写成空操作即可。
        self.schedule_stream.synchronize = lambda: None
    # 整个 scheduler 事件循环都运行在 schedule_stream 上下文中。
    with self.device_module.StreamContext(self.schedule_stream):
        dispatch_event_loop(self)
 
def dispatch_event_loop(scheduler: Scheduler):
    # 按运行模式分发到不同事件循环；overlap 主路径在这里切入。
    if disaggregation_mode == DisaggregationMode.NULL:
        if scheduler.enable_pdmux:
            scheduler.event_loop_pdmux()
        elif server_args.pp_size > 1:
            scheduler.event_loop_pp()
        elif scheduler.enable_overlap:
            scheduler.event_loop_overlap()
        else:
            scheduler.event_loop_normal()

def init_overlap(self):
    # 初始化与当前 device 绑定的 stream 上下文。
    self.device_module = torch.get_device_module(self.device)
 
    self.forward_stream_ctx = self.device_module.stream(self.forward_stream)
    self.copy_stream = self.device_module.Stream()
    self.copy_stream_ctx = self.device_module.stream(self.copy_stream)
 
    if not self.enable_overlap:
        # overlap 关闭时无需维护 future token 映射。
        self.future_map = None
        return
 
    # overlap 开启后，FutureMap 负责解析“未来 token 占位符”。
    self.future_map = FutureMap(
        self.max_running_requests,
        self.chunked_prefill_size,
        self.model_config.context_len,
        self.device,
        self.spec_algorithm,
    )
    # 保留最近两个 worker batch 的 Python 引用，避免底层 tensor 被过早回收。
    self.batch_record_buf = [None] * 2
    self.batch_record_ct = 0

def init_overlap(self):
    # 初始化与当前 device 绑定的 stream 上下文。
    self.device_module = torch.get_device_module(self.device)
 
    self.forward_stream_ctx = self.device_module.stream(self.forward_stream)
    self.copy_stream = self.device_module.Stream()
    self.copy_stream_ctx = self.device_module.stream(self.copy_stream)
 
    if not self.enable_overlap:
        # overlap 关闭时无需维护 future token 映射。
        self.future_map = None
        return
 
    # overlap 开启后，FutureMap 负责解析“未来 token 占位符”。
    self.future_map = FutureMap(
        self.max_running_requests,
        self.chunked_prefill_size,
        self.model_config.context_len,
        self.device,
        self.spec_algorithm,
    )
    # 保留最近两个 worker batch 的 Python 引用，避免底层 tensor 被过早回收。
    self.batch_record_buf = [None] * 2
    self.batch_record_ct = 0

def prepare_for_extend(self):
    # 进入 extend/prefill 模式，先从请求里整理 packed token 与长度信息。
    self.forward_mode = ForwardMode.EXTEND
 
    input_ids = [r.fill_ids[len(r.prefix_indices):] for r in reqs]
    extend_num_tokens = sum(len(ids) for ids in input_ids)
    seq_lens = [len(r.fill_ids) for r in reqs]
    prefix_lens = [len(r.prefix_indices) for r in reqs]
    extend_lens = [r.extend_input_len for r in reqs]
 
    # 异步 H2D，把输入 token 和长度元数据提前放到 schedule_stream。
    input_ids_tensor = torch.tensor(
        list(chain.from_iterable(input_ids)), dtype=torch.int64, pin_memory=_pin
    ).to(self.device, non_blocking=True)
    seq_lens_tensor = torch.tensor(seq_lens, dtype=torch.int64, pin_memory=_pin).to(
        self.device, non_blocking=True
    )
 
    # GPU 侧和 CPU 侧各保留一份长度信息，后续分别给 forward/backend 与宿主后处理使用。
    self.prefix_lens = prefix_lens
    self.extend_lens = extend_lens
    self.seq_lens = seq_lens_tensor
    self.seq_lens_cpu = torch.tensor(seq_lens, dtype=torch.int64)
    self.extend_num_tokens = extend_num_tokens
 
    # 提前为本轮 extend 分配 KV cache 写入位置和 req pool 槽位。
    out_cache_loc, req_pool_indices_tensor, req_pool_indices = alloc_for_extend(self)
 
    self.input_ids = input_ids_tensor
    self.req_pool_indices = req_pool_indices_tensor
    self.out_cache_loc = out_cache_loc
    # sampling_info 也在调度阶段准备好，forward 可直接消费。
    self.sampling_info = SamplingBatchInfo.from_schedule_batch(
        self, self.model_config.vocab_size
    )

def prepare_for_extend(self):
    # 进入 extend/prefill 模式，先从请求里整理 packed token 与长度信息。
    self.forward_mode = ForwardMode.EXTEND
 
    input_ids = [r.fill_ids[len(r.prefix_indices):] for r in reqs]
    extend_num_tokens = sum(len(ids) for ids in input_ids)
    seq_lens = [len(r.fill_ids) for r in reqs]
    prefix_lens = [len(r.prefix_indices) for r in reqs]
    extend_lens = [r.extend_input_len for r in reqs]
 
    # 异步 H2D，把输入 token 和长度元数据提前放到 schedule_stream。
    input_ids_tensor = torch.tensor(
        list(chain.from_iterable(input_ids)), dtype=torch.int64, pin_memory=_pin
    ).to(self.device, non_blocking=True)
    seq_lens_tensor = torch.tensor(seq_lens, dtype=torch.int64, pin_memory=_pin).to(
        self.device, non_blocking=True
    )
 
    # GPU 侧和 CPU 侧各保留一份长度信息，后续分别给 forward/backend 与宿主后处理使用。
    self.prefix_lens = prefix_lens
    self.extend_lens = extend_lens
    self.seq_lens = seq_lens_tensor
    self.seq_lens_cpu = torch.tensor(seq_lens, dtype=torch.int64)
    self.extend_num_tokens = extend_num_tokens
 
    # 提前为本轮 extend 分配 KV cache 写入位置和 req pool 槽位。
    out_cache_loc, req_pool_indices_tensor, req_pool_indices = alloc_for_extend(self)
 
    self.input_ids = input_ids_tensor
    self.req_pool_indices = req_pool_indices_tensor
    self.out_cache_loc = out_cache_loc
    # sampling_info 也在调度阶段准备好，forward 可直接消费。
    self.sampling_info = SamplingBatchInfo.from_schedule_batch(
        self, self.model_config.vocab_size
    )

def prepare_for_decode(self):
    # decode 阶段会直接消费上一拍留下的 output_ids。
    self.forward_mode = ForwardMode.DECODE
 
    if self.sampling_info.penalizer_orchestrator.is_required:
        if self.enable_overlap:
            # overlap 下真实 sampled token 还没落回 batch，只能先从 req 状态里取“已确认”的最后一个 token。
            delayed_output_ids = torch.tensor(
                [
                    req.output_ids[-1] if len(req.output_ids)
                    else req.origin_input_ids[-1]
                    for req in self.reqs
                ],
                dtype=torch.int64,
                device=self.device,
            )
            self.sampling_info.penalizer_orchestrator.cumulate_output_tokens(
                delayed_output_ids
            )
        else:
            self.sampling_info.penalizer_orchestrator.cumulate_output_tokens(
                self.output_ids.to(torch.int64)
            )
 
    # 下一拍 decode 的 input_ids 直接来自上一拍 output_ids；它此时可能仍是负数占位符。
    self.input_ids = self.output_ids
    self.output_ids = None
    # decode 每个 request 本轮只写 1 个新 token 的 KV 位置。
    self.out_cache_loc = alloc_for_decode(self, token_per_req=1)
 
    if self.enable_overlap:
        # overlap 下长度需要先行 +1，与“未来 token”保持一致。
        self.seq_lens = self.seq_lens + 1
        self.seq_lens_cpu = self.seq_lens_cpu + 1
        self.orig_seq_lens = self.orig_seq_lens + 1

def prepare_for_decode(self):
    # decode 阶段会直接消费上一拍留下的 output_ids。
    self.forward_mode = ForwardMode.DECODE
 
    if self.sampling_info.penalizer_orchestrator.is_required:
        if self.enable_overlap:
            # overlap 下真实 sampled token 还没落回 batch，只能先从 req 状态里取“已确认”的最后一个 token。
            delayed_output_ids = torch.tensor(
                [
                    req.output_ids[-1] if len(req.output_ids)
                    else req.origin_input_ids[-1]
                    for req in self.reqs
                ],
                dtype=torch.int64,
                device=self.device,
            )
            self.sampling_info.penalizer_orchestrator.cumulate_output_tokens(
                delayed_output_ids
            )
        else:
            self.sampling_info.penalizer_orchestrator.cumulate_output_tokens(
                self.output_ids.to(torch.int64)
            )
 
    # 下一拍 decode 的 input_ids 直接来自上一拍 output_ids；它此时可能仍是负数占位符。
    self.input_ids = self.output_ids
    self.output_ids = None
    # decode 每个 request 本轮只写 1 个新 token 的 KV 位置。
    self.out_cache_loc = alloc_for_decode(self, token_per_req=1)
 
    if self.enable_overlap:
        # overlap 下长度需要先行 +1，与“未来 token”保持一致。
        self.seq_lens = self.seq_lens + 1
        self.seq_lens_cpu = self.seq_lens_cpu + 1
        self.orig_seq_lens = self.orig_seq_lens + 1

def event_loop_overlap(self):
    # 队列里保存“上一拍 launch 的 batch 快照 + 结果”。
    self.result_queue = deque()
 
    def pop_and_process():
        # 处理队首上一拍结果。
        tmp_batch, tmp_result = self.result_queue.popleft()
        self.process_batch_result(tmp_batch, tmp_result)
 
    while True:
        # 1) 接收新请求并更新调度状态。
        recv_reqs = self.recv_requests()
        self.process_input_requests(recv_reqs)
        if self._engine_paused:
            continue
 
        # 2) 选择下一批，并判断本轮是否必须退回串行。
        batch = self.get_next_batch_to_run()
        self.cur_batch = batch
        disable_overlap_for_batch = self.is_disable_overlap_for_batch(batch)
 
        if disable_overlap_for_batch:
            # 特殊场景下先清掉上一拍，避免和当前 batch 发生顺序冲突。
            pop_and_process()
 
        if batch:
            # 3) 先 launch 当前批次，再把快照连同结果压入队列。
            batch_result = self.run_batch(batch)
            self.result_queue.append((batch.copy(), batch_result))
        else:
            batch_result = None
            self.cancel_bubble_timer()
 
        if self.last_batch:
            if not disable_overlap_for_batch:
                # 4) 稳态下在当前 forward 期间处理上一拍结果，实现 CPU/GPU overlap。
                pop_and_process()
        elif batch is None:
            self.self_check_during_idle()
 
        if self.is_generation:
            # grammar 延迟采样必须放在上一拍结果处理之后。
            self.launch_batch_sample_if_needed(batch_result)
 
        # 5) 把当前 batch 记成“上一拍”，供下一轮消费。
        self.last_batch = batch

def event_loop_overlap(self):
    # 队列里保存“上一拍 launch 的 batch 快照 + 结果”。
    self.result_queue = deque()
 
    def pop_and_process():
        # 处理队首上一拍结果。
        tmp_batch, tmp_result = self.result_queue.popleft()
        self.process_batch_result(tmp_batch, tmp_result)
 
    while True:
        # 1) 接收新请求并更新调度状态。
        recv_reqs = self.recv_requests()
        self.process_input_requests(recv_reqs)
        if self._engine_paused:
            continue
 
        # 2) 选择下一批，并判断本轮是否必须退回串行。
        batch = self.get_next_batch_to_run()
        self.cur_batch = batch
        disable_overlap_for_batch = self.is_disable_overlap_for_batch(batch)
 
        if disable_overlap_for_batch:
            # 特殊场景下先清掉上一拍，避免和当前 batch 发生顺序冲突。
            pop_and_process()
 
        if batch:
            # 3) 先 launch 当前批次，再把快照连同结果压入队列。
            batch_result = self.run_batch(batch)
            self.result_queue.append((batch.copy(), batch_result))
        else:
            batch_result = None
            self.cancel_bubble_timer()
 
        if self.last_batch:
            if not disable_overlap_for_batch:
                # 4) 稳态下在当前 forward 期间处理上一拍结果，实现 CPU/GPU overlap。
                pop_and_process()
        elif batch is None:
            self.self_check_during_idle()
 
        if self.is_generation:
            # grammar 延迟采样必须放在上一拍结果处理之后。
            self.launch_batch_sample_if_needed(batch_result)
 
        # 5) 把当前 batch 记成“上一拍”，供下一轮消费。
        self.last_batch = batch

def get_next_batch_to_run(self) -> Optional[ScheduleBatch]:
    if self.last_batch and self.last_batch.forward_mode.is_extend():
        # extend batch 可能会在下一轮开头继续被过滤或并入 running_batch。
        self.last_batch.filter_batch(
            chunked_req_to_exclude=list(chunked_req_to_exclude)
        )
        if not self.last_batch.is_empty():
            if self.running_batch.is_empty():
                self.running_batch = self.last_batch
            else:
                self.running_batch.merge_batch(self.last_batch)

def get_next_batch_to_run(self) -> Optional[ScheduleBatch]:
    if self.last_batch and self.last_batch.forward_mode.is_extend():
        # extend batch 可能会在下一轮开头继续被过滤或并入 running_batch。
        self.last_batch.filter_batch(
            chunked_req_to_exclude=list(chunked_req_to_exclude)
        )
        if not self.last_batch.is_empty():
            if self.running_batch.is_empty():
                self.running_batch = self.last_batch
            else:
                self.running_batch.merge_batch(self.last_batch)

def copy(self):
    # 这里只保留 process_batch_result() 真正会访问到的字段。
    # 对 reqs 列表做浅拷贝，避免原 batch 上的原地修改（如 filter_batch/merge_batch）
    # 污染这份快照；Req 对象本身仍与原 batch 共享。
    return ScheduleBatch(
        reqs=self.reqs[:],
        req_to_token_pool=self.req_to_token_pool,
        req_pool_indices=self.req_pool_indices,
        model_config=self.model_config,
        forward_mode=self.forward_mode,
        out_cache_loc=self.out_cache_loc,
        return_logprob=self.return_logprob,
        spec_algorithm=self.spec_algorithm,
        seq_lens_cpu=self.seq_lens_cpu,
        enable_overlap=self.enable_overlap,
        prefill_stats=self.prefill_stats,
    )

def copy(self):
    # 这里只保留 process_batch_result() 真正会访问到的字段。
    # 对 reqs 列表做浅拷贝，避免原 batch 上的原地修改（如 filter_batch/merge_batch）
    # 污染这份快照；Req 对象本身仍与原 batch 共享。
    return ScheduleBatch(
        reqs=self.reqs[:],
        req_to_token_pool=self.req_to_token_pool,
        req_pool_indices=self.req_pool_indices,
        model_config=self.model_config,
        forward_mode=self.forward_mode,
        out_cache_loc=self.out_cache_loc,
        return_logprob=self.return_logprob,
        spec_algorithm=self.spec_algorithm,
        seq_lens_cpu=self.seq_lens_cpu,
        enable_overlap=self.enable_overlap,
        prefill_stats=self.prefill_stats,
    )

@dataclass
class FutureIndices:
    indices: torch.Tensor
    interval: Optional[slice] = None
 
class FutureMap:
    def __init__(self, max_running_requests, chunked_prefill_size,
                 context_len, device, spec_algo=None):
        # future_ct 是环形分配指针，指向下一段 future 槽位。
        self.future_ct = 0
        self.device = device
        self.spec_algo = spec_algo
 
        # chunked prefill 会让同一 request 同时挂起多个 future token，先估算最坏 chunk 数。
        max_num_chunks = (
            (context_len + chunked_prefill_size - 1) // chunked_prefill_size
            if chunked_prefill_size else 0
        )
        # future_limit 是可轮转的逻辑槽位数；buffer_len 额外留出回旋空间。
        self.future_limit = max_running_requests * (3 + max_num_chunks)
        self.future_buffer_len = self.future_limit + 2 * max_running_requests
 
        if self.spec_algo.is_none():
            # 非 speculative 路径只需要一个 token id 环形缓冲区。
            self.token_ids_buf = torch.empty(
                (self.future_buffer_len,), dtype=torch.int64, device=self.device
            )
 
    def alloc_future_indices(self, bs: int) -> FutureIndices:
        # 为当前 batch 分配连续 future 槽位，并推进环形指针。
        cur_future_ct = self.future_ct
        self.future_ct = (cur_future_ct + bs) % self.future_limit
        start = cur_future_ct + 1
        end = cur_future_ct + 1 + bs
        indices = torch.arange(start, end, dtype=torch.int64, device=self.device)
        return FutureIndices(indices=indices, interval=slice(start, end))
 
    def resolve_future(self, model_worker_batch: ModelWorkerBatch):
        # forward 前原地解析负数占位符，把它们替换成真实 token。
        _resolve_future_token_ids(model_worker_batch.input_ids, self.token_ids_buf)
 
    def store_to_map(self, future_indices: FutureIndices,
                     batch_result: GenerationBatchResult):
        # 将本轮 sample 得到的真实 token 写回对应的 future 槽位。
        self.token_ids_buf[future_indices.interval] = batch_result.next_token_ids

@dataclass
class FutureIndices:
    indices: torch.Tensor
    interval: Optional[slice] = None
 
class FutureMap:
    def __init__(self, max_running_requests, chunked_prefill_size,
                 context_len, device, spec_algo=None):
        # future_ct 是环形分配指针，指向下一段 future 槽位。
        self.future_ct = 0
        self.device = device
        self.spec_algo = spec_algo
 
        # chunked prefill 会让同一 request 同时挂起多个 future token，先估算最坏 chunk 数。
        max_num_chunks = (
            (context_len + chunked_prefill_size - 1) // chunked_prefill_size
            if chunked_prefill_size else 0
        )
        # future_limit 是可轮转的逻辑槽位数；buffer_len 额外留出回旋空间。
        self.future_limit = max_running_requests * (3 + max_num_chunks)
        self.future_buffer_len = self.future_limit + 2 * max_running_requests
 
        if self.spec_algo.is_none():
            # 非 speculative 路径只需要一个 token id 环形缓冲区。
            self.token_ids_buf = torch.empty(
                (self.future_buffer_len,), dtype=torch.int64, device=self.device
            )
 
    def alloc_future_indices(self, bs: int) -> FutureIndices:
        # 为当前 batch 分配连续 future 槽位，并推进环形指针。
        cur_future_ct = self.future_ct
        self.future_ct = (cur_future_ct + bs) % self.future_limit
        start = cur_future_ct + 1
        end = cur_future_ct + 1 + bs
        indices = torch.arange(start, end, dtype=torch.int64, device=self.device)
        return FutureIndices(indices=indices, interval=slice(start, end))
 
    def resolve_future(self, model_worker_batch: ModelWorkerBatch):
        # forward 前原地解析负数占位符，把它们替换成真实 token。
        _resolve_future_token_ids(model_worker_batch.input_ids, self.token_ids_buf)
 
    def store_to_map(self, future_indices: FutureIndices,
                     batch_result: GenerationBatchResult):
        # 将本轮 sample 得到的真实 token 写回对应的 future 槽位。
        self.token_ids_buf[future_indices.interval] = batch_result.next_token_ids

def _resolve_future_token_ids_native(input_ids, future_token_ids_map):
    input_ids[:] = torch.where(
        # input_ids < 0 表示“未来 token 槽位编号”，需要查表替换。
        input_ids < 0,
        future_token_ids_map[torch.clamp(-input_ids, min=0)],
        # 正常 token id 保持不变。
        input_ids,
    )

def _resolve_future_token_ids_native(input_ids, future_token_ids_map):
    input_ids[:] = torch.where(
        # input_ids < 0 表示“未来 token 槽位编号”，需要查表替换。
        input_ids < 0,
        future_token_ids_map[torch.clamp(-input_ids, min=0)],
        # 正常 token id 保持不变。
        input_ids,
    )

# 为当前 batch 预留 future 槽位。
future_indices = self.future_map.alloc_future_indices(bs)
...
# scheduler 侧暂存的是负数占位符，下一拍 decode 再统一解析。
future_indices_or_next_token_ids = -future_indices.indices
batch.output_ids = future_indices_or_next_token_ids

# 为当前 batch 预留 future 槽位。
future_indices = self.future_map.alloc_future_indices(bs)
...
# scheduler 侧暂存的是负数占位符，下一拍 decode 再统一解析。
future_indices_or_next_token_ids = -future_indices.indices
batch.output_ids = future_indices_or_next_token_ids

def prepare_for_decode(self):
    ...
    # 把上一拍留下的 output_ids 直接变成当前拍 input_ids。
    self.input_ids = self.output_ids
    self.output_ids = None

def prepare_for_decode(self):
    ...
    # 把上一拍留下的 output_ids 直接变成当前拍 input_ids。
    self.input_ids = self.output_ids
    self.output_ids = None

with self.forward_stream_ctx:
    # forward 真正开始前，在 forward_stream 上统一解析 future 占位符。
    self.forward_stream.wait_stream(self.schedule_stream)
    self.future_map.resolve_future(model_worker_batch)
    batch_result = self.model_worker.forward_batch_generation(model_worker_batch)

with self.forward_stream_ctx:
    # forward 真正开始前，在 forward_stream 上统一解析 future 占位符。
    self.forward_stream.wait_stream(self.schedule_stream)
    self.future_map.resolve_future(model_worker_batch)
    batch_result = self.model_worker.forward_batch_generation(model_worker_batch)

def run_batch(self, batch: ScheduleBatch, pp_proxy_tensors=None):
    if self.is_generation:
        # overlap 路径会先把 ScheduleBatch 转成更贴近模型执行的 ModelWorkerBatch。
        worker_batch_or_batch = batch.get_model_worker_batch()
 
        if self.enable_overlap:
            model_worker_batch = worker_batch_or_batch
            # 保留最近两个 worker batch 的引用，避免其底层 GPU tensor 被提前释放。
            self.record_batch_in_overlap(model_worker_batch)
 
            # forward 专用 sampling_info 需要复制，避免和调度线程共享可变状态。
            model_worker_batch.sampling_info = (
                model_worker_batch.sampling_info.copy_for_forward()
            )
 
            bs = len(model_worker_batch.seq_lens)
            # 为本批每个 request 预分配 future 槽位。
            future_indices = self.future_map.alloc_future_indices(bs)
 
            with self.forward_stream_ctx, self.record_bubble_metrics(batch):
                # 先等 schedule_stream 上的 H2D / metadata 准备完成。
                self.forward_stream.wait_stream(self.schedule_stream)
                self.future_map.resolve_future(model_worker_batch)
                # 真正 launch 当前 batch 的 forward。
                batch_result = self.model_worker.forward_batch_generation(
                    model_worker_batch
                )
                # copy_done 用于通知 CPU：何时可以安全读取 D2H 结果。
                batch_result.copy_done = self.device_module.Event()
                if batch_result.delay_sample_func is None:
                    # 非延迟采样路径：立即把 token 写回 FutureMap，并发起 D2H。
                    self.future_map.store_to_map(future_indices, batch_result)
                    batch_result.copy_to_cpu(return_logprob=batch.return_logprob)
                else:
                    # grammar 延迟采样路径：先记住 future_indices，稍后再补写。
                    batch_result.future_indices = future_indices
 
            # scheduler 侧只保存负数占位符；真实 token 将在下一拍从 FutureMap 解析出来。
            future_indices_or_next_token_ids = -future_indices.indices
            batch.output_ids = future_indices_or_next_token_ids

def run_batch(self, batch: ScheduleBatch, pp_proxy_tensors=None):
    if self.is_generation:
        # overlap 路径会先把 ScheduleBatch 转成更贴近模型执行的 ModelWorkerBatch。
        worker_batch_or_batch = batch.get_model_worker_batch()
 
        if self.enable_overlap:
            model_worker_batch = worker_batch_or_batch
            # 保留最近两个 worker batch 的引用，避免其底层 GPU tensor 被提前释放。
            self.record_batch_in_overlap(model_worker_batch)
 
            # forward 专用 sampling_info 需要复制，避免和调度线程共享可变状态。
            model_worker_batch.sampling_info = (
                model_worker_batch.sampling_info.copy_for_forward()
            )
 
            bs = len(model_worker_batch.seq_lens)
            # 为本批每个 request 预分配 future 槽位。
            future_indices = self.future_map.alloc_future_indices(bs)
 
            with self.forward_stream_ctx, self.record_bubble_metrics(batch):
                # 先等 schedule_stream 上的 H2D / metadata 准备完成。
                self.forward_stream.wait_stream(self.schedule_stream)
                self.future_map.resolve_future(model_worker_batch)
                # 真正 launch 当前 batch 的 forward。
                batch_result = self.model_worker.forward_batch_generation(
                    model_worker_batch
                )
                # copy_done 用于通知 CPU：何时可以安全读取 D2H 结果。
                batch_result.copy_done = self.device_module.Event()
                if batch_result.delay_sample_func is None:
                    # 非延迟采样路径：立即把 token 写回 FutureMap，并发起 D2H。
                    self.future_map.store_to_map(future_indices, batch_result)
                    batch_result.copy_to_cpu(return_logprob=batch.return_logprob)
                else:
                    # grammar 延迟采样路径：先记住 future_indices，稍后再补写。
                    batch_result.future_indices = future_indices
 
            # scheduler 侧只保存负数占位符；真实 token 将在下一拍从 FutureMap 解析出来。
            future_indices_or_next_token_ids = -future_indices.indices
            batch.output_ids = future_indices_or_next_token_ids

def forward_batch_generation(self, model_worker_batch: ModelWorkerBatch, ...):
    # 先执行模型 forward，拿到 logits 与是否可复用 CUDA Graph 的信息。
    out = self.model_runner.forward(forward_batch, ...)
    logits_output, can_run_cuda_graph = out.logits_output, out.can_run_graph
    batch_result = GenerationBatchResult(
        logits_output=logits_output,
        can_run_cuda_graph=can_run_cuda_graph,
        expert_distribution_metrics=out.expert_distribution_metrics,
    )
 
    if (
        self.enable_overlap
        and not self.enable_spec
        and model_worker_batch.sampling_info.grammars is not None
    ):
        def sample_batch_func():
            # grammar + overlap 场景下，把真正的 sample 延迟到稍后执行。
            batch_result.next_token_ids = self.model_runner.sample(
                logits_output, forward_batch
            )
            return batch_result
 
        batch_result.delay_sample_func = sample_batch_func
        return batch_result
 
    if not model_worker_batch.is_prefill_only:
        # decode / mixed batch 直接采样出 next token。
        batch_result.next_token_ids = self.model_runner.sample(
            logits_output, forward_batch
        )
    else:
        # pure prefill 不需要真实 next token，这里放一个占位 tensor。
        batch_result.next_token_ids = torch.zeros(
            len(model_worker_batch.seq_lens),
            dtype=torch.long,
            device=model_worker_batch.input_ids.device,
        )

def forward_batch_generation(self, model_worker_batch: ModelWorkerBatch, ...):
    # 先执行模型 forward，拿到 logits 与是否可复用 CUDA Graph 的信息。
    out = self.model_runner.forward(forward_batch, ...)
    logits_output, can_run_cuda_graph = out.logits_output, out.can_run_graph
    batch_result = GenerationBatchResult(
        logits_output=logits_output,
        can_run_cuda_graph=can_run_cuda_graph,
        expert_distribution_metrics=out.expert_distribution_metrics,
    )
 
    if (
        self.enable_overlap
        and not self.enable_spec
        and model_worker_batch.sampling_info.grammars is not None
    ):
        def sample_batch_func():
            # grammar + overlap 场景下，把真正的 sample 延迟到稍后执行。
            batch_result.next_token_ids = self.model_runner.sample(
                logits_output, forward_batch
            )
            return batch_result
 
        batch_result.delay_sample_func = sample_batch_func
        return batch_result
 
    if not model_worker_batch.is_prefill_only:
        # decode / mixed batch 直接采样出 next token。
        batch_result.next_token_ids = self.model_runner.sample(
            logits_output, forward_batch
        )
    else:
        # pure prefill 不需要真实 next token，这里放一个占位 tensor。
        batch_result.next_token_ids = torch.zeros(
            len(model_worker_batch.seq_lens),
            dtype=torch.long,
            device=model_worker_batch.input_ids.device,
        )

def copy_to_cpu(self, return_logprob: bool):
    if return_logprob:
        ...
    if self.logits_output.hidden_states is not None:
        # 仅在需要时把 hidden states 异步拷回 CPU。
        self.logits_output.hidden_states = self.logits_output.hidden_states.to(
            "cpu", non_blocking=True
        )
    # next_token_ids 总是需要给 CPU 侧后处理消费。
    self.next_token_ids = self.next_token_ids.to("cpu", non_blocking=True)
 
    if self.accept_lens is not None:
        # speculative 路径的 acceptance 长度若存在，也一并异步拷回。
        self.accept_lens = self.accept_lens.to("cpu", non_blocking=True)
 
    # 在当前 stream 上记录拷贝完成事件，供 CPU 侧 synchronize。
    self.copy_done.record()

def copy_to_cpu(self, return_logprob: bool):
    if return_logprob:
        ...
    if self.logits_output.hidden_states is not None:
        # 仅在需要时把 hidden states 异步拷回 CPU。
        self.logits_output.hidden_states = self.logits_output.hidden_states.to(
            "cpu", non_blocking=True
        )
    # next_token_ids 总是需要给 CPU 侧后处理消费。
    self.next_token_ids = self.next_token_ids.to("cpu", non_blocking=True)
 
    if self.accept_lens is not None:
        # speculative 路径的 acceptance 长度若存在，也一并异步拷回。
        self.accept_lens = self.accept_lens.to("cpu", non_blocking=True)
 
    # 在当前 stream 上记录拷贝完成事件，供 CPU 侧 synchronize。
    self.copy_done.record()

if result.copy_done is not None:
    # CPU 真正读取结果前，先等待异步 D2H 完成。
    result.copy_done.synchronize()

if result.copy_done is not None:
    # CPU 真正读取结果前，先等待异步 D2H 完成。
    result.copy_done.synchronize()

def process_batch_result_prefill(self, batch, result):
    if result.copy_done is not None:
        # 上一拍 D2H 未完成时，CPU 这里会显式等待。
        result.copy_done.synchronize()
 
    logits_output = result.logits_output
    # 转成 Python list，便于逐 request 更新宿主侧状态。
    next_token_ids = result.next_token_ids.tolist()
 
    for i, (req, next_token_id) in enumerate(zip(batch.reqs, next_token_ids)):
        if req.finished() or req.is_retracted:
            # 已完成或已回收的请求不再处理。
            continue
 
        if req.is_chunked <= 0:
            # 非 chunked 情况：正式落 token，并更新 finish / cache / 返回值。
            req.output_ids.append(next_token_id)
            req.check_finished()
            if req.finished():
                release_kv_cache(req, self.tree_cache)
            elif not batch.decoding_reqs or req not in batch.decoding_reqs:
                self.tree_cache.cache_unfinished_req(req)
 
            if batch.return_logprob:
                self.add_logprob_return_values(...)
 
            if req.grammar is not None:
                # grammar 也要跟着消费这个 token，保持状态同步。
                req.grammar.accept_token(next_token_id)
                req.grammar.finished = req.finished()
        else:
            # chunked prefill 只推进剩余 chunk 计数，不立即输出 token。
            req.is_chunked -= 1

def process_batch_result_prefill(self, batch, result):
    if result.copy_done is not None:
        # 上一拍 D2H 未完成时，CPU 这里会显式等待。
        result.copy_done.synchronize()
 
    logits_output = result.logits_output
    # 转成 Python list，便于逐 request 更新宿主侧状态。
    next_token_ids = result.next_token_ids.tolist()
 
    for i, (req, next_token_id) in enumerate(zip(batch.reqs, next_token_ids)):
        if req.finished() or req.is_retracted:
            # 已完成或已回收的请求不再处理。
            continue
 
        if req.is_chunked <= 0:
            # 非 chunked 情况：正式落 token，并更新 finish / cache / 返回值。
            req.output_ids.append(next_token_id)
            req.check_finished()
            if req.finished():
                release_kv_cache(req, self.tree_cache)
            elif not batch.decoding_reqs or req not in batch.decoding_reqs:
                self.tree_cache.cache_unfinished_req(req)
 
            if batch.return_logprob:
                self.add_logprob_return_values(...)
 
            if req.grammar is not None:
                # grammar 也要跟着消费这个 token，保持状态同步。
                req.grammar.accept_token(next_token_id)
                req.grammar.finished = req.finished()
        else:
            # chunked prefill 只推进剩余 chunk 计数，不立即输出 token。
            req.is_chunked -= 1

def process_batch_result_decode(self, batch, result):
    if result.copy_done is not None:
        # decode 结果在 CPU 侧消费前同样要先等 D2H。
        result.copy_done.synchronize()
 
    logits_output, next_token_ids = result.logits_output, result.next_token_ids
 
    if batch.spec_algorithm.is_none() or batch.is_spec_v2:
        if batch.is_spec_v2:
            # spec v2 overlap 需要先把暂存结果解析回真实 token 序列。
            next_token_ids = self._resolve_spec_overlap_token_ids(result, batch)
        else:
            next_token_ids = next_token_ids.tolist()
 
    for i, req in enumerate(batch.reqs):
        if self.enable_overlap and (req.finished() or req.is_retracted):
            # overlap 下已结束 / 已回收的请求直接跳过。
            continue
 
        next_token_id = next_token_ids[i]
        if batch.spec_algorithm.is_none():
            # 普通 decode 每轮只追加一个 token。
            req.output_ids.append(next_token_id)
        else:
            # speculative 路径可能一次性接受多个 token。
            req.output_ids.extend(next_token_id)
 
        req.check_finished(...)
        # 统一处理 finished request 的 cache 回收及后续善后。
        self._handle_finished_req(req, i, logits_output)
 
        if req.return_logprob:
            ...
        if req.grammar is not None:
            # decode 路径同样要推进 grammar 状态。
            req.grammar.accept_token(next_token_id)
            req.grammar.finished = req.finished()
 
    # 最后把这一拍可输出的内容统一流式发出去。
    self.stream_output(batch.reqs, batch.return_logprob)

def process_batch_result_decode(self, batch, result):
    if result.copy_done is not None:
        # decode 结果在 CPU 侧消费前同样要先等 D2H。
        result.copy_done.synchronize()
 
    logits_output, next_token_ids = result.logits_output, result.next_token_ids
 
    if batch.spec_algorithm.is_none() or batch.is_spec_v2:
        if batch.is_spec_v2:
            # spec v2 overlap 需要先把暂存结果解析回真实 token 序列。
            next_token_ids = self._resolve_spec_overlap_token_ids(result, batch)
        else:
            next_token_ids = next_token_ids.tolist()
 
    for i, req in enumerate(batch.reqs):
        if self.enable_overlap and (req.finished() or req.is_retracted):
            # overlap 下已结束 / 已回收的请求直接跳过。
            continue
 
        next_token_id = next_token_ids[i]
        if batch.spec_algorithm.is_none():
            # 普通 decode 每轮只追加一个 token。
            req.output_ids.append(next_token_id)
        else:
            # speculative 路径可能一次性接受多个 token。
            req.output_ids.extend(next_token_id)
 
        req.check_finished(...)
        # 统一处理 finished request 的 cache 回收及后续善后。
        self._handle_finished_req(req, i, logits_output)
 
        if req.return_logprob:
            ...
        if req.grammar is not None:
            # decode 路径同样要推进 grammar 状态。
            req.grammar.accept_token(next_token_id)
            req.grammar.finished = req.finished()
 
    # 最后把这一拍可输出的内容统一流式发出去。
    self.stream_output(batch.reqs, batch.return_logprob)

def launch_batch_sample_if_needed(self, batch_result):
    # 只有 grammar 延迟采样路径才会走到这里。
    if batch_result is None or batch_result.delay_sample_func is None:
        return
 
    with self.forward_stream_ctx:
        # 采样前仍然要遵守 schedule_stream -> forward_stream 的依赖。
        self.forward_stream.wait_stream(self.schedule_stream)
        # 真正执行延迟 sample，并复用原来的 batch_result 容器。
        _batch_result = batch_result.delay_sample_func()
        assert _batch_result is batch_result
        # 采样完成后立刻补写 FutureMap，再发起 D2H。
        self.future_map.store_to_map(batch_result.future_indices, batch_result)
        batch_result.copy_to_cpu(return_logprob=self.cur_batch.return_logprob)
 
    # 清掉闭包和 logits 引用，避免大块 GPU 内存被闭包长期持有。
    batch_result.delay_sample_func = None
    if batch_result.logits_output is not None:
        batch_result.logits_output.next_token_logits = None

def launch_batch_sample_if_needed(self, batch_result):
    # 只有 grammar 延迟采样路径才会走到这里。
    if batch_result is None or batch_result.delay_sample_func is None:
        return
 
    with self.forward_stream_ctx:
        # 采样前仍然要遵守 schedule_stream -> forward_stream 的依赖。
        self.forward_stream.wait_stream(self.schedule_stream)
        # 真正执行延迟 sample，并复用原来的 batch_result 容器。
        _batch_result = batch_result.delay_sample_func()
        assert _batch_result is batch_result
        # 采样完成后立刻补写 FutureMap，再发起 D2H。
        self.future_map.store_to_map(batch_result.future_indices, batch_result)
        batch_result.copy_to_cpu(return_logprob=self.cur_batch.return_logprob)
 
    # 清掉闭包和 logits 引用，避免大块 GPU 内存被闭包长期持有。
    batch_result.delay_sample_func = None
    if batch_result.logits_output is not None:
        batch_result.logits_output.next_token_logits = None

def is_disable_overlap_for_batch(self, batch: ScheduleBatch) -> bool:
    if self.require_mlp_sync:
        # 某些路径下 extend 判断要改看 is_extend_in_batch。
        is_extend = lambda b: b and b.is_extend_in_batch
    else:
        is_extend = lambda b: b and b.forward_mode.is_extend()
 
    batch_is_extend = is_extend(batch)
    last_batch_is_extend = is_extend(self.last_batch)
 
    disable_overlap_for_batch = (
        # 可通过环境变量禁止连续两个 prefill 重叠，优先优化第一批 TTFT。
        envs.SGLANG_DISABLE_CONSECUTIVE_PREFILL_OVERLAP.get()
        and batch_is_extend
        and last_batch_is_extend
    )
 
    need_grammar_sync = (
        # spec v2 + grammar + decode 场景下，需要先把上一拍结果彻底处理完。
        batch
        and batch.is_spec_v2
        and batch.has_grammar
        and batch.forward_mode.is_decode()
        and len(self.result_queue) > 0
    )
 
    # 任一条件成立时，本轮都退回串行处理。
    return disable_overlap_for_batch or need_grammar_sync

def is_disable_overlap_for_batch(self, batch: ScheduleBatch) -> bool:
    if self.require_mlp_sync:
        # 某些路径下 extend 判断要改看 is_extend_in_batch。
        is_extend = lambda b: b and b.is_extend_in_batch
    else:
        is_extend = lambda b: b and b.forward_mode.is_extend()
 
    batch_is_extend = is_extend(batch)
    last_batch_is_extend = is_extend(self.last_batch)
 
    disable_overlap_for_batch = (
        # 可通过环境变量禁止连续两个 prefill 重叠，优先优化第一批 TTFT。
        envs.SGLANG_DISABLE_CONSECUTIVE_PREFILL_OVERLAP.get()
        and batch_is_extend
        and last_batch_is_extend
    )
 
    need_grammar_sync = (
        # spec v2 + grammar + decode 场景下，需要先把上一拍结果彻底处理完。
        batch
        and batch.is_spec_v2
        and batch.has_grammar
        and batch.forward_mode.is_decode()
        and len(self.result_queue) > 0
    )
 
    # 任一条件成立时，本轮都退回串行处理。
    return disable_overlap_for_batch or need_grammar_sync

def _handle_pipeline_parallelism(self):
    if self.pp_size > 1:
        # Pipeline Parallelism 当前直接禁用 overlap schedule。
        self.disable_overlap_schedule = True
 
def _handle_mps_backends(self):
    if self.device == "mps":
        # Apple MPS 后端不走这套 overlap。
        self.disable_overlap_schedule = True
 
if envs.SGLANG_EMBEDDINGS_SPARSE_HEAD.is_set():
    # 稀疏 embedding head 与 overlap 组合当前不支持。
    self.disable_overlap_schedule = True
 
if self.mamba_scheduler_strategy == "no_buffer" and not self.disable_radix_cache:
    # Mamba no_buffer + radix cache 组合会在启动时关闭 overlap。
    self.disable_overlap_schedule = True
 
if self.speculative_algorithm == "NGRAM":
    # NGRAM speculative decoding 当前不支持 overlap。
    self.disable_overlap_schedule = True
 
if not self.disable_overlap_schedule:
    self.disable_overlap_schedule = True
    # diffusion LLM 推理会统一关闭 overlap。
    logger.warning("Overlap schedule is disabled because of using diffusion LLM inference")

def _handle_pipeline_parallelism(self):
    if self.pp_size > 1:
        # Pipeline Parallelism 当前直接禁用 overlap schedule。
        self.disable_overlap_schedule = True
 
def _handle_mps_backends(self):
    if self.device == "mps":
        # Apple MPS 后端不走这套 overlap。
        self.disable_overlap_schedule = True
 
if envs.SGLANG_EMBEDDINGS_SPARSE_HEAD.is_set():
    # 稀疏 embedding head 与 overlap 组合当前不支持。
    self.disable_overlap_schedule = True
 
if self.mamba_scheduler_strategy == "no_buffer" and not self.disable_radix_cache:
    # Mamba no_buffer + radix cache 组合会在启动时关闭 overlap。
    self.disable_overlap_schedule = True
 
if self.speculative_algorithm == "NGRAM":
    # NGRAM speculative decoding 当前不支持 overlap。
    self.disable_overlap_schedule = True
 
if not self.disable_overlap_schedule:
    self.disable_overlap_schedule = True
    # diffusion LLM 推理会统一关闭 overlap。
    logger.warning("Overlap schedule is disabled because of using diffusion LLM inference")

# disaggregation prefill / decode 也沿用 PP、overlap、normal 三套事件循环骨架。
if disaggregation_mode == DisaggregationMode.PREFILL:
    if server_args.pp_size > 1:
        scheduler.event_loop_pp_disagg_prefill()
    elif scheduler.enable_overlap:
        scheduler.event_loop_overlap_disagg_prefill()
    else:
        scheduler.event_loop_normal_disagg_prefill()
elif disaggregation_mode == DisaggregationMode.DECODE:
    if server_args.pp_size > 1:
        scheduler.event_loop_pp_disagg_decode()
    elif scheduler.enable_overlap:
        scheduler.event_loop_overlap_disagg_decode()
    else:
        scheduler.event_loop_normal_disagg_decode()

# disaggregation prefill / decode 也沿用 PP、overlap、normal 三套事件循环骨架。
if disaggregation_mode == DisaggregationMode.PREFILL:
    if server_args.pp_size > 1:
        scheduler.event_loop_pp_disagg_prefill()
    elif scheduler.enable_overlap:
        scheduler.event_loop_overlap_disagg_prefill()
    else:
        scheduler.event_loop_normal_disagg_prefill()
elif disaggregation_mode == DisaggregationMode.DECODE:
    if server_args.pp_size > 1:
        scheduler.event_loop_pp_disagg_decode()
    elif scheduler.enable_overlap:
        scheduler.event_loop_overlap_disagg_decode()
    else:
        scheduler.event_loop_normal_disagg_decode()

hidden_states
  -> gate(router_logits)
  -> topk(topk_ids, topk_weights)
  -> dispatcher.dispatch(...)
  -> experts.run_moe_core(...)
       -> expert up/gate gemm
       -> down gemm
  -> dispatcher.combine(...)
  -> add shared experts output

hidden_states
  -> gate(router_logits)
  -> topk(topk_ids, topk_weights)
  -> dispatcher.dispatch(...)
  -> experts.run_moe_core(...)
       -> expert up/gate gemm
       -> down gemm
  -> dispatcher.combine(...)
  -> add shared experts output

# server_args.py
# SBO 总开关，默认关闭。
enable_single_batch_overlap: bool = False
 
# moe/utils.py
# 运行时用这个全局常量派生更细粒度的 SBO 模式判断。
IS_SBO_ENABLED = server_args.enable_single_batch_overlap

# server_args.py
# SBO 总开关，默认关闭。
enable_single_batch_overlap: bool = False
 
# moe/utils.py
# 运行时用这个全局常量派生更细粒度的 SBO 模式判断。
IS_SBO_ENABLED = server_args.enable_single_batch_overlap

class SboFlags:
    # TODO 后续可能还会补更多模式，例如 dispatch-gateup gemm 双流重叠等。
 
    @classmethod
    def enable_combine_down_gemm_two_stream_overlap(cls):
        return (
            is_sbo_enabled()
            # 目前只有 CuteDSL，或非 Blackwell 的 DeepGEMM 后端真正支持这一路。
            and (
                get_moe_runner_backend().is_flashinfer_cutedsl()
                or (get_moe_runner_backend().is_deep_gemm() and not is_blackwell())
            )
        )
 
    @classmethod
    def enable_combine_shared_two_stream_overlap(cls):
        return (
            is_sbo_enabled()
            # combine-shared 与 dispatch-shared 互斥；Blackwell 默认更接近这一路。
            and not cls.enable_dispatch_shared_one_stream_overlap()
            and not envs.SGLANG_BLACKWELL_OVERLAP_SHARED_EXPERTS_OUTSIDE_SBO.get()
        )
 
    @classmethod
    def enable_dispatch_shared_one_stream_overlap(cls):
        # 非 Blackwell 默认倾向把 shared experts 插到 dispatch 中间。
        return is_sbo_enabled() and not is_blackwell()
 
    @classmethod
    def fuse_shared_experts_inside_sbo(cls):
        # 只要任一 shared-experts overlap 模式开启，就把 shared experts 融到 SBO 内部。
        return (
            cls.enable_combine_shared_two_stream_overlap()
            or cls.enable_dispatch_shared_one_stream_overlap()
        )

class SboFlags:
    # TODO 后续可能还会补更多模式，例如 dispatch-gateup gemm 双流重叠等。
 
    @classmethod
    def enable_combine_down_gemm_two_stream_overlap(cls):
        return (
            is_sbo_enabled()
            # 目前只有 CuteDSL，或非 Blackwell 的 DeepGEMM 后端真正支持这一路。
            and (
                get_moe_runner_backend().is_flashinfer_cutedsl()
                or (get_moe_runner_backend().is_deep_gemm() and not is_blackwell())
            )
        )
 
    @classmethod
    def enable_combine_shared_two_stream_overlap(cls):
        return (
            is_sbo_enabled()
            # combine-shared 与 dispatch-shared 互斥；Blackwell 默认更接近这一路。
            and not cls.enable_dispatch_shared_one_stream_overlap()
            and not envs.SGLANG_BLACKWELL_OVERLAP_SHARED_EXPERTS_OUTSIDE_SBO.get()
        )
 
    @classmethod
    def enable_dispatch_shared_one_stream_overlap(cls):
        # 非 Blackwell 默认倾向把 shared experts 插到 dispatch 中间。
        return is_sbo_enabled() and not is_blackwell()
 
    @classmethod
    def fuse_shared_experts_inside_sbo(cls):
        # 只要任一 shared-experts overlap 模式开启，就把 shared experts 融到 SBO 内部。
        return (
            cls.enable_combine_shared_two_stream_overlap()
            or cls.enable_dispatch_shared_one_stream_overlap()
        )

# single_batch_overlap.py
hidden_states = dispatch_output.hidden_states
num_local_experts, num_tokens_static, hidden_dim = hidden_states.shape
 
# ep_moe/kernels.py
"""
input shape [expert_num, token_num_padded, hidden_dim]
...
masked_m shape [expert_num],
"""
assert input.shape[0] == masked_m.shape[0]

# single_batch_overlap.py
hidden_states = dispatch_output.hidden_states
num_local_experts, num_tokens_static, hidden_dim = hidden_states.shape
 
# ep_moe/kernels.py
"""
input shape [expert_num, token_num_padded, hidden_dim]
...
masked_m shape [expert_num],
"""
assert input.shape[0] == masked_m.shape[0]

@dataclass
class CombineOverlapArgs:
    # 这里的 overlap 标志只表示“是否与 down gemm 重叠”，
    # 不是泛指所有双 stream overlap。
    overlap: bool
    stream: torch.cuda.Stream
    wait_event: torch.cuda.Event
    num_sms: Optional[int] = None
    signal: Optional[torch.Tensor] = None
    block_m: Optional[int] = 64
    threshold: Optional[int] = 0
 
@dataclass
class DownGemmOverlapArgs:
    # 计算侧会用它拿到 SM 配额、signal tensor 与 start_event。
    num_sms: int
    signal: torch.Tensor
    start_event: torch.cuda.Event
 
def compute_overlap_args(dispatch_output, alt_stream):
    # 只有 combine-down 或 combine-shared 开启时，才需要构造 overlap 参数。
    if not (
        SboFlags.enable_combine_down_gemm_two_stream_overlap()
        or SboFlags.enable_combine_shared_two_stream_overlap()
    ):
        return None, None, {}
 
    # 这里假定 dispatch_output 已经是 DeepEP low-latency 的 3D packed 布局。
    hidden_states = dispatch_output.hidden_states
 
    num_local_experts, num_tokens_static, hidden_dim = hidden_states.shape
 
    # 读取设备总 SM 数，并切分给通信侧与计算侧。
    total_num_sms = torch.cuda.get_device_properties(
        device="cuda"
    ).multi_processor_count
 
    if envs.SGLANG_DEEPEP_LL_COMBINE_SEND_NUM_SMS.is_set():
        # 环境变量可以覆盖 combine 通信侧的默认 SM 预算。
        communicate_num_sms = envs.SGLANG_DEEPEP_LL_COMBINE_SEND_NUM_SMS.get()
    else:
        communicate_num_sms = 32 if is_blackwell() else 3
    compute_num_sms = total_num_sms - communicate_num_sms
 
    assert alt_stream is not None
    # combine 将切到 alt_stream 上执行，并在 wait_event 处等待起跑。
    combine_wait_event = torch.cuda.Event()
    combine_overlap_args = CombineOverlapArgs(
        overlap=False,
        num_sms=communicate_num_sms,
        stream=alt_stream,
        wait_event=combine_wait_event,
    )
    # 这个共享 dict 会在 down GEMM 与 combine 之间继续回填和传递。
    meta_overlap_args = dict(
        compute_num_sms=compute_num_sms,
    )
    down_gemm_overlap_args = None
 
    if SboFlags.enable_combine_down_gemm_two_stream_overlap():
        # combine-down 模式下，需要 signal tensor 让计算侧和通信侧做细粒度协作。
        # TODO 可改用 zero_allocator，避免这里额外的 torch.zeros 分配。
        # 注意：当前 v2 路径里 Blackwell 使用 uint32，而不是 int32。
        if is_blackwell():
            combine_signal = torch.zeros(
                num_local_experts, dtype=torch.uint32, device=hidden_states.device
            )
        else:
            MIN_BLOCK_M = 64
            combine_signal_size = num_local_experts * (
                (num_tokens_static + MIN_BLOCK_M - 1) // MIN_BLOCK_M
            )
            combine_signal = torch.zeros(
                combine_signal_size, dtype=torch.int32, device=hidden_states.device
            )
 
        down_gemm_overlap_args = DownGemmOverlapArgs(
            signal=combine_signal,
            start_event=combine_wait_event,
            num_sms=compute_num_sms,
        )
        combine_overlap_args.overlap = True
        combine_overlap_args.signal = combine_signal
        combine_overlap_args.threshold = compute_num_sms
    else:
        # combine-shared 模式下，不做 signal 协作，只在 down 之后打一条 event 边界。
        meta_overlap_args |= dict(
            record_event_after_down=combine_wait_event,
        )
 
    return combine_overlap_args, down_gemm_overlap_args, meta_overlap_args

@dataclass
class CombineOverlapArgs:
    # 这里的 overlap 标志只表示“是否与 down gemm 重叠”，
    # 不是泛指所有双 stream overlap。
    overlap: bool
    stream: torch.cuda.Stream
    wait_event: torch.cuda.Event
    num_sms: Optional[int] = None
    signal: Optional[torch.Tensor] = None
    block_m: Optional[int] = 64
    threshold: Optional[int] = 0
 
@dataclass
class DownGemmOverlapArgs:
    # 计算侧会用它拿到 SM 配额、signal tensor 与 start_event。
    num_sms: int
    signal: torch.Tensor
    start_event: torch.cuda.Event
 
def compute_overlap_args(dispatch_output, alt_stream):
    # 只有 combine-down 或 combine-shared 开启时，才需要构造 overlap 参数。
    if not (
        SboFlags.enable_combine_down_gemm_two_stream_overlap()
        or SboFlags.enable_combine_shared_two_stream_overlap()
    ):
        return None, None, {}
 
    # 这里假定 dispatch_output 已经是 DeepEP low-latency 的 3D packed 布局。
    hidden_states = dispatch_output.hidden_states
 
    num_local_experts, num_tokens_static, hidden_dim = hidden_states.shape
 
    # 读取设备总 SM 数，并切分给通信侧与计算侧。
    total_num_sms = torch.cuda.get_device_properties(
        device="cuda"
    ).multi_processor_count
 
    if envs.SGLANG_DEEPEP_LL_COMBINE_SEND_NUM_SMS.is_set():
        # 环境变量可以覆盖 combine 通信侧的默认 SM 预算。
        communicate_num_sms = envs.SGLANG_DEEPEP_LL_COMBINE_SEND_NUM_SMS.get()
    else:
        communicate_num_sms = 32 if is_blackwell() else 3
    compute_num_sms = total_num_sms - communicate_num_sms
 
    assert alt_stream is not None
    # combine 将切到 alt_stream 上执行，并在 wait_event 处等待起跑。
    combine_wait_event = torch.cuda.Event()
    combine_overlap_args = CombineOverlapArgs(
        overlap=False,
        num_sms=communicate_num_sms,
        stream=alt_stream,
        wait_event=combine_wait_event,
    )
    # 这个共享 dict 会在 down GEMM 与 combine 之间继续回填和传递。
    meta_overlap_args = dict(
        compute_num_sms=compute_num_sms,
    )
    down_gemm_overlap_args = None
 
    if SboFlags.enable_combine_down_gemm_two_stream_overlap():
        # combine-down 模式下，需要 signal tensor 让计算侧和通信侧做细粒度协作。
        # TODO 可改用 zero_allocator，避免这里额外的 torch.zeros 分配。
        # 注意：当前 v2 路径里 Blackwell 使用 uint32，而不是 int32。
        if is_blackwell():
            combine_signal = torch.zeros(
                num_local_experts, dtype=torch.uint32, device=hidden_states.device
            )
        else:
            MIN_BLOCK_M = 64
            combine_signal_size = num_local_experts * (
                (num_tokens_static + MIN_BLOCK_M - 1) // MIN_BLOCK_M
            )
            combine_signal = torch.zeros(
                combine_signal_size, dtype=torch.int32, device=hidden_states.device
            )
 
        down_gemm_overlap_args = DownGemmOverlapArgs(
            signal=combine_signal,
            start_event=combine_wait_event,
            num_sms=compute_num_sms,
        )
        combine_overlap_args.overlap = True
        combine_overlap_args.signal = combine_signal
        combine_overlap_args.threshold = compute_num_sms
    else:
        # combine-shared 模式下，不做 signal 协作，只在 down 之后打一条 event 边界。
        meta_overlap_args |= dict(
            record_event_after_down=combine_wait_event,
        )
 
    return combine_overlap_args, down_gemm_overlap_args, meta_overlap_args

class BaseDispatcher(ABC):
    def __init__(self):
        self.quant_config: Optional[dict] = None
 
        # overlap 参数由模型层在 post_dispatch hook 中注入。
        self.overlap_args: Optional[CombineOverlapArgs] = None
        self.meta_overlap_args: Optional[dict] = None
 
        # 四类 hook 分别钩住 dispatch/combine 的前后时机。
        self._pre_dispatch_hooks: Optional[_PreDispatchHooks] = None
        self._post_dispatch_hooks: Optional[_PostDispatchHooks] = None
        self._pre_combine_hooks: Optional[_PreCombineHooks] = None
        self._post_combine_hooks: Optional[_PostCombineHooks] = None
        self._original_dispatch_func: Optional[Callable] = None
        self._original_combine_func: Optional[Callable] = None
 
    def _dispatch_with_hook(
        self, hidden_states: torch.Tensor, topk_output: TopKOutput
    ) -> DispatchOutput:
        if self._pre_dispatch_hooks is not None:
            # pre hook 可以改写输入，再交给原始 dispatch 实现。
            hidden_states, topk_output = self._pre_dispatch_hooks(
                self, hidden_states, topk_output
            )
        dispatch_output = self._original_dispatch_func(
            hidden_states=hidden_states, topk_output=topk_output
        )
        if self._post_dispatch_hooks is not None:
            # post hook 可以基于完整 dispatch_output 继续注入 SBO 参数。
            dispatch_output = self._post_dispatch_hooks(self, dispatch_output)
        return dispatch_output
 
    def _combine_with_hook(self, combine_input: CombineInput) -> torch.Tensor:
        if self._pre_combine_hooks is not None:
            # combine 前 hook 常用于记录 event 或排 shared experts。
            combine_input = self._pre_combine_hooks(self, combine_input)
        hidden_states = self._original_combine_func(combine_input=combine_input)
        if self._post_combine_hooks is not None:
            # combine 后 hook 主要做清理，防止状态泄漏到下一层。
            hidden_states = self._post_combine_hooks(self, hidden_states)
        return hidden_states
 
    def set_overlap_args(
        self, combine_overlap_args: CombineOverlapArgs, meta_overlap_args: dict
    ) -> None:
        # 将 combine 侧 overlap 参数和共享元数据挂到 dispatcher 上。
        self.overlap_args = combine_overlap_args
        self.meta_overlap_args = meta_overlap_args
 
    def clear_overlap_args(self) -> None:
        # 每层 / 每批结束后清理，避免串到下一次调用。
        self.overlap_args = None
        self.meta_overlap_args = None

class BaseDispatcher(ABC):
    def __init__(self):
        self.quant_config: Optional[dict] = None
 
        # overlap 参数由模型层在 post_dispatch hook 中注入。
        self.overlap_args: Optional[CombineOverlapArgs] = None
        self.meta_overlap_args: Optional[dict] = None
 
        # 四类 hook 分别钩住 dispatch/combine 的前后时机。
        self._pre_dispatch_hooks: Optional[_PreDispatchHooks] = None
        self._post_dispatch_hooks: Optional[_PostDispatchHooks] = None
        self._pre_combine_hooks: Optional[_PreCombineHooks] = None
        self._post_combine_hooks: Optional[_PostCombineHooks] = None
        self._original_dispatch_func: Optional[Callable] = None
        self._original_combine_func: Optional[Callable] = None
 
    def _dispatch_with_hook(
        self, hidden_states: torch.Tensor, topk_output: TopKOutput
    ) -> DispatchOutput:
        if self._pre_dispatch_hooks is not None:
            # pre hook 可以改写输入，再交给原始 dispatch 实现。
            hidden_states, topk_output = self._pre_dispatch_hooks(
                self, hidden_states, topk_output
            )
        dispatch_output = self._original_dispatch_func(
            hidden_states=hidden_states, topk_output=topk_output
        )
        if self._post_dispatch_hooks is not None:
            # post hook 可以基于完整 dispatch_output 继续注入 SBO 参数。
            dispatch_output = self._post_dispatch_hooks(self, dispatch_output)
        return dispatch_output
 
    def _combine_with_hook(self, combine_input: CombineInput) -> torch.Tensor:
        if self._pre_combine_hooks is not None:
            # combine 前 hook 常用于记录 event 或排 shared experts。
            combine_input = self._pre_combine_hooks(self, combine_input)
        hidden_states = self._original_combine_func(combine_input=combine_input)
        if self._post_combine_hooks is not None:
            # combine 后 hook 主要做清理，防止状态泄漏到下一层。
            hidden_states = self._post_combine_hooks(self, hidden_states)
        return hidden_states
 
    def set_overlap_args(
        self, combine_overlap_args: CombineOverlapArgs, meta_overlap_args: dict
    ) -> None:
        # 将 combine 侧 overlap 参数和共享元数据挂到 dispatcher 上。
        self.overlap_args = combine_overlap_args
        self.meta_overlap_args = meta_overlap_args
 
    def clear_overlap_args(self) -> None:
        # 每层 / 每批结束后清理，避免串到下一次调用。
        self.overlap_args = None
        self.meta_overlap_args = None

def forward_impl(self, hidden_states: torch.Tensor, topk_output: TopKOutput):
    origin_hidden_states_dim = hidden_states.shape[-1]
    assert self.quant_method is not None
 
    # 先 dispatch，把 token 重新排布到专家计算所需的布局。
    dispatch_output = self.dispatcher.dispatch(
        hidden_states=hidden_states, topk_output=topk_output
    )
 
    # run_moe_core 负责真正的 routed experts 计算。
    combine_input = self.run_moe_core(
        dispatch_output=dispatch_output,
    )
 
    with use_symmetric_memory(
        get_tp_group(), disabled=not is_allocation_symmetric()
    ):
        # combine 再把专家输出收回原 token 布局。
        final_hidden_states = self.dispatcher.combine(combine_input=combine_input)
        final_hidden_states = final_hidden_states[
            ..., :origin_hidden_states_dim
        ].contiguous()
 
    if self.reduce_results and (self.moe_tp_size > 1 or self.moe_ep_size > 1):
        # TP / EP 场景下还要做一次 all-reduce。
        final_hidden_states = tensor_model_parallel_all_reduce(final_hidden_states)
 
    return final_hidden_states

def forward_impl(self, hidden_states: torch.Tensor, topk_output: TopKOutput):
    origin_hidden_states_dim = hidden_states.shape[-1]
    assert self.quant_method is not None
 
    # 先 dispatch，把 token 重新排布到专家计算所需的布局。
    dispatch_output = self.dispatcher.dispatch(
        hidden_states=hidden_states, topk_output=topk_output
    )
 
    # run_moe_core 负责真正的 routed experts 计算。
    combine_input = self.run_moe_core(
        dispatch_output=dispatch_output,
    )
 
    with use_symmetric_memory(
        get_tp_group(), disabled=not is_allocation_symmetric()
    ):
        # combine 再把专家输出收回原 token 布局。
        final_hidden_states = self.dispatcher.combine(combine_input=combine_input)
        final_hidden_states = final_hidden_states[
            ..., :origin_hidden_states_dim
        ].contiguous()
 
    if self.reduce_results and (self.moe_tp_size > 1 or self.moe_ep_size > 1):
        # TP / EP 场景下还要做一次 all-reduce。
        final_hidden_states = tensor_model_parallel_all_reduce(final_hidden_states)
 
    return final_hidden_states

def forward_deepep(
    self,
    hidden_states: torch.Tensor,
    forward_batch: ForwardBatch,
) -> torch.Tensor:
    shared_output = None
    # SBO 只在 shared experts 允许融入 SBO，且当前不是 nextn 路径时开启。
    sbo_enabled_flag = self._fuse_shared_experts_inside_sbo and not self.is_nextn
    # 非 Blackwell：shared experts 插到 dispatch 中间。
    sbo_overlap_dispatch_flag = (
        sbo_enabled_flag and SboFlags.enable_dispatch_shared_one_stream_overlap()
    )
    # Blackwell 默认更偏向把 shared experts 放到 combine 前后做双流重叠。
    sbo_overlap_combine_flag = (
        sbo_enabled_flag and SboFlags.enable_combine_shared_two_stream_overlap()
    )
 
    if hidden_states.shape[0] > 0:
        # router_logits 的 shape 是 (num_tokens, n_experts)。
        router_logits = self.gate(hidden_states, forward_batch=forward_batch)
        if not sbo_enabled_flag:
            if self.alt_stream is not None:
                # 非 SBO 分支下，可直接把 shared experts 提前丢到 alt_stream。
                self.alt_stream.wait_stream(torch.cuda.current_stream())
                with torch.cuda.stream(self.alt_stream):
                    shared_output = self._forward_shared_experts(hidden_states)
                    # 记录这份输出的归属 stream，并留下事件给主 stream 汇合。
                    shared_output.record_stream(self.alt_stream)
                    shared_event = self.alt_stream.record_event()
            else:
                shared_output = self._forward_shared_experts(hidden_states)
        # 计算 routed experts 的 top-k 路由结果。
        topk_output = self.topk(
            hidden_states,
            router_logits,
            num_token_non_padded=forward_batch.num_token_non_padded,
            expert_location_dispatch_info=ExpertLocationDispatchInfo.init_new(
                layer_id=self.layer_id,
            ),
        )
    else:
        # 空 batch 也返回一个空的 topk 结果，保持后续接口一致。
        topk_output = self.topk.empty_topk_output(hidden_states.device)
 
    if sbo_overlap_dispatch_flag:
        # 非 Blackwell：shared experts 插在 dispatch 中间，combine 再去和 down GEMM 重叠。
        shared_output = None
 
        def _deepep_dispatch_hook(dispatcher: BaseDispatcher):
            nonlocal shared_output
            # 这个 hook 发生在 dispatch_a() 与 dispatch_b() 之间。
            shared_output = self._forward_shared_experts(hidden_states)
            # 只对当前这一次 forward 生效，执行后立刻移除。
            for handle in deepep_dispatch_hook_handle:
                handle.remove()
 
        def _post_dispatch_hook(
            dispatcher: BaseDispatcher, dispatch_output: DispatchOutput
        ):
            # 基于 packed dispatch 输出生成 overlap 参数，并同时注入 dispatcher 与 experts。
            combine_overlap_args, down_gemm_overlap_args, meta_overlap_args = (
                compute_overlap_args(dispatch_output, self.alt_stream)
            )
            dispatcher.set_overlap_args(
                combine_overlap_args=combine_overlap_args,
                meta_overlap_args=meta_overlap_args,
            )
            self.experts.set_overlap_args(
                down_gemm_overlap_args=down_gemm_overlap_args,
                meta_overlap_args=meta_overlap_args,
            )
            # 只让这组参数作用于当前这一层当前这一拍。
            post_dispatch_hook_handle.remove()
 
        def _post_combine_hook(
            dispatcher: BaseDispatcher, hidden_states: torch.Tensor
        ):
            # combine 完成后立刻清理状态，避免泄漏到下一层 / 下一批。
            dispatcher.clear_overlap_args()
            self.experts.clear_overlap_args()
            post_combine_hook_handle.remove()
 
        # TBO 开启时这里拿到的是外层包装 dispatcher，但 overlap 参数会向内广播。
        assert isinstance(self.experts.dispatcher, MaybeTboDeepEPDispatcher)
        deepep_dispatch_hook_handle = (
            self.experts.dispatcher.register_deepep_dispatch_hook(
                _deepep_dispatch_hook
            )
        )
        post_dispatch_hook_handle = (
            self.experts.dispatcher.register_post_dispatch_hook(_post_dispatch_hook)
        )
        post_combine_hook_handle = (
            self.experts.dispatcher.register_post_combine_hook(_post_combine_hook)
        )
 
    elif sbo_overlap_combine_flag:
        # Blackwell 默认：shared experts 放到 pre_combine 阶段，与 combine 双流重叠。
        shared_output = None
 
        def _post_dispatch_hook(
            dispatcher: BaseDispatcher, dispatch_output: DispatchOutput
        ):
            # 先把 combine / down GEMM 所需的 overlap 参数注入进去。
            combine_overlap_args, down_gemm_overlap_args, meta_overlap_args = (
                compute_overlap_args(dispatch_output, self.alt_stream)
            )
            dispatcher.set_overlap_args(
                combine_overlap_args=combine_overlap_args,
                meta_overlap_args=meta_overlap_args,
            )
            self.experts.set_overlap_args(
                down_gemm_overlap_args=down_gemm_overlap_args,
                meta_overlap_args=meta_overlap_args,
            )
 
            post_dispatch_hook_handle.remove()
 
        def _pre_combine_hook(
            dispatcher: BaseDispatcher, combine_input: CombineInput
        ):
            nonlocal shared_output
 
            if (
                e := dispatcher.meta_overlap_args.get("record_event_after_down")
            ) is not None:
                # combine-shared 模式下，用它标记“当前 stream 已经推进到 down 之后”。
                e.record()
 
            # TODO 对非 deep_gemm 后端，也许还能进一步收缩 shared experts 可用 SM。
            with deep_gemm_wrapper.configure_deep_gemm_num_sms(
                # shared experts 也遵守 compute_num_sms 的 SM 预算。
                dispatcher.meta_overlap_args["compute_num_sms"]
            ):
                shared_output = self._forward_shared_experts(hidden_states)
 
            # 这个 hook 也只在当前 forward 中执行一次。
            pre_combine_hook_handle.remove()
 
        def _post_combine_hook(
            dispatcher: BaseDispatcher, hidden_states: torch.Tensor
        ):
            # combine 完成后清理 overlap 状态。
            dispatcher.clear_overlap_args()
            self.experts.clear_overlap_args()
            post_combine_hook_handle.remove()
 
        post_dispatch_hook_handle = (
            self.experts.dispatcher.register_post_dispatch_hook(_post_dispatch_hook)
        )
        pre_combine_hook_handle = self.experts.dispatcher.register_pre_combine_hook(
            _pre_combine_hook
        )
        post_combine_hook_handle = (
            self.experts.dispatcher.register_post_combine_hook(_post_combine_hook)
        )
    elif envs.SGLANG_BLACKWELL_OVERLAP_SHARED_EXPERTS_OUTSIDE_SBO.get():
        # 在 GB200 上：shared experts 走 alt_stream，down GEMM 与 DeepEP Combine 尝试重叠。
 
        def _post_dispatch_hook(
            dispatcher: BaseDispatcher, dispatch_output: DispatchOutput
        ):
            # 即便 shared experts 不走 SBO 内部融合，combine/down 的 overlap 参数仍在这里注入。
            combine_overlap_args, down_gemm_overlap_args, meta_overlap_args = (
                compute_overlap_args(dispatch_output, self.alt_stream)
            )
            dispatcher.set_overlap_args(
                combine_overlap_args=combine_overlap_args,
                meta_overlap_args=meta_overlap_args,
            )
            self.experts.set_overlap_args(
                down_gemm_overlap_args=down_gemm_overlap_args,
                meta_overlap_args=meta_overlap_args,
            )
 
            post_dispatch_hook_handle.remove()
 
        def _pre_combine_hook(
            dispatcher: BaseDispatcher, combine_input: CombineInput
        ):
            if (
                e := dispatcher.meta_overlap_args.get("record_event_after_down")
            ) is not None:
                # 若存在“down 之后再启动 combine”的事件，就在这里记录。
                e.record()
            pre_combine_hook_handle.remove()
 
        def _post_combine_hook(
            dispatcher: BaseDispatcher, hidden_states: torch.Tensor
        ):
            # 与其他分支一样，combine 结束后必须清理状态。
            dispatcher.clear_overlap_args()
            self.experts.clear_overlap_args()
            post_combine_hook_handle.remove()
 
        post_dispatch_hook_handle = (
            self.experts.dispatcher.register_post_dispatch_hook(_post_dispatch_hook)
        )
        pre_combine_hook_handle = self.experts.dispatcher.register_pre_combine_hook(
            _pre_combine_hook
        )
        post_combine_hook_handle = (
            self.experts.dispatcher.register_post_combine_hook(_post_combine_hook)
        )
 
    # 真正触发 routed experts 主路径；前面注册的 hook 都会在内部生效。
    final_hidden_states = self.experts(
        hidden_states=hidden_states,
        topk_output=topk_output,
    )
 
    if (
        hidden_states.shape[0] > 0
        and not sbo_enabled_flag
        and self.alt_stream is not None
    ):
        # 非 SBO 但 shared experts 走了 alt_stream 时，主 stream 需要在合并前等待。
        torch.cuda.current_stream().wait_event(shared_event)
 
    if shared_output is not None:
        # 有 shared_output 时，把 shared experts 输出和 routed experts 输出合并。
        x = shared_output
        if self.experts.should_fuse_routed_scaling_factor_in_topk or _use_aiter:
            # 某些后端已在 topk / runner 内部融合了缩放，这里直接相加。
            x.add_(final_hidden_states)
        else:
            # 否则按 routed_scaling_factor 做加权相加。
            x.add_(final_hidden_states, alpha=self.routed_scaling_factor)
        final_hidden_states = x
    else:
        if not (
            self.experts.should_fuse_routed_scaling_factor_in_topk or _use_aiter
        ):
            # 没有 shared_output 时，仅对 routed experts 输出补上缩放。
            final_hidden_states *= self.routed_scaling_factor
 
    return final_hidden_states

def forward_deepep(
    self,
    hidden_states: torch.Tensor,
    forward_batch: ForwardBatch,
) -> torch.Tensor:
    shared_output = None
    # SBO 只在 shared experts 允许融入 SBO，且当前不是 nextn 路径时开启。
    sbo_enabled_flag = self._fuse_shared_experts_inside_sbo and not self.is_nextn
    # 非 Blackwell：shared experts 插到 dispatch 中间。
    sbo_overlap_dispatch_flag = (
        sbo_enabled_flag and SboFlags.enable_dispatch_shared_one_stream_overlap()
    )
    # Blackwell 默认更偏向把 shared experts 放到 combine 前后做双流重叠。
    sbo_overlap_combine_flag = (
        sbo_enabled_flag and SboFlags.enable_combine_shared_two_stream_overlap()
    )
 
    if hidden_states.shape[0] > 0:
        # router_logits 的 shape 是 (num_tokens, n_experts)。
        router_logits = self.gate(hidden_states, forward_batch=forward_batch)
        if not sbo_enabled_flag:
            if self.alt_stream is not None:
                # 非 SBO 分支下，可直接把 shared experts 提前丢到 alt_stream。
                self.alt_stream.wait_stream(torch.cuda.current_stream())
                with torch.cuda.stream(self.alt_stream):
                    shared_output = self._forward_shared_experts(hidden_states)
                    # 记录这份输出的归属 stream，并留下事件给主 stream 汇合。
                    shared_output.record_stream(self.alt_stream)
                    shared_event = self.alt_stream.record_event()
            else:
                shared_output = self._forward_shared_experts(hidden_states)
        # 计算 routed experts 的 top-k 路由结果。
        topk_output = self.topk(
            hidden_states,
            router_logits,
            num_token_non_padded=forward_batch.num_token_non_padded,
            expert_location_dispatch_info=ExpertLocationDispatchInfo.init_new(
                layer_id=self.layer_id,
            ),
        )
    else:
        # 空 batch 也返回一个空的 topk 结果，保持后续接口一致。
        topk_output = self.topk.empty_topk_output(hidden_states.device)
 
    if sbo_overlap_dispatch_flag:
        # 非 Blackwell：shared experts 插在 dispatch 中间，combine 再去和 down GEMM 重叠。
        shared_output = None
 
        def _deepep_dispatch_hook(dispatcher: BaseDispatcher):
            nonlocal shared_output
            # 这个 hook 发生在 dispatch_a() 与 dispatch_b() 之间。
            shared_output = self._forward_shared_experts(hidden_states)
            # 只对当前这一次 forward 生效，执行后立刻移除。
            for handle in deepep_dispatch_hook_handle:
                handle.remove()
 
        def _post_dispatch_hook(
            dispatcher: BaseDispatcher, dispatch_output: DispatchOutput
        ):
            # 基于 packed dispatch 输出生成 overlap 参数，并同时注入 dispatcher 与 experts。
            combine_overlap_args, down_gemm_overlap_args, meta_overlap_args = (
                compute_overlap_args(dispatch_output, self.alt_stream)
            )
            dispatcher.set_overlap_args(
                combine_overlap_args=combine_overlap_args,
                meta_overlap_args=meta_overlap_args,
            )
            self.experts.set_overlap_args(
                down_gemm_overlap_args=down_gemm_overlap_args,
                meta_overlap_args=meta_overlap_args,
            )
            # 只让这组参数作用于当前这一层当前这一拍。
            post_dispatch_hook_handle.remove()
 
        def _post_combine_hook(
            dispatcher: BaseDispatcher, hidden_states: torch.Tensor
        ):
            # combine 完成后立刻清理状态，避免泄漏到下一层 / 下一批。
            dispatcher.clear_overlap_args()
            self.experts.clear_overlap_args()
            post_combine_hook_handle.remove()
 
        # TBO 开启时这里拿到的是外层包装 dispatcher，但 overlap 参数会向内广播。
        assert isinstance(self.experts.dispatcher, MaybeTboDeepEPDispatcher)
        deepep_dispatch_hook_handle = (
            self.experts.dispatcher.register_deepep_dispatch_hook(
                _deepep_dispatch_hook
            )
        )
        post_dispatch_hook_handle = (
            self.experts.dispatcher.register_post_dispatch_hook(_post_dispatch_hook)
        )
        post_combine_hook_handle = (
            self.experts.dispatcher.register_post_combine_hook(_post_combine_hook)
        )
 
    elif sbo_overlap_combine_flag:
        # Blackwell 默认：shared experts 放到 pre_combine 阶段，与 combine 双流重叠。
        shared_output = None
 
        def _post_dispatch_hook(
            dispatcher: BaseDispatcher, dispatch_output: DispatchOutput
        ):
            # 先把 combine / down GEMM 所需的 overlap 参数注入进去。
            combine_overlap_args, down_gemm_overlap_args, meta_overlap_args = (
                compute_overlap_args(dispatch_output, self.alt_stream)
            )
            dispatcher.set_overlap_args(
                combine_overlap_args=combine_overlap_args,
                meta_overlap_args=meta_overlap_args,
            )
            self.experts.set_overlap_args(
                down_gemm_overlap_args=down_gemm_overlap_args,
                meta_overlap_args=meta_overlap_args,
            )
 
            post_dispatch_hook_handle.remove()
 
        def _pre_combine_hook(
            dispatcher: BaseDispatcher, combine_input: CombineInput
        ):
            nonlocal shared_output
 
            if (
                e := dispatcher.meta_overlap_args.get("record_event_after_down")
            ) is not None:
                # combine-shared 模式下，用它标记“当前 stream 已经推进到 down 之后”。
                e.record()
 
            # TODO 对非 deep_gemm 后端，也许还能进一步收缩 shared experts 可用 SM。
            with deep_gemm_wrapper.configure_deep_gemm_num_sms(
                # shared experts 也遵守 compute_num_sms 的 SM 预算。
                dispatcher.meta_overlap_args["compute_num_sms"]
            ):
                shared_output = self._forward_shared_experts(hidden_states)
 
            # 这个 hook 也只在当前 forward 中执行一次。
            pre_combine_hook_handle.remove()
 
        def _post_combine_hook(
            dispatcher: BaseDispatcher, hidden_states: torch.Tensor
        ):
            # combine 完成后清理 overlap 状态。
            dispatcher.clear_overlap_args()
            self.experts.clear_overlap_args()
            post_combine_hook_handle.remove()
 
        post_dispatch_hook_handle = (
            self.experts.dispatcher.register_post_dispatch_hook(_post_dispatch_hook)
        )
        pre_combine_hook_handle = self.experts.dispatcher.register_pre_combine_hook(
            _pre_combine_hook
        )
        post_combine_hook_handle = (
            self.experts.dispatcher.register_post_combine_hook(_post_combine_hook)
        )
    elif envs.SGLANG_BLACKWELL_OVERLAP_SHARED_EXPERTS_OUTSIDE_SBO.get():
        # 在 GB200 上：shared experts 走 alt_stream，down GEMM 与 DeepEP Combine 尝试重叠。
 
        def _post_dispatch_hook(
            dispatcher: BaseDispatcher, dispatch_output: DispatchOutput
        ):
            # 即便 shared experts 不走 SBO 内部融合，combine/down 的 overlap 参数仍在这里注入。
            combine_overlap_args, down_gemm_overlap_args, meta_overlap_args = (
                compute_overlap_args(dispatch_output, self.alt_stream)
            )
            dispatcher.set_overlap_args(
                combine_overlap_args=combine_overlap_args,
                meta_overlap_args=meta_overlap_args,
            )
            self.experts.set_overlap_args(
                down_gemm_overlap_args=down_gemm_overlap_args,
                meta_overlap_args=meta_overlap_args,
            )
 
            post_dispatch_hook_handle.remove()
 
        def _pre_combine_hook(
            dispatcher: BaseDispatcher, combine_input: CombineInput
        ):
            if (
                e := dispatcher.meta_overlap_args.get("record_event_after_down")
            ) is not None:
                # 若存在“down 之后再启动 combine”的事件，就在这里记录。
                e.record()
            pre_combine_hook_handle.remove()
 
        def _post_combine_hook(
            dispatcher: BaseDispatcher, hidden_states: torch.Tensor
        ):
            # 与其他分支一样，combine 结束后必须清理状态。
            dispatcher.clear_overlap_args()
            self.experts.clear_overlap_args()
            post_combine_hook_handle.remove()
 
        post_dispatch_hook_handle = (
            self.experts.dispatcher.register_post_dispatch_hook(_post_dispatch_hook)
        )
        pre_combine_hook_handle = self.experts.dispatcher.register_pre_combine_hook(
            _pre_combine_hook
        )
        post_combine_hook_handle = (
            self.experts.dispatcher.register_post_combine_hook(_post_combine_hook)
        )
 
    # 真正触发 routed experts 主路径；前面注册的 hook 都会在内部生效。
    final_hidden_states = self.experts(
        hidden_states=hidden_states,
        topk_output=topk_output,
    )
 
    if (
        hidden_states.shape[0] > 0
        and not sbo_enabled_flag
        and self.alt_stream is not None
    ):
        # 非 SBO 但 shared experts 走了 alt_stream 时，主 stream 需要在合并前等待。
        torch.cuda.current_stream().wait_event(shared_event)
 
    if shared_output is not None:
        # 有 shared_output 时，把 shared experts 输出和 routed experts 输出合并。
        x = shared_output
        if self.experts.should_fuse_routed_scaling_factor_in_topk or _use_aiter:
            # 某些后端已在 topk / runner 内部融合了缩放，这里直接相加。
            x.add_(final_hidden_states)
        else:
            # 否则按 routed_scaling_factor 做加权相加。
            x.add_(final_hidden_states, alpha=self.routed_scaling_factor)
        final_hidden_states = x
    else:
        if not (
            self.experts.should_fuse_routed_scaling_factor_in_topk or _use_aiter
        ):
            # 没有 shared_output 时，仅对 routed experts 输出补上缩放。
            final_hidden_states *= self.routed_scaling_factor
 
    return final_hidden_states