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, 11:00 UTC
[Bugfix] Fix extract_hidden_states crash with quantized KV cache dtype (#39160)
Signed-off-by: Yubo Wang <yubowang2019@gmail.com>
Yubo Wang committed
Apr 7, 2026 at 11:18 UTC
08bfedc152f064d8e84f85c4f42b810e5a564229
1 file changed
+5
vllm/model_executor/models/extract_hidden_states.py
+5
@@ -9,6 +9,7 @@ extract_hidden_states speculative decoding method.
9
"""
10
11
from collections.abc import Iterable
12
+from dataclasses import replace
13
from typing import ClassVar
14
15
import torch
@@ -352,6 +353,10 @@ class ExtractHiddenStatesModel(nn.Module):
353
354
cache_config = vllm_config.cache_config
355
356
+ # Hidden states dtype should be independent of KV cache dtype.
357
+ if cache_config is not None and is_quantized_kv_cache(cache_config.cache_dtype):
358
+ cache_config = replace(cache_config, cache_dtype="auto")
359
+
360
# Create a single cache-only attention layer
361
# Note: We set num_heads <- self.num_hidden_states
362
# and head_size <- hidden_size so that we can insert