Fix import for vllm_flash_attn

Signed-off-by: Vinay Damodaran <[email protected]>

Author: vrdn-23

vllm/model_executor/models/qwen2_5_vl.py

@@ -321,7 +321,7 @@ def forward(
                                       "(b s) ... -> b s ...",
                                       b=batch_size)
         elif self.attn_backend == _Backend.FLASH_ATTN_VLLM_V1:
-            from vllm_flash_attn.flash_attn_interface import flash_attn_varlen_func
+            from vllm.vllm_flash_attn.flash_attn_interface import flash_attn_varlen_func
 
             q, k, v = (rearrange(x, "b s ... -> (b s) ...") for x in [q, k, v])
 

vllm/model_executor/models/qwen2_vl.py

@@ -343,7 +343,7 @@ def forward(
                                       "(b s) ... -> b s ...",
                                       b=batch_size)
         elif self.attn_backend == _Backend.FLASH_ATTN_VLLM_V1:
-            from vllm_flash_attn.flash_attn_interface import flash_attn_varlen_func
+            from vllm.vllm_flash_attn.flash_attn_interface import flash_attn_varlen_func
 
             q, k, v = (rearrange(x, "b s ... -> (b s) ...") for x in [q, k, v])