feat: Support HiCache for MiMo-V2 models (1/N) (#27378)
Co-authored-by: Zhangheng <hzh0425@apache.org> Co-authored-by: 晟海 <huangtingwei.htw@antgroup.com>
This commit is contained in:
Yingchun Lai
co-authored by
Zhangheng
晟海
Zhangheng
晟海
parent
60d4bd4c70
commit
806365e778
@@ -0,0 +1,153 @@
|
||||
import sys
|
||||
from types import SimpleNamespace
|
||||
|
||||
import pytest
|
||||
import torch
|
||||
|
||||
from sglang.srt.mem_cache.memory_pool_host import AsymmetricMHATokenToKVPoolHost
|
||||
from sglang.test.ci.ci_register import register_cuda_ci
|
||||
|
||||
register_cuda_ci(est_time=10, suite="base-b-kernel-unit-1-gpu-large")
|
||||
|
||||
# These tests use AsymmetricMHATokenToKVPoolHost methods and let that class call
|
||||
# the real sgl-kernel transfer ops. The asymmetric host pool is kernel-only;
|
||||
# direct/page_first_direct is intentionally rejected in the CPU dispatch tests.
|
||||
pytestmark = pytest.mark.skipif(
|
||||
not torch.cuda.is_available(), reason="asymmetric host-pool tests require CUDA."
|
||||
)
|
||||
|
||||
DEVICE = "cuda"
|
||||
PAGE_SIZE = 16
|
||||
NUM_LAYERS = 3
|
||||
TOTAL_ITEMS = PAGE_SIZE * 8
|
||||
HEAD_NUM = 4
|
||||
K_HEAD_DIM = 192
|
||||
V_HEAD_DIM = 128
|
||||
DTYPES = [torch.float16, torch.bfloat16]
|
||||
|
||||
|
||||
def token_indices_for_pages(pages, page_size=PAGE_SIZE, device=None):
|
||||
indices = torch.cat(
|
||||
[
|
||||
torch.arange(
|
||||
int(page) * page_size,
|
||||
(int(page) + 1) * page_size,
|
||||
dtype=torch.int64,
|
||||
)
|
||||
for page in pages.tolist()
|
||||
]
|
||||
)
|
||||
return indices if device is None else indices.to(device)
|
||||
|
||||
|
||||
def fill_with_offset(tensor, offset):
|
||||
data = torch.arange(tensor.numel(), device=tensor.device, dtype=tensor.dtype)
|
||||
tensor.copy_((data + offset).view_as(tensor))
|
||||
|
||||
|
||||
def make_host_pool(dtype):
|
||||
host = AsymmetricMHATokenToKVPoolHost.__new__(AsymmetricMHATokenToKVPoolHost)
|
||||
host.layout = "page_first"
|
||||
host.page_size = PAGE_SIZE
|
||||
host.layer_num = NUM_LAYERS
|
||||
host.head_num = HEAD_NUM
|
||||
host.head_dim = K_HEAD_DIM
|
||||
host.v_head_dim = V_HEAD_DIM
|
||||
host.dtype = dtype
|
||||
host.kv_buffer = (
|
||||
torch.zeros(
|
||||
TOTAL_ITEMS, NUM_LAYERS, HEAD_NUM, K_HEAD_DIM, dtype=dtype
|
||||
).pin_memory(),
|
||||
torch.zeros(
|
||||
TOTAL_ITEMS, NUM_LAYERS, HEAD_NUM, V_HEAD_DIM, dtype=dtype
|
||||
).pin_memory(),
|
||||
)
|
||||
return host
|
||||
|
||||
|
||||
def make_device_pool(dtype):
|
||||
k_buffer = [
|
||||
torch.empty(TOTAL_ITEMS, HEAD_NUM, K_HEAD_DIM, dtype=dtype, device=DEVICE)
|
||||
for _ in range(NUM_LAYERS)
|
||||
]
|
||||
v_buffer = [
|
||||
torch.empty(TOTAL_ITEMS, HEAD_NUM, V_HEAD_DIM, dtype=dtype, device=DEVICE)
|
||||
for _ in range(NUM_LAYERS)
|
||||
]
|
||||
for layer_id in range(NUM_LAYERS):
|
||||
fill_with_offset(k_buffer[layer_id], layer_id * 1000)
|
||||
fill_with_offset(v_buffer[layer_id], layer_id * 1000 + 100)
|
||||
|
||||
return SimpleNamespace(
|
||||
k_buffer=k_buffer,
|
||||
v_buffer=v_buffer,
|
||||
k_data_ptrs=torch.tensor(
|
||||
[x.data_ptr() for x in k_buffer], dtype=torch.uint64, device=DEVICE
|
||||
),
|
||||
v_data_ptrs=torch.tensor(
|
||||
[x.data_ptr() for x in v_buffer], dtype=torch.uint64, device=DEVICE
|
||||
),
|
||||
)
|
||||
|
||||
|
||||
def assert_backup_matches_device(host, device_pool, host_indices_host, device_indices):
|
||||
for layer_id in range(NUM_LAYERS):
|
||||
torch.testing.assert_close(
|
||||
host.k_buffer[host_indices_host, layer_id],
|
||||
device_pool.k_buffer[layer_id][device_indices].cpu(),
|
||||
)
|
||||
torch.testing.assert_close(
|
||||
host.v_buffer[host_indices_host, layer_id],
|
||||
device_pool.v_buffer[layer_id][device_indices].cpu(),
|
||||
)
|
||||
|
||||
|
||||
def assert_load_matches_host(host, device_pool, host_indices_host, load_indices):
|
||||
for layer_id in range(NUM_LAYERS):
|
||||
torch.testing.assert_close(
|
||||
device_pool.k_buffer[layer_id][load_indices],
|
||||
host.k_buffer[host_indices_host, layer_id].to(DEVICE),
|
||||
)
|
||||
torch.testing.assert_close(
|
||||
device_pool.v_buffer[layer_id][load_indices],
|
||||
host.v_buffer[host_indices_host, layer_id].to(DEVICE),
|
||||
)
|
||||
|
||||
|
||||
@pytest.mark.parametrize("dtype", DTYPES)
|
||||
def test_asymmetric_mha_kernel_page_first_roundtrip(dtype):
|
||||
# Covers D2H backup + H2D load through AsymmetricMHATokenToKVPoolHost using
|
||||
# MiMoV2's real K/V head dims and the real MLA single-buffer kernels.
|
||||
host = make_host_pool(dtype)
|
||||
device_pool = make_device_pool(dtype)
|
||||
|
||||
device_pages = torch.tensor([1, 2, 3], dtype=torch.int64)
|
||||
host_pages = torch.tensor([0, 1, 2], dtype=torch.int64)
|
||||
load_pages = torch.tensor([4, 5, 6], dtype=torch.int64)
|
||||
device_indices_host = token_indices_for_pages(device_pages)
|
||||
host_indices_host = token_indices_for_pages(host_pages)
|
||||
load_indices_host = token_indices_for_pages(load_pages)
|
||||
device_indices = device_indices_host.to(DEVICE)
|
||||
host_indices = host_indices_host.to(DEVICE)
|
||||
load_indices = load_indices_host.to(DEVICE)
|
||||
|
||||
host.backup_from_device_all_layer(
|
||||
device_pool, host_indices, device_indices, io_backend="kernel"
|
||||
)
|
||||
torch.cuda.synchronize()
|
||||
assert_backup_matches_device(
|
||||
host, device_pool, host_indices_host, device_indices_host
|
||||
)
|
||||
|
||||
for layer_id in range(NUM_LAYERS):
|
||||
device_pool.k_buffer[layer_id].zero_()
|
||||
device_pool.v_buffer[layer_id].zero_()
|
||||
host.load_to_device_per_layer(
|
||||
device_pool, host_indices, load_indices, layer_id, io_backend="kernel"
|
||||
)
|
||||
torch.cuda.synchronize()
|
||||
assert_load_matches_host(host, device_pool, host_indices_host, load_indices_host)
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
sys.exit(pytest.main([__file__, "-v", "-s"]))
|
||||
@@ -1,5 +1,6 @@
|
||||
import unittest
|
||||
|
||||
from sglang.srt.environ import envs
|
||||
from sglang.test.ci.ci_register import register_cuda_ci
|
||||
from sglang.test.kits.eval_accuracy_kit import GSM8KMixin
|
||||
from sglang.test.server_fixtures.mmmu_fixture import MMMUServerBase
|
||||
@@ -20,6 +21,13 @@ MIMO_V2_OTHER_ARGS = [
|
||||
"fa3",
|
||||
"--reasoning-parser",
|
||||
"mimo",
|
||||
"--enable-hierarchical-cache",
|
||||
"--hicache-ratio",
|
||||
"1.5",
|
||||
"--hicache-mem-layout",
|
||||
"page_first",
|
||||
"--hicache-io-backend",
|
||||
"kernel",
|
||||
]
|
||||
MIMO_V2_MTP_OTHER_ARGS = MIMO_V2_OTHER_ARGS + [
|
||||
"--speculative-algorithm",
|
||||
@@ -42,6 +50,11 @@ class TestMiMoV2(GSM8KMixin, MMMUServerBase):
|
||||
server_api_key = None
|
||||
other_args = MIMO_V2_MTP_OTHER_ARGS
|
||||
|
||||
@classmethod
|
||||
def setUpClass(cls):
|
||||
with envs.SGLANG_ENABLE_UNIFIED_RADIX_TREE.override(True):
|
||||
super().setUpClass()
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
unittest.main()
|
||||
|
||||
@@ -1,5 +1,6 @@
|
||||
import unittest
|
||||
|
||||
from sglang.srt.environ import envs
|
||||
from sglang.test.ci.ci_register import register_cuda_ci
|
||||
from sglang.test.kits.eval_accuracy_kit import GSM8KMixin
|
||||
from sglang.test.kits.spec_decoding_kit import SpecDecodingMixin
|
||||
@@ -42,11 +43,23 @@ class TestMiMoV2Flash(GSM8KMixin, SpecDecodingMixin, DefaultServerBase):
|
||||
"--enable-multi-layer-eagle",
|
||||
"--model-loader-extra-config",
|
||||
'{"enable_multithread_load": true,"num_threads": 64}',
|
||||
"--enable-hierarchical-cache",
|
||||
"--hicache-ratio",
|
||||
"1.5",
|
||||
"--hicache-mem-layout",
|
||||
"page_first",
|
||||
"--hicache-io-backend",
|
||||
"kernel",
|
||||
]
|
||||
|
||||
bs_1_speed_thres = 170
|
||||
accept_length_thres = 3.2
|
||||
|
||||
@classmethod
|
||||
def setUpClass(cls):
|
||||
with envs.SGLANG_ENABLE_UNIFIED_RADIX_TREE.override(True):
|
||||
super().setUpClass()
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
unittest.main()
|
||||
|
||||
@@ -0,0 +1,156 @@
|
||||
"""Unit tests for asymmetric MHA host KV pool transfer dispatch."""
|
||||
|
||||
import unittest
|
||||
from types import SimpleNamespace
|
||||
from unittest import mock
|
||||
|
||||
import torch
|
||||
|
||||
from sglang.srt.mem_cache.memory_pool_host import (
|
||||
AsymmetricMHATokenToKVPoolHost,
|
||||
MHATokenToKVPoolHost,
|
||||
get_mha_host_pool_cls,
|
||||
)
|
||||
from sglang.test.ci.ci_register import register_cpu_ci
|
||||
from sglang.test.test_utils import CustomTestCase
|
||||
|
||||
register_cpu_ci(est_time=5, suite="base-a-test-cpu")
|
||||
|
||||
|
||||
def _make_host(layout: str) -> AsymmetricMHATokenToKVPoolHost:
|
||||
host = AsymmetricMHATokenToKVPoolHost.__new__(AsymmetricMHATokenToKVPoolHost)
|
||||
host.layout = layout
|
||||
host.page_size = 2
|
||||
host.layer_num = 3
|
||||
host.head_num = 2
|
||||
host.head_dim = 4
|
||||
host.v_head_dim = 6
|
||||
host.dtype = torch.float16
|
||||
|
||||
if layout == "page_first":
|
||||
k_dims = (8, host.layer_num, host.head_num, host.head_dim)
|
||||
v_dims = (8, host.layer_num, host.head_num, host.v_head_dim)
|
||||
else:
|
||||
raise ValueError(f"Unsupported test layout: {layout}")
|
||||
|
||||
host.kv_buffer = (torch.empty(k_dims), torch.empty(v_dims))
|
||||
return host
|
||||
|
||||
|
||||
def _make_device_pool(host: AsymmetricMHATokenToKVPoolHost) -> SimpleNamespace:
|
||||
size = 8
|
||||
k_buffer = [
|
||||
torch.empty(size, host.head_num, host.head_dim) for _ in range(host.layer_num)
|
||||
]
|
||||
v_buffer = [
|
||||
torch.empty(size, host.head_num, host.v_head_dim) for _ in range(host.layer_num)
|
||||
]
|
||||
return SimpleNamespace(
|
||||
k_buffer=k_buffer,
|
||||
v_buffer=v_buffer,
|
||||
k_data_ptrs=torch.tensor([x.data_ptr() for x in k_buffer], dtype=torch.uint64),
|
||||
v_data_ptrs=torch.tensor([x.data_ptr() for x in v_buffer], dtype=torch.uint64),
|
||||
)
|
||||
|
||||
|
||||
class TestAsymmetricMHATokenToKVPoolHost(CustomTestCase):
|
||||
def test_factory_selects_asymmetric_pool_for_mismatched_kv_dims(self):
|
||||
symmetric_pool = SimpleNamespace(head_dim=4, v_head_dim=4)
|
||||
asymmetric_pool = SimpleNamespace(head_dim=4, v_head_dim=6)
|
||||
|
||||
self.assertIs(get_mha_host_pool_cls(symmetric_pool), MHATokenToKVPoolHost)
|
||||
self.assertIs(
|
||||
get_mha_host_pool_cls(asymmetric_pool), AsymmetricMHATokenToKVPoolHost
|
||||
)
|
||||
|
||||
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.
|
||||
host = _make_host("page_first")
|
||||
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.memory_pool_host.transfer_kv_per_layer_mla_pf_lf",
|
||||
create=True,
|
||||
) as transfer:
|
||||
host.load_to_device_per_layer(
|
||||
device_pool,
|
||||
host_indices,
|
||||
device_indices,
|
||||
layer_id=1,
|
||||
io_backend="kernel",
|
||||
)
|
||||
|
||||
self.assertEqual(transfer.call_count, 2)
|
||||
k_call, v_call = transfer.call_args_list
|
||||
self.assertIs(k_call.kwargs["src"], host.k_buffer)
|
||||
self.assertIs(k_call.kwargs["dst"], device_pool.k_buffer[1])
|
||||
self.assertEqual(k_call.kwargs["item_size"], 16)
|
||||
self.assertEqual(k_call.kwargs["src_layout_dim"], 48)
|
||||
self.assertIs(v_call.kwargs["src"], host.v_buffer)
|
||||
self.assertIs(v_call.kwargs["dst"], device_pool.v_buffer[1])
|
||||
self.assertEqual(v_call.kwargs["item_size"], 24)
|
||||
self.assertEqual(v_call.kwargs["src_layout_dim"], 72)
|
||||
|
||||
def test_kernel_backup_splits_k_and_v_with_separate_strides(self):
|
||||
# Dispatch-only test: D2H backup must pass separate K/V layer pointer
|
||||
# tables so the single-buffer MLA kernel gets the correct stride per side.
|
||||
host = _make_host("page_first")
|
||||
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.memory_pool_host.transfer_kv_all_layer_mla_lf_pf",
|
||||
create=True,
|
||||
) as transfer:
|
||||
host.backup_from_device_all_layer(
|
||||
device_pool, host_indices, device_indices, io_backend="kernel"
|
||||
)
|
||||
|
||||
self.assertEqual(transfer.call_count, 2)
|
||||
k_call, v_call = transfer.call_args_list
|
||||
self.assertIs(k_call.kwargs["src_layers"], device_pool.k_data_ptrs)
|
||||
self.assertIs(k_call.kwargs["dst"], host.k_buffer)
|
||||
self.assertEqual(k_call.kwargs["item_size"], 16)
|
||||
self.assertEqual(k_call.kwargs["dst_layout_dim"], 48)
|
||||
self.assertIs(v_call.kwargs["src_layers"], device_pool.v_data_ptrs)
|
||||
self.assertIs(v_call.kwargs["dst"], host.v_buffer)
|
||||
self.assertEqual(v_call.kwargs["item_size"], 24)
|
||||
self.assertEqual(v_call.kwargs["dst_layout_dim"], 72)
|
||||
|
||||
def test_direct_load_is_rejected(self):
|
||||
# Direct single-buffer D2H is not reliable for asymmetric K/V in the
|
||||
# current sgl-kernel fast path, so the asymmetric host pool is kernel-only.
|
||||
host = _make_host("page_first")
|
||||
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 self.assertRaisesRegex(ValueError, "expected 'kernel'"):
|
||||
host.load_to_device_per_layer(
|
||||
device_pool,
|
||||
host_indices,
|
||||
device_indices,
|
||||
layer_id=2,
|
||||
io_backend="direct",
|
||||
)
|
||||
|
||||
def test_direct_backup_is_rejected(self):
|
||||
# Same restriction for D2H backup: asymmetric MHA uses the kernel path
|
||||
# until the direct kernel has an explicit safe asymmetric mode.
|
||||
host = _make_host("page_first")
|
||||
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 self.assertRaisesRegex(ValueError, "expected 'kernel'"):
|
||||
host.backup_from_device_all_layer(
|
||||
device_pool, host_indices, device_indices, io_backend="direct"
|
||||
)
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
unittest.main()
|
||||
@@ -450,16 +450,24 @@ class TestUnifiedRadixCacheKVEvents(CustomTestCase):
|
||||
def _init_hicache(self, tree, *, write_policy: str = "write_through"):
|
||||
import sglang.srt.mem_cache.hybrid_cache.hybrid_pool_assembler as assembler
|
||||
|
||||
orig_kv_host_pool = assembler.MHATokenToKVPoolHost
|
||||
# Wrap the host-pool factory (not MHATokenToKVPoolHost directly)
|
||||
# because the assembler picks between MHATokenToKVPoolHost and
|
||||
# AsymmetricMHATokenToKVPoolHost via get_mha_host_pool_cls(device_pool).
|
||||
orig_get_mha_host_pool_cls = assembler.get_mha_host_pool_cls
|
||||
|
||||
def kv_host_pool_wrapper(*args, **kwargs):
|
||||
kwargs["pin_memory"] = False
|
||||
return orig_kv_host_pool(*args, **kwargs)
|
||||
def get_mha_host_pool_cls_wrapper(device_pool):
|
||||
host_pool_cls = orig_get_mha_host_pool_cls(device_pool)
|
||||
|
||||
def kv_host_pool_wrapper(*args, **kwargs):
|
||||
kwargs["pin_memory"] = False
|
||||
return host_pool_cls(*args, **kwargs)
|
||||
|
||||
return kv_host_pool_wrapper
|
||||
|
||||
patcher = mock.patch.object(
|
||||
assembler,
|
||||
"MHATokenToKVPoolHost",
|
||||
side_effect=kv_host_pool_wrapper,
|
||||
"get_mha_host_pool_cls",
|
||||
side_effect=get_mha_host_pool_cls_wrapper,
|
||||
)
|
||||
patcher.start()
|
||||
self.addCleanup(patcher.stop)
|
||||
@@ -2370,12 +2378,20 @@ class UnifiedRadixCacheSuite:
|
||||
):
|
||||
import sglang.srt.mem_cache.hybrid_cache.hybrid_pool_assembler as assembler
|
||||
|
||||
orig_kv_host_pool = assembler.MHATokenToKVPoolHost
|
||||
# See _init_hicache: wrap the factory rather than MHATokenToKVPoolHost
|
||||
# directly so the pin_memory=False override applies to both
|
||||
# MHATokenToKVPoolHost and AsymmetricMHATokenToKVPoolHost.
|
||||
orig_get_mha_host_pool_cls = assembler.get_mha_host_pool_cls
|
||||
orig_mamba_host_pool = assembler.MambaPoolHost
|
||||
|
||||
def kv_host_pool_wrapper(*args, **kwargs):
|
||||
kwargs["pin_memory"] = False
|
||||
return orig_kv_host_pool(*args, **kwargs)
|
||||
def get_mha_host_pool_cls_wrapper(device_pool):
|
||||
host_pool_cls = orig_get_mha_host_pool_cls(device_pool)
|
||||
|
||||
def kv_host_pool_wrapper(*args, **kwargs):
|
||||
kwargs["pin_memory"] = False
|
||||
return host_pool_cls(*args, **kwargs)
|
||||
|
||||
return kv_host_pool_wrapper
|
||||
|
||||
def mamba_host_pool_wrapper(*args, **kwargs):
|
||||
kwargs["pin_memory"] = False
|
||||
@@ -2384,8 +2400,8 @@ class UnifiedRadixCacheSuite:
|
||||
patchers = [
|
||||
mock.patch.object(
|
||||
assembler,
|
||||
"MHATokenToKVPoolHost",
|
||||
side_effect=kv_host_pool_wrapper,
|
||||
"get_mha_host_pool_cls",
|
||||
side_effect=get_mha_host_pool_cls_wrapper,
|
||||
),
|
||||
mock.patch.object(
|
||||
assembler,
|
||||
|
||||
Reference in New Issue
Block a user