[hisparse]: Adding ci for hisparse kvcache-swap-in jit-kernel (#22155)
This commit is contained in:
Zhangheng
@@ -0,0 +1,299 @@
|
||||
import sys
|
||||
|
||||
import pytest
|
||||
import torch
|
||||
|
||||
from sglang.jit_kernel.hisparse import load_cache_to_device_buffer_mla
|
||||
from sglang.srt.utils import is_cuda, is_hip, is_npu, is_xpu
|
||||
from sglang.test.ci.ci_register import register_cuda_ci
|
||||
|
||||
register_cuda_ci(est_time=10, suite="stage-b-kernel-unit-1-gpu-large")
|
||||
register_cuda_ci(est_time=120, suite="nightly-kernel-1-gpu", nightly=True)
|
||||
|
||||
pytestmark = pytest.mark.skipif(
|
||||
not torch.cuda.is_available()
|
||||
or is_npu()
|
||||
or is_xpu()
|
||||
or not (is_cuda() or is_hip()),
|
||||
reason="HiSparse JIT tests require CUDA/ROCm.",
|
||||
)
|
||||
|
||||
DEVICE = "cuda"
|
||||
DTYPE = torch.float32
|
||||
KV_DIM = 8
|
||||
HOT_BUFFER_SIZE = 4
|
||||
PADDED_BUFFER_SIZE = HOT_BUFFER_SIZE + 1
|
||||
HOST_CACHE_SIZE = 16
|
||||
DEVICE_CACHE_SIZE = 16
|
||||
ITEM_SIZE_BYTES = KV_DIM * torch.empty((), dtype=DTYPE).element_size()
|
||||
|
||||
|
||||
def _host_cache() -> torch.Tensor:
|
||||
host_cache = torch.empty(
|
||||
(HOST_CACHE_SIZE, 1, KV_DIM), dtype=DTYPE, device="cpu", pin_memory=True
|
||||
)
|
||||
host_cache.copy_(torch.arange(host_cache.numel(), dtype=DTYPE).view_as(host_cache))
|
||||
return host_cache
|
||||
|
||||
|
||||
def _run_kernel(
|
||||
*,
|
||||
top_k_tokens: torch.Tensor,
|
||||
device_buffer_tokens: torch.Tensor,
|
||||
host_cache_locs: torch.Tensor,
|
||||
device_buffer_locs: torch.Tensor,
|
||||
host_cache: torch.Tensor,
|
||||
device_buffer: torch.Tensor,
|
||||
lru_slots: torch.Tensor,
|
||||
seq_len: int | None = None,
|
||||
seq_lens: torch.Tensor | None = None,
|
||||
seq_lens_dtype: torch.dtype = torch.int32,
|
||||
req_pool_indices: torch.Tensor | None = None,
|
||||
num_real_reqs: int | None = None,
|
||||
) -> torch.Tensor:
|
||||
batch_size = top_k_tokens.shape[0]
|
||||
if req_pool_indices is None:
|
||||
req_pool_indices = torch.arange(batch_size, dtype=torch.int64, device=DEVICE)
|
||||
if seq_lens is None:
|
||||
seq_lens = torch.full(
|
||||
(batch_size,), seq_len, dtype=seq_lens_dtype, device=DEVICE
|
||||
)
|
||||
if num_real_reqs is None:
|
||||
num_real_reqs = batch_size
|
||||
|
||||
out = torch.full_like(top_k_tokens, -1)
|
||||
load_cache_to_device_buffer_mla(
|
||||
top_k_tokens=top_k_tokens,
|
||||
device_buffer_tokens=device_buffer_tokens,
|
||||
host_cache_locs=host_cache_locs,
|
||||
device_buffer_locs=device_buffer_locs,
|
||||
host_cache=host_cache,
|
||||
device_buffer=device_buffer,
|
||||
top_k_device_locs=out,
|
||||
req_pool_indices=req_pool_indices,
|
||||
seq_lens=seq_lens,
|
||||
lru_slots=lru_slots,
|
||||
item_size_bytes=ITEM_SIZE_BYTES,
|
||||
num_top_k=top_k_tokens.shape[1],
|
||||
hot_buffer_size=HOT_BUFFER_SIZE,
|
||||
page_size=1,
|
||||
block_size=256,
|
||||
num_real_reqs=torch.tensor([num_real_reqs], dtype=torch.int32, device=DEVICE),
|
||||
)
|
||||
torch.cuda.synchronize()
|
||||
return out
|
||||
|
||||
|
||||
def _make_state(
|
||||
device_buffer_locs_rows: list[list[int]],
|
||||
device_buffer_tokens_rows: list[list[int]],
|
||||
newest_tokens: list[int],
|
||||
):
|
||||
host_cache = _host_cache()
|
||||
device_buffer = torch.full(
|
||||
(DEVICE_CACHE_SIZE, 1, KV_DIM), -1, dtype=DTYPE, device=DEVICE
|
||||
)
|
||||
device_buffer_locs = torch.tensor(
|
||||
device_buffer_locs_rows, dtype=torch.int32, device=DEVICE
|
||||
)
|
||||
device_buffer_tokens = torch.tensor(
|
||||
device_buffer_tokens_rows, dtype=torch.int32, device=DEVICE
|
||||
)
|
||||
lru_slots = (
|
||||
torch.arange(HOT_BUFFER_SIZE, dtype=torch.int16, device=DEVICE)
|
||||
.view(1, -1)
|
||||
.repeat(device_buffer_locs.shape[0], 1)
|
||||
)
|
||||
host_cache_locs = (
|
||||
torch.arange(HOST_CACHE_SIZE, dtype=torch.int64, device=DEVICE)
|
||||
.view(1, -1)
|
||||
.repeat(device_buffer_locs.shape[0], 1)
|
||||
)
|
||||
|
||||
# Slots 0..3 participate in LRU; slot 4 is the reserved newest slot.
|
||||
for rid, newest_token in enumerate(newest_tokens):
|
||||
for slot, token in enumerate(device_buffer_tokens_rows[rid][:HOT_BUFFER_SIZE]):
|
||||
if token >= 0:
|
||||
device_buffer[device_buffer_locs[rid, slot]].copy_(
|
||||
host_cache[token].to(DEVICE, non_blocking=True)
|
||||
)
|
||||
device_buffer[device_buffer_locs[rid, HOT_BUFFER_SIZE]].copy_(
|
||||
host_cache[newest_token].to(DEVICE, non_blocking=True)
|
||||
)
|
||||
torch.cuda.synchronize()
|
||||
|
||||
return {
|
||||
"host_cache": host_cache,
|
||||
"device_buffer": device_buffer,
|
||||
"device_buffer_locs": device_buffer_locs,
|
||||
"device_buffer_tokens": device_buffer_tokens,
|
||||
"lru_slots": lru_slots,
|
||||
"host_cache_locs": host_cache_locs,
|
||||
}
|
||||
|
||||
|
||||
def _long_case():
|
||||
# One-request baseline used by the stateful cases below:
|
||||
# req 0 LRU slots : [0, 1, 2, 3]
|
||||
# req 0 cached tokens : slot0->1, slot1->4, slot2->2, slot3->5
|
||||
# req 0 physical locs : slot0->9, slot1->7, slot2->3, slot3->5
|
||||
# req 0 newest slot : slot4/newest -> token 7 at physical loc 11
|
||||
return _make_state([[9, 7, 3, 5, 11]], [[1, 4, 2, 5, -1]], [7])
|
||||
|
||||
|
||||
@pytest.mark.parametrize("seq_lens_dtype", [torch.int32, torch.int64])
|
||||
def test_load_cache_to_device_buffer_fast_path(seq_lens_dtype: torch.dtype) -> None:
|
||||
host_cache = _host_cache()
|
||||
device_buffer = torch.arange(
|
||||
DEVICE_CACHE_SIZE * KV_DIM, dtype=DTYPE, device=DEVICE
|
||||
).view(DEVICE_CACHE_SIZE, 1, KV_DIM)
|
||||
device_buffer_before = device_buffer.clone()
|
||||
device_buffer_locs = torch.tensor(
|
||||
[[13, 9, 5, 1, 15]], dtype=torch.int32, device=DEVICE
|
||||
)
|
||||
device_buffer_tokens = torch.tensor(
|
||||
[[10, 11, 12, 13, -1]], dtype=torch.int32, device=DEVICE
|
||||
)
|
||||
device_buffer_tokens_before = device_buffer_tokens.clone()
|
||||
lru_slots = torch.tensor([[0, 1, 2, 3]], dtype=torch.int16, device=DEVICE)
|
||||
lru_slots_before = lru_slots.clone()
|
||||
|
||||
# Short-sequence layout:
|
||||
# token position 0 -> physical loc 13
|
||||
# token position 1 -> physical loc 9
|
||||
# token position 2 -> physical loc 5
|
||||
#
|
||||
# seq_len <= HOT_BUFFER_SIZE should skip host loads and LRU mutations,
|
||||
# so top_k_tokens acts like direct indexing into device_buffer_locs.
|
||||
out = _run_kernel(
|
||||
top_k_tokens=torch.tensor([[2, 0, 1]], dtype=torch.int32, device=DEVICE),
|
||||
device_buffer_tokens=device_buffer_tokens,
|
||||
host_cache_locs=torch.arange(
|
||||
HOST_CACHE_SIZE, dtype=torch.int64, device=DEVICE
|
||||
).view(1, -1),
|
||||
device_buffer_locs=device_buffer_locs,
|
||||
host_cache=host_cache,
|
||||
device_buffer=device_buffer,
|
||||
lru_slots=lru_slots,
|
||||
seq_len=3,
|
||||
seq_lens_dtype=seq_lens_dtype,
|
||||
)
|
||||
|
||||
assert torch.equal(out.cpu(), torch.tensor([[5, 13, 9]], dtype=torch.int32))
|
||||
assert torch.equal(device_buffer_tokens.cpu(), device_buffer_tokens_before.cpu())
|
||||
assert torch.equal(lru_slots.cpu(), lru_slots_before.cpu())
|
||||
assert torch.equal(device_buffer.cpu(), device_buffer_before.cpu())
|
||||
|
||||
|
||||
def test_load_cache_to_device_buffer_hits_newest_and_updates_lru() -> None:
|
||||
state = _long_case()
|
||||
|
||||
# Query [4, 2, 7]:
|
||||
# 4 hits slot1 -> loc 7
|
||||
# 2 hits slot2 -> loc 3
|
||||
# 7 is the newest token -> reserved newest loc 11
|
||||
#
|
||||
# Hits move to the MRU tail, so [0, 1, 2, 3] becomes [0, 3, 1, 2].
|
||||
out = _run_kernel(
|
||||
top_k_tokens=torch.tensor([[4, 2, 7]], dtype=torch.int32, device=DEVICE),
|
||||
seq_len=8,
|
||||
**state,
|
||||
)
|
||||
|
||||
assert torch.equal(out.cpu(), torch.tensor([[7, 3, 11]], dtype=torch.int32))
|
||||
assert torch.equal(
|
||||
state["device_buffer_tokens"].cpu(),
|
||||
torch.tensor([[1, 4, 2, 5, -1]], dtype=torch.int32),
|
||||
)
|
||||
assert torch.equal(
|
||||
state["lru_slots"].cpu(), torch.tensor([[0, 3, 1, 2]], dtype=torch.int16)
|
||||
)
|
||||
|
||||
|
||||
def test_load_cache_to_device_buffer_miss_uses_updated_lru_slot() -> None:
|
||||
state = _long_case()
|
||||
|
||||
# Step 1: touch tokens [4, 2], so LRU becomes [0, 3, 1, 2].
|
||||
# Step 2: query token 6, which is a miss.
|
||||
# The kernel should reuse the new LRU head slot0, whose physical loc is 9.
|
||||
# This round has no regular hits, so the freshly loaded miss slot ends up at the tail.
|
||||
_run_kernel(
|
||||
top_k_tokens=torch.tensor([[4, 2]], dtype=torch.int32, device=DEVICE),
|
||||
seq_len=8,
|
||||
**state,
|
||||
)
|
||||
out = _run_kernel(
|
||||
top_k_tokens=torch.tensor([[6]], dtype=torch.int32, device=DEVICE),
|
||||
seq_len=8,
|
||||
**state,
|
||||
)
|
||||
|
||||
assert torch.equal(out.cpu(), torch.tensor([[9]], dtype=torch.int32))
|
||||
assert torch.equal(
|
||||
state["device_buffer_tokens"].cpu(),
|
||||
torch.tensor([[6, 4, 2, 5, -1]], dtype=torch.int32),
|
||||
)
|
||||
assert torch.equal(
|
||||
state["lru_slots"].cpu(), torch.tensor([[3, 1, 2, 0]], dtype=torch.int16)
|
||||
)
|
||||
assert torch.equal(state["device_buffer"][9].cpu(), state["host_cache"][6])
|
||||
|
||||
|
||||
def test_load_cache_to_device_buffer_batched_with_padding() -> None:
|
||||
state = _make_state(
|
||||
[
|
||||
[9, 7, 3, 5, 11],
|
||||
[12, 10, 8, 6, 14],
|
||||
[15, 4, 2, 1, 13],
|
||||
],
|
||||
[
|
||||
[1, 4, 2, 5, -1],
|
||||
[0, 1, 2, 3, -1],
|
||||
[9, 8, 7, 6, -1],
|
||||
],
|
||||
[7, 4, 5],
|
||||
)
|
||||
padded_tokens_before = state["device_buffer_tokens"][2].clone()
|
||||
padded_lru_before = state["lru_slots"][2].clone()
|
||||
|
||||
# req 0: long path
|
||||
# cached tokens/locs : 1@9, 4@7, 2@3, 5@5, newest 7@11
|
||||
# query [4, 6, 7] : hit loc 7, miss into slot0/loc 9, newest loc 11
|
||||
# LRU update : remaining evictables [2, 3], then miss [0], then hit [1]
|
||||
# : [0, 1, 2, 3] -> [2, 3, 0, 1]
|
||||
#
|
||||
# req 1: fast path
|
||||
# seq_len = 3 <= HOT_BUFFER_SIZE, so [2, 1, 0] maps directly to locs [8, 10, 12]
|
||||
#
|
||||
# req 2: padded block
|
||||
# num_real_reqs = 2 means this row must be ignored entirely.
|
||||
out = _run_kernel(
|
||||
top_k_tokens=torch.tensor(
|
||||
[[4, 6, 7], [2, 1, 0], [9, 8, 7]], dtype=torch.int32, device=DEVICE
|
||||
),
|
||||
seq_lens=torch.tensor([8, 3, 8], dtype=torch.int32, device=DEVICE),
|
||||
num_real_reqs=2,
|
||||
**state,
|
||||
)
|
||||
|
||||
assert torch.equal(
|
||||
out.cpu(),
|
||||
torch.tensor([[7, 9, 11], [8, 10, 12], [-1, -1, -1]], dtype=torch.int32),
|
||||
)
|
||||
assert torch.equal(
|
||||
state["device_buffer_tokens"][:2].cpu(),
|
||||
torch.tensor([[6, 4, 2, 5, -1], [0, 1, 2, 3, -1]], dtype=torch.int32),
|
||||
)
|
||||
assert torch.equal(
|
||||
state["lru_slots"][:2].cpu(),
|
||||
torch.tensor([[2, 3, 0, 1], [0, 1, 2, 3]], dtype=torch.int16),
|
||||
)
|
||||
assert torch.equal(
|
||||
state["device_buffer_tokens"][2].cpu(), padded_tokens_before.cpu()
|
||||
)
|
||||
assert torch.equal(state["lru_slots"][2].cpu(), padded_lru_before.cpu())
|
||||
assert torch.equal(state["device_buffer"][9].cpu(), state["host_cache"][6])
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
sys.exit(pytest.main([__file__, "-v", "-s"]))
|
||||
Reference in New Issue
Block a user