exact-gemm

exact-gemm computes the exact product of integer matrices on INT8 tensor cores, using the residue (Chinese Remainder Theorem) scheme, loadable through kernels. Entries to about 160 bits are supported; the pipeline is integer-only and the output is the exact product. The reference baselines are a Python-bigint oracle and FLINT fmpz_mat, against which it measures 43x to 740x. It uses the same residue engine as fp64-emu and fp128-emu.

Exact integer matrix multiplication, the arithmetic under lattice reduction, Gram matrices, and computer algebra, has been a CPU operation provided by multiprecision libraries. Floating point cannot do it: a single rounded bit invalidates the result. This kernel does it on GPU tensor cores by splitting every entry into residues modulo small primes, multiplying each residue plane as an INT8 GEMM, and reconstructing the exact integers, so a product that occupies a CPU core for seconds returns in milliseconds, digit for digit identical.

The Gram matrix X @ X^T at n = 256 with 128-bit entries, run in real time: 1.1 ms on the tensor cores against 1.16 s for GMP on one CPU core of the same machine, all 65,536 entries digit-for-digit equal, 1,094x.

Usage

import torch
from kernels import get_kernel

emu = get_kernel("phanerozoic/exact-gemm", version=1, trust_remote_code=True)

# plain int64 matrices (|entries| < 2^63); returns the exact product as limbs
A = torch.randint(-10**9, 10**9, (256, 256), dtype=torch.int64, device="cuda")
B = torch.randint(-10**9, 10**9, (256, 256), dtype=torch.int64, device="cuda")
C = emu.mm_i64(A, B)                       # [Lc, 256, 256] int64 limbs, C = A @ B^T

# big integers: pack python-int matrices, multiply, unpack
Alimbs = emu.pack(big_int_matrix)          # -> [L, M, K] int64 limb tensor on cuda
C = emu.mm(Alimbs, Blimbs)
result = emu.unpack(C)                      # -> list of lists of python ints

version selects the release branch; trust_remote_code is required by kernels for publishers without the trusted-publisher mark. The prime count is chosen automatically from operand size and K; torch._int_mm shape constraints (M > 16; K, N multiples of 8) are handled by zero-padding.

API

Method

Integer matrices are stored as little-endian int64 limbs in two's complement: [L, M, K], element (i, k) = sum_j limbs[j,i,k] * 2^(64j), the top limb carrying the sign. For each prime p < 256 (so balanced residues fit int8), the residue of an L-limb value is assembled from its 32-bit sub-limbs against a precomputed table 2^(32m) mod p, with a sign correction:

res(T) = ( sum_m subT_m * (2^(32m) mod p) ) - sign * (2^(64L) mod p)   (mod p)

One torch._int_mm INT8 tensor-core GEMM per prime gives C mod p. Balanced-Garner reconstruction produces the mixed-radix digits, and a Horner evaluation in signed 384-bit two's-complement integer arithmetic recovers the product, written out as [Lc, M, N] int64 limbs. No floating-point step and no rounding appear anywhere; the result is exact whenever the operands satisfy the range rule.

Range rule: the signed product must lie inside the CRT range. The wrapper selects the prime count so prod(primes) > 2|product|, with |product| < 2^(bitsA + bitsB + ceil(log2 K)). The full table (all 54 primes below 256) spans about 2^351, covering operands to ~160 bits at K = 2^17; the int32 GEMM accumulator sets that K bound (K * 125^2 < 2^31). Smaller entries use fewer GEMMs.

Measured

Against single-threaded FLINT (fmpz_mat) on the same host, computing the Gram matrix X @ X^T of a random integer basis; the two agreed in every case.

The same cases run 99x to 455x on H200 and 43x to 67x on L4.

Throughput (int-op rate = 2 n^3 / time, square operands):

Most of the runtime is the residue-extraction stage rather than the INT8 GEMMs; mm_timed reports the per-phase split.

Correctness

Tests assert exact integer equality against a python-bigint oracle and against FLINT fmpz_mat on every gated build: signed operands at 40, 100, 128, and 160 bits; asymmetric magnitudes with zeros; the aligned all-ones case at K = 4096; odd shapes (padded); and a lattice Gram matrix X @ X^T. The pipeline is deterministic.

Requirements and limits

• NVIDIA GPU with compute capability 8.0+ and torch with _int_mm.
• int64 limb tensors; up to 4 limbs; operands to ~160 bits; K <= 2^17.
• Working memory scales with the prime count: about P * (M*K + N*K) int8 plus P * M*N int32; roughly 50 GB at n = 16384 with 128-bit operands.
• The output limb count Lc is sized to the product; operands beyond the range rule raise before launch.

References

Ozaki, Ogita, Oishi, Rump 2012 (Numerical Algorithms 59); Ozaki, Uchino, Imamura 2025 (Ozaki Scheme II, arXiv:2504.08009; GEMMul8, RIKEN-RCCS); Hart et al., FLINT (Fast Library for Number Theory).

License

Apache-2.0.