Lightweight extract allocation logic from mem_cache/common.py to more clearly show nearly parallel variants (#29431)
This commit is contained in:
fzyzcjy
@@ -66,7 +66,7 @@ patches:
|
||||
|
||||
_PR_REVERT_YAML_26972 = """
|
||||
patches:
|
||||
- target: sglang.srt.mem_cache.common.get_req_to_token_extra_context_len
|
||||
- target: sglang.srt.mem_cache.allocation_sizing.get_req_to_token_extra_context_len
|
||||
edits:
|
||||
- match: |
|
||||
if (
|
||||
|
||||
@@ -33,8 +33,8 @@ from sglang.srt.hardware_backend.npu.allocator_npu import NPUPagedTokenToKVPoolA
|
||||
from sglang.srt.hardware_backend.npu.dsv4.dsv4_common_hooks import (
|
||||
maybe_write_dsv4_extend,
|
||||
)
|
||||
from sglang.srt.mem_cache.allocation import alloc_paged_token_slots_extend
|
||||
from sglang.srt.mem_cache.allocator.swa import SWATokenToKVPoolAllocator
|
||||
from sglang.srt.mem_cache.common import alloc_paged_token_slots_extend
|
||||
from sglang.srt.model_executor.forward_batch_info import DSV4OutCacheLoc, DSV4StateLens
|
||||
|
||||
if TYPE_CHECKING:
|
||||
|
||||
@@ -77,6 +77,11 @@ from sglang.srt.managers.embed_types import PositionalEmbeds
|
||||
from sglang.srt.managers.scheduler_components.new_token_ratio_tracker import (
|
||||
NewTokenRatioTracker,
|
||||
)
|
||||
from sglang.srt.mem_cache.allocation import (
|
||||
alloc_for_decode,
|
||||
alloc_for_extend,
|
||||
)
|
||||
from sglang.srt.mem_cache.allocation_sizing import get_alloc_reserve_per_decode
|
||||
from sglang.srt.mem_cache.allocator import BaseTokenToKVPoolAllocator
|
||||
from sglang.srt.mem_cache.base_prefix_cache import (
|
||||
BasePrefixCache,
|
||||
@@ -85,11 +90,8 @@ from sglang.srt.mem_cache.base_prefix_cache import (
|
||||
zero_match_result,
|
||||
)
|
||||
from sglang.srt.mem_cache.common import (
|
||||
alloc_for_decode,
|
||||
alloc_for_extend,
|
||||
evict_from_tree_cache,
|
||||
free_swa_out_of_window_slots,
|
||||
get_alloc_reserve_per_decode,
|
||||
release_kv_cache,
|
||||
)
|
||||
from sglang.srt.mem_cache.memory_pool import ReqToTokenPool
|
||||
|
||||
@@ -0,0 +1,469 @@
|
||||
from __future__ import annotations
|
||||
|
||||
import logging
|
||||
from typing import TYPE_CHECKING, Optional
|
||||
|
||||
import torch
|
||||
|
||||
from sglang.kernels.ops.memory.common import (
|
||||
get_last_loc_triton,
|
||||
get_last_loc_triton_safe,
|
||||
write_req_to_token_pool_triton,
|
||||
)
|
||||
from sglang.srt.hardware_backend.npu.dsv4.dsv4_common_hooks import (
|
||||
maybe_write_dsv4_decode,
|
||||
maybe_write_dsv4_extend,
|
||||
)
|
||||
from sglang.srt.mem_cache.base_prefix_cache import BasePrefixCache, EvictParams
|
||||
from sglang.srt.mem_cache.common import (
|
||||
MAMBA_STATE_PER_REQ_NO_CACHE,
|
||||
MAMBA_STATE_PER_REQ_PREFIX_CACHE,
|
||||
MAMBA_STATE_PER_REQ_PREFIX_CACHE_LAZY,
|
||||
available_and_evictable_str,
|
||||
evict_from_tree_cache,
|
||||
)
|
||||
from sglang.srt.mem_cache.memory_pool import HybridReqToTokenPool, ReqToTokenPool
|
||||
from sglang.srt.runtime_context import get_server_args
|
||||
from sglang.srt.utils import is_cuda, is_hip, is_npu, support_triton
|
||||
from sglang.srt.utils.common import is_pin_memory_available
|
||||
|
||||
_is_hip = is_hip()
|
||||
_is_npu = is_npu()
|
||||
_is_cuda = is_cuda()
|
||||
|
||||
if TYPE_CHECKING:
|
||||
from sglang.srt.managers.schedule_batch import Req, ScheduleBatch
|
||||
from sglang.srt.model_executor.forward_batch_info import DSV4StateLens
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
|
||||
def write_cache_indices(
|
||||
out_cache_loc: torch.Tensor,
|
||||
req_pool_indices_tensor: torch.Tensor,
|
||||
req_pool_indices_cpu: torch.Tensor,
|
||||
prefix_lens_tensor: torch.Tensor,
|
||||
prefix_lens_cpu: torch.Tensor,
|
||||
seq_lens_tensor: torch.Tensor,
|
||||
seq_lens_cpu: torch.Tensor,
|
||||
extend_lens_tensor: torch.Tensor,
|
||||
extend_lens_cpu: torch.Tensor,
|
||||
prefix_tensors: list[torch.Tensor],
|
||||
req_to_token_pool: ReqToTokenPool,
|
||||
):
|
||||
if support_triton(get_server_args().attention_backend):
|
||||
prefix_pointers = torch.tensor(
|
||||
[t.data_ptr() for t in prefix_tensors],
|
||||
dtype=torch.uint64,
|
||||
pin_memory=is_pin_memory_available(req_to_token_pool.device),
|
||||
).to(req_to_token_pool.device, non_blocking=True)
|
||||
# TODO: some tensors can be reused for ForwardBatchInfo (e.g., extend_lens, cumsum_start)
|
||||
write_req_to_token_pool_triton[(req_pool_indices_tensor.shape[0],)](
|
||||
req_to_token_pool.req_to_token,
|
||||
req_pool_indices_tensor,
|
||||
prefix_pointers,
|
||||
prefix_lens_tensor,
|
||||
seq_lens_tensor,
|
||||
extend_lens_tensor,
|
||||
out_cache_loc,
|
||||
req_to_token_pool.req_to_token.shape[1],
|
||||
)
|
||||
else:
|
||||
pt = 0
|
||||
for i in range(req_pool_indices_cpu.shape[0]):
|
||||
req_idx = req_pool_indices_cpu[i].item()
|
||||
prefix_len = prefix_lens_cpu[i].item()
|
||||
seq_len = seq_lens_cpu[i].item()
|
||||
extend_len = extend_lens_cpu[i].item()
|
||||
|
||||
req_to_token_pool.write(
|
||||
(req_idx, slice(0, prefix_len)),
|
||||
prefix_tensors[i],
|
||||
)
|
||||
req_to_token_pool.write(
|
||||
(req_idx, slice(prefix_len, seq_len)),
|
||||
out_cache_loc[pt : pt + extend_len],
|
||||
)
|
||||
pt += extend_len
|
||||
|
||||
|
||||
def get_last_loc(
|
||||
req_to_token: torch.Tensor,
|
||||
req_pool_indices_tensor: torch.Tensor,
|
||||
prefix_lens_tensor: torch.Tensor,
|
||||
) -> torch.Tensor:
|
||||
attn_backend = get_server_args().attention_backend
|
||||
uses_triton_dispatch = attn_backend not in ("ascend", "torch_native")
|
||||
|
||||
if _is_hip and uses_triton_dispatch:
|
||||
# HIP-only: the legacy get_last_loc_triton kernel emits a
|
||||
# mixed-width int32->int64 store that Triton mis-compiles on HIP,
|
||||
# producing out-of-range last_loc values under EAGLE +
|
||||
# page_size>1 (e.g. with aiter unified attention or the triton
|
||||
# attention backend). The bug is in the Triton HIP codegen, not
|
||||
# in any particular attention backend, so route every HIP path
|
||||
# that would otherwise use get_last_loc_triton through the
|
||||
# int32-safe variant. Non-HIP hardware keeps the original
|
||||
# dispatcher below.
|
||||
return get_last_loc_triton_safe(
|
||||
req_to_token, req_pool_indices_tensor, prefix_lens_tensor
|
||||
)
|
||||
|
||||
if uses_triton_dispatch:
|
||||
impl = get_last_loc_triton
|
||||
else:
|
||||
impl = get_last_loc_torch
|
||||
|
||||
return impl(req_to_token, req_pool_indices_tensor, prefix_lens_tensor)
|
||||
|
||||
|
||||
def get_last_loc_torch(
|
||||
req_to_token: torch.Tensor,
|
||||
req_pool_indices_tensor: torch.Tensor,
|
||||
prefix_lens_tensor: torch.Tensor,
|
||||
) -> torch.Tensor:
|
||||
return torch.where(
|
||||
prefix_lens_tensor > 0,
|
||||
req_to_token[req_pool_indices_tensor, prefix_lens_tensor - 1],
|
||||
torch.full_like(prefix_lens_tensor, -1),
|
||||
)
|
||||
|
||||
|
||||
def alloc_token_slots(
|
||||
tree_cache: BasePrefixCache,
|
||||
num_tokens: int,
|
||||
backup_state: bool = False,
|
||||
):
|
||||
allocator = tree_cache.token_to_kv_pool_allocator
|
||||
evict_from_tree_cache(tree_cache, num_tokens)
|
||||
|
||||
state = None
|
||||
if backup_state:
|
||||
state = allocator.backup_state()
|
||||
|
||||
out_cache_loc = allocator.alloc(num_tokens)
|
||||
|
||||
if out_cache_loc is None:
|
||||
error_msg = (
|
||||
f"Out of memory. Try to lower your batch size.\n"
|
||||
f"Try to allocate {num_tokens} tokens.\n"
|
||||
f"{available_and_evictable_str(tree_cache)}"
|
||||
)
|
||||
logger.error(error_msg)
|
||||
if tree_cache is not None:
|
||||
tree_cache.pretty_print()
|
||||
raise RuntimeError(error_msg)
|
||||
|
||||
return (out_cache_loc, state) if backup_state else out_cache_loc
|
||||
|
||||
|
||||
def _compute_dsv4_state_lens(batch, *, is_decode: bool):
|
||||
"""Per-req c{4,128}_state pool alloc lens (``DSV4StateLens``) for this step.
|
||||
None on CUDA / non-V4 paths (allocator has no ``compute_dsv4_state_lens_*``).
|
||||
"""
|
||||
allocator = batch.token_to_kv_pool_allocator
|
||||
if not hasattr(allocator, "compute_dsv4_state_lens_extend"):
|
||||
return None
|
||||
if is_decode:
|
||||
return allocator.compute_dsv4_state_lens_decode(batch.reqs)
|
||||
return allocator.compute_dsv4_state_lens_extend(
|
||||
batch.reqs, batch.seq_lens_cpu.tolist()
|
||||
)
|
||||
|
||||
|
||||
def alloc_paged_token_slots_extend(
|
||||
tree_cache: BasePrefixCache,
|
||||
prefix_lens: torch.Tensor,
|
||||
prefix_lens_cpu: torch.Tensor,
|
||||
seq_lens: torch.Tensor,
|
||||
seq_lens_cpu: torch.Tensor,
|
||||
last_loc: torch.Tensor,
|
||||
extend_num_tokens: int,
|
||||
backup_state: bool = False,
|
||||
req_pool_indices: Optional[torch.Tensor] = None,
|
||||
dsv4_state_lens: Optional[DSV4StateLens] = None,
|
||||
batch=None,
|
||||
):
|
||||
# Over estimate the number of tokens: assume each request needs a new page.
|
||||
allocator = tree_cache.token_to_kv_pool_allocator
|
||||
num_tokens = extend_num_tokens + len(seq_lens_cpu) * allocator.page_size
|
||||
evict_from_tree_cache(tree_cache, num_tokens)
|
||||
|
||||
state = None
|
||||
if backup_state:
|
||||
state = allocator.backup_state()
|
||||
|
||||
is_dsv4 = req_pool_indices is not None and hasattr(allocator, "c4_attn_allocator")
|
||||
extra_alloc_kwargs = {}
|
||||
if is_dsv4:
|
||||
extra_alloc_kwargs["req_pool_indices"] = req_pool_indices
|
||||
# Per-call per-req tables for the c-pool / state last_loc lookup.
|
||||
if batch is not None:
|
||||
extra_alloc_kwargs["req_to_token_pool"] = batch.req_to_token_pool
|
||||
if dsv4_state_lens is not None:
|
||||
extra_alloc_kwargs["dsv4_state_lens"] = dsv4_state_lens
|
||||
|
||||
out = allocator.alloc_extend(
|
||||
prefix_lens,
|
||||
prefix_lens_cpu,
|
||||
seq_lens,
|
||||
seq_lens_cpu,
|
||||
last_loc,
|
||||
extend_num_tokens,
|
||||
**extra_alloc_kwargs,
|
||||
)
|
||||
|
||||
if is_dsv4:
|
||||
bundle = out
|
||||
out_cache_loc = None if bundle is None else bundle.out_full_loc
|
||||
if batch is not None:
|
||||
batch.out_cache_loc_dsv4 = bundle
|
||||
else:
|
||||
out_cache_loc = out
|
||||
|
||||
if out_cache_loc is None:
|
||||
error_msg = (
|
||||
f"Prefill out of memory. Try to lower your batch size.\n"
|
||||
f"Try to allocate {extend_num_tokens} tokens.\n"
|
||||
f"{available_and_evictable_str(tree_cache)}"
|
||||
)
|
||||
logger.error(error_msg)
|
||||
if tree_cache is not None:
|
||||
tree_cache.pretty_print()
|
||||
raise RuntimeError(error_msg)
|
||||
|
||||
return (out_cache_loc, state) if backup_state else out_cache_loc
|
||||
|
||||
|
||||
def alloc_req_slots(
|
||||
req_to_token_pool: ReqToTokenPool,
|
||||
reqs: list[Req],
|
||||
tree_cache: BasePrefixCache | None,
|
||||
) -> list[int]:
|
||||
"""Allocate request slots from the pool.
|
||||
|
||||
Fail-loud: raises ``RuntimeError`` if the pool can't satisfy the batch. An
|
||||
alloc failure here means the admission budget (``PrefillAdder``) was wrong
|
||||
and should surface rather than be masked.
|
||||
"""
|
||||
num_reqs = len(reqs)
|
||||
if isinstance(req_to_token_pool, HybridReqToTokenPool):
|
||||
# Byte-coordinated for the shared allocator (accounts for the peer full
|
||||
# sub-pool's bytes); plain slot free count for the non-shared one.
|
||||
mamba_available_size = (
|
||||
req_to_token_pool.mamba_allocator.schedulable_available_size()
|
||||
)
|
||||
# Eviction headroom factor: 3x (or lazy variant) for radix COW, 1x for chunk.
|
||||
if tree_cache.supports_mamba():
|
||||
factor = (
|
||||
MAMBA_STATE_PER_REQ_PREFIX_CACHE_LAZY
|
||||
if req_to_token_pool.enable_mamba_extra_buffer_lazy
|
||||
else MAMBA_STATE_PER_REQ_PREFIX_CACHE
|
||||
)
|
||||
else:
|
||||
factor = MAMBA_STATE_PER_REQ_NO_CACHE
|
||||
mamba_state_needed = num_reqs * factor
|
||||
if mamba_available_size < mamba_state_needed:
|
||||
if tree_cache is not None and tree_cache.supports_mamba():
|
||||
mamba_num = max(0, mamba_state_needed - mamba_available_size)
|
||||
tree_cache.evict(EvictParams(num_tokens=0, mamba_num=mamba_num))
|
||||
req_pool_indices = req_to_token_pool.alloc(reqs)
|
||||
if req_pool_indices is None:
|
||||
raise RuntimeError(
|
||||
"alloc_req_slots runs out of memory. "
|
||||
"Please set a smaller number for `--max-running-requests`. "
|
||||
f"{req_to_token_pool.available_size()=}, {num_reqs=}, "
|
||||
)
|
||||
return req_pool_indices
|
||||
|
||||
|
||||
def _alloc_page_size(batch: ScheduleBatch) -> int:
|
||||
# DCP swaps in an allocator whose page_size is server_args.page_size *
|
||||
# dcp_size, so it can be > 1 even when tree_cache.page_size is 1; branch on
|
||||
# the real allocator's page_size there. Elsewhere the two are equal.
|
||||
if (_is_hip or _is_cuda) and get_server_args().dcp_size > 1:
|
||||
return batch.tree_cache.token_to_kv_pool_allocator.page_size
|
||||
return batch.tree_cache.page_size
|
||||
|
||||
|
||||
def alloc_for_extend(
|
||||
batch: ScheduleBatch,
|
||||
) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
|
||||
"""
|
||||
Allocate KV cache for extend batch and write to req_to_token_pool.
|
||||
|
||||
Returns ``(out_cache_loc, req_pool_indices_device, req_pool_indices_cpu)``
|
||||
(the last is the host/CPU mirror). ``alloc_req_slots`` raises ``RuntimeError``
|
||||
if the pool can't satisfy the batch (fail-loud — see its docstring).
|
||||
"""
|
||||
# free out-of-window swa tokens
|
||||
batch.maybe_evict_swa()
|
||||
|
||||
prefix_tensors = [r.prefix_indices for r in batch.reqs]
|
||||
|
||||
# Create tensors for allocation
|
||||
prefix_lens_cpu = torch.tensor(batch.prefix_lens, dtype=torch.int64)
|
||||
extend_lens_cpu = torch.tensor(batch.extend_lens, dtype=torch.int64)
|
||||
prefix_lens_device = prefix_lens_cpu.to(batch.device, non_blocking=True)
|
||||
extend_lens_device = extend_lens_cpu.to(batch.device, non_blocking=True)
|
||||
|
||||
# Allocate req slots (raises RuntimeError if the pool is exhausted)
|
||||
req_pool_indices = alloc_req_slots(
|
||||
batch.req_to_token_pool, batch.reqs, batch.tree_cache
|
||||
)
|
||||
req_pool_indices_cpu = torch.tensor(req_pool_indices, dtype=torch.int64)
|
||||
req_pool_indices_device = req_pool_indices_cpu.to(batch.device, non_blocking=True)
|
||||
|
||||
# Allocate KV cache (throws exception on failure)
|
||||
if _alloc_page_size(batch) == 1:
|
||||
out_cache_loc = alloc_token_slots(batch.tree_cache, batch.extend_num_tokens)
|
||||
else:
|
||||
# Paged allocation - build last_loc
|
||||
last_loc = [
|
||||
(t[-1:] if len(t) > 0 else torch.tensor([-1], device=batch.device))
|
||||
for t in prefix_tensors
|
||||
]
|
||||
out_cache_loc = alloc_paged_token_slots_extend(
|
||||
tree_cache=batch.tree_cache,
|
||||
prefix_lens=prefix_lens_device,
|
||||
prefix_lens_cpu=prefix_lens_cpu,
|
||||
seq_lens=batch.seq_lens,
|
||||
seq_lens_cpu=batch.seq_lens_cpu,
|
||||
last_loc=torch.cat(last_loc),
|
||||
extend_num_tokens=batch.extend_num_tokens,
|
||||
req_pool_indices=req_pool_indices_device,
|
||||
dsv4_state_lens=_compute_dsv4_state_lens(batch, is_decode=False),
|
||||
batch=batch,
|
||||
)
|
||||
|
||||
# Write to req_to_token_pool
|
||||
write_cache_indices(
|
||||
out_cache_loc,
|
||||
req_pool_indices_device,
|
||||
req_pool_indices_cpu,
|
||||
prefix_lens_device,
|
||||
prefix_lens_cpu,
|
||||
batch.seq_lens,
|
||||
batch.seq_lens_cpu,
|
||||
extend_lens_device,
|
||||
extend_lens_cpu,
|
||||
prefix_tensors,
|
||||
batch.req_to_token_pool,
|
||||
)
|
||||
|
||||
# DSV4-NPU hook: no-op on non-DSV4 paths.
|
||||
if _is_npu:
|
||||
maybe_write_dsv4_extend(
|
||||
batch,
|
||||
req_pool_indices_cpu,
|
||||
prefix_lens_cpu,
|
||||
batch.seq_lens_cpu,
|
||||
)
|
||||
|
||||
return out_cache_loc, req_pool_indices_device, req_pool_indices_cpu
|
||||
|
||||
|
||||
def alloc_paged_token_slots_decode(
|
||||
tree_cache: BasePrefixCache,
|
||||
seq_lens: torch.Tensor,
|
||||
seq_lens_cpu: torch.Tensor,
|
||||
last_loc: torch.Tensor,
|
||||
token_per_req: int = 1,
|
||||
req_pool_indices: Optional[torch.Tensor] = None,
|
||||
dsv4_state_lens: Optional[DSV4StateLens] = None,
|
||||
batch=None,
|
||||
) -> torch.Tensor:
|
||||
"""Allocate paged KV cache for decode batch."""
|
||||
allocator = tree_cache.token_to_kv_pool_allocator
|
||||
# Over estimate the number of tokens: assume each request needs a new page.
|
||||
num_tokens = len(seq_lens) * allocator.page_size
|
||||
evict_from_tree_cache(tree_cache, num_tokens)
|
||||
|
||||
# DSV4-NPU allocator also needs req_pool_indices + per-req state lens and
|
||||
# returns a DSV4OutCacheLoc bundle; hasattr-gated so others stay unchanged.
|
||||
is_dsv4 = req_pool_indices is not None and hasattr(allocator, "c4_attn_allocator")
|
||||
extra_alloc_kwargs = {}
|
||||
if is_dsv4:
|
||||
extra_alloc_kwargs["req_pool_indices"] = req_pool_indices
|
||||
# Per-call per-req tables for the last_loc lookup.
|
||||
if batch is not None:
|
||||
extra_alloc_kwargs["req_to_token_pool"] = batch.req_to_token_pool
|
||||
if dsv4_state_lens is not None:
|
||||
extra_alloc_kwargs["dsv4_state_lens"] = dsv4_state_lens
|
||||
|
||||
out = allocator.alloc_decode(seq_lens, seq_lens_cpu, last_loc, **extra_alloc_kwargs)
|
||||
|
||||
if is_dsv4:
|
||||
bundle = out
|
||||
out_cache_loc = None if bundle is None else bundle.out_full_loc
|
||||
if batch is not None:
|
||||
batch.out_cache_loc_dsv4 = bundle
|
||||
else:
|
||||
out_cache_loc = out
|
||||
|
||||
if out_cache_loc is None:
|
||||
error_msg = (
|
||||
f"Decode out of memory. Try to lower your batch size.\n"
|
||||
f"Try to allocate {len(seq_lens) * token_per_req} tokens.\n"
|
||||
f"{available_and_evictable_str(tree_cache)}"
|
||||
)
|
||||
logger.error(error_msg)
|
||||
if tree_cache is not None:
|
||||
tree_cache.pretty_print()
|
||||
raise RuntimeError(error_msg)
|
||||
|
||||
return out_cache_loc
|
||||
|
||||
|
||||
def alloc_for_decode(batch: ScheduleBatch, token_per_req: int) -> torch.Tensor:
|
||||
"""
|
||||
Allocate KV cache for decode batch and write to req_to_token_pool.
|
||||
|
||||
Returns:
|
||||
out_cache_loc: allocated cache locations
|
||||
"""
|
||||
|
||||
batch.maybe_evict_swa()
|
||||
|
||||
seq_lens_gpu = batch.seq_lens
|
||||
bs = seq_lens_gpu.shape[0]
|
||||
|
||||
if _alloc_page_size(batch) == 1:
|
||||
# Non-paged allocation
|
||||
out_cache_loc = alloc_token_slots(batch.tree_cache, bs * token_per_req)
|
||||
else:
|
||||
# Paged allocation
|
||||
last_loc = batch.req_to_token_pool.req_to_token[
|
||||
batch.req_pool_indices, seq_lens_gpu - 1
|
||||
]
|
||||
seq_lens_next = seq_lens_gpu + token_per_req
|
||||
out_cache_loc = alloc_paged_token_slots_decode(
|
||||
tree_cache=batch.tree_cache,
|
||||
seq_lens=seq_lens_next,
|
||||
seq_lens_cpu=batch.seq_lens_cpu + token_per_req,
|
||||
last_loc=last_loc,
|
||||
token_per_req=token_per_req,
|
||||
req_pool_indices=batch.req_pool_indices,
|
||||
dsv4_state_lens=_compute_dsv4_state_lens(batch, is_decode=True),
|
||||
batch=batch,
|
||||
)
|
||||
|
||||
# Write to req_to_token_pool
|
||||
if batch.model_config.is_encoder_decoder:
|
||||
locs = batch.encoder_lens + seq_lens_gpu
|
||||
else:
|
||||
locs = seq_lens_gpu.clone()
|
||||
|
||||
batch.req_to_token_pool.write(
|
||||
(batch.req_pool_indices, locs), out_cache_loc.to(torch.int32)
|
||||
)
|
||||
|
||||
# DSV4-NPU hook: no-op on non-DSV4 paths.
|
||||
if _is_npu:
|
||||
maybe_write_dsv4_decode(
|
||||
batch,
|
||||
batch.seq_lens_cpu + token_per_req,
|
||||
token_per_req,
|
||||
)
|
||||
|
||||
return out_cache_loc
|
||||
@@ -0,0 +1,60 @@
|
||||
from __future__ import annotations
|
||||
|
||||
from typing import Optional
|
||||
|
||||
from sglang.srt.runtime_context import get_server_args
|
||||
from sglang.srt.server_args import ServerArgs
|
||||
|
||||
|
||||
def get_alloc_len_per_decode(server_args: Optional[ServerArgs] = None) -> int:
|
||||
if server_args is None:
|
||||
server_args = get_server_args()
|
||||
|
||||
if server_args.speculative_algorithm is None:
|
||||
return 1
|
||||
|
||||
# Spec decoding allocates max(topk * num_steps, num_draft_tokens) per decode step.
|
||||
spec_steps = server_args.speculative_num_steps or 1
|
||||
spec_topk = server_args.speculative_eagle_topk or 1
|
||||
spec_tokens = server_args.max_speculative_num_draft_tokens
|
||||
page_size = server_args.page_size
|
||||
|
||||
from sglang.srt.speculative.spec_info import SpeculativeAlgorithm
|
||||
|
||||
spec_algo = SpeculativeAlgorithm.from_string(server_args.speculative_algorithm)
|
||||
if page_size == 1 or spec_topk == 1 or not spec_algo.has_draft_kv():
|
||||
return max(spec_steps * spec_topk, spec_tokens)
|
||||
else:
|
||||
# spec v2 tree (page>1, topk>1): worst-case page-aligned footprint per
|
||||
# topk branch is ceil((page_size-1 + num_steps) / page) pages, each branch
|
||||
# duplicated -- reserve for all topk branches.
|
||||
num_new_pages_per_topk = (
|
||||
(page_size - 1) + spec_steps + page_size - 1
|
||||
) // page_size
|
||||
return max(num_new_pages_per_topk * page_size * spec_topk, spec_tokens)
|
||||
|
||||
|
||||
def get_alloc_reserve_per_decode(server_args: Optional[ServerArgs] = None) -> int:
|
||||
"""KV length reserved per request at each decode step.
|
||||
|
||||
The 2x is a double-buffer that absorbs the kv_committed_len lag in overlap
|
||||
mode; see eagle_utils.eagle_prepare_for_decode.
|
||||
"""
|
||||
return 2 * get_alloc_len_per_decode(server_args)
|
||||
|
||||
|
||||
def get_req_to_token_extra_context_len(server_args: ServerArgs) -> int:
|
||||
"""req_to_token row headroom beyond the model context length.
|
||||
|
||||
Sized to hold the decode over-allocation; the spec v2 page>1 topk>1 holey
|
||||
draft footprint can outgrow the default num_draft_tokens headroom.
|
||||
"""
|
||||
# FIXME(lsyin): temporary fix for the context length issue under spec decoding
|
||||
extra = 4 + (server_args.max_speculative_num_draft_tokens or 0)
|
||||
if (
|
||||
server_args.speculative_algorithm is not None
|
||||
and server_args.page_size > 1
|
||||
and (server_args.speculative_eagle_topk or 1) > 1
|
||||
):
|
||||
extra = max(extra, get_alloc_reserve_per_decode(server_args))
|
||||
return extra
|
||||
@@ -1,7 +1,7 @@
|
||||
from __future__ import annotations
|
||||
|
||||
import logging
|
||||
from typing import TYPE_CHECKING, Optional
|
||||
from typing import TYPE_CHECKING
|
||||
|
||||
import numpy as np
|
||||
import torch
|
||||
@@ -10,34 +10,18 @@ from sglang.kernels.ops.memory.common import (
|
||||
_get_last_loc_safe_kernel as _get_last_loc_safe_kernel,
|
||||
)
|
||||
from sglang.kernels.ops.memory.common import get_last_loc_kernel as get_last_loc_kernel
|
||||
from sglang.kernels.ops.memory.common import (
|
||||
get_last_loc_triton,
|
||||
get_last_loc_triton_safe,
|
||||
write_req_to_token_pool_triton,
|
||||
)
|
||||
from sglang.srt.hardware_backend.npu.dsv4.dsv4_common_hooks import (
|
||||
maybe_evict_dsv4_state_on_swa,
|
||||
maybe_write_dsv4_decode,
|
||||
maybe_write_dsv4_extend,
|
||||
)
|
||||
from sglang.srt.mem_cache.allocator.swa import SWATokenToKVPoolAllocator
|
||||
from sglang.srt.mem_cache.base_prefix_cache import BasePrefixCache, EvictParams
|
||||
from sglang.srt.mem_cache.memory_pool import HybridReqToTokenPool, ReqToTokenPool
|
||||
from sglang.srt.runtime_context import get_server_args
|
||||
from sglang.srt.server_args import ServerArgs
|
||||
from sglang.srt.utils import is_cuda, is_hip, is_npu, support_triton
|
||||
from sglang.srt.utils.common import ceil_align, is_pin_memory_available
|
||||
|
||||
_is_npu = is_npu()
|
||||
|
||||
_is_hip = is_hip()
|
||||
|
||||
_is_cuda = is_cuda()
|
||||
from sglang.srt.utils.common import ceil_align
|
||||
|
||||
if TYPE_CHECKING:
|
||||
from sglang.srt.managers.schedule_batch import Req, ScheduleBatch
|
||||
from sglang.srt.managers.schedule_batch import Req
|
||||
from sglang.srt.mem_cache.allocator import BaseTokenToKVPoolAllocator
|
||||
from sglang.srt.model_executor.forward_batch_info import DSV4StateLens
|
||||
|
||||
# Needs 2 + 1 slots for mamba request with prefix cache. 2 for ping pong cache, 1 for running mamba state.
|
||||
MAMBA_STATE_PER_REQ_PREFIX_CACHE = 3
|
||||
@@ -118,179 +102,6 @@ def maybe_cache_unfinished_req(req: Req, tree_cache: BasePrefixCache, **kwargs):
|
||||
tree_cache.cache_unfinished_req(req, **kwargs)
|
||||
|
||||
|
||||
def write_cache_indices(
|
||||
out_cache_loc: torch.Tensor,
|
||||
req_pool_indices_tensor: torch.Tensor,
|
||||
req_pool_indices_cpu: torch.Tensor,
|
||||
prefix_lens_tensor: torch.Tensor,
|
||||
prefix_lens_cpu: torch.Tensor,
|
||||
seq_lens_tensor: torch.Tensor,
|
||||
seq_lens_cpu: torch.Tensor,
|
||||
extend_lens_tensor: torch.Tensor,
|
||||
extend_lens_cpu: torch.Tensor,
|
||||
prefix_tensors: list[torch.Tensor],
|
||||
req_to_token_pool: ReqToTokenPool,
|
||||
):
|
||||
if support_triton(get_server_args().attention_backend):
|
||||
prefix_pointers = torch.tensor(
|
||||
[t.data_ptr() for t in prefix_tensors],
|
||||
dtype=torch.uint64,
|
||||
pin_memory=is_pin_memory_available(req_to_token_pool.device),
|
||||
).to(req_to_token_pool.device, non_blocking=True)
|
||||
# TODO: some tensors can be reused for ForwardBatchInfo (e.g., extend_lens, cumsum_start)
|
||||
write_req_to_token_pool_triton[(req_pool_indices_tensor.shape[0],)](
|
||||
req_to_token_pool.req_to_token,
|
||||
req_pool_indices_tensor,
|
||||
prefix_pointers,
|
||||
prefix_lens_tensor,
|
||||
seq_lens_tensor,
|
||||
extend_lens_tensor,
|
||||
out_cache_loc,
|
||||
req_to_token_pool.req_to_token.shape[1],
|
||||
)
|
||||
else:
|
||||
pt = 0
|
||||
for i in range(req_pool_indices_cpu.shape[0]):
|
||||
req_idx = req_pool_indices_cpu[i].item()
|
||||
prefix_len = prefix_lens_cpu[i].item()
|
||||
seq_len = seq_lens_cpu[i].item()
|
||||
extend_len = extend_lens_cpu[i].item()
|
||||
|
||||
req_to_token_pool.write(
|
||||
(req_idx, slice(0, prefix_len)),
|
||||
prefix_tensors[i],
|
||||
)
|
||||
req_to_token_pool.write(
|
||||
(req_idx, slice(prefix_len, seq_len)),
|
||||
out_cache_loc[pt : pt + extend_len],
|
||||
)
|
||||
pt += extend_len
|
||||
|
||||
|
||||
def get_last_loc(
|
||||
req_to_token: torch.Tensor,
|
||||
req_pool_indices_tensor: torch.Tensor,
|
||||
prefix_lens_tensor: torch.Tensor,
|
||||
) -> torch.Tensor:
|
||||
attn_backend = get_server_args().attention_backend
|
||||
uses_triton_dispatch = attn_backend not in ("ascend", "torch_native")
|
||||
|
||||
if _is_hip and uses_triton_dispatch:
|
||||
# HIP-only: the legacy get_last_loc_triton kernel emits a
|
||||
# mixed-width int32->int64 store that Triton mis-compiles on HIP,
|
||||
# producing out-of-range last_loc values under EAGLE +
|
||||
# page_size>1 (e.g. with aiter unified attention or the triton
|
||||
# attention backend). The bug is in the Triton HIP codegen, not
|
||||
# in any particular attention backend, so route every HIP path
|
||||
# that would otherwise use get_last_loc_triton through the
|
||||
# int32-safe variant. Non-HIP hardware keeps the original
|
||||
# dispatcher below.
|
||||
return get_last_loc_triton_safe(
|
||||
req_to_token, req_pool_indices_tensor, prefix_lens_tensor
|
||||
)
|
||||
|
||||
if uses_triton_dispatch:
|
||||
impl = get_last_loc_triton
|
||||
else:
|
||||
impl = get_last_loc_torch
|
||||
|
||||
return impl(req_to_token, req_pool_indices_tensor, prefix_lens_tensor)
|
||||
|
||||
|
||||
def get_last_loc_torch(
|
||||
req_to_token: torch.Tensor,
|
||||
req_pool_indices_tensor: torch.Tensor,
|
||||
prefix_lens_tensor: torch.Tensor,
|
||||
) -> torch.Tensor:
|
||||
return torch.where(
|
||||
prefix_lens_tensor > 0,
|
||||
req_to_token[req_pool_indices_tensor, prefix_lens_tensor - 1],
|
||||
torch.full_like(prefix_lens_tensor, -1),
|
||||
)
|
||||
|
||||
|
||||
def get_alloc_len_per_decode(server_args: Optional[ServerArgs] = None) -> int:
|
||||
if server_args is None:
|
||||
server_args = get_server_args()
|
||||
|
||||
if server_args.speculative_algorithm is None:
|
||||
return 1
|
||||
|
||||
# Spec decoding allocates max(topk * num_steps, num_draft_tokens) per decode step.
|
||||
spec_steps = server_args.speculative_num_steps or 1
|
||||
spec_topk = server_args.speculative_eagle_topk or 1
|
||||
spec_tokens = server_args.max_speculative_num_draft_tokens
|
||||
page_size = server_args.page_size
|
||||
|
||||
from sglang.srt.speculative.spec_info import SpeculativeAlgorithm
|
||||
|
||||
spec_algo = SpeculativeAlgorithm.from_string(server_args.speculative_algorithm)
|
||||
if page_size == 1 or spec_topk == 1 or not spec_algo.has_draft_kv():
|
||||
return max(spec_steps * spec_topk, spec_tokens)
|
||||
else:
|
||||
# spec v2 tree (page>1, topk>1): worst-case page-aligned footprint per
|
||||
# topk branch is ceil((page_size-1 + num_steps) / page) pages, each branch
|
||||
# duplicated -- reserve for all topk branches.
|
||||
num_new_pages_per_topk = (
|
||||
(page_size - 1) + spec_steps + page_size - 1
|
||||
) // page_size
|
||||
return max(num_new_pages_per_topk * page_size * spec_topk, spec_tokens)
|
||||
|
||||
|
||||
def get_alloc_reserve_per_decode(server_args: Optional[ServerArgs] = None) -> int:
|
||||
"""KV length reserved per request at each decode step.
|
||||
|
||||
The 2x is a double-buffer that absorbs the kv_committed_len lag in overlap
|
||||
mode; see eagle_utils.eagle_prepare_for_decode.
|
||||
"""
|
||||
return 2 * get_alloc_len_per_decode(server_args)
|
||||
|
||||
|
||||
def get_req_to_token_extra_context_len(server_args: ServerArgs) -> int:
|
||||
"""req_to_token row headroom beyond the model context length.
|
||||
|
||||
Sized to hold the decode over-allocation; the spec v2 page>1 topk>1 holey
|
||||
draft footprint can outgrow the default num_draft_tokens headroom.
|
||||
"""
|
||||
# FIXME(lsyin): temporary fix for the context length issue under spec decoding
|
||||
extra = 4 + (server_args.max_speculative_num_draft_tokens or 0)
|
||||
if (
|
||||
server_args.speculative_algorithm is not None
|
||||
and server_args.page_size > 1
|
||||
and (server_args.speculative_eagle_topk or 1) > 1
|
||||
):
|
||||
extra = max(extra, get_alloc_reserve_per_decode(server_args))
|
||||
return extra
|
||||
|
||||
|
||||
def alloc_token_slots(
|
||||
tree_cache: BasePrefixCache,
|
||||
num_tokens: int,
|
||||
backup_state: bool = False,
|
||||
):
|
||||
allocator = tree_cache.token_to_kv_pool_allocator
|
||||
evict_from_tree_cache(tree_cache, num_tokens)
|
||||
|
||||
state = None
|
||||
if backup_state:
|
||||
state = allocator.backup_state()
|
||||
|
||||
out_cache_loc = allocator.alloc(num_tokens)
|
||||
|
||||
if out_cache_loc is None:
|
||||
error_msg = (
|
||||
f"Out of memory. Try to lower your batch size.\n"
|
||||
f"Try to allocate {num_tokens} tokens.\n"
|
||||
f"{available_and_evictable_str(tree_cache)}"
|
||||
)
|
||||
logger.error(error_msg)
|
||||
if tree_cache is not None:
|
||||
tree_cache.pretty_print()
|
||||
raise RuntimeError(error_msg)
|
||||
|
||||
return (out_cache_loc, state) if backup_state else out_cache_loc
|
||||
|
||||
|
||||
def evict_from_tree_cache(tree_cache: BasePrefixCache | None, num_tokens: int):
|
||||
if tree_cache is None:
|
||||
return
|
||||
@@ -317,318 +128,6 @@ def evict_from_tree_cache(tree_cache: BasePrefixCache | None, num_tokens: int):
|
||||
tree_cache.evict(EvictParams(num_tokens=num_tokens))
|
||||
|
||||
|
||||
def _compute_dsv4_state_lens(batch, *, is_decode: bool):
|
||||
"""Per-req c{4,128}_state pool alloc lens (``DSV4StateLens``) for this step.
|
||||
None on CUDA / non-V4 paths (allocator has no ``compute_dsv4_state_lens_*``).
|
||||
"""
|
||||
allocator = batch.token_to_kv_pool_allocator
|
||||
if not hasattr(allocator, "compute_dsv4_state_lens_extend"):
|
||||
return None
|
||||
if is_decode:
|
||||
return allocator.compute_dsv4_state_lens_decode(batch.reqs)
|
||||
return allocator.compute_dsv4_state_lens_extend(
|
||||
batch.reqs, batch.seq_lens_cpu.tolist()
|
||||
)
|
||||
|
||||
|
||||
def alloc_paged_token_slots_extend(
|
||||
tree_cache: BasePrefixCache,
|
||||
prefix_lens: torch.Tensor,
|
||||
prefix_lens_cpu: torch.Tensor,
|
||||
seq_lens: torch.Tensor,
|
||||
seq_lens_cpu: torch.Tensor,
|
||||
last_loc: torch.Tensor,
|
||||
extend_num_tokens: int,
|
||||
backup_state: bool = False,
|
||||
req_pool_indices: Optional[torch.Tensor] = None,
|
||||
dsv4_state_lens: Optional[DSV4StateLens] = None,
|
||||
batch=None,
|
||||
):
|
||||
# Over estimate the number of tokens: assume each request needs a new page.
|
||||
allocator = tree_cache.token_to_kv_pool_allocator
|
||||
num_tokens = extend_num_tokens + len(seq_lens_cpu) * allocator.page_size
|
||||
evict_from_tree_cache(tree_cache, num_tokens)
|
||||
|
||||
state = None
|
||||
if backup_state:
|
||||
state = allocator.backup_state()
|
||||
|
||||
is_dsv4 = req_pool_indices is not None and hasattr(allocator, "c4_attn_allocator")
|
||||
extra_alloc_kwargs = {}
|
||||
if is_dsv4:
|
||||
extra_alloc_kwargs["req_pool_indices"] = req_pool_indices
|
||||
# Per-call per-req tables for the c-pool / state last_loc lookup.
|
||||
if batch is not None:
|
||||
extra_alloc_kwargs["req_to_token_pool"] = batch.req_to_token_pool
|
||||
if dsv4_state_lens is not None:
|
||||
extra_alloc_kwargs["dsv4_state_lens"] = dsv4_state_lens
|
||||
|
||||
out = allocator.alloc_extend(
|
||||
prefix_lens,
|
||||
prefix_lens_cpu,
|
||||
seq_lens,
|
||||
seq_lens_cpu,
|
||||
last_loc,
|
||||
extend_num_tokens,
|
||||
**extra_alloc_kwargs,
|
||||
)
|
||||
|
||||
if is_dsv4:
|
||||
bundle = out
|
||||
out_cache_loc = None if bundle is None else bundle.out_full_loc
|
||||
if batch is not None:
|
||||
batch.out_cache_loc_dsv4 = bundle
|
||||
else:
|
||||
out_cache_loc = out
|
||||
|
||||
if out_cache_loc is None:
|
||||
error_msg = (
|
||||
f"Prefill out of memory. Try to lower your batch size.\n"
|
||||
f"Try to allocate {extend_num_tokens} tokens.\n"
|
||||
f"{available_and_evictable_str(tree_cache)}"
|
||||
)
|
||||
logger.error(error_msg)
|
||||
if tree_cache is not None:
|
||||
tree_cache.pretty_print()
|
||||
raise RuntimeError(error_msg)
|
||||
|
||||
return (out_cache_loc, state) if backup_state else out_cache_loc
|
||||
|
||||
|
||||
def alloc_req_slots(
|
||||
req_to_token_pool: ReqToTokenPool,
|
||||
reqs: list[Req],
|
||||
tree_cache: BasePrefixCache | None,
|
||||
) -> list[int]:
|
||||
"""Allocate request slots from the pool.
|
||||
|
||||
Fail-loud: raises ``RuntimeError`` if the pool can't satisfy the batch. An
|
||||
alloc failure here means the admission budget (``PrefillAdder``) was wrong
|
||||
and should surface rather than be masked.
|
||||
"""
|
||||
num_reqs = len(reqs)
|
||||
if isinstance(req_to_token_pool, HybridReqToTokenPool):
|
||||
# Byte-coordinated for the shared allocator (accounts for the peer full
|
||||
# sub-pool's bytes); plain slot free count for the non-shared one.
|
||||
mamba_available_size = (
|
||||
req_to_token_pool.mamba_allocator.schedulable_available_size()
|
||||
)
|
||||
# Eviction headroom factor: 3x (or lazy variant) for radix COW, 1x for chunk.
|
||||
if tree_cache.supports_mamba():
|
||||
factor = (
|
||||
MAMBA_STATE_PER_REQ_PREFIX_CACHE_LAZY
|
||||
if req_to_token_pool.enable_mamba_extra_buffer_lazy
|
||||
else MAMBA_STATE_PER_REQ_PREFIX_CACHE
|
||||
)
|
||||
else:
|
||||
factor = MAMBA_STATE_PER_REQ_NO_CACHE
|
||||
mamba_state_needed = num_reqs * factor
|
||||
if mamba_available_size < mamba_state_needed:
|
||||
if tree_cache is not None and tree_cache.supports_mamba():
|
||||
mamba_num = max(0, mamba_state_needed - mamba_available_size)
|
||||
tree_cache.evict(EvictParams(num_tokens=0, mamba_num=mamba_num))
|
||||
req_pool_indices = req_to_token_pool.alloc(reqs)
|
||||
if req_pool_indices is None:
|
||||
raise RuntimeError(
|
||||
"alloc_req_slots runs out of memory. "
|
||||
"Please set a smaller number for `--max-running-requests`. "
|
||||
f"{req_to_token_pool.available_size()=}, {num_reqs=}, "
|
||||
)
|
||||
return req_pool_indices
|
||||
|
||||
|
||||
def _alloc_page_size(batch: ScheduleBatch) -> int:
|
||||
# DCP swaps in an allocator whose page_size is server_args.page_size *
|
||||
# dcp_size, so it can be > 1 even when tree_cache.page_size is 1; branch on
|
||||
# the real allocator's page_size there. Elsewhere the two are equal.
|
||||
if (_is_hip or _is_cuda) and get_server_args().dcp_size > 1:
|
||||
return batch.tree_cache.token_to_kv_pool_allocator.page_size
|
||||
return batch.tree_cache.page_size
|
||||
|
||||
|
||||
def alloc_for_extend(
|
||||
batch: ScheduleBatch,
|
||||
) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
|
||||
"""
|
||||
Allocate KV cache for extend batch and write to req_to_token_pool.
|
||||
|
||||
Returns ``(out_cache_loc, req_pool_indices_device, req_pool_indices_cpu)``
|
||||
(the last is the host/CPU mirror). ``alloc_req_slots`` raises ``RuntimeError``
|
||||
if the pool can't satisfy the batch (fail-loud — see its docstring).
|
||||
"""
|
||||
# free out-of-window swa tokens
|
||||
batch.maybe_evict_swa()
|
||||
|
||||
prefix_tensors = [r.prefix_indices for r in batch.reqs]
|
||||
|
||||
# Create tensors for allocation
|
||||
prefix_lens_cpu = torch.tensor(batch.prefix_lens, dtype=torch.int64)
|
||||
extend_lens_cpu = torch.tensor(batch.extend_lens, dtype=torch.int64)
|
||||
prefix_lens_device = prefix_lens_cpu.to(batch.device, non_blocking=True)
|
||||
extend_lens_device = extend_lens_cpu.to(batch.device, non_blocking=True)
|
||||
|
||||
# Allocate req slots (raises RuntimeError if the pool is exhausted)
|
||||
req_pool_indices = alloc_req_slots(
|
||||
batch.req_to_token_pool, batch.reqs, batch.tree_cache
|
||||
)
|
||||
req_pool_indices_cpu = torch.tensor(req_pool_indices, dtype=torch.int64)
|
||||
req_pool_indices_device = req_pool_indices_cpu.to(batch.device, non_blocking=True)
|
||||
|
||||
# Allocate KV cache (throws exception on failure)
|
||||
if _alloc_page_size(batch) == 1:
|
||||
out_cache_loc = alloc_token_slots(batch.tree_cache, batch.extend_num_tokens)
|
||||
else:
|
||||
# Paged allocation - build last_loc
|
||||
last_loc = [
|
||||
(t[-1:] if len(t) > 0 else torch.tensor([-1], device=batch.device))
|
||||
for t in prefix_tensors
|
||||
]
|
||||
out_cache_loc = alloc_paged_token_slots_extend(
|
||||
tree_cache=batch.tree_cache,
|
||||
prefix_lens=prefix_lens_device,
|
||||
prefix_lens_cpu=prefix_lens_cpu,
|
||||
seq_lens=batch.seq_lens,
|
||||
seq_lens_cpu=batch.seq_lens_cpu,
|
||||
last_loc=torch.cat(last_loc),
|
||||
extend_num_tokens=batch.extend_num_tokens,
|
||||
req_pool_indices=req_pool_indices_device,
|
||||
dsv4_state_lens=_compute_dsv4_state_lens(batch, is_decode=False),
|
||||
batch=batch,
|
||||
)
|
||||
|
||||
# Write to req_to_token_pool
|
||||
write_cache_indices(
|
||||
out_cache_loc,
|
||||
req_pool_indices_device,
|
||||
req_pool_indices_cpu,
|
||||
prefix_lens_device,
|
||||
prefix_lens_cpu,
|
||||
batch.seq_lens,
|
||||
batch.seq_lens_cpu,
|
||||
extend_lens_device,
|
||||
extend_lens_cpu,
|
||||
prefix_tensors,
|
||||
batch.req_to_token_pool,
|
||||
)
|
||||
|
||||
# DSV4-NPU hook: no-op on non-DSV4 paths.
|
||||
if _is_npu:
|
||||
maybe_write_dsv4_extend(
|
||||
batch,
|
||||
req_pool_indices_cpu,
|
||||
prefix_lens_cpu,
|
||||
batch.seq_lens_cpu,
|
||||
)
|
||||
|
||||
return out_cache_loc, req_pool_indices_device, req_pool_indices_cpu
|
||||
|
||||
|
||||
def alloc_paged_token_slots_decode(
|
||||
tree_cache: BasePrefixCache,
|
||||
seq_lens: torch.Tensor,
|
||||
seq_lens_cpu: torch.Tensor,
|
||||
last_loc: torch.Tensor,
|
||||
token_per_req: int = 1,
|
||||
req_pool_indices: Optional[torch.Tensor] = None,
|
||||
dsv4_state_lens: Optional[DSV4StateLens] = None,
|
||||
batch=None,
|
||||
) -> torch.Tensor:
|
||||
"""Allocate paged KV cache for decode batch."""
|
||||
allocator = tree_cache.token_to_kv_pool_allocator
|
||||
# Over estimate the number of tokens: assume each request needs a new page.
|
||||
num_tokens = len(seq_lens) * allocator.page_size
|
||||
evict_from_tree_cache(tree_cache, num_tokens)
|
||||
|
||||
# DSV4-NPU allocator also needs req_pool_indices + per-req state lens and
|
||||
# returns a DSV4OutCacheLoc bundle; hasattr-gated so others stay unchanged.
|
||||
is_dsv4 = req_pool_indices is not None and hasattr(allocator, "c4_attn_allocator")
|
||||
extra_alloc_kwargs = {}
|
||||
if is_dsv4:
|
||||
extra_alloc_kwargs["req_pool_indices"] = req_pool_indices
|
||||
# Per-call per-req tables for the last_loc lookup.
|
||||
if batch is not None:
|
||||
extra_alloc_kwargs["req_to_token_pool"] = batch.req_to_token_pool
|
||||
if dsv4_state_lens is not None:
|
||||
extra_alloc_kwargs["dsv4_state_lens"] = dsv4_state_lens
|
||||
|
||||
out = allocator.alloc_decode(seq_lens, seq_lens_cpu, last_loc, **extra_alloc_kwargs)
|
||||
|
||||
if is_dsv4:
|
||||
bundle = out
|
||||
out_cache_loc = None if bundle is None else bundle.out_full_loc
|
||||
if batch is not None:
|
||||
batch.out_cache_loc_dsv4 = bundle
|
||||
else:
|
||||
out_cache_loc = out
|
||||
|
||||
if out_cache_loc is None:
|
||||
error_msg = (
|
||||
f"Decode out of memory. Try to lower your batch size.\n"
|
||||
f"Try to allocate {len(seq_lens) * token_per_req} tokens.\n"
|
||||
f"{available_and_evictable_str(tree_cache)}"
|
||||
)
|
||||
logger.error(error_msg)
|
||||
if tree_cache is not None:
|
||||
tree_cache.pretty_print()
|
||||
raise RuntimeError(error_msg)
|
||||
|
||||
return out_cache_loc
|
||||
|
||||
|
||||
def alloc_for_decode(batch: ScheduleBatch, token_per_req: int) -> torch.Tensor:
|
||||
"""
|
||||
Allocate KV cache for decode batch and write to req_to_token_pool.
|
||||
|
||||
Returns:
|
||||
out_cache_loc: allocated cache locations
|
||||
"""
|
||||
|
||||
batch.maybe_evict_swa()
|
||||
|
||||
seq_lens_gpu = batch.seq_lens
|
||||
bs = seq_lens_gpu.shape[0]
|
||||
|
||||
if _alloc_page_size(batch) == 1:
|
||||
# Non-paged allocation
|
||||
out_cache_loc = alloc_token_slots(batch.tree_cache, bs * token_per_req)
|
||||
else:
|
||||
# Paged allocation
|
||||
last_loc = batch.req_to_token_pool.req_to_token[
|
||||
batch.req_pool_indices, seq_lens_gpu - 1
|
||||
]
|
||||
seq_lens_next = seq_lens_gpu + token_per_req
|
||||
out_cache_loc = alloc_paged_token_slots_decode(
|
||||
tree_cache=batch.tree_cache,
|
||||
seq_lens=seq_lens_next,
|
||||
seq_lens_cpu=batch.seq_lens_cpu + token_per_req,
|
||||
last_loc=last_loc,
|
||||
token_per_req=token_per_req,
|
||||
req_pool_indices=batch.req_pool_indices,
|
||||
dsv4_state_lens=_compute_dsv4_state_lens(batch, is_decode=True),
|
||||
batch=batch,
|
||||
)
|
||||
|
||||
# Write to req_to_token_pool
|
||||
if batch.model_config.is_encoder_decoder:
|
||||
locs = batch.encoder_lens + seq_lens_gpu
|
||||
else:
|
||||
locs = seq_lens_gpu.clone()
|
||||
|
||||
batch.req_to_token_pool.write(
|
||||
(batch.req_pool_indices, locs), out_cache_loc.to(torch.int32)
|
||||
)
|
||||
|
||||
# DSV4-NPU hook: no-op on non-DSV4 paths.
|
||||
if _is_npu:
|
||||
maybe_write_dsv4_decode(
|
||||
batch,
|
||||
batch.seq_lens_cpu + token_per_req,
|
||||
token_per_req,
|
||||
)
|
||||
|
||||
return out_cache_loc
|
||||
|
||||
|
||||
def release_kv_cache(req: Req, tree_cache: BasePrefixCache, is_insert: bool = True):
|
||||
# the two resources currently have the same lifecycle, thus simplify logic below
|
||||
assert (req.req_pool_idx is None) == (req.kv is None)
|
||||
|
||||
@@ -21,6 +21,7 @@ from sglang.srt.configs.model_config import (
|
||||
)
|
||||
from sglang.srt.distributed.parallel_state import get_world_group
|
||||
from sglang.srt.environ import envs
|
||||
from sglang.srt.mem_cache.allocation_sizing import get_req_to_token_extra_context_len
|
||||
from sglang.srt.mem_cache.allocator import (
|
||||
BaseTokenToKVPoolAllocator,
|
||||
PagedTokenToKVPoolAllocator,
|
||||
@@ -34,7 +35,6 @@ from sglang.srt.mem_cache.allocator.swa import (
|
||||
PureSWATokenToKVPoolAllocator,
|
||||
SWATokenToKVPoolAllocator,
|
||||
)
|
||||
from sglang.srt.mem_cache.common import get_req_to_token_extra_context_len
|
||||
from sglang.srt.mem_cache.deepseek_v4_memory_pool import DeepSeekV4TokenToKVPool
|
||||
from sglang.srt.mem_cache.hisparse_memory_pool import HiSparseDSATokenToKVPool
|
||||
from sglang.srt.mem_cache.memory_pool import (
|
||||
|
||||
@@ -30,7 +30,7 @@ from sglang.srt.configs.model_config import (
|
||||
is_minimax_sparse,
|
||||
)
|
||||
from sglang.srt.environ import envs
|
||||
from sglang.srt.mem_cache.common import get_alloc_len_per_decode
|
||||
from sglang.srt.mem_cache.allocation_sizing import get_alloc_len_per_decode
|
||||
from sglang.srt.mem_cache.deepseek_v4_memory_pool import get_compress_state_ring_size
|
||||
from sglang.srt.mem_cache.memory_pool import DSATokenToKVPool
|
||||
from sglang.srt.runtime_context import get_parallel
|
||||
|
||||
@@ -8,7 +8,7 @@ import torch
|
||||
|
||||
from sglang.srt.environ import envs
|
||||
from sglang.srt.managers.schedule_batch import ScheduleBatch
|
||||
from sglang.srt.mem_cache.common import (
|
||||
from sglang.srt.mem_cache.allocation import (
|
||||
alloc_paged_token_slots_extend,
|
||||
alloc_token_slots,
|
||||
get_last_loc,
|
||||
|
||||
@@ -18,12 +18,12 @@ from sglang.srt.hardware_backend.npu.dsv4.dsv4_allocator import (
|
||||
from sglang.srt.hardware_backend.npu.dsv4.dsv4_common_hooks import (
|
||||
maybe_build_dsv4_verify_bundle,
|
||||
)
|
||||
from sglang.srt.mem_cache.common import (
|
||||
from sglang.srt.mem_cache.allocation import (
|
||||
alloc_paged_token_slots_extend,
|
||||
alloc_token_slots,
|
||||
get_alloc_reserve_per_decode,
|
||||
get_last_loc,
|
||||
)
|
||||
from sglang.srt.mem_cache.allocation_sizing import get_alloc_reserve_per_decode
|
||||
from sglang.srt.runtime_context import get_parallel
|
||||
from sglang.srt.utils import (
|
||||
is_cpu,
|
||||
|
||||
Reference in New Issue
Block a user