wudixzy stack docker build test: 2615-line qwen3_5.py + 12 prebuilt .so + 251-line patch_ops.sh
This commit is contained in:
project6-dev
291
qwen3_6_scripts/gdn_prefix.py
Normal file
291
qwen3_6_scripts/gdn_prefix.py
Normal file
@@ -0,0 +1,291 @@
|
||||
"""Shared GDN prefix-state cache contracts for the BI100 runtime."""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import os
|
||||
from collections import OrderedDict
|
||||
from dataclasses import dataclass
|
||||
from typing import Iterable, List, Optional, Sequence, Tuple
|
||||
|
||||
|
||||
GdnPrefixKey = Tuple[int, bytes]
|
||||
GdnCapturePoint = Tuple[int, GdnPrefixKey]
|
||||
|
||||
_VALID_POLICIES = {"fine32", "admission64", "off"}
|
||||
GDN_KERNEL_CHUNK_TOKENS = 64
|
||||
GDN_DIRECT_MIN_REPLAY_TOKENS = 2
|
||||
|
||||
_VALID_RESTORE_MODES = {"direct", "hybrid64", "chunk64", "aligned"}
|
||||
|
||||
|
||||
def _env_choice(name: str, default: str, choices: set[str]) -> str:
|
||||
value = os.getenv(name, default).strip().lower()
|
||||
if value not in choices:
|
||||
allowed = ", ".join(sorted(choices))
|
||||
raise RuntimeError(f"invalid {name}={value!r}; expected one of: {allowed}")
|
||||
return value
|
||||
|
||||
|
||||
def gdn_cache_policy_from_env() -> str:
|
||||
return _env_choice("BI100_GDN_CACHE_POLICY", "fine32", _VALID_POLICIES)
|
||||
|
||||
|
||||
def gdn_restore_mode_from_env() -> str:
|
||||
return _env_choice(
|
||||
"BI100_GDN_RESTORE_MODE", "direct", _VALID_RESTORE_MODES)
|
||||
|
||||
|
||||
def gdn_restore_alignment(restore_mode: str, block_size: int,
|
||||
scheduler_chunk_tokens: int) -> int:
|
||||
"""Return the content boundary required by a restore mode."""
|
||||
if block_size <= 0:
|
||||
raise ValueError("block_size must be positive")
|
||||
if restore_mode == "direct":
|
||||
return block_size
|
||||
if restore_mode in {"hybrid64", "chunk64"}:
|
||||
alignment = GDN_KERNEL_CHUNK_TOKENS
|
||||
elif restore_mode == "aligned":
|
||||
alignment = scheduler_chunk_tokens
|
||||
else:
|
||||
raise ValueError(f"unknown GDN restore mode: {restore_mode}")
|
||||
if alignment <= 0 or alignment % block_size != 0:
|
||||
raise ValueError(
|
||||
f"{restore_mode} GDN restore requires a positive alignment "
|
||||
f"divisible by block_size={block_size}; got {alignment}")
|
||||
return alignment
|
||||
|
||||
|
||||
def make_prefix_key(block_count: int, digest: bytes) -> GdnPrefixKey:
|
||||
if block_count <= 0:
|
||||
raise ValueError("GDN prefix key requires at least one complete block")
|
||||
if not isinstance(digest, bytes) or len(digest) != 32:
|
||||
raise ValueError("GDN prefix digest must be exactly 32 bytes")
|
||||
return block_count, digest
|
||||
|
||||
|
||||
def keys_from_block_hashes(block_hashes: Sequence[bytes]) -> List[GdnPrefixKey]:
|
||||
return [make_prefix_key(i + 1, digest)
|
||||
for i, digest in enumerate(block_hashes)]
|
||||
|
||||
|
||||
def strict_prefix_block_count(token_count: int, block_size: int) -> int:
|
||||
if block_size <= 0:
|
||||
raise ValueError("block_size must be positive")
|
||||
if token_count <= 1:
|
||||
return 0
|
||||
return (token_count - 1) // block_size
|
||||
|
||||
|
||||
def key_at_strict_boundary(block_hashes: Sequence[bytes], token_count: int,
|
||||
block_size: int) -> Optional[GdnPrefixKey]:
|
||||
block_count = min(
|
||||
len(block_hashes), strict_prefix_block_count(token_count, block_size))
|
||||
if block_count <= 0:
|
||||
return None
|
||||
return make_prefix_key(block_count, block_hashes[block_count - 1])
|
||||
|
||||
|
||||
def final_capture_key(
|
||||
block_hashes: Sequence[bytes], prompt_tokens: int, block_size: int,
|
||||
restore_mode: str, replay_alignment: int) -> Optional[GdnPrefixKey]:
|
||||
if restore_mode in {"direct", "hybrid64"}:
|
||||
block_count = min(
|
||||
len(block_hashes), strict_prefix_block_count(
|
||||
prompt_tokens, block_size))
|
||||
if (block_count > 0
|
||||
and prompt_tokens - block_count * block_size
|
||||
< GDN_DIRECT_MIN_REPLAY_TOKENS):
|
||||
block_count -= 1
|
||||
if block_count <= 0:
|
||||
return None
|
||||
return make_prefix_key(block_count, block_hashes[block_count - 1])
|
||||
if restore_mode not in {"chunk64", "aligned"}:
|
||||
raise ValueError(f"unknown GDN restore mode: {restore_mode}")
|
||||
if (replay_alignment <= 0 or replay_alignment % block_size != 0
|
||||
or prompt_tokens <= 1):
|
||||
return None
|
||||
boundary_tokens = ((prompt_tokens - 1) // replay_alignment
|
||||
* replay_alignment)
|
||||
block_count = min(len(block_hashes), boundary_tokens // block_size)
|
||||
if block_count <= 0:
|
||||
return None
|
||||
return make_prefix_key(block_count, block_hashes[block_count - 1])
|
||||
|
||||
|
||||
def restore_key_is_eligible(
|
||||
key: GdnPrefixKey, prompt_tokens: int, block_size: int,
|
||||
restore_mode: str, replay_alignment: int,
|
||||
direct_final_key: Optional[GdnPrefixKey] = None) -> bool:
|
||||
"""Return whether restoring ``key`` preserves the execution contract."""
|
||||
make_prefix_key(*key)
|
||||
if block_size <= 0:
|
||||
raise ValueError("block_size must be positive")
|
||||
boundary_tokens = key[0] * block_size
|
||||
remaining_tokens = prompt_tokens - boundary_tokens
|
||||
if remaining_tokens <= 0:
|
||||
return False
|
||||
if restore_mode == "direct":
|
||||
return remaining_tokens >= GDN_DIRECT_MIN_REPLAY_TOKENS
|
||||
if restore_mode == "hybrid64":
|
||||
if direct_final_key is not None:
|
||||
make_prefix_key(*direct_final_key)
|
||||
return (remaining_tokens >= GDN_DIRECT_MIN_REPLAY_TOKENS
|
||||
and replay_alignment > 0
|
||||
and (boundary_tokens % replay_alignment == 0
|
||||
or key == direct_final_key))
|
||||
if restore_mode not in {"chunk64", "aligned"}:
|
||||
raise ValueError(f"unknown GDN restore mode: {restore_mode}")
|
||||
return (replay_alignment > 0
|
||||
and boundary_tokens % replay_alignment == 0)
|
||||
|
||||
|
||||
def capture_points_for_step(
|
||||
targets: Iterable[GdnPrefixKey], physical_context_tokens: int,
|
||||
logical_end_tokens: int, block_size: int) -> Tuple[GdnCapturePoint, ...]:
|
||||
if physical_context_tokens < 0 or logical_end_tokens < 0:
|
||||
raise ValueError("token positions must be non-negative")
|
||||
if logical_end_tokens <= physical_context_tokens:
|
||||
return ()
|
||||
selected = {}
|
||||
for key in targets:
|
||||
make_prefix_key(*key)
|
||||
boundary_tokens = key[0] * block_size
|
||||
if physical_context_tokens < boundary_tokens <= logical_end_tokens:
|
||||
selected[boundary_tokens - physical_context_tokens] = key
|
||||
points = tuple(sorted(selected.items()))
|
||||
if len(points) > 2:
|
||||
raise ValueError("at most two GDN capture points are allowed per step")
|
||||
return points
|
||||
|
||||
|
||||
def cap_prefill_end_at_capture_boundary(
|
||||
logical_start_tokens: int, logical_end_tokens: int,
|
||||
targets: Iterable[GdnPrefixKey], block_size: int) -> int:
|
||||
"""Stop a physical prefill step at its earliest pending capture boundary."""
|
||||
if logical_start_tokens < 0 or logical_end_tokens < 0:
|
||||
raise ValueError("token positions must be non-negative")
|
||||
if logical_end_tokens < logical_start_tokens:
|
||||
raise ValueError("logical end must not precede logical start")
|
||||
if block_size <= 0:
|
||||
raise ValueError("block_size must be positive")
|
||||
|
||||
capped_end = logical_end_tokens
|
||||
for key in targets:
|
||||
make_prefix_key(*key)
|
||||
boundary_tokens = key[0] * block_size
|
||||
if logical_start_tokens < boundary_tokens < capped_end:
|
||||
capped_end = boundary_tokens
|
||||
return capped_end
|
||||
|
||||
|
||||
def canonical_direct_segment_offsets(
|
||||
block_hashes: Sequence[bytes], physical_context_tokens: int,
|
||||
logical_end_tokens: int, block_size: int,
|
||||
scheduler_chunk_tokens: int) -> Tuple[int, ...]:
|
||||
"""Reproduce cold fine32/direct segment boundaries after fast-forward."""
|
||||
if physical_context_tokens < 0 or logical_end_tokens < 0:
|
||||
raise ValueError("token positions must be non-negative")
|
||||
if block_size <= 0 or scheduler_chunk_tokens <= 0:
|
||||
raise ValueError("block and scheduler chunk sizes must be positive")
|
||||
if scheduler_chunk_tokens % block_size != 0:
|
||||
raise ValueError("scheduler chunk size must be divisible by block size")
|
||||
if logical_end_tokens <= physical_context_tokens:
|
||||
return ()
|
||||
|
||||
boundaries = set()
|
||||
step_ends = list(range(scheduler_chunk_tokens, logical_end_tokens,
|
||||
scheduler_chunk_tokens))
|
||||
for step_end in (*step_ends, logical_end_tokens):
|
||||
key = final_capture_key(block_hashes, step_end, block_size,
|
||||
"direct", block_size)
|
||||
if key is not None:
|
||||
boundaries.add(key[0] * block_size)
|
||||
boundaries.update(step_ends)
|
||||
return tuple(
|
||||
boundary - physical_context_tokens
|
||||
for boundary in sorted(boundaries)
|
||||
if physical_context_tokens < boundary < logical_end_tokens)
|
||||
|
||||
|
||||
@dataclass(frozen=True)
|
||||
class GdnCachePlan:
|
||||
restore_key: Optional[GdnPrefixKey] = None
|
||||
capture_points: Tuple[GdnCapturePoint, ...] = ()
|
||||
evict_keys: Tuple[GdnPrefixKey, ...] = ()
|
||||
|
||||
|
||||
class GdnPrefixStatePolicy:
|
||||
"""Scheduler-owned state index with deterministic worker actions."""
|
||||
|
||||
def __init__(self, policy: str) -> None:
|
||||
if policy not in _VALID_POLICIES:
|
||||
raise ValueError(f"unknown GDN cache policy: {policy}")
|
||||
self.policy = policy
|
||||
self.capacity = {"fine32": 32, "admission64": 64, "off": 0}[policy]
|
||||
self._resident: OrderedDict[GdnPrefixKey, None] = OrderedDict()
|
||||
|
||||
def __len__(self) -> int:
|
||||
return len(self._resident)
|
||||
|
||||
def resident_keys(self) -> Tuple[GdnPrefixKey, ...]:
|
||||
return tuple(self._resident)
|
||||
|
||||
def contains(self, key: GdnPrefixKey) -> bool:
|
||||
return key in self._resident
|
||||
|
||||
def should_capture_final(self, key: GdnPrefixKey) -> bool:
|
||||
"""Return whether a final state must be materialized on this request."""
|
||||
make_prefix_key(*key)
|
||||
if self.policy == "off":
|
||||
return False
|
||||
if self.policy == "admission64":
|
||||
return key not in self._resident
|
||||
return True
|
||||
|
||||
def select_restore(
|
||||
self, live_prefix_keys: Sequence[GdnPrefixKey],
|
||||
max_blocks: int) -> Optional[GdnPrefixKey]:
|
||||
if self.capacity == 0 or max_blocks <= 0:
|
||||
return None
|
||||
best = None
|
||||
for key in live_prefix_keys[:max_blocks]:
|
||||
if key in self._resident:
|
||||
best = key
|
||||
if best is not None:
|
||||
self._resident.move_to_end(best)
|
||||
return best
|
||||
|
||||
def repeated_branch_candidate(
|
||||
self, live_prefix_keys: Sequence[GdnPrefixKey],
|
||||
max_blocks: int) -> Optional[GdnPrefixKey]:
|
||||
"""Return a repeated raw-KV branch that lacks recurrent state.
|
||||
|
||||
A live KV hit proves that the content occurred in an earlier request;
|
||||
the current request is therefore the second or later occurrence.
|
||||
"""
|
||||
if (self.policy != "admission64" or max_blocks <= 0
|
||||
or not live_prefix_keys):
|
||||
return None
|
||||
candidate = live_prefix_keys[min(len(live_prefix_keys), max_blocks) - 1]
|
||||
if candidate in self._resident:
|
||||
return None
|
||||
return candidate
|
||||
|
||||
def admit(self, keys: Iterable[GdnPrefixKey]) -> Tuple[GdnPrefixKey, ...]:
|
||||
evicted: List[GdnPrefixKey] = []
|
||||
if self.capacity == 0:
|
||||
return ()
|
||||
for key in keys:
|
||||
make_prefix_key(*key)
|
||||
if key in self._resident:
|
||||
self._resident.move_to_end(key)
|
||||
else:
|
||||
self._resident[key] = None
|
||||
while len(self._resident) > self.capacity:
|
||||
evicted_key, _ = self._resident.popitem(last=False)
|
||||
evicted.append(evicted_key)
|
||||
return tuple(evicted)
|
||||
|
||||
def forget(self, keys: Iterable[GdnPrefixKey]) -> None:
|
||||
for key in keys:
|
||||
self._resident.pop(key, None)
|
||||
Reference in New Issue
Block a user