minke.yu/sglang
1
0
Fork 0
Code Issues Pull Requests Actions 12 Packages Projects Releases Wiki Activity

[Diffusion] glm-image apply flashinfer rope (#17689)

Browse Source
This commit is contained in:
Xiaoyu Zhang
2026-01-28 08:51:37 +08:00
committed by GitHub
parent 93423ff780
commit 331a22427c
1 changed files with 31 additions and 8 deletions
Download Patch File Download Diff File Expand all files Collapse all files
python/sglang/multimodal_gen/runtime/models/dits/glm_image.py
+31 -8
View File
@@ -25,15 +25,23 @@ from sglang.multimodal_gen.runtime.layers.layernorm import (
ScaleResidualLayerNormScaleShift, ScaleResidualLayerNormScaleShift,
) )
from sglang.multimodal_gen.runtime.layers.linear import ReplicatedLinear 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.layers.visual_embedding import Timesteps
from sglang.multimodal_gen.runtime.models.dits.base import CachableDiT 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.layerwise_offload import OffloadableDiTMixin
from sglang.multimodal_gen.runtime.utils.logging_utils import init_logger from sglang.multimodal_gen.runtime.utils.logging_utils import init_logger
logger = init_logger(__name__) logger = init_logger(__name__)
_is_cuda = current_platform.is_cuda()
class GlmImageLayerKVCache: class GlmImageLayerKVCache:
"""KV cache for GlmImage model.""" """KV cache for GlmImage model."""
@@ -383,12 +391,27 @@ class GlmImageAttention(torch.nn.Module):
if image_rotary_emb is not None: if image_rotary_emb is not None:
cos, sin = image_rotary_emb cos, sin = image_rotary_emb
query[:, text_seq_length:, :, :] = _apply_rotary_emb( if _is_cuda and cos.dim() == 2:
query[:, text_seq_length:, :, :], cos, sin, is_neox_style=True q_img = query[:, text_seq_length:, :, :]
) k_img = key[:, text_seq_length:, :, :]
key[:, text_seq_length:, :, :] = _apply_rotary_emb( cos_sin_cache = torch.cat(
key[:, text_seq_length:, :, :], cos, sin, is_neox_style=True [
) 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 is not None:
if kv_cache.mode == "write": if kv_cache.mode == "write":