[diffusion] weight-only FP8: dequantize linear weights once at first use (Ideogram-4 denoise -18.8% H200 / -7.8% H100, bit-exact) (#34305)

Co-authored-by: Claude Fable 5 <noreply@anthropic.com>

python/sglang/multimodal_gen/envs.py

@@ -69,6 +69,7 @@ if TYPE_CHECKING:
     SGLANG_LINGBOT_LAZY_VAE_ENCODE_BLACK_FRAMES: int | None = None
     SGLANG_DIFFUSION_FLASHINFER_FP4_GEMM_BACKEND: str | None = None
     SGLANG_DIFFUSION_ENABLE_W8A8_FP8_GEMM: bool = False
+    SGLANG_DIFFUSION_FP8_WEIGHT_DEQUANT_CACHE: bool = True
     SGLANG_DIFFUSION_VAE_CHANNELS_LAST_3D: str = "auto"
     SGLANG_USE_ROCM_VAE: bool = False
     SGLANG_USE_ROCM_CUDNN_BENCHMARK: bool = False
@@ -287,6 +288,13 @@ environment_variables: dict[str, Callable[[], Any]] = {
     "SGLANG_DIFFUSION_ENABLE_W8A8_FP8_GEMM": _lazy_bool(
         "SGLANG_DIFFUSION_ENABLE_W8A8_FP8_GEMM"
     ),
+    # Dequantize storage-only FP8 linear weights to the compute dtype once,
+    # at first use (bit-identical outputs; trades weight VRAM for skipping
+    # the per-forward dequant pass). Weights are kept FP8-resident when free
+    # memory is low or when this flag is disabled.
+    "SGLANG_DIFFUSION_FP8_WEIGHT_DEQUANT_CACHE": _lazy_bool(
+        "SGLANG_DIFFUSION_FP8_WEIGHT_DEQUANT_CACHE", "true"
+    ),
     # ROCm: use AITer GroupNorm in VAE for improved performance
     "SGLANG_USE_ROCM_VAE": _lazy_bool("SGLANG_USE_ROCM_VAE"),
     # ROCm: enable cudnn.benchmark (MIOpen auto-tuning) for VAE conv layers

python/sglang/multimodal_gen/runtime/layers/quantization/weight_only_fp8.py

@@ -76,6 +76,11 @@ def _apply_weight_only_fp8_linear(
     compute_dtype: torch.dtype,
     enable_fused_w8a8: bool,
 ) -> torch.Tensor:
+    if weight.dtype != FP8_WEIGHT_DTYPE:
+        # Weight was dequantized to the compute dtype once at load time
+        # (see dequantize_weight_only_fp8_linears_at_load).
+        bias = bias.to(weight.dtype) if bias is not None else None
+        return F.linear(x.to(weight.dtype), weight, bias)
     x = x.to(compute_dtype)
     bias = bias.to(compute_dtype) if bias is not None else None
     if enable_fused_w8a8 and _can_apply_fused_w8a8_fp8_linear(
@@ -117,6 +122,7 @@ class WeightOnlyFP8Linear(nn.Module):
         self.out_features = out_features
         self.compute_dtype = compute_dtype
         self.enable_fused_w8a8 = _resolve_enable_fused_w8a8(enable_fused_w8a8)
+        self._fp8_dequant_decided = False
         self.weight = nn.Parameter(
             torch.empty(out_features, in_features, dtype=FP8_WEIGHT_DTYPE),
             requires_grad=False,
@@ -137,6 +143,8 @@ class WeightOnlyFP8Linear(nn.Module):
             self.register_parameter("bias", None)
 
     def forward(self, x: torch.Tensor) -> torch.Tensor:
+        if not self._fp8_dequant_decided:
+            _maybe_promote_fp8_weight(self, x_dtype=x.dtype)
         compute_dtype = self.compute_dtype or x.dtype
         return _apply_weight_only_fp8_linear(
             x,
@@ -171,6 +179,7 @@ class WeightOnlyFP8ColumnParallelLinear(nn.Module):
         self.tp_size = get_group_size(self.tp_group)
         self.tp_rank = get_group_rank(self.tp_group)
         self.out_features_per_partition = divide(out_features, self.tp_size)
+        self._fp8_dequant_decided = False
         self.weight = nn.Parameter(
             torch.empty(
                 self.out_features_per_partition,
@@ -231,6 +240,8 @@ class WeightOnlyFP8ColumnParallelLinear(nn.Module):
         param.data.copy_(loaded_weight)
 
     def forward(self, x: torch.Tensor) -> torch.Tensor:
+        if not self._fp8_dequant_decided:
+            _maybe_promote_fp8_weight(self, x_dtype=x.dtype)
         compute_dtype = self.compute_dtype or x.dtype
         output_parallel = _apply_weight_only_fp8_linear(
             x,
@@ -321,6 +332,7 @@ class WeightOnlyFP8RowParallelLinear(nn.Module):
         self.tp_size = get_group_size(self.tp_group)
         self.tp_rank = get_group_rank(self.tp_group)
         self.in_features_per_partition = divide(in_features, self.tp_size)
+        self._fp8_dequant_decided = False
         self.weight = nn.Parameter(
             torch.empty(
                 out_features,
@@ -380,6 +392,8 @@ class WeightOnlyFP8RowParallelLinear(nn.Module):
                 x, num_partitions=self.tp_size
             )[self.tp_rank].contiguous()
 
+        if not self._fp8_dequant_decided:
+            _maybe_promote_fp8_weight(self, x_dtype=x.dtype)
         compute_dtype = self.compute_dtype or x.dtype
         bias = None if self.tp_rank > 0 else self.bias
         output_parallel = _apply_weight_only_fp8_linear(
@@ -416,6 +430,65 @@ def _log_w8a8_fp8_gemm_warning_once() -> None:
     _w8a8_fp8_gemm_warning_logged = True
 
 
+# Free-memory headroom to preserve when caching dequantized weights; the
+# activation workspace and NCCL buffers still grow after the first forward.
+_DEQUANT_CACHE_RESERVE_BYTES = 4 << 30
+_dequant_cache_logged = False
+_dequant_low_memory_logged = False
+
+
+def _maybe_promote_fp8_weight(module: nn.Module, x_dtype: torch.dtype) -> None:
+    """Dequantize the FP8 weight once, on the first device-resident forward.
+
+    Every later forward then runs a plain compute-dtype GEMM instead of
+    re-materializing the full weight matrix per call; outputs are
+    bit-identical because the same dequantized values feed the same GEMM.
+    The weight stays FP8-resident (per-forward dequant) when the flag is off,
+    under the W8A8 GEMM path, or when free device memory is low; the decision
+    is remembered per module. Runs on host weights leave the decision open so
+    CPU-offloaded modules decide once they land on the device.
+    """
+    global _dequant_cache_logged, _dequant_low_memory_logged
+    weight = module.weight
+    if weight.dtype != FP8_WEIGHT_DTYPE:
+        module._fp8_dequant_decided = True
+        return
+    if weight.device.type != "cuda":
+        return
+    if torch.cuda.is_current_stream_capturing():
+        # Never allocate inside a CUDA graph capture; retry on an eager call.
+        return
+    module._fp8_dequant_decided = True
+    if not envs.SGLANG_DIFFUSION_FP8_WEIGHT_DEQUANT_CACHE:
+        return
+    if module.enable_fused_w8a8:
+        # The W8A8 GEMM path consumes the FP8 weight directly.
+        return
+    dtype = module.compute_dtype or x_dtype
+    if dtype not in (torch.float16, torch.bfloat16, torch.float32):
+        return
+    needed_bytes = weight.numel() * dtype.itemsize
+    free_bytes, _ = torch.cuda.mem_get_info(weight.device)
+    if free_bytes < needed_bytes + _DEQUANT_CACHE_RESERVE_BYTES:
+        if not _dequant_low_memory_logged:
+            logger.warning(
+                "Keeping weight-only FP8 linear weights FP8-resident (low "
+                "free device memory); they dequantize on every forward."
+            )
+            _dequant_low_memory_logged = True
+        return
+    dequant = dequantize_rowwise_fp8_weight(weight, module.weight_scale, dtype)
+    module.weight = nn.Parameter(dequant, requires_grad=False)
+    if not _dequant_cache_logged:
+        logger.info(
+            "Dequantizing weight-only FP8 linear weights once at first use "
+            "(bit-identical outputs; set "
+            "SGLANG_DIFFUSION_FP8_WEIGHT_DEQUANT_CACHE=0 to keep weights "
+            "FP8-resident)."
+        )
+        _dequant_cache_logged = True
+
+
 def swap_linears_to_weight_only_fp8(module: nn.Module) -> None:
     """Recursively replace nn.Linear with WeightOnlyFP8Linear.
 

python/sglang/multimodal_gen/test/unit/test_weight_only_fp8_dequant_cache.py

@@ -0,0 +1,77 @@
+import os
+import unittest
+from unittest.mock import patch
+
+import torch
+
+from sglang.multimodal_gen.runtime.layers.quantization.weight_only_fp8 import (
+    FP8_WEIGHT_DTYPE,
+    WeightOnlyFP8Linear,
+    dequantize_rowwise_fp8_weight,
+)
+
+
+def _make_linear(device: torch.device) -> WeightOnlyFP8Linear:
+    torch.manual_seed(0)
+    linear = WeightOnlyFP8Linear(64, 32, bias=True, compute_dtype=torch.bfloat16)
+    linear.weight.data = (torch.randn(32, 64, device=device) * 0.1).to(FP8_WEIGHT_DTYPE)
+    linear.weight_scale.data = torch.rand(32, device=device, dtype=torch.float32) + 0.5
+    linear.bias.data = torch.randn(32, device=device, dtype=torch.bfloat16)
+    return linear.to(device)
+
+
+def _reference(linear: WeightOnlyFP8Linear, x: torch.Tensor) -> torch.Tensor:
+    dequant = dequantize_rowwise_fp8_weight(
+        linear.weight, linear.weight_scale, torch.bfloat16
+    )
+    return torch.nn.functional.linear(x, dequant, linear.bias)
+
+
+class TestWeightOnlyFP8DequantCache(unittest.TestCase):
+    def test_cpu_forward_stays_fp8(self):
+        linear = _make_linear(torch.device("cpu"))
+        x = torch.randn(4, 64, dtype=torch.bfloat16)
+        reference = _reference(linear, x)
+        self.assertTrue(torch.equal(reference, linear(x)))
+        self.assertEqual(linear.weight.dtype, FP8_WEIGHT_DTYPE)
+
+    @unittest.skipUnless(torch.cuda.is_available(), "requires CUDA")
+    def test_first_forward_promotes_bit_identically(self):
+        linear = _make_linear(torch.device("cuda"))
+        x = torch.randn(8, 64, device="cuda", dtype=torch.bfloat16)
+        reference = _reference(linear, x)
+
+        out = linear(x)
+        self.assertEqual(linear.weight.dtype, torch.bfloat16)
+        self.assertTrue(torch.equal(reference, out))
+        self.assertTrue(torch.equal(reference, linear(x)))
+
+    @unittest.skipUnless(torch.cuda.is_available(), "requires CUDA")
+    def test_promotion_dtype_follows_input_when_unset(self):
+        linear = _make_linear(torch.device("cuda"))
+        linear.compute_dtype = None
+        x = torch.randn(8, 64, device="cuda", dtype=torch.bfloat16)
+        linear(x)
+        self.assertEqual(linear.weight.dtype, torch.bfloat16)
+
+    @unittest.skipUnless(torch.cuda.is_available(), "requires CUDA")
+    def test_env_kill_switch(self):
+        linear = _make_linear(torch.device("cuda"))
+        x = torch.randn(8, 64, device="cuda", dtype=torch.bfloat16)
+        reference = _reference(linear, x)
+        with patch.dict(os.environ, {"SGLANG_DIFFUSION_FP8_WEIGHT_DEQUANT_CACHE": "0"}):
+            out = linear(x)
+        self.assertEqual(linear.weight.dtype, FP8_WEIGHT_DTYPE)
+        self.assertTrue(torch.equal(reference, out))
+
+    @unittest.skipUnless(torch.cuda.is_available(), "requires CUDA")
+    def test_low_memory_keeps_fp8(self):
+        linear = _make_linear(torch.device("cuda"))
+        x = torch.randn(8, 64, device="cuda", dtype=torch.bfloat16)
+        with patch("torch.cuda.mem_get_info", return_value=(0, 0)):
+            linear(x)
+        self.assertEqual(linear.weight.dtype, FP8_WEIGHT_DTYPE)
+
+
+if __name__ == "__main__":
+    unittest.main()