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, 18:30 UTC
Fix MLA runs when use_inductor_graph_partition=True (#38631)
Signed-off-by: ElizaWszola <ewszola@redhat.com>
ElizaWszola committed
Mar 31, 2026 at 15:37 UTC
202f147cf213b2b1f5407d90ab09209542304d9d
1 file changed
+6
-4
vllm/model_executor/layers/attention/mla_attention.py
+6
-4
@@ -929,13 +929,15 @@ def unified_mla_kv_cache_update(
929
the data dependency between them to ensure torch.compile preserves ordering.
930
"""
931
forward_context = get_forward_context()
932
- if forward_context.attn_metadata is None:
933
- # Dummy/profile forwards should not update live KV cache pages.
934
- return torch.empty(0, device=kv_c_normed.device, dtype=kv_c_normed.dtype)
935
-
932
attn_layer = forward_context.no_compile_layers[layer_name]
933
kv_cache = attn_layer.kv_cache
934
935
+ # This needs to run even when we don't have metadata yet, so that the op
936
+ # is correctly captured.
937
+ if kv_cache.numel() == 0:
938
+ # Can't update an empty KV cache.
939
+ return torch.empty(0, device=kv_c_normed.device, dtype=kv_c_normed.dtype)
940
+
941
slot_mapping = forward_context.slot_mapping
942
assert isinstance(slot_mapping, dict), (
943
f"Expected slot_mapping to be a dict, got {type(slot_mapping)}. "