执行摘要

• 一句话：回退 workspace buffer 初始化重构，修复残缺 wrappers 问题
• 推荐动作：建议立即跟进修复 init_mha_chunk_metadata 中的 AttributeError，在调用父类前增加 hasattr(self, 'mha_chunk_kv_cache') 检查或条件保护。长期而言，可重新设计 workspace 初始化方案，兼顾子类复用与父类完整性，但需确保所有代码路径下 wrapper 初始化完备。

功能与动机

PR #25321 引入的优化导致父类在 skip_init_workspace_buffer=True 时跳过所有 attention wrappers（如 prefill_wrapper_ragged）的创建，而子类并未补全这些功能，使得 TRTLLMMLABackend 在 speculative decoding 等场景下功能残缺。回退是最直接的修复方式。

实现拆解

实施以下几步反转：

1. flashinfer_mla_backend.py：移除构造函数中的 skip_init_workspace_buffer 形参和所有条件分支，始终调用全局 workspace buffer 的分配逻辑，并无条件初始化 prefill_wrapper_ragged、decode_wrapper 及 indices_updater_decode 等关键属性。
2. trtllm_mla_backend.py：移除构造函数中的 skip_init_workspace_buffer 形参，不再向父类传递 skip_init_workspace_buffer=True，恢复自身对 global_zero_init_workspace_buffer 的初始化代码（与回退前一致）。同时移除原先的空 init_mha_chunk_metadata 覆盖，改为调用父类的对应方法（传入 disable_flashinfer_ragged=True）。
3. tokenspeed_mla_backend.py：移除 super().__init__ 调用中的 skip_init_workspace_buffer=True 实参。
4. 测试与配置：本次回退未引入新的测试或配置项。

关键文件：

• python/sglang/srt/layers/attention/flashinfer_mla_backend.py（模块 注意力后端；类别 source；类型 core-logic）: 核心父类，移除了 skip_init_workspace_buffer 参数，简化初始化逻辑，统一 wrapper 创建路径
• python/sglang/srt/layers/attention/trtllm_mla_backend.py（模块 注意力后端；类别 source；类型 core-logic；符号 init_mha_chunk_metadata）: 恢复父类初始化调用，移除子类参数传递，重新添加 init_mha_chunk_metadata 方法
• python/sglang/srt/layers/attention/tokenspeed_mla_backend.py（模块 注意力后端；类别 source；类型 core-logic）: 仅移除 super() 调用中的 skip_init_workspace_buffer=True

关键符号：FlashInferMLAAttnBackend.init, TRTLLMMLABackend.init, TokenspeedMLABackend.init, TRTLLMMLABackend.init_mha_chunk_metadata

关键源码片段

python/sglang/srt/layers/attention/flashinfer_mla_backend.py

核心父类，移除了 skip_init_workspace_buffer 参数，简化初始化逻辑，统一 wrapper 创建路径

# FlashInferMLAAttnBackend.__init__（revert 后版本）
# 移除了 skip_init_workspace_buffer 参数，始终创建 wrappers
class FlashInferMLAAttnBackend(AttentionBackend):
    def __init__(
        self,
        model_runner: ModelRunner,
        skip_prefill: bool = False,
        kv_indptr_buf: Optional[torch.Tensor] = None,
        q_indptr_decode_buf: Optional[torch.Tensor] = None,
    ):
        super().__init__()
        self.max_context_len = model_runner.model_config.context_len
        self.device = model_runner.device
        self.skip_prefill = skip_prefill
        self.enable_chunk_kv = (
            not skip_prefill
            and get_global_server_args().disaggregation_mode != "decode"
            and not get_global_server_args().disable_chunked_prefix_cache
            and not get_global_server_args().flashinfer_mla_disable_ragged
        )
        self.page_size = model_runner.page_size
​
        # Allocate buffers — 不再有 skip_init_workspace_buffer 分支
        global global_workspace_buffer
        if global_workspace_buffer is None:
            global_workspace_buffer = torch.empty(
                envs.SGLANG_FLASHINFER_WORKSPACE_SIZE.get(),
                dtype=torch.uint8,
                device=model_runner.device,
            )
        self.workspace_buffer = global_workspace_buffer
​
        # 始终创建 attention wrappers（之前可能跳过）
        if is_sm100_supported():
            self.fmha_backend = "cutlass"
        else:
            self.fmha_backend = "auto"
​
        self.prefill_wrapper_ragged = BatchPrefillWithRaggedKVCacheWrapper(
            self.workspace_buffer, "NHD", backend=self.fmha_backend
        )
        if not self.skip_prefill:
            self.prefill_wrapper_paged = BatchMLAPagedAttentionWrapper(
                self.workspace_buffer, backend="auto"
            )
            self.prefill_wrapper_verify = BatchMLAPagedAttentionWrapper(
                self.workspace_buffer, backend="auto"
            )
        self.decode_wrapper = BatchMLAPagedAttentionWrapper(
            self.workspace_buffer, backend="auto"
        )
        # indices updater 无条件创建
        self.indices_updater_decode = FlashInferMLAIndicesUpdaterDecode(
            model_runner, self
        )

python/sglang/srt/layers/attention/trtllm_mla_backend.py

恢复父类初始化调用，移除子类参数传递，重新添加 init_mha_chunk_metadata 方法

# TRTLLMMLABackend.__init__ 与 init_mha_chunk_metadata（revert 后版本）
class TRTLLMMLABackend(FlashInferMLAAttnBackend):
    def __init__(
        self,
        model_runner: ModelRunner,
        skip_prefill: bool = False,
        kv_indptr_buf: Optional[torch.Tensor] = None,
        q_indptr_decode_buf: Optional[torch.Tensor] = None,
    ):
        # 不再传递 skip_init_workspace_buffer=True，父类会正常初始化 wrappers
        super().__init__(
            model_runner,
            skip_prefill,
            kv_indptr_buf,
            q_indptr_decode_buf,
        )
​
        # Workspace allocation — 子类额外维护一份全局 buffer
        self.workspace_size = DEFAULT_WORKSPACE_SIZE_MB * 1024 * 1024
        global global_zero_init_workspace_buffer
        if global_zero_init_workspace_buffer is None:
            global_zero_init_workspace_buffer = torch.zeros(
                self.workspace_size,
                dtype=torch.uint8,
                device=model_runner.device,
            )
        self.workspace_buffer = global_zero_init_workspace_buffer
        # ... 其余初始化不变
​
    def init_mha_chunk_metadata(self, forward_batch: ForwardBatch):
        # 调用父类方法，但未检查 self.mha_chunk_kv_cache 是否存在
        # 当 skip_prefill=True 或 enable_chunk_kv=False 时可能引发 AttributeError
        super().init_mha_chunk_metadata(forward_batch, disable_flashinfer_ragged=True)

评论区精华

review 中 gemini-code-assist[bot] 指出新增的 init_mha_chunk_metadata 方法存在潜在 AttributeError：当 skip_prefill=True（speculative decoding draft worker）或 enable_chunk_kv=False 时，父类不会创建 self.mha_chunk_kv_cache，而调用 super().init_mha_chunk_metadata 会试图访问该属性，导致崩溃。该评论为高优先级，但 PR 仍被合并，问题未解决。

• init_mha_chunk_metadata 潜在 AttributeError (correctness): PR 已合并，问题未解决，需后续修复。

风险与影响

• 风险：
  1. 新引入的崩溃风险：TRTLLMMLABackend.init_mha_chunk_metadata 无条件调用父类方法，若 mha_chunk_kv_cache 未初始化（如 draft 模型或禁用 chunked prefix cache），将抛出 AttributeError。
  2. 内存浪费：父类和子类各自维护一份 workspace buffer（global_workspace_buffer 与 global_zero_init_workspace_buffer），造成少量显存冗余，但功能上无影响。
  3. 无回归：该 revert 恢复了 PR #25321 之前的稳定行为，之前的功能缺陷得以修复。
     - 影响：影响范围：所有使用 MLA attention backend 的模型，尤其依赖 TRTLLMMLABackend 和 TokenspeedMLABackend 的 Blackwell SM100 推理场景。程度：回退后之前因缺少 wrappers 导致的解码失败问题消失，但新引入的 AttributeError 仅在特定配置下触发，属于潜在风险。
     - 风险标记：潜在 AttributeError, 缺少条件保护

关联脉络

• PR #25321 [attn backend] avoid initing parent class's workspace buffer: 本 PR 回退该 PR 引入的变更，恢复原有初始化行为。