[FEATURE] DuSwiftConnector support glm5 model PD （attention sparse_attn_indexer layer_name change ）

vllm/distributed/kv_transfer/kv_connector/v1/du/du_swift_connector.py

@@ -209,7 +209,7 @@ class DuSwiftConnector(KVConnectorBase_V1):
                 request_id (str): request id for log
             """
             dst_kv_cache_layer_shape = dst_kv_cache_layer.shape
-            if isinstance(attn_metadata, MLACommonMetadata) or all(isinstance(value, MLACommonMetadata) for value in attn_metadata.values()):
+            if isinstance(attn_metadata, MLACommonMetadata) or all(isinstance(value, MLACommonMetadata) for value in attn_metadata.values())  or dst_kv_cache_layer.ndim == 3:
                 num_pages = dst_kv_cache_layer_shape[0]
                 page_size = dst_kv_cache_layer_shape[1]
                 dst_kv_cache_layer = dst_kv_cache_layer.reshape(
@@ -379,7 +379,7 @@ class DuSwiftConnector(KVConnectorBase_V1):
 
         assert self.du_swift_engine is not None
 
-        is_mla = isinstance(attn_metadata, MLACommonMetadata)
+        is_mla = isinstance(attn_metadata, MLACommonMetadata) or kv_layer.ndim == 3
 
         def extract_kv_from_layer(
             layer: torch.Tensor,
@@ -390,7 +390,7 @@ class DuSwiftConnector(KVConnectorBase_V1):
             Assume the shape of the layer is (2, num_pages, page_size, xxx)
             if MLA is not used, and (num_pages, page_size, xxx) otherwise.
             """            
-            if isinstance(attn_metadata, MLACommonMetadata):
+            if isinstance(attn_metadata, MLACommonMetadata) or layer.ndim == 3:
                 num_pages, page_size = layer.shape[0], layer.shape[1]
                 return layer.reshape(num_pages * page_size, -1)[slot_mapping,
                                                                 ...]

vllm/model_executor/layers/sparse_attn_indexer.py

@@ -20,6 +20,10 @@ from vllm.v1.attention.ops.rocm_aiter_mla_sparse import indexer_k_bf16_cache_tri
 from vllm.v1.worker.workspace import current_workspace_manager
 from lightop import op, gemmopt
 
+from vllm.attention.utils.kv_transfer_utils import (
+    maybe_transfer_kv_layer,
+)
+
 if current_platform.is_cuda_alike():
     from vllm import _custom_ops as ops
 elif current_platform.is_xpu():
@@ -27,10 +31,10 @@ elif current_platform.is_xpu():
 
 logger = init_logger(__name__)
 
-
+@maybe_transfer_kv_layer
 def sparse_attn_indexer(
     hidden_states: torch.Tensor,
-    k_cache_prefix: str,
+    layer_name:str,
     kv_cache: torch.Tensor,
     q_fp8: torch.Tensor,
     k: torch.Tensor,
@@ -56,7 +60,7 @@ def sparse_attn_indexer(
         )
         return sparse_attn_indexer_fake(
             hidden_states,
-            k_cache_prefix,
+            layer_name,
             kv_cache,
             q_fp8,
             k,
@@ -69,7 +73,7 @@ def sparse_attn_indexer(
             total_seq_lens,
             topk_indices_buffer,
         )
-    attn_metadata = attn_metadata[k_cache_prefix]
+    attn_metadata = attn_metadata[layer_name]
     assert isinstance(attn_metadata, DeepseekV32IndexerMetadata)
     slot_mapping = attn_metadata.slot_mapping
     has_decode = attn_metadata.num_decodes > 0
@@ -282,7 +286,7 @@ def sparse_attn_indexer(
 
 def sparse_attn_indexer_fake(
     hidden_states: torch.Tensor,
-    k_cache_prefix: str,
+    layer_name: str,
     kv_cache: torch.Tensor,
     q_fp8: torch.Tensor,
     k: torch.Tensor,