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/weight-only-ffn
-
-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/weight-only-ffn
-
-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,460 @@
+#!/usr/bin/env python3
+"""Tile/variant benchmark for the production M<=4 weight-only domain."""
+
+from __future__ import annotations
+
+import argparse
+import importlib
+import json
+import os
+import statistics
+import sys
+from pathlib import Path
+
+import torch
+import torch.nn.functional as F
+
+
+SHAPES = {
+    "llm_m1": (1, 4096, 11008, 4096),
+    "llm_m2": (2, 4096, 11008, 4096),
+    "llm_m3": (3, 4096, 11008, 4096),
+    "llm_m4": (4, 4096, 11008, 4096),
+    "vla_m1": (1, 1024, 4096, 1024),
+    "vla_m2": (2, 1024, 4096, 1024),
+    "vla_m4": (4, 1024, 4096, 1024),
+    "vision_m1": (1, 1536, 6144, 1536),
+    "vision_m2": (2, 1536, 6144, 1536),
+    "vision_m4": (4, 1536, 6144, 1536),
+}
+
+LINEAR_SHAPES = {
+    "llm_square_m1": (1, 4096, 4096),
+    "llm_wide_m1": (1, 4096, 11008),
+    "vla_wide_m1": (1, 1024, 4096),
+    "vision_wide_m1": (1, 1536, 6144),
+    "llm_square_m2": (2, 4096, 4096),
+    "llm_wide_m2": (2, 4096, 11008),
+    "vla_wide_m2": (2, 1024, 4096),
+    "vision_wide_m2": (2, 1536, 6144),
+}
+
+
+def bench(fn, warmup: int, iterations: int, repeats: int = 3) -> float:
+    for _ in range(warmup):
+        fn()
+    torch.cuda.synchronize()
+    samples = []
+    for _ in range(repeats):
+        start = torch.cuda.Event(enable_timing=True)
+        end = torch.cuda.Event(enable_timing=True)
+        start.record()
+        for _ in range(iterations):
+            fn()
+        end.record()
+        torch.cuda.synchronize()
+        samples.append(float(start.elapsed_time(end) * 1000.0 / iterations))
+    return statistics.median(samples)
+
+
+def sfb_bytes(rows: int, cols: int) -> int:
+    return ((rows + 127) // 128) * (((cols // 16) + 3) // 4) * 512
+
+
+class SourceModule:
+    def __init__(self, ops) -> None:
+        self.ops = ops
+
+    def quantize_w4_weight_bf16(self, weight):
+        n, k = weight.shape
+        packed = torch.empty((n, k // 2), device="cuda", dtype=torch.uint8)
+        scale = torch.empty((sfb_bytes(n, k),), device="cuda", dtype=torch.uint8)
+        self.ops.quantize_w4_weight_bf16(weight, packed, scale)
+        return packed, scale
+
+    def dequantize_w4_weight_bf16(self, packed, scale, *, cols):
+        out = torch.empty((packed.shape[0], cols), device="cuda", dtype=torch.bfloat16)
+        self.ops.dequantize_w4_weight_bf16(packed, scale, out)
+        return out
+
+    def quantize_w8_weight_bf16(self, weight):
+        packed = torch.empty_like(weight, dtype=torch.int8)
+        scale = torch.empty((weight.shape[0],), device="cuda", dtype=torch.float32)
+        self.ops.quantize_w8_weight_bf16(weight, packed, scale)
+        return packed, scale
+
+    def dequantize_w8_weight_bf16(self, packed, scale):
+        out = torch.empty_like(packed, dtype=torch.bfloat16)
+        self.ops.dequantize_w8_weight_bf16(packed, scale, out)
+        return out
+
+    def w4a16_linear_bf16(self, x, weight, scale, *, variant=0, out=None):
+        if out is None:
+            out = torch.empty(
+                (x.shape[0], weight.shape[0]), device=x.device,
+                dtype=torch.bfloat16,
+            )
+        self.ops.w4a16_linear_bf16(
+            x, weight, scale, 1.0, variant, out,
+        )
+        return out
+
+    def w8a16_linear_bf16(self, x, weight, scale, *, variant=0, out=None):
+        if out is None:
+            out = torch.empty(
+                (x.shape[0], weight.shape[0]), device=x.device,
+                dtype=torch.bfloat16,
+            )
+        self.ops.w8a16_linear_bf16(x, weight, scale, variant, out)
+        return out
+
+    def _gated(self, bits, x, gu_w, gu_s, dn_w, dn_s, *, gelu,
+               gate_up_bias, down_bias, variant, workspace, out):
+        gu, hidden = workspace
+        if bits == 4:
+            self.ops.w4a16_gated_ffn_bf16(
+                x, gu_w, gu_s, dn_w, dn_s, gate_up_bias, down_bias,
+                gelu, 1.0, 1.0, variant, gu, hidden, out,
+            )
+        else:
+            self.ops.w8a16_gated_ffn_bf16(
+                x, gu_w, gu_s, dn_w, dn_s, gate_up_bias, down_bias,
+                gelu, variant, gu, hidden, out,
+            )
+        return out
+
+    def w4a16_swiglu_ffn_bf16(self, *args, **kwargs):
+        return self._gated(4, *args, gelu=False, **kwargs)
+
+    def w4a16_geglu_ffn_bf16(self, *args, **kwargs):
+        return self._gated(4, *args, gelu=True, **kwargs)
+
+    def w8a16_swiglu_ffn_bf16(self, *args, **kwargs):
+        return self._gated(8, *args, gelu=False, **kwargs)
+
+    def w8a16_geglu_ffn_bf16(self, *args, **kwargs):
+        return self._gated(8, *args, gelu=True, **kwargs)
+
+    def _gelu(self, bits, x, up_w, up_s, dn_w, dn_s, *, up_bias,
+              down_bias, variant, workspace, out):
+        up, hidden = workspace
+        if bits == 4:
+            self.ops.w4a16_gelu_ffn_bf16(
+                x, up_w, up_s, dn_w, dn_s, up_bias, down_bias,
+                1.0, 1.0, variant, up, hidden, out,
+            )
+        else:
+            self.ops.w8a16_gelu_ffn_bf16(
+                x, up_w, up_s, dn_w, dn_s, up_bias, down_bias,
+                variant, up, hidden, out,
+            )
+        return out
+
+    def w4a16_gelu_ffn_bf16(self, *args, **kwargs):
+        return self._gelu(4, *args, **kwargs)
+
+    def w8a16_gelu_ffn_bf16(self, *args, **kwargs):
+        return self._gelu(8, *args, **kwargs)
+
+
+def load_source_module():
+    from torch.utils.cpp_extension import load
+
+    root = Path(__file__).resolve().parents[1]
+    registration = root.parent.parent / "kernels" / "kernel-builder" / "src" / "pyproject" / "templates" / "torch"
+    major, minor = torch.cuda.get_device_capability(0)
+    if major == 11 and minor == 0:
+        arch = "11.0a"
+    elif major == 12 and minor == 1:
+        arch = "12.1"
+    elif major >= 12:
+        arch = "12.0a"
+    else:
+        raise RuntimeError("source benchmark requires Blackwell SM110/SM120/SM121")
+    os.environ.setdefault("TORCH_CUDA_ARCH_LIST", arch)
+    namespace = "weight_only_ffn_benchmark_source"
+    load(
+        name=namespace,
+        sources=[str(root / path) for path in [
+            "torch-ext/torch_binding.cpp", "csrc/w4_weight_only.cu",
+            "csrc/w4a16_gemm_sm120.cu", "csrc/w4a16_matvec_sm120.cu",
+            "csrc/w8_weight_only.cu", "csrc/ffn_epilogues.cu",
+        ]],
+        extra_include_paths=[str(root / "csrc"), str(registration)],
+        extra_cflags=["-O3", "-std=c++17", "-DCUDA_KERNEL"],
+        extra_cuda_cflags=["-O3", "-std=c++17", "-DCUDA_KERNEL", "--use_fast_math"],
+        is_python_module=False,
+        verbose=False,
+    )
+    return SourceModule(getattr(torch.ops, namespace))
+
+
+def load_module(backend: str, artifact: str | None):
+    if backend == "source":
+        return load_source_module()
+    if artifact:
+        sys.path.insert(0, artifact)
+    try:
+        return importlib.import_module("weight_only_ffn")
+    finally:
+        if artifact:
+            sys.path.remove(artifact)
+
+
+def quantize(module, bits: int, weight: torch.Tensor):
+    if bits == 4:
+        packed, scale = module.quantize_w4_weight_bf16(weight)
+        dequant = module.dequantize_w4_weight_bf16(packed, scale, cols=weight.shape[1])
+    else:
+        packed, scale = module.quantize_w8_weight_bf16(weight)
+        dequant = module.dequantize_w8_weight_bf16(packed, scale)
+    return packed, scale, dequant
+
+
+def run_case(module, name: str, shape, bits: int, activation: str,
+             warmup: int, iterations: int):
+    m, k, h, n = shape
+    gated = activation in {"swiglu", "geglu"}
+    up_rows = 2 * h if gated else h
+    generator = torch.Generator(device="cuda").manual_seed(
+        91000 + m + k + h + n + bits + len(activation)
+    )
+    x = (torch.randn((m, k), generator=generator, device="cuda") * 0.1).bfloat16()
+    up_weight = (torch.randn((up_rows, k), generator=generator, device="cuda") * 0.02).bfloat16()
+    down_weight = (torch.randn((n, h), generator=generator, device="cuda") * 0.02).bfloat16()
+    up_bias = (torch.randn((up_rows,), generator=generator, device="cuda") * 0.01).bfloat16()
+    down_bias = (torch.randn((n,), generator=generator, device="cuda") * 0.01).bfloat16()
+    up_packed, up_scale, up_dequant = quantize(module, bits, up_weight)
+    down_packed, down_scale, down_dequant = quantize(module, bits, down_weight)
+    first = torch.empty((m, up_rows), device="cuda", dtype=torch.bfloat16)
+    hidden = torch.empty((m, h), device="cuda", dtype=torch.bfloat16)
+    out = torch.empty((m, n), device="cuda", dtype=torch.bfloat16)
+
+    if gated:
+        fn_name = f"w{bits}a16_{activation}_ffn_bf16"
+        kernel_fn = getattr(module, fn_name)
+
+        def kernel(variant: int):
+            return kernel_fn(
+                x, up_packed, up_scale, down_packed, down_scale,
+                gate_up_bias=up_bias, down_bias=down_bias, variant=variant,
+                workspace=(first, hidden), out=out,
+            )
+
+        def reference():
+            merged = F.linear(x, up_dequant, up_bias)
+            gate, up = merged.split(h, dim=-1)
+            act = F.silu(gate) if activation == "swiglu" else F.gelu(gate, approximate="tanh")
+            return F.linear(act * up, down_dequant, down_bias)
+    else:
+        fn_name = f"w{bits}a16_gelu_ffn_bf16"
+        kernel_fn = getattr(module, fn_name)
+
+        def kernel(variant: int):
+            return kernel_fn(
+                x, up_packed, up_scale, down_packed, down_scale,
+                up_bias=up_bias, down_bias=down_bias, variant=variant,
+                workspace=(first, hidden), out=out,
+            )
+
+        def reference():
+            return F.linear(
+                F.gelu(F.linear(x, up_dequant, up_bias), approximate="tanh"),
+                down_dequant, down_bias,
+            )
+
+    eager_us = bench(reference, warmup, iterations)
+    compiled = torch.compile(reference, fullgraph=True, mode="max-autotune-no-cudagraphs")
+    compiled()
+    torch.cuda.synchronize()
+    compiled_us = bench(compiled, warmup, iterations)
+    variants = {
+        str(variant): bench(
+            lambda variant=variant: kernel(variant), warmup, iterations
+        )
+        for variant in (1, 2, 3)
+    }
+    auto_error = None
+    try:
+        auto_us = bench(lambda: kernel(0), warmup, iterations)
+        kernel(0)
+    except RuntimeError as exc:
+        if "not qualified" not in str(exc) and "no qualified fast path" not in str(exc):
+            raise
+        auto_us = None
+        auto_error = str(exc)
+        best_diagnostic_variant = min(variants, key=variants.get)
+        kernel(int(best_diagnostic_variant))
+    ref = reference()
+    torch.cuda.synchronize()
+    diff = (out.float() - ref.float()).abs().flatten()
+    cosine = F.cosine_similarity(out.float().flatten(), ref.float().flatten(), dim=0)
+    best_diagnostic_variant = min(variants, key=variants.get)
+    best_diagnostic_us = variants[best_diagnostic_variant]
+    if auto_us is not None and auto_us > best_diagnostic_us * 1.05:
+        raise AssertionError(
+            f"{name} W{bits}A16 {activation}: auto {auto_us:.3f} us is more "
+            f"than 5% slower than diagnostic variant {best_diagnostic_variant} "
+            f"at {best_diagnostic_us:.3f} us"
+        )
+    if auto_us is not None and auto_us * 1.02 >= min(eager_us, compiled_us):
+        raise AssertionError(
+            f"{name} W{bits}A16 {activation}: accepted auto path must beat "
+            f"the strongest eager/compile baseline by at least 2%; "
+            f"auto={auto_us:.3f} us, eager={eager_us:.3f} us, "
+            f"compile={compiled_us:.3f} us"
+        )
+    return {
+        "region": "ffn",
+        "shape": name,
+        "M": m,
+        "K": k,
+        "H": h,
+        "N": n,
+        "precision": f"W{bits}A16",
+        "op": activation,
+        "eager_us": eager_us,
+        "compile_us": compiled_us,
+        "variant_us": variants,
+        "auto_status": "accepted" if auto_us is not None else "rejected",
+        "auto_us": auto_us,
+        "auto_error": auto_error,
+        "auto_speedup_vs_eager": eager_us / auto_us if auto_us is not None else None,
+        "auto_speedup_vs_compile": compiled_us / auto_us if auto_us is not None else None,
+        "best_diagnostic_variant": int(best_diagnostic_variant),
+        "best_diagnostic_us": best_diagnostic_us,
+        "max_abs": float(diff.max()),
+        "mean_abs": float(diff.mean()),
+        "p99_abs": float(torch.quantile(diff, 0.99)),
+        "cosine": float(cosine),
+    }
+
+
+def run_linear_case(module, name: str, shape, bits: int, warmup: int,
+                    iterations: int):
+    m, k, n = shape
+    generator = torch.Generator(device="cuda").manual_seed(
+        92000 + m + k + n + bits
+    )
+    x = (torch.randn((m, k), generator=generator, device="cuda") * 0.1).bfloat16()
+    weight = (
+        torch.randn((n, k), generator=generator, device="cuda") * 0.02
+    ).bfloat16()
+    packed, scale, dequant = quantize(module, bits, weight)
+    out = torch.empty((m, n), device="cuda", dtype=torch.bfloat16)
+    linear = getattr(module, f"w{bits}a16_linear_bf16")
+
+    def kernel(variant: int):
+        return linear(x, packed, scale, variant=variant, out=out)
+
+    def reference():
+        return F.linear(x, dequant)
+
+    eager_us = bench(reference, warmup, iterations)
+    compiled = torch.compile(
+        reference, fullgraph=True, mode="max-autotune-no-cudagraphs"
+    )
+    compiled()
+    torch.cuda.synchronize()
+    compiled_us = bench(compiled, warmup, iterations)
+    variants = {
+        str(variant): bench(
+            lambda variant=variant: kernel(variant), warmup, iterations
+        )
+        for variant in (1, 2, 3)
+    }
+    auto_error = None
+    try:
+        auto_us = bench(lambda: kernel(0), warmup, iterations)
+        kernel(0)
+    except RuntimeError as exc:
+        if "no qualified fast path" not in str(exc):
+            raise
+        auto_us = None
+        auto_error = str(exc)
+        kernel(int(min(variants, key=variants.get)))
+    ref = reference()
+    torch.cuda.synchronize()
+    diff = (out.float() - ref.float()).abs().flatten()
+    cosine = F.cosine_similarity(
+        out.float().flatten(), ref.float().flatten(), dim=0
+    )
+    best_variant = min(variants, key=variants.get)
+    best_us = variants[best_variant]
+    if auto_us is not None and auto_us > best_us * 1.05:
+        raise AssertionError(
+            f"{name} W{bits}A16 linear: auto {auto_us:.3f} us is more than "
+            f"5% slower than diagnostic variant {best_variant} at {best_us:.3f} us"
+        )
+    if auto_us is not None and auto_us * 1.02 >= min(eager_us, compiled_us):
+        raise AssertionError(
+            f"{name} W{bits}A16 linear: accepted auto path must beat the "
+            f"strongest eager/compile baseline by at least 2%; "
+            f"auto={auto_us:.3f} us, eager={eager_us:.3f} us, "
+            f"compile={compiled_us:.3f} us"
+        )
+    return {
+        "region": "linear",
+        "shape": name,
+        "M": m,
+        "K": k,
+        "N": n,
+        "precision": f"W{bits}A16",
+        "op": "linear",
+        "eager_us": eager_us,
+        "compile_us": compiled_us,
+        "variant_us": variants,
+        "auto_status": "accepted" if auto_us is not None else "rejected",
+        "auto_us": auto_us,
+        "auto_error": auto_error,
+        "auto_speedup_vs_eager": eager_us / auto_us if auto_us else None,
+        "auto_speedup_vs_compile": compiled_us / auto_us if auto_us else None,
+        "best_diagnostic_variant": int(best_variant),
+        "best_diagnostic_us": best_us,
+        "max_abs": float(diff.max()),
+        "mean_abs": float(diff.mean()),
+        "p99_abs": float(torch.quantile(diff, 0.99)),
+        "cosine": float(cosine),
+    }
+
+
+def main() -> int:
+    parser = argparse.ArgumentParser()
+    parser.add_argument("--backend", choices=["source", "installed"], default="source")
+    parser.add_argument("--artifact")
+    parser.add_argument("--mode", choices=["smoke", "full"], default="smoke")
+    parser.add_argument("--warmup", type=int, default=20)
+    parser.add_argument("--iterations", type=int, default=100)
+    parser.add_argument("--json-out")
+    args = parser.parse_args()
+    module = load_module(args.backend, args.artifact)
+    names = ["llm_m1"] if args.mode == "smoke" else list(SHAPES)
+    rows = []
+    for name in names:
+        for bits in (4, 8):
+            for activation in ("swiglu", "geglu", "gelu"):
+                rows.append(run_case(module, name, SHAPES[name], bits, activation,
+                                     args.warmup, args.iterations))
+    linear_names = ["llm_square_m1"] if args.mode == "smoke" else list(LINEAR_SHAPES)
+    for name in linear_names:
+        for bits in (4, 8):
+            rows.append(run_linear_case(
+                module, name, LINEAR_SHAPES[name], bits,
+                args.warmup, args.iterations,
+            ))
+    payload = {
+        "device": torch.cuda.get_device_name(),
+        "capability": list(torch.cuda.get_device_capability()),
+        "torch": torch.__version__,
+        "rows": rows,
+    }
+    print(json.dumps(payload, indent=2))
+    if args.json_out:
+        path = Path(args.json_out)
+        path.parent.mkdir(parents=True, exist_ok=True)
+        path.write_text(json.dumps(payload, indent=2) + "\n")
+    return 0
+
+
+if __name__ == "__main__":
+    raise SystemExit(main())

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

@@ -0,0 +1,261 @@
+"""Weight-only BF16-activation linear and FFN regions."""
+
+from __future__ import annotations
+
+from typing import Optional
+
+import torch
+
+from ._ops import add_op_namespace_prefix, ops
+
+
+def w4_sfb_size_bytes(rows: int, cols: int) -> int:
+    if rows <= 0 or cols <= 0 or cols % 64:
+        raise ValueError("rows must be positive and cols must be divisible by 64")
+    return ((int(rows) + 127) // 128) * (((int(cols) // 16) + 3) // 4) * 512
+
+
+@torch.library.register_fake(add_op_namespace_prefix("quantize_w4_weight_bf16"))
+def _quantize_w4_fake(weight: torch.Tensor, packed: torch.Tensor, sfb: torch.Tensor) -> None:
+    return None
+
+
+@torch.library.register_fake(add_op_namespace_prefix("dequantize_w4_weight_bf16"))
+def _dequantize_w4_fake(packed: torch.Tensor, sfb: torch.Tensor, weight: torch.Tensor) -> None:
+    return None
+
+
+@torch.library.register_fake(add_op_namespace_prefix("w4a16_linear_bf16"))
+def _w4_linear_fake(
+    x: torch.Tensor,
+    packed: torch.Tensor,
+    sfb: torch.Tensor,
+    alpha: float,
+    variant: int,
+    out: torch.Tensor,
+) -> None:
+    return None
+
+
+@torch.library.register_fake(add_op_namespace_prefix("quantize_w8_weight_bf16"))
+def _quantize_w8_fake(weight: torch.Tensor, quantized: torch.Tensor, scales: torch.Tensor) -> None:
+    return None
+
+
+@torch.library.register_fake(add_op_namespace_prefix("dequantize_w8_weight_bf16"))
+def _dequantize_w8_fake(quantized: torch.Tensor, scales: torch.Tensor, weight: torch.Tensor) -> None:
+    return None
+
+
+@torch.library.register_fake(add_op_namespace_prefix("w8a16_linear_bf16"))
+def _w8_linear_fake(
+    x: torch.Tensor,
+    quantized: torch.Tensor,
+    scales: torch.Tensor,
+    variant: int,
+    out: torch.Tensor,
+) -> None:
+    return None
+
+
+@torch.library.register_fake(add_op_namespace_prefix("w4a16_gated_ffn_bf16"))
+def _w4_gated_fake(
+    x, gu_w, gu_s, dn_w, dn_s, gu_b, dn_b, gelu, gu_alpha, dn_alpha,
+    variant, gu, hidden, out,
+) -> None:
+    return None
+
+
+@torch.library.register_fake(add_op_namespace_prefix("w8a16_gated_ffn_bf16"))
+def _w8_gated_fake(
+    x, gu_w, gu_s, dn_w, dn_s, gu_b, dn_b, gelu, variant, gu, hidden, out,
+) -> None:
+    return None
+
+
+@torch.library.register_fake(add_op_namespace_prefix("w4a16_gelu_ffn_bf16"))
+def _w4_gelu_fake(
+    x, up_w, up_s, dn_w, dn_s, up_b, dn_b, up_alpha, dn_alpha,
+    variant, up, hidden, out,
+) -> None:
+    return None
+
+
+@torch.library.register_fake(add_op_namespace_prefix("w8a16_gelu_ffn_bf16"))
+def _w8_gelu_fake(
+    x, up_w, up_s, dn_w, dn_s, up_b, dn_b, variant, up, hidden, out,
+) -> None:
+    return None
+
+
+def quantize_w4_weight_bf16(weight: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor]:
+    """Pack a static BF16 weight matrix into NVFP4 data and SFB scales."""
+    n, k = weight.shape
+    packed = torch.empty((n, k // 2), device=weight.device, dtype=torch.uint8)
+    sfb = torch.empty((w4_sfb_size_bytes(n, k),), device=weight.device, dtype=torch.uint8)
+    ops.quantize_w4_weight_bf16(weight, packed, sfb)
+    return packed, sfb
+
+
+def dequantize_w4_weight_bf16(
+    packed: torch.Tensor, sfb: torch.Tensor, *, cols: int
+) -> torch.Tensor:
+    weight = torch.empty((packed.shape[0], int(cols)), device=packed.device, dtype=torch.bfloat16)
+    ops.dequantize_w4_weight_bf16(packed, sfb, weight)
+    return weight
+
+
+def quantize_w8_weight_bf16(weight: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor]:
+    """Symmetrically quantize a static BF16 weight matrix per output channel."""
+    quantized = torch.empty_like(weight, dtype=torch.int8)
+    scales = torch.empty((weight.shape[0],), device=weight.device, dtype=torch.float32)
+    ops.quantize_w8_weight_bf16(weight, quantized, scales)
+    return quantized, scales
+
+
+def dequantize_w8_weight_bf16(
+    quantized: torch.Tensor, scales: torch.Tensor
+) -> torch.Tensor:
+    weight = torch.empty_like(quantized, dtype=torch.bfloat16)
+    ops.dequantize_w8_weight_bf16(quantized, scales, weight)
+    return weight
+
+
+def w4a16_linear_bf16(
+    x: torch.Tensor,
+    weight_packed: torch.Tensor,
+    weight_sfb: torch.Tensor,
+    *,
+    alpha: float = 1.0,
+    variant: int = 0,
+    out: Optional[torch.Tensor] = None,
+) -> torch.Tensor:
+    if out is None:
+        out = torch.empty((x.shape[0], weight_packed.shape[0]), device=x.device, dtype=torch.bfloat16)
+    ops.w4a16_linear_bf16(x, weight_packed, weight_sfb, float(alpha), int(variant), out)
+    return out
+
+
+def w8a16_linear_bf16(
+    x: torch.Tensor,
+    weight_int8: torch.Tensor,
+    weight_scale: torch.Tensor,
+    *,
+    variant: int = 0,
+    out: Optional[torch.Tensor] = None,
+) -> torch.Tensor:
+    if out is None:
+        out = torch.empty((x.shape[0], weight_int8.shape[0]), device=x.device, dtype=torch.bfloat16)
+    ops.w8a16_linear_bf16(x, weight_int8, weight_scale, int(variant), out)
+    return out
+
+
+def _gated_workspace(x: torch.Tensor, gate_up_rows: int, out_rows: int, workspace, out):
+    hidden_size = gate_up_rows // 2
+    if workspace is None:
+        gate_up = torch.empty((x.shape[0], gate_up_rows), device=x.device, dtype=torch.bfloat16)
+        hidden = torch.empty((x.shape[0], hidden_size), device=x.device, dtype=torch.bfloat16)
+    else:
+        gate_up, hidden = workspace
+    if out is None:
+        out = torch.empty((x.shape[0], out_rows), device=x.device, dtype=torch.bfloat16)
+    return gate_up, hidden, out
+
+
+def _gelu_workspace(x: torch.Tensor, hidden_size: int, out_rows: int, workspace, out):
+    if workspace is None:
+        up = torch.empty((x.shape[0], hidden_size), device=x.device, dtype=torch.bfloat16)
+        hidden = torch.empty_like(up)
+    else:
+        up, hidden = workspace
+    if out is None:
+        out = torch.empty((x.shape[0], out_rows), device=x.device, dtype=torch.bfloat16)
+    return up, hidden, out
+
+
+def _w4_gated(
+    x, gate_up_weight, gate_up_sfb, down_weight, down_sfb, *, gelu,
+    gate_up_bias=None, down_bias=None, gate_up_alpha=1.0, down_alpha=1.0,
+    variant=0, workspace=None, out=None,
+):
+    gate_up, hidden, out = _gated_workspace(
+        x, gate_up_weight.shape[0], down_weight.shape[0], workspace, out
+    )
+    ops.w4a16_gated_ffn_bf16(
+        x, gate_up_weight, gate_up_sfb, down_weight, down_sfb,
+        gate_up_bias, down_bias, bool(gelu), float(gate_up_alpha),
+        float(down_alpha), int(variant), gate_up, hidden, out,
+    )
+    return out
+
+
+def w4a16_swiglu_ffn_bf16(x, gate_up_weight, gate_up_sfb, down_weight, down_sfb, **kwargs):
+    return _w4_gated(x, gate_up_weight, gate_up_sfb, down_weight, down_sfb, gelu=False, **kwargs)
+
+
+def w4a16_geglu_ffn_bf16(x, gate_up_weight, gate_up_sfb, down_weight, down_sfb, **kwargs):
+    return _w4_gated(x, gate_up_weight, gate_up_sfb, down_weight, down_sfb, gelu=True, **kwargs)
+
+
+def _w8_gated(
+    x, gate_up_weight, gate_up_scale, down_weight, down_scale, *, gelu,
+    gate_up_bias=None, down_bias=None, variant=0, workspace=None, out=None,
+):
+    gate_up, hidden, out = _gated_workspace(
+        x, gate_up_weight.shape[0], down_weight.shape[0], workspace, out
+    )
+    ops.w8a16_gated_ffn_bf16(
+        x, gate_up_weight, gate_up_scale, down_weight, down_scale,
+        gate_up_bias, down_bias, bool(gelu), int(variant), gate_up, hidden, out,
+    )
+    return out
+
+
+def w8a16_swiglu_ffn_bf16(x, gate_up_weight, gate_up_scale, down_weight, down_scale, **kwargs):
+    return _w8_gated(x, gate_up_weight, gate_up_scale, down_weight, down_scale, gelu=False, **kwargs)
+
+
+def w8a16_geglu_ffn_bf16(x, gate_up_weight, gate_up_scale, down_weight, down_scale, **kwargs):
+    return _w8_gated(x, gate_up_weight, gate_up_scale, down_weight, down_scale, gelu=True, **kwargs)
+
+
+def w4a16_gelu_ffn_bf16(
+    x, up_weight, up_sfb, down_weight, down_sfb, *, up_bias=None,
+    down_bias=None, up_alpha=1.0, down_alpha=1.0, variant=0,
+    workspace=None, out=None,
+):
+    up, hidden, out = _gelu_workspace(x, up_weight.shape[0], down_weight.shape[0], workspace, out)
+    ops.w4a16_gelu_ffn_bf16(
+        x, up_weight, up_sfb, down_weight, down_sfb, up_bias, down_bias,
+        float(up_alpha), float(down_alpha), int(variant), up, hidden, out,
+    )
+    return out
+
+
+def w8a16_gelu_ffn_bf16(
+    x, up_weight, up_scale, down_weight, down_scale, *, up_bias=None,
+    down_bias=None, variant=0, workspace=None, out=None,
+):
+    up, hidden, out = _gelu_workspace(x, up_weight.shape[0], down_weight.shape[0], workspace, out)
+    ops.w8a16_gelu_ffn_bf16(
+        x, up_weight, up_scale, down_weight, down_scale, up_bias, down_bias,
+        int(variant), up, hidden, out,
+    )
+    return out
+
+
+__all__ = [
+    "dequantize_w4_weight_bf16",
+    "dequantize_w8_weight_bf16",
+    "quantize_w4_weight_bf16",
+    "quantize_w8_weight_bf16",
+    "w4_sfb_size_bytes",
+    "w4a16_geglu_ffn_bf16",
+    "w4a16_gelu_ffn_bf16",
+    "w4a16_linear_bf16",
+    "w4a16_swiglu_ffn_bf16",
+    "w8a16_geglu_ffn_bf16",
+    "w8a16_gelu_ffn_bf16",
+    "w8a16_linear_bf16",
+    "w8a16_swiglu_ffn_bf16",
+]

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

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

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

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

build/torch211-cxx11-cu128-x86_64-linux/metadata.json

@@ -0,0 +1,33 @@
+{
+  "name": "weight-only-ffn",
+  "id": "_weight_only_ffn_cuda_38d37b6",
+  "version": 1,
+  "license": "Apache-2.0",
+  "python-depends": [],
+  "backend": {
+    "type": "cuda",
+    "archs": [
+      "12.0a"
+    ]
+  },
+  "digest": {
+    "algorithm": "sha256",
+    "files": {
+      "__init__.py": "0q6TMj65una1LdyHLC291Dwy2VzrBbAD8Wm8MEdmqwA=",
+      "_ops.py": "iqY+HoE4nGHqvVlEZTwEAcyu1/w+QrjwuuW9sdTI2uk=",
+      "_weight_only_ffn_cuda_38d37b6.abi3.so": "fuOR5RgCyiTtpsPHPLG3EeYYZNvpJyufa3nM5JHer/4=",
+      "weight_only_ffn/__init__.py": "DFYPlrhXwYjEqCl/8n0SmWGZV8NFml5DPhMjKfv98GY="
+    }
+  },
+  "provenance": {
+    "kernel-builder": {
+      "version": "0.17.0-dev0",
+      "sha": "19aaa6421e674e9fecc352bbae6eab81d19a6bf4",
+      "dirty": false
+    },
+    "kernel": {
+      "sha": "38d37b67acfe940c99d7b3db721fa67163b82cec",
+      "dirty": false
+    }
+  }
+}

build/torch211-cxx11-cu128-x86_64-linux/weight_only_ffn/__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-aarch64-linux/__init__.py

@@ -0,0 +1,261 @@
+"""Weight-only BF16-activation linear and FFN regions."""
+
+from __future__ import annotations
+
+from typing import Optional
+
+import torch
+
+from ._ops import add_op_namespace_prefix, ops
+
+
+def w4_sfb_size_bytes(rows: int, cols: int) -> int:
+    if rows <= 0 or cols <= 0 or cols % 64:
+        raise ValueError("rows must be positive and cols must be divisible by 64")
+    return ((int(rows) + 127) // 128) * (((int(cols) // 16) + 3) // 4) * 512
+
+
+@torch.library.register_fake(add_op_namespace_prefix("quantize_w4_weight_bf16"))
+def _quantize_w4_fake(weight: torch.Tensor, packed: torch.Tensor, sfb: torch.Tensor) -> None:
+    return None
+
+
+@torch.library.register_fake(add_op_namespace_prefix("dequantize_w4_weight_bf16"))
+def _dequantize_w4_fake(packed: torch.Tensor, sfb: torch.Tensor, weight: torch.Tensor) -> None:
+    return None
+
+
+@torch.library.register_fake(add_op_namespace_prefix("w4a16_linear_bf16"))
+def _w4_linear_fake(
+    x: torch.Tensor,
+    packed: torch.Tensor,
+    sfb: torch.Tensor,
+    alpha: float,
+    variant: int,
+    out: torch.Tensor,
+) -> None:
+    return None
+
+
+@torch.library.register_fake(add_op_namespace_prefix("quantize_w8_weight_bf16"))
+def _quantize_w8_fake(weight: torch.Tensor, quantized: torch.Tensor, scales: torch.Tensor) -> None:
+    return None
+
+
+@torch.library.register_fake(add_op_namespace_prefix("dequantize_w8_weight_bf16"))
+def _dequantize_w8_fake(quantized: torch.Tensor, scales: torch.Tensor, weight: torch.Tensor) -> None:
+    return None
+
+
+@torch.library.register_fake(add_op_namespace_prefix("w8a16_linear_bf16"))
+def _w8_linear_fake(
+    x: torch.Tensor,
+    quantized: torch.Tensor,
+    scales: torch.Tensor,
+    variant: int,
+    out: torch.Tensor,
+) -> None:
+    return None
+
+
+@torch.library.register_fake(add_op_namespace_prefix("w4a16_gated_ffn_bf16"))
+def _w4_gated_fake(
+    x, gu_w, gu_s, dn_w, dn_s, gu_b, dn_b, gelu, gu_alpha, dn_alpha,
+    variant, gu, hidden, out,
+) -> None:
+    return None
+
+
+@torch.library.register_fake(add_op_namespace_prefix("w8a16_gated_ffn_bf16"))
+def _w8_gated_fake(
+    x, gu_w, gu_s, dn_w, dn_s, gu_b, dn_b, gelu, variant, gu, hidden, out,
+) -> None:
+    return None
+
+
+@torch.library.register_fake(add_op_namespace_prefix("w4a16_gelu_ffn_bf16"))
+def _w4_gelu_fake(
+    x, up_w, up_s, dn_w, dn_s, up_b, dn_b, up_alpha, dn_alpha,
+    variant, up, hidden, out,
+) -> None:
+    return None
+
+
+@torch.library.register_fake(add_op_namespace_prefix("w8a16_gelu_ffn_bf16"))
+def _w8_gelu_fake(
+    x, up_w, up_s, dn_w, dn_s, up_b, dn_b, variant, up, hidden, out,
+) -> None:
+    return None
+
+
+def quantize_w4_weight_bf16(weight: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor]:
+    """Pack a static BF16 weight matrix into NVFP4 data and SFB scales."""
+    n, k = weight.shape
+    packed = torch.empty((n, k // 2), device=weight.device, dtype=torch.uint8)
+    sfb = torch.empty((w4_sfb_size_bytes(n, k),), device=weight.device, dtype=torch.uint8)
+    ops.quantize_w4_weight_bf16(weight, packed, sfb)
+    return packed, sfb
+
+
+def dequantize_w4_weight_bf16(
+    packed: torch.Tensor, sfb: torch.Tensor, *, cols: int
+) -> torch.Tensor:
+    weight = torch.empty((packed.shape[0], int(cols)), device=packed.device, dtype=torch.bfloat16)
+    ops.dequantize_w4_weight_bf16(packed, sfb, weight)
+    return weight
+
+
+def quantize_w8_weight_bf16(weight: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor]:
+    """Symmetrically quantize a static BF16 weight matrix per output channel."""
+    quantized = torch.empty_like(weight, dtype=torch.int8)
+    scales = torch.empty((weight.shape[0],), device=weight.device, dtype=torch.float32)
+    ops.quantize_w8_weight_bf16(weight, quantized, scales)
+    return quantized, scales
+
+
+def dequantize_w8_weight_bf16(
+    quantized: torch.Tensor, scales: torch.Tensor
+) -> torch.Tensor:
+    weight = torch.empty_like(quantized, dtype=torch.bfloat16)
+    ops.dequantize_w8_weight_bf16(quantized, scales, weight)
+    return weight
+
+
+def w4a16_linear_bf16(
+    x: torch.Tensor,
+    weight_packed: torch.Tensor,
+    weight_sfb: torch.Tensor,
+    *,
+    alpha: float = 1.0,
+    variant: int = 0,
+    out: Optional[torch.Tensor] = None,
+) -> torch.Tensor:
+    if out is None:
+        out = torch.empty((x.shape[0], weight_packed.shape[0]), device=x.device, dtype=torch.bfloat16)
+    ops.w4a16_linear_bf16(x, weight_packed, weight_sfb, float(alpha), int(variant), out)
+    return out
+
+
+def w8a16_linear_bf16(
+    x: torch.Tensor,
+    weight_int8: torch.Tensor,
+    weight_scale: torch.Tensor,
+    *,
+    variant: int = 0,
+    out: Optional[torch.Tensor] = None,
+) -> torch.Tensor:
+    if out is None:
+        out = torch.empty((x.shape[0], weight_int8.shape[0]), device=x.device, dtype=torch.bfloat16)
+    ops.w8a16_linear_bf16(x, weight_int8, weight_scale, int(variant), out)
+    return out
+
+
+def _gated_workspace(x: torch.Tensor, gate_up_rows: int, out_rows: int, workspace, out):
+    hidden_size = gate_up_rows // 2
+    if workspace is None:
+        gate_up = torch.empty((x.shape[0], gate_up_rows), device=x.device, dtype=torch.bfloat16)
+        hidden = torch.empty((x.shape[0], hidden_size), device=x.device, dtype=torch.bfloat16)
+    else:
+        gate_up, hidden = workspace
+    if out is None:
+        out = torch.empty((x.shape[0], out_rows), device=x.device, dtype=torch.bfloat16)
+    return gate_up, hidden, out
+
+
+def _gelu_workspace(x: torch.Tensor, hidden_size: int, out_rows: int, workspace, out):
+    if workspace is None:
+        up = torch.empty((x.shape[0], hidden_size), device=x.device, dtype=torch.bfloat16)
+        hidden = torch.empty_like(up)
+    else:
+        up, hidden = workspace
+    if out is None:
+        out = torch.empty((x.shape[0], out_rows), device=x.device, dtype=torch.bfloat16)
+    return up, hidden, out
+
+
+def _w4_gated(
+    x, gate_up_weight, gate_up_sfb, down_weight, down_sfb, *, gelu,
+    gate_up_bias=None, down_bias=None, gate_up_alpha=1.0, down_alpha=1.0,
+    variant=0, workspace=None, out=None,
+):
+    gate_up, hidden, out = _gated_workspace(
+        x, gate_up_weight.shape[0], down_weight.shape[0], workspace, out
+    )
+    ops.w4a16_gated_ffn_bf16(
+        x, gate_up_weight, gate_up_sfb, down_weight, down_sfb,
+        gate_up_bias, down_bias, bool(gelu), float(gate_up_alpha),
+        float(down_alpha), int(variant), gate_up, hidden, out,
+    )
+    return out
+
+
+def w4a16_swiglu_ffn_bf16(x, gate_up_weight, gate_up_sfb, down_weight, down_sfb, **kwargs):
+    return _w4_gated(x, gate_up_weight, gate_up_sfb, down_weight, down_sfb, gelu=False, **kwargs)
+
+
+def w4a16_geglu_ffn_bf16(x, gate_up_weight, gate_up_sfb, down_weight, down_sfb, **kwargs):
+    return _w4_gated(x, gate_up_weight, gate_up_sfb, down_weight, down_sfb, gelu=True, **kwargs)
+
+
+def _w8_gated(
+    x, gate_up_weight, gate_up_scale, down_weight, down_scale, *, gelu,
+    gate_up_bias=None, down_bias=None, variant=0, workspace=None, out=None,
+):
+    gate_up, hidden, out = _gated_workspace(
+        x, gate_up_weight.shape[0], down_weight.shape[0], workspace, out
+    )
+    ops.w8a16_gated_ffn_bf16(
+        x, gate_up_weight, gate_up_scale, down_weight, down_scale,
+        gate_up_bias, down_bias, bool(gelu), int(variant), gate_up, hidden, out,
+    )
+    return out
+
+
+def w8a16_swiglu_ffn_bf16(x, gate_up_weight, gate_up_scale, down_weight, down_scale, **kwargs):
+    return _w8_gated(x, gate_up_weight, gate_up_scale, down_weight, down_scale, gelu=False, **kwargs)
+
+
+def w8a16_geglu_ffn_bf16(x, gate_up_weight, gate_up_scale, down_weight, down_scale, **kwargs):
+    return _w8_gated(x, gate_up_weight, gate_up_scale, down_weight, down_scale, gelu=True, **kwargs)
+
+
+def w4a16_gelu_ffn_bf16(
+    x, up_weight, up_sfb, down_weight, down_sfb, *, up_bias=None,
+    down_bias=None, up_alpha=1.0, down_alpha=1.0, variant=0,
+    workspace=None, out=None,
+):
+    up, hidden, out = _gelu_workspace(x, up_weight.shape[0], down_weight.shape[0], workspace, out)
+    ops.w4a16_gelu_ffn_bf16(
+        x, up_weight, up_sfb, down_weight, down_sfb, up_bias, down_bias,
+        float(up_alpha), float(down_alpha), int(variant), up, hidden, out,
+    )
+    return out
+
+
+def w8a16_gelu_ffn_bf16(
+    x, up_weight, up_scale, down_weight, down_scale, *, up_bias=None,
+    down_bias=None, variant=0, workspace=None, out=None,
+):
+    up, hidden, out = _gelu_workspace(x, up_weight.shape[0], down_weight.shape[0], workspace, out)
+    ops.w8a16_gelu_ffn_bf16(
+        x, up_weight, up_scale, down_weight, down_scale, up_bias, down_bias,
+        int(variant), up, hidden, out,
+    )
+    return out
+
+
+__all__ = [
+    "dequantize_w4_weight_bf16",
+    "dequantize_w8_weight_bf16",
+    "quantize_w4_weight_bf16",
+    "quantize_w8_weight_bf16",
+    "w4_sfb_size_bytes",
+    "w4a16_geglu_ffn_bf16",
+    "w4a16_gelu_ffn_bf16",
+    "w4a16_linear_bf16",
+    "w4a16_swiglu_ffn_bf16",
+    "w8a16_geglu_ffn_bf16",
+    "w8a16_gelu_ffn_bf16",
+    "w8a16_linear_bf16",
+    "w8a16_swiglu_ffn_bf16",
+]

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

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

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

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

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

@@ -0,0 +1,26 @@
+{
+  "name": "weight-only-ffn",
+  "id": "_weight_only_ffn_cuda_38d37b6",
+  "version": 1,
+  "license": "Apache-2.0",
+  "python-depends": [],
+  "backend": {
+    "type": "cuda",
+    "archs": [
+      "11.0a",
+      "12.0a",
+      "12.1"
+    ]
+  },
+  "provenance": {
+    "kernel-builder": {
+      "version": "0.17.0-dev0",
+      "sha": "19aaa6421e674e9fecc352bbae6eab81d19a6bf4",
+      "dirty": false
+    },
+    "kernel": {
+      "sha": "38d37b67acfe940c99d7b3db721fa67163b82cec",
+      "dirty": false
+    }
+  }
+}

build/torch211-cxx11-cu130-aarch64-linux/weight_only_ffn/__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/__init__.py

@@ -0,0 +1,261 @@
+"""Weight-only BF16-activation linear and FFN regions."""
+
+from __future__ import annotations
+
+from typing import Optional
+
+import torch
+
+from ._ops import add_op_namespace_prefix, ops
+
+
+def w4_sfb_size_bytes(rows: int, cols: int) -> int:
+    if rows <= 0 or cols <= 0 or cols % 64:
+        raise ValueError("rows must be positive and cols must be divisible by 64")
+    return ((int(rows) + 127) // 128) * (((int(cols) // 16) + 3) // 4) * 512
+
+
+@torch.library.register_fake(add_op_namespace_prefix("quantize_w4_weight_bf16"))
+def _quantize_w4_fake(weight: torch.Tensor, packed: torch.Tensor, sfb: torch.Tensor) -> None:
+    return None
+
+
+@torch.library.register_fake(add_op_namespace_prefix("dequantize_w4_weight_bf16"))
+def _dequantize_w4_fake(packed: torch.Tensor, sfb: torch.Tensor, weight: torch.Tensor) -> None:
+    return None
+
+
+@torch.library.register_fake(add_op_namespace_prefix("w4a16_linear_bf16"))
+def _w4_linear_fake(
+    x: torch.Tensor,
+    packed: torch.Tensor,
+    sfb: torch.Tensor,
+    alpha: float,
+    variant: int,
+    out: torch.Tensor,
+) -> None:
+    return None
+
+
+@torch.library.register_fake(add_op_namespace_prefix("quantize_w8_weight_bf16"))
+def _quantize_w8_fake(weight: torch.Tensor, quantized: torch.Tensor, scales: torch.Tensor) -> None:
+    return None
+
+
+@torch.library.register_fake(add_op_namespace_prefix("dequantize_w8_weight_bf16"))
+def _dequantize_w8_fake(quantized: torch.Tensor, scales: torch.Tensor, weight: torch.Tensor) -> None:
+    return None
+
+
+@torch.library.register_fake(add_op_namespace_prefix("w8a16_linear_bf16"))
+def _w8_linear_fake(
+    x: torch.Tensor,
+    quantized: torch.Tensor,
+    scales: torch.Tensor,
+    variant: int,
+    out: torch.Tensor,
+) -> None:
+    return None
+
+
+@torch.library.register_fake(add_op_namespace_prefix("w4a16_gated_ffn_bf16"))
+def _w4_gated_fake(
+    x, gu_w, gu_s, dn_w, dn_s, gu_b, dn_b, gelu, gu_alpha, dn_alpha,
+    variant, gu, hidden, out,
+) -> None:
+    return None
+
+
+@torch.library.register_fake(add_op_namespace_prefix("w8a16_gated_ffn_bf16"))
+def _w8_gated_fake(
+    x, gu_w, gu_s, dn_w, dn_s, gu_b, dn_b, gelu, variant, gu, hidden, out,
+) -> None:
+    return None
+
+
+@torch.library.register_fake(add_op_namespace_prefix("w4a16_gelu_ffn_bf16"))
+def _w4_gelu_fake(
+    x, up_w, up_s, dn_w, dn_s, up_b, dn_b, up_alpha, dn_alpha,
+    variant, up, hidden, out,
+) -> None:
+    return None
+
+
+@torch.library.register_fake(add_op_namespace_prefix("w8a16_gelu_ffn_bf16"))
+def _w8_gelu_fake(
+    x, up_w, up_s, dn_w, dn_s, up_b, dn_b, variant, up, hidden, out,
+) -> None:
+    return None
+
+
+def quantize_w4_weight_bf16(weight: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor]:
+    """Pack a static BF16 weight matrix into NVFP4 data and SFB scales."""
+    n, k = weight.shape
+    packed = torch.empty((n, k // 2), device=weight.device, dtype=torch.uint8)
+    sfb = torch.empty((w4_sfb_size_bytes(n, k),), device=weight.device, dtype=torch.uint8)
+    ops.quantize_w4_weight_bf16(weight, packed, sfb)
+    return packed, sfb
+
+
+def dequantize_w4_weight_bf16(
+    packed: torch.Tensor, sfb: torch.Tensor, *, cols: int
+) -> torch.Tensor:
+    weight = torch.empty((packed.shape[0], int(cols)), device=packed.device, dtype=torch.bfloat16)
+    ops.dequantize_w4_weight_bf16(packed, sfb, weight)
+    return weight
+
+
+def quantize_w8_weight_bf16(weight: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor]:
+    """Symmetrically quantize a static BF16 weight matrix per output channel."""
+    quantized = torch.empty_like(weight, dtype=torch.int8)
+    scales = torch.empty((weight.shape[0],), device=weight.device, dtype=torch.float32)
+    ops.quantize_w8_weight_bf16(weight, quantized, scales)
+    return quantized, scales
+
+
+def dequantize_w8_weight_bf16(
+    quantized: torch.Tensor, scales: torch.Tensor
+) -> torch.Tensor:
+    weight = torch.empty_like(quantized, dtype=torch.bfloat16)
+    ops.dequantize_w8_weight_bf16(quantized, scales, weight)
+    return weight
+
+
+def w4a16_linear_bf16(
+    x: torch.Tensor,
+    weight_packed: torch.Tensor,
+    weight_sfb: torch.Tensor,
+    *,
+    alpha: float = 1.0,
+    variant: int = 0,
+    out: Optional[torch.Tensor] = None,
+) -> torch.Tensor:
+    if out is None:
+        out = torch.empty((x.shape[0], weight_packed.shape[0]), device=x.device, dtype=torch.bfloat16)
+    ops.w4a16_linear_bf16(x, weight_packed, weight_sfb, float(alpha), int(variant), out)
+    return out
+
+
+def w8a16_linear_bf16(
+    x: torch.Tensor,
+    weight_int8: torch.Tensor,
+    weight_scale: torch.Tensor,
+    *,
+    variant: int = 0,
+    out: Optional[torch.Tensor] = None,
+) -> torch.Tensor:
+    if out is None:
+        out = torch.empty((x.shape[0], weight_int8.shape[0]), device=x.device, dtype=torch.bfloat16)
+    ops.w8a16_linear_bf16(x, weight_int8, weight_scale, int(variant), out)
+    return out
+
+
+def _gated_workspace(x: torch.Tensor, gate_up_rows: int, out_rows: int, workspace, out):
+    hidden_size = gate_up_rows // 2
+    if workspace is None:
+        gate_up = torch.empty((x.shape[0], gate_up_rows), device=x.device, dtype=torch.bfloat16)
+        hidden = torch.empty((x.shape[0], hidden_size), device=x.device, dtype=torch.bfloat16)
+    else:
+        gate_up, hidden = workspace
+    if out is None:
+        out = torch.empty((x.shape[0], out_rows), device=x.device, dtype=torch.bfloat16)
+    return gate_up, hidden, out
+
+
+def _gelu_workspace(x: torch.Tensor, hidden_size: int, out_rows: int, workspace, out):
+    if workspace is None:
+        up = torch.empty((x.shape[0], hidden_size), device=x.device, dtype=torch.bfloat16)
+        hidden = torch.empty_like(up)
+    else:
+        up, hidden = workspace
+    if out is None:
+        out = torch.empty((x.shape[0], out_rows), device=x.device, dtype=torch.bfloat16)
+    return up, hidden, out
+
+
+def _w4_gated(
+    x, gate_up_weight, gate_up_sfb, down_weight, down_sfb, *, gelu,
+    gate_up_bias=None, down_bias=None, gate_up_alpha=1.0, down_alpha=1.0,
+    variant=0, workspace=None, out=None,
+):
+    gate_up, hidden, out = _gated_workspace(
+        x, gate_up_weight.shape[0], down_weight.shape[0], workspace, out
+    )
+    ops.w4a16_gated_ffn_bf16(
+        x, gate_up_weight, gate_up_sfb, down_weight, down_sfb,
+        gate_up_bias, down_bias, bool(gelu), float(gate_up_alpha),
+        float(down_alpha), int(variant), gate_up, hidden, out,
+    )
+    return out
+
+
+def w4a16_swiglu_ffn_bf16(x, gate_up_weight, gate_up_sfb, down_weight, down_sfb, **kwargs):
+    return _w4_gated(x, gate_up_weight, gate_up_sfb, down_weight, down_sfb, gelu=False, **kwargs)
+
+
+def w4a16_geglu_ffn_bf16(x, gate_up_weight, gate_up_sfb, down_weight, down_sfb, **kwargs):
+    return _w4_gated(x, gate_up_weight, gate_up_sfb, down_weight, down_sfb, gelu=True, **kwargs)
+
+
+def _w8_gated(
+    x, gate_up_weight, gate_up_scale, down_weight, down_scale, *, gelu,
+    gate_up_bias=None, down_bias=None, variant=0, workspace=None, out=None,
+):
+    gate_up, hidden, out = _gated_workspace(
+        x, gate_up_weight.shape[0], down_weight.shape[0], workspace, out
+    )
+    ops.w8a16_gated_ffn_bf16(
+        x, gate_up_weight, gate_up_scale, down_weight, down_scale,
+        gate_up_bias, down_bias, bool(gelu), int(variant), gate_up, hidden, out,
+    )
+    return out
+
+
+def w8a16_swiglu_ffn_bf16(x, gate_up_weight, gate_up_scale, down_weight, down_scale, **kwargs):
+    return _w8_gated(x, gate_up_weight, gate_up_scale, down_weight, down_scale, gelu=False, **kwargs)
+
+
+def w8a16_geglu_ffn_bf16(x, gate_up_weight, gate_up_scale, down_weight, down_scale, **kwargs):
+    return _w8_gated(x, gate_up_weight, gate_up_scale, down_weight, down_scale, gelu=True, **kwargs)
+
+
+def w4a16_gelu_ffn_bf16(
+    x, up_weight, up_sfb, down_weight, down_sfb, *, up_bias=None,
+    down_bias=None, up_alpha=1.0, down_alpha=1.0, variant=0,
+    workspace=None, out=None,
+):
+    up, hidden, out = _gelu_workspace(x, up_weight.shape[0], down_weight.shape[0], workspace, out)
+    ops.w4a16_gelu_ffn_bf16(
+        x, up_weight, up_sfb, down_weight, down_sfb, up_bias, down_bias,
+        float(up_alpha), float(down_alpha), int(variant), up, hidden, out,
+    )
+    return out
+
+
+def w8a16_gelu_ffn_bf16(
+    x, up_weight, up_scale, down_weight, down_scale, *, up_bias=None,
+    down_bias=None, variant=0, workspace=None, out=None,
+):
+    up, hidden, out = _gelu_workspace(x, up_weight.shape[0], down_weight.shape[0], workspace, out)
+    ops.w8a16_gelu_ffn_bf16(
+        x, up_weight, up_scale, down_weight, down_scale, up_bias, down_bias,
+        int(variant), up, hidden, out,
+    )
+    return out
+
+
+__all__ = [
+    "dequantize_w4_weight_bf16",
+    "dequantize_w8_weight_bf16",
+    "quantize_w4_weight_bf16",
+    "quantize_w8_weight_bf16",
+    "w4_sfb_size_bytes",
+    "w4a16_geglu_ffn_bf16",
+    "w4a16_gelu_ffn_bf16",
+    "w4a16_linear_bf16",
+    "w4a16_swiglu_ffn_bf16",
+    "w8a16_geglu_ffn_bf16",
+    "w8a16_gelu_ffn_bf16",
+    "w8a16_linear_bf16",
+    "w8a16_swiglu_ffn_bf16",
+]

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

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

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

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

build/torch211-cxx11-cu130-x86_64-linux/metadata.json

@@ -0,0 +1,35 @@
+{
+  "name": "weight-only-ffn",
+  "id": "_weight_only_ffn_cuda_38d37b6",
+  "version": 1,
+  "license": "Apache-2.0",
+  "python-depends": [],
+  "backend": {
+    "type": "cuda",
+    "archs": [
+      "11.0a",
+      "12.0a",
+      "12.1"
+    ]
+  },
+  "digest": {
+    "algorithm": "sha256",
+    "files": {
+      "__init__.py": "0q6TMj65una1LdyHLC291Dwy2VzrBbAD8Wm8MEdmqwA=",
+      "_ops.py": "iqY+HoE4nGHqvVlEZTwEAcyu1/w+QrjwuuW9sdTI2uk=",
+      "_weight_only_ffn_cuda_38d37b6.abi3.so": "bZgBCn5rvouSVonuYPECZRmDXS2EHW8aWPt5InFNJ00=",
+      "weight_only_ffn/__init__.py": "DFYPlrhXwYjEqCl/8n0SmWGZV8NFml5DPhMjKfv98GY="
+    }
+  },
+  "provenance": {
+    "kernel-builder": {
+      "version": "0.17.0-dev0",
+      "sha": "19aaa6421e674e9fecc352bbae6eab81d19a6bf4",
+      "dirty": false
+    },
+    "kernel": {
+      "sha": "38d37b67acfe940c99d7b3db721fa67163b82cec",
+      "dirty": false
+    }
+  }
+}

build/torch211-cxx11-cu130-x86_64-linux/weight_only_ffn/__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/__init__.py

@@ -0,0 +1,261 @@
+"""Weight-only BF16-activation linear and FFN regions."""
+
+from __future__ import annotations
+
+from typing import Optional
+
+import torch
+
+from ._ops import add_op_namespace_prefix, ops
+
+
+def w4_sfb_size_bytes(rows: int, cols: int) -> int:
+    if rows <= 0 or cols <= 0 or cols % 64:
+        raise ValueError("rows must be positive and cols must be divisible by 64")
+    return ((int(rows) + 127) // 128) * (((int(cols) // 16) + 3) // 4) * 512
+
+
+@torch.library.register_fake(add_op_namespace_prefix("quantize_w4_weight_bf16"))
+def _quantize_w4_fake(weight: torch.Tensor, packed: torch.Tensor, sfb: torch.Tensor) -> None:
+    return None
+
+
+@torch.library.register_fake(add_op_namespace_prefix("dequantize_w4_weight_bf16"))
+def _dequantize_w4_fake(packed: torch.Tensor, sfb: torch.Tensor, weight: torch.Tensor) -> None:
+    return None
+
+
+@torch.library.register_fake(add_op_namespace_prefix("w4a16_linear_bf16"))
+def _w4_linear_fake(
+    x: torch.Tensor,
+    packed: torch.Tensor,
+    sfb: torch.Tensor,
+    alpha: float,
+    variant: int,
+    out: torch.Tensor,
+) -> None:
+    return None
+
+
+@torch.library.register_fake(add_op_namespace_prefix("quantize_w8_weight_bf16"))
+def _quantize_w8_fake(weight: torch.Tensor, quantized: torch.Tensor, scales: torch.Tensor) -> None:
+    return None
+
+
+@torch.library.register_fake(add_op_namespace_prefix("dequantize_w8_weight_bf16"))
+def _dequantize_w8_fake(quantized: torch.Tensor, scales: torch.Tensor, weight: torch.Tensor) -> None:
+    return None
+
+
+@torch.library.register_fake(add_op_namespace_prefix("w8a16_linear_bf16"))
+def _w8_linear_fake(
+    x: torch.Tensor,
+    quantized: torch.Tensor,
+    scales: torch.Tensor,
+    variant: int,
+    out: torch.Tensor,
+) -> None:
+    return None
+
+
+@torch.library.register_fake(add_op_namespace_prefix("w4a16_gated_ffn_bf16"))
+def _w4_gated_fake(
+    x, gu_w, gu_s, dn_w, dn_s, gu_b, dn_b, gelu, gu_alpha, dn_alpha,
+    variant, gu, hidden, out,
+) -> None:
+    return None
+
+
+@torch.library.register_fake(add_op_namespace_prefix("w8a16_gated_ffn_bf16"))
+def _w8_gated_fake(
+    x, gu_w, gu_s, dn_w, dn_s, gu_b, dn_b, gelu, variant, gu, hidden, out,
+) -> None:
+    return None
+
+
+@torch.library.register_fake(add_op_namespace_prefix("w4a16_gelu_ffn_bf16"))
+def _w4_gelu_fake(
+    x, up_w, up_s, dn_w, dn_s, up_b, dn_b, up_alpha, dn_alpha,
+    variant, up, hidden, out,
+) -> None:
+    return None
+
+
+@torch.library.register_fake(add_op_namespace_prefix("w8a16_gelu_ffn_bf16"))
+def _w8_gelu_fake(
+    x, up_w, up_s, dn_w, dn_s, up_b, dn_b, variant, up, hidden, out,
+) -> None:
+    return None
+
+
+def quantize_w4_weight_bf16(weight: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor]:
+    """Pack a static BF16 weight matrix into NVFP4 data and SFB scales."""
+    n, k = weight.shape
+    packed = torch.empty((n, k // 2), device=weight.device, dtype=torch.uint8)
+    sfb = torch.empty((w4_sfb_size_bytes(n, k),), device=weight.device, dtype=torch.uint8)
+    ops.quantize_w4_weight_bf16(weight, packed, sfb)
+    return packed, sfb
+
+
+def dequantize_w4_weight_bf16(
+    packed: torch.Tensor, sfb: torch.Tensor, *, cols: int
+) -> torch.Tensor:
+    weight = torch.empty((packed.shape[0], int(cols)), device=packed.device, dtype=torch.bfloat16)
+    ops.dequantize_w4_weight_bf16(packed, sfb, weight)
+    return weight
+
+
+def quantize_w8_weight_bf16(weight: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor]:
+    """Symmetrically quantize a static BF16 weight matrix per output channel."""
+    quantized = torch.empty_like(weight, dtype=torch.int8)
+    scales = torch.empty((weight.shape[0],), device=weight.device, dtype=torch.float32)
+    ops.quantize_w8_weight_bf16(weight, quantized, scales)
+    return quantized, scales
+
+
+def dequantize_w8_weight_bf16(
+    quantized: torch.Tensor, scales: torch.Tensor
+) -> torch.Tensor:
+    weight = torch.empty_like(quantized, dtype=torch.bfloat16)
+    ops.dequantize_w8_weight_bf16(quantized, scales, weight)
+    return weight
+
+
+def w4a16_linear_bf16(
+    x: torch.Tensor,
+    weight_packed: torch.Tensor,
+    weight_sfb: torch.Tensor,
+    *,
+    alpha: float = 1.0,
+    variant: int = 0,
+    out: Optional[torch.Tensor] = None,
+) -> torch.Tensor:
+    if out is None:
+        out = torch.empty((x.shape[0], weight_packed.shape[0]), device=x.device, dtype=torch.bfloat16)
+    ops.w4a16_linear_bf16(x, weight_packed, weight_sfb, float(alpha), int(variant), out)
+    return out
+
+
+def w8a16_linear_bf16(
+    x: torch.Tensor,
+    weight_int8: torch.Tensor,
+    weight_scale: torch.Tensor,
+    *,
+    variant: int = 0,
+    out: Optional[torch.Tensor] = None,
+) -> torch.Tensor:
+    if out is None:
+        out = torch.empty((x.shape[0], weight_int8.shape[0]), device=x.device, dtype=torch.bfloat16)
+    ops.w8a16_linear_bf16(x, weight_int8, weight_scale, int(variant), out)
+    return out
+
+
+def _gated_workspace(x: torch.Tensor, gate_up_rows: int, out_rows: int, workspace, out):
+    hidden_size = gate_up_rows // 2
+    if workspace is None:
+        gate_up = torch.empty((x.shape[0], gate_up_rows), device=x.device, dtype=torch.bfloat16)
+        hidden = torch.empty((x.shape[0], hidden_size), device=x.device, dtype=torch.bfloat16)
+    else:
+        gate_up, hidden = workspace
+    if out is None:
+        out = torch.empty((x.shape[0], out_rows), device=x.device, dtype=torch.bfloat16)
+    return gate_up, hidden, out
+
+
+def _gelu_workspace(x: torch.Tensor, hidden_size: int, out_rows: int, workspace, out):
+    if workspace is None:
+        up = torch.empty((x.shape[0], hidden_size), device=x.device, dtype=torch.bfloat16)
+        hidden = torch.empty_like(up)
+    else:
+        up, hidden = workspace
+    if out is None:
+        out = torch.empty((x.shape[0], out_rows), device=x.device, dtype=torch.bfloat16)
+    return up, hidden, out
+
+
+def _w4_gated(
+    x, gate_up_weight, gate_up_sfb, down_weight, down_sfb, *, gelu,
+    gate_up_bias=None, down_bias=None, gate_up_alpha=1.0, down_alpha=1.0,
+    variant=0, workspace=None, out=None,
+):
+    gate_up, hidden, out = _gated_workspace(
+        x, gate_up_weight.shape[0], down_weight.shape[0], workspace, out
+    )
+    ops.w4a16_gated_ffn_bf16(
+        x, gate_up_weight, gate_up_sfb, down_weight, down_sfb,
+        gate_up_bias, down_bias, bool(gelu), float(gate_up_alpha),
+        float(down_alpha), int(variant), gate_up, hidden, out,
+    )
+    return out
+
+
+def w4a16_swiglu_ffn_bf16(x, gate_up_weight, gate_up_sfb, down_weight, down_sfb, **kwargs):
+    return _w4_gated(x, gate_up_weight, gate_up_sfb, down_weight, down_sfb, gelu=False, **kwargs)
+
+
+def w4a16_geglu_ffn_bf16(x, gate_up_weight, gate_up_sfb, down_weight, down_sfb, **kwargs):
+    return _w4_gated(x, gate_up_weight, gate_up_sfb, down_weight, down_sfb, gelu=True, **kwargs)
+
+
+def _w8_gated(
+    x, gate_up_weight, gate_up_scale, down_weight, down_scale, *, gelu,
+    gate_up_bias=None, down_bias=None, variant=0, workspace=None, out=None,
+):
+    gate_up, hidden, out = _gated_workspace(
+        x, gate_up_weight.shape[0], down_weight.shape[0], workspace, out
+    )
+    ops.w8a16_gated_ffn_bf16(
+        x, gate_up_weight, gate_up_scale, down_weight, down_scale,
+        gate_up_bias, down_bias, bool(gelu), int(variant), gate_up, hidden, out,
+    )
+    return out
+
+
+def w8a16_swiglu_ffn_bf16(x, gate_up_weight, gate_up_scale, down_weight, down_scale, **kwargs):
+    return _w8_gated(x, gate_up_weight, gate_up_scale, down_weight, down_scale, gelu=False, **kwargs)
+
+
+def w8a16_geglu_ffn_bf16(x, gate_up_weight, gate_up_scale, down_weight, down_scale, **kwargs):
+    return _w8_gated(x, gate_up_weight, gate_up_scale, down_weight, down_scale, gelu=True, **kwargs)
+
+
+def w4a16_gelu_ffn_bf16(
+    x, up_weight, up_sfb, down_weight, down_sfb, *, up_bias=None,
+    down_bias=None, up_alpha=1.0, down_alpha=1.0, variant=0,
+    workspace=None, out=None,
+):
+    up, hidden, out = _gelu_workspace(x, up_weight.shape[0], down_weight.shape[0], workspace, out)
+    ops.w4a16_gelu_ffn_bf16(
+        x, up_weight, up_sfb, down_weight, down_sfb, up_bias, down_bias,
+        float(up_alpha), float(down_alpha), int(variant), up, hidden, out,
+    )
+    return out
+
+
+def w8a16_gelu_ffn_bf16(
+    x, up_weight, up_scale, down_weight, down_scale, *, up_bias=None,
+    down_bias=None, variant=0, workspace=None, out=None,
+):
+    up, hidden, out = _gelu_workspace(x, up_weight.shape[0], down_weight.shape[0], workspace, out)
+    ops.w8a16_gelu_ffn_bf16(
+        x, up_weight, up_scale, down_weight, down_scale, up_bias, down_bias,
+        int(variant), up, hidden, out,
+    )
+    return out
+
+
+__all__ = [
+    "dequantize_w4_weight_bf16",
+    "dequantize_w8_weight_bf16",
+    "quantize_w4_weight_bf16",
+    "quantize_w8_weight_bf16",
+    "w4_sfb_size_bytes",
+    "w4a16_geglu_ffn_bf16",
+    "w4a16_gelu_ffn_bf16",
+    "w4a16_linear_bf16",
+    "w4a16_swiglu_ffn_bf16",
+    "w8a16_geglu_ffn_bf16",
+    "w8a16_gelu_ffn_bf16",
+    "w8a16_linear_bf16",
+    "w8a16_swiglu_ffn_bf16",
+]

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

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

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

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

build/torch212-cxx11-cu130-x86_64-linux/metadata.json

@@ -0,0 +1,35 @@
+{
+  "name": "weight-only-ffn",
+  "id": "_weight_only_ffn_cuda_38d37b6",
+  "version": 1,
+  "license": "Apache-2.0",
+  "python-depends": [],
+  "backend": {
+    "type": "cuda",
+    "archs": [
+      "11.0a",
+      "12.0a",
+      "12.1"
+    ]
+  },
+  "digest": {
+    "algorithm": "sha256",
+    "files": {
+      "__init__.py": "0q6TMj65una1LdyHLC291Dwy2VzrBbAD8Wm8MEdmqwA=",
+      "_ops.py": "iqY+HoE4nGHqvVlEZTwEAcyu1/w+QrjwuuW9sdTI2uk=",
+      "_weight_only_ffn_cuda_38d37b6.abi3.so": "689f8vMnCgMncSu4yRemqgNhQPeiwGWEndiXpYiYJhs=",
+      "weight_only_ffn/__init__.py": "DFYPlrhXwYjEqCl/8n0SmWGZV8NFml5DPhMjKfv98GY="
+    }
+  },
+  "provenance": {
+    "kernel-builder": {
+      "version": "0.17.0-dev0",
+      "sha": "19aaa6421e674e9fecc352bbae6eab81d19a6bf4",
+      "dirty": false
+    },
+    "kernel": {
+      "sha": "38d37b67acfe940c99d7b3db721fa67163b82cec",
+      "dirty": false
+    }
+  }
+}

build/torch212-cxx11-cu130-x86_64-linux/weight_only_ffn/__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/__init__.py

@@ -0,0 +1,261 @@
+"""Weight-only BF16-activation linear and FFN regions."""
+
+from __future__ import annotations
+
+from typing import Optional
+
+import torch
+
+from ._ops import add_op_namespace_prefix, ops
+
+
+def w4_sfb_size_bytes(rows: int, cols: int) -> int:
+    if rows <= 0 or cols <= 0 or cols % 64:
+        raise ValueError("rows must be positive and cols must be divisible by 64")
+    return ((int(rows) + 127) // 128) * (((int(cols) // 16) + 3) // 4) * 512
+
+
+@torch.library.register_fake(add_op_namespace_prefix("quantize_w4_weight_bf16"))
+def _quantize_w4_fake(weight: torch.Tensor, packed: torch.Tensor, sfb: torch.Tensor) -> None:
+    return None
+
+
+@torch.library.register_fake(add_op_namespace_prefix("dequantize_w4_weight_bf16"))
+def _dequantize_w4_fake(packed: torch.Tensor, sfb: torch.Tensor, weight: torch.Tensor) -> None:
+    return None
+
+
+@torch.library.register_fake(add_op_namespace_prefix("w4a16_linear_bf16"))
+def _w4_linear_fake(
+    x: torch.Tensor,
+    packed: torch.Tensor,
+    sfb: torch.Tensor,
+    alpha: float,
+    variant: int,
+    out: torch.Tensor,
+) -> None:
+    return None
+
+
+@torch.library.register_fake(add_op_namespace_prefix("quantize_w8_weight_bf16"))
+def _quantize_w8_fake(weight: torch.Tensor, quantized: torch.Tensor, scales: torch.Tensor) -> None:
+    return None
+
+
+@torch.library.register_fake(add_op_namespace_prefix("dequantize_w8_weight_bf16"))
+def _dequantize_w8_fake(quantized: torch.Tensor, scales: torch.Tensor, weight: torch.Tensor) -> None:
+    return None
+
+
+@torch.library.register_fake(add_op_namespace_prefix("w8a16_linear_bf16"))
+def _w8_linear_fake(
+    x: torch.Tensor,
+    quantized: torch.Tensor,
+    scales: torch.Tensor,
+    variant: int,
+    out: torch.Tensor,
+) -> None:
+    return None
+
+
+@torch.library.register_fake(add_op_namespace_prefix("w4a16_gated_ffn_bf16"))
+def _w4_gated_fake(
+    x, gu_w, gu_s, dn_w, dn_s, gu_b, dn_b, gelu, gu_alpha, dn_alpha,
+    variant, gu, hidden, out,
+) -> None:
+    return None
+
+
+@torch.library.register_fake(add_op_namespace_prefix("w8a16_gated_ffn_bf16"))
+def _w8_gated_fake(
+    x, gu_w, gu_s, dn_w, dn_s, gu_b, dn_b, gelu, variant, gu, hidden, out,
+) -> None:
+    return None
+
+
+@torch.library.register_fake(add_op_namespace_prefix("w4a16_gelu_ffn_bf16"))
+def _w4_gelu_fake(
+    x, up_w, up_s, dn_w, dn_s, up_b, dn_b, up_alpha, dn_alpha,
+    variant, up, hidden, out,
+) -> None:
+    return None
+
+
+@torch.library.register_fake(add_op_namespace_prefix("w8a16_gelu_ffn_bf16"))
+def _w8_gelu_fake(
+    x, up_w, up_s, dn_w, dn_s, up_b, dn_b, variant, up, hidden, out,
+) -> None:
+    return None
+
+
+def quantize_w4_weight_bf16(weight: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor]:
+    """Pack a static BF16 weight matrix into NVFP4 data and SFB scales."""
+    n, k = weight.shape
+    packed = torch.empty((n, k // 2), device=weight.device, dtype=torch.uint8)
+    sfb = torch.empty((w4_sfb_size_bytes(n, k),), device=weight.device, dtype=torch.uint8)
+    ops.quantize_w4_weight_bf16(weight, packed, sfb)
+    return packed, sfb
+
+
+def dequantize_w4_weight_bf16(
+    packed: torch.Tensor, sfb: torch.Tensor, *, cols: int
+) -> torch.Tensor:
+    weight = torch.empty((packed.shape[0], int(cols)), device=packed.device, dtype=torch.bfloat16)
+    ops.dequantize_w4_weight_bf16(packed, sfb, weight)
+    return weight
+
+
+def quantize_w8_weight_bf16(weight: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor]:
+    """Symmetrically quantize a static BF16 weight matrix per output channel."""
+    quantized = torch.empty_like(weight, dtype=torch.int8)
+    scales = torch.empty((weight.shape[0],), device=weight.device, dtype=torch.float32)
+    ops.quantize_w8_weight_bf16(weight, quantized, scales)
+    return quantized, scales
+
+
+def dequantize_w8_weight_bf16(
+    quantized: torch.Tensor, scales: torch.Tensor
+) -> torch.Tensor:
+    weight = torch.empty_like(quantized, dtype=torch.bfloat16)
+    ops.dequantize_w8_weight_bf16(quantized, scales, weight)
+    return weight
+
+
+def w4a16_linear_bf16(
+    x: torch.Tensor,
+    weight_packed: torch.Tensor,
+    weight_sfb: torch.Tensor,
+    *,
+    alpha: float = 1.0,
+    variant: int = 0,
+    out: Optional[torch.Tensor] = None,
+) -> torch.Tensor:
+    if out is None:
+        out = torch.empty((x.shape[0], weight_packed.shape[0]), device=x.device, dtype=torch.bfloat16)
+    ops.w4a16_linear_bf16(x, weight_packed, weight_sfb, float(alpha), int(variant), out)
+    return out
+
+
+def w8a16_linear_bf16(
+    x: torch.Tensor,
+    weight_int8: torch.Tensor,
+    weight_scale: torch.Tensor,
+    *,
+    variant: int = 0,
+    out: Optional[torch.Tensor] = None,
+) -> torch.Tensor:
+    if out is None:
+        out = torch.empty((x.shape[0], weight_int8.shape[0]), device=x.device, dtype=torch.bfloat16)
+    ops.w8a16_linear_bf16(x, weight_int8, weight_scale, int(variant), out)
+    return out
+
+
+def _gated_workspace(x: torch.Tensor, gate_up_rows: int, out_rows: int, workspace, out):
+    hidden_size = gate_up_rows // 2
+    if workspace is None:
+        gate_up = torch.empty((x.shape[0], gate_up_rows), device=x.device, dtype=torch.bfloat16)
+        hidden = torch.empty((x.shape[0], hidden_size), device=x.device, dtype=torch.bfloat16)
+    else:
+        gate_up, hidden = workspace
+    if out is None:
+        out = torch.empty((x.shape[0], out_rows), device=x.device, dtype=torch.bfloat16)
+    return gate_up, hidden, out
+
+
+def _gelu_workspace(x: torch.Tensor, hidden_size: int, out_rows: int, workspace, out):
+    if workspace is None:
+        up = torch.empty((x.shape[0], hidden_size), device=x.device, dtype=torch.bfloat16)
+        hidden = torch.empty_like(up)
+    else:
+        up, hidden = workspace
+    if out is None:
+        out = torch.empty((x.shape[0], out_rows), device=x.device, dtype=torch.bfloat16)
+    return up, hidden, out
+
+
+def _w4_gated(
+    x, gate_up_weight, gate_up_sfb, down_weight, down_sfb, *, gelu,
+    gate_up_bias=None, down_bias=None, gate_up_alpha=1.0, down_alpha=1.0,
+    variant=0, workspace=None, out=None,
+):
+    gate_up, hidden, out = _gated_workspace(
+        x, gate_up_weight.shape[0], down_weight.shape[0], workspace, out
+    )
+    ops.w4a16_gated_ffn_bf16(
+        x, gate_up_weight, gate_up_sfb, down_weight, down_sfb,
+        gate_up_bias, down_bias, bool(gelu), float(gate_up_alpha),
+        float(down_alpha), int(variant), gate_up, hidden, out,
+    )
+    return out
+
+
+def w4a16_swiglu_ffn_bf16(x, gate_up_weight, gate_up_sfb, down_weight, down_sfb, **kwargs):
+    return _w4_gated(x, gate_up_weight, gate_up_sfb, down_weight, down_sfb, gelu=False, **kwargs)
+
+
+def w4a16_geglu_ffn_bf16(x, gate_up_weight, gate_up_sfb, down_weight, down_sfb, **kwargs):
+    return _w4_gated(x, gate_up_weight, gate_up_sfb, down_weight, down_sfb, gelu=True, **kwargs)
+
+
+def _w8_gated(
+    x, gate_up_weight, gate_up_scale, down_weight, down_scale, *, gelu,
+    gate_up_bias=None, down_bias=None, variant=0, workspace=None, out=None,
+):
+    gate_up, hidden, out = _gated_workspace(
+        x, gate_up_weight.shape[0], down_weight.shape[0], workspace, out
+    )
+    ops.w8a16_gated_ffn_bf16(
+        x, gate_up_weight, gate_up_scale, down_weight, down_scale,
+        gate_up_bias, down_bias, bool(gelu), int(variant), gate_up, hidden, out,
+    )
+    return out
+
+
+def w8a16_swiglu_ffn_bf16(x, gate_up_weight, gate_up_scale, down_weight, down_scale, **kwargs):
+    return _w8_gated(x, gate_up_weight, gate_up_scale, down_weight, down_scale, gelu=False, **kwargs)
+
+
+def w8a16_geglu_ffn_bf16(x, gate_up_weight, gate_up_scale, down_weight, down_scale, **kwargs):
+    return _w8_gated(x, gate_up_weight, gate_up_scale, down_weight, down_scale, gelu=True, **kwargs)
+
+
+def w4a16_gelu_ffn_bf16(
+    x, up_weight, up_sfb, down_weight, down_sfb, *, up_bias=None,
+    down_bias=None, up_alpha=1.0, down_alpha=1.0, variant=0,
+    workspace=None, out=None,
+):
+    up, hidden, out = _gelu_workspace(x, up_weight.shape[0], down_weight.shape[0], workspace, out)
+    ops.w4a16_gelu_ffn_bf16(
+        x, up_weight, up_sfb, down_weight, down_sfb, up_bias, down_bias,
+        float(up_alpha), float(down_alpha), int(variant), up, hidden, out,
+    )
+    return out
+
+
+def w8a16_gelu_ffn_bf16(
+    x, up_weight, up_scale, down_weight, down_scale, *, up_bias=None,
+    down_bias=None, variant=0, workspace=None, out=None,
+):
+    up, hidden, out = _gelu_workspace(x, up_weight.shape[0], down_weight.shape[0], workspace, out)
+    ops.w8a16_gelu_ffn_bf16(
+        x, up_weight, up_scale, down_weight, down_scale, up_bias, down_bias,
+        int(variant), up, hidden, out,
+    )
+    return out
+
+
+__all__ = [
+    "dequantize_w4_weight_bf16",
+    "dequantize_w8_weight_bf16",
+    "quantize_w4_weight_bf16",
+    "quantize_w8_weight_bf16",
+    "w4_sfb_size_bytes",
+    "w4a16_geglu_ffn_bf16",
+    "w4a16_gelu_ffn_bf16",
+    "w4a16_linear_bf16",
+    "w4a16_swiglu_ffn_bf16",
+    "w8a16_geglu_ffn_bf16",
+    "w8a16_gelu_ffn_bf16",
+    "w8a16_linear_bf16",
+    "w8a16_swiglu_ffn_bf16",
+]

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

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

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

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

build/torch212-cxx11-cu132-x86_64-linux/metadata.json

@@ -0,0 +1,35 @@
+{
+  "name": "weight-only-ffn",
+  "id": "_weight_only_ffn_cuda_38d37b6",
+  "version": 1,
+  "license": "Apache-2.0",
+  "python-depends": [],
+  "backend": {
+    "type": "cuda",
+    "archs": [
+      "11.0a",
+      "12.0a",
+      "12.1"
+    ]
+  },
+  "digest": {
+    "algorithm": "sha256",
+    "files": {
+      "__init__.py": "0q6TMj65una1LdyHLC291Dwy2VzrBbAD8Wm8MEdmqwA=",
+      "_ops.py": "iqY+HoE4nGHqvVlEZTwEAcyu1/w+QrjwuuW9sdTI2uk=",
+      "_weight_only_ffn_cuda_38d37b6.abi3.so": "jvVCDKovIwl+Po+XdyUMm7uYYY/nZJPUkuXLWl8JM2w=",
+      "weight_only_ffn/__init__.py": "DFYPlrhXwYjEqCl/8n0SmWGZV8NFml5DPhMjKfv98GY="
+    }
+  },
+  "provenance": {
+    "kernel-builder": {
+      "version": "0.17.0-dev0",
+      "sha": "19aaa6421e674e9fecc352bbae6eab81d19a6bf4",
+      "dirty": false
+    },
+    "kernel": {
+      "sha": "38d37b67acfe940c99d7b3db721fa67163b82cec",
+      "dirty": false
+    }
+  }
+}

build/torch212-cxx11-cu132-x86_64-linux/weight_only_ffn/__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/torch213-cxx11-cu130-aarch64-linux/__init__.py

@@ -0,0 +1,261 @@
+"""Weight-only BF16-activation linear and FFN regions."""
+
+from __future__ import annotations
+
+from typing import Optional
+
+import torch
+
+from ._ops import add_op_namespace_prefix, ops
+
+
+def w4_sfb_size_bytes(rows: int, cols: int) -> int:
+    if rows <= 0 or cols <= 0 or cols % 64:
+        raise ValueError("rows must be positive and cols must be divisible by 64")
+    return ((int(rows) + 127) // 128) * (((int(cols) // 16) + 3) // 4) * 512
+
+
+@torch.library.register_fake(add_op_namespace_prefix("quantize_w4_weight_bf16"))
+def _quantize_w4_fake(weight: torch.Tensor, packed: torch.Tensor, sfb: torch.Tensor) -> None:
+    return None
+
+
+@torch.library.register_fake(add_op_namespace_prefix("dequantize_w4_weight_bf16"))
+def _dequantize_w4_fake(packed: torch.Tensor, sfb: torch.Tensor, weight: torch.Tensor) -> None:
+    return None
+
+
+@torch.library.register_fake(add_op_namespace_prefix("w4a16_linear_bf16"))
+def _w4_linear_fake(
+    x: torch.Tensor,
+    packed: torch.Tensor,
+    sfb: torch.Tensor,
+    alpha: float,
+    variant: int,
+    out: torch.Tensor,
+) -> None:
+    return None
+
+
+@torch.library.register_fake(add_op_namespace_prefix("quantize_w8_weight_bf16"))
+def _quantize_w8_fake(weight: torch.Tensor, quantized: torch.Tensor, scales: torch.Tensor) -> None:
+    return None
+
+
+@torch.library.register_fake(add_op_namespace_prefix("dequantize_w8_weight_bf16"))
+def _dequantize_w8_fake(quantized: torch.Tensor, scales: torch.Tensor, weight: torch.Tensor) -> None:
+    return None
+
+
+@torch.library.register_fake(add_op_namespace_prefix("w8a16_linear_bf16"))
+def _w8_linear_fake(
+    x: torch.Tensor,
+    quantized: torch.Tensor,
+    scales: torch.Tensor,
+    variant: int,
+    out: torch.Tensor,
+) -> None:
+    return None
+
+
+@torch.library.register_fake(add_op_namespace_prefix("w4a16_gated_ffn_bf16"))
+def _w4_gated_fake(
+    x, gu_w, gu_s, dn_w, dn_s, gu_b, dn_b, gelu, gu_alpha, dn_alpha,
+    variant, gu, hidden, out,
+) -> None:
+    return None
+
+
+@torch.library.register_fake(add_op_namespace_prefix("w8a16_gated_ffn_bf16"))
+def _w8_gated_fake(
+    x, gu_w, gu_s, dn_w, dn_s, gu_b, dn_b, gelu, variant, gu, hidden, out,
+) -> None:
+    return None
+
+
+@torch.library.register_fake(add_op_namespace_prefix("w4a16_gelu_ffn_bf16"))
+def _w4_gelu_fake(
+    x, up_w, up_s, dn_w, dn_s, up_b, dn_b, up_alpha, dn_alpha,
+    variant, up, hidden, out,
+) -> None:
+    return None
+
+
+@torch.library.register_fake(add_op_namespace_prefix("w8a16_gelu_ffn_bf16"))
+def _w8_gelu_fake(
+    x, up_w, up_s, dn_w, dn_s, up_b, dn_b, variant, up, hidden, out,
+) -> None:
+    return None
+
+
+def quantize_w4_weight_bf16(weight: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor]:
+    """Pack a static BF16 weight matrix into NVFP4 data and SFB scales."""
+    n, k = weight.shape
+    packed = torch.empty((n, k // 2), device=weight.device, dtype=torch.uint8)
+    sfb = torch.empty((w4_sfb_size_bytes(n, k),), device=weight.device, dtype=torch.uint8)
+    ops.quantize_w4_weight_bf16(weight, packed, sfb)
+    return packed, sfb
+
+
+def dequantize_w4_weight_bf16(
+    packed: torch.Tensor, sfb: torch.Tensor, *, cols: int
+) -> torch.Tensor:
+    weight = torch.empty((packed.shape[0], int(cols)), device=packed.device, dtype=torch.bfloat16)
+    ops.dequantize_w4_weight_bf16(packed, sfb, weight)
+    return weight
+
+
+def quantize_w8_weight_bf16(weight: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor]:
+    """Symmetrically quantize a static BF16 weight matrix per output channel."""
+    quantized = torch.empty_like(weight, dtype=torch.int8)
+    scales = torch.empty((weight.shape[0],), device=weight.device, dtype=torch.float32)
+    ops.quantize_w8_weight_bf16(weight, quantized, scales)
+    return quantized, scales
+
+
+def dequantize_w8_weight_bf16(
+    quantized: torch.Tensor, scales: torch.Tensor
+) -> torch.Tensor:
+    weight = torch.empty_like(quantized, dtype=torch.bfloat16)
+    ops.dequantize_w8_weight_bf16(quantized, scales, weight)
+    return weight
+
+
+def w4a16_linear_bf16(
+    x: torch.Tensor,
+    weight_packed: torch.Tensor,
+    weight_sfb: torch.Tensor,
+    *,
+    alpha: float = 1.0,
+    variant: int = 0,
+    out: Optional[torch.Tensor] = None,
+) -> torch.Tensor:
+    if out is None:
+        out = torch.empty((x.shape[0], weight_packed.shape[0]), device=x.device, dtype=torch.bfloat16)
+    ops.w4a16_linear_bf16(x, weight_packed, weight_sfb, float(alpha), int(variant), out)
+    return out
+
+
+def w8a16_linear_bf16(
+    x: torch.Tensor,
+    weight_int8: torch.Tensor,
+    weight_scale: torch.Tensor,
+    *,
+    variant: int = 0,
+    out: Optional[torch.Tensor] = None,
+) -> torch.Tensor:
+    if out is None:
+        out = torch.empty((x.shape[0], weight_int8.shape[0]), device=x.device, dtype=torch.bfloat16)
+    ops.w8a16_linear_bf16(x, weight_int8, weight_scale, int(variant), out)
+    return out
+
+
+def _gated_workspace(x: torch.Tensor, gate_up_rows: int, out_rows: int, workspace, out):
+    hidden_size = gate_up_rows // 2
+    if workspace is None:
+        gate_up = torch.empty((x.shape[0], gate_up_rows), device=x.device, dtype=torch.bfloat16)
+        hidden = torch.empty((x.shape[0], hidden_size), device=x.device, dtype=torch.bfloat16)
+    else:
+        gate_up, hidden = workspace
+    if out is None:
+        out = torch.empty((x.shape[0], out_rows), device=x.device, dtype=torch.bfloat16)
+    return gate_up, hidden, out
+
+
+def _gelu_workspace(x: torch.Tensor, hidden_size: int, out_rows: int, workspace, out):
+    if workspace is None:
+        up = torch.empty((x.shape[0], hidden_size), device=x.device, dtype=torch.bfloat16)
+        hidden = torch.empty_like(up)
+    else:
+        up, hidden = workspace
+    if out is None:
+        out = torch.empty((x.shape[0], out_rows), device=x.device, dtype=torch.bfloat16)
+    return up, hidden, out
+
+
+def _w4_gated(
+    x, gate_up_weight, gate_up_sfb, down_weight, down_sfb, *, gelu,
+    gate_up_bias=None, down_bias=None, gate_up_alpha=1.0, down_alpha=1.0,
+    variant=0, workspace=None, out=None,
+):
+    gate_up, hidden, out = _gated_workspace(
+        x, gate_up_weight.shape[0], down_weight.shape[0], workspace, out
+    )
+    ops.w4a16_gated_ffn_bf16(
+        x, gate_up_weight, gate_up_sfb, down_weight, down_sfb,
+        gate_up_bias, down_bias, bool(gelu), float(gate_up_alpha),
+        float(down_alpha), int(variant), gate_up, hidden, out,
+    )
+    return out
+
+
+def w4a16_swiglu_ffn_bf16(x, gate_up_weight, gate_up_sfb, down_weight, down_sfb, **kwargs):
+    return _w4_gated(x, gate_up_weight, gate_up_sfb, down_weight, down_sfb, gelu=False, **kwargs)
+
+
+def w4a16_geglu_ffn_bf16(x, gate_up_weight, gate_up_sfb, down_weight, down_sfb, **kwargs):
+    return _w4_gated(x, gate_up_weight, gate_up_sfb, down_weight, down_sfb, gelu=True, **kwargs)
+
+
+def _w8_gated(
+    x, gate_up_weight, gate_up_scale, down_weight, down_scale, *, gelu,
+    gate_up_bias=None, down_bias=None, variant=0, workspace=None, out=None,
+):
+    gate_up, hidden, out = _gated_workspace(
+        x, gate_up_weight.shape[0], down_weight.shape[0], workspace, out
+    )
+    ops.w8a16_gated_ffn_bf16(
+        x, gate_up_weight, gate_up_scale, down_weight, down_scale,
+        gate_up_bias, down_bias, bool(gelu), int(variant), gate_up, hidden, out,
+    )
+    return out
+
+
+def w8a16_swiglu_ffn_bf16(x, gate_up_weight, gate_up_scale, down_weight, down_scale, **kwargs):
+    return _w8_gated(x, gate_up_weight, gate_up_scale, down_weight, down_scale, gelu=False, **kwargs)
+
+
+def w8a16_geglu_ffn_bf16(x, gate_up_weight, gate_up_scale, down_weight, down_scale, **kwargs):
+    return _w8_gated(x, gate_up_weight, gate_up_scale, down_weight, down_scale, gelu=True, **kwargs)
+
+
+def w4a16_gelu_ffn_bf16(
+    x, up_weight, up_sfb, down_weight, down_sfb, *, up_bias=None,
+    down_bias=None, up_alpha=1.0, down_alpha=1.0, variant=0,
+    workspace=None, out=None,
+):
+    up, hidden, out = _gelu_workspace(x, up_weight.shape[0], down_weight.shape[0], workspace, out)
+    ops.w4a16_gelu_ffn_bf16(
+        x, up_weight, up_sfb, down_weight, down_sfb, up_bias, down_bias,
+        float(up_alpha), float(down_alpha), int(variant), up, hidden, out,
+    )
+    return out
+
+
+def w8a16_gelu_ffn_bf16(
+    x, up_weight, up_scale, down_weight, down_scale, *, up_bias=None,
+    down_bias=None, variant=0, workspace=None, out=None,
+):
+    up, hidden, out = _gelu_workspace(x, up_weight.shape[0], down_weight.shape[0], workspace, out)
+    ops.w8a16_gelu_ffn_bf16(
+        x, up_weight, up_scale, down_weight, down_scale, up_bias, down_bias,
+        int(variant), up, hidden, out,
+    )
+    return out
+
+
+__all__ = [
+    "dequantize_w4_weight_bf16",
+    "dequantize_w8_weight_bf16",
+    "quantize_w4_weight_bf16",
+    "quantize_w8_weight_bf16",
+    "w4_sfb_size_bytes",
+    "w4a16_geglu_ffn_bf16",
+    "w4a16_gelu_ffn_bf16",
+    "w4a16_linear_bf16",
+    "w4a16_swiglu_ffn_bf16",
+    "w8a16_geglu_ffn_bf16",
+    "w8a16_gelu_ffn_bf16",
+    "w8a16_linear_bf16",
+    "w8a16_swiglu_ffn_bf16",
+]

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

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

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

@@ -0,0 +1,32 @@
+{
+  "name": "weight-only-ffn",
+  "id": "weight_only_ffn_source_test",
+  "version": 1,
+  "license": "Apache-2.0",
+  "python-depends": [],
+  "backend": {
+    "type": "cuda",
+    "archs": [
+      "11.0a"
+    ]
+  },
+  "digest": {
+    "algorithm": "sha256",
+    "files": {
+      "__init__.py": "0q6TMj65una1LdyHLC291Dwy2VzrBbAD8Wm8MEdmqwA=",
+      "weight_only_ffn_source_test.abi3.so": "4bTmrdGoTStyhYzGGZQbYZmrMo/Hbq9EVXN9i//6Q3Q=",
+      "_ops.py": "sIcm8AtiXP5iJMFaDNUMdBLwzQULjS0N7fOxWIFF1DQ=",
+      "weight_only_ffn/__init__.py": "v6p5XMfQzddhi1fLSAw4HX9CyS0rQsidvu9VsT01xi4="
+    }
+  },
+  "provenance": {
+    "kernel": {
+      "sha": "456d297",
+      "dirty": false
+    },
+    "validation": {
+      "torch": "2.13.0+cu130",
+      "cuda": "13.0"
+    }
+  }
+}

build/torch213-cxx11-cu130-aarch64-linux/weight_only_ffn/__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/weight_only_ffn_source_test.abi3.so

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

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

@@ -0,0 +1,261 @@
+"""Weight-only BF16-activation linear and FFN regions."""
+
+from __future__ import annotations
+
+from typing import Optional
+
+import torch
+
+from ._ops import add_op_namespace_prefix, ops
+
+
+def w4_sfb_size_bytes(rows: int, cols: int) -> int:
+    if rows <= 0 or cols <= 0 or cols % 64:
+        raise ValueError("rows must be positive and cols must be divisible by 64")
+    return ((int(rows) + 127) // 128) * (((int(cols) // 16) + 3) // 4) * 512
+
+
+@torch.library.register_fake(add_op_namespace_prefix("quantize_w4_weight_bf16"))
+def _quantize_w4_fake(weight: torch.Tensor, packed: torch.Tensor, sfb: torch.Tensor) -> None:
+    return None
+
+
+@torch.library.register_fake(add_op_namespace_prefix("dequantize_w4_weight_bf16"))
+def _dequantize_w4_fake(packed: torch.Tensor, sfb: torch.Tensor, weight: torch.Tensor) -> None:
+    return None
+
+
+@torch.library.register_fake(add_op_namespace_prefix("w4a16_linear_bf16"))
+def _w4_linear_fake(
+    x: torch.Tensor,
+    packed: torch.Tensor,
+    sfb: torch.Tensor,
+    alpha: float,
+    variant: int,
+    out: torch.Tensor,
+) -> None:
+    return None
+
+
+@torch.library.register_fake(add_op_namespace_prefix("quantize_w8_weight_bf16"))
+def _quantize_w8_fake(weight: torch.Tensor, quantized: torch.Tensor, scales: torch.Tensor) -> None:
+    return None
+
+
+@torch.library.register_fake(add_op_namespace_prefix("dequantize_w8_weight_bf16"))
+def _dequantize_w8_fake(quantized: torch.Tensor, scales: torch.Tensor, weight: torch.Tensor) -> None:
+    return None
+
+
+@torch.library.register_fake(add_op_namespace_prefix("w8a16_linear_bf16"))
+def _w8_linear_fake(
+    x: torch.Tensor,
+    quantized: torch.Tensor,
+    scales: torch.Tensor,
+    variant: int,
+    out: torch.Tensor,
+) -> None:
+    return None
+
+
+@torch.library.register_fake(add_op_namespace_prefix("w4a16_gated_ffn_bf16"))
+def _w4_gated_fake(
+    x, gu_w, gu_s, dn_w, dn_s, gu_b, dn_b, gelu, gu_alpha, dn_alpha,
+    variant, gu, hidden, out,
+) -> None:
+    return None
+
+
+@torch.library.register_fake(add_op_namespace_prefix("w8a16_gated_ffn_bf16"))
+def _w8_gated_fake(
+    x, gu_w, gu_s, dn_w, dn_s, gu_b, dn_b, gelu, variant, gu, hidden, out,
+) -> None:
+    return None
+
+
+@torch.library.register_fake(add_op_namespace_prefix("w4a16_gelu_ffn_bf16"))
+def _w4_gelu_fake(
+    x, up_w, up_s, dn_w, dn_s, up_b, dn_b, up_alpha, dn_alpha,
+    variant, up, hidden, out,
+) -> None:
+    return None
+
+
+@torch.library.register_fake(add_op_namespace_prefix("w8a16_gelu_ffn_bf16"))
+def _w8_gelu_fake(
+    x, up_w, up_s, dn_w, dn_s, up_b, dn_b, variant, up, hidden, out,
+) -> None:
+    return None
+
+
+def quantize_w4_weight_bf16(weight: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor]:
+    """Pack a static BF16 weight matrix into NVFP4 data and SFB scales."""
+    n, k = weight.shape
+    packed = torch.empty((n, k // 2), device=weight.device, dtype=torch.uint8)
+    sfb = torch.empty((w4_sfb_size_bytes(n, k),), device=weight.device, dtype=torch.uint8)
+    ops.quantize_w4_weight_bf16(weight, packed, sfb)
+    return packed, sfb
+
+
+def dequantize_w4_weight_bf16(
+    packed: torch.Tensor, sfb: torch.Tensor, *, cols: int
+) -> torch.Tensor:
+    weight = torch.empty((packed.shape[0], int(cols)), device=packed.device, dtype=torch.bfloat16)
+    ops.dequantize_w4_weight_bf16(packed, sfb, weight)
+    return weight
+
+
+def quantize_w8_weight_bf16(weight: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor]:
+    """Symmetrically quantize a static BF16 weight matrix per output channel."""
+    quantized = torch.empty_like(weight, dtype=torch.int8)
+    scales = torch.empty((weight.shape[0],), device=weight.device, dtype=torch.float32)
+    ops.quantize_w8_weight_bf16(weight, quantized, scales)
+    return quantized, scales
+
+
+def dequantize_w8_weight_bf16(
+    quantized: torch.Tensor, scales: torch.Tensor
+) -> torch.Tensor:
+    weight = torch.empty_like(quantized, dtype=torch.bfloat16)
+    ops.dequantize_w8_weight_bf16(quantized, scales, weight)
+    return weight
+
+
+def w4a16_linear_bf16(
+    x: torch.Tensor,
+    weight_packed: torch.Tensor,
+    weight_sfb: torch.Tensor,
+    *,
+    alpha: float = 1.0,
+    variant: int = 0,
+    out: Optional[torch.Tensor] = None,
+) -> torch.Tensor:
+    if out is None:
+        out = torch.empty((x.shape[0], weight_packed.shape[0]), device=x.device, dtype=torch.bfloat16)
+    ops.w4a16_linear_bf16(x, weight_packed, weight_sfb, float(alpha), int(variant), out)
+    return out
+
+
+def w8a16_linear_bf16(
+    x: torch.Tensor,
+    weight_int8: torch.Tensor,
+    weight_scale: torch.Tensor,
+    *,
+    variant: int = 0,
+    out: Optional[torch.Tensor] = None,
+) -> torch.Tensor:
+    if out is None:
+        out = torch.empty((x.shape[0], weight_int8.shape[0]), device=x.device, dtype=torch.bfloat16)
+    ops.w8a16_linear_bf16(x, weight_int8, weight_scale, int(variant), out)
+    return out
+
+
+def _gated_workspace(x: torch.Tensor, gate_up_rows: int, out_rows: int, workspace, out):
+    hidden_size = gate_up_rows // 2
+    if workspace is None:
+        gate_up = torch.empty((x.shape[0], gate_up_rows), device=x.device, dtype=torch.bfloat16)
+        hidden = torch.empty((x.shape[0], hidden_size), device=x.device, dtype=torch.bfloat16)
+    else:
+        gate_up, hidden = workspace
+    if out is None:
+        out = torch.empty((x.shape[0], out_rows), device=x.device, dtype=torch.bfloat16)
+    return gate_up, hidden, out
+
+
+def _gelu_workspace(x: torch.Tensor, hidden_size: int, out_rows: int, workspace, out):
+    if workspace is None:
+        up = torch.empty((x.shape[0], hidden_size), device=x.device, dtype=torch.bfloat16)
+        hidden = torch.empty_like(up)
+    else:
+        up, hidden = workspace
+    if out is None:
+        out = torch.empty((x.shape[0], out_rows), device=x.device, dtype=torch.bfloat16)
+    return up, hidden, out
+
+
+def _w4_gated(
+    x, gate_up_weight, gate_up_sfb, down_weight, down_sfb, *, gelu,
+    gate_up_bias=None, down_bias=None, gate_up_alpha=1.0, down_alpha=1.0,
+    variant=0, workspace=None, out=None,
+):
+    gate_up, hidden, out = _gated_workspace(
+        x, gate_up_weight.shape[0], down_weight.shape[0], workspace, out
+    )
+    ops.w4a16_gated_ffn_bf16(
+        x, gate_up_weight, gate_up_sfb, down_weight, down_sfb,
+        gate_up_bias, down_bias, bool(gelu), float(gate_up_alpha),
+        float(down_alpha), int(variant), gate_up, hidden, out,
+    )
+    return out
+
+
+def w4a16_swiglu_ffn_bf16(x, gate_up_weight, gate_up_sfb, down_weight, down_sfb, **kwargs):
+    return _w4_gated(x, gate_up_weight, gate_up_sfb, down_weight, down_sfb, gelu=False, **kwargs)
+
+
+def w4a16_geglu_ffn_bf16(x, gate_up_weight, gate_up_sfb, down_weight, down_sfb, **kwargs):
+    return _w4_gated(x, gate_up_weight, gate_up_sfb, down_weight, down_sfb, gelu=True, **kwargs)
+
+
+def _w8_gated(
+    x, gate_up_weight, gate_up_scale, down_weight, down_scale, *, gelu,
+    gate_up_bias=None, down_bias=None, variant=0, workspace=None, out=None,
+):
+    gate_up, hidden, out = _gated_workspace(
+        x, gate_up_weight.shape[0], down_weight.shape[0], workspace, out
+    )
+    ops.w8a16_gated_ffn_bf16(
+        x, gate_up_weight, gate_up_scale, down_weight, down_scale,
+        gate_up_bias, down_bias, bool(gelu), int(variant), gate_up, hidden, out,
+    )
+    return out
+
+
+def w8a16_swiglu_ffn_bf16(x, gate_up_weight, gate_up_scale, down_weight, down_scale, **kwargs):
+    return _w8_gated(x, gate_up_weight, gate_up_scale, down_weight, down_scale, gelu=False, **kwargs)
+
+
+def w8a16_geglu_ffn_bf16(x, gate_up_weight, gate_up_scale, down_weight, down_scale, **kwargs):
+    return _w8_gated(x, gate_up_weight, gate_up_scale, down_weight, down_scale, gelu=True, **kwargs)
+
+
+def w4a16_gelu_ffn_bf16(
+    x, up_weight, up_sfb, down_weight, down_sfb, *, up_bias=None,
+    down_bias=None, up_alpha=1.0, down_alpha=1.0, variant=0,
+    workspace=None, out=None,
+):
+    up, hidden, out = _gelu_workspace(x, up_weight.shape[0], down_weight.shape[0], workspace, out)
+    ops.w4a16_gelu_ffn_bf16(
+        x, up_weight, up_sfb, down_weight, down_sfb, up_bias, down_bias,
+        float(up_alpha), float(down_alpha), int(variant), up, hidden, out,
+    )
+    return out
+
+
+def w8a16_gelu_ffn_bf16(
+    x, up_weight, up_scale, down_weight, down_scale, *, up_bias=None,
+    down_bias=None, variant=0, workspace=None, out=None,
+):
+    up, hidden, out = _gelu_workspace(x, up_weight.shape[0], down_weight.shape[0], workspace, out)
+    ops.w8a16_gelu_ffn_bf16(
+        x, up_weight, up_scale, down_weight, down_scale, up_bias, down_bias,
+        int(variant), up, hidden, out,
+    )
+    return out
+
+
+__all__ = [
+    "dequantize_w4_weight_bf16",
+    "dequantize_w8_weight_bf16",
+    "quantize_w4_weight_bf16",
+    "quantize_w8_weight_bf16",
+    "w4_sfb_size_bytes",
+    "w4a16_geglu_ffn_bf16",
+    "w4a16_gelu_ffn_bf16",
+    "w4a16_linear_bf16",
+    "w4a16_swiglu_ffn_bf16",
+    "w8a16_geglu_ffn_bf16",
+    "w8a16_gelu_ffn_bf16",
+    "w8a16_linear_bf16",
+    "w8a16_swiglu_ffn_bf16",
+]

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

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

build/torch213-cxx11-cu130-x86_64-linux/_weight_only_ffn_cuda_38d37b6.abi3.so

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

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

@@ -0,0 +1,35 @@
+{
+  "name": "weight-only-ffn",
+  "id": "_weight_only_ffn_cuda_38d37b6",
+  "version": 1,
+  "license": "Apache-2.0",
+  "python-depends": [],
+  "backend": {
+    "type": "cuda",
+    "archs": [
+      "11.0a",
+      "12.0a",
+      "12.1"
+    ]
+  },
+  "digest": {
+    "algorithm": "sha256",
+    "files": {
+      "__init__.py": "0q6TMj65una1LdyHLC291Dwy2VzrBbAD8Wm8MEdmqwA=",
+      "_ops.py": "iqY+HoE4nGHqvVlEZTwEAcyu1/w+QrjwuuW9sdTI2uk=",
+      "_weight_only_ffn_cuda_38d37b6.abi3.so": "LCCD2zzjxa1x+WRERQP5ymU5hunoeNIOoIVy9MLH7sk=",
+      "weight_only_ffn/__init__.py": "DFYPlrhXwYjEqCl/8n0SmWGZV8NFml5DPhMjKfv98GY="
+    }
+  },
+  "provenance": {
+    "kernel-builder": {
+      "version": "0.17.0-dev0",
+      "sha": "19aaa6421e674e9fecc352bbae6eab81d19a6bf4",
+      "dirty": false
+    },
+    "kernel": {
+      "sha": "38d37b67acfe940c99d7b3db721fa67163b82cec",
+      "dirty": false
+    }
+  }
+}

build/torch213-cxx11-cu130-x86_64-linux/weight_only_ffn/__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/torch213-cxx11-cu132-x86_64-linux/__init__.py

@@ -0,0 +1,261 @@
+"""Weight-only BF16-activation linear and FFN regions."""
+
+from __future__ import annotations
+
+from typing import Optional
+
+import torch
+
+from ._ops import add_op_namespace_prefix, ops
+
+
+def w4_sfb_size_bytes(rows: int, cols: int) -> int:
+    if rows <= 0 or cols <= 0 or cols % 64:
+        raise ValueError("rows must be positive and cols must be divisible by 64")
+    return ((int(rows) + 127) // 128) * (((int(cols) // 16) + 3) // 4) * 512
+
+
+@torch.library.register_fake(add_op_namespace_prefix("quantize_w4_weight_bf16"))
+def _quantize_w4_fake(weight: torch.Tensor, packed: torch.Tensor, sfb: torch.Tensor) -> None:
+    return None
+
+
+@torch.library.register_fake(add_op_namespace_prefix("dequantize_w4_weight_bf16"))
+def _dequantize_w4_fake(packed: torch.Tensor, sfb: torch.Tensor, weight: torch.Tensor) -> None:
+    return None
+
+
+@torch.library.register_fake(add_op_namespace_prefix("w4a16_linear_bf16"))
+def _w4_linear_fake(
+    x: torch.Tensor,
+    packed: torch.Tensor,
+    sfb: torch.Tensor,
+    alpha: float,
+    variant: int,
+    out: torch.Tensor,
+) -> None:
+    return None
+
+
+@torch.library.register_fake(add_op_namespace_prefix("quantize_w8_weight_bf16"))
+def _quantize_w8_fake(weight: torch.Tensor, quantized: torch.Tensor, scales: torch.Tensor) -> None:
+    return None
+
+
+@torch.library.register_fake(add_op_namespace_prefix("dequantize_w8_weight_bf16"))
+def _dequantize_w8_fake(quantized: torch.Tensor, scales: torch.Tensor, weight: torch.Tensor) -> None:
+    return None
+
+
+@torch.library.register_fake(add_op_namespace_prefix("w8a16_linear_bf16"))
+def _w8_linear_fake(
+    x: torch.Tensor,
+    quantized: torch.Tensor,
+    scales: torch.Tensor,
+    variant: int,
+    out: torch.Tensor,
+) -> None:
+    return None
+
+
+@torch.library.register_fake(add_op_namespace_prefix("w4a16_gated_ffn_bf16"))
+def _w4_gated_fake(
+    x, gu_w, gu_s, dn_w, dn_s, gu_b, dn_b, gelu, gu_alpha, dn_alpha,
+    variant, gu, hidden, out,
+) -> None:
+    return None
+
+
+@torch.library.register_fake(add_op_namespace_prefix("w8a16_gated_ffn_bf16"))
+def _w8_gated_fake(
+    x, gu_w, gu_s, dn_w, dn_s, gu_b, dn_b, gelu, variant, gu, hidden, out,
+) -> None:
+    return None
+
+
+@torch.library.register_fake(add_op_namespace_prefix("w4a16_gelu_ffn_bf16"))
+def _w4_gelu_fake(
+    x, up_w, up_s, dn_w, dn_s, up_b, dn_b, up_alpha, dn_alpha,
+    variant, up, hidden, out,
+) -> None:
+    return None
+
+
+@torch.library.register_fake(add_op_namespace_prefix("w8a16_gelu_ffn_bf16"))
+def _w8_gelu_fake(
+    x, up_w, up_s, dn_w, dn_s, up_b, dn_b, variant, up, hidden, out,
+) -> None:
+    return None
+
+
+def quantize_w4_weight_bf16(weight: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor]:
+    """Pack a static BF16 weight matrix into NVFP4 data and SFB scales."""
+    n, k = weight.shape
+    packed = torch.empty((n, k // 2), device=weight.device, dtype=torch.uint8)
+    sfb = torch.empty((w4_sfb_size_bytes(n, k),), device=weight.device, dtype=torch.uint8)
+    ops.quantize_w4_weight_bf16(weight, packed, sfb)
+    return packed, sfb
+
+
+def dequantize_w4_weight_bf16(
+    packed: torch.Tensor, sfb: torch.Tensor, *, cols: int
+) -> torch.Tensor:
+    weight = torch.empty((packed.shape[0], int(cols)), device=packed.device, dtype=torch.bfloat16)
+    ops.dequantize_w4_weight_bf16(packed, sfb, weight)
+    return weight
+
+
+def quantize_w8_weight_bf16(weight: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor]:
+    """Symmetrically quantize a static BF16 weight matrix per output channel."""
+    quantized = torch.empty_like(weight, dtype=torch.int8)
+    scales = torch.empty((weight.shape[0],), device=weight.device, dtype=torch.float32)
+    ops.quantize_w8_weight_bf16(weight, quantized, scales)
+    return quantized, scales
+
+
+def dequantize_w8_weight_bf16(
+    quantized: torch.Tensor, scales: torch.Tensor
+) -> torch.Tensor:
+    weight = torch.empty_like(quantized, dtype=torch.bfloat16)
+    ops.dequantize_w8_weight_bf16(quantized, scales, weight)
+    return weight
+
+
+def w4a16_linear_bf16(
+    x: torch.Tensor,
+    weight_packed: torch.Tensor,
+    weight_sfb: torch.Tensor,
+    *,
+    alpha: float = 1.0,
+    variant: int = 0,
+    out: Optional[torch.Tensor] = None,
+) -> torch.Tensor:
+    if out is None:
+        out = torch.empty((x.shape[0], weight_packed.shape[0]), device=x.device, dtype=torch.bfloat16)
+    ops.w4a16_linear_bf16(x, weight_packed, weight_sfb, float(alpha), int(variant), out)
+    return out
+
+
+def w8a16_linear_bf16(
+    x: torch.Tensor,
+    weight_int8: torch.Tensor,
+    weight_scale: torch.Tensor,
+    *,
+    variant: int = 0,
+    out: Optional[torch.Tensor] = None,
+) -> torch.Tensor:
+    if out is None:
+        out = torch.empty((x.shape[0], weight_int8.shape[0]), device=x.device, dtype=torch.bfloat16)
+    ops.w8a16_linear_bf16(x, weight_int8, weight_scale, int(variant), out)
+    return out
+
+
+def _gated_workspace(x: torch.Tensor, gate_up_rows: int, out_rows: int, workspace, out):
+    hidden_size = gate_up_rows // 2
+    if workspace is None:
+        gate_up = torch.empty((x.shape[0], gate_up_rows), device=x.device, dtype=torch.bfloat16)
+        hidden = torch.empty((x.shape[0], hidden_size), device=x.device, dtype=torch.bfloat16)
+    else:
+        gate_up, hidden = workspace
+    if out is None:
+        out = torch.empty((x.shape[0], out_rows), device=x.device, dtype=torch.bfloat16)
+    return gate_up, hidden, out
+
+
+def _gelu_workspace(x: torch.Tensor, hidden_size: int, out_rows: int, workspace, out):
+    if workspace is None:
+        up = torch.empty((x.shape[0], hidden_size), device=x.device, dtype=torch.bfloat16)
+        hidden = torch.empty_like(up)
+    else:
+        up, hidden = workspace
+    if out is None:
+        out = torch.empty((x.shape[0], out_rows), device=x.device, dtype=torch.bfloat16)
+    return up, hidden, out
+
+
+def _w4_gated(
+    x, gate_up_weight, gate_up_sfb, down_weight, down_sfb, *, gelu,
+    gate_up_bias=None, down_bias=None, gate_up_alpha=1.0, down_alpha=1.0,
+    variant=0, workspace=None, out=None,
+):
+    gate_up, hidden, out = _gated_workspace(
+        x, gate_up_weight.shape[0], down_weight.shape[0], workspace, out
+    )
+    ops.w4a16_gated_ffn_bf16(
+        x, gate_up_weight, gate_up_sfb, down_weight, down_sfb,
+        gate_up_bias, down_bias, bool(gelu), float(gate_up_alpha),
+        float(down_alpha), int(variant), gate_up, hidden, out,
+    )
+    return out
+
+
+def w4a16_swiglu_ffn_bf16(x, gate_up_weight, gate_up_sfb, down_weight, down_sfb, **kwargs):
+    return _w4_gated(x, gate_up_weight, gate_up_sfb, down_weight, down_sfb, gelu=False, **kwargs)
+
+
+def w4a16_geglu_ffn_bf16(x, gate_up_weight, gate_up_sfb, down_weight, down_sfb, **kwargs):
+    return _w4_gated(x, gate_up_weight, gate_up_sfb, down_weight, down_sfb, gelu=True, **kwargs)
+
+
+def _w8_gated(
+    x, gate_up_weight, gate_up_scale, down_weight, down_scale, *, gelu,
+    gate_up_bias=None, down_bias=None, variant=0, workspace=None, out=None,
+):
+    gate_up, hidden, out = _gated_workspace(
+        x, gate_up_weight.shape[0], down_weight.shape[0], workspace, out
+    )
+    ops.w8a16_gated_ffn_bf16(
+        x, gate_up_weight, gate_up_scale, down_weight, down_scale,
+        gate_up_bias, down_bias, bool(gelu), int(variant), gate_up, hidden, out,
+    )
+    return out
+
+
+def w8a16_swiglu_ffn_bf16(x, gate_up_weight, gate_up_scale, down_weight, down_scale, **kwargs):
+    return _w8_gated(x, gate_up_weight, gate_up_scale, down_weight, down_scale, gelu=False, **kwargs)
+
+
+def w8a16_geglu_ffn_bf16(x, gate_up_weight, gate_up_scale, down_weight, down_scale, **kwargs):
+    return _w8_gated(x, gate_up_weight, gate_up_scale, down_weight, down_scale, gelu=True, **kwargs)
+
+
+def w4a16_gelu_ffn_bf16(
+    x, up_weight, up_sfb, down_weight, down_sfb, *, up_bias=None,
+    down_bias=None, up_alpha=1.0, down_alpha=1.0, variant=0,
+    workspace=None, out=None,
+):
+    up, hidden, out = _gelu_workspace(x, up_weight.shape[0], down_weight.shape[0], workspace, out)
+    ops.w4a16_gelu_ffn_bf16(
+        x, up_weight, up_sfb, down_weight, down_sfb, up_bias, down_bias,
+        float(up_alpha), float(down_alpha), int(variant), up, hidden, out,
+    )
+    return out
+
+
+def w8a16_gelu_ffn_bf16(
+    x, up_weight, up_scale, down_weight, down_scale, *, up_bias=None,
+    down_bias=None, variant=0, workspace=None, out=None,
+):
+    up, hidden, out = _gelu_workspace(x, up_weight.shape[0], down_weight.shape[0], workspace, out)
+    ops.w8a16_gelu_ffn_bf16(
+        x, up_weight, up_scale, down_weight, down_scale, up_bias, down_bias,
+        int(variant), up, hidden, out,
+    )
+    return out
+
+
+__all__ = [
+    "dequantize_w4_weight_bf16",
+    "dequantize_w8_weight_bf16",
+    "quantize_w4_weight_bf16",
+    "quantize_w8_weight_bf16",
+    "w4_sfb_size_bytes",
+    "w4a16_geglu_ffn_bf16",
+    "w4a16_gelu_ffn_bf16",
+    "w4a16_linear_bf16",
+    "w4a16_swiglu_ffn_bf16",
+    "w8a16_geglu_ffn_bf16",
+    "w8a16_gelu_ffn_bf16",
+    "w8a16_linear_bf16",
+    "w8a16_swiglu_ffn_bf16",
+]

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

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

build/torch213-cxx11-cu132-x86_64-linux/_weight_only_ffn_cuda_38d37b6.abi3.so

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

build/torch213-cxx11-cu132-x86_64-linux/metadata.json

@@ -0,0 +1,35 @@
+{
+  "name": "weight-only-ffn",
+  "id": "_weight_only_ffn_cuda_38d37b6",
+  "version": 1,
+  "license": "Apache-2.0",
+  "python-depends": [],
+  "backend": {
+    "type": "cuda",
+    "archs": [
+      "11.0a",
+      "12.0a",
+      "12.1"
+    ]
+  },
+  "digest": {
+    "algorithm": "sha256",
+    "files": {
+      "__init__.py": "0q6TMj65una1LdyHLC291Dwy2VzrBbAD8Wm8MEdmqwA=",
+      "_ops.py": "iqY+HoE4nGHqvVlEZTwEAcyu1/w+QrjwuuW9sdTI2uk=",
+      "_weight_only_ffn_cuda_38d37b6.abi3.so": "k/8WxKvZBX+RZvst0LdlBT6o1L1p3ctRS0/D5t5WuxA=",
+      "weight_only_ffn/__init__.py": "DFYPlrhXwYjEqCl/8n0SmWGZV8NFml5DPhMjKfv98GY="
+    }
+  },
+  "provenance": {
+    "kernel-builder": {
+      "version": "0.17.0-dev0",
+      "sha": "19aaa6421e674e9fecc352bbae6eab81d19a6bf4",
+      "dirty": false
+    },
+    "kernel": {
+      "sha": "38d37b67acfe940c99d7b3db721fa67163b82cec",
+      "dirty": false
+    }
+  }
+}

build/torch213-cxx11-cu132-x86_64-linux/weight_only_ffn/__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")))