Preserve decode suffix KV locations after cache hits
Decode queue compaction receives req_to_token rows after the prefill side has already populated cached prefix slots. Cache-hit requests therefore need the extend/suffix slice, not the leading prefix slice, when building the prebuilt transfer chunk. Constraint: Prefill/decode disaggregation shares req_to_token rows across cached prefix and new suffix positions. Rejected: Keep slicing from zero | cache-hit requests would copy prefix KV locs into the prebuilt suffix chunk. Confidence: medium Scope-risk: narrow Directive: Do not change prepare_for_prebuilt slicing without testing cache-hit req_to_token layouts. Tested: python -m py_compile on changed runtime files. Not-tested: Local pytest blocked before collection by missing orjson dependency. (cherry picked from commit 416112b617fabe71e8cff7484794af73f3e84440)
This commit is contained in:
@@ -50,7 +50,7 @@ class ScheduleBatchDisaggregationDecodeMixin:
|
||||
|
||||
pre_len = len(req.prefix_indices)
|
||||
chunk = self.req_to_token_pool.req_to_token[req.req_pool_idx][
|
||||
: req.extend_input_len
|
||||
pre_len : pre_len + req.extend_input_len
|
||||
]
|
||||
assert (
|
||||
offset + req.extend_input_len <= total_size
|
||||
@@ -96,7 +96,7 @@ class ScheduleBatchDisaggregationDecodeMixin:
|
||||
seq_len,
|
||||
int(getattr(req, "cached_tokens", 0) or 0),
|
||||
req.req_pool_idx,
|
||||
0,
|
||||
pre_len,
|
||||
pre_len,
|
||||
int(chunk.numel()),
|
||||
)
|
||||
|
||||
Reference in New Issue
Block a user