diff --git a/tensorrt_llm/_torch/auto_deploy/transform/library/mlir_elementwise_fusion.py b/tensorrt_llm/_torch/auto_deploy/transform/library/mlir_elementwise_fusion.py index c16f0b30b803..cd9ce643f1df 100644 --- a/tensorrt_llm/_torch/auto_deploy/transform/library/mlir_elementwise_fusion.py +++ b/tensorrt_llm/_torch/auto_deploy/transform/library/mlir_elementwise_fusion.py @@ -166,15 +166,15 @@ def _min_output_rank(sg): back_converter = MLIRToFXConverter(gm) new_gm = back_converter.convert(mlir_module, converter.metadata) - # Step 6: Re-validate the graph and propagate shapes. - from ...utils._graph import canonicalize_graph, run_shape_prop - - canonicalize_graph(new_gm) - run_shape_prop(new_gm) - + # Step 6: Defer graph cleanup and shape-prop to the framework. Running + # fake-tensor shape-prop here is unsafe at post_load_fusion: the graph + # may be in a deliberately invalid intermediate state (e.g. + # fuse_rope_into_trtllm_attention rewires Q/K/V to a fused-QKV tensor + # whose rank does not match torch_attention.register_fake; the op swap + # happens later at cache_init). return new_gm, TransformInfo( skipped=False, num_matches=num_replaced, - is_clean=True, - has_valid_shapes=True, + is_clean=False, + has_valid_shapes=False, ) diff --git a/tests/integration/test_lists/waives.txt b/tests/integration/test_lists/waives.txt index 8ec46b10e0c9..113550cc441c 100644 --- a/tests/integration/test_lists/waives.txt +++ b/tests/integration/test_lists/waives.txt @@ -299,7 +299,6 @@ perf/test_perf_sanity.py::test_e2e[aggr_upload-k25_thinking_fp4_2_nodes_grace_bl perf/test_perf_sanity.py::test_e2e[aggr_upload-k25_thinking_fp4_blackwell-k25_thinking_fp4_dep8_32k8k] SKIP (https://nvbugs/6236094) perf/test_perf_sanity.py::test_e2e[aggr_upload-k25_thinking_fp4_blackwell-k25_thinking_fp4_dep8_8k1k] SKIP (https://nvbugs/6227472) perf/test_perf_sanity.py::test_e2e[aggr_upload-k25_thinking_fp4_blackwell-k25_thinking_fp4_tep8_32k8k] SKIP (https://nvbugs/6227472) -perf/test_perf_sanity.py::test_e2e[aggr_upload-llama3_1_8b_fp8_ad_hopper-llama3_1_8b_ad_ws1_1k1k] SKIP (https://nvbugs/6244474) perf/test_perf_sanity.py::test_e2e[aggr_upload-super_ad_blackwell-super_ad_ws1_1k1k] SKIP (https://nvbugs/6153575) perf/test_perf_sanity.py::test_e2e[disagg_upload-e2e-gb200_deepseek-r1-fp4_128k8k_con128_ctx1_pp8_gen1_dep16_eplb0_mtp1_ccb-NIXL] SKIP (https://nvbugs/6215844) perf/test_perf_sanity.py::test_e2e[disagg_upload-e2e-gb200_kimi-k25-thinking-fp4_8k1k_con4096_ctx1_dep4_gen1_dep16_eplb0_mtp0_ccb-NIXL] SKIP (https://nvbugs/6179661)