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

_cpp_source = r"""
#include <torch/extension.h>
torch::Tensor triplet_loss_cuda(
    torch::Tensor anchor,
    torch::Tensor positive,
    torch::Tensor negative,
    double margin);
"""

_cuda_source = r"""
#include <torch/extension.h>
#include <ATen/cuda/CUDAContext.h>
#include <c10/cuda/CUDAGuard.h>
#include <c10/cuda/CUDAException.h>
#include <cuda_runtime.h>
#include <cstdint>

__device__ __forceinline__ float warp_sum(float value) {
    #pragma unroll
    for (int offset = 16; offset > 0; offset >>= 1)
        value += __shfl_down_sync(0xffffffff, value, offset);
    return value;
}

__device__ __forceinline__ void accumulate(
    float a, float p, float n, float& sp, float& sn) {
    float dp = (a - p) + 1.0e-6f;
    float dn = (a - n) + 1.0e-6f;
    sp = fmaf(dp, dp, sp);
    sn = fmaf(dn, dn, sn);
}

template<bool Vectorized>
__global__ void triplet_rows(
    const float* __restrict__ anchor,
    const float* __restrict__ positive,
    const float* __restrict__ negative,
    float* __restrict__ losses,
    int64_t width,
    float margin) {

    constexpr int Warps = 8;
    __shared__ float shared_p[Warps];
    __shared__ float shared_n[Warps];

    const int tid = threadIdx.x;
    const int lane = tid & 31;
    const int warp = tid >> 5;
    const int64_t base = int64_t(blockIdx.x) * width;

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

    if (Vectorized) {
        const float4* a =
            reinterpret_cast<const float4*>(anchor + base);
        const float4* p =
            reinterpret_cast<const float4*>(positive + base);
        const float4* n =
            reinterpret_cast<const float4*>(negative + base);

        for (int64_t j = tid; j < width / 4; j += blockDim.x) {
            float4 av = a[j];
            float4 pv = p[j];
            float4 nv = n[j];
            accumulate(av.x, pv.x, nv.x, sp, sn);
            accumulate(av.y, pv.y, nv.y, sp, sn);
            accumulate(av.z, pv.z, nv.z, sp, sn);
            accumulate(av.w, pv.w, nv.w, sp, sn);
        }
    } else {
        for (int64_t j = tid; j < width; j += blockDim.x)
            accumulate(anchor[base + j], positive[base + j],
                       negative[base + j], sp, sn);
    }

    sp = warp_sum(sp);
    sn = warp_sum(sn);

    if (lane == 0) {
        shared_p[warp] = sp;
        shared_n[warp] = sn;
    }
    __syncthreads();

    if (warp == 0) {
        sp = warp_sum(lane < Warps ? shared_p[lane] : 0.0f);
        sn = warp_sum(lane < Warps ? shared_n[lane] : 0.0f);

        if (lane == 0) {
            float loss = (margin + sqrtf(sp)) - sqrtf(sn);
            losses[blockIdx.x] = loss < 0.0f ? 0.0f : loss;
        }
    }
}

__global__ void triplet_mean(
    const float* __restrict__ losses,
    float* __restrict__ output,
    int64_t rows) {

    __shared__ float partials[8];
    const int tid = threadIdx.x;
    const int lane = tid & 31;
    const int warp = tid >> 5;

    float sum = 0.0f;
    for (int64_t i = tid; i < rows; i += blockDim.x)
        sum += losses[i];

    sum = warp_sum(sum);
    if (lane == 0)
        partials[warp] = sum;
    __syncthreads();

    if (warp == 0) {
        sum = warp_sum(lane < 8 ? partials[lane] : 0.0f);
        if (lane == 0)
            output[0] = sum / static_cast<float>(rows);
    }
}

torch::Tensor triplet_loss_cuda(
    torch::Tensor anchor,
    torch::Tensor positive,
    torch::Tensor negative,
    double margin) {

    TORCH_CHECK(anchor.is_cuda() && positive.is_cuda() &&
                negative.is_cuda(), "Inputs must be CUDA tensors");
    TORCH_CHECK(anchor.device() == positive.device() &&
                anchor.device() == negative.device(),
                "Inputs must be on the same device");
    TORCH_CHECK(anchor.scalar_type() == at::kFloat &&
                positive.scalar_type() == at::kFloat &&
                negative.scalar_type() == at::kFloat,
                "Inputs must have float32 dtype");
    TORCH_CHECK(anchor.dim() == 2 &&
                anchor.sizes() == positive.sizes() &&
                anchor.sizes() == negative.sizes(),
                "Inputs must have matching two-dimensional shapes");
    TORCH_CHECK(anchor.is_contiguous() && positive.is_contiguous() &&
                negative.is_contiguous(), "Inputs must be contiguous");
    TORCH_CHECK(anchor.size(0) > 0, "Batch must be nonempty");

    const c10::cuda::CUDAGuard guard(anchor.device());
    const int64_t rows = anchor.size(0);
    const int64_t width = anchor.size(1);

    auto losses = torch::empty({rows}, anchor.options());
    auto output = torch::empty({}, anchor.options());
    cudaStream_t stream = at::cuda::getCurrentCUDAStream();

    const float* a = anchor.data_ptr<float>();
    const float* p = positive.data_ptr<float>();
    const float* n = negative.data_ptr<float>();

    const bool vectorized =
        width % 4 == 0 &&
        ((reinterpret_cast<uintptr_t>(a) |
          reinterpret_cast<uintptr_t>(p) |
          reinterpret_cast<uintptr_t>(n)) & 15) == 0;

    if (vectorized) {
        triplet_rows<true><<<rows, 256, 0, stream>>>(
            a, p, n, losses.data_ptr<float>(), width,
            static_cast<float>(margin));
    } else {
        triplet_rows<false><<<rows, 256, 0, stream>>>(
            a, p, n, losses.data_ptr<float>(), width,
            static_cast<float>(margin));
    }
    C10_CUDA_KERNEL_LAUNCH_CHECK();

    triplet_mean<<<1, 256, 0, stream>>>(
        losses.data_ptr<float>(), output.data_ptr<float>(), rows);
    C10_CUDA_KERNEL_LAUNCH_CHECK();

    return output;
}
"""

_triplet_extension = load_inline(
    name="fused_triplet_margin_loss_cuda",
    cpp_sources=_cpp_source,
    cuda_sources=_cuda_source,
    functions=["triplet_loss_cuda"],
    extra_cflags=["-O3"],
    extra_cuda_cflags=["-O3"],
    verbose=False,
)


class ModelNew(nn.Module):
    def __init__(self, margin=1.0):
        super().__init__()
        if margin <= 0:
            raise ValueError("margin must be greater than zero")
        self.margin = float(margin)

    def forward(self, anchor, positive, negative):
        inputs = (anchor, positive, negative)
        use_custom = (
            all(
                x.is_cuda
                and x.dtype == torch.float32
                and x.ndim == 2
                and x.is_contiguous()
                for x in inputs
            )
            and anchor.shape == positive.shape == negative.shape
            and anchor.device == positive.device == negative.device
            and anchor.shape[0] > 0
            and not (
                torch.is_grad_enabled()
                and any(x.requires_grad for x in inputs)
            )
        )

        if use_custom:
            return _triplet_extension.triplet_loss_cuda(
                anchor, positive, negative, self.margin
            )

        return F.triplet_margin_loss(
            anchor, positive, negative,
            margin=self.margin, p=2, eps=1e-6,
            swap=False, reduction="mean",
        )
```