chenchenghao/sglang
1
0
Fork 0
Code Issues Pull Requests Actions 11 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

39
python/sglang/multimodal_gen/runtime/models/dits/glm_image.py
View File

@@ -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":