Kernels:

kernels-community
/

flash-attn4

Trusted publisher

kernels-bot commited on 18 days ago

Commit

eca38cd

verified ·

1 Parent(s): 6337869

Uploaded using `kernel-builder`.

Files changed (2) hide show

build/torch-cuda/_ops.py CHANGED Viewed

@@ -1,8 +1,38 @@
 import torch
-ops = torch.ops._flash_attn4_efe2479
-def add_op_namespace_prefix(op_name: str):
     """
     Prefix op by namespace.
     """
-    return f"_flash_attn4_efe2479::{op_name}"

 import torch
+def get_backend() -> str:
+    """Detect the backend by inspecting torch."""
+    import torch
+    if hasattr(torch, "neuron"):
+        # Needs to be sorted before specific Torch builds, since Neuron
+        # extension can be loaded into e.g. CUDA Torch builds.
+        return "neuron"
+    elif torch.version.cuda is not None:
+        return "cuda"
+    elif torch.version.hip is not None:
+        return "rocm"
+    elif torch.backends.mps.is_available():
+        return "metal"
+    elif hasattr(torch.version, "xpu") and torch.version.xpu is not None:
+        return "xpu"
+    else:
+        return "cpu"
+def _find_ops_name() -> str:
+    kernel_name = "flash_attn4"
+    unique_id = "86f75d9"
+    backend = get_backend()
+    return f"_{kernel_name}_{backend}_{unique_id}"
+_OPS_NAME = _find_ops_name()
+ops = getattr(torch.ops, _OPS_NAME)
+def add_op_namespace_prefix(op_name: str) -> str:
     """
     Prefix op by namespace.
     """
+    return f"{_OPS_NAME}::{op_name}"

build/torch-cuda/metadata.json CHANGED Viewed

@@ -1,4 +1,6 @@
 {
   "version": 0,
   "license": "BSD-3-Clause",
   "python-depends": [

 {
+  "name": "flash-attn4",
+  "id": "_flash_attn4_cuda_86f75d9",
   "version": 0,
   "license": "BSD-3-Clause",
   "python-depends": [