Refactor Req.fill_ids into full_untruncated_fill_ids + fill_len with equivalence (#26637)
This commit is contained in:
@@ -371,9 +371,10 @@ def prepare_inputs_for_correctness_test(bench_args, tokenizer, custom_prompts):
|
||||
origin_input_ids=array("q", tmp_input_ids),
|
||||
sampling_params=sampling_params,
|
||||
)
|
||||
req.fill_ids = req.origin_input_ids
|
||||
req.full_untruncated_fill_ids = req.origin_input_ids
|
||||
req.fill_len = len(req.full_untruncated_fill_ids)
|
||||
req.logprob_start_len = -1
|
||||
req.set_extend_input_len(len(req.fill_ids) - len(req.prefix_indices))
|
||||
req.set_extend_input_len(req.fill_len - len(req.prefix_indices))
|
||||
reqs.append(req)
|
||||
|
||||
return input_ids, reqs
|
||||
@@ -384,14 +385,15 @@ def prepare_extend_inputs_for_correctness_test(
|
||||
):
|
||||
for i in range(len(reqs)):
|
||||
req: Req = reqs[i]
|
||||
req.fill_ids.extend(input_ids[i][bench_args.cut_len :])
|
||||
req.full_untruncated_fill_ids += input_ids[i][bench_args.cut_len :]
|
||||
req.fill_len = len(req.full_untruncated_fill_ids)
|
||||
if model_runner is not None:
|
||||
# Use req.req_pool_idx instead of i to handle slot 0 padding correctly
|
||||
req.prefix_indices = model_runner.req_to_token_pool.req_to_token[
|
||||
req.req_pool_idx, : bench_args.cut_len
|
||||
].to(req.prefix_indices.dtype)
|
||||
req.logprob_start_len = -1
|
||||
req.set_extend_input_len(len(req.fill_ids) - len(req.prefix_indices))
|
||||
req.set_extend_input_len(req.fill_len - len(req.prefix_indices))
|
||||
return reqs
|
||||
|
||||
|
||||
@@ -416,9 +418,10 @@ def prepare_synthetic_inputs_for_latency_test(
|
||||
origin_input_ids=array("q", input_ids[i]),
|
||||
sampling_params=sampling_params,
|
||||
)
|
||||
req.fill_ids = req.origin_input_ids
|
||||
req.full_untruncated_fill_ids = req.origin_input_ids
|
||||
req.fill_len = len(req.full_untruncated_fill_ids)
|
||||
req.logprob_start_len = -1
|
||||
req.set_extend_input_len(len(req.fill_ids) - len(req.prefix_indices))
|
||||
req.set_extend_input_len(req.fill_len - len(req.prefix_indices))
|
||||
reqs.append(req)
|
||||
|
||||
return reqs
|
||||
@@ -557,7 +560,7 @@ class _MlxBenchRunner:
|
||||
req_ids = [str(req.rid) for req in reqs]
|
||||
results = []
|
||||
for rid, req in zip(req_ids, reqs):
|
||||
token_ids = [int(t) for t in req.fill_ids]
|
||||
token_ids = [int(t) for t in req.get_fill_ids()]
|
||||
next_token = self.mlx_runner.prefill(
|
||||
req_id=rid,
|
||||
new_token_ids=token_ids,
|
||||
|
||||
@@ -1048,9 +1048,7 @@ class DecodePreallocQueue(DecodeHiCachePreallocMixin):
|
||||
|
||||
@property
|
||||
def num_tokens_pre_allocated(self):
|
||||
return sum(
|
||||
len(decode_req.req.fill_ids) for decode_req in self.transfer_queue.queue
|
||||
)
|
||||
return sum(decode_req.req.fill_len for decode_req in self.transfer_queue.queue)
|
||||
|
||||
def _need_space_for_single_req(
|
||||
self, retractable_tokens: Optional[int] = None
|
||||
@@ -1380,10 +1378,11 @@ class DecodePreallocQueue(DecodeHiCachePreallocMixin):
|
||||
kv_loc,
|
||||
)
|
||||
|
||||
# Truncate fill_ids to kv_committed_len so cache_unfinished_req only
|
||||
# Truncate fill_len to kv_committed_len so cache_unfinished_req only
|
||||
# inserts committed KV into the radix tree. The last output token
|
||||
# hasn't had KV committed yet (fill_ids is 1 ahead).
|
||||
req.fill_ids = (req.origin_input_ids + req.output_ids)[: req.kv_committed_len]
|
||||
# hasn't had KV committed yet (output_ids is 1 ahead).
|
||||
req.full_untruncated_fill_ids = req.origin_input_ids + req.output_ids
|
||||
req.fill_len = req.kv_committed_len
|
||||
# Set prefix_indices so downstream consumers (init_next_round_input,
|
||||
# prepare_for_extend) see the correct prefix length. In the agg path
|
||||
# this is done inside init_next_round_input, but decode-disagg needs
|
||||
@@ -1391,7 +1390,7 @@ class DecodePreallocQueue(DecodeHiCachePreallocMixin):
|
||||
req.prefix_indices = (
|
||||
prefix_indices if prefix_len > 0 else torch.empty((0,), dtype=torch.int64)
|
||||
)
|
||||
req.set_extend_input_len(len(req.fill_ids) - total_prefix_len)
|
||||
req.set_extend_input_len(req.fill_len - total_prefix_len)
|
||||
|
||||
# Return the transfer destination indices:
|
||||
if self.scheduler.enable_hisparse:
|
||||
@@ -1839,13 +1838,12 @@ class SchedulerDisaggregationDecodeMixin:
|
||||
else:
|
||||
tree_cache = self.tree_cache
|
||||
req.init_next_round_input(tree_cache)
|
||||
# Truncate fill_ids to kv_committed_len so cache_unfinished_req
|
||||
# only sees committed KV (fill_ids includes one uncommitted token).
|
||||
# Truncate fill_len to kv_committed_len so cache_unfinished_req
|
||||
# only sees committed KV (full array includes one uncommitted
|
||||
# token because init_next_round_input rebuilt it as full).
|
||||
if req.kv_committed_len is not None:
|
||||
req.fill_ids = req.fill_ids[: req.kv_committed_len]
|
||||
req.set_extend_input_len(
|
||||
len(req.fill_ids) - len(req.prefix_indices)
|
||||
)
|
||||
req.fill_len = req.kv_committed_len
|
||||
req.set_extend_input_len(req.fill_len - len(req.prefix_indices))
|
||||
else:
|
||||
waiting_queue.append(req)
|
||||
|
||||
|
||||
@@ -29,7 +29,7 @@ class ScheduleBatchDisaggregationDecodeMixin:
|
||||
|
||||
self.forward_mode = ForwardMode.PREBUILT
|
||||
reqs = self.reqs
|
||||
input_ids = [r.fill_ids[len(r.prefix_indices) :] for r in reqs]
|
||||
input_ids = [r.get_fill_ids()[len(r.prefix_indices) :] for r in reqs]
|
||||
extend_num_tokens = sum(len(ids) for ids in input_ids)
|
||||
seq_lens = []
|
||||
pre_lens = []
|
||||
|
||||
@@ -870,7 +870,7 @@ class SchedulerDisaggregationPrefillMixin:
|
||||
elif self.enable_overlap:
|
||||
# Delay KV transfer to process_batch_result_disagg_prefill when overlap is enabled to ensure results are resolved
|
||||
self.chunked_req.tmp_end_idx = min(
|
||||
len(self.chunked_req.fill_ids),
|
||||
self.chunked_req.fill_len,
|
||||
len(self.chunked_req.origin_input_ids),
|
||||
)
|
||||
else:
|
||||
@@ -906,7 +906,7 @@ class SchedulerDisaggregationPrefillMixin:
|
||||
end_idx = (
|
||||
end_idx
|
||||
if end_idx is not None
|
||||
else min(len(req.fill_ids), len(req.origin_input_ids))
|
||||
else min(req.fill_len, len(req.origin_input_ids))
|
||||
)
|
||||
|
||||
if not last_chunk:
|
||||
@@ -937,7 +937,7 @@ class SchedulerDisaggregationPrefillMixin:
|
||||
# length here avoids emitting an extra state page when the sampled
|
||||
# token crosses a page boundary, which mismatched src/dst lengths in
|
||||
# group_concurrent_contiguous.
|
||||
seq_len = min(len(req.fill_ids), len(req.origin_input_ids))
|
||||
seq_len = min(req.fill_len, len(req.origin_input_ids))
|
||||
|
||||
def _mamba_payload():
|
||||
return [
|
||||
|
||||
@@ -42,11 +42,11 @@ class ReqDllmMixin:
|
||||
prefix_length = len(self.prefix_indices)
|
||||
min_required_length = prefix_length + self.dllm_config.block_size
|
||||
|
||||
if len(self.fill_ids) < min_required_length:
|
||||
if self.fill_len < min_required_length:
|
||||
# still incoming stage
|
||||
return
|
||||
|
||||
input_block = self.fill_ids[prefix_length:min_required_length]
|
||||
input_block = self.get_fill_ids()[prefix_length:min_required_length]
|
||||
is_prefill_phase = self.dllm_config.mask_id not in input_block
|
||||
|
||||
if is_prefill_phase:
|
||||
@@ -57,14 +57,15 @@ class ReqDllmMixin:
|
||||
def _init_fill_ids_for_dllm(self: Req):
|
||||
self.dllm_block_offset = (
|
||||
0
|
||||
if not self.fill_ids
|
||||
if self.fill_len == 0
|
||||
else self.dllm_block_offset + self.dllm_config.block_size
|
||||
)
|
||||
self.fill_ids = (
|
||||
self.full_untruncated_fill_ids = (
|
||||
self.origin_input_ids
|
||||
+ self.output_ids
|
||||
+ array("q", [self.dllm_config.mask_id] * self.dllm_config.block_size)
|
||||
)
|
||||
self.fill_len = len(self.full_untruncated_fill_ids)
|
||||
|
||||
def _update_block_offset_for_dllm(self):
|
||||
prefix_len = len(self.prefix_indices)
|
||||
|
||||
@@ -80,7 +80,9 @@ class SchedulerDllmMixin:
|
||||
if new_tokens == 0:
|
||||
continue
|
||||
|
||||
req.fill_ids[-new_tokens:] = array("q", next_token_ids)
|
||||
req.full_untruncated_fill_ids[
|
||||
req.fill_len - new_tokens : req.fill_len
|
||||
] = array("q", next_token_ids)
|
||||
self.metrics_reporter.num_generated_tokens += new_tokens
|
||||
|
||||
req.output_ids.extend(next_token_ids)
|
||||
|
||||
@@ -185,7 +185,7 @@ class MlxTpModelWorker(TpModelWorker):
|
||||
else:
|
||||
# New prefill
|
||||
prefix_slot_ids = req.prefix_indices.tolist()
|
||||
full_token_ids = list(req.fill_ids)
|
||||
full_token_ids = list(req.get_fill_ids())
|
||||
next_token = self._mlx_runner.prefill(
|
||||
req_id=req.rid,
|
||||
new_token_ids=req_token_ids,
|
||||
@@ -330,7 +330,7 @@ class MlxTpModelWorker(TpModelWorker):
|
||||
else:
|
||||
# New prefill
|
||||
prefix_slot_ids = req.prefix_indices.tolist()
|
||||
full_token_ids = list(req.fill_ids)
|
||||
full_token_ids = list(req.get_fill_ids())
|
||||
pending_prefills.append(
|
||||
self._mlx_runner.prefill_start(
|
||||
req_id=req.rid,
|
||||
|
||||
@@ -207,7 +207,7 @@ class HiSparseCoordinator:
|
||||
req.hisparse_staging = True
|
||||
|
||||
full_kv_indices = self.req_to_token_pool.req_to_token[
|
||||
req.req_pool_idx, : len(req.fill_ids)
|
||||
req.req_pool_idx, : req.fill_len
|
||||
].to(dtype=torch.int64, copy=True)
|
||||
device_indices = (
|
||||
self.mem_pool_device.translate_loc_from_full_to_hisparse_device(
|
||||
@@ -298,7 +298,7 @@ class HiSparseCoordinator:
|
||||
|
||||
def alloc_device_buffer(self, req: Req) -> None:
|
||||
if self.is_dsv4_hisparse:
|
||||
allocated_len = len(req.fill_ids)
|
||||
allocated_len = req.fill_len
|
||||
alloc_size = self.padded_buffer_size
|
||||
else:
|
||||
allocated_len = req.kv_allocated_len
|
||||
@@ -707,7 +707,7 @@ class HiSparseCoordinator:
|
||||
# Wait for any in-flight staging DMA to complete before freeing
|
||||
self.write_staging_stream.synchronize()
|
||||
|
||||
prefill_len = len(req.fill_ids)
|
||||
prefill_len = req.fill_len
|
||||
allocated_locs = self.req_to_token_pool.req_to_token[
|
||||
req.req_pool_idx, :prefill_len
|
||||
]
|
||||
|
||||
@@ -696,8 +696,11 @@ class Req(ReqDllmMixin):
|
||||
) # Before image padding
|
||||
# Each decode stage's output ids
|
||||
self.output_ids = array("q")
|
||||
# fill_ids = origin_input_ids + output_ids. Updated if chunked.
|
||||
self.fill_ids = array("q")
|
||||
# Full untruncated sequence: origin + output (+ DLLM mask block).
|
||||
# Rebuilt at the top of each init_next_round_input; admission only
|
||||
# updates fill_len, never mutates this array's length.
|
||||
self.full_untruncated_fill_ids = array("q")
|
||||
self.fill_len: int = 0
|
||||
|
||||
self.session = session
|
||||
self.input_embeds = input_embeds
|
||||
@@ -959,8 +962,8 @@ class Req(ReqDllmMixin):
|
||||
# the start index of the sent kv cache
|
||||
# We want to send it chunk by chunk for chunked prefill.
|
||||
# After every chunk forward, we do the following:
|
||||
# kv_send(req.input_ids[req.start_send_idx:len(req.fill_ids)])
|
||||
# start_send_idx = len(req.fill_ids)
|
||||
# kv_send(req.input_ids[req.start_send_idx:req.fill_len])
|
||||
# start_send_idx = req.fill_len
|
||||
self.start_send_idx: int = 0
|
||||
|
||||
# For overlap schedule, we delay the kv transfer until `process_batch_result_disagg_prefill` rather than `process_prefill_chunk` in non-overlap
|
||||
@@ -1062,6 +1065,9 @@ class Req(ReqDllmMixin):
|
||||
# Whether request reached finished condition
|
||||
return self.finished_reason is not None
|
||||
|
||||
def get_fill_ids(self) -> array:
|
||||
return self.full_untruncated_fill_ids[: self.fill_len]
|
||||
|
||||
def init_next_round_input(
|
||||
self,
|
||||
tree_cache: Optional[BasePrefixCache] = None,
|
||||
@@ -1071,9 +1077,10 @@ class Req(ReqDllmMixin):
|
||||
self._init_fill_ids_for_dllm()
|
||||
self.determine_dllm_phase()
|
||||
else:
|
||||
self.fill_ids = self.origin_input_ids + self.output_ids
|
||||
self.full_untruncated_fill_ids = self.origin_input_ids + self.output_ids
|
||||
self.fill_len = len(self.full_untruncated_fill_ids)
|
||||
|
||||
input_len = len(self.fill_ids)
|
||||
input_len = self.fill_len
|
||||
|
||||
# Streaming sessions reuse committed KV from the session slot, so
|
||||
# custom logprob_start_len is not supported — override to -1.
|
||||
@@ -1091,7 +1098,9 @@ class Req(ReqDllmMixin):
|
||||
)
|
||||
self.logprob_start_len = -1
|
||||
|
||||
token_ids_to_match = self.fill_ids[: self._compute_max_prefix_len(input_len)]
|
||||
token_ids_to_match = self.get_fill_ids()[
|
||||
: self._compute_max_prefix_len(input_len)
|
||||
]
|
||||
|
||||
# Disable prefix caching when embed overrides are present: same token IDs
|
||||
# with different override vectors must not share cached KV values.
|
||||
@@ -1154,7 +1163,7 @@ class Req(ReqDllmMixin):
|
||||
)
|
||||
)
|
||||
|
||||
self.set_extend_input_len(len(self.fill_ids) - len(self.prefix_indices))
|
||||
self.set_extend_input_len(self.fill_len - len(self.prefix_indices))
|
||||
|
||||
def _compute_max_prefix_len(self, input_len: int) -> int:
|
||||
# NOTE: the matched length is at most 1 less than the input length to enable logprob computation
|
||||
@@ -1420,7 +1429,7 @@ class Req(ReqDllmMixin):
|
||||
# - extend_input_len: Number of tokens that need to be processed in this extend batch
|
||||
self.extend_input_len = extend_input_len
|
||||
if self.logprob_start_len == -1:
|
||||
logprob_start_len = len(self.fill_ids)
|
||||
logprob_start_len = self.fill_len
|
||||
else:
|
||||
# logprob_start_len should be at least the length of the prefix indices
|
||||
logprob_start_len = max(self.logprob_start_len, len(self.prefix_indices))
|
||||
@@ -1544,7 +1553,7 @@ def _compute_chunked_req_next_prompt_token(
|
||||
) -> Optional[int]:
|
||||
if chunked_req is None:
|
||||
return None
|
||||
fill_len = len(chunked_req.fill_ids)
|
||||
fill_len = chunked_req.fill_len
|
||||
if fill_len >= len(chunked_req.origin_input_ids):
|
||||
return None
|
||||
return int(chunked_req.origin_input_ids[fill_len])
|
||||
@@ -1892,10 +1901,10 @@ class ScheduleBatch(ScheduleBatchDisaggregationDecodeMixin):
|
||||
|
||||
# Init tensors
|
||||
reqs = self.reqs
|
||||
input_ids = [r.fill_ids[len(r.prefix_indices) :] for r in reqs]
|
||||
input_ids = [r.get_fill_ids()[len(r.prefix_indices) :] for r in reqs]
|
||||
extend_num_tokens = sum(len(ids) for ids in input_ids)
|
||||
seq_lens = [len(r.fill_ids) for r in reqs]
|
||||
orig_seq_lens = [max(len(r.fill_ids), len(r.origin_input_ids)) for r in reqs]
|
||||
seq_lens = [r.fill_len for r in reqs]
|
||||
orig_seq_lens = [max(r.fill_len, len(r.origin_input_ids)) for r in reqs]
|
||||
prefix_lens = [len(r.prefix_indices) for r in reqs]
|
||||
extend_lens = [r.extend_input_len for r in reqs]
|
||||
|
||||
@@ -1948,7 +1957,7 @@ class ScheduleBatch(ScheduleBatchDisaggregationDecodeMixin):
|
||||
# If input_embeds are available, store them
|
||||
if req.input_embeds is not None:
|
||||
# Slice to match extend_input_len — PrefillAdder truncates
|
||||
# fill_ids/extend_input_len on chunk overflow but not input_embeds.
|
||||
# fill_len/extend_input_len on chunk overflow but not input_embeds.
|
||||
input_embeds.extend(
|
||||
req.input_embeds[pre_len : pre_len + req.extend_input_len]
|
||||
)
|
||||
@@ -2022,16 +2031,16 @@ class ScheduleBatch(ScheduleBatchDisaggregationDecodeMixin):
|
||||
# to compute input logprobs. E.g., (chunk size 2)
|
||||
#
|
||||
# input_logprobs = [1, 2, 3, 4]
|
||||
# fill_ids = [1, 2]
|
||||
# get_fill_ids() = [1, 2]
|
||||
# extend_input_logprob_token_id = [2, 3]
|
||||
#
|
||||
# Note that it can also overflow. In this case, we pad it with 0.
|
||||
# input_logprobs = [1, 2, 3, 4]
|
||||
# fill_ids = [3, 4]
|
||||
# get_fill_ids() = [3, 4]
|
||||
# extend_input_logprob_token_id = [4, 0]
|
||||
global_start_idx, global_end_idx = (
|
||||
len(req.prefix_indices),
|
||||
len(req.fill_ids),
|
||||
req.fill_len,
|
||||
)
|
||||
if req.logprob_start_len == -1:
|
||||
logprob_start_len = len(req.origin_input_ids)
|
||||
@@ -2261,7 +2270,8 @@ class ScheduleBatch(ScheduleBatchDisaggregationDecodeMixin):
|
||||
running_bs = running_batch.batch_size()
|
||||
|
||||
for req in running_batch.reqs:
|
||||
req.fill_ids = req.origin_input_ids + req.output_ids
|
||||
req.full_untruncated_fill_ids = req.origin_input_ids + req.output_ids
|
||||
req.fill_len = len(req.full_untruncated_fill_ids)
|
||||
req.set_extend_input_len(1)
|
||||
|
||||
# Decode tokens of the running portion live in future_map.output_tokens_buf.
|
||||
|
||||
@@ -660,7 +660,7 @@ class PrefillAdder:
|
||||
)
|
||||
|
||||
req.extend_input_len = trunc_len
|
||||
req.fill_ids = req.fill_ids[: prefix_len + trunc_len]
|
||||
req.fill_len = prefix_len + trunc_len
|
||||
|
||||
self.can_run_list.append(req)
|
||||
|
||||
@@ -681,7 +681,7 @@ class PrefillAdder:
|
||||
# Truncate input length to available tokens and update request metadata
|
||||
truncated = req.extend_input_len > _rem_tokens
|
||||
req.extend_input_len = min(req.extend_input_len, _rem_tokens)
|
||||
req.fill_ids = req.fill_ids[: len(req.prefix_indices) + req.extend_input_len]
|
||||
req.fill_len = len(req.prefix_indices) + req.extend_input_len
|
||||
self.can_run_list.append(req)
|
||||
|
||||
# Update budget: reserve max_new_tokens only if not truncated
|
||||
@@ -721,7 +721,7 @@ class PrefillAdder:
|
||||
|
||||
truncated = req.extend_input_len > _rem_tokens
|
||||
req.set_extend_input_len(min(req.extend_input_len, _rem_tokens))
|
||||
req.fill_ids = req.fill_ids[: len(req.prefix_indices) + req.extend_input_len]
|
||||
req.fill_len = len(req.prefix_indices) + req.extend_input_len
|
||||
self.can_run_list.append(req)
|
||||
self._update_prefill_budget(
|
||||
0,
|
||||
@@ -843,7 +843,7 @@ class PrefillAdder:
|
||||
trunc_len = self.rem_chunk_tokens
|
||||
|
||||
req.set_extend_input_len(trunc_len)
|
||||
req.fill_ids = req.fill_ids[:trunc_len]
|
||||
req.fill_len = trunc_len
|
||||
self.can_run_list.append(req)
|
||||
self.new_chunked_req = req
|
||||
self._update_prefill_budget(0, trunc_len, 0, req.retracted_stain)
|
||||
@@ -932,7 +932,7 @@ class PrefillAdder:
|
||||
)
|
||||
)
|
||||
req.prefix_indices = torch.cat([req.prefix_indices, new_indices])
|
||||
req.set_extend_input_len(len(req.fill_ids) - len(req.prefix_indices))
|
||||
req.set_extend_input_len(req.fill_len - len(req.prefix_indices))
|
||||
prefix_len = len(req.prefix_indices)
|
||||
req.cache_protected_len = prefix_len
|
||||
|
||||
@@ -999,7 +999,7 @@ class PrefillAdder:
|
||||
|
||||
# Chunked prefill
|
||||
req.set_extend_input_len(trunc_len)
|
||||
req.fill_ids = req.fill_ids[: len(req.prefix_indices) + trunc_len]
|
||||
req.fill_len = len(req.prefix_indices) + trunc_len
|
||||
|
||||
self.can_run_list.append(req)
|
||||
self.new_chunked_req = req
|
||||
|
||||
@@ -941,9 +941,9 @@ class Scheduler(
|
||||
# req_pool_idx from prepare_for_extend). Used to gate the
|
||||
# stash_chunked_request call at the top of get_next_batch_to_run:
|
||||
# if add_chunked_req early-returned under hybrid-SWA pressure,
|
||||
# the req_pool_idx was already freed and fill_ids was reset by
|
||||
# init_next_round_input, so running stash would double-free and
|
||||
# corrupt prefix_indices.
|
||||
# the req_pool_idx was already freed and the full_untruncated_fill_ids
|
||||
# was rebuilt by init_next_round_input, so running stash would
|
||||
# double-free and corrupt prefix_indices.
|
||||
self._chunked_req_scheduled_last_iter = False
|
||||
self.is_mixed_chunk = (
|
||||
self.chunked_prefill_size is not None
|
||||
@@ -2159,7 +2159,7 @@ class Scheduler(
|
||||
if last_host_node.backuped or last_host_node is self.tree_cache.root_node:
|
||||
last_hash = last_host_node.get_last_hash_value()
|
||||
matched_len = len(req.prefix_indices) + req.host_hit_length
|
||||
new_input_tokens = req.fill_ids[matched_len:]
|
||||
new_input_tokens = req.get_fill_ids()[matched_len:]
|
||||
|
||||
prefix_keys = (
|
||||
last_host_node.get_prefix_hash_values(last_host_node.parent)
|
||||
|
||||
@@ -184,7 +184,7 @@ class SchedulerLogprobResultProcessor:
|
||||
i: The request index in a batch.
|
||||
req: The request. Input logprobs inside req are modified as a
|
||||
consequence of the API
|
||||
fill_ids: The prefill ids processed.
|
||||
logprob_pt: Pointer into the prefill ids processed.
|
||||
output: Logit processor output that's used to compute input logprobs
|
||||
last_prefill_chunk: True if it is the last prefill (when chunked).
|
||||
Some of input logprob operation should only happen at the last
|
||||
|
||||
@@ -606,9 +606,10 @@ class SchedulerPPMixin:
|
||||
origin_input_ids=input_ids,
|
||||
sampling_params=sampling_params,
|
||||
)
|
||||
req.fill_ids = req.origin_input_ids
|
||||
req.full_untruncated_fill_ids = req.origin_input_ids
|
||||
req.fill_len = len(req.full_untruncated_fill_ids)
|
||||
req.logprob_start_len = -1
|
||||
req.set_extend_input_len(len(req.fill_ids) - len(req.prefix_indices))
|
||||
req.set_extend_input_len(req.fill_len - len(req.prefix_indices))
|
||||
|
||||
# Prepare batch
|
||||
batch = ScheduleBatch.init_new(
|
||||
@@ -621,7 +622,7 @@ class SchedulerPPMixin:
|
||||
self.spec_algorithm,
|
||||
)
|
||||
|
||||
current_seq_len = len(req.fill_ids)
|
||||
current_seq_len = req.fill_len
|
||||
|
||||
if is_dp_attention_enabled():
|
||||
# For profiling, we only have one request on PP0
|
||||
@@ -685,7 +686,7 @@ class SchedulerPPMixin:
|
||||
# Release KV cache
|
||||
if req.req_pool_idx is not None:
|
||||
kv_indices = self.req_to_token_pool.req_to_token[
|
||||
req.req_pool_idx, : len(req.fill_ids)
|
||||
req.req_pool_idx, : req.fill_len
|
||||
]
|
||||
self.token_to_kv_pool_allocator.free(kv_indices)
|
||||
self.req_to_token_pool.free(req)
|
||||
|
||||
@@ -86,7 +86,7 @@ class ChunkCache(BasePrefixCache):
|
||||
|
||||
def cache_unfinished_req(self, req: Req, chunked=False):
|
||||
kv_indices = self.req_to_token_pool.req_to_token[
|
||||
req.req_pool_idx, : len(req.fill_ids)
|
||||
req.req_pool_idx, : req.fill_len
|
||||
]
|
||||
# `req.prefix_indices` will be used in `PrefillAdder::add_chunked_req` later
|
||||
req.prefix_indices = kv_indices.to(dtype=torch.int64, copy=True)
|
||||
|
||||
@@ -611,14 +611,14 @@ class MambaRadixCache(KVCacheEventMixin, BasePrefixCache):
|
||||
|
||||
def _skip_cache_unfinished_req(req: Req) -> None:
|
||||
kv_indices = self.req_to_token_pool.req_to_token[
|
||||
req.req_pool_idx, : len(req.fill_ids)
|
||||
req.req_pool_idx, : req.fill_len
|
||||
]
|
||||
|
||||
# `req.prefix_indices` will be used in `PrefillAdder::add_chunked_req` later
|
||||
req.prefix_indices = kv_indices.to(dtype=torch.int64, copy=True)
|
||||
return
|
||||
|
||||
token_ids = req.fill_ids
|
||||
token_ids = req.get_fill_ids()
|
||||
cache_len = (
|
||||
req.mamba_last_track_seqlen
|
||||
if self.enable_mamba_extra_buffer
|
||||
|
||||
@@ -466,7 +466,7 @@ class RadixCache(KVCacheEventMixin, BasePrefixCache):
|
||||
if self.disable:
|
||||
return
|
||||
|
||||
token_ids = req.fill_ids
|
||||
token_ids = req.get_fill_ids()
|
||||
kv_indices = self.req_to_token_pool.req_to_token[
|
||||
req.req_pool_idx, : len(token_ids)
|
||||
]
|
||||
|
||||
@@ -210,7 +210,7 @@ class RadixCacheCpp(BasePrefixCache):
|
||||
def cache_unfinished_req(self, req: Req, chunked=False):
|
||||
"""Cache request when it is unfinished."""
|
||||
assert req.req_pool_idx is not None
|
||||
token_ids = req.fill_ids
|
||||
token_ids = req.get_fill_ids()
|
||||
prefill_len = len(token_ids) # prefill only (maybe chunked)
|
||||
kv_indices = self.req_to_token_pool.req_to_token[
|
||||
req.req_pool_idx, :prefill_len
|
||||
|
||||
@@ -488,14 +488,14 @@ class SWARadixCache(KVCacheEventMixin, BasePrefixCache):
|
||||
"""Cache request when it is unfinished."""
|
||||
if self.disable:
|
||||
kv_indices = self.req_to_token_pool.req_to_token[
|
||||
req.req_pool_idx, : len(req.fill_ids)
|
||||
req.req_pool_idx, : req.fill_len
|
||||
]
|
||||
|
||||
# `req.prefix_indices` will be used in `PrefillAdder::add_chunked_req` later
|
||||
req.prefix_indices = kv_indices
|
||||
return
|
||||
|
||||
token_ids = req.fill_ids
|
||||
token_ids = req.get_fill_ids()
|
||||
kv_indices = self.req_to_token_pool.req_to_token[
|
||||
req.req_pool_idx, : len(token_ids)
|
||||
]
|
||||
|
||||
@@ -722,7 +722,7 @@ class UnifiedRadixCache(KVCacheEventMixin, BasePrefixCache):
|
||||
if self.session.try_cache_unfinished_req(req, chunked=chunked, **kwargs):
|
||||
return
|
||||
|
||||
token_ids = req.fill_ids
|
||||
token_ids = req.get_fill_ids()
|
||||
|
||||
if self.disable:
|
||||
kv_indices = self.req_to_token_pool.req_to_token[
|
||||
|
||||
@@ -244,9 +244,9 @@ class StreamingSession(BasePrefixCache):
|
||||
req = params.req
|
||||
slot.restore_to_req(req)
|
||||
|
||||
# token_ids = fill_ids[:input_len-1] (1-token logit reserve already
|
||||
# applied). min handles retract retry where committed_len can
|
||||
# exceed len(token_ids) by 1.
|
||||
# token_ids = get_fill_ids()[:input_len-1] (1-token logit reserve
|
||||
# already applied). min handles retract retry where committed_len
|
||||
# can exceed len(token_ids) by 1.
|
||||
prefix_len = min(req.kv_committed_len, len(params.key.token_ids))
|
||||
|
||||
# Streaming sessions are append-only (session_controller rollback
|
||||
@@ -353,7 +353,7 @@ class StreamingSession(BasePrefixCache):
|
||||
return False
|
||||
if chunked:
|
||||
kv_indices = self.req_to_token_pool.req_to_token[
|
||||
req.req_pool_idx, : len(req.fill_ids)
|
||||
req.req_pool_idx, : req.fill_len
|
||||
]
|
||||
req.prefix_indices = kv_indices.to(dtype=torch.int64, copy=True)
|
||||
return True
|
||||
|
||||
Reference in New Issue
Block a user