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

fix(hicache): support staged write-back for asymmetric MHA (#30981)

Browse Source 
Co-authored-by: 晟海 <huangtingwei.htw@antgroup.com>
Co-authored-by: Zhangheng <hzh0425@apache.org>
This commit is contained in:
HZY
2026-07-22 09:58:49 +08:00
committed by GitHub
co-authored by 晟海 Zhangheng
parent 4597dd4d88
commit 3217b7e3ce
5 changed files with 408 additions and 19 deletions
Download Patch File Download Diff File Expand all files Collapse all files
python/sglang/srt/mem_cache/pool_host/mha.py
+84 -18
View File
@@ -16,6 +16,9 @@ from sglang.jit_kernel.hicache import (
from sglang.jit_kernel.hicache import (
transfer_hicache_all_layer_mla as jit_transfer_hicache_all_layer_mla,
)
from sglang.jit_kernel.hicache import (
transfer_hicache_all_layer_mla_staged_lf_pf as jit_transfer_hicache_all_layer_mla_staged_lf_pf,
)
from sglang.jit_kernel.hicache import (
transfer_hicache_all_layer_staged_lf_pf as jit_transfer_hicache_all_layer_staged_lf_pf,
)
@@ -941,6 +944,51 @@ class AsymmetricMHATokenToKVPoolHost(MHATokenToKVPoolHost):
kernels derive copy sizes from each call's first tensor.
"""
def _init_write_back_staging_buffers(self):
self.staging_page_capacity = 0
self.staging_token_capacity = 0
self.staging_k_buffer = None
self.staging_v_buffer = None
self.can_use_write_back_jit = False
if self.layout != "page_first" or (_is_npu or _is_xpu or _is_mps):
return
# K and V have different element sizes. Use the single-buffer staged
# kernel for each side, which specializes to its native stride.
can_use_staged_jit = (_is_cuda or _is_hip) and all(
can_use_write_back_jit_kernel(element_size=element_size)
for element_size in (
self._k_token_stride_size(),
self._v_token_stride_size(),
)
)
if not can_use_staged_jit:
return
self.can_use_write_back_jit = True
self.staging_page_capacity = min(self.page_num, _WRITE_BACK_STAGING_PAGE_CHUNK)
self.staging_token_capacity = self.staging_page_capacity * self.page_size
self.staging_k_buffer = torch.empty(
(
self.staging_token_capacity,
self.layer_num,
self.head_num,
self.head_dim,
),
dtype=self.dtype,
device=self.device_pool.device,
)
self.staging_v_buffer = torch.empty(
(
self.staging_token_capacity,
self.layer_num,
self.head_num,
self.v_head_dim,
),
dtype=self.dtype,
device=self.device_pool.device,
)
def get_size_per_token(self):
self.head_num = self.device_pool.head_num
self.head_dim = self.device_pool.head_dim
@@ -1092,24 +1140,42 @@ class AsymmetricMHATokenToKVPoolHost(MHATokenToKVPoolHost):
f"Unsupported layout for models with head_dim != v_head_dim "
f"and io_backend='kernel': {self.layout}; expected 'page_first'."
)
transfer_kv_all_layer_mla_lf_pf(
src_layers=device_pool.k_data_ptrs,
dst=self.k_buffer,
src_indices=device_indices,
dst_indices=host_indices,
item_size=self._k_token_stride_size(),
dst_layout_dim=self._k_layout_dim(),
num_layers=self.layer_num,
)
transfer_kv_all_layer_mla_lf_pf(
src_layers=device_pool.v_data_ptrs,
dst=self.v_buffer,
src_indices=device_indices,
dst_indices=host_indices,
item_size=self._v_token_stride_size(),
dst_layout_dim=self._v_layout_dim(),
num_layers=self.layer_num,
)
if self.can_use_write_back_jit:
jit_transfer_hicache_all_layer_mla_staged_lf_pf(
ptr_src=device_pool.k_data_ptrs,
src_indices=device_indices,
dst_indices=host_indices,
staging=self.staging_k_buffer,
dst=self.k_buffer,
page_size=self.page_size,
)
jit_transfer_hicache_all_layer_mla_staged_lf_pf(
ptr_src=device_pool.v_data_ptrs,
src_indices=device_indices,
dst_indices=host_indices,
staging=self.staging_v_buffer,
dst=self.v_buffer,
page_size=self.page_size,
)
else:
transfer_kv_all_layer_mla_lf_pf(
src_layers=device_pool.k_data_ptrs,
dst=self.k_buffer,
src_indices=device_indices,
dst_indices=host_indices,
item_size=self._k_token_stride_size(),
dst_layout_dim=self._k_layout_dim(),
num_layers=self.layer_num,
)
transfer_kv_all_layer_mla_lf_pf(
src_layers=device_pool.v_data_ptrs,
dst=self.v_buffer,
src_indices=device_indices,
dst_indices=host_indices,
item_size=self._v_token_stride_size(),
dst_layout_dim=self._v_layout_dim(),
num_layers=self.layer_num,
)
elif io_backend == "direct":
if self.layout != "page_first_direct":
raise ValueError(
python/sglang/test/kl_test_utils.py
+11 -1
View File
@@ -137,6 +137,7 @@ def _generate(
return_logprob=False,
logprob_start_len=-1,
temperature=0.0,
routed_dp_rank=None,
):
"""Send generate request and return results."""
json_data = {
@@ -155,11 +156,19 @@ def _generate(
"logprob_start_len": logprob_start_len,
}
)
if routed_dp_rank is not None:
json_data["routed_dp_rank"] = routed_dp_rank
response = requests.post(base_url + "/generate", json=json_data)
return response.json()
def _get_input_logprobs(base_url, new_input_ids, output_logprobs, temperature=0.0):
def _get_input_logprobs(
base_url,
new_input_ids,
output_logprobs,
temperature=0.0,
routed_dp_rank=None,
):
"""Run prefill to get input logprobs matching output logprobs."""
_flush_cache(base_url)
results = _generate(
@@ -169,6 +178,7 @@ def _get_input_logprobs(base_url, new_input_ids, output_logprobs, temperature=0.
return_logprob=True,
logprob_start_len=0,
temperature=temperature,
routed_dp_rank=routed_dp_rank,
)
assert len(results) == len(new_input_ids)
test/registered/jit/test_kvcacheio_asymmetric.py
+1
View File
@@ -55,6 +55,7 @@ def make_host_pool(dtype, layout="page_first"):
host.head_dim = K_HEAD_DIM
host.v_head_dim = V_HEAD_DIM
host.dtype = dtype
host.can_use_write_back_jit = False
if layout == "page_first":
k_dims = (TOTAL_ITEMS, NUM_LAYERS, HEAD_NUM, K_HEAD_DIM)
v_dims = (TOTAL_ITEMS, NUM_LAYERS, HEAD_NUM, V_HEAD_DIM)
test/registered/radix_cache/unified_radix_tree/test_unified_radix_cache_kl_mimo.py
+165
View File
@@ -0,0 +1,165 @@
"""MiMo V2.5 HiCache host load-back accuracy regression test."""
import random
import unittest
from sglang.srt.utils import kill_process_tree
from sglang.test.ci.ci_register import register_cuda_ci
from sglang.test.kl_test_utils import (
_extract_output_logprobs,
_flush_cache,
_generate,
_get_input_logprobs,
compare_kl_divergence,
)
from sglang.test.test_utils import (
DEFAULT_URL_FOR_TEST,
CustomTestCase,
popen_launch_server,
)
MIMO_MODEL = "XiaomiMiMo/MiMo-V2.5"
MIMO_LAUNCH_TIMEOUT = 3600
# MiMo V2.5 is pre-cached on the eight-H200 runner. The H200-only nightly suite
# exercises the asymmetric MHA host pool end to end without adding PR CI cost.
register_cuda_ci(est_time=1200, suite="nightly-8-gpu-h200", nightly=True)
class TestUnifiedMiMoHiCacheLoadBackKL(CustomTestCase):
"""Verify KL accuracy after asymmetric MHA KV is evicted to and loaded from L2."""
page_size = 64
prompt_len = 1024
max_total_tokens = 4096
# EAGLE verification and full-prefill replay use different execution paths;
# allow their expected numerical drift while still catching regressions.
kl_threshold = 0.01
# DP ranks own independent radix trees, so cache pressure and load-back must
# target the same rank instead of following the round-robin default.
routed_dp_rank = 0
@classmethod
def setUpClass(cls):
cls.model = MIMO_MODEL
cls.base_url = DEFAULT_URL_FOR_TEST
cls.process = popen_launch_server(
cls.model,
cls.base_url,
timeout=MIMO_LAUNCH_TIMEOUT,
other_args=[
"--trust-remote-code",
"--cuda-graph-max-bs-decode",
"64",
"--enable-multimodal",
"--tp",
"8",
"--dp",
"2",
"--enable-dp-attention",
"--mm-enable-dp-encoder",
"--attention-backend",
"fa3",
"--mm-attention-backend",
"fa3",
"--mem-fraction-static",
"0.65",
"--speculative-algorithm",
"EAGLE",
"--speculative-num-steps",
"3",
"--speculative-eagle-topk",
"1",
"--speculative-num-draft-tokens",
"4",
"--enable-multi-layer-eagle",
"--reasoning-parser",
"mimo",
"--page-size",
str(cls.page_size),
"--max-total-tokens",
str(cls.max_total_tokens),
"--enable-hierarchical-cache",
"--hicache-ratio",
"1.2",
"--hicache-write-policy",
"write_through",
"--hicache-io-backend",
"kernel",
"--hicache-mem-layout",
"page_first",
],
env={
"SGLANG_ENABLE_UNIFIED_RADIX_TREE": "1",
"SGLANG_USE_CUDA_IPC_TRANSPORT": "1",
},
)
@classmethod
def tearDownClass(cls):
kill_process_tree(cls.process.pid)
@classmethod
def _prompt(cls, seed: int) -> list[int]:
rng = random.Random(seed)
return [rng.randint(1, 30000) for _ in range(cls.prompt_len)]
def _generate_one(self, input_ids, max_new_tokens, return_logprob=False):
results = _generate(
self.base_url,
[input_ids],
max_new_tokens=max_new_tokens,
return_logprob=return_logprob,
temperature=0,
routed_dp_rank=self.routed_dp_rank,
)
self.assertEqual(len(results), 1)
return results[0]
def test_host_load_back_logprobs_match_prefill_replay(self):
"""Force L2 eviction, then compare load-back output logprobs with replay."""
base_prompt = self._prompt(1)
pressure_prompts = [self._prompt(seed) for seed in range(2, 6)]
_flush_cache(self.base_url)
self._generate_one(base_prompt, max_new_tokens=1)
# Four unique page-aligned prefixes fill the 4096-token L1 cache and
# evict the oldest prefix (base_prompt) to the HiCache host tier.
for prompt in pressure_prompts:
self._generate_one(prompt, max_new_tokens=1)
load_back = self._generate_one(
base_prompt, max_new_tokens=8, return_logprob=True
)
meta_info = load_back["meta_info"]
cached_details = meta_info.get("cached_tokens_details") or {}
host_cached_tokens = int(cached_details.get("host", 0))
self.assertGreater(
host_cached_tokens,
0,
"Expected the original prefix to be restored from the HiCache host tier; "
f"got cached_tokens={meta_info.get('cached_tokens')}, "
f"cached_tokens_details={cached_details}",
)
output_logprobs = [_extract_output_logprobs(load_back)]
replay_input_ids = [base_prompt + load_back["output_ids"]]
input_logprobs = _get_input_logprobs(
self.base_url,
replay_input_ids,
output_logprobs,
temperature=0,
routed_dp_rank=self.routed_dp_rank,
)
compare_kl_divergence(
input_logprobs,
output_logprobs,
{self.model: {"kl_div": self.kl_threshold}},
self.model,
"hicache_host_load_back",
)
if __name__ == "__main__":
unittest.main()
test/registered/unit/mem_cache/test_asymmetric_mha_pool_host_unit.py
+147
View File
@@ -50,6 +50,7 @@ def _make_host(layout: str) -> AsymmetricMHATokenToKVPoolHost:
raise ValueError(f"Unsupported test layout: {layout}")
host.kv_buffer = (torch.empty(k_dims), torch.empty(v_dims))
host.can_use_write_back_jit = False
return host
@@ -79,6 +80,42 @@ class TestAsymmetricMHATokenToKVPoolHost(CustomTestCase):
get_mha_host_pool_cls(asymmetric_pool), AsymmetricMHATokenToKVPoolHost
)
def test_staged_write_back_jit_uses_separate_kv_buffers(self):
host = _make_host("page_first")
host.page_num = 4
host.v_head_dim = 8
host.device_pool = SimpleNamespace(device="cuda")
cpu_empty = torch.empty
def _cpu_empty(shape, *, dtype, device):
return cpu_empty(shape, dtype=dtype)
with (
mock.patch("sglang.srt.mem_cache.pool_host.mha._is_cuda", True),
mock.patch("sglang.srt.mem_cache.pool_host.mha._is_hip", False),
mock.patch("sglang.srt.mem_cache.pool_host.mha._is_npu", False),
mock.patch("sglang.srt.mem_cache.pool_host.mha._is_xpu", False),
mock.patch("sglang.srt.mem_cache.pool_host.mha._is_mps", False),
mock.patch(
"sglang.srt.mem_cache.pool_host.mha.can_use_write_back_jit_kernel",
return_value=True,
) as can_use,
mock.patch(
"sglang.srt.mem_cache.pool_host.mha.torch.empty",
side_effect=_cpu_empty,
),
):
host._init_write_back_staging_buffers()
self.assertTrue(host.can_use_write_back_jit)
self.assertEqual(host.staging_page_capacity, 4)
self.assertEqual(host.staging_token_capacity, 8)
self.assertEqual(host.staging_k_buffer.shape, (8, 3, 2, 4))
self.assertEqual(host.staging_v_buffer.shape, (8, 3, 2, 8))
self.assertEqual(
[call.kwargs["element_size"] for call in can_use.call_args_list], [16, 32]
)
def test_kernel_load_splits_k_and_v_with_separate_strides(self):
# Dispatch-only test: the CUDA kernel is mocked; this verifies that K and
# V are sent as separate single-buffer calls with their own byte strides.
@@ -137,6 +174,116 @@ class TestAsymmetricMHATokenToKVPoolHost(CustomTestCase):
self.assertEqual(v_call.kwargs["item_size"], 24)
self.assertEqual(v_call.kwargs["dst_layout_dim"], 72)
def test_kernel_backup_uses_staged_kernel_for_each_kv_buffer(self):
host = _make_host("page_first")
host.can_use_write_back_jit = True
host.staging_k_buffer = torch.empty(4, 3, 2, 4)
host.staging_v_buffer = torch.empty(4, 3, 2, 6)
device_pool = _make_device_pool(host)
host_indices = torch.tensor([0, 1, 2, 3], dtype=torch.int64)
device_indices = torch.tensor([4, 5, 6, 7], dtype=torch.int64)
with (
mock.patch(
"sglang.srt.mem_cache.pool_host.mha.jit_transfer_hicache_all_layer_mla_staged_lf_pf"
) as staged,
mock.patch(
"sglang.srt.mem_cache.pool_host.mha.transfer_kv_all_layer_mla_lf_pf",
create=True,
) as fallback,
):
host.backup_from_device_all_layer(
device_pool, host_indices, device_indices, io_backend="kernel"
)
self.assertEqual(staged.call_count, 2)
self.assertEqual(fallback.call_count, 0)
k_call, v_call = staged.call_args_list
self.assertIs(k_call.kwargs["ptr_src"], device_pool.k_data_ptrs)
self.assertIs(k_call.kwargs["staging"], host.staging_k_buffer)
self.assertIs(k_call.kwargs["dst"], host.k_buffer)
self.assertIs(v_call.kwargs["ptr_src"], device_pool.v_data_ptrs)
self.assertIs(v_call.kwargs["staging"], host.staging_v_buffer)
self.assertIs(v_call.kwargs["dst"], host.v_buffer)
def test_staged_kernel_backup_load_roundtrip_preserves_asymmetric_kv(self):
"""Staged write-back must preserve both K and V values across a round trip."""
host = _make_host("page_first")
host.can_use_write_back_jit = True
host.staging_k_buffer = torch.empty(4, 3, 2, 4)
host.staging_v_buffer = torch.empty(4, 3, 2, 6)
device_pool = _make_device_pool(host)
host_indices = torch.tensor([0, 1, 2, 3], dtype=torch.int64)
device_indices = torch.tensor([4, 5, 6, 7], dtype=torch.int64)
for layer_id in range(host.layer_num):
device_pool.k_buffer[layer_id].copy_(
torch.arange(device_pool.k_buffer[layer_id].numel()).reshape_as(
device_pool.k_buffer[layer_id]
)
+ layer_id * 1000
)
device_pool.v_buffer[layer_id].copy_(
torch.arange(device_pool.v_buffer[layer_id].numel()).reshape_as(
device_pool.v_buffer[layer_id]
)
+ layer_id * 10000
)
expected_k = [buffer[device_indices].clone() for buffer in device_pool.k_buffer]
expected_v = [buffer[device_indices].clone() for buffer in device_pool.v_buffer]
buffers_by_ptrs = {
tuple(device_pool.k_data_ptrs.tolist()): device_pool.k_buffer,
tuple(device_pool.v_data_ptrs.tolist()): device_pool.v_buffer,
}
def staged_copy(*, ptr_src, src_indices, dst_indices, dst, **_):
src_buffers = buffers_by_ptrs[tuple(ptr_src.tolist())]
for layer_id, src in enumerate(src_buffers):
dst[dst_indices, layer_id] = src[src_indices]
def per_layer_copy(*, src, dst, src_indices, dst_indices, layer_id, **_):
dst[dst_indices] = src[src_indices, layer_id]
with (
mock.patch(
"sglang.srt.mem_cache.pool_host.mha.jit_transfer_hicache_all_layer_mla_staged_lf_pf",
side_effect=staged_copy,
),
mock.patch(
"sglang.srt.mem_cache.pool_host.mha.transfer_kv_per_layer_mla_pf_lf",
side_effect=per_layer_copy,
create=True,
),
):
host.backup_from_device_all_layer(
device_pool, host_indices, device_indices, io_backend="kernel"
)
for buffer in device_pool.k_buffer + device_pool.v_buffer:
buffer.zero_()
for layer_id in range(host.layer_num):
host.load_to_device_per_layer(
device_pool,
host_indices,
device_indices,
layer_id=layer_id,
io_backend="kernel",
)
for layer_id in range(host.layer_num):
torch.testing.assert_close(
host.k_buffer[host_indices, layer_id], expected_k[layer_id]
)
torch.testing.assert_close(
host.v_buffer[host_indices, layer_id], expected_v[layer_id]
)
torch.testing.assert_close(
device_pool.k_buffer[layer_id][device_indices], expected_k[layer_id]
)
torch.testing.assert_close(
device_pool.v_buffer[layer_id][device_indices], expected_v[layer_id]
)
def test_direct_load_splits_k_and_v_for_page_first_direct(self):
# Direct kernels derive copy size from each call's first tensor, so K/V
# must be dispatched separately when their head dims differ.