liangsu9988 commited on 18 days ago

Commit

34562e4

verified ·

1 Parent(s): b208e84

Uploaded using `kernel-builder`.

Browse files

Files changed (37) hide show

benchmarks/benchmark.py +17 -5
build/torch211-cxx11-cu128-x86_64-linux/__init__.py +8 -2
build/torch211-cxx11-cu128-x86_64-linux/{_flashrt_fp8_ffn_cuda_5de4768.abi3.so → _flashrt_fp8_ffn_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 +8 -2
build/torch211-cxx11-cu130-x86_64-linux/{_flashrt_fp8_ffn_cuda_5de4768.abi3.so → _flashrt_fp8_ffn_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 +227 -0
build/{torch212-cxx11-cu130-x86_64-linux/_flashrt_fp8_ffn_cuda_5de4768.abi3.so → torch211-cxx11-rocm71-x86_64-linux/_flashrt_fp8_ffn_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_fp8_ffn/__init__.py +26 -0
build/torch211-cxx11-rocm71-x86_64-linux/metadata.json +13 -0
build/torch211-cxx11-rocm72-x86_64-linux/__init__.py +227 -0
build/{torch212-cxx11-cu132-x86_64-linux/_flashrt_fp8_ffn_cuda_5de4768.abi3.so → torch211-cxx11-rocm72-x86_64-linux/_flashrt_fp8_ffn_rocm_4a17df5.abi3.so} +2 -2
build/torch211-cxx11-rocm72-x86_64-linux/_ops.py +9 -0
build/torch211-cxx11-rocm72-x86_64-linux/flashrt_fp8_ffn/__init__.py +26 -0
build/torch211-cxx11-rocm72-x86_64-linux/metadata.json +13 -0
build/torch212-cxx11-cu130-x86_64-linux/__init__.py +8 -2
build/torch212-cxx11-cu130-x86_64-linux/_flashrt_fp8_ffn_cuda_4a17df5.abi3.so +3 -0
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 +8 -2
build/torch212-cxx11-cu132-x86_64-linux/_flashrt_fp8_ffn_cuda_4a17df5.abi3.so +3 -0
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 +227 -0
build/torch212-cxx11-rocm71-x86_64-linux/_flashrt_fp8_ffn_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_fp8_ffn/__init__.py +26 -0
build/torch212-cxx11-rocm71-x86_64-linux/metadata.json +13 -0
build/torch212-cxx11-rocm72-x86_64-linux/__init__.py +227 -0
build/torch212-cxx11-rocm72-x86_64-linux/_flashrt_fp8_ffn_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_fp8_ffn/__init__.py +26 -0
build/torch212-cxx11-rocm72-x86_64-linux/metadata.json +13 -0

benchmarks/benchmark.py CHANGED Viewed

@@ -133,7 +133,7 @@ class SourceOps:
         if hidden is None:
             hidden = torch.empty((x.shape[0], up_w.shape[0]), device=x.device, dtype=torch.bfloat16)
         if hidden_fp8 is None:
-            hidden_fp8 = torch.empty_like(hidden, dtype=torch.float8_e4m3fn)
         if out is None:
             out = torch.empty((x.shape[0], down_w.shape[0]), device=x.device, dtype=torch.bfloat16)
         self._ops.fp8_gelu_mlp_bf16(
@@ -204,11 +204,22 @@ def load_installed_ops(artifact: str | None):
 def load_hub_ops(repo_id: str, version: int):
     from kernels import get_kernel
-    return get_kernel(repo_id, version=version, trust_remote_code=True)
 def quantize_fp8(x: torch.Tensor, scale: torch.Tensor) -> torch.Tensor:
-    return torch.clamp(x.float() / scale.float(), -448.0, 448.0).to(torch.float8_e4m3fn)
 def dequant_fp8(x: torch.Tensor, scale: torch.Tensor) -> torch.Tensor:
@@ -242,7 +253,8 @@ stable_bf16_bias_add = compiler_disable(bf16_bias_add_boundary)
 def torch_mlp(x, up_w, up_b, down_w, down_b, x_s, up_s, hid_s, dn_s):
     hidden = (dequant_fp8(x, x_s) @ dequant_fp8(up_w, up_s).T).to(torch.bfloat16)
     hidden = torch.nn.functional.gelu(hidden + up_b.float(), approximate="tanh")
-    hidden_fp8 = torch.clamp(hidden / hid_s.float(), -448.0, 448.0).to(torch.float8_e4m3fn)
     out = (dequant_fp8(hidden_fp8, hid_s) @ dequant_fp8(down_w, dn_s).T).to(torch.bfloat16)
     return (out + down_b.float()).to(torch.bfloat16)
@@ -274,7 +286,7 @@ def make_inputs(M: int, K: int, H: int, N: int, layers: int):
     up_bs = [torch.randn((H,), device="cuda", dtype=torch.bfloat16) for _ in range(layers)]
     down_bs = [torch.randn((N,), device="cuda", dtype=torch.bfloat16) for _ in range(layers)]
     hidden = [torch.empty((M, H), device="cuda", dtype=torch.bfloat16) for _ in range(layers)]
-    hidden_fp8 = [torch.empty((M, H), device="cuda", dtype=torch.float8_e4m3fn) for _ in range(layers)]
     outs = [torch.empty((M, N), device="cuda", dtype=torch.bfloat16) for _ in range(layers)]
     return xs, up_ws, up_bs, down_ws, down_bs, x_scale, up_scale, hidden_scale, down_scale, hidden, hidden_fp8, outs

         if hidden is None:
             hidden = torch.empty((x.shape[0], up_w.shape[0]), device=x.device, dtype=torch.bfloat16)
         if hidden_fp8 is None:
+            hidden_fp8 = torch.empty_like(hidden, dtype=fp8_dtype())
         if out is None:
             out = torch.empty((x.shape[0], down_w.shape[0]), device=x.device, dtype=torch.bfloat16)
         self._ops.fp8_gelu_mlp_bf16(
 def load_hub_ops(repo_id: str, version: int):
     from kernels import get_kernel
+    return get_kernel(repo_id, version=version)
+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 quantize_fp8(x: torch.Tensor, scale: torch.Tensor) -> torch.Tensor:
+    limit = fp8_max()
+    return torch.clamp(x.float() / scale.float(), -limit, limit).to(fp8_dtype())
 def dequant_fp8(x: torch.Tensor, scale: torch.Tensor) -> torch.Tensor:
 def torch_mlp(x, up_w, up_b, down_w, down_b, x_s, up_s, hid_s, dn_s):
     hidden = (dequant_fp8(x, x_s) @ dequant_fp8(up_w, up_s).T).to(torch.bfloat16)
     hidden = torch.nn.functional.gelu(hidden + up_b.float(), approximate="tanh")
+    limit = fp8_max()
+    hidden_fp8 = torch.clamp(hidden / hid_s.float(), -limit, limit).to(fp8_dtype())
     out = (dequant_fp8(hidden_fp8, hid_s) @ dequant_fp8(down_w, dn_s).T).to(torch.bfloat16)
     return (out + down_b.float()).to(torch.bfloat16)
     up_bs = [torch.randn((H,), device="cuda", dtype=torch.bfloat16) for _ in range(layers)]
     down_bs = [torch.randn((N,), device="cuda", dtype=torch.bfloat16) for _ in range(layers)]
     hidden = [torch.empty((M, H), device="cuda", dtype=torch.bfloat16) for _ in range(layers)]
+    hidden_fp8 = [torch.empty((M, H), device="cuda", dtype=fp8_dtype()) for _ in range(layers)]
     outs = [torch.empty((M, N), device="cuda", dtype=torch.bfloat16) for _ in range(layers)]
     return xs, up_ws, up_bs, down_ws, down_bs, x_scale, up_scale, hidden_scale, down_scale, hidden, hidden_fp8, outs

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

@@ -34,6 +34,12 @@ _preload_cublaslt()
 from ._ops import add_op_namespace_prefix, ops
 @torch.library.register_fake(add_op_namespace_prefix("fp8_gemm_bf16"))
 def _fp8_gemm_bf16_fake(
     input: torch.Tensor,
@@ -146,7 +152,7 @@ def fp8_linear_bias_gelu_quant_bf16(
             dtype=torch.bfloat16,
         )
     if out_fp8 is None:
-        out_fp8 = torch.empty_like(hidden_bf16, dtype=torch.float8_e4m3fn)
     ops.fp8_linear_bias_gelu_quant_bf16(
         input,
         weight,
@@ -190,7 +196,7 @@ def fp8_gelu_mlp_bf16(
             dtype=torch.bfloat16,
         )
     if hidden_fp8 is None:
-        hidden_fp8 = torch.empty_like(hidden_bf16, dtype=torch.float8_e4m3fn)
     if out is None:
         out = torch.empty(
             (input.shape[0], down_weight.shape[0]),

 from ._ops import add_op_namespace_prefix, ops
+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
 @torch.library.register_fake(add_op_namespace_prefix("fp8_gemm_bf16"))
 def _fp8_gemm_bf16_fake(
     input: torch.Tensor,
             dtype=torch.bfloat16,
         )
     if out_fp8 is None:
+        out_fp8 = torch.empty_like(hidden_bf16, dtype=_fp8_dtype())
     ops.fp8_linear_bias_gelu_quant_bf16(
         input,
         weight,
             dtype=torch.bfloat16,
         )
     if hidden_fp8 is None:
+        hidden_fp8 = torch.empty_like(hidden_bf16, dtype=_fp8_dtype())
     if out is None:
         out = torch.empty(
             (input.shape[0], down_weight.shape[0]),

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

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

 version https://git-lfs.github.com/spec/v1
+oid sha256:49d584720e4739dbf34f58426b49bd60f9ce619de6956132177812c850af7479
+size 300056

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

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

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

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

@@ -1,6 +1,6 @@
 {
   "name": "flashrt-fp8-ffn",
-  "id": "_flashrt_fp8_ffn_cuda_5de4768",
   "version": 1,
   "license": "Apache-2.0",
   "python-depends": [],

 {
   "name": "flashrt-fp8-ffn",
+  "id": "_flashrt_fp8_ffn_cuda_4a17df5",
   "version": 1,
   "license": "Apache-2.0",
   "python-depends": [],

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

@@ -34,6 +34,12 @@ _preload_cublaslt()
 from ._ops import add_op_namespace_prefix, ops
 @torch.library.register_fake(add_op_namespace_prefix("fp8_gemm_bf16"))
 def _fp8_gemm_bf16_fake(
     input: torch.Tensor,
@@ -146,7 +152,7 @@ def fp8_linear_bias_gelu_quant_bf16(
             dtype=torch.bfloat16,
         )
     if out_fp8 is None:
-        out_fp8 = torch.empty_like(hidden_bf16, dtype=torch.float8_e4m3fn)
     ops.fp8_linear_bias_gelu_quant_bf16(
         input,
         weight,
@@ -190,7 +196,7 @@ def fp8_gelu_mlp_bf16(
             dtype=torch.bfloat16,
         )
     if hidden_fp8 is None:
-        hidden_fp8 = torch.empty_like(hidden_bf16, dtype=torch.float8_e4m3fn)
     if out is None:
         out = torch.empty(
             (input.shape[0], down_weight.shape[0]),

 from ._ops import add_op_namespace_prefix, ops
+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
 @torch.library.register_fake(add_op_namespace_prefix("fp8_gemm_bf16"))
 def _fp8_gemm_bf16_fake(
     input: torch.Tensor,
             dtype=torch.bfloat16,
         )
     if out_fp8 is None:
+        out_fp8 = torch.empty_like(hidden_bf16, dtype=_fp8_dtype())
     ops.fp8_linear_bias_gelu_quant_bf16(
         input,
         weight,
             dtype=torch.bfloat16,
         )
     if hidden_fp8 is None:
+        hidden_fp8 = torch.empty_like(hidden_bf16, dtype=_fp8_dtype())
     if out is None:
         out = torch.empty(
             (input.shape[0], down_weight.shape[0]),

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

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

 version https://git-lfs.github.com/spec/v1
+oid sha256:3113678b5434234f55790ee0a215dffc4da6936faa0af7f6e1822211a720372f
+size 298176

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

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

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

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

@@ -1,6 +1,6 @@
 {
   "name": "flashrt-fp8-ffn",
-  "id": "_flashrt_fp8_ffn_cuda_5de4768",
   "version": 1,
   "license": "Apache-2.0",
   "python-depends": [],

 {
   "name": "flashrt-fp8-ffn",
+  "id": "_flashrt_fp8_ffn_cuda_4a17df5",
   "version": 1,
   "license": "Apache-2.0",
   "python-depends": [],

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

	@@ -0,0 +1,227 @@

+"""FlashRT FP8 FFN 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
+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
+@torch.library.register_fake(add_op_namespace_prefix("fp8_gemm_bf16"))
+def _fp8_gemm_bf16_fake(
+    input: torch.Tensor,
+    weight: torch.Tensor,
+    input_scale: torch.Tensor,
+    weight_scale: torch.Tensor,
+    out: torch.Tensor,
+) -> None:
+    if input.dim() != 2 or weight.dim() != 2:
+        raise RuntimeError("input and weight must be rank-2 tensors")
+    if out.shape != (input.shape[0], weight.shape[0]):
+        raise RuntimeError("out shape must be (input.shape[0], weight.shape[0])")
+    return None
+@torch.library.register_fake(add_op_namespace_prefix("fp8_linear_bias_gelu_quant_bf16"))
+def _fp8_linear_bias_gelu_quant_bf16_fake(
+    input: torch.Tensor,
+    weight: torch.Tensor,
+    bias: torch.Tensor,
+    input_scale: torch.Tensor,
+    weight_scale: torch.Tensor,
+    output_scale: torch.Tensor,
+    hidden_bf16: torch.Tensor,
+    out_fp8: torch.Tensor,
+) -> None:
+    expected = (input.shape[0], weight.shape[0])
+    if hidden_bf16.shape != expected or out_fp8.shape != expected:
+        raise RuntimeError(
+            "hidden_bf16 and out_fp8 shapes must be "
+            "(input.shape[0], weight.shape[0])"
+        )
+    return None
+@torch.library.register_fake(add_op_namespace_prefix("fp8_gelu_mlp_bf16"))
+def _fp8_gelu_mlp_bf16_fake(
+    input: torch.Tensor,
+    up_weight: torch.Tensor,
+    up_bias: torch.Tensor,
+    down_weight: torch.Tensor,
+    down_bias: torch.Tensor,
+    input_scale: torch.Tensor,
+    up_weight_scale: torch.Tensor,
+    hidden_scale: torch.Tensor,
+    down_weight_scale: torch.Tensor,
+    hidden_bf16: torch.Tensor,
+    hidden_fp8: torch.Tensor,
+    out: torch.Tensor,
+) -> None:
+    hidden_shape = (input.shape[0], up_weight.shape[0])
+    out_shape = (input.shape[0], down_weight.shape[0])
+    if hidden_bf16.shape != hidden_shape or hidden_fp8.shape != hidden_shape:
+        raise RuntimeError(
+            "hidden buffers must be (input.shape[0], up_weight.shape[0])"
+        )
+    if out.shape != out_shape:
+        raise RuntimeError("out shape must be (input.shape[0], down_weight.shape[0])")
+    return None
+def _scalar_scale_like(input: torch.Tensor, value: float = 1.0) -> torch.Tensor:
+    return torch.tensor([value], device=input.device, dtype=torch.float32)
+def fp8_gemm_bf16(
+    input: torch.Tensor,
+    weight: torch.Tensor,
+    input_scale: torch.Tensor,
+    weight_scale: torch.Tensor,
+    out: torch.Tensor | None = None,
+) -> torch.Tensor:
+    """Compute ``(input * input_scale) @ (weight * weight_scale).T``.
+    ``input`` is FP8 E4M3 with shape ``(M, K)``. ``weight`` is FP8 E4M3 with
+    shape ``(N, K)``. ``input_scale`` and ``weight_scale`` are CUDA float32
+    scalar tensors. Output is BF16 with shape ``(M, N)``.
+    """
+    if out is None:
+        out = torch.empty(
+            (input.shape[0], weight.shape[0]),
+            device=input.device,
+            dtype=torch.bfloat16,
+        )
+    ops.fp8_gemm_bf16(input, weight, input_scale, weight_scale, out)
+    return out
+def fp8_linear_bias_gelu_quant_bf16(
+    input: torch.Tensor,
+    weight: torch.Tensor,
+    bias: torch.Tensor,
+    input_scale: torch.Tensor,
+    weight_scale: torch.Tensor,
+    output_scale: torch.Tensor,
+    hidden_bf16: torch.Tensor | None = None,
+    out_fp8: torch.Tensor | None = None,
+) -> tuple[torch.Tensor, torch.Tensor]:
+    """FP8 linear + BF16 bias/GELU + FP8 quantized output.
+    Returns ``(hidden_bf16, out_fp8)``. ``hidden_bf16`` is the post-GEMM
+    pre-activation scratch; ``out_fp8`` is the quantized activation.
+    """
+    if hidden_bf16 is None:
+        hidden_bf16 = torch.empty(
+            (input.shape[0], weight.shape[0]),
+            device=input.device,
+            dtype=torch.bfloat16,
+        )
+    if out_fp8 is None:
+        out_fp8 = torch.empty_like(hidden_bf16, dtype=_fp8_dtype())
+    ops.fp8_linear_bias_gelu_quant_bf16(
+        input,
+        weight,
+        bias,
+        input_scale,
+        weight_scale,
+        output_scale,
+        hidden_bf16,
+        out_fp8,
+    )
+    return hidden_bf16, out_fp8
+def fp8_gelu_mlp_bf16(
+    input: torch.Tensor,
+    up_weight: torch.Tensor,
+    up_bias: torch.Tensor,
+    down_weight: torch.Tensor,
+    down_bias: torch.Tensor,
+    input_scale: torch.Tensor,
+    up_weight_scale: torch.Tensor,
+    hidden_scale: torch.Tensor,
+    down_weight_scale: torch.Tensor,
+    hidden_bf16: torch.Tensor | None = None,
+    hidden_fp8: torch.Tensor | None = None,
+    out: torch.Tensor | None = None,
+) -> torch.Tensor:
+    """FP8 GELU MLP block with BF16 output.
+    Computes:
+    ``hidden = gelu(fp8_gemm(input, up_weight) + up_bias)``
+    ``hidden_fp8 = quantize_fp8(hidden, hidden_scale)``
+    ``out = fp8_gemm(hidden_fp8, down_weight) + down_bias``
+    """
+    if hidden_bf16 is None:
+        hidden_bf16 = torch.empty(
+            (input.shape[0], up_weight.shape[0]),
+            device=input.device,
+            dtype=torch.bfloat16,
+        )
+    if hidden_fp8 is None:
+        hidden_fp8 = torch.empty_like(hidden_bf16, dtype=_fp8_dtype())
+    if out is None:
+        out = torch.empty(
+            (input.shape[0], down_weight.shape[0]),
+            device=input.device,
+            dtype=torch.bfloat16,
+        )
+    ops.fp8_gelu_mlp_bf16(
+        input,
+        up_weight,
+        up_bias,
+        down_weight,
+        down_bias,
+        input_scale,
+        up_weight_scale,
+        hidden_scale,
+        down_weight_scale,
+        hidden_bf16,
+        hidden_fp8,
+        out,
+    )
+    return out
+__all__ = [
+    "fp8_gemm_bf16",
+    "fp8_gelu_mlp_bf16",
+    "fp8_linear_bias_gelu_quant_bf16",
+]

build/{torch212-cxx11-cu130-x86_64-linux/_flashrt_fp8_ffn_cuda_5de4768.abi3.so → torch211-cxx11-rocm71-x86_64-linux/_flashrt_fp8_ffn_rocm_4a17df5.abi3.so} RENAMED Viewed

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

 version https://git-lfs.github.com/spec/v1
+oid sha256:9d56b2ca87eae36419a8aeee3fe444a1894f7c1024c70eaabdea7848201374e3
+size 145192

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

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

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

build/torch211-cxx11-rocm71-x86_64-linux/flashrt_fp8_ffn/__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-fp8-ffn",
+  "id": "_flashrt_fp8_ffn_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,227 @@

+"""FlashRT FP8 FFN 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
+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
+@torch.library.register_fake(add_op_namespace_prefix("fp8_gemm_bf16"))
+def _fp8_gemm_bf16_fake(
+    input: torch.Tensor,
+    weight: torch.Tensor,
+    input_scale: torch.Tensor,
+    weight_scale: torch.Tensor,
+    out: torch.Tensor,
+) -> None:
+    if input.dim() != 2 or weight.dim() != 2:
+        raise RuntimeError("input and weight must be rank-2 tensors")
+    if out.shape != (input.shape[0], weight.shape[0]):
+        raise RuntimeError("out shape must be (input.shape[0], weight.shape[0])")
+    return None
+@torch.library.register_fake(add_op_namespace_prefix("fp8_linear_bias_gelu_quant_bf16"))
+def _fp8_linear_bias_gelu_quant_bf16_fake(
+    input: torch.Tensor,
+    weight: torch.Tensor,
+    bias: torch.Tensor,
+    input_scale: torch.Tensor,
+    weight_scale: torch.Tensor,
+    output_scale: torch.Tensor,
+    hidden_bf16: torch.Tensor,
+    out_fp8: torch.Tensor,
+) -> None:
+    expected = (input.shape[0], weight.shape[0])
+    if hidden_bf16.shape != expected or out_fp8.shape != expected:
+        raise RuntimeError(
+            "hidden_bf16 and out_fp8 shapes must be "
+            "(input.shape[0], weight.shape[0])"
+        )
+    return None
+@torch.library.register_fake(add_op_namespace_prefix("fp8_gelu_mlp_bf16"))
+def _fp8_gelu_mlp_bf16_fake(
+    input: torch.Tensor,
+    up_weight: torch.Tensor,
+    up_bias: torch.Tensor,
+    down_weight: torch.Tensor,
+    down_bias: torch.Tensor,
+    input_scale: torch.Tensor,
+    up_weight_scale: torch.Tensor,
+    hidden_scale: torch.Tensor,
+    down_weight_scale: torch.Tensor,
+    hidden_bf16: torch.Tensor,
+    hidden_fp8: torch.Tensor,
+    out: torch.Tensor,
+) -> None:
+    hidden_shape = (input.shape[0], up_weight.shape[0])
+    out_shape = (input.shape[0], down_weight.shape[0])
+    if hidden_bf16.shape != hidden_shape or hidden_fp8.shape != hidden_shape:
+        raise RuntimeError(
+            "hidden buffers must be (input.shape[0], up_weight.shape[0])"
+        )
+    if out.shape != out_shape:
+        raise RuntimeError("out shape must be (input.shape[0], down_weight.shape[0])")
+    return None
+def _scalar_scale_like(input: torch.Tensor, value: float = 1.0) -> torch.Tensor:
+    return torch.tensor([value], device=input.device, dtype=torch.float32)
+def fp8_gemm_bf16(
+    input: torch.Tensor,
+    weight: torch.Tensor,
+    input_scale: torch.Tensor,
+    weight_scale: torch.Tensor,
+    out: torch.Tensor | None = None,
+) -> torch.Tensor:
+    """Compute ``(input * input_scale) @ (weight * weight_scale).T``.
+    ``input`` is FP8 E4M3 with shape ``(M, K)``. ``weight`` is FP8 E4M3 with
+    shape ``(N, K)``. ``input_scale`` and ``weight_scale`` are CUDA float32
+    scalar tensors. Output is BF16 with shape ``(M, N)``.
+    """
+    if out is None:
+        out = torch.empty(
+            (input.shape[0], weight.shape[0]),
+            device=input.device,
+            dtype=torch.bfloat16,
+        )
+    ops.fp8_gemm_bf16(input, weight, input_scale, weight_scale, out)
+    return out
+def fp8_linear_bias_gelu_quant_bf16(
+    input: torch.Tensor,
+    weight: torch.Tensor,
+    bias: torch.Tensor,
+    input_scale: torch.Tensor,
+    weight_scale: torch.Tensor,
+    output_scale: torch.Tensor,
+    hidden_bf16: torch.Tensor | None = None,
+    out_fp8: torch.Tensor | None = None,
+) -> tuple[torch.Tensor, torch.Tensor]:
+    """FP8 linear + BF16 bias/GELU + FP8 quantized output.
+    Returns ``(hidden_bf16, out_fp8)``. ``hidden_bf16`` is the post-GEMM
+    pre-activation scratch; ``out_fp8`` is the quantized activation.
+    """
+    if hidden_bf16 is None:
+        hidden_bf16 = torch.empty(
+            (input.shape[0], weight.shape[0]),
+            device=input.device,
+            dtype=torch.bfloat16,
+        )
+    if out_fp8 is None:
+        out_fp8 = torch.empty_like(hidden_bf16, dtype=_fp8_dtype())
+    ops.fp8_linear_bias_gelu_quant_bf16(
+        input,
+        weight,
+        bias,
+        input_scale,
+        weight_scale,
+        output_scale,
+        hidden_bf16,
+        out_fp8,
+    )
+    return hidden_bf16, out_fp8
+def fp8_gelu_mlp_bf16(
+    input: torch.Tensor,
+    up_weight: torch.Tensor,
+    up_bias: torch.Tensor,
+    down_weight: torch.Tensor,
+    down_bias: torch.Tensor,
+    input_scale: torch.Tensor,
+    up_weight_scale: torch.Tensor,
+    hidden_scale: torch.Tensor,
+    down_weight_scale: torch.Tensor,
+    hidden_bf16: torch.Tensor | None = None,
+    hidden_fp8: torch.Tensor | None = None,
+    out: torch.Tensor | None = None,
+) -> torch.Tensor:
+    """FP8 GELU MLP block with BF16 output.
+    Computes:
+    ``hidden = gelu(fp8_gemm(input, up_weight) + up_bias)``
+    ``hidden_fp8 = quantize_fp8(hidden, hidden_scale)``
+    ``out = fp8_gemm(hidden_fp8, down_weight) + down_bias``
+    """
+    if hidden_bf16 is None:
+        hidden_bf16 = torch.empty(
+            (input.shape[0], up_weight.shape[0]),
+            device=input.device,
+            dtype=torch.bfloat16,
+        )
+    if hidden_fp8 is None:
+        hidden_fp8 = torch.empty_like(hidden_bf16, dtype=_fp8_dtype())
+    if out is None:
+        out = torch.empty(
+            (input.shape[0], down_weight.shape[0]),
+            device=input.device,
+            dtype=torch.bfloat16,
+        )
+    ops.fp8_gelu_mlp_bf16(
+        input,
+        up_weight,
+        up_bias,
+        down_weight,
+        down_bias,
+        input_scale,
+        up_weight_scale,
+        hidden_scale,
+        down_weight_scale,
+        hidden_bf16,
+        hidden_fp8,
+        out,
+    )
+    return out
+__all__ = [
+    "fp8_gemm_bf16",
+    "fp8_gelu_mlp_bf16",
+    "fp8_linear_bias_gelu_quant_bf16",
+]

build/{torch212-cxx11-cu132-x86_64-linux/_flashrt_fp8_ffn_cuda_5de4768.abi3.so → torch211-cxx11-rocm72-x86_64-linux/_flashrt_fp8_ffn_rocm_4a17df5.abi3.so} RENAMED Viewed

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

 version https://git-lfs.github.com/spec/v1
+oid sha256:74b3c142dc669c0e4354b31e05cdf0b080617a74cdf451761aa4c5520d86f651
+size 143848

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

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

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

build/torch211-cxx11-rocm72-x86_64-linux/flashrt_fp8_ffn/__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-fp8-ffn",
+  "id": "_flashrt_fp8_ffn_rocm_4a17df5",
+  "version": 1,
+  "license": "Apache-2.0",
+  "python-depends": [],
+  "backend": {
+    "type": "rocm",
+    "archs": [
+      "gfx942"
+    ]
+  }
+}

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

@@ -34,6 +34,12 @@ _preload_cublaslt()
 from ._ops import add_op_namespace_prefix, ops
 @torch.library.register_fake(add_op_namespace_prefix("fp8_gemm_bf16"))
 def _fp8_gemm_bf16_fake(
     input: torch.Tensor,
@@ -146,7 +152,7 @@ def fp8_linear_bias_gelu_quant_bf16(
             dtype=torch.bfloat16,
         )
     if out_fp8 is None:
-        out_fp8 = torch.empty_like(hidden_bf16, dtype=torch.float8_e4m3fn)
     ops.fp8_linear_bias_gelu_quant_bf16(
         input,
         weight,
@@ -190,7 +196,7 @@ def fp8_gelu_mlp_bf16(
             dtype=torch.bfloat16,
         )
     if hidden_fp8 is None:
-        hidden_fp8 = torch.empty_like(hidden_bf16, dtype=torch.float8_e4m3fn)
     if out is None:
         out = torch.empty(
             (input.shape[0], down_weight.shape[0]),

 from ._ops import add_op_namespace_prefix, ops
+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
 @torch.library.register_fake(add_op_namespace_prefix("fp8_gemm_bf16"))
 def _fp8_gemm_bf16_fake(
     input: torch.Tensor,
             dtype=torch.bfloat16,
         )
     if out_fp8 is None:
+        out_fp8 = torch.empty_like(hidden_bf16, dtype=_fp8_dtype())
     ops.fp8_linear_bias_gelu_quant_bf16(
         input,
         weight,
             dtype=torch.bfloat16,
         )
     if hidden_fp8 is None:
+        hidden_fp8 = torch.empty_like(hidden_bf16, dtype=_fp8_dtype())
     if out is None:
         out = torch.empty(
             (input.shape[0], down_weight.shape[0]),

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

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

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

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

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

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

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

@@ -1,6 +1,6 @@
 {
   "name": "flashrt-fp8-ffn",
-  "id": "_flashrt_fp8_ffn_cuda_5de4768",
   "version": 1,
   "license": "Apache-2.0",
   "python-depends": [],

 {
   "name": "flashrt-fp8-ffn",
+  "id": "_flashrt_fp8_ffn_cuda_4a17df5",
   "version": 1,
   "license": "Apache-2.0",
   "python-depends": [],

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

@@ -34,6 +34,12 @@ _preload_cublaslt()
 from ._ops import add_op_namespace_prefix, ops
 @torch.library.register_fake(add_op_namespace_prefix("fp8_gemm_bf16"))
 def _fp8_gemm_bf16_fake(
     input: torch.Tensor,
@@ -146,7 +152,7 @@ def fp8_linear_bias_gelu_quant_bf16(
             dtype=torch.bfloat16,
         )
     if out_fp8 is None:
-        out_fp8 = torch.empty_like(hidden_bf16, dtype=torch.float8_e4m3fn)
     ops.fp8_linear_bias_gelu_quant_bf16(
         input,
         weight,
@@ -190,7 +196,7 @@ def fp8_gelu_mlp_bf16(
             dtype=torch.bfloat16,
         )
     if hidden_fp8 is None:
-        hidden_fp8 = torch.empty_like(hidden_bf16, dtype=torch.float8_e4m3fn)
     if out is None:
         out = torch.empty(
             (input.shape[0], down_weight.shape[0]),

 from ._ops import add_op_namespace_prefix, ops
+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
 @torch.library.register_fake(add_op_namespace_prefix("fp8_gemm_bf16"))
 def _fp8_gemm_bf16_fake(
     input: torch.Tensor,
             dtype=torch.bfloat16,
         )
     if out_fp8 is None:
+        out_fp8 = torch.empty_like(hidden_bf16, dtype=_fp8_dtype())
     ops.fp8_linear_bias_gelu_quant_bf16(
         input,
         weight,
             dtype=torch.bfloat16,
         )
     if hidden_fp8 is None:
+        hidden_fp8 = torch.empty_like(hidden_bf16, dtype=_fp8_dtype())
     if out is None:
         out = torch.empty(
             (input.shape[0], down_weight.shape[0]),

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

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

+version https://git-lfs.github.com/spec/v1
+oid sha256:14e82facb9b426efcf66b0855d0f281192c806d20d4e4bd00ef63c1a94ea4d48
+size 298032

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

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

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

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

@@ -1,6 +1,6 @@
 {
   "name": "flashrt-fp8-ffn",
-  "id": "_flashrt_fp8_ffn_cuda_5de4768",
   "version": 1,
   "license": "Apache-2.0",
   "python-depends": [],

 {
   "name": "flashrt-fp8-ffn",
+  "id": "_flashrt_fp8_ffn_cuda_4a17df5",
   "version": 1,
   "license": "Apache-2.0",
   "python-depends": [],

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

	@@ -0,0 +1,227 @@

+"""FlashRT FP8 FFN 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
+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
+@torch.library.register_fake(add_op_namespace_prefix("fp8_gemm_bf16"))
+def _fp8_gemm_bf16_fake(
+    input: torch.Tensor,
+    weight: torch.Tensor,
+    input_scale: torch.Tensor,
+    weight_scale: torch.Tensor,
+    out: torch.Tensor,
+) -> None:
+    if input.dim() != 2 or weight.dim() != 2:
+        raise RuntimeError("input and weight must be rank-2 tensors")
+    if out.shape != (input.shape[0], weight.shape[0]):
+        raise RuntimeError("out shape must be (input.shape[0], weight.shape[0])")
+    return None
+@torch.library.register_fake(add_op_namespace_prefix("fp8_linear_bias_gelu_quant_bf16"))
+def _fp8_linear_bias_gelu_quant_bf16_fake(
+    input: torch.Tensor,
+    weight: torch.Tensor,
+    bias: torch.Tensor,
+    input_scale: torch.Tensor,
+    weight_scale: torch.Tensor,
+    output_scale: torch.Tensor,
+    hidden_bf16: torch.Tensor,
+    out_fp8: torch.Tensor,
+) -> None:
+    expected = (input.shape[0], weight.shape[0])
+    if hidden_bf16.shape != expected or out_fp8.shape != expected:
+        raise RuntimeError(
+            "hidden_bf16 and out_fp8 shapes must be "
+            "(input.shape[0], weight.shape[0])"
+        )
+    return None
+@torch.library.register_fake(add_op_namespace_prefix("fp8_gelu_mlp_bf16"))
+def _fp8_gelu_mlp_bf16_fake(
+    input: torch.Tensor,
+    up_weight: torch.Tensor,
+    up_bias: torch.Tensor,
+    down_weight: torch.Tensor,
+    down_bias: torch.Tensor,
+    input_scale: torch.Tensor,
+    up_weight_scale: torch.Tensor,
+    hidden_scale: torch.Tensor,
+    down_weight_scale: torch.Tensor,
+    hidden_bf16: torch.Tensor,
+    hidden_fp8: torch.Tensor,
+    out: torch.Tensor,
+) -> None:
+    hidden_shape = (input.shape[0], up_weight.shape[0])
+    out_shape = (input.shape[0], down_weight.shape[0])
+    if hidden_bf16.shape != hidden_shape or hidden_fp8.shape != hidden_shape:
+        raise RuntimeError(
+            "hidden buffers must be (input.shape[0], up_weight.shape[0])"
+        )
+    if out.shape != out_shape:
+        raise RuntimeError("out shape must be (input.shape[0], down_weight.shape[0])")
+    return None
+def _scalar_scale_like(input: torch.Tensor, value: float = 1.0) -> torch.Tensor:
+    return torch.tensor([value], device=input.device, dtype=torch.float32)
+def fp8_gemm_bf16(
+    input: torch.Tensor,
+    weight: torch.Tensor,
+    input_scale: torch.Tensor,
+    weight_scale: torch.Tensor,
+    out: torch.Tensor | None = None,
+) -> torch.Tensor:
+    """Compute ``(input * input_scale) @ (weight * weight_scale).T``.
+    ``input`` is FP8 E4M3 with shape ``(M, K)``. ``weight`` is FP8 E4M3 with
+    shape ``(N, K)``. ``input_scale`` and ``weight_scale`` are CUDA float32
+    scalar tensors. Output is BF16 with shape ``(M, N)``.
+    """
+    if out is None:
+        out = torch.empty(
+            (input.shape[0], weight.shape[0]),
+            device=input.device,
+            dtype=torch.bfloat16,
+        )
+    ops.fp8_gemm_bf16(input, weight, input_scale, weight_scale, out)
+    return out
+def fp8_linear_bias_gelu_quant_bf16(
+    input: torch.Tensor,
+    weight: torch.Tensor,
+    bias: torch.Tensor,
+    input_scale: torch.Tensor,
+    weight_scale: torch.Tensor,
+    output_scale: torch.Tensor,
+    hidden_bf16: torch.Tensor | None = None,
+    out_fp8: torch.Tensor | None = None,
+) -> tuple[torch.Tensor, torch.Tensor]:
+    """FP8 linear + BF16 bias/GELU + FP8 quantized output.
+    Returns ``(hidden_bf16, out_fp8)``. ``hidden_bf16`` is the post-GEMM
+    pre-activation scratch; ``out_fp8`` is the quantized activation.
+    """
+    if hidden_bf16 is None:
+        hidden_bf16 = torch.empty(
+            (input.shape[0], weight.shape[0]),
+            device=input.device,
+            dtype=torch.bfloat16,
+        )
+    if out_fp8 is None:
+        out_fp8 = torch.empty_like(hidden_bf16, dtype=_fp8_dtype())
+    ops.fp8_linear_bias_gelu_quant_bf16(
+        input,
+        weight,
+        bias,
+        input_scale,
+        weight_scale,
+        output_scale,
+        hidden_bf16,
+        out_fp8,
+    )
+    return hidden_bf16, out_fp8
+def fp8_gelu_mlp_bf16(
+    input: torch.Tensor,
+    up_weight: torch.Tensor,
+    up_bias: torch.Tensor,
+    down_weight: torch.Tensor,
+    down_bias: torch.Tensor,
+    input_scale: torch.Tensor,
+    up_weight_scale: torch.Tensor,
+    hidden_scale: torch.Tensor,
+    down_weight_scale: torch.Tensor,
+    hidden_bf16: torch.Tensor | None = None,
+    hidden_fp8: torch.Tensor | None = None,
+    out: torch.Tensor | None = None,
+) -> torch.Tensor:
+    """FP8 GELU MLP block with BF16 output.
+    Computes:
+    ``hidden = gelu(fp8_gemm(input, up_weight) + up_bias)``
+    ``hidden_fp8 = quantize_fp8(hidden, hidden_scale)``
+    ``out = fp8_gemm(hidden_fp8, down_weight) + down_bias``
+    """
+    if hidden_bf16 is None:
+        hidden_bf16 = torch.empty(
+            (input.shape[0], up_weight.shape[0]),
+            device=input.device,
+            dtype=torch.bfloat16,
+        )
+    if hidden_fp8 is None:
+        hidden_fp8 = torch.empty_like(hidden_bf16, dtype=_fp8_dtype())
+    if out is None:
+        out = torch.empty(
+            (input.shape[0], down_weight.shape[0]),
+            device=input.device,
+            dtype=torch.bfloat16,
+        )
+    ops.fp8_gelu_mlp_bf16(
+        input,
+        up_weight,
+        up_bias,
+        down_weight,
+        down_bias,
+        input_scale,
+        up_weight_scale,
+        hidden_scale,
+        down_weight_scale,
+        hidden_bf16,
+        hidden_fp8,
+        out,
+    )
+    return out
+__all__ = [
+    "fp8_gemm_bf16",
+    "fp8_gelu_mlp_bf16",
+    "fp8_linear_bias_gelu_quant_bf16",
+]

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

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

+version https://git-lfs.github.com/spec/v1
+oid sha256:3dc2093b25e985a0af258cd2702fb3e28ac98e4ff04a406c423ce48dd66cea11
+size 146416

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

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

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

build/torch212-cxx11-rocm71-x86_64-linux/flashrt_fp8_ffn/__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-fp8-ffn",
+  "id": "_flashrt_fp8_ffn_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,227 @@

+"""FlashRT FP8 FFN 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
+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
+@torch.library.register_fake(add_op_namespace_prefix("fp8_gemm_bf16"))
+def _fp8_gemm_bf16_fake(
+    input: torch.Tensor,
+    weight: torch.Tensor,
+    input_scale: torch.Tensor,
+    weight_scale: torch.Tensor,
+    out: torch.Tensor,
+) -> None:
+    if input.dim() != 2 or weight.dim() != 2:
+        raise RuntimeError("input and weight must be rank-2 tensors")
+    if out.shape != (input.shape[0], weight.shape[0]):
+        raise RuntimeError("out shape must be (input.shape[0], weight.shape[0])")
+    return None
+@torch.library.register_fake(add_op_namespace_prefix("fp8_linear_bias_gelu_quant_bf16"))
+def _fp8_linear_bias_gelu_quant_bf16_fake(
+    input: torch.Tensor,
+    weight: torch.Tensor,
+    bias: torch.Tensor,
+    input_scale: torch.Tensor,
+    weight_scale: torch.Tensor,
+    output_scale: torch.Tensor,
+    hidden_bf16: torch.Tensor,
+    out_fp8: torch.Tensor,
+) -> None:
+    expected = (input.shape[0], weight.shape[0])
+    if hidden_bf16.shape != expected or out_fp8.shape != expected:
+        raise RuntimeError(
+            "hidden_bf16 and out_fp8 shapes must be "
+            "(input.shape[0], weight.shape[0])"
+        )
+    return None
+@torch.library.register_fake(add_op_namespace_prefix("fp8_gelu_mlp_bf16"))
+def _fp8_gelu_mlp_bf16_fake(
+    input: torch.Tensor,
+    up_weight: torch.Tensor,
+    up_bias: torch.Tensor,
+    down_weight: torch.Tensor,
+    down_bias: torch.Tensor,
+    input_scale: torch.Tensor,
+    up_weight_scale: torch.Tensor,
+    hidden_scale: torch.Tensor,
+    down_weight_scale: torch.Tensor,
+    hidden_bf16: torch.Tensor,
+    hidden_fp8: torch.Tensor,
+    out: torch.Tensor,
+) -> None:
+    hidden_shape = (input.shape[0], up_weight.shape[0])
+    out_shape = (input.shape[0], down_weight.shape[0])
+    if hidden_bf16.shape != hidden_shape or hidden_fp8.shape != hidden_shape:
+        raise RuntimeError(
+            "hidden buffers must be (input.shape[0], up_weight.shape[0])"
+        )
+    if out.shape != out_shape:
+        raise RuntimeError("out shape must be (input.shape[0], down_weight.shape[0])")
+    return None
+def _scalar_scale_like(input: torch.Tensor, value: float = 1.0) -> torch.Tensor:
+    return torch.tensor([value], device=input.device, dtype=torch.float32)
+def fp8_gemm_bf16(
+    input: torch.Tensor,
+    weight: torch.Tensor,
+    input_scale: torch.Tensor,
+    weight_scale: torch.Tensor,
+    out: torch.Tensor | None = None,
+) -> torch.Tensor:
+    """Compute ``(input * input_scale) @ (weight * weight_scale).T``.
+    ``input`` is FP8 E4M3 with shape ``(M, K)``. ``weight`` is FP8 E4M3 with
+    shape ``(N, K)``. ``input_scale`` and ``weight_scale`` are CUDA float32
+    scalar tensors. Output is BF16 with shape ``(M, N)``.
+    """
+    if out is None:
+        out = torch.empty(
+            (input.shape[0], weight.shape[0]),
+            device=input.device,
+            dtype=torch.bfloat16,
+        )
+    ops.fp8_gemm_bf16(input, weight, input_scale, weight_scale, out)
+    return out
+def fp8_linear_bias_gelu_quant_bf16(
+    input: torch.Tensor,
+    weight: torch.Tensor,
+    bias: torch.Tensor,
+    input_scale: torch.Tensor,
+    weight_scale: torch.Tensor,
+    output_scale: torch.Tensor,
+    hidden_bf16: torch.Tensor | None = None,
+    out_fp8: torch.Tensor | None = None,
+) -> tuple[torch.Tensor, torch.Tensor]:
+    """FP8 linear + BF16 bias/GELU + FP8 quantized output.
+    Returns ``(hidden_bf16, out_fp8)``. ``hidden_bf16`` is the post-GEMM
+    pre-activation scratch; ``out_fp8`` is the quantized activation.
+    """
+    if hidden_bf16 is None:
+        hidden_bf16 = torch.empty(
+            (input.shape[0], weight.shape[0]),
+            device=input.device,
+            dtype=torch.bfloat16,
+        )
+    if out_fp8 is None:
+        out_fp8 = torch.empty_like(hidden_bf16, dtype=_fp8_dtype())
+    ops.fp8_linear_bias_gelu_quant_bf16(
+        input,
+        weight,
+        bias,
+        input_scale,
+        weight_scale,
+        output_scale,
+        hidden_bf16,
+        out_fp8,
+    )
+    return hidden_bf16, out_fp8
+def fp8_gelu_mlp_bf16(
+    input: torch.Tensor,
+    up_weight: torch.Tensor,
+    up_bias: torch.Tensor,
+    down_weight: torch.Tensor,
+    down_bias: torch.Tensor,
+    input_scale: torch.Tensor,
+    up_weight_scale: torch.Tensor,
+    hidden_scale: torch.Tensor,
+    down_weight_scale: torch.Tensor,
+    hidden_bf16: torch.Tensor | None = None,
+    hidden_fp8: torch.Tensor | None = None,
+    out: torch.Tensor | None = None,
+) -> torch.Tensor:
+    """FP8 GELU MLP block with BF16 output.
+    Computes:
+    ``hidden = gelu(fp8_gemm(input, up_weight) + up_bias)``
+    ``hidden_fp8 = quantize_fp8(hidden, hidden_scale)``
+    ``out = fp8_gemm(hidden_fp8, down_weight) + down_bias``
+    """
+    if hidden_bf16 is None:
+        hidden_bf16 = torch.empty(
+            (input.shape[0], up_weight.shape[0]),
+            device=input.device,
+            dtype=torch.bfloat16,
+        )
+    if hidden_fp8 is None:
+        hidden_fp8 = torch.empty_like(hidden_bf16, dtype=_fp8_dtype())
+    if out is None:
+        out = torch.empty(
+            (input.shape[0], down_weight.shape[0]),
+            device=input.device,
+            dtype=torch.bfloat16,
+        )
+    ops.fp8_gelu_mlp_bf16(
+        input,
+        up_weight,
+        up_bias,
+        down_weight,
+        down_bias,
+        input_scale,
+        up_weight_scale,
+        hidden_scale,
+        down_weight_scale,
+        hidden_bf16,
+        hidden_fp8,
+        out,
+    )
+    return out
+__all__ = [
+    "fp8_gemm_bf16",
+    "fp8_gelu_mlp_bf16",
+    "fp8_linear_bias_gelu_quant_bf16",
+]

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

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

+version https://git-lfs.github.com/spec/v1
+oid sha256:387239194294371025c6fa91fd16e3ac550052f1632b73ec37b7879371cc8b9d
+size 145040

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

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

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

build/torch212-cxx11-rocm72-x86_64-linux/flashrt_fp8_ffn/__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-fp8-ffn",
+  "id": "_flashrt_fp8_ffn_rocm_4a17df5",
+  "version": 1,
+  "license": "Apache-2.0",
+  "python-depends": [],
+  "backend": {
+    "type": "rocm",
+    "archs": [
+      "gfx942"
+    ]
+  }
+}