Promote latest kernel artifacts to main

README.md

@@ -1,9 +0,0 @@
-# flashrt/transformer-layout-primitives
-
-This repository is a compatibility mirror for older `kernels` clients
-that resolve repositories through the default Hugging Face model repo API.
-
-Canonical Kernel Hub repo: https://huggingface.co/kernels/flashrt/transformer-layout-primitives
-
-Do not edit this mirror by hand. It is generated from the Kernel Hub
-`vN` branches and contains the same `build/**` artifacts.

benchmarks/benchmark.py

@@ -12,7 +12,12 @@ import torch
 
 ROOT = Path(__file__).resolve().parents[2]
 sys.path.insert(0, str(ROOT / "transformer-layout-primitives" / "tests"))
-from test_transformer_layout_primitives import load_source_ops, qk_rmsnorm_rope_ref, rotate_half_ref  # noqa: E402
+from test_transformer_layout_primitives import (  # noqa: E402
+    load_source_ops,
+    qk_pair_rmsnorm_rope_ref,
+    qk_rmsnorm_rope_ref,
+    rotate_half_ref,
+)
 
 
 def load_ops(backend: str, artifact: str | None):
@@ -102,6 +107,44 @@ def main() -> int:
         eu = time_us(eager_rope, max(5, args.warmup // 2), max(20, args.iters // 2))
         print(f"{name},{seq}x{heads}x{dim},rope_rotate_half_bf16_,{fu:.3f},{eu:.3f},{eu/fu:.2f}x")
 
+    pair_shapes = [
+        ("groot_n17_llm", 277, 16, 8, 128),
+        ("qwen3vl_vision", 1024, 16, 16, 72),
+        ("lingbot_vision", 1024, 16, 16, 80),
+        ("video_transformer", 2520, 24, 24, 128),
+    ]
+    if args.mode == "full":
+        pair_shapes += [
+            ("action_boundary", 51, 16, 16, 64),
+            ("wan_partial_tile", 5070, 24, 24, 128),
+        ]
+    for name, rows, q_heads, k_heads, dim in pair_shapes:
+        q = torch.randn((rows, q_heads, dim), device="cuda", dtype=torch.bfloat16)
+        k = torch.randn((rows, k_heads, dim), device="cuda", dtype=torch.bfloat16)
+        q_weight = torch.randn((dim,), device="cuda", dtype=torch.bfloat16)
+        k_weight = torch.randn((dim,), device="cuda", dtype=torch.bfloat16)
+        cos = torch.randn((rows, dim), device="cuda", dtype=torch.bfloat16)
+        sin = torch.randn((rows, dim), device="cuda", dtype=torch.bfloat16)
+        q_out = torch.empty_like(q)
+        k_out = torch.empty_like(k)
+
+        def flash_pair():
+            ops.qk_pair_rmsnorm_rope_bf16(
+                q, k, q_weight, k_weight, cos, sin, q_out=q_out, k_out=k_out
+            )
+
+        def eager_pair():
+            qk_pair_rmsnorm_rope_ref(
+                q, k, q_weight, k_weight, cos, sin
+            )
+
+        fu = time_us(flash_pair, args.warmup, args.iters)
+        eu = time_us(eager_pair, max(5, args.warmup // 2), max(20, args.iters // 2))
+        print(
+            f"{name},{rows}x{q_heads}+{k_heads}x{dim},"
+            f"qk_pair_rmsnorm_rope_bf16,{fu:.3f},{eu:.3f},{eu/fu:.2f}x"
+        )
+
     batch, seq, dim = 8, 2048, 2048
     x = torch.randn((batch * seq, dim), device="cuda", dtype=torch.bfloat16)
     gathered = torch.empty((2 * batch, dim), device="cuda", dtype=torch.bfloat16)

build/torch211-cxx11-cu128-x86_64-linux/__init__.py

@@ -55,13 +55,75 @@ def _qk_rmsnorm_rope_bf16_fake(
     weight: torch.Tensor,
     cos: torch.Tensor,
     sin: torch.Tensor,
-    eps: float,
+    eps: float = 1e-6,
 ) -> None:
     if qk.dim() != 3 or weight.shape != (qk.shape[2],) or cos.shape != (qk.shape[0], qk.shape[2]) or sin.shape != cos.shape:
         raise RuntimeError("qk_rmsnorm_rope_bf16_ expects qk (rows, heads, dim), weight (dim,), cos/sin (rows, dim)")
     return None
 
 
+@torch.library.register_fake(add_op_namespace_prefix("qk_pair_rmsnorm_rope_bf16"))
+def _qk_pair_rmsnorm_rope_bf16_fake(
+    q: torch.Tensor,
+    k: torch.Tensor,
+    q_weight: torch.Tensor,
+    k_weight: torch.Tensor,
+    cos: torch.Tensor,
+    sin: torch.Tensor,
+    eps: float,
+    q_out: torch.Tensor,
+    k_out: torch.Tensor,
+) -> None:
+    if (
+        q.dim() != 3
+        or k.dim() != 3
+        or q.shape[0] != k.shape[0]
+        or q.shape[2] != k.shape[2]
+        or q.shape[2] < 8
+        or q.shape[2] > 256
+        or q.shape[2] % 2 != 0
+        or q_weight.shape != (q.shape[2],)
+        or k_weight.shape != q_weight.shape
+        or cos.shape != (q.shape[0], q.shape[2])
+        or sin.shape != cos.shape
+        or q_out.shape != q.shape
+        or k_out.shape != k.shape
+    ):
+        raise RuntimeError(
+            "qk_pair_rmsnorm_rope_bf16 expects q/k (rows, heads, even_dim), "
+            "weights (dim,), cos/sin (rows, dim), and matching outputs"
+        )
+    return None
+
+
+@torch.library.register_fake(add_op_namespace_prefix("gather_rows_bf16"))
+def _gather_rows_bf16_fake(
+    src: torch.Tensor,
+    row_indices: torch.Tensor,
+    dst: torch.Tensor,
+) -> None:
+    if src.dim() != 2 or row_indices.dim() != 1 or dst.shape != (row_indices.numel(), src.shape[1]):
+        raise RuntimeError(
+            "gather_rows_bf16 expects src (source_rows, hidden), "
+            "row_indices (rows,), dst (rows, hidden)"
+        )
+    return None
+
+
+@torch.library.register_fake(add_op_namespace_prefix("scatter_rows_bf16"))
+def _scatter_rows_bf16_fake(
+    src: torch.Tensor,
+    row_indices: torch.Tensor,
+    dst: torch.Tensor,
+) -> None:
+    if src.dim() != 2 or row_indices.dim() != 1 or src.shape != (row_indices.numel(), dst.shape[1]):
+        raise RuntimeError(
+            "scatter_rows_bf16 expects src (rows, hidden), "
+            "row_indices (rows,), dst (destination_rows, hidden)"
+        )
+    return None
+
+
 def fill_neginf_bf16(dst: torch.Tensor) -> torch.Tensor:
     ops.fill_neginf_bf16(dst)
     return dst
@@ -118,6 +180,61 @@ def qk_rmsnorm_rope_bf16_(
     return qk
 
 
+def qk_pair_rmsnorm_rope_bf16(
+    q: torch.Tensor,
+    k: torch.Tensor,
+    q_weight: torch.Tensor,
+    k_weight: torch.Tensor,
+    cos: torch.Tensor,
+    sin: torch.Tensor,
+    eps: float = 1e-6,
+    *,
+    q_out: Optional[torch.Tensor] = None,
+    k_out: Optional[torch.Tensor] = None,
+) -> tuple[torch.Tensor, torch.Tensor]:
+    if q_out is None:
+        q_out = torch.empty_like(q)
+    if k_out is None:
+        k_out = torch.empty_like(k)
+    ops.qk_pair_rmsnorm_rope_bf16(
+        q, k, q_weight, k_weight, cos, sin, float(eps), q_out, k_out
+    )
+    return q_out, k_out
+
+
+def gather_rows_bf16(
+    src: torch.Tensor,
+    row_indices: torch.Tensor,
+    *,
+    out: Optional[torch.Tensor] = None,
+) -> torch.Tensor:
+    """Gather rows using in-range CUDA int64 indices without host synchronization."""
+
+    if out is None:
+        out = torch.empty(
+            (row_indices.numel(), src.shape[1]), device=src.device, dtype=src.dtype
+        )
+    ops.gather_rows_bf16(src, row_indices, out)
+    return out
+
+
+def scatter_rows_bf16(
+    src: torch.Tensor,
+    row_indices: torch.Tensor,
+    destination_rows: int,
+    *,
+    out: Optional[torch.Tensor] = None,
+) -> torch.Tensor:
+    """Scatter rows to unique, in-range CUDA int64 indices."""
+
+    if out is None:
+        out = torch.zeros(
+            (destination_rows, src.shape[1]), device=src.device, dtype=src.dtype
+        )
+    ops.scatter_rows_bf16(src, row_indices, out)
+    return out
+
+
 __all__ = [
     "fill_neginf_bf16",
     "add_bias_bf16_",
@@ -126,4 +243,7 @@ __all__ = [
     "text_scatter_bf16",
     "rope_rotate_half_bf16_",
     "qk_rmsnorm_rope_bf16_",
+    "qk_pair_rmsnorm_rope_bf16",
+    "gather_rows_bf16",
+    "scatter_rows_bf16",
 ]

build/torch211-cxx11-cu128-x86_64-linux/_ops.py

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

build/{torch212-cxx11-cu132-x86_64-linux/_transformer_layout_primitives_cuda_c9b2fa0.abi3.so → torch211-cxx11-cu128-x86_64-linux/_transformer_layout_primitives_cuda_0a2c98e.abi3.so}

@@ -1,3 +1,3 @@
 version https://git-lfs.github.com/spec/v1
-oid sha256:1947aebc33c49a48c82a0da681b56e16c0517460121fa5ce7cbc1d39f1e9832b
-size 574520
+oid sha256:837e8b7e5c5f099bdad66ac56e375101f4e54ba47af08478f3f132222fadaa84
+size 816416

build/torch211-cxx11-cu128-x86_64-linux/metadata.json

@@ -1,6 +1,6 @@
 {
   "name": "transformer-layout-primitives",
-  "id": "_transformer_layout_primitives_cuda_c9b2fa0",
+  "id": "_transformer_layout_primitives_cuda_0a2c98e",
   "version": 1,
   "license": "Apache-2.0",
   "python-depends": [],

build/torch211-cxx11-cu130-aarch64-linux/__init__.py

@@ -0,0 +1,249 @@
+"""FlashRT transformer layout primitives."""
+
+from __future__ import annotations
+
+from typing import Optional
+
+import torch
+
+from ._ops import add_op_namespace_prefix, ops
+
+
+@torch.library.register_fake(add_op_namespace_prefix("fill_neginf_bf16"))
+def _fill_neginf_bf16_fake(dst: torch.Tensor) -> None:
+    return None
+
+
+@torch.library.register_fake(add_op_namespace_prefix("add_bias_bf16_"))
+def _add_bias_bf16_fake(data: torch.Tensor, bias: torch.Tensor) -> None:
+    if data.dim() != 2 or bias.shape != (data.shape[1],):
+        raise RuntimeError("add_bias_bf16_ expects data (rows, cols), bias (cols,)")
+    return None
+
+
+@torch.library.register_fake(add_op_namespace_prefix("repeat_interleave_heads_bf16"))
+def _repeat_interleave_heads_bf16_fake(src: torch.Tensor, repeat: int, dst: torch.Tensor) -> None:
+    if src.dim() != 3 or dst.shape != (src.shape[0], src.shape[1] * repeat, src.shape[2]):
+        raise RuntimeError("repeat_interleave_heads_bf16 expects src (seq, heads, dim), dst (seq, heads*repeat, dim)")
+    return None
+
+
+@torch.library.register_fake(add_op_namespace_prefix("text_gather_bf16"))
+def _text_gather_bf16_fake(src: torch.Tensor, batch: int, seq: int, dst: torch.Tensor) -> None:
+    if src.dim() != 2 or src.shape[0] != batch * seq or dst.shape != (2 * batch, src.shape[1]):
+        raise RuntimeError("text_gather_bf16 expects src (batch*seq, dim), dst (2*batch, dim)")
+    return None
+
+
+@torch.library.register_fake(add_op_namespace_prefix("text_scatter_bf16"))
+def _text_scatter_bf16_fake(dst: torch.Tensor, src: torch.Tensor, batch: int, seq: int) -> None:
+    if dst.dim() != 2 or dst.shape[0] != batch * seq or src.shape != (2 * batch, dst.shape[1]):
+        raise RuntimeError("text_scatter_bf16 expects dst (batch*seq, dim), src (2*batch, dim)")
+    return None
+
+
+@torch.library.register_fake(add_op_namespace_prefix("rope_rotate_half_bf16_"))
+def _rope_rotate_half_bf16_fake(x: torch.Tensor, cos: torch.Tensor, sin: torch.Tensor) -> None:
+    if x.dim() != 3 or x.shape[2] % 2 != 0 or cos.shape != (x.shape[0], x.shape[2]) or sin.shape != cos.shape:
+        raise RuntimeError("rope_rotate_half_bf16_ expects x (seq, heads, even_dim), cos/sin (seq, dim)")
+    return None
+
+
+@torch.library.register_fake(add_op_namespace_prefix("qk_rmsnorm_rope_bf16_"))
+def _qk_rmsnorm_rope_bf16_fake(
+    qk: torch.Tensor,
+    weight: torch.Tensor,
+    cos: torch.Tensor,
+    sin: torch.Tensor,
+    eps: float = 1e-6,
+) -> None:
+    if qk.dim() != 3 or weight.shape != (qk.shape[2],) or cos.shape != (qk.shape[0], qk.shape[2]) or sin.shape != cos.shape:
+        raise RuntimeError("qk_rmsnorm_rope_bf16_ expects qk (rows, heads, dim), weight (dim,), cos/sin (rows, dim)")
+    return None
+
+
+@torch.library.register_fake(add_op_namespace_prefix("qk_pair_rmsnorm_rope_bf16"))
+def _qk_pair_rmsnorm_rope_bf16_fake(
+    q: torch.Tensor,
+    k: torch.Tensor,
+    q_weight: torch.Tensor,
+    k_weight: torch.Tensor,
+    cos: torch.Tensor,
+    sin: torch.Tensor,
+    eps: float,
+    q_out: torch.Tensor,
+    k_out: torch.Tensor,
+) -> None:
+    if (
+        q.dim() != 3
+        or k.dim() != 3
+        or q.shape[0] != k.shape[0]
+        or q.shape[2] != k.shape[2]
+        or q.shape[2] < 8
+        or q.shape[2] > 256
+        or q.shape[2] % 2 != 0
+        or q_weight.shape != (q.shape[2],)
+        or k_weight.shape != q_weight.shape
+        or cos.shape != (q.shape[0], q.shape[2])
+        or sin.shape != cos.shape
+        or q_out.shape != q.shape
+        or k_out.shape != k.shape
+    ):
+        raise RuntimeError(
+            "qk_pair_rmsnorm_rope_bf16 expects q/k (rows, heads, even_dim), "
+            "weights (dim,), cos/sin (rows, dim), and matching outputs"
+        )
+    return None
+
+
+@torch.library.register_fake(add_op_namespace_prefix("gather_rows_bf16"))
+def _gather_rows_bf16_fake(
+    src: torch.Tensor,
+    row_indices: torch.Tensor,
+    dst: torch.Tensor,
+) -> None:
+    if src.dim() != 2 or row_indices.dim() != 1 or dst.shape != (row_indices.numel(), src.shape[1]):
+        raise RuntimeError(
+            "gather_rows_bf16 expects src (source_rows, hidden), "
+            "row_indices (rows,), dst (rows, hidden)"
+        )
+    return None
+
+
+@torch.library.register_fake(add_op_namespace_prefix("scatter_rows_bf16"))
+def _scatter_rows_bf16_fake(
+    src: torch.Tensor,
+    row_indices: torch.Tensor,
+    dst: torch.Tensor,
+) -> None:
+    if src.dim() != 2 or row_indices.dim() != 1 or src.shape != (row_indices.numel(), dst.shape[1]):
+        raise RuntimeError(
+            "scatter_rows_bf16 expects src (rows, hidden), "
+            "row_indices (rows,), dst (destination_rows, hidden)"
+        )
+    return None
+
+
+def fill_neginf_bf16(dst: torch.Tensor) -> torch.Tensor:
+    ops.fill_neginf_bf16(dst)
+    return dst
+
+
+def add_bias_bf16_(data: torch.Tensor, bias: torch.Tensor) -> torch.Tensor:
+    ops.add_bias_bf16_(data, bias)
+    return data
+
+
+def repeat_interleave_heads_bf16(
+    src: torch.Tensor,
+    repeat: int,
+    *,
+    out: Optional[torch.Tensor] = None,
+) -> torch.Tensor:
+    if out is None:
+        out = torch.empty((src.shape[0], src.shape[1] * repeat, src.shape[2]), device=src.device, dtype=src.dtype)
+    ops.repeat_interleave_heads_bf16(src, int(repeat), out)
+    return out
+
+
+def text_gather_bf16(
+    src: torch.Tensor,
+    batch: int,
+    seq: int,
+    *,
+    out: Optional[torch.Tensor] = None,
+) -> torch.Tensor:
+    if out is None:
+        out = torch.empty((2 * batch, src.shape[1]), device=src.device, dtype=src.dtype)
+    ops.text_gather_bf16(src, int(batch), int(seq), out)
+    return out
+
+
+def text_scatter_bf16(dst: torch.Tensor, src: torch.Tensor, batch: int, seq: int) -> torch.Tensor:
+    ops.text_scatter_bf16(dst, src, int(batch), int(seq))
+    return dst
+
+
+def rope_rotate_half_bf16_(x: torch.Tensor, cos: torch.Tensor, sin: torch.Tensor) -> torch.Tensor:
+    ops.rope_rotate_half_bf16_(x, cos, sin)
+    return x
+
+
+def qk_rmsnorm_rope_bf16_(
+    qk: torch.Tensor,
+    weight: torch.Tensor,
+    cos: torch.Tensor,
+    sin: torch.Tensor,
+    eps: float = 1e-6,
+) -> torch.Tensor:
+    ops.qk_rmsnorm_rope_bf16_(qk, weight, cos, sin, float(eps))
+    return qk
+
+
+def qk_pair_rmsnorm_rope_bf16(
+    q: torch.Tensor,
+    k: torch.Tensor,
+    q_weight: torch.Tensor,
+    k_weight: torch.Tensor,
+    cos: torch.Tensor,
+    sin: torch.Tensor,
+    eps: float = 1e-6,
+    *,
+    q_out: Optional[torch.Tensor] = None,
+    k_out: Optional[torch.Tensor] = None,
+) -> tuple[torch.Tensor, torch.Tensor]:
+    if q_out is None:
+        q_out = torch.empty_like(q)
+    if k_out is None:
+        k_out = torch.empty_like(k)
+    ops.qk_pair_rmsnorm_rope_bf16(
+        q, k, q_weight, k_weight, cos, sin, float(eps), q_out, k_out
+    )
+    return q_out, k_out
+
+
+def gather_rows_bf16(
+    src: torch.Tensor,
+    row_indices: torch.Tensor,
+    *,
+    out: Optional[torch.Tensor] = None,
+) -> torch.Tensor:
+    """Gather rows using in-range CUDA int64 indices without host synchronization."""
+
+    if out is None:
+        out = torch.empty(
+            (row_indices.numel(), src.shape[1]), device=src.device, dtype=src.dtype
+        )
+    ops.gather_rows_bf16(src, row_indices, out)
+    return out
+
+
+def scatter_rows_bf16(
+    src: torch.Tensor,
+    row_indices: torch.Tensor,
+    destination_rows: int,
+    *,
+    out: Optional[torch.Tensor] = None,
+) -> torch.Tensor:
+    """Scatter rows to unique, in-range CUDA int64 indices."""
+
+    if out is None:
+        out = torch.zeros(
+            (destination_rows, src.shape[1]), device=src.device, dtype=src.dtype
+        )
+    ops.scatter_rows_bf16(src, row_indices, out)
+    return out
+
+
+__all__ = [
+    "fill_neginf_bf16",
+    "add_bias_bf16_",
+    "repeat_interleave_heads_bf16",
+    "text_gather_bf16",
+    "text_scatter_bf16",
+    "rope_rotate_half_bf16_",
+    "qk_rmsnorm_rope_bf16_",
+    "qk_pair_rmsnorm_rope_bf16",
+    "gather_rows_bf16",
+    "scatter_rows_bf16",
+]

build/torch211-cxx11-cu130-aarch64-linux/_ops.py

@@ -0,0 +1,6 @@
+import torch
+from . import _transformer_layout_primitives_cuda_916f685
+ops = torch.ops._transformer_layout_primitives_cuda_916f685
+
+def add_op_namespace_prefix(op_name: str):
+    return f"_transformer_layout_primitives_cuda_916f685::{op_name}"

build/{torch211-cxx11-cu128-x86_64-linux/_transformer_layout_primitives_cuda_c9b2fa0.abi3.so → torch211-cxx11-cu130-aarch64-linux/_transformer_layout_primitives_cuda_916f685.abi3.so}

@@ -1,3 +1,3 @@
 version https://git-lfs.github.com/spec/v1
-oid sha256:8391a04a4384cc92d83c6df4c86ecdd205567c00105c13a01ca54c2b88fdd371
-size 572816
+oid sha256:c8756d718dd2b2d6ebb87618d7db67c81f18b842e7f6e4368b17e2bf60a1c9dc
+size 328480

build/torch211-cxx11-cu130-aarch64-linux/metadata.json

@@ -0,0 +1,22 @@
+{
+  "name": "transformer-layout-primitives",
+  "id": "_transformer_layout_primitives_cuda_916f685",
+  "version": 1,
+  "license": "Apache-2.0",
+  "python-depends": [],
+  "backend": {
+    "type": "cuda",
+    "archs": [
+      "11.0"
+    ]
+  },
+  "digest": {
+    "algorithm": "sha256",
+    "files": {
+      "__init__.py": "Le2NlMbv2iPKF5/hg87FKZwzjsVjK1dras7vbnPWYx8=",
+      "_transformer_layout_primitives_cuda_916f685.abi3.so": "yHVtcY3SstbruHYY19tnyB8YuELn9uQ2ixfiv2Chydw=",
+      "_ops.py": "YTOtZXEGILlA33lBYiZozOFNN4oQApx6STt61PAHcUc=",
+      "transformer_layout_primitives/__init__.py": "v6p5XMfQzddhi1fLSAw4HX9CyS0rQsidvu9VsT01xi4="
+    }
+  }
+}

build/torch211-cxx11-cu130-aarch64-linux/transformer_layout_primitives/__init__.py

@@ -0,0 +1,14 @@
+import ctypes
+import importlib.util
+import sys
+from pathlib import Path
+
+def _import_from_path(file_path: Path):
+    path_hash = '{:x}'.format(ctypes.c_size_t(hash(file_path.absolute())).value)
+    spec = importlib.util.spec_from_file_location(path_hash, file_path)
+    module = importlib.util.module_from_spec(spec)
+    sys.modules[path_hash] = module
+    spec.loader.exec_module(module)
+    return module
+
+globals().update(vars(_import_from_path(Path(__file__).parent.parent / '__init__.py')))

build/torch211-cxx11-cu130-x86_64-linux/__init__.py

@@ -55,13 +55,75 @@ def _qk_rmsnorm_rope_bf16_fake(
     weight: torch.Tensor,
     cos: torch.Tensor,
     sin: torch.Tensor,
-    eps: float,
+    eps: float = 1e-6,
 ) -> None:
     if qk.dim() != 3 or weight.shape != (qk.shape[2],) or cos.shape != (qk.shape[0], qk.shape[2]) or sin.shape != cos.shape:
         raise RuntimeError("qk_rmsnorm_rope_bf16_ expects qk (rows, heads, dim), weight (dim,), cos/sin (rows, dim)")
     return None
 
 
+@torch.library.register_fake(add_op_namespace_prefix("qk_pair_rmsnorm_rope_bf16"))
+def _qk_pair_rmsnorm_rope_bf16_fake(
+    q: torch.Tensor,
+    k: torch.Tensor,
+    q_weight: torch.Tensor,
+    k_weight: torch.Tensor,
+    cos: torch.Tensor,
+    sin: torch.Tensor,
+    eps: float,
+    q_out: torch.Tensor,
+    k_out: torch.Tensor,
+) -> None:
+    if (
+        q.dim() != 3
+        or k.dim() != 3
+        or q.shape[0] != k.shape[0]
+        or q.shape[2] != k.shape[2]
+        or q.shape[2] < 8
+        or q.shape[2] > 256
+        or q.shape[2] % 2 != 0
+        or q_weight.shape != (q.shape[2],)
+        or k_weight.shape != q_weight.shape
+        or cos.shape != (q.shape[0], q.shape[2])
+        or sin.shape != cos.shape
+        or q_out.shape != q.shape
+        or k_out.shape != k.shape
+    ):
+        raise RuntimeError(
+            "qk_pair_rmsnorm_rope_bf16 expects q/k (rows, heads, even_dim), "
+            "weights (dim,), cos/sin (rows, dim), and matching outputs"
+        )
+    return None
+
+
+@torch.library.register_fake(add_op_namespace_prefix("gather_rows_bf16"))
+def _gather_rows_bf16_fake(
+    src: torch.Tensor,
+    row_indices: torch.Tensor,
+    dst: torch.Tensor,
+) -> None:
+    if src.dim() != 2 or row_indices.dim() != 1 or dst.shape != (row_indices.numel(), src.shape[1]):
+        raise RuntimeError(
+            "gather_rows_bf16 expects src (source_rows, hidden), "
+            "row_indices (rows,), dst (rows, hidden)"
+        )
+    return None
+
+
+@torch.library.register_fake(add_op_namespace_prefix("scatter_rows_bf16"))
+def _scatter_rows_bf16_fake(
+    src: torch.Tensor,
+    row_indices: torch.Tensor,
+    dst: torch.Tensor,
+) -> None:
+    if src.dim() != 2 or row_indices.dim() != 1 or src.shape != (row_indices.numel(), dst.shape[1]):
+        raise RuntimeError(
+            "scatter_rows_bf16 expects src (rows, hidden), "
+            "row_indices (rows,), dst (destination_rows, hidden)"
+        )
+    return None
+
+
 def fill_neginf_bf16(dst: torch.Tensor) -> torch.Tensor:
     ops.fill_neginf_bf16(dst)
     return dst
@@ -118,6 +180,61 @@ def qk_rmsnorm_rope_bf16_(
     return qk
 
 
+def qk_pair_rmsnorm_rope_bf16(
+    q: torch.Tensor,
+    k: torch.Tensor,
+    q_weight: torch.Tensor,
+    k_weight: torch.Tensor,
+    cos: torch.Tensor,
+    sin: torch.Tensor,
+    eps: float = 1e-6,
+    *,
+    q_out: Optional[torch.Tensor] = None,
+    k_out: Optional[torch.Tensor] = None,
+) -> tuple[torch.Tensor, torch.Tensor]:
+    if q_out is None:
+        q_out = torch.empty_like(q)
+    if k_out is None:
+        k_out = torch.empty_like(k)
+    ops.qk_pair_rmsnorm_rope_bf16(
+        q, k, q_weight, k_weight, cos, sin, float(eps), q_out, k_out
+    )
+    return q_out, k_out
+
+
+def gather_rows_bf16(
+    src: torch.Tensor,
+    row_indices: torch.Tensor,
+    *,
+    out: Optional[torch.Tensor] = None,
+) -> torch.Tensor:
+    """Gather rows using in-range CUDA int64 indices without host synchronization."""
+
+    if out is None:
+        out = torch.empty(
+            (row_indices.numel(), src.shape[1]), device=src.device, dtype=src.dtype
+        )
+    ops.gather_rows_bf16(src, row_indices, out)
+    return out
+
+
+def scatter_rows_bf16(
+    src: torch.Tensor,
+    row_indices: torch.Tensor,
+    destination_rows: int,
+    *,
+    out: Optional[torch.Tensor] = None,
+) -> torch.Tensor:
+    """Scatter rows to unique, in-range CUDA int64 indices."""
+
+    if out is None:
+        out = torch.zeros(
+            (destination_rows, src.shape[1]), device=src.device, dtype=src.dtype
+        )
+    ops.scatter_rows_bf16(src, row_indices, out)
+    return out
+
+
 __all__ = [
     "fill_neginf_bf16",
     "add_bias_bf16_",
@@ -126,4 +243,7 @@ __all__ = [
     "text_scatter_bf16",
     "rope_rotate_half_bf16_",
     "qk_rmsnorm_rope_bf16_",
+    "qk_pair_rmsnorm_rope_bf16",
+    "gather_rows_bf16",
+    "scatter_rows_bf16",
 ]

build/torch211-cxx11-cu130-x86_64-linux/_ops.py

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

build/torch211-cxx11-cu130-x86_64-linux/{_transformer_layout_primitives_cuda_c9b2fa0.abi3.so → _transformer_layout_primitives_cuda_0a2c98e.abi3.so}

@@ -1,3 +1,3 @@
 version https://git-lfs.github.com/spec/v1
-oid sha256:99943084af51325b92b740243b59d224141161ee451749776e1c89ce7cceded2
-size 564160
+oid sha256:f184cbedf680649ed33eed7aeab4e5421ebb99ec37323b0cf21b57577c9d594d
+size 803048

build/torch211-cxx11-cu130-x86_64-linux/metadata.json

@@ -1,6 +1,6 @@
 {
   "name": "transformer-layout-primitives",
-  "id": "_transformer_layout_primitives_cuda_c9b2fa0",
+  "id": "_transformer_layout_primitives_cuda_0a2c98e",
   "version": 1,
   "license": "Apache-2.0",
   "python-depends": [],

build/torch212-cxx11-cu130-x86_64-linux/__init__.py

@@ -55,13 +55,75 @@ def _qk_rmsnorm_rope_bf16_fake(
     weight: torch.Tensor,
     cos: torch.Tensor,
     sin: torch.Tensor,
-    eps: float,
+    eps: float = 1e-6,
 ) -> None:
     if qk.dim() != 3 or weight.shape != (qk.shape[2],) or cos.shape != (qk.shape[0], qk.shape[2]) or sin.shape != cos.shape:
         raise RuntimeError("qk_rmsnorm_rope_bf16_ expects qk (rows, heads, dim), weight (dim,), cos/sin (rows, dim)")
     return None
 
 
+@torch.library.register_fake(add_op_namespace_prefix("qk_pair_rmsnorm_rope_bf16"))
+def _qk_pair_rmsnorm_rope_bf16_fake(
+    q: torch.Tensor,
+    k: torch.Tensor,
+    q_weight: torch.Tensor,
+    k_weight: torch.Tensor,
+    cos: torch.Tensor,
+    sin: torch.Tensor,
+    eps: float,
+    q_out: torch.Tensor,
+    k_out: torch.Tensor,
+) -> None:
+    if (
+        q.dim() != 3
+        or k.dim() != 3
+        or q.shape[0] != k.shape[0]
+        or q.shape[2] != k.shape[2]
+        or q.shape[2] < 8
+        or q.shape[2] > 256
+        or q.shape[2] % 2 != 0
+        or q_weight.shape != (q.shape[2],)
+        or k_weight.shape != q_weight.shape
+        or cos.shape != (q.shape[0], q.shape[2])
+        or sin.shape != cos.shape
+        or q_out.shape != q.shape
+        or k_out.shape != k.shape
+    ):
+        raise RuntimeError(
+            "qk_pair_rmsnorm_rope_bf16 expects q/k (rows, heads, even_dim), "
+            "weights (dim,), cos/sin (rows, dim), and matching outputs"
+        )
+    return None
+
+
+@torch.library.register_fake(add_op_namespace_prefix("gather_rows_bf16"))
+def _gather_rows_bf16_fake(
+    src: torch.Tensor,
+    row_indices: torch.Tensor,
+    dst: torch.Tensor,
+) -> None:
+    if src.dim() != 2 or row_indices.dim() != 1 or dst.shape != (row_indices.numel(), src.shape[1]):
+        raise RuntimeError(
+            "gather_rows_bf16 expects src (source_rows, hidden), "
+            "row_indices (rows,), dst (rows, hidden)"
+        )
+    return None
+
+
+@torch.library.register_fake(add_op_namespace_prefix("scatter_rows_bf16"))
+def _scatter_rows_bf16_fake(
+    src: torch.Tensor,
+    row_indices: torch.Tensor,
+    dst: torch.Tensor,
+) -> None:
+    if src.dim() != 2 or row_indices.dim() != 1 or src.shape != (row_indices.numel(), dst.shape[1]):
+        raise RuntimeError(
+            "scatter_rows_bf16 expects src (rows, hidden), "
+            "row_indices (rows,), dst (destination_rows, hidden)"
+        )
+    return None
+
+
 def fill_neginf_bf16(dst: torch.Tensor) -> torch.Tensor:
     ops.fill_neginf_bf16(dst)
     return dst
@@ -118,6 +180,61 @@ def qk_rmsnorm_rope_bf16_(
     return qk
 
 
+def qk_pair_rmsnorm_rope_bf16(
+    q: torch.Tensor,
+    k: torch.Tensor,
+    q_weight: torch.Tensor,
+    k_weight: torch.Tensor,
+    cos: torch.Tensor,
+    sin: torch.Tensor,
+    eps: float = 1e-6,
+    *,
+    q_out: Optional[torch.Tensor] = None,
+    k_out: Optional[torch.Tensor] = None,
+) -> tuple[torch.Tensor, torch.Tensor]:
+    if q_out is None:
+        q_out = torch.empty_like(q)
+    if k_out is None:
+        k_out = torch.empty_like(k)
+    ops.qk_pair_rmsnorm_rope_bf16(
+        q, k, q_weight, k_weight, cos, sin, float(eps), q_out, k_out
+    )
+    return q_out, k_out
+
+
+def gather_rows_bf16(
+    src: torch.Tensor,
+    row_indices: torch.Tensor,
+    *,
+    out: Optional[torch.Tensor] = None,
+) -> torch.Tensor:
+    """Gather rows using in-range CUDA int64 indices without host synchronization."""
+
+    if out is None:
+        out = torch.empty(
+            (row_indices.numel(), src.shape[1]), device=src.device, dtype=src.dtype
+        )
+    ops.gather_rows_bf16(src, row_indices, out)
+    return out
+
+
+def scatter_rows_bf16(
+    src: torch.Tensor,
+    row_indices: torch.Tensor,
+    destination_rows: int,
+    *,
+    out: Optional[torch.Tensor] = None,
+) -> torch.Tensor:
+    """Scatter rows to unique, in-range CUDA int64 indices."""
+
+    if out is None:
+        out = torch.zeros(
+            (destination_rows, src.shape[1]), device=src.device, dtype=src.dtype
+        )
+    ops.scatter_rows_bf16(src, row_indices, out)
+    return out
+
+
 __all__ = [
     "fill_neginf_bf16",
     "add_bias_bf16_",
@@ -126,4 +243,7 @@ __all__ = [
     "text_scatter_bf16",
     "rope_rotate_half_bf16_",
     "qk_rmsnorm_rope_bf16_",
+    "qk_pair_rmsnorm_rope_bf16",
+    "gather_rows_bf16",
+    "scatter_rows_bf16",
 ]

build/torch212-cxx11-cu130-x86_64-linux/_ops.py

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

build/torch212-cxx11-cu130-x86_64-linux/{_transformer_layout_primitives_cuda_c9b2fa0.abi3.so → _transformer_layout_primitives_cuda_0a2c98e.abi3.so}

@@ -1,3 +1,3 @@
 version https://git-lfs.github.com/spec/v1
-oid sha256:6c4ebdc5e5faa851f8c4f537f5a1e75e28cebad98a650c5bf72dbf60db981844
-size 570424
+oid sha256:cb4fd9c8372d2543efd339d10b2787539fc96db7368a509e175d001882d51e20
+size 809400

build/torch212-cxx11-cu130-x86_64-linux/metadata.json

@@ -1,6 +1,6 @@
 {
   "name": "transformer-layout-primitives",
-  "id": "_transformer_layout_primitives_cuda_c9b2fa0",
+  "id": "_transformer_layout_primitives_cuda_0a2c98e",
   "version": 1,
   "license": "Apache-2.0",
   "python-depends": [],

build/torch212-cxx11-cu132-x86_64-linux/__init__.py

@@ -55,13 +55,75 @@ def _qk_rmsnorm_rope_bf16_fake(
     weight: torch.Tensor,
     cos: torch.Tensor,
     sin: torch.Tensor,
-    eps: float,
+    eps: float = 1e-6,
 ) -> None:
     if qk.dim() != 3 or weight.shape != (qk.shape[2],) or cos.shape != (qk.shape[0], qk.shape[2]) or sin.shape != cos.shape:
         raise RuntimeError("qk_rmsnorm_rope_bf16_ expects qk (rows, heads, dim), weight (dim,), cos/sin (rows, dim)")
     return None
 
 
+@torch.library.register_fake(add_op_namespace_prefix("qk_pair_rmsnorm_rope_bf16"))
+def _qk_pair_rmsnorm_rope_bf16_fake(
+    q: torch.Tensor,
+    k: torch.Tensor,
+    q_weight: torch.Tensor,
+    k_weight: torch.Tensor,
+    cos: torch.Tensor,
+    sin: torch.Tensor,
+    eps: float,
+    q_out: torch.Tensor,
+    k_out: torch.Tensor,
+) -> None:
+    if (
+        q.dim() != 3
+        or k.dim() != 3
+        or q.shape[0] != k.shape[0]
+        or q.shape[2] != k.shape[2]
+        or q.shape[2] < 8
+        or q.shape[2] > 256
+        or q.shape[2] % 2 != 0
+        or q_weight.shape != (q.shape[2],)
+        or k_weight.shape != q_weight.shape
+        or cos.shape != (q.shape[0], q.shape[2])
+        or sin.shape != cos.shape
+        or q_out.shape != q.shape
+        or k_out.shape != k.shape
+    ):
+        raise RuntimeError(
+            "qk_pair_rmsnorm_rope_bf16 expects q/k (rows, heads, even_dim), "
+            "weights (dim,), cos/sin (rows, dim), and matching outputs"
+        )
+    return None
+
+
+@torch.library.register_fake(add_op_namespace_prefix("gather_rows_bf16"))
+def _gather_rows_bf16_fake(
+    src: torch.Tensor,
+    row_indices: torch.Tensor,
+    dst: torch.Tensor,
+) -> None:
+    if src.dim() != 2 or row_indices.dim() != 1 or dst.shape != (row_indices.numel(), src.shape[1]):
+        raise RuntimeError(
+            "gather_rows_bf16 expects src (source_rows, hidden), "
+            "row_indices (rows,), dst (rows, hidden)"
+        )
+    return None
+
+
+@torch.library.register_fake(add_op_namespace_prefix("scatter_rows_bf16"))
+def _scatter_rows_bf16_fake(
+    src: torch.Tensor,
+    row_indices: torch.Tensor,
+    dst: torch.Tensor,
+) -> None:
+    if src.dim() != 2 or row_indices.dim() != 1 or src.shape != (row_indices.numel(), dst.shape[1]):
+        raise RuntimeError(
+            "scatter_rows_bf16 expects src (rows, hidden), "
+            "row_indices (rows,), dst (destination_rows, hidden)"
+        )
+    return None
+
+
 def fill_neginf_bf16(dst: torch.Tensor) -> torch.Tensor:
     ops.fill_neginf_bf16(dst)
     return dst
@@ -118,6 +180,61 @@ def qk_rmsnorm_rope_bf16_(
     return qk
 
 
+def qk_pair_rmsnorm_rope_bf16(
+    q: torch.Tensor,
+    k: torch.Tensor,
+    q_weight: torch.Tensor,
+    k_weight: torch.Tensor,
+    cos: torch.Tensor,
+    sin: torch.Tensor,
+    eps: float = 1e-6,
+    *,
+    q_out: Optional[torch.Tensor] = None,
+    k_out: Optional[torch.Tensor] = None,
+) -> tuple[torch.Tensor, torch.Tensor]:
+    if q_out is None:
+        q_out = torch.empty_like(q)
+    if k_out is None:
+        k_out = torch.empty_like(k)
+    ops.qk_pair_rmsnorm_rope_bf16(
+        q, k, q_weight, k_weight, cos, sin, float(eps), q_out, k_out
+    )
+    return q_out, k_out
+
+
+def gather_rows_bf16(
+    src: torch.Tensor,
+    row_indices: torch.Tensor,
+    *,
+    out: Optional[torch.Tensor] = None,
+) -> torch.Tensor:
+    """Gather rows using in-range CUDA int64 indices without host synchronization."""
+
+    if out is None:
+        out = torch.empty(
+            (row_indices.numel(), src.shape[1]), device=src.device, dtype=src.dtype
+        )
+    ops.gather_rows_bf16(src, row_indices, out)
+    return out
+
+
+def scatter_rows_bf16(
+    src: torch.Tensor,
+    row_indices: torch.Tensor,
+    destination_rows: int,
+    *,
+    out: Optional[torch.Tensor] = None,
+) -> torch.Tensor:
+    """Scatter rows to unique, in-range CUDA int64 indices."""
+
+    if out is None:
+        out = torch.zeros(
+            (destination_rows, src.shape[1]), device=src.device, dtype=src.dtype
+        )
+    ops.scatter_rows_bf16(src, row_indices, out)
+    return out
+
+
 __all__ = [
     "fill_neginf_bf16",
     "add_bias_bf16_",
@@ -126,4 +243,7 @@ __all__ = [
     "text_scatter_bf16",
     "rope_rotate_half_bf16_",
     "qk_rmsnorm_rope_bf16_",
+    "qk_pair_rmsnorm_rope_bf16",
+    "gather_rows_bf16",
+    "scatter_rows_bf16",
 ]

build/torch212-cxx11-cu132-x86_64-linux/_ops.py

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

build/torch212-cxx11-cu132-x86_64-linux/_transformer_layout_primitives_cuda_0a2c98e.abi3.so

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

build/torch212-cxx11-cu132-x86_64-linux/metadata.json

@@ -1,6 +1,6 @@
 {
   "name": "transformer-layout-primitives",
-  "id": "_transformer_layout_primitives_cuda_c9b2fa0",
+  "id": "_transformer_layout_primitives_cuda_0a2c98e",
   "version": 1,
   "license": "Apache-2.0",
   "python-depends": [],

build/torch213-cxx11-cu130-aarch64-linux/__init__.py

@@ -0,0 +1,249 @@
+"""FlashRT transformer layout primitives."""
+
+from __future__ import annotations
+
+from typing import Optional
+
+import torch
+
+from ._ops import add_op_namespace_prefix, ops
+
+
+@torch.library.register_fake(add_op_namespace_prefix("fill_neginf_bf16"))
+def _fill_neginf_bf16_fake(dst: torch.Tensor) -> None:
+    return None
+
+
+@torch.library.register_fake(add_op_namespace_prefix("add_bias_bf16_"))
+def _add_bias_bf16_fake(data: torch.Tensor, bias: torch.Tensor) -> None:
+    if data.dim() != 2 or bias.shape != (data.shape[1],):
+        raise RuntimeError("add_bias_bf16_ expects data (rows, cols), bias (cols,)")
+    return None
+
+
+@torch.library.register_fake(add_op_namespace_prefix("repeat_interleave_heads_bf16"))
+def _repeat_interleave_heads_bf16_fake(src: torch.Tensor, repeat: int, dst: torch.Tensor) -> None:
+    if src.dim() != 3 or dst.shape != (src.shape[0], src.shape[1] * repeat, src.shape[2]):
+        raise RuntimeError("repeat_interleave_heads_bf16 expects src (seq, heads, dim), dst (seq, heads*repeat, dim)")
+    return None
+
+
+@torch.library.register_fake(add_op_namespace_prefix("text_gather_bf16"))
+def _text_gather_bf16_fake(src: torch.Tensor, batch: int, seq: int, dst: torch.Tensor) -> None:
+    if src.dim() != 2 or src.shape[0] != batch * seq or dst.shape != (2 * batch, src.shape[1]):
+        raise RuntimeError("text_gather_bf16 expects src (batch*seq, dim), dst (2*batch, dim)")
+    return None
+
+
+@torch.library.register_fake(add_op_namespace_prefix("text_scatter_bf16"))
+def _text_scatter_bf16_fake(dst: torch.Tensor, src: torch.Tensor, batch: int, seq: int) -> None:
+    if dst.dim() != 2 or dst.shape[0] != batch * seq or src.shape != (2 * batch, dst.shape[1]):
+        raise RuntimeError("text_scatter_bf16 expects dst (batch*seq, dim), src (2*batch, dim)")
+    return None
+
+
+@torch.library.register_fake(add_op_namespace_prefix("rope_rotate_half_bf16_"))
+def _rope_rotate_half_bf16_fake(x: torch.Tensor, cos: torch.Tensor, sin: torch.Tensor) -> None:
+    if x.dim() != 3 or x.shape[2] % 2 != 0 or cos.shape != (x.shape[0], x.shape[2]) or sin.shape != cos.shape:
+        raise RuntimeError("rope_rotate_half_bf16_ expects x (seq, heads, even_dim), cos/sin (seq, dim)")
+    return None
+
+
+@torch.library.register_fake(add_op_namespace_prefix("qk_rmsnorm_rope_bf16_"))
+def _qk_rmsnorm_rope_bf16_fake(
+    qk: torch.Tensor,
+    weight: torch.Tensor,
+    cos: torch.Tensor,
+    sin: torch.Tensor,
+    eps: float = 1e-6,
+) -> None:
+    if qk.dim() != 3 or weight.shape != (qk.shape[2],) or cos.shape != (qk.shape[0], qk.shape[2]) or sin.shape != cos.shape:
+        raise RuntimeError("qk_rmsnorm_rope_bf16_ expects qk (rows, heads, dim), weight (dim,), cos/sin (rows, dim)")
+    return None
+
+
+@torch.library.register_fake(add_op_namespace_prefix("qk_pair_rmsnorm_rope_bf16"))
+def _qk_pair_rmsnorm_rope_bf16_fake(
+    q: torch.Tensor,
+    k: torch.Tensor,
+    q_weight: torch.Tensor,
+    k_weight: torch.Tensor,
+    cos: torch.Tensor,
+    sin: torch.Tensor,
+    eps: float,
+    q_out: torch.Tensor,
+    k_out: torch.Tensor,
+) -> None:
+    if (
+        q.dim() != 3
+        or k.dim() != 3
+        or q.shape[0] != k.shape[0]
+        or q.shape[2] != k.shape[2]
+        or q.shape[2] < 8
+        or q.shape[2] > 256
+        or q.shape[2] % 2 != 0
+        or q_weight.shape != (q.shape[2],)
+        or k_weight.shape != q_weight.shape
+        or cos.shape != (q.shape[0], q.shape[2])
+        or sin.shape != cos.shape
+        or q_out.shape != q.shape
+        or k_out.shape != k.shape
+    ):
+        raise RuntimeError(
+            "qk_pair_rmsnorm_rope_bf16 expects q/k (rows, heads, even_dim), "
+            "weights (dim,), cos/sin (rows, dim), and matching outputs"
+        )
+    return None
+
+
+@torch.library.register_fake(add_op_namespace_prefix("gather_rows_bf16"))
+def _gather_rows_bf16_fake(
+    src: torch.Tensor,
+    row_indices: torch.Tensor,
+    dst: torch.Tensor,
+) -> None:
+    if src.dim() != 2 or row_indices.dim() != 1 or dst.shape != (row_indices.numel(), src.shape[1]):
+        raise RuntimeError(
+            "gather_rows_bf16 expects src (source_rows, hidden), "
+            "row_indices (rows,), dst (rows, hidden)"
+        )
+    return None
+
+
+@torch.library.register_fake(add_op_namespace_prefix("scatter_rows_bf16"))
+def _scatter_rows_bf16_fake(
+    src: torch.Tensor,
+    row_indices: torch.Tensor,
+    dst: torch.Tensor,
+) -> None:
+    if src.dim() != 2 or row_indices.dim() != 1 or src.shape != (row_indices.numel(), dst.shape[1]):
+        raise RuntimeError(
+            "scatter_rows_bf16 expects src (rows, hidden), "
+            "row_indices (rows,), dst (destination_rows, hidden)"
+        )
+    return None
+
+
+def fill_neginf_bf16(dst: torch.Tensor) -> torch.Tensor:
+    ops.fill_neginf_bf16(dst)
+    return dst
+
+
+def add_bias_bf16_(data: torch.Tensor, bias: torch.Tensor) -> torch.Tensor:
+    ops.add_bias_bf16_(data, bias)
+    return data
+
+
+def repeat_interleave_heads_bf16(
+    src: torch.Tensor,
+    repeat: int,
+    *,
+    out: Optional[torch.Tensor] = None,
+) -> torch.Tensor:
+    if out is None:
+        out = torch.empty((src.shape[0], src.shape[1] * repeat, src.shape[2]), device=src.device, dtype=src.dtype)
+    ops.repeat_interleave_heads_bf16(src, int(repeat), out)
+    return out
+
+
+def text_gather_bf16(
+    src: torch.Tensor,
+    batch: int,
+    seq: int,
+    *,
+    out: Optional[torch.Tensor] = None,
+) -> torch.Tensor:
+    if out is None:
+        out = torch.empty((2 * batch, src.shape[1]), device=src.device, dtype=src.dtype)
+    ops.text_gather_bf16(src, int(batch), int(seq), out)
+    return out
+
+
+def text_scatter_bf16(dst: torch.Tensor, src: torch.Tensor, batch: int, seq: int) -> torch.Tensor:
+    ops.text_scatter_bf16(dst, src, int(batch), int(seq))
+    return dst
+
+
+def rope_rotate_half_bf16_(x: torch.Tensor, cos: torch.Tensor, sin: torch.Tensor) -> torch.Tensor:
+    ops.rope_rotate_half_bf16_(x, cos, sin)
+    return x
+
+
+def qk_rmsnorm_rope_bf16_(
+    qk: torch.Tensor,
+    weight: torch.Tensor,
+    cos: torch.Tensor,
+    sin: torch.Tensor,
+    eps: float = 1e-6,
+) -> torch.Tensor:
+    ops.qk_rmsnorm_rope_bf16_(qk, weight, cos, sin, float(eps))
+    return qk
+
+
+def qk_pair_rmsnorm_rope_bf16(
+    q: torch.Tensor,
+    k: torch.Tensor,
+    q_weight: torch.Tensor,
+    k_weight: torch.Tensor,
+    cos: torch.Tensor,
+    sin: torch.Tensor,
+    eps: float = 1e-6,
+    *,
+    q_out: Optional[torch.Tensor] = None,
+    k_out: Optional[torch.Tensor] = None,
+) -> tuple[torch.Tensor, torch.Tensor]:
+    if q_out is None:
+        q_out = torch.empty_like(q)
+    if k_out is None:
+        k_out = torch.empty_like(k)
+    ops.qk_pair_rmsnorm_rope_bf16(
+        q, k, q_weight, k_weight, cos, sin, float(eps), q_out, k_out
+    )
+    return q_out, k_out
+
+
+def gather_rows_bf16(
+    src: torch.Tensor,
+    row_indices: torch.Tensor,
+    *,
+    out: Optional[torch.Tensor] = None,
+) -> torch.Tensor:
+    """Gather rows using in-range CUDA int64 indices without host synchronization."""
+
+    if out is None:
+        out = torch.empty(
+            (row_indices.numel(), src.shape[1]), device=src.device, dtype=src.dtype
+        )
+    ops.gather_rows_bf16(src, row_indices, out)
+    return out
+
+
+def scatter_rows_bf16(
+    src: torch.Tensor,
+    row_indices: torch.Tensor,
+    destination_rows: int,
+    *,
+    out: Optional[torch.Tensor] = None,
+) -> torch.Tensor:
+    """Scatter rows to unique, in-range CUDA int64 indices."""
+
+    if out is None:
+        out = torch.zeros(
+            (destination_rows, src.shape[1]), device=src.device, dtype=src.dtype
+        )
+    ops.scatter_rows_bf16(src, row_indices, out)
+    return out
+
+
+__all__ = [
+    "fill_neginf_bf16",
+    "add_bias_bf16_",
+    "repeat_interleave_heads_bf16",
+    "text_gather_bf16",
+    "text_scatter_bf16",
+    "rope_rotate_half_bf16_",
+    "qk_rmsnorm_rope_bf16_",
+    "qk_pair_rmsnorm_rope_bf16",
+    "gather_rows_bf16",
+    "scatter_rows_bf16",
+]

build/torch213-cxx11-cu130-aarch64-linux/_ops.py

@@ -0,0 +1,6 @@
+import torch
+from . import transformer_layout_primitives_source_test
+ops = torch.ops.transformer_layout_primitives_source_test
+
+def add_op_namespace_prefix(op_name: str):
+    return f"transformer_layout_primitives_source_test::{op_name}"

build/torch213-cxx11-cu130-aarch64-linux/metadata.json

@@ -0,0 +1,32 @@
+{
+  "name": "transformer-layout-primitives",
+  "id": "transformer_layout_primitives_source_test",
+  "version": 1,
+  "license": "Apache-2.0",
+  "python-depends": [],
+  "backend": {
+    "type": "cuda",
+    "archs": [
+      "11.0"
+    ]
+  },
+  "digest": {
+    "algorithm": "sha256",
+    "files": {
+      "__init__.py": "Le2NlMbv2iPKF5/hg87FKZwzjsVjK1dras7vbnPWYx8=",
+      "transformer_layout_primitives_source_test.abi3.so": "rTS3hp8n1EQCUhT21GetFkjokJSYi7tWSDu8POFV274=",
+      "_ops.py": "YSImh8cFlOo35GBoY1uX2AQlaapTWG5MigCcM3MXx6U=",
+      "transformer_layout_primitives/__init__.py": "v6p5XMfQzddhi1fLSAw4HX9CyS0rQsidvu9VsT01xi4="
+    }
+  },
+  "provenance": {
+    "kernel": {
+      "sha": "456d297",
+      "dirty": false
+    },
+    "validation": {
+      "torch": "2.13.0+cu130",
+      "cuda": "13.0"
+    }
+  }
+}

build/torch213-cxx11-cu130-aarch64-linux/transformer_layout_primitives/__init__.py

@@ -0,0 +1,14 @@
+import ctypes
+import importlib.util
+import sys
+from pathlib import Path
+
+def _import_from_path(file_path: Path):
+    path_hash = '{:x}'.format(ctypes.c_size_t(hash(file_path.absolute())).value)
+    spec = importlib.util.spec_from_file_location(path_hash, file_path)
+    module = importlib.util.module_from_spec(spec)
+    sys.modules[path_hash] = module
+    spec.loader.exec_module(module)
+    return module
+
+globals().update(vars(_import_from_path(Path(__file__).parent.parent / '__init__.py')))

build/torch213-cxx11-cu130-aarch64-linux/transformer_layout_primitives_source_test.abi3.so

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