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

[HiCache & JIT Kernel] Refactoring HiCache Write-Back Kernel (#21631)

Browse Source
This commit is contained in:
huangtingwei
2026-06-15 19:44:27 -07:00
committed by GitHub
parent a4a8a614b1
commit b5bcd76a41
12 changed files with 1207 additions and 81 deletions
Download Patch File Download Diff File Expand all files Collapse all files
test/registered/jit/test_hicache.py
+189 -4
View File
@@ -31,17 +31,21 @@ POOL_SIZE = PAGE_SIZE * 8
MHA_ELEMENT_DIMS = [128, 256, 512, 1024]
MLA_ELEMENT_DIMS = [576]
LAYOUTS = ["layer_first", "page_first"]
STAGED_WRITE_BACK_PAGE_COUNTS = [1, 63, 64, 65, 67, 128, 129]
def _token_indices_for_pages(
pages: torch.Tensor, page_size: int = PAGE_SIZE, device: str = DEVICE
pages: torch.Tensor,
page_size: int = PAGE_SIZE,
device: str = DEVICE,
dtype: torch.dtype = torch.int64,
) -> torch.Tensor:
parts = [
torch.arange(
int(page) * page_size,
(int(page) + 1) * page_size,
device=device,
dtype=torch.int64,
dtype=dtype,
)
for page in pages.tolist()
]
@@ -71,6 +75,12 @@ def _copy_tensor_with_offset(tensor: torch.Tensor, offset: int) -> None:
tensor.copy_(data + offset)
def _assert_page_filled(tensor: torch.Tensor, page: int, value: float) -> None:
page_slice = tensor[page * PAGE_SIZE : (page + 1) * PAGE_SIZE]
expected = torch.full_like(page_slice, value)
assert torch.equal(page_slice.cpu(), expected.cpu())
def _run_transfer_roundtrip_mha(layout: str, element_dim: int) -> None:
device_pool = MHATokenToKVPool(
size=POOL_SIZE,
@@ -99,9 +109,14 @@ def _run_transfer_roundtrip_mha(layout: str, element_dim: int) -> None:
host_pages = torch.tensor([0, 1, 2], device=DEVICE, dtype=torch.int64)
device_indices = _token_indices_for_pages(device_pages)
host_indices = _token_indices_for_pages(host_pages)
host_indices_backup = (
_token_indices_for_pages(host_pages, device="cpu")
if layout == "page_first"
else host_indices
)
host_pool.backup_from_device_all_layer(
device_pool, host_indices, device_indices, "kernel"
device_pool, host_indices_backup, device_indices, "kernel"
)
torch.cuda.synchronize()
@@ -187,9 +202,14 @@ def _run_transfer_roundtrip_mla(layout: str, element_dim: int) -> None:
host_pages = torch.tensor([0, 1, 2], device=DEVICE, dtype=torch.int64)
device_indices = _token_indices_for_pages(device_pages)
host_indices = _token_indices_for_pages(host_pages)
host_indices_backup = (
_token_indices_for_pages(host_pages, device="cpu")
if layout == "page_first"
else host_indices
)
host_pool.backup_from_device_all_layer(
device_pool, host_indices, device_indices, "kernel"
device_pool, host_indices_backup, device_indices, "kernel"
)
torch.cuda.synchronize()
@@ -231,6 +251,153 @@ def _run_transfer_roundtrip_mla(layout: str, element_dim: int) -> None:
)
def _run_page_first_staged_write_back_mha(
layout: str, element_dim: int, page_count: int
) -> None:
pool_size = PAGE_SIZE * (page_count + 8)
device_pool = MHATokenToKVPool(
size=pool_size,
page_size=PAGE_SIZE,
head_num=element_dim // 128,
head_dim=128,
dtype=torch.bfloat16,
layer_num=NUM_LAYERS,
device=DEVICE,
enable_memory_saver=False,
)
host_pool = _pinned_host_pool(
MHATokenToKVPoolHost,
device_pool=device_pool,
layout=layout,
)
assert host_pool.can_use_jit
assert host_pool.staging_page_capacity > 0
if page_count > 64:
assert host_pool.staging_page_capacity < page_count
for layer_id in range(NUM_LAYERS):
_copy_tensor_with_offset(device_pool.k_buffer[layer_id], layer_id)
_copy_tensor_with_offset(device_pool.v_buffer[layer_id], layer_id + 100)
host_pool.k_buffer.fill_(-7)
host_pool.v_buffer.fill_(-11)
device_pages = torch.arange(
2,
2 + page_count,
device=DEVICE,
dtype=torch.int64,
)
host_pages = torch.arange(
page_count,
0,
-1,
dtype=torch.int64,
)
src_index_dtype = torch.int32 if page_count == 64 else torch.int64
device_indices = _token_indices_for_pages(device_pages, dtype=src_index_dtype)
host_indices = _token_indices_for_pages(host_pages, device="cpu")
assert not host_indices.is_cuda
host_pool.backup_from_device_all_layer(
device_pool, host_indices, device_indices, "kernel"
)
torch.cuda.synchronize()
for layer_id in range(NUM_LAYERS):
for host_page, device_page in zip(host_pages.tolist(), device_pages.tolist()):
host_start = host_page * PAGE_SIZE
device_start = device_page * PAGE_SIZE
assert torch.equal(
host_pool.k_data_refs[layer_id][
host_start : host_start + PAGE_SIZE
].cpu(),
device_pool.k_buffer[layer_id][
device_start : device_start + PAGE_SIZE
].cpu(),
)
assert torch.equal(
host_pool.v_data_refs[layer_id][
host_start : host_start + PAGE_SIZE
].cpu(),
device_pool.v_buffer[layer_id][
device_start : device_start + PAGE_SIZE
].cpu(),
)
for layer_id in range(NUM_LAYERS):
for untouched_page in [0, page_count + 1]:
_assert_page_filled(host_pool.k_data_refs[layer_id], untouched_page, -7)
_assert_page_filled(host_pool.v_data_refs[layer_id], untouched_page, -11)
def _run_page_first_staged_write_back_mla(
layout: str, element_dim: int, page_count: int
) -> None:
pool_size = PAGE_SIZE * (page_count + 8)
device_pool = MLATokenToKVPool(
size=pool_size,
page_size=PAGE_SIZE,
kv_lora_rank=element_dim - 64,
qk_rope_head_dim=64,
dtype=torch.bfloat16,
layer_num=NUM_LAYERS,
device=DEVICE,
enable_memory_saver=False,
)
host_pool = _pinned_host_pool(
MLATokenToKVPoolHost,
device_pool=device_pool,
layout=layout,
)
assert host_pool.can_use_jit
assert host_pool.staging_page_capacity > 0
if page_count > 64:
assert host_pool.staging_page_capacity < page_count
for layer_id in range(NUM_LAYERS):
_copy_tensor_with_offset(device_pool.kv_buffer[layer_id], layer_id)
host_pool.kv_buffer.fill_(-13)
device_pages = torch.arange(
2,
2 + page_count,
device=DEVICE,
dtype=torch.int64,
)
host_pages = torch.arange(
page_count,
0,
-1,
dtype=torch.int64,
)
src_index_dtype = torch.int32 if page_count == 64 else torch.int64
device_indices = _token_indices_for_pages(device_pages, dtype=src_index_dtype)
host_indices = _token_indices_for_pages(host_pages, device="cpu")
assert not host_indices.is_cuda
host_pool.backup_from_device_all_layer(
device_pool, host_indices, device_indices, "kernel"
)
torch.cuda.synchronize()
for layer_id in range(NUM_LAYERS):
for host_page, device_page in zip(host_pages.tolist(), device_pages.tolist()):
host_start = host_page * PAGE_SIZE
device_start = device_page * PAGE_SIZE
assert torch.equal(
host_pool.data_refs[layer_id][
host_start : host_start + PAGE_SIZE
].cpu(),
device_pool.kv_buffer[layer_id][
device_start : device_start + PAGE_SIZE
].cpu(),
)
for layer_id in range(NUM_LAYERS):
for untouched_page in [0, page_count + 1]:
_assert_page_filled(host_pool.data_refs[layer_id], untouched_page, -13)
@pytest.mark.parametrize("layout", LAYOUTS)
@pytest.mark.parametrize("element_dim", MHA_ELEMENT_DIMS)
def test_hicache_transfer_mha(layout: str, element_dim: int) -> None:
@@ -243,5 +410,23 @@ def test_hicache_transfer_mla(layout: str, element_dim: int) -> None:
_run_transfer_roundtrip_mla(layout, element_dim)
@pytest.mark.parametrize("layout", ["page_first"])
@pytest.mark.parametrize("element_dim", MHA_ELEMENT_DIMS)
@pytest.mark.parametrize("page_count", STAGED_WRITE_BACK_PAGE_COUNTS)
def test_hicache_page_first_staged_write_back_mha(
layout: str, element_dim: int, page_count: int
) -> None:
_run_page_first_staged_write_back_mha(layout, element_dim, page_count)
@pytest.mark.parametrize("layout", ["page_first"])
@pytest.mark.parametrize("element_dim", MLA_ELEMENT_DIMS)
@pytest.mark.parametrize("page_count", STAGED_WRITE_BACK_PAGE_COUNTS)
def test_hicache_page_first_staged_write_back_mla(
layout: str, element_dim: int, page_count: int
) -> None:
_run_page_first_staged_write_back_mla(layout, element_dim, page_count)
if __name__ == "__main__":
sys.exit(pytest.main([__file__, "-v", "-s"]))
test/registered/unit/mem_cache/test_hiradix_cache_unit.py
+1
View File
@@ -37,6 +37,7 @@ class TestHiRadixCacheKVEvents(CustomTestCase):
model_path="dummy",
page_size=PAGE_SIZE,
hicache_io_backend="direct",
hicache_mem_layout="layer_first",
hicache_write_policy="write_through",
)
set_global_server_args_for_scheduler(server_args)
test/registered/unit/server_args/test_server_args.py
+16 -8
View File
@@ -620,6 +620,14 @@ class TestHiCacheArgs(unittest.TestCase):
def test_hicache_io_backend_and_mem_layout_compatibility(self):
cases = [
{
"name": "default_kernel_page_first",
"overrides": {
"enable_hierarchical_cache": True,
},
"expected_io_backend": "kernel",
"expected_mem_layout": "page_first",
},
{
"name": "kernel_with_page_first_direct",
"overrides": {
@@ -660,8 +668,9 @@ class TestHiCacheArgs(unittest.TestCase):
"attention_backend": "triton",
"decode_attention_backend": "fa3",
},
"expected_io_backend": "direct",
"expected_mem_layout": "page_first_direct",
"expected_io_backend": "kernel",
"expected_mem_layout": "page_first",
"expected_decode_backend": "fa3",
},
]
@@ -673,13 +682,10 @@ class TestHiCacheArgs(unittest.TestCase):
args,
expected_io_backend=case["expected_io_backend"],
expected_mem_layout=case["expected_mem_layout"],
expected_decode_backend=case.get("expected_decode_backend"),
)
@patch.object(ServerArgs, "use_mla_backend", return_value=False)
@patch("sglang.srt.server_args.is_flashinfer_available", return_value=False)
def test_decode_attention_backend_with_implicit_fa3(
self, _mock_flashinfer, _mock_use_mla_backend
):
def test_hicache_kernel_keeps_implicit_fa3_decode_backend(self):
args = self._make_args(
enable_hierarchical_cache=True,
hicache_io_backend="kernel",
@@ -689,7 +695,9 @@ class TestHiCacheArgs(unittest.TestCase):
args._handle_hicache()
self.assertEqual(args.decode_attention_backend, "triton")
self.assertEqual(args.hicache_io_backend, "kernel")
self.assertEqual(args.hicache_mem_layout, "page_first")
self.assertIsNone(args.decode_attention_backend)
class TestNgramExternalSamArgs(CustomTestCase):