已合并
[inductor]layernorm fix welford_reduce issue #45429
rain-666创建于 20 天前
[inductor]layernorm fix welford_reduce issue #45429
已合并
共 4 个文件变更+223-13
| @@ -236,6 +236,40 @@ class TestVarMean(TestUtils): | |||
| 236 | npu_config.enable_welford = previous | 236 | npu_config.enable_welford = previous |
| 237 | torch._dynamo.reset() | 237 | torch._dynamo.reset() |
| 238 | 238 | ||
| 239 | + def test_welford_simd_codegen_above_persistent_threshold(self): | ||
| 240 | + if not npu_config.is_ascend950: | ||
| 241 | + self.skipTest("Welford SIMD rollout is Ascend 950-specific") | ||
| 242 | + previous = npu_config.enable_welford | ||
| 243 | + npu_config.enable_welford = True | ||
| 244 | + torch._dynamo.reset() | ||
| 245 | + try: | ||
| 246 | + input_element = self._generate_tensor((200, 5036), "float16") | ||
| 247 | + weight = self._generate_tensor((5036,), "float16") | ||
| 248 | + bias = self._generate_tensor((5036,), "float16") | ||
| 249 | + | ||
| 250 | + def layer_norm(x, gamma, beta): | ||
| 251 | + return F.layer_norm(x, (5036,), gamma, beta, 1e-6) | ||
| 252 | + | ||
| 253 | + expected = layer_norm(input_element, weight, bias) | ||
| 254 | + compiled = torch.compile( | ||
| 255 | + layer_norm, | ||
| 256 | + backend="inductor", | ||
| 257 | + dynamic=False, | ||
| 258 | + options={"unroll_reductions_threshold": 1}, | ||
| 259 | + ) | ||
| 260 | + actual, codes = run_and_get_code( | ||
| 261 | + compiled, input_element, weight, bias | ||
| 262 | + ) | ||
| 263 | + | ||
| 264 | + self.assertEqual(expected, actual, atol=1e-1, rtol=1e-1) | ||
| 265 | + code = "\n".join(codes) | ||
| 266 | + self.assertIn("npu_kernel_type': 'simd'", code) | ||
| 267 | + self.assertIn("'vectorized_welford_axis':", code) | ||
| 268 | + self.assertNotIn("for loop_r", code) | ||
| 269 | + finally: | ||
| 270 | + npu_config.enable_welford = previous | ||
| 271 | + torch._dynamo.reset() | ||
| 272 | + | ||
| 239 | 273 | ||
| 240 | def test_welford_simd_low_precision_codegen(self, dtype): | 274 | def test_welford_simd_low_precision_codegen(self, dtype): |
| 241 | if not npu_config.is_ascend950: | 275 | if not npu_config.is_ascend950: |
| @@ -497,12 +497,22 @@ class IterationRangesEntryNPUIndex(IterationRangesEntry): | |||
| 497 | else: | 497 | else: |
| 498 | pass | 498 | pass |
| 499 | 499 | ||
| 500 | + def _is_scalar_welford_outer_axis(self): | ||
| 501 | + vector_axis = getattr(V.kernel, "vectorized_welford_axis", None) | ||
| 502 | + return ( | ||
| 503 | + vector_axis is not None | ||
| 504 | + and self.prefix != "r" | ||
| 505 | + and self.is_split_axis | ||
| 506 | + and not self.is_vectorized_split | ||
| 507 | + ) | ||
| 508 | + | ||
| 500 | def get_axis_direction(self): | 509 | def get_axis_direction(self): |
| 501 | # assume self.golden_var_list is to be correct axis order | 510 | # assume self.golden_var_list is to be correct axis order |
| 502 | if self.is_vectorized_split: | 511 | if self.is_vectorized_split: |
| 512 | + rank = self.kernel.vectorized_welford_rank() | ||
| 503 | return ( | 513 | return ( |
| 504 | "[" | 514 | "[" |
| 505 | - + ",".join([":"] + ["None"] * len(self.kernel.tiling_axis)) | 515 | + + ",".join([":"] + ["None"] * max(rank - 1, 0)) |
| 506 | + "]" | 516 | + "]" |
| 507 | ) | 517 | ) |
| 508 | 518 | ||
| @@ -527,6 +537,33 @@ class IterationRangesEntryNPUIndex(IterationRangesEntry): | |||
| 527 | def _codegen(self): | 537 | def _codegen(self): |
| 528 | self.indexing_code.clear() | 538 | self.indexing_code.clear() |
| 529 | index = None | 539 | index = None |
| 540 | + | ||
| 541 | + def initialize_welford_accumulators(): | ||
| 542 | + for acc_name in V.kernel.reduction_result_vars: | ||
| 543 | + acc_name = str(acc_name) | ||
| 544 | + if acc_name.endswith(("_acc_sum", "_acc_sum_sq", "_acc_count")): | ||
| 545 | + self.writeline( | ||
| 546 | + f"{acc_name} = tl.zeros(" | ||
| 547 | + f"{V.kernel.welford_acc_shape}, {V.kernel.welford_acc_type})" | ||
| 548 | + ) | ||
| 549 | + | ||
| 550 | + if self._is_scalar_welford_outer_axis(): | ||
| 551 | + self.writeline(f"{self.name}_mask = True") | ||
| 552 | + for var in self.var_directions: | ||
| 553 | + self.writeline(f"{var.name} = {self.name}") | ||
| 554 | + self.writeline(f"{var.name}_mask = True") | ||
| 555 | + for removed in V.kernel.range_tree_nodes_removed.values(): | ||
| 556 | + if ( | ||
| 557 | + removed.prefix == self.prefix | ||
| 558 | + and not removed.is_vectorized_split | ||
| 559 | + and removed.name != self.name | ||
| 560 | + and V.graph.sizevars.statically_known_equals( | ||
| 561 | + removed.length, self.length | ||
| 562 | + ) | ||
| 563 | + ): | ||
| 564 | + self.writeline(f"{removed.name} = {self.name}") | ||
| 565 | + self.writeline(f"{removed.name}_mask = True") | ||
| 566 | + return self.name | ||
| 530 | # for multiple reduce dims, don't need this | 567 | # for multiple reduce dims, don't need this |
| 531 | if not self.is_tiling_axis: | 568 | if not self.is_tiling_axis: |
| 532 | if self.is_vectorized_split: | 569 | if self.is_vectorized_split: |
| @@ -534,6 +571,7 @@ class IterationRangesEntryNPUIndex(IterationRangesEntry): | |||
| 534 | index = f"{self.name} = {self.codegen_index(direction)}" | 571 | index = f"{self.name} = {self.codegen_index(direction)}" |
| 535 | self.writeline(index) | 572 | self.writeline(index) |
| 536 | self._codegen_mask() | 573 | self._codegen_mask() |
| 574 | + initialize_welford_accumulators() | ||
| 537 | return self.name | 575 | return self.name |
| 538 | 576 | ||
| 539 | direction = self.get_axis_direction() | 577 | direction = self.get_axis_direction() |
| @@ -561,6 +599,13 @@ class IterationRangesEntryNPUIndex(IterationRangesEntry): | |||
| 561 | self.writeline(index) | 599 | self.writeline(index) |
| 562 | self._codegen_mask() | 600 | self._codegen_mask() |
| 563 | 601 | ||
| 602 | + # Each vectorized outer-axis tile contains independent LayerNorm | ||
| 603 | + # rows. Initialize Welford state after entering that tile, rather | ||
| 604 | + # than once per scalar outer axis, to prevent rows accumulating into | ||
| 605 | + # one another. | ||
| 606 | + if self.is_vectorized_split: | ||
| 607 | + initialize_welford_accumulators() | ||
| 608 | + | ||
| 564 | return self.name | 609 | return self.name |
| 565 | 610 | ||
| 566 | def writeline(self, line): | 611 | def writeline(self, line): |
| @@ -624,7 +669,11 @@ class IterationRangesEntryNPUIndex(IterationRangesEntry): | |||
| 624 | lines.append( | 669 | lines.append( |
| 625 | f"base_{self.name}= tl.arange(0, {BLOCK_NAME_SUB}){dtype_cast_str}" | 670 | f"base_{self.name}= tl.arange(0, {BLOCK_NAME_SUB}){dtype_cast_str}" |
| 626 | ) | 671 | ) |
| 627 | - elif self.is_tiling_axis: | 672 | + elif ( |
| 673 | + self.is_tiling_axis | ||
| 674 | + and not self._is_scalar_welford_outer_axis() | ||
| 675 | + and not self.is_vectorized_split | ||
| 676 | + ): | ||
| 628 | lines.append( | 677 | lines.append( |
| 629 | f"base_{self.name}= tl.arange(0, {BLOCK_NAME_SUB}){dtype_cast_str}" | 678 | f"base_{self.name}= tl.arange(0, {BLOCK_NAME_SUB}){dtype_cast_str}" |
| 630 | ) | 679 | ) |
| @@ -986,6 +1035,8 @@ class NPUIndexTritonKernel(TritonKernel): | |||
| 986 | self._original_compute_buf: IndentedBuffer | None = None | 1035 | self._original_compute_buf: IndentedBuffer | None = None |
| 987 | self._original_stores_buf: IndentedBuffer | None = None | 1036 | self._original_stores_buf: IndentedBuffer | None = None |
| 988 | self.vectorized_welford_axis = None | 1037 | self.vectorized_welford_axis = None |
| 1038 | + self.welford_acc_type = "tl.float32" | ||
| 1039 | + self.welford_acc_shape = "[1]" | ||
| 989 | self.full_static_welford_reduction = False | 1040 | self.full_static_welford_reduction = False |
| 990 | self.decide_codegen_dims_in_kernel() | 1041 | self.decide_codegen_dims_in_kernel() |
| 991 | 1042 | ||
| @@ -1024,7 +1075,20 @@ class NPUIndexTritonKernel(TritonKernel): | |||
| 1024 | return False | 1075 | return False |
| 1025 | if not config.triton.persistent_reductions: | 1076 | if not config.triton.persistent_reductions: |
| 1026 | return False | 1077 | return False |
| 1027 | - if npu_config.is_ascend950 : | 1078 | + reduction_node = self.find_reduction_node() |
| 1079 | + reduction_numel = self.features.reduction_numel | ||
| 1080 | + if isinstance(reduction_numel, NumelList): | ||
| 1081 | + reduction_numel = reduction_numel.numels() | ||
| 1082 | + if ( | ||
| 1083 | + npu_config.is_ascend950 | ||
| 1084 | + and npu_config.enable_welford | ||
| 1085 | + and getattr(reduction_node, "reduction_type", None) == "welford_reduce" | ||
| 1086 | + and V.graph.sizevars.statically_known_leq(reduction_numel, 8192) | ||
| 1087 | + ): | ||
| 1088 | + # A5 Welford SIMD consumes a static reduction tile. Its compiler | ||
| 1089 | + # UB checks still reject individual oversized tile configurations. | ||
| 1090 | + return True | ||
| 1091 | + if npu_config.is_ascend950: | ||
| 1028 | threshold = {ReductionHint.INNER: 4096, ReductionHint.DEFAULT: 4096}.get( | 1092 | threshold = {ReductionHint.INNER: 4096, ReductionHint.DEFAULT: 4096}.get( |
| 1029 | self.features.get_reduction_hint(), 64 | 1093 | self.features.get_reduction_hint(), 64 |
| 1030 | ) | 1094 | ) |
| @@ -1461,6 +1525,14 @@ class NPUIndexTritonKernel(TritonKernel): | |||
| 1461 | for reduction_axis in self.reduction_axis_list() | 1525 | for reduction_axis in self.reduction_axis_list() |
| 1462 | ) | 1526 | ) |
| 1463 | 1527 | ||
| 1528 | + def vectorized_welford_rank(self): | ||
| 1529 | + """Return the DSL rank after adding the vectorized outer axis.""" | ||
| 1530 | + vector_axis = self.vectorized_welford_axis | ||
| 1531 | + if vector_axis is None: | ||
| 1532 | + return len(self.golden_var_list or ()) | ||
| 1533 | + golden_vars = tuple(self.golden_var_list or ()) | ||
| 1534 | + return len(golden_vars) + (vector_axis.symbol() not in golden_vars) | ||
| 1535 | + | ||
| 1464 | def _static_welford_reduction_numel(self): | 1536 | def _static_welford_reduction_numel(self): |
| 1465 | if not self.full_static_welford_reduction: | 1537 | if not self.full_static_welford_reduction: |
| 1466 | return None | 1538 | return None |
| @@ -2677,9 +2749,36 @@ class NPUIndexTritonKernel(TritonKernel): | |||
| 2677 | ) and need_axis_loop: | 2749 | ) and need_axis_loop: |
| 2678 | if self.numof_reduction_axis() <= 1: | 2750 | if self.numof_reduction_axis() <= 1: |
| 2679 | self.body.splice(self.prefix) | 2751 | self.body.splice(self.prefix) |
| 2680 | - self.body.writeline( | 2752 | + self.prefix.clear() |
| 2681 | - f"for loop_{range_val.name} in range(loops_{range_val.name}):" | 2753 | + if range_val._is_scalar_welford_outer_axis(): |
| 2682 | - ) | 2754 | + if self.prefix._lines: |
| 2755 | + self.body.splice(self.prefix) | ||
| 2756 | + self.prefix.clear() | ||
| 2757 | + self.body.writeline( | ||
| 2758 | + f"for {range_val.name} in range(" | ||
| 2759 | + f"{range_val.name}_offset, min(" | ||
| 2760 | + f"{range_val.name}_offset + {range_val.name.upper()}BLOCK, " | ||
| 2761 | + f"{range_val.name}_numel)):" | ||
| 2762 | + ) | ||
| 2763 | + elif range_val.is_vectorized_split: | ||
| 2764 | + block_sub = f"{range_val.name.upper()}BLOCK_SUB" | ||
| 2765 | + if range_val.is_split_axis: | ||
| 2766 | + loop_start = f"{range_val.name}_offset" | ||
| 2767 | + loop_end = ( | ||
| 2768 | + f"min({loop_start} + {range_val.name.upper()}BLOCK, " | ||
| 2769 | + f"{range_val.name}_numel)" | ||
| 2770 | + ) | ||
| 2771 | + else: | ||
| 2772 | + loop_start = "0" | ||
| 2773 | + loop_end = f"{range_val.name}_numel" | ||
| 2774 | + self.body.writeline( | ||
| 2775 | + f"for {range_val.name}_loop_offset in range(" | ||
| 2776 | + f"{loop_start}, {loop_end}, {block_sub}):" | ||
| 2777 | + ) | ||
| 2778 | + else: | ||
| 2779 | + self.body.writeline( | ||
| 2780 | + f"for loop_{range_val.name} in range(loops_{range_val.name}):" | ||
| 2781 | + ) | ||
| 2683 | do_indent = True | 2782 | do_indent = True |
| 2684 | loop_body(index, indexing_code, is_last_axis, do_indent) | 2783 | loop_body(index, indexing_code, is_last_axis, do_indent) |
| 2685 | if use_outer_reduction_post_loop: | 2784 | if use_outer_reduction_post_loop: |
| @@ -2726,13 +2825,41 @@ class NPUIndexTritonKernel(TritonKernel): | |||
| 2726 | ) | 2825 | ) |
| 2727 | if elide_full_reduction_loop and is_first_reduction_tiling: | 2826 | if elide_full_reduction_loop and is_first_reduction_tiling: |
| 2728 | self.body.splice(self.prefix) | 2827 | self.body.splice(self.prefix) |
| 2828 | + self.prefix.clear() | ||
| 2729 | if not range_val.is_no_loop_axis and not elide_full_reduction_loop: | 2829 | if not range_val.is_no_loop_axis and not elide_full_reduction_loop: |
| 2730 | do_indent = True | 2830 | do_indent = True |
| 2731 | if is_first_reduction_tiling: | 2831 | if is_first_reduction_tiling: |
| 2732 | self.body.splice(self.prefix) | 2832 | self.body.splice(self.prefix) |
| 2733 | - self.body.writeline( | 2833 | + self.prefix.clear() |
| 2734 | - f"for loop_{range_val.name} in range(loops_{range_val.name}):" | 2834 | + if range_val._is_scalar_welford_outer_axis(): |
| 2735 | - ) | 2835 | + if self.prefix._lines: |
| 2836 | + self.body.splice(self.prefix) | ||
| 2837 | + self.prefix.clear() | ||
| 2838 | + self.body.writeline( | ||
| 2839 | + f"for {range_val.name} in range(" | ||
| 2840 | + f"{range_val.name}_offset, min(" | ||
| 2841 | + f"{range_val.name}_offset + {range_val.name.upper()}BLOCK, " | ||
| 2842 | + f"{range_val.name}_numel)):" | ||
| 2843 | + ) | ||
| 2844 | + elif range_val.is_vectorized_split: | ||
| 2845 | + block_sub = f"{range_val.name.upper()}BLOCK_SUB" | ||
| 2846 | + if range_val.is_split_axis: | ||
| 2847 | + loop_start = f"{range_val.name}_offset" | ||
| 2848 | + loop_end = ( | ||
| 2849 | + f"min({loop_start} + {range_val.name.upper()}BLOCK, " | ||
| 2850 | + f"{range_val.name}_numel)" | ||
| 2851 | + ) | ||
| 2852 | + else: | ||
| 2853 | + loop_start = "0" | ||
| 2854 | + loop_end = f"{range_val.name}_numel" | ||
| 2855 | + self.body.writeline( | ||
| 2856 | + f"for {range_val.name}_loop_offset in range(" | ||
| 2857 | + f"{loop_start}, {loop_end}, {block_sub}):" | ||
| 2858 | + ) | ||
| 2859 | + else: | ||
| 2860 | + self.body.writeline( | ||
| 2861 | + f"for loop_{range_val.name} in range(loops_{range_val.name}):" | ||
| 2862 | + ) | ||
| 2736 | loop_body(index, indexing_code, is_last_axis, do_indent=do_indent) | 2863 | loop_body(index, indexing_code, is_last_axis, do_indent=do_indent) |
| 2737 | if is_first_reduction_tiling and use_outer_reduction_post_loop: | 2864 | if is_first_reduction_tiling and use_outer_reduction_post_loop: |
| 2738 | self.body.splice(self.post_loop_combine) | 2865 | self.body.splice(self.post_loop_combine) |
| @@ -2988,7 +3115,10 @@ class NPUIndexTritonKernel(TritonKernel): | |||
| 2988 | if reduction is not None and isinstance(reduction, ir.Reduction): | 3115 | if reduction is not None and isinstance(reduction, ir.Reduction): |
| 2989 | return reduction | 3116 | return reduction |
| 2990 | 3117 | ||
| 2991 | - for node in self.node_schedule: | 3118 | + # Called during the base kernel constructor, before the subclass has |
| 3119 | + # assigned ``node_schedule``. | ||
| 3120 | + node_schedule = getattr(self, "node_schedule", self.features.node_schedule) | ||
| 3121 | + for node in node_schedule: | ||
| 2992 | if node in (EnableReduction, DisableReduction): | 3122 | if node in (EnableReduction, DisableReduction): |
| 2993 | continue | 3123 | continue |
| 2994 | reduction = node.node.data | 3124 | reduction = node.node.data |
| @@ -3733,7 +3863,9 @@ class NPUIndexTritonKernel(TritonKernel): | |||
| 3733 | def reduction_resize(self, value, dim): | 3863 | def reduction_resize(self, value, dim): |
| 3734 | if self.vectorized_welford_axis is not None: | 3864 | if self.vectorized_welford_axis is not None: |
| 3735 | block_sub = f"{self.vectorized_welford_axis.name.upper()}BLOCK_SUB" | 3865 | block_sub = f"{self.vectorized_welford_axis.name.upper()}BLOCK_SUB" |
| 3736 | - broadcast_dims = [block_sub] + ["1"] * len(self.tiling_axis) | 3866 | + broadcast_dims = [block_sub] + ["1"] * max( |
| 3867 | + self.vectorized_welford_rank() - 1, 0 | ||
| 3868 | + ) | ||
| 3737 | return f"{value}.reshape([{', '.join(broadcast_dims)}])" | 3869 | return f"{value}.reshape([{', '.join(broadcast_dims)}])" |
| 3738 | 3870 | ||
| 3739 | ndims = self.triton_tensor_ndim() | 3871 | ndims = self.triton_tensor_ndim() |
| @@ -4057,9 +4189,12 @@ class NPUIndexTritonKernel(TritonKernel): | |||
| 4057 | if vector_axis is not None: | 4189 | if vector_axis is not None: |
| 4058 | block_sub = f"{vector_axis.name.upper()}BLOCK_SUB" | 4190 | block_sub = f"{vector_axis.name.upper()}BLOCK_SUB" |
| 4059 | acc_shape = f"[{block_sub}, 1]" | 4191 | acc_shape = f"[{block_sub}, 1]" |
| 4192 | + self.welford_acc_type = acc_type | ||
| 4193 | + self.welford_acc_shape = acc_shape | ||
| 4060 | accumulator_shape = (block_sub, "1") | 4194 | accumulator_shape = (block_sub, "1") |
| 4061 | resized_shape = tuple( | 4195 | resized_shape = tuple( |
| 4062 | - [block_sub] + ["1"] * len(self.tiling_axis) | 4196 | + [block_sub] |
| 4197 | + + ["1"] * max(self.vectorized_welford_rank() - 1, 0) | ||
| 4063 | ) | 4198 | ) |
| 4064 | reduce_dim = 1 | 4199 | reduce_dim = 1 |
| 4065 | keep_dims = ", keep_dims=True" | 4200 | keep_dims = ", keep_dims=True" |
| @@ -172,6 +172,10 @@ def _read_env_bool(name: str, default: str = "False") -> bool: | |||
| 172 | # Keep it disabled by default while the new path is being rolled out. | 172 | # Keep it disabled by default while the new path is being rolled out. |
| 173 | enable_welford = os.getenv("TORCHINDUCTOR_ENABLE_WELFORD", "0") == "1" | 173 | enable_welford = os.getenv("TORCHINDUCTOR_ENABLE_WELFORD", "0") == "1" |
| 174 | 174 | ||
| 175 | +# Keep the A5 LayerNormV4 workaround disabled by default. Set this to 1 to | ||
| 176 | +# enable the width-512 LayerNormV4 fallback while Welford is enabled. | ||
| 177 | +enable_layernorm_v4 = os.getenv("TORCHINDUCTOR_ENABLE_LAYERNORM_V4", "0") == "1" | ||
| 178 | + | ||
| 175 | 179 | ||
| 176 | class catlass: | 180 | class catlass: |
| 177 | # Whether to enable debug info, e.g., line number | 181 | # Whether to enable debug info, e.g., line number |
| @@ -1052,12 +1052,49 @@ def _register_npu_inductor_fallbacks(): | |||
| 1052 | bias=None, | 1052 | bias=None, |
| 1053 | eps=1e-5 | 1053 | eps=1e-5 |
| 1054 | ): | 1054 | ): |
| 1055 | + def should_use_layer_norm_v4(): | ||
| 1056 | + """Keep the A5 W=512 large-row case on the fused CANN kernel. | ||
| 1057 | + | ||
| 1058 | + The Welford lowering currently materializes statistics and emits | ||
| 1059 | + separate pointwise post-processing kernels for this shape. Until | ||
| 1060 | + the reduction epilogue is fused into the Welford kernel, the | ||
| 1061 | + dedicated LayerNormV4 implementation is both faster and avoids | ||
| 1062 | + the SIMD-reduction/SIMT-post-processing split. | ||
| 1063 | + """ | ||
| 1064 | + if ( | ||
| 1065 | + not is_ascend950 | ||
| 1066 | + or not npu_config.enable_welford | ||
| 1067 | + or not npu_config.enable_layernorm_v4 | ||
| 1068 | + ): | ||
| 1069 | + return False | ||
| 1070 | + if x.dtype not in (torch.float16, torch.bfloat16): | ||
| 1071 | + return False | ||
| 1072 | + if not isinstance(normalized_shape, (list, tuple)): | ||
| 1073 | + shape = (normalized_shape,) | ||
| 1074 | + else: | ||
| 1075 | + shape = tuple(normalized_shape) | ||
| 1076 | + if len(shape) != 1 or shape[0] != 512: | ||
| 1077 | + return False | ||
| 1078 | + | ||
| 1079 | + input_shape = x.get_size() | ||
| 1080 | + if len(input_shape) < 2: | ||
| 1081 | + return False | ||
| 1082 | + row_numel = sympy_product(input_shape[:-1]) | ||
| 1083 | + # Unbacked dynamic rows such as ``u0 + 200`` cannot be proven to | ||
| 1084 | + # exceed the threshold at compile time even when the profiled | ||
| 1085 | + # runtime value is 17000. Keep Welford only when the compiler can | ||
| 1086 | + # prove this is a genuinely small-row case. | ||
| 1087 | + return not V.graph.sizevars.statically_known_lt(row_numel, 512) | ||
| 1088 | + | ||
| 1055 | # Keep the existing low-precision fallback unless Welford is explicitly | 1089 | # Keep the existing low-precision fallback unless Welford is explicitly |
| 1056 | # enabled for FP16/BF16 on Ascend 950. | 1090 | # enabled for FP16/BF16 on Ascend 950. |
| 1057 | if ( | 1091 | if ( |
| 1058 | is_ascend950 | 1092 | is_ascend950 |
| 1059 | and x.dtype in (torch.float16, torch.bfloat16) | 1093 | and x.dtype in (torch.float16, torch.bfloat16) |
| 1060 | - and not npu_config.enable_welford | 1094 | + and ( |
| 1095 | + not npu_config.enable_welford | ||
| 1096 | + or should_use_layer_norm_v4() | ||
| 1097 | + ) | ||
| 1061 | ): | 1098 | ): |
| 1062 | return fallback_handler(aten.native_layer_norm.default)(x, normalized_shape, weight, bias, eps) | 1099 | return fallback_handler(aten.native_layer_norm.default)(x, normalized_shape, weight, bias, eps) |
| 1063 | # Validate input | 1100 | # Validate input |