chenchenghao/sglang
1
0
Fork 0
Code Issues Pull Requests Actions 14 Packages Projects Releases Wiki Activity

[Diffusion] Fix torch.compile RMSNorm fallback for Z-Image (#20962)

Browse Source 
Co-authored-by: Mick <mickjagger19@icloud.com>
This commit is contained in:
Xiaoyu Zhang
2026-03-22 15:38:22 +08:00
committed by GitHub
co-authored by Mick
parent 1e97864d75
commit 1b65c0d259
12 changed files with 170 additions and 252 deletions
Download Patch File Download Diff File Expand all files Collapse all files
python/sglang/jit_kernel/benchmark/bench_norm_impls.py
+5 -7
View File
@@ -57,11 +57,11 @@ ACTUAL_DIFFUSION_GROUPS: list[
"qwen-edit",
"1 GPU",
[
("qwen_edit_ln_189x3072", "layernorm", (1, 189, 3072), SGL_LN_PAIR),
("qwen_edit_ln_192x3072", "layernorm", (1, 192, 3072), SGL_LN_PAIR),
("qwen_edit_ln_200x3072", "layernorm", (1, 200, 3072), SGL_LN_PAIR),
("qwen_edit_ln_203x3072", "layernorm", (1, 203, 3072), SGL_LN_PAIR),
("qwen_edit_ln_8308x3072", "layernorm", (1, 8308, 3072), TORCH_LN),
("qwen_edit_rms_189x3584", "rmsnorm", (1, 189, 3584), SGL_RMS),
("qwen_edit_rms_192x3584", "rmsnorm", (1, 192, 3584), SGL_RMS),
("qwen_edit_rms_200x3584", "rmsnorm", (1, 200, 3584), SGL_RMS),
("qwen_edit_rms_203x3584", "rmsnorm", (1, 203, 3584), SGL_RMS),
],
),
(
@@ -93,9 +93,7 @@ ACTUAL_DIFFUSION_GROUPS: list[
("zimage_rms_32x3840", "rmsnorm", (1, 32, 3840), SGL_RMS),
("zimage_rms_4096x3840", "rmsnorm", (1, 4096, 3840), SGL_RMS),
("zimage_rms_4128x3840", "rmsnorm", (1, 4128, 3840), SGL_RMS),
("zimage_rms_512x2560", "rmsnorm", (1, 512, 2560), SGL_RMS),
("zimage_rms_512x32x128", "rmsnorm", (1, 512, 32, 128), SGL_RMS),
("zimage_rms_512x8x128", "rmsnorm", (1, 512, 8, 128), SGL_RMS),
("zimage_rms_32x2560", "rmsnorm", (32, 2560), SGL_RMS),
],
),
(
python/sglang/jit_kernel/diffusion/triton/rmsnorm_onepass.py
+14 -5
View File
@@ -2,6 +2,8 @@ import torch
import triton # type: ignore
import triton.language as tl # type: ignore
from sglang.srt.utils.custom_op import register_custom_op
# Adapted from https://github.com/ModelTC/LightX2V/blob/main/lightx2v/common/ops/norm/triton_ops.py#L905-L956
@triton.jit
@@ -33,7 +35,10 @@ def _rms_norm_tiled_onepass(
tl.store(y_blk, x * rstd * w, mask=mask)
def triton_one_pass_rms_norm(x: torch.Tensor, w: torch.Tensor, eps: float = 1e-6):
@register_custom_op(op_name="triton_one_pass_rms_norm_cuda", out_shape="x")
def _triton_one_pass_rms_norm_cuda(
x: torch.Tensor, w: torch.Tensor, eps: float = 1e-6
) -> torch.Tensor:
shape = x.shape
x = x.contiguous()
y = torch.empty_like(x)
@@ -41,11 +46,11 @@ def triton_one_pass_rms_norm(x: torch.Tensor, w: torch.Tensor, eps: float = 1e-6
y_view = y.reshape(-1, shape[-1])
S, D = x_view.shape
BLOCK_SIZE_SEQ = min(16, triton.next_power_of_2(max(1, S // 512)))
grid = (triton.cdiv(S, BLOCK_SIZE_SEQ),)
block_size_seq = min(16, triton.next_power_of_2(max(1, S // 512)))
grid = (triton.cdiv(S, block_size_seq),)
with torch.get_device_module().device(x.device):
torch.library.wrap_triton(_rms_norm_tiled_onepass)[grid](
_rms_norm_tiled_onepass[grid](
y_view,
x_view,
w,
@@ -53,11 +58,15 @@ def triton_one_pass_rms_norm(x: torch.Tensor, w: torch.Tensor, eps: float = 1e-6
D,
eps,
BLOCK_SIZE_DIM=triton.next_power_of_2(D),
BLOCK_SIZE_SEQ=BLOCK_SIZE_SEQ,
BLOCK_SIZE_SEQ=block_size_seq,
)
return y
def triton_one_pass_rms_norm(x: torch.Tensor, w: torch.Tensor, eps: float = 1e-6):
return _triton_one_pass_rms_norm_cuda(x, w, eps)
from sglang.multimodal_gen.runtime.platforms import current_platform
if current_platform.is_mps():