[BUG] fix crash on flashinfer backend with cudagraph disabled, when attention group_size not in [1,2,4,8]

[learninmou]

when I use flashinfer backend and disable cuda graph, load a model with attention group_size=6, vllm crashs and shows the following log:

This error consistently occurs under the following conditions：

1. user use flashinfer attention backend explicitly, (set env VLLM_ATTENTION_BACKEND="FLASHINFER")
2. the model use GQA, and group size not in [1,2,4,8]. (I've reviewed flashinfer's source, it only support group_size [1,2,4,8] otherwise user need to use_tensor_cores=True to do decode)
3. cuda graph is disabled, it can happen when user set enforce_eager=True, or the decoding batch_size > _BATCH_SIZES_TO_CAPTURE[-1] (variable in model_runner.py, currently 256)

[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 consists a small and essential subset of CI tests to quickly catch errors. You can run other CI tests on top of default ones by unblocking the steps in your fast-check build on Buildkite UI.

Once the PR is approved and ready to go, please make sure to run full CI as it is required to merge (or just use auto-merge).

To run full CI, you can do one of these:

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

🚀

[comaniac]

LGTM. Can you add a unit test?

[learninmou]

LGTM. Can you add a unit test?

add utest case in the last commit which can cover this bug if user do not set BatchDecodeWithPagedKVCacheWrapper properly.

[learninmou]

/ready

[JaheimLee]

When will this pr be merged?

[comaniac]

@learninmou could you rebase to see if CI failure is fixed in main branch already?

[learninmou]

@learninmou could you rebase to see if CI failure is fixed in main branch already?

rebase finished, all checks have passed

[elfiegg]

Hi, we recently discover that this PR caused 11.4% perf regression:
Before: 5233.12 tokens/s
After: 4697.55 tokens/s

Tested on H200 machines with below command:
python benchmarks/benchmark_throughput.py --model=neuralmagic/Meta-Llama-3-70B-Instruct-FP8-KV \ --num-prompts=1024 --output-len=256 --input-len=256 --quantization=fp8 \ --max-num-batched-tokens=8192 --max-num-seqs=1024

Have we done perf test on this PR?
If this is the case - Can we revert it until further investigations are completed?

[comaniac]

Hmm this PR shouldn't introduce performance regression to existing workloads as it attempts to fix uncovered cases. Is your vLLM benchmark before and after exactly this PR with the same environment? If so is it possible for you to identify the root cause? It could be one of the following I could think of:

1. The use of use_tensor_cores in FlashInfer kernel.
2. The overhead of the logic checking the group size.

Also which model and GPU you're benchmarking? Thanks.

[elfiegg]

I'm not exactly sure why this would cause perf regression yet. But I reran the benchmark, it shows about 18% slowdown:

Before: Throughput: 10.56 requests/s, 5409.15 tokens/s
After: Throughput: 8.91 requests/s, 4563.67 tokens/s

Benchmark is performed on Llama3-70B on H200.
Command:
python benchmarks/benchmark_throughput.py --model=neuralmagic/Meta-Llama-3-70B-Instruct-FP8 --num-prompts=768 --output-len=256 --input-len=256 --quantization=fp8 --kv-cache-dtype=fp8 --gpu-memory-utilization=0.98 --max-num-batched-tokens=2048 --max-num-seqs=768 --max_model_len=512

Is it possible that we revert it for now until further investigations?

[comaniac]

Ok I probably know the reason. Before this PR:

use_tensor_cores = num_qo_heads // num_kv_heads >= 4

After this PR:

use_tensor_cores = (num_qo_heads // num_kv_heads) not in (1, 2, 4, 8)

• Llama-3-8B: qo_heads (32) // kv_heads (8) = 4. Before this PR we set use_tensor_cores=True but after this PR it sets to False.
• Llama-3-70B: qo_heads (64) // kv_heads (8) = 8. Before this PR we set use_tensor_cores=True but after this PR it sets to False.

I'll file a PR to fix this and you could test the PR to see if that helps.

[yzh119]

Sorry I just noticed this PR, this fix will degrade performance. The reason flashinfer compiles kernel with use_tensor_cores=False for group_size in (1,2,4,8) is not related to performance, it's because they are most popular group sizes and I keep the group_size 8 for use_tensor_cores=False mainly for correctness test. We should always enable tensor cores for large group size (more specifically, when group_size > 4), I should make it more clear in flashinfer documentation.

[yzh119]

It will be great if you can set me as a reviewer for pull requests related to flashinfer, I always miss mentions because of huge amount of notifications..