[V1] Optimize the CPU overheads in FlashAttention custom op

[WoosukKwon]

With piece-wise CUDA graphs, we have to make sure that the attention custom op causes minimal CPU overheads. This PR made a few changes to optimize the CPU overheads in the FlashAttention custom op:

1. We directly use torch.ops.vllm_flash_attn_c.varlen_fwd rather than flash_attn_varlen_func, since FlashAttnFunc which inherits torch.autograd.Function causes unnecessary overheads.
2. We move the reshapes and shape check logics to outside of the custom op, so that they can be done at the CUDA graph capture time.

Results of python benchmarks/benchmark_latency.py (opt-125m) on a single H100 GPU:

• V1 main: 227 ms
• V1 this PR: 192 ms
• V0 + 8-step: 130 ms

Next step: further reduce the unnecessary CPU ops inside the FlashAttention op.

[github-actions]

👋 Hi! Thank you for contributing to the vLLM project.
Just a reminder: PRs would not trigger full CI run by default. Instead, it would only run fastcheck CI which starts running only a small and essential subset of CI tests to quickly catch errors. You can run other CI tests on top of those by going to your fastcheck build on Buildkite UI (linked in the PR checks section) and unblock them. If you do not have permission to unblock, ping simon-mo or khluu to add you in our Buildkite org.

Once the PR is approved and ready to go, your PR reviewer(s) can run CI to test the changes comprehensively before merging.

To run CI, PR reviewers can do one of these:

• Add ready label to the PR
• Enable auto-merge.

🚀

[youkaichao]

well, I think I forgot to update the v1 flash attention file, after #10558 , you don't need the torch.ops.vllm.unified_v1_flash_attention call.

nvm

[youkaichao]

can you also update the corresponding v0 code?

[tlrmchlsmth]

Looking at profile results on #9856, this saves about 60µs off of the CPU time spent in each flash attention call (approx 300µs -> 240µs)

Thanks!

[mgoin]

LGTM with Kaichao's comment, thanks for quickly improving this. The failing test is due to neuralmagic/Phi-3-medium-128k-instruct-quantized.w4a16 and unrelated

[WoosukKwon]

@youkaichao @mgoin As we merged vllm-project/flash-attention#30, we don't have to directly use torch.ops.vllm_flash_attn_c.varlen_fwd. We can just use flash_attn_varlen_func as we currently do. Both V0 and V1 already gets the benefits after vllm-project/flash-attention#30.

[WoosukKwon]

One weird phenomenon I found is that V1 has a spike in latency:

Avg latency: 0.20093455887205589 seconds
10% percentile latency: 0.1931818482640665 seconds
25% percentile latency: 0.19354040725738741 seconds
50% percentile latency: 0.19391279752017 seconds
75% percentile latency: 0.19426249974640086 seconds
90% percentile latency: 0.1961068181961309 seconds
99% percentile latency: 0.3368887884780999 seconds

This is highly reproducible on my dev machine. Can this be because of Python gc or something like that?

[robertgshaw2-redhat]

One weird phenomenon I found is that V1 has a spike in latency:

Avg latency: 0.20093455887205589 seconds
10% percentile latency: 0.1931818482640665 seconds
25% percentile latency: 0.19354040725738741 seconds
50% percentile latency: 0.19391279752017 seconds
75% percentile latency: 0.19426249974640086 seconds
90% percentile latency: 0.1961068181961309 seconds
99% percentile latency: 0.3368887884780999 seconds

This is highly reproducible on my dev machine. Can this be because of Python gc or something like that?

It’s probably the prefix caching …

[comaniac]

Hmm but benchmark_latency.py does sample each prompts separately: https://github.com/vllm-project/vllm/blob/main/benchmarks/benchmark_latency.py#L36

[comaniac]

Hmm but benchmark_latency.py does sample each prompts separately: https://github.com/vllm-project/vllm/blob/main/benchmarks/benchmark_latency.py#L36

Just found that it has a warmup phase. It's still possible due to prefix caching if all prompts are cached then. Suggest to explicitly disable prefix caching to double check.

[WoosukKwon]

@comaniac @robertgshaw2-neuralmagic You're right. The latency becomes stable when prefix caching is turned off.

Avg latency: 0.1945609479948568 seconds
10% percentile latency: 0.19310778125654907 seconds
25% percentile latency: 0.19390572598786093 seconds
50% percentile latency: 0.19475348049309105 seconds
75% percentile latency: 0.195164829317946 seconds
90% percentile latency: 0.19570096801035106 seconds
99% percentile latency: 0.1962820820847992 seconds