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

[diffusion] fix: keep fused qk-norm-rope out of dynamo tracing (#31849)

Browse Source
This commit is contained in:
Mick
2026-07-28 19:51:00 +08:00
committed by GitHub
parent a24906a091
commit 75017c3fa0
2 changed files with 3 additions and 2 deletions
Download Patch File Download Diff File Expand all files Collapse all files
python/sglang/multimodal_gen/runtime/layers/layernorm.py
+1
View File
@@ -1021,6 +1021,7 @@ def apply_qk_norm_rope(
if ( if (
fused_enabled fused_enabled
and _is_cuda and _is_cuda
and not torch.compiler.is_compiling()
and allow_inplace and allow_inplace
and (q_eps == k_eps) and (q_eps == k_eps)
and q.dtype in (torch.float16, torch.bfloat16) and q.dtype in (torch.float16, torch.bfloat16)
python/sglang/multimodal_gen/runtime/models/dits/qwen_image.py
+2 -2
View File
@@ -737,7 +737,7 @@ class QwenImageCrossAttention(nn.Module):
k=img_key, k=img_key,
q_norm=self.norm_q, q_norm=self.norm_q,
k_norm=self.norm_k, k_norm=self.norm_k,
head_dim=img_query.shape[-1], head_dim=self.head_dim,
cos_sin_cache=img_cache, cos_sin_cache=img_cache,
is_neox=False, is_neox=False,
allow_inplace=True, allow_inplace=True,
@@ -747,7 +747,7 @@ class QwenImageCrossAttention(nn.Module):
k=txt_key, k=txt_key,
q_norm=self.norm_added_q, q_norm=self.norm_added_q,
k_norm=self.norm_added_k, k_norm=self.norm_added_k,
head_dim=txt_query.shape[-1], head_dim=self.head_dim,
cos_sin_cache=txt_cache, cos_sin_cache=txt_cache,
is_neox=False, is_neox=False,
allow_inplace=True, allow_inplace=True,