diff --git a/fastvideo/attention/backends/sage_attn3.py b/fastvideo/attention/backends/sage_attn3.py index 5d7b5b799..580c800af 100644 --- a/fastvideo/attention/backends/sage_attn3.py +++ b/fastvideo/attention/backends/sage_attn3.py @@ -418,7 +418,7 @@ class _SageAttnBlackwellWithTritonBwd(torch.autograd.Function): return dq_BHLD, dk_BHLD, dv_BHLD, None, None, None, None -def sageattn_blackwell_with_triton_bwd(q_BLHD, k_BLHD, v_BLHD, is_causal=False, per_block_mean=False): +def sageattn_blackwell_with_triton_bwd(q_BLHD, k_BLHD, v_BLHD, is_causal=False, per_block_mean=True): softmax_scale = 1.0 / sqrt(q_BLHD.shape[-1]) # FA forward only to get softmax_lse (which is shape BHL) # Use no_grad to prevent autograd from tracking this call diff --git a/fastvideo/attention/backends/sageattn/api.py b/fastvideo/attention/backends/sageattn/api.py index def4f50fe..fc1a5854b 100644 --- a/fastvideo/attention/backends/sageattn/api.py +++ b/fastvideo/attention/backends/sageattn/api.py @@ -92,7 +92,7 @@ def preprocess_qkv(q: torch.Tensor, k: torch.Tensor, v: torch.Tensor, per_block_ q, qm = triton_group_mean(q) else: qm = q.mean(dim=-2, keepdim=True) - q = q - qm + q = q - qm delta_s = torch.matmul(qm, k.transpose(-2, -1)).to(torch.float32).contiguous() if disable_delta_s: # used to disable q smoothing delta_s = torch.zeros_like(delta_s)