[KernelGen][Nvidia] Add fake_quantize_per_channel_affine operator with Triton kernel (flagos-ai#5450)

* Add fake_quantize_per_channel_affine operator

- Implement Triton kernel for per-channel fake quantization
- Add accuracy tests with multiple shapes and axis configurations
- Add performance benchmark
- Register in operators.yaml

All tests passed (14/14).

* fix: initialize per-channel benchmark config

Signed-off-by: xuanzhengdu-eng <xuanzhengdu@gmail.com>

---------

Signed-off-by: xuanzhengdu-eng <xuanzhengdu@gmail.com>

Author: xuanzhengdu-eng

benchmark/test_fake_quantize_per_channel_affine.py

@@ -0,0 +1,70 @@
+# Copyright 2026 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.
+
+import pytest
+import torch
+
+from . import base, consts
+
+
+@pytest.mark.fake_quantize_per_channel_affine
+def test_fake_quantize_per_channel_affine():
+    class BenchmarkFakeQuantizePerChannelAffine(base.Benchmark):
+        """
+        Benchmark fake_quantize_per_channel_affine operator
+        """
+
+        axis_configs = (0, 1)
+        DEFAULT_SHAPES = [
+            (4, 4),
+            (64, 64),
+            (128, 256),
+            (512, 512),
+            (1024, 1024),
+            (2, 3, 128, 128),
+            (8, 16, 64, 64),
+        ]
+
+        def set_shapes(self, shape_file_path=None):
+            self.shapes = self.DEFAULT_SHAPES
+
+        def get_input_iter(self, dtype):
+            for shape in self.shapes:
+                for axis in self.axis_configs:
+                    if axis >= len(shape):
+                        continue
+                    inp = torch.randn(shape, dtype=dtype, device="cuda")
+                    n_channels = shape[axis]
+                    scale = (
+                        torch.rand(n_channels, dtype=torch.float32, device="cuda") * 0.1
+                        + 0.01
+                    )
+                    zero_point = torch.zeros(
+                        n_channels, dtype=torch.int32, device="cuda"
+                    )
+                    quant_min = 0
+                    quant_max = 255
+                    yield inp, scale, zero_point, axis, quant_min, quant_max
+
+        def forward(self, inp, scale, zero_point, axis, quant_min, quant_max):
+            return torch.fake_quantize_per_channel_affine(
+                inp, scale, zero_point, axis, quant_min, quant_max
+            )
+
+    bench = BenchmarkFakeQuantizePerChannelAffine(
+        op_name="fake_quantize_per_channel_affine",
+        torch_op=torch.fake_quantize_per_channel_affine,
+        dtypes=consts.FLOAT_DTYPES,
+    )
+    bench.run()

conf/operators.yaml

@@ -3437,6 +3437,18 @@ ops:
       - NeuralNetwork
     stages:
       - alpha: '5.3'
+  - id: fake_quantize_per_channel_affine
+    description: Applies fake quantization per channel with affine parameters (scale and zero_point).
+    for:
+      - fake_quantize_per_channel_affine
+    labels:
+      - aten
+      - pointwise
+      - KernelGen
+    kind:
+      - Quantization
+    stages:
+      - alpha: '5.4'
   - id: fill_scalar
     description: Fills a scalar with the specified value.
     for:

src/flag_gems/__init__.py

@@ -484,6 +484,7 @@ def torch_ge(v):
     ("exponential_", exponential_),
     ("eye", eye),
     ("eye.m", eye_m),
+    ("fake_quantize_per_channel_affine", fake_quantize_per_channel_affine),
     ("feature_dropout", feature_dropout),
     ("feature_dropout_", feature_dropout_),
     ("fill.Scalar", fill_scalar),

src/flag_gems/ops/__init__.py

@@ -300,6 +300,9 @@
 from flag_gems.ops.exponential_ import exponential_
 from flag_gems.ops.eye import eye
 from flag_gems.ops.eye_m import eye_m
+from flag_gems.ops.fake_quantize_per_channel_affine import (
+    fake_quantize_per_channel_affine,
+)
 from flag_gems.ops.feature_dropout import feature_dropout, feature_dropout_
 from flag_gems.ops.fft import fft
 from flag_gems.ops.fill import (
@@ -1140,6 +1143,7 @@
     "exponential_",
     "eye",
     "eye_m",
+    "fake_quantize_per_channel_affine",
     "feature_dropout",
     "feature_dropout_",
     "fft",

src/flag_gems/ops/fake_quantize_per_channel_affine.py

@@ -0,0 +1,113 @@
+# Copyright 2026 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.
+
+import logging
+
+import torch
+import triton
+import triton.language as tl
+
+from flag_gems.runtime import torch_device_fn
+
+logger = logging.getLogger(__name__)
+
+
+@triton.jit
+def _round_half_to_even(x):
+    floor_x = tl.floor(x)
+    fraction = x - floor_x
+    floor_is_even = (floor_x % 2.0) == 0.0
+    return tl.where(
+        fraction > 0.5,
+        floor_x + 1.0,
+        tl.where((fraction < 0.5) | floor_is_even, floor_x, floor_x + 1.0),
+    )
+
+
+@triton.jit
+def fake_quantize_per_channel_affine_kernel(
+    input_ptr,
+    scale_ptr,
+    zero_point_ptr,
+    output_ptr,
+    n_elements,
+    n_channels,
+    channel_stride,
+    quant_min,
+    quant_max,
+    BLOCK_SIZE: tl.constexpr,
+):
+    pid = tl.program_id(axis=0)
+    block_start = pid * BLOCK_SIZE
+    offsets = block_start + tl.arange(0, BLOCK_SIZE)
+    mask = offsets < n_elements
+
+    x = tl.load(input_ptr + offsets, mask=mask, other=0.0)
+    channel_idx = (offsets // channel_stride) % n_channels
+    scale = tl.load(scale_ptr + channel_idx, mask=mask, other=1.0)
+    zero_point = tl.load(zero_point_ptr + channel_idx, mask=mask, other=0.0)
+
+    x_fp32 = x.to(tl.float32)
+    scale_fp32 = scale.to(tl.float32)
+    zero_point_fp32 = zero_point.to(tl.float32)
+    x_quantized = _round_half_to_even(x_fp32 / scale_fp32) + zero_point_fp32
+    x_clamped = tl.minimum(tl.maximum(x_quantized, quant_min), quant_max)
+    output = (x_clamped - zero_point_fp32) * scale_fp32
+
+    tl.store(output_ptr + offsets, output, mask=mask)
+
+
+def fake_quantize_per_channel_affine(
+    input, scale, zero_point, axis, quant_min, quant_max
+):
+    logger.debug("GEMS FAKE_QUANTIZE_PER_CHANNEL_AFFINE")
+
+    if not isinstance(input, torch.Tensor):
+        raise TypeError("input must be a torch.Tensor")
+
+    input = input.contiguous()
+    scale = scale.contiguous()
+    zero_point = zero_point.contiguous()
+
+    n_elements = input.numel()
+    if n_elements == 0:
+        return torch.empty_like(input)
+
+    shape = input.shape
+    n_channels = shape[axis]
+
+    channel_stride = 1
+    for i in range(axis + 1, len(shape)):
+        channel_stride *= shape[i]
+
+    output = torch.empty_like(input)
+
+    BLOCK_SIZE = 1024
+    grid = lambda meta: (triton.cdiv(n_elements, meta["BLOCK_SIZE"]),)
+
+    with torch_device_fn.device(input.device):
+        fake_quantize_per_channel_affine_kernel[grid](
+            input,
+            scale,
+            zero_point,
+            output,
+            n_elements,
+            n_channels,
+            channel_stride,
+            quant_min,
+            quant_max,
+            BLOCK_SIZE=BLOCK_SIZE,
+        )
+
+    return output

tests/test_fake_quantize_per_channel_affine.py

@@ -0,0 +1,129 @@
+# Copyright 2026 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.
+
+import pytest
+import torch
+
+import flag_gems
+
+from .accuracy_utils import gems_assert_close, to_reference
+
+QUANT_SHAPES = [(4, 4), (16, 32), (2, 3, 4), (8, 16, 32)]
+
+
+@pytest.mark.fake_quantize_per_channel_affine
+@pytest.mark.parametrize("shape", QUANT_SHAPES)
+@pytest.mark.parametrize("axis", [0, 1])
+@pytest.mark.parametrize("dtype", [torch.float16, torch.float32, torch.bfloat16])
+@pytest.mark.parametrize("quant_min, quant_max", [(0, 255), (-128, 127)])
+def test_accuracy_fake_quantize_per_channel_affine(
+    shape, axis, dtype, quant_min, quant_max
+):
+    if axis >= len(shape):
+        pytest.skip(f"axis {axis} >= ndim {len(shape)}")
+
+    inp = torch.randn(shape, dtype=dtype, device=flag_gems.device)
+    n_channels = shape[axis]
+    scale = (
+        torch.rand(n_channels, dtype=torch.float32, device=flag_gems.device) * 0.1
+        + 0.01
+    )
+    zero_point = torch.randint(
+        quant_min,
+        quant_max + 1,
+        (n_channels,),
+        dtype=torch.int32,
+        device=flag_gems.device,
+    )
+
+    ref_inp = to_reference(inp)
+    ref_scale = to_reference(scale)
+    ref_zero_point = to_reference(zero_point)
+
+    ref_out = torch.fake_quantize_per_channel_affine(
+        ref_inp, ref_scale, ref_zero_point, axis, quant_min, quant_max
+    )
+
+    with flag_gems.use_gems():
+        res_out = torch.fake_quantize_per_channel_affine(
+            inp, scale, zero_point, axis, quant_min, quant_max
+        )
+
+    gems_assert_close(res_out, ref_out, dtype=dtype)
+
+
+@pytest.mark.fake_quantize_per_channel_affine
+@pytest.mark.parametrize("shape", [(2, 3, 4, 5)])
+@pytest.mark.parametrize("axis", [0, 1, 2, 3])
+def test_accuracy_fake_quantize_per_channel_affine_multi_dim(shape, axis):
+    inp = torch.randn(shape, dtype=torch.float32, device=flag_gems.device)
+    n_channels = shape[axis]
+    scale = (
+        torch.rand(n_channels, dtype=torch.float32, device=flag_gems.device) * 0.1
+        + 0.01
+    )
+    zero_point = torch.randint(
+        0, 255, (n_channels,), dtype=torch.int32, device=flag_gems.device
+    )
+
+    ref_inp = to_reference(inp)
+    ref_scale = to_reference(scale)
+    ref_zero_point = to_reference(zero_point)
+
+    ref_out = torch.fake_quantize_per_channel_affine(
+        ref_inp, ref_scale, ref_zero_point, axis, 0, 255
+    )
+
+    with flag_gems.use_gems():
+        res_out = torch.fake_quantize_per_channel_affine(
+            inp, scale, zero_point, axis, 0, 255
+        )
+
+    gems_assert_close(res_out, ref_out, dtype=torch.float32)
+
+
+@pytest.mark.fake_quantize_per_channel_affine
+def test_accuracy_fake_quantize_per_channel_affine_half_to_even():
+    inp = torch.tensor(
+        [[-3.5, -2.5, -1.5, -0.5, 0.5, 1.5, 2.5, 3.5]],
+        dtype=torch.float32,
+        device=flag_gems.device,
+    )
+    scale = torch.ones(8, dtype=torch.float32, device=flag_gems.device)
+    zero_point = torch.zeros(8, dtype=torch.int32, device=flag_gems.device)
+    ref_out = torch.fake_quantize_per_channel_affine(
+        to_reference(inp), to_reference(scale), to_reference(zero_point), 1, -128, 127
+    )
+
+    with flag_gems.use_gems():
+        res_out = torch.fake_quantize_per_channel_affine(
+            inp, scale, zero_point, 1, -128, 127
+        )
+
+    gems_assert_close(res_out, ref_out, dtype=torch.float32)
+
+
+@pytest.mark.fake_quantize_per_channel_affine
+def test_accuracy_fake_quantize_per_channel_affine_empty():
+    inp = torch.empty((2, 0, 3), dtype=torch.float32, device=flag_gems.device)
+    scale = torch.empty(0, dtype=torch.float32, device=flag_gems.device)
+    zero_point = torch.empty(0, dtype=torch.int32, device=flag_gems.device)
+
+    with flag_gems.use_gems():
+        result = torch.fake_quantize_per_channel_affine(
+            inp, scale, zero_point, 1, 0, 255
+        )
+
+    assert result.shape == inp.shape
+    assert result.dtype == inp.dtype