liangsu9988 commited on Jun 22

Commit

c19a95e

verified ·

1 Parent(s): 5864a8f

Uploaded using `kernel-builder`.

Browse files

Files changed (21) hide show

benchmarks/benchmark.py +120 -0
build/torch211-cxx11-cu128-x86_64-linux/__init__.py +271 -0
build/torch211-cxx11-cu128-x86_64-linux/_adaptive_layernorm_producers_cuda_1798e7f.abi3.so +3 -0
build/torch211-cxx11-cu128-x86_64-linux/_ops.py +9 -0
build/torch211-cxx11-cu128-x86_64-linux/adaptive_layernorm_producers/__init__.py +26 -0
build/torch211-cxx11-cu128-x86_64-linux/metadata.json +23 -0
build/torch211-cxx11-cu130-x86_64-linux/__init__.py +271 -0
build/torch211-cxx11-cu130-x86_64-linux/_adaptive_layernorm_producers_cuda_1798e7f.abi3.so +3 -0
build/torch211-cxx11-cu130-x86_64-linux/_ops.py +9 -0
build/torch211-cxx11-cu130-x86_64-linux/adaptive_layernorm_producers/__init__.py +26 -0
build/torch211-cxx11-cu130-x86_64-linux/metadata.json +22 -0
build/torch212-cxx11-cu130-x86_64-linux/__init__.py +271 -0
build/torch212-cxx11-cu130-x86_64-linux/_adaptive_layernorm_producers_cuda_1798e7f.abi3.so +3 -0
build/torch212-cxx11-cu130-x86_64-linux/_ops.py +9 -0
build/torch212-cxx11-cu130-x86_64-linux/adaptive_layernorm_producers/__init__.py +26 -0
build/torch212-cxx11-cu130-x86_64-linux/metadata.json +22 -0
build/torch212-cxx11-cu132-x86_64-linux/__init__.py +271 -0
build/torch212-cxx11-cu132-x86_64-linux/_adaptive_layernorm_producers_cuda_1798e7f.abi3.so +3 -0
build/torch212-cxx11-cu132-x86_64-linux/_ops.py +9 -0
build/torch212-cxx11-cu132-x86_64-linux/adaptive_layernorm_producers/__init__.py +26 -0
build/torch212-cxx11-cu132-x86_64-linux/metadata.json +22 -0

benchmarks/benchmark.py ADDED Viewed

	@@ -0,0 +1,120 @@

+#!/usr/bin/env python3
+"""Benchmark adaptive-layernorm-producers against eager producer chains."""
+from __future__ import annotations
+import argparse
+import importlib
+import sys
+from pathlib import Path
+import torch
+ROOT = Path(__file__).resolve().parents[2]
+sys.path.insert(0, str(ROOT / "adaptive-layernorm-producers" / "tests"))
+from test_adaptive_layernorm_producers import (  # noqa: E402
+    load_source_ops,
+    make_case,
+    quant_fp8,
+    ref_adaln,
+    ref_layer_norm_no_affine,
+)
+def load_installed_ops(artifact: str | None):
+    if artifact:
+        sys.path.insert(0, artifact)
+    try:
+        return importlib.import_module("adaptive_layernorm_producers")
+    finally:
+        if artifact:
+            sys.path.remove(artifact)
+def time_cuda(fn, iters: int = 200, warmup: int = 50) -> float:
+    for _ in range(warmup):
+        fn()
+    torch.cuda.synchronize()
+    start = torch.cuda.Event(enable_timing=True)
+    end = torch.cuda.Event(enable_timing=True)
+    start.record()
+    for _ in range(iters):
+        fn()
+    end.record()
+    torch.cuda.synchronize()
+    return start.elapsed_time(end) * 1000.0 / iters
+def run_case(ops, name: str, rows: int, dim: int, eps: float, iters: int) -> dict[str, float | str | int]:
+    x, scale, shift, _inv_s, act_scale, _scale_fp8, _shift_fp8, _scale_deq, _shift_deq = make_case(rows, dim)
+    out = torch.empty_like(x, dtype=torch.float8_e4m3fn)
+    def fused():
+        ops.ada_layer_norm_quant_fp8_bf16(x, scale, shift, act_scale, eps, out=out)
+    def eager():
+        quant_fp8(ref_adaln(x, scale, shift, eps), act_scale)
+    no_affine_out = torch.empty_like(x, dtype=torch.float8_e4m3fn)
+    def fused_no_affine():
+        ops.layer_norm_no_affine_quant_fp8_static_bf16(x, act_scale, eps, out=no_affine_out)
+    def eager_no_affine():
+        quant_fp8(ref_layer_norm_no_affine(x, eps), act_scale)
+    fused_us = time_cuda(fused, iters=iters)
+    eager_us = time_cuda(eager, iters=iters)
+    fused_no_affine_us = time_cuda(fused_no_affine, iters=iters)
+    eager_no_affine_us = time_cuda(eager_no_affine, iters=iters)
+    return {
+        "shape": name,
+        "rows": rows,
+        "dim": dim,
+        "ada_fp8_us": fused_us,
+        "ada_eager_us": eager_us,
+        "ada_speedup": eager_us / fused_us,
+        "no_affine_fp8_us": fused_no_affine_us,
+        "no_affine_eager_us": eager_no_affine_us,
+        "no_affine_speedup": eager_no_affine_us / fused_no_affine_us,
+    }
+def main() -> None:
+    parser = argparse.ArgumentParser()
+    parser.add_argument("--backend", choices=["source", "installed"], default="source")
+    parser.add_argument("--artifact", default=None)
+    parser.add_argument("--iters", type=int, default=200)
+    parser.add_argument("--markdown", default=None)
+    args = parser.parse_args()
+    if not torch.cuda.is_available():
+        raise SystemExit("CUDA is required")
+    torch.manual_seed(2026)
+    ops = load_source_ops() if args.backend == "source" else load_installed_ops(args.artifact)
+    shapes = [
+        ("decode_action", 16, 2048),
+        ("wan_video_short", 64, 3072),
+        ("wan_video_ctx", 256, 3072),
+        ("wan_video_2k", 2520, 3072),
+        ("wan_video_4k", 4096, 3072),
+    ]
+    rows = [run_case(ops, name, r, d, 1e-5, args.iters) for name, r, d in shapes]
+    lines = [
+        "| Shape | Rows | Dim | AdaLN->FP8 us | Eager chain us | Speedup | LN->FP8 us | Eager LN chain us | Speedup |",
+        "|---|---:|---:|---:|---:|---:|---:|---:|---:|",
+    ]
+    for row in rows:
+        line = (
+            f"| {row['shape']} | {row['rows']} | {row['dim']} | "
+            f"{row['ada_fp8_us']:.3f} | {row['ada_eager_us']:.3f} | {row['ada_speedup']:.2f}x | "
+            f"{row['no_affine_fp8_us']:.3f} | {row['no_affine_eager_us']:.3f} | {row['no_affine_speedup']:.2f}x |"
+        )
+        lines.append(line)
+        print(line)
+    if args.markdown:
+        Path(args.markdown).write_text("\n".join(lines) + "\n", encoding="utf-8")
+if __name__ == "__main__":
+    main()

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

	@@ -0,0 +1,271 @@

+"""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
+__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",
+    "swizzled_sf_size",
+]

build/torch211-cxx11-cu128-x86_64-linux/_adaptive_layernorm_producers_cuda_1798e7f.abi3.so ADDED Viewed

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

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

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

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

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

build/torch211-cxx11-cu128-x86_64-linux/adaptive_layernorm_producers/__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-cu128-x86_64-linux/metadata.json ADDED Viewed

	@@ -0,0 +1,23 @@

+{
+  "name": "adaptive-layernorm-producers",
+  "id": "_adaptive_layernorm_producers_cuda_1798e7f",
+  "version": 1,
+  "license": "Apache-2.0",
+  "python-depends": [],
+  "backend": {
+    "type": "cuda",
+    "archs": [
+      "10.0",
+      "10.1",
+      "12.0+PTX",
+      "7.0",
+      "7.2",
+      "7.5",
+      "8.0",
+      "8.6",
+      "8.7",
+      "8.9",
+      "9.0"
+    ]
+  }
+}

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

	@@ -0,0 +1,271 @@

+"""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
+__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",
+    "swizzled_sf_size",
+]

build/torch211-cxx11-cu130-x86_64-linux/_adaptive_layernorm_producers_cuda_1798e7f.abi3.so ADDED Viewed

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

+version https://git-lfs.github.com/spec/v1
+oid sha256:66d3a7d56691415cd91a7b2b8694e10eb13bce362bb450e1599c047978be1971
+size 1313616

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

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

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

build/torch211-cxx11-cu130-x86_64-linux/adaptive_layernorm_producers/__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-cu130-x86_64-linux/metadata.json ADDED Viewed

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

+{
+  "name": "adaptive-layernorm-producers",
+  "id": "_adaptive_layernorm_producers_cuda_1798e7f",
+  "version": 1,
+  "license": "Apache-2.0",
+  "python-depends": [],
+  "backend": {
+    "type": "cuda",
+    "archs": [
+      "10.0",
+      "11.0",
+      "12.0",
+      "12.1+PTX",
+      "7.5",
+      "8.0",
+      "8.6",
+      "8.7",
+      "8.9",
+      "9.0"
+    ]
+  }
+}

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

	@@ -0,0 +1,271 @@

+"""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
+__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",
+    "swizzled_sf_size",
+]

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

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

+version https://git-lfs.github.com/spec/v1
+oid sha256:7e0327c2ffbc171f1b2b70447f0076a4d0acd80bb1eb716045ee973e867a030e
+size 1320088

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

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

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

build/torch212-cxx11-cu130-x86_64-linux/adaptive_layernorm_producers/__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-cu130-x86_64-linux/metadata.json ADDED Viewed

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

+{
+  "name": "adaptive-layernorm-producers",
+  "id": "_adaptive_layernorm_producers_cuda_1798e7f",
+  "version": 1,
+  "license": "Apache-2.0",
+  "python-depends": [],
+  "backend": {
+    "type": "cuda",
+    "archs": [
+      "10.0",
+      "11.0",
+      "12.0",
+      "12.1+PTX",
+      "7.5",
+      "8.0",
+      "8.6",
+      "8.7",
+      "8.9",
+      "9.0"
+    ]
+  }
+}

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

	@@ -0,0 +1,271 @@

+"""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
+__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",
+    "swizzled_sf_size",
+]

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

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

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

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

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

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

build/torch212-cxx11-cu132-x86_64-linux/adaptive_layernorm_producers/__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-cu132-x86_64-linux/metadata.json ADDED Viewed

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

+{
+  "name": "adaptive-layernorm-producers",
+  "id": "_adaptive_layernorm_producers_cuda_1798e7f",
+  "version": 1,
+  "license": "Apache-2.0",
+  "python-depends": [],
+  "backend": {
+    "type": "cuda",
+    "archs": [
+      "10.0",
+      "11.0",
+      "12.0",
+      "12.1+PTX",
+      "7.5",
+      "8.0",
+      "8.6",
+      "8.7",
+      "8.9",
+      "9.0"
+    ]
+  }
+}