[Bug] RuntimeError: CUDA error: an illegal memory access was encountered with batch_size>=1

[weixuansun]

I encountered illegal memory access issue when batch_size>=1 using NSA compression.
I have checked the code and guess the issue is from the ptr position of k and v:

native-sparse-attention/native_sparse_attention/ops/parallel.py

Line 86 in c776771

p_k = tl.make_block_ptr(k + (bos * H + i_h) * K, (K, TC), (1, H*K), (0, i_c), (BK, BC), (0, 1))

I haven't dive in the code too much , but here is a possible fix:

bos_k, eos_k = i_b * TC, i_b * TC + TC
p_k = tl.make_block_ptr(k + (bos_k * H + i_h) * K, (K, TC), (1, H*K), (0, i_c), (BK, BC), (0, 1))

[ChloeL19]

Hi, is this in any way related to the following error message?

ValueError: Pointer argument (at 0) cannot be accessed from Triton (cpu tensor?)

[ChloeL19]

In any event, I tried applying the fix @weixuansun suggested, and the unit tests for test_nsa_with_compression still seem to fail with CUDA illegal memory access errors