Hugging Face's logo

Kernels:
flashrt
/
transformer-layout-primitives

Kernel card Files
xet
 Community
liangsu9988 commited on
Commit
37c352a
·
verified ·
1 Parent(s): 072b453

Uploaded using `kernel-builder`.

Browse files
Files changed (17) hide show
  1. benchmarks/benchmark.py +44 -1
  2. build/torch211-cxx11-cu128-x86_64-linux/__init__.py +121 -1
  3. build/torch211-cxx11-cu128-x86_64-linux/_ops.py +3 -3
  4. build/torch211-cxx11-cu128-x86_64-linux/{_transformer_layout_primitives_cuda_c9b2fa0.abi3.so → _transformer_layout_primitives_cuda_7781728.abi3.so} +2 -2
  5. build/torch211-cxx11-cu128-x86_64-linux/metadata.json +1 -1
  6. build/torch211-cxx11-cu130-x86_64-linux/__init__.py +121 -1
  7. build/torch211-cxx11-cu130-x86_64-linux/_ops.py +3 -3
  8. build/torch211-cxx11-cu130-x86_64-linux/{_transformer_layout_primitives_cuda_c9b2fa0.abi3.so → _transformer_layout_primitives_cuda_7781728.abi3.so} +2 -2
  9. build/torch211-cxx11-cu130-x86_64-linux/metadata.json +1 -1
  10. build/torch212-cxx11-cu130-x86_64-linux/__init__.py +121 -1
  11. build/torch212-cxx11-cu130-x86_64-linux/_ops.py +3 -3
  12. build/torch212-cxx11-cu130-x86_64-linux/{_transformer_layout_primitives_cuda_c9b2fa0.abi3.so → _transformer_layout_primitives_cuda_7781728.abi3.so} +2 -2
  13. build/torch212-cxx11-cu130-x86_64-linux/metadata.json +1 -1
  14. build/torch212-cxx11-cu132-x86_64-linux/__init__.py +121 -1
  15. build/torch212-cxx11-cu132-x86_64-linux/_ops.py +3 -3
  16. build/torch212-cxx11-cu132-x86_64-linux/{_transformer_layout_primitives_cuda_c9b2fa0.abi3.so → _transformer_layout_primitives_cuda_7781728.abi3.so} +2 -2
  17. build/torch212-cxx11-cu132-x86_64-linux/metadata.json +1 -1
benchmarks/benchmark.py CHANGED
@@ -12,7 +12,12 @@ import torch
12
 
13
  ROOT = Path(__file__).resolve().parents[2]
14
  sys.path.insert(0, str(ROOT / "transformer-layout-primitives" / "tests"))
15
- from test_transformer_layout_primitives import load_source_ops, qk_rmsnorm_rope_ref, rotate_half_ref # noqa: E402
 
 
 
 
 
16
 
17
 
18
  def load_ops(backend: str, artifact: str | None):
@@ -102,6 +107,44 @@ def main() -> int:
102
  eu = time_us(eager_rope, max(5, args.warmup // 2), max(20, args.iters // 2))
103
  print(f"{name},{seq}x{heads}x{dim},rope_rotate_half_bf16_,{fu:.3f},{eu:.3f},{eu/fu:.2f}x")
104
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
105
  batch, seq, dim = 8, 2048, 2048
106
  x = torch.randn((batch * seq, dim), device="cuda", dtype=torch.bfloat16)
107
  gathered = torch.empty((2 * batch, dim), device="cuda", dtype=torch.bfloat16)
 
12
 
13
  ROOT = Path(__file__).resolve().parents[2]
14
  sys.path.insert(0, str(ROOT / "transformer-layout-primitives" / "tests"))
15
+ from test_transformer_layout_primitives import ( # noqa: E402
16
+ load_source_ops,
17
+ qk_pair_rmsnorm_rope_ref,
18
+ qk_rmsnorm_rope_ref,
19
+ rotate_half_ref,
20
+ )
21
 
22
 
23
  def load_ops(backend: str, artifact: str | None):
 
107
  eu = time_us(eager_rope, max(5, args.warmup // 2), max(20, args.iters // 2))
108
  print(f"{name},{seq}x{heads}x{dim},rope_rotate_half_bf16_,{fu:.3f},{eu:.3f},{eu/fu:.2f}x")
109
 
110
+ pair_shapes = [
111
+ ("groot_n17_llm", 277, 16, 8, 128),
112
+ ("qwen3vl_vision", 1024, 16, 16, 72),
113
+ ("lingbot_vision", 1024, 16, 16, 80),
114
+ ("video_transformer", 2520, 24, 24, 128),
115
+ ]
116
+ if args.mode == "full":
117
+ pair_shapes += [
118
+ ("action_boundary", 51, 16, 16, 64),
119
+ ("wan_partial_tile", 5070, 24, 24, 128),
120
+ ]
121
+ for name, rows, q_heads, k_heads, dim in pair_shapes:
122
+ q = torch.randn((rows, q_heads, dim), device="cuda", dtype=torch.bfloat16)
123
+ k = torch.randn((rows, k_heads, dim), device="cuda", dtype=torch.bfloat16)
124
+ q_weight = torch.randn((dim,), device="cuda", dtype=torch.bfloat16)
125
+ k_weight = torch.randn((dim,), device="cuda", dtype=torch.bfloat16)
126
+ cos = torch.randn((rows, dim), device="cuda", dtype=torch.bfloat16)
127
+ sin = torch.randn((rows, dim), device="cuda", dtype=torch.bfloat16)
128
+ q_out = torch.empty_like(q)
129
+ k_out = torch.empty_like(k)
130
+
131
+ def flash_pair():
132
+ ops.qk_pair_rmsnorm_rope_bf16(
133
+ q, k, q_weight, k_weight, cos, sin, q_out=q_out, k_out=k_out
134
+ )
135
+
136
+ def eager_pair():
137
+ qk_pair_rmsnorm_rope_ref(
138
+ q, k, q_weight, k_weight, cos, sin
139
+ )
140
+
141
+ fu = time_us(flash_pair, args.warmup, args.iters)
142
+ eu = time_us(eager_pair, max(5, args.warmup // 2), max(20, args.iters // 2))
143
+ print(
144
+ f"{name},{rows}x{q_heads}+{k_heads}x{dim},"
145
+ f"qk_pair_rmsnorm_rope_bf16,{fu:.3f},{eu:.3f},{eu/fu:.2f}x"
146
+ )
147
+
148
  batch, seq, dim = 8, 2048, 2048
149
  x = torch.randn((batch * seq, dim), device="cuda", dtype=torch.bfloat16)
150
  gathered = torch.empty((2 * batch, dim), device="cuda", dtype=torch.bfloat16)
build/torch211-cxx11-cu128-x86_64-linux/__init__.py CHANGED
@@ -55,13 +55,75 @@ def _qk_rmsnorm_rope_bf16_fake(
55
  weight: torch.Tensor,
56
  cos: torch.Tensor,
57
  sin: torch.Tensor,
58
- eps: float,
59
  ) -> None:
60
  if qk.dim() != 3 or weight.shape != (qk.shape[2],) or cos.shape != (qk.shape[0], qk.shape[2]) or sin.shape != cos.shape:
61
  raise RuntimeError("qk_rmsnorm_rope_bf16_ expects qk (rows, heads, dim), weight (dim,), cos/sin (rows, dim)")
62
  return None
63
 
64
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
65
  def fill_neginf_bf16(dst: torch.Tensor) -> torch.Tensor:
66
  ops.fill_neginf_bf16(dst)
67
  return dst
@@ -118,6 +180,61 @@ def qk_rmsnorm_rope_bf16_(
118
  return qk
119
 
120
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
121
  __all__ = [
122
  "fill_neginf_bf16",
123
  "add_bias_bf16_",
@@ -126,4 +243,7 @@ __all__ = [
126
  "text_scatter_bf16",
127
  "rope_rotate_half_bf16_",
128
  "qk_rmsnorm_rope_bf16_",
 
 
 
129
  ]
 
55
  weight: torch.Tensor,
56
  cos: torch.Tensor,
57
  sin: torch.Tensor,
58
+ eps: float = 1e-6,
59
  ) -> None:
60
  if qk.dim() != 3 or weight.shape != (qk.shape[2],) or cos.shape != (qk.shape[0], qk.shape[2]) or sin.shape != cos.shape:
61
  raise RuntimeError("qk_rmsnorm_rope_bf16_ expects qk (rows, heads, dim), weight (dim,), cos/sin (rows, dim)")
62
  return None
63
 
64
 
65
+ @torch.library.register_fake(add_op_namespace_prefix("qk_pair_rmsnorm_rope_bf16"))
66
+ def _qk_pair_rmsnorm_rope_bf16_fake(
67
+ q: torch.Tensor,
68
+ k: torch.Tensor,
69
+ q_weight: torch.Tensor,
70
+ k_weight: torch.Tensor,
71
+ cos: torch.Tensor,
72
+ sin: torch.Tensor,
73
+ eps: float,
74
+ q_out: torch.Tensor,
75
+ k_out: torch.Tensor,
76
+ ) -> None:
77
+ if (
78
+ q.dim() != 3
79
+ or k.dim() != 3
80
+ or q.shape[0] != k.shape[0]
81
+ or q.shape[2] != k.shape[2]
82
+ or q.shape[2] < 8
83
+ or q.shape[2] > 256
84
+ or q.shape[2] % 2 != 0
85
+ or q_weight.shape != (q.shape[2],)
86
+ or k_weight.shape != q_weight.shape
87
+ or cos.shape != (q.shape[0], q.shape[2])
88
+ or sin.shape != cos.shape
89
+ or q_out.shape != q.shape
90
+ or k_out.shape != k.shape
91
+ ):
92
+ raise RuntimeError(
93
+ "qk_pair_rmsnorm_rope_bf16 expects q/k (rows, heads, even_dim), "
94
+ "weights (dim,), cos/sin (rows, dim), and matching outputs"
95
+ )
96
+ return None
97
+
98
+
99
+ @torch.library.register_fake(add_op_namespace_prefix("gather_rows_bf16"))
100
+ def _gather_rows_bf16_fake(
101
+ src: torch.Tensor,
102
+ row_indices: torch.Tensor,
103
+ dst: torch.Tensor,
104
+ ) -> None:
105
+ if src.dim() != 2 or row_indices.dim() != 1 or dst.shape != (row_indices.numel(), src.shape[1]):
106
+ raise RuntimeError(
107
+ "gather_rows_bf16 expects src (source_rows, hidden), "
108
+ "row_indices (rows,), dst (rows, hidden)"
109
+ )
110
+ return None
111
+
112
+
113
+ @torch.library.register_fake(add_op_namespace_prefix("scatter_rows_bf16"))
114
+ def _scatter_rows_bf16_fake(
115
+ src: torch.Tensor,
116
+ row_indices: torch.Tensor,
117
+ dst: torch.Tensor,
118
+ ) -> None:
119
+ if src.dim() != 2 or row_indices.dim() != 1 or src.shape != (row_indices.numel(), dst.shape[1]):
120
+ raise RuntimeError(
121
+ "scatter_rows_bf16 expects src (rows, hidden), "
122
+ "row_indices (rows,), dst (destination_rows, hidden)"
123
+ )
124
+ return None
125
+
126
+
127
  def fill_neginf_bf16(dst: torch.Tensor) -> torch.Tensor:
128
  ops.fill_neginf_bf16(dst)
129
  return dst
 
180
  return qk
181
 
182
 
183
+ def qk_pair_rmsnorm_rope_bf16(
184
+ q: torch.Tensor,
185
+ k: torch.Tensor,
186
+ q_weight: torch.Tensor,
187
+ k_weight: torch.Tensor,
188
+ cos: torch.Tensor,
189
+ sin: torch.Tensor,
190
+ eps: float = 1e-6,
191
+ *,
192
+ q_out: Optional[torch.Tensor] = None,
193
+ k_out: Optional[torch.Tensor] = None,
194
+ ) -> tuple[torch.Tensor, torch.Tensor]:
195
+ if q_out is None:
196
+ q_out = torch.empty_like(q)
197
+ if k_out is None:
198
+ k_out = torch.empty_like(k)
199
+ ops.qk_pair_rmsnorm_rope_bf16(
200
+ q, k, q_weight, k_weight, cos, sin, float(eps), q_out, k_out
201
+ )
202
+ return q_out, k_out
203
+
204
+
205
+ def gather_rows_bf16(
206
+ src: torch.Tensor,
207
+ row_indices: torch.Tensor,
208
+ *,
209
+ out: Optional[torch.Tensor] = None,
210
+ ) -> torch.Tensor:
211
+ """Gather rows using in-range CUDA int64 indices without host synchronization."""
212
+
213
+ if out is None:
214
+ out = torch.empty(
215
+ (row_indices.numel(), src.shape[1]), device=src.device, dtype=src.dtype
216
+ )
217
+ ops.gather_rows_bf16(src, row_indices, out)
218
+ return out
219
+
220
+
221
+ def scatter_rows_bf16(
222
+ src: torch.Tensor,
223
+ row_indices: torch.Tensor,
224
+ destination_rows: int,
225
+ *,
226
+ out: Optional[torch.Tensor] = None,
227
+ ) -> torch.Tensor:
228
+ """Scatter rows to unique, in-range CUDA int64 indices."""
229
+
230
+ if out is None:
231
+ out = torch.zeros(
232
+ (destination_rows, src.shape[1]), device=src.device, dtype=src.dtype
233
+ )
234
+ ops.scatter_rows_bf16(src, row_indices, out)
235
+ return out
236
+
237
+
238
  __all__ = [
239
  "fill_neginf_bf16",
240
  "add_bias_bf16_",
 
243
  "text_scatter_bf16",
244
  "rope_rotate_half_bf16_",
245
  "qk_rmsnorm_rope_bf16_",
246
+ "qk_pair_rmsnorm_rope_bf16",
247
+ "gather_rows_bf16",
248
+ "scatter_rows_bf16",
249
  ]
build/torch211-cxx11-cu128-x86_64-linux/_ops.py CHANGED
@@ -1,9 +1,9 @@
1
  import torch
2
- from . import _transformer_layout_primitives_cuda_c9b2fa0
3
- ops = torch.ops._transformer_layout_primitives_cuda_c9b2fa0
4
 
5
  def add_op_namespace_prefix(op_name: str):
6
  """
7
  Prefix op by namespace.
8
  """
9
- return f"_transformer_layout_primitives_cuda_c9b2fa0::{op_name}"
 
1
  import torch
2
+ from . import _transformer_layout_primitives_cuda_7781728
3
+ ops = torch.ops._transformer_layout_primitives_cuda_7781728
4
 
5
  def add_op_namespace_prefix(op_name: str):
6
  """
7
  Prefix op by namespace.
8
  """
9
+ return f"_transformer_layout_primitives_cuda_7781728::{op_name}"
build/torch211-cxx11-cu128-x86_64-linux/{_transformer_layout_primitives_cuda_c9b2fa0.abi3.so → _transformer_layout_primitives_cuda_7781728.abi3.so} RENAMED
@@ -1,3 +1,3 @@
1
  version https://git-lfs.github.com/spec/v1
2
- oid sha256:8391a04a4384cc92d83c6df4c86ecdd205567c00105c13a01ca54c2b88fdd371
3
- size 572816
 
1
  version https://git-lfs.github.com/spec/v1
2
+ oid sha256:4019b88f701a6fed3a5be649ed4130cd412e10ca21731b164b1a1d6c89186240
3
+ size 816416
build/torch211-cxx11-cu128-x86_64-linux/metadata.json CHANGED
@@ -1,6 +1,6 @@
1
  {
2
  "name": "transformer-layout-primitives",
3
- "id": "_transformer_layout_primitives_cuda_c9b2fa0",
4
  "version": 1,
5
  "license": "Apache-2.0",
6
  "python-depends": [],
 
1
  {
2
  "name": "transformer-layout-primitives",
3
+ "id": "_transformer_layout_primitives_cuda_7781728",
4
  "version": 1,
5
  "license": "Apache-2.0",
6
  "python-depends": [],
build/torch211-cxx11-cu130-x86_64-linux/__init__.py CHANGED
@@ -55,13 +55,75 @@ def _qk_rmsnorm_rope_bf16_fake(
55
  weight: torch.Tensor,
56
  cos: torch.Tensor,
57
  sin: torch.Tensor,
58
- eps: float,
59
  ) -> None:
60
  if qk.dim() != 3 or weight.shape != (qk.shape[2],) or cos.shape != (qk.shape[0], qk.shape[2]) or sin.shape != cos.shape:
61
  raise RuntimeError("qk_rmsnorm_rope_bf16_ expects qk (rows, heads, dim), weight (dim,), cos/sin (rows, dim)")
62
  return None
63
 
64
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
65
  def fill_neginf_bf16(dst: torch.Tensor) -> torch.Tensor:
66
  ops.fill_neginf_bf16(dst)
67
  return dst
@@ -118,6 +180,61 @@ def qk_rmsnorm_rope_bf16_(
118
  return qk
119
 
120
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
121
  __all__ = [
122
  "fill_neginf_bf16",
123
  "add_bias_bf16_",
@@ -126,4 +243,7 @@ __all__ = [
126
  "text_scatter_bf16",
127
  "rope_rotate_half_bf16_",
128
  "qk_rmsnorm_rope_bf16_",
 
 
 
129
  ]
 
55
  weight: torch.Tensor,
56
  cos: torch.Tensor,
57
  sin: torch.Tensor,
58
+ eps: float = 1e-6,
59
  ) -> None:
60
  if qk.dim() != 3 or weight.shape != (qk.shape[2],) or cos.shape != (qk.shape[0], qk.shape[2]) or sin.shape != cos.shape:
61
  raise RuntimeError("qk_rmsnorm_rope_bf16_ expects qk (rows, heads, dim), weight (dim,), cos/sin (rows, dim)")
62
  return None
63
 
64
 
65
+ @torch.library.register_fake(add_op_namespace_prefix("qk_pair_rmsnorm_rope_bf16"))
66
+ def _qk_pair_rmsnorm_rope_bf16_fake(
67
+ q: torch.Tensor,
68
+ k: torch.Tensor,
69
+ q_weight: torch.Tensor,
70
+ k_weight: torch.Tensor,
71
+ cos: torch.Tensor,
72
+ sin: torch.Tensor,
73
+ eps: float,
74
+ q_out: torch.Tensor,
75
+ k_out: torch.Tensor,
76
+ ) -> None:
77
+ if (
78
+ q.dim() != 3
79
+ or k.dim() != 3
80
+ or q.shape[0] != k.shape[0]
81
+ or q.shape[2] != k.shape[2]
82
+ or q.shape[2] < 8
83
+ or q.shape[2] > 256
84
+ or q.shape[2] % 2 != 0
85
+ or q_weight.shape != (q.shape[2],)
86
+ or k_weight.shape != q_weight.shape
87
+ or cos.shape != (q.shape[0], q.shape[2])
88
+ or sin.shape != cos.shape
89
+ or q_out.shape != q.shape
90
+ or k_out.shape != k.shape
91
+ ):
92
+ raise RuntimeError(
93
+ "qk_pair_rmsnorm_rope_bf16 expects q/k (rows, heads, even_dim), "
94
+ "weights (dim,), cos/sin (rows, dim), and matching outputs"
95
+ )
96
+ return None
97
+
98
+
99
+ @torch.library.register_fake(add_op_namespace_prefix("gather_rows_bf16"))
100
+ def _gather_rows_bf16_fake(
101
+ src: torch.Tensor,
102
+ row_indices: torch.Tensor,
103
+ dst: torch.Tensor,
104
+ ) -> None:
105
+ if src.dim() != 2 or row_indices.dim() != 1 or dst.shape != (row_indices.numel(), src.shape[1]):
106
+ raise RuntimeError(
107
+ "gather_rows_bf16 expects src (source_rows, hidden), "
108
+ "row_indices (rows,), dst (rows, hidden)"
109
+ )
110
+ return None
111
+
112
+
113
+ @torch.library.register_fake(add_op_namespace_prefix("scatter_rows_bf16"))
114
+ def _scatter_rows_bf16_fake(
115
+ src: torch.Tensor,
116
+ row_indices: torch.Tensor,
117
+ dst: torch.Tensor,
118
+ ) -> None:
119
+ if src.dim() != 2 or row_indices.dim() != 1 or src.shape != (row_indices.numel(), dst.shape[1]):
120
+ raise RuntimeError(
121
+ "scatter_rows_bf16 expects src (rows, hidden), "
122
+ "row_indices (rows,), dst (destination_rows, hidden)"
123
+ )
124
+ return None
125
+
126
+
127
  def fill_neginf_bf16(dst: torch.Tensor) -> torch.Tensor:
128
  ops.fill_neginf_bf16(dst)
129
  return dst
 
180
  return qk
181
 
182
 
183
+ def qk_pair_rmsnorm_rope_bf16(
184
+ q: torch.Tensor,
185
+ k: torch.Tensor,
186
+ q_weight: torch.Tensor,
187
+ k_weight: torch.Tensor,
188
+ cos: torch.Tensor,
189
+ sin: torch.Tensor,
190
+ eps: float = 1e-6,
191
+ *,
192
+ q_out: Optional[torch.Tensor] = None,
193
+ k_out: Optional[torch.Tensor] = None,
194
+ ) -> tuple[torch.Tensor, torch.Tensor]:
195
+ if q_out is None:
196
+ q_out = torch.empty_like(q)
197
+ if k_out is None:
198
+ k_out = torch.empty_like(k)
199
+ ops.qk_pair_rmsnorm_rope_bf16(
200
+ q, k, q_weight, k_weight, cos, sin, float(eps), q_out, k_out
201
+ )
202
+ return q_out, k_out
203
+
204
+
205
+ def gather_rows_bf16(
206
+ src: torch.Tensor,
207
+ row_indices: torch.Tensor,
208
+ *,
209
+ out: Optional[torch.Tensor] = None,
210
+ ) -> torch.Tensor:
211
+ """Gather rows using in-range CUDA int64 indices without host synchronization."""
212
+
213
+ if out is None:
214
+ out = torch.empty(
215
+ (row_indices.numel(), src.shape[1]), device=src.device, dtype=src.dtype
216
+ )
217
+ ops.gather_rows_bf16(src, row_indices, out)
218
+ return out
219
+
220
+
221
+ def scatter_rows_bf16(
222
+ src: torch.Tensor,
223
+ row_indices: torch.Tensor,
224
+ destination_rows: int,
225
+ *,
226
+ out: Optional[torch.Tensor] = None,
227
+ ) -> torch.Tensor:
228
+ """Scatter rows to unique, in-range CUDA int64 indices."""
229
+
230
+ if out is None:
231
+ out = torch.zeros(
232
+ (destination_rows, src.shape[1]), device=src.device, dtype=src.dtype
233
+ )
234
+ ops.scatter_rows_bf16(src, row_indices, out)
235
+ return out
236
+
237
+
238
  __all__ = [
239
  "fill_neginf_bf16",
240
  "add_bias_bf16_",
 
243
  "text_scatter_bf16",
244
  "rope_rotate_half_bf16_",
245
  "qk_rmsnorm_rope_bf16_",
246
+ "qk_pair_rmsnorm_rope_bf16",
247
+ "gather_rows_bf16",
248
+ "scatter_rows_bf16",
249
  ]
build/torch211-cxx11-cu130-x86_64-linux/_ops.py CHANGED
@@ -1,9 +1,9 @@
1
  import torch
2
- from . import _transformer_layout_primitives_cuda_c9b2fa0
3
- ops = torch.ops._transformer_layout_primitives_cuda_c9b2fa0
4
 
5
  def add_op_namespace_prefix(op_name: str):
6
  """
7
  Prefix op by namespace.
8
  """
9
- return f"_transformer_layout_primitives_cuda_c9b2fa0::{op_name}"
 
1
  import torch
2
+ from . import _transformer_layout_primitives_cuda_7781728
3
+ ops = torch.ops._transformer_layout_primitives_cuda_7781728
4
 
5
  def add_op_namespace_prefix(op_name: str):
6
  """
7
  Prefix op by namespace.
8
  """
9
+ return f"_transformer_layout_primitives_cuda_7781728::{op_name}"
build/torch211-cxx11-cu130-x86_64-linux/{_transformer_layout_primitives_cuda_c9b2fa0.abi3.so → _transformer_layout_primitives_cuda_7781728.abi3.so} RENAMED
@@ -1,3 +1,3 @@
1
  version https://git-lfs.github.com/spec/v1
2
- oid sha256:99943084af51325b92b740243b59d224141161ee451749776e1c89ce7cceded2
3
- size 564160
 
1
  version https://git-lfs.github.com/spec/v1
2
+ oid sha256:f8e7ec238e03c3bfcd61470d74f50580b4eccd3c028c0a5915d3ef450384a4a3
3
+ size 803048
build/torch211-cxx11-cu130-x86_64-linux/metadata.json CHANGED
@@ -1,6 +1,6 @@
1
  {
2
  "name": "transformer-layout-primitives",
3
- "id": "_transformer_layout_primitives_cuda_c9b2fa0",
4
  "version": 1,
5
  "license": "Apache-2.0",
6
  "python-depends": [],
 
1
  {
2
  "name": "transformer-layout-primitives",
3
+ "id": "_transformer_layout_primitives_cuda_7781728",
4
  "version": 1,
5
  "license": "Apache-2.0",
6
  "python-depends": [],
build/torch212-cxx11-cu130-x86_64-linux/__init__.py CHANGED
@@ -55,13 +55,75 @@ def _qk_rmsnorm_rope_bf16_fake(
55
  weight: torch.Tensor,
56
  cos: torch.Tensor,
57
  sin: torch.Tensor,
58
- eps: float,
59
  ) -> None:
60
  if qk.dim() != 3 or weight.shape != (qk.shape[2],) or cos.shape != (qk.shape[0], qk.shape[2]) or sin.shape != cos.shape:
61
  raise RuntimeError("qk_rmsnorm_rope_bf16_ expects qk (rows, heads, dim), weight (dim,), cos/sin (rows, dim)")
62
  return None
63
 
64
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
65
  def fill_neginf_bf16(dst: torch.Tensor) -> torch.Tensor:
66
  ops.fill_neginf_bf16(dst)
67
  return dst
@@ -118,6 +180,61 @@ def qk_rmsnorm_rope_bf16_(
118
  return qk
119
 
120
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
121
  __all__ = [
122
  "fill_neginf_bf16",
123
  "add_bias_bf16_",
@@ -126,4 +243,7 @@ __all__ = [
126
  "text_scatter_bf16",
127
  "rope_rotate_half_bf16_",
128
  "qk_rmsnorm_rope_bf16_",
 
 
 
129
  ]
 
55
  weight: torch.Tensor,
56
  cos: torch.Tensor,
57
  sin: torch.Tensor,
58
+ eps: float = 1e-6,
59
  ) -> None:
60
  if qk.dim() != 3 or weight.shape != (qk.shape[2],) or cos.shape != (qk.shape[0], qk.shape[2]) or sin.shape != cos.shape:
61
  raise RuntimeError("qk_rmsnorm_rope_bf16_ expects qk (rows, heads, dim), weight (dim,), cos/sin (rows, dim)")
62
  return None
63
 
64
 
65
+ @torch.library.register_fake(add_op_namespace_prefix("qk_pair_rmsnorm_rope_bf16"))
66
+ def _qk_pair_rmsnorm_rope_bf16_fake(
67
+ q: torch.Tensor,
68
+ k: torch.Tensor,
69
+ q_weight: torch.Tensor,
70
+ k_weight: torch.Tensor,
71
+ cos: torch.Tensor,
72
+ sin: torch.Tensor,
73
+ eps: float,
74
+ q_out: torch.Tensor,
75
+ k_out: torch.Tensor,
76
+ ) -> None:
77
+ if (
78
+ q.dim() != 3
79
+ or k.dim() != 3
80
+ or q.shape[0] != k.shape[0]
81
+ or q.shape[2] != k.shape[2]
82
+ or q.shape[2] < 8
83
+ or q.shape[2] > 256
84
+ or q.shape[2] % 2 != 0
85
+ or q_weight.shape != (q.shape[2],)
86
+ or k_weight.shape != q_weight.shape
87
+ or cos.shape != (q.shape[0], q.shape[2])
88
+ or sin.shape != cos.shape
89
+ or q_out.shape != q.shape
90
+ or k_out.shape != k.shape
91
+ ):
92
+ raise RuntimeError(
93
+ "qk_pair_rmsnorm_rope_bf16 expects q/k (rows, heads, even_dim), "
94
+ "weights (dim,), cos/sin (rows, dim), and matching outputs"
95
+ )
96
+ return None
97
+
98
+
99
+ @torch.library.register_fake(add_op_namespace_prefix("gather_rows_bf16"))
100
+ def _gather_rows_bf16_fake(
101
+ src: torch.Tensor,
102
+ row_indices: torch.Tensor,
103
+ dst: torch.Tensor,
104
+ ) -> None:
105
+ if src.dim() != 2 or row_indices.dim() != 1 or dst.shape != (row_indices.numel(), src.shape[1]):
106
+ raise RuntimeError(
107
+ "gather_rows_bf16 expects src (source_rows, hidden), "
108
+ "row_indices (rows,), dst (rows, hidden)"
109
+ )
110
+ return None
111
+
112
+
113
+ @torch.library.register_fake(add_op_namespace_prefix("scatter_rows_bf16"))
114
+ def _scatter_rows_bf16_fake(
115
+ src: torch.Tensor,
116
+ row_indices: torch.Tensor,
117
+ dst: torch.Tensor,
118
+ ) -> None:
119
+ if src.dim() != 2 or row_indices.dim() != 1 or src.shape != (row_indices.numel(), dst.shape[1]):
120
+ raise RuntimeError(
121
+ "scatter_rows_bf16 expects src (rows, hidden), "
122
+ "row_indices (rows,), dst (destination_rows, hidden)"
123
+ )
124
+ return None
125
+
126
+
127
  def fill_neginf_bf16(dst: torch.Tensor) -> torch.Tensor:
128
  ops.fill_neginf_bf16(dst)
129
  return dst
 
180
  return qk
181
 
182
 
183
+ def qk_pair_rmsnorm_rope_bf16(
184
+ q: torch.Tensor,
185
+ k: torch.Tensor,
186
+ q_weight: torch.Tensor,
187
+ k_weight: torch.Tensor,
188
+ cos: torch.Tensor,
189
+ sin: torch.Tensor,
190
+ eps: float = 1e-6,
191
+ *,
192
+ q_out: Optional[torch.Tensor] = None,
193
+ k_out: Optional[torch.Tensor] = None,
194
+ ) -> tuple[torch.Tensor, torch.Tensor]:
195
+ if q_out is None:
196
+ q_out = torch.empty_like(q)
197
+ if k_out is None:
198
+ k_out = torch.empty_like(k)
199
+ ops.qk_pair_rmsnorm_rope_bf16(
200
+ q, k, q_weight, k_weight, cos, sin, float(eps), q_out, k_out
201
+ )
202
+ return q_out, k_out
203
+
204
+
205
+ def gather_rows_bf16(
206
+ src: torch.Tensor,
207
+ row_indices: torch.Tensor,
208
+ *,
209
+ out: Optional[torch.Tensor] = None,
210
+ ) -> torch.Tensor:
211
+ """Gather rows using in-range CUDA int64 indices without host synchronization."""
212
+
213
+ if out is None:
214
+ out = torch.empty(
215
+ (row_indices.numel(), src.shape[1]), device=src.device, dtype=src.dtype
216
+ )
217
+ ops.gather_rows_bf16(src, row_indices, out)
218
+ return out
219
+
220
+
221
+ def scatter_rows_bf16(
222
+ src: torch.Tensor,
223
+ row_indices: torch.Tensor,
224
+ destination_rows: int,
225
+ *,
226
+ out: Optional[torch.Tensor] = None,
227
+ ) -> torch.Tensor:
228
+ """Scatter rows to unique, in-range CUDA int64 indices."""
229
+
230
+ if out is None:
231
+ out = torch.zeros(
232
+ (destination_rows, src.shape[1]), device=src.device, dtype=src.dtype
233
+ )
234
+ ops.scatter_rows_bf16(src, row_indices, out)
235
+ return out
236
+
237
+
238
  __all__ = [
239
  "fill_neginf_bf16",
240
  "add_bias_bf16_",
 
243
  "text_scatter_bf16",
244
  "rope_rotate_half_bf16_",
245
  "qk_rmsnorm_rope_bf16_",
246
+ "qk_pair_rmsnorm_rope_bf16",
247
+ "gather_rows_bf16",
248
+ "scatter_rows_bf16",
249
  ]
build/torch212-cxx11-cu130-x86_64-linux/_ops.py CHANGED
@@ -1,9 +1,9 @@
1
  import torch
2
- from . import _transformer_layout_primitives_cuda_c9b2fa0
3
- ops = torch.ops._transformer_layout_primitives_cuda_c9b2fa0
4
 
5
  def add_op_namespace_prefix(op_name: str):
6
  """
7
  Prefix op by namespace.
8
  """
9
- return f"_transformer_layout_primitives_cuda_c9b2fa0::{op_name}"
 
1
  import torch
2
+ from . import _transformer_layout_primitives_cuda_7781728
3
+ ops = torch.ops._transformer_layout_primitives_cuda_7781728
4
 
5
  def add_op_namespace_prefix(op_name: str):
6
  """
7
  Prefix op by namespace.
8
  """
9
+ return f"_transformer_layout_primitives_cuda_7781728::{op_name}"
build/torch212-cxx11-cu130-x86_64-linux/{_transformer_layout_primitives_cuda_c9b2fa0.abi3.so → _transformer_layout_primitives_cuda_7781728.abi3.so} RENAMED
@@ -1,3 +1,3 @@
1
  version https://git-lfs.github.com/spec/v1
2
- oid sha256:6c4ebdc5e5faa851f8c4f537f5a1e75e28cebad98a650c5bf72dbf60db981844
3
- size 570424
 
1
  version https://git-lfs.github.com/spec/v1
2
+ oid sha256:926f6b87163a116fc90811edb58e041496595da240cf0d7e6b3ceb96eb15063a
3
+ size 809400
build/torch212-cxx11-cu130-x86_64-linux/metadata.json CHANGED
@@ -1,6 +1,6 @@
1
  {
2
  "name": "transformer-layout-primitives",
3
- "id": "_transformer_layout_primitives_cuda_c9b2fa0",
4
  "version": 1,
5
  "license": "Apache-2.0",
6
  "python-depends": [],
 
1
  {
2
  "name": "transformer-layout-primitives",
3
+ "id": "_transformer_layout_primitives_cuda_7781728",
4
  "version": 1,
5
  "license": "Apache-2.0",
6
  "python-depends": [],
build/torch212-cxx11-cu132-x86_64-linux/__init__.py CHANGED
@@ -55,13 +55,75 @@ def _qk_rmsnorm_rope_bf16_fake(
55
  weight: torch.Tensor,
56
  cos: torch.Tensor,
57
  sin: torch.Tensor,
58
- eps: float,
59
  ) -> None:
60
  if qk.dim() != 3 or weight.shape != (qk.shape[2],) or cos.shape != (qk.shape[0], qk.shape[2]) or sin.shape != cos.shape:
61
  raise RuntimeError("qk_rmsnorm_rope_bf16_ expects qk (rows, heads, dim), weight (dim,), cos/sin (rows, dim)")
62
  return None
63
 
64
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
65
  def fill_neginf_bf16(dst: torch.Tensor) -> torch.Tensor:
66
  ops.fill_neginf_bf16(dst)
67
  return dst
@@ -118,6 +180,61 @@ def qk_rmsnorm_rope_bf16_(
118
  return qk
119
 
120
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
121
  __all__ = [
122
  "fill_neginf_bf16",
123
  "add_bias_bf16_",
@@ -126,4 +243,7 @@ __all__ = [
126
  "text_scatter_bf16",
127
  "rope_rotate_half_bf16_",
128
  "qk_rmsnorm_rope_bf16_",
 
 
 
129
  ]
 
55
  weight: torch.Tensor,
56
  cos: torch.Tensor,
57
  sin: torch.Tensor,
58
+ eps: float = 1e-6,
59
  ) -> None:
60
  if qk.dim() != 3 or weight.shape != (qk.shape[2],) or cos.shape != (qk.shape[0], qk.shape[2]) or sin.shape != cos.shape:
61
  raise RuntimeError("qk_rmsnorm_rope_bf16_ expects qk (rows, heads, dim), weight (dim,), cos/sin (rows, dim)")
62
  return None
63
 
64
 
65
+ @torch.library.register_fake(add_op_namespace_prefix("qk_pair_rmsnorm_rope_bf16"))
66
+ def _qk_pair_rmsnorm_rope_bf16_fake(
67
+ q: torch.Tensor,
68
+ k: torch.Tensor,
69
+ q_weight: torch.Tensor,
70
+ k_weight: torch.Tensor,
71
+ cos: torch.Tensor,
72
+ sin: torch.Tensor,
73
+ eps: float,
74
+ q_out: torch.Tensor,
75
+ k_out: torch.Tensor,
76
+ ) -> None:
77
+ if (
78
+ q.dim() != 3
79
+ or k.dim() != 3
80
+ or q.shape[0] != k.shape[0]
81
+ or q.shape[2] != k.shape[2]
82
+ or q.shape[2] < 8
83
+ or q.shape[2] > 256
84
+ or q.shape[2] % 2 != 0
85
+ or q_weight.shape != (q.shape[2],)
86
+ or k_weight.shape != q_weight.shape
87
+ or cos.shape != (q.shape[0], q.shape[2])
88
+ or sin.shape != cos.shape
89
+ or q_out.shape != q.shape
90
+ or k_out.shape != k.shape
91
+ ):
92
+ raise RuntimeError(
93
+ "qk_pair_rmsnorm_rope_bf16 expects q/k (rows, heads, even_dim), "
94
+ "weights (dim,), cos/sin (rows, dim), and matching outputs"
95
+ )
96
+ return None
97
+
98
+
99
+ @torch.library.register_fake(add_op_namespace_prefix("gather_rows_bf16"))
100
+ def _gather_rows_bf16_fake(
101
+ src: torch.Tensor,
102
+ row_indices: torch.Tensor,
103
+ dst: torch.Tensor,
104
+ ) -> None:
105
+ if src.dim() != 2 or row_indices.dim() != 1 or dst.shape != (row_indices.numel(), src.shape[1]):
106
+ raise RuntimeError(
107
+ "gather_rows_bf16 expects src (source_rows, hidden), "
108
+ "row_indices (rows,), dst (rows, hidden)"
109
+ )
110
+ return None
111
+
112
+
113
+ @torch.library.register_fake(add_op_namespace_prefix("scatter_rows_bf16"))
114
+ def _scatter_rows_bf16_fake(
115
+ src: torch.Tensor,
116
+ row_indices: torch.Tensor,
117
+ dst: torch.Tensor,
118
+ ) -> None:
119
+ if src.dim() != 2 or row_indices.dim() != 1 or src.shape != (row_indices.numel(), dst.shape[1]):
120
+ raise RuntimeError(
121
+ "scatter_rows_bf16 expects src (rows, hidden), "
122
+ "row_indices (rows,), dst (destination_rows, hidden)"
123
+ )
124
+ return None
125
+
126
+
127
  def fill_neginf_bf16(dst: torch.Tensor) -> torch.Tensor:
128
  ops.fill_neginf_bf16(dst)
129
  return dst
 
180
  return qk
181
 
182
 
183
+ def qk_pair_rmsnorm_rope_bf16(
184
+ q: torch.Tensor,
185
+ k: torch.Tensor,
186
+ q_weight: torch.Tensor,
187
+ k_weight: torch.Tensor,
188
+ cos: torch.Tensor,
189
+ sin: torch.Tensor,
190
+ eps: float = 1e-6,
191
+ *,
192
+ q_out: Optional[torch.Tensor] = None,
193
+ k_out: Optional[torch.Tensor] = None,
194
+ ) -> tuple[torch.Tensor, torch.Tensor]:
195
+ if q_out is None:
196
+ q_out = torch.empty_like(q)
197
+ if k_out is None:
198
+ k_out = torch.empty_like(k)
199
+ ops.qk_pair_rmsnorm_rope_bf16(
200
+ q, k, q_weight, k_weight, cos, sin, float(eps), q_out, k_out
201
+ )
202
+ return q_out, k_out
203
+
204
+
205
+ def gather_rows_bf16(
206
+ src: torch.Tensor,
207
+ row_indices: torch.Tensor,
208
+ *,
209
+ out: Optional[torch.Tensor] = None,
210
+ ) -> torch.Tensor:
211
+ """Gather rows using in-range CUDA int64 indices without host synchronization."""
212
+
213
+ if out is None:
214
+ out = torch.empty(
215
+ (row_indices.numel(), src.shape[1]), device=src.device, dtype=src.dtype
216
+ )
217
+ ops.gather_rows_bf16(src, row_indices, out)
218
+ return out
219
+
220
+
221
+ def scatter_rows_bf16(
222
+ src: torch.Tensor,
223
+ row_indices: torch.Tensor,
224
+ destination_rows: int,
225
+ *,
226
+ out: Optional[torch.Tensor] = None,
227
+ ) -> torch.Tensor:
228
+ """Scatter rows to unique, in-range CUDA int64 indices."""
229
+
230
+ if out is None:
231
+ out = torch.zeros(
232
+ (destination_rows, src.shape[1]), device=src.device, dtype=src.dtype
233
+ )
234
+ ops.scatter_rows_bf16(src, row_indices, out)
235
+ return out
236
+
237
+
238
  __all__ = [
239
  "fill_neginf_bf16",
240
  "add_bias_bf16_",
 
243
  "text_scatter_bf16",
244
  "rope_rotate_half_bf16_",
245
  "qk_rmsnorm_rope_bf16_",
246
+ "qk_pair_rmsnorm_rope_bf16",
247
+ "gather_rows_bf16",
248
+ "scatter_rows_bf16",
249
  ]
build/torch212-cxx11-cu132-x86_64-linux/_ops.py CHANGED
@@ -1,9 +1,9 @@
1
  import torch
2
- from . import _transformer_layout_primitives_cuda_c9b2fa0
3
- ops = torch.ops._transformer_layout_primitives_cuda_c9b2fa0
4
 
5
  def add_op_namespace_prefix(op_name: str):
6
  """
7
  Prefix op by namespace.
8
  """
9
- return f"_transformer_layout_primitives_cuda_c9b2fa0::{op_name}"
 
1
  import torch
2
+ from . import _transformer_layout_primitives_cuda_7781728
3
+ ops = torch.ops._transformer_layout_primitives_cuda_7781728
4
 
5
  def add_op_namespace_prefix(op_name: str):
6
  """
7
  Prefix op by namespace.
8
  """
9
+ return f"_transformer_layout_primitives_cuda_7781728::{op_name}"
build/torch212-cxx11-cu132-x86_64-linux/{_transformer_layout_primitives_cuda_c9b2fa0.abi3.so → _transformer_layout_primitives_cuda_7781728.abi3.so} RENAMED
@@ -1,3 +1,3 @@
1
  version https://git-lfs.github.com/spec/v1
2
- oid sha256:1947aebc33c49a48c82a0da681b56e16c0517460121fa5ce7cbc1d39f1e9832b
3
- size 574520
 
1
  version https://git-lfs.github.com/spec/v1
2
+ oid sha256:6a0708065d5489fcac2b6ab817a845f604227a9a779c5ed433340eade428be00
3
+ size 813496
build/torch212-cxx11-cu132-x86_64-linux/metadata.json CHANGED
@@ -1,6 +1,6 @@
1
  {
2
  "name": "transformer-layout-primitives",
3
- "id": "_transformer_layout_primitives_cuda_c9b2fa0",
4
  "version": 1,
5
  "license": "Apache-2.0",
6
  "python-depends": [],
 
1
  {
2
  "name": "transformer-layout-primitives",
3
+ "id": "_transformer_layout_primitives_cuda_7781728",
4
  "version": 1,
5
  "license": "Apache-2.0",
6
  "python-depends": [],