format docment

Author: YY4994

src/flag_gems/logging_utils.py

@@ -11,12 +11,12 @@
 """
 
 import logging
-import functools
 import traceback
 from pathlib import Path
 
 import torch
 
+
 class LogOncePerLocationFilter(logging.Filter):
     def __init__(self):
         super().__init__()
@@ -127,7 +127,11 @@ def get_call_location():
 def compare_outputs(fg_out, pt_out, rtol, atol):
     if isinstance(fg_out, torch.Tensor) and isinstance(pt_out, torch.Tensor):
         if fg_out.shape != pt_out.shape:
-            return False, {"error": "shape_mismatch", "fg": tuple(fg_out.shape), "pt": tuple(pt_out.shape)}
+            return False, {
+                "error": "shape_mismatch",
+                "fg": tuple(fg_out.shape),
+                "pt": tuple(pt_out.shape),
+            }
         try:
             fg = fg_out.detach().float()
             pt = pt_out.detach().float()
@@ -146,18 +150,23 @@ def compare_outputs(fg_out, pt_out, rtol, atol):
                 return False, info
     return True, {}
 
-def enable_precision_check(rtol=1e-4, atol=1e-5, log_once=True, max_checks=10, path=None):
+
+def enable_precision_check(
+    rtol=1e-4, atol=1e-5, log_once=True, max_checks=10, path=None
+):
     setup_precision_logging(path)
-    precision_config.update({
-        "enabled": True,
-        "rtol": rtol,
-        "atol": atol,
-        "log_once": log_once,
-        "max_checks": max_checks,
-        "logged_ops": set(),
-    })
+    precision_config.update(
+        {
+            "enabled": True,
+            "rtol": rtol,
+            "atol": atol,
+            "log_once": log_once,
+            "max_checks": max_checks,
+            "logged_ops": set(),
+        }
+    )
 
 
 def disable_precision_check():
     """close precision log file and disable precision check."""
-    precision_config["enabled"] = False
+    precision_config["enabled"] = False

src/flag_gems/runtime/register.py

@@ -1,16 +1,16 @@
-import warnings
 import functools
+import warnings
+
 import torch
 
-from . import backend, common, error
-from .backend.device import DeviceDetector
 from ..logging_utils import (
-    precision_config,
+    compare_outputs,
     get_precision_logger,
     get_tensor_info,
-    get_call_location,
-    compare_outputs,
+    precision_config,
 )
+from . import backend, common, error
+from .backend.device import DeviceDetector
 
 
 class Register:
@@ -34,7 +34,6 @@ def __init__(
         self.all_ops = []
         self.all_keys = []
 
-
         # optional mapping func_name -> list of config entries
         self.full_config_by_func = full_config_by_func
         self.cpp_patched_ops = set(cpp_patched_ops or [])
@@ -144,27 +143,13 @@ def get_vendor_name(self):
     def get_current_device(self):
         return self.device.name
 
-# def _wrap_op_with_precision_check(op_key, fn):
-#     @functools.wraps(fn)
-#     def wrapper(*args, **kwargs):
-#         fg_result = fn(*args, **kwargs)
-        
-#         cfg = precision_config
-#         if cfg["enabled"] and op_key not in cfg["logged_ops"]:
-#             cfg["logged_ops"].add(op_key)
-#             logger = get_precision_logger()
-#             input_info = [get_tensor_info(a) for a in args if get_tensor_info(a)]
-#             logger.info(f"Op: {op_key} | in: {input_info}")
-        
-#         return fg_result
-#     return wrapper
-
-# 精度检查时允许的最大 tensor 元素数（超过则跳过，避免大 tensor 拷贝开销）
-_MAX_NUMEL_FOR_CHECK = 1 * 1024 * 1024  # 1M 元素
+
+# Maximum tensor element count allowed for precision check (skip if exceeded to avoid large tensor copy overhead)
+_MAX_NUMEL_FOR_CHECK = 1 * 1024 * 1024  # 1M elements
 
 
 def _get_dtype_tolerance(args, default_rtol, default_atol):
-    """根据输入 tensor 的 dtype 自动调整容差"""
+    """Automatically adjust tolerance based on the dtype of input tensors"""
     for a in args:
         if isinstance(a, torch.Tensor) and a.is_floating_point():
             if a.dtype in (torch.bfloat16, torch.float16):
@@ -174,7 +159,7 @@ def _get_dtype_tolerance(args, default_rtol, default_atol):
 
 
 def _to_cpu(x):
-    """递归地将 tensor 移到 CPU"""
+    """Recursively move tensors to CPU"""
     if isinstance(x, torch.Tensor):
         return x.detach().cpu()
     elif isinstance(x, (list, tuple)):
@@ -185,7 +170,7 @@ def _to_cpu(x):
 
 
 def _max_tensor_numel(args):
-    """返回参数中最大 tensor 的元素数"""
+    """Return the element count of the largest tensor in the arguments"""
     max_n = 0
     for a in args:
         if isinstance(a, torch.Tensor):
@@ -194,65 +179,88 @@ def _max_tensor_numel(args):
 
 
 def _wrap_op_with_precision_check(op_key, fn):
-    """包装 FlagGems 算子，在执行后与 PyTorch 原生 CPU 实现做精度对比。
+    """Wrap a FlagGems operator to compare its output against the native PyTorch CPU implementation after execution.
 
-    由于 FlagGems 替换了 CUDA dispatch，无法在 GPU 上调用原生实现，
-    因此将输入拷贝到 CPU 计算参考结果。通过以下方式控制性能开销：
-    - max_checks: 每个算子只检查前 N 次调用（默认 10）
-    - 跳过大 tensor（超过 1M 元素）
-    - 一旦记录过失败就不再检查该算子
+    Since FlagGems replaces the CUDA dispatch, the native implementation cannot be called on GPU,
+    so inputs are copied to CPU to compute the reference result. Performance overhead is controlled by:
+    - max_checks: only check the first N calls per operator (default 10)
+    - skip large tensors (over 1M elements)
+    - once a failure is logged, that operator is no longer checked
     """
     _call_count = 0
 
     @functools.wraps(fn)
     def wrapper(*args, **kwargs):
         nonlocal _call_count
 
-        # 先执行 FlagGems 实现
+        # Execute the FlagGems implementation first
         fg_result = fn(*args, **kwargs)
 
         cfg = precision_config
         if not cfg["enabled"]:
             return fg_result
 
-        # 已经记录过失败的算子不再检查
+        # Skip operators that have already logged a failure
         if op_key in cfg["logged_ops"]:
             return fg_result
 
-        # 采样：每个算子只检查前 N 次调用
+        # Sampling: only check the first N calls per operator
         _call_count += 1
         if _call_count > cfg.get("max_checks", 10):
             return fg_result
 
-        op_name = op_key.split("::")[-1].split(".")[0] if "::" in op_key else op_key.split(".")[0]
+        op_name = (
+            op_key.split("::")[-1].split(".")[0]
+            if "::" in op_key
+            else op_key.split(".")[0]
+        )
 
-        # 跳过 out variant（参数签名不同，CPU 调用容易出错）
+        # Skip out variants (different argument signatures, prone to errors when calling on CPU)
         overload_part = op_key.split(".")[-1] if "." in op_key else ""
         if overload_part == "out" or op_name.endswith("_out"):
             return fg_result
 
-        # 跳过不需要检查的算子
+        # Skip operators that do not need checking
         skip_ops = {
-            # 纯 layout / 内存操作
-            'copy_', '_to_copy', 'view', 'reshape', 'expand', 'permute',
-            'transpose', 'contiguous', 'clone', 'to', 'empty', 'zeros',
-            'ones', 'full', 'masked_fill_',
-            # 随机采样算子（GPU/CPU 随机数生成器不同，结果必然不一致）
-            'exponential_', 'normal_', 'uniform_', 'bernoulli_', 'random_',
-            'multinomial', 'randperm',
-            # 排序/选择算子（相同值的排序顺序可能不同，不适合逐元素比较）
-            'sort', 'topk', 'argsort',
+            # Pure layout / memory operations
+            "copy_",
+            "_to_copy",
+            "view",
+            "reshape",
+            "expand",
+            "permute",
+            "transpose",
+            "contiguous",
+            "clone",
+            "to",
+            "empty",
+            "zeros",
+            "ones",
+            "full",
+            "masked_fill_",
+            # Random sampling operators (GPU/CPU RNGs differ, results will inevitably mismatch)
+            "exponential_",
+            "normal_",
+            "uniform_",
+            "bernoulli_",
+            "random_",
+            "multinomial",
+            "randperm",
+            # Sorting/selection operators (order of equal values may differ, not suitable for element-wise comparison)
+            "sort",
+            "topk",
+            "argsort",
         }
         if op_name in skip_ops:
             return fg_result
 
-        # 跳过大 tensor，避免拷贝开销
+        # Skip large tensors to avoid copy overhead
         if _max_tensor_numel(args) > _MAX_NUMEL_FOR_CHECK:
             return fg_result
 
         try:
-            # 解析 op_key 获取正确的 overload
-            # op_key 形如 "add.Tensor", "mm.default", "softmax.int" 等
+            # Parse op_key to get the correct overload
+            # op_key is in the form "add.Tensor", "mm.default", "softmax.int", etc.
             parts = op_key.split(".")
             base_name = parts[0]
             overload_name = parts[1] if len(parts) > 1 else "default"
@@ -264,17 +272,17 @@ def wrapper(*args, **kwargs):
             if aten_overload is None:
                 return fg_result
 
-            # 将输入拷贝到 CPU，调用原生 aten 实现（CPU 上不受 FlagGems 影响）
+            # Copy inputs to CPU and call the native aten implementation (unaffected by FlagGems on CPU)
             cpu_args = [_to_cpu(a) for a in args]
             cpu_kwargs = {k: _to_cpu(v) for k, v in kwargs.items()}
 
             with torch.no_grad():
                 pt_result_cpu = aten_overload(*cpu_args, **cpu_kwargs)
 
-            # 将 FlagGems 结果也拷贝到 CPU 做比较（避免把 CPU 结果搬回 GPU）
+            # Also copy FlagGems result to CPU for comparison (avoid moving CPU result back to GPU)
             fg_result_cpu = _to_cpu(fg_result)
 
-            # 根据 dtype 自动调整容差
+            # Automatically adjust tolerance based on dtype
             rtol, atol = _get_dtype_tolerance(args, cfg["rtol"], cfg["atol"])
 
             is_close, info = compare_outputs(fg_result_cpu, pt_result_cpu, rtol, atol)
@@ -300,6 +308,7 @@ def wrapper(*args, **kwargs):
 
     return wrapper
 
+
 class PrecisionCheckRegister(Register):
     def register_impl(self, key, fn):
         if self.lib is None: