diff --git a/build/torch25-cxx11-cu118-x86_64-linux/flash_mla/__init__.py b/build/torch25-cxx11-cu118-x86_64-linux/flash_mla/__init__.py deleted file mode 100644 index 5d2a129688bf920047cca1ca70a9c3521a9b4182..0000000000000000000000000000000000000000 --- a/build/torch25-cxx11-cu118-x86_64-linux/flash_mla/__init__.py +++ /dev/null @@ -1,33 +0,0 @@ -import torch - -from ._ops import ops - - -def get_mla_metadata(seqlens_k: torch.Tensor, s_q: int, h_kv: int): - return ops.get_mla_metadata(seqlens_k, s_q, h_kv) - - -def mha_fwd_kvcache_mla( - q: torch.Tensor, - kcache: torch.Tensor, - vcache_: torch.Tensor, - head_size_v: int, - seqlens_k: torch.Tensor, - block_table: torch.Tensor, - softmax_scale: float, - is_causal_: bool, - tile_scheduler_metadata: torch.Tensor, - num_splits: torch.Tensor, -) -> torch.Tensor: - return ops.mha_fwd_kvcache_mla( - q, - kcache, - vcache_, - head_size_v, - seqlens_k, - block_table, - softmax_scale, - is_causal_, - tile_scheduler_metadata, - num_splits - ) diff --git a/build/torch25-cxx11-cu118-x86_64-linux/flash_mla/_flash_mla_d4f4195.abi3.so b/build/torch25-cxx11-cu118-x86_64-linux/flash_mla/_flash_mla_d4f4195.abi3.so deleted file mode 100755 index 495e029409a66d955828ec98a963d997e7f55803..0000000000000000000000000000000000000000 --- a/build/torch25-cxx11-cu118-x86_64-linux/flash_mla/_flash_mla_d4f4195.abi3.so +++ /dev/null @@ -1,3 +0,0 @@ -version https://git-lfs.github.com/spec/v1 -oid sha256:a9a2b99b276b5aa714b27d1f54cc5da2d451e65a9ed385c583daf528f2c030a9 -size 2564144 diff --git a/build/torch25-cxx11-cu118-x86_64-linux/flash_mla/_ops.py b/build/torch25-cxx11-cu118-x86_64-linux/flash_mla/_ops.py deleted file mode 100644 index 6e5aa4c56e24b75711edaa5d90f25828a6eb2484..0000000000000000000000000000000000000000 --- a/build/torch25-cxx11-cu118-x86_64-linux/flash_mla/_ops.py +++ /dev/null @@ -1,9 +0,0 @@ -import torch -from . import _flash_mla_d4f4195 -ops = torch.ops._flash_mla_d4f4195 - -def add_op_namespace_prefix(op_name: str): - """ - Prefix op by namespace. - """ - return f"_flash_mla_d4f4195::{op_name}" \ No newline at end of file diff --git a/build/torch25-cxx11-cu121-x86_64-linux/flash_mla/__init__.py b/build/torch25-cxx11-cu121-x86_64-linux/flash_mla/__init__.py deleted file mode 100644 index 5d2a129688bf920047cca1ca70a9c3521a9b4182..0000000000000000000000000000000000000000 --- a/build/torch25-cxx11-cu121-x86_64-linux/flash_mla/__init__.py +++ /dev/null @@ -1,33 +0,0 @@ -import torch - -from ._ops import ops - - -def get_mla_metadata(seqlens_k: torch.Tensor, s_q: int, h_kv: int): - return ops.get_mla_metadata(seqlens_k, s_q, h_kv) - - -def mha_fwd_kvcache_mla( - q: torch.Tensor, - kcache: torch.Tensor, - vcache_: torch.Tensor, - head_size_v: int, - seqlens_k: torch.Tensor, - block_table: torch.Tensor, - softmax_scale: float, - is_causal_: bool, - tile_scheduler_metadata: torch.Tensor, - num_splits: torch.Tensor, -) -> torch.Tensor: - return ops.mha_fwd_kvcache_mla( - q, - kcache, - vcache_, - head_size_v, - seqlens_k, - block_table, - softmax_scale, - is_causal_, - tile_scheduler_metadata, - num_splits - ) diff --git a/build/torch25-cxx11-cu121-x86_64-linux/flash_mla/_flash_mla_d4f4195.abi3.so b/build/torch25-cxx11-cu121-x86_64-linux/flash_mla/_flash_mla_d4f4195.abi3.so deleted file mode 100755 index 4576745bcbec270ebb59949a1aac7146e5754c46..0000000000000000000000000000000000000000 --- a/build/torch25-cxx11-cu121-x86_64-linux/flash_mla/_flash_mla_d4f4195.abi3.so +++ /dev/null @@ -1,3 +0,0 @@ -version https://git-lfs.github.com/spec/v1 -oid sha256:77e241f633fa5b103f379ba6ac58d2cc068e0c3fc4d4f20ac1e1c679fc19614f -size 2595176 diff --git a/build/torch25-cxx11-cu121-x86_64-linux/flash_mla/_ops.py b/build/torch25-cxx11-cu121-x86_64-linux/flash_mla/_ops.py deleted file mode 100644 index 6e5aa4c56e24b75711edaa5d90f25828a6eb2484..0000000000000000000000000000000000000000 --- a/build/torch25-cxx11-cu121-x86_64-linux/flash_mla/_ops.py +++ /dev/null @@ -1,9 +0,0 @@ -import torch -from . import _flash_mla_d4f4195 -ops = torch.ops._flash_mla_d4f4195 - -def add_op_namespace_prefix(op_name: str): - """ - Prefix op by namespace. - """ - return f"_flash_mla_d4f4195::{op_name}" \ No newline at end of file diff --git a/build/torch25-cxx11-cu124-x86_64-linux/flash_mla/__init__.py b/build/torch25-cxx11-cu124-x86_64-linux/flash_mla/__init__.py deleted file mode 100644 index 5d2a129688bf920047cca1ca70a9c3521a9b4182..0000000000000000000000000000000000000000 --- a/build/torch25-cxx11-cu124-x86_64-linux/flash_mla/__init__.py +++ /dev/null @@ -1,33 +0,0 @@ -import torch - -from ._ops import ops - - -def get_mla_metadata(seqlens_k: torch.Tensor, s_q: int, h_kv: int): - return ops.get_mla_metadata(seqlens_k, s_q, h_kv) - - -def mha_fwd_kvcache_mla( - q: torch.Tensor, - kcache: torch.Tensor, - vcache_: torch.Tensor, - head_size_v: int, - seqlens_k: torch.Tensor, - block_table: torch.Tensor, - softmax_scale: float, - is_causal_: bool, - tile_scheduler_metadata: torch.Tensor, - num_splits: torch.Tensor, -) -> torch.Tensor: - return ops.mha_fwd_kvcache_mla( - q, - kcache, - vcache_, - head_size_v, - seqlens_k, - block_table, - softmax_scale, - is_causal_, - tile_scheduler_metadata, - num_splits - ) diff --git a/build/torch25-cxx11-cu124-x86_64-linux/flash_mla/_flash_mla_d4f4195.abi3.so b/build/torch25-cxx11-cu124-x86_64-linux/flash_mla/_flash_mla_d4f4195.abi3.so deleted file mode 100755 index 60646291e7927b00f5a921ae85e16102a115fb52..0000000000000000000000000000000000000000 --- a/build/torch25-cxx11-cu124-x86_64-linux/flash_mla/_flash_mla_d4f4195.abi3.so +++ /dev/null @@ -1,3 +0,0 @@ -version https://git-lfs.github.com/spec/v1 -oid sha256:31aa895a57efbd29aeff693b65b02842926bf1788d6f98022c32470a60265f9e -size 2580248 diff --git a/build/torch25-cxx11-cu124-x86_64-linux/flash_mla/_ops.py b/build/torch25-cxx11-cu124-x86_64-linux/flash_mla/_ops.py deleted file mode 100644 index 6e5aa4c56e24b75711edaa5d90f25828a6eb2484..0000000000000000000000000000000000000000 --- a/build/torch25-cxx11-cu124-x86_64-linux/flash_mla/_ops.py +++ /dev/null @@ -1,9 +0,0 @@ -import torch -from . import _flash_mla_d4f4195 -ops = torch.ops._flash_mla_d4f4195 - -def add_op_namespace_prefix(op_name: str): - """ - Prefix op by namespace. - """ - return f"_flash_mla_d4f4195::{op_name}" \ No newline at end of file diff --git a/build/torch25-cxx98-cu118-x86_64-linux/flash_mla/__init__.py b/build/torch25-cxx98-cu118-x86_64-linux/flash_mla/__init__.py deleted file mode 100644 index 5d2a129688bf920047cca1ca70a9c3521a9b4182..0000000000000000000000000000000000000000 --- a/build/torch25-cxx98-cu118-x86_64-linux/flash_mla/__init__.py +++ /dev/null @@ -1,33 +0,0 @@ -import torch - -from ._ops import ops - - -def get_mla_metadata(seqlens_k: torch.Tensor, s_q: int, h_kv: int): - return ops.get_mla_metadata(seqlens_k, s_q, h_kv) - - -def mha_fwd_kvcache_mla( - q: torch.Tensor, - kcache: torch.Tensor, - vcache_: torch.Tensor, - head_size_v: int, - seqlens_k: torch.Tensor, - block_table: torch.Tensor, - softmax_scale: float, - is_causal_: bool, - tile_scheduler_metadata: torch.Tensor, - num_splits: torch.Tensor, -) -> torch.Tensor: - return ops.mha_fwd_kvcache_mla( - q, - kcache, - vcache_, - head_size_v, - seqlens_k, - block_table, - softmax_scale, - is_causal_, - tile_scheduler_metadata, - num_splits - ) diff --git a/build/torch25-cxx98-cu118-x86_64-linux/flash_mla/_flash_mla_d4f4195.abi3.so b/build/torch25-cxx98-cu118-x86_64-linux/flash_mla/_flash_mla_d4f4195.abi3.so deleted file mode 100755 index c22b643f0c606f54dced21fa2116796a02f23198..0000000000000000000000000000000000000000 --- a/build/torch25-cxx98-cu118-x86_64-linux/flash_mla/_flash_mla_d4f4195.abi3.so +++ /dev/null @@ -1,3 +0,0 @@ -version https://git-lfs.github.com/spec/v1 -oid sha256:7454c10a3b29128e035bdb3fa18d5fc3706f7970542a0bcb55d9714f0999d42f -size 2556792 diff --git a/build/torch25-cxx98-cu118-x86_64-linux/flash_mla/_ops.py b/build/torch25-cxx98-cu118-x86_64-linux/flash_mla/_ops.py deleted file mode 100644 index 6e5aa4c56e24b75711edaa5d90f25828a6eb2484..0000000000000000000000000000000000000000 --- a/build/torch25-cxx98-cu118-x86_64-linux/flash_mla/_ops.py +++ /dev/null @@ -1,9 +0,0 @@ -import torch -from . import _flash_mla_d4f4195 -ops = torch.ops._flash_mla_d4f4195 - -def add_op_namespace_prefix(op_name: str): - """ - Prefix op by namespace. - """ - return f"_flash_mla_d4f4195::{op_name}" \ No newline at end of file diff --git a/build/torch25-cxx98-cu121-x86_64-linux/flash_mla/__init__.py b/build/torch25-cxx98-cu121-x86_64-linux/flash_mla/__init__.py deleted file mode 100644 index 5d2a129688bf920047cca1ca70a9c3521a9b4182..0000000000000000000000000000000000000000 --- a/build/torch25-cxx98-cu121-x86_64-linux/flash_mla/__init__.py +++ /dev/null @@ -1,33 +0,0 @@ -import torch - -from ._ops import ops - - -def get_mla_metadata(seqlens_k: torch.Tensor, s_q: int, h_kv: int): - return ops.get_mla_metadata(seqlens_k, s_q, h_kv) - - -def mha_fwd_kvcache_mla( - q: torch.Tensor, - kcache: torch.Tensor, - vcache_: torch.Tensor, - head_size_v: int, - seqlens_k: torch.Tensor, - block_table: torch.Tensor, - softmax_scale: float, - is_causal_: bool, - tile_scheduler_metadata: torch.Tensor, - num_splits: torch.Tensor, -) -> torch.Tensor: - return ops.mha_fwd_kvcache_mla( - q, - kcache, - vcache_, - head_size_v, - seqlens_k, - block_table, - softmax_scale, - is_causal_, - tile_scheduler_metadata, - num_splits - ) diff --git a/build/torch25-cxx98-cu121-x86_64-linux/flash_mla/_flash_mla_d4f4195.abi3.so b/build/torch25-cxx98-cu121-x86_64-linux/flash_mla/_flash_mla_d4f4195.abi3.so deleted file mode 100755 index 5fb00b76fa56c3200f744e14a11afa5a3090dd7b..0000000000000000000000000000000000000000 --- a/build/torch25-cxx98-cu121-x86_64-linux/flash_mla/_flash_mla_d4f4195.abi3.so +++ /dev/null @@ -1,3 +0,0 @@ -version https://git-lfs.github.com/spec/v1 -oid sha256:c8cb9402f3091420227cbccf1ec4938a444765e26f5d34c356c76bf7c85630d0 -size 2587896 diff --git a/build/torch25-cxx98-cu121-x86_64-linux/flash_mla/_ops.py b/build/torch25-cxx98-cu121-x86_64-linux/flash_mla/_ops.py deleted file mode 100644 index 6e5aa4c56e24b75711edaa5d90f25828a6eb2484..0000000000000000000000000000000000000000 --- a/build/torch25-cxx98-cu121-x86_64-linux/flash_mla/_ops.py +++ /dev/null @@ -1,9 +0,0 @@ -import torch -from . import _flash_mla_d4f4195 -ops = torch.ops._flash_mla_d4f4195 - -def add_op_namespace_prefix(op_name: str): - """ - Prefix op by namespace. - """ - return f"_flash_mla_d4f4195::{op_name}" \ No newline at end of file diff --git a/build/torch25-cxx98-cu124-x86_64-linux/flash_mla/__init__.py b/build/torch25-cxx98-cu124-x86_64-linux/flash_mla/__init__.py deleted file mode 100644 index 5d2a129688bf920047cca1ca70a9c3521a9b4182..0000000000000000000000000000000000000000 --- a/build/torch25-cxx98-cu124-x86_64-linux/flash_mla/__init__.py +++ /dev/null @@ -1,33 +0,0 @@ -import torch - -from ._ops import ops - - -def get_mla_metadata(seqlens_k: torch.Tensor, s_q: int, h_kv: int): - return ops.get_mla_metadata(seqlens_k, s_q, h_kv) - - -def mha_fwd_kvcache_mla( - q: torch.Tensor, - kcache: torch.Tensor, - vcache_: torch.Tensor, - head_size_v: int, - seqlens_k: torch.Tensor, - block_table: torch.Tensor, - softmax_scale: float, - is_causal_: bool, - tile_scheduler_metadata: torch.Tensor, - num_splits: torch.Tensor, -) -> torch.Tensor: - return ops.mha_fwd_kvcache_mla( - q, - kcache, - vcache_, - head_size_v, - seqlens_k, - block_table, - softmax_scale, - is_causal_, - tile_scheduler_metadata, - num_splits - ) diff --git a/build/torch25-cxx98-cu124-x86_64-linux/flash_mla/_flash_mla_d4f4195.abi3.so b/build/torch25-cxx98-cu124-x86_64-linux/flash_mla/_flash_mla_d4f4195.abi3.so deleted file mode 100755 index 85a192d3f985a99cc8c984fabe3614ac2f44d9b2..0000000000000000000000000000000000000000 --- a/build/torch25-cxx98-cu124-x86_64-linux/flash_mla/_flash_mla_d4f4195.abi3.so +++ /dev/null @@ -1,3 +0,0 @@ -version https://git-lfs.github.com/spec/v1 -oid sha256:bb4be09cbde1979c1aa17e3bc93c1538f129b438d305bee0fe96f3c08efeee04 -size 2572968 diff --git a/build/torch25-cxx98-cu124-x86_64-linux/flash_mla/_ops.py b/build/torch25-cxx98-cu124-x86_64-linux/flash_mla/_ops.py deleted file mode 100644 index 6e5aa4c56e24b75711edaa5d90f25828a6eb2484..0000000000000000000000000000000000000000 --- a/build/torch25-cxx98-cu124-x86_64-linux/flash_mla/_ops.py +++ /dev/null @@ -1,9 +0,0 @@ -import torch -from . import _flash_mla_d4f4195 -ops = torch.ops._flash_mla_d4f4195 - -def add_op_namespace_prefix(op_name: str): - """ - Prefix op by namespace. - """ - return f"_flash_mla_d4f4195::{op_name}" \ No newline at end of file diff --git a/build/torch26-cxx11-cu118-x86_64-linux/flash_mla/__init__.py b/build/torch26-cxx11-cu118-x86_64-linux/flash_mla/__init__.py deleted file mode 100644 index 5d2a129688bf920047cca1ca70a9c3521a9b4182..0000000000000000000000000000000000000000 --- a/build/torch26-cxx11-cu118-x86_64-linux/flash_mla/__init__.py +++ /dev/null @@ -1,33 +0,0 @@ -import torch - -from ._ops import ops - - -def get_mla_metadata(seqlens_k: torch.Tensor, s_q: int, h_kv: int): - return ops.get_mla_metadata(seqlens_k, s_q, h_kv) - - -def mha_fwd_kvcache_mla( - q: torch.Tensor, - kcache: torch.Tensor, - vcache_: torch.Tensor, - head_size_v: int, - seqlens_k: torch.Tensor, - block_table: torch.Tensor, - softmax_scale: float, - is_causal_: bool, - tile_scheduler_metadata: torch.Tensor, - num_splits: torch.Tensor, -) -> torch.Tensor: - return ops.mha_fwd_kvcache_mla( - q, - kcache, - vcache_, - head_size_v, - seqlens_k, - block_table, - softmax_scale, - is_causal_, - tile_scheduler_metadata, - num_splits - ) diff --git a/build/torch26-cxx11-cu118-x86_64-linux/flash_mla/_flash_mla_d4f4195.abi3.so b/build/torch26-cxx11-cu118-x86_64-linux/flash_mla/_flash_mla_d4f4195.abi3.so deleted file mode 100755 index 513d74291a3c84f2b48b590a9fd6cec6c72f8f9e..0000000000000000000000000000000000000000 --- a/build/torch26-cxx11-cu118-x86_64-linux/flash_mla/_flash_mla_d4f4195.abi3.so +++ /dev/null @@ -1,3 +0,0 @@ -version https://git-lfs.github.com/spec/v1 -oid sha256:35c67c788220d8988e47cd4ad976495450b71cd682bd8ab08af3db066d625126 -size 2564496 diff --git a/build/torch26-cxx11-cu118-x86_64-linux/flash_mla/_ops.py b/build/torch26-cxx11-cu118-x86_64-linux/flash_mla/_ops.py deleted file mode 100644 index 6e5aa4c56e24b75711edaa5d90f25828a6eb2484..0000000000000000000000000000000000000000 --- a/build/torch26-cxx11-cu118-x86_64-linux/flash_mla/_ops.py +++ /dev/null @@ -1,9 +0,0 @@ -import torch -from . import _flash_mla_d4f4195 -ops = torch.ops._flash_mla_d4f4195 - -def add_op_namespace_prefix(op_name: str): - """ - Prefix op by namespace. - """ - return f"_flash_mla_d4f4195::{op_name}" \ No newline at end of file diff --git a/build/torch26-cxx11-cu124-x86_64-linux/flash_mla/__init__.py b/build/torch26-cxx11-cu124-x86_64-linux/flash_mla/__init__.py deleted file mode 100644 index 5d2a129688bf920047cca1ca70a9c3521a9b4182..0000000000000000000000000000000000000000 --- a/build/torch26-cxx11-cu124-x86_64-linux/flash_mla/__init__.py +++ /dev/null @@ -1,33 +0,0 @@ -import torch - -from ._ops import ops - - -def get_mla_metadata(seqlens_k: torch.Tensor, s_q: int, h_kv: int): - return ops.get_mla_metadata(seqlens_k, s_q, h_kv) - - -def mha_fwd_kvcache_mla( - q: torch.Tensor, - kcache: torch.Tensor, - vcache_: torch.Tensor, - head_size_v: int, - seqlens_k: torch.Tensor, - block_table: torch.Tensor, - softmax_scale: float, - is_causal_: bool, - tile_scheduler_metadata: torch.Tensor, - num_splits: torch.Tensor, -) -> torch.Tensor: - return ops.mha_fwd_kvcache_mla( - q, - kcache, - vcache_, - head_size_v, - seqlens_k, - block_table, - softmax_scale, - is_causal_, - tile_scheduler_metadata, - num_splits - ) diff --git a/build/torch26-cxx11-cu124-x86_64-linux/flash_mla/_flash_mla_d4f4195.abi3.so b/build/torch26-cxx11-cu124-x86_64-linux/flash_mla/_flash_mla_d4f4195.abi3.so deleted file mode 100755 index c5820e98edcaaf6652865037166c27a77cb8cdca..0000000000000000000000000000000000000000 --- a/build/torch26-cxx11-cu124-x86_64-linux/flash_mla/_flash_mla_d4f4195.abi3.so +++ /dev/null @@ -1,3 +0,0 @@ -version https://git-lfs.github.com/spec/v1 -oid sha256:069fb3e3a051c91e73390245c7463218829b8decf0f60bd6fc9a0ba8127b5bd2 -size 2580592 diff --git a/build/torch26-cxx11-cu124-x86_64-linux/flash_mla/_ops.py b/build/torch26-cxx11-cu124-x86_64-linux/flash_mla/_ops.py deleted file mode 100644 index 6e5aa4c56e24b75711edaa5d90f25828a6eb2484..0000000000000000000000000000000000000000 --- a/build/torch26-cxx11-cu124-x86_64-linux/flash_mla/_ops.py +++ /dev/null @@ -1,9 +0,0 @@ -import torch -from . import _flash_mla_d4f4195 -ops = torch.ops._flash_mla_d4f4195 - -def add_op_namespace_prefix(op_name: str): - """ - Prefix op by namespace. - """ - return f"_flash_mla_d4f4195::{op_name}" \ No newline at end of file diff --git a/build/torch26-cxx11-cu126-aarch64-linux/flash_mla/__init__.py b/build/torch26-cxx11-cu126-aarch64-linux/flash_mla/__init__.py deleted file mode 100644 index 5d2a129688bf920047cca1ca70a9c3521a9b4182..0000000000000000000000000000000000000000 --- a/build/torch26-cxx11-cu126-aarch64-linux/flash_mla/__init__.py +++ /dev/null @@ -1,33 +0,0 @@ -import torch - -from ._ops import ops - - -def get_mla_metadata(seqlens_k: torch.Tensor, s_q: int, h_kv: int): - return ops.get_mla_metadata(seqlens_k, s_q, h_kv) - - -def mha_fwd_kvcache_mla( - q: torch.Tensor, - kcache: torch.Tensor, - vcache_: torch.Tensor, - head_size_v: int, - seqlens_k: torch.Tensor, - block_table: torch.Tensor, - softmax_scale: float, - is_causal_: bool, - tile_scheduler_metadata: torch.Tensor, - num_splits: torch.Tensor, -) -> torch.Tensor: - return ops.mha_fwd_kvcache_mla( - q, - kcache, - vcache_, - head_size_v, - seqlens_k, - block_table, - softmax_scale, - is_causal_, - tile_scheduler_metadata, - num_splits - ) diff --git a/build/torch26-cxx11-cu126-aarch64-linux/flash_mla/_flash_mla_341ab77.abi3.so b/build/torch26-cxx11-cu126-aarch64-linux/flash_mla/_flash_mla_341ab77.abi3.so deleted file mode 100755 index b3622a1486b5cc9b43381b072ead09c68986493b..0000000000000000000000000000000000000000 --- a/build/torch26-cxx11-cu126-aarch64-linux/flash_mla/_flash_mla_341ab77.abi3.so +++ /dev/null @@ -1,3 +0,0 @@ -version https://git-lfs.github.com/spec/v1 -oid sha256:e1e97fef62f5ebbe6b19b0d5fbe700fcdf6b9acd7a54cba6f0b1d23665188fa9 -size 2643848 diff --git a/build/torch26-cxx11-cu126-aarch64-linux/flash_mla/_ops.py b/build/torch26-cxx11-cu126-aarch64-linux/flash_mla/_ops.py deleted file mode 100644 index 6ca4becc90e11facbc2ad156a8ef8bb23aeebed0..0000000000000000000000000000000000000000 --- a/build/torch26-cxx11-cu126-aarch64-linux/flash_mla/_ops.py +++ /dev/null @@ -1,9 +0,0 @@ -import torch -from . import _flash_mla_341ab77 -ops = torch.ops._flash_mla_341ab77 - -def add_op_namespace_prefix(op_name: str): - """ - Prefix op by namespace. - """ - return f"_flash_mla_341ab77::{op_name}" \ No newline at end of file diff --git a/build/torch26-cxx11-cu126-x86_64-linux/flash_mla/__init__.py b/build/torch26-cxx11-cu126-x86_64-linux/flash_mla/__init__.py deleted file mode 100644 index 5d2a129688bf920047cca1ca70a9c3521a9b4182..0000000000000000000000000000000000000000 --- a/build/torch26-cxx11-cu126-x86_64-linux/flash_mla/__init__.py +++ /dev/null @@ -1,33 +0,0 @@ -import torch - -from ._ops import ops - - -def get_mla_metadata(seqlens_k: torch.Tensor, s_q: int, h_kv: int): - return ops.get_mla_metadata(seqlens_k, s_q, h_kv) - - -def mha_fwd_kvcache_mla( - q: torch.Tensor, - kcache: torch.Tensor, - vcache_: torch.Tensor, - head_size_v: int, - seqlens_k: torch.Tensor, - block_table: torch.Tensor, - softmax_scale: float, - is_causal_: bool, - tile_scheduler_metadata: torch.Tensor, - num_splits: torch.Tensor, -) -> torch.Tensor: - return ops.mha_fwd_kvcache_mla( - q, - kcache, - vcache_, - head_size_v, - seqlens_k, - block_table, - softmax_scale, - is_causal_, - tile_scheduler_metadata, - num_splits - ) diff --git a/build/torch26-cxx11-cu126-x86_64-linux/flash_mla/_flash_mla_d4f4195.abi3.so b/build/torch26-cxx11-cu126-x86_64-linux/flash_mla/_flash_mla_d4f4195.abi3.so deleted file mode 100755 index 48e2389e6245bf0a3a220f3dddf52de715b00564..0000000000000000000000000000000000000000 --- a/build/torch26-cxx11-cu126-x86_64-linux/flash_mla/_flash_mla_d4f4195.abi3.so +++ /dev/null @@ -1,3 +0,0 @@ -version https://git-lfs.github.com/spec/v1 -oid sha256:50fad86fa7bc15096c2a1feadf8091b20e188e32b8c0633423ec26e4e8e8e7ce -size 2560552 diff --git a/build/torch26-cxx11-cu126-x86_64-linux/flash_mla/_ops.py b/build/torch26-cxx11-cu126-x86_64-linux/flash_mla/_ops.py deleted file mode 100644 index 6e5aa4c56e24b75711edaa5d90f25828a6eb2484..0000000000000000000000000000000000000000 --- a/build/torch26-cxx11-cu126-x86_64-linux/flash_mla/_ops.py +++ /dev/null @@ -1,9 +0,0 @@ -import torch -from . import _flash_mla_d4f4195 -ops = torch.ops._flash_mla_d4f4195 - -def add_op_namespace_prefix(op_name: str): - """ - Prefix op by namespace. - """ - return f"_flash_mla_d4f4195::{op_name}" \ No newline at end of file diff --git a/build/torch26-cxx98-cu118-x86_64-linux/flash_mla/__init__.py b/build/torch26-cxx98-cu118-x86_64-linux/flash_mla/__init__.py deleted file mode 100644 index 5d2a129688bf920047cca1ca70a9c3521a9b4182..0000000000000000000000000000000000000000 --- a/build/torch26-cxx98-cu118-x86_64-linux/flash_mla/__init__.py +++ /dev/null @@ -1,33 +0,0 @@ -import torch - -from ._ops import ops - - -def get_mla_metadata(seqlens_k: torch.Tensor, s_q: int, h_kv: int): - return ops.get_mla_metadata(seqlens_k, s_q, h_kv) - - -def mha_fwd_kvcache_mla( - q: torch.Tensor, - kcache: torch.Tensor, - vcache_: torch.Tensor, - head_size_v: int, - seqlens_k: torch.Tensor, - block_table: torch.Tensor, - softmax_scale: float, - is_causal_: bool, - tile_scheduler_metadata: torch.Tensor, - num_splits: torch.Tensor, -) -> torch.Tensor: - return ops.mha_fwd_kvcache_mla( - q, - kcache, - vcache_, - head_size_v, - seqlens_k, - block_table, - softmax_scale, - is_causal_, - tile_scheduler_metadata, - num_splits - ) diff --git a/build/torch26-cxx98-cu118-x86_64-linux/flash_mla/_flash_mla_d4f4195.abi3.so b/build/torch26-cxx98-cu118-x86_64-linux/flash_mla/_flash_mla_d4f4195.abi3.so deleted file mode 100755 index 8f8b0677194aee3da40b16b1d4ccef0bcb1c6a75..0000000000000000000000000000000000000000 --- a/build/torch26-cxx98-cu118-x86_64-linux/flash_mla/_flash_mla_d4f4195.abi3.so +++ /dev/null @@ -1,3 +0,0 @@ -version https://git-lfs.github.com/spec/v1 -oid sha256:ae937ddfbc3e6097b2fdd9197f2ddb5b9f66c65146a4de30ccab59dab6e18dd4 -size 2557136 diff --git a/build/torch26-cxx98-cu118-x86_64-linux/flash_mla/_ops.py b/build/torch26-cxx98-cu118-x86_64-linux/flash_mla/_ops.py deleted file mode 100644 index 6e5aa4c56e24b75711edaa5d90f25828a6eb2484..0000000000000000000000000000000000000000 --- a/build/torch26-cxx98-cu118-x86_64-linux/flash_mla/_ops.py +++ /dev/null @@ -1,9 +0,0 @@ -import torch -from . import _flash_mla_d4f4195 -ops = torch.ops._flash_mla_d4f4195 - -def add_op_namespace_prefix(op_name: str): - """ - Prefix op by namespace. - """ - return f"_flash_mla_d4f4195::{op_name}" \ No newline at end of file diff --git a/build/torch26-cxx98-cu124-x86_64-linux/flash_mla/__init__.py b/build/torch26-cxx98-cu124-x86_64-linux/flash_mla/__init__.py deleted file mode 100644 index 5d2a129688bf920047cca1ca70a9c3521a9b4182..0000000000000000000000000000000000000000 --- a/build/torch26-cxx98-cu124-x86_64-linux/flash_mla/__init__.py +++ /dev/null @@ -1,33 +0,0 @@ -import torch - -from ._ops import ops - - -def get_mla_metadata(seqlens_k: torch.Tensor, s_q: int, h_kv: int): - return ops.get_mla_metadata(seqlens_k, s_q, h_kv) - - -def mha_fwd_kvcache_mla( - q: torch.Tensor, - kcache: torch.Tensor, - vcache_: torch.Tensor, - head_size_v: int, - seqlens_k: torch.Tensor, - block_table: torch.Tensor, - softmax_scale: float, - is_causal_: bool, - tile_scheduler_metadata: torch.Tensor, - num_splits: torch.Tensor, -) -> torch.Tensor: - return ops.mha_fwd_kvcache_mla( - q, - kcache, - vcache_, - head_size_v, - seqlens_k, - block_table, - softmax_scale, - is_causal_, - tile_scheduler_metadata, - num_splits - ) diff --git a/build/torch26-cxx98-cu124-x86_64-linux/flash_mla/_flash_mla_d4f4195.abi3.so b/build/torch26-cxx98-cu124-x86_64-linux/flash_mla/_flash_mla_d4f4195.abi3.so deleted file mode 100755 index e086d51f613f85e3cdeb42100e4350be2a181b28..0000000000000000000000000000000000000000 --- a/build/torch26-cxx98-cu124-x86_64-linux/flash_mla/_flash_mla_d4f4195.abi3.so +++ /dev/null @@ -1,3 +0,0 @@ -version https://git-lfs.github.com/spec/v1 -oid sha256:503910324475f8bd9dab47687339005f58e5b623bf0c9e4234fabf099c08da33 -size 2573312 diff --git a/build/torch26-cxx98-cu124-x86_64-linux/flash_mla/_ops.py b/build/torch26-cxx98-cu124-x86_64-linux/flash_mla/_ops.py deleted file mode 100644 index 6e5aa4c56e24b75711edaa5d90f25828a6eb2484..0000000000000000000000000000000000000000 --- a/build/torch26-cxx98-cu124-x86_64-linux/flash_mla/_ops.py +++ /dev/null @@ -1,9 +0,0 @@ -import torch -from . import _flash_mla_d4f4195 -ops = torch.ops._flash_mla_d4f4195 - -def add_op_namespace_prefix(op_name: str): - """ - Prefix op by namespace. - """ - return f"_flash_mla_d4f4195::{op_name}" \ No newline at end of file diff --git a/build/torch26-cxx98-cu126-aarch64-linux/flash_mla/__init__.py b/build/torch26-cxx98-cu126-aarch64-linux/flash_mla/__init__.py deleted file mode 100644 index 5d2a129688bf920047cca1ca70a9c3521a9b4182..0000000000000000000000000000000000000000 --- a/build/torch26-cxx98-cu126-aarch64-linux/flash_mla/__init__.py +++ /dev/null @@ -1,33 +0,0 @@ -import torch - -from ._ops import ops - - -def get_mla_metadata(seqlens_k: torch.Tensor, s_q: int, h_kv: int): - return ops.get_mla_metadata(seqlens_k, s_q, h_kv) - - -def mha_fwd_kvcache_mla( - q: torch.Tensor, - kcache: torch.Tensor, - vcache_: torch.Tensor, - head_size_v: int, - seqlens_k: torch.Tensor, - block_table: torch.Tensor, - softmax_scale: float, - is_causal_: bool, - tile_scheduler_metadata: torch.Tensor, - num_splits: torch.Tensor, -) -> torch.Tensor: - return ops.mha_fwd_kvcache_mla( - q, - kcache, - vcache_, - head_size_v, - seqlens_k, - block_table, - softmax_scale, - is_causal_, - tile_scheduler_metadata, - num_splits - ) diff --git a/build/torch26-cxx98-cu126-aarch64-linux/flash_mla/_flash_mla_341ab77.abi3.so b/build/torch26-cxx98-cu126-aarch64-linux/flash_mla/_flash_mla_341ab77.abi3.so deleted file mode 100755 index fe8ba1fea4456f4132b5245ca04ffae76c2d43f4..0000000000000000000000000000000000000000 --- a/build/torch26-cxx98-cu126-aarch64-linux/flash_mla/_flash_mla_341ab77.abi3.so +++ /dev/null @@ -1,3 +0,0 @@ -version https://git-lfs.github.com/spec/v1 -oid sha256:6f15b3b0bd0bee56760bd6500175ca5a1fd17f2742ef9496c28ea3720d038c66 -size 2640208 diff --git a/build/torch26-cxx98-cu126-aarch64-linux/flash_mla/_ops.py b/build/torch26-cxx98-cu126-aarch64-linux/flash_mla/_ops.py deleted file mode 100644 index 6ca4becc90e11facbc2ad156a8ef8bb23aeebed0..0000000000000000000000000000000000000000 --- a/build/torch26-cxx98-cu126-aarch64-linux/flash_mla/_ops.py +++ /dev/null @@ -1,9 +0,0 @@ -import torch -from . import _flash_mla_341ab77 -ops = torch.ops._flash_mla_341ab77 - -def add_op_namespace_prefix(op_name: str): - """ - Prefix op by namespace. - """ - return f"_flash_mla_341ab77::{op_name}" \ No newline at end of file diff --git a/build/torch26-cxx98-cu126-x86_64-linux/flash_mla/__init__.py b/build/torch26-cxx98-cu126-x86_64-linux/flash_mla/__init__.py deleted file mode 100644 index 5d2a129688bf920047cca1ca70a9c3521a9b4182..0000000000000000000000000000000000000000 --- a/build/torch26-cxx98-cu126-x86_64-linux/flash_mla/__init__.py +++ /dev/null @@ -1,33 +0,0 @@ -import torch - -from ._ops import ops - - -def get_mla_metadata(seqlens_k: torch.Tensor, s_q: int, h_kv: int): - return ops.get_mla_metadata(seqlens_k, s_q, h_kv) - - -def mha_fwd_kvcache_mla( - q: torch.Tensor, - kcache: torch.Tensor, - vcache_: torch.Tensor, - head_size_v: int, - seqlens_k: torch.Tensor, - block_table: torch.Tensor, - softmax_scale: float, - is_causal_: bool, - tile_scheduler_metadata: torch.Tensor, - num_splits: torch.Tensor, -) -> torch.Tensor: - return ops.mha_fwd_kvcache_mla( - q, - kcache, - vcache_, - head_size_v, - seqlens_k, - block_table, - softmax_scale, - is_causal_, - tile_scheduler_metadata, - num_splits - ) diff --git a/build/torch26-cxx98-cu126-x86_64-linux/flash_mla/_flash_mla_d4f4195.abi3.so b/build/torch26-cxx98-cu126-x86_64-linux/flash_mla/_flash_mla_d4f4195.abi3.so deleted file mode 100755 index 2339738fe7aa81421ecdc9e619b68c1b5a2db07c..0000000000000000000000000000000000000000 --- a/build/torch26-cxx98-cu126-x86_64-linux/flash_mla/_flash_mla_d4f4195.abi3.so +++ /dev/null @@ -1,3 +0,0 @@ -version https://git-lfs.github.com/spec/v1 -oid sha256:2c41fa4058ee2bb5d3d90458a7f92f0ef1c10e8bc854329cf7c208025bb244b2 -size 2553280 diff --git a/build/torch26-cxx98-cu126-x86_64-linux/flash_mla/_ops.py b/build/torch26-cxx98-cu126-x86_64-linux/flash_mla/_ops.py deleted file mode 100644 index 6e5aa4c56e24b75711edaa5d90f25828a6eb2484..0000000000000000000000000000000000000000 --- a/build/torch26-cxx98-cu126-x86_64-linux/flash_mla/_ops.py +++ /dev/null @@ -1,9 +0,0 @@ -import torch -from . import _flash_mla_d4f4195 -ops = torch.ops._flash_mla_d4f4195 - -def add_op_namespace_prefix(op_name: str): - """ - Prefix op by namespace. - """ - return f"_flash_mla_d4f4195::{op_name}" \ No newline at end of file diff --git a/build/torch27-cxx11-cu118-x86_64-linux/flash_mla/__init__.py b/build/torch27-cxx11-cu118-x86_64-linux/flash_mla/__init__.py deleted file mode 100644 index 5d2a129688bf920047cca1ca70a9c3521a9b4182..0000000000000000000000000000000000000000 --- a/build/torch27-cxx11-cu118-x86_64-linux/flash_mla/__init__.py +++ /dev/null @@ -1,33 +0,0 @@ -import torch - -from ._ops import ops - - -def get_mla_metadata(seqlens_k: torch.Tensor, s_q: int, h_kv: int): - return ops.get_mla_metadata(seqlens_k, s_q, h_kv) - - -def mha_fwd_kvcache_mla( - q: torch.Tensor, - kcache: torch.Tensor, - vcache_: torch.Tensor, - head_size_v: int, - seqlens_k: torch.Tensor, - block_table: torch.Tensor, - softmax_scale: float, - is_causal_: bool, - tile_scheduler_metadata: torch.Tensor, - num_splits: torch.Tensor, -) -> torch.Tensor: - return ops.mha_fwd_kvcache_mla( - q, - kcache, - vcache_, - head_size_v, - seqlens_k, - block_table, - softmax_scale, - is_causal_, - tile_scheduler_metadata, - num_splits - ) diff --git a/build/torch27-cxx11-cu118-x86_64-linux/flash_mla/_flash_mla_d4f4195.abi3.so b/build/torch27-cxx11-cu118-x86_64-linux/flash_mla/_flash_mla_d4f4195.abi3.so deleted file mode 100755 index 518b48e7aaa58c369e3c721c64c4f2e5c7a88035..0000000000000000000000000000000000000000 --- a/build/torch27-cxx11-cu118-x86_64-linux/flash_mla/_flash_mla_d4f4195.abi3.so +++ /dev/null @@ -1,3 +0,0 @@ -version https://git-lfs.github.com/spec/v1 -oid sha256:59c4034880f4482b06e447a2c4810aaf8009b7d4c86a4fd71356f169df986535 -size 2564632 diff --git a/build/torch27-cxx11-cu118-x86_64-linux/flash_mla/_ops.py b/build/torch27-cxx11-cu118-x86_64-linux/flash_mla/_ops.py deleted file mode 100644 index 6e5aa4c56e24b75711edaa5d90f25828a6eb2484..0000000000000000000000000000000000000000 --- a/build/torch27-cxx11-cu118-x86_64-linux/flash_mla/_ops.py +++ /dev/null @@ -1,9 +0,0 @@ -import torch -from . import _flash_mla_d4f4195 -ops = torch.ops._flash_mla_d4f4195 - -def add_op_namespace_prefix(op_name: str): - """ - Prefix op by namespace. - """ - return f"_flash_mla_d4f4195::{op_name}" \ No newline at end of file diff --git a/build/torch27-cxx11-cu126-aarch64-linux/flash_mla/__init__.py b/build/torch27-cxx11-cu126-aarch64-linux/flash_mla/__init__.py deleted file mode 100644 index 5d2a129688bf920047cca1ca70a9c3521a9b4182..0000000000000000000000000000000000000000 --- a/build/torch27-cxx11-cu126-aarch64-linux/flash_mla/__init__.py +++ /dev/null @@ -1,33 +0,0 @@ -import torch - -from ._ops import ops - - -def get_mla_metadata(seqlens_k: torch.Tensor, s_q: int, h_kv: int): - return ops.get_mla_metadata(seqlens_k, s_q, h_kv) - - -def mha_fwd_kvcache_mla( - q: torch.Tensor, - kcache: torch.Tensor, - vcache_: torch.Tensor, - head_size_v: int, - seqlens_k: torch.Tensor, - block_table: torch.Tensor, - softmax_scale: float, - is_causal_: bool, - tile_scheduler_metadata: torch.Tensor, - num_splits: torch.Tensor, -) -> torch.Tensor: - return ops.mha_fwd_kvcache_mla( - q, - kcache, - vcache_, - head_size_v, - seqlens_k, - block_table, - softmax_scale, - is_causal_, - tile_scheduler_metadata, - num_splits - ) diff --git a/build/torch27-cxx11-cu126-aarch64-linux/flash_mla/_flash_mla_341ab77.abi3.so b/build/torch27-cxx11-cu126-aarch64-linux/flash_mla/_flash_mla_341ab77.abi3.so deleted file mode 100755 index 5770b43c7c5f6d6b5beff9fff23624279651b0aa..0000000000000000000000000000000000000000 --- a/build/torch27-cxx11-cu126-aarch64-linux/flash_mla/_flash_mla_341ab77.abi3.so +++ /dev/null @@ -1,3 +0,0 @@ -version https://git-lfs.github.com/spec/v1 -oid sha256:fb925b062d31034672a45d925a3767d953e97a3c6c483467e6b81833d42b5a27 -size 2644048 diff --git a/build/torch27-cxx11-cu126-aarch64-linux/flash_mla/_ops.py b/build/torch27-cxx11-cu126-aarch64-linux/flash_mla/_ops.py deleted file mode 100644 index 6ca4becc90e11facbc2ad156a8ef8bb23aeebed0..0000000000000000000000000000000000000000 --- a/build/torch27-cxx11-cu126-aarch64-linux/flash_mla/_ops.py +++ /dev/null @@ -1,9 +0,0 @@ -import torch -from . import _flash_mla_341ab77 -ops = torch.ops._flash_mla_341ab77 - -def add_op_namespace_prefix(op_name: str): - """ - Prefix op by namespace. - """ - return f"_flash_mla_341ab77::{op_name}" \ No newline at end of file diff --git a/build/torch27-cxx11-cu126-x86_64-linux/flash_mla/__init__.py b/build/torch27-cxx11-cu126-x86_64-linux/flash_mla/__init__.py deleted file mode 100644 index 5d2a129688bf920047cca1ca70a9c3521a9b4182..0000000000000000000000000000000000000000 --- a/build/torch27-cxx11-cu126-x86_64-linux/flash_mla/__init__.py +++ /dev/null @@ -1,33 +0,0 @@ -import torch - -from ._ops import ops - - -def get_mla_metadata(seqlens_k: torch.Tensor, s_q: int, h_kv: int): - return ops.get_mla_metadata(seqlens_k, s_q, h_kv) - - -def mha_fwd_kvcache_mla( - q: torch.Tensor, - kcache: torch.Tensor, - vcache_: torch.Tensor, - head_size_v: int, - seqlens_k: torch.Tensor, - block_table: torch.Tensor, - softmax_scale: float, - is_causal_: bool, - tile_scheduler_metadata: torch.Tensor, - num_splits: torch.Tensor, -) -> torch.Tensor: - return ops.mha_fwd_kvcache_mla( - q, - kcache, - vcache_, - head_size_v, - seqlens_k, - block_table, - softmax_scale, - is_causal_, - tile_scheduler_metadata, - num_splits - ) diff --git a/build/torch27-cxx11-cu126-x86_64-linux/flash_mla/_flash_mla_d4f4195.abi3.so b/build/torch27-cxx11-cu126-x86_64-linux/flash_mla/_flash_mla_d4f4195.abi3.so deleted file mode 100755 index 0b5c36b5a1b53830df46e7d3081382621e484bdc..0000000000000000000000000000000000000000 --- a/build/torch27-cxx11-cu126-x86_64-linux/flash_mla/_flash_mla_d4f4195.abi3.so +++ /dev/null @@ -1,3 +0,0 @@ -version https://git-lfs.github.com/spec/v1 -oid sha256:5db69ef4975e2eee001e6a9b7466c1fe40bc2228ed64eb8c24caf3e0fb6ed0b2 -size 2560584 diff --git a/build/torch27-cxx11-cu126-x86_64-linux/flash_mla/_ops.py b/build/torch27-cxx11-cu126-x86_64-linux/flash_mla/_ops.py deleted file mode 100644 index 6e5aa4c56e24b75711edaa5d90f25828a6eb2484..0000000000000000000000000000000000000000 --- a/build/torch27-cxx11-cu126-x86_64-linux/flash_mla/_ops.py +++ /dev/null @@ -1,9 +0,0 @@ -import torch -from . import _flash_mla_d4f4195 -ops = torch.ops._flash_mla_d4f4195 - -def add_op_namespace_prefix(op_name: str): - """ - Prefix op by namespace. - """ - return f"_flash_mla_d4f4195::{op_name}" \ No newline at end of file diff --git a/build/torch27-cxx11-cu128-aarch64-linux/flash_mla/__init__.py b/build/torch27-cxx11-cu128-aarch64-linux/flash_mla/__init__.py deleted file mode 100644 index 5d2a129688bf920047cca1ca70a9c3521a9b4182..0000000000000000000000000000000000000000 --- a/build/torch27-cxx11-cu128-aarch64-linux/flash_mla/__init__.py +++ /dev/null @@ -1,33 +0,0 @@ -import torch - -from ._ops import ops - - -def get_mla_metadata(seqlens_k: torch.Tensor, s_q: int, h_kv: int): - return ops.get_mla_metadata(seqlens_k, s_q, h_kv) - - -def mha_fwd_kvcache_mla( - q: torch.Tensor, - kcache: torch.Tensor, - vcache_: torch.Tensor, - head_size_v: int, - seqlens_k: torch.Tensor, - block_table: torch.Tensor, - softmax_scale: float, - is_causal_: bool, - tile_scheduler_metadata: torch.Tensor, - num_splits: torch.Tensor, -) -> torch.Tensor: - return ops.mha_fwd_kvcache_mla( - q, - kcache, - vcache_, - head_size_v, - seqlens_k, - block_table, - softmax_scale, - is_causal_, - tile_scheduler_metadata, - num_splits - ) diff --git a/build/torch27-cxx11-cu128-aarch64-linux/flash_mla/_flash_mla_341ab77.abi3.so b/build/torch27-cxx11-cu128-aarch64-linux/flash_mla/_flash_mla_341ab77.abi3.so deleted file mode 100755 index 12e46dd52eaa72ac67e34c8f60333c9a9d111c80..0000000000000000000000000000000000000000 --- a/build/torch27-cxx11-cu128-aarch64-linux/flash_mla/_flash_mla_341ab77.abi3.so +++ /dev/null @@ -1,3 +0,0 @@ -version https://git-lfs.github.com/spec/v1 -oid sha256:7776c629263bc0b32b82b8a094ead0749d6c393b6ca25c9ffa812bd8fbdb3002 -size 2709472 diff --git a/build/torch27-cxx11-cu128-aarch64-linux/flash_mla/_ops.py b/build/torch27-cxx11-cu128-aarch64-linux/flash_mla/_ops.py deleted file mode 100644 index 6ca4becc90e11facbc2ad156a8ef8bb23aeebed0..0000000000000000000000000000000000000000 --- a/build/torch27-cxx11-cu128-aarch64-linux/flash_mla/_ops.py +++ /dev/null @@ -1,9 +0,0 @@ -import torch -from . import _flash_mla_341ab77 -ops = torch.ops._flash_mla_341ab77 - -def add_op_namespace_prefix(op_name: str): - """ - Prefix op by namespace. - """ - return f"_flash_mla_341ab77::{op_name}" \ No newline at end of file diff --git a/build/torch27-cxx11-cu128-x86_64-linux/flash_mla/__init__.py b/build/torch27-cxx11-cu128-x86_64-linux/flash_mla/__init__.py deleted file mode 100644 index 5d2a129688bf920047cca1ca70a9c3521a9b4182..0000000000000000000000000000000000000000 --- a/build/torch27-cxx11-cu128-x86_64-linux/flash_mla/__init__.py +++ /dev/null @@ -1,33 +0,0 @@ -import torch - -from ._ops import ops - - -def get_mla_metadata(seqlens_k: torch.Tensor, s_q: int, h_kv: int): - return ops.get_mla_metadata(seqlens_k, s_q, h_kv) - - -def mha_fwd_kvcache_mla( - q: torch.Tensor, - kcache: torch.Tensor, - vcache_: torch.Tensor, - head_size_v: int, - seqlens_k: torch.Tensor, - block_table: torch.Tensor, - softmax_scale: float, - is_causal_: bool, - tile_scheduler_metadata: torch.Tensor, - num_splits: torch.Tensor, -) -> torch.Tensor: - return ops.mha_fwd_kvcache_mla( - q, - kcache, - vcache_, - head_size_v, - seqlens_k, - block_table, - softmax_scale, - is_causal_, - tile_scheduler_metadata, - num_splits - ) diff --git a/build/torch27-cxx11-cu128-x86_64-linux/flash_mla/_flash_mla_d4f4195.abi3.so b/build/torch27-cxx11-cu128-x86_64-linux/flash_mla/_flash_mla_d4f4195.abi3.so deleted file mode 100755 index 7769a2d82af13e218c14d783100f0c9e36090cbc..0000000000000000000000000000000000000000 --- a/build/torch27-cxx11-cu128-x86_64-linux/flash_mla/_flash_mla_d4f4195.abi3.so +++ /dev/null @@ -1,3 +0,0 @@ -version https://git-lfs.github.com/spec/v1 -oid sha256:3fc7eb9341c975d0e313d837977ca3ed13556e6fe63926e0bf117f62499ea052 -size 2615448 diff --git a/build/torch27-cxx11-cu128-x86_64-linux/flash_mla/_ops.py b/build/torch27-cxx11-cu128-x86_64-linux/flash_mla/_ops.py deleted file mode 100644 index 6e5aa4c56e24b75711edaa5d90f25828a6eb2484..0000000000000000000000000000000000000000 --- a/build/torch27-cxx11-cu128-x86_64-linux/flash_mla/_ops.py +++ /dev/null @@ -1,9 +0,0 @@ -import torch -from . import _flash_mla_d4f4195 -ops = torch.ops._flash_mla_d4f4195 - -def add_op_namespace_prefix(op_name: str): - """ - Prefix op by namespace. - """ - return f"_flash_mla_d4f4195::{op_name}" \ No newline at end of file diff --git a/build/torch29-cxx11-cu128-aarch64-linux/__init__.py b/build/torch29-cxx11-cu128-aarch64-linux/__init__.py deleted file mode 100644 index db300fe9b95176a20b27b3641d89be657d0c4319..0000000000000000000000000000000000000000 --- a/build/torch29-cxx11-cu128-aarch64-linux/__init__.py +++ /dev/null @@ -1,163 +0,0 @@ -from typing import Optional, Tuple -import torch - -from .flash_mla_interface import FlashMLASchedMeta -from . import flash_mla_interface as _impl - - -def get_mla_metadata(*args, **kwargs) -> Tuple[FlashMLASchedMeta, None]: - return _impl.get_mla_metadata(*args, **kwargs) - - -def flash_mla_with_kvcache( - q: torch.Tensor, - k_cache: torch.Tensor, - block_table: Optional[torch.Tensor], - cache_seqlens: Optional[torch.Tensor], - head_dim_v: int, - tile_scheduler_metadata: FlashMLASchedMeta, - num_splits: None = None, - softmax_scale: Optional[float] = None, - causal: bool = False, - is_fp8_kvcache: bool = False, - indices: Optional[torch.Tensor] = None, - attn_sink: Optional[torch.Tensor] = None, - extra_k_cache: Optional[torch.Tensor] = None, - extra_indices_in_kvcache: Optional[torch.Tensor] = None, - topk_length: Optional[torch.Tensor] = None, - extra_topk_length: Optional[torch.Tensor] = None, -) -> Tuple[torch.Tensor, torch.Tensor]: - return _impl.flash_mla_with_kvcache( - q=q, - k_cache=k_cache, - block_table=block_table, - cache_seqlens=cache_seqlens, - head_dim_v=head_dim_v, - tile_scheduler_metadata=tile_scheduler_metadata, - num_splits=num_splits, - softmax_scale=softmax_scale, - causal=causal, - is_fp8_kvcache=is_fp8_kvcache, - indices=indices, - attn_sink=attn_sink, - extra_k_cache=extra_k_cache, - extra_indices_in_kvcache=extra_indices_in_kvcache, - topk_length=topk_length, - extra_topk_length=extra_topk_length, - ) - - -def flash_mla_sparse_fwd( - q: torch.Tensor, - kv: torch.Tensor, - indices: torch.Tensor, - sm_scale: float, - d_v: int = 512, - attn_sink: Optional[torch.Tensor] = None, - topk_length: Optional[torch.Tensor] = None, -) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor]: - return _impl.flash_mla_sparse_fwd( - q=q, - kv=kv, - indices=indices, - sm_scale=sm_scale, - d_v=d_v, - attn_sink=attn_sink, - topk_length=topk_length, - ) - - -def flash_attn_varlen_func( - q: torch.Tensor, - k: torch.Tensor, - v: torch.Tensor, - cu_seqlens_qo: torch.Tensor, - cu_seqlens_kv: torch.Tensor, - max_seqlen_qo: int, - max_seqlen_kv: int, - dropout_p: float = 0.0, - softmax_scale: Optional[float] = None, - causal: bool = False, - deterministic: bool = False, - is_varlen: bool = True, -) -> Tuple[torch.Tensor, torch.Tensor]: - return _impl.flash_attn_varlen_func( - q=q, - k=k, - v=v, - cu_seqlens_qo=cu_seqlens_qo, - cu_seqlens_kv=cu_seqlens_kv, - max_seqlen_qo=max_seqlen_qo, - max_seqlen_kv=max_seqlen_kv, - dropout_p=dropout_p, - softmax_scale=softmax_scale, - causal=causal, - deterministic=deterministic, - is_varlen=is_varlen, - ) - - -def flash_attn_varlen_qkvpacked_func( - qkv: torch.Tensor, - cu_seqlens: torch.Tensor, - max_seqlen: int, - head_dim_qk: int, - dropout_p: float = 0.0, - softmax_scale: Optional[float] = None, - causal: bool = False, - deterministic: bool = False, - is_varlen: bool = True, -) -> Tuple[torch.Tensor, torch.Tensor]: - return _impl.flash_attn_varlen_qkvpacked_func( - qkv=qkv, - cu_seqlens=cu_seqlens, - max_seqlen=max_seqlen, - head_dim_qk=head_dim_qk, - dropout_p=dropout_p, - softmax_scale=softmax_scale, - causal=causal, - deterministic=deterministic, - is_varlen=is_varlen, - ) - - -def flash_attn_varlen_kvpacked_func( - q: torch.Tensor, - kv: torch.Tensor, - cu_seqlens_qo: torch.Tensor, - cu_seqlens_kv: torch.Tensor, - max_seqlen_qo: int, - max_seqlen_kv: int, - head_dim_qk: int, - dropout_p: float = 0.0, - softmax_scale: Optional[float] = None, - causal: bool = False, - deterministic: bool = False, - is_varlen: bool = True, -) -> Tuple[torch.Tensor, torch.Tensor]: - return _impl.flash_attn_varlen_kvpacked_func( - q=q, - kv=kv, - cu_seqlens_qo=cu_seqlens_qo, - cu_seqlens_kv=cu_seqlens_kv, - max_seqlen_qo=max_seqlen_qo, - max_seqlen_kv=max_seqlen_kv, - head_dim_qk=head_dim_qk, - dropout_p=dropout_p, - softmax_scale=softmax_scale, - causal=causal, - deterministic=deterministic, - is_varlen=is_varlen, - ) - - -__all__ = [ - "__version__", - "FlashMLASchedMeta", - "get_mla_metadata", - "flash_mla_with_kvcache", - "flash_attn_varlen_func", - "flash_attn_varlen_qkvpacked_func", - "flash_attn_varlen_kvpacked_func", - "flash_mla_sparse_fwd", -] diff --git a/build/torch29-cxx11-cu128-aarch64-linux/_flash_mla_cuda_09f70ef.abi3.so b/build/torch29-cxx11-cu128-aarch64-linux/_flash_mla_cuda_09f70ef.abi3.so deleted file mode 100644 index 8d043d5d2fc760bfae5abe1b6cbd8b9887e0aecd..0000000000000000000000000000000000000000 --- a/build/torch29-cxx11-cu128-aarch64-linux/_flash_mla_cuda_09f70ef.abi3.so +++ /dev/null @@ -1,3 +0,0 @@ -version https://git-lfs.github.com/spec/v1 -oid sha256:f3f3b9f82de962911a1b8467a3130b34ea10e0b8e7db32432bb42ac24862e52e -size 3667784 diff --git a/build/torch29-cxx11-cu128-aarch64-linux/_ops.py b/build/torch29-cxx11-cu128-aarch64-linux/_ops.py deleted file mode 100644 index ea7ed02f7680582f28bdb0d1e552de1dc177f7c5..0000000000000000000000000000000000000000 --- a/build/torch29-cxx11-cu128-aarch64-linux/_ops.py +++ /dev/null @@ -1,9 +0,0 @@ -import torch -from . import _flash_mla_cuda_09f70ef -ops = torch.ops._flash_mla_cuda_09f70ef - -def add_op_namespace_prefix(op_name: str): - """ - Prefix op by namespace. - """ - return f"_flash_mla_cuda_09f70ef::{op_name}" diff --git a/build/torch29-cxx11-cu128-aarch64-linux/flash_mla/__init__.py b/build/torch29-cxx11-cu128-aarch64-linux/flash_mla/__init__.py deleted file mode 100644 index 03dbc1afe1cf156661a2b1b22003cd5f599a0309..0000000000000000000000000000000000000000 --- a/build/torch29-cxx11-cu128-aarch64-linux/flash_mla/__init__.py +++ /dev/null @@ -1,26 +0,0 @@ -import ctypes -import sys - -import importlib -from pathlib import Path -from types import ModuleType - -def _import_from_path(file_path: Path) -> ModuleType: - # We cannot use the module name as-is, after adding it to `sys.modules`, - # it would also be used for other imports. So, we make a module name that - # depends on the path for it to be unique using the hex-encoded hash of - # the path. - path_hash = "{:x}".format(ctypes.c_size_t(hash(file_path.absolute())).value) - module_name = path_hash - spec = importlib.util.spec_from_file_location(module_name, file_path) - if spec is None: - raise ImportError(f"Cannot load spec for {module_name} from {file_path}") - module = importlib.util.module_from_spec(spec) - if module is None: - raise ImportError(f"Cannot load module {module_name} from spec") - sys.modules[module_name] = module - spec.loader.exec_module(module) # type: ignore - return module - - -globals().update(vars(_import_from_path(Path(__file__).parent.parent / "__init__.py"))) diff --git a/build/torch29-cxx11-cu128-aarch64-linux/flash_mla_interface.py b/build/torch29-cxx11-cu128-aarch64-linux/flash_mla_interface.py deleted file mode 100644 index a84e448ffe741bb6d3dafaf7888ed8cc94984467..0000000000000000000000000000000000000000 --- a/build/torch29-cxx11-cu128-aarch64-linux/flash_mla_interface.py +++ /dev/null @@ -1,435 +0,0 @@ -from typing import Optional, Tuple -import dataclasses - -import torch - -from ._ops import ops as flash_mla_cuda - -@dataclasses.dataclass -class FlashMLASchedMeta: - """ - A class that stores the tile scheduler metadata of FlashMLA - """ - - @dataclasses.dataclass - class Config: - b: int - s_q: int - h_q: int - page_block_size: int - h_k: int - - causal: bool - is_fp8_kvcache: bool - topk: Optional[int] - - extra_page_block_size: Optional[int] - extra_topk: Optional[int] - - have_initialized: bool = False - - config: Optional[Config] = None - - tile_scheduler_metadata: Optional[torch.Tensor] = None # (num_sm_parts, TileSchedulerMetaDataSize), dtype torch.int32. - num_splits: Optional[torch.Tensor] = None # (1), dtype torch.int32. - - -def get_mla_metadata( - *args, - **kwargs -) -> Tuple[FlashMLASchedMeta, None]: - """ - Returns an empty instance of FlashMLASchedMeta. The actual scheduling metadata will be generated during the first invocation of flash_mla_with_kvcache. - - Arguments: - This function does not need any arguments, but we keep *args and **kwargs to be compatible with the old interface. - - Return: - A tuple. Due to historical reasons, we return a tuple of (FlashMLASchedMeta, None) now. Only the first element is useful. - """ - return FlashMLASchedMeta(), None - - -def flash_mla_with_kvcache( - q: torch.Tensor, - k_cache: torch.Tensor, - block_table: Optional[torch.Tensor], - cache_seqlens: Optional[torch.Tensor], - head_dim_v: int, - tile_scheduler_metadata: FlashMLASchedMeta, - num_splits: None = None, - softmax_scale: Optional[float] = None, - causal: bool = False, - is_fp8_kvcache: bool = False, - indices: Optional[torch.Tensor] = None, - attn_sink: Optional[torch.Tensor] = None, - extra_k_cache: Optional[torch.Tensor] = None, - extra_indices_in_kvcache: Optional[torch.Tensor] = None, - topk_length: Optional[torch.Tensor] = None, - extra_topk_length: Optional[torch.Tensor] = None -) -> Tuple[torch.Tensor, torch.Tensor]: - """ - Arguments: - q: (batch_size, seq_len_q, num_heads_q, head_dim). - k_cache: (num_blocks, page_block_size, num_heads_k, head_dim). - Different modes (including fp8/bf16, and sparsity) has different KV cache layouts. See comments below for details. - The KV cache must be contiguously valid for sparse attention on sm100. Here "contiguously valid" means that every byte, from the very beginning of the KV cache, till the last byte in the KV cache, is valid memory address to visit (i.e. won't IMA). In other words, the KV cache could be a slice of a larger array, but cannot be a list of disjoint memory blocks. - block_table: (batch_size, max_num_blocks_per_seq), torch.int32. Can be None when sparse attention is used. - cache_seqlens: (batch_size), torch.int32. Can be None when sparse attention is used. - head_dim_v: Head_dim of v. Must be 512 - sched_meta: FlashMLASchedMeta, return by get_mla_metadata. You may reuse the same sched_meta across different invocations, but only when the tensor shapes and the values of cache_seqlens, topk_length, and extra_topk_length remain the same. - num_splits_placeholder: must be "None" (to be compatible with the old interface). - softmax_scale: float. The scaling of QK^T before applying softmax. Default to 1 / sqrt(head_dim_k). - causal: bool. Whether to apply causal attention mask. Only valid for dense attention - is_fp8_kvcache: bool. - indices: (batch_size, seq_len_q, topk). KV indices when sparse attention is enabled. - Pay attention that indices_in_kvcache[i][j][k] = (the index of the page block where token t resides) * block_size + (the offset of token t among the page block), - where t is the k-th token of the j-th q-sequence in the i-th batch. - attn_sink: Optional[torch.Tensor], (num_heads_q, ), torch.float32. If presented, the final output will be scaled by exp(lse) / (exp(lse) + exp(attn_sink)). Have no affect on the returned softmax_lse. +inf will cause the result to become 0. - extra_k_cache and extra_indices_in_kvcache: If provided, will attend to these extra tokens in addition to those in k_cache and indices_in_kvcache. Their format requirements are the same as k_cache and indices_in_kvcache respectively. - topk_length/extra_topk_length: (batch_size, ), torch.int32. If provided, only the leftmost topk_length indices will be processed. Useful when the actual topk for different queries are different so that we can save some computation, compared to masking. - - For DeepSeek V3, DeepSeek V3.1, and DeepSeek V3.2: - head_dim should be 576 while head_dim_v should be 512. - In FP8+sparse mode, each token's KV cache is 656 Bytes, structured as: - - The shape of the tensor `k_cache` is (num_blocks, page_block_size, num_heads_k, head_dim), and num_heads_k must be 1. - - First 512 bytes: The "quantized NoPE" part, containing 512 float8_e4m3 values. - - Next 16 bytes: Scale factors, containing 4 float32 values. The first float32 is the scale for the first 128 float8_e4m3 values, the second for the next 128, and so on. - - Last 128 bytes: The "RoPE" part, containing 64 bfloat16 values. This part is not quantized for accuracy. - - Return: - out: (batch_size, seq_len_q, num_heads_q, head_dim_v). - softmax_lse: (batch_size, num_heads_q, seq_len_q), torch.float32. - """ - sched_meta = tile_scheduler_metadata - indices_in_kvcache = indices - assert isinstance(sched_meta, FlashMLASchedMeta), "tile_scheduler_metadata must be of type FlashMLASchedMeta" - assert num_splits is None, "num_splits must be None" - - topk = indices_in_kvcache.shape[-1] if indices_in_kvcache is not None else None - extra_k_page_block_size = extra_k_cache.shape[1] if extra_k_cache is not None else None - extra_topk = extra_indices_in_kvcache.shape[-1] if extra_indices_in_kvcache is not None else None - if softmax_scale is None: - softmax_scale = q.shape[-1] ** (-0.5) - - if not sched_meta.have_initialized: - # Sanity check. We only perform sanity check during the first invocation to save CPU time. - if indices_in_kvcache is not None: - assert not causal, "causal must be False when indices_in_kvcache is not None (i.e. sparse attention is enabled)" - - # Initialize the tile scheduler metadata during the first invocation. - sched_meta.have_initialized = True - sched_meta.config = FlashMLASchedMeta.Config( - q.shape[0], - q.shape[1], - q.shape[2], - k_cache.shape[1], - k_cache.shape[2], - - causal, - is_fp8_kvcache, - topk, - - extra_k_page_block_size, - extra_topk, - ) - else: - # Check whether the input arguments are consistent with sched_meta - helper_msg = " Your input arguments are inconsistent with sched_meta. Please make sure the input arguments are consistent across different invocations of flash_mla_with_kvcache on the same sched_meta." - assert sched_meta.config is not None - assert sched_meta.config.b == q.shape[0], "sched_meta.config.b must be equal to batch_size." + helper_msg - assert sched_meta.config.s_q == q.shape[1], "sched_meta.config.s_q must be equal to seq_len_q." + helper_msg - assert sched_meta.config.h_q == q.shape[2], "sched_meta.config.h_q must be equal to num_heads_q." + helper_msg - assert sched_meta.config.page_block_size == k_cache.shape[1], "sched_meta.config.page_block_size must be equal to page_block_size." + helper_msg - assert sched_meta.config.h_k == k_cache.shape[2], "sched_meta.config.h_k must be equal to num_heads_k." + helper_msg - assert sched_meta.config.causal == causal, "sched_meta.config.causal must be equal to causal." + helper_msg - assert sched_meta.config.is_fp8_kvcache == is_fp8_kvcache, "sched_meta.config.is_fp8_kvcache must be equal to is_fp8_kvcache." + helper_msg - assert sched_meta.config.topk == topk, "sched_meta.config.topk must be equal to the last dim of indices_in_kvcache." + helper_msg - assert sched_meta.config.extra_page_block_size == extra_k_page_block_size, "sched_meta.config.extra_page_block_size must be equal to the page_block_size of extra_k_cache." + helper_msg - assert sched_meta.config.extra_topk == extra_topk, "sched_meta.config.extra_topk must be equal to the last dim of extra_indices_in_kvcache." + helper_msg - - if topk is not None: - # Sparse attention - assert not causal, "causal must be False when sparse attention is enabled" - assert is_fp8_kvcache, "is_fp8_kvcache must be True when sparse attention is enabled" - out, lse, new_tile_scheduler_metadata, new_num_splits = flash_mla_cuda.sparse_decode_fwd( - q, k_cache, indices_in_kvcache, topk_length, attn_sink, - sched_meta.tile_scheduler_metadata, sched_meta.num_splits, - extra_k_cache, extra_indices_in_kvcache, extra_topk_length, - head_dim_v, softmax_scale - ) - else: - # Dense attention - assert indices_in_kvcache is None and attn_sink is None and extra_k_cache is None and extra_indices_in_kvcache is None and topk_length is None and extra_topk_length is None, "indices_in_kvcache, attn_sink, extra_k_cache, extra_indices_in_kvcache, topk_length and extra_topk_length must be None when dense attention is used." - assert block_table is not None and cache_seqlens is not None, "block_table and cache_seqlens must be provided when dense attention is used." - out, lse, new_tile_scheduler_metadata, new_num_splits = flash_mla_cuda.dense_decode_fwd( - q, k_cache, head_dim_v, - cache_seqlens, block_table, - softmax_scale, causal, - sched_meta.tile_scheduler_metadata, sched_meta.num_splits - ) - sched_meta.tile_scheduler_metadata = new_tile_scheduler_metadata - sched_meta.num_splits = new_num_splits - return (out, lse) - - -def flash_mla_sparse_fwd( - q: torch.Tensor, - kv: torch.Tensor, - indices: torch.Tensor, - sm_scale: float, - d_v: int = 512, - attn_sink: Optional[torch.Tensor] = None, - topk_length: Optional[torch.Tensor] = None, -) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor]: - """ - Sparse attention prefill kernel - - Args: - q: [s_q, h_q, d_qk], bfloat16 - kv: [s_kv, h_kv, d_qk], bfloat16 - indices: [s_q, h_kv, topk], int32. Invalid indices should be set to -1 or numbers >= s_kv - sm_scale: float - d_v: The dimension of value vectors. Can only be 512 - attn_sink: optional, [h_q], float32. - If attn_sink is provided, when computing output, output will be additionally multiplied by exp(lse) / (exp(lse) + exp(attn_sink)). - +-inf in attn_sink will be handled normally (i.e., -inf has no effect, +inf will make corresponding output all zeros). - This argument has no effect on lse and max_logits. - topk_length: optional, [s_q], int32. If provided, the i-th q token will only attend to k tokens specified by indices[i, :, :topk_length[i]], ignoring later k/v tokens (even if provided in indices). - In extremely rare cases (topk_length provided, there is a valid topk index between topk_length[i] ~ s_kv, and that topk index points to a k token containing NaN), operator output will contain NaN, so please avoid this situation. - - Returns: - (output, max_logits, lse) - Please refer to tests/ref.py for the precise definitions of these parameters. - - output: [s_q, h_q, d_v], bfloat16 - - max_logits: [s_q, h_q], float - - lse: [s_q, h_q], float, log-sum-exp of attention scores - """ - results = flash_mla_cuda.sparse_prefill_fwd( - q, kv, indices, sm_scale, d_v, attn_sink, topk_length - ) - return results - - -def _flash_attn_varlen_forward( - q: torch.Tensor, - k: torch.Tensor, - v: torch.Tensor, - cu_seqlens_qo: torch.Tensor, - cu_seqlens_kv: torch.Tensor, - max_seqlen_qo: int, - max_seqlen_kv: int, - out: Optional[torch.Tensor] = None, - lse: Optional[torch.Tensor] = None, - causal: bool = False, - softmax_scale: Optional[float] = None, - is_varlen: bool = True, -) -> Tuple[torch.Tensor, torch.Tensor]: - qo_total_len, num_qo_heads, head_dim_qk = q.shape - kv_total_len, num_kv_heads, head_dim_vo = v.shape - - mask_mode_code = 1 if causal else 0 - if softmax_scale is None: - softmax_scale = head_dim_qk ** (-0.5) - - if out is None: - out = torch.empty(qo_total_len, num_qo_heads, head_dim_vo, device=q.device, dtype=q.dtype) - if lse is None: - # Make lse contiguous on seqlen dim - lse = torch.empty(num_qo_heads, qo_total_len, device=q.device, dtype=torch.float32).T - - workspace_buffer = torch.empty(32 * 1024 * 1024, dtype=torch.uint8, device=q.device) - flash_mla_cuda.dense_prefill_fwd( - workspace_buffer, - q, - k, - v, - cu_seqlens_qo, - cu_seqlens_kv, - out, - lse, - mask_mode_code, - softmax_scale, - max_seqlen_qo, - max_seqlen_kv, - is_varlen, - ) - - return out, lse - - -def _flash_attn_varlen_backward( - do: torch.Tensor, - q: torch.Tensor, - k: torch.Tensor, - v: torch.Tensor, - out: torch.Tensor, - lse: torch.Tensor, - cu_seqlens_qo: torch.Tensor, - cu_seqlens_kv: torch.Tensor, - max_seqlen_qo: int, - max_seqlen_kv: int, - dq: Optional[torch.Tensor] = None, - dk: Optional[torch.Tensor] = None, - dv: Optional[torch.Tensor] = None, - causal: bool = False, - softmax_scale: Optional[float] = None, - is_varlen: bool = True, -) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor]: - qo_total_len, num_qo_heads, head_dim_qk = q.shape - kv_total_len, num_kv_heads, head_dim_vo = v.shape - - # TODO: fix bwd GQA - if num_qo_heads != num_kv_heads: - raise ValueError(f"SM100 bwd doesn't support GQA now. num_qo_heads: {num_qo_heads}, num_kv_heads: {num_kv_heads}.") - - mask_mode_code = 1 if causal else 0 - if softmax_scale is None: - softmax_scale = head_dim_qk ** (-0.5) - - if dq is None: - dq = torch.empty(qo_total_len, num_qo_heads, head_dim_qk, device=q.device, dtype=q.dtype) - if dk is None: - dk = torch.empty(kv_total_len, num_kv_heads, head_dim_qk, device=q.device, dtype=q.dtype) - if dv is None: - dv = torch.empty(kv_total_len, num_kv_heads, head_dim_vo, device=q.device, dtype=q.dtype) - - max_seqlen_qo_aligned = (max_seqlen_qo + 7) // 8 * 8 - bs = cu_seqlens_qo.shape[0] - 1 - workspace_bytes = 0 - workspace_bytes += 4 * bs * max_seqlen_qo_aligned * num_qo_heads * head_dim_qk # dQ_acc - workspace_bytes += 4 * max_seqlen_qo_aligned * bs * num_qo_heads * 2 # sum_OdO and scaled_lse - if num_qo_heads != num_kv_heads: - workspace_bytes += 2 * kv_total_len * num_qo_heads * (head_dim_qk + head_dim_vo) # dKV_acc - workspace_buffer = torch.empty(workspace_bytes, dtype=torch.uint8, device=q.device) - flash_mla_cuda.dense_prefill_bwd( - workspace_buffer, - do, - q, - k, - v, - out, - lse, - cu_seqlens_qo, - cu_seqlens_kv, - dq, - dk, - dv, - mask_mode_code, - softmax_scale, - max_seqlen_qo, - max_seqlen_kv, - is_varlen, - ) - - return dq, dk, dv - - -class FlashAttnVarlenFunc(torch.autograd.Function): - def forward( - ctx, - q: torch.Tensor, - k: torch.Tensor, - v: torch.Tensor, - cu_seqlens_qo: torch.Tensor, - cu_seqlens_kv: torch.Tensor, - max_seqlen_qo: int, - max_seqlen_kv: int, - causal: bool = False, - softmax_scale: Optional[float] = None, - is_varlen: bool = True, - ) -> Tuple[torch.Tensor, torch.Tensor]: - out, lse = _flash_attn_varlen_forward( - q, k, v, - cu_seqlens_qo, cu_seqlens_kv, max_seqlen_qo, max_seqlen_kv, - causal=causal, softmax_scale=softmax_scale, - is_varlen=is_varlen, - ) - ctx.save_for_backward(q, k, v, out, lse, cu_seqlens_qo, cu_seqlens_kv) - ctx.max_seqlen_qo = max_seqlen_qo - ctx.max_seqlen_kv = max_seqlen_kv - ctx.causal = causal - ctx.softmax_scale = softmax_scale - ctx.is_varlen = is_varlen - return out, lse - - def backward( - ctx, - do: torch.Tensor, - dlse: torch.Tensor, - ): - del dlse # LSE doesn't support backward currently - q, k, v, out, lse, cu_seqlens_qo, cu_seqlens_kv = ctx.saved_tensors - dq, dk, dv = _flash_attn_varlen_backward( - do, q, k, v, out, lse, - cu_seqlens_qo, cu_seqlens_kv, ctx.max_seqlen_qo, ctx.max_seqlen_kv, - causal=ctx.causal, softmax_scale=ctx.softmax_scale, - is_varlen=ctx.is_varlen, - ) - return dq, dk, dv, None, None, None, None, None, None, None - - -def flash_attn_varlen_func( - q: torch.Tensor, - k: torch.Tensor, - v: torch.Tensor, - cu_seqlens_qo: torch.Tensor, - cu_seqlens_kv: torch.Tensor, - max_seqlen_qo: int, - max_seqlen_kv: int, - dropout_p: float = 0.0, - softmax_scale: Optional[float] = None, - causal: bool = False, - deterministic: bool = False, - is_varlen: bool = True, -) -> Tuple[torch.Tensor, torch.Tensor]: - assert dropout_p == 0.0 - assert not deterministic - return FlashAttnVarlenFunc.apply( - q, k, v, - cu_seqlens_qo, cu_seqlens_kv, max_seqlen_qo, max_seqlen_kv, - causal, softmax_scale, is_varlen, - ) - - -def flash_attn_varlen_qkvpacked_func( - qkv: torch.Tensor, - cu_seqlens: torch.Tensor, - max_seqlen: int, - head_dim_qk: int, - dropout_p: float = 0.0, - softmax_scale: Optional[float] = None, - causal: bool = False, - deterministic: bool = False, - is_varlen: bool = True, -) -> Tuple[torch.Tensor, torch.Tensor]: - assert dropout_p == 0.0 - assert not deterministic - return FlashAttnVarlenFunc.apply( - qkv[:, :, :head_dim_qk], qkv[:, :, head_dim_qk:head_dim_qk * 2], qkv[:, :, head_dim_qk * 2:], - cu_seqlens, cu_seqlens, max_seqlen, max_seqlen, - causal, softmax_scale, is_varlen, - ) - - -def flash_attn_varlen_kvpacked_func( - q: torch.Tensor, - kv: torch.Tensor, - cu_seqlens_qo: torch.Tensor, - cu_seqlens_kv: torch.Tensor, - max_seqlen_qo: int, - max_seqlen_kv: int, - head_dim_qk: int, - dropout_p: float = 0.0, - softmax_scale: Optional[float] = None, - causal: bool = False, - deterministic: bool = False, - is_varlen: bool = True, -) -> Tuple[torch.Tensor, torch.Tensor]: - assert dropout_p == 0.0 - assert not deterministic - return FlashAttnVarlenFunc.apply( - q, kv[:, :, :head_dim_qk], kv[:, :, head_dim_qk:], - cu_seqlens_qo, cu_seqlens_kv, max_seqlen_qo, max_seqlen_kv, - causal, softmax_scale, is_varlen, - ) diff --git a/build/torch29-cxx11-cu128-aarch64-linux/metadata.json b/build/torch29-cxx11-cu128-aarch64-linux/metadata.json deleted file mode 100644 index 4899badb63d45293425e2164944268b6058af95d..0000000000000000000000000000000000000000 --- a/build/torch29-cxx11-cu128-aarch64-linux/metadata.json +++ /dev/null @@ -1,11 +0,0 @@ -{ - "version": 1, - "license": "MIT", - "python-depends": [], - "backend": { - "type": "cuda", - "archs": [ - "9.0a" - ] - } -} diff --git a/build/torch29-cxx11-cu128-x86_64-linux/__init__.py b/build/torch29-cxx11-cu128-x86_64-linux/__init__.py deleted file mode 100644 index db300fe9b95176a20b27b3641d89be657d0c4319..0000000000000000000000000000000000000000 --- a/build/torch29-cxx11-cu128-x86_64-linux/__init__.py +++ /dev/null @@ -1,163 +0,0 @@ -from typing import Optional, Tuple -import torch - -from .flash_mla_interface import FlashMLASchedMeta -from . import flash_mla_interface as _impl - - -def get_mla_metadata(*args, **kwargs) -> Tuple[FlashMLASchedMeta, None]: - return _impl.get_mla_metadata(*args, **kwargs) - - -def flash_mla_with_kvcache( - q: torch.Tensor, - k_cache: torch.Tensor, - block_table: Optional[torch.Tensor], - cache_seqlens: Optional[torch.Tensor], - head_dim_v: int, - tile_scheduler_metadata: FlashMLASchedMeta, - num_splits: None = None, - softmax_scale: Optional[float] = None, - causal: bool = False, - is_fp8_kvcache: bool = False, - indices: Optional[torch.Tensor] = None, - attn_sink: Optional[torch.Tensor] = None, - extra_k_cache: Optional[torch.Tensor] = None, - extra_indices_in_kvcache: Optional[torch.Tensor] = None, - topk_length: Optional[torch.Tensor] = None, - extra_topk_length: Optional[torch.Tensor] = None, -) -> Tuple[torch.Tensor, torch.Tensor]: - return _impl.flash_mla_with_kvcache( - q=q, - k_cache=k_cache, - block_table=block_table, - cache_seqlens=cache_seqlens, - head_dim_v=head_dim_v, - tile_scheduler_metadata=tile_scheduler_metadata, - num_splits=num_splits, - softmax_scale=softmax_scale, - causal=causal, - is_fp8_kvcache=is_fp8_kvcache, - indices=indices, - attn_sink=attn_sink, - extra_k_cache=extra_k_cache, - extra_indices_in_kvcache=extra_indices_in_kvcache, - topk_length=topk_length, - extra_topk_length=extra_topk_length, - ) - - -def flash_mla_sparse_fwd( - q: torch.Tensor, - kv: torch.Tensor, - indices: torch.Tensor, - sm_scale: float, - d_v: int = 512, - attn_sink: Optional[torch.Tensor] = None, - topk_length: Optional[torch.Tensor] = None, -) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor]: - return _impl.flash_mla_sparse_fwd( - q=q, - kv=kv, - indices=indices, - sm_scale=sm_scale, - d_v=d_v, - attn_sink=attn_sink, - topk_length=topk_length, - ) - - -def flash_attn_varlen_func( - q: torch.Tensor, - k: torch.Tensor, - v: torch.Tensor, - cu_seqlens_qo: torch.Tensor, - cu_seqlens_kv: torch.Tensor, - max_seqlen_qo: int, - max_seqlen_kv: int, - dropout_p: float = 0.0, - softmax_scale: Optional[float] = None, - causal: bool = False, - deterministic: bool = False, - is_varlen: bool = True, -) -> Tuple[torch.Tensor, torch.Tensor]: - return _impl.flash_attn_varlen_func( - q=q, - k=k, - v=v, - cu_seqlens_qo=cu_seqlens_qo, - cu_seqlens_kv=cu_seqlens_kv, - max_seqlen_qo=max_seqlen_qo, - max_seqlen_kv=max_seqlen_kv, - dropout_p=dropout_p, - softmax_scale=softmax_scale, - causal=causal, - deterministic=deterministic, - is_varlen=is_varlen, - ) - - -def flash_attn_varlen_qkvpacked_func( - qkv: torch.Tensor, - cu_seqlens: torch.Tensor, - max_seqlen: int, - head_dim_qk: int, - dropout_p: float = 0.0, - softmax_scale: Optional[float] = None, - causal: bool = False, - deterministic: bool = False, - is_varlen: bool = True, -) -> Tuple[torch.Tensor, torch.Tensor]: - return _impl.flash_attn_varlen_qkvpacked_func( - qkv=qkv, - cu_seqlens=cu_seqlens, - max_seqlen=max_seqlen, - head_dim_qk=head_dim_qk, - dropout_p=dropout_p, - softmax_scale=softmax_scale, - causal=causal, - deterministic=deterministic, - is_varlen=is_varlen, - ) - - -def flash_attn_varlen_kvpacked_func( - q: torch.Tensor, - kv: torch.Tensor, - cu_seqlens_qo: torch.Tensor, - cu_seqlens_kv: torch.Tensor, - max_seqlen_qo: int, - max_seqlen_kv: int, - head_dim_qk: int, - dropout_p: float = 0.0, - softmax_scale: Optional[float] = None, - causal: bool = False, - deterministic: bool = False, - is_varlen: bool = True, -) -> Tuple[torch.Tensor, torch.Tensor]: - return _impl.flash_attn_varlen_kvpacked_func( - q=q, - kv=kv, - cu_seqlens_qo=cu_seqlens_qo, - cu_seqlens_kv=cu_seqlens_kv, - max_seqlen_qo=max_seqlen_qo, - max_seqlen_kv=max_seqlen_kv, - head_dim_qk=head_dim_qk, - dropout_p=dropout_p, - softmax_scale=softmax_scale, - causal=causal, - deterministic=deterministic, - is_varlen=is_varlen, - ) - - -__all__ = [ - "__version__", - "FlashMLASchedMeta", - "get_mla_metadata", - "flash_mla_with_kvcache", - "flash_attn_varlen_func", - "flash_attn_varlen_qkvpacked_func", - "flash_attn_varlen_kvpacked_func", - "flash_mla_sparse_fwd", -] diff --git a/build/torch29-cxx11-cu128-x86_64-linux/_flash_mla_cuda_09f70ef.abi3.so b/build/torch29-cxx11-cu128-x86_64-linux/_flash_mla_cuda_09f70ef.abi3.so deleted file mode 100644 index 17cee4c7563a933314a32d9c5e074a58f2443baf..0000000000000000000000000000000000000000 --- a/build/torch29-cxx11-cu128-x86_64-linux/_flash_mla_cuda_09f70ef.abi3.so +++ /dev/null @@ -1,3 +0,0 @@ -version https://git-lfs.github.com/spec/v1 -oid sha256:22efc944dccf34d020a87fe713033bfae026d91eda6ecc86a4d491abe38edc51 -size 3523096 diff --git a/build/torch29-cxx11-cu128-x86_64-linux/_ops.py b/build/torch29-cxx11-cu128-x86_64-linux/_ops.py deleted file mode 100644 index ea7ed02f7680582f28bdb0d1e552de1dc177f7c5..0000000000000000000000000000000000000000 --- a/build/torch29-cxx11-cu128-x86_64-linux/_ops.py +++ /dev/null @@ -1,9 +0,0 @@ -import torch -from . import _flash_mla_cuda_09f70ef -ops = torch.ops._flash_mla_cuda_09f70ef - -def add_op_namespace_prefix(op_name: str): - """ - Prefix op by namespace. - """ - return f"_flash_mla_cuda_09f70ef::{op_name}" diff --git a/build/torch29-cxx11-cu128-x86_64-linux/flash_mla/__init__.py b/build/torch29-cxx11-cu128-x86_64-linux/flash_mla/__init__.py deleted file mode 100644 index 03dbc1afe1cf156661a2b1b22003cd5f599a0309..0000000000000000000000000000000000000000 --- a/build/torch29-cxx11-cu128-x86_64-linux/flash_mla/__init__.py +++ /dev/null @@ -1,26 +0,0 @@ -import ctypes -import sys - -import importlib -from pathlib import Path -from types import ModuleType - -def _import_from_path(file_path: Path) -> ModuleType: - # We cannot use the module name as-is, after adding it to `sys.modules`, - # it would also be used for other imports. So, we make a module name that - # depends on the path for it to be unique using the hex-encoded hash of - # the path. - path_hash = "{:x}".format(ctypes.c_size_t(hash(file_path.absolute())).value) - module_name = path_hash - spec = importlib.util.spec_from_file_location(module_name, file_path) - if spec is None: - raise ImportError(f"Cannot load spec for {module_name} from {file_path}") - module = importlib.util.module_from_spec(spec) - if module is None: - raise ImportError(f"Cannot load module {module_name} from spec") - sys.modules[module_name] = module - spec.loader.exec_module(module) # type: ignore - return module - - -globals().update(vars(_import_from_path(Path(__file__).parent.parent / "__init__.py"))) diff --git a/build/torch29-cxx11-cu128-x86_64-linux/flash_mla_interface.py b/build/torch29-cxx11-cu128-x86_64-linux/flash_mla_interface.py deleted file mode 100644 index a84e448ffe741bb6d3dafaf7888ed8cc94984467..0000000000000000000000000000000000000000 --- a/build/torch29-cxx11-cu128-x86_64-linux/flash_mla_interface.py +++ /dev/null @@ -1,435 +0,0 @@ -from typing import Optional, Tuple -import dataclasses - -import torch - -from ._ops import ops as flash_mla_cuda - -@dataclasses.dataclass -class FlashMLASchedMeta: - """ - A class that stores the tile scheduler metadata of FlashMLA - """ - - @dataclasses.dataclass - class Config: - b: int - s_q: int - h_q: int - page_block_size: int - h_k: int - - causal: bool - is_fp8_kvcache: bool - topk: Optional[int] - - extra_page_block_size: Optional[int] - extra_topk: Optional[int] - - have_initialized: bool = False - - config: Optional[Config] = None - - tile_scheduler_metadata: Optional[torch.Tensor] = None # (num_sm_parts, TileSchedulerMetaDataSize), dtype torch.int32. - num_splits: Optional[torch.Tensor] = None # (1), dtype torch.int32. - - -def get_mla_metadata( - *args, - **kwargs -) -> Tuple[FlashMLASchedMeta, None]: - """ - Returns an empty instance of FlashMLASchedMeta. The actual scheduling metadata will be generated during the first invocation of flash_mla_with_kvcache. - - Arguments: - This function does not need any arguments, but we keep *args and **kwargs to be compatible with the old interface. - - Return: - A tuple. Due to historical reasons, we return a tuple of (FlashMLASchedMeta, None) now. Only the first element is useful. - """ - return FlashMLASchedMeta(), None - - -def flash_mla_with_kvcache( - q: torch.Tensor, - k_cache: torch.Tensor, - block_table: Optional[torch.Tensor], - cache_seqlens: Optional[torch.Tensor], - head_dim_v: int, - tile_scheduler_metadata: FlashMLASchedMeta, - num_splits: None = None, - softmax_scale: Optional[float] = None, - causal: bool = False, - is_fp8_kvcache: bool = False, - indices: Optional[torch.Tensor] = None, - attn_sink: Optional[torch.Tensor] = None, - extra_k_cache: Optional[torch.Tensor] = None, - extra_indices_in_kvcache: Optional[torch.Tensor] = None, - topk_length: Optional[torch.Tensor] = None, - extra_topk_length: Optional[torch.Tensor] = None -) -> Tuple[torch.Tensor, torch.Tensor]: - """ - Arguments: - q: (batch_size, seq_len_q, num_heads_q, head_dim). - k_cache: (num_blocks, page_block_size, num_heads_k, head_dim). - Different modes (including fp8/bf16, and sparsity) has different KV cache layouts. See comments below for details. - The KV cache must be contiguously valid for sparse attention on sm100. Here "contiguously valid" means that every byte, from the very beginning of the KV cache, till the last byte in the KV cache, is valid memory address to visit (i.e. won't IMA). In other words, the KV cache could be a slice of a larger array, but cannot be a list of disjoint memory blocks. - block_table: (batch_size, max_num_blocks_per_seq), torch.int32. Can be None when sparse attention is used. - cache_seqlens: (batch_size), torch.int32. Can be None when sparse attention is used. - head_dim_v: Head_dim of v. Must be 512 - sched_meta: FlashMLASchedMeta, return by get_mla_metadata. You may reuse the same sched_meta across different invocations, but only when the tensor shapes and the values of cache_seqlens, topk_length, and extra_topk_length remain the same. - num_splits_placeholder: must be "None" (to be compatible with the old interface). - softmax_scale: float. The scaling of QK^T before applying softmax. Default to 1 / sqrt(head_dim_k). - causal: bool. Whether to apply causal attention mask. Only valid for dense attention - is_fp8_kvcache: bool. - indices: (batch_size, seq_len_q, topk). KV indices when sparse attention is enabled. - Pay attention that indices_in_kvcache[i][j][k] = (the index of the page block where token t resides) * block_size + (the offset of token t among the page block), - where t is the k-th token of the j-th q-sequence in the i-th batch. - attn_sink: Optional[torch.Tensor], (num_heads_q, ), torch.float32. If presented, the final output will be scaled by exp(lse) / (exp(lse) + exp(attn_sink)). Have no affect on the returned softmax_lse. +inf will cause the result to become 0. - extra_k_cache and extra_indices_in_kvcache: If provided, will attend to these extra tokens in addition to those in k_cache and indices_in_kvcache. Their format requirements are the same as k_cache and indices_in_kvcache respectively. - topk_length/extra_topk_length: (batch_size, ), torch.int32. If provided, only the leftmost topk_length indices will be processed. Useful when the actual topk for different queries are different so that we can save some computation, compared to masking. - - For DeepSeek V3, DeepSeek V3.1, and DeepSeek V3.2: - head_dim should be 576 while head_dim_v should be 512. - In FP8+sparse mode, each token's KV cache is 656 Bytes, structured as: - - The shape of the tensor `k_cache` is (num_blocks, page_block_size, num_heads_k, head_dim), and num_heads_k must be 1. - - First 512 bytes: The "quantized NoPE" part, containing 512 float8_e4m3 values. - - Next 16 bytes: Scale factors, containing 4 float32 values. The first float32 is the scale for the first 128 float8_e4m3 values, the second for the next 128, and so on. - - Last 128 bytes: The "RoPE" part, containing 64 bfloat16 values. This part is not quantized for accuracy. - - Return: - out: (batch_size, seq_len_q, num_heads_q, head_dim_v). - softmax_lse: (batch_size, num_heads_q, seq_len_q), torch.float32. - """ - sched_meta = tile_scheduler_metadata - indices_in_kvcache = indices - assert isinstance(sched_meta, FlashMLASchedMeta), "tile_scheduler_metadata must be of type FlashMLASchedMeta" - assert num_splits is None, "num_splits must be None" - - topk = indices_in_kvcache.shape[-1] if indices_in_kvcache is not None else None - extra_k_page_block_size = extra_k_cache.shape[1] if extra_k_cache is not None else None - extra_topk = extra_indices_in_kvcache.shape[-1] if extra_indices_in_kvcache is not None else None - if softmax_scale is None: - softmax_scale = q.shape[-1] ** (-0.5) - - if not sched_meta.have_initialized: - # Sanity check. We only perform sanity check during the first invocation to save CPU time. - if indices_in_kvcache is not None: - assert not causal, "causal must be False when indices_in_kvcache is not None (i.e. sparse attention is enabled)" - - # Initialize the tile scheduler metadata during the first invocation. - sched_meta.have_initialized = True - sched_meta.config = FlashMLASchedMeta.Config( - q.shape[0], - q.shape[1], - q.shape[2], - k_cache.shape[1], - k_cache.shape[2], - - causal, - is_fp8_kvcache, - topk, - - extra_k_page_block_size, - extra_topk, - ) - else: - # Check whether the input arguments are consistent with sched_meta - helper_msg = " Your input arguments are inconsistent with sched_meta. Please make sure the input arguments are consistent across different invocations of flash_mla_with_kvcache on the same sched_meta." - assert sched_meta.config is not None - assert sched_meta.config.b == q.shape[0], "sched_meta.config.b must be equal to batch_size." + helper_msg - assert sched_meta.config.s_q == q.shape[1], "sched_meta.config.s_q must be equal to seq_len_q." + helper_msg - assert sched_meta.config.h_q == q.shape[2], "sched_meta.config.h_q must be equal to num_heads_q." + helper_msg - assert sched_meta.config.page_block_size == k_cache.shape[1], "sched_meta.config.page_block_size must be equal to page_block_size." + helper_msg - assert sched_meta.config.h_k == k_cache.shape[2], "sched_meta.config.h_k must be equal to num_heads_k." + helper_msg - assert sched_meta.config.causal == causal, "sched_meta.config.causal must be equal to causal." + helper_msg - assert sched_meta.config.is_fp8_kvcache == is_fp8_kvcache, "sched_meta.config.is_fp8_kvcache must be equal to is_fp8_kvcache." + helper_msg - assert sched_meta.config.topk == topk, "sched_meta.config.topk must be equal to the last dim of indices_in_kvcache." + helper_msg - assert sched_meta.config.extra_page_block_size == extra_k_page_block_size, "sched_meta.config.extra_page_block_size must be equal to the page_block_size of extra_k_cache." + helper_msg - assert sched_meta.config.extra_topk == extra_topk, "sched_meta.config.extra_topk must be equal to the last dim of extra_indices_in_kvcache." + helper_msg - - if topk is not None: - # Sparse attention - assert not causal, "causal must be False when sparse attention is enabled" - assert is_fp8_kvcache, "is_fp8_kvcache must be True when sparse attention is enabled" - out, lse, new_tile_scheduler_metadata, new_num_splits = flash_mla_cuda.sparse_decode_fwd( - q, k_cache, indices_in_kvcache, topk_length, attn_sink, - sched_meta.tile_scheduler_metadata, sched_meta.num_splits, - extra_k_cache, extra_indices_in_kvcache, extra_topk_length, - head_dim_v, softmax_scale - ) - else: - # Dense attention - assert indices_in_kvcache is None and attn_sink is None and extra_k_cache is None and extra_indices_in_kvcache is None and topk_length is None and extra_topk_length is None, "indices_in_kvcache, attn_sink, extra_k_cache, extra_indices_in_kvcache, topk_length and extra_topk_length must be None when dense attention is used." - assert block_table is not None and cache_seqlens is not None, "block_table and cache_seqlens must be provided when dense attention is used." - out, lse, new_tile_scheduler_metadata, new_num_splits = flash_mla_cuda.dense_decode_fwd( - q, k_cache, head_dim_v, - cache_seqlens, block_table, - softmax_scale, causal, - sched_meta.tile_scheduler_metadata, sched_meta.num_splits - ) - sched_meta.tile_scheduler_metadata = new_tile_scheduler_metadata - sched_meta.num_splits = new_num_splits - return (out, lse) - - -def flash_mla_sparse_fwd( - q: torch.Tensor, - kv: torch.Tensor, - indices: torch.Tensor, - sm_scale: float, - d_v: int = 512, - attn_sink: Optional[torch.Tensor] = None, - topk_length: Optional[torch.Tensor] = None, -) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor]: - """ - Sparse attention prefill kernel - - Args: - q: [s_q, h_q, d_qk], bfloat16 - kv: [s_kv, h_kv, d_qk], bfloat16 - indices: [s_q, h_kv, topk], int32. Invalid indices should be set to -1 or numbers >= s_kv - sm_scale: float - d_v: The dimension of value vectors. Can only be 512 - attn_sink: optional, [h_q], float32. - If attn_sink is provided, when computing output, output will be additionally multiplied by exp(lse) / (exp(lse) + exp(attn_sink)). - +-inf in attn_sink will be handled normally (i.e., -inf has no effect, +inf will make corresponding output all zeros). - This argument has no effect on lse and max_logits. - topk_length: optional, [s_q], int32. If provided, the i-th q token will only attend to k tokens specified by indices[i, :, :topk_length[i]], ignoring later k/v tokens (even if provided in indices). - In extremely rare cases (topk_length provided, there is a valid topk index between topk_length[i] ~ s_kv, and that topk index points to a k token containing NaN), operator output will contain NaN, so please avoid this situation. - - Returns: - (output, max_logits, lse) - Please refer to tests/ref.py for the precise definitions of these parameters. - - output: [s_q, h_q, d_v], bfloat16 - - max_logits: [s_q, h_q], float - - lse: [s_q, h_q], float, log-sum-exp of attention scores - """ - results = flash_mla_cuda.sparse_prefill_fwd( - q, kv, indices, sm_scale, d_v, attn_sink, topk_length - ) - return results - - -def _flash_attn_varlen_forward( - q: torch.Tensor, - k: torch.Tensor, - v: torch.Tensor, - cu_seqlens_qo: torch.Tensor, - cu_seqlens_kv: torch.Tensor, - max_seqlen_qo: int, - max_seqlen_kv: int, - out: Optional[torch.Tensor] = None, - lse: Optional[torch.Tensor] = None, - causal: bool = False, - softmax_scale: Optional[float] = None, - is_varlen: bool = True, -) -> Tuple[torch.Tensor, torch.Tensor]: - qo_total_len, num_qo_heads, head_dim_qk = q.shape - kv_total_len, num_kv_heads, head_dim_vo = v.shape - - mask_mode_code = 1 if causal else 0 - if softmax_scale is None: - softmax_scale = head_dim_qk ** (-0.5) - - if out is None: - out = torch.empty(qo_total_len, num_qo_heads, head_dim_vo, device=q.device, dtype=q.dtype) - if lse is None: - # Make lse contiguous on seqlen dim - lse = torch.empty(num_qo_heads, qo_total_len, device=q.device, dtype=torch.float32).T - - workspace_buffer = torch.empty(32 * 1024 * 1024, dtype=torch.uint8, device=q.device) - flash_mla_cuda.dense_prefill_fwd( - workspace_buffer, - q, - k, - v, - cu_seqlens_qo, - cu_seqlens_kv, - out, - lse, - mask_mode_code, - softmax_scale, - max_seqlen_qo, - max_seqlen_kv, - is_varlen, - ) - - return out, lse - - -def _flash_attn_varlen_backward( - do: torch.Tensor, - q: torch.Tensor, - k: torch.Tensor, - v: torch.Tensor, - out: torch.Tensor, - lse: torch.Tensor, - cu_seqlens_qo: torch.Tensor, - cu_seqlens_kv: torch.Tensor, - max_seqlen_qo: int, - max_seqlen_kv: int, - dq: Optional[torch.Tensor] = None, - dk: Optional[torch.Tensor] = None, - dv: Optional[torch.Tensor] = None, - causal: bool = False, - softmax_scale: Optional[float] = None, - is_varlen: bool = True, -) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor]: - qo_total_len, num_qo_heads, head_dim_qk = q.shape - kv_total_len, num_kv_heads, head_dim_vo = v.shape - - # TODO: fix bwd GQA - if num_qo_heads != num_kv_heads: - raise ValueError(f"SM100 bwd doesn't support GQA now. num_qo_heads: {num_qo_heads}, num_kv_heads: {num_kv_heads}.") - - mask_mode_code = 1 if causal else 0 - if softmax_scale is None: - softmax_scale = head_dim_qk ** (-0.5) - - if dq is None: - dq = torch.empty(qo_total_len, num_qo_heads, head_dim_qk, device=q.device, dtype=q.dtype) - if dk is None: - dk = torch.empty(kv_total_len, num_kv_heads, head_dim_qk, device=q.device, dtype=q.dtype) - if dv is None: - dv = torch.empty(kv_total_len, num_kv_heads, head_dim_vo, device=q.device, dtype=q.dtype) - - max_seqlen_qo_aligned = (max_seqlen_qo + 7) // 8 * 8 - bs = cu_seqlens_qo.shape[0] - 1 - workspace_bytes = 0 - workspace_bytes += 4 * bs * max_seqlen_qo_aligned * num_qo_heads * head_dim_qk # dQ_acc - workspace_bytes += 4 * max_seqlen_qo_aligned * bs * num_qo_heads * 2 # sum_OdO and scaled_lse - if num_qo_heads != num_kv_heads: - workspace_bytes += 2 * kv_total_len * num_qo_heads * (head_dim_qk + head_dim_vo) # dKV_acc - workspace_buffer = torch.empty(workspace_bytes, dtype=torch.uint8, device=q.device) - flash_mla_cuda.dense_prefill_bwd( - workspace_buffer, - do, - q, - k, - v, - out, - lse, - cu_seqlens_qo, - cu_seqlens_kv, - dq, - dk, - dv, - mask_mode_code, - softmax_scale, - max_seqlen_qo, - max_seqlen_kv, - is_varlen, - ) - - return dq, dk, dv - - -class FlashAttnVarlenFunc(torch.autograd.Function): - def forward( - ctx, - q: torch.Tensor, - k: torch.Tensor, - v: torch.Tensor, - cu_seqlens_qo: torch.Tensor, - cu_seqlens_kv: torch.Tensor, - max_seqlen_qo: int, - max_seqlen_kv: int, - causal: bool = False, - softmax_scale: Optional[float] = None, - is_varlen: bool = True, - ) -> Tuple[torch.Tensor, torch.Tensor]: - out, lse = _flash_attn_varlen_forward( - q, k, v, - cu_seqlens_qo, cu_seqlens_kv, max_seqlen_qo, max_seqlen_kv, - causal=causal, softmax_scale=softmax_scale, - is_varlen=is_varlen, - ) - ctx.save_for_backward(q, k, v, out, lse, cu_seqlens_qo, cu_seqlens_kv) - ctx.max_seqlen_qo = max_seqlen_qo - ctx.max_seqlen_kv = max_seqlen_kv - ctx.causal = causal - ctx.softmax_scale = softmax_scale - ctx.is_varlen = is_varlen - return out, lse - - def backward( - ctx, - do: torch.Tensor, - dlse: torch.Tensor, - ): - del dlse # LSE doesn't support backward currently - q, k, v, out, lse, cu_seqlens_qo, cu_seqlens_kv = ctx.saved_tensors - dq, dk, dv = _flash_attn_varlen_backward( - do, q, k, v, out, lse, - cu_seqlens_qo, cu_seqlens_kv, ctx.max_seqlen_qo, ctx.max_seqlen_kv, - causal=ctx.causal, softmax_scale=ctx.softmax_scale, - is_varlen=ctx.is_varlen, - ) - return dq, dk, dv, None, None, None, None, None, None, None - - -def flash_attn_varlen_func( - q: torch.Tensor, - k: torch.Tensor, - v: torch.Tensor, - cu_seqlens_qo: torch.Tensor, - cu_seqlens_kv: torch.Tensor, - max_seqlen_qo: int, - max_seqlen_kv: int, - dropout_p: float = 0.0, - softmax_scale: Optional[float] = None, - causal: bool = False, - deterministic: bool = False, - is_varlen: bool = True, -) -> Tuple[torch.Tensor, torch.Tensor]: - assert dropout_p == 0.0 - assert not deterministic - return FlashAttnVarlenFunc.apply( - q, k, v, - cu_seqlens_qo, cu_seqlens_kv, max_seqlen_qo, max_seqlen_kv, - causal, softmax_scale, is_varlen, - ) - - -def flash_attn_varlen_qkvpacked_func( - qkv: torch.Tensor, - cu_seqlens: torch.Tensor, - max_seqlen: int, - head_dim_qk: int, - dropout_p: float = 0.0, - softmax_scale: Optional[float] = None, - causal: bool = False, - deterministic: bool = False, - is_varlen: bool = True, -) -> Tuple[torch.Tensor, torch.Tensor]: - assert dropout_p == 0.0 - assert not deterministic - return FlashAttnVarlenFunc.apply( - qkv[:, :, :head_dim_qk], qkv[:, :, head_dim_qk:head_dim_qk * 2], qkv[:, :, head_dim_qk * 2:], - cu_seqlens, cu_seqlens, max_seqlen, max_seqlen, - causal, softmax_scale, is_varlen, - ) - - -def flash_attn_varlen_kvpacked_func( - q: torch.Tensor, - kv: torch.Tensor, - cu_seqlens_qo: torch.Tensor, - cu_seqlens_kv: torch.Tensor, - max_seqlen_qo: int, - max_seqlen_kv: int, - head_dim_qk: int, - dropout_p: float = 0.0, - softmax_scale: Optional[float] = None, - causal: bool = False, - deterministic: bool = False, - is_varlen: bool = True, -) -> Tuple[torch.Tensor, torch.Tensor]: - assert dropout_p == 0.0 - assert not deterministic - return FlashAttnVarlenFunc.apply( - q, kv[:, :, :head_dim_qk], kv[:, :, head_dim_qk:], - cu_seqlens_qo, cu_seqlens_kv, max_seqlen_qo, max_seqlen_kv, - causal, softmax_scale, is_varlen, - ) diff --git a/build/torch29-cxx11-cu128-x86_64-linux/metadata.json b/build/torch29-cxx11-cu128-x86_64-linux/metadata.json deleted file mode 100644 index 4899badb63d45293425e2164944268b6058af95d..0000000000000000000000000000000000000000 --- a/build/torch29-cxx11-cu128-x86_64-linux/metadata.json +++ /dev/null @@ -1,11 +0,0 @@ -{ - "version": 1, - "license": "MIT", - "python-depends": [], - "backend": { - "type": "cuda", - "archs": [ - "9.0a" - ] - } -} diff --git a/build/torch29-cxx11-cu129-aarch64-linux/__init__.py b/build/torch29-cxx11-cu129-aarch64-linux/__init__.py deleted file mode 100644 index db300fe9b95176a20b27b3641d89be657d0c4319..0000000000000000000000000000000000000000 --- a/build/torch29-cxx11-cu129-aarch64-linux/__init__.py +++ /dev/null @@ -1,163 +0,0 @@ -from typing import Optional, Tuple -import torch - -from .flash_mla_interface import FlashMLASchedMeta -from . import flash_mla_interface as _impl - - -def get_mla_metadata(*args, **kwargs) -> Tuple[FlashMLASchedMeta, None]: - return _impl.get_mla_metadata(*args, **kwargs) - - -def flash_mla_with_kvcache( - q: torch.Tensor, - k_cache: torch.Tensor, - block_table: Optional[torch.Tensor], - cache_seqlens: Optional[torch.Tensor], - head_dim_v: int, - tile_scheduler_metadata: FlashMLASchedMeta, - num_splits: None = None, - softmax_scale: Optional[float] = None, - causal: bool = False, - is_fp8_kvcache: bool = False, - indices: Optional[torch.Tensor] = None, - attn_sink: Optional[torch.Tensor] = None, - extra_k_cache: Optional[torch.Tensor] = None, - extra_indices_in_kvcache: Optional[torch.Tensor] = None, - topk_length: Optional[torch.Tensor] = None, - extra_topk_length: Optional[torch.Tensor] = None, -) -> Tuple[torch.Tensor, torch.Tensor]: - return _impl.flash_mla_with_kvcache( - q=q, - k_cache=k_cache, - block_table=block_table, - cache_seqlens=cache_seqlens, - head_dim_v=head_dim_v, - tile_scheduler_metadata=tile_scheduler_metadata, - num_splits=num_splits, - softmax_scale=softmax_scale, - causal=causal, - is_fp8_kvcache=is_fp8_kvcache, - indices=indices, - attn_sink=attn_sink, - extra_k_cache=extra_k_cache, - extra_indices_in_kvcache=extra_indices_in_kvcache, - topk_length=topk_length, - extra_topk_length=extra_topk_length, - ) - - -def flash_mla_sparse_fwd( - q: torch.Tensor, - kv: torch.Tensor, - indices: torch.Tensor, - sm_scale: float, - d_v: int = 512, - attn_sink: Optional[torch.Tensor] = None, - topk_length: Optional[torch.Tensor] = None, -) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor]: - return _impl.flash_mla_sparse_fwd( - q=q, - kv=kv, - indices=indices, - sm_scale=sm_scale, - d_v=d_v, - attn_sink=attn_sink, - topk_length=topk_length, - ) - - -def flash_attn_varlen_func( - q: torch.Tensor, - k: torch.Tensor, - v: torch.Tensor, - cu_seqlens_qo: torch.Tensor, - cu_seqlens_kv: torch.Tensor, - max_seqlen_qo: int, - max_seqlen_kv: int, - dropout_p: float = 0.0, - softmax_scale: Optional[float] = None, - causal: bool = False, - deterministic: bool = False, - is_varlen: bool = True, -) -> Tuple[torch.Tensor, torch.Tensor]: - return _impl.flash_attn_varlen_func( - q=q, - k=k, - v=v, - cu_seqlens_qo=cu_seqlens_qo, - cu_seqlens_kv=cu_seqlens_kv, - max_seqlen_qo=max_seqlen_qo, - max_seqlen_kv=max_seqlen_kv, - dropout_p=dropout_p, - softmax_scale=softmax_scale, - causal=causal, - deterministic=deterministic, - is_varlen=is_varlen, - ) - - -def flash_attn_varlen_qkvpacked_func( - qkv: torch.Tensor, - cu_seqlens: torch.Tensor, - max_seqlen: int, - head_dim_qk: int, - dropout_p: float = 0.0, - softmax_scale: Optional[float] = None, - causal: bool = False, - deterministic: bool = False, - is_varlen: bool = True, -) -> Tuple[torch.Tensor, torch.Tensor]: - return _impl.flash_attn_varlen_qkvpacked_func( - qkv=qkv, - cu_seqlens=cu_seqlens, - max_seqlen=max_seqlen, - head_dim_qk=head_dim_qk, - dropout_p=dropout_p, - softmax_scale=softmax_scale, - causal=causal, - deterministic=deterministic, - is_varlen=is_varlen, - ) - - -def flash_attn_varlen_kvpacked_func( - q: torch.Tensor, - kv: torch.Tensor, - cu_seqlens_qo: torch.Tensor, - cu_seqlens_kv: torch.Tensor, - max_seqlen_qo: int, - max_seqlen_kv: int, - head_dim_qk: int, - dropout_p: float = 0.0, - softmax_scale: Optional[float] = None, - causal: bool = False, - deterministic: bool = False, - is_varlen: bool = True, -) -> Tuple[torch.Tensor, torch.Tensor]: - return _impl.flash_attn_varlen_kvpacked_func( - q=q, - kv=kv, - cu_seqlens_qo=cu_seqlens_qo, - cu_seqlens_kv=cu_seqlens_kv, - max_seqlen_qo=max_seqlen_qo, - max_seqlen_kv=max_seqlen_kv, - head_dim_qk=head_dim_qk, - dropout_p=dropout_p, - softmax_scale=softmax_scale, - causal=causal, - deterministic=deterministic, - is_varlen=is_varlen, - ) - - -__all__ = [ - "__version__", - "FlashMLASchedMeta", - "get_mla_metadata", - "flash_mla_with_kvcache", - "flash_attn_varlen_func", - "flash_attn_varlen_qkvpacked_func", - "flash_attn_varlen_kvpacked_func", - "flash_mla_sparse_fwd", -] diff --git a/build/torch29-cxx11-cu129-aarch64-linux/_flash_mla_cuda_a7c2f77.abi3.so b/build/torch29-cxx11-cu129-aarch64-linux/_flash_mla_cuda_a7c2f77.abi3.so deleted file mode 100644 index 12c2057932b18b438c2065c3928f9d2ed02a13b0..0000000000000000000000000000000000000000 --- a/build/torch29-cxx11-cu129-aarch64-linux/_flash_mla_cuda_a7c2f77.abi3.so +++ /dev/null @@ -1,3 +0,0 @@ -version https://git-lfs.github.com/spec/v1 -oid sha256:31d32a58f4621e1e24e5f939d30e9eb1f06deb5c781df1eac063a772e9416022 -size 9363528 diff --git a/build/torch29-cxx11-cu129-aarch64-linux/_ops.py b/build/torch29-cxx11-cu129-aarch64-linux/_ops.py deleted file mode 100644 index 4534f81f82900669a0d8cdb6a9b3a2eeebb7b601..0000000000000000000000000000000000000000 --- a/build/torch29-cxx11-cu129-aarch64-linux/_ops.py +++ /dev/null @@ -1,9 +0,0 @@ -import torch -from . import _flash_mla_cuda_a7c2f77 -ops = torch.ops._flash_mla_cuda_a7c2f77 - -def add_op_namespace_prefix(op_name: str): - """ - Prefix op by namespace. - """ - return f"_flash_mla_cuda_a7c2f77::{op_name}" diff --git a/build/torch29-cxx11-cu129-aarch64-linux/flash_mla/__init__.py b/build/torch29-cxx11-cu129-aarch64-linux/flash_mla/__init__.py deleted file mode 100644 index a9b2672c1cd85b74c1b3ded0fc0b2100e1aeac23..0000000000000000000000000000000000000000 --- a/build/torch29-cxx11-cu129-aarch64-linux/flash_mla/__init__.py +++ /dev/null @@ -1,26 +0,0 @@ -import ctypes -import importlib.util -import sys -from pathlib import Path -from types import ModuleType - - -def _import_from_path(file_path: Path) -> ModuleType: - # We cannot use the module name as-is, after adding it to `sys.modules`, - # it would also be used for other imports. So, we make a module name that - # depends on the path for it to be unique using the hex-encoded hash of - # the path. - path_hash = "{:x}".format(ctypes.c_size_t(hash(file_path.absolute())).value) - module_name = path_hash - spec = importlib.util.spec_from_file_location(module_name, file_path) - if spec is None: - raise ImportError(f"Cannot load spec for {module_name} from {file_path}") - module = importlib.util.module_from_spec(spec) - if module is None: - raise ImportError(f"Cannot load module {module_name} from spec") - sys.modules[module_name] = module - spec.loader.exec_module(module) # type: ignore - return module - - -globals().update(vars(_import_from_path(Path(__file__).parent.parent / "__init__.py"))) diff --git a/build/torch29-cxx11-cu129-aarch64-linux/flash_mla_interface.py b/build/torch29-cxx11-cu129-aarch64-linux/flash_mla_interface.py deleted file mode 100644 index a84e448ffe741bb6d3dafaf7888ed8cc94984467..0000000000000000000000000000000000000000 --- a/build/torch29-cxx11-cu129-aarch64-linux/flash_mla_interface.py +++ /dev/null @@ -1,435 +0,0 @@ -from typing import Optional, Tuple -import dataclasses - -import torch - -from ._ops import ops as flash_mla_cuda - -@dataclasses.dataclass -class FlashMLASchedMeta: - """ - A class that stores the tile scheduler metadata of FlashMLA - """ - - @dataclasses.dataclass - class Config: - b: int - s_q: int - h_q: int - page_block_size: int - h_k: int - - causal: bool - is_fp8_kvcache: bool - topk: Optional[int] - - extra_page_block_size: Optional[int] - extra_topk: Optional[int] - - have_initialized: bool = False - - config: Optional[Config] = None - - tile_scheduler_metadata: Optional[torch.Tensor] = None # (num_sm_parts, TileSchedulerMetaDataSize), dtype torch.int32. - num_splits: Optional[torch.Tensor] = None # (1), dtype torch.int32. - - -def get_mla_metadata( - *args, - **kwargs -) -> Tuple[FlashMLASchedMeta, None]: - """ - Returns an empty instance of FlashMLASchedMeta. The actual scheduling metadata will be generated during the first invocation of flash_mla_with_kvcache. - - Arguments: - This function does not need any arguments, but we keep *args and **kwargs to be compatible with the old interface. - - Return: - A tuple. Due to historical reasons, we return a tuple of (FlashMLASchedMeta, None) now. Only the first element is useful. - """ - return FlashMLASchedMeta(), None - - -def flash_mla_with_kvcache( - q: torch.Tensor, - k_cache: torch.Tensor, - block_table: Optional[torch.Tensor], - cache_seqlens: Optional[torch.Tensor], - head_dim_v: int, - tile_scheduler_metadata: FlashMLASchedMeta, - num_splits: None = None, - softmax_scale: Optional[float] = None, - causal: bool = False, - is_fp8_kvcache: bool = False, - indices: Optional[torch.Tensor] = None, - attn_sink: Optional[torch.Tensor] = None, - extra_k_cache: Optional[torch.Tensor] = None, - extra_indices_in_kvcache: Optional[torch.Tensor] = None, - topk_length: Optional[torch.Tensor] = None, - extra_topk_length: Optional[torch.Tensor] = None -) -> Tuple[torch.Tensor, torch.Tensor]: - """ - Arguments: - q: (batch_size, seq_len_q, num_heads_q, head_dim). - k_cache: (num_blocks, page_block_size, num_heads_k, head_dim). - Different modes (including fp8/bf16, and sparsity) has different KV cache layouts. See comments below for details. - The KV cache must be contiguously valid for sparse attention on sm100. Here "contiguously valid" means that every byte, from the very beginning of the KV cache, till the last byte in the KV cache, is valid memory address to visit (i.e. won't IMA). In other words, the KV cache could be a slice of a larger array, but cannot be a list of disjoint memory blocks. - block_table: (batch_size, max_num_blocks_per_seq), torch.int32. Can be None when sparse attention is used. - cache_seqlens: (batch_size), torch.int32. Can be None when sparse attention is used. - head_dim_v: Head_dim of v. Must be 512 - sched_meta: FlashMLASchedMeta, return by get_mla_metadata. You may reuse the same sched_meta across different invocations, but only when the tensor shapes and the values of cache_seqlens, topk_length, and extra_topk_length remain the same. - num_splits_placeholder: must be "None" (to be compatible with the old interface). - softmax_scale: float. The scaling of QK^T before applying softmax. Default to 1 / sqrt(head_dim_k). - causal: bool. Whether to apply causal attention mask. Only valid for dense attention - is_fp8_kvcache: bool. - indices: (batch_size, seq_len_q, topk). KV indices when sparse attention is enabled. - Pay attention that indices_in_kvcache[i][j][k] = (the index of the page block where token t resides) * block_size + (the offset of token t among the page block), - where t is the k-th token of the j-th q-sequence in the i-th batch. - attn_sink: Optional[torch.Tensor], (num_heads_q, ), torch.float32. If presented, the final output will be scaled by exp(lse) / (exp(lse) + exp(attn_sink)). Have no affect on the returned softmax_lse. +inf will cause the result to become 0. - extra_k_cache and extra_indices_in_kvcache: If provided, will attend to these extra tokens in addition to those in k_cache and indices_in_kvcache. Their format requirements are the same as k_cache and indices_in_kvcache respectively. - topk_length/extra_topk_length: (batch_size, ), torch.int32. If provided, only the leftmost topk_length indices will be processed. Useful when the actual topk for different queries are different so that we can save some computation, compared to masking. - - For DeepSeek V3, DeepSeek V3.1, and DeepSeek V3.2: - head_dim should be 576 while head_dim_v should be 512. - In FP8+sparse mode, each token's KV cache is 656 Bytes, structured as: - - The shape of the tensor `k_cache` is (num_blocks, page_block_size, num_heads_k, head_dim), and num_heads_k must be 1. - - First 512 bytes: The "quantized NoPE" part, containing 512 float8_e4m3 values. - - Next 16 bytes: Scale factors, containing 4 float32 values. The first float32 is the scale for the first 128 float8_e4m3 values, the second for the next 128, and so on. - - Last 128 bytes: The "RoPE" part, containing 64 bfloat16 values. This part is not quantized for accuracy. - - Return: - out: (batch_size, seq_len_q, num_heads_q, head_dim_v). - softmax_lse: (batch_size, num_heads_q, seq_len_q), torch.float32. - """ - sched_meta = tile_scheduler_metadata - indices_in_kvcache = indices - assert isinstance(sched_meta, FlashMLASchedMeta), "tile_scheduler_metadata must be of type FlashMLASchedMeta" - assert num_splits is None, "num_splits must be None" - - topk = indices_in_kvcache.shape[-1] if indices_in_kvcache is not None else None - extra_k_page_block_size = extra_k_cache.shape[1] if extra_k_cache is not None else None - extra_topk = extra_indices_in_kvcache.shape[-1] if extra_indices_in_kvcache is not None else None - if softmax_scale is None: - softmax_scale = q.shape[-1] ** (-0.5) - - if not sched_meta.have_initialized: - # Sanity check. We only perform sanity check during the first invocation to save CPU time. - if indices_in_kvcache is not None: - assert not causal, "causal must be False when indices_in_kvcache is not None (i.e. sparse attention is enabled)" - - # Initialize the tile scheduler metadata during the first invocation. - sched_meta.have_initialized = True - sched_meta.config = FlashMLASchedMeta.Config( - q.shape[0], - q.shape[1], - q.shape[2], - k_cache.shape[1], - k_cache.shape[2], - - causal, - is_fp8_kvcache, - topk, - - extra_k_page_block_size, - extra_topk, - ) - else: - # Check whether the input arguments are consistent with sched_meta - helper_msg = " Your input arguments are inconsistent with sched_meta. Please make sure the input arguments are consistent across different invocations of flash_mla_with_kvcache on the same sched_meta." - assert sched_meta.config is not None - assert sched_meta.config.b == q.shape[0], "sched_meta.config.b must be equal to batch_size." + helper_msg - assert sched_meta.config.s_q == q.shape[1], "sched_meta.config.s_q must be equal to seq_len_q." + helper_msg - assert sched_meta.config.h_q == q.shape[2], "sched_meta.config.h_q must be equal to num_heads_q." + helper_msg - assert sched_meta.config.page_block_size == k_cache.shape[1], "sched_meta.config.page_block_size must be equal to page_block_size." + helper_msg - assert sched_meta.config.h_k == k_cache.shape[2], "sched_meta.config.h_k must be equal to num_heads_k." + helper_msg - assert sched_meta.config.causal == causal, "sched_meta.config.causal must be equal to causal." + helper_msg - assert sched_meta.config.is_fp8_kvcache == is_fp8_kvcache, "sched_meta.config.is_fp8_kvcache must be equal to is_fp8_kvcache." + helper_msg - assert sched_meta.config.topk == topk, "sched_meta.config.topk must be equal to the last dim of indices_in_kvcache." + helper_msg - assert sched_meta.config.extra_page_block_size == extra_k_page_block_size, "sched_meta.config.extra_page_block_size must be equal to the page_block_size of extra_k_cache." + helper_msg - assert sched_meta.config.extra_topk == extra_topk, "sched_meta.config.extra_topk must be equal to the last dim of extra_indices_in_kvcache." + helper_msg - - if topk is not None: - # Sparse attention - assert not causal, "causal must be False when sparse attention is enabled" - assert is_fp8_kvcache, "is_fp8_kvcache must be True when sparse attention is enabled" - out, lse, new_tile_scheduler_metadata, new_num_splits = flash_mla_cuda.sparse_decode_fwd( - q, k_cache, indices_in_kvcache, topk_length, attn_sink, - sched_meta.tile_scheduler_metadata, sched_meta.num_splits, - extra_k_cache, extra_indices_in_kvcache, extra_topk_length, - head_dim_v, softmax_scale - ) - else: - # Dense attention - assert indices_in_kvcache is None and attn_sink is None and extra_k_cache is None and extra_indices_in_kvcache is None and topk_length is None and extra_topk_length is None, "indices_in_kvcache, attn_sink, extra_k_cache, extra_indices_in_kvcache, topk_length and extra_topk_length must be None when dense attention is used." - assert block_table is not None and cache_seqlens is not None, "block_table and cache_seqlens must be provided when dense attention is used." - out, lse, new_tile_scheduler_metadata, new_num_splits = flash_mla_cuda.dense_decode_fwd( - q, k_cache, head_dim_v, - cache_seqlens, block_table, - softmax_scale, causal, - sched_meta.tile_scheduler_metadata, sched_meta.num_splits - ) - sched_meta.tile_scheduler_metadata = new_tile_scheduler_metadata - sched_meta.num_splits = new_num_splits - return (out, lse) - - -def flash_mla_sparse_fwd( - q: torch.Tensor, - kv: torch.Tensor, - indices: torch.Tensor, - sm_scale: float, - d_v: int = 512, - attn_sink: Optional[torch.Tensor] = None, - topk_length: Optional[torch.Tensor] = None, -) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor]: - """ - Sparse attention prefill kernel - - Args: - q: [s_q, h_q, d_qk], bfloat16 - kv: [s_kv, h_kv, d_qk], bfloat16 - indices: [s_q, h_kv, topk], int32. Invalid indices should be set to -1 or numbers >= s_kv - sm_scale: float - d_v: The dimension of value vectors. Can only be 512 - attn_sink: optional, [h_q], float32. - If attn_sink is provided, when computing output, output will be additionally multiplied by exp(lse) / (exp(lse) + exp(attn_sink)). - +-inf in attn_sink will be handled normally (i.e., -inf has no effect, +inf will make corresponding output all zeros). - This argument has no effect on lse and max_logits. - topk_length: optional, [s_q], int32. If provided, the i-th q token will only attend to k tokens specified by indices[i, :, :topk_length[i]], ignoring later k/v tokens (even if provided in indices). - In extremely rare cases (topk_length provided, there is a valid topk index between topk_length[i] ~ s_kv, and that topk index points to a k token containing NaN), operator output will contain NaN, so please avoid this situation. - - Returns: - (output, max_logits, lse) - Please refer to tests/ref.py for the precise definitions of these parameters. - - output: [s_q, h_q, d_v], bfloat16 - - max_logits: [s_q, h_q], float - - lse: [s_q, h_q], float, log-sum-exp of attention scores - """ - results = flash_mla_cuda.sparse_prefill_fwd( - q, kv, indices, sm_scale, d_v, attn_sink, topk_length - ) - return results - - -def _flash_attn_varlen_forward( - q: torch.Tensor, - k: torch.Tensor, - v: torch.Tensor, - cu_seqlens_qo: torch.Tensor, - cu_seqlens_kv: torch.Tensor, - max_seqlen_qo: int, - max_seqlen_kv: int, - out: Optional[torch.Tensor] = None, - lse: Optional[torch.Tensor] = None, - causal: bool = False, - softmax_scale: Optional[float] = None, - is_varlen: bool = True, -) -> Tuple[torch.Tensor, torch.Tensor]: - qo_total_len, num_qo_heads, head_dim_qk = q.shape - kv_total_len, num_kv_heads, head_dim_vo = v.shape - - mask_mode_code = 1 if causal else 0 - if softmax_scale is None: - softmax_scale = head_dim_qk ** (-0.5) - - if out is None: - out = torch.empty(qo_total_len, num_qo_heads, head_dim_vo, device=q.device, dtype=q.dtype) - if lse is None: - # Make lse contiguous on seqlen dim - lse = torch.empty(num_qo_heads, qo_total_len, device=q.device, dtype=torch.float32).T - - workspace_buffer = torch.empty(32 * 1024 * 1024, dtype=torch.uint8, device=q.device) - flash_mla_cuda.dense_prefill_fwd( - workspace_buffer, - q, - k, - v, - cu_seqlens_qo, - cu_seqlens_kv, - out, - lse, - mask_mode_code, - softmax_scale, - max_seqlen_qo, - max_seqlen_kv, - is_varlen, - ) - - return out, lse - - -def _flash_attn_varlen_backward( - do: torch.Tensor, - q: torch.Tensor, - k: torch.Tensor, - v: torch.Tensor, - out: torch.Tensor, - lse: torch.Tensor, - cu_seqlens_qo: torch.Tensor, - cu_seqlens_kv: torch.Tensor, - max_seqlen_qo: int, - max_seqlen_kv: int, - dq: Optional[torch.Tensor] = None, - dk: Optional[torch.Tensor] = None, - dv: Optional[torch.Tensor] = None, - causal: bool = False, - softmax_scale: Optional[float] = None, - is_varlen: bool = True, -) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor]: - qo_total_len, num_qo_heads, head_dim_qk = q.shape - kv_total_len, num_kv_heads, head_dim_vo = v.shape - - # TODO: fix bwd GQA - if num_qo_heads != num_kv_heads: - raise ValueError(f"SM100 bwd doesn't support GQA now. num_qo_heads: {num_qo_heads}, num_kv_heads: {num_kv_heads}.") - - mask_mode_code = 1 if causal else 0 - if softmax_scale is None: - softmax_scale = head_dim_qk ** (-0.5) - - if dq is None: - dq = torch.empty(qo_total_len, num_qo_heads, head_dim_qk, device=q.device, dtype=q.dtype) - if dk is None: - dk = torch.empty(kv_total_len, num_kv_heads, head_dim_qk, device=q.device, dtype=q.dtype) - if dv is None: - dv = torch.empty(kv_total_len, num_kv_heads, head_dim_vo, device=q.device, dtype=q.dtype) - - max_seqlen_qo_aligned = (max_seqlen_qo + 7) // 8 * 8 - bs = cu_seqlens_qo.shape[0] - 1 - workspace_bytes = 0 - workspace_bytes += 4 * bs * max_seqlen_qo_aligned * num_qo_heads * head_dim_qk # dQ_acc - workspace_bytes += 4 * max_seqlen_qo_aligned * bs * num_qo_heads * 2 # sum_OdO and scaled_lse - if num_qo_heads != num_kv_heads: - workspace_bytes += 2 * kv_total_len * num_qo_heads * (head_dim_qk + head_dim_vo) # dKV_acc - workspace_buffer = torch.empty(workspace_bytes, dtype=torch.uint8, device=q.device) - flash_mla_cuda.dense_prefill_bwd( - workspace_buffer, - do, - q, - k, - v, - out, - lse, - cu_seqlens_qo, - cu_seqlens_kv, - dq, - dk, - dv, - mask_mode_code, - softmax_scale, - max_seqlen_qo, - max_seqlen_kv, - is_varlen, - ) - - return dq, dk, dv - - -class FlashAttnVarlenFunc(torch.autograd.Function): - def forward( - ctx, - q: torch.Tensor, - k: torch.Tensor, - v: torch.Tensor, - cu_seqlens_qo: torch.Tensor, - cu_seqlens_kv: torch.Tensor, - max_seqlen_qo: int, - max_seqlen_kv: int, - causal: bool = False, - softmax_scale: Optional[float] = None, - is_varlen: bool = True, - ) -> Tuple[torch.Tensor, torch.Tensor]: - out, lse = _flash_attn_varlen_forward( - q, k, v, - cu_seqlens_qo, cu_seqlens_kv, max_seqlen_qo, max_seqlen_kv, - causal=causal, softmax_scale=softmax_scale, - is_varlen=is_varlen, - ) - ctx.save_for_backward(q, k, v, out, lse, cu_seqlens_qo, cu_seqlens_kv) - ctx.max_seqlen_qo = max_seqlen_qo - ctx.max_seqlen_kv = max_seqlen_kv - ctx.causal = causal - ctx.softmax_scale = softmax_scale - ctx.is_varlen = is_varlen - return out, lse - - def backward( - ctx, - do: torch.Tensor, - dlse: torch.Tensor, - ): - del dlse # LSE doesn't support backward currently - q, k, v, out, lse, cu_seqlens_qo, cu_seqlens_kv = ctx.saved_tensors - dq, dk, dv = _flash_attn_varlen_backward( - do, q, k, v, out, lse, - cu_seqlens_qo, cu_seqlens_kv, ctx.max_seqlen_qo, ctx.max_seqlen_kv, - causal=ctx.causal, softmax_scale=ctx.softmax_scale, - is_varlen=ctx.is_varlen, - ) - return dq, dk, dv, None, None, None, None, None, None, None - - -def flash_attn_varlen_func( - q: torch.Tensor, - k: torch.Tensor, - v: torch.Tensor, - cu_seqlens_qo: torch.Tensor, - cu_seqlens_kv: torch.Tensor, - max_seqlen_qo: int, - max_seqlen_kv: int, - dropout_p: float = 0.0, - softmax_scale: Optional[float] = None, - causal: bool = False, - deterministic: bool = False, - is_varlen: bool = True, -) -> Tuple[torch.Tensor, torch.Tensor]: - assert dropout_p == 0.0 - assert not deterministic - return FlashAttnVarlenFunc.apply( - q, k, v, - cu_seqlens_qo, cu_seqlens_kv, max_seqlen_qo, max_seqlen_kv, - causal, softmax_scale, is_varlen, - ) - - -def flash_attn_varlen_qkvpacked_func( - qkv: torch.Tensor, - cu_seqlens: torch.Tensor, - max_seqlen: int, - head_dim_qk: int, - dropout_p: float = 0.0, - softmax_scale: Optional[float] = None, - causal: bool = False, - deterministic: bool = False, - is_varlen: bool = True, -) -> Tuple[torch.Tensor, torch.Tensor]: - assert dropout_p == 0.0 - assert not deterministic - return FlashAttnVarlenFunc.apply( - qkv[:, :, :head_dim_qk], qkv[:, :, head_dim_qk:head_dim_qk * 2], qkv[:, :, head_dim_qk * 2:], - cu_seqlens, cu_seqlens, max_seqlen, max_seqlen, - causal, softmax_scale, is_varlen, - ) - - -def flash_attn_varlen_kvpacked_func( - q: torch.Tensor, - kv: torch.Tensor, - cu_seqlens_qo: torch.Tensor, - cu_seqlens_kv: torch.Tensor, - max_seqlen_qo: int, - max_seqlen_kv: int, - head_dim_qk: int, - dropout_p: float = 0.0, - softmax_scale: Optional[float] = None, - causal: bool = False, - deterministic: bool = False, - is_varlen: bool = True, -) -> Tuple[torch.Tensor, torch.Tensor]: - assert dropout_p == 0.0 - assert not deterministic - return FlashAttnVarlenFunc.apply( - q, kv[:, :, :head_dim_qk], kv[:, :, head_dim_qk:], - cu_seqlens_qo, cu_seqlens_kv, max_seqlen_qo, max_seqlen_kv, - causal, softmax_scale, is_varlen, - ) diff --git a/build/torch29-cxx11-cu129-aarch64-linux/metadata.json b/build/torch29-cxx11-cu129-aarch64-linux/metadata.json deleted file mode 100644 index 8190d75efa8fd6449ddcd73de2072f17086e0842..0000000000000000000000000000000000000000 --- a/build/torch29-cxx11-cu129-aarch64-linux/metadata.json +++ /dev/null @@ -1,12 +0,0 @@ -{ - "version": 1, - "license": "MIT", - "python-depends": [], - "backend": { - "type": "cuda", - "archs": [ - "10.0f", - "9.0a" - ] - } -} diff --git a/build/torch29-cxx11-cu129-x86_64-linux/__init__.py b/build/torch29-cxx11-cu129-x86_64-linux/__init__.py deleted file mode 100644 index db300fe9b95176a20b27b3641d89be657d0c4319..0000000000000000000000000000000000000000 --- a/build/torch29-cxx11-cu129-x86_64-linux/__init__.py +++ /dev/null @@ -1,163 +0,0 @@ -from typing import Optional, Tuple -import torch - -from .flash_mla_interface import FlashMLASchedMeta -from . import flash_mla_interface as _impl - - -def get_mla_metadata(*args, **kwargs) -> Tuple[FlashMLASchedMeta, None]: - return _impl.get_mla_metadata(*args, **kwargs) - - -def flash_mla_with_kvcache( - q: torch.Tensor, - k_cache: torch.Tensor, - block_table: Optional[torch.Tensor], - cache_seqlens: Optional[torch.Tensor], - head_dim_v: int, - tile_scheduler_metadata: FlashMLASchedMeta, - num_splits: None = None, - softmax_scale: Optional[float] = None, - causal: bool = False, - is_fp8_kvcache: bool = False, - indices: Optional[torch.Tensor] = None, - attn_sink: Optional[torch.Tensor] = None, - extra_k_cache: Optional[torch.Tensor] = None, - extra_indices_in_kvcache: Optional[torch.Tensor] = None, - topk_length: Optional[torch.Tensor] = None, - extra_topk_length: Optional[torch.Tensor] = None, -) -> Tuple[torch.Tensor, torch.Tensor]: - return _impl.flash_mla_with_kvcache( - q=q, - k_cache=k_cache, - block_table=block_table, - cache_seqlens=cache_seqlens, - head_dim_v=head_dim_v, - tile_scheduler_metadata=tile_scheduler_metadata, - num_splits=num_splits, - softmax_scale=softmax_scale, - causal=causal, - is_fp8_kvcache=is_fp8_kvcache, - indices=indices, - attn_sink=attn_sink, - extra_k_cache=extra_k_cache, - extra_indices_in_kvcache=extra_indices_in_kvcache, - topk_length=topk_length, - extra_topk_length=extra_topk_length, - ) - - -def flash_mla_sparse_fwd( - q: torch.Tensor, - kv: torch.Tensor, - indices: torch.Tensor, - sm_scale: float, - d_v: int = 512, - attn_sink: Optional[torch.Tensor] = None, - topk_length: Optional[torch.Tensor] = None, -) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor]: - return _impl.flash_mla_sparse_fwd( - q=q, - kv=kv, - indices=indices, - sm_scale=sm_scale, - d_v=d_v, - attn_sink=attn_sink, - topk_length=topk_length, - ) - - -def flash_attn_varlen_func( - q: torch.Tensor, - k: torch.Tensor, - v: torch.Tensor, - cu_seqlens_qo: torch.Tensor, - cu_seqlens_kv: torch.Tensor, - max_seqlen_qo: int, - max_seqlen_kv: int, - dropout_p: float = 0.0, - softmax_scale: Optional[float] = None, - causal: bool = False, - deterministic: bool = False, - is_varlen: bool = True, -) -> Tuple[torch.Tensor, torch.Tensor]: - return _impl.flash_attn_varlen_func( - q=q, - k=k, - v=v, - cu_seqlens_qo=cu_seqlens_qo, - cu_seqlens_kv=cu_seqlens_kv, - max_seqlen_qo=max_seqlen_qo, - max_seqlen_kv=max_seqlen_kv, - dropout_p=dropout_p, - softmax_scale=softmax_scale, - causal=causal, - deterministic=deterministic, - is_varlen=is_varlen, - ) - - -def flash_attn_varlen_qkvpacked_func( - qkv: torch.Tensor, - cu_seqlens: torch.Tensor, - max_seqlen: int, - head_dim_qk: int, - dropout_p: float = 0.0, - softmax_scale: Optional[float] = None, - causal: bool = False, - deterministic: bool = False, - is_varlen: bool = True, -) -> Tuple[torch.Tensor, torch.Tensor]: - return _impl.flash_attn_varlen_qkvpacked_func( - qkv=qkv, - cu_seqlens=cu_seqlens, - max_seqlen=max_seqlen, - head_dim_qk=head_dim_qk, - dropout_p=dropout_p, - softmax_scale=softmax_scale, - causal=causal, - deterministic=deterministic, - is_varlen=is_varlen, - ) - - -def flash_attn_varlen_kvpacked_func( - q: torch.Tensor, - kv: torch.Tensor, - cu_seqlens_qo: torch.Tensor, - cu_seqlens_kv: torch.Tensor, - max_seqlen_qo: int, - max_seqlen_kv: int, - head_dim_qk: int, - dropout_p: float = 0.0, - softmax_scale: Optional[float] = None, - causal: bool = False, - deterministic: bool = False, - is_varlen: bool = True, -) -> Tuple[torch.Tensor, torch.Tensor]: - return _impl.flash_attn_varlen_kvpacked_func( - q=q, - kv=kv, - cu_seqlens_qo=cu_seqlens_qo, - cu_seqlens_kv=cu_seqlens_kv, - max_seqlen_qo=max_seqlen_qo, - max_seqlen_kv=max_seqlen_kv, - head_dim_qk=head_dim_qk, - dropout_p=dropout_p, - softmax_scale=softmax_scale, - causal=causal, - deterministic=deterministic, - is_varlen=is_varlen, - ) - - -__all__ = [ - "__version__", - "FlashMLASchedMeta", - "get_mla_metadata", - "flash_mla_with_kvcache", - "flash_attn_varlen_func", - "flash_attn_varlen_qkvpacked_func", - "flash_attn_varlen_kvpacked_func", - "flash_mla_sparse_fwd", -] diff --git a/build/torch29-cxx11-cu129-x86_64-linux/_flash_mla_cuda_a7c2f77.abi3.so b/build/torch29-cxx11-cu129-x86_64-linux/_flash_mla_cuda_a7c2f77.abi3.so deleted file mode 100644 index 961149a477dcbbd74ae0065bbbfb85eeb2466750..0000000000000000000000000000000000000000 --- a/build/torch29-cxx11-cu129-x86_64-linux/_flash_mla_cuda_a7c2f77.abi3.so +++ /dev/null @@ -1,3 +0,0 @@ -version https://git-lfs.github.com/spec/v1 -oid sha256:2310ab5cf6b5d1f535f9d71ecca1ee0e1f098c57b3cb332f786fe75f409a1e9a -size 9283880 diff --git a/build/torch29-cxx11-cu129-x86_64-linux/_ops.py b/build/torch29-cxx11-cu129-x86_64-linux/_ops.py deleted file mode 100644 index 4534f81f82900669a0d8cdb6a9b3a2eeebb7b601..0000000000000000000000000000000000000000 --- a/build/torch29-cxx11-cu129-x86_64-linux/_ops.py +++ /dev/null @@ -1,9 +0,0 @@ -import torch -from . import _flash_mla_cuda_a7c2f77 -ops = torch.ops._flash_mla_cuda_a7c2f77 - -def add_op_namespace_prefix(op_name: str): - """ - Prefix op by namespace. - """ - return f"_flash_mla_cuda_a7c2f77::{op_name}" diff --git a/build/torch29-cxx11-cu129-x86_64-linux/flash_mla/__init__.py b/build/torch29-cxx11-cu129-x86_64-linux/flash_mla/__init__.py deleted file mode 100644 index a9b2672c1cd85b74c1b3ded0fc0b2100e1aeac23..0000000000000000000000000000000000000000 --- a/build/torch29-cxx11-cu129-x86_64-linux/flash_mla/__init__.py +++ /dev/null @@ -1,26 +0,0 @@ -import ctypes -import importlib.util -import sys -from pathlib import Path -from types import ModuleType - - -def _import_from_path(file_path: Path) -> ModuleType: - # We cannot use the module name as-is, after adding it to `sys.modules`, - # it would also be used for other imports. So, we make a module name that - # depends on the path for it to be unique using the hex-encoded hash of - # the path. - path_hash = "{:x}".format(ctypes.c_size_t(hash(file_path.absolute())).value) - module_name = path_hash - spec = importlib.util.spec_from_file_location(module_name, file_path) - if spec is None: - raise ImportError(f"Cannot load spec for {module_name} from {file_path}") - module = importlib.util.module_from_spec(spec) - if module is None: - raise ImportError(f"Cannot load module {module_name} from spec") - sys.modules[module_name] = module - spec.loader.exec_module(module) # type: ignore - return module - - -globals().update(vars(_import_from_path(Path(__file__).parent.parent / "__init__.py"))) diff --git a/build/torch29-cxx11-cu129-x86_64-linux/flash_mla_interface.py b/build/torch29-cxx11-cu129-x86_64-linux/flash_mla_interface.py deleted file mode 100644 index a84e448ffe741bb6d3dafaf7888ed8cc94984467..0000000000000000000000000000000000000000 --- a/build/torch29-cxx11-cu129-x86_64-linux/flash_mla_interface.py +++ /dev/null @@ -1,435 +0,0 @@ -from typing import Optional, Tuple -import dataclasses - -import torch - -from ._ops import ops as flash_mla_cuda - -@dataclasses.dataclass -class FlashMLASchedMeta: - """ - A class that stores the tile scheduler metadata of FlashMLA - """ - - @dataclasses.dataclass - class Config: - b: int - s_q: int - h_q: int - page_block_size: int - h_k: int - - causal: bool - is_fp8_kvcache: bool - topk: Optional[int] - - extra_page_block_size: Optional[int] - extra_topk: Optional[int] - - have_initialized: bool = False - - config: Optional[Config] = None - - tile_scheduler_metadata: Optional[torch.Tensor] = None # (num_sm_parts, TileSchedulerMetaDataSize), dtype torch.int32. - num_splits: Optional[torch.Tensor] = None # (1), dtype torch.int32. - - -def get_mla_metadata( - *args, - **kwargs -) -> Tuple[FlashMLASchedMeta, None]: - """ - Returns an empty instance of FlashMLASchedMeta. The actual scheduling metadata will be generated during the first invocation of flash_mla_with_kvcache. - - Arguments: - This function does not need any arguments, but we keep *args and **kwargs to be compatible with the old interface. - - Return: - A tuple. Due to historical reasons, we return a tuple of (FlashMLASchedMeta, None) now. Only the first element is useful. - """ - return FlashMLASchedMeta(), None - - -def flash_mla_with_kvcache( - q: torch.Tensor, - k_cache: torch.Tensor, - block_table: Optional[torch.Tensor], - cache_seqlens: Optional[torch.Tensor], - head_dim_v: int, - tile_scheduler_metadata: FlashMLASchedMeta, - num_splits: None = None, - softmax_scale: Optional[float] = None, - causal: bool = False, - is_fp8_kvcache: bool = False, - indices: Optional[torch.Tensor] = None, - attn_sink: Optional[torch.Tensor] = None, - extra_k_cache: Optional[torch.Tensor] = None, - extra_indices_in_kvcache: Optional[torch.Tensor] = None, - topk_length: Optional[torch.Tensor] = None, - extra_topk_length: Optional[torch.Tensor] = None -) -> Tuple[torch.Tensor, torch.Tensor]: - """ - Arguments: - q: (batch_size, seq_len_q, num_heads_q, head_dim). - k_cache: (num_blocks, page_block_size, num_heads_k, head_dim). - Different modes (including fp8/bf16, and sparsity) has different KV cache layouts. See comments below for details. - The KV cache must be contiguously valid for sparse attention on sm100. Here "contiguously valid" means that every byte, from the very beginning of the KV cache, till the last byte in the KV cache, is valid memory address to visit (i.e. won't IMA). In other words, the KV cache could be a slice of a larger array, but cannot be a list of disjoint memory blocks. - block_table: (batch_size, max_num_blocks_per_seq), torch.int32. Can be None when sparse attention is used. - cache_seqlens: (batch_size), torch.int32. Can be None when sparse attention is used. - head_dim_v: Head_dim of v. Must be 512 - sched_meta: FlashMLASchedMeta, return by get_mla_metadata. You may reuse the same sched_meta across different invocations, but only when the tensor shapes and the values of cache_seqlens, topk_length, and extra_topk_length remain the same. - num_splits_placeholder: must be "None" (to be compatible with the old interface). - softmax_scale: float. The scaling of QK^T before applying softmax. Default to 1 / sqrt(head_dim_k). - causal: bool. Whether to apply causal attention mask. Only valid for dense attention - is_fp8_kvcache: bool. - indices: (batch_size, seq_len_q, topk). KV indices when sparse attention is enabled. - Pay attention that indices_in_kvcache[i][j][k] = (the index of the page block where token t resides) * block_size + (the offset of token t among the page block), - where t is the k-th token of the j-th q-sequence in the i-th batch. - attn_sink: Optional[torch.Tensor], (num_heads_q, ), torch.float32. If presented, the final output will be scaled by exp(lse) / (exp(lse) + exp(attn_sink)). Have no affect on the returned softmax_lse. +inf will cause the result to become 0. - extra_k_cache and extra_indices_in_kvcache: If provided, will attend to these extra tokens in addition to those in k_cache and indices_in_kvcache. Their format requirements are the same as k_cache and indices_in_kvcache respectively. - topk_length/extra_topk_length: (batch_size, ), torch.int32. If provided, only the leftmost topk_length indices will be processed. Useful when the actual topk for different queries are different so that we can save some computation, compared to masking. - - For DeepSeek V3, DeepSeek V3.1, and DeepSeek V3.2: - head_dim should be 576 while head_dim_v should be 512. - In FP8+sparse mode, each token's KV cache is 656 Bytes, structured as: - - The shape of the tensor `k_cache` is (num_blocks, page_block_size, num_heads_k, head_dim), and num_heads_k must be 1. - - First 512 bytes: The "quantized NoPE" part, containing 512 float8_e4m3 values. - - Next 16 bytes: Scale factors, containing 4 float32 values. The first float32 is the scale for the first 128 float8_e4m3 values, the second for the next 128, and so on. - - Last 128 bytes: The "RoPE" part, containing 64 bfloat16 values. This part is not quantized for accuracy. - - Return: - out: (batch_size, seq_len_q, num_heads_q, head_dim_v). - softmax_lse: (batch_size, num_heads_q, seq_len_q), torch.float32. - """ - sched_meta = tile_scheduler_metadata - indices_in_kvcache = indices - assert isinstance(sched_meta, FlashMLASchedMeta), "tile_scheduler_metadata must be of type FlashMLASchedMeta" - assert num_splits is None, "num_splits must be None" - - topk = indices_in_kvcache.shape[-1] if indices_in_kvcache is not None else None - extra_k_page_block_size = extra_k_cache.shape[1] if extra_k_cache is not None else None - extra_topk = extra_indices_in_kvcache.shape[-1] if extra_indices_in_kvcache is not None else None - if softmax_scale is None: - softmax_scale = q.shape[-1] ** (-0.5) - - if not sched_meta.have_initialized: - # Sanity check. We only perform sanity check during the first invocation to save CPU time. - if indices_in_kvcache is not None: - assert not causal, "causal must be False when indices_in_kvcache is not None (i.e. sparse attention is enabled)" - - # Initialize the tile scheduler metadata during the first invocation. - sched_meta.have_initialized = True - sched_meta.config = FlashMLASchedMeta.Config( - q.shape[0], - q.shape[1], - q.shape[2], - k_cache.shape[1], - k_cache.shape[2], - - causal, - is_fp8_kvcache, - topk, - - extra_k_page_block_size, - extra_topk, - ) - else: - # Check whether the input arguments are consistent with sched_meta - helper_msg = " Your input arguments are inconsistent with sched_meta. Please make sure the input arguments are consistent across different invocations of flash_mla_with_kvcache on the same sched_meta." - assert sched_meta.config is not None - assert sched_meta.config.b == q.shape[0], "sched_meta.config.b must be equal to batch_size." + helper_msg - assert sched_meta.config.s_q == q.shape[1], "sched_meta.config.s_q must be equal to seq_len_q." + helper_msg - assert sched_meta.config.h_q == q.shape[2], "sched_meta.config.h_q must be equal to num_heads_q." + helper_msg - assert sched_meta.config.page_block_size == k_cache.shape[1], "sched_meta.config.page_block_size must be equal to page_block_size." + helper_msg - assert sched_meta.config.h_k == k_cache.shape[2], "sched_meta.config.h_k must be equal to num_heads_k." + helper_msg - assert sched_meta.config.causal == causal, "sched_meta.config.causal must be equal to causal." + helper_msg - assert sched_meta.config.is_fp8_kvcache == is_fp8_kvcache, "sched_meta.config.is_fp8_kvcache must be equal to is_fp8_kvcache." + helper_msg - assert sched_meta.config.topk == topk, "sched_meta.config.topk must be equal to the last dim of indices_in_kvcache." + helper_msg - assert sched_meta.config.extra_page_block_size == extra_k_page_block_size, "sched_meta.config.extra_page_block_size must be equal to the page_block_size of extra_k_cache." + helper_msg - assert sched_meta.config.extra_topk == extra_topk, "sched_meta.config.extra_topk must be equal to the last dim of extra_indices_in_kvcache." + helper_msg - - if topk is not None: - # Sparse attention - assert not causal, "causal must be False when sparse attention is enabled" - assert is_fp8_kvcache, "is_fp8_kvcache must be True when sparse attention is enabled" - out, lse, new_tile_scheduler_metadata, new_num_splits = flash_mla_cuda.sparse_decode_fwd( - q, k_cache, indices_in_kvcache, topk_length, attn_sink, - sched_meta.tile_scheduler_metadata, sched_meta.num_splits, - extra_k_cache, extra_indices_in_kvcache, extra_topk_length, - head_dim_v, softmax_scale - ) - else: - # Dense attention - assert indices_in_kvcache is None and attn_sink is None and extra_k_cache is None and extra_indices_in_kvcache is None and topk_length is None and extra_topk_length is None, "indices_in_kvcache, attn_sink, extra_k_cache, extra_indices_in_kvcache, topk_length and extra_topk_length must be None when dense attention is used." - assert block_table is not None and cache_seqlens is not None, "block_table and cache_seqlens must be provided when dense attention is used." - out, lse, new_tile_scheduler_metadata, new_num_splits = flash_mla_cuda.dense_decode_fwd( - q, k_cache, head_dim_v, - cache_seqlens, block_table, - softmax_scale, causal, - sched_meta.tile_scheduler_metadata, sched_meta.num_splits - ) - sched_meta.tile_scheduler_metadata = new_tile_scheduler_metadata - sched_meta.num_splits = new_num_splits - return (out, lse) - - -def flash_mla_sparse_fwd( - q: torch.Tensor, - kv: torch.Tensor, - indices: torch.Tensor, - sm_scale: float, - d_v: int = 512, - attn_sink: Optional[torch.Tensor] = None, - topk_length: Optional[torch.Tensor] = None, -) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor]: - """ - Sparse attention prefill kernel - - Args: - q: [s_q, h_q, d_qk], bfloat16 - kv: [s_kv, h_kv, d_qk], bfloat16 - indices: [s_q, h_kv, topk], int32. Invalid indices should be set to -1 or numbers >= s_kv - sm_scale: float - d_v: The dimension of value vectors. Can only be 512 - attn_sink: optional, [h_q], float32. - If attn_sink is provided, when computing output, output will be additionally multiplied by exp(lse) / (exp(lse) + exp(attn_sink)). - +-inf in attn_sink will be handled normally (i.e., -inf has no effect, +inf will make corresponding output all zeros). - This argument has no effect on lse and max_logits. - topk_length: optional, [s_q], int32. If provided, the i-th q token will only attend to k tokens specified by indices[i, :, :topk_length[i]], ignoring later k/v tokens (even if provided in indices). - In extremely rare cases (topk_length provided, there is a valid topk index between topk_length[i] ~ s_kv, and that topk index points to a k token containing NaN), operator output will contain NaN, so please avoid this situation. - - Returns: - (output, max_logits, lse) - Please refer to tests/ref.py for the precise definitions of these parameters. - - output: [s_q, h_q, d_v], bfloat16 - - max_logits: [s_q, h_q], float - - lse: [s_q, h_q], float, log-sum-exp of attention scores - """ - results = flash_mla_cuda.sparse_prefill_fwd( - q, kv, indices, sm_scale, d_v, attn_sink, topk_length - ) - return results - - -def _flash_attn_varlen_forward( - q: torch.Tensor, - k: torch.Tensor, - v: torch.Tensor, - cu_seqlens_qo: torch.Tensor, - cu_seqlens_kv: torch.Tensor, - max_seqlen_qo: int, - max_seqlen_kv: int, - out: Optional[torch.Tensor] = None, - lse: Optional[torch.Tensor] = None, - causal: bool = False, - softmax_scale: Optional[float] = None, - is_varlen: bool = True, -) -> Tuple[torch.Tensor, torch.Tensor]: - qo_total_len, num_qo_heads, head_dim_qk = q.shape - kv_total_len, num_kv_heads, head_dim_vo = v.shape - - mask_mode_code = 1 if causal else 0 - if softmax_scale is None: - softmax_scale = head_dim_qk ** (-0.5) - - if out is None: - out = torch.empty(qo_total_len, num_qo_heads, head_dim_vo, device=q.device, dtype=q.dtype) - if lse is None: - # Make lse contiguous on seqlen dim - lse = torch.empty(num_qo_heads, qo_total_len, device=q.device, dtype=torch.float32).T - - workspace_buffer = torch.empty(32 * 1024 * 1024, dtype=torch.uint8, device=q.device) - flash_mla_cuda.dense_prefill_fwd( - workspace_buffer, - q, - k, - v, - cu_seqlens_qo, - cu_seqlens_kv, - out, - lse, - mask_mode_code, - softmax_scale, - max_seqlen_qo, - max_seqlen_kv, - is_varlen, - ) - - return out, lse - - -def _flash_attn_varlen_backward( - do: torch.Tensor, - q: torch.Tensor, - k: torch.Tensor, - v: torch.Tensor, - out: torch.Tensor, - lse: torch.Tensor, - cu_seqlens_qo: torch.Tensor, - cu_seqlens_kv: torch.Tensor, - max_seqlen_qo: int, - max_seqlen_kv: int, - dq: Optional[torch.Tensor] = None, - dk: Optional[torch.Tensor] = None, - dv: Optional[torch.Tensor] = None, - causal: bool = False, - softmax_scale: Optional[float] = None, - is_varlen: bool = True, -) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor]: - qo_total_len, num_qo_heads, head_dim_qk = q.shape - kv_total_len, num_kv_heads, head_dim_vo = v.shape - - # TODO: fix bwd GQA - if num_qo_heads != num_kv_heads: - raise ValueError(f"SM100 bwd doesn't support GQA now. num_qo_heads: {num_qo_heads}, num_kv_heads: {num_kv_heads}.") - - mask_mode_code = 1 if causal else 0 - if softmax_scale is None: - softmax_scale = head_dim_qk ** (-0.5) - - if dq is None: - dq = torch.empty(qo_total_len, num_qo_heads, head_dim_qk, device=q.device, dtype=q.dtype) - if dk is None: - dk = torch.empty(kv_total_len, num_kv_heads, head_dim_qk, device=q.device, dtype=q.dtype) - if dv is None: - dv = torch.empty(kv_total_len, num_kv_heads, head_dim_vo, device=q.device, dtype=q.dtype) - - max_seqlen_qo_aligned = (max_seqlen_qo + 7) // 8 * 8 - bs = cu_seqlens_qo.shape[0] - 1 - workspace_bytes = 0 - workspace_bytes += 4 * bs * max_seqlen_qo_aligned * num_qo_heads * head_dim_qk # dQ_acc - workspace_bytes += 4 * max_seqlen_qo_aligned * bs * num_qo_heads * 2 # sum_OdO and scaled_lse - if num_qo_heads != num_kv_heads: - workspace_bytes += 2 * kv_total_len * num_qo_heads * (head_dim_qk + head_dim_vo) # dKV_acc - workspace_buffer = torch.empty(workspace_bytes, dtype=torch.uint8, device=q.device) - flash_mla_cuda.dense_prefill_bwd( - workspace_buffer, - do, - q, - k, - v, - out, - lse, - cu_seqlens_qo, - cu_seqlens_kv, - dq, - dk, - dv, - mask_mode_code, - softmax_scale, - max_seqlen_qo, - max_seqlen_kv, - is_varlen, - ) - - return dq, dk, dv - - -class FlashAttnVarlenFunc(torch.autograd.Function): - def forward( - ctx, - q: torch.Tensor, - k: torch.Tensor, - v: torch.Tensor, - cu_seqlens_qo: torch.Tensor, - cu_seqlens_kv: torch.Tensor, - max_seqlen_qo: int, - max_seqlen_kv: int, - causal: bool = False, - softmax_scale: Optional[float] = None, - is_varlen: bool = True, - ) -> Tuple[torch.Tensor, torch.Tensor]: - out, lse = _flash_attn_varlen_forward( - q, k, v, - cu_seqlens_qo, cu_seqlens_kv, max_seqlen_qo, max_seqlen_kv, - causal=causal, softmax_scale=softmax_scale, - is_varlen=is_varlen, - ) - ctx.save_for_backward(q, k, v, out, lse, cu_seqlens_qo, cu_seqlens_kv) - ctx.max_seqlen_qo = max_seqlen_qo - ctx.max_seqlen_kv = max_seqlen_kv - ctx.causal = causal - ctx.softmax_scale = softmax_scale - ctx.is_varlen = is_varlen - return out, lse - - def backward( - ctx, - do: torch.Tensor, - dlse: torch.Tensor, - ): - del dlse # LSE doesn't support backward currently - q, k, v, out, lse, cu_seqlens_qo, cu_seqlens_kv = ctx.saved_tensors - dq, dk, dv = _flash_attn_varlen_backward( - do, q, k, v, out, lse, - cu_seqlens_qo, cu_seqlens_kv, ctx.max_seqlen_qo, ctx.max_seqlen_kv, - causal=ctx.causal, softmax_scale=ctx.softmax_scale, - is_varlen=ctx.is_varlen, - ) - return dq, dk, dv, None, None, None, None, None, None, None - - -def flash_attn_varlen_func( - q: torch.Tensor, - k: torch.Tensor, - v: torch.Tensor, - cu_seqlens_qo: torch.Tensor, - cu_seqlens_kv: torch.Tensor, - max_seqlen_qo: int, - max_seqlen_kv: int, - dropout_p: float = 0.0, - softmax_scale: Optional[float] = None, - causal: bool = False, - deterministic: bool = False, - is_varlen: bool = True, -) -> Tuple[torch.Tensor, torch.Tensor]: - assert dropout_p == 0.0 - assert not deterministic - return FlashAttnVarlenFunc.apply( - q, k, v, - cu_seqlens_qo, cu_seqlens_kv, max_seqlen_qo, max_seqlen_kv, - causal, softmax_scale, is_varlen, - ) - - -def flash_attn_varlen_qkvpacked_func( - qkv: torch.Tensor, - cu_seqlens: torch.Tensor, - max_seqlen: int, - head_dim_qk: int, - dropout_p: float = 0.0, - softmax_scale: Optional[float] = None, - causal: bool = False, - deterministic: bool = False, - is_varlen: bool = True, -) -> Tuple[torch.Tensor, torch.Tensor]: - assert dropout_p == 0.0 - assert not deterministic - return FlashAttnVarlenFunc.apply( - qkv[:, :, :head_dim_qk], qkv[:, :, head_dim_qk:head_dim_qk * 2], qkv[:, :, head_dim_qk * 2:], - cu_seqlens, cu_seqlens, max_seqlen, max_seqlen, - causal, softmax_scale, is_varlen, - ) - - -def flash_attn_varlen_kvpacked_func( - q: torch.Tensor, - kv: torch.Tensor, - cu_seqlens_qo: torch.Tensor, - cu_seqlens_kv: torch.Tensor, - max_seqlen_qo: int, - max_seqlen_kv: int, - head_dim_qk: int, - dropout_p: float = 0.0, - softmax_scale: Optional[float] = None, - causal: bool = False, - deterministic: bool = False, - is_varlen: bool = True, -) -> Tuple[torch.Tensor, torch.Tensor]: - assert dropout_p == 0.0 - assert not deterministic - return FlashAttnVarlenFunc.apply( - q, kv[:, :, :head_dim_qk], kv[:, :, head_dim_qk:], - cu_seqlens_qo, cu_seqlens_kv, max_seqlen_qo, max_seqlen_kv, - causal, softmax_scale, is_varlen, - ) diff --git a/build/torch29-cxx11-cu129-x86_64-linux/metadata.json b/build/torch29-cxx11-cu129-x86_64-linux/metadata.json deleted file mode 100644 index 8190d75efa8fd6449ddcd73de2072f17086e0842..0000000000000000000000000000000000000000 --- a/build/torch29-cxx11-cu129-x86_64-linux/metadata.json +++ /dev/null @@ -1,12 +0,0 @@ -{ - "version": 1, - "license": "MIT", - "python-depends": [], - "backend": { - "type": "cuda", - "archs": [ - "10.0f", - "9.0a" - ] - } -} diff --git a/build/torch29-cxx11-cu130-aarch64-linux/__init__.py b/build/torch29-cxx11-cu130-aarch64-linux/__init__.py deleted file mode 100644 index db300fe9b95176a20b27b3641d89be657d0c4319..0000000000000000000000000000000000000000 --- a/build/torch29-cxx11-cu130-aarch64-linux/__init__.py +++ /dev/null @@ -1,163 +0,0 @@ -from typing import Optional, Tuple -import torch - -from .flash_mla_interface import FlashMLASchedMeta -from . import flash_mla_interface as _impl - - -def get_mla_metadata(*args, **kwargs) -> Tuple[FlashMLASchedMeta, None]: - return _impl.get_mla_metadata(*args, **kwargs) - - -def flash_mla_with_kvcache( - q: torch.Tensor, - k_cache: torch.Tensor, - block_table: Optional[torch.Tensor], - cache_seqlens: Optional[torch.Tensor], - head_dim_v: int, - tile_scheduler_metadata: FlashMLASchedMeta, - num_splits: None = None, - softmax_scale: Optional[float] = None, - causal: bool = False, - is_fp8_kvcache: bool = False, - indices: Optional[torch.Tensor] = None, - attn_sink: Optional[torch.Tensor] = None, - extra_k_cache: Optional[torch.Tensor] = None, - extra_indices_in_kvcache: Optional[torch.Tensor] = None, - topk_length: Optional[torch.Tensor] = None, - extra_topk_length: Optional[torch.Tensor] = None, -) -> Tuple[torch.Tensor, torch.Tensor]: - return _impl.flash_mla_with_kvcache( - q=q, - k_cache=k_cache, - block_table=block_table, - cache_seqlens=cache_seqlens, - head_dim_v=head_dim_v, - tile_scheduler_metadata=tile_scheduler_metadata, - num_splits=num_splits, - softmax_scale=softmax_scale, - causal=causal, - is_fp8_kvcache=is_fp8_kvcache, - indices=indices, - attn_sink=attn_sink, - extra_k_cache=extra_k_cache, - extra_indices_in_kvcache=extra_indices_in_kvcache, - topk_length=topk_length, - extra_topk_length=extra_topk_length, - ) - - -def flash_mla_sparse_fwd( - q: torch.Tensor, - kv: torch.Tensor, - indices: torch.Tensor, - sm_scale: float, - d_v: int = 512, - attn_sink: Optional[torch.Tensor] = None, - topk_length: Optional[torch.Tensor] = None, -) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor]: - return _impl.flash_mla_sparse_fwd( - q=q, - kv=kv, - indices=indices, - sm_scale=sm_scale, - d_v=d_v, - attn_sink=attn_sink, - topk_length=topk_length, - ) - - -def flash_attn_varlen_func( - q: torch.Tensor, - k: torch.Tensor, - v: torch.Tensor, - cu_seqlens_qo: torch.Tensor, - cu_seqlens_kv: torch.Tensor, - max_seqlen_qo: int, - max_seqlen_kv: int, - dropout_p: float = 0.0, - softmax_scale: Optional[float] = None, - causal: bool = False, - deterministic: bool = False, - is_varlen: bool = True, -) -> Tuple[torch.Tensor, torch.Tensor]: - return _impl.flash_attn_varlen_func( - q=q, - k=k, - v=v, - cu_seqlens_qo=cu_seqlens_qo, - cu_seqlens_kv=cu_seqlens_kv, - max_seqlen_qo=max_seqlen_qo, - max_seqlen_kv=max_seqlen_kv, - dropout_p=dropout_p, - softmax_scale=softmax_scale, - causal=causal, - deterministic=deterministic, - is_varlen=is_varlen, - ) - - -def flash_attn_varlen_qkvpacked_func( - qkv: torch.Tensor, - cu_seqlens: torch.Tensor, - max_seqlen: int, - head_dim_qk: int, - dropout_p: float = 0.0, - softmax_scale: Optional[float] = None, - causal: bool = False, - deterministic: bool = False, - is_varlen: bool = True, -) -> Tuple[torch.Tensor, torch.Tensor]: - return _impl.flash_attn_varlen_qkvpacked_func( - qkv=qkv, - cu_seqlens=cu_seqlens, - max_seqlen=max_seqlen, - head_dim_qk=head_dim_qk, - dropout_p=dropout_p, - softmax_scale=softmax_scale, - causal=causal, - deterministic=deterministic, - is_varlen=is_varlen, - ) - - -def flash_attn_varlen_kvpacked_func( - q: torch.Tensor, - kv: torch.Tensor, - cu_seqlens_qo: torch.Tensor, - cu_seqlens_kv: torch.Tensor, - max_seqlen_qo: int, - max_seqlen_kv: int, - head_dim_qk: int, - dropout_p: float = 0.0, - softmax_scale: Optional[float] = None, - causal: bool = False, - deterministic: bool = False, - is_varlen: bool = True, -) -> Tuple[torch.Tensor, torch.Tensor]: - return _impl.flash_attn_varlen_kvpacked_func( - q=q, - kv=kv, - cu_seqlens_qo=cu_seqlens_qo, - cu_seqlens_kv=cu_seqlens_kv, - max_seqlen_qo=max_seqlen_qo, - max_seqlen_kv=max_seqlen_kv, - head_dim_qk=head_dim_qk, - dropout_p=dropout_p, - softmax_scale=softmax_scale, - causal=causal, - deterministic=deterministic, - is_varlen=is_varlen, - ) - - -__all__ = [ - "__version__", - "FlashMLASchedMeta", - "get_mla_metadata", - "flash_mla_with_kvcache", - "flash_attn_varlen_func", - "flash_attn_varlen_qkvpacked_func", - "flash_attn_varlen_kvpacked_func", - "flash_mla_sparse_fwd", -] diff --git a/build/torch29-cxx11-cu130-aarch64-linux/_flash_mla_cuda_09f70ef.abi3.so b/build/torch29-cxx11-cu130-aarch64-linux/_flash_mla_cuda_09f70ef.abi3.so deleted file mode 100644 index 69ac8d6d7b2a2b86ac945f814f49c1c032296812..0000000000000000000000000000000000000000 --- a/build/torch29-cxx11-cu130-aarch64-linux/_flash_mla_cuda_09f70ef.abi3.so +++ /dev/null @@ -1,3 +0,0 @@ -version https://git-lfs.github.com/spec/v1 -oid sha256:97dd39bb9a543eef84a778853629313312dfb6af07ae5a01d9d5e7c10dc7df16 -size 9441960 diff --git a/build/torch29-cxx11-cu130-aarch64-linux/_ops.py b/build/torch29-cxx11-cu130-aarch64-linux/_ops.py deleted file mode 100644 index ea7ed02f7680582f28bdb0d1e552de1dc177f7c5..0000000000000000000000000000000000000000 --- a/build/torch29-cxx11-cu130-aarch64-linux/_ops.py +++ /dev/null @@ -1,9 +0,0 @@ -import torch -from . import _flash_mla_cuda_09f70ef -ops = torch.ops._flash_mla_cuda_09f70ef - -def add_op_namespace_prefix(op_name: str): - """ - Prefix op by namespace. - """ - return f"_flash_mla_cuda_09f70ef::{op_name}" diff --git a/build/torch29-cxx11-cu130-aarch64-linux/flash_mla/__init__.py b/build/torch29-cxx11-cu130-aarch64-linux/flash_mla/__init__.py deleted file mode 100644 index 03dbc1afe1cf156661a2b1b22003cd5f599a0309..0000000000000000000000000000000000000000 --- a/build/torch29-cxx11-cu130-aarch64-linux/flash_mla/__init__.py +++ /dev/null @@ -1,26 +0,0 @@ -import ctypes -import sys - -import importlib -from pathlib import Path -from types import ModuleType - -def _import_from_path(file_path: Path) -> ModuleType: - # We cannot use the module name as-is, after adding it to `sys.modules`, - # it would also be used for other imports. So, we make a module name that - # depends on the path for it to be unique using the hex-encoded hash of - # the path. - path_hash = "{:x}".format(ctypes.c_size_t(hash(file_path.absolute())).value) - module_name = path_hash - spec = importlib.util.spec_from_file_location(module_name, file_path) - if spec is None: - raise ImportError(f"Cannot load spec for {module_name} from {file_path}") - module = importlib.util.module_from_spec(spec) - if module is None: - raise ImportError(f"Cannot load module {module_name} from spec") - sys.modules[module_name] = module - spec.loader.exec_module(module) # type: ignore - return module - - -globals().update(vars(_import_from_path(Path(__file__).parent.parent / "__init__.py"))) diff --git a/build/torch29-cxx11-cu130-aarch64-linux/flash_mla_interface.py b/build/torch29-cxx11-cu130-aarch64-linux/flash_mla_interface.py deleted file mode 100644 index a84e448ffe741bb6d3dafaf7888ed8cc94984467..0000000000000000000000000000000000000000 --- a/build/torch29-cxx11-cu130-aarch64-linux/flash_mla_interface.py +++ /dev/null @@ -1,435 +0,0 @@ -from typing import Optional, Tuple -import dataclasses - -import torch - -from ._ops import ops as flash_mla_cuda - -@dataclasses.dataclass -class FlashMLASchedMeta: - """ - A class that stores the tile scheduler metadata of FlashMLA - """ - - @dataclasses.dataclass - class Config: - b: int - s_q: int - h_q: int - page_block_size: int - h_k: int - - causal: bool - is_fp8_kvcache: bool - topk: Optional[int] - - extra_page_block_size: Optional[int] - extra_topk: Optional[int] - - have_initialized: bool = False - - config: Optional[Config] = None - - tile_scheduler_metadata: Optional[torch.Tensor] = None # (num_sm_parts, TileSchedulerMetaDataSize), dtype torch.int32. - num_splits: Optional[torch.Tensor] = None # (1), dtype torch.int32. - - -def get_mla_metadata( - *args, - **kwargs -) -> Tuple[FlashMLASchedMeta, None]: - """ - Returns an empty instance of FlashMLASchedMeta. The actual scheduling metadata will be generated during the first invocation of flash_mla_with_kvcache. - - Arguments: - This function does not need any arguments, but we keep *args and **kwargs to be compatible with the old interface. - - Return: - A tuple. Due to historical reasons, we return a tuple of (FlashMLASchedMeta, None) now. Only the first element is useful. - """ - return FlashMLASchedMeta(), None - - -def flash_mla_with_kvcache( - q: torch.Tensor, - k_cache: torch.Tensor, - block_table: Optional[torch.Tensor], - cache_seqlens: Optional[torch.Tensor], - head_dim_v: int, - tile_scheduler_metadata: FlashMLASchedMeta, - num_splits: None = None, - softmax_scale: Optional[float] = None, - causal: bool = False, - is_fp8_kvcache: bool = False, - indices: Optional[torch.Tensor] = None, - attn_sink: Optional[torch.Tensor] = None, - extra_k_cache: Optional[torch.Tensor] = None, - extra_indices_in_kvcache: Optional[torch.Tensor] = None, - topk_length: Optional[torch.Tensor] = None, - extra_topk_length: Optional[torch.Tensor] = None -) -> Tuple[torch.Tensor, torch.Tensor]: - """ - Arguments: - q: (batch_size, seq_len_q, num_heads_q, head_dim). - k_cache: (num_blocks, page_block_size, num_heads_k, head_dim). - Different modes (including fp8/bf16, and sparsity) has different KV cache layouts. See comments below for details. - The KV cache must be contiguously valid for sparse attention on sm100. Here "contiguously valid" means that every byte, from the very beginning of the KV cache, till the last byte in the KV cache, is valid memory address to visit (i.e. won't IMA). In other words, the KV cache could be a slice of a larger array, but cannot be a list of disjoint memory blocks. - block_table: (batch_size, max_num_blocks_per_seq), torch.int32. Can be None when sparse attention is used. - cache_seqlens: (batch_size), torch.int32. Can be None when sparse attention is used. - head_dim_v: Head_dim of v. Must be 512 - sched_meta: FlashMLASchedMeta, return by get_mla_metadata. You may reuse the same sched_meta across different invocations, but only when the tensor shapes and the values of cache_seqlens, topk_length, and extra_topk_length remain the same. - num_splits_placeholder: must be "None" (to be compatible with the old interface). - softmax_scale: float. The scaling of QK^T before applying softmax. Default to 1 / sqrt(head_dim_k). - causal: bool. Whether to apply causal attention mask. Only valid for dense attention - is_fp8_kvcache: bool. - indices: (batch_size, seq_len_q, topk). KV indices when sparse attention is enabled. - Pay attention that indices_in_kvcache[i][j][k] = (the index of the page block where token t resides) * block_size + (the offset of token t among the page block), - where t is the k-th token of the j-th q-sequence in the i-th batch. - attn_sink: Optional[torch.Tensor], (num_heads_q, ), torch.float32. If presented, the final output will be scaled by exp(lse) / (exp(lse) + exp(attn_sink)). Have no affect on the returned softmax_lse. +inf will cause the result to become 0. - extra_k_cache and extra_indices_in_kvcache: If provided, will attend to these extra tokens in addition to those in k_cache and indices_in_kvcache. Their format requirements are the same as k_cache and indices_in_kvcache respectively. - topk_length/extra_topk_length: (batch_size, ), torch.int32. If provided, only the leftmost topk_length indices will be processed. Useful when the actual topk for different queries are different so that we can save some computation, compared to masking. - - For DeepSeek V3, DeepSeek V3.1, and DeepSeek V3.2: - head_dim should be 576 while head_dim_v should be 512. - In FP8+sparse mode, each token's KV cache is 656 Bytes, structured as: - - The shape of the tensor `k_cache` is (num_blocks, page_block_size, num_heads_k, head_dim), and num_heads_k must be 1. - - First 512 bytes: The "quantized NoPE" part, containing 512 float8_e4m3 values. - - Next 16 bytes: Scale factors, containing 4 float32 values. The first float32 is the scale for the first 128 float8_e4m3 values, the second for the next 128, and so on. - - Last 128 bytes: The "RoPE" part, containing 64 bfloat16 values. This part is not quantized for accuracy. - - Return: - out: (batch_size, seq_len_q, num_heads_q, head_dim_v). - softmax_lse: (batch_size, num_heads_q, seq_len_q), torch.float32. - """ - sched_meta = tile_scheduler_metadata - indices_in_kvcache = indices - assert isinstance(sched_meta, FlashMLASchedMeta), "tile_scheduler_metadata must be of type FlashMLASchedMeta" - assert num_splits is None, "num_splits must be None" - - topk = indices_in_kvcache.shape[-1] if indices_in_kvcache is not None else None - extra_k_page_block_size = extra_k_cache.shape[1] if extra_k_cache is not None else None - extra_topk = extra_indices_in_kvcache.shape[-1] if extra_indices_in_kvcache is not None else None - if softmax_scale is None: - softmax_scale = q.shape[-1] ** (-0.5) - - if not sched_meta.have_initialized: - # Sanity check. We only perform sanity check during the first invocation to save CPU time. - if indices_in_kvcache is not None: - assert not causal, "causal must be False when indices_in_kvcache is not None (i.e. sparse attention is enabled)" - - # Initialize the tile scheduler metadata during the first invocation. - sched_meta.have_initialized = True - sched_meta.config = FlashMLASchedMeta.Config( - q.shape[0], - q.shape[1], - q.shape[2], - k_cache.shape[1], - k_cache.shape[2], - - causal, - is_fp8_kvcache, - topk, - - extra_k_page_block_size, - extra_topk, - ) - else: - # Check whether the input arguments are consistent with sched_meta - helper_msg = " Your input arguments are inconsistent with sched_meta. Please make sure the input arguments are consistent across different invocations of flash_mla_with_kvcache on the same sched_meta." - assert sched_meta.config is not None - assert sched_meta.config.b == q.shape[0], "sched_meta.config.b must be equal to batch_size." + helper_msg - assert sched_meta.config.s_q == q.shape[1], "sched_meta.config.s_q must be equal to seq_len_q." + helper_msg - assert sched_meta.config.h_q == q.shape[2], "sched_meta.config.h_q must be equal to num_heads_q." + helper_msg - assert sched_meta.config.page_block_size == k_cache.shape[1], "sched_meta.config.page_block_size must be equal to page_block_size." + helper_msg - assert sched_meta.config.h_k == k_cache.shape[2], "sched_meta.config.h_k must be equal to num_heads_k." + helper_msg - assert sched_meta.config.causal == causal, "sched_meta.config.causal must be equal to causal." + helper_msg - assert sched_meta.config.is_fp8_kvcache == is_fp8_kvcache, "sched_meta.config.is_fp8_kvcache must be equal to is_fp8_kvcache." + helper_msg - assert sched_meta.config.topk == topk, "sched_meta.config.topk must be equal to the last dim of indices_in_kvcache." + helper_msg - assert sched_meta.config.extra_page_block_size == extra_k_page_block_size, "sched_meta.config.extra_page_block_size must be equal to the page_block_size of extra_k_cache." + helper_msg - assert sched_meta.config.extra_topk == extra_topk, "sched_meta.config.extra_topk must be equal to the last dim of extra_indices_in_kvcache." + helper_msg - - if topk is not None: - # Sparse attention - assert not causal, "causal must be False when sparse attention is enabled" - assert is_fp8_kvcache, "is_fp8_kvcache must be True when sparse attention is enabled" - out, lse, new_tile_scheduler_metadata, new_num_splits = flash_mla_cuda.sparse_decode_fwd( - q, k_cache, indices_in_kvcache, topk_length, attn_sink, - sched_meta.tile_scheduler_metadata, sched_meta.num_splits, - extra_k_cache, extra_indices_in_kvcache, extra_topk_length, - head_dim_v, softmax_scale - ) - else: - # Dense attention - assert indices_in_kvcache is None and attn_sink is None and extra_k_cache is None and extra_indices_in_kvcache is None and topk_length is None and extra_topk_length is None, "indices_in_kvcache, attn_sink, extra_k_cache, extra_indices_in_kvcache, topk_length and extra_topk_length must be None when dense attention is used." - assert block_table is not None and cache_seqlens is not None, "block_table and cache_seqlens must be provided when dense attention is used." - out, lse, new_tile_scheduler_metadata, new_num_splits = flash_mla_cuda.dense_decode_fwd( - q, k_cache, head_dim_v, - cache_seqlens, block_table, - softmax_scale, causal, - sched_meta.tile_scheduler_metadata, sched_meta.num_splits - ) - sched_meta.tile_scheduler_metadata = new_tile_scheduler_metadata - sched_meta.num_splits = new_num_splits - return (out, lse) - - -def flash_mla_sparse_fwd( - q: torch.Tensor, - kv: torch.Tensor, - indices: torch.Tensor, - sm_scale: float, - d_v: int = 512, - attn_sink: Optional[torch.Tensor] = None, - topk_length: Optional[torch.Tensor] = None, -) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor]: - """ - Sparse attention prefill kernel - - Args: - q: [s_q, h_q, d_qk], bfloat16 - kv: [s_kv, h_kv, d_qk], bfloat16 - indices: [s_q, h_kv, topk], int32. Invalid indices should be set to -1 or numbers >= s_kv - sm_scale: float - d_v: The dimension of value vectors. Can only be 512 - attn_sink: optional, [h_q], float32. - If attn_sink is provided, when computing output, output will be additionally multiplied by exp(lse) / (exp(lse) + exp(attn_sink)). - +-inf in attn_sink will be handled normally (i.e., -inf has no effect, +inf will make corresponding output all zeros). - This argument has no effect on lse and max_logits. - topk_length: optional, [s_q], int32. If provided, the i-th q token will only attend to k tokens specified by indices[i, :, :topk_length[i]], ignoring later k/v tokens (even if provided in indices). - In extremely rare cases (topk_length provided, there is a valid topk index between topk_length[i] ~ s_kv, and that topk index points to a k token containing NaN), operator output will contain NaN, so please avoid this situation. - - Returns: - (output, max_logits, lse) - Please refer to tests/ref.py for the precise definitions of these parameters. - - output: [s_q, h_q, d_v], bfloat16 - - max_logits: [s_q, h_q], float - - lse: [s_q, h_q], float, log-sum-exp of attention scores - """ - results = flash_mla_cuda.sparse_prefill_fwd( - q, kv, indices, sm_scale, d_v, attn_sink, topk_length - ) - return results - - -def _flash_attn_varlen_forward( - q: torch.Tensor, - k: torch.Tensor, - v: torch.Tensor, - cu_seqlens_qo: torch.Tensor, - cu_seqlens_kv: torch.Tensor, - max_seqlen_qo: int, - max_seqlen_kv: int, - out: Optional[torch.Tensor] = None, - lse: Optional[torch.Tensor] = None, - causal: bool = False, - softmax_scale: Optional[float] = None, - is_varlen: bool = True, -) -> Tuple[torch.Tensor, torch.Tensor]: - qo_total_len, num_qo_heads, head_dim_qk = q.shape - kv_total_len, num_kv_heads, head_dim_vo = v.shape - - mask_mode_code = 1 if causal else 0 - if softmax_scale is None: - softmax_scale = head_dim_qk ** (-0.5) - - if out is None: - out = torch.empty(qo_total_len, num_qo_heads, head_dim_vo, device=q.device, dtype=q.dtype) - if lse is None: - # Make lse contiguous on seqlen dim - lse = torch.empty(num_qo_heads, qo_total_len, device=q.device, dtype=torch.float32).T - - workspace_buffer = torch.empty(32 * 1024 * 1024, dtype=torch.uint8, device=q.device) - flash_mla_cuda.dense_prefill_fwd( - workspace_buffer, - q, - k, - v, - cu_seqlens_qo, - cu_seqlens_kv, - out, - lse, - mask_mode_code, - softmax_scale, - max_seqlen_qo, - max_seqlen_kv, - is_varlen, - ) - - return out, lse - - -def _flash_attn_varlen_backward( - do: torch.Tensor, - q: torch.Tensor, - k: torch.Tensor, - v: torch.Tensor, - out: torch.Tensor, - lse: torch.Tensor, - cu_seqlens_qo: torch.Tensor, - cu_seqlens_kv: torch.Tensor, - max_seqlen_qo: int, - max_seqlen_kv: int, - dq: Optional[torch.Tensor] = None, - dk: Optional[torch.Tensor] = None, - dv: Optional[torch.Tensor] = None, - causal: bool = False, - softmax_scale: Optional[float] = None, - is_varlen: bool = True, -) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor]: - qo_total_len, num_qo_heads, head_dim_qk = q.shape - kv_total_len, num_kv_heads, head_dim_vo = v.shape - - # TODO: fix bwd GQA - if num_qo_heads != num_kv_heads: - raise ValueError(f"SM100 bwd doesn't support GQA now. num_qo_heads: {num_qo_heads}, num_kv_heads: {num_kv_heads}.") - - mask_mode_code = 1 if causal else 0 - if softmax_scale is None: - softmax_scale = head_dim_qk ** (-0.5) - - if dq is None: - dq = torch.empty(qo_total_len, num_qo_heads, head_dim_qk, device=q.device, dtype=q.dtype) - if dk is None: - dk = torch.empty(kv_total_len, num_kv_heads, head_dim_qk, device=q.device, dtype=q.dtype) - if dv is None: - dv = torch.empty(kv_total_len, num_kv_heads, head_dim_vo, device=q.device, dtype=q.dtype) - - max_seqlen_qo_aligned = (max_seqlen_qo + 7) // 8 * 8 - bs = cu_seqlens_qo.shape[0] - 1 - workspace_bytes = 0 - workspace_bytes += 4 * bs * max_seqlen_qo_aligned * num_qo_heads * head_dim_qk # dQ_acc - workspace_bytes += 4 * max_seqlen_qo_aligned * bs * num_qo_heads * 2 # sum_OdO and scaled_lse - if num_qo_heads != num_kv_heads: - workspace_bytes += 2 * kv_total_len * num_qo_heads * (head_dim_qk + head_dim_vo) # dKV_acc - workspace_buffer = torch.empty(workspace_bytes, dtype=torch.uint8, device=q.device) - flash_mla_cuda.dense_prefill_bwd( - workspace_buffer, - do, - q, - k, - v, - out, - lse, - cu_seqlens_qo, - cu_seqlens_kv, - dq, - dk, - dv, - mask_mode_code, - softmax_scale, - max_seqlen_qo, - max_seqlen_kv, - is_varlen, - ) - - return dq, dk, dv - - -class FlashAttnVarlenFunc(torch.autograd.Function): - def forward( - ctx, - q: torch.Tensor, - k: torch.Tensor, - v: torch.Tensor, - cu_seqlens_qo: torch.Tensor, - cu_seqlens_kv: torch.Tensor, - max_seqlen_qo: int, - max_seqlen_kv: int, - causal: bool = False, - softmax_scale: Optional[float] = None, - is_varlen: bool = True, - ) -> Tuple[torch.Tensor, torch.Tensor]: - out, lse = _flash_attn_varlen_forward( - q, k, v, - cu_seqlens_qo, cu_seqlens_kv, max_seqlen_qo, max_seqlen_kv, - causal=causal, softmax_scale=softmax_scale, - is_varlen=is_varlen, - ) - ctx.save_for_backward(q, k, v, out, lse, cu_seqlens_qo, cu_seqlens_kv) - ctx.max_seqlen_qo = max_seqlen_qo - ctx.max_seqlen_kv = max_seqlen_kv - ctx.causal = causal - ctx.softmax_scale = softmax_scale - ctx.is_varlen = is_varlen - return out, lse - - def backward( - ctx, - do: torch.Tensor, - dlse: torch.Tensor, - ): - del dlse # LSE doesn't support backward currently - q, k, v, out, lse, cu_seqlens_qo, cu_seqlens_kv = ctx.saved_tensors - dq, dk, dv = _flash_attn_varlen_backward( - do, q, k, v, out, lse, - cu_seqlens_qo, cu_seqlens_kv, ctx.max_seqlen_qo, ctx.max_seqlen_kv, - causal=ctx.causal, softmax_scale=ctx.softmax_scale, - is_varlen=ctx.is_varlen, - ) - return dq, dk, dv, None, None, None, None, None, None, None - - -def flash_attn_varlen_func( - q: torch.Tensor, - k: torch.Tensor, - v: torch.Tensor, - cu_seqlens_qo: torch.Tensor, - cu_seqlens_kv: torch.Tensor, - max_seqlen_qo: int, - max_seqlen_kv: int, - dropout_p: float = 0.0, - softmax_scale: Optional[float] = None, - causal: bool = False, - deterministic: bool = False, - is_varlen: bool = True, -) -> Tuple[torch.Tensor, torch.Tensor]: - assert dropout_p == 0.0 - assert not deterministic - return FlashAttnVarlenFunc.apply( - q, k, v, - cu_seqlens_qo, cu_seqlens_kv, max_seqlen_qo, max_seqlen_kv, - causal, softmax_scale, is_varlen, - ) - - -def flash_attn_varlen_qkvpacked_func( - qkv: torch.Tensor, - cu_seqlens: torch.Tensor, - max_seqlen: int, - head_dim_qk: int, - dropout_p: float = 0.0, - softmax_scale: Optional[float] = None, - causal: bool = False, - deterministic: bool = False, - is_varlen: bool = True, -) -> Tuple[torch.Tensor, torch.Tensor]: - assert dropout_p == 0.0 - assert not deterministic - return FlashAttnVarlenFunc.apply( - qkv[:, :, :head_dim_qk], qkv[:, :, head_dim_qk:head_dim_qk * 2], qkv[:, :, head_dim_qk * 2:], - cu_seqlens, cu_seqlens, max_seqlen, max_seqlen, - causal, softmax_scale, is_varlen, - ) - - -def flash_attn_varlen_kvpacked_func( - q: torch.Tensor, - kv: torch.Tensor, - cu_seqlens_qo: torch.Tensor, - cu_seqlens_kv: torch.Tensor, - max_seqlen_qo: int, - max_seqlen_kv: int, - head_dim_qk: int, - dropout_p: float = 0.0, - softmax_scale: Optional[float] = None, - causal: bool = False, - deterministic: bool = False, - is_varlen: bool = True, -) -> Tuple[torch.Tensor, torch.Tensor]: - assert dropout_p == 0.0 - assert not deterministic - return FlashAttnVarlenFunc.apply( - q, kv[:, :, :head_dim_qk], kv[:, :, head_dim_qk:], - cu_seqlens_qo, cu_seqlens_kv, max_seqlen_qo, max_seqlen_kv, - causal, softmax_scale, is_varlen, - ) diff --git a/build/torch29-cxx11-cu130-aarch64-linux/metadata.json b/build/torch29-cxx11-cu130-aarch64-linux/metadata.json deleted file mode 100644 index 8190d75efa8fd6449ddcd73de2072f17086e0842..0000000000000000000000000000000000000000 --- a/build/torch29-cxx11-cu130-aarch64-linux/metadata.json +++ /dev/null @@ -1,12 +0,0 @@ -{ - "version": 1, - "license": "MIT", - "python-depends": [], - "backend": { - "type": "cuda", - "archs": [ - "10.0f", - "9.0a" - ] - } -} diff --git a/build/torch29-cxx11-cu130-x86_64-linux/__init__.py b/build/torch29-cxx11-cu130-x86_64-linux/__init__.py deleted file mode 100644 index db300fe9b95176a20b27b3641d89be657d0c4319..0000000000000000000000000000000000000000 --- a/build/torch29-cxx11-cu130-x86_64-linux/__init__.py +++ /dev/null @@ -1,163 +0,0 @@ -from typing import Optional, Tuple -import torch - -from .flash_mla_interface import FlashMLASchedMeta -from . import flash_mla_interface as _impl - - -def get_mla_metadata(*args, **kwargs) -> Tuple[FlashMLASchedMeta, None]: - return _impl.get_mla_metadata(*args, **kwargs) - - -def flash_mla_with_kvcache( - q: torch.Tensor, - k_cache: torch.Tensor, - block_table: Optional[torch.Tensor], - cache_seqlens: Optional[torch.Tensor], - head_dim_v: int, - tile_scheduler_metadata: FlashMLASchedMeta, - num_splits: None = None, - softmax_scale: Optional[float] = None, - causal: bool = False, - is_fp8_kvcache: bool = False, - indices: Optional[torch.Tensor] = None, - attn_sink: Optional[torch.Tensor] = None, - extra_k_cache: Optional[torch.Tensor] = None, - extra_indices_in_kvcache: Optional[torch.Tensor] = None, - topk_length: Optional[torch.Tensor] = None, - extra_topk_length: Optional[torch.Tensor] = None, -) -> Tuple[torch.Tensor, torch.Tensor]: - return _impl.flash_mla_with_kvcache( - q=q, - k_cache=k_cache, - block_table=block_table, - cache_seqlens=cache_seqlens, - head_dim_v=head_dim_v, - tile_scheduler_metadata=tile_scheduler_metadata, - num_splits=num_splits, - softmax_scale=softmax_scale, - causal=causal, - is_fp8_kvcache=is_fp8_kvcache, - indices=indices, - attn_sink=attn_sink, - extra_k_cache=extra_k_cache, - extra_indices_in_kvcache=extra_indices_in_kvcache, - topk_length=topk_length, - extra_topk_length=extra_topk_length, - ) - - -def flash_mla_sparse_fwd( - q: torch.Tensor, - kv: torch.Tensor, - indices: torch.Tensor, - sm_scale: float, - d_v: int = 512, - attn_sink: Optional[torch.Tensor] = None, - topk_length: Optional[torch.Tensor] = None, -) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor]: - return _impl.flash_mla_sparse_fwd( - q=q, - kv=kv, - indices=indices, - sm_scale=sm_scale, - d_v=d_v, - attn_sink=attn_sink, - topk_length=topk_length, - ) - - -def flash_attn_varlen_func( - q: torch.Tensor, - k: torch.Tensor, - v: torch.Tensor, - cu_seqlens_qo: torch.Tensor, - cu_seqlens_kv: torch.Tensor, - max_seqlen_qo: int, - max_seqlen_kv: int, - dropout_p: float = 0.0, - softmax_scale: Optional[float] = None, - causal: bool = False, - deterministic: bool = False, - is_varlen: bool = True, -) -> Tuple[torch.Tensor, torch.Tensor]: - return _impl.flash_attn_varlen_func( - q=q, - k=k, - v=v, - cu_seqlens_qo=cu_seqlens_qo, - cu_seqlens_kv=cu_seqlens_kv, - max_seqlen_qo=max_seqlen_qo, - max_seqlen_kv=max_seqlen_kv, - dropout_p=dropout_p, - softmax_scale=softmax_scale, - causal=causal, - deterministic=deterministic, - is_varlen=is_varlen, - ) - - -def flash_attn_varlen_qkvpacked_func( - qkv: torch.Tensor, - cu_seqlens: torch.Tensor, - max_seqlen: int, - head_dim_qk: int, - dropout_p: float = 0.0, - softmax_scale: Optional[float] = None, - causal: bool = False, - deterministic: bool = False, - is_varlen: bool = True, -) -> Tuple[torch.Tensor, torch.Tensor]: - return _impl.flash_attn_varlen_qkvpacked_func( - qkv=qkv, - cu_seqlens=cu_seqlens, - max_seqlen=max_seqlen, - head_dim_qk=head_dim_qk, - dropout_p=dropout_p, - softmax_scale=softmax_scale, - causal=causal, - deterministic=deterministic, - is_varlen=is_varlen, - ) - - -def flash_attn_varlen_kvpacked_func( - q: torch.Tensor, - kv: torch.Tensor, - cu_seqlens_qo: torch.Tensor, - cu_seqlens_kv: torch.Tensor, - max_seqlen_qo: int, - max_seqlen_kv: int, - head_dim_qk: int, - dropout_p: float = 0.0, - softmax_scale: Optional[float] = None, - causal: bool = False, - deterministic: bool = False, - is_varlen: bool = True, -) -> Tuple[torch.Tensor, torch.Tensor]: - return _impl.flash_attn_varlen_kvpacked_func( - q=q, - kv=kv, - cu_seqlens_qo=cu_seqlens_qo, - cu_seqlens_kv=cu_seqlens_kv, - max_seqlen_qo=max_seqlen_qo, - max_seqlen_kv=max_seqlen_kv, - head_dim_qk=head_dim_qk, - dropout_p=dropout_p, - softmax_scale=softmax_scale, - causal=causal, - deterministic=deterministic, - is_varlen=is_varlen, - ) - - -__all__ = [ - "__version__", - "FlashMLASchedMeta", - "get_mla_metadata", - "flash_mla_with_kvcache", - "flash_attn_varlen_func", - "flash_attn_varlen_qkvpacked_func", - "flash_attn_varlen_kvpacked_func", - "flash_mla_sparse_fwd", -] diff --git a/build/torch29-cxx11-cu130-x86_64-linux/_flash_mla_cuda_09f70ef.abi3.so b/build/torch29-cxx11-cu130-x86_64-linux/_flash_mla_cuda_09f70ef.abi3.so deleted file mode 100644 index c5f00086319d5168c2913aad7489d8d83204b054..0000000000000000000000000000000000000000 --- a/build/torch29-cxx11-cu130-x86_64-linux/_flash_mla_cuda_09f70ef.abi3.so +++ /dev/null @@ -1,3 +0,0 @@ -version https://git-lfs.github.com/spec/v1 -oid sha256:243747242a38da9e2034590bfd52c3f883683b4faec3ccfd6ec76a0f97addf43 -size 9380384 diff --git a/build/torch29-cxx11-cu130-x86_64-linux/_ops.py b/build/torch29-cxx11-cu130-x86_64-linux/_ops.py deleted file mode 100644 index ea7ed02f7680582f28bdb0d1e552de1dc177f7c5..0000000000000000000000000000000000000000 --- a/build/torch29-cxx11-cu130-x86_64-linux/_ops.py +++ /dev/null @@ -1,9 +0,0 @@ -import torch -from . import _flash_mla_cuda_09f70ef -ops = torch.ops._flash_mla_cuda_09f70ef - -def add_op_namespace_prefix(op_name: str): - """ - Prefix op by namespace. - """ - return f"_flash_mla_cuda_09f70ef::{op_name}" diff --git a/build/torch29-cxx11-cu130-x86_64-linux/flash_mla/__init__.py b/build/torch29-cxx11-cu130-x86_64-linux/flash_mla/__init__.py deleted file mode 100644 index 03dbc1afe1cf156661a2b1b22003cd5f599a0309..0000000000000000000000000000000000000000 --- a/build/torch29-cxx11-cu130-x86_64-linux/flash_mla/__init__.py +++ /dev/null @@ -1,26 +0,0 @@ -import ctypes -import sys - -import importlib -from pathlib import Path -from types import ModuleType - -def _import_from_path(file_path: Path) -> ModuleType: - # We cannot use the module name as-is, after adding it to `sys.modules`, - # it would also be used for other imports. So, we make a module name that - # depends on the path for it to be unique using the hex-encoded hash of - # the path. - path_hash = "{:x}".format(ctypes.c_size_t(hash(file_path.absolute())).value) - module_name = path_hash - spec = importlib.util.spec_from_file_location(module_name, file_path) - if spec is None: - raise ImportError(f"Cannot load spec for {module_name} from {file_path}") - module = importlib.util.module_from_spec(spec) - if module is None: - raise ImportError(f"Cannot load module {module_name} from spec") - sys.modules[module_name] = module - spec.loader.exec_module(module) # type: ignore - return module - - -globals().update(vars(_import_from_path(Path(__file__).parent.parent / "__init__.py"))) diff --git a/build/torch29-cxx11-cu130-x86_64-linux/flash_mla_interface.py b/build/torch29-cxx11-cu130-x86_64-linux/flash_mla_interface.py deleted file mode 100644 index a84e448ffe741bb6d3dafaf7888ed8cc94984467..0000000000000000000000000000000000000000 --- a/build/torch29-cxx11-cu130-x86_64-linux/flash_mla_interface.py +++ /dev/null @@ -1,435 +0,0 @@ -from typing import Optional, Tuple -import dataclasses - -import torch - -from ._ops import ops as flash_mla_cuda - -@dataclasses.dataclass -class FlashMLASchedMeta: - """ - A class that stores the tile scheduler metadata of FlashMLA - """ - - @dataclasses.dataclass - class Config: - b: int - s_q: int - h_q: int - page_block_size: int - h_k: int - - causal: bool - is_fp8_kvcache: bool - topk: Optional[int] - - extra_page_block_size: Optional[int] - extra_topk: Optional[int] - - have_initialized: bool = False - - config: Optional[Config] = None - - tile_scheduler_metadata: Optional[torch.Tensor] = None # (num_sm_parts, TileSchedulerMetaDataSize), dtype torch.int32. - num_splits: Optional[torch.Tensor] = None # (1), dtype torch.int32. - - -def get_mla_metadata( - *args, - **kwargs -) -> Tuple[FlashMLASchedMeta, None]: - """ - Returns an empty instance of FlashMLASchedMeta. The actual scheduling metadata will be generated during the first invocation of flash_mla_with_kvcache. - - Arguments: - This function does not need any arguments, but we keep *args and **kwargs to be compatible with the old interface. - - Return: - A tuple. Due to historical reasons, we return a tuple of (FlashMLASchedMeta, None) now. Only the first element is useful. - """ - return FlashMLASchedMeta(), None - - -def flash_mla_with_kvcache( - q: torch.Tensor, - k_cache: torch.Tensor, - block_table: Optional[torch.Tensor], - cache_seqlens: Optional[torch.Tensor], - head_dim_v: int, - tile_scheduler_metadata: FlashMLASchedMeta, - num_splits: None = None, - softmax_scale: Optional[float] = None, - causal: bool = False, - is_fp8_kvcache: bool = False, - indices: Optional[torch.Tensor] = None, - attn_sink: Optional[torch.Tensor] = None, - extra_k_cache: Optional[torch.Tensor] = None, - extra_indices_in_kvcache: Optional[torch.Tensor] = None, - topk_length: Optional[torch.Tensor] = None, - extra_topk_length: Optional[torch.Tensor] = None -) -> Tuple[torch.Tensor, torch.Tensor]: - """ - Arguments: - q: (batch_size, seq_len_q, num_heads_q, head_dim). - k_cache: (num_blocks, page_block_size, num_heads_k, head_dim). - Different modes (including fp8/bf16, and sparsity) has different KV cache layouts. See comments below for details. - The KV cache must be contiguously valid for sparse attention on sm100. Here "contiguously valid" means that every byte, from the very beginning of the KV cache, till the last byte in the KV cache, is valid memory address to visit (i.e. won't IMA). In other words, the KV cache could be a slice of a larger array, but cannot be a list of disjoint memory blocks. - block_table: (batch_size, max_num_blocks_per_seq), torch.int32. Can be None when sparse attention is used. - cache_seqlens: (batch_size), torch.int32. Can be None when sparse attention is used. - head_dim_v: Head_dim of v. Must be 512 - sched_meta: FlashMLASchedMeta, return by get_mla_metadata. You may reuse the same sched_meta across different invocations, but only when the tensor shapes and the values of cache_seqlens, topk_length, and extra_topk_length remain the same. - num_splits_placeholder: must be "None" (to be compatible with the old interface). - softmax_scale: float. The scaling of QK^T before applying softmax. Default to 1 / sqrt(head_dim_k). - causal: bool. Whether to apply causal attention mask. Only valid for dense attention - is_fp8_kvcache: bool. - indices: (batch_size, seq_len_q, topk). KV indices when sparse attention is enabled. - Pay attention that indices_in_kvcache[i][j][k] = (the index of the page block where token t resides) * block_size + (the offset of token t among the page block), - where t is the k-th token of the j-th q-sequence in the i-th batch. - attn_sink: Optional[torch.Tensor], (num_heads_q, ), torch.float32. If presented, the final output will be scaled by exp(lse) / (exp(lse) + exp(attn_sink)). Have no affect on the returned softmax_lse. +inf will cause the result to become 0. - extra_k_cache and extra_indices_in_kvcache: If provided, will attend to these extra tokens in addition to those in k_cache and indices_in_kvcache. Their format requirements are the same as k_cache and indices_in_kvcache respectively. - topk_length/extra_topk_length: (batch_size, ), torch.int32. If provided, only the leftmost topk_length indices will be processed. Useful when the actual topk for different queries are different so that we can save some computation, compared to masking. - - For DeepSeek V3, DeepSeek V3.1, and DeepSeek V3.2: - head_dim should be 576 while head_dim_v should be 512. - In FP8+sparse mode, each token's KV cache is 656 Bytes, structured as: - - The shape of the tensor `k_cache` is (num_blocks, page_block_size, num_heads_k, head_dim), and num_heads_k must be 1. - - First 512 bytes: The "quantized NoPE" part, containing 512 float8_e4m3 values. - - Next 16 bytes: Scale factors, containing 4 float32 values. The first float32 is the scale for the first 128 float8_e4m3 values, the second for the next 128, and so on. - - Last 128 bytes: The "RoPE" part, containing 64 bfloat16 values. This part is not quantized for accuracy. - - Return: - out: (batch_size, seq_len_q, num_heads_q, head_dim_v). - softmax_lse: (batch_size, num_heads_q, seq_len_q), torch.float32. - """ - sched_meta = tile_scheduler_metadata - indices_in_kvcache = indices - assert isinstance(sched_meta, FlashMLASchedMeta), "tile_scheduler_metadata must be of type FlashMLASchedMeta" - assert num_splits is None, "num_splits must be None" - - topk = indices_in_kvcache.shape[-1] if indices_in_kvcache is not None else None - extra_k_page_block_size = extra_k_cache.shape[1] if extra_k_cache is not None else None - extra_topk = extra_indices_in_kvcache.shape[-1] if extra_indices_in_kvcache is not None else None - if softmax_scale is None: - softmax_scale = q.shape[-1] ** (-0.5) - - if not sched_meta.have_initialized: - # Sanity check. We only perform sanity check during the first invocation to save CPU time. - if indices_in_kvcache is not None: - assert not causal, "causal must be False when indices_in_kvcache is not None (i.e. sparse attention is enabled)" - - # Initialize the tile scheduler metadata during the first invocation. - sched_meta.have_initialized = True - sched_meta.config = FlashMLASchedMeta.Config( - q.shape[0], - q.shape[1], - q.shape[2], - k_cache.shape[1], - k_cache.shape[2], - - causal, - is_fp8_kvcache, - topk, - - extra_k_page_block_size, - extra_topk, - ) - else: - # Check whether the input arguments are consistent with sched_meta - helper_msg = " Your input arguments are inconsistent with sched_meta. Please make sure the input arguments are consistent across different invocations of flash_mla_with_kvcache on the same sched_meta." - assert sched_meta.config is not None - assert sched_meta.config.b == q.shape[0], "sched_meta.config.b must be equal to batch_size." + helper_msg - assert sched_meta.config.s_q == q.shape[1], "sched_meta.config.s_q must be equal to seq_len_q." + helper_msg - assert sched_meta.config.h_q == q.shape[2], "sched_meta.config.h_q must be equal to num_heads_q." + helper_msg - assert sched_meta.config.page_block_size == k_cache.shape[1], "sched_meta.config.page_block_size must be equal to page_block_size." + helper_msg - assert sched_meta.config.h_k == k_cache.shape[2], "sched_meta.config.h_k must be equal to num_heads_k." + helper_msg - assert sched_meta.config.causal == causal, "sched_meta.config.causal must be equal to causal." + helper_msg - assert sched_meta.config.is_fp8_kvcache == is_fp8_kvcache, "sched_meta.config.is_fp8_kvcache must be equal to is_fp8_kvcache." + helper_msg - assert sched_meta.config.topk == topk, "sched_meta.config.topk must be equal to the last dim of indices_in_kvcache." + helper_msg - assert sched_meta.config.extra_page_block_size == extra_k_page_block_size, "sched_meta.config.extra_page_block_size must be equal to the page_block_size of extra_k_cache." + helper_msg - assert sched_meta.config.extra_topk == extra_topk, "sched_meta.config.extra_topk must be equal to the last dim of extra_indices_in_kvcache." + helper_msg - - if topk is not None: - # Sparse attention - assert not causal, "causal must be False when sparse attention is enabled" - assert is_fp8_kvcache, "is_fp8_kvcache must be True when sparse attention is enabled" - out, lse, new_tile_scheduler_metadata, new_num_splits = flash_mla_cuda.sparse_decode_fwd( - q, k_cache, indices_in_kvcache, topk_length, attn_sink, - sched_meta.tile_scheduler_metadata, sched_meta.num_splits, - extra_k_cache, extra_indices_in_kvcache, extra_topk_length, - head_dim_v, softmax_scale - ) - else: - # Dense attention - assert indices_in_kvcache is None and attn_sink is None and extra_k_cache is None and extra_indices_in_kvcache is None and topk_length is None and extra_topk_length is None, "indices_in_kvcache, attn_sink, extra_k_cache, extra_indices_in_kvcache, topk_length and extra_topk_length must be None when dense attention is used." - assert block_table is not None and cache_seqlens is not None, "block_table and cache_seqlens must be provided when dense attention is used." - out, lse, new_tile_scheduler_metadata, new_num_splits = flash_mla_cuda.dense_decode_fwd( - q, k_cache, head_dim_v, - cache_seqlens, block_table, - softmax_scale, causal, - sched_meta.tile_scheduler_metadata, sched_meta.num_splits - ) - sched_meta.tile_scheduler_metadata = new_tile_scheduler_metadata - sched_meta.num_splits = new_num_splits - return (out, lse) - - -def flash_mla_sparse_fwd( - q: torch.Tensor, - kv: torch.Tensor, - indices: torch.Tensor, - sm_scale: float, - d_v: int = 512, - attn_sink: Optional[torch.Tensor] = None, - topk_length: Optional[torch.Tensor] = None, -) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor]: - """ - Sparse attention prefill kernel - - Args: - q: [s_q, h_q, d_qk], bfloat16 - kv: [s_kv, h_kv, d_qk], bfloat16 - indices: [s_q, h_kv, topk], int32. Invalid indices should be set to -1 or numbers >= s_kv - sm_scale: float - d_v: The dimension of value vectors. Can only be 512 - attn_sink: optional, [h_q], float32. - If attn_sink is provided, when computing output, output will be additionally multiplied by exp(lse) / (exp(lse) + exp(attn_sink)). - +-inf in attn_sink will be handled normally (i.e., -inf has no effect, +inf will make corresponding output all zeros). - This argument has no effect on lse and max_logits. - topk_length: optional, [s_q], int32. If provided, the i-th q token will only attend to k tokens specified by indices[i, :, :topk_length[i]], ignoring later k/v tokens (even if provided in indices). - In extremely rare cases (topk_length provided, there is a valid topk index between topk_length[i] ~ s_kv, and that topk index points to a k token containing NaN), operator output will contain NaN, so please avoid this situation. - - Returns: - (output, max_logits, lse) - Please refer to tests/ref.py for the precise definitions of these parameters. - - output: [s_q, h_q, d_v], bfloat16 - - max_logits: [s_q, h_q], float - - lse: [s_q, h_q], float, log-sum-exp of attention scores - """ - results = flash_mla_cuda.sparse_prefill_fwd( - q, kv, indices, sm_scale, d_v, attn_sink, topk_length - ) - return results - - -def _flash_attn_varlen_forward( - q: torch.Tensor, - k: torch.Tensor, - v: torch.Tensor, - cu_seqlens_qo: torch.Tensor, - cu_seqlens_kv: torch.Tensor, - max_seqlen_qo: int, - max_seqlen_kv: int, - out: Optional[torch.Tensor] = None, - lse: Optional[torch.Tensor] = None, - causal: bool = False, - softmax_scale: Optional[float] = None, - is_varlen: bool = True, -) -> Tuple[torch.Tensor, torch.Tensor]: - qo_total_len, num_qo_heads, head_dim_qk = q.shape - kv_total_len, num_kv_heads, head_dim_vo = v.shape - - mask_mode_code = 1 if causal else 0 - if softmax_scale is None: - softmax_scale = head_dim_qk ** (-0.5) - - if out is None: - out = torch.empty(qo_total_len, num_qo_heads, head_dim_vo, device=q.device, dtype=q.dtype) - if lse is None: - # Make lse contiguous on seqlen dim - lse = torch.empty(num_qo_heads, qo_total_len, device=q.device, dtype=torch.float32).T - - workspace_buffer = torch.empty(32 * 1024 * 1024, dtype=torch.uint8, device=q.device) - flash_mla_cuda.dense_prefill_fwd( - workspace_buffer, - q, - k, - v, - cu_seqlens_qo, - cu_seqlens_kv, - out, - lse, - mask_mode_code, - softmax_scale, - max_seqlen_qo, - max_seqlen_kv, - is_varlen, - ) - - return out, lse - - -def _flash_attn_varlen_backward( - do: torch.Tensor, - q: torch.Tensor, - k: torch.Tensor, - v: torch.Tensor, - out: torch.Tensor, - lse: torch.Tensor, - cu_seqlens_qo: torch.Tensor, - cu_seqlens_kv: torch.Tensor, - max_seqlen_qo: int, - max_seqlen_kv: int, - dq: Optional[torch.Tensor] = None, - dk: Optional[torch.Tensor] = None, - dv: Optional[torch.Tensor] = None, - causal: bool = False, - softmax_scale: Optional[float] = None, - is_varlen: bool = True, -) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor]: - qo_total_len, num_qo_heads, head_dim_qk = q.shape - kv_total_len, num_kv_heads, head_dim_vo = v.shape - - # TODO: fix bwd GQA - if num_qo_heads != num_kv_heads: - raise ValueError(f"SM100 bwd doesn't support GQA now. num_qo_heads: {num_qo_heads}, num_kv_heads: {num_kv_heads}.") - - mask_mode_code = 1 if causal else 0 - if softmax_scale is None: - softmax_scale = head_dim_qk ** (-0.5) - - if dq is None: - dq = torch.empty(qo_total_len, num_qo_heads, head_dim_qk, device=q.device, dtype=q.dtype) - if dk is None: - dk = torch.empty(kv_total_len, num_kv_heads, head_dim_qk, device=q.device, dtype=q.dtype) - if dv is None: - dv = torch.empty(kv_total_len, num_kv_heads, head_dim_vo, device=q.device, dtype=q.dtype) - - max_seqlen_qo_aligned = (max_seqlen_qo + 7) // 8 * 8 - bs = cu_seqlens_qo.shape[0] - 1 - workspace_bytes = 0 - workspace_bytes += 4 * bs * max_seqlen_qo_aligned * num_qo_heads * head_dim_qk # dQ_acc - workspace_bytes += 4 * max_seqlen_qo_aligned * bs * num_qo_heads * 2 # sum_OdO and scaled_lse - if num_qo_heads != num_kv_heads: - workspace_bytes += 2 * kv_total_len * num_qo_heads * (head_dim_qk + head_dim_vo) # dKV_acc - workspace_buffer = torch.empty(workspace_bytes, dtype=torch.uint8, device=q.device) - flash_mla_cuda.dense_prefill_bwd( - workspace_buffer, - do, - q, - k, - v, - out, - lse, - cu_seqlens_qo, - cu_seqlens_kv, - dq, - dk, - dv, - mask_mode_code, - softmax_scale, - max_seqlen_qo, - max_seqlen_kv, - is_varlen, - ) - - return dq, dk, dv - - -class FlashAttnVarlenFunc(torch.autograd.Function): - def forward( - ctx, - q: torch.Tensor, - k: torch.Tensor, - v: torch.Tensor, - cu_seqlens_qo: torch.Tensor, - cu_seqlens_kv: torch.Tensor, - max_seqlen_qo: int, - max_seqlen_kv: int, - causal: bool = False, - softmax_scale: Optional[float] = None, - is_varlen: bool = True, - ) -> Tuple[torch.Tensor, torch.Tensor]: - out, lse = _flash_attn_varlen_forward( - q, k, v, - cu_seqlens_qo, cu_seqlens_kv, max_seqlen_qo, max_seqlen_kv, - causal=causal, softmax_scale=softmax_scale, - is_varlen=is_varlen, - ) - ctx.save_for_backward(q, k, v, out, lse, cu_seqlens_qo, cu_seqlens_kv) - ctx.max_seqlen_qo = max_seqlen_qo - ctx.max_seqlen_kv = max_seqlen_kv - ctx.causal = causal - ctx.softmax_scale = softmax_scale - ctx.is_varlen = is_varlen - return out, lse - - def backward( - ctx, - do: torch.Tensor, - dlse: torch.Tensor, - ): - del dlse # LSE doesn't support backward currently - q, k, v, out, lse, cu_seqlens_qo, cu_seqlens_kv = ctx.saved_tensors - dq, dk, dv = _flash_attn_varlen_backward( - do, q, k, v, out, lse, - cu_seqlens_qo, cu_seqlens_kv, ctx.max_seqlen_qo, ctx.max_seqlen_kv, - causal=ctx.causal, softmax_scale=ctx.softmax_scale, - is_varlen=ctx.is_varlen, - ) - return dq, dk, dv, None, None, None, None, None, None, None - - -def flash_attn_varlen_func( - q: torch.Tensor, - k: torch.Tensor, - v: torch.Tensor, - cu_seqlens_qo: torch.Tensor, - cu_seqlens_kv: torch.Tensor, - max_seqlen_qo: int, - max_seqlen_kv: int, - dropout_p: float = 0.0, - softmax_scale: Optional[float] = None, - causal: bool = False, - deterministic: bool = False, - is_varlen: bool = True, -) -> Tuple[torch.Tensor, torch.Tensor]: - assert dropout_p == 0.0 - assert not deterministic - return FlashAttnVarlenFunc.apply( - q, k, v, - cu_seqlens_qo, cu_seqlens_kv, max_seqlen_qo, max_seqlen_kv, - causal, softmax_scale, is_varlen, - ) - - -def flash_attn_varlen_qkvpacked_func( - qkv: torch.Tensor, - cu_seqlens: torch.Tensor, - max_seqlen: int, - head_dim_qk: int, - dropout_p: float = 0.0, - softmax_scale: Optional[float] = None, - causal: bool = False, - deterministic: bool = False, - is_varlen: bool = True, -) -> Tuple[torch.Tensor, torch.Tensor]: - assert dropout_p == 0.0 - assert not deterministic - return FlashAttnVarlenFunc.apply( - qkv[:, :, :head_dim_qk], qkv[:, :, head_dim_qk:head_dim_qk * 2], qkv[:, :, head_dim_qk * 2:], - cu_seqlens, cu_seqlens, max_seqlen, max_seqlen, - causal, softmax_scale, is_varlen, - ) - - -def flash_attn_varlen_kvpacked_func( - q: torch.Tensor, - kv: torch.Tensor, - cu_seqlens_qo: torch.Tensor, - cu_seqlens_kv: torch.Tensor, - max_seqlen_qo: int, - max_seqlen_kv: int, - head_dim_qk: int, - dropout_p: float = 0.0, - softmax_scale: Optional[float] = None, - causal: bool = False, - deterministic: bool = False, - is_varlen: bool = True, -) -> Tuple[torch.Tensor, torch.Tensor]: - assert dropout_p == 0.0 - assert not deterministic - return FlashAttnVarlenFunc.apply( - q, kv[:, :, :head_dim_qk], kv[:, :, head_dim_qk:], - cu_seqlens_qo, cu_seqlens_kv, max_seqlen_qo, max_seqlen_kv, - causal, softmax_scale, is_varlen, - ) diff --git a/build/torch29-cxx11-cu130-x86_64-linux/metadata.json b/build/torch29-cxx11-cu130-x86_64-linux/metadata.json deleted file mode 100644 index 8190d75efa8fd6449ddcd73de2072f17086e0842..0000000000000000000000000000000000000000 --- a/build/torch29-cxx11-cu130-x86_64-linux/metadata.json +++ /dev/null @@ -1,12 +0,0 @@ -{ - "version": 1, - "license": "MIT", - "python-depends": [], - "backend": { - "type": "cuda", - "archs": [ - "10.0f", - "9.0a" - ] - } -}