diff --git a/python/sglang/jit_kernel/tests/test_hisparse.py b/python/sglang/jit_kernel/tests/test_hisparse.py new file mode 100644 index 000000000..190f22355 --- /dev/null +++ b/python/sglang/jit_kernel/tests/test_hisparse.py @@ -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"]))