[Diffusion] Fix Ring Parallel bug with FA4 (#18062)

2026-02-02 17:06:51 +08:00
parent 750ad0d290
commit a0757c9624
1 changed files with 1 additions and 0 deletions
@@ -1364,6 +1364,7 @@ class FlashAttnVarlenFunc(torch.autograd.Function):
            pack_gqa=pack_gqa,
            score_mod=score_mod,
            aux_tensors=aux_tensors,
            return_lse=True,
        )
        ctx.save_for_backward(q, k, v, out, lse, cu_seqlens_q, cu_seqlens_k, seqused_q, seqused_k)
        ctx.softmax_scale = softmax_scale