diff --git a/python/sglang/multimodal_gen/runtime/models/dits/cosmos3video.py b/python/sglang/multimodal_gen/runtime/models/dits/cosmos3video.py index 6ff98af59..9f76363ea 100644 --- a/python/sglang/multimodal_gen/runtime/models/dits/cosmos3video.py +++ b/python/sglang/multimodal_gen/runtime/models/dits/cosmos3video.py @@ -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, ) diff --git a/python/sglang/multimodal_gen/runtime/pipelines_core/stages/model_specific_stages/cosmos3.py b/python/sglang/multimodal_gen/runtime/pipelines_core/stages/model_specific_stages/cosmos3.py index 6d8e0cb59..697ce3bbf 100644 --- a/python/sglang/multimodal_gen/runtime/pipelines_core/stages/model_specific_stages/cosmos3.py +++ b/python/sglang/multimodal_gen/runtime/pipelines_core/stages/model_specific_stages/cosmos3.py @@ -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: diff --git a/test/registered/kernels/ops/diffusion/test_qknorm_rope.py b/test/registered/kernels/ops/diffusion/test_qknorm_rope.py index be3b67d6e..867d23c78 100644 --- a/test/registered/kernels/ops/diffusion/test_qknorm_rope.py +++ b/test/registered/kernels/ops/diffusion/test_qknorm_rope.py @@ -413,6 +413,99 @@ def test_qknorm_rope_pack_kv_matches_separate_ops() -> None: assert torch.equal(packed_v_ref, packed_kv[1]) +def test_qknorm_rope_pack_kv_preserves_split_bf16_rounding() -> None: + from sgl_kernel import rotary_embedding + + from sglang.kernels.ops.diffusion.qknorm_rope import ( + fused_qknorm_rope_pack_kv, + ) + from sglang.kernels.ops.layernorm.norm import fused_inplace_qknorm + + batch_size = 1 + prefix_tokens, suffix_tokens = 17, 1024 + num_q_heads, num_kv_heads, head_dim = 32, 8, 64 + num_heads = num_q_heads + 2 * num_kv_heads + qkv = torch.randn( + batch_size, + suffix_tokens, + num_heads, + head_dim, + device=DEVICE, + dtype=DTYPE, + ) + prefix_qkv = torch.randn( + batch_size, + prefix_tokens, + num_heads, + head_dim, + device=DEVICE, + dtype=DTYPE, + ) + k_prefix = prefix_qkv[:, :, num_q_heads : num_q_heads + num_kv_heads] + v_prefix = prefix_qkv[:, :, num_q_heads + num_kv_heads :] + q_weight = torch.randn(head_dim, device=DEVICE, dtype=DTYPE) + k_weight = torch.randn(head_dim, device=DEVICE, dtype=DTYPE) + positions = torch.arange( + batch_size * suffix_tokens, device=DEVICE, dtype=torch.int64 + ) + cos_sin_cache = create_cos_sin_cache(head_dim, batch_size * suffix_tokens).to(DTYPE) + + qkv_ref = qkv.clone() + q_ref = qkv_ref[:, :, :num_q_heads] + k_ref = qkv_ref[:, :, num_q_heads : num_q_heads + num_kv_heads] + v_ref = qkv_ref[:, :, num_q_heads + num_kv_heads :] + fused_inplace_qknorm( + q_ref.view(-1, num_q_heads, head_dim), + k_ref.view(-1, num_kv_heads, head_dim), + q_weight, + k_weight, + eps=1e-6, + ) + rotary_embedding( + positions, + q_ref.view(-1, num_q_heads * head_dim), + k_ref.view(-1, num_kv_heads * head_dim), + head_dim, + cos_sin_cache, + True, + ) + packed_k_ref = torch.cat([k_prefix, k_ref], dim=1) + packed_v_ref = torch.cat([v_prefix, v_ref], dim=1) + + qkv_fused = qkv.clone() + q_fused = qkv_fused[:, :, :num_q_heads] + k_fused = qkv_fused[:, :, num_q_heads : num_q_heads + num_kv_heads] + v_fused = qkv_fused[:, :, num_q_heads + num_kv_heads :] + packed_kv = torch.empty( + 2, + batch_size, + prefix_tokens + suffix_tokens, + num_kv_heads, + head_dim, + device=DEVICE, + dtype=DTYPE, + ) + fused_qknorm_rope_pack_kv( + q_fused, + k_fused, + v_fused, + k_prefix, + v_prefix, + packed_kv, + q_weight, + k_weight, + cos_sin_cache, + positions, + is_neox=True, + rope_dim=head_dim, + round_norm_before_rope=True, + ) + + assert torch.equal(q_ref, q_fused) + assert torch.equal(packed_k_ref, packed_kv[0]) + assert torch.equal(packed_v_ref, packed_kv[1]) + + def test_qknorm_rope_accepts_empty_token_dimension() -> None: from sglang.kernels.ops.diffusion.qknorm_rope import fused_inplace_qknorm_rope