From 75411ce0cc26d595c22df457030b9a2e0581dac2 Mon Sep 17 00:00:00 2001 From: ViperEkura <3081035982@qq.com> Date: Fri, 31 Jul 2026 15:40:38 +0800 Subject: [PATCH] fix: skip CUDA rotary kernel when grad is enabled MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit - apply_rotary_emb now checks torch.is_grad_enabled() before dispatching to CUDA kernel - Training (grad enabled) uses torch complex multiply path which supports autograd backward - Inference (inference_mode/no_grad) uses CUDA kernel as before - Without this fix, training backward would crash — the CUDA kernel has no autograd backward() --- astrai/extension/rotary_backend.py | 7 ++++++- 1 file changed, 6 insertions(+), 1 deletion(-) diff --git a/astrai/extension/rotary_backend.py b/astrai/extension/rotary_backend.py index 3803a32..e5de677 100644 --- a/astrai/extension/rotary_backend.py +++ b/astrai/extension/rotary_backend.py @@ -43,6 +43,11 @@ def apply_rotary_emb(x: Tensor, rotary_emb: tuple[Tensor, Tensor]) -> Tensor: [batch, seq_len, n_heads, head_dim] (bf16) """ cos, sin = rotary_emb - if _cuda_available() and x.is_cuda and x.dtype == torch.bfloat16: + if ( + _cuda_available() + and not torch.is_grad_enabled() + and x.is_cuda + and x.dtype == torch.bfloat16 + ): return _cuda_rotary(x, cos, sin) return _torch_apply(x, cos, sin)