Publish torch213-cxx11-cu130-x86_64-linux

Browse files

Files changed (4) hide show

build/torch213-cxx11-cu130-x86_64-linux/__init__.py +52 -0
build/torch213-cxx11-cu130-x86_64-linux/_fp4_gemm_cuda_d0297c9.abi3.so +3 -0
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-cu130-x86_64-linux/__init__.py CHANGED Viewed

@@ -127,6 +127,13 @@ def _gemm_warpsplit_mrows_fake(a_packed, b_packed, sfa, sfb, out, alpha: float =
     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:
@@ -257,6 +264,11 @@ def _quant_bf16_fake(x: torch.Tensor, packed: torch.Tensor, sfa: torch.Tensor, i
     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
@@ -468,6 +480,19 @@ def quantize_fp4_sfa_bf16(
     return packed, sfa
 def dequantize_fp4_sfa_fp16(
     packed: torch.Tensor,
     sfa: torch.Tensor,
@@ -683,6 +708,31 @@ def fp4_w4a4_gemm_warpsplit_mrows_bf16(
     return out
 def fp4_w4a4_gemv_warpsplit_bf16(
     a_packed: torch.Tensor,
     b_packed: torch.Tensor,
@@ -1062,6 +1112,7 @@ __all__ = [
     "fp4_w4a4_gemv_warpsplit_bf16",
     "fp4_w4a4_gemv_warpsplit_interleaved_bf16",
     "fp4_w4a4_gemm_warpsplit_mrows_bf16",
     "nvfp4_gemm_bf16",
     "nvfp4_gemm_fp16",
     "nvfp4_gemm_variant_bf16",
@@ -1087,5 +1138,6 @@ __all__ = [
     "pack_nvfp4_weight_bf16",
     "quantize_e0m3_sfa_fp16",
     "quantize_fp4_sfa_bf16",
     "sfa_size_bytes",
 ]

     return None
+@torch.library.register_fake(add_op_namespace_prefix("fp4_w4a4_gemm_warpsplit_mrows_pdl_bf16"))
+def _gemm_warpsplit_mrows_pdl_fake(a_packed, b_packed, sfa, sfb, out, alpha: float = 1.0, warps: int = 2, stages: int = 6) -> None:
+    return _gemm_warpsplit_mrows_fake(
+        a_packed, b_packed, sfa, sfb, out, alpha, warps, stages
+    )
 @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:
     return None
+@torch.library.register_fake(add_op_namespace_prefix("quantize_fp4_sfa_bf16_pdl"))
+def _quant_bf16_pdl_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
     return packed, sfa
+def quantize_fp4_sfa_bf16_pdl(
+    x: torch.Tensor,
+    packed: torch.Tensor | None = None,
+    sfa: torch.Tensor | None = None,
+    is_sfb: bool = False,
+):
+    """PDL-enabled twin of :func:`quantize_fp4_sfa_bf16`."""
+    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_pdl(x, packed, sfa, bool(is_sfb))
+    return packed, sfa
 def dequantize_fp4_sfa_fp16(
     packed: torch.Tensor,
     sfa: torch.Tensor,
     return out
+def fp4_w4a4_gemm_warpsplit_mrows_pdl_bf16(
+    a_packed: torch.Tensor,
+    b_packed: torch.Tensor,
+    sfa: torch.Tensor,
+    sfb: torch.Tensor,
+    *,
+    alpha: float = 1.0,
+    warps: int = 2,
+    stages: int = 6,
+    out: Optional[torch.Tensor] = None,
+) -> torch.Tensor:
+    """PDL-enabled twin of the SM120 M<=16 warp-split GEMM."""
+    if out is None:
+        out = torch.empty(
+            (a_packed.shape[0], b_packed.shape[0]),
+            device=a_packed.device,
+            dtype=torch.bfloat16,
+        )
+    ops.fp4_w4a4_gemm_warpsplit_mrows_pdl_bf16(
+        a_packed, b_packed, sfa, sfb, out,
+        float(alpha), int(warps), int(stages),
+    )
+    return out
 def fp4_w4a4_gemv_warpsplit_bf16(
     a_packed: torch.Tensor,
     b_packed: torch.Tensor,
     "fp4_w4a4_gemv_warpsplit_bf16",
     "fp4_w4a4_gemv_warpsplit_interleaved_bf16",
     "fp4_w4a4_gemm_warpsplit_mrows_bf16",
+    "fp4_w4a4_gemm_warpsplit_mrows_pdl_bf16",
     "nvfp4_gemm_bf16",
     "nvfp4_gemm_fp16",
     "nvfp4_gemm_variant_bf16",
     "pack_nvfp4_weight_bf16",
     "quantize_e0m3_sfa_fp16",
     "quantize_fp4_sfa_bf16",
+    "quantize_fp4_sfa_bf16_pdl",
     "sfa_size_bytes",
 ]

build/torch213-cxx11-cu130-x86_64-linux/_fp4_gemm_cuda_d0297c9.abi3.so ADDED Viewed

	@@ -0,0 +1,3 @@

+version https://git-lfs.github.com/spec/v1
+oid sha256:53fcef343db9da31b616727c11c8c71d298383eaca65d95547b0fab907bff372
+size 11046856

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

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

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

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

@@ -1,6 +1,6 @@
 {
   "name": "fp4-gemm",
-  "id": "_fp4_gemm_cuda_6971f7a",
   "version": 1,
   "license": "Apache-2.0",
   "python-depends": [],
@@ -14,9 +14,9 @@
   "digest": {
     "algorithm": "sha256",
     "files": {
-      "__init__.py": "0hDN2W6X3IMoWIi5BBufrnqrq2yj5N1/q9wBhCwBRxA=",
-      "_fp4_gemm_cuda_6971f7a.abi3.so": "0qnksHZNwFkyIXuU7z00XYofRqosbRcU45ISTPSLouk=",
-      "_ops.py": "D7InzLughnNL5JQowMPNYpSpNPRFeeD8mlgs38Xw03g=",
       "fp4_gemm/__init__.py": "DFYPlrhXwYjEqCl/8n0SmWGZV8NFml5DPhMjKfv98GY="
     }
   },
@@ -27,7 +27,7 @@
       "dirty": false
     },
     "kernel": {
-      "sha": "6971f7a37410ba0d6edae93b389dbcbe3ba924d5",
       "dirty": false
     }
   }

 {
   "name": "fp4-gemm",
+  "id": "_fp4_gemm_cuda_d0297c9",
   "version": 1,
   "license": "Apache-2.0",
   "python-depends": [],
   "digest": {
     "algorithm": "sha256",
     "files": {
+      "__init__.py": "vJAc910iyL+GomTQLPDvl5MM5X0d6a8RzdQ/SV/BiQ8=",
+      "_fp4_gemm_cuda_d0297c9.abi3.so": "U/zvND252jG2FnJ8EcjHHSmDg+rKZdlVR7D6uQe/83I=",
+      "_ops.py": "mFWrViw+mMqyMmt/TuyUUNjzunMWYkC4JH0mqbPp6KQ=",
       "fp4_gemm/__init__.py": "DFYPlrhXwYjEqCl/8n0SmWGZV8NFml5DPhMjKfv98GY="
     }
   },
       "dirty": false
     },
     "kernel": {
+      "sha": "d0297c93d91ab430899d09ba9f9bd1c05fbfc9f9",
       "dirty": false
     }
   }