执行摘要

• 一句话：新增 AMD 持久 CTA 工作窃取解码注意力内核
• 推荐动作：值得精读：WCA 的 persistent-CTA + work-stealing + capture-time 决策三个设计点对注意力 kernel 工程很有借鉴价值；建议重点看自动门控在 eager 与 CUDA graph 下的两套决策逻辑，以及 1xCU 网格在 MI300X/MI355X 上的 A/B 验证过程。若团队后续要支持更长上下文或引入更多 AMD 内核，本 PR 的分层（kernel / dispatch / gate / bench / test）是一个很好的模板。

功能与动机

PR 指出标准 flash-decoding（SplitK）在长上下文低 batch 时按 batch × head-tiles × num_splits 启动固定网格，长序列难以填满所有 CU，导致大量 CU 空闲而 decode 延迟随上下文近乎线性增长；在 ragged batch 下也无法把长请求的工作量再平衡到短请求空出的 CU 上。为此引入 Work-Centric Attention（WCA / Lean）：固定 CU 数量的持久网格 + 设备端工作窃取，使延迟几乎平坦并吸收 raggedness（ITL 最高降低 3.62 倍），同时保持与标准内核数值一致。

实现拆解

1. 内核与启动辅助（python/sglang/kernels/ops/attention/decode_attention.py）：新增 _lean_attention_decode_kernel、_decode_lean_attention_fwd、_lean_decode_launch_params、_lean_head_dim_ok、_lean_num_cus 等。decode_attention_fwd 在入口处增加 enable_lean、lean_Mp 等参数，并在 _is_hip、head_dim 合法、gate 通过时提前分派到 Lean 路径。持久网格按设备 CU 数确定（SGLANG_FORCE_LEAN_GRID_CU_MULT 默认 1.0），内核从 kv_indptr 推导全局 tile 计划并在线做跨 CTA 归约。

2. 后端接线（python/sglang/srt/layers/attention/triton_backend.py）：ForwardMetadata 新增 lean_Mp/lean_Lp/lean_Op/lean_locks 字段；__init__ 从 model_runner.server_args.enable_lean_attention 读取开关并预计算 lean_total_programs；init_forward_metadata 在 decode 路径按持久网格大小分配四块缓冲。决策链为 explicit override → 自动门控 → kill-switch（SGLANG_DISABLE_LEAN_ATTENTION）。

3. 配置入口（python/sglang/srt/server_args.py、environ.py）：--enable-lean-attention 是 Optional[bool]（默认 None 自动），归入 exec.kernel 命名空间；SGLANG_DISABLE_LEAN_ATTENTION 作为紧急隔离开关，SGLANG_FORCE_LEAN_GRID_CU_MULT 作为持久网格大小 A/B 调参旋钮，均无需 rebuild。

4. 自动门控：eager 模式用 lean_decode_seqlen_gate（基于 query-head 并行度 tiles = ceil(num_q_heads / min(16, kv_group)) 与 batch 放宽阈值）；CUDA graph 捕获期用 lean_capture_policy（基于 batch、tiles、is_mla 预决：GQA/MHA batch ≥ 16、MLA batch ≥ 8、tiles < 4 不启用），保证捕获图在 replay 时按真实 seq_len 工作窃取，无 host sync。

5. 测试与基准：test/registered/kernels/test_lean_attention.py 通过 _run_pair 系列在 bf16、fp8、paged KV 上对比 SplitK 与 Lean，验证余弦接近 1（分别为 0.999、0.99、0.999）并锁定 gate/policy 行为；benchmark/lean_kernel_sweep.py 对 Qwen2.5-7B 与 Llama-3.1-8B，在 batch {1..32} × context {8K..128K} 上输出延迟、加速比、余弦与 gate 决策 CSV。

关键文件：

• python/sglang/kernels/ops/attention/decode_attention.py（模块 解码内核；类别 source；类型 core-logic；符号 _lean_head_dim_ok, _lean_num_cus, _lean_decode_block_n, remap_xcd）: Lean 内核本体所在文件（+832 行）：新增 _lean_attention_decode_kernel、_decode_lean_attention_fwd、_lean_decode_launch_params 以及 lean_decode_seqlen_gate / lean_capture_policy 自动门控；decode_attention_fwd 增加 enable_lean 分派。
• test/registered/kernels/test_lean_attention.py（模块 CI 测试；类别 test；类型 test-coverage；符号 _run_pair, _lean_scratch, _run_pair_fp8, _call）: CI 门控测试（stage-b-test-1-gpu-small-amd-mi35x）：校验 Lean 与标准 SplitK 在 bf16/fp8/paged KV 上的余弦一致性，并锁定 eager gate 与 capture policy 的行为，防止内核回归。
• python/sglang/srt/layers/attention/triton_backend.py（模块 注意力后端；类别 source；类型 dependency-wiring）: 注意力后端接线：ForwardMetadata 新增 lean_Mp/Lp/Op/locks 字段；init 解析 enable_lean_attention 并缓存 _lean_decode_launch_params 结果；init_forward_metadata 按持久网格分配解码缓冲，forward_decode 按 explicit → auto-gate → kill-switch 顺序裁决。
• benchmark/lean_kernel_sweep.py（模块 性能基准；类别 source；类型 dependency-wiring；符号 bench, run, main）: Kernel 级 sweep 基准：对 Qwen2.5-7B（28Q/4KV）与 Llama-3.1-8B（32Q/8KV），在 batch {1..32} × context {8K..128K} 上对比 SplitK 与 Lean 的延迟、加速比、余弦一致性与 eager gate 决策，输出 CSV 供 PR 表格复现。
• python/sglang/srt/environ.py（模块 环境配置；类别 source；类型 core-logic）: 新增 SGLANG_DISABLE_LEAN_ATTENTION 紧急开关与 SGLANG_FORCE_LEAN_GRID_CU_MULT 持久网格扩展倍率，为线上隔离与网格 A/B 提供无需 rebuild 的旋钮。
• python/sglang/srt/server_args.py（模块 启动参数；类别 config；类型 configuration；符号 enable_lean_attention）: 新增 --enable-lean-attention 参数（Optional[bool]，默认 None 走自动门控），并归入 exec.kernel 命名空间，满足 server_args namespace CI 检查。

关键符号：_lean_attention_decode_kernel, _decode_lean_attention_fwd, _lean_decode_launch_params, _lean_head_dim_ok, lean_decode_seqlen_gate, lean_capture_policy, _should_use_lean_decode, decode_attention_fwd

关键源码片段

python/sglang/kernels/ops/attention/decode_attention.py

Lean 内核本体所在文件（+832 行）：新增 _lean_attention_decode_kernel、_decode_lean_attention_fwd、_lean_decode_launch_params 以及 lean_decode_seqlen_gate / lean_capture_policy 自动门控；decode_attention_fwd 增加 enable_lean 分派。

# decode_attention.py —— Lean 分派入口（位于 decode_attention_fwd 顶部）
# 仅在 ROCm/AMD 上开启（_is_hip），且 head_dim 必须能装进共享内存；
# logit_cap / sink / xAI 温度 / score_mod 等未实现特性则退回标准 SplitK。
if (
    _is_hip
    and _lean_head_dim_ok(k_buffer.shape[-1], v_buffer.shape[-1])
    and _should_use_lean_decode(
        enable_lean, logit_cap, sinks, xai_temperature_len, score_mod
    )
):
    # 持久网格大小只取决于设备 CU 数与 KV 头分组，与 batch 无关，
    # 因此可以放入 CUDA graph，并在每次 replay 时按 kv_indptr 推导 tile 计划。
    total_programs, XCD_REMAP, NUM_XCDS = _lean_decode_launch_params(
        v_buffer.shape[-2], kv_group_num
    )
    _decode_lean_attention_fwd(
        q, k_buffer, v_buffer, o,
        kv_indptr, kv_indices, total_programs,
        # 将 k_scale 折入 sm_scale，并单独传 v_scale：
        # 与标准 grouped kernel 的 fp8 反量化完全一致，保证数值等价。
        sm_scale * k_scale, v_scale,
        XCD_REMAP, NUM_XCDS,
        lean_Mp, lean_Lp, lean_Op, lean_locks,
        page_size=page_size,
    )
    return

评论区精华

PR 无独立 code review 评论，HaiShaw 直接 APPROVED 并无备注。Issue 侧只有 CI 命令与机器人提示（/tag-and-rerun-ci、@amd-bot ci-status、Gemini 停运通知），不涉及技术讨论。技术决策全部沉淀在 23 个 commit 中：fp8 支持、paged KV 地址数学、CUDA graph capture 策略、seq_lens_sum = None 的 EAGLE crash 修复、tl.static_range 编译崩溃修复、1xCU 网格默认值。

• CI 与合并流程 (question): 合并者认可后直接合并，无未解决技术疑虑。

风险与影响

• 风险：风险集中在默认自动门控上：AMD 上长上下文（gate 命中）的 decode 路径会切换为新内核，虽然数值等价性由 CI 测试锁定，但对 gate 阈值和网格倍率的改动未覆盖所有硬件组合。CUDA graph 捕获期曾出现 tl.static_range 编译崩溃（commit 070d1f3 改为 while 循环），说明 Triton 编译器边界仍可能在其他组合复现。fp8 路径把 k_scale 折入 sm_scale 并 cast q → K.dtype，任何改动能轻易破坏一致性，依赖 test_fp8_kv_parity 守住。_lean_head_dim_ok 对 head_dim 256（如 Gemma-2/3）会走标准 kernel，属于防护性降级，但若 LDS 预算判断不准确仍可能触发 OutOfResources。EAGLE 场景 seq_lens_sum = None 曾引发 TypeError（commit 8693652 修复），说明 speculative decoding 与 gate 的组合仍需关注。
• 影响：对 AMD 用户：在 MI355X/MI300X 上，长上下文（64K–128K）低 batch 与 ragged 场景 decode 延迟显著下降（E2E 最高 1.52×，ITL 最高 3.62×），且默认保守 gate 保证短上下文与 NVIDIA 不受影响。对系统：新增 1 个 server arg、2 个环境变量和 1 个 CI 门控测试；标准 SplitK 路径在 Lean 未启用时完全不变。对团队：内核与门控集中沉淀在 decode_attention.py，后续调参只需动 gate 阈值或环境变量，无需 rebuild，便于在 AMD 硬件上做长上下文解码的持续优化。
• 风险标记：AMD-only 分派保护, 默认自动门控影响 AMD decode 路径, CUDA graph 捕获期 Triton 编译兼容, fp8 反量化数值等价依赖, head_dim 256 共享内存溢出防护

关联脉络

• PR #36094 [AMD][DSV4] perf: retune decode split-K heuristic for MI355X: 同属 AMD 解码注意力路径的性能调优：该 PR 调优 DSv4 split-K 启发式并配套拆分启发式测试，与 Lean 共享 decode_attention 性能优化语境；两者都涉及 AMD decode 内核的 split/调度决策，可对照阅读。