fix: correct finding the kv cache shape for mha

Signed-off-by: Jade Zheng <zheng.shoujian@outlook.com>

Author: jianzs

vllm_ascend/distributed/llmdatadist_connector_v1.py

@@ -785,11 +785,12 @@ def _extract_kv_from_layer(
         """
         if is_mla:
             num_heads, head_dim = kv_cache_layer.shape[
-                2], kv_cache_layer.shape[3]
+                -2], kv_cache_layer.shape[-1]
             return kv_cache_layer.view(-1, num_heads, head_dim)[slot_mapping,
                                                                 ...]
 
-        num_heads, head_dim = kv_cache_layer.shape[2], kv_cache_layer.shape[3]
+        num_heads, head_dim = kv_cache_layer.shape[-2], kv_cache_layer.shape[
+            -1]
         return kv_cache_layer.view(2, -1, num_heads, head_dim)[:, slot_mapping,
                                                                ...]