liangsu9988 commited on 26 days ago

Commit

b34d99d

verified ·

1 Parent(s): ed6548a

Uploaded using `kernel-builder`.

Browse files

Files changed (48) hide show

benchmarks/benchmark.py +25 -7
benchmarks/benchmark_channel_scale.py +21 -4
build/torch211-cxx11-cu126-x86_64-linux/__init__.py +7 -1
build/torch211-cxx11-cu126-x86_64-linux/{_flashrt_gemm_epilogues_cuda_5de4768.abi3.so → _flashrt_gemm_epilogues_cuda_4a17df5.abi3.so} +2 -2
build/torch211-cxx11-cu126-x86_64-linux/_ops.py +3 -3
build/torch211-cxx11-cu126-x86_64-linux/metadata.json +1 -1
build/torch211-cxx11-cu128-x86_64-linux/__init__.py +7 -1
build/torch211-cxx11-cu128-x86_64-linux/{_flashrt_gemm_epilogues_cuda_5de4768.abi3.so → _flashrt_gemm_epilogues_cuda_4a17df5.abi3.so} +2 -2
build/torch211-cxx11-cu128-x86_64-linux/_ops.py +3 -3
build/torch211-cxx11-cu128-x86_64-linux/metadata.json +1 -1
build/torch211-cxx11-cu130-x86_64-linux/__init__.py +7 -1
build/torch211-cxx11-cu130-x86_64-linux/{_flashrt_gemm_epilogues_cuda_5de4768.abi3.so → _flashrt_gemm_epilogues_cuda_4a17df5.abi3.so} +2 -2
build/torch211-cxx11-cu130-x86_64-linux/_ops.py +3 -3
build/torch211-cxx11-cu130-x86_64-linux/metadata.json +1 -1
build/torch211-cxx11-rocm71-x86_64-linux/__init__.py +277 -0
build/{torch212-cxx11-cu126-x86_64-linux/_flashrt_gemm_epilogues_cuda_5de4768.abi3.so → torch211-cxx11-rocm71-x86_64-linux/_flashrt_gemm_epilogues_rocm_4a17df5.abi3.so} +2 -2
build/torch211-cxx11-rocm71-x86_64-linux/_ops.py +9 -0
build/torch211-cxx11-rocm71-x86_64-linux/flashrt_gemm_epilogues/__init__.py +26 -0
build/torch211-cxx11-rocm71-x86_64-linux/metadata.json +13 -0
build/torch211-cxx11-rocm72-x86_64-linux/__init__.py +277 -0
build/torch211-cxx11-rocm72-x86_64-linux/_flashrt_gemm_epilogues_rocm_4a17df5.abi3.so +3 -0
build/torch211-cxx11-rocm72-x86_64-linux/_ops.py +9 -0
build/torch211-cxx11-rocm72-x86_64-linux/flashrt_gemm_epilogues/__init__.py +26 -0
build/torch211-cxx11-rocm72-x86_64-linux/metadata.json +13 -0
build/torch212-cxx11-cu126-x86_64-linux/__init__.py +7 -1
build/torch212-cxx11-cu126-x86_64-linux/_flashrt_gemm_epilogues_cuda_4a17df5.abi3.so +3 -0
build/torch212-cxx11-cu126-x86_64-linux/_ops.py +3 -3
build/torch212-cxx11-cu126-x86_64-linux/metadata.json +1 -1
build/torch212-cxx11-cu130-x86_64-linux/__init__.py +7 -1
build/torch212-cxx11-cu130-x86_64-linux/_flashrt_gemm_epilogues_cuda_4a17df5.abi3.so +3 -0
build/torch212-cxx11-cu130-x86_64-linux/_flashrt_gemm_epilogues_cuda_5de4768.abi3.so +0 -3
build/torch212-cxx11-cu130-x86_64-linux/_ops.py +3 -3
build/torch212-cxx11-cu130-x86_64-linux/metadata.json +1 -1
build/torch212-cxx11-cu132-x86_64-linux/__init__.py +7 -1
build/torch212-cxx11-cu132-x86_64-linux/_flashrt_gemm_epilogues_cuda_4a17df5.abi3.so +3 -0
build/torch212-cxx11-cu132-x86_64-linux/_flashrt_gemm_epilogues_cuda_5de4768.abi3.so +0 -3
build/torch212-cxx11-cu132-x86_64-linux/_ops.py +3 -3
build/torch212-cxx11-cu132-x86_64-linux/metadata.json +1 -1
build/torch212-cxx11-rocm71-x86_64-linux/__init__.py +277 -0
build/torch212-cxx11-rocm71-x86_64-linux/_flashrt_gemm_epilogues_rocm_4a17df5.abi3.so +3 -0
build/torch212-cxx11-rocm71-x86_64-linux/_ops.py +9 -0
build/torch212-cxx11-rocm71-x86_64-linux/flashrt_gemm_epilogues/__init__.py +26 -0
build/torch212-cxx11-rocm71-x86_64-linux/metadata.json +13 -0
build/torch212-cxx11-rocm72-x86_64-linux/__init__.py +277 -0
build/torch212-cxx11-rocm72-x86_64-linux/_flashrt_gemm_epilogues_rocm_4a17df5.abi3.so +3 -0
build/torch212-cxx11-rocm72-x86_64-linux/_ops.py +9 -0
build/torch212-cxx11-rocm72-x86_64-linux/flashrt_gemm_epilogues/__init__.py +26 -0
build/torch212-cxx11-rocm72-x86_64-linux/metadata.json +13 -0

benchmarks/benchmark.py CHANGED Viewed

@@ -6,8 +6,24 @@ from kernels.benchmark import Benchmark
 _original_allclose = torch.allclose
 def _flashrt_allclose(input, other, rtol=1e-05, atol=1e-08, equal_nan=False):
-    if input.dtype == torch.float8_e4m3fn or other.dtype == torch.float8_e4m3fn:
         return _original_allclose(
             input.float(),
             other.float(),
@@ -47,13 +63,14 @@ class BiasGeluFp8QuantizeBenchmark(Benchmark):
         self.input = torch.randn((m, n), device=self.device, dtype=torch.bfloat16)
         self.bias = torch.randn((n,), device=self.device, dtype=torch.bfloat16)
         self.scale = torch.tensor([0.25], device=self.device, dtype=torch.float32)
-        self.out = torch.empty_like(self.input, dtype=torch.float8_e4m3fn)
     def _reference(self) -> torch.Tensor:
         y = self.input.float() + self.bias.float()
         y = torch.nn.functional.gelu(y, approximate="tanh")
-        y = torch.clamp(y / self.scale.float(), -448.0, 448.0)
-        return y.to(torch.float8_e4m3fn)
     def setup_decode_m1(self) -> None:
         self._setup_shape(1, 4096)
@@ -158,12 +175,13 @@ class GeluFp8QuantizeBenchmark(Benchmark):
     def _setup_shape(self, m: int, n: int) -> None:
         self.input = torch.randn((m, n), device=self.device, dtype=torch.bfloat16)
         self.scale = torch.tensor([0.25], device=self.device, dtype=torch.float32)
-        self.out = torch.empty_like(self.input, dtype=torch.float8_e4m3fn)
     def _reference(self) -> torch.Tensor:
         y = torch.nn.functional.gelu(self.input.float(), approximate="tanh")
-        y = torch.clamp(y / self.scale.float(), -448.0, 448.0)
-        return y.to(torch.float8_e4m3fn)
     def setup_decode_m1(self) -> None:
         self._setup_shape(1, 4096)

 _original_allclose = torch.allclose
+def _fp8_dtype() -> torch.dtype:
+    if torch.version.hip is not None and hasattr(torch, "float8_e4m3fnuz"):
+        return torch.float8_e4m3fnuz
+    return torch.float8_e4m3fn
+def _fp8_max() -> float:
+    return 240.0 if torch.version.hip is not None else 448.0
+def _is_fp8_dtype(dtype: torch.dtype) -> bool:
+    return dtype == torch.float8_e4m3fn or (
+        hasattr(torch, "float8_e4m3fnuz") and dtype == torch.float8_e4m3fnuz
+    )
 def _flashrt_allclose(input, other, rtol=1e-05, atol=1e-08, equal_nan=False):
+    if _is_fp8_dtype(input.dtype) or _is_fp8_dtype(other.dtype):
         return _original_allclose(
             input.float(),
             other.float(),
         self.input = torch.randn((m, n), device=self.device, dtype=torch.bfloat16)
         self.bias = torch.randn((n,), device=self.device, dtype=torch.bfloat16)
         self.scale = torch.tensor([0.25], device=self.device, dtype=torch.float32)
+        self.out = torch.empty_like(self.input, dtype=_fp8_dtype())
     def _reference(self) -> torch.Tensor:
         y = self.input.float() + self.bias.float()
         y = torch.nn.functional.gelu(y, approximate="tanh")
+        limit = _fp8_max()
+        y = torch.clamp(y / self.scale.float(), -limit, limit)
+        return y.to(_fp8_dtype())
     def setup_decode_m1(self) -> None:
         self._setup_shape(1, 4096)
     def _setup_shape(self, m: int, n: int) -> None:
         self.input = torch.randn((m, n), device=self.device, dtype=torch.bfloat16)
         self.scale = torch.tensor([0.25], device=self.device, dtype=torch.float32)
+        self.out = torch.empty_like(self.input, dtype=_fp8_dtype())
     def _reference(self) -> torch.Tensor:
         y = torch.nn.functional.gelu(self.input.float(), approximate="tanh")
+        limit = _fp8_max()
+        y = torch.clamp(y / self.scale.float(), -limit, limit)
+        return y.to(_fp8_dtype())
     def setup_decode_m1(self) -> None:
         self._setup_shape(1, 4096)

benchmarks/benchmark_channel_scale.py CHANGED Viewed

@@ -6,8 +6,24 @@ from kernels.benchmark import Benchmark
 _original_allclose = torch.allclose
 def _flashrt_allclose(input, other, rtol=1e-05, atol=1e-08, equal_nan=False):
-    if input.dtype == torch.float8_e4m3fn or other.dtype == torch.float8_e4m3fn:
         return _original_allclose(
             input.float(),
             other.float(),
@@ -49,12 +65,13 @@ class ChannelScaleFp8QuantizeBenchmark(Benchmark):
             (n,), device=self.device, dtype=torch.bfloat16
         )
         self.scale = torch.tensor([0.25], device=self.device, dtype=torch.float32)
-        self.out = torch.empty_like(self.input, dtype=torch.float8_e4m3fn)
     def _reference(self) -> torch.Tensor:
         y = self.input.float() * self.channel_scale.float()
-        y = torch.clamp(y / self.scale.float(), -448.0, 448.0)
-        return y.to(torch.float8_e4m3fn)
     def setup_decode_m1(self) -> None:
         self._setup_shape(1, 4096)

 _original_allclose = torch.allclose
+def _fp8_dtype() -> torch.dtype:
+    if torch.version.hip is not None and hasattr(torch, "float8_e4m3fnuz"):
+        return torch.float8_e4m3fnuz
+    return torch.float8_e4m3fn
+def _fp8_max() -> float:
+    return 240.0 if torch.version.hip is not None else 448.0
+def _is_fp8_dtype(dtype: torch.dtype) -> bool:
+    return dtype == torch.float8_e4m3fn or (
+        hasattr(torch, "float8_e4m3fnuz") and dtype == torch.float8_e4m3fnuz
+    )
 def _flashrt_allclose(input, other, rtol=1e-05, atol=1e-08, equal_nan=False):
+    if _is_fp8_dtype(input.dtype) or _is_fp8_dtype(other.dtype):
         return _original_allclose(
             input.float(),
             other.float(),
             (n,), device=self.device, dtype=torch.bfloat16
         )
         self.scale = torch.tensor([0.25], device=self.device, dtype=torch.float32)
+        self.out = torch.empty_like(self.input, dtype=_fp8_dtype())
     def _reference(self) -> torch.Tensor:
         y = self.input.float() * self.channel_scale.float()
+        limit = _fp8_max()
+        y = torch.clamp(y / self.scale.float(), -limit, limit)
+        return y.to(_fp8_dtype())
     def setup_decode_m1(self) -> None:
         self._setup_shape(1, 4096)

build/torch211-cxx11-cu126-x86_64-linux/__init__.py CHANGED Viewed

@@ -126,8 +126,14 @@ def _channel_scale_quantize_fp8_static_bf16_fake(
     return None
 def _allocate_fp8_like(input: torch.Tensor) -> torch.Tensor:
-    return torch.empty(input.shape, device=input.device, dtype=torch.float8_e4m3fn)
 def bf16_linear_bf16(

     return None
+def _fp8_dtype() -> torch.dtype:
+    if torch.version.hip is not None and hasattr(torch, "float8_e4m3fnuz"):
+        return torch.float8_e4m3fnuz
+    return torch.float8_e4m3fn
 def _allocate_fp8_like(input: torch.Tensor) -> torch.Tensor:
+    return torch.empty(input.shape, device=input.device, dtype=_fp8_dtype())
 def bf16_linear_bf16(

build/torch211-cxx11-cu126-x86_64-linux/{_flashrt_gemm_epilogues_cuda_5de4768.abi3.so → _flashrt_gemm_epilogues_cuda_4a17df5.abi3.so} RENAMED Viewed

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

 version https://git-lfs.github.com/spec/v1
+oid sha256:fb5e2bf3a388a8f86dd83dcbfdbe02e35bbe1aafca724dc6d833f5e00513ad5d
+size 302272

build/torch211-cxx11-cu126-x86_64-linux/_ops.py CHANGED Viewed

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

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

build/torch211-cxx11-cu126-x86_64-linux/metadata.json CHANGED Viewed

@@ -1,6 +1,6 @@
 {
   "name": "flashrt-gemm-epilogues",
-  "id": "_flashrt_gemm_epilogues_cuda_5de4768",
   "version": 1,
   "license": "Apache-2.0",
   "python-depends": [],

 {
   "name": "flashrt-gemm-epilogues",
+  "id": "_flashrt_gemm_epilogues_cuda_4a17df5",
   "version": 1,
   "license": "Apache-2.0",
   "python-depends": [],

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

@@ -126,8 +126,14 @@ def _channel_scale_quantize_fp8_static_bf16_fake(
     return None
 def _allocate_fp8_like(input: torch.Tensor) -> torch.Tensor:
-    return torch.empty(input.shape, device=input.device, dtype=torch.float8_e4m3fn)
 def bf16_linear_bf16(

     return None
+def _fp8_dtype() -> torch.dtype:
+    if torch.version.hip is not None and hasattr(torch, "float8_e4m3fnuz"):
+        return torch.float8_e4m3fnuz
+    return torch.float8_e4m3fn
 def _allocate_fp8_like(input: torch.Tensor) -> torch.Tensor:
+    return torch.empty(input.shape, device=input.device, dtype=_fp8_dtype())
 def bf16_linear_bf16(

build/torch211-cxx11-cu128-x86_64-linux/{_flashrt_gemm_epilogues_cuda_5de4768.abi3.so → _flashrt_gemm_epilogues_cuda_4a17df5.abi3.so} RENAMED Viewed

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

 version https://git-lfs.github.com/spec/v1
+oid sha256:8bf38d750282de45d0bbe2dbd9cca35740f97d2f1d0541dae68761c2113079b8
+size 388208

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

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

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

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

@@ -1,6 +1,6 @@
 {
   "name": "flashrt-gemm-epilogues",
-  "id": "_flashrt_gemm_epilogues_cuda_5de4768",
   "version": 1,
   "license": "Apache-2.0",
   "python-depends": [],

 {
   "name": "flashrt-gemm-epilogues",
+  "id": "_flashrt_gemm_epilogues_cuda_4a17df5",
   "version": 1,
   "license": "Apache-2.0",
   "python-depends": [],

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

@@ -126,8 +126,14 @@ def _channel_scale_quantize_fp8_static_bf16_fake(
     return None
 def _allocate_fp8_like(input: torch.Tensor) -> torch.Tensor:
-    return torch.empty(input.shape, device=input.device, dtype=torch.float8_e4m3fn)
 def bf16_linear_bf16(

     return None
+def _fp8_dtype() -> torch.dtype:
+    if torch.version.hip is not None and hasattr(torch, "float8_e4m3fnuz"):
+        return torch.float8_e4m3fnuz
+    return torch.float8_e4m3fn
 def _allocate_fp8_like(input: torch.Tensor) -> torch.Tensor:
+    return torch.empty(input.shape, device=input.device, dtype=_fp8_dtype())
 def bf16_linear_bf16(

build/torch211-cxx11-cu130-x86_64-linux/{_flashrt_gemm_epilogues_cuda_5de4768.abi3.so → _flashrt_gemm_epilogues_cuda_4a17df5.abi3.so} RENAMED Viewed

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

 version https://git-lfs.github.com/spec/v1
+oid sha256:c96cf9ac67cd49f8d7b710e4bdd68e8f517e8f3c1c787de265fd45d292f3c53b
+size 390816

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

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

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

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

@@ -1,6 +1,6 @@
 {
   "name": "flashrt-gemm-epilogues",
-  "id": "_flashrt_gemm_epilogues_cuda_5de4768",
   "version": 1,
   "license": "Apache-2.0",
   "python-depends": [],

 {
   "name": "flashrt-gemm-epilogues",
+  "id": "_flashrt_gemm_epilogues_cuda_4a17df5",
   "version": 1,
   "license": "Apache-2.0",
   "python-depends": [],

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

	@@ -0,0 +1,277 @@

+"""FlashRT fused GEMM epilogue kernels."""
+from __future__ import annotations
+import ctypes
+import ctypes.util
+from pathlib import Path
+from typing import Optional
+import torch
+def _torch_bundled_cublaslt() -> Optional[Path]:
+    for parent in Path(torch.__file__).resolve().parents:
+        candidate = parent / "nvidia" / "cublas" / "lib" / "libcublasLt.so.12"
+        if candidate.exists():
+            return candidate
+    return None
+def _preload_cublaslt() -> None:
+    bundled = _torch_bundled_cublaslt()
+    library = str(bundled) if bundled is not None else (
+        ctypes.util.find_library("cublasLt") or "libcublasLt.so"
+    )
+    try:
+        ctypes.CDLL(library, mode=ctypes.RTLD_GLOBAL)
+    except OSError:
+        pass
+_preload_cublaslt()
+from ._ops import add_op_namespace_prefix, ops
+@torch.library.register_fake(add_op_namespace_prefix("bf16_linear_bf16"))
+def _bf16_linear_bf16_fake(
+    x: torch.Tensor,
+    w: torch.Tensor,
+    out: torch.Tensor,
+) -> None:
+    if x.dim() != 2 or w.dim() != 2:
+        raise RuntimeError("x and w must be rank-2 tensors")
+    if out.shape != (x.shape[0], w.shape[1]):
+        raise RuntimeError("out shape must be (x.shape[0], w.shape[1])")
+    return None
+@torch.library.register_fake(add_op_namespace_prefix("bf16_linear_bias_bf16"))
+def _bf16_linear_bias_bf16_fake(
+    x: torch.Tensor,
+    w: torch.Tensor,
+    bias: torch.Tensor,
+    out: torch.Tensor,
+) -> None:
+    if x.dim() != 2 or w.dim() != 2:
+        raise RuntimeError("x and w must be rank-2 tensors")
+    if bias.dim() != 1:
+        raise RuntimeError("bias must be rank-1")
+    if out.shape != (x.shape[0], w.shape[1]):
+        raise RuntimeError("out shape must be (x.shape[0], w.shape[1])")
+    return None
+@torch.library.register_fake(add_op_namespace_prefix("bf16_gemm_bias_gelu"))
+def _bf16_gemm_bias_gelu_fake(
+    a: torch.Tensor,
+    b: torch.Tensor,
+    bias: torch.Tensor,
+    out: torch.Tensor,
+) -> None:
+    if a.dim() != 2 or b.dim() != 2:
+        raise RuntimeError("a and b must be rank-2 tensors")
+    if out.shape != (a.shape[0], b.shape[1]):
+        raise RuntimeError("out shape must be (a.shape[0], b.shape[1])")
+    return None
+@torch.library.register_fake(add_op_namespace_prefix("bf16_gemm_bias"))
+def _bf16_gemm_bias_fake(
+    a: torch.Tensor,
+    b: torch.Tensor,
+    bias: torch.Tensor,
+    out: torch.Tensor,
+) -> None:
+    if a.dim() != 2 or b.dim() != 2:
+        raise RuntimeError("a and b must be rank-2 tensors")
+    if out.shape != (a.shape[0], b.shape[1]):
+        raise RuntimeError("out shape must be (a.shape[0], b.shape[1])")
+    return None
+@torch.library.register_fake(add_op_namespace_prefix("bias_gelu_quantize_fp8_static_bf16"))
+def _bias_gelu_quantize_fp8_static_bf16_fake(
+    input: torch.Tensor,
+    bias: torch.Tensor,
+    scale: torch.Tensor,
+    out: torch.Tensor,
+) -> None:
+    if out.shape != input.shape:
+        raise RuntimeError("out shape must match input shape")
+    return None
+@torch.library.register_fake(add_op_namespace_prefix("gelu_quantize_fp8_static_bf16"))
+def _gelu_quantize_fp8_static_bf16_fake(
+    input: torch.Tensor,
+    scale: torch.Tensor,
+    out: torch.Tensor,
+) -> None:
+    if out.shape != input.shape:
+        raise RuntimeError("out shape must match input shape")
+    return None
+@torch.library.register_fake(add_op_namespace_prefix("channel_scale_quantize_fp8_static_bf16"))
+def _channel_scale_quantize_fp8_static_bf16_fake(
+    input: torch.Tensor,
+    channel_scale: torch.Tensor,
+    scale: torch.Tensor,
+    out: torch.Tensor,
+) -> None:
+    if out.shape != input.shape:
+        raise RuntimeError("out shape must match input shape")
+    return None
+def _fp8_dtype() -> torch.dtype:
+    if torch.version.hip is not None and hasattr(torch, "float8_e4m3fnuz"):
+        return torch.float8_e4m3fnuz
+    return torch.float8_e4m3fn
+def _allocate_fp8_like(input: torch.Tensor) -> torch.Tensor:
+    return torch.empty(input.shape, device=input.device, dtype=_fp8_dtype())
+def bf16_linear_bf16(
+    x: torch.Tensor,
+    w: torch.Tensor,
+    out: Optional[torch.Tensor] = None,
+) -> torch.Tensor:
+    """Compute ``x @ w`` and store BF16 output.
+    ``x`` must be contiguous BF16 with shape ``(M, K)``. ``w`` must be
+    contiguous BF16 with shape ``(K, N)``. If ``out`` is omitted, a BF16
+    ``(M, N)`` tensor is allocated.
+    """
+    if out is None:
+        out = torch.empty(
+            (x.shape[0], w.shape[1]), device=x.device, dtype=torch.bfloat16
+        )
+    ops.bf16_linear_bf16(x, w, out)
+    return out
+def bf16_linear_bias_bf16(
+    x: torch.Tensor,
+    w: torch.Tensor,
+    bias: torch.Tensor,
+    out: Optional[torch.Tensor] = None,
+) -> torch.Tensor:
+    """Compute ``x @ w + bias`` and store BF16 output.
+    ``x`` must be contiguous BF16 with shape ``(M, K)``. ``w`` must be
+    contiguous BF16 with shape ``(K, N)``. ``bias`` must be contiguous BF16
+    with shape ``(N,)``. If ``out`` is omitted, a BF16 ``(M, N)`` tensor is
+    allocated.
+    """
+    if out is None:
+        out = torch.empty(
+            (x.shape[0], w.shape[1]), device=x.device, dtype=torch.bfloat16
+        )
+    ops.bf16_linear_bias_bf16(x, w, bias, out)
+    return out
+def bf16_gemm_bias_gelu(
+    a: torch.Tensor,
+    b: torch.Tensor,
+    bias: torch.Tensor,
+    out: Optional[torch.Tensor] = None,
+) -> torch.Tensor:
+    """Compute ``gelu(a @ b + bias)`` and store BF16 output.
+    ``a`` must be contiguous BF16 with shape ``(M, K)``. ``b`` must be
+    contiguous BF16 with shape ``(K, N)``. ``bias`` must be contiguous BF16
+    with shape ``(N,)``. If ``out`` is omitted, a BF16 ``(M, N)`` tensor is
+    allocated.
+    """
+    if out is None:
+        out = torch.empty(
+            (a.shape[0], b.shape[1]), device=a.device, dtype=torch.bfloat16
+        )
+    ops.bf16_gemm_bias_gelu(a, b, bias, out)
+    return out
+def bf16_gemm_bias(
+    a: torch.Tensor,
+    b: torch.Tensor,
+    bias: torch.Tensor,
+    out: Optional[torch.Tensor] = None,
+) -> torch.Tensor:
+    """Compute ``a @ b + bias`` and store BF16 output."""
+    if out is None:
+        out = torch.empty(
+            (a.shape[0], b.shape[1]), device=a.device, dtype=torch.bfloat16
+        )
+    ops.bf16_gemm_bias(a, b, bias, out)
+    return out
+def bias_gelu_quantize_fp8_static_bf16(
+    input: torch.Tensor,
+    bias: torch.Tensor,
+    scale: torch.Tensor,
+    out: Optional[torch.Tensor] = None,
+) -> torch.Tensor:
+    """Compute ``gelu(input + bias) / scale`` and store FP8 e4m3 output.
+    ``input`` must be contiguous BF16 with shape ``(..., N)``. ``bias`` must be
+    contiguous BF16 with shape ``(N,)``. ``scale`` must be a CUDA float32 scalar
+    tensor. If ``out`` is omitted, an FP8 e4m3 output tensor is allocated.
+    """
+    if out is None:
+        out = _allocate_fp8_like(input)
+    ops.bias_gelu_quantize_fp8_static_bf16(input, bias, scale, out)
+    return out
+def gelu_quantize_fp8_static_bf16(
+    input: torch.Tensor,
+    scale: torch.Tensor,
+    out: Optional[torch.Tensor] = None,
+) -> torch.Tensor:
+    """Compute ``gelu(input) / scale`` and store FP8 e4m3 output."""
+    if out is None:
+        out = _allocate_fp8_like(input)
+    ops.gelu_quantize_fp8_static_bf16(input, scale, out)
+    return out
+def channel_scale_quantize_fp8_static_bf16(
+    input: torch.Tensor,
+    channel_scale: torch.Tensor,
+    scale: torch.Tensor,
+    out: Optional[torch.Tensor] = None,
+) -> torch.Tensor:
+    """Compute ``input * channel_scale / scale`` and store FP8 e4m3 output.
+    ``channel_scale`` must be contiguous BF16 with shape ``(input.shape[-1],)``.
+    It is broadcast over all leading input dimensions.
+    """
+    if out is None:
+        out = _allocate_fp8_like(input)
+    ops.channel_scale_quantize_fp8_static_bf16(input, channel_scale, scale, out)
+    return out
+__all__ = [
+    "bf16_linear_bf16",
+    "bf16_linear_bias_bf16",
+    "bf16_gemm_bias_gelu",
+    "bf16_gemm_bias",
+    "bias_gelu_quantize_fp8_static_bf16",
+    "channel_scale_quantize_fp8_static_bf16",
+    "gelu_quantize_fp8_static_bf16",
+]

build/{torch212-cxx11-cu126-x86_64-linux/_flashrt_gemm_epilogues_cuda_5de4768.abi3.so → torch211-cxx11-rocm71-x86_64-linux/_flashrt_gemm_epilogues_rocm_4a17df5.abi3.so} RENAMED Viewed

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

 version https://git-lfs.github.com/spec/v1
+oid sha256:466c2f4bd1a6a50f798707736e5d2f2bbfda6acfd216a9ffaa2b686a781c4a25
+size 157472

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

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

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

build/torch211-cxx11-rocm71-x86_64-linux/flashrt_gemm_epilogues/__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-rocm71-x86_64-linux/metadata.json ADDED Viewed

	@@ -0,0 +1,13 @@

+{
+  "name": "flashrt-gemm-epilogues",
+  "id": "_flashrt_gemm_epilogues_rocm_4a17df5",
+  "version": 1,
+  "license": "Apache-2.0",
+  "python-depends": [],
+  "backend": {
+    "type": "rocm",
+    "archs": [
+      "gfx942"
+    ]
+  }
+}

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

	@@ -0,0 +1,277 @@

+"""FlashRT fused GEMM epilogue kernels."""
+from __future__ import annotations
+import ctypes
+import ctypes.util
+from pathlib import Path
+from typing import Optional
+import torch
+def _torch_bundled_cublaslt() -> Optional[Path]:
+    for parent in Path(torch.__file__).resolve().parents:
+        candidate = parent / "nvidia" / "cublas" / "lib" / "libcublasLt.so.12"
+        if candidate.exists():
+            return candidate
+    return None
+def _preload_cublaslt() -> None:
+    bundled = _torch_bundled_cublaslt()
+    library = str(bundled) if bundled is not None else (
+        ctypes.util.find_library("cublasLt") or "libcublasLt.so"
+    )
+    try:
+        ctypes.CDLL(library, mode=ctypes.RTLD_GLOBAL)
+    except OSError:
+        pass
+_preload_cublaslt()
+from ._ops import add_op_namespace_prefix, ops
+@torch.library.register_fake(add_op_namespace_prefix("bf16_linear_bf16"))
+def _bf16_linear_bf16_fake(
+    x: torch.Tensor,
+    w: torch.Tensor,
+    out: torch.Tensor,
+) -> None:
+    if x.dim() != 2 or w.dim() != 2:
+        raise RuntimeError("x and w must be rank-2 tensors")
+    if out.shape != (x.shape[0], w.shape[1]):
+        raise RuntimeError("out shape must be (x.shape[0], w.shape[1])")
+    return None
+@torch.library.register_fake(add_op_namespace_prefix("bf16_linear_bias_bf16"))
+def _bf16_linear_bias_bf16_fake(
+    x: torch.Tensor,
+    w: torch.Tensor,
+    bias: torch.Tensor,
+    out: torch.Tensor,
+) -> None:
+    if x.dim() != 2 or w.dim() != 2:
+        raise RuntimeError("x and w must be rank-2 tensors")
+    if bias.dim() != 1:
+        raise RuntimeError("bias must be rank-1")
+    if out.shape != (x.shape[0], w.shape[1]):
+        raise RuntimeError("out shape must be (x.shape[0], w.shape[1])")
+    return None
+@torch.library.register_fake(add_op_namespace_prefix("bf16_gemm_bias_gelu"))
+def _bf16_gemm_bias_gelu_fake(
+    a: torch.Tensor,
+    b: torch.Tensor,
+    bias: torch.Tensor,
+    out: torch.Tensor,
+) -> None:
+    if a.dim() != 2 or b.dim() != 2:
+        raise RuntimeError("a and b must be rank-2 tensors")
+    if out.shape != (a.shape[0], b.shape[1]):
+        raise RuntimeError("out shape must be (a.shape[0], b.shape[1])")
+    return None
+@torch.library.register_fake(add_op_namespace_prefix("bf16_gemm_bias"))
+def _bf16_gemm_bias_fake(
+    a: torch.Tensor,
+    b: torch.Tensor,
+    bias: torch.Tensor,
+    out: torch.Tensor,
+) -> None:
+    if a.dim() != 2 or b.dim() != 2:
+        raise RuntimeError("a and b must be rank-2 tensors")
+    if out.shape != (a.shape[0], b.shape[1]):
+        raise RuntimeError("out shape must be (a.shape[0], b.shape[1])")
+    return None
+@torch.library.register_fake(add_op_namespace_prefix("bias_gelu_quantize_fp8_static_bf16"))
+def _bias_gelu_quantize_fp8_static_bf16_fake(
+    input: torch.Tensor,
+    bias: torch.Tensor,
+    scale: torch.Tensor,
+    out: torch.Tensor,
+) -> None:
+    if out.shape != input.shape:
+        raise RuntimeError("out shape must match input shape")
+    return None
+@torch.library.register_fake(add_op_namespace_prefix("gelu_quantize_fp8_static_bf16"))
+def _gelu_quantize_fp8_static_bf16_fake(
+    input: torch.Tensor,
+    scale: torch.Tensor,
+    out: torch.Tensor,
+) -> None:
+    if out.shape != input.shape:
+        raise RuntimeError("out shape must match input shape")
+    return None
+@torch.library.register_fake(add_op_namespace_prefix("channel_scale_quantize_fp8_static_bf16"))
+def _channel_scale_quantize_fp8_static_bf16_fake(
+    input: torch.Tensor,
+    channel_scale: torch.Tensor,
+    scale: torch.Tensor,
+    out: torch.Tensor,
+) -> None:
+    if out.shape != input.shape:
+        raise RuntimeError("out shape must match input shape")
+    return None
+def _fp8_dtype() -> torch.dtype:
+    if torch.version.hip is not None and hasattr(torch, "float8_e4m3fnuz"):
+        return torch.float8_e4m3fnuz
+    return torch.float8_e4m3fn
+def _allocate_fp8_like(input: torch.Tensor) -> torch.Tensor:
+    return torch.empty(input.shape, device=input.device, dtype=_fp8_dtype())
+def bf16_linear_bf16(
+    x: torch.Tensor,
+    w: torch.Tensor,
+    out: Optional[torch.Tensor] = None,
+) -> torch.Tensor:
+    """Compute ``x @ w`` and store BF16 output.
+    ``x`` must be contiguous BF16 with shape ``(M, K)``. ``w`` must be
+    contiguous BF16 with shape ``(K, N)``. If ``out`` is omitted, a BF16
+    ``(M, N)`` tensor is allocated.
+    """
+    if out is None:
+        out = torch.empty(
+            (x.shape[0], w.shape[1]), device=x.device, dtype=torch.bfloat16
+        )
+    ops.bf16_linear_bf16(x, w, out)
+    return out
+def bf16_linear_bias_bf16(
+    x: torch.Tensor,
+    w: torch.Tensor,
+    bias: torch.Tensor,
+    out: Optional[torch.Tensor] = None,
+) -> torch.Tensor:
+    """Compute ``x @ w + bias`` and store BF16 output.
+    ``x`` must be contiguous BF16 with shape ``(M, K)``. ``w`` must be
+    contiguous BF16 with shape ``(K, N)``. ``bias`` must be contiguous BF16
+    with shape ``(N,)``. If ``out`` is omitted, a BF16 ``(M, N)`` tensor is
+    allocated.
+    """
+    if out is None:
+        out = torch.empty(
+            (x.shape[0], w.shape[1]), device=x.device, dtype=torch.bfloat16
+        )
+    ops.bf16_linear_bias_bf16(x, w, bias, out)
+    return out
+def bf16_gemm_bias_gelu(
+    a: torch.Tensor,
+    b: torch.Tensor,
+    bias: torch.Tensor,
+    out: Optional[torch.Tensor] = None,
+) -> torch.Tensor:
+    """Compute ``gelu(a @ b + bias)`` and store BF16 output.
+    ``a`` must be contiguous BF16 with shape ``(M, K)``. ``b`` must be
+    contiguous BF16 with shape ``(K, N)``. ``bias`` must be contiguous BF16
+    with shape ``(N,)``. If ``out`` is omitted, a BF16 ``(M, N)`` tensor is
+    allocated.
+    """
+    if out is None:
+        out = torch.empty(
+            (a.shape[0], b.shape[1]), device=a.device, dtype=torch.bfloat16
+        )
+    ops.bf16_gemm_bias_gelu(a, b, bias, out)
+    return out
+def bf16_gemm_bias(
+    a: torch.Tensor,
+    b: torch.Tensor,
+    bias: torch.Tensor,
+    out: Optional[torch.Tensor] = None,
+) -> torch.Tensor:
+    """Compute ``a @ b + bias`` and store BF16 output."""
+    if out is None:
+        out = torch.empty(
+            (a.shape[0], b.shape[1]), device=a.device, dtype=torch.bfloat16
+        )
+    ops.bf16_gemm_bias(a, b, bias, out)
+    return out
+def bias_gelu_quantize_fp8_static_bf16(
+    input: torch.Tensor,
+    bias: torch.Tensor,
+    scale: torch.Tensor,
+    out: Optional[torch.Tensor] = None,
+) -> torch.Tensor:
+    """Compute ``gelu(input + bias) / scale`` and store FP8 e4m3 output.
+    ``input`` must be contiguous BF16 with shape ``(..., N)``. ``bias`` must be
+    contiguous BF16 with shape ``(N,)``. ``scale`` must be a CUDA float32 scalar
+    tensor. If ``out`` is omitted, an FP8 e4m3 output tensor is allocated.
+    """
+    if out is None:
+        out = _allocate_fp8_like(input)
+    ops.bias_gelu_quantize_fp8_static_bf16(input, bias, scale, out)
+    return out
+def gelu_quantize_fp8_static_bf16(
+    input: torch.Tensor,
+    scale: torch.Tensor,
+    out: Optional[torch.Tensor] = None,
+) -> torch.Tensor:
+    """Compute ``gelu(input) / scale`` and store FP8 e4m3 output."""
+    if out is None:
+        out = _allocate_fp8_like(input)
+    ops.gelu_quantize_fp8_static_bf16(input, scale, out)
+    return out
+def channel_scale_quantize_fp8_static_bf16(
+    input: torch.Tensor,
+    channel_scale: torch.Tensor,
+    scale: torch.Tensor,
+    out: Optional[torch.Tensor] = None,
+) -> torch.Tensor:
+    """Compute ``input * channel_scale / scale`` and store FP8 e4m3 output.
+    ``channel_scale`` must be contiguous BF16 with shape ``(input.shape[-1],)``.
+    It is broadcast over all leading input dimensions.
+    """
+    if out is None:
+        out = _allocate_fp8_like(input)
+    ops.channel_scale_quantize_fp8_static_bf16(input, channel_scale, scale, out)
+    return out
+__all__ = [
+    "bf16_linear_bf16",
+    "bf16_linear_bias_bf16",
+    "bf16_gemm_bias_gelu",
+    "bf16_gemm_bias",
+    "bias_gelu_quantize_fp8_static_bf16",
+    "channel_scale_quantize_fp8_static_bf16",
+    "gelu_quantize_fp8_static_bf16",
+]

build/torch211-cxx11-rocm72-x86_64-linux/_flashrt_gemm_epilogues_rocm_4a17df5.abi3.so ADDED Viewed

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

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

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

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

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

build/torch211-cxx11-rocm72-x86_64-linux/flashrt_gemm_epilogues/__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-rocm72-x86_64-linux/metadata.json ADDED Viewed

	@@ -0,0 +1,13 @@

+{
+  "name": "flashrt-gemm-epilogues",
+  "id": "_flashrt_gemm_epilogues_rocm_4a17df5",
+  "version": 1,
+  "license": "Apache-2.0",
+  "python-depends": [],
+  "backend": {
+    "type": "rocm",
+    "archs": [
+      "gfx942"
+    ]
+  }
+}

build/torch212-cxx11-cu126-x86_64-linux/__init__.py CHANGED Viewed

@@ -126,8 +126,14 @@ def _channel_scale_quantize_fp8_static_bf16_fake(
     return None
 def _allocate_fp8_like(input: torch.Tensor) -> torch.Tensor:
-    return torch.empty(input.shape, device=input.device, dtype=torch.float8_e4m3fn)
 def bf16_linear_bf16(

     return None
+def _fp8_dtype() -> torch.dtype:
+    if torch.version.hip is not None and hasattr(torch, "float8_e4m3fnuz"):
+        return torch.float8_e4m3fnuz
+    return torch.float8_e4m3fn
 def _allocate_fp8_like(input: torch.Tensor) -> torch.Tensor:
+    return torch.empty(input.shape, device=input.device, dtype=_fp8_dtype())
 def bf16_linear_bf16(

build/torch212-cxx11-cu126-x86_64-linux/_flashrt_gemm_epilogues_cuda_4a17df5.abi3.so ADDED Viewed

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

+version https://git-lfs.github.com/spec/v1
+oid sha256:d45f8c4f76ed67ebd7e0b63d839f2d60a8fb7d32e0056ba18527fbee08e83c6c
+size 301968

build/torch212-cxx11-cu126-x86_64-linux/_ops.py CHANGED Viewed

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

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

build/torch212-cxx11-cu126-x86_64-linux/metadata.json CHANGED Viewed

@@ -1,6 +1,6 @@
 {
   "name": "flashrt-gemm-epilogues",
-  "id": "_flashrt_gemm_epilogues_cuda_5de4768",
   "version": 1,
   "license": "Apache-2.0",
   "python-depends": [],

 {
   "name": "flashrt-gemm-epilogues",
+  "id": "_flashrt_gemm_epilogues_cuda_4a17df5",
   "version": 1,
   "license": "Apache-2.0",
   "python-depends": [],

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

@@ -126,8 +126,14 @@ def _channel_scale_quantize_fp8_static_bf16_fake(
     return None
 def _allocate_fp8_like(input: torch.Tensor) -> torch.Tensor:
-    return torch.empty(input.shape, device=input.device, dtype=torch.float8_e4m3fn)
 def bf16_linear_bf16(

     return None
+def _fp8_dtype() -> torch.dtype:
+    if torch.version.hip is not None and hasattr(torch, "float8_e4m3fnuz"):
+        return torch.float8_e4m3fnuz
+    return torch.float8_e4m3fn
 def _allocate_fp8_like(input: torch.Tensor) -> torch.Tensor:
+    return torch.empty(input.shape, device=input.device, dtype=_fp8_dtype())
 def bf16_linear_bf16(

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

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

+version https://git-lfs.github.com/spec/v1
+oid sha256:25b9331cadbe889a5ac03a849e4d37b1e155ddccdb9c00c2750ab9865c2deac6
+size 390632

build/torch212-cxx11-cu130-x86_64-linux/_flashrt_gemm_epilogues_cuda_5de4768.abi3.so DELETED Viewed

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

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

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

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

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

@@ -1,6 +1,6 @@
 {
   "name": "flashrt-gemm-epilogues",
-  "id": "_flashrt_gemm_epilogues_cuda_5de4768",
   "version": 1,
   "license": "Apache-2.0",
   "python-depends": [],

 {
   "name": "flashrt-gemm-epilogues",
+  "id": "_flashrt_gemm_epilogues_cuda_4a17df5",
   "version": 1,
   "license": "Apache-2.0",
   "python-depends": [],

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

@@ -126,8 +126,14 @@ def _channel_scale_quantize_fp8_static_bf16_fake(
     return None
 def _allocate_fp8_like(input: torch.Tensor) -> torch.Tensor:
-    return torch.empty(input.shape, device=input.device, dtype=torch.float8_e4m3fn)
 def bf16_linear_bf16(

     return None
+def _fp8_dtype() -> torch.dtype:
+    if torch.version.hip is not None and hasattr(torch, "float8_e4m3fnuz"):
+        return torch.float8_e4m3fnuz
+    return torch.float8_e4m3fn
 def _allocate_fp8_like(input: torch.Tensor) -> torch.Tensor:
+    return torch.empty(input.shape, device=input.device, dtype=_fp8_dtype())
 def bf16_linear_bf16(

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

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

+version https://git-lfs.github.com/spec/v1
+oid sha256:4119535ec007c4b1a76db5bdd1f8841415b9e6322f8cdd97f40bb191057e7efe
+size 390632

build/torch212-cxx11-cu132-x86_64-linux/_flashrt_gemm_epilogues_cuda_5de4768.abi3.so DELETED Viewed

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

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

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

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

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

@@ -1,6 +1,6 @@
 {
   "name": "flashrt-gemm-epilogues",
-  "id": "_flashrt_gemm_epilogues_cuda_5de4768",
   "version": 1,
   "license": "Apache-2.0",
   "python-depends": [],

 {
   "name": "flashrt-gemm-epilogues",
+  "id": "_flashrt_gemm_epilogues_cuda_4a17df5",
   "version": 1,
   "license": "Apache-2.0",
   "python-depends": [],

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

	@@ -0,0 +1,277 @@

+"""FlashRT fused GEMM epilogue kernels."""
+from __future__ import annotations
+import ctypes
+import ctypes.util
+from pathlib import Path
+from typing import Optional
+import torch
+def _torch_bundled_cublaslt() -> Optional[Path]:
+    for parent in Path(torch.__file__).resolve().parents:
+        candidate = parent / "nvidia" / "cublas" / "lib" / "libcublasLt.so.12"
+        if candidate.exists():
+            return candidate
+    return None
+def _preload_cublaslt() -> None:
+    bundled = _torch_bundled_cublaslt()
+    library = str(bundled) if bundled is not None else (
+        ctypes.util.find_library("cublasLt") or "libcublasLt.so"
+    )
+    try:
+        ctypes.CDLL(library, mode=ctypes.RTLD_GLOBAL)
+    except OSError:
+        pass
+_preload_cublaslt()
+from ._ops import add_op_namespace_prefix, ops
+@torch.library.register_fake(add_op_namespace_prefix("bf16_linear_bf16"))
+def _bf16_linear_bf16_fake(
+    x: torch.Tensor,
+    w: torch.Tensor,
+    out: torch.Tensor,
+) -> None:
+    if x.dim() != 2 or w.dim() != 2:
+        raise RuntimeError("x and w must be rank-2 tensors")
+    if out.shape != (x.shape[0], w.shape[1]):
+        raise RuntimeError("out shape must be (x.shape[0], w.shape[1])")
+    return None
+@torch.library.register_fake(add_op_namespace_prefix("bf16_linear_bias_bf16"))
+def _bf16_linear_bias_bf16_fake(
+    x: torch.Tensor,
+    w: torch.Tensor,
+    bias: torch.Tensor,
+    out: torch.Tensor,
+) -> None:
+    if x.dim() != 2 or w.dim() != 2:
+        raise RuntimeError("x and w must be rank-2 tensors")
+    if bias.dim() != 1:
+        raise RuntimeError("bias must be rank-1")
+    if out.shape != (x.shape[0], w.shape[1]):
+        raise RuntimeError("out shape must be (x.shape[0], w.shape[1])")
+    return None
+@torch.library.register_fake(add_op_namespace_prefix("bf16_gemm_bias_gelu"))
+def _bf16_gemm_bias_gelu_fake(
+    a: torch.Tensor,
+    b: torch.Tensor,
+    bias: torch.Tensor,
+    out: torch.Tensor,
+) -> None:
+    if a.dim() != 2 or b.dim() != 2:
+        raise RuntimeError("a and b must be rank-2 tensors")
+    if out.shape != (a.shape[0], b.shape[1]):
+        raise RuntimeError("out shape must be (a.shape[0], b.shape[1])")
+    return None
+@torch.library.register_fake(add_op_namespace_prefix("bf16_gemm_bias"))
+def _bf16_gemm_bias_fake(
+    a: torch.Tensor,
+    b: torch.Tensor,
+    bias: torch.Tensor,
+    out: torch.Tensor,
+) -> None:
+    if a.dim() != 2 or b.dim() != 2:
+        raise RuntimeError("a and b must be rank-2 tensors")
+    if out.shape != (a.shape[0], b.shape[1]):
+        raise RuntimeError("out shape must be (a.shape[0], b.shape[1])")
+    return None
+@torch.library.register_fake(add_op_namespace_prefix("bias_gelu_quantize_fp8_static_bf16"))
+def _bias_gelu_quantize_fp8_static_bf16_fake(
+    input: torch.Tensor,
+    bias: torch.Tensor,
+    scale: torch.Tensor,
+    out: torch.Tensor,
+) -> None:
+    if out.shape != input.shape:
+        raise RuntimeError("out shape must match input shape")
+    return None
+@torch.library.register_fake(add_op_namespace_prefix("gelu_quantize_fp8_static_bf16"))
+def _gelu_quantize_fp8_static_bf16_fake(
+    input: torch.Tensor,
+    scale: torch.Tensor,
+    out: torch.Tensor,
+) -> None:
+    if out.shape != input.shape:
+        raise RuntimeError("out shape must match input shape")
+    return None
+@torch.library.register_fake(add_op_namespace_prefix("channel_scale_quantize_fp8_static_bf16"))
+def _channel_scale_quantize_fp8_static_bf16_fake(
+    input: torch.Tensor,
+    channel_scale: torch.Tensor,
+    scale: torch.Tensor,
+    out: torch.Tensor,
+) -> None:
+    if out.shape != input.shape:
+        raise RuntimeError("out shape must match input shape")
+    return None
+def _fp8_dtype() -> torch.dtype:
+    if torch.version.hip is not None and hasattr(torch, "float8_e4m3fnuz"):
+        return torch.float8_e4m3fnuz
+    return torch.float8_e4m3fn
+def _allocate_fp8_like(input: torch.Tensor) -> torch.Tensor:
+    return torch.empty(input.shape, device=input.device, dtype=_fp8_dtype())
+def bf16_linear_bf16(
+    x: torch.Tensor,
+    w: torch.Tensor,
+    out: Optional[torch.Tensor] = None,
+) -> torch.Tensor:
+    """Compute ``x @ w`` and store BF16 output.
+    ``x`` must be contiguous BF16 with shape ``(M, K)``. ``w`` must be
+    contiguous BF16 with shape ``(K, N)``. If ``out`` is omitted, a BF16
+    ``(M, N)`` tensor is allocated.
+    """
+    if out is None:
+        out = torch.empty(
+            (x.shape[0], w.shape[1]), device=x.device, dtype=torch.bfloat16
+        )
+    ops.bf16_linear_bf16(x, w, out)
+    return out
+def bf16_linear_bias_bf16(
+    x: torch.Tensor,
+    w: torch.Tensor,
+    bias: torch.Tensor,
+    out: Optional[torch.Tensor] = None,
+) -> torch.Tensor:
+    """Compute ``x @ w + bias`` and store BF16 output.
+    ``x`` must be contiguous BF16 with shape ``(M, K)``. ``w`` must be
+    contiguous BF16 with shape ``(K, N)``. ``bias`` must be contiguous BF16
+    with shape ``(N,)``. If ``out`` is omitted, a BF16 ``(M, N)`` tensor is
+    allocated.
+    """
+    if out is None:
+        out = torch.empty(
+            (x.shape[0], w.shape[1]), device=x.device, dtype=torch.bfloat16
+        )
+    ops.bf16_linear_bias_bf16(x, w, bias, out)
+    return out
+def bf16_gemm_bias_gelu(
+    a: torch.Tensor,
+    b: torch.Tensor,
+    bias: torch.Tensor,
+    out: Optional[torch.Tensor] = None,
+) -> torch.Tensor:
+    """Compute ``gelu(a @ b + bias)`` and store BF16 output.
+    ``a`` must be contiguous BF16 with shape ``(M, K)``. ``b`` must be
+    contiguous BF16 with shape ``(K, N)``. ``bias`` must be contiguous BF16
+    with shape ``(N,)``. If ``out`` is omitted, a BF16 ``(M, N)`` tensor is
+    allocated.
+    """
+    if out is None:
+        out = torch.empty(
+            (a.shape[0], b.shape[1]), device=a.device, dtype=torch.bfloat16
+        )
+    ops.bf16_gemm_bias_gelu(a, b, bias, out)
+    return out
+def bf16_gemm_bias(
+    a: torch.Tensor,
+    b: torch.Tensor,
+    bias: torch.Tensor,
+    out: Optional[torch.Tensor] = None,
+) -> torch.Tensor:
+    """Compute ``a @ b + bias`` and store BF16 output."""
+    if out is None:
+        out = torch.empty(
+            (a.shape[0], b.shape[1]), device=a.device, dtype=torch.bfloat16
+        )
+    ops.bf16_gemm_bias(a, b, bias, out)
+    return out
+def bias_gelu_quantize_fp8_static_bf16(
+    input: torch.Tensor,
+    bias: torch.Tensor,
+    scale: torch.Tensor,
+    out: Optional[torch.Tensor] = None,
+) -> torch.Tensor:
+    """Compute ``gelu(input + bias) / scale`` and store FP8 e4m3 output.
+    ``input`` must be contiguous BF16 with shape ``(..., N)``. ``bias`` must be
+    contiguous BF16 with shape ``(N,)``. ``scale`` must be a CUDA float32 scalar
+    tensor. If ``out`` is omitted, an FP8 e4m3 output tensor is allocated.
+    """
+    if out is None:
+        out = _allocate_fp8_like(input)
+    ops.bias_gelu_quantize_fp8_static_bf16(input, bias, scale, out)
+    return out
+def gelu_quantize_fp8_static_bf16(
+    input: torch.Tensor,
+    scale: torch.Tensor,
+    out: Optional[torch.Tensor] = None,
+) -> torch.Tensor:
+    """Compute ``gelu(input) / scale`` and store FP8 e4m3 output."""
+    if out is None:
+        out = _allocate_fp8_like(input)
+    ops.gelu_quantize_fp8_static_bf16(input, scale, out)
+    return out
+def channel_scale_quantize_fp8_static_bf16(
+    input: torch.Tensor,
+    channel_scale: torch.Tensor,
+    scale: torch.Tensor,
+    out: Optional[torch.Tensor] = None,
+) -> torch.Tensor:
+    """Compute ``input * channel_scale / scale`` and store FP8 e4m3 output.
+    ``channel_scale`` must be contiguous BF16 with shape ``(input.shape[-1],)``.
+    It is broadcast over all leading input dimensions.
+    """
+    if out is None:
+        out = _allocate_fp8_like(input)
+    ops.channel_scale_quantize_fp8_static_bf16(input, channel_scale, scale, out)
+    return out
+__all__ = [
+    "bf16_linear_bf16",
+    "bf16_linear_bias_bf16",
+    "bf16_gemm_bias_gelu",
+    "bf16_gemm_bias",
+    "bias_gelu_quantize_fp8_static_bf16",
+    "channel_scale_quantize_fp8_static_bf16",
+    "gelu_quantize_fp8_static_bf16",
+]

build/torch212-cxx11-rocm71-x86_64-linux/_flashrt_gemm_epilogues_rocm_4a17df5.abi3.so ADDED Viewed

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

+version https://git-lfs.github.com/spec/v1
+oid sha256:31963b516aa5c67cad683bb0a2f7caa3250e2adb100f97df649043eb8f2f4fb7
+size 155688

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

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

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

build/torch212-cxx11-rocm71-x86_64-linux/flashrt_gemm_epilogues/__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-rocm71-x86_64-linux/metadata.json ADDED Viewed

	@@ -0,0 +1,13 @@

+{
+  "name": "flashrt-gemm-epilogues",
+  "id": "_flashrt_gemm_epilogues_rocm_4a17df5",
+  "version": 1,
+  "license": "Apache-2.0",
+  "python-depends": [],
+  "backend": {
+    "type": "rocm",
+    "archs": [
+      "gfx942"
+    ]
+  }
+}

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

	@@ -0,0 +1,277 @@

+"""FlashRT fused GEMM epilogue kernels."""
+from __future__ import annotations
+import ctypes
+import ctypes.util
+from pathlib import Path
+from typing import Optional
+import torch
+def _torch_bundled_cublaslt() -> Optional[Path]:
+    for parent in Path(torch.__file__).resolve().parents:
+        candidate = parent / "nvidia" / "cublas" / "lib" / "libcublasLt.so.12"
+        if candidate.exists():
+            return candidate
+    return None
+def _preload_cublaslt() -> None:
+    bundled = _torch_bundled_cublaslt()
+    library = str(bundled) if bundled is not None else (
+        ctypes.util.find_library("cublasLt") or "libcublasLt.so"
+    )
+    try:
+        ctypes.CDLL(library, mode=ctypes.RTLD_GLOBAL)
+    except OSError:
+        pass
+_preload_cublaslt()
+from ._ops import add_op_namespace_prefix, ops
+@torch.library.register_fake(add_op_namespace_prefix("bf16_linear_bf16"))
+def _bf16_linear_bf16_fake(
+    x: torch.Tensor,
+    w: torch.Tensor,
+    out: torch.Tensor,
+) -> None:
+    if x.dim() != 2 or w.dim() != 2:
+        raise RuntimeError("x and w must be rank-2 tensors")
+    if out.shape != (x.shape[0], w.shape[1]):
+        raise RuntimeError("out shape must be (x.shape[0], w.shape[1])")
+    return None
+@torch.library.register_fake(add_op_namespace_prefix("bf16_linear_bias_bf16"))
+def _bf16_linear_bias_bf16_fake(
+    x: torch.Tensor,
+    w: torch.Tensor,
+    bias: torch.Tensor,
+    out: torch.Tensor,
+) -> None:
+    if x.dim() != 2 or w.dim() != 2:
+        raise RuntimeError("x and w must be rank-2 tensors")
+    if bias.dim() != 1:
+        raise RuntimeError("bias must be rank-1")
+    if out.shape != (x.shape[0], w.shape[1]):
+        raise RuntimeError("out shape must be (x.shape[0], w.shape[1])")
+    return None
+@torch.library.register_fake(add_op_namespace_prefix("bf16_gemm_bias_gelu"))
+def _bf16_gemm_bias_gelu_fake(
+    a: torch.Tensor,
+    b: torch.Tensor,
+    bias: torch.Tensor,
+    out: torch.Tensor,
+) -> None:
+    if a.dim() != 2 or b.dim() != 2:
+        raise RuntimeError("a and b must be rank-2 tensors")
+    if out.shape != (a.shape[0], b.shape[1]):
+        raise RuntimeError("out shape must be (a.shape[0], b.shape[1])")
+    return None
+@torch.library.register_fake(add_op_namespace_prefix("bf16_gemm_bias"))
+def _bf16_gemm_bias_fake(
+    a: torch.Tensor,
+    b: torch.Tensor,
+    bias: torch.Tensor,
+    out: torch.Tensor,
+) -> None:
+    if a.dim() != 2 or b.dim() != 2:
+        raise RuntimeError("a and b must be rank-2 tensors")
+    if out.shape != (a.shape[0], b.shape[1]):
+        raise RuntimeError("out shape must be (a.shape[0], b.shape[1])")
+    return None
+@torch.library.register_fake(add_op_namespace_prefix("bias_gelu_quantize_fp8_static_bf16"))
+def _bias_gelu_quantize_fp8_static_bf16_fake(
+    input: torch.Tensor,
+    bias: torch.Tensor,
+    scale: torch.Tensor,
+    out: torch.Tensor,
+) -> None:
+    if out.shape != input.shape:
+        raise RuntimeError("out shape must match input shape")
+    return None
+@torch.library.register_fake(add_op_namespace_prefix("gelu_quantize_fp8_static_bf16"))
+def _gelu_quantize_fp8_static_bf16_fake(
+    input: torch.Tensor,
+    scale: torch.Tensor,
+    out: torch.Tensor,
+) -> None:
+    if out.shape != input.shape:
+        raise RuntimeError("out shape must match input shape")
+    return None
+@torch.library.register_fake(add_op_namespace_prefix("channel_scale_quantize_fp8_static_bf16"))
+def _channel_scale_quantize_fp8_static_bf16_fake(
+    input: torch.Tensor,
+    channel_scale: torch.Tensor,
+    scale: torch.Tensor,
+    out: torch.Tensor,
+) -> None:
+    if out.shape != input.shape:
+        raise RuntimeError("out shape must match input shape")
+    return None
+def _fp8_dtype() -> torch.dtype:
+    if torch.version.hip is not None and hasattr(torch, "float8_e4m3fnuz"):
+        return torch.float8_e4m3fnuz
+    return torch.float8_e4m3fn
+def _allocate_fp8_like(input: torch.Tensor) -> torch.Tensor:
+    return torch.empty(input.shape, device=input.device, dtype=_fp8_dtype())
+def bf16_linear_bf16(
+    x: torch.Tensor,
+    w: torch.Tensor,
+    out: Optional[torch.Tensor] = None,
+) -> torch.Tensor:
+    """Compute ``x @ w`` and store BF16 output.
+    ``x`` must be contiguous BF16 with shape ``(M, K)``. ``w`` must be
+    contiguous BF16 with shape ``(K, N)``. If ``out`` is omitted, a BF16
+    ``(M, N)`` tensor is allocated.
+    """
+    if out is None:
+        out = torch.empty(
+            (x.shape[0], w.shape[1]), device=x.device, dtype=torch.bfloat16
+        )
+    ops.bf16_linear_bf16(x, w, out)
+    return out
+def bf16_linear_bias_bf16(
+    x: torch.Tensor,
+    w: torch.Tensor,
+    bias: torch.Tensor,
+    out: Optional[torch.Tensor] = None,
+) -> torch.Tensor:
+    """Compute ``x @ w + bias`` and store BF16 output.
+    ``x`` must be contiguous BF16 with shape ``(M, K)``. ``w`` must be
+    contiguous BF16 with shape ``(K, N)``. ``bias`` must be contiguous BF16
+    with shape ``(N,)``. If ``out`` is omitted, a BF16 ``(M, N)`` tensor is
+    allocated.
+    """
+    if out is None:
+        out = torch.empty(
+            (x.shape[0], w.shape[1]), device=x.device, dtype=torch.bfloat16
+        )
+    ops.bf16_linear_bias_bf16(x, w, bias, out)
+    return out
+def bf16_gemm_bias_gelu(
+    a: torch.Tensor,
+    b: torch.Tensor,
+    bias: torch.Tensor,
+    out: Optional[torch.Tensor] = None,
+) -> torch.Tensor:
+    """Compute ``gelu(a @ b + bias)`` and store BF16 output.
+    ``a`` must be contiguous BF16 with shape ``(M, K)``. ``b`` must be
+    contiguous BF16 with shape ``(K, N)``. ``bias`` must be contiguous BF16
+    with shape ``(N,)``. If ``out`` is omitted, a BF16 ``(M, N)`` tensor is
+    allocated.
+    """
+    if out is None:
+        out = torch.empty(
+            (a.shape[0], b.shape[1]), device=a.device, dtype=torch.bfloat16
+        )
+    ops.bf16_gemm_bias_gelu(a, b, bias, out)
+    return out
+def bf16_gemm_bias(
+    a: torch.Tensor,
+    b: torch.Tensor,
+    bias: torch.Tensor,
+    out: Optional[torch.Tensor] = None,
+) -> torch.Tensor:
+    """Compute ``a @ b + bias`` and store BF16 output."""
+    if out is None:
+        out = torch.empty(
+            (a.shape[0], b.shape[1]), device=a.device, dtype=torch.bfloat16
+        )
+    ops.bf16_gemm_bias(a, b, bias, out)
+    return out
+def bias_gelu_quantize_fp8_static_bf16(
+    input: torch.Tensor,
+    bias: torch.Tensor,
+    scale: torch.Tensor,
+    out: Optional[torch.Tensor] = None,
+) -> torch.Tensor:
+    """Compute ``gelu(input + bias) / scale`` and store FP8 e4m3 output.
+    ``input`` must be contiguous BF16 with shape ``(..., N)``. ``bias`` must be
+    contiguous BF16 with shape ``(N,)``. ``scale`` must be a CUDA float32 scalar
+    tensor. If ``out`` is omitted, an FP8 e4m3 output tensor is allocated.
+    """
+    if out is None:
+        out = _allocate_fp8_like(input)
+    ops.bias_gelu_quantize_fp8_static_bf16(input, bias, scale, out)
+    return out
+def gelu_quantize_fp8_static_bf16(
+    input: torch.Tensor,
+    scale: torch.Tensor,
+    out: Optional[torch.Tensor] = None,
+) -> torch.Tensor:
+    """Compute ``gelu(input) / scale`` and store FP8 e4m3 output."""
+    if out is None:
+        out = _allocate_fp8_like(input)
+    ops.gelu_quantize_fp8_static_bf16(input, scale, out)
+    return out
+def channel_scale_quantize_fp8_static_bf16(
+    input: torch.Tensor,
+    channel_scale: torch.Tensor,
+    scale: torch.Tensor,
+    out: Optional[torch.Tensor] = None,
+) -> torch.Tensor:
+    """Compute ``input * channel_scale / scale`` and store FP8 e4m3 output.
+    ``channel_scale`` must be contiguous BF16 with shape ``(input.shape[-1],)``.
+    It is broadcast over all leading input dimensions.
+    """
+    if out is None:
+        out = _allocate_fp8_like(input)
+    ops.channel_scale_quantize_fp8_static_bf16(input, channel_scale, scale, out)
+    return out
+__all__ = [
+    "bf16_linear_bf16",
+    "bf16_linear_bias_bf16",
+    "bf16_gemm_bias_gelu",
+    "bf16_gemm_bias",
+    "bias_gelu_quantize_fp8_static_bf16",
+    "channel_scale_quantize_fp8_static_bf16",
+    "gelu_quantize_fp8_static_bf16",
+]

build/torch212-cxx11-rocm72-x86_64-linux/_flashrt_gemm_epilogues_rocm_4a17df5.abi3.so ADDED Viewed

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

+version https://git-lfs.github.com/spec/v1
+oid sha256:c21f5a7cb50ef45ce45bd825e801629525ee5a1cbcbcc7d71a835499e964e691
+size 154344

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

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

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

build/torch212-cxx11-rocm72-x86_64-linux/flashrt_gemm_epilogues/__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-rocm72-x86_64-linux/metadata.json ADDED Viewed

	@@ -0,0 +1,13 @@

+{
+  "name": "flashrt-gemm-epilogues",
+  "id": "_flashrt_gemm_epilogues_rocm_4a17df5",
+  "version": 1,
+  "license": "Apache-2.0",
+  "python-depends": [],
+  "backend": {
+    "type": "rocm",
+    "archs": [
+      "gfx942"
+    ]
+  }
+}