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

[Diffusion] Make the apply_qknorm function easier to use (#17537)

Browse Source
This commit is contained in:
Xiaoyu Zhang
2026-01-22 22:32:15 +08:00
committed by GitHub
parent 3705f90629
commit 5324027007
4 changed files with 40 additions and 98 deletions
Download Patch File Download Diff File Expand all files Collapse all files
python/sglang/multimodal_gen/runtime/layers/layernorm.py
-7
View File
@@ -471,13 +471,6 @@ def apply_qk_norm(
) )
return q, k return q, k
# Fallback for AMD/ROCm: apply RMSNorm separately to q and k
import warnings
warnings.warn(
"Fused QK-norm not available, using RMSNorm fallback",
stacklevel=2,
)
q_shape = q.shape q_shape = q.shape
k_shape = k.shape k_shape = k.shape
q_out = q_norm(q.view(-1, head_dim)).view(q_shape) q_out = q_norm(q.view(-1, head_dim)).view(q_shape)
python/sglang/multimodal_gen/runtime/models/dits/flux.py
+16 -37
View File
@@ -27,7 +27,6 @@ from diffusers.models.normalization import (
) )
from torch.nn import LayerNorm as LayerNorm from torch.nn import LayerNorm as LayerNorm
from sglang.jit_kernel.norm import can_use_fused_inplace_qknorm
from sglang.multimodal_gen.configs.models.dits.flux import FluxConfig from sglang.multimodal_gen.configs.models.dits.flux import FluxConfig
from sglang.multimodal_gen.runtime.layers.attention import USPAttention from sglang.multimodal_gen.runtime.layers.attention import USPAttention
@@ -49,7 +48,6 @@ from sglang.multimodal_gen.runtime.utils.layerwise_offload import OffloadableDiT
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__) # pylint: disable=invalid-name logger = init_logger(__name__) # pylint: disable=invalid-name
_is_cuda = current_platform.is_cuda()
def _get_qkv_projections( def _get_qkv_projections(
@@ -167,47 +165,28 @@ class FluxAttention(torch.nn.Module, AttentionModuleMixin):
query = query.unflatten(-1, (self.heads, -1)) query = query.unflatten(-1, (self.heads, -1))
key = key.unflatten(-1, (self.heads, -1)) key = key.unflatten(-1, (self.heads, -1))
value = value.unflatten(-1, (self.heads, -1)) value = value.unflatten(-1, (self.heads, -1))
if ( query, key = apply_qk_norm(
_is_cuda q=query,
and (self.norm_q.variance_epsilon == self.norm_k.variance_epsilon) k=key,
and can_use_fused_inplace_qknorm(self.head_dim, query.dtype) q_norm=self.norm_q,
): k_norm=self.norm_k,
query, key = apply_qk_norm( head_dim=self.head_dim,
q=query, allow_inplace=True,
k=key, )
q_norm=self.norm_q,
k_norm=self.norm_k,
head_dim=self.head_dim,
allow_inplace=True,
)
else:
query = self.norm_q(query)
key = self.norm_k(key)
if self.added_kv_proj_dim is not None: if self.added_kv_proj_dim is not None:
encoder_query = encoder_query.unflatten(-1, (self.heads, -1)) encoder_query = encoder_query.unflatten(-1, (self.heads, -1))
encoder_key = encoder_key.unflatten(-1, (self.heads, -1)) encoder_key = encoder_key.unflatten(-1, (self.heads, -1))
encoder_value = encoder_value.unflatten(-1, (self.heads, -1)) encoder_value = encoder_value.unflatten(-1, (self.heads, -1))
if ( encoder_query, encoder_key = apply_qk_norm(
_is_cuda q=encoder_query,
and ( k=encoder_key,
self.norm_added_q.variance_epsilon q_norm=self.norm_added_q,
== self.norm_added_k.variance_epsilon k_norm=self.norm_added_k,
) head_dim=self.head_dim,
and can_use_fused_inplace_qknorm(self.head_dim, encoder_query.dtype) allow_inplace=True,
): )
encoder_query, encoder_key = apply_qk_norm(
q=encoder_query,
k=encoder_key,
q_norm=self.norm_added_q,
k_norm=self.norm_added_k,
head_dim=self.head_dim,
allow_inplace=True,
)
else:
encoder_query = self.norm_added_q(encoder_query)
encoder_key = self.norm_added_k(encoder_key)
bsz, seq_len, _, _ = query.shape bsz, seq_len, _, _ = query.shape
query = torch.cat([encoder_query, query], dim=1) query = torch.cat([encoder_query, query], dim=1)
python/sglang/multimodal_gen/runtime/models/dits/flux_2.py
+16 -37
View File
@@ -20,7 +20,6 @@ from diffusers.models.attention import AttentionModuleMixin
from diffusers.models.embeddings import TimestepEmbedding, Timesteps from diffusers.models.embeddings import TimestepEmbedding, Timesteps
from diffusers.models.normalization import AdaLayerNormContinuous from diffusers.models.normalization import AdaLayerNormContinuous
from sglang.jit_kernel.norm import can_use_fused_inplace_qknorm
from sglang.multimodal_gen.configs.models.dits.flux import FluxConfig from sglang.multimodal_gen.configs.models.dits.flux import FluxConfig
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, apply_qk_norm from sglang.multimodal_gen.runtime.layers.layernorm import RMSNorm, apply_qk_norm
@@ -35,7 +34,6 @@ from sglang.multimodal_gen.runtime.utils.layerwise_offload import OffloadableDiT
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__) # pylint: disable=invalid-name logger = init_logger(__name__) # pylint: disable=invalid-name
_is_cuda = current_platform.is_cuda()
def _get_qkv_projections( def _get_qkv_projections(
@@ -198,47 +196,28 @@ class Flux2Attention(torch.nn.Module, AttentionModuleMixin):
key = key.unflatten(-1, (self.heads, -1)) key = key.unflatten(-1, (self.heads, -1))
value = value.unflatten(-1, (self.heads, -1)) value = value.unflatten(-1, (self.heads, -1))
if ( query, key = apply_qk_norm(
_is_cuda q=query,
and (self.norm_q.variance_epsilon == self.norm_k.variance_epsilon) k=key,
and can_use_fused_inplace_qknorm(self.head_dim, query.dtype) q_norm=self.norm_q,
): k_norm=self.norm_k,
query, key = apply_qk_norm( head_dim=self.head_dim,
q=query, allow_inplace=True,
k=key, )
q_norm=self.norm_q,
k_norm=self.norm_k,
head_dim=self.head_dim,
allow_inplace=True,
)
else:
query = self.norm_q(query)
key = self.norm_k(key)
if self.added_kv_proj_dim is not None: if self.added_kv_proj_dim is not None:
encoder_query = encoder_query.unflatten(-1, (self.heads, -1)) encoder_query = encoder_query.unflatten(-1, (self.heads, -1))
encoder_key = encoder_key.unflatten(-1, (self.heads, -1)) encoder_key = encoder_key.unflatten(-1, (self.heads, -1))
encoder_value = encoder_value.unflatten(-1, (self.heads, -1)) encoder_value = encoder_value.unflatten(-1, (self.heads, -1))
if ( encoder_query, encoder_key = apply_qk_norm(
_is_cuda q=encoder_query,
and ( k=encoder_key,
self.norm_added_q.variance_epsilon q_norm=self.norm_added_q,
== self.norm_added_k.variance_epsilon k_norm=self.norm_added_k,
) head_dim=self.head_dim,
and can_use_fused_inplace_qknorm(self.head_dim, encoder_query.dtype) allow_inplace=True,
): )
encoder_query, encoder_key = apply_qk_norm(
q=encoder_query,
k=encoder_key,
q_norm=self.norm_added_q,
k_norm=self.norm_added_k,
head_dim=self.head_dim,
allow_inplace=True,
)
else:
encoder_query = self.norm_added_q(encoder_query)
encoder_key = self.norm_added_k(encoder_key)
query = torch.cat([encoder_query, query], dim=1) query = torch.cat([encoder_query, query], dim=1)
key = torch.cat([encoder_key, key], dim=1) key = torch.cat([encoder_key, key], dim=1)
python/sglang/multimodal_gen/runtime/models/dits/zimage.py
+8 -17
View File
@@ -4,7 +4,6 @@ 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.distributed import get_tp_world_size from sglang.multimodal_gen.runtime.distributed import get_tp_world_size
from sglang.multimodal_gen.runtime.layers.activation import SiluAndMul from sglang.multimodal_gen.runtime.layers.activation import SiluAndMul
@@ -171,22 +170,14 @@ class ZImageAttention(nn.Module):
v = v.view(*v.shape[:-1], self.local_num_kv_heads, self.head_dim) v = v.view(*v.shape[:-1], self.local_num_kv_heads, self.head_dim)
if self.qk_norm: if self.qk_norm:
if ( q, k = apply_qk_norm(
_is_cuda q=q,
and (self.norm_q.variance_epsilon == self.norm_k.variance_epsilon) k=k,
and can_use_fused_inplace_qknorm(self.head_dim, q.dtype) q_norm=self.norm_q,
): k_norm=self.norm_k,
q, k = apply_qk_norm( head_dim=self.head_dim,
q=q, allow_inplace=True,
k=k, )
q_norm=self.norm_q,
k_norm=self.norm_k,
head_dim=self.head_dim,
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:
cos, sin = freqs_cis cos, sin = freqs_cis