batchsize dim for flashinfer api

[arthursunbao]

Hi,

I am trying to figure out how to replace the interface from flash attention to flashinfer

In flashattention, the q,k,v in flash_attn_func has the batchsize dimension: q: (batch_size, seqlen, nheads, headdim). However, I found that in flashinfer, it seems that it only has [kv_len, num_kv_heads, head_dim] for both prefill and decode api.

So how can I replace the flash_attn_func to FlashInfer Python API?

Jason