Add torch213-cxx11-cu130-aarch64-linux artifact

Browse files

Files changed (4) hide show

build/torch213-cxx11-cu130-aarch64-linux/__init__.py +195 -2
build/torch213-cxx11-cu130-aarch64-linux/_ops.py +3 -3
build/torch213-cxx11-cu130-aarch64-linux/fp4_gemm_source_test.abi3.so +3 -0
build/torch213-cxx11-cu130-aarch64-linux/metadata.json +9 -5

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

@@ -19,7 +19,8 @@ def sfa_size_bytes(rows: int, dim: int) -> int:
 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),
     )
@@ -36,6 +37,26 @@ def _linear_fake(
     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:
@@ -73,6 +94,21 @@ 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
@@ -122,6 +158,62 @@ def quantize_fp4_sfa_fp16(
     return packed, sfa
 def quantize_fp4_sfa_bf16(
     x: torch.Tensor,
     packed: torch.Tensor | None = None,
@@ -162,6 +254,100 @@ def nvfp4_gemm_bf16(
     return out
 def fp4_w4a4_gemv_warpsplit_bf16(
     a_packed: torch.Tensor,
     b_packed: torch.Tensor,
@@ -292,7 +478,7 @@ def nvfp4_gemm_bias_gelu_nvfp4(
     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)
     )
@@ -342,17 +528,24 @@ def nvfp4_gemm_streamk_bias_bf16(
 __all__ = [
     "dequantize_fp4_sfa_fp16",
     "fp4_w4a16_linear_bf16",
     "fp4_w4a4_gemv_warpsplit_bf16",
     "nvfp4_gemm_bf16",
     "nvfp4_gemm_bias_bf16",
     "nvfp4_gemm_bias_gelu_bf16",
     "nvfp4_gemm_bias_gelu_nvfp4",
     "nvfp4_gemm_bias_residual_bf16",
     "nvfp4_gemm_residual_bf16",
     "nvfp4_gemm_streamk_bf16",
     "nvfp4_gemm_streamk_bias_bf16",
     "quantize_fp4_sfa_fp16",
     "quantize_fp4_sfa_bf16",
     "sfa_size_bytes",
 ]

 def _alloc_fp4(rows: int, dim: int, device: torch.device | str):
     return (
         torch.empty((rows, dim // 2), device=device, dtype=torch.uint8),
+        # Tile-layout padding entries are not written by every quantizer.
+        torch.zeros((sfa_size_bytes(rows, dim),), device=device, dtype=torch.uint8),
     )
     return None
+@torch.library.register_fake(add_op_namespace_prefix("nvfp4_gemm_fp16"))
+def _linear_fp16_fake(a_packed, b_packed, sfa, sfb, out, alpha: float = 1.0, variant: int = -1) -> None:
+    return None
+@torch.library.register_fake(add_op_namespace_prefix("nvfp4_gemm_geglu_nvfp4_fp16"))
+def _geglu_fp4_fake(a, b, sfa, sfb, scratch, out_packed, out_sfa, skinny: bool = False) -> None:
+    return None
+@torch.library.register_fake(add_op_namespace_prefix("nvfp4_gemm_bias_gelu_nvfp4_fp16"))
+def _bias_gelu_fp4_fp16_fake(a, b, sfa, sfb, bias, out_packed, out_sfa) -> None:
+    return None
+@torch.library.register_fake(add_op_namespace_prefix("nvfp4_gemm_bias_residual_fp16"))
+def _bias_residual_fp16_fake(a, b, sfa, sfb, bias, residual, out) -> 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:
     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
+@torch.library.register_fake(add_op_namespace_prefix("e0m3_weight_gemm_fp16"))
+def _e0m3_gemm_fake(a, b, sfa, sfb, out, alpha: float = 1.0, a_format: int = 1) -> None:
+    return None
+@torch.library.register_fake(add_op_namespace_prefix("nvfp4_gemm_relu2_nvfp4"))
+def _relu2_gemm_fake(a, b, sfa, sfb, out_packed, out_sfa) -> 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
     return packed, sfa
+def quantize_e0m3_sfa_fp16(
+    x: torch.Tensor,
+    packed: torch.Tensor | None = None,
+    sfa: torch.Tensor | None = None,
+    is_sfb: bool = False,
+):
+    """Quantize FP16 to packed E0M3 with CUTLASS SFA/SFB layout."""
+    if packed is None or sfa is None:
+        packed, sfa = _alloc_fp4(x.shape[0], x.shape[1], x.device)
+    ops.quantize_e0m3_sfa_fp16(x, packed, sfa, bool(is_sfb))
+    return packed, sfa
+def e0m3_weight_gemm_fp16(
+    a_packed: torch.Tensor,
+    b_packed: torch.Tensor,
+    sfa: torch.Tensor,
+    sfb: torch.Tensor,
+    *,
+    alpha: float = 1.0,
+    a_format: int = 1,
+    out: torch.Tensor | None = None,
+):
+    """SM110 GEMM with E0M3 weights and E2M1 or E0M3 activations."""
+    if out is None:
+        out = torch.empty(
+            (a_packed.shape[0], b_packed.shape[0]),
+            device=a_packed.device,
+            dtype=torch.float16,
+        )
+    ops.e0m3_weight_gemm_fp16(
+        a_packed, b_packed, sfa, sfb, out, float(alpha), int(a_format)
+    )
+    return out
+def nvfp4_gemm_relu2_nvfp4(
+    a_packed: torch.Tensor,
+    b_packed: torch.Tensor,
+    sfa: torch.Tensor,
+    sfb: torch.Tensor,
+    *,
+    out_packed: torch.Tensor | None = None,
+    out_sfa: torch.Tensor | None = None,
+):
+    """SM110 NVFP4 GEMM with fused ReLU-squared and NVFP4 output."""
+    if out_packed is None or out_sfa is None:
+        out_packed, out_sfa = _alloc_fp4(
+            a_packed.shape[0], b_packed.shape[0], a_packed.device
+        )
+    ops.nvfp4_gemm_relu2_nvfp4(
+        a_packed, b_packed, sfa, sfb, out_packed, out_sfa
+    )
+    return out_packed, out_sfa
 def quantize_fp4_sfa_bf16(
     x: torch.Tensor,
     packed: torch.Tensor | None = None,
     return out
+def nvfp4_gemm_fp16(
+    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:
+    """SM110 native NVFP4 GEMM with FP16 output."""
+    if out is None:
+        out = torch.empty(
+            (a_packed.shape[0], b_packed.shape[0]),
+            device=a_packed.device,
+            dtype=torch.float16,
+        )
+    ops.nvfp4_gemm_fp16(
+        a_packed, b_packed, sfa, sfb, out, float(alpha), int(variant)
+    )
+    return out
+def nvfp4_gemm_geglu_nvfp4_fp16(
+    a_packed: torch.Tensor,
+    b_interleaved_packed: torch.Tensor,
+    sfa: torch.Tensor,
+    sfb: torch.Tensor,
+    *,
+    skinny: bool = False,
+    scratch: torch.Tensor | None = None,
+    out_packed: torch.Tensor | None = None,
+    out_sfa: torch.Tensor | None = None,
+) -> tuple[torch.Tensor, torch.Tensor]:
+    """GEMM with fused GeGLU and compact NVFP4 output on SM110.
+    ``b_interleaved_packed`` stores gate/up rows pairwise, so its first
+    dimension is twice the logical hidden width.
+    """
+    m, n_twice = a_packed.shape[0], b_interleaved_packed.shape[0]
+    hidden = n_twice // 2
+    if scratch is None:
+        scratch = torch.empty((m, hidden), device=a_packed.device, dtype=torch.uint8)
+    if out_packed is None:
+        out_packed = torch.empty((m, hidden // 2), device=a_packed.device, dtype=torch.uint8)
+    if out_sfa is None:
+        out_sfa = torch.zeros((sfa_size_bytes(m, hidden),), device=a_packed.device, dtype=torch.uint8)
+    ops.nvfp4_gemm_geglu_nvfp4_fp16(
+        a_packed, b_interleaved_packed, sfa, sfb, scratch,
+        out_packed, out_sfa, bool(skinny)
+    )
+    return out_packed, out_sfa
+def nvfp4_gemm_bias_gelu_nvfp4_fp16(
+    a_packed: torch.Tensor,
+    b_packed: torch.Tensor,
+    sfa: torch.Tensor,
+    sfb: torch.Tensor,
+    bias: torch.Tensor,
+    *,
+    out_packed: torch.Tensor | None = None,
+    out_sfa: torch.Tensor | None = None,
+) -> tuple[torch.Tensor, torch.Tensor]:
+    """FP16-bias GEMM with fused GELU and NVFP4 output on SM110."""
+    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.zeros((sfa_size_bytes(m, n),), device=a_packed.device, dtype=torch.uint8)
+    ops.nvfp4_gemm_bias_gelu_nvfp4_fp16(
+        a_packed, b_packed, sfa, sfb, bias, out_packed, out_sfa
+    )
+    return out_packed, out_sfa
+def nvfp4_gemm_bias_residual_fp16(
+    a_packed: torch.Tensor,
+    b_packed: torch.Tensor,
+    sfa: torch.Tensor,
+    sfb: torch.Tensor,
+    bias: torch.Tensor,
+    residual: torch.Tensor,
+    *,
+    out: torch.Tensor | None = None,
+) -> torch.Tensor:
+    """FP16-output GEMM with fused FP16 bias and residual on SM110."""
+    if out is None:
+        out = torch.empty_like(residual)
+    ops.nvfp4_gemm_bias_residual_fp16(
+        a_packed, b_packed, sfa, sfb, bias, residual, out
+    )
+    return out
 def fp4_w4a4_gemv_warpsplit_bf16(
     a_packed: torch.Tensor,
     b_packed: torch.Tensor,
     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.zeros((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)
     )
 __all__ = [
     "dequantize_fp4_sfa_fp16",
+    "e0m3_weight_gemm_fp16",
     "fp4_w4a16_linear_bf16",
     "fp4_w4a4_gemv_warpsplit_bf16",
     "nvfp4_gemm_bf16",
+    "nvfp4_gemm_fp16",
+    "nvfp4_gemm_geglu_nvfp4_fp16",
+    "nvfp4_gemm_bias_gelu_nvfp4_fp16",
+    "nvfp4_gemm_bias_residual_fp16",
     "nvfp4_gemm_bias_bf16",
     "nvfp4_gemm_bias_gelu_bf16",
     "nvfp4_gemm_bias_gelu_nvfp4",
     "nvfp4_gemm_bias_residual_bf16",
     "nvfp4_gemm_residual_bf16",
+    "nvfp4_gemm_relu2_nvfp4",
     "nvfp4_gemm_streamk_bf16",
     "nvfp4_gemm_streamk_bias_bf16",
     "quantize_fp4_sfa_fp16",
+    "quantize_e0m3_sfa_fp16",
     "quantize_fp4_sfa_bf16",
     "sfa_size_bytes",
 ]

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

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

 import torch
+from . import fp4_gemm_source_test
+ops = torch.ops.fp4_gemm_source_test
 def add_op_namespace_prefix(op_name: str):
+    return f"fp4_gemm_source_test::{op_name}"

build/torch213-cxx11-cu130-aarch64-linux/fp4_gemm_source_test.abi3.so ADDED Viewed

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

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

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

@@ -1,6 +1,6 @@
 {
   "name": "fp4-gemm",
-  "id": "_fp4_gemm_cuda_8a66d8b",
   "version": 1,
   "license": "Apache-2.0",
   "python-depends": [],
@@ -13,16 +13,20 @@
   "digest": {
     "algorithm": "sha256",
     "files": {
-      "__init__.py": "+Kk/VNnIWwe9nszQYvL2qtU3DpKxbqHVfklg5WtVuB4=",
-      "_fp4_gemm_cuda_8a66d8b.abi3.so": "R542jHXp/QBVPiBucDjx4J0Ls3zi97wQyKqsCAH4NIo=",
-      "_ops.py": "/hna1MwmSjR4pewWf12X5smsDRcaiHuoR2Gri/PrBFc=",
       "fp4_gemm/__init__.py": "v6p5XMfQzddhi1fLSAw4HX9CyS0rQsidvu9VsT01xi4="
     }
   },
   "provenance": {
     "kernel": {
-      "sha": "8a66d8b",
       "dirty": false
     }
   }
 }

 {
   "name": "fp4-gemm",
+  "id": "fp4_gemm_source_test",
   "version": 1,
   "license": "Apache-2.0",
   "python-depends": [],
   "digest": {
     "algorithm": "sha256",
     "files": {
+      "__init__.py": "Qfdvg/9EedWLGh34xAalDfC0IsAHCOL44L4ShaP/w5Y=",
+      "fp4_gemm_source_test.abi3.so": "twiQsoD7EQX+2vIWpWOL8aN6hm/l9gwLKapUfhhsix0=",
+      "_ops.py": "ZPjRxrYk7lCA0ZHoCjg8BaJyFY7a5i4NC3jIL6VMRJ4=",
       "fp4_gemm/__init__.py": "v6p5XMfQzddhi1fLSAw4HX9CyS0rQsidvu9VsT01xi4="
     }
   },
   "provenance": {
     "kernel": {
+      "sha": "3a505c0",
       "dirty": false
+    },
+    "validation": {
+      "torch": "2.13.0+cu130",
+      "cuda": "13.0"
     }
   }
 }