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


_cpp_source = """
#include <torch/extension.h>

torch::Tensor triplet_margin_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>

__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;
}

__global__ void triplet_rows_kernel(
    const float4* __restrict__ anchor,
    const float4* __restrict__ positive,
    const float4* __restrict__ negative,
    float* __restrict__ losses,
    const int64_t vectors_per_row,
    const float margin)
{
    const int tid = threadIdx.x;
    const int lane = tid & 31;
    const int warp = tid >> 5;
    const int64_t base = int64_t(blockIdx.x) * vectors_per_row;

    float p0 = 0.0f, p1 = 0.0f, p2 = 0.0f, p3 = 0.0f;
    float n0 = 0.0f, n1 = 0.0f, n2 = 0.0f, n3 = 0.0f;

    for (int64_t i = tid; i < vectors_per_row; i += blockDim.x) {
        const float4 a = anchor[base + i];
        const float4 p = positive[base + i];
        const float4 n = negative[base + i];

        const float dx0 = (a.x - p.x) + 1.0e-6f;
        const float dx1 = (a.y - p.y) + 1.0e-6f;
        const float dx2 = (a.z - p.z) + 1.0e-6f;
        const float dx3 = (a.w - p.w) + 1.0e-6f;

        const float dy0 = (a.x - n.x) + 1.0e-6f;
        const float dy1 = (a.y - n.y) + 1.0e-6f;
        const float dy2 = (a.z - n.z) + 1.0e-6f;
        const float dy3 = (a.w - n.w) + 1.0e-6f;

        p0 += dx0 * dx0;
        p1 += dx1 * dx1;
        p2 += dx2 * dx2;
        p3 += dx3 * dx3;

        n0 += dy0 * dy0;
        n1 += dy1 * dy1;
        n2 += dy2 * dy2;
        n3 += dy3 * dy3;
    }

    float ps = warp_sum((p0 + p1) + (p2 + p3));
    float ns = warp_sum((n0 + n1) + (n2 + n3));

    __shared__ float positive_sums[8];
    __shared__ float negative_sums[8];

    if (lane == 0) {
        positive_sums[warp] = ps;
        negative_sums[warp] = ns;
    }
    __syncthreads();

    if (warp == 0) {
        ps = warp_sum(lane < 8 ? positive_sums[lane] : 0.0f);
        ns = warp_sum(lane < 8 ? negative_sums[lane] : 0.0f);

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

__global__ void mean_loss_kernel(
    const float* __restrict__ losses,
    float* __restrict__ output,
    const int64_t count)
{
    float sum = 0.0f;
    for (int64_t i = threadIdx.x; i < count; i += blockDim.x) {
        sum += losses[i];
    }

    sum = warp_sum(sum);

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

    if (lane == 0) {
        warp_sums[warp] = sum;
    }
    __syncthreads();

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

torch::Tensor triplet_margin_cuda(
    torch::Tensor anchor,
    torch::Tensor positive,
    torch::Tensor negative,
    double margin)
{
    const c10::cuda::CUDAGuard device_guard(anchor.device());

    const int64_t rows = anchor.size(0);
    const int64_t vectors_per_row = anchor.size(1) / 4;

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

    triplet_rows_kernel<<<rows, 256, 0, stream>>>(
        reinterpret_cast<const float4*>(anchor.data_ptr<float>()),
        reinterpret_cast<const float4*>(positive.data_ptr<float>()),
        reinterpret_cast<const float4*>(negative.data_ptr<float>()),
        losses.data_ptr<float>(),
        vectors_per_row,
        static_cast<float>(margin));
    C10_CUDA_KERNEL_LAUNCH_CHECK();

    mean_loss_kernel<<<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_float4",
    cpp_sources=_cpp_source,
    cuda_sources=_cuda_source,
    functions=["triplet_margin_cuda"],
    extra_cflags=["-O3"],
    extra_cuda_cflags=["-O3", "--fmad=false"],
    verbose=False,
)


class ModelNew(nn.Module):
    def __init__(self, margin=1.0):
        super().__init__()
        self.loss_fn = nn.TripletMarginLoss(margin=margin)

    def forward(self, anchor, positive, negative):
        tensors = (anchor, positive, negative)

        use_custom = (
            anchor.is_cuda
            and anchor.ndim == 2
            and 0 < anchor.shape[0] <= 2147483647
            and anchor.shape[1] > 0
            and anchor.shape[1] % 4 == 0
            and all(
                t.device == anchor.device
                and t.dtype == torch.float32
                and t.shape == anchor.shape
                and t.is_contiguous()
                and t.data_ptr() % 16 == 0
                for t in tensors
            )
            and not (
                torch.is_grad_enabled()
                and any(t.requires_grad for t in tensors)
            )
        )

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

        return self.loss_fn(anchor, positive, negative)
```