diff --git a/python/sglang/multimodal_gen/runtime/models/dits/glm_image.py b/python/sglang/multimodal_gen/runtime/models/dits/glm_image.py index 439cc7242..68e6ebefa 100644 --- a/python/sglang/multimodal_gen/runtime/models/dits/glm_image.py +++ b/python/sglang/multimodal_gen/runtime/models/dits/glm_image.py @@ -25,15 +25,23 @@ from sglang.multimodal_gen.runtime.layers.layernorm import ( ScaleResidualLayerNormScaleShift, ) from sglang.multimodal_gen.runtime.layers.linear import ReplicatedLinear -from sglang.multimodal_gen.runtime.layers.rotary_embedding import _apply_rotary_emb +from sglang.multimodal_gen.runtime.layers.rotary_embedding import ( + _apply_rotary_emb, + apply_flashinfer_rope_qk_inplace, +) from sglang.multimodal_gen.runtime.layers.visual_embedding import Timesteps from sglang.multimodal_gen.runtime.models.dits.base import CachableDiT -from sglang.multimodal_gen.runtime.platforms import AttentionBackendEnum +from sglang.multimodal_gen.runtime.platforms import ( + AttentionBackendEnum, + current_platform, +) from sglang.multimodal_gen.runtime.utils.layerwise_offload import OffloadableDiTMixin from sglang.multimodal_gen.runtime.utils.logging_utils import init_logger logger = init_logger(__name__) +_is_cuda = current_platform.is_cuda() + class GlmImageLayerKVCache: """KV cache for GlmImage model.""" @@ -383,12 +391,27 @@ class GlmImageAttention(torch.nn.Module): if image_rotary_emb is not None: cos, sin = image_rotary_emb - query[:, text_seq_length:, :, :] = _apply_rotary_emb( - query[:, text_seq_length:, :, :], cos, sin, is_neox_style=True - ) - key[:, text_seq_length:, :, :] = _apply_rotary_emb( - key[:, text_seq_length:, :, :], cos, sin, is_neox_style=True - ) + if _is_cuda and cos.dim() == 2: + q_img = query[:, text_seq_length:, :, :] + k_img = key[:, text_seq_length:, :, :] + cos_sin_cache = torch.cat( + [ + cos.to(dtype=torch.float32).contiguous(), + sin.to(dtype=torch.float32).contiguous(), + ], + dim=-1, + ) + # apply_flashinfer_rope_qk_inplace is inplace kernel and q_img/k_img are views of query/key, so we need not copy back + q_out, k_out = apply_flashinfer_rope_qk_inplace( + q_img, k_img, cos_sin_cache, is_neox=True + ) + else: + query[:, text_seq_length:, :, :] = _apply_rotary_emb( + query[:, text_seq_length:, :, :], cos, sin, is_neox_style=True + ) + key[:, text_seq_length:, :, :] = _apply_rotary_emb( + key[:, text_seq_length:, :, :], cos, sin, is_neox_style=True + ) if kv_cache is not None: if kv_cache.mode == "write":