Segmentation Fault with Qwen-7B and Specific Attention Mechanisms

[RuiZiHan]

First off, I want to express my appreciation for your excellent work on this project! It's been incredibly helpful.

I'm encountering a Segmentation fault (core dumped) error when running the Qwen-7B model on the ruler-niah task. This issue specifically occurs when using the quest, flexprefill, vertival_and_slash, and block_sparse attention mechanisms. Interestingly, the dense, snapkv, and ada_snapkv attention types do not produce this error and run successfully.

I've traced the segmentation fault to line 96 in sparse-frontier/sparse_frontier/modelling/models/vllm_model.py, within the section wrapped by vllm. My environment is nearly identical to yours, including the versions of CUDA, PyTorch, Python, FlashAttention, and vLLM.

Could you please provide some guidance on what might be causing this or suggest any potential workarounds? Any insights would be greatly appreciated.

Thank you for your time and help!

[PiotrNawrot]

Thanks for kind words!

Hmm, so something must be happening inside the model execution. What about triton version? Could you try re-doing the environment too maybe? I double checked it and it works on my end :/

Could you try calling block_sparse or vertical_and_slash with some randomly generated inputs of similar shapes and dtype and see if the error still occurs? You could use functions vertical_and_slash_kernel or block_sparse_attention from sparse-frontier/sparse_frontier/modelling/attention/minference or the functions from sparse-frontier/sparse_frontier/modelling/attention/efficient_prefilling.py

Does the problem occur for other tasks as well?
Have you overwritten any other parameters or functions, would you mind sharing the exact command you've run?

[PiotrNawrot]

Another idea is that maybe your GPU doesn't support some of the operations? What GPU are you using?