glm53-spark-tp4 / kv_cache /prefix_cache.py

GLM-5.3 EXL3 3.38bpw on 4x DGX Spark: serving kit + incremental KV prefix cache

dfc39f1 verified 12 days ago

3.47 kB

"""The prompt-prefix record for the GLM lane's full prefix cache (audit 253).
WHY THIS IS NOT `cuda/streams.PrefixCache`: that class keeps a LIST of `(ids, state, snap)` entries, i.e. one
whole KV state per entry. On this lane the KV is ONE shared full-capacity `fused.State` (~15-20 GB at 1M) that
`Runner` allocates once and every request writes into, and the lane serves one stream at a time with DCP=4 --
so keeping several states is neither possible nor needed. What can be kept is a small host-side RECORD of what
the shared State currently holds: the committed ids and its length (tens of KB, not GB).
The offset rule mirrors `streams.PrefixCache.longest`: reuse only a STRICT prefix and always leave at least one
token to prefill (the engine needs a row to produce a logit). Anything else -- a divergence, a shorter prompt, a
new session -- leaves the offset at 0, so a stale KV is never read as if it were the current context.
"""
from __future__ import annotations

import os

# DEFAULT OFF (audit 253): the resume path is NOT yet bit-identical in every case. Measured: reuse with a
# 1-token suffix matched cold exactly, but reuse with a 3-token suffix did not -- and BOTH draft modes
# agreed within each case, so the divergence is in the TARGET path, not the drafter. Reuse is proven
# (cached=N, up to ~113x faster on a repeated prefix) but correctness is not. Opt in with
# TF_GLM53_PREFIX_CACHE=1 to keep testing; the served lane stays bit-exact until this is fixed.
TF_PREFIX_CACHE = os.environ.get("TF_GLM53_PREFIX_CACHE", "1") != "0"   # audit 259: ON by default
KEEP_MIN = int(os.environ.get("TF_GLM53_PREFIX_CACHE_MIN", "64"))   # below this a reuse saves too little to be worth it


def common_prefix(a: list, b: list) -> int:
    """How many leading tokens ``a`` and ``b`` share."""
    n = min(len(a), len(b))
    i = 0
    while i < n and a[i] == b[i]:
        i += 1
    return i


class PromptPrefix:
    """What the shared State holds: the committed ids and their count, from the last request to fill it."""

    def __init__(self) -> None:
        self.ids: list = []
        self.canonical_end: int = 0        # audit 260: highest row a CANONICAL PREFILL chunk wrote
        self.enabled = TF_PREFIX_CACHE

    def offset_for(self, prompt: list) -> int:
        """The number of leading prompt tokens already valid in the State; 0 when nothing may be reused.
        Rules: reuse only when the prompt STRICTLY extends what is held (so the held KV is a true prefix) and the
        reused length is at least KEEP_MIN; never reuse the whole prompt (one token is always left to prefill).
        """
        if not self.enabled or not self.ids:
            return 0
        n = common_prefix(self.ids, prompt)
        # audit 260: never offer a row decode wrote. canonical_end is the previous request's prompt
        # length -- the highest position a full prefill chunk produced.
        n = min(n, self.canonical_end, len(prompt) - 1)
        if n < KEEP_MIN:
            return 0
        return n

    def record(self, ids: list, canonical_end: int) -> None:
        """Publish the record. canonical_end = this request PROMPT length, the highest row a canonical
        prefill chunk wrote. Generated tokens past it are decode-written and must never be reused."""
        if self.enabled:
            self.ids = list(ids)
            self.canonical_end = max(0, int(canonical_end))

    def clear(self) -> None:
        self.ids = []

	"""The prompt-prefix record for the GLM lane's full prefix cache (audit 253).

	WHY THIS IS NOT `cuda/streams.PrefixCache`: that class keeps a LIST of `(ids, state, snap)` entries, i.e. one
	whole KV state per entry. On this lane the KV is ONE shared full-capacity `fused.State` (~15-20 GB at 1M) that
	`Runner` allocates once and every request writes into, and the lane serves one stream at a time with DCP=4 --
	so keeping several states is neither possible nor needed. What can be kept is a small host-side RECORD of what
	the shared State currently holds: the committed ids and its length (tens of KB, not GB).

	The offset rule mirrors `streams.PrefixCache.longest`: reuse only a STRICT prefix and always leave at least one
	token to prefill (the engine needs a row to produce a logit). Anything else -- a divergence, a shorter prompt, a
	new session -- leaves the offset at 0, so a stale KV is never read as if it were the current context.
	"""
	from __future__ import annotations

	import os

	# DEFAULT OFF (audit 253): the resume path is NOT yet bit-identical in every case. Measured: reuse with a
	# 1-token suffix matched cold exactly, but reuse with a 3-token suffix did not -- and BOTH draft modes
	# agreed within each case, so the divergence is in the TARGET path, not the drafter. Reuse is proven
	# (cached=N, up to ~113x faster on a repeated prefix) but correctness is not. Opt in with
	# TF_GLM53_PREFIX_CACHE=1 to keep testing; the served lane stays bit-exact until this is fixed.
	TF_PREFIX_CACHE = os.environ.get("TF_GLM53_PREFIX_CACHE", "1") != "0" # audit 259: ON by default
	KEEP_MIN = int(os.environ.get("TF_GLM53_PREFIX_CACHE_MIN", "64")) # below this a reuse saves too little to be worth it


	def common_prefix(a: list, b: list) -> int:
	"""How many leading tokens ``a`` and ``b`` share."""
	n = min(len(a), len(b))
	i = 0
	while i < n and a[i] == b[i]:
	i += 1
	return i


	class PromptPrefix:
	"""What the shared State holds: the committed ids and their count, from the last request to fill it."""

	def __init__(self) -> None:
	self.ids: list = []
	self.canonical_end: int = 0 # audit 260: highest row a CANONICAL PREFILL chunk wrote
	self.enabled = TF_PREFIX_CACHE

	def offset_for(self, prompt: list) -> int:
	"""The number of leading prompt tokens already valid in the State; 0 when nothing may be reused.

	Rules: reuse only when the prompt STRICTLY extends what is held (so the held KV is a true prefix) and the
	reused length is at least KEEP_MIN; never reuse the whole prompt (one token is always left to prefill).
	"""
	if not self.enabled or not self.ids:
	return 0
	n = common_prefix(self.ids, prompt)
	# audit 260: never offer a row decode wrote. canonical_end is the previous request's prompt
	# length -- the highest position a full prefill chunk produced.
	n = min(n, self.canonical_end, len(prompt) - 1)
	if n < KEEP_MIN:
	return 0
	return n

	def record(self, ids: list, canonical_end: int) -> None:
	"""Publish the record. canonical_end = this request PROMPT length, the highest row a canonical
	prefill chunk wrote. Generated tokens past it are decode-written and must never be reused."""
	if self.enabled:
	self.ids = list(ids)
	self.canonical_end = max(0, int(canonical_end))

	def clear(self) -> None:
	self.ids = []