[graph_trainer][precompile] Fix CooR precompile FlexAttention input mismatch (#4014)