tests/test_cross_kv_boundary.py

import importlib.util
import json
import sys
import types
from pathlib import Path
from types import SimpleNamespace

import pytest
import torch

REPO_ROOT = Path(__file__).resolve().parents[1]
PACKAGE_NAME = "moss_vl_checkpoint_test"


def _load_checkpoint_module(module_name):
    if PACKAGE_NAME not in sys.modules:
        package = types.ModuleType(PACKAGE_NAME)
        package.__path__ = [str(REPO_ROOT)]
        package.__package__ = PACKAGE_NAME
        sys.modules[PACKAGE_NAME] = package

    qualified_name = f"{PACKAGE_NAME}.{module_name}"
    if qualified_name in sys.modules:
        return sys.modules[qualified_name]

    spec = importlib.util.spec_from_file_location(
        qualified_name,
        REPO_ROOT / f"{module_name}.py",
    )
    module = importlib.util.module_from_spec(spec)
    sys.modules[qualified_name] = module
    spec.loader.exec_module(module)
    return module


configuration = _load_checkpoint_module("configuration_moss_vl")
modeling = _load_checkpoint_module("modeling_moss_vl")


def _vision_token_info():
    return [
        {
            "medias": [
                {
                    "length": 8,
                    "num_frames": 2,
                    "vision_tokens_per_frame": 3,
                    "has_separator": True,
                },
                {
                    "length": 2,
                    "num_frames": 1,
                    "vision_tokens_per_frame": 1,
                    "has_separator": True,
                },
            ],
            "pad_end": 16,
        },
        {
            "medias": [
                {
                    "length": 12,
                    "num_frames": 4,
                    "has_separator": False,
                }
            ],
            "pad_end": 16,
        },
    ]


def _frame_mask(dtype=torch.bool):
    visible_frame_counts = torch.tensor(
        [
            [0, 1, 3, 3],
            [0, 2, 3, 4],
        ]
    )
    frame_ids = torch.arange(4).view(1, 1, 1, 4)
    mask = frame_ids >= visible_frame_counts[:, None, :, None]
    if dtype == torch.bool:
        return mask

    additive_mask = torch.zeros(mask.shape, dtype=dtype)
    return additive_mask.masked_fill(mask, torch.finfo(dtype).min)


@pytest.mark.parametrize("mask_dtype", [torch.bool, torch.float32])
def test_boundary_matches_dense_mask(mask_dtype):
    frame_mask = _frame_mask(mask_dtype)
    vision_token_info = _vision_token_info()

    boundary = modeling.MossVLModel._build_cross_kv_boundary(frame_mask, vision_token_info)

    assert boundary.dtype == torch.int32
    assert boundary.is_contiguous()
    torch.testing.assert_close(
        boundary,
        torch.tensor(
            [
                [0, 4, 10, 10],
                [0, 6, 9, 12],
            ],
            dtype=torch.int32,
        ),
    )

    model = modeling.MossVLModel.__new__(modeling.MossVLModel)
    torch.nn.Module.__init__(model)
    dense_mask = model._expand_cross_attention_mask(
        frame_mask,
        vision_token_info,
        target_dtype=torch.float32,
    )
    dense_visible = dense_mask[:, 0] != torch.finfo(dense_mask.dtype).min
    token_ids = torch.arange(dense_mask.shape[-1]).view(1, 1, -1)
    boundary_visible = token_ids < boundary.unsqueeze(-1)

    torch.testing.assert_close(boundary_visible, dense_visible)


def test_boundary_rejects_non_prefix_visibility():
    frame_mask = _frame_mask()
    frame_mask[0, 0, 1] = torch.tensor([False, True, False, True])

    with pytest.raises(ValueError, match="prefix of vision frames"):
        modeling.MossVLModel._build_cross_kv_boundary(frame_mask, _vision_token_info())


def test_boundary_rejects_arbitrary_additive_bias():
    frame_mask = _frame_mask(torch.float32)
    frame_mask[0, 0, 1, 0] = -1.0

    with pytest.raises(ValueError, match="only 0 and the dtype minimum value"):
        modeling.MossVLModel._build_cross_kv_boundary(frame_mask, _vision_token_info())


def test_cross_attention_backend_override_reaches_specialized_fa3(monkeypatch):
    config = configuration.MossVLTextConfig(
        hidden_size=8,
        intermediate_size=16,
        num_hidden_layers=1,
        num_attention_heads=2,
        num_key_value_heads=1,
        head_dim=4,
        cross_attention_layers=[0],
        cross_attention_implementation="flash_attention_3",
    )
    config._attn_implementation = "flash_attention_2"
    attention = modeling.MossVLTextCrossAttention(config, layer_idx=0)

    captured = {}

    def fake_flash_attn_func(
        query,
        key,
        value,
        softmax_scale=None,
        causal=None,
        cross_kv_boundary=None,
    ):
        captured.update(
            query=query,
            key=key,
            value=value,
            softmax_scale=softmax_scale,
            causal=causal,
            cross_kv_boundary=cross_kv_boundary,
        )
        return torch.zeros_like(query)

    monkeypatch.setattr(modeling, "_get_moss_flash_attn_3_func", lambda: fake_flash_attn_func)

    hidden_states = torch.randn(2, 3, 8)
    vision_states = torch.randn(2, 5, 8)
    boundary = torch.tensor([[0, 2, 5], [1, 3, 5]], dtype=torch.int64)
    output, attention_weights = attention(
        hidden_states,
        vision_states,
        cross_kv_boundary=boundary,
        cache_position=torch.arange(5),
    )

    assert output.shape == hidden_states.shape
    assert attention_weights is None
    assert captured["query"].shape == (2, 3, 2, 4)
    assert captured["key"].shape == (2, 5, 1, 4)
    assert captured["value"].shape == (2, 5, 1, 4)
    assert captured["softmax_scale"] == 0.5
    assert captured["causal"] is False
    assert captured["cross_kv_boundary"].dtype == torch.int32
    assert captured["cross_kv_boundary"].is_contiguous()
    torch.testing.assert_close(
        captured["cross_kv_boundary"],
        boundary.to(torch.int32),
    )


def test_global_fa3_does_not_opt_into_specialized_cross_attention(monkeypatch):
    config = configuration.MossVLTextConfig(
        hidden_size=8,
        intermediate_size=16,
        num_hidden_layers=1,
        num_attention_heads=2,
        num_key_value_heads=1,
        head_dim=4,
        cross_attention_layers=[0],
    )
    config._attn_implementation = "flash_attention_3"
    attention = modeling.MossVLTextCrossAttention(config, layer_idx=0)
    monkeypatch.setattr(
        modeling,
        "_get_moss_flash_attn_3_func",
        lambda: pytest.fail("global FA3 must keep the existing SDPA cross-attention path"),
    )

    hidden_states = torch.randn(2, 3, 8)
    output, _ = attention(
        hidden_states,
        torch.randn(2, 5, 8),
        cross_kv_boundary=torch.full((2, 3), 5, dtype=torch.int32),
        cache_position=torch.arange(5),
    )

    assert output.shape == hidden_states.shape


def test_model_routes_coarse_mask_to_boundary_without_dense_expansion(monkeypatch):
    model = modeling.MossVLModel.__new__(modeling.MossVLModel)
    torch.nn.Module.__init__(model)
    model.config = SimpleNamespace(
        text_config=SimpleNamespace(
            _attn_implementation="flash_attention_2",
            cross_attention_implementation="flash_attention_3",
        )
    )
    model.vision_token_info = _vision_token_info()
    model.rope_deltas = None

    captured = {}

    class FakeLanguageModel(torch.nn.Module):
        def forward(self, **kwargs):
            captured.update(kwargs)
            return SimpleNamespace(
                last_hidden_state=kwargs["inputs_embeds"],
                past_key_values=None,
                hidden_states=None,
                attentions=None,
            )

    model.language_model = FakeLanguageModel()
    monkeypatch.setattr(
        model,
        "_expand_cross_attention_mask",
        lambda *args, **kwargs: pytest.fail("dense mask expansion should not run"),
    )

    inputs_embeds = torch.randn(2, 4, 8)
    model(
        inputs_embeds=inputs_embeds,
        position_ids=torch.arange(4).view(1, 1, 4).expand(3, 2, 4),
        cross_attention_mask=_frame_mask(),
        cache_position=torch.arange(4),
    )

    expected_boundary = torch.tensor(
        [
            [0, 4, 10, 10],
            [0, 6, 9, 12],
        ],
        dtype=torch.int32,
    )
    assert captured["cross_attention_mask"] is None
    torch.testing.assert_close(captured["cross_kv_boundary"], expected_boundary)
    torch.testing.assert_close(
        captured["full_text_row_masked_out_mask"],
        (expected_boundary > 0)[:, None, :, None].to(inputs_embeds.dtype),
    )


def test_global_fa3_keeps_dense_cross_attention_mask(monkeypatch):
    model = modeling.MossVLModel.__new__(modeling.MossVLModel)
    torch.nn.Module.__init__(model)
    model.config = SimpleNamespace(
        text_config=SimpleNamespace(
            _attn_implementation="flash_attention_3",
            cross_attention_implementation=None,
        )
    )
    model.vision_token_info = _vision_token_info()
    model.rope_deltas = None

    captured = {}

    class FakeLanguageModel(torch.nn.Module):
        def forward(self, **kwargs):
            captured.update(kwargs)
            return SimpleNamespace(
                last_hidden_state=kwargs["inputs_embeds"],
                past_key_values=None,
                hidden_states=None,
                attentions=None,
            )

    model.language_model = FakeLanguageModel()
    original_expand = model._expand_cross_attention_mask
    expand_calls = 0

    def tracking_expand(*args, **kwargs):
        nonlocal expand_calls
        expand_calls += 1
        return original_expand(*args, **kwargs)

    monkeypatch.setattr(model, "_expand_cross_attention_mask", tracking_expand)

    inputs_embeds = torch.randn(2, 4, 8)
    model(
        inputs_embeds=inputs_embeds,
        position_ids=torch.arange(4).view(1, 1, 4).expand(3, 2, 4),
        cross_attention_mask=_frame_mask(),
        cache_position=torch.arange(4),
    )

    assert expand_calls == 1
    assert captured["cross_attention_mask"] is not None
    assert captured["cross_kv_boundary"] is None


def test_loader_keeps_fa2_for_other_attention_and_overrides_cross_attention():
    config, unused_kwargs = configuration.MossVLConfig.from_pretrained(
        REPO_ROOT,
        attn_implementation="flash_attention_2",
        cross_attention_implementation="flash_attention_3",
        return_unused_kwargs=True,
    )

    assert unused_kwargs == {}
    assert config._attn_implementation == "flash_attention_2"
    assert config.vision_config._attn_implementation == "flash_attention_2"
    assert config.text_config._attn_implementation == "flash_attention_2"
    assert config.cross_attention_implementation == "flash_attention_3"
    assert config.text_config.cross_attention_implementation == "flash_attention_3"


def test_loader_default_does_not_change_existing_cross_attention_dispatch():
    config = configuration.MossVLConfig.from_pretrained(
        REPO_ROOT,
        attn_implementation="flash_attention_2",
    )

    assert config.text_config._attn_implementation == "flash_attention_2"
    assert config.cross_attention_implementation is None
    assert config.text_config.cross_attention_implementation is None


def test_cross_attention_implementation_round_trips_through_text_config(tmp_path):
    config = configuration.MossVLConfig.from_pretrained(
        REPO_ROOT,
        cross_attention_implementation="flash_attention_3",
    )

    config.save_pretrained(tmp_path)
    serialized = json.loads((tmp_path / "config.json").read_text())

    assert "cross_attention_implementation" not in serialized
    assert serialized["text_config"]["cross_attention_implementation"] == "flash_attention_3"

    reloaded = configuration.MossVLConfig.from_pretrained(tmp_path)
    assert reloaded.cross_attention_implementation == "flash_attention_3"
    assert reloaded.text_config.cross_attention_implementation == "flash_attention_3"