[diffusion] optimize: cosmos3 fused qknorm rope (#27096)

python/sglang/multimodal_gen/runtime/layers/layernorm.py

@@ -869,15 +869,27 @@ def apply_qk_norm_rope(
         raise ValueError(
             f"apply_qk_norm_rope expects 4D q/k tensors, got q:{tuple(q.shape)} k:{tuple(k.shape)}"
         )
-    if q.shape != k.shape:
+    if q.shape[:2] != k.shape[:2] or q.shape[-1] != k.shape[-1]:
         raise ValueError(
-            f"apply_qk_norm_rope expects q/k to have the same shape, got {q.shape} vs {k.shape}"
+            "apply_qk_norm_rope expects q/k to share batch, sequence, and head size, "
+            f"got {q.shape} vs {k.shape}"
+        )
+    if not (isinstance(cos_sin_cache, torch.Tensor) and cos_sin_cache.dim() == 2):
+        raise ValueError("cos_sin_cache must be a 2D torch.Tensor")
+    if k.device != q.device or cos_sin_cache.device != q.device:
+        raise ValueError(
+            "q, k, and cos_sin_cache must be on the same device, "
+            f"got q={q.device}, k={k.device}, cos_sin_cache={cos_sin_cache.device}"
         )
 
     batch_size, seq_len, _, _ = q.shape
     q_eps = q_norm.variance_epsilon
     k_eps = k_norm.variance_epsilon
     rope_dim = cos_sin_cache.size(-1)
+    if rope_dim % 2 != 0 or rope_dim > head_dim:
+        raise ValueError(
+            f"cos_sin_cache width must be even and <= head_dim, got {rope_dim} vs {head_dim}"
+        )
     fused_enabled = os.getenv("SGLANG_ENABLE_FUSED_QKNORM_ROPE", "1").lower() not in {
         "0",
         "false",
@@ -898,6 +910,7 @@ def apply_qk_norm_rope(
             raise ValueError(
                 f"positions must be 1D of length {batch_size * seq_len}, got shape={tuple(positions.shape)}"
             )
+        positions = positions.to(device=q.device, dtype=torch.long)
 
     if (
         fused_enabled

python/sglang/multimodal_gen/runtime/layers/rotary_embedding/mrope.py

@@ -123,11 +123,7 @@ class Qwen3VLTextRotaryEmbedding(torch.nn.Module):
             freqs_t[..., idx] = freqs[dim, ..., idx]
         return freqs_t
 
-    @torch.no_grad()
-    def forward(
-        self, x: torch.Tensor, position_ids: torch.Tensor
-    ) -> tuple[torch.Tensor, torch.Tensor]:
-        """Return cos/sin for position IDs shaped [3, B, S], [B, S, 3], or [B, S]."""
+    def _normalize_position_ids(self, position_ids: torch.Tensor) -> torch.Tensor:
         if position_ids.ndim == 3 and position_ids.shape[-1] == 3:
             position_ids = position_ids.permute(2, 0, 1)
         elif position_ids.ndim == 2:
@@ -137,6 +133,10 @@ class Qwen3VLTextRotaryEmbedding(torch.nn.Module):
                 "Qwen3 mRoPE position_ids must have shape [3, B, S], [B, S, 3], "
                 f"or [B, S], got {tuple(position_ids.shape)}"
             )
+        return position_ids
+
+    def _compute_interleaved_freqs(self, position_ids: torch.Tensor) -> torch.Tensor:
+        position_ids = self._normalize_position_ids(position_ids)
 
         inv_freq_expanded = (
             self.inv_freq[None, None, :, None]
@@ -147,7 +147,31 @@ class Qwen3VLTextRotaryEmbedding(torch.nn.Module):
         position_ids_expanded = position_ids[:, :, None, :].float()
 
         freqs = (inv_freq_expanded @ position_ids_expanded).transpose(2, 3)
-        freqs = self.apply_interleaved_mrope(freqs, self.mrope_section)
+        return self.apply_interleaved_mrope(freqs, self.mrope_section)
+
+    @torch.no_grad()
+    def build_rope_cache_inputs(
+        self, position_ids: torch.Tensor, *, cache_dtype: torch.dtype | None = None
+    ) -> tuple[torch.Tensor, torch.Tensor]:
+        freqs = self._compute_interleaved_freqs(position_ids)
+        cos = freqs.cos() * self.attention_scaling
+        sin = freqs.sin() * self.attention_scaling
+        if cache_dtype is not None and cache_dtype != torch.float32:
+            cos = cos.to(cache_dtype).float()
+            sin = sin.to(cache_dtype).float()
+        cos_sin_cache = torch.cat((cos, sin), dim=-1).reshape(-1, self.head_dim)
+        cos_sin_cache = cos_sin_cache.contiguous()
+        cache_positions = torch.arange(
+            cos_sin_cache.shape[0], device=cos_sin_cache.device, dtype=torch.long
+        )
+        return cos_sin_cache, cache_positions
+
+    @torch.no_grad()
+    def forward(
+        self, x: torch.Tensor, position_ids: torch.Tensor
+    ) -> tuple[torch.Tensor, torch.Tensor]:
+        """Return cos/sin for position IDs shaped [3, B, S], [B, S, 3], or [B, S]."""
+        freqs = self._compute_interleaved_freqs(position_ids)
         emb = torch.cat((freqs, freqs), dim=-1)
         cos = emb.cos() * self.attention_scaling
         sin = emb.sin() * self.attention_scaling

python/sglang/multimodal_gen/runtime/layers/rotary_embedding/utils.py

@@ -80,58 +80,91 @@ def apply_flashinfer_rope_qk_inplace(
             f"Expected q/k to be 4D [bsz, seqlen, nheads, head_size], "
             f"got q:{tuple(q.shape)} k:{tuple(k.shape)}"
         )
-    if q.shape != k.shape:
+    if q.shape[:2] != k.shape[:2] or q.shape[-1] != k.shape[-1]:
         raise ValueError(
-            f"q and k must have the same shape, got {q.shape} vs {k.shape}"
+            f"q and k must share batch, sequence, and head size, got {q.shape} vs {k.shape}"
         )
 
     if not (isinstance(cos_sin_cache, torch.Tensor) and cos_sin_cache.dim() == 2):
         raise ValueError("cos_sin_cache must be a 2D torch.Tensor")
 
-    bsz, seqlen, nheads, d = q.shape
+    bsz, seqlen, q_heads, d = q.shape
+    k_heads = k.shape[2]
+    rope_dim = cos_sin_cache.shape[-1]
+    if k.device != q.device or cos_sin_cache.device != q.device:
+        raise ValueError(
+            "q, k, and cos_sin_cache must be on the same device, "
+            f"got q={q.device}, k={k.device}, cos_sin_cache={cos_sin_cache.device}"
+        )
+    if rope_dim % 2 != 0 or rope_dim > d:
+        raise ValueError(
+            f"cos_sin_cache width must be even and <= head_size, got {rope_dim} vs {d}"
+        )
     if head_size is None:
         head_size = d
     if head_size != d:
         raise ValueError(f"head_size mismatch: inferred {d}, but head_size={head_size}")
 
-    if flashinfer_apply_rope_inplace is None:
-        # Triton fallback for AMD/ROCm where FlashInfer is not available
+    use_flashinfer = (
+        flashinfer_apply_rope_inplace is not None
+        and q.is_cuda
+        and k.is_cuda
+        and cos_sin_cache.is_cuda
+        and q_heads == k_heads
+    )
 
-        _warn_about_missing_flashinfer()
+    if not use_flashinfer:
+        if flashinfer_apply_rope_inplace is None:
+            _warn_about_missing_flashinfer()
 
-        half_size = cos_sin_cache.shape[-1] // 2
+        half_size = rope_dim // 2
         if positions is None:
             cos = cos_sin_cache[:seqlen, :half_size].to(q.dtype)
             sin = cos_sin_cache[:seqlen, half_size:].to(q.dtype)
             cos = cos.unsqueeze(0).expand(bsz, -1, -1).reshape(bsz * seqlen, -1)
             sin = sin.unsqueeze(0).expand(bsz, -1, -1).reshape(bsz * seqlen, -1)
         else:
-            positions = positions.to(cos_sin_cache.device).view(-1)
+            positions = positions.to(device=q.device, dtype=torch.long).view(-1)
             cos = cos_sin_cache[positions, :half_size].to(q.dtype)
             sin = cos_sin_cache[positions, half_size:].to(q.dtype)
-        q_flat = q.reshape(bsz * seqlen, nheads, d)
-        k_flat = k.reshape(bsz * seqlen, nheads, d)
-        q_rot = apply_rotary_embedding(q_flat, cos, sin, interleaved=not is_neox)
-        k_rot = apply_rotary_embedding(k_flat, cos, sin, interleaved=not is_neox)
-        return q_rot.view(bsz, seqlen, nheads, d), k_rot.view(bsz, seqlen, nheads, d)
+
+        def apply_rope_prefix(x: torch.Tensor, num_heads: int) -> torch.Tensor:
+            x_flat = x.reshape(bsz * seqlen, num_heads, d)
+            x_rot = x_flat[..., :rope_dim]
+            out_rot = torch.empty_like(x_rot)
+            cos_b = cos.unsqueeze(-2)
+            sin_b = sin.unsqueeze(-2)
+            if is_neox:
+                x1, x2 = torch.chunk(x_rot, 2, dim=-1)
+                out_rot[..., :half_size] = x1 * cos_b - x2 * sin_b
+                out_rot[..., half_size:] = x2 * cos_b + x1 * sin_b
+            else:
+                x1 = x_rot[..., ::2]
+                x2 = x_rot[..., 1::2]
+                out_rot[..., ::2] = x1 * cos_b - x2 * sin_b
+                out_rot[..., 1::2] = x2 * cos_b + x1 * sin_b
+            if rope_dim == d:
+                return out_rot.view(bsz, seqlen, num_heads, d)
+            out = x_flat.clone()
+            out[..., :rope_dim] = out_rot
+            return out.view(bsz, seqlen, num_heads, d)
+
+        return apply_rope_prefix(q, q_heads), apply_rope_prefix(k, k_heads)
 
     if positions is None:
         pos_1d = torch.arange(seqlen, device=q.device, dtype=torch.long)
         positions = pos_1d if bsz == 1 else pos_1d.repeat(bsz)
     else:
-        if not (
-            isinstance(positions, torch.Tensor)
-            and positions.dtype == torch.long
-            and positions.dim() == 1
-        ):
-            raise ValueError("positions must be a 1D torch.long Tensor")
+        if not (isinstance(positions, torch.Tensor) and positions.dim() == 1):
+            raise ValueError("positions must be a 1D Tensor")
         if positions.numel() != bsz * seqlen:
             raise ValueError(
                 f"positions length must be bsz*seqlen={bsz*seqlen}, got {positions.numel()}"
             )
+        positions = positions.to(device=q.device, dtype=torch.long)
 
-    q_flat = q.reshape(bsz * seqlen, nheads * d).contiguous()
-    k_flat = k.reshape(bsz * seqlen, nheads * d).contiguous()
+    q_flat = q.reshape(bsz * seqlen, q_heads * d).contiguous()
+    k_flat = k.reshape(bsz * seqlen, k_heads * d).contiguous()
     flashinfer_apply_rope_inplace(
         positions=positions,
         query=q_flat,
@@ -140,7 +173,7 @@ def apply_flashinfer_rope_qk_inplace(
         cos_sin_cache=cos_sin_cache,
         is_neox=is_neox,
     )
-    return q_flat.view(bsz, seqlen, nheads, d), k_flat.view(bsz, seqlen, nheads, d)
+    return q_flat.view(bsz, seqlen, q_heads, d), k_flat.view(bsz, seqlen, k_heads, d)
 
 
 @torch.compiler.assume_constant_result

python/sglang/multimodal_gen/runtime/loader/component_loaders/vocoder_loader.py

@@ -33,10 +33,10 @@ class VocoderLoader(ComponentLoader):
         self, component_model_path: str, server_args: ServerArgs, component_name: str
     ):
         config = get_diffusers_component_config(component_path=component_model_path)
-        class_name = config.pop("_class_name", None)
+        class_name = config.pop("_class_name", None) or self.component_architecture
         assert (
             class_name is not None
-        ), "Model config does not contain a _class_name attribute. Only diffusers format is supported."
+        ), "Vocoder class name must be available from component config or pipeline config."
 
         server_args.model_paths[component_name] = component_model_path
 

python/sglang/multimodal_gen/runtime/models/dits/cosmos3video.py

@@ -22,7 +22,11 @@ from sglang.multimodal_gen.runtime.distributed import (
 )
 from sglang.multimodal_gen.runtime.layers.activation import SiluAndMul
 from sglang.multimodal_gen.runtime.layers.attention import USPAttention
-from sglang.multimodal_gen.runtime.layers.layernorm import RMSNorm, apply_qk_norm
+from sglang.multimodal_gen.runtime.layers.layernorm import (
+    RMSNorm,
+    apply_qk_norm,
+    apply_qk_norm_rope,
+)
 from sglang.multimodal_gen.runtime.layers.linear import (
     MergedColumnParallelLinear,
     ReplicatedLinear,
@@ -34,7 +38,6 @@ from sglang.multimodal_gen.runtime.layers.quantization.configs.base_config impor
 )
 from sglang.multimodal_gen.runtime.layers.rotary_embedding import (
     Qwen3VLTextRotaryEmbedding,
-    qwen3_apply_rotary_pos_emb,
 )
 from sglang.multimodal_gen.runtime.layers.visual_embedding import timestep_embedding
 from sglang.multimodal_gen.runtime.layers.vocab_parallel_embedding import (
@@ -131,6 +134,66 @@ def compute_mrope_position_ids_vision(
     return mrope_ids, next_offset
 
 
+# -----------------------------------------------------------------------------
+# Qwen3-style RoPE functions
+# -----------------------------------------------------------------------------
+
+
+def _apply_qwen3_qk_norm_rope(
+    q: torch.Tensor,
+    k: torch.Tensor,
+    q_norm: RMSNorm,
+    k_norm: RMSNorm,
+    head_dim: int,
+    cos_sin_cache: torch.Tensor,
+    rope_cache_positions: torch.Tensor,
+) -> tuple[torch.Tensor, torch.Tensor]:
+    return apply_qk_norm_rope(
+        q=q.contiguous(),
+        k=k.contiguous(),
+        q_norm=q_norm,
+        k_norm=k_norm,
+        head_dim=head_dim,
+        cos_sin_cache=cos_sin_cache,
+        is_neox=True,
+        positions=rope_cache_positions,
+    )
+
+
+def _apply_qwen3_rope_from_cache(
+    q: torch.Tensor, k: torch.Tensor, cos_sin_cache: torch.Tensor
+) -> tuple[torch.Tensor, torch.Tensor]:
+    batch_size, seq_len = q.shape[:2]
+    half = q.shape[-1] // 2
+    cos = cos_sin_cache[:, :half].view(batch_size, seq_len, 1, half).to(q.dtype)
+    sin = cos_sin_cache[:, half:].view(batch_size, seq_len, 1, half).to(q.dtype)
+
+    q1 = q[..., :half]
+    q2 = q[..., half:]
+    q_out = torch.empty_like(q)
+    q_out[..., :half] = q1 * cos - q2 * sin
+    q_out[..., half:] = q2 * cos + q1 * sin
+
+    k1 = k[..., :half]
+    k2 = k[..., half:]
+    k_out = torch.empty_like(k)
+    k_out[..., :half] = k1 * cos - k2 * sin
+    k_out[..., half:] = k2 * cos + k1 * sin
+    return q_out, k_out
+
+
+def _apply_qwen3_qk_norm_rope_split(
+    q: torch.Tensor,
+    k: torch.Tensor,
+    q_norm: RMSNorm,
+    k_norm: RMSNorm,
+    head_dim: int,
+    cos_sin_cache: torch.Tensor,
+) -> tuple[torch.Tensor, torch.Tensor]:
+    q, k = apply_qk_norm(q.contiguous(), k.contiguous(), q_norm, k_norm, head_dim)
+    return _apply_qwen3_rope_from_cache(q, k, cos_sin_cache)
+
+
 # -----------------------------------------------------------------------------
 # Cosmos3 Timestep Embedder
 # -----------------------------------------------------------------------------
@@ -295,15 +358,15 @@ class Cosmos3CausalAttention(nn.Module):
             prefix=add_prefix("to_out", prefix),
         )
 
-        # Per-head QK norm. Modules hold the weights; F.rms_norm in forward.
+        # Per-head QK norm.
         self.norm_q = RMSNorm(head_dim, eps=1e-6)
         self.norm_k = RMSNorm(head_dim, eps=1e-6)
 
     def forward(
         self,
         hidden_states: torch.Tensor,
-        freqs_cos: torch.Tensor,
-        freqs_sin: torch.Tensor,
+        cos_sin_cache: torch.Tensor,
+        rope_cache_positions: torch.Tensor,
     ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
         """Forward with KV cache return.
 
@@ -335,7 +398,7 @@ class Cosmos3CausalAttention(nn.Module):
         k = F.rms_norm(
             k, (self.head_dim,), self.norm_k.weight, self.norm_k.variance_epsilon
         )
-        q, k = qwen3_apply_rotary_pos_emb(q, k, freqs_cos, freqs_sin)
+        q, k = _apply_qwen3_rope_from_cache(q, k, cos_sin_cache)
 
         out = F.scaled_dot_product_attention(
             q.transpose(1, 2),
@@ -409,8 +472,9 @@ class Cosmos3CrossAttention(nn.Module):
         hidden_states: torch.Tensor,
         k_und: torch.Tensor,
         v_und: torch.Tensor,
-        freqs_cos: torch.Tensor,
-        freqs_sin: torch.Tensor,
+        cos_sin_cache: torch.Tensor,
+        rope_cache_positions: torch.Tensor,
+        use_fused_qk_norm_rope: bool,
     ) -> torch.Tensor:
         """Cross-attention from GEN to cached UND K/V.
 
@@ -418,8 +482,8 @@ class Cosmos3CrossAttention(nn.Module):
             hidden_states: [B, S_gen_local, hidden_size] visual tokens (may be sharded)
             k_und: [B, S_und, H_kv, D] pre-computed UND keys (always full/replicated)
             v_und: [B, S_und, H_kv, D] pre-computed UND values (always full/replicated)
-            freqs_cos: [B, S_gen_local, 1, D] cosine part of RoPE (for local shard)
-            freqs_sin: [B, S_gen_local, 1, D] sine part of RoPE (for local shard)
+            cos_sin_cache: [B*S_gen_local, D] local rows of [cos, sin]
+            rope_cache_positions: identity row positions into cos_sin_cache
         """
         batch_size, seq_len_gen = hidden_states.shape[:2]
 
@@ -440,10 +504,20 @@ class Cosmos3CrossAttention(nn.Module):
         ]
         v = qkv[:, :, self.num_attention_heads + self.num_key_value_heads :, :]
 
-        q, k = apply_qk_norm(
-            q.contiguous(), k.contiguous(), self.norm_q, self.norm_k, self.head_dim
-        )
-        q, k = qwen3_apply_rotary_pos_emb(q, k, freqs_cos, freqs_sin)
+        if use_fused_qk_norm_rope:
+            q, k = _apply_qwen3_qk_norm_rope(
+                q,
+                k,
+                self.norm_q,
+                self.norm_k,
+                self.head_dim,
+                cos_sin_cache,
+                rope_cache_positions,
+            )
+        else:
+            q, k = _apply_qwen3_qk_norm_rope_split(
+                q, k, self.norm_q, self.norm_k, self.head_dim, cos_sin_cache
+            )
 
         # K/V = [text (replicated full on every SP rank) | image (sharded same as Q)].
         # USPAttention routes through the registered attention backend (FA, sage,
@@ -497,8 +571,8 @@ class Cosmos3UndDecoderLayer(nn.Module):
     def forward(
         self,
         hidden_states: torch.Tensor,
-        freqs_cos: torch.Tensor,
-        freqs_sin: torch.Tensor,
+        cos_sin_cache: torch.Tensor,
+        rope_cache_positions: torch.Tensor,
     ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
         """Forward pass.
 
@@ -508,7 +582,9 @@ class Cosmos3UndDecoderLayer(nn.Module):
         residual = hidden_states
         hidden_states = self.input_layernorm(hidden_states)
 
-        attn_out, k, v = self.self_attn(hidden_states, freqs_cos, freqs_sin)
+        attn_out, k, v = self.self_attn(
+            hidden_states, cos_sin_cache, rope_cache_positions
+        )
         hidden_states = residual + attn_out
 
         residual = hidden_states
@@ -565,8 +641,9 @@ class Cosmos3GenDecoderLayer(nn.Module):
         hidden_states: torch.Tensor,
         k_und: torch.Tensor,
         v_und: torch.Tensor,
-        freqs_cos: torch.Tensor,
-        freqs_sin: torch.Tensor,
+        cos_sin_cache: torch.Tensor,
+        rope_cache_positions: torch.Tensor,
+        use_fused_qk_norm_rope: bool,
         residual: torch.Tensor | None = None,
     ) -> tuple[torch.Tensor, torch.Tensor]:
         # Fused add+rmsnorm: each `(hidden_states, residual) = norm(...)`
@@ -580,7 +657,12 @@ class Cosmos3GenDecoderLayer(nn.Module):
             hidden_states, residual = self.input_layernorm(hidden_states, residual)
 
         hidden_states = self.cross_attention(
-            hidden_states, k_und, v_und, freqs_cos, freqs_sin
+            hidden_states,
+            k_und,
+            v_und,
+            cos_sin_cache,
+            rope_cache_positions,
+            use_fused_qk_norm_rope,
         )
 
         hidden_states, residual = self.post_attention_layernorm(hidden_states, residual)
@@ -646,31 +728,28 @@ class Cosmos3LanguageModel(nn.Module):
         self,
         text_ids: torch.Tensor,
         text_mask: torch.Tensor,
-        freqs_cos: torch.Tensor,
-        freqs_sin: torch.Tensor,
+        position_ids: torch.Tensor,
     ) -> list[tuple[torch.Tensor, torch.Tensor]]:
         """Process text tokens and return per-layer K/V cache.
 
         Args:
             text_ids: [B, S] token IDs
             text_mask: [B, S] float mask (1=real, 0=pad)
-            freqs_cos: [B, S, D] RoPE cosines
-            freqs_sin: [B, S, D] RoPE sines
+            position_ids: [3, B, S] mRoPE position IDs
 
         Returns:
             List of (K, V) per layer for GEN cross-attention
         """
         hidden = self.embed_tokens(text_ids)
         mask_3d = text_mask.unsqueeze(-1)
-
-        # Add dimension for per-head broadcast
-        freqs_cos = freqs_cos.unsqueeze(2)  # [B, S, 1, D]
-        freqs_sin = freqs_sin.unsqueeze(2)
+        cos_sin_cache, rope_cache_positions = self.rotary_emb.build_rope_cache_inputs(
+            position_ids, cache_dtype=hidden.dtype
+        )
 
         cached_kv: list[tuple[torch.Tensor, torch.Tensor]] = []
         for layer in self.layers:
             hidden = hidden * mask_3d
-            hidden, k, v = layer(hidden, freqs_cos, freqs_sin)
+            hidden, k, v = layer(hidden, cos_sin_cache, rope_cache_positions)
             cached_kv.append((k, v))
 
         return cached_kv
@@ -800,7 +879,7 @@ class Cosmos3OmniTransformer(CachableDiT):
         # This allows maintaining separate caches for conditional and unconditional
         # prompts, avoiding recomputation on every denoising step
         self.cached_kv: dict[str, list[tuple[torch.Tensor, torch.Tensor]]] = {}
-        self.cached_freqs_gen: dict[str, tuple[torch.Tensor, torch.Tensor]] = {}
+        self.cached_gen_rope_inputs: dict[str, tuple[torch.Tensor, torch.Tensor]] = {}
 
         self.__post_init__()
 
@@ -840,7 +919,7 @@ class Cosmos3OmniTransformer(CachableDiT):
             x = x[:, :, :, :H, :W]
         return x
 
-    def _compute_rope_freqs(
+    def _compute_rope_position_ids(
         self,
         text_mask: torch.Tensor,
         T: int,
@@ -848,12 +927,8 @@ class Cosmos3OmniTransformer(CachableDiT):
         Wp: int,
         fps: float | None,
         device: torch.device,
-        dtype: torch.dtype,
-    ) -> tuple[
-        tuple[torch.Tensor, torch.Tensor],
-        tuple[torch.Tensor, torch.Tensor],
-    ]:
-        """Compute mRoPE cos/sin for UND (text) and GEN (visual) pathways."""
+    ) -> tuple[torch.Tensor, torch.Tensor]:
+        """Compute mRoPE position IDs for UND text and GEN visual tokens."""
         B = text_mask.shape[0]
         S_text = text_mask.shape[1]
         text_lengths = text_mask.sum(dim=1).long()
@@ -892,14 +967,7 @@ class Cosmos3OmniTransformer(CachableDiT):
         text_pos_ids = torch.stack(text_pos_list, dim=1).to(device)  # [3, B, S_text]
         vis_pos_ids = torch.stack(vis_pos_list, dim=1).to(device)  # [3, B, S_vis]
 
-        rotary_emb = self.language_model.rotary_emb
-        _dummy = torch.tensor([], dtype=dtype, device=device)
-        cos_und, sin_und = rotary_emb(_dummy, position_ids=text_pos_ids)
-        cos_gen, sin_gen = rotary_emb(_dummy, position_ids=vis_pos_ids)
-
-        freqs_und = (cos_und, sin_und)
-        freqs_gen = (cos_gen, sin_gen)
-        return freqs_und, freqs_gen
+        return text_pos_ids, vis_pos_ids
 
     def reset_cache(self, cache_key: str | None = None):
         """Reset cached K/V from UND pathway.
@@ -911,20 +979,20 @@ class Cosmos3OmniTransformer(CachableDiT):
         if cache_key is None:
             # Reset all caches
             self.cached_kv = {}
-            self.cached_freqs_gen = {}
+            self.cached_gen_rope_inputs = {}
         else:
             # Reset specific cache
             if cache_key in self.cached_kv:
                 del self.cached_kv[cache_key]
-            if cache_key in self.cached_freqs_gen:
-                del self.cached_freqs_gen[cache_key]
+            if cache_key in self.cached_gen_rope_inputs:
+                del self.cached_gen_rope_inputs[cache_key]
 
     def _ensure_cache_dicts(self):
         """Ensure cache dictionaries exist (for backwards compatibility)."""
         if not isinstance(self.cached_kv, dict):
             self.cached_kv = {}
-        if not isinstance(self.cached_freqs_gen, dict):
-            self.cached_freqs_gen = {}
+        if not isinstance(self.cached_gen_rope_inputs, dict):
+            self.cached_gen_rope_inputs = {}
 
     def forward(
         self,
@@ -1035,47 +1103,49 @@ class Cosmos3OmniTransformer(CachableDiT):
 
         # Compute UND K/V cache for this cache_key if not already cached
         # This allows reusing the cache across denoising steps for the same text
-        if cache_key not in self.cached_kv:
-            freqs_und, freqs_gen = self._compute_rope_freqs(
-                text_mask, T, Hp, Wp, fps, hidden_states.device, hidden_states.dtype
+        if (
+            cache_key not in self.cached_kv
+            or cache_key not in self.cached_gen_rope_inputs
+        ):
+            text_pos_ids, vis_pos_ids = self._compute_rope_position_ids(
+                text_mask, T, Hp, Wp, fps, hidden_states.device
             )
             # UND K/V cache is kept FULL on all ranks (not sharded). Text
             # sequence is short, so memory impact is minimal, and the GEN
             # cross-attention needs the full K/V on every SP rank.
             self.cached_kv[cache_key] = self.language_model(
-                text_ids, text_mask, freqs_und[0], freqs_und[1]
+                text_ids, text_mask, text_pos_ids
             )
-            cos_gen, sin_gen = freqs_gen
             if sequence_shard_enabled:
                 if seq_shard_pad > 0:
-                    pad_cos = cos_gen[:, -1:].expand(-1, seq_shard_pad, -1)
-                    pad_sin = sin_gen[:, -1:].expand(-1, seq_shard_pad, -1)
-                    cos_gen = torch.cat([cos_gen, pad_cos], dim=1)
-                    sin_gen = torch.cat([sin_gen, pad_sin], dim=1)
-                cos_gen = cos_gen.view(batch_size, self.sp_size, local_seq_len, -1)
-                sin_gen = sin_gen.view(batch_size, self.sp_size, local_seq_len, -1)
-                cos_gen = cos_gen[:, self.sp_rank, :, :]
-                sin_gen = sin_gen[:, self.sp_rank, :, :]
-            cos_gen = cos_gen.unsqueeze(2)  # [B, S, 1, D]
-            sin_gen = sin_gen.unsqueeze(2)
-            self.cached_freqs_gen[cache_key] = (cos_gen, sin_gen)
+                    pad_pos = vis_pos_ids[:, :, -1:].expand(-1, -1, seq_shard_pad)
+                    vis_pos_ids = torch.cat([vis_pos_ids, pad_pos], dim=2)
+                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] = (
+                self.language_model.rotary_emb.build_rope_cache_inputs(
+                    vis_pos_ids, cache_dtype=hidden_gen.dtype
+                )
+            )
 
-        freqs_gen = self.cached_freqs_gen[cache_key]
-        cos_gen, sin_gen = freqs_gen
+        cos_sin_gen, gen_rope_cache_positions = self.cached_gen_rope_inputs[cache_key]
 
         # Run GEN layers. `residual` is threaded so each layer's
         # input_layernorm and post_attention_layernorm can use the
         # 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
         for i, layer in enumerate(self.gen_layers):
             k_und, v_und = cached_kv_for_key[i]
             hidden_gen, residual = layer(
                 hidden_gen,
                 k_und,
                 v_und,
-                cos_gen,
-                sin_gen,
+                cos_sin_gen,
+                gen_rope_cache_positions,
+                use_fused_qk_norm_rope,
                 residual=residual,
             )