执行摘要

• 一句话：提取 NgramEmbeddingManager 组件，集中管理 ngram embedding 状态
• 推荐动作：该 PR 是典型的组件提取重构，值得阅读 NgramEmbeddingManager 的实现，特别是 from_model 和 prepare_for_forward 的逻辑。注意其如何通过 dataclass 集中管理状态，并通过注入依赖减少与 ModelRunner 的耦合。对于从事类似大型类拆分的开发者有借鉴意义。

功能与动机

将分散的 ngram embedding 管理逻辑集中到一个组件中，降低 ModelRunner 的职责复杂度，为后续进一步分离和测试做准备。PR body 明确标识为 ngram embedding migration (PR 1/3)。

实现拆解

1. 创建 NgramEmbeddingManager 组件 (python/sglang/srt/model_executor/model_runner_components/ngram_embedding_manager.py)：一个 frozen dataclass，包含 enabled、table、n、k 字段。提供 from_model 类方法执行初始化逻辑（创建 token table、验证 chunked prefill、初始化 NgramEmbedding 模块的 buffers）。提供 update_after_decode 和 prepare_for_forward 实例方法（原对应 ModelRunner 的 maybe_update_ngram_token_table 和 Scheduler 的 _maybe_prepare_ngram_embedding）。
2. 在 ModelRunner 中集成 (python/sglang/srt/model_executor/model_runner.py)：添加 init_ngram_embedding_manager 方法替换旧的 maybe_init_ngram_embedding。移除旧的 maybe_init_ngram_embedding 和 maybe_update_ngram_token_table 方法，修改 alloc_memory_pool 中调用点。删除 ngram_token_table.py 的导入。
3. 在 Scheduler 中集成 (python/sglang/srt/managers/scheduler.py)：将 _maybe_prepare_ngram_embedding 方法替换为调用 model_runner.ngram_embedding_manager.prepare_for_forward。调整 maybe_init_ngram_embedding 以通过 manager 获取 table 和 enabled 状态。
4. 更新其他消费者：base_runner.py 和 decode_cuda_graph_runner.py 改用 ngram_embedding_manager 访问 token table 和 enabled 标志；forward_batch_info.py 调整了某个导入；mlx model_runner_stub.py 添加了 ngram_embedding_manager 的占位。
5. 调整测试：将 test_ngram_token_table.py 重命名为 test_ngram_embedding_manager.py 并更新导入；更新 test_scheduler_chunked_req_gate.py 和 test_gdn_prefill_backend_policy.py 以反映从 manager 获取字段的变更。

关键文件：

• python/sglang/srt/model_executor/model_runner_components/ngram_embedding_manager.py（模块 嵌入管理；类别 source；类型 data-contract；符号 NgramEmbeddingManager, from_model, update_after_decode, prepare_for_forward）: 新增的核心文件，定义了 NgramEmbeddingManager 类及其所有方法，是本次重构的中心。
• python/sglang/srt/model_executor/model_runner.py（模块 模型运行器；类别 source；类型 data-contract；符号 init_ngram_embedding_manager, maybe_init_ngram_embedding, maybe_update_ngram_token_table）: 集成 NgramEmbeddingManager，替代原有的 ngram embedding 方法，是主变更入口之一。
• python/sglang/srt/model_executor/ngram_token_table.py（模块 嵌入管理；类别 source；类型 deletion；符号 update_ngram_token_table_after_sampling）: 被删除的旧文件，其逻辑已完全迁移至 NgramEmbeddingManager 中。
• python/sglang/srt/managers/scheduler.py（模块 调度器；类别 source；类型 core-logic；符号 _maybe_prepare_ngram_embedding）: 调度器中 ngram embedding 相关逻辑改为调用 NgramEmbeddingManager，是核心逻辑变更之一。
• python/sglang/srt/model_executor/runner/base_runner.py（模块 基类运行器；类别 source；类型 data-contract）: 作为消费者更新为使用 ngram_embedding_manager 访问 token table 和 enabled。
• python/sglang/srt/model_executor/runner/decode_cuda_graph_runner.py（模块 解码图运行器；类别 source；类型 data-contract）: 作为消费者更新为使用 ngram_embedding_manager 访问 token table 和 enabled。
• test/registered/unit/model_executor/model_runner_components/test_ngram_embedding_manager.py（模块 单元测试；类别 test；类型 rename-or-move）: 测试文件从原 test_ngram_token_table.py 重命名并更新导入，反映组件变化。

关键符号：NgramEmbeddingManager.from_model, NgramEmbeddingManager.update_after_decode, NgramEmbeddingManager.prepare_for_forward, update_ngram_token_table_after_sampling, maybe_init_ngram_embedding, maybe_update_ngram_token_table, init_ngram_embedding_manager, _maybe_prepare_ngram_embedding

关键源码片段

python/sglang/srt/model_executor/model_runner_components/ngram_embedding_manager.py

新增的核心文件，定义了 NgramEmbeddingManager 类及其所有方法，是本次重构的中心。

"""Utilities for updating LongCat ngram embedding token tables."""
​
from __future__ import annotations
​
from dataclasses import dataclass
from typing import TYPE_CHECKING, Optional
​
import torch
​
from sglang.jit_kernel.ngram_embedding import update_token_table
from sglang.srt.configs.model_config import ModelConfig
from sglang.srt.managers.schedule_batch import ForwardMode
from sglang.srt.mem_cache.memory_pool import ReqToTokenPool
from sglang.srt.server_args import ServerArgs
​
if TYPE_CHECKING:
    from sglang.srt.managers.schedule_batch import Req, ScheduleBatch
    from sglang.srt.model_executor.forward_batch_info import ForwardBatch
​
​
@dataclass(frozen=True, slots=True, kw_only=True)
class NgramEmbeddingManager:
    """将 ngram embedding 的状态（enabled、table、n、k）以及相关操作封装为不可变组件。"""
​
    enabled: bool
    table: Optional[torch.Tensor] # 与 req_to_token 同大小的 token 表
    n: int
    k: int
​
    @classmethod
    def from_model(
        cls,
        *,
        model: torch.nn.Module,
        model_config: ModelConfig,
        req_to_token_pool: ReqToTokenPool,
        server_args: ServerArgs,
        max_running_requests: int,
        device: str,
    ) -> "NgramEmbeddingManager":
        """工厂方法，从模型和服务参数构造管理器。
​
        如果启用 ngram embedding，则分配 token 表、验证 chunked prefill、初始化
        模型中每个 NgramEmbedding 模块的预分配缓冲区。
        """
        token_table = None
        ngram_embedding_n = 0
        ngram_embedding_k = 0
        use_ngram_embedding = model_config.use_ngram_embedding
        if use_ngram_embedding:
            from sglang.srt.layers.n_gram_embedding import NgramEmbedding
​
            # 为 req_to_token 分配镜像大小的 token 表（按 req_pool_idx 索引）
            token_table = torch.empty(
                req_to_token_pool.req_to_token.shape[0],
                model_config.context_len,
                dtype=torch.int32,
                device=device,
            )
            chunked_prefill_size = server_args.chunked_prefill_size
            # ngram embedding 需要启用 chunked prefill
            assert (
                chunked_prefill_size is not None and chunked_prefill_size > 0
            ), "Ngram embedding requires chunked prefill to be enabled (chunked_prefill_size > 0)"
            # 遍历模型中的所有模块，初始化 NgramEmbedding 缓冲区
            for module in model.modules():
                if isinstance(module, NgramEmbedding):
                    module.init_buffers(max_running_requests, chunked_prefill_size, device)
            hf_config = model_config.hf_config
            ngram_embedding_n = hf_config.ngram_embedding_n
            ngram_embedding_k = hf_config.ngram_embedding_k
        return cls(
            enabled=use_ngram_embedding,
            table=token_table,
            n=ngram_embedding_n,
            k=ngram_embedding_k,
        )
​
    # update_after_decode() 和 prepare_for_forward() 两种实例方法分别封装了
    # 原 ModelRunner.maybe_update_ngram_token_table 和 Scheduler._maybe_prepare_ngram_embedding 的逻辑。
    # 它们通过访问 self.table、self.enabled、self.n 等字段，不再依赖 ModelRunner 状态。

评论区精华

无实质性 review 讨论。仅有一个自动代码审查机器人评论，未提供反馈。

• 代码审查机器人反馈 (other): 无需处理。

风险与影响

• 风险：核心路径变更：修改了 ModelRunner 和 Scheduler 的初始化与前向逻辑，可能影响 ngram embedding 功能。提供了完整的消费者迁移，但需确保所有使用 use_ngram_embedding 和 token_table 的地方已全部切换。删除 ngram_token_table.py 后，若外部依赖该模块则报错（但仓库内已无引用）。mlx 硬件后端的 model_runner_stub 仅添加了 manager 引用，未实现，若在 mlx 上使用 ngram embedding 可能异常（但 ngram embedding 本身可能不支持 mlx）。测试覆盖了基本场景，但未覆盖所有边缘情况（如 enable_ngram_embedding 关闭时的行为）。
• 影响：用户无直接影响，功能保持不变。内部重构使代码结构更清晰，降低维护成本。团队可以从该组件提取模式中获得模块化设计的参考。
• 风险标记：核心路径变更, 删除原文件可能遗漏引用, mlx 后端仅添加占位

关联脉络

• PR #31167 Extract attention-backend setup into a module: 类似的组件提取重构，将注意力后端设置从 ModelRunner 中分离。
• PR #31155 Extract load_model helpers into a load_model_utils module: 同样的 ModelRunner 职责拆分模式，将模型加载辅助函数抽取为独立模块。