sgemm-bi for PyTorch

Deterministic, batch-invariant CUDA GEMM for PyTorch training: bit-identical matmuls (forward / dW / dX) in float32, bfloat16, and float16, with an opt-in tensor-core tier that is faster than cuBLAS PEDANTIC on transformer-class shapes.

Built on the sgemm-bi Rust engine. No libtorch linkage — tensors cross as raw device pointers, so one wheel works with any PyTorch build. Kernels compile once at engine creation via NVRTC; the runtime needs only the NVIDIA driver.

Requirements: Linux x86_64, NVIDIA Ampere or newer (sm_80+), PyTorch with CUDA. (No macOS/Windows — there is no CUDA there; the package refuses to import with a clear message.)

Install

pip install maturin
cd python && maturin build --release
pip install target/wheels/sgemm_bi-*.whl

Use

import torch, sgemm_bi

# Drop-in layer (weight stored [in, out] — GEMM-natural; convert
# existing layers with Linear.from_torch):
layer = sgemm_bi.Linear(768, 3072, device="cuda", dtype=torch.bfloat16,
                        tensor_cores=True)
y = layer(x)
y.sum().backward()   # deterministic dW (f32-accumulated) and dX

# Functional form:
y = sgemm_bi.deterministic_linear(x, weight, bias, tensor_cores=True)

# Low-level engine (raw pointers, explicit stream):
eng = sgemm_bi.Engine(0)
eng.forward(y.data_ptr(), x.data_ptr(), w.data_ptr(), None,
            m, k, n, "bfloat16",
            torch.cuda.current_stream().cuda_stream, True)

Determinism contracts

tier guarantee
f32 / typed scalar bit-identical across runs; bf16/f16 ≡ "upcast → f32 kernel → RNE downcast"; full shape coverage
tensor cores (tensor_cores=True) own deterministic contract (mma.sync, f32 accumulate); bit-identical across runs; forward strictly batch-invariant across all M; falls back to the scalar tier when an output dim is < 64 (shape-only dispatch — still deterministic; two bit-identical tile families, 128×128 and 64×64, cover everything above)

Bias gradient is a plain f32 column sum done in torch (deterministic run-to-run for a fixed shape; not part of the engine's batch-invariance contract).

Examples

examples/train_deterministic.py trains twice from one seed and asserts bit-identical weights, then checks batch invariance. Typed stubs ship in the wheel (py.typed) — IDEs autocomplete and document the native Engine class.

Tests

pip install pytest && pytest tests/ -v   # needs a CUDA GPU

Project links

Data verified by PyPI on Aug 2, 2026
Data provided by the project maintainers, verified at the time the release was uploaded to PyPI.

Key dates

PyPI data
Data sourced directly from PyPI's database.
  • Released:
Latest release

1 maintainer

PyPI data
Data sourced directly from PyPI's database.
Avatar for silvermpx from gravatar.com silvermpx

Credits

Author: silvermpx

GitHub Statistics

Data verified by PyPI on Aug 2, 2026
The GitHub source repository was provided by the project maintainers and verified by PyPI at the time of upload. Stars, forks, and open issues/PRs are derived from that repository and have not been independently verified.

License expression

MIT OR Apache-2.0
View SPDX License List

Requires

Python >=3.9

Tags

cuda gemm deterministic pytorch reproducibility

Classifiers

Report project as malware

Download files

Download the file for your platform. If you're not sure which to choose, learn more about installing packages.

Source Distribution

sgemm_bi-0.1.2.tar.gz (121.0 kB view details)

Uploaded Source

Built Distribution

If you're not sure about the file name format, learn more about wheel file names.

sgemm_bi-0.1.2-cp39-abi3-manylinux_2_34_x86_64.whl (469.1 kB view details)

Uploaded CPython 3.9+manylinux: glibc 2.34+ x86-64

File details

Details for the file sgemm_bi-0.1.2.tar.gz.

File metadata

  • Download URL: sgemm_bi-0.1.2.tar.gz
  • Upload date:
  • Size: 121.0 kB
  • Tags: Source
  • Uploaded using Trusted Publishing? Yes
  • Uploaded via: twine/7.0.0 CPython/3.13.14

File hashes

Hashes for sgemm_bi-0.1.2.tar.gz
Algorithm Hash digest
SHA256 31414a9594dba0ccbd4558d6a871b8c4f122eda6744ec15919bc3f4a38d80d84
MD5 cdde932c2dfcf04c0d5afd1b35ef4f1a
BLAKE2b-256 118203979af9909c3647fef603f5ab6514c65fd588b463d6bd2ed65fce1d0561

See more details on using hashes here.

Provenance

The following attestation bundles were made for sgemm_bi-0.1.2.tar.gz:

Publisher: release.yml on silvermpx/sgemm-bi

Attestations: Values shown here reflect the state when the release was signed and may no longer be current.

File details

Details for the file sgemm_bi-0.1.2-cp39-abi3-manylinux_2_34_x86_64.whl.

File metadata

File hashes

Hashes for sgemm_bi-0.1.2-cp39-abi3-manylinux_2_34_x86_64.whl
Algorithm Hash digest
SHA256 dbee301eb11c0df70daa2f7f3d2c09697a91375bbc85ad0dd291f282a9721265
MD5 933d7de37af4382bb7fd10c48808f054
BLAKE2b-256 ef88a50f51dbb1d943aa85d5a5344ef21fba76cd6152010e17dde6aed0e21c10

See more details on using hashes here.

Provenance

The following attestation bundles were made for sgemm_bi-0.1.2-cp39-abi3-manylinux_2_34_x86_64.whl:

Publisher: release.yml on silvermpx/sgemm-bi

Attestations: Values shown here reflect the state when the release was signed and may no longer be current.

Release history Release notifications | RSS feed

This release

0.1.2 This release

2 files

0.1.1.post2

2 files

0.1.1.post1

2 files

0.1.1

2 files

Anthropic, PBC Visionary sponsor Bloomberg Visionary sponsor Hudson River Trading Visionary sponsor Meta Visionary sponsor NVIDIA Visionary sponsor Microsoft Sustainability sponsor Depot Continuous Integration AWS Cloud computing and Security Sponsor Datadog Monitoring Fastly CDN Google Download Analytics Sentry Error logging StatusPage Status page