Add torch213-cxx11-cu130-aarch64-linux Thor artifact

Browse files

Files changed (5) hide show

build/torch213-cxx11-cu130-aarch64-linux/__init__.py +455 -0
build/torch213-cxx11-cu130-aarch64-linux/_adaptive_layernorm_producers_cuda_8a66d8b.abi3.so +3 -0
build/torch213-cxx11-cu130-aarch64-linux/_ops.py +6 -0
build/torch213-cxx11-cu130-aarch64-linux/adaptive_layernorm_producers/__init__.py +14 -0
build/torch213-cxx11-cu130-aarch64-linux/metadata.json +28 -0

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

	@@ -0,0 +1,455 @@

+"""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_ptok_bf16"))
+def _ada_layer_norm_quant_fp8_ptok_bf16_fake(
+    x: torch.Tensor,
+    scale: torch.Tensor,
+    shift: torch.Tensor,
+    act_scale: torch.Tensor,
+    eps: float,
+    out: torch.Tensor,
+) -> None:
+    _check_x(x)
+    if scale.shape != x.shape or shift.shape != x.shape:
+        raise RuntimeError(
+            "per-token scale and shift must have shape (seq_len, 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_fp8_ptok_table_bf16"))
+def _ada_layer_norm_quant_fp8_ptok_table_bf16_fake(
+    x: torch.Tensor,
+    temb: torch.Tensor,
+    table: torch.Tensor,
+    act_scale: torch.Tensor,
+    shift_idx: int,
+    scale_idx: int,
+    eps: float,
+    out: torch.Tensor,
+) -> None:
+    _check_x(x)
+    if temb.dim() != 3 or temb.shape[0] != x.shape[0] or temb.shape[2] != x.shape[1]:
+        raise RuntimeError("temb must have shape (seq_len, n_chunks, dim)")
+    if table.dim() != 2 or table.shape[0] != temb.shape[1] or table.shape[1] != x.shape[1]:
+        raise RuntimeError("table must have shape (n_chunks, 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_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
+@torch.library.register_fake(add_op_namespace_prefix("layer_norm_no_affine_quant_nvfp4_swizzled_bf16"))
+def _layer_norm_no_affine_quant_nvfp4_swizzled_bf16_fake(
+    x: torch.Tensor,
+    eps: float,
+    packed: torch.Tensor,
+    sf_swizzled: torch.Tensor,
+) -> None:
+    _check_x(x)
+    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
+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_ptok_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:
+    """Per-token AdaLN producer: no-affine LayerNorm + [M, D] scale/shift
+    modulation + static FP8 quantize, one pass.
+    The per-token form serves hosts whose timestep embedding carries one
+    modulation vector per activation row (video DiTs with [B, S, 6, D]
+    tables); the broadcast form above cannot express that seam.
+    """
+    if out is None:
+        out = _alloc_fp8(x)
+    ops.ada_layer_norm_quant_fp8_ptok_bf16(
+        x, scale, shift, act_scale, float(eps), out)
+    return out
+def ada_layer_norm_quant_fp8_ptok_table_bf16(
+    x: torch.Tensor,
+    temb: torch.Tensor,
+    table: torch.Tensor,
+    act_scale: torch.Tensor,
+    shift_idx: int,
+    scale_idx: int,
+    eps: float = 1e-5,
+    out: torch.Tensor | None = None,
+) -> torch.Tensor:
+    """Per-token AdaLN producer, table form.
+    ``temb`` is the host's per-token modulation embedding
+    ``[seq_len, n_chunks, dim]`` (cast to BF16 once per call) and
+    ``table`` the block's own ``[n_chunks, dim]`` FP32 modulation
+    parameter. The table add, chunk selection, no-affine LayerNorm,
+    modulation and static FP8 quantize all happen in one pass — the
+    per-block six-chunk materialization the host would otherwise pay
+    never exists.
+    """
+    if out is None:
+        out = _alloc_fp8(x)
+    ops.ada_layer_norm_quant_fp8_ptok_table_bf16(
+        x, temb, table, act_scale, int(shift_idx), int(scale_idx),
+        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
+def layer_norm_no_affine_quant_nvfp4_swizzled_bf16(
+    x: torch.Tensor,
+    eps: float = 1e-5,
+    packed: torch.Tensor | None = None,
+    sf_swizzled: torch.Tensor | None = None,
+) -> tuple[torch.Tensor, torch.Tensor]:
+    """SM110 no-affine LayerNorm producing NVFP4 and CUTLASS SFA."""
+    if packed is None or sf_swizzled is None:
+        packed, sf_swizzled = _alloc_nvfp4(x)
+    ops.layer_norm_no_affine_quant_nvfp4_swizzled_bf16(
+        x, float(eps), packed, sf_swizzled
+    )
+    return packed, sf_swizzled
+@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_ptok_bf16",
+    "ada_layer_norm_quant_fp8_ptok_table_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",
+    "layer_norm_no_affine_quant_nvfp4_swizzled_bf16",
+    "adaln_modulation6_bf16",
+    "swizzled_sf_size",
+]

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

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

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

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

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

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

build/torch213-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/torch213-cxx11-cu130-aarch64-linux/metadata.json ADDED Viewed

	@@ -0,0 +1,28 @@

+{
+  "name": "adaptive-layernorm-producers",
+  "id": "_adaptive_layernorm_producers_cuda_8a66d8b",
+  "version": 1,
+  "license": "Apache-2.0",
+  "python-depends": [],
+  "backend": {
+    "type": "cuda",
+    "archs": [
+      "11.0a"
+    ]
+  },
+  "digest": {
+    "algorithm": "sha256",
+    "files": {
+      "__init__.py": "/6Ld6WSyL7PjILcYt4P7x3fnXi2C+ovbSP/7nNFetb8=",
+      "_adaptive_layernorm_producers_cuda_8a66d8b.abi3.so": "t6LbLgB9vIDaJiwgaJEFjH24wcfySbkq1hhcHOEbgU8=",
+      "_ops.py": "8CjL8CxdQouAOnaox+O0AKwLERGSDc8mao+V3ArMuTk=",
+      "adaptive_layernorm_producers/__init__.py": "v6p5XMfQzddhi1fLSAw4HX9CyS0rQsidvu9VsT01xi4="
+    }
+  },
+  "provenance": {
+    "kernel": {
+      "sha": "8a66d8b",
+      "dirty": false
+    }
+  }
+}