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

[HiSparse] Support hisparse multi-step swap io kernel (#32162)

Browse Source
This commit is contained in:
huangtingwei
2026-08-24 17:29:44 -07:00
committed by GitHub
parent 1ec20fd25d
commit 0f7ba3d115
5 changed files with 2182 additions and 5 deletions
Download Patch File Download Diff File Expand all files Collapse all files
python/sglang/kernels/jit/csrc/hisparse.cuh → python/sglang/kernels/jit/csrc/kvcacheio/hisparse.cuh
View File
python/sglang/kernels/jit/csrc/kvcacheio/hisparse_spec.cuh
+1113
View File
File diff suppressed because it is too large Load Diff
python/sglang/kernels/ops/kvcache/hisparse.py
+160 -5
View File
@@ -1,16 +1,171 @@
from __future__ import annotations
import functools
from typing import TYPE_CHECKING
from typing import TYPE_CHECKING, NamedTuple
import torch
from sglang.kernels.jit.utils import load_jit, make_cpp_args
from sglang.kernels.jit.utils import (
cache_once,
is_arch_support_pdl,
load_jit,
make_cpp_args,
)
if TYPE_CHECKING:
from tvm_ffi.module import Module
_GATHER_BLOCK_SIZE = 64
class HiSparseSpecState(NamedTuple):
"""Persistent cache state and reusable miss workspace for speculative swap.
``cache_index`` stores the two int64 hash banks as
``[num_requests, 2, hash_size]``. ``cache_policy`` uses a control-plane row
for the packed CLOCK states followed by one reference-epoch row per
request: ``[1 + num_requests, hot_buffer_size]``.
``scratch_locs`` and ``scratch_state`` hold reusable miss locations,
counters, and metadata shared by all layers.
"""
cache_index: torch.Tensor
cache_policy: torch.Tensor
scratch_locs: torch.Tensor
scratch_state: torch.Tensor
@cache_once
def _jit_spec_module(
item_size_bytes: int,
block_size: int,
num_top_k: int,
hot_buffer_size: int,
num_steps: int,
record_miss_plan: bool,
) -> Module:
template_args = make_cpp_args(
block_size,
num_top_k,
hot_buffer_size,
item_size_bytes,
num_steps,
record_miss_plan,
is_arch_support_pdl(),
)
return load_jit(
"hisparse_spec",
*template_args,
cuda_files=["kvcacheio/hisparse_spec.cuh"],
cuda_wrappers=[
(
"load_cache_to_device_buffer_spec",
f"load_cache_to_device_buffer_spec<{template_args}>",
)
],
)
def load_cache_to_device_buffer_spec_mla(
*,
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,
top_k_device_locs: torch.Tensor,
req_pool_indices: torch.Tensor,
seq_lens: torch.Tensor,
state: HiSparseSpecState,
num_real_reqs: torch.Tensor,
miss_src: torch.Tensor | None = None,
miss_dst: torch.Tensor | None = None,
miss_count: torch.Tensor | None = None,
) -> None:
"""Resolve all speculative steps and swap unique misses in one launch pair.
Optional miss-plan outputs use the same protocol as the single-step HiSparse
kernel, so shared-index layers can replay only the Host-to-GPU copies with
``copy_cache_planned_mla``.
"""
_, num_steps, num_top_k = top_k_tokens.shape
if not 2 <= num_steps <= 4:
raise ValueError(
f"HiSparse speculative swap requires 2-4 steps, got {num_steps}."
)
hot_buffer_size = state.cache_policy.size(1)
page_size = device_buffer_tokens.size(1) - hot_buffer_size
item_size_bytes = host_cache.stride(0) * host_cache.element_size()
record_miss_plan = miss_src is not None
if record_miss_plan:
if miss_dst is None or miss_count is None:
raise ValueError(
"miss_src, miss_dst, and miss_count must be provided together."
)
if miss_src.dtype != torch.int64 or miss_dst.dtype != torch.int32:
raise ValueError("miss_src must be int64 and miss_dst must be int32.")
if miss_count.dtype != torch.int32:
raise ValueError("miss_count must be int32.")
plan_capacity = num_steps * num_top_k
batch_size = top_k_tokens.size(0)
if (
miss_src.ndim != 2
or miss_dst.ndim != 2
or miss_src.size(0) < batch_size
or miss_dst.size(0) < batch_size
or miss_src.size(1) < plan_capacity
or miss_dst.size(1) < plan_capacity
):
raise ValueError(
"speculative miss_src/miss_dst must have shape "
f"[batch, >= steps * top_k] (capacity {plan_capacity})."
)
if miss_count.ndim != 1 or miss_count.numel() < batch_size:
raise ValueError("speculative miss_count must have shape [batch].")
if miss_src.stride(0) != miss_dst.stride(0):
raise ValueError("miss_src/miss_dst row strides must match.")
else:
if miss_dst is not None or miss_count is not None:
raise ValueError(
"miss_src, miss_dst, and miss_count must be provided together."
)
empty = torch.empty(0)
miss_src = miss_dst = miss_count = empty
module = _jit_spec_module(
item_size_bytes,
_GATHER_BLOCK_SIZE,
num_top_k,
hot_buffer_size,
num_steps,
record_miss_plan,
)
module.load_cache_to_device_buffer_spec(
top_k_tokens,
device_buffer_tokens,
host_cache_locs,
device_buffer_locs,
host_cache,
device_buffer,
top_k_device_locs,
req_pool_indices,
seq_lens,
state.cache_index,
state.cache_policy,
state.scratch_locs,
state.scratch_state,
num_real_reqs,
page_size,
miss_src,
miss_dst,
miss_count,
)
@functools.cache
def _jit_sparse_module(
item_size_bytes: int,
@@ -46,7 +201,7 @@ def _jit_sparse_module(
return load_jit(
"sparse_cache",
*cache_args,
cuda_files=["hisparse.cuh"],
cuda_files=["kvcacheio/hisparse.cuh"],
cuda_wrappers=[
(
"load_cache_to_device_buffer",
@@ -70,7 +225,7 @@ def _jit_copy_planned_module(
is_mla,
is_dsv4_layout,
skip_io,
cuda_files=["hisparse.cuh"],
cuda_files=["kvcacheio/hisparse.cuh"],
cuda_wrappers=[
(
"copy_cache_planned",
@@ -86,7 +241,7 @@ def _jit_dsv4_transfer_module(block_size: int) -> Module:
return load_jit(
"sparse_cache_dsv4_transfer",
block_size,
cuda_files=["hisparse.cuh"],
cuda_files=["kvcacheio/hisparse.cuh"],
cuda_wrappers=[
(
"transfer_cache_dsv4_mla",
test/registered/jit/benchmark/bench_hisparse_spec.py
+359
View File
@@ -0,0 +1,359 @@
from __future__ import annotations
from typing import NamedTuple
import torch
from sglang.kernels.jit.benchmark import marker
from sglang.kernels.ops.kvcache.hisparse import (
HiSparseSpecState,
copy_cache_planned_mla,
load_cache_to_device_buffer_mla,
load_cache_to_device_buffer_spec_mla,
)
from sglang.test.ci.ci_register import register_cuda_ci
register_cuda_ci(
est_time=45, stage="base-b-kernel-benchmark", runner_config="1-gpu-large"
)
DEVICE = "cuda"
NUM_STEPS = 4
NUM_INDEX_LAYERS = 4
NUM_SHARED_LAYERS = 3
TOP_K = 2048
HOT_BUFFER_SIZE = 4096
PAGE_SIZE = 64
ITEM_WORDS = 72
ITEM_SIZE_BYTES = ITEM_WORDS * 8
MISS_COUNT_PER_STEP = 196
UNIQUE_MISS_COUNT = 782
# CUDA Graph benchmarking replays 100 swap calls per graph. Keep every replay
# on a fresh miss range while staying inside the native GLM-5.2 context length.
MISS_ADVANCE = 800
SEQ_LEN = 1_048_576
TOKEN_SCALE = 1_000_003
class _BenchmarkState(NamedTuple):
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_buffers: torch.Tensor
req_pool_indices: torch.Tensor
seq_lens: torch.Tensor
num_real_reqs: torch.Tensor
lru_slots: torch.Tensor
spec_out: torch.Tensor
lru_out: torch.Tensor
miss_src: torch.Tensor
miss_dst: torch.Tensor
miss_count: torch.Tensor
swap_states: tuple[HiSparseSpecState, ...]
def _make_top_k_tokens(batch_size: int) -> torch.Tensor:
hit_count = TOP_K - MISS_COUNT_PER_STEP
steps = []
next_miss = HOT_BUFFER_SIZE
duplicate_tokens = []
for step in range(NUM_STEPS):
hits = torch.roll(
torch.arange(HOT_BUFFER_SIZE, dtype=torch.int32, device=DEVICE),
step * 137,
)[:hit_count]
if step < 2:
misses = torch.arange(
next_miss,
next_miss + MISS_COUNT_PER_STEP,
dtype=torch.int32,
device=DEVICE,
)
duplicate_tokens.append(misses[step])
next_miss += MISS_COUNT_PER_STEP
else:
unique_misses = torch.arange(
next_miss,
next_miss + MISS_COUNT_PER_STEP - 1,
dtype=torch.int32,
device=DEVICE,
)
misses = torch.cat((duplicate_tokens[step - 2].view(1), unique_misses))
next_miss += MISS_COUNT_PER_STEP - 1
steps.append(torch.cat((hits, misses)))
top_k_tokens = torch.stack(steps).unsqueeze(0)
top_k_tokens = top_k_tokens.repeat(batch_size, 1, 1).contiguous()
assert torch.unique(top_k_tokens[0, :, -MISS_COUNT_PER_STEP:]).numel() == (
UNIQUE_MISS_COUNT
)
return top_k_tokens
def _make_cache_index(batch_size: int) -> torch.Tensor:
hash_size = 1 << (2 * HOT_BUFFER_SIZE - 1).bit_length()
cache_index = torch.full(
(batch_size, 2, hash_size), -1, dtype=torch.int64, device=DEVICE
)
tokens = torch.arange(HOT_BUFFER_SIZE, dtype=torch.int64, device=DEVICE)
hash_slots = ((tokens * 2654435761) & (hash_size - 1)).to(torch.long)
cache_index[:, 0, hash_slots] = (tokens << 32) | tokens
return cache_index
def _build_state(batch_size: int) -> _BenchmarkState:
buffer_size = HOT_BUFFER_SIZE + PAGE_SIZE
scratch_size = HOT_BUFFER_SIZE
physical_tokens_per_req = buffer_size + scratch_size
top_k_tokens = _make_top_k_tokens(batch_size)
device_buffer_tokens = torch.full(
(batch_size, buffer_size), -1, dtype=torch.int32, device=DEVICE
)
device_buffer_tokens[:, :HOT_BUFFER_SIZE] = torch.arange(
HOT_BUFFER_SIZE, dtype=torch.int32, device=DEVICE
)
device_buffer_tokens = (
device_buffer_tokens.unsqueeze(0).repeat(NUM_INDEX_LAYERS, 1, 1).contiguous()
)
request_bases = (
torch.arange(batch_size, dtype=torch.int32, device=DEVICE).view(-1, 1)
* physical_tokens_per_req
)
device_buffer_locs = (
request_bases
+ torch.arange(buffer_size, dtype=torch.int32, device=DEVICE).view(1, -1)
).contiguous()
scratch_locs = (
request_bases
+ buffer_size
+ torch.arange(scratch_size, dtype=torch.int32, device=DEVICE).view(1, -1)
).contiguous()
host_cache_locs = torch.arange(SEQ_LEN, dtype=torch.int64, device=DEVICE)
host_cache_locs = host_cache_locs.view(1, -1).repeat(batch_size, 1).contiguous()
host_cache = torch.empty((SEQ_LEN, ITEM_WORDS), dtype=torch.int64, pin_memory=True)
host_cache.copy_(
torch.arange(SEQ_LEN, dtype=torch.int64).view(-1, 1) * TOKEN_SCALE
+ torch.arange(ITEM_WORDS, dtype=torch.int64).view(1, -1)
)
device_buffers = torch.empty(
(NUM_INDEX_LAYERS, batch_size * physical_tokens_per_req, ITEM_WORDS),
dtype=torch.int64,
device=DEVICE,
)
hot_locs = device_buffer_locs[:, :HOT_BUFFER_SIZE].to(torch.long)
hot_values = host_cache[:HOT_BUFFER_SIZE].to(DEVICE)
for device_buffer in device_buffers:
device_buffer[hot_locs] = hot_values
total_occurrences = NUM_STEPS * TOP_K
swap_states = []
for _ in range(NUM_INDEX_LAYERS):
scratch_state = torch.full(
(batch_size + 1, max(4 * batch_size, 5 * total_occurrences)),
-1,
dtype=torch.int32,
device=DEVICE,
)
scratch_state[0].zero_()
swap_states.append(
HiSparseSpecState(
cache_index=_make_cache_index(batch_size),
cache_policy=torch.zeros(
(batch_size + 1, HOT_BUFFER_SIZE),
dtype=torch.int32,
device=DEVICE,
),
scratch_locs=scratch_locs,
scratch_state=scratch_state,
)
)
return _BenchmarkState(
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_buffers=device_buffers,
req_pool_indices=torch.arange(batch_size, dtype=torch.int64, device=DEVICE),
seq_lens=torch.full(
(batch_size * NUM_STEPS,), SEQ_LEN, dtype=torch.int32, device=DEVICE
),
num_real_reqs=torch.tensor([batch_size], dtype=torch.int32, device=DEVICE),
lru_slots=torch.arange(HOT_BUFFER_SIZE, dtype=torch.int16, device=DEVICE)
.view(1, -1)
.repeat(batch_size, 1),
spec_out=torch.full(
(NUM_INDEX_LAYERS, *top_k_tokens.shape),
-1,
dtype=top_k_tokens.dtype,
device=DEVICE,
),
lru_out=torch.full_like(top_k_tokens, -1),
miss_src=torch.full(
(batch_size, total_occurrences), -1, dtype=torch.int64, device=DEVICE
),
miss_dst=torch.full(
(batch_size, total_occurrences), -1, dtype=torch.int32, device=DEVICE
),
miss_count=torch.zeros(batch_size, dtype=torch.int32, device=DEVICE),
swap_states=tuple(swap_states),
)
def _run_spec_layer(
state: _BenchmarkState, layer_idx: int, *, record_plan: bool = False
) -> None:
load_cache_to_device_buffer_spec_mla(
top_k_tokens=state.top_k_tokens,
device_buffer_tokens=state.device_buffer_tokens[layer_idx],
host_cache_locs=state.host_cache_locs,
device_buffer_locs=state.device_buffer_locs,
host_cache=state.host_cache,
device_buffer=state.device_buffers[layer_idx],
top_k_device_locs=state.spec_out[layer_idx],
req_pool_indices=state.req_pool_indices,
seq_lens=state.seq_lens,
state=state.swap_states[layer_idx],
num_real_reqs=state.num_real_reqs,
miss_src=state.miss_src if record_plan else None,
miss_dst=state.miss_dst if record_plan else None,
miss_count=state.miss_count if record_plan else None,
)
def _run_four_full_layers(state: _BenchmarkState) -> None:
for layer_idx in range(NUM_INDEX_LAYERS):
_run_spec_layer(state, layer_idx)
def _run_planned_copies(state: _BenchmarkState) -> None:
for layer_idx in range(1, NUM_SHARED_LAYERS + 1):
copy_cache_planned_mla(
miss_src=state.miss_src,
miss_dst=state.miss_dst,
miss_count=state.miss_count,
num_real_reqs=state.num_real_reqs,
host_cache=state.host_cache,
device_buffer=state.device_buffers[layer_idx],
item_size_bytes=ITEM_SIZE_BYTES,
num_blocks=8,
)
def _run_lru_step(state: _BenchmarkState, step: int) -> None:
batch_size = state.top_k_tokens.size(0)
load_cache_to_device_buffer_mla(
top_k_tokens=state.top_k_tokens[:, step],
device_buffer_tokens=state.device_buffer_tokens[0],
host_cache_locs=state.host_cache_locs,
device_buffer_locs=state.device_buffer_locs,
host_cache=state.host_cache,
device_buffer=state.device_buffers[0],
top_k_device_locs=state.lru_out[:, step],
req_pool_indices=state.req_pool_indices,
# All benchmark steps use the same logical length. Reuse a contiguous
# slice so the measured path does not allocate a temporary tensor.
seq_lens=state.seq_lens[:batch_size],
lru_slots=state.lru_slots,
item_size_bytes=ITEM_SIZE_BYTES,
num_top_k=TOP_K,
hot_buffer_size=HOT_BUFFER_SIZE,
page_size=PAGE_SIZE,
block_size=1024,
num_real_reqs=state.num_real_reqs,
)
def _assert_current_result(state: _BenchmarkState, impl: str) -> None:
torch.cuda.synchronize()
expected = state.top_k_tokens.to(torch.int64) * TOKEN_SCALE
if impl == "spec_4_full":
for layer_idx in range(NUM_INDEX_LAYERS):
actual = state.device_buffers[layer_idx, :, 0][
state.spec_out[layer_idx].to(torch.long)
]
torch.testing.assert_close(actual, expected)
else:
out = state.lru_out if impl == "lru_1_full" else state.spec_out[0]
actual = state.device_buffers[0, :, 0][out.to(torch.long)]
torch.testing.assert_close(actual, expected)
if impl == "spec_1_full_3_shared":
for bid in range(state.top_k_tokens.size(0)):
count = int(state.miss_count[bid].item())
src = state.miss_src[bid, :count].to(torch.long)
dst = state.miss_dst[bid, :count].to(torch.long)
expected = state.host_cache[src.cpu(), 0].to(DEVICE)
for layer_idx in range(1, NUM_SHARED_LAYERS + 1):
torch.testing.assert_close(
state.device_buffers[layer_idx, :, 0][dst],
expected,
)
def _check_initial_result(state: _BenchmarkState, impl: str) -> None:
if impl == "lru_1_full":
for step in range(NUM_STEPS):
_run_lru_step(state, step)
elif impl == "spec_1_full":
_run_spec_layer(state, 0)
elif impl == "spec_4_full":
_run_four_full_layers(state)
else:
_run_spec_layer(state, 0, record_plan=True)
_run_planned_copies(state)
_assert_current_result(state, impl)
@marker.parametrize("batch_size", [1, 2, 4, 8, 16, 32, 64], [1])
@marker.benchmark(
"impl",
["lru_1_full", "spec_1_full", "spec_4_full", "spec_1_full_3_shared"],
)
def benchmark(batch_size: int, impl: str):
state = _build_state(batch_size)
_check_initial_result(state, impl)
if impl == "lru_1_full":
def run() -> None:
state.top_k_tokens[..., -MISS_COUNT_PER_STEP:].add_(MISS_ADVANCE)
for step in range(NUM_STEPS):
_run_lru_step(state, step)
elif impl == "spec_1_full":
def run() -> None:
state.top_k_tokens[..., -MISS_COUNT_PER_STEP:].add_(MISS_ADVANCE)
_run_spec_layer(state, 0)
elif impl == "spec_4_full":
def run() -> None:
state.top_k_tokens[..., -MISS_COUNT_PER_STEP:].add_(MISS_ADVANCE)
_run_four_full_layers(state)
else:
def run() -> None:
state.top_k_tokens[..., -MISS_COUNT_PER_STEP:].add_(MISS_ADVANCE)
_run_spec_layer(state, 0, record_plan=True)
_run_planned_copies(state)
result = marker.do_bench(
run,
use_cuda_graph=True,
warmup_iters=5,
replay_iters=40,
disable_log_bandwidth=True,
memory_args=None,
memory_output=None,
)
_assert_current_result(state, impl)
return result
if __name__ == "__main__":
benchmark.run()
test/registered/jit/test_hisparse_spec.py
+550
View File
@@ -0,0 +1,550 @@
from __future__ import annotations
from typing import NamedTuple
import pytest
import torch
from sglang.kernels.ops.kvcache.hisparse import (
HiSparseSpecState,
copy_cache_planned_mla,
load_cache_to_device_buffer_spec_mla,
)
from sglang.srt.utils import is_npu, is_xpu
from sglang.test.ci.ci_register import register_cuda_ci
from sglang.test.test_utils import CustomTestCase
register_cuda_ci(est_time=120, stage="base-b-kernel-unit", runner_config="1-gpu-large")
pytestmark = pytest.mark.skipif(
not torch.cuda.is_available() or is_npu() or is_xpu(),
reason="HiSparse speculative swap tests require a CUDA GPU.",
)
DEVICE = "cuda"
TOKEN_SCALE = 1_000_003
class _SwapState(NamedTuple):
device_buffer_tokens: torch.Tensor
device_buffer_locs: torch.Tensor
host_cache_locs: torch.Tensor
host_cache: torch.Tensor
device_buffer: torch.Tensor
swap_state: HiSparseSpecState
def _make_cache_index(num_reqs: int, hot_buffer_size: int) -> torch.Tensor:
hash_size = 1 << (2 * hot_buffer_size - 1).bit_length()
cache_index = torch.full(
(num_reqs, 2, hash_size), -1, dtype=torch.int64, device=DEVICE
)
tokens = torch.arange(hot_buffer_size, dtype=torch.int64, device=DEVICE)
hash_slots = ((tokens * 2654435761) & (hash_size - 1)).to(torch.long)
packed_entries = (tokens << 32) | tokens
cache_index[:, 0, hash_slots] = packed_entries
return cache_index
def _make_state(
*,
num_reqs: int,
hot_buffer_size: int,
page_size: int,
scratch_size: int,
seq_len: int,
item_words: int,
metadata_occurrences: int,
) -> _SwapState:
buffer_size = hot_buffer_size + page_size
device_buffer_tokens = torch.full(
(num_reqs, buffer_size), -1, dtype=torch.int32, device=DEVICE
)
device_buffer_tokens[:, :hot_buffer_size] = torch.arange(
hot_buffer_size, dtype=torch.int32, device=DEVICE
)
physical_tokens_per_req = buffer_size + scratch_size
request_bases = (
torch.arange(num_reqs, dtype=torch.int32, device=DEVICE).view(-1, 1)
* physical_tokens_per_req
)
device_buffer_locs = (
request_bases
+ torch.arange(buffer_size, dtype=torch.int32, device=DEVICE).view(1, -1)
).contiguous()
scratch_locs = (
request_bases
+ buffer_size
+ torch.arange(scratch_size, dtype=torch.int32, device=DEVICE).view(1, -1)
).contiguous()
host_cache_locs = torch.arange(seq_len, dtype=torch.int64, device=DEVICE)
host_cache_locs = host_cache_locs.view(1, -1).repeat(num_reqs, 1).contiguous()
host_cache = torch.empty((seq_len, item_words), dtype=torch.int64, pin_memory=True)
host_cache.copy_(
torch.arange(seq_len, dtype=torch.int64).view(-1, 1) * TOKEN_SCALE
+ torch.arange(item_words, dtype=torch.int64).view(1, -1)
)
device_buffer = torch.full(
(num_reqs * physical_tokens_per_req, item_words),
-1,
dtype=torch.int64,
device=DEVICE,
)
hot_locs = device_buffer_locs[:, :hot_buffer_size].to(torch.long)
device_buffer[hot_locs] = host_cache[:hot_buffer_size].to(DEVICE)
scratch_state = torch.full(
(num_reqs + 1, max(4 * num_reqs, 5 * metadata_occurrences)),
-1,
dtype=torch.int32,
device=DEVICE,
)
scratch_state[0].zero_()
swap_state = HiSparseSpecState(
cache_index=_make_cache_index(num_reqs, hot_buffer_size),
cache_policy=torch.zeros(
(num_reqs + 1, hot_buffer_size),
dtype=torch.int32,
device=DEVICE,
),
scratch_locs=scratch_locs,
scratch_state=scratch_state,
)
return _SwapState(
device_buffer_tokens=device_buffer_tokens,
device_buffer_locs=device_buffer_locs,
host_cache_locs=host_cache_locs,
host_cache=host_cache,
device_buffer=device_buffer,
swap_state=swap_state,
)
def _run_swap(
*,
top_k_tokens: torch.Tensor,
seq_lens: torch.Tensor,
state: _SwapState,
out: torch.Tensor | None = None,
req_pool_indices: torch.Tensor | None = None,
num_real_reqs: torch.Tensor | None = None,
miss_src: torch.Tensor | None = None,
miss_dst: torch.Tensor | None = None,
miss_count: torch.Tensor | None = None,
) -> torch.Tensor:
if out is None:
out = torch.full_like(top_k_tokens, -1)
else:
out.fill_(-1)
num_reqs = top_k_tokens.size(0)
if req_pool_indices is None:
req_pool_indices = torch.arange(num_reqs, dtype=torch.int64, device=DEVICE)
if num_real_reqs is None:
num_real_reqs = torch.tensor([num_reqs], dtype=torch.int32, device=DEVICE)
load_cache_to_device_buffer_spec_mla(
top_k_tokens=top_k_tokens,
device_buffer_tokens=state.device_buffer_tokens,
host_cache_locs=state.host_cache_locs,
device_buffer_locs=state.device_buffer_locs,
host_cache=state.host_cache,
device_buffer=state.device_buffer,
top_k_device_locs=out,
req_pool_indices=req_pool_indices,
seq_lens=seq_lens,
state=state.swap_state,
num_real_reqs=num_real_reqs,
miss_src=miss_src,
miss_dst=miss_dst,
miss_count=miss_count,
)
return out
def _assert_output_matches_tokens(
state: _SwapState, out: torch.Tensor, tokens: torch.Tensor
) -> None:
actual = state.device_buffer[out.to(torch.long)]
expected = tokens.to(torch.int64).unsqueeze(-1) * TOKEN_SCALE + torch.arange(
state.device_buffer.size(-1), dtype=torch.int64, device=DEVICE
)
torch.testing.assert_close(actual, expected)
class TestHiSparseSpec(CustomTestCase):
def test_deduplicates_repeated_misses_and_copies_full_items(self) -> None:
hot_size, page_size = 4096, 64
num_steps, top_k, item_words = 4, 2048, 72
total_occurrences = num_steps * top_k
state = _make_state(
num_reqs=1,
hot_buffer_size=hot_size,
page_size=page_size,
scratch_size=hot_size,
seq_len=16384,
item_words=item_words,
metadata_occurrences=total_occurrences,
)
miss_count = 196
hits = torch.arange(top_k - miss_count, dtype=torch.int32, device=DEVICE)
shared_misses = hot_size + torch.arange(
miss_count, dtype=torch.int32, device=DEVICE
)
step = torch.cat((hits, shared_misses))
top_k_tokens = step.view(1, 1, -1).repeat(1, num_steps, 1)
seq_lens = torch.full((num_steps,), 16384, dtype=torch.int32, device=DEVICE)
out = _run_swap(top_k_tokens=top_k_tokens, seq_lens=seq_lens, state=state)
torch.cuda.synchronize()
_assert_output_matches_tokens(state, out, top_k_tokens)
self.assertEqual(int(state.swap_state.scratch_state[0, 0].item()), miss_count)
repeated_miss_locs = out[0, :, -miss_count:]
self.assertTrue(torch.all(repeated_miss_locs == repeated_miss_locs[0]).item())
def test_copies_782_cross_step_unique_misses(self) -> None:
hot_size, page_size = 4096, 64
num_steps, top_k, item_words = 4, 2048, 72
total_occurrences = num_steps * top_k
state = _make_state(
num_reqs=1,
hot_buffer_size=hot_size,
page_size=page_size,
scratch_size=hot_size,
seq_len=16384,
item_words=item_words,
metadata_occurrences=total_occurrences,
)
steps = []
next_miss = hot_size
for step_idx, miss_count in enumerate((196, 196, 195, 195)):
hits = torch.roll(
torch.arange(hot_size, dtype=torch.int32, device=DEVICE),
step_idx * 137,
)[: top_k - miss_count]
misses = torch.arange(
next_miss,
next_miss + miss_count,
dtype=torch.int32,
device=DEVICE,
)
next_miss += miss_count
steps.append(torch.cat((hits, misses)))
top_k_tokens = torch.stack(steps).unsqueeze(0).contiguous()
seq_lens = torch.full((num_steps,), 16384, dtype=torch.int32, device=DEVICE)
out = _run_swap(
top_k_tokens=top_k_tokens,
seq_lens=seq_lens,
state=state,
)
torch.cuda.synchronize()
_assert_output_matches_tokens(state, out, top_k_tokens)
self.assertEqual(int(state.swap_state.scratch_state[0, 0].item()), 782)
def test_records_union_plan_for_shared_layer_io(self) -> None:
hot_size, page_size = 4096, 64
num_steps, top_k, item_words = 4, 2048, 72
total_occurrences = num_steps * top_k
state = _make_state(
num_reqs=1,
hot_buffer_size=hot_size,
page_size=page_size,
scratch_size=hot_size,
seq_len=16384,
item_words=item_words,
metadata_occurrences=total_occurrences,
)
steps = []
next_miss = hot_size
for step_idx, step_miss_count in enumerate((196, 196, 195, 195)):
hits = torch.roll(
torch.arange(hot_size, dtype=torch.int32, device=DEVICE),
step_idx * 137,
)[: top_k - step_miss_count]
misses = torch.arange(
next_miss,
next_miss + step_miss_count,
dtype=torch.int32,
device=DEVICE,
)
next_miss += step_miss_count
steps.append(torch.cat((hits, misses)))
top_k_tokens = torch.stack(steps).unsqueeze(0).contiguous()
seq_lens = torch.full((num_steps,), 16384, dtype=torch.int32, device=DEVICE)
miss_src = torch.full(
(1, total_occurrences), -1, dtype=torch.int64, device=DEVICE
)
miss_dst = torch.full(
(1, total_occurrences), -1, dtype=torch.int32, device=DEVICE
)
miss_count = torch.full((1,), -1, dtype=torch.int32, device=DEVICE)
_run_swap(
top_k_tokens=top_k_tokens,
seq_lens=seq_lens,
state=state,
miss_src=miss_src,
miss_dst=miss_dst,
miss_count=miss_count,
)
shared_layer_buffer = torch.full_like(state.device_buffer, -1)
copy_cache_planned_mla(
miss_src=miss_src,
miss_dst=miss_dst,
miss_count=miss_count,
num_real_reqs=torch.ones(1, dtype=torch.int32, device=DEVICE),
host_cache=state.host_cache,
device_buffer=shared_layer_buffer,
item_size_bytes=state.host_cache.stride(0)
* state.host_cache.element_size(),
)
torch.cuda.synchronize()
self.assertEqual(int(miss_count.item()), 782)
count = int(miss_count.item())
src = miss_src[0, :count].to(torch.long)
dst = miss_dst[0, :count].to(torch.long)
torch.testing.assert_close(
shared_layer_buffer[dst], state.host_cache[src.cpu()].to(DEVICE)
)
torch.testing.assert_close(shared_layer_buffer[dst], state.device_buffer[dst])
def test_padded_request_clears_stale_plan_count(self) -> None:
hot_size, page_size = 4096, 64
num_steps, top_k = 4, 2048
total_occurrences = num_steps * top_k
state = _make_state(
num_reqs=2,
hot_buffer_size=hot_size,
page_size=page_size,
scratch_size=hot_size,
seq_len=8192,
item_words=1,
metadata_occurrences=total_occurrences,
)
top_k_tokens = torch.arange(top_k, dtype=torch.int32, device=DEVICE).view(
1, 1, -1
)
top_k_tokens = top_k_tokens.repeat(2, num_steps, 1).contiguous()
seq_lens = torch.full((2 * num_steps,), 8192, dtype=torch.int32, device=DEVICE)
miss_src = torch.full(
(2, total_occurrences), -1, dtype=torch.int64, device=DEVICE
)
miss_dst = torch.full(
(2, total_occurrences), -1, dtype=torch.int32, device=DEVICE
)
miss_count = torch.full((2,), 123, dtype=torch.int32, device=DEVICE)
_run_swap(
top_k_tokens=top_k_tokens,
seq_lens=seq_lens,
state=state,
num_real_reqs=torch.ones(1, dtype=torch.int32, device=DEVICE),
miss_src=miss_src,
miss_dst=miss_dst,
miss_count=miss_count,
)
torch.cuda.synchronize()
self.assertEqual(int(miss_count[1].item()), 0)
def test_resolves_all_speculative_extra_page_slots_without_host_io(self) -> None:
hot_size, page_size = 4096, 64
num_steps, top_k = 4, 2048
total_occurrences = num_steps * top_k
seq_len = 8192
state = _make_state(
num_reqs=1,
hot_buffer_size=hot_size,
page_size=page_size,
scratch_size=hot_size,
seq_len=seq_len,
item_words=72,
metadata_occurrences=total_occurrences,
)
draft_tokens = torch.arange(
seq_len - num_steps, seq_len, dtype=torch.int32, device=DEVICE
)
extra_offsets = torch.tensor([0, 7, 31, 63], device=DEVICE)
extra_locs = state.device_buffer_locs[0, hot_size + extra_offsets].to(
torch.long
)
state.device_buffer_tokens[0, hot_size + extra_offsets] = draft_tokens
state.device_buffer[extra_locs] = state.host_cache[
draft_tokens.to(device="cpu", dtype=torch.long)
].to(DEVICE)
state.host_cache_locs[0, draft_tokens.to(torch.long)] = -1
hits = torch.arange(top_k - 1, dtype=torch.int32, device=DEVICE)
top_k_tokens = torch.stack(
[torch.cat((hits, draft_tokens[step : step + 1])) for step in range(4)]
).unsqueeze(0)
seq_lens = draft_tokens + 1
out = _run_swap(top_k_tokens=top_k_tokens, seq_lens=seq_lens, state=state)
torch.cuda.synchronize()
_assert_output_matches_tokens(state, out, top_k_tokens)
torch.testing.assert_close(out[0, :, -1].to(torch.long), extra_locs)
self.assertEqual(int(state.swap_state.scratch_state[0, 0].item()), 0)
def test_full_union_overflow_preserves_all_8192_outputs(self) -> None:
hot_size, page_size = 4096, 64
num_steps, top_k = 4, 2048
total_occurrences = num_steps * top_k
state = _make_state(
num_reqs=1,
hot_buffer_size=hot_size,
page_size=page_size,
scratch_size=total_occurrences - hot_size,
seq_len=16384,
item_words=72,
metadata_occurrences=total_occurrences,
)
top_k_tokens = (
hot_size + torch.arange(total_occurrences, dtype=torch.int32, device=DEVICE)
).view(1, num_steps, top_k)
seq_lens = torch.full((num_steps,), 16384, dtype=torch.int32, device=DEVICE)
miss_src = torch.full(
(1, total_occurrences), -1, dtype=torch.int64, device=DEVICE
)
miss_dst = torch.full(
(1, total_occurrences), -1, dtype=torch.int32, device=DEVICE
)
miss_count = torch.full((1,), -1, dtype=torch.int32, device=DEVICE)
out = _run_swap(
top_k_tokens=top_k_tokens,
seq_lens=seq_lens,
state=state,
miss_src=miss_src,
miss_dst=miss_dst,
miss_count=miss_count,
)
torch.cuda.synchronize()
_assert_output_matches_tokens(state, out, top_k_tokens)
self.assertEqual(torch.unique(out).numel(), total_occurrences)
self.assertEqual(
int(state.swap_state.scratch_state[0, 0].item()), total_occurrences
)
self.assertEqual(int(miss_count.item()), total_occurrences)
self.assertTrue(miss_src.ge(0).all().item())
self.assertTrue(miss_dst.ge(0).all().item())
def test_packed_ring_supports_glm52_native_context_length(self) -> None:
hot_size, page_size = 4096, 64
num_steps, top_k = 4, 2048
total_occurrences = num_steps * top_k
seq_len = 1_048_648
state = _make_state(
num_reqs=1,
hot_buffer_size=hot_size,
page_size=page_size,
scratch_size=hot_size,
seq_len=seq_len,
item_words=1,
metadata_occurrences=total_occurrences,
)
top_k_tokens = torch.arange(top_k, dtype=torch.int32, device=DEVICE).view(
1, 1, -1
)
top_k_tokens = top_k_tokens.repeat(1, num_steps, 1)
high_token = seq_len - 1
top_k_tokens[:, :, -1] = high_token
seq_lens = torch.full((num_steps,), seq_len, dtype=torch.int32, device=DEVICE)
out = _run_swap(top_k_tokens=top_k_tokens, seq_lens=seq_lens, state=state)
torch.cuda.synchronize()
_assert_output_matches_tokens(state, out, top_k_tokens)
self.assertTrue(out.ge(0).all().item())
# The first call admits the high token into the packed hash. The
# second call must resolve it as a hot hit rather than truncating the
# packed int64 entry and repeating Host-to-GPU IO.
out = _run_swap(top_k_tokens=top_k_tokens, seq_lens=seq_lens, state=state)
torch.cuda.synchronize()
_assert_output_matches_tokens(state, out, top_k_tokens)
self.assertEqual(int(state.swap_state.scratch_state[0, 0].item()), 0)
def test_cuda_graph_replay_preserves_valid_locations(self) -> None:
hot_size, page_size = 4096, 64
num_steps, top_k = 4, 2048
total_occurrences = num_steps * top_k
state = _make_state(
num_reqs=1,
hot_buffer_size=hot_size,
page_size=page_size,
scratch_size=hot_size,
seq_len=65536,
item_words=72,
metadata_occurrences=total_occurrences,
)
top_k_tokens = torch.arange(top_k, dtype=torch.int32, device=DEVICE).repeat(
num_steps, 1
)
for step, miss_count in enumerate((164, 102, 61, 20)):
top_k_tokens[step, -miss_count:] = torch.arange(
8192 + step * top_k,
8192 + step * top_k + miss_count,
dtype=torch.int32,
device=DEVICE,
)
top_k_tokens = top_k_tokens.unsqueeze(0).contiguous()
seq_lens = torch.tensor(
[65533, 65534, 65535, 65536], dtype=torch.int32, device=DEVICE
)
_run_swap(top_k_tokens=top_k_tokens, seq_lens=seq_lens, state=state)
torch.cuda.synchronize()
graph_out = torch.full_like(top_k_tokens, -1)
req_pool_indices = torch.arange(1, dtype=torch.int64, device=DEVICE)
num_real_reqs = torch.tensor([1], dtype=torch.int32, device=DEVICE)
graph = torch.cuda.CUDAGraph()
with torch.cuda.graph(graph):
_run_swap(
top_k_tokens=top_k_tokens,
seq_lens=seq_lens,
state=state,
out=graph_out,
req_pool_indices=req_pool_indices,
num_real_reqs=num_real_reqs,
)
for _ in range(4):
graph.replay()
torch.cuda.synchronize()
_assert_output_matches_tokens(state, graph_out, top_k_tokens)
self.assertTrue(graph_out.ge(0).all().item())
def test_rejects_invalid_step_shape_before_compilation(self) -> None:
state = _make_state(
num_reqs=1,
hot_buffer_size=4096,
page_size=64,
scratch_size=4096,
seq_len=8192,
item_words=1,
metadata_occurrences=8192,
)
with self.assertRaisesRegex(ValueError, "2-4 steps"):
_run_swap(
top_k_tokens=torch.zeros(
(1, 1, 2048), dtype=torch.int32, device=DEVICE
),
seq_lens=torch.tensor([8192], dtype=torch.int32, device=DEVICE),
state=state,
)
if __name__ == "__main__":
raise SystemExit(pytest.main([__file__, "-v", "-s"]))