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

[CPU] Support FP8 KV cache (#32733)

Browse Source 
Co-authored-by: Zheng, Beilei <beilei.zheng@intel.com>
Co-authored-by: mingfeima <mingfei.ma@intel.com>
This commit is contained in:
Xuan Liao
2026-09-02 10:20:54 +08:00
committed by GitHub
co-authored by Zheng, Beilei mingfeima
parent e874ae64cd
commit 26f760d5c0
14 changed files with 952 additions and 262 deletions
Download Patch File Download Diff File Expand all files Collapse all files
test/registered/cpu/test_decode.py
+57 -5
View File
@@ -158,7 +158,18 @@ class TestDecodeAttention(CustomTestCase):
return output
def _test_grouped_decode_attention_once(
self, B, H_Q, H_KV, D, D_V, sliding_window, sink, is_cross_attn, dtype, device
self,
B,
H_Q,
H_KV,
D,
D_V,
sliding_window,
sink,
is_cross_attn,
dtype,
device,
kvcache_dtype=torch.bfloat16,
):
# This represents the number of tokens already in the sequence
seq_len = 1024
@@ -176,14 +187,32 @@ class TestDecodeAttention(CustomTestCase):
# k_buffer and v_buffer represent all previous tokens
k_buffer = torch.randn(total_tokens, H_KV, D, dtype=dtype, device=device)
v_buffer = torch.randn(total_tokens, H_KV, D_V, dtype=dtype, device=device)
k_scale = 1.0
v_scale = 1.0
if kvcache_dtype == torch.float8_e4m3fn:
k_scale = 0.5
v_scale = 0.25
k_buffer_fp8 = (k_buffer / k_scale).to(torch.float8_e4m3fn)
v_buffer_fp8 = (v_buffer / v_scale).to(torch.float8_e4m3fn)
k_buffer = (k_buffer_fp8.float() * k_scale).to(dtype)
v_buffer = (v_buffer_fp8.float() * v_scale).to(dtype)
key = torch.randn(B, H_KV, D, dtype=dtype)
value = torch.randn(B, H_KV, D_V, dtype=dtype)
loc = torch.randint(0, 10, (B,)).to(torch.int64)
# set kv cache
k_buffer[loc] = key
v_buffer[loc] = value
if not is_cross_attn:
if kvcache_dtype == torch.float8_e4m3fn:
k_buffer[loc] = (
(key / k_scale).to(torch.float8_e4m3fn).float() * k_scale
).to(dtype)
v_buffer[loc] = (
(value / v_scale).to(torch.float8_e4m3fn).float() * v_scale
).to(dtype)
else:
k_buffer[loc] = key
v_buffer[loc] = value
# o will have the same shape as q
o = torch.zeros(B, H_Q, D_V, dtype=dtype, device=device)
@@ -212,8 +241,10 @@ class TestDecodeAttention(CustomTestCase):
value = value.transpose(0, 1).contiguous().transpose(0, 1)
torch.ops.sgl_kernel.decode_attention_cpu(
q,
k_buffer,
v_buffer,
(k_buffer if kvcache_dtype != torch.float8_e4m3fn else k_buffer_fp8),
(v_buffer if kvcache_dtype != torch.float8_e4m3fn else v_buffer_fp8),
k_scale,
v_scale,
o,
key if not is_cross_attn else None,
value if not is_cross_attn else None,
@@ -305,6 +336,27 @@ class TestDecodeAttention(CustomTestCase):
B, H_Q, H_KV, D, D_V, None, False, True, dtype=dtype, device=device
)
fp8_configs = [
(2, 32, 8, 33, 55, None, False, False),
(1, 16, 1, 576, 512, None, False, False),
(2, 16, 16, 64, 64, 10, True, False),
(2, 16, 1, 64, 64, None, False, True),
]
for B, H_Q, H_KV, D, D_V, sliding_window, sink, is_cross_attn in fp8_configs:
self._test_grouped_decode_attention_once(
B,
H_Q,
H_KV,
D,
D_V,
sliding_window,
sink,
is_cross_attn,
dtype=torch.bfloat16,
device=device,
kvcache_dtype=torch.float8_e4m3fn,
)
def test_grouped_decode_attention(self):
self._test_grouped_decode_attention("cpu")
test/registered/cpu/test_extend.py
+24 -3
View File
@@ -194,6 +194,7 @@ class TestExtendAttention(CustomTestCase):
has_sink=False,
mla=False,
is_cross_attn=False,
kvcache_dtype=torch.bfloat16,
*,
b_seq_len_prefix=None,
b_seq_len_extend=None,
@@ -243,6 +244,15 @@ class TestExtendAttention(CustomTestCase):
H_BUF = 1 if mla else H_KV
k_buffer = torch.randn((total_token_num, H_BUF, D), dtype=dtype)
v_buffer = torch.randn((total_token_num, H_BUF, DV), dtype=dtype)
k_scale = 1.0
v_scale = 1.0
if kvcache_dtype == torch.float8_e4m3fn:
k_scale = 0.5
v_scale = 0.25
k_buffer_fp8 = (k_buffer / k_scale).to(torch.float8_e4m3fn)
v_buffer_fp8 = (v_buffer / v_scale).to(torch.float8_e4m3fn)
k_buffer = (k_buffer_fp8.float() * k_scale).to(dtype)
v_buffer = (v_buffer_fp8.float() * v_scale).to(dtype)
k_extend = torch.empty((extend_token_num, H_KV, D), dtype=dtype)
v_extend = torch.empty((extend_token_num, H_KV, DV), dtype=dtype)
@@ -328,8 +338,10 @@ class TestExtendAttention(CustomTestCase):
None if kv_from_cache else k_extend,
None if kv_from_cache else v_extend,
o_extend,
k_buffer,
v_buffer,
(k_buffer if kvcache_dtype != torch.float8_e4m3fn else k_buffer_fp8),
(v_buffer if kvcache_dtype != torch.float8_e4m3fn else v_buffer_fp8),
k_scale,
v_scale,
req_to_tokens,
b_req_idx,
b_seq_len,
@@ -346,7 +358,8 @@ class TestExtendAttention(CustomTestCase):
is_causal,
)
torch.testing.assert_close(o_ref, o_extend, atol=1e-2, rtol=1e-2)
tolerance = 2e-2 if kv_from_cache else 1e-2
torch.testing.assert_close(o_ref, o_extend, atol=tolerance, rtol=tolerance)
def test_extend_attention(self):
for is_mla in [True, False]:
@@ -379,6 +392,14 @@ class TestExtendAttention(CustomTestCase):
1, 20, 1, 1, 64, 64, sliding_window, has_sink, False, False
)
fp8_configs = [
(1, 123, 16, 1, 128, 96, None, False, False, False),
(1, 123, 16, 1, 128, 96, None, False, False, True),
(1, 20, 1, 1, 64, 64, 10, True, False, False),
]
for config in fp8_configs:
self._test_extend_attention_once(*config, kvcache_dtype=torch.float8_e4m3fn)
def test_extend_attention_kv_from_cache(self):
# KV-shared layers pass no extend K/V, so the kernel masks the extend
# range causally itself; sizes straddle several BLOCK_N.
test/registered/cpu/test_mla.py
+2
View File
@@ -109,6 +109,8 @@ class TestMLA(CustomTestCase):
q,
k_buffer2,
v_buffer2,
1.0,
1.0,
o,
key,
value,
test/registered/cpu/test_spec_kernels.py
+2
View File
@@ -1114,6 +1114,8 @@ class TestExtendAttentionTreeMask(CustomTestCase):
o_extend,
k_buffer,
v_buffer,
1.0, # k_buf_scale
1.0, # v_buf_scale
req_to_token,
req_pool_indices,
seq_lens,
test/registered/cpu/test_store_cache.py
+89
View File
@@ -1,9 +1,14 @@
import sys
from types import SimpleNamespace
import pytest
import sgl_kernel # noqa: F401
import torch
from sglang.srt.layers.quantization.fp4_kv_cache_quant_method import (
CPUFP8KVCacheMethod,
)
from sglang.srt.mem_cache.memory_pool import MHATokenToKVPool
from sglang.test.ci.ci_register import register_cpu_ci
register_cpu_ci(est_time=14, suite="base-b-test-cpu")
@@ -80,5 +85,89 @@ def test_store_cache_int32_indices(batch_size, num_heads, head_dim, dtype):
assert torch.equal(v_cache, v_cache_ref)
@pytest.mark.parametrize(
("k_scale", "v_scale"),
[(None, None), (0.5, 0.25)],
ids=["unit-scale-default", "non-unit-static-scale"],
)
def test_mha_fp8_e4m3_pool_decode_numerics(k_scale, v_scale):
seq_len = 16
num_heads = 2
head_dim = 64
num_kv_splits = 8
sm_scale = head_dim**-0.5
pool = MHATokenToKVPool(
size=seq_len,
page_size=1,
dtype=torch.float8_e4m3fn,
head_num=num_heads,
head_dim=head_dim,
layer_num=1,
device=DEVICE,
enable_memory_saver=False,
quant_method=CPUFP8KVCacheMethod(),
)
layer = SimpleNamespace(layer_id=0)
loc = torch.arange(seq_len, dtype=torch.int64, device=DEVICE)
cache_k = torch.randn(
(seq_len, num_heads, head_dim), dtype=torch.bfloat16, device=DEVICE
)
cache_v = torch.randn(
(seq_len, num_heads, head_dim), dtype=torch.bfloat16, device=DEVICE
)
pool.set_kv_buffer(layer, loc, cache_k, cache_v, k_scale=k_scale, v_scale=v_scale)
effective_k_scale = 1.0 if k_scale is None else k_scale
effective_v_scale = 1.0 if v_scale is None else v_scale
k_dequant = (pool.get_key_buffer(0).float() * effective_k_scale).to(torch.bfloat16)
v_dequant = (pool.get_value_buffer(0).float() * effective_v_scale).to(
torch.bfloat16
)
query = torch.randn((1, num_heads, head_dim), dtype=torch.bfloat16, device=DEVICE)
output = torch.empty_like(query)
req_to_token = loc.to(torch.int32).unsqueeze(0)
req_pool_indices = torch.zeros(1, dtype=torch.int64, device=DEVICE)
seq_lens = torch.full((1,), seq_len, dtype=torch.int64, device=DEVICE)
attn_logits = torch.empty(
(1, num_heads, num_kv_splits, head_dim + 1),
dtype=torch.float32,
device=DEVICE,
)
torch.ops.sgl_kernel.decode_attention_cpu(
query,
pool.get_key_buffer(0),
pool.get_value_buffer(0),
effective_k_scale,
effective_v_scale,
output,
None,
None,
None,
attn_logits,
req_to_token,
req_pool_indices,
seq_lens,
sm_scale,
0.0,
False,
0,
None,
None,
)
output_ref = (
torch.nn.functional.scaled_dot_product_attention(
query.movedim(0, 1).unsqueeze(0),
k_dequant[:seq_len].movedim(0, 1).unsqueeze(0),
v_dequant[:seq_len].movedim(0, 1).unsqueeze(0),
scale=sm_scale,
)
.squeeze(0)
.movedim(1, 0)
)
torch.testing.assert_close(output, output_ref, atol=3e-2, rtol=1e-6)
if __name__ == "__main__":
sys.exit(pytest.main([__file__]))
test/registered/unit/layers/quantization/test_fp4_kv_cache_quant_method.py
+61 -1
View File
@@ -1,4 +1,4 @@
"""Unit tests for FP4 KV cache quantization strategy pattern - no server, no model loading."""
"""Unit tests for KV cache quantization strategies - no server, no model loading."""
from sglang.test.ci.ci_register import register_cpu_ci
@@ -30,6 +30,7 @@ class TestKVCacheQuantRegistry(CustomTestCase):
self.assertIn("nvfp4", KV_CACHE_QUANT_REGISTRY)
self.assertIn("fp4_mx_block16", KV_CACHE_QUANT_REGISTRY)
self.assertIn("cpu_fp8_e4m3", KV_CACHE_QUANT_REGISTRY)
def test_factory_nvfp4(self):
from sglang.srt.layers.quantization.fp4_kv_cache_quant_method import (
@@ -102,6 +103,65 @@ class TestKVCacheQuantRegistry(CustomTestCase):
get_kv_cache_quant_method("unknown_method")
class TestCPUFP8KVCacheMethod(CustomTestCase):
def test_static_scale_quantize_and_store(self):
from sglang.srt.layers.quantization.fp4_kv_cache_quant_method import (
CPUFP8KVCacheMethod,
)
method = CPUFP8KVCacheMethod()
buffers = method.create_buffers(4, 2, 8, 1, "cpu")
loc = torch.tensor([1, 3])
cache_k = torch.randn(2, 2, 8, dtype=torch.bfloat16)
cache_v = torch.randn(2, 2, 8, dtype=torch.bfloat16)
method.quantize_and_store(
buffers["k_buffer"][0],
buffers["v_buffer"][0],
buffers["k_scale_buffer"],
buffers["v_scale_buffer"],
loc,
cache_k,
cache_v,
k_scale=0.5,
v_scale=0.25,
)
torch.testing.assert_close(
buffers["k_buffer"][0][loc].float(),
(cache_k / 0.5).to(torch.float8_e4m3fn).float(),
)
torch.testing.assert_close(
buffers["v_buffer"][0][loc].float(),
(cache_v / 0.25).to(torch.float8_e4m3fn).float(),
)
self.assertIsNone(buffers["k_scale_buffer"])
self.assertIsNone(buffers["v_scale_buffer"])
self.assertEqual(method.compute_cell_size(2, 8, 1, 4), 128)
def test_defaults_to_unit_scales(self):
from sglang.srt.layers.quantization.fp4_kv_cache_quant_method import (
CPUFP8KVCacheMethod,
)
method = CPUFP8KVCacheMethod()
buffers = method.create_buffers(1, 1, 8, 1, "cpu")
cache = torch.ones(1, 1, 8, dtype=torch.bfloat16)
method.quantize_and_store(
buffers["k_buffer"][0],
buffers["v_buffer"][0],
buffers["k_scale_buffer"],
buffers["v_scale_buffer"],
torch.tensor([0]),
cache,
cache,
)
expected = cache.to(torch.float8_e4m3fn)
torch.testing.assert_close(buffers["k_buffer"][0][0], expected[0])
torch.testing.assert_close(buffers["v_buffer"][0][0], expected[0])
class TestNVFP4KVCacheMethod(CustomTestCase):
"""Test NVFP4KVCacheMethod buffer creation and properties."""