opengemm 0.1.0__tar.gz

opengemm-0.1.0/LICENSE

@@ -0,0 +1,21 @@
+MIT License
+
+Copyright (c) 2026 Aditya Ramesh
+
+Permission is hereby granted, free of charge, to any person obtaining a copy
+of this software and associated documentation files (the "Software"), to deal
+in the Software without restriction, including without limitation the rights
+to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
+copies of the Software, and to permit persons to whom the Software is
+furnished to do so, subject to the following conditions:
+
+The above copyright notice and this permission notice shall be included in all
+copies or substantial portions of the Software.
+
+THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
+IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
+FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
+AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
+LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
+OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
+SOFTWARE.

opengemm-0.1.0/MANIFEST.in

@@ -0,0 +1,4 @@
+include LICENSE README.md
+recursive-include opengemm *.cu *.cuh *.json *.jsonc
+global-exclude *.pyc
+global-exclude __pycache__/*

opengemm-0.1.0/PKG-INFO

@@ -0,0 +1,158 @@
+Metadata-Version: 2.4
+Name: opengemm
+Version: 0.1.0
+Summary: Open GEMM kernels for NVIDIA B200 (sm_100a) in CUDA
+Author: Aditya Ramesh
+License: MIT License
+        
+        Copyright (c) 2026 Aditya Ramesh
+        
+        Permission is hereby granted, free of charge, to any person obtaining a copy
+        of this software and associated documentation files (the "Software"), to deal
+        in the Software without restriction, including without limitation the rights
+        to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
+        copies of the Software, and to permit persons to whom the Software is
+        furnished to do so, subject to the following conditions:
+        
+        The above copyright notice and this permission notice shall be included in all
+        copies or substantial portions of the Software.
+        
+        THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
+        IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
+        FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
+        AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
+        LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
+        OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
+        SOFTWARE.
+        
+Project-URL: Homepage, https://github.com/aramesh10/OpenGEMM
+Project-URL: Repository, https://github.com/aramesh10/OpenGEMM
+Project-URL: Issues, https://github.com/aramesh10/OpenGEMM/issues
+Keywords: gemm,cuda,blackwell,sm_100a,kernels,agents
+Classifier: Development Status :: 3 - Alpha
+Classifier: Intended Audience :: Developers
+Classifier: License :: OSI Approved :: MIT License
+Classifier: Programming Language :: Python :: 3
+Classifier: Programming Language :: C++
+Classifier: Environment :: GPU :: NVIDIA CUDA
+Classifier: Topic :: Scientific/Engineering
+Requires-Python: >=3.10
+Description-Content-Type: text/markdown
+License-File: LICENSE
+Requires-Dist: torch>=2.8
+Dynamic: license-file
+
+# OpenGEMM
+
+GEMM kernels for NVIDIA B200 (sm_100a) in CUDA.
+
+```python
+import opengemm as og
+
+c = og.gemm(a, b)                    # C[M, N] = A[M, K] @ B[N, K].T
+c = og.gemm(a, b, sfa, sfb)          # block-scaled: nvfp4, mxfp8, mxfp4
+
+og.emit_kernel(a, b, file="k.cu")    # emits .cu/.cuh for this shape
+c = og.run_kernel("k.cu", a, b)      # compiles emitted kernel and runs it
+```
+
+## Install
+
+From PyPI
+```bash
+pip install opengemm
+```
+
+From a clone:
+```bash
+git clone https://github.com/aramesh10/OpenGEMM.git
+cd OpenGEMM
+pip install -e .
+```
+
+Requirements: 
+- sm_100a
+- CUDA 12.9+ with `nvcc` on the path
+- PyTorch 2.8+.
+
+The first `gemm()` call builds the extension (a few minutes, then cached by torch).
+
+## Agent Quickstart
+
+Give your agent this prompt to use OpenGEMM as a tool:
+
+```
+OpenGEMM emits standalone CUDA GEMM kernels for B200 (sm_100a), no GPU
+needed to emit:
+
+python -c "
+import opengemm as og
+S = dict(m=1024, n=1024, k=1024)
+
+og.emit_kernel(**S, atype='bf16', file='k')        # writes k.cu and k.cuh
+og.emit_kernel(**S, atype='e4m3', btype='e5m2')   # mixed, names itself
+og.emit_kernel(**S, atype='e2m1', sftype='ue4m3') # block-scaled (nvfp4)
+src, hdr = og.emit_kernel(**S, atype='bf16')       # the text, always returned
+print(src, hdr)
+"
+atype / btype: bf16 f16 tf32 s8 u8 e4m3 e5m2 e3m2 e2m3 e2m1
+sftype (block-scaled): ue4m3 (nvfp4) or ue8m0 (mxfp8, mxfp4)
+dtype (output): f32, s32 for s8/u8, bf16 when scaled — inferred, optional.
+```
+
+## Dense and block-scaled
+
+`C[M, N] = A[M, K] @ B[N, K].T`. Both operands are row-major with K innermost.
+
+| GEMM | `atype` / `btype` | `sftype` | `dtype` | `torch.dtype` (in → out) |
+| --- | --- | --- | --- | --- |
+| bfloat16 | bf16 | — | f32 | `bfloat16` → `float32` |
+| float16 | f16 | — | f32 | `float16` → `float32` |
+| tf32 | tf32 | — | f32 | `float32` → `float32` |
+| int8 | s8 | — | s32 | `int8` → `int32` |
+| uint8 | u8 | — | s32 | `uint8` → `int32` |
+| fp8 | e4m3 | — | f32 | `float8_e4m3fn` → `float32` |
+| fp8 | e5m2 | — | f32 | `float8_e5m2` → `float32` |
+| mixed fp8 | e4m3, e5m2 | — | f32 | `float8_e4m3fn`, `float8_e5m2` → `float32` |
+| fp6 | e3m2 | — | f32 | `uint8` → `float32` |
+| fp6 | e2m3 | — | f32 | `uint8` → `float32` |
+| fp4 | e2m1 | — | f32 | `uint8` → `float32` |
+| nvfp4 | e2m1 | ue4m3 (per 16) | bf16 | `float4_e2m1fn_x2`, `float8_e4m3fn` → `bfloat16` |
+| mxfp8 | e4m3 | ue8m0 (per 32) | bf16 | `float8_e4m3fn`, `float8_e8m0fnu` → `bfloat16` |
+| mxfp4 | e2m1 | ue8m0 (per 32) | bf16 | `float4_e2m1fn_x2`, `float8_e8m0fnu` → `bfloat16` |
+
+Note: fp6 and fp4 have no torch dtype. They arrive densely packed in `uint8` and are named - `gemm(a, b, atype="e2m1")` 
+Use `btype=` when the two operands differ.
+
+Input is `[M, K]` and `[N, K]` with column-major strides `(1, M)` and `(1, N)`
+Output is `[M, N]` with column-major strides `(1, M)`
+
+## Tuning and performance
+
+There is no heursitic to choose the config. Optimized configs are stored in `configs.json`. 
+If a particular shape has not been optimized, the library autotunes and returns and saves the best config locally to `./opengemm-configs/tuned_configs.json` or to `OPENGEMM_CONFIGS` env variable.
+
+```bash
+CUDA_VISIBLE_DEVICES=0 python scripts/tune.py --dtype f16 --shape 4096 4096 4096
+CUDA_VISIBLE_DEVICES=0 python scripts/benchmark.py --dtype bf16 e4m3    # vs cuBLAS
+CUDA_VISIBLE_DEVICES=0 python scripts/test.py                           # correctness
+```
+
+`tune.py` ablates every compiled configuration for a shape and records the best performing config to `configs.json`
+
+## Standalone kernels
+
+```bash
+python scripts/emit_kernel.py --dtype e4m3 --shape 4096 4096 4096 --file emitted/e4m3_4k.cu
+python scripts/run_kernel.py emitted/e4m3_4k.cu       # correctness, then timing vs cuBLAS
+```
+
+OpenGEMM can also emit the optimized CUDA files for a kernel given a shape and dtype. It can be ran with `scripts/run_kernel.py` or built with `nvcc`:
+
+```bash
+nvcc -O3 -std=c++20 -gencode=arch=compute_100a,code=sm_100a --expt-relaxed-constexpr -shared -Xcompiler -fPIC -lcuda <KERNEL_FILE>.cu -o <KERNEL_FILE>.so
+```
+
+`emit_kernel` reads only shapes and dtypes, so meta tensors work:
+`emit_kernel(torch.empty(4096, 4096, dtype=torch.bfloat16, device="meta"), ...)`.
+

opengemm-0.1.0/README.md

@@ -0,0 +1,114 @@
+# OpenGEMM
+
+GEMM kernels for NVIDIA B200 (sm_100a) in CUDA.
+
+```python
+import opengemm as og
+
+c = og.gemm(a, b)                    # C[M, N] = A[M, K] @ B[N, K].T
+c = og.gemm(a, b, sfa, sfb)          # block-scaled: nvfp4, mxfp8, mxfp4
+
+og.emit_kernel(a, b, file="k.cu")    # emits .cu/.cuh for this shape
+c = og.run_kernel("k.cu", a, b)      # compiles emitted kernel and runs it
+```
+
+## Install
+
+From PyPI
+```bash
+pip install opengemm
+```
+
+From a clone:
+```bash
+git clone https://github.com/aramesh10/OpenGEMM.git
+cd OpenGEMM
+pip install -e .
+```
+
+Requirements: 
+- sm_100a
+- CUDA 12.9+ with `nvcc` on the path
+- PyTorch 2.8+.
+
+The first `gemm()` call builds the extension (a few minutes, then cached by torch).
+
+## Agent Quickstart
+
+Give your agent this prompt to use OpenGEMM as a tool:
+
+```
+OpenGEMM emits standalone CUDA GEMM kernels for B200 (sm_100a), no GPU
+needed to emit:
+
+python -c "
+import opengemm as og
+S = dict(m=1024, n=1024, k=1024)
+
+og.emit_kernel(**S, atype='bf16', file='k')        # writes k.cu and k.cuh
+og.emit_kernel(**S, atype='e4m3', btype='e5m2')   # mixed, names itself
+og.emit_kernel(**S, atype='e2m1', sftype='ue4m3') # block-scaled (nvfp4)
+src, hdr = og.emit_kernel(**S, atype='bf16')       # the text, always returned
+print(src, hdr)
+"
+atype / btype: bf16 f16 tf32 s8 u8 e4m3 e5m2 e3m2 e2m3 e2m1
+sftype (block-scaled): ue4m3 (nvfp4) or ue8m0 (mxfp8, mxfp4)
+dtype (output): f32, s32 for s8/u8, bf16 when scaled — inferred, optional.
+```
+
+## Dense and block-scaled
+
+`C[M, N] = A[M, K] @ B[N, K].T`. Both operands are row-major with K innermost.
+
+| GEMM | `atype` / `btype` | `sftype` | `dtype` | `torch.dtype` (in → out) |
+| --- | --- | --- | --- | --- |
+| bfloat16 | bf16 | — | f32 | `bfloat16` → `float32` |
+| float16 | f16 | — | f32 | `float16` → `float32` |
+| tf32 | tf32 | — | f32 | `float32` → `float32` |
+| int8 | s8 | — | s32 | `int8` → `int32` |
+| uint8 | u8 | — | s32 | `uint8` → `int32` |
+| fp8 | e4m3 | — | f32 | `float8_e4m3fn` → `float32` |
+| fp8 | e5m2 | — | f32 | `float8_e5m2` → `float32` |
+| mixed fp8 | e4m3, e5m2 | — | f32 | `float8_e4m3fn`, `float8_e5m2` → `float32` |
+| fp6 | e3m2 | — | f32 | `uint8` → `float32` |
+| fp6 | e2m3 | — | f32 | `uint8` → `float32` |
+| fp4 | e2m1 | — | f32 | `uint8` → `float32` |
+| nvfp4 | e2m1 | ue4m3 (per 16) | bf16 | `float4_e2m1fn_x2`, `float8_e4m3fn` → `bfloat16` |
+| mxfp8 | e4m3 | ue8m0 (per 32) | bf16 | `float8_e4m3fn`, `float8_e8m0fnu` → `bfloat16` |
+| mxfp4 | e2m1 | ue8m0 (per 32) | bf16 | `float4_e2m1fn_x2`, `float8_e8m0fnu` → `bfloat16` |
+
+Note: fp6 and fp4 have no torch dtype. They arrive densely packed in `uint8` and are named - `gemm(a, b, atype="e2m1")` 
+Use `btype=` when the two operands differ.
+
+Input is `[M, K]` and `[N, K]` with column-major strides `(1, M)` and `(1, N)`
+Output is `[M, N]` with column-major strides `(1, M)`
+
+## Tuning and performance
+
+There is no heursitic to choose the config. Optimized configs are stored in `configs.json`. 
+If a particular shape has not been optimized, the library autotunes and returns and saves the best config locally to `./opengemm-configs/tuned_configs.json` or to `OPENGEMM_CONFIGS` env variable.
+
+```bash
+CUDA_VISIBLE_DEVICES=0 python scripts/tune.py --dtype f16 --shape 4096 4096 4096
+CUDA_VISIBLE_DEVICES=0 python scripts/benchmark.py --dtype bf16 e4m3    # vs cuBLAS
+CUDA_VISIBLE_DEVICES=0 python scripts/test.py                           # correctness
+```
+
+`tune.py` ablates every compiled configuration for a shape and records the best performing config to `configs.json`
+
+## Standalone kernels
+
+```bash
+python scripts/emit_kernel.py --dtype e4m3 --shape 4096 4096 4096 --file emitted/e4m3_4k.cu
+python scripts/run_kernel.py emitted/e4m3_4k.cu       # correctness, then timing vs cuBLAS
+```
+
+OpenGEMM can also emit the optimized CUDA files for a kernel given a shape and dtype. It can be ran with `scripts/run_kernel.py` or built with `nvcc`:
+
+```bash
+nvcc -O3 -std=c++20 -gencode=arch=compute_100a,code=sm_100a --expt-relaxed-constexpr -shared -Xcompiler -fPIC -lcuda <KERNEL_FILE>.cu -o <KERNEL_FILE>.so
+```
+
+`emit_kernel` reads only shapes and dtypes, so meta tensors work:
+`emit_kernel(torch.empty(4096, 4096, dtype=torch.bfloat16, device="meta"), ...)`.
+

opengemm-0.1.0/opengemm/__init__.py

@@ -0,0 +1,17 @@
+"""opengemm: open GEMM kernels for NVIDIA B200 (sm_100a) in plain CUDA.
+
+    import opengemm as og
+    c = og.gemm(a, b)                      # dense
+    c = og.gemm(a, b, sfa, sfb)            # block-scaled
+    og.emit_kernel(a, b, file="k.cu")      # a standalone .cu/.cuh pair
+    og.run_kernel("k.cu", a, b)            # compile and run it
+"""
+
+from .python.api import dtype_name, gemm
+from .python.dtypes import DTYPES, pack_e2m1, quantize, to_blocked, unpack_e2m1
+from .python.emit import emit_kernel
+from .python.run import run_kernel
+
+__all__ = ["gemm", "emit_kernel", "run_kernel", "dtype_name", "DTYPES",
+           "quantize", "to_blocked", "pack_e2m1", "unpack_e2m1"]
+__version__ = "0.1.0"

opengemm-0.1.0/opengemm/python/__init__.py

@@ -0,0 +1,2 @@
+"""Python side of opengemm: the API, the tuner, the emitter and the
+benchmark harness. `opengemm/__init__.py` re-exports the public names."""

opengemm-0.1.0/opengemm/python/api.py

@@ -0,0 +1,132 @@
+"""C[M, N] = A[M, K] @ B[N, K].T on B200."""
+
+import torch
+
+from .bench import config_args
+from .build import extension
+from .dtypes import DENSE_OF_TORCH, DTYPES, ELEM_INDEX, SCALED_OF_TORCH
+from .tune import resolve_config
+
+# Keyed by (dtype name, M, N, K).
+_launchers = {}
+
+
+def _dense_elem(t, name, which):
+    if name is not None:
+        if name not in ELEM_INDEX:
+            raise ValueError(f"{which}type={name!r}; have {list(ELEM_INDEX)}")
+        return name
+    elem = DENSE_OF_TORCH.get(t.dtype)
+    if elem is None:
+        raise TypeError(
+            f"{which} is {t.dtype}, which names no element: pass "
+            f"{which}type= (uint8 carries u8, e3m2, e2m3 and e2m1 alike; "
+            f"the rest are inferred)")
+    return elem
+
+
+def dtype_name(a, b, sfa=None, atype=None, btype=None):
+    """Return the dtype name for these operands, as configs.json spells it.
+
+    Args:
+        a: Left operand.
+        b: Right operand.
+        sfa: Scales of `a` for a block-scaled call, else None.
+        atype: Element name for a `uint8` operand: `"u8"`, `"e3m2"`, `"e2m3"`
+            or `"e2m1"`.
+        btype: Element name for `b`; defaults to `atype` when the dtypes
+            match.
+
+    Returns:
+        A name such as `"bf16"`, `"e4m3xe5m2"` or `"nvfp4"`.
+
+    Raises:
+        TypeError: If the dtypes name no element or no block-scaled format.
+    """
+    if sfa is not None:
+        if (a.dtype, sfa.dtype) not in SCALED_OF_TORCH:
+            raise TypeError(
+                f"{a.dtype} operands with {sfa.dtype} scales is not a "
+                f"format; have { {v: k for k, v in SCALED_OF_TORCH.items()} }")
+        return SCALED_OF_TORCH[(a.dtype, sfa.dtype)]
+    ea = _dense_elem(a, atype, "a")
+    eb = _dense_elem(b, btype or (atype if b.dtype == a.dtype else None), "b")
+    return ea if ea == eb else f"{ea}x{eb}"
+
+
+def launcher(dtype, m, n, k):
+    """Return the launcher bound to the stored configuration for a shape.
+
+    A shape with no stored configuration is tuned first and the winner stored.
+
+    Args:
+        dtype: Dtype name as configs.json spells it.
+        m: Rows of A and C.
+        n: Rows of B and columns of C.
+        k: Reduction length in elements, not bytes.
+
+    Returns:
+        A callable taking the operands and an optional `out`.
+    """
+    key = (dtype, m, n, k)
+    launch = _launchers.get(key)
+    if launch is None:
+        d = DTYPES[dtype]
+        config = resolve_config(dtype, m, n, k)
+        args = config_args(config, d.impl)
+        if d.impl == "mm":
+            args += (ELEM_INDEX[d.elem_a], ELEM_INDEX[d.elem_b])
+        launch = _launchers[key] = extension(d.impl).launcher(*args)
+    return launch
+
+
+def gemm(a, b, sfa=None, sfb=None, out=None, atype=None, btype=None):
+    """Compute C[M, N] = A[M, K] @ B[N, K].T.
+
+    Both operands are row-major with K innermost. Every call runs a measured
+    configuration from configs.json; a (dtype, shape) with no entry is tuned on
+    its first call, a few minutes on the GPU, and the winner is stored. The
+    first call builds the extension.
+
+    Dense: the element is inferred from the dtype (bfloat16, float16, float32
+    computed as tf32, int8, float8_e4m3fn, float8_e5m2). uint8 holds u8, e3m2,
+    e2m3 or e2m1, fp6 and fp4 packed densely along K with K a multiple of 128,
+    and is named with `atype` and `btype`. Mixed e4m3 x e5m2 is supported.
+
+    Block-scaled: nvfp4 is float4_e2m1fn_x2 operands with float8_e4m3fn scales
+    per 16, mxfp8 is float8_e4m3fn with float8_e8m0fnu per 32, mxfp4 is
+    float4_e2m1fn_x2 with float8_e8m0fnu per 32.
+
+    Args:
+        a: `(M, K)` operand.
+        b: `(N, K)` operand.
+        sfa: Scales of `a` in the 128x4 blocked layout `torch._scaled_mm` takes
+            (`to_blocked` builds it); block-scaled calls only.
+        sfb: Scales of `b`, likewise.
+        out: Output to write into instead of allocating one.
+        atype: Element name for a `uint8` operand.
+        btype: Element name for `b`; defaults to `atype` when the dtypes
+            match.
+
+    Returns:
+        Dense: `(M, N)` float32, or int32 for int8, with column-major strides
+        `(1, M)`, which is how the accumulator leaves tensor memory; call
+        `.contiguous()` for row-major. Block-scaled: `(M, N)` bfloat16,
+        row-major.
+    """
+    if (sfa is None) != (sfb is None):
+        raise ValueError("sfa and sfb go together: both or neither")
+    dtype = dtype_name(a, b, sfa, atype, btype)
+    d = DTYPES[dtype]
+    key = (dtype, a.size(0), b.size(0), d.k_values(a.size(1)))
+    launch = _launchers.get(key)
+    if launch is None:
+        if not a.is_cuda:
+            raise TypeError("a and b must be CUDA tensors")
+        # A sweep allocates and times on the current device; make it the
+        # operands'.
+        with torch.cuda.device(a.device):
+            launch = launcher(*key)
+    if sfa is None:
+        return launch(a, b, out=out)
+    return launch(a, b, sfa, sfb, out=out)