[Question]: RuntimeError encountered when trying to reproduce results in needle in a haystack

Describe the issue

Hi,

Thanks again for your help. I encountered an error while reproducing results in needle_in_a_haystack by running bash experiments/needle_in_a_haystack/run_needle.sh and would appreciate any insights:

[   1000   72357  143714  215071  286429  357786  429143  500500  571857
  643214  714571  785929  857286  928643 1000000]
[ 286429  357786  429143  500500  571857  643214  714571  785929  857286
  928643 1000000]
# Too long, ignore some logs
 File "/home/far/MInference/minference/modules/minference_forward.py", line 656, in forward
    part_o = self.gather_last_q_vertical_slash_topk_v4(part_q, part_k, part_v, head)
  File "/home/far/MInference/minference/modules/minference_forward.py", line 463, in gather_last_q_vertical_slash_topk_v4
    return fc(q, k, v, vertical_size, slash_size)
  File "/home/far/MInference/minference/modules/minference_forward.py", line 383, in vertical_and_slash_kernel
    slash = sum_all_diagonal_matrix(qk)[...,:-last_q + 1]
  File "/home/far/MInference/minference/modules/minference_forward.py", line 103, in sum_all_diagonal_matrix
    zero_mat = torch.zeros((b, h, n, n)).to(mat.device) # Zero matrix used for padding
  File "/home/far/MInference/minference/modules/minference_forward.py", line 103, in sum_all_diagonal_matrix
    zero_mat = torch.zeros((b, h, n, n)).to(mat.device) # Zero matrix used for padding
RuntimeError: CUDA error: an illegal memory access was encountered
CUDA kernel errors might be asynchronously reported at some other API call, so the stacktrace below might be incorrect.

I noticed the error only occurs when starting from job 4 with the --kv_cache_cpu argument. Jobs in the range [0-4) work fine. Any suggestions on this?

Additionally, I found that the vllm module is required when performing the needle_in_a_haystack experiment. In my opinion, vllm isn't necessary for minference. Is there a specific reason for this, or something I might have missed?

Looking forward to your response!