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:
@@ -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)
|
||||
|
||||
Reference in New Issue
Block a user