[Sampler] Adapt to FlashInfer 0.2.3 sampler API

[abmfy]

FlashInfer 0.2.3 introduced some breaking changes to its sampler API, this PR updates the calling sites in vLLM to adapt to the update.

FIX #14815
FIX #15666

[github-actions]

👋 Hi! Thank you for contributing to the vLLM project.

💬 Join our developer Slack at https://slack.vllm.ai to discuss your PR in #pr-reviews, coordinate on features in #feat- channels, or join special interest groups in #sig- channels.

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 either: Add ready label to the PR or enable auto-merge.

🚀

[yzh119]

Hi @youkaichao , are failed test related to this PR?

I saw some error message like this:

[2025-03-30T18:53:05Z] INFO 03-30 11:53:05 [backends.py:144] Compiling a graph for general shape takes 25.10 s
--
  | [2025-03-30T18:53:05Z] DEBUG 03-30 11:53:05 [backends.py:469] Computation graph saved to /root/.cache/vllm/torch_compile_cache/da20f97f50/rank_0_0/computation_graph.py
  | [2025-03-30T18:53:05Z] DEBUG 03-30 11:53:05 [wrapper.py:105] Dynamo transformed code saved to /root/.cache/vllm/torch_compile_cache/da20f97f50/rank_0_0/transformed_code.py
  | [2025-03-30T18:53:16Z] INFO 03-30 11:53:16 [monitor.py:33] torch.compile takes 30.91 s in total
  | [2025-03-30T18:53:17Z] ERROR 03-30 11:53:17 [core.py:367] EngineCore hit an exception: Traceback (most recent call last):
  | [2025-03-30T18:53:17Z] ERROR 03-30 11:53:17 [core.py:367]   File "/usr/local/lib/python3.12/dist-packages/vllm/v1/engine/core.py", line 355, in run_engine_core
  | [2025-03-30T18:53:17Z] ERROR 03-30 11:53:17 [core.py:367]     engine_core = EngineCoreProc(*args, **kwargs)
  | [2025-03-30T18:53:17Z] ERROR 03-30 11:53:17 [core.py:367]                   ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  | [2025-03-30T18:53:17Z] ERROR 03-30 11:53:17 [core.py:367]   File "/usr/local/lib/python3.12/dist-packages/vllm/v1/engine/core.py", line 296, in __init__
  | [2025-03-30T18:53:17Z] ERROR 03-30 11:53:17 [core.py:367]     super().__init__(vllm_config, executor_class, log_stats)
  | [2025-03-30T18:53:17Z] ERROR 03-30 11:53:17 [core.py:367]   File "/usr/local/lib/python3.12/dist-packages/vllm/v1/engine/core.py", line 67, in __init__
  | [2025-03-30T18:53:17Z] ERROR 03-30 11:53:17 [core.py:367]     num_gpu_blocks, num_cpu_blocks = self._initialize_kv_caches(
  | [2025-03-30T18:53:17Z] ERROR 03-30 11:53:17 [core.py:367]                                      ^^^^^^^^^^^^^^^^^^^^^^^^^^^
  | [2025-03-30T18:53:17Z] ERROR 03-30 11:53:17 [core.py:367]   File "/usr/local/lib/python3.12/dist-packages/vllm/v1/engine/core.py", line 133, in _initialize_kv_caches
  | [2025-03-30T18:53:17Z] ERROR 03-30 11:53:17 [core.py:367]     get_kv_cache_config(vllm_config, kv_cache_spec_one_worker,
  | [2025-03-30T18:53:17Z] ERROR 03-30 11:53:17 [core.py:367]   File "/usr/local/lib/python3.12/dist-packages/vllm/v1/core/kv_cache_utils.py", line 615, in get_kv_cache_config
  | [2025-03-30T18:53:17Z] ERROR 03-30 11:53:17 [core.py:367]     check_enough_kv_cache_memory(vllm_config, kv_cache_spec, available_memory)
  | [2025-03-30T18:53:17Z] ERROR 03-30 11:53:17 [core.py:367]   File "/usr/local/lib/python3.12/dist-packages/vllm/v1/core/kv_cache_utils.py", line 489, in check_enough_kv_cache_memory
  | [2025-03-30T18:53:17Z] ERROR 03-30 11:53:17 [core.py:367]     raise ValueError(
  | [2025-03-30T18:53:17Z] ERROR 03-30 11:53:17 [core.py:367] ValueError: To serve at least one request with the models's max seq len (8192), (0.81 GB KV cache is needed, which is larger than the available KV cache memory (0.55 GB). Try increasing `gpu_memory_utilization` or decreasing `max_model_len` when initializing the engine.
  | [2025-03-30T18:53:17Z] ERROR 03-30 11:53:17 [core.py:367]
  | [2025-03-30T18:53:17Z] CRITICAL 03-30 11:53:17 [core_client.py:319] Got fatal signal from worker processes, shutting down. See stack trace above for root cause issue.
  | [2025-03-30T18:53:17Z] bash: line 1:   510 Killed                  pytest -v -s basic_correctness/test_basic_correctness.py

[youkaichao]

do you need to update the dockerfile to install the new flashinfer wheel?

[abmfy]

do you need to update the dockerfile to install the new flashinfer wheel?

Let me have a try

[mergify]

This pull request has merge conflicts that must be resolved before it can be
merged. Please rebase the PR, @abmfy.

https://docs.github.com/en/pull-requests/collaborating-with-pull-requests/working-with-forks/syncing-a-fork

[mgoin]

Can we revive this? I would like to update flashinfer to latest now that we have it integrated with V1 as an attention backend

[WoosukKwon]

@mgoin +1. Let's update this.

@abmfy Sorry for the late review. Could you please add an accuracy test for the new kernel?

[abmfy]

Sure, I'll add some tests soon

[chenyang78]

FYI - synced with @mgoin @luccafong @houseroad offline, there was some numerical issue with flashinfer 0.2.5. Fortunately, we found that the issue was fixed already by the upstream (flashinfer-ai/flashinfer#1029)

before the fix:

$ VLLM_USE_V1=1 VLLM_ATTENTION_BACKEND=FLASHINFER lm_eval --model vllm --model_args pretrained=meta-llama/Llama-3.1-8B-Instruct --trust_remote_code --tasks gsm8k --num_fewshot 5 --batch_size auto
...
vllm (pretrained=meta-llama/Llama-3.1-8B-Instruct,trust_remote_code=True), gen_kwargs: (None), limit: None, num_fewshot: 5, batch_size: auto
|Tasks|Version|     Filter     |n-shot|  Metric   |   |Value |   |Stderr|
|-----|------:|----------------|-----:|-----------|---|-----:|---|-----:|
|gsm8k|      3|flexible-extract|     5|exact_match|↑  |0.0675|±  |0.0069|
|     |       |strict-match    |     5|exact_match|↑  |0.0629|±  |0.0067|

after the fix:

$ VLLM_USE_V1=1 VLLM_ATTENTION_BACKEND=FLASHINFER lm_eval --model vllm --model_args pretrained=meta-llama/Llama-3.1-8B-Instruct --trust_remote_code --tasks gsm8k --num_fewshot 5 --batch_size auto
...
|Tasks|Version|     Filter     |n-shot|  Metric   |   |Value |   |Stderr|
|-----|------:|----------------|-----:|-----------|---|-----:|---|-----:|
|gsm8k|      3|flexible-extract|     5|exact_match|↑  |0.7892|±  |0.0112|
|     |       |strict-match    |     5|exact_match|↑  |0.7650|±  |0.0117|

[mgoin]

Hi @abmfy do you plan to have test updates soon? We can help make them if you don't have time right now

[abmfy]

Hi @abmfy do you plan to have test updates soon? We can help make them if you don't have time right now

Hi, sorry for the delayed response. I’ve just settled in the U.S. recently and was dealing with limited bandwidth. I’ll focus on adding the test today.

[mergify]

This pull request has merge conflicts that must be resolved before it can be
merged. Please rebase the PR, @abmfy.

https://docs.github.com/en/pull-requests/collaborating-with-pull-requests/working-with-forks/syncing-a-fork

[WoosukKwon]

LGTM. Thanks a lot for the PR!

[mgoin]

I don't think we can use this wheel since it is built for cu124torch2.6. We need cu128 and torch 2.8

[abmfy]

Sure. However, FlashInfer 0.2.4 only provides wheels up to cu124torch2.6, so we may need to build from source in CI for now—at least until a new release of FlashInfer becomes available.

[abmfy]

Some sampler tests of v0 seem to be failing due to the removal of the ability to pass pre-generated uniform samples to FlashInfer kernels.

v1 sampler tests all passed. Will get back to this later.

[simon-mo]

@abmfy Thank you can you fix this PR? This is part of the release blocker now.

[abmfy]

@abmfy Thank you can you fix this PR? This is part of the release blocker now.

Sure, I'll fix the tests that seems to be related to this PR

[DarkLight1337]

Thanks for the contribution, however your PR appears to have broken many tests (see #18462), can those involved fix them to unblock the release?

[abmfy]

Thanks for the contribution, however your PR appears to have broken many tests (see #18462), can those involved fix them to unblock the release?

That's strange, will look into it

[DarkLight1337]

See #18416 #18417 for more info