From 6ac1bc660dac09167807b8d829ad799c5c5c73dc Mon Sep 17 00:00:00 2001 From: merceod Date: Thu, 29 Oct 2026 00:23:46 -0700 Subject: [PATCH 01/13] Cosmos3: the batched-step warmup hands the step the rotary table (cherry picked from commit bada56bcea2de856fc32bb4a8e5de7bde3b11ef8) --- mstar/model/cosmos3/submodules.py | 5 +++-- mstar/model/cosmos3/tests/test_denoise_compile.py | 12 +++++++++++- 2 files changed, 14 insertions(+), 3 deletions(-) diff --git a/mstar/model/cosmos3/submodules.py b/mstar/model/cosmos3/submodules.py index 29dad4781f..c2dc901d85 100644 --- a/mstar/model/cosmos3/submodules.py +++ b/mstar/model/cosmos3/submodules.py @@ -398,8 +398,9 @@ def warmup_compiles(self) -> None: dtype=self.transformer.proj_in.weight.dtype, ), "vision_timesteps": torch.zeros(static["num_noisy_vision_tokens"], device=device), - "position_ids_cond": static["vision_mrope_ids"], - "position_ids_uncond": static["vision_mrope_ids"], + # the step takes the mRoPE table the statics built, not the position ids + "rotary_cond": static["rotary"], + "rotary_uncond": static["rotary"], "vision_token_shapes": static["vision_token_shapes"], "vision_noisy_frame_indexes": static["vision_noisy_frame_indexes"], "vision_mse_loss_indexes": static["mse_gen_indexes"], diff --git a/mstar/model/cosmos3/tests/test_denoise_compile.py b/mstar/model/cosmos3/tests/test_denoise_compile.py index a5fd4cd2fa..8f4e8ca068 100644 --- a/mstar/model/cosmos3/tests/test_denoise_compile.py +++ b/mstar/model/cosmos3/tests/test_denoise_compile.py @@ -90,7 +90,14 @@ def __init__(self): self.proj_in = torch.nn.Linear(2, 2) self.batched_calls = [] + # what the real step indexes per request; a warmup row missing one would raise here as it did on a server + REQ_KEYS = ("latents", "vision_timesteps", "rotary_cond", "rotary_uncond", "vision_token_shapes", + "vision_noisy_frame_indexes", "vision_mse_loss_indexes") + def denoise_step_batched(self, requests, label, attn): + for req in requests: + for key in self.REQ_KEYS: + req[key] self.batched_calls.append((requests, label, attn)) return [(None, None)] * len(requests) @@ -103,11 +110,12 @@ def test_warmup_compiles_runs_one_two_request_step_at_the_default_size(monkeypat t = _RecordingTransformer() sub = Cosmos3DiTSubmodule(t, cfg) built = [] + rotary = (torch.zeros(4, 2), torch.zeros(4, 2)) def fake_static(ids, height, width, num_frames, fps, has_image_condition, device): built.append((height, width, num_frames, has_image_condition)) return { - "num_noisy_vision_tokens": 4, "vision_mrope_ids": torch.zeros(3, 4), + "num_noisy_vision_tokens": 4, "rotary": rotary, "vision_token_shapes": [(1, 2, 2)], "vision_noisy_frame_indexes": [torch.zeros(1)], "mse_gen_indexes": torch.arange(4), } @@ -121,6 +129,8 @@ def fake_static(ids, height, width, num_frames, fps, has_image_condition, device [(requests, label, attn)] = t.batched_calls assert len(requests) == 2 and label == submodules.CFG_BATCHED_LABEL assert requests[0]["latents"].shape == sub._latent_shape(480, 832, 121) + # both branches take the statics' table, as a served window's rows do + assert requests[0]["rotary_cond"] is rotary and requests[1]["rotary_uncond"] is rotary # the stand-in reads no plan and answers with the query's shape q = torch.randn(5, 2, 3) assert not attn.requires_kv_write and torch.equal(attn.run(q, k=q, v=q), torch.zeros_like(q)) From 2e3ecc57c4cecf16356e2deabcf929bcc18d7660 Mon Sep 17 00:00:00 2001 From: merceod Date: Thu, 29 Oct 2026 19:41:03 +0000 Subject: [PATCH 02/13] wan22: per-frame modulation through a one-pass adaLN component Every token of a latent frame shares its timestep, so the time embedding and the six modulation vectors are computed once per frame (bit-identical rows) instead of per token. The norm+modulation+gated-residual steps go through components/diffusion/modulation.py, which runs a single-pass Triton kernel inside the compiled DiT and the plain eager ops otherwise. --- .../model/components/diffusion/modulation.py | 191 ++++++++++++++++++ mstar/model/wan22/components/dit.py | 65 +++--- mstar/model/wan22/submodules.py | 6 +- .../test_wan22_native_dit_equivalence.py | 33 ++- .../test_wan22_reference_equivalence.py | 4 +- 5 files changed, 255 insertions(+), 44 deletions(-) create mode 100644 mstar/model/components/diffusion/modulation.py 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/wan22/components/dit.py b/mstar/model/wan22/components/dit.py index 8a8fb97515..803043d8dd 100644 --- a/mstar/model/wan22/components/dit.py +++ b/mstar/model/wan22/components/dit.py @@ -2,8 +2,12 @@ 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). +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. Per-frame rows +are the reference's per-token rows bit for bit (the rows of a GEMM do not +depend on how many there are). Numerics contract (bf16 weights with fp32 islands): @@ -34,6 +38,7 @@ import torch.nn.functional as F from torch import nn +from mstar.model.components.diffusion.modulation import ada_layer_norm from mstar.model.components.mlp import MLP from mstar.model.wan22.config import Wan22Config @@ -214,12 +219,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 +257,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,15 +340,15 @@ 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) @@ -353,10 +363,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/submodules.py b/mstar/model/wan22/submodules.py index 2413c41cd4..d18de914a0 100644 --- a/mstar/model/wan22/submodules.py +++ b/mstar/model/wan22/submodules.py @@ -478,9 +478,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 ) diff --git a/test/modular/test_wan22_native_dit_equivalence.py b/test/modular/test_wan22_native_dit_equivalence.py index 406dadfd64..099b6bde85 100644 --- a/test/modular/test_wan22_native_dit_equivalence.py +++ b/test/modular/test_wan22_native_dit_equivalence.py @@ -12,7 +12,9 @@ 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 and +computes its time embedding per frame; its rows are compared per token. Skipped +without CUDA or the checkpoint in the local HF cache. """ from __future__ import annotations @@ -60,6 +62,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 +95,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,21 +172,22 @@ def _first_tensor(output): def test_per_layer_outputs_bitwise_identical(wrapped, native, inputs): - hidden, timestep, text = inputs + hidden, timestep, frame_timestep, text = inputs stages: dict[str, list[torch.Tensor]] = {"wrapped": [], "native": []} hooks = [] - def capture(side, module): + def capture(side, module, transform=None): def hook(_module, _args, output): - stages[side].append(_first_tensor(output).detach()) + out = _first_tensor(output).detach() + stages[side].append(transform(out) if transform else out) 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. for side, model in (("wrapped", wrapped), ("native", native)): capture(side, model.patch_embedding) - capture(side, model.condition_embedder) + capture(side, model.condition_embedder, _per_token if side == "native" else None) assert len(model.blocks) == 30 for block in model.blocks: capture(side, block) @@ -187,7 +198,7 @@ def hook(_module, _args, output): 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) + native_out = native(hidden_states=hidden, timestep=frame_timestep, encoder_hidden_states=text) finally: for hook in hooks: hook.remove() @@ -216,16 +227,16 @@ 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 + _, timestep, frame_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] ) ref_proj = ref_proj.unflatten(2, (6, -1)) - temb, proj, text_out = native.condition_embedder(timestep, text) + temb, proj, text_out = native.condition_embedder(frame_timestep, text) for name, ours, theirs in ( - ("temb", temb, ref_temb), - ("timestep_proj", proj, ref_proj), + ("temb", _per_token(temb), ref_temb), + ("timestep_proj", _per_token(proj), ref_proj), ("text_proj", text_out, ref_text), ): diff = (ours.float() - theirs.float()).abs().max().item() diff --git a/test/modular/test_wan22_reference_equivalence.py b/test/modular/test_wan22_reference_equivalence.py index 250db664c1..d780a4cb83 100644 --- a/test/modular/test_wan22_reference_equivalence.py +++ b/test/modular/test_wan22_reference_equivalence.py @@ -449,7 +449,7 @@ def test_sequential_cfg_trajectory_is_bitexact_vs_oracle(oracle_meta, text_embed """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 + per-frame 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).""" meta = _step_metadata(oracle_meta) @@ -507,7 +507,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(): From 7a7a2b1808f51b9b76ee199dec1739c5616888dc Mon Sep 17 00:00:00 2001 From: merceod Date: Thu, 29 Oct 2026 19:41:04 +0000 Subject: [PATCH 03/13] wan22: fuse the self-attn q/k RMSNorm and RoPE The slice-assign rope made Inductor write eight fp32 intermediates per block. Eager now uses the shared stacked rotation (same bits), and the compiled DiT runs q and k through one kernel. --- .../components/diffusion/qk_norm_rope.py | 95 +++++++++++++++++++ mstar/model/wan22/components/dit.py | 40 +++----- 2 files changed, 107 insertions(+), 28 deletions(-) create mode 100644 mstar/model/components/diffusion/qk_norm_rope.py 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 803043d8dd..ba8b467b5d 100644 --- a/mstar/model/wan22/components/dit.py +++ b/mstar/model/wan22/components/dit.py @@ -39,6 +39,7 @@ 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 @@ -121,7 +122,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) @@ -132,27 +133,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 @@ -177,17 +162,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). @@ -351,7 +335,7 @@ def forward( 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) From 97494425b0ef5474e7c079d00d2c424c96c0a085 Mon Sep 17 00:00:00 2001 From: merceod Date: Thu, 29 Oct 2026 19:41:04 +0000 Subject: [PATCH 04/13] wan22: keep the denoise step free of host syncs The step index comes from the loop counter, the UniPC tables go up through pinned memory, and the corrector solve skips its host-side info check, so the worker queues each step behind the previous one. --- mstar/model/wan22/components/unipc.py | 4 +++- mstar/model/wan22/submodules.py | 22 +++++++++++++++++----- 2 files changed, 20 insertions(+), 6 deletions(-) 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/submodules.py b/mstar/model/wan22/submodules.py index d18de914a0..663530ca2d 100644 --- a/mstar/model/wan22/submodules.py +++ b/mstar/model/wan22/submodules.py @@ -169,6 +169,16 @@ def _loop_step(fwd_info: CurrentForwardPassInfo, time_index: torch.Tensor) -> in return int(time_index.reshape(-1)[0].item()) +def _to_device_without_sync(t: torch.Tensor, device: torch.device) -> torch.Tensor: + """``t.to(device)`` for a small host table. Nothing in the denoise step + may wait on the GPU, or the worker cannot queue the step behind the one + still running: a blocking copy from pageable memory syncs the stream, so + the table goes up through pinned memory instead (the same bits).""" + if device.type == "cuda": + return t.pin_memory().to(device, non_blocking=True) + return t.to(device) + + def _latent_grid(config: Wan22Config, step_metadata: dict) -> tuple[int, int, int]: """(T_lat, H_lat, W_lat) for the request's pixel-space dims (diffusers ``prepare_latents``: T=(F-1)//4+1, spatial /16).""" @@ -459,15 +469,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, time_index) 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_without_sync(sigmas, device) + t = _to_device_without_sync(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) @@ -493,7 +505,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 +513,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), ) From 9a7b20be4cc66268d20dfe0f854478f6e52e21c0 Mon Sep 17 00:00:00 2001 From: merceod Date: Thu, 29 Oct 2026 19:41:04 +0000 Subject: [PATCH 05/13] dense attn: don't replan an unchanged layout --- mstar/engine/resources/attn/dense.py | 32 ++++++++++++++++--------- test/modular/test_dense_attention.py | 36 ++++++++++++++++++++++++++-- 2 files changed, 55 insertions(+), 13 deletions(-) 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/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)])) From d4871e866dc1daf90fcfbcf3a4d23e2798c73720 Mon Sep 17 00:00:00 2001 From: merceod Date: Thu, 29 Oct 2026 19:41:04 +0000 Subject: [PATCH 06/13] tests and docs for the DiT kernel components --- docs/adding_models.rst | 10 ++ test/modular/test_diffusion_dit_kernels.py | 128 +++++++++++++++++++++ 2 files changed, 138 insertions(+) create mode 100644 test/modular/test_diffusion_dit_kernels.py diff --git a/docs/adding_models.rst b/docs/adding_models.rst index a83a719014..d4192771a6 100644 --- a/docs/adding_models.rst +++ b/docs/adding_models.rst @@ -2295,6 +2295,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/test/modular/test_diffusion_dit_kernels.py b/test/modular/test_diffusion_dit_kernels.py new file mode 100644 index 0000000000..57534e1f21 --- /dev/null +++ b/test/modular/test_diffusion_dit_kernels.py @@ -0,0 +1,128 @@ +"""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.components.diffusion.rope import apply_rotary_interleaved + +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 test_eager_rms_norm_rope_is_the_modules_bitwise(): + g = torch.Generator(device="cpu").manual_seed(0) + heads, head_dim, seq = 3, 32, 20 + q = torch.randn(1, seq, heads * head_dim, generator=g).bfloat16() + k = torch.randn(1, seq, heads * head_dim, generator=g).bfloat16() + q_norm = nn.RMSNorm(heads * head_dim, eps=EPS, dtype=torch.bfloat16) + k_norm = nn.RMSNorm(heads * head_dim, eps=EPS, dtype=torch.bfloat16) + ang = torch.rand(seq, head_dim // 2, generator=g) * 6 + cos, sin = ang.cos().repeat_interleave(2, 1), ang.sin().repeat_interleave(2, 1) + q_out, k_out = qk_norm_rope.rms_norm_rope(q, k, q_norm, k_norm, cos, sin, heads) + assert torch.equal(q_out, apply_rotary_interleaved(q_norm(q).unflatten(2, (heads, -1)), cos, sin)) + assert torch.equal(k_out, apply_rotary_interleaved(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) From 93bc40cc41247ebdd848d86218bb9596ec271b04 Mon Sep 17 00:00:00 2001 From: merceod Date: Fri, 30 Oct 2026 12:17:21 +0000 Subject: [PATCH 07/13] wan22: forward stages its tables with to_device_nonblocking --- mstar/model/wan22/submodules.py | 15 +++------------ 1 file changed, 3 insertions(+), 12 deletions(-) diff --git a/mstar/model/wan22/submodules.py b/mstar/model/wan22/submodules.py index dbdddd0bf9..db06616a97 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__) @@ -169,16 +170,6 @@ def _loop_step(fwd_info: CurrentForwardPassInfo, time_index: torch.Tensor) -> in return int(time_index.reshape(-1)[0].item()) -def _to_device_without_sync(t: torch.Tensor, device: torch.device) -> torch.Tensor: - """``t.to(device)`` for a small host table. Nothing in the denoise step - may wait on the GPU, or the worker cannot queue the step behind the one - still running: a blocking copy from pageable memory syncs the stream, so - the table goes up through pinned memory instead (the same bits).""" - if device.type == "cuda": - return t.pin_memory().to(device, non_blocking=True) - return t.to(device) - - def _latent_grid(config: Wan22Config, step_metadata: dict) -> tuple[int, int, int]: """(T_lat, H_lat, W_lat) for the request's pixel-space dims (diffusers ``prepare_latents``: T=(F-1)//4+1, spatial /16).""" @@ -478,8 +469,8 @@ def forward( with _inference_ctx(): sigmas, timesteps = make_unipc_tables(num_steps, self.config.flow_shift) - sigmas_device = _to_device_without_sync(sigmas, device) - t = _to_device_without_sync(timesteps, device)[k] + 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) From 1af7dfe1ba2b874dee07ceae648b309bcfb8a8e7 Mon Sep 17 00:00:00 2001 From: merceod Date: Fri, 30 Oct 2026 12:17:21 +0000 Subject: [PATCH 08/13] wan22: k comes from the loop counter, time_index edge goes Nothing on the served path read the edge anymore, it was only there for the .item() fallback. A caller with no counter now gets a KeyError instead of the sync. --- mstar/model/wan22/submodules.py | 32 ++++------ mstar/model/wan22/wan22_model.py | 7 +-- test/modular/test_step_prep_host_syncs.py | 58 +++++++++++-------- test/modular/test_wan22_dit_compile_sizes.py | 7 ++- test/modular/test_wan22_model.py | 13 ++--- .../test_wan22_reference_equivalence.py | 16 +++-- 6 files changed, 68 insertions(+), 65 deletions(-) diff --git a/mstar/model/wan22/submodules.py b/mstar/model/wan22/submodules.py index db06616a97..aaf4e1343f 100644 --- a/mstar/model/wan22/submodules.py +++ b/mstar/model/wan22/submodules.py @@ -156,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]: @@ -336,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 @@ -419,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 @@ -442,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) @@ -452,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, @@ -464,7 +453,7 @@ def forward( step_metadata = fwd_info.step_metadata num_steps = int(step_metadata["num_inference_steps"]) guidance_scale = float(step_metadata["guidance_scale"]) - k = _loop_step(fwd_info, time_index) + k = _loop_step(fwd_info) device = latents.device with _inference_ctx(): @@ -516,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_step_prep_host_syncs.py b/test/modular/test_step_prep_host_syncs.py index cadebdade1..d664fded93 100644 --- a/test/modular/test_step_prep_host_syncs.py +++ b/test/modular/test_step_prep_host_syncs.py @@ -1,19 +1,19 @@ """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. """ @@ -100,29 +100,41 @@ 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)) # ---------------------------------------------------------------- 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_reference_equivalence.py b/test/modular/test_wan22_reference_equivalence.py index 9a51167e40..dd180f861e 100644 --- a/test/modular/test_wan22_reference_equivalence.py +++ b/test/modular/test_wan22_reference_equivalence.py @@ -44,6 +44,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, ) @@ -332,10 +333,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()) @@ -467,10 +469,11 @@ def test_sequential_cfg_trajectory_is_bitexact_vs_oracle(oracle_meta, text_embed 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") } 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()) @@ -574,10 +577,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]) @@ -776,10 +780,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() @@ -855,10 +860,11 @@ def test_no_cfg_batch1_trajectory_is_bitexact(oracle_meta, text_embeds, dit_subm 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") } 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()) From 4090bb6dcb8053184496d66ce5fe0be077c3ee98 Mon Sep 17 00:00:00 2001 From: merceod Date: Fri, 30 Oct 2026 12:17:21 +0000 Subject: [PATCH 09/13] wan22: CUDA test for the sync-free step forward with a stub transformer under sync-debug mode, so linalg.solve, a host read of a device tensor or a pageable table copy fails it. --- test/modular/test_step_prep_host_syncs.py | 62 ++++++++++++++++++++++- 1 file changed, 61 insertions(+), 1 deletion(-) diff --git a/test/modular/test_step_prep_host_syncs.py b/test/modular/test_step_prep_host_syncs.py index d664fded93..145a5fbb43 100644 --- a/test/modular/test_step_prep_host_syncs.py +++ b/test/modular/test_step_prep_host_syncs.py @@ -15,7 +15,8 @@ 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 @@ -137,6 +139,64 @@ def test_wan22_step_prep_fails_without_the_loop_counter(): 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 def _omnivoice(): From f2b23ab4de1919915a3682d57b88cffeabf043dc Mon Sep 17 00:00:00 2001 From: merceod Date: Fri, 30 Oct 2026 12:17:21 +0000 Subject: [PATCH 10/13] wan22: per-frame time embedding is rounding-close, not bitwise cuBLAS can pick another kernel at 18 rows than at 7020. The native test holds those rows to that rounding and checks the blocks bitwise on the same rows. The reference suite prices the trajectory and keeps a 0.0 gate with the embedder at the reference's row count. Also fixes its I2V test, which still expected a per-token timestep. --- mstar/model/wan22/components/dit.py | 21 +-- mstar/model/wan22/config.py | 2 +- .../test_wan22_native_dit_equivalence.py | 73 ++++++---- .../test_wan22_reference_equivalence.py | 126 +++++++++++------- 4 files changed, 137 insertions(+), 85 deletions(-) diff --git a/mstar/model/wan22/components/dit.py b/mstar/model/wan22/components/dit.py index ba8b467b5d..f990e9716f 100644 --- a/mstar/model/wan22/components/dit.py +++ b/mstar/model/wan22/components/dit.py @@ -1,13 +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 (``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. Per-frame rows -are the reference's per-token rows bit for bit (the rows of a GEMM do not -depend on how many there are). +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): @@ -184,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 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/test/modular/test_wan22_native_dit_equivalence.py b/test/modular/test_wan22_native_dit_equivalence.py index 099b6bde85..11c7d4f48e 100644 --- a/test/modular/test_wan22_native_dit_equivalence.py +++ b/test/modular/test_wan22_native_dit_equivalence.py @@ -1,19 +1,24 @@ """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 time embedding and modulation rows are the exception. The native module computes +them per latent frame and the reference per token, and cuBLAS may pick another GEMM +kernel at the other row count, so those rows are held to that kernel's rounding. The +stage comparison gives both modules the same rows (the native ones, expanded per token +for the reference), so it stays bitwise. 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. The native module takes that grid one value per latent frame and -computes its time embedding per frame; its rows are compared per token. Skipped +special case. The native module takes that grid one value per latent frame. Skipped without CUDA or the checkpoint in the local HF cache. """ @@ -64,6 +69,14 @@ def _hf_cache_has_checkpoint() -> bool: LATENT_SHAPE = (2, 48, 9, 30, 52) TOKENS_PER_FRAME = (30 // 2) * (52 // 2) +# Per-frame (18) vs per-token (7020) rows, about twice the worst seen. temb moves by +# 9.8e-4 on H100, where cuBLAS picks another kernel for the time embedder's second +# Linear. On B300 the rows are bitwise under the default numerics and move by 3.1e-5 +# (modulation 3.9e-3) with NVIDIA_TF32_OVERRIDE=0; switching TF32 off in the time embedder +# alone moves temb by the same 9.8e-4 and the modulation rows by 1.6e-2. +TEMB_ROWS_MAX_ABS = 2e-3 +TIMESTEP_PROJ_ROWS_MAX_ABS = 3.2e-2 + @pytest.fixture(scope="module") def wrapped(): @@ -172,40 +185,52 @@ def _first_tensor(output): def test_per_layer_outputs_bitwise_identical(wrapped, native, 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 = [] - def capture(side, module, transform=None): + def capture(side, module): def hook(_module, _args, output): - out = _first_tensor(output).detach() - stages[side].append(transform(out) if transform else out) + 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, _per_token if side == "native" else None) 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=frame_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):") @@ -224,9 +249,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.""" +def test_condition_embedder_rows_match_within_kernel_rounding(wrapped, native, inputs): + """Per-frame rows against the reference's per-token rows, to the rounding of a + different GEMM kernel. The projected text stream has the same row count on both + sides, so it stays bitwise.""" _, timestep, frame_timestep, text = inputs with torch.no_grad(): ref_temb, ref_proj, ref_text, _ = wrapped.condition_embedder( @@ -234,11 +260,12 @@ def test_condition_embedder_all_outputs_bitwise_identical(wrapped, native, input ) ref_proj = ref_proj.unflatten(2, (6, -1)) temb, proj, text_out = native.condition_embedder(frame_timestep, text) - for name, ours, theirs in ( - ("temb", _per_token(temb), ref_temb), - ("timestep_proj", _per_token(proj), ref_proj), - ("text_proj", text_out, ref_text), + for name, ours, theirs, bound in ( + ("temb", _per_token(temb), ref_temb, TEMB_ROWS_MAX_ABS), + ("timestep_proj", _per_token(proj), ref_proj, TIMESTEP_PROJ_ROWS_MAX_ABS), + ("text_proj", text_out, ref_text, 0.0), ): + 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})" + assert diff <= bound, f"condition_embedder {name} diverges (max_abs {diff:.3e}, bound {bound:.1e})" diff --git a/test/modular/test_wan22_reference_equivalence.py b/test/modular/test_wan22_reference_equivalence.py index dd180f861e..57bd58e45d 100644 --- a/test/modular/test_wan22_reference_equivalence.py +++ b/test/modular/test_wan22_reference_equivalence.py @@ -65,18 +65,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. @@ -95,11 +103,12 @@ # 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_TRAJECTORY_MAX_ABS = 0.0 +# No-CFG (guidance_scale <= 1.0) batch-1 path vs a no-CFG oracle: one batch-1 forward +# per step, so the same departure as the sequential-CFG trajectory, over the oracle's +# 8 small-grid steps. 0.48 on B300 and 0.76 with the perturbation above; ~1.3x that. +NO_CFG_TRAJECTORY_MAX_ABS = 1.0 CUDA_AVAILABLE = torch.cuda.is_available() @@ -240,13 +249,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") @@ -447,28 +453,22 @@ 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-frame 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 = { @@ -479,15 +479,46 @@ def test_sequential_cfg_trajectory_is_bitexact_vs_oracle(oracle_meta, text_embed diffs.append((outputs["latents"][0].cpu().float() - ref.float()).abs().max().item()) finally: Wan22DitSubmodule._noise_prediction = original + return diffs + +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 per-frame time embedding's rows are held to + the reference's 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.""" + embedder = dit_submodule.transformer.condition_embedder + per_frame_rows = embedder.forward + tokens_per_frame = (oracle_meta["height"] // 32) * (oracle_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: + diffs = _sequential_trajectory_diffs(oracle_meta, text_embeds, dit_submodule) + finally: + del embedder.forward # restore class-method lookup + line = _report_trajectory("sequential-CFG trajectory, per-token rows", diffs) + assert max(diffs) <= SEQ_PER_TOKEN_ROWS_MAX_ABS, line # --------------------------------------------------------------------------- @@ -551,7 +582,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. """ @@ -602,11 +633,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 # --------------------------------------------------------------------------- @@ -738,8 +769,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 @@ -810,7 +841,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. # --------------------------------------------------------------------------- @@ -822,12 +853,11 @@ 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_matches_oracle(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 + and hold every step to the same bound as the sequential-CFG gate, 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. """ @@ -879,11 +909,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 From 3f492412d98040bf91d4aa0777776295a0d6540d Mon Sep 17 00:00:00 2001 From: merceod Date: Fri, 30 Oct 2026 12:17:21 +0000 Subject: [PATCH 11/13] test: eager rope against the slice-write rotation it replaced --- test/modular/test_diffusion_dit_kernels.py | 40 ++++++++++++++++------ 1 file changed, 29 insertions(+), 11 deletions(-) diff --git a/test/modular/test_diffusion_dit_kernels.py b/test/modular/test_diffusion_dit_kernels.py index 57534e1f21..35f2a719ab 100644 --- a/test/modular/test_diffusion_dit_kernels.py +++ b/test/modular/test_diffusion_dit_kernels.py @@ -16,7 +16,8 @@ from torch import nn from mstar.model.components.diffusion import modulation, qk_norm_rope -from mstar.model.components.diffusion.rope import apply_rotary_interleaved +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( @@ -61,18 +62,35 @@ def test_eager_formulation_is_the_inline_block_math_bitwise(): assert torch.equal(x_out, plain) and torch.equal(out, plain_mod) -def test_eager_rms_norm_rope_is_the_modules_bitwise(): +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) - heads, head_dim, seq = 3, 32, 20 - q = torch.randn(1, seq, heads * head_dim, generator=g).bfloat16() - k = torch.randn(1, seq, heads * head_dim, generator=g).bfloat16() - q_norm = nn.RMSNorm(heads * head_dim, eps=EPS, dtype=torch.bfloat16) - k_norm = nn.RMSNorm(heads * head_dim, eps=EPS, dtype=torch.bfloat16) - ang = torch.rand(seq, head_dim // 2, generator=g) * 6 - cos, sin = ang.cos().repeat_interleave(2, 1), ang.sin().repeat_interleave(2, 1) + 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, apply_rotary_interleaved(q_norm(q).unflatten(2, (heads, -1)), cos, sin)) - assert torch.equal(k_out, apply_rotary_interleaved(k_norm(k).unflatten(2, (heads, -1)), cos, sin)) + 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): From da66a31830f5b3aa63d1a3dccd8aab58d6c46e38 Mon Sep 17 00:00:00 2001 From: Gaurav Shah Date: Fri, 30 Oct 2026 16:02:56 -0400 Subject: [PATCH 12/13] test: hold the wan22 time embedder bitwise on per-token rows The two bounds were set from the worst case seen: one sat below a bf16 step at temb's top, and the other let a moved bf16 boundary through. --- .../test_wan22_native_dit_equivalence.py | 40 ++++++++----------- 1 file changed, 16 insertions(+), 24 deletions(-) diff --git a/test/modular/test_wan22_native_dit_equivalence.py b/test/modular/test_wan22_native_dit_equivalence.py index 11c7d4f48e..d3cfaa5220 100644 --- a/test/modular/test_wan22_native_dit_equivalence.py +++ b/test/modular/test_wan22_native_dit_equivalence.py @@ -7,11 +7,11 @@ pass — a nonzero there is a port bug, and the printed table names the first divergent stage. -The time embedding and modulation rows are the exception. The native module computes -them per latent frame and the reference per token, and cuBLAS may pick another GEMM -kernel at the other row count, so those rows are held to that kernel's rounding. The -stage comparison gives both modules the same rows (the native ones, expanded per token -for the reference), so it stays bitwise. +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. @@ -69,14 +69,6 @@ def _hf_cache_has_checkpoint() -> bool: LATENT_SHAPE = (2, 48, 9, 30, 52) TOKENS_PER_FRAME = (30 // 2) * (52 // 2) -# Per-frame (18) vs per-token (7020) rows, about twice the worst seen. temb moves by -# 9.8e-4 on H100, where cuBLAS picks another kernel for the time embedder's second -# Linear. On B300 the rows are bitwise under the default numerics and move by 3.1e-5 -# (modulation 3.9e-3) with NVIDIA_TF32_OVERRIDE=0; switching TF32 off in the time embedder -# alone moves temb by the same 9.8e-4 and the modulation rows by 1.6e-2. -TEMB_ROWS_MAX_ABS = 2e-3 -TIMESTEP_PROJ_ROWS_MAX_ABS = 3.2e-2 - @pytest.fixture(scope="module") def wrapped(): @@ -249,23 +241,23 @@ def use_native_rows(_module, _args, output): ) -def test_condition_embedder_rows_match_within_kernel_rounding(wrapped, native, inputs): - """Per-frame rows against the reference's per-token rows, to the rounding of a - different GEMM kernel. The projected text stream has the same row count on both - sides, so it stays bitwise.""" - _, timestep, frame_timestep, text = inputs +def test_condition_embedder_all_outputs_bitwise_identical(wrapped, native, 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] ) ref_proj = ref_proj.unflatten(2, (6, -1)) - temb, proj, text_out = native.condition_embedder(frame_timestep, text) - for name, ours, theirs, bound in ( - ("temb", _per_token(temb), ref_temb, TEMB_ROWS_MAX_ABS), - ("timestep_proj", _per_token(proj), ref_proj, TIMESTEP_PROJ_ROWS_MAX_ABS), - ("text_proj", text_out, ref_text, 0.0), + temb, proj, text_out = native.condition_embedder(timestep, text) + for name, ours, theirs in ( + ("temb", temb, ref_temb), + ("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 diff <= bound, f"condition_embedder {name} diverges (max_abs {diff:.3e}, bound {bound:.1e})" + assert torch.equal(ours, theirs), f"condition_embedder {name} diverges (max_abs {diff:.3e})" From f10967434ffa588d087c8e92df1c67223102385a Mon Sep 17 00:00:00 2001 From: Gaurav Shah Date: Fri, 30 Oct 2026 16:02:56 -0400 Subject: [PATCH 13/13] test: no-cfg wan22 trajectory is bitwise on per-token rows --- .../test_wan22_reference_equivalence.py | 52 +++++++++++-------- 1 file changed, 31 insertions(+), 21 deletions(-) diff --git a/test/modular/test_wan22_reference_equivalence.py b/test/modular/test_wan22_reference_equivalence.py index 57bd58e45d..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 @@ -105,10 +106,10 @@ # 3.675 (H200). Bound = worst observed x ~1.3 headroom. The eager path stays the # 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: one batch-1 forward -# per step, so the same departure as the sequential-CFG trajectory, over the oracle's -# 8 small-grid steps. 0.48 on B300 and 0.76 with the perturbation above; ~1.3x that. -NO_CFG_TRAJECTORY_MAX_ABS = 1.0 +# 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() @@ -482,6 +483,25 @@ def _sequential_trajectory_diffs(oracle_meta, text_embeds, dit_submodule) -> lis 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] line = f"{label} vs oracle: max_abs={max(diffs):.3e}, first nonzero step={nonzero[0] if nonzero else None}" @@ -492,8 +512,8 @@ def _report_trajectory(label: str, diffs: list[float]) -> str: 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 per-frame time embedding's rows are held to - the reference's in test_wan22_native_dit_equivalence.py.""" + -> 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 @@ -504,19 +524,8 @@ def test_sequential_cfg_trajectory_is_bitexact_on_per_token_rows(oracle_meta, te 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.""" - embedder = dit_submodule.transformer.condition_embedder - per_frame_rows = embedder.forward - tokens_per_frame = (oracle_meta["height"] // 32) * (oracle_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: + with _per_token_rows(dit_submodule.transformer, oracle_meta): diffs = _sequential_trajectory_diffs(oracle_meta, text_embeds, dit_submodule) - finally: - del embedder.forward # restore class-method lookup line = _report_trajectory("sequential-CFG trajectory, per-token rows", diffs) assert max(diffs) <= SEQ_PER_TOKEN_ROWS_MAX_ABS, line @@ -853,10 +862,11 @@ 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_matches_oracle(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 hold every step to the same bound as the sequential-CFG gate, since a no-CFG + 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. @@ -888,7 +898,7 @@ def test_no_cfg_batch1_trajectory_matches_oracle(oracle_meta, text_embeds, dit_s 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})