sglang/test/registered/unit/configs/test_nsa_index_layers.py

from types import SimpleNamespace

import pytest

from sglang.srt.configs.nsa_index_layers import (
    build_nsa_index_layer_plan,
    nsa_indexer_layer_needs_weights,
    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)


def test_indexer_weights_needed_only_for_active_target_layers():
    cfg = SimpleNamespace(index_topk_freq=1, index_topk_pattern="FSF")

    assert nsa_indexer_layer_needs_weights(cfg, 0) is True
    assert nsa_indexer_layer_needs_weights(cfg, 1) is False
    assert nsa_indexer_layer_needs_weights(cfg, 2) is True


def test_indexer_weights_kept_for_nextn_even_when_topk_is_shared():
    cfg = SimpleNamespace(index_topk_freq=4, index_skip_topk_offset=1)

    assert nsa_index_skip_flags(cfg, 0, is_nextn=True) == (True, True)
    assert nsa_indexer_layer_needs_weights(cfg, 0, is_nextn=True) is True