init v0.23.0

Signed-off-by: Sun Ruoxi <sunruoxi@4paradigm.com>

tests/ut/worker/test_encoder_acl_graph.py

@@ -0,0 +1,190 @@
+from types import SimpleNamespace
+from unittest.mock import MagicMock, patch
+
+import pytest
+import torch
+from vllm.config import CompilationConfig, VllmConfig
+
+from vllm_ascend.utils import vllm_version_is
+from vllm_ascend.worker import encoder_acl_graph
+from vllm_ascend.worker.encoder_acl_graph import (
+    EncoderAclGraphManager,
+    get_encoder_forward_context,
+    get_encoder_graph_params,
+    maybe_compute_actual_seq_lengths,
+    set_encoder_graph_params,
+    update_encoder_graph_params,
+)
+
+
+def _reset_encoder_acl_graph_state() -> None:
+    encoder_acl_graph._encoder_graph_params = None
+    encoder_acl_graph._reset_encoder_forward_context()
+
+
+@pytest.fixture(autouse=True)
+def _reset_state():
+    _reset_encoder_acl_graph_state()
+    yield
+    _reset_encoder_acl_graph_state()
+
+
+@pytest.mark.parametrize(
+    "cu_seqlens, num_tokens, expected",
+    [
+        (torch.tensor([0, 4, 16], dtype=torch.int32), 8, [4, 16]),
+    ],
+)
+def test_maybe_compute_actual_seq_lengths_eager(cu_seqlens, num_tokens, expected):
+    actual_q, actual_kv = maybe_compute_actual_seq_lengths(
+        cu_seqlens,
+        num_tokens,
+        num_tokens,
+        cudagraph_mm_encoder=False,
+    )
+    assert actual_q == expected
+    assert actual_kv == expected
+
+
+def test_maybe_compute_actual_seq_lengths_eager_unequal_q_kv():
+    """Molmo-style uniform cross-attention: scale KV endpoints by kv/q ratio."""
+    cu_seqlens = torch.tensor([0, 1, 2], dtype=torch.int32)
+    actual_q, actual_kv = maybe_compute_actual_seq_lengths(
+        cu_seqlens,
+        2,
+        8,
+        cudagraph_mm_encoder=False,
+    )
+    assert actual_q == [1, 2]
+    assert actual_kv == [4, 8]
+
+
+@pytest.mark.parametrize(
+    "cu_seqlens, num_tokens, expected",
+    [
+        (torch.tensor([0, 4, 8], dtype=torch.int32), 8, [4, 8]),
+        (torch.tensor([0, 4, 16], dtype=torch.int32), 8, [4, 8]),
+    ],
+)
+def test_maybe_compute_actual_seq_lengths_graph(cu_seqlens, num_tokens, expected):
+    actual_q, actual_kv = maybe_compute_actual_seq_lengths(
+        cu_seqlens,
+        num_tokens,
+        num_tokens,
+        cudagraph_mm_encoder=True,
+    )
+    assert actual_q == expected
+    assert actual_kv == expected
+
+
+def test_update_encoder_graph_params_cu_seqlens():
+    set_encoder_graph_params([2048])
+    params = get_encoder_graph_params()
+    query = MagicMock()
+    query.shape = [8, 4, 72]
+    key = MagicMock()
+    key.shape = [8, 4, 72]
+    packed = (
+        query,
+        key,
+        MagicMock(),
+        None,
+        None,
+        128,
+        4,
+        4,
+        0.125,
+        MagicMock(),
+        MagicMock(),
+    )
+    params.handles[2048] = [1]
+    params.events[2048] = [MagicMock()]
+    params.attn_params[2048] = [packed]
+    params.workspaces[2048] = MagicMock()
+
+    ctx = get_encoder_forward_context()
+    ctx.cu_seqlens_cpu = torch.tensor([0, 4, 8], dtype=torch.int32)
+
+    captured = {}
+
+    def fake_out(**kwargs):
+        captured["actual_seq_lengths"] = kwargs["actual_seq_lengths"]
+
+    fake_fia = SimpleNamespace(out=fake_out)
+    with (
+        patch("vllm_ascend.worker.encoder_acl_graph.torch.npu.stream"),
+        patch("vllm_ascend.worker.encoder_acl_graph.torch.npu.graph_task_update_begin"),
+        patch("vllm_ascend.worker.encoder_acl_graph.torch.npu.graph_task_update_end"),
+        patch(
+            "vllm_ascend.worker.encoder_acl_graph.torch_npu.npu_fused_infer_attention_score",
+            fake_fia,
+        ),
+    ):
+        update_encoder_graph_params(MagicMock(), 2048)
+
+    assert captured["actual_seq_lengths"] == [4, 8]
+
+
+def _make_manager():
+    vllm_config = MagicMock(spec=VllmConfig)
+    vllm_config.compilation_config = CompilationConfig()
+    mm_config = MagicMock()
+    mm_config.get_limit_per_prompt.return_value = 0
+    mm_config.mm_encoder_tp_mode = "tensor"
+    vllm_config.model_config = MagicMock()
+    vllm_config.model_config.multimodal_config = mm_config
+    vllm_config.parallel_config = MagicMock()
+    vllm_config.parallel_config.tensor_parallel_size = 1
+
+    model = MagicMock()
+    model.get_encoder_cudagraph_config.return_value = MagicMock(
+        modalities=["image"],
+        buffer_keys=["cu_seqlens"],
+        out_hidden_size=64,
+        enable_dual_path_graph=False,
+        padding_logics={},
+        max_frames_per_video=1,
+    )
+    model.get_encoder_cudagraph_budget_range.return_value = (64, 2048)
+    return EncoderAclGraphManager(vllm_config, "npu", "bfloat16", model), model
+
+
+def test_capture_graph_params():
+    mgr, _ = _make_manager()
+    mgr.token_budgets = [2048]
+
+    with patch("vllm.v1.worker.encoder_cudagraph.EncoderCudaGraphManager.capture", return_value=None):
+        mgr.capture()
+
+    params = get_encoder_graph_params()
+    assert params is not None
+    assert 2048 in params.events
+
+
+def test_capture_budget_graph_npu():
+    mgr, model = _make_manager()
+    mgr.max_batch_size = 2
+    mgr.max_frames_per_batch = 0
+    capture_values = {"cu_seqlens": torch.zeros(3, dtype=torch.int32)}
+    model.prepare_encoder_cudagraph_capture_inputs.return_value = MagicMock(
+        values=capture_values,
+    )
+    model.encoder_cudagraph_forward.return_value = torch.zeros(2, 64)
+
+    fake_graph = MagicMock()
+    with (
+        patch("vllm_ascend.worker.encoder_acl_graph.torch.npu.NPUGraph", return_value=fake_graph),
+        patch("vllm_ascend.worker.encoder_acl_graph.torch.npu.graph"),
+        patch(
+            "vllm_ascend.worker.encoder_acl_graph.weak_ref_tensors",
+            side_effect=lambda tensors: tensors,
+        ),
+    ):
+        mgr._capture_budget_graph(2048)
+
+    if vllm_version_is("0.23.0"):
+        graph_meta = mgr.budget_graphs[2048]
+    else:
+        graph_meta = mgr._get_graph_set("default")[2048]
+    assert graph_meta.graph is fake_graph
+    assert graph_meta.input_buffers is capture_values