[Spec] Rename num_tokens_per_bs to num_tokens_per_req (#30977)

This commit is contained in:
Liangsheng Yin
2026-07-13 13:47:53 -05:00
committed by GitHub
parent b677babc62
commit c0f1f7e062
38 changed files with 282 additions and 246 deletions
@@ -68,12 +68,12 @@ class EagerRunner(BaseRunner):
sa = mr.server_args
# Built first so the cg runners coalesce onto its buffers via the shared
# input pool; size to the largest tokens/req across modes the worker hits.
num_tokens_per_bs = 1
num_tokens_per_req = 1
if mr.spec_algorithm.is_speculative():
# speculative_adaptive can grow draft tokens at runtime; size to the max.
num_draft_tokens = sa.max_speculative_num_draft_tokens or 1
if mr.is_draft_worker:
num_tokens_per_bs = max(
num_tokens_per_req = max(
sa.speculative_eagle_topk or 1,
num_draft_tokens,
(
@@ -83,8 +83,8 @@ class EagerRunner(BaseRunner):
),
)
else:
num_tokens_per_bs = (
mr.spec_algorithm.get_num_tokens_per_bs_for_target_verify(
num_tokens_per_req = (
mr.spec_algorithm.get_num_tokens_per_req_for_target_verify(
num_draft_tokens, mr.is_draft_worker
)
)
@@ -92,7 +92,7 @@ class EagerRunner(BaseRunner):
dllm_config = DllmConfig.from_server_args(sa)
if dllm_config is not None:
# dLLM runs block_size tokens/request (DLLM_EXTEND).
num_tokens_per_bs = dllm_config.block_size
num_tokens_per_req = dllm_config.block_size
max_bs = mr.max_running_requests
if (
mr.is_draft_worker
@@ -109,12 +109,12 @@ class EagerRunner(BaseRunner):
max_bs = ceil_align(max_bs, self.attn_tp_size)
max_bs = ceil_align(max_bs, get_cp_padding_align_size())
prefill_ceiling = max(mr.max_total_num_tokens, sa.max_prefill_buffer_tokens())
max_num_token = max(prefill_ceiling, max_bs * num_tokens_per_bs)
max_num_token = max(prefill_ceiling, max_bs * num_tokens_per_req)
if require_mlp_sync(sa):
max_num_token = ceil_align(max_num_token, self.attn_tp_size)
max_num_token = ceil_align(max_num_token, get_cp_padding_align_size())
self._eager_max_bs = max_bs
self._eager_num_tokens_per_bs = num_tokens_per_bs
self._eager_num_tokens_per_req = num_tokens_per_req
is_encoder_decoder = mr.model_config.is_encoder_decoder
self._eager_registry = build_eager_registry(
device=mr.device,
@@ -139,7 +139,7 @@ class EagerRunner(BaseRunner):
self.warmup()
def _autotune_buffers(self) -> Tuple[Any, int]:
"""Decode-shaped dummy buffers (bs * num_tokens_per_bs) for the warmup
"""Decode-shaped dummy buffers (bs * num_tokens_per_req) for the warmup
flashinfer-autotune forward.
flashinfer's MoE autotuner times candidate tactics against the buffer it
@@ -148,16 +148,16 @@ class EagerRunner(BaseRunner):
ceiling; the dummy run only needs the decode-sized slice.
"""
mr = self.model_runner
num_tokens_per_bs = 1
num_tokens_per_req = 1
if mr.spec_algorithm.is_speculative():
num_tokens_per_bs = (
mr.spec_algorithm.get_num_tokens_per_bs_for_target_verify(
num_tokens_per_req = (
mr.spec_algorithm.get_num_tokens_per_req_for_target_verify(
mr.server_args.speculative_num_draft_tokens, mr.is_draft_worker
)
)
return (
self._alloc_dummy_decode_buffers(
self._eager_max_bs, num_tokens_per_bs=num_tokens_per_bs
self._eager_max_bs, num_tokens_per_req=num_tokens_per_req
),
self._eager_max_bs,
)