nodes/hao-ai-lab-FastVideo
1
0
Fork 0
Code Issues Pull Requests Packages Projects Releases Wiki Activity

try sage3 finetuning with aligned backward pass

Browse Source
This commit is contained in:
RandNMR73
2026-01-07 01:37:09 +00:00
parent 30eeec0411
commit 3adfb07465
4 changed files with 99 additions and 33 deletions
Download Patch File Download Diff File Expand all files Collapse all files
examples/training/finetune/wan_t2v_1.3B/crush_smol/finetune_t2v_sage3_5090.sh
+1 -1
View File
@@ -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"
fastvideo/attention/backends/sage_attn3.py
+60 -12
View File
@@ -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()
qat_attn.py
+29 -14
View File
@@ -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
quant_utils.py
+9 -6
View File
@@ -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])