A high-throughput and memory-efficient inference and serving engine for LLMs
[BUG] fix hidden states nan for hybrid attention models (#45849)
Signed-off-by: shanjiaz <hezhao@redhat.com> Co-authored-by: shanjiaz <hezhao@redhat.com>
shanjiaz committed
Jun 17, 2026 at 18:02 UTC
5fd21eb0b291e16dad66da740b0faf50666a78a5
2 files changed
+9
vllm/distributed/kv_transfer/kv_connector/v1/example_hidden_states_connector.py
+7
@@ -150,7 +150,14 @@ class ExampleHiddenStatesConnector(KVConnectorBase_V1, SupportsHMA):
150
151
# Worker-side state (set by register_kv_caches).
152
self._kv_cache: torch.Tensor | None = None
153
+
154
+ # Identify which KV cache group holds the hidden-states layer.
155
self._hs_group_idx: int = 0
156
+ if self._kv_cache_config is not None:
157
+ for i, group in enumerate(self._kv_cache_config.kv_cache_groups):
158
+ if any("cache_only_layers" in n for n in group.layer_names):
159
+ self._hs_group_idx = i
160
+ break
161
# Only TP rank 0 writes hidden states to disk; other TP ranks no-op.
162
# Set in register_kv_caches (after distributed init).
163
self._is_tp_rank_zero: bool = True
vllm/v1/core/single_type_kv_cache_manager.py
+2
@@ -271,6 +271,7 @@ class SingleTypeKVCacheManager(ABC):
271
FullAttentionSpec,
272
TQFullAttentionSpec,
273
MLAAttentionSpec,
274
+ HiddenStateCacheSpec,
275
):
276
self.new_block_ids.extend(b.block_id for b in allocated_blocks)
277
@@ -303,6 +304,7 @@ class SingleTypeKVCacheManager(ABC):
304
FullAttentionSpec,
305
TQFullAttentionSpec,
306
MLAAttentionSpec,
307
+ HiddenStateCacheSpec,
308
):
309
self.new_block_ids.extend(b.block_id for b in new_blocks)
310
return new_blocks