#36094 [AMD][DSV4] perf: retune decode split-K heuristic for MI355X

原始 PR 作者 karverma-amd 合并时间 2026-08-29 04:48 文件变更 2 提交数 5 评论 3 代码增减 +158 / -7

执行摘要

调优 DSv4 split-K 启发式，MI355X 解码内核提速 20.8%

_kv_splits_heuristic 在 MI355X 上对整个 decode 区间过度拆分：它在 H=128/block_h=64 时为 T=32/64/128 选择了 8/4/2 个 split，而实测 4/2/1 更快。PR body 指出 Split-K 只在基础网格未填满设备时才有收益，每次额外 split 都会增加 partial buffer 写入和 reduce 内核工作；当每个 split 的 K 变短后，这些开销不再被摊还。旧常量 (2.0, 64) 的几何平均遗憾为 33.5%（最坏 119%），新常量 (1.5, 16) 降到 3.7%（最坏 36%）。

值得精读。这是一个小而完整的性能调优案例：用几何平均遗憾量化启发式偏差、识别 CUDAGraph 捕获期无法感知 per-token K 的根本限制，并用契约测试固定不变量而非数值。对 AMD 性能工作和 attention 内核开发者有直接参考价值，尤其是 test_tuned_operating_points 显式传参跨运行器的做法。

讨论亮点

核心讨论围绕平台安全性展开。HaiShaw 在 review 中要求：“Please make parameter change under _is_hip”，即重调参数不能影响未测量的 CUDA 路径；作者在 cbfa616 提交中实现 _is_hip 门控，并补充说明 test_tuned_operating_points 固定的是 MI355X 数字、在 CPU CI 上 is_hip() 为 False，因此测试改为显式传入 MI355X 常量，表格在任何运行器上保持有效。HaiShaw 随后批准并确认 “HIP specific”。没有遗留未解决问题。

实现拆解

定位过度拆分：分析 python/sglang/kernels/ops/attention/dsv4/unified_kv_kernels/paged_decode.py 中 _kv_splits_heuristic 的填充逻辑，其用 base_ctas = T * ceil(H / block_h) 与 target_wg = target_wg_per_cu * num_cu 比较决定是否拆分；MI355X 扫描（T in 1..256 × kv_len in 128/512/1024，H=128）显示旧 (2.0, 64) 的几何平均遗憾 33.5%、最坏 119%，而 (1.5, 16) 只有 3.7% / 36%。
调整常量并加入 HIP 门控：引入 from sglang.srt.utils import is_hip，定义模块级 _is_hip，_MAX_KV_SPLITS 在 HIP 上收紧为 16（CUDA 保持 64），_TARGET_WG_PER_CU 在 HIP 上降为 1.5（CUDA 保持 2.0），并让 _kv_splits_heuristic 的默认参数指向这两个常量；同步更新 docstring，记录 MI355X 测量依据与拆分开销原理。CUDA 路径不受影响。
新增契约测试：新增 test/registered/unit/layers/test_dsv4_kv_splits_heuristic.py，8 组 38 个用例，通过 register_cpu_ci(est_time=5, suite="base-a-test-cpu") 注册 CPU CI。测试断言的是不变量（pow2 上限、单调不增、饱和网格不拆、CUDAGraph 安全）而非常量值，只有 test_tuned_operating_points 用显式传入的 MI355X 常量钉住六档操作点，未来重新调优只需改这一张表。
配套提交：e3ecc36 满足 pre-commit 的 assert 换行格式；fbcc65d 注册 CPU CI（否则 test/registered/ 下的文件不会被调度）；dc3cb4e 修复 codespell 对 retuned 的误报；cbfa616 按 review 要求把常量改动收进 _is_hip 门控。

文件	模块	状态	重要度
`python/sglang/kernels/ops/attention/dsv4/unified_kv_kernels/paged_decode.py`	解码内核	modified	5.64
`test/registered/unit/layers/test_dsv4_kv_splits_heuristic.py`	启发式测试	added	7.25

关键符号

_kv_splits_heuristic _prev_pow2 test_tuned_operating_points test_splits_never_increase_with_token_count test_saturated_grid_does_not_split test_does_not_read_tensors_only_capture_time_scalars test_result_is_positive_power_of_two_within_cap

关键源码片段

python/sglang/kernels/ops/attention/dsv4/unified_kv_kernels/paged_decode.py core-logic

核心参数调优位置：`_kv_splits_heuristic` 的默认常量 `target_wg_per_cu` 与 `_MAX_KV_SPLITS` 按平台切换，并新增 `_is_hip` 门控保证 CUDA 不受影响。

# Split-K 启发式常量。 (1.5, 16) 是针对 MI355X 实测调优的组合；
# CUDA 路径没有测量，保留原有 (2.0, 64)，避免未经验证的回归。
_is_hip = is_hip()
_MAX_KV_SPLITS = 16 if _is_hip else 64 # 每个 token 的 split 硬上限
_TARGET_WG_PER_CU = 1.5 if _is_hip else 2.0


def _kv_splits_heuristic(
    T: int,
    H: int,
    block_h: int,
    num_cu: int | None = None,
    target_wg_per_cu: float = _TARGET_WG_PER_CU,
    max_kv_splits: int = _MAX_KV_SPLITS,
) -> int:
    '挑选 KV_SPLITS 来填满 GPU。 CUDAGraph 安全：只依赖捕获时标量。\n\n    原 2.0 会在整个 decode 区间多拆一个 2 的幂次（H=128 时 T=32/64/128\n    选 8/4/2，实测 4/2/1 更快）。 Split-K 只在基础网格填不满设备时才有\n    收益，额外 split 的 partial buffer 写入与 reduce 开销会在 K 变短后\n    占据主导。\n    '
    # base_ctas = T * ceil(H / block_h)：每个 token 的头部块数 × token 数
    base_ctas = T * ceil(H / block_h)
    # 目标工作项数略高于 CU 数，用来掩盖负载不均衡
    target_wg = target_wg_per_cu * num_cu
    if base_ctas >= target_wg:
        # 基础网格已饱和，再拆只会增加 reduce 开销
        return 1
    # 拆到刚好能填满设备并向下取 2 的幂次；reduce 内核按 pow2 步长
    # 索引 partial buffer，非 2 的幂次会破坏归约结果而不是变慢
    return _prev_pow2(int(min(target_wg / base_ctas, max_kv_splits)))

test/registered/unit/layers/test_dsv4_kv_splits_heuristic.py test-coverage

为之前完全没有测试的 `_kv_splits_heuristic` 新增 38 个契约测试，固定启发式不变量与 MI355X 调优操作点，并注册 CPU CI。

@pytest.mark.parametrize(
    'tokens,expected',
    [
        # MI355X（256 CU、H=128、block_h=64）上调优的操作点。
        # base_ctas = tokens * 2；target_wg = int(1.5 * 256) = 384。
        (8, 16), # 384 // 16 = 24 -> 受上限 16 约束
        (16, 8), # 384 // 32 = 12 -> prev_pow2 -> 8
        (32, 4), # 384 // 64 = 6 -> prev_pow2 -> 4
        (64, 2), # 384 // 128 = 3 -> prev_pow2 -> 2
        (128, 1), # 384 // 256 = 1
        (256, 1), # 基础网格已饱和，不再拆分
    ],
)
def test_tuned_operating_points(tokens, expected):
    '显式传入 MI355X 常量，使该表在任何 CI 运行器上都成立；\n    重新调优时只需更新这一张表。\n    '
    assert (
        _kv_splits_heuristic(
            tokens,
            HEADS,
            BLOCK_H,
            num_cu=NUM_CU,
            target_wg_per_cu=1.5,
            max_kv_splits=16,
        )
        == expected
    )


def test_does_not_read_tensors_only_capture_time_scalars():
    'CUDAGraph 安全：启发式必须能用纯 int 调用且不创建 CUDA 上下文。\n    若在捕获阶段读取 kv_indices / kv_indptr，会把捕获时的值烘焙进每次回放。\n    '
    splits = _kv_splits_heuristic(32, HEADS, BLOCK_H, num_cu=NUM_CU)
    assert isinstance(splits, int)

评论区精华

将重调参数限定在 HIP 路径 设计

HaiShaw 要求把参数变更放到 `_is_hip` 条件下，避免影响未测量的 CUDA 路径；作者在 `cbfa616` 提交中实现门控，并调整 `test_tuned_operating_points` 显式传 MI355X 常量以保证 CPU CI 上仍然有效。

结论：已实现 `is_hip()` 门控，HIP 使用 (1.5, 16)，CUDA 保持 (2.0, 64)；HaiShaw 批准并确认 'HIP specific'。 · 已解决

CPU CI 上测试与 HIP 常量的脱节 测试

`test_tuned_operating_points` 固定 MI355X 数字但运行在 CPU CI（`is_hip()` 为 False），模块默认常量会是 CUDA 的 (2.0, 64)。

结论：测试显式传入 target_wg_per_cu=1.5 与 max_kv_splits=16，表在任何运行器上有效；属性测试继续使用模块默认值。 · 已解决

风险与影响

平台覆盖不全：新常量只在 MI355X 实测，所有 HIP 设备（包括 MI300 系列）都会走 (1.5, 16)；虽然启发式本身对任意常量都安全，但其他 ROCm 设备上可能不是最优。
性能数据噪声：端到端基准每个并发点只跑一次，存在运行噪声；不过吞吐与 TPOT 在每个并发点都同向变化（一升一降），比单一指标更可信。
测试运行环境：CPU CI 上模块默认常量是 CUDA 的 (2.0, 64)，属性测试因此没有直接覆盖 HIP 常量；但属性对两组常量都成立，操作点表单独显式传参，不会造成回归盲区。
正确性风险低：选择逻辑未改，只是两个常量；GSM8K 两次运行 0.958 / 0.946、invalid 0，且拆分归约仅引入 bf16 重关联差异。

用户影响：MI355X + DP attention 部署 DeepSeek-V4-Pro 时端到端吞吐提升 0.8%～4.3%（并发越高越明显），TPOT 下降 0.8%～4.2%，解码内核生产形状（T=32、top-k 1024、H=128）从 55.7 us 降至 44.1 us（-20.8%）。普通 TP8 下效果持平。系统影响：改动集中于 ROCm-only 的 dsv4 unified_kv 后端，CUDA 保持原状；无 API、配置或部署变更。团队影响：为 _kv_splits_heuristic 补齐了此前缺失的测试覆盖，建立了契约测试模式，未来架构调参的成本显著降低。

仅 HIP 路径生效其他 ROCm 设备未验证基准单次运行 CUDAGraph 捕获期静态启发式

关联 Issue

未识别关联 Issue

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

完整报告

执行摘要

一句话：调优 DSv4 split-K 启发式，MI355X 解码内核提速 20.8%
推荐动作：值得精读。这是一个小而完整的性能调优案例：用几何平均遗憾量化启发式偏差、识别 CUDAGraph 捕获期无法感知 per-token K 的根本限制，并用契约测试固定不变量而非数值。对 AMD 性能工作和 attention 内核开发者有直接参考价值，尤其是 test_tuned_operating_points 显式传参跨运行器的做法。

功能与动机

实现拆解

定位过度拆分：分析 python/sglang/kernels/ops/attention/dsv4/unified_kv_kernels/paged_decode.py 中 _kv_splits_heuristic 的填充逻辑，其用 base_ctas = T * ceil(H / block_h) 与 target_wg = target_wg_per_cu * num_cu 比较决定是否拆分；MI355X 扫描（T in 1..256 × kv_len in 128/512/1024，H=128）显示旧 (2.0, 64) 的几何平均遗憾 33.5%、最坏 119%，而 (1.5, 16) 只有 3.7% / 36%。
调整常量并加入 HIP 门控：引入 from sglang.srt.utils import is_hip，定义模块级 _is_hip，_MAX_KV_SPLITS 在 HIP 上收紧为 16（CUDA 保持 64），_TARGET_WG_PER_CU 在 HIP 上降为 1.5（CUDA 保持 2.0），并让 _kv_splits_heuristic 的默认参数指向这两个常量；同步更新 docstring，记录 MI355X 测量依据与拆分开销原理。CUDA 路径不受影响。
新增契约测试：新增 test/registered/unit/layers/test_dsv4_kv_splits_heuristic.py，8 组 38 个用例，通过 register_cpu_ci(est_time=5, suite="base-a-test-cpu") 注册 CPU CI。测试断言的是不变量（pow2 上限、单调不增、饱和网格不拆、CUDAGraph 安全）而非常量值，只有 test_tuned_operating_points 用显式传入的 MI355X 常量钉住六档操作点，未来重新调优只需改这一张表。
配套提交：e3ecc36 满足 pre-commit 的 assert 换行格式；fbcc65d 注册 CPU CI（否则 test/registered/ 下的文件不会被调度）；dc3cb4e 修复 codespell 对 retuned 的误报；cbfa616 按 review 要求把常量改动收进 _is_hip 门控。

关键文件：

python/sglang/kernels/ops/attention/dsv4/unified_kv_kernels/paged_decode.py（模块解码内核；类别 source；类型 core-logic；符号 _kv_splits_heuristic, _MAX_KV_SPLITS, _TARGET_WG_PER_CU, _is_hip）: 核心参数调优位置：_kv_splits_heuristic 的默认常量 target_wg_per_cu 与 _MAX_KV_SPLITS 按平台切换，并新增 _is_hip 门控保证 CUDA 不受影响。
test/registered/unit/layers/test_dsv4_kv_splits_heuristic.py（模块启发式测试；类别 test；类型 test-coverage；符号 test_prev_pow2_is_largest_power_of_two_not_exceeding, test_prev_pow2_clamps_non_positive, test_result_is_positive_power_of_two_within_cap, test_splits_never_increase_with_token_count）: 为之前完全没有测试的 _kv_splits_heuristic 新增 38 个契约测试，固定启发式不变量与 MI355X 调优操作点，并注册 CPU CI。

关键符号：_kv_splits_heuristic, _prev_pow2, test_tuned_operating_points, test_splits_never_increase_with_token_count, test_saturated_grid_does_not_split, test_does_not_read_tensors_only_capture_time_scalars, test_result_is_positive_power_of_two_within_cap

关键源码片段

`python/sglang/kernels/ops/attention/dsv4/unified_kv_kernels/paged_decode.py`

核心参数调优位置：_kv_splits_heuristic 的默认常量 target_wg_per_cu 与 _MAX_KV_SPLITS 按平台切换，并新增 _is_hip 门控保证 CUDA 不受影响。

# Split-K 启发式常量。 (1.5, 16) 是针对 MI355X 实测调优的组合；
# CUDA 路径没有测量，保留原有 (2.0, 64)，避免未经验证的回归。
_is_hip = is_hip()
_MAX_KV_SPLITS = 16 if _is_hip else 64 # 每个 token 的 split 硬上限
_TARGET_WG_PER_CU = 1.5 if _is_hip else 2.0


def _kv_splits_heuristic(
    T: int,
    H: int,
    block_h: int,
    num_cu: int | None = None,
    target_wg_per_cu: float = _TARGET_WG_PER_CU,
    max_kv_splits: int = _MAX_KV_SPLITS,
) -> int:
    '挑选 KV_SPLITS 来填满 GPU。 CUDAGraph 安全：只依赖捕获时标量。

    原 2.0 会在整个 decode 区间多拆一个 2 的幂次（H=128 时 T=32/64/128\n 选 8/4/2，实测 4/2/1 更快）。 Split-K 只在基础网格填不满设备时才有\n 收益，额外 split 的 partial buffer 写入与 reduce 开销会在 K 变短后\n 占据主导。\n '
    # base_ctas = T * ceil(H / block_h)：每个 token 的头部块数 × token 数
    base_ctas = T * ceil(H / block_h)
    # 目标工作项数略高于 CU 数，用来掩盖负载不均衡
    target_wg = target_wg_per_cu * num_cu
    if base_ctas >= target_wg:
        # 基础网格已饱和，再拆只会增加 reduce 开销
        return 1
    # 拆到刚好能填满设备并向下取 2 的幂次；reduce 内核按 pow2 步长
    # 索引 partial buffer，非 2 的幂次会破坏归约结果而不是变慢
    return _prev_pow2(int(min(target_wg / base_ctas, max_kv_splits)))

`test/registered/unit/layers/test_dsv4_kv_splits_heuristic.py`

为之前完全没有测试的 _kv_splits_heuristic 新增 38 个契约测试，固定启发式不变量与 MI355X 调优操作点，并注册 CPU CI。

@pytest.mark.parametrize(
    'tokens,expected',
    [
        # MI355X（256 CU、H=128、block_h=64）上调优的操作点。
        # base_ctas = tokens * 2；target_wg = int(1.5 * 256) = 384。
        (8, 16), # 384 // 16 = 24 -> 受上限 16 约束
        (16, 8), # 384 // 32 = 12 -> prev_pow2 -> 8
        (32, 4), # 384 // 64 = 6 -> prev_pow2 -> 4
        (64, 2), # 384 // 128 = 3 -> prev_pow2 -> 2
        (128, 1), # 384 // 256 = 1
        (256, 1), # 基础网格已饱和，不再拆分
    ],
)
def test_tuned_operating_points(tokens, expected):
    '显式传入 MI355X 常量，使该表在任何 CI 运行器上都成立；\n    重新调优时只需更新这一张表。\n    '
    assert (
        _kv_splits_heuristic(
            tokens,
            HEADS,
            BLOCK_H,
            num_cu=NUM_CU,
            target_wg_per_cu=1.5,
            max_kv_splits=16,
        )
        == expected
    )


def test_does_not_read_tensors_only_capture_time_scalars():
    'CUDAGraph 安全：启发式必须能用纯 int 调用且不创建 CUDA 上下文。\n    若在捕获阶段读取 kv_indices / kv_indptr，会把捕获时的值烘焙进每次回放。\n    '
    splits = _kv_splits_heuristic(32, HEADS, BLOCK_H, num_cu=NUM_CU)
    assert isinstance(splits, int)

评论区精华

将重调参数限定在 HIP 路径 (design): 已实现 is_hip() 门控，HIP 使用 (1.5, 16)，CUDA 保持 (2.0, 64)；HaiShaw 批准并确认 'HIP specific'。
CPU CI 上测试与 HIP 常量的脱节 (testing): 测试显式传入 target_wg_per_cu=1.5 与 max_kv_splits=16，表在任何运行器上有效；属性测试继续使用模块默认值。

风险与影响

风险：
1. 平台覆盖不全：新常量只在 MI355X 实测，所有 HIP 设备（包括 MI300 系列）都会走 (1.5, 16)；虽然启发式本身对任意常量都安全，但其他 ROCm 设备上可能不是最优。
2. 性能数据噪声：端到端基准每个并发点只跑一次，存在运行噪声；不过吞吐与 TPOT 在每个并发点都同向变化（一升一降），比单一指标更可信。
3. 测试运行环境：CPU CI 上模块默认常量是 CUDA 的 (2.0, 64)，属性测试因此没有直接覆盖 HIP 常量；但属性对两组常量都成立，操作点表单独显式传参，不会造成回归盲区。
4. 正确性风险低：选择逻辑未改，只是两个常量；GSM8K 两次运行 0.958 / 0.946、invalid 0，且拆分归约仅引入 bf16 重关联差异。
  - 影响：用户影响：MI355X + DP attention 部署 DeepSeek-V4-Pro 时端到端吞吐提升 0.8%～4.3%（并发越高越明显），TPOT 下降 0.8%～4.2%，解码内核生产形状（T=32、top-k 1024、H=128）从 55.7 us 降至 44.1 us（-20.8%）。普通 TP8 下效果持平。系统影响：改动集中于 ROCm-only 的 dsv4 unified_kv 后端，CUDA 保持原状；无 API、配置或部署变更。团队影响：为 _kv_splits_heuristic 补齐了此前缺失的测试覆盖，建立了契约测试模式，未来架构调参的成本显著降低。
  - 风险标记：仅 HIP 路径生效, 其他 ROCm 设备未验证, 基准单次运行, CUDAGraph 捕获期静态启发式

关联脉络

暂无明显关联 PR

#36094 [AMD][DSV4] perf: retune decode split-K heuristic for MI355X

执行摘要

调优 DSv4 split-K 启发式，MI355X 解码内核提速 20.8%

实现拆解

评论区精华

风险与影响

关联 Issue

未识别关联 Issue

完整报告

参与讨论