From 87a09494fa3fbd685bd7c88d6a2dbdd3135de602 Mon Sep 17 00:00:00 2001 From: Liangsheng Yin Date: Tue, 18 Aug 2026 14:39:41 -0700 Subject: [PATCH] [Refactor] Share the page-aligned decode alloc lens between EAGLE and DFLASH (#35382) --- .../sglang/srt/mem_cache/allocation_sizing.py | 23 +++++++++++++ .../sglang/srt/speculative/dflash_info_v2.py | 33 ++++++++----------- python/sglang/srt/speculative/eagle_utils.py | 31 ++++++----------- 3 files changed, 47 insertions(+), 40 deletions(-) diff --git a/python/sglang/srt/mem_cache/allocation_sizing.py b/python/sglang/srt/mem_cache/allocation_sizing.py index 510123bdc..fc3c463da 100644 --- a/python/sglang/srt/mem_cache/allocation_sizing.py +++ b/python/sglang/srt/mem_cache/allocation_sizing.py @@ -47,6 +47,29 @@ def get_alloc_reserve_per_decode() -> int: return 2 * get_alloc_len_per_decode() +def page_aligned_decode_alloc_lens( + reqs, + *, + reserve: int, + page_size: int, +): + """Whole-page decode alloc lens: nxt rounds committed up to page so allocated + == recorded (unaligned tails leak at ps>1).""" + cur_kv_lens = [0] * len(reqs) + nxt_kv_lens = [0] * len(reqs) + num_needed_tokens = 0 + for i, r in enumerate(reqs): + cur = r.kv.kv_allocated_len + nxt = max( + cur, + (r.kv_committed_len + reserve + page_size - 1) // page_size * page_size, + ) + cur_kv_lens[i] = cur + nxt_kv_lens[i] = nxt + num_needed_tokens += nxt - cur + return cur_kv_lens, nxt_kv_lens, num_needed_tokens + + def get_req_to_token_extra_context_len() -> int: """req_to_token row headroom beyond the model context length. diff --git a/python/sglang/srt/speculative/dflash_info_v2.py b/python/sglang/srt/speculative/dflash_info_v2.py index b1afb6227..7dba1d74b 100644 --- a/python/sglang/srt/speculative/dflash_info_v2.py +++ b/python/sglang/srt/speculative/dflash_info_v2.py @@ -9,6 +9,7 @@ import torch from sglang.srt.environ import envs from sglang.srt.managers.schedule_batch import ScheduleBatch from sglang.srt.mem_cache.allocation import alloc_for_spec_decode +from sglang.srt.mem_cache.allocation_sizing import page_aligned_decode_alloc_lens from sglang.srt.runtime_context import get_spec from sglang.srt.speculative.spec_info import SpecInput, SpecInputType from sglang.srt.utils.common import is_pin_memory_available @@ -135,6 +136,7 @@ class DFlashDraftInputV2(SpecInput): assert self._prepare_nxt_kv_lens_gpu_buf is not None batch_seq_lens_cpu_t = self._prepare_batch_seq_lens_cpu_buf[:bs] cur_kv_lens_cpu_t = self._prepare_cur_kv_lens_cpu_buf[:bs] + nxt_kv_lens_cpu_t = self._prepare_nxt_kv_lens_cpu_buf[:bs] # For DFLASH, each decode step needs a fixed-size verify block. block_size = int(get_spec().speculative_num_draft_tokens) @@ -142,37 +144,30 @@ class DFlashDraftInputV2(SpecInput): raise ValueError( f"DFLASH invalid speculative_num_draft_tokens={block_size}." ) + reserve = 2 * block_size page_size = batch.token_to_kv_pool_allocator.page_size - nxt_kv_lens_cpu_t = self._prepare_nxt_kv_lens_cpu_buf[:bs] - committed_seq_lens_sum = 0 - nxt_kv_lens_sum = 0 - num_needed_tokens = 0 + + cur_kv_lens, nxt_kv_lens, num_needed_tokens = page_aligned_decode_alloc_lens( + batch.reqs, + reserve=reserve, + page_size=page_size, + ) + max_top_k = 1 uniform_top_k_value = None uniform_top_k = True - for i, req in enumerate(batch.reqs): + nxt_kv_lens_sum = 0 + committed_seq_lens_sum = 0 + for i, (req, cur, nxt) in enumerate(zip(batch.reqs, cur_kv_lens, nxt_kv_lens)): committed_len = int(req.kv_committed_len) - # Read the allocation watermark from the req object like EAGLE. - cur = int(req.kv.kv_allocated_len) - # Whole-page accounting (same as eagle_prepare_for_decode): the - # paged allocator hands out full pages, so an unaligned reserve - # strands the tail of the last page -- allocated but never recorded. - nxt = max( - cur, - (committed_len + 2 * block_size + page_size - 1) - // page_size - * page_size, - ) + committed_seq_lens_sum += committed_len top_k = int(req.sampling_params.top_k) batch_seq_lens_cpu_t[i] = committed_len cur_kv_lens_cpu_t[i] = cur nxt_kv_lens_cpu_t[i] = nxt - committed_seq_lens_sum += committed_len nxt_kv_lens_sum += nxt - num_needed_tokens += nxt - cur - if top_k > max_top_k: max_top_k = top_k if i == 0: diff --git a/python/sglang/srt/speculative/eagle_utils.py b/python/sglang/srt/speculative/eagle_utils.py index c4e23ec62..662ca8091 100644 --- a/python/sglang/srt/speculative/eagle_utils.py +++ b/python/sglang/srt/speculative/eagle_utils.py @@ -15,7 +15,10 @@ from sglang.srt.hardware_backend.npu.dsv4.dsv4_common_hooks import ( maybe_build_dsv4_verify_bundle, ) from sglang.srt.mem_cache.allocation import alloc_for_spec_decode -from sglang.srt.mem_cache.allocation_sizing import get_alloc_reserve_per_decode +from sglang.srt.mem_cache.allocation_sizing import ( + get_alloc_reserve_per_decode, + page_aligned_decode_alloc_lens, +) from sglang.srt.runtime_context import get_parallel, get_spec from sglang.srt.utils import ( is_cpu, @@ -907,26 +910,12 @@ def eagle_prepare_for_decode(batch: ScheduleBatch): page_size = batch.token_to_kv_pool_allocator.page_size double_alloc = get_alloc_reserve_per_decode() - cur_kv_lens = [0] * bs - nxt_kv_lens = [0] * bs - num_needed_tokens = 0 - for i, r in enumerate(batch.reqs): - cur = r.kv.kv_allocated_len - # max(cur, ...) clamps so adaptive downswitch cannot make nxt < cur. - # kv_committed_len is honest (bonus committed in resolve, not here), - # so it lags batch.seq_lens by ~1 verify in overlap; 2*alloc absorbs. - # Whole-page accounting: the paged allocator hands out full pages, so - # round nxt up to the page boundary or the unaligned tail is allocated - # but never recorded — a stranded-tail leak at page_size > 1. - nxt = max( - cur, - (r.kv_committed_len + double_alloc + page_size - 1) - // page_size - * page_size, - ) - cur_kv_lens[i] = cur - nxt_kv_lens[i] = nxt - num_needed_tokens += nxt - cur + cur_kv_lens, nxt_kv_lens, num_needed_tokens = page_aligned_decode_alloc_lens( + batch.reqs, + reserve=double_alloc, + page_size=page_size, + ) + for r in batch.reqs: r.decode_batch_idx += 1 cur_kv_lens_cpu = torch.tensor(cur_kv_lens, dtype=torch.int32, device="cpu")