leavelet
3eef989739
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>
137 lines
4.5 KiB
Python
137 lines
4.5 KiB
Python
from types import SimpleNamespace
|
|
|
|
import pytest
|
|
|
|
from sglang.srt.configs.nsa_index_layers import (
|
|
build_nsa_index_layer_plan,
|
|
nsa_index_skip_flags,
|
|
)
|
|
|
|
|
|
def test_default_freq_one_marks_all_target_layers_active():
|
|
cfg = SimpleNamespace(index_topk_freq=1)
|
|
plan = build_nsa_index_layer_plan(cfg, 0, 6)
|
|
assert plan.active_layer_ids == (0, 1, 2, 3, 4, 5)
|
|
assert [nsa_index_skip_flags(cfg, i)[0] for i in range(6)] == [False] * 6
|
|
|
|
|
|
def test_freq_four_without_offset_matches_current_model_formula():
|
|
cfg = SimpleNamespace(index_topk_freq=4)
|
|
plan = build_nsa_index_layer_plan(cfg, 0, 12)
|
|
assert plan.active_layer_ids == (0, 1, 5, 9)
|
|
assert [nsa_index_skip_flags(cfg, i)[0] for i in range(10)] == [
|
|
False,
|
|
False,
|
|
True,
|
|
True,
|
|
True,
|
|
False,
|
|
True,
|
|
True,
|
|
True,
|
|
False,
|
|
]
|
|
|
|
|
|
def test_freq_four_with_offset_one_uses_layers_zero_four_eight():
|
|
cfg = SimpleNamespace(index_topk_freq=4, index_skip_topk_offset=1)
|
|
plan = build_nsa_index_layer_plan(cfg, 0, 12)
|
|
assert plan.active_layer_ids == (0, 4, 8)
|
|
assert [nsa_index_skip_flags(cfg, i)[0] for i in range(9)] == [
|
|
False,
|
|
True,
|
|
True,
|
|
True,
|
|
False,
|
|
True,
|
|
True,
|
|
True,
|
|
False,
|
|
]
|
|
|
|
|
|
def test_pattern_marks_non_shared_layers_active():
|
|
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)
|
|
assert plan.active_layer_ids == (0,)
|
|
assert nsa_index_skip_flags(cfg, 0, is_nextn=True) == (True, True)
|
|
|
|
|
|
def test_nonzero_start_layer_preserves_logical_layer_ids():
|
|
cfg = SimpleNamespace(index_topk_freq=4, index_skip_topk_offset=1)
|
|
plan = build_nsa_index_layer_plan(cfg, 4, 13)
|
|
assert plan.active_layer_ids == (4, 8, 12)
|
|
assert plan.slot_for_layer(8) == 1
|
|
|
|
|
|
def test_inactive_slot_lookup_fails_fast():
|
|
cfg = SimpleNamespace(index_topk_freq=4, index_skip_topk_offset=1)
|
|
plan = build_nsa_index_layer_plan(cfg, 0, 8)
|
|
with pytest.raises(RuntimeError, match="inactive index layer requested"):
|
|
plan.slot_for_layer(1)
|
|
|
|
|
|
def test_invalid_offset_fails_before_layer_zero_can_skip_without_prior_topk():
|
|
cfg = SimpleNamespace(index_topk_freq=4, index_skip_topk_offset=0)
|
|
with pytest.raises(ValueError, match="index_skip_topk_offset"):
|
|
nsa_index_skip_flags(cfg, 0)
|