From 7b85b29df53a64117923b6cdbf15da5fd7dbd2e8 Mon Sep 17 00:00:00 2001 From: Zhiyao Jiang Date: Thu, 6 Aug 2026 19:48:59 +0000 Subject: [PATCH] [ROCm] Compile CP softmax LSE corrections with dynamic shapes --- .../dot_product_attention/context_parallel.py | 6 +++--- transformer_engine/pytorch/jit.py | 21 +++++++++++++++++++ 2 files changed, 24 insertions(+), 3 deletions(-) diff --git a/transformer_engine/pytorch/attention/dot_product_attention/context_parallel.py b/transformer_engine/pytorch/attention/dot_product_attention/context_parallel.py index cca013461d..0b3dd1a0e1 100644 --- a/transformer_engine/pytorch/attention/dot_product_attention/context_parallel.py +++ b/transformer_engine/pytorch/attention/dot_product_attention/context_parallel.py @@ -27,7 +27,7 @@ from transformer_engine.pytorch.tensor.float8_tensor import Float8Tensor from transformer_engine.pytorch.tensor.storage.float8_tensor_storage import Float8TensorStorage from transformer_engine.pytorch.quantized_tensor import QuantizedTensorStorage -from transformer_engine.pytorch.jit import jit_fuser +from transformer_engine.pytorch.jit import jit_fuser, jit_fuser_dynamic from transformer_engine.pytorch.graph import is_graph_capturing from transformer_engine.pytorch.constants import dist_group_type from transformer_engine.pytorch.distributed import ( @@ -184,7 +184,7 @@ def flash_attn_fwd_second_half_out_correction( out_.add_(out_corrected) -@jit_fuser +@jit_fuser_dynamic def flash_attn_fwd_softmax_lse_correction( softmax_lse: torch.Tensor, softmax_lse_per_step: torch.Tensor, @@ -196,7 +196,7 @@ def flash_attn_fwd_softmax_lse_correction( softmax_lse.copy_(new_scale) -@jit_fuser +@jit_fuser_dynamic def flash_attn_fwd_second_half_softmax_lse_correction( softmax_lse: torch.Tensor, softmax_lse_per_step: torch.Tensor, diff --git a/transformer_engine/pytorch/jit.py b/transformer_engine/pytorch/jit.py index 4d52d9b925..25089d3bbd 100644 --- a/transformer_engine/pytorch/jit.py +++ b/transformer_engine/pytorch/jit.py @@ -36,9 +36,30 @@ def wrapper(*args, **kwargs): return wrapper +# See: https://github.com/ROCm/TransformerEngine/issues/693 +def lazy_compile_dynamic(func): + """Lazy compile a function with torch.compile, using dynamic shapes + + Compiling dynamic up front skips the static specialization dynamo would otherwise build on + the first shape it sees, leaving one compiled artifact for the process instead of two. + """ + compiled_func = None + + @wraps(func) + def wrapper(*args, **kwargs): + nonlocal compiled_func + if compiled_func is None: + compiled_func = torch.compile(func, dynamic=True) + return compiled_func(*args, **kwargs) + + return wrapper + + jit_fuser = lambda func: func +jit_fuser_dynamic = lambda func: func if torch_version() >= (2, 0, 0) and bool(int(os.getenv("NVTE_TORCH_COMPILE", "1"))): jit_fuser = lazy_compile + jit_fuser_dynamic = lazy_compile_dynamic # See: https://github.com/NVIDIA/TransformerEngine/issues/597