Performance drop on large input tensor

Hi community, I'm trying to implement the following pytorch function with triton:

// input tensor shape: [row, 4]
// output tensor shape: [row, 4]
def torch_xyxy2xywh(x):
    # Convert nx4 boxes from [x1, y1, x2, y2] to [x, y, w, h] where xy1=top-left, xy2=bottom-right
    y = x.clone() if isinstance(x, torch.Tensor) else np.copy(x)
    y[:, 0] = (x[:, 0] + x[:, 2]) / 2  # x center
    y[:, 1] = (x[:, 1] + x[:, 3]) / 2  # y center
    y[:, 2] = x[:, 2] - x[:, 0]  # width
    y[:, 3] = x[:, 3] - x[:, 1]  # height
    return y

Here's my triton implementation with autotune enabled:

@triton.autotune(
    configs=[
        triton.Config({'BLOCK_SIZE': 64}),
        triton.Config({'BLOCK_SIZE': 128}),
        triton.Config({'BLOCK_SIZE': 256}),
        triton.Config({'BLOCK_SIZE': 512}),
    ],
    key=['n_rows'],
)
@triton.jit
def triton_xyxy2xywh_kernel(
    input_ptr,
    output_ptr,
    n_rows,
    BLOCK_SIZE: tl.constexpr,
):
    # one program process `[BLOCK_SIZE, 4]` block
    pid = tl.program_id(0)
    stride_row = 4 # always have 4 cols
    offset_row = pid * BLOCK_SIZE + tl.arange(0, BLOCK_SIZE)
    offset_col = tl.arange(0, 4)
    row_mask = offset_row[:, None] < n_rows
    x0_ptrs = input_ptr + offset_row[:, None] * stride_row
    x1_ptrs = input_ptr + offset_row[:, None] * stride_row + 1
    x2_ptrs = input_ptr + offset_row[:, None] * stride_row + 2
    x3_ptrs = input_ptr + offset_row[:, None] * stride_row + 3

    x0 = tl.load(x0_ptrs, mask=row_mask)
    x1 = tl.load(x1_ptrs, mask=row_mask)
    x2 = tl.load(x2_ptrs, mask=row_mask)
    x3 = tl.load(x3_ptrs, mask=row_mask)

    y0 = (x0 + x2) / 2   # x center
    y1 = (x1 + x3) / 2   # y center
    y2 = x2 - x0         # width
    y3 = x3 - x1         # height

    tl.store(output_ptr + offset_row[:, None] * stride_row    , y0, mask=row_mask)
    tl.store(output_ptr + offset_row[:, None] * stride_row + 1, y1, mask=row_mask)
    tl.store(output_ptr + offset_row[:, None] * stride_row + 2, y2, mask=row_mask)
    tl.store(output_ptr + offset_row[:, None] * stride_row + 3, y3, mask=row_mask)

I've also written a simple CUDA kernel for reference, and benchmarked the code. Here's the performance:

Both triton and CUDA implementation out-performs the naive pytorch function, which is great. The triton code performs as-good-as CUDA kernel until the row of input exceeds ~100K.

I wonder if the performance drop is related to my implementation? The 4 column-wise load/store seems to be the bottleneck to me. How do I optimize the implementation? Didn't find any good tutorials on such performance issues...