@tduncan / Nvidia-Optimization-Stack-VLLM Mirror · read-only / commits / 191e3fdaa1
Fork 0

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 4, 22:30 UTC

Update flashinfer to 0.6.8 (#39959)

Signed-off-by: bai <v@gor.io>

bai committed Apr 20, 2026 at 20:37 UTC 191e3fdaa1fd3dd09441e7b22d4f2ddef51c012c
9 files changed +31 -61
docker/Dockerfile
+3 -24
@@ -580,33 +580,12 @@ RUN --mount=type=cache,target=/root/.cache/uv \
580 # Install FlashInfer JIT cache (requires CUDA-version-specific index URL)
581 # https://docs.flashinfer.ai/installation.html
582 # From versions.json: .flashinfer.version
583 -# 0.6.7: CUTLASS 4.4.2 bump, fixes TMA grouped GEMM on SM12x (flashinfer#2798)
584 -# TODO: bump to 0.6.8 when released for NVFP4/MXFP4 group GEMMs on
585 -# SM120/SM121 (RTX 50 / DGX Spark) via flashinfer#2738
586 -ARG FLASHINFER_VERSION=0.6.7
583 +ARG FLASHINFER_VERSION=0.6.8.post1
584 RUN --mount=type=cache,target=/root/.cache/uv \
585 uv pip install --system flashinfer-jit-cache==${FLASHINFER_VERSION} \
586 --extra-index-url https://flashinfer.ai/whl/cu$(echo $CUDA_VERSION | cut -d. -f1,2 | tr -d '.') \
590 - && flashinfer show-config
591 -
592 -# Pre-download FlashInfer TRTLLM BMM headers for air-gapped environments.
593 -# At runtime, MoE JIT compilation downloads these from edge.urm.nvidia.com
594 -# which fails without internet. This step caches them at build time.
595 -RUN python3 <<'PYEOF'
596 -from flashinfer.jit import env as jit_env
597 -from flashinfer.jit.cubin_loader import download_trtllm_headers, get_cubin
598 -from flashinfer.artifacts import ArtifactPath, CheckSumHash
599 -
600 -download_trtllm_headers(
601 - 'bmm',
602 - jit_env.FLASHINFER_CUBIN_DIR / 'flashinfer' / 'trtllm' / 'batched_gemm' / 'trtllmGen_bmm_export',
603 - f'{ArtifactPath.TRTLLM_GEN_BMM}/include/trtllmGen_bmm_export',
604 - ArtifactPath.TRTLLM_GEN_BMM,
605 - get_cubin(f'{ArtifactPath.TRTLLM_GEN_BMM}/checksums.txt', CheckSumHash.TRTLLM_GEN_BMM),
606 -)
607 -
608 -print('FlashInfer TRTLLM BMM headers downloaded successfully')
609 -PYEOF
587 + && flashinfer show-config \
588 + && flashinfer download-cubin
589
590 # ============================================================
591 # OPENAI API SERVER DEPENDENCIES
docker/Dockerfile.nightly_torch
+2 -5
@@ -217,16 +217,13 @@ RUN pip install setuptools==75.6.0 packaging==23.2 ninja==1.11.1.3 build==1.2.2.
217
218
219 # build flashinfer for torch nightly from source around 10 mins
220 -# release version: v0.6.7
221 -# 0.6.7: CUTLASS 4.4.2 bump, fixes TMA grouped GEMM on SM12x (flashinfer#2798)
222 -# TODO: bump to 0.6.8 when released for NVFP4/MXFP4 group GEMMs on
223 -# SM120/SM121 (RTX 50 / DGX Spark) via flashinfer#2738
220 +# release version: v0.6.8.post1
221 # todo(elainewy): cache flashinfer build result for faster build
222 ENV CCACHE_DIR=/root/.cache/ccache
223 RUN --mount=type=cache,target=/root/.cache/ccache \
224 --mount=type=cache,target=/root/.cache/uv \
225 echo "git clone flashinfer..." \
229 - && git clone --depth 1 --branch v0.6.7 --recursive https://github.com/flashinfer-ai/flashinfer.git \
226 + && git clone --depth 1 --branch v0.6.8.post1 --recursive https://github.com/flashinfer-ai/flashinfer.git \
227 && cd flashinfer \
228 && git submodule update --init --recursive \
229 && echo "finish git clone flashinfer..." \
docker/versions.json
+1 -1
@@ -65,7 +65,7 @@
65 "default": "true"
66 },
67 "FLASHINFER_VERSION": {
68 - "default": "0.6.7"
68 + "default": "0.6.8.post1"
69 },
70 "GDRCOPY_CUDA_VERSION": {
71 "default": "12.8"
docs/design/attention_backends.md
+1 -1
@@ -169,7 +169,7 @@ Priority is **1 = highest** (tried first).
169 | ------- | ------- | ------ | --------- | ----------- | ---------- | ---- | --------- | --- | --------------- | ------------ |
170 | `CPU_ATTN` | | fp16, bf16, fp32 | `auto` | Any | 32, 64, 80, 96, 112, 128, 160, 192, 224, 256, 512 | ❌ | ❌ | ❌ | All | N/A |
171 | `FLASHINFER` | Native† | fp16, bf16 | `auto`, `float16`, `bfloat16`, `fp8`, `fp8_e4m3`, `fp8_e5m2` | 16, 32, 64 | 64, 128, 256 | ❌ | ❌ | ✅ | Decoder | 7.x-9.x |
172 -| `FLASHINFER` | TRTLLM† | fp16, bf16 | `auto`, `float16`, `bfloat16`, `fp8`, `fp8_e4m3`, `fp8_e5m2` | 16, 32, 64 | 64, 128, 256 | ✅ | ❌ | ✅ | Decoder | 10.0 |
172 +| `FLASHINFER` | TRTLLM† | fp16, bf16 | `auto`, `float16`, `bfloat16`, `fp8`, `fp8_e4m3`, `fp8_e5m2` | 16, 32, 64 | 64, 128, 256 | ✅ | ❌ | ✅ | Decoder | 10.x |
173 | `FLASH_ATTN` | FA2* | fp16, bf16 | `auto`, `float16`, `bfloat16` | %16 | Any | ❌ | ❌ | ✅ | All | ≥8.0 |
174 | `FLASH_ATTN` | FA3* | fp16, bf16 | `auto`, `float16`, `bfloat16`, `fp8`, `fp8_e4m3`, `fp8_e5m2` | %16 | Any | ✅ | ❌ | ✅ | All | 9.x |
175 | `FLASH_ATTN` | FA4* | fp16, bf16 | `auto`, `float16`, `bfloat16` | %16 | Any | ❌ | ❌ | ✅ | All | ≥10.0 |
requirements/cuda.txt
+2 -2
@@ -9,8 +9,8 @@ torchaudio==2.11.0
9 # These must be updated alongside torch
10 torchvision==0.26.0 # Required for phi3v processor. See https://github.com/pytorch/vision?tab=readme-ov-file#installation for corresponding version
11 # FlashInfer should be updated together with the Dockerfile
12 -flashinfer-python==0.6.7
13 -flashinfer-cubin==0.6.7
12 +flashinfer-python==0.6.8.post1
13 +flashinfer-cubin==0.6.8.post1
14 # Cap nvidia-cudnn-frontend (transitive dep of flashinfer) due to
15 # breaking changes in 1.19.0
16 nvidia-cudnn-frontend>=1.13.0,<1.19.0
tests/kernels/attention/test_use_trtllm_attention.py
+1 -1
@@ -62,7 +62,7 @@ def test_supports_batch_invariant_disables():
62
63 @patch("vllm.envs.VLLM_BATCH_INVARIANT", False)
64 @patch(
65 - "vllm.utils.flashinfer.current_platform.is_device_capability",
65 + "vllm.utils.flashinfer.current_platform.is_device_capability_family",
66 return_value=True,
67 )
68 @patch("vllm.utils.flashinfer.has_nvidia_artifactory", return_value=True)
tests/kernels/moe/test_ocp_mx_moe.py
+17 -15
@@ -548,7 +548,9 @@ def test_trtllm_gen_mxfp4_fused_moe(
548 hidden_states, hidden_states_scale = mxfp8_quantize(
549 hidden_states, is_sf_swizzled_layout=False
550 )
551 - hidden_states_scale = hidden_states_scale.view(torch.float8_e4m3fn).reshape(-1)
551 + hidden_states_scale = hidden_states_scale.view(torch.float8_e4m3fn).reshape(
552 + *hidden_states.shape[:-1], -1
553 + )
554 else:
555 hidden_states_scale = None
556
@@ -595,20 +597,20 @@ def test_trtllm_gen_mxfp4_fused_moe(
597 if beta is not None:
598 beta = torch.full((num_experts,), beta, device=hidden_states.device)
599 tg_result = tg_mxfp4_moe(
598 - router_logits,
599 - topk,
600 - num_experts,
601 - intermediate_size,
602 - hidden_size,
603 - hidden_states,
604 - hidden_states_scale,
605 - w13,
606 - w13_scale,
607 - bias13,
608 - w2,
609 - w2_scale,
610 - bias2,
611 - act_type,
600 + router_logits=router_logits,
601 + topk=topk,
602 + num_experts=num_experts,
603 + intermediate_size=intermediate_size,
604 + hidden_size=hidden_size,
605 + hidden_states=hidden_states,
606 + hidden_states_scale=hidden_states_scale,
607 + w13_weight=w13,
608 + w13_weight_scale=w13_scale,
609 + w13_bias=bias13,
610 + w2_weight=w2,
611 + w2_weight_scale=w2_scale,
612 + w2_bias=bias2,
613 + act_type=act_type,
614 alpha=alpha,
615 beta=beta,
616 limit=limit,
vllm/model_executor/layers/fused_moe/flashinfer_cutlass_moe.py
+1 -8
@@ -130,14 +130,7 @@ class FlashInferExperts(mk.FusedMoEExpertsModular):
130 p.is_device_capability(90)
131 or p.is_device_capability_family(100)
132 or p.is_device_capability_family(110)
133 - or p.is_device_capability(120)
134 - # NOTE: SM121 (DGX Spark) is excluded because the bf16
135 - # unquantized CUTLASS MoE GEMM in flashinfer <= 0.6.7 has no
136 - # Relu2 template instantiation and throws "Invalid activation
137 - # type" on Nemotron-H. Fixed upstream by
138 - # https://github.com/flashinfer-ai/flashinfer/pull/2926
139 - # (merged 2026-04-01, not yet in a stable release); lift this
140 - # restriction once flashinfer >= 0.6.8 is the minimum.
133 + or p.is_device_capability_family(120)
134 )
135 and has_flashinfer_cutlass_fused_moe()
136 )
vllm/utils/flashinfer.py
+3 -4
@@ -305,10 +305,9 @@ def supports_trtllm_attention() -> bool:
305 if envs.VLLM_BATCH_INVARIANT:
306 return False
307
308 - # TRTLLM attention is currently only validated on SM100 (CC 10.0).
309 - # SM103 (GB300) hangs with FlashInfer >= 0.6.7.
310 - # See: https://github.com/flashinfer-ai/flashinfer/issues/2939
311 - return current_platform.is_device_capability(100) and has_nvidia_artifactory()
308 + return (
309 + current_platform.is_device_capability_family(100) and has_nvidia_artifactory()
310 + )
311
312
313 def force_use_trtllm_attention() -> bool | None: