feat: pass vars into fwd and include build

.gitattributes

@@ -33,3 +33,4 @@ saved_model/**/* filter=lfs diff=lfs merge=lfs -text
 *.zip filter=lfs diff=lfs merge=lfs -text
 *.zst filter=lfs diff=lfs merge=lfs -text
 *tfevents* filter=lfs diff=lfs merge=lfs -text
+*.so filter=lfs diff=lfs merge=lfs -text

build.toml

@@ -13,13 +13,24 @@ src = [
   "flash_attn/src/hardware_info.h",
   "flash_attn/src/flash.h",
   "flash_attn/src/static_switch.h",
-  #
   "flash_attn/src/alibi.h",
   "flash_attn/src/block_info.h",
   "flash_attn/src/dropout.h",
+  "flash_attn/src/flash.h",
+  "flash_attn/src/generate_kernels.py",
+  "flash_attn/src/hardware_info.h",
+  "flash_attn/src/kernel_traits.h",
+  "flash_attn/src/mask.h",
+  "flash_attn/src/namespace_config.h",
+  "flash_attn/src/philox.cuh",
+  "flash_attn/src/philox_unpack.cuh",
+  "flash_attn/src/rotary.h",
+  "flash_attn/src/softmax.h",
+  "flash_attn/src/static_switch.h",
+  "flash_attn/src/utils.h",
+
+  ## TODO: include bwd kernels
 
-  # TODO: dont skip bwd kernels
-  
   # "flash_attn/src/flash_bwd_hdim128_bf16_causal_sm80.cu",
   # "flash_attn/src/flash_bwd_hdim128_bf16_sm80.cu",
   # "flash_attn/src/flash_bwd_hdim128_fp16_causal_sm80.cu",
@@ -52,75 +63,65 @@ src = [
   # "flash_attn/src/flash_bwd_launch_template.h",
   # "flash_attn/src/flash_bwd_preprocess_kernel.h",
 
-  "flash_attn/src/flash_fwd_hdim128_bf16_causal_sm80.cu",
-  "flash_attn/src/flash_fwd_hdim128_bf16_sm80.cu",
-  "flash_attn/src/flash_fwd_hdim128_fp16_causal_sm80.cu",
-  "flash_attn/src/flash_fwd_hdim128_fp16_sm80.cu",
-  "flash_attn/src/flash_fwd_hdim160_bf16_causal_sm80.cu",
-  "flash_attn/src/flash_fwd_hdim160_bf16_sm80.cu",
-  "flash_attn/src/flash_fwd_hdim160_fp16_causal_sm80.cu",
-  "flash_attn/src/flash_fwd_hdim160_fp16_sm80.cu",
-  "flash_attn/src/flash_fwd_hdim192_bf16_causal_sm80.cu",
-  "flash_attn/src/flash_fwd_hdim192_bf16_sm80.cu",
-  "flash_attn/src/flash_fwd_hdim192_fp16_causal_sm80.cu",
-  "flash_attn/src/flash_fwd_hdim192_fp16_sm80.cu",
-  "flash_attn/src/flash_fwd_hdim256_bf16_causal_sm80.cu",
-  "flash_attn/src/flash_fwd_hdim256_bf16_sm80.cu",
-  "flash_attn/src/flash_fwd_hdim256_fp16_causal_sm80.cu",
-  "flash_attn/src/flash_fwd_hdim256_fp16_sm80.cu",
-  "flash_attn/src/flash_fwd_hdim32_bf16_causal_sm80.cu",
-  "flash_attn/src/flash_fwd_hdim32_bf16_sm80.cu",
-  "flash_attn/src/flash_fwd_hdim32_fp16_causal_sm80.cu",
-  "flash_attn/src/flash_fwd_hdim32_fp16_sm80.cu",
-  "flash_attn/src/flash_fwd_hdim64_bf16_causal_sm80.cu",
-  "flash_attn/src/flash_fwd_hdim64_bf16_sm80.cu",
-  "flash_attn/src/flash_fwd_hdim64_fp16_causal_sm80.cu",
-  "flash_attn/src/flash_fwd_hdim64_fp16_sm80.cu",
-  "flash_attn/src/flash_fwd_hdim96_bf16_causal_sm80.cu",
-  "flash_attn/src/flash_fwd_hdim96_bf16_sm80.cu",
-  "flash_attn/src/flash_fwd_hdim96_fp16_causal_sm80.cu",
+  ## TODO: include fwd kernels
+
+  # "flash_attn/src/flash_fwd_hdim128_bf16_causal_sm80.cu",
+  # "flash_attn/src/flash_fwd_hdim128_bf16_sm80.cu",
+  # "flash_attn/src/flash_fwd_hdim128_fp16_causal_sm80.cu",
+  # "flash_attn/src/flash_fwd_hdim128_fp16_sm80.cu",
+  # "flash_attn/src/flash_fwd_hdim160_bf16_causal_sm80.cu",
+  # "flash_attn/src/flash_fwd_hdim160_bf16_sm80.cu",
+  # "flash_attn/src/flash_fwd_hdim160_fp16_causal_sm80.cu",
+  # "flash_attn/src/flash_fwd_hdim160_fp16_sm80.cu",
+  # "flash_attn/src/flash_fwd_hdim192_bf16_causal_sm80.cu",
+  # "flash_attn/src/flash_fwd_hdim192_bf16_sm80.cu",
+  # "flash_attn/src/flash_fwd_hdim192_fp16_causal_sm80.cu",
+  # "flash_attn/src/flash_fwd_hdim192_fp16_sm80.cu",
+  # "flash_attn/src/flash_fwd_hdim256_bf16_causal_sm80.cu",
+  # "flash_attn/src/flash_fwd_hdim256_bf16_sm80.cu",
+  # "flash_attn/src/flash_fwd_hdim256_fp16_causal_sm80.cu",
+  # "flash_attn/src/flash_fwd_hdim256_fp16_sm80.cu",
+  # "flash_attn/src/flash_fwd_hdim32_bf16_causal_sm80.cu",
+  # "flash_attn/src/flash_fwd_hdim32_bf16_sm80.cu",
+  # "flash_attn/src/flash_fwd_hdim32_fp16_causal_sm80.cu",
+  # "flash_attn/src/flash_fwd_hdim32_fp16_sm80.cu",
+  # "flash_attn/src/flash_fwd_hdim64_bf16_causal_sm80.cu",
+  # "flash_attn/src/flash_fwd_hdim64_bf16_sm80.cu",
+  # "flash_attn/src/flash_fwd_hdim64_fp16_causal_sm80.cu",
+  # "flash_attn/src/flash_fwd_hdim64_fp16_sm80.cu",
+  # "flash_attn/src/flash_fwd_hdim96_bf16_causal_sm80.cu",
+  # "flash_attn/src/flash_fwd_hdim96_bf16_sm80.cu",
+  # "flash_attn/src/flash_fwd_hdim96_fp16_causal_sm80.cu",
   "flash_attn/src/flash_fwd_hdim96_fp16_sm80.cu",
   "flash_attn/src/flash_fwd_kernel.h",
   "flash_attn/src/flash_fwd_launch_template.h",
-  "flash_attn/src/flash_fwd_split_hdim128_bf16_causal_sm80.cu",
-  "flash_attn/src/flash_fwd_split_hdim128_bf16_sm80.cu",
-  "flash_attn/src/flash_fwd_split_hdim128_fp16_causal_sm80.cu",
-  "flash_attn/src/flash_fwd_split_hdim128_fp16_sm80.cu",
-  "flash_attn/src/flash_fwd_split_hdim160_bf16_causal_sm80.cu",
-  "flash_attn/src/flash_fwd_split_hdim160_bf16_sm80.cu",
-  "flash_attn/src/flash_fwd_split_hdim160_fp16_causal_sm80.cu",
-  "flash_attn/src/flash_fwd_split_hdim160_fp16_sm80.cu",
-  "flash_attn/src/flash_fwd_split_hdim192_bf16_causal_sm80.cu",
-  "flash_attn/src/flash_fwd_split_hdim192_bf16_sm80.cu",
-  "flash_attn/src/flash_fwd_split_hdim192_fp16_causal_sm80.cu",
-  "flash_attn/src/flash_fwd_split_hdim192_fp16_sm80.cu",
-  "flash_attn/src/flash_fwd_split_hdim256_bf16_causal_sm80.cu",
-  "flash_attn/src/flash_fwd_split_hdim256_bf16_sm80.cu",
-  "flash_attn/src/flash_fwd_split_hdim256_fp16_causal_sm80.cu",
-  "flash_attn/src/flash_fwd_split_hdim256_fp16_sm80.cu",
-  "flash_attn/src/flash_fwd_split_hdim32_bf16_causal_sm80.cu",
-  "flash_attn/src/flash_fwd_split_hdim32_bf16_sm80.cu",
-  "flash_attn/src/flash_fwd_split_hdim32_fp16_causal_sm80.cu",
-  "flash_attn/src/flash_fwd_split_hdim32_fp16_sm80.cu",
-  "flash_attn/src/flash_fwd_split_hdim64_bf16_causal_sm80.cu",
-  "flash_attn/src/flash_fwd_split_hdim64_bf16_sm80.cu",
-  "flash_attn/src/flash_fwd_split_hdim64_fp16_causal_sm80.cu",
-  "flash_attn/src/flash_fwd_split_hdim64_fp16_sm80.cu",
-  "flash_attn/src/flash_fwd_split_hdim96_bf16_causal_sm80.cu",
-  "flash_attn/src/flash_fwd_split_hdim96_bf16_sm80.cu",
-  "flash_attn/src/flash_fwd_split_hdim96_fp16_causal_sm80.cu",
-  "flash_attn/src/flash_fwd_split_hdim96_fp16_sm80.cu",
-  "flash_attn/src/flash.h",
-  "flash_attn/src/generate_kernels.py",
-  "flash_attn/src/hardware_info.h",
-  "flash_attn/src/kernel_traits.h",
-  "flash_attn/src/mask.h",
-  "flash_attn/src/namespace_config.h",
-  "flash_attn/src/philox.cuh",
-  "flash_attn/src/philox_unpack.cuh",
-  "flash_attn/src/rotary.h",
-  "flash_attn/src/softmax.h",
-  "flash_attn/src/static_switch.h",
-  "flash_attn/src/utils.h",
+  # "flash_attn/src/flash_fwd_split_hdim128_bf16_causal_sm80.cu",
+  # "flash_attn/src/flash_fwd_split_hdim128_bf16_sm80.cu",
+  # "flash_attn/src/flash_fwd_split_hdim128_fp16_causal_sm80.cu",
+  # "flash_attn/src/flash_fwd_split_hdim128_fp16_sm80.cu",
+  # "flash_attn/src/flash_fwd_split_hdim160_bf16_causal_sm80.cu",
+  # "flash_attn/src/flash_fwd_split_hdim160_bf16_sm80.cu",
+  # "flash_attn/src/flash_fwd_split_hdim160_fp16_causal_sm80.cu",
+  # "flash_attn/src/flash_fwd_split_hdim160_fp16_sm80.cu",
+  # "flash_attn/src/flash_fwd_split_hdim192_bf16_causal_sm80.cu",
+  # "flash_attn/src/flash_fwd_split_hdim192_bf16_sm80.cu",
+  # "flash_attn/src/flash_fwd_split_hdim192_fp16_causal_sm80.cu",
+  # "flash_attn/src/flash_fwd_split_hdim192_fp16_sm80.cu",
+  # "flash_attn/src/flash_fwd_split_hdim256_bf16_causal_sm80.cu",
+  # "flash_attn/src/flash_fwd_split_hdim256_bf16_sm80.cu",
+  # "flash_attn/src/flash_fwd_split_hdim256_fp16_causal_sm80.cu",
+  # "flash_attn/src/flash_fwd_split_hdim256_fp16_sm80.cu",
+  # "flash_attn/src/flash_fwd_split_hdim32_bf16_causal_sm80.cu",
+  # "flash_attn/src/flash_fwd_split_hdim32_bf16_sm80.cu",
+  # "flash_attn/src/flash_fwd_split_hdim32_fp16_causal_sm80.cu",
+  # "flash_attn/src/flash_fwd_split_hdim32_fp16_sm80.cu",
+  # "flash_attn/src/flash_fwd_split_hdim64_bf16_causal_sm80.cu",
+  # "flash_attn/src/flash_fwd_split_hdim64_bf16_sm80.cu",
+  # "flash_attn/src/flash_fwd_split_hdim64_fp16_causal_sm80.cu",
+  # "flash_attn/src/flash_fwd_split_hdim64_fp16_sm80.cu",
+  # "flash_attn/src/flash_fwd_split_hdim96_bf16_causal_sm80.cu",
+  # "flash_attn/src/flash_fwd_split_hdim96_bf16_sm80.cu",
+  # "flash_attn/src/flash_fwd_split_hdim96_fp16_causal_sm80.cu",
+  # "flash_attn/src/flash_fwd_split_hdim96_fp16_sm80.cu",
 ]
 depends = ["torch", "cutlass_3_6"]

build/torch25-cxx11-cu118-x86_64-linux/flash_attn/__init__.py

@@ -0,0 +1,37 @@
+from typing import Optional
+
+import torch
+
+from ._ops import ops
+
+def mha_fwd(
+    q: torch.Tensor,
+    k: torch.Tensor,
+    v: torch.Tensor,
+    out: torch.Tensor,
+    alibi_slopes: torch.Tensor,
+    p_dropout: float,
+    softmax_scale: float,
+    is_causal: bool,
+    window_size_left: int,
+    window_size_right: int,
+    softcap: float,
+    return_softmax: bool,
+    gen: Optional[torch.Generator],
+) -> torch.Tensor:
+    return ops.mha_fwd(
+        q,
+        k,
+        v,
+        out,
+        alibi_slopes,
+        p_dropout,
+        softmax_scale,
+        is_causal,
+        window_size_left,
+        window_size_right,
+        softcap,
+        return_softmax,
+        gen,
+    )
+    return out

build/torch25-cxx11-cu118-x86_64-linux/flash_attn/_flash_attn_a7165c8_dirty.abi3.so

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:c2a6f11f1665f62c8f3b96cd843c806b737966575c28804c602bc68d089c1759
+size 17469320

build/torch25-cxx11-cu118-x86_64-linux/flash_attn/_ops.py

@@ -0,0 +1,9 @@
+import torch
+from . import _flash_attn_a7165c8_dirty
+ops = torch.ops._flash_attn_a7165c8_dirty
+
+def add_op_namespace_prefix(op_name: str):
+    """
+    Prefix op by namespace.
+    """
+    return f"_flash_attn_a7165c8_dirty::{op_name}"

build/torch25-cxx11-cu121-x86_64-linux/flash_attn/__init__.py

@@ -0,0 +1,37 @@
+from typing import Optional
+
+import torch
+
+from ._ops import ops
+
+def mha_fwd(
+    q: torch.Tensor,
+    k: torch.Tensor,
+    v: torch.Tensor,
+    out: torch.Tensor,
+    alibi_slopes: torch.Tensor,
+    p_dropout: float,
+    softmax_scale: float,
+    is_causal: bool,
+    window_size_left: int,
+    window_size_right: int,
+    softcap: float,
+    return_softmax: bool,
+    gen: Optional[torch.Generator],
+) -> torch.Tensor:
+    return ops.mha_fwd(
+        q,
+        k,
+        v,
+        out,
+        alibi_slopes,
+        p_dropout,
+        softmax_scale,
+        is_causal,
+        window_size_left,
+        window_size_right,
+        softcap,
+        return_softmax,
+        gen,
+    )
+    return out

build/torch25-cxx11-cu121-x86_64-linux/flash_attn/_flash_attn_a7165c8_dirty.abi3.so

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:2e481af6967a53e2017631ade57897e3ef32e1a13e8badb11310df46e8748dab
+size 17561616

build/torch25-cxx11-cu121-x86_64-linux/flash_attn/_ops.py

@@ -0,0 +1,9 @@
+import torch
+from . import _flash_attn_a7165c8_dirty
+ops = torch.ops._flash_attn_a7165c8_dirty
+
+def add_op_namespace_prefix(op_name: str):
+    """
+    Prefix op by namespace.
+    """
+    return f"_flash_attn_a7165c8_dirty::{op_name}"

build/torch25-cxx11-cu124-x86_64-linux/flash_attn/__init__.py

@@ -0,0 +1,37 @@
+from typing import Optional
+
+import torch
+
+from ._ops import ops
+
+def mha_fwd(
+    q: torch.Tensor,
+    k: torch.Tensor,
+    v: torch.Tensor,
+    out: torch.Tensor,
+    alibi_slopes: torch.Tensor,
+    p_dropout: float,
+    softmax_scale: float,
+    is_causal: bool,
+    window_size_left: int,
+    window_size_right: int,
+    softcap: float,
+    return_softmax: bool,
+    gen: Optional[torch.Generator],
+) -> torch.Tensor:
+    return ops.mha_fwd(
+        q,
+        k,
+        v,
+        out,
+        alibi_slopes,
+        p_dropout,
+        softmax_scale,
+        is_causal,
+        window_size_left,
+        window_size_right,
+        softcap,
+        return_softmax,
+        gen,
+    )
+    return out

build/torch25-cxx11-cu124-x86_64-linux/flash_attn/_flash_attn_a7165c8_dirty.abi3.so

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:df22c84c094e57e3e08c4adb615637c8e1a10fc914f9601a372eb1749ffcda12
+size 17820800

build/torch25-cxx11-cu124-x86_64-linux/flash_attn/_ops.py

@@ -0,0 +1,9 @@
+import torch
+from . import _flash_attn_a7165c8_dirty
+ops = torch.ops._flash_attn_a7165c8_dirty
+
+def add_op_namespace_prefix(op_name: str):
+    """
+    Prefix op by namespace.
+    """
+    return f"_flash_attn_a7165c8_dirty::{op_name}"

build/torch25-cxx98-cu118-x86_64-linux/flash_attn/__init__.py

@@ -0,0 +1,37 @@
+from typing import Optional
+
+import torch
+
+from ._ops import ops
+
+def mha_fwd(
+    q: torch.Tensor,
+    k: torch.Tensor,
+    v: torch.Tensor,
+    out: torch.Tensor,
+    alibi_slopes: torch.Tensor,
+    p_dropout: float,
+    softmax_scale: float,
+    is_causal: bool,
+    window_size_left: int,
+    window_size_right: int,
+    softcap: float,
+    return_softmax: bool,
+    gen: Optional[torch.Generator],
+) -> torch.Tensor:
+    return ops.mha_fwd(
+        q,
+        k,
+        v,
+        out,
+        alibi_slopes,
+        p_dropout,
+        softmax_scale,
+        is_causal,
+        window_size_left,
+        window_size_right,
+        softcap,
+        return_softmax,
+        gen,
+    )
+    return out

build/torch25-cxx98-cu118-x86_64-linux/flash_attn/_flash_attn_a7165c8_dirty.abi3.so

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:dd9920d56ee47082c06be48f07d20a869864954713bb8d05991dfcf01992cc6b
+size 17461960

build/torch25-cxx98-cu118-x86_64-linux/flash_attn/_ops.py

@@ -0,0 +1,9 @@
+import torch
+from . import _flash_attn_a7165c8_dirty
+ops = torch.ops._flash_attn_a7165c8_dirty
+
+def add_op_namespace_prefix(op_name: str):
+    """
+    Prefix op by namespace.
+    """
+    return f"_flash_attn_a7165c8_dirty::{op_name}"

build/torch25-cxx98-cu121-x86_64-linux/flash_attn/__init__.py

@@ -0,0 +1,37 @@
+from typing import Optional
+
+import torch
+
+from ._ops import ops
+
+def mha_fwd(
+    q: torch.Tensor,
+    k: torch.Tensor,
+    v: torch.Tensor,
+    out: torch.Tensor,
+    alibi_slopes: torch.Tensor,
+    p_dropout: float,
+    softmax_scale: float,
+    is_causal: bool,
+    window_size_left: int,
+    window_size_right: int,
+    softcap: float,
+    return_softmax: bool,
+    gen: Optional[torch.Generator],
+) -> torch.Tensor:
+    return ops.mha_fwd(
+        q,
+        k,
+        v,
+        out,
+        alibi_slopes,
+        p_dropout,
+        softmax_scale,
+        is_causal,
+        window_size_left,
+        window_size_right,
+        softcap,
+        return_softmax,
+        gen,
+    )
+    return out

build/torch25-cxx98-cu121-x86_64-linux/flash_attn/_flash_attn_a7165c8_dirty.abi3.so

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:806a71437827eb1724e80bbaf1cee7f1ef0242cd7c9a34b7e6ff696a8536f16a
+size 17558544

build/torch25-cxx98-cu121-x86_64-linux/flash_attn/_ops.py

@@ -0,0 +1,9 @@
+import torch
+from . import _flash_attn_a7165c8_dirty
+ops = torch.ops._flash_attn_a7165c8_dirty
+
+def add_op_namespace_prefix(op_name: str):
+    """
+    Prefix op by namespace.
+    """
+    return f"_flash_attn_a7165c8_dirty::{op_name}"

build/torch25-cxx98-cu124-x86_64-linux/flash_attn/__init__.py

@@ -0,0 +1,37 @@
+from typing import Optional
+
+import torch
+
+from ._ops import ops
+
+def mha_fwd(
+    q: torch.Tensor,
+    k: torch.Tensor,
+    v: torch.Tensor,
+    out: torch.Tensor,
+    alibi_slopes: torch.Tensor,
+    p_dropout: float,
+    softmax_scale: float,
+    is_causal: bool,
+    window_size_left: int,
+    window_size_right: int,
+    softcap: float,
+    return_softmax: bool,
+    gen: Optional[torch.Generator],
+) -> torch.Tensor:
+    return ops.mha_fwd(
+        q,
+        k,
+        v,
+        out,
+        alibi_slopes,
+        p_dropout,
+        softmax_scale,
+        is_causal,
+        window_size_left,
+        window_size_right,
+        softcap,
+        return_softmax,
+        gen,
+    )
+    return out

build/torch25-cxx98-cu124-x86_64-linux/flash_attn/_flash_attn_a7165c8_dirty.abi3.so

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:098cc28e134482be440715e9df0fe5b3e4023c1b5ca2c562da39571b630c4d73
+size 17817728

build/torch25-cxx98-cu124-x86_64-linux/flash_attn/_ops.py

@@ -0,0 +1,9 @@
+import torch
+from . import _flash_attn_a7165c8_dirty
+ops = torch.ops._flash_attn_a7165c8_dirty
+
+def add_op_namespace_prefix(op_name: str):
+    """
+    Prefix op by namespace.
+    """
+    return f"_flash_attn_a7165c8_dirty::{op_name}"

build/torch26-cxx11-cu118-x86_64-linux/flash_attn/__init__.py

@@ -0,0 +1,37 @@
+from typing import Optional
+
+import torch
+
+from ._ops import ops
+
+def mha_fwd(
+    q: torch.Tensor,
+    k: torch.Tensor,
+    v: torch.Tensor,
+    out: torch.Tensor,
+    alibi_slopes: torch.Tensor,
+    p_dropout: float,
+    softmax_scale: float,
+    is_causal: bool,
+    window_size_left: int,
+    window_size_right: int,
+    softcap: float,
+    return_softmax: bool,
+    gen: Optional[torch.Generator],
+) -> torch.Tensor:
+    return ops.mha_fwd(
+        q,
+        k,
+        v,
+        out,
+        alibi_slopes,
+        p_dropout,
+        softmax_scale,
+        is_causal,
+        window_size_left,
+        window_size_right,
+        softcap,
+        return_softmax,
+        gen,
+    )
+    return out

build/torch26-cxx11-cu118-x86_64-linux/flash_attn/_flash_attn_a7165c8_dirty.abi3.so

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:db8e9a06cafa5dffe988c22df459745deb3ee1b22b084e53ed6429e49867aae7
+size 17469464

build/torch26-cxx11-cu118-x86_64-linux/flash_attn/_ops.py

@@ -0,0 +1,9 @@
+import torch
+from . import _flash_attn_a7165c8_dirty
+ops = torch.ops._flash_attn_a7165c8_dirty
+
+def add_op_namespace_prefix(op_name: str):
+    """
+    Prefix op by namespace.
+    """
+    return f"_flash_attn_a7165c8_dirty::{op_name}"

build/torch26-cxx11-cu124-x86_64-linux/flash_attn/__init__.py

@@ -0,0 +1,37 @@
+from typing import Optional
+
+import torch
+
+from ._ops import ops
+
+def mha_fwd(
+    q: torch.Tensor,
+    k: torch.Tensor,
+    v: torch.Tensor,
+    out: torch.Tensor,
+    alibi_slopes: torch.Tensor,
+    p_dropout: float,
+    softmax_scale: float,
+    is_causal: bool,
+    window_size_left: int,
+    window_size_right: int,
+    softcap: float,
+    return_softmax: bool,
+    gen: Optional[torch.Generator],
+) -> torch.Tensor:
+    return ops.mha_fwd(
+        q,
+        k,
+        v,
+        out,
+        alibi_slopes,
+        p_dropout,
+        softmax_scale,
+        is_causal,
+        window_size_left,
+        window_size_right,
+        softcap,
+        return_softmax,
+        gen,
+    )
+    return out

build/torch26-cxx11-cu124-x86_64-linux/flash_attn/_flash_attn_a7165c8_dirty.abi3.so

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:3177cf407996b4f51ee139bfa4dcaf647fd659429cf9901ade2ac08117e20f9d
+size 17821096

build/torch26-cxx11-cu124-x86_64-linux/flash_attn/_ops.py

@@ -0,0 +1,9 @@
+import torch
+from . import _flash_attn_a7165c8_dirty
+ops = torch.ops._flash_attn_a7165c8_dirty
+
+def add_op_namespace_prefix(op_name: str):
+    """
+    Prefix op by namespace.
+    """
+    return f"_flash_attn_a7165c8_dirty::{op_name}"

build/torch26-cxx11-cu126-x86_64-linux/flash_attn/__init__.py

@@ -0,0 +1,37 @@
+from typing import Optional
+
+import torch
+
+from ._ops import ops
+
+def mha_fwd(
+    q: torch.Tensor,
+    k: torch.Tensor,
+    v: torch.Tensor,
+    out: torch.Tensor,
+    alibi_slopes: torch.Tensor,
+    p_dropout: float,
+    softmax_scale: float,
+    is_causal: bool,
+    window_size_left: int,
+    window_size_right: int,
+    softcap: float,
+    return_softmax: bool,
+    gen: Optional[torch.Generator],
+) -> torch.Tensor:
+    return ops.mha_fwd(
+        q,
+        k,
+        v,
+        out,
+        alibi_slopes,
+        p_dropout,
+        softmax_scale,
+        is_causal,
+        window_size_left,
+        window_size_right,
+        softcap,
+        return_softmax,
+        gen,
+    )
+    return out

build/torch26-cxx11-cu126-x86_64-linux/flash_attn/_flash_attn_a7165c8_dirty.abi3.so

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:43f21a0f290a6f42e004303c760e5aacc851ad55bd9093cea4752c0a7d6b202e
+size 17981304

build/torch26-cxx11-cu126-x86_64-linux/flash_attn/_ops.py

@@ -0,0 +1,9 @@
+import torch
+from . import _flash_attn_a7165c8_dirty
+ops = torch.ops._flash_attn_a7165c8_dirty
+
+def add_op_namespace_prefix(op_name: str):
+    """
+    Prefix op by namespace.
+    """
+    return f"_flash_attn_a7165c8_dirty::{op_name}"

build/torch26-cxx98-cu118-x86_64-linux/flash_attn/__init__.py

@@ -0,0 +1,37 @@
+from typing import Optional
+
+import torch
+
+from ._ops import ops
+
+def mha_fwd(
+    q: torch.Tensor,
+    k: torch.Tensor,
+    v: torch.Tensor,
+    out: torch.Tensor,
+    alibi_slopes: torch.Tensor,
+    p_dropout: float,
+    softmax_scale: float,
+    is_causal: bool,
+    window_size_left: int,
+    window_size_right: int,
+    softcap: float,
+    return_softmax: bool,
+    gen: Optional[torch.Generator],
+) -> torch.Tensor:
+    return ops.mha_fwd(
+        q,
+        k,
+        v,
+        out,
+        alibi_slopes,
+        p_dropout,
+        softmax_scale,
+        is_causal,
+        window_size_left,
+        window_size_right,
+        softcap,
+        return_softmax,
+        gen,
+    )
+    return out

build/torch26-cxx98-cu118-x86_64-linux/flash_attn/_flash_attn_a7165c8_dirty.abi3.so

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:434702696304310402d3ce50496e7f9f113b632ebc90ef602e255562a54d480a
+size 17462256

build/torch26-cxx98-cu118-x86_64-linux/flash_attn/_ops.py

@@ -0,0 +1,9 @@
+import torch
+from . import _flash_attn_a7165c8_dirty
+ops = torch.ops._flash_attn_a7165c8_dirty
+
+def add_op_namespace_prefix(op_name: str):
+    """
+    Prefix op by namespace.
+    """
+    return f"_flash_attn_a7165c8_dirty::{op_name}"

build/torch26-cxx98-cu124-x86_64-linux/flash_attn/__init__.py

@@ -0,0 +1,37 @@
+from typing import Optional
+
+import torch
+
+from ._ops import ops
+
+def mha_fwd(
+    q: torch.Tensor,
+    k: torch.Tensor,
+    v: torch.Tensor,
+    out: torch.Tensor,
+    alibi_slopes: torch.Tensor,
+    p_dropout: float,
+    softmax_scale: float,
+    is_causal: bool,
+    window_size_left: int,
+    window_size_right: int,
+    softcap: float,
+    return_softmax: bool,
+    gen: Optional[torch.Generator],
+) -> torch.Tensor:
+    return ops.mha_fwd(
+        q,
+        k,
+        v,
+        out,
+        alibi_slopes,
+        p_dropout,
+        softmax_scale,
+        is_causal,
+        window_size_left,
+        window_size_right,
+        softcap,
+        return_softmax,
+        gen,
+    )
+    return out

build/torch26-cxx98-cu124-x86_64-linux/flash_attn/_flash_attn_a7165c8_dirty.abi3.so

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:d90d1be6a4a87ec538a3b009356af76b6c1a1b5b18ce1e69b0fe8b0316972090
+size 17817920

build/torch26-cxx98-cu124-x86_64-linux/flash_attn/_ops.py

@@ -0,0 +1,9 @@
+import torch
+from . import _flash_attn_a7165c8_dirty
+ops = torch.ops._flash_attn_a7165c8_dirty
+
+def add_op_namespace_prefix(op_name: str):
+    """
+    Prefix op by namespace.
+    """
+    return f"_flash_attn_a7165c8_dirty::{op_name}"

build/torch26-cxx98-cu126-x86_64-linux/flash_attn/__init__.py

@@ -0,0 +1,37 @@
+from typing import Optional
+
+import torch
+
+from ._ops import ops
+
+def mha_fwd(
+    q: torch.Tensor,
+    k: torch.Tensor,
+    v: torch.Tensor,
+    out: torch.Tensor,
+    alibi_slopes: torch.Tensor,
+    p_dropout: float,
+    softmax_scale: float,
+    is_causal: bool,
+    window_size_left: int,
+    window_size_right: int,
+    softcap: float,
+    return_softmax: bool,
+    gen: Optional[torch.Generator],
+) -> torch.Tensor:
+    return ops.mha_fwd(
+        q,
+        k,
+        v,
+        out,
+        alibi_slopes,
+        p_dropout,
+        softmax_scale,
+        is_causal,
+        window_size_left,
+        window_size_right,
+        softcap,
+        return_softmax,
+        gen,
+    )
+    return out

build/torch26-cxx98-cu126-x86_64-linux/flash_attn/_flash_attn_a7165c8_dirty.abi3.so

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:ff763dccb46211a07fab8e63cfd96f76984cd994525d6c8ce0e274489e8099ca
+size 17978128

build/torch26-cxx98-cu126-x86_64-linux/flash_attn/_ops.py

@@ -0,0 +1,9 @@
+import torch
+from . import _flash_attn_a7165c8_dirty
+ops = torch.ops._flash_attn_a7165c8_dirty
+
+def add_op_namespace_prefix(op_name: str):
+    """
+    Prefix op by namespace.
+    """
+    return f"_flash_attn_a7165c8_dirty::{op_name}"

flake.lock

@@ -41,11 +41,11 @@
         "rocm-nix": "rocm-nix"
       },
       "locked": {
-        "lastModified": 1742582705,
-        "narHash": "sha256-1Vq5IauC/8fjBqcnMbDzckLN/XLIGwWr3/c2Wt3I2vs=",
+        "lastModified": 1742905006,
+        "narHash": "sha256-SCi1f5Lti4AM0kNPlAidcgN/5YM4HgJP4KwCsMrB0IE=",
         "ref": "refs/heads/main",
-        "rev": "e06e3e72947fad8bfd2c1eb5d8e7f5ec01d359d6",
-        "revCount": 103,
+        "rev": "517a2bf2d0a3f1faf058ab995b6ca280b0999e7c",
+        "revCount": 105,
         "type": "git",
         "url": "ssh://[email protected]/huggingface/kernel-builder"
       },

flash_attn/flash_api.cpp

@@ -1490,7 +1490,23 @@ mha_fwd(const at::Tensor &q,         // batch_size x seqlen_q x num_heads x roun
         const double softcap,
         const bool return_softmax,
         const c10::optional<at::Generator> gen_) {
-    // return FLASH_NAMESPACE::mha_fwd(q, k, v, out_, alibi_slopes_, p_dropout, softmax_scale, is_causal, window_size_left, window_size_right, softcap, return_softmax, gen_);
-    // return dummy value for now
-    return {};
-};
+
+    auto gen = gen_.value_or(at::cuda::detail::getDefaultCUDAGenerator());
+
+    // Prepare the optional arguments as non-const references.
+    std::optional<at::Tensor> out = out_.has_value() ? std::optional<at::Tensor>(const_cast<at::Tensor &>(out_.value())) : std::nullopt;
+    std::optional<at::Tensor> alibi_slopes = alibi_slopes_.has_value() ? std::optional<at::Tensor>(const_cast<at::Tensor &>(alibi_slopes_.value())) : std::nullopt;
+   
+    if (!out.has_value()){
+        out = torch::empty_like(q);
+    }
+
+    // Convert double to float and int64_t to int.
+    float p_dropout_float = static_cast<float>(p_dropout);
+    float softmax_scale_float = static_cast<float>(softmax_scale);
+    float softcap_float = static_cast<float>(softcap);
+    int window_size_left_int = static_cast<int>(window_size_left);
+    int window_size_right_int = static_cast<int>(window_size_right);
+    
+    return FLASH_NAMESPACE::mha_fwd(const_cast<at::Tensor &>(q), k, v, out, alibi_slopes, p_dropout_float, softmax_scale_float, is_causal, window_size_left_int, window_size_right_int, softcap_float, return_softmax, gen);
+}

torch-ext/flash_attn/__init__.py

@@ -19,7 +19,7 @@ def mha_fwd(
     return_softmax: bool,
     gen: Optional[torch.Generator],
 ) -> torch.Tensor:
-    return ops.mha_fwd(
+    ops.mha_fwd(
         q,
         k,
         v,
@@ -34,4 +34,4 @@ def mha_fwd(
         return_softmax,
         gen,
     )
-    return out
+    return out