Promote latest kernel artifacts to main

README.md

@@ -1,9 +0,0 @@
-# flashrt/fp4-fused-ops
-
-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/fp4-fused-ops
-
-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

@@ -6,6 +6,7 @@ from __future__ import annotations
 import argparse
 import importlib.util
 import json
+import os
 import sys
 from dataclasses import asdict, dataclass
 from pathlib import Path
@@ -24,8 +25,13 @@ class BenchResult:
     dim: int
     workload: str
     reference_us: float | None
+    native_us: float | None
     flashrt_us: float
+    native_graph_us: float | None
+    flashrt_graph_us: float | None
     speedup: float | None
+    wrapper_over_native: float | None
+    graph_over_native: float | None
     status: str
 
 
@@ -53,7 +59,17 @@ def measure(fn, warmup: int, iters: int) -> float:
     return float(start.elapsed_time(end) * 1000.0 / iters)
 
 
-def bench_case(helpers, ops, rows: int, dim: int, warmup: int, iters: int) -> list[BenchResult]:
+def measure_graph(fn, warmup: int, iters: int) -> float:
+    for _ in range(warmup):
+        fn()
+    torch.cuda.synchronize()
+    graph = torch.cuda.CUDAGraph()
+    with torch.cuda.graph(graph):
+        fn()
+    return measure(graph.replay, warmup, iters)
+
+
+def bench_case(helpers, ops, native, rows: int, dim: int, warmup: int, iters: int) -> list[BenchResult]:
     make_fp16 = helpers.make_fp16
     results: list[BenchResult] = []
 
@@ -76,6 +92,29 @@ def bench_case(helpers, ops, rows: int, dim: int, warmup: int, iters: int) -> li
         warmup,
         iters,
     )
+    stream = torch.cuda.current_stream().cuda_stream
+    native_f3_us = measure(
+        lambda: native.residual_add_rms_norm_fp4_sfa_v2_fp16(
+            residual_v2.copy_(residual).data_ptr(), x.data_ptr(),
+            packed.data_ptr(), sfa.data_ptr(), rows, dim, stream
+        ),
+        warmup,
+        iters,
+    )
+    graph_f3_us = measure_graph(
+        lambda: ops.residual_add_rms_norm_fp4_sfa_v2_fp16(residual_v2.copy_(residual), x, packed, sfa),
+        warmup,
+        iters,
+    )
+    native_graph_f3_us = measure_graph(
+        lambda: native.residual_add_rms_norm_fp4_sfa_v2_fp16(
+            residual_v2.copy_(residual).data_ptr(), x.data_ptr(),
+            packed.data_ptr(), sfa.data_ptr(), rows, dim,
+            torch.cuda.current_stream().cuda_stream
+        ),
+        warmup,
+        iters,
+    )
     results.append(
         BenchResult(
             case=f"rows{rows}_dim{dim}",
@@ -83,8 +122,13 @@ def bench_case(helpers, ops, rows: int, dim: int, warmup: int, iters: int) -> li
             dim=dim,
             workload="residual_add_rms_norm_fp4_sfa_v2",
             reference_us=ref_us,
+            native_us=native_f3_us,
             flashrt_us=f3_us,
+            native_graph_us=native_graph_f3_us,
+            flashrt_graph_us=graph_f3_us,
             speedup=(ref_us / f3_us) if ref_us else None,
+            wrapper_over_native=f3_us / native_f3_us,
+            graph_over_native=graph_f3_us / native_graph_f3_us,
             status="v1 reference" if ref_us else "v2 only; v1 rejects this dim",
         )
     )
@@ -94,6 +138,25 @@ def bench_case(helpers, ops, rows: int, dim: int, warmup: int, iters: int) -> li
     packed_v2, sfa_v2 = ops.alloc(rows, dim)
     ref_us = measure(lambda: ops.silu_mul_fp4_sfa_fp16(merged, packed_v1, sfa_v1), warmup, iters)
     f4_us = measure(lambda: ops.silu_mul_fp4_sfa_v2_fp16(merged, packed_v2, sfa_v2), warmup, iters)
+    native_f4_us = measure(
+        lambda: native.gate_geglu_fp4_sfa_v2_fp16(
+            merged.data_ptr(), packed_v2.data_ptr(), sfa_v2.data_ptr(),
+            rows, dim, stream
+        ),
+        warmup,
+        iters,
+    )
+    graph_f4_us = measure_graph(
+        lambda: ops.silu_mul_fp4_sfa_v2_fp16(merged, packed_v2, sfa_v2), warmup, iters
+    )
+    native_graph_f4_us = measure_graph(
+        lambda: native.gate_geglu_fp4_sfa_v2_fp16(
+            merged.data_ptr(), packed_v2.data_ptr(), sfa_v2.data_ptr(), rows, dim,
+            torch.cuda.current_stream().cuda_stream
+        ),
+        warmup,
+        iters,
+    )
     results.append(
         BenchResult(
             case=f"rows{rows}_dim{dim}",
@@ -101,14 +164,40 @@ def bench_case(helpers, ops, rows: int, dim: int, warmup: int, iters: int) -> li
             dim=dim,
             workload="silu_mul_fp4_sfa_v2",
             reference_us=ref_us,
+            native_us=native_f4_us,
             flashrt_us=f4_us,
+            native_graph_us=native_graph_f4_us,
+            flashrt_graph_us=graph_f4_us,
             speedup=ref_us / f4_us,
+            wrapper_over_native=f4_us / native_f4_us,
+            graph_over_native=graph_f4_us / native_graph_f4_us,
             status="v1 reference",
         )
     )
 
     inv_s = (torch.rand((dim,), device="cuda") * 0.25 + 0.875).to(torch.float16).contiguous()
     awq_us = measure(lambda: ops.silu_mul_mul_fp4_sfa_v2_fp16(merged, inv_s, packed_v2, sfa_v2), warmup, iters)
+    native_awq_us = measure(
+        lambda: native.gate_geglu_mul_fp4_sfa_v2_fp16(
+            merged.data_ptr(), inv_s.data_ptr(), packed_v2.data_ptr(),
+            sfa_v2.data_ptr(), rows, dim, stream
+        ),
+        warmup,
+        iters,
+    )
+    graph_awq_us = measure_graph(
+        lambda: ops.silu_mul_mul_fp4_sfa_v2_fp16(merged, inv_s, packed_v2, sfa_v2),
+        warmup,
+        iters,
+    )
+    native_graph_awq_us = measure_graph(
+        lambda: native.gate_geglu_mul_fp4_sfa_v2_fp16(
+            merged.data_ptr(), inv_s.data_ptr(), packed_v2.data_ptr(),
+            sfa_v2.data_ptr(), rows, dim, torch.cuda.current_stream().cuda_stream
+        ),
+        warmup,
+        iters,
+    )
     results.append(
         BenchResult(
             case=f"rows{rows}_dim{dim}",
@@ -116,8 +205,13 @@ def bench_case(helpers, ops, rows: int, dim: int, warmup: int, iters: int) -> li
             dim=dim,
             workload="silu_mul_mul_fp4_sfa_v2",
             reference_us=None,
+            native_us=native_awq_us,
             flashrt_us=awq_us,
+            native_graph_us=native_graph_awq_us,
+            flashrt_graph_us=graph_awq_us,
             speedup=None,
+            wrapper_over_native=awq_us / native_awq_us,
+            graph_over_native=graph_awq_us / native_graph_awq_us,
             status="fused AWQ producer latency",
         )
     )
@@ -132,11 +226,57 @@ def bench_case(helpers, ops, rows: int, dim: int, warmup: int, iters: int) -> li
         warmup,
         iters,
     )
+    native_two_us = measure(
+        lambda: native.geglu_two_fp4_to_fp4(
+            gate_packed.data_ptr(), gate_sfa.data_ptr(), up_packed.data_ptr(),
+            up_sfa.data_ptr(), out_packed.data_ptr(), out_sfa.data_ptr(),
+            rows, dim, stream
+        ),
+        warmup,
+        iters,
+    )
+    graph_two_us = measure_graph(
+        lambda: ops.silu_mul_two_fp4_to_fp4(gate_packed, gate_sfa, up_packed, up_sfa, out_packed, out_sfa),
+        warmup,
+        iters,
+    )
+    native_graph_two_us = measure_graph(
+        lambda: native.geglu_two_fp4_to_fp4(
+            gate_packed.data_ptr(), gate_sfa.data_ptr(), up_packed.data_ptr(),
+            up_sfa.data_ptr(), out_packed.data_ptr(), out_sfa.data_ptr(), rows, dim,
+            torch.cuda.current_stream().cuda_stream
+        ),
+        warmup,
+        iters,
+    )
     two_mul_us = measure(
         lambda: ops.silu_mul_two_mul_fp4_to_fp4(gate_packed, gate_sfa, up_packed, up_sfa, inv_s, out_packed, out_sfa),
         warmup,
         iters,
     )
+    native_two_mul_us = measure(
+        lambda: native.geglu_two_mul_fp4_to_fp4(
+            gate_packed.data_ptr(), gate_sfa.data_ptr(), up_packed.data_ptr(),
+            up_sfa.data_ptr(), inv_s.data_ptr(), out_packed.data_ptr(),
+            out_sfa.data_ptr(), rows, dim, stream
+        ),
+        warmup,
+        iters,
+    )
+    graph_two_mul_us = measure_graph(
+        lambda: ops.silu_mul_two_mul_fp4_to_fp4(gate_packed, gate_sfa, up_packed, up_sfa, inv_s, out_packed, out_sfa),
+        warmup,
+        iters,
+    )
+    native_graph_two_mul_us = measure_graph(
+        lambda: native.geglu_two_mul_fp4_to_fp4(
+            gate_packed.data_ptr(), gate_sfa.data_ptr(), up_packed.data_ptr(),
+            up_sfa.data_ptr(), inv_s.data_ptr(), out_packed.data_ptr(),
+            out_sfa.data_ptr(), rows, dim, torch.cuda.current_stream().cuda_stream
+        ),
+        warmup,
+        iters,
+    )
     results.append(
         BenchResult(
             case=f"rows{rows}_dim{dim}",
@@ -144,8 +284,13 @@ def bench_case(helpers, ops, rows: int, dim: int, warmup: int, iters: int) -> li
             dim=dim,
             workload="silu_mul_two_fp4_to_fp4",
             reference_us=None,
+            native_us=native_two_us,
             flashrt_us=two_us,
+            native_graph_us=native_graph_two_us,
+            flashrt_graph_us=graph_two_us,
             speedup=None,
+            wrapper_over_native=two_us / native_two_us,
+            graph_over_native=graph_two_us / native_graph_two_us,
             status="FP4-to-FP4 combiner latency",
         )
     )
@@ -156,8 +301,13 @@ def bench_case(helpers, ops, rows: int, dim: int, warmup: int, iters: int) -> li
             dim=dim,
             workload="silu_mul_two_mul_fp4_to_fp4",
             reference_us=None,
+            native_us=native_two_mul_us,
             flashrt_us=two_mul_us,
+            native_graph_us=native_graph_two_mul_us,
+            flashrt_graph_us=graph_two_mul_us,
             speedup=None,
+            wrapper_over_native=two_mul_us / native_two_mul_us,
+            graph_over_native=graph_two_mul_us / native_graph_two_mul_us,
             status="FP4-to-FP4 AWQ combiner latency",
         )
     )
@@ -166,7 +316,11 @@ def bench_case(helpers, ops, rows: int, dim: int, warmup: int, iters: int) -> li
 
 def main() -> int:
     parser = argparse.ArgumentParser()
-    parser.add_argument("--mode", choices=["smoke", "headline"], default="headline")
+    parser.add_argument("--backend", choices=["source", "installed"], default="source")
+    parser.add_argument("--artifact", default=None)
+    parser.add_argument(
+        "--mode", choices=["smoke", "headline", "thor-models"], default="headline"
+    )
     parser.add_argument("--warmup", type=int, default=20)
     parser.add_argument("--iterations", type=int, default=100)
     parser.add_argument("--json-out", default=None)
@@ -175,15 +329,28 @@ def main() -> int:
     if not torch.cuda.is_available():
         raise RuntimeError("CUDA is required")
     helpers = load_test_module()
-    ops = helpers.load_source_ops()
+    ops = (
+        helpers.load_source_ops()
+        if args.backend == "source"
+        else helpers.load_installed_ops(args.artifact)
+    )
+    native_root = Path(os.environ.get("FLASHRT_NATIVE_ROOT", str(ROOT.parent / "official" / "FlashRT")))
+    sys.path.insert(0, str(native_root))
+    try:
+        import flash_rt.flash_rt_fp4 as native
+    finally:
+        sys.path.pop(0)
 
     shapes = [(1, 1024), (10, 2048)] if args.mode == "smoke" else [(1, 1024), (10, 2048), (64, 2048), (128, 4096)]
+    if args.mode == "thor-models":
+        shapes = list(helpers.THOR_MODEL_SHAPES.values())
     results: list[BenchResult] = []
     for rows, dim in shapes:
-        results.extend(bench_case(helpers, ops, rows, dim, args.warmup, args.iterations))
+        results.extend(bench_case(helpers, ops, native, rows, dim, args.warmup, args.iterations))
 
     payload = {
         "mode": args.mode,
+        "backend": args.backend,
         "device": torch.cuda.get_device_name(),
         "torch": torch.__version__,
         "results": [asdict(item) for item in results],

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

@@ -91,6 +91,57 @@ def _dequant_fake(packed: torch.Tensor, sfa: torch.Tensor, out: torch.Tensor) ->
     return None
 
 
+@torch.library.register_fake(add_op_namespace_prefix("rms_silu_nvfp4_ndhwc_bf16"))
+def _rms_silu_nvfp4_ndhwc_bf16_fake(
+    x: torch.Tensor,
+    gamma: torch.Tensor,
+    awq_inv_scale: torch.Tensor | None,
+    eps: float,
+    packed: torch.Tensor,
+    scale_factors: torch.Tensor,
+) -> None:
+    del eps
+    if x.dim() != 5:
+        raise RuntimeError("x must have shape (B,C,T,H,W)")
+    b, c, t, h, w = x.shape
+    if (
+        c % 128 != 0
+        or c > 1024
+        or gamma.shape != (c,)
+        or packed.shape != (b, t, h, w, c // 2)
+        or scale_factors.shape != (b, t, h, w, c // 16)
+        or (awq_inv_scale is not None and awq_inv_scale.shape != (c,))
+    ):
+        raise RuntimeError(
+            "rms_silu_nvfp4_ndhwc_bf16 shape contract failed"
+        )
+    return None
+
+
+@torch.library.register_fake(
+    add_op_namespace_prefix("quantize_bf16_to_nvfp4_linear")
+)
+def _quantize_bf16_to_nvfp4_linear_fake(
+    input: torch.Tensor,
+    packed: torch.Tensor,
+    scale_factors: torch.Tensor,
+) -> None:
+    if input.dim() != 2:
+        raise RuntimeError("input must have shape (rows, cols)")
+    rows, cols = input.shape
+    if (
+        rows <= 0
+        or cols <= 0
+        or cols % 16 != 0
+        or packed.shape != (rows, cols // 2)
+        or scale_factors.shape != (rows, cols // 16)
+    ):
+        raise RuntimeError(
+            "quantize_bf16_to_nvfp4_linear shape contract failed"
+        )
+    return None
+
+
 def rms_norm_fp4_sfa_fp16(x: torch.Tensor, packed: torch.Tensor | None = None, sfa: torch.Tensor | None = None):
     if packed is None or sfa is None:
         packed, sfa = _alloc_fp4(x.shape[0], x.shape[1], x.device)
@@ -204,3 +255,156 @@ def dequantize_fp4_sfa_fp16(
         out = torch.empty((packed.shape[0], packed.shape[1] * 2), device=packed.device, dtype=torch.float16)
     ops.dequantize_fp4_sfa_fp16(packed, sfa, out)
     return out
+
+
+def rms_silu_nvfp4_ndhwc_bf16(
+    x: torch.Tensor,
+    gamma: torch.Tensor,
+    awq_inv_scale: torch.Tensor | None = None,
+    eps: float = 1e-6,
+    *,
+    packed: torch.Tensor | None = None,
+    scale_factors: torch.Tensor | None = None,
+) -> tuple[torch.Tensor, torch.Tensor]:
+    """Fuse BF16 NCDHW RMSNorm, SiLU and linear-SF NVFP4 quantization."""
+
+    b, c, t, h, w = x.shape
+    if packed is None:
+        packed = torch.empty(
+            (b, t, h, w, c // 2), device=x.device, dtype=torch.uint8
+        )
+    if scale_factors is None:
+        scale_factors = torch.empty(
+            (b, t, h, w, c // 16), device=x.device, dtype=torch.uint8
+        )
+    ops.rms_silu_nvfp4_ndhwc_bf16(
+        x,
+        gamma,
+        awq_inv_scale,
+        float(eps),
+        packed,
+        scale_factors,
+    )
+    return packed, scale_factors
+
+
+def quantize_bf16_to_nvfp4_linear(
+    input: torch.Tensor,
+    *,
+    packed: torch.Tensor | None = None,
+    scale_factors: torch.Tensor | None = None,
+) -> tuple[torch.Tensor, torch.Tensor]:
+    """Quantize BF16 rows to packed E2M1 with linear UE4M3 block-16 scales."""
+
+    rows, cols = input.shape
+    if packed is None:
+        packed = torch.empty(
+            (rows, cols // 2), device=input.device, dtype=torch.uint8
+        )
+    if scale_factors is None:
+        scale_factors = torch.empty(
+            (rows, cols // 16), device=input.device, dtype=torch.uint8
+        )
+    ops.quantize_bf16_to_nvfp4_linear(input, packed, scale_factors)
+    return packed, scale_factors
+
+
+def _check_ncdhw(x: torch.Tensor, gamma: torch.Tensor) -> None:
+    if x.dim() != 5:
+        raise RuntimeError("x must have shape (B,C,T,H,W)")
+    if x.shape[1] <= 0 or x.shape[1] % 2 != 0 or x.shape[1] > 1024:
+        raise RuntimeError("C must be even and at most 1024")
+    if gamma.shape != (x.shape[1],):
+        raise RuntimeError("gamma must have shape (C,)")
+
+
+@torch.library.register_fake(add_op_namespace_prefix("bf16_rms_silu_ncdhw"))
+def _bf16_rms_silu_ncdhw_fake(
+    x: torch.Tensor,
+    gamma: torch.Tensor,
+    prev_cache: torch.Tensor | None,
+    eps: float,
+    out: torch.Tensor,
+    next_cache: torch.Tensor | None,
+) -> None:
+    del eps
+    _check_ncdhw(x, gamma)
+    if out.shape != x.shape:
+        raise RuntimeError("out must have the same shape as x")
+    expected_cache = (x.shape[0], x.shape[1], 2, x.shape[3], x.shape[4])
+    if prev_cache is not None and prev_cache.shape != expected_cache:
+        raise RuntimeError("prev_cache must have shape (B,C,2,H,W)")
+    if next_cache is not None and next_cache.shape != expected_cache:
+        raise RuntimeError("next_cache must have shape (B,C,2,H,W)")
+    return None
+
+
+@torch.library.register_fake(add_op_namespace_prefix("bf16_rms_norm_ncdhw"))
+def _bf16_rms_norm_ncdhw_fake(
+    x: torch.Tensor,
+    gamma: torch.Tensor,
+    bias: torch.Tensor | None,
+    eps: float,
+    out: torch.Tensor,
+) -> None:
+    del eps
+    _check_ncdhw(x, gamma)
+    if bias is not None and bias.shape != (x.shape[1],):
+        raise RuntimeError("bias must have shape (C,)")
+    if out.shape != x.shape:
+        raise RuntimeError("out must have the same shape as x")
+    return None
+
+
+def bf16_rms_silu_ncdhw(
+    x: torch.Tensor,
+    gamma: torch.Tensor,
+    prev_cache: torch.Tensor | None = None,
+    eps: float = 1e-6,
+    *,
+    out: torch.Tensor | None = None,
+    next_cache: torch.Tensor | None = None,
+) -> tuple[torch.Tensor, torch.Tensor | None]:
+    """Apply channel RMSNorm and SiLU to BF16 NCDHW with optional cache update."""
+
+    if out is None:
+        out = torch.empty_like(x)
+    ops.bf16_rms_silu_ncdhw(
+        x, gamma, prev_cache, float(eps), out, next_cache
+    )
+    return out, next_cache
+
+
+def bf16_rms_norm_ncdhw(
+    x: torch.Tensor,
+    gamma: torch.Tensor,
+    bias: torch.Tensor | None = None,
+    eps: float = 1e-6,
+    *,
+    out: torch.Tensor | None = None,
+) -> torch.Tensor:
+    """Apply channel RMSNorm and optional bias to BF16 NCDHW."""
+
+    if out is None:
+        out = torch.empty_like(x)
+    ops.bf16_rms_norm_ncdhw(x, gamma, bias, float(eps), out)
+    return out
+
+
+__all__ = [
+    "bf16_rms_norm_ncdhw",
+    "bf16_rms_silu_ncdhw",
+    "dequantize_fp4_sfa_fp16",
+    "quantize_bf16_to_nvfp4_linear",
+    "residual_add_rms_norm_fp4_sfa_fp16",
+    "residual_add_rms_norm_fp4_sfa_v2_fp16",
+    "residual_add_rms_norm_mul_fp4_sfa_fp16",
+    "rms_norm_fp4_sfa_fp16",
+    "rms_silu_nvfp4_ndhwc_bf16",
+    "sfa_size_bytes",
+    "silu_mul_fp4_sfa_fp16",
+    "silu_mul_fp4_sfa_v2_fp16",
+    "silu_mul_mul_fp4_sfa_v2_fp16",
+    "silu_mul_two_fp4_to_fp4",
+    "silu_mul_two_mul_fp4_to_fp4",
+]

build/{torch212-cxx11-cu132-x86_64-linux/_fp4_fused_ops_cuda_7e2e5b9.abi3.so → torch211-cxx11-cu128-x86_64-linux/_fp4_fused_ops_cuda_8e0b1cd.abi3.so}

@@ -1,3 +1,3 @@
 version https://git-lfs.github.com/spec/v1
-oid sha256:fbb803d2760f5ff420a311dc741b11b284819e558d00601d9a3c2f5e921d3288
-size 545896
+oid sha256:65ea1e5a247cdb9abb385600a3208c1b55d7d4564a9dffe7b31e6c2f0adda3fc
+size 592896

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

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

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

@@ -1,6 +1,6 @@
 {
   "name": "fp4-fused-ops",
-  "id": "_fp4_fused_ops_cuda_7e2e5b9",
+  "id": "_fp4_fused_ops_cuda_8e0b1cd",
   "version": 1,
   "license": "Apache-2.0",
   "python-depends": [],
@@ -13,9 +13,9 @@
   "digest": {
     "algorithm": "sha256",
     "files": {
-      "__init__.py": "aB7foEHGoXVq7xDlENZhAsIjQvzWhMpiKhanJhtzmk8=",
-      "_fp4_fused_ops_cuda_7e2e5b9.abi3.so": "3Q9S5xqFrpzHoz7d+WQkoZeRqkj5kfrZo1jRKjE7rgs=",
-      "_ops.py": "qtD0ljSFCLkwC71lYCEP/FcSj+NFL+0sWWugI+4CH4s=",
+      "__init__.py": "511/VkU+HAwubcrzG/Ck9wB3ukPbjKrWxG+9EgtcP00=",
+      "_fp4_fused_ops_cuda_8e0b1cd.abi3.so": "ZeoeWiR825q7OFYAoyCMG1XX1FZKnf/nsx5sLwrdo/w=",
+      "_ops.py": "L5M3920Afl87xMKszpaz5zL32AN6JoBJM6Aw1J19k2I=",
       "fp4_fused_ops/__init__.py": "DFYPlrhXwYjEqCl/8n0SmWGZV8NFml5DPhMjKfv98GY="
     }
   }

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

@@ -0,0 +1,410 @@
+"""FlashRT FP4 fused producer kernels."""
+
+from __future__ import annotations
+
+import torch
+
+from ._ops import add_op_namespace_prefix, ops
+
+
+def sfa_size_bytes(rows: int, dim: int, is_sfb: bool = False, device: torch.device | str | None = None) -> int:
+    if device is None:
+        device = torch.device("cuda", torch.cuda.current_device())
+    anchor = torch.empty((1,), device=device, dtype=torch.uint8)
+    return int(ops.sfa_size_bytes_for(anchor, int(rows), int(dim), bool(is_sfb)))
+
+
+def _alloc_fp4(rows: int, dim: int, device: torch.device | str):
+    packed = torch.empty((rows, dim // 2), device=device, dtype=torch.uint8)
+    sfa = torch.empty((sfa_size_bytes(rows, dim, False, device=device),), device=device, dtype=torch.uint8)
+    return packed, sfa
+
+
+@torch.library.register_fake(add_op_namespace_prefix("rms_norm_fp4_sfa_fp16"))
+def _rms_norm_fake(x: torch.Tensor, packed: torch.Tensor, sfa: torch.Tensor) -> None:
+    return None
+
+
+@torch.library.register_fake(add_op_namespace_prefix("residual_add_rms_norm_fp4_sfa_fp16"))
+def _residual_fake(residual: torch.Tensor, x: torch.Tensor, packed: torch.Tensor, sfa: torch.Tensor) -> None:
+    return None
+
+
+@torch.library.register_fake(add_op_namespace_prefix("residual_add_rms_norm_fp4_sfa_v2_fp16"))
+def _residual_v2_fake(residual: torch.Tensor, x: torch.Tensor, packed: torch.Tensor, sfa: torch.Tensor) -> None:
+    return None
+
+
+@torch.library.register_fake(add_op_namespace_prefix("residual_add_rms_norm_mul_fp4_sfa_fp16"))
+def _residual_mul_fake(
+    residual: torch.Tensor,
+    x: torch.Tensor,
+    inv_s: torch.Tensor,
+    packed: torch.Tensor,
+    sfa: torch.Tensor,
+) -> None:
+    return None
+
+
+@torch.library.register_fake(add_op_namespace_prefix("silu_mul_fp4_sfa_fp16"))
+def _silu_fake(merged: torch.Tensor, packed: torch.Tensor, sfa: torch.Tensor) -> None:
+    return None
+
+
+@torch.library.register_fake(add_op_namespace_prefix("silu_mul_fp4_sfa_v2_fp16"))
+def _silu_v2_fake(merged: torch.Tensor, packed: torch.Tensor, sfa: torch.Tensor) -> None:
+    return None
+
+
+@torch.library.register_fake(add_op_namespace_prefix("silu_mul_mul_fp4_sfa_v2_fp16"))
+def _silu_mul_v2_fake(merged: torch.Tensor, inv_s: torch.Tensor, packed: torch.Tensor, sfa: torch.Tensor) -> None:
+    return None
+
+
+@torch.library.register_fake(add_op_namespace_prefix("silu_mul_two_fp4_to_fp4"))
+def _silu_two_fake(
+    gate_packed: torch.Tensor,
+    gate_sfa: torch.Tensor,
+    up_packed: torch.Tensor,
+    up_sfa: torch.Tensor,
+    out_packed: torch.Tensor,
+    out_sfa: torch.Tensor,
+) -> None:
+    return None
+
+
+@torch.library.register_fake(add_op_namespace_prefix("silu_mul_two_mul_fp4_to_fp4"))
+def _silu_two_mul_fake(
+    gate_packed: torch.Tensor,
+    gate_sfa: torch.Tensor,
+    up_packed: torch.Tensor,
+    up_sfa: torch.Tensor,
+    inv_s: torch.Tensor,
+    out_packed: torch.Tensor,
+    out_sfa: torch.Tensor,
+) -> None:
+    return None
+
+
+@torch.library.register_fake(add_op_namespace_prefix("dequantize_fp4_sfa_fp16"))
+def _dequant_fake(packed: torch.Tensor, sfa: torch.Tensor, out: torch.Tensor) -> None:
+    return None
+
+
+@torch.library.register_fake(add_op_namespace_prefix("rms_silu_nvfp4_ndhwc_bf16"))
+def _rms_silu_nvfp4_ndhwc_bf16_fake(
+    x: torch.Tensor,
+    gamma: torch.Tensor,
+    awq_inv_scale: torch.Tensor | None,
+    eps: float,
+    packed: torch.Tensor,
+    scale_factors: torch.Tensor,
+) -> None:
+    del eps
+    if x.dim() != 5:
+        raise RuntimeError("x must have shape (B,C,T,H,W)")
+    b, c, t, h, w = x.shape
+    if (
+        c % 128 != 0
+        or c > 1024
+        or gamma.shape != (c,)
+        or packed.shape != (b, t, h, w, c // 2)
+        or scale_factors.shape != (b, t, h, w, c // 16)
+        or (awq_inv_scale is not None and awq_inv_scale.shape != (c,))
+    ):
+        raise RuntimeError(
+            "rms_silu_nvfp4_ndhwc_bf16 shape contract failed"
+        )
+    return None
+
+
+@torch.library.register_fake(
+    add_op_namespace_prefix("quantize_bf16_to_nvfp4_linear")
+)
+def _quantize_bf16_to_nvfp4_linear_fake(
+    input: torch.Tensor,
+    packed: torch.Tensor,
+    scale_factors: torch.Tensor,
+) -> None:
+    if input.dim() != 2:
+        raise RuntimeError("input must have shape (rows, cols)")
+    rows, cols = input.shape
+    if (
+        rows <= 0
+        or cols <= 0
+        or cols % 16 != 0
+        or packed.shape != (rows, cols // 2)
+        or scale_factors.shape != (rows, cols // 16)
+    ):
+        raise RuntimeError(
+            "quantize_bf16_to_nvfp4_linear shape contract failed"
+        )
+    return None
+
+
+def rms_norm_fp4_sfa_fp16(x: torch.Tensor, packed: torch.Tensor | None = None, sfa: torch.Tensor | None = None):
+    if packed is None or sfa is None:
+        packed, sfa = _alloc_fp4(x.shape[0], x.shape[1], x.device)
+    ops.rms_norm_fp4_sfa_fp16(x, packed, sfa)
+    return packed, sfa
+
+
+def residual_add_rms_norm_fp4_sfa_fp16(
+    residual: torch.Tensor,
+    x: torch.Tensor,
+    packed: torch.Tensor | None = None,
+    sfa: torch.Tensor | None = None,
+):
+    if packed is None or sfa is None:
+        packed, sfa = _alloc_fp4(x.shape[0], x.shape[1], x.device)
+    ops.residual_add_rms_norm_fp4_sfa_fp16(residual, x, packed, sfa)
+    return packed, sfa
+
+
+def residual_add_rms_norm_fp4_sfa_v2_fp16(
+    residual: torch.Tensor,
+    x: torch.Tensor,
+    packed: torch.Tensor | None = None,
+    sfa: torch.Tensor | None = None,
+):
+    if packed is None or sfa is None:
+        packed, sfa = _alloc_fp4(x.shape[0], x.shape[1], x.device)
+    ops.residual_add_rms_norm_fp4_sfa_v2_fp16(residual, x, packed, sfa)
+    return packed, sfa
+
+
+def residual_add_rms_norm_mul_fp4_sfa_fp16(
+    residual: torch.Tensor,
+    x: torch.Tensor,
+    inv_s: torch.Tensor,
+    packed: torch.Tensor | None = None,
+    sfa: torch.Tensor | None = None,
+):
+    if packed is None or sfa is None:
+        packed, sfa = _alloc_fp4(x.shape[0], x.shape[1], x.device)
+    ops.residual_add_rms_norm_mul_fp4_sfa_fp16(residual, x, inv_s, packed, sfa)
+    return packed, sfa
+
+
+def silu_mul_fp4_sfa_fp16(merged: torch.Tensor, packed: torch.Tensor | None = None, sfa: torch.Tensor | None = None):
+    hidden = merged.shape[1] // 2
+    if packed is None or sfa is None:
+        packed, sfa = _alloc_fp4(merged.shape[0], hidden, merged.device)
+    ops.silu_mul_fp4_sfa_fp16(merged, packed, sfa)
+    return packed, sfa
+
+
+def silu_mul_fp4_sfa_v2_fp16(merged: torch.Tensor, packed: torch.Tensor | None = None, sfa: torch.Tensor | None = None):
+    hidden = merged.shape[1] // 2
+    if packed is None or sfa is None:
+        packed, sfa = _alloc_fp4(merged.shape[0], hidden, merged.device)
+    ops.silu_mul_fp4_sfa_v2_fp16(merged, packed, sfa)
+    return packed, sfa
+
+
+def silu_mul_mul_fp4_sfa_v2_fp16(
+    merged: torch.Tensor,
+    inv_s: torch.Tensor,
+    packed: torch.Tensor | None = None,
+    sfa: torch.Tensor | None = None,
+):
+    hidden = merged.shape[1] // 2
+    if packed is None or sfa is None:
+        packed, sfa = _alloc_fp4(merged.shape[0], hidden, merged.device)
+    ops.silu_mul_mul_fp4_sfa_v2_fp16(merged, inv_s, packed, sfa)
+    return packed, sfa
+
+
+def silu_mul_two_fp4_to_fp4(
+    gate_packed: torch.Tensor,
+    gate_sfa: torch.Tensor,
+    up_packed: torch.Tensor,
+    up_sfa: torch.Tensor,
+    out_packed: torch.Tensor | None = None,
+    out_sfa: torch.Tensor | None = None,
+):
+    hidden = gate_packed.shape[1] * 2
+    if out_packed is None or out_sfa is None:
+        out_packed, out_sfa = _alloc_fp4(gate_packed.shape[0], hidden, gate_packed.device)
+    ops.silu_mul_two_fp4_to_fp4(gate_packed, gate_sfa, up_packed, up_sfa, out_packed, out_sfa)
+    return out_packed, out_sfa
+
+
+def silu_mul_two_mul_fp4_to_fp4(
+    gate_packed: torch.Tensor,
+    gate_sfa: torch.Tensor,
+    up_packed: torch.Tensor,
+    up_sfa: torch.Tensor,
+    inv_s: torch.Tensor,
+    out_packed: torch.Tensor | None = None,
+    out_sfa: torch.Tensor | None = None,
+):
+    hidden = gate_packed.shape[1] * 2
+    if out_packed is None or out_sfa is None:
+        out_packed, out_sfa = _alloc_fp4(gate_packed.shape[0], hidden, gate_packed.device)
+    ops.silu_mul_two_mul_fp4_to_fp4(gate_packed, gate_sfa, up_packed, up_sfa, inv_s, out_packed, out_sfa)
+    return out_packed, out_sfa
+
+
+def dequantize_fp4_sfa_fp16(
+    packed: torch.Tensor,
+    sfa: torch.Tensor,
+    out: torch.Tensor | None = None,
+) -> torch.Tensor:
+    if out is None:
+        out = torch.empty((packed.shape[0], packed.shape[1] * 2), device=packed.device, dtype=torch.float16)
+    ops.dequantize_fp4_sfa_fp16(packed, sfa, out)
+    return out
+
+
+def rms_silu_nvfp4_ndhwc_bf16(
+    x: torch.Tensor,
+    gamma: torch.Tensor,
+    awq_inv_scale: torch.Tensor | None = None,
+    eps: float = 1e-6,
+    *,
+    packed: torch.Tensor | None = None,
+    scale_factors: torch.Tensor | None = None,
+) -> tuple[torch.Tensor, torch.Tensor]:
+    """Fuse BF16 NCDHW RMSNorm, SiLU and linear-SF NVFP4 quantization."""
+
+    b, c, t, h, w = x.shape
+    if packed is None:
+        packed = torch.empty(
+            (b, t, h, w, c // 2), device=x.device, dtype=torch.uint8
+        )
+    if scale_factors is None:
+        scale_factors = torch.empty(
+            (b, t, h, w, c // 16), device=x.device, dtype=torch.uint8
+        )
+    ops.rms_silu_nvfp4_ndhwc_bf16(
+        x,
+        gamma,
+        awq_inv_scale,
+        float(eps),
+        packed,
+        scale_factors,
+    )
+    return packed, scale_factors
+
+
+def quantize_bf16_to_nvfp4_linear(
+    input: torch.Tensor,
+    *,
+    packed: torch.Tensor | None = None,
+    scale_factors: torch.Tensor | None = None,
+) -> tuple[torch.Tensor, torch.Tensor]:
+    """Quantize BF16 rows to packed E2M1 with linear UE4M3 block-16 scales."""
+
+    rows, cols = input.shape
+    if packed is None:
+        packed = torch.empty(
+            (rows, cols // 2), device=input.device, dtype=torch.uint8
+        )
+    if scale_factors is None:
+        scale_factors = torch.empty(
+            (rows, cols // 16), device=input.device, dtype=torch.uint8
+        )
+    ops.quantize_bf16_to_nvfp4_linear(input, packed, scale_factors)
+    return packed, scale_factors
+
+
+def _check_ncdhw(x: torch.Tensor, gamma: torch.Tensor) -> None:
+    if x.dim() != 5:
+        raise RuntimeError("x must have shape (B,C,T,H,W)")
+    if x.shape[1] <= 0 or x.shape[1] % 2 != 0 or x.shape[1] > 1024:
+        raise RuntimeError("C must be even and at most 1024")
+    if gamma.shape != (x.shape[1],):
+        raise RuntimeError("gamma must have shape (C,)")
+
+
+@torch.library.register_fake(add_op_namespace_prefix("bf16_rms_silu_ncdhw"))
+def _bf16_rms_silu_ncdhw_fake(
+    x: torch.Tensor,
+    gamma: torch.Tensor,
+    prev_cache: torch.Tensor | None,
+    eps: float,
+    out: torch.Tensor,
+    next_cache: torch.Tensor | None,
+) -> None:
+    del eps
+    _check_ncdhw(x, gamma)
+    if out.shape != x.shape:
+        raise RuntimeError("out must have the same shape as x")
+    expected_cache = (x.shape[0], x.shape[1], 2, x.shape[3], x.shape[4])
+    if prev_cache is not None and prev_cache.shape != expected_cache:
+        raise RuntimeError("prev_cache must have shape (B,C,2,H,W)")
+    if next_cache is not None and next_cache.shape != expected_cache:
+        raise RuntimeError("next_cache must have shape (B,C,2,H,W)")
+    return None
+
+
+@torch.library.register_fake(add_op_namespace_prefix("bf16_rms_norm_ncdhw"))
+def _bf16_rms_norm_ncdhw_fake(
+    x: torch.Tensor,
+    gamma: torch.Tensor,
+    bias: torch.Tensor | None,
+    eps: float,
+    out: torch.Tensor,
+) -> None:
+    del eps
+    _check_ncdhw(x, gamma)
+    if bias is not None and bias.shape != (x.shape[1],):
+        raise RuntimeError("bias must have shape (C,)")
+    if out.shape != x.shape:
+        raise RuntimeError("out must have the same shape as x")
+    return None
+
+
+def bf16_rms_silu_ncdhw(
+    x: torch.Tensor,
+    gamma: torch.Tensor,
+    prev_cache: torch.Tensor | None = None,
+    eps: float = 1e-6,
+    *,
+    out: torch.Tensor | None = None,
+    next_cache: torch.Tensor | None = None,
+) -> tuple[torch.Tensor, torch.Tensor | None]:
+    """Apply channel RMSNorm and SiLU to BF16 NCDHW with optional cache update."""
+
+    if out is None:
+        out = torch.empty_like(x)
+    ops.bf16_rms_silu_ncdhw(
+        x, gamma, prev_cache, float(eps), out, next_cache
+    )
+    return out, next_cache
+
+
+def bf16_rms_norm_ncdhw(
+    x: torch.Tensor,
+    gamma: torch.Tensor,
+    bias: torch.Tensor | None = None,
+    eps: float = 1e-6,
+    *,
+    out: torch.Tensor | None = None,
+) -> torch.Tensor:
+    """Apply channel RMSNorm and optional bias to BF16 NCDHW."""
+
+    if out is None:
+        out = torch.empty_like(x)
+    ops.bf16_rms_norm_ncdhw(x, gamma, bias, float(eps), out)
+    return out
+
+
+__all__ = [
+    "bf16_rms_norm_ncdhw",
+    "bf16_rms_silu_ncdhw",
+    "dequantize_fp4_sfa_fp16",
+    "quantize_bf16_to_nvfp4_linear",
+    "residual_add_rms_norm_fp4_sfa_fp16",
+    "residual_add_rms_norm_fp4_sfa_v2_fp16",
+    "residual_add_rms_norm_mul_fp4_sfa_fp16",
+    "rms_norm_fp4_sfa_fp16",
+    "rms_silu_nvfp4_ndhwc_bf16",
+    "sfa_size_bytes",
+    "silu_mul_fp4_sfa_fp16",
+    "silu_mul_fp4_sfa_v2_fp16",
+    "silu_mul_mul_fp4_sfa_v2_fp16",
+    "silu_mul_two_fp4_to_fp4",
+    "silu_mul_two_mul_fp4_to_fp4",
+]

build/{torch211-cxx11-cu128-x86_64-linux/_fp4_fused_ops_cuda_7e2e5b9.abi3.so → torch211-cxx11-cu130-aarch64-linux/_fp4_fused_ops_cuda_916f685.abi3.so}

@@ -1,3 +1,3 @@
 version https://git-lfs.github.com/spec/v1
-oid sha256:dd0f52e71a85ae9cc7a33eddf96424a19791aa48f991fad9a358d12a313bae0b
-size 502488
+oid sha256:1da63bd51d661880b027e0efaa004f80f1a4ecb1c60e67395143a5e07ebc90ff
+size 689560

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

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

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

@@ -0,0 +1,14 @@
+import ctypes
+import importlib.util
+import sys
+from pathlib import Path
+
+def _import_from_path(file_path: Path):
+    path_hash = '{:x}'.format(ctypes.c_size_t(hash(file_path.absolute())).value)
+    spec = importlib.util.spec_from_file_location(path_hash, file_path)
+    module = importlib.util.module_from_spec(spec)
+    sys.modules[path_hash] = module
+    spec.loader.exec_module(module)
+    return module
+
+globals().update(vars(_import_from_path(Path(__file__).parent.parent / '__init__.py')))

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

@@ -0,0 +1,32 @@
+{
+  "name": "fp4-fused-ops",
+  "id": "_fp4_fused_ops_cuda_916f685",
+  "version": 1,
+  "license": "Apache-2.0",
+  "python-depends": [],
+  "backend": {
+    "type": "cuda",
+    "archs": [
+      "11.0a"
+    ]
+  },
+  "digest": {
+    "algorithm": "sha256",
+    "files": {
+      "__init__.py": "511/VkU+HAwubcrzG/Ck9wB3ukPbjKrWxG+9EgtcP00=",
+      "_fp4_fused_ops_cuda_916f685.abi3.so": "HaY71R1mGICwJ+DvqgBPgPGk7LHGDmc5UUOl4H68kP8=",
+      "_ops.py": "IFiCUaM6wO9cSSX0QcNr+RWvWYMiOy8A+hgr+8LClME=",
+      "fp4_fused_ops/__init__.py": "v6p5XMfQzddhi1fLSAw4HX9CyS0rQsidvu9VsT01xi4="
+    }
+  },
+  "provenance": {
+    "kernel": {
+      "sha": "4c95454360b18c79bed75ec6346a7bac47ae2ca2",
+      "dirty": false
+    },
+    "validation": {
+      "torch": "2.11.0+cu130",
+      "cuda": "13.0"
+    }
+  }
+}

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

@@ -91,6 +91,57 @@ def _dequant_fake(packed: torch.Tensor, sfa: torch.Tensor, out: torch.Tensor) ->
     return None
 
 
+@torch.library.register_fake(add_op_namespace_prefix("rms_silu_nvfp4_ndhwc_bf16"))
+def _rms_silu_nvfp4_ndhwc_bf16_fake(
+    x: torch.Tensor,
+    gamma: torch.Tensor,
+    awq_inv_scale: torch.Tensor | None,
+    eps: float,
+    packed: torch.Tensor,
+    scale_factors: torch.Tensor,
+) -> None:
+    del eps
+    if x.dim() != 5:
+        raise RuntimeError("x must have shape (B,C,T,H,W)")
+    b, c, t, h, w = x.shape
+    if (
+        c % 128 != 0
+        or c > 1024
+        or gamma.shape != (c,)
+        or packed.shape != (b, t, h, w, c // 2)
+        or scale_factors.shape != (b, t, h, w, c // 16)
+        or (awq_inv_scale is not None and awq_inv_scale.shape != (c,))
+    ):
+        raise RuntimeError(
+            "rms_silu_nvfp4_ndhwc_bf16 shape contract failed"
+        )
+    return None
+
+
+@torch.library.register_fake(
+    add_op_namespace_prefix("quantize_bf16_to_nvfp4_linear")
+)
+def _quantize_bf16_to_nvfp4_linear_fake(
+    input: torch.Tensor,
+    packed: torch.Tensor,
+    scale_factors: torch.Tensor,
+) -> None:
+    if input.dim() != 2:
+        raise RuntimeError("input must have shape (rows, cols)")
+    rows, cols = input.shape
+    if (
+        rows <= 0
+        or cols <= 0
+        or cols % 16 != 0
+        or packed.shape != (rows, cols // 2)
+        or scale_factors.shape != (rows, cols // 16)
+    ):
+        raise RuntimeError(
+            "quantize_bf16_to_nvfp4_linear shape contract failed"
+        )
+    return None
+
+
 def rms_norm_fp4_sfa_fp16(x: torch.Tensor, packed: torch.Tensor | None = None, sfa: torch.Tensor | None = None):
     if packed is None or sfa is None:
         packed, sfa = _alloc_fp4(x.shape[0], x.shape[1], x.device)
@@ -204,3 +255,156 @@ def dequantize_fp4_sfa_fp16(
         out = torch.empty((packed.shape[0], packed.shape[1] * 2), device=packed.device, dtype=torch.float16)
     ops.dequantize_fp4_sfa_fp16(packed, sfa, out)
     return out
+
+
+def rms_silu_nvfp4_ndhwc_bf16(
+    x: torch.Tensor,
+    gamma: torch.Tensor,
+    awq_inv_scale: torch.Tensor | None = None,
+    eps: float = 1e-6,
+    *,
+    packed: torch.Tensor | None = None,
+    scale_factors: torch.Tensor | None = None,
+) -> tuple[torch.Tensor, torch.Tensor]:
+    """Fuse BF16 NCDHW RMSNorm, SiLU and linear-SF NVFP4 quantization."""
+
+    b, c, t, h, w = x.shape
+    if packed is None:
+        packed = torch.empty(
+            (b, t, h, w, c // 2), device=x.device, dtype=torch.uint8
+        )
+    if scale_factors is None:
+        scale_factors = torch.empty(
+            (b, t, h, w, c // 16), device=x.device, dtype=torch.uint8
+        )
+    ops.rms_silu_nvfp4_ndhwc_bf16(
+        x,
+        gamma,
+        awq_inv_scale,
+        float(eps),
+        packed,
+        scale_factors,
+    )
+    return packed, scale_factors
+
+
+def quantize_bf16_to_nvfp4_linear(
+    input: torch.Tensor,
+    *,
+    packed: torch.Tensor | None = None,
+    scale_factors: torch.Tensor | None = None,
+) -> tuple[torch.Tensor, torch.Tensor]:
+    """Quantize BF16 rows to packed E2M1 with linear UE4M3 block-16 scales."""
+
+    rows, cols = input.shape
+    if packed is None:
+        packed = torch.empty(
+            (rows, cols // 2), device=input.device, dtype=torch.uint8
+        )
+    if scale_factors is None:
+        scale_factors = torch.empty(
+            (rows, cols // 16), device=input.device, dtype=torch.uint8
+        )
+    ops.quantize_bf16_to_nvfp4_linear(input, packed, scale_factors)
+    return packed, scale_factors
+
+
+def _check_ncdhw(x: torch.Tensor, gamma: torch.Tensor) -> None:
+    if x.dim() != 5:
+        raise RuntimeError("x must have shape (B,C,T,H,W)")
+    if x.shape[1] <= 0 or x.shape[1] % 2 != 0 or x.shape[1] > 1024:
+        raise RuntimeError("C must be even and at most 1024")
+    if gamma.shape != (x.shape[1],):
+        raise RuntimeError("gamma must have shape (C,)")
+
+
+@torch.library.register_fake(add_op_namespace_prefix("bf16_rms_silu_ncdhw"))
+def _bf16_rms_silu_ncdhw_fake(
+    x: torch.Tensor,
+    gamma: torch.Tensor,
+    prev_cache: torch.Tensor | None,
+    eps: float,
+    out: torch.Tensor,
+    next_cache: torch.Tensor | None,
+) -> None:
+    del eps
+    _check_ncdhw(x, gamma)
+    if out.shape != x.shape:
+        raise RuntimeError("out must have the same shape as x")
+    expected_cache = (x.shape[0], x.shape[1], 2, x.shape[3], x.shape[4])
+    if prev_cache is not None and prev_cache.shape != expected_cache:
+        raise RuntimeError("prev_cache must have shape (B,C,2,H,W)")
+    if next_cache is not None and next_cache.shape != expected_cache:
+        raise RuntimeError("next_cache must have shape (B,C,2,H,W)")
+    return None
+
+
+@torch.library.register_fake(add_op_namespace_prefix("bf16_rms_norm_ncdhw"))
+def _bf16_rms_norm_ncdhw_fake(
+    x: torch.Tensor,
+    gamma: torch.Tensor,
+    bias: torch.Tensor | None,
+    eps: float,
+    out: torch.Tensor,
+) -> None:
+    del eps
+    _check_ncdhw(x, gamma)
+    if bias is not None and bias.shape != (x.shape[1],):
+        raise RuntimeError("bias must have shape (C,)")
+    if out.shape != x.shape:
+        raise RuntimeError("out must have the same shape as x")
+    return None
+
+
+def bf16_rms_silu_ncdhw(
+    x: torch.Tensor,
+    gamma: torch.Tensor,
+    prev_cache: torch.Tensor | None = None,
+    eps: float = 1e-6,
+    *,
+    out: torch.Tensor | None = None,
+    next_cache: torch.Tensor | None = None,
+) -> tuple[torch.Tensor, torch.Tensor | None]:
+    """Apply channel RMSNorm and SiLU to BF16 NCDHW with optional cache update."""
+
+    if out is None:
+        out = torch.empty_like(x)
+    ops.bf16_rms_silu_ncdhw(
+        x, gamma, prev_cache, float(eps), out, next_cache
+    )
+    return out, next_cache
+
+
+def bf16_rms_norm_ncdhw(
+    x: torch.Tensor,
+    gamma: torch.Tensor,
+    bias: torch.Tensor | None = None,
+    eps: float = 1e-6,
+    *,
+    out: torch.Tensor | None = None,
+) -> torch.Tensor:
+    """Apply channel RMSNorm and optional bias to BF16 NCDHW."""
+
+    if out is None:
+        out = torch.empty_like(x)
+    ops.bf16_rms_norm_ncdhw(x, gamma, bias, float(eps), out)
+    return out
+
+
+__all__ = [
+    "bf16_rms_norm_ncdhw",
+    "bf16_rms_silu_ncdhw",
+    "dequantize_fp4_sfa_fp16",
+    "quantize_bf16_to_nvfp4_linear",
+    "residual_add_rms_norm_fp4_sfa_fp16",
+    "residual_add_rms_norm_fp4_sfa_v2_fp16",
+    "residual_add_rms_norm_mul_fp4_sfa_fp16",
+    "rms_norm_fp4_sfa_fp16",
+    "rms_silu_nvfp4_ndhwc_bf16",
+    "sfa_size_bytes",
+    "silu_mul_fp4_sfa_fp16",
+    "silu_mul_fp4_sfa_v2_fp16",
+    "silu_mul_mul_fp4_sfa_v2_fp16",
+    "silu_mul_two_fp4_to_fp4",
+    "silu_mul_two_mul_fp4_to_fp4",
+]

build/torch211-cxx11-cu130-x86_64-linux/{_fp4_fused_ops_cuda_7e2e5b9.abi3.so → _fp4_fused_ops_cuda_8e0b1cd.abi3.so}

@@ -1,3 +1,3 @@
 version https://git-lfs.github.com/spec/v1
-oid sha256:18ee567014dff126b5d85586845d69c6463950d1fc1ef19d528215a71ea3f6be
-size 535352
+oid sha256:8b0f8e8fe5865d1bdc245b48c0fb71e0910acd89e4405d91c9e4c1ae08af938a
+size 966768

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

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

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

@@ -1,21 +1,22 @@
 {
   "name": "fp4-fused-ops",
-  "id": "_fp4_fused_ops_cuda_7e2e5b9",
+  "id": "_fp4_fused_ops_cuda_8e0b1cd",
   "version": 1,
   "license": "Apache-2.0",
   "python-depends": [],
   "backend": {
     "type": "cuda",
     "archs": [
+      "11.0a",
       "12.0a"
     ]
   },
   "digest": {
     "algorithm": "sha256",
     "files": {
-      "__init__.py": "aB7foEHGoXVq7xDlENZhAsIjQvzWhMpiKhanJhtzmk8=",
-      "_fp4_fused_ops_cuda_7e2e5b9.abi3.so": "GO5WcBTf8Sa12FWGhF1pxkY5UNH8HvGdUoIVpx6j9r4=",
-      "_ops.py": "qtD0ljSFCLkwC71lYCEP/FcSj+NFL+0sWWugI+4CH4s=",
+      "__init__.py": "511/VkU+HAwubcrzG/Ck9wB3ukPbjKrWxG+9EgtcP00=",
+      "_fp4_fused_ops_cuda_8e0b1cd.abi3.so": "iw+Oj+WGXRvcJFtIwPtx4JEKzYnkQF2RyeTBrgivk4o=",
+      "_ops.py": "L5M3920Afl87xMKszpaz5zL32AN6JoBJM6Aw1J19k2I=",
       "fp4_fused_ops/__init__.py": "DFYPlrhXwYjEqCl/8n0SmWGZV8NFml5DPhMjKfv98GY="
     }
   }

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

@@ -91,6 +91,57 @@ def _dequant_fake(packed: torch.Tensor, sfa: torch.Tensor, out: torch.Tensor) ->
     return None
 
 
+@torch.library.register_fake(add_op_namespace_prefix("rms_silu_nvfp4_ndhwc_bf16"))
+def _rms_silu_nvfp4_ndhwc_bf16_fake(
+    x: torch.Tensor,
+    gamma: torch.Tensor,
+    awq_inv_scale: torch.Tensor | None,
+    eps: float,
+    packed: torch.Tensor,
+    scale_factors: torch.Tensor,
+) -> None:
+    del eps
+    if x.dim() != 5:
+        raise RuntimeError("x must have shape (B,C,T,H,W)")
+    b, c, t, h, w = x.shape
+    if (
+        c % 128 != 0
+        or c > 1024
+        or gamma.shape != (c,)
+        or packed.shape != (b, t, h, w, c // 2)
+        or scale_factors.shape != (b, t, h, w, c // 16)
+        or (awq_inv_scale is not None and awq_inv_scale.shape != (c,))
+    ):
+        raise RuntimeError(
+            "rms_silu_nvfp4_ndhwc_bf16 shape contract failed"
+        )
+    return None
+
+
+@torch.library.register_fake(
+    add_op_namespace_prefix("quantize_bf16_to_nvfp4_linear")
+)
+def _quantize_bf16_to_nvfp4_linear_fake(
+    input: torch.Tensor,
+    packed: torch.Tensor,
+    scale_factors: torch.Tensor,
+) -> None:
+    if input.dim() != 2:
+        raise RuntimeError("input must have shape (rows, cols)")
+    rows, cols = input.shape
+    if (
+        rows <= 0
+        or cols <= 0
+        or cols % 16 != 0
+        or packed.shape != (rows, cols // 2)
+        or scale_factors.shape != (rows, cols // 16)
+    ):
+        raise RuntimeError(
+            "quantize_bf16_to_nvfp4_linear shape contract failed"
+        )
+    return None
+
+
 def rms_norm_fp4_sfa_fp16(x: torch.Tensor, packed: torch.Tensor | None = None, sfa: torch.Tensor | None = None):
     if packed is None or sfa is None:
         packed, sfa = _alloc_fp4(x.shape[0], x.shape[1], x.device)
@@ -204,3 +255,156 @@ def dequantize_fp4_sfa_fp16(
         out = torch.empty((packed.shape[0], packed.shape[1] * 2), device=packed.device, dtype=torch.float16)
     ops.dequantize_fp4_sfa_fp16(packed, sfa, out)
     return out
+
+
+def rms_silu_nvfp4_ndhwc_bf16(
+    x: torch.Tensor,
+    gamma: torch.Tensor,
+    awq_inv_scale: torch.Tensor | None = None,
+    eps: float = 1e-6,
+    *,
+    packed: torch.Tensor | None = None,
+    scale_factors: torch.Tensor | None = None,
+) -> tuple[torch.Tensor, torch.Tensor]:
+    """Fuse BF16 NCDHW RMSNorm, SiLU and linear-SF NVFP4 quantization."""
+
+    b, c, t, h, w = x.shape
+    if packed is None:
+        packed = torch.empty(
+            (b, t, h, w, c // 2), device=x.device, dtype=torch.uint8
+        )
+    if scale_factors is None:
+        scale_factors = torch.empty(
+            (b, t, h, w, c // 16), device=x.device, dtype=torch.uint8
+        )
+    ops.rms_silu_nvfp4_ndhwc_bf16(
+        x,
+        gamma,
+        awq_inv_scale,
+        float(eps),
+        packed,
+        scale_factors,
+    )
+    return packed, scale_factors
+
+
+def quantize_bf16_to_nvfp4_linear(
+    input: torch.Tensor,
+    *,
+    packed: torch.Tensor | None = None,
+    scale_factors: torch.Tensor | None = None,
+) -> tuple[torch.Tensor, torch.Tensor]:
+    """Quantize BF16 rows to packed E2M1 with linear UE4M3 block-16 scales."""
+
+    rows, cols = input.shape
+    if packed is None:
+        packed = torch.empty(
+            (rows, cols // 2), device=input.device, dtype=torch.uint8
+        )
+    if scale_factors is None:
+        scale_factors = torch.empty(
+            (rows, cols // 16), device=input.device, dtype=torch.uint8
+        )
+    ops.quantize_bf16_to_nvfp4_linear(input, packed, scale_factors)
+    return packed, scale_factors
+
+
+def _check_ncdhw(x: torch.Tensor, gamma: torch.Tensor) -> None:
+    if x.dim() != 5:
+        raise RuntimeError("x must have shape (B,C,T,H,W)")
+    if x.shape[1] <= 0 or x.shape[1] % 2 != 0 or x.shape[1] > 1024:
+        raise RuntimeError("C must be even and at most 1024")
+    if gamma.shape != (x.shape[1],):
+        raise RuntimeError("gamma must have shape (C,)")
+
+
+@torch.library.register_fake(add_op_namespace_prefix("bf16_rms_silu_ncdhw"))
+def _bf16_rms_silu_ncdhw_fake(
+    x: torch.Tensor,
+    gamma: torch.Tensor,
+    prev_cache: torch.Tensor | None,
+    eps: float,
+    out: torch.Tensor,
+    next_cache: torch.Tensor | None,
+) -> None:
+    del eps
+    _check_ncdhw(x, gamma)
+    if out.shape != x.shape:
+        raise RuntimeError("out must have the same shape as x")
+    expected_cache = (x.shape[0], x.shape[1], 2, x.shape[3], x.shape[4])
+    if prev_cache is not None and prev_cache.shape != expected_cache:
+        raise RuntimeError("prev_cache must have shape (B,C,2,H,W)")
+    if next_cache is not None and next_cache.shape != expected_cache:
+        raise RuntimeError("next_cache must have shape (B,C,2,H,W)")
+    return None
+
+
+@torch.library.register_fake(add_op_namespace_prefix("bf16_rms_norm_ncdhw"))
+def _bf16_rms_norm_ncdhw_fake(
+    x: torch.Tensor,
+    gamma: torch.Tensor,
+    bias: torch.Tensor | None,
+    eps: float,
+    out: torch.Tensor,
+) -> None:
+    del eps
+    _check_ncdhw(x, gamma)
+    if bias is not None and bias.shape != (x.shape[1],):
+        raise RuntimeError("bias must have shape (C,)")
+    if out.shape != x.shape:
+        raise RuntimeError("out must have the same shape as x")
+    return None
+
+
+def bf16_rms_silu_ncdhw(
+    x: torch.Tensor,
+    gamma: torch.Tensor,
+    prev_cache: torch.Tensor | None = None,
+    eps: float = 1e-6,
+    *,
+    out: torch.Tensor | None = None,
+    next_cache: torch.Tensor | None = None,
+) -> tuple[torch.Tensor, torch.Tensor | None]:
+    """Apply channel RMSNorm and SiLU to BF16 NCDHW with optional cache update."""
+
+    if out is None:
+        out = torch.empty_like(x)
+    ops.bf16_rms_silu_ncdhw(
+        x, gamma, prev_cache, float(eps), out, next_cache
+    )
+    return out, next_cache
+
+
+def bf16_rms_norm_ncdhw(
+    x: torch.Tensor,
+    gamma: torch.Tensor,
+    bias: torch.Tensor | None = None,
+    eps: float = 1e-6,
+    *,
+    out: torch.Tensor | None = None,
+) -> torch.Tensor:
+    """Apply channel RMSNorm and optional bias to BF16 NCDHW."""
+
+    if out is None:
+        out = torch.empty_like(x)
+    ops.bf16_rms_norm_ncdhw(x, gamma, bias, float(eps), out)
+    return out
+
+
+__all__ = [
+    "bf16_rms_norm_ncdhw",
+    "bf16_rms_silu_ncdhw",
+    "dequantize_fp4_sfa_fp16",
+    "quantize_bf16_to_nvfp4_linear",
+    "residual_add_rms_norm_fp4_sfa_fp16",
+    "residual_add_rms_norm_fp4_sfa_v2_fp16",
+    "residual_add_rms_norm_mul_fp4_sfa_fp16",
+    "rms_norm_fp4_sfa_fp16",
+    "rms_silu_nvfp4_ndhwc_bf16",
+    "sfa_size_bytes",
+    "silu_mul_fp4_sfa_fp16",
+    "silu_mul_fp4_sfa_v2_fp16",
+    "silu_mul_mul_fp4_sfa_v2_fp16",
+    "silu_mul_two_fp4_to_fp4",
+    "silu_mul_two_mul_fp4_to_fp4",
+]

build/torch212-cxx11-cu130-x86_64-linux/{_fp4_fused_ops_cuda_7e2e5b9.abi3.so → _fp4_fused_ops_cuda_8e0b1cd.abi3.so}

@@ -1,3 +1,3 @@
 version https://git-lfs.github.com/spec/v1
-oid sha256:f24cde57e006bc22f3e087331b3bf75f223e796eb8526444792d779ac5c191b0
-size 545840
+oid sha256:3b9aede5c54c9d74330dcd256ebecf8b77ef96ce5d0561e03f1996b625395c85
+size 973216

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

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

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

@@ -1,21 +1,22 @@
 {
   "name": "fp4-fused-ops",
-  "id": "_fp4_fused_ops_cuda_7e2e5b9",
+  "id": "_fp4_fused_ops_cuda_8e0b1cd",
   "version": 1,
   "license": "Apache-2.0",
   "python-depends": [],
   "backend": {
     "type": "cuda",
     "archs": [
+      "11.0a",
       "12.0a"
     ]
   },
   "digest": {
     "algorithm": "sha256",
     "files": {
-      "__init__.py": "aB7foEHGoXVq7xDlENZhAsIjQvzWhMpiKhanJhtzmk8=",
-      "_fp4_fused_ops_cuda_7e2e5b9.abi3.so": "8kzeV+AGvCLz4IczGzv3XyI+eW64UmREeS13msXBkbA=",
-      "_ops.py": "qtD0ljSFCLkwC71lYCEP/FcSj+NFL+0sWWugI+4CH4s=",
+      "__init__.py": "511/VkU+HAwubcrzG/Ck9wB3ukPbjKrWxG+9EgtcP00=",
+      "_fp4_fused_ops_cuda_8e0b1cd.abi3.so": "O5rt5cVMnXQzDc0lbr7Pi3fvls5dBWHgPxmWtiU5XIU=",
+      "_ops.py": "L5M3920Afl87xMKszpaz5zL32AN6JoBJM6Aw1J19k2I=",
       "fp4_fused_ops/__init__.py": "DFYPlrhXwYjEqCl/8n0SmWGZV8NFml5DPhMjKfv98GY="
     }
   }

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

@@ -91,6 +91,57 @@ def _dequant_fake(packed: torch.Tensor, sfa: torch.Tensor, out: torch.Tensor) ->
     return None
 
 
+@torch.library.register_fake(add_op_namespace_prefix("rms_silu_nvfp4_ndhwc_bf16"))
+def _rms_silu_nvfp4_ndhwc_bf16_fake(
+    x: torch.Tensor,
+    gamma: torch.Tensor,
+    awq_inv_scale: torch.Tensor | None,
+    eps: float,
+    packed: torch.Tensor,
+    scale_factors: torch.Tensor,
+) -> None:
+    del eps
+    if x.dim() != 5:
+        raise RuntimeError("x must have shape (B,C,T,H,W)")
+    b, c, t, h, w = x.shape
+    if (
+        c % 128 != 0
+        or c > 1024
+        or gamma.shape != (c,)
+        or packed.shape != (b, t, h, w, c // 2)
+        or scale_factors.shape != (b, t, h, w, c // 16)
+        or (awq_inv_scale is not None and awq_inv_scale.shape != (c,))
+    ):
+        raise RuntimeError(
+            "rms_silu_nvfp4_ndhwc_bf16 shape contract failed"
+        )
+    return None
+
+
+@torch.library.register_fake(
+    add_op_namespace_prefix("quantize_bf16_to_nvfp4_linear")
+)
+def _quantize_bf16_to_nvfp4_linear_fake(
+    input: torch.Tensor,
+    packed: torch.Tensor,
+    scale_factors: torch.Tensor,
+) -> None:
+    if input.dim() != 2:
+        raise RuntimeError("input must have shape (rows, cols)")
+    rows, cols = input.shape
+    if (
+        rows <= 0
+        or cols <= 0
+        or cols % 16 != 0
+        or packed.shape != (rows, cols // 2)
+        or scale_factors.shape != (rows, cols // 16)
+    ):
+        raise RuntimeError(
+            "quantize_bf16_to_nvfp4_linear shape contract failed"
+        )
+    return None
+
+
 def rms_norm_fp4_sfa_fp16(x: torch.Tensor, packed: torch.Tensor | None = None, sfa: torch.Tensor | None = None):
     if packed is None or sfa is None:
         packed, sfa = _alloc_fp4(x.shape[0], x.shape[1], x.device)
@@ -204,3 +255,156 @@ def dequantize_fp4_sfa_fp16(
         out = torch.empty((packed.shape[0], packed.shape[1] * 2), device=packed.device, dtype=torch.float16)
     ops.dequantize_fp4_sfa_fp16(packed, sfa, out)
     return out
+
+
+def rms_silu_nvfp4_ndhwc_bf16(
+    x: torch.Tensor,
+    gamma: torch.Tensor,
+    awq_inv_scale: torch.Tensor | None = None,
+    eps: float = 1e-6,
+    *,
+    packed: torch.Tensor | None = None,
+    scale_factors: torch.Tensor | None = None,
+) -> tuple[torch.Tensor, torch.Tensor]:
+    """Fuse BF16 NCDHW RMSNorm, SiLU and linear-SF NVFP4 quantization."""
+
+    b, c, t, h, w = x.shape
+    if packed is None:
+        packed = torch.empty(
+            (b, t, h, w, c // 2), device=x.device, dtype=torch.uint8
+        )
+    if scale_factors is None:
+        scale_factors = torch.empty(
+            (b, t, h, w, c // 16), device=x.device, dtype=torch.uint8
+        )
+    ops.rms_silu_nvfp4_ndhwc_bf16(
+        x,
+        gamma,
+        awq_inv_scale,
+        float(eps),
+        packed,
+        scale_factors,
+    )
+    return packed, scale_factors
+
+
+def quantize_bf16_to_nvfp4_linear(
+    input: torch.Tensor,
+    *,
+    packed: torch.Tensor | None = None,
+    scale_factors: torch.Tensor | None = None,
+) -> tuple[torch.Tensor, torch.Tensor]:
+    """Quantize BF16 rows to packed E2M1 with linear UE4M3 block-16 scales."""
+
+    rows, cols = input.shape
+    if packed is None:
+        packed = torch.empty(
+            (rows, cols // 2), device=input.device, dtype=torch.uint8
+        )
+    if scale_factors is None:
+        scale_factors = torch.empty(
+            (rows, cols // 16), device=input.device, dtype=torch.uint8
+        )
+    ops.quantize_bf16_to_nvfp4_linear(input, packed, scale_factors)
+    return packed, scale_factors
+
+
+def _check_ncdhw(x: torch.Tensor, gamma: torch.Tensor) -> None:
+    if x.dim() != 5:
+        raise RuntimeError("x must have shape (B,C,T,H,W)")
+    if x.shape[1] <= 0 or x.shape[1] % 2 != 0 or x.shape[1] > 1024:
+        raise RuntimeError("C must be even and at most 1024")
+    if gamma.shape != (x.shape[1],):
+        raise RuntimeError("gamma must have shape (C,)")
+
+
+@torch.library.register_fake(add_op_namespace_prefix("bf16_rms_silu_ncdhw"))
+def _bf16_rms_silu_ncdhw_fake(
+    x: torch.Tensor,
+    gamma: torch.Tensor,
+    prev_cache: torch.Tensor | None,
+    eps: float,
+    out: torch.Tensor,
+    next_cache: torch.Tensor | None,
+) -> None:
+    del eps
+    _check_ncdhw(x, gamma)
+    if out.shape != x.shape:
+        raise RuntimeError("out must have the same shape as x")
+    expected_cache = (x.shape[0], x.shape[1], 2, x.shape[3], x.shape[4])
+    if prev_cache is not None and prev_cache.shape != expected_cache:
+        raise RuntimeError("prev_cache must have shape (B,C,2,H,W)")
+    if next_cache is not None and next_cache.shape != expected_cache:
+        raise RuntimeError("next_cache must have shape (B,C,2,H,W)")
+    return None
+
+
+@torch.library.register_fake(add_op_namespace_prefix("bf16_rms_norm_ncdhw"))
+def _bf16_rms_norm_ncdhw_fake(
+    x: torch.Tensor,
+    gamma: torch.Tensor,
+    bias: torch.Tensor | None,
+    eps: float,
+    out: torch.Tensor,
+) -> None:
+    del eps
+    _check_ncdhw(x, gamma)
+    if bias is not None and bias.shape != (x.shape[1],):
+        raise RuntimeError("bias must have shape (C,)")
+    if out.shape != x.shape:
+        raise RuntimeError("out must have the same shape as x")
+    return None
+
+
+def bf16_rms_silu_ncdhw(
+    x: torch.Tensor,
+    gamma: torch.Tensor,
+    prev_cache: torch.Tensor | None = None,
+    eps: float = 1e-6,
+    *,
+    out: torch.Tensor | None = None,
+    next_cache: torch.Tensor | None = None,
+) -> tuple[torch.Tensor, torch.Tensor | None]:
+    """Apply channel RMSNorm and SiLU to BF16 NCDHW with optional cache update."""
+
+    if out is None:
+        out = torch.empty_like(x)
+    ops.bf16_rms_silu_ncdhw(
+        x, gamma, prev_cache, float(eps), out, next_cache
+    )
+    return out, next_cache
+
+
+def bf16_rms_norm_ncdhw(
+    x: torch.Tensor,
+    gamma: torch.Tensor,
+    bias: torch.Tensor | None = None,
+    eps: float = 1e-6,
+    *,
+    out: torch.Tensor | None = None,
+) -> torch.Tensor:
+    """Apply channel RMSNorm and optional bias to BF16 NCDHW."""
+
+    if out is None:
+        out = torch.empty_like(x)
+    ops.bf16_rms_norm_ncdhw(x, gamma, bias, float(eps), out)
+    return out
+
+
+__all__ = [
+    "bf16_rms_norm_ncdhw",
+    "bf16_rms_silu_ncdhw",
+    "dequantize_fp4_sfa_fp16",
+    "quantize_bf16_to_nvfp4_linear",
+    "residual_add_rms_norm_fp4_sfa_fp16",
+    "residual_add_rms_norm_fp4_sfa_v2_fp16",
+    "residual_add_rms_norm_mul_fp4_sfa_fp16",
+    "rms_norm_fp4_sfa_fp16",
+    "rms_silu_nvfp4_ndhwc_bf16",
+    "sfa_size_bytes",
+    "silu_mul_fp4_sfa_fp16",
+    "silu_mul_fp4_sfa_v2_fp16",
+    "silu_mul_mul_fp4_sfa_v2_fp16",
+    "silu_mul_two_fp4_to_fp4",
+    "silu_mul_two_mul_fp4_to_fp4",
+]

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

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

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

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

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

@@ -1,21 +1,22 @@
 {
   "name": "fp4-fused-ops",
-  "id": "_fp4_fused_ops_cuda_7e2e5b9",
+  "id": "_fp4_fused_ops_cuda_8e0b1cd",
   "version": 1,
   "license": "Apache-2.0",
   "python-depends": [],
   "backend": {
     "type": "cuda",
     "archs": [
+      "11.0a",
       "12.0a"
     ]
   },
   "digest": {
     "algorithm": "sha256",
     "files": {
-      "__init__.py": "aB7foEHGoXVq7xDlENZhAsIjQvzWhMpiKhanJhtzmk8=",
-      "_fp4_fused_ops_cuda_7e2e5b9.abi3.so": "+7gD0nYPX/QgoxHcdBsRsoSBnlWNAGAdmjwvXpIdMog=",
-      "_ops.py": "qtD0ljSFCLkwC71lYCEP/FcSj+NFL+0sWWugI+4CH4s=",
+      "__init__.py": "511/VkU+HAwubcrzG/Ck9wB3ukPbjKrWxG+9EgtcP00=",
+      "_fp4_fused_ops_cuda_8e0b1cd.abi3.so": "mPvyxKJaC0dtZT8C62tQKJVIIAc0jhDO7aV8oS+4d4g=",
+      "_ops.py": "L5M3920Afl87xMKszpaz5zL32AN6JoBJM6Aw1J19k2I=",
       "fp4_fused_ops/__init__.py": "DFYPlrhXwYjEqCl/8n0SmWGZV8NFml5DPhMjKfv98GY="
     }
   }

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

@@ -0,0 +1,410 @@
+"""FlashRT FP4 fused producer kernels."""
+
+from __future__ import annotations
+
+import torch
+
+from ._ops import add_op_namespace_prefix, ops
+
+
+def sfa_size_bytes(rows: int, dim: int, is_sfb: bool = False, device: torch.device | str | None = None) -> int:
+    if device is None:
+        device = torch.device("cuda", torch.cuda.current_device())
+    anchor = torch.empty((1,), device=device, dtype=torch.uint8)
+    return int(ops.sfa_size_bytes_for(anchor, int(rows), int(dim), bool(is_sfb)))
+
+
+def _alloc_fp4(rows: int, dim: int, device: torch.device | str):
+    packed = torch.empty((rows, dim // 2), device=device, dtype=torch.uint8)
+    sfa = torch.empty((sfa_size_bytes(rows, dim, False, device=device),), device=device, dtype=torch.uint8)
+    return packed, sfa
+
+
+@torch.library.register_fake(add_op_namespace_prefix("rms_norm_fp4_sfa_fp16"))
+def _rms_norm_fake(x: torch.Tensor, packed: torch.Tensor, sfa: torch.Tensor) -> None:
+    return None
+
+
+@torch.library.register_fake(add_op_namespace_prefix("residual_add_rms_norm_fp4_sfa_fp16"))
+def _residual_fake(residual: torch.Tensor, x: torch.Tensor, packed: torch.Tensor, sfa: torch.Tensor) -> None:
+    return None
+
+
+@torch.library.register_fake(add_op_namespace_prefix("residual_add_rms_norm_fp4_sfa_v2_fp16"))
+def _residual_v2_fake(residual: torch.Tensor, x: torch.Tensor, packed: torch.Tensor, sfa: torch.Tensor) -> None:
+    return None
+
+
+@torch.library.register_fake(add_op_namespace_prefix("residual_add_rms_norm_mul_fp4_sfa_fp16"))
+def _residual_mul_fake(
+    residual: torch.Tensor,
+    x: torch.Tensor,
+    inv_s: torch.Tensor,
+    packed: torch.Tensor,
+    sfa: torch.Tensor,
+) -> None:
+    return None
+
+
+@torch.library.register_fake(add_op_namespace_prefix("silu_mul_fp4_sfa_fp16"))
+def _silu_fake(merged: torch.Tensor, packed: torch.Tensor, sfa: torch.Tensor) -> None:
+    return None
+
+
+@torch.library.register_fake(add_op_namespace_prefix("silu_mul_fp4_sfa_v2_fp16"))
+def _silu_v2_fake(merged: torch.Tensor, packed: torch.Tensor, sfa: torch.Tensor) -> None:
+    return None
+
+
+@torch.library.register_fake(add_op_namespace_prefix("silu_mul_mul_fp4_sfa_v2_fp16"))
+def _silu_mul_v2_fake(merged: torch.Tensor, inv_s: torch.Tensor, packed: torch.Tensor, sfa: torch.Tensor) -> None:
+    return None
+
+
+@torch.library.register_fake(add_op_namespace_prefix("silu_mul_two_fp4_to_fp4"))
+def _silu_two_fake(
+    gate_packed: torch.Tensor,
+    gate_sfa: torch.Tensor,
+    up_packed: torch.Tensor,
+    up_sfa: torch.Tensor,
+    out_packed: torch.Tensor,
+    out_sfa: torch.Tensor,
+) -> None:
+    return None
+
+
+@torch.library.register_fake(add_op_namespace_prefix("silu_mul_two_mul_fp4_to_fp4"))
+def _silu_two_mul_fake(
+    gate_packed: torch.Tensor,
+    gate_sfa: torch.Tensor,
+    up_packed: torch.Tensor,
+    up_sfa: torch.Tensor,
+    inv_s: torch.Tensor,
+    out_packed: torch.Tensor,
+    out_sfa: torch.Tensor,
+) -> None:
+    return None
+
+
+@torch.library.register_fake(add_op_namespace_prefix("dequantize_fp4_sfa_fp16"))
+def _dequant_fake(packed: torch.Tensor, sfa: torch.Tensor, out: torch.Tensor) -> None:
+    return None
+
+
+@torch.library.register_fake(add_op_namespace_prefix("rms_silu_nvfp4_ndhwc_bf16"))
+def _rms_silu_nvfp4_ndhwc_bf16_fake(
+    x: torch.Tensor,
+    gamma: torch.Tensor,
+    awq_inv_scale: torch.Tensor | None,
+    eps: float,
+    packed: torch.Tensor,
+    scale_factors: torch.Tensor,
+) -> None:
+    del eps
+    if x.dim() != 5:
+        raise RuntimeError("x must have shape (B,C,T,H,W)")
+    b, c, t, h, w = x.shape
+    if (
+        c % 128 != 0
+        or c > 1024
+        or gamma.shape != (c,)
+        or packed.shape != (b, t, h, w, c // 2)
+        or scale_factors.shape != (b, t, h, w, c // 16)
+        or (awq_inv_scale is not None and awq_inv_scale.shape != (c,))
+    ):
+        raise RuntimeError(
+            "rms_silu_nvfp4_ndhwc_bf16 shape contract failed"
+        )
+    return None
+
+
+@torch.library.register_fake(
+    add_op_namespace_prefix("quantize_bf16_to_nvfp4_linear")
+)
+def _quantize_bf16_to_nvfp4_linear_fake(
+    input: torch.Tensor,
+    packed: torch.Tensor,
+    scale_factors: torch.Tensor,
+) -> None:
+    if input.dim() != 2:
+        raise RuntimeError("input must have shape (rows, cols)")
+    rows, cols = input.shape
+    if (
+        rows <= 0
+        or cols <= 0
+        or cols % 16 != 0
+        or packed.shape != (rows, cols // 2)
+        or scale_factors.shape != (rows, cols // 16)
+    ):
+        raise RuntimeError(
+            "quantize_bf16_to_nvfp4_linear shape contract failed"
+        )
+    return None
+
+
+def rms_norm_fp4_sfa_fp16(x: torch.Tensor, packed: torch.Tensor | None = None, sfa: torch.Tensor | None = None):
+    if packed is None or sfa is None:
+        packed, sfa = _alloc_fp4(x.shape[0], x.shape[1], x.device)
+    ops.rms_norm_fp4_sfa_fp16(x, packed, sfa)
+    return packed, sfa
+
+
+def residual_add_rms_norm_fp4_sfa_fp16(
+    residual: torch.Tensor,
+    x: torch.Tensor,
+    packed: torch.Tensor | None = None,
+    sfa: torch.Tensor | None = None,
+):
+    if packed is None or sfa is None:
+        packed, sfa = _alloc_fp4(x.shape[0], x.shape[1], x.device)
+    ops.residual_add_rms_norm_fp4_sfa_fp16(residual, x, packed, sfa)
+    return packed, sfa
+
+
+def residual_add_rms_norm_fp4_sfa_v2_fp16(
+    residual: torch.Tensor,
+    x: torch.Tensor,
+    packed: torch.Tensor | None = None,
+    sfa: torch.Tensor | None = None,
+):
+    if packed is None or sfa is None:
+        packed, sfa = _alloc_fp4(x.shape[0], x.shape[1], x.device)
+    ops.residual_add_rms_norm_fp4_sfa_v2_fp16(residual, x, packed, sfa)
+    return packed, sfa
+
+
+def residual_add_rms_norm_mul_fp4_sfa_fp16(
+    residual: torch.Tensor,
+    x: torch.Tensor,
+    inv_s: torch.Tensor,
+    packed: torch.Tensor | None = None,
+    sfa: torch.Tensor | None = None,
+):
+    if packed is None or sfa is None:
+        packed, sfa = _alloc_fp4(x.shape[0], x.shape[1], x.device)
+    ops.residual_add_rms_norm_mul_fp4_sfa_fp16(residual, x, inv_s, packed, sfa)
+    return packed, sfa
+
+
+def silu_mul_fp4_sfa_fp16(merged: torch.Tensor, packed: torch.Tensor | None = None, sfa: torch.Tensor | None = None):
+    hidden = merged.shape[1] // 2
+    if packed is None or sfa is None:
+        packed, sfa = _alloc_fp4(merged.shape[0], hidden, merged.device)
+    ops.silu_mul_fp4_sfa_fp16(merged, packed, sfa)
+    return packed, sfa
+
+
+def silu_mul_fp4_sfa_v2_fp16(merged: torch.Tensor, packed: torch.Tensor | None = None, sfa: torch.Tensor | None = None):
+    hidden = merged.shape[1] // 2
+    if packed is None or sfa is None:
+        packed, sfa = _alloc_fp4(merged.shape[0], hidden, merged.device)
+    ops.silu_mul_fp4_sfa_v2_fp16(merged, packed, sfa)
+    return packed, sfa
+
+
+def silu_mul_mul_fp4_sfa_v2_fp16(
+    merged: torch.Tensor,
+    inv_s: torch.Tensor,
+    packed: torch.Tensor | None = None,
+    sfa: torch.Tensor | None = None,
+):
+    hidden = merged.shape[1] // 2
+    if packed is None or sfa is None:
+        packed, sfa = _alloc_fp4(merged.shape[0], hidden, merged.device)
+    ops.silu_mul_mul_fp4_sfa_v2_fp16(merged, inv_s, packed, sfa)
+    return packed, sfa
+
+
+def silu_mul_two_fp4_to_fp4(
+    gate_packed: torch.Tensor,
+    gate_sfa: torch.Tensor,
+    up_packed: torch.Tensor,
+    up_sfa: torch.Tensor,
+    out_packed: torch.Tensor | None = None,
+    out_sfa: torch.Tensor | None = None,
+):
+    hidden = gate_packed.shape[1] * 2
+    if out_packed is None or out_sfa is None:
+        out_packed, out_sfa = _alloc_fp4(gate_packed.shape[0], hidden, gate_packed.device)
+    ops.silu_mul_two_fp4_to_fp4(gate_packed, gate_sfa, up_packed, up_sfa, out_packed, out_sfa)
+    return out_packed, out_sfa
+
+
+def silu_mul_two_mul_fp4_to_fp4(
+    gate_packed: torch.Tensor,
+    gate_sfa: torch.Tensor,
+    up_packed: torch.Tensor,
+    up_sfa: torch.Tensor,
+    inv_s: torch.Tensor,
+    out_packed: torch.Tensor | None = None,
+    out_sfa: torch.Tensor | None = None,
+):
+    hidden = gate_packed.shape[1] * 2
+    if out_packed is None or out_sfa is None:
+        out_packed, out_sfa = _alloc_fp4(gate_packed.shape[0], hidden, gate_packed.device)
+    ops.silu_mul_two_mul_fp4_to_fp4(gate_packed, gate_sfa, up_packed, up_sfa, inv_s, out_packed, out_sfa)
+    return out_packed, out_sfa
+
+
+def dequantize_fp4_sfa_fp16(
+    packed: torch.Tensor,
+    sfa: torch.Tensor,
+    out: torch.Tensor | None = None,
+) -> torch.Tensor:
+    if out is None:
+        out = torch.empty((packed.shape[0], packed.shape[1] * 2), device=packed.device, dtype=torch.float16)
+    ops.dequantize_fp4_sfa_fp16(packed, sfa, out)
+    return out
+
+
+def rms_silu_nvfp4_ndhwc_bf16(
+    x: torch.Tensor,
+    gamma: torch.Tensor,
+    awq_inv_scale: torch.Tensor | None = None,
+    eps: float = 1e-6,
+    *,
+    packed: torch.Tensor | None = None,
+    scale_factors: torch.Tensor | None = None,
+) -> tuple[torch.Tensor, torch.Tensor]:
+    """Fuse BF16 NCDHW RMSNorm, SiLU and linear-SF NVFP4 quantization."""
+
+    b, c, t, h, w = x.shape
+    if packed is None:
+        packed = torch.empty(
+            (b, t, h, w, c // 2), device=x.device, dtype=torch.uint8
+        )
+    if scale_factors is None:
+        scale_factors = torch.empty(
+            (b, t, h, w, c // 16), device=x.device, dtype=torch.uint8
+        )
+    ops.rms_silu_nvfp4_ndhwc_bf16(
+        x,
+        gamma,
+        awq_inv_scale,
+        float(eps),
+        packed,
+        scale_factors,
+    )
+    return packed, scale_factors
+
+
+def quantize_bf16_to_nvfp4_linear(
+    input: torch.Tensor,
+    *,
+    packed: torch.Tensor | None = None,
+    scale_factors: torch.Tensor | None = None,
+) -> tuple[torch.Tensor, torch.Tensor]:
+    """Quantize BF16 rows to packed E2M1 with linear UE4M3 block-16 scales."""
+
+    rows, cols = input.shape
+    if packed is None:
+        packed = torch.empty(
+            (rows, cols // 2), device=input.device, dtype=torch.uint8
+        )
+    if scale_factors is None:
+        scale_factors = torch.empty(
+            (rows, cols // 16), device=input.device, dtype=torch.uint8
+        )
+    ops.quantize_bf16_to_nvfp4_linear(input, packed, scale_factors)
+    return packed, scale_factors
+
+
+def _check_ncdhw(x: torch.Tensor, gamma: torch.Tensor) -> None:
+    if x.dim() != 5:
+        raise RuntimeError("x must have shape (B,C,T,H,W)")
+    if x.shape[1] <= 0 or x.shape[1] % 2 != 0 or x.shape[1] > 1024:
+        raise RuntimeError("C must be even and at most 1024")
+    if gamma.shape != (x.shape[1],):
+        raise RuntimeError("gamma must have shape (C,)")
+
+
+@torch.library.register_fake(add_op_namespace_prefix("bf16_rms_silu_ncdhw"))
+def _bf16_rms_silu_ncdhw_fake(
+    x: torch.Tensor,
+    gamma: torch.Tensor,
+    prev_cache: torch.Tensor | None,
+    eps: float,
+    out: torch.Tensor,
+    next_cache: torch.Tensor | None,
+) -> None:
+    del eps
+    _check_ncdhw(x, gamma)
+    if out.shape != x.shape:
+        raise RuntimeError("out must have the same shape as x")
+    expected_cache = (x.shape[0], x.shape[1], 2, x.shape[3], x.shape[4])
+    if prev_cache is not None and prev_cache.shape != expected_cache:
+        raise RuntimeError("prev_cache must have shape (B,C,2,H,W)")
+    if next_cache is not None and next_cache.shape != expected_cache:
+        raise RuntimeError("next_cache must have shape (B,C,2,H,W)")
+    return None
+
+
+@torch.library.register_fake(add_op_namespace_prefix("bf16_rms_norm_ncdhw"))
+def _bf16_rms_norm_ncdhw_fake(
+    x: torch.Tensor,
+    gamma: torch.Tensor,
+    bias: torch.Tensor | None,
+    eps: float,
+    out: torch.Tensor,
+) -> None:
+    del eps
+    _check_ncdhw(x, gamma)
+    if bias is not None and bias.shape != (x.shape[1],):
+        raise RuntimeError("bias must have shape (C,)")
+    if out.shape != x.shape:
+        raise RuntimeError("out must have the same shape as x")
+    return None
+
+
+def bf16_rms_silu_ncdhw(
+    x: torch.Tensor,
+    gamma: torch.Tensor,
+    prev_cache: torch.Tensor | None = None,
+    eps: float = 1e-6,
+    *,
+    out: torch.Tensor | None = None,
+    next_cache: torch.Tensor | None = None,
+) -> tuple[torch.Tensor, torch.Tensor | None]:
+    """Apply channel RMSNorm and SiLU to BF16 NCDHW with optional cache update."""
+
+    if out is None:
+        out = torch.empty_like(x)
+    ops.bf16_rms_silu_ncdhw(
+        x, gamma, prev_cache, float(eps), out, next_cache
+    )
+    return out, next_cache
+
+
+def bf16_rms_norm_ncdhw(
+    x: torch.Tensor,
+    gamma: torch.Tensor,
+    bias: torch.Tensor | None = None,
+    eps: float = 1e-6,
+    *,
+    out: torch.Tensor | None = None,
+) -> torch.Tensor:
+    """Apply channel RMSNorm and optional bias to BF16 NCDHW."""
+
+    if out is None:
+        out = torch.empty_like(x)
+    ops.bf16_rms_norm_ncdhw(x, gamma, bias, float(eps), out)
+    return out
+
+
+__all__ = [
+    "bf16_rms_norm_ncdhw",
+    "bf16_rms_silu_ncdhw",
+    "dequantize_fp4_sfa_fp16",
+    "quantize_bf16_to_nvfp4_linear",
+    "residual_add_rms_norm_fp4_sfa_fp16",
+    "residual_add_rms_norm_fp4_sfa_v2_fp16",
+    "residual_add_rms_norm_mul_fp4_sfa_fp16",
+    "rms_norm_fp4_sfa_fp16",
+    "rms_silu_nvfp4_ndhwc_bf16",
+    "sfa_size_bytes",
+    "silu_mul_fp4_sfa_fp16",
+    "silu_mul_fp4_sfa_v2_fp16",
+    "silu_mul_mul_fp4_sfa_v2_fp16",
+    "silu_mul_two_fp4_to_fp4",
+    "silu_mul_two_mul_fp4_to_fp4",
+]

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

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

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

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

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

@@ -0,0 +1,32 @@
+{
+  "name": "fp4-fused-ops",
+  "id": "fp4_fused_ops_source_test",
+  "version": 1,
+  "license": "Apache-2.0",
+  "python-depends": [],
+  "backend": {
+    "type": "cuda",
+    "archs": [
+      "11.0a"
+    ]
+  },
+  "digest": {
+    "algorithm": "sha256",
+    "files": {
+      "__init__.py": "511/VkU+HAwubcrzG/Ck9wB3ukPbjKrWxG+9EgtcP00=",
+      "fp4_fused_ops_source_test.abi3.so": "Onfr4DZkn9dXFmSq6uH+Zhkhj9lNXHB/GHmXI5mEHEE=",
+      "_ops.py": "uZrqOE8sQDwTb0/km0N/KGsHFOC8TmQOIPHR2bVdj24=",
+      "fp4_fused_ops/__init__.py": "v6p5XMfQzddhi1fLSAw4HX9CyS0rQsidvu9VsT01xi4="
+    }
+  },
+  "provenance": {
+    "kernel": {
+      "sha": "456d297",
+      "dirty": false
+    },
+    "validation": {
+      "torch": "2.13.0+cu130",
+      "cuda": "13.0"
+    }
+  }
+}