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

[Diffusion] Zimage opt with qknorm and flashinfer rope (#16161)

Browse Source 
Co-authored-by: gemini-code-assist[bot] <176961590+gemini-code-assist[bot]@users.noreply.github.com>
This commit is contained in:
Xiaoyu Zhang
2025-12-30 23:39:32 +08:00
committed by GitHub
co-authored by gemini-code-assist[bot]
parent b369aaa23f
commit 733a0c1a37
1 changed files with 38 additions and 17 deletions
Download Patch File Download Diff File Expand all files Collapse all files
python/sglang/multimodal_gen/runtime/models/dits/zimage.py
+38 -17
View File
@@ -4,17 +4,21 @@ from typing import Any, List, Optional, Tuple
import torch import torch
import torch.nn as nn import torch.nn as nn
from sglang.jit_kernel.norm import can_use_fused_inplace_qknorm
from sglang.multimodal_gen.configs.models.dits.zimage import ZImageDitConfig from sglang.multimodal_gen.configs.models.dits.zimage import ZImageDitConfig
from sglang.multimodal_gen.runtime.layers.activation import SiluAndMul from sglang.multimodal_gen.runtime.layers.activation import SiluAndMul
from sglang.multimodal_gen.runtime.layers.attention import USPAttention from sglang.multimodal_gen.runtime.layers.attention import USPAttention
from sglang.multimodal_gen.runtime.layers.layernorm import RMSNorm from sglang.multimodal_gen.runtime.layers.layernorm import RMSNorm, apply_qk_norm
from sglang.multimodal_gen.runtime.layers.linear import ( from sglang.multimodal_gen.runtime.layers.linear import (
ColumnParallelLinear, ColumnParallelLinear,
MergedColumnParallelLinear, MergedColumnParallelLinear,
ReplicatedLinear, ReplicatedLinear,
RowParallelLinear, RowParallelLinear,
) )
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.models.dits.base import CachableDiT from sglang.multimodal_gen.runtime.models.dits.base import CachableDiT
from sglang.multimodal_gen.runtime.platforms import current_platform from sglang.multimodal_gen.runtime.platforms import current_platform
from sglang.multimodal_gen.runtime.utils.layerwise_offload import OffloadableDiTMixin from sglang.multimodal_gen.runtime.utils.layerwise_offload import OffloadableDiTMixin
@@ -151,23 +155,40 @@ class ZImageAttention(nn.Module):
k = k.view(*k.shape[:-1], self.num_kv_heads, self.head_dim) k = k.view(*k.shape[:-1], self.num_kv_heads, self.head_dim)
v = v.view(*v.shape[:-1], self.num_kv_heads, self.head_dim) v = v.view(*v.shape[:-1], self.num_kv_heads, self.head_dim)
if self.norm_q is not None: if self.qk_norm:
q = self.norm_q(q) if (
if self.norm_k is not None: q.is_cuda
k = self.norm_k(k) and (self.norm_q.variance_epsilon == self.norm_k.variance_epsilon)
and can_use_fused_inplace_qknorm(self.head_dim)
# Apply RoPE ):
def apply_rotary_emb( q, k = apply_qk_norm(
x_in: torch.Tensor, q=q,
freqs_cis: torch.Tensor, k=k,
) -> Tuple[torch.Tensor, torch.Tensor]: q_norm=self.norm_q,
cos, sin = freqs_cis k_norm=self.norm_k,
x_out = _apply_rotary_emb(x_in, cos, sin, is_neox_style=False) head_dim=self.head_dim,
return x_out allow_inplace=True,
)
else:
q = self.norm_q(q)
k = self.norm_k(k)
if freqs_cis is not None: if freqs_cis is not None:
q = apply_rotary_emb(q, freqs_cis) cos, sin = freqs_cis
k = apply_rotary_emb(k, freqs_cis) if q.is_cuda and q.shape == k.shape:
cos_sin_cache = torch.cat(
[
cos.to(dtype=torch.float32).contiguous(),
sin.to(dtype=torch.float32).contiguous(),
],
dim=-1,
)
q, k = apply_flashinfer_rope_qk_inplace(
q, k, cos_sin_cache, is_neox=False
)
else:
q = _apply_rotary_emb(q, cos, sin, is_neox_style=False)
k = _apply_rotary_emb(k, cos, sin, is_neox_style=False)
hidden_states = self.attn(q, k, v) hidden_states = self.attn(q, k, v)
hidden_states = hidden_states.flatten(2) hidden_states = hidden_states.flatten(2)