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

perf: migrate Req token-id storage to array.array('q') in Scheduler (#25098)

Browse Source 
Co-authored-by: jialino <jialino@fb.com>
This commit is contained in:
Jialin Ouyang
2026-05-22 10:51:07 -07:00
committed by GitHub
co-authored by jialino
parent 5e9bd21979
commit 06c23d55b5
34 changed files with 833 additions and 299 deletions
Download Patch File Download Diff File Expand all files Collapse all files
benchmark/scheduler/bench_token_storage.py
+334
View File
@@ -0,0 +1,334 @@
"""Benchmark `list[int]` vs `array.array('q')` storage for
`Req.origin_input_ids` / `Req.output_ids` over one request lifecycle.
Simulated steps (per batch):
1. ingest -- tokenizer list[int] -> storage container.
2. prefix_match -- scheduler radix-tree lookup; RadixKey.match()
zip+!= walk. Exposes the per-element PyLong-boxing
cost array.array introduces (list[int] iterates
existing PyLongs and pays nothing).
3. prefill -- (a) fill_ids = origin + output,
(b) per-req slice fill_ids[prefix_len:],
(c) cross-req flatten + pinned cuda tensor build.
4. decode -- per-step output.append(next_token) for n_decode steps.
5. finish -- cache_finished_req:
(a) concat (origin + output)[:kv_committed_len]
for the radix-tree insert.
(b) RadixKey.match() zip+!= walk during insert's
tree traversal — second PyLong-boxing hotspot
on the array.array path.
Usage:
python benchmark/scheduler/bench_token_storage.py
"""
from __future__ import annotations
import time
from array import array
from collections import defaultdict
from contextlib import contextmanager
from itertools import chain
from typing import Any, Callable, Iterator
import numpy as np
import torch
# Per-req stages accumulate across reqs in a batch; batch_torch_tensor
# is the single cross-req prepare_for_extend tensor build.
STAGES = (
"ingest",
"prefix_match",
"prefill_concat",
"prefill_perreq_slice",
"batch_torch_tensor",
"decode_append",
"finish_concat",
"cache_finished_req",
)
def _ingest_list(seed: list[int]) -> list[int]:
return seed
def _ingest_pyarray(seed: list[int]) -> array:
return array("q", seed)
def _empty_list() -> list[int]:
return []
def _empty_pyarray() -> array:
return array("q")
def _zip_iterate(t0: Any, t1: Any) -> int:
"""Simulate zip iteration which surface PyLong boxing cost in array scenario"""
i = 0
for a, b in zip(t0, t1):
if a != b:
break
i += 1
return i
def _batch_tensor_from_lists(parts: list[list[int]]) -> torch.Tensor:
flat = list(chain.from_iterable(parts))
return torch.tensor(flat, dtype=torch.int64, pin_memory=True).to(
"cuda", non_blocking=True
)
def _batch_tensor_from_pyarrays(parts: list[array]) -> torch.Tensor:
# np.frombuffer gives a zero-copy view; np.concatenate is one C-level
# memcpy. This bypasses the per-element PyLong->int64 walk that
# torch.tensor(array('q')) would otherwise do.
views = [np.frombuffer(p, dtype=np.int64) for p in parts]
combined = np.concatenate(views) if len(views) > 1 else views[0]
return torch.from_numpy(combined).pin_memory().to("cuda", non_blocking=True)
LIST_KIT = {
"ingest_fn": _ingest_list,
"empty_fn": _empty_list,
"batch_torch_fn": _batch_tensor_from_lists,
}
PYARRAY_KIT = {
"ingest_fn": _ingest_pyarray,
"empty_fn": _empty_pyarray,
"batch_torch_fn": _batch_tensor_from_pyarrays,
}
@contextmanager
def timed(timings: dict[str, float], stage: str) -> Iterator[None]:
t0 = time.monotonic_ns()
try:
yield
finally:
timings[stage] += time.monotonic_ns() - t0
def simulate(
seeds: list[list[int]],
n_decode: int,
*,
ingest_fn: Callable[[list[int]], Any],
empty_fn: Callable[[], Any],
batch_torch_fn: Callable[[list[Any]], torch.Tensor],
) -> dict[str, float]:
"""One scheduling-round lifecycle. Returns per-stage cumulative ns."""
timings: dict[str, float] = defaultdict(float)
n_reqs = len(seeds)
n_origins = [len(s) for s in seeds]
origins: list[Any] = [None] * n_reqs
outputs: list[Any] = [None] * n_reqs
# 1. ingest
for i, seed in enumerate(seeds):
with timed(timings, "ingest"):
origins[i] = ingest_fn(seed)
outputs[i] = empty_fn()
# 2. prefix_match: simulating the worse scenario of PyLong-boxing overhead during prefix_match
for i in range(n_reqs):
with timed(timings, "prefix_match"):
_ = _zip_iterate(origins[i], origins[i])
# 3. prefill
per_req_slices: list[Any] = [None] * n_reqs
for i in range(n_reqs):
# 3a. fill_ids = origin_input_ids + output_ids
with timed(timings, "prefill_concat"):
fill_ids = origins[i] + outputs[i]
# 3b. input_ids = fill_ids[len(prefix_indices):]; prefix_len=0 here.
with timed(timings, "prefill_perreq_slice"):
per_req_slices[i] = fill_ids[0:]
# 3c. prepare_for_extend tensor build: flatten per-req slices, then
# build the pinned GPU tensor (kit-specific path).
with timed(timings, "batch_torch_tensor"):
_ = batch_torch_fn(per_req_slices)
# 4. decode
for i in range(n_reqs):
with timed(timings, "decode_append"):
for j in range(n_decode):
outputs[i].append(j)
# 5. finish: cache_finished_req -> insert -> _insert_helper tree walk.
for i in range(n_reqs):
# 5a. (origin + output)[:kv_committed_len] for the radix-tree insert.
with timed(timings, "finish_concat"):
committed = (origins[i] + outputs[i])[: n_origins[i] + n_decode]
# 5b. simulating the worse scenario of PyLong-boxing overhead during cache_finished_req
with timed(timings, "cache_finished_req"):
_ = _zip_iterate(committed, committed)
return timings
def bench_lifecycle(
seeds: list[list[int]],
n_decode: int,
iterations: int,
*,
ingest_fn: Callable[[list[int]], Any],
empty_fn: Callable[[], Any],
batch_torch_fn: Callable[[list[Any]], torch.Tensor],
warmup: int = 5,
) -> dict[str, float]:
"""Run simulate() N times, return mean per-stage us per batch.
GPU sync is excluded from per-iteration timing: production issues
`to(device, non_blocking=True)` and continues, so we measure issue
cost rather than H2D completion.
"""
kit = {
"ingest_fn": ingest_fn,
"empty_fn": empty_fn,
"batch_torch_fn": batch_torch_fn,
}
torch.cuda.synchronize()
for _ in range(warmup):
simulate(seeds, n_decode, **kit)
torch.cuda.synchronize()
accum: dict[str, float] = defaultdict(float)
for _ in range(iterations):
t = simulate(seeds, n_decode, **kit)
for k, v in t.items():
accum[k] += v
torch.cuda.synchronize()
return {k: accum[k] / iterations / 1000.0 for k in STAGES} # ns -> us
def print_breakdown(title: str, results: dict[str, dict[str, float]]) -> None:
"""Print per-stage timings with delta us vs the first (baseline) column."""
labels = list(results.keys())
baseline_label = labels[0]
baseline = results[baseline_label]
width = max(len(s) for s in STAGES)
header_cells = [f"{baseline_label + ' us':>11s}"]
for lbl in labels[1:]:
header_cells.append(f"{lbl + ' us':>11s}")
header_cells.append(f"{'delta':>10s}")
print(f"=== {title} ===")
print(f"{'Stage':<{width}s} " + " ".join(header_cells))
print("-" * (width + 2 + sum(len(c) + 2 for c in header_cells)))
for s in STAGES:
cells = [f"{baseline[s]:>11.3f}"]
for lbl in labels[1:]:
v = results[lbl][s]
cells.append(f"{v:>11.3f}")
d = v - baseline[s]
cells.append(f"{d:>+10.3f}")
print(f"{s:<{width}s} " + " ".join(cells))
print("-" * (width + 2 + sum(len(c) + 2 for c in header_cells)))
base_total = sum(baseline.values())
total_cells = [f"{base_total:>11.3f}"]
for lbl in labels[1:]:
v = sum(results[lbl].values())
total_cells.append(f"{v:>11.3f}")
d = v - base_total
total_cells.append(f"{d:>+10.3f}")
print(f"{'TOTAL':<{width}s} " + " ".join(total_cells))
print()
for lbl in labels[1:]:
v = sum(results[lbl].values())
d = v - base_total
speedup = base_total / v if v > 0 else 0.0
verdict = "LOSES" if d > 0 else "WINS"
print(
f" {lbl:<14s} vs {baseline_label}: {verdict} by {abs(d):>8.2f} us ({speedup:.2f}x)"
)
print()
def microbench_torch_tensor_paths(
sizes: tuple[int, ...] = (1_000, 10_000, 100_000)
) -> None:
"""Compare three CPU-buffer -> pinned cuda tensor paths.
A. torch.tensor(list, pin) -> cuda
B. torch.tensor(array('q'), pin) -> cuda
C. torch.from_numpy(np.frombuffer(array('q'))).pin() -> cuda
"""
def t(fn, iterations: int) -> float:
for _ in range(20):
fn()
torch.cuda.synchronize()
t0 = time.monotonic_ns()
for _ in range(iterations):
fn()
torch.cuda.synchronize()
return (time.monotonic_ns() - t0) / iterations / 1000.0
print("=== microbench: CPU-buffer -> pinned cuda tensor (us/op) ===\n")
width = 56
print(f"{'Path':<{width}s} " + " ".join(f"{f'N={n}':>10s}" for n in sizes))
print("-" * (width + 2 + 12 * len(sizes)))
for label, build in [
(
"(A) torch.tensor(list, pin) -> cuda",
lambda x: torch.tensor(x, dtype=torch.int64, pin_memory=True).to(
"cuda", non_blocking=True
),
),
(
"(B) torch.tensor(array('q'), pin) -> cuda (naive)",
lambda x: torch.tensor(x, dtype=torch.int64, pin_memory=True).to(
"cuda", non_blocking=True
),
),
(
"(C) from_numpy(frombuf(array('q'))).pin() -> cuda",
lambda x: torch.from_numpy(np.frombuffer(x, dtype=np.int64))
.pin_memory()
.to("cuda", non_blocking=True),
),
]:
cells = []
for n in sizes:
iters = max(50, 200_000 // max(n, 1))
if "(A)" in label:
src = list(range(n))
else:
src = array("q", range(n))
us = t(lambda src=src, build=build: build(src), iters)
cells.append(f"{us:>10.2f}")
print(f"{label:<{width}s} " + " ".join(cells))
print()
def main() -> None:
microbench_torch_tensor_paths()
n_reqs = 2
cases = [
("short prompt N_origin=1K N_decode=1K", 1_000, 1_000, 1_000),
("medium prompt N_origin=10K N_decode=1K", 10_000, 1_000, 200),
("long prompt N_origin=100K N_decode=1K", 100_000, 1_000, 30),
]
print(f"Batch size = {n_reqs} reqs/batch (per-req stages accumulate)\n")
for label, n_origin, n_decode, iters in cases:
seeds = [list(range(n_origin)) for _ in range(n_reqs)]
results = {
"list": bench_lifecycle(seeds, n_decode, iters, **LIST_KIT),
"pyarray": bench_lifecycle(seeds, n_decode, iters, **PYARRAY_KIT),
}
print_breakdown(label, results)
if __name__ == "__main__":
main()
python/sglang/srt/disaggregation/decode_schedule_batch_mixin.py
+2 -1
View File
@@ -1,6 +1,7 @@
from __future__ import annotations from __future__ import annotations
import logging import logging
from array import array
from http import HTTPStatus from http import HTTPStatus
from typing import TYPE_CHECKING, List from typing import TYPE_CHECKING, List
@@ -71,7 +72,7 @@ class ScheduleBatchDisaggregationDecodeMixin:
# Set fields # Set fields
self.input_ids = torch.tensor( self.input_ids = torch.tensor(
sum(input_ids, []), dtype=torch.int32, device=self.device sum(input_ids, array("q")), dtype=torch.int32, device=self.device
) )
self.req_pool_indices = torch.tensor( self.req_pool_indices = torch.tensor(
req_pool_indices, dtype=torch.int64, device=self.device req_pool_indices, dtype=torch.int64, device=self.device
python/sglang/srt/disaggregation/encode_receiver.py
+2 -1
View File
@@ -7,6 +7,7 @@ import threading
import time import time
import uuid import uuid
from abc import ABC, abstractmethod from abc import ABC, abstractmethod
from array import array
from collections import OrderedDict, defaultdict from collections import OrderedDict, defaultdict
from enum import IntEnum from enum import IntEnum
from http import HTTPStatus from http import HTTPStatus
@@ -588,7 +589,7 @@ class WaitingImageRequest:
**self.recv_embedding_data.get_mm_extra_meta(), **self.recv_embedding_data.get_mm_extra_meta(),
) )
self.recv_req.mm_inputs = mm_inputs self.recv_req.mm_inputs = mm_inputs
self.recv_req.input_ids = mm_inputs.input_ids self.recv_req.input_ids = array("q", mm_inputs.input_ids)
self.status = WaitingImageRequestStatus.SUCCESS self.status = WaitingImageRequestStatus.SUCCESS
self.recv_socket.close() self.recv_socket.close()
python/sglang/srt/dllm/mixin/req.py
+2 -1
View File
@@ -1,6 +1,7 @@
from __future__ import annotations from __future__ import annotations
import enum import enum
from array import array
from typing import TYPE_CHECKING, Optional from typing import TYPE_CHECKING, Optional
from sglang.srt.dllm.config import DllmConfig from sglang.srt.dllm.config import DllmConfig
@@ -62,7 +63,7 @@ class ReqDllmMixin:
self.fill_ids = ( self.fill_ids = (
self.origin_input_ids self.origin_input_ids
+ self.output_ids + self.output_ids
+ [self.dllm_config.mask_id] * self.dllm_config.block_size + array("q", [self.dllm_config.mask_id] * self.dllm_config.block_size)
) )
def _update_block_offset_for_dllm(self): def _update_block_offset_for_dllm(self):
python/sglang/srt/dllm/mixin/scheduler.py
+2 -1
View File
@@ -1,6 +1,7 @@
from __future__ import annotations from __future__ import annotations
import logging import logging
from array import array
from typing import TYPE_CHECKING, List, Optional, Set, Union from typing import TYPE_CHECKING, List, Optional, Set, Union
from sglang.srt.dllm.config import DllmConfig from sglang.srt.dllm.config import DllmConfig
@@ -79,7 +80,7 @@ class SchedulerDllmMixin:
if new_tokens == 0: if new_tokens == 0:
continue continue
req.fill_ids[-new_tokens:] = next_token_ids[:] req.fill_ids[-new_tokens:] = array("q", next_token_ids)
self.metrics_reporter.num_generated_tokens += new_tokens self.metrics_reporter.num_generated_tokens += new_tokens
req.output_ids.extend(next_token_ids) req.output_ids.extend(next_token_ids)
python/sglang/srt/managers/detokenizer_manager.py
+1 -1
View File
@@ -238,7 +238,7 @@ class DetokenizerManager(MultiHttpWorkerDetokenizerMixin):
if rid not in self.decode_status: if rid not in self.decode_status:
s = DecodeStatus( s = DecodeStatus(
decoded_text=recv_obj.decoded_texts[i], decoded_text=recv_obj.decoded_texts[i],
decode_ids=recv_obj.decode_ids[i], decode_ids=list(recv_obj.decode_ids[i]),
surr_offset=0, surr_offset=0,
read_offset=recv_obj.read_offsets[i], read_offset=recv_obj.read_offsets[i],
) )
python/sglang/srt/managers/io_struct.py
+5 -4
View File
@@ -21,6 +21,7 @@ from __future__ import annotations
import copy import copy
import uuid import uuid
from abc import ABC from abc import ABC
from array import array
from collections import Counter from collections import Counter
from dataclasses import dataclass, field from dataclasses import dataclass, field
from enum import Enum from enum import Enum
@@ -712,7 +713,7 @@ class TokenizedGenerateReqInput(BaseReq):
# The input text # The input text
input_text: str input_text: str
# The input token ids # The input token ids
input_ids: List[int] input_ids: Optional[array[int]]
# The multimodal inputs # The multimodal inputs
mm_inputs: object mm_inputs: object
# The sampling parameters # The sampling parameters
@@ -1027,7 +1028,7 @@ class TokenizedEmbeddingReqInput(BaseReq):
# The input text # The input text
input_text: str input_text: str
# The input token ids # The input token ids
input_ids: List[int] input_ids: array[int]
# The image inputs # The image inputs
image_inputs: dict image_inputs: dict
# The token type ids # The token type ids
@@ -1075,10 +1076,10 @@ class BatchTokenIDOutput(BaseBatchReq, SpeculativeDecodingMetricsMixin):
finished_reasons: List[BaseFinishReason] finished_reasons: List[BaseFinishReason]
# For incremental decoding # For incremental decoding
decoded_texts: List[str] decoded_texts: List[str]
decode_ids: List[int] decode_ids: List[array[int]]
read_offsets: List[int] read_offsets: List[int]
# Only used when `--skip-tokenizer-init` is on # Only used when `--skip-tokenizer-init` is on
output_ids: Optional[List[int]] output_ids: Optional[List[array[int]]]
# Detokenization configs # Detokenization configs
skip_special_tokens: List[bool] skip_special_tokens: List[bool]
spaces_between_special_tokens: List[bool] spaces_between_special_tokens: List[bool]
python/sglang/srt/managers/schedule_batch.py
+22 -17
View File
@@ -2,7 +2,11 @@ from __future__ import annotations
from sglang.srt.dllm.config import DllmConfig from sglang.srt.dllm.config import DllmConfig
from sglang.srt.model_executor.forward_batch_info import ForwardBatch from sglang.srt.model_executor.forward_batch_info import ForwardBatch
from sglang.srt.utils.common import ceil_align, is_pin_memory_available from sglang.srt.utils.common import (
ceil_align,
flatten_arrays_to_int64_tensor,
is_pin_memory_available,
)
# Copyright 2023-2024 SGLang Team # Copyright 2023-2024 SGLang Team
# Licensed under the Apache License, Version 2.0 (the "License"); # Licensed under the Apache License, Version 2.0 (the "License");
@@ -35,11 +39,11 @@ import copy
import dataclasses import dataclasses
import logging import logging
import re import re
from array import array
from concurrent.futures import Future from concurrent.futures import Future
from enum import Enum, auto from enum import Enum, auto
from functools import lru_cache from functools import lru_cache
from http import HTTPStatus from http import HTTPStatus
from itertools import chain
from typing import ( from typing import (
TYPE_CHECKING, TYPE_CHECKING,
Any, Any,
@@ -611,14 +615,14 @@ class Req(ReqDllmMixin):
self, self,
rid: str, rid: str,
origin_input_text: str, origin_input_text: str,
origin_input_ids: List[int], origin_input_ids: array[int],
sampling_params: SamplingParams, sampling_params: SamplingParams,
return_logprob: bool = False, return_logprob: bool = False,
top_logprobs_num: int = 0, top_logprobs_num: int = 0,
dllm_config: Optional[DllmConfig] = None, dllm_config: Optional[DllmConfig] = None,
token_ids_logprob: List[int] = None, token_ids_logprob: List[int] = None,
stream: bool = False, stream: bool = False,
origin_input_ids_unpadded: Optional[Tuple[int]] = None, origin_input_ids_unpadded: Optional[array[int]] = None,
lora_id: Optional[str] = None, lora_id: Optional[str] = None,
input_embeds: Optional[List[List[float]]] = None, input_embeds: Optional[List[List[float]]] = None,
positional_embed_overrides: Optional[PositionalEmbeds] = None, positional_embed_overrides: Optional[PositionalEmbeds] = None,
@@ -659,9 +663,10 @@ class Req(ReqDllmMixin):
) )
self.origin_input_ids = origin_input_ids self.origin_input_ids = origin_input_ids
# Each decode stage's output ids # Each decode stage's output ids
self.output_ids = [] self.output_ids = array("q")
# fill_ids = origin_input_ids + output_ids. Updated if chunked. # fill_ids = origin_input_ids + output_ids. Updated if chunked.
self.fill_ids = [] self.fill_ids = array("q")
self.session = session self.session = session
self.input_embeds = input_embeds self.input_embeds = input_embeds
self.positional_embed_overrides = positional_embed_overrides self.positional_embed_overrides = positional_embed_overrides
@@ -948,7 +953,7 @@ class Req(ReqDllmMixin):
return self.sampling_params.max_new_tokens == 0 and spec_alg is None return self.sampling_params.max_new_tokens == 0 and spec_alg is None
@property @property
def output_ids_through_stop(self) -> List[int]: def output_ids_through_stop(self) -> array[int]:
"""Get the output ids through the stop condition. Stop position is included.""" """Get the output ids through the stop condition. Stop position is included."""
if self.finished_len is not None: if self.finished_len is not None:
return self.output_ids[: self.finished_len] return self.output_ids[: self.finished_len]
@@ -1037,7 +1042,7 @@ class Req(ReqDllmMixin):
# Disable prefix caching when embed overrides are present: same token IDs # Disable prefix caching when embed overrides are present: same token IDs
# with different override vectors must not share cached KV values. # with different override vectors must not share cached KV values.
if self.positional_embed_overrides is not None: if self.positional_embed_overrides is not None:
token_ids_to_match = [] token_ids_to_match = array("q")
if tree_cache is not None: if tree_cache is not None:
if cow_mamba is None: if cow_mamba is None:
@@ -1305,7 +1310,7 @@ class Req(ReqDllmMixin):
# Therefore, we discard the generated output_ids and restart prefill and generation # Therefore, we discard the generated output_ids and restart prefill and generation
# to ensure shape consistency in KV cache. # to ensure shape consistency in KV cache.
if self.input_embeds is not None: if self.input_embeds is not None:
self.output_ids = [] self.output_ids = array("q")
def offload_kv_cache(self, req_to_token_pool, token_to_kv_pool_allocator): def offload_kv_cache(self, req_to_token_pool, token_to_kv_pool_allocator):
token_indices = req_to_token_pool.req_to_token[ token_indices = req_to_token_pool.req_to_token[
@@ -1370,7 +1375,9 @@ class Req(ReqDllmMixin):
logger.error(f"{error_msg}, {self.rid=}") logger.error(f"{error_msg}, {self.rid=}")
self.multimodal_inputs = None self.multimodal_inputs = None
self.grammar = None self.grammar = None
self.origin_input_ids = [0] # set it to one token to skip the long prefill self.origin_input_ids = array(
"q", [0]
) # set it to one token to skip the long prefill
self.return_logprob = False self.return_logprob = False
self.logprob_start_len = -1 self.logprob_start_len = -1
self.to_finish = FINISH_ABORT( self.to_finish = FINISH_ABORT(
@@ -1629,7 +1636,9 @@ class ScheduleBatch(ScheduleBatchDisaggregationDecodeMixin):
def is_dllm(self): def is_dllm(self):
return self.dllm_config is not None return self.dllm_config is not None
def prepare_encoder_info_extend(self, input_ids: List[int], seq_lens: List[int]): def prepare_encoder_info_extend(
self, input_ids: List[array[int]], seq_lens: List[int]
):
_pin = is_pin_memory_available(self.device) _pin = is_pin_memory_available(self.device)
self.encoder_lens_cpu = [] self.encoder_lens_cpu = []
self.encoder_cached = [] self.encoder_cached = []
@@ -1678,9 +1687,7 @@ class ScheduleBatch(ScheduleBatchDisaggregationDecodeMixin):
pt += req.extend_input_len pt += req.extend_input_len
# Reassign # Reassign
self.input_ids = torch.tensor( self.input_ids = flatten_arrays_to_int64_tensor(input_ids, self.device, _pin)
sum(input_ids, []), dtype=torch.int64, pin_memory=_pin
).to(self.device, non_blocking=True)
self.seq_lens = torch.tensor(seq_lens, dtype=torch.int64, pin_memory=_pin).to( self.seq_lens = torch.tensor(seq_lens, dtype=torch.int64, pin_memory=_pin).to(
self.device, non_blocking=True self.device, non_blocking=True
) )
@@ -1783,9 +1790,7 @@ class ScheduleBatch(ScheduleBatchDisaggregationDecodeMixin):
] ]
_pin = is_pin_memory_available(self.device) _pin = is_pin_memory_available(self.device)
input_ids_tensor = torch.tensor( input_ids_tensor = flatten_arrays_to_int64_tensor(input_ids, self.device, _pin)
list(chain.from_iterable(input_ids)), dtype=torch.int64, pin_memory=_pin
).to(self.device, non_blocking=True)
seq_lens_tensor = torch.tensor(seq_lens, dtype=torch.int64, pin_memory=_pin).to( seq_lens_tensor = torch.tensor(seq_lens, dtype=torch.int64, pin_memory=_pin).to(
self.device, non_blocking=True self.device, non_blocking=True
) )
python/sglang/srt/managers/schedule_policy.py
+2 -1
View File
@@ -1,6 +1,7 @@
from __future__ import annotations from __future__ import annotations
import logging import logging
from array import array
from sglang.srt.environ import envs from sglang.srt.environ import envs
from sglang.srt.managers.prefill_delayer import PrefillDelayerSinglePassExecutor from sglang.srt.managers.prefill_delayer import PrefillDelayerSinglePassExecutor
@@ -84,7 +85,7 @@ IGNORE_EOS_RESERVE_TOKENS = 1
def match_prefix_for_req( def match_prefix_for_req(
tree_cache: BasePrefixCache, tree_cache: BasePrefixCache,
req: Req, req: Req,
token_ids: Optional[List[int]] = None, token_ids: Optional[array[int]] = None,
*, *,
cow_mamba: bool = False, cow_mamba: bool = False,
include_req: bool = False, include_req: bool = False,
python/sglang/srt/managers/scheduler.py
+8 -6
View File
@@ -20,6 +20,7 @@ import os
import signal import signal
import sys import sys
import time import time
from array import array
from collections import deque from collections import deque
from contextlib import contextmanager, nullcontext from contextlib import contextmanager, nullcontext
from functools import partial from functools import partial
@@ -959,6 +960,7 @@ class Scheduler(
) )
self.dp_tp_cpu_group = self.dp_tp_group.cpu_group self.dp_tp_cpu_group = self.dp_tp_group.cpu_group
# TODO(Jialin): Migrate pad_input_ids implementations to return array.
self.pad_input_ids_func = self.tp_worker.get_pad_input_ids_func() self.pad_input_ids_func = self.tp_worker.get_pad_input_ids_func()
set_random_seed(self.random_seed) set_random_seed(self.random_seed)
@@ -1782,8 +1784,7 @@ class Scheduler(
if recv_req.input_embeds is not None: if recv_req.input_embeds is not None:
# Generate fake input_ids based on the length of input_embeds # Generate fake input_ids based on the length of input_embeds
seq_length = len(recv_req.input_embeds) seq_length = len(recv_req.input_embeds)
fake_input_ids = [1] * seq_length recv_req.input_ids = array("q", [1]) * seq_length
recv_req.input_ids = fake_input_ids
if recv_req.bootstrap_port is None: if recv_req.bootstrap_port is None:
# Use default bootstrap port # Use default bootstrap port
@@ -1909,8 +1910,8 @@ class Scheduler(
# Expand a single image token into multiple dummy tokens for receiving image embeddings. # Expand a single image token into multiple dummy tokens for receiving image embeddings.
# The pad function is model-specific and can be None for some backends. # The pad function is model-specific and can be None for some backends.
if self.pad_input_ids_func: if self.pad_input_ids_func:
req.origin_input_ids = self.pad_input_ids_func( req.origin_input_ids = array(
req.origin_input_ids, image_inputs "q", self.pad_input_ids_func(req.origin_input_ids, image_inputs)
) )
req.extend_image_inputs(image_inputs) req.extend_image_inputs(image_inputs)
self._maybe_compute_mrope_positions(req) self._maybe_compute_mrope_positions(req)
@@ -2182,8 +2183,9 @@ class Scheduler(
# embedding models or models not requiring special padding. # embedding models or models not requiring special padding.
# If None, `req.origin_input_ids` is expected to be correctly populated already. # If None, `req.origin_input_ids` is expected to be correctly populated already.
if self.pad_input_ids_func: if self.pad_input_ids_func:
req.origin_input_ids = self.pad_input_ids_func( # See companion call site above for the array.array wrap rationale.
req.origin_input_ids, image_inputs req.origin_input_ids = array(
"q", self.pad_input_ids_func(req.origin_input_ids, image_inputs)
) )
req.extend_image_inputs(image_inputs) req.extend_image_inputs(image_inputs)
python/sglang/srt/managers/scheduler_pp_mixin.py
+8 -4
View File
@@ -3,6 +3,7 @@ from __future__ import annotations
import logging import logging
import math import math
import time import time
from array import array
from collections import defaultdict, deque from collections import defaultdict, deque
from dataclasses import dataclass from dataclasses import dataclass
from typing import TYPE_CHECKING, Dict, List, Optional, Tuple from typing import TYPE_CHECKING, Dict, List, Optional, Tuple
@@ -556,7 +557,7 @@ class SchedulerPPMixin:
if self.pp_group.is_first_rank: if self.pp_group.is_first_rank:
model_runner = self.tp_worker.model_runner model_runner = self.tp_worker.model_runner
model_config = model_runner.model_config model_config = model_runner.model_config
input_ids_list = [] input_ids_list: List[array[int]] = []
for i in range(128): for i in range(128):
chunk_size = int( chunk_size = int(
self.chunked_prefill_size * 1.25 self.chunked_prefill_size * 1.25
@@ -564,9 +565,12 @@ class SchedulerPPMixin:
) )
if chunk_size <= 0: if chunk_size <= 0:
break break
input_ids = np.random.randint( input_ids = array(
0, 10000, size=chunk_size, dtype=np.int64 "q",
).tolist() np.random.randint(
0, 10000, size=chunk_size, dtype=np.int64
).tobytes(),
)
input_ids_list.append(input_ids) input_ids_list.append(input_ids)
sampling_params = SamplingParams( sampling_params = SamplingParams(
python/sglang/srt/managers/tokenizer_manager.py
+13 -7
View File
@@ -13,6 +13,8 @@
# ============================================================================== # ==============================================================================
"""TokenizerManager is a process that tokenizes the text.""" """TokenizerManager is a process that tokenizes the text."""
from __future__ import annotations
import asyncio import asyncio
import copy import copy
import dataclasses import dataclasses
@@ -24,6 +26,7 @@ import signal
import socket import socket
import sys import sys
import threading import threading
from array import array
from collections import deque from collections import deque
from contextlib import nullcontext from contextlib import nullcontext
from datetime import datetime from datetime import datetime
@@ -987,12 +990,15 @@ class TokenizerManager(TokenizerControlMixin, TokenizerManagerScoreMixin):
self, self,
obj: Union[GenerateReqInput, EmbeddingReqInput], obj: Union[GenerateReqInput, EmbeddingReqInput],
input_text: str, input_text: str,
input_ids: List[int], input_ids: Optional[List[int]],
input_embeds: Optional[Union[List[float], None]] = None, input_embeds: Optional[Union[List[float], None]] = None,
mm_inputs=None, mm_inputs=None,
token_type_ids: Optional[List[int]] = None, token_type_ids: Optional[List[int]] = None,
) -> Union[TokenizedGenerateReqInput, TokenizedEmbeddingReqInput]: ) -> Union[TokenizedGenerateReqInput, TokenizedEmbeddingReqInput]:
"""Create a tokenized request object from common parameters.""" """Create a tokenized request object from common parameters."""
input_ids_arr: Optional[array[int]] = (
array("q", input_ids) if input_ids is not None else None
)
# Parse sampling parameters # Parse sampling parameters
# Note: if there are preferred sampling params, we use them if they are not # Note: if there are preferred sampling params, we use them if they are not
# explicitly passed in sampling_params # explicitly passed in sampling_params
@@ -1020,7 +1026,7 @@ class TokenizerManager(TokenizerControlMixin, TokenizerManagerScoreMixin):
tokenized_obj = TokenizedGenerateReqInput( tokenized_obj = TokenizedGenerateReqInput(
input_text, input_text,
input_ids, input_ids_arr,
mm_inputs, mm_inputs,
sampling_params, sampling_params,
obj.return_logprob, obj.return_logprob,
@@ -1062,12 +1068,12 @@ class TokenizerManager(TokenizerControlMixin, TokenizerManagerScoreMixin):
and obj.embed_override_token_id is not None and obj.embed_override_token_id is not None
): ):
positional_embed_overrides = self._resolve_embed_overrides( positional_embed_overrides = self._resolve_embed_overrides(
input_ids, obj.embed_override_token_id, obj.embed_overrides input_ids_arr, obj.embed_override_token_id, obj.embed_overrides
) )
tokenized_obj = TokenizedEmbeddingReqInput( tokenized_obj = TokenizedEmbeddingReqInput(
input_text, input_text,
input_ids, input_ids_arr,
mm_inputs, mm_inputs,
token_type_ids, token_type_ids,
sampling_params, sampling_params,
@@ -1088,7 +1094,7 @@ class TokenizerManager(TokenizerControlMixin, TokenizerManagerScoreMixin):
@staticmethod @staticmethod
def _resolve_embed_overrides( def _resolve_embed_overrides(
input_ids: List[int], input_ids: array[int],
token_id: int, token_id: int,
embeds: List[torch.Tensor], embeds: List[torch.Tensor],
) -> PositionalEmbeds: ) -> PositionalEmbeds:
@@ -1787,7 +1793,7 @@ class TokenizerManager(TokenizerControlMixin, TokenizerManagerScoreMixin):
self.server_args.incremental_streaming_output and is_stream self.server_args.incremental_streaming_output and is_stream
) )
delta_text = recv_obj.output_strs[i] delta_text = recv_obj.output_strs[i]
delta_output_ids = recv_obj.output_ids[i] delta_output_ids = list(recv_obj.output_ids[i])
output_offset = state.last_output_offset output_offset = state.last_output_offset
state.append_text(delta_text) state.append_text(delta_text)
state.output_ids.extend(delta_output_ids) state.output_ids.extend(delta_output_ids)
@@ -1830,7 +1836,7 @@ class TokenizerManager(TokenizerControlMixin, TokenizerManagerScoreMixin):
incremental = ( incremental = (
self.server_args.incremental_streaming_output and is_stream self.server_args.incremental_streaming_output and is_stream
) )
delta_output_ids = recv_obj.output_ids[i] delta_output_ids = list(recv_obj.output_ids[i])
output_offset = state.last_output_offset output_offset = state.last_output_offset
state.output_ids.extend(delta_output_ids) state.output_ids.extend(delta_output_ids)
python/sglang/srt/mem_cache/events.py
+1 -1
View File
@@ -65,7 +65,7 @@ class KVCacheEventMixin:
if is_bigram: if is_bigram:
page_tokens = [(raw[j], raw[j + 1]) for j in range(start, end)] page_tokens = [(raw[j], raw[j + 1]) for j in range(start, end)]
else: else:
page_tokens = raw[start:end] page_tokens = list(raw[start:end])
block_hash = hash_str_to_int64(node.hash_value[page_index]) block_hash = hash_str_to_int64(node.hash_value[page_index])
python/sglang/srt/mem_cache/hiradix_cache.py
-4
View File
@@ -804,10 +804,6 @@ class HiRadixCache(RadixCache):
def evictable_size(self): def evictable_size(self):
return self.evictable_size_ return self.evictable_size_
def _to_radix_key(self, token_ids: List[int]) -> RadixKey:
"""Convert raw token_ids to a RadixKey; must be list (not tuple) for paged match."""
return RadixKey(token_ids=list(token_ids))
def inc_lock_ref(self, node: TreeNode) -> IncLockRefResult: def inc_lock_ref(self, node: TreeNode) -> IncLockRefResult:
if self.disable: if self.disable:
return IncLockRefResult(delta=0) return IncLockRefResult(delta=0)
python/sglang/srt/mem_cache/mamba_radix_cache.py
+2 -1
View File
@@ -20,6 +20,7 @@ The radix tree data structure for managing the hybrid (full and Mamba) KV cache.
""" """
import heapq import heapq
from array import array
from collections import defaultdict from collections import defaultdict
from functools import lru_cache from functools import lru_cache
from typing import TYPE_CHECKING, List, Optional, Tuple from typing import TYPE_CHECKING, List, Optional, Tuple
@@ -456,7 +457,7 @@ class MambaRadixCache(KVCacheEventMixin, BasePrefixCache):
def reset(self) -> None: def reset(self) -> None:
self.root_node = TreeNode() self.root_node = TreeNode()
self.root_node.key = RadixKey([], None) self.root_node.key = RadixKey(array("q"), None)
self.root_node.value = [] self.root_node.value = []
self.root_node.hash_value = [] self.root_node.hash_value = []
self.root_node.full_lock_ref = 1 self.root_node.full_lock_ref = 1
python/sglang/srt/mem_cache/radix_cache.py
+12 -14
View File
@@ -26,6 +26,7 @@ import heapq
import logging import logging
import sys import sys
import time import time
from array import array
from collections import defaultdict from collections import defaultdict
from functools import lru_cache from functools import lru_cache
from typing import TYPE_CHECKING, Any, Iterator, List, Optional, Tuple, Union from typing import TYPE_CHECKING, Any, Iterator, List, Optional, Tuple, Union
@@ -70,7 +71,7 @@ class RadixKey:
def __init__( def __init__(
self, self,
token_ids: List[int], token_ids: array[int],
extra_key: Optional[str] = None, extra_key: Optional[str] = None,
is_bigram: bool = False, is_bigram: bool = False,
): ):
@@ -87,6 +88,7 @@ class RadixKey:
return n - 1 if n > 0 else 0 return n - 1 if n > 0 else 0
return len(self.token_ids) return len(self.token_ids)
# TODO(Jialin): vectorize with numpy without PyLong boxing
def __iter__(self) -> Iterator: def __iter__(self) -> Iterator:
if self.is_bigram: if self.is_bigram:
t = self.token_ids t = self.token_ids
@@ -110,7 +112,7 @@ class RadixKey:
if self.is_bigram: if self.is_bigram:
# bigrams [start, stop) span raw tokens [start, stop + 1); # bigrams [start, stop) span raw tokens [start, stop + 1);
# empty slice -> empty raw tokens (not a dangling boundary token). # empty slice -> empty raw tokens (not a dangling boundary token).
raw = self.token_ids[start : stop + 1] if stop > start else [] raw = self.token_ids[start : stop + 1] if stop > start else array("q")
return RadixKey(raw, self.extra_key, is_bigram=True) return RadixKey(raw, self.extra_key, is_bigram=True)
return RadixKey(self.token_ids[start:stop], self.extra_key) return RadixKey(self.token_ids[start:stop], self.extra_key)
@@ -144,6 +146,7 @@ class RadixKey:
f"{self.extra_key=} != {other.extra_key=}" f"{self.extra_key=} != {other.extra_key=}"
) )
# TODO(Jialin): replace zip with numpy to skip per-element PyLong boxing
def match(self, other: "RadixKey", page_size: int = 1) -> int: def match(self, other: "RadixKey", page_size: int = 1) -> int:
"""Logical-unit prefix length shared with ``other``. Result is rounded down to ``page_size``.""" """Logical-unit prefix length shared with ``other``. Result is rounded down to ``page_size``."""
self._check_compatible(other) self._check_compatible(other)
@@ -337,7 +340,7 @@ class RadixCache(KVCacheEventMixin, BasePrefixCache):
def reset(self): def reset(self):
# Initialize root with minimum priority so any real priority overrides it # Initialize root with minimum priority so any real priority overrides it
self.root_node = TreeNode(priority=-sys.maxsize) self.root_node = TreeNode(priority=-sys.maxsize)
self.root_node.key = RadixKey(token_ids=[], extra_key=None) self.root_node.key = RadixKey(token_ids=array("q"), extra_key=None)
self.root_node.value = [] self.root_node.value = []
self.root_node.host_value = [] self.root_node.host_value = []
self.root_node.lock_ref = 1 self.root_node.lock_ref = 1
@@ -811,20 +814,15 @@ class RadixCache(KVCacheEventMixin, BasePrefixCache):
if __name__ == "__main__": if __name__ == "__main__":
tree = RadixCache.create_simulated() tree = RadixCache.create_simulated()
# Example token id sequences (as lists of ints) tree.insert(InsertParams(key=RadixKey(token_ids=array("q", [1, 2, 3]))))
tree.insert(InsertParams(key=RadixKey(token_ids=[1, 2, 3], extra_key=None))) tree.insert(InsertParams(key=RadixKey(token_ids=array("q", [1, 2, 3]))))
tree.insert(InsertParams(key=RadixKey(token_ids=[1, 2, 3], extra_key=None))) tree.insert(InsertParams(key=RadixKey(token_ids=array("q", [1, 2, 4, 5]))))
tree.insert(InsertParams(key=RadixKey(token_ids=[1, 2, 4, 5], extra_key=None))) tree.insert(InsertParams(key=RadixKey(token_ids=array("q", [1, 2, 4, 5, 6, 7]))))
tree.insert( tree.insert(InsertParams(key=RadixKey(token_ids=array("q", [8, 9, 10, 11, 12]))))
InsertParams(key=RadixKey(token_ids=[1, 2, 4, 5, 6, 7], extra_key=None))
)
tree.insert(
InsertParams(key=RadixKey(token_ids=[8, 9, 10, 11, 12], extra_key=None))
)
tree.pretty_print() tree.pretty_print()
print( print(
tree.match_prefix( tree.match_prefix(
MatchPrefixParams(key=RadixKey(token_ids=[1, 2, 3, 13, 14], extra_key=None)) MatchPrefixParams(key=RadixKey(token_ids=array("q", [1, 2, 3, 13, 14])))
) )
) )
python/sglang/srt/mem_cache/unified_radix_cache.py
+2 -1
View File
@@ -3,6 +3,7 @@ from __future__ import annotations
import logging import logging
import threading import threading
import time import time
from array import array
from collections import defaultdict from collections import defaultdict
from functools import partial from functools import partial
from typing import TYPE_CHECKING, Any, Optional from typing import TYPE_CHECKING, Any, Optional
@@ -257,7 +258,7 @@ class UnifiedRadixCache(BasePrefixCache):
def _reset_full(self) -> None: def _reset_full(self) -> None:
"""Full reset: destroy entire tree and all state.""" """Full reset: destroy entire tree and all state."""
self.root_node = UnifiedTreeNode(self.tree_components) self.root_node = UnifiedTreeNode(self.tree_components)
self.root_node.key = RadixKey([], None) self.root_node.key = RadixKey(array("q"), None)
self.root_node.component_data[BASE_COMPONENT_TYPE].value = [] self.root_node.component_data[BASE_COMPONENT_TYPE].value = []
for ct in self.tree_components: for ct in self.tree_components:
self.root_node.component_data[ct].lock_ref = 1 self.root_node.component_data[ct].lock_ref = 1
python/sglang/srt/models/llava.py
+8 -2
View File
@@ -13,8 +13,11 @@
# ============================================================================== # ==============================================================================
"""Inference-only LLaVa model compatible with HuggingFace weights.""" """Inference-only LLaVa model compatible with HuggingFace weights."""
from __future__ import annotations
import math import math
import re import re
from array import array
from functools import lru_cache from functools import lru_cache
from typing import Dict, Iterable, List, Optional, Tuple, Type, Union from typing import Dict, Iterable, List, Optional, Tuple, Type, Union
@@ -73,7 +76,9 @@ class LlavaBaseForCausalLM(nn.Module):
return "pad" return "pad"
return "anyres" return "anyres"
def pad_input_ids(self, input_ids: List[int], image_inputs: MultimodalInputs): def pad_input_ids(
self, input_ids: array[int], image_inputs: MultimodalInputs
) -> array[int]:
image_sizes = flatten_nested_list( image_sizes = flatten_nested_list(
[item.image_sizes for item in image_inputs.mm_items] [item.image_sizes for item in image_inputs.mm_items]
) )
@@ -125,9 +130,10 @@ class LlavaBaseForCausalLM(nn.Module):
except ValueError: except ValueError:
offset = 0 offset = 0
# old_len + pad_len - 1, because we need to remove image_token_id # old_len + pad_len - 1, because we need to remove image_token_id
pad_token = pad_values[image_idx % len(pad_values)]
input_ids = ( input_ids = (
input_ids[:offset] input_ids[:offset]
+ [pad_values[image_idx % len(pad_values)]] * new_image_feature_len + array("q", [pad_token]) * new_image_feature_len
+ input_ids[offset + 1 :] + input_ids[offset + 1 :]
) )
offset_list.append(offset) offset_list.append(offset)
python/sglang/srt/models/llavavid.py
+8 -3
View File
@@ -13,7 +13,10 @@
# ============================================================================== # ==============================================================================
"""Inference-only LLaVa video model compatible with HuggingFace weights.""" """Inference-only LLaVa video model compatible with HuggingFace weights."""
from typing import Iterable, List, Optional, Tuple from __future__ import annotations
from array import array
from typing import Iterable, Optional, Tuple
import numpy as np import numpy as np
import torch import torch
@@ -57,8 +60,10 @@ class LlavaVidForCausalLM(nn.Module):
torch.empty(config.text_config.hidden_size, dtype=torch.float16) torch.empty(config.text_config.hidden_size, dtype=torch.float16)
) )
def pad_input_ids(self, input_ids: List[int], image_inputs: MultimodalInputs): def pad_input_ids(
pad_values = [item.pad_value for item in image_inputs.mm_items] self, input_ids: array[int], image_inputs: MultimodalInputs
) -> array[int]:
pad_values = array("q", (item.pad_value for item in image_inputs.mm_items))
new_image_feature_len = self.image_feature_len new_image_feature_len = self.image_feature_len
pad_ids = pad_values * ( pad_ids = pad_values * (
python/sglang/srt/models/mllama.py
+7 -2
View File
@@ -4,7 +4,10 @@
# https://github.com/vllm-project/vllm/blob/7193774b1ff8603ad5bf4598e5efba0d9a39b436/vllm/model_executor/models/mllama.py # https://github.com/vllm-project/vllm/blob/7193774b1ff8603ad5bf4598e5efba0d9a39b436/vllm/model_executor/models/mllama.py
"""PyTorch Mllama model.""" """PyTorch Mllama model."""
from __future__ import annotations
import math import math
from array import array
from typing import Iterable, List, Optional, Tuple, Union from typing import Iterable, List, Optional, Tuple, Union
import torch import torch
@@ -823,9 +826,11 @@ class MllamaForConditionalGeneration(nn.Module):
) )
self.logits_processor = LogitsProcessor(config.text_config) self.logits_processor = LogitsProcessor(config.text_config)
def pad_input_ids(self, input_ids: List[int], mm_inputs: MultimodalInputs): def pad_input_ids(
self, input_ids: array[int], mm_inputs: MultimodalInputs
) -> array[int]:
pixel_values = torch.cat([item.feature for item in mm_inputs.mm_items], dim=0) pixel_values = torch.cat([item.feature for item in mm_inputs.mm_items], dim=0)
pad_values = [item.pad_value for item in mm_inputs.mm_items] pad_values = array("q", (item.pad_value for item in mm_inputs.mm_items))
num_concurrent_media, num_tiles = pixel_values.shape[1:3] num_concurrent_media, num_tiles = pixel_values.shape[1:3]
num_patches = self.vision_model.num_patches num_patches = self.vision_model.num_patches
python/sglang/srt/models/moss_vl.py
+7 -4
View File
@@ -3,6 +3,7 @@
from __future__ import annotations from __future__ import annotations
import logging import logging
from array import array
from functools import partial from functools import partial
from typing import Iterable, List, Optional, Tuple from typing import Iterable, List, Optional, Tuple
@@ -1122,15 +1123,17 @@ class MossVLForConditionalGeneration(nn.Module):
return total_len return total_len
def _build_encoder_prefix_pad_ids(self, mm_inputs: MultimodalInputs) -> List[int]: def _build_encoder_prefix_pad_ids(self, mm_inputs: MultimodalInputs) -> array[int]:
encoder_len = self._get_encoder_len(mm_inputs) encoder_len = self._get_encoder_len(mm_inputs)
if encoder_len == 0 or not mm_inputs.mm_items: if encoder_len == 0 or not mm_inputs.mm_items:
return [] return array("q")
pad_value = mm_inputs.mm_items[0].pad_value pad_value = mm_inputs.mm_items[0].pad_value
return [pad_value] * encoder_len return array("q", [pad_value]) * encoder_len
def pad_input_ids(self, input_ids: List[int], mm_inputs: MultimodalInputs): def pad_input_ids(
self, input_ids: array[int], mm_inputs: MultimodalInputs
) -> array[int]:
encoder_len = self._get_encoder_len(mm_inputs) encoder_len = self._get_encoder_len(mm_inputs)
mm_inputs.num_image_tokens = encoder_len mm_inputs.num_image_tokens = encoder_len
if encoder_len == 0: if encoder_len == 0:
python/sglang/srt/models/whisper.py
+8 -4
View File
@@ -1,4 +1,7 @@
from typing import Any, Iterable, List, Optional, Tuple from __future__ import annotations
from array import array
from typing import Any, Iterable, Optional, Tuple
import torch import torch
from transformers import WhisperConfig from transformers import WhisperConfig
@@ -418,14 +421,15 @@ class WhisperForConditionalGeneration(torch.nn.Module):
weight_loader = getattr(param, "weight_loader", default_weight_loader) weight_loader = getattr(param, "weight_loader", default_weight_loader)
weight_loader(param, loaded_weight) weight_loader(param, loaded_weight)
def pad_input_ids(self, input_ids: List[int], mm_inputs: MultimodalInputs): def pad_input_ids(
self, input_ids: array[int], mm_inputs: MultimodalInputs
) -> array[int]:
# Prepend dummy encoder tokens so that prepare_encoder_info_extend # Prepend dummy encoder tokens so that prepare_encoder_info_extend
# correctly allocates encoder KV cache locations in the KV pool. # correctly allocates encoder KV cache locations in the KV pool.
# These dummy tokens are stripped before the model forward receives input_ids. # These dummy tokens are stripped before the model forward receives input_ids.
encoder_len = self.config.max_source_positions encoder_len = self.config.max_source_positions
mm_inputs.num_image_tokens = encoder_len mm_inputs.num_image_tokens = encoder_len
pad_ids = [0] * encoder_len return array("q", [0]) * encoder_len + input_ids
return pad_ids + input_ids
def forward( def forward(
self, self,
python/sglang/srt/utils/common.py
+16
View File
@@ -45,6 +45,7 @@ import traceback
import types import types
import uuid import uuid
import warnings import warnings
from array import array
from collections import OrderedDict, defaultdict from collections import OrderedDict, defaultdict
from contextlib import contextmanager from contextlib import contextmanager
from dataclasses import dataclass from dataclasses import dataclass
@@ -101,6 +102,21 @@ logger = logging.getLogger(__name__)
torch_release = pkg_version.parse(torch.__version__).release torch_release = pkg_version.parse(torch.__version__).release
def flatten_arrays_to_int64_tensor(
parts: List[array[int]], device, pin: bool
) -> torch.Tensor:
"""Flatten a list of array.array('q') buffers into one int64 tensor.
Uses NumPy here to speed up the conversion by using memcpy
instead of a per-element PyLong-to-int64 walk.
"""
combined = np.concatenate([np.frombuffer(p, dtype=np.int64) for p in parts])
cpu_t = torch.from_numpy(combined)
if pin:
cpu_t = cpu_t.pin_memory()
return cpu_t.to(device, non_blocking=True)
# https://pytorch.org/docs/stable/notes/hip.html#checking-for-hip # https://pytorch.org/docs/stable/notes/hip.html#checking-for-hip
@lru_cache(maxsize=1) @lru_cache(maxsize=1)
def is_hip() -> bool: def is_hip() -> bool:
test/registered/unit/managers/test_scheduler_chunked_req_gate.py
+4 -3
View File
@@ -1,6 +1,7 @@
"""Regression tests for the SWA chunked-req stash gate (#24252).""" """Regression tests for the SWA chunked-req stash gate (#24252)."""
import unittest import unittest
from array import array
from types import SimpleNamespace from types import SimpleNamespace
from unittest.mock import MagicMock from unittest.mock import MagicMock
@@ -27,9 +28,9 @@ def _make_req(
) -> Req: ) -> Req:
req = Req.__new__(Req) req = Req.__new__(Req)
req.rid = "test-req" req.rid = "test-req"
req.origin_input_ids = list(fill_ids) req.origin_input_ids = array("q", fill_ids)
req.output_ids = [] req.output_ids = array("q")
req.fill_ids = list(fill_ids) req.fill_ids = array("q", fill_ids)
req.prefix_indices = prefix_indices req.prefix_indices = prefix_indices
req.req_pool_idx = req_pool_idx req.req_pool_idx = req_pool_idx
req.extend_input_len = extend_input_len req.extend_input_len = extend_input_len
test/registered/unit/mem_cache/test_decode_radix_lock_ref.py
+17 -10
View File
@@ -26,6 +26,7 @@ register_cuda_ci(est_time=10, stage="base-b", runner_config="1-gpu-small")
register_amd_ci(est_time=10, suite="stage-b-test-1-gpu-small-amd") register_amd_ci(est_time=10, suite="stage-b-test-1-gpu-small-amd")
import unittest import unittest
from array import array
from unittest.mock import MagicMock from unittest.mock import MagicMock
import torch import torch
@@ -65,11 +66,11 @@ class MockReq:
"""Minimal mock Req with fields needed by cache_unfinished/finished_req.""" """Minimal mock Req with fields needed by cache_unfinished/finished_req."""
def __init__(self, fill_ids, req_pool_idx=0, cache_protected_len=0, last_node=None): def __init__(self, fill_ids, req_pool_idx=0, cache_protected_len=0, last_node=None):
self.fill_ids = list(fill_ids) self.fill_ids = array("q", fill_ids)
self.origin_input_ids = ( self.origin_input_ids = array(
list(fill_ids[:-1]) if len(fill_ids) > 1 else list(fill_ids) "q", fill_ids[:-1] if len(fill_ids) > 1 else fill_ids
) )
self.output_ids = [fill_ids[-1]] if len(fill_ids) > 1 else [] self.output_ids = array("q", [fill_ids[-1]] if len(fill_ids) > 1 else [])
self.req_pool_idx = req_pool_idx self.req_pool_idx = req_pool_idx
self.cache_protected_len = cache_protected_len self.cache_protected_len = cache_protected_len
self.last_node = last_node self.last_node = last_node
@@ -99,7 +100,7 @@ class TestDecodeLockRefScenarios(unittest.TestCase):
"""Insert a prefix into the tree so future requests can match it.""" """Insert a prefix into the tree so future requests can match it."""
cache.insert( cache.insert(
InsertParams( InsertParams(
key=RadixKey(prefix_ids), key=RadixKey(array("q", prefix_ids)),
value=torch.tensor(prefix_values, dtype=torch.int64), value=torch.tensor(prefix_values, dtype=torch.int64),
) )
) )
@@ -119,7 +120,7 @@ class TestDecodeLockRefScenarios(unittest.TestCase):
self._populate_prefix(cache, prefix, prefix_vals) self._populate_prefix(cache, prefix, prefix_vals)
# Match prefix (simulates _match_prefix_and_lock in pop_preallocated) # Match prefix (simulates _match_prefix_and_lock in pop_preallocated)
result = cache.match_prefix(MatchPrefixParams(key=RadixKey(prefix))) result = cache.match_prefix(MatchPrefixParams(key=RadixKey(array("q", prefix))))
matched_node = result.last_device_node matched_node = result.last_device_node
prefix_len = len(result.device_indices) prefix_len = len(result.device_indices)
self.assertEqual(prefix_len, 3) self.assertEqual(prefix_len, 3)
@@ -164,7 +165,9 @@ class TestDecodeLockRefScenarios(unittest.TestCase):
# No prefix in tree -- match returns root # No prefix in tree -- match returns root
full_ids = [10, 20, 30] full_ids = [10, 20, 30]
result = cache.match_prefix(MatchPrefixParams(key=RadixKey(full_ids))) result = cache.match_prefix(
MatchPrefixParams(key=RadixKey(array("q", full_ids)))
)
matched_node = result.last_device_node matched_node = result.last_device_node
self.assertEqual(len(result.device_indices), 0) # no match self.assertEqual(len(result.device_indices), 0) # no match
# matched_node is root # matched_node is root
@@ -212,7 +215,7 @@ class TestDecodeLockRefScenarios(unittest.TestCase):
self._populate_prefix(cache, prefix, prefix_vals) self._populate_prefix(cache, prefix, prefix_vals)
# Match and lock # Match and lock
result = cache.match_prefix(MatchPrefixParams(key=RadixKey(prefix))) result = cache.match_prefix(MatchPrefixParams(key=RadixKey(array("q", prefix))))
matched_node = result.last_device_node matched_node = result.last_device_node
prefix_len = len(result.device_indices) prefix_len = len(result.device_indices)
@@ -256,7 +259,9 @@ class TestDecodeLockRefScenarios(unittest.TestCase):
# No prefix in tree -- match returns root (simulates _match_prefix_and_lock) # No prefix in tree -- match returns root (simulates _match_prefix_and_lock)
full_ids = [10, 20, 30] full_ids = [10, 20, 30]
result = cache.match_prefix(MatchPrefixParams(key=RadixKey(full_ids))) result = cache.match_prefix(
MatchPrefixParams(key=RadixKey(array("q", full_ids)))
)
matched_node = result.last_device_node matched_node = result.last_device_node
self.assertIs(matched_node, cache.root_node) self.assertIs(matched_node, cache.root_node)
@@ -356,7 +361,9 @@ class TestDecodeLockRefScenarios(unittest.TestCase):
self._populate_prefix(cache, prefix, prefix_vals) self._populate_prefix(cache, prefix, prefix_vals)
for iteration in range(5): for iteration in range(5):
result = cache.match_prefix(MatchPrefixParams(key=RadixKey(prefix))) result = cache.match_prefix(
MatchPrefixParams(key=RadixKey(array("q", prefix)))
)
matched_node = result.last_device_node matched_node = result.last_device_node
prefix_len = len(result.device_indices) prefix_len = len(result.device_indices)
test/registered/unit/mem_cache/test_mamba_unittest.py
+34 -23
View File
@@ -1,4 +1,5 @@
import unittest import unittest
from array import array
import torch import torch
@@ -116,7 +117,7 @@ class TestMamba(unittest.TestCase):
req = Req( req = Req(
rid=0, rid=0,
origin_input_text="", origin_input_text="",
origin_input_ids=[], origin_input_ids=array("q"),
sampling_params=sampling_params, sampling_params=sampling_params,
) )
@@ -158,7 +159,7 @@ class TestMamba(unittest.TestCase):
print( print(
f"req1: inserting, req1_token_ids: {req1_token_ids}, req1_kv_indices: {req1_kv_indices}" f"req1: inserting, req1_token_ids: {req1_token_ids}, req1_kv_indices: {req1_kv_indices}"
) )
key = RadixKey(req1_token_ids) key = RadixKey(array("q", req1_token_ids))
result = tree.insert( result = tree.insert(
InsertParams( InsertParams(
key=key, key=key,
@@ -176,7 +177,7 @@ class TestMamba(unittest.TestCase):
print( print(
f"req2: inserting, req2_token_ids: {req2_token_ids}, req2_kv_indices: {req2_kv_indices}" f"req2: inserting, req2_token_ids: {req2_token_ids}, req2_kv_indices: {req2_kv_indices}"
) )
key = RadixKey(req2_token_ids) key = RadixKey(array("q", req2_token_ids))
result = tree.insert( result = tree.insert(
InsertParams( InsertParams(
key=key, key=key,
@@ -195,7 +196,7 @@ class TestMamba(unittest.TestCase):
print( print(
f"req3: inserting, req3_token_ids: {req3_token_ids}, req3_kv_indices: {req3_kv_indices}" f"req3: inserting, req3_token_ids: {req3_token_ids}, req3_kv_indices: {req3_kv_indices}"
) )
key = RadixKey(req3_token_ids) key = RadixKey(array("q", req3_token_ids))
result = tree.insert( result = tree.insert(
InsertParams( InsertParams(
key=key, key=key,
@@ -213,7 +214,7 @@ class TestMamba(unittest.TestCase):
print( print(
f"req4: inserting, req4_token_ids: {req4_token_ids}, req4_kv_indices: {req4_kv_indices}" f"req4: inserting, req4_token_ids: {req4_token_ids}, req4_kv_indices: {req4_kv_indices}"
) )
key = RadixKey(req4_token_ids) key = RadixKey(array("q", req4_token_ids))
result = tree.insert( result = tree.insert(
InsertParams( InsertParams(
key=key, key=key,
@@ -244,7 +245,9 @@ class TestMamba(unittest.TestCase):
tree.pretty_print() tree.pretty_print()
req5_token_ids = [1, 2, 3, 4, 5] req5_token_ids = [1, 2, 3, 4, 5]
result = tree.match_prefix(MatchPrefixParams(key=RadixKey(req5_token_ids))) result = tree.match_prefix(
MatchPrefixParams(key=RadixKey(array("q", req5_token_ids)))
)
kv_indices, last_node = result.device_indices, result.last_device_node kv_indices, last_node = result.device_indices, result.last_device_node
print( print(
f"req5: token_ids: {req5_token_ids}, matched kv_indices: {kv_indices}, last_node.key: {last_node.key}" f"req5: token_ids: {req5_token_ids}, matched kv_indices: {kv_indices}, last_node.key: {last_node.key}"
@@ -252,7 +255,9 @@ class TestMamba(unittest.TestCase):
assert len(kv_indices) == 0 assert len(kv_indices) == 0
req6_token_ids = [1, 2, 3, 4, 5, 60, 70] req6_token_ids = [1, 2, 3, 4, 5, 60, 70]
result = tree.match_prefix(MatchPrefixParams(key=RadixKey(req6_token_ids))) result = tree.match_prefix(
MatchPrefixParams(key=RadixKey(array("q", req6_token_ids)))
)
kv_indices, last_node = result.device_indices, result.last_device_node kv_indices, last_node = result.device_indices, result.last_device_node
print( print(
f"req6: token_ids: {req6_token_ids}, matched kv_indices: {kv_indices}, last_node.key: {last_node.key}" f"req6: token_ids: {req6_token_ids}, matched kv_indices: {kv_indices}, last_node.key: {last_node.key}"
@@ -261,7 +266,9 @@ class TestMamba(unittest.TestCase):
assert len(last_node.key) == 2 assert len(last_node.key) == 2
req7_token_ids = [1, 2, 3, 4, 5, 6, 7] req7_token_ids = [1, 2, 3, 4, 5, 6, 7]
result = tree.match_prefix(MatchPrefixParams(key=RadixKey(req7_token_ids))) result = tree.match_prefix(
MatchPrefixParams(key=RadixKey(array("q", req7_token_ids)))
)
kv_indices, last_node = result.device_indices, result.last_device_node kv_indices, last_node = result.device_indices, result.last_device_node
print( print(
f"req7: token_ids: {req7_token_ids}, matched kv_indices: {kv_indices}, last_node.key: {last_node.key}" f"req7: token_ids: {req7_token_ids}, matched kv_indices: {kv_indices}, last_node.key: {last_node.key}"
@@ -278,7 +285,9 @@ class TestMamba(unittest.TestCase):
tree.pretty_print() tree.pretty_print()
req8_token_ids = [1, 2, 3, 4, 5, 60, 70] req8_token_ids = [1, 2, 3, 4, 5, 60, 70]
result = tree.match_prefix(MatchPrefixParams(key=RadixKey(req8_token_ids))) result = tree.match_prefix(
MatchPrefixParams(key=RadixKey(array("q", req8_token_ids)))
)
kv_indices, last_node = result.device_indices, result.last_device_node kv_indices, last_node = result.device_indices, result.last_device_node
print( print(
f"req8: token_ids: {req8_token_ids}, matched kv_indices: {kv_indices}, last_node.key: {last_node.key}" f"req8: token_ids: {req8_token_ids}, matched kv_indices: {kv_indices}, last_node.key: {last_node.key}"
@@ -289,7 +298,9 @@ class TestMamba(unittest.TestCase):
req9_token_ids = [1, 2, 3, 4, 5, 6, 7] req9_token_ids = [1, 2, 3, 4, 5, 6, 7]
req9 = make_dummy_req() req9 = make_dummy_req()
result = tree.match_prefix( result = tree.match_prefix(
MatchPrefixParams(key=RadixKey(req9_token_ids), req=req9, cow_mamba=True) MatchPrefixParams(
key=RadixKey(array("q", req9_token_ids)), req=req9, cow_mamba=True
)
) )
kv_indices, last_node = result.device_indices, result.last_device_node kv_indices, last_node = result.device_indices, result.last_device_node
assert req9.mamba_pool_idx is not None assert req9.mamba_pool_idx is not None
@@ -315,7 +326,7 @@ class TestMamba(unittest.TestCase):
stored_hashes = [] stored_hashes = []
req1 = make_dummy_req() req1 = make_dummy_req()
key1 = RadixKey([1, 2, 3]) key1 = RadixKey(array("q", [1, 2, 3]))
tree.insert( tree.insert(
InsertParams( InsertParams(
key=key1, key=key1,
@@ -330,7 +341,7 @@ class TestMamba(unittest.TestCase):
stored_hashes.extend(e.block_hashes[0] for e in stored_events) stored_hashes.extend(e.block_hashes[0] for e in stored_events)
req2 = make_dummy_req() req2 = make_dummy_req()
key2 = RadixKey([1, 2, 3, 4, 5]) key2 = RadixKey(array("q", [1, 2, 3, 4, 5]))
tree.insert( tree.insert(
InsertParams( InsertParams(
key=key2, key=key2,
@@ -367,7 +378,7 @@ class TestMamba(unittest.TestCase):
tree.take_events() # Clear the reset event. tree.take_events() # Clear the reset event.
req1 = make_dummy_req() req1 = make_dummy_req()
key1 = RadixKey([1, 2, 3, 4]) key1 = RadixKey(array("q", [1, 2, 3, 4]))
tree.insert( tree.insert(
InsertParams( InsertParams(
key=key1, key=key1,
@@ -382,7 +393,7 @@ class TestMamba(unittest.TestCase):
split_parent_hash = first_insert_events[1].block_hashes[0] split_parent_hash = first_insert_events[1].block_hashes[0]
req2 = make_dummy_req() req2 = make_dummy_req()
key2 = RadixKey([1, 2, 5, 6]) key2 = RadixKey(array("q", [1, 2, 5, 6]))
tree.insert( tree.insert(
InsertParams( InsertParams(
key=key2, key=key2,
@@ -394,7 +405,7 @@ class TestMamba(unittest.TestCase):
e for e in tree.take_events() if isinstance(e, BlockStored) e for e in tree.take_events() if isinstance(e, BlockStored)
] ]
self.assertEqual(len(second_insert_events), 2) self.assertEqual(len(second_insert_events), 2)
self.assertEqual(second_insert_events[0].token_ids, [5]) self.assertEqual(list(second_insert_events[0].token_ids), [5])
self.assertEqual(second_insert_events[0].parent_block_hash, split_parent_hash) self.assertEqual(second_insert_events[0].parent_block_hash, split_parent_hash)
def _setup_tree_and_allocator(self, enable_kv_cache_events=False): def _setup_tree_and_allocator(self, enable_kv_cache_events=False):
@@ -478,7 +489,7 @@ class TestMamba(unittest.TestCase):
req = Req( req = Req(
rid=0, rid=0,
origin_input_text="", origin_input_text="",
origin_input_ids=[], origin_input_ids=array("q"),
sampling_params=sampling_params, sampling_params=sampling_params,
) )
req_to_token_pool.alloc([req]) req_to_token_pool.alloc([req])
@@ -492,9 +503,9 @@ class TestMamba(unittest.TestCase):
parent = TreeNode() parent = TreeNode()
deleted = TreeNode() deleted = TreeNode()
root.key = RadixKey([]) root.key = RadixKey(array("q", []))
parent.key = RadixKey([1]) parent.key = RadixKey(array("q", [1]))
deleted.key = RadixKey([2]) deleted.key = RadixKey(array("q", [2]))
parent.parent = root parent.parent = root
deleted.parent = parent deleted.parent = parent
parent.value = torch.tensor([1], dtype=torch.int64) parent.value = torch.tensor([1], dtype=torch.int64)
@@ -668,7 +679,7 @@ class TestMamba(unittest.TestCase):
# Step 1: Insert [1,2,3] to create first node # Step 1: Insert [1,2,3] to create first node
req1 = make_dummy_req() req1 = make_dummy_req()
key1 = RadixKey([1, 2, 3]) key1 = RadixKey(array("q", [1, 2, 3]))
tree.insert( tree.insert(
InsertParams( InsertParams(
key=key1, key=key1,
@@ -681,7 +692,7 @@ class TestMamba(unittest.TestCase):
# Step 2: Insert [1,2,3,4,5,6,7] with prev_prefix_len=0 (free all matched) # Step 2: Insert [1,2,3,4,5,6,7] with prev_prefix_len=0 (free all matched)
# Creates tree: [1,2,3] -> [4,5,6,7] # Creates tree: [1,2,3] -> [4,5,6,7]
req2 = make_dummy_req() req2 = make_dummy_req()
key2 = RadixKey([1, 2, 3, 4, 5, 6, 7]) key2 = RadixKey(array("q", [1, 2, 3, 4, 5, 6, 7]))
result = tree.insert( result = tree.insert(
InsertParams( InsertParams(
key=key2, key=key2,
@@ -699,7 +710,7 @@ class TestMamba(unittest.TestCase):
# Matched prefix = 7 (across two nodes: [1,2,3] len=3, [4,5,6,7] len=4) # Matched prefix = 7 (across two nodes: [1,2,3] len=3, [4,5,6,7] len=4)
# Protected [0..1], freed [2..6] = 5 slots, new [7] = 1 slot stored # Protected [0..1], freed [2..6] = 5 slots, new [7] = 1 slot stored
req3 = make_dummy_req() req3 = make_dummy_req()
key3 = RadixKey([1, 2, 3, 4, 5, 6, 7, 8]) key3 = RadixKey(array("q", [1, 2, 3, 4, 5, 6, 7, 8]))
result = tree.insert( result = tree.insert(
InsertParams( InsertParams(
key=key3, key=key3,
@@ -716,7 +727,7 @@ class TestMamba(unittest.TestCase):
# Step 4: Insert [1,2,3,4,5,6,7,8,9] with prev_prefix_len=8 (covers all matched) # Step 4: Insert [1,2,3,4,5,6,7,8,9] with prev_prefix_len=8 (covers all matched)
# Matched prefix = 8, prev_prefix_len=8 => nothing freed # Matched prefix = 8, prev_prefix_len=8 => nothing freed
req4 = make_dummy_req() req4 = make_dummy_req()
key4 = RadixKey([1, 2, 3, 4, 5, 6, 7, 8, 9]) key4 = RadixKey(array("q", [1, 2, 3, 4, 5, 6, 7, 8, 9]))
result = tree.insert( result = tree.insert(
InsertParams( InsertParams(
key=key4, key=key4,
test/registered/unit/mem_cache/test_radix_cache_slru_accuracy.py
+7 -4
View File
@@ -1,4 +1,5 @@
import unittest import unittest
from array import array
import torch import torch
@@ -63,7 +64,7 @@ class TestSLRUAccuracy(unittest.TestCase):
"""Test that SLRU eviction mechanism works correctly""" """Test that SLRU eviction mechanism works correctly"""
# Insert one key-value three times (high frequency access) # Insert one key-value three times (high frequency access)
frequent_key = RadixKey([1, 2]) # High hit rate, should be retained frequent_key = RadixKey(array("q", [1, 2])) # High hit rate, should be retained
frequent_val = torch.tensor([10, 20], dtype=torch.int64) frequent_val = torch.tensor([10, 20], dtype=torch.int64)
# Insert the frequent key multiple times to increase its hit count # Insert the frequent key multiple times to increase its hit count
@@ -71,7 +72,9 @@ class TestSLRUAccuracy(unittest.TestCase):
self.cache.insert(InsertParams(key=frequent_key, value=frequent_val)) self.cache.insert(InsertParams(key=frequent_key, value=frequent_val))
# Insert first low-frequency key-value pair that should be evicted # Insert first low-frequency key-value pair that should be evicted
first_low_freq_key = RadixKey([5, 6]) # Low hit rate, should be evicted first_low_freq_key = RadixKey(
array("q", [5, 6])
) # Low hit rate, should be evicted
first_low_freq_val = torch.tensor([50, 60], dtype=torch.int64) first_low_freq_val = torch.tensor([50, 60], dtype=torch.int64)
self.cache.insert( self.cache.insert(
@@ -81,14 +84,14 @@ class TestSLRUAccuracy(unittest.TestCase):
# Insert other key-values once each (low frequency access) - fill up the cache # Insert other key-values once each (low frequency access) - fill up the cache
other_keys = [] other_keys = []
for i in range(4): # Reduce the number to fit in our smaller cache for i in range(4): # Reduce the number to fit in our smaller cache
key = RadixKey([i + 10]) # Unique keys for low-frequency items key = RadixKey(array("q", [i + 10])) # Unique keys for low-frequency items
val = torch.tensor([i + 100], dtype=torch.int64) val = torch.tensor([i + 100], dtype=torch.int64)
self.cache.insert(InsertParams(key=key, value=val)) self.cache.insert(InsertParams(key=key, value=val))
other_keys.append(key) other_keys.append(key)
# Now insert more items to trigger evictions # Now insert more items to trigger evictions
for i in range(6, 10): # Add more items to definitely exceed capacity for i in range(6, 10): # Add more items to definitely exceed capacity
key = RadixKey([i * 2]) # Different pattern to avoid conflicts key = RadixKey(array("q", [i * 2])) # Different pattern to avoid conflicts
val = torch.tensor([i * 200], dtype=torch.int64) val = torch.tensor([i * 200], dtype=torch.int64)
self.cache.insert(InsertParams(key=key, value=val)) self.cache.insert(InsertParams(key=key, value=val))
test/registered/unit/mem_cache/test_radix_cache_unit.py
+85 -53
View File
@@ -28,6 +28,7 @@ import random
import time import time
import unittest import unittest
import unittest.mock import unittest.mock
from array import array
import torch import torch
@@ -50,30 +51,30 @@ class TestRadixKey(unittest.TestCase):
def test_init_basic(self): def test_init_basic(self):
"""Test basic initialization of RadixKey.""" """Test basic initialization of RadixKey."""
token_ids = [1, 2, 3, 4] token_ids = [1, 2, 3, 4]
key = RadixKey(token_ids) key = RadixKey(array("q", token_ids))
self.assertEqual(key.token_ids, token_ids) self.assertEqual(list(key.token_ids), token_ids)
self.assertIsNone(key.extra_key) self.assertIsNone(key.extra_key)
def test_init_with_extra_key(self): def test_init_with_extra_key(self):
"""Test initialization with extra_key.""" """Test initialization with extra_key."""
token_ids = [1, 2, 3] token_ids = [1, 2, 3]
extra_key = "test_key" extra_key = "test_key"
key = RadixKey(token_ids, extra_key) key = RadixKey(array("q", token_ids), extra_key)
self.assertEqual(key.token_ids, token_ids) self.assertEqual(list(key.token_ids), token_ids)
self.assertEqual(key.extra_key, extra_key) self.assertEqual(key.extra_key, extra_key)
def test_len(self): def test_len(self):
"""Test __len__ method.""" """Test __len__ method."""
key = RadixKey([1, 2, 3]) key = RadixKey(array("q", [1, 2, 3]))
self.assertEqual(len(key), 3) self.assertEqual(len(key), 3)
empty_key = RadixKey([]) empty_key = RadixKey(array("q", []))
self.assertEqual(len(empty_key), 0) self.assertEqual(len(empty_key), 0)
def test_iter(self): def test_iter(self):
"""Test __iter__ method.""" """Test __iter__ method."""
token_ids = [1, 2, 3, 4] token_ids = [1, 2, 3, 4]
key = RadixKey(token_ids) key = RadixKey(array("q", token_ids))
self.assertEqual(list(key), token_ids) self.assertEqual(list(key), token_ids)
def test_len_and_iter(self): def test_len_and_iter(self):
@@ -86,7 +87,7 @@ class TestRadixKey(unittest.TestCase):
for tokens, expected in test_cases: for tokens, expected in test_cases:
with self.subTest(tokens=tokens): with self.subTest(tokens=tokens):
key = RadixKey(tokens) key = RadixKey(array("q", tokens))
self.assertEqual(len(key), expected) self.assertEqual(len(key), expected)
self.assertEqual(list(key), tokens) self.assertEqual(list(key), tokens)
@@ -100,34 +101,34 @@ class TestRadixKey(unittest.TestCase):
for tokens, index, expected in test_cases: for tokens, index, expected in test_cases:
with self.subTest(tokens=tokens, index=index): with self.subTest(tokens=tokens, index=index):
key = RadixKey(tokens) key = RadixKey(array("q", tokens))
result = key[index] result = key[index]
self.assertIsInstance(result, RadixKey) self.assertIsInstance(result, RadixKey)
self.assertEqual(result.token_ids, expected) self.assertEqual(list(result.token_ids), expected)
def test_getitem_slice(self): def test_getitem_slice(self):
"""Test __getitem__ with slice and edge cases.""" """Test __getitem__ with slice and edge cases."""
key = RadixKey([1, 2, 3, 4, 5], "extra") key = RadixKey(array("q", [1, 2, 3, 4, 5]), "extra")
# Basic slice # Basic slice
sliced = key[1:4] sliced = key[1:4]
self.assertIsInstance(sliced, RadixKey) self.assertIsInstance(sliced, RadixKey)
self.assertEqual(sliced.token_ids, [2, 3, 4]) self.assertEqual(list(sliced.token_ids), [2, 3, 4])
self.assertEqual(sliced.extra_key, "extra") self.assertEqual(sliced.extra_key, "extra")
# Edge cases # Edge cases
self.assertEqual(key[2:2].token_ids, []) # Empty slice self.assertEqual(list(key[2:2].token_ids), []) # Empty slice
self.assertEqual(key[:].token_ids, [1, 2, 3, 4, 5]) # Full slice self.assertEqual(list(key[:].token_ids), [1, 2, 3, 4, 5]) # Full slice
def test_getitem_invalid_index(self): def test_getitem_invalid_index(self):
"""Test __getitem__ with invalid indices.""" """Test __getitem__ with invalid indices."""
key = RadixKey([1, 2, 3]) key = RadixKey(array("q", [1, 2, 3]))
with self.assertRaises(IndexError): with self.assertRaises(IndexError):
_ = key[10] # Out of bounds _ = key[10] # Out of bounds
def test_repr(self): def test_repr(self):
"""Test __repr__ method.""" """Test __repr__ method."""
key = RadixKey([1, 2, 3], "test") key = RadixKey(array("q", [1, 2, 3]), "test")
repr_str = repr(key) repr_str = repr(key)
self.assertIn("RadixKey", repr_str) self.assertIn("RadixKey", repr_str)
self.assertIn("extra_key='test'", repr_str) self.assertIn("extra_key='test'", repr_str)
@@ -136,7 +137,7 @@ class TestRadixKey(unittest.TestCase):
def test_repr_long_token_ids(self): def test_repr_long_token_ids(self):
"""Test __repr__ with long token_ids.""" """Test __repr__ with long token_ids."""
long_tokens = list(range(15)) long_tokens = list(range(15))
key = RadixKey(long_tokens) key = RadixKey(array("q", long_tokens))
repr_str = repr(key) repr_str = repr(key)
self.assertIn("...", repr_str) # Should be truncated self.assertIn("...", repr_str) # Should be truncated
@@ -274,7 +275,7 @@ class TestRadixCache(unittest.TestCase):
# Insert some data # Insert some data
cache.insert( cache.insert(
InsertParams( InsertParams(
key=RadixKey([1, 2, 3]), key=RadixKey(array("q", [1, 2, 3])),
value=torch.tensor([10, 20, 30], dtype=torch.int64), value=torch.tensor([10, 20, 30], dtype=torch.int64),
) )
) )
@@ -292,7 +293,7 @@ class TestRadixCache(unittest.TestCase):
with self.subTest(disable_cache=disable_cache): with self.subTest(disable_cache=disable_cache):
cache = RadixCache.create_simulated(disable=disable_cache) cache = RadixCache.create_simulated(disable=disable_cache)
key = RadixKey([1, 2, 3]) key = RadixKey(array("q", [1, 2, 3]))
value = torch.tensor([10, 20, 30], dtype=torch.int64) value = torch.tensor([10, 20, 30], dtype=torch.int64)
result = cache.insert(InsertParams(key=key, value=value)) result = cache.insert(InsertParams(key=key, value=value))
prefix_len = result.prefix_len prefix_len = result.prefix_len
@@ -307,12 +308,16 @@ class TestRadixCache(unittest.TestCase):
self.assertEqual(cache.evictable_size(), 3) self.assertEqual(cache.evictable_size(), 3)
# Test match_prefix # Test match_prefix
result = cache.match_prefix(MatchPrefixParams(key=RadixKey([1, 2, 3]))) result = cache.match_prefix(
MatchPrefixParams(key=RadixKey(array("q", [1, 2, 3])))
)
self.assertEqual(len(result.device_indices), 3) self.assertEqual(len(result.device_indices), 3)
torch.testing.assert_close(result.device_indices, value) torch.testing.assert_close(result.device_indices, value)
# Test partial match # Test partial match
result = cache.match_prefix(MatchPrefixParams(key=RadixKey([1, 2]))) result = cache.match_prefix(
MatchPrefixParams(key=RadixKey(array("q", [1, 2])))
)
self.assertEqual(len(result.device_indices), 2) self.assertEqual(len(result.device_indices), 2)
torch.testing.assert_close( torch.testing.assert_close(
result.device_indices, torch.tensor([10, 20], dtype=torch.int64) result.device_indices, torch.tensor([10, 20], dtype=torch.int64)
@@ -322,7 +327,7 @@ class TestRadixCache(unittest.TestCase):
"""Test insert with None value (should use token_ids as list).""" """Test insert with None value (should use token_ids as list)."""
cache = RadixCache.create_simulated() cache = RadixCache.create_simulated()
key = RadixKey([1, 2, 3]) key = RadixKey(array("q", [1, 2, 3]))
result = cache.insert(InsertParams(key=key, value=None)) result = cache.insert(InsertParams(key=key, value=None))
prefix_len = result.prefix_len prefix_len = result.prefix_len
@@ -338,7 +343,7 @@ class TestRadixCache(unittest.TestCase):
cache.insert( cache.insert(
InsertParams( InsertParams(
key=RadixKey([1, 2, 3]), key=RadixKey(array("q", [1, 2, 3])),
value=torch.tensor([10, 20, 30], dtype=torch.int64), value=torch.tensor([10, 20, 30], dtype=torch.int64),
) )
) )
@@ -346,7 +351,8 @@ class TestRadixCache(unittest.TestCase):
cache.insert( cache.insert(
InsertParams( InsertParams(
key=RadixKey([4, 5]), value=torch.tensor([40, 50], dtype=torch.int64) key=RadixKey(array("q", [4, 5])),
value=torch.tensor([40, 50], dtype=torch.int64),
) )
) )
self.assertEqual(cache.total_size(), 5) self.assertEqual(cache.total_size(), 5)
@@ -366,7 +372,9 @@ class TestRadixCache(unittest.TestCase):
) )
# Insert data # Insert data
cache.insert(InsertParams(key=RadixKey([1, 2, 3, 4, 5]), value=None)) cache.insert(
InsertParams(key=RadixKey(array("q", [1, 2, 3, 4, 5])), value=None)
)
# Take events # Take events
events = cache.take_events() events = cache.take_events()
@@ -395,7 +403,7 @@ class TestRadixCache(unittest.TestCase):
# Insert and then evict data # Insert and then evict data
cache.insert( cache.insert(
InsertParams( InsertParams(
key=RadixKey([1, 2, 3]), key=RadixKey(array("q", [1, 2, 3])),
value=torch.tensor([10, 20, 30], dtype=torch.int64), value=torch.tensor([10, 20, 30], dtype=torch.int64),
) )
) )
@@ -427,29 +435,35 @@ class TestRadixCache(unittest.TestCase):
# Insert same token sequence with different extra keys # Insert same token sequence with different extra keys
cache.insert( cache.insert(
InsertParams( InsertParams(
key=RadixKey([1, 2, 3], "key1"), key=RadixKey(array("q", [1, 2, 3]), "key1"),
value=torch.tensor([10, 20, 30], dtype=torch.int64), value=torch.tensor([10, 20, 30], dtype=torch.int64),
) )
) )
cache.insert( cache.insert(
InsertParams( InsertParams(
key=RadixKey([1, 2, 3], "key2"), key=RadixKey(array("q", [1, 2, 3]), "key2"),
value=torch.tensor([40, 50, 60], dtype=torch.int64), value=torch.tensor([40, 50, 60], dtype=torch.int64),
) )
) )
cache.insert( cache.insert(
InsertParams( InsertParams(
key=RadixKey([1, 2, 3], None), key=RadixKey(array("q", [1, 2, 3]), None),
value=torch.tensor([70, 80, 90], dtype=torch.int64), value=torch.tensor([70, 80, 90], dtype=torch.int64),
) )
) )
# Keys with different extra_key should not match each other # Keys with different extra_key should not match each other
result1 = cache.match_prefix(MatchPrefixParams(key=RadixKey([1, 2, 3], "key1"))) result1 = cache.match_prefix(
result2 = cache.match_prefix(MatchPrefixParams(key=RadixKey([1, 2, 3], "key2"))) MatchPrefixParams(key=RadixKey(array("q", [1, 2, 3]), "key1"))
result3 = cache.match_prefix(MatchPrefixParams(key=RadixKey([1, 2, 3], None))) )
result2 = cache.match_prefix(
MatchPrefixParams(key=RadixKey(array("q", [1, 2, 3]), "key2"))
)
result3 = cache.match_prefix(
MatchPrefixParams(key=RadixKey(array("q", [1, 2, 3]), None))
)
result4 = cache.match_prefix( result4 = cache.match_prefix(
MatchPrefixParams(key=RadixKey([1, 2, 3], "nonexistent")) MatchPrefixParams(key=RadixKey(array("q", [1, 2, 3]), "nonexistent"))
) )
# Each should match only its own data # Each should match only its own data
@@ -478,13 +492,15 @@ class TestRadixCache(unittest.TestCase):
# Insert sequence # Insert sequence
cache.insert( cache.insert(
InsertParams( InsertParams(
key=RadixKey([1, 2, 3]), key=RadixKey(array("q", [1, 2, 3])),
value=torch.tensor([10, 20, 30], dtype=torch.int64), value=torch.tensor([10, 20, 30], dtype=torch.int64),
) )
) )
# Get node # Get node
result = cache.match_prefix(MatchPrefixParams(key=RadixKey([1, 2, 3]))) result = cache.match_prefix(
MatchPrefixParams(key=RadixKey(array("q", [1, 2, 3])))
)
node = result.last_device_node node = result.last_device_node
initial_evictable = cache.evictable_size() initial_evictable = cache.evictable_size()
@@ -510,12 +526,14 @@ class TestRadixCache(unittest.TestCase):
# Insert sequences # Insert sequences
cache.insert( cache.insert(
InsertParams( InsertParams(
key=RadixKey([1, 2]), value=torch.tensor([10, 20], dtype=torch.int64) key=RadixKey(array("q", [1, 2])),
value=torch.tensor([10, 20], dtype=torch.int64),
) )
) )
cache.insert( cache.insert(
InsertParams( InsertParams(
key=RadixKey([3, 4]), value=torch.tensor([30, 40], dtype=torch.int64) key=RadixKey(array("q", [3, 4])),
value=torch.tensor([30, 40], dtype=torch.int64),
) )
) )
@@ -547,7 +565,7 @@ class TestRadixCache(unittest.TestCase):
cache = RadixCache.create_simulated(page_size=page_size) cache = RadixCache.create_simulated(page_size=page_size)
tokens = list(range(sequence_length)) tokens = list(range(sequence_length))
key = RadixKey(tokens) key = RadixKey(array("q", tokens))
cache.insert( cache.insert(
InsertParams( InsertParams(
key=key, key=key,
@@ -555,7 +573,9 @@ class TestRadixCache(unittest.TestCase):
) )
) )
result = cache.match_prefix(MatchPrefixParams(key=RadixKey(tokens))) result = cache.match_prefix(
MatchPrefixParams(key=RadixKey(array("q", tokens)))
)
self.assertGreater(len(result.device_indices), 0) self.assertGreater(len(result.device_indices), 0)
# Match length should be page-aligned # Match length should be page-aligned
@@ -568,7 +588,7 @@ class TestRadixCache(unittest.TestCase):
cache.insert( cache.insert(
InsertParams( InsertParams(
key=RadixKey([1, 2, 3]), key=RadixKey(array("q", [1, 2, 3])),
value=torch.tensor([10, 20, 30], dtype=torch.int64), value=torch.tensor([10, 20, 30], dtype=torch.int64),
) )
) )
@@ -585,12 +605,14 @@ class TestRadixCache(unittest.TestCase):
cache.insert( cache.insert(
InsertParams( InsertParams(
key=RadixKey([1, 2]), value=torch.tensor([10, 20], dtype=torch.int64) key=RadixKey(array("q", [1, 2])),
value=torch.tensor([10, 20], dtype=torch.int64),
) )
) )
cache.insert( cache.insert(
InsertParams( InsertParams(
key=RadixKey([3, 4]), value=torch.tensor([30, 40], dtype=torch.int64) key=RadixKey(array("q", [3, 4])),
value=torch.tensor([30, 40], dtype=torch.int64),
) )
) )
@@ -609,12 +631,12 @@ class TestRadixCache(unittest.TestCase):
# Insert a long sequence that will be split later. # Insert a long sequence that will be split later.
seq1 = [1, 2, 3, 4, 5, 6, 7, 8] seq1 = [1, 2, 3, 4, 5, 6, 7, 8]
val1 = torch.tensor([x * 10 for x in seq1], dtype=torch.int64) val1 = torch.tensor([x * 10 for x in seq1], dtype=torch.int64)
cache.insert(InsertParams(key=RadixKey(seq1), value=val1)) cache.insert(InsertParams(key=RadixKey(array("q", seq1)), value=val1))
# Insert a diverging branch to create an internal node on the path. # Insert a diverging branch to create an internal node on the path.
seq2 = [1, 2, 9, 10] seq2 = [1, 2, 9, 10]
val2 = torch.tensor([x * 10 for x in seq2], dtype=torch.int64) val2 = torch.tensor([x * 10 for x in seq2], dtype=torch.int64)
cache.insert(InsertParams(key=RadixKey(seq2), value=val2)) cache.insert(InsertParams(key=RadixKey(array("q", seq2)), value=val2))
print(cache.pretty_print()) print(cache.pretty_print())
baseline_total = cache.total_size() baseline_total = cache.total_size()
@@ -624,24 +646,30 @@ class TestRadixCache(unittest.TestCase):
# Match that causes a split inside an existing node: # Match that causes a split inside an existing node:
# take first 4 tokens of seq1, then diverge. # take first 4 tokens of seq1, then diverge.
query1 = [1, 2, 3, 4, 999, 1000] query1 = [1, 2, 3, 4, 999, 1000]
result1 = cache.match_prefix(MatchPrefixParams(key=RadixKey(query1))) result1 = cache.match_prefix(
MatchPrefixParams(key=RadixKey(array("q", query1)))
)
torch.testing.assert_close(result1.device_indices, val1[:4]) torch.testing.assert_close(result1.device_indices, val1[:4])
# No data change after structural split during matching. # No data change after structural split during matching.
self.assertEqual(cache.total_size(), baseline_total) self.assertEqual(cache.total_size(), baseline_total)
# Full match of the long sequence still returns the full indices. # Full match of the long sequence still returns the full indices.
result_full = cache.match_prefix(MatchPrefixParams(key=RadixKey(seq1))) result_full = cache.match_prefix(
MatchPrefixParams(key=RadixKey(array("q", seq1)))
)
torch.testing.assert_close(result_full.device_indices, val1) torch.testing.assert_close(result_full.device_indices, val1)
# Another split deeper on the path (after matching 6 tokens, then diverge). # Another split deeper on the path (after matching 6 tokens, then diverge).
query2 = [1, 2, 3, 4, 5, 6, 777, 888] query2 = [1, 2, 3, 4, 5, 6, 777, 888]
result2 = cache.match_prefix(MatchPrefixParams(key=RadixKey(query2))) result2 = cache.match_prefix(
MatchPrefixParams(key=RadixKey(array("q", query2)))
)
torch.testing.assert_close(result2.device_indices, val1[:6]) torch.testing.assert_close(result2.device_indices, val1[:6])
self.assertEqual(cache.total_size(), baseline_total) self.assertEqual(cache.total_size(), baseline_total)
# Matching the short diverging branch should return exactly its indices. # Matching the short diverging branch should return exactly its indices.
result_branch = cache.match_prefix( result_branch = cache.match_prefix(
MatchPrefixParams(key=RadixKey(seq2)) MatchPrefixParams(key=RadixKey(array("q", seq2)))
) )
torch.testing.assert_close(result_branch.device_indices, val2) torch.testing.assert_close(result_branch.device_indices, val2)
@@ -653,7 +681,9 @@ class TestRadixCache(unittest.TestCase):
) )
# Insert a sequence # Insert a sequence
cache.insert(InsertParams(key=RadixKey([1, 2, 3, 4, 5, 6, 7, 8]), value=None)) cache.insert(
InsertParams(key=RadixKey(array("q", [1, 2, 3, 4, 5, 6, 7, 8])), value=None)
)
# Trigger event emission to compute hash_value lazily # Trigger event emission to compute hash_value lazily
cache.take_events() cache.take_events()
@@ -679,7 +709,9 @@ class TestRadixCache(unittest.TestCase):
) )
# Insert a sequence with repeating token pattern: [1,2,3,4, 1,2,3,4] # Insert a sequence with repeating token pattern: [1,2,3,4, 1,2,3,4]
cache.insert(InsertParams(key=RadixKey([1, 2, 3, 4, 1, 2, 3, 4]), value=None)) cache.insert(
InsertParams(key=RadixKey(array("q", [1, 2, 3, 4, 1, 2, 3, 4])), value=None)
)
events = cache.take_events() events = cache.take_events()
block_stored_events = [e for e in events if isinstance(e, BlockStored)] block_stored_events = [e for e in events if isinstance(e, BlockStored)]
@@ -713,11 +745,11 @@ class TestRadixCache(unittest.TestCase):
) )
# Insert a sequence that will cause a split # Insert a sequence that will cause a split
cache.insert(InsertParams(key=RadixKey([1, 2, 3, 4]), value=None)) cache.insert(InsertParams(key=RadixKey(array("q", [1, 2, 3, 4])), value=None))
cache.take_events() # Clear events and compute hash_value for first node cache.take_events() # Clear events and compute hash_value for first node
# Insert a diverging sequence that will cause a split at page boundary # Insert a diverging sequence that will cause a split at page boundary
cache.insert(InsertParams(key=RadixKey([1, 2, 5, 6]), value=None)) cache.insert(InsertParams(key=RadixKey(array("q", [1, 2, 5, 6])), value=None))
cache.take_events() # Trigger event emission to compute hash_value cache.take_events() # Trigger event emission to compute hash_value
# Find the split node # Find the split node
@@ -754,7 +786,7 @@ class TestRadixCache(unittest.TestCase):
cache: RadixCache = RadixCache.create_simulated() cache: RadixCache = RadixCache.create_simulated()
for key, value in zip(keys, values): for key, value in zip(keys, values):
cache.insert(InsertParams(key=RadixKey(key), value=value)) cache.insert(InsertParams(key=RadixKey(array("q", key)), value=value))
del values del values
test/registered/unit/mem_cache/test_radix_force_miss.py
+8 -5
View File
@@ -11,6 +11,7 @@ register_cpu_ci(est_time=5, suite="base-a-test-cpu")
import unittest import unittest
import unittest.mock import unittest.mock
from array import array
import torch import torch
@@ -27,8 +28,8 @@ from sglang.srt.mem_cache.radix_cache import RadixCache, RadixKey
class _StubReq: class _StubReq:
def __init__(self, token_ids): def __init__(self, token_ids):
self.origin_input_ids = list(token_ids) self.origin_input_ids = array("q", token_ids)
self.output_ids = [] self.output_ids = array("q")
self.extra_key = None self.extra_key = None
self.prefix_indices = None self.prefix_indices = None
self.last_node = None self.last_node = None
@@ -42,9 +43,9 @@ class _StubReq:
class TestZeroMatchResult(unittest.TestCase): class TestZeroMatchResult(unittest.TestCase):
def test_zero_replaces_indices_and_nodes(self): def test_zero_replaces_indices_and_nodes(self):
tree = RadixCache.create_simulated() tree = RadixCache.create_simulated()
tree.insert(InsertParams(key=RadixKey(token_ids=[1, 2, 3, 4, 5]))) tree.insert(InsertParams(key=RadixKey(token_ids=array("q", [1, 2, 3, 4, 5]))))
match = tree.match_prefix( match = tree.match_prefix(
MatchPrefixParams(key=RadixKey(token_ids=[1, 2, 3, 9])) MatchPrefixParams(key=RadixKey(token_ids=array("q", [1, 2, 3, 9])))
) )
self.assertGreater(len(match.device_indices), 0) self.assertGreater(len(match.device_indices), 0)
zeroed = zero_match_result(tree, match) zeroed = zero_match_result(tree, match)
@@ -76,7 +77,9 @@ class TestMatchPrefixForReqForceMiss(unittest.TestCase):
def test_force_miss_zeros_req_prefix(self): def test_force_miss_zeros_req_prefix(self):
tree = RadixCache.create_simulated() tree = RadixCache.create_simulated()
tree.insert( tree.insert(
InsertParams(key=RadixKey(token_ids=[10, 11, 12, 13, 14, 15, 16, 17])) InsertParams(
key=RadixKey(token_ids=array("q", [10, 11, 12, 13, 14, 15, 16, 17]))
)
) )
# Sanity: without the flag, the same lookup hits. # Sanity: without the flag, the same lookup hits.
test/registered/unit/mem_cache/test_swa_lock_release_lifecycle.py
+2
View File
@@ -12,6 +12,7 @@ Covers:
""" """
import unittest import unittest
from array import array
import torch import torch
@@ -110,6 +111,7 @@ def _swa_alloc(allocator, need_size):
def _insert_chain(tree, allocator, token_ids): def _insert_chain(tree, allocator, token_ids):
token_ids = array("q", token_ids)
indices = _swa_alloc(allocator, len(token_ids)) indices = _swa_alloc(allocator, len(token_ids))
assert indices is not None assert indices is not None
tree.insert(InsertParams(key=RadixKey(token_ids), value=indices)) tree.insert(InsertParams(key=RadixKey(token_ids), value=indices))
test/registered/unit/mem_cache/test_swa_unittest.py
+31 -20
View File
@@ -1,4 +1,5 @@
import unittest import unittest
from array import array
import torch import torch
@@ -113,12 +114,12 @@ def _swa_alloc(allocator, need_size):
def _insert(tree, allocator, token_ids): def _insert(tree, allocator, token_ids):
indices = _swa_alloc(allocator, len(token_ids)) indices = _swa_alloc(allocator, len(token_ids))
assert indices is not None assert indices is not None
tree.insert(InsertParams(key=RadixKey(token_ids), value=indices)) tree.insert(InsertParams(key=RadixKey(array("q", token_ids)), value=indices))
def _insert_chain(tree, allocator, token_ids): def _insert_chain(tree, allocator, token_ids):
_insert(tree, allocator, token_ids) _insert(tree, allocator, token_ids)
match = tree.match_prefix(MatchPrefixParams(key=RadixKey(token_ids))) match = tree.match_prefix(MatchPrefixParams(key=RadixKey(array("q", token_ids))))
return match.last_device_node return match.last_device_node
@@ -193,7 +194,7 @@ class TestSWA(unittest.TestCase):
e for e in tree.take_events() if isinstance(e, BlockStored) e for e in tree.take_events() if isinstance(e, BlockStored)
] ]
self.assertEqual(len(second_insert_events), 2) self.assertEqual(len(second_insert_events), 2)
self.assertEqual(second_insert_events[0].token_ids, [5]) self.assertEqual(list(second_insert_events[0].token_ids), [5])
self.assertEqual(second_insert_events[0].parent_block_hash, split_parent_hash) self.assertEqual(second_insert_events[0].parent_block_hash, split_parent_hash)
def test_swa_memory_pool(self): def test_swa_memory_pool(self):
@@ -313,7 +314,7 @@ class TestSWA(unittest.TestCase):
print( print(
f"req1: inserting, req1_token_ids: {req1_token_ids}, req1_kv_indices: {req1_kv_indices}" f"req1: inserting, req1_token_ids: {req1_token_ids}, req1_kv_indices: {req1_kv_indices}"
) )
key = RadixKey(req1_token_ids) key = RadixKey(array("q", req1_token_ids))
result = tree.insert(InsertParams(key=key, value=req1_kv_indices[: len(key)])) result = tree.insert(InsertParams(key=key, value=req1_kv_indices[: len(key)]))
prefix_len = result.prefix_len prefix_len = result.prefix_len
print( print(
@@ -324,7 +325,7 @@ class TestSWA(unittest.TestCase):
print( print(
f"req2: inserting, req2_token_ids: {req2_token_ids}, req2_kv_indices: {req2_kv_indices}" f"req2: inserting, req2_token_ids: {req2_token_ids}, req2_kv_indices: {req2_kv_indices}"
) )
key = RadixKey(req2_token_ids) key = RadixKey(array("q", req2_token_ids))
result = tree.insert(InsertParams(key=key, value=req2_kv_indices[: len(key)])) result = tree.insert(InsertParams(key=key, value=req2_kv_indices[: len(key)]))
prefix_len = result.prefix_len prefix_len = result.prefix_len
print( print(
@@ -335,7 +336,7 @@ class TestSWA(unittest.TestCase):
print( print(
f"req3: inserting, req3_token_ids: {req3_token_ids}, req3_kv_indices: {req3_kv_indices}" f"req3: inserting, req3_token_ids: {req3_token_ids}, req3_kv_indices: {req3_kv_indices}"
) )
key = RadixKey(req3_token_ids) key = RadixKey(array("q", req3_token_ids))
result = tree.insert(InsertParams(key=key, value=req3_kv_indices[: len(key)])) result = tree.insert(InsertParams(key=key, value=req3_kv_indices[: len(key)]))
prefix_len = result.prefix_len prefix_len = result.prefix_len
print( print(
@@ -346,7 +347,7 @@ class TestSWA(unittest.TestCase):
print( print(
f"req4: inserting, req4_token_ids: {req4_token_ids}, req4_kv_indices: {req4_kv_indices}" f"req4: inserting, req4_token_ids: {req4_token_ids}, req4_kv_indices: {req4_kv_indices}"
) )
key = RadixKey(req4_token_ids) key = RadixKey(array("q", req4_token_ids))
result = tree.insert(InsertParams(key=key, value=req4_kv_indices[: len(key)])) result = tree.insert(InsertParams(key=key, value=req4_kv_indices[: len(key)]))
prefix_len = result.prefix_len prefix_len = result.prefix_len
print( print(
@@ -376,7 +377,9 @@ class TestSWA(unittest.TestCase):
tree.pretty_print() tree.pretty_print()
req5_token_ids = [1, 2, 3, 4, 5] req5_token_ids = [1, 2, 3, 4, 5]
result = tree.match_prefix(MatchPrefixParams(key=RadixKey(req5_token_ids))) result = tree.match_prefix(
MatchPrefixParams(key=RadixKey(array("q", req5_token_ids)))
)
kv_indices, last_node = result.device_indices, result.last_device_node kv_indices, last_node = result.device_indices, result.last_device_node
print( print(
f"req5: token_ids: {req5_token_ids}, matched kv_indices: {kv_indices}, last_node.key: {last_node.key}" f"req5: token_ids: {req5_token_ids}, matched kv_indices: {kv_indices}, last_node.key: {last_node.key}"
@@ -384,7 +387,9 @@ class TestSWA(unittest.TestCase):
self.assertEqual(len(kv_indices), 0) self.assertEqual(len(kv_indices), 0)
req6_token_ids = [1, 2, 3, 4, 5, 60, 70] req6_token_ids = [1, 2, 3, 4, 5, 60, 70]
result = tree.match_prefix(MatchPrefixParams(key=RadixKey(req6_token_ids))) result = tree.match_prefix(
MatchPrefixParams(key=RadixKey(array("q", req6_token_ids)))
)
kv_indices, last_node = result.device_indices, result.last_device_node kv_indices, last_node = result.device_indices, result.last_device_node
print( print(
f"req6: token_ids: {req6_token_ids}, matched kv_indices: {kv_indices}, last_node.key: {last_node.key}" f"req6: token_ids: {req6_token_ids}, matched kv_indices: {kv_indices}, last_node.key: {last_node.key}"
@@ -468,7 +473,7 @@ class TestSWA(unittest.TestCase):
print( print(
f"req1: inserting, req1_token_ids: {req1_token_ids}, req1_kv_indices: {req1_kv_indices}" f"req1: inserting, req1_token_ids: {req1_token_ids}, req1_kv_indices: {req1_kv_indices}"
) )
key = RadixKey(req1_token_ids) key = RadixKey(array("q", req1_token_ids))
result = tree.insert(InsertParams(key=key, value=req1_kv_indices[: len(key)])) result = tree.insert(InsertParams(key=key, value=req1_kv_indices[: len(key)]))
prefix_len = result.prefix_len prefix_len = result.prefix_len
self.assertEqual(prefix_len, 0) self.assertEqual(prefix_len, 0)
@@ -480,7 +485,7 @@ class TestSWA(unittest.TestCase):
print( print(
f"req2: inserting, req2_token_ids: {req2_token_ids}, req2_kv_indices: {req2_kv_indices}" f"req2: inserting, req2_token_ids: {req2_token_ids}, req2_kv_indices: {req2_kv_indices}"
) )
key = RadixKey(req2_token_ids) key = RadixKey(array("q", req2_token_ids))
result = tree.insert(InsertParams(key=key, value=req2_kv_indices[: len(key)])) result = tree.insert(InsertParams(key=key, value=req2_kv_indices[: len(key)]))
prefix_len = result.prefix_len prefix_len = result.prefix_len
self.assertEqual(prefix_len, 2) self.assertEqual(prefix_len, 2)
@@ -492,7 +497,7 @@ class TestSWA(unittest.TestCase):
print( print(
f"req3: inserting, req3_token_ids: {req3_token_ids}, req3_kv_indices: {req3_kv_indices}" f"req3: inserting, req3_token_ids: {req3_token_ids}, req3_kv_indices: {req3_kv_indices}"
) )
key = RadixKey(req3_token_ids) key = RadixKey(array("q", req3_token_ids))
result = tree.insert(InsertParams(key=key, value=req3_kv_indices[: len(key)])) result = tree.insert(InsertParams(key=key, value=req3_kv_indices[: len(key)]))
prefix_len = result.prefix_len prefix_len = result.prefix_len
self.assertEqual(prefix_len, 0) self.assertEqual(prefix_len, 0)
@@ -504,7 +509,7 @@ class TestSWA(unittest.TestCase):
print( print(
f"req4: inserting, req4_token_ids: {req4_token_ids}, req4_kv_indices: {req4_kv_indices}" f"req4: inserting, req4_token_ids: {req4_token_ids}, req4_kv_indices: {req4_kv_indices}"
) )
key = RadixKey(req4_token_ids) key = RadixKey(array("q", req4_token_ids))
result = tree.insert(InsertParams(key=key, value=req4_kv_indices[: len(key)])) result = tree.insert(InsertParams(key=key, value=req4_kv_indices[: len(key)]))
prefix_len = result.prefix_len prefix_len = result.prefix_len
self.assertEqual(prefix_len, 4) self.assertEqual(prefix_len, 4)
@@ -553,7 +558,9 @@ class TestSWA(unittest.TestCase):
tree.pretty_print() tree.pretty_print()
req5_token_ids = [1, 2, 3, 4, 5] req5_token_ids = [1, 2, 3, 4, 5]
result = tree.match_prefix(MatchPrefixParams(key=RadixKey(req5_token_ids))) result = tree.match_prefix(
MatchPrefixParams(key=RadixKey(array("q", req5_token_ids)))
)
kv_indices, last_node = result.device_indices, result.last_device_node kv_indices, last_node = result.device_indices, result.last_device_node
print( print(
f"req5: token_ids: {req5_token_ids}, matched kv_indices: {kv_indices}, last_node.key: {last_node.key}" f"req5: token_ids: {req5_token_ids}, matched kv_indices: {kv_indices}, last_node.key: {last_node.key}"
@@ -561,7 +568,9 @@ class TestSWA(unittest.TestCase):
self.assertEqual(len(kv_indices), 0) # no swa prefix matched self.assertEqual(len(kv_indices), 0) # no swa prefix matched
req6_token_ids = [1, 2, 3, 4, 5, 60, 70] req6_token_ids = [1, 2, 3, 4, 5, 60, 70]
result = tree.match_prefix(MatchPrefixParams(key=RadixKey(req6_token_ids))) result = tree.match_prefix(
MatchPrefixParams(key=RadixKey(array("q", req6_token_ids)))
)
kv_indices, last_node = result.device_indices, result.last_device_node kv_indices, last_node = result.device_indices, result.last_device_node
print( print(
f"req6: token_ids: {req6_token_ids}, matched kv_indices: {kv_indices}, last_node.key: {last_node.key}" f"req6: token_ids: {req6_token_ids}, matched kv_indices: {kv_indices}, last_node.key: {last_node.key}"
@@ -578,8 +587,8 @@ class TestSWA(unittest.TestCase):
# Case 1: is_insert=True should pass bigram key and use cache_protected_len. # Case 1: is_insert=True should pass bigram key and use cache_protected_len.
req = _DummyReq() req = _DummyReq()
req.req_pool_idx = 0 req.req_pool_idx = 0
req.origin_input_ids = [1, 2, 3, 4, 5, 6] req.origin_input_ids = array("q", [1, 2, 3, 4, 5, 6])
req.output_ids = [] req.output_ids = array("q")
req._kv_committed_len = len(req.origin_input_ids) req._kv_committed_len = len(req.origin_input_ids)
kv_indices = allocator.alloc(req._kv_committed_len) kv_indices = allocator.alloc(req._kv_committed_len)
req_to_token_pool.write( req_to_token_pool.write(
@@ -613,8 +622,8 @@ class TestSWA(unittest.TestCase):
# even when len(prefix_indices) is intentionally larger. # even when len(prefix_indices) is intentionally larger.
req2 = _DummyReq() req2 = _DummyReq()
req2.req_pool_idx = 1 req2.req_pool_idx = 1
req2.origin_input_ids = [11, 12, 13, 14, 15, 16] req2.origin_input_ids = array("q", [11, 12, 13, 14, 15, 16])
req2.output_ids = [] req2.output_ids = array("q")
req2._kv_committed_len = len(req2.origin_input_ids) req2._kv_committed_len = len(req2.origin_input_ids)
kv_indices2 = allocator.alloc(req2._kv_committed_len) kv_indices2 = allocator.alloc(req2._kv_committed_len)
req_to_token_pool.write( req_to_token_pool.write(
@@ -730,7 +739,9 @@ class TestSWASplitLeafOnInsert(CustomTestCase):
with envs.SGLANG_OPT_SWA_SPLIT_LEAF_ON_INSERT.override(True): with envs.SGLANG_OPT_SWA_SPLIT_LEAF_ON_INSERT.override(True):
inserted_leaf = _insert_chain(tree, allocator, token_ids) inserted_leaf = _insert_chain(tree, allocator, token_ids)
self.assertEqual(len(inserted_leaf.value), 4) self.assertEqual(len(inserted_leaf.value), 4)
match = tree.match_prefix(MatchPrefixParams(key=RadixKey(token_ids))) match = tree.match_prefix(
MatchPrefixParams(key=RadixKey(array("q", token_ids)))
)
self.assertEqual(match.device_indices.shape[0], 12) self.assertEqual(match.device_indices.shape[0], 12)
self.assertIs(match.last_device_node, inserted_leaf) self.assertIs(match.last_device_node, inserted_leaf)
test/registered/unit/mem_cache/test_unified_radix_cache_bench.py
+10 -9
View File
@@ -12,6 +12,7 @@ import random
import statistics import statistics
import time import time
import unittest import unittest
from array import array
from contextlib import contextmanager from contextlib import contextmanager
from dataclasses import dataclass from dataclasses import dataclass
from typing import Callable from typing import Callable
@@ -335,7 +336,7 @@ def _insert_seq(env, seq):
if env.has_mamba: if env.has_mamba:
req = env.make_req() req = env.make_req()
mamba_val = req.mamba_pool_idx.unsqueeze(0) mamba_val = req.mamba_pool_idx.unsqueeze(0)
key = RadixKey(seq) key = RadixKey(array("q", seq))
env.tree.insert(InsertParams(key=key, value=v[: len(key)], mamba_value=mamba_val)) env.tree.insert(InsertParams(key=key, value=v[: len(key)], mamba_value=mamba_val))
return True return True
@@ -357,7 +358,7 @@ def _fill_no_evict(env):
if env.has_mamba: if env.has_mamba:
req = env.make_req() req = env.make_req()
mamba_val = req.mamba_pool_idx.unsqueeze(0) mamba_val = req.mamba_pool_idx.unsqueeze(0)
key = RadixKey(seq) key = RadixKey(array("q", seq))
env.tree.insert( env.tree.insert(
InsertParams(key=key, value=v[: len(key)], mamba_value=mamba_val) InsertParams(key=key, value=v[: len(key)], mamba_value=mamba_val)
) )
@@ -505,7 +506,7 @@ def bench_match_prefix(
queries.append([rng.randint(1, 32000)] * rng.randint(50, 300)) queries.append([rng.randint(1, 32000)] * rng.randint(50, 300))
def verify_fn(q): def verify_fn(q):
k = RadixKey(q) k = RadixKey(array("q", q))
r1 = env.tree.match_prefix(MatchPrefixParams(key=k)) r1 = env.tree.match_prefix(MatchPrefixParams(key=k))
r2 = env.tree.match_prefix(MatchPrefixParams(key=k)) r2 = env.tree.match_prefix(MatchPrefixParams(key=k))
assert len(r1.device_indices) == len(r2.device_indices), "match not idempotent" assert len(r1.device_indices) == len(r2.device_indices), "match not idempotent"
@@ -514,7 +515,7 @@ def bench_match_prefix(
return bench_api( return bench_api(
"match_prefix", "match_prefix",
lambda: queries, lambda: queries,
lambda q: env.tree.match_prefix(MatchPrefixParams(key=RadixKey(q))), lambda q: env.tree.match_prefix(MatchPrefixParams(key=RadixKey(array("q", q)))),
min(len(queries) - warmup, num_seqs), min(len(queries) - warmup, num_seqs),
env.avg_tokens, env.avg_tokens,
warmup, warmup,
@@ -566,7 +567,7 @@ def bench_lock_unlock(
nodes = [] nodes = []
for seq in env.seqs[: num_seqs // 2]: for seq in env.seqs[: num_seqs // 2]:
r = env.tree.match_prefix(MatchPrefixParams(key=RadixKey(seq))) r = env.tree.match_prefix(MatchPrefixParams(key=RadixKey(array("q", seq))))
if r.last_device_node != env.tree.root_node: if r.last_device_node != env.tree.root_node:
nodes.append(r.last_device_node) nodes.append(r.last_device_node)
if not nodes: if not nodes:
@@ -613,7 +614,7 @@ def bench_cache_finished(
# Pre-build Req objects with token IDs filled into req_to_token # Pre-build Req objects with token IDs filled into req_to_token
req_items: list = [] req_items: list = []
for seq in env.seqs: for seq in env.seqs:
key = RadixKey(seq) key = RadixKey(array("q", seq))
mr = env.tree.match_prefix(MatchPrefixParams(key=key)) mr = env.tree.match_prefix(MatchPrefixParams(key=key))
matched_len = len(mr.device_indices) matched_len = len(mr.device_indices)
node = mr.last_device_node node = mr.last_device_node
@@ -635,9 +636,9 @@ def bench_cache_finished(
kv_indices = mr.device_indices kv_indices = mr.device_indices
req = env.make_req() req = env.make_req()
req.origin_input_ids = list(seq) req.origin_input_ids = array("q", seq)
req.output_ids = [] req.output_ids = array("q")
req.fill_ids = list(seq) req.fill_ids = array("q", seq)
req.last_node = node req.last_node = node
req.cache_protected_len = matched_len req.cache_protected_len = matched_len
req.kv_committed_len = len(seq) req.kv_committed_len = len(seq)
test/registered/unit/mem_cache/test_unified_radix_cache_unittest.py
+113 -92
View File
@@ -1,6 +1,7 @@
"""Unit tests for UnifiedRadixCache""" """Unit tests for UnifiedRadixCache"""
import unittest import unittest
from array import array
from dataclasses import dataclass from dataclasses import dataclass
from typing import Optional from typing import Optional
from unittest import mock from unittest import mock
@@ -272,7 +273,7 @@ class UnifiedRadixCacheSuite:
def _insert(self, tree, allocator, req_to_token_pool, tokens): def _insert(self, tree, allocator, req_to_token_pool, tokens):
"""Insert tokens, attaching mamba data when the config has mamba.""" """Insert tokens, attaching mamba data when the config has mamba."""
key = RadixKey(tokens) key = RadixKey(array("q", tokens))
value = self._alloc(allocator, len(tokens)) value = self._alloc(allocator, len(tokens))
params = InsertParams(key=key, value=value[: len(key)]) params = InsertParams(key=key, value=value[: len(key)])
if self.cfg.has_mamba: if self.cfg.has_mamba:
@@ -290,15 +291,17 @@ class UnifiedRadixCacheSuite:
result = self._insert(tree, allocator, req_to_token_pool, seq_b) result = self._insert(tree, allocator, req_to_token_pool, seq_b)
self.assertEqual(result.prefix_len, len(seq_a)) self.assertEqual(result.prefix_len, len(seq_a))
m = tree.match_prefix(MatchPrefixParams(key=RadixKey(seq_b))) m = tree.match_prefix(MatchPrefixParams(key=RadixKey(array("q", seq_b))))
self.assertEqual(len(m.device_indices), len(seq_b)) self.assertEqual(len(m.device_indices), len(seq_b))
m = tree.match_prefix( m = tree.match_prefix(
MatchPrefixParams(key=RadixKey(seq_a + self._make_seq(9000, 1))) MatchPrefixParams(key=RadixKey(array("q", seq_a + self._make_seq(9000, 1))))
) )
self.assertEqual(len(m.device_indices), len(seq_a)) self.assertEqual(len(m.device_indices), len(seq_a))
m = tree.match_prefix(MatchPrefixParams(key=RadixKey(self._make_seq(5000, 2)))) m = tree.match_prefix(
MatchPrefixParams(key=RadixKey(array("q", self._make_seq(5000, 2))))
)
self.assertEqual(len(m.device_indices), 0) self.assertEqual(len(m.device_indices), 0)
tree.sanity_check() tree.sanity_check()
@@ -317,11 +320,11 @@ class UnifiedRadixCacheSuite:
self.assertEqual(result_b.prefix_len, len(base)) self.assertEqual(result_b.prefix_len, len(base))
for seq in (branch_a, branch_b): for seq in (branch_a, branch_b):
m = tree.match_prefix(MatchPrefixParams(key=RadixKey(seq))) m = tree.match_prefix(MatchPrefixParams(key=RadixKey(array("q", seq))))
self.assertEqual(len(m.device_indices), len(seq)) self.assertEqual(len(m.device_indices), len(seq))
m = tree.match_prefix( m = tree.match_prefix(
MatchPrefixParams(key=RadixKey(base + self._make_seq(999, 1))) MatchPrefixParams(key=RadixKey(array("q", base + self._make_seq(999, 1))))
) )
self.assertEqual(len(m.device_indices), len(base)) self.assertEqual(len(m.device_indices), len(base))
tree.sanity_check() tree.sanity_check()
@@ -350,13 +353,13 @@ class UnifiedRadixCacheSuite:
self._insert(tree, allocator, req_to_token_pool, seq_a) self._insert(tree, allocator, req_to_token_pool, seq_a)
self._insert(tree, allocator, req_to_token_pool, seq_b) self._insert(tree, allocator, req_to_token_pool, seq_b)
m = tree.match_prefix(MatchPrefixParams(key=RadixKey(seq_a))) m = tree.match_prefix(MatchPrefixParams(key=RadixKey(array("q", seq_a))))
lock_result = tree.inc_lock_ref(m.last_device_node) lock_result = tree.inc_lock_ref(m.last_device_node)
result = tree.evict(EvictParams(num_tokens=len(seq_a) + len(seq_b))) result = tree.evict(EvictParams(num_tokens=len(seq_a) + len(seq_b)))
self.assertGreaterEqual(result.num_tokens_evicted, len(seq_b)) self.assertGreaterEqual(result.num_tokens_evicted, len(seq_b))
m = tree.match_prefix(MatchPrefixParams(key=RadixKey(seq_a))) m = tree.match_prefix(MatchPrefixParams(key=RadixKey(array("q", seq_a))))
self.assertEqual(len(m.device_indices), len(seq_a)) self.assertEqual(len(m.device_indices), len(seq_a))
# Unlock -> should now be evictable # Unlock -> should now be evictable
@@ -395,7 +398,7 @@ class UnifiedRadixCacheSuite:
if self.cfg.has_mamba: if self.cfg.has_mamba:
self.assertEqual(tree.mamba_evictable_size(), 0) self.assertEqual(tree.mamba_evictable_size(), 0)
m = tree.match_prefix(MatchPrefixParams(key=RadixKey(seqs[0]))) m = tree.match_prefix(MatchPrefixParams(key=RadixKey(array("q", seqs[0]))))
self.assertEqual(len(m.device_indices), 0) self.assertEqual(len(m.device_indices), 0)
tree.sanity_check() tree.sanity_check()
@@ -413,7 +416,7 @@ class UnifiedRadixCacheSuite:
self.assertEqual(allocator.available_size(), initial_avail - len(seq_1p)) self.assertEqual(allocator.available_size(), initial_avail - len(seq_1p))
# Step 2: insert 2 pages with prev_prefix_len=0 → frees overlap of 1 page # Step 2: insert 2 pages with prev_prefix_len=0 → frees overlap of 1 page
key_2p = RadixKey(seq_2p) key_2p = RadixKey(array("q", seq_2p))
value_2p = self._alloc(allocator, len(seq_2p)) value_2p = self._alloc(allocator, len(seq_2p))
params = InsertParams( params = InsertParams(
key=key_2p, key=key_2p,
@@ -432,7 +435,7 @@ class UnifiedRadixCacheSuite:
# Step 3: insert 3 pages with prev_prefix_len=len(seq_2p) → nothing freed # Step 3: insert 3 pages with prev_prefix_len=len(seq_2p) → nothing freed
avail_before = allocator.available_size() avail_before = allocator.available_size()
key_3p = RadixKey(seq_3p) key_3p = RadixKey(array("q", seq_3p))
value_3p = self._alloc(allocator, len(seq_3p)) value_3p = self._alloc(allocator, len(seq_3p))
params = InsertParams( params = InsertParams(
key=key_3p, key=key_3p,
@@ -461,11 +464,11 @@ class UnifiedRadixCacheSuite:
self.assertEqual(result.prefix_len, len(base)) self.assertEqual(result.prefix_len, len(base))
for seq in (fork_a, fork_b): for seq in (fork_a, fork_b):
m = tree.match_prefix(MatchPrefixParams(key=RadixKey(seq))) m = tree.match_prefix(MatchPrefixParams(key=RadixKey(array("q", seq))))
self.assertEqual(len(m.device_indices), len(seq)) self.assertEqual(len(m.device_indices), len(seq))
m = tree.match_prefix( m = tree.match_prefix(
MatchPrefixParams(key=RadixKey(base + self._make_seq(999, 1))) MatchPrefixParams(key=RadixKey(array("q", base + self._make_seq(999, 1))))
) )
self.assertEqual(len(m.device_indices), len(base)) self.assertEqual(len(m.device_indices), len(base))
tree.sanity_check() tree.sanity_check()
@@ -477,8 +480,8 @@ class UnifiedRadixCacheSuite:
req = self._make_req(req_to_token_pool) req = self._make_req(req_to_token_pool)
input_ids = self._make_seq(1, 3) input_ids = self._make_seq(1, 3)
output_ids = self._make_seq(2000, 1) output_ids = self._make_seq(2000, 1)
req.origin_input_ids = input_ids req.origin_input_ids = array("q", input_ids)
req.output_ids = output_ids req.output_ids = array("q", output_ids)
kv_len = len(input_ids) + len(output_ids) kv_len = len(input_ids) + len(output_ids)
kv_indices = self._alloc(allocator, kv_len) kv_indices = self._alloc(allocator, kv_len)
req_to_token_pool.write((req.req_pool_idx, slice(0, kv_len)), kv_indices) req_to_token_pool.write((req.req_pool_idx, slice(0, kv_len)), kv_indices)
@@ -487,7 +490,7 @@ class UnifiedRadixCacheSuite:
req.cache_protected_len = 0 req.cache_protected_len = 0
req.swa_uuid_for_lock = None req.swa_uuid_for_lock = None
req.extra_key = None req.extra_key = None
req.fill_ids = input_ids + output_ids req.fill_ids = array("q", input_ids + output_ids)
if self.cfg.has_mamba: if self.cfg.has_mamba:
req.mamba_last_track_seqlen = kv_len req.mamba_last_track_seqlen = kv_len
@@ -495,7 +498,9 @@ class UnifiedRadixCacheSuite:
all_ids = input_ids + output_ids all_ids = input_ids + output_ids
aligned_len = (len(all_ids) // ps) * ps aligned_len = (len(all_ids) // ps) * ps
m = tree.match_prefix(MatchPrefixParams(key=RadixKey(all_ids[:aligned_len]))) m = tree.match_prefix(
MatchPrefixParams(key=RadixKey(array("q", all_ids[:aligned_len])))
)
self.assertEqual(len(m.device_indices), aligned_len) self.assertEqual(len(m.device_indices), aligned_len)
tree.sanity_check() tree.sanity_check()
@@ -506,9 +511,9 @@ class UnifiedRadixCacheSuite:
req = self._make_req(req_to_token_pool) req = self._make_req(req_to_token_pool)
prompt_ids = self._make_seq(1, 3) prompt_ids = self._make_seq(1, 3)
output_ids = self._make_seq(2000, 7) output_ids = self._make_seq(2000, 7)
req.origin_input_ids = prompt_ids req.origin_input_ids = array("q", prompt_ids)
req.output_ids = output_ids req.output_ids = array("q", output_ids)
req.fill_ids = prompt_ids + output_ids req.fill_ids = array("q", prompt_ids + output_ids)
kv_len = len(req.fill_ids) kv_len = len(req.fill_ids)
kv_indices = self._alloc(allocator, kv_len) kv_indices = self._alloc(allocator, kv_len)
req_to_token_pool.write((req.req_pool_idx, slice(0, kv_len)), kv_indices) req_to_token_pool.write((req.req_pool_idx, slice(0, kv_len)), kv_indices)
@@ -538,7 +543,9 @@ class UnifiedRadixCacheSuite:
prompt_aligned = (len(prompt_ids) // ps) * ps prompt_aligned = (len(prompt_ids) // ps) * ps
# Thinking+answer must not be reachable past the prompt. # Thinking+answer must not be reachable past the prompt.
m = tree.match_prefix(MatchPrefixParams(key=RadixKey(prompt_ids + output_ids))) m = tree.match_prefix(
MatchPrefixParams(key=RadixKey(array("q", prompt_ids + output_ids)))
)
self.assertEqual(len(m.device_indices), prompt_aligned) self.assertEqual(len(m.device_indices), prompt_aligned)
# Only prompt-aligned pages remain owned by the tree. # Only prompt-aligned pages remain owned by the tree.
self.assertEqual( self.assertEqual(
@@ -550,8 +557,8 @@ class UnifiedRadixCacheSuite:
tree, allocator, req_to_token_pool = build_fixture(self.cfg) tree, allocator, req_to_token_pool = build_fixture(self.cfg)
req = self._make_req(req_to_token_pool) req = self._make_req(req_to_token_pool)
tokens = self._make_seq(1, 2) tokens = self._make_seq(1, 2)
req.origin_input_ids = tokens req.origin_input_ids = array("q", tokens)
req.output_ids = [] req.output_ids = array("q")
kv_len = len(tokens) kv_len = len(tokens)
kv_indices = self._alloc(allocator, kv_len) kv_indices = self._alloc(allocator, kv_len)
req_to_token_pool.write((req.req_pool_idx, slice(0, kv_len)), kv_indices) req_to_token_pool.write((req.req_pool_idx, slice(0, kv_len)), kv_indices)
@@ -560,13 +567,13 @@ class UnifiedRadixCacheSuite:
req.cache_protected_len = 0 req.cache_protected_len = 0
req.swa_uuid_for_lock = None req.swa_uuid_for_lock = None
req.extra_key = None req.extra_key = None
req.fill_ids = tokens req.fill_ids = array("q", tokens)
avail_before = allocator.available_size() avail_before = allocator.available_size()
tree.cache_finished_req(req, is_insert=False) tree.cache_finished_req(req, is_insert=False)
self.assertEqual(allocator.available_size(), avail_before + kv_len) self.assertEqual(allocator.available_size(), avail_before + kv_len)
m = tree.match_prefix(MatchPrefixParams(key=RadixKey(tokens))) m = tree.match_prefix(MatchPrefixParams(key=RadixKey(array("q", tokens))))
self.assertEqual(len(m.device_indices), 0) self.assertEqual(len(m.device_indices), 0)
tree.sanity_check() tree.sanity_check()
@@ -575,9 +582,9 @@ class UnifiedRadixCacheSuite:
req = self._make_req(req_to_token_pool) req = self._make_req(req_to_token_pool)
tokens = self._make_seq(1, 3) tokens = self._make_seq(1, 3)
req.origin_input_ids = tokens req.origin_input_ids = array("q", tokens)
req.output_ids = [] req.output_ids = array("q")
req.fill_ids = tokens[:] req.fill_ids = array("q", tokens)
kv_len = len(tokens) kv_len = len(tokens)
kv_indices = self._alloc(allocator, kv_len) kv_indices = self._alloc(allocator, kv_len)
req_to_token_pool.write((req.req_pool_idx, slice(0, kv_len)), kv_indices) req_to_token_pool.write((req.req_pool_idx, slice(0, kv_len)), kv_indices)
@@ -628,11 +635,11 @@ class UnifiedRadixCacheSuite:
for suffix_start in [100, 200, 300]: for suffix_start in [100, 200, 300]:
seq = base + self._make_seq(suffix_start, 2) seq = base + self._make_seq(suffix_start, 2)
m = tree.match_prefix(MatchPrefixParams(key=RadixKey(seq))) m = tree.match_prefix(MatchPrefixParams(key=RadixKey(array("q", seq))))
self.assertEqual(len(m.device_indices), len(seq)) self.assertEqual(len(m.device_indices), len(seq))
m = tree.match_prefix( m = tree.match_prefix(
MatchPrefixParams(key=RadixKey(base + self._make_seq(999, 1))) MatchPrefixParams(key=RadixKey(array("q", base + self._make_seq(999, 1))))
) )
self.assertEqual(len(m.device_indices), len(base)) self.assertEqual(len(m.device_indices), len(base))
tree.sanity_check() tree.sanity_check()
@@ -641,7 +648,7 @@ class UnifiedRadixCacheSuite:
if self.cfg.page_size == 1: if self.cfg.page_size == 1:
self.skipTest("page_size > 1 only") self.skipTest("page_size > 1 only")
tree, _, _ = build_fixture(self.cfg) tree, _, _ = build_fixture(self.cfg)
key = RadixKey(self._make_seq(1, 1)) key = RadixKey(array("q", self._make_seq(1, 1)))
child_key = key.child_key(tree.page_size) child_key = key.child_key(tree.page_size)
self.assertIsInstance(child_key, tuple) self.assertIsInstance(child_key, tuple)
@@ -656,11 +663,13 @@ class UnifiedRadixCacheSuite:
# Tree truncates unaligned tail internally, so it matches the seq prefix. # Tree truncates unaligned tail internally, so it matches the seq prefix.
unaligned = seq + list(range(9000, 9000 + ps - 1)) unaligned = seq + list(range(9000, 9000 + ps - 1))
m = tree.match_prefix(MatchPrefixParams(key=RadixKey(unaligned))) m = tree.match_prefix(MatchPrefixParams(key=RadixKey(array("q", unaligned))))
self.assertEqual(len(m.device_indices), len(seq)) self.assertEqual(len(m.device_indices), len(seq))
# Below-page-size key aligns to 0 -> no match. # Below-page-size key aligns to 0 -> no match.
m = tree.match_prefix(MatchPrefixParams(key=RadixKey(seq[: ps - 1]))) m = tree.match_prefix(
MatchPrefixParams(key=RadixKey(array("q", seq[: ps - 1])))
)
self.assertEqual(len(m.device_indices), 0) self.assertEqual(len(m.device_indices), 0)
tree.sanity_check() tree.sanity_check()
@@ -679,12 +688,12 @@ class UnifiedRadixCacheSuite:
# Mismatch in second page → only first page matches # Mismatch in second page → only first page matches
bad_page2 = seq[:ps] + [9999] * ps bad_page2 = seq[:ps] + [9999] * ps
m = tree.match_prefix(MatchPrefixParams(key=RadixKey(bad_page2))) m = tree.match_prefix(MatchPrefixParams(key=RadixKey(array("q", bad_page2))))
self.assertEqual(len(m.device_indices), ps) self.assertEqual(len(m.device_indices), ps)
# Mismatch in first page → 0 match # Mismatch in first page → 0 match
bad_page1 = [9999] + seq[1:] bad_page1 = [9999] + seq[1:]
m = tree.match_prefix(MatchPrefixParams(key=RadixKey(bad_page1))) m = tree.match_prefix(MatchPrefixParams(key=RadixKey(array("q", bad_page1))))
self.assertEqual(len(m.device_indices), 0) self.assertEqual(len(m.device_indices), 0)
tree.sanity_check() tree.sanity_check()
@@ -699,8 +708,8 @@ class UnifiedRadixCacheSuite:
tail_extra = ps // 2 tail_extra = ps // 2
input_ids = self._make_seq(1, 1) + list(range(8000, 8000 + tail_extra)) input_ids = self._make_seq(1, 1) + list(range(8000, 8000 + tail_extra))
req = self._make_req(req_to_token_pool) req = self._make_req(req_to_token_pool)
req.origin_input_ids = input_ids req.origin_input_ids = array("q", input_ids)
req.output_ids = [] req.output_ids = array("q")
kv_len = len(input_ids) kv_len = len(input_ids)
kv_indices = self._alloc(allocator, kv_len) kv_indices = self._alloc(allocator, kv_len)
req_to_token_pool.write((req.req_pool_idx, slice(0, kv_len)), kv_indices) req_to_token_pool.write((req.req_pool_idx, slice(0, kv_len)), kv_indices)
@@ -709,7 +718,7 @@ class UnifiedRadixCacheSuite:
req.cache_protected_len = 0 req.cache_protected_len = 0
req.swa_uuid_for_lock = None req.swa_uuid_for_lock = None
req.extra_key = None req.extra_key = None
req.fill_ids = input_ids req.fill_ids = array("q", input_ids)
if self.cfg.has_mamba: if self.cfg.has_mamba:
req.mamba_last_track_seqlen = kv_len req.mamba_last_track_seqlen = kv_len
@@ -718,7 +727,7 @@ class UnifiedRadixCacheSuite:
self.assertEqual(allocator.available_size(), avail_before + tail_extra) self.assertEqual(allocator.available_size(), avail_before + tail_extra)
aligned = input_ids[: (len(input_ids) // ps) * ps] aligned = input_ids[: (len(input_ids) // ps) * ps]
m = tree.match_prefix(MatchPrefixParams(key=RadixKey(aligned))) m = tree.match_prefix(MatchPrefixParams(key=RadixKey(array("q", aligned))))
self.assertEqual(len(m.device_indices), len(aligned)) self.assertEqual(len(m.device_indices), len(aligned))
tree.sanity_check() tree.sanity_check()
@@ -749,7 +758,7 @@ class UnifiedRadixCacheSuite:
tree.evict(EvictParams(num_tokens=0, mamba_num=10)) tree.evict(EvictParams(num_tokens=0, mamba_num=10))
self.assertEqual(tree.mamba_evictable_size(), 0) self.assertEqual(tree.mamba_evictable_size(), 0)
m = tree.match_prefix(MatchPrefixParams(key=RadixKey(seq_long))) m = tree.match_prefix(MatchPrefixParams(key=RadixKey(array("q", seq_long))))
self.assertEqual(len(m.device_indices), 0) self.assertEqual(len(m.device_indices), 0)
tree.sanity_check() tree.sanity_check()
@@ -788,7 +797,7 @@ class UnifiedRadixCacheSuite:
req2 = self._make_req(req_to_token_pool) req2 = self._make_req(req_to_token_pool)
m = tree.match_prefix( m = tree.match_prefix(
MatchPrefixParams(key=RadixKey(seq), cow_mamba=True, req=req2) MatchPrefixParams(key=RadixKey(array("q", seq)), cow_mamba=True, req=req2)
) )
self.assertEqual(len(m.device_indices), len(seq)) self.assertEqual(len(m.device_indices), len(seq))
self.assertIsNotNone(req2.mamba_pool_idx) self.assertIsNotNone(req2.mamba_pool_idx)
@@ -809,7 +818,7 @@ class UnifiedRadixCacheSuite:
seq = self._make_seq(1, 3) seq = self._make_seq(1, 3)
self._insert(tree, allocator, req_to_token_pool, seq) self._insert(tree, allocator, req_to_token_pool, seq)
m = tree.match_prefix(MatchPrefixParams(key=RadixKey(seq))) m = tree.match_prefix(MatchPrefixParams(key=RadixKey(array("q", seq))))
self.assertEqual(len(m.device_indices), len(seq)) self.assertEqual(len(m.device_indices), len(seq))
tree.sanity_check() tree.sanity_check()
@@ -866,13 +875,13 @@ class UnifiedRadixCacheSuite:
self._insert(tree, allocator, req_to_token_pool, seq_a) self._insert(tree, allocator, req_to_token_pool, seq_a)
self._insert(tree, allocator, req_to_token_pool, seq_b) self._insert(tree, allocator, req_to_token_pool, seq_b)
m = tree.match_prefix(MatchPrefixParams(key=RadixKey(seq_a))) m = tree.match_prefix(MatchPrefixParams(key=RadixKey(array("q", seq_a))))
lock_result = tree.inc_lock_ref(m.last_device_node) lock_result = tree.inc_lock_ref(m.last_device_node)
result = tree.evict(EvictParams(num_tokens=len(seq_a) + len(seq_b))) result = tree.evict(EvictParams(num_tokens=len(seq_a) + len(seq_b)))
self.assertGreaterEqual(result.num_tokens_evicted, len(seq_b)) self.assertGreaterEqual(result.num_tokens_evicted, len(seq_b))
m = tree.match_prefix(MatchPrefixParams(key=RadixKey(seq_a))) m = tree.match_prefix(MatchPrefixParams(key=RadixKey(array("q", seq_a))))
self.assertEqual(len(m.device_indices), len(seq_a)) self.assertEqual(len(m.device_indices), len(seq_a))
tree.dec_lock_ref( tree.dec_lock_ref(
@@ -886,8 +895,8 @@ class UnifiedRadixCacheSuite:
parent = UnifiedTreeNode(self.cfg.components) parent = UnifiedTreeNode(self.cfg.components)
deleted = UnifiedTreeNode(self.cfg.components) deleted = UnifiedTreeNode(self.cfg.components)
parent.key = RadixKey(self._make_seq(1, 1)) parent.key = RadixKey(array("q", self._make_seq(1, 1)))
deleted.key = RadixKey(self._make_seq(1000, 1)) deleted.key = RadixKey(array("q", self._make_seq(1000, 1)))
parent.parent = tree.root_node parent.parent = tree.root_node
deleted.parent = parent deleted.parent = parent
parent.component_data[ComponentType.FULL].value = torch.arange( parent.component_data[ComponentType.FULL].value = torch.arange(
@@ -924,15 +933,15 @@ class UnifiedRadixCacheSuite:
node_count_before = count_nodes(tree.root_node) node_count_before = count_nodes(tree.root_node)
self.assertEqual(node_count_before, 2) self.assertEqual(node_count_before, 2)
tree._match_prefix_helper(RadixKey([1, 2])) tree._match_prefix_helper(RadixKey(array("q", [1, 2])))
( (
value, value,
best_match_node, best_match_node,
best_match_device_node, best_match_device_node,
best_value_len, best_value_len,
) = tree._match_prefix_helper(RadixKey([1, 2, 3, 4])) ) = tree._match_prefix_helper(RadixKey(array("q", [1, 2, 3, 4])))
self.assertEqual(best_value_len, 2) self.assertEqual(best_value_len, 2)
self.assertEqual(best_match_node.key.token_ids, [3, 4]) self.assertEqual(list(best_match_node.key.token_ids), [3, 4])
self.assertIs(best_match_device_node, best_match_node) self.assertIs(best_match_device_node, best_match_node)
node_count_after_regular = count_nodes(tree.root_node) node_count_after_regular = count_nodes(tree.root_node)
self.assertEqual(node_count_after_regular, node_count_before + 2) self.assertEqual(node_count_after_regular, node_count_before + 2)
@@ -942,9 +951,9 @@ class UnifiedRadixCacheSuite:
best_match_node, best_match_node,
best_match_device_node, best_match_device_node,
best_value_len, best_value_len,
) = tree._match_prefix_helper_readonly(RadixKey([1, 2, 3])) ) = tree._match_prefix_helper_readonly(RadixKey(array("q", [1, 2, 3])))
self.assertEqual(best_value_len, 1) self.assertEqual(best_value_len, 1)
self.assertEqual(best_match_node.key.token_ids, [1, 2]) self.assertEqual(list(best_match_node.key.token_ids), [1, 2])
self.assertIs(best_match_device_node, best_match_node) self.assertIs(best_match_device_node, best_match_node)
node_count_after_readonly = count_nodes(tree.root_node) node_count_after_readonly = count_nodes(tree.root_node)
self.assertEqual(node_count_after_readonly, node_count_after_regular) self.assertEqual(node_count_after_readonly, node_count_after_regular)
@@ -971,7 +980,7 @@ class UnifiedRadixCacheSuite:
seq = self._make_seq(1, 2) seq = self._make_seq(1, 2)
self._insert(tree, allocator, req_to_token_pool, seq) self._insert(tree, allocator, req_to_token_pool, seq)
match = tree.match_prefix(MatchPrefixParams(key=RadixKey(seq))) match = tree.match_prefix(MatchPrefixParams(key=RadixKey(array("q", seq))))
node = match.last_device_node node = match.last_device_node
full_cd = node.component_data[ComponentType.FULL] full_cd = node.component_data[ComponentType.FULL]
aux_cd = node.component_data[aux] aux_cd = node.component_data[aux]
@@ -1040,7 +1049,7 @@ class UnifiedRadixCacheSuite:
self._insert(tree, allocator, req_to_token_pool, leaf) self._insert(tree, allocator, req_to_token_pool, leaf)
# Lock the base node to prevent it from being evicted # Lock the base node to prevent it from being evicted
m_base = tree.match_prefix(MatchPrefixParams(key=RadixKey(base))) m_base = tree.match_prefix(MatchPrefixParams(key=RadixKey(array("q", base))))
lock_result = tree.inc_lock_ref(m_base.last_device_node) lock_result = tree.inc_lock_ref(m_base.last_device_node)
# Evict the leaf — parent (base) should become D-leaf after unlock # Evict the leaf — parent (base) should become D-leaf after unlock
@@ -1089,14 +1098,14 @@ class UnifiedRadixCacheSuite:
self._insert(tree, allocator, req_to_token_pool, seq_new) self._insert(tree, allocator, req_to_token_pool, seq_new)
# Touch seq_new to make it MRU # Touch seq_new to make it MRU
tree.match_prefix(MatchPrefixParams(key=RadixKey(seq_new))) tree.match_prefix(MatchPrefixParams(key=RadixKey(array("q", seq_new))))
# Evict just enough for one sequence # Evict just enough for one sequence
tree.evict(EvictParams(num_tokens=len(seq_old))) tree.evict(EvictParams(num_tokens=len(seq_old)))
# seq_old should be gone (LRU), seq_new should remain # seq_old should be gone (LRU), seq_new should remain
m_old = tree.match_prefix(MatchPrefixParams(key=RadixKey(seq_old))) m_old = tree.match_prefix(MatchPrefixParams(key=RadixKey(array("q", seq_old))))
m_new = tree.match_prefix(MatchPrefixParams(key=RadixKey(seq_new))) m_new = tree.match_prefix(MatchPrefixParams(key=RadixKey(array("q", seq_new))))
self.assertEqual(len(m_old.device_indices), 0) self.assertEqual(len(m_old.device_indices), 0)
self.assertEqual(len(m_new.device_indices), len(seq_new)) self.assertEqual(len(m_new.device_indices), len(seq_new))
tree.sanity_check() tree.sanity_check()
@@ -1136,13 +1145,13 @@ class UnifiedRadixCacheSuite:
self._insert(tree, allocator, req_to_token_pool, branch_b) self._insert(tree, allocator, req_to_token_pool, branch_b)
# Lock branch_b # Lock branch_b
m = tree.match_prefix(MatchPrefixParams(key=RadixKey(branch_b))) m = tree.match_prefix(MatchPrefixParams(key=RadixKey(array("q", branch_b))))
lr = tree.inc_lock_ref(m.last_device_node) lr = tree.inc_lock_ref(m.last_device_node)
# Evict — branch_a should go, base + branch_b stay # Evict — branch_a should go, base + branch_b stay
tree.evict(EvictParams(num_tokens=len(branch_a))) tree.evict(EvictParams(num_tokens=len(branch_a)))
m_b = tree.match_prefix(MatchPrefixParams(key=RadixKey(branch_b))) m_b = tree.match_prefix(MatchPrefixParams(key=RadixKey(array("q", branch_b))))
self.assertEqual(len(m_b.device_indices), len(branch_b)) self.assertEqual(len(m_b.device_indices), len(branch_b))
tree.dec_lock_ref( tree.dec_lock_ref(
@@ -1173,7 +1182,7 @@ class UnifiedRadixCacheSuite:
self._insert(tree, allocator, req_to_token_pool, seq_b) self._insert(tree, allocator, req_to_token_pool, seq_b)
# Lock seq_a # Lock seq_a
m = tree.match_prefix(MatchPrefixParams(key=RadixKey(seq_a))) m = tree.match_prefix(MatchPrefixParams(key=RadixKey(array("q", seq_a))))
lr = tree.inc_lock_ref(m.last_device_node) lr = tree.inc_lock_ref(m.last_device_node)
# Try to evict everything # Try to evict everything
@@ -1181,7 +1190,7 @@ class UnifiedRadixCacheSuite:
result = tree.evict(EvictParams(num_tokens=total)) result = tree.evict(EvictParams(num_tokens=total))
# seq_a should still be matchable (protected) # seq_a should still be matchable (protected)
m2 = tree.match_prefix(MatchPrefixParams(key=RadixKey(seq_a))) m2 = tree.match_prefix(MatchPrefixParams(key=RadixKey(array("q", seq_a))))
self.assertEqual(len(m2.device_indices), len(seq_a)) self.assertEqual(len(m2.device_indices), len(seq_a))
tree.dec_lock_ref( tree.dec_lock_ref(
@@ -1220,7 +1229,7 @@ class UnifiedRadixCacheSuite:
# Re-insert # Re-insert
seq_b = self._make_seq(500, 2) seq_b = self._make_seq(500, 2)
self._insert(tree, allocator, req_to_token_pool, seq_b) self._insert(tree, allocator, req_to_token_pool, seq_b)
m = tree.match_prefix(MatchPrefixParams(key=RadixKey(seq_b))) m = tree.match_prefix(MatchPrefixParams(key=RadixKey(array("q", seq_b))))
self.assertEqual(len(m.device_indices), len(seq_b)) self.assertEqual(len(m.device_indices), len(seq_b))
tree.sanity_check() tree.sanity_check()
@@ -1247,7 +1256,7 @@ class UnifiedRadixCacheSuite:
self._insert(tree, allocator, req_to_token_pool, s) self._insert(tree, allocator, req_to_token_pool, s)
# Lock some, evict some, unlock # Lock some, evict some, unlock
m = tree.match_prefix(MatchPrefixParams(key=RadixKey(seqs[0]))) m = tree.match_prefix(MatchPrefixParams(key=RadixKey(array("q", seqs[0]))))
lr = tree.inc_lock_ref(m.last_device_node) lr = tree.inc_lock_ref(m.last_device_node)
tree.evict(EvictParams(num_tokens=len(seqs[1]))) tree.evict(EvictParams(num_tokens=len(seqs[1])))
@@ -1409,7 +1418,7 @@ class UnifiedRadixCacheSuite:
self._insert(tree, allocator, req_to_token_pool, seq) self._insert(tree, allocator, req_to_token_pool, seq)
# Find the leaf node # Find the leaf node
m = tree.match_prefix(MatchPrefixParams(key=RadixKey(seq))) m = tree.match_prefix(MatchPrefixParams(key=RadixKey(array("q", seq))))
node = m.last_device_node node = m.last_device_node
self.assertIsNot(node, tree.root_node) self.assertIsNot(node, tree.root_node)
@@ -1435,7 +1444,7 @@ class UnifiedRadixCacheSuite:
seq = self._make_seq(1, 2) seq = self._make_seq(1, 2)
self._insert(tree, allocator, req_to_token_pool, seq) self._insert(tree, allocator, req_to_token_pool, seq)
m = tree.match_prefix(MatchPrefixParams(key=RadixKey(seq))) m = tree.match_prefix(MatchPrefixParams(key=RadixKey(array("q", seq))))
node = m.last_device_node node = m.last_device_node
self._backup_node(tree, node) self._backup_node(tree, node)
@@ -1469,7 +1478,7 @@ class UnifiedRadixCacheSuite:
self._backup_tree(tree) self._backup_tree(tree)
# Lock leaf so only base can be evicted # Lock leaf so only base can be evicted
m = tree.match_prefix(MatchPrefixParams(key=RadixKey(leaf))) m = tree.match_prefix(MatchPrefixParams(key=RadixKey(array("q", leaf))))
lr = tree.inc_lock_ref(m.last_device_node) lr = tree.inc_lock_ref(m.last_device_node)
# Evict base (inner node won't be evicted while child is locked) # Evict base (inner node won't be evicted while child is locked)
@@ -1479,7 +1488,7 @@ class UnifiedRadixCacheSuite:
m.last_device_node, m.last_device_node,
DecLockRefParams(swa_uuid_for_lock=getattr(lr, "swa_uuid_for_lock", None)), DecLockRefParams(swa_uuid_for_lock=getattr(lr, "swa_uuid_for_lock", None)),
) )
m = tree.match_prefix(MatchPrefixParams(key=RadixKey(leaf))) m = tree.match_prefix(MatchPrefixParams(key=RadixKey(array("q", leaf))))
self.assertGreaterEqual(len(m.device_indices), len(base)) self.assertGreaterEqual(len(m.device_indices), len(base))
tree.sanity_check() tree.sanity_check()
@@ -1495,7 +1504,7 @@ class UnifiedRadixCacheSuite:
query = expected_prefix + self._make_seq(9000, 1) query = expected_prefix + self._make_seq(9000, 1)
self._insert(tree, allocator, req_to_token_pool, seq) self._insert(tree, allocator, req_to_token_pool, seq)
m = tree.match_prefix(MatchPrefixParams(key=RadixKey(seq))) m = tree.match_prefix(MatchPrefixParams(key=RadixKey(array("q", seq))))
node = m.last_device_node node = m.last_device_node
self._backup_node(tree, node) self._backup_node(tree, node)
@@ -1503,7 +1512,7 @@ class UnifiedRadixCacheSuite:
self.assertTrue(node.evicted) self.assertTrue(node.evicted)
self.assertTrue(node.backuped) self.assertTrue(node.backuped)
m = tree.match_prefix(MatchPrefixParams(key=RadixKey(query))) m = tree.match_prefix(MatchPrefixParams(key=RadixKey(array("q", query))))
self.assertEqual(len(m.device_indices), 0) self.assertEqual(len(m.device_indices), 0)
self.assertIs(m.last_device_node, tree.root_node) self.assertIs(m.last_device_node, tree.root_node)
@@ -1512,8 +1521,8 @@ class UnifiedRadixCacheSuite:
self.assertIsNot(split_parent, tree.root_node) self.assertIsNot(split_parent, tree.root_node)
self.assertTrue(split_parent.evicted) self.assertTrue(split_parent.evicted)
self.assertTrue(split_parent.backuped) self.assertTrue(split_parent.backuped)
self.assertEqual(split_parent.key.token_ids, expected_prefix) self.assertEqual(list(split_parent.key.token_ids), expected_prefix)
self.assertEqual(node.key.token_ids, expected_suffix) self.assertEqual(list(node.key.token_ids), expected_suffix)
if self.cfg.has_mamba: if self.cfg.has_mamba:
self.assertEqual(m.host_hit_length, 0) self.assertEqual(m.host_hit_length, 0)
@@ -1536,7 +1545,7 @@ class UnifiedRadixCacheSuite:
self._insert(tree, allocator, req_to_token_pool, s) self._insert(tree, allocator, req_to_token_pool, s)
for i in range(2): for i in range(2):
m = tree.match_prefix(MatchPrefixParams(key=RadixKey(seqs[i]))) m = tree.match_prefix(MatchPrefixParams(key=RadixKey(array("q", seqs[i]))))
self._backup_node(tree, m.last_device_node) self._backup_node(tree, m.last_device_node)
# Evict one backed-up node # Evict one backed-up node
@@ -1555,7 +1564,7 @@ class UnifiedRadixCacheSuite:
seq = self._make_seq(1, 2) seq = self._make_seq(1, 2)
self._insert(tree, allocator, req_to_token_pool, seq) self._insert(tree, allocator, req_to_token_pool, seq)
m = tree.match_prefix(MatchPrefixParams(key=RadixKey(seq))) m = tree.match_prefix(MatchPrefixParams(key=RadixKey(array("q", seq))))
node = m.last_device_node node = m.last_device_node
self._backup_node(tree, node) self._backup_node(tree, node)
@@ -1580,7 +1589,7 @@ class UnifiedRadixCacheSuite:
base = self._make_seq(1, 2) base = self._make_seq(1, 2)
self._insert(tree, allocator, req_to_token_pool, base) self._insert(tree, allocator, req_to_token_pool, base)
m = tree.match_prefix(MatchPrefixParams(key=RadixKey(base))) m = tree.match_prefix(MatchPrefixParams(key=RadixKey(array("q", base))))
node = m.last_device_node node = m.last_device_node
original_device_indices = m.device_indices.clone() original_device_indices = m.device_indices.clone()
self._fill_full_kv(allocator, original_device_indices, marker=3) self._fill_full_kv(allocator, original_device_indices, marker=3)
@@ -1662,7 +1671,7 @@ class UnifiedRadixCacheSuite:
seq = self._make_seq(1, 2) seq = self._make_seq(1, 2)
self._insert(tree, allocator, req_to_token_pool, seq) self._insert(tree, allocator, req_to_token_pool, seq)
m = tree.match_prefix(MatchPrefixParams(key=RadixKey(seq))) m = tree.match_prefix(MatchPrefixParams(key=RadixKey(array("q", seq))))
node = m.last_device_node node = m.last_device_node
for aux in aux_types: for aux in aux_types:
@@ -1688,7 +1697,7 @@ class UnifiedRadixCacheSuite:
for i in range(num_pages): for i in range(num_pages):
seq = seq + self._make_seq(1000 * (i + 1), 1) seq = seq + self._make_seq(1000 * (i + 1), 1)
self._insert(tree, allocator, req_to_token_pool, seq) self._insert(tree, allocator, req_to_token_pool, seq)
m = tree.match_prefix(MatchPrefixParams(key=RadixKey(seq))) m = tree.match_prefix(MatchPrefixParams(key=RadixKey(array("q", seq))))
chain: list = [] chain: list = []
cur = m.last_device_node cur = m.last_device_node
while cur is not tree.root_node: while cur is not tree.root_node:
@@ -1738,7 +1747,7 @@ class UnifiedRadixCacheSuite:
seq = self._make_seq(1, (min_tokens + ps - 1) // ps) seq = self._make_seq(1, (min_tokens + ps - 1) // ps)
self._insert(tree, allocator, req_to_token_pool, seq) self._insert(tree, allocator, req_to_token_pool, seq)
result = tree.match_prefix(MatchPrefixParams(key=RadixKey(seq))) result = tree.match_prefix(MatchPrefixParams(key=RadixKey(array("q", seq))))
self.assertEqual(len(result.device_indices), len(seq)) self.assertEqual(len(result.device_indices), len(seq))
self.assertIs(result.best_match_node, result.last_device_node) self.assertIs(result.best_match_node, result.last_device_node)
@@ -1760,7 +1769,7 @@ class UnifiedRadixCacheSuite:
tree.evict(EvictParams(num_tokens=len(leaf.key))) tree.evict(EvictParams(num_tokens=len(leaf.key)))
self.assertTrue(leaf.evicted) self.assertTrue(leaf.evicted)
result = tree.match_prefix(MatchPrefixParams(key=RadixKey(tokens))) result = tree.match_prefix(MatchPrefixParams(key=RadixKey(array("q", tokens))))
self.assertIs(result.best_match_node, leaf) self.assertIs(result.best_match_node, leaf)
self.assertIs(result.last_device_node, parent) self.assertIs(result.last_device_node, parent)
@@ -1782,7 +1791,7 @@ class UnifiedRadixCacheSuite:
tree.evict(EvictParams(num_tokens=len(leaf.key))) tree.evict(EvictParams(num_tokens=len(leaf.key)))
self.assertTrue(leaf.evicted) self.assertTrue(leaf.evicted)
result = tree.match_prefix(MatchPrefixParams(key=RadixKey(tokens))) result = tree.match_prefix(MatchPrefixParams(key=RadixKey(array("q", tokens))))
self.assertIs(result.best_match_node, leaf) self.assertIs(result.best_match_node, leaf)
self.assertIs(result.last_device_node, parent) self.assertIs(result.last_device_node, parent)
@@ -1797,12 +1806,14 @@ class UnifiedRadixCacheSuite:
tokens = self._make_seq(1, chunk_size + 1) tokens = self._make_seq(1, chunk_size + 1)
self._insert(tree, allocator, req_to_token_pool, tokens) self._insert(tree, allocator, req_to_token_pool, tokens)
leaf = tree.match_prefix( leaf = tree.match_prefix(
MatchPrefixParams(key=RadixKey(tokens)) MatchPrefixParams(key=RadixKey(array("q", tokens)))
).last_device_node ).last_device_node
mamba_cd = leaf.component_data[ComponentType.MAMBA] mamba_cd = leaf.component_data[ComponentType.MAMBA]
mamba_cd.value = None mamba_cd.value = None
no_hicache = tree.match_prefix(MatchPrefixParams(key=RadixKey(tokens))) no_hicache = tree.match_prefix(
MatchPrefixParams(key=RadixKey(array("q", tokens)))
)
self.assertIs(no_hicache.best_match_node, tree.root_node) self.assertIs(no_hicache.best_match_node, tree.root_node)
self.assertIs(no_hicache.last_device_node, tree.root_node) self.assertIs(no_hicache.last_device_node, tree.root_node)
self.assertEqual(no_hicache.mamba_branching_seqlen, chunk_size) self.assertEqual(no_hicache.mamba_branching_seqlen, chunk_size)
@@ -1810,11 +1821,13 @@ class UnifiedRadixCacheSuite:
tree_h, allocator_h, req_to_token_pool_h = self._build_hicache_fixture() tree_h, allocator_h, req_to_token_pool_h = self._build_hicache_fixture()
self._insert(tree_h, allocator_h, req_to_token_pool_h, tokens) self._insert(tree_h, allocator_h, req_to_token_pool_h, tokens)
leaf_h = tree_h.match_prefix( leaf_h = tree_h.match_prefix(
MatchPrefixParams(key=RadixKey(tokens)) MatchPrefixParams(key=RadixKey(array("q", tokens)))
).last_device_node ).last_device_node
self._backup_node(tree_h, leaf_h) self._backup_node(tree_h, leaf_h)
tree_h.evict(EvictParams(num_tokens=len(tokens))) tree_h.evict(EvictParams(num_tokens=len(tokens)))
with_hicache = tree_h.match_prefix(MatchPrefixParams(key=RadixKey(tokens))) with_hicache = tree_h.match_prefix(
MatchPrefixParams(key=RadixKey(array("q", tokens)))
)
self.assertIs(with_hicache.best_match_node, leaf_h) self.assertIs(with_hicache.best_match_node, leaf_h)
self.assertIs(with_hicache.last_device_node, tree_h.root_node) self.assertIs(with_hicache.last_device_node, tree_h.root_node)
self.assertIsNone(with_hicache.mamba_branching_seqlen) self.assertIsNone(with_hicache.mamba_branching_seqlen)
@@ -1834,7 +1847,9 @@ class UnifiedRadixCacheSuite:
self.assertTrue(leaf.evicted) self.assertTrue(leaf.evicted)
req = self._make_req(req_to_token_pool) req = self._make_req(req_to_token_pool)
match = tree.match_prefix(MatchPrefixParams(key=RadixKey(tokens), req=req)) match = tree.match_prefix(
MatchPrefixParams(key=RadixKey(array("q", tokens)), req=req)
)
req.prefix_indices = match.device_indices req.prefix_indices = match.device_indices
req.last_node = match.last_device_node req.last_node = match.last_device_node
req.best_match_node = match.best_match_node req.best_match_node = match.best_match_node
@@ -1876,7 +1891,9 @@ class UnifiedRadixCacheSuite:
self._set_aux_host_tombstone(tree, leaf, aux) self._set_aux_host_tombstone(tree, leaf, aux)
req = self._make_req(req_to_token_pool) req = self._make_req(req_to_token_pool)
match = tree.match_prefix(MatchPrefixParams(key=RadixKey(tokens), req=req)) match = tree.match_prefix(
MatchPrefixParams(key=RadixKey(array("q", tokens)), req=req)
)
req.prefix_indices = match.device_indices req.prefix_indices = match.device_indices
req.last_node = match.last_device_node req.last_node = match.last_device_node
req.best_match_node = match.best_match_node req.best_match_node = match.best_match_node
@@ -1915,7 +1932,9 @@ class UnifiedRadixCacheSuite:
tree.evict(EvictParams(num_tokens=len(leaf.key))) tree.evict(EvictParams(num_tokens=len(leaf.key)))
req = self._make_req(req_to_token_pool) req = self._make_req(req_to_token_pool)
match = tree.match_prefix(MatchPrefixParams(key=RadixKey(tokens), req=req)) match = tree.match_prefix(
MatchPrefixParams(key=RadixKey(array("q", tokens)), req=req)
)
req.prefix_indices = match.device_indices req.prefix_indices = match.device_indices
req.last_node = match.last_device_node req.last_node = match.last_device_node
req.best_match_node = match.best_match_node req.best_match_node = match.best_match_node
@@ -1986,7 +2005,7 @@ class UnifiedRadixCacheSuite:
tree, allocator, req_to_token_pool = build_fixture(self.cfg) tree, allocator, req_to_token_pool = build_fixture(self.cfg)
seq = self._make_seq(1, 2) seq = self._make_seq(1, 2)
self._insert(tree, allocator, req_to_token_pool, seq) self._insert(tree, allocator, req_to_token_pool, seq)
m = tree.match_prefix(MatchPrefixParams(key=RadixKey(seq))) m = tree.match_prefix(MatchPrefixParams(key=RadixKey(array("q", seq))))
node = m.last_device_node node = m.last_device_node
self._simulate_backup(tree, node) self._simulate_backup(tree, node)
@@ -2077,7 +2096,9 @@ class UnifiedRadixCacheSuite:
) )
result = swa_comp.finalize_match_result( result = swa_comp.finalize_match_result(
result=result, result=result,
params=MatchPrefixParams(key=RadixKey(self._make_seq(1, 1))), params=MatchPrefixParams(
key=RadixKey(array("q", self._make_seq(1, 1)))
),
value_chunks=[], value_chunks=[],
best_value_len=0, best_value_len=0,
) )
@@ -2214,7 +2235,7 @@ class UnifiedRadixCacheSuite:
def test_hicache_swa_match_prefix_picks_best_match_node_above_last_host(self): def test_hicache_swa_match_prefix_picks_best_match_node_above_last_host(self):
tree, _, n, y, x, tokens = self._swa_anchor_setup() tree, _, n, y, x, tokens = self._swa_anchor_setup()
result = tree.match_prefix(MatchPrefixParams(key=RadixKey(tokens))) result = tree.match_prefix(MatchPrefixParams(key=RadixKey(array("q", tokens))))
self.assertIs(result.best_match_node, x) self.assertIs(result.best_match_node, x)
self.assertIs(result.last_device_node, n.parent) self.assertIs(result.last_device_node, n.parent)
self.assertIs(result.last_host_node, y) self.assertIs(result.last_host_node, y)
@@ -2247,7 +2268,7 @@ class UnifiedRadixCacheSuite:
) )
result = swa_comp.finalize_match_result( result = swa_comp.finalize_match_result(
result=base, result=base,
params=MatchPrefixParams(key=RadixKey(self._make_seq(1, 1))), params=MatchPrefixParams(key=RadixKey(array("q", self._make_seq(1, 1)))),
value_chunks=[], value_chunks=[],
best_value_len=0, best_value_len=0,
) )
@@ -2365,7 +2386,7 @@ class UnifiedRadixCacheSuite:
tree, allocator, req_to_token_pool = build_fixture(self.cfg) tree, allocator, req_to_token_pool = build_fixture(self.cfg)
seq = self._make_seq(1, 2) seq = self._make_seq(1, 2)
self._insert(tree, allocator, req_to_token_pool, seq) self._insert(tree, allocator, req_to_token_pool, seq)
m = tree.match_prefix(MatchPrefixParams(key=RadixKey(seq))) m = tree.match_prefix(MatchPrefixParams(key=RadixKey(array("q", seq))))
node = m.last_device_node node = m.last_device_node
cd = node.component_data[ComponentType.MAMBA] cd = node.component_data[ComponentType.MAMBA]
old_mamba = cd.value old_mamba = cd.value
@@ -2414,7 +2435,7 @@ class UnifiedRadixCacheSuite:
tree.sanity_check() tree.sanity_check()
for i in range(3): for i in range(3):
m = tree.match_prefix(MatchPrefixParams(key=RadixKey(seqs[i]))) m = tree.match_prefix(MatchPrefixParams(key=RadixKey(array("q", seqs[i]))))
self._backup_node(tree, m.last_device_node) self._backup_node(tree, m.last_device_node)
# Evict to free some tokens # Evict to free some tokens
@@ -2443,7 +2464,7 @@ class UnifiedRadixCacheSuite:
self._insert(tree, allocator, req_to_token_pool, base) self._insert(tree, allocator, req_to_token_pool, base)
self._insert(tree, allocator, req_to_token_pool, leaf_seq) self._insert(tree, allocator, req_to_token_pool, leaf_seq)
m = tree.match_prefix(MatchPrefixParams(key=RadixKey(leaf_seq))) m = tree.match_prefix(MatchPrefixParams(key=RadixKey(array("q", leaf_seq))))
leaf = m.last_device_node leaf = m.last_device_node
parent = leaf.parent parent = leaf.parent
self.assertIsNot(parent, tree.root_node) self.assertIsNot(parent, tree.root_node)
test/registered/unit/utils/test_common.py
+50
View File
@@ -0,0 +1,50 @@
import unittest
from array import array
import torch
from sglang.srt.utils.common import flatten_arrays_to_int64_tensor
from sglang.test.ci.ci_register import register_cuda_ci
from sglang.test.test_utils import CustomTestCase
register_cuda_ci(est_time=5, stage="base-b", runner_config="1-gpu-small")
@unittest.skipUnless(torch.cuda.is_available(), "requires CUDA")
class TestFlattenArraysToInt64Tensor(CustomTestCase):
"""`flatten_arrays_to_int64_tensor` is invoked by `prepare_for_extend`
to build the per-batch input_ids tensor (pinned, async H2D) from a
list of array.array('q') per-req fill_ids slices. Tests the full
matrix of (device, pin) the production code paths through.
"""
DEVICES = ("cpu", "cuda")
PIN_OPTIONS = (False, True)
def _check(self, parts: list, expected: list[int]) -> None:
for device in self.DEVICES:
for pin in self.PIN_OPTIONS:
with self.subTest(device=device, pin=pin):
out = flatten_arrays_to_int64_tensor(parts, device, pin)
if device == "cuda":
torch.cuda.synchronize()
self.assertEqual(out.dtype, torch.int64)
self.assertEqual(out.device.type, device)
self.assertEqual(out.shape, (len(expected),))
self.assertEqual(out.cpu().tolist(), expected)
def test_single_part(self):
parts = [array("q", [1, 2, 3, 4, 5])]
self._check(parts, [1, 2, 3, 4, 5])
def test_multiple_parts(self):
parts = [
array("q", [10, 20, 30]),
array("q", [100, 200]),
array("q", [1000]),
]
self._check(parts, [10, 20, 30, 100, 200, 1000])
if __name__ == "__main__":
unittest.main()