623 lines
21 KiB
Python
623 lines
21 KiB
Python
"""Inference scheduler for continuous batching."""
|
|
|
|
import threading
|
|
import time
|
|
import uuid
|
|
from typing import Any, Callable, Dict, List, Optional, Tuple
|
|
|
|
import torch
|
|
from torch import Tensor
|
|
|
|
from astrai.model.automodel import AutoModel
|
|
from astrai.tokenize import AutoTokenizer
|
|
|
|
|
|
class RadixNode:
|
|
"""Radix tree node for prefix cache."""
|
|
|
|
def __init__(self):
|
|
self.children: Dict[int, "RadixNode"] = {} # token_id -> child node
|
|
self.hash: Optional[int] = None # 64-bit hash of the prefix
|
|
self.slot: int = -1 # KV Cache slot, valid only for leaf nodes
|
|
self.ref_count: int = 0 # number of tasks referencing this prefix
|
|
self.last_access: float = 0.0 # timestamp for LRU
|
|
self.token_sequence: list = [] # full token sequence from root to this node
|
|
|
|
|
|
class PrefixCacheManager:
|
|
"""Prefix cache manager using Radix tree with LRU eviction."""
|
|
|
|
def __init__(self, max_capacity: int = 1000, base: int = 131, mod: int = 10**9 + 7):
|
|
self.root = RadixNode()
|
|
self.base = base
|
|
self.mod = mod
|
|
self.max_capacity = max_capacity
|
|
self.lru: List[Tuple[float, RadixNode]] = [] # (timestamp, node) for LRU
|
|
|
|
def insert(self, token_ids: Tuple[int, ...], slot: int) -> None:
|
|
"""Insert a prefix, increase ref_count if already exists, otherwise create new node."""
|
|
node = self.root
|
|
path = []
|
|
h = 0
|
|
for i, token_id in enumerate(token_ids):
|
|
if token_id not in node.children:
|
|
node.children[token_id] = RadixNode()
|
|
node = node.children[token_id]
|
|
h = (h * self.base + token_id) % self.mod
|
|
node.hash = h
|
|
path.append(token_id)
|
|
node.token_sequence = list(
|
|
path
|
|
) # store full sequence for exact verification
|
|
|
|
# Leaf node: set slot and increase ref_count
|
|
if node.slot == -1:
|
|
node.slot = slot
|
|
node.ref_count += 1
|
|
node.last_access = time.time()
|
|
self._update_lru(node)
|
|
self._evict_if_needed()
|
|
|
|
def find_longest_prefix(self, token_ids: List[int]) -> Optional[Tuple[int, int]]:
|
|
"""Find longest matching prefix, return (prefix_len, slot).
|
|
|
|
During traversal, compute hash per token and compare with node hash.
|
|
If hash matches, perform full token sequence verification to avoid
|
|
hash collision errors.
|
|
"""
|
|
node = self.root
|
|
best_len = 0
|
|
best_slot = -1
|
|
h = 0
|
|
|
|
for i, token_id in enumerate(token_ids):
|
|
if token_id not in node.children:
|
|
break
|
|
node = node.children[token_id]
|
|
h = (h * self.base + token_id) % self.mod
|
|
if node.hash == h: # hash matches
|
|
# Exact verification: compare full token sequence
|
|
if node.token_sequence == token_ids[: i + 1]:
|
|
best_len = i + 1
|
|
best_slot = node.slot
|
|
node.last_access = time.time()
|
|
self._update_lru(node)
|
|
|
|
if best_len > 0:
|
|
return (best_len, best_slot)
|
|
return None
|
|
|
|
def release(self, token_ids: Tuple[int, ...]) -> None:
|
|
"""Release reference to a prefix, decrease ref_count. If zero, mark as evictable."""
|
|
node = self.root
|
|
for token_id in token_ids:
|
|
if token_id not in node.children:
|
|
return
|
|
node = node.children[token_id]
|
|
if node.ref_count > 0:
|
|
node.ref_count -= 1
|
|
if node.ref_count == 0:
|
|
node.slot = -1 # slot can be reused
|
|
|
|
def _update_lru(self, node: RadixNode) -> None:
|
|
"""Update LRU list, move node to most recently used position."""
|
|
self.lru = [(ts, n) for (ts, n) in self.lru if n is not node]
|
|
self.lru.append((node.last_access, node))
|
|
|
|
def _evict_if_needed(self) -> None:
|
|
"""If cache entries exceed capacity, evict least recently used leaf nodes (ref_count must be 0)."""
|
|
if len(self.lru) <= self.max_capacity:
|
|
return
|
|
# Sort by timestamp
|
|
self.lru.sort(key=lambda x: x[0])
|
|
for ts, node in self.lru:
|
|
if node.ref_count == 0:
|
|
# Remove leaf node from tree (need to recursively delete empty branches)
|
|
self._remove_node(node)
|
|
self.lru.remove((ts, node))
|
|
if len(self.lru) <= self.max_capacity:
|
|
break
|
|
|
|
def _remove_node(self, node: RadixNode) -> None:
|
|
"""Remove node from tree (simplified implementation)."""
|
|
# Clear the node's leaf properties
|
|
node.slot = -1
|
|
node.hash = None
|
|
node.token_sequence = []
|
|
|
|
|
|
class TaskStatus:
|
|
"""Task state for continuous batching."""
|
|
|
|
PENDING = "pending"
|
|
RUNNING = "running"
|
|
FINISHED = "finished"
|
|
ABORTED = "aborted"
|
|
|
|
|
|
class Task:
|
|
"""Individual task for continuous batching."""
|
|
|
|
def __init__(
|
|
self,
|
|
task_id: str,
|
|
prompt_ids: List[int],
|
|
max_tokens: int = 1024,
|
|
temperature: float = 1.0,
|
|
top_p: float = 1.0,
|
|
top_k: int = 50,
|
|
stream_callback: Optional[Callable[[str], None]] = None,
|
|
):
|
|
self.task_id = task_id
|
|
self.prompt_ids = prompt_ids
|
|
self.max_tokens = max_tokens
|
|
self.temperature = temperature
|
|
self.top_p = top_p
|
|
self.top_k = top_k
|
|
|
|
self.status = TaskStatus.PENDING
|
|
self.output_ids: List[int] = []
|
|
self.input_tokens: int = 0
|
|
self.output_tokens: int = 0
|
|
self.slot: int = -1
|
|
self.prefix_len: int = 0 # prefix cache matched length
|
|
self.arrival_time = time.time()
|
|
self.finish_time: Optional[float] = None
|
|
|
|
self.stream_callback = stream_callback
|
|
|
|
def is_finished(self, stop_ids: List[int]) -> bool:
|
|
"""Check if task is finished."""
|
|
return (
|
|
bool(self.output_ids and self.output_ids[-1] in stop_ids)
|
|
or self.output_tokens >= self.max_tokens
|
|
)
|
|
|
|
|
|
def apply_sampling_strategies(
|
|
logits: Tensor,
|
|
temperature: float,
|
|
top_k: int,
|
|
top_p: float,
|
|
filter_value: float = -float("inf"),
|
|
) -> Tensor:
|
|
"""Apply sampling strategies to the logits tensor."""
|
|
# Clone logits to avoid inplace updates on inference tensor
|
|
logits = logits.clone()
|
|
|
|
if temperature != 1.0:
|
|
logits = logits / temperature
|
|
|
|
if top_k > 0:
|
|
top_k = min(top_k, logits.size(-1))
|
|
indices_to_remove = logits < torch.topk(logits, top_k, dim=-1)[0][..., -1, None]
|
|
logits[indices_to_remove] = filter_value
|
|
|
|
if top_p < 1.0:
|
|
sorted_logits, sorted_indices = torch.sort(logits, descending=True, dim=-1)
|
|
cumulative_probs = torch.cumsum(torch.softmax(sorted_logits, dim=-1), dim=-1)
|
|
|
|
sorted_indices_to_remove = cumulative_probs > top_p
|
|
sorted_indices_to_remove[..., 1:] = sorted_indices_to_remove[..., :-1].clone()
|
|
sorted_indices_to_remove[..., 0] = 0
|
|
|
|
indices_to_remove = torch.zeros_like(logits, dtype=torch.bool)
|
|
indices_to_remove.scatter_(
|
|
dim=1, index=sorted_indices, src=sorted_indices_to_remove
|
|
)
|
|
|
|
logits[indices_to_remove] = filter_value
|
|
|
|
return logits
|
|
|
|
|
|
class InferenceScheduler:
|
|
"""Inference scheduler with continuous batching support."""
|
|
|
|
def __init__(
|
|
self,
|
|
model: AutoModel,
|
|
tokenizer: AutoTokenizer,
|
|
max_batch_size: int = 16,
|
|
max_seq_len: Optional[int] = None,
|
|
max_prefix_len: int = 512,
|
|
cache_capacity: int = 1000,
|
|
device: str = "cuda",
|
|
dtype: torch.dtype = torch.bfloat16,
|
|
):
|
|
config = model.config
|
|
|
|
self.model = model
|
|
self.tokenizer = tokenizer
|
|
self.max_batch_size = max_batch_size
|
|
self.max_seq_len = max_seq_len or config.max_len
|
|
self.max_prefix_len = max_prefix_len
|
|
self.device = device or next(model.parameters()).device
|
|
self.dtype = dtype or next(model.parameters()).dtype
|
|
|
|
# Initialize prefix cache
|
|
self.prefix_cache = PrefixCacheManager(max_capacity=cache_capacity)
|
|
|
|
num_kv_heads = config.n_kv_heads
|
|
head_dim = config.dim // config.n_heads
|
|
n_layers = config.n_layers
|
|
|
|
k_cache = torch.empty(
|
|
(
|
|
max_batch_size,
|
|
self.max_seq_len,
|
|
n_layers,
|
|
num_kv_heads,
|
|
head_dim,
|
|
),
|
|
device=self.device,
|
|
dtype=self.dtype,
|
|
)
|
|
v_cache = torch.empty(
|
|
(
|
|
max_batch_size,
|
|
self.max_seq_len,
|
|
n_layers,
|
|
num_kv_heads,
|
|
head_dim,
|
|
),
|
|
device=self.device,
|
|
dtype=self.dtype,
|
|
)
|
|
self.kv_cache = (k_cache, v_cache)
|
|
self.seq_mask = torch.ones(
|
|
(max_batch_size, self.max_seq_len), device=self.device, dtype=torch.bool
|
|
)
|
|
|
|
self.waiting_queue: List[Task] = []
|
|
self.active_tasks: List[Task] = []
|
|
|
|
self._running = False
|
|
self._task_event = threading.Event()
|
|
self._lock = threading.Lock()
|
|
|
|
self._total_tasks = 0
|
|
self._total_tokens = 0
|
|
|
|
def add_task(
|
|
self,
|
|
prompt: str,
|
|
max_tokens: int = 1024,
|
|
temperature: float = 1.0,
|
|
top_p: float = 1.0,
|
|
top_k: int = 50,
|
|
stream_callback: Optional[Callable[[str], None]] = None,
|
|
) -> str:
|
|
"""Add a new task to the waiting queue."""
|
|
task_id = f"task_{int(time.time())}_{uuid.uuid4().hex[:8]}"
|
|
prompt_ids = self.tokenizer.encode(prompt)
|
|
|
|
# Truncate if exceeds max_prefix_len
|
|
if len(prompt_ids) > self.max_prefix_len:
|
|
prompt_ids = prompt_ids[: self.max_prefix_len]
|
|
|
|
task = Task(
|
|
task_id=task_id,
|
|
prompt_ids=prompt_ids,
|
|
max_tokens=max_tokens,
|
|
temperature=temperature,
|
|
top_p=top_p,
|
|
top_k=top_k,
|
|
stream_callback=stream_callback,
|
|
)
|
|
|
|
# Find longest matching prefix from cache
|
|
match = self.prefix_cache.find_longest_prefix(prompt_ids)
|
|
if match:
|
|
prefix_len, slot = match
|
|
task.prefix_len = prefix_len
|
|
task.slot = slot
|
|
else:
|
|
task.prefix_len = 0
|
|
task.slot = -1
|
|
|
|
with self._lock:
|
|
self.waiting_queue.append(task)
|
|
self._total_tasks += 1
|
|
|
|
self._task_event.set()
|
|
return task_id
|
|
|
|
def remove_task(self, task_id: str) -> None:
|
|
"""Remove a task from the scheduler."""
|
|
with self._lock:
|
|
self.waiting_queue = [t for t in self.waiting_queue if t.task_id != task_id]
|
|
self.active_tasks = [t for t in self.active_tasks if t.task_id != task_id]
|
|
|
|
def _remove_finished_tasks(self) -> None:
|
|
"""Remove finished tasks from active batch."""
|
|
finished = []
|
|
for task in self.active_tasks:
|
|
if task.is_finished(self.tokenizer.stop_ids):
|
|
task.status = TaskStatus.FINISHED
|
|
task.finish_time = time.time()
|
|
finished.append(task)
|
|
self._total_tokens += task.output_tokens
|
|
|
|
for task in finished:
|
|
slot = task.slot
|
|
if slot >= 0 and slot < len(self.active_tasks):
|
|
self.seq_mask[slot, :] = False
|
|
|
|
# Release prefix cache reference
|
|
if task.prefix_len > 0:
|
|
self.prefix_cache.release(tuple(task.prompt_ids[: task.prefix_len]))
|
|
|
|
task.slot = -1
|
|
|
|
self.active_tasks = [
|
|
t for t in self.active_tasks if t.status != TaskStatus.FINISHED
|
|
]
|
|
|
|
def _refill_active_batch(self) -> None:
|
|
"""Refill active batch with waiting tasks."""
|
|
available_slots = self.max_batch_size - len(self.active_tasks)
|
|
if available_slots <= 0:
|
|
return
|
|
|
|
with self._lock:
|
|
to_add = [
|
|
self.waiting_queue.pop(0)
|
|
for _ in range(min(available_slots, len(self.waiting_queue)))
|
|
]
|
|
for task in to_add:
|
|
task.slot = self._allocate_slot()
|
|
task.status = TaskStatus.RUNNING
|
|
self.active_tasks.append(task)
|
|
|
|
def _allocate_slot(self) -> int:
|
|
"""Allocate an available slot for a task."""
|
|
for i in range(self.max_batch_size):
|
|
if not any(t.slot == i for t in self.active_tasks):
|
|
return i
|
|
return -1
|
|
|
|
def _execute_prefill(self, tasks: List[Task]) -> None:
|
|
"""Execute Prefill phase with incremental prefill support."""
|
|
if not tasks:
|
|
return
|
|
|
|
# Group tasks by prefix cache status
|
|
fully_cached, partial, full = [], [], []
|
|
for task in tasks:
|
|
total_len, prefix_len = len(task.prompt_ids), task.prefix_len
|
|
if prefix_len == total_len:
|
|
fully_cached.append(task)
|
|
elif prefix_len > 0:
|
|
partial.append(task)
|
|
else:
|
|
full.append(task)
|
|
|
|
# Handle fully cached tasks
|
|
for t in fully_cached:
|
|
t.input_tokens, t.output_tokens = len(t.prompt_ids), 0
|
|
if t.slot >= 0:
|
|
self.seq_mask[t.slot, : t.input_tokens] = True
|
|
|
|
if full:
|
|
self._execute_full_prefill(full)
|
|
if partial:
|
|
self._execute_partial_prefill(partial)
|
|
|
|
def _execute_full_prefill(self, tasks: List[Task]) -> None:
|
|
"""Execute full prefill for tasks without prefix cache."""
|
|
if not tasks:
|
|
return
|
|
|
|
tasks = sorted(tasks, key=lambda t: t.slot)
|
|
|
|
prompt_lens = [len(task.prompt_ids) for task in tasks]
|
|
max_len = max(prompt_lens)
|
|
|
|
input_ids = torch.zeros(
|
|
len(tasks), max_len, dtype=torch.long, device=self.device
|
|
)
|
|
for i, task in enumerate(tasks):
|
|
if len(task.prompt_ids) > 0:
|
|
input_ids[i, : len(task.prompt_ids)] = torch.tensor(
|
|
task.prompt_ids, device=self.device
|
|
)
|
|
|
|
if self.tokenizer.pad_id is not None:
|
|
input_mask = torch.ne(input_ids, self.tokenizer.pad_id)
|
|
else:
|
|
input_mask = torch.ones(
|
|
input_ids.shape, dtype=torch.bool, device=self.device
|
|
)
|
|
|
|
with torch.inference_mode():
|
|
self.model(
|
|
input_ids,
|
|
input_mask=input_mask,
|
|
start_pos=0,
|
|
persistent_key_values=self.kv_cache,
|
|
)
|
|
|
|
for i, task in enumerate(tasks):
|
|
task.input_tokens = prompt_lens[i]
|
|
task.output_tokens = 0
|
|
# Insert new prefix into cache
|
|
self.prefix_cache.insert(tuple(task.prompt_ids), task.slot)
|
|
|
|
for task in tasks:
|
|
if task.slot >= 0:
|
|
self.seq_mask[task.slot, : task.input_tokens] = True
|
|
|
|
def _execute_partial_prefill(self, tasks: List[Task]) -> None:
|
|
"""Execute incremental prefill for tasks with partial prefix cache match."""
|
|
for task in tasks:
|
|
total_len = len(task.prompt_ids)
|
|
prefix_len = task.prefix_len
|
|
|
|
if prefix_len >= total_len:
|
|
task.input_tokens = total_len
|
|
task.output_tokens = 0
|
|
continue
|
|
|
|
# Get new tokens that need prefill
|
|
new_ids = task.prompt_ids[prefix_len:]
|
|
new_len = len(new_ids)
|
|
|
|
if new_len == 0:
|
|
task.input_tokens = total_len
|
|
task.output_tokens = 0
|
|
continue
|
|
|
|
# Build input for incremental prefill
|
|
input_ids = torch.tensor([new_ids], dtype=torch.long, device=self.device)
|
|
|
|
# Input mask should cover from position 0 to prefix_len + new_len
|
|
# The prefix part uses cached KV, new part needs computation
|
|
input_mask = torch.ones(
|
|
(1, prefix_len + new_len), dtype=torch.bool, device=self.device
|
|
)
|
|
|
|
with torch.inference_mode():
|
|
self.model(
|
|
input_ids,
|
|
input_mask=input_mask,
|
|
start_pos=prefix_len,
|
|
persistent_key_values=self.kv_cache,
|
|
)
|
|
|
|
task.input_tokens = total_len
|
|
task.output_tokens = 0
|
|
|
|
# Insert full prefix into cache (ref_count already increased in add_task)
|
|
self.prefix_cache.insert(tuple(task.prompt_ids), task.slot)
|
|
|
|
if task.slot >= 0:
|
|
self.seq_mask[task.slot, : task.input_tokens] = True
|
|
|
|
def _execute_decode(self, tasks: List[Task], start_pos: int) -> None:
|
|
"""Execute Decode phase."""
|
|
if not tasks:
|
|
return
|
|
|
|
tasks = sorted(tasks, key=lambda t: t.slot)
|
|
|
|
input_ids = torch.zeros(len(tasks), dtype=torch.long, device=self.device)
|
|
for i, task in enumerate(tasks):
|
|
if task.output_ids:
|
|
input_ids[i] = task.output_ids[-1]
|
|
else:
|
|
input_ids[i] = task.prompt_ids[-1]
|
|
|
|
input_tensor = input_ids.unsqueeze(1)
|
|
active_mask = torch.ones((len(tasks), 1), dtype=torch.bool, device=self.device)
|
|
|
|
with torch.inference_mode():
|
|
outputs = self.model(
|
|
input_tensor,
|
|
input_mask=active_mask,
|
|
persistent_key_values=self.kv_cache,
|
|
start_pos=start_pos,
|
|
)
|
|
logits = outputs["logits"][:, -1, :]
|
|
|
|
next_token_ids = []
|
|
for i, task in enumerate(tasks):
|
|
logit = logits[i : i + 1]
|
|
logit = apply_sampling_strategies(
|
|
logit,
|
|
task.temperature,
|
|
task.top_k,
|
|
task.top_p,
|
|
)
|
|
probs = torch.softmax(logit, dim=-1)
|
|
next_token = torch.multinomial(probs, num_samples=1)
|
|
next_token_ids.append(next_token.item())
|
|
|
|
for task, next_token in zip(tasks, next_token_ids):
|
|
task.output_ids.append(next_token)
|
|
task.output_tokens += 1
|
|
|
|
pos = task.input_tokens + task.output_tokens
|
|
if task.slot >= 0 and pos < self.max_seq_len:
|
|
self.seq_mask[task.slot, pos] = True
|
|
|
|
if task.stream_callback:
|
|
token_str = self.tokenizer.decode([next_token])
|
|
task.stream_callback(token_str)
|
|
|
|
for task in tasks:
|
|
if task.output_tokens >= task.max_tokens or (
|
|
task.output_ids and task.output_ids[-1] in self.tokenizer.stop_ids
|
|
):
|
|
if task.stream_callback:
|
|
task.stream_callback("[DONE]")
|
|
|
|
def _run_generation_loop(self) -> None:
|
|
"""Main generation loop."""
|
|
while self._running:
|
|
self._remove_finished_tasks()
|
|
self._refill_active_batch()
|
|
|
|
if not self.active_tasks:
|
|
self._task_event.wait(timeout=0.01)
|
|
self._task_event.clear()
|
|
continue
|
|
|
|
new_tasks = [t for t in self.active_tasks if t.output_tokens == 0]
|
|
decode_tasks = [t for t in self.active_tasks if t.output_tokens > 0]
|
|
|
|
if decode_tasks:
|
|
start_pos = max(t.input_tokens + t.output_tokens for t in decode_tasks)
|
|
else:
|
|
start_pos = 0
|
|
|
|
if new_tasks:
|
|
self._execute_prefill(new_tasks)
|
|
decode_tasks = new_tasks
|
|
start_pos = max(t.input_tokens for t in decode_tasks)
|
|
|
|
if decode_tasks:
|
|
self._execute_decode(decode_tasks, start_pos)
|
|
|
|
if not self.active_tasks and not self.waiting_queue:
|
|
self._task_event.wait(timeout=0.05)
|
|
self._task_event.clear()
|
|
|
|
def start(self) -> None:
|
|
"""Start the generation loop."""
|
|
if not self._running:
|
|
self._running = True
|
|
self._loop_thread = threading.Thread(target=self._run_generation_loop)
|
|
self._loop_thread.daemon = True
|
|
self._loop_thread.start()
|
|
|
|
def stop(self) -> None:
|
|
"""Stop the generation loop."""
|
|
self._running = False
|
|
if hasattr(self, "_loop_thread"):
|
|
self._loop_thread.join(timeout=1.0)
|
|
|
|
# Clear KV cache to free GPU memory
|
|
if self.kv_cache is not None:
|
|
k_cache, v_cache = self.kv_cache
|
|
if k_cache is not None:
|
|
k_cache.detach()
|
|
if v_cache is not None:
|
|
v_cache.detach()
|
|
|
|
# Clear seq mask
|
|
self.seq_mask.detach()
|
|
|
|
# Clear task lists
|
|
self.waiting_queue.clear()
|
|
self.active_tasks.clear()
|
|
|
|
def get_stats(self) -> Dict[str, Any]:
|
|
"""Get scheduler statistics."""
|
|
return {
|
|
"total_tasks": self._total_tasks,
|
|
"total_tokens": self._total_tokens,
|
|
"active_tasks": len(self.active_tasks),
|
|
"waiting_queue": len(self.waiting_queue),
|
|
}
|