From 48270612478adfba1274ebc1f7bd084402bfc17a Mon Sep 17 00:00:00 2001 From: Xiaoyu Zhang <1182563586@qq.com> Date: Wed, 12 Aug 2026 16:26:31 +0800 Subject: [PATCH] [Diffusion] Make weight-only FP8 dequant cache torch.compile-safe (#34506) --- .../layers/quantization/weight_only_fp8.py | 5 ++++- .../test_weight_only_fp8_dequant_cache.py | 19 ++++++++++++++++++- 2 files changed, 22 insertions(+), 2 deletions(-) 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 ac2fa0f70..99ed9a948 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 @@ -477,7 +477,10 @@ def _maybe_promote_fp8_weight(module: nn.Module, x_dtype: torch.dtype) -> None: ) _dequant_low_memory_logged = True return - dequant = dequantize_rowwise_fp8_weight(weight, module.weight_scale, dtype) + # Server warmup commonly runs under inference_mode. A cached inference + # tensor has no version counter, so Dynamo cannot later guard it. + with torch.inference_mode(False), torch.no_grad(): + 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( 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 index a4335252b..8417d174b 100644 --- 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 @@ -9,6 +9,7 @@ from sglang.multimodal_gen.runtime.layers.quantization.weight_only_fp8 import ( WeightOnlyFP8Linear, dequantize_rowwise_fp8_weight, ) +from sglang.test.test_utils import CustomTestCase def _make_linear(device: torch.device) -> WeightOnlyFP8Linear: @@ -27,7 +28,7 @@ def _reference(linear: WeightOnlyFP8Linear, x: torch.Tensor) -> torch.Tensor: return torch.nn.functional.linear(x, dequant, linear.bias) -class TestWeightOnlyFP8DequantCache(unittest.TestCase): +class TestWeightOnlyFP8DequantCache(CustomTestCase): def test_cpu_forward_stays_fp8(self): linear = _make_linear(torch.device("cpu")) x = torch.randn(4, 64, dtype=torch.bfloat16) @@ -54,6 +55,22 @@ class TestWeightOnlyFP8DequantCache(unittest.TestCase): linear(x) self.assertEqual(linear.weight.dtype, torch.bfloat16) + @unittest.skipUnless(torch.cuda.is_available(), "requires CUDA") + def test_inference_mode_promotion_supports_torch_compile(self): + linear = _make_linear(torch.device("cuda")) + x = torch.randn(8, 64, device="cuda", dtype=torch.bfloat16) + reference = _reference(linear, x) + + with torch.inference_mode(): + eager_out = linear(x) + self.assertFalse(linear.weight.is_inference()) + + compiled = torch.compile(linear, fullgraph=True) + with torch.inference_mode(): + compiled_out = compiled(x) + self.assertTrue(torch.equal(reference, eager_out)) + self.assertTrue(torch.equal(reference, compiled_out)) + @unittest.skipUnless(torch.cuda.is_available(), "requires CUDA") def test_env_kill_switch(self): linear = _make_linear(torch.device("cuda"))