# Copyright (c) Microsoft Corporation.
# SPDX-License-Identifier: Apache-2.0

# DeepSpeed Team
"""
Unit tests for AutoSP multimodal sequence parallelism:
  - autosp_detector: model scanning
  - UlyssesSPViTAttention: ViT SP wrapper
  - auto_wrap_model_for_sp: end-to-end wrapping
  - ModalityFusionSPAdapter: cross-modal gather/scatter
  - LlavaFusionAdapter: LLaVA-style visual token splice
  - InternVLFusionAdapter: InternVL-style IMG_CONTEXT token splice
  - Qwen2VLFusionAdapter: Qwen2-VL vision_start/end bounded splice
"""

import pytest
import torch
import torch.nn as nn

from deepspeed.sequence.autosp_detector import (SPModelInfo, _LLM_ATTN_CLASSNAMES, _VIT_ATTN_CLASSNAMES,
                                                detect_model_sp_info)
from deepspeed.sequence.autosp_fusion import (InternVLFusionAdapter, LlavaFusionAdapter, ModalityFusionSPAdapter,
                                              Qwen2VLFusionAdapter)
from deepspeed.sequence.autosp_vit import UlyssesSPViTAttention
from deepspeed.sequence.auto_sp import _set_module_by_name, auto_wrap_model_for_sp
from deepspeed.sequence.layer import DistributedAttention

# ---------------------------------------------------------------------------
# Minimal fake modules that mimic the interface of real attention layers
# without requiring a GPU or a real transformer model.
# ---------------------------------------------------------------------------


class _FakeViTAttn(nn.Module):
    """Identity ViT attention — returns hidden_states unchanged."""

    def forward(self, hidden_states, **kwargs):
        return hidden_states


class _FakeViTAttnTuple(nn.Module):
    """ViT attention that returns a (output, weights) tuple."""

    def forward(self, hidden_states, **kwargs):
        weights = torch.zeros(hidden_states.shape[0], 1, hidden_states.shape[1], hidden_states.shape[1])
        return hidden_states, weights


class _FakeLLMAttn(nn.Module):
    """Identity LLM attention."""

    def forward(self, query, key, value, *args, **kwargs):
        return query


# Register fake class names so the detector recognises them
_VIT_ATTN_CLASSNAMES.add("_FakeViTAttn")
_VIT_ATTN_CLASSNAMES.add("_FakeViTAttnTuple")
_LLM_ATTN_CLASSNAMES.add("_FakeLLMAttn")


class _FakeMultimodalModel(nn.Module):
    """Minimal multimodal model with one ViT and one LLM attention layer."""

    def __init__(self):
        super().__init__()
        self.vision_encoder = nn.ModuleList([_FakeViTAttn()])
        self.mm_projector = nn.Linear(64, 64)
        self.llm = nn.ModuleList([_FakeLLMAttn()])


class _FakeViTOnlyModel(nn.Module):

    def __init__(self, num_layers=3):
        super().__init__()
        self.layers = nn.ModuleList([_FakeViTAttn() for _ in range(num_layers)])


class _FakeLLMOnlyModel(nn.Module):
    """Minimal LLM-only model with multiple decoder attention layers."""

    def __init__(self, num_layers=2):
        super().__init__()
        self.layers = nn.ModuleList([_FakeLLMAttn() for _ in range(num_layers)])


# ---------------------------------------------------------------------------
# Helpers
# ---------------------------------------------------------------------------


def _make_mock_process_group(world_size: int, rank: int):
    """Return a mock object that satisfies dist.get_world_size / get_rank."""
    import unittest.mock as mock
    import deepspeed.comm as dist

    pg = mock.MagicMock()
    dist.get_world_size = mock.MagicMock(return_value=world_size)
    dist.get_rank = mock.MagicMock(return_value=rank)

    def _fake_all_gather(tensor_list, tensor, group=None):
        for t in tensor_list:
            t.copy_(tensor)

    dist.all_gather = _fake_all_gather
    return pg


# ---------------------------------------------------------------------------
# autosp_detector tests
# ---------------------------------------------------------------------------


class TestAutospDetector:

    def test_detects_vit_and_llm(self):
        model = _FakeMultimodalModel()
        info = detect_model_sp_info(model)
        assert len(info.vit_attn_modules) == 1
        assert len(info.llm_attn_modules) == 1

    def test_detects_vision_projection(self):
        model = _FakeMultimodalModel()
        info = detect_model_sp_info(model)
        assert info.vision_projection_module is not None
        name, module = info.vision_projection_module
        assert "mm_projector" in name

    def test_detects_multiple_vit_layers(self):
        model = _FakeViTOnlyModel(num_layers=4)
        info = detect_model_sp_info(model)
        assert len(info.vit_attn_modules) == 4
        assert len(info.llm_attn_modules) == 0
        assert info.vision_projection_module is None

    def test_empty_model_returns_empty_info(self):
        model = nn.Sequential(nn.Linear(8, 8))
        info = detect_model_sp_info(model)
        assert isinstance(info, SPModelInfo)
        assert len(info.vit_attn_modules) == 0
        assert len(info.llm_attn_modules) == 0

    def test_only_first_projection_is_recorded(self):
        """Multiple projection-like names → only the outermost is recorded."""

        class _M(nn.Module):

            def __init__(self):
                super().__init__()
                self.mm_projector = nn.Sequential(nn.Linear(8, 8))
                self.mm_projector.visual_projection = nn.Linear(8, 8)

        model = _M()
        info = detect_model_sp_info(model)
        assert info.vision_projection_module is not None
        # Should be the outermost "mm_projector", not the nested one
        name, _ = info.vision_projection_module
        assert name == "mm_projector"


# ---------------------------------------------------------------------------
# UlyssesSPViTAttention tests (CPU, rank-0 simulation via mocks)
# ---------------------------------------------------------------------------


class TestUlyssesSPViTAttention:

    @pytest.mark.parametrize("has_cls_token", [True, False])
    @pytest.mark.parametrize("num_patches,world_size", [
        (16, 4),
        (16, 2),
        (9, 3),
    ])
    def test_output_shape_matches_input(self, has_cls_token, num_patches, world_size):
        """Output shape must equal input shape for any padding scenario."""
        pg = _make_mock_process_group(world_size=world_size, rank=0)
        attn = _FakeViTAttn()
        wrapper = UlyssesSPViTAttention(attn, pg, has_cls_token=has_cls_token)

        local_patches = num_patches // world_size
        seq_len = (1 + local_patches) if has_cls_token else local_patches
        x = torch.randn(2, seq_len, 32)

        out = wrapper(x)
        assert out.shape == x.shape, f"Expected {x.shape}, got {out.shape}"

    def test_tuple_output_unwrapped_correctly(self):
        """Wrappers that return (output, weights) tuples are handled."""
        pg = _make_mock_process_group(world_size=2, rank=0)
        attn = _FakeViTAttnTuple()
        wrapper = UlyssesSPViTAttention(attn, pg, has_cls_token=False)

        x = torch.randn(1, 8, 16)  # 8 patches, 2 ranks → 4 local each
        result = wrapper(x)
        # Should return a tuple: (attention_output, attention_weights)
        assert isinstance(result, tuple)
        assert result[0].shape == x.shape

    def test_identity_attn_preserves_values(self):
        """When attn is identity, output values should match input values."""
        world_size = 2
        pg = _make_mock_process_group(world_size=world_size, rank=0)
        attn = _FakeViTAttn()
        wrapper = UlyssesSPViTAttention(attn, pg, has_cls_token=True)

        # Each rank holds cls + 4 local patches
        x = torch.arange(2 * 5 * 4, dtype=torch.float).reshape(2, 5, 4)
        out = wrapper(x)
        # CLS token should be identical
        assert torch.allclose(out[:, :1, :], x[:, :1, :])
        # Local patch slice should match input patches for identity attn
        assert torch.allclose(out[:, 1:, :], x[:, 1:, :])


# ---------------------------------------------------------------------------
# auto_wrap_model_for_sp tests
# ---------------------------------------------------------------------------


class TestAutoWrapModelForSP:

    def test_vit_layers_replaced(self):
        pg = _make_mock_process_group(world_size=2, rank=0)
        model = _FakeViTOnlyModel(num_layers=2)
        auto_wrap_model_for_sp(model, pg)
        for layer in model.layers:
            assert isinstance(layer, UlyssesSPViTAttention)

    def test_raises_on_unknown_model(self):
        pg = _make_mock_process_group(world_size=2, rank=0)
        model = nn.Sequential(nn.Linear(8, 8))
        with pytest.raises(ValueError, match="no recognisable attention"):
            auto_wrap_model_for_sp(model, pg)

    def test_set_module_by_name_shallow(self):
        model = _FakeViTOnlyModel(num_layers=1)
        new_mod = nn.Linear(4, 4)
        _set_module_by_name(model, "layers.0", new_mod)
        assert model.layers[0] is new_mod

    def test_set_module_by_name_deep(self):
        model = _FakeMultimodalModel()
        new_mod = nn.Identity()
        _set_module_by_name(model, "vision_encoder.0", new_mod)
        assert model.vision_encoder[0] is new_mod

    def test_llm_layers_replaced_with_distributed_attention(self):
        """LLM attention layers must be wrapped with DistributedAttention."""
        pg = _make_mock_process_group(world_size=2, rank=0)
        model = _FakeLLMOnlyModel(num_layers=3)
        auto_wrap_model_for_sp(model, pg)
        for layer in model.layers:
            assert isinstance(layer, DistributedAttention)

    def test_multimodal_model_wraps_both_branches(self):
        """Both ViT and LLM attention layers must be replaced in a combined model."""
        pg = _make_mock_process_group(world_size=2, rank=0)
        model = _FakeMultimodalModel()
        returned = auto_wrap_model_for_sp(model, pg)
        # auto_wrap_model_for_sp must return the same object (in-place)
        assert returned is model
        assert isinstance(model.vision_encoder[0], UlyssesSPViTAttention)
        assert isinstance(model.llm[0], DistributedAttention)

    def test_original_module_preserved_inside_wrapper(self):
        """The wrapped module should still be accessible inside the wrapper."""
        pg = _make_mock_process_group(world_size=2, rank=0)
        model = _FakeViTOnlyModel(num_layers=1)
        original_attn = model.layers[0]
        auto_wrap_model_for_sp(model, pg)
        assert model.layers[0].attn is original_attn


# ---------------------------------------------------------------------------
# ModalityFusionSPAdapter tests
# ---------------------------------------------------------------------------


class _ConcatFusionAdapter(ModalityFusionSPAdapter):
    """Concrete subclass that appends visual tokens after text tokens."""

    def _splice_visual_into_text(self, text_embeds, visual_embeds, input_ids):
        return torch.cat([text_embeds, visual_embeds], dim=1)


class TestModalityFusionSPAdapter:

    def test_base_class_raises_not_implemented(self):
        """The base _splice_visual_into_text must raise NotImplementedError."""
        pg = _make_mock_process_group(world_size=2, rank=0)
        adapter = ModalityFusionSPAdapter(nn.Identity(), pg)
        with pytest.raises(NotImplementedError):
            adapter._splice_visual_into_text(None, None, None)

    @pytest.mark.parametrize("world_size,local_v,text_len,hidden", [
        (2, 4, 6, 8),
        (4, 3, 5, 16),
        (1, 8, 8, 4),
    ])
    def test_output_shape(self, world_size, local_v, text_len, hidden):
        """Output local_len must equal ceil(fused_len / world_size)."""
        pg = _make_mock_process_group(world_size=world_size, rank=0)
        adapter = _ConcatFusionAdapter(nn.Identity(), pg)

        bs = 2
        visual = torch.randn(bs, local_v, hidden)
        text = torch.randn(bs, text_len, hidden)
        ids = torch.zeros(bs, text_len, dtype=torch.long)

        out = adapter(visual, text, ids)

        # all_gather mock copies local_v to each of world_size slots
        fused_len = text_len + local_v * world_size
        pad = (world_size - fused_len % world_size) % world_size
        expected_local = (fused_len + pad) // world_size
        assert out.shape == (bs, expected_local, hidden), f"Expected ({bs},{expected_local},{hidden}), got {out.shape}"

    def test_padding_produces_valid_output_when_not_divisible(self):
        """When fused_len % world_size != 0, padding must not raise and output is well-formed."""
        world_size = 4
        # text_len=5, local_v=3 → fused_len = 5 + 3*4 = 17, needs padding of 3
        pg = _make_mock_process_group(world_size=world_size, rank=0)
        adapter = _ConcatFusionAdapter(nn.Identity(), pg)

        bs, local_v, text_len, hidden = 1, 3, 5, 4
        out = adapter(
            torch.randn(bs, local_v, hidden),
            torch.randn(bs, text_len, hidden),
            torch.zeros(bs, text_len, dtype=torch.long),
        )
        # padded_len = 20, local_len = 5
        assert out.shape == (bs, 5, hidden)

    def test_no_padding_when_divisible(self):
        """When fused_len is already divisible, no extra tokens should be added."""
        world_size = 4
        # text_len=4, local_v=4 → fused_len = 4 + 4*4 = 20, divisible by 4
        pg = _make_mock_process_group(world_size=world_size, rank=0)
        adapter = _ConcatFusionAdapter(nn.Identity(), pg)

        bs, local_v, text_len, hidden = 1, 4, 4, 8
        out = adapter(
            torch.randn(bs, local_v, hidden),
            torch.randn(bs, text_len, hidden),
            torch.zeros(bs, text_len, dtype=torch.long),
        )
        assert out.shape == (bs, 5, hidden)  # 20 // 4 = 5

    def test_different_ranks_return_different_slices(self):
        """Rank 0 and rank 1 must return different slices of the fused sequence."""
        world_size = 2
        bs, local_v, text_len, hidden = 1, 4, 4, 8
        # Use distinct text vs visual values so slices clearly differ
        text = torch.zeros(bs, text_len, hidden)
        visual = torch.ones(bs, local_v, hidden)
        ids = torch.zeros(bs, text_len, dtype=torch.long)

        outputs = {}
        for rank in range(world_size):
            pg = _make_mock_process_group(world_size=world_size, rank=rank)
            adapter = _ConcatFusionAdapter(nn.Identity(), pg)
            outputs[rank] = adapter(visual.clone(), text.clone(), ids.clone())

        # fused = [0,0,0,0, 1,1,1,1, 1,1,1,1]  (text zeros then visual ones x2)
        # rank 0: indices 0-5, rank 1: indices 6-11
        assert not torch.allclose(outputs[0], outputs[1])

    def test_projection_is_applied(self):
        """Projection layer must transform visual features before gather."""
        world_size = 2
        pg = _make_mock_process_group(world_size=world_size, rank=0)

        # Use a projection that doubles all values
        class _DoubleProjection(nn.Module):

            def forward(self, x):
                return x * 2.0

        adapter = _ConcatFusionAdapter(_DoubleProjection(), pg)
        bs, local_v, text_len, hidden = 1, 4, 4, 8
        visual = torch.ones(bs, local_v, hidden)
        text = torch.zeros(bs, text_len, hidden)
        ids = torch.zeros(bs, text_len, dtype=torch.long)

        out = adapter(visual, text, ids)
        # The visual part of the output should have value 2.0 (doubled), not 1.0
        # rank 0 gets the first local_len tokens; fused = [text(0)*4, visual(2)*8]
        # Since text_len=4 and local_len=6, rank0 slice starts with text zeros
        # and ends with some visual twos.
        assert out.max().item() == pytest.approx(2.0)


# ---------------------------------------------------------------------------
# LlavaFusionAdapter tests  (tests _splice_visual_into_text directly)
# ---------------------------------------------------------------------------

_IMAGE_ID = -200  # matches ModalityFusionSPAdapter default


def _make_llava_adapter(world_size=2, rank=0):
    pg = _make_mock_process_group(world_size=world_size, rank=rank)
    return LlavaFusionAdapter(nn.Identity(), pg, image_token_id=_IMAGE_ID)


class TestLlavaFusionAdapter:

    def test_single_image_fused_shape(self):
        """One image placeholder per sample → fused length = text_len - 1 + num_visual."""
        adapter = _make_llava_adapter()
        bs, text_len, num_vis, hidden = 2, 6, 4, 8
        # Place a single image placeholder at position 2.
        ids = torch.zeros(bs, text_len, dtype=torch.long)
        ids[:, 2] = _IMAGE_ID
        text = torch.randn(bs, text_len, hidden)
        visual = torch.randn(bs, num_vis, hidden)

        fused = adapter._splice_visual_into_text(text, visual, ids)
        # placeholder is removed and replaced by num_vis tokens
        assert fused.shape == (bs, text_len - 1 + num_vis, hidden)

    def test_text_values_preserved_around_image(self):
        """Text tokens before and after the placeholder must be numerically intact."""
        adapter = _make_llava_adapter()
        bs, text_len, num_vis, hidden = 1, 5, 3, 4
        # Placeholder at index 2: text = [A, B, <img>, C, D]
        ids = torch.zeros(bs, text_len, dtype=torch.long)
        ids[0, 2] = _IMAGE_ID
        text = torch.arange(bs * text_len * hidden, dtype=torch.float).reshape(bs, text_len, hidden)
        visual = torch.ones(bs, num_vis, hidden) * 99.0

        fused = adapter._splice_visual_into_text(text, visual, ids)
        # fused = [A, B, vis0, vis1, vis2, C, D]
        assert torch.allclose(fused[0, :2], text[0, :2])  # A, B preserved
        assert torch.allclose(fused[0, 5:], text[0, 3:])  # C, D preserved
        assert torch.allclose(fused[0, 2:5], visual[0])  # visual inserted

    def test_no_image_token_returns_text_unchanged(self):
        """When input_ids has no placeholder, output equals text_embeds exactly."""
        adapter = _make_llava_adapter()
        bs, text_len, hidden = 2, 6, 8
        ids = torch.zeros(bs, text_len, dtype=torch.long)  # no -200
        text = torch.randn(bs, text_len, hidden)
        visual = torch.randn(bs, 4, hidden)

        fused = adapter._splice_visual_into_text(text, visual, ids)
        assert fused.shape == (bs, text_len, hidden)
        assert torch.allclose(fused, text)

    def test_multi_image_splice(self):
        """Two placeholders per sample → visual tokens split evenly between them."""
        adapter = _make_llava_adapter()
        bs, text_len, num_vis, hidden = 1, 7, 6, 4
        # Placeholders at index 1 and 4: [t0, <img>, t2, t3, <img>, t5, t6]
        ids = torch.zeros(bs, text_len, dtype=torch.long)
        ids[0, 1] = _IMAGE_ID
        ids[0, 4] = _IMAGE_ID
        text = torch.zeros(bs, text_len, hidden)
        # First 3 visual tokens = 1.0, last 3 = 2.0 (so we can tell them apart)
        visual = torch.cat([torch.ones(bs, 3, hidden), torch.full((bs, 3, hidden), 2.0)], dim=1)

        fused = adapter._splice_visual_into_text(text, visual, ids)
        # Expected fused length: 7 - 2 placeholders + 6 visual = 11
        assert fused.shape == (bs, 11, hidden)
        # First chunk (indices 1-3) should be 1.0
        assert torch.allclose(fused[0, 1:4], torch.ones(3, hidden))
        # Second chunk (indices 6-8) should be 2.0
        assert torch.allclose(fused[0, 6:9], torch.full((3, hidden), 2.0))

    def test_batch_padding_when_lengths_differ(self):
        """Samples with different numbers of image tokens are padded to max length."""
        adapter = _make_llava_adapter()
        hidden = 4
        # Sample 0: 1 placeholder in a 4-token sequence + 2 visual → fused len = 5
        # Sample 1: no placeholder in a 4-token sequence → fused len = 4
        ids = torch.zeros(2, 4, dtype=torch.long)
        ids[0, 1] = _IMAGE_ID
        text = torch.ones(2, 4, hidden)
        visual = torch.ones(2, 2, hidden) * 3.0

        fused = adapter._splice_visual_into_text(text, visual, ids)
        # Max fused length is 5; sample 1 padded with zeros at the end.
        assert fused.shape == (2, 5, hidden)
        assert torch.all(fused[1, 4:] == 0)  # padding tokens are zero

    def test_forward_end_to_end_shape(self):
        """Full forward pass through LlavaFusionAdapter returns the correct shard shape."""
        world_size = 2
        pg = _make_mock_process_group(world_size=world_size, rank=0)
        adapter = LlavaFusionAdapter(nn.Identity(), pg, image_token_id=_IMAGE_ID)

        bs, local_v, text_len, hidden = 1, 4, 6, 8
        ids = torch.zeros(bs, text_len, dtype=torch.long)
        ids[0, 2] = _IMAGE_ID  # one placeholder
        visual = torch.randn(bs, local_v, hidden)
        text = torch.randn(bs, text_len, hidden)

        out = adapter(visual, text, ids)
        # fused_len = text_len - 1 + local_v * world_size = 5 + 8 = 13
        # padded to 14 (next multiple of 2), local = 7
        assert out.shape == (bs, 7, hidden)


# ---------------------------------------------------------------------------
# InternVLFusionAdapter tests  (tests _splice_visual_into_text directly)
# ---------------------------------------------------------------------------

_CONTEXT_ID = 92546  # arbitrary IMG_CONTEXT token id for tests
_START_ID = 92545
_END_ID = 92547


def _make_internvl_adapter(world_size=2, rank=0):
    pg = _make_mock_process_group(world_size=world_size, rank=rank)
    return InternVLFusionAdapter(nn.Identity(), pg, image_token_id=_CONTEXT_ID)


class TestInternVLFusionAdapter:

    def test_context_tokens_replaced_with_visual(self):
        """IMG_CONTEXT positions must carry visual embeddings after splice."""
        adapter = _make_internvl_adapter()
        bs, text_len, hidden = 1, 7, 4
        # Layout: [t0, START, ctx, ctx, ctx, END, t6]
        ids = torch.zeros(bs, text_len, dtype=torch.long)
        ids[0, 2] = _CONTEXT_ID
        ids[0, 3] = _CONTEXT_ID
        ids[0, 4] = _CONTEXT_ID

        text = torch.zeros(bs, text_len, hidden)
        visual = torch.ones(bs, 3, hidden) * 7.0

        fused = adapter._splice_visual_into_text(text, visual, ids)
        assert torch.allclose(fused[0, 2:5], visual[0])

    def test_sequence_length_preserved(self):
        """Output length must equal input length (1-to-1 replacement)."""
        adapter = _make_internvl_adapter()
        bs, text_len, hidden = 2, 10, 8
        ids = torch.zeros(bs, text_len, dtype=torch.long)
        ids[:, 3:7] = _CONTEXT_ID  # 4 context tokens per sample
        text = torch.randn(bs, text_len, hidden)
        visual = torch.randn(bs, 4, hidden)

        fused = adapter._splice_visual_into_text(text, visual, ids)
        assert fused.shape == (bs, text_len, hidden)

    def test_boundary_tokens_preserved(self):
        """IMG_START and IMG_END embeddings must be unchanged after splice."""
        adapter = _make_internvl_adapter()
        bs, text_len, hidden = 1, 5, 4
        # [START, ctx, ctx, END, text]
        ids = torch.zeros(bs, text_len, dtype=torch.long)
        ids[0, 1] = _CONTEXT_ID
        ids[0, 2] = _CONTEXT_ID

        text = torch.arange(bs * text_len * hidden, dtype=torch.float).reshape(bs, text_len, hidden)
        visual = torch.ones(bs, 2, hidden) * 99.0

        fused = adapter._splice_visual_into_text(text, visual, ids)
        # Position 0 (START) and 3 (END) must be unchanged.
        assert torch.allclose(fused[0, 0], text[0, 0])
        assert torch.allclose(fused[0, 3], text[0, 3])

    def test_no_context_tokens_returns_text_unchanged(self):
        """When there are no IMG_CONTEXT tokens the output must equal text_embeds."""
        adapter = _make_internvl_adapter()
        bs, text_len, hidden = 2, 6, 8
        ids = torch.zeros(bs, text_len, dtype=torch.long)
        text = torch.randn(bs, text_len, hidden)
        visual = torch.randn(bs, 4, hidden)

        fused = adapter._splice_visual_into_text(text, visual, ids)
        assert torch.allclose(fused, text)

    def test_multi_image_replacement(self):
        """Two separate runs of context tokens correspond to two images."""
        adapter = _make_internvl_adapter()
        bs, text_len, hidden = 1, 10, 4
        # Image 1: positions 1-2, Image 2: positions 6-7
        ids = torch.zeros(bs, text_len, dtype=torch.long)
        ids[0, 1] = _CONTEXT_ID
        ids[0, 2] = _CONTEXT_ID
        ids[0, 6] = _CONTEXT_ID
        ids[0, 7] = _CONTEXT_ID

        text = torch.zeros(bs, text_len, hidden)
        # First 2 visual tokens = 1.0, next 2 = 2.0
        visual = torch.cat([torch.ones(bs, 2, hidden), torch.full((bs, 2, hidden), 2.0)], dim=1)

        fused = adapter._splice_visual_into_text(text, visual, ids)
        assert fused.shape == (bs, text_len, hidden)
        assert torch.allclose(fused[0, 1:3], torch.ones(2, hidden))
        assert torch.allclose(fused[0, 6:8], torch.full((2, hidden), 2.0))

    def test_forward_end_to_end_shape(self):
        """Full forward pass returns the correct shard shape."""
        world_size = 2
        pg = _make_mock_process_group(world_size=world_size, rank=0)
        adapter = InternVLFusionAdapter(nn.Identity(), pg, image_token_id=_CONTEXT_ID)

        bs, local_v, text_len, hidden = 1, 3, 8, 4
        ids = torch.zeros(bs, text_len, dtype=torch.long)
        ids[0, 2:5] = _CONTEXT_ID  # 3 context tokens; local_v * world_size = 6 total
        visual = torch.randn(bs, local_v, hidden)
        text = torch.randn(bs, text_len, hidden)

        out = adapter(visual, text, ids)
        # fused_len == text_len == 8 (length-preserving); padded to 8 (divisible by 2); local = 4
        assert out.shape == (bs, 4, hidden)


# ---------------------------------------------------------------------------
# Qwen2VLFusionAdapter tests  (tests _splice_visual_into_text directly)
# ---------------------------------------------------------------------------

_VIS_START_ID = 151652
_VIS_END_ID = 151653


def _make_qwen2vl_adapter(world_size=2, rank=0):
    pg = _make_mock_process_group(world_size=world_size, rank=rank)
    return Qwen2VLFusionAdapter(nn.Identity(),
                                pg,
                                vision_start_token_id=_VIS_START_ID,
                                vision_end_token_id=_VIS_END_ID)


class TestQwen2VLFusionAdapter:

    def test_inner_tokens_replaced_with_visual(self):
        """Tokens between vision_start and vision_end must become visual embeddings."""
        adapter = _make_qwen2vl_adapter()
        bs, text_len, hidden = 1, 7, 4
        # [t0, t1, <vis_start>, pad, pad, <vis_end>, t6]
        ids = torch.zeros(bs, text_len, dtype=torch.long)
        ids[0, 2] = _VIS_START_ID
        ids[0, 5] = _VIS_END_ID

        text = torch.zeros(bs, text_len, hidden)
        visual = torch.ones(bs, 2, hidden) * 5.0

        fused = adapter._splice_visual_into_text(text, visual, ids)
        assert torch.allclose(fused[0, 3:5], visual[0])

    def test_sequence_length_preserved(self):
        """Output length must equal input length (1-to-1 replacement)."""
        adapter = _make_qwen2vl_adapter()
        bs, text_len, hidden = 2, 12, 8
        ids = torch.zeros(bs, text_len, dtype=torch.long)
        ids[:, 2] = _VIS_START_ID
        ids[:, 8] = _VIS_END_ID  # 5 inner placeholder tokens
        text = torch.randn(bs, text_len, hidden)
        visual = torch.randn(bs, 5, hidden)

        fused = adapter._splice_visual_into_text(text, visual, ids)
        assert fused.shape == (bs, text_len, hidden)

    def test_boundary_tokens_preserved(self):
        """vision_start and vision_end embeddings must be unchanged after splice."""
        adapter = _make_qwen2vl_adapter()
        bs, text_len, hidden = 1, 6, 4
        # [t0, <vis_start>, pad, pad, <vis_end>, t5]
        ids = torch.zeros(bs, text_len, dtype=torch.long)
        ids[0, 1] = _VIS_START_ID
        ids[0, 4] = _VIS_END_ID

        text = torch.arange(bs * text_len * hidden, dtype=torch.float).reshape(bs, text_len, hidden)
        visual = torch.ones(bs, 2, hidden) * 99.0

        fused = adapter._splice_visual_into_text(text, visual, ids)
        assert torch.allclose(fused[0, 1], text[0, 1])  # vision_start preserved
        assert torch.allclose(fused[0, 4], text[0, 4])  # vision_end preserved

    def test_no_vision_tokens_returns_text_unchanged(self):
        """When there are no vision_start/end tokens the output must equal text_embeds."""
        adapter = _make_qwen2vl_adapter()
        bs, text_len, hidden = 2, 8, 4
        ids = torch.zeros(bs, text_len, dtype=torch.long)
        text = torch.randn(bs, text_len, hidden)
        visual = torch.randn(bs, 4, hidden)

        fused = adapter._splice_visual_into_text(text, visual, ids)
        assert torch.allclose(fused, text)

    def test_multi_image_replacement(self):
        """Two vision blocks are handled independently."""
        adapter = _make_qwen2vl_adapter()
        bs, text_len, hidden = 1, 14, 4
        # Block 1: positions 1 (start) .. 4 (end), 2 inner tokens at 2-3
        # Block 2: positions 8 (start) .. 12 (end), 3 inner tokens at 9-11
        ids = torch.zeros(bs, text_len, dtype=torch.long)
        ids[0, 1] = _VIS_START_ID
        ids[0, 4] = _VIS_END_ID
        ids[0, 8] = _VIS_START_ID
        ids[0, 12] = _VIS_END_ID

        text = torch.zeros(bs, text_len, hidden)
        visual = torch.cat([torch.ones(bs, 2, hidden), torch.full((bs, 3, hidden), 2.0)], dim=1)

        fused = adapter._splice_visual_into_text(text, visual, ids)
        assert fused.shape == (bs, text_len, hidden)
        assert torch.allclose(fused[0, 2:4], torch.ones(2, hidden))
        assert torch.allclose(fused[0, 9:12], torch.full((3, hidden), 2.0))

    def test_forward_end_to_end_shape(self):
        """Full forward pass returns the correct shard shape."""
        world_size = 2
        pg = _make_mock_process_group(world_size=world_size, rank=0)
        adapter = Qwen2VLFusionAdapter(nn.Identity(),
                                       pg,
                                       vision_start_token_id=_VIS_START_ID,
                                       vision_end_token_id=_VIS_END_ID)

        bs, local_v, text_len, hidden = 1, 3, 10, 4
        ids = torch.zeros(bs, text_len, dtype=torch.long)
        # 6 inner placeholder tokens (local_v * world_size = 6)
        ids[0, 1] = _VIS_START_ID
        ids[0, 8] = _VIS_END_ID
        visual = torch.randn(bs, local_v, hidden)
        text = torch.randn(bs, text_len, hidden)

        out = adapter(visual, text, ids)
        # fused_len == text_len == 10 (length-preserving); padded to 10; local = 5
        assert out.shape == (bs, 5, hidden)