Prhub
← 返回仓库详情

#22159 Add MLX profiling to bench_one_batch.py

sgl-project/sglang · 作者 Jonahcb · 合并时间 2026-04-09 20:45

分析状态 已生成
文件变更 1提交数 5 · 评论 7
代码增减 +40 / -9
feature performance run-ci macos

执行摘要

为 bench_one_batch.py 脚本添加 MLX 性能分析支持，支持 GPU 和系统范围分析。

根据PR body描述，动机是"添加MLX性能分析支持到bench_one_batch.py脚本，以便轻松进行LLM推理的prefill和decode阶段的GPU和系统范围性能分析"。

建议技术管理者和工程师：

  • 值得快速浏览以了解MLX profiling集成模式，特别是在条件处理和文件名适配方面的设计决策。
  • 关注gemini-code-assist[bot]提出的兼容性讨论，学习如何在跨平台代码中维护正确性。
  • 对于涉及性能分析或macOS开发的工程师，推荐精读以复用类似实现。
讨论亮点

review讨论中的核心要点：

  • gemini-code-assist[bot]指出关键问题：初始实现中硬编码attn_tp_size=1并移除attn_cp_size会破坏CUDA兼容性，建议使用条件逻辑（如基于use_mlx()）来正确派生参数，确保CUDA路径不受影响。结论是PR作者通过传递实际文件名来解决，但这个问题被强调为重要权衡。
  • gemini-code-assist[bot]还建议清理临时trace目录以避免冲突，但alexnails提出更简单的解决方案：在文件名中添加随机salt。最终实现通过将trace_filename传递给start_profile函数来管理，避免了临时文件冲突。
  • changminbark确认测试通过，LGTM（Looks Good To Me），验证了功能正确性。

实现拆解

实现主要集中在修改python/sglang/bench_one_batch.py文件：

  1. 在start_profile函数中添加MLX分支：当use_mlx()为真时，调用mx.metal.start_capture开始性能捕获，并处理trace文件名从.trace.json.gz到.gputrace的转换。
  2. 在stop_profile函数中添加MLX分支：停止捕获并保存文件。
  3. 在latency_test_run_once函数中，调整trace_filename的传递逻辑，确保在prefill和decode阶段正确传递给MLX路径。
  4. 修复了gemini-code-assist[bot]指出的CUDA兼容性问题：在prepare_mlp_sync_batch_raw调用中，使用条件逻辑避免硬编码attn_tp_size，保持对CUDA和MLX的支持。
文件 模块 状态 重要度
python/sglang/bench_one_batch.py 性能分析工具 modified 8.0

分析完成后，这里会展示 LLM 生成的相对完整源码片段和详细注释。

关键符号

start_profile stop_profile latency_test_run_once

评论区精华

CUDA 兼容性问题 正确性

gemini-code-assist[bot] 指出在 prepare_mlp_sync_batch_raw 调用中硬编码 attn_tp_size=1 并移除 attn_cp_size 会破坏 CUDA 兼容性，可能导致 TypeError 或错误行为。

结论：建议使用条件逻辑（如基于 use_mlx()）正确派生参数，以维护跨平台支持；PR 作者通过整体逻辑调整解决。 · 已解决

临时目录清理与文件名冲突 设计

gemini-code-assist[bot] 建议清理临时 trace 目录以避免冲突，alexnails 提出在文件名中添加 salt 作为替代方案。

结论：最终实现通过将 trace_filename 参数传递给 start_profile 函数来处理文件命名，避免了固定临时文件名的冲突风险。 · 已解决

功能验证 测试

changminbark 评论测试通过，LGTM，确认了输出长度和 prefill/decode 阶段计数正确。

结论：变更在测试中表现正常，未发现明显问题，增强了信心。 · 已解决

风险与影响

技术风险：

  1. 初始实现有CUDA兼容性风险：硬编码attn_tp_size可能在其他平台上导致TypeError或错误行为，但已通过条件逻辑修复。
  2. 缺少测试覆盖：PR body中没有提及添加单元测试，可能影响变更的回归测试和长期稳定性。
  3. 潜在文件命名冲突：使用固定文件名如"sglang_tmp.gputrace"可能导致并行运行时冲突，但最终实现通过传递参数缓解。
  4. 依赖MLX库：新增对mlx.core的导入，如果环境中未安装MLX，可能会影响脚本使用。

影响范围和程度：

  1. 对用户：在macOS或使用MLX后端的开发者受益，可以更方便地进行LLM推理性能分析；对CUDA用户无影响，因为修改通过use_mlx()条件检查隔离。
  2. 对系统：仅影响bench_one_batch.py脚本，不影响核心推理引擎或其他模块，影响范围小。
  3. 对团队：提升性能分析能力，有助于优化MLX平台性能；由于变更较简单，维护成本低。
CUDA 兼容性风险 缺少测试覆盖 临时文件冲突

关联 Issue

未识别关联 Issue

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

完整报告

Markdown PDF

执行摘要

  • 一句话：为bench_one_batch.py脚本添加MLX性能分析支持，支持GPU和系统范围分析。
  • 推荐动作：建议技术管理者和工程师：
    • 值得快速浏览以了解MLX profiling集成模式，特别是在条件处理和文件名适配方面的设计决策。
    • 关注gemini-code-assist[bot]提出的兼容性讨论，学习如何在跨平台代码中维护正确性。
    • 对于涉及性能分析或macOS开发的工程师，推荐精读以复用类似实现。

功能与动机

根据PR body描述，动机是"添加MLX性能分析支持到bench_one_batch.py脚本，以便轻松进行LLM推理的prefill和decode阶段的GPU和系统范围性能分析"。

实现拆解

实现主要集中在修改python/sglang/bench_one_batch.py文件：

  1. 在start_profile函数中添加MLX分支：当use_mlx()为真时，调用mx.metal.start_capture开始性能捕获，并处理trace文件名从.trace.json.gz到.gputrace的转换。
  2. 在stop_profile函数中添加MLX分支：停止捕获并保存文件。
  3. 在latency_test_run_once函数中，调整trace_filename的传递逻辑，确保在prefill和decode阶段正确传递给MLX路径。
  4. 修复了gemini-code-assist[bot]指出的CUDA兼容性问题：在prepare_mlp_sync_batch_raw调用中，使用条件逻辑避免硬编码attn_tp_size，保持对CUDA和MLX的支持。

关键文件：

  • python/sglang/bench_one_batch.py（模块 性能分析工具）: 核心变更文件，修改了profiling函数以支持MLX Metal捕获，是整个PR的唯一修改文件，直接影响性能分析功能。

关键符号：start_profile, stop_profile, latency_test_run_once

评论区精华

review讨论中的核心要点：

  • gemini-code-assist[bot]指出关键问题：初始实现中硬编码attn_tp_size=1并移除attn_cp_size会破坏CUDA兼容性，建议使用条件逻辑（如基于use_mlx()）来正确派生参数，确保CUDA路径不受影响。结论是PR作者通过传递实际文件名来解决，但这个问题被强调为重要权衡。
  • gemini-code-assist[bot]还建议清理临时trace目录以避免冲突，但alexnails提出更简单的解决方案：在文件名中添加随机salt。最终实现通过将trace_filename传递给start_profile函数来管理，避免了临时文件冲突。

  • changminbark确认测试通过，LGTM（Looks Good To Me），验证了功能正确性。

  • CUDA兼容性问题 (correctness): 建议使用条件逻辑（如基于use_mlx()）正确派生参数，以维护跨平台支持；PR作者通过整体逻辑调整解决。

  • 临时目录清理与文件名冲突 (design): 最终实现通过将trace_filename参数传递给start_profile函数来处理文件命名，避免了固定临时文件名的冲突风险。
  • 功能验证 (testing): 变更在测试中表现正常，未发现明显问题，增强了信心。

风险与影响

  • 风险：技术风险：
    1. 初始实现有CUDA兼容性风险：硬编码attn_tp_size可能在其他平台上导致TypeError或错误行为，但已通过条件逻辑修复。
    2. 缺少测试覆盖：PR body中没有提及添加单元测试，可能影响变更的回归测试和长期稳定性。
    3. 潜在文件命名冲突：使用固定文件名如"sglang_tmp.gputrace"可能导致并行运行时冲突，但最终实现通过传递参数缓解。
    4. 依赖MLX库：新增对mlx.core的导入，如果环境中未安装MLX，可能会影响脚本使用。
  • 影响：影响范围和程度：
    1. 对用户：在macOS或使用MLX后端的开发者受益，可以更方便地进行LLM推理性能分析；对CUDA用户无影响，因为修改通过use_mlx()条件检查隔离。
    2. 对系统：仅影响bench_one_batch.py脚本，不影响核心推理引擎或其他模块，影响范围小。
    3. 对团队：提升性能分析能力，有助于优化MLX平台性能；由于变更较简单，维护成本低。
  • 风险标记：CUDA兼容性风险, 缺少测试覆盖, 临时文件冲突

关联脉络

  • PR #22440 Upgrade sglang-torch-profiler-analysis SKILLS: 同样涉及性能分析工具改进，可以对比学习SGLang中性能分析功能的演进。
  • PR #22424 [AMD] Use aiter CK layernorm2d for LayerNorm to reduce NSA indexer kernel launches: 涉及性能优化和平台特定优化，与本PR的MLX profiling有相似的跨平台考量。

参与讨论