[Diffusion] Refactor diffusion is_cuda check (#17498)

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

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)

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

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

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():

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(),

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(),