@triton.jit
def store_1d(vals, ptr, sz: tl.constexpr, n, max, stride=1):
    """Store 1d block into nth chunk of vector (defined by ptr), where each chunk has size sz"""
    offs = get_1d_offest(sz, n)
    mask = get_1d_mask(offs, max)
    tl.store(ptr + offs, vals, mask)
