liangsu9988 commited on Jun 5

Commit

fcc5ab1

verified ·

1 Parent(s): 9843dc0

Uploaded using `kernel-builder`.

Browse files

Files changed (21) hide show

benchmarks/benchmark_nvfp4_w4a4_decode_matvec.py +204 -0
build/torch211-cxx11-cu128-x86_64-linux/__init__.py +46 -0
build/torch211-cxx11-cu128-x86_64-linux/_flashrt_smallm_gemm_cuda_e9a1fe0.abi3.so +3 -0
build/torch211-cxx11-cu128-x86_64-linux/_ops.py +9 -0
build/torch211-cxx11-cu128-x86_64-linux/flashrt_smallm_gemm/__init__.py +26 -0
build/torch211-cxx11-cu128-x86_64-linux/metadata.json +13 -0
build/torch211-cxx11-cu130-x86_64-linux/__init__.py +46 -0
build/torch211-cxx11-cu130-x86_64-linux/_flashrt_smallm_gemm_cuda_e9a1fe0.abi3.so +3 -0
build/torch211-cxx11-cu130-x86_64-linux/_ops.py +9 -0
build/torch211-cxx11-cu130-x86_64-linux/flashrt_smallm_gemm/__init__.py +26 -0
build/torch211-cxx11-cu130-x86_64-linux/metadata.json +13 -0
build/torch212-cxx11-cu130-x86_64-linux/__init__.py +46 -0
build/torch212-cxx11-cu130-x86_64-linux/_flashrt_smallm_gemm_cuda_e9a1fe0.abi3.so +3 -0
build/torch212-cxx11-cu130-x86_64-linux/_ops.py +9 -0
build/torch212-cxx11-cu130-x86_64-linux/flashrt_smallm_gemm/__init__.py +26 -0
build/torch212-cxx11-cu130-x86_64-linux/metadata.json +13 -0
build/torch212-cxx11-cu132-x86_64-linux/__init__.py +46 -0
build/torch212-cxx11-cu132-x86_64-linux/_flashrt_smallm_gemm_cuda_e9a1fe0.abi3.so +3 -0
build/torch212-cxx11-cu132-x86_64-linux/_ops.py +9 -0
build/torch212-cxx11-cu132-x86_64-linux/flashrt_smallm_gemm/__init__.py +26 -0
build/torch212-cxx11-cu132-x86_64-linux/metadata.json +13 -0

benchmarks/benchmark_nvfp4_w4a4_decode_matvec.py ADDED Viewed

	@@ -0,0 +1,204 @@

+import torch
+from kernels.benchmark import Benchmark
+_original_allclose = torch.allclose
+def _bf16_max_ulp(input: torch.Tensor, other: torch.Tensor) -> int:
+    got_bits = input.detach().cpu().view(torch.int16).to(torch.int32) & 0xFFFF
+    exp_bits = other.detach().cpu().view(torch.int16).to(torch.int32) & 0xFFFF
+    got_ordered = torch.where((got_bits & 0x8000) != 0, 0x8000 - (got_bits & 0x7FFF), got_bits)
+    exp_ordered = torch.where((exp_bits & 0x8000) != 0, 0x8000 - (exp_bits & 0x7FFF), exp_bits)
+    return int((got_ordered - exp_ordered).abs().max().item())
+def _flashrt_allclose(input, other, rtol=1e-05, atol=1e-08, equal_nan=False):
+    if input.dtype == torch.bfloat16 and other.dtype == torch.bfloat16:
+        return _bf16_max_ulp(input, other) <= 5
+    return _original_allclose(input, other, rtol=rtol, atol=atol, equal_nan=equal_nan)
+torch.allclose = _flashrt_allclose
+DECODE_SHAPES = [
+    ("k4096_n1024", 4096, 1024),
+    ("k4096_n4096", 4096, 4096),
+    ("k4096_n12288", 4096, 12288),
+    ("k12288_n1024", 12288, 1024),
+    ("k12288_n4096", 12288, 4096),
+    ("k12288_n12288", 12288, 12288),
+]
+def _swizzled_bytes(rows: int, cols: int) -> int:
+    n_blocks = cols // 16
+    return ((rows + 127) // 128) * ((n_blocks + 3) // 4) * 512
+def _swizzle_constant_scale(rows: int, cols: int, value: int) -> torch.Tensor:
+    return torch.full((_swizzled_bytes(rows, cols),), value, dtype=torch.uint8)
+def _reference_swizzle(scales: torch.Tensor) -> torch.Tensor:
+    rows, n_blocks = scales.shape
+    n_col_super = (n_blocks + 3) // 4
+    src = scales.cpu()
+    out = torch.zeros(
+        ((rows + 127) // 128) * n_col_super * 512,
+        dtype=torch.uint8,
+    )
+    for row in range(rows):
+        rb = row // 128
+        ri = row % 128
+        for block in range(n_blocks):
+            cb = block // 4
+            ci = block % 4
+            super_idx = rb * n_col_super + cb
+            inner_off = (ri % 32) * 16 + (ri // 32) * 4 + ci
+            out[super_idx * 512 + inner_off] = src[row, block]
+    return out
+def _ue4m3_to_float(byte: int) -> float:
+    sign = -1.0 if (byte & 0x80) else 1.0
+    exp = (byte >> 3) & 0x0F
+    mant = byte & 0x07
+    if exp == 0:
+        return sign * (mant / 8.0) * (2.0 ** -6)
+    if exp == 15 and mant == 7:
+        return 0.0
+    return sign * (1.0 + mant / 8.0) * (2.0 ** (exp - 7))
+def _ue4m3_lut() -> torch.Tensor:
+    return torch.tensor([_ue4m3_to_float(i) for i in range(256)], dtype=torch.float32)
+def _fp4_codebook() -> torch.Tensor:
+    return torch.tensor(
+        [
+            0.0,
+            0.5,
+            1.0,
+            1.5,
+            2.0,
+            3.0,
+            4.0,
+            6.0,
+            -0.0,
+            -0.5,
+            -1.0,
+            -1.5,
+            -2.0,
+            -3.0,
+            -4.0,
+            -6.0,
+        ],
+        dtype=torch.float32,
+    )
+def _unpack_fp4(packed: torch.Tensor) -> torch.Tensor:
+    codebook = _fp4_codebook().to(packed.device)
+    lo = packed & 0x0F
+    hi = packed >> 4
+    out = torch.empty(
+        (packed.shape[0], packed.shape[1] * 2),
+        device=packed.device,
+        dtype=torch.float32,
+    )
+    out[:, 0::2] = codebook[lo.long()]
+    out[:, 1::2] = codebook[hi.long()]
+    return out
+def _reference_smallm(
+    a_packed: torch.Tensor,
+    b_packed: torch.Tensor,
+    sfa_linear: torch.Tensor,
+    sfb_linear: torch.Tensor,
+    K: int,
+    alpha: float,
+    chunk_rows: int = 256,
+) -> torch.Tensor:
+    device = b_packed.device
+    N = b_packed.shape[0]
+    lut = _ue4m3_lut().to(device)
+    a = _unpack_fp4(a_packed.reshape(1, -1)).reshape(K)
+    a_scale = lut[sfa_linear.reshape(-1).to(device).long()].repeat_interleave(16)
+    a = a * a_scale
+    sfb_linear = sfb_linear.to(device)
+    out = torch.empty((N,), device=device, dtype=torch.bfloat16)
+    for start in range(0, N, chunk_rows):
+        end = min(start + chunk_rows, N)
+        b = _unpack_fp4(b_packed[start:end])
+        b_scale = lut[sfb_linear[start:end].long()].repeat_interleave(16, dim=1)
+        expected = (b * b_scale * a.reshape(1, K)).sum(dim=1) * alpha
+        out[start:end] = expected.to(torch.bfloat16)
+    return out
+class Nvfp4W4A4DecodeMatvecBenchmark(Benchmark):
+    seed = 23
+    def _setup_shape(self, K: int, N: int) -> None:
+        torch.manual_seed(600 + K + N)
+        if torch.cuda.is_available():
+            torch.cuda.manual_seed_all(600 + K + N)
+        self.K = K
+        self.N = N
+        self.alpha = 0.5
+        self.a_packed = torch.randint(
+            0, 256, (K // 2,), device=self.device, dtype=torch.uint8
+        )
+        self.b_packed = torch.randint(
+            0, 256, (N, K // 2), device=self.device, dtype=torch.uint8
+        )
+        self.sfa_linear = torch.randint(0, 0x78, (1, K // 16), dtype=torch.uint8)
+        self.sfb_linear = torch.randint(0, 0x78, (N, K // 16), dtype=torch.uint8)
+        self.sfa = _reference_swizzle(self.sfa_linear).to(self.device)
+        self.sfb = _reference_swizzle(self.sfb_linear).to(self.device)
+        self.out = torch.empty((N,), device=self.device, dtype=torch.bfloat16)
+    def _benchmark(self) -> None:
+        self.kernel.nvfp4_w4a4_decode_matvec_bf16out(
+            self.a_packed,
+            self.b_packed,
+            self.sfa,
+            self.sfb,
+            alpha=self.alpha,
+            out=self.out,
+        )
+    def _reference(self) -> torch.Tensor:
+        return _reference_smallm(
+            self.a_packed,
+            self.b_packed,
+            self.sfa_linear,
+            self.sfb_linear,
+            self.K,
+            self.alpha,
+        )
+def _register_shapes() -> None:
+    for label, K, N in DECODE_SHAPES:
+        def setup(self, K=K, N=N) -> None:
+            self._setup_shape(K, N)
+        def benchmark(self) -> None:
+            self._benchmark()
+        def verify(self) -> torch.Tensor:
+            return self._reference()
+        setattr(Nvfp4W4A4DecodeMatvecBenchmark, f"setup_{label}", setup)
+        setattr(Nvfp4W4A4DecodeMatvecBenchmark, f"benchmark_{label}", benchmark)
+        setattr(Nvfp4W4A4DecodeMatvecBenchmark, f"verify_{label}", verify)
+_register_shapes()

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

	@@ -0,0 +1,46 @@

+"""FlashRT small-M GEMM kernels."""
+from __future__ import annotations
+from typing import Optional
+import torch
+from ._ops import ops
+def nvfp4_w4a4_decode_matvec_bf16out(
+    a_packed: torch.Tensor,
+    b_packed: torch.Tensor,
+    sfa: torch.Tensor,
+    sfb: torch.Tensor,
+    *,
+    alpha: float = 1.0,
+    out: Optional[torch.Tensor] = None,
+) -> torch.Tensor:
+    """Compute an SM120 NVFP4 W4A4 M=1 matvec with BF16 output.
+    ``a_packed`` contains one packed activation row with shape ``(K / 2,)`` or
+    ``(1, K / 2)``. ``b_packed`` is row-major with shape ``(N, K / 2)``. ``sfa``
+    and ``sfb`` are CUTLASS Sm1xx swizzled UE4M3 scale-factor byte buffers.
+    The current kernel supports ``K in {4096, 12288}``.
+    """
+    if b_packed.dim() != 2:
+        raise ValueError("b_packed must have shape (N, K / 2)")
+    if out is None:
+        out = torch.empty((b_packed.shape[0],), device=b_packed.device, dtype=torch.bfloat16)
+    ops.nvfp4_w4a4_decode_matvec_bf16out(
+        a_packed,
+        b_packed,
+        sfa,
+        sfb,
+        out,
+        float(alpha),
+    )
+    return out
+__all__ = [
+    "nvfp4_w4a4_decode_matvec_bf16out",
+]

build/torch211-cxx11-cu128-x86_64-linux/_flashrt_smallm_gemm_cuda_e9a1fe0.abi3.so ADDED Viewed

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

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

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

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

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

build/torch211-cxx11-cu128-x86_64-linux/flashrt_smallm_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-cu128-x86_64-linux/metadata.json ADDED Viewed

	@@ -0,0 +1,13 @@

+{
+  "name": "flashrt-smallm-gemm",
+  "id": "_flashrt_smallm_gemm_cuda_e9a1fe0",
+  "version": 1,
+  "license": "Apache-2.0",
+  "python-depends": [],
+  "backend": {
+    "type": "cuda",
+    "archs": [
+      "12.0"
+    ]
+  }
+}

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

	@@ -0,0 +1,46 @@

+"""FlashRT small-M GEMM kernels."""
+from __future__ import annotations
+from typing import Optional
+import torch
+from ._ops import ops
+def nvfp4_w4a4_decode_matvec_bf16out(
+    a_packed: torch.Tensor,
+    b_packed: torch.Tensor,
+    sfa: torch.Tensor,
+    sfb: torch.Tensor,
+    *,
+    alpha: float = 1.0,
+    out: Optional[torch.Tensor] = None,
+) -> torch.Tensor:
+    """Compute an SM120 NVFP4 W4A4 M=1 matvec with BF16 output.
+    ``a_packed`` contains one packed activation row with shape ``(K / 2,)`` or
+    ``(1, K / 2)``. ``b_packed`` is row-major with shape ``(N, K / 2)``. ``sfa``
+    and ``sfb`` are CUTLASS Sm1xx swizzled UE4M3 scale-factor byte buffers.
+    The current kernel supports ``K in {4096, 12288}``.
+    """
+    if b_packed.dim() != 2:
+        raise ValueError("b_packed must have shape (N, K / 2)")
+    if out is None:
+        out = torch.empty((b_packed.shape[0],), device=b_packed.device, dtype=torch.bfloat16)
+    ops.nvfp4_w4a4_decode_matvec_bf16out(
+        a_packed,
+        b_packed,
+        sfa,
+        sfb,
+        out,
+        float(alpha),
+    )
+    return out
+__all__ = [
+    "nvfp4_w4a4_decode_matvec_bf16out",
+]

build/torch211-cxx11-cu130-x86_64-linux/_flashrt_smallm_gemm_cuda_e9a1fe0.abi3.so ADDED Viewed

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

+version https://git-lfs.github.com/spec/v1
+oid sha256:8b9a7f6a78cb3784d9efa7cde5349c220f183799254f23326aba75cd1accee4e
+size 122624

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

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

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

build/torch211-cxx11-cu130-x86_64-linux/flashrt_smallm_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-x86_64-linux/metadata.json ADDED Viewed

	@@ -0,0 +1,13 @@

+{
+  "name": "flashrt-smallm-gemm",
+  "id": "_flashrt_smallm_gemm_cuda_e9a1fe0",
+  "version": 1,
+  "license": "Apache-2.0",
+  "python-depends": [],
+  "backend": {
+    "type": "cuda",
+    "archs": [
+      "12.0"
+    ]
+  }
+}

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

	@@ -0,0 +1,46 @@

+"""FlashRT small-M GEMM kernels."""
+from __future__ import annotations
+from typing import Optional
+import torch
+from ._ops import ops
+def nvfp4_w4a4_decode_matvec_bf16out(
+    a_packed: torch.Tensor,
+    b_packed: torch.Tensor,
+    sfa: torch.Tensor,
+    sfb: torch.Tensor,
+    *,
+    alpha: float = 1.0,
+    out: Optional[torch.Tensor] = None,
+) -> torch.Tensor:
+    """Compute an SM120 NVFP4 W4A4 M=1 matvec with BF16 output.
+    ``a_packed`` contains one packed activation row with shape ``(K / 2,)`` or
+    ``(1, K / 2)``. ``b_packed`` is row-major with shape ``(N, K / 2)``. ``sfa``
+    and ``sfb`` are CUTLASS Sm1xx swizzled UE4M3 scale-factor byte buffers.
+    The current kernel supports ``K in {4096, 12288}``.
+    """
+    if b_packed.dim() != 2:
+        raise ValueError("b_packed must have shape (N, K / 2)")
+    if out is None:
+        out = torch.empty((b_packed.shape[0],), device=b_packed.device, dtype=torch.bfloat16)
+    ops.nvfp4_w4a4_decode_matvec_bf16out(
+        a_packed,
+        b_packed,
+        sfa,
+        sfb,
+        out,
+        float(alpha),
+    )
+    return out
+__all__ = [
+    "nvfp4_w4a4_decode_matvec_bf16out",
+]

build/torch212-cxx11-cu130-x86_64-linux/_flashrt_smallm_gemm_cuda_e9a1fe0.abi3.so ADDED Viewed

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

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

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

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

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

build/torch212-cxx11-cu130-x86_64-linux/flashrt_smallm_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-x86_64-linux/metadata.json ADDED Viewed

	@@ -0,0 +1,13 @@

+{
+  "name": "flashrt-smallm-gemm",
+  "id": "_flashrt_smallm_gemm_cuda_e9a1fe0",
+  "version": 1,
+  "license": "Apache-2.0",
+  "python-depends": [],
+  "backend": {
+    "type": "cuda",
+    "archs": [
+      "12.0"
+    ]
+  }
+}

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

	@@ -0,0 +1,46 @@

+"""FlashRT small-M GEMM kernels."""
+from __future__ import annotations
+from typing import Optional
+import torch
+from ._ops import ops
+def nvfp4_w4a4_decode_matvec_bf16out(
+    a_packed: torch.Tensor,
+    b_packed: torch.Tensor,
+    sfa: torch.Tensor,
+    sfb: torch.Tensor,
+    *,
+    alpha: float = 1.0,
+    out: Optional[torch.Tensor] = None,
+) -> torch.Tensor:
+    """Compute an SM120 NVFP4 W4A4 M=1 matvec with BF16 output.
+    ``a_packed`` contains one packed activation row with shape ``(K / 2,)`` or
+    ``(1, K / 2)``. ``b_packed`` is row-major with shape ``(N, K / 2)``. ``sfa``
+    and ``sfb`` are CUTLASS Sm1xx swizzled UE4M3 scale-factor byte buffers.
+    The current kernel supports ``K in {4096, 12288}``.
+    """
+    if b_packed.dim() != 2:
+        raise ValueError("b_packed must have shape (N, K / 2)")
+    if out is None:
+        out = torch.empty((b_packed.shape[0],), device=b_packed.device, dtype=torch.bfloat16)
+    ops.nvfp4_w4a4_decode_matvec_bf16out(
+        a_packed,
+        b_packed,
+        sfa,
+        sfb,
+        out,
+        float(alpha),
+    )
+    return out
+__all__ = [
+    "nvfp4_w4a4_decode_matvec_bf16out",
+]

build/torch212-cxx11-cu132-x86_64-linux/_flashrt_smallm_gemm_cuda_e9a1fe0.abi3.so ADDED Viewed

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

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

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

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

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

build/torch212-cxx11-cu132-x86_64-linux/flashrt_smallm_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-x86_64-linux/metadata.json ADDED Viewed

	@@ -0,0 +1,13 @@

+{
+  "name": "flashrt-smallm-gemm",
+  "id": "_flashrt_smallm_gemm_cuda_e9a1fe0",
+  "version": 1,
+  "license": "Apache-2.0",
+  "python-depends": [],
+  "backend": {
+    "type": "cuda",
+    "archs": [
+      "12.0"
+    ]
+  }
+}