执行摘要

• 一句话：修复 Mamba 混合模型中投机解码重放陈旧初始化元数据的问题。
• 推荐动作：此 PR 值得精读。虽然代码改动很小（仅 4 行核心修复），但它揭示了混合模型状态管理中的一个微妙边界问题。PR 的描述和测试计划非常清晰，展示了严谨的调试和验证方法。对于理解 SGLang 中 Mamba 状态、投机解码流程以及 ScheduleBatch 生命周期管理非常有帮助。

功能与动机

Deferred Mamba clear/COW 索引是为一次 EXTEND 前向传播准备的临时元数据。目标前向传播从 ForwardBatch 中消费它们，但原始的 ScheduleBatch 仍保留着这些索引。对于使用内置多层 EAGLE/MTP 的混合 Mamba 模型，后续的 DRAFT_EXTEND_V2 批次会复制这些陈旧的索引，并将其对应的已完成操作重放到活动的草稿池状态上，导致错误和性能下降（如 PR body 中所示，草稿接受率从 47.65% 降至 57.38%）。

实现拆解

1. 核心逻辑修复 (schedule_batch.py): 在 ScheduleBatch.prepare_for_decode() 方法中，在设置 forward_mode 并清空 input_embeds 之后，立即添加三行代码，将 self.mamba_cow_src_indices、self.mamba_cow_dst_indices 和 self.mamba_clear_indices 这三个属性显式设置为 None。此操作在判断是否为投机解码（spec_algorithm.is_none()）的分支之前执行，确保无论后续走投机解码路径还是普通解码路径，这些陈旧的一次性元数据都会被清除。
2. 单元测试覆盖 (test_schedule_batch_prepare_for_decode.py): 新增 TestPrepareForDecodeMambaInit 测试类及 test_spec_decode_drops_extend_mamba_init_metadata 方法。该测试模拟一个设置了 spec_algorithm（表示处于投机解码模式）和三个 mamba 初始化索引的 ScheduleBatch，调用 prepare_for_decode 后断言这三个索引被正确清空。测试通过模拟 spec_prepare_for_decode 来隔离测试目标逻辑。
3. 测试依赖更新: 测试文件导入了 CustomTestCase 以替换原有的 unittest.TestCase，用于与新增的测试类。

关键文件：

• python/sglang/srt/managers/schedule_batch.py（模块 调度批次管理；类别 source；类型 core-logic）: 核心修复点，在 prepare_for_decode 方法中清除陈旧的 Mamba 初始化元数据。
• test/registered/unit/managers/test_schedule_batch_prepare_for_decode.py（模块 调度批次测试；类别 test；类型 test-coverage；符号 TestPrepareForDecodeMambaInit, test_spec_decode_drops_extend_mamba_init_metadata）: 新增针对投机解码路径下 Mamba 初始化元数据清除的单元测试。

关键符号：ScheduleBatch.prepare_for_decode

关键源码片段

python/sglang/srt/managers/schedule_batch.py

核心修复点，在 prepare_for_decode 方法中清除陈旧的 Mamba 初始化元数据。

# python/sglang/srt/managers/schedule_batch.py
    def prepare_for_decode(self):
        self.forward_mode = ForwardMode.DECODE
        # Decode embeds the last output token via embed_tokens; clear the stale
        # prefill-time tensor so it doesn't leak into ForwardBatch.
        self.input_embeds = None
​
        # 修复核心：在投机解码路径判断之前，无条件地清除 EXTEND 阶段遗留的
        # Mamba 初始化元数据。这些是一次性数据，目标前向传播消费后应丢弃。
        # 若不清除，后续的 DRAFT_EXTEND_V2 批次会错误地重放这些已完成的操作。
        self.mamba_cow_src_indices = None
        self.mamba_cow_dst_indices = None
        self.mamba_clear_indices = None
​
        # Clear context parallel metadata - CP is only for prefill, not decode
        if hasattr(self, "attn_cp_metadata") and self.attn_cp_metadata is not None:
            self.attn_cp_metadata = None
​
        if not self.spec_algorithm.is_none():
            # Spec decoding owns decode preparation (allocation, seq-lens bookkeeping).
            from sglang.srt.speculative.spec_utils import spec_prepare_for_decode
​
            self.mamba_track_mask_cpu = None
            self.mamba_track_mask_next_cpu = None
            self.mamba_decode_batch_idx_cpu = None
            spec_prepare_for_decode(self)
            return

test/registered/unit/managers/test_schedule_batch_prepare_for_decode.py

新增针对投机解码路径下 Mamba 初始化元数据清除的单元测试。

# test/registered/unit/managers/test_schedule_batch_prepare_for_decode.py
class TestPrepareForDecodeMambaInit(CustomTestCase):
    def test_spec_decode_drops_extend_mamba_init_metadata(self):
        # 模拟一个处于投机解码模式（spec_algorithm 非 None）且设置了
        # EXTEND 阶段 Mamba 初始化索引的 ScheduleBatch。
        batch = ScheduleBatch(reqs=[])
        batch.spec_algorithm = types.SimpleNamespace(is_none=lambda: False)
        batch.mamba_cow_src_indices = torch.tensor([1])
        batch.mamba_cow_dst_indices = torch.tensor([2])
        batch.mamba_clear_indices = torch.tensor([3])
​
        # 模拟投机解码的 prepare_for_decode，以隔离测试目标逻辑。
        with patch(
            "sglang.srt.speculative.spec_utils.spec_prepare_for_decode"
        ) as prepare:
            batch.prepare_for_decode()
​
        # 验证投机解码的 prepare_for_decode 被调用，且三个元数据字段被清空。
        prepare.assert_called_once_with(batch)
        self.assertIsNone(batch.mamba_cow_src_indices)
        self.assertIsNone(batch.mamba_cow_dst_indices)
        self.assertIsNone(batch.mamba_clear_indices)

评论区精华

该 PR 的 review 过程非常简洁，没有技术讨论。审核人 Qiaolin-Yu 直接批准了更改。PR body 中提供了详尽的动机解释、问题复现方法（通过 Inkling 集成）以及量化的 A/B 测试结果，这本身构成了高质量的自审内容。

• 投机解码路径下 Mamba 初始化元数据的生命周期管理 (correctness): 修复方案是在 prepare_for_decode 入口点无条件清空这三个字段，确保无论后续路径如何，陈旧元数据都不会泄漏。

风险与影响

• 风险：风险较低。修改是局部且防御性的，在 prepare_for_decode 的入口点清除不再需要的元数据，不会影响正常的 EXTEND 流程。单元测试明确验证了投机解码路径下的清除行为。潜在风险在于是否还有其他地方依赖这些字段在 prepare_for_decode 后仍被保留，但根据 PR 描述，这些是一次性元数据，且清理位置合理。
• 影响：该修复主要影响使用混合 Mamba 模型并启用多层 EAGLE/MTP 投机解码（如 Inkling）的场景。它解决了因陈旧元数据重放导致的草稿接受率下降和无效的集合通信操作问题。对于不使用 Mamba 或投机解码的常规工作流，影响为零。修复提升了投机解码子系统的正确性和性能效率。
• 风险标记：投机解码状态管理

关联脉络

• PR #38418 [CI] test: initialize disable_radix_cache in scheduler stub: 同属调度器/批次管理（schedule_batch）相关的 bugfix 和测试修复，都涉及 scheduler 新增或修改字段的初始化问题。
• PR #38313 [Unified Tree] fix: make dfs weight ordering iterative: 虽然主题不同，但都涉及对核心调度/缓存模块（base_prefix_cache.py vs schedule_batch.py）的逻辑修复，体现了对内存和状态管理的持续优化。
• PR #38389 [Scheduler] Unify per-iteration request intake into ingest_requests(): 同一调度器重构线的一部分，目标是提升代码一致性和健壮性。本 PR 的修复点 prepare_for_decode 也是调度器请求处理的关键环节。