From f4c17fed0767626cd81ada078ed2fda5d3187f25 Mon Sep 17 00:00:00 2001 From: Xiaoyu Zhang <1182563586@qq.com> Date: Wed, 2 Sep 2026 08:21:14 +0800 Subject: [PATCH] [Diffusion] Fuse FLUX.2 NVFP4 FC1, SwiGLU, and FC2 quantization (#37096) Co-authored-by: Cursor --- .../sglang/kernels/ops/diffusion/__init__.py | 4 + .../sites/flux2_nvfp4_swiglu_quant_site.py | 36 +++++ .../layers/quantization/modelopt_quant.py | 141 ++++++++++++++++-- .../runtime/models/dits/flux_2.py | 34 +++++ .../pipelines_core/stages/denoising.py | 7 + .../ops/diffusion/test_model_fast_paths.py | 7 + .../kernels/ops/diffusion/test_sites.py | 16 ++ .../test_diffusion_nvfp4_scaled_mm.py | 101 ++++++++++++- 8 files changed, 333 insertions(+), 13 deletions(-) create mode 100644 python/sglang/kernels/ops/diffusion/sites/flux2_nvfp4_swiglu_quant_site.py diff --git a/python/sglang/kernels/ops/diffusion/__init__.py b/python/sglang/kernels/ops/diffusion/__init__.py index 4cd0e1b57..5e2a31262 100644 --- a/python/sglang/kernels/ops/diffusion/__init__.py +++ b/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", diff --git a/python/sglang/kernels/ops/diffusion/sites/flux2_nvfp4_swiglu_quant_site.py b/python/sglang/kernels/ops/diffusion/sites/flux2_nvfp4_swiglu_quant_site.py new file mode 100644 index 000000000..ae73cca3a --- /dev/null +++ b/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) diff --git a/python/sglang/multimodal_gen/runtime/layers/quantization/modelopt_quant.py b/python/sglang/multimodal_gen/runtime/layers/quantization/modelopt_quant.py index 8b31410a6..eebcc1c7a 100755 --- a/python/sglang/multimodal_gen/runtime/layers/quantization/modelopt_quant.py +++ b/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]) diff --git a/python/sglang/multimodal_gen/runtime/models/dits/flux_2.py b/python/sglang/multimodal_gen/runtime/models/dits/flux_2.py index 0a3c1b46d..ecfdaed92 100644 --- a/python/sglang/multimodal_gen/runtime/models/dits/flux_2.py +++ b/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) 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 fbda98243..1d79a4a47 100644 --- a/python/sglang/multimodal_gen/runtime/pipelines_core/stages/denoising.py +++ b/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, diff --git a/test/registered/kernels/ops/diffusion/test_model_fast_paths.py b/test/registered/kernels/ops/diffusion/test_model_fast_paths.py index 0335b15c3..85d6b984e 100644 --- a/test/registered/kernels/ops/diffusion/test_model_fast_paths.py +++ b/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:] diff --git a/test/registered/kernels/ops/diffusion/test_sites.py b/test/registered/kernels/ops/diffusion/test_sites.py index b00e90d3d..7146125dc 100644 --- a/test/registered/kernels/ops/diffusion/test_sites.py +++ b/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() diff --git a/test/registered/kernels/ops/quantization/test_diffusion_nvfp4_scaled_mm.py b/test/registered/kernels/ops/quantization/test_diffusion_nvfp4_scaled_mm.py index d99be8193..73a759746 100644 --- a/test/registered/kernels/ops/quantization/test_diffusion_nvfp4_scaled_mm.py +++ b/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"