Add torch211-cxx11-cu130-aarch64-linux SM110 artifact

Browse files

Files changed (8) hide show

build/torch211-cxx11-cu130-aarch64-linux/__init__.py +328 -0
build/torch211-cxx11-cu130-aarch64-linux/__pycache__/__init__.cpython-312.pyc +0 -0
build/torch211-cxx11-cu130-aarch64-linux/__pycache__/_ops.cpython-312.pyc +0 -0
build/torch211-cxx11-cu130-aarch64-linux/_adaptive_layernorm_producers_cuda_7781728.abi3.so +3 -0
build/torch211-cxx11-cu130-aarch64-linux/_ops.py +6 -0
build/torch211-cxx11-cu130-aarch64-linux/adaptive_layernorm_producers/__init__.py +14 -0
build/torch211-cxx11-cu130-aarch64-linux/adaptive_layernorm_producers/__pycache__/__init__.cpython-312.pyc +0 -0
build/torch211-cxx11-cu130-aarch64-linux/metadata.json +22 -0

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

	@@ -0,0 +1,328 @@

+"""FlashRT adaptive LayerNorm producer kernels for DiT/Wan-style blocks."""
+from __future__ import annotations
+import torch
+from ._ops import add_op_namespace_prefix, ops
+def _check_x(x: torch.Tensor) -> None:
+    if x.dim() != 2:
+        raise RuntimeError("x must have shape (rows, dim)")
+    if x.shape[1] % 2 != 0:
+        raise RuntimeError("x.shape[1] must be even")
+def _check_mod(x: torch.Tensor, scale: torch.Tensor, shift: torch.Tensor) -> None:
+    _check_x(x)
+    dim = x.shape[1]
+    if scale.shape != (dim,) or shift.shape != (dim,):
+        raise RuntimeError("scale and shift must have shape (dim,)")
+def _check_scalar(scale: torch.Tensor, name: str) -> None:
+    if scale.numel() != 1:
+        raise RuntimeError(f"{name} must be a scalar tensor")
+def swizzled_sf_size(rows: int, dim: int) -> int:
+    """Return bytes required by the FlashRT/CUTLASS NVFP4 swizzled scale layout."""
+    if dim % 16 != 0:
+        raise RuntimeError("dim must be divisible by 16 for NVFP4 swizzled output")
+    n_blocks = dim // 16
+    n_row_super = (rows + 127) // 128
+    n_col_super = (n_blocks + 3) // 4
+    return n_row_super * n_col_super * 128 * 64
+def _alloc_fp8(x: torch.Tensor) -> torch.Tensor:
+    return torch.empty_like(x, dtype=torch.float8_e4m3fn)
+def _alloc_nvfp4(x: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor]:
+    rows, dim = x.shape
+    packed = torch.empty((rows, dim // 2), device=x.device, dtype=torch.uint8)
+    sf = torch.zeros((swizzled_sf_size(rows, dim),), device=x.device, dtype=torch.uint8)
+    return packed, sf
+@torch.library.register_fake(add_op_namespace_prefix("ada_layer_norm_quant_fp8_bf16"))
+def _ada_layer_norm_quant_fp8_bf16_fake(
+    x: torch.Tensor,
+    scale: torch.Tensor,
+    shift: torch.Tensor,
+    act_scale: torch.Tensor,
+    eps: float,
+    out: torch.Tensor,
+) -> None:
+    _check_mod(x, scale, shift)
+    _check_scalar(act_scale, "act_scale")
+    if out.shape != x.shape:
+        raise RuntimeError("out must have the same shape as x")
+    return None
+@torch.library.register_fake(add_op_namespace_prefix("ada_layer_norm_quant_fp8_modfp8_bf16"))
+def _ada_layer_norm_quant_fp8_modfp8_bf16_fake(
+    x: torch.Tensor,
+    scale_fp8: torch.Tensor,
+    shift_fp8: torch.Tensor,
+    scale_deq: torch.Tensor,
+    shift_deq: torch.Tensor,
+    act_scale: torch.Tensor,
+    eps: float,
+    out: torch.Tensor,
+) -> None:
+    _check_mod(x, scale_fp8, shift_fp8)
+    _check_scalar(scale_deq, "scale_deq")
+    _check_scalar(shift_deq, "shift_deq")
+    _check_scalar(act_scale, "act_scale")
+    if out.shape != x.shape:
+        raise RuntimeError("out must have the same shape as x")
+    return None
+@torch.library.register_fake(add_op_namespace_prefix("awq_ada_layer_norm_quant_fp8_bf16"))
+def _awq_ada_layer_norm_quant_fp8_bf16_fake(
+    x: torch.Tensor,
+    scale: torch.Tensor,
+    shift: torch.Tensor,
+    inv_s: torch.Tensor,
+    act_scale: torch.Tensor,
+    eps: float,
+    out: torch.Tensor,
+) -> None:
+    _check_mod(x, scale, shift)
+    if inv_s.shape != (x.shape[1],):
+        raise RuntimeError("inv_s must have shape (dim,)")
+    _check_scalar(act_scale, "act_scale")
+    if out.shape != x.shape:
+        raise RuntimeError("out must have the same shape as x")
+    return None
+@torch.library.register_fake(add_op_namespace_prefix("ada_layer_norm_quant_nvfp4_swizzled_bf16"))
+def _ada_layer_norm_quant_nvfp4_swizzled_bf16_fake(
+    x: torch.Tensor,
+    scale: torch.Tensor,
+    shift: torch.Tensor,
+    eps: float,
+    packed: torch.Tensor,
+    sf_swizzled: torch.Tensor,
+) -> None:
+    _check_mod(x, scale, shift)
+    rows, dim = x.shape
+    if packed.shape != (rows, dim // 2):
+        raise RuntimeError("packed must have shape (rows, dim // 2)")
+    if sf_swizzled.numel() < swizzled_sf_size(rows, dim):
+        raise RuntimeError("sf_swizzled is too small")
+    return None
+@torch.library.register_fake(add_op_namespace_prefix("ada_layer_norm_quant_nvfp4_swizzled_modfp8_bf16"))
+def _ada_layer_norm_quant_nvfp4_swizzled_modfp8_bf16_fake(
+    x: torch.Tensor,
+    scale_fp8: torch.Tensor,
+    shift_fp8: torch.Tensor,
+    scale_deq: torch.Tensor,
+    shift_deq: torch.Tensor,
+    eps: float,
+    packed: torch.Tensor,
+    sf_swizzled: torch.Tensor,
+) -> None:
+    _check_mod(x, scale_fp8, shift_fp8)
+    _check_scalar(scale_deq, "scale_deq")
+    _check_scalar(shift_deq, "shift_deq")
+    rows, dim = x.shape
+    if packed.shape != (rows, dim // 2):
+        raise RuntimeError("packed must have shape (rows, dim // 2)")
+    if sf_swizzled.numel() < swizzled_sf_size(rows, dim):
+        raise RuntimeError("sf_swizzled is too small")
+    return None
+@torch.library.register_fake(add_op_namespace_prefix("layer_norm_no_affine_quant_fp8_static_bf16"))
+def _layer_norm_no_affine_quant_fp8_static_bf16_fake(
+    x: torch.Tensor,
+    act_scale: torch.Tensor,
+    eps: float,
+    out: torch.Tensor,
+) -> None:
+    _check_x(x)
+    _check_scalar(act_scale, "act_scale")
+    if out.shape != x.shape:
+        raise RuntimeError("out must have the same shape as x")
+    return None
+def ada_layer_norm_quant_fp8_bf16(
+    x: torch.Tensor,
+    scale: torch.Tensor,
+    shift: torch.Tensor,
+    act_scale: torch.Tensor,
+    eps: float = 1e-5,
+    out: torch.Tensor | None = None,
+) -> torch.Tensor:
+    """Fused no-affine LayerNorm + AdaLN scale/shift + static FP8 quantize."""
+    if out is None:
+        out = _alloc_fp8(x)
+    ops.ada_layer_norm_quant_fp8_bf16(x, scale, shift, act_scale, float(eps), out)
+    return out
+def ada_layer_norm_quant_fp8_modfp8_bf16(
+    x: torch.Tensor,
+    scale_fp8: torch.Tensor,
+    shift_fp8: torch.Tensor,
+    scale_deq: torch.Tensor,
+    shift_deq: torch.Tensor,
+    act_scale: torch.Tensor,
+    eps: float = 1e-5,
+    out: torch.Tensor | None = None,
+) -> torch.Tensor:
+    """Fused AdaLN producer when modulation vectors are stored in FP8."""
+    if out is None:
+        out = _alloc_fp8(x)
+    ops.ada_layer_norm_quant_fp8_modfp8_bf16(
+        x, scale_fp8, shift_fp8, scale_deq, shift_deq, act_scale, float(eps), out
+    )
+    return out
+def awq_ada_layer_norm_quant_fp8_bf16(
+    x: torch.Tensor,
+    scale: torch.Tensor,
+    shift: torch.Tensor,
+    inv_s: torch.Tensor,
+    act_scale: torch.Tensor,
+    eps: float = 1e-5,
+    out: torch.Tensor | None = None,
+) -> torch.Tensor:
+    """Fused AdaLN producer with AWQ/SmoothQuant per-channel activation scale."""
+    if out is None:
+        out = _alloc_fp8(x)
+    ops.awq_ada_layer_norm_quant_fp8_bf16(x, scale, shift, inv_s, act_scale, float(eps), out)
+    return out
+def ada_layer_norm_quant_nvfp4_swizzled_bf16(
+    x: torch.Tensor,
+    scale: torch.Tensor,
+    shift: torch.Tensor,
+    eps: float = 1e-5,
+    packed: torch.Tensor | None = None,
+    sf_swizzled: torch.Tensor | None = None,
+) -> tuple[torch.Tensor, torch.Tensor]:
+    """Fused AdaLN producer to NVFP4 packed activations and swizzled scale factors."""
+    if packed is None or sf_swizzled is None:
+        packed, sf_swizzled = _alloc_nvfp4(x)
+    ops.ada_layer_norm_quant_nvfp4_swizzled_bf16(x, scale, shift, float(eps), packed, sf_swizzled)
+    return packed, sf_swizzled
+def ada_layer_norm_quant_nvfp4_swizzled_modfp8_bf16(
+    x: torch.Tensor,
+    scale_fp8: torch.Tensor,
+    shift_fp8: torch.Tensor,
+    scale_deq: torch.Tensor,
+    shift_deq: torch.Tensor,
+    eps: float = 1e-5,
+    packed: torch.Tensor | None = None,
+    sf_swizzled: torch.Tensor | None = None,
+) -> tuple[torch.Tensor, torch.Tensor]:
+    """Fused FP8-modulated AdaLN producer to NVFP4 packed activations."""
+    if packed is None or sf_swizzled is None:
+        packed, sf_swizzled = _alloc_nvfp4(x)
+    ops.ada_layer_norm_quant_nvfp4_swizzled_modfp8_bf16(
+        x, scale_fp8, shift_fp8, scale_deq, shift_deq, float(eps), packed, sf_swizzled
+    )
+    return packed, sf_swizzled
+def layer_norm_no_affine_quant_fp8_static_bf16(
+    x: torch.Tensor,
+    act_scale: torch.Tensor,
+    eps: float = 1e-5,
+    out: torch.Tensor | None = None,
+) -> torch.Tensor:
+    """Fused no-affine LayerNorm + static FP8 quantize."""
+    if out is None:
+        out = _alloc_fp8(x)
+    ops.layer_norm_no_affine_quant_fp8_static_bf16(x, act_scale, float(eps), out)
+    return out
+@torch.library.register_fake(add_op_namespace_prefix("adaln_modulation6_bf16"))
+def _adaln_modulation6_bf16_fake(
+    adaln_params: torch.Tensor,
+    layer_modulation: torch.Tensor,
+    out0: torch.Tensor,
+    out1: torch.Tensor,
+    out2: torch.Tensor,
+    out3: torch.Tensor,
+    out4: torch.Tensor,
+    out5: torch.Tensor,
+) -> None:
+    if adaln_params.dim() != 4 or adaln_params.shape[2] != 6:
+        raise RuntimeError(
+            "adaln_params must have shape (batch, sequence, 6, dim)"
+        )
+    expected = (
+        adaln_params.shape[0],
+        adaln_params.shape[1],
+        adaln_params.shape[3],
+    )
+    if layer_modulation.shape != (6, expected[2]):
+        raise RuntimeError("layer_modulation must have shape (6, dim)")
+    for output in (out0, out1, out2, out3, out4, out5):
+        if output.shape != expected:
+            raise RuntimeError(
+                "each output must have shape (batch, sequence, dim)"
+            )
+    return None
+def adaln_modulation6_bf16(
+    adaln_params: torch.Tensor,
+    layer_modulation: torch.Tensor,
+    *,
+    out: tuple[torch.Tensor, ...] | None = None,
+) -> tuple[torch.Tensor, ...]:
+    """Add six FP32 AdaLN modulation vectors and return six BF16 tensors."""
+    batch, sequence, _, dim = adaln_params.shape
+    if out is None:
+        out = tuple(
+            torch.empty(
+                (batch, sequence, dim),
+                device=adaln_params.device,
+                dtype=torch.bfloat16,
+            )
+            for _ in range(6)
+        )
+    if len(out) != 6:
+        raise RuntimeError("out must contain six tensors")
+    ops.adaln_modulation6_bf16(
+        adaln_params, layer_modulation, *out
+    )
+    return out
+__all__ = [
+    "ada_layer_norm_quant_fp8_bf16",
+    "ada_layer_norm_quant_fp8_modfp8_bf16",
+    "awq_ada_layer_norm_quant_fp8_bf16",
+    "ada_layer_norm_quant_nvfp4_swizzled_bf16",
+    "ada_layer_norm_quant_nvfp4_swizzled_modfp8_bf16",
+    "layer_norm_no_affine_quant_fp8_static_bf16",
+    "adaln_modulation6_bf16",
+    "swizzled_sf_size",
+]

build/torch211-cxx11-cu130-aarch64-linux/__pycache__/__init__.cpython-312.pyc ADDED Viewed

Binary file (13.8 kB). View file

build/torch211-cxx11-cu130-aarch64-linux/__pycache__/_ops.cpython-312.pyc ADDED Viewed

Binary file (599 Bytes). View file

build/torch211-cxx11-cu130-aarch64-linux/_adaptive_layernorm_producers_cuda_7781728.abi3.so ADDED Viewed

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

+version https://git-lfs.github.com/spec/v1
+oid sha256:34e6e6e5a82ce354180b41746d8e703825a0bca39a247e0eb326710951f6b497
+size 389352

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

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

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

build/torch211-cxx11-cu130-aarch64-linux/adaptive_layernorm_producers/__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/torch211-cxx11-cu130-aarch64-linux/adaptive_layernorm_producers/__pycache__/__init__.cpython-312.pyc ADDED Viewed

Binary file (1.39 kB). View file

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

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

+{
+  "name": "adaptive-layernorm-producers",
+  "id": "_adaptive_layernorm_producers_cuda_7781728",
+  "version": 1,
+  "license": "Apache-2.0",
+  "python-depends": [],
+  "backend": {
+    "type": "cuda",
+    "archs": [
+      "11.0"
+    ]
+  },
+  "digest": {
+    "algorithm": "sha256",
+    "files": {
+      "__init__.py": "4NZMN7oQ+stR735F4NXwBZ1dbTlAj1Q67+X/X6B/8zI=",
+      "_adaptive_layernorm_producers_cuda_7781728.abi3.so": "NObm5ags41QYC0F0bY5wOCWgvKOaJH4OsyZxCVH2tJc=",
+      "_ops.py": "kG1EL+EuLDtXPuUtyj/ygYLPUZg6eszWjMCBtbLC0pw=",
+      "adaptive_layernorm_producers/__init__.py": "v6p5XMfQzddhi1fLSAw4HX9CyS0rQsidvu9VsT01xi4="
+    }
+  }
+}