Add torch213-cxx11-cu130-aarch64-linux artifact

Browse files

Files changed (5) hide show

build/torch213-cxx11-cu130-aarch64-linux/__init__.py +115 -0
build/torch213-cxx11-cu130-aarch64-linux/_ops.py +6 -0
build/torch213-cxx11-cu130-aarch64-linux/flashrt_adaptive_norms/__init__.py +14 -0
build/torch213-cxx11-cu130-aarch64-linux/flashrt_adaptive_norms_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,115 @@

+"""FlashRT adaptive norm kernels."""
+from __future__ import annotations
+import torch
+from ._ops import add_op_namespace_prefix, ops
+def _check_adaptive_shapes(
+    x: torch.Tensor,
+    weight: torch.Tensor,
+    style: torch.Tensor,
+    out: torch.Tensor,
+    gate_out: torch.Tensor,
+) -> None:
+    if x.dim() != 2:
+        raise RuntimeError("x must have shape (rows, dim)")
+    rows, dim = x.shape
+    if dim % 2 != 0:
+        raise RuntimeError("x.shape[1] must be even")
+    if weight.shape != (dim,):
+        raise RuntimeError("weight must have shape (dim,)")
+    if style.shape != (rows, 3 * dim):
+        raise RuntimeError("style must have shape (rows, 3 * dim)")
+    if out.shape != x.shape or gate_out.shape != x.shape:
+        raise RuntimeError("out and gate_out must have the same shape as x")
+@torch.library.register_fake(add_op_namespace_prefix("ada_rms_norm_style_bf16"))
+def _ada_rms_norm_style_bf16_fake(
+    x: torch.Tensor,
+    weight: torch.Tensor,
+    style: torch.Tensor,
+    eps: float,
+    out: torch.Tensor,
+    gate_out: torch.Tensor,
+) -> None:
+    _check_adaptive_shapes(x, weight, style, out, gate_out)
+    return None
+@torch.library.register_fake(add_op_namespace_prefix("gate_residual_ada_norm_fp8_static_bf16"))
+def _gate_residual_ada_norm_fp8_static_bf16_fake(
+    residual: torch.Tensor,
+    x: torch.Tensor,
+    gate: torch.Tensor,
+    weight: torch.Tensor,
+    style: torch.Tensor,
+    scale: torch.Tensor,
+    eps: float,
+    out: torch.Tensor,
+    gate_out: torch.Tensor,
+) -> None:
+    _check_adaptive_shapes(residual, weight, style, out, gate_out)
+    if x.shape != residual.shape or gate.shape != residual.shape:
+        raise RuntimeError("x and gate must have the same shape as residual")
+    if scale.numel() != 1:
+        raise RuntimeError("scale must be a scalar tensor")
+    return None
+def ada_rms_norm_style_bf16(
+    x: torch.Tensor,
+    weight: torch.Tensor,
+    style: torch.Tensor,
+    eps: float = 1e-6,
+    out: torch.Tensor | None = None,
+    gate_out: torch.Tensor | None = None,
+) -> tuple[torch.Tensor, torch.Tensor]:
+    """Apply RMSNorm, style scale/shift, and return the style gate."""
+    if out is None:
+        out = torch.empty_like(x)
+    if gate_out is None:
+        gate_out = torch.empty_like(x)
+    ops.ada_rms_norm_style_bf16(x, weight, style, float(eps), out, gate_out)
+    return out, gate_out
+def gate_residual_ada_norm_fp8_static_bf16(
+    residual: torch.Tensor,
+    x: torch.Tensor,
+    gate: torch.Tensor,
+    weight: torch.Tensor,
+    style: torch.Tensor,
+    scale: torch.Tensor,
+    eps: float = 1e-6,
+    out: torch.Tensor | None = None,
+    gate_out: torch.Tensor | None = None,
+) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
+    """Update residual in place, apply AdaRMSNorm, and emit static-scale FP8."""
+    if out is None:
+        out = torch.empty_like(residual, dtype=torch.float8_e4m3fn)
+    if gate_out is None:
+        gate_out = torch.empty_like(residual)
+    ops.gate_residual_ada_norm_fp8_static_bf16(
+        residual,
+        x,
+        gate,
+        weight,
+        style,
+        scale,
+        float(eps),
+        out,
+        gate_out,
+    )
+    return residual, out, gate_out
+__all__ = [
+    "ada_rms_norm_style_bf16",
+    "gate_residual_ada_norm_fp8_static_bf16",
+]

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

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

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

build/torch213-cxx11-cu130-aarch64-linux/flashrt_adaptive_norms/__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_adaptive_norms_test.abi3.so ADDED Viewed

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

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

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

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

+{
+  "name": "flashrt-adaptive-norms",
+  "id": "flashrt_adaptive_norms_test",
+  "version": 1,
+  "license": "Apache-2.0",
+  "python-depends": [],
+  "backend": {
+    "type": "cuda",
+    "archs": [
+      "11.0"
+    ]
+  },
+  "digest": {
+    "algorithm": "sha256",
+    "files": {
+      "__init__.py": "kv6yHmnnIUpOpDGE4dHU9J32michTE3ygV62h/+D+RI=",
+      "flashrt_adaptive_norms_test.abi3.so": "3gmpuredKMap3RCcR8jleSqnc8XQaUVb2PyAivtMMMs=",
+      "_ops.py": "JOw3+PHqZUGoT770CFlNJojgiXwMOTvFMrvsxZXy7Xc=",
+      "flashrt_adaptive_norms/__init__.py": "v6p5XMfQzddhi1fLSAw4HX9CyS0rQsidvu9VsT01xi4="
+    }
+  },
+  "provenance": {
+    "kernel": {
+      "sha": "456d297",
+      "dirty": false
+    },
+    "validation": {
+      "torch": "2.13.0+cu130",
+      "cuda": "13.0"
+    }
+  }
+}