There are precision errors compared with flash_attn_2_cuda.varlen_fwd

[Amanda-Barbara]

There are precision errors compared with flash_attn_2_cuda.varlen_fwd when I use flashinfer.single_prefill_with_kv_cache function to run cohere_plus model, below is the code I used:
fi_fwd_out = flashinfer.single_prefill_with_kv_cache(q.contiguous(), k.contiguous(), v.contiguous(), causal=True, sm_scale=softmax_scale, allow_fp16_qk_reduction=False)
fa2_fwd_out = flash_attn_2_cuda.varlen_fwd(
q,
k,
v,
out,
cu_seqlens,
cu_seqlens,
max_s,
max_s,
0.0,
softmax_scale,
False,
True,
False,
None,
)
torch.allclose(fi_fwd_out, fa2_fwd_out, rtol=1e-3, atol=1e-3)
It is worth noting that the first half of the layers are same, but second half are different.
can you give an official example code for precision comparison with flash_attn_2_cuda.varlen_fwd? Thanks!

[Amanda-Barbara]

It seems the problem with the accuracy error is that flashinfer.single_prefill_with_kv_cache doesn't support cu_seqlens_q and cu_seqlens_k, If I want to use flashinfer's prefill function， How to call it like flash_attn_2_cuda.varlen_fwd?

[yzh119]

single_prefill_with_kv_cache is only designed for single request (no batching and variable length).

For batch prefill with variable length, you have to use https://docs.flashinfer.ai/api/python/prefill.html#flashinfer.prefill.BatchPrefillWithRaggedKVCacheWrapper, where query and key/value cache are organized as ragged tensor. (see our layout documentation: https://docs.flashinfer.ai/tutorials/kv_layout.html).

If use you paged kv-cache, you should use https://docs.flashinfer.ai/api/python/prefill.html#flashinfer.prefill.BatchPrefillWithPagedKVCacheWrapper

[Amanda-Barbara]

@yzh119 thanks very much, I will try it.