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

Signed-off-by: Xiaodong Ye <xiaodong.ye@mthreads.com>
This commit is contained in:
R0CKSTAR
2026-02-04 04:42:58 +08:00
committed by GitHub
parent e166ca8758
commit acf724b036
2 changed files with 12 additions and 2 deletions

View File

@@ -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

View File

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