From 25d8f431d161204792c06a37d702d2d90dabbbd3 Mon Sep 17 00:00:00 2001 From: Chi McIsaac <153383231+qimcis@users.noreply.github.com> Date: Fri, 5 Jun 2026 18:15:42 -0700 Subject: [PATCH] [diffusion] optimize: cosmos3 fused qknorm rope (#27096) --- .../runtime/layers/layernorm.py | 17 +- .../runtime/layers/rotary_embedding/mrope.py | 36 ++- .../runtime/layers/rotary_embedding/utils.py | 77 +++++-- .../component_loaders/vocoder_loader.py | 4 +- .../runtime/models/dits/cosmos3video.py | 210 ++++++++++++------ 5 files changed, 242 insertions(+), 102 deletions(-) diff --git a/python/sglang/multimodal_gen/runtime/layers/layernorm.py b/python/sglang/multimodal_gen/runtime/layers/layernorm.py index d4ae19ee1..ebd803e59 100755 --- a/python/sglang/multimodal_gen/runtime/layers/layernorm.py +++ b/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 diff --git a/python/sglang/multimodal_gen/runtime/layers/rotary_embedding/mrope.py b/python/sglang/multimodal_gen/runtime/layers/rotary_embedding/mrope.py index e1ad8300f..bfe77777f 100644 --- a/python/sglang/multimodal_gen/runtime/layers/rotary_embedding/mrope.py +++ b/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 diff --git a/python/sglang/multimodal_gen/runtime/layers/rotary_embedding/utils.py b/python/sglang/multimodal_gen/runtime/layers/rotary_embedding/utils.py index 3647b1a7e..66789b087 100644 --- a/python/sglang/multimodal_gen/runtime/layers/rotary_embedding/utils.py +++ b/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 diff --git a/python/sglang/multimodal_gen/runtime/loader/component_loaders/vocoder_loader.py b/python/sglang/multimodal_gen/runtime/loader/component_loaders/vocoder_loader.py index 8e8d6abc9..b7675eff1 100644 --- a/python/sglang/multimodal_gen/runtime/loader/component_loaders/vocoder_loader.py +++ b/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 diff --git a/python/sglang/multimodal_gen/runtime/models/dits/cosmos3video.py b/python/sglang/multimodal_gen/runtime/models/dits/cosmos3video.py index c1ba691eb..0f3f28333 100644 --- a/python/sglang/multimodal_gen/runtime/models/dits/cosmos3video.py +++ b/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, )