A bug when run run_chunkllama_100k with flash decoding

[lwj2001]

When i use Llama3 (with flash decoding ) to run run_chunkllama_100k, and it can successful start. But when i input prompt. then encounter a TypeError that:

File "ChunkLlama/flash_decoding_chunkllama.py", line 140, in new_flash_attn_with_kvcache
    out, softmax_lse = flash_attn_cuda.fwd_kvcache(
TypeError: fwd_kvcache(): incompatible function arguments. The following argument types are supported:
    1. (arg0: torch.Tensor, arg1: torch.Tensor, arg2: torch.Tensor, arg3: Optional[torch.Tensor], arg4: Optional[torch.Tensor], arg5: Optional[torch.Tensor], arg6: Optional[torch.Tensor], arg7: Optional[torch.Tensor], arg8: Optional[torch.Tensor], arg9: Optional[torch.Tensor], arg10: Optional[torch.Tensor], arg11: Optional[torch.Tensor], arg12: Optional[torch.Tensor], 
    arg13: float, arg14: bool, arg15: int, arg16: int, arg17: float, arg18: bool, arg19: int) -> list[torch.Tensor]
Invoked with: tensor(xxx) , ...... , None, None, None, None, None, None, 0.08838834764831845, False, -1, -1, True, 0

emmm, this is hard to me...

[ChenxinAn-fdu]

I have tested the code again and it works. python run_chunkllama_100k.py
Please make sure you have flash_attn >=2.5.3, transformers==4.37.2 and attn_implementation="flash_attention_2",

[hsiehjackson]

I think the error is because we should install flash_attn >=2.5.3,<2.6.0

[ChenxinAn-fdu]

Thank you !!! I have updated this in Readme