A high-throughput and memory-efficient inference and serving engine for LLMs
Mirrored from https://github.com/tduncan-lum/Nvidia-Optimization-Stack---VLLM.git
Synced Nov 14, 05:30 UTC
Bump Flashinfer to v0.4.0 (#26326)
Signed-off-by: elvischenv <219235043+elvischenv@users.noreply.github.com>
elvischenv committed
Oct 9, 2025 at 14:58 UTC
5e49c3e777b59312890da254c71ba22a84f28553
7 files changed
+25
-23
docker/Dockerfile
+3
-3
@@ -15,7 +15,7 @@ ARG PYTHON_VERSION=3.12
15
# Example:
16
# docker build --build-arg BUILD_BASE_IMAGE=registry.acme.org/mirror/nvidia/cuda:${CUDA_VERSION}-devel-ubuntu20.04
17
18
-# Important: We build with an old version of Ubuntu to maintain broad
18
+# Important: We build with an old version of Ubuntu to maintain broad
19
# compatibility with other Linux OSes. The main reason for this is that the
20
# glibc version is baked into the distro, and binaries built with one glibc
21
# version are not backwards compatible with OSes that use an earlier version.
@@ -371,7 +371,7 @@ RUN --mount=type=bind,from=build,src=/workspace/dist,target=/vllm-workspace/dist
371
# Install FlashInfer from source
372
ARG FLASHINFER_GIT_REPO="https://github.com/flashinfer-ai/flashinfer.git"
373
# Keep this in sync with "flashinfer" extra in setup.py
374
-ARG FLASHINFER_GIT_REF="v0.3.1"
374
+ARG FLASHINFER_GIT_REF="v0.4.0"
375
# Flag to control whether to compile FlashInfer AOT kernels
376
# Set to "true" to enable AOT compilation:
377
# docker build --build-arg FLASHINFER_AOT_COMPILE=true ...
@@ -392,7 +392,7 @@ RUN --mount=type=cache,target=/root/.cache/uv bash - <<'BASH'
392
FI_TORCH_CUDA_ARCH_LIST="7.5 8.0 8.9 9.0a 10.0a 12.0"
393
fi
394
pushd flashinfer
395
- if [[ "${CUDA_VERSION}" == 12.8.* ]] && [ "$TARGETPLATFORM" = "linux/amd64" ]; then
395
+ if [[ "${CUDA_VERSION}" == 12.8.* ]] && [ "$TARGETPLATFORM" = "linux/amd64" ] && [ "${FLASHINFER_GIT_REF}" = "v0.3.1" ]; then
396
# NOTE: To make new precompiled wheels, see tools/flashinfer-build.sh
397
echo "🏗️ Installing FlashInfer from pre-compiled wheel"
398
uv pip install --system https://wheels.vllm.ai/flashinfer-python/flashinfer_python-0.3.1-cp39-abi3-manylinux1_x86_64.whl \
docker/Dockerfile.nightly_torch
+2
-2
@@ -246,7 +246,7 @@ RUN pip install setuptools==75.6.0 packaging==23.2 ninja==1.11.1.3 build==1.2.2.
246
247
248
# build flashinfer for torch nightly from source around 10 mins
249
-# release version: v0.3.1
249
+# release version: v0.4.0
250
# todo(elainewy): cache flashinfer build result for faster build
251
ENV CCACHE_DIR=/root/.cache/ccache
252
RUN --mount=type=cache,target=/root/.cache/ccache \
@@ -254,7 +254,7 @@ RUN --mount=type=cache,target=/root/.cache/ccache \
254
echo "git clone flashinfer..." \
255
&& git clone --recursive https://github.com/flashinfer-ai/flashinfer.git \
256
&& cd flashinfer \
257
- && git checkout v0.3.1 \
257
+ && git checkout v0.4.0 \
258
&& git submodule update --init --recursive \
259
&& echo "finish git clone flashinfer..." \
260
&& rm -rf build \
setup.py
+1
-1
@@ -715,7 +715,7 @@ setup(
715
], # Required for audio processing
716
"video": [], # Kept for backwards compatibility
717
# FlashInfer should be updated together with the Dockerfile
718
- "flashinfer": ["flashinfer-python==0.3.1"],
718
+ "flashinfer": ["flashinfer-python==0.4.0"],
719
# Optional deps for AMD FP4 quantization support
720
"petit-kernel": ["petit-kernel"],
721
},
tests/kernels/attention/test_flashinfer_trtllm_attention.py
+9
-12
@@ -7,9 +7,8 @@ import pytest
7
import torch
8
9
from tests.kernels.quantization.nvfp4_utils import (
10
- FLOAT4_E2M1_MAX,
11
- FLOAT8_E4M3_MAX,
10
dequantize_nvfp4_to_dtype,
11
+ get_nvfp4_global_scale,
12
)
13
from vllm.platforms import current_platform
14
from vllm.utils import round_up
@@ -171,13 +170,12 @@ def test_flashinfer_trtllm_decode_with_baseline(
170
output = torch.empty(ref_query.shape, dtype=dtype)
171
wrapper.run(ref_query, ref_kv_cache, out=output)
172
o_scale = 1.0
174
- o_sf_scale = None
173
+ o_sf_scale_float = None
174
if o_quant_dtype == FP8_DTYPE:
175
_, o_scale = to_float8(output)
176
elif o_quant_dtype == FP4_DTYPE:
178
- o_sf_scale = (
179
- (FLOAT8_E4M3_MAX * FLOAT4_E2M1_MAX) / torch.amax(output.flatten(), dim=-1)
180
- ).to(torch.float32)
177
+ o_sf_scale = get_nvfp4_global_scale(output)
178
+ o_sf_scale_float = o_sf_scale.item()
179
180
# TRTLLM Decode
181
if o_quant_dtype == FP4_DTYPE:
@@ -204,7 +202,7 @@ def test_flashinfer_trtllm_decode_with_baseline(
202
bmm1_scale=q_scale * k_scale * sm_scale,
203
bmm2_scale=v_scale / o_scale,
204
window_left=window_left,
207
- o_sf_scale=o_sf_scale,
205
+ o_sf_scale=o_sf_scale_float,
206
out=output_trtllm,
207
)
208
if o_quant_dtype == FP8_DTYPE:
@@ -361,13 +359,12 @@ def test_flashinfer_trtllm_prefill_with_baseline(
359
output = torch.empty(ref_query.shape, dtype=dtype)
360
wrapper.run(ref_query, ref_kv_cache, out=output)
361
o_scale = 1.0
364
- o_sf_scale = None
362
+ o_sf_scale_float = None
363
if o_quant_dtype == FP8_DTYPE:
364
_, o_scale = to_float8(output)
365
elif o_quant_dtype == FP4_DTYPE:
368
- o_sf_scale = (
369
- (FLOAT8_E4M3_MAX * FLOAT4_E2M1_MAX) / torch.amax(output.flatten(), dim=-1)
370
- ).to(torch.float32)
366
+ o_sf_scale = get_nvfp4_global_scale(output)
367
+ o_sf_scale_float = o_sf_scale.item()
368
369
# TRTLLM Prefill
370
if o_quant_dtype == FP4_DTYPE:
@@ -398,7 +395,7 @@ def test_flashinfer_trtllm_prefill_with_baseline(
395
cum_seq_lens_q=q_indptr,
396
cum_seq_lens_kv=kv_indptr,
397
window_left=window_left,
401
- o_sf_scale=o_sf_scale,
398
+ o_sf_scale=o_sf_scale_float,
399
out=output_trtllm,
400
)
401
if o_quant_dtype == FP8_DTYPE:
tests/kernels/quantization/nvfp4_utils.py
+5
-3
@@ -66,9 +66,11 @@ def break_fp4_bytes(a, dtype):
66
return values.reshape(m, n * 2).to(dtype=dtype)
67
68
69
+def get_nvfp4_global_scale(a: torch.Tensor):
70
+ return (FLOAT8_E4M3_MAX * FLOAT4_E2M1_MAX) / torch.abs(a).max().to(torch.float32)
71
+
72
+
73
def quant_nvfp4_tensor(a: torch.Tensor):
70
- a_global_scale = (FLOAT8_E4M3_MAX * FLOAT4_E2M1_MAX) / torch.abs(a).max().to(
71
- torch.float32
72
- )
74
+ a_global_scale = get_nvfp4_global_scale(a)
75
a_quant, a_block_scale = scaled_fp4_quant(a, a_global_scale)
76
return a_quant, a_block_scale, a_global_scale
tests/quantization/test_blackwell_moe.py
+1
-1
@@ -50,7 +50,7 @@ def can_initialize(model: str, extra_args: Optional[list[str]] = None):
50
with RemoteOpenAIServer(
51
model,
52
server_args,
53
- max_wait_seconds=1000, # Due to FlashInfer compile
53
+ max_wait_seconds=1500, # Due to FlashInfer compile
54
override_hf_configs=dummy_hf_overrides,
55
) as server:
56
client = server.get_client()
vllm/v1/attention/backends/flashinfer.py
+4
-1
@@ -1199,7 +1199,7 @@ def fast_plan_decode(
1199
qo_indptr_host = _get_range_buf(batch_size + 1, "cpu")
1200
1201
try:
1202
- # Make sure we pass exactly 15 arguments for tensor core version
1202
+ # Make sure we pass exactly 18 arguments for tensor core version
1203
self._plan_info = self._cached_module.plan(
1204
self._float_workspace_buffer,
1205
self._int_workspace_buffer,
@@ -1216,6 +1216,9 @@ def fast_plan_decode(
1216
head_dim,
1217
head_dim,
1218
False, # causal
1219
+ window_left,
1220
+ -1, # fixed_split_size
1221
+ False, # disable_split_kv
1222
)
1223
except Exception as e:
1224
raise RuntimeError(f"Error in tensor core plan: {e}") from e