From 8c28cdd1165b006b5d49f2b2dc1e300309197ce7 Mon Sep 17 00:00:00 2001 From: Xiaoyu Zhang <1182563586@qq.com> Date: Mon, 31 Aug 2026 01:37:35 +0800 Subject: [PATCH] [Diffusion][Kernel] Fuse Wan2.2 NVFP4 bias + GELU on Blackwell (#37075) --- .../jit/csrc/elementwise/bias_gelu.cuh | 105 ++++++++++++++++++ .../sglang/kernels/ops/diffusion/__init__.py | 4 + .../diffusion/sites/nvfp4_bias_gelu_site.py | 67 +++++++++++ .../kernels/ops/elementwise/bias_gelu.py | 46 ++++++++ .../multimodal_gen/runtime/layers/mlp.py | 26 ++++- .../runtime/models/dits/wanvideo.py | 21 +++- .../pipelines_core/stages/denoising.py | 7 ++ .../test/unit/test_wan_gelu_mlp.py | 51 +++++++++ .../benchmark/elementwise/bench_bias_gelu.py | 41 +++++++ .../kernels/ops/elementwise/test_bias_gelu.py | 48 ++++++++ 10 files changed, 411 insertions(+), 5 deletions(-) create mode 100644 python/sglang/kernels/jit/csrc/elementwise/bias_gelu.cuh create mode 100644 python/sglang/kernels/ops/diffusion/sites/nvfp4_bias_gelu_site.py create mode 100644 python/sglang/kernels/ops/elementwise/bias_gelu.py create mode 100644 test/registered/kernels/benchmark/elementwise/bench_bias_gelu.py create mode 100644 test/registered/kernels/ops/elementwise/test_bias_gelu.py diff --git a/python/sglang/kernels/jit/csrc/elementwise/bias_gelu.cuh b/python/sglang/kernels/jit/csrc/elementwise/bias_gelu.cuh new file mode 100644 index 000000000..5f697383d --- /dev/null +++ b/python/sglang/kernels/jit/csrc/elementwise/bias_gelu.cuh @@ -0,0 +1,105 @@ +#include +#include + +#include +#include +#include +#include + +#include + +#include +#include +#include + +namespace sglang { + +SGL_DEVICE float gelu_tanh(float x) { + constexpr float kAlpha = 0.044715f; + constexpr float kBeta = 0.7978845608028654f; + const float cdf = 0.5f * (1.0f + tanhf(kBeta * (x + kAlpha * x * x * x))); + return x * cdf; +} + +/** + * \brief Add a row-wise bias and apply approximate GELU. + * + * The intermediate bias result is rounded to the input dtype before GELU to + * preserve the eager add-then-GELU numerical boundary. + */ +template +__global__ void bias_gelu_tanh_kernel( + const T* __restrict__ input, + const T* __restrict__ bias, + T* __restrict__ output, + int64_t num_vecs, + int64_t row_vecs) { + using vec_t = device::AlignedVector; + + device::PDLWaitPrimary(); + const int64_t stride = static_cast(blockDim.x) * gridDim.x; + for (int64_t vec_id = static_cast(blockIdx.x) * blockDim.x + threadIdx.x; vec_id < num_vecs; + vec_id += stride) { + vec_t x; + vec_t b; + x.load(input, vec_id); + b.load(bias, vec_id % row_vecs); + + vec_t result; +#pragma unroll + for (int i = 0; i < kVecN; ++i) { + const float x_f32 = device::cast(x[i]); + const float bias_f32 = device::cast(b[i]); + const T biased = device::cast(x_f32 + bias_f32); + result[i] = device::cast(gelu_tanh(device::cast(biased))); + } + result.store(output, vec_id); + } + device::PDLTriggerSecondary(); +} + +/** + * \brief Validate and launch row-wise bias plus approximate GELU. + * + * \param input Contiguous two-dimensional input tensor. + * \param bias Contiguous bias matching the final input dimension. + * \param output Contiguous two-dimensional output tensor. + */ +template +void bias_gelu_tanh(tvm::ffi::TensorView input, tvm::ffi::TensorView bias, tvm::ffi::TensorView output) { + using namespace host; + + auto num_rows = SymbolicSize{"num_rows"}; + auto hidden_dim = SymbolicSize{"hidden_dim"}; + auto device_ = SymbolicDevice{}; + device_.set_options(); + + TensorMatcher({num_rows, hidden_dim}).with_dtype().with_device(device_).verify(input); + TensorMatcher({hidden_dim}).with_dtype().with_device(device_).verify(bias); + TensorMatcher({num_rows, hidden_dim}).with_dtype().with_device(device_).verify(output); + + constexpr int kVecN = device::kMaxVecBytes / sizeof(T); + const int64_t rows = num_rows.unwrap(); + const int64_t width = hidden_dim.unwrap(); + CHECK_HOST(rows > 0) << "bias_gelu_tanh: num_rows must be positive"; + CHECK_HOST(width > 0 && width % kVecN == 0) + << "bias_gelu_tanh: hidden_dim must be positive and divisible by " << kVecN; + + const int64_t row_vecs = width / kVecN; + const int64_t num_vecs = rows * row_vecs; + constexpr int64_t kBlockSize = 256; + const auto kernel = bias_gelu_tanh_kernel; + const int64_t occupancy = runtime::get_blocks_per_sm(kernel, kBlockSize); + const int64_t num_sms = runtime::get_sm_count(device_.unwrap().device_id); + const int64_t grid = std::min(num_sms * occupancy, div_ceil(num_vecs, kBlockSize)); + LaunchKernel(grid, kBlockSize, device_.unwrap()) + .enable_pdl(kUsePDL)( + kernel, + static_cast(input.data_ptr()), + static_cast(bias.data_ptr()), + static_cast(output.data_ptr()), + num_vecs, + row_vecs); +} + +} // namespace sglang diff --git a/python/sglang/kernels/ops/diffusion/__init__.py b/python/sglang/kernels/ops/diffusion/__init__.py index 67793a51d..ea8486b5d 100644 --- a/python/sglang/kernels/ops/diffusion/__init__.py +++ b/python/sglang/kernels/ops/diffusion/__init__.py @@ -463,6 +463,10 @@ _EXPORTS: dict[str, str] = { "mark_fused_gelu_site": "sites.fused_linear_gelu_site", "mount_fused_linear_gelu": "sites.fused_linear_gelu_site", "unmount_fused_linear_gelu": "sites.fused_linear_gelu_site", + "mark_nvfp4_bias_gelu_site": "sites.nvfp4_bias_gelu_site", + "mount_nvfp4_bias_gelu": "sites.nvfp4_bias_gelu_site", + "nvfp4_bias_gelu_active": "sites.nvfp4_bias_gelu_site", + "unmount_nvfp4_bias_gelu": "sites.nvfp4_bias_gelu_site", "can_use_ln_modulate": "sites.fused_ln_modulate_site", "fused_ln_modulate": "sites.fused_ln_modulate_site", "fused_ln_modulate_active": "sites.fused_ln_modulate_site", diff --git a/python/sglang/kernels/ops/diffusion/sites/nvfp4_bias_gelu_site.py b/python/sglang/kernels/ops/diffusion/sites/nvfp4_bias_gelu_site.py new file mode 100644 index 000000000..0d800a117 --- /dev/null +++ b/python/sglang/kernels/ops/diffusion/sites/nvfp4_bias_gelu_site.py @@ -0,0 +1,67 @@ +"""Request-scoped Wan NVFP4 bias+GELU fusion. + +The fused JIT kernel is bit-exact with the local eager ``add + GELU`` chain for +eligible ModelOpt FP4 linears. It still changes the model's kernel schedule, so +keep the default ``quality="lossless"`` path unchanged and mount this fast path +only for the existing ``quality="high"`` contract. +""" + +from __future__ import annotations + +import logging + +import torch.nn as nn + +from sglang.kernels.ops.diffusion.sites.quality_gate import QualityGatedFusion + +logger = logging.getLogger(__name__) + +_FUSION = QualityGatedFusion( + name="Wan NVFP4 bias+GELU", + marker_attr="_sgl_nvfp4_bias_gelu_site", + enabled_attr="_sgl_nvfp4_bias_gelu_enabled", +) + + +def mark_nvfp4_bias_gelu_site(module: nn.Module) -> None: + """Mark an MLP whose ``fc_in`` bias can be deferred for this fusion.""" + _FUSION.mark(module) + + +def nvfp4_bias_gelu_active(module: nn.Module) -> bool: + """Whether the quality-gated fusion is mounted on ``module``.""" + return _FUSION.is_enabled(module) + + +def _site_reject_reason(site: nn.Module) -> str | None: + if not getattr(site, "fuse_bias_gelu_tanh", False): + return "site is not an NVFP4 fused-GELU target" + linear = getattr(site, "fc_in", None) + if linear is None: + return "missing fc_in" + + from sglang.multimodal_gen.runtime.layers.quantization.modelopt_quant import ( + ModelOptFp4LinearMethod, + ) + + if not isinstance(getattr(linear, "quant_method", None), ModelOptFp4LinearMethod): + return "fc_in is not ModelOpt NVFP4" + if getattr(linear, "bias", None) is None: + return "fc_in has no bias" + return None + + +def mount_nvfp4_bias_gelu(root: nn.Module) -> bool: + """Enable every eligible marked site under ``root``.""" + sites = list(_FUSION.iter_sites(root)) + mounted = _FUSION.mount(root, reject_reason=_site_reject_reason, logger=logger) + for site in sites: + site.fc_in.skip_bias_add = mounted + return mounted + + +def unmount_nvfp4_bias_gelu(root: nn.Module) -> None: + """Restore every marked site to its original linear+GELU path.""" + _FUSION.unmount(root) + for site in _FUSION.iter_sites(root): + site.fc_in.skip_bias_add = False diff --git a/python/sglang/kernels/ops/elementwise/bias_gelu.py b/python/sglang/kernels/ops/elementwise/bias_gelu.py new file mode 100644 index 000000000..ef3da192d --- /dev/null +++ b/python/sglang/kernels/ops/elementwise/bias_gelu.py @@ -0,0 +1,46 @@ +from __future__ import annotations + +from typing import TYPE_CHECKING + +import torch + +from sglang.kernels.jit.utils import ( + cache_once, + is_arch_support_pdl, + load_jit, + make_cpp_args, +) +from sglang.srt.utils.custom_op import register_custom_op + +if TYPE_CHECKING: + from tvm_ffi.module import Module + + +@cache_once +def _jit_bias_gelu_tanh_module(dtype: torch.dtype) -> Module: + if dtype not in (torch.float16, torch.bfloat16): + raise RuntimeError(f"bias_gelu_tanh does not support {dtype}") + args = make_cpp_args(dtype, is_arch_support_pdl()) + return load_jit( + "bias_gelu_tanh", + *args, + cuda_files=["elementwise/bias_gelu.cuh"], + cuda_wrappers=[("bias_gelu_tanh", f"bias_gelu_tanh<{args}>")], + ) + + +@register_custom_op(mutates_args=["output"]) +def _bias_gelu_tanh( + input: torch.Tensor, bias: torch.Tensor, output: torch.Tensor +) -> None: + input_2d = input.view(-1, input.shape[-1]) + output_2d = output.view_as(input_2d) + module = _jit_bias_gelu_tanh_module(input.dtype) + module.bias_gelu_tanh(input_2d, bias, output_2d) + + +def bias_gelu_tanh(input: torch.Tensor, bias: torch.Tensor) -> torch.Tensor: + """Add a row-wise bias and apply approximate GELU.""" + output = torch.empty_like(input) + _bias_gelu_tanh(input, bias, output) + return output diff --git a/python/sglang/multimodal_gen/runtime/layers/mlp.py b/python/sglang/multimodal_gen/runtime/layers/mlp.py index 11c1739fe..d582d0cc8 100644 --- a/python/sglang/multimodal_gen/runtime/layers/mlp.py +++ b/python/sglang/multimodal_gen/runtime/layers/mlp.py @@ -38,12 +38,15 @@ class MLP(nn.Module): dtype: torch.dtype | None = None, prefix: str = "", quant_config: QuantizationConfig = None, + fuse_bias_gelu_tanh: bool = False, ): super().__init__() + self.fuse_bias_gelu_tanh = fuse_bias_gelu_tanh self.fc_in = ColumnParallelLinear( input_dim, mlp_hidden_dim, bias=True, + skip_bias_add=False, gather_output=False, quant_config=quant_config, prefix=add_prefix("fc_in", prefix), @@ -61,9 +64,28 @@ class MLP(nn.Module): prefix=add_prefix("fc_out", prefix), ) + def _apply_activation( + self, + x: torch.Tensor, + bias: torch.Tensor | None, + *, + use_fused_bias_gelu: bool = False, + ) -> torch.Tensor: + if self.fuse_bias_gelu_tanh and bias is not None: + if ( + use_fused_bias_gelu + and x.is_cuda + and x.dtype in (torch.float16, torch.bfloat16) + ): + from sglang.kernels.ops.elementwise.bias_gelu import bias_gelu_tanh + + return bias_gelu_tanh(x, bias) + return self.act(x + bias) + return self.act(x) + def forward(self, x: torch.Tensor) -> torch.Tensor: - x, _ = self.fc_in(x) - x = self.act(x) + x, bias = self.fc_in(x) + x = self._apply_activation(x, bias) x, _ = self.fc_out(x) return x diff --git a/python/sglang/multimodal_gen/runtime/models/dits/wanvideo.py b/python/sglang/multimodal_gen/runtime/models/dits/wanvideo.py index 3333f73fb..b80b22702 100755 --- a/python/sglang/multimodal_gen/runtime/models/dits/wanvideo.py +++ b/python/sglang/multimodal_gen/runtime/models/dits/wanvideo.py @@ -17,6 +17,8 @@ from sglang.kernels.ops.diffusion import ( fused_linear_gelu_tanh, fused_temb_table_slices, mark_fused_gelu_site, + mark_nvfp4_bias_gelu_site, + nvfp4_bias_gelu_active, tensors_equal, ) from sglang.multimodal_gen.configs.models.dits import WanVideoConfig @@ -49,6 +51,9 @@ from sglang.multimodal_gen.runtime.layers.mlp import MLP from sglang.multimodal_gen.runtime.layers.quantization.configs.base_config import ( QuantizationConfig, ) +from sglang.multimodal_gen.runtime.layers.quantization.modelopt_quant import ( + ModelOptFp4LinearMethod, +) from sglang.multimodal_gen.runtime.layers.rotary_embedding import ( NDRotaryEmbedding, _apply_rotary_emb, @@ -81,7 +86,7 @@ if USE_AITER: class _WanGELUMLP(MLP): - """Wan FFN with a quality-gated cublasLt GELU epilogue.""" + """Wan FFN with request-scoped GELU fast paths.""" def __init__( self, @@ -96,15 +101,25 @@ class _WanGELUMLP(MLP): act_type="gelu_pytorch_tanh", prefix=prefix, quant_config=quant_config, + fuse_bias_gelu_tanh=False, ) mark_fused_gelu_site(self, "fc_in") + self.fuse_bias_gelu_tanh = isinstance( + self.fc_in.quant_method, ModelOptFp4LinearMethod + ) + if self.fuse_bias_gelu_tanh: + mark_nvfp4_bias_gelu_site(self) def forward(self, x: torch.Tensor) -> torch.Tensor: if fused_gelu_active(self) and can_use_linear_gelu(self.fc_in, x): x = fused_linear_gelu_tanh(x, self.fc_in.weight, self.fc_in.bias) else: - x, _ = self.fc_in(x) - x = self.act(x) + x, bias = self.fc_in(x) + x = self._apply_activation( + x, + bias, + use_fused_bias_gelu=nvfp4_bias_gelu_active(self), + ) x, _ = self.fc_out(x) return x diff --git a/python/sglang/multimodal_gen/runtime/pipelines_core/stages/denoising.py b/python/sglang/multimodal_gen/runtime/pipelines_core/stages/denoising.py index 24f36549e..39b179915 100644 --- a/python/sglang/multimodal_gen/runtime/pipelines_core/stages/denoising.py +++ b/python/sglang/multimodal_gen/runtime/pipelines_core/stages/denoising.py @@ -27,6 +27,7 @@ from sglang.kernels.ops.diffusion import ( mount_hunyuan_qknorm, mount_lingbot_video_rmsnorm, mount_ltx2_rms_norm_modulate, + mount_nvfp4_bias_gelu, mount_sana_video_linear_attention, unmount_fused_gate_rmsnorm, unmount_fused_linear_gelu, @@ -34,6 +35,7 @@ from sglang.kernels.ops.diffusion import ( unmount_hunyuan_qknorm, unmount_lingbot_video_rmsnorm, unmount_ltx2_rms_norm_modulate, + unmount_nvfp4_bias_gelu, unmount_sana_video_linear_attention, ) from sglang.multimodal_gen import envs @@ -166,6 +168,11 @@ _QUALITY_FUSION_HANDLERS: tuple[ mount_fused_linear_gelu, unmount_fused_linear_gelu, ), + ( + "Wan NVFP4 fused bias+GELU", + mount_nvfp4_bias_gelu, + unmount_nvfp4_bias_gelu, + ), ( "fused LN+modulate (affine folding)", mount_fused_ln_modulate, diff --git a/python/sglang/multimodal_gen/test/unit/test_wan_gelu_mlp.py b/python/sglang/multimodal_gen/test/unit/test_wan_gelu_mlp.py index 5fd84f2f2..c6bc41325 100644 --- a/python/sglang/multimodal_gen/test/unit/test_wan_gelu_mlp.py +++ b/python/sglang/multimodal_gen/test/unit/test_wan_gelu_mlp.py @@ -4,12 +4,18 @@ import torch from sglang.kernels.ops.diffusion import ( fused_gelu_active, mount_fused_linear_gelu, + mount_nvfp4_bias_gelu, + nvfp4_bias_gelu_active, unmount_fused_linear_gelu, + unmount_nvfp4_bias_gelu, ) from sglang.multimodal_gen.runtime.distributed.parallel_state import ( maybe_init_distributed_environment_and_model_parallel, model_parallel_is_initialized, ) +from sglang.multimodal_gen.runtime.layers.quantization.modelopt_quant import ( + ModelOptFp4Config, +) from sglang.multimodal_gen.runtime.models.dits.wanvideo import _WanGELUMLP from sglang.multimodal_gen.test.single_test_file.component_accuracy.utils import ( ensure_distributed_env_defaults, @@ -23,6 +29,12 @@ def _ensure_single_process_parallel_runtime() -> None: maybe_init_distributed_environment_and_model_parallel(tp_size=1, sp_size=1) +requires_blackwell = pytest.mark.skipif( + not torch.cuda.is_available() or torch.cuda.get_device_capability()[0] < 10, + reason="requires a Blackwell CUDA GPU", +) + + @pytest.mark.skipif(not torch.cuda.is_available(), reason="requires CUDA") @torch.no_grad() def test_wan_gelu_mlp_quality_path_and_lossless_restore(): @@ -44,3 +56,42 @@ def test_wan_gelu_mlp_quality_path_and_lossless_restore(): unmount_fused_linear_gelu(mlp) assert not fused_gelu_active(mlp) assert torch.equal(mlp(x), reference) + + +@requires_blackwell +def test_wan_nvfp4_mlp_defers_bias_for_gelu_fusion(): + _ensure_single_process_parallel_runtime() + quant_config = ModelOptFp4Config( + is_checkpoint_nvfp4_serialized=True, + group_size=16, + ) + + mlp = _WanGELUMLP(64, 256, prefix="", quant_config=quant_config) + + assert mlp.fuse_bias_gelu_tanh + assert not mlp.fc_in.skip_bias_add + assert not mlp.fc_out.skip_bias_add + assert not nvfp4_bias_gelu_active(mlp) + + assert mount_nvfp4_bias_gelu(mlp) + assert nvfp4_bias_gelu_active(mlp) + assert mlp.fc_in.skip_bias_add + unmount_nvfp4_bias_gelu(mlp) + assert not nvfp4_bias_gelu_active(mlp) + assert not mlp.fc_in.skip_bias_add + + +@requires_blackwell +def test_wan_nvfp4_mlp_does_not_mark_excluded_linear(): + _ensure_single_process_parallel_runtime() + quant_config = ModelOptFp4Config( + is_checkpoint_nvfp4_serialized=True, + group_size=16, + exclude_modules=["fc_in"], + ) + + mlp = _WanGELUMLP(64, 256, prefix="", quant_config=quant_config) + + assert not mlp.fuse_bias_gelu_tanh + assert not mlp.fc_in.skip_bias_add + assert not mount_nvfp4_bias_gelu(mlp) diff --git a/test/registered/kernels/benchmark/elementwise/bench_bias_gelu.py b/test/registered/kernels/benchmark/elementwise/bench_bias_gelu.py new file mode 100644 index 000000000..5c622b11d --- /dev/null +++ b/test/registered/kernels/benchmark/elementwise/bench_bias_gelu.py @@ -0,0 +1,41 @@ +import torch +import torch.nn.functional as F + +from sglang.kernels.jit.benchmark import marker +from sglang.kernels.ops.elementwise.bias_gelu import bias_gelu_tanh +from sglang.test.ci.ci_register import register_cuda_ci + +register_cuda_ci( + est_time=12, stage="base-b-kernel-benchmark", runner_config="1-gpu-large" +) + + +def torch_bias_gelu(input: torch.Tensor, bias: torch.Tensor) -> torch.Tensor: + return F.gelu(input + bias, approximate="tanh") + + +FN_MAP = { + "jit": bias_gelu_tanh, + "torch": torch_bias_gelu, +} + + +@marker.parametrize( + "rows,hidden_dim", + [(32760, 5120), (32760, 13824)], + [(4096, 13824)], +) +@marker.benchmark("impl", ["jit", "torch"]) +def benchmark(rows: int, hidden_dim: int, impl: str): + input = torch.randn(rows, hidden_dim, dtype=torch.bfloat16, device="cuda") + bias = torch.randn(hidden_dim, dtype=torch.bfloat16, device="cuda") + return marker.do_bench( + FN_MAP[impl], + input_args=(input, bias), + memory_args=(input, bias), + memory_output="out", + ) + + +if __name__ == "__main__": + benchmark.run() diff --git a/test/registered/kernels/ops/elementwise/test_bias_gelu.py b/test/registered/kernels/ops/elementwise/test_bias_gelu.py new file mode 100644 index 000000000..6ff815958 --- /dev/null +++ b/test/registered/kernels/ops/elementwise/test_bias_gelu.py @@ -0,0 +1,48 @@ +import sys + +import pytest +import torch +import torch.nn.functional as F + +from sglang.kernels.ops.elementwise.bias_gelu import bias_gelu_tanh +from sglang.test.ci.ci_register import register_cuda_ci + +register_cuda_ci(est_time=20, stage="base-b-kernel-unit", runner_config="1-gpu-large") +register_cuda_ci(est_time=20, stage="base-b-kernel-unit", runner_config="4-gpu-b200") + + +@pytest.mark.parametrize("dtype", [torch.float16, torch.bfloat16]) +@pytest.mark.parametrize("shape", [(1, 128), (2, 7, 512), (1, 4096, 13824)]) +def test_bias_gelu_tanh_is_bit_exact(dtype: torch.dtype, shape: tuple[int, ...]): + torch.manual_seed(0) + input = torch.randn(shape, device="cuda", dtype=dtype) + bias = torch.randn(shape[-1], device="cuda", dtype=dtype) + original_input = input.clone() + + expected = F.gelu(input + bias, approximate="tanh") + actual = bias_gelu_tanh(input, bias) + + assert actual.shape == input.shape + assert actual.data_ptr() != input.data_ptr() + assert torch.equal(input, original_input) + assert torch.equal(actual, expected) + + +def test_bias_gelu_tanh_rejects_unsupported_width(): + input = torch.randn(2, 127, device="cuda", dtype=torch.bfloat16) + bias = torch.randn(127, device="cuda", dtype=torch.bfloat16) + + with pytest.raises(RuntimeError, match="hidden_dim"): + bias_gelu_tanh(input, bias) + + +def test_bias_gelu_tanh_rejects_unsupported_dtype(): + input = torch.ones(2, 128, device="cuda", dtype=torch.int32) + bias = torch.ones(128, device="cuda", dtype=torch.int32) + + with pytest.raises(RuntimeError, match="does not support"): + bias_gelu_tanh(input, bias) + + +if __name__ == "__main__": + sys.exit(pytest.main([__file__, "-v", "-s"]))