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

[Feature] Add FP4 KV Cache Design and support SM120 GPUs (#21601)

Browse Source
This commit is contained in:
Sam (Kesen Li)
2026-07-17 14:49:43 -07:00
committed by GitHub
parent 7fc3fb9657
commit ec6a3163b7
19 changed files with 1829 additions and 327 deletions
Download Patch File Download Diff File Expand all files Collapse all files
test/registered/models_e2e/test_llama8b_nvfp4_kv_cache_sm120.py
+64
View File
@@ -0,0 +1,64 @@
import unittest
from sglang.srt.utils.common import is_sm120_supported
from sglang.test.accuracy_test_runner import AccuracyTestParams
from sglang.test.ci.ci_register import register_cuda_ci
from sglang.test.run_combined_tests import run_combined_tests
from sglang.test.test_utils import CustomTestCase, ModelLaunchSettings
register_cuda_ci(est_time=300, stage="extra-a", runner_config="1-gpu-small")
LLAMA8B_NVFP4_MODEL = "nvidia/Llama-3.1-8B-Instruct-NVFP4"
TP_SIZE = 1
@unittest.skipUnless(
is_sm120_supported(), "requires at least 1 SM120 GPU with CUDA 12.8+"
)
class TestLlama8BNVFP4KVCacheSM120(CustomTestCase):
"""Llama-3.1-8B-Instruct-NVFP4 with NVFP4 KV cache on SM120."""
def test_gsm8k(self):
variants = [
ModelLaunchSettings(
LLAMA8B_NVFP4_MODEL,
tp_size=TP_SIZE,
extra_args=[
"--quantization",
"modelopt_fp4",
"--fp4-gemm-backend",
"auto",
"--kv-cache-dtype",
"nvfp4",
"--prefill-attention-backend",
"flashinfer",
"--decode-attention-backend",
"trtllm_mha",
"--page-size",
"64",
"--cuda-graph-backend-prefill=disabled",
],
variant="NVFP4-GEMM+NVFP4-KV+SM120-XQA",
)
]
run_combined_tests(
models=variants,
test_name="Llama-3.1-8B-Instruct-NVFP4-KV-SM120",
accuracy_params=AccuracyTestParams(
dataset="gsm8k",
# Full GSM8K measured locally with 1319 requested / 1314 scored:
# - FP8 KV: 0.6461187214611872
# - NVFP4 KV: 0.632420091324201
# Keep the threshold 0.015 below the NVFP4 KV score.
baseline_accuracy=0.632420091324201 - 0.015,
num_examples=1319,
num_threads=200,
max_tokens=512,
api="completion",
),
)
if __name__ == "__main__":
unittest.main()
test/manual/quant/test_kvfp4_quant_dequant.py → test/registered/quant/test_kvfp4_quant_dequant.py
+18 -14
View File
@@ -1,12 +1,16 @@
#!/usr/bin/env python3
import sys
import time
import numpy as np
import pytest
import torch
from sglang.srt.layers.quantization.kvfp4_tensor import BlockFP4KVQuantizeUtil
from sglang.srt.layers.quantization.kvfp4_tensor import FP4MXBlock16KVQuantizeUtil
from sglang.test.ci.ci_register import register_cuda_ci
register_cuda_ci(est_time=30, stage="base-b", runner_config="1-gpu-large")
def calculate_accuracy_metrics(
@@ -28,7 +32,7 @@ def calculate_accuracy_metrics(
return {"MSE": mse, "MAE": mae, "PSNR": psnr, "Relative Error": rel_error}
def run_benchmark(m, n, k, num_runs=100) -> dict[str, dict[str, float]]:
def run_benchmark(m, n, k, num_runs=10) -> dict[str, dict[str, float]]:
"""Run FP8 vs KVFP4 quantization benchmark and return metrics."""
tensor_bf16 = torch.randn(m, n, k, dtype=torch.bfloat16, device="cuda")
@@ -52,18 +56,20 @@ 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 = BlockFP4KVQuantizeUtil.batched_quantize(tensor_bf16)
_ = BlockFP4KVQuantizeUtil.batched_dequantize(tensor_fp4, scale_factors)
tensor_fp4, scale_factors = FP4MXBlock16KVQuantizeUtil.batched_quantize(tensor_bf16)
_ = FP4MXBlock16KVQuantizeUtil.batched_dequantize(tensor_fp4, scale_factors)
start = time.time()
for _ in range(num_runs):
tensor_fp4, scale_factors = BlockFP4KVQuantizeUtil.batched_quantize(tensor_bf16)
tensor_fp4, scale_factors = FP4MXBlock16KVQuantizeUtil.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 = BlockFP4KVQuantizeUtil.batched_dequantize(
tensor_fp4_dequant = FP4MXBlock16KVQuantizeUtil.batched_dequantize(
tensor_fp4, scale_factors
)
torch.cuda.synchronize()
@@ -91,14 +97,8 @@ def run_benchmark(m, n, k, num_runs=100) -> dict[str, dict[str, float]]:
MNK_FACTORS = [
(64, 1, 576),
(512, 1, 576),
(1024, 1, 576),
(4096, 1, 576),
(2868672, 1, 576),
(64, 8, 64),
(512, 8, 64),
(1024, 8, 64),
(4096, 8, 64),
(2868672, 8, 64),
]
@@ -112,5 +112,9 @@ def test_kvfp4_quant_dequant(m, n, k):
print("FP4:", results["fp4"])
# Basic assertions to make sure metrics are reasonable
assert results["fp4"]["MSE"] < 1.0
assert results["fp8"]["MSE"] < 1.0
assert results["fp4"]["MSE"] < 0.1
assert results["fp8"]["MSE"] < 0.1
if __name__ == "__main__":
sys.exit(pytest.main([__file__, "-v"]))
test/registered/unit/layers/quantization/test_fp4_kv_cache_quant_method.py
+103 -61
View File
@@ -1,4 +1,4 @@
"""Unit tests for FP4 KV cache quantization strategy pattern — no server, no model loading."""
"""Unit tests for FP4 KV cache quantization strategy pattern - no server, no model loading."""
from sglang.test.ci.ci_register import register_cpu_ci
@@ -23,54 +23,87 @@ def skip_if_no_blackwell_nvfp4(func):
class TestKVCacheQuantRegistry(CustomTestCase):
"""Test the registry and factory function."""
def test_registry_contains_nvfp4_and_mxfp4(self):
def test_registry_contains_nvfp4_and_blockfp4(self):
from sglang.srt.layers.quantization.fp4_kv_cache_quant_method import (
FP4_KV_CACHE_QUANT_REGISTRY,
KV_CACHE_QUANT_REGISTRY,
)
self.assertIn("nvfp4", FP4_KV_CACHE_QUANT_REGISTRY)
self.assertIn("blockfp4", FP4_KV_CACHE_QUANT_REGISTRY)
self.assertIn("nvfp4", KV_CACHE_QUANT_REGISTRY)
self.assertIn("fp4_mx_block16", 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,
NVFP4KVCacheMethod,
get_kv_cache_quant_method,
)
method = get_fp4_kv_cache_quant_method(
"nvfp4", num_layers=4, device="cpu", sm_version=120
)
self.assertIsInstance(method, NVFP4KVMethod)
method = get_kv_cache_quant_method("nvfp4", num_layers=4, device="cpu")
self.assertIsInstance(method, NVFP4KVCacheMethod)
self.assertEqual(method.name, "nvfp4")
def test_factory_mxfp4(self):
def test_factory_blockfp4(self):
from sglang.srt.layers.quantization.fp4_kv_cache_quant_method import (
BlockFP4KVMethod,
get_fp4_kv_cache_quant_method,
FP4MXBlock16KVCacheMethod,
get_kv_cache_quant_method,
)
method = get_fp4_kv_cache_quant_method("blockfp4")
self.assertIsInstance(method, BlockFP4KVMethod)
self.assertEqual(method.name, "blockfp4")
method = get_kv_cache_quant_method("fp4_mx_block16")
self.assertIsInstance(method, FP4MXBlock16KVCacheMethod)
self.assertEqual(method.name, "fp4_mx_block16")
def test_factory_unknown_raises(self):
def test_resolve_explicit_recipes(self):
from sglang.srt.layers.quantization.fp4_kv_cache_quant_method import (
get_fp4_kv_cache_quant_method,
resolve_kv_cache_quant,
)
self.assertEqual(resolve_kv_cache_quant("nvfp4"), "nvfp4")
self.assertEqual(resolve_kv_cache_quant("fp4_mx_block16"), "fp4_mx_block16")
self.assertIsNone(resolve_kv_cache_quant("fp8_e4m3"))
def test_resolve_legacy_fp4_alias_raises(self):
from sglang.srt.layers.quantization.fp4_kv_cache_quant_method import (
resolve_kv_cache_quant,
)
with self.assertRaisesRegex(ValueError, "fp4_mx_block16"):
resolve_kv_cache_quant("fp4_e2m1")
def test_model_runner_rejects_legacy_fp4_alias(self):
from types import SimpleNamespace
from sglang.srt.model_executor.model_runner import ModelRunner
runner = object.__new__(ModelRunner)
runner.server_args = SimpleNamespace(kv_cache_dtype="fp4_e2m1")
with self.assertRaisesRegex(ValueError, "fp4_mx_block16"):
runner.configure_kv_cache_dtype()
def test_resolve_mxfp4_name_raises(self):
from sglang.srt.layers.quantization.fp4_kv_cache_quant_method import (
resolve_kv_cache_quant,
)
with self.assertRaises(ValueError):
get_fp4_kv_cache_quant_method("unknown_method")
resolve_kv_cache_quant("mxfp4")
def test_factory_unknown_raises(self):
from sglang.srt.layers.quantization.fp4_kv_cache_quant_method import (
get_kv_cache_quant_method,
)
with self.assertRaises(ValueError):
get_kv_cache_quant_method("unknown_method")
class TestNVFP4KVMethod(CustomTestCase):
"""Test NVFP4KVMethod buffer creation and properties."""
class TestNVFP4KVCacheMethod(CustomTestCase):
"""Test NVFP4KVCacheMethod buffer creation and properties."""
def test_properties(self):
from sglang.srt.layers.quantization.fp4_kv_cache_quant_method import (
NVFP4KVMethod,
NVFP4KVCacheMethod,
)
m = NVFP4KVMethod(num_layers=4, device="cpu", sm_version=120)
m = NVFP4KVCacheMethod(num_layers=4, device="cpu")
self.assertEqual(m.name, "nvfp4")
self.assertEqual(m.SCALE_BLOCK_SIZE, 16)
self.assertTrue(m.needs_dequant_workspace())
@@ -78,10 +111,10 @@ class TestNVFP4KVMethod(CustomTestCase):
def test_create_buffers_shapes(self):
from sglang.srt.layers.quantization.fp4_kv_cache_quant_method import (
NVFP4KVMethod,
NVFP4KVCacheMethod,
)
m = NVFP4KVMethod(num_layers=4, device="cpu", sm_version=120)
m = NVFP4KVCacheMethod(num_layers=4, device="cpu")
size, heads, dim, layers = 64, 8, 128, 4
bufs = m.create_buffers(size, heads, dim, layers, "cpu")
@@ -101,20 +134,20 @@ class TestNVFP4KVMethod(CustomTestCase):
def test_compute_cell_size(self):
from sglang.srt.layers.quantization.fp4_kv_cache_quant_method import (
NVFP4KVMethod,
NVFP4KVCacheMethod,
)
m = NVFP4KVMethod(num_layers=4, device="cpu")
m = NVFP4KVCacheMethod(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,
NVFP4KVCacheMethod,
)
m = NVFP4KVMethod(num_layers=4, device="cpu")
m = NVFP4KVCacheMethod(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))
@@ -122,13 +155,13 @@ class TestNVFP4KVMethod(CustomTestCase):
@skip_if_no_blackwell_nvfp4
def test_quantize_dequantize_roundtrip(self):
"""Test NVFP4 quantize→dequantize roundtrip on CUDA."""
"""Test NVFP4 quantize->dequantize roundtrip on CUDA."""
from sglang.srt.layers.quantization.fp4_kv_cache_quant_method import (
NVFP4KVMethod,
NVFP4KVCacheMethod,
)
major, minor = torch.cuda.get_device_capability()
m = NVFP4KVMethod(num_layers=1, device="cuda", sm_version=major * 10 + minor)
m = NVFP4KVCacheMethod(num_layers=1, device="cuda")
size, heads, dim = 32, 8, 128
bufs = m.create_buffers(size, heads, dim, 1, "cuda")
@@ -171,40 +204,55 @@ class TestNVFP4KVMethod(CustomTestCase):
)
class TestBlockFP4KVMethod(CustomTestCase):
"""Test BlockFP4KVMethod buffer creation and roundtrip."""
class TestFP4MXBlock16KVCacheMethod(CustomTestCase):
"""Test FP4MXBlock16KVCacheMethod buffer creation and roundtrip."""
def test_properties(self):
from sglang.srt.layers.quantization.fp4_kv_cache_quant_method import (
BlockFP4KVMethod,
FP4MXBlock16KVCacheMethod,
KVCacheAttentionAccessKind,
)
m = BlockFP4KVMethod()
self.assertEqual(m.name, "blockfp4")
self.assertTrue(m.needs_dequant_workspace())
m = FP4MXBlock16KVCacheMethod()
self.assertEqual(m.name, "fp4_mx_block16")
self.assertFalse(m.needs_dequant_workspace())
self.assertTrue(m.needs_plain_kv_dequant_read())
self.assertFalse(m.needs_global_scale())
self.assertEqual(m.plain_attention_kv_dtype(), torch.bfloat16)
self.assertEqual(
m.resolve_attention_access("prefill", "triton").kind,
KVCacheAttentionAccessKind.PLAIN,
)
self.assertEqual(
m.resolve_attention_access("decode", "trtllm_mha").kind,
KVCacheAttentionAccessKind.PLAIN,
)
self.assertIsNone(m.resolve_attention_access("prefill", "flashinfer"))
self.assertIsNone(m.resolve_attention_access("decode", "flashinfer"))
def test_create_buffers_shapes(self):
from sglang.srt.layers.quantization.fp4_kv_cache_quant_method import (
BlockFP4KVMethod,
FP4MXBlock16KVCacheMethod,
)
m = BlockFP4KVMethod()
m = FP4MXBlock16KVCacheMethod()
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
# Block-16 FP4 flattens head dims for scales
self.assertEqual(bufs["k_scale_buffer"][0].shape, (size, (heads * dim) // 16))
self.assertIsNone(bufs["dq_k_buffer"])
self.assertIsNone(bufs["dq_v_buffer"])
def test_quantize_dequantize_roundtrip_cpu(self):
"""Test MXFP4 quantize→dequantize roundtrip on CPU."""
"""Test block-16 FP4 quantize->dequantize roundtrip on CPU."""
from sglang.srt.layers.quantization.fp4_kv_cache_quant_method import (
BlockFP4KVMethod,
FP4MXBlock16KVCacheMethod,
)
m = BlockFP4KVMethod()
m = FP4MXBlock16KVCacheMethod()
size, heads, dim = 32, 8, 128
bufs = m.create_buffers(size, heads, dim, 1, "cpu")
@@ -231,18 +279,22 @@ class TestBlockFP4KVMethod(CustomTestCase):
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)
self.assertEqual(v_out.shape, (4, heads, dim))
self.assertEqual(k_out.dtype, torch.bfloat16)
self.assertEqual(v_out.dtype, torch.bfloat16)
class TestBlockFP4KVQuantizeUtil(CustomTestCase):
"""Test the existing MXFP4 BlockFP4KVQuantizeUtil roundtrip."""
class TestFP4MXBlock16KVQuantizeUtil(CustomTestCase):
"""Test the existing block-16 FP4 FP4MXBlock16KVQuantizeUtil roundtrip."""
def test_roundtrip_cpu(self):
from sglang.srt.layers.quantization.kvfp4_tensor import BlockFP4KVQuantizeUtil
from sglang.srt.layers.quantization.kvfp4_tensor import (
FP4MXBlock16KVQuantizeUtil,
)
x = torch.randn(4, 8, 128, dtype=torch.bfloat16)
packed, scales = BlockFP4KVQuantizeUtil.batched_quantize(x)
reconstructed = BlockFP4KVQuantizeUtil.batched_dequantize(packed, scales)
packed, scales = FP4MXBlock16KVQuantizeUtil.batched_quantize(x)
reconstructed = FP4MXBlock16KVQuantizeUtil.batched_dequantize(packed, scales)
self.assertEqual(reconstructed.shape, x.shape)
rel_error = (
@@ -251,15 +303,5 @@ class TestBlockFP4KVQuantizeUtil(CustomTestCase):
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()
test/registered/unit/mem_cache/test_quantized_kv_pool.py
+124
View File
@@ -0,0 +1,124 @@
# Copyright 2026 SGLang Team
# Licensed under the Apache License, Version 2.0
import types
import unittest
import torch
from sglang.test.ci.ci_register import register_cpu_ci
register_cpu_ci(est_time=5, suite="base-a-test-cpu")
class _FakeQuantMethod:
name = "fake_quant"
def __init__(self):
self.k_scales_gpu = torch.tensor([2.0], dtype=torch.float32)
self.v_scales_gpu = torch.tensor([3.0], dtype=torch.float32)
self.store_calls = []
def dequant_workspace_dtype(self):
return torch.float32
def create_buffers(self, size, head_num, head_dim, layer_num, device):
return {
"k_buffer": [
torch.zeros(
(size, head_num, head_dim), dtype=torch.uint8, device=device
)
for _ in range(layer_num)
],
"v_buffer": [
torch.zeros(
(size, head_num, head_dim), dtype=torch.uint8, device=device
)
for _ in range(layer_num)
],
"k_scale_buffer": [
torch.zeros((size, head_num, 1), dtype=torch.uint8, device=device)
for _ in range(layer_num)
],
"v_scale_buffer": [
torch.zeros((size, head_num, 1), dtype=torch.uint8, device=device)
for _ in range(layer_num)
],
"dq_k_buffer": torch.zeros(
(size, head_num, head_dim), dtype=torch.float32, device=device
),
"dq_v_buffer": torch.zeros(
(size, head_num, head_dim), dtype=torch.float32, device=device
),
"store_dtype": torch.uint8,
}
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,
):
self.store_calls.append(
{
"loc": loc,
"k_scale": k_scale,
"v_scale": v_scale,
"k_scale_buffer": k_scale_buffer,
"v_scale_buffer": v_scale_buffer,
}
)
k_buffer[loc] = 1
v_buffer[loc] = 2
k_scale_buffer[loc] = 3
v_scale_buffer[loc] = 4
class TestQuantizedKVPool(unittest.TestCase):
def test_quant_method_owns_buffers_and_store_path(self):
from sglang.srt.mem_cache.memory_pool import MHATokenToKVPool
quant_method = _FakeQuantMethod()
pool = MHATokenToKVPool(
size=4,
page_size=1,
dtype=torch.bfloat16,
head_num=1,
head_dim=8,
layer_num=1,
device="cpu",
enable_memory_saver=False,
quant_method=quant_method,
)
self.assertTrue(pool.is_quantized_kv_cache)
self.assertIs(pool.quant_method, quant_method)
self.assertIsNotNone(pool.k_scale_buffer)
self.assertIs(pool.get_dequant_workspace()[0], pool.dq_k_buffer)
loc = torch.tensor([0, 1], dtype=torch.int64)
layer = types.SimpleNamespace(layer_id=0)
pool.set_kv_buffer(
layer,
loc,
torch.zeros((2, 1, 8), dtype=torch.bfloat16),
torch.zeros((2, 1, 8), dtype=torch.bfloat16),
)
self.assertEqual(len(quant_method.store_calls), 1)
call = quant_method.store_calls[0]
self.assertIs(call["loc"], loc)
self.assertTrue(torch.equal(call["k_scale"], quant_method.k_scales_gpu[0:1]))
self.assertTrue(torch.equal(call["v_scale"], quant_method.v_scales_gpu[0:1]))
self.assertEqual(pool.k_buffer[0][loc].unique().tolist(), [1])
self.assertEqual(pool.v_buffer[0][loc].unique().tolist(), [2])
if __name__ == "__main__":
unittest.main()
test/registered/unit/server_args/test_server_args.py
+5 -3
View File
@@ -1136,7 +1136,7 @@ class TestPrefillOnlyDisableKvCache(unittest.TestCase):
- disable_radix_cache (radix cache otherwise indexes empty pool slots),
- no context-parallel attention (CP writes to the pool via set_kv_buffer),
- no HiSparse (uses a different pool family),
- kv_cache_dtype != fp4_e2m1 (FP4 pool is a separate allocation path).
- kv_cache_dtype is not nvfp4/fp4_mx_block16 (FP4 pool is a separate allocation path).
All other configurations must be rejected before model load.
"""
@@ -1186,8 +1186,10 @@ class TestPrefillOnlyDisableKvCache(unittest.TestCase):
self._validate_prefill_only_args(enable_hisparse=True)
def test_rejects_fp4_kv_cache(self):
with self.assertRaisesRegex(ValueError, "fp4_e2m1"):
self._validate_prefill_only_args(kv_cache_dtype="fp4_e2m1")
for kv_cache_dtype in ("nvfp4", "fp4_mx_block16"):
with self.subTest(kv_cache_dtype=kv_cache_dtype):
with self.assertRaisesRegex(ValueError, "nvfp4.*fp4_mx_block16"):
self._validate_prefill_only_args(kv_cache_dtype=kv_cache_dtype)
class TestSessionRadixCacheServerArgs(unittest.TestCase):