Wan2.2 step: one-pass adaLN per frame, fused q/k norm+RoPE, no host syncs; Cosmos3: keep the dense attention plan

docs/adding_models.rst

@@ -2346,6 +2346,16 @@ attention over packed tokens: SDPA eagerly, the ragged FlashInfer resource under
 family), ``image_io.py`` (latent packing, PNG) and ``lora.py`` (static LoRA merge at load,
 also into fused ``q``/``k``/``v`` shards via ``FusedColumnLinear.shard_slice``).
 
+Two pieces swap a one-pass kernel in for what Inductor would emit, inside a compiled
+transformer only; eager calls run the plain PyTorch formulation, which is what a port's
+bit-exact tests hold. ``modulation.ada_layer_norm`` is a block's gated residual add,
+LayerNorm and shift/scale modulation, with the modulation given per group of tokens
+(``[B, G, D]``: one group per sample, or per latent frame as in Wan2.2); Inductor reads a
+LayerNorm row wider than 1024 twice. ``qk_norm_rope.rms_norm_rope`` is the q/k RMSNorm
+across heads followed by the interleaved-pair RoPE, q and k in one launch. Both agree
+with eager to bf16 rounding rather than bitwise. Wan2.2 is the first user;
+``test/modular/test_diffusion_dit_kernels.py`` holds them to the eager form.
+
 On the API side ``DiffusionImageAdapter`` (``mstar/api_server/openai/adapters.py``) serves
 ``/v1/images/generations`` and ``/v1/images/edits`` for any such model (``size`` becomes
 ``width``/``height``, ``seed`` and other knobs pass through); register it under the

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

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,
+    )

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)