EngineX/xc-llm-ascend
3
0
Fork 0
Code Issues Pull Requests Actions Projects Releases Wiki Activity

[Misc]Remove PD v0 code (#2047)

Browse Source 
Cleanup V0 disaggregated prefill code for V0 Engine.

part of https://github.com/vllm-project/vllm-ascend/issues/1620

TODO: enable v1 e2e test.

- vLLM version: v0.10.0
- vLLM main:
2cc571199b

Signed-off-by: wangxiyuan <wangxiyuan1007@gmail.com>
This commit is contained in:
wangxiyuan
2025-07-28 19:09:22 +08:00
committed by GitHub
parent 935e9d4c9d
commit 0190b68f51
16 changed files with 0 additions and 2619 deletions
Download Patch File Download Diff File Expand all files Collapse all files

8
vllm_ascend/distributed/__init__.py
View File

@@ -18,14 +18,6 @@
from vllm.distributed.kv_transfer.kv_connector.factory import \
KVConnectorFactory
KVConnectorFactory.register_connector(
"AscendHcclConnector", "vllm_ascend.distributed.llmdatadist_connector",
"LLMDataDistConnector")
KVConnectorFactory.register_connector(
"AscendSimpleConnector",
"vllm_ascend.distributed.kv_transfer.simple_connector", "SimpleConnector")
KVConnectorFactory.register_connector(
"LLMDataDistCMgrConnector",
"vllm_ascend.distributed.llmdatadist_c_mgr_connector",

0
vllm_ascend/distributed/kv_transfer/__init__.py
View File

207
vllm_ascend/distributed/kv_transfer/simple_buffer.py
View File

@@ -1,207 +0,0 @@
#
# Copyright (c) 2025 Huawei Technologies Co., Ltd. All Rights Reserved.
# This file is a part of the vllm-ascend project.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
# http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
#
import zlib
from typing import List, Optional
import llm_datadist # type: ignore
import torch
from vllm.distributed.kv_transfer.kv_lookup_buffer.base import \
KVLookupBufferBase
from vllm.logger import logger
from vllm_ascend.distributed.kv_transfer.simple_pipe import SimplePipe
from vllm_ascend.distributed.kv_transfer.utils import TORCH_DTYPE_TO_NPU_DTYPE
# Hash a string into a int32 value.
def int32_hash(data):
assert isinstance(data, str)
data = data.encode("utf-8")
return zlib.adler32(data)
class SimpleBuffer(KVLookupBufferBase):
def __init__(self, data_pipe: SimplePipe):
self.data_pipe = data_pipe
# Consumer buffer need these information to construct receiving buffer.
self.num_layers = None
self.num_heads = None
self.head_size = None
self.dtype = None
self.hidden_size = None
self.key_buffer = None
self.value_buffer = None
self.hidden_buffer = None
def insert(
self,
input_tokens: torch.Tensor,
roi: torch.Tensor,
key: torch.Tensor,
value: torch.Tensor,
hidden: torch.Tensor,
req_id: str,
) -> None:
"""
seq_len: num_tokens of current request.
input_tokens: [seq_len]
roi: [seq_len]
key: [num_layers, seq_len, num_kv_heads, head_size]
value: [num_layers, seq_len, num_kv_heads, head_size]
hidden: [seq_len, hidden_size]
"""
orig_k_shape = key.shape
num_layers = orig_k_shape[0]
# unsequeeze all tensors to make first dim to 1.
# This is because D node can only pull one batch data from P.
# So we make first dim to 1 here in order to pull full data.
key = key.view(num_layers, -1).unsqueeze(0)
value = value.view(num_layers, -1).unsqueeze(0)
hidden = hidden.unsqueeze(0)
hidden_dtype = key.dtype
# initialize LLMDatadist data structure
key_desc = llm_datadist.CacheDesc(
1,
key.shape,
TORCH_DTYPE_TO_NPU_DTYPE[hidden_dtype],
seq_len_dim_index=1,
)
value_desc = llm_datadist.CacheDesc(
1,
value.shape,
TORCH_DTYPE_TO_NPU_DTYPE[hidden_dtype],
seq_len_dim_index=1,
)
hidden_desc = llm_datadist.CacheDesc(
1,
hidden.shape,
TORCH_DTYPE_TO_NPU_DTYPE[hidden_dtype],
seq_len_dim_index=-1,
)
req_id = int32_hash(req_id)
key_cache_key = llm_datadist.CacheKey(self.data_pipe.cluster_id,
req_id, 1)
value_cache_key = llm_datadist.CacheKey(self.data_pipe.cluster_id,
req_id, 2)
hidden_cache_key = llm_datadist.CacheKey(self.data_pipe.cluster_id,
req_id, 3)
# Currently we use hash value of request id as key, so no need to send input_tokens
self.key_buffer = self.data_pipe.send_tensor(key, key_desc,
key_cache_key)
self.value_buffer = self.data_pipe.send_tensor(value, value_desc,
value_cache_key)
self.hidden_buffer = self.data_pipe.send_tensor(
hidden, hidden_desc, hidden_cache_key)
def drop_select(
self,
input_tokens: torch.Tensor,
roi: Optional[torch.Tensor],
req_id: str,
) -> List[Optional[torch.Tensor]]:
"""Select and *drop* KV cache entries from the lookup buffer.
The functionality is similar to the following python statements
```
ret = buffer.pop(input_tokens, roi)
return ret
```
Args:
input_tokens (torch.Tensor): token IDs.
roi (torch.Tensor): A binary mask on top of the input tokens
Returns:
A list of tensors including:
key: [num_layers, num_tokens, num_heads, head_size]
value: [num_layers, num_tokens, num_heads, head_size]
hidden_or_intermediate_states: [num_tokens, hidden_size]
roi: None (Currently we don't supported roi)
"""
orig_req_id = req_id
req_id = int32_hash(req_id)
num_tokens = input_tokens.shape[0]
kv_shape = (
1,
self.num_layers,
num_tokens * self.num_heads * self.head_size,
)
hidden_shape = (1, num_tokens, self.hidden_size)
key_desc = llm_datadist.CacheDesc(
1,
kv_shape,
TORCH_DTYPE_TO_NPU_DTYPE[self.dtype],
seq_len_dim_index=-1,
)
value_desc = llm_datadist.CacheDesc(
1,
kv_shape,
TORCH_DTYPE_TO_NPU_DTYPE[self.dtype],
seq_len_dim_index=-1,
)
hidden_desc = llm_datadist.CacheDesc(
1,
hidden_shape,
TORCH_DTYPE_TO_NPU_DTYPE[self.dtype],
seq_len_dim_index=-1,
)
key_cache_key = llm_datadist.CacheKey(self.data_pipe.cluster_id,
req_id, 1)
value_cache_key = llm_datadist.CacheKey(self.data_pipe.cluster_id,
req_id, 2)
hidden_cache_key = llm_datadist.CacheKey(self.data_pipe.cluster_id,
req_id, 3)
# Deallocate buffer allocated in last round.
if self.key_buffer:
try:
self.data_pipe.deallocate_buffer(self.key_buffer)
self.data_pipe.deallocate_buffer(self.value_buffer)
self.data_pipe.deallocate_buffer(self.hidden_buffer)
except Exception as e:
logger.warning(
f"Failed to free kv cache buffer, Error code: {str(e)}")
try:
self.key_buffer, key = self.data_pipe.recv_tensor(
key_desc, key_cache_key)
self.value_buffer, value = self.data_pipe.recv_tensor(
value_desc, value_cache_key)
self.hidden_buffer, hidden = self.data_pipe.recv_tensor(
hidden_desc, hidden_cache_key)
key = key.view(self.num_layers, num_tokens, self.num_heads,
self.head_size)
value = value.view(self.num_layers, num_tokens, self.num_heads,
self.head_size)
hidden = hidden.view(num_tokens, self.hidden_size)
except Exception as e:
logger.warning(
f"Fail to receive kv cache and hidden states of request: {orig_req_id} "
f"Error is {str(e)}")
return [None, None, None, None]
return [key, value, hidden, roi]
def close(self):
pass

379
vllm_ascend/distributed/kv_transfer/simple_connector.py
View File

@@ -1,379 +0,0 @@
#
# Copyright (c) 2025 Huawei Technologies Co., Ltd. All Rights Reserved.
# This file is a part of the vllm-ascend project.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
# http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
#
from typing import TYPE_CHECKING, List, Optional, Tuple, Union
import torch
import torch_npu
import vllm.envs as vllm_envs
from vllm.config import VllmConfig
from vllm.distributed.kv_transfer.kv_connector.base import KVConnectorBase
from vllm.distributed.parallel_state import get_dp_group
from vllm.logger import logger
from vllm.sequence import IntermediateTensors
from vllm_ascend.distributed.kv_transfer.simple_buffer import SimpleBuffer
from vllm_ascend.distributed.kv_transfer.simple_pipe import SimplePipe
if TYPE_CHECKING:
from vllm.worker.model_runner import ModelInputForGPUWithSamplingMetadata
class SimpleConnector(KVConnectorBase):
def __init__(
self,
rank: int,
local_rank: int,
config: VllmConfig,
):
self.config = config
self.model_config = config.model_config.hf_config
self.tp_size = config.parallel_config.tensor_parallel_size
self.rank = rank
self.local_rank = local_rank
self.is_deepseek_mla = config.model_config.is_deepseek_mla
self.use_mla_opt = not vllm_envs.VLLM_MLA_DISABLE
self.n_layer = self.config.model_config.get_num_layers(
self.config.parallel_config)
self.producer_data_pipe: Optional[SimplePipe]
self.consumer_data_pipe: Optional[SimplePipe]
self.producer_buffer: Optional[SimpleBuffer]
self.consumer_buffer: Optional[SimpleBuffer]
if self.config.kv_transfer_config.is_kv_producer:
self.producer_data_pipe = SimplePipe(
rank=rank,
local_rank=local_rank,
kv_transfer_config=config.kv_transfer_config,
hostname="",
port_offset=rank,
)
self.producer_buffer = SimpleBuffer(self.producer_data_pipe)
else:
self.consumer_data_pipe = SimplePipe(
rank=rank,
local_rank=local_rank,
kv_transfer_config=config.kv_transfer_config,
hostname="",
port_offset=rank,
)
self.consumer_buffer = SimpleBuffer(self.consumer_data_pipe)
def select(
self,
input_tokens: Optional[torch.Tensor],
roi: Optional[torch.Tensor],
req_id: str,
) -> List[Optional[torch.Tensor]]:
assert self.consumer_buffer is not None, (
"Please initialize the "
"consumer buffer before calling select.")
return self.consumer_buffer.drop_select(input_tokens, roi, req_id)
def insert(
self,
input_tokens: torch.Tensor,
roi: torch.Tensor,
keys: torch.Tensor,
values: torch.Tensor,
hidden: torch.Tensor,
req_id: str,
) -> None:
assert self.producer_buffer is not None, (
"Please initialize the "
"producer buffer before calling insert.")
self.producer_buffer.insert(input_tokens, roi, keys, values, hidden,
req_id)
def send_kv_caches_and_hidden_states(
self,
model_executable: torch.nn.Module,
model_input: "ModelInputForGPUWithSamplingMetadata",
kv_caches: List[torch.Tensor],
hidden_or_intermediate_states: Union[torch.Tensor,
IntermediateTensors],
) -> None:
input_tokens_tensor = model_input.input_tokens
seq_lens = model_input.attn_metadata.seq_lens
slot_mapping_flat = model_input.attn_metadata.slot_mapping.flatten()
num_prefill_tokens = model_input.attn_metadata.num_prefill_tokens
start_layer = model_executable.model.start_layer
end_layer = model_executable.model.end_layer
model_config = self.model_config
num_heads = int(model_config.num_key_value_heads / self.tp_size)
hidden_size = model_config.hidden_size
num_attention_heads = model_config.num_attention_heads
# Deepseek's MLA (Multi-head Latent Attention) uses two different
# kv_cache shapes based on whether VLLM_MLA_DISABLE is set to 0.
# When VLLM_MLA_DISABLE=0 (default), forward absorb is applied,
# resulting in a kv_cache shape of [num_blks, blk_size, 1,
# kv_lora_rank + qk_rope_head_dim].
# When VLLM_MLA_DISABLE=1, standard FA is used instead, leading
# to a kv_cache shape of [2, num_blks, blk_size,
# num_key_value_heads / tp, qk_nope_head_dim + qk_rope_head_dim].
# For more details, see vllm/attention/backends/mla/common.py.
if self.is_deepseek_mla and self.use_mla_opt:
head_size = (model_config.kv_lora_rank +
model_config.qk_rope_head_dim)
num_heads = 1
elif self.is_deepseek_mla and not self.use_mla_opt:
head_size = (model_config.qk_nope_head_dim +
model_config.qk_rope_head_dim)
else:
head_size = getattr(
model_config,
"head_dim",
int(hidden_size // num_attention_heads),
)
# Enumerate over all requests and insert them one by one.
for idx, slen in enumerate(seq_lens):
start_pos = sum(seq_lens[:idx])
end_pos = start_pos + slen
if start_pos >= num_prefill_tokens:
# vllm/worker/model_runner.py::_prepare_model_input_tensors:
# - input_tokens[:num_prefill_tokens] contains prefill tokens.
# - input_tokens[num_prefill_tokens:] contains decode tokens.
logger.warning("You have some decode requests while using "
"SimpleConnector. Their KVCache won't be sent.")
break
current_tokens = input_tokens_tensor[start_pos:end_pos]
keys, values = [], []
for layer_id in range(start_layer, end_layer):
kv_cache = kv_caches[layer_id - start_layer]
if self.is_deepseek_mla and self.use_mla_opt:
key_cache = kv_cache.reshape(-1, num_heads, head_size)
value_cache = kv_cache.reshape(-1, num_heads, head_size)
else:
key_cache = kv_cache[0].reshape(-1, num_heads, head_size)
value_cache = kv_cache[1].reshape(-1, num_heads, head_size)
current_slot_mapping = slot_mapping_flat[start_pos:end_pos]
keys.append(key_cache[current_slot_mapping].unsqueeze(0))
values.append(value_cache[current_slot_mapping].unsqueeze(0))
# shape: [num_layers, num_tokens, num_heads, head_size]
keys = torch.cat(keys, dim=0)
values = torch.cat(values, dim=0)
cur_req_id = list(model_input.request_ids_to_seq_ids.keys())[idx]
# Currently we haven't considered situation of roi, pass None here.
self.insert(
current_tokens,
None,
keys,
values,
hidden_or_intermediate_states[start_pos:end_pos],
cur_req_id,
)
logger.info("[rank%d][P]: KV send DONE.", torch.distributed.get_rank())
def recv_kv_caches_and_hidden_states(
self,
model_executable: torch.nn.Module,
model_input: "ModelInputForGPUWithSamplingMetadata",
kv_caches: List[torch.Tensor],
) -> Tuple[
Union[torch.Tensor, IntermediateTensors],
bool,
"ModelInputForGPUWithSamplingMetadata",
]:
bypass_model_exec = True
model_config = self.model_config
# get model config
start_layer = model_executable.model.start_layer
end_layer = model_executable.model.end_layer
num_heads, head_dim = kv_caches[0].shape[-2:]
hidden_size = model_config.hidden_size
num_attention_heads = model_config.num_attention_heads
num_layers = end_layer - start_layer
if self.is_deepseek_mla and self.use_mla_opt:
head_size = (model_config.kv_lora_rank +
model_config.qk_rope_head_dim)
num_heads = 1
elif self.is_deepseek_mla and not self.use_mla_opt:
head_size = (model_config.qk_nope_head_dim +
model_config.qk_rope_head_dim)
else:
head_size = getattr(
model_config,
"head_dim",
int(hidden_size // num_attention_heads),
)
self.consumer_buffer.num_heads = num_heads # type: ignore
self.consumer_buffer.num_layers = num_layers # type: ignore
self.consumer_buffer.head_size = head_size # type: ignore
self.consumer_buffer.dtype = kv_caches[0].dtype # type: ignore
self.consumer_buffer.hidden_size = hidden_size # type: ignore
input_tokens_tensor = model_input.input_tokens
seq_lens = model_input.attn_metadata.seq_lens
num_prefill_tokens = model_input.attn_metadata.num_prefill_tokens
slot_mapping = model_input.attn_metadata.slot_mapping.flatten()
total_tokens = model_input.attn_metadata.num_prefill_tokens + model_input.attn_metadata.num_decode_tokens
hidden_or_intermediate_states_for_one_req = []
input_tokens_list = []
num_computed_tokens_list = []
start_pos_list = []
# enumerate different requests
for idx, slen in enumerate(seq_lens):
start_pos = sum(seq_lens[:idx])
end_pos = start_pos + slen
if start_pos >= num_prefill_tokens:
logger.warning("You should set --enable_chunked_prefill=False "
"and --max_num_batched_tokens "
"should be equal to --max_seq_len_to_capture")
bypass_model_exec = False
assert start_pos == num_prefill_tokens
break
current_tokens = input_tokens_tensor[start_pos:end_pos]
num_tokens = slen
# collecting data for rebuilding the input
input_tokens_list.append(current_tokens)
start_pos_list.append(start_pos)
cur_req_id = list(model_input.request_ids_to_seq_ids.keys())[idx]
ret = self.select(
current_tokens,
torch.ones_like(current_tokens, dtype=bool),
cur_req_id,
)
if ret[0] is None:
# didn't find any match.
bypass_model_exec = False
num_computed_tokens_list.append(0)
continue
keys: torch.Tensor = ret[0]
values: torch.Tensor = ret[1]
hidden: torch.Tensor = ret[2]
num_computed_tokens = keys.shape[1]
num_computed_tokens_list.append(num_computed_tokens)
# check if both KV cache and the hidden states are received
# If not, need to redo the forwarding to compute missing states
if not all([(num_computed_tokens == num_tokens), hidden is not None
]):
bypass_model_exec = False
# update the end position based on how many tokens are cached.
end_pos = start_pos + num_computed_tokens
# put received KV caches into paged memory
for i in range(
model_executable.model.start_layer,
model_executable.model.end_layer,
):
kv_cache = kv_caches[i - model_executable.model.start_layer]
layer = model_executable.model.layers[i]
if self.is_deepseek_mla and self.use_mla_opt:
layer.self_attn.attn = layer.self_attn.mla_attn
key_cache = kv_cache
slots = slot_mapping[start_pos:end_pos]
sliced_key = keys[i - model_executable.model.start_layer]
torch_npu._npu_reshape_and_cache_siso(key=sliced_key,
key_cache=key_cache,
slot_indices=slots)
else:
key_cache, value_cache = kv_cache[0], kv_cache[1]
sliced_key = keys[i - model_executable.model.start_layer]
sliced_value = values[i -
model_executable.model.start_layer]
torch_npu._npu_reshape_and_cache(
key=sliced_key,
value=sliced_value,
key_cache=key_cache,
value_cache=value_cache,
slot_indices=slot_mapping[start_pos:end_pos],
)
hidden_or_intermediate_states_for_one_req.append(hidden)
if not bypass_model_exec:
# Some of the KV cache is not retrieved
# Here we will fall back to normal model forwarding
# But optionally you can adjust model_input so that you only do
# prefilling on those tokens that are missing KV caches.
if get_dp_group().world_size > 1:
bypass_model_exec = True
hidden_or_intermediate_states = torch.empty(
[total_tokens, hidden_size],
dtype=kv_caches[0].dtype,
device=kv_caches[0].device)
logger.warning(
"[Detect there is more one DP rank in this decode node, in this scenario, no recompute is expected when kv cache dose not received.]"
)
else:
logger.warning(
"[rank%d]: Failed to receive all KVs and hidden "
"states, redo model forwarding.",
torch.distributed.get_rank())
hidden_or_intermediate_states = None
else:
logger.debug(
"[rank%d]: Successfully received all KVs and hidden "
"states, skip model forwarding.",
torch.distributed.get_rank(),
)
# Can't directly concat here which might cause error when bs = 1.
# hidden_or_intermediate_states = torch.empty(total_num_tokens, hidden_size, dtype=kv_caches[0].dtype, device=kv_caches[0].device)
if len(hidden_or_intermediate_states_for_one_req) == 1:
hidden = hidden_or_intermediate_states_for_one_req[0]
tmp_indice = torch.tensor([0] * hidden.shape[0],
dtype=torch.int64).npu()
hidden_or_intermediate_states = torch.empty_like(hidden)
torch_npu.scatter_update_(
hidden_or_intermediate_states,
tmp_indice,
hidden,
axis=-1,
)
else:
hidden_or_intermediate_states = torch.cat(
hidden_or_intermediate_states_for_one_req, dim=0)
return hidden_or_intermediate_states, bypass_model_exec, model_input
def close(self):
self.producer_data_pipe.close() # type: ignore
self.consumer_data_pipe.close() # type: ignore
self.producer_buffer.close() # type: ignore
self.consumer_buffer.close() # type: ignore

207
vllm_ascend/distributed/kv_transfer/simple_pipe.py
View File

@@ -1,207 +0,0 @@
#
# Copyright (c) 2025 Huawei Technologies Co., Ltd. All Rights Reserved.
# This file is a part of the vllm-ascend project.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
# http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
#
import threading
import time
from typing import Optional
import llm_datadist # type: ignore
import msgpack # type: ignore
import torch
import torch_npu
import torchair # type: ignore
import zmq # type: ignore
from vllm.distributed.kv_transfer.kv_pipe.base import KVPipeBase
from vllm.logger import logger
from vllm.utils import get_ip
import vllm_ascend.envs as envs
from vllm_ascend.distributed.kv_transfer.utils import NPU_DTYPE_TO_TORCH_DTYPE
class SimplePipe(KVPipeBase):
def __init__(
self,
rank,
local_rank,
kv_transfer_config,
hostname: str = "",
port_offset: int = 0, # NPU offset in current P/D instance.
):
self.rank = rank
self.local_rank = local_rank
# Currently for 1P1D situation, we use cluster_id=0 for both Prefill and Decode
# Will change here in the future to support xPyD.
self.cluster_id = 0
self.config = kv_transfer_config
kv_connector_extra_config = kv_transfer_config.kv_connector_extra_config
kv_role = kv_transfer_config.kv_role
if kv_role == "kv_producer":
self.role = llm_datadist.LLMRole.PROMPT
elif kv_role == "kv_consumer":
self.role = llm_datadist.LLMRole.DECODER
else:
raise NotImplementedError(
"kv_role should be inside [kv_producer, kv_consumer]")
prefill_device_ips = kv_connector_extra_config.get(
"prefill_device_ips", None)
decode_device_ips = kv_connector_extra_config.get(
"decode_device_ips", None)
if prefill_device_ips is None or decode_device_ips is None:
raise ValueError(
"Please specify prefill_device_ips and decode_device_ips"
"in kv_transfer_config.kv_connector_extra_config")
p_device_num = len(prefill_device_ips)
d_device_num = len(decode_device_ips)
# When number of devices in P and D is not equal,
# we assume that device in D can be mapped to any device in P.
self.p_device_rank = self.rank % p_device_num
self.d_device_rank = self.rank % d_device_num
self.prompt_ip_list = prefill_device_ips
self.decode_ip_list = decode_device_ips
self.llmdatadist_comm_port = kv_connector_extra_config.get(
"llmdatadist_comm_port", 26000)
# LLMDataDist initializing.
self.data_dist = llm_datadist.LLMDataDist(self.role, self.cluster_id)
self._prepare_data_dist()
# Decoder needs to initialize and link cluster
if self.role == llm_datadist.LLMRole.DECODER:
self.cluster = self._make_cluster()
_, ret = self.data_dist.link_clusters([self.cluster], 20000)
logger.info(
f"rank {self.rank}, local_rank {self.local_rank} link, ret={ret}"
)
# If `proxy_ip` or `proxy_port` is `""`,
# then the ping thread will not be enabled.
proxy_ip = self.config.get_from_extra_config("proxy_ip", "")
proxy_port = self.config.get_from_extra_config("proxy_port", "")
if proxy_ip == "" or proxy_port == "":
self.proxy_address = ""
else:
self.proxy_address = proxy_ip + ":" + str(proxy_port)
self._register_thread = None
if port_offset == 0 and self.proxy_address != "":
# Initialize zmq socket and register to proxy.
# Note that only NPU 0 of each P/D instance register to proxy.
if not hostname:
hostname = get_ip() # Get ip of current host.
port = int(kv_transfer_config.kv_port) + port_offset
if port == 0:
raise ValueError("Port cannot be 0")
self._hostname = hostname
self._port = port
# Each card corresponds to a ZMQ address.
self.zmq_address = f"{self._hostname}:{self._port}"
self.context = zmq.Context() # type: ignore
self.router_socket = self.context.socket(
zmq.ROUTER) # type: ignore
self.router_socket.bind(f"tcp://{self.zmq_address}")
# The `http_port` must be consistent with the serving port of OpenAI.
self.http_address = (
f"{self._hostname}:"
f"{self.config.kv_connector_extra_config['http_port']}")
self._register_thread = threading.Thread(
target=self._register_to_proxy, daemon=True)
self._register_thread.start()
def _prepare_data_dist(self):
options = {
"llm.SyncKvCacheWaitTime": envs.LLMDATADIST_SYNC_CACHE_WAIT_TIME,
}
if self.role == llm_datadist.LLMRole.PROMPT:
options["ge.exec.deviceId"] = str(self.local_rank)
options["llm.listenIpInfo"] = (
f"{self.prompt_ip_list[self.p_device_rank]}:{self.llmdatadist_comm_port}"
)
else:
options["ge.exec.deviceId"] = str(self.local_rank)
print(f"prepare datadist, options: {options}")
self.data_dist.init(options)
self.kv_transfer = self.data_dist.kv_cache_manager
print(f"{self.rank} rank data dist is ready")
def _make_cluster(self):
cluster = llm_datadist.LLMClusterInfo()
cluster.remote_cluster_id = self.cluster_id
local_ip = self.decode_ip_list[self.d_device_rank]
remote_ip = self.prompt_ip_list[self.p_device_rank]
cluster.append_local_ip_info(local_ip, 0)
cluster.append_remote_ip_info(remote_ip, self.llmdatadist_comm_port)
return cluster
def _register_to_proxy(self):
sock = self.context.socket(zmq.DEALER) # type: ignore
sock.setsockopt_string(zmq.IDENTITY, self.zmq_address) # type: ignore
logger.debug("ping start, zmq_address:%s", self.zmq_address)
sock.connect(f"tcp://{self.proxy_address}")
data = {
"type": "P" if self.config.is_kv_producer else "D",
"http_address": self.http_address,
"zmq_address": self.zmq_address,
}
while True:
sock.send(msgpack.dumps(data))
time.sleep(3)
def send_tensor(
self,
tensor: Optional[torch.Tensor],
tensor_desc: llm_datadist.CacheDesc,
tensor_key: llm_datadist.CacheKey,
) -> llm_datadist.Cache:
buffer = self.kv_transfer.allocate_cache(tensor_desc, [tensor_key])
buffer_addr = buffer.per_device_tensor_addrs[0]
data_tensor = torchair.llm_datadist.create_npu_tensors(
tensor_desc.shape, tensor.dtype, buffer_addr)[0] # type: ignore
update_indices = torch.tensor(
[0] * tensor.shape[0], # type: ignore
dtype=torch.int64).npu()
torch_npu.scatter_update_(data_tensor, update_indices, tensor, axis=-1)
# Free cache_id of buffer, actual deallocate will happen after consumer performing pull_cache.
self.kv_transfer.deallocate_cache(buffer)
return buffer
def recv_tensor(
self,
tensor_desc: llm_datadist.CacheDesc,
tensor_key: llm_datadist.CacheKey,
) -> llm_datadist.Cache:
"""Note that this function only creates empty tensor on buffer addr and returns it."""
tmp_buffer = self.kv_transfer.allocate_cache(tensor_desc)
buffer_addr = tmp_buffer.per_device_tensor_addrs[0]
data_tensor = torchair.llm_datadist.create_npu_tensors(
tensor_desc.shape,
NPU_DTYPE_TO_TORCH_DTYPE[tensor_desc.data_type],
buffer_addr,
)[0]
self.kv_transfer.pull_cache(tensor_key, tmp_buffer, 0)
# tmp_buffer is allocated without key and will be deallocated here immediately.
# Free buffer here will cause accuracy problem.
# self.kv_transfer.deallocate_cache(tmp_buffer)
return tmp_buffer, data_tensor
def deallocate_buffer(self, buffer: llm_datadist.Cache):
self.kv_transfer.deallocate_cache(buffer)
def close(self):
self.data_dist.unlink_clusters([self.cluster], 5000)

40
vllm_ascend/distributed/kv_transfer/utils.py
View File

@@ -1,40 +0,0 @@
#
# Copyright (c) 2025 Huawei Technologies Co., Ltd. All Rights Reserved.
# This file is a part of the vllm-ascend project.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
# http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
#
import llm_datadist # type: ignore
import torch
TORCH_DTYPE_TO_NPU_DTYPE = {
torch.half: llm_datadist.DataType.DT_FLOAT16,
torch.float16: llm_datadist.DataType.DT_FLOAT16,
torch.bfloat16: llm_datadist.DataType.DT_BF16,
torch.float: llm_datadist.DataType.DT_FLOAT,
torch.float32: llm_datadist.DataType.DT_FLOAT,
torch.int8: llm_datadist.DataType.DT_INT8,
torch.int64: llm_datadist.DataType.DT_INT64,
torch.int32: llm_datadist.DataType.DT_INT32,
}
NPU_DTYPE_TO_TORCH_DTYPE = {
llm_datadist.DataType.DT_FLOAT16: torch.half,
llm_datadist.DataType.DT_FLOAT16: torch.float16,
llm_datadist.DataType.DT_BF16: torch.bfloat16,
llm_datadist.DataType.DT_FLOAT: torch.float,
llm_datadist.DataType.DT_FLOAT: torch.float32,
llm_datadist.DataType.DT_INT8: torch.int8,
llm_datadist.DataType.DT_INT64: torch.int64,
llm_datadist.DataType.DT_INT32: torch.int32,
}

470
vllm_ascend/distributed/llmdatadist_connector.py
View File

@@ -1,470 +0,0 @@
#
# Copyright (c) 2025 Huawei Technologies Co., Ltd. All Rights Reserved.
# This file is a part of the vllm-ascend project.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
# http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
#
import os
import re
import subprocess
from typing import TYPE_CHECKING, List, Tuple, Union
import torch
import torch_npu
import torchair # type: ignore
from vllm.config import VllmConfig
from vllm.distributed.kv_transfer.kv_connector.base import KVConnectorBase
from vllm.logger import logger
from vllm.sequence import IntermediateTensors
import vllm_ascend.envs as envs
if TYPE_CHECKING:
from vllm.worker.model_runner import ModelInputForGPUWithSamplingMetadata
import llm_datadist # type: ignore
TORCH_DTYPE_TO_NPU_DTYPE = {
torch.half: llm_datadist.DataType.DT_FLOAT16,
torch.float16: llm_datadist.DataType.DT_FLOAT16,
torch.bfloat16: llm_datadist.DataType.DT_BF16,
torch.float: llm_datadist.DataType.DT_FLOAT,
torch.float32: llm_datadist.DataType.DT_FLOAT,
torch.int8: llm_datadist.DataType.DT_INT8,
torch.int64: llm_datadist.DataType.DT_INT64,
torch.int32: llm_datadist.DataType.DT_INT32
}
# Get all device ips using hccn_tool
HCCN_TOOL_PATH = envs.HCCN_PATH
def get_device_ips():
world_size = 8
npu_info = subprocess.run(['npu-smi', 'info', '-m'],
stdout=subprocess.PIPE,
stderr=subprocess.PIPE,
universal_newlines=True)
if npu_info.returncode != 0 or not os.path.exists(HCCN_TOOL_PATH):
raise RuntimeError("No npu-smi/hccn_tool tools provided for NPU.")
re_result = re.match(r'.*\n\t([0-9]+).*', npu_info.stdout)
if re_result is None:
raise RuntimeError("Can't find npu start index")
npu_start_idx = int(re_result.group(1))
device_ip_list = []
for ip_offset in range(world_size):
cmd = [
HCCN_TOOL_PATH, '-i', f'{npu_start_idx + ip_offset}', '-ip', '-g'
]
device_ip_info = subprocess.run(cmd,
stdout=subprocess.PIPE,
stderr=subprocess.PIPE,
universal_newlines=True)
re_result = re.match(r'ipaddr:(.*)\n', device_ip_info.stdout)
if re_result is None:
raise RuntimeError("Can't find npu ip")
device_ip = re_result.group(1)
device_ip_list.append(device_ip)
return device_ip_list
class KVTransferEngine:
def __init__(self, world_size, n_layer, role, local_rank):
self.world_size = world_size
self.n_layer = n_layer
self.role = role
self.device_ip_list = get_device_ips()
self.local_rank = local_rank
self.cluster_id = local_rank
self.data_dist = llm_datadist.LLMDataDist(self.role, self.cluster_id)
prompt_device_ids = envs.PROMPT_DEVICE_ID
decode_device_ids = envs.DECODE_DEVICE_ID
if prompt_device_ids is None or decode_device_ids is None:
raise ValueError(
"Please specify env PROMPT_DEVICE_ID or DECODE_DEVICE_ID")
prompt_ids = [
int(x.strip()) for x in prompt_device_ids.split(",") if x.strip()
]
decode_ids = [
int(x.strip()) for x in decode_device_ids.split(",") if x.strip()
]
self.prompt_ip_list = [self.device_ip_list[i] for i in prompt_ids]
self.decode_ip_list = [self.device_ip_list[i] for i in decode_ids]
def prepare_data_dist(self):
options = {
"llm.SyncKvCacheWaitTime": envs.LLMDATADIST_SYNC_CACHE_WAIT_TIME,
}
if self.role == llm_datadist.LLMRole.PROMPT:
options["ge.exec.deviceId"] = str(self.local_rank)
options[
"llm.listenIpInfo"] = f"{self.prompt_ip_list[self.local_rank]}:{envs.LLMDATADIST_COMM_PORT}"
else:
options["ge.exec.deviceId"] = str(self.local_rank)
self.data_dist.init(options)
self.kv_transfer = self.data_dist.kv_cache_manager
logger.info(
f"{self.local_rank}/{self.world_size} rank data dist is ready")
def make_cluster(self, prefill_ip, cluster_id=-1):
cluster = llm_datadist.LLMClusterInfo()
cluster.remote_cluster_id = cluster_id
local_ip = self.decode_ip_list[self.local_rank]
remote_ip = prefill_ip
cluster.append_local_ip_info(local_ip, 0)
cluster.append_remote_ip_info(remote_ip, 26000)
return cluster
class LLMDataDistConnector(KVConnectorBase):
def __init__(
self,
rank: int,
local_rank: int,
config: VllmConfig,
):
self.config = config
self.tp_size = config.parallel_config.tensor_parallel_size
self.rank = rank
self.local_rank = local_rank
if self.config.kv_transfer_config.kv_role == "kv_producer":
self.role = llm_datadist.LLMRole.PROMPT
elif self.config.kv_transfer_config.kv_role == "kv_consumer":
self.role = llm_datadist.LLMRole.DECODER
else:
raise NotImplementedError(
"kv_role should be inside [kv_producer, kv_consumer]")
self.world_size = self.config.parallel_config.world_size
self.n_layer = self.config.model_config.get_num_layers(
self.config.parallel_config)
self.llm_datadist_engine = KVTransferEngine(self.world_size,
self.n_layer, self.role,
self.local_rank)
if self.role == llm_datadist.LLMRole.PROMPT:
self.llm_datadist_engine.prepare_data_dist()
else:
self.llm_datadist_engine.prepare_data_dist()
self.cluster = self.llm_datadist_engine.make_cluster(
self.llm_datadist_engine.prompt_ip_list[self.local_rank],
self.llm_datadist_engine.cluster_id)
_, ret = self.llm_datadist_engine.data_dist.link_clusters(
[self.cluster], 20000)
logger.info(f"local_rank {self.local_rank} link, ret={ret}")
def send_kv_caches_and_hidden_states(
self, model_executable: torch.nn.Module,
model_input: "ModelInputForGPUWithSamplingMetadata",
kv_caches: List[torch.Tensor],
hidden_or_intermediate_states: Union[torch.Tensor, IntermediateTensors]
) -> None:
input_tokens_tensor = model_input.input_tokens
seq_lens = model_input.attn_metadata.seq_lens
slot_mapping_flat = model_input.attn_metadata.slot_mapping.flatten()
start_layer = model_executable.model.start_layer
end_layer = model_executable.model.end_layer
model_config = model_executable.model.config
num_heads = int(model_config.num_key_value_heads / self.tp_size)
hidden_size = model_config.hidden_size
num_attention_heads = model_config.num_attention_heads
head_size = int(hidden_size / num_attention_heads)
num_layer = end_layer - start_layer
# Get shape of input_tokens_tensor and kv_cache
input_shape = (1, input_tokens_tensor.shape[0], 1, 1)
hidden_shape = (1, input_tokens_tensor.shape[0], 1, hidden_size)
kv_shape = (1, input_tokens_tensor.shape[0], num_heads, head_size)
assert kv_caches[0].dtype == hidden_or_intermediate_states.dtype
kv_hidden_dtype = kv_caches[0].dtype
input_dtype = torch.int32
# initialize LLMDatadist data structure
key_desc = llm_datadist.CacheDesc(
num_layer,
kv_shape,
TORCH_DTYPE_TO_NPU_DTYPE[kv_hidden_dtype],
seq_len_dim_index=1)
value_desc = llm_datadist.CacheDesc(
num_layer,
kv_shape,
TORCH_DTYPE_TO_NPU_DTYPE[kv_hidden_dtype],
seq_len_dim_index=1)
input_desc = llm_datadist.CacheDesc(
1,
input_shape,
TORCH_DTYPE_TO_NPU_DTYPE[input_dtype],
seq_len_dim_index=-1)
hidden_desc = llm_datadist.CacheDesc(
1,
hidden_shape,
TORCH_DTYPE_TO_NPU_DTYPE[kv_hidden_dtype],
seq_len_dim_index=-1)
key_cache_keys = [
llm_datadist.CacheKey(self.llm_datadist_engine.cluster_id, 0, 1)
]
value_cache_keys = [
llm_datadist.CacheKey(self.llm_datadist_engine.cluster_id, 0, 2)
]
input_cache_keys = [
llm_datadist.CacheKey(self.llm_datadist_engine.cluster_id, 0, 3)
]
hidden_cache_keys = [
llm_datadist.CacheKey(self.llm_datadist_engine.cluster_id, 0, 4)
]
self.key_buffer = self.llm_datadist_engine.kv_transfer.allocate_cache(
key_desc, key_cache_keys)
self.value_buffer = self.llm_datadist_engine.kv_transfer.allocate_cache(
value_desc, value_cache_keys)
self.input_buffer = self.llm_datadist_engine.kv_transfer.allocate_cache(
input_desc, input_cache_keys)
self.hidden_buffer = self.llm_datadist_engine.kv_transfer.allocate_cache(
hidden_desc, hidden_cache_keys)
key_buffer_addr = self.key_buffer.per_device_tensor_addrs[0]
value_buffer_addr = self.value_buffer.per_device_tensor_addrs[0]
input_buffer_addr = self.input_buffer.per_device_tensor_addrs[0]
hidden_buffer_addr = self.hidden_buffer.per_device_tensor_addrs[0]
self.key_cache = torchair.llm_datadist.create_npu_tensors(
key_desc.shape, kv_hidden_dtype, key_buffer_addr)
self.value_cache = torchair.llm_datadist.create_npu_tensors(
value_desc.shape, kv_hidden_dtype, value_buffer_addr)
self.input_cache = torchair.llm_datadist.create_npu_tensors(
input_desc.shape, input_dtype, input_buffer_addr)
self.hidden_cache = torchair.llm_datadist.create_npu_tensors(
hidden_desc.shape, kv_hidden_dtype, hidden_buffer_addr)
indices = torch.tensor([0], dtype=torch.int64).npu()
# copy cache data into llm datadist cache using scatter update
for idx, slen in enumerate(seq_lens):
start_pos = sum(seq_lens[:idx])
end_pos = start_pos + slen
current_tokens = input_tokens_tensor[start_pos:end_pos].to(
torch.int32)
for layer_id in range(start_layer, end_layer):
kv_cache = kv_caches[layer_id - start_layer]
key_cache = kv_cache[0].view(-1, num_heads, head_size)
value_cache = kv_cache[1].view(-1, num_heads, head_size)
current_slot_mapping = slot_mapping_flat[start_pos:end_pos]
# copy key into datadist
k = self.key_cache[layer_id][:, start_pos:end_pos, :, :]
new_k = key_cache[current_slot_mapping].unsqueeze(0)
torch_npu.scatter_update_(k, indices, new_k, axis=-2)
# copy value into datadist
val = self.value_cache[layer_id][:, start_pos:end_pos, :, :]
new_val = value_cache[current_slot_mapping].unsqueeze(0)
torch_npu.scatter_update_(val, indices, new_val, axis=-2)
# copy input into datadist
inp = self.input_cache[0][:, start_pos:end_pos, :, :]
new_inp = current_tokens.view(1, current_tokens.shape[0], 1, 1)
torch_npu.scatter_update_(inp, indices, new_inp, axis=-2)
# copy hidden into datadist
hid = self.hidden_cache[0][:, start_pos:end_pos, :, :]
hid_shape0, hid_shape1 = hidden_or_intermediate_states[
start_pos:end_pos].shape
new_hid = hidden_or_intermediate_states[start_pos:end_pos].view(
1, hid_shape0, 1, hid_shape1)
torch_npu.scatter_update_(hid, indices, new_hid, axis=-2)
logger.info("[rank%d][P]: KV send DONE.", torch.distributed.get_rank())
def recv_kv_caches_and_hidden_states(
self, model_executable: torch.nn.Module,
model_input: "ModelInputForGPUWithSamplingMetadata",
kv_caches: List[torch.Tensor]
) -> Tuple[Union[torch.Tensor, IntermediateTensors], bool,
"ModelInputForGPUWithSamplingMetadata"]:
bypass_model_exec = True
input_tokens_tensor = model_input.input_tokens
seq_lens = model_input.attn_metadata.seq_lens
slot_mapping = model_input.attn_metadata.slot_mapping.flatten()
hidden_or_intermediate_states_for_one_req = []
input_tokens_list = []
num_computed_tokens_list = []
start_pos_list = []
# get model config
start_layer = model_executable.model.start_layer
end_layer = model_executable.model.end_layer
model_config = model_executable.model.config
num_heads = int(model_config.num_key_value_heads / self.tp_size)
hidden_size = model_config.hidden_size
num_attention_heads = model_config.num_attention_heads
head_size = int(hidden_size / num_attention_heads)
num_layer = end_layer - start_layer
# get input_tensor_shape and hidden_shape
input_shape = (1, input_tokens_tensor.shape[0], 1, 1)
hidden_shape = (1, input_tokens_tensor.shape[0], 1, hidden_size)
kv_shape = (1, input_tokens_tensor.shape[0], num_heads, head_size)
kv_hidden_dtype = kv_caches[0].dtype
input_dtype = torch.int32
# Add LLM DataDist initialization
key_desc = llm_datadist.CacheDesc(
num_layer,
kv_shape,
TORCH_DTYPE_TO_NPU_DTYPE[kv_hidden_dtype],
seq_len_dim_index=-1)
value_desc = llm_datadist.CacheDesc(
num_layer,
kv_shape,
TORCH_DTYPE_TO_NPU_DTYPE[kv_hidden_dtype],
seq_len_dim_index=-1)
input_desc = llm_datadist.CacheDesc(
1,
input_shape,
TORCH_DTYPE_TO_NPU_DTYPE[input_dtype],
seq_len_dim_index=-1)
hidden_desc = llm_datadist.CacheDesc(
1,
hidden_shape,
TORCH_DTYPE_TO_NPU_DTYPE[kv_hidden_dtype],
seq_len_dim_index=-1)
self.decode_key_buffer = self.llm_datadist_engine.kv_transfer.allocate_cache(
key_desc)
self.decode_value_buffer = self.llm_datadist_engine.kv_transfer.allocate_cache(
value_desc)
self.decode_input_buffer = self.llm_datadist_engine.kv_transfer.allocate_cache(
input_desc)
self.decode_hidden_buffer = self.llm_datadist_engine.kv_transfer.allocate_cache(
hidden_desc)
key_buffer_addrs = self.decode_key_buffer.per_device_tensor_addrs[0]
value_buffer_addrs = self.decode_value_buffer.per_device_tensor_addrs[
0]
input_buffer_addrs = self.decode_input_buffer.per_device_tensor_addrs[
0]
hidden_buffer_addrs = self.decode_hidden_buffer.per_device_tensor_addrs[
0]
self.key_cache = torchair.llm_datadist.create_npu_tensors(
key_desc.shape, kv_hidden_dtype, key_buffer_addrs)
self.value_cache = torchair.llm_datadist.create_npu_tensors(
value_desc.shape, kv_hidden_dtype, value_buffer_addrs)
self.input_cache = torchair.llm_datadist.create_npu_tensors(
input_desc.shape, input_dtype, input_buffer_addrs)
self.hidden_cache = torchair.llm_datadist.create_npu_tensors(
hidden_desc.shape, kv_hidden_dtype, hidden_buffer_addrs)
key_cache_key = llm_datadist.CacheKeyByIdAndIndex(
self.cluster.remote_cluster_id, 1, 0)
value_cache_key = llm_datadist.CacheKeyByIdAndIndex(
self.cluster.remote_cluster_id, 2, 0)
input_cache_key = llm_datadist.CacheKeyByIdAndIndex(
self.cluster.remote_cluster_id, 3, 0)
hidden_cache_key = llm_datadist.CacheKeyByIdAndIndex(
self.cluster.remote_cluster_id, 4, 0)
self.llm_datadist_engine.kv_transfer.pull_cache(
key_cache_key, self.decode_key_buffer, 0)
self.llm_datadist_engine.kv_transfer.pull_cache(
value_cache_key, self.decode_value_buffer, 0)
self.llm_datadist_engine.kv_transfer.pull_cache(
input_cache_key, self.decode_input_buffer, 0)
self.llm_datadist_engine.kv_transfer.pull_cache(
hidden_cache_key, self.decode_hidden_buffer, 0)
keys = self.key_cache
values = self.value_cache
inputs = self.input_cache
hidden = self.hidden_cache
# enumerate different requests
for idx, slen in enumerate(seq_lens):
start_pos = sum(seq_lens[:idx])
end_pos = start_pos + slen
current_tokens = input_tokens_tensor[start_pos:end_pos]
num_tokens = slen
# collecting data for rebuilding the input
input_tokens_list.append(current_tokens)
start_pos_list.append(start_pos)
num_computed_tokens = inputs[0][0, start_pos:end_pos, 0,
0].shape[0]
num_computed_tokens_list.append(num_computed_tokens)
# check if both KV cache and the hidden states are received
# If not, need to redo the forwarding to compute missing states
if not all([(num_computed_tokens == num_tokens), hidden is not None
]):
bypass_model_exec = False
# update the end position based on how many tokens are cached.
end_pos = start_pos + num_computed_tokens
# put received KV caches into paged memory
for i in range(model_executable.model.start_layer,
model_executable.model.end_layer):
kv_cache = kv_caches[i - model_executable.model.start_layer]
key_cache, value_cache = kv_cache[0], kv_cache[1]
sliced_key = keys[i - model_executable.model.start_layer][
0, start_pos:end_pos, :, :]
sliced_value = values[i - model_executable.model.start_layer][
0, start_pos:end_pos, :, :]
torch_npu._npu_reshape_and_cache(
key=sliced_key,
value=sliced_value,
key_cache=key_cache,
value_cache=value_cache,
slot_indices=slot_mapping[start_pos:end_pos])
hidden_or_intermediate_states_for_one_req.append(
hidden[0][0, start_pos:end_pos, 0, :])
if not bypass_model_exec:
# Some of the KV cache is not retrieved
# Here we will fall back to normal model forwarding
# But optionally you can adjust model_input so that you only do
# prefilling on those tokens that are missing KV caches.
logger.info(
"[rank%d][D]: Failed to receive all KVs and hidden "
"states, redo model forwarding.", torch.distributed.get_rank())
hidden_or_intermediate_states = None
else:
logger.info(
"[rank%d][D]: Successfully received all KVs and hidden "
"states, skip model forwarding.", torch.distributed.get_rank())
hidden_or_intermediate_states = torch.cat(
hidden_or_intermediate_states_for_one_req, dim=0)
return hidden_or_intermediate_states, bypass_model_exec, model_input
def close(self, ):
self.llm_datadist_engine.data_dist.unlink_clusters([self.cluster],
5000)