Exclude partial prompt blocks from prefix cache replay
This commit is contained in:
@@ -275,6 +275,13 @@ def materialize(args: argparse.Namespace, *, tokenizer: Any | None = None) -> di
|
|||||||
identity_collisions += int(previous != witness)
|
identity_collisions += int(previous != witness)
|
||||||
prompt_kind = "synthetic_missing_prompt_fallback"
|
prompt_kind = "synthetic_missing_prompt_fallback"
|
||||||
synthetic_fallback_requests += 1
|
synthetic_fallback_requests += 1
|
||||||
|
# vLLM only inserts complete physical blocks into the prefix cache.
|
||||||
|
# Keep the final partial block identity in selected-remapped.jsonl
|
||||||
|
# for exact token/block validation, but never expose it as a
|
||||||
|
# cacheable Frontier block. Otherwise an exact repeated prompt can
|
||||||
|
# report ceil(ISL / 16) cached blocks, exceeding ISL and leaving the
|
||||||
|
# request permanently unschedulable.
|
||||||
|
cacheable_block_ids = block_ids[: isl // TARGET_BLOCK_SIZE]
|
||||||
parent = row.get("parent_chat_id")
|
parent = row.get("parent_chat_id")
|
||||||
if args.validate_parents and parent not in (None, "", -1, "-1"):
|
if args.validate_parents and parent not in (None, "", -1, "-1"):
|
||||||
parent_blocks = parent_sequences.get(parent)
|
parent_blocks = parent_sequences.get(parent)
|
||||||
@@ -290,7 +297,7 @@ def materialize(args: argparse.Namespace, *, tokenizer: Any | None = None) -> di
|
|||||||
# unaligned block. So we require a contiguous common prefix and
|
# unaligned block. So we require a contiguous common prefix and
|
||||||
# only tolerate a bounded mismatch confined to the parent tail.
|
# only tolerate a bounded mismatch confined to the parent tail.
|
||||||
common = 0
|
common = 0
|
||||||
for a, b in zip(block_ids, parent_blocks):
|
for a, b in zip(cacheable_block_ids, parent_blocks):
|
||||||
if a != b:
|
if a != b:
|
||||||
break
|
break
|
||||||
common += 1
|
common += 1
|
||||||
@@ -305,7 +312,7 @@ def materialize(args: argparse.Namespace, *, tokenizer: Any | None = None) -> di
|
|||||||
)
|
)
|
||||||
parent_prefix_common_blocks += common
|
parent_prefix_common_blocks += common
|
||||||
parent_tail_nonreused_blocks += len(parent_blocks) - common
|
parent_tail_nonreused_blocks += len(parent_blocks) - common
|
||||||
parent_sequences[row_identity(row, source_index)] = block_ids
|
parent_sequences[row_identity(row, source_index)] = cacheable_block_ids
|
||||||
|
|
||||||
if first_selected_timestamp is None:
|
if first_selected_timestamp is None:
|
||||||
first_selected_timestamp = ts
|
first_selected_timestamp = ts
|
||||||
@@ -318,7 +325,9 @@ def materialize(args: argparse.Namespace, *, tokenizer: Any | None = None) -> di
|
|||||||
"num_prefill_tokens": isl,
|
"num_prefill_tokens": isl,
|
||||||
"num_decode_tokens": osl,
|
"num_decode_tokens": osl,
|
||||||
"session_id": sid,
|
"session_id": sid,
|
||||||
"block_hash_ids": "|".join(str(value) for value in block_ids),
|
"block_hash_ids": "|".join(
|
||||||
|
str(value) for value in cacheable_block_ids
|
||||||
|
),
|
||||||
}
|
}
|
||||||
)
|
)
|
||||||
row_vector_digest.update(
|
row_vector_digest.update(
|
||||||
@@ -329,7 +338,7 @@ def materialize(args: argparse.Namespace, *, tokenizer: Any | None = None) -> di
|
|||||||
isl,
|
isl,
|
||||||
osl,
|
osl,
|
||||||
sid,
|
sid,
|
||||||
block_ids,
|
cacheable_block_ids,
|
||||||
],
|
],
|
||||||
separators=(",", ":"),
|
separators=(",", ":"),
|
||||||
).encode()
|
).encode()
|
||||||
@@ -343,7 +352,7 @@ def materialize(args: argparse.Namespace, *, tokenizer: Any | None = None) -> di
|
|||||||
"input_length": isl,
|
"input_length": isl,
|
||||||
"output_length": osl,
|
"output_length": osl,
|
||||||
"session_id": sid,
|
"session_id": sid,
|
||||||
"runtime_block_ids": block_ids,
|
"runtime_block_ids": cacheable_block_ids,
|
||||||
"body": {
|
"body": {
|
||||||
"model": args.served_model,
|
"model": args.served_model,
|
||||||
"prompt": request_prompt,
|
"prompt": request_prompt,
|
||||||
@@ -375,7 +384,7 @@ def materialize(args: argparse.Namespace, *, tokenizer: Any | None = None) -> di
|
|||||||
prefill_tokens += isl
|
prefill_tokens += isl
|
||||||
decode_tokens += osl
|
decode_tokens += osl
|
||||||
source_blocks += len(original_hashes)
|
source_blocks += len(original_hashes)
|
||||||
target_blocks += len(block_ids)
|
target_blocks += len(cacheable_block_ids)
|
||||||
|
|
||||||
if request_count == 0:
|
if request_count == 0:
|
||||||
raise ValueError("rho/window selected no requests")
|
raise ValueError("rho/window selected no requests")
|
||||||
@@ -419,6 +428,7 @@ def materialize(args: argparse.Namespace, *, tokenizer: Any | None = None) -> di
|
|||||||
"paired_row_vector_sha256": row_vector_digest.hexdigest(),
|
"paired_row_vector_sha256": row_vector_digest.hexdigest(),
|
||||||
"source_block_size": source_block_size,
|
"source_block_size": source_block_size,
|
||||||
"target_block_size": TARGET_BLOCK_SIZE,
|
"target_block_size": TARGET_BLOCK_SIZE,
|
||||||
|
"prefix_cache_blocks": "complete 16-token blocks only; final partial block excluded",
|
||||||
"workload_mode": workload_mode,
|
"workload_mode": workload_mode,
|
||||||
"mapping": (
|
"mapping": (
|
||||||
"real prompt tokens: BLAKE2b-128(parent runtime identity, exact 16-token block); "
|
"real prompt tokens: BLAKE2b-128(parent runtime identity, exact 16-token block); "
|
||||||
|
|||||||
@@ -2,8 +2,9 @@
|
|||||||
|
|
||||||
from __future__ import annotations
|
from __future__ import annotations
|
||||||
|
|
||||||
import math
|
import csv
|
||||||
import json
|
import json
|
||||||
|
import math
|
||||||
import sys
|
import sys
|
||||||
import tempfile
|
import tempfile
|
||||||
import unittest
|
import unittest
|
||||||
@@ -124,15 +125,21 @@ class StrictRemapTest(unittest.TestCase):
|
|||||||
manifest = materialize(args)
|
manifest = materialize(args)
|
||||||
request = json.loads((output / "real_requests.jsonl").read_text())
|
request = json.loads((output / "real_requests.jsonl").read_text())
|
||||||
mapped = json.loads((output / "selected-remapped.jsonl").read_text())
|
mapped = json.loads((output / "selected-remapped.jsonl").read_text())
|
||||||
|
with (output / "frontier.csv").open(newline="") as stream:
|
||||||
|
frontier_row = next(csv.DictReader(stream))
|
||||||
|
|
||||||
self.assertEqual(manifest["source_block_size"], 512)
|
self.assertEqual(manifest["source_block_size"], 512)
|
||||||
self.assertEqual(manifest["workload_mode"], "prefill_only")
|
self.assertEqual(manifest["workload_mode"], "prefill_only")
|
||||||
self.assertEqual(manifest["target_16_blocks"], math.ceil(513 / 16))
|
self.assertEqual(manifest["target_16_blocks"], math.floor(513 / 16))
|
||||||
self.assertEqual(request["input_length"], 513)
|
self.assertEqual(request["input_length"], 513)
|
||||||
self.assertEqual(request["output_length"], 1)
|
self.assertEqual(request["output_length"], 1)
|
||||||
self.assertEqual(request["body"]["max_tokens"], 1)
|
self.assertEqual(request["body"]["max_tokens"], 1)
|
||||||
|
self.assertEqual(len(request["runtime_block_ids"]), math.floor(513 / 16))
|
||||||
self.assertEqual(mapped["output_length"], 1)
|
self.assertEqual(mapped["output_length"], 1)
|
||||||
self.assertEqual(len(mapped["hash_ids_16"]), math.ceil(513 / 16))
|
self.assertEqual(len(mapped["hash_ids_16"]), math.ceil(513 / 16))
|
||||||
|
self.assertEqual(
|
||||||
|
len(frontier_row["block_hash_ids"].split("|")), math.floor(513 / 16)
|
||||||
|
)
|
||||||
|
|
||||||
def test_real_prompt_tokens_preserve_parent_and_four_to_one_contract(self) -> None:
|
def test_real_prompt_tokens_preserve_parent_and_four_to_one_contract(self) -> None:
|
||||||
tokenizer = CharacterTokenizer()
|
tokenizer = CharacterTokenizer()
|
||||||
|
|||||||
Reference in New Issue
Block a user