update prefix cache interface and vdim padding strategy

vllm/attention/backends/mla/utils.py

@@ -535,23 +535,37 @@ class MLACommonImpl(MLAAttentionImpl[T], Generic[T]):
         # 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)
-
-        attn_output = flash_attn_varlen_func(
-            q=q,
-            k=k,
-            # v=v_padded,
-            v=v,
-            cu_seqlens_q=seq_start_loc,
-            cu_seqlens_k=seq_start_loc,
-            max_seqlen_q=max_prefill_seq_len,
-            max_seqlen_k=max_prefill_seq_len,
-            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\
-                .reshape(-1, self.num_heads * v.shape[-1])
+        v_padded = torch.nn.functional.pad(v, [0, (q.shape[-1] - v.shape[-1]-32)],
+                                           value=0)
+
+        if torch.cuda.get_device_properties(torch.cuda.current_device()).multi_processor_count == 120:
+            attn_output = flash_attn_varlen_func(
+                q=q,
+                k=k,
+                v=v_padded,
+                cu_seqlens_q=seq_start_loc,
+                cu_seqlens_k=seq_start_loc,
+                max_seqlen_q=max_prefill_seq_len,
+                max_seqlen_k=max_prefill_seq_len,
+                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])
+        else:
+            attn_output = flash_attn_varlen_func(
+                q=q,
+                k=k,
+                v=v,
+                cu_seqlens_q=seq_start_loc,
+                cu_seqlens_k=seq_start_loc,
+                max_seqlen_q=max_prefill_seq_len,
+                max_seqlen_k=max_prefill_seq_len,
+                softmax_scale=self.scale,
+                causal=True,
+            )
+            attn_output = attn_output\
+                    .reshape(-1, self.num_heads * v.shape[-1])
 
         return self.o_proj(attn_output)[0]

vllm/attention/backends/rocm_flash_attn.py

@@ -790,6 +790,7 @@ class ROCmFlashAttentionImpl(AttentionImpl):
                                prefill_meta.block_tables,
                                prefill_meta.query_start_loc,
                                prefill_meta.seq_lens_tensor,
+                               prefill_meta.context_lens_tensor,
                                prefill_meta.max_query_len,
                                self.alibi_slopes,
                                self.sliding_window[0],