Promote latest kernel artifacts to main

.gitattributes

@@ -1,35 +1,3 @@
-*.7z filter=lfs diff=lfs merge=lfs -text
-*.arrow filter=lfs diff=lfs merge=lfs -text
-*.bin filter=lfs diff=lfs merge=lfs -text
-*.bz2 filter=lfs diff=lfs merge=lfs -text
-*.ckpt filter=lfs diff=lfs merge=lfs -text
-*.ftz filter=lfs diff=lfs merge=lfs -text
-*.gz filter=lfs diff=lfs merge=lfs -text
-*.h5 filter=lfs diff=lfs merge=lfs -text
-*.joblib filter=lfs diff=lfs merge=lfs -text
-*.lfs.* filter=lfs diff=lfs merge=lfs -text
-*.mlmodel filter=lfs diff=lfs merge=lfs -text
-*.model filter=lfs diff=lfs merge=lfs -text
-*.msgpack filter=lfs diff=lfs merge=lfs -text
-*.npy filter=lfs diff=lfs merge=lfs -text
-*.npz filter=lfs diff=lfs merge=lfs -text
-*.onnx filter=lfs diff=lfs merge=lfs -text
-*.ot filter=lfs diff=lfs merge=lfs -text
-*.parquet filter=lfs diff=lfs merge=lfs -text
-*.pb filter=lfs diff=lfs merge=lfs -text
-*.pickle filter=lfs diff=lfs merge=lfs -text
-*.pkl filter=lfs diff=lfs merge=lfs -text
-*.pt filter=lfs diff=lfs merge=lfs -text
-*.pth filter=lfs diff=lfs merge=lfs -text
-*.rar filter=lfs diff=lfs merge=lfs -text
-*.safetensors filter=lfs diff=lfs merge=lfs -text
-saved_model/**/* filter=lfs diff=lfs merge=lfs -text
-*.tar.* filter=lfs diff=lfs merge=lfs -text
-*.tar filter=lfs diff=lfs merge=lfs -text
-*.tflite filter=lfs diff=lfs merge=lfs -text
-*.tgz filter=lfs diff=lfs merge=lfs -text
-*.wasm filter=lfs diff=lfs merge=lfs -text
-*.xz filter=lfs diff=lfs merge=lfs -text
-*.zip filter=lfs diff=lfs merge=lfs -text
-*.zst filter=lfs diff=lfs merge=lfs -text
-*tfevents* filter=lfs diff=lfs merge=lfs -text
+*.so filter=lfs diff=lfs merge=lfs -text
+*.pyd filter=lfs diff=lfs merge=lfs -text
+*.dylib filter=lfs diff=lfs merge=lfs -text

README.md

@@ -1,9 +0,0 @@
-# flashrt/blockwise-fp8-producers
-
-This repository is a compatibility mirror for older `kernels` clients
-that resolve repositories through the default Hugging Face model repo API.
-
-Canonical Kernel Hub repo: https://huggingface.co/kernels/flashrt/blockwise-fp8-producers
-
-Do not edit this mirror by hand. It is generated from the Kernel Hub
-`vN` branches and contains the same `build/**` artifacts.

benchmarks/benchmark.py

@@ -0,0 +1,123 @@
+#!/usr/bin/env python3
+"""Benchmark blockwise FP8 producer APIs."""
+
+from __future__ import annotations
+
+import argparse
+import importlib
+import sys
+from pathlib import Path
+
+import torch
+import torch.nn.functional as F
+
+ROOT = Path(__file__).resolve().parents[2]
+sys.path.insert(0, str(ROOT / "blockwise-fp8-producers" / "tests"))
+from test_blockwise_fp8_producers import load_source_ops  # noqa: E402
+
+
+def load_ops(backend: str, artifact: str | None):
+    if backend == "source":
+        return load_source_ops()
+    if artifact:
+        sys.path.insert(0, artifact)
+    try:
+        return importlib.import_module("blockwise_fp8_producers")
+    finally:
+        if artifact:
+            sys.path.remove(artifact)
+
+
+def reference(kind, x, weight, bias):
+    if kind == "layer_norm":
+        produced = F.layer_norm(
+            x.float(), (x.shape[1],), weight.float(), bias.float(), 1e-6
+        )
+    elif kind == "rms_norm":
+        produced = (
+            x.float()
+            * torch.rsqrt(x.float().square().mean(-1, keepdim=True) + 1e-6)
+            * weight.float()
+        )
+    elif kind == "gelu_bias":
+        produced = F.gelu(x.float() + bias.float(), approximate="tanh")
+    else:
+        produced = x.float()
+    blocks = produced.reshape(produced.shape[0], produced.shape[1] // 128, 128)
+    scale = torch.clamp(blocks.abs().amax(-1) / 448.0, min=1.0e-12)
+    quantized = torch.clamp(
+        blocks / scale.unsqueeze(-1), -448.0, 448.0
+    ).to(torch.float8_e4m3fn)
+    return quantized.reshape_as(x), scale
+
+
+def time_us(fn, warmup: int, iters: int) -> 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 main() -> int:
+    parser = argparse.ArgumentParser()
+    parser.add_argument("--backend", choices=["source", "installed"], default="source")
+    parser.add_argument("--artifact", default=None)
+    parser.add_argument("--mode", choices=["headline", "full"], default="headline")
+    parser.add_argument("--warmup", type=int, default=30)
+    parser.add_argument("--iters", type=int, default=200)
+    args = parser.parse_args()
+    ops = load_ops(args.backend, args.artifact)
+    shapes = [(51, 4096), (277, 9216), (1024, 1152)]
+    if args.mode == "full":
+        shapes = [(1, 4096), (17, 1152), (51, 4096), (65, 4352), (277, 9216), (1024, 1152)]
+
+    print("kind,shape,artifact_us,eager_us,compile_us,eager_speedup,compile_speedup")
+    for rows, dim in shapes:
+        x = torch.randn((rows, dim), device="cuda", dtype=torch.bfloat16)
+        weight = torch.randn((dim,), device="cuda", dtype=torch.bfloat16)
+        bias = torch.randn((dim,), device="cuda", dtype=torch.bfloat16)
+        output = torch.empty_like(x, dtype=torch.float8_e4m3fn)
+        scale = torch.empty((rows, dim // 128), device="cuda", dtype=torch.float32)
+        for kind in ("quantize", "layer_norm", "rms_norm", "gelu_bias"):
+            if kind == "quantize":
+                artifact_fn = lambda: ops.quantize_fp8_block128_bf16(
+                    x, output=output, scale=scale
+                )
+            elif kind == "layer_norm":
+                artifact_fn = lambda: ops.layer_norm_fp8_block128_bf16(
+                    x, weight, bias, output=output, scale=scale
+                )
+            elif kind == "rms_norm":
+                artifact_fn = lambda: ops.rms_norm_fp8_block128_bf16(
+                    x, weight, output=output, scale=scale
+                )
+            else:
+                artifact_fn = lambda: ops.gelu_tanh_bias_fp8_block128_bf16(
+                    x, bias, output=output, scale=scale
+                )
+            eager_fn = lambda: reference(kind, x, weight, bias)
+            torch._dynamo.reset()
+            compiled = torch.compile(
+                lambda a, w, b: reference(kind, a, w, b), fullgraph=True
+            )
+            compiled_fn = lambda: compiled(x, weight, bias)
+            artifact_us = time_us(artifact_fn, args.warmup, args.iters)
+            eager_us = time_us(eager_fn, max(10, args.warmup // 2), max(50, args.iters // 2))
+            compile_us = time_us(compiled_fn, max(10, args.warmup // 2), max(50, args.iters // 2))
+            print(
+                f"{kind},{rows}x{dim},{artifact_us:.3f},{eager_us:.3f},"
+                f"{compile_us:.3f},{eager_us/artifact_us:.2f}x,"
+                f"{compile_us/artifact_us:.2f}x"
+            )
+    return 0
+
+
+if __name__ == "__main__":
+    raise SystemExit(main())

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

@@ -0,0 +1,297 @@
+"""Blockwise FP8 producers for transformer and world-model regions."""
+
+from __future__ import annotations
+
+from typing import Optional
+
+import torch
+
+from ._ops import add_op_namespace_prefix, ops
+
+
+def _check_matrix(input: torch.Tensor, output: torch.Tensor, scale: torch.Tensor) -> None:
+    if (
+        input.dim() != 2
+        or input.shape[0] <= 0
+        or input.shape[1] <= 0
+        or input.shape[1] % 128 != 0
+        or output.shape != input.shape
+        or scale.shape != (input.shape[0], input.shape[1] // 128)
+    ):
+        raise RuntimeError(
+            "expected input/output (rows, dim) with dim a positive multiple "
+            "of 128 and scale (rows, dim / 128)"
+        )
+
+
+@torch.library.register_fake(add_op_namespace_prefix("quantize_fp8_block128_bf16"))
+def _quantize_fake(
+    input: torch.Tensor, output: torch.Tensor, scale: torch.Tensor
+) -> None:
+    _check_matrix(input, output, scale)
+    return None
+
+
+@torch.library.register_fake(add_op_namespace_prefix("layer_norm_fp8_block128_bf16"))
+def _layer_norm_fake(
+    input: torch.Tensor,
+    weight: torch.Tensor,
+    bias: torch.Tensor,
+    eps: float,
+    output: torch.Tensor,
+    scale: torch.Tensor,
+) -> None:
+    _check_matrix(input, output, scale)
+    if weight.shape != (input.shape[1],) or bias.shape != weight.shape:
+        raise RuntimeError("weight and bias must have shape (dim,)")
+    return None
+
+
+@torch.library.register_fake(add_op_namespace_prefix("rms_norm_fp8_block128_bf16"))
+def _rms_norm_fake(
+    input: torch.Tensor,
+    weight: torch.Tensor,
+    eps: float,
+    output: torch.Tensor,
+    scale: torch.Tensor,
+) -> None:
+    _check_matrix(input, output, scale)
+    if weight.shape != (input.shape[1],):
+        raise RuntimeError("weight must have shape (dim,)")
+    return None
+
+
+@torch.library.register_fake(
+    add_op_namespace_prefix("residual_add_rms_norm_fp8_block128_bf16")
+)
+def _residual_rms_norm_fake(
+    residual: torch.Tensor,
+    input: torch.Tensor,
+    weight: torch.Tensor,
+    eps: float,
+    residual_out: torch.Tensor,
+    output: torch.Tensor,
+    scale: torch.Tensor,
+) -> None:
+    _check_matrix(input, output, scale)
+    if (
+        residual.shape != input.shape
+        or residual_out.shape != input.shape
+        or weight.shape != (input.shape[1],)
+    ):
+        raise RuntimeError("residual/output must match input and weight must be (dim,)")
+    return None
+
+
+@torch.library.register_fake(add_op_namespace_prefix("gelu_tanh_fp8_block128_bf16"))
+def _gelu_fake(
+    input: torch.Tensor, output: torch.Tensor, scale: torch.Tensor
+) -> None:
+    _check_matrix(input, output, scale)
+    return None
+
+
+@torch.library.register_fake(
+    add_op_namespace_prefix("gelu_tanh_bias_fp8_block128_bf16")
+)
+def _gelu_bias_fake(
+    input: torch.Tensor,
+    bias: torch.Tensor,
+    output: torch.Tensor,
+    scale: torch.Tensor,
+) -> None:
+    _check_matrix(input, output, scale)
+    if bias.shape != (input.shape[1],):
+        raise RuntimeError("bias must have shape (dim,)")
+    return None
+
+
+@torch.library.register_fake(add_op_namespace_prefix("silu_mul_fp8_block128_bf16"))
+def _silu_mul_fake(
+    gate: torch.Tensor,
+    up: torch.Tensor,
+    output: torch.Tensor,
+    scale: torch.Tensor,
+) -> None:
+    _check_matrix(gate, output, scale)
+    if up.shape != gate.shape:
+        raise RuntimeError("up must match gate")
+    return None
+
+
+@torch.library.register_fake(
+    add_op_namespace_prefix("silu_mul_merged_fp8_block128_bf16")
+)
+def _silu_merged_fake(
+    gate_up: torch.Tensor, output: torch.Tensor, scale: torch.Tensor
+) -> None:
+    if (
+        gate_up.dim() != 2
+        or gate_up.shape[0] <= 0
+        or gate_up.shape[1] <= 0
+        or gate_up.shape[1] % 256 != 0
+        or output.shape != (gate_up.shape[0], gate_up.shape[1] // 2)
+        or scale.shape != (gate_up.shape[0], gate_up.shape[1] // 256)
+    ):
+        raise RuntimeError(
+            "gate_up must be (rows, 2 * dim), dim multiple of 128"
+        )
+    return None
+
+
+def _allocate(input: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor]:
+    return (
+        torch.empty_like(input, dtype=torch.float8_e4m3fn),
+        torch.empty(
+            (input.shape[0], input.shape[1] // 128),
+            device=input.device,
+            dtype=torch.float32,
+        ),
+    )
+
+
+def quantize_fp8_block128_bf16(
+    input: torch.Tensor,
+    *,
+    output: Optional[torch.Tensor] = None,
+    scale: Optional[torch.Tensor] = None,
+) -> tuple[torch.Tensor, torch.Tensor]:
+    if output is None or scale is None:
+        allocated_output, allocated_scale = _allocate(input)
+        output = allocated_output if output is None else output
+        scale = allocated_scale if scale is None else scale
+    ops.quantize_fp8_block128_bf16(input, output, scale)
+    return output, scale
+
+
+def layer_norm_fp8_block128_bf16(
+    input: torch.Tensor,
+    weight: torch.Tensor,
+    bias: torch.Tensor,
+    eps: float = 1e-6,
+    *,
+    output: Optional[torch.Tensor] = None,
+    scale: Optional[torch.Tensor] = None,
+) -> tuple[torch.Tensor, torch.Tensor]:
+    if output is None or scale is None:
+        allocated_output, allocated_scale = _allocate(input)
+        output = allocated_output if output is None else output
+        scale = allocated_scale if scale is None else scale
+    ops.layer_norm_fp8_block128_bf16(
+        input, weight, bias, float(eps), output, scale
+    )
+    return output, scale
+
+
+def rms_norm_fp8_block128_bf16(
+    input: torch.Tensor,
+    weight: torch.Tensor,
+    eps: float = 1e-6,
+    *,
+    output: Optional[torch.Tensor] = None,
+    scale: Optional[torch.Tensor] = None,
+) -> tuple[torch.Tensor, torch.Tensor]:
+    if output is None or scale is None:
+        allocated_output, allocated_scale = _allocate(input)
+        output = allocated_output if output is None else output
+        scale = allocated_scale if scale is None else scale
+    ops.rms_norm_fp8_block128_bf16(input, weight, float(eps), output, scale)
+    return output, scale
+
+
+def residual_add_rms_norm_fp8_block128_bf16(
+    residual: torch.Tensor,
+    input: torch.Tensor,
+    weight: torch.Tensor,
+    eps: float = 1e-6,
+    *,
+    residual_out: Optional[torch.Tensor] = None,
+    output: Optional[torch.Tensor] = None,
+    scale: Optional[torch.Tensor] = None,
+) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
+    if residual_out is None:
+        residual_out = torch.empty_like(input)
+    if output is None or scale is None:
+        allocated_output, allocated_scale = _allocate(input)
+        output = allocated_output if output is None else output
+        scale = allocated_scale if scale is None else scale
+    ops.residual_add_rms_norm_fp8_block128_bf16(
+        residual, input, weight, float(eps), residual_out, output, scale
+    )
+    return residual_out, output, scale
+
+
+def gelu_tanh_fp8_block128_bf16(
+    input: torch.Tensor,
+    *,
+    output: Optional[torch.Tensor] = None,
+    scale: Optional[torch.Tensor] = None,
+) -> tuple[torch.Tensor, torch.Tensor]:
+    if output is None or scale is None:
+        allocated_output, allocated_scale = _allocate(input)
+        output = allocated_output if output is None else output
+        scale = allocated_scale if scale is None else scale
+    ops.gelu_tanh_fp8_block128_bf16(input, output, scale)
+    return output, scale
+
+
+def gelu_tanh_bias_fp8_block128_bf16(
+    input: torch.Tensor,
+    bias: torch.Tensor,
+    *,
+    output: Optional[torch.Tensor] = None,
+    scale: Optional[torch.Tensor] = None,
+) -> tuple[torch.Tensor, torch.Tensor]:
+    if output is None or scale is None:
+        allocated_output, allocated_scale = _allocate(input)
+        output = allocated_output if output is None else output
+        scale = allocated_scale if scale is None else scale
+    ops.gelu_tanh_bias_fp8_block128_bf16(input, bias, output, scale)
+    return output, scale
+
+
+def silu_mul_fp8_block128_bf16(
+    gate: torch.Tensor,
+    up: torch.Tensor,
+    *,
+    output: Optional[torch.Tensor] = None,
+    scale: Optional[torch.Tensor] = None,
+) -> tuple[torch.Tensor, torch.Tensor]:
+    if output is None or scale is None:
+        allocated_output, allocated_scale = _allocate(gate)
+        output = allocated_output if output is None else output
+        scale = allocated_scale if scale is None else scale
+    ops.silu_mul_fp8_block128_bf16(gate, up, output, scale)
+    return output, scale
+
+
+def silu_mul_merged_fp8_block128_bf16(
+    gate_up: torch.Tensor,
+    *,
+    output: Optional[torch.Tensor] = None,
+    scale: Optional[torch.Tensor] = None,
+) -> tuple[torch.Tensor, torch.Tensor]:
+    rows, merged_dim = gate_up.shape
+    dim = merged_dim // 2
+    if output is None:
+        output = torch.empty(
+            (rows, dim), device=gate_up.device, dtype=torch.float8_e4m3fn
+        )
+    if scale is None:
+        scale = torch.empty(
+            (rows, dim // 128), device=gate_up.device, dtype=torch.float32
+        )
+    ops.silu_mul_merged_fp8_block128_bf16(gate_up, output, scale)
+    return output, scale
+
+
+__all__ = [
+    "quantize_fp8_block128_bf16",
+    "layer_norm_fp8_block128_bf16",
+    "rms_norm_fp8_block128_bf16",
+    "residual_add_rms_norm_fp8_block128_bf16",
+    "gelu_tanh_fp8_block128_bf16",
+    "gelu_tanh_bias_fp8_block128_bf16",
+    "silu_mul_fp8_block128_bf16",
+    "silu_mul_merged_fp8_block128_bf16",
+]

build/torch211-cxx11-cu128-x86_64-linux/_blockwise_fp8_producers_cuda_7781728.abi3.so

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:03a00a579f79e06cf2047011f707091466083eaa57930fa6d2f5dbff919f4119
+size 2044952

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

@@ -0,0 +1,9 @@
+import torch
+from . import _blockwise_fp8_producers_cuda_7781728
+ops = torch.ops._blockwise_fp8_producers_cuda_7781728
+
+def add_op_namespace_prefix(op_name: str):
+    """
+    Prefix op by namespace.
+    """
+    return f"_blockwise_fp8_producers_cuda_7781728::{op_name}"

build/torch211-cxx11-cu128-x86_64-linux/blockwise_fp8_producers/__init__.py

@@ -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

@@ -0,0 +1,23 @@
+{
+  "name": "blockwise-fp8-producers",
+  "id": "_blockwise_fp8_producers_cuda_7781728",
+  "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-aarch64-linux/__init__.py

@@ -0,0 +1,297 @@
+"""Blockwise FP8 producers for transformer and world-model regions."""
+
+from __future__ import annotations
+
+from typing import Optional
+
+import torch
+
+from ._ops import add_op_namespace_prefix, ops
+
+
+def _check_matrix(input: torch.Tensor, output: torch.Tensor, scale: torch.Tensor) -> None:
+    if (
+        input.dim() != 2
+        or input.shape[0] <= 0
+        or input.shape[1] <= 0
+        or input.shape[1] % 128 != 0
+        or output.shape != input.shape
+        or scale.shape != (input.shape[0], input.shape[1] // 128)
+    ):
+        raise RuntimeError(
+            "expected input/output (rows, dim) with dim a positive multiple "
+            "of 128 and scale (rows, dim / 128)"
+        )
+
+
+@torch.library.register_fake(add_op_namespace_prefix("quantize_fp8_block128_bf16"))
+def _quantize_fake(
+    input: torch.Tensor, output: torch.Tensor, scale: torch.Tensor
+) -> None:
+    _check_matrix(input, output, scale)
+    return None
+
+
+@torch.library.register_fake(add_op_namespace_prefix("layer_norm_fp8_block128_bf16"))
+def _layer_norm_fake(
+    input: torch.Tensor,
+    weight: torch.Tensor,
+    bias: torch.Tensor,
+    eps: float,
+    output: torch.Tensor,
+    scale: torch.Tensor,
+) -> None:
+    _check_matrix(input, output, scale)
+    if weight.shape != (input.shape[1],) or bias.shape != weight.shape:
+        raise RuntimeError("weight and bias must have shape (dim,)")
+    return None
+
+
+@torch.library.register_fake(add_op_namespace_prefix("rms_norm_fp8_block128_bf16"))
+def _rms_norm_fake(
+    input: torch.Tensor,
+    weight: torch.Tensor,
+    eps: float,
+    output: torch.Tensor,
+    scale: torch.Tensor,
+) -> None:
+    _check_matrix(input, output, scale)
+    if weight.shape != (input.shape[1],):
+        raise RuntimeError("weight must have shape (dim,)")
+    return None
+
+
+@torch.library.register_fake(
+    add_op_namespace_prefix("residual_add_rms_norm_fp8_block128_bf16")
+)
+def _residual_rms_norm_fake(
+    residual: torch.Tensor,
+    input: torch.Tensor,
+    weight: torch.Tensor,
+    eps: float,
+    residual_out: torch.Tensor,
+    output: torch.Tensor,
+    scale: torch.Tensor,
+) -> None:
+    _check_matrix(input, output, scale)
+    if (
+        residual.shape != input.shape
+        or residual_out.shape != input.shape
+        or weight.shape != (input.shape[1],)
+    ):
+        raise RuntimeError("residual/output must match input and weight must be (dim,)")
+    return None
+
+
+@torch.library.register_fake(add_op_namespace_prefix("gelu_tanh_fp8_block128_bf16"))
+def _gelu_fake(
+    input: torch.Tensor, output: torch.Tensor, scale: torch.Tensor
+) -> None:
+    _check_matrix(input, output, scale)
+    return None
+
+
+@torch.library.register_fake(
+    add_op_namespace_prefix("gelu_tanh_bias_fp8_block128_bf16")
+)
+def _gelu_bias_fake(
+    input: torch.Tensor,
+    bias: torch.Tensor,
+    output: torch.Tensor,
+    scale: torch.Tensor,
+) -> None:
+    _check_matrix(input, output, scale)
+    if bias.shape != (input.shape[1],):
+        raise RuntimeError("bias must have shape (dim,)")
+    return None
+
+
+@torch.library.register_fake(add_op_namespace_prefix("silu_mul_fp8_block128_bf16"))
+def _silu_mul_fake(
+    gate: torch.Tensor,
+    up: torch.Tensor,
+    output: torch.Tensor,
+    scale: torch.Tensor,
+) -> None:
+    _check_matrix(gate, output, scale)
+    if up.shape != gate.shape:
+        raise RuntimeError("up must match gate")
+    return None
+
+
+@torch.library.register_fake(
+    add_op_namespace_prefix("silu_mul_merged_fp8_block128_bf16")
+)
+def _silu_merged_fake(
+    gate_up: torch.Tensor, output: torch.Tensor, scale: torch.Tensor
+) -> None:
+    if (
+        gate_up.dim() != 2
+        or gate_up.shape[0] <= 0
+        or gate_up.shape[1] <= 0
+        or gate_up.shape[1] % 256 != 0
+        or output.shape != (gate_up.shape[0], gate_up.shape[1] // 2)
+        or scale.shape != (gate_up.shape[0], gate_up.shape[1] // 256)
+    ):
+        raise RuntimeError(
+            "gate_up must be (rows, 2 * dim), dim multiple of 128"
+        )
+    return None
+
+
+def _allocate(input: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor]:
+    return (
+        torch.empty_like(input, dtype=torch.float8_e4m3fn),
+        torch.empty(
+            (input.shape[0], input.shape[1] // 128),
+            device=input.device,
+            dtype=torch.float32,
+        ),
+    )
+
+
+def quantize_fp8_block128_bf16(
+    input: torch.Tensor,
+    *,
+    output: Optional[torch.Tensor] = None,
+    scale: Optional[torch.Tensor] = None,
+) -> tuple[torch.Tensor, torch.Tensor]:
+    if output is None or scale is None:
+        allocated_output, allocated_scale = _allocate(input)
+        output = allocated_output if output is None else output
+        scale = allocated_scale if scale is None else scale
+    ops.quantize_fp8_block128_bf16(input, output, scale)
+    return output, scale
+
+
+def layer_norm_fp8_block128_bf16(
+    input: torch.Tensor,
+    weight: torch.Tensor,
+    bias: torch.Tensor,
+    eps: float = 1e-6,
+    *,
+    output: Optional[torch.Tensor] = None,
+    scale: Optional[torch.Tensor] = None,
+) -> tuple[torch.Tensor, torch.Tensor]:
+    if output is None or scale is None:
+        allocated_output, allocated_scale = _allocate(input)
+        output = allocated_output if output is None else output
+        scale = allocated_scale if scale is None else scale
+    ops.layer_norm_fp8_block128_bf16(
+        input, weight, bias, float(eps), output, scale
+    )
+    return output, scale
+
+
+def rms_norm_fp8_block128_bf16(
+    input: torch.Tensor,
+    weight: torch.Tensor,
+    eps: float = 1e-6,
+    *,
+    output: Optional[torch.Tensor] = None,
+    scale: Optional[torch.Tensor] = None,
+) -> tuple[torch.Tensor, torch.Tensor]:
+    if output is None or scale is None:
+        allocated_output, allocated_scale = _allocate(input)
+        output = allocated_output if output is None else output
+        scale = allocated_scale if scale is None else scale
+    ops.rms_norm_fp8_block128_bf16(input, weight, float(eps), output, scale)
+    return output, scale
+
+
+def residual_add_rms_norm_fp8_block128_bf16(
+    residual: torch.Tensor,
+    input: torch.Tensor,
+    weight: torch.Tensor,
+    eps: float = 1e-6,
+    *,
+    residual_out: Optional[torch.Tensor] = None,
+    output: Optional[torch.Tensor] = None,
+    scale: Optional[torch.Tensor] = None,
+) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
+    if residual_out is None:
+        residual_out = torch.empty_like(input)
+    if output is None or scale is None:
+        allocated_output, allocated_scale = _allocate(input)
+        output = allocated_output if output is None else output
+        scale = allocated_scale if scale is None else scale
+    ops.residual_add_rms_norm_fp8_block128_bf16(
+        residual, input, weight, float(eps), residual_out, output, scale
+    )
+    return residual_out, output, scale
+
+
+def gelu_tanh_fp8_block128_bf16(
+    input: torch.Tensor,
+    *,
+    output: Optional[torch.Tensor] = None,
+    scale: Optional[torch.Tensor] = None,
+) -> tuple[torch.Tensor, torch.Tensor]:
+    if output is None or scale is None:
+        allocated_output, allocated_scale = _allocate(input)
+        output = allocated_output if output is None else output
+        scale = allocated_scale if scale is None else scale
+    ops.gelu_tanh_fp8_block128_bf16(input, output, scale)
+    return output, scale
+
+
+def gelu_tanh_bias_fp8_block128_bf16(
+    input: torch.Tensor,
+    bias: torch.Tensor,
+    *,
+    output: Optional[torch.Tensor] = None,
+    scale: Optional[torch.Tensor] = None,
+) -> tuple[torch.Tensor, torch.Tensor]:
+    if output is None or scale is None:
+        allocated_output, allocated_scale = _allocate(input)
+        output = allocated_output if output is None else output
+        scale = allocated_scale if scale is None else scale
+    ops.gelu_tanh_bias_fp8_block128_bf16(input, bias, output, scale)
+    return output, scale
+
+
+def silu_mul_fp8_block128_bf16(
+    gate: torch.Tensor,
+    up: torch.Tensor,
+    *,
+    output: Optional[torch.Tensor] = None,
+    scale: Optional[torch.Tensor] = None,
+) -> tuple[torch.Tensor, torch.Tensor]:
+    if output is None or scale is None:
+        allocated_output, allocated_scale = _allocate(gate)
+        output = allocated_output if output is None else output
+        scale = allocated_scale if scale is None else scale
+    ops.silu_mul_fp8_block128_bf16(gate, up, output, scale)
+    return output, scale
+
+
+def silu_mul_merged_fp8_block128_bf16(
+    gate_up: torch.Tensor,
+    *,
+    output: Optional[torch.Tensor] = None,
+    scale: Optional[torch.Tensor] = None,
+) -> tuple[torch.Tensor, torch.Tensor]:
+    rows, merged_dim = gate_up.shape
+    dim = merged_dim // 2
+    if output is None:
+        output = torch.empty(
+            (rows, dim), device=gate_up.device, dtype=torch.float8_e4m3fn
+        )
+    if scale is None:
+        scale = torch.empty(
+            (rows, dim // 128), device=gate_up.device, dtype=torch.float32
+        )
+    ops.silu_mul_merged_fp8_block128_bf16(gate_up, output, scale)
+    return output, scale
+
+
+__all__ = [
+    "quantize_fp8_block128_bf16",
+    "layer_norm_fp8_block128_bf16",
+    "rms_norm_fp8_block128_bf16",
+    "residual_add_rms_norm_fp8_block128_bf16",
+    "gelu_tanh_fp8_block128_bf16",
+    "gelu_tanh_bias_fp8_block128_bf16",
+    "silu_mul_fp8_block128_bf16",
+    "silu_mul_merged_fp8_block128_bf16",
+]

build/torch211-cxx11-cu130-aarch64-linux/_blockwise_fp8_producers_cuda_7781728.abi3.so

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:293d539efbdddbf49f647c269e2fd34d619e35ae6eb462dac4434780075d0bf6
+size 455696

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

@@ -0,0 +1,6 @@
+import torch
+from . import _blockwise_fp8_producers_cuda_7781728
+ops = torch.ops._blockwise_fp8_producers_cuda_7781728
+
+def add_op_namespace_prefix(op_name: str):
+    return f"_blockwise_fp8_producers_cuda_7781728::{op_name}"

build/torch211-cxx11-cu130-aarch64-linux/blockwise_fp8_producers/__init__.py

@@ -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/metadata.json

@@ -0,0 +1,22 @@
+{
+  "name": "blockwise-fp8-producers",
+  "id": "_blockwise_fp8_producers_cuda_7781728",
+  "version": 1,
+  "license": "Apache-2.0",
+  "python-depends": [],
+  "backend": {
+    "type": "cuda",
+    "archs": [
+      "11.0"
+    ]
+  },
+  "digest": {
+    "algorithm": "sha256",
+    "files": {
+      "__init__.py": "jTa/yZ4lH7PUVKeENwdd7iZgdhrc9AOOJ8+L2usA23M=",
+      "_blockwise_fp8_producers_cuda_7781728.abi3.so": "KT1Tnvvd2/SfZHwmni/TTWGeNa5utGLaxENHgAddC/Y=",
+      "_ops.py": "pl0Svm/AyQFPOD+QjsTr7CaugfBN69uRVroTvi3cHrE=",
+      "blockwise_fp8_producers/__init__.py": "v6p5XMfQzddhi1fLSAw4HX9CyS0rQsidvu9VsT01xi4="
+    }
+  }
+}

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

@@ -0,0 +1,297 @@
+"""Blockwise FP8 producers for transformer and world-model regions."""
+
+from __future__ import annotations
+
+from typing import Optional
+
+import torch
+
+from ._ops import add_op_namespace_prefix, ops
+
+
+def _check_matrix(input: torch.Tensor, output: torch.Tensor, scale: torch.Tensor) -> None:
+    if (
+        input.dim() != 2
+        or input.shape[0] <= 0
+        or input.shape[1] <= 0
+        or input.shape[1] % 128 != 0
+        or output.shape != input.shape
+        or scale.shape != (input.shape[0], input.shape[1] // 128)
+    ):
+        raise RuntimeError(
+            "expected input/output (rows, dim) with dim a positive multiple "
+            "of 128 and scale (rows, dim / 128)"
+        )
+
+
+@torch.library.register_fake(add_op_namespace_prefix("quantize_fp8_block128_bf16"))
+def _quantize_fake(
+    input: torch.Tensor, output: torch.Tensor, scale: torch.Tensor
+) -> None:
+    _check_matrix(input, output, scale)
+    return None
+
+
+@torch.library.register_fake(add_op_namespace_prefix("layer_norm_fp8_block128_bf16"))
+def _layer_norm_fake(
+    input: torch.Tensor,
+    weight: torch.Tensor,
+    bias: torch.Tensor,
+    eps: float,
+    output: torch.Tensor,
+    scale: torch.Tensor,
+) -> None:
+    _check_matrix(input, output, scale)
+    if weight.shape != (input.shape[1],) or bias.shape != weight.shape:
+        raise RuntimeError("weight and bias must have shape (dim,)")
+    return None
+
+
+@torch.library.register_fake(add_op_namespace_prefix("rms_norm_fp8_block128_bf16"))
+def _rms_norm_fake(
+    input: torch.Tensor,
+    weight: torch.Tensor,
+    eps: float,
+    output: torch.Tensor,
+    scale: torch.Tensor,
+) -> None:
+    _check_matrix(input, output, scale)
+    if weight.shape != (input.shape[1],):
+        raise RuntimeError("weight must have shape (dim,)")
+    return None
+
+
+@torch.library.register_fake(
+    add_op_namespace_prefix("residual_add_rms_norm_fp8_block128_bf16")
+)
+def _residual_rms_norm_fake(
+    residual: torch.Tensor,
+    input: torch.Tensor,
+    weight: torch.Tensor,
+    eps: float,
+    residual_out: torch.Tensor,
+    output: torch.Tensor,
+    scale: torch.Tensor,
+) -> None:
+    _check_matrix(input, output, scale)
+    if (
+        residual.shape != input.shape
+        or residual_out.shape != input.shape
+        or weight.shape != (input.shape[1],)
+    ):
+        raise RuntimeError("residual/output must match input and weight must be (dim,)")
+    return None
+
+
+@torch.library.register_fake(add_op_namespace_prefix("gelu_tanh_fp8_block128_bf16"))
+def _gelu_fake(
+    input: torch.Tensor, output: torch.Tensor, scale: torch.Tensor
+) -> None:
+    _check_matrix(input, output, scale)
+    return None
+
+
+@torch.library.register_fake(
+    add_op_namespace_prefix("gelu_tanh_bias_fp8_block128_bf16")
+)
+def _gelu_bias_fake(
+    input: torch.Tensor,
+    bias: torch.Tensor,
+    output: torch.Tensor,
+    scale: torch.Tensor,
+) -> None:
+    _check_matrix(input, output, scale)
+    if bias.shape != (input.shape[1],):
+        raise RuntimeError("bias must have shape (dim,)")
+    return None
+
+
+@torch.library.register_fake(add_op_namespace_prefix("silu_mul_fp8_block128_bf16"))
+def _silu_mul_fake(
+    gate: torch.Tensor,
+    up: torch.Tensor,
+    output: torch.Tensor,
+    scale: torch.Tensor,
+) -> None:
+    _check_matrix(gate, output, scale)
+    if up.shape != gate.shape:
+        raise RuntimeError("up must match gate")
+    return None
+
+
+@torch.library.register_fake(
+    add_op_namespace_prefix("silu_mul_merged_fp8_block128_bf16")
+)
+def _silu_merged_fake(
+    gate_up: torch.Tensor, output: torch.Tensor, scale: torch.Tensor
+) -> None:
+    if (
+        gate_up.dim() != 2
+        or gate_up.shape[0] <= 0
+        or gate_up.shape[1] <= 0
+        or gate_up.shape[1] % 256 != 0
+        or output.shape != (gate_up.shape[0], gate_up.shape[1] // 2)
+        or scale.shape != (gate_up.shape[0], gate_up.shape[1] // 256)
+    ):
+        raise RuntimeError(
+            "gate_up must be (rows, 2 * dim), dim multiple of 128"
+        )
+    return None
+
+
+def _allocate(input: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor]:
+    return (
+        torch.empty_like(input, dtype=torch.float8_e4m3fn),
+        torch.empty(
+            (input.shape[0], input.shape[1] // 128),
+            device=input.device,
+            dtype=torch.float32,
+        ),
+    )
+
+
+def quantize_fp8_block128_bf16(
+    input: torch.Tensor,
+    *,
+    output: Optional[torch.Tensor] = None,
+    scale: Optional[torch.Tensor] = None,
+) -> tuple[torch.Tensor, torch.Tensor]:
+    if output is None or scale is None:
+        allocated_output, allocated_scale = _allocate(input)
+        output = allocated_output if output is None else output
+        scale = allocated_scale if scale is None else scale
+    ops.quantize_fp8_block128_bf16(input, output, scale)
+    return output, scale
+
+
+def layer_norm_fp8_block128_bf16(
+    input: torch.Tensor,
+    weight: torch.Tensor,
+    bias: torch.Tensor,
+    eps: float = 1e-6,
+    *,
+    output: Optional[torch.Tensor] = None,
+    scale: Optional[torch.Tensor] = None,
+) -> tuple[torch.Tensor, torch.Tensor]:
+    if output is None or scale is None:
+        allocated_output, allocated_scale = _allocate(input)
+        output = allocated_output if output is None else output
+        scale = allocated_scale if scale is None else scale
+    ops.layer_norm_fp8_block128_bf16(
+        input, weight, bias, float(eps), output, scale
+    )
+    return output, scale
+
+
+def rms_norm_fp8_block128_bf16(
+    input: torch.Tensor,
+    weight: torch.Tensor,
+    eps: float = 1e-6,
+    *,
+    output: Optional[torch.Tensor] = None,
+    scale: Optional[torch.Tensor] = None,
+) -> tuple[torch.Tensor, torch.Tensor]:
+    if output is None or scale is None:
+        allocated_output, allocated_scale = _allocate(input)
+        output = allocated_output if output is None else output
+        scale = allocated_scale if scale is None else scale
+    ops.rms_norm_fp8_block128_bf16(input, weight, float(eps), output, scale)
+    return output, scale
+
+
+def residual_add_rms_norm_fp8_block128_bf16(
+    residual: torch.Tensor,
+    input: torch.Tensor,
+    weight: torch.Tensor,
+    eps: float = 1e-6,
+    *,
+    residual_out: Optional[torch.Tensor] = None,
+    output: Optional[torch.Tensor] = None,
+    scale: Optional[torch.Tensor] = None,
+) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
+    if residual_out is None:
+        residual_out = torch.empty_like(input)
+    if output is None or scale is None:
+        allocated_output, allocated_scale = _allocate(input)
+        output = allocated_output if output is None else output
+        scale = allocated_scale if scale is None else scale
+    ops.residual_add_rms_norm_fp8_block128_bf16(
+        residual, input, weight, float(eps), residual_out, output, scale
+    )
+    return residual_out, output, scale
+
+
+def gelu_tanh_fp8_block128_bf16(
+    input: torch.Tensor,
+    *,
+    output: Optional[torch.Tensor] = None,
+    scale: Optional[torch.Tensor] = None,
+) -> tuple[torch.Tensor, torch.Tensor]:
+    if output is None or scale is None:
+        allocated_output, allocated_scale = _allocate(input)
+        output = allocated_output if output is None else output
+        scale = allocated_scale if scale is None else scale
+    ops.gelu_tanh_fp8_block128_bf16(input, output, scale)
+    return output, scale
+
+
+def gelu_tanh_bias_fp8_block128_bf16(
+    input: torch.Tensor,
+    bias: torch.Tensor,
+    *,
+    output: Optional[torch.Tensor] = None,
+    scale: Optional[torch.Tensor] = None,
+) -> tuple[torch.Tensor, torch.Tensor]:
+    if output is None or scale is None:
+        allocated_output, allocated_scale = _allocate(input)
+        output = allocated_output if output is None else output
+        scale = allocated_scale if scale is None else scale
+    ops.gelu_tanh_bias_fp8_block128_bf16(input, bias, output, scale)
+    return output, scale
+
+
+def silu_mul_fp8_block128_bf16(
+    gate: torch.Tensor,
+    up: torch.Tensor,
+    *,
+    output: Optional[torch.Tensor] = None,
+    scale: Optional[torch.Tensor] = None,
+) -> tuple[torch.Tensor, torch.Tensor]:
+    if output is None or scale is None:
+        allocated_output, allocated_scale = _allocate(gate)
+        output = allocated_output if output is None else output
+        scale = allocated_scale if scale is None else scale
+    ops.silu_mul_fp8_block128_bf16(gate, up, output, scale)
+    return output, scale
+
+
+def silu_mul_merged_fp8_block128_bf16(
+    gate_up: torch.Tensor,
+    *,
+    output: Optional[torch.Tensor] = None,
+    scale: Optional[torch.Tensor] = None,
+) -> tuple[torch.Tensor, torch.Tensor]:
+    rows, merged_dim = gate_up.shape
+    dim = merged_dim // 2
+    if output is None:
+        output = torch.empty(
+            (rows, dim), device=gate_up.device, dtype=torch.float8_e4m3fn
+        )
+    if scale is None:
+        scale = torch.empty(
+            (rows, dim // 128), device=gate_up.device, dtype=torch.float32
+        )
+    ops.silu_mul_merged_fp8_block128_bf16(gate_up, output, scale)
+    return output, scale
+
+
+__all__ = [
+    "quantize_fp8_block128_bf16",
+    "layer_norm_fp8_block128_bf16",
+    "rms_norm_fp8_block128_bf16",
+    "residual_add_rms_norm_fp8_block128_bf16",
+    "gelu_tanh_fp8_block128_bf16",
+    "gelu_tanh_bias_fp8_block128_bf16",
+    "silu_mul_fp8_block128_bf16",
+    "silu_mul_merged_fp8_block128_bf16",
+]

build/torch211-cxx11-cu130-x86_64-linux/_blockwise_fp8_producers_cuda_7781728.abi3.so

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:e93a29699c24fd814affd52be775e1fd46f77c9b0a7b080e0ff68f132b03add2
+size 1965336

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

@@ -0,0 +1,9 @@
+import torch
+from . import _blockwise_fp8_producers_cuda_7781728
+ops = torch.ops._blockwise_fp8_producers_cuda_7781728
+
+def add_op_namespace_prefix(op_name: str):
+    """
+    Prefix op by namespace.
+    """
+    return f"_blockwise_fp8_producers_cuda_7781728::{op_name}"

build/torch211-cxx11-cu130-x86_64-linux/blockwise_fp8_producers/__init__.py

@@ -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

@@ -0,0 +1,22 @@
+{
+  "name": "blockwise-fp8-producers",
+  "id": "_blockwise_fp8_producers_cuda_7781728",
+  "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

@@ -0,0 +1,297 @@
+"""Blockwise FP8 producers for transformer and world-model regions."""
+
+from __future__ import annotations
+
+from typing import Optional
+
+import torch
+
+from ._ops import add_op_namespace_prefix, ops
+
+
+def _check_matrix(input: torch.Tensor, output: torch.Tensor, scale: torch.Tensor) -> None:
+    if (
+        input.dim() != 2
+        or input.shape[0] <= 0
+        or input.shape[1] <= 0
+        or input.shape[1] % 128 != 0
+        or output.shape != input.shape
+        or scale.shape != (input.shape[0], input.shape[1] // 128)
+    ):
+        raise RuntimeError(
+            "expected input/output (rows, dim) with dim a positive multiple "
+            "of 128 and scale (rows, dim / 128)"
+        )
+
+
+@torch.library.register_fake(add_op_namespace_prefix("quantize_fp8_block128_bf16"))
+def _quantize_fake(
+    input: torch.Tensor, output: torch.Tensor, scale: torch.Tensor
+) -> None:
+    _check_matrix(input, output, scale)
+    return None
+
+
+@torch.library.register_fake(add_op_namespace_prefix("layer_norm_fp8_block128_bf16"))
+def _layer_norm_fake(
+    input: torch.Tensor,
+    weight: torch.Tensor,
+    bias: torch.Tensor,
+    eps: float,
+    output: torch.Tensor,
+    scale: torch.Tensor,
+) -> None:
+    _check_matrix(input, output, scale)
+    if weight.shape != (input.shape[1],) or bias.shape != weight.shape:
+        raise RuntimeError("weight and bias must have shape (dim,)")
+    return None
+
+
+@torch.library.register_fake(add_op_namespace_prefix("rms_norm_fp8_block128_bf16"))
+def _rms_norm_fake(
+    input: torch.Tensor,
+    weight: torch.Tensor,
+    eps: float,
+    output: torch.Tensor,
+    scale: torch.Tensor,
+) -> None:
+    _check_matrix(input, output, scale)
+    if weight.shape != (input.shape[1],):
+        raise RuntimeError("weight must have shape (dim,)")
+    return None
+
+
+@torch.library.register_fake(
+    add_op_namespace_prefix("residual_add_rms_norm_fp8_block128_bf16")
+)
+def _residual_rms_norm_fake(
+    residual: torch.Tensor,
+    input: torch.Tensor,
+    weight: torch.Tensor,
+    eps: float,
+    residual_out: torch.Tensor,
+    output: torch.Tensor,
+    scale: torch.Tensor,
+) -> None:
+    _check_matrix(input, output, scale)
+    if (
+        residual.shape != input.shape
+        or residual_out.shape != input.shape
+        or weight.shape != (input.shape[1],)
+    ):
+        raise RuntimeError("residual/output must match input and weight must be (dim,)")
+    return None
+
+
+@torch.library.register_fake(add_op_namespace_prefix("gelu_tanh_fp8_block128_bf16"))
+def _gelu_fake(
+    input: torch.Tensor, output: torch.Tensor, scale: torch.Tensor
+) -> None:
+    _check_matrix(input, output, scale)
+    return None
+
+
+@torch.library.register_fake(
+    add_op_namespace_prefix("gelu_tanh_bias_fp8_block128_bf16")
+)
+def _gelu_bias_fake(
+    input: torch.Tensor,
+    bias: torch.Tensor,
+    output: torch.Tensor,
+    scale: torch.Tensor,
+) -> None:
+    _check_matrix(input, output, scale)
+    if bias.shape != (input.shape[1],):
+        raise RuntimeError("bias must have shape (dim,)")
+    return None
+
+
+@torch.library.register_fake(add_op_namespace_prefix("silu_mul_fp8_block128_bf16"))
+def _silu_mul_fake(
+    gate: torch.Tensor,
+    up: torch.Tensor,
+    output: torch.Tensor,
+    scale: torch.Tensor,
+) -> None:
+    _check_matrix(gate, output, scale)
+    if up.shape != gate.shape:
+        raise RuntimeError("up must match gate")
+    return None
+
+
+@torch.library.register_fake(
+    add_op_namespace_prefix("silu_mul_merged_fp8_block128_bf16")
+)
+def _silu_merged_fake(
+    gate_up: torch.Tensor, output: torch.Tensor, scale: torch.Tensor
+) -> None:
+    if (
+        gate_up.dim() != 2
+        or gate_up.shape[0] <= 0
+        or gate_up.shape[1] <= 0
+        or gate_up.shape[1] % 256 != 0
+        or output.shape != (gate_up.shape[0], gate_up.shape[1] // 2)
+        or scale.shape != (gate_up.shape[0], gate_up.shape[1] // 256)
+    ):
+        raise RuntimeError(
+            "gate_up must be (rows, 2 * dim), dim multiple of 128"
+        )
+    return None
+
+
+def _allocate(input: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor]:
+    return (
+        torch.empty_like(input, dtype=torch.float8_e4m3fn),
+        torch.empty(
+            (input.shape[0], input.shape[1] // 128),
+            device=input.device,
+            dtype=torch.float32,
+        ),
+    )
+
+
+def quantize_fp8_block128_bf16(
+    input: torch.Tensor,
+    *,
+    output: Optional[torch.Tensor] = None,
+    scale: Optional[torch.Tensor] = None,
+) -> tuple[torch.Tensor, torch.Tensor]:
+    if output is None or scale is None:
+        allocated_output, allocated_scale = _allocate(input)
+        output = allocated_output if output is None else output
+        scale = allocated_scale if scale is None else scale
+    ops.quantize_fp8_block128_bf16(input, output, scale)
+    return output, scale
+
+
+def layer_norm_fp8_block128_bf16(
+    input: torch.Tensor,
+    weight: torch.Tensor,
+    bias: torch.Tensor,
+    eps: float = 1e-6,
+    *,
+    output: Optional[torch.Tensor] = None,
+    scale: Optional[torch.Tensor] = None,
+) -> tuple[torch.Tensor, torch.Tensor]:
+    if output is None or scale is None:
+        allocated_output, allocated_scale = _allocate(input)
+        output = allocated_output if output is None else output
+        scale = allocated_scale if scale is None else scale
+    ops.layer_norm_fp8_block128_bf16(
+        input, weight, bias, float(eps), output, scale
+    )
+    return output, scale
+
+
+def rms_norm_fp8_block128_bf16(
+    input: torch.Tensor,
+    weight: torch.Tensor,
+    eps: float = 1e-6,
+    *,
+    output: Optional[torch.Tensor] = None,
+    scale: Optional[torch.Tensor] = None,
+) -> tuple[torch.Tensor, torch.Tensor]:
+    if output is None or scale is None:
+        allocated_output, allocated_scale = _allocate(input)
+        output = allocated_output if output is None else output
+        scale = allocated_scale if scale is None else scale
+    ops.rms_norm_fp8_block128_bf16(input, weight, float(eps), output, scale)
+    return output, scale
+
+
+def residual_add_rms_norm_fp8_block128_bf16(
+    residual: torch.Tensor,
+    input: torch.Tensor,
+    weight: torch.Tensor,
+    eps: float = 1e-6,
+    *,
+    residual_out: Optional[torch.Tensor] = None,
+    output: Optional[torch.Tensor] = None,
+    scale: Optional[torch.Tensor] = None,
+) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
+    if residual_out is None:
+        residual_out = torch.empty_like(input)
+    if output is None or scale is None:
+        allocated_output, allocated_scale = _allocate(input)
+        output = allocated_output if output is None else output
+        scale = allocated_scale if scale is None else scale
+    ops.residual_add_rms_norm_fp8_block128_bf16(
+        residual, input, weight, float(eps), residual_out, output, scale
+    )
+    return residual_out, output, scale
+
+
+def gelu_tanh_fp8_block128_bf16(
+    input: torch.Tensor,
+    *,
+    output: Optional[torch.Tensor] = None,
+    scale: Optional[torch.Tensor] = None,
+) -> tuple[torch.Tensor, torch.Tensor]:
+    if output is None or scale is None:
+        allocated_output, allocated_scale = _allocate(input)
+        output = allocated_output if output is None else output
+        scale = allocated_scale if scale is None else scale
+    ops.gelu_tanh_fp8_block128_bf16(input, output, scale)
+    return output, scale
+
+
+def gelu_tanh_bias_fp8_block128_bf16(
+    input: torch.Tensor,
+    bias: torch.Tensor,
+    *,
+    output: Optional[torch.Tensor] = None,
+    scale: Optional[torch.Tensor] = None,
+) -> tuple[torch.Tensor, torch.Tensor]:
+    if output is None or scale is None:
+        allocated_output, allocated_scale = _allocate(input)
+        output = allocated_output if output is None else output
+        scale = allocated_scale if scale is None else scale
+    ops.gelu_tanh_bias_fp8_block128_bf16(input, bias, output, scale)
+    return output, scale
+
+
+def silu_mul_fp8_block128_bf16(
+    gate: torch.Tensor,
+    up: torch.Tensor,
+    *,
+    output: Optional[torch.Tensor] = None,
+    scale: Optional[torch.Tensor] = None,
+) -> tuple[torch.Tensor, torch.Tensor]:
+    if output is None or scale is None:
+        allocated_output, allocated_scale = _allocate(gate)
+        output = allocated_output if output is None else output
+        scale = allocated_scale if scale is None else scale
+    ops.silu_mul_fp8_block128_bf16(gate, up, output, scale)
+    return output, scale
+
+
+def silu_mul_merged_fp8_block128_bf16(
+    gate_up: torch.Tensor,
+    *,
+    output: Optional[torch.Tensor] = None,
+    scale: Optional[torch.Tensor] = None,
+) -> tuple[torch.Tensor, torch.Tensor]:
+    rows, merged_dim = gate_up.shape
+    dim = merged_dim // 2
+    if output is None:
+        output = torch.empty(
+            (rows, dim), device=gate_up.device, dtype=torch.float8_e4m3fn
+        )
+    if scale is None:
+        scale = torch.empty(
+            (rows, dim // 128), device=gate_up.device, dtype=torch.float32
+        )
+    ops.silu_mul_merged_fp8_block128_bf16(gate_up, output, scale)
+    return output, scale
+
+
+__all__ = [
+    "quantize_fp8_block128_bf16",
+    "layer_norm_fp8_block128_bf16",
+    "rms_norm_fp8_block128_bf16",
+    "residual_add_rms_norm_fp8_block128_bf16",
+    "gelu_tanh_fp8_block128_bf16",
+    "gelu_tanh_bias_fp8_block128_bf16",
+    "silu_mul_fp8_block128_bf16",
+    "silu_mul_merged_fp8_block128_bf16",
+]

build/torch212-cxx11-cu130-x86_64-linux/_blockwise_fp8_producers_cuda_7781728.abi3.so

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:7c7611d3c85039b78b8d4270343b02a026e8f069068fb38a0c7caac2a1a384ce
+size 1975904

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

@@ -0,0 +1,9 @@
+import torch
+from . import _blockwise_fp8_producers_cuda_7781728
+ops = torch.ops._blockwise_fp8_producers_cuda_7781728
+
+def add_op_namespace_prefix(op_name: str):
+    """
+    Prefix op by namespace.
+    """
+    return f"_blockwise_fp8_producers_cuda_7781728::{op_name}"

build/torch212-cxx11-cu130-x86_64-linux/blockwise_fp8_producers/__init__.py

@@ -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

@@ -0,0 +1,22 @@
+{
+  "name": "blockwise-fp8-producers",
+  "id": "_blockwise_fp8_producers_cuda_7781728",
+  "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

@@ -0,0 +1,297 @@
+"""Blockwise FP8 producers for transformer and world-model regions."""
+
+from __future__ import annotations
+
+from typing import Optional
+
+import torch
+
+from ._ops import add_op_namespace_prefix, ops
+
+
+def _check_matrix(input: torch.Tensor, output: torch.Tensor, scale: torch.Tensor) -> None:
+    if (
+        input.dim() != 2
+        or input.shape[0] <= 0
+        or input.shape[1] <= 0
+        or input.shape[1] % 128 != 0
+        or output.shape != input.shape
+        or scale.shape != (input.shape[0], input.shape[1] // 128)
+    ):
+        raise RuntimeError(
+            "expected input/output (rows, dim) with dim a positive multiple "
+            "of 128 and scale (rows, dim / 128)"
+        )
+
+
+@torch.library.register_fake(add_op_namespace_prefix("quantize_fp8_block128_bf16"))
+def _quantize_fake(
+    input: torch.Tensor, output: torch.Tensor, scale: torch.Tensor
+) -> None:
+    _check_matrix(input, output, scale)
+    return None
+
+
+@torch.library.register_fake(add_op_namespace_prefix("layer_norm_fp8_block128_bf16"))
+def _layer_norm_fake(
+    input: torch.Tensor,
+    weight: torch.Tensor,
+    bias: torch.Tensor,
+    eps: float,
+    output: torch.Tensor,
+    scale: torch.Tensor,
+) -> None:
+    _check_matrix(input, output, scale)
+    if weight.shape != (input.shape[1],) or bias.shape != weight.shape:
+        raise RuntimeError("weight and bias must have shape (dim,)")
+    return None
+
+
+@torch.library.register_fake(add_op_namespace_prefix("rms_norm_fp8_block128_bf16"))
+def _rms_norm_fake(
+    input: torch.Tensor,
+    weight: torch.Tensor,
+    eps: float,
+    output: torch.Tensor,
+    scale: torch.Tensor,
+) -> None:
+    _check_matrix(input, output, scale)
+    if weight.shape != (input.shape[1],):
+        raise RuntimeError("weight must have shape (dim,)")
+    return None
+
+
+@torch.library.register_fake(
+    add_op_namespace_prefix("residual_add_rms_norm_fp8_block128_bf16")
+)
+def _residual_rms_norm_fake(
+    residual: torch.Tensor,
+    input: torch.Tensor,
+    weight: torch.Tensor,
+    eps: float,
+    residual_out: torch.Tensor,
+    output: torch.Tensor,
+    scale: torch.Tensor,
+) -> None:
+    _check_matrix(input, output, scale)
+    if (
+        residual.shape != input.shape
+        or residual_out.shape != input.shape
+        or weight.shape != (input.shape[1],)
+    ):
+        raise RuntimeError("residual/output must match input and weight must be (dim,)")
+    return None
+
+
+@torch.library.register_fake(add_op_namespace_prefix("gelu_tanh_fp8_block128_bf16"))
+def _gelu_fake(
+    input: torch.Tensor, output: torch.Tensor, scale: torch.Tensor
+) -> None:
+    _check_matrix(input, output, scale)
+    return None
+
+
+@torch.library.register_fake(
+    add_op_namespace_prefix("gelu_tanh_bias_fp8_block128_bf16")
+)
+def _gelu_bias_fake(
+    input: torch.Tensor,
+    bias: torch.Tensor,
+    output: torch.Tensor,
+    scale: torch.Tensor,
+) -> None:
+    _check_matrix(input, output, scale)
+    if bias.shape != (input.shape[1],):
+        raise RuntimeError("bias must have shape (dim,)")
+    return None
+
+
+@torch.library.register_fake(add_op_namespace_prefix("silu_mul_fp8_block128_bf16"))
+def _silu_mul_fake(
+    gate: torch.Tensor,
+    up: torch.Tensor,
+    output: torch.Tensor,
+    scale: torch.Tensor,
+) -> None:
+    _check_matrix(gate, output, scale)
+    if up.shape != gate.shape:
+        raise RuntimeError("up must match gate")
+    return None
+
+
+@torch.library.register_fake(
+    add_op_namespace_prefix("silu_mul_merged_fp8_block128_bf16")
+)
+def _silu_merged_fake(
+    gate_up: torch.Tensor, output: torch.Tensor, scale: torch.Tensor
+) -> None:
+    if (
+        gate_up.dim() != 2
+        or gate_up.shape[0] <= 0
+        or gate_up.shape[1] <= 0
+        or gate_up.shape[1] % 256 != 0
+        or output.shape != (gate_up.shape[0], gate_up.shape[1] // 2)
+        or scale.shape != (gate_up.shape[0], gate_up.shape[1] // 256)
+    ):
+        raise RuntimeError(
+            "gate_up must be (rows, 2 * dim), dim multiple of 128"
+        )
+    return None
+
+
+def _allocate(input: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor]:
+    return (
+        torch.empty_like(input, dtype=torch.float8_e4m3fn),
+        torch.empty(
+            (input.shape[0], input.shape[1] // 128),
+            device=input.device,
+            dtype=torch.float32,
+        ),
+    )
+
+
+def quantize_fp8_block128_bf16(
+    input: torch.Tensor,
+    *,
+    output: Optional[torch.Tensor] = None,
+    scale: Optional[torch.Tensor] = None,
+) -> tuple[torch.Tensor, torch.Tensor]:
+    if output is None or scale is None:
+        allocated_output, allocated_scale = _allocate(input)
+        output = allocated_output if output is None else output
+        scale = allocated_scale if scale is None else scale
+    ops.quantize_fp8_block128_bf16(input, output, scale)
+    return output, scale
+
+
+def layer_norm_fp8_block128_bf16(
+    input: torch.Tensor,
+    weight: torch.Tensor,
+    bias: torch.Tensor,
+    eps: float = 1e-6,
+    *,
+    output: Optional[torch.Tensor] = None,
+    scale: Optional[torch.Tensor] = None,
+) -> tuple[torch.Tensor, torch.Tensor]:
+    if output is None or scale is None:
+        allocated_output, allocated_scale = _allocate(input)
+        output = allocated_output if output is None else output
+        scale = allocated_scale if scale is None else scale
+    ops.layer_norm_fp8_block128_bf16(
+        input, weight, bias, float(eps), output, scale
+    )
+    return output, scale
+
+
+def rms_norm_fp8_block128_bf16(
+    input: torch.Tensor,
+    weight: torch.Tensor,
+    eps: float = 1e-6,
+    *,
+    output: Optional[torch.Tensor] = None,
+    scale: Optional[torch.Tensor] = None,
+) -> tuple[torch.Tensor, torch.Tensor]:
+    if output is None or scale is None:
+        allocated_output, allocated_scale = _allocate(input)
+        output = allocated_output if output is None else output
+        scale = allocated_scale if scale is None else scale
+    ops.rms_norm_fp8_block128_bf16(input, weight, float(eps), output, scale)
+    return output, scale
+
+
+def residual_add_rms_norm_fp8_block128_bf16(
+    residual: torch.Tensor,
+    input: torch.Tensor,
+    weight: torch.Tensor,
+    eps: float = 1e-6,
+    *,
+    residual_out: Optional[torch.Tensor] = None,
+    output: Optional[torch.Tensor] = None,
+    scale: Optional[torch.Tensor] = None,
+) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
+    if residual_out is None:
+        residual_out = torch.empty_like(input)
+    if output is None or scale is None:
+        allocated_output, allocated_scale = _allocate(input)
+        output = allocated_output if output is None else output
+        scale = allocated_scale if scale is None else scale
+    ops.residual_add_rms_norm_fp8_block128_bf16(
+        residual, input, weight, float(eps), residual_out, output, scale
+    )
+    return residual_out, output, scale
+
+
+def gelu_tanh_fp8_block128_bf16(
+    input: torch.Tensor,
+    *,
+    output: Optional[torch.Tensor] = None,
+    scale: Optional[torch.Tensor] = None,
+) -> tuple[torch.Tensor, torch.Tensor]:
+    if output is None or scale is None:
+        allocated_output, allocated_scale = _allocate(input)
+        output = allocated_output if output is None else output
+        scale = allocated_scale if scale is None else scale
+    ops.gelu_tanh_fp8_block128_bf16(input, output, scale)
+    return output, scale
+
+
+def gelu_tanh_bias_fp8_block128_bf16(
+    input: torch.Tensor,
+    bias: torch.Tensor,
+    *,
+    output: Optional[torch.Tensor] = None,
+    scale: Optional[torch.Tensor] = None,
+) -> tuple[torch.Tensor, torch.Tensor]:
+    if output is None or scale is None:
+        allocated_output, allocated_scale = _allocate(input)
+        output = allocated_output if output is None else output
+        scale = allocated_scale if scale is None else scale
+    ops.gelu_tanh_bias_fp8_block128_bf16(input, bias, output, scale)
+    return output, scale
+
+
+def silu_mul_fp8_block128_bf16(
+    gate: torch.Tensor,
+    up: torch.Tensor,
+    *,
+    output: Optional[torch.Tensor] = None,
+    scale: Optional[torch.Tensor] = None,
+) -> tuple[torch.Tensor, torch.Tensor]:
+    if output is None or scale is None:
+        allocated_output, allocated_scale = _allocate(gate)
+        output = allocated_output if output is None else output
+        scale = allocated_scale if scale is None else scale
+    ops.silu_mul_fp8_block128_bf16(gate, up, output, scale)
+    return output, scale
+
+
+def silu_mul_merged_fp8_block128_bf16(
+    gate_up: torch.Tensor,
+    *,
+    output: Optional[torch.Tensor] = None,
+    scale: Optional[torch.Tensor] = None,
+) -> tuple[torch.Tensor, torch.Tensor]:
+    rows, merged_dim = gate_up.shape
+    dim = merged_dim // 2
+    if output is None:
+        output = torch.empty(
+            (rows, dim), device=gate_up.device, dtype=torch.float8_e4m3fn
+        )
+    if scale is None:
+        scale = torch.empty(
+            (rows, dim // 128), device=gate_up.device, dtype=torch.float32
+        )
+    ops.silu_mul_merged_fp8_block128_bf16(gate_up, output, scale)
+    return output, scale
+
+
+__all__ = [
+    "quantize_fp8_block128_bf16",
+    "layer_norm_fp8_block128_bf16",
+    "rms_norm_fp8_block128_bf16",
+    "residual_add_rms_norm_fp8_block128_bf16",
+    "gelu_tanh_fp8_block128_bf16",
+    "gelu_tanh_bias_fp8_block128_bf16",
+    "silu_mul_fp8_block128_bf16",
+    "silu_mul_merged_fp8_block128_bf16",
+]

build/torch212-cxx11-cu132-x86_64-linux/_blockwise_fp8_producers_cuda_7781728.abi3.so

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:2c97dac88cca54b24510190aa866b5495d96ab666ebae5449d3c40dbf2c18584
+size 1980000

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

@@ -0,0 +1,9 @@
+import torch
+from . import _blockwise_fp8_producers_cuda_7781728
+ops = torch.ops._blockwise_fp8_producers_cuda_7781728
+
+def add_op_namespace_prefix(op_name: str):
+    """
+    Prefix op by namespace.
+    """
+    return f"_blockwise_fp8_producers_cuda_7781728::{op_name}"

build/torch212-cxx11-cu132-x86_64-linux/blockwise_fp8_producers/__init__.py

@@ -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

@@ -0,0 +1,22 @@
+{
+  "name": "blockwise-fp8-producers",
+  "id": "_blockwise_fp8_producers_cuda_7781728",
+  "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/torch213-cxx11-cu130-aarch64-linux/__init__.py

@@ -0,0 +1,297 @@
+"""Blockwise FP8 producers for transformer and world-model regions."""
+
+from __future__ import annotations
+
+from typing import Optional
+
+import torch
+
+from ._ops import add_op_namespace_prefix, ops
+
+
+def _check_matrix(input: torch.Tensor, output: torch.Tensor, scale: torch.Tensor) -> None:
+    if (
+        input.dim() != 2
+        or input.shape[0] <= 0
+        or input.shape[1] <= 0
+        or input.shape[1] % 128 != 0
+        or output.shape != input.shape
+        or scale.shape != (input.shape[0], input.shape[1] // 128)
+    ):
+        raise RuntimeError(
+            "expected input/output (rows, dim) with dim a positive multiple "
+            "of 128 and scale (rows, dim / 128)"
+        )
+
+
+@torch.library.register_fake(add_op_namespace_prefix("quantize_fp8_block128_bf16"))
+def _quantize_fake(
+    input: torch.Tensor, output: torch.Tensor, scale: torch.Tensor
+) -> None:
+    _check_matrix(input, output, scale)
+    return None
+
+
+@torch.library.register_fake(add_op_namespace_prefix("layer_norm_fp8_block128_bf16"))
+def _layer_norm_fake(
+    input: torch.Tensor,
+    weight: torch.Tensor,
+    bias: torch.Tensor,
+    eps: float,
+    output: torch.Tensor,
+    scale: torch.Tensor,
+) -> None:
+    _check_matrix(input, output, scale)
+    if weight.shape != (input.shape[1],) or bias.shape != weight.shape:
+        raise RuntimeError("weight and bias must have shape (dim,)")
+    return None
+
+
+@torch.library.register_fake(add_op_namespace_prefix("rms_norm_fp8_block128_bf16"))
+def _rms_norm_fake(
+    input: torch.Tensor,
+    weight: torch.Tensor,
+    eps: float,
+    output: torch.Tensor,
+    scale: torch.Tensor,
+) -> None:
+    _check_matrix(input, output, scale)
+    if weight.shape != (input.shape[1],):
+        raise RuntimeError("weight must have shape (dim,)")
+    return None
+
+
+@torch.library.register_fake(
+    add_op_namespace_prefix("residual_add_rms_norm_fp8_block128_bf16")
+)
+def _residual_rms_norm_fake(
+    residual: torch.Tensor,
+    input: torch.Tensor,
+    weight: torch.Tensor,
+    eps: float,
+    residual_out: torch.Tensor,
+    output: torch.Tensor,
+    scale: torch.Tensor,
+) -> None:
+    _check_matrix(input, output, scale)
+    if (
+        residual.shape != input.shape
+        or residual_out.shape != input.shape
+        or weight.shape != (input.shape[1],)
+    ):
+        raise RuntimeError("residual/output must match input and weight must be (dim,)")
+    return None
+
+
+@torch.library.register_fake(add_op_namespace_prefix("gelu_tanh_fp8_block128_bf16"))
+def _gelu_fake(
+    input: torch.Tensor, output: torch.Tensor, scale: torch.Tensor
+) -> None:
+    _check_matrix(input, output, scale)
+    return None
+
+
+@torch.library.register_fake(
+    add_op_namespace_prefix("gelu_tanh_bias_fp8_block128_bf16")
+)
+def _gelu_bias_fake(
+    input: torch.Tensor,
+    bias: torch.Tensor,
+    output: torch.Tensor,
+    scale: torch.Tensor,
+) -> None:
+    _check_matrix(input, output, scale)
+    if bias.shape != (input.shape[1],):
+        raise RuntimeError("bias must have shape (dim,)")
+    return None
+
+
+@torch.library.register_fake(add_op_namespace_prefix("silu_mul_fp8_block128_bf16"))
+def _silu_mul_fake(
+    gate: torch.Tensor,
+    up: torch.Tensor,
+    output: torch.Tensor,
+    scale: torch.Tensor,
+) -> None:
+    _check_matrix(gate, output, scale)
+    if up.shape != gate.shape:
+        raise RuntimeError("up must match gate")
+    return None
+
+
+@torch.library.register_fake(
+    add_op_namespace_prefix("silu_mul_merged_fp8_block128_bf16")
+)
+def _silu_merged_fake(
+    gate_up: torch.Tensor, output: torch.Tensor, scale: torch.Tensor
+) -> None:
+    if (
+        gate_up.dim() != 2
+        or gate_up.shape[0] <= 0
+        or gate_up.shape[1] <= 0
+        or gate_up.shape[1] % 256 != 0
+        or output.shape != (gate_up.shape[0], gate_up.shape[1] // 2)
+        or scale.shape != (gate_up.shape[0], gate_up.shape[1] // 256)
+    ):
+        raise RuntimeError(
+            "gate_up must be (rows, 2 * dim), dim multiple of 128"
+        )
+    return None
+
+
+def _allocate(input: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor]:
+    return (
+        torch.empty_like(input, dtype=torch.float8_e4m3fn),
+        torch.empty(
+            (input.shape[0], input.shape[1] // 128),
+            device=input.device,
+            dtype=torch.float32,
+        ),
+    )
+
+
+def quantize_fp8_block128_bf16(
+    input: torch.Tensor,
+    *,
+    output: Optional[torch.Tensor] = None,
+    scale: Optional[torch.Tensor] = None,
+) -> tuple[torch.Tensor, torch.Tensor]:
+    if output is None or scale is None:
+        allocated_output, allocated_scale = _allocate(input)
+        output = allocated_output if output is None else output
+        scale = allocated_scale if scale is None else scale
+    ops.quantize_fp8_block128_bf16(input, output, scale)
+    return output, scale
+
+
+def layer_norm_fp8_block128_bf16(
+    input: torch.Tensor,
+    weight: torch.Tensor,
+    bias: torch.Tensor,
+    eps: float = 1e-6,
+    *,
+    output: Optional[torch.Tensor] = None,
+    scale: Optional[torch.Tensor] = None,
+) -> tuple[torch.Tensor, torch.Tensor]:
+    if output is None or scale is None:
+        allocated_output, allocated_scale = _allocate(input)
+        output = allocated_output if output is None else output
+        scale = allocated_scale if scale is None else scale
+    ops.layer_norm_fp8_block128_bf16(
+        input, weight, bias, float(eps), output, scale
+    )
+    return output, scale
+
+
+def rms_norm_fp8_block128_bf16(
+    input: torch.Tensor,
+    weight: torch.Tensor,
+    eps: float = 1e-6,
+    *,
+    output: Optional[torch.Tensor] = None,
+    scale: Optional[torch.Tensor] = None,
+) -> tuple[torch.Tensor, torch.Tensor]:
+    if output is None or scale is None:
+        allocated_output, allocated_scale = _allocate(input)
+        output = allocated_output if output is None else output
+        scale = allocated_scale if scale is None else scale
+    ops.rms_norm_fp8_block128_bf16(input, weight, float(eps), output, scale)
+    return output, scale
+
+
+def residual_add_rms_norm_fp8_block128_bf16(
+    residual: torch.Tensor,
+    input: torch.Tensor,
+    weight: torch.Tensor,
+    eps: float = 1e-6,
+    *,
+    residual_out: Optional[torch.Tensor] = None,
+    output: Optional[torch.Tensor] = None,
+    scale: Optional[torch.Tensor] = None,
+) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
+    if residual_out is None:
+        residual_out = torch.empty_like(input)
+    if output is None or scale is None:
+        allocated_output, allocated_scale = _allocate(input)
+        output = allocated_output if output is None else output
+        scale = allocated_scale if scale is None else scale
+    ops.residual_add_rms_norm_fp8_block128_bf16(
+        residual, input, weight, float(eps), residual_out, output, scale
+    )
+    return residual_out, output, scale
+
+
+def gelu_tanh_fp8_block128_bf16(
+    input: torch.Tensor,
+    *,
+    output: Optional[torch.Tensor] = None,
+    scale: Optional[torch.Tensor] = None,
+) -> tuple[torch.Tensor, torch.Tensor]:
+    if output is None or scale is None:
+        allocated_output, allocated_scale = _allocate(input)
+        output = allocated_output if output is None else output
+        scale = allocated_scale if scale is None else scale
+    ops.gelu_tanh_fp8_block128_bf16(input, output, scale)
+    return output, scale
+
+
+def gelu_tanh_bias_fp8_block128_bf16(
+    input: torch.Tensor,
+    bias: torch.Tensor,
+    *,
+    output: Optional[torch.Tensor] = None,
+    scale: Optional[torch.Tensor] = None,
+) -> tuple[torch.Tensor, torch.Tensor]:
+    if output is None or scale is None:
+        allocated_output, allocated_scale = _allocate(input)
+        output = allocated_output if output is None else output
+        scale = allocated_scale if scale is None else scale
+    ops.gelu_tanh_bias_fp8_block128_bf16(input, bias, output, scale)
+    return output, scale
+
+
+def silu_mul_fp8_block128_bf16(
+    gate: torch.Tensor,
+    up: torch.Tensor,
+    *,
+    output: Optional[torch.Tensor] = None,
+    scale: Optional[torch.Tensor] = None,
+) -> tuple[torch.Tensor, torch.Tensor]:
+    if output is None or scale is None:
+        allocated_output, allocated_scale = _allocate(gate)
+        output = allocated_output if output is None else output
+        scale = allocated_scale if scale is None else scale
+    ops.silu_mul_fp8_block128_bf16(gate, up, output, scale)
+    return output, scale
+
+
+def silu_mul_merged_fp8_block128_bf16(
+    gate_up: torch.Tensor,
+    *,
+    output: Optional[torch.Tensor] = None,
+    scale: Optional[torch.Tensor] = None,
+) -> tuple[torch.Tensor, torch.Tensor]:
+    rows, merged_dim = gate_up.shape
+    dim = merged_dim // 2
+    if output is None:
+        output = torch.empty(
+            (rows, dim), device=gate_up.device, dtype=torch.float8_e4m3fn
+        )
+    if scale is None:
+        scale = torch.empty(
+            (rows, dim // 128), device=gate_up.device, dtype=torch.float32
+        )
+    ops.silu_mul_merged_fp8_block128_bf16(gate_up, output, scale)
+    return output, scale
+
+
+__all__ = [
+    "quantize_fp8_block128_bf16",
+    "layer_norm_fp8_block128_bf16",
+    "rms_norm_fp8_block128_bf16",
+    "residual_add_rms_norm_fp8_block128_bf16",
+    "gelu_tanh_fp8_block128_bf16",
+    "gelu_tanh_bias_fp8_block128_bf16",
+    "silu_mul_fp8_block128_bf16",
+    "silu_mul_merged_fp8_block128_bf16",
+]

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

@@ -0,0 +1,6 @@
+import torch
+from . import blockwise_fp8_producers_source_test
+ops = torch.ops.blockwise_fp8_producers_source_test
+
+def add_op_namespace_prefix(op_name: str):
+    return f"blockwise_fp8_producers_source_test::{op_name}"

build/torch213-cxx11-cu130-aarch64-linux/blockwise_fp8_producers/__init__.py

@@ -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/blockwise_fp8_producers_source_test.abi3.so

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:482bf87ff6a9a2a9d92b8608e41f8255e151853f0de46dc95594801236cfe335
+size 458120

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

@@ -0,0 +1,32 @@
+{
+  "name": "blockwise-fp8-producers",
+  "id": "blockwise_fp8_producers_source_test",
+  "version": 1,
+  "license": "Apache-2.0",
+  "python-depends": [],
+  "backend": {
+    "type": "cuda",
+    "archs": [
+      "11.0"
+    ]
+  },
+  "digest": {
+    "algorithm": "sha256",
+    "files": {
+      "__init__.py": "jTa/yZ4lH7PUVKeENwdd7iZgdhrc9AOOJ8+L2usA23M=",
+      "blockwise_fp8_producers_source_test.abi3.so": "SCv4f/apoqnZK4YI5B+CVeFRhT8N5G3JVZSAEjbP4zU=",
+      "_ops.py": "gWnlQMRiABwvEGwVmY0l2A/6y1/DB/ruVb0A4WOgziE=",
+      "blockwise_fp8_producers/__init__.py": "v6p5XMfQzddhi1fLSAw4HX9CyS0rQsidvu9VsT01xi4="
+    }
+  },
+  "provenance": {
+    "kernel": {
+      "sha": "456d297",
+      "dirty": false
+    },
+    "validation": {
+      "torch": "2.13.0+cu130",
+      "cuda": "13.0"
+    }
+  }
+}