| 1 | # SPDX-License-Identifier: Apache-2.0 |
| 2 | # SPDX-FileCopyrightText: Copyright contributors to the vLLM project |
| 3 | from types import SimpleNamespace |
| 4 | |
| 5 | import pytest |
| 6 | import torch |
| 7 | |
| 8 | from vllm.v1.attention.backends.utils import get_dcp_local_seq_lens |
| 9 | from vllm.v1.worker.cp_utils import should_skip_dcp_context_attention |
| 10 | |
| 11 | |
| 12 | def test_cp_checks_allow_a_replicated_draft_backend_without_dcp(monkeypatch): |
| 13 | from vllm.v1.worker import cp_utils |
| 14 | |
| 15 | config = SimpleNamespace( |
| 16 | parallel_config=SimpleNamespace( |
| 17 | prefill_context_parallel_size=1, |
| 18 | decode_context_parallel_size=4, |
| 19 | cp_kv_cache_interleave_size=16, |
| 20 | ), |
| 21 | speculative_config=object(), |
| 22 | ) |
| 23 | target = SimpleNamespace( |
| 24 | impl=SimpleNamespace( |
| 25 | dcp_world_size=4, |
| 26 | supports_mtp_with_cp_non_trivial_interleave_size=True, |
| 27 | need_to_return_lse_for_decode=True, |
| 28 | ) |
| 29 | ) |
| 30 | draft = SimpleNamespace( |
| 31 | impl=SimpleNamespace( |
| 32 | dcp_world_size=1, |
| 33 | supports_mtp_with_cp_non_trivial_interleave_size=False, |
| 34 | need_to_return_lse_for_decode=False, |
| 35 | ) |
| 36 | ) |
| 37 | layers = {"target": target, "draft": draft} |
| 38 | monkeypatch.setattr(cp_utils, "get_layers_from_vllm_config", lambda *args: layers) |
| 39 | cp_utils.check_attention_cp_compatibility(config, target_layer_names={"target"}) |
| 40 | with pytest.raises(AssertionError, match="MTP with cp_kv_cache_interleave_size"): |
| 41 | cp_utils.check_attention_cp_compatibility(config) |
| 42 | |
| 43 | |
| 44 | def test_skip_gate_only_for_zero_context(): |
| 45 | assert should_skip_dcp_context_attention(torch.zeros(3, dtype=torch.int32)) |
| 46 | assert not should_skip_dcp_context_attention( |
| 47 | torch.tensor([0, 5, 0], dtype=torch.int32) |
| 48 | ) |
| 49 | |
| 50 | |
| 51 | @pytest.mark.parametrize( |
| 52 | "dcp_world_size,interleave_size,context_len", |
| 53 | [(2, 16, 10), (4, 16, 10), (8, 16, 10), (4, 1, 2)], |
| 54 | ) |
| 55 | def test_skip_gate_rank_invariant_with_divergent_local_context( |
| 56 | dcp_world_size: int, interleave_size: int, context_len: int |
| 57 | ): |
| 58 | """Contexts shorter than a full interleave round land entirely on a |
| 59 | subset of DCP ranks, so the per-rank local context lengths diverge: |
| 60 | some ranks hold zero local context while others hold all of it. Ranks |
| 61 | with zero local context must still take the collective (non-skip) path, |
| 62 | otherwise the query all-gather in _forward_with_dcp deadlocks across |
| 63 | ranks. The skip gate must therefore depend only on the rank-invariant |
| 64 | global context lengths, never on get_dcp_local_seq_lens output. |
| 65 | """ |
| 66 | context_kv_lens = torch.tensor([context_len], dtype=torch.int32) |
| 67 | local_maxes = [ |
| 68 | int( |
| 69 | get_dcp_local_seq_lens( |
| 70 | context_kv_lens, dcp_world_size, rank, interleave_size |
| 71 | ).max() |
| 72 | ) |
| 73 | for rank in range(dcp_world_size) |
| 74 | ] |
| 75 | # Precondition: the local view diverges across ranks. |
| 76 | assert 0 in local_maxes |
| 77 | assert max(local_maxes) > 0 |
| 78 | # The batch still has context globally, so no rank may skip. |
| 79 | assert not should_skip_dcp_context_attention(context_kv_lens) |