From a5c6313e5f099cde9f06f688dbc6757eae6d9da3 Mon Sep 17 00:00:00 2001 From: Anthony Shoumikhin Date: Tue, 25 Aug 2026 19:48:33 -0700 Subject: [PATCH] fix(tests): compare baddbmm lowering against TensorRT in fp32, not TF32 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 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. On the standard runtime, disable_tf32 makes both sides run fp32 and the strict DECIMALS_OF_AGREEMENT check holds. TensorRT-RTX ignores the TF32 builder flag and always runs the matmul in TF32, so the fp32 reference can still differ by one TF32 rounding step. Keep the strict check on the standard runtime and relax it to one decimal (0.05, comfortably above the 0.0054 TF32 drift) on RTX, which still catches any real breakage of the decomposition. --- tests/py/dynamo/lowering/test_decompositions.py | 12 +++++++++++- 1 file changed, 11 insertions(+), 1 deletion(-) diff --git a/tests/py/dynamo/lowering/test_decompositions.py b/tests/py/dynamo/lowering/test_decompositions.py index 5beea8c1a8..89aacfa262 100644 --- a/tests/py/dynamo/lowering/test_decompositions.py +++ b/tests/py/dynamo/lowering/test_decompositions.py @@ -2507,16 +2507,26 @@ def forward(self, bias, batch1, batch2): inputs, min_block_size=1, pass_through_build_failures=True, + # The reference below runs the matmul in full fp32, so TensorRT has to + # as well. TF32 is on by default and rounds the operands to 10 mantissa + # bits, which costs far more than DECIMALS_OF_AGREEMENT allows. + disable_tf32=True, ) with torch.no_grad(): trt_results = optimized_model(*inputs).detach().cpu() torch_results = fx_graph(*inputs).detach().cpu() max_diff = float(torch.max(torch.abs(trt_results - torch_results))) + # TensorRT-RTX ignores disable_tf32 and always runs the matmul in TF32, so + # the fp32 reference can differ by one TF32 rounding step. Only the standard + # runtime can meet DECIMALS_OF_AGREEMENT here. + decimals_of_agreement = ( + 1 if torch_tensorrt.ENABLED_FEATURES.tensorrt_rtx else DECIMALS_OF_AGREEMENT + ) self.assertAlmostEqual( max_diff, 0, - DECIMALS_OF_AGREEMENT, + decimals_of_agreement, f"baddbmm TRT outputs don't match with the original model. (diff={max_diff})", )