[diffusion] Accelerate Cosmos3 T2I QKNorm+RoPE (#34932)

2026-08-16 20:15:32 +08:00
parent 095ec6c997
commit 41abbb0d32
3 changed files with 135 additions and 2 deletions
@@ -231,6 +231,8 @@ def _apply_qwen3_qk_norm_rope(
    head_dim: int,
    cos_sin_cache: torch.Tensor,
    rope_cache_positions: torch.Tensor,
+    *,
+    round_norm_before_rope: bool = False,
 ) -> tuple[torch.Tensor, torch.Tensor]:
    return apply_qk_norm_rope(
        q=q,
@@ -242,6 +244,7 @@ def _apply_qwen3_qk_norm_rope(
        is_neox=True,
        positions=rope_cache_positions,
        allow_strided_qk=True,
+        round_norm_before_rope=round_norm_before_rope,
    )


@@ -256,6 +259,8 @@ def _apply_qwen3_qk_norm_rope_pack_kv(
    head_dim: int,
    cos_sin_cache: torch.Tensor,
    rope_cache_positions: torch.Tensor,
+    *,
+    round_norm_before_rope: bool = False,
 ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
    batch_size, suffix_tokens, _, _ = q.shape
    prefix_tokens = k_prefix.shape[1]
@@ -283,6 +288,7 @@ def _apply_qwen3_qk_norm_rope_pack_kv(
        eps=q_norm.variance_epsilon,
        head_dim=head_dim,
        rope_dim=cos_sin_cache.shape[-1],
+        round_norm_before_rope=round_norm_before_rope,
    )
    return q, packed_kv[0], packed_kv[1]

@@ -746,6 +752,7 @@ class Cosmos3CrossAttention(nn.Module):
        cos_sin_cache: torch.Tensor,
        rope_cache_positions: torch.Tensor,
        use_fused_qk_norm_rope: bool,
+        round_norm_before_rope: bool = False,
    ) -> torch.Tensor:
        """Cross-attention from GEN to cached UND K/V.

@@ -795,6 +802,7 @@ class Cosmos3CrossAttention(nn.Module):
                True,
                q.dtype,
                cos_sin_cache.dtype,
+                round_norm_before_rope=round_norm_before_rope,
                pack_kv=True,
            )
        )
@@ -810,6 +818,7 @@ class Cosmos3CrossAttention(nn.Module):
                self.head_dim,
                cos_sin_cache,
                rope_cache_positions,
+                round_norm_before_rope=round_norm_before_rope,
            )
            out = self.attn.forward(q, packed_k, packed_v)
        elif use_fused_qk_norm_rope:
@@ -821,6 +830,7 @@ class Cosmos3CrossAttention(nn.Module):
                self.head_dim,
                cos_sin_cache,
                rope_cache_positions,
+                round_norm_before_rope=round_norm_before_rope,
            )
        else:
            q, k = _apply_qwen3_qk_norm_rope_split(
@@ -960,6 +970,7 @@ class Cosmos3GenDecoderLayer(nn.Module):
        cos_sin_cache: torch.Tensor,
        rope_cache_positions: torch.Tensor,
        use_fused_qk_norm_rope: bool,
+        round_norm_before_rope: bool = False,
        residual: torch.Tensor | None = None,
    ) -> tuple[torch.Tensor, torch.Tensor]:
        # Fused add+rmsnorm: each `(hidden_states, residual) = norm(...)`
@@ -979,6 +990,7 @@ class Cosmos3GenDecoderLayer(nn.Module):
            cos_sin_cache,
            rope_cache_positions,
            use_fused_qk_norm_rope,
+            round_norm_before_rope,
        )

        hidden_states, residual = self.post_attention_layernorm(hidden_states, residual)
@@ -1126,6 +1138,7 @@ class Cosmos3OmniTransformer(CachableDiT, LayerwiseOffloadableModuleMixin):
        self.rms_norm_eps = arch.rms_norm_eps
        self.hidden_act = arch.hidden_act
        self.rope_theta = arch.rope_theta
+        self._gen_layers_torch_compiled = False

        # The checkpoint may override the activation (and thus the MLP weight
        # layout), so bind the arch-derived mappings on the instance.
@@ -1618,11 +1631,21 @@ class Cosmos3OmniTransformer(CachableDiT, LayerwiseOffloadableModuleMixin):
                vis_pos_ids = vis_pos_ids.view(
                    3, batch_size, self.sp_size, local_seq_len
                )[:, :, self.sp_rank, :]
-            self.cached_gen_rope_inputs[cache_key] = (
+            cos_sin_gen, gen_rope_cache_positions = (
                self.language_model.rotary_emb.build_rope_cache_inputs(
                    vis_pos_ids, cache_dtype=hidden_gen.dtype
                )
            )
+            if T == 1 and not self._gen_layers_torch_compiled:
+                # build_rope_cache_inputs already rounds through cache_dtype
+                # before returning FP32 storage. Keep that rounded cache in the
+                # activation dtype so the exact fused QKNorm+RoPE kernel can
+                # consume it without repeating the cast in every GEN layer.
+                cos_sin_gen = cos_sin_gen.to(hidden_gen.dtype)
+            self.cached_gen_rope_inputs[cache_key] = (
+                cos_sin_gen,
+                gen_rope_cache_positions,
+            )

        cos_sin_gen, gen_rope_cache_positions = self.cached_gen_rope_inputs[cache_key]

@@ -1631,7 +1654,22 @@ class Cosmos3OmniTransformer(CachableDiT, LayerwiseOffloadableModuleMixin):
        # fused add+rmsnorm path instead of separate add + norm kernels.
        cached_kv_for_key = self.cached_kv[cache_key]
        residual: torch.Tensor | None = None
-        use_fused_qk_norm_rope = T > 1
+        round_norm_before_rope = T == 1
+        use_fused_qk_norm_rope = T > 1 or (
+            hidden_gen.device.type == "cuda"
+            and not torch.compiler.is_compiling()
+            and not self._gen_layers_torch_compiled
+            and get_sp_world_size() == 1
+            and can_use_fused_inplace_qknorm_rope(
+                self.head_dim,
+                cos_sin_gen.shape[-1],
+                True,
+                hidden_gen.dtype,
+                cos_sin_gen.dtype,
+                round_norm_before_rope=True,
+                pack_kv=True,
+            )
+        )
        for i, layer in enumerate(self.gen_layers):
            k_und, v_und = cached_kv_for_key[i]
            hidden_gen, residual = layer(
@@ -1641,6 +1679,7 @@ class Cosmos3OmniTransformer(CachableDiT, LayerwiseOffloadableModuleMixin):
                cos_sin_gen,
                gen_rope_cache_positions,
                use_fused_qk_norm_rope,
+                round_norm_before_rope,
                residual=residual,
            )

@@ -870,6 +870,7 @@ class Cosmos3DenoisingStage(PipelineStage):
                len(gen_layers),
                compile_kwargs,
            )
+            transformer._gen_layers_torch_compiled = True
            for i, layer in enumerate(gen_layers):
                gen_layers[i] = torch.compile(layer, **compile_kwargs)
        else: