From 627ed3476b0567609ca81a15d584eb4abb036e97 Mon Sep 17 00:00:00 2001 From: Zach Zhu Date: Sat, 13 Jun 2026 04:44:11 +0800 Subject: [PATCH] Fix invalid KVFP4QuantizeUtil references (#28013) Signed-off-by: Zach Zhu --- python/sglang/srt/mem_cache/memory_pool.py | 40 ++++++++++++------- test/manual/quant/test_kvfp4_quant_dequant.py | 10 ++--- 2 files changed, 31 insertions(+), 19 deletions(-) diff --git a/python/sglang/srt/mem_cache/memory_pool.py b/python/sglang/srt/mem_cache/memory_pool.py index 85bd9a798..31db39e25 100644 --- a/python/sglang/srt/mem_cache/memory_pool.py +++ b/python/sglang/srt/mem_cache/memory_pool.py @@ -1663,9 +1663,11 @@ class MHATokenToKVPoolFP4(MHATokenToKVPool): ) cache_k_nope_fp4_sf = self.k_scale_buffer[layer_id - self.start_layer] - from sglang.srt.layers.quantization.kvfp4_tensor import KVFP4QuantizeUtil + from sglang.srt.layers.quantization.kvfp4_tensor import ( + BlockFP4KVQuantizeUtil, + ) - cache_k_nope_fp4_dequant = KVFP4QuantizeUtil.batched_dequantize( + cache_k_nope_fp4_dequant = BlockFP4KVQuantizeUtil.batched_dequantize( cache_k_nope_fp4, cache_k_nope_fp4_sf ) return cache_k_nope_fp4_dequant @@ -1679,9 +1681,11 @@ class MHATokenToKVPoolFP4(MHATokenToKVPool): ) cache_v_nope_fp4_sf = self.v_scale_buffer[layer_id - self.start_layer] - from sglang.srt.layers.quantization.kvfp4_tensor import KVFP4QuantizeUtil + from sglang.srt.layers.quantization.kvfp4_tensor import ( + BlockFP4KVQuantizeUtil, + ) - cache_v_nope_fp4_dequant = KVFP4QuantizeUtil.batched_dequantize( + cache_v_nope_fp4_dequant = BlockFP4KVQuantizeUtil.batched_dequantize( cache_v_nope_fp4, cache_v_nope_fp4_sf ) return cache_v_nope_fp4_dequant @@ -1711,10 +1715,12 @@ class MHATokenToKVPoolFP4(MHATokenToKVPool): if v_scale is not None: cache_v.div_(v_scale) - from sglang.srt.layers.quantization.kvfp4_tensor import KVFP4QuantizeUtil + from sglang.srt.layers.quantization.kvfp4_tensor import ( + BlockFP4KVQuantizeUtil, + ) - cache_k, cache_k_fp4_sf = KVFP4QuantizeUtil.batched_quantize(cache_k) - cache_v, cache_v_fp4_sf = KVFP4QuantizeUtil.batched_quantize(cache_v) + cache_k, cache_k_fp4_sf = BlockFP4KVQuantizeUtil.batched_quantize(cache_k) + cache_v, cache_v_fp4_sf = BlockFP4KVQuantizeUtil.batched_quantize(cache_v) if self.store_dtype != self.dtype: cache_k = cache_k.view(self.store_dtype) @@ -2268,9 +2274,11 @@ class MLATokenToKVPoolFP4(MLATokenToKVPool): ) cache_k_nope_fp4_sf = self.kv_scale_buffer[layer_id - self.start_layer] - from sglang.srt.layers.quantization.kvfp4_tensor import KVFP4QuantizeUtil + from sglang.srt.layers.quantization.kvfp4_tensor import ( + BlockFP4KVQuantizeUtil, + ) - cache_k_nope_fp4_dequant = KVFP4QuantizeUtil.batched_dequantize( + cache_k_nope_fp4_dequant = BlockFP4KVQuantizeUtil.batched_dequantize( cache_k_nope_fp4, cache_k_nope_fp4_sf ) return cache_k_nope_fp4_dequant @@ -2290,9 +2298,13 @@ class MLATokenToKVPoolFP4(MLATokenToKVPool): layer_id = layer.layer_id assert not self.dsa_kv_cache_store_fp8 if cache_k.dtype != self.dtype: - from sglang.srt.layers.quantization.kvfp4_tensor import KVFP4QuantizeUtil + from sglang.srt.layers.quantization.kvfp4_tensor import ( + BlockFP4KVQuantizeUtil, + ) - cache_k_fp4, cache_k_fp4_sf = KVFP4QuantizeUtil.batched_quantize(cache_k) + cache_k_fp4, cache_k_fp4_sf = BlockFP4KVQuantizeUtil.batched_quantize( + cache_k + ) if self.store_dtype != self.dtype: self.kv_buffer[layer_id - self.start_layer][loc] = cache_k_fp4.view( @@ -2326,14 +2338,14 @@ class MLATokenToKVPoolFP4(MLATokenToKVPool): else: if cache_k_nope.dtype != self.dtype: from sglang.srt.layers.quantization.kvfp4_tensor import ( - KVFP4QuantizeUtil, + BlockFP4KVQuantizeUtil, ) cache_k_nope_fp4, cache_k_nope_fp4_sf = ( - KVFP4QuantizeUtil.batched_quantize(cache_k_nope) + BlockFP4KVQuantizeUtil.batched_quantize(cache_k_nope) ) cache_k_rope_fp4, cache_k_rope_fp4_sf = ( - KVFP4QuantizeUtil.batched_quantize(cache_k_rope) + BlockFP4KVQuantizeUtil.batched_quantize(cache_k_rope) ) if self.store_dtype != self.dtype: diff --git a/test/manual/quant/test_kvfp4_quant_dequant.py b/test/manual/quant/test_kvfp4_quant_dequant.py index 122694a41..3f7d5b6a8 100755 --- a/test/manual/quant/test_kvfp4_quant_dequant.py +++ b/test/manual/quant/test_kvfp4_quant_dequant.py @@ -6,7 +6,7 @@ import numpy as np import pytest import torch -from sglang.srt.layers.quantization.kvfp4_tensor import KVFP4QuantizeUtil +from sglang.srt.layers.quantization.kvfp4_tensor import BlockFP4KVQuantizeUtil def calculate_accuracy_metrics( @@ -52,18 +52,18 @@ def run_benchmark(m, n, k, num_runs=100) -> dict[str, dict[str, float]]: fp8_metrics = calculate_accuracy_metrics(tensor_bf16, tensor_fp8_dequant) # --- KVFP4 --- - tensor_fp4, scale_factors = KVFP4QuantizeUtil.batched_quantize(tensor_bf16) - _ = KVFP4QuantizeUtil.batched_dequantize(tensor_fp4, scale_factors) + tensor_fp4, scale_factors = BlockFP4KVQuantizeUtil.batched_quantize(tensor_bf16) + _ = BlockFP4KVQuantizeUtil.batched_dequantize(tensor_fp4, scale_factors) start = time.time() for _ in range(num_runs): - tensor_fp4, scale_factors = KVFP4QuantizeUtil.batched_quantize(tensor_bf16) + tensor_fp4, scale_factors = BlockFP4KVQuantizeUtil.batched_quantize(tensor_bf16) torch.cuda.synchronize() fp4_quant_time = (time.time() - start) / num_runs start = time.time() for _ in range(num_runs): - tensor_fp4_dequant = KVFP4QuantizeUtil.batched_dequantize( + tensor_fp4_dequant = BlockFP4KVQuantizeUtil.batched_dequantize( tensor_fp4, scale_factors ) torch.cuda.synchronize()