```python
# Copyright (c) Meta Platforms, Inc. and affiliates.
# All rights reserved.
#
# This source code is licensed under the license found in the
# LICENSE file in the root directory of this source tree.

import torch
import torch.nn as nn
import torch.nn.functional as F
from torch.utils.cpp_extension import load_inline

# Fused Triplet Margin Loss CUDA kernel.
# Computes, in a single pass over memory:
#   dist_pos = ||a - p + eps||_2   (per row)
#   dist_neg = ||a - n + eps||_2   (per row)
#   loss     = mean( max(dist_pos - dist_neg + margin, 0) )
# This avoids materializing intermediate difference tensors and reads `anchor`
# only once, making the op purely memory-bandwidth bound at minimal traffic.
triplet_margin_loss_cuda_source = r"""
#include <torch/extension.h>
#include <cuda_runtime.h>

__global__ void triplet_margin_loss_vec4_kernel(
    const float4* __restrict__ anchor,
    const float4* __restrict__ positive,
    const float4* __restrict__ negative,
    float* __restrict__ out,
    long row_len4,
    float margin,
    float eps
) {
    long row = blockIdx.x;
    const float4* a_row = anchor + row * row_len4;
    const float4* p_row = positive + row * row_len4;
    const float4* n_row = negative + row * row_len4;

    float sp = 0.0f;
    float sn = 0.0f;

    for (long j = threadIdx.x; j < row_len4; j += blockDim.x) {
        float4 av = a_row[j];
        float4 pv = p_row[j];
        float4 nv = n_row[j];

        float d0 = av.x - pv.x + eps;
        float d1 = av.y - pv.y + eps;
        float d2 = av.z - pv.z + eps;
        float d3 = av.w - pv.w + eps;
        sp += d0 * d0 + d1 * d1 + d2 * d2 + d3 * d3;

        float e0 = av.x - nv.x + eps;
        float e1 = av.y - nv.y + eps;
        float e2 = av.z - nv.z + eps;
        float e3 = av.w - nv.w + eps;
        sn += e0 * e0 + e1 * e1 + e2 * e2 + e3 * e3;
    }

    __shared__ float shp[32];
    __shared__ float shn[32];

    for (int offset = 16; offset > 0; offset >>= 1) {
        sp += __shfl_down_sync(0xffffffffu, sp, offset);
        sn += __shfl_down_sync(0xffffffffu, sn, offset);
    }

    int lane = threadIdx.x & 31;
    int warp = threadIdx.x >> 5;
    if (lane == 0) {
        shp[warp] = sp;
        shn[warp] = sn;
    }
    __syncthreads();

    if (warp == 0) {
        int nwarp = (blockDim.x + 31) >> 5;
        sp = (lane < nwarp) ? shp[lane] : 0.0f;
        sn = (lane < nwarp) ? shn[lane] : 0.0f;
        for (int offset = 16; offset > 0; offset >>= 1) {
            sp += __shfl_down_sync(0xffffffffu, sp, offset);
            sn += __shfl_down_sync(0xffffffffu, sn, offset);
        }
        if (lane == 0) {
            float dist_pos = sqrtf(sp);
            float dist_neg = sqrtf(sn);
            float loss = dist_pos - dist_neg + margin;
            loss = loss > 0.0f ? loss : 0.0f;
            atomicAdd(out, loss);
        }
    }
}

__global__ void triplet_margin_loss_scalar_kernel(
    const float* __restrict__ anchor,
    const float* __restrict__ positive,
    const float* __restrict__ negative,
    float* __restrict__ out,
    long row_len,
    float margin,
    float eps
) {
    long row = blockIdx.x;
    const float* a_row = anchor + row * row_len;
    const float* p_row = positive + row * row_len;
    const float* n_row = negative + row * row_len;

    float sp = 0.0f;
    float sn = 0.0f;

    for (long j = threadIdx.x; j < row_len; j += blockDim.x) {
        float a = a_row[j];
        float d = a - p_row[j] + eps;
        float e = a - n_row[j] + eps;
        sp += d * d;
        sn += e * e;
    }

    __shared__ float shp[32];
    __shared__ float shn[32];

    for (int offset = 16; offset > 0; offset >>= 1) {
        sp += __shfl_down_sync(0xffffffffu, sp, offset);
        sn += __shfl_down_sync(0xffffffffu, sn, offset);
    }

    int lane = threadIdx.x & 31;
    int warp = threadIdx.x >> 5;
    if (lane == 0) {
        shp[warp] = sp;
        shn[warp] = sn;
    }
    __syncthreads();

    if (warp == 0) {
        int nwarp = (blockDim.x + 31) >> 5;
        sp = (lane < nwarp) ? shp[lane] : 0.0f;
        sn = (lane < nwarp) ? shn[lane] : 0.0f;
        for (int offset = 16; offset > 0; offset >>= 1) {
            sp += __shfl_down_sync(0xffffffffu, sp, offset);
            sn += __shfl_down_sync(0xffffffffu, sn, offset);
        }
        if (lane == 0) {
            float dist_pos = sqrtf(sp);
            float dist_neg = sqrtf(sn);
            float loss = dist_pos - dist_neg + margin;
            loss = loss > 0.0f ? loss : 0.0f;
            atomicAdd(out, loss);
        }
    }
}

torch::Tensor triplet_margin_loss_cuda(
    torch::Tensor anchor,
    torch::Tensor positive,
    torch::Tensor negative,
    double margin,
    double eps
) {
    TORCH_CHECK(anchor.is_cuda() && positive.is_cuda() && negative.is_cuda(),
                "All inputs must be CUDA tensors");
    TORCH_CHECK(anchor.scalar_type() == at::kFloat &&
                positive.scalar_type() == at::kFloat &&
                negative.scalar_type() == at::kFloat,
                "All inputs must be float32");

    auto a = anchor.contiguous();
    auto p = positive.contiguous();
    auto n = negative.contiguous();

    long batch = a.size(0);
    long row_len = (batch > 0) ? (a.numel() / batch) : 0;

    auto out = torch::zeros({}, a.options());

    if (batch == 0) {
        return out;
    }

    const int threads = 256;
    dim3 blocks((unsigned int)batch);

    if ((row_len & 3) == 0 && row_len > 0) {
        long row_len4 = row_len >> 2;
        triplet_margin_loss_vec4_kernel<<<blocks, threads>>>(
            reinterpret_cast<const float4*>(a.data_ptr<float>()),
            reinterpret_cast<const float4*>(p.data_ptr<float>()),
            reinterpret_cast<const float4*>(n.data_ptr<float>()),
            out.data_ptr<float>(),
            row_len4,
            (float)margin,
            (float)eps
        );
    } else {
        triplet_margin_loss_scalar_kernel<<<blocks, threads>>>(
            a.data_ptr<float>(),
            p.data_ptr<float>(),
            n.data_ptr<float>(),
            out.data_ptr<float>(),
            row_len,
            (float)margin,
            (float)eps
        );
    }

    // mean reduction over the batch
    return out.div((double)batch);
}
"""

triplet_margin_loss_cpp_source = (
    "torch::Tensor triplet_margin_loss_cuda(torch::Tensor anchor, torch::Tensor positive, "
    "torch::Tensor negative, double margin, double eps);"
)

# Compile the fused triplet margin loss CUDA extension
triplet_margin_loss_ext = load_inline(
    name="fused_triplet_margin_loss",
    cpp_sources=triplet_margin_loss_cpp_source,
    cuda_sources=triplet_margin_loss_cuda_source,
    functions=["triplet_margin_loss_cuda"],
    verbose=False,
    extra_cuda_cflags=["-O3"],
)


class ModelNew(nn.Module):
    """
    Optimized Triplet Margin Loss using a single fused CUDA kernel.

    Parameters:
        margin (float): The margin between the positive and negative samples.
    """
    def __init__(self, margin=1.0):
        super(ModelNew, self).__init__()
        self.margin = float(margin)
        self.p = 2.0
        self.eps = 1e-6
        self.swap = False
        self.reduction = 'mean'

    def forward(self, anchor, positive, negative):
        if (anchor.is_cuda and anchor.dtype == torch.float32 and anchor.dim() == 2
                and anchor.shape == positive.shape == negative.shape):
            return triplet_margin_loss_ext.triplet_margin_loss_cuda(
                anchor, positive, negative, self.margin, self.eps
            )
        # Generic fallback for non-CUDA / non-float32 / non-2D inputs
        return F.triplet_margin_loss(
            anchor, positive, negative,
            margin=self.margin, p=self.p, eps=self.eps,
            swap=self.swap, reduction=self.reduction
        )
```