Uploaded using `kernel-builder`.

Browse files

Files changed (14) hide show

benchmarks/benchmark.py +46 -0
build/torch211-cxx11-cu128-x86_64-linux/__init__.py +154 -0
build/torch211-cxx11-cu128-x86_64-linux/_int4_blackwell_cuda_7293754.abi3.so +3 -0
build/torch211-cxx11-cu128-x86_64-linux/_ops.py +9 -0
build/torch211-cxx11-cu128-x86_64-linux/cubin/sm120/probe.cubin +0 -0
build/torch211-cxx11-cu128-x86_64-linux/cubin/sm120/probe_int4.cubin +0 -0
build/torch211-cxx11-cu128-x86_64-linux/cubin/sm120/probe_int4a.cubin +0 -0
build/torch211-cxx11-cu128-x86_64-linux/cubin/sm120/probe_int4b.cubin +0 -0
build/torch211-cxx11-cu128-x86_64-linux/cubin/sm121/probe.cubin +0 -0
build/torch211-cxx11-cu128-x86_64-linux/cubin/sm121/probe_int4.cubin +0 -0
build/torch211-cxx11-cu128-x86_64-linux/cubin/sm121/probe_int4a.cubin +0 -0
build/torch211-cxx11-cu128-x86_64-linux/cubin/sm121/probe_int4b.cubin +0 -0
build/torch211-cxx11-cu128-x86_64-linux/int4_blackwell/__init__.py +26 -0
build/torch211-cxx11-cu128-x86_64-linux/metadata.json +41 -0

benchmarks/benchmark.py ADDED Viewed

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

+import argparse
+import torch
+import int4_blackwell
+def main() -> None:
+    parser = argparse.ArgumentParser()
+    parser.add_argument("--iterations", type=int, default=8192)
+    parser.add_argument("--repeats", type=int, default=10)
+    parser.add_argument("--launches", type=int, default=20)
+    args = parser.parse_args()
+    props = torch.cuda.get_device_properties(0)
+    blocks = props.multi_processor_count * 4
+    warps = blocks * 8
+    flops = warps * 4 * args.iterations * 2 * 16 * 8 * 64
+    out = torch.empty((blocks, 256), device="cuda", dtype=torch.float32)
+    for mode in ("e2m1", "a", "b", "ab"):
+        int4_blackwell.mma_probe(
+            mode, iterations=args.iterations, blocks=blocks, out=out
+        )
+        torch.cuda.synchronize()
+        start = torch.cuda.Event(enable_timing=True)
+        end = torch.cuda.Event(enable_timing=True)
+        best_ms = float("inf")
+        for _ in range(args.repeats):
+            start.record()
+            int4_blackwell.mma_probe(
+                mode,
+                iterations=args.iterations,
+                blocks=blocks,
+                launches=args.launches,
+                out=out,
+            )
+            end.record()
+            end.synchronize()
+            best_ms = min(best_ms, start.elapsed_time(end))
+        per_launch_ms = best_ms / args.launches
+        tflops = flops / (per_launch_ms * 1e-3) / 1e12
+        print(f"{mode:5s} {per_launch_ms * 1e3:9.3f} us {tflops:8.1f} TFLOPS")
+if __name__ == "__main__":
+    main()

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

	@@ -0,0 +1,154 @@

+"""Experimental native E0M3/INT4 tensor-core primitives for Blackwell."""
+from __future__ import annotations
+from importlib.resources import files
+from typing import Literal
+import torch
+from ._ops import ops
+OperandMode = Literal["e2m1", "a", "b", "ab"]
+_CUBIN_NAMES = {
+    "e2m1": "probe.cubin",
+    "a": "probe_int4a.cubin",
+    "b": "probe_int4b.cubin",
+    "ab": "probe_int4.cubin",
+}
+_SUPPORTED_ARCHES = {(12, 0): "sm120", (12, 1): "sm121"}
+_CACHE: dict[tuple[str, str], torch.Tensor] = {}
+def _architecture(device: int) -> str:
+    capability = torch.cuda.get_device_capability(device)
+    try:
+        return _SUPPORTED_ARCHES[capability]
+    except KeyError as error:
+        supported = ", ".join(name.upper() for name in _SUPPORTED_ARCHES.values())
+        raise RuntimeError(
+            f"int4-blackwell supports {supported}; got SM{capability[0]}{capability[1]}"
+        ) from error
+def _capability(device: int) -> tuple[int, int]:
+    return tuple(torch.cuda.get_device_capability(device))
+def _is_tcgen05(device: int) -> bool:
+    return _capability(device) in {(10, 0), (10, 3), (11, 0)}
+def _cubin(mode: OperandMode, device: int) -> torch.Tensor:
+    if mode not in _CUBIN_NAMES:
+        raise ValueError(
+            f"mode must be one of {tuple(_CUBIN_NAMES)}, got {mode!r}"
+        )
+    architecture = _architecture(device)
+    key = (architecture, mode)
+    if key not in _CACHE:
+        data = (
+            files(__package__)
+            .joinpath("cubin", architecture, _CUBIN_NAMES[mode])
+            .read_bytes()
+        )
+        _CACHE[key] = torch.frombuffer(bytearray(data), dtype=torch.uint8)
+    return _CACHE[key]
+def _device_index(device: int | torch.device | None) -> int:
+    if device is None:
+        return torch.cuda.current_device()
+    if isinstance(device, int):
+        return device
+    parsed = torch.device(device)
+    if parsed.type != "cuda":
+        raise ValueError(f"device must be CUDA, got {parsed}")
+    return torch.cuda.current_device() if parsed.index is None else parsed.index
+def codebook_probe(
+    mode: OperandMode = "ab", *, device: int | torch.device | None = None
+) -> torch.Tensor:
+    """Return the 16-value A-operand decode table measured by native MMA.
+    ``mode`` selects standard E2M1 or the patched INT4 format independently
+    for the A and B operands. The result is synchronized and returned on CPU.
+    """
+    dev = _device_index(device)
+    if _is_tcgen05(dev):
+        if mode != "ab":
+            raise ValueError(
+                "the tcgen05 backend currently exposes the native INT4 x INT4 "
+                "descriptor; use mode='ab'"
+            )
+        m = n = k = 128
+        b_packed = torch.full(
+            (n, k // 2), 0x11, device=f"cuda:{dev}", dtype=torch.uint8
+        )
+        # Constant-one UE4M3 storage is layout-invariant. Deliberately
+        # overallocate the physical CUTLASS scale-factor tensors for this
+        # instruction canary so no private layout helper leaks into the API.
+        sfa = torch.full((m * k,), 0x38, device=f"cuda:{dev}", dtype=torch.uint8)
+        sfb = torch.full((n * k,), 0x38, device=f"cuda:{dev}", dtype=torch.uint8)
+        values = []
+        for value in range(16):
+            packed = value | (value << 4)
+            a_packed = torch.full(
+                (m, k // 2), packed, device=f"cuda:{dev}", dtype=torch.uint8
+            )
+            tile = ops.tcgen05_int4_gemm_bf16(a_packed, sfa, b_packed, sfb)
+            first = tile[0, 0]
+            if not torch.equal(tile, first.expand_as(tile)):
+                raise RuntimeError("tcgen05 INT4 codebook output is not uniform")
+            values.append(first.float() / k)
+        return torch.stack(values).cpu()
+    tile = ops.run_codebook_probe(_cubin(mode, dev), dev)
+    torch.cuda.synchronize(dev)
+    if not torch.equal(tile, tile[:, :1].expand_as(tile)):
+        raise RuntimeError("native MMA output tile is not uniform")
+    return (tile[:, 0] / 64.0).cpu()
+def mma_probe(
+    mode: OperandMode = "ab",
+    *,
+    iterations: int = 8192,
+    blocks: int | None = None,
+    launches: int = 1,
+    device: int | torch.device | None = None,
+    out: torch.Tensor | None = None,
+) -> torch.Tensor:
+    """Launch the register-resident MMA throughput probe asynchronously."""
+    dev = _device_index(device)
+    if _is_tcgen05(dev):
+        raise RuntimeError(
+            "mma_probe is the SM120/SM121 register-resident OMMA probe; "
+            "benchmark tcgen05_int4_gemm_bf16 on SM100/SM103/SM110"
+        )
+    if blocks is None:
+        blocks = torch.cuda.get_device_properties(dev).multi_processor_count * 4
+    if out is None:
+        out = torch.empty((blocks, 256), device=f"cuda:{dev}", dtype=torch.float32)
+    ops.run_mma_probe(_cubin(mode, dev), out, iterations, blocks, launches, dev)
+    return out
+def tcgen05_int4_gemm_bf16(
+    a_packed: torch.Tensor,
+    sfa_physical: torch.Tensor,
+    b_packed: torch.Tensor,
+    sfb_physical: torch.Tensor,
+) -> torch.Tensor:
+    """Run native E0M3 x E0M3 block-scaled GEMM on SM100/SM103/SM110.
+    ``a_packed`` and ``b_packed`` contain two sign-magnitude INT4 values per
+    byte. Scale tensors are the physical CUTLASS UE4M3 block-16 layouts.
+    """
+    return ops.tcgen05_int4_gemm_bf16(
+        a_packed, sfa_physical, b_packed, sfb_physical
+    )
+__all__ = ["codebook_probe", "mma_probe", "tcgen05_int4_gemm_bf16"]

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

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

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

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

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

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

build/torch211-cxx11-cu128-x86_64-linux/cubin/sm120/probe.cubin ADDED Viewed