[RFC] Upstream Chunked Prefill

[rkooo567]

Progress

• 2D -> 1D query refactoring [RFC] Upstream Chunked Prefill #3130
• Sequence group metadata API update https://github.com/vllm-project/vllm/pull/3538/files
• Scheduler refactoring [3/N] Refactor scheduler for chunked prefill scheduling #3550
• Chunked prefill scheduler [Chunked Prefill][4/n] Chunked prefill scheduler. #3853
• Chunked prefill attention update [Core][5/N] Fully working chunked prefill e2e #3884
• logprob fix [Core] Refactoring sampler and support prompt logprob for chunked prefill  #4309
• Enable cuda-graph on prefill path -> Won't do

Future Extension (not a scope of this RFC)

• Make it work with sliding window attention
• Use same mechanism for prefix caching
• Use better kernel than context attention forward (ideally flash infer)

Chunked prefill chunks prefill requests into multiple chunks and batch them with decoding requests. Since prefill requests are compute-bound and decoding requests are memory-bound, it can greatly improve system efficiency by overlapping these two. See the linked papers for more details.

• https://arxiv.org/pdf/2401.08671.pdf
• https://arxiv.org/pdf/2308.16369.pdf

We are planning to upstream chunked prefill internally used at Anyscale to OSS Vllm repo. We turned on this feature in the Anyscale production endpoint.

Benchmark Result

See #3130 (comment)

The following diagram is the benchmark result of Llama 13B x 2 A100 for different QPS (it is the result from Anyscale forked vLLM). Chunked prefill greatly improves latency when QPS is high, but has competitive performance at low QPS. The detailed benchmark results for different parameters with OSS vLLM is in progress (ETA end of this week).

Design

Kernel

For chunked prefill, we are internally using flash attn with paged attention enabled. so we will also upstream flash attn integration first. We decided to go with existing context attention first, and migrate to flash infer. Note that the kernel choice is not finalized, and we are actively investigating other kernels like FlashInfer.

Scheduler & Batch Layout

• Since chunked prefill mixes the prefill and decoding, having 2D query with batch dimensions is inefficient. To get around this, we will move 2D query back to 1D query.
• We change Sequence so that it keeps track of the number prompting tokens have been prefilled.
• In Scheduler, if chunked prefill is enabled, it will first prioritize decoding requests, and then pick prompting request. If prompting request is too long, the prompting request will be chunked.
• The worker now will create metadata for both decoding requests and prompting requests, including their block tables, context lengths, and number of decoding and prefiling tokens and requests.

    Chunked Prefill support:
    If chunked prefill is enabled, the input will include both prompt tokens
    and generation tokens. The layout is as follows:
    |<---------------------- num_valid_tokens -------------------------->|
    |<--------- num_prompt_tokens ----->|<--- num_generation_tokens >|
    |<-prompt_0->|<-prompt_1->|...............|<-gen_0->|<-gen_1->|..............|
    Notice that both num_valid_tokens, num_prompt_tokens and num_generation_tokens
    includes padding.
    The actual prompt length and offeset are stored in cum_prompt_context_lens.
    The actual num generation tokens are stored in num_generation_tokens_tensor.
    To support chunked prefill, where the prompt and context might have different
    length, we stored the context's length in cum_prompt_context_lens.

• Once request is processed and sampled, the result belongs to chunked prefilling requests will be ignored.

Cuda Graph

Cuda graph won't be supported for chunked prefill because the performance benefit of cuda graph only happens with smaller # of tokens.

Following is how we make everything cuda-graphable. We pad both prompt_tokens and generation tokens, and the cuda-graph is cached based number of prompt-tokens, number of generation tokens, and prompt batch size. Following are the details.

        # The key for caching cuda-graph.
        # We ensure that with the same total number of tokens,
        # same number of prompts, and same number of prompt tokens,
        # the intermediate tensors have the same shapes thus can be
        # replayed.
        #
        # Refer to following layout for the shapes of intermediate tensors.
        # during attention calculation.
        #
        #|<---------------------- padded_num_tokens ------------------------->|
        #|<--------- num_prompt_tokens ----->|<--- num_generation_tokens----->|
        #|<-prompt_0->|<-prompt_1->|...|<pad>|<-gen_0->|<-gen_1->|......|<pad>|

[tdene]

The "Benchmark Result" screenshot links directly to anyscale's private fork of vllm, and thus is not viewable by the general public.

[rkooo567]

@tdene thank you for pointing that out! I believe it is fixed

[rkooo567]

E2e PR is here; #3204 (comment) (porting internal code to OSS takes some time)

[rkooo567]

For kernel, we will try using context_attention_forward and then switch to flash infer once it is available

[rkooo567]

I got e2e working. Benchmarking result with OSS will be displayed very soon. I am also start focusing merging PRs starting from #3236.

[rkooo567]

Here's the initial benchmark result for chunked prefill. For the fairness of the benchmark, I used context attention in the default path.

I only had a throughput benchmark. As expected chunked prefill greatly improves the e2e throughput. Based on our internal benchmark, it works better as there is more load.

Note that chunked prefill has innately slower e2e latency as called out in this blog.

• ITL (inter token latency) as decoding tokens are batched with prefill tokens. We alleviated it by always batching decoding tokens first when scheduling.
• TTFT is supposed to be a bit slower because we chunk prefill (which adds per step overhead and KV cache memory load overhead).

We can alleviate this by not chunking prefill when there are no decoding requests.

In Anyscale’s internal benchmark we found it slower when QPS is 1 specifically. We will need more tuning and profiling for the optimal performance in this case.

The benchmark was performed on A100, 80GB, TP8, llama 70B

python benchmark_throughput.py --backend vllm --model "meta-llama/Llama-2-70b-hf" --dataset ../../data/ShareGPT_V3_unfiltered_cleaned_split.json --num-prompts 2000 -tp 8

Default OSS (on a 1dquery branch)
Throughput: 6.61 requests/s, 3198.02 tokens/s

Default OSS + use context attention (on a 1dquery branch)
Throughput: 6.51 requests/s, 3149.89 tokens/s

Chunk size 512 (cuda graph is decoding only)
Throughput: 7.72 requests/s, 3736.73 tokens/s

Chunk size 768 (cuda graph is decoding only)
Throughput: 7.51 requests/s, 3632.23 tokens/s

[njhill]

In Anyscale’s internal benchmark we found it slower when QPS is 1 specifically. We will need more tuning and profiling for the optimal performance in this case.

@rkooo567 does this mean one request at a time i.e. batch size 1? I would expect there to be no latency impact for this case if we're not chunking prefill when there are no decoding requests? Or does QPS 1 still entail overlapping requests since they take > 1 sec each?

[rkooo567]

It means 1 request per 1 second. So if e2e takes more than 1 second, BS can be > 1.

We haven't taken a deep look here (I think we will after this PR is merged to oss) for the root cause. But it could be becasue for longer prompt they are chunked when they don't have to be when QPS == 1.

[rkooo567]

Also, depending on scheudling strategy, you can optimize perf/throughput. For example, if you only allow to chunk 1 prompt per batch, it can actually improve ITL (because there's less interrupt from longer prompt for decoding) vs no chunked prefill (in our internal benchmark, with this strategy, we saw QPS>=2 always much better than non chunked prefill). Or if you maximize chunk batches, it can maximize throughput (which is the benchmark above), but it may have TTFT, ITL penalty