变更摘要

该 PR 旨在防止分组内核自动调优（grouped autotune）在 NPU 设备上因基准测试内存占用过大而导致 OOM。核心思路是：在生成分组基准测试参数后，估算其所需的内存占用量（estimate_grouped_benchmark_footprint），并与设备总内存的一定比例（由新配置项 symbolic_group_max_benchmark_memory_ratio 控制，默认 25%）进行比较；若超出预算则自动禁用分组调优并回退到非分组路径（同时启用 auto_blockify）。该守卫逻辑嵌入在 NPUIndexTritonKernel 的分组调优流程中。

主要改动

• 新增分组基准内存估算核心逻辑：在 symbolic_grouping.py 中新增 GroupedBenchmarkFootprint / GroupedBenchmarkArgFootprint 数据结构，以及 evaluate_grouped_benchmark_expr、required_storage_numel、estimate_grouped_benchmark_footprint 等函数，用于根据分组代表值和参数规格估算基准测试的内存占用量（含合成张量、变参克隆等峰值）。
• 新增内存预算守卫并集成到 Triton 内核：在 triton.py 的 NPUIndexTritonKernel 中新增 _disable_grouped_autotune_if_benchmark_too_large 方法，并在分组调优参数构建后调用，当估算占用超过 total_memory * ratio 时自动禁用分组并输出诊断日志。
• 新增可配置内存比率：在 config.py 中新增 symbolic_group_max_benchmark_memory_ratio 配置项（环境变量 INDUCTOR_ASCEND_SYMBOLIC_GROUP_MAX_BENCHMARK_MEMORY_RATIO，默认值 0.25），用于控制基准测试可使用的最大设备内存比例。
• 新增单元测试与集成测试：新增 test_symbolic_grouping.py 测试文件，覆盖存储容量计算、表达式求值、内存占用估算及异常边界；在 test_inductor_dynamic_shapes.py 中新增集成测试，验证宽张量场景下分组被正确禁用并启用 auto_blockify 回退机制。

openLiBing可接管PR, PR详情链接:
https://www.openlibing.com/apps/transitionPrDetail?owner=Ascend&repo=pytorch&number=43884&platform=gitcode&codeHostingPlatformFlag=gitcode

代码审查

经过全面审查所有 5 个变更文件后，我已完成深入分析。以下是最终结论：

审查总结

变更文件审查结果

1. torch_npu/_inductor/runtime/symbolic_grouping.py — 无问题

   • 新增了 GroupedBenchmarkArgFootprint、GroupedBenchmarkFootprint 数据类，以及 evaluate_grouped_benchmark_expr、required_storage_numel、_grouped_benchmark_dtype_itemsize、estimate_grouped_benchmark_footprint 函数。所有边界条件（零尺寸张量、负 stride、rank 不匹配、除零、运行时依赖表达式、未知 dtype）均通过 UnsupportedGroupedPlan 正确触发 fail-closed 行为。required_storage_numel 公式与 PyTorch 的 computeStorageNbytes 语义一致。

2. torch_npu/_inductor/codegen/triton.py — 无问题

   • 新增了 _disable_grouped_autotune_if_benchmark_too_large 方法，在 codegen_kernel 中的 build_grouped_benchmark_arg_specs 之后、_enable_auto_blockify_for_grouped_fallback_if_needed 之前被调用。异常处理通过宽泛的 except 子句实现 fail-safe（捕获 UnsupportedGroupedPlan / RuntimeError 等，禁用分组 autotune）。与已有的 _disable_grouped_autotune_if_unsupported 配合正确（先快速检查，后内存估算）。

3. torch_npu/_inductor/config.py — 无问题

   • 新增了 symbolic_group_max_benchmark_memory_ratio 配置项，默认值 0.25，有效范围 (0.0, 1.0]，与现有 costmodel_ratio 的 _parse_float_env 语义一致。可通过环境变量 INDUCTOR_ASCEND_SYMBOLIC_GROUP_MAX_BENCHMARK_MEMORY_RATIO 覆盖。

4. test/_inductor/test_symbolic_grouping.py（新文件）— 无问题

   • 新增了 TestGroupedBenchmarkFootprint 测试类，覆盖：required_storage_numel 正确性（含 broadcast stride）、非法布局 fail-closed、表达式求值、含 mutated clone 的足迹求和、运行时依赖表达式 fail-closed。断言数值经过手动验证，与实现一致。

5. test/_inductor/test_inductor_dynamic_shapes.py — 无问题

   • 新增了 TestPointwiseSymbolicGrouping 类中的 4 个测试方法：参数化宽度边界测试、预算边界包含性测试、auto_blockify 启用测试、端到端集成测试。使用 object.__new__ 创建未初始化的 kernel（因为被测试方法不依赖实例属性），通过 unittest.mock.patch 隔离外部依赖。try/finally 中正确调用 torch._dynamo.reset()。

各优先级问题数量

• P0：0
• P1：0
• P2：0
• P3：0

整体风险评估

低风险。 此变更实现了对 grouped autotune benchmark 内存占用的保守估计，在所有无法确定边界的情况下均采用 fail-closed 策略（禁用分组 autotune 并回退到 auto-blockify）。实现中包含全面的边界条件检查（零尺寸、负 stride、除零、运行时依赖表达式、非法 dtype），测试覆盖了单元级别和端到端级别的正确性。守卫逻辑在 codegen_kernel 中的位置合理（在构建 arg specs 之后、生成 kernel 代码之前），与已有的 _disable_grouped_autotune_if_unsupported 和 _enable_auto_blockify_for_grouped_fallback_if_needed 正确协作。

💬 仅评论

Pull Request 已合并或已关闭。

If you want to solve this problem, you can click here to do it in the FAQs.