Uploaded using `kernel-builder`.

Browse files

Files changed (16) hide show

build/torch212-cxx11-cu130-x86_64-linux/__init__.py +118 -0
build/torch212-cxx11-cu130-x86_64-linux/{_fp4_gemm_cuda_b6588aa.abi3.so → _fp4_gemm_cuda_92c4b99.abi3.so} +2 -2
build/torch212-cxx11-cu130-x86_64-linux/_ops.py +3 -3
build/torch212-cxx11-cu130-x86_64-linux/metadata.json +5 -5
build/torch212-cxx11-cu132-x86_64-linux/__init__.py +118 -0
build/torch212-cxx11-cu132-x86_64-linux/{_fp4_gemm_cuda_b6588aa.abi3.so → _fp4_gemm_cuda_92c4b99.abi3.so} +2 -2
build/torch212-cxx11-cu132-x86_64-linux/_ops.py +3 -3
build/torch212-cxx11-cu132-x86_64-linux/metadata.json +5 -5
build/torch213-cxx11-cu130-x86_64-linux/__init__.py +118 -0
build/torch213-cxx11-cu130-x86_64-linux/{_fp4_gemm_cuda_b6588aa.abi3.so → _fp4_gemm_cuda_92c4b99.abi3.so} +2 -2
build/torch213-cxx11-cu130-x86_64-linux/_ops.py +3 -3
build/torch213-cxx11-cu130-x86_64-linux/metadata.json +5 -5
build/torch213-cxx11-cu132-x86_64-linux/__init__.py +118 -0
build/torch213-cxx11-cu132-x86_64-linux/{_fp4_gemm_cuda_b6588aa.abi3.so → _fp4_gemm_cuda_92c4b99.abi3.so} +2 -2
build/torch213-cxx11-cu132-x86_64-linux/_ops.py +3 -3
build/torch213-cxx11-cu132-x86_64-linux/metadata.json +5 -5

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

@@ -109,6 +109,11 @@ def _quant_mse_fake(x, packed, sfa, is_sfb: bool = False) -> None:
     return None
 @torch.library.register_fake(add_op_namespace_prefix("quantize_e0m3_sfa_fp16"))
 def _quant_e0m3_fake(x, packed, sfa, is_sfb: bool = False) -> None:
     return None
@@ -186,6 +191,91 @@ def quantize_fp4_sfa_mse_fp16(
     return packed, sfa
 def quantize_e0m3_sfa_fp16(
     x: torch.Tensor,
     packed: torch.Tensor | None = None,
@@ -381,6 +471,29 @@ def nvfp4_gemm_geglu_nvfp4_fp16(
     return out_packed, out_sfa
 def nvfp4_gemm_bias_gelu_nvfp4_fp16(
     a_packed: torch.Tensor,
     b_packed: torch.Tensor,
@@ -601,6 +714,8 @@ def nvfp4_gemm_streamk_bias_bf16(
 __all__ = [
     "dequantize_fp4_sfa_fp16",
     "e0m3_weight_gemm_fp16",
     "fp4_w4a16_linear_bf16",
@@ -622,6 +737,9 @@ __all__ = [
     "nvfp4_gemm_streamk_bias_bf16",
     "quantize_fp4_sfa_fp16",
     "quantize_fp4_sfa_mse_fp16",
     "quantize_e0m3_sfa_fp16",
     "quantize_fp4_sfa_bf16",
     "sfa_size_bytes",

     return None
+@torch.library.register_fake(add_op_namespace_prefix("quantize_fp4_sfa_mse_bf16"))
+def _quant_mse_bf16_fake(x, packed, sfa, is_sfb: bool = False) -> None:
+    return None
 @torch.library.register_fake(add_op_namespace_prefix("quantize_e0m3_sfa_fp16"))
 def _quant_e0m3_fake(x, packed, sfa, is_sfb: bool = False) -> None:
     return None
     return packed, sfa
+def quantize_fp4_sfa_mse_bf16(
+    x: torch.Tensor,
+    packed: torch.Tensor | None = None,
+    sfa: torch.Tensor | None = None,
+    is_sfb: bool = False,
+):
+    """Pack BF16 weights using per-block reconstruction-MSE scale search."""
+    if packed is None or sfa is None:
+        packed, sfa = _alloc_fp4(x.shape[0], x.shape[1], x.device)
+    ops.quantize_fp4_sfa_mse_bf16(x, packed, sfa, bool(is_sfb))
+    return packed, sfa
+def aligned_fp4_dim(dim: int, alignment: int = 32) -> int:
+    """Return the physical FP4 GEMM dimension for a logical dimension."""
+    if dim <= 0 or alignment <= 0 or alignment % 16 != 0:
+        raise ValueError("dim must be positive and alignment divisible by 16")
+    return ((dim + alignment - 1) // alignment) * alignment
+def quantize_fp4_sfa_padded_bf16(
+    x: torch.Tensor,
+    *,
+    alignment: int = 32,
+    is_sfb: bool = False,
+):
+    """Bind-time BF16 pack with zero padding on the final dimension.
+    The returned packed tensor uses the physical aligned width. Keep the
+    logical width returned as the third item for slicing final model outputs.
+    This helper allocates and is intentionally not a runtime hot-path op.
+    """
+    if x.ndim != 2 or x.dtype != torch.bfloat16:
+        raise ValueError("x must be a 2-D BF16 tensor")
+    logical_dim = int(x.shape[1])
+    physical_dim = aligned_fp4_dim(logical_dim, alignment)
+    if physical_dim == logical_dim:
+        padded = x.contiguous()
+    else:
+        padded = torch.zeros(
+            (x.shape[0], physical_dim), device=x.device, dtype=x.dtype
+        )
+        padded[:, :logical_dim].copy_(x)
+    packed, sfa = quantize_fp4_sfa_bf16(padded, is_sfb=is_sfb)
+    return packed, sfa, logical_dim
+def pack_nvfp4_weight_bf16(
+    weight: torch.Tensor,
+    bias: torch.Tensor | None = None,
+    *,
+    alignment: int = 32,
+    mse: bool = False,
+):
+    """Zero-pad and pack a logical ``(N, K)`` BF16 weight at bind time.
+    Both physical ``N`` and ``K`` are aligned. This is the model-neutral
+    contract used by widths such as SigLIP's logical 4304, which becomes 4320
+    for Blackwell NVFP4 TMA. The hot path consumes the returned static tensors
+    directly and performs no padding or allocation.
+    """
+    if weight.ndim != 2 or weight.dtype != torch.bfloat16:
+        raise ValueError("weight must be a 2-D BF16 tensor")
+    logical_n, logical_k = map(int, weight.shape)
+    physical_n = aligned_fp4_dim(logical_n, alignment)
+    physical_k = aligned_fp4_dim(logical_k, alignment)
+    padded_weight = torch.zeros(
+        (physical_n, physical_k), device=weight.device, dtype=weight.dtype
+    )
+    padded_weight[:logical_n, :logical_k].copy_(weight)
+    quantize = quantize_fp4_sfa_mse_bf16 if mse else quantize_fp4_sfa_bf16
+    packed, sfb = quantize(padded_weight, is_sfb=True)
+    padded_bias = None
+    if bias is not None:
+        if bias.ndim != 1 or bias.numel() != logical_n:
+            raise ValueError("bias must have shape (N,)")
+        if bias.dtype != torch.bfloat16 or bias.device != weight.device:
+            raise ValueError("bias must match the weight dtype and device")
+        padded_bias = torch.zeros(
+            (physical_n,), device=bias.device, dtype=bias.dtype
+        )
+        padded_bias[:logical_n].copy_(bias)
+    return packed, sfb, padded_bias, (logical_n, logical_k)
 def quantize_e0m3_sfa_fp16(
     x: torch.Tensor,
     packed: torch.Tensor | None = None,
     return out_packed, out_sfa
+def cutlass_fp4_gemm_geglu_il_hw_v10(
+    a_packed: torch.Tensor,
+    b_interleaved_packed: torch.Tensor,
+    sfa: torch.Tensor,
+    sfb: torch.Tensor,
+    *,
+    scratch: torch.Tensor | None = None,
+    out_packed: torch.Tensor | None = None,
+    out_sfa: torch.Tensor | None = None,
+) -> tuple[torch.Tensor, torch.Tensor]:
+    """Explicit public alias for the native SM110 v10 GeGLU schedule."""
+    return nvfp4_gemm_geglu_nvfp4_fp16(
+        a_packed,
+        b_interleaved_packed,
+        sfa,
+        sfb,
+        skinny=True,
+        scratch=scratch,
+        out_packed=out_packed,
+        out_sfa=out_sfa,
+    )
 def nvfp4_gemm_bias_gelu_nvfp4_fp16(
     a_packed: torch.Tensor,
     b_packed: torch.Tensor,
 __all__ = [
+    "aligned_fp4_dim",
+    "cutlass_fp4_gemm_geglu_il_hw_v10",
     "dequantize_fp4_sfa_fp16",
     "e0m3_weight_gemm_fp16",
     "fp4_w4a16_linear_bf16",
     "nvfp4_gemm_streamk_bias_bf16",
     "quantize_fp4_sfa_fp16",
     "quantize_fp4_sfa_mse_fp16",
+    "quantize_fp4_sfa_mse_bf16",
+    "quantize_fp4_sfa_padded_bf16",
+    "pack_nvfp4_weight_bf16",
     "quantize_e0m3_sfa_fp16",
     "quantize_fp4_sfa_bf16",
     "sfa_size_bytes",

build/torch212-cxx11-cu130-x86_64-linux/{_fp4_gemm_cuda_b6588aa.abi3.so → _fp4_gemm_cuda_92c4b99.abi3.so} RENAMED Viewed

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

 version https://git-lfs.github.com/spec/v1
+oid sha256:fe2db33c4dc640281bcfccce7ae0c98a53a42efeef388079b65149bb7e3031a5
+size 7895896

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

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

 import torch
+from . import _fp4_gemm_cuda_92c4b99
+ops = torch.ops._fp4_gemm_cuda_92c4b99
 def add_op_namespace_prefix(op_name: str):
     """
     Prefix op by namespace.
     """
+    return f"_fp4_gemm_cuda_92c4b99::{op_name}"

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

@@ -1,6 +1,6 @@
 {
   "name": "fp4-gemm",
-  "id": "_fp4_gemm_cuda_b6588aa",
   "version": 1,
   "license": "Apache-2.0",
   "python-depends": [],
@@ -14,9 +14,9 @@
   "digest": {
     "algorithm": "sha256",
     "files": {
-      "__init__.py": "Fb/FDYOpDw+YGlevblUsv4AkILDQCJdB807qdDUDI1c=",
-      "_fp4_gemm_cuda_b6588aa.abi3.so": "LZvlS4WJ5RMe+lpP+d3SaZUZ5BqlYD1UtbCoQQx8Kgc=",
-      "_ops.py": "KKClchvuzBz2g1A148Gi2R6QIx0Ipj0CKlTT511GrUY="
     }
   },
   "provenance": {
@@ -26,7 +26,7 @@
       "dirty": false
     },
     "kernel": {
-      "sha": "b6588aa59dd4f8c646f90b58312a0a66886e061f",
       "dirty": false
     }
   }

 {
   "name": "fp4-gemm",
+  "id": "_fp4_gemm_cuda_92c4b99",
   "version": 1,
   "license": "Apache-2.0",
   "python-depends": [],
   "digest": {
     "algorithm": "sha256",
     "files": {
+      "__init__.py": "Mr1/bAUzs+af68D+bRXX6iONDof//dzKgn6QU5pKdu8=",
+      "_fp4_gemm_cuda_92c4b99.abi3.so": "/i2zPE3GQCgbz8zOeuDJilOkLv7vOIB5tlFJu34wMaU=",
+      "_ops.py": "c5GHeHUOPDgpedue7MgdtRjCwHOmUsBo4OfqoFVegq0="
     }
   },
   "provenance": {
       "dirty": false
     },
     "kernel": {
+      "sha": "92c4b991b6f1fbd89744aa6f46f62b1bfac7ca7e",
       "dirty": false
     }
   }

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

@@ -109,6 +109,11 @@ def _quant_mse_fake(x, packed, sfa, is_sfb: bool = False) -> None:
     return None
 @torch.library.register_fake(add_op_namespace_prefix("quantize_e0m3_sfa_fp16"))
 def _quant_e0m3_fake(x, packed, sfa, is_sfb: bool = False) -> None:
     return None
@@ -186,6 +191,91 @@ def quantize_fp4_sfa_mse_fp16(
     return packed, sfa
 def quantize_e0m3_sfa_fp16(
     x: torch.Tensor,
     packed: torch.Tensor | None = None,
@@ -381,6 +471,29 @@ def nvfp4_gemm_geglu_nvfp4_fp16(
     return out_packed, out_sfa
 def nvfp4_gemm_bias_gelu_nvfp4_fp16(
     a_packed: torch.Tensor,
     b_packed: torch.Tensor,
@@ -601,6 +714,8 @@ def nvfp4_gemm_streamk_bias_bf16(
 __all__ = [
     "dequantize_fp4_sfa_fp16",
     "e0m3_weight_gemm_fp16",
     "fp4_w4a16_linear_bf16",
@@ -622,6 +737,9 @@ __all__ = [
     "nvfp4_gemm_streamk_bias_bf16",
     "quantize_fp4_sfa_fp16",
     "quantize_fp4_sfa_mse_fp16",
     "quantize_e0m3_sfa_fp16",
     "quantize_fp4_sfa_bf16",
     "sfa_size_bytes",

     return None
+@torch.library.register_fake(add_op_namespace_prefix("quantize_fp4_sfa_mse_bf16"))
+def _quant_mse_bf16_fake(x, packed, sfa, is_sfb: bool = False) -> None:
+    return None
 @torch.library.register_fake(add_op_namespace_prefix("quantize_e0m3_sfa_fp16"))
 def _quant_e0m3_fake(x, packed, sfa, is_sfb: bool = False) -> None:
     return None
     return packed, sfa
+def quantize_fp4_sfa_mse_bf16(
+    x: torch.Tensor,
+    packed: torch.Tensor | None = None,
+    sfa: torch.Tensor | None = None,
+    is_sfb: bool = False,
+):
+    """Pack BF16 weights using per-block reconstruction-MSE scale search."""
+    if packed is None or sfa is None:
+        packed, sfa = _alloc_fp4(x.shape[0], x.shape[1], x.device)
+    ops.quantize_fp4_sfa_mse_bf16(x, packed, sfa, bool(is_sfb))
+    return packed, sfa
+def aligned_fp4_dim(dim: int, alignment: int = 32) -> int:
+    """Return the physical FP4 GEMM dimension for a logical dimension."""
+    if dim <= 0 or alignment <= 0 or alignment % 16 != 0:
+        raise ValueError("dim must be positive and alignment divisible by 16")
+    return ((dim + alignment - 1) // alignment) * alignment
+def quantize_fp4_sfa_padded_bf16(
+    x: torch.Tensor,
+    *,
+    alignment: int = 32,
+    is_sfb: bool = False,
+):
+    """Bind-time BF16 pack with zero padding on the final dimension.
+    The returned packed tensor uses the physical aligned width. Keep the
+    logical width returned as the third item for slicing final model outputs.
+    This helper allocates and is intentionally not a runtime hot-path op.
+    """
+    if x.ndim != 2 or x.dtype != torch.bfloat16:
+        raise ValueError("x must be a 2-D BF16 tensor")
+    logical_dim = int(x.shape[1])
+    physical_dim = aligned_fp4_dim(logical_dim, alignment)
+    if physical_dim == logical_dim:
+        padded = x.contiguous()
+    else:
+        padded = torch.zeros(
+            (x.shape[0], physical_dim), device=x.device, dtype=x.dtype
+        )
+        padded[:, :logical_dim].copy_(x)
+    packed, sfa = quantize_fp4_sfa_bf16(padded, is_sfb=is_sfb)
+    return packed, sfa, logical_dim
+def pack_nvfp4_weight_bf16(
+    weight: torch.Tensor,
+    bias: torch.Tensor | None = None,
+    *,
+    alignment: int = 32,
+    mse: bool = False,
+):
+    """Zero-pad and pack a logical ``(N, K)`` BF16 weight at bind time.
+    Both physical ``N`` and ``K`` are aligned. This is the model-neutral
+    contract used by widths such as SigLIP's logical 4304, which becomes 4320
+    for Blackwell NVFP4 TMA. The hot path consumes the returned static tensors
+    directly and performs no padding or allocation.
+    """
+    if weight.ndim != 2 or weight.dtype != torch.bfloat16:
+        raise ValueError("weight must be a 2-D BF16 tensor")
+    logical_n, logical_k = map(int, weight.shape)
+    physical_n = aligned_fp4_dim(logical_n, alignment)
+    physical_k = aligned_fp4_dim(logical_k, alignment)
+    padded_weight = torch.zeros(
+        (physical_n, physical_k), device=weight.device, dtype=weight.dtype
+    )
+    padded_weight[:logical_n, :logical_k].copy_(weight)
+    quantize = quantize_fp4_sfa_mse_bf16 if mse else quantize_fp4_sfa_bf16
+    packed, sfb = quantize(padded_weight, is_sfb=True)
+    padded_bias = None
+    if bias is not None:
+        if bias.ndim != 1 or bias.numel() != logical_n:
+            raise ValueError("bias must have shape (N,)")
+        if bias.dtype != torch.bfloat16 or bias.device != weight.device:
+            raise ValueError("bias must match the weight dtype and device")
+        padded_bias = torch.zeros(
+            (physical_n,), device=bias.device, dtype=bias.dtype
+        )
+        padded_bias[:logical_n].copy_(bias)
+    return packed, sfb, padded_bias, (logical_n, logical_k)
 def quantize_e0m3_sfa_fp16(
     x: torch.Tensor,
     packed: torch.Tensor | None = None,
     return out_packed, out_sfa
+def cutlass_fp4_gemm_geglu_il_hw_v10(
+    a_packed: torch.Tensor,
+    b_interleaved_packed: torch.Tensor,
+    sfa: torch.Tensor,
+    sfb: torch.Tensor,
+    *,
+    scratch: torch.Tensor | None = None,
+    out_packed: torch.Tensor | None = None,
+    out_sfa: torch.Tensor | None = None,
+) -> tuple[torch.Tensor, torch.Tensor]:
+    """Explicit public alias for the native SM110 v10 GeGLU schedule."""
+    return nvfp4_gemm_geglu_nvfp4_fp16(
+        a_packed,
+        b_interleaved_packed,
+        sfa,
+        sfb,
+        skinny=True,
+        scratch=scratch,
+        out_packed=out_packed,
+        out_sfa=out_sfa,
+    )
 def nvfp4_gemm_bias_gelu_nvfp4_fp16(
     a_packed: torch.Tensor,
     b_packed: torch.Tensor,
 __all__ = [
+    "aligned_fp4_dim",
+    "cutlass_fp4_gemm_geglu_il_hw_v10",
     "dequantize_fp4_sfa_fp16",
     "e0m3_weight_gemm_fp16",
     "fp4_w4a16_linear_bf16",
     "nvfp4_gemm_streamk_bias_bf16",
     "quantize_fp4_sfa_fp16",
     "quantize_fp4_sfa_mse_fp16",
+    "quantize_fp4_sfa_mse_bf16",
+    "quantize_fp4_sfa_padded_bf16",
+    "pack_nvfp4_weight_bf16",
     "quantize_e0m3_sfa_fp16",
     "quantize_fp4_sfa_bf16",
     "sfa_size_bytes",

build/torch212-cxx11-cu132-x86_64-linux/{_fp4_gemm_cuda_b6588aa.abi3.so → _fp4_gemm_cuda_92c4b99.abi3.so} RENAMED Viewed

@@ -1,3 +1,3 @@
 version https://git-lfs.github.com/spec/v1
-oid sha256:308bf04520e36c5f5cbdcb45644cc86b01449729e246a10f87cc835b2935be25
-size 7706544

 version https://git-lfs.github.com/spec/v1
+oid sha256:68ba4bc31b2467a824e3e0a2daeb3ddacac43f4e16b5ec0779b90c283194f0cb
+size 7887648

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

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

 import torch
+from . import _fp4_gemm_cuda_92c4b99
+ops = torch.ops._fp4_gemm_cuda_92c4b99
 def add_op_namespace_prefix(op_name: str):
     """
     Prefix op by namespace.
     """
+    return f"_fp4_gemm_cuda_92c4b99::{op_name}"

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

@@ -1,6 +1,6 @@
 {
   "name": "fp4-gemm",
-  "id": "_fp4_gemm_cuda_b6588aa",
   "version": 1,
   "license": "Apache-2.0",
   "python-depends": [],
@@ -14,9 +14,9 @@
   "digest": {
     "algorithm": "sha256",
     "files": {
-      "__init__.py": "Fb/FDYOpDw+YGlevblUsv4AkILDQCJdB807qdDUDI1c=",
-      "_fp4_gemm_cuda_b6588aa.abi3.so": "MIvwRSDjbF9cvctFZEzIawFElyniRqEPh8yDWyk1viU=",
-      "_ops.py": "KKClchvuzBz2g1A148Gi2R6QIx0Ipj0CKlTT511GrUY="
     }
   },
   "provenance": {
@@ -26,7 +26,7 @@
       "dirty": false
     },
     "kernel": {
-      "sha": "b6588aa59dd4f8c646f90b58312a0a66886e061f",
       "dirty": false
     }
   }

 {
   "name": "fp4-gemm",
+  "id": "_fp4_gemm_cuda_92c4b99",
   "version": 1,
   "license": "Apache-2.0",
   "python-depends": [],
   "digest": {
     "algorithm": "sha256",
     "files": {
+      "__init__.py": "Mr1/bAUzs+af68D+bRXX6iONDof//dzKgn6QU5pKdu8=",
+      "_fp4_gemm_cuda_92c4b99.abi3.so": "aLpLwxskZ6gk4+Ci2us92srEP04WtewHebkMKDGU8Ms=",
+      "_ops.py": "c5GHeHUOPDgpedue7MgdtRjCwHOmUsBo4OfqoFVegq0="
     }
   },
   "provenance": {
       "dirty": false
     },
     "kernel": {
+      "sha": "92c4b991b6f1fbd89744aa6f46f62b1bfac7ca7e",
       "dirty": false
     }
   }

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

@@ -109,6 +109,11 @@ def _quant_mse_fake(x, packed, sfa, is_sfb: bool = False) -> None:
     return None
 @torch.library.register_fake(add_op_namespace_prefix("quantize_e0m3_sfa_fp16"))
 def _quant_e0m3_fake(x, packed, sfa, is_sfb: bool = False) -> None:
     return None
@@ -186,6 +191,91 @@ def quantize_fp4_sfa_mse_fp16(
     return packed, sfa
 def quantize_e0m3_sfa_fp16(
     x: torch.Tensor,
     packed: torch.Tensor | None = None,
@@ -381,6 +471,29 @@ def nvfp4_gemm_geglu_nvfp4_fp16(
     return out_packed, out_sfa
 def nvfp4_gemm_bias_gelu_nvfp4_fp16(
     a_packed: torch.Tensor,
     b_packed: torch.Tensor,
@@ -601,6 +714,8 @@ def nvfp4_gemm_streamk_bias_bf16(
 __all__ = [
     "dequantize_fp4_sfa_fp16",
     "e0m3_weight_gemm_fp16",
     "fp4_w4a16_linear_bf16",
@@ -622,6 +737,9 @@ __all__ = [
     "nvfp4_gemm_streamk_bias_bf16",
     "quantize_fp4_sfa_fp16",
     "quantize_fp4_sfa_mse_fp16",
     "quantize_e0m3_sfa_fp16",
     "quantize_fp4_sfa_bf16",
     "sfa_size_bytes",

     return None
+@torch.library.register_fake(add_op_namespace_prefix("quantize_fp4_sfa_mse_bf16"))
+def _quant_mse_bf16_fake(x, packed, sfa, is_sfb: bool = False) -> None:
+    return None
 @torch.library.register_fake(add_op_namespace_prefix("quantize_e0m3_sfa_fp16"))
 def _quant_e0m3_fake(x, packed, sfa, is_sfb: bool = False) -> None:
     return None
     return packed, sfa
+def quantize_fp4_sfa_mse_bf16(
+    x: torch.Tensor,
+    packed: torch.Tensor | None = None,
+    sfa: torch.Tensor | None = None,
+    is_sfb: bool = False,
+):
+    """Pack BF16 weights using per-block reconstruction-MSE scale search."""
+    if packed is None or sfa is None:
+        packed, sfa = _alloc_fp4(x.shape[0], x.shape[1], x.device)
+    ops.quantize_fp4_sfa_mse_bf16(x, packed, sfa, bool(is_sfb))
+    return packed, sfa
+def aligned_fp4_dim(dim: int, alignment: int = 32) -> int:
+    """Return the physical FP4 GEMM dimension for a logical dimension."""
+    if dim <= 0 or alignment <= 0 or alignment % 16 != 0:
+        raise ValueError("dim must be positive and alignment divisible by 16")
+    return ((dim + alignment - 1) // alignment) * alignment
+def quantize_fp4_sfa_padded_bf16(
+    x: torch.Tensor,
+    *,
+    alignment: int = 32,
+    is_sfb: bool = False,
+):
+    """Bind-time BF16 pack with zero padding on the final dimension.
+    The returned packed tensor uses the physical aligned width. Keep the
+    logical width returned as the third item for slicing final model outputs.
+    This helper allocates and is intentionally not a runtime hot-path op.
+    """
+    if x.ndim != 2 or x.dtype != torch.bfloat16:
+        raise ValueError("x must be a 2-D BF16 tensor")
+    logical_dim = int(x.shape[1])
+    physical_dim = aligned_fp4_dim(logical_dim, alignment)
+    if physical_dim == logical_dim:
+        padded = x.contiguous()
+    else:
+        padded = torch.zeros(
+            (x.shape[0], physical_dim), device=x.device, dtype=x.dtype
+        )
+        padded[:, :logical_dim].copy_(x)
+    packed, sfa = quantize_fp4_sfa_bf16(padded, is_sfb=is_sfb)
+    return packed, sfa, logical_dim
+def pack_nvfp4_weight_bf16(
+    weight: torch.Tensor,
+    bias: torch.Tensor | None = None,
+    *,
+    alignment: int = 32,
+    mse: bool = False,
+):
+    """Zero-pad and pack a logical ``(N, K)`` BF16 weight at bind time.
+    Both physical ``N`` and ``K`` are aligned. This is the model-neutral
+    contract used by widths such as SigLIP's logical 4304, which becomes 4320
+    for Blackwell NVFP4 TMA. The hot path consumes the returned static tensors
+    directly and performs no padding or allocation.
+    """
+    if weight.ndim != 2 or weight.dtype != torch.bfloat16:
+        raise ValueError("weight must be a 2-D BF16 tensor")
+    logical_n, logical_k = map(int, weight.shape)
+    physical_n = aligned_fp4_dim(logical_n, alignment)
+    physical_k = aligned_fp4_dim(logical_k, alignment)
+    padded_weight = torch.zeros(
+        (physical_n, physical_k), device=weight.device, dtype=weight.dtype
+    )
+    padded_weight[:logical_n, :logical_k].copy_(weight)
+    quantize = quantize_fp4_sfa_mse_bf16 if mse else quantize_fp4_sfa_bf16
+    packed, sfb = quantize(padded_weight, is_sfb=True)
+    padded_bias = None
+    if bias is not None:
+        if bias.ndim != 1 or bias.numel() != logical_n:
+            raise ValueError("bias must have shape (N,)")
+        if bias.dtype != torch.bfloat16 or bias.device != weight.device:
+            raise ValueError("bias must match the weight dtype and device")
+        padded_bias = torch.zeros(
+            (physical_n,), device=bias.device, dtype=bias.dtype
+        )
+        padded_bias[:logical_n].copy_(bias)
+    return packed, sfb, padded_bias, (logical_n, logical_k)
 def quantize_e0m3_sfa_fp16(
     x: torch.Tensor,
     packed: torch.Tensor | None = None,
     return out_packed, out_sfa
+def cutlass_fp4_gemm_geglu_il_hw_v10(
+    a_packed: torch.Tensor,
+    b_interleaved_packed: torch.Tensor,
+    sfa: torch.Tensor,
+    sfb: torch.Tensor,
+    *,
+    scratch: torch.Tensor | None = None,
+    out_packed: torch.Tensor | None = None,
+    out_sfa: torch.Tensor | None = None,
+) -> tuple[torch.Tensor, torch.Tensor]:
+    """Explicit public alias for the native SM110 v10 GeGLU schedule."""
+    return nvfp4_gemm_geglu_nvfp4_fp16(
+        a_packed,
+        b_interleaved_packed,
+        sfa,
+        sfb,
+        skinny=True,
+        scratch=scratch,
+        out_packed=out_packed,
+        out_sfa=out_sfa,
+    )
 def nvfp4_gemm_bias_gelu_nvfp4_fp16(
     a_packed: torch.Tensor,
     b_packed: torch.Tensor,
 __all__ = [
+    "aligned_fp4_dim",
+    "cutlass_fp4_gemm_geglu_il_hw_v10",
     "dequantize_fp4_sfa_fp16",
     "e0m3_weight_gemm_fp16",
     "fp4_w4a16_linear_bf16",
     "nvfp4_gemm_streamk_bias_bf16",
     "quantize_fp4_sfa_fp16",
     "quantize_fp4_sfa_mse_fp16",
+    "quantize_fp4_sfa_mse_bf16",
+    "quantize_fp4_sfa_padded_bf16",
+    "pack_nvfp4_weight_bf16",
     "quantize_e0m3_sfa_fp16",
     "quantize_fp4_sfa_bf16",
     "sfa_size_bytes",

build/torch213-cxx11-cu130-x86_64-linux/{_fp4_gemm_cuda_b6588aa.abi3.so → _fp4_gemm_cuda_92c4b99.abi3.so} RENAMED Viewed

@@ -1,3 +1,3 @@
 version https://git-lfs.github.com/spec/v1
-oid sha256:67ad5a71cafc0766b46303485eb638fcd03d048a6b48e97f3a0a90e17585ee2a
-size 7714632

 version https://git-lfs.github.com/spec/v1
+oid sha256:6f97f9a442fba5511c0bcded0b9d388f387b9fd26fe56951698a79ded706ca52
+size 7895736

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

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

 import torch
+from . import _fp4_gemm_cuda_92c4b99
+ops = torch.ops._fp4_gemm_cuda_92c4b99
 def add_op_namespace_prefix(op_name: str):
     """
     Prefix op by namespace.
     """
+    return f"_fp4_gemm_cuda_92c4b99::{op_name}"

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

@@ -1,6 +1,6 @@
 {
   "name": "fp4-gemm",
-  "id": "_fp4_gemm_cuda_b6588aa",
   "version": 1,
   "license": "Apache-2.0",
   "python-depends": [],
@@ -14,9 +14,9 @@
   "digest": {
     "algorithm": "sha256",
     "files": {
-      "__init__.py": "Fb/FDYOpDw+YGlevblUsv4AkILDQCJdB807qdDUDI1c=",
-      "_fp4_gemm_cuda_b6588aa.abi3.so": "Z61accr8B2a0YwNIXrY4/NA9BIprSOl/OgqQ4XWF7io=",
-      "_ops.py": "KKClchvuzBz2g1A148Gi2R6QIx0Ipj0CKlTT511GrUY="
     }
   },
   "provenance": {
@@ -26,7 +26,7 @@
       "dirty": false
     },
     "kernel": {
-      "sha": "b6588aa59dd4f8c646f90b58312a0a66886e061f",
       "dirty": false
     }
   }

 {
   "name": "fp4-gemm",
+  "id": "_fp4_gemm_cuda_92c4b99",
   "version": 1,
   "license": "Apache-2.0",
   "python-depends": [],
   "digest": {
     "algorithm": "sha256",
     "files": {
+      "__init__.py": "Mr1/bAUzs+af68D+bRXX6iONDof//dzKgn6QU5pKdu8=",
+      "_fp4_gemm_cuda_92c4b99.abi3.so": "b5f5pEL7pVEcC83tC504jzh7n9Jv5WlRaYp53tcGylI=",
+      "_ops.py": "c5GHeHUOPDgpedue7MgdtRjCwHOmUsBo4OfqoFVegq0="
     }
   },
   "provenance": {
       "dirty": false
     },
     "kernel": {
+      "sha": "92c4b991b6f1fbd89744aa6f46f62b1bfac7ca7e",
       "dirty": false
     }
   }

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

@@ -109,6 +109,11 @@ def _quant_mse_fake(x, packed, sfa, is_sfb: bool = False) -> None:
     return None
 @torch.library.register_fake(add_op_namespace_prefix("quantize_e0m3_sfa_fp16"))
 def _quant_e0m3_fake(x, packed, sfa, is_sfb: bool = False) -> None:
     return None
@@ -186,6 +191,91 @@ def quantize_fp4_sfa_mse_fp16(
     return packed, sfa
 def quantize_e0m3_sfa_fp16(
     x: torch.Tensor,
     packed: torch.Tensor | None = None,
@@ -381,6 +471,29 @@ def nvfp4_gemm_geglu_nvfp4_fp16(
     return out_packed, out_sfa
 def nvfp4_gemm_bias_gelu_nvfp4_fp16(
     a_packed: torch.Tensor,
     b_packed: torch.Tensor,
@@ -601,6 +714,8 @@ def nvfp4_gemm_streamk_bias_bf16(
 __all__ = [
     "dequantize_fp4_sfa_fp16",
     "e0m3_weight_gemm_fp16",
     "fp4_w4a16_linear_bf16",
@@ -622,6 +737,9 @@ __all__ = [
     "nvfp4_gemm_streamk_bias_bf16",
     "quantize_fp4_sfa_fp16",
     "quantize_fp4_sfa_mse_fp16",
     "quantize_e0m3_sfa_fp16",
     "quantize_fp4_sfa_bf16",
     "sfa_size_bytes",

     return None
+@torch.library.register_fake(add_op_namespace_prefix("quantize_fp4_sfa_mse_bf16"))
+def _quant_mse_bf16_fake(x, packed, sfa, is_sfb: bool = False) -> None:
+    return None
 @torch.library.register_fake(add_op_namespace_prefix("quantize_e0m3_sfa_fp16"))
 def _quant_e0m3_fake(x, packed, sfa, is_sfb: bool = False) -> None:
     return None
     return packed, sfa
+def quantize_fp4_sfa_mse_bf16(
+    x: torch.Tensor,
+    packed: torch.Tensor | None = None,
+    sfa: torch.Tensor | None = None,
+    is_sfb: bool = False,
+):
+    """Pack BF16 weights using per-block reconstruction-MSE scale search."""
+    if packed is None or sfa is None:
+        packed, sfa = _alloc_fp4(x.shape[0], x.shape[1], x.device)
+    ops.quantize_fp4_sfa_mse_bf16(x, packed, sfa, bool(is_sfb))
+    return packed, sfa
+def aligned_fp4_dim(dim: int, alignment: int = 32) -> int:
+    """Return the physical FP4 GEMM dimension for a logical dimension."""
+    if dim <= 0 or alignment <= 0 or alignment % 16 != 0:
+        raise ValueError("dim must be positive and alignment divisible by 16")
+    return ((dim + alignment - 1) // alignment) * alignment
+def quantize_fp4_sfa_padded_bf16(
+    x: torch.Tensor,
+    *,
+    alignment: int = 32,
+    is_sfb: bool = False,
+):
+    """Bind-time BF16 pack with zero padding on the final dimension.
+    The returned packed tensor uses the physical aligned width. Keep the
+    logical width returned as the third item for slicing final model outputs.
+    This helper allocates and is intentionally not a runtime hot-path op.
+    """
+    if x.ndim != 2 or x.dtype != torch.bfloat16:
+        raise ValueError("x must be a 2-D BF16 tensor")
+    logical_dim = int(x.shape[1])
+    physical_dim = aligned_fp4_dim(logical_dim, alignment)
+    if physical_dim == logical_dim:
+        padded = x.contiguous()
+    else:
+        padded = torch.zeros(
+            (x.shape[0], physical_dim), device=x.device, dtype=x.dtype
+        )
+        padded[:, :logical_dim].copy_(x)
+    packed, sfa = quantize_fp4_sfa_bf16(padded, is_sfb=is_sfb)
+    return packed, sfa, logical_dim
+def pack_nvfp4_weight_bf16(
+    weight: torch.Tensor,
+    bias: torch.Tensor | None = None,
+    *,
+    alignment: int = 32,
+    mse: bool = False,
+):
+    """Zero-pad and pack a logical ``(N, K)`` BF16 weight at bind time.
+    Both physical ``N`` and ``K`` are aligned. This is the model-neutral
+    contract used by widths such as SigLIP's logical 4304, which becomes 4320
+    for Blackwell NVFP4 TMA. The hot path consumes the returned static tensors
+    directly and performs no padding or allocation.
+    """
+    if weight.ndim != 2 or weight.dtype != torch.bfloat16:
+        raise ValueError("weight must be a 2-D BF16 tensor")
+    logical_n, logical_k = map(int, weight.shape)
+    physical_n = aligned_fp4_dim(logical_n, alignment)
+    physical_k = aligned_fp4_dim(logical_k, alignment)
+    padded_weight = torch.zeros(
+        (physical_n, physical_k), device=weight.device, dtype=weight.dtype
+    )
+    padded_weight[:logical_n, :logical_k].copy_(weight)
+    quantize = quantize_fp4_sfa_mse_bf16 if mse else quantize_fp4_sfa_bf16
+    packed, sfb = quantize(padded_weight, is_sfb=True)
+    padded_bias = None
+    if bias is not None:
+        if bias.ndim != 1 or bias.numel() != logical_n:
+            raise ValueError("bias must have shape (N,)")
+        if bias.dtype != torch.bfloat16 or bias.device != weight.device:
+            raise ValueError("bias must match the weight dtype and device")
+        padded_bias = torch.zeros(
+            (physical_n,), device=bias.device, dtype=bias.dtype
+        )
+        padded_bias[:logical_n].copy_(bias)
+    return packed, sfb, padded_bias, (logical_n, logical_k)
 def quantize_e0m3_sfa_fp16(
     x: torch.Tensor,
     packed: torch.Tensor | None = None,
     return out_packed, out_sfa
+def cutlass_fp4_gemm_geglu_il_hw_v10(
+    a_packed: torch.Tensor,
+    b_interleaved_packed: torch.Tensor,
+    sfa: torch.Tensor,
+    sfb: torch.Tensor,
+    *,
+    scratch: torch.Tensor | None = None,
+    out_packed: torch.Tensor | None = None,
+    out_sfa: torch.Tensor | None = None,
+) -> tuple[torch.Tensor, torch.Tensor]:
+    """Explicit public alias for the native SM110 v10 GeGLU schedule."""
+    return nvfp4_gemm_geglu_nvfp4_fp16(
+        a_packed,
+        b_interleaved_packed,
+        sfa,
+        sfb,
+        skinny=True,
+        scratch=scratch,
+        out_packed=out_packed,
+        out_sfa=out_sfa,
+    )
 def nvfp4_gemm_bias_gelu_nvfp4_fp16(
     a_packed: torch.Tensor,
     b_packed: torch.Tensor,
 __all__ = [
+    "aligned_fp4_dim",
+    "cutlass_fp4_gemm_geglu_il_hw_v10",
     "dequantize_fp4_sfa_fp16",
     "e0m3_weight_gemm_fp16",
     "fp4_w4a16_linear_bf16",
     "nvfp4_gemm_streamk_bias_bf16",
     "quantize_fp4_sfa_fp16",
     "quantize_fp4_sfa_mse_fp16",
+    "quantize_fp4_sfa_mse_bf16",
+    "quantize_fp4_sfa_padded_bf16",
+    "pack_nvfp4_weight_bf16",
     "quantize_e0m3_sfa_fp16",
     "quantize_fp4_sfa_bf16",
     "sfa_size_bytes",

build/torch213-cxx11-cu132-x86_64-linux/{_fp4_gemm_cuda_b6588aa.abi3.so → _fp4_gemm_cuda_92c4b99.abi3.so} RENAMED Viewed

@@ -1,3 +1,3 @@
 version https://git-lfs.github.com/spec/v1
-oid sha256:4b27d4b88246c79fd017c8845f6f99cd1c9e7946039fda8f6abec2635a4edd65
-size 7706384

 version https://git-lfs.github.com/spec/v1
+oid sha256:d1cfd40cecbd45d8a9e0958a5f9af7508f4b8deb0910eda2ad078c5e82690c55
+size 7887496

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

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

 import torch
+from . import _fp4_gemm_cuda_92c4b99
+ops = torch.ops._fp4_gemm_cuda_92c4b99
 def add_op_namespace_prefix(op_name: str):
     """
     Prefix op by namespace.
     """
+    return f"_fp4_gemm_cuda_92c4b99::{op_name}"

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

@@ -1,6 +1,6 @@
 {
   "name": "fp4-gemm",
-  "id": "_fp4_gemm_cuda_b6588aa",
   "version": 1,
   "license": "Apache-2.0",
   "python-depends": [],
@@ -14,9 +14,9 @@
   "digest": {
     "algorithm": "sha256",
     "files": {
-      "__init__.py": "Fb/FDYOpDw+YGlevblUsv4AkILDQCJdB807qdDUDI1c=",
-      "_fp4_gemm_cuda_b6588aa.abi3.so": "SyfUuIJGx5/QF8iEX2+ZzRyeeUYDn9qPar7CY1pO3WU=",
-      "_ops.py": "KKClchvuzBz2g1A148Gi2R6QIx0Ipj0CKlTT511GrUY="
     }
   },
   "provenance": {
@@ -26,7 +26,7 @@
       "dirty": false
     },
     "kernel": {
-      "sha": "b6588aa59dd4f8c646f90b58312a0a66886e061f",
       "dirty": false
     }
   }

 {
   "name": "fp4-gemm",
+  "id": "_fp4_gemm_cuda_92c4b99",
   "version": 1,
   "license": "Apache-2.0",
   "python-depends": [],
   "digest": {
     "algorithm": "sha256",
     "files": {
+      "__init__.py": "Mr1/bAUzs+af68D+bRXX6iONDof//dzKgn6QU5pKdu8=",
+      "_fp4_gemm_cuda_92c4b99.abi3.so": "0c/UDOy9Rdip4JWKX5r3UI9LjesJEO2irQeMXoJpDFU=",
+      "_ops.py": "c5GHeHUOPDgpedue7MgdtRjCwHOmUsBo4OfqoFVegq0="
     }
   },
   "provenance": {
       "dirty": false
     },
     "kernel": {
+      "sha": "92c4b991b6f1fbd89744aa6f46f62b1bfac7ca7e",
       "dirty": false
     }
   }