# PR #26112 完整报告

- 仓库：`sgl-project/sglang`
- 标题：[Kernel] Reuse WNA16 Marlin MoE workspace
- 合并时间：2026-05-26 13:19
- 原文链接：http://prhub.com.cn/sgl-project/sglang/pull/26112

---

# 执行摘要

- 一句话：WNA16 MoE Marlin 工作空间复用优化
- 推荐动作：该 PR 是一个小而有效的性能优化，值得合并。建议关注 workspace 大小参数是否应可配置，以及是否可推广到其他 MoE 后端。

# 功能与动机

Marlin MoE 内核在每次推理调用时都需要分配临时 workspace，导致不必要的内存分配开销，影响首 token 延迟和每个 token 的生成时间。通过复用持久化 workspace 来消除这部分开销。

# 实现拆解

1. **导入新增函数**：在 `compressed_tensors_wNa16_moe.py` 中从 `sglang.srt.layers.quantization.marlin_utils` 导入 `marlin_make_workspace`。
2. **创建持久 workspace**：在 `process_weights_after_loading` 方法末尾调用 `marlin_make_workspace`，将 workspace 存储为 `layer.workspace`，仅需创建一次。
3. **传递 workspace**：在 `apply_weights` 方法中，将 `layer.workspace` 作为参数传递给 `fused_marlin_moe` 内核，确保每次推理时复用已有 workspace。

关键文件：
- `python/sglang/srt/layers/quantization/compressed_tensors/schemes/compressed_tensors_wNa16_moe.py`（模块 量化；类别 source；类型 dependency-wiring；符号 marlin_make_workspace, process_weights_after_loading, apply_weights）: 核心变更文件：导入 marlin_make_workspace，创建持久 workspace，并在 forward 中传递。

关键符号：marlin_make_workspace, process_weights_after_loading, apply_weights

## 关键源码片段

### `python/sglang/srt/layers/quantization/compressed_tensors/schemes/compressed_tensors_wNa16_moe.py`

核心变更文件：导入 marlin_make_workspace，创建持久 workspace，并在 forward 中传递。

```python
# 导入新增 : 新增 marlin_make_workspace 用于创建持久化工作空间
from sglang.srt.layers.quantization.marlin_utils import (
    marlin_make_workspace,
    marlin_moe_permute_scales,
)

# 在 process_weights_after_loading 末尾创建 workspace
# 此处硬编码 size=4，需与 Marlin 内核要求的 workspace 大小一致
layer.workspace = marlin_make_workspace(layer.w13_weight_packed.device, 4)

# 在 apply_weights 中传递 workspace 给 fused_marlin_moe
fused_marlin_moe(
    x,
    ...,
    workspace=layer.workspace,
)

```

# 评论区精华

PR 无 review 讨论，仅有一条自动化评论警告配额限制。由于变更较小且经过基准测试，未产生争议。

- 暂无高价值评论线程

# 风险与影响

- 风险：风险较低。变更仅添加了 6 行代码，核心逻辑不变，workspace 管理由库函数 `marlin_make_workspace` 负责。可能的风险是不同配置下的 workspace 大小假设不一致，但当前 hardcode 为 4 应与内核需求匹配。无安全或兼容性问题。
- 影响：对用户：降低使用 WNA16 Marlin MoE 量化模型的推理延迟，尤其是 TTFT 改善明显（-9%）。对系统：减少显存分配次数，提升 GPU 利用率。对团队：仅影响一个源文件，后续需确保其他 MoE 后端（如 Triton）也能类似优化。影响范围为所有使用 `CompressedTensorsWNA16MoE` 量化方案的模型（如 Kimi K2.5）。
- 风险标记：hardcode workspace size

# 关联脉络

- PR #25910 WNA16 Marlin MoE workspace refactor: PR #25910 引入了 WNA16 Marlin MoE 的基础支持，本 PR 在其之上添加了 workspace 复用优化。