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 7, 11:30 UTC
[Performance] Extract kv update ops from MLA attention backends (#34627)
Signed-off-by: ElizaWszola <ewszola@redhat.com> Signed-off-by: Luka Govedič <ProExpertProg@users.noreply.github.com> Co-authored-by: Di Wu <dw2761@nyu.edu> Co-authored-by: Luka Govedič <ProExpertProg@users.noreply.github.com>
ElizaWszola committed
Mar 2, 2026 at 16:43 UTC
d9c77308776b4d31f03fad8d4129a3d539154166
3 files changed
+128
-11
vllm/config/compilation.py
+1
@@ -1007,6 +1007,7 @@ class CompilationConfig:
1007
# https://github.com/vllm-project/vllm/issues/33267
1008
if not self.use_inductor_graph_partition:
1009
self.splitting_ops.append("vllm::unified_kv_cache_update")
1010
+ self.splitting_ops.append("vllm::unified_mla_kv_cache_update")
1011
1012
elif len(self.splitting_ops) == 0:
1013
if (
vllm/model_executor/layers/attention/mla_attention.py
+83
-11
@@ -434,7 +434,19 @@ class MLAAttention(nn.Module, AttentionLayerBase):
434
if isinstance(attn_metadata, dict):
435
attn_metadata = attn_metadata[self.layer_name]
436
self_kv_cache = self.kv_cache[forward_context.virtual_engine]
437
+ slot_mapping = forward_context.slot_mapping
438
439
+ assert isinstance(slot_mapping, dict), (
440
+ f"Expected slot_mapping to be a dict, got {type(slot_mapping)}. "
441
+ )
442
+ self.impl.do_kv_cache_update(
443
+ kv_c_normed,
444
+ k_pe,
445
+ self_kv_cache,
446
+ slot_mapping.get(self.layer_name),
447
+ self.kv_cache_dtype,
448
+ self._k_scale,
449
+ )
450
if self.attn_backend.accept_output_buffer:
451
output = torch.empty(output_shape, dtype=q.dtype, device=q.device)
452
self.forward_impl(
@@ -451,6 +463,13 @@ class MLAAttention(nn.Module, AttentionLayerBase):
463
q, kv_c_normed, k_pe, self_kv_cache, attn_metadata
464
)
465
else:
466
+ kv_cache_dummy_dep = torch.ops.vllm.unified_mla_kv_cache_update(
467
+ kv_c_normed,
468
+ k_pe,
469
+ self.layer_name,
470
+ self.kv_cache_dtype,
471
+ self._k_scale,
472
+ )
473
if self.attn_backend.accept_output_buffer:
474
output = torch.empty(output_shape, dtype=q.dtype, device=q.device)
475
torch.ops.vllm.unified_mla_attention_with_output(
@@ -459,6 +478,7 @@ class MLAAttention(nn.Module, AttentionLayerBase):
478
k_pe,
479
output,
480
self.layer_name,
481
+ kv_cache_dummy_dep=kv_cache_dummy_dep,
482
)
483
return output
484
else:
@@ -467,6 +487,7 @@ class MLAAttention(nn.Module, AttentionLayerBase):
487
kv_c_normed,
488
k_pe,
489
self.layer_name,
490
+ kv_cache_dummy_dep=kv_cache_dummy_dep,
491
)
492
493
def forward_impl(
@@ -520,17 +541,6 @@ class MLAAttention(nn.Module, AttentionLayerBase):
541
k_c_normed = k_c_normed[:num_actual_toks, ...]
542
k_pe = k_pe[:num_actual_toks, ...]
543
523
- # write the latent and rope to kv cache
524
- if kv_cache.numel() > 0:
525
- ops.concat_and_cache_mla(
526
- k_c_normed,
527
- k_pe.squeeze(1),
528
- kv_cache,
529
- attn_metadata.slot_mapping.flatten(),
530
- kv_cache_dtype=self.kv_cache_dtype,
531
- scale=self._k_scale,
532
- )
533
-
544
if fp8_attention and self.kv_cache_dtype != "fp8_ds_mla":
545
kv_cache = kv_cache.view(current_platform.fp8_dtype())
546
@@ -827,7 +837,12 @@ def unified_mla_attention(
837
kv_c_normed: torch.Tensor,
838
k_pe: torch.Tensor,
839
layer_name: str,
840
+ kv_cache_dummy_dep: torch.Tensor | None = None,
841
) -> torch.Tensor:
842
+ # kv_cache_dummy_dep is not used but accepting it creates a data dependency
843
+ # that ensures torch.compile preserves ordering between KV cache update and
844
+ # attention forward.
845
+ del kv_cache_dummy_dep
846
attn_metadata, layer, kv_cache, _ = get_attention_context(layer_name)
847
output = layer.forward_impl(q, kv_c_normed, k_pe, kv_cache, attn_metadata)
848
@@ -839,6 +854,7 @@ def unified_mla_attention_fake(
854
kv_c_normed: torch.Tensor,
855
k_pe: torch.Tensor,
856
layer_name: str,
857
+ kv_cache_dummy_dep: torch.Tensor | None = None,
858
) -> torch.Tensor:
859
return torch.empty_like(q).contiguous()
860
@@ -852,6 +868,56 @@ direct_register_custom_op(
868
)
869
870
871
+def unified_mla_kv_cache_update(
872
+ kv_c_normed: torch.Tensor,
873
+ k_pe: torch.Tensor,
874
+ layer_name: str,
875
+ kv_cache_dtype: str,
876
+ k_scale: torch.Tensor,
877
+) -> torch.Tensor:
878
+ """
879
+ Returns a dummy that is passed to unified_attention to signal a side effect and
880
+ the data dependency between them to ensure torch.compile preserves ordering.
881
+ """
882
+ forward_context = get_forward_context()
883
+ attn_layer = forward_context.no_compile_layers[layer_name]
884
+ kv_cache = attn_layer.kv_cache[forward_context.virtual_engine]
885
+
886
+ slot_mapping = forward_context.slot_mapping
887
+ assert isinstance(slot_mapping, dict), (
888
+ f"Expected slot_mapping to be a dict, got {type(slot_mapping)}. "
889
+ )
890
+ layer_slot_mapping = slot_mapping.get(layer_name)
891
+ if layer_slot_mapping is not None:
892
+ attn_layer.impl.do_kv_cache_update(
893
+ kv_c_normed,
894
+ k_pe,
895
+ kv_cache,
896
+ layer_slot_mapping,
897
+ kv_cache_dtype,
898
+ k_scale,
899
+ )
900
+
901
+ return torch.empty(0, device=kv_c_normed.device, dtype=kv_c_normed.dtype)
902
+
903
+
904
+def unified_mla_kv_cache_update_fake(
905
+ kv_c_normed: torch.Tensor,
906
+ k_pe: torch.Tensor,
907
+ layer_name: str,
908
+ kv_cache_dtype: str,
909
+ k_scale: torch.Tensor,
910
+) -> torch.Tensor:
911
+ return torch.empty(0, device=kv_c_normed.device, dtype=kv_c_normed.dtype)
912
+
913
+
914
+direct_register_custom_op(
915
+ op_name="unified_mla_kv_cache_update",
916
+ op_func=unified_mla_kv_cache_update,
917
+ fake_impl=unified_mla_kv_cache_update_fake,
918
+)
919
+
920
+
921
@maybe_transfer_kv_layer
922
def unified_mla_attention_with_output(
923
q: torch.Tensor,
@@ -861,7 +927,12 @@ def unified_mla_attention_with_output(
927
layer_name: str,
928
output_scale: torch.Tensor | None = None,
929
output_block_scale: torch.Tensor | None = None,
930
+ kv_cache_dummy_dep: torch.Tensor | None = None,
931
) -> None:
932
+ # kv_cache_dummy_dep is not used but accepting it creates a data dependency
933
+ # that ensures torch.compile preserves ordering between KV cache update and
934
+ # attention forward.
935
+ del kv_cache_dummy_dep
936
attn_metadata, layer, kv_cache, _ = get_attention_context(layer_name)
937
layer.forward_impl(
938
q,
@@ -883,6 +954,7 @@ def unified_mla_attention_with_output_fake(
954
layer_name: str,
955
output_scale: torch.Tensor | None = None,
956
output_block_scale: torch.Tensor | None = None,
957
+ kv_cache_dummy_dep: torch.Tensor | None = None,
958
) -> None:
959
return
960
vllm/v1/attention/backend.py
+44
@@ -811,6 +811,28 @@ class MLAAttentionImpl(AttentionImplBase[T], Generic[T]):
811
"""MQA-style decode forward pass."""
812
raise NotImplementedError
813
814
+ def do_kv_cache_update(
815
+ self,
816
+ kv_c_normed: torch.Tensor,
817
+ k_pe: torch.Tensor,
818
+ kv_cache: torch.Tensor,
819
+ slot_mapping: torch.Tensor,
820
+ kv_cache_dtype: str,
821
+ k_scale: torch.Tensor,
822
+ ) -> None:
823
+ if kv_cache.numel() == 0:
824
+ return
825
+ from vllm import _custom_ops as ops
826
+
827
+ ops.concat_and_cache_mla(
828
+ kv_c_normed,
829
+ k_pe.squeeze(1),
830
+ kv_cache,
831
+ slot_mapping.flatten(),
832
+ kv_cache_dtype=kv_cache_dtype,
833
+ scale=k_scale,
834
+ )
835
+
836
837
class SparseMLAAttentionImpl(AttentionImplBase[T], Generic[T]):
838
"""Sparse MLA attention implementation with only forward_mqa method.
@@ -856,6 +878,28 @@ class SparseMLAAttentionImpl(AttentionImplBase[T], Generic[T]):
878
"""MQA-style decode forward pass."""
879
raise NotImplementedError
880
881
+ def do_kv_cache_update(
882
+ self,
883
+ kv_c_normed: torch.Tensor,
884
+ k_pe: torch.Tensor,
885
+ kv_cache: torch.Tensor,
886
+ slot_mapping: torch.Tensor,
887
+ kv_cache_dtype: str,
888
+ k_scale: torch.Tensor,
889
+ ) -> None:
890
+ if kv_cache.numel() == 0:
891
+ return
892
+ from vllm import _custom_ops as ops
893
+
894
+ ops.concat_and_cache_mla(
895
+ kv_c_normed,
896
+ k_pe.squeeze(1),
897
+ kv_cache,
898
+ slot_mapping.flatten(),
899
+ kv_cache_dtype=kv_cache_dtype,
900
+ scale=k_scale,
901
+ )
902
+
903
904
def is_quantized_kv_cache(kv_cache_dtype: str) -> bool:
905
return kv_cache_dtype.startswith("fp8")