执行摘要

• 一句话：FP8 模型 PCG inductor 路径性能优化
• 推荐动作：值得精读，特别是 _reshape_for_qk_norm 函数中对不同后端的条件分支设计，以及 FP8 量化路径中如何利用 PyTorch 原生操作促进编译器融合。

功能与动机

原始 PR #21734 的 inductor 融合优化提升了 FP8 模型在 H100/H200 上的吞吐，但因在 AMD GPU 上 view() 操作导致 Memory access fault 而被 revert（#23159）。本次 PR 旨在保留 inductor 融合收益，同时修复 AMD 兼容性问题。

实现拆解

1. 新增 _reshape_for_qk_norm 函数（python/sglang/srt/models/utils.py）：根据后端和编译器策略选择不同的 reshape 方式。
   - CUDA + inductor 路径：使用 view(*x.shape[:-1], -1, head_dim)，保持多维 shape 信息，让 inductor 能将 reshape 与后续 RMSNorm、FP8 量化融合为一个 Triton kernel。
   - 其他路径（ROCm、eager PCG fallback）：使用 reshape(-1, head_dim)，对非连续的 QKV-split 步长视图会触发拷贝，避免 ROCm RMSNorm kernel 因步长不连续而报错。
2. 更新 apply_qk_norm 中的调用点：将原有的 q.reshape(-1, head_dim) 和 k.reshape(-1, head_dim) 替换为 _reshape_for_qk_norm(q, head_dim) 和 _reshape_for_qk_norm(k, head_dim)，覆盖了 if alt_stream 和 else 两个分支。
3. 禁止 inductor 路径使用 fused_inplace_qknorm：在 apply_qk_norm 的条件判断中增加了 get_global_server_args().piecewise_cuda_graph_compiler != "inductor" 条件，使得 inductor 路径跳过自定义的融合 kernel，让 inductor 自动融合 QK norm。
4. 优化 apply_fp8_linear 中的 static per-tensor 量化路径（python/sglang/srt/layers/quantization/fp8_utils.py）：
   • 在 compressed_tensor_quant 分支中，当 input_scale 为 per-tensor scalar 且使用 inductor 编译器时，用纯 PyTorch 的 multiply、clamp、type-convert 替换自定义 scaled_fp8_quant kernel。
   • 这样 inductor 可以将量化与上游 RMSNorm、residual add 融合，减少 kernel launch 次数。
5. 新增导入：两个文件中均导入了 from sglang.srt.server_args import get_global_server_args 以读取编译器配置。

关键文件：

• python/sglang/srt/models/utils.py（模块 模型工具；类别 source；类型 data-contract；符号 _reshape_for_qk_norm）: 核心变更文件，新增 _reshape_for_qk_norm 函数并修改 apply_qk_norm 调用点，实现不同后端/编译器策略的分支逻辑。
• python/sglang/srt/layers/quantization/fp8_utils.py（模块 量化层；类别 source；类型 dependency-wiring）: 在 compressed_tensor_quant 分支中新增 inductor 路径的纯 PyTorch 量化实现，使 FP8 量化与周边 op 融合。

关键符号：_reshape_for_qk_norm, apply_qk_norm, apply_fp8_linear

评论区精华

Gemini Code Assist bot 建议将 apply_qk_norm 中 if 和 else 分支重复的 reshape+norm 逻辑提取为辅助函数 _normalize_tensor，以提高代码清晰度。该建议未被采纳，PR 最终以当前实现合并。

• 重复逻辑抽取为辅助函数 (style): 未采纳，PR 以当前实现合并。

风险与影响

• 风险：
  1. 回归风险（CUDA 路径）：_reshape_for_qk_norm 在 CUDA+inductor 路径使用 view()，若输入张量并非真正的步长兼容（尽管条件判断为 _is_cuda and piecewise_cuda_graph_compiler == "inductor"），仍可能引发错误，但该概率较低。
  2. 数值精度风险：FP8 量化路径中使用纯 PyTorch 的乘法、clamp 和类型转换可能引入微小的数值差异，但预期差异在量化误差范围内。
  3. 覆盖范围局限：优化仅对 compress_tensor_quant 分支生效，非该分支仍使用原有 kernel。
• 影响：
  1. 用户影响：使用 FP8 精度模型且启用 piecewise_cuda_graph_compiler=inductor 的用户将获得吞吐提升（约 10-20%，参考 #21734 数据）；AMD 用户不再遇到 crash。
  2. 系统影响：减少了 inductor 编译路径下的 kernel launch 数量，优化了 GPU 利用率。
  3. 团队影响：该改动是 #21734 的延续，表明团队对 inductor 融合策略的持续投入。
     - 风险标记：核心路径变更, 缺少测试覆盖

关联脉络

• PR #21734 perf: optimize PCG inductor path for FP8 models: 原始优化 PR，此 PR 是它的重新应用，修复了 AMD crash。
• PR #23159 Revert inductor fusion optimization due to AMD crash: 之前的 revert PR，此 PR 修复了其中提到的 AMD 问题。