minke.yu/sglang
1
0
Fork 0
Code Issues Pull Requests Actions 12 Packages Projects Releases Wiki Activity

[Diffusion] Fuse FLUX.2 NVFP4 FC1, SwiGLU, and FC2 quantization (#37096)

Browse Source 
Co-authored-by: Cursor <cursoragent@cursor.com>
This commit is contained in:
Xiaoyu Zhang
2026-09-02 08:21:14 +08:00
committed by GitHub
co-authored by Cursor
parent c593527f33
commit f4c17fed07
8 changed files with 333 additions and 13 deletions
Download Patch File Download Diff File Expand all files Collapse all files
python/sglang/kernels/ops/diffusion/__init__.py
+4
View File
@@ -548,6 +548,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",
"flux2_nvfp4_swiglu_quant_active": "sites.flux2_nvfp4_swiglu_quant_site",
"mark_flux2_nvfp4_swiglu_quant_site": "sites.flux2_nvfp4_swiglu_quant_site",
"mount_flux2_nvfp4_swiglu_quant": "sites.flux2_nvfp4_swiglu_quant_site",
"unmount_flux2_nvfp4_swiglu_quant": "sites.flux2_nvfp4_swiglu_quant_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",
python/sglang/kernels/ops/diffusion/sites/flux2_nvfp4_swiglu_quant_site.py
+36
View File
@@ -0,0 +1,36 @@
"""Request-scoped gate for the FLUX.2 NVFP4 SwiGLU fusion.
The fused FC1 + SwiGLU + FC2-input quantization path changes the rounding
order by quantizing before the reference BF16 intermediate is materialized.
Keep it disabled for the lossless default and mount it only for
``quality="high"`` requests at denoising batch boundaries.
"""
from __future__ import annotations
from torch import nn
from sglang.kernels.ops.diffusion.sites.quality_gate import QualityGatedFusion
_FUSION = QualityGatedFusion(
name="FLUX.2 NVFP4 FC1+SwiGLU+quant",
marker_attr="_sgl_flux2_nvfp4_swiglu_quant_site",
enabled_attr="_sgl_flux2_nvfp4_swiglu_quant_enabled",
)
def mark_flux2_nvfp4_swiglu_quant_site(module: nn.Module) -> None:
"""Mark an eligible FLUX.2 feed-forward site, disabled by default."""
_FUSION.mark(module)
def flux2_nvfp4_swiglu_quant_active(module: nn.Module) -> bool:
return _FUSION.is_enabled(module)
def mount_flux2_nvfp4_swiglu_quant(root: nn.Module) -> bool:
return _FUSION.mount(root)
def unmount_flux2_nvfp4_swiglu_quant(root: nn.Module) -> None:
_FUSION.unmount(root)
python/sglang/multimodal_gen/runtime/layers/quantization/modelopt_quant.py
+129 -12
View File
@@ -94,6 +94,69 @@ def _swizzled_nvfp4_scales_to_linear(scales: torch.Tensor) -> torch.Tensor:
return linear.squeeze(0) if scale_ndim == 2 else linear
def _prepare_nvfp4_swiglu_fusion_weights(
layer: torch.nn.Module,
weight: torch.Tensor,
scales: torch.Tensor,
) -> None:
from sglang.kernels.ops.quantization.nvfp4_gemm_swiglu_nvfp4_quant import (
interleave_linear_and_gate,
swizzle_blockscale_2d,
)
weight, weights_padding_cols = pad_nvfp4_weight(weight)
if weights_padding_cols != 0:
raise ValueError(
"Fused NVFP4 SwiGLU does not support K-padded weights; "
f"got weights_padding_cols={weights_padding_cols}."
)
if scales.ndim != 2:
raise ValueError(
"Fused NVFP4 SwiGLU expects a 2D weight scale, "
f"got shape={tuple(scales.shape)}."
)
if scales.shape[0] != weight.shape[0]:
raise ValueError(
"Fused NVFP4 SwiGLU does not support N padding; "
f"scale rows={scales.shape[0]} vs weight rows={weight.shape[0]}."
)
if weight.shape[0] % 128 != 0:
raise ValueError(
"Fused NVFP4 SwiGLU requires FC1 N % 128 == 0, " f"got N={weight.shape[0]}."
)
# FLUX.2 stores [gate; up]. The kernel consumes 64-row groups in
# [up; gate] order and applies SiLU to the gate half.
gate_weight, up_weight = weight.chunk(2, dim=0)
weight_swiglu_interleaved = interleave_linear_and_gate(
torch.cat((up_weight, gate_weight), dim=0), group_size=64, dim=0
)
gate_scale, up_scale = scales.chunk(2, dim=0)
weight_scale_swiglu_interleaved = swizzle_blockscale_2d(
interleave_linear_and_gate(
torch.cat((up_scale, gate_scale), dim=0), group_size=64, dim=0
)
)
for name, value in (
("weight_swiglu_interleaved", weight_swiglu_interleaved),
("weight_scale_swiglu_interleaved", weight_scale_swiglu_interleaved),
):
value = value.detach()
existing = layer._buffers.get(name)
if (
existing is not None
and existing.shape == value.shape
and existing.dtype == value.dtype
and existing.device == value.device
):
existing.copy_(value)
elif name in layer._buffers:
layer._buffers[name] = value
else:
layer.register_buffer(name, value, persistent=False)
layer._swiglu_fusion_ready = True
def _require_flashinfer():
if flashinfer is None:
raise RuntimeError(
@@ -658,6 +721,14 @@ class ModelOptFp4LinearMethod(LinearMethodBase):
):
scales = _swizzled_nvfp4_scales_to_linear(scales)
if getattr(layer, "_interleave_for_swiglu_fusion", False):
# The regular GEMM path swizzles this tensor below. The fused
# GEMM+SwiGLU epilogue needs the logical row order so gate/up rows
# can first be paired and then swizzled as one matrix.
_prepare_nvfp4_swiglu_fusion_weights(
layer, w_swapped, scales.detach().clone()
)
_, flashinfer_backend = _get_fp4_gemm_op()
if flashinfer_backend == "trtllm":
flashinfer_ops = _require_flashinfer()
@@ -730,23 +801,31 @@ class ModelOptFp4LinearMethod(LinearMethodBase):
def apply(
self,
layer: torch.nn.Module,
x: torch.Tensor,
x: torch.Tensor | tuple[torch.Tensor, torch.Tensor],
bias: Optional[torch.Tensor] = None,
) -> torch.Tensor:
output_dtype = x.dtype
input_shape = x.shape
x = x.view(-1, input_shape[-1])
output_size = layer.output_size_per_partition
output_shape = list(input_shape[:-1]) + [output_size]
if getattr(layer, "_accepts_prequantized_fp4", False) and isinstance(x, tuple):
x_fp4, x_scale_interleaved = x
output_dtype = layer.params_dtype
output_shape = [x_fp4.shape[0], output_size]
else:
if isinstance(x, tuple):
raise TypeError(
"Prequantized NVFP4 input requires the linear layer to opt in."
)
output_dtype = x.dtype
input_shape = x.shape
x = x.view(-1, input_shape[-1])
output_shape = list(input_shape[:-1]) + [output_size]
fp4_quantize = _get_fp4_quantize_op()
if fp4_quantize is None:
raise RuntimeError(
"No FP4 quantization kernel available. Install flashinfer."
)
fp4_quantize = _get_fp4_quantize_op()
if fp4_quantize is None:
raise RuntimeError(
"No FP4 quantization kernel available. Install flashinfer."
)
x_fp4, x_scale_interleaved = fp4_quantize(x, layer.input_scale_inv)
x_fp4, x_scale_interleaved = fp4_quantize(x, layer.input_scale_inv)
weights_padding_cols = getattr(layer, "weights_padding_cols", 0)
x_fp4 = pad_nvfp4_activation_for_cutlass(x_fp4, weights_padding_cols)
@@ -809,3 +888,41 @@ def apply_nvfp4_gemm_prequantized(
)
out = slice_nvfp4_output(out, layer.output_size_per_partition)
return out + bias if bias is not None else out
def apply_nvfp4_gemm_swiglu_quant(
linear_in: torch.nn.Module,
linear_out: torch.nn.Module,
x: torch.Tensor,
) -> torch.Tensor:
"""Run FLUX.2 FC1 GEMM + SwiGLU + FC2-input NVFP4 quantization."""
if not getattr(linear_in, "_swiglu_fusion_ready", False):
raise RuntimeError("NVFP4 SwiGLU weights were not prepared for fusion.")
if not getattr(linear_out, "_accepts_prequantized_fp4", False):
raise RuntimeError("NVFP4 output projection did not opt into packed input.")
fp4_quantize = _get_fp4_quantize_op()
if fp4_quantize is None:
raise RuntimeError("No FP4 quantization kernel available. Install flashinfer.")
from sglang.kernels.ops.quantization.nvfp4_gemm_swiglu_nvfp4_quant import (
nvfp4_gemm_swiglu_nvfp4_quant,
)
input_shape = x.shape
x_2d = x.view(-1, input_shape[-1])
x_fp4, x_scale_interleaved = fp4_quantize(x_2d, linear_in.input_scale_inv)
if x_scale_interleaved.dtype == torch.uint8:
x_scale_interleaved = x_scale_interleaved.view(torch.float8_e4m3fn)
out_fp4, out_scale_interleaved = nvfp4_gemm_swiglu_nvfp4_quant(
x_fp4,
x_scale_interleaved,
linear_in.weight_swiglu_interleaved,
linear_in.weight_scale_swiglu_interleaved,
linear_in.alpha,
linear_out.input_scale_inv,
enable_pdl=True,
)
out, _ = linear_out((out_fp4, out_scale_interleaved))
return out.view(*input_shape[:-1], out.shape[-1])
python/sglang/multimodal_gen/runtime/models/dits/flux_2.py
+34
View File
@@ -27,10 +27,12 @@ from sglang.kernels.ops.diffusion import (
can_use_flux2_gated_resnorm,
can_use_fused_layernorm_modulate,
flux2_gated_resnorm_raw,
flux2_nvfp4_swiglu_quant_active,
fused_layernorm_modulate_fp8_quant_raw,
fused_layernorm_modulate_raw,
fused_packed_silu_mul_bitexact,
is_plain_layer_norm,
mark_flux2_nvfp4_swiglu_quant_site,
residual_gate_add,
try_flux2_token_cat_fp8,
try_flux2_token_cat_nvfp4,
@@ -70,6 +72,7 @@ from sglang.multimodal_gen.runtime.layers.quantization.modelopt_quant import (
ModelOptFp8Config,
ModelOptFp8LinearMethod,
apply_nvfp4_gemm_prequantized,
apply_nvfp4_gemm_swiglu_quant,
)
from sglang.multimodal_gen.runtime.layers.rotary_embedding import (
NDRotaryEmbedding,
@@ -216,6 +219,13 @@ def _flux2_gated_resnorm(
return _flux2_norm_modulate(norm, residual, scale, shift), residual
def _can_use_nvfp4_swiglu_quant_fusion(capability: Any) -> bool:
# The end-to-end accuracy and performance validation for this fusion was
# done on SM103. SM100 currently produces a deterministic but materially
# different FLUX.2 image, so keep B200/GB200 on the existing unfused path.
return capability is not None and (capability.major, capability.minor) == (10, 3)
def _flux2_norm_maybe_fp8(
norm: nn.Module,
hidden_states: torch.Tensor | PendingGatedResidual,
@@ -405,7 +415,31 @@ class Flux2FeedForward(nn.Module):
prefix=f"{prefix}.linear_out" if prefix else "linear_out",
)
capability = current_platform.get_device_capability()
if (
_can_use_nvfp4_swiglu_quant_fusion(capability)
and isinstance(self.linear_in.quant_method, ModelOptFp4LinearMethod)
and isinstance(self.linear_out.quant_method, ModelOptFp4LinearMethod)
and self.linear_in.output_size_per_partition % 128 == 0
and self.linear_in.input_size_per_partition % 16 == 0
and self.linear_in.bias is None
and self.linear_out.bias is None
):
# These flags are consumed after checkpoint loading. Keep the
# regular weight layout as well so graph/compile fallback remains
# available.
self.linear_in._interleave_for_swiglu_fusion = True
self.linear_out._accepts_prequantized_fp4 = True
mark_flux2_nvfp4_swiglu_quant_site(self)
def forward(self, x: torch.Tensor) -> torch.Tensor:
if (
flux2_nvfp4_swiglu_quant_active(self)
and getattr(self.linear_in, "_swiglu_fusion_ready", False)
and not torch.compiler.is_compiling()
and not torch.cuda.is_current_stream_capturing()
):
return apply_nvfp4_gemm_swiglu_quant(self.linear_in, self.linear_out, x)
x, _ = self.linear_in(x)
x = self.act_fn(x)
x, _ = self.linear_out(x)
python/sglang/multimodal_gen/runtime/pipelines_core/stages/denoising.py
+7
View File
@@ -21,6 +21,7 @@ import torch
import torch.nn as nn
from sglang.kernels.ops.diffusion import (
mount_flux2_nvfp4_swiglu_quant,
mount_fused_gate_rmsnorm,
mount_fused_linear_gelu,
mount_fused_ln_modulate,
@@ -30,6 +31,7 @@ from sglang.kernels.ops.diffusion import (
mount_nvfp4_bias_gelu,
mount_qwen_image_added_qkv,
mount_sana_video_linear_attention,
unmount_flux2_nvfp4_swiglu_quant,
unmount_fused_gate_rmsnorm,
unmount_fused_linear_gelu,
unmount_fused_ln_modulate,
@@ -165,6 +167,11 @@ logger = init_logger(__name__)
_QUALITY_FUSION_HANDLERS: tuple[
tuple[str, Callable[[nn.Module], bool], Callable[[nn.Module], None]], ...
] = (
(
"FLUX.2 NVFP4 FC1+SwiGLU+quant",
mount_flux2_nvfp4_swiglu_quant,
unmount_flux2_nvfp4_swiglu_quant,
),
(
"fused linear+GELU (cublasLt epilogue)",
mount_fused_linear_gelu,