Add torch213-cxx11-cu130-aarch64-linux artifact

Browse files

Files changed (5) hide show

build/torch213-cxx11-cu130-aarch64-linux/__init__.py +443 -0
build/torch213-cxx11-cu130-aarch64-linux/_ops.py +6 -0
build/torch213-cxx11-cu130-aarch64-linux/flashrt_fp8_ffn/__init__.py +14 -0
build/torch213-cxx11-cu130-aarch64-linux/flashrt_fp8_ffn_test.abi3.so +3 -0
build/torch213-cxx11-cu130-aarch64-linux/metadata.json +32 -0

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

	@@ -0,0 +1,443 @@

+"""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_bf16"))
+def _fp8_linear_bias_bf16_fake(
+    input: torch.Tensor,
+    weight: torch.Tensor,
+    bias: 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 bias.shape != (weight.shape[0],):
+        raise RuntimeError("bias shape must be (weight.shape[0],)")
+    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("bf16_fp8_linear_bias_bf16"))
+def _bf16_fp8_linear_bias_bf16_fake(
+    input: torch.Tensor,
+    weight: torch.Tensor,
+    bias: torch.Tensor,
+    input_scale: torch.Tensor,
+    weight_scale: torch.Tensor,
+    input_fp8: 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 input.shape[1] != weight.shape[1]:
+        raise RuntimeError("input.shape[1] must match weight.shape[1]")
+    padded_m = input_fp8.shape[0]
+    if padded_m < input.shape[0]:
+        raise RuntimeError("input_fp8 padded rows must cover input rows")
+    if input_fp8.shape != (padded_m, input.shape[1]):
+        raise RuntimeError("input_fp8 must have shape (padded_m, input.shape[1])")
+    if bias.shape != (weight.shape[0],):
+        raise RuntimeError("bias shape must be (weight.shape[0],)")
+    if out.shape != (padded_m, weight.shape[0]):
+        raise RuntimeError("out must have shape (padded_m, 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
+@torch.library.register_fake(add_op_namespace_prefix("bf16_fp8_gelu_mlp_bf16"))
+def _bf16_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,
+    input_fp8: torch.Tensor,
+    hidden_bf16: torch.Tensor,
+    hidden_fp8: torch.Tensor,
+    out: torch.Tensor,
+) -> None:
+    padded_m = input_fp8.shape[0]
+    if input_fp8.shape != (padded_m, input.shape[1]):
+        raise RuntimeError("input_fp8 must have shape (padded_m, input.shape[1])")
+    if hidden_bf16.shape != (padded_m, up_weight.shape[0]):
+        raise RuntimeError("hidden_bf16 has an invalid padded shape")
+    if hidden_fp8.shape != hidden_bf16.shape:
+        raise RuntimeError("hidden_fp8 must match hidden_bf16")
+    if out.shape != (padded_m, down_weight.shape[0]):
+        raise RuntimeError("out has an invalid padded shape")
+    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_bf16(
+    input: torch.Tensor,
+    weight: torch.Tensor,
+    bias: torch.Tensor,
+    input_scale: torch.Tensor,
+    weight_scale: torch.Tensor,
+    out: torch.Tensor | None = None,
+) -> torch.Tensor:
+    """Run an FP8-weight linear projection with BF16 bias and output.
+    ``input`` and ``weight`` use FP8 E4M3 with shapes ``(M, K)`` and
+    ``(N, K)``. Scales are device float32 scalar tensors. Pass ``out`` to keep
+    the call allocation-free for CUDA Graph capture.
+    """
+    if out is None:
+        out = torch.empty(
+            (input.shape[0], weight.shape[0]),
+            device=input.device,
+            dtype=torch.bfloat16,
+        )
+    ops.fp8_linear_bias_bf16(
+        input, weight, bias, 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
+def _midm_padded_rows(input: torch.Tensor) -> int:
+    rows = input.shape[0]
+    if (
+        torch.version.hip is None
+        and input.is_cuda
+        and torch.cuda.get_device_capability(input.device) == (11, 0)
+        and 9 <= rows <= 128
+    ):
+        return ((rows + 63) // 64) * 64
+    return rows
+def bf16_fp8_linear_bias_bf16(
+    input: torch.Tensor,
+    weight: torch.Tensor,
+    bias: torch.Tensor,
+    input_scale: torch.Tensor,
+    weight_scale: torch.Tensor,
+    input_fp8: torch.Tensor | None = None,
+    out: torch.Tensor | None = None,
+    *,
+    pad_to: int | None = None,
+) -> torch.Tensor:
+    """Run a static-scale FP8 linear projection from a BF16 boundary.
+    This traceable region entry quantizes ``input`` and computes the FP8
+    projection with BF16 bias/output. Thor SM110 pads M=9..128 to a multiple of
+    64 by default. Preallocate ``input_fp8`` and ``out`` for allocation-free
+    CUDA Graph replay.
+    """
+    logical_m = input.shape[0]
+    padded_m = _midm_padded_rows(input) if pad_to is None else pad_to
+    if padded_m < logical_m:
+        raise ValueError("pad_to must be >= input.shape[0]")
+    if input_fp8 is None:
+        input_fp8 = torch.empty(
+            (padded_m, input.shape[1]), device=input.device, dtype=_fp8_dtype()
+        )
+    if out is None:
+        out = torch.empty(
+            (padded_m, weight.shape[0]),
+            device=input.device,
+            dtype=torch.bfloat16,
+        )
+    ops.bf16_fp8_linear_bias_bf16(
+        input, weight, bias, input_scale, weight_scale, input_fp8, out
+    )
+    return out[:logical_m]
+def bf16_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,
+    input_fp8: torch.Tensor | None = None,
+    hidden_bf16: torch.Tensor | None = None,
+    hidden_fp8: torch.Tensor | None = None,
+    out: torch.Tensor | None = None,
+    *,
+    pad_to: int | None = None,
+) -> torch.Tensor:
+    """Run a static-scale FP8 GELU MLP from a BF16 region boundary.
+    The op quantizes ``input`` and executes both FP8-weight GEMMs in one
+    traceable package call. Thor SM110 pads rows 9..128 to a multiple of 64;
+    other targets keep the logical M unless ``pad_to`` is supplied. Pass
+    preallocated buffers for allocation-free CUDA Graph replay. This is one
+    custom-op boundary, not a single CUDA launch.
+    """
+    logical_m = input.shape[0]
+    padded_m = _midm_padded_rows(input) if pad_to is None else pad_to
+    if padded_m < logical_m:
+        raise ValueError("pad_to must be >= input.shape[0]")
+    device = input.device
+    if input_fp8 is None:
+        input_fp8 = torch.empty(
+            (padded_m, input.shape[1]), device=device, dtype=_fp8_dtype()
+        )
+    if hidden_bf16 is None:
+        hidden_bf16 = torch.empty(
+            (padded_m, up_weight.shape[0]), device=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(
+            (padded_m, down_weight.shape[0]), device=device, dtype=torch.bfloat16
+        )
+    ops.bf16_fp8_gelu_mlp_bf16(
+        input,
+        up_weight,
+        up_bias,
+        down_weight,
+        down_bias,
+        input_scale,
+        up_weight_scale,
+        hidden_scale,
+        down_weight_scale,
+        input_fp8,
+        hidden_bf16,
+        hidden_fp8,
+        out,
+    )
+    return out[:logical_m]
+__all__ = [
+    "bf16_fp8_linear_bias_bf16",
+    "bf16_fp8_gelu_mlp_bf16",
+    "fp8_gemm_bf16",
+    "fp8_gelu_mlp_bf16",
+    "fp8_linear_bias_bf16",
+    "fp8_linear_bias_gelu_quant_bf16",
+]

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

	@@ -0,0 +1,6 @@

+import torch
+from . import flashrt_fp8_ffn_test
+ops = torch.ops.flashrt_fp8_ffn_test
+def add_op_namespace_prefix(op_name: str):
+    return f"flashrt_fp8_ffn_test::{op_name}"

build/torch213-cxx11-cu130-aarch64-linux/flashrt_fp8_ffn/__init__.py ADDED Viewed

	@@ -0,0 +1,14 @@

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

build/torch213-cxx11-cu130-aarch64-linux/flashrt_fp8_ffn_test.abi3.so ADDED Viewed

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

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

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

	@@ -0,0 +1,32 @@

+{
+  "name": "flashrt-fp8-ffn",
+  "id": "flashrt_fp8_ffn_test",
+  "version": 1,
+  "license": "Apache-2.0",
+  "python-depends": [],
+  "backend": {
+    "type": "cuda",
+    "archs": [
+      "11.0"
+    ]
+  },
+  "digest": {
+    "algorithm": "sha256",
+    "files": {
+      "__init__.py": "Lfpi4OMG3jS0AMFqtyCpL5H/nNuAzqtpwFaUnv2wgqo=",
+      "flashrt_fp8_ffn_test.abi3.so": "jWgOFoEmUXoRPn1kEuHZw1ldV+E4iAn/VdWjZui/KFc=",
+      "_ops.py": "F8Tu59YGRLKuw0Q2GWUMHzFHDDk3uAor7D3+cCubz5E=",
+      "flashrt_fp8_ffn/__init__.py": "v6p5XMfQzddhi1fLSAw4HX9CyS0rQsidvu9VsT01xi4="
+    }
+  },
+  "provenance": {
+    "kernel": {
+      "sha": "456d297",
+      "dirty": false
+    },
+    "validation": {
+      "torch": "2.13.0+cu130",
+      "cuda": "13.0"
+    }
+  }
+}