From 3606b8502373d7c9c43d6e6e5f57ef76ef934bee Mon Sep 17 00:00:00 2001 From: Rahul Chalamala <22563365+rchalamala@users.noreply.github.com> Date: Fri, 18 Sep 2026 19:54:39 +0000 Subject: [PATCH 1/2] quant: reject non-finite fp8 K/V scales at load (sibling of the zero-scale check) * quant: reject non-finite fp8 K/V scales at load (sibling of the zero-scale check) * kv_cache: raise ValueError for non-finite scales instead of assert * kv_cache: validate the final converted scales as well; FNUZ doubling can push a finite float32 scale past the dtype max, so reject non-finite or overflowing values before mutating the layer --------- Co-authored-by: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- .../srt/layers/quantization/kv_cache.py | 21 +++++ .../test_kv_cache_scale_validation.py | 92 +++++++++++++++++++ 2 files changed, 113 insertions(+) create mode 100644 test/registered/unit/layers/quantization/test_kv_cache_scale_validation.py diff --git a/python/sglang/srt/layers/quantization/kv_cache.py b/python/sglang/srt/layers/quantization/kv_cache.py index c8e1f005361f..1adc548d4843 100644 --- a/python/sglang/srt/layers/quantization/kv_cache.py +++ b/python/sglang/srt/layers/quantization/kv_cache.py @@ -3,6 +3,7 @@ # Adapted from https://github.com/vllm-project/vllm/blob/main/vllm/model_executor/layers/quantization/kv_cache.py import logging +import math import torch @@ -49,6 +50,11 @@ def apply(self, layer: torch.nn.Module) -> torch.Tensor: raise RuntimeError(f"{self.__class__.__name__}.apply should not be called.") def process_weights_after_loading(self, layer) -> None: + if not torch.isfinite(layer.k_scale) or not torch.isfinite(layer.v_scale): + raise ValueError( + "fp8 KV cache scales must be finite, got " + f"k_scale={layer.k_scale.item()} v_scale={layer.v_scale.item()}" + ) if layer.k_scale > 0.0 and layer.v_scale > 0.0: # We prefer to use separate k_scale and v_scale if present k_scale = layer.k_scale.to("cpu").tolist() @@ -76,6 +82,21 @@ def process_weights_after_loading(self, layer) -> None: if not isinstance(k_scale, float) or not isinstance(v_scale, float): raise ValueError("Only support per-tensor scaling factor for fp8 KV cache") + # The FNUZ conversion above doubles loaded scales, so a finite + # checkpoint value can still overflow the destination dtype; validate + # the final scales before mutating the layer. + max_scale = torch.finfo(layer.k_scale.dtype).max + if not ( + math.isfinite(k_scale) + and math.isfinite(v_scale) + and k_scale <= max_scale + and v_scale <= max_scale + ): + raise ValueError( + "fp8 KV cache scales must be finite and fit in " + f"{layer.k_scale.dtype}, got k_scale={k_scale} v_scale={v_scale}" + ) + # These are used in the final Attention.forward() layer.k_scale.copy_(k_scale) layer.v_scale.copy_(v_scale) diff --git a/test/registered/unit/layers/quantization/test_kv_cache_scale_validation.py b/test/registered/unit/layers/quantization/test_kv_cache_scale_validation.py new file mode 100644 index 000000000000..be248a16cac0 --- /dev/null +++ b/test/registered/unit/layers/quantization/test_kv_cache_scale_validation.py @@ -0,0 +1,92 @@ +"""Unit tests for fp8 KV cache k/v scale validation at load - no server, no model.""" + +import math +import unittest +from unittest import mock + +import torch + +from sglang.srt.layers.quantization.kv_cache import BaseKVCacheMethod +from sglang.test.ci.ci_register import register_cpu_ci +from sglang.test.test_utils import CustomTestCase + +register_cpu_ci(est_time=5, suite="base-a-test-cpu") + + +def _layer_with_scales(k_scale: float, v_scale: float) -> torch.nn.Module: + layer = torch.nn.Module() + BaseKVCacheMethod(quant_config=mock.MagicMock()).create_weights(layer) + layer.k_scale.copy_(torch.tensor(k_scale, dtype=torch.float32)) + layer.v_scale.copy_(torch.tensor(v_scale, dtype=torch.float32)) + return layer + + +class TestKVCacheScaleValidation(CustomTestCase): + def setUp(self): + self.method = BaseKVCacheMethod(quant_config=mock.MagicMock()) + patcher = mock.patch( + "sglang.srt.layers.quantization.kv_cache.is_fp8_fnuz", + return_value=False, + ) + patcher.start() + self.addCleanup(patcher.stop) + + def test_zero_k_scale_rejected(self): + layer = _layer_with_scales(0.0, 1.0) + with self.assertRaises(AssertionError): + self.method.process_weights_after_loading(layer) + + def test_inf_k_scale_rejected(self): + layer = _layer_with_scales(math.inf, 1.0) + with self.assertRaisesRegex(ValueError, "must be finite"): + self.method.process_weights_after_loading(layer) + + def test_neg_inf_v_scale_rejected(self): + layer = _layer_with_scales(1.0, -math.inf) + with self.assertRaisesRegex(ValueError, "must be finite"): + self.method.process_weights_after_loading(layer) + + def test_nan_k_scale_rejected(self): + layer = _layer_with_scales(math.nan, 1.0) + with self.assertRaisesRegex(ValueError, "must be finite"): + self.method.process_weights_after_loading(layer) + + def test_finite_scales_accepted(self): + layer = _layer_with_scales(1.0, 1.0) + self.method.process_weights_after_loading(layer) + self.assertEqual(layer.k_scale_float, 1.0) + self.assertEqual(layer.v_scale_float, 1.0) + + def test_fnuz_doubling_overflow_rejected(self): + # 3e38 is finite float32 but doubles past the float32 maximum on FNUZ. + with mock.patch( + "sglang.srt.layers.quantization.kv_cache.is_fp8_fnuz", + return_value=True, + ): + layer = _layer_with_scales(3e38, 1.0) + with self.assertRaisesRegex(ValueError, "must be finite"): + self.method.process_weights_after_loading(layer) + + def test_fnuz_doubling_within_float32_accepted(self): + with mock.patch( + "sglang.srt.layers.quantization.kv_cache.is_fp8_fnuz", + return_value=True, + ): + layer = _layer_with_scales(1e38, 1.0) + self.method.process_weights_after_loading(layer) + # 1e38 is not exactly representable in float32, so the doubled + # value is twice the float32-rounded checkpoint value. + self.assertEqual( + layer.k_scale_float, torch.tensor(1e38, dtype=torch.float32).item() * 2 + ) + self.assertEqual(layer.v_scale_float, 2.0) + + def test_unset_scales_default_to_one(self): + layer = _layer_with_scales(-1.0, -1.0) + self.method.process_weights_after_loading(layer) + self.assertEqual(layer.k_scale_float, 1.0) + self.assertEqual(layer.v_scale_float, 1.0) + + +if __name__ == "__main__": + unittest.main() From df9ed1c9e4e414339df22d3b9807cf15756af8fb Mon Sep 17 00:00:00 2001 From: Rahul Chalamala <22563365+rchalamala@users.noreply.github.com> Date: Fri, 18 Sep 2026 19:54:49 +0000 Subject: [PATCH 2/2] quant: reject non-per-tensor fp8 K/V scales explicitly at load Check numel before the finite check so multi-element scales raise a clear ValueError instead of an ambiguous RuntimeError, and evaluate isfinite on CPU copies. Co-authored-by: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- python/sglang/srt/layers/quantization/kv_cache.py | 14 ++++++++++++-- .../test_kv_cache_scale_validation.py | 15 +++++++++++++++ 2 files changed, 27 insertions(+), 2 deletions(-) diff --git a/python/sglang/srt/layers/quantization/kv_cache.py b/python/sglang/srt/layers/quantization/kv_cache.py index 1adc548d4843..5f4c7b73338a 100644 --- a/python/sglang/srt/layers/quantization/kv_cache.py +++ b/python/sglang/srt/layers/quantization/kv_cache.py @@ -50,10 +50,20 @@ def apply(self, layer: torch.nn.Module) -> torch.Tensor: raise RuntimeError(f"{self.__class__.__name__}.apply should not be called.") def process_weights_after_loading(self, layer) -> None: - if not torch.isfinite(layer.k_scale) or not torch.isfinite(layer.v_scale): + if layer.k_scale.numel() != 1 or layer.v_scale.numel() != 1: + raise ValueError( + "Only support per-tensor scaling factor for fp8 KV cache, got " + f"k_scale shape={tuple(layer.k_scale.shape)} " + f"v_scale shape={tuple(layer.v_scale.shape)}" + ) + k_scale_cpu = layer.k_scale.to("cpu") + v_scale_cpu = layer.v_scale.to("cpu") + if not ( + torch.isfinite(k_scale_cpu).all() and torch.isfinite(v_scale_cpu).all() + ): raise ValueError( "fp8 KV cache scales must be finite, got " - f"k_scale={layer.k_scale.item()} v_scale={layer.v_scale.item()}" + f"k_scale={k_scale_cpu.item()} v_scale={v_scale_cpu.item()}" ) if layer.k_scale > 0.0 and layer.v_scale > 0.0: # We prefer to use separate k_scale and v_scale if present diff --git a/test/registered/unit/layers/quantization/test_kv_cache_scale_validation.py b/test/registered/unit/layers/quantization/test_kv_cache_scale_validation.py index be248a16cac0..ebbfa4b45a9a 100644 --- a/test/registered/unit/layers/quantization/test_kv_cache_scale_validation.py +++ b/test/registered/unit/layers/quantization/test_kv_cache_scale_validation.py @@ -51,6 +51,21 @@ def test_nan_k_scale_rejected(self): with self.assertRaisesRegex(ValueError, "must be finite"): self.method.process_weights_after_loading(layer) + def test_multi_element_scale_rejected_with_value_error(self): + layer = _layer_with_scales(1.0, 1.0) + layer.k_scale = torch.nn.Parameter( + torch.tensor([1.0, math.nan]), requires_grad=False + ) + with self.assertRaisesRegex(ValueError, "per-tensor"): + self.method.process_weights_after_loading(layer) + + layer = _layer_with_scales(1.0, 1.0) + layer.k_scale = torch.nn.Parameter( + torch.tensor([1.0, 2.0]), requires_grad=False + ) + with self.assertRaisesRegex(ValueError, "per-tensor"): + self.method.process_weights_after_loading(layer) + def test_finite_scales_accepted(self): layer = _layer_with_scales(1.0, 1.0) self.method.process_weights_after_loading(layer)