support vdim 128

README.md

@@ -12,7 +12,7 @@ vLLM是一个快速且易于使用的LLM推理和服务库,使用PageAttention
 | :------: | :------: | :------: | :------: |:------: |
 | LlamaForCausalLM      | Llama 3.1,Llama 3,Llama 2,Llama,Yi,Codellama,DeepSeek-R1-Distill-Llama     | Yes | Yes | Yes |  
 | QWenLMHeadModel       | QWen,Qwen-VL                                                | Yes | Yes | Yes |
-| Qwen2ForCausalLM      | QWen2,QWen1.5,CodeQwen1.5,DeepSeek-R1-Distill-Qwen          | Yes | Yes | Yes |
+| Qwen2ForCausalLM      | QWen2,QWen1.5,CodeQwen1.5,DeepSeek-R1-Distill-Qwen,gte_Qwen2-1.5B-instruct          | Yes | Yes | Yes |
 | ChatGLMModel          | glm-4v-9b,chatglm3,chatglm2                                 | Yes | No  | Yes |
 | DeepseekForCausalLM   | Deepseek                                                    | Yes | No  | -   |
 | DeepseekV2ForCausalLM | DeepSeek-V2                                                 | Yes | No  | -   |
@@ -31,6 +31,9 @@ vLLM是一个快速且易于使用的LLM推理和服务库,使用PageAttention
 | Qwen2VLForConditionalGeneration     | Qwen2-VL                                      | Yes | No  | Yes |
 | MiniCPMV                            | MiniCPM-V                                     | Yes | No  | -   |
 | Phi3VForCausalLM                    | Phi-3.5-vision                                | Yes | No  | -   |
+| BertModel                           | bge-large-zh-v1.5                             | Yes | No  | -   |
+| XLMRobertaModel                     | bge-m3                                        | Yes | No  | -   |
+| XLMRobertaForSequenceClassification | bge-reranker-v2-m3                            | Yes | No  | -   |
 
 
 ## 安装

vllm/attention/backends/mla/utils.py

@@ -533,13 +533,14 @@ class MLACommonImpl(MLAAttentionImpl[T], Generic[T]):
 
         # For MLA the v head dim is smaller than qk head dim so we pad out
         # v with 0s to match the qk head dim
-        v_padded = torch.nn.functional.pad(v, [0, q.shape[-1] - v.shape[-1]],
-                                           value=0)
+        # v_padded = torch.nn.functional.pad(v, [0, q.shape[-1] - v.shape[-1]],
+        #                                    value=0)
 
         attn_output = flash_attn_varlen_func(
             q=q,
             k=k,
-            v=v_padded,
+            # v=v_padded,
+            v=v,
             cu_seqlens_q=seq_start_loc,
             cu_seqlens_k=seq_start_loc,
             max_seqlen_q=max_prefill_seq_len,
@@ -547,8 +548,10 @@ class MLACommonImpl(MLAAttentionImpl[T], Generic[T]):
             softmax_scale=self.scale,
             causal=True,
         )
+        # attn_output = attn_output\
+        #     .view(-1, self.num_heads, q.shape[-1])[..., :v.shape[-1]]\
+        #         .reshape(-1, self.num_heads * v.shape[-1])
         attn_output = attn_output\
-            .view(-1, self.num_heads, q.shape[-1])[..., :v.shape[-1]]\
                 .reshape(-1, self.num_heads * v.shape[-1])
 
         return self.o_proj(attn_output)[0]