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 12, 02:15 UTC
[KVConnector] remove unused code (the model aware kv ops class) (#29709)
Signed-off-by: KuntaiDu <kuntai@uchicago.edu>
Kuntai Du committed
Dec 4, 2025 at 14:00 UTC
8aaa81b35f96a3b5c56d3dccf58f48129ee34126
1 file changed
+1
-86
vllm/distributed/kv_transfer/kv_connector/utils.py
+1
-86
@@ -8,9 +8,7 @@ from typing import TYPE_CHECKING, Literal
8
9
import torch
10
11
-import vllm.envs as envs
12
-from vllm import _custom_ops as ops
13
-from vllm.config import VllmConfig, get_current_vllm_config
11
+from vllm.config import get_current_vllm_config
12
from vllm.distributed.kv_transfer.kv_connector.factory import KVConnectorFactory
13
from vllm.logger import init_logger
14
from vllm.v1.outputs import KVConnectorOutput, ModelRunnerOutput
@@ -21,89 +19,6 @@ if TYPE_CHECKING:
19
logger = init_logger(__name__)
20
21
24
-class model_aware_kv_ops_helper:
25
- def __init__(self, config: VllmConfig):
26
- self.is_deepseek_mla = config.model_config.is_deepseek_mla
27
- self.use_mla_opt = not envs.VLLM_MLA_DISABLE
28
- self.tp_size = config.parallel_config.tensor_parallel_size
29
-
30
- def get_model_args(self, model_executable: torch.nn.Module):
31
- model_config = model_executable.model.config
32
- self.model_executable = model_executable
33
- num_heads = int(model_config.num_key_value_heads / self.tp_size)
34
- hidden_size = model_config.hidden_size
35
- num_attention_heads = model_config.num_attention_heads
36
-
37
- # Deepseek's MLA (Multi-head Latent Attention) uses two different
38
- # kv_cache shapes based on whether VLLM_MLA_DISABLE is set to 0.
39
- # When VLLM_MLA_DISABLE=0 (default), forward absorb is applied,
40
- # resulting in a kv_cache shape of [num_blks, blk_size, 1,
41
- # kv_lora_rank + qk_rope_head_dim].
42
- # When VLLM_MLA_DISABLE=1, standard FA is used instead, leading
43
- # to a kv_cache shape of [2, num_blks, blk_size,
44
- # num_key_value_heads / tp, qk_nope_head_dim + qk_rope_head_dim].
45
- # For more details, see vllm/v1/attention/backends/mla/common.py.
46
- if self.is_deepseek_mla and self.use_mla_opt:
47
- head_size = model_config.kv_lora_rank + model_config.qk_rope_head_dim
48
- num_heads = 1
49
- elif self.is_deepseek_mla and not self.use_mla_opt:
50
- head_size = model_config.qk_nope_head_dim + model_config.qk_rope_head_dim
51
- else:
52
- head_size = getattr(model_config, "head_dim", None)
53
- if head_size is None:
54
- head_size = int(hidden_size // num_attention_heads)
55
-
56
- return num_heads, head_size
57
-
58
- def get_kv_from_cache(self, kv_cache, num_heads, head_size):
59
- if self.is_deepseek_mla and self.use_mla_opt:
60
- key_cache = kv_cache.reshape(-1, num_heads, head_size)
61
- value_cache = kv_cache.reshape(-1, num_heads, head_size)
62
- else:
63
- key_cache = kv_cache[0].reshape(-1, num_heads, head_size)
64
- value_cache = kv_cache[1].reshape(-1, num_heads, head_size)
65
- return key_cache, value_cache
66
-
67
- def put_kv_to_cache(
68
- self,
69
- model_executable: torch.nn.Module,
70
- keys,
71
- values,
72
- layer,
73
- kv_cache,
74
- slot_mapping,
75
- start_pos,
76
- end_pos,
77
- ):
78
- model_config = model_executable.model.config
79
-
80
- if self.is_deepseek_mla and self.use_mla_opt:
81
- layer.self_attn.attn = layer.self_attn.mla_attn
82
- k_c_normed_k_pe = keys.squeeze(1)
83
- k_c_normed = k_c_normed_k_pe[:, : model_config.kv_lora_rank]
84
- k_pe = k_c_normed_k_pe[:, model_config.kv_lora_rank :]
85
- ops.concat_and_cache_mla(
86
- k_c_normed.to(kv_cache.device),
87
- k_pe.to(kv_cache.device),
88
- kv_cache,
89
- slot_mapping[start_pos:end_pos],
90
- layer.self_attn.attn.kv_cache_dtype,
91
- layer.self_attn.attn._k_scale,
92
- )
93
- else:
94
- key_cache, value_cache = kv_cache[0], kv_cache[1]
95
- ops.reshape_and_cache_flash(
96
- keys.to(key_cache.device),
97
- values.to(value_cache.device),
98
- key_cache,
99
- value_cache,
100
- slot_mapping[start_pos:end_pos],
101
- layer.self_attn.attn.kv_cache_dtype,
102
- layer.self_attn.attn._k_scale,
103
- layer.self_attn.attn._v_scale,
104
- )
105
-
106
-
22
def get_kv_connector_cache_layout():
23
# NOTE (NickLucche) When running disaggregated PD with NIXL, HND layout is
24
# used for faster transfer.