Simple 2D reduction

[NightWinkle]

Hello,

As an exercise, I wanted to implement simple 2D reductions in Triton, over two vectors x of size N and y of size M, of the following form :
s[i] = sum(exp(-(x[i] - y[j])/(2*sigma)) for j in range(M))

I wrote the following code :

@triton.autotune(
    configs=[
        triton.Config({'BLOCK_SIZE': 256}, num_stages=3, num_warps=8),
        triton.Config({'BLOCK_SIZE': 256}, num_stages=4, num_warps=4),
        triton.Config({'BLOCK_SIZE': 128}, num_stages=4, num_warps=4),
        triton.Config({'BLOCK_SIZE': 64}, num_stages=4, num_warps=4),
        triton.Config({'BLOCK_SIZE': 64}, num_stages=5, num_warps=2),
        triton.Config({'BLOCK_SIZE': 32}, num_stages=5, num_warps=2),
    ],
    key=['M', 'N', 'D'],
)
@triton.jit
def sumgausskernel(
    X, Y, sigma,
    N:tl.constexpr, M:tl.constexpr, D: tl.constexpr,
    stride_x_n, stride_x_d,
    stride_y_m, stride_y_d,
    Out,
    stride_x_o,
    BLOCK_SIZE: tl.constexpr
):
    block_id = tl.program_id(0)
    
    Xblockptr = tl.make_block_ptr(base=X, shape=(N, D), strides=(stride_x_n, stride_x_d), offsets=(block_id*BLOCK_SIZE, 0), block_shape=(BLOCK_SIZE, D), order=(1, 0))
    Yblockptr = tl.make_block_ptr(base=Y, shape=(M, D), strides=(stride_y_m, stride_y_d), offsets=(0, 0), block_shape=(BLOCK_SIZE, D), order=(0, 1))

    Xblock = tl.load(Xblockptr)
    acc = tl.zeros([BLOCK_SIZE], dtype=tl.float32)
    
    for start_m in range(0, M, BLOCK_SIZE):
        Yblock = tl.load(Yblockptr)
        
        diff = tl.expand_dims(Yblock, axis=0) - tl.expand_dims(Xblock, axis=1)
        diff = diff*diff
        
        norm = tl.sum(diff, axis=2)
        norm = tl.exp(-norm/(2*sigma))
        acc += tl.sum(norm, axis=1)
        
        Yblockptr = tl.advance(Yblockptr, (0, BLOCK_SIZE))
        
    Outblockptr = tl.make_block_ptr(base=Out, shape=(N,), strides=(stride_x_o,), offsets=(block_id*BLOCK_SIZE,), block_shape=(BLOCK_SIZE,), order=(0,))
    tl.store(Outblockptr, acc) 

def sumgauss(x, y, sigma=0.2):
    N, M, D = int(x.shape[0]), int(y.shape[0]), int(x.shape[1])

    out = torch.zeros((N,), device="cuda")
    
    grid = lambda META: (
        triton.cdiv(N, META["BLOCK_SIZE"]),
    )
    sumgausskernel[grid](
        x, y, sigma,
        N, M, D,
        x.stride(0), x.stride(1),
        y.stride(0), y.stride(1),
        out,
        out.stride(0),
    )
    return out

I've ignored boundary effects in this example. The performance I get is worse than PyKeops implementation for the similar sum (https://www.kernel-operations.io/keops/engine/map_reduce_schemes.html), while it should be the same results from my understanding (even though there is no fine control over local and shared memory).

Is there something i'm doing wrong ?
Thanks