Harden IndexCache for first e2e: validation, fail-fast, prefetch gating

Pre-e2e audit fixes for NSA index-topk sharing (index_topk_pattern):

- Validate the pattern at init: F/S charset only, must start with F
  (layer 0 has nothing to share from), and length must equal
  num_hidden_layers — a short pattern previously crashed deep in pool
  init with an obscure per-layer ValueError, a long one silently
  ignored tail characters. Warn when a pattern shadows a configured
  index_topk_freq (pattern takes precedence silently otherwise).
- Fail fast when a shared (S) layer receives prev_topk_indices=None:
  both gated indexer call sites previously fell through to omitting
  topk_indices entirely, running sparse attention without indices and
  silently corrupting output if threading were ever dropped.
- Reject pipeline parallelism with index-topk sharing (same rationale
  as the existing TBO guard): topk indices are not threaded across
  PP-stage boundaries, and each stage's loop resets them to None.
- Gate the CP shared-KV index prefetcher on the next layer having an
  index-cache slot: prefetching an S layer's nonexistent buffer
  tripped the inactive-layer fail-fast and disabled the prefetcher
  for the entire run on the first F->S transition.
- Log the resolved active index layer plan at pool init so an e2e run
  can confirm IndexCache is actually on.

Tests: pattern validation cases incl. the GLM-5 reference 78-char
pattern, next_skip == skip-of-next invariant; existing 'C' pattern
test updated to upstream 'F' charset.

Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
This commit is contained in:
2026-06-10 06:17:52 +00:00
co-authored by Claude Fable 5
parent 5bec69f40a
commit 3eef989739
6 changed files with 160 additions and 9 deletions
@@ -51,11 +51,64 @@ def test_freq_four_with_offset_one_uses_layers_zero_four_eight():
def test_pattern_marks_non_shared_layers_active():
cfg = SimpleNamespace(index_topk_freq=1, index_topk_pattern="CSSSCSS")
cfg = SimpleNamespace(index_topk_freq=1, index_topk_pattern="FSSSFSS")
plan = build_nsa_index_layer_plan(cfg, 0, len(cfg.index_topk_pattern))
assert plan.active_layer_ids == (0, 4)
def test_pattern_next_skip_matches_skip_of_next_layer():
cfg = SimpleNamespace(index_topk_freq=1, index_topk_pattern="FFSFSSSF")
flags = [nsa_index_skip_flags(cfg, i) for i in range(8)]
for i in range(7):
assert flags[i][1] == flags[i + 1][0], f"layer {i}"
assert flags[0][0] is False
assert flags[7][1] is False # last layer: nothing follows
def test_pattern_rejects_invalid_characters():
cfg = SimpleNamespace(index_topk_freq=1, index_topk_pattern="FCSSF")
with pytest.raises(ValueError, match="invalid characters"):
build_nsa_index_layer_plan(cfg, 0, 5)
def test_pattern_rejects_leading_shared_layer():
cfg = SimpleNamespace(index_topk_freq=1, index_topk_pattern="SFFFF")
with pytest.raises(ValueError, match="must start with 'F'"):
build_nsa_index_layer_plan(cfg, 0, 5)
def test_pattern_length_must_match_num_hidden_layers():
cfg = SimpleNamespace(
index_topk_freq=1, index_topk_pattern="FFSF", num_hidden_layers=6
)
with pytest.raises(ValueError, match="does not match"):
build_nsa_index_layer_plan(cfg, 0, 6)
def test_pattern_length_matching_num_hidden_layers_accepted():
cfg = SimpleNamespace(
index_topk_freq=1, index_topk_pattern="FFSFSS", num_hidden_layers=6
)
plan = build_nsa_index_layer_plan(cfg, 0, 6)
assert plan.active_layer_ids == (0, 1, 3)
def test_glm5_reference_pattern_accepted():
pattern = (
"FFSFSSSFSSFFFSSSFFFSFSSSSSSFFSFFSFFSSFFFFFFSFFFFFSFFSSSSSSFSFFFSFSSSFSFFSFFSSS"
)
cfg = SimpleNamespace(
index_topk_freq=1,
index_topk_pattern=pattern,
num_hidden_layers=len(pattern),
)
plan = build_nsa_index_layer_plan(cfg, 0, len(pattern))
assert len(plan.active_layer_ids) == pattern.count("F")
for layer_id in range(len(pattern)):
skip, _ = nsa_index_skip_flags(cfg, layer_id)
assert skip == (pattern[layer_id] == "S")
def test_nextn_keeps_all_draft_layers_active_for_state_safety():
cfg = SimpleNamespace(index_topk_freq=4, index_skip_topk_offset=1)
plan = build_nsa_index_layer_plan(cfg, 0, 1, is_nextn=True)