liangsu9988 commited on Jun 20

Commit

4fb53e5

verified ·

1 Parent(s): 755a1f2

Uploaded using `kernel-builder`.

Browse files

Files changed (21) hide show

benchmarks/benchmark.py +140 -0
build/torch211-cxx11-cu128-x86_64-linux/__init__.py +86 -0
build/torch211-cxx11-cu128-x86_64-linux/_fp4_gemm_cuda_d8a589a.abi3.so +3 -0
build/torch211-cxx11-cu128-x86_64-linux/_ops.py +9 -0
build/torch211-cxx11-cu128-x86_64-linux/fp4_gemm/__init__.py +26 -0
build/torch211-cxx11-cu128-x86_64-linux/metadata.json +22 -0
build/torch211-cxx11-cu130-x86_64-linux/__init__.py +86 -0
build/torch211-cxx11-cu130-x86_64-linux/_fp4_gemm_cuda_d8a589a.abi3.so +3 -0
build/torch211-cxx11-cu130-x86_64-linux/_ops.py +9 -0
build/torch211-cxx11-cu130-x86_64-linux/fp4_gemm/__init__.py +26 -0
build/torch211-cxx11-cu130-x86_64-linux/metadata.json +22 -0
build/torch212-cxx11-cu130-x86_64-linux/__init__.py +86 -0
build/torch212-cxx11-cu130-x86_64-linux/_fp4_gemm_cuda_d8a589a.abi3.so +3 -0
build/torch212-cxx11-cu130-x86_64-linux/_ops.py +9 -0
build/torch212-cxx11-cu130-x86_64-linux/fp4_gemm/__init__.py +26 -0
build/torch212-cxx11-cu130-x86_64-linux/metadata.json +22 -0
build/torch212-cxx11-cu132-x86_64-linux/__init__.py +86 -0
build/torch212-cxx11-cu132-x86_64-linux/_fp4_gemm_cuda_d8a589a.abi3.so +3 -0
build/torch212-cxx11-cu132-x86_64-linux/_ops.py +9 -0
build/torch212-cxx11-cu132-x86_64-linux/fp4_gemm/__init__.py +26 -0
build/torch212-cxx11-cu132-x86_64-linux/metadata.json +22 -0

benchmarks/benchmark.py ADDED Viewed

	@@ -0,0 +1,140 @@

+#!/usr/bin/env python3
+"""Benchmark fp4-gemm."""
+from __future__ import annotations
+import argparse
+import importlib.util
+import json
+import sys
+from dataclasses import asdict, dataclass
+from pathlib import Path
+import torch
+ROOT = Path(__file__).resolve().parents[2]
+TEST_FILE = ROOT / "fp4-gemm" / "tests" / "test_fp4_gemm.py"
+@dataclass
+class BenchResult:
+    shape: str
+    M: int
+    N: int
+    K: int
+    variant: int
+    flashrt_us: float
+    torch_reference_us: float
+    speedup_vs_reference: float
+    max_abs: float
+    mean_abs: float
+    p99_abs: float
+    cosine: float
+    status: str
+def load_helpers():
+    spec = importlib.util.spec_from_file_location("fp4_gemm_test_helpers", TEST_FILE)
+    if spec is None or spec.loader is None:
+        raise RuntimeError(f"cannot load helpers from {TEST_FILE}")
+    module = importlib.util.module_from_spec(spec)
+    sys.modules["fp4_gemm_test_helpers"] = module
+    spec.loader.exec_module(module)
+    return module
+def measure(fn, warmup: int, iters: int) -> float:
+    for _ in range(warmup):
+        fn()
+    torch.cuda.synchronize()
+    start = torch.cuda.Event(enable_timing=True)
+    end = torch.cuda.Event(enable_timing=True)
+    start.record()
+    for _ in range(iters):
+        fn()
+    end.record()
+    torch.cuda.synchronize()
+    return float(start.elapsed_time(end) * 1000.0 / iters)
+def bench_case(helpers, ops, name: str, shape: tuple[int, int, int], warmup: int, iters: int) -> list[BenchResult]:
+    m, n, k = shape
+    a_packed, b_packed, sfa, sfb, expected = helpers.prepare_quantized(ops, m, n, k)
+    a_deq = torch.empty((m, k), device="cuda", dtype=torch.float16)
+    b_deq = torch.empty((n, k), device="cuda", dtype=torch.float16)
+    ops.dequantize_fp4_sfa_fp16(a_packed, sfa, a_deq, False)
+    ops.dequantize_fp4_sfa_fp16(b_packed, sfb, b_deq, True)
+    torch.cuda.synchronize()
+    def torch_ref():
+        return (a_deq.float() @ b_deq.float().T).to(torch.bfloat16)
+    torch_us = measure(torch_ref, warmup, iters)
+    results: list[BenchResult] = []
+    for variant in (0, 1, 2):
+        out = torch.empty((m, n), device="cuda", dtype=torch.bfloat16)
+        ops.fp4_w4a16_linear_bf16(a_packed, b_packed, sfa, sfb, out, 1.0, variant)
+        torch.cuda.synchronize()
+        max_abs, mean_abs, p99_abs, cosine = helpers.metrics(out, expected)
+        flashrt_us = measure(
+            lambda: ops.fp4_w4a16_linear_bf16(a_packed, b_packed, sfa, sfb, out, 1.0, variant),
+            warmup,
+            iters,
+        )
+        results.append(
+            BenchResult(
+                shape=name,
+                M=m,
+                N=n,
+                K=k,
+                variant=variant,
+                flashrt_us=flashrt_us,
+                torch_reference_us=torch_us,
+                speedup_vs_reference=torch_us / flashrt_us,
+                max_abs=max_abs,
+                mean_abs=mean_abs,
+                p99_abs=p99_abs,
+                cosine=cosine,
+                status="ok",
+            )
+        )
+    return results
+def main() -> int:
+    parser = argparse.ArgumentParser()
+    parser.add_argument("--mode", choices=["smoke", "headline"], default="headline")
+    parser.add_argument("--warmup", type=int, default=20)
+    parser.add_argument("--iterations", type=int, default=100)
+    parser.add_argument("--json-out", default=None)
+    args = parser.parse_args()
+    helpers = load_helpers()
+    ops = helpers.load_source_ops()
+    shapes = {
+        "small_m16_n128_k128": (16, 128, 128),
+        "small_m32_n256_k256": (32, 256, 256),
+        "mlp_tile_m64_n512_k512": (64, 512, 512),
+    }
+    if args.mode == "smoke":
+        shapes = {"small_m16_n128_k128": shapes["small_m16_n128_k128"]}
+    results: list[BenchResult] = []
+    for name, shape in shapes.items():
+        results.extend(bench_case(helpers, ops, name, shape, args.warmup, args.iterations))
+    payload = {
+        "mode": args.mode,
+        "device": torch.cuda.get_device_name(),
+        "torch": torch.__version__,
+        "results": [asdict(item) for item in results],
+    }
+    print(json.dumps(payload, indent=2))
+    if args.json_out:
+        out = Path(args.json_out)
+        out.parent.mkdir(parents=True, exist_ok=True)
+        out.write_text(json.dumps(payload, indent=2) + "\n")
+    return 0
+if __name__ == "__main__":
+    raise SystemExit(main())

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

	@@ -0,0 +1,86 @@

+"""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("fp4_w4a16_linear_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 = 0,
+) -> 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("dequantize_fp4_sfa_fp16"))
+def _dequant_fake(packed: torch.Tensor, sfa: torch.Tensor, out: torch.Tensor, is_sfb: bool = False) -> 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 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 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 = 0,
+) -> 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.fp4_w4a16_linear_bf16(a_packed, b_packed, sfa, sfb, out, float(alpha), int(variant))
+    return out

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

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

+version https://git-lfs.github.com/spec/v1
+oid sha256:5c5bb3e58d2e4dcdee5807ca455fafeb1a1e72515cbdd0bd9c130ea3bfacb9d1
+size 671608

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

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

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

build/torch211-cxx11-cu128-x86_64-linux/fp4_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,22 @@

+{
+  "name": "fp4-gemm",
+  "id": "_fp4_gemm_cuda_d8a589a",
+  "version": 1,
+  "license": "Apache-2.0",
+  "python-depends": [],
+  "backend": {
+    "type": "cuda",
+    "archs": [
+      "12.0a"
+    ]
+  },
+  "digest": {
+    "algorithm": "sha256",
+    "files": {
+      "__init__.py": "WSLdLQD92ulRtOhSLoAhWgFYN6oucbbkTg19GAaltBs=",
+      "_fp4_gemm_cuda_d8a589a.abi3.so": "XFuz5Y0uTc3uWAfKRV+v6xoeclFcvdC9nBMOo7+sudE=",
+      "_ops.py": "JCogli/U0X3ICMajcGwPdPXOfLjDX5Izclp70jtHdJM=",
+      "fp4_gemm/__init__.py": "DFYPlrhXwYjEqCl/8n0SmWGZV8NFml5DPhMjKfv98GY="
+    }
+  }
+}

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

	@@ -0,0 +1,86 @@

+"""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("fp4_w4a16_linear_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 = 0,
+) -> 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("dequantize_fp4_sfa_fp16"))
+def _dequant_fake(packed: torch.Tensor, sfa: torch.Tensor, out: torch.Tensor, is_sfb: bool = False) -> 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 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 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 = 0,
+) -> 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.fp4_w4a16_linear_bf16(a_packed, b_packed, sfa, sfb, out, float(alpha), int(variant))
+    return out

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

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

+version https://git-lfs.github.com/spec/v1
+oid sha256:5fc0542c6250a3290db008f65939284a1fc789c05ff972945905d647673f3b76
+size 714176

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

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

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

build/torch211-cxx11-cu130-x86_64-linux/fp4_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,22 @@

+{
+  "name": "fp4-gemm",
+  "id": "_fp4_gemm_cuda_d8a589a",
+  "version": 1,
+  "license": "Apache-2.0",
+  "python-depends": [],
+  "backend": {
+    "type": "cuda",
+    "archs": [
+      "12.0a"
+    ]
+  },
+  "digest": {
+    "algorithm": "sha256",
+    "files": {
+      "__init__.py": "WSLdLQD92ulRtOhSLoAhWgFYN6oucbbkTg19GAaltBs=",
+      "_fp4_gemm_cuda_d8a589a.abi3.so": "X8BULGJQoykNsAj2WTkoSh/HicBf+XKUWQXWR2c/O3Y=",
+      "_ops.py": "JCogli/U0X3ICMajcGwPdPXOfLjDX5Izclp70jtHdJM=",
+      "fp4_gemm/__init__.py": "DFYPlrhXwYjEqCl/8n0SmWGZV8NFml5DPhMjKfv98GY="
+    }
+  }
+}

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

	@@ -0,0 +1,86 @@

+"""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("fp4_w4a16_linear_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 = 0,
+) -> 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("dequantize_fp4_sfa_fp16"))
+def _dequant_fake(packed: torch.Tensor, sfa: torch.Tensor, out: torch.Tensor, is_sfb: bool = False) -> 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 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 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 = 0,
+) -> 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.fp4_w4a16_linear_bf16(a_packed, b_packed, sfa, sfb, out, float(alpha), int(variant))
+    return out

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

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

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

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

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

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

build/torch212-cxx11-cu130-x86_64-linux/fp4_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,22 @@

+{
+  "name": "fp4-gemm",
+  "id": "_fp4_gemm_cuda_d8a589a",
+  "version": 1,
+  "license": "Apache-2.0",
+  "python-depends": [],
+  "backend": {
+    "type": "cuda",
+    "archs": [
+      "12.0a"
+    ]
+  },
+  "digest": {
+    "algorithm": "sha256",
+    "files": {
+      "__init__.py": "WSLdLQD92ulRtOhSLoAhWgFYN6oucbbkTg19GAaltBs=",
+      "_fp4_gemm_cuda_d8a589a.abi3.so": "1hdbbg2bHzOAxPc8xe4pI28kF4+rFnZH4v4CAaTOCh4=",
+      "_ops.py": "JCogli/U0X3ICMajcGwPdPXOfLjDX5Izclp70jtHdJM=",
+      "fp4_gemm/__init__.py": "DFYPlrhXwYjEqCl/8n0SmWGZV8NFml5DPhMjKfv98GY="
+    }
+  }
+}

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

	@@ -0,0 +1,86 @@

+"""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("fp4_w4a16_linear_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 = 0,
+) -> 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("dequantize_fp4_sfa_fp16"))
+def _dequant_fake(packed: torch.Tensor, sfa: torch.Tensor, out: torch.Tensor, is_sfb: bool = False) -> 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 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 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 = 0,
+) -> 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.fp4_w4a16_linear_bf16(a_packed, b_packed, sfa, sfb, out, float(alpha), int(variant))
+    return out

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

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

+version https://git-lfs.github.com/spec/v1
+oid sha256:9358a9bee95a43f46c5ba7716c25513c5c01d11ba020ebcca042bcc5daff807e
+size 724632

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

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

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

build/torch212-cxx11-cu132-x86_64-linux/fp4_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,22 @@

+{
+  "name": "fp4-gemm",
+  "id": "_fp4_gemm_cuda_d8a589a",
+  "version": 1,
+  "license": "Apache-2.0",
+  "python-depends": [],
+  "backend": {
+    "type": "cuda",
+    "archs": [
+      "12.0a"
+    ]
+  },
+  "digest": {
+    "algorithm": "sha256",
+    "files": {
+      "__init__.py": "WSLdLQD92ulRtOhSLoAhWgFYN6oucbbkTg19GAaltBs=",
+      "_fp4_gemm_cuda_d8a589a.abi3.so": "k1ipvulaQ/RsW6dxbCVRPFwB0RugIOvMoEK8xdr/gH4=",
+      "_ops.py": "JCogli/U0X3ICMajcGwPdPXOfLjDX5Izclp70jtHdJM=",
+      "fp4_gemm/__init__.py": "DFYPlrhXwYjEqCl/8n0SmWGZV8NFml5DPhMjKfv98GY="
+    }
+  }
+}