Hugging Face's logo

Kernels:
flashrt
/
masked-mha-runtime

Kernel card Files
xet
 Community
liangsu9988 commited on
Commit
1a7e37f
·
verified ·
1 Parent(s): 03b8715

Uploaded using `kernel-builder`.

Browse files
Files changed (16) hide show
  1. build/torch212-cxx11-cu130-x86_64-linux/__init__.py +143 -0
  2. build/torch212-cxx11-cu130-x86_64-linux/_masked_mha_runtime_cuda_ce210b3.abi3.so +3 -0
  3. build/torch212-cxx11-cu130-x86_64-linux/_ops.py +9 -0
  4. build/torch212-cxx11-cu130-x86_64-linux/metadata.json +34 -0
  5. build/torch212-cxx11-cu132-x86_64-linux/__init__.py +143 -0
  6. build/torch212-cxx11-cu132-x86_64-linux/_masked_mha_runtime_cuda_ce210b3.abi3.so +3 -0
  7. build/torch212-cxx11-cu132-x86_64-linux/_ops.py +9 -0
  8. build/torch212-cxx11-cu132-x86_64-linux/metadata.json +34 -0
  9. build/torch213-cxx11-cu130-x86_64-linux/__init__.py +143 -0
  10. build/torch213-cxx11-cu130-x86_64-linux/_masked_mha_runtime_cuda_ce210b3.abi3.so +3 -0
  11. build/torch213-cxx11-cu130-x86_64-linux/_ops.py +9 -0
  12. build/torch213-cxx11-cu130-x86_64-linux/metadata.json +34 -0
  13. build/torch213-cxx11-cu132-x86_64-linux/__init__.py +143 -0
  14. build/torch213-cxx11-cu132-x86_64-linux/_masked_mha_runtime_cuda_ce210b3.abi3.so +3 -0
  15. build/torch213-cxx11-cu132-x86_64-linux/_ops.py +9 -0
  16. build/torch213-cxx11-cu132-x86_64-linux/metadata.json +34 -0
build/torch212-cxx11-cu130-x86_64-linux/__init__.py ADDED
@@ -0,0 +1,143 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ """Allocation-free masked MHA runtime operators."""
2
+
3
+ from __future__ import annotations
4
+
5
+ import math
6
+ from typing import Optional
7
+
8
+ import torch
9
+
10
+ from ._ops import add_op_namespace_prefix, ops
11
+
12
+
13
+ @torch.library.register_fake(add_op_namespace_prefix("forward_static"))
14
+ def _forward_static_fake(q, k, v, logits, out, scale: float) -> None:
15
+ del scale
16
+ if q.dim() != 3 or k.dim() != 3 or v.dim() != 3:
17
+ raise RuntimeError("q/k/v must have shape (sequence, heads, head_dim)")
18
+ if out.shape != q.shape:
19
+ raise RuntimeError("out must match q")
20
+ if logits.dim() != 3 or logits.shape[:2] != (q.shape[1], q.shape[0]):
21
+ raise RuntimeError("logits must have shape (heads, sequence_q, stride)")
22
+ if logits.shape[2] < k.shape[0]:
23
+ raise RuntimeError("logits stride must cover sequence_kv")
24
+
25
+
26
+ @torch.library.register_fake(add_op_namespace_prefix("attention_mha_fp16_masked"))
27
+ def _attention_mha_fp16_masked_fake(q, k, v, logits, out, scale: float) -> None:
28
+ _forward_static_fake(q, k, v, logits, out, scale)
29
+
30
+
31
+ @torch.library.register_fake(add_op_namespace_prefix("attention_mha_bf16_masked"))
32
+ def _attention_mha_bf16_masked_fake(
33
+ q, k, v, logits, out, scale: float, qkv_token_stride: int
34
+ ) -> None:
35
+ _forward_static_fake(q, k, v, logits, out, scale)
36
+ if qkv_token_stride != q.stride(0):
37
+ raise RuntimeError("qkv_token_stride must match q.stride(0)")
38
+
39
+
40
+ @torch.library.register_fake(add_op_namespace_prefix("forward_seqused_static"))
41
+ def _forward_seqused_static_fake(q, k, v, valid_k, logits, out, scale: float) -> None:
42
+ del scale
43
+ if q.dim() != 3 or k.dim() != 2 or v.dim() != 2:
44
+ raise RuntimeError("q must be (sequence, heads, head_dim); k/v must be (max_sequence, head_dim)")
45
+ if valid_k.numel() != 1 or valid_k.dtype != torch.int32:
46
+ raise RuntimeError("valid_k must be an int32 scalar tensor")
47
+ if out.shape != q.shape or logits.shape != (q.shape[0] * q.shape[1], k.shape[0]):
48
+ raise RuntimeError("caller-owned output shapes do not match q/k")
49
+
50
+
51
+ def allocate_workspace(q: torch.Tensor, k: torch.Tensor) -> torch.Tensor:
52
+ """Allocate padded logits scratch once, outside the hot path."""
53
+ stride = (k.shape[0] + 7) // 8 * 8
54
+ return torch.empty(
55
+ (q.shape[1], q.shape[0], stride), device=q.device, dtype=q.dtype
56
+ )
57
+
58
+
59
+ def forward_static(
60
+ q: torch.Tensor,
61
+ k: torch.Tensor,
62
+ v: torch.Tensor,
63
+ *,
64
+ logits: torch.Tensor,
65
+ out: torch.Tensor,
66
+ scale: Optional[float] = None,
67
+ ) -> torch.Tensor:
68
+ """Run MHA without pre-filling padded logits; all buffers are caller-owned."""
69
+ if scale is None:
70
+ scale = 1.0 / math.sqrt(q.shape[-1])
71
+ ops.forward_static(q, k, v, logits, out, float(scale))
72
+ return out
73
+
74
+
75
+ def attention_mha_fp16_masked(
76
+ q: torch.Tensor,
77
+ k: torch.Tensor,
78
+ v: torch.Tensor,
79
+ *,
80
+ logits: torch.Tensor,
81
+ out: torch.Tensor,
82
+ scale: Optional[float] = None,
83
+ ) -> torch.Tensor:
84
+ """Run the deterministic valid-column FP16 masked MHA kernel."""
85
+ if scale is None:
86
+ scale = 1.0 / math.sqrt(q.shape[-1])
87
+ ops.attention_mha_fp16_masked(q, k, v, logits, out, float(scale))
88
+ return out
89
+
90
+
91
+ def attention_mha_bf16_masked(
92
+ q: torch.Tensor,
93
+ k: torch.Tensor,
94
+ v: torch.Tensor,
95
+ *,
96
+ logits: torch.Tensor,
97
+ out: torch.Tensor,
98
+ qkv_token_stride: Optional[int] = None,
99
+ scale: Optional[float] = None,
100
+ ) -> torch.Tensor:
101
+ """Run BF16 masked MHA directly from contiguous or fused-QKV views."""
102
+ if scale is None:
103
+ scale = 1.0 / math.sqrt(q.shape[-1])
104
+ if qkv_token_stride is None:
105
+ qkv_token_stride = q.stride(0)
106
+ ops.attention_mha_bf16_masked(
107
+ q, k, v, logits, out, float(scale), int(qkv_token_stride)
108
+ )
109
+ return out
110
+
111
+
112
+ def forward(q, k, v, *, scale: Optional[float] = None):
113
+ """Convenience allocation wrapper; use ``forward_static`` in hot paths."""
114
+ logits = allocate_workspace(q, k)
115
+ out = torch.empty_like(q, memory_format=torch.contiguous_format)
116
+ return forward_static(q, k, v, logits=logits, out=out, scale=scale)
117
+
118
+
119
+ def forward_seqused_static(
120
+ q: torch.Tensor,
121
+ k: torch.Tensor,
122
+ v: torch.Tensor,
123
+ valid_k: torch.Tensor,
124
+ *,
125
+ logits: torch.Tensor,
126
+ out: torch.Tensor,
127
+ scale: Optional[float] = None,
128
+ ) -> torch.Tensor:
129
+ """Run FP16 shared-KV attention with a graph-resident valid-key count."""
130
+ if scale is None:
131
+ scale = 1.0 / math.sqrt(q.shape[-1])
132
+ ops.forward_seqused_static(q, k, v, valid_k, logits, out, float(scale))
133
+ return out
134
+
135
+
136
+ __all__ = [
137
+ "attention_mha_bf16_masked",
138
+ "attention_mha_fp16_masked",
139
+ "allocate_workspace",
140
+ "forward",
141
+ "forward_static",
142
+ "forward_seqused_static",
143
+ ]
build/torch212-cxx11-cu130-x86_64-linux/_masked_mha_runtime_cuda_ce210b3.abi3.so ADDED
@@ -0,0 +1,3 @@
 
 
 
 
1
+ version https://git-lfs.github.com/spec/v1
2
+ oid sha256:88d8d18a68bcb7494bafde0ee0dc24b708a35539ff20608642d11874c4e7536c
3
+ size 438384
build/torch212-cxx11-cu130-x86_64-linux/_ops.py ADDED
@@ -0,0 +1,9 @@
 
 
 
 
 
 
 
 
 
 
1
+ import torch
2
+ from . import _masked_mha_runtime_cuda_ce210b3
3
+ ops = torch.ops._masked_mha_runtime_cuda_ce210b3
4
+
5
+ def add_op_namespace_prefix(op_name: str):
6
+ """
7
+ Prefix op by namespace.
8
+ """
9
+ return f"_masked_mha_runtime_cuda_ce210b3::{op_name}"
build/torch212-cxx11-cu130-x86_64-linux/metadata.json ADDED
@@ -0,0 +1,34 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ {
2
+ "name": "masked-mha-runtime",
3
+ "id": "_masked_mha_runtime_cuda_ce210b3",
4
+ "version": 1,
5
+ "license": "Apache-2.0",
6
+ "python-depends": [],
7
+ "kernel-depends": [],
8
+ "backend": {
9
+ "type": "cuda",
10
+ "archs": [
11
+ "11.0a",
12
+ "12.0"
13
+ ]
14
+ },
15
+ "digest": {
16
+ "algorithm": "sha256",
17
+ "files": {
18
+ "__init__.py": "w4T07fNWmll9bTURpoHVj+llGGGmT6f7hRpFFKZjBoQ=",
19
+ "_masked_mha_runtime_cuda_ce210b3.abi3.so": "iNjRimi8t0lLr94O4NwktwijVTn/IGCGQtEYdMTnU2w=",
20
+ "_ops.py": "sGta1hB+MbwsaTBv+TPRDifw8Ztrf0Q1bJ2usT/karo="
21
+ }
22
+ },
23
+ "provenance": {
24
+ "kernel-builder": {
25
+ "version": "0.17.0-dev0",
26
+ "sha": "2c40e10d3da1392245ee81bd05b55fb3103b2b34",
27
+ "dirty": false
28
+ },
29
+ "kernel": {
30
+ "sha": "ce210b38ac208aa6dbe87d2028c75e028e1b8ef3",
31
+ "dirty": false
32
+ }
33
+ }
34
+ }
build/torch212-cxx11-cu132-x86_64-linux/__init__.py ADDED
@@ -0,0 +1,143 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ """Allocation-free masked MHA runtime operators."""
2
+
3
+ from __future__ import annotations
4
+
5
+ import math
6
+ from typing import Optional
7
+
8
+ import torch
9
+
10
+ from ._ops import add_op_namespace_prefix, ops
11
+
12
+
13
+ @torch.library.register_fake(add_op_namespace_prefix("forward_static"))
14
+ def _forward_static_fake(q, k, v, logits, out, scale: float) -> None:
15
+ del scale
16
+ if q.dim() != 3 or k.dim() != 3 or v.dim() != 3:
17
+ raise RuntimeError("q/k/v must have shape (sequence, heads, head_dim)")
18
+ if out.shape != q.shape:
19
+ raise RuntimeError("out must match q")
20
+ if logits.dim() != 3 or logits.shape[:2] != (q.shape[1], q.shape[0]):
21
+ raise RuntimeError("logits must have shape (heads, sequence_q, stride)")
22
+ if logits.shape[2] < k.shape[0]:
23
+ raise RuntimeError("logits stride must cover sequence_kv")
24
+
25
+
26
+ @torch.library.register_fake(add_op_namespace_prefix("attention_mha_fp16_masked"))
27
+ def _attention_mha_fp16_masked_fake(q, k, v, logits, out, scale: float) -> None:
28
+ _forward_static_fake(q, k, v, logits, out, scale)
29
+
30
+
31
+ @torch.library.register_fake(add_op_namespace_prefix("attention_mha_bf16_masked"))
32
+ def _attention_mha_bf16_masked_fake(
33
+ q, k, v, logits, out, scale: float, qkv_token_stride: int
34
+ ) -> None:
35
+ _forward_static_fake(q, k, v, logits, out, scale)
36
+ if qkv_token_stride != q.stride(0):
37
+ raise RuntimeError("qkv_token_stride must match q.stride(0)")
38
+
39
+
40
+ @torch.library.register_fake(add_op_namespace_prefix("forward_seqused_static"))
41
+ def _forward_seqused_static_fake(q, k, v, valid_k, logits, out, scale: float) -> None:
42
+ del scale
43
+ if q.dim() != 3 or k.dim() != 2 or v.dim() != 2:
44
+ raise RuntimeError("q must be (sequence, heads, head_dim); k/v must be (max_sequence, head_dim)")
45
+ if valid_k.numel() != 1 or valid_k.dtype != torch.int32:
46
+ raise RuntimeError("valid_k must be an int32 scalar tensor")
47
+ if out.shape != q.shape or logits.shape != (q.shape[0] * q.shape[1], k.shape[0]):
48
+ raise RuntimeError("caller-owned output shapes do not match q/k")
49
+
50
+
51
+ def allocate_workspace(q: torch.Tensor, k: torch.Tensor) -> torch.Tensor:
52
+ """Allocate padded logits scratch once, outside the hot path."""
53
+ stride = (k.shape[0] + 7) // 8 * 8
54
+ return torch.empty(
55
+ (q.shape[1], q.shape[0], stride), device=q.device, dtype=q.dtype
56
+ )
57
+
58
+
59
+ def forward_static(
60
+ q: torch.Tensor,
61
+ k: torch.Tensor,
62
+ v: torch.Tensor,
63
+ *,
64
+ logits: torch.Tensor,
65
+ out: torch.Tensor,
66
+ scale: Optional[float] = None,
67
+ ) -> torch.Tensor:
68
+ """Run MHA without pre-filling padded logits; all buffers are caller-owned."""
69
+ if scale is None:
70
+ scale = 1.0 / math.sqrt(q.shape[-1])
71
+ ops.forward_static(q, k, v, logits, out, float(scale))
72
+ return out
73
+
74
+
75
+ def attention_mha_fp16_masked(
76
+ q: torch.Tensor,
77
+ k: torch.Tensor,
78
+ v: torch.Tensor,
79
+ *,
80
+ logits: torch.Tensor,
81
+ out: torch.Tensor,
82
+ scale: Optional[float] = None,
83
+ ) -> torch.Tensor:
84
+ """Run the deterministic valid-column FP16 masked MHA kernel."""
85
+ if scale is None:
86
+ scale = 1.0 / math.sqrt(q.shape[-1])
87
+ ops.attention_mha_fp16_masked(q, k, v, logits, out, float(scale))
88
+ return out
89
+
90
+
91
+ def attention_mha_bf16_masked(
92
+ q: torch.Tensor,
93
+ k: torch.Tensor,
94
+ v: torch.Tensor,
95
+ *,
96
+ logits: torch.Tensor,
97
+ out: torch.Tensor,
98
+ qkv_token_stride: Optional[int] = None,
99
+ scale: Optional[float] = None,
100
+ ) -> torch.Tensor:
101
+ """Run BF16 masked MHA directly from contiguous or fused-QKV views."""
102
+ if scale is None:
103
+ scale = 1.0 / math.sqrt(q.shape[-1])
104
+ if qkv_token_stride is None:
105
+ qkv_token_stride = q.stride(0)
106
+ ops.attention_mha_bf16_masked(
107
+ q, k, v, logits, out, float(scale), int(qkv_token_stride)
108
+ )
109
+ return out
110
+
111
+
112
+ def forward(q, k, v, *, scale: Optional[float] = None):
113
+ """Convenience allocation wrapper; use ``forward_static`` in hot paths."""
114
+ logits = allocate_workspace(q, k)
115
+ out = torch.empty_like(q, memory_format=torch.contiguous_format)
116
+ return forward_static(q, k, v, logits=logits, out=out, scale=scale)
117
+
118
+
119
+ def forward_seqused_static(
120
+ q: torch.Tensor,
121
+ k: torch.Tensor,
122
+ v: torch.Tensor,
123
+ valid_k: torch.Tensor,
124
+ *,
125
+ logits: torch.Tensor,
126
+ out: torch.Tensor,
127
+ scale: Optional[float] = None,
128
+ ) -> torch.Tensor:
129
+ """Run FP16 shared-KV attention with a graph-resident valid-key count."""
130
+ if scale is None:
131
+ scale = 1.0 / math.sqrt(q.shape[-1])
132
+ ops.forward_seqused_static(q, k, v, valid_k, logits, out, float(scale))
133
+ return out
134
+
135
+
136
+ __all__ = [
137
+ "attention_mha_bf16_masked",
138
+ "attention_mha_fp16_masked",
139
+ "allocate_workspace",
140
+ "forward",
141
+ "forward_static",
142
+ "forward_seqused_static",
143
+ ]
build/torch212-cxx11-cu132-x86_64-linux/_masked_mha_runtime_cuda_ce210b3.abi3.so ADDED
@@ -0,0 +1,3 @@
 
 
 
 
1
+ version https://git-lfs.github.com/spec/v1
2
+ oid sha256:4c9c9938f4f9eec4542ed0b98d65c1ac2b8950608a98883c17a2a51457cec5c4
3
+ size 438384
build/torch212-cxx11-cu132-x86_64-linux/_ops.py ADDED
@@ -0,0 +1,9 @@
 
 
 
 
 
 
 
 
 
 
1
+ import torch
2
+ from . import _masked_mha_runtime_cuda_ce210b3
3
+ ops = torch.ops._masked_mha_runtime_cuda_ce210b3
4
+
5
+ def add_op_namespace_prefix(op_name: str):
6
+ """
7
+ Prefix op by namespace.
8
+ """
9
+ return f"_masked_mha_runtime_cuda_ce210b3::{op_name}"
build/torch212-cxx11-cu132-x86_64-linux/metadata.json ADDED
@@ -0,0 +1,34 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ {
2
+ "name": "masked-mha-runtime",
3
+ "id": "_masked_mha_runtime_cuda_ce210b3",
4
+ "version": 1,
5
+ "license": "Apache-2.0",
6
+ "python-depends": [],
7
+ "kernel-depends": [],
8
+ "backend": {
9
+ "type": "cuda",
10
+ "archs": [
11
+ "11.0a",
12
+ "12.0"
13
+ ]
14
+ },
15
+ "digest": {
16
+ "algorithm": "sha256",
17
+ "files": {
18
+ "__init__.py": "w4T07fNWmll9bTURpoHVj+llGGGmT6f7hRpFFKZjBoQ=",
19
+ "_masked_mha_runtime_cuda_ce210b3.abi3.so": "TJyZOPT57sRULtC5jWXBrCuJUGCKmIg8F6KlFFfOxcQ=",
20
+ "_ops.py": "sGta1hB+MbwsaTBv+TPRDifw8Ztrf0Q1bJ2usT/karo="
21
+ }
22
+ },
23
+ "provenance": {
24
+ "kernel-builder": {
25
+ "version": "0.17.0-dev0",
26
+ "sha": "2c40e10d3da1392245ee81bd05b55fb3103b2b34",
27
+ "dirty": false
28
+ },
29
+ "kernel": {
30
+ "sha": "ce210b38ac208aa6dbe87d2028c75e028e1b8ef3",
31
+ "dirty": false
32
+ }
33
+ }
34
+ }
build/torch213-cxx11-cu130-x86_64-linux/__init__.py ADDED
@@ -0,0 +1,143 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ """Allocation-free masked MHA runtime operators."""
2
+
3
+ from __future__ import annotations
4
+
5
+ import math
6
+ from typing import Optional
7
+
8
+ import torch
9
+
10
+ from ._ops import add_op_namespace_prefix, ops
11
+
12
+
13
+ @torch.library.register_fake(add_op_namespace_prefix("forward_static"))
14
+ def _forward_static_fake(q, k, v, logits, out, scale: float) -> None:
15
+ del scale
16
+ if q.dim() != 3 or k.dim() != 3 or v.dim() != 3:
17
+ raise RuntimeError("q/k/v must have shape (sequence, heads, head_dim)")
18
+ if out.shape != q.shape:
19
+ raise RuntimeError("out must match q")
20
+ if logits.dim() != 3 or logits.shape[:2] != (q.shape[1], q.shape[0]):
21
+ raise RuntimeError("logits must have shape (heads, sequence_q, stride)")
22
+ if logits.shape[2] < k.shape[0]:
23
+ raise RuntimeError("logits stride must cover sequence_kv")
24
+
25
+
26
+ @torch.library.register_fake(add_op_namespace_prefix("attention_mha_fp16_masked"))
27
+ def _attention_mha_fp16_masked_fake(q, k, v, logits, out, scale: float) -> None:
28
+ _forward_static_fake(q, k, v, logits, out, scale)
29
+
30
+
31
+ @torch.library.register_fake(add_op_namespace_prefix("attention_mha_bf16_masked"))
32
+ def _attention_mha_bf16_masked_fake(
33
+ q, k, v, logits, out, scale: float, qkv_token_stride: int
34
+ ) -> None:
35
+ _forward_static_fake(q, k, v, logits, out, scale)
36
+ if qkv_token_stride != q.stride(0):
37
+ raise RuntimeError("qkv_token_stride must match q.stride(0)")
38
+
39
+
40
+ @torch.library.register_fake(add_op_namespace_prefix("forward_seqused_static"))
41
+ def _forward_seqused_static_fake(q, k, v, valid_k, logits, out, scale: float) -> None:
42
+ del scale
43
+ if q.dim() != 3 or k.dim() != 2 or v.dim() != 2:
44
+ raise RuntimeError("q must be (sequence, heads, head_dim); k/v must be (max_sequence, head_dim)")
45
+ if valid_k.numel() != 1 or valid_k.dtype != torch.int32:
46
+ raise RuntimeError("valid_k must be an int32 scalar tensor")
47
+ if out.shape != q.shape or logits.shape != (q.shape[0] * q.shape[1], k.shape[0]):
48
+ raise RuntimeError("caller-owned output shapes do not match q/k")
49
+
50
+
51
+ def allocate_workspace(q: torch.Tensor, k: torch.Tensor) -> torch.Tensor:
52
+ """Allocate padded logits scratch once, outside the hot path."""
53
+ stride = (k.shape[0] + 7) // 8 * 8
54
+ return torch.empty(
55
+ (q.shape[1], q.shape[0], stride), device=q.device, dtype=q.dtype
56
+ )
57
+
58
+
59
+ def forward_static(
60
+ q: torch.Tensor,
61
+ k: torch.Tensor,
62
+ v: torch.Tensor,
63
+ *,
64
+ logits: torch.Tensor,
65
+ out: torch.Tensor,
66
+ scale: Optional[float] = None,
67
+ ) -> torch.Tensor:
68
+ """Run MHA without pre-filling padded logits; all buffers are caller-owned."""
69
+ if scale is None:
70
+ scale = 1.0 / math.sqrt(q.shape[-1])
71
+ ops.forward_static(q, k, v, logits, out, float(scale))
72
+ return out
73
+
74
+
75
+ def attention_mha_fp16_masked(
76
+ q: torch.Tensor,
77
+ k: torch.Tensor,
78
+ v: torch.Tensor,
79
+ *,
80
+ logits: torch.Tensor,
81
+ out: torch.Tensor,
82
+ scale: Optional[float] = None,
83
+ ) -> torch.Tensor:
84
+ """Run the deterministic valid-column FP16 masked MHA kernel."""
85
+ if scale is None:
86
+ scale = 1.0 / math.sqrt(q.shape[-1])
87
+ ops.attention_mha_fp16_masked(q, k, v, logits, out, float(scale))
88
+ return out
89
+
90
+
91
+ def attention_mha_bf16_masked(
92
+ q: torch.Tensor,
93
+ k: torch.Tensor,
94
+ v: torch.Tensor,
95
+ *,
96
+ logits: torch.Tensor,
97
+ out: torch.Tensor,
98
+ qkv_token_stride: Optional[int] = None,
99
+ scale: Optional[float] = None,
100
+ ) -> torch.Tensor:
101
+ """Run BF16 masked MHA directly from contiguous or fused-QKV views."""
102
+ if scale is None:
103
+ scale = 1.0 / math.sqrt(q.shape[-1])
104
+ if qkv_token_stride is None:
105
+ qkv_token_stride = q.stride(0)
106
+ ops.attention_mha_bf16_masked(
107
+ q, k, v, logits, out, float(scale), int(qkv_token_stride)
108
+ )
109
+ return out
110
+
111
+
112
+ def forward(q, k, v, *, scale: Optional[float] = None):
113
+ """Convenience allocation wrapper; use ``forward_static`` in hot paths."""
114
+ logits = allocate_workspace(q, k)
115
+ out = torch.empty_like(q, memory_format=torch.contiguous_format)
116
+ return forward_static(q, k, v, logits=logits, out=out, scale=scale)
117
+
118
+
119
+ def forward_seqused_static(
120
+ q: torch.Tensor,
121
+ k: torch.Tensor,
122
+ v: torch.Tensor,
123
+ valid_k: torch.Tensor,
124
+ *,
125
+ logits: torch.Tensor,
126
+ out: torch.Tensor,
127
+ scale: Optional[float] = None,
128
+ ) -> torch.Tensor:
129
+ """Run FP16 shared-KV attention with a graph-resident valid-key count."""
130
+ if scale is None:
131
+ scale = 1.0 / math.sqrt(q.shape[-1])
132
+ ops.forward_seqused_static(q, k, v, valid_k, logits, out, float(scale))
133
+ return out
134
+
135
+
136
+ __all__ = [
137
+ "attention_mha_bf16_masked",
138
+ "attention_mha_fp16_masked",
139
+ "allocate_workspace",
140
+ "forward",
141
+ "forward_static",
142
+ "forward_seqused_static",
143
+ ]
build/torch213-cxx11-cu130-x86_64-linux/_masked_mha_runtime_cuda_ce210b3.abi3.so ADDED
@@ -0,0 +1,3 @@
 
 
 
 
1
+ version https://git-lfs.github.com/spec/v1
2
+ oid sha256:8e252d170e56e86838f0e4de6250311668e9eeccdb8a88b0abb9e385ee2f0d8f
3
+ size 438224
build/torch213-cxx11-cu130-x86_64-linux/_ops.py ADDED
@@ -0,0 +1,9 @@
 
 
 
 
 
 
 
 
 
 
1
+ import torch
2
+ from . import _masked_mha_runtime_cuda_ce210b3
3
+ ops = torch.ops._masked_mha_runtime_cuda_ce210b3
4
+
5
+ def add_op_namespace_prefix(op_name: str):
6
+ """
7
+ Prefix op by namespace.
8
+ """
9
+ return f"_masked_mha_runtime_cuda_ce210b3::{op_name}"
build/torch213-cxx11-cu130-x86_64-linux/metadata.json ADDED
@@ -0,0 +1,34 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ {
2
+ "name": "masked-mha-runtime",
3
+ "id": "_masked_mha_runtime_cuda_ce210b3",
4
+ "version": 1,
5
+ "license": "Apache-2.0",
6
+ "python-depends": [],
7
+ "kernel-depends": [],
8
+ "backend": {
9
+ "type": "cuda",
10
+ "archs": [
11
+ "11.0a",
12
+ "12.0"
13
+ ]
14
+ },
15
+ "digest": {
16
+ "algorithm": "sha256",
17
+ "files": {
18
+ "__init__.py": "w4T07fNWmll9bTURpoHVj+llGGGmT6f7hRpFFKZjBoQ=",
19
+ "_masked_mha_runtime_cuda_ce210b3.abi3.so": "jiUtFw5W6Gg48OTeYlAxFmjp7szbioiwq7njhe4vDY8=",
20
+ "_ops.py": "sGta1hB+MbwsaTBv+TPRDifw8Ztrf0Q1bJ2usT/karo="
21
+ }
22
+ },
23
+ "provenance": {
24
+ "kernel-builder": {
25
+ "version": "0.17.0-dev0",
26
+ "sha": "2c40e10d3da1392245ee81bd05b55fb3103b2b34",
27
+ "dirty": false
28
+ },
29
+ "kernel": {
30
+ "sha": "ce210b38ac208aa6dbe87d2028c75e028e1b8ef3",
31
+ "dirty": false
32
+ }
33
+ }
34
+ }
build/torch213-cxx11-cu132-x86_64-linux/__init__.py ADDED
@@ -0,0 +1,143 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ """Allocation-free masked MHA runtime operators."""
2
+
3
+ from __future__ import annotations
4
+
5
+ import math
6
+ from typing import Optional
7
+
8
+ import torch
9
+
10
+ from ._ops import add_op_namespace_prefix, ops
11
+
12
+
13
+ @torch.library.register_fake(add_op_namespace_prefix("forward_static"))
14
+ def _forward_static_fake(q, k, v, logits, out, scale: float) -> None:
15
+ del scale
16
+ if q.dim() != 3 or k.dim() != 3 or v.dim() != 3:
17
+ raise RuntimeError("q/k/v must have shape (sequence, heads, head_dim)")
18
+ if out.shape != q.shape:
19
+ raise RuntimeError("out must match q")
20
+ if logits.dim() != 3 or logits.shape[:2] != (q.shape[1], q.shape[0]):
21
+ raise RuntimeError("logits must have shape (heads, sequence_q, stride)")
22
+ if logits.shape[2] < k.shape[0]:
23
+ raise RuntimeError("logits stride must cover sequence_kv")
24
+
25
+
26
+ @torch.library.register_fake(add_op_namespace_prefix("attention_mha_fp16_masked"))
27
+ def _attention_mha_fp16_masked_fake(q, k, v, logits, out, scale: float) -> None:
28
+ _forward_static_fake(q, k, v, logits, out, scale)
29
+
30
+
31
+ @torch.library.register_fake(add_op_namespace_prefix("attention_mha_bf16_masked"))
32
+ def _attention_mha_bf16_masked_fake(
33
+ q, k, v, logits, out, scale: float, qkv_token_stride: int
34
+ ) -> None:
35
+ _forward_static_fake(q, k, v, logits, out, scale)
36
+ if qkv_token_stride != q.stride(0):
37
+ raise RuntimeError("qkv_token_stride must match q.stride(0)")
38
+
39
+
40
+ @torch.library.register_fake(add_op_namespace_prefix("forward_seqused_static"))
41
+ def _forward_seqused_static_fake(q, k, v, valid_k, logits, out, scale: float) -> None:
42
+ del scale
43
+ if q.dim() != 3 or k.dim() != 2 or v.dim() != 2:
44
+ raise RuntimeError("q must be (sequence, heads, head_dim); k/v must be (max_sequence, head_dim)")
45
+ if valid_k.numel() != 1 or valid_k.dtype != torch.int32:
46
+ raise RuntimeError("valid_k must be an int32 scalar tensor")
47
+ if out.shape != q.shape or logits.shape != (q.shape[0] * q.shape[1], k.shape[0]):
48
+ raise RuntimeError("caller-owned output shapes do not match q/k")
49
+
50
+
51
+ def allocate_workspace(q: torch.Tensor, k: torch.Tensor) -> torch.Tensor:
52
+ """Allocate padded logits scratch once, outside the hot path."""
53
+ stride = (k.shape[0] + 7) // 8 * 8
54
+ return torch.empty(
55
+ (q.shape[1], q.shape[0], stride), device=q.device, dtype=q.dtype
56
+ )
57
+
58
+
59
+ def forward_static(
60
+ q: torch.Tensor,
61
+ k: torch.Tensor,
62
+ v: torch.Tensor,
63
+ *,
64
+ logits: torch.Tensor,
65
+ out: torch.Tensor,
66
+ scale: Optional[float] = None,
67
+ ) -> torch.Tensor:
68
+ """Run MHA without pre-filling padded logits; all buffers are caller-owned."""
69
+ if scale is None:
70
+ scale = 1.0 / math.sqrt(q.shape[-1])
71
+ ops.forward_static(q, k, v, logits, out, float(scale))
72
+ return out
73
+
74
+
75
+ def attention_mha_fp16_masked(
76
+ q: torch.Tensor,
77
+ k: torch.Tensor,
78
+ v: torch.Tensor,
79
+ *,
80
+ logits: torch.Tensor,
81
+ out: torch.Tensor,
82
+ scale: Optional[float] = None,
83
+ ) -> torch.Tensor:
84
+ """Run the deterministic valid-column FP16 masked MHA kernel."""
85
+ if scale is None:
86
+ scale = 1.0 / math.sqrt(q.shape[-1])
87
+ ops.attention_mha_fp16_masked(q, k, v, logits, out, float(scale))
88
+ return out
89
+
90
+
91
+ def attention_mha_bf16_masked(
92
+ q: torch.Tensor,
93
+ k: torch.Tensor,
94
+ v: torch.Tensor,
95
+ *,
96
+ logits: torch.Tensor,
97
+ out: torch.Tensor,
98
+ qkv_token_stride: Optional[int] = None,
99
+ scale: Optional[float] = None,
100
+ ) -> torch.Tensor:
101
+ """Run BF16 masked MHA directly from contiguous or fused-QKV views."""
102
+ if scale is None:
103
+ scale = 1.0 / math.sqrt(q.shape[-1])
104
+ if qkv_token_stride is None:
105
+ qkv_token_stride = q.stride(0)
106
+ ops.attention_mha_bf16_masked(
107
+ q, k, v, logits, out, float(scale), int(qkv_token_stride)
108
+ )
109
+ return out
110
+
111
+
112
+ def forward(q, k, v, *, scale: Optional[float] = None):
113
+ """Convenience allocation wrapper; use ``forward_static`` in hot paths."""
114
+ logits = allocate_workspace(q, k)
115
+ out = torch.empty_like(q, memory_format=torch.contiguous_format)
116
+ return forward_static(q, k, v, logits=logits, out=out, scale=scale)
117
+
118
+
119
+ def forward_seqused_static(
120
+ q: torch.Tensor,
121
+ k: torch.Tensor,
122
+ v: torch.Tensor,
123
+ valid_k: torch.Tensor,
124
+ *,
125
+ logits: torch.Tensor,
126
+ out: torch.Tensor,
127
+ scale: Optional[float] = None,
128
+ ) -> torch.Tensor:
129
+ """Run FP16 shared-KV attention with a graph-resident valid-key count."""
130
+ if scale is None:
131
+ scale = 1.0 / math.sqrt(q.shape[-1])
132
+ ops.forward_seqused_static(q, k, v, valid_k, logits, out, float(scale))
133
+ return out
134
+
135
+
136
+ __all__ = [
137
+ "attention_mha_bf16_masked",
138
+ "attention_mha_fp16_masked",
139
+ "allocate_workspace",
140
+ "forward",
141
+ "forward_static",
142
+ "forward_seqused_static",
143
+ ]
build/torch213-cxx11-cu132-x86_64-linux/_masked_mha_runtime_cuda_ce210b3.abi3.so ADDED
@@ -0,0 +1,3 @@
 
 
 
 
1
+ version https://git-lfs.github.com/spec/v1
2
+ oid sha256:e4b0b75ead28d998d37f7cbddb3045a7803961d8298892953ce562ac791835ce
3
+ size 438224
build/torch213-cxx11-cu132-x86_64-linux/_ops.py ADDED
@@ -0,0 +1,9 @@
 
 
 
 
 
 
 
 
 
 
1
+ import torch
2
+ from . import _masked_mha_runtime_cuda_ce210b3
3
+ ops = torch.ops._masked_mha_runtime_cuda_ce210b3
4
+
5
+ def add_op_namespace_prefix(op_name: str):
6
+ """
7
+ Prefix op by namespace.
8
+ """
9
+ return f"_masked_mha_runtime_cuda_ce210b3::{op_name}"
build/torch213-cxx11-cu132-x86_64-linux/metadata.json ADDED
@@ -0,0 +1,34 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ {
2
+ "name": "masked-mha-runtime",
3
+ "id": "_masked_mha_runtime_cuda_ce210b3",
4
+ "version": 1,
5
+ "license": "Apache-2.0",
6
+ "python-depends": [],
7
+ "kernel-depends": [],
8
+ "backend": {
9
+ "type": "cuda",
10
+ "archs": [
11
+ "11.0a",
12
+ "12.0"
13
+ ]
14
+ },
15
+ "digest": {
16
+ "algorithm": "sha256",
17
+ "files": {
18
+ "__init__.py": "w4T07fNWmll9bTURpoHVj+llGGGmT6f7hRpFFKZjBoQ=",
19
+ "_masked_mha_runtime_cuda_ce210b3.abi3.so": "5LC3Xq0o2ZjTf3y92zBFp4A5YdgpiJKVPOVirHkYNc4=",
20
+ "_ops.py": "sGta1hB+MbwsaTBv+TPRDifw8Ztrf0Q1bJ2usT/karo="
21
+ }
22
+ },
23
+ "provenance": {
24
+ "kernel-builder": {
25
+ "version": "0.17.0-dev0",
26
+ "sha": "2c40e10d3da1392245ee81bd05b55fb3103b2b34",
27
+ "dirty": false
28
+ },
29
+ "kernel": {
30
+ "sha": "ce210b38ac208aa6dbe87d2028c75e028e1b8ef3",
31
+ "dirty": false
32
+ }
33
+ }
34
+ }