[Diffusion] Only import sgl_kernel in custom op cuda path (SiluAndMul and RMSNorm) (#15592)

Signed-off-by: Xiaodong Ye <xiaodong.ye@mthreads.com>

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
 

python/sglang/multimodal_gen/runtime/layers/layernorm.py

@@ -8,6 +8,11 @@ from typing import Optional, Tuple, Union
 import torch
 import torch.nn as nn
 import torch.nn.functional as F
+
+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