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 5, 07:15 UTC
[Bugfix] Fix FlashInfer crash with kv_cache_dtype_skip_layers (#39002)
Signed-off-by: Yifan Zong <yzong@redhat.com>
yzong-rh committed
Apr 10, 2026 at 14:50 UTC
e816a8811f2fbe9a8fa7f5d5f37a25cd410805d8
2 files changed
+15
-11
tests/compile/passes/test_fusion_attn.py
+5
-8
@@ -39,7 +39,7 @@ from vllm.platforms import current_platform
39
from vllm.utils.flashinfer import has_flashinfer
40
from vllm.v1.attention.backend import AttentionMetadata
41
from vllm.v1.attention.backends.registry import AttentionBackendEnum
42
-from vllm.v1.kv_cache_interface import AttentionSpec
42
+from vllm.v1.kv_cache_interface import AttentionSpec, get_kv_quant_mode
43
44
FP8_DTYPE = current_platform.fp8_dtype()
45
FP4_DTYPE = torch.uint8
@@ -53,7 +53,6 @@ class AttentionQuantPatternModel(torch.nn.Module):
53
num_qo_heads: int,
54
num_kv_heads: int,
55
head_size: int,
56
- kv_cache_dtype: torch.dtype,
56
device: torch.device,
57
vllm_config: VllmConfig,
58
block_size: int,
@@ -63,7 +62,6 @@ class AttentionQuantPatternModel(torch.nn.Module):
62
self.num_qo_heads = num_qo_heads
63
self.num_kv_heads = num_kv_heads
64
self.head_size = head_size
66
- self.kv_cache_dtype = kv_cache_dtype
65
self.device = device
66
self.vllm_config = vllm_config
67
self.dtype = vllm_config.model_config.dtype
@@ -81,13 +79,14 @@ class AttentionQuantPatternModel(torch.nn.Module):
79
80
self.block_size = block_size
81
84
- # Initialize attn MetadataBuilder
82
+ # Initialize attn MetadataBuilder (match Attention.get_kv_cache_spec)
83
self.builder = self.attn.attn_backend.get_builder_cls()(
84
kv_cache_spec=AttentionSpec(
85
block_size=self.block_size,
86
num_kv_heads=self.num_kv_heads,
87
head_size=self.head_size,
90
- dtype=self.kv_cache_dtype,
88
+ dtype=self.attn.kv_cache_torch_dtype,
89
+ kv_quant_mode=get_kv_quant_mode(self.attn.kv_cache_dtype),
90
),
91
layer_names=[self.attn.layer_name],
92
vllm_config=self.vllm_config,
@@ -126,7 +125,7 @@ class AttentionQuantPatternModel(torch.nn.Module):
125
# Create dummy KV cache
126
raw_tensor = torch.zeros(
127
2 * num_blocks * self.block_size * self.num_kv_heads * self.head_size,
129
- dtype=self.kv_cache_dtype,
128
+ dtype=self.attn.kv_cache_torch_dtype,
129
device=self.device,
130
)
131
raw_tensor = raw_tensor.view(kv_cache_shape)
@@ -348,7 +347,6 @@ def test_attention_quant_pattern(
347
num_qo_heads=num_qo_heads,
348
num_kv_heads=num_kv_heads,
349
head_size=head_size,
351
- kv_cache_dtype=FP8_DTYPE,
350
device=device,
351
vllm_config=vllm_config_unfused,
352
block_size=block_size,
@@ -376,7 +374,6 @@ def test_attention_quant_pattern(
374
num_qo_heads=num_qo_heads,
375
num_kv_heads=num_kv_heads,
376
head_size=head_size,
379
- kv_cache_dtype=FP8_DTYPE,
377
device=device,
378
vllm_config=vllm_config,
379
w=model_unfused.w,
vllm/v1/attention/backends/flashinfer.py
+10
-3
@@ -63,7 +63,11 @@ from vllm.v1.attention.backends.utils import (
63
from vllm.v1.attention.ops.common import cp_lse_ag_out_rs
64
from vllm.v1.attention.ops.dcp_alltoall import dcp_a2a_lse_reduce
65
from vllm.v1.attention.ops.merge_attn_states import merge_attn_states
66
-from vllm.v1.kv_cache_interface import AttentionSpec, UniformTypeKVCacheSpecs
66
+from vllm.v1.kv_cache_interface import (
67
+ AttentionSpec,
68
+ KVQuantMode,
69
+ UniformTypeKVCacheSpecs,
70
+)
71
from vllm.v1.utils import CpuGpuBuffer
72
73
FLASHINFER_WORKSPACE_BUFFER_SIZE_BATCH_INVARIANT = 2048 * 1024 * 1024
@@ -600,12 +604,15 @@ class FlashInferMetadataBuilder(AttentionMetadataBuilder[FlashInferMetadata]):
604
self.head_dim = self.kv_cache_spec.head_size
605
self.page_size = self.kv_cache_spec.block_size
606
603
- self.cache_dtype = self.cache_config.cache_dtype
604
- if is_quantized_kv_cache(self.cache_dtype):
607
+ if self.kv_cache_spec.kv_quant_mode != KVQuantMode.NONE:
608
+ self.cache_dtype = self.cache_config.cache_dtype
609
+ # Cannot use self.kv_cache_spec.dtype here because kv_cache_spec
610
+ # storage dtype may not be the same as the op dtype (uint8 vs fp8_e4m3)
611
self.kv_cache_dtype = FlashInferBackend.get_fp8_dtype_for_flashinfer(
612
self.cache_dtype
613
)
614
else:
615
+ self.cache_dtype = "auto"
616
assert self.kv_cache_spec.dtype == self.model_config.dtype
617
self.kv_cache_dtype = self.kv_cache_spec.dtype
618