Build (aarch64)

build/torch26-cxx11-cu126-aarch64-linux/flash_mla/__init__.py

@@ -0,0 +1,33 @@
+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
+    )

build/torch26-cxx11-cu126-aarch64-linux/flash_mla/_flash_mla_341ab77.abi3.so

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

build/torch26-cxx11-cu126-aarch64-linux/flash_mla/_ops.py

@@ -0,0 +1,9 @@
+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}"

build/torch26-cxx98-cu126-aarch64-linux/flash_mla/__init__.py

@@ -0,0 +1,33 @@
+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
+    )

build/torch26-cxx98-cu126-aarch64-linux/flash_mla/_flash_mla_341ab77.abi3.so

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

build/torch26-cxx98-cu126-aarch64-linux/flash_mla/_ops.py

@@ -0,0 +1,9 @@
+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}"

build/torch27-cxx11-cu126-aarch64-linux/flash_mla/__init__.py

@@ -0,0 +1,33 @@
+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
+    )

build/torch27-cxx11-cu126-aarch64-linux/flash_mla/_flash_mla_341ab77.abi3.so

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

build/torch27-cxx11-cu126-aarch64-linux/flash_mla/_ops.py

@@ -0,0 +1,9 @@
+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}"

build/torch27-cxx11-cu128-aarch64-linux/flash_mla/__init__.py

@@ -0,0 +1,33 @@
+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
+    )

build/torch27-cxx11-cu128-aarch64-linux/flash_mla/_flash_mla_341ab77.abi3.so

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

build/torch27-cxx11-cu128-aarch64-linux/flash_mla/_ops.py

@@ -0,0 +1,9 @@
+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}"