diff --git a/.buildkite/scripts/unit_test.sh b/.buildkite/scripts/unit_test.sh index 227bafdc5..0e950f0af 100644 --- a/.buildkite/scripts/unit_test.sh +++ b/.buildkite/scripts/unit_test.sh @@ -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 \ diff --git a/docs/contributing/env_vars.md b/docs/contributing/env_vars.md index 036d4e6ad..f5f9a6dcf 100644 --- a/docs/contributing/env_vars.md +++ b/docs/contributing/env_vars.md @@ -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. | diff --git a/fastvideo/envs.py b/fastvideo/envs.py index b52d7866b..b43562941 100644 --- a/fastvideo/envs.py +++ b/fastvideo/envs.py @@ -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 ================== diff --git a/fastvideo/layers/layernorm.py b/fastvideo/layers/layernorm.py index b74059888..553b50779 100644 --- a/fastvideo/layers/layernorm.py +++ b/fastvideo/layers/layernorm.py @@ -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) diff --git a/fastvideo/layers/triton_fused_norm.py b/fastvideo/layers/triton_fused_norm.py new file mode 100644 index 000000000..5d88ae775 --- /dev/null +++ b/fastvideo/layers/triton_fused_norm.py @@ -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 diff --git a/fastvideo/models/wan/transformer.py b/fastvideo/models/wan/transformer.py index 35d811d43..2f7b39710 100644 --- a/fastvideo/models/wan/transformer.py +++ b/fastvideo/models/wan/transformer.py @@ -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 diff --git a/fastvideo/tests/layers/test_triton_fused_norm.py b/fastvideo/tests/layers/test_triton_fused_norm.py new file mode 100644 index 000000000..3f3e3b2ea --- /dev/null +++ b/fastvideo/tests/layers/test_triton_fused_norm.py @@ -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