diff --git a/python/sglang/srt/layers/quantization/fp4_kv_cache_quant_method.py b/python/sglang/srt/layers/quantization/fp4_kv_cache_quant_method.py new file mode 100644 index 000000000..b6b1d6b8a --- /dev/null +++ b/python/sglang/srt/layers/quantization/fp4_kv_cache_quant_method.py @@ -0,0 +1,426 @@ +# Copyright 2025 SGLang Team +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +# ============================================================================== +""" +KV cache quantization strategy pattern. + +Three-player design: + quant_method (pure compute) ► Pool (buffer + batch dequant) ► Backend (view adaptation) +""" + +from abc import ABC, abstractmethod +from typing import Optional + +import torch +from torch import Tensor + +from sglang.srt.layers.quantization.kvfp4_tensor import E2M1_MAX + + +class FP4KVCacheQuantMethod(ABC): + """Abstract base for FP4 KV cache quantization strategies. + + Owns the quantize/dequantize computation. The Pool owns the buffers and + orchestrates the batch dequant loop. Backends only do view/reshape. + + All operations (quantize_and_store, dequantize_prev_kv) use FlashInfer + kernels or pure tensor ops, so they are CUDA-graph compatible. + """ + + name: str + SCALE_BLOCK_SIZE: int = 1 + + def needs_dequant_workspace(self) -> bool: + """Whether the pool should allocate dq_k_buffer / dq_v_buffer for prefill.""" + return False + + def needs_global_scale(self) -> bool: + """Whether this method uses a per-layer global FP32 scale.""" + return False + + @abstractmethod + def create_buffers( + self, size: int, head_num: int, head_dim: int, layer_num: int, device: str + ) -> dict: + """Allocate and return a buffer dict: + { + "k_buffer": list[Tensor], # per-layer, shape (size, head_num, head_dim//2) + "v_buffer": list[Tensor], + "k_scale_buffer": list[Tensor] | None, + "v_scale_buffer": list[Tensor] | None, + "dq_k_buffer": Tensor | None, # shared across layers (FP8 E4M3) + "dq_v_buffer": Tensor | None, + "store_dtype": torch.dtype, + } + """ + + @abstractmethod + def quantize_and_store( + self, + k_buffer: Tensor, + v_buffer: Tensor, + k_scale_buffer: Optional[Tensor], + v_scale_buffer: Optional[Tensor], + loc: Tensor, + cache_k: Tensor, + cache_v: Tensor, + k_scale=None, + v_scale=None, + ) -> None: + """Quantize cache_k / cache_v and write into buffers at loc.""" + + @abstractmethod + def dequantize_prev_kv( + self, + k_fp4: Tensor, + k_scales: Tensor, + v_fp4: Tensor, + v_scales: Tensor, + layer_id: int, + ) -> tuple[Tensor, Tensor]: + """Dequantize stored FP4 KV (selected token indices already applied). + + Returns: + (k_fp8, v_fp8): Both in torch.float8_e4m3fn dtype with shape + matching the input (after unpacking). These are written into the + shared dequant workspace buffer for the FlashInfer FP8 prefill kernel. + """ + + @abstractmethod + def compute_cell_size( + self, head_num: int, head_dim: int, num_layers: int, kv_size: int + ) -> int: + """Per-token memory footprint in bytes (for capacity estimation).""" + + def load_scales_from_model(self, model_runner, sm_version: int = None) -> None: + """Load per-layer global scales from model weights (no-op by default).""" + pass + + +class NVFP4KVMethod(FP4KVCacheQuantMethod): + """NVFP4 two-level scaling: global FP32 + per-block FP8 E4M3. + + Supported on SM100 and SM120. + """ + + name = "nvfp4" + SCALE_BLOCK_SIZE = 16 + + def __init__(self, num_layers: int, device: str, sm_version: int = 120): + self.num_layers = num_layers + self.device = device + self.sm_version = sm_version + # Per-layer global FP32 scales; filled by load_scales_from_model() + self.k_scales_gpu = torch.ones(num_layers, dtype=torch.float32, device=device) + self.v_scales_gpu = torch.ones(num_layers, dtype=torch.float32, device=device) + + def needs_dequant_workspace(self) -> bool: + return ( + True # prefill uses FP8 dequant workspace; future native FP4 kernel → False + ) + + def needs_global_scale(self) -> bool: + return True + + def load_scales_from_model(self, model_runner, sm_version: int = None) -> None: + if sm_version is not None: + self.sm_version = sm_version + + from sglang.srt.model_executor.model_runner import resolve_language_model + + language_model = resolve_language_model(model_runner.model) + + attention_layers = [] + for layer in language_model.layers: + if hasattr(layer, "self_attn"): + if hasattr(layer.self_attn, "attn"): + attention_layers.append(layer.self_attn.attn) + elif hasattr(layer.self_attn, "attn_mqa"): + attention_layers.append(layer.self_attn.attn_mqa) + elif hasattr(layer, "attn"): + attention_layers.append(layer.attn) + elif hasattr(layer, "attention"): + if hasattr(layer.attention, "attn"): + attention_layers.append(layer.attention.attn) + + if not attention_layers: + return + + # k_scales_gpu is indexed by global (absolute) layer_id. Resize if the model + # has layers with global IDs larger than what was pre-allocated. + # This happens in hybrid models (e.g., GDN) where only a subset of layers + # are full-attention, but their layer_ids are non-contiguous. + max_global_id = max(layer.layer_id for layer in attention_layers) + required_size = max_global_id + 1 + if required_size > len(self.k_scales_gpu): + self.k_scales_gpu = torch.ones( + required_size, dtype=torch.float32, device=self.device + ) + self.v_scales_gpu = torch.ones( + required_size, dtype=torch.float32, device=self.device + ) + + k_scales_cpu = self.k_scales_gpu.cpu().clone() + v_scales_cpu = self.v_scales_gpu.cpu().clone() + + for layer in attention_layers: + layer_id = layer.layer_id # global id + k_scale = ( + float(layer.k_scale) + if hasattr(layer, "k_scale") and layer.k_scale is not None + else 1.0 + ) + v_scale = ( + float(layer.v_scale) + if hasattr(layer, "v_scale") and layer.v_scale is not None + else 1.0 + ) + # SM100 uses TRT-LLM XQA kernels that expect KV scales as + # amax / 448, but the calibrated checkpoint stores amax / (6 * 448). + # We multiply by E2M1_MAX (6.0) to bridge the gap. SM120 uses a + # different kernel path where scales already include this factor. + # The FP4 data type itself is identical on both architectures. + # Reference: TRT-LLM FP8QDQLinearMethod.process_weights_after_loading_fused_qkv_linear + # https://github.com/NVIDIA/TensorRT-LLM/blob/main/tensorrt_llm/_torch/modules/linear.py + if self.sm_version == 100: + k_scale *= E2M1_MAX + v_scale *= E2M1_MAX + k_scales_cpu[layer_id] = k_scale + v_scales_cpu[layer_id] = v_scale + + self.k_scales_gpu.copy_(k_scales_cpu, non_blocking=True) + self.v_scales_gpu.copy_(v_scales_cpu, non_blocking=True) + + def create_buffers( + self, size: int, head_num: int, head_dim: int, layer_num: int, device: str + ) -> dict: + m = size + n = head_num + k = head_dim + store_dtype = torch.uint8 + dq_dtype = torch.float8_e4m3fn + + k_buffer = [ + torch.zeros((m, n, k // 2), dtype=store_dtype, device=device) + for _ in range(layer_num) + ] + v_buffer = [ + torch.zeros((m, n, k // 2), dtype=store_dtype, device=device) + for _ in range(layer_num) + ] + k_scale_buffer = [ + torch.zeros( + (m, n, k // self.SCALE_BLOCK_SIZE), dtype=store_dtype, device=device + ) + for _ in range(layer_num) + ] + v_scale_buffer = [ + torch.zeros( + (m, n, k // self.SCALE_BLOCK_SIZE), dtype=store_dtype, device=device + ) + for _ in range(layer_num) + ] + # Shared dequant workspace — one copy, reused per layer during prefill + dq_k_buffer = torch.zeros((m, n, k), dtype=dq_dtype, device=device) + dq_v_buffer = torch.zeros((m, n, k), dtype=dq_dtype, device=device) + + return { + "k_buffer": k_buffer, + "v_buffer": v_buffer, + "k_scale_buffer": k_scale_buffer, + "v_scale_buffer": v_scale_buffer, + "dq_k_buffer": dq_k_buffer, + "dq_v_buffer": dq_v_buffer, + "store_dtype": store_dtype, + } + + def quantize_and_store( + self, + k_buffer: Tensor, + v_buffer: Tensor, + k_scale_buffer: Optional[Tensor], + v_scale_buffer: Optional[Tensor], + loc: Tensor, + cache_k: Tensor, + cache_v: Tensor, + k_scale=None, + v_scale=None, + ) -> None: + from sglang.srt.layers.quantization.kvfp4_tensor import NVFP4KVQuantizeUtil + + cache_k, cache_k_fp4_sf, _ = NVFP4KVQuantizeUtil.quantize( + cache_k.contiguous(), k_scale + ) + cache_v, cache_v_fp4_sf, _ = NVFP4KVQuantizeUtil.quantize( + cache_v.contiguous(), v_scale + ) + + k_buffer[loc] = cache_k.view(torch.uint8) + v_buffer[loc] = cache_v.view(torch.uint8) + k_scale_buffer[loc] = cache_k_fp4_sf.view(torch.uint8) + v_scale_buffer[loc] = cache_v_fp4_sf.view(torch.uint8) + + def dequantize_prev_kv( + self, + k_fp4: Tensor, + k_scales: Tensor, + v_fp4: Tensor, + v_scales: Tensor, + layer_id: int, + ) -> tuple[Tensor, Tensor]: + """Dequantize FP4 KV (indexed tokens) → FP8 E4M3.""" + from sglang.srt.layers.quantization.kvfp4_tensor import NVFP4KVQuantizeUtil + + cur_k_scale = self.k_scales_gpu[layer_id : layer_id + 1] + cur_v_scale = self.v_scales_gpu[layer_id : layer_id + 1] + k_bf16 = NVFP4KVQuantizeUtil.dequantize( + k_fp4.view(torch.uint8), k_scales, cur_k_scale + ) + v_bf16 = NVFP4KVQuantizeUtil.dequantize( + v_fp4.view(torch.uint8), v_scales, cur_v_scale + ) + return k_bf16.to(torch.float8_e4m3fn), v_bf16.to(torch.float8_e4m3fn) + + def compute_cell_size( + self, head_num: int, head_dim: int, num_layers: int, kv_size: int + ) -> int: + # FP4 data: per-layer, K+V + fp4_size = head_num * (head_dim // 2) * num_layers * 2 * kv_size + # Block scales: per-layer, K+V (uint8) + scale_size = ( + head_num * (head_dim // self.SCALE_BLOCK_SIZE) * num_layers * 2 * kv_size + ) + # Dequant workspace: shared across layers (not multiplied by num_layers), FP8 + dq_size = head_num * head_dim * 2 * kv_size + return fp4_size + scale_size + dq_size + + +class BlockFP4KVMethod(FP4KVCacheQuantMethod): + """Block-wise FP4 single-level scaling (similar to MXFP4 but block_size=16).""" + + name = "blockfp4" + SCALE_BLOCK_SIZE = 16 + + def needs_dequant_workspace(self) -> bool: + return True + + def create_buffers( + self, size: int, head_num: int, head_dim: int, layer_num: int, device: str + ) -> dict: + m = size + store_dtype = torch.uint8 + dq_dtype = torch.float8_e4m3fn + + k_buffer = [ + torch.zeros((m, head_num, head_dim // 2), dtype=store_dtype, device=device) + for _ in range(layer_num) + ] + v_buffer = [ + torch.zeros((m, head_num, head_dim // 2), dtype=store_dtype, device=device) + for _ in range(layer_num) + ] + # MXFP4 flattens head dimensions for scale storage + k_scale_buffer = [ + torch.zeros( + (m, (head_num * head_dim) // self.SCALE_BLOCK_SIZE), + dtype=store_dtype, + device=device, + ) + for _ in range(layer_num) + ] + v_scale_buffer = [ + torch.zeros( + (m, (head_num * head_dim) // self.SCALE_BLOCK_SIZE), + dtype=store_dtype, + device=device, + ) + for _ in range(layer_num) + ] + dq_k_buffer = torch.zeros( + (m, head_num, head_dim), dtype=dq_dtype, device=device + ) + dq_v_buffer = torch.zeros( + (m, head_num, head_dim), dtype=dq_dtype, device=device + ) + + return { + "k_buffer": k_buffer, + "v_buffer": v_buffer, + "k_scale_buffer": k_scale_buffer, + "v_scale_buffer": v_scale_buffer, + "dq_k_buffer": dq_k_buffer, + "dq_v_buffer": dq_v_buffer, + "store_dtype": store_dtype, + } + + def quantize_and_store( + self, + k_buffer, + v_buffer, + k_scale_buffer, + v_scale_buffer, + loc, + cache_k, + cache_v, + k_scale=None, + v_scale=None, + ) -> None: + from sglang.srt.layers.quantization.kvfp4_tensor import BlockFP4KVQuantizeUtil + + cache_k_fp4, cache_k_sf = BlockFP4KVQuantizeUtil.batched_quantize(cache_k) + cache_v_fp4, cache_v_sf = BlockFP4KVQuantizeUtil.batched_quantize(cache_v) + k_buffer[loc] = cache_k_fp4 + v_buffer[loc] = cache_v_fp4 + k_scale_buffer[loc] = cache_k_sf + v_scale_buffer[loc] = cache_v_sf + + def dequantize_prev_kv( + self, + k_fp4: Tensor, + k_scales: Tensor, + v_fp4: Tensor, + v_scales: Tensor, + layer_id: int, + ) -> tuple[Tensor, Tensor]: + from sglang.srt.layers.quantization.kvfp4_tensor import BlockFP4KVQuantizeUtil + + k_bf16 = BlockFP4KVQuantizeUtil.batched_dequantize(k_fp4, k_scales) + v_bf16 = BlockFP4KVQuantizeUtil.batched_dequantize(v_fp4, v_scales) + return k_bf16.to(torch.float8_e4m3fn), v_bf16.to(torch.float8_e4m3fn) + + def compute_cell_size( + self, head_num: int, head_dim: int, num_layers: int, kv_size: int + ) -> int: + fp4_size = head_num * (head_dim // 2) * num_layers * 2 * kv_size + scale_size = ( + (head_num * head_dim // self.SCALE_BLOCK_SIZE) * num_layers * 2 * kv_size + ) + dq_size = head_num * head_dim * 2 * kv_size + return fp4_size + scale_size + dq_size + + +# Registry: name → class. Only classes for fp4_e2m1 dtype need to be listed. +FP4_KV_CACHE_QUANT_REGISTRY: dict[str, type[FP4KVCacheQuantMethod]] = { + "nvfp4": NVFP4KVMethod, + "blockfp4": BlockFP4KVMethod, +} + + +def get_fp4_kv_cache_quant_method(name: str, **kwargs) -> FP4KVCacheQuantMethod: + """Instantiate a FP4KVCacheQuantMethod by recipe name.""" + if name not in FP4_KV_CACHE_QUANT_REGISTRY: + raise ValueError( + f"Unknown fp4_kv_cache_recipe: '{name}'. " + f"Available: {list(FP4_KV_CACHE_QUANT_REGISTRY)}" + ) + return FP4_KV_CACHE_QUANT_REGISTRY[name](**kwargs) diff --git a/python/sglang/srt/layers/quantization/kvfp4_tensor.py b/python/sglang/srt/layers/quantization/kvfp4_tensor.py index 545199ff0..2d5eb45d1 100644 --- a/python/sglang/srt/layers/quantization/kvfp4_tensor.py +++ b/python/sglang/srt/layers/quantization/kvfp4_tensor.py @@ -12,21 +12,58 @@ # limitations under the License. # ============================================================================== +# Define a enum class for FP4 formats, including MXFP4, NVFP4 and future formats +from enum import Enum + import torch + +class FP4KVCacheRecipe(Enum): + MXFP4 = 1 # KVFP4: block-wise scaling + NVFP4 = 2 # two-level scaling: global FP32 + block FP8 E4M3 + + E2M1_MAX = 6.0 +MAX_BLOCK_SCALE_FP8 = 448.0 # Maximum FP8 E4M3 value # Put constants directly on CUDA if available _device = "cuda" if torch.cuda.is_available() else "cpu" +# E2M1 format: 1 sign bit + 2 exponent bits + 1 mantissa bit = 4 bits +# 16 possible values: 0x0-0xF +# Negative values: 0x8-0xF (sign bit = 1) +# Positive values: 0x0-0x7 (sign bit = 0) E2M1_VALUES = torch.tensor( - [0, 0.5, 1, 1.5, 2, 3, 4, 6], dtype=torch.float32, device=_device + [ + 0, + 0.5, + 1, + 1.5, + 2, + 3, + 4, + 6, # 0x0-0x7: positive values + -0, + -0.5, + -1, + -1.5, + -2, + -3, + -4, + -6, + ], # 0x8-0xF: negative values + dtype=torch.float32, + device=_device, ) E2M1_BOUNDS = torch.tensor( [0.25, 0.75, 1.25, 1.75, 2.5, 3.5, 5], dtype=torch.float32, device=_device ) -class KVFP4QuantizeUtil: - """Utility class for MXFP4 quantization and dequantization operations.""" +class BlockFP4KVQuantizeUtil: + """Block-wise FP4 (E2M1) quantization for KV cache. + + Similar to MXFP4 but uses block_size=16 (MXFP4 spec defines block_size=32). + Each block of 16 elements shares one uint8 exponent-only scale factor. + """ @staticmethod @torch.compile @@ -110,3 +147,127 @@ class KVFP4QuantizeUtil: scaled = reshaped * torch.exp2(scale_exp.unsqueeze(-1)) return scaled.view(b, m, n).to(dtype) + + +class NVFP4KVQuantizeUtil: + """Utility class for NVFP4 quantization and dequantization with two-level scaling + (global FP32 + block FP8 E4M3). + + Quantize formula: x_fp4 * block_scale * global_scale = x_bf16 + - Quantize: ``nvfp4_kv_quantize`` (SM100+), fallback ``fp4_quantize`` (SM90) + - Dequantize: ``nvfp4_kv_dequantize`` (SM100+) + """ + + @staticmethod + def quantize( + tensor: torch.Tensor, global_scale: torch.Tensor + ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]: + """Quantize BF16/FP16 tensor to NVFP4 format. + + Requires SM90+. Uses ``nvfp4_kv_quantize`` on SM100+ (native PTX), + falls back to ``fp4_quantize`` on SM90. + + Args: + tensor: Input tensor of shape [B, M, N] + global_scale: Global scale factor (float32 scalar or 1-element tensor) + + Returns: + (fp4_data, block_scales, global_scale): + fp4_data: shape [B, M, N/2], dtype uint8 + block_scales: shape [B, M, N/16], dtype float8_e4m3fn + global_scale: passthrough + """ + from sglang.srt.utils import is_sm90_supported, is_sm100_supported + + assert is_sm90_supported(), "NVFP4 KV cache quantize requires SM90+ GPU" + + b, m, n = tensor.shape + tensor_2d = tensor.reshape(b * m, n) + + if isinstance(global_scale, (int, float)): + global_scale = torch.tensor( + [global_scale], dtype=torch.float32, device=tensor.device + ) + elif global_scale.dim() == 0: + global_scale = global_scale.unsqueeze(0) + + if is_sm100_supported(): + from flashinfer import nvfp4_kv_quantize + + # nvfp4_kv_quantize takes global_scale directly (not inverted) + fp4_2d, scales_2d = nvfp4_kv_quantize(tensor_2d, global_scale) + else: + # SM90: fp4_quantize takes inverted global_scale + from flashinfer import fp4_quantize + + global_scale_inv = 1.0 / global_scale + fp4_2d, scales_2d = fp4_quantize( + tensor_2d, + global_scale_inv, + sf_vec_size=16, + sf_use_ue8m0=False, + is_sf_swizzled_layout=False, + is_sf_8x4_layout=False, + enable_pdl=None, + ) + + fp4_data = fp4_2d.view(b, m, fp4_2d.shape[-1]) + block_scales = scales_2d.view(b, m, scales_2d.shape[-1]).view( + torch.float8_e4m3fn + ) + return fp4_data, block_scales, global_scale + + @staticmethod + def dequantize( + quant_tensor: torch.Tensor, + block_scales: torch.Tensor, + global_scale: torch.Tensor, + dtype: torch.dtype = torch.bfloat16, + ) -> torch.Tensor: + """Dequantize NVFP4 tensor to BF16/FP16. + + Uses ``nvfp4_kv_dequantize`` on SM100+, falls back to pure PyTorch + E2M1 LUT on SM90. + + Args: + quant_tensor: Packed FP4 data of shape [B, M, N/2] (uint8) + block_scales: Per-block FP8 E4M3 scales of shape [B, M, N/16] + global_scale: Global scale factor (float32) + dtype: Output dtype (bfloat16 or float16) + + Returns: + Dequantized tensor of shape [B, M, N] + """ + from sglang.srt.utils import is_sm100_supported + + b, m, n_half = quant_tensor.shape + + if isinstance(global_scale, (int, float)): + global_scale = torch.tensor( + [global_scale], dtype=torch.float32, device=quant_tensor.device + ) + elif global_scale.dim() == 0: + global_scale = global_scale.unsqueeze(0) + + if is_sm100_supported(): + from flashinfer import nvfp4_kv_dequantize + + quant_2d = quant_tensor.view(torch.uint8).reshape(b * m, n_half) + scales_2d = block_scales.view(torch.uint8).reshape(b * m, -1) + output_2d = nvfp4_kv_dequantize( + quant_2d, scales_2d, global_scale, output_dtype=dtype + ) + return output_2d.reshape(b, m, -1) + else: + # Pure PyTorch fallback for SM90 + n = n_half * 2 + fp4_vals = torch.empty( + b, m, n, dtype=torch.uint8, device=quant_tensor.device + ) + fp4_vals[..., 0::2] = quant_tensor & 0x0F + fp4_vals[..., 1::2] = (quant_tensor >> 4) & 0x0F + float_vals = E2M1_VALUES[fp4_vals.long()] + reshaped = float_vals.view(b, m * n // 16, 16) + block_scales_float = block_scales.float().unsqueeze(-1) + scaled = reshaped * block_scales_float + return (scaled.view(b, m, n) * global_scale).to(dtype) diff --git a/test/registered/unit/layers/quantization/test_fp4_kv_cache_quant_method.py b/test/registered/unit/layers/quantization/test_fp4_kv_cache_quant_method.py new file mode 100644 index 000000000..2c31bf238 --- /dev/null +++ b/test/registered/unit/layers/quantization/test_fp4_kv_cache_quant_method.py @@ -0,0 +1,259 @@ +"""Unit tests for FP4 KV cache quantization strategy pattern — no server, no model loading.""" + +from sglang.test.ci.ci_register import register_cpu_ci + +register_cpu_ci(est_time=5, suite="stage-a-test-cpu") + +import unittest + +import torch + +from sglang.test.test_utils import CustomTestCase + + +def skip_if_no_cuda(func): + """Skip test if CUDA is not available.""" + return unittest.skipUnless(torch.cuda.is_available(), "CUDA not available")(func) + + +class TestKVCacheQuantRegistry(CustomTestCase): + """Test the registry and factory function.""" + + def test_registry_contains_nvfp4_and_mxfp4(self): + from sglang.srt.layers.quantization.fp4_kv_cache_quant_method import ( + FP4_KV_CACHE_QUANT_REGISTRY, + ) + + self.assertIn("nvfp4", FP4_KV_CACHE_QUANT_REGISTRY) + self.assertIn("blockfp4", FP4_KV_CACHE_QUANT_REGISTRY) + + def test_factory_nvfp4(self): + from sglang.srt.layers.quantization.fp4_kv_cache_quant_method import ( + NVFP4KVMethod, + get_fp4_kv_cache_quant_method, + ) + + method = get_fp4_kv_cache_quant_method( + "nvfp4", num_layers=4, device="cpu", sm_version=120 + ) + self.assertIsInstance(method, NVFP4KVMethod) + self.assertEqual(method.name, "nvfp4") + + def test_factory_mxfp4(self): + from sglang.srt.layers.quantization.fp4_kv_cache_quant_method import ( + BlockFP4KVMethod, + get_fp4_kv_cache_quant_method, + ) + + method = get_fp4_kv_cache_quant_method("blockfp4") + self.assertIsInstance(method, BlockFP4KVMethod) + self.assertEqual(method.name, "blockfp4") + + def test_factory_unknown_raises(self): + from sglang.srt.layers.quantization.fp4_kv_cache_quant_method import ( + get_fp4_kv_cache_quant_method, + ) + + with self.assertRaises(ValueError): + get_fp4_kv_cache_quant_method("unknown_method") + + +class TestNVFP4KVMethod(CustomTestCase): + """Test NVFP4KVMethod buffer creation and properties.""" + + def test_properties(self): + from sglang.srt.layers.quantization.fp4_kv_cache_quant_method import ( + NVFP4KVMethod, + ) + + m = NVFP4KVMethod(num_layers=4, device="cpu", sm_version=120) + self.assertEqual(m.name, "nvfp4") + self.assertEqual(m.SCALE_BLOCK_SIZE, 16) + self.assertTrue(m.needs_dequant_workspace()) + self.assertTrue(m.needs_global_scale()) + + def test_create_buffers_shapes(self): + from sglang.srt.layers.quantization.fp4_kv_cache_quant_method import ( + NVFP4KVMethod, + ) + + m = NVFP4KVMethod(num_layers=4, device="cpu", sm_version=120) + size, heads, dim, layers = 64, 8, 128, 4 + bufs = m.create_buffers(size, heads, dim, layers, "cpu") + + self.assertEqual(len(bufs["k_buffer"]), layers) + self.assertEqual(len(bufs["v_buffer"]), layers) + self.assertEqual(len(bufs["k_scale_buffer"]), layers) + self.assertEqual(len(bufs["v_scale_buffer"]), layers) + + # FP4 packed: (size, heads, dim//2) + self.assertEqual(bufs["k_buffer"][0].shape, (size, heads, dim // 2)) + # Block scales: (size, heads, dim//16) + self.assertEqual(bufs["k_scale_buffer"][0].shape, (size, heads, dim // 16)) + # Dequant workspace: (size, heads, dim), FP8 + self.assertEqual(bufs["dq_k_buffer"].shape, (size, heads, dim)) + self.assertEqual(bufs["dq_k_buffer"].dtype, torch.float8_e4m3fn) + self.assertEqual(bufs["store_dtype"], torch.uint8) + + def test_compute_cell_size(self): + from sglang.srt.layers.quantization.fp4_kv_cache_quant_method import ( + NVFP4KVMethod, + ) + + m = NVFP4KVMethod(num_layers=4, device="cpu") + cell = m.compute_cell_size(head_num=8, head_dim=128, num_layers=4, kv_size=1) + # FP4: 8*64*4*2 = 4096, scales: 8*8*4*2 = 512, dq: 8*128*2 = 2048 + self.assertEqual(cell, 4096 + 512 + 2048) + + def test_scales_init(self): + from sglang.srt.layers.quantization.fp4_kv_cache_quant_method import ( + NVFP4KVMethod, + ) + + m = NVFP4KVMethod(num_layers=4, device="cpu") + # Default scales should be 1.0 + self.assertTrue(torch.all(m.k_scales_gpu == 1.0)) + self.assertTrue(torch.all(m.v_scales_gpu == 1.0)) + self.assertEqual(len(m.k_scales_gpu), 4) + + @skip_if_no_cuda + def test_quantize_dequantize_roundtrip(self): + """Test NVFP4 quantize→dequantize roundtrip on CUDA.""" + from sglang.srt.layers.quantization.fp4_kv_cache_quant_method import ( + NVFP4KVMethod, + ) + + m = NVFP4KVMethod(num_layers=1, device="cuda", sm_version=120) + size, heads, dim = 32, 8, 128 + bufs = m.create_buffers(size, heads, dim, 1, "cuda") + + # Create random input + k = torch.randn(4, heads, dim, dtype=torch.bfloat16, device="cuda") + v = torch.randn(4, heads, dim, dtype=torch.bfloat16, device="cuda") + loc = torch.arange(4, device="cuda") + + # Quantize + m.quantize_and_store( + bufs["k_buffer"][0], + bufs["v_buffer"][0], + bufs["k_scale_buffer"][0], + bufs["v_scale_buffer"][0], + loc, + k, + v, + k_scale=m.k_scales_gpu[0:1], + v_scale=m.v_scales_gpu[0:1], + ) + + # Dequantize + k_fp4 = bufs["k_buffer"][0][loc] + k_scales = bufs["k_scale_buffer"][0][loc] + v_fp4 = bufs["v_buffer"][0][loc] + v_scales = bufs["v_scale_buffer"][0][loc] + k_out, v_out = m.dequantize_prev_kv(k_fp4, k_scales, v_fp4, v_scales, 0) + + # Check shapes + self.assertEqual(k_out.shape, (4, heads, dim)) + self.assertEqual(k_out.dtype, torch.float8_e4m3fn) + + # Check roundtrip error is bounded (FP4 is very lossy, ~20% relative error) + k_ref = k.float() + k_rec = k_out.float() + rel_error = (k_ref - k_rec).abs().mean() / k_ref.abs().mean() + self.assertLess( + rel_error, 0.5, f"NVFP4 roundtrip error too high: {rel_error:.3f}" + ) + + +class TestBlockFP4KVMethod(CustomTestCase): + """Test BlockFP4KVMethod buffer creation and roundtrip.""" + + def test_properties(self): + from sglang.srt.layers.quantization.fp4_kv_cache_quant_method import ( + BlockFP4KVMethod, + ) + + m = BlockFP4KVMethod() + self.assertEqual(m.name, "blockfp4") + self.assertTrue(m.needs_dequant_workspace()) + self.assertFalse(m.needs_global_scale()) + + def test_create_buffers_shapes(self): + from sglang.srt.layers.quantization.fp4_kv_cache_quant_method import ( + BlockFP4KVMethod, + ) + + m = BlockFP4KVMethod() + size, heads, dim, layers = 64, 8, 128, 4 + bufs = m.create_buffers(size, heads, dim, layers, "cpu") + + self.assertEqual(len(bufs["k_buffer"]), layers) + self.assertEqual(bufs["k_buffer"][0].shape, (size, heads, dim // 2)) + # MXFP4 flattens head dims for scales + self.assertEqual(bufs["k_scale_buffer"][0].shape, (size, (heads * dim) // 16)) + + def test_quantize_dequantize_roundtrip_cpu(self): + """Test MXFP4 quantize→dequantize roundtrip on CPU.""" + from sglang.srt.layers.quantization.fp4_kv_cache_quant_method import ( + BlockFP4KVMethod, + ) + + m = BlockFP4KVMethod() + size, heads, dim = 32, 8, 128 + bufs = m.create_buffers(size, heads, dim, 1, "cpu") + + k = torch.randn(4, heads, dim, dtype=torch.bfloat16) + v = torch.randn(4, heads, dim, dtype=torch.bfloat16) + loc = torch.arange(4) + + # Quantize + m.quantize_and_store( + bufs["k_buffer"][0], + bufs["v_buffer"][0], + bufs["k_scale_buffer"][0], + bufs["v_scale_buffer"][0], + loc, + k, + v, + ) + + # Dequantize + k_fp4 = bufs["k_buffer"][0][loc] + k_scales = bufs["k_scale_buffer"][0][loc] + v_fp4 = bufs["v_buffer"][0][loc] + v_scales = bufs["v_scale_buffer"][0][loc] + k_out, v_out = m.dequantize_prev_kv(k_fp4, k_scales, v_fp4, v_scales, 0) + + self.assertEqual(k_out.shape, (4, heads, dim)) + self.assertEqual(k_out.dtype, torch.float8_e4m3fn) + + +class TestBlockFP4KVQuantizeUtil(CustomTestCase): + """Test the existing MXFP4 BlockFP4KVQuantizeUtil roundtrip.""" + + def test_roundtrip_cpu(self): + from sglang.srt.layers.quantization.kvfp4_tensor import BlockFP4KVQuantizeUtil + + x = torch.randn(4, 8, 128, dtype=torch.bfloat16) + packed, scales = BlockFP4KVQuantizeUtil.batched_quantize(x) + reconstructed = BlockFP4KVQuantizeUtil.batched_dequantize(packed, scales) + + self.assertEqual(reconstructed.shape, x.shape) + rel_error = ( + x.float() - reconstructed.float() + ).abs().mean() / x.float().abs().mean() + self.assertLess(rel_error, 0.5) + + +class TestFP4KVCacheRecipe(CustomTestCase): + """Test enum.""" + + def test_enum_values(self): + from sglang.srt.layers.quantization.kvfp4_tensor import FP4KVCacheRecipe + + self.assertEqual(FP4KVCacheRecipe.MXFP4.value, 1) + self.assertEqual(FP4KVCacheRecipe.NVFP4.value, 2) + + +if __name__ == "__main__": + unittest.main()