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
[bugfix] Fix Llama3/4 issues caused by FlashInfer 0.2.10 (#22426)
Signed-off-by: Po-Han Huang <pohanh@nvidia.com>
Po-Han Huang (NVIDIA) committed
Aug 8, 2025 at 11:25 UTC
af473f0a85731c17d9cf708deec3e864e674feb0
2 files changed
+18
-9
vllm/model_executor/layers/quantization/utils/flashinfer_utils.py
+16
-8
@@ -6,14 +6,22 @@ import torch
6
7
8
def calculate_tile_tokens_dim(num_tokens, top_k, num_experts):
9
- from flashinfer import next_positive_power_of_2
10
-
11
- # Guess tokens per expert assuming perfect expert distribution first.
12
- num_tokens_per_expert = (num_tokens * top_k) // num_experts
13
- # And pad the number to the next power of 2.
14
- tile_tokens_dim = next_positive_power_of_2(num_tokens_per_expert)
15
- # Cap to 8-64 tokens per CTA tile as it's the range supported by the kernel.
16
- tile_tokens_dim = min(max(tile_tokens_dim, 8), 64)
9
+
10
+ # FlashInfer 0.2.10 has issues with larger tile sizes. Set to 8 for now.
11
+ # TODO: Revert this to dynamic calculation once a new version of FlashInfer
12
+ # with the necessary kernels is released.
13
+ tile_tokens_dim = 8
14
+
15
+ # from flashinfer import next_positive_power_of_2
16
+
17
+ # # Guess tokens per expert assuming perfect expert distribution first.
18
+ # num_tokens_per_expert = (num_tokens * top_k) // num_experts
19
+ # # And pad the number to the next power of 2.
20
+ # tile_tokens_dim = next_positive_power_of_2(num_tokens_per_expert)
21
+ # # Cap to 8-64 tokens per CTA tile as it's the range supported by the
22
+ # # kernel.
23
+ # tile_tokens_dim = min(max(tile_tokens_dim, 8), 64)
24
+
25
return tile_tokens_dim
26
27
vllm/v1/attention/backends/flashinfer.py
+2
-1
@@ -524,7 +524,8 @@ class FlashInferMetadataBuilder(AttentionMetadataBuilder[FlashInferMetadata]):
524
head_dim = self.kv_cache_spec.head_size
525
526
# currently prefill trtllm attention does not support fp8 kv cache
527
- prefill_use_trtllm = use_trtllm_attention(
527
+ prefill_use_trtllm = not cache_dtype.startswith("fp8") \
528
+ and use_trtllm_attention(
529
num_prefill_tokens, max_seq_len, cache_dtype,
530
num_qo_heads, num_kv_heads, head_dim)
531
decode_use_trtllm = use_trtllm_attention(