Skip to content

[blas][cublas] Support int8 inputs with float output in gemm_batch - #761

Open
zjin-lcf wants to merge 2 commits into
uxlfoundation:developfrom
zjin-lcf:feature/cublas-int8-gemm-batch
Open

[blas][cublas] Support int8 inputs with float output in gemm_batch#761
zjin-lcf wants to merge 2 commits into
uxlfoundation:developfrom
zjin-lcf:feature/cublas-int8-gemm-batch

Conversation

@zjin-lcf

Copy link
Copy Markdown
Contributor

Description

Enables the (int8, int8, float, float) gemm_batch combination on the cuBLAS backend for the column-major layout, in all three entry points: buffer strided, USM strided, and USM group.

No new implementation code was required. cublasGemmStridedBatchedEx and cublasGemmBatchedEx already accept the datatypes that the existing launchers forward, so this combination only needed to be routed to gemm_batch_impl instead of to the macro that throws unimplemented.

(int8, int8, int32, float) remains unimplemented, and the reason is now recorded in a comment next to it: cuBLAS produces an int32 output only under CUBLAS_COMPUTE_32I, which requires int32 alpha and beta, whereas oneMath specifies float scalars for this combination. Row-major remains unimplemented as before.

Test tolerance

Enabling the type combination alone makes the existing gemm_batch tests fail on a handful of output entries, and the cause is the tolerance model rather than the backend.

A float output accumulated from int8 inputs is rounded at the magnitude of the terms being summed, |alpha| * sum|a*b|, which for int8 data runs orders of magnitude above the output entries themselves. Where the products mostly cancel, the entry is far smaller than the terms that produced it, and no relative bound can cover it: this is ordinary cancellation, not a backend defect. Verified against an exact int64 reference, the largest observed discrepancy is well inside what single precision permits for an accumulation of that size.

The shared checker therefore takes an optional absolute tolerance:

  • check_equal accepts abs_bound, defaulted to 0.0 and threaded through check_equal_matrix and check_almost_equal_matrix. With the default, abs_bound cannot loosen anything, so all existing callers behave exactly as before. The integral overload accepts and ignores it, since an integer result must still match exactly.
  • The three gemm_batch tests compute eps * |alpha| * k * 128 * 128 under if constexpr, so only the int8-to-float instantiation is affected. Since every int8 magnitude is below 128, k * 128 * 128 is an upper bound on sum|a*b|. The relative bound stays at 10 * k * eps.

The constant in this tolerance is calibrated from measurements rather than derived, so it is a scale-based bound supported by evidence. For what it is worth, it sits far below the deterministic gamma_k * sum|a*b| worst case, and the entries that rely on it use only a few percent of it.

Regression test

Int8Int8SinglePrecisionErrorModel covers this path deterministically. Test sizes and data come from an in-test generator rather than std::rand(), so it does not depend on the order the tests run in, and the expected result is accumulated exactly in integers, so it does not depend on the reference BLAS either. The leading rows of A and columns of B are built from triples whose products are 5x, -3x and -2x: their exact dot product is zero while the terms summed stay large, and no ratio is a power of two, so rounding does not cancel along with the terms. Each entry is checked against the relative and absolute bounds, and against the accumulation error model the absolute bound is calibrated from.

On the machine below the test reports 122 entries missing the relative bound, 120 of them the constructed cancelling ones, using 18% of the absolute tolerance. Shrinking that tolerance by 1000x makes the test fail, so the assertion is doing work rather than passing vacuously.

One caveat worth flagging for reviewers: shapes are not interchangeable in this test, because a backend may accumulate a given shape exactly, in which case no tolerance is needed and nothing exercises this one. The test prints how many entries relied on the absolute bound so that this stays visible.

Fixes #506

Checklist

All Submissions

  • Do all unit tests pass locally? Attach a log.

Built with the DPC++ compiler and the cuBLAS backend, run on an NVIDIA A100-SXM4-40GB.

gemm_batch tests, compile-time and run-time dispatch:

$ ./bin/test_main_blas_ct --gtest_filter=*GemmBatch*
int8 accumulation error reached 0.496482 of the accumulated magnitude the model allows and
0.178236 of the absolute tolerance; 122 entries missed the relative bound, 120 of them cancelling
[==========] 52 tests from 4 test suites ran. (9122 ms total)
[  PASSED  ] 23 tests.
[  SKIPPED ] 29 tests

$ ./bin/test_main_blas_rt --gtest_filter=*GemmBatch*
[==========] 52 tests from 4 test suites ran. (10176 ms total)
[  PASSED  ] 23 tests.
[  SKIPPED ] 29 tests

The skips are the row-major and other unimplemented cuBLAS combinations, unchanged by this PR.

Wider BLAS suite, both dispatch modes:

$ ./bin/test_main_blas_ct --gtest_filter=-*ComplexDouble*
[  PASSED  ] 277 tests.
[  FAILED  ] 0 tests.

$ ./bin/test_main_blas_rt --gtest_filter=-*ComplexDouble*
[  PASSED  ] 277 tests.
[  FAILED  ] 0 tests.

Complex double level 1 tests are excluded because they crash identically on this setup with and without these changes, in an untouched build of develop as well, so the failure predates this PR and is unrelated to it.

Because the suite does not seed std::rand(), test sizes depend on execution order, so the filtered and wider runs above exercise two different draws of sizes and data.

New features

  • Have you provided motivation for adding a new feature? See Implement CuBlas/MKL int8, float mixed precision gemm_batch #506.
  • Have you added relevant tests? The existing gemm_batch tests already cover this type combination and now run against the cuBLAS backend rather than skipping; Int8Int8SinglePrecisionErrorModel was added for the tolerance itself.

Bug fixes

Made with Cursor

cuBLAS reaches this combination through cublasGemmStridedBatchedEx and
cublasGemmBatchedEx, which already accept the datatypes the existing
launchers forward, so the column-major buffer, USM strided and USM group
entry points only needed routing to the implementation instead of
throwing unimplemented. The int32 output combination stays unimplemented
because cuBLAS produces it only under CUBLAS_COMPUTE_32I, which takes
int32 alpha and beta, whereas oneMath specifies float scalars.

A float output accumulated from int8 inputs is rounded at the magnitude
of the terms summed rather than at the magnitude of the output entry, so
an entry whose sum cancels cannot meet any relative bound. The shared
checker takes an optional absolute tolerance, defaulted to zero so that
existing callers are unaffected, and the int8-to-float gemm_batch tests
pass eps times k * 128 * 128, an upper bound on the accumulated
magnitude sum|a*b|. Int8Int8SinglePrecisionErrorModel covers that path
with fixed data whose leading rows and columns cancel exactly.

Co-authored-by: Cursor <cursoragent@cursor.com>
@zjin-lcf
zjin-lcf requested a review from a team as a code owner August 14, 2026 21:45
…del bound

An entry whose terms and stored C value are all zero gives a zero model
bound, which the reported usage ratio would divide by.

Co-authored-by: Cursor <cursoragent@cursor.com>
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

Implement CuBlas/MKL int8, float mixed precision gemm_batch

1 participant