Add validated torch213-cxx11-cu130-aarch64-linux Thor artifact

Browse files

Files changed (5) hide show

build/torch213-cxx11-cu130-aarch64-linux/__init__.py +358 -0
build/torch213-cxx11-cu130-aarch64-linux/_fp4_gemm_cuda_8a66d8b.abi3.so +3 -0
build/torch213-cxx11-cu130-aarch64-linux/_ops.py +6 -0
build/torch213-cxx11-cu130-aarch64-linux/fp4_gemm/__init__.py +14 -0
build/torch213-cxx11-cu130-aarch64-linux/metadata.json +28 -0

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

	@@ -0,0 +1,358 @@

+"""FlashRT FP4 GEMM kernels."""
+from __future__ import annotations
+import torch
+from ._ops import add_op_namespace_prefix, ops
+def sfa_size_bytes(rows: int, dim: int) -> int:
+    if rows <= 0 or dim <= 0 or dim % 16 != 0:
+        raise ValueError("rows must be positive and dim must be positive/divisible by 16")
+    n_blocks = dim // 16
+    n_row_super = (rows + 127) // 128
+    n_col_super = (n_blocks + 3) // 4
+    return n_row_super * n_col_super * 512
+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),
+    )
+@torch.library.register_fake(add_op_namespace_prefix("nvfp4_gemm_bf16"))
+def _linear_fake(
+    a_packed: torch.Tensor,
+    b_packed: torch.Tensor,
+    sfa: torch.Tensor,
+    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,
+    b_packed: torch.Tensor,
+    sfa: torch.Tensor,
+    sfb: torch.Tensor,
+    out: torch.Tensor,
+    alpha: float = 1.0,
+    variant: int = -1,
+) -> None:
+    return None
+@torch.library.register_fake(add_op_namespace_prefix("nvfp4_gemm_bias_bf16"))
+def _bias_fake(a, b, sfa, sfb, bias, out) -> None:
+    return None
+@torch.library.register_fake(add_op_namespace_prefix("nvfp4_gemm_bias_residual_bf16"))
+def _bias_residual_fake(a, b, sfa, sfb, bias, residual, out) -> None:
+    return None
+@torch.library.register_fake(add_op_namespace_prefix("quantize_fp4_sfa_fp16"))
+def _quant_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("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
+@torch.library.register_fake(add_op_namespace_prefix("nvfp4_gemm_residual_bf16"))
+def _residual_fake(a, b, sfa, sfb, residual, out, alpha: float = 1.0) -> None:
+    return None
+@torch.library.register_fake(add_op_namespace_prefix("nvfp4_gemm_bias_gelu_bf16"))
+def _bias_gelu_fake(a, b, sfa, sfb, bias, out, alpha: float = 1.0) -> None:
+    return None
+@torch.library.register_fake(add_op_namespace_prefix("nvfp4_gemm_bias_gelu_nvfp4"))
+def _bias_gelu_nvfp4_fake(
+    a, b, sfa, sfb, bias, out_packed, out_sfa, alpha: float = 1.0
+) -> None:
+    return None
+@torch.library.register_fake(add_op_namespace_prefix("nvfp4_gemm_streamk_bf16"))
+def _streamk_fake(a, b, sfa, sfb, out, alpha: float = 1.0) -> None:
+    return None
+@torch.library.register_fake(add_op_namespace_prefix("nvfp4_gemm_streamk_bias_bf16"))
+def _streamk_bias_fake(a, b, sfa, sfb, bias, out, alpha: float = 1.0) -> None:
+    return None
+def quantize_fp4_sfa_fp16(
+    x: torch.Tensor,
+    packed: torch.Tensor | None = None,
+    sfa: torch.Tensor | None = None,
+    is_sfb: bool = False,
+):
+    if packed is None or sfa is None:
+        packed, sfa = _alloc_fp4(x.shape[0], x.shape[1], x.device)
+    ops.quantize_fp4_sfa_fp16(x, packed, sfa, bool(is_sfb))
+    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,
+    out: torch.Tensor | None = None,
+    is_sfb: bool = False,
+) -> torch.Tensor:
+    if out is None:
+        out = torch.empty((packed.shape[0], packed.shape[1] * 2), device=packed.device, dtype=torch.float16)
+    ops.dequantize_fp4_sfa_fp16(packed, sfa, out, bool(is_sfb))
+    return out
+def nvfp4_gemm_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 = -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)
+    ops.nvfp4_gemm_bf16(a_packed, b_packed, sfa, sfb, out, float(alpha), int(variant))
+    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,
+    sfa: 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(
+        a_packed, b_packed, sfa, sfb, alpha=alpha, out=out, variant=variant
+    )
+def nvfp4_gemm_bias_bf16(
+    a_packed: torch.Tensor,
+    b_packed: torch.Tensor,
+    sfa: torch.Tensor,
+    sfb: torch.Tensor,
+    bias: torch.Tensor,
+    *,
+    out: torch.Tensor | None = None,
+) -> torch.Tensor:
+    """SM110 NVFP4 GEMM with a fused per-column BF16 bias."""
+    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_bias_bf16(a_packed, b_packed, sfa, sfb, bias, out)
+    return out
+def nvfp4_gemm_bias_residual_bf16(
+    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:
+    """SM110 NVFP4 GEMM with fused BF16 bias and residual add."""
+    if out is None:
+        out = torch.empty_like(residual)
+    ops.nvfp4_gemm_bias_residual_bf16(
+        a_packed, b_packed, sfa, sfb, bias, residual, out
+    )
+    return out
+def nvfp4_gemm_residual_bf16(
+    a_packed: torch.Tensor,
+    b_packed: torch.Tensor,
+    sfa: torch.Tensor,
+    sfb: torch.Tensor,
+    residual: torch.Tensor,
+    alpha: float = 1.0,
+    out: torch.Tensor | None = None,
+) -> torch.Tensor:
+    if out is None:
+        out = torch.empty_like(residual)
+    ops.nvfp4_gemm_residual_bf16(
+        a_packed, b_packed, sfa, sfb, residual, out, float(alpha)
+    )
+    return out
+def nvfp4_gemm_bias_gelu_bf16(
+    a_packed: torch.Tensor,
+    b_packed: torch.Tensor,
+    sfa: torch.Tensor,
+    sfb: torch.Tensor,
+    bias: torch.Tensor,
+    alpha: float = 1.0,
+    out: torch.Tensor | None = None,
+) -> torch.Tensor:
+    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_bias_gelu_bf16(
+        a_packed, b_packed, sfa, sfb, bias, out, float(alpha)
+    )
+    return out
+def nvfp4_gemm_bias_gelu_nvfp4(
+    a_packed: torch.Tensor,
+    b_packed: torch.Tensor,
+    sfa: torch.Tensor,
+    sfb: torch.Tensor,
+    bias: torch.Tensor,
+    alpha: float = 1.0,
+    out_packed: torch.Tensor | None = None,
+    out_sfa: torch.Tensor | None = None,
+) -> tuple[torch.Tensor, torch.Tensor]:
+    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.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)
+    )
+    return out_packed, out_sfa
+def nvfp4_gemm_streamk_bf16(
+    a_packed: torch.Tensor,
+    b_packed: torch.Tensor,
+    sfa: torch.Tensor,
+    sfb: torch.Tensor,
+    alpha: float = 1.0,
+    out: torch.Tensor | None = None,
+) -> torch.Tensor:
+    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_streamk_bf16(
+        a_packed, b_packed, sfa, sfb, out, float(alpha)
+    )
+    return out
+def nvfp4_gemm_streamk_bias_bf16(
+    a_packed: torch.Tensor,
+    b_packed: torch.Tensor,
+    sfa: torch.Tensor,
+    sfb: torch.Tensor,
+    bias: torch.Tensor,
+    alpha: float = 1.0,
+    out: torch.Tensor | None = None,
+) -> torch.Tensor:
+    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_streamk_bias_bf16(
+        a_packed, b_packed, sfa, sfb, bias, out, float(alpha)
+    )
+    return out
+__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",
+]

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

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

+version https://git-lfs.github.com/spec/v1
+oid sha256:479e368c75e9fd00553e206e7038f1e09d0bb37ce2f7bc10c8aaac0801f8348a
+size 1122320

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

	@@ -0,0 +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}"

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

	@@ -0,0 +1,14 @@

+import ctypes
+import importlib.util
+import sys
+from pathlib import Path
+def _import_from_path(file_path: Path):
+    path_hash = '{:x}'.format(ctypes.c_size_t(hash(file_path.absolute())).value)
+    spec = importlib.util.spec_from_file_location(path_hash, file_path)
+    module = importlib.util.module_from_spec(spec)
+    sys.modules[path_hash] = module
+    spec.loader.exec_module(module)
+    return module
+globals().update(vars(_import_from_path(Path(__file__).parent.parent / '__init__.py')))

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

	@@ -0,0 +1,28 @@

+{
+  "name": "fp4-gemm",
+  "id": "_fp4_gemm_cuda_8a66d8b",
+  "version": 1,
+  "license": "Apache-2.0",
+  "python-depends": [],
+  "backend": {
+    "type": "cuda",
+    "archs": [
+      "11.0a"
+    ]
+  },
+  "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
+    }
+  }
+}