minke.yu/sglang
1
0
Fork 0
Code Issues Pull Requests Actions 12 Packages Projects Releases Wiki Activity

Fix invalid KVFP4QuantizeUtil references (#28013)

Browse Source 
Signed-off-by: Zach Zhu <zzqshu@126.com>
This commit is contained in:
Zach Zhu
2026-06-12 13:44:11 -07:00
committed by GitHub
parent 3be5a7ec89
commit 627ed3476b
2 changed files with 31 additions and 19 deletions
Download Patch File Download Diff File Expand all files Collapse all files
python/sglang/srt/mem_cache/memory_pool.py
+26 -14
View File
@@ -1663,9 +1663,11 @@ class MHATokenToKVPoolFP4(MHATokenToKVPool):
) )
cache_k_nope_fp4_sf = self.k_scale_buffer[layer_id - self.start_layer] 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 cache_k_nope_fp4, cache_k_nope_fp4_sf
) )
return cache_k_nope_fp4_dequant 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] 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 cache_v_nope_fp4, cache_v_nope_fp4_sf
) )
return cache_v_nope_fp4_dequant return cache_v_nope_fp4_dequant
@@ -1711,10 +1715,12 @@ class MHATokenToKVPoolFP4(MHATokenToKVPool):
if v_scale is not None: if v_scale is not None:
cache_v.div_(v_scale) 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_k, cache_k_fp4_sf = BlockFP4KVQuantizeUtil.batched_quantize(cache_k)
cache_v, cache_v_fp4_sf = KVFP4QuantizeUtil.batched_quantize(cache_v) cache_v, cache_v_fp4_sf = BlockFP4KVQuantizeUtil.batched_quantize(cache_v)
if self.store_dtype != self.dtype: if self.store_dtype != self.dtype:
cache_k = cache_k.view(self.store_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] 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 cache_k_nope_fp4, cache_k_nope_fp4_sf
) )
return cache_k_nope_fp4_dequant return cache_k_nope_fp4_dequant
@@ -2290,9 +2298,13 @@ class MLATokenToKVPoolFP4(MLATokenToKVPool):
layer_id = layer.layer_id layer_id = layer.layer_id
assert not self.dsa_kv_cache_store_fp8 assert not self.dsa_kv_cache_store_fp8
if cache_k.dtype != self.dtype: 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: if self.store_dtype != self.dtype:
self.kv_buffer[layer_id - self.start_layer][loc] = cache_k_fp4.view( self.kv_buffer[layer_id - self.start_layer][loc] = cache_k_fp4.view(
@@ -2326,14 +2338,14 @@ class MLATokenToKVPoolFP4(MLATokenToKVPool):
else: else:
if cache_k_nope.dtype != self.dtype: if cache_k_nope.dtype != self.dtype:
from sglang.srt.layers.quantization.kvfp4_tensor import ( from sglang.srt.layers.quantization.kvfp4_tensor import (
KVFP4QuantizeUtil, BlockFP4KVQuantizeUtil,
) )
cache_k_nope_fp4, cache_k_nope_fp4_sf = ( 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 = ( 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: if self.store_dtype != self.dtype:
test/manual/quant/test_kvfp4_quant_dequant.py
+5 -5
View File
@@ -6,7 +6,7 @@ import numpy as np
import pytest import pytest
import torch import torch
from sglang.srt.layers.quantization.kvfp4_tensor import KVFP4QuantizeUtil from sglang.srt.layers.quantization.kvfp4_tensor import BlockFP4KVQuantizeUtil
def calculate_accuracy_metrics( 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) fp8_metrics = calculate_accuracy_metrics(tensor_bf16, tensor_fp8_dequant)
# --- KVFP4 --- # --- KVFP4 ---
tensor_fp4, scale_factors = KVFP4QuantizeUtil.batched_quantize(tensor_bf16) tensor_fp4, scale_factors = BlockFP4KVQuantizeUtil.batched_quantize(tensor_bf16)
_ = KVFP4QuantizeUtil.batched_dequantize(tensor_fp4, scale_factors) _ = BlockFP4KVQuantizeUtil.batched_dequantize(tensor_fp4, scale_factors)
start = time.time() start = time.time()
for _ in range(num_runs): 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() torch.cuda.synchronize()
fp4_quant_time = (time.time() - start) / num_runs fp4_quant_time = (time.time() - start) / num_runs
start = time.time() start = time.time()
for _ in range(num_runs): for _ in range(num_runs):
tensor_fp4_dequant = KVFP4QuantizeUtil.batched_dequantize( tensor_fp4_dequant = BlockFP4KVQuantizeUtil.batched_dequantize(
tensor_fp4, scale_factors tensor_fp4, scale_factors
) )
torch.cuda.synchronize() torch.cuda.synchronize()