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

[dLLM] Reuse block KV/req slots in place across FDFO rounds (#27877)

Browse Source 
Co-authored-by: Xiaoyu Zhang <1182563586@qq.com>
This commit is contained in:
Chenchen Hong
2026-07-20 13:36:13 +08:00
committed by GitHub
co-authored by Xiaoyu Zhang
parent b15a83983c
commit 49b9c46f41
5 changed files with 360 additions and 14 deletions
Download Patch File Download Diff File Expand all files Collapse all files
python/sglang/srt/dllm/mixin/scheduler.py
+21 -12
View File
@@ -78,8 +78,7 @@ class SchedulerDllmMixin:
not fdfo_mode or result.accept_length_per_req_cpu is not None
), "FDFO dLLM result is missing accept lengths."
# Sync mode emits tokens only once a block fully resolves; FDFO always
# commits (resolved blocks decode, unresolved blocks stash + free KV).
# FDFO also commits unresolved blocks so their KV can be reused.
if fdfo_mode or result.next_token_ids:
block_size = self.dllm_config.block_size
algo_states = result.dllm_algo_state
@@ -111,20 +110,11 @@ class SchedulerDllmMixin:
assert len(next_token_ids) == block_size
if result.accept_length_per_req_cpu[idx] == 0:
# Block unresolved: stash partial state and free the KV slots
# of the still-masked block so the next FDFO round can
# re-denoise it without leaking the previous allocation.
# Unresolved: keep partial state and KV for the next FDFO round.
req.dllm_incomplete_ids = array("q", next_token_ids)
req.dllm_algo_state = (
algo_states[idx] if algo_states is not None else None
)
old_prefix_len = len(req.prefix_indices)
new_fill_len = req.extend_range.end
if new_fill_len > old_prefix_len:
kv_indices_to_free = self.req_to_token_pool.req_to_token[
req.req_pool_idx, old_prefix_len:new_fill_len
]
self.token_to_kv_pool_allocator.free(kv_indices_to_free)
continue
req.dllm_incomplete_ids = array("q")
@@ -426,6 +416,25 @@ class DllmManager:
self.waiting_queue = [req for req in self.waiting_queue if not req.finished()]
self.staging_queue = [req for req in self.staging_queue if not req.finished()]
def pop_aborted_reqs(self, abort_all: bool, rid: str) -> List[Req]:
aborted_reqs: List[Req] = []
seen: Set[int] = set()
for queue_name in ("waiting_queue", "staging_queue"):
queue = getattr(self, queue_name)
kept_queue = []
for req in queue:
if abort_all or req.rid.startswith(rid):
req_id = id(req)
if req_id not in seen:
aborted_reqs.append(req)
seen.add(req_id)
else:
kept_queue.append(req)
setattr(self, queue_name, kept_queue)
return aborted_reqs
def init_next_round(self) -> None:
"""Initialize staging requests for next round and clear staging queue."""
for req in self.staging_queue:
python/sglang/srt/managers/schedule_policy.py
+2
View File
@@ -776,6 +776,8 @@ class PrefillAdder:
cand_extend_input_len = len(req.full_untruncated_fill_ids) - len(
req.prefix_indices
)
if req.dllm_incomplete_ids and cand_extend_input_len > _rem_tokens:
return AddReqResult.NO_TOKEN
truncated = cand_extend_input_len > _rem_tokens
new_len = min(cand_extend_input_len, _rem_tokens)
req.set_extend_range(len(req.prefix_indices), len(req.prefix_indices) + new_len)
python/sglang/srt/managers/scheduler.py
+18 -1
View File
@@ -2693,7 +2693,8 @@ class Scheduler(
if self.dllm_config.first_done_first_out_mode:
if not req.dllm_incomplete_ids:
self.stash_chunked_request(req)
self.req_to_token_pool.free(req)
self.req_to_token_pool.free(req)
# Otherwise, keep req slot/KV for reuse.
else:
self.stash_chunked_request(req)
@@ -4061,6 +4062,22 @@ class Scheduler(
release_kv_cache(req, self.tree_cache, is_insert=False)
logger.debug(f"Abort queued request. {req.rid=}")
if self.dllm_config is not None:
for req in self.dllm_manager.pop_aborted_reqs(
recv_req.abort_all, recv_req.rid
):
if self.enable_hicache_storage:
self.tree_cache.release_aborted_request(req.rid)
self.ipc_channels.send_to_tokenizer.send_output(
AbortReq(rid=req.rid), req
)
if (
req.req_pool_idx is not None
or getattr(req, "mamba_pool_idx", None) is not None
):
release_kv_cache(req, self.tree_cache, is_insert=False)
logger.debug(f"Abort dLLM queued request. {req.rid=}")
# Delete the requests in the grammar queue
# Abort method 2: call `set_finish_with_abort`
# The request will still run one prefill forward pass.
python/sglang/srt/mem_cache/allocation.py
+100 -1
View File
@@ -315,6 +315,13 @@ def alloc_for_extend(
prefix_tensors = [r.prefix_indices for r in batch.reqs]
reuse_kv = None
if batch.is_dllm():
reuse_kv = [
r.req_pool_idx is not None and bool(r.dllm_incomplete_ids)
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)
@@ -329,7 +336,18 @@ def alloc_for_extend(
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:
alloc_page_size = _alloc_page_size(batch)
if reuse_kv is not None and any(reuse_kv):
out_cache_loc = _alloc_extend_loc_with_kv_reuse(
batch,
reuse_kv,
req_pool_indices_cpu,
prefix_lens_cpu,
extend_lens_cpu,
req_pool_indices_device,
alloc_page_size,
)
elif alloc_page_size == 1:
out_cache_loc = alloc_token_slots(batch.tree_cache, batch.extend_num_tokens)
else:
# Paged allocation - build last_loc
@@ -385,6 +403,87 @@ def alloc_for_extend(
return out_cache_loc, req_pool_indices_device, req_pool_indices_cpu
def _alloc_extend_loc_with_kv_reuse(
batch: ScheduleBatch,
reuse_kv: list[bool],
req_pool_indices_cpu: torch.Tensor,
prefix_lens_cpu: torch.Tensor,
extend_lens_cpu: torch.Tensor,
req_pool_indices_device: torch.Tensor,
alloc_page_size: int,
) -> torch.Tensor:
device = batch.device
req_to_token = batch.req_to_token_pool.req_to_token
for i, req in enumerate(batch.reqs):
if not reuse_kv[i]:
continue
prefix_len = int(prefix_lens_cpu[i])
extend_len = int(extend_lens_cpu[i])
retained_len = len(req.dllm_incomplete_ids)
if extend_len != retained_len:
raise RuntimeError("dLLM FDFO retained KV must be reused as a full block.")
if req.kv is None or prefix_len + extend_len > req.kv.kv_allocated_len:
raise RuntimeError("dLLM FDFO retained KV is missing.")
alloc_extend_lens = [
0 if reuse_kv[i] else int(extend_lens_cpu[i]) for i in range(len(reuse_kv))
]
alloc_extend_num_tokens = sum(alloc_extend_lens)
fresh_slots = None
if alloc_extend_num_tokens > 0:
if alloc_page_size == 1:
fresh_slots = alloc_token_slots(batch.tree_cache, alloc_extend_num_tokens)
else:
alloc_seq_lens_cpu = torch.tensor(
[
(
int(prefix_lens_cpu[i])
if reuse_kv[i]
else int(batch.seq_lens_cpu[i])
)
for i in range(len(reuse_kv))
],
dtype=torch.int64,
)
last_loc = [
(t[-1:] if len(t) > 0 else torch.tensor([-1], device=device))
for t in (r.prefix_indices for r in batch.reqs)
]
fresh_slots = alloc_paged_token_slots_extend(
tree_cache=batch.tree_cache,
prefix_lens=prefix_lens_cpu.to(device, non_blocking=True),
prefix_lens_cpu=prefix_lens_cpu,
seq_lens=alloc_seq_lens_cpu.to(device, non_blocking=True),
seq_lens_cpu=alloc_seq_lens_cpu,
last_loc=torch.cat(last_loc),
extend_num_tokens=alloc_extend_num_tokens,
req_pool_indices=req_pool_indices_device,
dsv4_state_lens=_compute_dsv4_state_lens(batch, is_decode=False),
batch=batch,
)
reuse_dtype = fresh_slots.dtype if fresh_slots is not None else torch.int64
parts: list[torch.Tensor] = []
fresh_ptr = 0
for i in range(len(reuse_kv)):
prefix_len = int(prefix_lens_cpu[i])
extend_len = int(extend_lens_cpu[i])
if reuse_kv[i]:
req_idx = int(req_pool_indices_cpu[i])
parts.append(
req_to_token[req_idx, prefix_len : prefix_len + extend_len].to(
reuse_dtype
)
)
else:
parts.append(fresh_slots[fresh_ptr : fresh_ptr + extend_len])
fresh_ptr += extend_len
return torch.cat(parts)
def alloc_paged_token_slots_decode(
tree_cache: BasePrefixCache,
seq_lens: torch.Tensor,