[diffusion] fix: Fix FLUX.1/2 graph breaks (#23648)

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

@@ -41,6 +41,7 @@ from sglang.multimodal_gen.runtime.utils.logging_utils import init_logger
 
 logger = init_logger(__name__)
 
+IS_AMP_SUPPORTED = current_platform.is_amp_supported()
 WEIGHT_LOADER_V2_SUPPORTED = [
     "CompressedTensorsLinearMethod",
     "AWQMarlinLinearMethod",
@@ -156,7 +157,7 @@ class UnquantizedLinearMethod(LinearMethodBase):
     ) -> torch.Tensor:
         output = (
             F.linear(x, layer.weight, bias)
-            if current_platform.is_amp_supported() or bias is None
+            if IS_AMP_SUPPORTED or bias is None
             else F.linear(x, layer.weight, bias.to(x.dtype))
         )  # NOTE: explicit dtype cast for bias is needed on platforms where amp isn't supported
         return output

python/sglang/multimodal_gen/runtime/layers/rotary_embedding/utils.py

@@ -7,8 +7,11 @@ import torch
 from sglang.jit_kernel.diffusion.triton.rotary import apply_rotary_embedding
 from sglang.kernel_api_logging import debug_kernel_api
 from sglang.multimodal_gen.runtime.platforms import current_platform
+from sglang.multimodal_gen.runtime.utils.logging_utils import init_logger
 from sglang.srt.utils.custom_op import register_custom_op_from_extern
 
+logger = init_logger(__name__)
+
 _is_cuda = current_platform.is_cuda()
 if _is_cuda:
     try:
@@ -93,12 +96,9 @@ def apply_flashinfer_rope_qk_inplace(
 
     if flashinfer_apply_rope_inplace is None:
         # Triton fallback for AMD/ROCm where FlashInfer is not available
-        import warnings
 
-        warnings.warn(
-            "FlashInfer not available, using Triton fallback for RoPE",
-            stacklevel=2,
-        )
+        _warn_about_missing_flashinfer()
+
         half_size = cos_sin_cache.shape[-1] // 2
         if positions is None:
             cos = cos_sin_cache[:seqlen, :half_size].to(q.dtype)
@@ -141,3 +141,14 @@ def apply_flashinfer_rope_qk_inplace(
         is_neox=is_neox,
     )
     return q_flat.view(bsz, seqlen, nheads, d), k_flat.view(bsz, seqlen, nheads, d)
+
+
+@torch.compiler.assume_constant_result
+def _warn_about_missing_flashinfer():
+    """
+    Function to warn about the missing FlashInfer.
+    Exists to not cause a graph break during the compilation.
+    """
+    logger.warning_once(
+        "FlashInfer not available, using Triton fallback for RoPE",
+    )