Add fake_quantize_per_channel_affine operator with Triton kernel

[xuanzhengdu-eng]

Summary

Add a Triton implementation of aten.fake_quantize_per_channel_affine for per-channel affine fake quantization.

The implementation adds the required ATen-compatible behavior:

• Apply an independent scale and zero_point to every channel along the selected axis.
• Compute quantized values with round-half-to-even semantics.
• Clamp quantized values to the requested [quant_min, quant_max] range before dequantization.
• Preserve the input dtype for FP16, FP32, and BF16 tensors.
• Support signed and unsigned quantization ranges.
• Handle arbitrary channel axes, nonzero zero points, non-contiguous inputs, and empty tensors.
• Register the operator in the ATen dispatch table.
• Add a corresponding entry to conf/operators.yaml.
• Add accuracy tests and a standard FlagGems benchmark.

The kernel was generated and iteratively corrected with KernelGen.

Testing

Tested on an NVIDIA H20-3e with PyTorch reference results.

Accuracy tests:

pytest tests/test_fake_quantize_per_channel_affine.py -q

Result:

54 passed

Coverage includes:

• FP16, FP32, and BF16 inputs
• 2D, 3D, and 4D input shapes
• Channel axes from the first through the last dimension
• Unsigned [0, 255] and signed [-128, 127] quantization ranges
• Random nonzero per-channel zero points
• Exact positive and negative half-integer cases for round-half-to-even behavior
• Empty tensors
• Output values and output dtype against the native PyTorch operator

Comprehensive benchmark:

pytest benchmark/test_fake_quantize_per_channel_affine.py \
  --level comprehensive \
  --mode kernel \
  --record json \
  -q --disable-warnings

Result:

1 passed

Pre-commit checks:

python -m pre_commit run --files \
  benchmark/test_fake_quantize_per_channel_affine.py \
  conf/operators.yaml \
  src/flag_gems/__init__.py \
  src/flag_gems/ops/__init__.py \
  src/flag_gems/ops/fake_quantize_per_channel_affine.py \
  tests/test_fake_quantize_per_channel_affine.py

Result:

Check Yaml................Passed
Fix End of Files..........Passed
Trim Trailing Whitespace..Passed
Flake8....................Passed
isort.....................Passed
black.....................Passed
black-jupyter.............Passed

git diff --check upstream/master...HEAD also passed.

Performance

Performance was measured on an NVIDIA H20-3e using the standard FlagGems benchmark framework in kernel mode. All 42 dtype/shape/axis cases completed successfully.

FP16

Shape	Axis	Torch latency (ms)	Gems latency (ms)	Speedup
[4, 4]	0	0.053888	0.005504	9.791x
[4, 4]	1	0.052416	0.005472	9.579x
[64, 64]	0	0.054144	0.005792	9.348x
[64, 64]	1	0.053760	0.005344	10.060x
[128, 256]	0	0.055152	0.005472	10.079x
[128, 256]	1	0.055056	0.005824	9.453x
[512, 512]	0	0.054752	0.006368	8.598x
[512, 512]	1	0.055344	0.006176	8.961x
[1024, 1024]	0	0.065408	0.008576	7.627x
[1024, 1024]	1	0.068000	0.008704	7.813x
[2, 3, 128, 128]	0	0.053792	0.006016	8.941x
[2, 3, 128, 128]	1	0.054336	0.006016	9.032x
[8, 16, 64, 64]	0	0.056000	0.006912	8.102x
[8, 16, 64, 64]	1	0.057760	0.006912	8.356x

Average FP16 speedup across these cases: 8.981x.

FP32

Shape	Axis	Torch latency (ms)	Gems latency (ms)	Speedup
[4, 4]	0	0.052544	0.005696	9.225x
[4, 4]	1	0.052416	0.005664	9.254x
[64, 64]	0	0.053536	0.006176	8.668x
[64, 64]	1	0.053280	0.006016	8.856x
[128, 256]	0	0.053760	0.006368	8.442x
[128, 256]	1	0.053504	0.006208	8.619x
[512, 512]	0	0.054080	0.006976	7.752x
[512, 512]	1	0.054432	0.006784	8.024x
[1024, 1024]	0	0.066912	0.009184	7.286x
[1024, 1024]	1	0.066752	0.008480	7.872x
[2, 3, 128, 128]	0	0.053568	0.006560	8.166x
[2, 3, 128, 128]	1	0.053792	0.006560	8.200x
[8, 16, 64, 64]	0	0.056352	0.007584	7.430x
[8, 16, 64, 64]	1	0.058848	0.007200	8.173x

Average FP32 speedup across these cases: 8.283x.

BF16

Shape	Axis	Torch latency (ms)	Gems latency (ms)	Speedup
[4, 4]	0	0.053488	0.005568	9.606x
[4, 4]	1	0.052576	0.005472	9.608x
[64, 64]	0	0.054400	0.005824	9.341x
[64, 64]	1	0.054816	0.005632	9.733x
[128, 256]	0	0.054688	0.005856	9.339x
[128, 256]	1	0.054304	0.005792	9.376x
[512, 512]	0	0.054976	0.006400	8.590x
[512, 512]	1	0.056000	0.006368	8.794x
[1024, 1024]	0	0.070720	0.008480	8.340x
[1024, 1024]	1	0.070592	0.008832	7.993x
[2, 3, 128, 128]	0	0.054720	0.006016	9.096x
[2, 3, 128, 128]	1	0.055072	0.006016	9.154x
[8, 16, 64, 64]	0	0.057232	0.006912	8.280x
[8, 16, 64, 64]	1	0.060832	0.006976	8.720x

Average BF16 speedup across these cases: 8.998x.

Combined arithmetic average speedup across all 42 cases: 8.754x. Every measured case is faster than the PyTorch reference, with speedups ranging from 7.286x to 10.079x.

Files changed

• src/flag_gems/ops/fake_quantize_per_channel_affine.py

  • Implements the per-channel affine fake-quantization Triton kernel.
  • Maps flattened offsets to the selected channel and loads the corresponding scale and zero point.
  • Implements round-half-to-even, clamping, and dequantization in FP32 while preserving the input/output dtype.

• tests/test_fake_quantize_per_channel_affine.py

  • Adds accuracy tests for supported floating-point dtypes, shapes, axes, quantization ranges, and nonzero zero points.
  • Adds explicit regression coverage for half-to-even rounding and empty tensors.

• benchmark/test_fake_quantize_per_channel_affine.py

  • Adds a standard FlagGems benchmark over multiple shapes, axes, and floating-point dtypes.

• src/flag_gems/ops/__init__.py

  • Exports the implementation.

• src/flag_gems/__init__.py

  • Registers fake_quantize_per_channel_affine in _FULL_CONFIG.

• conf/operators.yaml

  • Adds the operator entry with the KernelGen label and alpha: '5.4' stage.