基准测试支持随机种子控制#15
Conversation
![gemini-code-assist[bot]](https://avatars.githubusercontent.com/in/956858?s=60&v=4){kind=small}
There was a problem hiding this comment.
Code Review
This pull request introduces a configurable random seed for benchmarking in bench_flash_mla.py via a new --seed command-line argument and a set_benchmark_seed helper function. It also adds a unit test to verify the seeding behavior. The reviewer recommends refactoring the unit test to avoid using fragile AST parsing and exec to bypass CUDA dependencies, suggesting instead to mock the CUDA-related modules in sys.modules to allow a standard import.
Important
The consumer version of Gemini Code Assist on GitHub is being sunset. Starting June 18, 2026, new organization installations will be blocked, and all code review activity will officially cease on July 17, 2026.
For more details on the timeline and next steps, please review the Help Documentation.
| import ast | ||
| import random | ||
| from pathlib import Path | ||
| from unittest.mock import Mock | ||
|
|
||
|
|
||
| def _load_seed_helper(): | ||
| source_path = Path(__file__).resolve().parents[1] / "benchmark" / "bench_flash_mla.py" | ||
| tree = ast.parse(source_path.read_text(encoding="utf-8")) | ||
| nodes = [ | ||
| node | ||
| for node in tree.body | ||
| if ( | ||
| isinstance(node, ast.Import) | ||
| and any(alias.name == "random" for alias in node.names) | ||
| ) | ||
| or (isinstance(node, ast.FunctionDef) and node.name == "set_benchmark_seed") | ||
| ] | ||
| module = ast.Module(body=nodes, type_ignores=[]) | ||
| ast.fix_missing_locations(module) | ||
| namespace = {"torch": Mock()} | ||
| exec(compile(module, str(source_path), "exec"), namespace) | ||
| return namespace | ||
|
|
||
|
|
||
| def test_set_benchmark_seed_updates_torch_and_python_random(): | ||
| namespace = _load_seed_helper() | ||
| namespace["set_benchmark_seed"](123) | ||
|
|
||
| namespace["torch"].manual_seed.assert_called_once_with(123) | ||
| first = random.random() | ||
| random.seed(123) | ||
| assert first == random.random() | ||
|
|
||
|
|
||
| if __name__ == "__main__": | ||
| test_set_benchmark_seed_updates_torch_and_python_random() |
There was a problem hiding this comment.
使用 ast 解析和 exec 执行来动态加载测试目标函数（set_benchmark_seed）虽然规避了导入整个基准测试文件时的 CUDA 依赖问题，但这种做法非常脆弱且难以维护：\n\n1. 脆弱性高：如果 set_benchmark_seed 的实现或 bench_flash_mla.py 的模块级导入发生变化，AST 解析逻辑很容易中断。\n2. 工具链支持差：IDE 无法进行代码跳转和自动补全，静态分析工具（如 linter/type checker）无法正常工作，且单测覆盖率工具（如 coverage.py）无法正确追踪 exec 执行的代码。\n3. 测试逻辑间接：通过断言 random.random() 的返回值来间接验证 random.seed 是否被调用，不如直接 mock random.seed 并断言其调用参数来得直观和可靠。\n\n改进建议：\n利用 unittest.mock.patch.dict 临时 mock 掉 sys.modules 中的 CUDA 相关依赖（如 torch、triton、flashinfer 等），从而可以直接安全地导入 set_benchmark_seed 进行标准单元测试。这样既能避免 CUDA 环境依赖，又能保持测试代码的简洁和高可维护性。
import sys\nfrom unittest.mock import MagicMock, patch\n\n\ndef test_set_benchmark_seed_updates_torch_and_python_random():\n mocks = {\n \"triton\": MagicMock(),\n \"triton.language\": MagicMock(),\n \"flashinfer\": MagicMock(),\n \"flash_mla\": MagicMock(),\n \"torch\": MagicMock(),\n }\n with patch.dict(sys.modules, mocks):\n sys.modules.pop(\"benchmark.bench_flash_mla\", None)\n from benchmark.bench_flash_mla import set_benchmark_seed\n import torch\n\n with patch(\"random.seed\") as mock_random_seed:\n set_benchmark_seed(123)\n torch.manual_seed.assert_called_once_with(123)\n mock_random_seed.assert_called_once_with(123)\n\n\nif __name__ == \"__main__\":\n test_set_benchmark_seed_updates_torch_and_python_random()\n1 participant
该 PR 为 benchmark 增加随机种子控制，减少输入随机性对性能和正确性验证的干扰，让沐曦算力环境中的复测结果更稳定。
这个修改面向沐曦 GPU 适配场景中比较容易影响开发、构建或验证稳定性的环节，把原来需要人工排查的问题前移到工具链、运行前检查或基准脚本中处理。实现上保持对现有默认行为的兼容，只在检测到明确配置、输入或环境异常时给出更直接的诊断，避免引入额外运行依赖，也方便维护者独立审阅该分支。
已在沐曦算力环境中完成对应分支验证，验证记录包含真实运行日志、命令输出和失败路径检查，本地归档目录为：E:/Documents/muxi/测试报告/FlashMLA_real_maca_validation_20260608。提交分支：
mengz/benchmark-seed-control，目标仓库：MetaX-MACA/FlashMLA。