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, 23:45 UTC
update flashinfer to v0.2.9rc1 (#21485)
Signed-off-by: Weiliang Liu <weiliangl@nvidia.com>
weiliang committed
Jul 25, 2025 at 05:06 UTC
2dd72d23d965c6951e14029c9540e0b0e57bc9ee
3 files changed
+6
-15
docker/Dockerfile
+1
-1
@@ -386,7 +386,7 @@ RUN --mount=type=bind,from=build,src=/workspace/dist,target=/vllm-workspace/dist
386
387
# Install FlashInfer from source
388
ARG FLASHINFER_GIT_REPO="https://github.com/flashinfer-ai/flashinfer.git"
389
-ARG FLASHINFER_GIT_REF="v0.2.8"
389
+ARG FLASHINFER_GIT_REF="v0.2.9rc1"
390
RUN --mount=type=cache,target=/root/.cache/uv bash - <<'BASH'
391
. /etc/environment
392
git clone --depth 1 --recursive --shallow-submodules \
vllm/attention/backends/flashinfer.py
+3
-7
@@ -1169,16 +1169,12 @@ class FlashInferImpl(AttentionImpl):
1169
query=decode_query,
1170
kv_cache=kv_cache.permute(*stride_order),
1171
workspace_buffer=workspace_buffer,
1172
- num_heads=num_heads,
1173
- num_kv_heads=num_kv_heads,
1174
- scale=softmax_scale,
1172
block_tables=attn_metadata.block_tables,
1173
seq_lens=decode_meta.seq_lens_tensor,
1177
- block_size=attn_metadata.page_size,
1174
max_seq_len=attn_metadata.max_decode_seq_len,
1179
- kv_cache_dtype=kv_cache_dtype,
1180
- k_scale=layer._k_scale_float,
1181
- v_scale=layer._v_scale_float)
1175
+ bmm1_scale=layer._k_scale_float * softmax_scale,
1176
+ bmm2_scale=layer._v_scale_float,
1177
+ )
1178
1179
if prefill_output is None and decode_output is not None:
1180
# Decode only batch.
vllm/v1/attention/backends/flashinfer.py
+2
-7
@@ -678,15 +678,10 @@ class FlashInferImpl(AttentionImpl):
678
query=decode_query,
679
kv_cache=kv_cache_permute,
680
workspace_buffer=attn_metadata.workspace_buffer,
681
- num_heads=self.num_heads,
682
- num_kv_heads=self.num_kv_heads,
683
- scale=self.scale,
681
block_tables=block_tables_decode,
682
seq_lens=seq_lens_decode,
686
- block_size=attn_metadata.page_size,
683
max_seq_len=attn_metadata.max_seq_len,
688
- kv_cache_dtype=self.kv_cache_dtype,
689
- k_scale=layer._k_scale_float,
690
- v_scale=layer._v_scale_float,
684
+ bmm1_scale=layer._k_scale_float * self.scale,
685
+ bmm2_scale=layer._v_scale_float,
686
))
687
return output_padded