vLLM/tests/v1/worker/test_gpu_model_runner.py

# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project

import gc
from contextlib import nullcontext
from types import SimpleNamespace
from unittest.mock import Mock

import numpy as np
import pytest
import torch

import vllm.model_executor.model_loader.reload.inplace as reload_inplace_module
import vllm.v1.worker.gpu_model_runner as gpu_model_runner_module
from vllm.config import (
    AttentionConfig,
    CacheConfig,
    ModelConfig,
    ParallelConfig,
    SchedulerConfig,
    VllmConfig,
    set_current_vllm_config,
)
from vllm.config.reasoning import ReasoningConfig
from vllm.distributed.parallel_state import (
    init_distributed_environment,
    initialize_model_parallel,
)
from vllm.lora.layers import LoRAMappingType
from vllm.lora.request import LoRARequest
from vllm.model_executor.layers.attention import Attention
from vllm.model_executor.layers.mamba.mamba_mixer2 import MambaMixer2
from vllm.multimodal.inputs import MultiModalFeatureSpec, PlaceholderRange
from vllm.platforms import current_platform
from vllm.platforms.interface import Platform
from vllm.sampling_params import SamplingParams
from vllm.utils.mem_constants import GiB_bytes
from vllm.utils.system_utils import update_environment_variables
from vllm.utils.torch_utils import set_random_seed
from vllm.v1.attention.backend import MultipleOf
from vllm.v1.attention.backends.mla.indexer import DeepseekV32IndexerBackend
from vllm.v1.attention.backends.mla.rocm_aiter_mla_sparse import (
    ROCMAiterMLASparseBackend,
)
from vllm.v1.attention.backends.registry import AttentionBackendEnum
from vllm.v1.core.kv_cache_utils import estimate_max_model_len, get_kv_cache_configs
from vllm.v1.core.sched.output import CachedRequestData, NewRequestData, SchedulerOutput
from vllm.v1.kv_cache_interface import (
    FullAttentionSpec,
    KVCacheConfig,
    KVCacheGroupSpec,
    KVCacheTensor,
)
from vllm.v1.outputs import EMPTY_MODEL_RUNNER_OUTPUT
from vllm.v1.sample.metadata import SamplingMetadata
from vllm.v1.spec_decode.metadata import SpecDecodeMetadata
from vllm.v1.worker.block_table import (
    MultiGroupBlockTable,
    SlotMappingMode,
    get_block_table_width,
)
from vllm.v1.worker.gpu.lora_utils import LoraState
from vllm.v1.worker.gpu.mm.encoder_cache import EncoderCache
from vllm.v1.worker.gpu.mm.lora import set_active_mm_loras
from vllm.v1.worker.gpu_input_batch import InputBatch
from vllm.v1.worker.gpu_model_runner import GPUModelRunner
from vllm.v1.worker.gpu_worker import Worker
from vllm.v1.worker.utils import select_common_block_size

BLOCK_SIZE = 16
NUM_BLOCKS = 10
DEVICE_TYPE = current_platform.device_type


@pytest.fixture(autouse=True)
def _restore_default_dtype():
    """Several tests here set the process-wide default dtype to float16 and
    previously leaked it, corrupting later float-sensitive tests in the same
    pytest process (torch.randn silently produced fp16)."""
    old = torch.get_default_dtype()
    yield
    torch.set_default_dtype(old)


def initialize_kv_cache(runner: GPUModelRunner):
    """Only perform necessary steps in GPUModelRunner.initialize_kv_cache()."""
    attn_spec = FullAttentionSpec(
        block_size=BLOCK_SIZE,
        num_kv_heads=runner.model_config.get_num_kv_heads(runner.parallel_config),
        head_size=runner.model_config.get_head_size(),
        dtype=runner.kv_cache_dtype,
    )
    tensor_size = attn_spec.page_size_bytes * NUM_BLOCKS
    kv_cache_config = KVCacheConfig(
        num_blocks=NUM_BLOCKS,
        kv_cache_tensors=[
            KVCacheTensor(
                size=tensor_size,
                layers=["layer.0"],
                layer_stride=tensor_size,
                block_stride=attn_spec.page_size_bytes,
            ),
        ],
        kv_cache_groups=[
            KVCacheGroupSpec(layer_names=["layer.0"], kv_cache_spec=attn_spec)
        ],
    )
    runner.kv_cache_config = kv_cache_config
    runner.input_batch = InputBatch(
        max_num_reqs=runner.max_num_reqs,
        max_model_len=runner.max_model_len,
        max_num_batched_tokens=runner.max_num_tokens,
        device=runner.device,
        vocab_size=runner.model_config.get_vocab_size(),
        block_sizes=[kv_cache_config.kv_cache_groups[0].kv_cache_spec.block_size],
        max_num_blocks_per_req=[NUM_BLOCKS],
    )
    runner.initialize_attn_backend(kv_cache_config)


def get_vllm_config():
    model_config = ModelConfig(
        model="facebook/opt-125m",
        dtype="float16",
        seed=42,
    )
    scheduler_config = SchedulerConfig(
        max_num_seqs=10,
        max_num_batched_tokens=512,
        max_model_len=512,
        is_encoder_decoder=model_config.is_encoder_decoder,
    )
    cache_config = CacheConfig(
        block_size=BLOCK_SIZE,
        gpu_memory_utilization=0.9,
        cache_dtype="auto",
    )
    cache_config.kv_cache_layout = "LBNHC"
    parallel_config = ParallelConfig()
    vllm_config = VllmConfig(
        model_config=model_config,
        cache_config=cache_config,
        scheduler_config=scheduler_config,
        parallel_config=parallel_config,
    )
    return vllm_config


@pytest.mark.parametrize("gc_initially_enabled", [True, False])
def test_freeze_gc_disables_and_restores_automatic_gc(
    monkeypatch: pytest.MonkeyPatch,
    gc_initially_enabled: bool,
):
    monkeypatch.setenv("VLLM_ENABLE_CUDAGRAPH_GC", "0")
    original_gc_state = gc.isenabled()
    try:
        if gc_initially_enabled:
            gc.enable()
        else:
            gc.disable()

        with (
            pytest.raises(RuntimeError, match="capture failed"),
            GPUModelRunner._freeze_gc(),
        ):
            assert not gc.isenabled()
            raise RuntimeError("capture failed")

        assert gc.isenabled() is gc_initially_enabled
    finally:
        if original_gc_state:
            gc.enable()
        else:
            gc.disable()


def test_prepare_padding_mask_marks_sequence_parallel_padding():
    runner = GPUModelRunner.__new__(GPUModelRunner)
    runner.is_padding = torch.empty(8, dtype=torch.bool)

    mask = runner._prepare_padding_mask(1, 8)

    assert mask.tolist() == [False, True, True, True, True, True, True, True]
    assert mask.data_ptr() == runner.is_padding.data_ptr()

    mask = runner._prepare_padding_mask(0, 8)

    assert mask.all()


@pytest.fixture
def model_runner():
    vllm_config = get_vllm_config()
    with set_current_vllm_config(vllm_config):
        model_config = vllm_config.model_config
        num_heads = model_config.get_num_kv_heads(vllm_config.parallel_config)
        head_size = model_config.get_head_size()
        vllm_config.compilation_config.static_forward_context["layer.0"] = Attention(
            num_heads, head_size, 0.1
        )
        runner = GPUModelRunner(vllm_config, DEVICE_TYPE)
        initialize_kv_cache(runner)
        yield runner


model_runner_2 = model_runner


def _schedule_new_request(*req_ids: str) -> SchedulerOutput:
    new_reqs = []
    num_scheduled_tokens = {}
    total_num_scheduled_tokens = 0
    for req_id in req_ids:
        new_reqs.append(
            NewRequestData(
                req_id=req_id,
                prompt_token_ids=[1, 2, 3],
                mm_features=[],
                sampling_params=SamplingParams(),
                pooling_params=None,
                block_ids=([0],),
                num_computed_tokens=0,
                lora_request=None,
            )
        )
        num_scheduled_tokens[req_id] = 3
        total_num_scheduled_tokens += num_scheduled_tokens[req_id]

    return SchedulerOutput(
        scheduled_new_reqs=new_reqs,
        scheduled_cached_reqs=CachedRequestData.make_empty(),
        num_scheduled_tokens=num_scheduled_tokens,
        total_num_scheduled_tokens=total_num_scheduled_tokens,
        scheduled_spec_decode_tokens={},
        scheduled_encoder_inputs={},
        finished_req_ids=set(),
        free_encoder_mm_hashes=[],
    )


def _schedule_cached_requests(
    req_ids: list[str],
    num_scheduled_tokens: dict[str, int],
    new_token_ids: list[list[int]],
    num_computed_tokens: list[int],
    num_output_tokens: list[int],
) -> SchedulerOutput:
    return SchedulerOutput(
        scheduled_new_reqs=[],
        scheduled_cached_reqs=CachedRequestData(
            req_ids=req_ids,
            resumed_req_ids=set(),
            new_token_ids=new_token_ids,
            all_token_ids={},
            new_block_ids=[None] * len(req_ids),
            num_computed_tokens=num_computed_tokens,
            num_output_tokens=num_output_tokens,
        ),
        num_scheduled_tokens=num_scheduled_tokens,
        total_num_scheduled_tokens=sum(num_scheduled_tokens.values()),
        scheduled_spec_decode_tokens={},
        scheduled_encoder_inputs={},
        finished_req_ids=set(),
        free_encoder_mm_hashes=[],
    )


def _is_req_scheduled(model_runner, req_id: str) -> bool:
    return req_id in model_runner.input_batch.req_id_to_index


def _is_req_added(model_runner, req_id: str) -> bool:
    return req_id in model_runner.requests


def _is_sampling_metadata_changed(
    model_runner, sampling_metadata_before: SamplingMetadata
):
    return model_runner.input_batch.sampling_metadata is not (sampling_metadata_before)


def _is_req_state_block_table_match(model_runner, req_id: str) -> bool:
    req_index = model_runner.input_batch.req_id_to_index[req_id]
    block_table = model_runner.input_batch.block_table[0]
    req_state = model_runner.requests[req_id]
    if block_table.num_blocks_per_row[req_index] != len(req_state.block_ids[0]):
        return False
    num_blocks = block_table.num_blocks_per_row[req_index]
    return (
        block_table.block_table.np[req_index, :num_blocks] == req_state.block_ids[0]
    ).all()


def _make_mock_backend_for_kernel_block_size(
    supported_sizes: list[int | MultipleOf],
):
    class _MockBackend:
        @staticmethod
        def get_name():
            return "MOCK"

        @staticmethod
        def get_supported_kernel_block_sizes(kv_cache_spec=None):
            return supported_sizes

    return _MockBackend()


def _make_kv_cache_spec() -> FullAttentionSpec:
    return FullAttentionSpec(block_size=1, num_kv_heads=1, head_size=1, dtype="float16")


def test_select_common_block_size_prefers_manager_block_size():
    backend_a = _make_mock_backend_for_kernel_block_size([MultipleOf(32)])
    backend_b = _make_mock_backend_for_kernel_block_size([64, MultipleOf(16)])

    selected_size = select_common_block_size(128, [backend_a, backend_b])
    assert selected_size == 128


def test_select_common_block_size_uses_largest_shared_int():
    backend_a = _make_mock_backend_for_kernel_block_size([128, 64])
    backend_b = _make_mock_backend_for_kernel_block_size([64, 32])

    selected_size = select_common_block_size(256, [backend_a, backend_b])
    assert selected_size == 64


def test_select_common_block_size_without_active_backends_uses_manager_size():
    assert select_common_block_size(256, []) == 256


def test_select_common_block_size_accepts_rocm_sparse_block_size_16(monkeypatch):
    monkeypatch.setattr(current_platform, "is_rocm", lambda: True)

    selected_size = select_common_block_size(
        16,
        [DeepseekV32IndexerBackend, ROCMAiterMLASparseBackend],
    )
    assert selected_size == 16


def test_reasoning_config_without_custom_logitsprocs_does_not_need_output_token_ids(
    dist_init,
):
    vllm_config = get_vllm_config()
    assert vllm_config.model_config.logits_processors is None
    reasoning_config = ReasoningConfig(
        reasoning_start_str="<think>", reasoning_end_str="</think>"
    )
    reasoning_config._reasoning_start_token_ids = [1]
    reasoning_config._reasoning_end_token_ids = [2]
    vllm_config.reasoning_config = reasoning_config

    with set_current_vllm_config(vllm_config):
        model_config = vllm_config.model_config
        num_heads = model_config.get_num_kv_heads(vllm_config.parallel_config)
        head_size = model_config.get_head_size()
        vllm_config.compilation_config.static_forward_context["layer.0"] = Attention(
            num_heads, head_size, 0.1
        )
        runner = GPUModelRunner(vllm_config, torch.device("cpu"))

    assert runner.input_batch.thinking_budget_state_holder is not None
    assert runner.input_batch.logitsprocs_need_output_token_ids is False


@pytest.mark.skip_global_cleanup
@pytest.mark.parametrize(
    ("world_size", "is_last_rank", "expected_calls"),
    [(1, True, 0), (2, True, 0), (2, False, 1)],
)
def test_sample_tokens_receives_pp_sampled_ids_only_on_non_last_rank(
    monkeypatch: pytest.MonkeyPatch,
    world_size: int,
    is_last_rank: bool,
    expected_calls: int,
):
    runner = GPUModelRunner.__new__(GPUModelRunner)
    runner.execute_model_state = None
    runner.kv_connector_output = None
    runner.use_async_scheduling = True
    receive_calls = 0

    def receive_prev_sampled_token_ids():
        nonlocal receive_calls
        receive_calls += 1

    runner._pp_receive_prev_sampled_token_ids_to_input_batch = (
        receive_prev_sampled_token_ids
    )
    monkeypatch.setattr(
        gpu_model_runner_module,
        "get_pp_group",
        lambda: SimpleNamespace(world_size=world_size, is_last_rank=is_last_rank),
    )

    output = GPUModelRunner.sample_tokens(runner, None)
    assert output in (EMPTY_MODEL_RUNNER_OUTPUT, None)
    assert receive_calls == expected_calls


@pytest.mark.skip_global_cleanup
def test_sample_tokens_skips_pp_group_lookup_without_async_scheduling(
    monkeypatch: pytest.MonkeyPatch,
):
    runner = GPUModelRunner.__new__(GPUModelRunner)
    runner.execute_model_state = None
    runner.kv_connector_output = None
    runner.use_async_scheduling = False

    monkeypatch.setattr(
        gpu_model_runner_module,
        "get_pp_group",
        pytest.fail,
    )

    output = GPUModelRunner.sample_tokens(runner, None)
    assert output in (EMPTY_MODEL_RUNNER_OUTPUT, None)


def test_select_common_block_size_no_valid_option():
    backend_a = _make_mock_backend_for_kernel_block_size([64])
    backend_b = _make_mock_backend_for_kernel_block_size([MultipleOf(16)])

    with pytest.raises(ValueError):
        select_common_block_size(48, [backend_a, backend_b])


def _mock_backend(supported: list, *, exact: bool = False):
    """Backend accepting multiples of ``supported``, or only those exact sizes
    when ``exact`` (as CPU_MLA does)."""
    from vllm.v1.attention.backend import AttentionBackend

    class _MockBackendCls(AttentionBackend):
        @staticmethod
        def get_name() -> str:
            return "MOCK_EXACT" if exact else "MOCK"

        @staticmethod
        def get_supported_kernel_block_sizes(kv_cache_spec=None):
            return list(supported)

        if exact:

            @classmethod
            def supports_block_size(cls, block_size: int | None) -> bool:
                return block_size is None or block_size in supported

    return _MockBackendCls


@pytest.mark.parametrize(
    "backends,expected",
    [
        # A lone backend keeps the default; MultipleOf(8) makes a
        # minimum-based regression return 8 instead of 16.
        ([[MultipleOf(8)]], 16),
        # Sparse-MLA main backend (32 or 64) beside an indexer supporting only
        # 64: the main backend alone picks 32, which the indexer rejects at
        # select_common_block_size time.
        ([[32, 64], [64]], 64),
        ([[32, 64]], 32),
        # Neither divides the other, so only the LCM satisfies both.
        ([[32], [48]], 96),
    ],
)
def test_preferred_block_size_satisfies_every_backend(backends, expected):
    classes = [_mock_backend(s) for s in backends]
    assert Platform._preferred_block_size_for_backends(classes, 16, None) == expected


def test_preferred_block_size_searches_past_an_exact_size_backend():
    # Extending greedily picks lcm(16, 32) = 32, which the exact backend
    # rejects; 96 is accepted by both.
    classes = [_mock_backend([16, 96], exact=True), _mock_backend([MultipleOf(32)])]
    assert Platform._preferred_block_size_for_backends(classes, 16, None) == 96


def test_preferred_block_size_rejects_backends_with_no_common_size():
    classes = [_mock_backend([16], exact=True), _mock_backend([MultipleOf(64)])]
    with pytest.raises(ValueError, match="share no supported KV cache block size"):
        Platform._preferred_block_size_for_backends(classes, 16, None)


@pytest.mark.parametrize("cpu", [False, True])
def test_alignment_rejected_by_a_sibling_backend_raises(monkeypatch, cpu):
    # Alignment that lands on a size an exact-size sibling rejects must fail
    # here, not later in select_common_block_size().
    backends = [_mock_backend([MultipleOf(16)]), _mock_backend([16], exact=True)]
    monkeypatch.setattr(
        Platform, "_find_non_ssm_backends", classmethod(lambda cls, c: backends)
    )

    def bump(cls, vllm_config, backend_cls):
        vllm_config.cache_config.block_size = 32

    monkeypatch.setattr(Platform, "_align_hybrid_block_size", classmethod(bump))
    cache_config = SimpleNamespace(
        block_size=16, user_specified_block_size=True, kv_cache_dtype_skip_layers=None
    )
    vllm_config = SimpleNamespace(
        cache_config=cache_config, model_config=SimpleNamespace(is_hybrid=True)
    )
    from vllm.platforms.cpu import CpuPlatform  # overrides the method; must check too

    platform = CpuPlatform if cpu else Platform
    with pytest.raises(ValueError, match="MOCK_EXACT"):
        platform.update_block_size_for_backend(vllm_config)


def test_xpu_gdn_rounding_rejected_by_a_sibling_backend_raises(monkeypatch):
    # XPU's GDN rounding runs after super()'s check, so it must check again.
    try:
        from vllm.platforms.xpu import XPUPlatform
    except ImportError:
        pytest.skip("vllm_xpu_kernels not importable outside an XPU stack")

    backends = [_mock_backend([MultipleOf(16)]), _mock_backend([16], exact=True)]
    monkeypatch.setattr(
        Platform, "update_block_size_for_backend", classmethod(lambda cls, c: None)
    )
    monkeypatch.setattr(
        Platform, "_find_non_ssm_backends", classmethod(lambda cls, c: backends)
    )
    gdn_layer = SimpleNamespace(
        get_attn_backend=lambda: SimpleNamespace(get_name=lambda: "GDN_ATTN")
    )
    monkeypatch.setattr(
        "vllm.config.vllm.get_layers_from_vllm_config",
        lambda vllm_config, layer_type: {"gdn": gdn_layer},
    )
    cache_config = SimpleNamespace(
        block_size=16,
        mamba_cache_mode="none",
        mamba_page_size_padded=None,
        user_specified_block_size=True,
        kv_cache_dtype_skip_layers=None,
    )
    vllm_config = SimpleNamespace(
        cache_config=cache_config, model_config=SimpleNamespace(is_hybrid=True)
    )
    with pytest.raises(ValueError, match="MOCK_EXACT"):
        XPUPlatform.update_block_size_for_backend(vllm_config)


def test_set_active_mm_loras_builds_tower_and_connector_mappings():
    model = Mock()
    model.get_mm_lora_token_counts.side_effect = (
        lambda *, modality, mm_kwargs, num_mm_embeds: (
            num_mm_embeds + 1,
            num_mm_embeds + 11,
        )
    )
    model.get_mm_mapping.return_value = SimpleNamespace(connector=True)

    lora_manager = Mock()
    lora_manager.supports_tower_connector_lora.return_value = True

    encoder_cache = EncoderCache()
    encoder_cache.mm_features["req-with-lora"] = [
        MultiModalFeatureSpec(
            data=None,
            modality="image",
            identifier="img-0",
            mm_position=PlaceholderRange(offset=0, length=2),
        ),
        MultiModalFeatureSpec(
            data=None,
            modality="image",
            identifier="img-1",
            mm_position=PlaceholderRange(offset=2, length=3),
        ),
    ]
    encoder_cache.mm_features["req-no-lora"] = [
        MultiModalFeatureSpec(
            data=None,
            modality="image",
            identifier="img-2",
            mm_position=PlaceholderRange(offset=0, length=1),
        )
    ]

    lora_state = LoraState(max_num_reqs=4)
    lora_request = LoRARequest("vision-lora", 7, "/tmp/vision-lora")
    lora_state.add_request("req-with-lora", 0, lora_request)
    lora_state.add_request("req-no-lora", 1, None)

    set_active_mm_loras(
        model=model,
        lora_manager=lora_manager,
        encoder_cache=encoder_cache,
        req_id_to_index={
            "req-with-lora": 0,
            "req-no-lora": 1,
        },
        lora_state=lora_state,
        scheduled_encoder_inputs={
            "req-with-lora": [1, 0],
            "req-no-lora": [0],
            "missing-req": [0],
        },
    )

    assert lora_manager.set_active_adapters.call_count == 2

    tower_requests, tower_mapping = lora_manager.set_active_adapters.call_args_list[
        0
    ].args
    assert tower_requests == {lora_request}
    assert tower_mapping.type is LoRAMappingType.TOWER
    assert tower_mapping.prompt_mapping == (7, 7, 0)
    assert tower_mapping.index_mapping == (7, 7, 7, 7, 7, 7, 7, 0, 0)

    connector_requests, connector_mapping = (
        lora_manager.set_active_adapters.call_args_list[1].args
    )
    assert connector_requests == {lora_request}
    assert connector_mapping.type is LoRAMappingType.CONNECTOR
    assert connector_mapping.prompt_mapping == (7, 7, 0)
    assert connector_mapping.index_mapping == ((7,) * 14 + (7,) * 13 + (0,) * 12)


def test_set_active_mm_loras_follows_modality_sorted_encoder_order():
    model = Mock()
    model.get_mm_lora_token_counts.side_effect = (
        lambda *, modality, mm_kwargs, num_mm_embeds: (num_mm_embeds, None)
    )
    lora_manager = Mock()
    lora_manager.supports_tower_connector_lora.return_value = True

    encoder_cache = EncoderCache()
    for req_id, modality, length in [("req-a", "video", 1), ("req-b", "image", 2)]:
        encoder_cache.mm_features[req_id] = [
            MultiModalFeatureSpec(
                data=None,
                modality=modality,
                identifier=f"{req_id}-0",
                mm_position=PlaceholderRange(offset=0, length=length),
            )
        ]
    lora_state = LoraState(max_num_reqs=2)
    lora_request = LoRARequest("vision-lora", 7, "/tmp/vision-lora")
    lora_state.add_request("req-a", 0, lora_request)
    lora_state.add_request("req-b", 1, None)

    set_active_mm_loras(
        model=model,
        lora_manager=lora_manager,
        encoder_cache=encoder_cache,
        req_id_to_index={"req-a": 0, "req-b": 1},
        lora_state=lora_state,
        scheduled_encoder_inputs={"req-a": [0], "req-b": [0]},
    )

    # The image is encoded before the video, so its tokens come first.
    _, tower_mapping = lora_manager.set_active_adapters.call_args_list[0].args
    assert tower_mapping.index_mapping == (0, 0, 7)


def test_batch_mm_inputs_from_scheduler_sorts_by_modality():
    features = {
        f"req{i}": MultiModalFeatureSpec(
            data=Mock(),
            modality=m,
            identifier=f"hash{i}",
            mm_position=PlaceholderRange(offset=i, length=1),
        )
        for i, m in enumerate(["video", "image", "video"])
    }
    runner = SimpleNamespace(
        requests={r: SimpleNamespace(mm_features=[f]) for r, f in features.items()}
    )
    scheduler_output = SimpleNamespace(
        scheduled_encoder_inputs={r: [0] for r in features}
    )

    mm_hashes, mm_kwargs, mm_lora_refs = GPUModelRunner._batch_mm_inputs_from_scheduler(
        runner, scheduler_output
    )

    order = ["req1", "req0", "req2"]
    assert mm_hashes == [features[r].identifier for r in order]
    assert mm_kwargs == [(features[r].modality, features[r].data) for r in order]
    assert mm_lora_refs == [(r, features[r].mm_position) for r in order]


def test_update_states_new_request(model_runner, dist_init):
    req_id = "req_0"

    # new req
    scheduler_output = _schedule_new_request(req_id)

    metadata_before = model_runner.input_batch.sampling_metadata
    model_runner._update_states(scheduler_output)
    assert _is_sampling_metadata_changed(model_runner, metadata_before)
    assert _is_req_added(model_runner, req_id)
    assert _is_req_scheduled(model_runner, req_id)
    assert _is_req_state_block_table_match(model_runner, req_id)


def test_update_states_request_finished(model_runner, dist_init):
    req_id = "req_0"

    # new req
    scheduler_output = _schedule_new_request(req_id)

    model_runner._update_states(scheduler_output)
    assert _is_req_added(model_runner, req_id)
    assert _is_req_scheduled(model_runner, req_id)

    # finish req
    scheduler_output = SchedulerOutput(
        scheduled_new_reqs=[],
        scheduled_cached_reqs=CachedRequestData.make_empty(),
        num_scheduled_tokens={},
        total_num_scheduled_tokens=0,
        scheduled_spec_decode_tokens={},
        scheduled_encoder_inputs={},
        finished_req_ids={req_id},
        free_encoder_mm_hashes=[],
    )

    metadata_before = model_runner.input_batch.sampling_metadata
    model_runner._update_states(scheduler_output)
    assert _is_sampling_metadata_changed(model_runner, metadata_before)
    assert not _is_req_added(model_runner, req_id)
    assert not _is_req_scheduled(model_runner, req_id)


def test_update_states_request_resumed(model_runner, dist_init):
    req_id = "req_0"

    # new req
    scheduler_output = _schedule_new_request(req_id)

    model_runner._update_states(scheduler_output)
    assert _is_req_added(model_runner, req_id)
    assert _is_req_scheduled(model_runner, req_id)

    # unschedule req
    scheduler_output = SchedulerOutput(
        scheduled_new_reqs=[],
        scheduled_cached_reqs=CachedRequestData.make_empty(),
        num_scheduled_tokens={},
        total_num_scheduled_tokens=0,
        scheduled_spec_decode_tokens={},
        scheduled_encoder_inputs={},
        finished_req_ids=set(),
        free_encoder_mm_hashes=[],
    )

    model_runner._update_states(scheduler_output)
    assert _is_req_added(model_runner, req_id)
    assert not _is_req_scheduled(model_runner, req_id)

    # resume req
    cached_req_data = CachedRequestData(
        req_ids=[req_id],
        resumed_req_ids=set(),
        new_token_ids=[[]],
        all_token_ids={},
        new_block_ids=[([0],)],
        num_computed_tokens=[0],
        num_output_tokens=[0],
    )

    scheduler_output = SchedulerOutput(
        scheduled_new_reqs=[],
        scheduled_cached_reqs=cached_req_data,
        num_scheduled_tokens={req_id: 1},
        total_num_scheduled_tokens=1,
        scheduled_spec_decode_tokens={},
        scheduled_encoder_inputs={},
        finished_req_ids=set(),
        free_encoder_mm_hashes=[],
    )

    metadata_before = model_runner.input_batch.sampling_metadata
    model_runner._update_states(scheduler_output)
    assert _is_sampling_metadata_changed(model_runner, metadata_before)
    assert _is_req_added(model_runner, req_id)
    assert _is_req_scheduled(model_runner, req_id)
    assert _is_req_state_block_table_match(model_runner, req_id)


def test_get_nans_in_logits(model_runner, dist_init, monkeypatch):
    monkeypatch.setattr(gpu_model_runner_module.envs, "VLLM_RAISE_ON_LOGIT_NANS", False)
    req_ids = ("req_0", "req_1")

    scheduler_output = _schedule_new_request(*req_ids)
    model_runner._update_states(scheduler_output)

    logits = torch.tensor(
        [
            [1.0, 2.0, 3.0],
            [3.0, 2.0, 1.0],
        ],
        device=DEVICE_TYPE,
    )
    result = model_runner._get_nans_in_logits(logits)
    assert result == {"req_0": 0, "req_1": 0}

    logits = torch.tensor(
        [
            [1.0, float("nan"), 3.0],
            [4.0, float("nan"), float("nan")],
        ],
        device=DEVICE_TYPE,
    )
    result = model_runner._get_nans_in_logits(logits)
    assert result == {"req_0": 1, "req_1": 2}

    logits = torch.tensor(
        [
            [1.0, 2.0, 3.0],
            [4.0, float("nan"), float("nan")],
        ],
        device=DEVICE_TYPE,
    )
    result = model_runner._get_nans_in_logits(logits)
    assert result == {"req_0": 0, "req_1": 2}

    result = model_runner._get_nans_in_logits(logits=None)
    assert result == {"req_0": 0, "req_1": 0}

    logits = torch.tensor(
        [
            [1.0, float("nan"), 3.0],
        ],
        device=DEVICE_TYPE,
    )
    result = model_runner._get_nans_in_logits(logits)
    assert result == {"req_0": 1, "req_1": 0}

    logits = torch.tensor(
        [
            [float("nan"), float("nan"), 2.0],
            [1.0, 2.0, 3.0],
            [float("nan"), 2.0, 3.0],
        ],
        device=DEVICE_TYPE,
    )
    result = model_runner._get_nans_in_logits(logits)
    assert result == {"req_0": 2, "req_1": 0}


def test_update_states_no_changes(model_runner, dist_init):
    req_id = "req_0"

    # new req
    scheduler_output = _schedule_new_request(req_id)

    model_runner._update_states(scheduler_output)
    assert _is_req_added(model_runner, req_id)
    assert _is_req_scheduled(model_runner, req_id)

    # schedule req
    scheduler_output = SchedulerOutput(
        scheduled_new_reqs=[],
        scheduled_cached_reqs=CachedRequestData.make_empty(),
        num_scheduled_tokens={req_id: 1},
        total_num_scheduled_tokens=1,
        scheduled_spec_decode_tokens={},
        scheduled_encoder_inputs={},
        finished_req_ids=set(),
        free_encoder_mm_hashes=[],
    )

    metadata_before = model_runner.input_batch.sampling_metadata
    model_runner._update_states(scheduler_output)
    assert not _is_sampling_metadata_changed(model_runner, metadata_before)
    assert _is_req_added(model_runner, req_id)
    assert _is_req_scheduled(model_runner, req_id)
    assert _is_req_state_block_table_match(model_runner, req_id)


def test_update_states_request_unscheduled(model_runner, dist_init):
    req_ids = ("req_0", "req_1")

    # new reqs
    scheduler_output = _schedule_new_request(*req_ids)

    model_runner._update_states(scheduler_output)

    assert _is_req_added(model_runner, req_ids[0])
    assert _is_req_scheduled(model_runner, req_ids[0])

    assert _is_req_added(model_runner, req_ids[1])
    assert _is_req_scheduled(model_runner, req_ids[1])

    # unschedule req_1
    scheduler_output = SchedulerOutput(
        scheduled_new_reqs=[],
        scheduled_cached_reqs=CachedRequestData.make_empty(),
        num_scheduled_tokens={req_ids[0]: 1},
        total_num_scheduled_tokens=1,
        scheduled_spec_decode_tokens={},
        scheduled_encoder_inputs={},
        finished_req_ids=set(),
        free_encoder_mm_hashes=[],
    )

    metadata_before = model_runner._update_states(scheduler_output)
    assert _is_sampling_metadata_changed(model_runner, metadata_before)

    assert _is_req_added(model_runner, req_ids[0])
    assert _is_req_scheduled(model_runner, req_ids[0])

    assert _is_req_added(model_runner, req_ids[1])
    assert not _is_req_scheduled(model_runner, req_ids[1])


def test_update_states_pp_non_async_multi_request_keeps_token_buffers_consistent(
    model_runner, model_runner_2, dist_init, monkeypatch
):
    req_ids = ["req_0", "req_1"]
    non_last_runner = model_runner
    last_runner = model_runner_2
    non_last_runner.use_async_scheduling = False
    last_runner.use_async_scheduling = False

    # Both ranks start from the same request set.
    monkeypatch.setattr(
        "vllm.v1.worker.gpu_model_runner.get_pp_group",
        lambda: SimpleNamespace(is_last_rank=False, world_size=2),
    )
    non_last_runner._update_states(_schedule_new_request(*req_ids))
    last_runner._update_states(_schedule_new_request(*req_ids))

    sampled_by_last_rank = {req_ids[0]: 101, req_ids[1]: 201}
    # Emulate last-rank bookkeeping result from previous step:
    # sampled tokens already cached in CPU token buffers.
    for req_id, token_id in sampled_by_last_rank.items():
        req_index = last_runner.input_batch.req_id_to_index[req_id]
        start_idx = int(last_runner.input_batch.num_tokens_no_spec[req_index])
        end_idx = start_idx + 1
        last_runner.input_batch.token_ids_cpu[req_index, start_idx:end_idx] = [token_id]
        last_runner.input_batch.is_token_ids[req_index, start_idx:end_idx] = True
        last_runner.input_batch.num_tokens_no_spec[req_index] = end_idx
        last_runner.requests[req_id].output_token_ids.append(token_id)

    scheduler_output = _schedule_cached_requests(
        req_ids=req_ids,
        num_scheduled_tokens={req_ids[0]: 1, req_ids[1]: 1},
        new_token_ids=[[101], [201]],
        num_computed_tokens=[3, 3],  # prompt tokens only
        num_output_tokens=[1, 1],
    )
    # non-last rank appends new_token_ids in _update_states.
    monkeypatch.setattr(
        "vllm.v1.worker.gpu_model_runner.get_pp_group",
        lambda: SimpleNamespace(is_last_rank=False, world_size=2),
    )
    non_last_runner._update_states(scheduler_output)
    # last rank should keep its already-bookkept CPU buffers unchanged.
    monkeypatch.setattr(
        "vllm.v1.worker.gpu_model_runner.get_pp_group",
        lambda: SimpleNamespace(is_last_rank=True, world_size=2),
    )
    last_runner._update_states(scheduler_output)

    # Verify consistency between PP ranks after _update_states.
    for req_id in req_ids:
        non_last_idx = non_last_runner.input_batch.req_id_to_index[req_id]
        last_idx = last_runner.input_batch.req_id_to_index[req_id]
        non_last_len = int(non_last_runner.input_batch.num_tokens_no_spec[non_last_idx])
        last_len = int(last_runner.input_batch.num_tokens_no_spec[last_idx])
        assert non_last_len == last_len
        assert (
            non_last_runner.input_batch.token_ids_cpu[
                non_last_idx, :non_last_len
            ].tolist()
            == last_runner.input_batch.token_ids_cpu[last_idx, :last_len].tolist()
        )


def test_update_states_pp_async_multi_request_keeps_rank_state_consistent(
    model_runner, model_runner_2, dist_init, monkeypatch
):
    req_ids = ["req_0", "req_1"]
    non_last_runner = model_runner
    last_runner = model_runner_2
    non_last_runner.use_async_scheduling = True
    last_runner.use_async_scheduling = True

    # Both ranks start from the same request set.
    monkeypatch.setattr(
        "vllm.v1.worker.gpu_model_runner.get_pp_group",
        lambda: SimpleNamespace(is_last_rank=False, world_size=2),
    )
    non_last_runner._update_states(_schedule_new_request(*req_ids))
    last_runner._update_states(_schedule_new_request(*req_ids))

    # Simulate async previous-step sampled tokens known on both ranks.
    # non-last rank may receive them via PP communication; last rank has
    # them from local sampling/bookkeeping.
    sampled_by_last_rank = {req_ids[0]: 111, req_ids[1]: 222}
    for runner in (non_last_runner, last_runner):
        for req_id, token_id in sampled_by_last_rank.items():
            req_index = runner.input_batch.req_id_to_index[req_id]
            start_idx = int(runner.input_batch.num_tokens_no_spec[req_index])
            end_idx = start_idx + 1
            runner.input_batch.token_ids_cpu[req_index, start_idx:end_idx] = [token_id]
            runner.input_batch.is_token_ids[req_index, start_idx:end_idx] = True
            runner.input_batch.num_tokens_no_spec[req_index] = end_idx
            runner.requests[req_id].output_token_ids.append(token_id)

    scheduler_output = _schedule_cached_requests(
        req_ids=req_ids,
        num_scheduled_tokens={req_ids[0]: 1, req_ids[1]: 1},
        new_token_ids=[],
        num_computed_tokens=[4, 4],
        num_output_tokens=[1, 1],
    )
    # non-last rank: async PP branch (new_token_ids empty).
    monkeypatch.setattr(
        "vllm.v1.worker.gpu_model_runner.get_pp_group",
        lambda: SimpleNamespace(is_last_rank=False, world_size=2),
    )
    non_last_runner._update_states(scheduler_output)
    # last rank: keep already-bookkept state aligned with scheduler view.
    monkeypatch.setattr(
        "vllm.v1.worker.gpu_model_runner.get_pp_group",
        lambda: SimpleNamespace(is_last_rank=True, world_size=2),
    )
    last_runner._update_states(scheduler_output)

    for req_id in req_ids:
        non_last_idx = non_last_runner.input_batch.req_id_to_index[req_id]
        last_idx = last_runner.input_batch.req_id_to_index[req_id]
        non_last_len = int(non_last_runner.input_batch.num_tokens_no_spec[non_last_idx])
        last_len = int(last_runner.input_batch.num_tokens_no_spec[last_idx])
        assert non_last_len == last_len
        assert (
            non_last_runner.input_batch.token_ids_cpu[
                non_last_idx, :non_last_len
            ].tolist()
            == last_runner.input_batch.token_ids_cpu[last_idx, :last_len].tolist()
        )


def _update_config(model_runner: GPUModelRunner, overrides: dict) -> None:
    # Apply overrides via the worker without constructing a full worker.
    worker = object.__new__(Worker)
    worker.vllm_config = model_runner.vllm_config
    worker.model_runner = model_runner
    worker.update_config(overrides)


def test_update_config(model_runner):
    # Simple update
    _update_config(model_runner, {"load_config": {"load_format": "dummy"}})
    assert model_runner.load_config.load_format == "dummy"
    assert model_runner.vllm_config.load_config is model_runner.load_config
    # Raise error on non-existing config
    with pytest.raises(ValueError, match="do_not_exist_config"):
        _update_config(model_runner, {"do_not_exist_config": "dummy"})


def test_load_model_weights_inplace(dist_init, model_runner, model_runner_2):
    # In this test, model_runner loads model + weights in one go, while
    # model_runner_2 loads dummy weights first then load real weights inplace
    model_runner.load_model()
    original_load_format = model_runner_2.load_config.load_format
    _update_config(model_runner_2, {"load_config": {"load_format": "dummy"}})
    model_runner_2.load_model()  # Initial model loading with dummy weights
    assert str(model_runner.get_model().state_dict()) != str(
        model_runner_2.get_model().state_dict()
    )
    _update_config(
        model_runner_2, {"load_config": {"load_format": original_load_format}}
    )
    model_runner_2.reload_weights()  # Load real weights inplace
    assert str(model_runner.get_model().state_dict()) == str(
        model_runner_2.get_model().state_dict()
    )


def test_reload_weights_before_load_model(model_runner):
    with pytest.raises(ValueError):
        model_runner.reload_weights()


def test_reload_weights_path_replaces_object_storage_source(monkeypatch):
    # An engine started from an object-storage URI keeps that URI in
    # model_config.model_weights, which the runai_streamer loader prefers over
    # model_config.model. A new weights_path must replace it, or the loader
    # silently re-streams the original checkpoint.
    loader = Mock()
    loader.get_all_weights.return_value = iter(())
    monkeypatch.setattr(reload_inplace_module, "get_model_loader", lambda _: loader)
    monkeypatch.setattr(reload_inplace_module, "initialize_layerwise_reload", Mock())
    monkeypatch.setattr(reload_inplace_module, "finalize_layerwise_reload", Mock())

    vllm_config = SimpleNamespace(
        lora_config=None,
        load_config=SimpleNamespace(load_format="runai_streamer"),
        model_config=SimpleNamespace(
            model="/tmp/pulled-config-files",
            model_weights="s3://bucket/original",
            revision="abc123",
            quantization=None,
        ),
    )
    model = Mock()
    model.named_parameters.return_value = []
    model.load_weights.return_value = None

    reload_inplace_module.reload_weights(
        vllm_config,
        model,
        weights_iterator=None,
        weights_path="org/new-model",
        is_checkpoint_format=True,
    )

    loader.get_all_weights.assert_called_once_with(vllm_config.model_config, model)
    cfg = vllm_config.model_config
    assert (cfg.model, cfg.model_weights, cfg.revision) == ("org/new-model", "", None)


def test_sample_passes_reordered_draft_probs_to_rejection_sampler():
    runner = object.__new__(GPUModelRunner)
    runner.use_async_scheduling = False
    runner.input_batch = SimpleNamespace(
        sampling_metadata=Mock(spec=SamplingMetadata),
        update_async_output_token_ids=Mock(),
        req_ids=["req_a", "req_b", "req_c"],
    )
    runner.rejection_sampler = Mock(return_value="sampler_output")
    runner.sampler = Mock()
    runner._draft_prob_req_ids = ["req_c", "req_a", "req_b"]
    runner._draft_probs = torch.arange(3 * 3 * 4, dtype=torch.float32).reshape(3, 3, 4)

    spec_decode_metadata = SpecDecodeMetadata.make_dummy(
        [[1, 2], [], [3]],
        device=torch.device("cpu"),
    )
    logits = torch.randn(6, 4)

    output = GPUModelRunner._sample(runner, logits, spec_decode_metadata)

    assert output == "sampler_output"
    passed_draft_probs = runner.rejection_sampler.call_args.args[1]
    expected_draft_probs = torch.cat(
        [
            runner._draft_probs[1, :2],
            runner._draft_probs[0, :1],
        ],
        dim=0,
    )
    assert torch.equal(passed_draft_probs, expected_draft_probs)


def test_dummy_sampler_run_warms_all_greedy_rejection_sampler(monkeypatch):
    runner = object.__new__(GPUModelRunner)
    runner.device = torch.device("cpu")
    runner.vllm_config = SimpleNamespace(
        model_config=SimpleNamespace(multimodal_config=None)
    )
    runner.model = SimpleNamespace(compute_logits=Mock(return_value=torch.randn(3, 8)))
    runner.sampler = Mock(return_value="sampler_output")
    runner.sampler.logprobs_mode = "processed_logprobs"
    runner.speculative_config = SimpleNamespace(
        rejection_sample_method="standard",
        draft_sample_method="greedy",
    )
    runner.rejection_sampler = Mock()
    synchronize = Mock()
    monkeypatch.setattr(torch.accelerator, "synchronize", synchronize)

    output = GPUModelRunner._dummy_sampler_run(runner, torch.randn(3, 4))

    assert output == "sampler_output"
    assert runner.rejection_sampler.call_count == 2
    mixed_metadata = runner.rejection_sampler.call_args_list[0].args[3]
    all_greedy_metadata = runner.rejection_sampler.call_args_list[1].args[3]
    assert not mixed_metadata.all_greedy
    assert mixed_metadata.temperature is not None
    assert all_greedy_metadata.all_greedy
    assert not all_greedy_metadata.all_random
    assert all_greedy_metadata.temperature is None
    synchronize.assert_called_once_with()


def test_invalid_draft_suffixes_remain_rejected_in_metadata():
    runner = object.__new__(GPUModelRunner)
    runner.device = torch.device("cpu")
    runner.arange_np = np.arange(64, dtype=np.int64)
    runner._arange_scratch = np.empty(64, dtype=np.int64)
    # Placeholder (-1) drafts are kept in input_ids (clamped to 0 only at the
    # embedding boundary). For num_draft_tokens=[2, 1, 2] the draft positions
    # are [1, 2, 4, 6, 7], so the gather carries the -1s straight into the
    # rejection-sampling metadata.
    runner.input_ids = SimpleNamespace(
        gpu=torch.tensor([99, 10, -1, 99, 12, 99, 13, -1], dtype=torch.int32),
    )

    metadata = GPUModelRunner._calc_spec_decode_metadata(
        runner,
        np.array([2, 1, 2], dtype=np.int32),
        np.array([3, 5, 8], dtype=np.int32),
    )

    assert metadata.draft_token_ids.tolist() == [10, -1, 12, 13, -1]


def test_init_kv_cache_with_kv_sharing_invalid_target_layer_order(default_vllm_config):
    torch.set_default_dtype(torch.float16)
    layer_0 = "model.layers.0.self_attn.attn"
    layer_1 = "model.layers.1.self_attn.attn"
    error_msg = f"{layer_1} must come before the current layer"
    with pytest.raises(ValueError, match=error_msg):
        fwd_context = {
            # initialization below will fail because target layer is invalid;
            # the target layer needs to come before layer 1
            layer_0: Attention(
                num_heads=8,
                head_size=64,
                scale=1.0,
                prefix=layer_0,
                kv_sharing_target_layer_name=layer_1,
            ),
            layer_1: Attention(
                num_heads=8,
                head_size=64,
                scale=1.0,
                prefix=layer_1,
            ),
        }
        # suppress var not used error
        assert fwd_context is not None


def test_init_kv_cache_with_kv_sharing_target_layer_not_exist(default_vllm_config):
    torch.set_default_dtype(torch.float16)
    layer_0 = "model.layers.0.self_attn.attn"
    layer_1 = "model.layers.1.self_attn.attn"
    invalid_layer = "model.layers.0.cross_attn.attn"
    error_msg = f"{invalid_layer} is not a valid Attention layer in the model"
    with pytest.raises(ValueError, match=error_msg):
        fwd_context = {
            layer_0: Attention(
                num_heads=8,
                head_size=64,
                scale=1.0,
                prefix=layer_0,
            ),
            layer_1: Attention(
                num_heads=8,
                head_size=64,
                scale=1.0,
                prefix=layer_1,
                # invalid layer: cross_attn.atn doesn't exist!
                kv_sharing_target_layer_name=invalid_layer,
            ),
        }
        # suppress var not used error
        assert fwd_context is not None


def test_init_kv_cache_with_kv_sharing_target_same_as_current(default_vllm_config):
    torch.set_default_dtype(torch.float16)
    layer_0 = "model.layers.0.self_attn.attn"
    layer_1 = "model.layers.1.self_attn.attn"
    error_msg = f"{layer_1} cannot be the same as the current layer"
    with pytest.raises(ValueError, match=error_msg):
        fwd_context = {
            # initialization below will fail because target layer is invalid;
            # the target layer needs to come before layer 1
            layer_0: Attention(
                num_heads=8,
                head_size=64,
                scale=1.0,
                prefix=layer_0,
            ),
            layer_1: Attention(
                num_heads=8,
                head_size=64,
                scale=1.0,
                prefix=layer_1,
                kv_sharing_target_layer_name=layer_1,
            ),
        }
        # suppress var not used error
        assert fwd_context is not None


def test_init_kv_cache_without_kv_sharing(default_vllm_config):
    torch.set_default_dtype(torch.float16)
    layer_0 = "model.layers.0.self_attn.attn"
    layer_1 = "model.layers.1.self_attn.attn"
    vllm_config = get_vllm_config()
    with set_current_vllm_config(vllm_config):
        fwd_context = {
            layer_0: Attention(
                num_heads=8,
                head_size=64,
                scale=1.0,
                prefix=layer_0,
            ),
            layer_1: Attention(
                num_heads=8,
                head_size=64,
                scale=1.0,
                prefix=layer_1,
            ),
        }
        # suppress var not used error
        assert fwd_context is not None
    # Set high context length to test max context length estimation
    vllm_config.model_config.max_model_len = 3_000_000
    vllm_ctx = vllm_config.compilation_config.static_forward_context
    runner = GPUModelRunner(vllm_config, DEVICE_TYPE)
    kv_cache_spec = runner.get_kv_cache_spec()
    assert len(kv_cache_spec) == 2
    assert len(runner.shared_kv_cache_layers) == 0

    available_memory = 20 * GiB_bytes
    # page size for layer 0's kv_cache_spec is 32KB
    num_expected_blocks = 327680  # 20GB / 32KB / 2 (num layers)
    kv_cache_config = get_kv_cache_configs(
        vllm_config, [kv_cache_spec], [available_memory]
    )[0]
    assert kv_cache_config.num_blocks == num_expected_blocks
    assert len(kv_cache_config.kv_cache_tensors) == 1
    assert kv_cache_config.kv_cache_tensors[0].size == available_memory

    max_context_len = estimate_max_model_len(vllm_config, kv_cache_spec, 5 * GiB_bytes)
    # max context len with KV sharing should be 2x as large as without
    assert max_context_len == 1310720

    # important: override tensor size to prevent large mem alloc during test
    # this will only allocate 1 block worth of memory per layer (2 layers * 32kb)
    kv_cache_config.num_blocks = 1
    for kv_cache_tensor in kv_cache_config.kv_cache_tensors:
        page_size = kv_cache_spec[kv_cache_tensor.layers[0]].page_size_bytes
        kv_cache_tensor.size = page_size * len(kv_cache_tensor.layers)
        # One block per layer: a layer's region is exactly one page.
        kv_cache_tensor.layer_stride = page_size
        kv_cache_tensor.block_stride = page_size

    runner.initialize_kv_cache(kv_cache_config)

    layer_0_kv = vllm_ctx[layer_0].kv_cache
    layer_1_kv = vllm_ctx[layer_1].kv_cache
    # check layer 1 kv cache does NOT share memory with layer 0
    assert id(layer_1_kv) != id(layer_0_kv)

    # check layer 1 added to kv cache group's layer names
    assert len(kv_cache_config.kv_cache_groups) == 1
    assert len(kv_cache_config.kv_cache_groups[0].layer_names) == 2
    assert kv_cache_config.kv_cache_groups[0].layer_names[0] == layer_0
    assert kv_cache_config.kv_cache_groups[0].layer_names[1] == layer_1


def test_init_kv_cache_with_kv_sharing_valid(default_vllm_config):
    torch.set_default_dtype(torch.float16)
    layer_0 = "model.layers.0.self_attn.attn"
    layer_1 = "model.layers.1.self_attn.attn"
    vllm_config = get_vllm_config()
    with set_current_vllm_config(vllm_config):
        fwd_context = {
            layer_0: Attention(
                num_heads=8,
                head_size=64,
                scale=1.0,
                prefix=layer_0,
            ),
            layer_1: Attention(
                num_heads=8,
                head_size=64,
                scale=1.0,
                prefix=layer_1,
                kv_sharing_target_layer_name="model.layers.0.self_attn.attn",
            ),
        }
        # suppress var not used error
        assert fwd_context is not None
    # Set high context length to test max context length estimation
    vllm_config.model_config.max_model_len = 3_000_000
    vllm_ctx = vllm_config.compilation_config.static_forward_context
    runner = GPUModelRunner(vllm_config, DEVICE_TYPE)
    kv_cache_spec = runner.get_kv_cache_spec()
    assert len(kv_cache_spec) == 1
    assert layer_0 in kv_cache_spec
    assert runner.shared_kv_cache_layers[layer_1] == layer_0

    available_memory = 20 * GiB_bytes
    # page size for layer 0's kv_cache_spec is 32KB
    # with KV sharing, we can allocate (available_mem//page_size//1) blocks
    # which is twice as many as without KV sharing
    num_expected_blocks = 655360  # 20GB / 32KB
    kv_cache_config = get_kv_cache_configs(
        vllm_config, [kv_cache_spec], [available_memory]
    )[0]
    assert kv_cache_config.num_blocks == num_expected_blocks
    assert len(kv_cache_config.kv_cache_tensors) == 1
    # Each layer now has twice the available memory for KV cache
    # compared to no KV sharing
    assert kv_cache_config.kv_cache_tensors[0].size == available_memory

    max_context_len = estimate_max_model_len(vllm_config, kv_cache_spec, 5 * GiB_bytes)
    # max context len with KV sharing should be 2x as large as without
    assert max_context_len == 2 * 1310720

    # important: override tensor size to prevent large mem alloc during test
    # this will only allocate 1 block worth of memory (32kb)
    kv_cache_config.num_blocks = 1
    page_size = kv_cache_spec[layer_0].page_size_bytes
    kv_cache_config.kv_cache_tensors[0].size = page_size
    kv_cache_config.kv_cache_tensors[0].layer_stride = page_size

    runner.initialize_kv_cache(kv_cache_config)
    kv_cache_config_after_init = runner.kv_cache_config

    layer_0_kv = vllm_ctx[layer_0].kv_cache
    layer_1_kv = vllm_ctx[layer_1].kv_cache
    # check layer 1 kv cache shares memory with layer 0
    assert id(layer_1_kv) == id(layer_0_kv)

    # check layer 1 added to kv cache group's layer names
    assert len(kv_cache_config_after_init.kv_cache_groups) == 1
    assert len(kv_cache_config_after_init.kv_cache_groups[0].layer_names) == 2
    assert kv_cache_config_after_init.kv_cache_groups[0].layer_names[0] == layer_0
    assert kv_cache_config_after_init.kv_cache_groups[0].layer_names[1] == layer_1


@pytest.mark.skipif(
    not current_platform.is_cuda(),
    reason="Attention backend FLASHINFER is only supported on CUDA.",
)
def test_hybrid_attention_mamba_tensor_shapes():
    """The GPU model runner creates different views into the
    KVCacheTensors for the attention and mamba layers
    (via _allocate_kv_caches). This test verifies
    that the views are compatible: writing a mamba block
    will not corrupt an attention block and vice versa
    """
    set_random_seed(42)

    update_environment_variables(
        {
            "RANK": "0",
            "LOCAL_RANK": "0",
            "WORLD_SIZE": "1",
            "MASTER_ADDR": "localhost",
            "MASTER_PORT": "12345",
        }
    )
    from tests.utils import ensure_current_vllm_config

    with ensure_current_vllm_config():
        init_distributed_environment()
        initialize_model_parallel(tensor_model_parallel_size=1)
    torch.set_default_dtype(torch.float16)

    model_config = ModelConfig(
        model="ibm-granite/granite-4.0-tiny-preview",
        dtype="float16",
    )
    scheduler_config = SchedulerConfig(
        max_num_seqs=10,
        max_num_batched_tokens=512,
        max_model_len=512,
        is_encoder_decoder=model_config.is_encoder_decoder,
    )
    cache_config = CacheConfig(
        block_size=BLOCK_SIZE,
        gpu_memory_utilization=0.9,
        cache_dtype="auto",
    )
    cache_config.kv_cache_layout = "LBNHC"
    parallel_config = ParallelConfig()
    attention_config = AttentionConfig(backend=AttentionBackendEnum.FLASHINFER)
    vllm_config = VllmConfig(
        model_config=model_config,
        cache_config=cache_config,
        scheduler_config=scheduler_config,
        parallel_config=parallel_config,
        attention_config=attention_config,
    )

    layer_0 = "model.layers.0.self_attn.attn"
    layer_1 = "model.layers.1.self_attn.attn"
    layer_2 = "model.layers.2.mixer"
    layer_3 = "model.layers.3.mixer"
    layer_4 = "model.layers.4.mixer"
    layer_5 = "model.layers.5.mixer"

    with set_current_vllm_config(vllm_config):
        hf_config = vllm_config.model_config.hf_config
        fwd_context = {}
        for key in [layer_0, layer_1]:
            fwd_context[key] = Attention(
                num_heads=model_config.get_num_attention_heads(parallel_config),
                num_kv_heads=model_config.get_num_kv_heads(parallel_config),
                head_size=model_config.get_head_size(),
                scale=1.0,
                prefix=key,
            )
        for key in [layer_2, layer_3, layer_4, layer_5]:
            fwd_context[key] = MambaMixer2(
                hidden_size=hf_config.hidden_size,
                ssm_state_size=hf_config.mamba_d_state,
                conv_kernel_size=hf_config.mamba_d_conv,
                intermediate_size=hf_config.mamba_expand * hf_config.hidden_size,
                use_conv_bias=hf_config.mamba_conv_bias,
                use_bias=hf_config.mamba_proj_bias,
                n_groups=hf_config.mamba_n_groups,
                num_heads=hf_config.mamba_n_heads,
                head_dim=hf_config.mamba_d_head,
                rms_norm_eps=hf_config.rms_norm_eps,
                activation=hf_config.hidden_act,
                cache_config=cache_config,
                model_config=model_config,
                prefix=key,
            )
        # suppress var not used error
        assert fwd_context is not None
        vllm_ctx = vllm_config.compilation_config.static_forward_context

        runner = GPUModelRunner(vllm_config, DEVICE_TYPE)
        current_platform.update_block_size_for_backend(vllm_config)
        kv_cache_spec = runner.get_kv_cache_spec()

        available_memory = 5 * GiB_bytes
        kv_cache_config = get_kv_cache_configs(
            vllm_config, [kv_cache_spec], [available_memory]
        )[0]
        runner.initialize_kv_cache(kv_cache_config)

    # random partition of blocks
    # blocks0 will be assigned to attention layers
    # blocks1 will be assigned to mamba layers
    num_blocks = kv_cache_config.num_blocks
    ind = np.arange(num_blocks)
    np.random.shuffle(ind)
    blocks0, blocks1 = ind[: (num_blocks // 2)], ind[(num_blocks // 2) :]

    attn_shape = vllm_ctx[layer_0].kv_cache.shape
    conv_shape = vllm_ctx[layer_2].kv_cache[0].shape
    ssm_shape = vllm_ctx[layer_2].kv_cache[1].shape

    # assert we are using FlashInfer
    assert attn_shape[0] % num_blocks == 0
    block_split_ratio = attn_shape[0] // num_blocks

    # use small blocks for testing to avoid memory issues
    test_block_size = min(2, len(blocks0), len(blocks1))

    # use non-overlapping blocks to avoid data contamination
    # Split kernel blocks: first half for attention, second half for mamba
    mid_point = num_blocks // 2

    # attention uses kernel blocks from first half (mapped to logical blocks)
    kv_blocks_for_attention = np.array([0, 1])[:test_block_size]

    # mamba uses kernel blocks from second half
    kv_blocks_for_mamba = np.array([mid_point, mid_point + 1])[:test_block_size]

    # create small constant tensors for testing with corrected shapes
    # attention: [block_size, ...] starting from dimension 2
    attn_constant_shape = attn_shape[2:]
    conv_constant_shape = conv_shape[1:]
    ssm_constant_shape = ssm_shape[1:]

    attn_blocks_constant = torch.full(
        (test_block_size, *attn_constant_shape), device=DEVICE_TYPE, fill_value=3.33
    )
    conv_blocks_constant = torch.full(
        (test_block_size, *conv_constant_shape), device=DEVICE_TYPE, fill_value=6.66
    )
    ssm_blocks_constant = torch.full(
        (test_block_size, *ssm_constant_shape), device=DEVICE_TYPE, fill_value=9.99
    )

    # Fill attention blocks with constants using kv block indices
    kernel_blocks_for_attention = kv_blocks_for_attention * block_split_ratio

    for layer in [layer_0, layer_1]:
        # attention: kv_cache[kernel_block_idx, kv_idx, ...]
        for i, kernel_block in enumerate(kernel_blocks_for_attention):
            vllm_ctx[layer].kv_cache[kernel_block, :] = attn_blocks_constant[i]

    # fill mamba blocks with constants using kernel block indices
    for layer in [layer_2, layer_3, layer_4, layer_5]:
        # mamba: kv_cache[component][kernel_block_idx, ...]
        for i, kv_block in enumerate(kv_blocks_for_mamba):
            vllm_ctx[layer].kv_cache[0][kv_block, :] = conv_blocks_constant[i]
            vllm_ctx[layer].kv_cache[1][kv_block, :] = ssm_blocks_constant[i]

    # verify attention and mamba contents are correct
    for layer in [layer_0, layer_1]:
        for i, kernel_block in enumerate(kernel_blocks_for_attention):
            actual_kv = vllm_ctx[layer].kv_cache[kernel_block, :]
            expected = attn_blocks_constant[i]

            # Packed layout: (num_kv_heads, block_size, 2*head_size). Every
            # head in the block was filled with the same constant.
            for head_idx in range(actual_kv.shape[0]):
                assert torch.equal(actual_kv[head_idx], expected)

    for layer in [layer_2, layer_3, layer_4, layer_5]:
        for i, kv_block in enumerate(kv_blocks_for_mamba):
            actual_conv = vllm_ctx[layer].kv_cache[0][kv_block, :]
            actual_ssm = vllm_ctx[layer].kv_cache[1][kv_block, :]
            expected_conv = conv_blocks_constant[i]
            expected_ssm = ssm_blocks_constant[i]

            assert torch.equal(actual_conv, expected_conv)
            assert torch.equal(actual_ssm, expected_ssm)

    for layer in [layer_2, layer_3, layer_4, layer_5]:
        for i, kv_block in enumerate(kv_blocks_for_mamba):
            actual_conv = vllm_ctx[layer].kv_cache[0][kv_block, :]
            actual_ssm = vllm_ctx[layer].kv_cache[1][kv_block, :]
            expected_conv = conv_blocks_constant[i]
            expected_ssm = ssm_blocks_constant[i]
            assert torch.equal(actual_conv, expected_conv)
            assert torch.equal(actual_ssm, expected_ssm)


def test_input_batch_reinitialized_after_late_interleave_adjustment(monkeypatch):
    runner = object.__new__(GPUModelRunner)
    runner.vllm_config = SimpleNamespace(reasoning_config=None)
    runner.parallel_config = SimpleNamespace(cp_kv_cache_interleave_size=16)
    runner.cache_config = SimpleNamespace(use_replayssm=False)
    runner.model_config = SimpleNamespace(get_vocab_size=lambda: 32)
    runner.max_model_len = 64
    runner.max_encoder_len = 0
    runner.max_num_reqs = 1
    runner.max_num_tokens = 64
    runner.num_spec_tokens = 0
    runner.device = torch.device("cpu")
    runner.is_pooling_model = False
    runner._init_block_sizes = [16]
    runner._init_max_num_blocks = [4]
    runner._init_slot_mapping_modes = [
        gpu_model_runner_module.SlotMappingMode.TOKEN_TO_KV_SLOT
    ]
    runner.cp_kv_cache_interleave_size = 1
    runner.input_batch = SimpleNamespace(
        logitsprocs=None,
        logitsprocs_need_output_token_ids=False,
    )
    runner.jit_warmup_registry = Mock()
    runner.jit_warmup_registry.activate.return_value = nullcontext()

    spec = SimpleNamespace(
        block_size=16,
        max_num_blocks_per_req=lambda *_: 4,
    )
    kv_cache_config = SimpleNamespace(
        kv_cache_groups=[SimpleNamespace(kv_cache_spec=spec)]
    )
    input_batch_cls = Mock(return_value=SimpleNamespace())
    monkeypatch.setattr(gpu_model_runner_module, "InputBatch", input_batch_cls)
    monkeypatch.setattr(
        gpu_model_runner_module,
        "get_kv_cache_spec_kind",
        lambda _: gpu_model_runner_module.KVCacheSpecKind.FULL_ATTENTION,
    )

    runner.may_reinitialize_input_batch(kv_cache_config)

    assert input_batch_cls.call_count == 1
    assert input_batch_cls.call_args.kwargs["cp_kv_cache_interleave_size"] == 16


def test_v2_runner_snapshots_late_interleave_adjustment(monkeypatch):
    from vllm.v1.worker.gpu import model_runner as v2_model_runner_module

    runner = object.__new__(v2_model_runner_module.GPUModelRunner)
    runner.parallel_config = SimpleNamespace(cp_kv_cache_interleave_size=16)
    runner.cp_interleave = 1

    class StopInitialization(Exception):
        pass

    monkeypatch.setattr(
        v2_model_runner_module,
        "deepcopy",
        Mock(side_effect=StopInitialization),
    )

    with pytest.raises(StopInitialization):
        runner.initialize_kv_cache(SimpleNamespace())

    assert runner.cp_interleave == 16


def test_get_block_table_width_aligns_to_128_tokens():
    assert get_block_table_width(1875, 64) == 1876


def test_mamba_state_table_width_is_not_aligned():
    block_tables = MultiGroupBlockTable(
        max_num_reqs=1,
        max_num_batched_tokens=1,
        pin_memory=False,
        device=torch.device("cpu"),
        block_sizes=[39664],
        max_num_blocks=[1],
        slot_mapping_modes=[SlotMappingMode.NONE],
    )

    assert block_tables[0].max_num_blocks_per_req == 1


def test_hybrid_cache_integration(default_vllm_config, dist_init):
    """Test hybrid cache architecture integration with GPUModelRunner."""
    # Create a new model runner with hybrid cache configuration
    vllm_config = get_vllm_config()

    # Configure hybrid cache with different kvcache_manager block size
    vllm_config.cache_config.block_size = 32

    model_config = vllm_config.model_config
    num_heads = model_config.get_num_kv_heads(vllm_config.parallel_config)
    head_size = model_config.get_head_size()
    vllm_config.compilation_config.static_forward_context["layer.0"] = Attention(
        num_heads, head_size, 0.1
    )

    runner = GPUModelRunner(vllm_config, DEVICE_TYPE)

    # Initialize KV cache with configuration
    attn_spec = FullAttentionSpec(
        block_size=16,  # Use kernel block size directly
        num_kv_heads=runner.model_config.get_num_kv_heads(runner.parallel_config),
        head_size=runner.model_config.get_head_size(),
        dtype=runner.kv_cache_dtype,
    )
    tensor_size = attn_spec.page_size_bytes * NUM_BLOCKS
    kv_cache_config = KVCacheConfig(
        num_blocks=NUM_BLOCKS,
        kv_cache_tensors=[
            KVCacheTensor(
                size=tensor_size,
                layers=["layer.0"],
                layer_stride=tensor_size,
                block_stride=attn_spec.page_size_bytes,
            ),
        ],
        kv_cache_groups=[
            KVCacheGroupSpec(layer_names=["layer.0"], kv_cache_spec=attn_spec)
        ],
    )
    runner.kv_cache_config = kv_cache_config

    # Initialize input batch with kernel block sizes
    runner.input_batch = InputBatch(
        max_num_reqs=runner.max_num_reqs,
        max_model_len=runner.max_model_len,
        max_num_batched_tokens=runner.max_num_tokens,
        device=runner.device,
        vocab_size=runner.model_config.get_vocab_size(),
        block_sizes=[kv_cache_config.kv_cache_groups[0].kv_cache_spec.block_size],
        max_num_blocks_per_req=[NUM_BLOCKS],
    )

    runner.initialize_attn_backend(kv_cache_config)

    # Verify hybrid block table configuration
    block_table = runner.input_batch.block_table.block_tables[0]
    assert block_table.block_size == (
        kv_cache_config.kv_cache_groups[0].kv_cache_spec.block_size
    )

    # Test request processing with hybrid blocks
    req_id = "hybrid_req_0"
    scheduler_output = _schedule_new_request(req_id)

    # Update states should work with hybrid blocks
    runner._update_states(scheduler_output)
    assert _is_req_scheduled(runner, req_id)
    assert _is_req_state_block_table_match(runner, req_id)


def test_is_uniform_decode() -> None:
    # Normal
    assert GPUModelRunner._is_uniform_decode(
        max_num_scheduled_tokens=1,
        uniform_decode_query_len=1,
        num_tokens=16,
        num_reqs=16,
    )
    assert not GPUModelRunner._is_uniform_decode(
        max_num_scheduled_tokens=2,
        uniform_decode_query_len=1,
        num_tokens=16,
        num_reqs=16,
    )
    assert not GPUModelRunner._is_uniform_decode(
        max_num_scheduled_tokens=1,
        uniform_decode_query_len=1,
        num_tokens=16,
        num_reqs=15,
    )
    # Spec decoding
    assert GPUModelRunner._is_uniform_decode(
        max_num_scheduled_tokens=5,
        uniform_decode_query_len=5,
        num_tokens=30,
        num_reqs=6,
    )
    assert not GPUModelRunner._is_uniform_decode(
        max_num_scheduled_tokens=5,
        uniform_decode_query_len=4,
        num_tokens=30,
        num_reqs=6,
    )
    assert not GPUModelRunner._is_uniform_decode(
        max_num_scheduled_tokens=5,
        uniform_decode_query_len=5,
        num_tokens=30,
        num_reqs=7,
    )
    # Force uniform decode
    assert GPUModelRunner._is_uniform_decode(
        max_num_scheduled_tokens=1,
        uniform_decode_query_len=1,
        num_tokens=16,
        num_reqs=16,
        force_uniform_decode=True,
    )
    assert GPUModelRunner._is_uniform_decode(
        max_num_scheduled_tokens=2,
        uniform_decode_query_len=1,
        num_tokens=16,
        num_reqs=16,
        force_uniform_decode=True,
    )
    assert GPUModelRunner._is_uniform_decode(
        max_num_scheduled_tokens=1,
        uniform_decode_query_len=1,
        num_tokens=16,
        num_reqs=15,
        force_uniform_decode=True,
    )
    assert not GPUModelRunner._is_uniform_decode(
        max_num_scheduled_tokens=1,
        uniform_decode_query_len=1,
        num_tokens=16,
        num_reqs=16,
        force_uniform_decode=False,
    )
    assert not GPUModelRunner._is_uniform_decode(
        max_num_scheduled_tokens=2,
        uniform_decode_query_len=1,
        num_tokens=16,
        num_reqs=16,
        force_uniform_decode=False,
    )
    assert not GPUModelRunner._is_uniform_decode(
        max_num_scheduled_tokens=1,
        uniform_decode_query_len=1,
        num_tokens=16,
        num_reqs=15,
        force_uniform_decode=False,
    )


@pytest.mark.skipif(
    not current_platform.is_cuda(),
    reason="Attention backend FLASHINFER is only supported on CUDA.",
)
def test_mamba_cache_raises_when_max_num_seqs_exceeds_blocks():
    """Test that a ValueError is raised when max_num_seqs exceeds the
    available Mamba cache blocks for hybrid models with FULL cudagraphs.

    See: https://github.com/vllm-project/vllm/issues/34094
    """
    set_random_seed(42)

    update_environment_variables(
        {
            "RANK": "0",
            "LOCAL_RANK": "0",
            "WORLD_SIZE": "1",
            "MASTER_ADDR": "localhost",
            "MASTER_PORT": "12345",
        }
    )
    from tests.utils import ensure_current_vllm_config

    with ensure_current_vllm_config():
        init_distributed_environment()
        initialize_model_parallel(tensor_model_parallel_size=1)
    torch.set_default_dtype(torch.float16)

    model_config = ModelConfig(
        model="ibm-granite/granite-4.0-tiny-preview",
        dtype="float16",
    )
    scheduler_config = SchedulerConfig(
        max_num_seqs=10,
        max_num_batched_tokens=512,
        max_model_len=512,
        is_encoder_decoder=model_config.is_encoder_decoder,
    )
    cache_config = CacheConfig(
        block_size=BLOCK_SIZE,
        gpu_memory_utilization=0.9,
        cache_dtype="auto",
    )
    cache_config.kv_cache_layout = "LBNHC"
    parallel_config = ParallelConfig()
    attention_config = AttentionConfig(backend=AttentionBackendEnum.FLASHINFER)
    vllm_config = VllmConfig(
        model_config=model_config,
        cache_config=cache_config,
        scheduler_config=scheduler_config,
        parallel_config=parallel_config,
        attention_config=attention_config,
    )

    with set_current_vllm_config(vllm_config):
        hf_config = vllm_config.model_config.hf_config
        fwd_context = {}
        for key in ["model.layers.0.self_attn.attn", "model.layers.1.self_attn.attn"]:
            fwd_context[key] = Attention(
                num_heads=model_config.get_num_attention_heads(parallel_config),
                num_kv_heads=model_config.get_num_kv_heads(parallel_config),
                head_size=model_config.get_head_size(),
                scale=1.0,
                prefix=key,
            )
        for key in [
            "model.layers.2.mixer",
            "model.layers.3.mixer",
            "model.layers.4.mixer",
            "model.layers.5.mixer",
        ]:
            fwd_context[key] = MambaMixer2(
                hidden_size=hf_config.hidden_size,
                ssm_state_size=hf_config.mamba_d_state,
                conv_kernel_size=hf_config.mamba_d_conv,
                intermediate_size=hf_config.mamba_expand * hf_config.hidden_size,
                use_conv_bias=hf_config.mamba_conv_bias,
                use_bias=hf_config.mamba_proj_bias,
                n_groups=hf_config.mamba_n_groups,
                num_heads=hf_config.mamba_n_heads,
                head_dim=hf_config.mamba_d_head,
                rms_norm_eps=hf_config.rms_norm_eps,
                activation=hf_config.hidden_act,
                cache_config=cache_config,
                model_config=model_config,
                prefix=key,
            )
        assert fwd_context is not None

        runner = GPUModelRunner(vllm_config, DEVICE_TYPE)
        current_platform.update_block_size_for_backend(vllm_config)
        kv_cache_spec = runner.get_kv_cache_spec()

        available_memory = 5 * GiB_bytes
        kv_cache_config = get_kv_cache_configs(
            vllm_config, [kv_cache_spec], [available_memory]
        )[0]
        num_blocks = kv_cache_config.num_blocks

        # Force max_num_seqs to exceed num_blocks so the check triggers.
        runner.max_num_reqs = num_blocks + 100

        with pytest.raises(ValueError, match="max_num_seqs"):
            runner.initialize_kv_cache(kv_cache_config)