set per_block_mean=True and modify api

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

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)