diff --git a/python/sglang/multimodal_gen/runtime/layers/custom_op.py b/python/sglang/multimodal_gen/runtime/layers/custom_op.py index f4cd030d6..f6c797c64 100644 --- a/python/sglang/multimodal_gen/runtime/layers/custom_op.py +++ b/python/sglang/multimodal_gen/runtime/layers/custom_op.py @@ -12,6 +12,7 @@ from sglang.multimodal_gen.runtime.platforms import current_platform from sglang.multimodal_gen.runtime.utils.logging_utils import init_logger logger = init_logger(__name__) +_is_cuda = current_platform.is_cuda() class CustomOp(nn.Module): @@ -58,7 +59,7 @@ class CustomOp(nn.Module): return self.forward_native(*args, **kwargs) def dispatch_forward(self) -> Callable: - if current_platform.is_cuda(): + if _is_cuda: return self.forward_cuda elif current_platform.is_hip(): return self.forward_hip diff --git a/python/sglang/multimodal_gen/runtime/layers/layernorm.py b/python/sglang/multimodal_gen/runtime/layers/layernorm.py index 27f040054..86622ad17 100644 --- a/python/sglang/multimodal_gen/runtime/layers/layernorm.py +++ b/python/sglang/multimodal_gen/runtime/layers/layernorm.py @@ -23,8 +23,11 @@ from sglang.multimodal_gen.runtime.layers.triton_ops import ( rms_norm_fn, triton_one_pass_rms_norm, ) +from sglang.multimodal_gen.runtime.platforms import current_platform from sglang.multimodal_gen.runtime.utils.common import get_bool_env_var +_is_cuda = current_platform.is_cuda() + # Copied and adapted from sglang @CustomOp.register("rms_norm") @@ -453,7 +456,7 @@ def apply_qk_norm( k_eps = k_norm.variance_epsilon # Only try fused path on CUDA and when it won't introduce implicit copies. if ( - q.is_cuda + _is_cuda and allow_inplace and (q_eps == k_eps) and can_use_fused_inplace_qknorm(head_dim, q.dtype) diff --git a/python/sglang/multimodal_gen/runtime/models/dits/flux.py b/python/sglang/multimodal_gen/runtime/models/dits/flux.py index f46818ca3..e6891cd6f 100644 --- a/python/sglang/multimodal_gen/runtime/models/dits/flux.py +++ b/python/sglang/multimodal_gen/runtime/models/dits/flux.py @@ -49,6 +49,7 @@ from sglang.multimodal_gen.runtime.utils.layerwise_offload import OffloadableDiT from sglang.multimodal_gen.runtime.utils.logging_utils import init_logger logger = init_logger(__name__) # pylint: disable=invalid-name +_is_cuda = current_platform.is_cuda() def _get_qkv_projections( @@ -167,7 +168,7 @@ class FluxAttention(torch.nn.Module, AttentionModuleMixin): key = key.unflatten(-1, (self.heads, -1)) value = value.unflatten(-1, (self.heads, -1)) if ( - query.is_cuda + _is_cuda and (self.norm_q.variance_epsilon == self.norm_k.variance_epsilon) and can_use_fused_inplace_qknorm(self.head_dim, query.dtype) ): @@ -189,7 +190,7 @@ class FluxAttention(torch.nn.Module, AttentionModuleMixin): encoder_value = encoder_value.unflatten(-1, (self.heads, -1)) if ( - encoder_query.is_cuda + _is_cuda and ( self.norm_added_q.variance_epsilon == self.norm_added_k.variance_epsilon diff --git a/python/sglang/multimodal_gen/runtime/models/dits/flux_2.py b/python/sglang/multimodal_gen/runtime/models/dits/flux_2.py index c1b88f154..907c39888 100644 --- a/python/sglang/multimodal_gen/runtime/models/dits/flux_2.py +++ b/python/sglang/multimodal_gen/runtime/models/dits/flux_2.py @@ -35,6 +35,7 @@ from sglang.multimodal_gen.runtime.utils.layerwise_offload import OffloadableDiT from sglang.multimodal_gen.runtime.utils.logging_utils import init_logger logger = init_logger(__name__) # pylint: disable=invalid-name +_is_cuda = current_platform.is_cuda() def _get_qkv_projections( @@ -198,7 +199,7 @@ class Flux2Attention(torch.nn.Module, AttentionModuleMixin): value = value.unflatten(-1, (self.heads, -1)) if ( - query.is_cuda + _is_cuda and (self.norm_q.variance_epsilon == self.norm_k.variance_epsilon) and can_use_fused_inplace_qknorm(self.head_dim, query.dtype) ): @@ -220,7 +221,7 @@ class Flux2Attention(torch.nn.Module, AttentionModuleMixin): encoder_value = encoder_value.unflatten(-1, (self.heads, -1)) if ( - encoder_query.is_cuda + _is_cuda and ( self.norm_added_q.variance_epsilon == self.norm_added_k.variance_epsilon diff --git a/python/sglang/multimodal_gen/runtime/models/dits/qwen_image.py b/python/sglang/multimodal_gen/runtime/models/dits/qwen_image.py index 7d0e3847b..5d5a795d2 100644 --- a/python/sglang/multimodal_gen/runtime/models/dits/qwen_image.py +++ b/python/sglang/multimodal_gen/runtime/models/dits/qwen_image.py @@ -31,11 +31,15 @@ from sglang.multimodal_gen.runtime.layers.triton_ops import ( fuse_scale_shift_kernel, ) from sglang.multimodal_gen.runtime.models.dits.base import CachableDiT -from sglang.multimodal_gen.runtime.platforms import AttentionBackendEnum +from sglang.multimodal_gen.runtime.platforms import ( + AttentionBackendEnum, + current_platform, +) from sglang.multimodal_gen.runtime.utils.layerwise_offload import OffloadableDiTMixin from sglang.multimodal_gen.runtime.utils.logging_utils import init_logger logger = init_logger(__name__) # pylint: disable=invalid-name +_is_cuda = current_platform.is_cuda() def _get_qkv_projections( @@ -684,7 +688,7 @@ class QwenImageTransformerBlock(nn.Module): ) gate0, gate1 = gate[:actual_batch], gate[actual_batch : 2 * actual_batch] - if x.is_cuda: + if _is_cuda: if not x.is_contiguous(): x = x.contiguous() if not index.is_contiguous(): diff --git a/python/sglang/multimodal_gen/runtime/models/dits/wanvideo.py b/python/sglang/multimodal_gen/runtime/models/dits/wanvideo.py index 99e88a849..d90d96b1c 100644 --- a/python/sglang/multimodal_gen/runtime/models/dits/wanvideo.py +++ b/python/sglang/multimodal_gen/runtime/models/dits/wanvideo.py @@ -54,6 +54,7 @@ from sglang.multimodal_gen.runtime.utils.layerwise_offload import OffloadableDiT from sglang.multimodal_gen.runtime.utils.logging_utils import init_logger logger = init_logger(__name__) +_is_cuda = current_platform.is_cuda() class WanImageEmbedding(torch.nn.Module): @@ -444,7 +445,7 @@ class WanTransformerBlock(nn.Module): # Apply rotary embeddings cos, sin = freqs_cis - if query.is_cuda and query.shape == key.shape: + if _is_cuda and query.shape == key.shape: cos_sin_cache = torch.cat( [ cos.to(dtype=torch.float32).contiguous(), @@ -626,7 +627,7 @@ class WanTransformerBlock_VSA(nn.Module): # Apply rotary embeddings cos, sin = freqs_cis - if query.is_cuda and query.shape == key.shape: + if _is_cuda and query.shape == key.shape: cos_sin_cache = torch.cat( [ cos.to(dtype=torch.float32).contiguous(), diff --git a/python/sglang/multimodal_gen/runtime/models/dits/zimage.py b/python/sglang/multimodal_gen/runtime/models/dits/zimage.py index 2addd365d..8bdac4807 100644 --- a/python/sglang/multimodal_gen/runtime/models/dits/zimage.py +++ b/python/sglang/multimodal_gen/runtime/models/dits/zimage.py @@ -26,6 +26,7 @@ from sglang.multimodal_gen.runtime.utils.layerwise_offload import OffloadableDiT from sglang.multimodal_gen.runtime.utils.logging_utils import init_logger logger = init_logger(__name__) +_is_cuda = current_platform.is_cuda() ADALN_EMBED_DIM = 256 SEQ_MULTI_OF = 32 @@ -171,7 +172,7 @@ class ZImageAttention(nn.Module): if self.qk_norm: if ( - q.is_cuda + _is_cuda and (self.norm_q.variance_epsilon == self.norm_k.variance_epsilon) and can_use_fused_inplace_qknorm(self.head_dim, q.dtype) ): @@ -189,7 +190,7 @@ class ZImageAttention(nn.Module): if freqs_cis is not None: cos, sin = freqs_cis - if q.is_cuda and q.shape == k.shape: + if _is_cuda and q.shape == k.shape: cos_sin_cache = torch.cat( [ cos.to(dtype=torch.float32).contiguous(),