fix(tests): compare baddbmm lowering against TensorRT in fp32, not TF32 - #4586
Open
shoumikhin wants to merge 2 commits into
Open
fix(tests): compare baddbmm lowering against TensorRT in fp32, not TF32#4586shoumikhin wants to merge 2 commits into
shoumikhin wants to merge 2 commits into
Conversation
test_lowering_baddbmm builds the engine with TF32 left at its default (on), then checks the result against an eager fp32 reference to DECIMALS_OF_AGREEMENT. TF32 rounds the matmul operands to 10 mantissa bits, so a passing run only means TensorRT happened to pick a non-TF32 kernel. TensorRT-RTX picks a TF32 one and the test fails on main with a difference of 0.0019378662109375. The decomposition itself is exact: over 500 random draws of the shapes the test uses, bias + bmm(batch1, batch2) in fp32 matches aten.baddbmm bit for bit, while the same math with the operands rounded to TF32 drifts up to 0.0054. Ask TensorRT for fp32 so both sides run the same arithmetic. The tolerance and the coverage stay as they are.
shoumikhin
force-pushed
the
fix/baddbmm-tf32
branch
from
August 26, 2026 02:51
b6c603a to
f72e37a
Compare
The repository lint job runs `black --check .` across the whole tree, so any file that does not match the formatter fails CI for every open pull request, not only the one that touched it. `tests/py/dynamo/conversion/test_cumsum_aten.py` is currently not black-conformant on main, which turns the Python Linting check red here. Reformat that one file with black. This is a formatting-only change: two statements that fit on a single line are un-wrapped. No test logic changes. Verified by running `black --check .` on the full tree: all files pass.
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Problem
dynamo-lowering-rtxis red on main:The same test passes on standard TensorRT, so it went in green and only turned red once the RTX lane ran on main.
Cause
The test compares a TensorRT engine against an eager fp32 reference to
DECIMALS_OF_AGREEMENT, but builds the engine with TF32 left at its default, which is on. TF32 rounds both matmul operands to 10 mantissa bits before multiplying. A passing run therefore only means TensorRT happened to pick a kernel that was not TF32; it is not something the test asks for. TensorRT-RTX picks a TF32 kernel and the check fails.The decomposition is not at fault. Over 500 random draws of the shapes the test uses,
bias + bmm(batch1, batch2)in fp32 matchesaten.baddbmmbit for bit (max difference 0.0). Repeating the same math with the two operands rounded to 10 mantissa bits drifts by up to 0.0054, which brackets the 0.0019 seen in CI.Fix
Pass
disable_tf32=Trueso both sides run the same arithmetic. The tolerance is unchanged and nothing is skipped. This matches how other numeric tests in the repo handle it, for exampletests/py/dynamo/conversion/test_index_put_aten.pyandtests/py/dynamo/models/test_models.py.Testing
test_lowering_baddbmmpasses locally on standard TensorRT with the change, and the debug log confirmsdisable_tf32=Truereaches the compilation settings through thetorch_compilepath. The RTX lane needsci: fullor an approving review to run here.