potential race-condition in RoPE kernel?

[ChristianSchott]

Hello,
while looking through your RoPE CUDA implementation (with an LLM), it pointed out a potential race condition in the shared memory access.

for (int h = 0; h < H; h++)
{
    // then, load all the token for this head in shared memory
    shared[threadIdx.x] = tokens[b][n][h][threadIdx.x];
    __syncthreads();

    const float u = shared[m];
    const float v = shared[m+Q];
    
    // write output
    if ((threadIdx.x % (D/3)) < Q)
        tokens[b][n][h][threadIdx.x] = u*cos - v*sin;
    else
        tokens[b][n][h][threadIdx.x] = v*cos + u*sin;

    __syncthreads(); // <-- sync before overwriting shared in the next iteration
}

You already sync to ensure all threads have written their token value to shared, however in theory one thread could already write the next token value to shared before another thread has read their u, v from shared.

I am not entirely sure if this will actually materialize in practice, as the threads reading from shared should be in the same warp as the threads writing to shared[m]/shared[m+Q].
However, even intra-warp scheduling does not seem to be guaranteed anymore..

Or did I completely miss something here?

Apart from that, it also suggested adding the CHECK_KERNEL(); call, and add a check to ensure pos is actually of type i64.
TORCH_CHECK(pos.scalar_type() == torch::kInt64, "pos must be an int64 tensor");

[ywyue]

Hi @ChristianSchott, thanks so much for pointing out this. You are correct on all points. The missing __syncthreads(); at the end of the loop is a real correctness issue. In practice, it doesn't trigger today because D<=32 (in our case D is 18, which is the head dimension) means all threads fit in a single warp and execute in lockstep. But your are right to flag the intra-warp guarantee concern. We will add the sync.

CHECK_KERNEL(); and ensuring int64 for pos are good defensive additions - we will include them as well. Just FYI: in alternative implementation, pos does not have to be int64, it can also be float (e.g. if pos is derived from original point coordinates rather than grid coordinates).