[perf] Wan kernel fusion: Triton-fused residual+LayerNorm+modulate inference path (#1708)
Co-authored-by: SolitaryThinker <wlsaidhi@gmail.com> Co-authored-by: Claude Opus 5.5 (1M context) <noreply@anthropic.com>
This commit is contained in:
Haochen Jiang
co-authored by
SolitaryThinker
Claude Opus 5.5
SolitaryThinker
Claude Opus 5.5
parent
bc0191d775
commit
e722023bd6
@@ -28,6 +28,7 @@ exec pytest \
|
||||
./fastvideo/tests/inference/test_basic_fasth3_omniref_pdd.py \
|
||||
./fastvideo/tests/attention/ \
|
||||
./fastvideo/tests/layers/test_pdd_linear.py \
|
||||
./fastvideo/tests/layers/test_triton_fused_norm.py \
|
||||
./fastvideo/tests/modal/test_kernel_build_cache.py \
|
||||
./fastvideo/tests/modal/test_pr_test.py \
|
||||
./fastvideo/tests/modal/test_ssim_test.py \
|
||||
|
||||
@@ -214,6 +214,7 @@ longer exists also fails the test, so the fixing pull request deletes its entry.
|
||||
| `FASTVIDEO_LTX2_VAE_CHANNELS_LAST_3D` | bool | `1` | performance | LTX-2 VAE uses the channels_last_3d memory format. |
|
||||
| `FASTVIDEO_LTX2_DISABLE_AUDIO_AUTOCAST` | bool | `1` | performance | LTX-2 audio decoding runs without CUDA autocast. Deprecated names: `LTX2_DISABLE_AUDIO_AUTOCAST`. |
|
||||
| `FASTVIDEO_FLUX2_DISABLE_BF16_REDUCED_PRECISION_REDUCTION` | bool | `0` | performance | Flux denoising disables reduced-precision reductions in bf16 matmuls, which tightens accumulation for the 4-step Klein model. |
|
||||
| `FASTVIDEO_DISABLE_FUSED_NORM` | bool | `0` | performance | Turn off the Triton-fused residual + LayerNorm + modulate inference path that Wan blocks opt into, and use the eager path. |
|
||||
| `FASTVIDEO_FFMPEG_BIN` | str | `ffmpeg` | output | ffmpeg executable used to save video with audio. |
|
||||
| `FASTVIDEO_VIDEO_CODEC` | str | `libx264` | output | ffmpeg video codec for saved videos. |
|
||||
| `FASTVIDEO_NVENC_PRESET` | str | `p1` | output | NVENC preset when the codec is an \*_nvenc codec. |
|
||||
|
||||
@@ -436,6 +436,12 @@ FASTVIDEO_FLUX2_DISABLE_BF16_REDUCED_PRECISION_REDUCTION = EnvBool(
|
||||
category="performance",
|
||||
doc="Flux denoising disables reduced-precision reductions in bf16 matmuls, which tightens accumulation for "
|
||||
"the 4-step Klein model.")
|
||||
# Global kill switch for the opt-in, inference-only Triton fusion in
|
||||
# fastvideo/layers/triton_fused_norm.py, for debugging numerics or Triton issues.
|
||||
FASTVIDEO_DISABLE_FUSED_NORM = EnvBool(False,
|
||||
category="performance",
|
||||
doc="Turn off the Triton-fused residual + LayerNorm + modulate inference "
|
||||
"path that Wan blocks opt into, and use the eager path.")
|
||||
|
||||
# ================== Output encoding ==================
|
||||
|
||||
|
||||
@@ -144,8 +144,15 @@ class ScaleResidualLayerNormScaleShift(nn.Module):
|
||||
dtype: torch.dtype = torch.float32,
|
||||
compute_dtype: torch.dtype | None = None,
|
||||
prefix: str = "",
|
||||
fuse_inference: bool = False,
|
||||
):
|
||||
super().__init__()
|
||||
# Opt-in Triton fusion for the no-grad path. Off by default: the fused
|
||||
# kernel returns both outputs in the *stream* dtype (so caller-side
|
||||
# .to(orig_dtype) casts become no-ops), which changes the return dtype
|
||||
# contract relative to eager type promotion. Only callers that follow
|
||||
# every call with a cast to the stream dtype (e.g. Wan) should enable.
|
||||
self.fuse_inference = fuse_inference
|
||||
if norm_type == "rms":
|
||||
self.norm = RMSNorm(hidden_size, has_weight=elementwise_affine, eps=eps, dtype=dtype)
|
||||
elif norm_type == "layer":
|
||||
@@ -160,19 +167,32 @@ class ScaleResidualLayerNormScaleShift(nn.Module):
|
||||
residual: torch.Tensor,
|
||||
x: torch.Tensor,
|
||||
gate: torch.Tensor | int,
|
||||
shift: torch.Tensor,
|
||||
scale: torch.Tensor,
|
||||
shift: torch.Tensor | None,
|
||||
scale: torch.Tensor | None,
|
||||
convert_modulation_dtype: bool = False) -> tuple[torch.Tensor, torch.Tensor]:
|
||||
"""
|
||||
Apply gated residual connection, followed by layernorm and
|
||||
Apply gated residual connection, followed by layernorm and
|
||||
scale/shift in a single fused operation.
|
||||
|
||||
|
||||
Pass ``shift=None, scale=None`` (both, not just one) to skip the
|
||||
modulation entirely and return the normalized tensor as-is. Callers
|
||||
that need norm-only behavior (e.g. Wan's post-self-attn site) would
|
||||
otherwise pay two full-tensor elementwise passes computing
|
||||
``normalized * (1.0 + 0) + 0``.
|
||||
|
||||
Returns:
|
||||
Tuple containing:
|
||||
- normalized and modulated output
|
||||
- residual value (value after residual connection
|
||||
- residual value (value after residual connection
|
||||
but before normalization)
|
||||
"""
|
||||
if (shift is None) != (scale is None):
|
||||
raise ValueError("shift and scale must be None together or both be tensors")
|
||||
|
||||
if self.fuse_inference:
|
||||
from fastvideo.layers.triton_fused_norm import (fused_path_supported, fused_residual_norm_mod)
|
||||
if fused_path_supported(residual, x, gate, shift, scale, self.norm):
|
||||
return fused_residual_norm_mod(residual, x, gate, shift, scale, self.norm)
|
||||
# x.shape: [batch_size, seq_len, inner_dim]
|
||||
# Apply residual connection with gating
|
||||
if isinstance(gate, int):
|
||||
@@ -196,6 +216,10 @@ class ScaleResidualLayerNormScaleShift(nn.Module):
|
||||
# Apply normalization
|
||||
normalized = self.norm(residual_output)
|
||||
|
||||
# Identity modulation: skip the scale/shift passes entirely.
|
||||
if shift is None or scale is None:
|
||||
return normalized, residual_output
|
||||
|
||||
if convert_modulation_dtype:
|
||||
scale = scale.to(normalized.dtype)
|
||||
shift = shift.to(normalized.dtype)
|
||||
|
||||
@@ -0,0 +1,288 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
"""Triton-fused residual + LayerNorm + scale/shift for inference.
|
||||
|
||||
Collapses the eager chain
|
||||
|
||||
residual_out = residual + x * gate (or residual + x)
|
||||
normalized = FP32LayerNorm(residual_out)
|
||||
modulated = normalized * (1 + scale) + shift (optional)
|
||||
... followed by the caller's .to(orig_dtype) casts ...
|
||||
|
||||
into a single kernel with two outputs, both already in the stream dtype, so the
|
||||
caller-side casts become no-ops. All arithmetic is fp32 regardless of I/O dtype.
|
||||
|
||||
Numerics: for the two Wan input classes the eager path's rounding points are
|
||||
replicated exactly. When the gate is a fp32 tensor, eager type promotion keeps
|
||||
the whole chain in fp32 with one final round -- the kernel does the same. When
|
||||
the gate is the scalar 1 and the stream is bf16, eager materializes bf16
|
||||
intermediates (the residual sum and the norm output), so the kernel round-trips
|
||||
through bf16 at the same two points. A bf16 *tensor* gate with a bf16 stream is
|
||||
only half-modelled: eager rounds the product ``x * gate`` to bf16 before the
|
||||
add, while the kernel keeps the product in fp32 and rounds the sum once, so the
|
||||
fused residual can differ from eager by up to one bf16 ulp of the product (the
|
||||
fused value is slightly more accurate, not wrong). The remaining differences
|
||||
from eager are the reduction order inside mean/variance (last-ulp fp32) and
|
||||
that product-rounding case.
|
||||
|
||||
Inference-only: callers must gate on ``torch.is_grad_enabled()``; there is no
|
||||
backward. Disable globally with FASTVIDEO_DISABLE_FUSED_NORM=1.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import torch
|
||||
|
||||
import fastvideo.envs as envs
|
||||
|
||||
try:
|
||||
import triton
|
||||
import triton.language as tl
|
||||
_HAS_TRITON = True
|
||||
except ImportError: # pragma: no cover - CPU-only installs
|
||||
_HAS_TRITON = False
|
||||
|
||||
_SUPPORTED_STREAM_DTYPES = (torch.bfloat16, torch.float32)
|
||||
_MAX_HIDDEN = 8192
|
||||
|
||||
if _HAS_TRITON:
|
||||
|
||||
@triton.jit
|
||||
def _fused_residual_norm_mod_kernel(
|
||||
X,
|
||||
RES,
|
||||
GATE,
|
||||
W,
|
||||
B,
|
||||
SHIFT,
|
||||
SCALE,
|
||||
OUT,
|
||||
RES_OUT,
|
||||
S,
|
||||
H,
|
||||
x_sb,
|
||||
x_ss,
|
||||
r_sb,
|
||||
r_ss,
|
||||
g_sb,
|
||||
g_ss,
|
||||
sh_sb,
|
||||
sh_ss,
|
||||
sc_sb,
|
||||
sc_ss,
|
||||
o_sb,
|
||||
o_ss,
|
||||
ro_sb,
|
||||
ro_ss,
|
||||
eps,
|
||||
HAS_GATE_TENSOR: tl.constexpr,
|
||||
INTERMEDIATE_ROUND: tl.constexpr,
|
||||
HAS_AFFINE: tl.constexpr,
|
||||
HAS_BIAS: tl.constexpr,
|
||||
HAS_MOD: tl.constexpr,
|
||||
OUT_BF16: tl.constexpr,
|
||||
BLOCK_H: tl.constexpr,
|
||||
):
|
||||
row = tl.program_id(0)
|
||||
b = row // S
|
||||
s = row % S
|
||||
cols = tl.arange(0, BLOCK_H)
|
||||
mask = cols < H
|
||||
|
||||
x = tl.load(X + b * x_sb + s * x_ss + cols, mask=mask, other=0.0).to(tl.float32)
|
||||
r = tl.load(RES + b * r_sb + s * r_ss + cols, mask=mask, other=0.0).to(tl.float32)
|
||||
if HAS_GATE_TENSOR:
|
||||
g = tl.load(GATE + b * g_sb + s * g_ss + cols, mask=mask, other=0.0).to(tl.float32)
|
||||
acc = r + x * g
|
||||
else:
|
||||
acc = r + x
|
||||
|
||||
if OUT_BF16:
|
||||
res_bf16 = acc.to(tl.bfloat16)
|
||||
tl.store(RES_OUT + b * ro_sb + s * ro_ss + cols, res_bf16, mask=mask)
|
||||
if INTERMEDIATE_ROUND:
|
||||
# eager produced a bf16 residual tensor here; the norm consumed
|
||||
# the rounded values.
|
||||
acc = res_bf16.to(tl.float32)
|
||||
else:
|
||||
tl.store(RES_OUT + b * ro_sb + s * ro_ss + cols, acc, mask=mask)
|
||||
|
||||
mean = tl.sum(tl.where(mask, acc, 0.0), axis=0) / H
|
||||
diff = tl.where(mask, acc - mean, 0.0)
|
||||
var = tl.sum(diff * diff, axis=0) / H
|
||||
xhat = diff * tl.rsqrt(var + eps)
|
||||
if HAS_AFFINE:
|
||||
w = tl.load(W + cols, mask=mask, other=0.0).to(tl.float32)
|
||||
xhat = xhat * w
|
||||
if HAS_BIAS:
|
||||
bias = tl.load(B + cols, mask=mask, other=0.0).to(tl.float32)
|
||||
xhat = xhat + bias
|
||||
|
||||
if HAS_MOD:
|
||||
if INTERMEDIATE_ROUND and OUT_BF16:
|
||||
# eager rounded the norm output to bf16 before modulating.
|
||||
xhat = xhat.to(tl.bfloat16).to(tl.float32)
|
||||
sc = tl.load(SCALE + b * sc_sb + s * sc_ss + cols, mask=mask, other=0.0).to(tl.float32)
|
||||
sh = tl.load(SHIFT + b * sh_sb + s * sh_ss + cols, mask=mask, other=0.0).to(tl.float32)
|
||||
y = xhat * (1.0 + sc) + sh
|
||||
else:
|
||||
y = xhat
|
||||
|
||||
if OUT_BF16:
|
||||
tl.store(OUT + b * o_sb + s * o_ss + cols, y.to(tl.bfloat16), mask=mask)
|
||||
else:
|
||||
tl.store(OUT + b * o_sb + s * o_ss + cols, y, mask=mask)
|
||||
|
||||
|
||||
def _broadcast_strides(t: torch.Tensor, batch: int, seq: int, hidden: int) -> tuple[int, int] | None:
|
||||
"""(batch_stride, seq_stride) for a [B|1, S|1, H] tensor, or None if unsupported.
|
||||
|
||||
The last dim must equal ``hidden`` exactly: the kernel indexes columns
|
||||
``0..hidden-1`` of every modulation tensor, so a narrower (which eager would
|
||||
broadcast) or wider last dim would silently read the wrong elements instead
|
||||
of raising.
|
||||
"""
|
||||
if t.dim() != 3 or t.stride(-1) != 1 or t.shape[-1] != hidden:
|
||||
return None
|
||||
tb, ts, _ = t.shape
|
||||
if tb not in (1, batch) or ts not in (1, seq):
|
||||
return None
|
||||
return (t.stride(0) if tb == batch else 0, t.stride(1) if ts == seq else 0)
|
||||
|
||||
|
||||
def fused_path_supported(
|
||||
residual: torch.Tensor,
|
||||
x: torch.Tensor,
|
||||
gate: torch.Tensor | int,
|
||||
shift: torch.Tensor | None,
|
||||
scale: torch.Tensor | None,
|
||||
norm: torch.nn.Module,
|
||||
) -> bool:
|
||||
"""Cheap eligibility check; any False falls back to the eager path."""
|
||||
if not _HAS_TRITON or envs.FASTVIDEO_DISABLE_FUSED_NORM.get():
|
||||
return False
|
||||
if torch.is_grad_enabled():
|
||||
return False
|
||||
if torch.compiler.is_compiling():
|
||||
# Dynamo must not trace into the Triton launcher (see the custom-op
|
||||
# boundaries in fastvideo/models/dits/minimax_h3_fusions/modulation.py).
|
||||
# Until this kernel is wrapped in torch.library.custom_op + register_fake
|
||||
# the way that module does, compiled forwards take the eager path, so
|
||||
# enable_torch_compile keeps its pre-fusion behavior instead of hitting
|
||||
# an unregistered launch inside the compiled graph.
|
||||
return False
|
||||
# LayerNorm family only (RMS keeps its own path); weight/bias fp32 as
|
||||
# FP32LayerNorm guarantees after .float().
|
||||
if not isinstance(norm, torch.nn.LayerNorm):
|
||||
return False
|
||||
if residual.dim() != 3 or residual.shape != x.shape:
|
||||
return False
|
||||
if residual.dtype != x.dtype or residual.dtype not in _SUPPORTED_STREAM_DTYPES:
|
||||
return False
|
||||
if not (residual.is_cuda and x.is_cuda):
|
||||
return False
|
||||
if not (residual.is_contiguous() and x.is_contiguous()):
|
||||
return False
|
||||
batch, seq, hidden = residual.shape
|
||||
if hidden > _MAX_HIDDEN or norm.normalized_shape != (hidden, ):
|
||||
return False
|
||||
if isinstance(gate, torch.Tensor):
|
||||
if gate.dtype not in (torch.float32, residual.dtype) or not gate.is_cuda:
|
||||
return False
|
||||
if _broadcast_strides(gate, batch, seq, hidden) is None:
|
||||
return False
|
||||
elif gate != 1:
|
||||
return False
|
||||
if (shift is None) != (scale is None):
|
||||
return False
|
||||
if shift is not None and scale is not None:
|
||||
for t in (shift, scale):
|
||||
if t.dtype not in (torch.float32, residual.dtype) or not t.is_cuda:
|
||||
return False
|
||||
if _broadcast_strides(t, batch, seq, hidden) is None:
|
||||
return False
|
||||
return True
|
||||
|
||||
|
||||
def fused_residual_norm_mod(
|
||||
residual: torch.Tensor,
|
||||
x: torch.Tensor,
|
||||
gate: torch.Tensor | int,
|
||||
shift: torch.Tensor | None,
|
||||
scale: torch.Tensor | None,
|
||||
norm: torch.nn.Module,
|
||||
) -> tuple[torch.Tensor, torch.Tensor]:
|
||||
"""Run the fused kernel. Caller must have checked fused_path_supported().
|
||||
|
||||
Returns (modulated_or_normalized, residual_out), both in the stream dtype.
|
||||
"""
|
||||
batch, seq, hidden = residual.shape
|
||||
stream_dtype = residual.dtype
|
||||
out = torch.empty_like(residual)
|
||||
res_out = torch.empty_like(residual)
|
||||
|
||||
gate_is_tensor = isinstance(gate, torch.Tensor)
|
||||
has_mod = shift is not None
|
||||
# Eager promotion keeps the chain fp32 (no intermediate rounding) exactly
|
||||
# when the gate is a fp32 tensor; the scalar-1 gate path materializes
|
||||
# stream-dtype intermediates.
|
||||
intermediate_round = (not gate_is_tensor or gate.dtype == stream_dtype) and stream_dtype == torch.bfloat16
|
||||
|
||||
if gate_is_tensor:
|
||||
g_sb, g_ss = _broadcast_strides(gate, batch, seq, hidden) # type: ignore[misc]
|
||||
gate_arg = gate
|
||||
else:
|
||||
g_sb = g_ss = 0
|
||||
gate_arg = residual # unused placeholder pointer
|
||||
if has_mod:
|
||||
sh_sb, sh_ss = _broadcast_strides(shift, batch, seq, hidden) # type: ignore[arg-type,misc]
|
||||
sc_sb, sc_ss = _broadcast_strides(scale, batch, seq, hidden) # type: ignore[arg-type,misc]
|
||||
shift_arg, scale_arg = shift, scale
|
||||
else:
|
||||
sh_sb = sh_ss = sc_sb = sc_ss = 0
|
||||
shift_arg = scale_arg = residual # unused placeholder pointers
|
||||
|
||||
weight = norm.weight
|
||||
bias = norm.bias
|
||||
has_affine = weight is not None
|
||||
has_bias = bias is not None
|
||||
|
||||
block_h = triton.next_power_of_2(hidden)
|
||||
num_warps = 4 if block_h <= 2048 else 8
|
||||
|
||||
_fused_residual_norm_mod_kernel[(batch * seq, )](
|
||||
x,
|
||||
residual,
|
||||
gate_arg,
|
||||
weight if has_affine else residual,
|
||||
bias if has_bias else residual,
|
||||
shift_arg,
|
||||
scale_arg,
|
||||
out,
|
||||
res_out,
|
||||
seq,
|
||||
hidden,
|
||||
x.stride(0),
|
||||
x.stride(1),
|
||||
residual.stride(0),
|
||||
residual.stride(1),
|
||||
g_sb,
|
||||
g_ss,
|
||||
sh_sb,
|
||||
sh_ss,
|
||||
sc_sb,
|
||||
sc_ss,
|
||||
out.stride(0),
|
||||
out.stride(1),
|
||||
res_out.stride(0),
|
||||
res_out.stride(1),
|
||||
norm.eps,
|
||||
HAS_GATE_TENSOR=gate_is_tensor,
|
||||
INTERMEDIATE_ROUND=intermediate_round,
|
||||
HAS_AFFINE=has_affine,
|
||||
HAS_BIAS=has_bias,
|
||||
HAS_MOD=has_mod,
|
||||
OUT_BF16=stream_dtype == torch.bfloat16,
|
||||
BLOCK_H=block_h,
|
||||
num_warps=num_warps,
|
||||
)
|
||||
return out, res_out
|
||||
@@ -326,7 +326,8 @@ class WanTransformerBlock(nn.Module):
|
||||
eps=eps,
|
||||
elementwise_affine=True,
|
||||
dtype=torch.float32,
|
||||
compute_dtype=torch.float32)
|
||||
compute_dtype=torch.float32,
|
||||
fuse_inference=True)
|
||||
|
||||
# 2. Cross-attention
|
||||
if added_kv_proj_dim is not None:
|
||||
@@ -350,7 +351,8 @@ class WanTransformerBlock(nn.Module):
|
||||
eps=eps,
|
||||
elementwise_affine=False,
|
||||
dtype=torch.float32,
|
||||
compute_dtype=torch.float32)
|
||||
compute_dtype=torch.float32,
|
||||
fuse_inference=True)
|
||||
|
||||
# 3. Feed-forward
|
||||
self.ffn = MLP(dim, ffn_dim, act_type="gelu_pytorch_tanh", quant_config=quant_config, prefix=f"{prefix}.ffn")
|
||||
@@ -415,9 +417,10 @@ class WanTransformerBlock(nn.Module):
|
||||
attn_output, _ = self.to_out(attn_output)
|
||||
attn_output = attn_output.squeeze(1)
|
||||
|
||||
null_shift = null_scale = torch.tensor([0], device=hidden_states.device)
|
||||
norm_hidden_states, hidden_states = self.self_attn_residual_norm(hidden_states, attn_output, gate_msa,
|
||||
null_shift, null_scale)
|
||||
# Norm-only site: shift/scale=None skips the identity modulation
|
||||
# (normalized * 1.0 + 0) that two full-tensor passes used to compute.
|
||||
norm_hidden_states, hidden_states = self.self_attn_residual_norm(hidden_states, attn_output, gate_msa, None,
|
||||
None)
|
||||
norm_hidden_states, hidden_states = norm_hidden_states.to(orig_dtype), hidden_states.to(orig_dtype)
|
||||
|
||||
# 2. Cross-attention
|
||||
@@ -485,7 +488,8 @@ class WanTransformerBlock_VSA(nn.Module):
|
||||
eps=eps,
|
||||
elementwise_affine=True,
|
||||
dtype=torch.float32,
|
||||
compute_dtype=torch.float32)
|
||||
compute_dtype=torch.float32,
|
||||
fuse_inference=True)
|
||||
|
||||
# 2. Cross-attention
|
||||
if added_kv_proj_dim is not None:
|
||||
@@ -509,7 +513,8 @@ class WanTransformerBlock_VSA(nn.Module):
|
||||
eps=eps,
|
||||
elementwise_affine=False,
|
||||
dtype=torch.float32,
|
||||
compute_dtype=torch.float32)
|
||||
compute_dtype=torch.float32,
|
||||
fuse_inference=True)
|
||||
|
||||
# 3. Feed-forward
|
||||
self.ffn = MLP(dim, ffn_dim, act_type="gelu_pytorch_tanh", quant_config=quant_config, prefix=f"{prefix}.ffn")
|
||||
@@ -563,9 +568,10 @@ class WanTransformerBlock_VSA(nn.Module):
|
||||
attn_output, _ = self.to_out(attn_output)
|
||||
attn_output = attn_output.squeeze(1)
|
||||
|
||||
null_shift = null_scale = torch.tensor([0], device=hidden_states.device)
|
||||
norm_hidden_states, hidden_states = self.self_attn_residual_norm(hidden_states, attn_output, gate_msa,
|
||||
null_shift, null_scale)
|
||||
# Norm-only site: shift/scale=None skips the identity modulation
|
||||
# (normalized * 1.0 + 0) that two full-tensor passes used to compute.
|
||||
norm_hidden_states, hidden_states = self.self_attn_residual_norm(hidden_states, attn_output, gate_msa, None,
|
||||
None)
|
||||
norm_hidden_states, hidden_states = norm_hidden_states.to(orig_dtype), hidden_states.to(orig_dtype)
|
||||
|
||||
# 2. Cross-attention
|
||||
|
||||
@@ -0,0 +1,196 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
"""Parity tests for the Triton-fused residual+norm+modulate inference path.
|
||||
|
||||
The fused path claims to replicate the eager rounding sequence exactly, with
|
||||
the layer-norm reduction order as the only remaining difference. These tests
|
||||
pin that claim: outputs must match eager to within a bf16 ulp everywhere, with
|
||||
the overwhelming majority of elements bit-identical.
|
||||
"""
|
||||
import os
|
||||
|
||||
import pytest
|
||||
import torch
|
||||
|
||||
from fastvideo.layers.layernorm import ScaleResidualLayerNormScaleShift
|
||||
|
||||
cuda_only = pytest.mark.skipif(not torch.cuda.is_available(), reason="requires CUDA")
|
||||
|
||||
B, S, H = 1, 9360, 1536 # the Wan2.1-T2V-1.3B sp2 shapes
|
||||
|
||||
|
||||
def _make_modules(elementwise_affine: bool, seed: int = 0):
|
||||
"""One eager module and one fused module with identical weights."""
|
||||
torch.manual_seed(seed)
|
||||
kwargs = dict(
|
||||
norm_type="layer",
|
||||
eps=1e-6,
|
||||
elementwise_affine=elementwise_affine,
|
||||
dtype=torch.float32,
|
||||
compute_dtype=torch.float32,
|
||||
)
|
||||
eager = ScaleResidualLayerNormScaleShift(H, **kwargs).cuda()
|
||||
fused = ScaleResidualLayerNormScaleShift(H, fuse_inference=True, **kwargs).cuda()
|
||||
if elementwise_affine:
|
||||
with torch.no_grad():
|
||||
w = torch.randn(H) * 0.02 + 1.0
|
||||
bias = torch.randn(H) * 0.01
|
||||
for m in (eager, fused):
|
||||
m.norm.weight.copy_(w)
|
||||
m.norm.bias.copy_(bias)
|
||||
return eager, fused
|
||||
|
||||
|
||||
def _inputs(stream_dtype: torch.dtype, seed: int = 1):
|
||||
torch.manual_seed(seed)
|
||||
residual = (torch.randn(B, S, H, device="cuda") * 3).to(stream_dtype)
|
||||
x = torch.randn(B, S, H, device="cuda").to(stream_dtype)
|
||||
return residual, x
|
||||
|
||||
|
||||
def _assert_close_to_eager(fused_out, eager_out, stream_dtype, shift=None):
|
||||
"""Fused must match eager to within the provable error bound.
|
||||
|
||||
The fused kernel differs from eager only in reduction order, rsqrt, and FMA
|
||||
contraction -- all fp32-ulp effects on the *norm output*. After bf16
|
||||
rounding that is at most one bf16 step in xhat (bf16 has a 7-bit
|
||||
mantissa, so one step is 2^-7 of the local binade), which modulation then
|
||||
amplifies by |1 + scale|. Since xhat * (1 + scale) == y - shift, the bound
|
||||
on the final output is
|
||||
|
||||
|f - e| <= 2^-7 * (|e - shift| + |e|) + 1e-6
|
||||
|
||||
(one amplified xhat step plus one final-rounding step; the 1e-6 floor
|
||||
covers near-cancellation elements whose magnitude is below bf16
|
||||
resolution). Any structural defect -- wrong eps, missing bias, bad mean --
|
||||
violates this by orders of magnitude.
|
||||
"""
|
||||
f = fused_out.float()
|
||||
e = eager_out.float()
|
||||
assert not torch.isnan(f).any()
|
||||
if stream_dtype == torch.bfloat16:
|
||||
sh = 0.0 if shift is None else shift.float()
|
||||
tol = 2**-7 * ((e - sh).abs() + e.abs()) + 1e-6
|
||||
bad = (f - e).abs() > tol
|
||||
assert not bad.any(), (f"{bad.sum().item()} elements exceed the ulp bound; "
|
||||
f"max abs diff {(f - e).abs().max().item():.3e}")
|
||||
bit_equal = (fused_out.view(torch.int16) == eager_out.view(torch.int16)).float().mean().item()
|
||||
assert bit_equal > 0.999, f"only {bit_equal:.4%} of elements bit-identical"
|
||||
else:
|
||||
torch.testing.assert_close(f, e, rtol=1e-5, atol=1e-5)
|
||||
|
||||
|
||||
@cuda_only
|
||||
@pytest.mark.parametrize("stream_dtype", [torch.bfloat16, torch.float32])
|
||||
def test_site2_gate_vector_norm_only(stream_dtype):
|
||||
"""Wan post-self-attn: fp32 tensor gate, affine norm, no modulation."""
|
||||
eager, fused = _make_modules(elementwise_affine=True)
|
||||
residual, x = _inputs(stream_dtype)
|
||||
gate = torch.randn(1, 1, H, device="cuda", dtype=torch.float32)
|
||||
|
||||
with torch.no_grad():
|
||||
e_mod, e_res = eager(residual, x, gate, None, None)
|
||||
f_mod, f_res = fused(residual, x, gate, None, None)
|
||||
|
||||
# fused returns the stream dtype; eager promotes to fp32 -- compare after
|
||||
# the caller-side cast eager relies on.
|
||||
assert f_mod.dtype == stream_dtype and f_res.dtype == stream_dtype
|
||||
_assert_close_to_eager(f_mod, e_mod.to(stream_dtype), stream_dtype)
|
||||
_assert_close_to_eager(f_res, e_res.to(stream_dtype), stream_dtype)
|
||||
|
||||
|
||||
@cuda_only
|
||||
@pytest.mark.parametrize("stream_dtype", [torch.bfloat16, torch.float32])
|
||||
def test_site3_scalar_gate_with_modulation(stream_dtype):
|
||||
"""Wan post-cross-attn: gate==1, no-affine norm, scale/shift modulation."""
|
||||
eager, fused = _make_modules(elementwise_affine=False)
|
||||
residual, x = _inputs(stream_dtype, seed=2)
|
||||
shift = torch.randn(1, 1, H, device="cuda", dtype=torch.float32)
|
||||
scale = torch.randn(1, 1, H, device="cuda", dtype=torch.float32)
|
||||
|
||||
with torch.no_grad():
|
||||
e_mod, e_res = eager(residual, x, 1, shift, scale)
|
||||
f_mod, f_res = fused(residual, x, 1, shift, scale)
|
||||
|
||||
assert f_mod.dtype == stream_dtype and f_res.dtype == stream_dtype
|
||||
_assert_close_to_eager(f_mod, e_mod.to(stream_dtype), stream_dtype, shift=shift)
|
||||
_assert_close_to_eager(f_res, e_res.to(stream_dtype), stream_dtype)
|
||||
|
||||
|
||||
@cuda_only
|
||||
def test_per_token_modulation():
|
||||
"""ti2v-style per-token scale/shift ([B, S, H]) goes down the fused path."""
|
||||
eager, fused = _make_modules(elementwise_affine=False)
|
||||
residual, x = _inputs(torch.bfloat16, seed=3)
|
||||
shift = torch.randn(B, S, H, device="cuda", dtype=torch.float32)
|
||||
scale = torch.randn(B, S, H, device="cuda", dtype=torch.float32)
|
||||
|
||||
with torch.no_grad():
|
||||
e_mod, e_res = eager(residual, x, 1, shift, scale)
|
||||
f_mod, f_res = fused(residual, x, 1, shift, scale)
|
||||
|
||||
_assert_close_to_eager(f_mod, e_mod.to(torch.bfloat16), torch.bfloat16, shift=shift)
|
||||
_assert_close_to_eager(f_res, e_res.to(torch.bfloat16), torch.bfloat16)
|
||||
|
||||
|
||||
@cuda_only
|
||||
def test_residual_out_feeds_next_site_bit_exactly():
|
||||
"""Site-3 chaining: the bf16 residual the fused kernel emits must be
|
||||
bit-identical to eager's, because it feeds the next block's sum."""
|
||||
eager, fused = _make_modules(elementwise_affine=False)
|
||||
residual, x = _inputs(torch.bfloat16, seed=4)
|
||||
|
||||
with torch.no_grad():
|
||||
_, e_res = eager(residual, x, 1, None, None)
|
||||
_, f_res = fused(residual, x, 1, None, None)
|
||||
|
||||
# eager: bf16 + bf16 -> bf16 tensor; fused: fp32 sum rounded to bf16.
|
||||
# PyTorch's eager bf16 add also computes in fp32, so these must agree
|
||||
# bitwise -- no reduction is involved in the residual sum.
|
||||
assert torch.equal(e_res, f_res)
|
||||
|
||||
|
||||
@cuda_only
|
||||
def test_grad_enabled_falls_back_to_eager():
|
||||
"""Training path must not touch the fused kernel (it has no backward)."""
|
||||
_, fused = _make_modules(elementwise_affine=True)
|
||||
residual, x = _inputs(torch.bfloat16, seed=5)
|
||||
gate = torch.randn(1, 1, H, device="cuda", dtype=torch.float32)
|
||||
|
||||
out, res = fused(residual, x, gate, None, None)
|
||||
# eager promotion produces fp32 outputs; the fused kernel would have
|
||||
# produced bf16. dtype is the witness for which path ran.
|
||||
assert out.dtype == torch.float32
|
||||
assert res.dtype == torch.float32
|
||||
|
||||
|
||||
@cuda_only
|
||||
def test_env_kill_switch(env_overrides):
|
||||
import fastvideo.envs as envs
|
||||
env_overrides.enter_context(envs.FASTVIDEO_DISABLE_FUSED_NORM.override(True))
|
||||
_, fused = _make_modules(elementwise_affine=True)
|
||||
residual, x = _inputs(torch.bfloat16, seed=6)
|
||||
gate = torch.randn(1, 1, H, device="cuda", dtype=torch.float32)
|
||||
with torch.no_grad():
|
||||
out, _ = fused(residual, x, gate, None, None)
|
||||
assert out.dtype == torch.float32 # eager promotion, not the fused bf16
|
||||
|
||||
|
||||
@cuda_only
|
||||
def test_mixed_none_still_raises():
|
||||
_, fused = _make_modules(elementwise_affine=True)
|
||||
residual, x = _inputs(torch.bfloat16, seed=7)
|
||||
with pytest.raises(ValueError):
|
||||
with torch.no_grad():
|
||||
fused(residual, x, 1, None, torch.randn(1, 1, H, device="cuda"))
|
||||
|
||||
|
||||
@cuda_only
|
||||
def test_default_module_never_fuses():
|
||||
"""Existing users (fuse_inference unset) keep eager dtype semantics."""
|
||||
eager, _ = _make_modules(elementwise_affine=True)
|
||||
residual, x = _inputs(torch.bfloat16, seed=8)
|
||||
gate = torch.randn(1, 1, H, device="cuda", dtype=torch.float32)
|
||||
with torch.no_grad():
|
||||
out, res = eager(residual, x, gate, None, None)
|
||||
assert out.dtype == torch.float32
|
||||
assert res.dtype == torch.float32
|
||||
Reference in New Issue
Block a user