[blas][cublas] Support int8 inputs with float output in gemm_batch - #761
Conversation
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>
…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
left a comment
There was a problem hiding this comment.
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:
- It's opt-in and narrowly scoped.
abs_bounddefaults to0.0andabs_limit = max(bound, abs_bound), so every existing caller's behavior is bit-for-bit unchanged. The wider bound is gated behindif constexpr (int8_to_float)on the exact four-type combination — not applied to half, not to float, not to the int32 path. - The bound is derived, not chosen.
eps × k × 128 × 128follows 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. - You wrote a test that validates the model instead of just consuming it.
Int8Int8SinglePrecisionErrorModelconstructs deliberate catastrophic cancellation —Aentries5g, 3g, 2gagainstBentriesh, −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
left a comment
There was a problem hiding this comment.
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.
Description
Enables the
(int8, int8, float, float)gemm_batchcombination 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.
cublasGemmStridedBatchedExandcublasGemmBatchedExalready accept the datatypes that the existing launchers forward, so this combination only needed to be routed togemm_batch_implinstead of to the macro that throwsunimplemented.(int8, int8, int32, float)remains unimplemented, and the reason is now recorded in a comment next to it: cuBLAS produces an int32 output only underCUBLAS_COMPUTE_32I, which requires int32alphaandbeta, 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_batchtests 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 exactint64reference, 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_equalacceptsabs_bound, defaulted to0.0and threaded throughcheck_equal_matrixandcheck_almost_equal_matrix. With the default,abs_boundcannot 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.gemm_batchtests computeeps * |alpha| * k * 128 * 128underif constexpr, so only the int8-to-float instantiation is affected. Since every int8 magnitude is below 128,k * 128 * 128is an upper bound onsum|a*b|. The relative bound stays at10 * 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
Int8Int8SinglePrecisionErrorModelcovers this path deterministically. Test sizes and data come from an in-test generator rather thanstd::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 are5x,-3xand-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
Built with the DPC++ compiler and the cuBLAS backend, run on an NVIDIA A100-SXM4-40GB.
gemm_batchtests, compile-time and run-time dispatch:The skips are the row-major and other unimplemented cuBLAS combinations, unchanged by this PR.
Wider BLAS suite, both dispatch modes:
Complex double level 1 tests are excluded because they crash identically on this setup with and without these changes, in an untouched build of
developas 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
gemm_batchtests already cover this type combination and now run against the cuBLAS backend rather than skipping;Int8Int8SinglePrecisionErrorModelwas added for the tolerance itself.Bug fixes
Made with Cursor