add addcdiv.out, addcmul.out, baddbmm.out, cat.out etc; refactor op name according to native_functions.yaml

Author: taozhiwei

benchmark/test_attention_perf.py

@@ -46,7 +46,7 @@ def torch_flash_attention_forward(
 def gems_flash_attention_forward(
     q, k, v, scale, is_causal, dropout_p=0.0, return_debug_mask=False, **extra_kwargs
 ):
-    return flag_gems.ops.flash_attention_forward(
+    return flag_gems.ops._flash_attention_forward(
         q,
         k,
         v,

benchmark/test_blas_perf.py

@@ -360,7 +360,7 @@ def test_blas_benchmark(op_name, torch_op, input_fn, bench_cls):
     )
 
     if op_name == "groupmm":
-        gems_op = flag_gems.group_mm
+        gems_op = flag_gems._grouped_mm
         bench.set_gems(gems_op)
 
     bench.run()

src/flag_gems/__init__.py

@@ -29,40 +29,48 @@ def torch_ge(v):
 
 _FULL_CONFIG = (
     ("_assert_async", _assert_async),
-    ("_flash_attention_forward", flash_attention_forward),
+    ("_flash_attention_forward", _flash_attention_forward),
     (
         "_functional_sym_constrain_range_for_size",
         _functional_sym_constrain_range_for_size,
     ),
-    ("_grouped_mm", group_mm),
-    ("_log_softmax", log_softmax),
-    ("_log_softmax_backward_data", log_softmax_backward),
+    ("_grouped_mm", _grouped_mm),
+    ("_log_softmax", _log_softmax),
+    ("_log_softmax.out", _log_softmax_out),
+    ("_log_softmax_backward_data", _log_softmax_backward_data),
+    ("_log_softmax_backward_data.out", _log_softmax_backward_data_out),
     ("_safe_softmax", _safe_softmax),
-    ("_softmax", softmax),
-    ("_softmax_backward_data", softmax_backward),
+    ("_softmax", _softmax),
+    ("_softmax.out", _softmax_out),
+    ("_softmax_backward_data", _softmax_backward_data),
+    ("_softmax_backward_data.out", _softmax_backward_data_out),
     (
         "_to_copy",
-        to_copy,
+        _to_copy,
         lambda: version.parse(torch.__version__) >= version.parse("2.4"),
     ),
     ("_unique2", _unique2),
     ("_upsample_bicubic2d_aa", _upsample_bicubic2d_aa),
     ("_upsample_bicubic2d_aa_backward", _upsample_bicubic2d_aa_backward),
     ("_upsample_nearest_exact1d", _upsample_nearest_exact1d),
-    ("_weight_norm_interface", weight_norm_interface),
-    ("_weight_norm_interface_backward", weight_norm_interface_backward),
+    ("_weight_norm_interface", _weight_norm_interface),
+    ("_weight_norm_interface_backward", _weight_norm_interface_backward),
     ("abs", abs),
     ("abs_", abs_),
     ("absolute", absolute),
     ("acos", acos),
     ("add.Tensor", add),
     ("add_.Tensor", add_),
     ("addcdiv", addcdiv),
+    ("addcdiv.out", addcdiv_out),
     ("addcmul", addcmul),
+    ("addcmul.out", addcmul_out),
     ("addmv", addmv),
     ("addmv.out", addmv_out),
     ("addmm", addmm),
     ("addmm.out", addmm_out),
+    ("addmm.dtype", addmm_dtype),
+    ("addmm.dtype_out", addmm_dtype_out),
     ("addr", addr),
     ("alias_copy", alias_copy),
     ("all", all),
@@ -92,6 +100,7 @@ def torch_ge(v):
     ("avg_pool2d", avg_pool2d),
     ("avg_pool2d_backward", avg_pool2d_backward),
     ("baddbmm", baddbmm),
+    ("baddbmm.out", baddbmm_out),
     ("bincount", bincount),
     ("bitwise_and.Scalar", bitwise_and_scalar),
     ("bitwise_and.Scalar_Tensor", bitwise_and_scalar_tensor),
@@ -110,6 +119,7 @@ def torch_ge(v):
     ("bmm", bmm),
     ("bmm.out", bmm_out),
     ("cat", cat),
+    ("cat.out", cat_out),
     ("celu", celu),
     ("celu_", celu_),
     ("ceil", ceil),
@@ -364,8 +374,8 @@ def torch_ge(v):
     ("resolve_neg", resolve_neg),
     ("rms_norm", rms_norm),
     ("round", round),
-    ("round.out", round_out),
     ("round_", round_),
+    ("round.out", round_out),
     ("rrelu_with_noise_backward", rrelu_with_noise_backward),
     ("rsqrt", rsqrt),
     ("rsqrt_", rsqrt_),
@@ -463,6 +473,21 @@ def torch_ge(v):
     func_name = fn.__name__ if hasattr(fn, "__name__") else str(fn)
     FULL_CONFIG_BY_FUNC.setdefault(func_name, []).append(_item)
 
+# Friendly names for only_enable(include=[...]) when the registered impl is *.out
+for _alias, _target in (
+    ("softmax", "_softmax_out"),
+    ("softmax_backward", "_softmax_backward_data_out"),
+    ("log_softmax", "_log_softmax_out"),
+    ("log_softmax_backward", "_log_softmax_backward_data_out"),
+    ("flash_attention_forward", "_flash_attention_forward"),
+    ("group_mm", "_grouped_mm"),
+    ("to_copy", "_to_copy"),
+    ("weight_norm_interface", "_weight_norm_interface"),
+    ("weight_norm_interface_backward", "_weight_norm_interface_backward"),
+):
+    if _target in FULL_CONFIG_BY_FUNC:
+        FULL_CONFIG_BY_FUNC.setdefault(_alias, []).extend(FULL_CONFIG_BY_FUNC[_target])
+
 
 def enable(
     lib=aten_lib,

src/flag_gems/fused/weight_norm.py

@@ -6,7 +6,7 @@
 import triton.language as tl
 
 from flag_gems import runtime
-from flag_gems.ops import weight_norm_interface, weight_norm_interface_backward
+from flag_gems.ops import _weight_norm_interface, _weight_norm_interface_backward
 from flag_gems.runtime import torch_device_fn
 from flag_gems.utils import libentry
 from flag_gems.utils import triton_lang_extension as tle
@@ -195,7 +195,7 @@ def forward(ctx, v, g, dim=0):
         dim = dim % v.ndim
         can_use_fused = dim == 0 or dim == v.ndim - 1
         if can_use_fused:
-            output, norm = weight_norm_interface(v, g, dim)
+            output, norm = _weight_norm_interface(v, g, dim)
         else:
             output, norm = weight_norm_except_dim(v, g, dim)
         ctx.save_for_backward(v, g, norm)
@@ -209,7 +209,7 @@ def backward(ctx, grad):
         v, g, norm = ctx.saved_tensors
         dim = ctx.dim
         if ctx.can_use_fused:
-            v_grad, g_grad = weight_norm_interface_backward(grad, v, g, norm, dim)
+            v_grad, g_grad = _weight_norm_interface_backward(grad, v, g, norm, dim)
         else:
             v_grad, g_grad = weight_norm_except_dim_backward(grad, v, g, norm, dim)
         return v_grad, g_grad, None

src/flag_gems/ops/__init__.py

@@ -7,9 +7,9 @@
 from flag_gems.ops.absolute import absolute
 from flag_gems.ops.acos import acos
 from flag_gems.ops.add import add, add_
-from flag_gems.ops.addcdiv import addcdiv
-from flag_gems.ops.addcmul import addcmul
-from flag_gems.ops.addmm import addmm, addmm_out
+from flag_gems.ops.addcdiv import addcdiv, addcdiv_out
+from flag_gems.ops.addcmul import addcmul, addcmul_out
+from flag_gems.ops.addmm import addmm, addmm_dtype, addmm_dtype_out, addmm_out
 from flag_gems.ops.addmv import addmv, addmv_out
 from flag_gems.ops.addr import addr
 from flag_gems.ops.alias_copy import alias_copy, alias_copy_out
@@ -30,15 +30,15 @@
 from flag_gems.ops.atan2 import atan2, atan2_out
 from flag_gems.ops.attention import (
     ScaleDotProductAttention,
-    flash_attention_forward,
+    _flash_attention_forward,
     flash_attn_varlen_func,
     flash_attn_varlen_opt_func,
     scaled_dot_product_attention,
     scaled_dot_product_attention_backward,
     scaled_dot_product_attention_forward,
 )
 from flag_gems.ops.avg_pool2d import avg_pool2d, avg_pool2d_backward
-from flag_gems.ops.baddbmm import baddbmm
+from flag_gems.ops.baddbmm import baddbmm, baddbmm_out
 from flag_gems.ops.batch_norm import batch_norm, batch_norm_backward
 from flag_gems.ops.bitwise_and import (
     bitwise_and_scalar,
@@ -58,7 +58,7 @@
 )
 from flag_gems.ops.bitwise_right_shift import bitwise_right_shift
 from flag_gems.ops.bmm import bmm, bmm_out
-from flag_gems.ops.cat import cat
+from flag_gems.ops.cat import cat, cat_out
 from flag_gems.ops.ceil import ceil, ceil_, ceil_out
 from flag_gems.ops.celu import celu, celu_
 from flag_gems.ops.clamp import (
@@ -135,7 +135,7 @@
     greater_scalar,
     greater_scalar_out,
 )
-from flag_gems.ops.group_gemm import group_mm
+from flag_gems.ops.group_gemm import _grouped_mm
 from flag_gems.ops.groupnorm import group_norm, group_norm_backward
 from flag_gems.ops.gt import gt, gt_scalar
 from flag_gems.ops.hardsigmoid import hardsigmoid, hardsigmoid_out
@@ -163,7 +163,12 @@
 from flag_gems.ops.log import log
 from flag_gems.ops.log1p_ import log1p_
 from flag_gems.ops.log_sigmoid import log_sigmoid
-from flag_gems.ops.log_softmax import log_softmax, log_softmax_backward
+from flag_gems.ops.log_softmax import (
+    _log_softmax,
+    _log_softmax_backward_data,
+    _log_softmax_backward_data_out,
+    _log_softmax_out,
+)
 from flag_gems.ops.logaddexp import logaddexp, logaddexp_out
 from flag_gems.ops.logical_and import logical_and, logical_and_
 from flag_gems.ops.logical_not import logical_not
@@ -269,7 +274,12 @@
 from flag_gems.ops.slice_backward import slice_backward
 from flag_gems.ops.slice_scatter import slice_scatter
 from flag_gems.ops.soft_margin_loss import soft_margin_loss, soft_margin_loss_out
-from flag_gems.ops.softmax import softmax, softmax_backward
+from flag_gems.ops.softmax import (
+    _softmax,
+    _softmax_backward_data,
+    _softmax_backward_data_out,
+    _softmax_out,
+)
 from flag_gems.ops.softplus import softplus
 from flag_gems.ops.softshrink import softshrink, softshrink_out
 from flag_gems.ops.sort import sort, sort_stable
@@ -286,7 +296,7 @@
 from flag_gems.ops.tanh import tanh, tanh_, tanh_backward
 from flag_gems.ops.threshold import threshold, threshold_backward
 from flag_gems.ops.tile import tile
-from flag_gems.ops.to import to_copy
+from flag_gems.ops.to import _to_copy
 from flag_gems.ops.topk import topk
 from flag_gems.ops.trace import trace
 from flag_gems.ops.tril import tril, tril_out
@@ -307,8 +317,8 @@
 from flag_gems.ops.vstack import vstack
 from flag_gems.ops.w8a8_block_fp8_matmul import w8a8_block_fp8_matmul
 from flag_gems.ops.weightnorm import (
-    weight_norm_interface,
-    weight_norm_interface_backward,
+    _weight_norm_interface,
+    _weight_norm_interface_backward,
 )
 from flag_gems.ops.where import (
     where_scalar_other,
@@ -336,8 +346,12 @@
     "add",
     "add_",
     "addcdiv",
+    "addcdiv_out",
     "addcmul",
+    "addcmul_out",
     "addmm",
+    "addmm_dtype",
+    "addmm_dtype_out",
     "addmm_out",
     "addmv",
     "addmv_out",
@@ -370,6 +384,7 @@
     "avg_pool2d",
     "avg_pool2d_backward",
     "baddbmm",
+    "baddbmm_out",
     "batch_norm",
     "batch_norm_backward",
     "bitwise_and_scalar",
@@ -389,6 +404,7 @@
     "bmm",
     "bmm_out",
     "cat",
+    "cat_out",
     "ceil",
     "ceil_",
     "ceil_out",
@@ -454,7 +470,7 @@
     "fill_tensor",
     "fill_tensor_",
     "fill_tensor_out",
-    "flash_attention_forward",
+    "_flash_attention_forward",
     "flash_attn_varlen_func",
     "flash_attn_varlen_opt_func",
     "flip",
@@ -480,7 +496,7 @@
     "greater_out",
     "greater_scalar",
     "greater_scalar_out",
-    "group_mm",
+    "_grouped_mm",
     "group_norm",
     "group_norm_backward",
     "gt",
@@ -521,8 +537,10 @@
     "linspace",
     "log",
     "log_sigmoid",
-    "log_softmax",
-    "log_softmax_backward",
+    "_log_softmax",
+    "_log_softmax_backward_data",
+    "_log_softmax_backward_data_out",
+    "_log_softmax_out",
     "log1p_",
     "logaddexp",
     "logaddexp_out",
@@ -658,8 +676,10 @@
     "slice_scatter",
     "soft_margin_loss",
     "soft_margin_loss_out",
-    "softmax",
-    "softmax_backward",
+    "_softmax",
+    "_softmax_backward_data",
+    "_softmax_backward_data_out",
+    "_softmax_out",
     "softplus",
     "softshrink",
     "softshrink_out",
@@ -694,7 +714,7 @@
     "threshold",
     "threshold_backward",
     "tile",
-    "to_copy",
+    "_to_copy",
     "topk",
     "trace",
     "tril",
@@ -716,8 +736,8 @@
     "vector_norm",
     "vstack",
     "w8a8_block_fp8_matmul",
-    "weight_norm_interface",
-    "weight_norm_interface_backward",
+    "_weight_norm_interface",
+    "_weight_norm_interface_backward",
     "where_scalar_other",
     "where_scalar_self",
     "where_self",

src/flag_gems/ops/addcdiv.py

@@ -16,12 +16,14 @@ def addcdiv_kernel(x, t1, t2, value):
     return x + value * (t1 / t2)
 
 
-def addcdiv(inp, tensor1, tensor2, value=1.0, out=None):
-    logger.debug("GEMS ADDCDIV FORWARD")
-
-    if out is None:
-        out = torch.empty_like(inp)
-
+def addcdiv_out(inp, tensor1, tensor2, *, value=1.0, out):
+    logger.debug("GEMS ADDCDIV_OUT")
     addcdiv_kernel(inp, tensor1, tensor2, value, out0=out)
-
     return out
+
+
+def addcdiv(inp, tensor1, tensor2, value=1.0):
+    """Functional entry; CUDA may dispatch here without hitting ``addcdiv.out``."""
+    logger.debug("GEMS ADDCDIV")
+    out = torch.empty_like(inp)
+    return addcdiv_out(inp, tensor1, tensor2, value=value, out=out)

src/flag_gems/ops/addcmul.py

@@ -16,15 +16,21 @@ def addcmul_forward(x, t1, t2, value):
     return x + value * t1 * t2
 
 
-def addcmul(inp, tensor1, tensor2, *, value=1.0, out=None):
-    logger.debug("GEMS ADDCMUL FORWARD")
-    if out is not None:
-        broadcast_shape = torch.broadcast_shapes(
-            inp.shape, tensor1.shape, tensor2.shape
-        )
-        if list(out.shape) != list(broadcast_shape):
-            out.resize_(broadcast_shape)
-        addcmul_forward(inp, tensor1, tensor2, value, out0=out)
-        return out
-    else:
-        return addcmul_forward(inp, tensor1, tensor2, value)
+def addcmul_out(inp, tensor1, tensor2, *, value=1.0, out):
+    logger.debug("GEMS ADDCMUL_OUT")
+    broadcast_shape = torch.broadcast_shapes(inp.shape, tensor1.shape, tensor2.shape)
+    if list(out.shape) != list(broadcast_shape):
+        out.resize_(broadcast_shape)
+    addcmul_forward(inp, tensor1, tensor2, value, out0=out)
+    return out
+
+
+def addcmul(inp, tensor1, tensor2, *, value=1.0):
+    """Functional entry; keep alongside ``addcmul.out`` for dispatch coverage."""
+    logger.debug("GEMS ADDCMUL")
+    broadcast_shape = torch.broadcast_shapes(inp.shape, tensor1.shape, tensor2.shape)
+    dtype = torch.promote_types(
+        inp.dtype, torch.promote_types(tensor1.dtype, tensor2.dtype)
+    )
+    out = torch.empty(broadcast_shape, device=inp.device, dtype=dtype)
+    return addcmul_out(inp, tensor1, tensor2, value=value, out=out)