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

Co-authored-by: Cursor <cursoragent@cursor.com>

python/sglang/kernels/ops/diffusion/__init__.py

@@ -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

@@ -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

@@ -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

@@ -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

@@ -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,

test/registered/kernels/ops/diffusion/test_model_fast_paths.py

@@ -82,6 +82,7 @@ from sglang.multimodal_gen.runtime.models.dits.flux import (
     _flux_norm_modulate,
 )
 from sglang.multimodal_gen.runtime.models.dits.flux_2 import (
+    _can_use_nvfp4_swiglu_quant_fusion,
     _flux2_norm_modulate,
     _flux2_swiglu,
 )
@@ -121,6 +122,7 @@ from sglang.multimodal_gen.runtime.models.vaes.wan_vae_cuda_opt import (
     VaeFastPathGate,
 )
 from sglang.multimodal_gen.runtime.models.vaes.wanvae import WanRMS_norm
+from sglang.multimodal_gen.runtime.platforms.interface import DeviceCapability
 from sglang.test.ci.ci_register import register_amd_ci, register_cuda_ci
 from sglang.test.test_utils import CustomTestCase
 
@@ -342,6 +344,11 @@ class TestFlux2EagerFusions(CustomTestCase):
         self.assertFalse(flux2._FLUX2_SWIGLU.disabled)
         self.assertEqual(len(flux2._FLUX2_SWIGLU_SIGS), 2)
 
+    def test_nvfp4_swiglu_quant_fusion_is_sm103_only(self):
+        self.assertFalse(_can_use_nvfp4_swiglu_quant_fusion(DeviceCapability(10, 0)))
+        self.assertTrue(_can_use_nvfp4_swiglu_quant_fusion(DeviceCapability(10, 3)))
+        self.assertFalse(_can_use_nvfp4_swiglu_quant_fusion(DeviceCapability(12, 0)))
+
     def test_fp16_preserves_reference_path(self):
         x = torch.randn(1, 17, 512, device="cuda", dtype=torch.float16)
         expected = F.silu(x[..., :256]) * x[..., 256:]

test/registered/kernels/ops/diffusion/test_sites.py

@@ -35,11 +35,15 @@ from sglang.kernels.ops.diffusion import (
     QualityGatedFusion,
     can_use_ln_modulate,
     flashinfer_rmsnorm_diagnostic_hint,
+    flux2_nvfp4_swiglu_quant_active,
     fused_ln_modulate,
     fused_ln_modulate_active,
+    mark_flux2_nvfp4_swiglu_quant_site,
     mark_fused_ln_modulate_site,
+    mount_flux2_nvfp4_swiglu_quant,
     mount_fused_ln_modulate,
     tensors_equal,
+    unmount_flux2_nvfp4_swiglu_quant,
     unmount_fused_ln_modulate,
 )
 from sglang.test.ci.ci_register import register_cpu_ci, register_cuda_ci
@@ -92,6 +96,18 @@ def test_quality_gate_rejection_is_all_or_nothing():
     assert not fusion.mount(nn.Module())
 
 
+def test_flux2_nvfp4_swiglu_quant_is_disabled_until_quality_gate_mounts():
+    site = nn.Module()
+    root = nn.ModuleList([site])
+    mark_flux2_nvfp4_swiglu_quant_site(site)
+
+    assert not flux2_nvfp4_swiglu_quant_active(site)
+    assert mount_flux2_nvfp4_swiglu_quant(root)
+    assert flux2_nvfp4_swiglu_quant_active(site)
+    unmount_flux2_nvfp4_swiglu_quant(root)
+    assert not flux2_nvfp4_swiglu_quant_active(site)
+
+
 def test_qwen_image_added_qkv_site_is_request_scoped():
     site = nn.Module()
     site.to_added_qkv = nn.Module()

test/registered/kernels/ops/quantization/test_diffusion_nvfp4_scaled_mm.py

@@ -3,6 +3,7 @@ import sys
 import flashinfer
 import pytest
 import torch
+import torch.nn.functional as F
 
 from sglang.kernels.ops.diffusion import (
     fused_scale_residual_norm_scale_shift,
@@ -15,6 +16,7 @@ from sglang.multimodal_gen.runtime.layers.quantization.modelopt_quant import (
     ModelOptFp4Config,
     ModelOptFp4LinearMethod,
     apply_nvfp4_gemm_prequantized,
+    apply_nvfp4_gemm_swiglu_quant,
 )
 from sglang.multimodal_gen.runtime.platforms import current_platform
 from sglang.srt.layers.quantization.modelopt_quant import pad_nvfp4_weight
@@ -37,6 +39,11 @@ TEST_CASES = [
 FLUX2_PROJECTION_SHAPE = (512, 6144, 128)
 
 
+class _TestLinear(torch.nn.Module):
+    def forward(self, x):
+        return self.quant_method.apply(self, x), None
+
+
 def _nvfp4_supported() -> bool:
     return torch.cuda.is_available() and torch.cuda.get_device_capability() >= (10, 0)
 
@@ -197,6 +204,8 @@ def _build_layer(
     *,
     weight_scale_device: torch.device | str | None = None,
     checkpoint_weight_scale_layout: str = "linear",
+    prepare_swiglu_fusion: bool = False,
+    accepts_prequantized_fp4: bool = False,
 ) -> tuple[ModelOptFp4LinearMethod, torch.nn.Module]:
     output_size, input_size_half = weight_fp4.shape
     input_size = input_size_half * 2
@@ -208,7 +217,9 @@ def _build_layer(
             checkpoint_weight_scale_layout=checkpoint_weight_scale_layout,
         )
     )
-    layer = torch.nn.Module()
+    layer = _TestLinear()
+    layer.quant_method = method
+    layer.params_dtype = DTYPE
     method.create_weights(
         layer,
         input_size_per_partition=input_size,
@@ -234,6 +245,9 @@ def _build_layer(
             layer.weight_scale.detach().to(weight_scale_device), requires_grad=False
         )
 
+    layer._interleave_for_swiglu_fusion = prepare_swiglu_fusion
+    layer._accepts_prequantized_fp4 = accepts_prequantized_fp4
+
     method.process_weights_after_loading(layer)
 
     _, flashinfer_backend = current_platform.get_modelopt_fp4_gemm_op()
@@ -300,6 +314,91 @@ def _build_layer(
     return method, layer
 
 
+@pytest.mark.skipif(
+    not _nvfp4_supported(),
+    reason="Diffusion NVFP4 fused SwiGLU correctness requires Blackwell GPUs",
+)
+def test_flux2_fused_nvfp4_swiglu_quant_matches_unfused() -> None:
+    batch, seq_len, hidden_size, inner_size, output_size = 2, 32, 128, 128, 128
+    generator = torch.Generator(device=DEVICE)
+    generator.manual_seed(20260830)
+
+    x = torch.randn(
+        (batch, seq_len, hidden_size),
+        device=DEVICE,
+        dtype=DTYPE,
+        generator=generator,
+    )
+    weight_in = torch.randn(
+        (2 * inner_size, hidden_size),
+        device=DEVICE,
+        dtype=DTYPE,
+        generator=generator,
+    )
+    weight_out = torch.randn(
+        (output_size, inner_size),
+        device=DEVICE,
+        dtype=DTYPE,
+        generator=generator,
+    )
+
+    input_global_scale = _make_global_scale(x)
+    weight_in_global_scale = _make_global_scale(weight_in)
+    weight_out_global_scale = _make_global_scale(weight_out)
+    output_input_global_scale = torch.tensor(512.0, device=DEVICE, dtype=torch.float32)
+
+    weight_in_fp4, weight_in_scale = _quantize_weight_for_checkpoint(
+        weight_in, weight_in_global_scale
+    )
+    weight_out_fp4, weight_out_scale = _quantize_weight_for_checkpoint(
+        weight_out, weight_out_global_scale
+    )
+    method_in, layer_in = _build_layer(
+        weight_in_fp4,
+        weight_in_scale,
+        input_global_scale,
+        weight_in_global_scale,
+        prepare_swiglu_fusion=True,
+    )
+    method_out, layer_out = _build_layer(
+        weight_out_fp4,
+        weight_out_scale,
+        output_input_global_scale,
+        weight_out_global_scale,
+        accepts_prequantized_fp4=True,
+    )
+
+    projected = method_in.apply(layer_in, x)
+    expected = method_out.apply(
+        layer_out,
+        F.silu(projected[..., :inner_size]) * projected[..., inner_size:],
+    )
+    actual = apply_nvfp4_gemm_swiglu_quant(layer_in, layer_out, x)
+
+    assert actual.shape == expected.shape == (batch, seq_len, output_size)
+    assert torch.isfinite(actual).all()
+    assert "weight_swiglu_interleaved" in dict(layer_in.named_buffers())
+    assert "weight_scale_swiglu_interleaved" in dict(layer_in.named_buffers())
+    assert "weight_swiglu_interleaved" not in layer_in.state_dict()
+    assert "weight_scale_swiglu_interleaved" not in layer_in.state_dict()
+    diff = _calc_diff(actual, expected)
+    assert diff < DEEPGEMM_FP4_MAX_DIFF, f"{diff=:.6f}"
+
+    weight_ptr = layer_in.weight_swiglu_interleaved.data_ptr()
+    scale_ptr = layer_in.weight_scale_swiglu_interleaved.data_ptr()
+    previous_interleaved_weight = layer_in.weight_swiglu_interleaved.clone()
+    reloaded_weight = _swap_fp4_nibbles(weight_in_fp4).clone()
+    reloaded_weight.view(torch.uint8).flatten()[0] ^= 0x11
+    layer_in.weight.data.copy_(reloaded_weight)
+    layer_in.weight_scale.data.copy_(weight_in_scale)
+    method_in.process_weights_after_loading(layer_in)
+    assert layer_in.weight_swiglu_interleaved.data_ptr() == weight_ptr
+    assert layer_in.weight_scale_swiglu_interleaved.data_ptr() == scale_ptr
+    assert not torch.equal(
+        layer_in.weight_swiglu_interleaved, previous_interleaved_weight
+    )
+
+
 def _resolve_mode(mode: str):
     if mode == "flashinfer2":
         return flashinfer.fp4_quantize, flashinfer.mm_fp4, "cudnn"