Add torch213-cxx11-cu130-aarch64-linux artifact

Browse files

Files changed (4) hide show

build/torch213-cxx11-cu130-aarch64-linux/__init__.py +77 -0
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 +10 -6

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

@@ -42,6 +42,16 @@ def _linear_fp16_fake(a_packed, b_packed, sfa, sfb, out, alpha: float = 1.0, var
     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
@@ -94,6 +104,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("quantize_e0m3_sfa_fp16"))
 def _quant_e0m3_fake(x, packed, sfa, is_sfb: bool = False) -> None:
     return None
@@ -158,6 +173,19 @@ def quantize_fp4_sfa_fp16(
     return packed, sfa
 def quantize_e0m3_sfa_fp16(
     x: torch.Tensor,
     packed: torch.Tensor | None = None,
@@ -276,6 +304,52 @@ def nvfp4_gemm_fp16(
     return out
 def nvfp4_gemm_geglu_nvfp4_fp16(
     a_packed: torch.Tensor,
     b_interleaved_packed: torch.Tensor,
@@ -533,6 +607,8 @@ __all__ = [
     "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",
@@ -545,6 +621,7 @@ __all__ = [
     "nvfp4_gemm_streamk_bf16",
     "nvfp4_gemm_streamk_bias_bf16",
     "quantize_fp4_sfa_fp16",
     "quantize_e0m3_sfa_fp16",
     "quantize_fp4_sfa_bf16",
     "sfa_size_bytes",

     return None
+@torch.library.register_fake(add_op_namespace_prefix("nvfp4_gemm_variant_bf16"))
+def _linear_variant_bf16_fake(a, b, sfa, sfb, out, alpha: float = 1.0, variant: int = 7) -> None:
+    return None
+@torch.library.register_fake(add_op_namespace_prefix("nvfp4_gemm_nvfp4"))
+def _linear_fp4out_fake(a, b, sfa, sfb, out_packed, out_sfa) -> 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
     return None
+@torch.library.register_fake(add_op_namespace_prefix("quantize_fp4_sfa_mse_fp16"))
+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
     return packed, sfa
+def quantize_fp4_sfa_mse_fp16(
+    x: torch.Tensor,
+    packed: torch.Tensor | None = None,
+    sfa: torch.Tensor | None = None,
+    is_sfb: bool = False,
+):
+    """Pack FP16 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_fp16(x, packed, sfa, bool(is_sfb))
+    return packed, sfa
 def quantize_e0m3_sfa_fp16(
     x: torch.Tensor,
     packed: torch.Tensor | None = None,
     return out
+def nvfp4_gemm_variant_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 = 7,
+) -> torch.Tensor:
+    """SM110 native NVFP4 GEMM with BF16 output.
+    ``variant=7`` serves large-M encoder/SigLIP shapes; ``variant=10`` is
+    the narrow-N PI0.5 decoder tile. The indices match FlashRT native.
+    """
+    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_variant_bf16(
+        a_packed, b_packed, sfa, sfb, out, float(alpha), int(variant)
+    )
+    return out
+def nvfp4_gemm_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,
+) -> tuple[torch.Tensor, torch.Tensor]:
+    """SM110 NVFP4 GEMM with an NVFP4 output and CUTLASS SFA layout."""
+    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_nvfp4(
+        a_packed, b_packed, sfa, sfb, out_packed, out_sfa
+    )
+    return out_packed, out_sfa
 def nvfp4_gemm_geglu_nvfp4_fp16(
     a_packed: torch.Tensor,
     b_interleaved_packed: torch.Tensor,
     "fp4_w4a4_gemv_warpsplit_bf16",
     "nvfp4_gemm_bf16",
     "nvfp4_gemm_fp16",
+    "nvfp4_gemm_variant_bf16",
+    "nvfp4_gemm_nvfp4",
     "nvfp4_gemm_geglu_nvfp4_fp16",
     "nvfp4_gemm_bias_gelu_nvfp4_fp16",
     "nvfp4_gemm_bias_residual_fp16",
     "nvfp4_gemm_streamk_bf16",
     "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",

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

@@ -1,6 +1,6 @@
 import torch
-from . import _fp4_gemm_cuda_230e467
-ops = torch.ops._fp4_gemm_cuda_230e467
 def add_op_namespace_prefix(op_name: str):
-    return f"_fp4_gemm_cuda_230e467::{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:73b4cbcb15d13cc6452007e59961f3b217d21771f0c0937cc121cc8e0e6b1e12
+size 6529016

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

@@ -1,28 +1,32 @@
 {
   "name": "fp4-gemm",
-  "id": "_fp4_gemm_cuda_230e467",
   "version": 1,
   "license": "Apache-2.0",
   "python-depends": [],
   "backend": {
     "type": "cuda",
     "archs": [
-      "11.0a"
     ]
   },
   "digest": {
     "algorithm": "sha256",
     "files": {
-      "__init__.py": "Qfdvg/9EedWLGh34xAalDfC0IsAHCOL44L4ShaP/w5Y=",
-      "_fp4_gemm_cuda_230e467.abi3.so": "VNqOTzsMcoxjCdy0NPtgjPUAhd+/MFE4H/axaHMWEDw=",
-      "_ops.py": "IfU2iMVrWYRnoe6q8HI3hU8TPnAlLrXDerlv117PcXI=",
       "fp4_gemm/__init__.py": "v6p5XMfQzddhi1fLSAw4HX9CyS0rQsidvu9VsT01xi4="
     }
   },
   "provenance": {
     "kernel": {
-      "sha": "230e467",
       "dirty": false
     }
   }
 }

 {
   "name": "fp4-gemm",
+  "id": "fp4_gemm_source_test",
   "version": 1,
   "license": "Apache-2.0",
   "python-depends": [],
   "backend": {
     "type": "cuda",
     "archs": [
+      "11.0"
     ]
   },
   "digest": {
     "algorithm": "sha256",
     "files": {
+      "__init__.py": "Fb/FDYOpDw+YGlevblUsv4AkILDQCJdB807qdDUDI1c=",
+      "fp4_gemm_source_test.abi3.so": "c7TLyxXRPMZFIAflmWHzshfSF3HwwJN8wSHMjg5rHhI=",
+      "_ops.py": "ZPjRxrYk7lCA0ZHoCjg8BaJyFY7a5i4NC3jIL6VMRJ4=",
       "fp4_gemm/__init__.py": "v6p5XMfQzddhi1fLSAw4HX9CyS0rQsidvu9VsT01xi4="
     }
   },
   "provenance": {
     "kernel": {
+      "sha": "b6588aa59dd4f8c646f90b58312a0a66886e061f",
       "dirty": false
+    },
+    "validation": {
+      "torch": "2.13.0+cu130",
+      "cuda": "13.0"
     }
   }
 }