diff --git a/python/sglang/multimodal_gen/envs.py b/python/sglang/multimodal_gen/envs.py index 3943f0164..3ff7b41af 100644 --- a/python/sglang/multimodal_gen/envs.py +++ b/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 diff --git a/python/sglang/multimodal_gen/runtime/layers/quantization/weight_only_fp8.py b/python/sglang/multimodal_gen/runtime/layers/quantization/weight_only_fp8.py index 0e19ee07b..ac2fa0f70 100644 --- a/python/sglang/multimodal_gen/runtime/layers/quantization/weight_only_fp8.py +++ b/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. diff --git a/python/sglang/multimodal_gen/test/unit/test_weight_only_fp8_dequant_cache.py b/python/sglang/multimodal_gen/test/unit/test_weight_only_fp8_dequant_cache.py new file mode 100644 index 000000000..a4335252b --- /dev/null +++ b/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()