Promote latest kernel artifacts to main

README.md

@@ -1,9 +0,0 @@
-# flashrt/fp4-gemm
-
-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-gemm
-
-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,9 +25,13 @@ class BenchResult:
     N: int
     K: int
     variant: int
+    native_us: float
     flashrt_us: float
-    torch_reference_us: float
-    speedup_vs_reference: float
+    torch_eager_us: float
+    torch_compile_us: float
+    speedup_vs_eager: float
+    speedup_vs_compile: float
+    wrapper_over_native: float
     max_abs: float
     mean_abs: float
     p99_abs: float
@@ -58,7 +63,7 @@ def measure(fn, warmup: int, iters: int) -> float:
     return float(start.elapsed_time(end) * 1000.0 / iters)
 
 
-def bench_case(helpers, ops, name: str, shape: tuple[int, int, int], warmup: int, iters: int) -> list[BenchResult]:
+def bench_case(helpers, ops, native, name: str, shape: tuple[int, int, int], warmup: int, iters: int) -> list[BenchResult]:
     m, n, k = shape
     a_packed, b_packed, sfa, sfb, expected = helpers.prepare_quantized(ops, m, n, k)
     a_deq = torch.empty((m, k), device="cuda", dtype=torch.float16)
@@ -70,15 +75,45 @@ def bench_case(helpers, ops, name: str, shape: tuple[int, int, int], warmup: int
     def torch_ref():
         return (a_deq.float() @ b_deq.float().T).to(torch.bfloat16)
 
-    torch_us = measure(torch_ref, warmup, iters)
+    torch_eager_us = measure(torch_ref, warmup, iters)
+    compiled_ref = torch.compile(torch_ref, mode="max-autotune-no-cudagraphs")
+    torch_compile_us = measure(compiled_ref, warmup, iters)
+    stream = torch.cuda.current_stream().cuda_stream
     results: list[BenchResult] = []
-    for variant in (0, 1, 2):
+    variants = (-1, 0, 1, 2) if torch.cuda.get_device_capability(0) == (11, 0) else (0, 1, 2)
+    for variant in variants:
         out = torch.empty((m, n), device="cuda", dtype=torch.bfloat16)
-        ops.fp4_w4a16_linear_bf16(a_packed, b_packed, sfa, sfb, out, 1.0, variant)
+        ops.nvfp4_gemm_bf16(a_packed, b_packed, sfa, sfb, out, 1.0, variant)
         torch.cuda.synchronize()
         max_abs, mean_abs, p99_abs, cosine = helpers.metrics(out, expected)
         flashrt_us = measure(
-            lambda: ops.fp4_w4a16_linear_bf16(a_packed, b_packed, sfa, sfb, out, 1.0, variant),
+            lambda: ops.nvfp4_gemm_bf16(a_packed, b_packed, sfa, sfb, out, 1.0, variant),
+            warmup,
+            iters,
+        )
+        native_variant = variant
+        if native_variant < 0:
+            native_variant = helpers.select_sm110_variant(shape)
+        native_function = (
+            native.fp4_w4a16_gemm_sm120_bf16out
+            if native_variant == 0
+            else native.fp4_w4a16_gemm_sm120_bf16out_widen
+            if native_variant == 1
+            else native.fp4_w4a16_gemm_sm120_bf16out_pingpong
+        )
+        native_us = measure(
+            lambda: native_function(
+                a_packed.data_ptr(),
+                b_packed.data_ptr(),
+                out.data_ptr(),
+                m,
+                n,
+                k,
+                sfa.data_ptr(),
+                sfb.data_ptr(),
+                1.0,
+                stream,
+            ),
             warmup,
             iters,
         )
@@ -89,9 +124,13 @@ def bench_case(helpers, ops, name: str, shape: tuple[int, int, int], warmup: int
                 N=n,
                 K=k,
                 variant=variant,
+                native_us=native_us,
                 flashrt_us=flashrt_us,
-                torch_reference_us=torch_us,
-                speedup_vs_reference=torch_us / flashrt_us,
+                torch_eager_us=torch_eager_us,
+                torch_compile_us=torch_compile_us,
+                speedup_vs_eager=torch_eager_us / flashrt_us,
+                speedup_vs_compile=torch_compile_us / flashrt_us,
+                wrapper_over_native=flashrt_us / native_us,
                 max_abs=max_abs,
                 mean_abs=mean_abs,
                 p99_abs=p99_abs,
@@ -102,31 +141,107 @@ def bench_case(helpers, ops, name: str, shape: tuple[int, int, int], warmup: int
     return results
 
 
+def bench_bf16_producer(ops, native, k: int, warmup: int, iters: int):
+    x = torch.randn((1, k), device="cuda", dtype=torch.bfloat16)
+    direct_packed, direct_sfa = ops.alloc_fp4(1, k)
+    compat_packed, compat_sfa = ops.alloc_fp4(1, k)
+    native_packed, native_sfa = ops.alloc_fp4(1, k)
+    stream = torch.cuda.current_stream().cuda_stream
+
+    def direct():
+        ops.quantize_fp4_sfa_bf16(
+            x, direct_packed, direct_sfa, False
+        )
+
+    def compat():
+        ops.quantize_fp4_sfa_fp16(
+            x.to(torch.float16), compat_packed, compat_sfa, False
+        )
+
+    def native_direct():
+        native.quantize_bf16_to_nvfp4_swizzled(
+            x.data_ptr(), native_packed.data_ptr(), native_sfa.data_ptr(),
+            1, k, stream,
+        )
+
+    direct()
+    compat()
+    torch.cuda.synchronize()
+    direct_us = measure(direct, warmup, iters)
+    compat_us = measure(compat, warmup, iters)
+    native_us = measure(native_direct, warmup, iters)
+    return {
+        "M": 1,
+        "K": k,
+        "direct_bf16_us": direct_us,
+        "cast_plus_fp16_us": compat_us,
+        "native_bf16_us": native_us,
+        "speedup_vs_cast_plus_fp16": compat_us / direct_us,
+        "wrapper_over_native": direct_us / native_us,
+        "packed_exact_vs_fp16_contract": bool(
+            torch.equal(direct_packed, compat_packed)
+        ),
+        "note": "native_bf16 uses a distinct FlashRT quantization strategy",
+    }
+
+
 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)
     args = parser.parse_args()
 
     helpers = load_helpers()
-    ops = helpers.load_source_ops()
+    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_kernels as native
+    finally:
+        sys.path.pop(0)
+    ops = (
+        helpers.load_source_ops()
+        if args.backend == "source"
+        else helpers.load_installed_ops(args.artifact)
+    )
     shapes = {
         "small_m16_n128_k128": (16, 128, 128),
         "small_m32_n256_k256": (32, 256, 256),
         "mlp_tile_m64_n512_k512": (64, 512, 512),
+        "groot_dit_projection": (51, 1536, 1536),
+        "vla_projection": (105, 2048, 2048),
+        "motus_up": (360, 14336, 3072),
+        "motus_down": (360, 3072, 14336),
     }
     if args.mode == "smoke":
         shapes = {"small_m16_n128_k128": shapes["small_m16_n128_k128"]}
+    elif args.mode == "thor-models":
+        shapes = dict(helpers.SM110_SHAPES)
     results: list[BenchResult] = []
     for name, shape in shapes.items():
-        results.extend(bench_case(helpers, ops, name, shape, args.warmup, args.iterations))
+        results.extend(
+            bench_case(
+                helpers, ops, native, name, shape, args.warmup, args.iterations
+            )
+        )
+    producer_results = [
+        bench_bf16_producer(ops, native, k, args.warmup, args.iterations)
+        for k in (5120, 6144, 17408)
+    ]
     payload = {
         "mode": args.mode,
+        "backend": args.backend,
         "device": torch.cuda.get_device_name(),
         "torch": torch.__version__,
         "results": [asdict(item) for item in results],
+        "bf16_producer_results": producer_results,
     }
     print(json.dumps(payload, indent=2))
     if args.json_out:

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

@@ -23,7 +23,7 @@ def _alloc_fp4(rows: int, dim: int, device: torch.device | str):
     )
 
 
-@torch.library.register_fake(add_op_namespace_prefix("fp4_w4a16_linear_bf16"))
+@torch.library.register_fake(add_op_namespace_prefix("nvfp4_gemm_bf16"))
 def _linear_fake(
     a_packed: torch.Tensor,
     b_packed: torch.Tensor,
@@ -36,6 +36,19 @@ def _linear_fake(
     return None
 
 
+@torch.library.register_fake(add_op_namespace_prefix("fp4_w4a16_linear_bf16"))
+def _legacy_linear_fake(
+    a_packed: torch.Tensor,
+    b_packed: torch.Tensor,
+    sfa: torch.Tensor,
+    sfb: torch.Tensor,
+    out: torch.Tensor,
+    alpha: float = 1.0,
+    variant: int = 0,
+) -> None:
+    return None
+
+
 @torch.library.register_fake(add_op_namespace_prefix("quantize_fp4_sfa_fp16"))
 def _quant_fake(x: torch.Tensor, packed: torch.Tensor, sfa: torch.Tensor, is_sfb: bool = False) -> None:
     return None
@@ -46,6 +59,33 @@ def _dequant_fake(packed: torch.Tensor, sfa: torch.Tensor, out: torch.Tensor, is
     return None
 
 
+@torch.library.register_fake(add_op_namespace_prefix("nvfp4_gemm_residual_bf16"))
+def _residual_fake(a, b, sfa, sfb, residual, out, alpha: float = 1.0) -> None:
+    return None
+
+
+@torch.library.register_fake(add_op_namespace_prefix("nvfp4_gemm_bias_gelu_bf16"))
+def _bias_gelu_fake(a, b, sfa, sfb, bias, out, alpha: float = 1.0) -> None:
+    return None
+
+
+@torch.library.register_fake(add_op_namespace_prefix("nvfp4_gemm_bias_gelu_nvfp4"))
+def _bias_gelu_nvfp4_fake(
+    a, b, sfa, sfb, bias, out_packed, out_sfa, alpha: float = 1.0
+) -> None:
+    return None
+
+
+@torch.library.register_fake(add_op_namespace_prefix("nvfp4_gemm_streamk_bf16"))
+def _streamk_fake(a, b, sfa, sfb, out, alpha: float = 1.0) -> None:
+    return None
+
+
+@torch.library.register_fake(add_op_namespace_prefix("nvfp4_gemm_streamk_bias_bf16"))
+def _streamk_bias_fake(a, b, sfa, sfb, bias, out, alpha: float = 1.0) -> None:
+    return None
+
+
 def quantize_fp4_sfa_fp16(
     x: torch.Tensor,
     packed: torch.Tensor | None = None,
@@ -70,7 +110,7 @@ def dequantize_fp4_sfa_fp16(
     return out
 
 
-def fp4_w4a16_linear_bf16(
+def nvfp4_gemm_bf16(
     a_packed: torch.Tensor,
     b_packed: torch.Tensor,
     sfa: torch.Tensor,
@@ -81,6 +121,134 @@ def fp4_w4a16_linear_bf16(
 ) -> torch.Tensor:
     if out is None:
         out = torch.empty((a_packed.shape[0], b_packed.shape[0]), device=a_packed.device, dtype=torch.bfloat16)
-    ops.fp4_w4a16_linear_bf16(a_packed, b_packed, sfa, sfb, out, float(alpha), int(variant))
+    ops.nvfp4_gemm_bf16(a_packed, b_packed, sfa, sfb, out, float(alpha), int(variant))
     return out
 
+
+def fp4_w4a16_linear_bf16(
+    a_packed: torch.Tensor,
+    b_packed: torch.Tensor,
+    sfa: torch.Tensor,
+    sfb: torch.Tensor,
+    alpha: float = 1.0,
+    out: torch.Tensor | None = None,
+    variant: int = 0,
+) -> torch.Tensor:
+    """Compatibility alias for :func:`nvfp4_gemm_bf16`."""
+    return nvfp4_gemm_bf16(
+        a_packed, b_packed, sfa, sfb, alpha=alpha, out=out, variant=variant
+    )
+
+
+def nvfp4_gemm_residual_bf16(
+    a_packed: torch.Tensor,
+    b_packed: torch.Tensor,
+    sfa: torch.Tensor,
+    sfb: torch.Tensor,
+    residual: torch.Tensor,
+    alpha: float = 1.0,
+    out: torch.Tensor | None = None,
+) -> torch.Tensor:
+    if out is None:
+        out = torch.empty_like(residual)
+    ops.nvfp4_gemm_residual_bf16(
+        a_packed, b_packed, sfa, sfb, residual, out, float(alpha)
+    )
+    return out
+
+
+def nvfp4_gemm_bias_gelu_bf16(
+    a_packed: torch.Tensor,
+    b_packed: torch.Tensor,
+    sfa: torch.Tensor,
+    sfb: torch.Tensor,
+    bias: torch.Tensor,
+    alpha: float = 1.0,
+    out: torch.Tensor | None = None,
+) -> torch.Tensor:
+    if out is None:
+        out = torch.empty(
+            (a_packed.shape[0], b_packed.shape[0]),
+            device=a_packed.device,
+            dtype=torch.bfloat16,
+        )
+    ops.nvfp4_gemm_bias_gelu_bf16(
+        a_packed, b_packed, sfa, sfb, bias, out, float(alpha)
+    )
+    return out
+
+
+def nvfp4_gemm_bias_gelu_nvfp4(
+    a_packed: torch.Tensor,
+    b_packed: torch.Tensor,
+    sfa: torch.Tensor,
+    sfb: torch.Tensor,
+    bias: torch.Tensor,
+    alpha: float = 1.0,
+    out_packed: torch.Tensor | None = None,
+    out_sfa: torch.Tensor | None = None,
+) -> tuple[torch.Tensor, torch.Tensor]:
+    m, n = a_packed.shape[0], b_packed.shape[0]
+    if out_packed is None:
+        out_packed = torch.empty((m, n // 2), device=a_packed.device, dtype=torch.uint8)
+    if out_sfa is None:
+        out_sfa = torch.empty((sfa_size_bytes(m, n),), device=a_packed.device, dtype=torch.uint8)
+    ops.nvfp4_gemm_bias_gelu_nvfp4(
+        a_packed, b_packed, sfa, sfb, bias, out_packed, out_sfa, float(alpha)
+    )
+    return out_packed, out_sfa
+
+
+def nvfp4_gemm_streamk_bf16(
+    a_packed: torch.Tensor,
+    b_packed: torch.Tensor,
+    sfa: torch.Tensor,
+    sfb: torch.Tensor,
+    alpha: float = 1.0,
+    out: torch.Tensor | None = None,
+) -> torch.Tensor:
+    if out is None:
+        out = torch.empty(
+            (a_packed.shape[0], b_packed.shape[0]),
+            device=a_packed.device,
+            dtype=torch.bfloat16,
+        )
+    ops.nvfp4_gemm_streamk_bf16(
+        a_packed, b_packed, sfa, sfb, out, float(alpha)
+    )
+    return out
+
+
+def nvfp4_gemm_streamk_bias_bf16(
+    a_packed: torch.Tensor,
+    b_packed: torch.Tensor,
+    sfa: torch.Tensor,
+    sfb: torch.Tensor,
+    bias: torch.Tensor,
+    alpha: float = 1.0,
+    out: torch.Tensor | None = None,
+) -> torch.Tensor:
+    if out is None:
+        out = torch.empty(
+            (a_packed.shape[0], b_packed.shape[0]),
+            device=a_packed.device,
+            dtype=torch.bfloat16,
+        )
+    ops.nvfp4_gemm_streamk_bias_bf16(
+        a_packed, b_packed, sfa, sfb, bias, out, float(alpha)
+    )
+    return out
+
+
+__all__ = [
+    "dequantize_fp4_sfa_fp16",
+    "fp4_w4a16_linear_bf16",
+    "nvfp4_gemm_bf16",
+    "nvfp4_gemm_bias_gelu_bf16",
+    "nvfp4_gemm_bias_gelu_nvfp4",
+    "nvfp4_gemm_residual_bf16",
+    "nvfp4_gemm_streamk_bf16",
+    "nvfp4_gemm_streamk_bias_bf16",
+    "quantize_fp4_sfa_fp16",
+    "sfa_size_bytes",
+]

build/torch211-cxx11-cu128-x86_64-linux/{_fp4_gemm_cuda_7e2e5b9.abi3.so → _fp4_gemm_cuda_7781728.abi3.so}

@@ -1,3 +1,3 @@
 version https://git-lfs.github.com/spec/v1
-oid sha256:d031c5d46ab5e14b13f68565ad53dff18a7f9ee76296b8158f626b11b50f61ef
-size 671608
+oid sha256:43c2606ea93854faa4bd43f78835977beed5ddd1a005411ab41dc79988d99920
+size 1633016

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

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

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

@@ -1,6 +1,6 @@
 {
   "name": "fp4-gemm",
-  "id": "_fp4_gemm_cuda_7e2e5b9",
+  "id": "_fp4_gemm_cuda_7781728",
   "version": 1,
   "license": "Apache-2.0",
   "python-depends": [],
@@ -13,10 +13,21 @@
   "digest": {
     "algorithm": "sha256",
     "files": {
-      "__init__.py": "WSLdLQD92ulRtOhSLoAhWgFYN6oucbbkTg19GAaltBs=",
-      "_fp4_gemm_cuda_7e2e5b9.abi3.so": "0DHF1Gq14UsT9oVlrVPf8Yp/nudilrgVj2JrEbUPYe8=",
-      "_ops.py": "UnZlP4anh/kfneg4T/V+ueYOufY4kzFv6zxKMKlaRuA=",
+      "__init__.py": "FD6tCd6u2NRYfLKq1Uz7OYi+tZV5o9GDAb+KN3GX2oU=",
+      "_fp4_gemm_cuda_7781728.abi3.so": "Q8Jgbqk4VPqkvUP3iDWXe+7V3dGgBUEatB3HmYjZmSA=",
+      "_ops.py": "4wRBLd5HIJqgJskLFCRBmKJpvQfYeE9AmKTqIPQtPnw=",
       "fp4_gemm/__init__.py": "DFYPlrhXwYjEqCl/8n0SmWGZV8NFml5DPhMjKfv98GY="
     }
+  },
+  "provenance": {
+    "kernel-builder": {
+      "version": "0.17.0-dev0",
+      "sha": "19aaa6421e674e9fecc352bbae6eab81d19a6bf4",
+      "dirty": false
+    },
+    "kernel": {
+      "sha": "77817286c522d8f7abf8d2cd873e1f84a79357b9",
+      "dirty": false
+    }
   }
 }

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

@@ -0,0 +1,254 @@
+"""FlashRT FP4 GEMM kernels."""
+
+from __future__ import annotations
+
+import torch
+
+from ._ops import add_op_namespace_prefix, ops
+
+
+def sfa_size_bytes(rows: int, dim: int) -> int:
+    if rows <= 0 or dim <= 0 or dim % 16 != 0:
+        raise ValueError("rows must be positive and dim must be positive/divisible by 16")
+    n_blocks = dim // 16
+    n_row_super = (rows + 127) // 128
+    n_col_super = (n_blocks + 3) // 4
+    return n_row_super * n_col_super * 512
+
+
+def _alloc_fp4(rows: int, dim: int, device: torch.device | str):
+    return (
+        torch.empty((rows, dim // 2), device=device, dtype=torch.uint8),
+        torch.empty((sfa_size_bytes(rows, dim),), device=device, dtype=torch.uint8),
+    )
+
+
+@torch.library.register_fake(add_op_namespace_prefix("nvfp4_gemm_bf16"))
+def _linear_fake(
+    a_packed: torch.Tensor,
+    b_packed: torch.Tensor,
+    sfa: torch.Tensor,
+    sfb: torch.Tensor,
+    out: torch.Tensor,
+    alpha: float = 1.0,
+    variant: int = -1,
+) -> None:
+    return None
+
+
+@torch.library.register_fake(add_op_namespace_prefix("fp4_w4a16_linear_bf16"))
+def _legacy_linear_fake(
+    a_packed: torch.Tensor,
+    b_packed: torch.Tensor,
+    sfa: torch.Tensor,
+    sfb: torch.Tensor,
+    out: torch.Tensor,
+    alpha: float = 1.0,
+    variant: int = -1,
+) -> None:
+    return None
+
+
+@torch.library.register_fake(add_op_namespace_prefix("quantize_fp4_sfa_fp16"))
+def _quant_fake(x: torch.Tensor, packed: torch.Tensor, sfa: torch.Tensor, is_sfb: bool = False) -> 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, is_sfb: bool = False) -> None:
+    return None
+
+
+@torch.library.register_fake(add_op_namespace_prefix("nvfp4_gemm_residual_bf16"))
+def _residual_fake(a, b, sfa, sfb, residual, out, alpha: float = 1.0) -> None:
+    return None
+
+
+@torch.library.register_fake(add_op_namespace_prefix("nvfp4_gemm_bias_gelu_bf16"))
+def _bias_gelu_fake(a, b, sfa, sfb, bias, out, alpha: float = 1.0) -> None:
+    return None
+
+
+@torch.library.register_fake(add_op_namespace_prefix("nvfp4_gemm_bias_gelu_nvfp4"))
+def _bias_gelu_nvfp4_fake(
+    a, b, sfa, sfb, bias, out_packed, out_sfa, alpha: float = 1.0
+) -> None:
+    return None
+
+
+@torch.library.register_fake(add_op_namespace_prefix("nvfp4_gemm_streamk_bf16"))
+def _streamk_fake(a, b, sfa, sfb, out, alpha: float = 1.0) -> None:
+    return None
+
+
+@torch.library.register_fake(add_op_namespace_prefix("nvfp4_gemm_streamk_bias_bf16"))
+def _streamk_bias_fake(a, b, sfa, sfb, bias, out, alpha: float = 1.0) -> None:
+    return None
+
+
+def quantize_fp4_sfa_fp16(
+    x: torch.Tensor,
+    packed: torch.Tensor | None = None,
+    sfa: torch.Tensor | None = None,
+    is_sfb: bool = False,
+):
+    if packed is None or sfa is None:
+        packed, sfa = _alloc_fp4(x.shape[0], x.shape[1], x.device)
+    ops.quantize_fp4_sfa_fp16(x, packed, sfa, bool(is_sfb))
+    return packed, sfa
+
+
+def dequantize_fp4_sfa_fp16(
+    packed: torch.Tensor,
+    sfa: torch.Tensor,
+    out: torch.Tensor | None = None,
+    is_sfb: bool = False,
+) -> 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, bool(is_sfb))
+    return out
+
+
+def nvfp4_gemm_bf16(
+    a_packed: torch.Tensor,
+    b_packed: torch.Tensor,
+    sfa: torch.Tensor,
+    sfb: torch.Tensor,
+    alpha: float = 1.0,
+    out: torch.Tensor | None = None,
+    variant: int = -1,
+) -> torch.Tensor:
+    if out is None:
+        out = torch.empty((a_packed.shape[0], b_packed.shape[0]), device=a_packed.device, dtype=torch.bfloat16)
+    ops.nvfp4_gemm_bf16(a_packed, b_packed, sfa, sfb, out, float(alpha), int(variant))
+    return out
+
+
+def fp4_w4a16_linear_bf16(
+    a_packed: torch.Tensor,
+    b_packed: torch.Tensor,
+    sfa: torch.Tensor,
+    sfb: torch.Tensor,
+    alpha: float = 1.0,
+    out: torch.Tensor | None = None,
+    variant: int = -1,
+) -> torch.Tensor:
+    """Compatibility alias for :func:`nvfp4_gemm_bf16`."""
+    return nvfp4_gemm_bf16(
+        a_packed, b_packed, sfa, sfb, alpha=alpha, out=out, variant=variant
+    )
+
+
+def nvfp4_gemm_residual_bf16(
+    a_packed: torch.Tensor,
+    b_packed: torch.Tensor,
+    sfa: torch.Tensor,
+    sfb: torch.Tensor,
+    residual: torch.Tensor,
+    alpha: float = 1.0,
+    out: torch.Tensor | None = None,
+) -> torch.Tensor:
+    if out is None:
+        out = torch.empty_like(residual)
+    ops.nvfp4_gemm_residual_bf16(
+        a_packed, b_packed, sfa, sfb, residual, out, float(alpha)
+    )
+    return out
+
+
+def nvfp4_gemm_bias_gelu_bf16(
+    a_packed: torch.Tensor,
+    b_packed: torch.Tensor,
+    sfa: torch.Tensor,
+    sfb: torch.Tensor,
+    bias: torch.Tensor,
+    alpha: float = 1.0,
+    out: torch.Tensor | None = None,
+) -> torch.Tensor:
+    if out is None:
+        out = torch.empty(
+            (a_packed.shape[0], b_packed.shape[0]),
+            device=a_packed.device,
+            dtype=torch.bfloat16,
+        )
+    ops.nvfp4_gemm_bias_gelu_bf16(
+        a_packed, b_packed, sfa, sfb, bias, out, float(alpha)
+    )
+    return out
+
+
+def nvfp4_gemm_bias_gelu_nvfp4(
+    a_packed: torch.Tensor,
+    b_packed: torch.Tensor,
+    sfa: torch.Tensor,
+    sfb: torch.Tensor,
+    bias: torch.Tensor,
+    alpha: float = 1.0,
+    out_packed: torch.Tensor | None = None,
+    out_sfa: torch.Tensor | None = None,
+) -> tuple[torch.Tensor, torch.Tensor]:
+    m, n = a_packed.shape[0], b_packed.shape[0]
+    if out_packed is None:
+        out_packed = torch.empty((m, n // 2), device=a_packed.device, dtype=torch.uint8)
+    if out_sfa is None:
+        out_sfa = torch.empty((sfa_size_bytes(m, n),), device=a_packed.device, dtype=torch.uint8)
+    ops.nvfp4_gemm_bias_gelu_nvfp4(
+        a_packed, b_packed, sfa, sfb, bias, out_packed, out_sfa, float(alpha)
+    )
+    return out_packed, out_sfa
+
+
+def nvfp4_gemm_streamk_bf16(
+    a_packed: torch.Tensor,
+    b_packed: torch.Tensor,
+    sfa: torch.Tensor,
+    sfb: torch.Tensor,
+    alpha: float = 1.0,
+    out: torch.Tensor | None = None,
+) -> torch.Tensor:
+    if out is None:
+        out = torch.empty(
+            (a_packed.shape[0], b_packed.shape[0]),
+            device=a_packed.device,
+            dtype=torch.bfloat16,
+        )
+    ops.nvfp4_gemm_streamk_bf16(
+        a_packed, b_packed, sfa, sfb, out, float(alpha)
+    )
+    return out
+
+
+def nvfp4_gemm_streamk_bias_bf16(
+    a_packed: torch.Tensor,
+    b_packed: torch.Tensor,
+    sfa: torch.Tensor,
+    sfb: torch.Tensor,
+    bias: torch.Tensor,
+    alpha: float = 1.0,
+    out: torch.Tensor | None = None,
+) -> torch.Tensor:
+    if out is None:
+        out = torch.empty(
+            (a_packed.shape[0], b_packed.shape[0]),
+            device=a_packed.device,
+            dtype=torch.bfloat16,
+        )
+    ops.nvfp4_gemm_streamk_bias_bf16(
+        a_packed, b_packed, sfa, sfb, bias, out, float(alpha)
+    )
+    return out
+
+
+__all__ = [
+    "dequantize_fp4_sfa_fp16",
+    "fp4_w4a16_linear_bf16",
+    "nvfp4_gemm_bf16",
+    "nvfp4_gemm_bias_gelu_bf16",
+    "nvfp4_gemm_bias_gelu_nvfp4",
+    "nvfp4_gemm_residual_bf16",
+    "nvfp4_gemm_streamk_bf16",
+    "nvfp4_gemm_streamk_bias_bf16",
+    "quantize_fp4_sfa_fp16",
+    "sfa_size_bytes",
+]

build/{torch212-cxx11-cu132-x86_64-linux/_fp4_gemm_cuda_7e2e5b9.abi3.so → torch211-cxx11-cu130-aarch64-linux/_fp4_gemm_cuda_7781728.abi3.so}

@@ -1,3 +1,3 @@
 version https://git-lfs.github.com/spec/v1
-oid sha256:fd5e66f4faeddfb45f42d8467508f13b2e065fb5d3c52a8404dc6a765a1b891f
-size 724632
+oid sha256:ca5e0135a51929d6d3fdbb486ddfaf1de4b0626150eeb93a00de328036501acb
+size 680560

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

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

build/torch211-cxx11-cu130-aarch64-linux/fp4_gemm/__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-gemm",
+  "id": "_fp4_gemm_cuda_7781728",
+  "version": 1,
+  "license": "Apache-2.0",
+  "python-depends": [],
+  "backend": {
+    "type": "cuda",
+    "archs": [
+      "11.0a"
+    ]
+  },
+  "digest": {
+    "algorithm": "sha256",
+    "files": {
+      "__init__.py": "EJ8HPhVwOKhifotSgdwCmB6MRbtnP1S0QfUGOVascy0=",
+      "_fp4_gemm_cuda_7781728.abi3.so": "yl4BNaUZKdbT/btIbd+vHeSwYmFQ7rk6AN4ygDZQGss=",
+      "_ops.py": "V90sTDYZObC3Ga+eak9fZiO8p6DsMdNZaWpcnYd61Yo=",
+      "fp4_gemm/__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

@@ -23,7 +23,7 @@ def _alloc_fp4(rows: int, dim: int, device: torch.device | str):
     )
 
 
-@torch.library.register_fake(add_op_namespace_prefix("fp4_w4a16_linear_bf16"))
+@torch.library.register_fake(add_op_namespace_prefix("nvfp4_gemm_bf16"))
 def _linear_fake(
     a_packed: torch.Tensor,
     b_packed: torch.Tensor,
@@ -36,6 +36,19 @@ def _linear_fake(
     return None
 
 
+@torch.library.register_fake(add_op_namespace_prefix("fp4_w4a16_linear_bf16"))
+def _legacy_linear_fake(
+    a_packed: torch.Tensor,
+    b_packed: torch.Tensor,
+    sfa: torch.Tensor,
+    sfb: torch.Tensor,
+    out: torch.Tensor,
+    alpha: float = 1.0,
+    variant: int = 0,
+) -> None:
+    return None
+
+
 @torch.library.register_fake(add_op_namespace_prefix("quantize_fp4_sfa_fp16"))
 def _quant_fake(x: torch.Tensor, packed: torch.Tensor, sfa: torch.Tensor, is_sfb: bool = False) -> None:
     return None
@@ -46,6 +59,33 @@ def _dequant_fake(packed: torch.Tensor, sfa: torch.Tensor, out: torch.Tensor, is
     return None
 
 
+@torch.library.register_fake(add_op_namespace_prefix("nvfp4_gemm_residual_bf16"))
+def _residual_fake(a, b, sfa, sfb, residual, out, alpha: float = 1.0) -> None:
+    return None
+
+
+@torch.library.register_fake(add_op_namespace_prefix("nvfp4_gemm_bias_gelu_bf16"))
+def _bias_gelu_fake(a, b, sfa, sfb, bias, out, alpha: float = 1.0) -> None:
+    return None
+
+
+@torch.library.register_fake(add_op_namespace_prefix("nvfp4_gemm_bias_gelu_nvfp4"))
+def _bias_gelu_nvfp4_fake(
+    a, b, sfa, sfb, bias, out_packed, out_sfa, alpha: float = 1.0
+) -> None:
+    return None
+
+
+@torch.library.register_fake(add_op_namespace_prefix("nvfp4_gemm_streamk_bf16"))
+def _streamk_fake(a, b, sfa, sfb, out, alpha: float = 1.0) -> None:
+    return None
+
+
+@torch.library.register_fake(add_op_namespace_prefix("nvfp4_gemm_streamk_bias_bf16"))
+def _streamk_bias_fake(a, b, sfa, sfb, bias, out, alpha: float = 1.0) -> None:
+    return None
+
+
 def quantize_fp4_sfa_fp16(
     x: torch.Tensor,
     packed: torch.Tensor | None = None,
@@ -70,7 +110,7 @@ def dequantize_fp4_sfa_fp16(
     return out
 
 
-def fp4_w4a16_linear_bf16(
+def nvfp4_gemm_bf16(
     a_packed: torch.Tensor,
     b_packed: torch.Tensor,
     sfa: torch.Tensor,
@@ -81,6 +121,134 @@ def fp4_w4a16_linear_bf16(
 ) -> torch.Tensor:
     if out is None:
         out = torch.empty((a_packed.shape[0], b_packed.shape[0]), device=a_packed.device, dtype=torch.bfloat16)
-    ops.fp4_w4a16_linear_bf16(a_packed, b_packed, sfa, sfb, out, float(alpha), int(variant))
+    ops.nvfp4_gemm_bf16(a_packed, b_packed, sfa, sfb, out, float(alpha), int(variant))
     return out
 
+
+def fp4_w4a16_linear_bf16(
+    a_packed: torch.Tensor,
+    b_packed: torch.Tensor,
+    sfa: torch.Tensor,
+    sfb: torch.Tensor,
+    alpha: float = 1.0,
+    out: torch.Tensor | None = None,
+    variant: int = 0,
+) -> torch.Tensor:
+    """Compatibility alias for :func:`nvfp4_gemm_bf16`."""
+    return nvfp4_gemm_bf16(
+        a_packed, b_packed, sfa, sfb, alpha=alpha, out=out, variant=variant
+    )
+
+
+def nvfp4_gemm_residual_bf16(
+    a_packed: torch.Tensor,
+    b_packed: torch.Tensor,
+    sfa: torch.Tensor,
+    sfb: torch.Tensor,
+    residual: torch.Tensor,
+    alpha: float = 1.0,
+    out: torch.Tensor | None = None,
+) -> torch.Tensor:
+    if out is None:
+        out = torch.empty_like(residual)
+    ops.nvfp4_gemm_residual_bf16(
+        a_packed, b_packed, sfa, sfb, residual, out, float(alpha)
+    )
+    return out
+
+
+def nvfp4_gemm_bias_gelu_bf16(
+    a_packed: torch.Tensor,
+    b_packed: torch.Tensor,
+    sfa: torch.Tensor,
+    sfb: torch.Tensor,
+    bias: torch.Tensor,
+    alpha: float = 1.0,
+    out: torch.Tensor | None = None,
+) -> torch.Tensor:
+    if out is None:
+        out = torch.empty(
+            (a_packed.shape[0], b_packed.shape[0]),
+            device=a_packed.device,
+            dtype=torch.bfloat16,
+        )
+    ops.nvfp4_gemm_bias_gelu_bf16(
+        a_packed, b_packed, sfa, sfb, bias, out, float(alpha)
+    )
+    return out
+
+
+def nvfp4_gemm_bias_gelu_nvfp4(
+    a_packed: torch.Tensor,
+    b_packed: torch.Tensor,
+    sfa: torch.Tensor,
+    sfb: torch.Tensor,
+    bias: torch.Tensor,
+    alpha: float = 1.0,
+    out_packed: torch.Tensor | None = None,
+    out_sfa: torch.Tensor | None = None,
+) -> tuple[torch.Tensor, torch.Tensor]:
+    m, n = a_packed.shape[0], b_packed.shape[0]
+    if out_packed is None:
+        out_packed = torch.empty((m, n // 2), device=a_packed.device, dtype=torch.uint8)
+    if out_sfa is None:
+        out_sfa = torch.empty((sfa_size_bytes(m, n),), device=a_packed.device, dtype=torch.uint8)
+    ops.nvfp4_gemm_bias_gelu_nvfp4(
+        a_packed, b_packed, sfa, sfb, bias, out_packed, out_sfa, float(alpha)
+    )
+    return out_packed, out_sfa
+
+
+def nvfp4_gemm_streamk_bf16(
+    a_packed: torch.Tensor,
+    b_packed: torch.Tensor,
+    sfa: torch.Tensor,
+    sfb: torch.Tensor,
+    alpha: float = 1.0,
+    out: torch.Tensor | None = None,
+) -> torch.Tensor:
+    if out is None:
+        out = torch.empty(
+            (a_packed.shape[0], b_packed.shape[0]),
+            device=a_packed.device,
+            dtype=torch.bfloat16,
+        )
+    ops.nvfp4_gemm_streamk_bf16(
+        a_packed, b_packed, sfa, sfb, out, float(alpha)
+    )
+    return out
+
+
+def nvfp4_gemm_streamk_bias_bf16(
+    a_packed: torch.Tensor,
+    b_packed: torch.Tensor,
+    sfa: torch.Tensor,
+    sfb: torch.Tensor,
+    bias: torch.Tensor,
+    alpha: float = 1.0,
+    out: torch.Tensor | None = None,
+) -> torch.Tensor:
+    if out is None:
+        out = torch.empty(
+            (a_packed.shape[0], b_packed.shape[0]),
+            device=a_packed.device,
+            dtype=torch.bfloat16,
+        )
+    ops.nvfp4_gemm_streamk_bias_bf16(
+        a_packed, b_packed, sfa, sfb, bias, out, float(alpha)
+    )
+    return out
+
+
+__all__ = [
+    "dequantize_fp4_sfa_fp16",
+    "fp4_w4a16_linear_bf16",
+    "nvfp4_gemm_bf16",
+    "nvfp4_gemm_bias_gelu_bf16",
+    "nvfp4_gemm_bias_gelu_nvfp4",
+    "nvfp4_gemm_residual_bf16",
+    "nvfp4_gemm_streamk_bf16",
+    "nvfp4_gemm_streamk_bias_bf16",
+    "quantize_fp4_sfa_fp16",
+    "sfa_size_bytes",
+]

build/torch211-cxx11-cu130-x86_64-linux/{_fp4_gemm_cuda_7e2e5b9.abi3.so → _fp4_gemm_cuda_7781728.abi3.so}

@@ -1,3 +1,3 @@
 version https://git-lfs.github.com/spec/v1
-oid sha256:e211cafb6cf95d3ba32015492d7de6bdb29cb6df0cff26400146ed1099ce1cb3
-size 714176
+oid sha256:4c48b26279edf27c786fa4acba75e8b66e410e2f0fbb6fa1c318bfa84e23c392
+size 1721016

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

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

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

@@ -1,6 +1,6 @@
 {
   "name": "fp4-gemm",
-  "id": "_fp4_gemm_cuda_7e2e5b9",
+  "id": "_fp4_gemm_cuda_7781728",
   "version": 1,
   "license": "Apache-2.0",
   "python-depends": [],
@@ -13,10 +13,21 @@
   "digest": {
     "algorithm": "sha256",
     "files": {
-      "__init__.py": "WSLdLQD92ulRtOhSLoAhWgFYN6oucbbkTg19GAaltBs=",
-      "_fp4_gemm_cuda_7e2e5b9.abi3.so": "4hHK+2z5XTujIBVJLX3mvbKctt8M/yZAAUbtEJnOHLM=",
-      "_ops.py": "UnZlP4anh/kfneg4T/V+ueYOufY4kzFv6zxKMKlaRuA=",
+      "__init__.py": "FD6tCd6u2NRYfLKq1Uz7OYi+tZV5o9GDAb+KN3GX2oU=",
+      "_fp4_gemm_cuda_7781728.abi3.so": "TEiyYnnt8nx4b6SsunXotm5BDi8Pu2+hwxi/qE4jw5I=",
+      "_ops.py": "4wRBLd5HIJqgJskLFCRBmKJpvQfYeE9AmKTqIPQtPnw=",
       "fp4_gemm/__init__.py": "DFYPlrhXwYjEqCl/8n0SmWGZV8NFml5DPhMjKfv98GY="
     }
+  },
+  "provenance": {
+    "kernel-builder": {
+      "version": "0.17.0-dev0",
+      "sha": "19aaa6421e674e9fecc352bbae6eab81d19a6bf4",
+      "dirty": false
+    },
+    "kernel": {
+      "sha": "77817286c522d8f7abf8d2cd873e1f84a79357b9",
+      "dirty": false
+    }
   }
 }

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

@@ -23,7 +23,7 @@ def _alloc_fp4(rows: int, dim: int, device: torch.device | str):
     )
 
 
-@torch.library.register_fake(add_op_namespace_prefix("fp4_w4a16_linear_bf16"))
+@torch.library.register_fake(add_op_namespace_prefix("nvfp4_gemm_bf16"))
 def _linear_fake(
     a_packed: torch.Tensor,
     b_packed: torch.Tensor,
@@ -31,7 +31,29 @@ def _linear_fake(
     sfb: torch.Tensor,
     out: torch.Tensor,
     alpha: float = 1.0,
-    variant: int = 0,
+    variant: int = -1,
+) -> None:
+    return None
+
+
+@torch.library.register_fake(add_op_namespace_prefix("fp4_w4a4_gemv_warpsplit_bf16"))
+def _gemv_warpsplit_fake(a_packed, b_packed, sfa, sfb, out, alpha: float = 1.0, warps: int = 4, stages: int = 4) -> None:
+    if a_packed.shape[0] != 1:
+        raise RuntimeError("warp-split GEMV serves M=1 only")
+    if out.shape != (1, b_packed.shape[0]):
+        raise RuntimeError("out must have shape (1, N)")
+    return None
+
+
+@torch.library.register_fake(add_op_namespace_prefix("fp4_w4a16_linear_bf16"))
+def _legacy_linear_fake(
+    a_packed: torch.Tensor,
+    b_packed: torch.Tensor,
+    sfa: torch.Tensor,
+    sfb: torch.Tensor,
+    out: torch.Tensor,
+    alpha: float = 1.0,
+    variant: int = -1,
 ) -> None:
     return None
 
@@ -41,11 +63,43 @@ def _quant_fake(x: torch.Tensor, packed: torch.Tensor, sfa: torch.Tensor, is_sfb
     return None
 
 
+@torch.library.register_fake(add_op_namespace_prefix("quantize_fp4_sfa_bf16"))
+def _quant_bf16_fake(x: torch.Tensor, packed: torch.Tensor, sfa: torch.Tensor, is_sfb: bool = False) -> 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, is_sfb: bool = False) -> None:
     return None
 
 
+@torch.library.register_fake(add_op_namespace_prefix("nvfp4_gemm_residual_bf16"))
+def _residual_fake(a, b, sfa, sfb, residual, out, alpha: float = 1.0) -> None:
+    return None
+
+
+@torch.library.register_fake(add_op_namespace_prefix("nvfp4_gemm_bias_gelu_bf16"))
+def _bias_gelu_fake(a, b, sfa, sfb, bias, out, alpha: float = 1.0) -> None:
+    return None
+
+
+@torch.library.register_fake(add_op_namespace_prefix("nvfp4_gemm_bias_gelu_nvfp4"))
+def _bias_gelu_nvfp4_fake(
+    a, b, sfa, sfb, bias, out_packed, out_sfa, alpha: float = 1.0
+) -> None:
+    return None
+
+
+@torch.library.register_fake(add_op_namespace_prefix("nvfp4_gemm_streamk_bf16"))
+def _streamk_fake(a, b, sfa, sfb, out, alpha: float = 1.0) -> None:
+    return None
+
+
+@torch.library.register_fake(add_op_namespace_prefix("nvfp4_gemm_streamk_bias_bf16"))
+def _streamk_bias_fake(a, b, sfa, sfb, bias, out, alpha: float = 1.0) -> None:
+    return None
+
+
 def quantize_fp4_sfa_fp16(
     x: torch.Tensor,
     packed: torch.Tensor | None = None,
@@ -58,6 +112,19 @@ def quantize_fp4_sfa_fp16(
     return packed, sfa
 
 
+def quantize_fp4_sfa_bf16(
+    x: torch.Tensor,
+    packed: torch.Tensor | None = None,
+    sfa: torch.Tensor | None = None,
+    is_sfb: bool = False,
+):
+    """Quantize BF16 directly to packed E2M1 and CUTLASS SFA/SFB."""
+    if packed is None or sfa is None:
+        packed, sfa = _alloc_fp4(x.shape[0], x.shape[1], x.device)
+    ops.quantize_fp4_sfa_bf16(x, packed, sfa, bool(is_sfb))
+    return packed, sfa
+
+
 def dequantize_fp4_sfa_fp16(
     packed: torch.Tensor,
     sfa: torch.Tensor,
@@ -70,17 +137,171 @@ def dequantize_fp4_sfa_fp16(
     return out
 
 
-def fp4_w4a16_linear_bf16(
+def nvfp4_gemm_bf16(
     a_packed: torch.Tensor,
     b_packed: torch.Tensor,
     sfa: torch.Tensor,
     sfb: torch.Tensor,
     alpha: float = 1.0,
     out: torch.Tensor | None = None,
-    variant: int = 0,
+    variant: int = -1,
 ) -> torch.Tensor:
     if out is None:
         out = torch.empty((a_packed.shape[0], b_packed.shape[0]), device=a_packed.device, dtype=torch.bfloat16)
-    ops.fp4_w4a16_linear_bf16(a_packed, b_packed, sfa, sfb, out, float(alpha), int(variant))
+    ops.nvfp4_gemm_bf16(a_packed, b_packed, sfa, sfb, out, float(alpha), int(variant))
+    return out
+
+
+def fp4_w4a4_gemv_warpsplit_bf16(
+    a_packed: torch.Tensor,
+    b_packed: torch.Tensor,
+    sfa: torch.Tensor,
+    sfb: torch.Tensor,
+    *,
+    alpha: float = 1.0,
+    warps: int = 4,
+    stages: int = 4,
+    out: Optional[torch.Tensor] = None,
+) -> torch.Tensor:
+    """Warp-split-K NVFP4 W4A4 GEMV for the M=1 decode row (SM120).
+
+    Splits K across warps inside one block with a shared-memory reduce -
+    no cross-block intermediate, so it stays safe under CUDA-graph
+    replay - and fills the SMs the tiled GEMM underfills at long-K
+    small-M decode shapes. Same packed/scale layouts as the linear
+    entry points."""
+    if out is None:
+        out = torch.empty((1, b_packed.shape[0]), device=a_packed.device, dtype=torch.bfloat16)
+    ops.fp4_w4a4_gemv_warpsplit_bf16(a_packed, b_packed, sfa, sfb, out, float(alpha), int(warps), int(stages))
+    return out
+
+
+def fp4_w4a16_linear_bf16(
+    a_packed: torch.Tensor,
+    b_packed: torch.Tensor,
+    sfa: torch.Tensor,
+    sfb: torch.Tensor,
+    alpha: float = 1.0,
+    out: torch.Tensor | None = None,
+    variant: int = -1,
+) -> torch.Tensor:
+    """Compatibility alias for :func:`nvfp4_gemm_bf16`."""
+    return nvfp4_gemm_bf16(
+        a_packed, b_packed, sfa, sfb, alpha=alpha, out=out, variant=variant
+    )
+
+
+def nvfp4_gemm_residual_bf16(
+    a_packed: torch.Tensor,
+    b_packed: torch.Tensor,
+    sfa: torch.Tensor,
+    sfb: torch.Tensor,
+    residual: torch.Tensor,
+    alpha: float = 1.0,
+    out: torch.Tensor | None = None,
+) -> torch.Tensor:
+    if out is None:
+        out = torch.empty_like(residual)
+    ops.nvfp4_gemm_residual_bf16(
+        a_packed, b_packed, sfa, sfb, residual, out, float(alpha)
+    )
+    return out
+
+
+def nvfp4_gemm_bias_gelu_bf16(
+    a_packed: torch.Tensor,
+    b_packed: torch.Tensor,
+    sfa: torch.Tensor,
+    sfb: torch.Tensor,
+    bias: torch.Tensor,
+    alpha: float = 1.0,
+    out: torch.Tensor | None = None,
+) -> torch.Tensor:
+    if out is None:
+        out = torch.empty(
+            (a_packed.shape[0], b_packed.shape[0]),
+            device=a_packed.device,
+            dtype=torch.bfloat16,
+        )
+    ops.nvfp4_gemm_bias_gelu_bf16(
+        a_packed, b_packed, sfa, sfb, bias, out, float(alpha)
+    )
+    return out
+
+
+def nvfp4_gemm_bias_gelu_nvfp4(
+    a_packed: torch.Tensor,
+    b_packed: torch.Tensor,
+    sfa: torch.Tensor,
+    sfb: torch.Tensor,
+    bias: torch.Tensor,
+    alpha: float = 1.0,
+    out_packed: torch.Tensor | None = None,
+    out_sfa: torch.Tensor | None = None,
+) -> tuple[torch.Tensor, torch.Tensor]:
+    m, n = a_packed.shape[0], b_packed.shape[0]
+    if out_packed is None:
+        out_packed = torch.empty((m, n // 2), device=a_packed.device, dtype=torch.uint8)
+    if out_sfa is None:
+        out_sfa = torch.empty((sfa_size_bytes(m, n),), device=a_packed.device, dtype=torch.uint8)
+    ops.nvfp4_gemm_bias_gelu_nvfp4(
+        a_packed, b_packed, sfa, sfb, bias, out_packed, out_sfa, float(alpha)
+    )
+    return out_packed, out_sfa
+
+
+def nvfp4_gemm_streamk_bf16(
+    a_packed: torch.Tensor,
+    b_packed: torch.Tensor,
+    sfa: torch.Tensor,
+    sfb: torch.Tensor,
+    alpha: float = 1.0,
+    out: torch.Tensor | None = None,
+) -> torch.Tensor:
+    if out is None:
+        out = torch.empty(
+            (a_packed.shape[0], b_packed.shape[0]),
+            device=a_packed.device,
+            dtype=torch.bfloat16,
+        )
+    ops.nvfp4_gemm_streamk_bf16(
+        a_packed, b_packed, sfa, sfb, out, float(alpha)
+    )
     return out
 
+
+def nvfp4_gemm_streamk_bias_bf16(
+    a_packed: torch.Tensor,
+    b_packed: torch.Tensor,
+    sfa: torch.Tensor,
+    sfb: torch.Tensor,
+    bias: torch.Tensor,
+    alpha: float = 1.0,
+    out: torch.Tensor | None = None,
+) -> torch.Tensor:
+    if out is None:
+        out = torch.empty(
+            (a_packed.shape[0], b_packed.shape[0]),
+            device=a_packed.device,
+            dtype=torch.bfloat16,
+        )
+    ops.nvfp4_gemm_streamk_bias_bf16(
+        a_packed, b_packed, sfa, sfb, bias, out, float(alpha)
+    )
+    return out
+
+
+__all__ = [
+    "dequantize_fp4_sfa_fp16",
+    "fp4_w4a16_linear_bf16",
+    "fp4_w4a4_gemv_warpsplit_bf16",
+    "nvfp4_gemm_bf16",
+    "nvfp4_gemm_bias_gelu_bf16",
+    "nvfp4_gemm_bias_gelu_nvfp4",
+    "nvfp4_gemm_residual_bf16",
+    "nvfp4_gemm_streamk_bf16",
+    "nvfp4_gemm_streamk_bias_bf16",
+    "quantize_fp4_sfa_fp16",
+    "quantize_fp4_sfa_bf16",
+    "sfa_size_bytes",
+]

build/torch212-cxx11-cu130-x86_64-linux/{_fp4_gemm_cuda_7e2e5b9.abi3.so → _fp4_gemm_cuda_b46a817.abi3.so}

@@ -1,3 +1,3 @@
 version https://git-lfs.github.com/spec/v1
-oid sha256:63e8e675f69f6cdb3cf1c0a78fe0032f72638db29e8e4844501a2204cca77e08
-size 724584
+oid sha256:1234af9f97c77fdec732392890c5b8fa889a416b3034c8c7491edfbecb4cd4d8
+size 2428552

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

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

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

@@ -1,26 +0,0 @@
-import ctypes
-import importlib.util
-import sys
-from pathlib import Path
-from types import ModuleType
-
-
-def _import_from_path(file_path: Path) -> ModuleType:
-    # We cannot use the module name as-is, after adding it to `sys.modules`,
-    # it would also be used for other imports. So, we make a module name that
-    # depends on the path for it to be unique using the hex-encoded hash of
-    # the path.
-    path_hash = "{:x}".format(ctypes.c_size_t(hash(file_path.absolute())).value)
-    module_name = path_hash
-    spec = importlib.util.spec_from_file_location(module_name, file_path)
-    if spec is None:
-        raise ImportError(f"Cannot load spec for {module_name} from {file_path}")
-    module = importlib.util.module_from_spec(spec)
-    if module is None:
-        raise ImportError(f"Cannot load module {module_name} from spec")
-    sys.modules[module_name] = module
-    spec.loader.exec_module(module)  # type: ignore
-    return module
-
-
-globals().update(vars(_import_from_path(Path(__file__).parent.parent / "__init__.py")))

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

@@ -1,22 +1,33 @@
 {
   "name": "fp4-gemm",
-  "id": "_fp4_gemm_cuda_7e2e5b9",
+  "id": "_fp4_gemm_cuda_b46a817",
   "version": 1,
   "license": "Apache-2.0",
   "python-depends": [],
   "backend": {
     "type": "cuda",
     "archs": [
+      "11.0a",
       "12.0a"
     ]
   },
   "digest": {
     "algorithm": "sha256",
     "files": {
-      "__init__.py": "WSLdLQD92ulRtOhSLoAhWgFYN6oucbbkTg19GAaltBs=",
-      "_fp4_gemm_cuda_7e2e5b9.abi3.so": "Y+jmdfafbNs88cCnj+ADL3JjjbKejkhEUBoiBMynfgg=",
-      "_ops.py": "UnZlP4anh/kfneg4T/V+ueYOufY4kzFv6zxKMKlaRuA=",
-      "fp4_gemm/__init__.py": "DFYPlrhXwYjEqCl/8n0SmWGZV8NFml5DPhMjKfv98GY="
+      "__init__.py": "pJlGLDlOKnju+m59WyZJjpynMTvkPQTF6w+z4YfCwSY=",
+      "_fp4_gemm_cuda_b46a817.abi3.so": "EjSvn5fHf97HMjkokMW4+oiaQWswNMjHSR7fvstM1Ng=",
+      "_ops.py": "tD2Fh7MGjwwYSATplu1PwYdWuhn0eHO42VKMA5pSjgU="
+    }
+  },
+  "provenance": {
+    "kernel-builder": {
+      "version": "0.17.0-dev0",
+      "sha": "d720fa90fb9cd92d1bc60a9dc5c55bef2aafabb8",
+      "dirty": false
+    },
+    "kernel": {
+      "sha": "b46a81771fdabe26f2c2494dc2ecf8492c88c45a",
+      "dirty": false
     }
   }
 }

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

@@ -23,7 +23,7 @@ def _alloc_fp4(rows: int, dim: int, device: torch.device | str):
     )
 
 
-@torch.library.register_fake(add_op_namespace_prefix("fp4_w4a16_linear_bf16"))
+@torch.library.register_fake(add_op_namespace_prefix("nvfp4_gemm_bf16"))
 def _linear_fake(
     a_packed: torch.Tensor,
     b_packed: torch.Tensor,
@@ -31,7 +31,29 @@ def _linear_fake(
     sfb: torch.Tensor,
     out: torch.Tensor,
     alpha: float = 1.0,
-    variant: int = 0,
+    variant: int = -1,
+) -> None:
+    return None
+
+
+@torch.library.register_fake(add_op_namespace_prefix("fp4_w4a4_gemv_warpsplit_bf16"))
+def _gemv_warpsplit_fake(a_packed, b_packed, sfa, sfb, out, alpha: float = 1.0, warps: int = 4, stages: int = 4) -> None:
+    if a_packed.shape[0] != 1:
+        raise RuntimeError("warp-split GEMV serves M=1 only")
+    if out.shape != (1, b_packed.shape[0]):
+        raise RuntimeError("out must have shape (1, N)")
+    return None
+
+
+@torch.library.register_fake(add_op_namespace_prefix("fp4_w4a16_linear_bf16"))
+def _legacy_linear_fake(
+    a_packed: torch.Tensor,
+    b_packed: torch.Tensor,
+    sfa: torch.Tensor,
+    sfb: torch.Tensor,
+    out: torch.Tensor,
+    alpha: float = 1.0,
+    variant: int = -1,
 ) -> None:
     return None
 
@@ -41,11 +63,43 @@ def _quant_fake(x: torch.Tensor, packed: torch.Tensor, sfa: torch.Tensor, is_sfb
     return None
 
 
+@torch.library.register_fake(add_op_namespace_prefix("quantize_fp4_sfa_bf16"))
+def _quant_bf16_fake(x: torch.Tensor, packed: torch.Tensor, sfa: torch.Tensor, is_sfb: bool = False) -> 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, is_sfb: bool = False) -> None:
     return None
 
 
+@torch.library.register_fake(add_op_namespace_prefix("nvfp4_gemm_residual_bf16"))
+def _residual_fake(a, b, sfa, sfb, residual, out, alpha: float = 1.0) -> None:
+    return None
+
+
+@torch.library.register_fake(add_op_namespace_prefix("nvfp4_gemm_bias_gelu_bf16"))
+def _bias_gelu_fake(a, b, sfa, sfb, bias, out, alpha: float = 1.0) -> None:
+    return None
+
+
+@torch.library.register_fake(add_op_namespace_prefix("nvfp4_gemm_bias_gelu_nvfp4"))
+def _bias_gelu_nvfp4_fake(
+    a, b, sfa, sfb, bias, out_packed, out_sfa, alpha: float = 1.0
+) -> None:
+    return None
+
+
+@torch.library.register_fake(add_op_namespace_prefix("nvfp4_gemm_streamk_bf16"))
+def _streamk_fake(a, b, sfa, sfb, out, alpha: float = 1.0) -> None:
+    return None
+
+
+@torch.library.register_fake(add_op_namespace_prefix("nvfp4_gemm_streamk_bias_bf16"))
+def _streamk_bias_fake(a, b, sfa, sfb, bias, out, alpha: float = 1.0) -> None:
+    return None
+
+
 def quantize_fp4_sfa_fp16(
     x: torch.Tensor,
     packed: torch.Tensor | None = None,
@@ -58,6 +112,19 @@ def quantize_fp4_sfa_fp16(
     return packed, sfa
 
 
+def quantize_fp4_sfa_bf16(
+    x: torch.Tensor,
+    packed: torch.Tensor | None = None,
+    sfa: torch.Tensor | None = None,
+    is_sfb: bool = False,
+):
+    """Quantize BF16 directly to packed E2M1 and CUTLASS SFA/SFB."""
+    if packed is None or sfa is None:
+        packed, sfa = _alloc_fp4(x.shape[0], x.shape[1], x.device)
+    ops.quantize_fp4_sfa_bf16(x, packed, sfa, bool(is_sfb))
+    return packed, sfa
+
+
 def dequantize_fp4_sfa_fp16(
     packed: torch.Tensor,
     sfa: torch.Tensor,
@@ -70,17 +137,171 @@ def dequantize_fp4_sfa_fp16(
     return out
 
 
-def fp4_w4a16_linear_bf16(
+def nvfp4_gemm_bf16(
     a_packed: torch.Tensor,
     b_packed: torch.Tensor,
     sfa: torch.Tensor,
     sfb: torch.Tensor,
     alpha: float = 1.0,
     out: torch.Tensor | None = None,
-    variant: int = 0,
+    variant: int = -1,
 ) -> torch.Tensor:
     if out is None:
         out = torch.empty((a_packed.shape[0], b_packed.shape[0]), device=a_packed.device, dtype=torch.bfloat16)
-    ops.fp4_w4a16_linear_bf16(a_packed, b_packed, sfa, sfb, out, float(alpha), int(variant))
+    ops.nvfp4_gemm_bf16(a_packed, b_packed, sfa, sfb, out, float(alpha), int(variant))
+    return out
+
+
+def fp4_w4a4_gemv_warpsplit_bf16(
+    a_packed: torch.Tensor,
+    b_packed: torch.Tensor,
+    sfa: torch.Tensor,
+    sfb: torch.Tensor,
+    *,
+    alpha: float = 1.0,
+    warps: int = 4,
+    stages: int = 4,
+    out: Optional[torch.Tensor] = None,
+) -> torch.Tensor:
+    """Warp-split-K NVFP4 W4A4 GEMV for the M=1 decode row (SM120).
+
+    Splits K across warps inside one block with a shared-memory reduce -
+    no cross-block intermediate, so it stays safe under CUDA-graph
+    replay - and fills the SMs the tiled GEMM underfills at long-K
+    small-M decode shapes. Same packed/scale layouts as the linear
+    entry points."""
+    if out is None:
+        out = torch.empty((1, b_packed.shape[0]), device=a_packed.device, dtype=torch.bfloat16)
+    ops.fp4_w4a4_gemv_warpsplit_bf16(a_packed, b_packed, sfa, sfb, out, float(alpha), int(warps), int(stages))
+    return out
+
+
+def fp4_w4a16_linear_bf16(
+    a_packed: torch.Tensor,
+    b_packed: torch.Tensor,
+    sfa: torch.Tensor,
+    sfb: torch.Tensor,
+    alpha: float = 1.0,
+    out: torch.Tensor | None = None,
+    variant: int = -1,
+) -> torch.Tensor:
+    """Compatibility alias for :func:`nvfp4_gemm_bf16`."""
+    return nvfp4_gemm_bf16(
+        a_packed, b_packed, sfa, sfb, alpha=alpha, out=out, variant=variant
+    )
+
+
+def nvfp4_gemm_residual_bf16(
+    a_packed: torch.Tensor,
+    b_packed: torch.Tensor,
+    sfa: torch.Tensor,
+    sfb: torch.Tensor,
+    residual: torch.Tensor,
+    alpha: float = 1.0,
+    out: torch.Tensor | None = None,
+) -> torch.Tensor:
+    if out is None:
+        out = torch.empty_like(residual)
+    ops.nvfp4_gemm_residual_bf16(
+        a_packed, b_packed, sfa, sfb, residual, out, float(alpha)
+    )
+    return out
+
+
+def nvfp4_gemm_bias_gelu_bf16(
+    a_packed: torch.Tensor,
+    b_packed: torch.Tensor,
+    sfa: torch.Tensor,
+    sfb: torch.Tensor,
+    bias: torch.Tensor,
+    alpha: float = 1.0,
+    out: torch.Tensor | None = None,
+) -> torch.Tensor:
+    if out is None:
+        out = torch.empty(
+            (a_packed.shape[0], b_packed.shape[0]),
+            device=a_packed.device,
+            dtype=torch.bfloat16,
+        )
+    ops.nvfp4_gemm_bias_gelu_bf16(
+        a_packed, b_packed, sfa, sfb, bias, out, float(alpha)
+    )
+    return out
+
+
+def nvfp4_gemm_bias_gelu_nvfp4(
+    a_packed: torch.Tensor,
+    b_packed: torch.Tensor,
+    sfa: torch.Tensor,
+    sfb: torch.Tensor,
+    bias: torch.Tensor,
+    alpha: float = 1.0,
+    out_packed: torch.Tensor | None = None,
+    out_sfa: torch.Tensor | None = None,
+) -> tuple[torch.Tensor, torch.Tensor]:
+    m, n = a_packed.shape[0], b_packed.shape[0]
+    if out_packed is None:
+        out_packed = torch.empty((m, n // 2), device=a_packed.device, dtype=torch.uint8)
+    if out_sfa is None:
+        out_sfa = torch.empty((sfa_size_bytes(m, n),), device=a_packed.device, dtype=torch.uint8)
+    ops.nvfp4_gemm_bias_gelu_nvfp4(
+        a_packed, b_packed, sfa, sfb, bias, out_packed, out_sfa, float(alpha)
+    )
+    return out_packed, out_sfa
+
+
+def nvfp4_gemm_streamk_bf16(
+    a_packed: torch.Tensor,
+    b_packed: torch.Tensor,
+    sfa: torch.Tensor,
+    sfb: torch.Tensor,
+    alpha: float = 1.0,
+    out: torch.Tensor | None = None,
+) -> torch.Tensor:
+    if out is None:
+        out = torch.empty(
+            (a_packed.shape[0], b_packed.shape[0]),
+            device=a_packed.device,
+            dtype=torch.bfloat16,
+        )
+    ops.nvfp4_gemm_streamk_bf16(
+        a_packed, b_packed, sfa, sfb, out, float(alpha)
+    )
     return out
 
+
+def nvfp4_gemm_streamk_bias_bf16(
+    a_packed: torch.Tensor,
+    b_packed: torch.Tensor,
+    sfa: torch.Tensor,
+    sfb: torch.Tensor,
+    bias: torch.Tensor,
+    alpha: float = 1.0,
+    out: torch.Tensor | None = None,
+) -> torch.Tensor:
+    if out is None:
+        out = torch.empty(
+            (a_packed.shape[0], b_packed.shape[0]),
+            device=a_packed.device,
+            dtype=torch.bfloat16,
+        )
+    ops.nvfp4_gemm_streamk_bias_bf16(
+        a_packed, b_packed, sfa, sfb, bias, out, float(alpha)
+    )
+    return out
+
+
+__all__ = [
+    "dequantize_fp4_sfa_fp16",
+    "fp4_w4a16_linear_bf16",
+    "fp4_w4a4_gemv_warpsplit_bf16",
+    "nvfp4_gemm_bf16",
+    "nvfp4_gemm_bias_gelu_bf16",
+    "nvfp4_gemm_bias_gelu_nvfp4",
+    "nvfp4_gemm_residual_bf16",
+    "nvfp4_gemm_streamk_bf16",
+    "nvfp4_gemm_streamk_bias_bf16",
+    "quantize_fp4_sfa_fp16",
+    "quantize_fp4_sfa_bf16",
+    "sfa_size_bytes",
+]

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

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

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

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

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

@@ -1,26 +0,0 @@
-import ctypes
-import importlib.util
-import sys
-from pathlib import Path
-from types import ModuleType
-
-
-def _import_from_path(file_path: Path) -> ModuleType:
-    # We cannot use the module name as-is, after adding it to `sys.modules`,
-    # it would also be used for other imports. So, we make a module name that
-    # depends on the path for it to be unique using the hex-encoded hash of
-    # the path.
-    path_hash = "{:x}".format(ctypes.c_size_t(hash(file_path.absolute())).value)
-    module_name = path_hash
-    spec = importlib.util.spec_from_file_location(module_name, file_path)
-    if spec is None:
-        raise ImportError(f"Cannot load spec for {module_name} from {file_path}")
-    module = importlib.util.module_from_spec(spec)
-    if module is None:
-        raise ImportError(f"Cannot load module {module_name} from spec")
-    sys.modules[module_name] = module
-    spec.loader.exec_module(module)  # type: ignore
-    return module
-
-
-globals().update(vars(_import_from_path(Path(__file__).parent.parent / "__init__.py")))

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

@@ -1,22 +1,33 @@
 {
   "name": "fp4-gemm",
-  "id": "_fp4_gemm_cuda_7e2e5b9",
+  "id": "_fp4_gemm_cuda_b46a817",
   "version": 1,
   "license": "Apache-2.0",
   "python-depends": [],
   "backend": {
     "type": "cuda",
     "archs": [
+      "11.0a",
       "12.0a"
     ]
   },
   "digest": {
     "algorithm": "sha256",
     "files": {
-      "__init__.py": "WSLdLQD92ulRtOhSLoAhWgFYN6oucbbkTg19GAaltBs=",
-      "_fp4_gemm_cuda_7e2e5b9.abi3.so": "/V5m9Prt37RfQthGdQjxOy4GX7XTxSqEBNxqdlobiR8=",
-      "_ops.py": "UnZlP4anh/kfneg4T/V+ueYOufY4kzFv6zxKMKlaRuA=",
-      "fp4_gemm/__init__.py": "DFYPlrhXwYjEqCl/8n0SmWGZV8NFml5DPhMjKfv98GY="
+      "__init__.py": "pJlGLDlOKnju+m59WyZJjpynMTvkPQTF6w+z4YfCwSY=",
+      "_fp4_gemm_cuda_b46a817.abi3.so": "uIA5FtAYSEX0w0H845Hc31escVtviFMZ5/KBkGMp3Ug=",
+      "_ops.py": "tD2Fh7MGjwwYSATplu1PwYdWuhn0eHO42VKMA5pSjgU="
+    }
+  },
+  "provenance": {
+    "kernel-builder": {
+      "version": "0.17.0-dev0",
+      "sha": "d720fa90fb9cd92d1bc60a9dc5c55bef2aafabb8",
+      "dirty": false
+    },
+    "kernel": {
+      "sha": "b46a81771fdabe26f2c2494dc2ecf8492c88c45a",
+      "dirty": false
     }
   }
 }

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

@@ -0,0 +1,307 @@
+"""FlashRT FP4 GEMM kernels."""
+
+from __future__ import annotations
+
+import torch
+
+from ._ops import add_op_namespace_prefix, ops
+
+
+def sfa_size_bytes(rows: int, dim: int) -> int:
+    if rows <= 0 or dim <= 0 or dim % 16 != 0:
+        raise ValueError("rows must be positive and dim must be positive/divisible by 16")
+    n_blocks = dim // 16
+    n_row_super = (rows + 127) // 128
+    n_col_super = (n_blocks + 3) // 4
+    return n_row_super * n_col_super * 512
+
+
+def _alloc_fp4(rows: int, dim: int, device: torch.device | str):
+    return (
+        torch.empty((rows, dim // 2), device=device, dtype=torch.uint8),
+        torch.empty((sfa_size_bytes(rows, dim),), device=device, dtype=torch.uint8),
+    )
+
+
+@torch.library.register_fake(add_op_namespace_prefix("nvfp4_gemm_bf16"))
+def _linear_fake(
+    a_packed: torch.Tensor,
+    b_packed: torch.Tensor,
+    sfa: torch.Tensor,
+    sfb: torch.Tensor,
+    out: torch.Tensor,
+    alpha: float = 1.0,
+    variant: int = -1,
+) -> None:
+    return None
+
+
+@torch.library.register_fake(add_op_namespace_prefix("fp4_w4a4_gemv_warpsplit_bf16"))
+def _gemv_warpsplit_fake(a_packed, b_packed, sfa, sfb, out, alpha: float = 1.0, warps: int = 4, stages: int = 4) -> None:
+    if a_packed.shape[0] != 1:
+        raise RuntimeError("warp-split GEMV serves M=1 only")
+    if out.shape != (1, b_packed.shape[0]):
+        raise RuntimeError("out must have shape (1, N)")
+    return None
+
+
+@torch.library.register_fake(add_op_namespace_prefix("fp4_w4a16_linear_bf16"))
+def _legacy_linear_fake(
+    a_packed: torch.Tensor,
+    b_packed: torch.Tensor,
+    sfa: torch.Tensor,
+    sfb: torch.Tensor,
+    out: torch.Tensor,
+    alpha: float = 1.0,
+    variant: int = -1,
+) -> None:
+    return None
+
+
+@torch.library.register_fake(add_op_namespace_prefix("quantize_fp4_sfa_fp16"))
+def _quant_fake(x: torch.Tensor, packed: torch.Tensor, sfa: torch.Tensor, is_sfb: bool = False) -> None:
+    return None
+
+
+@torch.library.register_fake(add_op_namespace_prefix("quantize_fp4_sfa_bf16"))
+def _quant_bf16_fake(x: torch.Tensor, packed: torch.Tensor, sfa: torch.Tensor, is_sfb: bool = False) -> 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, is_sfb: bool = False) -> None:
+    return None
+
+
+@torch.library.register_fake(add_op_namespace_prefix("nvfp4_gemm_residual_bf16"))
+def _residual_fake(a, b, sfa, sfb, residual, out, alpha: float = 1.0) -> None:
+    return None
+
+
+@torch.library.register_fake(add_op_namespace_prefix("nvfp4_gemm_bias_gelu_bf16"))
+def _bias_gelu_fake(a, b, sfa, sfb, bias, out, alpha: float = 1.0) -> None:
+    return None
+
+
+@torch.library.register_fake(add_op_namespace_prefix("nvfp4_gemm_bias_gelu_nvfp4"))
+def _bias_gelu_nvfp4_fake(
+    a, b, sfa, sfb, bias, out_packed, out_sfa, alpha: float = 1.0
+) -> None:
+    return None
+
+
+@torch.library.register_fake(add_op_namespace_prefix("nvfp4_gemm_streamk_bf16"))
+def _streamk_fake(a, b, sfa, sfb, out, alpha: float = 1.0) -> None:
+    return None
+
+
+@torch.library.register_fake(add_op_namespace_prefix("nvfp4_gemm_streamk_bias_bf16"))
+def _streamk_bias_fake(a, b, sfa, sfb, bias, out, alpha: float = 1.0) -> None:
+    return None
+
+
+def quantize_fp4_sfa_fp16(
+    x: torch.Tensor,
+    packed: torch.Tensor | None = None,
+    sfa: torch.Tensor | None = None,
+    is_sfb: bool = False,
+):
+    if packed is None or sfa is None:
+        packed, sfa = _alloc_fp4(x.shape[0], x.shape[1], x.device)
+    ops.quantize_fp4_sfa_fp16(x, packed, sfa, bool(is_sfb))
+    return packed, sfa
+
+
+def quantize_fp4_sfa_bf16(
+    x: torch.Tensor,
+    packed: torch.Tensor | None = None,
+    sfa: torch.Tensor | None = None,
+    is_sfb: bool = False,
+):
+    """Quantize BF16 directly to packed E2M1 and CUTLASS SFA/SFB."""
+    if packed is None or sfa is None:
+        packed, sfa = _alloc_fp4(x.shape[0], x.shape[1], x.device)
+    ops.quantize_fp4_sfa_bf16(x, packed, sfa, bool(is_sfb))
+    return packed, sfa
+
+
+def dequantize_fp4_sfa_fp16(
+    packed: torch.Tensor,
+    sfa: torch.Tensor,
+    out: torch.Tensor | None = None,
+    is_sfb: bool = False,
+) -> 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, bool(is_sfb))
+    return out
+
+
+def nvfp4_gemm_bf16(
+    a_packed: torch.Tensor,
+    b_packed: torch.Tensor,
+    sfa: torch.Tensor,
+    sfb: torch.Tensor,
+    alpha: float = 1.0,
+    out: torch.Tensor | None = None,
+    variant: int = -1,
+) -> torch.Tensor:
+    if out is None:
+        out = torch.empty((a_packed.shape[0], b_packed.shape[0]), device=a_packed.device, dtype=torch.bfloat16)
+    ops.nvfp4_gemm_bf16(a_packed, b_packed, sfa, sfb, out, float(alpha), int(variant))
+    return out
+
+
+def fp4_w4a4_gemv_warpsplit_bf16(
+    a_packed: torch.Tensor,
+    b_packed: torch.Tensor,
+    sfa: torch.Tensor,
+    sfb: torch.Tensor,
+    *,
+    alpha: float = 1.0,
+    warps: int = 4,
+    stages: int = 4,
+    out: Optional[torch.Tensor] = None,
+) -> torch.Tensor:
+    """Warp-split-K NVFP4 W4A4 GEMV for the M=1 decode row (SM120).
+
+    Splits K across warps inside one block with a shared-memory reduce -
+    no cross-block intermediate, so it stays safe under CUDA-graph
+    replay - and fills the SMs the tiled GEMM underfills at long-K
+    small-M decode shapes. Same packed/scale layouts as the linear
+    entry points."""
+    if out is None:
+        out = torch.empty((1, b_packed.shape[0]), device=a_packed.device, dtype=torch.bfloat16)
+    ops.fp4_w4a4_gemv_warpsplit_bf16(a_packed, b_packed, sfa, sfb, out, float(alpha), int(warps), int(stages))
+    return out
+
+
+def fp4_w4a16_linear_bf16(
+    a_packed: torch.Tensor,
+    b_packed: torch.Tensor,
+    sfa: torch.Tensor,
+    sfb: torch.Tensor,
+    alpha: float = 1.0,
+    out: torch.Tensor | None = None,
+    variant: int = -1,
+) -> torch.Tensor:
+    """Compatibility alias for :func:`nvfp4_gemm_bf16`."""
+    return nvfp4_gemm_bf16(
+        a_packed, b_packed, sfa, sfb, alpha=alpha, out=out, variant=variant
+    )
+
+
+def nvfp4_gemm_residual_bf16(
+    a_packed: torch.Tensor,
+    b_packed: torch.Tensor,
+    sfa: torch.Tensor,
+    sfb: torch.Tensor,
+    residual: torch.Tensor,
+    alpha: float = 1.0,
+    out: torch.Tensor | None = None,
+) -> torch.Tensor:
+    if out is None:
+        out = torch.empty_like(residual)
+    ops.nvfp4_gemm_residual_bf16(
+        a_packed, b_packed, sfa, sfb, residual, out, float(alpha)
+    )
+    return out
+
+
+def nvfp4_gemm_bias_gelu_bf16(
+    a_packed: torch.Tensor,
+    b_packed: torch.Tensor,
+    sfa: torch.Tensor,
+    sfb: torch.Tensor,
+    bias: torch.Tensor,
+    alpha: float = 1.0,
+    out: torch.Tensor | None = None,
+) -> torch.Tensor:
+    if out is None:
+        out = torch.empty(
+            (a_packed.shape[0], b_packed.shape[0]),
+            device=a_packed.device,
+            dtype=torch.bfloat16,
+        )
+    ops.nvfp4_gemm_bias_gelu_bf16(
+        a_packed, b_packed, sfa, sfb, bias, out, float(alpha)
+    )
+    return out
+
+
+def nvfp4_gemm_bias_gelu_nvfp4(
+    a_packed: torch.Tensor,
+    b_packed: torch.Tensor,
+    sfa: torch.Tensor,
+    sfb: torch.Tensor,
+    bias: torch.Tensor,
+    alpha: float = 1.0,
+    out_packed: torch.Tensor | None = None,
+    out_sfa: torch.Tensor | None = None,
+) -> tuple[torch.Tensor, torch.Tensor]:
+    m, n = a_packed.shape[0], b_packed.shape[0]
+    if out_packed is None:
+        out_packed = torch.empty((m, n // 2), device=a_packed.device, dtype=torch.uint8)
+    if out_sfa is None:
+        out_sfa = torch.empty((sfa_size_bytes(m, n),), device=a_packed.device, dtype=torch.uint8)
+    ops.nvfp4_gemm_bias_gelu_nvfp4(
+        a_packed, b_packed, sfa, sfb, bias, out_packed, out_sfa, float(alpha)
+    )
+    return out_packed, out_sfa
+
+
+def nvfp4_gemm_streamk_bf16(
+    a_packed: torch.Tensor,
+    b_packed: torch.Tensor,
+    sfa: torch.Tensor,
+    sfb: torch.Tensor,
+    alpha: float = 1.0,
+    out: torch.Tensor | None = None,
+) -> torch.Tensor:
+    if out is None:
+        out = torch.empty(
+            (a_packed.shape[0], b_packed.shape[0]),
+            device=a_packed.device,
+            dtype=torch.bfloat16,
+        )
+    ops.nvfp4_gemm_streamk_bf16(
+        a_packed, b_packed, sfa, sfb, out, float(alpha)
+    )
+    return out
+
+
+def nvfp4_gemm_streamk_bias_bf16(
+    a_packed: torch.Tensor,
+    b_packed: torch.Tensor,
+    sfa: torch.Tensor,
+    sfb: torch.Tensor,
+    bias: torch.Tensor,
+    alpha: float = 1.0,
+    out: torch.Tensor | None = None,
+) -> torch.Tensor:
+    if out is None:
+        out = torch.empty(
+            (a_packed.shape[0], b_packed.shape[0]),
+            device=a_packed.device,
+            dtype=torch.bfloat16,
+        )
+    ops.nvfp4_gemm_streamk_bias_bf16(
+        a_packed, b_packed, sfa, sfb, bias, out, float(alpha)
+    )
+    return out
+
+
+__all__ = [
+    "dequantize_fp4_sfa_fp16",
+    "fp4_w4a16_linear_bf16",
+    "fp4_w4a4_gemv_warpsplit_bf16",
+    "nvfp4_gemm_bf16",
+    "nvfp4_gemm_bias_gelu_bf16",
+    "nvfp4_gemm_bias_gelu_nvfp4",
+    "nvfp4_gemm_residual_bf16",
+    "nvfp4_gemm_streamk_bf16",
+    "nvfp4_gemm_streamk_bias_bf16",
+    "quantize_fp4_sfa_fp16",
+    "quantize_fp4_sfa_bf16",
+    "sfa_size_bytes",
+]

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

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

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

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

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

@@ -0,0 +1,33 @@
+{
+  "name": "fp4-gemm",
+  "id": "_fp4_gemm_cuda_b46a817",
+  "version": 1,
+  "license": "Apache-2.0",
+  "python-depends": [],
+  "backend": {
+    "type": "cuda",
+    "archs": [
+      "11.0a",
+      "12.0a"
+    ]
+  },
+  "digest": {
+    "algorithm": "sha256",
+    "files": {
+      "__init__.py": "pJlGLDlOKnju+m59WyZJjpynMTvkPQTF6w+z4YfCwSY=",
+      "_fp4_gemm_cuda_b46a817.abi3.so": "W58aYR6tBPhbMcYu1GCwKRzpSOAozVUB0mCSJZtojZ4=",
+      "_ops.py": "tD2Fh7MGjwwYSATplu1PwYdWuhn0eHO42VKMA5pSjgU="
+    }
+  },
+  "provenance": {
+    "kernel-builder": {
+      "version": "0.17.0-dev0",
+      "sha": "d720fa90fb9cd92d1bc60a9dc5c55bef2aafabb8",
+      "dirty": false
+    },
+    "kernel": {
+      "sha": "b46a81771fdabe26f2c2494dc2ecf8492c88c45a",
+      "dirty": false
+    }
+  }
+}

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

@@ -0,0 +1,307 @@
+"""FlashRT FP4 GEMM kernels."""
+
+from __future__ import annotations
+
+import torch
+
+from ._ops import add_op_namespace_prefix, ops
+
+
+def sfa_size_bytes(rows: int, dim: int) -> int:
+    if rows <= 0 or dim <= 0 or dim % 16 != 0:
+        raise ValueError("rows must be positive and dim must be positive/divisible by 16")
+    n_blocks = dim // 16
+    n_row_super = (rows + 127) // 128
+    n_col_super = (n_blocks + 3) // 4
+    return n_row_super * n_col_super * 512
+
+
+def _alloc_fp4(rows: int, dim: int, device: torch.device | str):
+    return (
+        torch.empty((rows, dim // 2), device=device, dtype=torch.uint8),
+        torch.empty((sfa_size_bytes(rows, dim),), device=device, dtype=torch.uint8),
+    )
+
+
+@torch.library.register_fake(add_op_namespace_prefix("nvfp4_gemm_bf16"))
+def _linear_fake(
+    a_packed: torch.Tensor,
+    b_packed: torch.Tensor,
+    sfa: torch.Tensor,
+    sfb: torch.Tensor,
+    out: torch.Tensor,
+    alpha: float = 1.0,
+    variant: int = -1,
+) -> None:
+    return None
+
+
+@torch.library.register_fake(add_op_namespace_prefix("fp4_w4a4_gemv_warpsplit_bf16"))
+def _gemv_warpsplit_fake(a_packed, b_packed, sfa, sfb, out, alpha: float = 1.0, warps: int = 4, stages: int = 4) -> None:
+    if a_packed.shape[0] != 1:
+        raise RuntimeError("warp-split GEMV serves M=1 only")
+    if out.shape != (1, b_packed.shape[0]):
+        raise RuntimeError("out must have shape (1, N)")
+    return None
+
+
+@torch.library.register_fake(add_op_namespace_prefix("fp4_w4a16_linear_bf16"))
+def _legacy_linear_fake(
+    a_packed: torch.Tensor,
+    b_packed: torch.Tensor,
+    sfa: torch.Tensor,
+    sfb: torch.Tensor,
+    out: torch.Tensor,
+    alpha: float = 1.0,
+    variant: int = -1,
+) -> None:
+    return None
+
+
+@torch.library.register_fake(add_op_namespace_prefix("quantize_fp4_sfa_fp16"))
+def _quant_fake(x: torch.Tensor, packed: torch.Tensor, sfa: torch.Tensor, is_sfb: bool = False) -> None:
+    return None
+
+
+@torch.library.register_fake(add_op_namespace_prefix("quantize_fp4_sfa_bf16"))
+def _quant_bf16_fake(x: torch.Tensor, packed: torch.Tensor, sfa: torch.Tensor, is_sfb: bool = False) -> 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, is_sfb: bool = False) -> None:
+    return None
+
+
+@torch.library.register_fake(add_op_namespace_prefix("nvfp4_gemm_residual_bf16"))
+def _residual_fake(a, b, sfa, sfb, residual, out, alpha: float = 1.0) -> None:
+    return None
+
+
+@torch.library.register_fake(add_op_namespace_prefix("nvfp4_gemm_bias_gelu_bf16"))
+def _bias_gelu_fake(a, b, sfa, sfb, bias, out, alpha: float = 1.0) -> None:
+    return None
+
+
+@torch.library.register_fake(add_op_namespace_prefix("nvfp4_gemm_bias_gelu_nvfp4"))
+def _bias_gelu_nvfp4_fake(
+    a, b, sfa, sfb, bias, out_packed, out_sfa, alpha: float = 1.0
+) -> None:
+    return None
+
+
+@torch.library.register_fake(add_op_namespace_prefix("nvfp4_gemm_streamk_bf16"))
+def _streamk_fake(a, b, sfa, sfb, out, alpha: float = 1.0) -> None:
+    return None
+
+
+@torch.library.register_fake(add_op_namespace_prefix("nvfp4_gemm_streamk_bias_bf16"))
+def _streamk_bias_fake(a, b, sfa, sfb, bias, out, alpha: float = 1.0) -> None:
+    return None
+
+
+def quantize_fp4_sfa_fp16(
+    x: torch.Tensor,
+    packed: torch.Tensor | None = None,
+    sfa: torch.Tensor | None = None,
+    is_sfb: bool = False,
+):
+    if packed is None or sfa is None:
+        packed, sfa = _alloc_fp4(x.shape[0], x.shape[1], x.device)
+    ops.quantize_fp4_sfa_fp16(x, packed, sfa, bool(is_sfb))
+    return packed, sfa
+
+
+def quantize_fp4_sfa_bf16(
+    x: torch.Tensor,
+    packed: torch.Tensor | None = None,
+    sfa: torch.Tensor | None = None,
+    is_sfb: bool = False,
+):
+    """Quantize BF16 directly to packed E2M1 and CUTLASS SFA/SFB."""
+    if packed is None or sfa is None:
+        packed, sfa = _alloc_fp4(x.shape[0], x.shape[1], x.device)
+    ops.quantize_fp4_sfa_bf16(x, packed, sfa, bool(is_sfb))
+    return packed, sfa
+
+
+def dequantize_fp4_sfa_fp16(
+    packed: torch.Tensor,
+    sfa: torch.Tensor,
+    out: torch.Tensor | None = None,
+    is_sfb: bool = False,
+) -> 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, bool(is_sfb))
+    return out
+
+
+def nvfp4_gemm_bf16(
+    a_packed: torch.Tensor,
+    b_packed: torch.Tensor,
+    sfa: torch.Tensor,
+    sfb: torch.Tensor,
+    alpha: float = 1.0,
+    out: torch.Tensor | None = None,
+    variant: int = -1,
+) -> torch.Tensor:
+    if out is None:
+        out = torch.empty((a_packed.shape[0], b_packed.shape[0]), device=a_packed.device, dtype=torch.bfloat16)
+    ops.nvfp4_gemm_bf16(a_packed, b_packed, sfa, sfb, out, float(alpha), int(variant))
+    return out
+
+
+def fp4_w4a4_gemv_warpsplit_bf16(
+    a_packed: torch.Tensor,
+    b_packed: torch.Tensor,
+    sfa: torch.Tensor,
+    sfb: torch.Tensor,
+    *,
+    alpha: float = 1.0,
+    warps: int = 4,
+    stages: int = 4,
+    out: Optional[torch.Tensor] = None,
+) -> torch.Tensor:
+    """Warp-split-K NVFP4 W4A4 GEMV for the M=1 decode row (SM120).
+
+    Splits K across warps inside one block with a shared-memory reduce -
+    no cross-block intermediate, so it stays safe under CUDA-graph
+    replay - and fills the SMs the tiled GEMM underfills at long-K
+    small-M decode shapes. Same packed/scale layouts as the linear
+    entry points."""
+    if out is None:
+        out = torch.empty((1, b_packed.shape[0]), device=a_packed.device, dtype=torch.bfloat16)
+    ops.fp4_w4a4_gemv_warpsplit_bf16(a_packed, b_packed, sfa, sfb, out, float(alpha), int(warps), int(stages))
+    return out
+
+
+def fp4_w4a16_linear_bf16(
+    a_packed: torch.Tensor,
+    b_packed: torch.Tensor,
+    sfa: torch.Tensor,
+    sfb: torch.Tensor,
+    alpha: float = 1.0,
+    out: torch.Tensor | None = None,
+    variant: int = -1,
+) -> torch.Tensor:
+    """Compatibility alias for :func:`nvfp4_gemm_bf16`."""
+    return nvfp4_gemm_bf16(
+        a_packed, b_packed, sfa, sfb, alpha=alpha, out=out, variant=variant
+    )
+
+
+def nvfp4_gemm_residual_bf16(
+    a_packed: torch.Tensor,
+    b_packed: torch.Tensor,
+    sfa: torch.Tensor,
+    sfb: torch.Tensor,
+    residual: torch.Tensor,
+    alpha: float = 1.0,
+    out: torch.Tensor | None = None,
+) -> torch.Tensor:
+    if out is None:
+        out = torch.empty_like(residual)
+    ops.nvfp4_gemm_residual_bf16(
+        a_packed, b_packed, sfa, sfb, residual, out, float(alpha)
+    )
+    return out
+
+
+def nvfp4_gemm_bias_gelu_bf16(
+    a_packed: torch.Tensor,
+    b_packed: torch.Tensor,
+    sfa: torch.Tensor,
+    sfb: torch.Tensor,
+    bias: torch.Tensor,
+    alpha: float = 1.0,
+    out: torch.Tensor | None = None,
+) -> torch.Tensor:
+    if out is None:
+        out = torch.empty(
+            (a_packed.shape[0], b_packed.shape[0]),
+            device=a_packed.device,
+            dtype=torch.bfloat16,
+        )
+    ops.nvfp4_gemm_bias_gelu_bf16(
+        a_packed, b_packed, sfa, sfb, bias, out, float(alpha)
+    )
+    return out
+
+
+def nvfp4_gemm_bias_gelu_nvfp4(
+    a_packed: torch.Tensor,
+    b_packed: torch.Tensor,
+    sfa: torch.Tensor,
+    sfb: torch.Tensor,
+    bias: torch.Tensor,
+    alpha: float = 1.0,
+    out_packed: torch.Tensor | None = None,
+    out_sfa: torch.Tensor | None = None,
+) -> tuple[torch.Tensor, torch.Tensor]:
+    m, n = a_packed.shape[0], b_packed.shape[0]
+    if out_packed is None:
+        out_packed = torch.empty((m, n // 2), device=a_packed.device, dtype=torch.uint8)
+    if out_sfa is None:
+        out_sfa = torch.empty((sfa_size_bytes(m, n),), device=a_packed.device, dtype=torch.uint8)
+    ops.nvfp4_gemm_bias_gelu_nvfp4(
+        a_packed, b_packed, sfa, sfb, bias, out_packed, out_sfa, float(alpha)
+    )
+    return out_packed, out_sfa
+
+
+def nvfp4_gemm_streamk_bf16(
+    a_packed: torch.Tensor,
+    b_packed: torch.Tensor,
+    sfa: torch.Tensor,
+    sfb: torch.Tensor,
+    alpha: float = 1.0,
+    out: torch.Tensor | None = None,
+) -> torch.Tensor:
+    if out is None:
+        out = torch.empty(
+            (a_packed.shape[0], b_packed.shape[0]),
+            device=a_packed.device,
+            dtype=torch.bfloat16,
+        )
+    ops.nvfp4_gemm_streamk_bf16(
+        a_packed, b_packed, sfa, sfb, out, float(alpha)
+    )
+    return out
+
+
+def nvfp4_gemm_streamk_bias_bf16(
+    a_packed: torch.Tensor,
+    b_packed: torch.Tensor,
+    sfa: torch.Tensor,
+    sfb: torch.Tensor,
+    bias: torch.Tensor,
+    alpha: float = 1.0,
+    out: torch.Tensor | None = None,
+) -> torch.Tensor:
+    if out is None:
+        out = torch.empty(
+            (a_packed.shape[0], b_packed.shape[0]),
+            device=a_packed.device,
+            dtype=torch.bfloat16,
+        )
+    ops.nvfp4_gemm_streamk_bias_bf16(
+        a_packed, b_packed, sfa, sfb, bias, out, float(alpha)
+    )
+    return out
+
+
+__all__ = [
+    "dequantize_fp4_sfa_fp16",
+    "fp4_w4a16_linear_bf16",
+    "fp4_w4a4_gemv_warpsplit_bf16",
+    "nvfp4_gemm_bf16",
+    "nvfp4_gemm_bias_gelu_bf16",
+    "nvfp4_gemm_bias_gelu_nvfp4",
+    "nvfp4_gemm_residual_bf16",
+    "nvfp4_gemm_streamk_bf16",
+    "nvfp4_gemm_streamk_bias_bf16",
+    "quantize_fp4_sfa_fp16",
+    "quantize_fp4_sfa_bf16",
+    "sfa_size_bytes",
+]

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

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

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

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

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

@@ -0,0 +1,33 @@
+{
+  "name": "fp4-gemm",
+  "id": "_fp4_gemm_cuda_b46a817",
+  "version": 1,
+  "license": "Apache-2.0",
+  "python-depends": [],
+  "backend": {
+    "type": "cuda",
+    "archs": [
+      "11.0a",
+      "12.0a"
+    ]
+  },
+  "digest": {
+    "algorithm": "sha256",
+    "files": {
+      "__init__.py": "pJlGLDlOKnju+m59WyZJjpynMTvkPQTF6w+z4YfCwSY=",
+      "_fp4_gemm_cuda_b46a817.abi3.so": "UIK1OQaZPeFDr00zSZWrmdk2ygbGsBTOnVavHvEBuNQ=",
+      "_ops.py": "tD2Fh7MGjwwYSATplu1PwYdWuhn0eHO42VKMA5pSjgU="
+    }
+  },
+  "provenance": {
+    "kernel-builder": {
+      "version": "0.17.0-dev0",
+      "sha": "d720fa90fb9cd92d1bc60a9dc5c55bef2aafabb8",
+      "dirty": false
+    },
+    "kernel": {
+      "sha": "b46a81771fdabe26f2c2494dc2ecf8492c88c45a",
+      "dirty": false
+    }
+  }
+}