8030a11b96
project_7 docker build 已在竞赛平台验证成功。 完整搬运 wudixzy/competition stack: - qwen3_5.py 2615 行 (12 个 corex .so 调用) - patch_ops.sh 251 行 (set -eo pipefail + cd dirname) - 12 prebuilt corex .so (SHA256 verified) - 13 CUDA .cu 源码 + 11 build scripts - 9 vendor overrides (block/sampler/scheduler) - transformers-4.55.3 offline wheel - computility-run.yaml: 262144 max-model-len, BI100 env vars - Dockerfile 结构不变 (COPY qwen3_6_scripts + RUN patch_ops.sh)
292 lines
11 KiB
Python
292 lines
11 KiB
Python
"""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)
|