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

Co-authored-by: Mick <mickjagger19@icloud.com>

python/sglang/jit_kernel/diffusion/triton/rmsnorm_onepass.py

@@ -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():