[KernelGen][Nvidia] Add _amp_foreach_non_finite_check_and_unscale_ operator with Triton kernel (flagos-ai#3799)

* [KernelGen][Nvidia] Add _amp_foreach_non_finite_check_and_unscale_ operator with Triton kernel

* fix: use to_reference for inv_scale in reference path

When running with --ref cpu, ref_tensors and ref_found_inf are moved to CPU
but inv_scale stayed on CUDA, causing a device mismatch error in the
reference call to torch._amp_foreach_non_finite_check_and_unscale_.

* Update operator stage to alpha 5.4

* fix: address PR check findings for _amp_foreach_non_finite_check_and_unscale_

* fix: sort ops __all__ for amp foreach

---------

Co-authored-by: yangy0906 <yangyang0906c@163.com>

Author: XDYuanzhuLee

benchmark/test_amp_foreach_non_finite_check_and_unscale_.py

@@ -0,0 +1,53 @@
+from typing import Generator
+
+import pytest
+import torch
+
+from . import base
+
+
+class AmpForeachNonFiniteCheckAndUnscaleBenchmark(base.Benchmark):
+    """
+    Benchmark class for _amp_foreach_non_finite_check_and_unscale_ operation.
+    This operation takes a list of tensors, a found_inf tensor, and an inv_scale
+    scalar tensor, and checks for non-finite values while unscaling.
+    """
+
+    # Common sizes for gradient tensors in training scenarios covering small to medium workloads
+    DEFAULT_SHAPES = [
+        (1024, 1024),
+        (2048, 2048),
+        (4096, 4096),
+    ]
+    # shapes describe the first tensor; a second tensor at half size is generated
+    DEFAULT_SHAPE_DESC = "M, N"
+
+    def set_more_shapes(self):
+        # larger shapes for comprehensive benchmark level
+        more_shapes_2d = [(1024, 2**i) for i in range(2, 14, 4)]
+        more_shapes_3d = [(64, 2**i, 64) for i in range(2, 10, 4)]
+        return more_shapes_2d + more_shapes_3d
+
+    def get_input_iter(self, dtype) -> Generator:
+        for shape in self.shapes:
+            # generate 2 tensors: one at shape, one at half size in the first dim
+            second_shape = (max(1, shape[0] // 2),) + shape[1:]
+            tensors = [
+                torch.randn(shape, device=self.device, dtype=dtype),
+                torch.randn(second_shape, device=self.device, dtype=dtype),
+            ]
+            # PyTorch expects inv_scale and found_inf as float32
+            inv_scale = torch.tensor(2.0, device=self.device, dtype=torch.float32)
+            found_inf = torch.tensor(0.0, device=self.device, dtype=torch.float32)
+            yield tensors, found_inf, inv_scale
+
+
+@pytest.mark.amp_foreach_non_finite_check_and_unscale_
+def test_amp_foreach_non_finite_check_and_unscale_():
+    bench = AmpForeachNonFiniteCheckAndUnscaleBenchmark(
+        op_name="amp_foreach_non_finite_check_and_unscale_",
+        torch_op=torch._amp_foreach_non_finite_check_and_unscale_,
+        # bfloat16 is not supported by the CUDA kernel for this operator
+        dtypes=[torch.float16, torch.float32],
+    )
+    bench.run()

conf/operators.yaml

@@ -436,6 +436,20 @@ ops:
       - Tensor
     stages:
       - beta: '5.3'
+  - id: amp_foreach_non_finite_check_and_unscale_
+    description: |
+      Checks gradient tensors for non-finite values (inf or nan) and
+      unscales each tensor in-place by `inv_scale`. If any non-finite
+      value is detected, `found_inf` is set to 1.0.
+    for:
+      - _amp_foreach_non_finite_check_and_unscale_
+    labels:
+      - aten
+      - KernelGen
+    kind:
+      - Math
+    stages:
+      - alpha: '5.4'
   - id: angle
     description: Computes the element-wise angle (in radians) of the given `input` tensor.
     for:

src/flag_gems/__init__.py

@@ -43,6 +43,10 @@ def torch_ge(v):
     ("__or__.Scalar", bitwise_or_scalar),
     ("__or__.Tensor", bitwise_or_tensor),
     ("_adaptive_avg_pool2d", adaptive_avg_pool2d),
+    (
+        "_amp_foreach_non_finite_check_and_unscale_",
+        _amp_foreach_non_finite_check_and_unscale_,
+    ),
     ("_assert_async", _assert_async),
     ("_cdist_backward", _cdist_backward),
     ("_conv_depthwise2d", _conv_depthwise2d),

src/flag_gems/ops/__init__.py

@@ -1,4 +1,7 @@
 from flag_gems.ops.__ilshift__ import __ilshift__
+from flag_gems.ops._amp_foreach_non_finite_check_and_unscale_ import (
+    _amp_foreach_non_finite_check_and_unscale_,
+)
 from flag_gems.ops._euclidean_dist import _euclidean_dist
 from flag_gems.ops._functional_sym_constrain_range_for_size import (
     _functional_sym_constrain_range_for_size,
@@ -482,6 +485,7 @@
     "SUPPORTED_FP8_DTYPE",
     "ScaleDotProductAttention",
     "__ilshift__",
+    "_amp_foreach_non_finite_check_and_unscale_",
     "_assert_async",
     "_cdist_backward",
     "_conv_depthwise2d",
@@ -674,13 +678,13 @@
     "feature_dropout",
     "feature_dropout_",
     "fft",
-    "fix",
     "fill_scalar",
     "fill_scalar_",
     "fill_scalar_out",
     "fill_tensor",
     "fill_tensor_",
     "fill_tensor_out",
+    "fix",
     "flash_attention_backward",
     "flash_attention_forward",
     "flash_attn_varlen_func",
@@ -857,12 +861,12 @@
     "nonzero",
     "nonzero_numpy",
     "normal_",
-    "not_equal",
-    "not_equal_scalar",
     "normal_float_tensor",
     "normal_tensor_float",
     "normal_tensor_tensor",
     "normed_cumsum",
+    "not_equal",
+    "not_equal_scalar",
     "one_hot",
     "ones",
     "ones_like",

src/flag_gems/ops/_amp_foreach_non_finite_check_and_unscale_.py

@@ -0,0 +1,136 @@
+# Copyright 2026, The FlagOS Contributors.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+#
+# Generated by KernelGen: https://github.com/flagos-ai/KernelGen
+import logging
+from typing import List
+
+import torch
+import triton
+import triton.language as tl
+
+from flag_gems.utils import tl_extra_shim
+
+logger = logging.getLogger(__name__)
+
+
+@triton.jit
+def _amp_foreach_non_finite_check_and_unscale_kernel(
+    inp_ptr,
+    inv_scale_ptr,
+    output_ptr,
+    num_elements,
+    BLOCK_SIZE: tl.constexpr,
+):
+    pid = tl.program_id(0)
+    block_start = pid * BLOCK_SIZE
+    offsets = tl.arange(0, BLOCK_SIZE)
+    mask = block_start + offsets < num_elements
+
+    # Load the input value
+    inp = tl.load(inp_ptr + block_start + offsets, mask=mask, other=0.0)
+
+    # Load the inv_scale value
+    scale = tl.load(inv_scale_ptr)
+
+    # Convert to float32 for computation to handle float16/bfloat16
+    inp_fp32 = inp.to(tl.float32)
+    scale_fp32 = scale.to(tl.float32)
+
+    # Check for non-finite values (inf or nan) using float32
+    is_non_finite = ~tl_extra_shim.finitef(inp_fp32)
+
+    # Scale the values: only finite values are scaled
+    # Non-finite values stay as-is (in original dtype)
+    scaled_fp32 = tl.where(is_non_finite, inp_fp32, inp_fp32 * scale_fp32)
+
+    # Convert back to original dtype
+    # We need to determine the output dtype from the input pointer
+    # For now, assume input dtype is preserved
+    scaled = scaled_fp32.to(inp.dtype)
+
+    # Store the result
+    tl.store(output_ptr + block_start + offsets, scaled, mask=mask)
+
+
+def _amp_foreach_non_finite_check_and_unscale_(
+    tensors: List[torch.Tensor],
+    found_inf: torch.Tensor,
+    inv_scale: torch.Tensor,
+):
+    """
+    Check for non-finite values in tensors and unscale them.
+
+    For each tensor in the list:
+    - Scale finite values by inv_scale
+    - Non-finite values (inf, nan) remain unchanged
+
+    If any tensor has non-finite values, set found_inf to 1.0.
+    """
+    logger.debug("GEMS AMP_FOREACH_NON_FINITE_CHECK_AND_UNSCALE")
+
+    if not isinstance(tensors, (list, tuple)):
+        raise TypeError(f"Expected list or tuple of tensors, got {type(tensors)}")
+
+    if len(tensors) == 0:
+        return
+
+    # Ensure inv_scale is a float32 tensor as PyTorch expects
+    inv_scale = inv_scale.to(dtype=torch.float32)
+
+    # Process each tensor in the list
+    for tensor in tensors:
+        if not tensor.is_floating_point():
+            # Skip non-floating point tensors
+            continue
+
+        num_elements = tensor.numel()
+        if num_elements == 0:
+            continue
+
+        # Allocate output tensor (in-place operation)
+        output = torch.empty_like(tensor)
+
+        # 1024 offers a good balance of occupancy and parallelism for typical gradient tensor sizes
+        BLOCK_SIZE = 1024
+        grid = (triton.cdiv(num_elements, BLOCK_SIZE),)
+
+        _amp_foreach_non_finite_check_and_unscale_kernel[grid](
+            tensor,
+            inv_scale,
+            output,
+            num_elements,
+            BLOCK_SIZE=BLOCK_SIZE,
+        )
+
+        # Copy output back to input (in-place)
+        tensor.copy_(output)
+
+    # Check if any tensor has non-finite values
+    # This needs to be done after scaling since scaling might produce inf from overflow
+    for tensor in tensors:
+        if not tensor.is_floating_point():
+            continue
+
+        if tensor.is_complex():
+            # Check real and imaginary parts separately
+            if not torch.all(torch.isfinite(tensor.real)) or not torch.all(
+                torch.isfinite(tensor.imag)
+            ):
+                found_inf.fill_(1.0)
+                return
+        else:
+            if not torch.all(torch.isfinite(tensor)):
+                found_inf.fill_(1.0)
+                return

tests/test_amp_foreach_non_finite_check_and_unscale_.py

@@ -0,0 +1,125 @@
+import pytest
+import torch
+
+import flag_gems
+
+from . import accuracy_utils as utils
+
+# worktree explicitly excludes bfloat16 via PRIMARY_FLOAT_DTYPES
+PRIMARY_FLOAT_DTYPES = [torch.float16, torch.float32]
+
+
+@pytest.mark.amp_foreach_non_finite_check_and_unscale_
+@pytest.mark.parametrize("dtype", PRIMARY_FLOAT_DTYPES)
+def test_amp_foreach_non_finite_check_and_unscale_(dtype):
+    """Test _amp_foreach_non_finite_check_and_unscale_ with normal tensors."""
+    # PyTorch expects inv_scale and found_inf to be float32
+    inv_scale = torch.tensor(2.0, device=flag_gems.device, dtype=torch.float32)
+    found_inf = torch.tensor(0.0, device=flag_gems.device, dtype=torch.float32)
+
+    tensors = [
+        torch.randn(16, 32, device=flag_gems.device, dtype=dtype),
+        torch.randn(8, 16, device=flag_gems.device, dtype=dtype),
+    ]
+
+    # Reference
+    ref_tensors = [utils.to_reference(t.clone()) for t in tensors]
+    ref_found_inf = utils.to_reference(found_inf.clone())
+    ref_inv_scale = utils.to_reference(inv_scale.clone())
+    getattr(torch, "_amp_foreach_non_finite_check_and_unscale_")(
+        ref_tensors, ref_found_inf, ref_inv_scale
+    )
+
+    # GEMS
+    res_tensors = [t.clone() for t in tensors]
+    res_found_inf = found_inf.clone()
+    with flag_gems.use_gems():
+        getattr(torch, "_amp_foreach_non_finite_check_and_unscale_")(
+            res_tensors, res_found_inf, inv_scale
+        )
+
+    # Compare mutated inputs (in-place operation)
+    for i, (inp, ref_inp) in enumerate(zip(res_tensors, ref_tensors)):
+        utils.gems_assert_close(inp, ref_inp, dtype)
+
+    # Compare found_inf (also mutated in-place)
+    utils.gems_assert_equal(res_found_inf, ref_found_inf)
+
+
+@pytest.mark.amp_foreach_non_finite_check_and_unscale_
+@pytest.mark.parametrize("dtype", PRIMARY_FLOAT_DTYPES)
+def test_amp_foreach_non_finite_check_and_unscale__inf(dtype):
+    """Test _amp_foreach_non_finite_check_and_unscale_ with inf values."""
+    # PyTorch expects inv_scale and found_inf to be float32
+    inv_scale = torch.tensor(2.0, device=flag_gems.device, dtype=torch.float32)
+    found_inf = torch.tensor(0.0, device=flag_gems.device, dtype=torch.float32)
+
+    tensors = [
+        torch.tensor(
+            [1.0, 2.0, float("inf"), 4.0], device=flag_gems.device, dtype=dtype
+        ),
+        torch.tensor([5.0, 6.0, 7.0], device=flag_gems.device, dtype=dtype),
+    ]
+
+    # Reference
+    ref_tensors = [utils.to_reference(t.clone()) for t in tensors]
+    ref_found_inf = utils.to_reference(found_inf.clone())
+    ref_inv_scale = utils.to_reference(inv_scale.clone())
+    getattr(torch, "_amp_foreach_non_finite_check_and_unscale_")(
+        ref_tensors, ref_found_inf, ref_inv_scale
+    )
+
+    # GEMS
+    res_tensors = [t.clone() for t in tensors]
+    res_found_inf = found_inf.clone()
+    with flag_gems.use_gems():
+        getattr(torch, "_amp_foreach_non_finite_check_and_unscale_")(
+            res_tensors, res_found_inf, inv_scale
+        )
+
+    # Compare mutated inputs (in-place operation)
+    # Note: inf values remain unchanged, only finite values are scaled
+    for i, (inp, ref_inp) in enumerate(zip(res_tensors, ref_tensors)):
+        utils.gems_assert_close(inp, ref_inp, dtype)
+
+    # Compare found_inf (also mutated in-place) - should be 1.0 when inf is present
+    utils.gems_assert_equal(res_found_inf, ref_found_inf)
+
+
+@pytest.mark.amp_foreach_non_finite_check_and_unscale_
+@pytest.mark.parametrize("dtype", PRIMARY_FLOAT_DTYPES)
+def test_amp_foreach_non_finite_check_and_unscale__nan(dtype):
+    """Test _amp_foreach_non_finite_check_and_unscale_ with nan values."""
+    # PyTorch expects inv_scale and found_inf to be float32
+    inv_scale = torch.tensor(2.0, device=flag_gems.device, dtype=torch.float32)
+    found_inf = torch.tensor(0.0, device=flag_gems.device, dtype=torch.float32)
+
+    tensors = [
+        torch.tensor(
+            [1.0, 2.0, float("nan"), 4.0], device=flag_gems.device, dtype=dtype
+        ),
+        torch.tensor([5.0, 6.0, 7.0], device=flag_gems.device, dtype=dtype),
+    ]
+
+    # Reference
+    ref_tensors = [utils.to_reference(t.clone()) for t in tensors]
+    ref_found_inf = utils.to_reference(found_inf.clone())
+    ref_inv_scale = utils.to_reference(inv_scale.clone())
+    getattr(torch, "_amp_foreach_non_finite_check_and_unscale_")(
+        ref_tensors, ref_found_inf, ref_inv_scale
+    )
+
+    # GEMS
+    res_tensors = [t.clone() for t in tensors]
+    res_found_inf = found_inf.clone()
+    with flag_gems.use_gems():
+        getattr(torch, "_amp_foreach_non_finite_check_and_unscale_")(
+            res_tensors, res_found_inf, inv_scale
+        )
+
+    # Compare mutated inputs with equal_nan=True since tensors contain NaN
+    for i, (inp, ref_inp) in enumerate(zip(res_tensors, ref_tensors)):
+        utils.gems_assert_close(inp, ref_inp, dtype, equal_nan=True)
+
+    # Compare found_inf - should be 1.0 when nan is present
+    utils.gems_assert_equal(res_found_inf, ref_found_inf)