From 22e4b3a81f6362123faac44d87e548a29e8f679f Mon Sep 17 00:00:00 2001 From: Xiaoyu Zhang <1182563586@qq.com> Date: Wed, 12 Aug 2026 12:12:00 +0800 Subject: [PATCH] [Diffusion] Avoid slow cuBLASLt GELU epilogue on SM120 (#34350) --- python/sglang/kernels/ops/diffusion/fused_linear_gelu.py | 8 ++++++++ 1 file changed, 8 insertions(+) diff --git a/python/sglang/kernels/ops/diffusion/fused_linear_gelu.py b/python/sglang/kernels/ops/diffusion/fused_linear_gelu.py index 64b2f33dc..59a678746 100644 --- a/python/sglang/kernels/ops/diffusion/fused_linear_gelu.py +++ b/python/sglang/kernels/ops/diffusion/fused_linear_gelu.py @@ -31,6 +31,7 @@ from typing import Any import torch import torch.nn as nn +from sglang.kernels.jit.utils import get_jit_cuda_arch from sglang.kernels.ops.diffusion.quality_gate import QualityGatedFusion from sglang.srt.utils.custom_op import register_custom_op @@ -131,6 +132,13 @@ def can_fuse_linear_gelu(linear: Any, x: torch.Tensor) -> bool: """Whether ``gelu(linear(x))`` can use the fused cublasLt epilogue now.""" if not (x.is_cuda and x.dtype in (torch.bfloat16, torch.float16)): return False + arch = get_jit_cuda_arch() + if arch.major * 10 + arch.minor >= 120: + # The cublasLt GELU epilogue selected by current SM120 PyTorch/CUDA + # builds is slower than GEMM + the native GELU kernel for the FLUX + # production shape (1, 512, 3072 -> 12288). Keep quality-gated sites + # on their existing eager path on RTX 5090; SM90 dispatch is unchanged. + return False if getattr(linear, "weight", None) is None or x.dtype != linear.weight.dtype: return False return can_fuse_linear_gelu_static(linear)