INT8 activation × INT8 weight (W8A8) GEMM with scales grouped along the reduction dimension, K.
Python package: grain_gemm.
For A[M, K] @ B[K, N]:
- Activations: one scale per row and K group.
- Weights: one scale per column and K group.
- Group sizes: 32, 64, 128, and 256.
- Recommended balance of accuracy and inference latency on current targets: G256. See Group-size recommendation.
- INT32 accumulation within each group, followed by scaling and FP32 accumulation across groups.
- Target GPUs: NVIDIA A100, H100, and Jetson Thor (T5000).
For group size G, the scale shapes are [M, ceil(K/G)] for A and [ceil(K/G), N] for B. The intended operation is:
Here, scale_A[m, g] and scale_B[g, n]. The outer weighted sum accumulates in FP32.
Each group covers up to G consecutive elements along K. GrainGEMM uses a configurable block scaling scheme and does not currently claim compatibility with the standard MXINT8 format.
GrainGEMM's INT8 design is motivated by INT v.s. FP: A Comprehensive Study of Fine-Grained Low-bit Quantization Formats, Mengzhao Chen et al. (2025). See also the authors' implementation.
- Theory (Section 4): Smaller blocks typically reduce the peak-to-RMS ratio (crest factor), improving uniform INT8 quantization under the paper's approximate quantization signal-to-noise ratio (QSNR) model.
- 8-bit evidence (Section 5.2): With 32-element blocks and UE8M0 scales, MXINT8 achieves lower KL divergence from BF16 than MXFP8 on all 12 evaluated models.
- Groups 32/64/128/256 (Table 10): These sizes appear in INT8 training ablations of scale precision and symmetric
[-127, 127]clipping. This table compares INT8 quantization recipes, without an FP8 baseline. - Hardware motivation (Section 6): Gate-level modeling estimates lower MXINT8 energy and area than MXFP8 at matched throughput.
These results motivate GrainGEMM's custom K-group scaling. The practical recommendation below also accounts for the cost of applying group scales on existing GPUs.
Based on project experiments, G256 (G = 256) offers a good balance between quantization accuracy and inference latency on A100, H100, and Jetson Thor (T5000), and is our recommended starting point for these GPUs.
These GPUs support INT8 Tensor Core matrix multiplication, but do not expose a native MXINT8 instruction that combines INT8 dot products with per-group scaling. Group scales therefore have to be applied in software. NVIDIA's PTX instruction reference documents the available integer and block-scaled MMA variants.
The software scaling path maintains two accumulator states:
- INT32 partial sums for the current K-group, produced by INT8 Tensor Core operations.
- FP32 running sums that accumulate the scaled contributions from all K-groups.
In a register-based implementation, these occupy two sets of accumulator registers, increasing register pressure. At each group boundary, the INT32 partial sums are converted to FP32, scaled by scale_A[m, g] * scale_B[g, n], and added to the FP32 running sums using CUDA-core arithmetic. The INT32 accumulator is then reused for the next group. The storage depends on the backend: a Thor tcgen05 path can keep MMA accumulators in Tensor Memory and move fragments into registers for software scaling.
For a fixed K, this conversion-and-scaling work occurs ceil(K / G) times. G256 reduces scale loads, conversions, and CUDA-core updates by performing more INT8 Tensor Core work between scaling steps. It does not eliminate the separate accumulator states. Smaller groups can improve quantization accuracy, but require more frequent software scaling on these GPUs; G256 balances these competing costs in our experiments.
On hardware with native MXINT8 support, prefer finer quantization groups supported by the instruction, such as G32. Applying block scales within the matrix-multiply operation reduces the software overhead that motivates G256 on the current targets. The final choice should reflect the hardware's supported granularity, matrix shapes, and the model's accuracy and latency requirements.
GrainGEMM includes a custom Triton kernel and optional native CUDA/CuTe and CUTLASS kernels, with measured G256 dispatch tables for GB10 (SM121). The public API selects a backend and launch configuration by architecture, shape, dtype, and layout. A100, H100, and Thor use the portable Triton path pending validation and dedicated tuning on those GPUs. SM120 native builds are experimental and require explicit backend selection; RTX 50-series correctness and performance have not yet been validated, and auto continues to use Triton on SM120.
Native sources are organized by architecture under
kernels/csrc/sm12x, while measured tuning
tables use GPU-specific names such as gb10_sm121_g256.json. Future H100 and
Thor specializations will live in separate architecture directories; they are
not implemented yet. See source organization.
Run from the repository root with CUDA-enabled PyTorch and an SM80-or-newer NVIDIA GPU:
python -m pip install -e ".[runtime]"This complete example starts with BF16 activations x[M, K] and a linear-layer
weight weight[N, K]. Replace the random tensors with your own activations and
weights. The helper quantizes each row independently in groups along K; it is
example code, not part of the grain_gemm API.
import torch
from grain_gemm import int8_gemm
def quantize_k_groups(x, group_size):
# This simple helper requires complete K groups.
assert x.ndim == 2 and x.shape[1] % group_size == 0
blocks = x.float().reshape(x.shape[0], -1, group_size)
scales = blocks.abs().amax(dim=-1).clamp_min(1e-12) / 127
quantized = (blocks / scales[..., None]).round().clamp(-127, 127)
return quantized.to(torch.int8).reshape(x.shape).contiguous(), scales.contiguous()
M, N, K, G = 1024, 1024, 1024, 256
torch.manual_seed(2026)
x = torch.randn((M, K), device="cuda", dtype=torch.bfloat16)
weight = torch.randn((N, K), device="cuda", dtype=torch.bfloat16)
a, scale_a = quantize_k_groups(x, G) # [M, K], [M, K/G]
weight_int8, weight_scales = quantize_k_groups(weight, G)
b = weight_int8.T # [K, N], column-major view
scale_b = weight_scales.T.contiguous() # [K/G, N], FP32
c = int8_gemm(
a, b, scale_a, scale_b, group_size=G,
output_dtype=torch.bfloat16,
)
print(c.shape, c.dtype) # torch.Size([1024, 1024]) torch.bfloat16int8_gemm returns BF16 by default; output_dtype=torch.bfloat16 above is
optional. Each K-group accumulates in INT32, and the scaled group results
accumulate in FP32 before conversion to the output dtype. Set
output_dtype=torch.float32 or output_dtype=torch.float16 explicitly for
FP32 or FP16 output; these output types use the Triton path with backend="auto".
For each group, the helper uses symmetric INT8 quantization with
scale = max(abs(values)) / 127, a positive floor for zero groups, and
q = round(values / scale) clipped to [-127, 127]. The FP32 scales restore
the approximate original values through values ≈ q * scale.
| Tensor | Dtype | Shape in this example | Meaning |
|---|---|---|---|
a |
INT8 | [M, K] |
Quantized activations from x |
b |
INT8 | [K, N] |
Transposed quantized weights; each column corresponds to one row of weight |
scale_a |
FP32 | [M, K/G] |
One activation scale per row and K group |
scale_b |
FP32 | [K/G, N] |
One weight scale per K group and column of b |
c |
BF16 | [M, N] |
Output allocated and returned by int8_gemm, approximating x @ weight.T |
c includes quantization error relative to multiplying the original BF16 inputs.
Keeping b = weight_int8.T preserves the column-major layout used by the native
backend; the transposed weight scales are made contiguous separately. The helper
above requires K % G == 0; the Triton GEMM implementation also supports partial
final groups with scale shapes [M, ceil(K/G)] and [ceil(K/G), N].
backend="auto" uses the measured GB10 configuration where applicable. The
optional CuTe backend requires a local CUDA build; otherwise the API uses Triton.
backend="triton" or backend="cuda" explicitly selects an implementation.
backend="cutlass" selects the separately built CUTLASS collective implementation;
see build and comparison instructions.
Group sizes 32/64/128/256, positive-stride views, K tails, and FP32/FP16/BF16
outputs are supported through the Triton path. The native fast path currently
supports G256 and BF16 output with aligned row-major A, column-major B,
contiguous scales, M/N multiples of 64, and K a multiple of 256.
See Kernel design and native build for the implementation, CUTLASS build instructions, and dispatch behavior. Inputs must already be quantized; the API does not perform quantization or provide autograd.
To try the existing native kernels on an SM120 device, build the desired backend with the explicit target:
python tools/build_cuda.py --cutlass-dir /path/to/cutlass --arch sm_120
python tools/build_cutlass.py --cutlass-dir /path/to/cutlass --arch sm_120Then pass backend="cuda" for CuTe or backend="cutlass" to int8_gemm.
The native G256, BF16 and layout constraints above still apply. These builds
use fixed compatible defaults instead of the GB10 tuning tables;
backend="auto" stays on Triton for SM120 until measurements justify native
dispatch. The default build target remains sm_121.
Each backend installs one target at a time; rebuilding it replaces its previous library. Runtime requires the build's GPU target to match the device. Builds can run without a GPU, but the shared library must also match the execution machine's CPU architecture and host ABI; a GB10 ARM64 host build is not an x86-64 desktop library. See native build requirements and CUTLASS build details.
RTX 50-series runtime correctness and performance remain unvalidated. Published GB10 benchmark results and experimental measurement scripts retain their GB10-only scope.
See Tests and benchmarks for validation, performance results, plots, and reproduction instructions.
python -m pip install -e .
python -c "import grain_gemm; print(grain_gemm.__version__)"@article{int_vs_fp_2025,
title={INT v.s. FP: A Comprehensive Study of Fine-Grained Low-bit Quantization Formats},
author={Chen, Mengzhao and Wu, Meng and Jin, Hui and Yuan, Zhihang and Liu, Jing and Zhang, Chaoyi and Li, Yunshui and Huang, Jie and Ma, Jin and Xue, Zeyue and Liu, Zhiheng and Bin, Xingyan and Luo, Ping},
journal={arXiv preprint arXiv:2510.25602},
year={2025},
url={https://arxiv.org/abs/2510.25602}
}