Commit 05c3c3bb authored by zhuwenwen's avatar zhuwenwen
Browse files

[Bugfix][P/D]skip fusedmoe layer for start_load_kv

[Bugfix][P/D]Slove the problem that attn_medadata is not MLACommonMetadata
parent 711aa9d5
...@@ -142,7 +142,7 @@ class P2pNcclConnector(KVConnectorBase_V1): ...@@ -142,7 +142,7 @@ class P2pNcclConnector(KVConnectorBase_V1):
request_id (str): request id for log request_id (str): request id for log
""" """
dst_kv_cache_layer_shape = dst_kv_cache_layer.shape dst_kv_cache_layer_shape = dst_kv_cache_layer.shape
if isinstance(attn_metadata, MLACommonMetadata): if isinstance(attn_metadata, MLACommonMetadata) or all(isinstance(value, MLACommonMetadata) for value in attn_metadata.values()):
num_pages = dst_kv_cache_layer_shape[0] num_pages = dst_kv_cache_layer_shape[0]
page_size = dst_kv_cache_layer_shape[1] page_size = dst_kv_cache_layer_shape[1]
dst_kv_cache_layer = dst_kv_cache_layer.reshape( dst_kv_cache_layer = dst_kv_cache_layer.reshape(
...@@ -192,8 +192,16 @@ class P2pNcclConnector(KVConnectorBase_V1): ...@@ -192,8 +192,16 @@ class P2pNcclConnector(KVConnectorBase_V1):
# Load the KV for each request each layer # Load the KV for each request each layer
for request in metadata.requests: for request in metadata.requests:
for layer_name in forward_context.no_compile_layers: for layer_name in forward_context.no_compile_layers:
attn_layer = forward_context.no_compile_layers[layer_name] layer = forward_context.no_compile_layers[layer_name]
kv_cache_layer = attn_layer.kv_cache[ \
# Only process layers that have kv_cache
# attribute (attention layers) Skip non-attention
# layers like FusedMoE
kv_cache = getattr(layer, 'kv_cache', None)
if kv_cache is None:
continue
kv_cache_layer = kv_cache[ \
forward_context.virtual_engine] forward_context.virtual_engine]
kv_cache = self.p2p_nccl_engine.recv_tensor( kv_cache = self.p2p_nccl_engine.recv_tensor(
......
Markdown is supported
0% or .
You are about to add 0 people to the discussion. Proceed with caution.
Finish editing this message first!
Please register or to comment