update prefix cache interface and vdim padding strategy

vllm/attention/backends/mla/utils.py

@@ -535,11 +535,28 @@ 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)
+        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_padded,
                 v=v,
                 cu_seqlens_q=seq_start_loc,
                 cu_seqlens_k=seq_start_loc,
@@ -548,9 +565,6 @@ 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\
                     .reshape(-1, self.num_heads * v.shape[-1])
 

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],