diff --git a/docs/adding_models.rst b/docs/adding_models.rst index 689e337d06..54f0117c79 100644 --- a/docs/adding_models.rst +++ b/docs/adding_models.rst @@ -2346,6 +2346,16 @@ attention over packed tokens: SDPA eagerly, the ragged FlashInfer resource under family), ``image_io.py`` (latent packing, PNG) and ``lora.py`` (static LoRA merge at load, also into fused ``q``/``k``/``v`` shards via ``FusedColumnLinear.shard_slice``). +Two pieces swap a one-pass kernel in for what Inductor would emit, inside a compiled +transformer only; eager calls run the plain PyTorch formulation, which is what a port's +bit-exact tests hold. ``modulation.ada_layer_norm`` is a block's gated residual add, +LayerNorm and shift/scale modulation, with the modulation given per group of tokens +(``[B, G, D]``: one group per sample, or per latent frame as in Wan2.2); Inductor reads a +LayerNorm row wider than 1024 twice. ``qk_norm_rope.rms_norm_rope`` is the q/k RMSNorm +across heads followed by the interleaved-pair RoPE, q and k in one launch. Both agree +with eager to bf16 rounding rather than bitwise. Wan2.2 is the first user; +``test/modular/test_diffusion_dit_kernels.py`` holds them to the eager form. + On the API side ``DiffusionImageAdapter`` (``mstar/api_server/openai/adapters.py``) serves ``/v1/images/generations`` and ``/v1/images/edits`` for any such model (``size`` becomes ``width``/``height``, ``seed`` and other knobs pass through); register it under the diff --git a/mstar/engine/resources/attn/dense.py b/mstar/engine/resources/attn/dense.py index 6eb9613761..74e8909bc7 100644 --- a/mstar/engine/resources/attn/dense.py +++ b/mstar/engine/resources/attn/dense.py @@ -197,6 +197,8 @@ def __init__( self._wrapper_class = wrapper_class self._wrappers: dict[EagerSlotKey, PrefixedPrefillWrapper] = {} + # wrapper -> the (cu_q, cu_k, causal) it was last planned with + self._planned: dict[EagerSlotKey, tuple] = {} self._wrapper_kwargs = dict( num_qo_heads=kv_config.num_qo_heads, num_kv_heads=kv_config.num_kv_heads, @@ -298,11 +300,11 @@ def plan(self, step: AttentionStep, ctx: StepContext): ) self._current_plans = { - plan_label: self._build_plan(kv_out, step.causal, self._wrapper(plan_label, ctx.slot)) + plan_label: self._build_plan(kv_out, step.causal, EagerSlotKey(label=plan_label, slot=ctx.slot)) for plan_label, kv_out in plan_outputs.items() } - def _build_plan(self, kv_out: KVPlanOutput, causal: bool, wrapper: PrefixedPrefillWrapper) -> DensePlan: + def _build_plan(self, kv_out: KVPlanOutput, causal: bool, slot_key: EagerSlotKey) -> DensePlan: """The per-segment gather layout plus the cumulative lengths one varlen kernel needs, in the packed order the KV plan defined.""" page_size = self._kv_config.page_size @@ -331,15 +333,23 @@ def _build_plan(self, kv_out: KVPlanOutput, causal: bool, wrapper: PrefixedPrefi max_q = max(max_q, view.to_compute) max_k = max(max_k, view.length) - # Fresh pinned host layouts per plan: the FA3 wrapper H2Ds them - # non-blocking (from pageable memory that copy would sync the stream), - # FlashInfer's plan reads them on the host. - pin = self._on_cuda - wrapper.plan( - torch.tensor(cu_q, dtype=torch.int32, pin_memory=pin), - torch.tensor(cu_k, dtype=torch.int32, pin_memory=pin), - causal=causal, - ) + # A wrapper holds the schedule of its last plan until it is planned + # again (``run`` only reads it), and a denoising request keeps its + # layout from step to step, so an unchanged layout is not replanned: + # FlashInfer's plan is a device kernel plus a stream sync. + wrapper = self._wrapper(slot_key.label, slot_key.slot) + layout = (tuple(cu_q), tuple(cu_k), causal) + if self._planned.get(slot_key) != layout: + # Fresh pinned host layouts per plan: the FA3 wrapper H2Ds them + # non-blocking (from pageable memory that copy would sync the + # stream), FlashInfer's plan reads them on the host. + pin = self._on_cuda + wrapper.plan( + torch.tensor(cu_q, dtype=torch.int32, pin_memory=pin), + torch.tensor(cu_k, dtype=torch.int32, pin_memory=pin), + causal=causal, + ) + self._planned[slot_key] = layout return DensePlan( segments=tuple(segments), cu_q=tuple(cu_q), diff --git a/mstar/model/components/diffusion/modulation.py b/mstar/model/components/diffusion/modulation.py new file mode 100644 index 0000000000..f42235bc1e --- /dev/null +++ b/mstar/model/components/diffusion/modulation.py @@ -0,0 +1,191 @@ +"""adaLN for DiT blocks: an optional gated residual add, LayerNorm and the +shift/scale modulation in one pass over the hidden stream. + +A DiT block normalizes its residual stream three times, each time right after +adding a branch output back into it, and the result feeds a bf16 GEMM: + + x = (x.float() + y.float() * gate).to(x.dtype) # the residual add + out = (LN(x.float()) * (1 + scale) + shift).to(x.dtype) # adaLN + +The modulation vectors (``shift``, ``scale``, ``gate``) are per group of +tokens: ``[B, G, D]``, group ``g`` covering tokens ``[g * S/G, (g+1) * S/G)``. +``G = 1`` is a per-sample modulation (FLUX), ``G = frames`` a per-frame one +(Wan's per-token timestep is constant over a latent frame). + +Compiled, Inductor lowers this to a looped two-pass reduction once ``D`` is +past its persistent-reduction width (1024), reading the row twice, so inside a +compiled region the kernel below does it in one: the row is read once into +registers, the mean and variance are taken there, and both outputs are +written once. Eager calls run the same math as plain PyTorch ops, which is +what the reference implementations hold bit-exact. The kernel rounds the +residual to the stream dtype before normalizing it, as eager does; its mean and +variance are reduced in a different order, so it agrees with eager to fp32 +rounding rather than bitwise. +""" + +from __future__ import annotations + +import torch + +try: + import triton + import triton.language as tl + from torch.library import triton_op, wrap_triton +except ImportError: # CPU-only installs + triton = None + + +def reference_ada_layer_norm( + x: torch.Tensor, + shift: torch.Tensor | None, + scale: torch.Tensor | None, + *, + eps: float, + weight: torch.Tensor | None = None, + bias: torch.Tensor | None = None, + residual: torch.Tensor | None = None, + gate: torch.Tensor | None = None, +) -> tuple[torch.Tensor, torch.Tensor]: + """The eager formulation; see ``ada_layer_norm`` for the arguments.""" + batch, seq, dim = x.shape + groups = next((t.shape[1] for t in (shift, scale, gate) if t is not None), 1) + + def grouped(t: torch.Tensor) -> torch.Tensor: + return t.unflatten(1, (groups, seq // groups)) + + def per_group(t: torch.Tensor) -> torch.Tensor: + return t.unsqueeze(2) + + xg = grouped(x) + if residual is not None: + branch = grouped(residual).float() + if gate is not None: + branch = branch * per_group(gate) + xg = (xg.float() + branch).to(x.dtype) + normed = torch.nn.functional.layer_norm( + xg.float(), (dim,), weight.float() if weight is not None else None, + bias.float() if bias is not None else None, eps, + ) + if scale is not None: + normed = normed * (1 + per_group(scale)) + if shift is not None: + normed = normed + per_group(shift) + return xg.flatten(1, 2), normed.to(x.dtype).flatten(1, 2) + + +if triton is not None: + + @triton.jit + def _ada_layer_norm_kernel( + x_ptr, res_ptr, gate_ptr, shift_ptr, scale_ptr, w_ptr, b_ptr, x_out_ptr, out_ptr, + seq, tokens_per_group, groups, mod_batch_stride, mod_group_stride, dim, eps, + HAS_RES: tl.constexpr, HAS_GATE: tl.constexpr, HAS_SHIFT: tl.constexpr, HAS_SCALE: tl.constexpr, + HAS_W: tl.constexpr, HAS_B: tl.constexpr, BLOCK_D: tl.constexpr, + ): + row = tl.program_id(0).to(tl.int64) + cols = tl.arange(0, BLOCK_D) + mask = cols < dim + offs = row * dim + cols + batch = row // seq + group = (row % seq) // tokens_per_group + mod = batch * mod_batch_stride + group * mod_group_stride + cols + + x = tl.load(x_ptr + offs, mask=mask, other=0.0) + if HAS_RES: + y = tl.load(res_ptr + offs, mask=mask, other=0.0).to(tl.float32) + if HAS_GATE: + y = y * tl.load(gate_ptr + mod, mask=mask, other=0.0) + x = (x.to(tl.float32) + y).to(x_ptr.dtype.element_ty) + tl.store(x_out_ptr + offs, x, mask=mask) + xf = x.to(tl.float32) + mean = tl.sum(xf, axis=0) / dim + diff = tl.where(mask, xf - mean, 0.0) + var = tl.sum(diff * diff, axis=0) / dim + # eps arrives as fp64 from a compiled caller; keep the row in fp32 + out = diff * tl.math.rsqrt(var + eps.to(tl.float32)) + if HAS_W: + out = out * tl.load(w_ptr + cols, mask=mask, other=0.0) + if HAS_B: + out = out + tl.load(b_ptr + cols, mask=mask, other=0.0) + if HAS_SCALE: + out = out * (1 + tl.load(scale_ptr + mod, mask=mask, other=0.0)) + if HAS_SHIFT: + out = out + tl.load(shift_ptr + mod, mask=mask, other=0.0) + tl.store(out_ptr + offs, out.to(out_ptr.dtype.element_ty), mask=mask) + + def _launch(x, residual, gate, shift, scale, weight, bias, eps, x_out, out): + batch, seq, dim = x.shape + mods = [t for t in (gate, shift, scale) if t is not None] + groups = mods[0].shape[1] if mods else 1 + for t in mods: + assert t.shape == mods[0].shape and t.stride() == mods[0].stride() and t.stride(2) == 1 + mod_strides = (mods[0].stride(0), mods[0].stride(1)) if mods else (0, 0) + # absent tensors go in as None (constexpr), never as an alias of x: a + # compiled caller treats every pointer the kernel could store through as + # mutated and copies it + wrap_triton(_ada_layer_norm_kernel)[(batch * seq,)]( + x, residual, gate, shift, scale, weight, bias, + x_out, out, seq, seq // groups, groups, mod_strides[0], mod_strides[1], dim, eps, + HAS_RES=residual is not None, HAS_GATE=gate is not None, HAS_SHIFT=shift is not None, + HAS_SCALE=scale is not None, HAS_W=weight is not None, HAS_B=bias is not None, + BLOCK_D=triton.next_power_of_2(dim), num_warps=8, + ) + + @triton_op("mstar::ada_layer_norm", mutates_args=()) + def _ada_layer_norm_op( + x: torch.Tensor, + shift: torch.Tensor | None, + scale: torch.Tensor | None, + weight: torch.Tensor | None, + bias: torch.Tensor | None, + eps: float, + ) -> torch.Tensor: + out = torch.empty_like(x) + _launch(x, None, None, shift, scale, weight, bias, eps, None, out) + return out + + @triton_op("mstar::residual_ada_layer_norm", mutates_args=()) + def _residual_ada_layer_norm_op( + x: torch.Tensor, + residual: torch.Tensor, + gate: torch.Tensor | None, + shift: torch.Tensor | None, + scale: torch.Tensor | None, + weight: torch.Tensor | None, + bias: torch.Tensor | None, + eps: float, + ) -> tuple[torch.Tensor, torch.Tensor]: + x_out = torch.empty_like(x) + out = torch.empty_like(x) + _launch(x, residual, gate, shift, scale, weight, bias, eps, x_out, out) + return x_out, out + + +def ada_layer_norm( + x: torch.Tensor, + shift: torch.Tensor | None, + scale: torch.Tensor | None, + *, + eps: float, + weight: torch.Tensor | None = None, + bias: torch.Tensor | None = None, + residual: torch.Tensor | None = None, + gate: torch.Tensor | None = None, +) -> tuple[torch.Tensor, torch.Tensor]: + """``(x', LN(x') * (1 + scale) + shift)`` with ``x' = x + residual * gate``. + + ``x`` and ``residual`` are ``[B, S, D]`` in the stream dtype; ``shift``, + ``scale`` and ``gate`` are fp32 ``[B, G, D]`` with ``G`` dividing ``S`` + (any of them may be None, as may ``residual``); ``weight`` and ``bias`` are + the LayerNorm affine, if any. Math is fp32; both outputs are in ``x``'s + dtype, and ``x'`` is ``x`` itself when there is no residual. + """ + if triton is not None and x.is_cuda and torch.compiler.is_compiling(): + if residual is None: + return x, _ada_layer_norm_op(x.contiguous(), shift, scale, weight, bias, eps) + return _residual_ada_layer_norm_op( + x.contiguous(), residual.contiguous(), gate, shift, scale, weight, bias, eps, + ) + return reference_ada_layer_norm( + x, shift, scale, eps=eps, weight=weight, bias=bias, residual=residual, gate=gate, + ) diff --git a/mstar/model/components/diffusion/qk_norm_rope.py b/mstar/model/components/diffusion/qk_norm_rope.py new file mode 100644 index 0000000000..0e65edd88a --- /dev/null +++ b/mstar/model/components/diffusion/qk_norm_rope.py @@ -0,0 +1,95 @@ +"""Query/key RMSNorm followed by the interleaved-pair rotary embedding, for DiTs +whose q/k norm spans the whole projection (one RMSNorm across all heads, as in +Wan's ``rms_norm_across_heads``). + +Eager, this is the model's own ``nn.RMSNorm`` and ``rope.apply_rotary_interleaved``. +Compiled, Inductor lowers the pair to several passes over each row (the norm's +reduction, then the rotated copy of the row, then the rotation), so inside a +compiled region one kernel does both q and k: each program holds a row in +registers, normalizes it, rounds the normalized row to the activation dtype +where the eager norm's output does, and rotates it in fp32. Its mean square is +reduced in a different order than the eager norm's, so the two agree to fp32 +rounding rather than bitwise. +""" + +from __future__ import annotations + +import torch +from torch import nn + +from mstar.model.components.diffusion.rope import apply_rotary_interleaved + +try: + import triton + import triton.language as tl + from torch.library import triton_op, wrap_triton +except ImportError: # CPU-only installs + triton = None + + +if triton is not None: + + @triton.jit + def _rms_norm_rope_kernel( + q_ptr, k_ptr, qw_ptr, kw_ptr, cos_ptr, sin_ptr, q_out_ptr, k_out_ptr, + seq, width, q_eps, k_eps, HEAD_DIM: tl.constexpr, BLOCK_PAIRS: tl.constexpr, + ): + row = tl.program_id(0).to(tl.int64) + if tl.program_id(1) == 0: + x_ptr, w_ptr, out_ptr, eps = q_ptr, qw_ptr, q_out_ptr, q_eps + else: + x_ptr, w_ptr, out_ptr, eps = k_ptr, kw_ptr, k_out_ptr, k_eps + pos = row % seq + pairs = tl.arange(0, BLOCK_PAIRS) + pmask = pairs < width // 2 + offs = pairs[:, None] * 2 + tl.arange(0, 2)[None, :] + mask = pmask[:, None] + + x = tl.load(x_ptr + row * width + offs, mask=mask, other=0.0).to(tl.float32) + mean_sq = tl.sum(tl.sum(x * x, axis=1), axis=0) / width + # eps arrives as fp64 from a compiled caller; keep the row in fp32 + rstd = tl.math.rsqrt(mean_sq + eps.to(tl.float32)) + w = tl.load(w_ptr + offs, mask=mask, other=0.0).to(tl.float32) + x = (x * rstd * w).to(out_ptr.dtype.element_ty).to(tl.float32) + + even, odd = tl.split(x) + # the tables repeat each frequency over its pair, so the even column serves both + col = pos * HEAD_DIM + (pairs % (HEAD_DIM // 2)) * 2 + cos = tl.load(cos_ptr + col, mask=pmask, other=0.0) + sin = tl.load(sin_ptr + col, mask=pmask, other=0.0) + out = tl.join(even * cos - odd * sin, odd * cos + even * sin) + tl.store(out_ptr + row * width + offs, out.to(out_ptr.dtype.element_ty), mask=mask) + + @triton_op("mstar::rms_norm_rope_pair", mutates_args=()) + def _rms_norm_rope_op( + q: torch.Tensor, k: torch.Tensor, q_weight: torch.Tensor, k_weight: torch.Tensor, + cos: torch.Tensor, sin: torch.Tensor, q_eps: float, k_eps: float, + ) -> tuple[torch.Tensor, torch.Tensor]: + batch, seq, width = q.shape + q_out, k_out = torch.empty_like(q), torch.empty_like(k) + wrap_triton(_rms_norm_rope_kernel)[(batch * seq, 2)]( + q, k, q_weight, k_weight, cos, sin, q_out, k_out, seq, width, q_eps, k_eps, + HEAD_DIM=cos.shape[-1], BLOCK_PAIRS=triton.next_power_of_2(width // 2), num_warps=8, + ) + return q_out, k_out + + +def rms_norm_rope( + q: torch.Tensor, k: torch.Tensor, q_norm: nn.RMSNorm, k_norm: nn.RMSNorm, + cos: torch.Tensor, sin: torch.Tensor, num_heads: int, +) -> tuple[torch.Tensor, torch.Tensor]: + """``rope(q_norm(q)), rope(k_norm(k))`` for ``q, k [B, S, H * D]`` and fp32 + tables ``cos, sin [S, D]`` (each frequency repeated over its pair), returned + as ``[B, S, H, D]``. Both norms span the whole ``H * D`` row.""" + + def eager(x: torch.Tensor, norm: nn.RMSNorm) -> torch.Tensor: + return apply_rotary_interleaved(norm(x).unflatten(2, (num_heads, -1)), cos, sin) + + if triton is not None and q.is_cuda and torch.compiler.is_compiling(): + assert q.shape == k.shape and cos.shape == (q.shape[1], q.shape[2] // num_heads) + q_out, k_out = _rms_norm_rope_op( + q.contiguous(), k.contiguous(), q_norm.weight, k_norm.weight, cos.contiguous(), sin.contiguous(), + q_norm.eps, k_norm.eps, + ) + return q_out.unflatten(2, (num_heads, -1)), k_out.unflatten(2, (num_heads, -1)) + return eager(q, q_norm), eager(k, k_norm) diff --git a/mstar/model/wan22/components/dit.py b/mstar/model/wan22/components/dit.py index 8a8fb97515..f990e9716f 100644 --- a/mstar/model/wan22/components/dit.py +++ b/mstar/model/wan22/components/dit.py @@ -1,9 +1,14 @@ """Native Wan2.2-TI2V-5B video DiT (dense 5B transformer). -Exact port of diffusers 0.39.0 ``WanTransformer3DModel``, restricted to the -TI2V-5B checkpoint's configuration: no image-conditioning branch, and the -per-token timestep path (the dit submodule always feeds a ``[batch, seq]`` -timestep grid). +Port of diffusers 0.39.0 ``WanTransformer3DModel``, restricted to the TI2V-5B +checkpoint's configuration: no image-conditioning branch, and the per-token +timestep path (``expand_timesteps``). Every token of a latent frame shares its +frame's timestep (I2V zeroes frame 0 only), so the timestep comes in per frame, +``[batch, frames]``, and the time embedding and the modulation are computed once +per frame and broadcast over the frame's tokens. cuBLAS may pick a different +GEMM kernel at the per-frame row count than at the reference's per-token one, so +these rows match the reference's to rounding rather than bitwise; given the same +rows, the blocks and the head are the reference's bits. Numerics contract (bf16 weights with fp32 islands): @@ -34,6 +39,8 @@ import torch.nn.functional as F from torch import nn +from mstar.model.components.diffusion.modulation import ada_layer_norm +from mstar.model.components.diffusion.qk_norm_rope import rms_norm_rope from mstar.model.components.mlp import MLP from mstar.model.wan22.config import Wan22Config @@ -116,7 +123,7 @@ def tables(self, device: torch.device) -> tuple[torch.Tensor, torch.Tensor]: return self._tables[device] def __call__(self, grid: tuple[int, int, int], device: torch.device) -> tuple[torch.Tensor, torch.Tensor]: - """Per-grid frequencies, each ``[1, t*h*w, 1, head_dim]`` fp32 + """Per-grid frequencies, each ``[t*h*w, head_dim]`` fp32 (reference ``WanRotaryPosEmbed.forward``).""" ppf, pph, ppw = grid freqs_cos, freqs_sin = self.tables(device) @@ -127,27 +134,11 @@ def expand_axes(split: tuple[torch.Tensor, ...]) -> torch.Tensor: f = split[0][:ppf].view(ppf, 1, 1, -1).expand(ppf, pph, ppw, -1) h = split[1][:pph].view(1, pph, 1, -1).expand(ppf, pph, ppw, -1) w = split[2][:ppw].view(1, 1, ppw, -1).expand(ppf, pph, ppw, -1) - return torch.cat([f, h, w], dim=-1).reshape(1, ppf * pph * ppw, 1, -1) + return torch.cat([f, h, w], dim=-1).reshape(ppf * pph * ppw, -1) return expand_axes(split_cos), expand_axes(split_sin) -def _apply_rotary_emb( - hidden_states: torch.Tensor, freqs_cos: torch.Tensor, freqs_sin: torch.Tensor -) -> torch.Tensor: - """Interleaved-pair rotary application (reference ``WanAttnProcessor`` - inner ``apply_rotary_emb``): fp32 math via type promotion against the - fp32 tables, cast back to the input dtype at the output-buffer - assignment.""" - x1, x2 = hidden_states.unflatten(-1, (-1, 2)).unbind(-1) - cos = freqs_cos[..., 0::2] - sin = freqs_sin[..., 1::2] - out = torch.empty_like(hidden_states) - out[..., 0::2] = x1 * cos - x2 * sin - out[..., 1::2] = x1 * sin + x2 * cos - return out.type_as(hidden_states) - - class Wan22DiTAttention(nn.Module): """Multi-head attention with across-heads RMSNorm on q/k (the checkpoint's ``rms_norm_across_heads``: one 3072-wide norm before the @@ -172,17 +163,16 @@ def forward( rotary_emb: tuple[torch.Tensor, torch.Tensor] | None = None, ) -> torch.Tensor: kv_source = hidden_states if encoder_hidden_states is None else encoder_hidden_states - query = self.q_norm(self.q_proj(hidden_states)) - key = self.k_norm(self.k_proj(kv_source)) - value = self.v_proj(kv_source) - - query = query.unflatten(2, (self.num_heads, -1)) - key = key.unflatten(2, (self.num_heads, -1)) - value = value.unflatten(2, (self.num_heads, -1)) + query = self.q_proj(hidden_states) + key = self.k_proj(kv_source) + value = self.v_proj(kv_source).unflatten(2, (self.num_heads, -1)) if rotary_emb is not None: - query = _apply_rotary_emb(query, *rotary_emb) - key = _apply_rotary_emb(key, *rotary_emb) + # across-heads RMSNorm, then the interleaved-pair rotation + query, key = rms_norm_rope(query, key, self.q_norm, self.k_norm, *rotary_emb, self.num_heads) + else: + query = self.q_norm(query).unflatten(2, (self.num_heads, -1)) + key = self.k_norm(key).unflatten(2, (self.num_heads, -1)) # [B, S, H, D] -> [B, H, S, D] SDPA -> back (the reference's # dispatch_attention_fn native path; bidirectional, no mask). @@ -195,9 +185,9 @@ def forward( class Wan22TimeTextEmbedding(nn.Module): - """Condition embedder: per-token timestep embedding + text projection + """Condition embedder: per-frame timestep embedding + text projection (reference ``WanTimeTextImageEmbedding``, image branch absent on - TI2V-5B). + TI2V-5B, its per-token rows computed once per frame). ``time_embedder`` is an fp32 island: the sinusoidal embedding and both its Linears run fp32; ``temb`` crosses to bf16 at ``type_as(text)``, so @@ -214,12 +204,12 @@ def __init__(self, dim: int, freq_dim: int, text_dim: int): def forward( self, timestep: torch.Tensor, encoder_hidden_states: torch.Tensor ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]: - """``timestep`` ``[B, seq]`` (per-token, ``expand_timesteps``); - returns ``temb [B, seq, dim]``, ``timestep_proj [B, seq, 6, dim]``, + """``timestep`` ``[B, frames]``, one per latent frame; returns + ``temb [B, frames, dim]``, ``timestep_proj [B, frames, 6, dim]``, projected text ``[B, 512, dim]`` — all in the text embeds' dtype.""" - batch, seq_len = timestep.shape + batch, frames = timestep.shape sincos = _sinusoidal_timestep_embedding(timestep.flatten(), self.freq_dim) - temb = self.time_embedder(sincos.unflatten(0, (batch, seq_len))) + temb = self.time_embedder(sincos.unflatten(0, (batch, frames))) temb = temb.type_as(encoder_hidden_states) timestep_proj = self.time_proj(F.silu(temb)).unflatten(2, (6, -1)) return temb, timestep_proj, self.text_embedder(encoder_hidden_states) @@ -252,24 +242,29 @@ def forward( timestep_proj: torch.Tensor, rotary_emb: tuple[torch.Tensor, torch.Tensor], ) -> torch.Tensor: - # [1, 1, 6, dim] + [B, seq, 6, dim] -> six [B, seq, dim] fp32 tensors. + # [1, 1, 6, dim] + [B, frames, 6, dim] -> six [B, frames, dim] fp32 + # tensors, each applied to all of a frame's tokens. shift_msa, scale_msa, gate_msa, c_shift_msa, c_scale_msa, c_gate_msa = ( - (self.scale_shift_table.unsqueeze(0) + timestep_proj.float()).chunk(6, dim=2) - ) - shift_msa, scale_msa, gate_msa = shift_msa.squeeze(2), scale_msa.squeeze(2), gate_msa.squeeze(2) - c_shift_msa, c_scale_msa, c_gate_msa = c_shift_msa.squeeze(2), c_scale_msa.squeeze(2), c_gate_msa.squeeze(2) + self.scale_shift_table.unsqueeze(0) + timestep_proj.float() + ).unbind(2) - norm_hidden = (self.norm1(hidden_states.float()) * (1 + scale_msa) + shift_msa).type_as(hidden_states) + _, norm_hidden = ada_layer_norm(hidden_states, shift_msa, scale_msa, eps=self.norm1.eps) attn_output = self.self_attn(norm_hidden, rotary_emb=rotary_emb) - hidden_states = (hidden_states.float() + attn_output * gate_msa).type_as(hidden_states) - - norm_hidden = self.norm2(hidden_states.float()).type_as(hidden_states) + hidden_states, norm_hidden = ada_layer_norm( + hidden_states, None, None, eps=self.norm2.eps, weight=self.norm2.weight, bias=self.norm2.bias, + residual=attn_output, gate=gate_msa, + ) attn_output = self.cross_attn(norm_hidden, encoder_hidden_states=encoder_hidden_states) - hidden_states = hidden_states + attn_output - - norm_hidden = (self.norm3(hidden_states.float()) * (1 + c_scale_msa) + c_shift_msa).type_as(hidden_states) + # the cross-attention residual is a plain add in the stream dtype (no gate) + hidden_states, norm_hidden = ada_layer_norm( + hidden_states, c_shift_msa, c_scale_msa, eps=self.norm3.eps, residual=attn_output, + ) ff_output = self.ffn(norm_hidden) - return (hidden_states.float() + ff_output.float() * c_gate_msa).type_as(hidden_states) + frames = timestep_proj.shape[1] + return ( + hidden_states.unflatten(1, (frames, -1)).float() + + ff_output.unflatten(1, (frames, -1)).float() * c_gate_msa.unsqueeze(2) + ).type_as(hidden_states).flatten(1, 2) class Wan22DiT(nn.Module): @@ -330,18 +325,18 @@ def forward( encoder_hidden_states: torch.Tensor, ) -> torch.Tensor: """``hidden_states`` ``[B, C, F, H, W]`` (bf16 latents), ``timestep`` - ``[B, post-patch seq]`` per-token grid (``expand_timesteps`` — the - only timestep form mstar's dit submodule produces), text embeds - ``[B, 512, text_dim]`` bf16. Returns ``[B, C_out, F, H, W]`` bf16.""" + ``[B, F // p_t]`` per post-patch frame (``expand_timesteps``: each + token takes its frame's value), text embeds ``[B, 512, text_dim]`` + bf16. Returns ``[B, C_out, F, H, W]`` bf16.""" batch_size, _, num_frames, height, width = hidden_states.shape p_t, p_h, p_w = self.config.patch_size grid = (num_frames // p_t, height // p_h, width // p_w) - if timestep.ndim != 2: + if timestep.ndim != 2 or timestep.shape[1] != grid[0]: raise ValueError( - f"Wan22DiT requires a per-token [batch, seq] timestep grid; got ndim={timestep.ndim}." + f"Wan22DiT requires a per-frame [batch, {grid[0]}] timestep; got {tuple(timestep.shape)}." ) - rotary_emb = self.rope(grid, hidden_states.device) + rotary_emb = tuple(t.contiguous() for t in self.rope(grid, hidden_states.device)) hidden_states = self.patch_embedding(hidden_states) hidden_states = hidden_states.flatten(2).transpose(1, 2) @@ -353,10 +348,9 @@ def forward( hidden_states = block(hidden_states, encoder_hidden_states, timestep_proj, rotary_emb) # Output head: 2-chunk adaLN in fp32 ([1, 1, 2, dim] island table + - # bf16 temb promotes), then project and unpatchify. - shift, scale = (self.scale_shift_table.unsqueeze(0) + temb.unsqueeze(2)).chunk(2, dim=2) - shift, scale = shift.squeeze(2), scale.squeeze(2) - hidden_states = (self.norm_out(hidden_states.float()) * (1 + scale) + shift).type_as(hidden_states) + # bf16 temb promotes, per frame), then project and unpatchify. + shift, scale = (self.scale_shift_table.unsqueeze(0) + temb.unsqueeze(2)).unbind(2) + _, hidden_states = ada_layer_norm(hidden_states, shift, scale, eps=self.norm_out.eps) hidden_states = self.proj_out(hidden_states) hidden_states = hidden_states.reshape(batch_size, *grid, p_t, p_h, p_w, -1) diff --git a/mstar/model/wan22/components/unipc.py b/mstar/model/wan22/components/unipc.py index 2ff7da2555..a2671ea9ea 100644 --- a/mstar/model/wan22/components/unipc.py +++ b/mstar/model/wan22/components/unipc.py @@ -111,7 +111,9 @@ def _bh2_rhos(rks: list[torch.Tensor], hh: torch.Tensor, order: int, device: tor b.append(h_phi_k * factorial_i / b_h) factorial_i *= i + 1 h_phi_k = h_phi_k / hh - 1 / factorial_i - return torch.linalg.solve(torch.stack(rows), torch.stack(b)) + # solve_ex: the same factorization as solve, without solve's host-side + # check of the LAPACK info (a device sync every step) + return torch.linalg.solve_ex(torch.stack(rows), torch.stack(b)).result def unipc_predictor_step( diff --git a/mstar/model/wan22/config.py b/mstar/model/wan22/config.py index 6e66703a5f..be99b3fc75 100644 --- a/mstar/model/wan22/config.py +++ b/mstar/model/wan22/config.py @@ -75,7 +75,7 @@ class Wan22Config: # torch.compile the inner DiT region (patchify -> blocks -> head). The UniPC # solver stays eager — its CPU-resident sigma trips Inductor (see # Wan22DitSubmodule). Default ON in serving; the reference-equivalence gates - # build with this False so the eager path stays the bit-exact reference. + # build with this False to hold the eager path against the reference. compile_dit: bool = True # "WxHxF" request shapes (width x height x num_frames) the compiled DiT # serves; None is the default shape (1280x704x81). Every other shape runs diff --git a/mstar/model/wan22/submodules.py b/mstar/model/wan22/submodules.py index 94d035f327..aaf4e1343f 100644 --- a/mstar/model/wan22/submodules.py +++ b/mstar/model/wan22/submodules.py @@ -38,6 +38,7 @@ unipc_predictor_step, ) from mstar.model.wan22.config import Wan22Config +from mstar.utils.h2d import to_device_nonblocking logger = logging.getLogger(__name__) @@ -155,18 +156,14 @@ def _set_buffer(module: nn.Module, dotted_name: str, value: torch.Tensor): setattr(owner, attr, value) -def _loop_step(fwd_info: CurrentForwardPassInfo, time_index: torch.Tensor) -> int: +def _loop_step(fwd_info: CurrentForwardPassInfo) -> int: """The denoise iteration this step runs, from the engine's loop counter. - ``time_index`` holds the same value, but it is the previous step's device - output: reading it on the host waits for that step to finish, which stalls - the speculative preparation of this one. Callers that drive the submodule - without a loop counter (tests, harnesses) fall back to the edge. + A value carried on a loop-back edge would be the previous step's device + output, and reading it on the host waits for that step to finish, which + stalls the speculative preparation of this one. """ - counts = getattr(fwd_info, "dynamic_loop_iter_counts", None) or {} - if DENOISE_LOOP_NAME in counts: - return int(counts[DENOISE_LOOP_NAME]) - return int(time_index.reshape(-1)[0].item()) + return int(fwd_info.dynamic_loop_iter_counts[DENOISE_LOOP_NAME]) def _latent_grid(config: Wan22Config, step_metadata: dict) -> tuple[int, int, int]: @@ -335,16 +332,15 @@ class Wan22DitSubmodule(_SingleRequestMixin, _Fp32IslandMixin, NodeSubmodule): Loop-carried edges, all float32 over the latent grid: latents current sample x_t - time_index [1] int64, the 0-based step index k unipc_model_outputs ring buffer of the last two converted outputs; see ``UniPCState`` unipc_last_sample the sample the previous predictor was given Nothing else is carried. The sigmas, the order ramp and the timestep pair all - follow from ``time_index`` and the per-request tables, so they are recomputed - rather than shipped across an edge. + follow from the engine's loop counter and the per-request tables, so they are + recomputed rather than shipped across an edge. - At iteration 0 the conductor sends the four edges empty and this submodule + At iteration 0 the conductor sends the three edges empty and this submodule seeds them (noise from the request's seed, on a CPU generator as diffusers does). I2V also consumes the persisted ``image_latent``: frame 0 of the input is replaced by it and its per-token timestep zeroed, and the conditioning is @@ -418,16 +414,12 @@ def prepare_inputs( generator = torch.Generator(device="cpu").manual_seed(fwd_info.random_seed) latents = torch.randn(shape, generator=generator, dtype=torch.float32).to(device) tensor_inputs["latents"] = latents - # 1-D, not 0-dim: the worker's output fanout reads dims[0] of every - # routed edge. - tensor_inputs["time_index"] = torch.zeros(1, dtype=torch.int64, device=device) tensor_inputs["unipc_model_outputs"] = torch.zeros( (SOLVER_ORDER, *shape), dtype=torch.float32, device=device ) tensor_inputs["unipc_last_sample"] = torch.zeros(shape, dtype=torch.float32, device=device) else: - time_index = inputs["time_index"][0] - k = _loop_step(fwd_info, time_index) + k = _loop_step(fwd_info) num_steps = int(fwd_info.step_metadata["num_inference_steps"]) if k >= num_steps: # Async scheduling dispatched an iteration past this request's @@ -441,7 +433,6 @@ def prepare_inputs( ) return None tensor_inputs["latents"] = inputs["latents"][0] - tensor_inputs["time_index"] = time_index tensor_inputs["unipc_model_outputs"] = inputs["unipc_model_outputs"][0] tensor_inputs["unipc_last_sample"] = inputs["unipc_last_sample"][0] return NodeInputs(tensor_inputs=tensor_inputs) @@ -451,7 +442,6 @@ def forward( graph_walk: str, engine_inputs: ModelInputsFromEngine, latents: torch.Tensor, - time_index: torch.Tensor, unipc_model_outputs: torch.Tensor, unipc_last_sample: torch.Tensor, text_embeds_pos: torch.Tensor, @@ -459,15 +449,17 @@ def forward( image_latent: torch.Tensor | None = None, **kwargs, ) -> NameToTensorList: - step_metadata = engine_inputs.single_request_info.step_metadata + fwd_info = engine_inputs.single_request_info + step_metadata = fwd_info.step_metadata num_steps = int(step_metadata["num_inference_steps"]) guidance_scale = float(step_metadata["guidance_scale"]) - k = int(time_index.item()) + k = _loop_step(fwd_info) device = latents.device with _inference_ctx(): sigmas, timesteps = make_unipc_tables(num_steps, self.config.flow_shift) - t = timesteps[k].to(device) + sigmas_device = to_device_nonblocking(sigmas, device) + t = to_device_nonblocking(timesteps, device)[k] # All-ones for T2V; frame 0 zeroed for I2V. mask = torch.ones(1, 1, *latents.shape[2:], dtype=torch.float32, device=device) @@ -478,9 +470,9 @@ def forward( model_input = latents model_input = model_input.to(self.transformer.dtype) - # expand_timesteps: per-token timestep over the post-patch grid - # ([:, ::2, ::2] == patch_size (1, 2, 2)); I2V frame-0 tokens get 0. - temp_ts = (mask[0, 0][:, ::2, ::2] * t).flatten() + # expand_timesteps: every token of a latent frame takes its frame's + # timestep (patch_size (1, 2, 2) keeps the frame axis); I2V frame 0 gets 0. + temp_ts = mask[0, 0, :, 0, 0] * t noise_pred = self._noise_prediction( model_input, temp_ts, text_embeds_pos, text_embeds_neg, guidance_scale ) @@ -493,7 +485,7 @@ def forward( UniPCState(model_outputs=unipc_model_outputs, last_sample=unipc_last_sample), this_model_output=m_k, this_sample=sample, - sigmas=sigmas, + sigmas=sigmas_device, step_index=k, order=unipc_effective_order(k - 1, num_steps), ) @@ -501,7 +493,7 @@ def forward( new_latents = unipc_predictor_step( UniPCState(model_outputs=new_ring, last_sample=sample), sample=sample, - sigmas=sigmas, + sigmas=sigmas_device, step_index=k, order=unipc_effective_order(k, num_steps), ) @@ -513,7 +505,6 @@ def forward( return { "latents": [new_latents], - "time_index": [time_index + 1], "unipc_model_outputs": [new_ring], "unipc_last_sample": [sample], } diff --git a/mstar/model/wan22/wan22_model.py b/mstar/model/wan22/wan22_model.py index 3ddc408882..40d3d2e1c1 100644 --- a/mstar/model/wan22/wan22_model.py +++ b/mstar/model/wan22/wan22_model.py @@ -175,11 +175,10 @@ def _build_video_gen_walk(self, i2v: bool) -> GraphSection: if i2v: dit_inputs.append("image_latent") # Loop-carried state: see Wan22DitSubmodule's docstring for the - # definitive shape/dtype inventory (latents, 0-based step index, the - # order-2 UniPC ring buffer, and the corrector's last_sample). + # definitive shape/dtype inventory (latents, the order-2 UniPC ring + # buffer, and the corrector's last_sample). dit_inputs += [ "latents", - "time_index", "unipc_model_outputs", "unipc_last_sample", ] @@ -191,7 +190,6 @@ def _build_video_gen_walk(self, i2v: bool) -> GraphSection: input_names=dit_inputs, outputs=[ GraphEdge(next_node="dit", name="latents"), - GraphEdge(next_node="dit", name="time_index"), GraphEdge(next_node="dit", name="unipc_model_outputs"), GraphEdge(next_node="dit", name="unipc_last_sample"), ], @@ -623,7 +621,6 @@ def _video_gen_inputs( inputs.append(edge) inputs += [ GraphEdge(next_node="dit", name="latents"), - GraphEdge(next_node="dit", name="time_index"), GraphEdge(next_node="dit", name="unipc_model_outputs"), GraphEdge(next_node="dit", name="unipc_last_sample"), ] diff --git a/test/modular/test_dense_attention.py b/test/modular/test_dense_attention.py index c8ab3f271f..efbc08a0d4 100644 --- a/test/modular/test_dense_attention.py +++ b/test/modular/test_dense_attention.py @@ -69,11 +69,11 @@ def _manager(wrapper_class=DensePrefixedPrefillWrapper, **kwargs) -> DenseAttent ) -def _ctx(views: list[SequenceView], **kwargs) -> StepContext: +def _ctx(views: list[SequenceView], slot: int = 0, **kwargs) -> StepContext: return StepContext( request_ids=tuple(dict.fromkeys(view.request_id for view in views)), graph_walk="gen", - slot=0, + slot=slot, capture=False, plan_results={ "kv": {"main": KVPlanOutput(cpu_indptrs=None, cuda_indptrs=None, views=views)} @@ -140,6 +140,38 @@ def test_cumulative_lengths_pack_in_kv_plan_order(self): assert [s.request_id for s in plan.segments] == ["r0", "r1"] assert plan.causal is False + def test_an_unchanged_layout_is_not_replanned(self, monkeypatch): + """A wrapper keeps its last schedule, so a step whose layout matches + what that wrapper was last planned with skips the plan; another + layout, another causal flag or the other slot's wrapper plans again.""" + calls = [] + real_plan = DensePrefixedPrefillWrapper.plan + + def counting_plan(self, *args, **kwargs): + calls.append(self) + return real_plan(self, *args, **kwargs) + + monkeypatch.setattr(DensePrefixedPrefillWrapper, "plan", counting_plan) + manager = _manager() + views = [_view("r0", [0, 1], prefix=3, fresh=2)] + step = AttentionStep(causal=False) + + manager.plan(step, _ctx(views)) + manager.plan(step, _ctx(views)) + assert len(calls) == 1 + # the skipped plan still hands the layer the planned wrapper + assert manager._current_plans["main"].wrapper is calls[0] + + manager.plan(step, _ctx(views, slot=1)) + assert len(calls) == 2 and calls[1] is not calls[0] + + manager.plan(AttentionStep(causal=True), _ctx(views)) + assert len(calls) == 3 + manager.plan(step, _ctx([_view("r0", [0, 1], prefix=3, fresh=3)])) + assert len(calls) == 4 + manager.plan(step, _ctx([_view("r0", [0, 1], prefix=3, fresh=3)])) + assert len(calls) == 4 + def test_causal_rides_on_the_step(self): manager = _manager() manager.plan(AttentionStep(causal=True), _ctx([_view("r0", [0], prefix=1, fresh=1)])) diff --git a/test/modular/test_diffusion_dit_kernels.py b/test/modular/test_diffusion_dit_kernels.py new file mode 100644 index 0000000000..35f2a719ab --- /dev/null +++ b/test/modular/test_diffusion_dit_kernels.py @@ -0,0 +1,146 @@ +"""The DiT components that swap in one-pass kernels inside a compiled region +(``modulation.ada_layer_norm``, ``qk_norm_rope.rms_norm_rope``), against their +eager formulations. + +The eager formulation is what the reference ports are held bit-exact in, so the +CPU tests pin it to the expressions a DiT block writes inline. The kernels +reduce in a different order and are held to bf16 rounding of the eager result: +at most one bf16 step anywhere, and almost every element equal. +""" + +from __future__ import annotations + +import pytest +import torch +import torch.nn.functional as F +from torch import nn + +from mstar.model.components.diffusion import modulation, qk_norm_rope +from mstar.model.wan22.components.dit import Wan22RoPE3D +from mstar.model.wan22.config import Wan22Config + +CUDA = torch.cuda.is_available() +needs_kernels = pytest.mark.skipif( + not CUDA or modulation.triton is None, reason="needs CUDA and Triton" +) + +B, FRAMES, HW, D = 2, 3, 40, 384 +EPS = 1e-6 + + +def _stream(device, dtype=torch.bfloat16, seed=0): + g = torch.Generator(device="cpu").manual_seed(seed) + x = (torch.randn(B, FRAMES * HW, D, generator=g) * 3).to(device, dtype) + y = torch.randn(B, FRAMES * HW, D, generator=g).to(device, dtype) + # strided per-frame views out of one [B, F, 6, D] table, as a block unbinds them + mod = (torch.randn(B, FRAMES, 6, D, generator=g) * 0.2).to(device) + w = (torch.randn(D, generator=g) * 0.1 + 1).to(device) + b = (torch.randn(D, generator=g) * 0.1).to(device) + return x, y, mod.unbind(2), w, b + + +def _inline_block_math(x, y, shift, scale, gate, w, b): + """The adaLN steps as a Wan block writes them on [B, F, HW, D] views.""" + xs, ys = x.unflatten(1, (FRAMES, HW)), y.unflatten(1, (FRAMES, HW)) + gated = (xs.float() + ys * gate.unsqueeze(2)).type_as(xs) + mod = (F.layer_norm(xs.float(), (D,), None, None, EPS) * (1 + scale.unsqueeze(2)) + shift.unsqueeze(2)).type_as(xs) + affine = F.layer_norm(gated.float(), (D,), w, b, EPS).type_as(xs) + plain = xs + ys + plain_mod = (F.layer_norm(plain.float(), (D,), None, None, EPS) * (1 + scale.unsqueeze(2)) + shift.unsqueeze(2)) + return [t.flatten(1, 2) for t in (gated, mod, affine, plain, plain_mod.type_as(xs))] + + +def test_eager_formulation_is_the_inline_block_math_bitwise(): + x, y, (shift, scale, gate, *_), w, b = _stream("cpu") + gated, mod, affine, plain, plain_mod = _inline_block_math(x, y, shift, scale, gate, w, b) + + _, out = modulation.ada_layer_norm(x, shift, scale, eps=EPS) + assert torch.equal(out, mod) + x_out, out = modulation.ada_layer_norm(x, None, None, eps=EPS, weight=w, bias=b, residual=y, gate=gate) + assert torch.equal(x_out, gated) and torch.equal(out, affine) + x_out, out = modulation.ada_layer_norm(x, shift, scale, eps=EPS, residual=y) + assert torch.equal(x_out, plain) and torch.equal(out, plain_mod) + + +def _slice_write_rotation(x, cos, sin): + """The reference's interleaved-pair rotation, written the way the Wan port did + before it moved to the shared stacked form: fp32 math through the promotion + against the tables, rounded to the stream dtype at the slice assignment.""" + x1, x2 = x.unflatten(-1, (-1, 2)).unbind(-1) + cos = cos[None, :, None, 0::2] + sin = sin[None, :, None, 1::2] + out = torch.empty_like(x) + out[..., 0::2] = x1 * cos - x2 * sin + out[..., 1::2] = x1 * sin + x2 * cos + return out + + +@pytest.mark.parametrize( + "device", ["cpu", pytest.param("cuda", marks=pytest.mark.skipif(not CUDA, reason="needs CUDA"))] +) +def test_eager_rms_norm_rope_is_the_slice_write_rotation_bitwise(device): + cfg = Wan22Config() + heads, head_dim = cfg.num_attention_heads, cfg.attention_head_dim + cos, sin = Wan22RoPE3D(head_dim, cfg.rope_max_seq_len)((3, 5, 8), torch.device(device)) + seq = cos.shape[0] + g = torch.Generator(device="cpu").manual_seed(0) + q = (torch.randn(2, seq, heads * head_dim, generator=g) * 3).to(device, torch.bfloat16) + k = (torch.randn(2, seq, heads * head_dim, generator=g) * 3).to(device, torch.bfloat16) + q_norm = nn.RMSNorm(heads * head_dim, eps=EPS, device=device, dtype=torch.bfloat16) + k_norm = nn.RMSNorm(heads * head_dim, eps=EPS, device=device, dtype=torch.bfloat16) + q_out, k_out = qk_norm_rope.rms_norm_rope(q, k, q_norm, k_norm, cos, sin, heads) + assert torch.equal(q_out, _slice_write_rotation(q_norm(q).unflatten(2, (heads, -1)), cos, sin)) + assert torch.equal(k_out, _slice_write_rotation(k_norm(k).unflatten(2, (heads, -1)), cos, sin)) + + +def _assert_bf16_close(got: torch.Tensor, want: torch.Tensor, min_equal: float = 0.999): + """Within one bf16 step of the eager result everywhere, equal almost everywhere. + The step is taken at unit scale too: an output that cancels toward zero + keeps the absolute error of its O(1) terms.""" + assert got.shape == want.shape and got.dtype == want.dtype + step = (want.float().abs() + 1.0) * 2.0**-7 + assert ((got.float() - want.float()).abs() <= step).all() + assert (got == want).float().mean().item() >= min_equal + + +@needs_kernels +@pytest.mark.parametrize("case", ["mod", "gated_affine", "plain_mod", "per_sample"]) +def test_ada_layer_norm_kernel_matches_eager(case): + x, y, (shift, scale, gate, *_), w, b = _stream("cuda") + if case == "per_sample": + shift, scale = shift[:, :1], scale[:, :1] + kwargs = { + "mod": dict(shift=shift, scale=scale), + "per_sample": dict(shift=shift, scale=scale), + "gated_affine": dict(shift=None, scale=None, weight=w, bias=b, residual=y, gate=gate), + "plain_mod": dict(shift=shift, scale=scale, residual=y), + }[case] + want_x, want = modulation.reference_ada_layer_norm(x, eps=EPS, **kwargs) + compiled = torch.compile( + lambda: modulation.ada_layer_norm(x, eps=EPS, **kwargs), fullgraph=True, dynamic=False + ) + got_x, got = compiled() + _assert_bf16_close(got_x, want_x) + _assert_bf16_close(got, want) + + +@needs_kernels +def test_rms_norm_rope_kernel_matches_eager(): + g = torch.Generator(device="cpu").manual_seed(1) + heads, head_dim, seq = 6, 128, 90 + q = torch.randn(2, seq, heads * head_dim, generator=g).cuda().bfloat16() + k = torch.randn(2, seq, heads * head_dim, generator=g).cuda().bfloat16() + q_norm = nn.RMSNorm(heads * head_dim, eps=EPS, device="cuda", dtype=torch.bfloat16) + k_norm = nn.RMSNorm(heads * head_dim, eps=EPS, device="cuda", dtype=torch.bfloat16) + with torch.no_grad(): + q_norm.weight.copy_(torch.randn(heads * head_dim, generator=g) * 0.1 + 1) + k_norm.weight.copy_(torch.randn(heads * head_dim, generator=g) * 0.1 + 1) + ang = torch.rand(seq, head_dim // 2, generator=g).cuda() * 6 + cos, sin = ang.cos().repeat_interleave(2, 1), ang.sin().repeat_interleave(2, 1) + with torch.inference_mode(): + want = qk_norm_rope.rms_norm_rope(q, k, q_norm, k_norm, cos, sin, heads) + got = torch.compile( + lambda: qk_norm_rope.rms_norm_rope(q, k, q_norm, k_norm, cos, sin, heads), fullgraph=True, dynamic=False + )() + for g_t, w_t in zip(got, want, strict=True): + _assert_bf16_close(g_t, w_t, min_equal=0.99) diff --git a/test/modular/test_step_prep_host_syncs.py b/test/modular/test_step_prep_host_syncs.py index cadebdade1..145a5fbb43 100644 --- a/test/modular/test_step_prep_host_syncs.py +++ b/test/modular/test_step_prep_host_syncs.py @@ -1,21 +1,22 @@ """Denoise step prep reads no tensor on the host (Cosmos3, Wan 2.2, OmniVoice). -The loop-carried step counter (``time_index`` / ``step_index``) is the -previous step's device output, so an ``.item()`` on it in ``prepare_inputs`` -waits for that step to finish and stalls the speculative preparation of the -next one. These loops read the step from the engine's per-batch loop counter -instead, as the DiT scaffold does, and keep Cosmos3's timesteps as a host list. +A loop-carried step counter (``time_index`` / ``step_index``) is the previous +step's device output, so an ``.item()`` on it in ``prepare_inputs`` waits for +that step to finish and stalls the speculative preparation of the next one. +These loops read the step from the engine's per-batch loop counter instead, as +the DiT scaffold does, and keep Cosmos3's timesteps as a host list. Each test walks a short schedule (including the async-overshoot step the prepare vetoes) and checks, per step: * no host read of any tensor (``item`` / ``tolist`` / ``int()`` / ...) happens while preparing with the loop counter set, as the engine sets it; -* the prepared inputs equal those of the edge-reading path, which is the - derivation the submodules used before (and still use for callers that pass - no loop counter). +* for Cosmos3 and OmniVoice, the prepared inputs equal those of the edge-reading + path, which callers that pass no loop counter still use. Wan 2.2 carries no + such edge, so a missing counter is an error there. -CPU only, no weights. +No weights. Wan 2.2's whole step also runs under CUDA's sync-debug mode, which +needs a GPU; everything else is CPU only. """ from __future__ import annotations @@ -24,6 +25,7 @@ import pytest import torch +from torch import nn from torch.overrides import TorchFunctionMode from mstar.engine.windowing import WindowSchedule @@ -100,29 +102,99 @@ def _check_schedule(prepare, walk_fwd, edge_inputs, loop, steps): # ---------------------------------------------------------------- Wan 2.2 +def _wan22_loop_inputs(shape, k): + return { + "text_embeds_pos": [torch.zeros(1, 8, 16)], + "text_embeds_neg": [torch.zeros(1, 8, 16)], + "latents": [torch.full(shape, float(k))], + "unipc_model_outputs": [torch.zeros((2, *shape))], + "unipc_last_sample": [torch.zeros(shape)], + } + + def test_wan22_step_prep_reads_no_tensor(): sub = Wan22DitSubmodule(transformer=None, config=Wan22Config()) shape = (1, 4, 2, 2, 2) steps = 3 - def inputs(k): - return { - "text_embeds_pos": [torch.zeros(1, 8, 16)], - "text_embeds_neg": [torch.zeros(1, 8, 16)], - "latents": [torch.full(shape, float(k))], - "time_index": [torch.tensor([k], dtype=torch.int64)], - "unipc_model_outputs": [torch.zeros((2, *shape))], - "unipc_last_sample": [torch.zeros(shape)], - } + def prepare(k): + info = _fwd("r", "video_gen", {"num_inference_steps": steps}, DENOISE_LOOP_NAME, k) + return sub.prepare_inputs("video_gen", info, _wan22_loop_inputs(shape, k)) - def fwd(loop, k): - return _fwd("r", "video_gen", {"num_inference_steps": steps}, loop, k) + for k in range(1, steps): + with _HostReads() as host: + got = prepare(k) + assert host.reads == [], f"step {k}: host reads {host.reads}" + assert got.tensor_inputs.keys() == _wan22_loop_inputs(shape, k).keys() + assert torch.equal(got.tensor_inputs["latents"], torch.full(shape, float(k))) + with _HostReads() as host: + assert prepare(steps) is None # overshoot vetoed + assert host.reads == [] - def prepare(fwd_info, inp): - return sub.prepare_inputs("video_gen", fwd_info, inp) - _check_schedule(prepare, fwd, inputs, DENOISE_LOOP_NAME, steps) - assert prepare(fwd(DENOISE_LOOP_NAME, steps), inputs(steps)) is None # overshoot vetoed +def test_wan22_step_prep_fails_without_the_loop_counter(): + sub = Wan22DitSubmodule(transformer=None, config=Wan22Config()) + info = _fwd("r", "video_gen", {"num_inference_steps": 3}) + with pytest.raises(KeyError, match=DENOISE_LOOP_NAME): + sub.prepare_inputs("video_gen", info, _wan22_loop_inputs((1, 4, 2, 2, 2), 1)) + + +class _WanStubTransformer(nn.Module): + """The Wan DiT's call signature and output on device ops alone, so the only + host syncs a step can take are the submodule's and the solver's.""" + + dtype = torch.bfloat16 + + def forward(self, hidden_states, timestep, encoder_hidden_states): + return hidden_states * (timestep[:, None, :, None, None] * 1e-3 + encoder_hidden_states.mean()) + + +def _wan22_steps(sub, walk, steps, latents, image_latent): + """``forward`` over a whole schedule as the engine drives it, the step from the + loop counter and the solver state over the loop-back edges.""" + device = latents.device + ring, last = torch.zeros((2, *latents.shape), device=device), torch.zeros_like(latents) + text_pos = torch.full((1, 8, 16), 1.0, device=device, dtype=torch.bfloat16) + text_neg = torch.full((1, 8, 16), 0.5, device=device, dtype=torch.bfloat16) + for k in range(steps): + info = SimpleNamespace( + step_metadata={"num_inference_steps": steps, "guidance_scale": 5.0}, + dynamic_loop_iter_counts={DENOISE_LOOP_NAME: k}, + ) + out = sub.forward( + walk, engine_inputs=SimpleNamespace(single_request_info=info), + latents=latents, unipc_model_outputs=ring, unipc_last_sample=last, + text_embeds_pos=text_pos, text_embeds_neg=text_neg, image_latent=image_latent, + ) + latents, ring, last = out["latents"][0], out["unipc_model_outputs"][0], out["unipc_last_sample"][0] + return latents + + +@pytest.mark.skipif(not torch.cuda.is_available(), reason="the host sync is a CUDA stream sync") +@pytest.mark.parametrize("walk", ["video_gen", "video_gen_i2v"]) +def test_wan22_forward_takes_no_host_sync(walk): + """A step that waits on the GPU keeps the worker from queueing it behind the one + still running. Sync-debug mode raises on one: a host read of a device tensor, a + blocking copy from pageable memory, or ``linalg.solve``'s check of its info. + Four steps cover the predictor alone, both corrector orders and the last step.""" + sub = Wan22DitSubmodule(_WanStubTransformer(), Wan22Config(in_channels=4, compile_dit=False)) + gen = torch.Generator().manual_seed(0) + shape = (1, 4, 3, 4, 4) + latents = torch.randn(shape, generator=gen) + image_latent = torch.randn(shape, generator=gen) if walk == sub._I2V_WALK else None + want = _wan22_steps(sub, walk, 4, latents, image_latent) + + latents = latents.cuda() + image_latent = None if image_latent is None else image_latent.cuda() + torch.cuda.synchronize() + torch.cuda.set_sync_debug_mode("error") + try: + got = _wan22_steps(sub, walk, 4, latents, image_latent) + finally: + torch.cuda.set_sync_debug_mode("default") + torch.cuda.synchronize() + # the tables staged without a sync are the host's: the run lands where the host's does + assert torch.allclose(got.cpu(), want, atol=1e-2), (got.cpu() - want).abs().max() # ---------------------------------------------------------------- OmniVoice diff --git a/test/modular/test_wan22_dit_compile_sizes.py b/test/modular/test_wan22_dit_compile_sizes.py index 2d90061cee..fdd3cf934d 100644 --- a/test/modular/test_wan22_dit_compile_sizes.py +++ b/test/modular/test_wan22_dit_compile_sizes.py @@ -13,7 +13,7 @@ from torch import nn from mstar.model.wan22.config import Wan22Config -from mstar.model.wan22.submodules import Wan22DitSubmodule +from mstar.model.wan22.submodules import DENOISE_LOOP_NAME, Wan22DitSubmodule from mstar.model.wan22.wan22_model import Wan22Model CHANNELS = 4 @@ -55,9 +55,10 @@ def _step(node: Wan22DitSubmodule, width: int, height: int, frames: int, guidanc meta = {"num_inference_steps": 4, "guidance_scale": guidance} out = node.forward( "video_gen", - engine_inputs=SimpleNamespace(single_request_info=SimpleNamespace(step_metadata=meta)), + engine_inputs=SimpleNamespace(single_request_info=SimpleNamespace( + step_metadata=meta, dynamic_loop_iter_counts={DENOISE_LOOP_NAME: 1}, + )), latents=torch.randn(shape, generator=gen), - time_index=torch.tensor([1]), unipc_model_outputs=torch.randn((2, *shape), generator=gen), unipc_last_sample=torch.randn(shape, generator=gen), text_embeds_pos=torch.randn(1, TEXT_LEN, TEXT_DIM, generator=gen), diff --git a/test/modular/test_wan22_model.py b/test/modular/test_wan22_model.py index 0c27a1a71b..cbeff77a60 100644 --- a/test/modular/test_wan22_model.py +++ b/test/modular/test_wan22_model.py @@ -42,7 +42,7 @@ ) # The dit node's loop-back edges (see Wan22DitSubmodule for the inventory). -LOOP_BACK_NAMES = {"latents", "time_index", "unipc_model_outputs", "unipc_last_sample"} +LOOP_BACK_NAMES = {"latents", "unipc_model_outputs", "unipc_last_sample"} def _make_model() -> Wan22Model: @@ -472,14 +472,13 @@ def test_wan22_check_stop_without_step_metadata_never_stops(): # ---------------------------------------------------------------------- -def _dit_loop_inputs(k: int) -> dict: - """The dit's inputs as they arrive on loop iteration k >= 1.""" +def _dit_loop_inputs() -> dict: + """The dit's inputs as they arrive on any loop iteration past the first.""" shape = (1, 4, 2, 2, 2) return { "text_embeds_pos": [torch.zeros(1, 8, 16)], "text_embeds_neg": [torch.zeros(1, 8, 16)], "latents": [torch.zeros(shape)], - "time_index": [torch.tensor([k], dtype=torch.int64)], "unipc_model_outputs": [torch.zeros((2, *shape))], "unipc_last_sample": [torch.zeros(shape)], } @@ -492,7 +491,7 @@ def test_wan22_dit_prepare_inputs_vetoes_async_overshoot(): result = submodule.prepare_inputs( graph_walk=Wan22Model.VIDEO_GEN_WALK, fwd_info=_fwd_info(requested_steps=2, iter_count=2), - inputs=_dit_loop_inputs(k=2), + inputs=_dit_loop_inputs(), ) assert result is None @@ -502,10 +501,10 @@ def test_wan22_dit_prepare_inputs_passes_in_range_iteration(): node_inputs = submodule.prepare_inputs( graph_walk=Wan22Model.VIDEO_GEN_WALK, fwd_info=_fwd_info(requested_steps=2, iter_count=1), - inputs=_dit_loop_inputs(k=1), + inputs=_dit_loop_inputs(), ) assert node_inputs is not None - assert int(node_inputs.tensor_inputs["time_index"][0]) == 1 + assert node_inputs.tensor_inputs.keys() == {"text_embeds_pos", "text_embeds_neg"} | LOOP_BACK_NAMES @pytest.mark.parametrize("guidance,expect_negative", [ diff --git a/test/modular/test_wan22_native_dit_equivalence.py b/test/modular/test_wan22_native_dit_equivalence.py index 406dadfd64..d3cfaa5220 100644 --- a/test/modular/test_wan22_native_dit_equivalence.py +++ b/test/modular/test_wan22_native_dit_equivalence.py @@ -1,18 +1,25 @@ """Per-layer equivalence: the native ``Wan22DiT`` vs the wrapped diffusers ``WanTransformer3DModel``, both loaded in the serving dtypes (bf16 + fp32 islands). -The native module is an exact port, so the bar is BIT-EXACT at every stage: the -weights (which also proves the loader's remap is a bijection), the RoPE tables, the -condition-embedder outputs, and the activations at patchify, at each of the 30 block -boundaries, and at the head. Never widen a bound to pass — a nonzero anywhere is a -port bug, and the printed table names the first divergent stage. +The bar is BIT-EXACT for the weights (which also proves the loader's remap is a +bijection), the RoPE tables, the projected text stream, and the activations at +patchify, at each of the 30 block boundaries, and at the head. Never widen a bound to +pass — a nonzero there is a port bug, and the printed table names the first divergent +stage. + +The native module computes the time embedding and modulation rows per latent frame and +the reference per token, and cuBLAS may pick another GEMM kernel at the other row count. +So the embedder is compared on the reference's per-token timesteps, where both run the +same GEMMs, and the stage comparison gives both modules the native per-frame rows +(expanded per token for the reference). Stage outputs are captured with forward hooks in execution order on BOTH modules, so a block-indexing mistake shifts every later row rather than passing silently. Inputs use the batch-2 CFG serving shape with an I2V-style per-token timestep grid (frame-0 tokens zeroed), which exercises the general path; T2V's uniform grid is the -special case. Skipped without CUDA or the checkpoint in the local HF cache. +special case. The native module takes that grid one value per latent frame. Skipped +without CUDA or the checkpoint in the local HF cache. """ from __future__ import annotations @@ -60,6 +67,7 @@ def _hf_cache_has_checkpoint() -> bool: # Oracle-run latent grid: 480x832x33 frames -> (9, 30, 52), post-patch # (9, 15, 26) = 3510 tokens. LATENT_SHAPE = (2, 48, 9, 30, 52) +TOKENS_PER_FRAME = (30 // 2) * (52 // 2) @pytest.fixture(scope="module") @@ -92,7 +100,14 @@ def inputs(): t = torch.tensor(999, dtype=torch.int64, device=DEVICE) temp_ts = (mask[0, 0][:, ::2, ::2] * t).flatten() timestep = temp_ts.unsqueeze(0).expand(2, -1) - return hidden, timestep, text + # the native module takes the same grid one value per frame + frame_timestep = (mask[0, 0, :, 0, 0] * t).unsqueeze(0).expand(2, -1) + return hidden, timestep, frame_timestep, text + + +def _per_token(per_frame: torch.Tensor) -> torch.Tensor: + """Native per-frame embedder rows expanded to the reference's per-token rows.""" + return per_frame.repeat_interleave(TOKENS_PER_FRAME, dim=1) # --------------------------------------------------------------------------- @@ -162,7 +177,9 @@ def _first_tensor(output): def test_per_layer_outputs_bitwise_identical(wrapped, native, inputs): - hidden, timestep, text = inputs + """Everything past the condition embedder, on the same time embedding and + modulation rows: the native module's, expanded per token for the reference.""" + hidden, timestep, frame_timestep, text = inputs stages: dict[str, list[torch.Tensor]] = {"wrapped": [], "native": []} hooks = [] @@ -172,29 +189,40 @@ def hook(_module, _args, output): stages[side].append(_first_tensor(output).detach()) hooks.append(module.register_forward_hook(hook)) - # Identical capture order on both sides: patchify, embedder temb, - # blocks 0..29 (execution order), then the forward return as the head. + rows = {} + + def keep_native_rows(_module, _args, output): + rows["temb"], rows["proj"], _ = output + + def use_native_rows(_module, _args, output): + _, _, text_out, image_out = output + return _per_token(rows["temb"]), _per_token(rows["proj"]).flatten(2), text_out, image_out + + hooks.append(native.condition_embedder.register_forward_hook(keep_native_rows)) + hooks.append(wrapped.condition_embedder.register_forward_hook(use_native_rows)) + + # Identical capture order on both sides: patchify, blocks 0..29 (execution + # order), then the forward return as the head. for side, model in (("wrapped", wrapped), ("native", native)): capture(side, model.patch_embedding) - capture(side, model.condition_embedder) assert len(model.blocks) == 30 for block in model.blocks: capture(side, block) try: with torch.no_grad(): + native_out = native(hidden_states=hidden, timestep=frame_timestep, encoder_hidden_states=text) wrapped_out = wrapped( hidden_states=hidden, timestep=timestep, encoder_hidden_states=text, return_dict=False, )[0] - native_out = native(hidden_states=hidden, timestep=timestep, encoder_hidden_states=text) finally: for hook in hooks: hook.remove() stages["wrapped"].append(wrapped_out) stages["native"].append(native_out) - labels = ["patchify", "time_embed"] + [f"block_{i:02d}" for i in range(30)] + ["head_out"] + labels = ["patchify"] + [f"block_{i:02d}" for i in range(30)] + ["head_out"] assert len(stages["wrapped"]) == len(stages["native"]) == len(labels) print("\nper-stage max-abs diff (wrapped vs native):") @@ -214,9 +242,10 @@ def hook(_module, _args, output): def test_condition_embedder_all_outputs_bitwise_identical(wrapped, native, inputs): - """The per-stage table sees only the embedder's first output (temb); pin - the modulation projection and the projected text stream too.""" - _, timestep, text = inputs + """The embedder given the reference's per-token timesteps, so both sides run the + same GEMM shapes: every output, the modulation projection and the projected text + stream included, must match exactly.""" + _, timestep, _, text = inputs with torch.no_grad(): ref_temb, ref_proj, ref_text, _ = wrapped.condition_embedder( timestep.flatten(), text, None, timestep_seq_len=timestep.shape[1] @@ -228,6 +257,7 @@ def test_condition_embedder_all_outputs_bitwise_identical(wrapped, native, input ("timestep_proj", proj, ref_proj), ("text_proj", text_out, ref_text), ): + assert ours.shape == theirs.shape and ours.dtype == theirs.dtype diff = (ours.float() - theirs.float()).abs().max().item() print(f"condition_embedder {name}: max_abs={diff:.3e}") assert torch.equal(ours, theirs), f"condition_embedder {name} diverges (max_abs {diff:.3e})" diff --git a/test/modular/test_wan22_reference_equivalence.py b/test/modular/test_wan22_reference_equivalence.py index 7c00524301..9120fe8f1c 100644 --- a/test/modular/test_wan22_reference_equivalence.py +++ b/test/modular/test_wan22_reference_equivalence.py @@ -34,6 +34,7 @@ import json import os import sys +from contextlib import contextmanager from pathlib import Path import pytest @@ -44,6 +45,7 @@ from mstar.conductor.request_info import CurrentForwardPassInfo from mstar.model.submodule_base import ModelInputsFromEngine, NodeSubmodule from mstar.model.wan22.submodules import ( + DENOISE_LOOP_NAME, Wan22DitSubmodule, Wan22TextEncoderSubmodule, ) @@ -64,18 +66,26 @@ # --------------------------------------------------------------------------- # Thresholds. NEVER loosen one to make a run pass — a regression is a finding. -# mstar's only deliberate departure is batching the two CFG forwards into one -# batch-2 call, so each bound below is either 0.0 or prices that batching's noise. +# mstar's deliberate departures are batching the two CFG forwards into one batch-2 +# call and computing the time embedding per frame instead of per token, so each bound +# below is either 0.0 or prices the noise of one of those. # --------------------------------------------------------------------------- # Same tokenizer, same UMT5 instance, same padded batch. TEXT_EMBEDS_MAX_ABS = 0.0 # Identical fp32 ops on identical inputs. UNIPC_PORT_MAX_ABS = 0.0 -# Run CFG the reference's way and the whole path must match exactly at every step. -SEQ_TRAJECTORY_MAX_ABS = 0.0 # Batched-CFG trajectory: batch-2 kernel noise compounds over the 50 steps (1.6e-3 # at step 0, 3.7 by step 49). ~1.2x that; a math bug passes 1e1 within a few steps. PER_STEP_MAX_ABS = 4.5 +# Run CFG the reference's way and only the per-frame time embedding is left. cuBLAS can +# round its rows differently from the per-token ones (see +# test_wan22_native_dit_equivalence.py), and that compounds like the batched-CFG noise +# above: 2.1 by step 49 on B300 under this file's numerics (temb off by 3.1e-5), 3.7 +# with temb's fp32 rows perturbed by 1.1e-4, the size of a TF32 / non-TF32 kernel change. +SEQ_TRAJECTORY_MAX_ABS = PER_STEP_MAX_ABS +# The same with the time embedding at the reference's row count, which leaves nothing to +# round differently: the whole path must match exactly at every step. +SEQ_PER_TOKEN_ROWS_MAX_ABS = 0.0 FINAL_LATENTS_MAX_ABS = 4.5 # One step, batched vs sequential CFG (3.9e-1 on outputs whose abs max is 8.6): bf16 # kernel noise, amplified up to 9x by the guidance combine. ~1.5x observed. @@ -94,10 +104,11 @@ # that compounds exactly like the batched-CFG kernel noise (1.6e-3 at step 0 growing # through step 49). Observed max over the trajectory: 2.66 (RTX 5090), 2.997 (H100), # 3.675 (H200). Bound = worst observed x ~1.3 headroom. The eager path stays the -# bit-exact reference (SEQ/UNIPC gates); this only prices compile. +# reference (SEQ/UNIPC gates); this only prices compile. COMPILED_VS_EAGER_MAX_ABS = 4.8 -# No-CFG (guidance_scale <= 1.0) batch-1 path vs a no-CFG oracle. Same discipline -# as the sequential-CFG trajectory: one batch-1 forward per step, so BIT-EXACT. +# No-CFG (guidance_scale <= 1.0) batch-1 path vs a no-CFG oracle, with the time +# embedding at the reference's row count as for SEQ_PER_TOKEN_ROWS_MAX_ABS: one batch-1 +# forward per step like the reference's, so every step must match exactly. NO_CFG_TRAJECTORY_MAX_ABS = 0.0 CUDA_AVAILABLE = torch.cuda.is_available() @@ -239,13 +250,10 @@ def text_embeds(oracle_meta) -> dict[str, torch.Tensor]: return result -# The bit-exact gates are defined against the EAGER transformer, so the shared +# The trajectory gates are defined against the EAGER transformer, so the shared # dit_submodule is built with torch.compile OFF. Set WAN22_TEST_COMPILE_DIT=1 to # rebuild it compiled and re-run the whole suite once — a "nothing downstream -# breaks" check: the priced/decode/structural tests still pass, while the bit-exact -# 0.0 gates (sequential-CFG, no-CFG) then FAIL by exactly the compile fusion delta, -# which is precisely why the default fixture keeps them on the eager path. The -# compiled path's own numeric bound is pinned by +# breaks" check. The compiled path's own numeric bound is pinned by # test_compiled_dit_matches_eager_per_step regardless of this switch. _TEST_COMPILE_DIT = os.environ.get("WAN22_TEST_COMPILE_DIT", "0").lower() in ("1", "true", "yes", "on") @@ -332,10 +340,11 @@ def t2v_trajectory(oracle_meta, text_embeds, dit_submodule) -> dict: try: with torch.no_grad(): for step in range(num_steps): + fwd_info.dynamic_loop_iter_counts[DENOISE_LOOP_NAME] = step outputs = _drive(dit_submodule, "video_gen", fwd_info, {**persisted, **carried}) carried = { name: outputs[name] - for name in ("latents", "time_index", "unipc_model_outputs", "unipc_last_sample") + for name in ("latents", "unipc_model_outputs", "unipc_last_sample") } latents = outputs["latents"][0] latents_per_step.append(latents.detach().cpu().float()) @@ -445,46 +454,80 @@ def _sequential_noise_prediction(self, model_input, temp_ts, text_embeds_pos, te return uncond + guidance_scale * (cond - uncond) -def test_sequential_cfg_trajectory_is_bitexact_vs_oracle(oracle_meta, text_embeds, dit_submodule): - """Drive the engine path with the reference's sequential batch-1 CFG: - every one of the 50 per-step latents must be BIT-EXACT vs the oracle. - This proves the whole mstar implementation (tokenize -> encode -> - per-token timesteps -> transformer -> inline UniPC) identical to the - stock pipeline, leaving the batch-2 CFG kernel shape as the suite's ONLY - priced deviation (tests 2/3/4).""" +def _sequential_trajectory_diffs(oracle_meta, text_embeds, dit_submodule) -> list[float]: + """Max-abs vs the oracle after each step, driving the engine path with the + reference's sequential batch-1 CFG.""" meta = _step_metadata(oracle_meta) fwd_info = _fwd_info("video_gen", meta, seed=oracle_meta["seed"]) persisted = { "text_embeds_pos": [text_embeds["ours_pos"]], "text_embeds_neg": [text_embeds["ours_neg"]], } - num_steps = oracle_meta["num_inference_steps"] - original = Wan22DitSubmodule._noise_prediction Wan22DitSubmodule._noise_prediction = _sequential_noise_prediction diffs: list[float] = [] try: carried: dict[str, list[torch.Tensor]] = {} with torch.no_grad(): - for step in range(num_steps): + for step in range(oracle_meta["num_inference_steps"]): + fwd_info.dynamic_loop_iter_counts[DENOISE_LOOP_NAME] = step outputs = _drive(dit_submodule, "video_gen", fwd_info, {**persisted, **carried}) carried = { name: outputs[name] - for name in ("latents", "time_index", "unipc_model_outputs", "unipc_last_sample") + for name in ("latents", "unipc_model_outputs", "unipc_last_sample") } ref = torch.load(ORACLE_DIR / "latents" / f"step_{step:03d}.pt", map_location="cpu") diffs.append((outputs["latents"][0].cpu().float() - ref.float()).abs().max().item()) finally: Wan22DitSubmodule._noise_prediction = original + return diffs + + +@contextmanager +def _per_token_rows(transformer: torch.nn.Module, meta: dict): + """Run ``transformer``'s time embedding at the reference's row count, one row per + token, and hand the blocks one row per frame as usual.""" + embedder = transformer.condition_embedder + per_frame_rows = embedder.forward + tokens_per_frame = (meta["height"] // 32) * (meta["width"] // 32) + + def per_token_rows(timestep, encoder_hidden_states): + temb, proj, text = per_frame_rows(timestep.repeat_interleave(tokens_per_frame, dim=1), encoder_hidden_states) + return temb[:, ::tokens_per_frame], proj[:, ::tokens_per_frame], text + + embedder.forward = per_token_rows + try: + yield + finally: + del embedder.forward # restore class-method lookup + +def _report_trajectory(label: str, diffs: list[float]) -> str: nonzero = [(i, d) for i, d in enumerate(diffs) if d > 0] - print(f"sequential-CFG trajectory vs oracle: max_abs={max(diffs):.3e}, " - f"first nonzero step={nonzero[0] if nonzero else None}") - assert max(diffs) <= SEQ_TRAJECTORY_MAX_ABS, ( - f"sequential-CFG trajectory diverged from the oracle: first nonzero " - f"step {nonzero[0][0]} (max_abs {nonzero[0][1]:.3e}) — implementation " - "regression, not kernel noise" - ) + line = f"{label} vs oracle: max_abs={max(diffs):.3e}, first nonzero step={nonzero[0] if nonzero else None}" + print(line) + return line + + +def test_sequential_cfg_trajectory_matches_oracle(oracle_meta, text_embeds, dit_submodule): + """Each of the 50 per-step latents must stay within ``SEQ_TRAJECTORY_MAX_ABS`` of + the oracle, over the whole mstar path (tokenize -> encode -> per-frame timesteps + -> transformer -> inline UniPC). The time embedder itself is held bitwise to the + reference's on per-token rows in test_wan22_native_dit_equivalence.py.""" + diffs = _sequential_trajectory_diffs(oracle_meta, text_embeds, dit_submodule) + line = _report_trajectory("sequential-CFG trajectory", diffs) + assert max(diffs) <= SEQ_TRAJECTORY_MAX_ABS, line + + +def test_sequential_cfg_trajectory_is_bitexact_on_per_token_rows(oracle_meta, text_embeds, dit_submodule): + """With the time embedding run at the reference's row count, each of the 50 + per-step latents must be BIT-EXACT vs the oracle. This holds everything but the + embedder's row count identical to the stock pipeline, leaving that and the + batch-2 CFG kernel shape as the suite's ONLY priced deviations.""" + with _per_token_rows(dit_submodule.transformer, oracle_meta): + diffs = _sequential_trajectory_diffs(oracle_meta, text_embeds, dit_submodule) + line = _report_trajectory("sequential-CFG trajectory, per-token rows", diffs) + assert max(diffs) <= SEQ_PER_TOKEN_ROWS_MAX_ABS, line # --------------------------------------------------------------------------- @@ -507,7 +550,7 @@ def test_cfg_batched_matches_sequential(oracle_meta, text_embeds, dit_submodule) model_input = latents.to(transformer.dtype) t = torch.tensor(999, dtype=torch.int64, device=DEVICE) mask = torch.ones(1, 1, *latents.shape[2:], dtype=torch.float32, device=DEVICE) - temp_ts = (mask[0, 0][:, ::2, ::2] * t).flatten() + temp_ts = mask[0, 0, :, 0, 0] * t # one per latent frame, as the submodule builds it pos, neg = text_embeds["ours_pos"], text_embeds["ours_neg"] with torch.no_grad(): @@ -548,7 +591,7 @@ def test_i2v_first_frame_injection_semantics(oracle_meta, text_embeds, dit_submo """8-step small-resolution I2V run asserting the expand_timesteps conditioning contract: * the transformer INPUT's frame 0 equals the condition at every step, with - frame-0 per-token timesteps zeroed; + frame 0's timestep zeroed; * the loop-carried latents stay un-injected mid-loop; * the OUTPUT latents get the condition only at the final iteration. """ @@ -574,10 +617,11 @@ def test_i2v_first_frame_injection_semantics(oracle_meta, text_embeds, dit_submo carried: dict[str, list[torch.Tensor]] = {} with torch.no_grad(): for step in range(num_steps): + fwd_info.dynamic_loop_iter_counts[DENOISE_LOOP_NAME] = step outputs = _drive(dit_submodule, "video_gen_i2v", fwd_info, {**persisted, **carried}) carried = { name: outputs[name] - for name in ("latents", "time_index", "unipc_model_outputs", "unipc_last_sample") + for name in ("latents", "unipc_model_outputs", "unipc_last_sample") } out_latents = outputs["latents"][0] frame0_injected = torch.equal(out_latents[:, :, 0], condition[:, :, 0]) @@ -598,11 +642,11 @@ def test_i2v_first_frame_injection_semantics(oracle_meta, text_embeds, dit_submo assert torch.equal(hidden_states[b : b + 1, :, 0], cond_bf16[:, :, 0]), ( f"step {step} batch {b}: transformer input frame 0 != condition" ) - # Per-token timesteps: frame-0 tokens 0, all later frames uniform t. - grid = timestep[0].reshape(t_lat, (h_lat + 1) // 2, (w_lat + 1) // 2) - assert torch.all(grid[0] == 0), f"step {step}: frame-0 timesteps not zeroed" - assert torch.all(grid[1:] == grid[1, 0, 0]), f"step {step}: non-uniform later timesteps" - assert grid[1, 0, 0].item() > 0 + # One timestep per latent frame: frame 0 at 0, all later frames at a uniform t. + assert timestep.shape == (2, t_lat) + assert timestep[0, 0] == 0, f"step {step}: frame-0 timestep not zeroed" + assert torch.all(timestep[0, 1:] == timestep[0, 1]), f"step {step}: non-uniform later timesteps" + assert timestep[0, 1].item() > 0 # --------------------------------------------------------------------------- @@ -734,8 +778,8 @@ def test_compiled_dit_matches_eager_per_step(oracle_meta, text_embeds, dit_submo * the compiled trajectory stays within ``COMPILED_VS_EAGER_MAX_ABS`` of eager, a bound derived from the observed per-step max. - The eager path stays the bit-exact reference (the SEQ/UNIPC gates); this test - prices compile and nothing else. + The eager path stays the reference (the SEQ/UNIPC gates); this test prices + compile and nothing else. """ import torch._dynamo from torch._inductor.compile_fx import compile_fx @@ -776,10 +820,11 @@ def _counting_inductor(gm, example_inputs): carried: dict[str, list[torch.Tensor]] = {} with torch.no_grad(): for step in range(num_steps): + fwd_info.dynamic_loop_iter_counts[DENOISE_LOOP_NAME] = step outputs = _drive(dit_submodule, "video_gen", fwd_info, {**persisted, **carried}) carried = { name: outputs[name] - for name in ("latents", "time_index", "unipc_model_outputs", "unipc_last_sample") + for name in ("latents", "unipc_model_outputs", "unipc_last_sample") } diffs.append( (outputs["latents"][0].cpu().float() - eager_latents[step]).abs().max().item() @@ -805,7 +850,7 @@ def _counting_inductor(gm, example_inputs): # --------------------------------------------------------------------------- # 8. No-CFG batch-1 path — Feature 4. guidance_scale <= 1.0 takes the batch-1 -# branch no other test exercises; it is bit-exact vs a no-CFG oracle. +# branch no other test exercises, against a no-CFG oracle. # --------------------------------------------------------------------------- @@ -817,12 +862,12 @@ def _counting_inductor(gm, example_inputs): "--height 224 --width 384 --num-frames 9 --out-dir ` (oracle venv, same torch)" ), ) -def test_no_cfg_batch1_trajectory_is_bitexact(oracle_meta, text_embeds, dit_submodule): +def test_no_cfg_batch1_trajectory_is_bitexact_on_per_token_rows(oracle_meta, text_embeds, dit_submodule): """guidance_scale <= 1.0 runs one batch-1 forward per step (no CFG), the branch no other test covered. Drive the full no-CFG oracle trajectory at guidance 1.0 - and require BIT-EXACT latents at every step — the same discipline as the - sequential-CFG gate, since a no-CFG step is one batch-1 forward exactly like the - reference at guidance 1.0. A recorder PROVES the batch-1 branch is taken (the + with the time embedding at the reference's row count and hold every step + BIT-EXACT, like the sequential-CFG gate on per-token rows, since a no-CFG + step is one batch-1 forward like the reference at guidance 1.0. A recorder PROVES the batch-1 branch is taken (the transformer sees batch 1, not the CFG batch-2), so the test cannot pass while silently exercising the wrong path. """ @@ -853,12 +898,13 @@ def test_no_cfg_batch1_trajectory_is_bitexact(oracle_meta, text_embeds, dit_subm diffs: list[float] = [] try: carried: dict[str, list[torch.Tensor]] = {} - with torch.no_grad(): + with torch.no_grad(), _per_token_rows(original, nocfg_meta): for step in range(num_steps): + fwd_info.dynamic_loop_iter_counts[DENOISE_LOOP_NAME] = step outputs = _drive(dit_submodule, "video_gen", fwd_info, {**persisted, **carried}) carried = { name: outputs[name] - for name in ("latents", "time_index", "unipc_model_outputs", "unipc_last_sample") + for name in ("latents", "unipc_model_outputs", "unipc_last_sample") } ref = torch.load(NOCFG_ORACLE_DIR / "latents" / f"step_{step:03d}.pt", map_location="cpu") diffs.append((outputs["latents"][0].cpu().float() - ref.float()).abs().max().item()) @@ -873,11 +919,5 @@ def test_no_cfg_batch1_trajectory_is_bitexact(oracle_meta, text_embeds, dit_subm "batch-2 path was taken, so this test would not be exercising no-CFG" ) - nonzero = [(i, d) for i, d in enumerate(diffs) if d > 0] - print(f"no-CFG batch-1 trajectory vs oracle: max_abs={max(diffs):.3e}, " - f"first nonzero step={nonzero[0] if nonzero else None}") - assert max(diffs) <= NO_CFG_TRAJECTORY_MAX_ABS, ( - f"no-CFG trajectory diverged from the oracle: first nonzero step " - f"{nonzero[0][0]} (max_abs {nonzero[0][1]:.3e}) — the batch-1 path is not " - "bit-exact against the reference" - ) + line = _report_trajectory("no-CFG batch-1 trajectory", diffs) + assert max(diffs) <= NO_CFG_TRAJECTORY_MAX_ABS, line