fix: added stubs to address vllm-project#17159

Signed-off-by: Aaron Pham <[email protected]>

Author: aarnphm

pyproject.toml

@@ -58,7 +58,8 @@ ignore_patterns = [
 line-length = 80
 exclude = [
     # External file, leaving license intact
-    "examples/other/fp8/quantizer/quantize.py"
+    "examples/other/fp8/quantizer/quantize.py",
+    "vllm/vllm_flash_attn/flash_attn_interface.pyi"
 ]
 
 [tool.ruff.lint.per-file-ignores]

setup.py

@@ -378,7 +378,6 @@ def run(self) -> None:
                 "vllm/vllm_flash_attn/_vllm_fa2_C.abi3.so",
                 "vllm/vllm_flash_attn/_vllm_fa3_C.abi3.so",
                 "vllm/vllm_flash_attn/flash_attn_interface.py",
-                "vllm/vllm_flash_attn/__init__.py",
                 "vllm/cumem_allocator.abi3.so",
                 # "vllm/_version.py", # not available in nightly wheels yet
             ]

vllm/vllm_flash_attn/__init__.py

@@ -0,0 +1,14 @@
+# SPDX-License-Identifier: Apache-2.0
+
+from .flash_attn_interface import (fa_version_unsupported_reason,
+                                   flash_attn_varlen_func,
+                                   flash_attn_with_kvcache,
+                                   get_scheduler_metadata,
+                                   is_fa_version_supported, sparse_attn_func,
+                                   sparse_attn_varlen_func)
+
+__all__ = [
+    'flash_attn_varlen_func', 'flash_attn_with_kvcache',
+    'get_scheduler_metadata', 'sparse_attn_func', 'sparse_attn_varlen_func',
+    'is_fa_version_supported', 'fa_version_unsupported_reason'
+]