diff --git a/examples/inference/optimizations/quantization_example.py b/examples/inference/optimizations/quantization_example.py index 171c3f6ce..d5518a489 100644 --- a/examples/inference/optimizations/quantization_example.py +++ b/examples/inference/optimizations/quantization_example.py @@ -1,9 +1,6 @@ import torch from fastvideo import VideoGenerator from fastvideo.configs.pipelines.base import PipelineConfig -from fastvideo.layers.quantization.fp4_config import FP4Config -# from fastvideo.layers.quantization.fp8_config import FP8Config -from fastvideo.layers.linear import ReplicatedLinear OUTPUT_PATH = "video_samples" @@ -25,19 +22,16 @@ def main(): model_id = "Wan-AI/Wan2.1-T2V-1.3B-Diffusers" # model_id = "Wan-AI/Wan2.1-T2V-14B-Diffusers" pipeline_config = PipelineConfig.from_pretrained(model_id) - pipeline_config.dit_precision = "bf16" - pipeline_config.dit_config.quant_config = FP4Config() - # pipeline_config.dit_config.quant_config = FP8Config() + pipeline_config.dit_precision = "bf16" - # ReplicatedLinear.print_shape_summary() - print("\nLoading model with FP4 quantization...") generator = VideoGenerator.from_pretrained( model_id, pipeline_config=pipeline_config, num_gpus=1, use_fsdp_inference=True, - dit_cpu_offload=False, + transformer_quant="fp4", + dit_cpu_offload=False, vae_cpu_offload=False, text_encoder_cpu_offload=True, pin_cpu_memory=False, @@ -91,4 +85,4 @@ def main(): if __name__ == "__main__": - main() \ No newline at end of file + main() diff --git a/examples/inference/optimizations/sageattn_inference_example_custom_weights_5090_32.py b/examples/inference/optimizations/sageattn_inference_example_custom_weights_5090_32.py deleted file mode 100644 index f69ee031a..000000000 --- a/examples/inference/optimizations/sageattn_inference_example_custom_weights_5090_32.py +++ /dev/null @@ -1,54 +0,0 @@ -from fastvideo import VideoGenerator -import os -from pathlib import Path -# from fastvideo.configs.sample import SamplingParam - -OUTPUT_PATH = "video_samples" -os.environ["FASTVIDEO_ATTENTION_BACKEND"] = "SAGE_ATTN_THREE" - -CHECKPOINT_PATH = Path(__file__).parent.parent.parent - -def main(): - # FastVideo will automatically use the optimal default arguments for the - # model. - # If a local path is provided, FastVideo will make a best effort - # attempt to identify the optimal arguments. - generator = VideoGenerator.from_pretrained( - "Wan-AI/Wan2.1-T2V-1.3B-Diffusers", - # FastVideo will automatically handle distributed setup - num_gpus=1, - use_fsdp_inference=True, - dit_cpu_offload=False, - vae_cpu_offload=False, - text_encoder_cpu_offload=True, - pin_cpu_memory=False, # set to false if low CPU RAM or hit obscure "CUDA error: Invalid argument" - # image_encoder_cpu_offload=False, - # Load custom weights from checkpoint-5000 - init_weights_from_safetensors="/workspace/FastVideo-Quantization/M2_32" - ) - - # sampling_param = SamplingParam.from_pretrained("Wan-AI/Wan2.1-T2V-1.3B-Diffusers") - # sampling_param.num_frames = 45 - # sampling_param.image_path = "https://huggingface.co/datasets/huggingface/documentation-images/resolve/main/diffusers/astronaut.jpg" - # Generate videos with the same simple API, regardless of GPU count - prompt = ( - "A curious raccoon peers through a vibrant field of yellow sunflowers, its eyes " - "wide with interest. The playful yet serene atmosphere is complemented by soft " - "natural light filtering through the petals. Mid-shot, warm and cheerful tones." - ) - video = generator.generate_video(prompt, output_path=OUTPUT_PATH, save_video=True) - # video = generator.generate_video(prompt, sampling_param=sampling_param, output_path="wan_t2v_videos/") - - # Generate another video with a different prompt, without reloading the - # model! - prompt2 = ( - "A majestic lion strides across the golden savanna, its powerful frame " - "glistening under the warm afternoon sun. The tall grass ripples gently in " - "the breeze, enhancing the lion's commanding presence. The tone is vibrant, " - "embodying the raw energy of the wild. Low angle, steady tracking shot, " - "cinematic.") - video2 = generator.generate_video(prompt2, output_path=OUTPUT_PATH, save_video=True) - - -if __name__ == "__main__": - main() diff --git a/examples/inference/sageattn+fp4linear.py b/examples/inference/sageattn+fp4linear.py index 1964c4763..6074ef2fc 100644 --- a/examples/inference/sageattn+fp4linear.py +++ b/examples/inference/sageattn+fp4linear.py @@ -3,7 +3,7 @@ import os # from fastvideo.configs.sample import SamplingParam OUTPUT_PATH = "video_samples" -os.environ["FASTVIDEO_ATTENTION_BACKEND"] = "SAGE_ATTN_THREE" +os.environ["FASTVIDEO_ATTENTION_BACKEND"] = "MODIFIED_SAGE_ATTN_THREE" def main(): # FastVideo will automatically use the optimal default arguments for the diff --git a/fastvideo/attention/backends/modified_sage_attn3.py b/fastvideo/attention/backends/modified_sage_attn3.py new file mode 100644 index 000000000..db251f5e6 --- /dev/null +++ b/fastvideo/attention/backends/modified_sage_attn3.py @@ -0,0 +1,495 @@ +# SPDX-License-Identifier: Apache-2.0 + +import os +import sys +from pathlib import Path + +# Add local repo roots to sys.path for kernel-side imports during development. +_project_root = Path(__file__).resolve().parent.parent.parent.parent +_kernel_root = _project_root / "fastvideo-kernel" +_kernel_python_root = _kernel_root / "python" +for _path in (_project_root, _kernel_root, _kernel_python_root): + if str(_path) not in sys.path: + sys.path.insert(0, str(_path)) + +import torch +import triton +try: + # Prefer the in-repo kernel implementation during local development. + from modified_sageattn import sageattn_blackwell +except ImportError: + # Fall back to installed package if local version not available + try: + from sageattn import sageattn_blackwell + except ImportError: + sageattn_blackwell = None + +from fastvideo.attention.backends.abstract import (AttentionBackend, AttentionImpl, AttentionMetadata, + AttentionMetadataBuilder) +from fastvideo.logger import init_logger +from fastvideo_kernel.triton_kernels.qat_attn import attention +from fastvideo_kernel.triton_kernels.qat_attn import ( + _attn_bwd, + _attn_bwd_dkdv_cross, + _attn_bwd_dq_cross, + _attn_bwd_preprocess, +) +logger = init_logger(__name__) + +from math import sqrt +try: + from flash_attn.flash_attn_interface import _wrapped_flash_attn_forward, _wrapped_flash_attn_backward +except ImportError: + _wrapped_flash_attn_forward = None + _wrapped_flash_attn_backward = None + +from fastvideo_kernel.triton_kernels.quant_utils import ( + fake_quantize_kv, + fake_quantize_q, +) + +from flashinfer import SfLayout, mm_fp4, nvfp4_quantize + +def matmul_3d_4bit(a, b): + H, M, K1 = a.shape + H, N, K2 = b.shape + assert K1 == K2 + padded_K = (K2 + 127) // 128 * 128 + a = torch.nn.functional.pad(a, (0, padded_K - K2)) + b = torch.nn.functional.pad(b, (0, padded_K - K2)) + out = torch.empty((H, M, N), device=a.device, dtype=a.dtype) + one = torch.ones(1, device=a.device, dtype=torch.float32) + for h in range(H): + q_fp4, q_inv_s = nvfp4_quantize(a[h], one, sfLayout=SfLayout.layout_128x4, do_shuffle=False) + k_fp4, k_inv_s = nvfp4_quantize(b[h], one, sfLayout=SfLayout.layout_128x4, do_shuffle=False) + mm_fp4( + q_fp4, k_fp4.T, q_inv_s, k_inv_s.T, one, + a.dtype, out[h], + block_size=16, + use_8x4_sf_layout=False, + backend="cutlass", + ) + return out + + +class _Matmul3d4bitFWD16bitBWD(torch.autograd.Function): + @staticmethod + def forward(ctx, a, b, scale=None): + ctx.save_for_backward(a, b) + if scale is not None: + a = a / scale + return matmul_3d_4bit(a, b) * scale + else: + return matmul_3d_4bit(a, b) + + @staticmethod + def backward(ctx, grad_out): + a, b = ctx.saved_tensors + grad_a = grad_out.matmul(b) + grad_b = grad_out.transpose(1, 2).matmul(a) + # None for the three extra forward args + return grad_a, grad_b, None + + +class _Matmul3d4bitFWD16bitBWD_with_smoothing(torch.autograd.Function): + @staticmethod + def forward(ctx, q, k): + ctx.save_for_backward(q, k) + k = k - k.mean(dim=-2, keepdim=True) # [H, L, D] + qm = q.mean(dim=-2, keepdim=True) # [H, 1, D] + q = q - qm # [B, H, L, D] + delta_s = torch.matmul(qm, k.transpose(-2, -1)).contiguous() # [H, 1, L] + return matmul_3d_4bit(q, k) + delta_s + + @staticmethod + def backward(ctx, grad_out): + a, b = ctx.saved_tensors + grad_a = grad_out.matmul(b) + grad_b = grad_out.transpose(1, 2).matmul(a) + # None for the three extra forward args + return grad_a, grad_b, None + +USE_SMOOTHING = False +def sage_attn_qk_torch(q, k, per_block_mean=False): + """ + Inputs: q/k/v in [B, H, L, D] + Returns: out in [B, H, L, D] + """ + B, H, qL, D = q.shape + B, H, kL, D = k.shape + assert B == 1 + q, k = q.squeeze(0), k.squeeze(0) + if USE_SMOOTHING: + qk = _Matmul3d4bitFWD16bitBWD_with_smoothing.apply(q, k) + else: + qk = _Matmul3d4bitFWD16bitBWD.apply(q, k) + qk = qk / torch.sqrt(torch.tensor(D, device=q.device, dtype=q.dtype)) + p = torch.softmax(qk, dim=-1) + return p.unsqueeze(0) + +def attn_backward_16bit(q_BLHD, k_BLHD, v_BLHD, out_BLHD, grad_out_BLHD, is_causal=False): + assert is_causal == False + q, k, v, o, do = q_BLHD.transpose(1 ,2).contiguous(), k_BLHD.transpose(1 ,2).contiguous(), v_BLHD.transpose(1 ,2).contiguous(), out_BLHD.transpose(1 ,2).contiguous(), grad_out_BLHD.transpose(1 ,2).contiguous() + B, H, _, D = q.shape + p = sage_attn_qk_torch(q, k, per_block_mean=False).to(do.dtype) + dV = p.transpose(-2, -1) @ do # [B,H,L,D] + + dP = do @ v.transpose(-2, -1) # [B,H,L,L] + + # do_dot_o = (do.float() * o.float()).sum(dim=-1, keepdim=True).to(dP.dtype) # [B,H,L,1] + # do_dot_o = (do * o).sum(dim=-1, keepdim=True) # [B,H,L,1] + dp_dot_p = (dP * p).sum(dim=-1, keepdim=True) # [B,H,L,1] + dS = p * (dP - dp_dot_p) # [B,H,L,L] + + + dQ = (dS @ k) / torch.sqrt(torch.tensor(D, device=q.device, dtype=q.dtype)) # [B,H,L,D] + dK = (dS.transpose(-2, -1) @ q) / torch.sqrt(torch.tensor(D, device=q.device, dtype=q.dtype)) # [B,H,L,D] + + # grads w.r.t inputs, None for non-tensor inputs + return dQ.transpose(1, 2).contiguous(), dK.transpose(1, 2).contiguous(), dV.transpose(1, 2).contiguous(), None, None, None + +def attn_forward_4bit(q_BHLD, k_BHLD, v_BHLD, is_causal=False): + assert is_causal == False + p_bf16 = sage_attn_qk_torch(q_BHLD, k_BHLD, per_block_mean=False).to(q_BHLD.dtype) + row_max_p = p_bf16.max(dim=-1, keepdim=True).values + scale_row = row_max_p / (448 * 6) + p = p_bf16 / scale_row + out = matmul_3d_4bit(p.squeeze(0), v_BHLD.squeeze(0).transpose(1, 2).contiguous()).unsqueeze(0) + out = out * scale_row + return out, p_bf16 + + +def flash_attn_backward(q_BLHD, k_BLHD, v_BLHD, out_BLHD, grad_out_BLHD, is_causal=False): + D = q_BLHD.shape[-1] + softmax_scale = 1.0 / sqrt(D) + + # FA forward only to get softmax_lse (rng_state is also returned) + _, softmax_lse, _S_dmask, rng_state = _wrapped_flash_attn_forward( + q_BLHD, k_BLHD, v_BLHD, + 0.0, # dropout_p + softmax_scale, # softmax_scale + is_causal, # causal + -1, -1, # window_size_left/right + 0.0, # softcap + None, # alibi_slopes + False # return_softmax + ) + + # Allocate grads and call FA backward using the SAGE output + dq_BLHD = torch.empty_like(q_BLHD) + dk_BLHD = torch.empty_like(k_BLHD) + dv_BLHD = torch.empty_like(v_BLHD) + + _wrapped_flash_attn_backward( + grad_out_BLHD, # dout (BLHD, 16-bit) + q_BLHD, k_BLHD, v_BLHD, + out_BLHD, # use SAGE forward output here + softmax_lse, + dq_BLHD, dk_BLHD, dv_BLHD, + 0.0, # dropout_p + softmax_scale, # softmax_scale + is_causal, # causal + -1, -1, # window_size_left/right + 0.0, # softcap + None, # alibi_slopes + False, # deterministic + rng_state=rng_state, + ) + return dq_BLHD, dk_BLHD, dv_BLHD, None, None + +USE_TORCH_4BIT_FWD = False +USE_TORCH_16BIT_BWD = False +USE_QAT_ATTN = False +class _SageAttnBlackwellWith16bitBwd(torch.autograd.Function): + @staticmethod + def forward(ctx, q_BLHD, k_BLHD, v_BLHD, is_causal=False, per_block_mean=True): + """ + Inputs: q/k/v in [B, L, H, D] + Returns: out in [B, L, H, D] + """ + # Save originals for backward (we'll recompute FA in 16-bit there) + ctx.is_causal = bool(is_causal) + + # Convert to BHLD for your sageattn forward + q_BHLD = q_BLHD.permute(0, 2, 1, 3).contiguous() + k_BHLD = k_BLHD.permute(0, 2, 1, 3).contiguous() + v_BHLD = v_BLHD.permute(0, 2, 1, 3).contiguous() + + with torch.no_grad(): + out_BHLD = sageattn_blackwell( + q_BHLD, k_BHLD, v_BHLD, + attn_mask=None, + is_causal=is_causal, + per_block_mean=per_block_mean, + ) + # Convert to BLHD for saving + out_BLHD = out_BHLD.permute(0, 2, 1, 3).contiguous() + ctx.save_for_backward(q_BLHD, k_BLHD, v_BLHD, out_BLHD) + + return out_BLHD + + @staticmethod + def backward(ctx, grad_out_BLHD): + is_causal = ctx.is_causal + q_BLHD, k_BLHD, v_BLHD, out_BLHD = ctx.saved_tensors + return flash_attn_backward(q_BLHD, k_BLHD, v_BLHD, out_BLHD, grad_out_BLHD, is_causal) + + + +def sageattn_blackwell_with_16bit_bwd(q_BLHD, k_BLHD, v_BLHD, is_causal=False, per_block_mean=True): + """ + Forward: uses sageattn_blackwell under the hood. + Backward: recomputes FlashAttention fwd+bwd in 16bit directly. + """ + return _SageAttnBlackwellWith16bitBwd.apply(q_BLHD, k_BLHD, v_BLHD, is_causal, per_block_mean) + +class _SageAttnBlackwellWithTritonBwd(torch.autograd.Function): + @staticmethod + def forward(ctx, q_BHLD, k_BHLD, v_BHLD, softmax_scale, softmax_lse, is_causal=False, per_block_mean=True, use_global_sf=True): + k_mean = k_BHLD.mean(dim=(0, 1, 2), keepdim=True).view(-1) + + out_BHLD = sageattn_blackwell( + q_BHLD, k_BHLD, v_BHLD, + attn_mask=None, + is_causal=is_causal, + per_block_mean=per_block_mean, + ) + + ctx.HEAD_DIM = q_BHLD.shape[-1] + ctx.sm_scale = softmax_scale # Use sm_scale to match backward access + ctx.causal = is_causal + ctx.IS_QAT = True + ctx.k_mean = k_mean + ctx.use_global_sf = use_global_sf + ctx.save_for_backward(q_BHLD, k_BHLD, v_BHLD, out_BHLD, softmax_lse) + return out_BHLD + + @staticmethod + def backward(ctx, do): + # note this assumes that we do everything in algorithm 1 from SageAttention3 except adding delta s + q_BHLD, k_BHLD, v_BHLD, o_for_bwd_BHLD, M_BHLD = ctx.saved_tensors + do = do.contiguous() + dq_BHLD = torch.empty_like(q_BHLD) + dk_BHLD = torch.empty_like(k_BHLD) + dv_BHLD = torch.empty_like(v_BHLD) + BATCH, N_HEAD, N_CTX_Q = q_BHLD.shape[:3] + N_CTX_KV = k_BHLD.shape[2] + assert k_BHLD.shape[2] == v_BHLD.shape[2], "k and v must have the same sequence length" + PRE_BLOCK = 128 + NUM_STAGES = 1 + NUM_WARPS = 4 + BLOCK_M1, BLOCK_N1, BLOCK_M2, BLOCK_N2 = 64, 64, 64, 64 + BLK_SLICE_FACTOR = 1 + pre_grid = ((N_CTX_Q + PRE_BLOCK - 1) // PRE_BLOCK, BATCH * N_HEAD) + delta = torch.empty_like(M_BHLD) + _attn_bwd_preprocess[pre_grid]( + o_for_bwd_BHLD, do, + delta, + BATCH, N_HEAD, N_CTX_Q, + BLOCK_M=PRE_BLOCK, HEAD_DIM=ctx.HEAD_DIM + ) + + if ctx.IS_QAT: + fake_q = torch.empty_like(q_BHLD) + fake_k = torch.empty_like(k_BHLD) + fake_v = torch.empty_like(v_BHLD) + + grid_1 = (triton.cdiv(q_BHLD.shape[2], BLOCK_M1), q_BHLD.shape[0] * q_BHLD.shape[1], 1) + grid_2 = (triton.cdiv(k_BHLD.shape[2], BLOCK_N2), q_BHLD.shape[0] * q_BHLD.shape[1], 1) + + fake_quantize_q[grid_1]( + q_BHLD, fake_q, + q_BHLD.stride(0), q_BHLD.stride(1), + q_BHLD.stride(2), q_BHLD.stride(3), + fake_q.stride(0), fake_q.stride(1), + fake_q.stride(2), fake_q.stride(3), + N_HEAD, N_CTX_Q, + BLOCK_M=BLOCK_M1, HEAD_DIM=ctx.HEAD_DIM, use_global_sf=ctx.use_global_sf + ) + + fake_quantize_kv[grid_2]( + k_BHLD, v_BHLD, fake_k, fake_v, + k_BHLD.stride(0), k_BHLD.stride(1), + k_BHLD.stride(2), k_BHLD.stride(3), + fake_k.stride(0), fake_k.stride(1), + fake_k.stride(2), fake_k.stride(3), + N_HEAD, N_CTX_KV, + BLOCK_N=BLOCK_N2, HEAD_DIM=ctx.HEAD_DIM, use_global_sf=ctx.use_global_sf + ) + + # Use fake quantized tensors for gradient computation (QAT) + # NOTE: The reassignment below only changes local variable references. + # The gradients (dq_BHLD, dk_BHLD, dv_BHLD) computed using fake_q/fake_k/fake_v + # will still correctly backpropagate to the original q_BHLD/k_BHLD/v_BHLD tensors + # from ctx.saved_tensors, because those original tensors remain in the computation graph. + q_BHLD = fake_q + k_BHLD = fake_k + v_BHLD = fake_v + + # NOTE: K is NOT pre-scaled here - scaling is applied AFTER qk dot product in kernels + # This improves precision by avoiding rounding errors in K before the dot product + arg_k = k_BHLD + + if N_CTX_Q == N_CTX_KV: + # Use existing kernel for self-attention (same sequence lengths) + grid = ((N_CTX_KV + BLOCK_N1 - 1) // BLOCK_N1, 1, BATCH * N_HEAD) + # Q_MEAN is required by the kernel signature but only used if SMOOTH_Q=True + # Since we're not using SMOOTH_Q, create a dummy tensor with the same shape as q_BHLD + q_mean = torch.empty_like(q_BHLD) + _attn_bwd[grid]( + q_BHLD, arg_k, v_BHLD, ctx.sm_scale, do, dq_BHLD, dk_BHLD, dv_BHLD, + M_BHLD, delta, q_mean, + q_BHLD.stride(0), q_BHLD.stride(1), q_BHLD.stride(2), q_BHLD.stride(3), + N_HEAD, N_CTX_KV, + ctx.k_mean, + BLOCK_M1=BLOCK_M1, BLOCK_N1=BLOCK_N1, + BLOCK_M2=BLOCK_M2, BLOCK_N2=BLOCK_N2, + BLK_SLICE_FACTOR=BLK_SLICE_FACTOR, + HEAD_DIM=ctx.HEAD_DIM, + CAUSAL=ctx.causal, + IS_QAT=ctx.IS_QAT, + SMOOTH_K=True, + two_level_quant_P=True, + fake_quant_P=True, + use_global_sf=ctx.use_global_sf, + num_warps=NUM_WARPS, + num_stages=NUM_STAGES + ) + else: + # Use separate kernels for cross-attention (different sequence lengths) + grid_dq = ((N_CTX_Q + BLOCK_M2 - 1) // BLOCK_M2, 1, BATCH * N_HEAD) + _attn_bwd_dq_cross[grid_dq]( + q_BHLD, arg_k, v_BHLD, ctx.sm_scale, do, dq_BHLD, M_BHLD, delta, + q_BHLD.stride(0), k_BHLD.stride(0), q_BHLD.stride(1), k_BHLD.stride(1), q_BHLD.stride(2), k_BHLD.stride(2), q_BHLD.stride(3), k_BHLD.stride(3), + N_HEAD, N_CTX_Q, N_CTX_KV, + ctx.k_mean, + BLOCK_M2=BLOCK_M2, BLOCK_N2=BLOCK_N2, + HEAD_DIM=ctx.HEAD_DIM, + num_warps=NUM_WARPS, + num_stages=NUM_STAGES, + SMOOTH_K=True, + ) + grid_dkdv = ((N_CTX_KV + BLOCK_N1 - 1) // BLOCK_N1, 1, BATCH * N_HEAD) + # Q_MEAN is required by the kernel signature but only used if SMOOTH_Q=True + # Since we're not using SMOOTH_Q, create a dummy tensor with the same shape as q_BHLD + q_mean = torch.empty_like(q_BHLD) + _attn_bwd_dkdv_cross[grid_dkdv]( + q_BHLD, arg_k, v_BHLD, ctx.sm_scale, do, dk_BHLD, dv_BHLD, M_BHLD, delta, q_mean, + q_BHLD.stride(0), k_BHLD.stride(0), q_BHLD.stride(1), k_BHLD.stride(1), q_BHLD.stride(2), k_BHLD.stride(2), q_BHLD.stride(3), k_BHLD.stride(3), + N_HEAD, N_CTX_Q, N_CTX_KV, + BLOCK_M1=BLOCK_M1, BLOCK_N1=BLOCK_N1, + HEAD_DIM=ctx.HEAD_DIM, + IS_QAT=ctx.IS_QAT, + two_level_quant_P=True, + fake_quant_P=True, + use_global_sf=ctx.use_global_sf, + num_warps=NUM_WARPS, + num_stages=NUM_STAGES + ) + + return dq_BHLD, dk_BHLD, dv_BHLD, None, None, None, None, None + +def sageattn_blackwell_with_triton_bwd(q_BLHD, k_BLHD, v_BLHD, is_causal=False, per_block_mean=True, use_global_sf=False): + 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 + with torch.no_grad(): + _, softmax_lse, _S_dmask, rng_state = _wrapped_flash_attn_forward( + q_BLHD, k_BLHD, v_BLHD, + 0.0, # dropout_p + softmax_scale, # softmax_scale + is_causal, # causal + -1, -1, # window_size_left/right + 0.0, # softcap + None, # alibi_slopes + False # return_softmax + ) + + q_BHLD = q_BLHD.permute(0, 2, 1, 3).contiguous() + k_BHLD = k_BLHD.permute(0, 2, 1, 3).contiguous() + v_BHLD = v_BLHD.permute(0, 2, 1, 3).contiguous() + out_BHLD = _SageAttnBlackwellWithTritonBwd.apply(q_BHLD, k_BHLD, v_BHLD, softmax_scale, softmax_lse, is_causal, per_block_mean, use_global_sf) + return out_BHLD.permute(0, 2, 1, 3).contiguous() + +def qat_attn(q_BLHD, k_BLHD, v_BLHD, is_causal=False): + q_BHLD = q_BLHD.permute(0, 2, 1, 3).contiguous() + k_BHLD = k_BLHD.permute(0, 2, 1, 3).contiguous() + v_BHLD = v_BLHD.permute(0, 2, 1, 3).contiguous() + + use_qat_qkv_backward = True + smooth_k = False + warp_specialize = True + IS_QAT = True + two_level_quant_P_sage3 = False + fake_quant_P_bwd = True + use_high_prec_o = True + smooth_q = False + sm_scale = 1.0 / sqrt(q_BHLD.shape[-1]) + use_global_sf_qkv = False + use_global_sf_p = False + o_BHLD = attention(q_BHLD, k_BHLD, v_BHLD, is_causal, sm_scale, + use_qat_qkv_backward, smooth_k, warp_specialize, IS_QAT, + two_level_quant_P_sage3, fake_quant_P_bwd, use_high_prec_o, smooth_q, use_global_sf_p, use_global_sf_qkv) + return o_BHLD.permute(0, 2, 1, 3).contiguous() + + +class ModifiedSageAttention3Backend(AttentionBackend): + + accept_output_buffer: bool = True + + @staticmethod + def get_supported_head_sizes() -> list[int]: + return [64, 128] + + @staticmethod + def get_name() -> str: + return "MODIFIED_SAGE_ATTN_THREE" + + @staticmethod + def get_impl_cls() -> type["ModifiedSageAttention3Impl"]: + return ModifiedSageAttention3Impl + + @staticmethod + def get_metadata_cls() -> type["AttentionMetadata"]: + raise NotImplementedError + + @staticmethod + def get_builder_cls() -> type["AttentionMetadataBuilder"]: + raise NotImplementedError + + # @staticmethod + # def get_metadata_cls() -> Type["AttentionMetadata"]: + # return FlashAttentionMetadata + + +class ModifiedSageAttention3Impl(AttentionImpl): + + def __init__( + self, + num_heads: int, + head_size: int, + causal: bool, + softmax_scale: float, + num_kv_heads: int | None = None, + prefix: str = "", + **extra_impl_args, + ) -> None: + self.causal = causal + self.softmax_scale = softmax_scale + self.dropout = extra_impl_args.get("dropout_p", 0.0) + + def forward( + self, + query: torch.Tensor, + key: torch.Tensor, + value: torch.Tensor, + attn_metadata: AttentionMetadata, + ) -> torch.Tensor: + # output = sageattn_blackwell_with_16bit_bwd(query, key, value, is_causal=self.causal) + # output = sageattn_blackwell_with_triton_bwd(query, key, value, is_causal=self.causal) + output = qat_attn(query, key, value, is_causal=self.causal) + return output + + +SageAttention3Backend = ModifiedSageAttention3Backend +SageAttention3Impl = ModifiedSageAttention3Impl diff --git a/fastvideo/attention/backends/sage_attn3.py b/fastvideo/attention/backends/sage_attn3.py index 83d40855d..3e2aa1eba 100644 --- a/fastvideo/attention/backends/sage_attn3.py +++ b/fastvideo/attention/backends/sage_attn3.py @@ -1,437 +1,14 @@ # SPDX-License-Identifier: Apache-2.0 -import os -import sys -from pathlib import Path - -# Add local repo roots to sys.path for kernel-side imports during development. -_project_root = Path(__file__).resolve().parent.parent.parent.parent -_kernel_root = _project_root / "fastvideo-kernel" -_kernel_python_root = _kernel_root / "python" -for _path in (_project_root, _kernel_root, _kernel_python_root): - if str(_path) not in sys.path: - sys.path.insert(0, str(_path)) - import torch -import triton -try: - # Prefer the in-repo kernel implementation during local development. - from modified_sageattn import sageattn_blackwell -except ImportError: - # Fall back to installed package if local version not available - try: - from sageattn import sageattn_blackwell - except ImportError: - sageattn_blackwell = None +from sageattn3 import sageattn3_blackwell from fastvideo.attention.backends.abstract import (AttentionBackend, AttentionImpl, AttentionMetadata, AttentionMetadataBuilder) from fastvideo.logger import init_logger -from fastvideo_kernel.triton_kernels.qat_attn import attention -from fastvideo_kernel.triton_kernels.qat_attn import ( - _attn_bwd, - _attn_bwd_dkdv_cross, - _attn_bwd_dq_cross, - _attn_bwd_preprocess, -) + logger = init_logger(__name__) -from math import sqrt -try: - from flash_attn.flash_attn_interface import _wrapped_flash_attn_forward, _wrapped_flash_attn_backward -except ImportError: - _wrapped_flash_attn_forward = None - _wrapped_flash_attn_backward = None - -from fastvideo_kernel.triton_kernels.quant_utils import ( - fake_quantize_kv, - fake_quantize_q, -) - -from flashinfer import SfLayout, mm_fp4, nvfp4_quantize - -def matmul_3d_4bit(a, b): - H, M, K1 = a.shape - H, N, K2 = b.shape - assert K1 == K2 - padded_K = (K2 + 127) // 128 * 128 - a = torch.nn.functional.pad(a, (0, padded_K - K2)) - b = torch.nn.functional.pad(b, (0, padded_K - K2)) - out = torch.empty((H, M, N), device=a.device, dtype=a.dtype) - one = torch.ones(1, device=a.device, dtype=torch.float32) - for h in range(H): - q_fp4, q_inv_s = nvfp4_quantize(a[h], one, sfLayout=SfLayout.layout_128x4, do_shuffle=False) - k_fp4, k_inv_s = nvfp4_quantize(b[h], one, sfLayout=SfLayout.layout_128x4, do_shuffle=False) - mm_fp4( - q_fp4, k_fp4.T, q_inv_s, k_inv_s.T, one, - a.dtype, out[h], - block_size=16, - use_8x4_sf_layout=False, - backend="cutlass", - ) - return out - - -class _Matmul3d4bitFWD16bitBWD(torch.autograd.Function): - @staticmethod - def forward(ctx, a, b, scale=None): - ctx.save_for_backward(a, b) - if scale is not None: - a = a / scale - return matmul_3d_4bit(a, b) * scale - else: - return matmul_3d_4bit(a, b) - - @staticmethod - def backward(ctx, grad_out): - a, b = ctx.saved_tensors - grad_a = grad_out.matmul(b) - grad_b = grad_out.transpose(1, 2).matmul(a) - # None for the three extra forward args - return grad_a, grad_b, None - - -class _Matmul3d4bitFWD16bitBWD_with_smoothing(torch.autograd.Function): - @staticmethod - def forward(ctx, q, k): - ctx.save_for_backward(q, k) - k = k - k.mean(dim=-2, keepdim=True) # [H, L, D] - qm = q.mean(dim=-2, keepdim=True) # [H, 1, D] - q = q - qm # [B, H, L, D] - delta_s = torch.matmul(qm, k.transpose(-2, -1)).contiguous() # [H, 1, L] - return matmul_3d_4bit(q, k) + delta_s - - @staticmethod - def backward(ctx, grad_out): - a, b = ctx.saved_tensors - grad_a = grad_out.matmul(b) - grad_b = grad_out.transpose(1, 2).matmul(a) - # None for the three extra forward args - return grad_a, grad_b, None - -USE_SMOOTHING = False -def sage_attn_qk_torch(q, k, per_block_mean=False): - """ - Inputs: q/k/v in [B, H, L, D] - Returns: out in [B, H, L, D] - """ - B, H, qL, D = q.shape - B, H, kL, D = k.shape - assert B == 1 - q, k = q.squeeze(0), k.squeeze(0) - if USE_SMOOTHING: - qk = _Matmul3d4bitFWD16bitBWD_with_smoothing.apply(q, k) - else: - qk = _Matmul3d4bitFWD16bitBWD.apply(q, k) - qk = qk / torch.sqrt(torch.tensor(D, device=q.device, dtype=q.dtype)) - p = torch.softmax(qk, dim=-1) - return p.unsqueeze(0) - -def attn_backward_16bit(q_BLHD, k_BLHD, v_BLHD, out_BLHD, grad_out_BLHD, is_causal=False): - assert is_causal == False - q, k, v, o, do = q_BLHD.transpose(1 ,2).contiguous(), k_BLHD.transpose(1 ,2).contiguous(), v_BLHD.transpose(1 ,2).contiguous(), out_BLHD.transpose(1 ,2).contiguous(), grad_out_BLHD.transpose(1 ,2).contiguous() - B, H, _, D = q.shape - p = sage_attn_qk_torch(q, k, per_block_mean=False).to(do.dtype) - dV = p.transpose(-2, -1) @ do # [B,H,L,D] - - dP = do @ v.transpose(-2, -1) # [B,H,L,L] - - # do_dot_o = (do.float() * o.float()).sum(dim=-1, keepdim=True).to(dP.dtype) # [B,H,L,1] - # do_dot_o = (do * o).sum(dim=-1, keepdim=True) # [B,H,L,1] - dp_dot_p = (dP * p).sum(dim=-1, keepdim=True) # [B,H,L,1] - dS = p * (dP - dp_dot_p) # [B,H,L,L] - - - dQ = (dS @ k) / torch.sqrt(torch.tensor(D, device=q.device, dtype=q.dtype)) # [B,H,L,D] - dK = (dS.transpose(-2, -1) @ q) / torch.sqrt(torch.tensor(D, device=q.device, dtype=q.dtype)) # [B,H,L,D] - - # grads w.r.t inputs, None for non-tensor inputs - return dQ.transpose(1, 2).contiguous(), dK.transpose(1, 2).contiguous(), dV.transpose(1, 2).contiguous(), None, None, None - -def attn_forward_4bit(q_BHLD, k_BHLD, v_BHLD, is_causal=False): - assert is_causal == False - p_bf16 = sage_attn_qk_torch(q_BHLD, k_BHLD, per_block_mean=False).to(q_BHLD.dtype) - row_max_p = p_bf16.max(dim=-1, keepdim=True).values - scale_row = row_max_p / (448 * 6) - p = p_bf16 / scale_row - out = matmul_3d_4bit(p.squeeze(0), v_BHLD.squeeze(0).transpose(1, 2).contiguous()).unsqueeze(0) - out = out * scale_row - return out, p_bf16 - - -def flash_attn_backward(q_BLHD, k_BLHD, v_BLHD, out_BLHD, grad_out_BLHD, is_causal=False): - D = q_BLHD.shape[-1] - softmax_scale = 1.0 / sqrt(D) - - # FA forward only to get softmax_lse (rng_state is also returned) - _, softmax_lse, _S_dmask, rng_state = _wrapped_flash_attn_forward( - q_BLHD, k_BLHD, v_BLHD, - 0.0, # dropout_p - softmax_scale, # softmax_scale - is_causal, # causal - -1, -1, # window_size_left/right - 0.0, # softcap - None, # alibi_slopes - False # return_softmax - ) - - # Allocate grads and call FA backward using the SAGE output - dq_BLHD = torch.empty_like(q_BLHD) - dk_BLHD = torch.empty_like(k_BLHD) - dv_BLHD = torch.empty_like(v_BLHD) - - _wrapped_flash_attn_backward( - grad_out_BLHD, # dout (BLHD, 16-bit) - q_BLHD, k_BLHD, v_BLHD, - out_BLHD, # use SAGE forward output here - softmax_lse, - dq_BLHD, dk_BLHD, dv_BLHD, - 0.0, # dropout_p - softmax_scale, # softmax_scale - is_causal, # causal - -1, -1, # window_size_left/right - 0.0, # softcap - None, # alibi_slopes - False, # deterministic - rng_state=rng_state, - ) - return dq_BLHD, dk_BLHD, dv_BLHD, None, None - -USE_TORCH_4BIT_FWD = False -USE_TORCH_16BIT_BWD = False -USE_QAT_ATTN = False -class _SageAttnBlackwellWith16bitBwd(torch.autograd.Function): - @staticmethod - def forward(ctx, q_BLHD, k_BLHD, v_BLHD, is_causal=False, per_block_mean=True): - """ - Inputs: q/k/v in [B, L, H, D] - Returns: out in [B, L, H, D] - """ - # Save originals for backward (we'll recompute FA in 16-bit there) - ctx.is_causal = bool(is_causal) - - # Convert to BHLD for your sageattn forward - q_BHLD = q_BLHD.permute(0, 2, 1, 3).contiguous() - k_BHLD = k_BLHD.permute(0, 2, 1, 3).contiguous() - v_BHLD = v_BLHD.permute(0, 2, 1, 3).contiguous() - - with torch.no_grad(): - out_BHLD = sageattn_blackwell( - q_BHLD, k_BHLD, v_BHLD, - attn_mask=None, - is_causal=is_causal, - per_block_mean=per_block_mean, - ) - # Convert to BLHD for saving - out_BLHD = out_BHLD.permute(0, 2, 1, 3).contiguous() - ctx.save_for_backward(q_BLHD, k_BLHD, v_BLHD, out_BLHD) - - return out_BLHD - - @staticmethod - def backward(ctx, grad_out_BLHD): - is_causal = ctx.is_causal - q_BLHD, k_BLHD, v_BLHD, out_BLHD = ctx.saved_tensors - return flash_attn_backward(q_BLHD, k_BLHD, v_BLHD, out_BLHD, grad_out_BLHD, is_causal) - - - -def sageattn_blackwell_with_16bit_bwd(q_BLHD, k_BLHD, v_BLHD, is_causal=False, per_block_mean=True): - """ - Forward: uses sageattn_blackwell under the hood. - Backward: recomputes FlashAttention fwd+bwd in 16bit directly. - """ - return _SageAttnBlackwellWith16bitBwd.apply(q_BLHD, k_BLHD, v_BLHD, is_causal, per_block_mean) - -class _SageAttnBlackwellWithTritonBwd(torch.autograd.Function): - @staticmethod - def forward(ctx, q_BHLD, k_BHLD, v_BHLD, softmax_scale, softmax_lse, is_causal=False, per_block_mean=True, use_global_sf=True): - k_mean = k_BHLD.mean(dim=(0, 1, 2), keepdim=True).view(-1) - - out_BHLD = sageattn_blackwell( - q_BHLD, k_BHLD, v_BHLD, - attn_mask=None, - is_causal=is_causal, - per_block_mean=per_block_mean, - ) - - ctx.HEAD_DIM = q_BHLD.shape[-1] - ctx.sm_scale = softmax_scale # Use sm_scale to match backward access - ctx.causal = is_causal - ctx.IS_QAT = True - ctx.k_mean = k_mean - ctx.use_global_sf = use_global_sf - ctx.save_for_backward(q_BHLD, k_BHLD, v_BHLD, out_BHLD, softmax_lse) - return out_BHLD - - @staticmethod - def backward(ctx, do): - # note this assumes that we do everything in algorithm 1 from SageAttention3 except adding delta s - q_BHLD, k_BHLD, v_BHLD, o_for_bwd_BHLD, M_BHLD = ctx.saved_tensors - do = do.contiguous() - dq_BHLD = torch.empty_like(q_BHLD) - dk_BHLD = torch.empty_like(k_BHLD) - dv_BHLD = torch.empty_like(v_BHLD) - BATCH, N_HEAD, N_CTX_Q = q_BHLD.shape[:3] - N_CTX_KV = k_BHLD.shape[2] - assert k_BHLD.shape[2] == v_BHLD.shape[2], "k and v must have the same sequence length" - PRE_BLOCK = 128 - NUM_STAGES = 1 - NUM_WARPS = 4 - BLOCK_M1, BLOCK_N1, BLOCK_M2, BLOCK_N2 = 64, 64, 64, 64 - BLK_SLICE_FACTOR = 1 - pre_grid = ((N_CTX_Q + PRE_BLOCK - 1) // PRE_BLOCK, BATCH * N_HEAD) - delta = torch.empty_like(M_BHLD) - _attn_bwd_preprocess[pre_grid]( - o_for_bwd_BHLD, do, - delta, - BATCH, N_HEAD, N_CTX_Q, - BLOCK_M=PRE_BLOCK, HEAD_DIM=ctx.HEAD_DIM - ) - - if ctx.IS_QAT: - fake_q = torch.empty_like(q_BHLD) - fake_k = torch.empty_like(k_BHLD) - fake_v = torch.empty_like(v_BHLD) - - grid_1 = (triton.cdiv(q_BHLD.shape[2], BLOCK_M1), q_BHLD.shape[0] * q_BHLD.shape[1], 1) - grid_2 = (triton.cdiv(k_BHLD.shape[2], BLOCK_N2), q_BHLD.shape[0] * q_BHLD.shape[1], 1) - - fake_quantize_q[grid_1]( - q_BHLD, fake_q, - q_BHLD.stride(0), q_BHLD.stride(1), - q_BHLD.stride(2), q_BHLD.stride(3), - fake_q.stride(0), fake_q.stride(1), - fake_q.stride(2), fake_q.stride(3), - N_HEAD, N_CTX_Q, - BLOCK_M=BLOCK_M1, HEAD_DIM=ctx.HEAD_DIM, use_global_sf=ctx.use_global_sf - ) - - fake_quantize_kv[grid_2]( - k_BHLD, v_BHLD, fake_k, fake_v, - k_BHLD.stride(0), k_BHLD.stride(1), - k_BHLD.stride(2), k_BHLD.stride(3), - fake_k.stride(0), fake_k.stride(1), - fake_k.stride(2), fake_k.stride(3), - N_HEAD, N_CTX_KV, - BLOCK_N=BLOCK_N2, HEAD_DIM=ctx.HEAD_DIM, use_global_sf=ctx.use_global_sf - ) - - # Use fake quantized tensors for gradient computation (QAT) - # NOTE: The reassignment below only changes local variable references. - # The gradients (dq_BHLD, dk_BHLD, dv_BHLD) computed using fake_q/fake_k/fake_v - # will still correctly backpropagate to the original q_BHLD/k_BHLD/v_BHLD tensors - # from ctx.saved_tensors, because those original tensors remain in the computation graph. - q_BHLD = fake_q - k_BHLD = fake_k - v_BHLD = fake_v - - # NOTE: K is NOT pre-scaled here - scaling is applied AFTER qk dot product in kernels - # This improves precision by avoiding rounding errors in K before the dot product - arg_k = k_BHLD - - if N_CTX_Q == N_CTX_KV: - # Use existing kernel for self-attention (same sequence lengths) - grid = ((N_CTX_KV + BLOCK_N1 - 1) // BLOCK_N1, 1, BATCH * N_HEAD) - # Q_MEAN is required by the kernel signature but only used if SMOOTH_Q=True - # Since we're not using SMOOTH_Q, create a dummy tensor with the same shape as q_BHLD - q_mean = torch.empty_like(q_BHLD) - _attn_bwd[grid]( - q_BHLD, arg_k, v_BHLD, ctx.sm_scale, do, dq_BHLD, dk_BHLD, dv_BHLD, - M_BHLD, delta, q_mean, - q_BHLD.stride(0), q_BHLD.stride(1), q_BHLD.stride(2), q_BHLD.stride(3), - N_HEAD, N_CTX_KV, - ctx.k_mean, - BLOCK_M1=BLOCK_M1, BLOCK_N1=BLOCK_N1, - BLOCK_M2=BLOCK_M2, BLOCK_N2=BLOCK_N2, - BLK_SLICE_FACTOR=BLK_SLICE_FACTOR, - HEAD_DIM=ctx.HEAD_DIM, - CAUSAL=ctx.causal, - IS_QAT=ctx.IS_QAT, - SMOOTH_K=True, - two_level_quant_P=True, - fake_quant_P=True, - use_global_sf=ctx.use_global_sf, - num_warps=NUM_WARPS, - num_stages=NUM_STAGES - ) - else: - # Use separate kernels for cross-attention (different sequence lengths) - grid_dq = ((N_CTX_Q + BLOCK_M2 - 1) // BLOCK_M2, 1, BATCH * N_HEAD) - _attn_bwd_dq_cross[grid_dq]( - q_BHLD, arg_k, v_BHLD, ctx.sm_scale, do, dq_BHLD, M_BHLD, delta, - q_BHLD.stride(0), k_BHLD.stride(0), q_BHLD.stride(1), k_BHLD.stride(1), q_BHLD.stride(2), k_BHLD.stride(2), q_BHLD.stride(3), k_BHLD.stride(3), - N_HEAD, N_CTX_Q, N_CTX_KV, - ctx.k_mean, - BLOCK_M2=BLOCK_M2, BLOCK_N2=BLOCK_N2, - HEAD_DIM=ctx.HEAD_DIM, - num_warps=NUM_WARPS, - num_stages=NUM_STAGES, - SMOOTH_K=True, - ) - grid_dkdv = ((N_CTX_KV + BLOCK_N1 - 1) // BLOCK_N1, 1, BATCH * N_HEAD) - # Q_MEAN is required by the kernel signature but only used if SMOOTH_Q=True - # Since we're not using SMOOTH_Q, create a dummy tensor with the same shape as q_BHLD - q_mean = torch.empty_like(q_BHLD) - _attn_bwd_dkdv_cross[grid_dkdv]( - q_BHLD, arg_k, v_BHLD, ctx.sm_scale, do, dk_BHLD, dv_BHLD, M_BHLD, delta, q_mean, - q_BHLD.stride(0), k_BHLD.stride(0), q_BHLD.stride(1), k_BHLD.stride(1), q_BHLD.stride(2), k_BHLD.stride(2), q_BHLD.stride(3), k_BHLD.stride(3), - N_HEAD, N_CTX_Q, N_CTX_KV, - BLOCK_M1=BLOCK_M1, BLOCK_N1=BLOCK_N1, - HEAD_DIM=ctx.HEAD_DIM, - IS_QAT=ctx.IS_QAT, - two_level_quant_P=True, - fake_quant_P=True, - use_global_sf=ctx.use_global_sf, - num_warps=NUM_WARPS, - num_stages=NUM_STAGES - ) - - return dq_BHLD, dk_BHLD, dv_BHLD, None, None, None, None, None - -def sageattn_blackwell_with_triton_bwd(q_BLHD, k_BLHD, v_BLHD, is_causal=False, per_block_mean=True, use_global_sf=False): - 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 - with torch.no_grad(): - _, softmax_lse, _S_dmask, rng_state = _wrapped_flash_attn_forward( - q_BLHD, k_BLHD, v_BLHD, - 0.0, # dropout_p - softmax_scale, # softmax_scale - is_causal, # causal - -1, -1, # window_size_left/right - 0.0, # softcap - None, # alibi_slopes - False # return_softmax - ) - - q_BHLD = q_BLHD.permute(0, 2, 1, 3).contiguous() - k_BHLD = k_BLHD.permute(0, 2, 1, 3).contiguous() - v_BHLD = v_BLHD.permute(0, 2, 1, 3).contiguous() - out_BHLD = _SageAttnBlackwellWithTritonBwd.apply(q_BHLD, k_BHLD, v_BHLD, softmax_scale, softmax_lse, is_causal, per_block_mean, use_global_sf) - return out_BHLD.permute(0, 2, 1, 3).contiguous() - -def qat_attn(q_BLHD, k_BLHD, v_BLHD, is_causal=False): - q_BHLD = q_BLHD.permute(0, 2, 1, 3).contiguous() - k_BHLD = k_BLHD.permute(0, 2, 1, 3).contiguous() - v_BHLD = v_BLHD.permute(0, 2, 1, 3).contiguous() - - use_qat_qkv_backward = True - smooth_k = False - warp_specialize = True - IS_QAT = True - two_level_quant_P_sage3 = False - fake_quant_P_bwd = True - use_high_prec_o = True - smooth_q = False - sm_scale = 1.0 / sqrt(q_BHLD.shape[-1]) - use_global_sf_qkv = False - use_global_sf_p = False - o_BHLD = attention(q_BHLD, k_BHLD, v_BHLD, is_causal, sm_scale, - use_qat_qkv_backward, smooth_k, warp_specialize, IS_QAT, - two_level_quant_P_sage3, fake_quant_P_bwd, use_high_prec_o, smooth_q, use_global_sf_p, use_global_sf_qkv) - return o_BHLD.permute(0, 2, 1, 3).contiguous() - class SageAttention3Backend(AttentionBackend): @@ -439,7 +16,7 @@ class SageAttention3Backend(AttentionBackend): @staticmethod def get_supported_head_sizes() -> list[int]: - return [64, 128] + return [64, 128, 256] @staticmethod def get_name() -> str: @@ -485,7 +62,9 @@ class SageAttention3Impl(AttentionImpl): value: torch.Tensor, attn_metadata: AttentionMetadata, ) -> torch.Tensor: - # output = sageattn_blackwell_with_16bit_bwd(query, key, value, is_causal=self.causal) - # output = sageattn_blackwell_with_triton_bwd(query, key, value, is_causal=self.causal) - output = qat_attn(query, key, value, is_causal=self.causal) + query = query.transpose(1, 2) + key = key.transpose(1, 2) + value = value.transpose(1, 2) + output = sageattn3_blackwell(query, key, value, is_causal=self.causal) + output = output.transpose(1, 2) return output \ No newline at end of file diff --git a/fastvideo/configs/models/dits/base.py b/fastvideo/configs/models/dits/base.py index 3b70c33aa..b99995dd2 100644 --- a/fastvideo/configs/models/dits/base.py +++ b/fastvideo/configs/models/dits/base.py @@ -19,6 +19,7 @@ class DiTArchConfig(ArchConfig): AttentionBackendEnum.TORCH_SDPA, AttentionBackendEnum.VIDEO_SPARSE_ATTN, AttentionBackendEnum.VMOBA_ATTN, AttentionBackendEnum.SAGE_ATTN_THREE, + AttentionBackendEnum.MODIFIED_SAGE_ATTN_THREE, AttentionBackendEnum.SLA_ATTN, AttentionBackendEnum.SAGE_SLA_ATTN) hidden_size: int = 0 diff --git a/fastvideo/envs.py b/fastvideo/envs.py index b6de4d588..539c3fe0f 100644 --- a/fastvideo/envs.py +++ b/fastvideo/envs.py @@ -201,6 +201,7 @@ environment_variables: dict[str, Callable[[], Any]] = { # - "VIDEO_SPARSE_ATTN": use Video Sparse Attention # - "SAGE_ATTN": use Sage Attention # - "SAGE_ATTN_THREE": use Sage Attention 3 + # - "MODIFIED_SAGE_ATTN_THREE": use the in-repo modified Sage Attention 3 "FASTVIDEO_ATTENTION_BACKEND": lambda: os.getenv("FASTVIDEO_ATTENTION_BACKEND", None), diff --git a/fastvideo/pipelines/stages/denoising.py b/fastvideo/pipelines/stages/denoising.py index 64bdfdffa..c085aaf7d 100644 --- a/fastvideo/pipelines/stages/denoising.py +++ b/fastvideo/pipelines/stages/denoising.py @@ -63,7 +63,8 @@ class DenoisingStage(PipelineStage): dtype=torch.float16, # TODO(will): hack supported_attention_backends=(AttentionBackendEnum.VIDEO_SPARSE_ATTN, AttentionBackendEnum.BSA_ATTN, AttentionBackendEnum.VMOBA_ATTN, AttentionBackendEnum.FLASH_ATTN, - AttentionBackendEnum.TORCH_SDPA, AttentionBackendEnum.SAGE_ATTN_THREE) # hack + AttentionBackendEnum.TORCH_SDPA, AttentionBackendEnum.SAGE_ATTN_THREE, + AttentionBackendEnum.MODIFIED_SAGE_ATTN_THREE) # hack ) def forward( diff --git a/fastvideo/pipelines/stages/sr_denoising.py b/fastvideo/pipelines/stages/sr_denoising.py index 0102770d7..2fae06df1 100644 --- a/fastvideo/pipelines/stages/sr_denoising.py +++ b/fastvideo/pipelines/stages/sr_denoising.py @@ -63,7 +63,8 @@ class SRDenoisingStage(PipelineStage): dtype=torch.float16, # TODO(will): hack supported_attention_backends=(AttentionBackendEnum.VIDEO_SPARSE_ATTN, AttentionBackendEnum.VMOBA_ATTN, AttentionBackendEnum.FLASH_ATTN, AttentionBackendEnum.TORCH_SDPA, - AttentionBackendEnum.SAGE_ATTN_THREE) # hack + AttentionBackendEnum.SAGE_ATTN_THREE, + AttentionBackendEnum.MODIFIED_SAGE_ATTN_THREE) # hack ) def add_noise_to_lq(self, lq_latents: torch.Tensor, strength: float = 0.7) -> torch.Tensor: diff --git a/fastvideo/platforms/cuda.py b/fastvideo/platforms/cuda.py index 1b048fcda..582ff477a 100644 --- a/fastvideo/platforms/cuda.py +++ b/fastvideo/platforms/cuda.py @@ -130,7 +130,7 @@ class CudaPlatformBase(Platform): logger.info("Sage Attention backend is not installed. Fall back to Flash Attention.") elif selected_backend == AttentionBackendEnum.SAGE_ATTN_THREE: try: - from modified_sageattn import sageattn_blackwell # noqa: F401 + from sageattn3 import sageattn3_blackwell # noqa: F401 from fastvideo.attention.backends.sage_attn3 import ( # noqa: F401 SageAttention3Backend) @@ -140,6 +140,18 @@ class CudaPlatformBase(Platform): except ImportError as e: logger.info(e) logger.info("Sage Attention 3 backend is not installed. Fall back to Flash Attention.") + elif selected_backend == AttentionBackendEnum.MODIFIED_SAGE_ATTN_THREE: + try: + from modified_sageattn import sageattn_blackwell # noqa: F401 + + from fastvideo.attention.backends.modified_sage_attn3 import ( # noqa: F401 + ModifiedSageAttention3Backend) + logger.info("Using modified Sage Attention 3 backend.") + + return "fastvideo.attention.backends.modified_sage_attn3.ModifiedSageAttention3Backend" + except ImportError as e: + logger.info(e) + logger.info("Modified Sage Attention 3 backend is not installed. Fall back to Flash Attention.") elif selected_backend == AttentionBackendEnum.VIDEO_SPARSE_ATTN: try: from fastvideo_kernel import video_sparse_attn # noqa: F401 diff --git a/fastvideo/platforms/interface.py b/fastvideo/platforms/interface.py index 60e2441f6..385b4afbb 100644 --- a/fastvideo/platforms/interface.py +++ b/fastvideo/platforms/interface.py @@ -15,6 +15,7 @@ class AttentionBackendEnum(enum.Enum): TORCH_SDPA = enum.auto() SAGE_ATTN = enum.auto() SAGE_ATTN_THREE = enum.auto() + MODIFIED_SAGE_ATTN_THREE = enum.auto() VIDEO_SPARSE_ATTN = enum.auto() BSA_ATTN = enum.auto() VMOBA_ATTN = enum.auto() diff --git a/fastvideo/training/training_pipeline.py b/fastvideo/training/training_pipeline.py index 7f0279e41..2945baebf 100644 --- a/fastvideo/training/training_pipeline.py +++ b/fastvideo/training/training_pipeline.py @@ -55,12 +55,24 @@ except Exception: logger = init_logger(__name__) +def _get_generator_sage_attn3_impl_cls() -> type[Any]: + from fastvideo.attention.backends.modified_sage_attn3 import ( + ModifiedSageAttention3Impl, + ) + + return ModifiedSageAttention3Impl + + def swap_sage_attn3(obj: Any, obj_path: str) -> int: """ - If `obj` has an attribute named `attn_impl`, replace it with SageAttention3Impl, - carrying over `.causal` and `.softmax_scale`. Everything else set to None. + If `obj` has an attribute named `attn_impl`, replace it with the modified + SageAttention3 implementation used by generator 4-bit attention, carrying + over `.causal` and `.softmax_scale`. + Returns number of swaps performed (0 or 1). """ + impl_cls = _get_generator_sage_attn3_impl_cls() + # Quick reject: no attribute if not hasattr(obj, "attn_impl"): return 0 @@ -72,7 +84,7 @@ def swap_sage_attn3(obj: Any, obj_path: str) -> int: return 0 # Already None or already of the desired type => skip - if old_impl is None or isinstance(old_impl, SageAttention3Impl): + if old_impl is None or isinstance(old_impl, impl_cls): return 0 # Extract fields with sensible defaults @@ -80,7 +92,7 @@ def swap_sage_attn3(obj: Any, obj_path: str) -> int: softmax_scale = getattr(old_impl, "softmax_scale", 1.0) - new_impl = SageAttention3Impl( + new_impl = impl_cls( num_heads=None, head_size=None, causal=causal,