diff --git a/turbodiffusion/ops/core.py b/turbodiffusion/ops/core.py index 6d74fe5..37527dd 100644 --- a/turbodiffusion/ops/core.py +++ b/turbodiffusion/ops/core.py @@ -377,6 +377,101 @@ def layernorm_noparam(x, eps): return y +@triton.jit +def _layer_norm_modulate_noparam_fwd_fused( + X, + Y, + Mean, + Rstd, + Scale, + Shift, + x_stride, + y_stride, + scale_b_stride, + shift_b_stride, + M: tl.constexpr, + L: tl.constexpr, + N: tl.constexpr, + N2: tl.constexpr, + eps, + BLOCK_M: tl.constexpr, +): + pid = tl.program_id(0) + rows = pid * BLOCK_M + tl.arange(0, BLOCK_M) + cols = tl.arange(0, N2) + row_mask = rows < M + col_mask = cols < N + mask = row_mask[:, None] & col_mask[None, :] + + x_ptr = X + rows[:, None] * x_stride + cols[None, :] + y_ptr = Y + rows[:, None] * y_stride + cols[None, :] + + x = tl.load(x_ptr, mask=mask, other=0.0).to(tl.float32) + mean = tl.sum(x, axis=1, keep_dims=True) / N + var = tl.sum((x - mean) * (x - mean), axis=1, keep_dims=True) / N + rstd = 1 / tl.sqrt(var + eps) + + _mean = tl.reshape(mean, (BLOCK_M)) + _rstd = tl.reshape(rstd, (BLOCK_M)) + tl.store(Mean + rows, _mean, mask=row_mask) + tl.store(Rstd + rows, _rstd, mask=row_mask) + + batch = rows // L + scale = tl.load(Scale + batch[:, None] * scale_b_stride + cols[None, :], mask=mask, other=0.0).to(tl.float32) + shift = tl.load(Shift + batch[:, None] * shift_b_stride + cols[None, :], mask=mask, other=0.0).to(tl.float32) + + x_hat = (x - mean) * rstd + x_hat = x_hat.to(X.type.element_ty).to(tl.float32) + y = x_hat * (1.0 + scale) + shift + tl.store(y_ptr, y.to(Y.type.element_ty), mask=mask) + + +def layernorm_modulate_noparam(x, scale, shift, eps, output_dtype=None): + assert x.dim() == 3, "FastLayerNorm.modulate expects [B, L, C] input" + if not x.is_contiguous(): + x = x.contiguous() + if scale.stride(-1) != 1: + scale = scale.contiguous() + if shift.stride(-1) != 1: + shift = shift.contiguous() + + B, L, N = x.shape + M = B * L + output_dtype = output_dtype or torch.promote_types(torch.promote_types(x.dtype, scale.dtype), shift.dtype) + y = torch.empty_like(x, dtype=output_dtype) + mean = torch.empty((M,), dtype=torch.float32, device=x.device) + rstd = torch.empty((M,), dtype=torch.float32, device=x.device) + + num_warps = 8 + N2 = triton.next_power_of_2(N) + BLOCK_M = 32 if N <= 512 else 1 + + scale_b_stride = scale.stride(0) if scale.dim() > 1 else 0 + shift_b_stride = shift.stride(0) if shift.dim() > 1 else 0 + + _layer_norm_modulate_noparam_fwd_fused[(triton.cdiv(M, BLOCK_M),)]( + x, + y, + mean, + rstd, + scale, + shift, + x.stride(1), + y.stride(1), + scale_b_stride, + shift_b_stride, + M, + L, + N, + N2, + eps, + num_warps=num_warps, + BLOCK_M=BLOCK_M, + ) + + return y + + def layernorm(x, w, b, eps, elementwise_affine=True): if elementwise_affine: assert w is not None and b is not None @@ -476,6 +571,13 @@ def __init__( def forward(self, x): return layernorm(x.float(), self.weight, self.bias, self.eps, self.elementwise_affine).to(x.dtype) + + def modulate(self, x, scale, shift, out_dtype=None): + output_dtype = out_dtype or torch.promote_types(torch.promote_types(x.dtype, scale.dtype), shift.dtype) + if self.elementwise_affine: + y = self.forward(x).float() * (1 + scale) + shift + return y.to(output_dtype) + return layernorm_modulate_noparam(x, scale, shift, self.eps, output_dtype=output_dtype) @classmethod def from_layernorm(cls, original_layernorm): diff --git a/turbodiffusion/rcm/networks/wan2pt1.py b/turbodiffusion/rcm/networks/wan2pt1.py index 56ee646..4342cac 100644 --- a/turbodiffusion/rcm/networks/wan2pt1.py +++ b/turbodiffusion/rcm/networks/wan2pt1.py @@ -169,11 +169,11 @@ def rope_apply(x, freqs): # freqs is already sharded to local seq_len under flattened CP freqs = freqs.view(seq_len, head_dim // 2) - cos = torch.cos(freqs).to(torch.float32) - sin = torch.sin(freqs).to(torch.float32) + cos = torch.cos(freqs) + sin = torch.sin(freqs) # Apply the rotation - rotated = flash_apply_rotary_emb(x.to(torch.float32), cos, sin, interleaved=True, inplace=False) + rotated = flash_apply_rotary_emb(x, cos, sin, interleaved=True, inplace=False) return rotated.to(x.dtype) @@ -212,6 +212,14 @@ def forward(self, x): return super().forward(x.float()).type_as(x) +def _modulate_norm(norm, x, scale, shift, out_dtype=None): + modulate = getattr(norm, "modulate", None) + if modulate is not None: + return modulate(x, scale, shift, out_dtype=out_dtype) + y = norm(x).float() * (1 + scale) + shift + return y if out_dtype is None else y.to(out_dtype) + + class WanSelfAttention(nn.Module): def __init__(self, dim, num_heads, qk_norm=True, eps=1e-6): assert dim % num_heads == 0 @@ -401,14 +409,14 @@ def forward(self, x, e, seq_lens, freqs, context, context_lens): assert e[0].dtype == torch.float32 # self-attention - y = self.self_attn((self.norm1(x).float() * (1 + e[1]) + e[0]).type_as(x), seq_lens, freqs) + y = self.self_attn(_modulate_norm(self.norm1, x, e[1], e[0], out_dtype=x.dtype), seq_lens, freqs) with amp.autocast("cuda", dtype=torch.float32): x = x + y * e[2].type_as(x) # cross-attention & ffn function def cross_attn_ffn(x, context, context_lens, e): x = x + self.cross_attn(self.norm3(x), context, context_lens) - y = self.ffn((self.norm2(x).float() * (1 + e[4]) + e[3]).type_as(x)) + y = self.ffn(_modulate_norm(self.norm2, x, e[4], e[3], out_dtype=x.dtype)) with amp.autocast("cuda", dtype=torch.float32): x = x + y * e[5].type_as(x) return x @@ -450,7 +458,7 @@ def forward(self, x, e): assert e.dtype == torch.float32 with amp.autocast("cuda", dtype=torch.float32): e = (self.modulation + e.unsqueeze(1)).chunk(2, dim=1) - x = self.head(self.norm(x) * (1 + e[1]) + e[0]) + x = self.head(_modulate_norm(self.norm, x, e[1], e[0])) return x diff --git a/turbodiffusion/rcm/networks/wan2pt2.py b/turbodiffusion/rcm/networks/wan2pt2.py index 6463110..e0b1056 100644 --- a/turbodiffusion/rcm/networks/wan2pt2.py +++ b/turbodiffusion/rcm/networks/wan2pt2.py @@ -212,6 +212,14 @@ def forward(self, x): return super().forward(x.float()).type_as(x) +def _modulate_norm(norm, x, scale, shift, out_dtype=None): + modulate = getattr(norm, "modulate", None) + if modulate is not None: + return modulate(x, scale, shift, out_dtype=out_dtype) + y = norm(x).float() * (1 + scale) + shift + return y if out_dtype is None else y.to(out_dtype) + + class WanSelfAttention(nn.Module): def __init__(self, dim, num_heads, qk_norm=True, eps=1e-6): assert dim % num_heads == 0 @@ -351,14 +359,14 @@ def forward(self, x, e, seq_lens, freqs, context, context_lens): assert e[0].dtype == torch.float32 # self-attention - y = self.self_attn((self.norm1(x).float() * (1 + e[1]) + e[0]).type_as(x), seq_lens, freqs) + y = self.self_attn(_modulate_norm(self.norm1, x, e[1], e[0], out_dtype=x.dtype), seq_lens, freqs) with amp.autocast("cuda", dtype=torch.float32): x = x + y * e[2].type_as(x) # cross-attention & ffn function def cross_attn_ffn(x, context, context_lens, e): x = x + self.cross_attn(self.norm3(x), context, context_lens) - y = self.ffn((self.norm2(x).float() * (1 + e[4]) + e[3]).type_as(x)) + y = self.ffn(_modulate_norm(self.norm2, x, e[4], e[3], out_dtype=x.dtype)) with amp.autocast("cuda", dtype=torch.float32): x = x + y * e[5].type_as(x) return x @@ -400,7 +408,7 @@ def forward(self, x, e): assert e.dtype == torch.float32 with amp.autocast("cuda", dtype=torch.float32): e = (self.modulation + e.unsqueeze(1)).chunk(2, dim=1) - x = self.head(self.norm(x) * (1 + e[1]) + e[0]) + x = self.head(_modulate_norm(self.norm, x, e[1], e[0])) return x