nologik commited on 3 days ago

Commit

b43edbb

verified ·

1 Parent(s): 69bd8ab

Uploaded using `kernel-builder`.

Browse files

Files changed (20) hide show

build/torch210-cxx11-cu130-aarch64-linux/__init__.py +57 -0
build/torch210-cxx11-cu130-aarch64-linux/_nvfp4_dense_gemm_cuda_bcba8ad.abi3.so +3 -0
build/torch210-cxx11-cu130-aarch64-linux/_ops.py +9 -0
build/torch210-cxx11-cu130-aarch64-linux/metadata.json +10 -0
build/torch210-cxx11-cu130-aarch64-linux/nvfp4_dense_gemm/__init__.py +26 -0
build/torch211-cxx11-cu130-aarch64-linux/__init__.py +57 -0
build/torch211-cxx11-cu130-aarch64-linux/_nvfp4_dense_gemm_cuda_bcba8ad.abi3.so +3 -0
build/torch211-cxx11-cu130-aarch64-linux/_ops.py +9 -0
build/torch211-cxx11-cu130-aarch64-linux/metadata.json +10 -0
build/torch211-cxx11-cu130-aarch64-linux/nvfp4_dense_gemm/__init__.py +26 -0
build/torch212-cxx11-cu130-aarch64-linux/__init__.py +57 -0
build/torch212-cxx11-cu130-aarch64-linux/_nvfp4_dense_gemm_cuda_bcba8ad.abi3.so +3 -0
build/torch212-cxx11-cu130-aarch64-linux/_ops.py +9 -0
build/torch212-cxx11-cu130-aarch64-linux/metadata.json +10 -0
build/torch212-cxx11-cu130-aarch64-linux/nvfp4_dense_gemm/__init__.py +26 -0
build/torch212-cxx11-cu132-aarch64-linux/__init__.py +57 -0
build/torch212-cxx11-cu132-aarch64-linux/_nvfp4_dense_gemm_cuda_bcba8ad.abi3.so +3 -0
build/torch212-cxx11-cu132-aarch64-linux/_ops.py +9 -0
build/torch212-cxx11-cu132-aarch64-linux/metadata.json +10 -0
build/torch212-cxx11-cu132-aarch64-linux/nvfp4_dense_gemm/__init__.py +26 -0

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

	@@ -0,0 +1,57 @@

+# SPDX-License-Identifier: AGPL-3.0-only
+"""Atlas dense NVFP4/FP8 GEMM kernels for Qwen3.6-27B on GB10 (SM121).
+Powers the attention Q/K/V/O projections and the dense FFN
+(intermediate_size=17408, hidden=5120). Two precisions are supported:
+  - W4A16 (NVFP4 weight, BF16 activation): direct path
+  - FP8 W/A: via ``predequant_nvfp4_to_fp8`` + ``fp8_gemm_t`` for
+    KV-cache scenarios where the weight is loaded as NVFP4 but the
+    GEMM runs in FP8 against an FP8-packed activation.
+"""
+import torch
+from ._ops import ops
+__all__ = [
+    "w4a16_gemm",
+    "w4a16_gemm_t",
+    "predequant_nvfp4_to_fp8",
+    "fp8_gemm_t",
+    "bf16_to_fp8",
+]
+def w4a16_gemm(
+    A: torch.Tensor, B_packed: torch.Tensor, B_scale: torch.Tensor,
+    scale2: float, C: torch.Tensor,
+) -> None:
+    """C = A @ B^T where B is NVFP4. Standard layout."""
+    ops.w4a16_gemm(A, B_packed, B_scale, scale2, C)
+def w4a16_gemm_t(
+    A: torch.Tensor, B_packed: torch.Tensor, B_scale: torch.Tensor,
+    scale2: float, C: torch.Tensor,
+) -> None:
+    """Transposed-B variant (N-major) — typically faster for [N>K] shapes."""
+    ops.w4a16_gemm_t(A, B_packed, B_scale, scale2, C)
+def predequant_nvfp4_to_fp8(
+    B_packed: torch.Tensor, B_scale: torch.Tensor,
+    scale2: float, B_fp8_out: torch.Tensor,
+) -> None:
+    """Materialize an NVFP4-packed weight into a dense FP8 E4M3 tensor."""
+    ops.predequant_nvfp4_to_fp8(B_packed, B_scale, scale2, B_fp8_out)
+def fp8_gemm_t(A: torch.Tensor, B_fp8: torch.Tensor, C: torch.Tensor) -> None:
+    """C = A @ B^T where both A is BF16 and B is FP8 E4M3."""
+    ops.fp8_gemm_t(A, B_fp8, C)
+def bf16_to_fp8(src: torch.Tensor, dst: torch.Tensor) -> None:
+    """Pair-wise BF16 → FP8 E4M3 quantization (no scaling, uses PTX cvt)."""
+    ops.bf16_to_fp8(src, dst)

build/torch210-cxx11-cu130-aarch64-linux/_nvfp4_dense_gemm_cuda_bcba8ad.abi3.so ADDED Viewed

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

+version https://git-lfs.github.com/spec/v1
+oid sha256:6e1407400342d579a6e86fbe821bcd506fda3ad0216671cf2101e9dcbe49f434
+size 2505208

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

	@@ -0,0 +1,9 @@

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

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

	@@ -0,0 +1,10 @@

+{
+  "name": "nvfp4-dense-gemm",
+  "id": "_nvfp4_dense_gemm_cuda_bcba8ad",
+  "version": 0,
+  "license": "AGPL-3.0-only",
+  "python-depends": [],
+  "backend": {
+    "type": "cuda"
+  }
+}

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

	@@ -0,0 +1,26 @@

+import ctypes
+import importlib.util
+import sys
+from pathlib import Path
+from types import ModuleType
+def _import_from_path(file_path: Path) -> ModuleType:
+    # We cannot use the module name as-is, after adding it to `sys.modules`,
+    # it would also be used for other imports. So, we make a module name that
+    # depends on the path for it to be unique using the hex-encoded hash of
+    # the path.
+    path_hash = "{:x}".format(ctypes.c_size_t(hash(file_path.absolute())).value)
+    module_name = path_hash
+    spec = importlib.util.spec_from_file_location(module_name, file_path)
+    if spec is None:
+        raise ImportError(f"Cannot load spec for {module_name} from {file_path}")
+    module = importlib.util.module_from_spec(spec)
+    if module is None:
+        raise ImportError(f"Cannot load module {module_name} from spec")
+    sys.modules[module_name] = module
+    spec.loader.exec_module(module)  # type: ignore
+    return module
+globals().update(vars(_import_from_path(Path(__file__).parent.parent / "__init__.py")))

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

	@@ -0,0 +1,57 @@

+# SPDX-License-Identifier: AGPL-3.0-only
+"""Atlas dense NVFP4/FP8 GEMM kernels for Qwen3.6-27B on GB10 (SM121).
+Powers the attention Q/K/V/O projections and the dense FFN
+(intermediate_size=17408, hidden=5120). Two precisions are supported:
+  - W4A16 (NVFP4 weight, BF16 activation): direct path
+  - FP8 W/A: via ``predequant_nvfp4_to_fp8`` + ``fp8_gemm_t`` for
+    KV-cache scenarios where the weight is loaded as NVFP4 but the
+    GEMM runs in FP8 against an FP8-packed activation.
+"""
+import torch
+from ._ops import ops
+__all__ = [
+    "w4a16_gemm",
+    "w4a16_gemm_t",
+    "predequant_nvfp4_to_fp8",
+    "fp8_gemm_t",
+    "bf16_to_fp8",
+]
+def w4a16_gemm(
+    A: torch.Tensor, B_packed: torch.Tensor, B_scale: torch.Tensor,
+    scale2: float, C: torch.Tensor,
+) -> None:
+    """C = A @ B^T where B is NVFP4. Standard layout."""
+    ops.w4a16_gemm(A, B_packed, B_scale, scale2, C)
+def w4a16_gemm_t(
+    A: torch.Tensor, B_packed: torch.Tensor, B_scale: torch.Tensor,
+    scale2: float, C: torch.Tensor,
+) -> None:
+    """Transposed-B variant (N-major) — typically faster for [N>K] shapes."""
+    ops.w4a16_gemm_t(A, B_packed, B_scale, scale2, C)
+def predequant_nvfp4_to_fp8(
+    B_packed: torch.Tensor, B_scale: torch.Tensor,
+    scale2: float, B_fp8_out: torch.Tensor,
+) -> None:
+    """Materialize an NVFP4-packed weight into a dense FP8 E4M3 tensor."""
+    ops.predequant_nvfp4_to_fp8(B_packed, B_scale, scale2, B_fp8_out)
+def fp8_gemm_t(A: torch.Tensor, B_fp8: torch.Tensor, C: torch.Tensor) -> None:
+    """C = A @ B^T where both A is BF16 and B is FP8 E4M3."""
+    ops.fp8_gemm_t(A, B_fp8, C)
+def bf16_to_fp8(src: torch.Tensor, dst: torch.Tensor) -> None:
+    """Pair-wise BF16 → FP8 E4M3 quantization (no scaling, uses PTX cvt)."""
+    ops.bf16_to_fp8(src, dst)

build/torch211-cxx11-cu130-aarch64-linux/_nvfp4_dense_gemm_cuda_bcba8ad.abi3.so ADDED Viewed

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

+version https://git-lfs.github.com/spec/v1
+oid sha256:3d646b4764445cdd173802d7c5074664bac8ce4c97df5124075f7e44ac51c808
+size 2505208

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

	@@ -0,0 +1,9 @@

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

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

	@@ -0,0 +1,10 @@

+{
+  "name": "nvfp4-dense-gemm",
+  "id": "_nvfp4_dense_gemm_cuda_bcba8ad",
+  "version": 0,
+  "license": "AGPL-3.0-only",
+  "python-depends": [],
+  "backend": {
+    "type": "cuda"
+  }
+}

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

	@@ -0,0 +1,26 @@

+import ctypes
+import importlib.util
+import sys
+from pathlib import Path
+from types import ModuleType
+def _import_from_path(file_path: Path) -> ModuleType:
+    # We cannot use the module name as-is, after adding it to `sys.modules`,
+    # it would also be used for other imports. So, we make a module name that
+    # depends on the path for it to be unique using the hex-encoded hash of
+    # the path.
+    path_hash = "{:x}".format(ctypes.c_size_t(hash(file_path.absolute())).value)
+    module_name = path_hash
+    spec = importlib.util.spec_from_file_location(module_name, file_path)
+    if spec is None:
+        raise ImportError(f"Cannot load spec for {module_name} from {file_path}")
+    module = importlib.util.module_from_spec(spec)
+    if module is None:
+        raise ImportError(f"Cannot load module {module_name} from spec")
+    sys.modules[module_name] = module
+    spec.loader.exec_module(module)  # type: ignore
+    return module
+globals().update(vars(_import_from_path(Path(__file__).parent.parent / "__init__.py")))

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

	@@ -0,0 +1,57 @@

+# SPDX-License-Identifier: AGPL-3.0-only
+"""Atlas dense NVFP4/FP8 GEMM kernels for Qwen3.6-27B on GB10 (SM121).
+Powers the attention Q/K/V/O projections and the dense FFN
+(intermediate_size=17408, hidden=5120). Two precisions are supported:
+  - W4A16 (NVFP4 weight, BF16 activation): direct path
+  - FP8 W/A: via ``predequant_nvfp4_to_fp8`` + ``fp8_gemm_t`` for
+    KV-cache scenarios where the weight is loaded as NVFP4 but the
+    GEMM runs in FP8 against an FP8-packed activation.
+"""
+import torch
+from ._ops import ops
+__all__ = [
+    "w4a16_gemm",
+    "w4a16_gemm_t",
+    "predequant_nvfp4_to_fp8",
+    "fp8_gemm_t",
+    "bf16_to_fp8",
+]
+def w4a16_gemm(
+    A: torch.Tensor, B_packed: torch.Tensor, B_scale: torch.Tensor,
+    scale2: float, C: torch.Tensor,
+) -> None:
+    """C = A @ B^T where B is NVFP4. Standard layout."""
+    ops.w4a16_gemm(A, B_packed, B_scale, scale2, C)
+def w4a16_gemm_t(
+    A: torch.Tensor, B_packed: torch.Tensor, B_scale: torch.Tensor,
+    scale2: float, C: torch.Tensor,
+) -> None:
+    """Transposed-B variant (N-major) — typically faster for [N>K] shapes."""
+    ops.w4a16_gemm_t(A, B_packed, B_scale, scale2, C)
+def predequant_nvfp4_to_fp8(
+    B_packed: torch.Tensor, B_scale: torch.Tensor,
+    scale2: float, B_fp8_out: torch.Tensor,
+) -> None:
+    """Materialize an NVFP4-packed weight into a dense FP8 E4M3 tensor."""
+    ops.predequant_nvfp4_to_fp8(B_packed, B_scale, scale2, B_fp8_out)
+def fp8_gemm_t(A: torch.Tensor, B_fp8: torch.Tensor, C: torch.Tensor) -> None:
+    """C = A @ B^T where both A is BF16 and B is FP8 E4M3."""
+    ops.fp8_gemm_t(A, B_fp8, C)
+def bf16_to_fp8(src: torch.Tensor, dst: torch.Tensor) -> None:
+    """Pair-wise BF16 → FP8 E4M3 quantization (no scaling, uses PTX cvt)."""
+    ops.bf16_to_fp8(src, dst)

build/torch212-cxx11-cu130-aarch64-linux/_nvfp4_dense_gemm_cuda_bcba8ad.abi3.so ADDED Viewed

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

+version https://git-lfs.github.com/spec/v1
+oid sha256:7f3695fa1e52677900262e486733b0454ce8312ec1ab834a2026237097733bfb
+size 2505216

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

	@@ -0,0 +1,9 @@

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

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

	@@ -0,0 +1,10 @@

+{
+  "name": "nvfp4-dense-gemm",
+  "id": "_nvfp4_dense_gemm_cuda_bcba8ad",
+  "version": 0,
+  "license": "AGPL-3.0-only",
+  "python-depends": [],
+  "backend": {
+    "type": "cuda"
+  }
+}

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

	@@ -0,0 +1,26 @@

+import ctypes
+import importlib.util
+import sys
+from pathlib import Path
+from types import ModuleType
+def _import_from_path(file_path: Path) -> ModuleType:
+    # We cannot use the module name as-is, after adding it to `sys.modules`,
+    # it would also be used for other imports. So, we make a module name that
+    # depends on the path for it to be unique using the hex-encoded hash of
+    # the path.
+    path_hash = "{:x}".format(ctypes.c_size_t(hash(file_path.absolute())).value)
+    module_name = path_hash
+    spec = importlib.util.spec_from_file_location(module_name, file_path)
+    if spec is None:
+        raise ImportError(f"Cannot load spec for {module_name} from {file_path}")
+    module = importlib.util.module_from_spec(spec)
+    if module is None:
+        raise ImportError(f"Cannot load module {module_name} from spec")
+    sys.modules[module_name] = module
+    spec.loader.exec_module(module)  # type: ignore
+    return module
+globals().update(vars(_import_from_path(Path(__file__).parent.parent / "__init__.py")))

build/torch212-cxx11-cu132-aarch64-linux/__init__.py ADDED Viewed

	@@ -0,0 +1,57 @@

+# SPDX-License-Identifier: AGPL-3.0-only
+"""Atlas dense NVFP4/FP8 GEMM kernels for Qwen3.6-27B on GB10 (SM121).
+Powers the attention Q/K/V/O projections and the dense FFN
+(intermediate_size=17408, hidden=5120). Two precisions are supported:
+  - W4A16 (NVFP4 weight, BF16 activation): direct path
+  - FP8 W/A: via ``predequant_nvfp4_to_fp8`` + ``fp8_gemm_t`` for
+    KV-cache scenarios where the weight is loaded as NVFP4 but the
+    GEMM runs in FP8 against an FP8-packed activation.
+"""
+import torch
+from ._ops import ops
+__all__ = [
+    "w4a16_gemm",
+    "w4a16_gemm_t",
+    "predequant_nvfp4_to_fp8",
+    "fp8_gemm_t",
+    "bf16_to_fp8",
+]
+def w4a16_gemm(
+    A: torch.Tensor, B_packed: torch.Tensor, B_scale: torch.Tensor,
+    scale2: float, C: torch.Tensor,
+) -> None:
+    """C = A @ B^T where B is NVFP4. Standard layout."""
+    ops.w4a16_gemm(A, B_packed, B_scale, scale2, C)
+def w4a16_gemm_t(
+    A: torch.Tensor, B_packed: torch.Tensor, B_scale: torch.Tensor,
+    scale2: float, C: torch.Tensor,
+) -> None:
+    """Transposed-B variant (N-major) — typically faster for [N>K] shapes."""
+    ops.w4a16_gemm_t(A, B_packed, B_scale, scale2, C)
+def predequant_nvfp4_to_fp8(
+    B_packed: torch.Tensor, B_scale: torch.Tensor,
+    scale2: float, B_fp8_out: torch.Tensor,
+) -> None:
+    """Materialize an NVFP4-packed weight into a dense FP8 E4M3 tensor."""
+    ops.predequant_nvfp4_to_fp8(B_packed, B_scale, scale2, B_fp8_out)
+def fp8_gemm_t(A: torch.Tensor, B_fp8: torch.Tensor, C: torch.Tensor) -> None:
+    """C = A @ B^T where both A is BF16 and B is FP8 E4M3."""
+    ops.fp8_gemm_t(A, B_fp8, C)
+def bf16_to_fp8(src: torch.Tensor, dst: torch.Tensor) -> None:
+    """Pair-wise BF16 → FP8 E4M3 quantization (no scaling, uses PTX cvt)."""
+    ops.bf16_to_fp8(src, dst)

build/torch212-cxx11-cu132-aarch64-linux/_nvfp4_dense_gemm_cuda_bcba8ad.abi3.so ADDED Viewed

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

+version https://git-lfs.github.com/spec/v1
+oid sha256:286187ea055712565459c5316a345e1f76fb4ef2027c98bdaf20a316cab1ab55
+size 2505216

build/torch212-cxx11-cu132-aarch64-linux/_ops.py ADDED Viewed

	@@ -0,0 +1,9 @@

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

build/torch212-cxx11-cu132-aarch64-linux/metadata.json ADDED Viewed

	@@ -0,0 +1,10 @@

+{
+  "name": "nvfp4-dense-gemm",
+  "id": "_nvfp4_dense_gemm_cuda_bcba8ad",
+  "version": 0,
+  "license": "AGPL-3.0-only",
+  "python-depends": [],
+  "backend": {
+    "type": "cuda"
+  }
+}

build/torch212-cxx11-cu132-aarch64-linux/nvfp4_dense_gemm/__init__.py ADDED Viewed

	@@ -0,0 +1,26 @@

+import ctypes
+import importlib.util
+import sys
+from pathlib import Path
+from types import ModuleType
+def _import_from_path(file_path: Path) -> ModuleType:
+    # We cannot use the module name as-is, after adding it to `sys.modules`,
+    # it would also be used for other imports. So, we make a module name that
+    # depends on the path for it to be unique using the hex-encoded hash of
+    # the path.
+    path_hash = "{:x}".format(ctypes.c_size_t(hash(file_path.absolute())).value)
+    module_name = path_hash
+    spec = importlib.util.spec_from_file_location(module_name, file_path)
+    if spec is None:
+        raise ImportError(f"Cannot load spec for {module_name} from {file_path}")
+    module = importlib.util.module_from_spec(spec)
+    if module is None:
+        raise ImportError(f"Cannot load module {module_name} from spec")
+    sys.modules[module_name] = module
+    spec.loader.exec_module(module)  # type: ignore
+    return module
+globals().update(vars(_import_from_path(Path(__file__).parent.parent / "__init__.py")))