Avoiding Numba CUDA Jit Race Condition

I have a simple example:

from numba import cuda
import numpy as np
import math

@cuda.jit
def func(i, y, z):
    start = cuda.grid(1)
    stride = cuda.gridsize(1)

    for j in range(start, y.shape[0], stride):
        # Note that these aren't my real functions but they demo the point
        if i < j:
            y[j, 0] = i
            z[j, 0] = i + j
        if i == j:
            y[j, 1] = i
            z[j, 1] = i * j
        if i > j:
            y[j, 2] = i
            z[j, 2] = j


if __name__ == '__main__':
    n = 30
    y = np.ones((n, 3))
    z = np.ones((n, 3)) * -1
    device_y = cuda.to_device(y)
    device_z = cuda.to_device(z)
    max_i = 5
    threads_per_block = 10
    blocks_per_grid = math.ceil(y.shape[0]/threads_per_block[1])

    for i in range(max_i):
        func[blocks_per_grid, threads_per_block](i, device_y, device_z)

    out = device_y.copy_to_host()
    print(out)

And the output should look like this:

[[1. 0. 4.]
 [0. 1. 4.]
 [1. 2. 4.]
 [2. 3. 4.]
 [3. 4. 1.]
 [4. 1. 1.]
 [4. 1. 1.]
 [4. 1. 1.]
 [4. 1. 1.]
 [4. 1. 1.]
 [4. 1. 1.]
 [4. 1. 1.]
 [4. 1. 1.]
 [4. 1. 1.]
 [4. 1. 1.]
 [4. 1. 1.]
 [4. 1. 1.]
 [4. 1. 1.]
 [4. 1. 1.]
 [4. 1. 1.]
 [4. 1. 1.]
 [4. 1. 1.]
 [4. 1. 1.]
 [4. 1. 1.]
 [4. 1. 1.]
 [4. 1. 1.]
 [4. 1. 1.]
 [4. 1. 1.]
 [4. 1. 1.]
 [4. 1. 1.]]

However, when max_i is large then most of the time is spend calling the CUDA kernel and I want to make this kernel as fast as possible. So, I'm trying to figure out how to move the max_i for loop into the kernel but it looks like I am getting into race conditions. Here's what I currently have:

from numba import cuda
import numpy as np
import math

@cuda.jit
def func(max_i, y, z):
    a, b = cuda.grid(2)
    a_stride, b_stride = cuda.gridsize(2)

    for i in range(a, max_i, a_stride):
        for j in range(b, y.shape[0], b_stride):
            if i < j:
                y[j, 0] = i
                z[j, 0] = i + j
            if i == j:
                y[j, 1] = i
                z[j, 1] = i * j
            if i > j:
                y[j, 2] = i
                z[j, 2] = j

if __name__ == '__main__':
    n = 30
    y = np.ones((n, 3))
    z = np.ones((n, 3)) * -1
    device_y = cuda.to_device(y)
    device_z = cuda.to_device(z)
    max_i = 5
    threads_per_block = (1, 10)
    blocks_per_grid = (max_i, math.ceil(y.shape[0]/threads_per_block[1]))

    func[blocks_per_grid, threads_per_block](max_i, device_y, device_z)

    out = device_y.copy_to_host()
    print(out)

And this (incorrect) output looks like:

[[1. 0. 4.]
 [0. 1. 4.]
 [1. 2. 4.]
 [1. 3. 4.]  # Should be [2. 3. 4.]
 [3. 4. 1.]
 [4. 1. 1.]
 [3. 1. 1.]  # Should be [4. 1. 1.]
 [3. 1. 1.]  # Should be [4. 1. 1.]
 [3. 1. 1.]  # Should be [4. 1. 1.]
 [3. 1. 1.]  # Should be [4. 1. 1.]
 [4. 1. 1.]
 [4. 1. 1.]
 [4. 1. 1.]
 [4. 1. 1.]
 [4. 1. 1.]
 [4. 1. 1.]
 [4. 1. 1.]
 [4. 1. 1.]
 [4. 1. 1.]
 [4. 1. 1.]
 [0. 1. 1.]  # Should be [4. 1. 1.]
 [0. 1. 1.]  # Should be [4. 1. 1.]
 [0. 1. 1.]  # Should be [4. 1. 1.]
 [0. 1. 1.]  # Should be [4. 1. 1.]
 [0. 1. 1.]  # Should be [4. 1. 1.]
 [0. 1. 1.]  # Should be [4. 1. 1.]
 [0. 1. 1.]  # Should be [4. 1. 1.]
 [0. 1. 1.]  # Should be [4. 1. 1.]
 [0. 1. 1.]  # Should be [4. 1. 1.]
 [0. 1. 1.]]  # Should be [4. 1. 1.]

As stated above, how do I obtain the right answer by using a single kernel while making this kernel as fast as possible (i.e., avoiding atomic operations)?

from Avoiding Numba CUDA Jit Race Condition