diff --git a/examples/training/finetune/wan_t2v_1.3B/crush_smol/finetune_t2v_sage3_5090.sh b/examples/training/finetune/wan_t2v_1.3B/crush_smol/finetune_t2v_sage3_5090.sh index c50feae70..a24cacbed 100644 --- a/examples/training/finetune/wan_t2v_1.3B/crush_smol/finetune_t2v_sage3_5090.sh +++ b/examples/training/finetune/wan_t2v_1.3B/crush_smol/finetune_t2v_sage3_5090.sh @@ -53,7 +53,7 @@ dataset_args=( # Validation arguments validation_args=( - --log_validation + # --log_validation --validation_dataset_file $VALIDATION_DATASET_FILE --validation_steps 100 --validation_sampling_steps "50" diff --git a/fastvideo/attention/backends/sage_attn3.py b/fastvideo/attention/backends/sage_attn3.py index f24844138..73875b73d 100644 --- a/fastvideo/attention/backends/sage_attn3.py +++ b/fastvideo/attention/backends/sage_attn3.py @@ -10,6 +10,7 @@ if str(_project_root) not in sys.path: sys.path.insert(0, str(_project_root)) import torch +import triton try: # Use local version first (avoids issues with installed package) from fastvideo.attention.backends.sageattn import sageattn_blackwell @@ -40,6 +41,7 @@ except ImportError: logger = init_logger(__name__) from flashinfer import SfLayout, mm_fp4, nvfp4_quantize +from quant_utils import fake_quantize_q, fake_quantize_kv def fp_16_qk(q, k): qk = torch.matmul(q, k) / torch.sqrt(torch.tensor(q.shape[-1], device=q.device, dtype=q.dtype)) @@ -321,7 +323,7 @@ def sageattn_blackwell_with_16bit_bwd(q_BLHD, k_BLHD, v_BLHD, is_causal=False, p 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): + 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( @@ -336,6 +338,7 @@ class _SageAttnBlackwellWithTritonBwd(torch.autograd.Function): 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 @@ -355,9 +358,6 @@ class _SageAttnBlackwellWithTritonBwd(torch.autograd.Function): NUM_WARPS = 4 BLOCK_M1, BLOCK_N1, BLOCK_M2, BLOCK_N2 = 64, 64, 64, 64 BLK_SLICE_FACTOR = 1 - # 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 pre_grid = ((N_CTX_Q + PRE_BLOCK - 1) // PRE_BLOCK, BATCH * N_HEAD) delta = torch.empty_like(M_BHLD) _attn_bwd_preprocess[pre_grid]( @@ -367,12 +367,55 @@ class _SageAttnBlackwellWithTritonBwd(torch.autograd.Function): BLOCK_M=PRE_BLOCK, HEAD_DIM=ctx.HEAD_DIM ) + 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, + 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, @@ -385,6 +428,7 @@ class _SageAttnBlackwellWithTritonBwd(torch.autograd.Function): 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 ) @@ -403,8 +447,11 @@ class _SageAttnBlackwellWithTritonBwd(torch.autograd.Function): 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_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, @@ -412,13 +459,14 @@ class _SageAttnBlackwellWithTritonBwd(torch.autograd.Function): 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 + + 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): +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 @@ -437,14 +485,14 @@ def sageattn_blackwell_with_triton_bwd(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() - out_BHLD = _SageAttnBlackwellWithTritonBwd.apply(q_BHLD, k_BHLD, v_BHLD, softmax_scale, softmax_lse, is_causal, per_block_mean) + 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): +def qat_attn(q_BLHD, k_BLHD, v_BLHD, is_causal=False, use_global_sf=True): 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() - o_BHLD = attention(q_BHLD, k_BHLD, v_BHLD, is_causal, 1.0 / sqrt(q_BLHD.shape[-1])) + o_BHLD = attention(q_BHLD, k_BHLD, v_BHLD, is_causal, 1.0 / sqrt(q_BLHD.shape[-1]), use_global_sf=use_global_sf) return o_BHLD.permute(0, 2, 1, 3).contiguous() diff --git a/qat_attn.py b/qat_attn.py index ce9dac3c4..a05502f68 100644 --- a/qat_attn.py +++ b/qat_attn.py @@ -45,7 +45,8 @@ def _attn_fwd_inner(acc, high_prec_acc, l_i, m_i, q, q_valid, STAGE: tl.constexpr, offs_m: tl.constexpr, offs_n: tl.constexpr, N_CTX: tl.constexpr, warp_specialize: tl.constexpr, IS_HOPPER: tl.constexpr, IS_QAT: tl.constexpr, - two_level_quant_P: tl.constexpr = False): + two_level_quant_P: tl.constexpr = False, + use_global_sf: tl.constexpr = True): # range of values handled by this stage (kv blocks) if STAGE == 1: lo, hi = 0, start_m * BLOCK_M @@ -85,7 +86,8 @@ def _attn_fwd_inner(acc, high_prec_acc, l_i, m_i, q, q_valid, BLOCK_SIZE_OUT_DIM=BLOCK_M, BLOCK_SIZE_QUANT_DIM=BLOCK_N, dst_dtype=dtype, - two_level_quant_P=two_level_quant_P + two_level_quant_P=two_level_quant_P, + use_global_sf=use_global_sf ) l_ij = tl.sum(high_prec_p, 1) else: @@ -199,6 +201,7 @@ def _attn_fwd(sm_scale, M, IS_HOPPER: tl.constexpr, IS_QAT: tl.constexpr, two_level_quant_P: tl.constexpr = False, + use_global_sf: tl.constexpr = True, ): dtype = tl.float8e5 if FP8_OUTPUT else tl.bfloat16 tl.static_assert(BLOCK_N <= HEAD_DIM) @@ -263,7 +266,7 @@ def _attn_fwd(sm_scale, M, offset_y_kv, dtype, start_m, qk_scale, BLOCK_M, HEAD_DIM, BLOCK_N, 4 - STAGE, offs_m, offs_n, N_CTX_KV, - warp_specialize, IS_HOPPER, IS_QAT, two_level_quant_P + warp_specialize, IS_HOPPER, IS_QAT, two_level_quant_P, use_global_sf ) # stage 2: on-band if STAGE & 2: @@ -273,7 +276,7 @@ def _attn_fwd(sm_scale, M, offset_y_kv, dtype, start_m, qk_scale, BLOCK_M, HEAD_DIM, BLOCK_N, 2, offs_m, offs_n, N_CTX_KV, - warp_specialize, IS_HOPPER, IS_QAT, two_level_quant_P + warp_specialize, IS_HOPPER, IS_QAT, two_level_quant_P, use_global_sf ) # epilogue m_i += tl.math.log2(l_i) @@ -322,7 +325,8 @@ def _attn_bwd_dkdv(dk, dv, IS_QAT: tl.constexpr, two_level_quant_P: tl.constexpr = False, fake_quant_P: tl.constexpr = True, - SMOOTH_Q: tl.constexpr = False): + SMOOTH_Q: tl.constexpr = False, + use_global_sf: tl.constexpr = True): offs_m = start_m + tl.arange(0, BLOCK_M1) offs_n = start_n + tl.arange(0, BLOCK_N1) offs_k = tl.arange(0, HEAD_DIM) @@ -359,7 +363,8 @@ def _attn_bwd_dkdv(dk, dv, BLOCK_SIZE_OUT_DIM=BLOCK_M1, BLOCK_SIZE_QUANT_DIM=BLOCK_N1, dst_dtype=p.dtype, - two_level_quant_P=two_level_quant_P + two_level_quant_P=two_level_quant_P, + use_global_sf=use_global_sf ) dv += tl.dot(tl.trans(p_quant.to(tl.bfloat16)), do) # D (= delta) is pre-divided by ds_scale. @@ -531,7 +536,8 @@ def _attn_bwd_dkdv_cross(Q, K, V, sm_scale, IS_QAT: tl.constexpr, two_level_quant_P: tl.constexpr = False, fake_quant_P: tl.constexpr = True, - SMOOTH_Q: tl.constexpr = False + SMOOTH_Q: tl.constexpr = False, + use_global_sf: tl.constexpr = True ): # Apply scale AFTER dot product for better precision RCP_LN2: tl.constexpr = 1.4426950408889634 # = 1.0 / ln(2) @@ -586,7 +592,8 @@ def _attn_bwd_dkdv_cross(Q, K, V, sm_scale, BLOCK_SIZE_OUT_DIM=BLOCK_M1, BLOCK_SIZE_QUANT_DIM=BLOCK_N1, dst_dtype=p.dtype, - two_level_quant_P=two_level_quant_P + two_level_quant_P=two_level_quant_P, + use_global_sf=use_global_sf ) dv += tl.dot(tl.trans(p_quant.to(tl.bfloat16)), do) @@ -628,7 +635,8 @@ def _attn_bwd(Q, K, V, sm_scale, SMOOTH_K: tl.constexpr, two_level_quant_P: tl.constexpr = False, fake_quant_P: tl.constexpr = True, - SMOOTH_Q: tl.constexpr = False): + SMOOTH_Q: tl.constexpr = False, + use_global_sf: tl.constexpr = True): bhid = tl.program_id(2) off_chz = (bhid * N_CTX).to(tl.int64) @@ -691,7 +699,8 @@ def _attn_bwd(Q, K, V, sm_scale, IS_QAT=IS_QAT, two_level_quant_P=two_level_quant_P, fake_quant_P=fake_quant_P, - SMOOTH_Q=SMOOTH_Q + SMOOTH_Q=SMOOTH_Q, + use_global_sf=use_global_sf ) start_m += num_steps * MASK_BLOCK_M1 @@ -715,7 +724,8 @@ def _attn_bwd(Q, K, V, sm_scale, IS_QAT=IS_QAT, two_level_quant_P=two_level_quant_P, fake_quant_P=fake_quant_P, - SMOOTH_Q=SMOOTH_Q + SMOOTH_Q=SMOOTH_Q, + use_global_sf=use_global_sf ) dv_ptrs = DV + offs_n[:, None] * stride_tok + offs_k[None, :] * stride_d @@ -803,7 +813,8 @@ class _attention(torch.autograd.Function): two_level_quant_P=True, fake_quant_P=True, use_high_prec_o=False, - smooth_q=False + smooth_q=False, + use_global_sf=True ): # shape constraints HEAD_DIM_Q, HEAD_DIM_K = q.shape[-1], k.shape[-1] @@ -968,6 +979,7 @@ class _attention(torch.autograd.Function): IS_HOPPER=is_hopper(), IS_QAT=IS_QAT, two_level_quant_P=two_level_quant_P, + use_global_sf=use_global_sf, num_warps=4, num_stages=2, **extra_kern_args @@ -989,6 +1001,7 @@ class _attention(torch.autograd.Function): ctx.two_level_quant_P = two_level_quant_P ctx.fake_quant_P = fake_quant_P ctx.smooth_q = smooth_q + ctx.use_global_sf = use_global_sf return o @staticmethod @@ -1047,6 +1060,7 @@ class _attention(torch.autograd.Function): two_level_quant_P=ctx.two_level_quant_P, fake_quant_P=ctx.fake_quant_P, SMOOTH_Q=ctx.smooth_q, + use_global_sf=ctx.use_global_sf, num_warps=NUM_WARPS, num_stages=NUM_STAGES ) @@ -1074,12 +1088,13 @@ class _attention(torch.autograd.Function): IS_QAT=ctx.IS_QAT, two_level_quant_P=ctx.two_level_quant_P, fake_quant_P=ctx.fake_quant_P, + SMOOTH_Q=ctx.smooth_q, + use_global_sf=ctx.use_global_sf, num_warps=NUM_WARPS, num_stages=NUM_STAGES, - SMOOTH_Q=ctx.smooth_q, ) - return dq, dk, dv, None, None, None, None, None, None, None, None, None + return dq, dk, dv, None, None, None, None, None, None, None, None, None, None attention = _attention.apply \ No newline at end of file diff --git a/quant_utils.py b/quant_utils.py index ee5696a14..8b824f4b3 100644 --- a/quant_utils.py +++ b/quant_utils.py @@ -6,11 +6,12 @@ from nvfp4_utils import _compute_quant_and_scale, _compute_dequant def fake_quantize(src_tensor, valid_src_mask, BLOCK_SIZE_OUT_DIM: tl.constexpr, BLOCK_SIZE_QUANT_DIM: tl.constexpr, dst_dtype: tl.constexpr, + use_global_sf: tl.constexpr = True, two_level_quant_P: tl.constexpr = False): high_prec_src_tensor = src_tensor src_tensor, src_scale, src_s_dec = _compute_quant_and_scale(src_tensor=src_tensor, valid_src_mask=valid_src_mask, - use_global_sf=True, + use_global_sf=use_global_sf, two_level_quant_P=two_level_quant_P) src_tensor = _compute_dequant(mx_tensor=src_tensor, scale=src_scale, @@ -27,7 +28,8 @@ def fake_quantize_q(Q, fake_Q, stride_z_q, stride_h_q, fake_stride_tok_q, fake_stride_d_q, H, N_CTX_Q, BLOCK_M: tl.constexpr, - HEAD_DIM: tl.constexpr): + HEAD_DIM: tl.constexpr, + use_global_sf: tl.constexpr = True): bhid = tl.program_id(1) adj_q = (stride_h_q * (bhid % H) + stride_z_q * (bhid // H)) fake_adj_q = (fake_stride_h_q * (bhid % H) + fake_stride_z_q * (bhid // H)) @@ -41,7 +43,7 @@ def fake_quantize_q(Q, fake_Q, stride_z_q, stride_h_q, q_valid = offs_m < N_CTX_Q q = tl.load(Q + offs_m[:, None] * stride_tok_q + offs_k[None, :] * stride_d_q, mask=q_valid[:, None], other=0.0) - q, _ = fake_quantize(src_tensor=q, valid_src_mask=q_valid[:, None], BLOCK_SIZE_OUT_DIM=BLOCK_M, BLOCK_SIZE_QUANT_DIM=HEAD_DIM, dst_dtype=q.dtype) + q, _ = fake_quantize(src_tensor=q, valid_src_mask=q_valid[:, None], BLOCK_SIZE_OUT_DIM=BLOCK_M, BLOCK_SIZE_QUANT_DIM=HEAD_DIM, dst_dtype=q.dtype, use_global_sf=use_global_sf) tl.store(fake_Q + offs_m[:, None] * fake_stride_tok_q + offs_k[None, :] * fake_stride_d_q, q, mask=q_valid[:, None]) @triton.jit @@ -51,7 +53,8 @@ def fake_quantize_kv(K, V, fake_K, fake_V, stride_z_kv, stride_h_kv, fake_stride_tok_kv, fake_stride_d_kv, H, N_CTX_KV, BLOCK_N: tl.constexpr, - HEAD_DIM: tl.constexpr): + HEAD_DIM: tl.constexpr, + use_global_sf: tl.constexpr = True): bhid = tl.program_id(1) adj_kv = (stride_h_kv * (bhid % H) + stride_z_kv * (bhid // H)) fake_adj_kv = (fake_stride_h_kv * (bhid % H) + fake_stride_z_kv * (bhid // H)) @@ -68,7 +71,7 @@ def fake_quantize_kv(K, V, fake_K, fake_V, stride_z_kv, stride_h_kv, kv_valid = offs_n < N_CTX_KV k_block = tl.load(K + offs_n[:, None] * stride_tok_kv + offs_k[None, :] * stride_d_kv, mask=kv_valid[:, None], other=0.0) v_block = tl.load(V + offs_n[:, None] * stride_tok_kv + offs_k[None, :] * stride_d_kv, mask=kv_valid[:, None], other=0.0) - k, _ = fake_quantize(src_tensor=k_block, valid_src_mask=kv_valid[:, None], BLOCK_SIZE_OUT_DIM=BLOCK_N, BLOCK_SIZE_QUANT_DIM=HEAD_DIM, dst_dtype=k_block.dtype) - v, _ = fake_quantize(src_tensor=v_block, valid_src_mask=kv_valid[:, None], BLOCK_SIZE_OUT_DIM=BLOCK_N, BLOCK_SIZE_QUANT_DIM=HEAD_DIM, dst_dtype=v_block.dtype) + k, _ = fake_quantize(src_tensor=k_block, valid_src_mask=kv_valid[:, None], BLOCK_SIZE_OUT_DIM=BLOCK_N, BLOCK_SIZE_QUANT_DIM=HEAD_DIM, dst_dtype=k_block.dtype, use_global_sf=use_global_sf) + v, _ = fake_quantize(src_tensor=v_block, valid_src_mask=kv_valid[:, None], BLOCK_SIZE_OUT_DIM=BLOCK_N, BLOCK_SIZE_QUANT_DIM=HEAD_DIM, dst_dtype=v_block.dtype, use_global_sf=use_global_sf) tl.store(fake_K + offs_n[:, None] * fake_stride_tok_kv + offs_k[None, :] * fake_stride_d_kv, k, mask=kv_valid[:, None]) tl.store(fake_V + offs_n[:, None] * fake_stride_tok_kv + offs_k[None, :] * fake_stride_d_kv, v, mask=kv_valid[:, None])