[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:
@@ -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
|
||||
|
||||
|
||||
@@ -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 (
|
||||
|
||||
Reference in New Issue
Block a user