[Bugfix] Fix vllm_flash_attn rotary import #17247

jeejeelee · 2025-04-27T08:14:02Z

Using qwen2-vl can reproduce this issue.
FIX #17263

Co-authored-by: Aaron Pham [email protected]

Signed-off-by: Jee Jee Li <[email protected]>

github-actions · 2025-04-27T08:14:12Z

👋 Hi! Thank you for contributing to the vLLM project.

💬 Join our developer Slack at https://slack.vllm.ai to discuss your PR in #pr-reviews, coordinate on features in #feat- channels, or join special interest groups in #sig- channels.

Just a reminder: PRs would not trigger full CI run by default. Instead, it would only run fastcheck CI which starts running only a small and essential subset of CI tests to quickly catch errors. You can run other CI tests on top of those by going to your fastcheck build on Buildkite UI (linked in the PR checks section) and unblock them. If you do not have permission to unblock, ping simon-mo or khluu to add you in our Buildkite org.

Once the PR is approved and ready to go, your PR reviewer(s) can run CI to test the changes comprehensively before merging.

To run CI, PR reviewers can either: Add ready label to the PR or enable auto-merge.

🚀

aarnphm

One slight change

aarnphm · 2025-04-27T13:25:31Z

setup.py

+        src_layers = os.path.join(self.build_lib, "vllm", "vllm_flash_attn",
+                                  "layers")
+        dst_layers = os.path.join("vllm", "vllm_flash_attn", "layers")
+        if os.path.exists(src_layers):
+            print(f"Copying directory {src_layers} to {dst_layers}")
+            if os.path.exists(dst_layers):
+                shutil.rmtree(dst_layers)
+            shutil.copytree(src_layers, dst_layers)
+
+        # copy vllm_flash_attn/ops dir
+        src_ops = os.path.join(self.build_lib, "vllm", "vllm_flash_attn",
+                               "ops")
+        dst_ops = os.path.join("vllm", "vllm_flash_attn", "ops")
+        if os.path.exists(src_ops):
+            print(f"Copying directory {src_ops} to {dst_ops}")
+            if os.path.exists(dst_ops):
+                shutil.rmtree(dst_ops)
+            shutil.copytree(src_ops, dst_ops)


You should use self.copy_tree instead

Suggested change

src_layers = os.path.join(self.build_lib, "vllm", "vllm_flash_attn",

"layers")

dst_layers = os.path.join("vllm", "vllm_flash_attn", "layers")

if os.path.exists(src_layers):

print(f"Copying directory {src_layers} to {dst_layers}")

if os.path.exists(dst_layers):

shutil.rmtree(dst_layers)

shutil.copytree(src_layers, dst_layers)

# copy vllm_flash_attn/ops dir

src_ops = os.path.join(self.build_lib, "vllm", "vllm_flash_attn",

"ops")

dst_ops = os.path.join("vllm", "vllm_flash_attn", "ops")

if os.path.exists(src_ops):

print(f"Copying directory {src_ops} to {dst_ops}")

if os.path.exists(dst_ops):

shutil.rmtree(dst_ops)

shutil.copytree(src_ops, dst_ops)

for folder in ("layers", "ops"):

src = os.path.join(self.build_lib, "vllm", "vllm_flash_attn", folder)

out = os.path.join("vllm", "vllm_flash_attn", folder)

self.copy_tree(src, out)

You can use this patch here

diff --git a/setup.py b/setup.py index 671ec7739..f5b3f5165 100755 --- a/setup.py +++ b/setup.py @@ -283,6 +283,13 @@ class cmake_build_ext(build_ext): dst_file = os.path.join("vllm/vllm_flash_attn", basename) print(f"Copying {file} to {dst_file}") self.copy_file(file, dst_file) + # We will need to copy these folders to use its ops for qwen-vl + for folder in ("layers", "ops"): + src = os.path.join(self.build_lib, 'vllm', 'vllm_flash_attn', + folder) + out = os.path.join('vllm', 'vllm_flash_attn', folder) + print(f"Copying {folder} from vllm/vllm_flash_attn") + self.copy_tree(src, out) class repackage_wheel(build_ext): @@ -402,6 +409,32 @@ class repackage_wheel(build_ext): package_data[package_name].append(file_name) + # Extract and include the layers and ops folders for qwen-vl + folders_to_copy = {"layers", "ops"} + for folder in folders_to_copy: + folder_path = f"vllm/vllm_flash_attn/{folder}" + folder_files = [ + f for f in wheel.filelist + if f.filename.startswith(folder_path) + ] + + if folder_files: + print(f"Include {folder} folder from vllm/vllm_flash_attn") + for file in folder_files: + wheel.extract(file) + + # Add the file to package_data if it's not a Python file + rel_path = file.filename.split("/") + # vllm/vllm_flash_attn/folder/file + if len(rel_path) >= 4: + package_name = "vllm.vllm_flash_attn." + folder + file_name = rel_path[-1] + + if not file_name.endswith(".py"): + if package_name not in package_data: + package_data[package_name] = [] + package_data[package_name].append(file_name) + def _is_hpu() -> bool: # if VLLM_TARGET_DEVICE env var was set explicitly, skip HPU autodetection

Signed-off-by: Jee Jee Li <[email protected]>

aarnphm

Thanks. LGTM

Can you also add a Co-authored-by: Aaron Pham <[email protected]> to the description.

aarnphm · 2025-04-27T14:51:33Z

Hmm, @jeejeelee there is a issue with building on CPU here?

aarnphm · 2025-04-27T14:53:47Z

Ah we need to gated the copy ovee in _is_cuda() only here.

27fdbeea7 - chore: only gated in CUDA (HEAD -> fix-flash-att-rotray)

Signed-off-by: Aaron Pham <[email protected]>
(G) Aaron Pham <[email protected]> (Sun Apr 27 14:57:44 2025 +0000)


diff --git a/setup.py b/setup.py
index 34be6a641..8e18c8215 100755
--- a/setup.py
+++ b/setup.py
@@ -280,13 +280,14 @@ class cmake_build_ext(build_ext):
             print(f"Copying {file} to {dst_file}")
             self.copy_file(file, dst_file)
 
-        # copy these folders to use the vllm_flash_attn rotary_kernel.
-        for folder in ("layers", "ops"):
-            src = os.path.join(self.build_lib, "vllm", "vllm_flash_attn",
-                               folder)
-            out = os.path.join("vllm", "vllm_flash_attn", folder)
-            print(f"Copying {folder} from vllm/vllm_flash_attn")
-            self.copy_tree(src, out)
+        if _is_cuda():
+            # copy these folders to use the vllm_flash_attn rotary_kernel.
+            for folder in ("layers", "ops"):
+                src = os.path.join(self.build_lib, "vllm", "vllm_flash_attn",
+                                   folder)
+                out = os.path.join("vllm", "vllm_flash_attn", folder)
+                print(f"Copying {folder} from vllm/vllm_flash_attn")
+                self.copy_tree(src, out)
 
 
 class repackage_wheel(build_ext):
@@ -409,29 +410,34 @@ class repackage_wheel(build_ext):
                 package_data[package_name].append(file_name)
 
             # Extract and include the layers and ops of rotary embedding.
-            folders_to_copy = {"layers", "ops"}
-            for folder in folders_to_copy:
-                folder_path = f"vllm/vllm_flash_attn/{folder}"
-                folder_files = [
-                    f for f in wheel.filelist
-                    if f.filename.startswith(folder_path)
-                ]
-
-                if folder_files:
-                    print(f"Include {folder} folder from vllm/vllm_flash_attn")
-                    for file in folder_files:
-                        wheel.extract(file)
-
-                        # Add the file to package_data if it's not a Python file
-                        rel_path = file.filename.split("/")
-                        # vllm/vllm_flash_attn/folder/file
-                        if len(rel_path) >= 4:
-                            package_name = "vllm.vllm_flash_attn." + folder
-                            file_name = rel_path[-1]
-                            if not file_name.endswith(".py"):
-                                if package_name not in package_data:
-                                    package_data[package_name] = []
-                                package_data[package_name].append(file_name)
+            if _is_cuda():
+                folders_to_copy = {"layers", "ops"}
+                for folder in folders_to_copy:
+                    folder_path = f"vllm/vllm_flash_attn/{folder}"
+                    folder_files = [
+                        f for f in wheel.filelist
+                        if f.filename.startswith(folder_path)
+                    ]
+
+                    if folder_files:
+                        print(
+                            f"Include {folder} folder from vllm/vllm_flash_attn"
+                        )
+                        for file in folder_files:
+                            wheel.extract(file)
+
+                            # Add the file to package_data if
+                            # it's not a Python file
+                            rel_path = file.filename.split("/")
+                            # vllm/vllm_flash_attn/folder/file
+                            if len(rel_path) >= 4:
+                                package_name = "vllm.vllm_flash_attn." + folder
+                                file_name = rel_path[-1]
+                                if not file_name.endswith(".py"):
+                                    if package_name not in package_data:
+                                        package_data[package_name] = []
+                                    package_data[package_name].append(
+                                        file_name)
 
 
 def _is_hpu() -> bool:

LucasWilkinson

Can you please add a description of the bug to the PR comment

LucasWilkinson · 2025-04-27T15:31:34Z

setup.py

@@ -280,6 +280,14 @@ def run(self):
            print(f"Copying {file} to {dst_file}")
            self.copy_file(file, dst_file)

+        # copy these folders to use the vllm_flash_attn rotary_kernel.
+        for folder in ("layers", "ops"):


this seems unnecessary, why not just do:

files = glob.glob( os.path.join(self.build_lib, "vllm", "vllm_flash_attn", "**", "*.py"))

on line 275-276 instead?

this will only grab all of the python files, then we still need to check if the folder exists or not, then copy over here.

Would this conflict with #17260?

LucasWilkinson · 2025-04-27T15:35:49Z

setup.py

@@ -400,6 +408,31 @@ def run(self) -> None:

                package_data[package_name].append(file_name)

+            # Extract and include the layers and ops of rotary embedding.
+            folders_to_copy = {"layers", "ops"}
+            for folder in folders_to_copy:


This also seems unnecessarily messing / complicated, I think we can just use glob.translate (https://docs.python.org/3/library/glob.html#glob.translate) and grab the file list like use the same glob pattern mentioned above

i agree, overlook on my end.

jeejeelee · 2025-04-27T15:43:02Z

Thanks. LGTM

Can you also add a Co-authored-by: Aaron Pham <[email protected]> to the description.

Done

jeejeelee · 2025-04-27T15:54:37Z

Ah we need to gated the copy ovee in _is_cuda() only here.

27fdbeea7 - chore: only gated in CUDA (HEAD -> fix-flash-att-rotray)

Signed-off-by: Aaron Pham <[email protected]>
(G) Aaron Pham <[email protected]> (Sun Apr 27 14:57:44 2025 +0000)


diff --git a/setup.py b/setup.py
index 34be6a641..8e18c8215 100755
--- a/setup.py
+++ b/setup.py
@@ -280,13 +280,14 @@ class cmake_build_ext(build_ext):
             print(f"Copying {file} to {dst_file}")
             self.copy_file(file, dst_file)
 
-        # copy these folders to use the vllm_flash_attn rotary_kernel.
-        for folder in ("layers", "ops"):
-            src = os.path.join(self.build_lib, "vllm", "vllm_flash_attn",
-                               folder)
-            out = os.path.join("vllm", "vllm_flash_attn", folder)
-            print(f"Copying {folder} from vllm/vllm_flash_attn")
-            self.copy_tree(src, out)
+        if _is_cuda():
+            # copy these folders to use the vllm_flash_attn rotary_kernel.
+            for folder in ("layers", "ops"):
+                src = os.path.join(self.build_lib, "vllm", "vllm_flash_attn",
+                                   folder)
+                out = os.path.join("vllm", "vllm_flash_attn", folder)
+                print(f"Copying {folder} from vllm/vllm_flash_attn")
+                self.copy_tree(src, out)
 
 
 class repackage_wheel(build_ext):
@@ -409,29 +410,34 @@ class repackage_wheel(build_ext):
                 package_data[package_name].append(file_name)
 
             # Extract and include the layers and ops of rotary embedding.
-            folders_to_copy = {"layers", "ops"}
-            for folder in folders_to_copy:
-                folder_path = f"vllm/vllm_flash_attn/{folder}"
-                folder_files = [
-                    f for f in wheel.filelist
-                    if f.filename.startswith(folder_path)
-                ]
-
-                if folder_files:
-                    print(f"Include {folder} folder from vllm/vllm_flash_attn")
-                    for file in folder_files:
-                        wheel.extract(file)
-
-                        # Add the file to package_data if it's not a Python file
-                        rel_path = file.filename.split("/")
-                        # vllm/vllm_flash_attn/folder/file
-                        if len(rel_path) >= 4:
-                            package_name = "vllm.vllm_flash_attn." + folder
-                            file_name = rel_path[-1]
-                            if not file_name.endswith(".py"):
-                                if package_name not in package_data:
-                                    package_data[package_name] = []
-                                package_data[package_name].append(file_name)
+            if _is_cuda():
+                folders_to_copy = {"layers", "ops"}
+                for folder in folders_to_copy:
+                    folder_path = f"vllm/vllm_flash_attn/{folder}"
+                    folder_files = [
+                        f for f in wheel.filelist
+                        if f.filename.startswith(folder_path)
+                    ]
+
+                    if folder_files:
+                        print(
+                            f"Include {folder} folder from vllm/vllm_flash_attn"
+                        )
+                        for file in folder_files:
+                            wheel.extract(file)
+
+                            # Add the file to package_data if
+                            # it's not a Python file
+                            rel_path = file.filename.split("/")
+                            # vllm/vllm_flash_attn/folder/file
+                            if len(rel_path) >= 4:
+                                package_name = "vllm.vllm_flash_attn." + folder
+                                file_name = rel_path[-1]
+                                if not file_name.endswith(".py"):
+                                    if package_name not in package_data:
+                                        package_data[package_name] = []
+                                    package_data[package_name].append(
+                                        file_name)
 
 
 def _is_hpu() -> bool:

IIUC, no need to add gate, these changes are just a supplement to vllm_flash_attn copy

aarnphm · 2025-04-27T16:46:45Z

This is superceded by #17267

Done

eddc994

Signed-off-by: Jee Jee Li <[email protected]>

jeejeelee requested a review from DarkLight1337 April 27, 2025 08:14

mergify bot added the ci/build label Apr 27, 2025

DarkLight1337 requested a review from aarnphm April 27, 2025 08:32

DarkLight1337 added the ready ONLY add when PR is ready to merge/full CI is needed label Apr 27, 2025

jeejeelee mentioned this pull request Apr 27, 2025

[Bug]: Cannot found flash_attn_interface after adding stub files #17246

Closed

1 task

DarkLight1337 mentioned this pull request Apr 27, 2025

[Chore] ignore override default __init__.py when building from source #17260

Closed

aarnphm suggested changes Apr 27, 2025

View reviewed changes

Modify code

09507fc

Signed-off-by: Jee Jee Li <[email protected]>

jeejeelee requested a review from aarnphm April 27, 2025 14:47

aarnphm self-assigned this Apr 27, 2025

aarnphm approved these changes Apr 27, 2025

View reviewed changes

DarkLight1337 approved these changes Apr 27, 2025

View reviewed changes

LucasWilkinson requested changes Apr 27, 2025

View reviewed changes

jeejeelee requested a review from LucasWilkinson April 27, 2025 16:03

aarnphm self-requested a review April 27, 2025 16:07

LucasWilkinson mentioned this pull request Apr 27, 2025

[BugFix] Fix vllm_flash_attn install issues #17267

Merged

aarnphm closed this Apr 27, 2025

jeejeelee deleted the fix-flash-att-rotray branch April 28, 2025 00:21

Provide feedback

Saved searches

Use saved searches to filter your results more quickly

[Bugfix] Fix vllm_flash_attn rotary import #17247

[Bugfix] Fix vllm_flash_attn rotary import #17247

jeejeelee commented Apr 27, 2025 •

edited by github-actions bot

Loading

github-actions bot commented Apr 27, 2025

aarnphm left a comment •

edited

Loading

aarnphm Apr 27, 2025

aarnphm Apr 27, 2025

jeejeelee Apr 27, 2025

aarnphm left a comment

aarnphm commented Apr 27, 2025

aarnphm commented Apr 27, 2025 •

edited

Loading

LucasWilkinson left a comment

LucasWilkinson Apr 27, 2025 •

edited

Loading

aarnphm Apr 27, 2025

jeejeelee Apr 27, 2025

LucasWilkinson Apr 27, 2025

aarnphm Apr 27, 2025

jeejeelee commented Apr 27, 2025

jeejeelee commented Apr 27, 2025

aarnphm commented Apr 27, 2025

[Bugfix] Fix vllm_flash_attn rotary import #17247

[Bugfix] Fix vllm_flash_attn rotary import #17247

Conversation

jeejeelee commented Apr 27, 2025 • edited by github-actions bot Loading

github-actions bot commented Apr 27, 2025

aarnphm left a comment • edited Loading

Choose a reason for hiding this comment

aarnphm Apr 27, 2025

Choose a reason for hiding this comment

aarnphm Apr 27, 2025

Choose a reason for hiding this comment

jeejeelee Apr 27, 2025

Choose a reason for hiding this comment

aarnphm left a comment

Choose a reason for hiding this comment

aarnphm commented Apr 27, 2025

aarnphm commented Apr 27, 2025 • edited Loading

LucasWilkinson left a comment

Choose a reason for hiding this comment

LucasWilkinson Apr 27, 2025 • edited Loading

Choose a reason for hiding this comment

aarnphm Apr 27, 2025

Choose a reason for hiding this comment

jeejeelee Apr 27, 2025

Choose a reason for hiding this comment

LucasWilkinson Apr 27, 2025

Choose a reason for hiding this comment

aarnphm Apr 27, 2025

Choose a reason for hiding this comment

jeejeelee commented Apr 27, 2025

jeejeelee commented Apr 27, 2025

aarnphm commented Apr 27, 2025

jeejeelee commented Apr 27, 2025 •

edited by github-actions bot

Loading

aarnphm left a comment •

edited

Loading

aarnphm commented Apr 27, 2025 •

edited

Loading

LucasWilkinson Apr 27, 2025 •

edited

Loading