[FA4] Compilation too long for backward on TransformerEngine 2.15

Hi, when I use Megatron, TransformerEngine 2.15.0 with support for FA4, I encounter significant time for backward. Later, I find that it is due to the re-compiling at every layer for every micro-batch. My environment, - Megatron 0.18.0 - TransformerEngine 2.15.0 - fa4 4.0.0.b13 - Training Setting: Varying length, packing=True - Model: Qwen2.5-7B While the ordinary path that I used to run is like, ``` Thread 115408 (active) __call__ (torch/_ops.py:1269) backward (flash_attn/flash_attn_interface.py:968) apply (torch/autograd/function.py:317) _engine_run_backward (torch/autograd/graph.py:869) backward (torch/autograd/__init__.py:381) backward (core/tensor_parallel/random.py:630) apply (torch/autograd/function.py:317) ``` Now it runs like ``` Thread 57902 (active+gil) compile (cutlass/base_dsl/compiler.py:148) compile_and_jit (cutlass/base_dsl/compiler.py:186) compile_and_jit (cutlass/base_dsl/dsl.py:976) compile_and_cache (cutlass/base_dsl/dsl.py:1229) compile_and_cache (cutlass/cutlass_dsl/cutlass.py:541) generate_mlir (cutlass/base_dsl/dsl.py:1382) _func (cutlass/base_dsl/dsl.py:1567) _compile (cutlass/base_dsl/compiler.py:641) __call__ (cutlass/base_dsl/compiler.py:562) _flash_attn_bwd (flash_attn/cute/interface.py:1812) backward (flash_attn/cute/interface.py:2100) apply (torch/autograd/function.py:317) _engine_run_backward (torch/autograd/graph.py:869) backward (torch/autograd/__init__.py:381) backward (core/tensor_parallel/random.py:630) apply (torch/autograd/function.py:317) ``` and it takes forever to accomplish the backward. Am I missing any switches that should help accelerate the JiT process or simply avoid JiT? Thank you!