From 6d5e16fb1c3fed24bed04c2ca08091f9a79e09df Mon Sep 17 00:00:00 2001 From: "Ho-Ren (Jack) Chuang" Date: Fri, 14 Nov 2025 22:31:35 -0800 Subject: [PATCH] feat: Add FP4 (E2M1) KV Cache Support for MHA (#12612) Signed-off-by: Ho-Ren (Jack) Chuang Co-authored-by: Yichen Wang --- python/sglang/srt/mem_cache/memory_pool.py | 135 +++++++++++++++--- .../sglang/srt/model_executor/model_runner.py | 18 +++ 2 files changed, 131 insertions(+), 22 deletions(-) diff --git a/python/sglang/srt/mem_cache/memory_pool.py b/python/sglang/srt/mem_cache/memory_pool.py index c89f5f6d0..8cd3865b2 100644 --- a/python/sglang/srt/mem_cache/memory_pool.py +++ b/python/sglang/srt/mem_cache/memory_pool.py @@ -634,24 +634,65 @@ class MHATokenToKVPool(KVCache): if self.enable_custom_mem_pool else nullcontext() ): - # [size, head_num, head_dim] for each layer - # The padded slot 0 is used for writing dummy outputs from padded tokens. - self.k_buffer = [ - torch.zeros( - (self.size + self.page_size, self.head_num, self.head_dim), - dtype=self.store_dtype, - device=self.device, - ) - for _ in range(self.layer_num) - ] - self.v_buffer = [ - torch.zeros( - (self.size + self.page_size, self.head_num, self.head_dim), - dtype=self.store_dtype, - device=self.device, - ) - for _ in range(self.layer_num) - ] + if is_float4_e2m1fn_x2(self.dtype): + m = self.size + self.page_size + n = self.head_num + k = self.head_dim + + scale_block_size = 16 + self.store_dtype = torch.uint8 + self.k_buffer = [ + torch.zeros( + (m, n, k // 2), + dtype=self.store_dtype, + device=self.device, + ) + for _ in range(self.layer_num) + ] + self.v_buffer = [ + torch.zeros( + (m, n, k // 2), + dtype=self.store_dtype, + device=self.device, + ) + for _ in range(self.layer_num) + ] + + self.k_scale_buffer = [ + torch.zeros( + (m, (n * k) // scale_block_size), + dtype=self.store_dtype, + device=self.device, + ) + for _ in range(self.layer_num) + ] + self.v_scale_buffer = [ + torch.zeros( + (m, (n * k) // scale_block_size), + dtype=self.store_dtype, + device=self.device, + ) + for _ in range(self.layer_num) + ] + else: + # [size, head_num, head_dim] for each layer + # The padded slot 0 is used for writing dummy outputs from padded tokens. + self.k_buffer = [ + torch.zeros( + (self.size + self.page_size, self.head_num, self.head_dim), + dtype=self.store_dtype, + device=self.device, + ) + for _ in range(self.layer_num) + ] + self.v_buffer = [ + torch.zeros( + (self.size + self.page_size, self.head_num, self.head_dim), + dtype=self.store_dtype, + device=self.device, + ) + for _ in range(self.layer_num) + ] self.k_data_ptrs = torch.tensor( [x.data_ptr() for x in self.k_buffer], @@ -752,7 +793,22 @@ class MHATokenToKVPool(KVCache): def _get_key_buffer(self, layer_id: int): # for internal use of referencing if self.store_dtype != self.dtype: - return self.k_buffer[layer_id - self.start_layer].view(self.dtype) + if is_float4_e2m1fn_x2(self.dtype): + cache_k_nope_fp4 = self.k_buffer[layer_id - self.start_layer].view( + torch.uint8 + ) + cache_k_nope_fp4_sf = self.k_scale_buffer[layer_id - self.start_layer] + + from sglang.srt.layers.quantization.kvfp4_tensor import ( + KVFP4QuantizeUtil, + ) + + cache_k_nope_fp4_dequant = KVFP4QuantizeUtil.batched_dequantize( + cache_k_nope_fp4, cache_k_nope_fp4_sf + ) + return cache_k_nope_fp4_dequant + else: + return self.k_buffer[layer_id - self.start_layer].view(self.dtype) return self.k_buffer[layer_id - self.start_layer] def get_key_buffer(self, layer_id: int): @@ -766,7 +822,22 @@ class MHATokenToKVPool(KVCache): def _get_value_buffer(self, layer_id: int): # for internal use of referencing if self.store_dtype != self.dtype: - return self.v_buffer[layer_id - self.start_layer].view(self.dtype) + if is_float4_e2m1fn_x2(self.dtype): + cache_v_nope_fp4 = self.v_buffer[layer_id - self.start_layer].view( + torch.uint8 + ) + cache_v_nope_fp4_sf = self.v_scale_buffer[layer_id - self.start_layer] + + from sglang.srt.layers.quantization.kvfp4_tensor import ( + KVFP4QuantizeUtil, + ) + + cache_v_nope_fp4_dequant = KVFP4QuantizeUtil.batched_dequantize( + cache_v_nope_fp4, cache_v_nope_fp4_sf + ) + return cache_v_nope_fp4_dequant + else: + return self.v_buffer[layer_id - self.start_layer].view(self.dtype) return self.v_buffer[layer_id - self.start_layer] def get_value_buffer(self, layer_id: int): @@ -798,24 +869,44 @@ class MHATokenToKVPool(KVCache): cache_k.div_(k_scale) if v_scale is not None: cache_v.div_(v_scale) - cache_k = cache_k.to(self.dtype) - cache_v = cache_v.to(self.dtype) + if is_float4_e2m1fn_x2(self.dtype): + from sglang.srt.layers.quantization.kvfp4_tensor import ( + KVFP4QuantizeUtil, + ) + + cache_k, cache_k_fp4_sf = KVFP4QuantizeUtil.batched_quantize(cache_k) + cache_v, cache_v_fp4_sf = KVFP4QuantizeUtil.batched_quantize(cache_v) + else: + cache_k = cache_k.to(self.dtype) + cache_v = cache_v.to(self.dtype) if self.store_dtype != self.dtype: cache_k = cache_k.view(self.store_dtype) cache_v = cache_v.view(self.store_dtype) + if is_float4_e2m1fn_x2(self.dtype): + cache_k_fp4_sf = cache_k_fp4_sf.view(self.store_dtype) + cache_v_fp4_sf = cache_v_fp4_sf.view(self.store_dtype) if get_is_capture_mode() and self.alt_stream is not None: # Overlap the copy of K and V cache for small batch size current_stream = self.device_module.current_stream() self.alt_stream.wait_stream(current_stream) self.k_buffer[layer_id - self.start_layer][loc] = cache_k + if is_float4_e2m1fn_x2(self.dtype): + self.k_scale_buffer[layer_id - self.start_layer][loc] = cache_k_fp4_sf with self.device_module.stream(self.alt_stream): self.v_buffer[layer_id - self.start_layer][loc] = cache_v + if is_float4_e2m1fn_x2(self.dtype): + self.v_scale_buffer[layer_id - self.start_layer][ + loc + ] = cache_v_fp4_sf current_stream.wait_stream(self.alt_stream) else: self.k_buffer[layer_id - self.start_layer][loc] = cache_k self.v_buffer[layer_id - self.start_layer][loc] = cache_v + if is_float4_e2m1fn_x2(self.dtype): + self.k_scale_buffer[layer_id - self.start_layer][loc] = cache_k_fp4_sf + self.v_scale_buffer[layer_id - self.start_layer][loc] = cache_v_fp4_sf def move_kv_cache(self, tgt_loc: torch.Tensor, src_loc: torch.Tensor): N = tgt_loc.numel() diff --git a/python/sglang/srt/model_executor/model_runner.py b/python/sglang/srt/model_executor/model_runner.py index 19e029d60..3ba5d352b 100644 --- a/python/sglang/srt/model_executor/model_runner.py +++ b/python/sglang/srt/model_executor/model_runner.py @@ -1325,6 +1325,24 @@ class ModelRunner: * 2 * torch._utils._element_size(self.kv_cache_dtype) ) + + if is_float4_e2m1fn_x2(self.kv_cache_dtype): + # kv_scale_buffer + scale_block_size = 16 + + n = self.model_config.get_num_kv_heads(get_attention_tp_size()) + k = self.model_config.head_dim + cell_size = (cell_size // 2) + ( + ( + n + * k + * num_layers + * 2 + * torch._utils._element_size(self.kv_cache_dtype) + ) + // scale_block_size + ) + rest_memory = available_gpu_memory - total_gpu_memory * ( 1 - self.mem_fraction_static )