已合并
[Inductor] fix rms norm reduction #42071
luqichao创建于 7月18日
[Inductor] fix rms norm reduction #42071
已合并
共 3 个文件变更+153-28
| @@ -25,6 +25,23 @@ class TestReduction(TestUtils): | |||
| 25 | view: "f32[9600, 2304]" = torch.ops.aten.view.default(add_3, [9600, 2304]) | 25 | view: "f32[9600, 2304]" = torch.ops.aten.view.default(add_3, [9600, 2304]) |
| 26 | return [None, primals_5, getitem_3, rsqrt, add_2, view, primals_2] | 26 | return [None, primals_5, getitem_3, rsqrt, add_2, view, primals_2] |
| 27 | 27 | ||
| 28 | + def rms_norm_weight_grad(self, grad_out_base, q, q_square_sum, permute_order): | ||
| 29 | + grad_out = grad_out_base.permute(*permute_order) | ||
| 30 | + inv_rms = torch.rsqrt(q_square_sum.unsqueeze(-1) / q.shape[-1] + 1e-6) | ||
| 31 | + grad_weight = (grad_out * q.float() * inv_rms).sum(dim=(0, 1, 2)) | ||
| 32 | + return grad_weight.to(torch.bfloat16) | ||
| 33 | + | ||
| 34 | + def check_rms_norm_weight_grad(self, grad_out_base, q, q_square_sum, permute_order): | ||
| 35 | + expected = self.rms_norm_weight_grad( | ||
| 36 | + grad_out_base, q, q_square_sum, permute_order | ||
| 37 | + ) | ||
| 38 | + compiled = torch.compile( | ||
| 39 | + self.rms_norm_weight_grad, backend="inductor", dynamic=False | ||
| 40 | + ) | ||
| 41 | + actual = compiled(grad_out_base, q, q_square_sum, permute_order) | ||
| 42 | + | ||
| 43 | + self.assertEqual(expected, actual, atol=1e-3, rtol=1e-3) | ||
| 44 | + | ||
| 28 | def test_reduction_cases_shapes(self): | 45 | def test_reduction_cases_shapes(self): |
| 29 | device = 'npu' | 46 | device = 'npu' |
| 30 | primals_2: "f32[32, 2304]" = torch.randn((32, 2304), device=device, dtype=torch.float32) | 47 | primals_2: "f32[32, 2304]" = torch.randn((32, 2304), device=device, dtype=torch.float32) |
| @@ -43,6 +60,22 @@ class TestReduction(TestUtils): | |||
| 43 | self.assertEqual(view_ref, view, atol=1e-3, rtol=1e-3, equal_nan=True) | 60 | self.assertEqual(view_ref, view, atol=1e-3, rtol=1e-3, equal_nan=True) |
| 44 | self.assertEqual(primals_2_ref, primals_2, atol=1e-3, rtol=1e-3, equal_nan=True) | 61 | self.assertEqual(primals_2_ref, primals_2, atol=1e-3, rtol=1e-3, equal_nan=True) |
| 45 | 62 | ||
| 63 | + def test_rms_norm_weight_grad_head_seq_permute(self): | ||
| 64 | + device = "npu" | ||
| 65 | + grad_out_base = torch.randn((2, 3, 4, 8), device=device, dtype=torch.bfloat16) | ||
| 66 | + q = torch.randn((2, 4, 3, 8), device=device, dtype=torch.bfloat16) | ||
| 67 | + q_square_sum = torch.rand((2, 4, 3), device=device, dtype=torch.float32) | ||
| 68 | + | ||
| 69 | + self.check_rms_norm_weight_grad(grad_out_base, q, q_square_sum, (0, 2, 1, 3)) | ||
| 70 | + | ||
| 71 | + def test_rms_norm_weight_grad_batch_seq_permute(self): | ||
| 72 | + device = "npu" | ||
| 73 | + grad_out_base = torch.randn((4, 2, 3, 8), device=device, dtype=torch.bfloat16) | ||
| 74 | + q = torch.randn((2, 4, 3, 8), device=device, dtype=torch.bfloat16) | ||
| 75 | + q_square_sum = torch.rand((2, 4, 3), device=device, dtype=torch.float32) | ||
| 76 | + | ||
| 77 | + self.check_rms_norm_weight_grad(grad_out_base, q, q_square_sum, (1, 0, 2, 3)) | ||
| 78 | + | ||
| 46 | 79 | ||
| 47 | if __name__ == "__main__": | 80 | if __name__ == "__main__": |
| 48 | run_tests() | 81 | run_tests() |
| @@ -455,6 +455,11 @@ class ReductionAnalysis: | |||
| 455 | if not reduction_layout_var_list: | 455 | if not reduction_layout_var_list: |
| 456 | raise RuntimeError("assert reduction_layout_var_list is not empty") | 456 | raise RuntimeError("assert reduction_layout_var_list is not empty") |
| 457 | 457 | ||
| 458 | + if self.numof_reduction_axis() > 1 and self.contiguous_reduction: | ||
| 459 | + if not self.kernel.golden_var_list: | ||
| 460 | + self.kernel.select_golden_varlist() | ||
| 461 | + return sum(1 for x in self.kernel.golden_var_list if x.name[0] != 'r') | ||
| 462 | + | ||
| 458 | dim = -1 | 463 | dim = -1 |
| 459 | for i, x in enumerate(reversed(reduction_layout_var_list)): | 464 | for i, x in enumerate(reversed(reduction_layout_var_list)): |
| 460 | if x.name[0] == 'r': | 465 | if x.name[0] == 'r': |
| @@ -2547,12 +2547,15 @@ class NPUIndexTritonKernel(TritonKernel): | |||
| 2547 | if not (self.loads or self.stores or self.compute or self.post_loop_store): | 2547 | if not (self.loads or self.stores or self.compute or self.post_loop_store): |
| 2548 | return | 2548 | return |
| 2549 | 2549 | ||
| 2550 | - def write_pointwise(): | 2550 | + def write_pointwise(allow_stores=None): |
| 2551 | + if allow_stores is None: | ||
| 2552 | + allow_stores = self.numof_reduction_axis() <= 1 | ||
| 2551 | self._emit_coordinate_transforms() | 2553 | self._emit_coordinate_transforms() |
| 2552 | self.body.splice(self.indexing_code) | 2554 | self.body.splice(self.indexing_code) |
| 2553 | self.body.splice(self.loads) | 2555 | self.body.splice(self.loads) |
| 2554 | self.body.splice(self.compute) | 2556 | self.body.splice(self.compute) |
| 2555 | - self.body.splice(self.stores) | 2557 | + if allow_stores: |
| 2558 | + self.body.splice(self.stores) | ||
| 2556 | 2559 | ||
| 2557 | def collect_store_unified_vars(): | 2560 | def collect_store_unified_vars(): |
| 2558 | """ | 2561 | """ |
| @@ -2666,6 +2669,18 @@ class NPUIndexTritonKernel(TritonKernel): | |||
| 2666 | 2669 | ||
| 2667 | reduction_1d = is_1d_reduction() | 2670 | reduction_1d = is_1d_reduction() |
| 2668 | do_indent = False | 2671 | do_indent = False |
| 2672 | + is_first_reduction_tiling = ( | ||
| 2673 | + self.numof_reduction_axis() > 1 | ||
| 2674 | + and range_val.is_tiling_axis | ||
| 2675 | + and range_val.prefix == "r" | ||
| 2676 | + and not any(ax.prefix == "r" for ax in self.sorted_axis[:index]) | ||
| 2677 | + ) | ||
| 2678 | + use_outer_reduction_post_loop = ( | ||
| 2679 | + self.numof_reduction_axis() > 1 | ||
| 2680 | + and range_val.prefix == "r" | ||
| 2681 | + and bool(self.prefix._lines) | ||
| 2682 | + ) | ||
| 2683 | + | ||
| 2669 | # tiling axis and last tiling | 2684 | # tiling axis and last tiling |
| 2670 | if range_val.is_tiling_axis and last_tiling: | 2685 | if range_val.is_tiling_axis and last_tiling: |
| 2671 | do_indent = False | 2686 | do_indent = False |
| @@ -2676,23 +2691,31 @@ class NPUIndexTritonKernel(TritonKernel): | |||
| 2676 | if ( | 2691 | if ( |
| 2677 | range_val.prefix != "r" or not self.persistent_reduction | 2692 | range_val.prefix != "r" or not self.persistent_reduction |
| 2678 | ) and need_axis_loop: | 2693 | ) and need_axis_loop: |
| 2679 | - self.body.splice(self.prefix) | 2694 | + if self.numof_reduction_axis() <= 1: |
| 2695 | + self.body.splice(self.prefix) | ||
| 2680 | self.body.writeline( | 2696 | self.body.writeline( |
| 2681 | f"for loop_{range_val.name} in range(loops_{range_val.name}):" | 2697 | f"for loop_{range_val.name} in range(loops_{range_val.name}):" |
| 2682 | ) | 2698 | ) |
| 2683 | do_indent = True | 2699 | do_indent = True |
| 2684 | loop_body(index, indexing_code, is_last_axis, do_indent) | 2700 | loop_body(index, indexing_code, is_last_axis, do_indent) |
| 2685 | - self.body.splice(self.post_loop_combine) | 2701 | + if use_outer_reduction_post_loop: |
| 2686 | - self.body.splice(self.post_loop_store) | 2702 | + pass |
 | |||
| 2687 | - # Output deferred reduction stores here (outside the loop). | 2703 | + else: |
| 2688 | - # body.writeline() controls indentation uniformly, keeping | 2704 | + if self.numof_reduction_axis() <= 1 or range_val.prefix != "r": |
| 2689 | - # these stores consistent with post_loop_store in both | 2705 | + self.body.splice(self.post_loop_combine) |
| 2690 | - # static and dynamic modes. | 2706 | + self.body.splice(self.post_loop_store) |
| 2691 | - for store_line in self._deferred_reduction_stores: | 2707 | + # Output deferred reduction stores here (outside the loop). |
| 2692 | - self.body.writeline(store_line) | 2708 | + # body.writeline() controls indentation uniformly, keeping |
| 2693 | - self._deferred_reduction_stores.clear() | 2709 | + # these stores consistent with post_loop_store in both |
| 2694 | - self.post_loop_combine.clear() | 2710 | + # static and dynamic modes. |
| 2695 | - self.post_loop_store.clear() | 2711 | + for store_line in self._deferred_reduction_stores: |
| 2712 | + self.body.writeline(store_line) | ||
| 2713 | + self._deferred_reduction_stores.clear() | ||
| 2714 | + if self.numof_reduction_axis() > 1 and range_val.prefix == "r": | ||
| 2715 | + self.body.splice(self.stores) | ||
| 2716 | + self.stores.clear() | ||
| 2717 | + self.post_loop_combine.clear() | ||
| 2718 | + self.post_loop_store.clear() | ||
| 2696 | 2719 | ||
| 2697 | # tiling axis and but not last tiling | 2720 | # tiling axis and but not last tiling |
| 2698 | elif range_val.is_tiling_axis: | 2721 | elif range_val.is_tiling_axis: |
| @@ -2703,10 +2726,24 @@ class NPUIndexTritonKernel(TritonKernel): | |||
| 2703 | indexing_code = None | 2726 | indexing_code = None |
| 2704 | if not range_val.is_no_loop_axis: | 2727 | if not range_val.is_no_loop_axis: |
| 2705 | do_indent = True | 2728 | do_indent = True |
| 2729 | + if is_first_reduction_tiling: | ||
| 2730 | + self.body.splice(self.prefix) | ||
| 2706 | self.body.writeline( | 2731 | self.body.writeline( |
| 2707 | f"for loop_{range_val.name} in range(loops_{range_val.name}):" | 2732 | f"for loop_{range_val.name} in range(loops_{range_val.name}):" |
| 2708 | ) | 2733 | ) |
| 2709 | loop_body(index, indexing_code, is_last_axis, do_indent=do_indent) | 2734 | loop_body(index, indexing_code, is_last_axis, do_indent=do_indent) |
| 2735 | + if is_first_reduction_tiling and use_outer_reduction_post_loop: | ||
| 2736 | + self.body.splice(self.post_loop_combine) | ||
| 2737 | + for store_line in self.post_loop_store._lines: | ||
| 2738 | + self.body.writeline(store_line) | ||
| 2739 | + for store_line in self.stores._lines: | ||
| 2740 | + self.body.writeline(store_line) | ||
| 2741 | + for store_line in self._deferred_reduction_stores: | ||
| 2742 | + self.body.writeline(store_line) | ||
| 2743 | + self._deferred_reduction_stores.clear() | ||
| 2744 | + self.stores.clear() | ||
| 2745 | + self.post_loop_combine.clear() | ||
| 2746 | + self.post_loop_store.clear() | ||
| 2710 | 2747 | ||
| 2711 | elif not is_last_axis: | 2748 | elif not is_last_axis: |
| 2712 | do_indent = True | 2749 | do_indent = True |
| @@ -2747,11 +2784,20 @@ class NPUIndexTritonKernel(TritonKernel): | |||
| 2747 | codegen_range(0) | 2784 | codegen_range(0) |
| 2748 | else: | 2785 | else: |
| 2749 | last_axis_order = self.tiling_axis[-1].sorted_order | 2786 | last_axis_order = self.tiling_axis[-1].sorted_order |
| 2750 | - if self.persistent_reduction and self.numof_reduction_axis() > 1: | 2787 | + skip_reduction_axes = False |
| 2788 | + if self.numof_reduction_axis() > 1: | ||
| 2751 | last_axis_order = last_axis_order - self.numof_reduction_axis() + 1 | 2789 | last_axis_order = last_axis_order - self.numof_reduction_axis() + 1 |
| 2790 | + skip_reduction_axes = not any( | ||
| 2791 | + self.find_axis_in_load_store(axis) | ||
| 2792 | + for axis in self.sorted_axis[last_axis_order:] | ||
| 2793 | + if axis.prefix == "r" | ||
| 2794 | + ) | ||
| 2752 | for _ in range(last_axis_order): | 2795 | for _ in range(last_axis_order): |
| 2753 | self.body.do_indent() | 2796 | self.body.do_indent() |
| 2754 | - codegen_range(last_axis_order) | 2797 | + if skip_reduction_axes: |
| 2798 | + write_pointwise(allow_stores=True) | ||
| 2799 | + else: | ||
| 2800 | + codegen_range(last_axis_order) | ||
| 2755 | for _ in range(last_axis_order): | 2801 | for _ in range(last_axis_order): |
| 2756 | self.body.do_unindent() | 2802 | self.body.do_unindent() |
| 2757 | 2803 | ||
| @@ -3648,6 +3694,19 @@ class NPUIndexTritonKernel(TritonKernel): | |||
| 3648 | ndims = self.triton_tensor_ndim() | 3694 | ndims = self.triton_tensor_ndim() |
| 3649 | if ndims == 1: | 3695 | if ndims == 1: |
| 3650 | return f"triton_helpers.promote_to_tensor({value})" | 3696 | return f"triton_helpers.promote_to_tensor({value})" |
| 3697 | + | ||
| 3698 | + if self.numof_reduction_axis() > 1 and self.is_contiguous_reduction(): | ||
| 3699 | + if not self.golden_var_list: | ||
| 3700 | + self.select_golden_varlist() | ||
| 3701 | + | ||
| 3702 | + dense_list = self.reduce_analysis.dense_size_list() | ||
| 3703 | + for i, axis in enumerate(reversed(self.golden_var_list)): | ||
| 3704 | + if axis.name[0] == "r": | ||
| 3705 | + dense_list[i] = "1" | ||
| 3706 | + | ||
| 3707 | + expand_str = ", ".join(dense_list) | ||
| 3708 | + return f"{value}.reshape({expand_str})" | ||
| 3709 | + | ||
| 3651 | dense_list = self.dense_size_list() | 3710 | dense_list = self.dense_size_list() |
| 3652 | dense_list[dim] = "1" | 3711 | dense_list[dim] = "1" |
| 3653 | contiguous_reduction = self.is_contiguous_reduction() | 3712 | contiguous_reduction = self.is_contiguous_reduction() |
| @@ -3773,6 +3832,8 @@ class NPUIndexTritonKernel(TritonKernel): | |||
| 3773 | 3832 | ||
| 3774 | dense_size_str = self.dense_size_str() | 3833 | dense_size_str = self.dense_size_str() |
| 3775 | value_shape = tuple(self.dense_size_list()) | 3834 | value_shape = tuple(self.dense_size_list()) |
| 3835 | + permute_order = None | ||
| 3836 | + need_permute = False | ||
| 3776 | axis_list = [] | 3837 | axis_list = [] |
| 3777 | for index in self.load_store_indexing: | 3838 | for index in self.load_store_indexing: |
| 3778 | for axis in V.kernel.range_tree_nodes: | 3839 | for axis in V.kernel.range_tree_nodes: |
| @@ -3789,18 +3850,41 @@ class NPUIndexTritonKernel(TritonKernel): | |||
| 3789 | ), | 3850 | ), |
| 3790 | value, | 3851 | value, |
| 3791 | ) | 3852 | ) |
| 3792 | - if len(dense_size_str) > 2 and ( | 3853 | + if ( |
| 3793 | - not self.persistent_reduction or self.numof_reduction_axis() != 1 | 3854 | + len(dense_size_str) > 2 |
| 3794 | - ): | 3855 | + and ( |
| 3795 | - value = self._map_tuple_or_scalar( | 3856 | + not self.persistent_reduction or self.numof_reduction_axis() != 1 |
| 3796 | - lambda v: self.cse.generate( | ||
| 3797 | - self.compute, | ||
| 3798 | - f"tl.reshape({v}, {dense_size_str})", | ||
| 3799 | - dtype=v.dtype, | ||
| 3800 | - shape=value_shape, | ||
| 3801 | - ), | ||
| 3802 | - value, | ||
| 3803 | ) | 3857 | ) |
| 3858 | + ): | ||
| 3859 | + if self.numof_reduction_axis() > 1 and self.is_contiguous_reduction(): | ||
| 3860 | + value_order = list(reversed(self.golden_var_list)) | ||
| 3861 | + target_order = [x for x in value_order if x.name[0] != "r"] + [ | ||
| 3862 | + x for x in value_order if x.name[0] == "r" | ||
| 3863 | + ] | ||
| 3864 | + permute_order = [value_order.index(x) for x in target_order] | ||
| 3865 | + current_order = list(range(len(value_order))) | ||
| 3866 | + need_permute = permute_order != current_order | ||
| 3867 | + | ||
| 3868 | + if need_permute: | ||
| 3869 | + value = self._map_tuple_or_scalar( | ||
| 3870 | + lambda v: self.cse.generate( | ||
| 3871 | + self.compute, | ||
| 3872 | + f"tl.reshape({v}.permute({permute_order}), {dense_size_str})", | ||
| 3873 | + dtype=v.dtype, | ||
| 3874 | + shape=value_shape, | ||
| 3875 | + ), | ||
| 3876 | + value, | ||
| 3877 | + ) | ||
| 3878 | + else: | ||
| 3879 | + value = self._map_tuple_or_scalar( | ||
| 3880 | + lambda v: self.cse.generate( | ||
| 3881 | + self.compute, | ||
| 3882 | + f"tl.reshape({v}, {dense_size_str})", | ||
| 3883 | + dtype=v.dtype, | ||
| 3884 | + shape=value_shape, | ||
| 3885 | + ), | ||
| 3886 | + value, | ||
| 3887 | + ) | ||
| 3804 | 3888 | ||
| 3805 | dim: int | 3889 | dim: int |
| 3806 | root_op: str | 3890 | root_op: str |
| @@ -3842,7 +3926,10 @@ class NPUIndexTritonKernel(TritonKernel): | |||
| 3842 | result_shape[dim] = "1" | 3926 | result_shape[dim] = "1" |
| 3843 | result_var: Any = self.cse.newvar(dtype=torch_acc_type, shape=tuple(result_shape)) | 3927 | result_var: Any = self.cse.newvar(dtype=torch_acc_type, shape=tuple(result_shape)) |
| 3844 | result_var.mask_vars = {var for var in masks if var[0] != "r"} # noqa: set_linter | 3928 | result_var.mask_vars = {var for var in masks if var[0] != "r"} # noqa: set_linter |
| 3845 | - cond = f"({' & '.join(masks)}).reshape({dense_size_str})" | 3929 | + cond_expr = f"({' & '.join(masks)})" |
| 3930 | + if need_permute: | ||
| 3931 | + cond_expr = f"{cond_expr}.permute({permute_order})" | ||
| 3932 | + cond = f"{cond_expr}.reshape({dense_size_str})" | ||
| 3846 | 3933 | ||
| 3847 | def where_cond(tval, fval): | 3934 | def where_cond(tval, fval): |
| 3848 | if not cond: | 3935 | if not cond: |
🟠 High Priority
在
codegen_body()的codegen_range内，当以下条件同时满足时，最后 tiling 轴分支（range_val.is_tiling_axis and last_tiling，第 2684-2718 行）会静默跳过 prefix 和所有 post-loop 代码（post_loop_combine、post_loop_store、stores、deferred_reduction_stores）：self.numof_reduction_axis() > 1（多 reduction 轴）not self.persistent_reduction导致self.prefix._lines非空，因此use_outer_reduction_post_loop= True此时：
而由于
is_first_reduction_tiling在此场景下同样为 True 但代码进入了第一分支（last_tiling），不会进入第二分支（elif range_val.is_tiling_axis:），因此第二分支中处理 post-loop 的补救逻辑（第 2735-2746 行）也不会执行。最终产生的 Triton kernel 缺少累加器初始化和最终 reduction 结果写回，是不完整的代码。该场景正是 rms_norm weight grad（对 4D tensor 的 dim (0,1,2) 连续 reduction，可能被合并为单个 tiling 轴）的典型情况，与 Issue #2543 描述的 "rms norm codegen error" 高度相关。
pass