Promote latest kernel artifacts to main

README.md

@@ -1,9 +1,128 @@
-# flashrt/fp4-fused-ops
+---
+library_name: kernels
+license: apache-2.0
+tags:
+  - cuda
+  - native-cuda
+  - blackwell
+  - nvfp4
+  - transformers
+  - diffusers
+---
 
-This repository is a compatibility mirror for older `kernels` clients
-that resolve repositories through the default Hugging Face model repo API.
+# fp4-fused-ops
 
-Canonical Kernel Hub repo: https://huggingface.co/kernels/flashrt/fp4-fused-ops
+FlashRT fused FP16-to-NVFP4 producer kernels for keeping low-bit transformer
+and diffuser paths continuous.
 
-Do not edit this mirror by hand. It is generated from the Kernel Hub
-`vN` branches and contains the same `build/**` artifacts.
+The source kernels use the CUTLASS SM100-family block-scale layout shared by
+SM110 and SM120. Published artifacts are architecture-specific. Native Thor
+artifacts cover Torch 2.11 and 2.13 with CUDA 13.0 and are validated on Jetson
+AGX Thor before upload.
+
+These kernels turn FP16 residual/norm/gated activations directly into NVFP4
+packed tensors plus CUTLASS-compatible SFA scale-factor buffers. They are meant
+to feed adjacent FP4 GEMM kernels without returning to PyTorch elementwise
+chains.
+
+## Available Functions
+
+- `sfa_size_bytes(rows, dim, is_sfb=False)`
+- `rms_norm_fp4_sfa_fp16(x, packed=None, sfa=None)`
+- `residual_add_rms_norm_fp4_sfa_fp16(residual, x, packed=None, sfa=None)`
+- `residual_add_rms_norm_fp4_sfa_v2_fp16(residual, x, packed=None, sfa=None)`
+- `residual_add_rms_norm_mul_fp4_sfa_fp16(residual, x, inv_s, packed=None, sfa=None)`
+- `silu_mul_fp4_sfa_fp16(merged, packed=None, sfa=None)`
+- `silu_mul_fp4_sfa_v2_fp16(merged, packed=None, sfa=None)`
+- `silu_mul_mul_fp4_sfa_v2_fp16(merged, inv_s, packed=None, sfa=None)`
+- `silu_mul_two_fp4_to_fp4(gate_packed, gate_sfa, up_packed, up_sfa, out_packed=None, out_sfa=None)`
+- `silu_mul_two_mul_fp4_to_fp4(gate_packed, gate_sfa, up_packed, up_sfa, inv_s, out_packed=None, out_sfa=None)`
+- `adaptive_rms_norm_nvfp4_fp16(x, style, packed=None, sfa=None, gate=None)`
+- `gated_residual_adaptive_rms_norm_nvfp4_fp16(x, previous_gate, residual, style, ...)`
+- `adaptive_rms_norm_nvfp4_bf16(x, style, packed=None, sfa=None, gate=None)`
+- `gated_residual_adaptive_rms_norm_nvfp4_bf16(x, previous_gate, residual, style, ...)`
+- `ada_rms_norm_quant_nvfp4_swizzled_fp16(...)` / `ada_rms_norm_quant_nvfp4_swizzled_bf16(...)`
+- `gate_res_ada_rms_norm_quant_nvfp4_swizzled_fp16(...)` / `gate_res_ada_rms_norm_quant_nvfp4_swizzled_bf16(...)`
+- `adaptive_rms_norm_fp8_static_fp16(x, style, scale, out=None, gate=None)`
+- `gate_res_ada_rms_norm_quant_fp8_static_fp16(x, previous_gate, residual, style, scale, ...)`
+- `adaptive_rms_norm_e0m3_fp16(x, style, use_rht=False, ...)`
+- `gated_residual_adaptive_rms_norm_e0m3_fp16(x, previous_gate, residual, style, use_rht=False, ...)`
+- `gelu_mul_e0m3_fp16(merged, use_rht=False, ...)`
+- `residual_add_rms_norm_quant_nvfp4_swizzled_bf16(residual, x, weight, eps=1e-6, ...)`
+- `relu2_quant_nvfp4_swizzled_fp16(x, ...)`
+- `layer_norm_fp8_fp16(x, gamma, beta, eps=1e-5, out=None)`
+- `layer_norm_nvfp4_fp16(x, gamma, beta, inv_s=None, eps=1e-5, ...)`
+- `gelu_mul_nvfp4_fp16(merged, packed=None, sfa=None)`
+- `dequantize_fp4_sfa_fp16(packed, sfa, out=None)`
+- `quantize_bf16_to_nvfp4_linear(input, packed=None, scale_factors=None)`
+- `rms_silu_nvfp4_ndhwc_bf16(x, gamma, awq_inv_scale=None, eps=1e-6, packed=None, scale_factors=None)`
+- `bf16_rms_norm_ncdhw(x, gamma, bias=None, eps=1e-6, out=None)`
+- `bf16_rms_silu_ncdhw(x, gamma, prev_cache=None, eps=1e-6, out=None, next_cache=None)`
+
+Tensor contract:
+
+- FP16 producer inputs are contiguous CUDA tensors.
+- Packed FP4 tensors are `torch.uint8` with shape `(rows, dim / 2)`.
+- SFA tensors are one-dimensional `torch.uint8` buffers sized by
+  `sfa_size_bytes(rows, dim, False, device=...)`.
+- `residual_add_*` updates `residual` in place.
+- `silu_mul_*` expects `merged` shape `(rows, 2 * hidden)`.
+- v1 shared-memory RMS producers support `dim <= 2048`; larger RMS producer
+  shapes should use `residual_add_rms_norm_fp4_sfa_v2_fp16`.
+- All dimensions must be divisible by 16. Unsupported shapes raise instead of
+  silently taking a slow or partial path.
+- Adaptive RMSNorm producers require `dim=1024`; `style` is `(rows, 3 * dim)`
+  containing scale, shift, and gate. FP16 and BF16 input/gate contracts are
+  separate additive entry points; neither silently casts.
+- Automatically allocated SFA buffers are zero-initialized because CUTLASS
+  tile padding is intentionally not written by every producer.
+- Linear NVFP4 uses E2M1 values and linear UE4M3 scale bytes per 16 channels.
+- E0M3 APIs use the uniform signed-int4 codebook with optional per-16
+  orthonormal Hadamard rotation; their packed/SFA layout is GEMM-compatible
+  but their values must not be decoded as E2M1.
+- NCDHW RMS kernels accept BF16 `(B,C,T,H,W)`, even `C <= 1024`; the fused
+  NVFP4 producer requires `C % 128 == 0`.
+
+## Minimal Usage
+
+```python
+from kernels import get_kernel
+import torch
+
+ops = get_kernel("flashrt/fp4-fused-ops", version=1, trust_remote_code=True)
+
+merged = torch.randn((16, 4096), device="cuda", dtype=torch.float16)
+packed, sfa = ops.silu_mul_fp4_sfa_v2_fp16(merged)
+
+# Optional debug/validation path. The hot path should pass packed/SFA directly
+# to the adjacent FP4 GEMM instead of dequantizing.
+dequant = ops.dequantize_fp4_sfa_fp16(packed, sfa)
+
+video = torch.randn((1, 128, 5, 9, 11), device="cuda", dtype=torch.bfloat16)
+gamma = torch.ones((128,), device="cuda", dtype=torch.bfloat16)
+packed_video, linear_scales = ops.rms_silu_nvfp4_ndhwc_bf16(video, gamma)
+```
+
+## Validation
+
+```bash
+python fp4-fused-ops/tests/test_fp4_fused_ops.py --backend source --mode full
+python fp4-fused-ops/benchmarks/benchmark.py --mode headline
+
+# Thor production-model rows
+python fp4-fused-ops/tests/test_fp4_fused_ops.py --backend installed \
+  --mode thor-models \
+  --artifact fp4-fused-ops/build/torch211-cxx11-cu130-aarch64-linux
+python fp4-fused-ops/benchmarks/benchmark.py --backend installed \
+  --mode thor-models \
+  --artifact fp4-fused-ops/build/torch211-cxx11-cu130-aarch64-linux
+```
+
+Validation checks:
+
+- in-place residual updates match the FP16 math contract exactly;
+- FP4/SFA outputs dequantize within the documented NVFP4 quantization error
+  envelope;
+- v2 fast paths are checked against v1 where the v1 kernel is valid;
+- unsupported shapes are rejected explicitly.
+- new NCDHW/cache outputs and linear NVFP4 bytes match raw launchers bitwise.

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

@@ -87,6 +87,36 @@ def _silu_two_mul_fake(
     return None
 
 
+@torch.library.register_fake(add_op_namespace_prefix("geglu_two_mul_nvfp4_native"))
+def _geglu_two_native_fake(gate_packed, gate_sfa, up_packed, up_sfa, inv_s, out_packed, out_sfa) -> None:
+    return None
+
+
+@torch.library.register_fake(add_op_namespace_prefix("gelu_mul_nvfp4_bf16"))
+def _gelu_mul_bf16_fake(merged, inv_s, packed, sfa) -> None:
+    return None
+
+
+@torch.library.register_fake(add_op_namespace_prefix("rms_norm_mul_nvfp4_bf16"))
+def _rms_mul_bf16_fake(x, inv_s, eps: float, packed, sfa) -> None:
+    return None
+
+
+@torch.library.register_fake(add_op_namespace_prefix("residual_add_rms_norm_nvfp4_bf16"))
+def _res_rms_bf16_fake(residual, x, inv_s, eps: float, packed, sfa) -> None:
+    return None
+
+
+@torch.library.register_fake(add_op_namespace_prefix("layer_norm_fp8_bf16"))
+def _layer_norm_fp8_bf16_fake(x, gamma, beta, eps: float, out) -> None:
+    return None
+
+
+@torch.library.register_fake(add_op_namespace_prefix("layer_norm_nvfp4_bf16"))
+def _layer_norm_nvfp4_bf16_fake(x, gamma, beta, inv_s, eps: float, packed, sfa) -> 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
@@ -319,6 +349,110 @@ def silu_mul_two_mul_fp4_to_fp4(
     return out_packed, out_sfa
 
 
+def geglu_two_mul_nvfp4_native(
+    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,
+):
+    """Native-E2M1 LUT GeGLU combiner for split gate/up FP4 outputs."""
+    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.geglu_two_mul_nvfp4_native(
+        gate_packed, gate_sfa, up_packed, up_sfa, inv_s,
+        out_packed, out_sfa
+    )
+    return out_packed, out_sfa
+
+
+def gelu_mul_nvfp4_bf16(
+    merged: torch.Tensor,
+    inv_s: torch.Tensor | None = None,
+    *,
+    packed: torch.Tensor | None = None,
+    sfa: torch.Tensor | None = None,
+):
+    """BF16 GeGLU producer with optional per-channel scaling and NVFP4 output."""
+    hidden = merged.shape[1] // 2
+    if packed is None or sfa is None:
+        packed, sfa = _alloc_fp4(merged.shape[0], hidden, merged.device)
+    ops.gelu_mul_nvfp4_bf16(merged, inv_s, packed, sfa)
+    return packed, sfa
+
+
+def rms_norm_mul_nvfp4_bf16(
+    x: torch.Tensor,
+    inv_s: torch.Tensor,
+    eps: float = 1e-6,
+    *,
+    packed: torch.Tensor | None = None,
+    sfa: torch.Tensor | None = None,
+):
+    """BF16 RMSNorm, channel scaling, and native NVFP4 production."""
+    if packed is None or sfa is None:
+        packed, sfa = _alloc_fp4(x.shape[0], x.shape[1], x.device)
+    ops.rms_norm_mul_nvfp4_bf16(x, inv_s, float(eps), packed, sfa)
+    return packed, sfa
+
+
+def residual_add_rms_norm_nvfp4_bf16(
+    residual: torch.Tensor,
+    x: torch.Tensor,
+    inv_s: torch.Tensor | None = None,
+    eps: float = 1e-6,
+    *,
+    packed: torch.Tensor | None = None,
+    sfa: torch.Tensor | None = None,
+):
+    """Update a BF16 residual, RMS-normalize, and produce native NVFP4."""
+    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_nvfp4_bf16(
+        residual, x, inv_s, float(eps), packed, sfa
+    )
+    return residual, packed, sfa
+
+
+def layer_norm_fp8_bf16(
+    x: torch.Tensor,
+    gamma: torch.Tensor,
+    beta: torch.Tensor,
+    eps: float = 1e-6,
+    *,
+    out: torch.Tensor | None = None,
+) -> torch.Tensor:
+    """Vectorized BF16 LayerNorm producer with FP8 E4M3 output."""
+    if out is None:
+        out = torch.empty_like(x, dtype=torch.float8_e4m3fn)
+    ops.layer_norm_fp8_bf16(x, gamma, beta, float(eps), out)
+    return out
+
+
+def layer_norm_nvfp4_bf16(
+    x: torch.Tensor,
+    gamma: torch.Tensor,
+    beta: torch.Tensor,
+    inv_s: torch.Tensor | None = None,
+    eps: float = 1e-6,
+    *,
+    packed: torch.Tensor | None = None,
+    sfa: torch.Tensor | None = None,
+):
+    """Vectorized BF16 LayerNorm producer with native NVFP4 output."""
+    if packed is None or sfa is None:
+        packed, sfa = _alloc_fp4(x.shape[0], x.shape[1], x.device)
+    ops.layer_norm_nvfp4_bf16(
+        x, gamma, beta, inv_s, float(eps), packed, sfa
+    )
+    return packed, sfa
+
+
 def adaptive_rms_norm_nvfp4_fp16(
     x: torch.Tensor,
     style: torch.Tensor,
@@ -802,13 +936,19 @@ __all__ = [
     "gate_res_ada_rms_norm_quant_nvfp4_swizzled_bf16",
     "gate_res_ada_rms_norm_quant_fp8_static_fp16",
     "gelu_mul_nvfp4_fp16",
+    "gelu_mul_nvfp4_bf16",
+    "geglu_two_mul_nvfp4_native",
     "gelu_mul_e0m3_fp16",
     "layer_norm_fp8_fp16",
+    "layer_norm_fp8_bf16",
     "layer_norm_nvfp4_fp16",
+    "layer_norm_nvfp4_bf16",
     "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",
+    "residual_add_rms_norm_nvfp4_bf16",
+    "rms_norm_mul_nvfp4_bf16",
     "residual_add_rms_norm_quant_nvfp4_swizzled_bf16",
     "rms_norm_fp4_sfa_fp16",
     "rms_silu_nvfp4_ndhwc_bf16",

build/torch211-cxx11-cu128-x86_64-linux/{_fp4_fused_ops_cuda_8cf8928_dirty.abi3.so → _fp4_fused_ops_cuda_92c4b99.abi3.so}

@@ -1,3 +1,3 @@
 version https://git-lfs.github.com/spec/v1
-oid sha256:d9217ecaa0bdcd4b9f998e934d87cfba926850c127b32c3d885dadc2e3f6172a
-size 1271200
+oid sha256:1f0dfb36611d5ce875a5e0296023bb47687c45ce00ebdff93bbb0a039642dceb
+size 1470080

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

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

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

@@ -2,25 +2,13 @@ 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}")
+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)
-    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
+    sys.modules[path_hash] = module
+    spec.loader.exec_module(module)
     return module
 
-
-globals().update(vars(_import_from_path(Path(__file__).parent.parent / "__init__.py")))
+globals().update(vars(_import_from_path(Path(__file__).parent.parent / '__init__.py')))

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

@@ -1,6 +1,6 @@
 {
   "name": "fp4-fused-ops",
-  "id": "_fp4_fused_ops_cuda_8cf8928_dirty",
+  "id": "_fp4_fused_ops_cuda_92c4b99",
   "version": 1,
   "license": "Apache-2.0",
   "python-depends": [],
@@ -13,10 +13,20 @@
   "digest": {
     "algorithm": "sha256",
     "files": {
-      "__init__.py": "KxY15agQ6cnqkslzItrun85+7hL4OqX3FSFpolNUq1Q=",
-      "_fp4_fused_ops_cuda_8cf8928_dirty.abi3.so": "2SF+yqC9zUufmY6TTYfPupJoUMEnsyw9iF2twuP2Fyo=",
-      "_ops.py": "5bB7feK8B7pytTbr2Ty52irfjV7U/tJBINU996zlQtI=",
-      "fp4_fused_ops/__init__.py": "DFYPlrhXwYjEqCl/8n0SmWGZV8NFml5DPhMjKfv98GY="
+      "__init__.py": "3Bkz49FpbwcatmJVYSwTnKEeRWr1RkkYZowj/l9fJOY=",
+      "_fp4_fused_ops_cuda_92c4b99.abi3.so": "Hw37NmEdXOh1peApYCO7R2h8Rc4A69/5O7sKA5ZC3Os=",
+      "_ops.py": "I13aTXDYLStDS3cO/5A8MkCfRnBg5EhUFfcIq8Sz8yk=",
+      "fp4_fused_ops/__init__.py": "v6p5XMfQzddhi1fLSAw4HX9CyS0rQsidvu9VsT01xi4="
+    }
+  },
+  "provenance": {
+    "kernel": {
+      "sha": "92c4b991b6f1fbd89744aa6f46f62b1bfac7ca7e",
+      "dirty": false
+    },
+    "validation": {
+      "torch": "2.11.0+cu128",
+      "cuda": "12.8"
     }
   }
-}
+}

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

@@ -87,6 +87,36 @@ def _silu_two_mul_fake(
     return None
 
 
+@torch.library.register_fake(add_op_namespace_prefix("geglu_two_mul_nvfp4_native"))
+def _geglu_two_native_fake(gate_packed, gate_sfa, up_packed, up_sfa, inv_s, out_packed, out_sfa) -> None:
+    return None
+
+
+@torch.library.register_fake(add_op_namespace_prefix("gelu_mul_nvfp4_bf16"))
+def _gelu_mul_bf16_fake(merged, inv_s, packed, sfa) -> None:
+    return None
+
+
+@torch.library.register_fake(add_op_namespace_prefix("rms_norm_mul_nvfp4_bf16"))
+def _rms_mul_bf16_fake(x, inv_s, eps: float, packed, sfa) -> None:
+    return None
+
+
+@torch.library.register_fake(add_op_namespace_prefix("residual_add_rms_norm_nvfp4_bf16"))
+def _res_rms_bf16_fake(residual, x, inv_s, eps: float, packed, sfa) -> None:
+    return None
+
+
+@torch.library.register_fake(add_op_namespace_prefix("layer_norm_fp8_bf16"))
+def _layer_norm_fp8_bf16_fake(x, gamma, beta, eps: float, out) -> None:
+    return None
+
+
+@torch.library.register_fake(add_op_namespace_prefix("layer_norm_nvfp4_bf16"))
+def _layer_norm_nvfp4_bf16_fake(x, gamma, beta, inv_s, eps: float, packed, sfa) -> 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
@@ -319,6 +349,110 @@ def silu_mul_two_mul_fp4_to_fp4(
     return out_packed, out_sfa
 
 
+def geglu_two_mul_nvfp4_native(
+    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,
+):
+    """Native-E2M1 LUT GeGLU combiner for split gate/up FP4 outputs."""
+    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.geglu_two_mul_nvfp4_native(
+        gate_packed, gate_sfa, up_packed, up_sfa, inv_s,
+        out_packed, out_sfa
+    )
+    return out_packed, out_sfa
+
+
+def gelu_mul_nvfp4_bf16(
+    merged: torch.Tensor,
+    inv_s: torch.Tensor | None = None,
+    *,
+    packed: torch.Tensor | None = None,
+    sfa: torch.Tensor | None = None,
+):
+    """BF16 GeGLU producer with optional per-channel scaling and NVFP4 output."""
+    hidden = merged.shape[1] // 2
+    if packed is None or sfa is None:
+        packed, sfa = _alloc_fp4(merged.shape[0], hidden, merged.device)
+    ops.gelu_mul_nvfp4_bf16(merged, inv_s, packed, sfa)
+    return packed, sfa
+
+
+def rms_norm_mul_nvfp4_bf16(
+    x: torch.Tensor,
+    inv_s: torch.Tensor,
+    eps: float = 1e-6,
+    *,
+    packed: torch.Tensor | None = None,
+    sfa: torch.Tensor | None = None,
+):
+    """BF16 RMSNorm, channel scaling, and native NVFP4 production."""
+    if packed is None or sfa is None:
+        packed, sfa = _alloc_fp4(x.shape[0], x.shape[1], x.device)
+    ops.rms_norm_mul_nvfp4_bf16(x, inv_s, float(eps), packed, sfa)
+    return packed, sfa
+
+
+def residual_add_rms_norm_nvfp4_bf16(
+    residual: torch.Tensor,
+    x: torch.Tensor,
+    inv_s: torch.Tensor | None = None,
+    eps: float = 1e-6,
+    *,
+    packed: torch.Tensor | None = None,
+    sfa: torch.Tensor | None = None,
+):
+    """Update a BF16 residual, RMS-normalize, and produce native NVFP4."""
+    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_nvfp4_bf16(
+        residual, x, inv_s, float(eps), packed, sfa
+    )
+    return residual, packed, sfa
+
+
+def layer_norm_fp8_bf16(
+    x: torch.Tensor,
+    gamma: torch.Tensor,
+    beta: torch.Tensor,
+    eps: float = 1e-6,
+    *,
+    out: torch.Tensor | None = None,
+) -> torch.Tensor:
+    """Vectorized BF16 LayerNorm producer with FP8 E4M3 output."""
+    if out is None:
+        out = torch.empty_like(x, dtype=torch.float8_e4m3fn)
+    ops.layer_norm_fp8_bf16(x, gamma, beta, float(eps), out)
+    return out
+
+
+def layer_norm_nvfp4_bf16(
+    x: torch.Tensor,
+    gamma: torch.Tensor,
+    beta: torch.Tensor,
+    inv_s: torch.Tensor | None = None,
+    eps: float = 1e-6,
+    *,
+    packed: torch.Tensor | None = None,
+    sfa: torch.Tensor | None = None,
+):
+    """Vectorized BF16 LayerNorm producer with native NVFP4 output."""
+    if packed is None or sfa is None:
+        packed, sfa = _alloc_fp4(x.shape[0], x.shape[1], x.device)
+    ops.layer_norm_nvfp4_bf16(
+        x, gamma, beta, inv_s, float(eps), packed, sfa
+    )
+    return packed, sfa
+
+
 def adaptive_rms_norm_nvfp4_fp16(
     x: torch.Tensor,
     style: torch.Tensor,
@@ -802,13 +936,19 @@ __all__ = [
     "gate_res_ada_rms_norm_quant_nvfp4_swizzled_bf16",
     "gate_res_ada_rms_norm_quant_fp8_static_fp16",
     "gelu_mul_nvfp4_fp16",
+    "gelu_mul_nvfp4_bf16",
+    "geglu_two_mul_nvfp4_native",
     "gelu_mul_e0m3_fp16",
     "layer_norm_fp8_fp16",
+    "layer_norm_fp8_bf16",
     "layer_norm_nvfp4_fp16",
+    "layer_norm_nvfp4_bf16",
     "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",
+    "residual_add_rms_norm_nvfp4_bf16",
+    "rms_norm_mul_nvfp4_bf16",
     "residual_add_rms_norm_quant_nvfp4_swizzled_bf16",
     "rms_norm_fp4_sfa_fp16",
     "rms_silu_nvfp4_ndhwc_bf16",

build/torch211-cxx11-cu130-aarch64-linux/{fp4_fused_ops_source_test.abi3.so → _fp4_fused_ops_cuda_92c4b99.abi3.so}

@@ -1,3 +1,3 @@
 version https://git-lfs.github.com/spec/v1
-oid sha256:fcc7c79b9342830a7b6b367246107d03aa1c9a95b0a5dfd21515ed53dbe6657b
-size 1361344
+oid sha256:98c2fcc607d2e1ac0f784306a0330b792a3441af1cbc7aa4effc2d954b360bfa
+size 1581160

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

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

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

@@ -1,6 +1,6 @@
 {
   "name": "fp4-fused-ops",
-  "id": "fp4_fused_ops_source_test",
+  "id": "_fp4_fused_ops_cuda_92c4b99",
   "version": 1,
   "license": "Apache-2.0",
   "python-depends": [],
@@ -13,15 +13,15 @@
   "digest": {
     "algorithm": "sha256",
     "files": {
-      "__init__.py": "KxY15agQ6cnqkslzItrun85+7hL4OqX3FSFpolNUq1Q=",
-      "fp4_fused_ops_source_test.abi3.so": "/MfHm5NCgwp7azZyRhB9A6ocmpWwpd/SFRXtU9vmZXs=",
-      "_ops.py": "uZrqOE8sQDwTb0/km0N/KGsHFOC8TmQOIPHR2bVdj24=",
+      "__init__.py": "3Bkz49FpbwcatmJVYSwTnKEeRWr1RkkYZowj/l9fJOY=",
+      "_fp4_fused_ops_cuda_92c4b99.abi3.so": "mML8xgfS4awPeEMGoDMLeSo0Qa8cvHqk7/wtlUs2C/o=",
+      "_ops.py": "I13aTXDYLStDS3cO/5A8MkCfRnBg5EhUFfcIq8Sz8yk=",
       "fp4_fused_ops/__init__.py": "v6p5XMfQzddhi1fLSAw4HX9CyS0rQsidvu9VsT01xi4="
     }
   },
   "provenance": {
     "kernel": {
-      "sha": "1ee78f4",
+      "sha": "92c4b991b6f1fbd89744aa6f46f62b1bfac7ca7e",
       "dirty": false
     },
     "validation": {

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

@@ -87,6 +87,36 @@ def _silu_two_mul_fake(
     return None
 
 
+@torch.library.register_fake(add_op_namespace_prefix("geglu_two_mul_nvfp4_native"))
+def _geglu_two_native_fake(gate_packed, gate_sfa, up_packed, up_sfa, inv_s, out_packed, out_sfa) -> None:
+    return None
+
+
+@torch.library.register_fake(add_op_namespace_prefix("gelu_mul_nvfp4_bf16"))
+def _gelu_mul_bf16_fake(merged, inv_s, packed, sfa) -> None:
+    return None
+
+
+@torch.library.register_fake(add_op_namespace_prefix("rms_norm_mul_nvfp4_bf16"))
+def _rms_mul_bf16_fake(x, inv_s, eps: float, packed, sfa) -> None:
+    return None
+
+
+@torch.library.register_fake(add_op_namespace_prefix("residual_add_rms_norm_nvfp4_bf16"))
+def _res_rms_bf16_fake(residual, x, inv_s, eps: float, packed, sfa) -> None:
+    return None
+
+
+@torch.library.register_fake(add_op_namespace_prefix("layer_norm_fp8_bf16"))
+def _layer_norm_fp8_bf16_fake(x, gamma, beta, eps: float, out) -> None:
+    return None
+
+
+@torch.library.register_fake(add_op_namespace_prefix("layer_norm_nvfp4_bf16"))
+def _layer_norm_nvfp4_bf16_fake(x, gamma, beta, inv_s, eps: float, packed, sfa) -> 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
@@ -319,6 +349,110 @@ def silu_mul_two_mul_fp4_to_fp4(
     return out_packed, out_sfa
 
 
+def geglu_two_mul_nvfp4_native(
+    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,
+):
+    """Native-E2M1 LUT GeGLU combiner for split gate/up FP4 outputs."""
+    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.geglu_two_mul_nvfp4_native(
+        gate_packed, gate_sfa, up_packed, up_sfa, inv_s,
+        out_packed, out_sfa
+    )
+    return out_packed, out_sfa
+
+
+def gelu_mul_nvfp4_bf16(
+    merged: torch.Tensor,
+    inv_s: torch.Tensor | None = None,
+    *,
+    packed: torch.Tensor | None = None,
+    sfa: torch.Tensor | None = None,
+):
+    """BF16 GeGLU producer with optional per-channel scaling and NVFP4 output."""
+    hidden = merged.shape[1] // 2
+    if packed is None or sfa is None:
+        packed, sfa = _alloc_fp4(merged.shape[0], hidden, merged.device)
+    ops.gelu_mul_nvfp4_bf16(merged, inv_s, packed, sfa)
+    return packed, sfa
+
+
+def rms_norm_mul_nvfp4_bf16(
+    x: torch.Tensor,
+    inv_s: torch.Tensor,
+    eps: float = 1e-6,
+    *,
+    packed: torch.Tensor | None = None,
+    sfa: torch.Tensor | None = None,
+):
+    """BF16 RMSNorm, channel scaling, and native NVFP4 production."""
+    if packed is None or sfa is None:
+        packed, sfa = _alloc_fp4(x.shape[0], x.shape[1], x.device)
+    ops.rms_norm_mul_nvfp4_bf16(x, inv_s, float(eps), packed, sfa)
+    return packed, sfa
+
+
+def residual_add_rms_norm_nvfp4_bf16(
+    residual: torch.Tensor,
+    x: torch.Tensor,
+    inv_s: torch.Tensor | None = None,
+    eps: float = 1e-6,
+    *,
+    packed: torch.Tensor | None = None,
+    sfa: torch.Tensor | None = None,
+):
+    """Update a BF16 residual, RMS-normalize, and produce native NVFP4."""
+    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_nvfp4_bf16(
+        residual, x, inv_s, float(eps), packed, sfa
+    )
+    return residual, packed, sfa
+
+
+def layer_norm_fp8_bf16(
+    x: torch.Tensor,
+    gamma: torch.Tensor,
+    beta: torch.Tensor,
+    eps: float = 1e-6,
+    *,
+    out: torch.Tensor | None = None,
+) -> torch.Tensor:
+    """Vectorized BF16 LayerNorm producer with FP8 E4M3 output."""
+    if out is None:
+        out = torch.empty_like(x, dtype=torch.float8_e4m3fn)
+    ops.layer_norm_fp8_bf16(x, gamma, beta, float(eps), out)
+    return out
+
+
+def layer_norm_nvfp4_bf16(
+    x: torch.Tensor,
+    gamma: torch.Tensor,
+    beta: torch.Tensor,
+    inv_s: torch.Tensor | None = None,
+    eps: float = 1e-6,
+    *,
+    packed: torch.Tensor | None = None,
+    sfa: torch.Tensor | None = None,
+):
+    """Vectorized BF16 LayerNorm producer with native NVFP4 output."""
+    if packed is None or sfa is None:
+        packed, sfa = _alloc_fp4(x.shape[0], x.shape[1], x.device)
+    ops.layer_norm_nvfp4_bf16(
+        x, gamma, beta, inv_s, float(eps), packed, sfa
+    )
+    return packed, sfa
+
+
 def adaptive_rms_norm_nvfp4_fp16(
     x: torch.Tensor,
     style: torch.Tensor,
@@ -802,13 +936,19 @@ __all__ = [
     "gate_res_ada_rms_norm_quant_nvfp4_swizzled_bf16",
     "gate_res_ada_rms_norm_quant_fp8_static_fp16",
     "gelu_mul_nvfp4_fp16",
+    "gelu_mul_nvfp4_bf16",
+    "geglu_two_mul_nvfp4_native",
     "gelu_mul_e0m3_fp16",
     "layer_norm_fp8_fp16",
+    "layer_norm_fp8_bf16",
     "layer_norm_nvfp4_fp16",
+    "layer_norm_nvfp4_bf16",
     "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",
+    "residual_add_rms_norm_nvfp4_bf16",
+    "rms_norm_mul_nvfp4_bf16",
     "residual_add_rms_norm_quant_nvfp4_swizzled_bf16",
     "rms_norm_fp4_sfa_fp16",
     "rms_silu_nvfp4_ndhwc_bf16",

build/torch211-cxx11-cu130-x86_64-linux/{_fp4_fused_ops_cuda_8cf8928_dirty.abi3.so → _fp4_fused_ops_cuda_92c4b99.abi3.so}

@@ -1,3 +1,3 @@
 version https://git-lfs.github.com/spec/v1
-oid sha256:be2b14ef9a2fc9f4e0fcd2f7fb140420ec03d3467bd8f11f4fc78edbf9416efd
-size 2180064
+oid sha256:d0519722786f9bf6ced8f22ca95801289b860275b1ff1f96ccb1cf10e8047de0
+size 1569904

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

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

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

@@ -2,25 +2,13 @@ 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}")
+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)
-    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
+    sys.modules[path_hash] = module
+    spec.loader.exec_module(module)
     return module
 
-
-globals().update(vars(_import_from_path(Path(__file__).parent.parent / "__init__.py")))
+globals().update(vars(_import_from_path(Path(__file__).parent.parent / '__init__.py')))

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

@@ -1,23 +1,32 @@
 {
   "name": "fp4-fused-ops",
-  "id": "_fp4_fused_ops_cuda_8cf8928_dirty",
+  "id": "_fp4_fused_ops_cuda_92c4b99",
   "version": 1,
   "license": "Apache-2.0",
   "python-depends": [],
   "backend": {
     "type": "cuda",
     "archs": [
-      "11.0a",
       "12.0a"
     ]
   },
   "digest": {
     "algorithm": "sha256",
     "files": {
-      "__init__.py": "KxY15agQ6cnqkslzItrun85+7hL4OqX3FSFpolNUq1Q=",
-      "_fp4_fused_ops_cuda_8cf8928_dirty.abi3.so": "visU75ovyfTg/NL3+xQEIOwD00Z72PEfT8eO2/lBbv0=",
-      "_ops.py": "5bB7feK8B7pytTbr2Ty52irfjV7U/tJBINU996zlQtI=",
-      "fp4_fused_ops/__init__.py": "DFYPlrhXwYjEqCl/8n0SmWGZV8NFml5DPhMjKfv98GY="
+      "__init__.py": "3Bkz49FpbwcatmJVYSwTnKEeRWr1RkkYZowj/l9fJOY=",
+      "_fp4_fused_ops_cuda_92c4b99.abi3.so": "0FGXInhvm/bO2PIsqVgBKJuGAnWx/x+WzLHPEOgEfeA=",
+      "_ops.py": "I13aTXDYLStDS3cO/5A8MkCfRnBg5EhUFfcIq8Sz8yk=",
+      "fp4_fused_ops/__init__.py": "v6p5XMfQzddhi1fLSAw4HX9CyS0rQsidvu9VsT01xi4="
+    }
+  },
+  "provenance": {
+    "kernel": {
+      "sha": "92c4b991b6f1fbd89744aa6f46f62b1bfac7ca7e",
+      "dirty": false
+    },
+    "validation": {
+      "torch": "2.11.0+cu130",
+      "cuda": "13.0"
     }
   }
-}
+}

build/torch213-cxx11-cu130-aarch64-linux/{fp4_fused_ops_source_test.abi3.so → _fp4_fused_ops_cuda_92c4b99.abi3.so}

@@ -1,3 +1,3 @@
 version https://git-lfs.github.com/spec/v1
-oid sha256:a80eac067fb4184c162f5a7521b09faf7cb0b1098d5f28c59e49b12d8c918e87
-size 1583520
+oid sha256:1856d76e72049f097d9e240fdda018263dbec848b991216167cc38da048d205c
+size 1583528

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

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

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

@@ -1,27 +1,27 @@
 {
   "name": "fp4-fused-ops",
-  "id": "fp4_fused_ops_source_test",
+  "id": "_fp4_fused_ops_cuda_92c4b99",
   "version": 1,
   "license": "Apache-2.0",
   "python-depends": [],
   "backend": {
     "type": "cuda",
     "archs": [
-      "11.0"
+      "11.0a"
     ]
   },
   "digest": {
     "algorithm": "sha256",
     "files": {
       "__init__.py": "3Bkz49FpbwcatmJVYSwTnKEeRWr1RkkYZowj/l9fJOY=",
-      "fp4_fused_ops_source_test.abi3.so": "qA6sBn+0GEwWL1p1IbCfr3ywsQmNXyjFnkmxLYyRjoc=",
-      "_ops.py": "uZrqOE8sQDwTb0/km0N/KGsHFOC8TmQOIPHR2bVdj24=",
+      "_fp4_fused_ops_cuda_92c4b99.abi3.so": "GFbXbnIEnwl9niQP3aAYJj2+yEi5kSFhZ8w42gSNIFw=",
+      "_ops.py": "I13aTXDYLStDS3cO/5A8MkCfRnBg5EhUFfcIq8Sz8yk=",
       "fp4_fused_ops/__init__.py": "v6p5XMfQzddhi1fLSAw4HX9CyS0rQsidvu9VsT01xi4="
     }
   },
   "provenance": {
     "kernel": {
-      "sha": "b6588aa59dd4f8c646f90b58312a0a66886e061f",
+      "sha": "92c4b991b6f1fbd89744aa6f46f62b1bfac7ca7e",
       "dirty": false
     },
     "validation": {