[Spec] Rename num_tokens_per_bs to num_tokens_per_req (#30977)
This commit is contained in:
@@ -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,
|
||||
)
|
||||
|
||||
Reference in New Issue
Block a user