Fuse LayerNorm modulation into Triton kernel and remove RoPE dtype casts

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):

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
 
 

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