liangsu9988 commited on Jul 7

Commit

072b453

verified ·

1 Parent(s): 94da621

Uploaded using `kernel-builder`.

Browse files

Files changed (21) hide show

benchmarks/benchmark.py +122 -0
build/torch211-cxx11-cu128-x86_64-linux/__init__.py +129 -0
build/torch211-cxx11-cu128-x86_64-linux/_ops.py +9 -0
build/torch211-cxx11-cu128-x86_64-linux/_transformer_layout_primitives_cuda_c9b2fa0.abi3.so +3 -0
build/torch211-cxx11-cu128-x86_64-linux/metadata.json +23 -0
build/torch211-cxx11-cu128-x86_64-linux/transformer_layout_primitives/__init__.py +26 -0
build/torch211-cxx11-cu130-x86_64-linux/__init__.py +129 -0
build/torch211-cxx11-cu130-x86_64-linux/_ops.py +9 -0
build/torch211-cxx11-cu130-x86_64-linux/_transformer_layout_primitives_cuda_c9b2fa0.abi3.so +3 -0
build/torch211-cxx11-cu130-x86_64-linux/metadata.json +22 -0
build/torch211-cxx11-cu130-x86_64-linux/transformer_layout_primitives/__init__.py +26 -0
build/torch212-cxx11-cu130-x86_64-linux/__init__.py +129 -0
build/torch212-cxx11-cu130-x86_64-linux/_ops.py +9 -0
build/torch212-cxx11-cu130-x86_64-linux/_transformer_layout_primitives_cuda_c9b2fa0.abi3.so +3 -0
build/torch212-cxx11-cu130-x86_64-linux/metadata.json +22 -0
build/torch212-cxx11-cu130-x86_64-linux/transformer_layout_primitives/__init__.py +26 -0
build/torch212-cxx11-cu132-x86_64-linux/__init__.py +129 -0
build/torch212-cxx11-cu132-x86_64-linux/_ops.py +9 -0
build/torch212-cxx11-cu132-x86_64-linux/_transformer_layout_primitives_cuda_c9b2fa0.abi3.so +3 -0
build/torch212-cxx11-cu132-x86_64-linux/metadata.json +22 -0
build/torch212-cxx11-cu132-x86_64-linux/transformer_layout_primitives/__init__.py +26 -0

benchmarks/benchmark.py ADDED Viewed

	@@ -0,0 +1,122 @@

+#!/usr/bin/env python3
+"""Benchmark transformer layout primitives."""
+from __future__ import annotations
+import argparse
+import importlib
+import sys
+from pathlib import Path
+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
+def load_ops(backend: str, artifact: str | None):
+    if backend == "source":
+        return load_source_ops()
+    if artifact:
+        sys.path.insert(0, artifact)
+    try:
+        return importlib.import_module("transformer_layout_primitives")
+    finally:
+        if artifact:
+            sys.path.remove(artifact)
+def time_us(fn, warmup: int, iters: int) -> float:
+    for _ in range(warmup):
+        fn()
+    torch.cuda.synchronize()
+    start = torch.cuda.Event(enable_timing=True)
+    end = torch.cuda.Event(enable_timing=True)
+    start.record()
+    for _ in range(iters):
+        fn()
+    end.record()
+    torch.cuda.synchronize()
+    return start.elapsed_time(end) * 1000.0 / iters
+def main() -> int:
+    parser = argparse.ArgumentParser()
+    parser.add_argument("--backend", choices=["source", "installed"], default="source")
+    parser.add_argument("--artifact", default=None)
+    parser.add_argument("--mode", choices=["headline", "full"], default="headline")
+    parser.add_argument("--warmup", type=int, default=20)
+    parser.add_argument("--iters", type=int, default=100)
+    args = parser.parse_args()
+    ops = load_ops(args.backend, args.artifact)
+    print("workload,shape,op,flashrt_us,torch_eager_us,speedup")
+    repeat_shapes = [("gqa_prefill", 2520, 8, 128, 4), ("decode_gqa", 1, 8, 128, 4)]
+    if args.mode == "full":
+        repeat_shapes += [("short_prefill", 128, 8, 128, 4), ("vl_prefill", 4096, 8, 128, 4)]
+    for name, seq, heads, dim, repeat in repeat_shapes:
+        src = torch.randn((seq, heads, dim), device="cuda", dtype=torch.bfloat16)
+        out = torch.empty((seq, heads * repeat, dim), device="cuda", dtype=torch.bfloat16)
+        def flash():
+            ops.repeat_interleave_heads_bf16(src, repeat, out=out)
+        def eager():
+            src.repeat_interleave(repeat, dim=1)
+        fu = time_us(flash, args.warmup, args.iters)
+        eu = time_us(eager, max(5, args.warmup // 2), max(20, args.iters // 2))
+        print(f"{name},{seq}x{heads}x{dim}x{repeat},repeat_interleave_heads_bf16,{fu:.3f},{eu:.3f},{eu/fu:.2f}x")
+    rope_shapes = [("qwen_prefill", 4096, 32, 128), ("video_prefill", 2520, 24, 128)]
+    if args.mode == "full":
+        rope_shapes += [("decode", 1, 32, 128), ("short_prefill", 128, 32, 128)]
+    for name, seq, heads, dim in rope_shapes:
+        x = torch.randn((seq, heads, dim), device="cuda", dtype=torch.bfloat16)
+        weight = torch.randn((dim,), device="cuda", dtype=torch.bfloat16)
+        cos = torch.randn((seq, dim), device="cuda", dtype=torch.bfloat16)
+        sin = torch.randn((seq, dim), device="cuda", dtype=torch.bfloat16)
+        out = x.clone()
+        def flash():
+            out.copy_(x)
+            ops.qk_rmsnorm_rope_bf16_(out, weight, cos, sin)
+        def eager():
+            qk_rmsnorm_rope_ref(x, weight, cos, sin)
+        fu = time_us(flash, args.warmup, args.iters)
+        eu = time_us(eager, max(5, args.warmup // 2), max(20, args.iters // 2))
+        print(f"{name},{seq}x{heads}x{dim},qk_rmsnorm_rope_bf16_,{fu:.3f},{eu:.3f},{eu/fu:.2f}x")
+        def flash_rope():
+            out.copy_(x)
+            ops.rope_rotate_half_bf16_(out, cos, sin)
+        def eager_rope():
+            rotate_half_ref(x, cos, sin)
+        fu = time_us(flash_rope, args.warmup, args.iters)
+        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")
+    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)
+    def flash_gather():
+        ops.text_gather_bf16(x, batch, seq, out=gathered)
+    def eager_gather():
+        torch.stack([x[b * seq + offset] for b in range(batch) for offset in (0, seq - 1)], dim=0)
+    fu = time_us(flash_gather, args.warmup, args.iters)
+    eu = time_us(eager_gather, max(5, args.warmup // 2), max(20, args.iters // 2))
+    print(f"text_tokens,{batch}x{seq}x{dim},text_gather_bf16,{fu:.3f},{eu:.3f},{eu/fu:.2f}x")
+    return 0
+if __name__ == "__main__":
+    raise SystemExit(main())

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

	@@ -0,0 +1,129 @@

+"""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,
+) -> 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
+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
+__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_",
+]

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

	@@ -0,0 +1,9 @@

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

build/torch211-cxx11-cu128-x86_64-linux/_transformer_layout_primitives_cuda_c9b2fa0.abi3.so ADDED Viewed

	@@ -0,0 +1,3 @@

+version https://git-lfs.github.com/spec/v1
+oid sha256:8391a04a4384cc92d83c6df4c86ecdd205567c00105c13a01ca54c2b88fdd371
+size 572816

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

	@@ -0,0 +1,23 @@

+{
+  "name": "transformer-layout-primitives",
+  "id": "_transformer_layout_primitives_cuda_c9b2fa0",
+  "version": 1,
+  "license": "Apache-2.0",
+  "python-depends": [],
+  "backend": {
+    "type": "cuda",
+    "archs": [
+      "10.0",
+      "10.1",
+      "12.0+PTX",
+      "7.0",
+      "7.2",
+      "7.5",
+      "8.0",
+      "8.6",
+      "8.7",
+      "8.9",
+      "9.0"
+    ]
+  }
+}

build/torch211-cxx11-cu128-x86_64-linux/transformer_layout_primitives/__init__.py ADDED Viewed

	@@ -0,0 +1,26 @@

+import ctypes
+import importlib.util
+import sys
+from pathlib import Path
+from types import ModuleType
+def _import_from_path(file_path: Path) -> ModuleType:
+    # We cannot use the module name as-is, after adding it to `sys.modules`,
+    # it would also be used for other imports. So, we make a module name that
+    # depends on the path for it to be unique using the hex-encoded hash of
+    # the path.
+    path_hash = "{:x}".format(ctypes.c_size_t(hash(file_path.absolute())).value)
+    module_name = path_hash
+    spec = importlib.util.spec_from_file_location(module_name, file_path)
+    if spec is None:
+        raise ImportError(f"Cannot load spec for {module_name} from {file_path}")
+    module = importlib.util.module_from_spec(spec)
+    if module is None:
+        raise ImportError(f"Cannot load module {module_name} from spec")
+    sys.modules[module_name] = module
+    spec.loader.exec_module(module)  # type: ignore
+    return module
+globals().update(vars(_import_from_path(Path(__file__).parent.parent / "__init__.py")))

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

	@@ -0,0 +1,129 @@

+"""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,
+) -> 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
+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
+__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_",
+]

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

	@@ -0,0 +1,9 @@

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

build/torch211-cxx11-cu130-x86_64-linux/_transformer_layout_primitives_cuda_c9b2fa0.abi3.so ADDED Viewed

	@@ -0,0 +1,3 @@

+version https://git-lfs.github.com/spec/v1
+oid sha256:99943084af51325b92b740243b59d224141161ee451749776e1c89ce7cceded2
+size 564160

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

	@@ -0,0 +1,22 @@

+{
+  "name": "transformer-layout-primitives",
+  "id": "_transformer_layout_primitives_cuda_c9b2fa0",
+  "version": 1,
+  "license": "Apache-2.0",
+  "python-depends": [],
+  "backend": {
+    "type": "cuda",
+    "archs": [
+      "10.0",
+      "11.0",
+      "12.0",
+      "12.1+PTX",
+      "7.5",
+      "8.0",
+      "8.6",
+      "8.7",
+      "8.9",
+      "9.0"
+    ]
+  }
+}

build/torch211-cxx11-cu130-x86_64-linux/transformer_layout_primitives/__init__.py ADDED Viewed

	@@ -0,0 +1,26 @@

+import ctypes
+import importlib.util
+import sys
+from pathlib import Path
+from types import ModuleType
+def _import_from_path(file_path: Path) -> ModuleType:
+    # We cannot use the module name as-is, after adding it to `sys.modules`,
+    # it would also be used for other imports. So, we make a module name that
+    # depends on the path for it to be unique using the hex-encoded hash of
+    # the path.
+    path_hash = "{:x}".format(ctypes.c_size_t(hash(file_path.absolute())).value)
+    module_name = path_hash
+    spec = importlib.util.spec_from_file_location(module_name, file_path)
+    if spec is None:
+        raise ImportError(f"Cannot load spec for {module_name} from {file_path}")
+    module = importlib.util.module_from_spec(spec)
+    if module is None:
+        raise ImportError(f"Cannot load module {module_name} from spec")
+    sys.modules[module_name] = module
+    spec.loader.exec_module(module)  # type: ignore
+    return module
+globals().update(vars(_import_from_path(Path(__file__).parent.parent / "__init__.py")))

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

	@@ -0,0 +1,129 @@

+"""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,
+) -> 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
+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
+__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_",
+]

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

	@@ -0,0 +1,9 @@

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

build/torch212-cxx11-cu130-x86_64-linux/_transformer_layout_primitives_cuda_c9b2fa0.abi3.so ADDED Viewed

	@@ -0,0 +1,3 @@

+version https://git-lfs.github.com/spec/v1
+oid sha256:6c4ebdc5e5faa851f8c4f537f5a1e75e28cebad98a650c5bf72dbf60db981844
+size 570424

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

	@@ -0,0 +1,22 @@

+{
+  "name": "transformer-layout-primitives",
+  "id": "_transformer_layout_primitives_cuda_c9b2fa0",
+  "version": 1,
+  "license": "Apache-2.0",
+  "python-depends": [],
+  "backend": {
+    "type": "cuda",
+    "archs": [
+      "10.0",
+      "11.0",
+      "12.0",
+      "12.1+PTX",
+      "7.5",
+      "8.0",
+      "8.6",
+      "8.7",
+      "8.9",
+      "9.0"
+    ]
+  }
+}

build/torch212-cxx11-cu130-x86_64-linux/transformer_layout_primitives/__init__.py ADDED Viewed

	@@ -0,0 +1,26 @@

+import ctypes
+import importlib.util
+import sys
+from pathlib import Path
+from types import ModuleType
+def _import_from_path(file_path: Path) -> ModuleType:
+    # We cannot use the module name as-is, after adding it to `sys.modules`,
+    # it would also be used for other imports. So, we make a module name that
+    # depends on the path for it to be unique using the hex-encoded hash of
+    # the path.
+    path_hash = "{:x}".format(ctypes.c_size_t(hash(file_path.absolute())).value)
+    module_name = path_hash
+    spec = importlib.util.spec_from_file_location(module_name, file_path)
+    if spec is None:
+        raise ImportError(f"Cannot load spec for {module_name} from {file_path}")
+    module = importlib.util.module_from_spec(spec)
+    if module is None:
+        raise ImportError(f"Cannot load module {module_name} from spec")
+    sys.modules[module_name] = module
+    spec.loader.exec_module(module)  # type: ignore
+    return module
+globals().update(vars(_import_from_path(Path(__file__).parent.parent / "__init__.py")))

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

	@@ -0,0 +1,129 @@

+"""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,
+) -> 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
+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
+__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_",
+]

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

	@@ -0,0 +1,9 @@

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

build/torch212-cxx11-cu132-x86_64-linux/_transformer_layout_primitives_cuda_c9b2fa0.abi3.so ADDED Viewed

	@@ -0,0 +1,3 @@

+version https://git-lfs.github.com/spec/v1
+oid sha256:1947aebc33c49a48c82a0da681b56e16c0517460121fa5ce7cbc1d39f1e9832b
+size 574520

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

	@@ -0,0 +1,22 @@

+{
+  "name": "transformer-layout-primitives",
+  "id": "_transformer_layout_primitives_cuda_c9b2fa0",
+  "version": 1,
+  "license": "Apache-2.0",
+  "python-depends": [],
+  "backend": {
+    "type": "cuda",
+    "archs": [
+      "10.0",
+      "11.0",
+      "12.0",
+      "12.1+PTX",
+      "7.5",
+      "8.0",
+      "8.6",
+      "8.7",
+      "8.9",
+      "9.0"
+    ]
+  }
+}

build/torch212-cxx11-cu132-x86_64-linux/transformer_layout_primitives/__init__.py ADDED Viewed

	@@ -0,0 +1,26 @@

+import ctypes
+import importlib.util
+import sys
+from pathlib import Path
+from types import ModuleType
+def _import_from_path(file_path: Path) -> ModuleType:
+    # We cannot use the module name as-is, after adding it to `sys.modules`,
+    # it would also be used for other imports. So, we make a module name that
+    # depends on the path for it to be unique using the hex-encoded hash of
+    # the path.
+    path_hash = "{:x}".format(ctypes.c_size_t(hash(file_path.absolute())).value)
+    module_name = path_hash
+    spec = importlib.util.spec_from_file_location(module_name, file_path)
+    if spec is None:
+        raise ImportError(f"Cannot load spec for {module_name} from {file_path}")
+    module = importlib.util.module_from_spec(spec)
+    if module is None:
+        raise ImportError(f"Cannot load module {module_name} from spec")
+    sys.modules[module_name] = module
+    spec.loader.exec_module(module)  # type: ignore
+    return module
+globals().update(vars(_import_from_path(Path(__file__).parent.parent / "__init__.py")))