支持通过环境变量覆盖可用 SM 数量#11
Conversation
There was a problem hiding this comment.
Code Review
This pull request introduces support for overriding the number of SMs using the TILE_KERNELS_NUM_SMS environment variable, along with corresponding unit tests. A review comment suggests caching the parsed environment variable value in the global _num_sms variable to avoid redundant parsing and validation on subsequent calls.
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.
| env_num_sms = os.environ.get('TILE_KERNELS_NUM_SMS') | ||
| if env_num_sms: | ||
| try: | ||
| value = int(env_num_sms) | ||
| except ValueError as exc: | ||
| raise ValueError('TILE_KERNELS_NUM_SMS must be an integer') from exc | ||
| if not 0 < value <= get_device_num_sms(): | ||
| raise ValueError( | ||
| f'TILE_KERNELS_NUM_SMS must be between 1 and {get_device_num_sms()}, got {value}' | ||
| ) | ||
| return value | ||
| return get_device_num_sms() |
There was a problem hiding this comment.
在当前的实现中,如果 _num_sms 为 0 且设置了环境变量 TILE_KERNELS_NUM_SMS,每次调用 get_num_sms() 都会重新读取环境变量、进行字符串到整数的转换以及范围校验。这在频繁调用该函数时会带来不必要的性能开销。\n\n建议将解析后的环境变量值或默认的设备 SM 数量缓存到全局变量 _num_sms 中,这样后续的调用可以直接返回缓存值,避免重复解析。
env_num_sms = os.environ.get('TILE_KERNELS_NUM_SMS')\n if env_num_sms:\n try:\n value = int(env_num_sms)\n except ValueError as exc:\n raise ValueError('TILE_KERNELS_NUM_SMS must be an integer') from exc\n if not 0 < value <= get_device_num_sms():\n raise ValueError(\n f'TILE_KERNELS_NUM_SMS must be between 1 and {get_device_num_sms()}, got {value}'\n )\n _num_sms = value\n else:\n _num_sms = get_device_num_sms()?? API ?? 2 ??????
这个改动让可用 SM 数量的环境变量覆盖更加可靠,解析结果会被缓存,避免同一次进程中重复读取环境变量造成行为不一致,也让调优与压测脚本更容易稳定复现配置。
对应测试已经补齐,重点验证了环境变量解析结果在首次读取后的复用行为,本地加载配置模块后的断言结果通过。