[KernelGen][Nvidia] Add cosine_embedding_loss operator with Triton kernel

Author: ShawnsYing

.pr_gate/cosine_embedding_loss.passed

@@ -0,0 +1,2 @@
+check=pass test=pass benchmark=pass
+time=2026-08-17T10:43:35.797838

benchmark/test_cosine_embedding_loss.py

@@ -0,0 +1,39 @@
+# 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
+
+
+def _input_fn(shape, dtype, device):
+    inp1 = torch.randn(shape, dtype=dtype, device=device)
+    inp2 = torch.randn(shape, dtype=dtype, device=device)
+    target = (torch.randint(0, 2, (shape[0],), device=device).to(dtype) * 2) - 1
+    # aten::cosine_embedding_loss(input1, input2, target, margin, reduction)
+    yield inp1, inp2, target, 0.0, 1
+
+
+@pytest.mark.cosine_embedding_loss
+def test_cosine_embedding_loss():
+    # 2D-only: cosine_embedding_loss needs (N, D) inputs with a 1D (N,) target.
+    bench = base.GenericBenchmark2DOnly(
+        input_fn=_input_fn,
+        op_name="cosine_embedding_loss",
+        torch_op=torch.ops.aten.cosine_embedding_loss,
+        dtypes=consts.FLOAT_DTYPES,
+    )
+
+    bench.run()

conf/operators.yaml

@@ -2340,6 +2340,18 @@ ops:
       - Math
     stages:
       - stable: '5.3'
+  - id: cosine_embedding_loss
+    description: Compute the cosine embedding loss.
+    for:
+      - cosine_embedding_loss
+    labels:
+      - aten
+      - nn.functional
+      - KernelGen
+    kind:
+      - NeuralNetwork
+    stages:
+      - alpha: '5.4'
   - id: count_nonzero
     description: |
       Counts the number of non-zero values in the tensor `input` along the given `dim`.

src/flag_gems/__init__.py

@@ -393,6 +393,7 @@ def torch_ge(v):
     ("cosh", cosh),
     ("cosh.out", cosh_out),
     ("cosh_", cosh_),
+    ("cosine_embedding_loss", cosine_embedding_loss),
     ("count_nonzero", count_nonzero),
     ("ctc_loss.IntList", ctc_loss, None, (AUTOGRAD_DISPATCH_KEY,)),
     ("ctc_loss.Tensor", ctc_loss, None, (AUTOGRAD_DISPATCH_KEY,)),

src/flag_gems/ops/__init__.py

@@ -238,6 +238,7 @@
 from flag_gems.ops.copysign_ import copysign_
 from flag_gems.ops.cos import cos, cos_
 from flag_gems.ops.cosh import cosh, cosh_, cosh_out
+from flag_gems.ops.cosine_embedding_loss import cosine_embedding_loss
 from flag_gems.ops.count_nonzero import count_nonzero
 from flag_gems.ops.ctc_loss import ctc_loss
 from flag_gems.ops.cudnn_attention_forward import cudnn_attention_forward
@@ -1041,6 +1042,7 @@
     "cosh",
     "cosh_",
     "cosh_out",
+    "cosine_embedding_loss",
     "count_nonzero",
     "ctc_loss",
     "cudnn_attention_forward",

src/flag_gems/ops/cosine_embedding_loss.py

@@ -0,0 +1,250 @@
+# 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
+
+import torch
+import triton
+import triton.language as tl
+
+import flag_gems
+
+logger = logging.getLogger(__name__)
+
+# The denominator is guarded against zero-norm inputs with a literal 1e-8 inside
+# each kernel (Triton @jit bodies cannot read module-level Python globals). This
+# matches the PyTorch reference on the tested float16/bfloat16/float32 workloads.
+
+
+@triton.jit
+def _cosine_embedding_loss_elementwise_kernel(
+    x1_ptr,
+    x2_ptr,
+    y_ptr,
+    out_ptr,
+    n_rows,
+    d,
+    x1_stride0,
+    x2_stride0,
+    margin,
+    BLOCK_D: tl.constexpr,
+):
+    # One program per row: compute cosine similarity along the last dim, then
+    # the per-sample loss, and write it to out[row].
+    row = tl.program_id(axis=0)
+    if row >= n_rows:
+        return
+
+    cols = tl.arange(0, BLOCK_D)
+    mask = cols < d
+
+    x1 = tl.load(x1_ptr + row * x1_stride0 + cols, mask=mask, other=0.0).to(tl.float32)
+    x2 = tl.load(x2_ptr + row * x2_stride0 + cols, mask=mask, other=0.0).to(tl.float32)
+
+    dot = tl.sum(x1 * x2, axis=0)
+    norm1 = tl.sqrt(tl.sum(x1 * x1, axis=0))
+    norm2 = tl.sqrt(tl.sum(x2 * x2, axis=0))
+    denom = norm1 * norm2
+    cos = dot / (denom + 1e-8)
+
+    y = tl.load(y_ptr + row).to(tl.float32)
+    # target == 1  -> 1 - cos ; target == -1 -> max(0, cos - margin)
+    is_pos = y > 0.0
+    loss = tl.where(is_pos, 1.0 - cos, tl.maximum(cos - margin, 0.0))
+
+    tl.store(out_ptr + row, loss)
+
+
+@triton.jit
+def _cosine_embedding_loss_per_row_kernel(
+    x1_ptr,
+    x2_ptr,
+    y_ptr,
+    mid_ptr,
+    n_rows,
+    d,
+    x1_stride0,
+    x2_stride0,
+    margin,
+    BLOCK_D: tl.constexpr,
+):
+    # One program per sample-row: compute the per-row loss in float32 and write
+    # it to mid[row]. A 1D (BLOCK_D,) tile keeps every load well under Triton's
+    # per-tile numel limit regardless of the row count. The final reduction over
+    # mid is done by _cosine_embedding_loss_final_reduce_kernel.
+    row = tl.program_id(axis=0)
+    if row >= n_rows:
+        return
+    cols = tl.arange(0, BLOCK_D)
+    col_mask = cols < d
+
+    x1 = tl.load(x1_ptr + row * x1_stride0 + cols, mask=col_mask, other=0.0).to(
+        tl.float32
+    )
+    x2 = tl.load(x2_ptr + row * x2_stride0 + cols, mask=col_mask, other=0.0).to(
+        tl.float32
+    )
+    dot = tl.sum(x1 * x2, axis=0)
+    norm1 = tl.sqrt(tl.sum(x1 * x1, axis=0))
+    norm2 = tl.sqrt(tl.sum(x2 * x2, axis=0))
+    cos = dot / (norm1 * norm2 + 1e-8)
+
+    y = tl.load(y_ptr + row).to(tl.float32)
+    is_pos = y > 0.0
+    loss = tl.where(is_pos, 1.0 - cos, tl.maximum(cos - margin, 0.0))
+    tl.store(mid_ptr + row, loss)
+
+
+@triton.jit
+def _cosine_embedding_loss_final_reduce_kernel(
+    mid_ptr, out_ptr, mid_size, inv_scale, BLOCK_MID: tl.constexpr
+):
+    # Grid-stride reduction so a large partial-sum buffer does not require a
+    # single oversized block. Accumulate in float32. inv_scale folds the mean
+    # division into the kernel (inv_scale=1/n_rows for mean, 1.0 for sum) so the
+    # host never issues a tensor/scalar div that use_gems would re-dispatch.
+    offsets = tl.arange(0, BLOCK_MID)
+    acc = 0.0
+    for start in range(0, mid_size, BLOCK_MID):
+        idx = start + offsets
+        mask = idx < mid_size
+        vals = tl.load(mid_ptr + idx, mask=mask, other=0.0).to(tl.float32)
+        acc += tl.sum(vals, axis=0)
+    # Store into out_ptr's own dtype so the host never issues a `.to()` cast
+    # (which use_gems would re-dispatch to a pointwise _to_copy kernel).
+    tl.store(out_ptr, (acc * inv_scale).to(out_ptr.dtype.element_ty))
+
+
+def _normalize_reduction(reduction):
+    # Accept both string and enum/int forms: 0=none,1=mean,2=sum
+    if isinstance(reduction, str):
+        r = reduction.lower()
+        if r == "none":
+            return 0
+        if r == "mean":
+            return 1
+        if r == "sum":
+            return 2
+        raise ValueError(f"Invalid reduction: {reduction}")
+    if isinstance(reduction, int):
+        if reduction in (0, 1, 2):
+            return reduction
+        raise ValueError(f"Invalid reduction int: {reduction}")
+    raise ValueError(f"Unsupported reduction type: {type(reduction)}")
+
+
+def _check_tensors(input1: torch.Tensor, input2: torch.Tensor, target: torch.Tensor):
+    if input1.device.type != flag_gems.device or input2.device.type != flag_gems.device:
+        raise AssertionError(
+            f"cosine_embedding_loss: inputs must be {flag_gems.device} tensors for Triton kernel."
+        )
+    if input1.device != input2.device or input1.device != target.device:
+        raise AssertionError(
+            "cosine_embedding_loss: all inputs must be on the same device."
+        )
+    if input1.shape != input2.shape:
+        raise AssertionError(
+            "cosine_embedding_loss: input1 and input2 must have the same shape."
+        )
+    if input1.dim() < 1:
+        raise AssertionError(
+            "cosine_embedding_loss: inputs must have at least one dimension."
+        )
+    # target is (N,) for batched inputs (input.ndim >= 2); for 1-D inputs target is ().
+    expected_target_shape = input1.shape[:-1] if input1.dim() >= 2 else ()
+    if target.shape != expected_target_shape:
+        raise AssertionError(
+            f"cosine_embedding_loss: target must have shape {expected_target_shape}, "
+            f"got {tuple(target.shape)}."
+        )
+    if not input1.is_contiguous():
+        input1 = input1.contiguous()
+    if not input2.is_contiguous():
+        input2 = input2.contiguous()
+    if not target.is_contiguous():
+        target = target.contiguous()
+    return input1, input2, target
+
+
+def cosine_embedding_loss(
+    input1: torch.Tensor,
+    input2: torch.Tensor,
+    target: torch.Tensor,
+    margin: float = 0.0,
+    reduction=1,
+):
+    logger.debug("GEMS COSINE_EMBEDDING_LOSS")
+    input1, input2, target = _check_tensors(input1, input2, target)
+    red = _normalize_reduction(reduction)
+    margin = float(margin)
+
+    # Flatten leading dims so the last dim is D and the row count is N.
+    x1 = input1.reshape(-1, input1.shape[-1])
+    x2 = input2.reshape(-1, input2.shape[-1])
+    y = target.reshape(-1) if target.numel() > 0 else target
+    n_rows = x1.shape[0]
+    d = x1.shape[1]
+    x1_stride0 = x1.stride(0)
+    x2_stride0 = x2.stride(0)
+
+    if red == 0:
+        # reduction = 'none'
+        out = torch.empty((n_rows,), device=input1.device, dtype=input1.dtype)
+        if n_rows == 0:
+            return out
+        BLOCK_D = triton.next_power_of_2(d)
+        _cosine_embedding_loss_elementwise_kernel[(n_rows,)](
+            x1, x2, y, out, n_rows, d, x1_stride0, x2_stride0, margin, BLOCK_D=BLOCK_D
+        )
+        # Restore the leading shape (target shape).
+        return out.reshape(target.shape) if target.numel() != n_rows else out
+    else:
+        # reduction = 'mean' (1) or 'sum' (2)
+        if n_rows == 0:
+            # Follow PyTorch behavior: sum -> 0, mean -> NaN
+            if red == 2:
+                return torch.zeros((), device=input1.device, dtype=input1.dtype)
+            else:
+                return torch.full(
+                    (), float("nan"), device=input1.device, dtype=input1.dtype
+                )
+        BLOCK_D = triton.next_power_of_2(d)
+        # Stage 1: one program per row writes its float32 loss into mid.
+        mid_size = n_rows
+        block_mid = min(1024, triton.next_power_of_2(mid_size))
+        mid = torch.empty((mid_size,), device=input1.device, dtype=torch.float32)
+        _cosine_embedding_loss_per_row_kernel[(n_rows,)](
+            x1,
+            x2,
+            y,
+            mid,
+            n_rows,
+            d,
+            x1_stride0,
+            x2_stride0,
+            margin,
+            BLOCK_D=BLOCK_D,
+        )
+        # Allocate the output directly in the input dtype and let the kernel
+        # cast on store. This avoids any host-side arithmetic / `.to()` that
+        # use_gems would re-dispatch to pointwise kernels.
+        out = torch.zeros((), device=input1.device, dtype=input1.dtype)
+        # inv_scale=1/n_rows folds the mean division into the reduction kernel;
+        # for sum it is 1.0.
+        inv_scale = 1.0 if red == 2 else (1.0 / float(n_rows))
+        _cosine_embedding_loss_final_reduce_kernel[(1,)](
+            mid, out, mid_size, inv_scale, BLOCK_MID=block_mid
+        )
+        return out

tests/test_cosine_embedding_loss.py

@@ -0,0 +1,60 @@
+# 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 . import accuracy_utils as utils
+from . import conftest as cfg
+
+# Shapes mirror the sibling hinge_embedding_loss test to exercise the two-stage
+# reduction across block counts. Inputs are (N, D); target is (N,).
+if cfg.QUICK_MODE:
+    # Single tiny shape keeps QUICK_MODE smoke runs fast.
+    COSINE_EMBEDDING_LOSS_SHAPES = [(2, 3)]
+else:
+    # Small/medium/large trio covers single- and multi-block reductions.
+    COSINE_EMBEDDING_LOSS_SHAPES = [(2, 3), (128, 256), (512, 512)]
+
+
+@pytest.mark.cosine_embedding_loss
+@pytest.mark.parametrize("shape", COSINE_EMBEDDING_LOSS_SHAPES)
+@pytest.mark.parametrize("dtype", utils.FLOAT_DTYPES)
+@pytest.mark.parametrize("reduction", [0, 1, 2])
+@pytest.mark.parametrize("margin", [0.0, 0.5])
+def test_cosine_embedding_loss(shape, dtype, reduction, margin):
+    inp1 = torch.randn(shape, dtype=dtype, device=flag_gems.device)
+    inp2 = torch.randn(shape, dtype=dtype, device=flag_gems.device)
+    target = (
+        torch.randint(0, 2, (shape[0],), device=flag_gems.device).to(dtype) * 2
+    ) - 1
+
+    # The kernel upcasts to float32 for the dot/norm reductions, so the
+    # reference must upcast too — otherwise a pure-fp16 cosine is numerically
+    # ill-conditioned near cos==0 and diverges from the (more accurate) result.
+    ref_inp1 = utils.to_reference(inp1, upcast=True)
+    ref_inp2 = utils.to_reference(inp2, upcast=True)
+    ref_target = utils.to_reference(target, upcast=True)
+    ref_out = torch.ops.aten.cosine_embedding_loss(
+        ref_inp1, ref_inp2, ref_target, margin, reduction
+    )
+
+    with flag_gems.use_gems():
+        res_out = torch.ops.aten.cosine_embedding_loss(
+            inp1, inp2, target, margin, reduction
+        )
+
+    utils.gems_assert_close(res_out, ref_out, dtype)