[Misc]: a question about chunked-prefill in flash-attn backends

[HarryWu99]

Anything you want to discuss about vllm.

vllm/vllm/attention/backends/flash_attn.py

Line 282 in 99caa49

if prefill_meta := attn_metadata.prefill_metadata:

I noticed that in flash-attn backends. forward_prefix and forward_decode seem to be executed serially. Does forward_decode wait for forward_prefix to finish before running? Can this take advantage of the performance provided by chunked-prefill? I mean the tokens of prefill are in the same batch as the tokens of decode.

if prefill_meta := attn_metadata.prefill_metadata:
    output[:num_prefill_tokens] = PagedAttention.forward_prefix(...)

if decode_meta := attn_metadata.decode_metadata:
    output[num_prefill_tokens:] = PagedAttention.forward_decode(...)

[rkooo567]

I noticed that in flash-attn backends. forward_prefix and forward_decode seem to be executed serially. Does forward_decode wait for forward_prefix to finish before running? Can this take advantage of the performance provided by chunked-prefill? I mean the tokens of prefill are in the same batch as the tokens of decode.

Yeah right now, it is running serially. I think after #4681, it should be possible to run them in the same attn kernel, but based on our past internal benchmark before, it didn't make much difference (we can definitely try to see how much perf improvement it will have). But this could be different now.

Note that this should be done after we re-revert #4820 because we should use prefix kernel to run both in the same attn kernel, and existing prefix kernel is too slow (flash attn varlen has at least 3X faster than this kernel)

[CrimsonDump]

I noticed that in flash-attn backends. forward_prefix and forward_decode seem to be executed serially. Does forward_decode wait for forward_prefix to finish before running? Can this take advantage of the performance provided by chunked-prefill? I mean the tokens of prefill are in the same batch as the tokens of decode.

Yeah right now, it is running serially. I think after #4681, it should be possible to run them in the same attn kernel, but based on our past internal benchmark before, it didn't make much difference (we can definitely try to see how much perf improvement it will have). But this could be different now.

Note that this should be done after we re-revert #4820 because we should use prefix kernel to run both in the same attn kernel, and existing prefix kernel is too slow (flash attn varlen has at least 3X faster than this kernel)

Is there a Issue/PR to "re-revert #4820 " for us to track?