Update Torch 2.11 CUDA 12.8 BF16 FP4 producer artifact

Browse files

Files changed (4) hide show

build/torch211-cxx11-cu128-x86_64-linux/__init__.py +57 -4
build/torch211-cxx11-cu128-x86_64-linux/{_fp4_gemm_cuda_7781728.abi3.so → _fp4_gemm_cuda_83a7cd0.abi3.so} +2 -2
build/torch211-cxx11-cu128-x86_64-linux/_ops.py +3 -3
build/torch211-cxx11-cu128-x86_64-linux/metadata.json +2 -22

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

@@ -31,11 +31,20 @@ def _linear_fake(
     sfb: torch.Tensor,
     out: torch.Tensor,
     alpha: float = 1.0,
-    variant: int = 0,
 ) -> None:
     return None
 @torch.library.register_fake(add_op_namespace_prefix("fp4_w4a16_linear_bf16"))
 def _legacy_linear_fake(
     a_packed: torch.Tensor,
@@ -44,7 +53,7 @@ def _legacy_linear_fake(
     sfb: torch.Tensor,
     out: torch.Tensor,
     alpha: float = 1.0,
-    variant: int = 0,
 ) -> None:
     return None
@@ -54,6 +63,11 @@ def _quant_fake(x: torch.Tensor, packed: torch.Tensor, sfa: torch.Tensor, is_sfb
     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
@@ -98,6 +112,19 @@ def quantize_fp4_sfa_fp16(
     return packed, sfa
 def dequantize_fp4_sfa_fp16(
     packed: torch.Tensor,
     sfa: torch.Tensor,
@@ -117,7 +144,7 @@ def nvfp4_gemm_bf16(
     sfb: torch.Tensor,
     alpha: float = 1.0,
     out: torch.Tensor | None = None,
-    variant: int = 0,
 ) -> torch.Tensor:
     if out is None:
         out = torch.empty((a_packed.shape[0], b_packed.shape[0]), device=a_packed.device, dtype=torch.bfloat16)
@@ -125,6 +152,30 @@ def nvfp4_gemm_bf16(
     return out
 def fp4_w4a16_linear_bf16(
     a_packed: torch.Tensor,
     b_packed: torch.Tensor,
@@ -132,7 +183,7 @@ def fp4_w4a16_linear_bf16(
     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(
@@ -243,6 +294,7 @@ def nvfp4_gemm_streamk_bias_bf16(
 __all__ = [
     "dequantize_fp4_sfa_fp16",
     "fp4_w4a16_linear_bf16",
     "nvfp4_gemm_bf16",
     "nvfp4_gemm_bias_gelu_bf16",
     "nvfp4_gemm_bias_gelu_nvfp4",
@@ -250,5 +302,6 @@ __all__ = [
     "nvfp4_gemm_streamk_bf16",
     "nvfp4_gemm_streamk_bias_bf16",
     "quantize_fp4_sfa_fp16",
     "sfa_size_bytes",
 ]

     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,
     sfb: torch.Tensor,
     out: torch.Tensor,
     alpha: float = 1.0,
+    variant: int = -1,
 ) -> None:
     return 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
     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,
     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)
     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,
     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(
 __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_streamk_bf16",
     "nvfp4_gemm_streamk_bias_bf16",
     "quantize_fp4_sfa_fp16",
+    "quantize_fp4_sfa_bf16",
     "sfa_size_bytes",
 ]

build/torch211-cxx11-cu128-x86_64-linux/{_fp4_gemm_cuda_7781728.abi3.so → _fp4_gemm_cuda_83a7cd0.abi3.so} RENAMED Viewed

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

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

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

@@ -1,9 +1,9 @@
 import torch
-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_7781728::{op_name}"

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

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

@@ -1,6 +1,6 @@
 {
   "name": "fp4-gemm",
-  "id": "_fp4_gemm_cuda_7781728",
   "version": 1,
   "license": "Apache-2.0",
   "python-depends": [],
@@ -9,25 +9,5 @@
     "archs": [
       "12.0a"
     ]
-  },
-  "digest": {
-    "algorithm": "sha256",
-    "files": {
-      "__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
-    }
   }
-}

 {
   "name": "fp4-gemm",
+  "id": "_fp4_gemm_cuda_83a7cd0",
   "version": 1,
   "license": "Apache-2.0",
   "python-depends": [],
     "archs": [
       "12.0a"
     ]
   }
+}