[1/4] NVFP4 KV cache: quantization strategy abstraction and kernel (#21954)

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)

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)

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()