Restore each indexer's K cache on incomplete compression steps
I made each K-owning indexer select its own cache on every call.
With ratio-2 compression, the first token in a group produces latent=None. After an even-length prefix, the encoder's next indexer call reads the decoder's index K left in the shared slot by the preceding forward. This affects the source layers at 2, 8, and 14 in the released configuration. I moved cache selection into its own if self.owns_k block. Cache writes keep their existing latent is not None condition, and Reindex consumers keep reading their source's cache.
I reproduced the selection with the official Compressor.forward and Indexer.forward extracted from fb2764a and executed on CPU. The original selected [1]; the encoder's own keys select [0]. This change selected [0]. I made RoPE and FP4 quantization identity operations in that check to isolate cache selection. GPU kernels and full-model output quality remain outside this check's scope.
Dependency-free cache-owner check
This smaller check extracts Indexer.forward, stops at query projection, and checks the active cache reference. It uses Python's standard library. Save it as check_indexer_owner.py and run python check_indexer_owner.py inference/model.py.
"""Inspect Indexer.forward cache selection before query projection.
Run with an inference/model.py path. The probe stops before tensor operations.
"""
import ast
import sys
from pathlib import Path
from types import SimpleNamespace
source = ast.parse(Path(sys.argv[1]).read_text())
indexer = next(n for n in source.body if isinstance(n, ast.ClassDef) and n.name == "Indexer")
forward = next(n for n in indexer.body if isinstance(n, ast.FunctionDef) and n.name == "forward")
shared = SimpleNamespace(index_k=None)
namespace = {"torch": SimpleNamespace(Tensor=object), "shared_attn": shared}
exec(compile(ast.Module(body=[forward], type_ignores=[]), "Indexer.forward", "exec"), namespace)
class QueryProjectionReached(Exception):
pass
def stop_at_query(qr):
raise QueryProjectionReached
def check(owns_k, active_cache, expected_cache):
indexer = SimpleNamespace(
owns_k=owns_k,
compress_ratio=2,
rope_head_dim=64,
freqs_cis=object(),
wq_b=stop_at_query,
)
if owns_k:
indexer.k_cache = expected_cache
shared.index_k = active_cache
try:
namespace["forward"](indexer, SimpleNamespace(size=lambda: (1, 1, 128)), None, None, 4, 0)
except QueryProjectionReached:
return shared.index_k is expected_cache
raise AssertionError("Query projection was not reached")
decoder_cache = object()
encoder_caches = [object() for _ in range(3)]
active_cache = decoder_cache
for layer, cache in zip((2, 8, 14), encoder_caches):
selected = check(True, active_cache, cache)
print(f"layer {layer}: own cache selected = {selected}")
assert selected, f"layer {layer} retained the previous source's index K"
assert check(False, cache, cache), "Reindex consumer changed the shared cache"
active_cache = cache
print("Incomplete-group owners and Reindex consumers selected the expected caches.")
The original fails at layer 2 with own cache selected = False. The patched code selects each encoder source's cache and preserves the cache for Reindex consumers.
I extended the H100 check through the official Attention.forward calls for Full, Reindex and Reuse, using initialized synthetic weights and fixed per-layer hidden inputs to isolate cache scheduling. It runs PyTorch 2.13.0, CUDA 13.0, TileLang 0.1.8 and TVM FFI 0.1.6, including the actual FP8 window quantization, RoPE, FP4 quantization, sparse attention and output projections.
After the decoder publishes its keys, the encoder's incomplete compression step selects position 1 on current main and changes all three roles' projected outputs; this patch selects position 0 and matches the isolated encoder's indices and outputs bit for bit. The check also covers sliding-window wraparound, preservation of the encoder's cached values and the next complete compression group, where both versions match the isolated run.
Full-checkpoint output quality remains a separate evaluation.