From acf724b036c595292b036942db69fb169efffc45 Mon Sep 17 00:00:00 2001 From: R0CKSTAR Date: Wed, 4 Feb 2026 04:42:58 +0800 Subject: [PATCH] [Diffusion] Only import sgl_kernel in custom op cuda path (SiluAndMul and RMSNorm) (#15592) Signed-off-by: Xiaodong Ye --- python/sglang/multimodal_gen/runtime/layers/activation.py | 7 ++++++- python/sglang/multimodal_gen/runtime/layers/layernorm.py | 7 ++++++- 2 files changed, 12 insertions(+), 2 deletions(-) diff --git a/python/sglang/multimodal_gen/runtime/layers/activation.py b/python/sglang/multimodal_gen/runtime/layers/activation.py index afbad4ed3..668a5b2f1 100644 --- a/python/sglang/multimodal_gen/runtime/layers/activation.py +++ b/python/sglang/multimodal_gen/runtime/layers/activation.py @@ -9,8 +9,13 @@ from typing import Any import torch import torch.nn as nn import torch.nn.functional as F -from sgl_kernel import silu_and_mul +from sglang.multimodal_gen.runtime.platforms import current_platform + +_is_cuda = current_platform.is_cuda() +_is_hip = current_platform.is_hip() +if _is_cuda or _is_hip: + from sgl_kernel import silu_and_mul # TODO (will): remove this dependency from sglang.multimodal_gen.runtime.layers.custom_op import CustomOp diff --git a/python/sglang/multimodal_gen/runtime/layers/layernorm.py b/python/sglang/multimodal_gen/runtime/layers/layernorm.py index 0fb825261..03989a55e 100644 --- a/python/sglang/multimodal_gen/runtime/layers/layernorm.py +++ b/python/sglang/multimodal_gen/runtime/layers/layernorm.py @@ -8,7 +8,12 @@ from typing import Optional, Tuple, Union import torch import torch.nn as nn import torch.nn.functional as F -from sgl_kernel import fused_add_rmsnorm, rmsnorm + +from sglang.multimodal_gen.runtime.platforms import current_platform + +_is_cuda = current_platform.is_cuda() +if _is_cuda: + from sgl_kernel import fused_add_rmsnorm, rmsnorm from sglang.jit_kernel.norm import can_use_fused_inplace_qknorm, fused_inplace_qknorm from sglang.multimodal_gen.runtime.distributed.parallel_state import (