Skip to content

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

Merged
ndingle-arm merged 2 commits into
uxlfoundation:developfrom
zjin-lcf:feature/cublas-int8-gemm-batch
Aug 27, 2026
Merged

[blas][cublas] Support int8 inputs with float output in gemm_batch#761
ndingle-arm merged 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>

@melonakos melonakos left a comment

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Approving. This is the most careful piece of work in your current queue, and I want to be specific about why, because the shape of this PR — "enable a type combination, loosen a test tolerance" — is normally a red flag and here it isn't.

The library change is appropriately tiny

All this really does is move GEMM_STRIDED_BATCH_LAUNCHER(std::int8_t, std::int8_t, float, float) out of the throwing macro block into the working one, leaving <int8, int8, int32, float> unimplemented. And you documented why it stays unimplemented:

cuBLAS computes an int32 output only with CUBLAS_COMPUTE_32I, which requires int32 alpha and beta, whereas oneMath specifies float scalars for this combination.

That's the right level of detail. It's a genuine API mismatch rather than an oversight, and the next person to wonder why int32 output is missing now has their answer in the code instead of having to rediscover it.

The tolerance change is principled, not a fudge

Loosening a correctness check to make a new type combination pass is the classic way a wrong kernel gets merged looking right, so I read this closely. Three things convinced me:

  1. It's opt-in and narrowly scoped. abs_bound defaults to 0.0 and abs_limit = max(bound, abs_bound), so every existing caller's behavior is bit-for-bit unchanged. The wider bound is gated behind if constexpr (int8_to_float) on the exact four-type combination — not applied to half, not to float, not to the int32 path.
  2. The bound is derived, not chosen. eps × k × 128 × 128 follows from the actual failure mode: an int8 product is bounded by 128², k of them are summed, and a float accumulator rounds at the magnitude of the running sum rather than of the final entry. Your comment explaining that an entry whose sum cancels can't be covered by any relative bound is exactly right, and it's the real reason a relative-only check is wrong here.
  3. You wrote a test that validates the model instead of just consuming it. Int8Int8SinglePrecisionErrorModel constructs deliberate catastrophic cancellation — A entries 5g, 3g, 2g against B entries h, −h, −h, summing to zero — computes the expected result exactly in integers so it doesn't depend on the reference BLAS, and then reports how much of the allowed tolerance was actually consumed. That last part is what separates this from tolerance-padding: if a future change starts eating the whole budget, the output says so. The note that the shapes aren't interchangeable because some backends will accumulate them exactly shows you thought about the case where the test can't fail.

Good work. I'd point other contributors at this as the way to introduce a mixed-precision tolerance.

One suggestion that would make the check stronger

For this specific combination the exact result is always an integer — it's a sum of products of integers. So whenever the exact value is representable, the tightest possible check isn't the accumulation model at all, it's:

|x - x_ref| < 0.5

i.e. the float output must round to the correct integer. For moderate k that's dramatically tighter than eps × k × 128 × 128. Concretely at k = 256 the model allows roughly 1.19e-7 × 256 × 16384 ≈ 0.5, so the two coincide there — but below that the model is looser than it needs to be, and a kernel that was wrong by a whole unit at small k would still pass.

The caveat is the ceiling: Σ|a·b| can reach k × 16384, which crosses 2²⁴ around k ≈ 1024, and past that the exact integer isn't representable in float and the accumulation model is the only correct bound. So the strengthening would be to take min of the two — round-to-nearest-integer while the result fits in 2²⁴, falling back to the accumulation bound above it. Worth doing if it's cheap; not a blocker, since what you have is sound and conservative in the right direction.

Minor

abs_bound is a double cast to fp_real. Harmless for the only caller that passes it (where fp_real is float), but if anyone ever passes a large bound with a half-precision type that cast overflows to infinity and silently disables the absolute check. A comment or a static_assert on the type would prevent that being discovered the hard way.

Note on the stack

#763 builds on this — it carries this same cublas_batch.cpp change and the same test-harness edits, plus the rocBLAS side. Once this lands, please rebase #763 down to just the rocBLAS delta so it isn't re-reviewed from scratch.

@ndingle-arm ndingle-arm left a comment

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Approved. The implementation change is nicely focused, the commit messages clearly explain both the supported and unsupported combinations, and the tolerance is narrowly scoped with a thoughtful regression test. I also like that cancelling_missing_relative_bound remains diagnostic: a backend that satisfies the relative bound—or accumulates exactly—should still pass.

Overall assessment

  • The three cuBLAS entry points are routed correctly. NVIDIA documents int8 A/B with float C and float compute/scalars as supported for both batched APIs.
  • The int32-output combination correctly remains unsupported because cuBLAS requires int32 scalars for CUBLAS_COMPUTE_32I.
  • The shared checker’s new argument defaults to zero, preserving existing callers.
  • Changes are focused: five files, with most additions comprising the deterministic regression test.
  • git diff --check passes.
  • All 11 GitHub checks pass; the PR supplies separate A100 results for the affected backend.

@ndingle-arm
ndingle-arm merged commit 53e36a7 into uxlfoundation:develop Aug 27, 2026
11 checks passed
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

3 participants