Prhub

#29228 [Spec] Merge dflash triton kernels into a single `dflash.py`

原始 PR 作者 hnyls2002 合并时间 2026-06-25 09:57 文件变更 3 提交数 1 评论 3 代码增减 +107 / -126

执行摘要

合并分散的 DFlash Triton 内核到单一 dflash.py

PR body 指出:dflash_accept_bonus.pydflash_prepare_block.py 各自持有一个 Triton 内核及其包装器,外加 _pick_num_warps_is_row_major_contiguous_2d 的字节相同副本。合并为单个 dflash.py(遵循 eagle.py / multi_layer_eagle.py 的按算法文件命名),消除重复辅助函数,提升可维护性。

值得精读以学习代码清理和消除重复的最佳实践。该 PR 展示了如何在不修改逻辑的前提下通过合并分散文件来提升可维护性。

讨论亮点

该 PR 无 Review 评论,合并没有公开讨论。

实现拆解

  1. 创建统一模块 python/sglang/srt/speculative/triton_ops/dflash.py:将原 dflash_accept_bonus.py 中的 _compute_dflash_accept_bonus_triton_unchecked 内核及其包装器,以及原 dflash_prepare_block.py 中的 _prepare_dflash_draft_block_contig_kernel_prepare_dflash_draft_block_unchecked 全部移入,同时只保留一份 _pick_num_warps_is_row_major_contiguous_2d 辅助函数。
  2. 删除旧文件dflash_accept_bonus.pydflash_prepare_block.py 被移除,不再存在两个分散文件。
  3. 更新导入关系:在 python/sglang/srt/speculative/dflash_worker_v2.py 中,将原来从两个单独模块的导入改为统一从 sglang.srt.speculative.triton_ops.dflash 导入 _compute_dflash_accept_bonus_triton_unchecked_prepare_dflash_draft_block_unchecked
    本次变更不包含测试、配置或部署配套改动,属于纯代码组织重构。
文件 模块 状态 重要度
python/sglang/srt/speculative/triton_ops/dflash.py Triton 操作 renamed 6.03
python/sglang/srt/speculative/triton_ops/dflash_prepare_block.py Triton 操作 removed 6.56
python/sglang/srt/speculative/dflash_worker_v2.py 推测解码 modified 4.27

关键符号

_compute_dflash_accept_bonus_triton_unchecked _prepare_dflash_draft_block_contig_kernel _prepare_dflash_draft_block_unchecked _pick_num_warps _is_row_major_contiguous_2d

关键源码片段

python/sglang/srt/speculative/triton_ops/dflash.py rename-or-move

核心输出文件:将原本分散在两个文件中的 Triton 内核和辅助函数统一合并到此文件,成为 DFlash 的单一 Triton 操作入口。

# python/sglang/srt/speculative/triton_ops/dflash.py
# 合并后的单一模块,包含两个内核及其统一辅助函数import torch
import triton
import triton.language as tl
​
​
# -------- 以下为原 dflash_accept_bonus.py 的内核 --------
@triton.jit
def _compute_dflash_accept_bonus_triton_kernel(
    # ... 参数列表与原文一致,省略以节省篇幅
):
    # 具体实现拷贝自原文件
    passdef _compute_dflash_accept_bonus_triton_unchecked(
    # ... 包装器逻辑
):
    # 从原文件直接移入
    pass# -------- 以下为原 dflash_prepare_block.py 的内核 --------
@triton.jit
def _prepare_dflash_draft_block_contig_kernel(
    # ... 参数列表
):
    # 从原文件直接移入
    passdef _prepare_dflash_draft_block_unchecked(
    # ... 包装器逻辑,包含校验
):
    # 从原文件直接移入
    pass# -------- 合并后的辅助函数(仅一份) --------
def _pick_num_warps(block_size: int) -> int:
    if block_size <= 16: return 1
    if block_size <= 32: return 2
    if block_size <= 64: return 4
    return 8def _is_row_major_contiguous_2d(x: torch.Tensor) -> bool:
    return x.ndim == 2 and x.is_contiguous()
python/sglang/srt/speculative/dflash_worker_v2.py dependency-wiring

唯一需要更新导入的文件:将两个独立导入合并为从新的 dflash.py 导入两个函数。

# python/sglang/srt/speculative/dflash_worker_v2.py
# 变更前(顶部集中导入部分):
# from sglang.srt.speculative.triton_ops.dflash_accept_bonus import (
# _compute_dflash_accept_bonus_triton_unchecked,
# )
# from sglang.srt.speculative.triton_ops.dflash_prepare_block import (
# _prepare_dflash_draft_block_unchecked,
# )# 变更后:
from sglang.srt.speculative.triton_ops.dflash import (
    _compute_dflash_accept_bonus_triton_unchecked,
    _prepare_dflash_draft_block_unchecked,
)
# 其余导入保持不变

评论区精华

没有提炼出高价值讨论线程

当前评论区没有形成足够清晰的争议点或结论,后续有更多讨论时会体现在这里。

风险与影响

极低风险。变更仅为代码移动和删除重复,不涉及逻辑修改。唯一风险是导入路径写错或遗漏符号,但 CI 中两项相关测试(test_dflash.pytest_pcg_with_speculative_decoding_dflash.py)已通过验证,确保行为不变。

影响范围:仅限 DFlash 推测解码模块。影响程度:无功能或性能变化,仅改善代码组织,降低未来维护成本。对用户无感知,对团队而言减少了重复代码,增强了与 Eagle 模块的命名一致性。

无逻辑变更 CI 测试已通过

关联 Issue

未识别关联 Issue

当前没有检测到明确关联的 Issue 链接,后续同步到相关引用后会出现在这里。

完整报告

参与讨论