Promote latest kernel artifacts to main

.gitattributes

@@ -1,35 +1,3 @@
-*.7z filter=lfs diff=lfs merge=lfs -text
-*.arrow filter=lfs diff=lfs merge=lfs -text
-*.bin filter=lfs diff=lfs merge=lfs -text
-*.bz2 filter=lfs diff=lfs merge=lfs -text
-*.ckpt filter=lfs diff=lfs merge=lfs -text
-*.ftz filter=lfs diff=lfs merge=lfs -text
-*.gz filter=lfs diff=lfs merge=lfs -text
-*.h5 filter=lfs diff=lfs merge=lfs -text
-*.joblib filter=lfs diff=lfs merge=lfs -text
-*.lfs.* filter=lfs diff=lfs merge=lfs -text
-*.mlmodel filter=lfs diff=lfs merge=lfs -text
-*.model filter=lfs diff=lfs merge=lfs -text
-*.msgpack filter=lfs diff=lfs merge=lfs -text
-*.npy filter=lfs diff=lfs merge=lfs -text
-*.npz filter=lfs diff=lfs merge=lfs -text
-*.onnx filter=lfs diff=lfs merge=lfs -text
-*.ot filter=lfs diff=lfs merge=lfs -text
-*.parquet filter=lfs diff=lfs merge=lfs -text
-*.pb filter=lfs diff=lfs merge=lfs -text
-*.pickle filter=lfs diff=lfs merge=lfs -text
-*.pkl filter=lfs diff=lfs merge=lfs -text
-*.pt filter=lfs diff=lfs merge=lfs -text
-*.pth filter=lfs diff=lfs merge=lfs -text
-*.rar filter=lfs diff=lfs merge=lfs -text
-*.safetensors filter=lfs diff=lfs merge=lfs -text
-saved_model/**/* filter=lfs diff=lfs merge=lfs -text
-*.tar.* filter=lfs diff=lfs merge=lfs -text
-*.tar filter=lfs diff=lfs merge=lfs -text
-*.tflite filter=lfs diff=lfs merge=lfs -text
-*.tgz filter=lfs diff=lfs merge=lfs -text
-*.wasm filter=lfs diff=lfs merge=lfs -text
-*.xz filter=lfs diff=lfs merge=lfs -text
-*.zip filter=lfs diff=lfs merge=lfs -text
-*.zst filter=lfs diff=lfs merge=lfs -text
-*tfevents* filter=lfs diff=lfs merge=lfs -text
+*.so filter=lfs diff=lfs merge=lfs -text
+*.pyd filter=lfs diff=lfs merge=lfs -text
+*.dylib filter=lfs diff=lfs merge=lfs -text

README.md

@@ -1,9 +0,0 @@
-# flashrt/grouped-moe-gemm
-
-This repository is a compatibility mirror for older `kernels` clients
-that resolve repositories through the default Hugging Face model repo API.
-
-Canonical Kernel Hub repo: https://huggingface.co/kernels/flashrt/grouped-moe-gemm
-
-Do not edit this mirror by hand. It is generated from the Kernel Hub
-`vN` branches and contains the same `build/**` artifacts.

benchmarks/benchmark.py

@@ -0,0 +1,139 @@
+#!/usr/bin/env python3
+"""35B-A3B grouped expert prefill GEMM package/native parity benchmark."""
+
+from __future__ import annotations
+
+import argparse
+import importlib
+import json
+import sys
+from pathlib import Path
+
+import torch
+
+
+PACKAGE = Path(__file__).resolve().parents[1]
+sys.path.insert(0, str(PACKAGE / "tests"))
+from _source_loader import load_source_ops  # noqa: E402
+from test_grouped_moe_gemm import make_sf  # noqa: E402
+
+
+WORKLOADS = {
+    "35b_gate_up": (64, 1024, 2048, 16, 256),
+    "35b_down": (64, 2048, 512, 16, 256),
+}
+
+
+def time_us(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 start.elapsed_time(end) * 1000.0 / iters
+
+
+def load_ops(backend: str, artifact: str | None):
+    if backend == "source":
+        return load_source_ops()
+    if artifact:
+        sys.path.insert(0, artifact)
+    try:
+        return importlib.import_module("grouped_moe_gemm")
+    finally:
+        if artifact:
+            sys.path.remove(artifact)
+
+
+def make_case(tile_rows: int, n: int, k: int, tiles: int, experts: int):
+    rows = tile_rows * tiles
+    dev = "cuda"
+    packed = torch.randint(0, 256, (rows, k // 2), device=dev, dtype=torch.uint8)
+    weights = torch.randint(
+        0, 256, (experts, n, k // 2), device=dev, dtype=torch.uint8
+    )
+    input_scale, _ = make_sf(rows, k, dev)
+    weight_scales = torch.stack([make_sf(n, k, dev)[0] for _ in range(experts)])
+    alpha = torch.rand(experts, device=dev, dtype=torch.float32) + 0.5
+    tile_expert = torch.arange(tiles, device=dev, dtype=torch.int32) % experts
+    out = torch.empty(rows, n, device=dev, dtype=torch.bfloat16)
+    return packed, weights, input_scale, weight_scales, alpha, tile_expert, out
+
+
+def main() -> int:
+    parser = argparse.ArgumentParser()
+    parser.add_argument("--backend", choices=("source", "installed"), default="source")
+    parser.add_argument("--artifact")
+    parser.add_argument("--warmup", type=int, default=20)
+    parser.add_argument("--iters", type=int, default=100)
+    parser.add_argument("--json-out")
+    args = parser.parse_args()
+
+    torch.manual_seed(9102)
+    ops = load_ops(args.backend, args.artifact)
+    rows = []
+    for name, (tile_rows, n, k, tiles, experts) in WORKLOADS.items():
+        packed, weights, input_scale, weight_scales, alpha, tile_expert, out = make_case(
+            tile_rows, n, k, tiles, experts
+        )
+        kwargs = dict(
+            tile_rows=tile_rows,
+            weight_stride=weights[0].numel(),
+            weight_scale_stride=weight_scales[0].numel(),
+        )
+
+        def public_call():
+            return ops.grouped_nvfp4_gemm_bf16(
+                packed, weights, input_scale, weight_scales, alpha,
+                tile_expert, out=out, **kwargs
+            )
+
+        def raw_call():
+            ops.ops.grouped_nvfp4_gemm_bf16_out(
+                packed, weights, input_scale, weight_scales, alpha,
+                tile_expert, tile_rows, 0, kwargs["weight_stride"],
+                kwargs["weight_scale_stride"], out,
+            )
+
+        expected = public_call().clone()
+        raw_call()
+        torch.cuda.synchronize()
+        if not torch.equal(out, expected):
+            raise AssertionError(f"{name}: public/raw output mismatch")
+        public_us = time_us(public_call, args.warmup, args.iters)
+        raw_us = time_us(raw_call, args.warmup, args.iters)
+        ratio = public_us / raw_us
+        accepted = public_us - raw_us <= max(0.75, raw_us * 0.05)
+        row = {
+            "workload": name,
+            "tile_rows": tile_rows,
+            "rows": tile_rows * tiles,
+            "n": n,
+            "k": k,
+            "experts": experts,
+            "raw_native_us": raw_us,
+            "public_wrapper_us": public_us,
+            "wrapper_native": ratio,
+            "bit_exact": True,
+            "accepted": accepted,
+        }
+        rows.append(row)
+        print(
+            f"{name}: raw={raw_us:.3f}us public={public_us:.3f}us "
+            f"ratio={ratio:.4f} accepted={accepted}"
+        )
+    if args.json_out:
+        Path(args.json_out).parent.mkdir(parents=True, exist_ok=True)
+        Path(args.json_out).write_text(json.dumps(rows, indent=2) + "\n")
+    if not all(row["accepted"] for row in rows):
+        raise AssertionError("grouped MoE wrapper/native parity failed")
+    return 0
+
+
+if __name__ == "__main__":
+    raise SystemExit(main())

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

@@ -0,0 +1,112 @@
+"""Grouped NVFP4 MoE GEMM for Blackwell prefill workloads."""
+
+from __future__ import annotations
+import torch
+from ._ops import add_op_namespace_prefix, ops
+
+
+@torch.library.custom_op(
+    add_op_namespace_prefix("_grouped_nvfp4_gemm_bf16"),
+    mutates_args=(),
+    device_types="cuda",
+)
+def _compileable(
+    input: torch.Tensor,
+    weight: torch.Tensor,
+    input_scale: torch.Tensor,
+    weight_scale: torch.Tensor,
+    alpha: torch.Tensor,
+    tile_expert: torch.Tensor,
+    tile_rows: int,
+    input_scale_stride: int,
+    weight_stride: int,
+    weight_scale_stride: int,
+) -> torch.Tensor:
+    output = torch.empty(
+        (input.shape[0], weight.shape[1]), device=input.device, dtype=torch.bfloat16
+    )
+    ops.grouped_nvfp4_gemm_bf16_out(
+        input,
+        weight,
+        input_scale,
+        weight_scale,
+        alpha,
+        tile_expert,
+        tile_rows,
+        input_scale_stride,
+        weight_stride,
+        weight_scale_stride,
+        output,
+    )
+    return output
+
+
+@torch.library.register_fake(add_op_namespace_prefix("_grouped_nvfp4_gemm_bf16"))
+def _fake(
+    input,
+    weight,
+    input_scale,
+    weight_scale,
+    alpha,
+    tile_expert,
+    tile_rows,
+    input_scale_stride,
+    weight_stride,
+    weight_scale_stride,
+):
+    if input.ndim != 2 or weight.ndim != 3:
+        raise RuntimeError("invalid grouped GEMM tensor shapes")
+    return torch.empty(
+        (input.shape[0], weight.shape[1]), device=input.device, dtype=torch.bfloat16
+    )
+
+
+def grouped_nvfp4_gemm_bf16(
+    input,
+    weight,
+    input_scale,
+    weight_scale,
+    alpha,
+    tile_expert,
+    *,
+    tile_rows,
+    input_scale_stride=0,
+    weight_stride=None,
+    weight_scale_stride=None,
+    out=None,
+):
+    """Compute expert-selected packed NVFP4 GEMM tiles with BF16 output."""
+    if weight_stride is None:
+        weight_stride = weight[0].numel()
+    if weight_scale_stride is None:
+        weight_scale_stride = weight_scale[0].numel()
+    if out is None:
+        return _compileable(
+            input,
+            weight,
+            input_scale,
+            weight_scale,
+            alpha,
+            tile_expert,
+            int(tile_rows),
+            int(input_scale_stride),
+            int(weight_stride),
+            int(weight_scale_stride),
+        )
+    ops.grouped_nvfp4_gemm_bf16_out(
+        input,
+        weight,
+        input_scale,
+        weight_scale,
+        alpha,
+        tile_expert,
+        int(tile_rows),
+        int(input_scale_stride),
+        int(weight_stride),
+        int(weight_scale_stride),
+        out,
+    )
+    return out
+
+
+__all__ = ["grouped_nvfp4_gemm_bf16"]

build/torch211-cxx11-cu128-x86_64-linux/_grouped_moe_gemm_cuda_32e9b2c.abi3.so

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:a19fb3043235055688d9e01dda3990ca2c00e9a5f8aec3656429a67bbdbac6fd
+size 175392

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

@@ -0,0 +1,9 @@
+import torch
+from . import _grouped_moe_gemm_cuda_32e9b2c
+ops = torch.ops._grouped_moe_gemm_cuda_32e9b2c
+
+def add_op_namespace_prefix(op_name: str):
+    """
+    Prefix op by namespace.
+    """
+    return f"_grouped_moe_gemm_cuda_32e9b2c::{op_name}"

build/torch211-cxx11-cu128-x86_64-linux/grouped_moe_gemm/__init__.py

@@ -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

@@ -0,0 +1,33 @@
+{
+  "name": "grouped-moe-gemm",
+  "id": "_grouped_moe_gemm_cuda_32e9b2c",
+  "version": 1,
+  "license": "Apache-2.0",
+  "python-depends": [],
+  "backend": {
+    "type": "cuda",
+    "archs": [
+      "12.0a"
+    ]
+  },
+  "digest": {
+    "algorithm": "sha256",
+    "files": {
+      "__init__.py": "/pd/JqQCw5Ie4nddrEvVeqaOFtlbkqNDqykgQLlZJF0=",
+      "_grouped_moe_gemm_cuda_32e9b2c.abi3.so": "oZ+zBDI1BVaI2eAd2jmQyiwA6aX4rsNlZCmme726xv0=",
+      "_ops.py": "AxmxTU/RwyGFkTVT0IV6Y0/bjHdcILOm6O5JNy/zf6k=",
+      "grouped_moe_gemm/__init__.py": "DFYPlrhXwYjEqCl/8n0SmWGZV8NFml5DPhMjKfv98GY="
+    }
+  },
+  "provenance": {
+    "kernel-builder": {
+      "version": "0.17.0-dev0",
+      "sha": "19aaa6421e674e9fecc352bbae6eab81d19a6bf4",
+      "dirty": false
+    },
+    "kernel": {
+      "sha": "32e9b2c39ee268d8a76fcfef667974a13022c283",
+      "dirty": false
+    }
+  }
+}

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

@@ -0,0 +1,112 @@
+"""Grouped NVFP4 MoE GEMM for Blackwell prefill workloads."""
+
+from __future__ import annotations
+import torch
+from ._ops import add_op_namespace_prefix, ops
+
+
+@torch.library.custom_op(
+    add_op_namespace_prefix("_grouped_nvfp4_gemm_bf16"),
+    mutates_args=(),
+    device_types="cuda",
+)
+def _compileable(
+    input: torch.Tensor,
+    weight: torch.Tensor,
+    input_scale: torch.Tensor,
+    weight_scale: torch.Tensor,
+    alpha: torch.Tensor,
+    tile_expert: torch.Tensor,
+    tile_rows: int,
+    input_scale_stride: int,
+    weight_stride: int,
+    weight_scale_stride: int,
+) -> torch.Tensor:
+    output = torch.empty(
+        (input.shape[0], weight.shape[1]), device=input.device, dtype=torch.bfloat16
+    )
+    ops.grouped_nvfp4_gemm_bf16_out(
+        input,
+        weight,
+        input_scale,
+        weight_scale,
+        alpha,
+        tile_expert,
+        tile_rows,
+        input_scale_stride,
+        weight_stride,
+        weight_scale_stride,
+        output,
+    )
+    return output
+
+
+@torch.library.register_fake(add_op_namespace_prefix("_grouped_nvfp4_gemm_bf16"))
+def _fake(
+    input,
+    weight,
+    input_scale,
+    weight_scale,
+    alpha,
+    tile_expert,
+    tile_rows,
+    input_scale_stride,
+    weight_stride,
+    weight_scale_stride,
+):
+    if input.ndim != 2 or weight.ndim != 3:
+        raise RuntimeError("invalid grouped GEMM tensor shapes")
+    return torch.empty(
+        (input.shape[0], weight.shape[1]), device=input.device, dtype=torch.bfloat16
+    )
+
+
+def grouped_nvfp4_gemm_bf16(
+    input,
+    weight,
+    input_scale,
+    weight_scale,
+    alpha,
+    tile_expert,
+    *,
+    tile_rows,
+    input_scale_stride=0,
+    weight_stride=None,
+    weight_scale_stride=None,
+    out=None,
+):
+    """Compute expert-selected packed NVFP4 GEMM tiles with BF16 output."""
+    if weight_stride is None:
+        weight_stride = weight[0].numel()
+    if weight_scale_stride is None:
+        weight_scale_stride = weight_scale[0].numel()
+    if out is None:
+        return _compileable(
+            input,
+            weight,
+            input_scale,
+            weight_scale,
+            alpha,
+            tile_expert,
+            int(tile_rows),
+            int(input_scale_stride),
+            int(weight_stride),
+            int(weight_scale_stride),
+        )
+    ops.grouped_nvfp4_gemm_bf16_out(
+        input,
+        weight,
+        input_scale,
+        weight_scale,
+        alpha,
+        tile_expert,
+        int(tile_rows),
+        int(input_scale_stride),
+        int(weight_stride),
+        int(weight_scale_stride),
+        out,
+    )
+    return out
+
+
+__all__ = ["grouped_nvfp4_gemm_bf16"]

build/torch211-cxx11-cu130-x86_64-linux/_grouped_moe_gemm_cuda_32e9b2c.abi3.so

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:43c9de9e12cfb9e84101046e351c8c9279f580d69a0f4dca9c98c5d6b5b3f0bd
+size 183176

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

@@ -0,0 +1,9 @@
+import torch
+from . import _grouped_moe_gemm_cuda_32e9b2c
+ops = torch.ops._grouped_moe_gemm_cuda_32e9b2c
+
+def add_op_namespace_prefix(op_name: str):
+    """
+    Prefix op by namespace.
+    """
+    return f"_grouped_moe_gemm_cuda_32e9b2c::{op_name}"

build/torch211-cxx11-cu130-x86_64-linux/grouped_moe_gemm/__init__.py

@@ -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

@@ -0,0 +1,33 @@
+{
+  "name": "grouped-moe-gemm",
+  "id": "_grouped_moe_gemm_cuda_32e9b2c",
+  "version": 1,
+  "license": "Apache-2.0",
+  "python-depends": [],
+  "backend": {
+    "type": "cuda",
+    "archs": [
+      "12.0a"
+    ]
+  },
+  "digest": {
+    "algorithm": "sha256",
+    "files": {
+      "__init__.py": "/pd/JqQCw5Ie4nddrEvVeqaOFtlbkqNDqykgQLlZJF0=",
+      "_grouped_moe_gemm_cuda_32e9b2c.abi3.so": "Q8nenhLPuehBAQRuNRyMknn1gNaaD03KnJjF1rWz8L0=",
+      "_ops.py": "AxmxTU/RwyGFkTVT0IV6Y0/bjHdcILOm6O5JNy/zf6k=",
+      "grouped_moe_gemm/__init__.py": "DFYPlrhXwYjEqCl/8n0SmWGZV8NFml5DPhMjKfv98GY="
+    }
+  },
+  "provenance": {
+    "kernel-builder": {
+      "version": "0.17.0-dev0",
+      "sha": "19aaa6421e674e9fecc352bbae6eab81d19a6bf4",
+      "dirty": false
+    },
+    "kernel": {
+      "sha": "32e9b2c39ee268d8a76fcfef667974a13022c283",
+      "dirty": false
+    }
+  }
+}

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

@@ -0,0 +1,112 @@
+"""Grouped NVFP4 MoE GEMM for Blackwell prefill workloads."""
+
+from __future__ import annotations
+import torch
+from ._ops import add_op_namespace_prefix, ops
+
+
+@torch.library.custom_op(
+    add_op_namespace_prefix("_grouped_nvfp4_gemm_bf16"),
+    mutates_args=(),
+    device_types="cuda",
+)
+def _compileable(
+    input: torch.Tensor,
+    weight: torch.Tensor,
+    input_scale: torch.Tensor,
+    weight_scale: torch.Tensor,
+    alpha: torch.Tensor,
+    tile_expert: torch.Tensor,
+    tile_rows: int,
+    input_scale_stride: int,
+    weight_stride: int,
+    weight_scale_stride: int,
+) -> torch.Tensor:
+    output = torch.empty(
+        (input.shape[0], weight.shape[1]), device=input.device, dtype=torch.bfloat16
+    )
+    ops.grouped_nvfp4_gemm_bf16_out(
+        input,
+        weight,
+        input_scale,
+        weight_scale,
+        alpha,
+        tile_expert,
+        tile_rows,
+        input_scale_stride,
+        weight_stride,
+        weight_scale_stride,
+        output,
+    )
+    return output
+
+
+@torch.library.register_fake(add_op_namespace_prefix("_grouped_nvfp4_gemm_bf16"))
+def _fake(
+    input,
+    weight,
+    input_scale,
+    weight_scale,
+    alpha,
+    tile_expert,
+    tile_rows,
+    input_scale_stride,
+    weight_stride,
+    weight_scale_stride,
+):
+    if input.ndim != 2 or weight.ndim != 3:
+        raise RuntimeError("invalid grouped GEMM tensor shapes")
+    return torch.empty(
+        (input.shape[0], weight.shape[1]), device=input.device, dtype=torch.bfloat16
+    )
+
+
+def grouped_nvfp4_gemm_bf16(
+    input,
+    weight,
+    input_scale,
+    weight_scale,
+    alpha,
+    tile_expert,
+    *,
+    tile_rows,
+    input_scale_stride=0,
+    weight_stride=None,
+    weight_scale_stride=None,
+    out=None,
+):
+    """Compute expert-selected packed NVFP4 GEMM tiles with BF16 output."""
+    if weight_stride is None:
+        weight_stride = weight[0].numel()
+    if weight_scale_stride is None:
+        weight_scale_stride = weight_scale[0].numel()
+    if out is None:
+        return _compileable(
+            input,
+            weight,
+            input_scale,
+            weight_scale,
+            alpha,
+            tile_expert,
+            int(tile_rows),
+            int(input_scale_stride),
+            int(weight_stride),
+            int(weight_scale_stride),
+        )
+    ops.grouped_nvfp4_gemm_bf16_out(
+        input,
+        weight,
+        input_scale,
+        weight_scale,
+        alpha,
+        tile_expert,
+        int(tile_rows),
+        int(input_scale_stride),
+        int(weight_stride),
+        int(weight_scale_stride),
+        out,
+    )
+    return out
+
+
+__all__ = ["grouped_nvfp4_gemm_bf16"]

build/torch212-cxx11-cu130-x86_64-linux/_grouped_moe_gemm_cuda_32e9b2c.abi3.so

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:b0b5048bcc55227f23f3606ffc6fe52cd8359293104c6f0e5cee8865813aaf62
+size 194088

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

@@ -0,0 +1,9 @@
+import torch
+from . import _grouped_moe_gemm_cuda_32e9b2c
+ops = torch.ops._grouped_moe_gemm_cuda_32e9b2c
+
+def add_op_namespace_prefix(op_name: str):
+    """
+    Prefix op by namespace.
+    """
+    return f"_grouped_moe_gemm_cuda_32e9b2c::{op_name}"

build/torch212-cxx11-cu130-x86_64-linux/grouped_moe_gemm/__init__.py

@@ -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

@@ -0,0 +1,33 @@
+{
+  "name": "grouped-moe-gemm",
+  "id": "_grouped_moe_gemm_cuda_32e9b2c",
+  "version": 1,
+  "license": "Apache-2.0",
+  "python-depends": [],
+  "backend": {
+    "type": "cuda",
+    "archs": [
+      "12.0a"
+    ]
+  },
+  "digest": {
+    "algorithm": "sha256",
+    "files": {
+      "__init__.py": "/pd/JqQCw5Ie4nddrEvVeqaOFtlbkqNDqykgQLlZJF0=",
+      "_grouped_moe_gemm_cuda_32e9b2c.abi3.so": "sLUEi8xVIn8j82Bv/G/lLNg1kpMQTG8OXO6IZYE6r2I=",
+      "_ops.py": "AxmxTU/RwyGFkTVT0IV6Y0/bjHdcILOm6O5JNy/zf6k=",
+      "grouped_moe_gemm/__init__.py": "DFYPlrhXwYjEqCl/8n0SmWGZV8NFml5DPhMjKfv98GY="
+    }
+  },
+  "provenance": {
+    "kernel-builder": {
+      "version": "0.17.0-dev0",
+      "sha": "19aaa6421e674e9fecc352bbae6eab81d19a6bf4",
+      "dirty": false
+    },
+    "kernel": {
+      "sha": "32e9b2c39ee268d8a76fcfef667974a13022c283",
+      "dirty": false
+    }
+  }
+}

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

@@ -0,0 +1,112 @@
+"""Grouped NVFP4 MoE GEMM for Blackwell prefill workloads."""
+
+from __future__ import annotations
+import torch
+from ._ops import add_op_namespace_prefix, ops
+
+
+@torch.library.custom_op(
+    add_op_namespace_prefix("_grouped_nvfp4_gemm_bf16"),
+    mutates_args=(),
+    device_types="cuda",
+)
+def _compileable(
+    input: torch.Tensor,
+    weight: torch.Tensor,
+    input_scale: torch.Tensor,
+    weight_scale: torch.Tensor,
+    alpha: torch.Tensor,
+    tile_expert: torch.Tensor,
+    tile_rows: int,
+    input_scale_stride: int,
+    weight_stride: int,
+    weight_scale_stride: int,
+) -> torch.Tensor:
+    output = torch.empty(
+        (input.shape[0], weight.shape[1]), device=input.device, dtype=torch.bfloat16
+    )
+    ops.grouped_nvfp4_gemm_bf16_out(
+        input,
+        weight,
+        input_scale,
+        weight_scale,
+        alpha,
+        tile_expert,
+        tile_rows,
+        input_scale_stride,
+        weight_stride,
+        weight_scale_stride,
+        output,
+    )
+    return output
+
+
+@torch.library.register_fake(add_op_namespace_prefix("_grouped_nvfp4_gemm_bf16"))
+def _fake(
+    input,
+    weight,
+    input_scale,
+    weight_scale,
+    alpha,
+    tile_expert,
+    tile_rows,
+    input_scale_stride,
+    weight_stride,
+    weight_scale_stride,
+):
+    if input.ndim != 2 or weight.ndim != 3:
+        raise RuntimeError("invalid grouped GEMM tensor shapes")
+    return torch.empty(
+        (input.shape[0], weight.shape[1]), device=input.device, dtype=torch.bfloat16
+    )
+
+
+def grouped_nvfp4_gemm_bf16(
+    input,
+    weight,
+    input_scale,
+    weight_scale,
+    alpha,
+    tile_expert,
+    *,
+    tile_rows,
+    input_scale_stride=0,
+    weight_stride=None,
+    weight_scale_stride=None,
+    out=None,
+):
+    """Compute expert-selected packed NVFP4 GEMM tiles with BF16 output."""
+    if weight_stride is None:
+        weight_stride = weight[0].numel()
+    if weight_scale_stride is None:
+        weight_scale_stride = weight_scale[0].numel()
+    if out is None:
+        return _compileable(
+            input,
+            weight,
+            input_scale,
+            weight_scale,
+            alpha,
+            tile_expert,
+            int(tile_rows),
+            int(input_scale_stride),
+            int(weight_stride),
+            int(weight_scale_stride),
+        )
+    ops.grouped_nvfp4_gemm_bf16_out(
+        input,
+        weight,
+        input_scale,
+        weight_scale,
+        alpha,
+        tile_expert,
+        int(tile_rows),
+        int(input_scale_stride),
+        int(weight_stride),
+        int(weight_scale_stride),
+        out,
+    )
+    return out
+
+
+__all__ = ["grouped_nvfp4_gemm_bf16"]

build/torch212-cxx11-cu132-x86_64-linux/_grouped_moe_gemm_cuda_32e9b2c.abi3.so

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:bceb0dc19dbe93563d630998a23384b2fe27896c9b72f5488886bac02c2cc471
+size 194144

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

@@ -0,0 +1,9 @@
+import torch
+from . import _grouped_moe_gemm_cuda_32e9b2c
+ops = torch.ops._grouped_moe_gemm_cuda_32e9b2c
+
+def add_op_namespace_prefix(op_name: str):
+    """
+    Prefix op by namespace.
+    """
+    return f"_grouped_moe_gemm_cuda_32e9b2c::{op_name}"

build/torch212-cxx11-cu132-x86_64-linux/grouped_moe_gemm/__init__.py

@@ -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

@@ -0,0 +1,33 @@
+{
+  "name": "grouped-moe-gemm",
+  "id": "_grouped_moe_gemm_cuda_32e9b2c",
+  "version": 1,
+  "license": "Apache-2.0",
+  "python-depends": [],
+  "backend": {
+    "type": "cuda",
+    "archs": [
+      "12.0a"
+    ]
+  },
+  "digest": {
+    "algorithm": "sha256",
+    "files": {
+      "__init__.py": "/pd/JqQCw5Ie4nddrEvVeqaOFtlbkqNDqykgQLlZJF0=",
+      "_grouped_moe_gemm_cuda_32e9b2c.abi3.so": "vOsNwZ2+k1Y9YwmYojOEsv4niWybcvVIiIa6wCwsxHE=",
+      "_ops.py": "AxmxTU/RwyGFkTVT0IV6Y0/bjHdcILOm6O5JNy/zf6k=",
+      "grouped_moe_gemm/__init__.py": "DFYPlrhXwYjEqCl/8n0SmWGZV8NFml5DPhMjKfv98GY="
+    }
+  },
+  "provenance": {
+    "kernel-builder": {
+      "version": "0.17.0-dev0",
+      "sha": "19aaa6421e674e9fecc352bbae6eab81d19a6bf4",
+      "dirty": false
+    },
+    "kernel": {
+      "sha": "32e9b2c39ee268d8a76fcfef667974a13022c283",
+      "dirty": false
+    }
+  }
+}

build/torch213-cxx11-cu130-x86_64-linux/__init__.py

@@ -0,0 +1,112 @@
+"""Grouped NVFP4 MoE GEMM for Blackwell prefill workloads."""
+
+from __future__ import annotations
+import torch
+from ._ops import add_op_namespace_prefix, ops
+
+
+@torch.library.custom_op(
+    add_op_namespace_prefix("_grouped_nvfp4_gemm_bf16"),
+    mutates_args=(),
+    device_types="cuda",
+)
+def _compileable(
+    input: torch.Tensor,
+    weight: torch.Tensor,
+    input_scale: torch.Tensor,
+    weight_scale: torch.Tensor,
+    alpha: torch.Tensor,
+    tile_expert: torch.Tensor,
+    tile_rows: int,
+    input_scale_stride: int,
+    weight_stride: int,
+    weight_scale_stride: int,
+) -> torch.Tensor:
+    output = torch.empty(
+        (input.shape[0], weight.shape[1]), device=input.device, dtype=torch.bfloat16
+    )
+    ops.grouped_nvfp4_gemm_bf16_out(
+        input,
+        weight,
+        input_scale,
+        weight_scale,
+        alpha,
+        tile_expert,
+        tile_rows,
+        input_scale_stride,
+        weight_stride,
+        weight_scale_stride,
+        output,
+    )
+    return output
+
+
+@torch.library.register_fake(add_op_namespace_prefix("_grouped_nvfp4_gemm_bf16"))
+def _fake(
+    input,
+    weight,
+    input_scale,
+    weight_scale,
+    alpha,
+    tile_expert,
+    tile_rows,
+    input_scale_stride,
+    weight_stride,
+    weight_scale_stride,
+):
+    if input.ndim != 2 or weight.ndim != 3:
+        raise RuntimeError("invalid grouped GEMM tensor shapes")
+    return torch.empty(
+        (input.shape[0], weight.shape[1]), device=input.device, dtype=torch.bfloat16
+    )
+
+
+def grouped_nvfp4_gemm_bf16(
+    input,
+    weight,
+    input_scale,
+    weight_scale,
+    alpha,
+    tile_expert,
+    *,
+    tile_rows,
+    input_scale_stride=0,
+    weight_stride=None,
+    weight_scale_stride=None,
+    out=None,
+):
+    """Compute expert-selected packed NVFP4 GEMM tiles with BF16 output."""
+    if weight_stride is None:
+        weight_stride = weight[0].numel()
+    if weight_scale_stride is None:
+        weight_scale_stride = weight_scale[0].numel()
+    if out is None:
+        return _compileable(
+            input,
+            weight,
+            input_scale,
+            weight_scale,
+            alpha,
+            tile_expert,
+            int(tile_rows),
+            int(input_scale_stride),
+            int(weight_stride),
+            int(weight_scale_stride),
+        )
+    ops.grouped_nvfp4_gemm_bf16_out(
+        input,
+        weight,
+        input_scale,
+        weight_scale,
+        alpha,
+        tile_expert,
+        int(tile_rows),
+        int(input_scale_stride),
+        int(weight_stride),
+        int(weight_scale_stride),
+        out,
+    )
+    return out
+
+
+__all__ = ["grouped_nvfp4_gemm_bf16"]

build/torch213-cxx11-cu130-x86_64-linux/_grouped_moe_gemm_cuda_32e9b2c.abi3.so

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:3cb23ee178877d6ab527694d23c227f9ed6272fd48f8befa080b15107ef6e2a6
+size 193928

build/torch213-cxx11-cu130-x86_64-linux/_ops.py

@@ -0,0 +1,9 @@
+import torch
+from . import _grouped_moe_gemm_cuda_32e9b2c
+ops = torch.ops._grouped_moe_gemm_cuda_32e9b2c
+
+def add_op_namespace_prefix(op_name: str):
+    """
+    Prefix op by namespace.
+    """
+    return f"_grouped_moe_gemm_cuda_32e9b2c::{op_name}"

build/torch213-cxx11-cu130-x86_64-linux/grouped_moe_gemm/__init__.py

@@ -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/torch213-cxx11-cu130-x86_64-linux/metadata.json

@@ -0,0 +1,33 @@
+{
+  "name": "grouped-moe-gemm",
+  "id": "_grouped_moe_gemm_cuda_32e9b2c",
+  "version": 1,
+  "license": "Apache-2.0",
+  "python-depends": [],
+  "backend": {
+    "type": "cuda",
+    "archs": [
+      "12.0a"
+    ]
+  },
+  "digest": {
+    "algorithm": "sha256",
+    "files": {
+      "__init__.py": "/pd/JqQCw5Ie4nddrEvVeqaOFtlbkqNDqykgQLlZJF0=",
+      "_grouped_moe_gemm_cuda_32e9b2c.abi3.so": "PLI+4XiHfWq1J2lNI8In+e1icv1I+L76CAsVEH724qY=",
+      "_ops.py": "AxmxTU/RwyGFkTVT0IV6Y0/bjHdcILOm6O5JNy/zf6k=",
+      "grouped_moe_gemm/__init__.py": "DFYPlrhXwYjEqCl/8n0SmWGZV8NFml5DPhMjKfv98GY="
+    }
+  },
+  "provenance": {
+    "kernel-builder": {
+      "version": "0.17.0-dev0",
+      "sha": "19aaa6421e674e9fecc352bbae6eab81d19a6bf4",
+      "dirty": false
+    },
+    "kernel": {
+      "sha": "32e9b2c39ee268d8a76fcfef667974a13022c283",
+      "dirty": false
+    }
+  }
+}

build/torch213-cxx11-cu132-x86_64-linux/__init__.py

@@ -0,0 +1,112 @@
+"""Grouped NVFP4 MoE GEMM for Blackwell prefill workloads."""
+
+from __future__ import annotations
+import torch
+from ._ops import add_op_namespace_prefix, ops
+
+
+@torch.library.custom_op(
+    add_op_namespace_prefix("_grouped_nvfp4_gemm_bf16"),
+    mutates_args=(),
+    device_types="cuda",
+)
+def _compileable(
+    input: torch.Tensor,
+    weight: torch.Tensor,
+    input_scale: torch.Tensor,
+    weight_scale: torch.Tensor,
+    alpha: torch.Tensor,
+    tile_expert: torch.Tensor,
+    tile_rows: int,
+    input_scale_stride: int,
+    weight_stride: int,
+    weight_scale_stride: int,
+) -> torch.Tensor:
+    output = torch.empty(
+        (input.shape[0], weight.shape[1]), device=input.device, dtype=torch.bfloat16
+    )
+    ops.grouped_nvfp4_gemm_bf16_out(
+        input,
+        weight,
+        input_scale,
+        weight_scale,
+        alpha,
+        tile_expert,
+        tile_rows,
+        input_scale_stride,
+        weight_stride,
+        weight_scale_stride,
+        output,
+    )
+    return output
+
+
+@torch.library.register_fake(add_op_namespace_prefix("_grouped_nvfp4_gemm_bf16"))
+def _fake(
+    input,
+    weight,
+    input_scale,
+    weight_scale,
+    alpha,
+    tile_expert,
+    tile_rows,
+    input_scale_stride,
+    weight_stride,
+    weight_scale_stride,
+):
+    if input.ndim != 2 or weight.ndim != 3:
+        raise RuntimeError("invalid grouped GEMM tensor shapes")
+    return torch.empty(
+        (input.shape[0], weight.shape[1]), device=input.device, dtype=torch.bfloat16
+    )
+
+
+def grouped_nvfp4_gemm_bf16(
+    input,
+    weight,
+    input_scale,
+    weight_scale,
+    alpha,
+    tile_expert,
+    *,
+    tile_rows,
+    input_scale_stride=0,
+    weight_stride=None,
+    weight_scale_stride=None,
+    out=None,
+):
+    """Compute expert-selected packed NVFP4 GEMM tiles with BF16 output."""
+    if weight_stride is None:
+        weight_stride = weight[0].numel()
+    if weight_scale_stride is None:
+        weight_scale_stride = weight_scale[0].numel()
+    if out is None:
+        return _compileable(
+            input,
+            weight,
+            input_scale,
+            weight_scale,
+            alpha,
+            tile_expert,
+            int(tile_rows),
+            int(input_scale_stride),
+            int(weight_stride),
+            int(weight_scale_stride),
+        )
+    ops.grouped_nvfp4_gemm_bf16_out(
+        input,
+        weight,
+        input_scale,
+        weight_scale,
+        alpha,
+        tile_expert,
+        int(tile_rows),
+        int(input_scale_stride),
+        int(weight_stride),
+        int(weight_scale_stride),
+        out,
+    )
+    return out
+
+
+__all__ = ["grouped_nvfp4_gemm_bf16"]

build/torch213-cxx11-cu132-x86_64-linux/_grouped_moe_gemm_cuda_32e9b2c.abi3.so

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:aaa4d78336ddd1b195d8cb02d379da8e4b253d756a0624228f910a79382dae95
+size 193984

build/torch213-cxx11-cu132-x86_64-linux/_ops.py

@@ -0,0 +1,9 @@
+import torch
+from . import _grouped_moe_gemm_cuda_32e9b2c
+ops = torch.ops._grouped_moe_gemm_cuda_32e9b2c
+
+def add_op_namespace_prefix(op_name: str):
+    """
+    Prefix op by namespace.
+    """
+    return f"_grouped_moe_gemm_cuda_32e9b2c::{op_name}"

build/torch213-cxx11-cu132-x86_64-linux/grouped_moe_gemm/__init__.py

@@ -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/torch213-cxx11-cu132-x86_64-linux/metadata.json

@@ -0,0 +1,33 @@
+{
+  "name": "grouped-moe-gemm",
+  "id": "_grouped_moe_gemm_cuda_32e9b2c",
+  "version": 1,
+  "license": "Apache-2.0",
+  "python-depends": [],
+  "backend": {
+    "type": "cuda",
+    "archs": [
+      "12.0a"
+    ]
+  },
+  "digest": {
+    "algorithm": "sha256",
+    "files": {
+      "__init__.py": "/pd/JqQCw5Ie4nddrEvVeqaOFtlbkqNDqykgQLlZJF0=",
+      "_grouped_moe_gemm_cuda_32e9b2c.abi3.so": "qqTXgzbd0bGV2MsC03najkslPXVqBiQij5EKeTgtrpU=",
+      "_ops.py": "AxmxTU/RwyGFkTVT0IV6Y0/bjHdcILOm6O5JNy/zf6k=",
+      "grouped_moe_gemm/__init__.py": "DFYPlrhXwYjEqCl/8n0SmWGZV8NFml5DPhMjKfv98GY="
+    }
+  },
+  "provenance": {
+    "kernel-builder": {
+      "version": "0.17.0-dev0",
+      "sha": "19aaa6421e674e9fecc352bbae6eab81d19a6bf4",
+      "dirty": false
+    },
+    "kernel": {
+      "sha": "32e9b2c39ee268d8a76fcfef667974a13022c283",
+      "dirty": false
+    }
+  }
+}